#!/usr/bin/python
# ------------------------------------------------------------------------------------------------------------------ #
#
#   DYNARDO example
#
#   Ten-Bar-Truss example
#
#   For background information see ../ten_bar_truss.pdf
# ------------------------------------------------------------------------------------------------------------------ #

# ------------------------------------------------------------------------------------------------------------------ #
# --- Input parameters                                                                                           --- #

# Check if Python is called outside optiSLang
from math import sqrt, atan2, cos, sin, pi
from numpy.linalg import solve
import numpy as np
if not 'OSL_REGULAR_EXECUTION' in locals():
    OSL_REGULAR_EXECUTION = False

# values to default if not existing yet
#if not OSL_REGULAR_EXECUTION:  # test run mode
area01 = 10.0  # optimal value= 7.9
area02 = 10.0  # optimal value= 0.1
area03 = 10.0  # optimal value= 3.9
area04 = 10.0  # optimal value= 8.1
area05 = 10.0  # optimal value= 0.1
area06 = 10.0  # optimal value= 0.1
area07 = 10.0  # optimal value= 5.8
area08 = 10.0  # optimal value= 5.51
area09 = 10.0  # optimal value= 3.68
area10 = 10.0  # optimal value= 0.14

# ------------------------------------------------------------------------------------------------------------------ #
# define system
L = 360
Youngs = 10000000
density = 0.100000
F1 = -100000.000
F2 = -100000.000
g = 0

areas = [area01, area02, area03, area04, area05, area06, area07, area08, area09, area10]
# ------------------------------------------------------------------------------------------------------------------ #


# Function definitions

# build up the global stiffness matrix of 1 truss element
def globStiffTruss(x1, x2, A, E, density):
    """
    Purpose:
        build up the global stiffness matrix of a linear truss element
        in the global coordinate system and its mass
    input:
        x1,x2  ...    coordinates of the 2 nodes (float vector[2])
        A      ...    area of the cross section (float scalar)
        E      ...    Youngs modulus of the material (float scalar)
        density   ...    density of the material (float scalar)
    output:
        stiff ...    global element stiffness matrix (float matrix[4,4])
        mass  ...    mass of the truss element  (float scalar)
    """
    # the length of the truss
    length = sqrt((x2[0] - x1[0]) ** 2+(x2[1] - x1[1]) ** 2)

    # the angle w.r.t. x-axis
    alpha = atan2(x2[1]-x1[1], x2[0]-x1[0])

    locStiff = locStiffTruss(length, A, E)

    trafo = trafoMatrix(alpha)

    stiff = trafo.transpose() * locStiff * trafo

    mass = A * length * density

    return stiff, mass


# build up the local stiffness matrix
def locStiffTruss(L, A, E):
    r"""
    Purpose:
        build up the local stiffness matrix of a linear truss element
        in the 2DOF local system
        \f{align*}{
         \boldsymbol{E} = E*\dfrac{A}{L}*\begin{bmatrix}
           1 & -1\\
           -1 & 1
         \end{bmatrix}
        \f}
    input:
        L ...    element lenght (float scalar)
        A ...    area of the cross section (float scalar)
        E ...    Youngs modulus of the material (float scalar)
    output:
        stiff ...    local stiffness matrix (float matrix[2,2])
    """
    stiff = E*A/L*np.matrix(((1, -1), (-1, 1)))

    return stiff


# build up transformation matrix
def trafoMatrix(alpha):
    r"""
    Purpose:
        build up transformation matrix to rotate from 2DOF local system
        to global coordinate system
        \f{align*}{
         \boldsymbol{T} = \begin{bmatrix}
           \cos(\alpha) & \sin(\alpha) & 0 & 0\\
           0 & 0 & \cos(\alpha) & \sin(\alpha)
         \end{bmatrix}
        \f}
    input:
        alpha ...    angle w.r.t. global x-axis of the truss (float scalar)
    output:
        T .   ..    transformation matrix (float matrix[2,4])
    """
    trafo = np.matrix((
        (cos(alpha), sin(alpha), 0, 0),
        (0, 0, cos(alpha), sin(alpha))
    ))

    return trafo


# assemble an element matrix to a globalDOF matrix
def assembleMatrix2D(K, k, nodes):
    """
    Purpose:
        assemble an element matrix to a globalDOF matrix
    input:
        K .   ..    globelDOF matrix(float matrix[numDOF,numDOF])
        k     ...    element matrix(float matrix[4,4])
        nodes ...    node IDs (float vector[2])
    output:
        K .   ..    globelDOF matrix(float matrix[numDOF,numDOF])
    """
    for i in range(len(nodes)):
        for j in range(len(nodes)):
            K[(2*nodes[i]-1)-1, (2*nodes[j]-1) - 1] = K[(2*nodes[i]-1)-1, (2*nodes[j]-1)-1]+k[2*i, 2*j]
            K[(2*nodes[i]-1)-1, (2*nodes[j]-1)] = K[(2*nodes[i]-1)-1, (2*nodes[j]-1)]+k[2*i, 2*j+1]
            K[(2*nodes[i]-1), (2*nodes[j]-1)-1] = K[(2*nodes[i]-1), (2*nodes[j]-1)-1]+k[2*i+1, 2*j]
            K[(2*nodes[i]-1), (2*nodes[j]-1)] = K[(2*nodes[i]-1), (2*nodes[j]-1)]+k[2*i+1, 2*j+1]
    return K


# assemble a global element vector to global DOFs
def assembleVector2D(V, v, nodes):
    """
    Purpose:
        assemble an element vector to a globalDOF vector
    input:
        V .   ..    globalDOF vector(float vector[numDOF])
        v     ...    element vector(float vector[4])
        nodes ...    node IDs (float vector[2])
    output:
        V ...    globelDOF vector(float vector[numDOF])
    """
    for i in range(len(nodes)):
        V[(2*nodes[i]-1)-1] = V[(2*nodes[i]-1)-1]+v[2*i]
        V[(2*nodes[i]-1)] = V[(2*nodes[i]-1)]+v[2*i+1]
    return V


# calculate the stress of a deformed truss
def stressTruss(x1, x2, E, u1, u2):
    """
    Purpose:
        compute the stress of a truss
    input:
        x1,x2  ...    coordinates of the 2 nodes (float vector[2])
        u1,u2  ...    displacements of the 2 nodes (float vector[2])
        E      ...    Youngs modulus of the material (float scalar)
    output:
        stress ...    stress  (float scalar)
    """
    # the length of the truss
    length = sqrt((x2[0]-x1[0])**2+(x2[1]-x1[1])**2)

    # the angle w.r.t. x-axis
    alpha = atan2(x2[1]-x1[1], x2[0]-x1[0])

    trafo = trafoMatrix(alpha)

    disp = np.matrix((u1[0], u1[1], u2[0], u2[1]))
    disp = trafo*disp.transpose()

    # the stress
    strain = (disp[1]-disp[0])/length

    # the stress
    stress = E * strain

    return stress


# calculate total mass of the truss system
def tenBarTruss(L, areas, young, F1, F2, g):
    """
    Purpose:
        calculate the total mass, the nodal displacements and the stresses
        for a Ten-Bar-Truss problem
    input:
        l     ...    grid lenght (float scalar)
        h     ...    grid height (float scalar)
        areas ...    vector of the areas (float vector[10])
        youngs ...    Youngs modulus of the material (float scalar)
        density   ...    density of the material (float scalar)
        force ...    value of the force (float scalar)
        g     ...    gravity acceleration (float scalar)
    output:
        totalMass ...    total mass of the system  (float scalar)
        globDisp  ...    global displacement of the nodes  (float vector[12])
        stresses  ...    stress values of the truss elements  (float vector[10])
    """

    # define structure
    node_coordinates = [[2*L, L], [2*L, 0], [L, L], [L, 0], [0, L], [0, 0]]
    element_incidences = [[3, 5], [1, 3], [4, 6], [2, 4], [3, 4], [1, 2], [4, 5], [3, 6], [2, 3], [1, 4]]
    restrainedDOFs = [8, 9, 10, 11]

    numNodes = len(node_coordinates)
    numElems = len(element_incidences)

    # initialize stiffness matrix and load vectors
    numDOF = 2*numNodes
    globStiff = np.zeros((numDOF, numDOF))
    deadLoad = np.zeros((numDOF))
    global_load = np.zeros(numDOF)
    totalMass = 0
    stresses = np.zeros(numElems)*np.nan
    globDisp = np.zeros(numNodes*2)*np.nan
    # test for invalid input - areas[i] < 0
    is_nan = False

    for i in range(len(element_incidences)):
        if areas[i] < 0.0:
            is_nan = True

    if is_nan == True:
        totalMass = np.nan
        max_disp = np.nan
        max_stress = np.nan
        return totalMass, globDisp.tolist(), stresses.tolist(), max_disp, max_stress

    # define the load
    global_load[3] = F2
    global_load[7] = F1

    for i in range(len(element_incidences)):
        thisNodes = element_incidences[i]
        elStiff, elmass = globStiffTruss(node_coordinates[thisNodes[0]-1], node_coordinates[thisNodes[1]-1], areas[i], young, density)
        globStiff = assembleMatrix2D(globStiff, elStiff, thisNodes)
        elDeadLoad = [0, -0.5*elmass*g, 0, -0.5*elmass*g]
        deadLoad = assembleVector2D(deadLoad, elDeadLoad, thisNodes)
        totalMass = totalMass + elmass

    # building up loadcase
    globLoad = global_load+deadLoad

    # building up availDOF matrix
    availDOF = np.identity(numDOF)
    availDOF = np.delete(availDOF, restrainedDOFs, axis=0)

    # strike out restrained DOFs
    globStiff = np.dot(np.dot(availDOF, globStiff), availDOF.T)
    globLoad = np.dot(availDOF, globLoad)

    # solve the PDE
    globDisp = solve(globStiff, globLoad)
    # transform solution to global DOFs
    globDisp = np.dot(availDOF.T, globDisp)
    # displacement matrix
    globDispMat = np.zeros((numNodes, 2))
    for i in range(numNodes):
        globDispMat[i, 0] = globDisp[i*2]
        globDispMat[i, 1] = globDisp[i*2+1]

    # evaluate stresses of the trusses
    for i in range(len(element_incidences)):
        thisNodes = element_incidences[i]
        stresses[i] = stressTruss(node_coordinates[thisNodes[0]-1], node_coordinates[thisNodes[1]-1],
                                  young, globDispMat[thisNodes[0]-1], globDispMat[thisNodes[1]-1])

    max_disp = abs(globDisp).max()
    max_stress = abs(stresses).max()

    return totalMass, globDisp.tolist(), stresses.tolist(), max_disp, max_stress


# ------------------------------------------------------------------------------------------------------------------ #
# --- responses                                                                                                  --- #
(mass, global_disp, element_stresses, max_disp, max_stress) = tenBarTruss(L, areas, Youngs, F1, F2, g)
# ------------------------------------------------------------------------------------------------------------------ #

try:
    from pyvariant import list_2_variant_vector
    globDisp = list_2_variant_vector(global_disp)
    stresses = list_2_variant_vector(element_stresses)
except:
    print("Did not export result to optiSLang")
    print("================================================================= ")
    print("mass = ", mass)
    print("----------------------------------------------------------------- ")
    print("global_disp      = %9.6f %9.6f ; %9.6f %9.6f ; %9.6f %9.6f ; %9.6f %9.6f ; %9.6f %9.6f ; %9.6f %9.6f" % (
        global_disp[0], global_disp[1], global_disp[2], global_disp[3], global_disp[4], global_disp[5], global_disp[6], global_disp[7], global_disp[8], global_disp[9], global_disp[10], global_disp[11]))
    print("----------------------------------------------------------------- ")
    print("element_stresses = %11.4e %11.4e %11.4e %11.4e %11.4e %11.4e %11.4e %11.4e %11.4e %11.4e" % (
        element_stresses[0], element_stresses[1], element_stresses[2], element_stresses[3], element_stresses[4], element_stresses[5], element_stresses[6], element_stresses[7], element_stresses[8], element_stresses[9]))
    print("----------------------------------------------------------------- ")
    print("max_disp = ", max_disp)
    print("----------------------------------------------------------------- ")
    print("max_stress = %11.4e" % (max_stress))
    print("================================================================= ")
