import pychrono as chrono
import pychrono.fea as fea
import pychrono.pardisomkl as mkl
import math
import numpy as np

#%% Create the Chrono system
system = chrono.ChSystemNSC()

#%% Material properties (used for all sections)
E = 210E9           # Young Modulus [N/m^2]
rho = 8500         # Density [kg/m^3]
nu = 0.3            # Poisson ratio [-]
G = E/(2*(1+nu))    # Shear Modulus [N/m^2]

#%% Define properties for all sections:
height = [0, 30, 40, 50.51, 61.01, 71.52, 82.02, 92.53, 103.03, 113.54, 124.04, 134.55, 145.63]  
# Properties of the segments (straight beam, not tapered).
D_o = [9,    9,    8.16, 7.88,  7.60,  7.325, 7.05,  6.77,  6.49,  6.21,  5.93,  5.645]    
D_i = [8.78, 8.78, 8.02, 7.745, 7.475, 7.205, 6.935, 6.665, 6.395, 6.130, 5.865, 5.585]
    
mesh = fea.ChMesh()
builder = fea.ChBuilderBeamTaperedTimoshenko()

beam_nodes = [] # To store all beam nodes
num_sections = len(height)-1
for i in range(num_sections):
    beam_section = fea.ChBeamSectionTimoshenkoAdvancedGeneric()      # Timoshenko beam formulation.
    beam_section.SetAxialRigidity(E*math.pi*((D_o[i]/2)**2-(D_i[i]/2)**2))  # EA
    beam_section.SetBendingRigidityY(E*(math.pi/4)*((D_o[i]/2)**4-(D_i[i]/2)**4)) # EIy
    beam_section.SetBendingRigidityZ(E*(math.pi/4)*((D_o[i]/2)**4-(D_i[i]/2)**4)) # EIz
    beam_section.SetTorsionRigidityX(G*(math.pi/2)*((D_o[i]/2)**4-(D_i[i]/2)**4)) # GJ
    beam_section.SetShearRigidityY(G*math.pi*((D_o[i]/2)**2-(D_i[i]/2)**2)* (10/12)) # GAk (shear correction factor)
    beam_section.SetShearRigidityZ(G*math.pi*((D_o[i]/2)**2-(D_i[i]/2)**2)* (10/12))
    
    beam_section.SetMassPerUnitLength(rho*math.pi*((D_o[i]/2)**2-(D_i[i]/2)**2)) # rho*A
    beam_section.SetDrawCircularRadius(D_o[i]/2)   # Visualization only. Radius [m]
    tapered_section = fea.ChBeamSectionTaperedTimoshenkoAdvancedGeneric()
    # For the moment, we use same section at both ends. No tapering.
    tapered_section.SetSectionA(beam_section)
    tapered_section.SetSectionB(beam_section)
    if i == 0:      # First beam
        builder.BuildBeam(mesh,
                    tapered_section,                        # Beam section properties
                    1,                                      # Elements per section  
                    chrono.ChVector3d(0, 0, 0),             # Start point
                    chrono.ChVector3d(0, 0, height[i+1]),   # End point
                    chrono.ChVector3d(0, 1, 0))
    elif i > 0:     # Successive beams. In this case we use the last node created by the last beam 
        builder.BuildBeam(mesh,
                    tapered_section,                        # Beam section properties
                    1,                                      # Elements per section 
                    builder.GetLastBeamNodes().back(),      # Start point (same as end of previous beam)
                    chrono.ChVector3d(0, 0, height[i+1]),   # End point
                    chrono.ChVector3d(0, 1, 0))

    # Storing the beam nodes:
    nodes = builder.GetLastBeamNodes()
    n_nodes = len(nodes)
    if i == 0: # First beam
        for n in nodes: # Storing all the nodes 
            beam_nodes.append(n)
    else: # Succesive beams: storing all the nodes except the first one (included in the previous beam)
        for n in nodes[1:]: # skip the first node
            beam_nodes.append(n)

#%% Boundary conditions:
first_node = beam_nodes[0]
first_node.SetFixed(True)
system.SetGravitationalAcceleration(chrono.ChVector3d(0, 0, 0)) # No gravity
last_node = beam_nodes[-1]
last_node.SetForce(chrono.ChVector3d(3E6, 0, 0)) # Force along x

#%% Solve the system:
system.Add(mesh)
solver = mkl.ChSolverPardisoMKL()
system.SetSolver(solver)
system.DoStaticNonlinear()

# Get positions of all nodes:
disp = np.zeros((len(beam_nodes), 3))
for i, node in enumerate(beam_nodes):
    pos = node.GetPos()
    disp[i, 0] = pos.x
    disp[i, 1] = pos.y
    disp[i, 2] = pos.z

print("Displaced positions (x, y, z) in meters:")
print(disp)