# =============================================================================
# PROJECT CHRONO - http:#projectchrono.org
#
# Copyright (c) 2014 projectchrono.org
# All rights reserved.
#
# Use of this source code is governed by a BSD-style license that can be found
# in the LICENSE file at the top level of the distribution and at
# http:#projectchrono.org/license-chrono.txt.
#
# =============================================================================
# Authors: Alessandro Tasora
# =============================================================================
#
# FEA for shells of Reissner 6-field type
#
# =============================================================================

import pychrono as chrono
import pychrono.fea as fea
import pychrono.pardisomkl as mkl
import pychrono.irrlicht as chronoirr
import os
import numpy as np
import matplotlib.pyplot as plt


# Output directory
out_dir = chrono.GetChronoOutputPath() + "FEA_SHELLS"

print( "Copyright (c) 2017 projectchrono.org")

# Create (if needed) output directory
if not os.path.isdir(out_dir):
    os.mkdir(out_dir)
if not os.path.isdir(out_dir):
	raise NameError("Error creating directory " )


# Create a Chrono::Engine physical system
my_system = chrono.ChSystemSMC()

# Create the Irrlicht visualization (open the Irrlicht device,
# bind a simple user interface, etc. etc.)
application = chronoirr.ChIrrApp(my_system, "Shells FEA", chronoirr.dimension2du(800, 600), False, True)

# Easy shortcuts to add camera, lights, logo and sky in Irrlicht scene:
application.AddTypicalLogo()
application.AddTypicalSky()
application.AddTypicalLights()
application.AddTypicalCamera(chronoirr.vector3df(0, 6.0, -10))

# Create a mesh, that is a container for groups
# of elements and their referenced nodes.
my_mesh = fea.ChMesh()

# Remember to add the mesh to the system!
my_system.Add(my_mesh)

# my_system.Set_G_acc(VNULL) or
my_mesh.SetAutomaticGravity(False)

nodePlotA = fea.ChNodeFEAxyzrot()
nodePlotB = fea.ChNodeFEAxyzrot()
nodesLoad = []

ref_X = chrono.ChFunction_Recorder()
ref_Y = chrono.ChFunction_Recorder()

load_torque = chrono.ChVectorD()
load_force = chrono.ChVectorD()

bench1 = False
bench2 = True
bench3 = False
#
# BENCHMARK n.1
#
# Add an EANS SHELL cantilever:
#

if (bench1):  # set as 'True' to execute this
    rect_thickness = 0.10
    rect_L = 10.0
    rect_W = 1
    # Create a material
    rho = 0.0
    E = 1.2e6
    nu = 0.0
    melasticity =fea.ChElasticityReissnerIsothropic(E, nu, 1.0, 0.01)
    mat = fea.ChMaterialShellReissner(melasticity)
    # In case you need also damping it would add...
    #mdamping = chrono_types::make_shared<ChDampingReissnerRayleigh>(melasticity,0.01)
    #mat = chrono_types::make_shared<ChMaterialShellReissner>(melasticity, nullptr, mdamping)
    mat.SetDensity(rho)
    # Create the nodes
    nels_L = 12
    nels_W = 1
    elarray = [fea.ChElementShellReissner4]*(nels_L * nels_W)
    nodearray = [fea.ChNodeFEAxyzrot]*((nels_L + 1) * (nels_W + 1))
    nodes_start = [fea.ChNodeFEAxyzrot]*(nels_W + 1)
    nodes_end = [fea.ChNodeFEAxyzrot]*(nels_W + 1)
    for il in range(nels_L+1) :
        for iw in range(nels_W +1):           
            # Make nodes
            nodepos = chrono.ChVectorD(rect_L * (il / nels_L), 0, rect_W * (iw / nels_W))
            noderot = chrono.ChQuaternionD(chrono.QUNIT)
            nodeframe = chrono.ChFrameD(nodepos, noderot)
            mnode = fea.ChNodeFEAxyzrot(nodeframe)
            my_mesh.AddNode(mnode)
            for i in range(3):
                mnode.GetInertia()[i,i] = 0   # approx]
            mnode.SetMass(0)
            nodearray[il * (nels_W + 1) + iw] = mnode
            if (il == 0):
                nodes_start[iw] = mnode
            if (il == nels_L):
                nodes_end[iw] = mnode
            # Make elements
            if (il > 0 and iw > 0) :
                melement = fea.ChElementShellReissner4()
                my_mesh.AddElement(melement)
                melement.SetNodes(nodearray[(il - 1) * (nels_W + 1) + (iw - 1)],
								   nodearray[(il) * (nels_W + 1) + (iw - 1)], nodearray[(il) * (nels_W + 1) + (iw)],
								   nodearray[(il - 1) * (nels_W + 1) + (iw)])

                melement.AddLayer(rect_thickness, 0 * chrono.CH_C_DEG_TO_RAD, mat)
                elarray[(il - 1) * (nels_W) + (iw - 1)] = melement
                
    nodesLoad = nodes_end
    nodePlotA = nodes_end[0]
    nodePlotB = nodes_end[-1]
    for startnode in( nodes_start):
        startnode.SetFixed(True)

    # applied load
    # load_force = chrono.ChVectorD(200000,0, 20000)
    load_force = chrono.ChVectorD(0, 4, 0)
    # load_torque = chrono.ChVectorD(0, 0, 50*CH_C_PI/3.0)

    # reference solution for (0, 4, 0) shear to plot
    ref_Y.AddPoint(0.10, 1.309)
    ref_X.AddPoint(0.40, 0.103)
    ref_Y.AddPoint(0.20, 2.493)
    ref_X.AddPoint(0.80, 0.381)
    ref_Y.AddPoint(0.30, 3.488)
    ref_X.AddPoint(1.20, 0.763)
    ref_Y.AddPoint(0.40, 4.292)
    ref_X.AddPoint(1.60, 1.184)
    ref_Y.AddPoint(0.50, 4.933)
    ref_X.AddPoint(2.00, 1.604)
    ref_Y.AddPoint(0.60, 5.444)