import panel as pn
import scipy.optimize
import scipy.linalg
import pandas as pd
import numpy as np
from matplotlib.figure import Figure

x_points = [0.] * 13
d_points = [0.] * 13

def elm_k(L, d1, d2, alpha):
    b = (d2 - d1) / L
    return 2 * alpha / (35*L**3) * np.array(
        [
            [
                6*(11*L**4*b**4 + 49*L**3*b**3*d1 + 84*L**2*b**2*d1**2 + 70*L*b*d1**3 + 35*d1**4),
                L*(19*L**4*b**4 + 84*L**3*b**3*d1 + 147*L**2*b**2*d1**2 + 140*L*b*d1**3 + 105*d1**4),
                6*(-11*L**4*b**4 - 49*L**3*b**3*d1 - 84*L**2*b**2*d1**2 - 70*L*b*d1**3 - 35*d1**4),
                L*(47*L**4*b**4 + 210*L**3*b**3*d1 + 357*L**2*b**2*d1**2 + 280*L*b*d1**3 + 105*d1**4)
            ],
            [
                L*(19*L**4*b**4 + 84*L**3*b**3*d1 + 147*L**2*b**2*d1**2 + 140*L*b*d1**3 + 105*d1**4),
                2*L**2*(3*L**4*b**4 + 14*L**3*b**3*d1 + 28*L**2*b**2*d1**2 + 35*L*b*d1**3 + 35*d1**4),
                L*(-19*L**4*b**4 - 84*L**3*b**3*d1 - 147*L**2*b**2*d1**2 - 140*L*b*d1**3 - 105*d1**4), 
                L**2*(13*L**4*b**4 + 56*L**3*b**3*d1 + 91*L**2*b**2*d1**2 + 70*L*b*d1**3 + 35*d1**4)
            ],
            [
                6*(-11*L**4*b**4 - 49*L**3*b**3*d1 - 84*L**2*b**2*d1**2 - 70*L*b*d1**3 - 35*d1**4),
                L*(-19*L**4*b**4 - 84*L**3*b**3*d1 - 147*L**2*b**2*d1**2 - 140*L*b*d1**3 - 105*d1**4),
                6*(11*L**4*b**4 + 49*L**3*b**3*d1 + 84*L**2*b**2*d1**2 + 70*L*b*d1**3 + 35*d1**4),
                L*(-47*L**4*b**4 - 210*L**3*b**3*d1 - 357*L**2*b**2*d1**2 - 280*L*b*d1**3 - 105*d1**4)
            ],
            [
                L*(47*L**4*b**4 + 210*L**3*b**3*d1 + 357*L**2*b**2*d1**2 + 280*L*b*d1**3 + 105*d1**4),
                L**2*(13*L**4*b**4 + 56*L**3*b**3*d1 + 91*L**2*b**2*d1**2 + 70*L*b*d1**3 + 35*d1**4),
                L*(-47*L**4*b**4 - 210*L**3*b**3*d1 - 357*L**2*b**2*d1**2 - 280*L*b*d1**3 - 105*d1**4),
                2*L**2*(17*L**4*b**4 + 77*L**3*b**3*d1 + 133*L**2*b**2*d1**2 + 105*L*b*d1**3 + 35*d1**4)
            ]
        ]
    )

lengthInput = pn.widgets.FloatInput(name="Bow Length (mm)",
    value=700, step=1, start=600, end=750)
lengthConstInput = pn.widgets.FloatInput(name="Length Constant Diameter/Width (mm)",
    value=110, step=1, start=60, end=200)
dButtInput = pn.widgets.FloatInput(name="Diameter/Width at Butt (mm)",
    value=8.6, step=0.1, start=6, end=12)
dTipInput = pn.widgets.FloatInput(name="Diameter/Width at Tip (mm)",
    value=5.6, step=0.1, start=3, end=10)
shapeInput = pn.widgets.Select(name="Shape", options={
    "Round": 0.0490874,
    "Octagonal": 0.0547379
})
modulusInput = pn.widgets.FloatInput(name="Elastic Modulus (GPa)",
    value=30, step=0.1, start=10, end=40
)
taperOutput = pn.pane.DataFrame(
    pd.DataFrame(data = {"x": x_points, "d": d_points}),
    border=3, width=400, col_space=75,
    justify="left",
    formatters=[
        lambda x: f"{x:.0f}",
        lambda x: f"{x:.2f}"
    ]
)
deflectionOutput = pn.pane.Str("")
fig0 = Figure(figsize=(6, 4))
plotOutput = pn.pane.Matplotlib(fig0, dpi=144)

calculateButton = pn.widgets.Button(name="Calculate", width=300)

def calc_taper():
    length = lengthInput.value
    length_constant = lengthConstInput.value
    d_butt = dButtInput.value
    d_head = dTipInput.value

    r = scipy.optimize.root(
        lambda r: length_constant * (r**12 - 1) / (r - 1) - length,
        22
    ).x[0]

    for i in range(13):
        if i == 0:
            d_points[i] = d_butt
        else:
            x_points[i] = length_constant * (r**i - 1) / (r - 1)
            d_points[i] = d_butt + (d_head - d_butt) * (i - 1.) / 10.
        if i == 12:
            d_points[i] = d_head
    df = pd.DataFrame(data = {
        "x": x_points,
        "d": d_points
    })
    taperOutput.object = df

def remesh():
    x_nodes = []
    d_nodes = []

    x_s2 = x_points[11]
    x_s1 = x_s2 - 575
    x_l = x_s2 - 575 / 2

    nid_s1 = -1  # storage for node ID of support #1
    nid_s2 = -1  # storage for node ID of support #2
    nid_l = -1   # storage for node ID of load application

    tol = lambda xa, xb: abs(xa - xb) < 1e-3
    inside = lambda x, xa, xb: (x - xa) * (x - xb) < 0

    for x1, x2, d1, d2 in zip(
        x_points, x_points[1:], d_points, d_points[1:]):
        x_nodes.append(x1)
        d_nodes.append(d1)

        if tol(x_s1, x1):
            nid_s1 = len(x_nodes) - 1
        elif inside(x_s1, x1, x2):
            x_nodes.append(x_s1)
            d_nodes.append(d1 + (x_s1 - x1) / (x2 - x1) * (d2 - d1))
            nid_s1 = len(x_nodes) - 1

        if tol(x_s2, x1):
            nid_s2 = len(x_nodes) - 1
        elif inside(x_s2, x1, x2):
            x_nodes.append(x_s2)
            d_nodes.append(d1 + (x_s2 - x1) / (x2 - x1) * (d2 - d1))
            nid_s2 = len(x_nodes) - 1
        
        if tol(x_l, x1):
            nid_l = len(x_nodes) - 1
        elif inside(x_l, x1, x2):
            x_nodes.append(x_l)
            d_nodes.append(d1 + (x_l - x1) / (x2 - x1) * (d2 - d1))
            nid_l = len(x_nodes) - 1
        
    x_nodes.append(x_points[-1])
    d_nodes.append(d_points[-1])
    return x_nodes, d_nodes, nid_s1, nid_s2, nid_l

def calc_stiffness_matrix(x_nodes, d_nodes):
    k_model= np.zeros((2 * len(x_nodes), 2 * len(x_nodes)))

    modulus = modulusInput.value * 1000
    shape_constant = shapeInput.value

    for i, (x1, x2, d1, d2) in enumerate(
        zip(x_nodes, x_nodes[1:], d_nodes, d_nodes[1:])):
        # Each element connects the two adjacent nodes
        k_elm = elm_k(
            L = x2 - x1,
            d1 = d1,
            d2 = d2,
            alpha = shape_constant * modulus
        )
        for ii in range(4):
            for jj in range(4):
                k_model[i * 2 + ii, i * 2 + jj] += k_elm[ii,jj]
    return k_model

def solve(k_model, p_model, x_nodes, nid_s1, nid_s2):
    mask = [i for i, _ in enumerate(p_model)
            if i != nid_s1 * 2 and i != nid_s2 * 2]

    p_const = p_model[mask]
    k_const = k_model[mask, :]
    k_const = k_const[:, mask]

    d_const = scipy.linalg.solve(k_const, p_const)

    d_model = np.zeros(2 * len(x_nodes))
    d_model[mask] = d_const
    return d_model

def draw_deflection(x_nodes, d_model):
    fig0.clear()
    ax0 = fig0.subplots()
    ax0.clear()
    ax0.plot(x_nodes, d_model[0::2])
    ax0.grid()
    ax0.set_title("Deflection, 2 lb load, 575mm span")
    ax0.set_xlabel("x")
    ax0.set_ylabel("Vertical Deflection")
    plotOutput.object = fig0

def calc_deflection():
    x_nodes, d_nodes, nid_s1, nid_s2, nid_l = remesh()
    k_model = calc_stiffness_matrix(x_nodes, d_nodes)
    
    # Calculate the load vector
    p_model = np.zeros(2 * len(x_nodes))
    p_model[nid_l * 2] = -8.9075  # 2 lb in N

    d_model = solve(k_model, p_model, x_nodes, nid_s1, nid_s2)
    draw_deflection(x_nodes, d_model)

    deflect_load = -d_model[nid_l * 2] / 25.4 * 1000
    deflectionOutput.object = f"Deflection: {deflect_load:.3f}"

def calculate(event):
    calc_taper()
    calc_deflection()

calculateButton.on_click(calculate)

pn.Column(
    pn.FlexBox(*[lengthInput, lengthConstInput,
        dButtInput, dTipInput,
        shapeInput, modulusInput]),
    calculateButton,
    taperOutput,
    deflectionOutput,
    plotOutput
).servable(target='taper_input_panel');
