from G2AnalysisNonlinearStatic import *
from G2Model import *
import math

elementType   = 2
materialType  = 'Bilinear'  # 'Bilinear'/'Plasticity'/'BoucWen'
P             = 100e3
A             = 0.04**2
E             = 200e9
fy            = 350e6
alpha         = 0.02
eta           = 2.0
gamma         = 0.5
beta          = 0.5
Hk            = alpha*E/(1-alpha)
K             = 0.0
delta         = 0.0
fy_inf        = 100.0
dt            = 1/20
nsteps        = 25
KcalcFrequency= 1
maxIter       = 100
tol           = 1e-5

def build_crane(Hm=30.0, Lj=45.0, Lc=15.0, width=5.0, panel=5.0):

    r6 = lambda x: round(x, 6)
    nodes = []                 
    idx  = {}                  

    def add_node(x, y):
        key = (r6(x), r6(y))
        if key not in idx:
            idx[key] = len(nodes)
            nodes.append([float(x), float(y)])
        return idx[key]


    # Mast chords: x=0 and x=width; y = 0..Hm
    n_mast = int(round(Hm/panel))
    mast_left_ids  = [add_node(0.0,   i*panel) for i in range(n_mast+1)]
    mast_right_ids = [add_node(width, i*panel) for i in range(n_mast+1)]

    # Jib chords: y = Hm (top), y = Hm - width (bottom)
    n_jib = int(round(Lj/panel))
    # top chord starts at x >= width to avoid duplicate at mast-right top
    jib_top_ids = [add_node(i*panel, Hm) for i in range(n_jib+1) if i*panel >= width - 1e-9]
    # bottom chord from x = 0 .. Lj
    jib_bot_ids = [add_node(i*panel, Hm - width) for i in range(n_jib+1)]

    # Counter-jib chords at negative x: y=Hm (top), y=Hm-width (bottom)
    n_cj  = int(round(Lc/panel))
    cj_top_ids = [add_node(-i*panel, Hm)         for i in range(1, n_cj+1)]   # -panel .. -Lc
    cj_bot_ids = [add_node(-i*panel, Hm - width) for i in range(1, n_cj+1)]

    top_by_x = { r6(nodes[n][0]): n for n in jib_top_ids }
    bot_by_x = { r6(nodes[n][0]): n for n in jib_bot_ids }
    cj_top_by_x = { r6(nodes[n][0]): n for n in cj_top_ids }
    cj_bot_by_x = { r6(nodes[n][0]): n for n in cj_bot_ids }

    # Elements (0-connectivity) 
    ELEMENTS_0B = []
    def add(sec, i, j):
        if i != j:
            ELEMENTS_0B.append([elementType, sec, i, j])

    # 0: Mast chords & horizontals
    for a,b in zip(mast_left_ids[:-1], mast_left_ids[1:]):   add(0, a, b)
    for a,b in zip(mast_right_ids[:-1], mast_right_ids[1:]): add(0, a, b)
    for a,b in zip(mast_left_ids, mast_right_ids):           add(0, a, b)

    # 1: Mast diagonals (alternating)
    for k in range(len(mast_left_ids)-1):
        if k % 2 == 0: add(1, mast_left_ids[k],  mast_right_ids[k+1])
        else:          add(1, mast_right_ids[k], mast_left_ids[k+1])

    # 2: Jib chords
    for a,b in zip(jib_top_ids[:-1], jib_top_ids[1:]): add(2, a, b)
    for a,b in zip(jib_bot_ids[:-1], jib_bot_ids[1:]): add(2, a, b)

    # 2: Jib verticals (skip x=0 and x=width to avoid duplicates)
    for x,t in top_by_x.items():
        if x in bot_by_x and x > r6(width):
            add(2, t, bot_by_x[x])

    # 3: Jib diagonals
    # even i: bottom-left -> top-right ; odd i: top-left -> bottom-right
    xs = [i*panel for i in range(n_jib)]  # left edges of panels
    for i, xL in enumerate(xs):
        xR = r6(xL + panel)
        xLr = r6(xL)
        tL, tR = top_by_x.get(xLr), top_by_x.get(xR)
        bL, bR = bot_by_x.get(xLr), bot_by_x.get(xR)
        if i % 2 == 0:
            if bL is not None and tR is not None:
                add(3, bL, tR)
        else:
            if tL is not None and bR is not None:
                add(3, tL, bR)

    if bot_by_x.get(r6(0.0)) is not None and len(top_by_x) > 0:
        first_top_x = min(top_by_x.keys())
        if first_top_x > r6(width) - 1e-9 and first_top_x not in {r6(panel)}:
            add(3, bot_by_x[r6(0.0)], top_by_x[first_top_x])

    # 4: Counter-jib chords & verticals (rooted at mast-left top/bottom)
    mast_top_left  = mast_left_ids[-1]         # (0, Hm)
    mast_bot_left  = mast_left_ids[-1-1] if width > 0 else mast_left_ids[-1]
    # top chord: 0 -> -panel -> ... -> -Lc
    prev = mast_top_left
    for x in sorted(cj_top_by_x.keys(), key=lambda v: -v):  
        add(4, prev, cj_top_by_x[x]); prev = cj_top_by_x[x]
    # bottom chord
    prev = mast_left_ids[n_mast-1]  # (0, Hm-width)
    for x in sorted(cj_bot_by_x.keys(), key=lambda v: -v):
        add(4, prev, cj_bot_by_x[x]); prev = cj_bot_by_x[x]
    # verticals
    for x in cj_top_by_x:
        if x in cj_bot_by_x:
            add(4, cj_top_by_x[x], cj_bot_by_x[x])

    # 5: Counter-jib diagonals 
    cj_xs = sorted(cj_top_by_x.keys())  
    for i in range(len(cj_xs)-1):
        xL, xR = cj_xs[i], cj_xs[i+1]    
        tL, tR = cj_top_by_x[xL], cj_top_by_x[xR]
        bL, bR = cj_bot_by_x.get(xL), cj_bot_by_x.get(xR)
        if bL is None or bR is None:  
            continue
        if i % 2 == 0:
            add(5, bL, tR)   # bottom-left -> top-right
        else:
            add(5, tL, bR)   # top-left -> bottom-right

    # PATCH
    if len(cj_xs) >= 1:
        x_first = r6(-panel)  # nearest counter-jib x to the mast
        if x_first in cj_top_by_x and x_first in cj_bot_by_x:
            # The next virtual index would be i = len(cj_xs)-1
            mast_top = mast_top_left  # (0, Hm)
            # bottom root of counter-jib at the mast (same node used to start cj bottom chord)
            mast_bot_cj_root = mast_left_ids[n_mast-1]  # (0, Hm - width) when width==panel

            i_virtual = len(cj_xs) - 1
            if i_virtual % 2 == 0:
                # bottom-left (-panel, Hm-width) -> top-right (0, Hm)
                add(5, cj_bot_by_x[x_first], mast_top)
            else:
                # top-left (-panel, Hm) -> bottom-right (0, Hm-width)
                add(5, cj_top_by_x[x_first], mast_bot_cj_root)


    # fix base nodes at (0,0) and (width,0)
    CONSTRAINTS = [[0,0] for _ in nodes]
    base_L = idx[(r6(0.0),   r6(0.0))]
    base_R = idx[(r6(width), r6(0.0))]
    CONSTRAINTS[base_L] = [1,1]
    CONSTRAINTS[base_R] = [1,1]

    # Loads: load downward at jib tip (top chord, max x at y=Hm)
    LOADS = [[0.0, 0.0] for _ in nodes]
    tip_top_id = max((n for n in jib_top_ids), key=lambda n: nodes[n][0])
    LOADS[tip_top_id][1] -= P

    # Section
    nel = len(ELEMENTS_0B)
    SECTIONS  = [['Truss', A] for _ in range(nel)]
    MATERIALS = []
    for _ in range(nel):
        if materialType == 'BoucWen':
            MATERIALS.append(['BoucWen', E, fy, alpha, eta, beta, gamma])
        elif materialType == 'Bilinear':
            MATERIALS.append(['Bilinear', E, fy, alpha])
        elif materialType == 'Plasticity':
            MATERIALS.append(['Plasticity', E, fy, Hk, K, delta, fy_inf])
        else:
            raise ValueError("Unknown materialType")

    # 1-based elements 
    ELEMENTS_1B = [[etype, sec, i+1, j+1] for (etype, sec, i, j) in ELEMENTS_0B]

    return nodes, CONSTRAINTS, ELEMENTS_0B, ELEMENTS_1B, SECTIONS, MATERIALS, LOADS, tip_top_id


Hm, Lj, Lc, width, panel = 20.0, 75.0, 12.0, 5.0, 5.0
NODES, CONSTRAINTS, ELEMENTS_0B, ELEMENTS, SECTIONS, MATERIALS, LOADS, tip_id = build_crane(
    Hm=Hm, Lj=Lj, Lc=Lc, width=width, panel=panel
)

# Track Node
trackNode = tip_id              # 0-based tip node
trackDOF  = 2

# model 
MASS = [[0.0, 0.0] for _ in NODES]   
input = [NODES, CONSTRAINTS, ELEMENTS, SECTIONS, MATERIALS, LOADS, MASS]
m = model(input)

#import matplotlib.pyplot as plt
#[plt.plot([NODES[i][0], NODES[j][0]], [NODES[i][1], NODES[j][1]]) for _,_,i,j in ELEMENTS_0B]
#plt.gca().set_aspect('equal'); plt.grid(True); plt.show()

# Response with DDM sensitivities
DDMparameters = [['Element', 'E', [1]],
                 ['Element', 'fy', [1]],
                 ['Element', 'alpha', [1]],
                 ['Element', 'A', [1]]]

# Run
t, loadFactor, u, dudx, ddm2 = nonlinearStaticAnalysis(m, nsteps, dt, maxIter, KcalcFrequency, tol, trackNode, trackDOF, DDMparameters, True)
 

# Plot the DDM sensitivities
plt.ion()
plt.figure()
plt.autoscale(True)
plt.plot(t, dudx[0, :], 'k-')
plt.plot(t, dudx[1, :], 'r-')
plt.plot(t, dudx[2, :], 'b-')
plt.plot(t, dudx[3, :], 'g-')


print('\n'"Click somewhere in the plot to continue...")
plt.waitforbuttonpress()