# =========================
# Shear & Moment Diagram Generator (Simply Supported Beam)
# + Add (max 2) & Delete controls for loads
# =========================

import math
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
import pandas as pd

plt.rcParams.update({"figure.dpi": 120})

# ---------- Core statics ----------
def reactions_simply_supported(L, point_loads, udls):
    Wp = sum(P for (_x, P) in point_loads)
    Wu = sum(w * (x2 - x1) for (x1, x2, w) in udls)
    Wtot = Wp + Wu

    Mp = sum(P * x for (x, P) in point_loads)
    Mu = sum((w * (x2 - x1)) * ((x1 + x2) / 2.0) for (x1, x2, w) in udls)

    RB = (Mp + Mu) / L
    RA = Wtot - RB
    return RA, RB

def vm_diagrams(L, point_loads, udls, n=1001):
    RA, RB = reactions_simply_supported(L, point_loads, udls)
    x = np.linspace(0.0, L, n)
    V = np.zeros_like(x, dtype=float)
    M = np.zeros_like(x, dtype=float)

    def udl_active_length(xi, x1, x2):
        if xi <= x1:
            return 0.0
        elif xi >= x2:
            return (x2 - x1)
        else:
            return (xi - x1)

    for i, xi in enumerate(x):
        v = RA
        for (xp, P) in point_loads:
            if xp <= xi + 1e-12:
                v -= P
        for (x1, x2, w) in udls:
            a = udl_active_length(xi, x1, x2)
            v -= w * max(0.0, a)
        V[i] = v

        m = RA * xi
        for (xp, P) in point_loads:
            a = max(0.0, xi - xp)
            m -= P * a
        for (x1, x2, w) in udls:
            a = udl_active_length(xi, x1, x2)
            if a > 0:
                centroid = x1 + a / 2.0
                m -= w * a * (xi - centroid)
        M[i] = m

    return x, V, M, (RA, RB)

# ---------- Plot helpers ----------
def draw_beam_sketch(L, point_loads, udls, RA, RB):
    fig, ax = plt.subplots(figsize=(7, 1.8))
    ax.set_xlim(-0.02*L, 1.02*L)
    ax.set_ylim(-1.2, 1.2)

    ax.plot([0, L], [0, 0], lw=4, color="black")
    triA = plt.Polygon([[0,0],[ -0.05*L, -0.5],[ 0.05*L, -0.5]], color="gray")
    triB = plt.Polygon([[L,0],[ L-0.05*L, -0.5],[ L+0.05*L, -0.5]], color="gray")
    ax.add_patch(triA); ax.add_patch(triB)

    ax.annotate("", xy=(0,0.6), xytext=(0,0.05),
                arrowprops=dict(arrowstyle="->", lw=2))
    ax.text(0, 0.7, f"R_A={RA:.1f} N", ha="center", va="bottom", fontsize=9)
    ax.annotate("", xy=(L,0.6), xytext=(L,0.05),
                arrowprops=dict(arrowstyle="->", lw=2))
    ax.text(L, 0.7, f"R_B={RB:.1f} N", ha="center", va="bottom", fontsize=9)

    for (xp, P) in point_loads:
        ax.annotate("", xy=(xp,-0.7), xytext=(xp,0.05),
                    arrowprops=dict(arrowstyle="->", lw=2))
        ax.text(xp, -0.8, f"P={P:.0f} N", ha="center", va="top", fontsize=9)

    for (x1, x2, w) in udls:
        ax.plot([x1, x2], [0.25, 0.25], lw=6, color="tab:blue")
        n_ar = max(2, int((x2-x1)/(L/10)) )
        xs = np.linspace(x1, x2, n_ar)
        for xk in xs:
            ax.annotate("", xy=(xk,0.05), xytext=(xk,0.4),
                        arrowprops=dict(arrowstyle="->", lw=1.5))
        ax.text((x1+x2)/2, 0.45, f"w={w:.0f} N/m", ha="center", va="bottom", fontsize=9, color="tab:blue")

    ax.axis("off")
    ax.set_title("Beam & Loads (simply supported)")
    fig.tight_layout()
    return fig

def plot_line(x, y, title, ylbl):
    fig, ax = plt.subplots(figsize=(7, 2.5))
    ax.axhline(0, color="k", lw=1)
    ax.plot(x, y, lw=2)
    ax.set_xlabel("x [m]")
    ax.set_ylabel(ylbl)
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    return fig

# ---------- Defaults & Help ----------
EX_POINT = pd.DataFrame(
    [[2.0, 8000],
     [5.0, 6000]],
    columns=["x (m)", "P (N, down +)"]
)
EX_UDL = pd.DataFrame(
    [[1.0, 3.0, 1500],
     [6.0, 8.5, 1000]],
    columns=["x1 (m)", "x2 (m)", "w (N/m, down +)"]
)

HELP_MD = """
**How to use**

1) Set **span L** and enter loads:
   - **Point loads:** `(x, P)`; `P>0` is downward.
   - **UDLs:** `(x1, x2, w)`; `w>0` is downward, active on `[x1,x2]`.

2) Use **Add** (limited to 2 rows per type), **Delete rows**, **Clear**, or **Reset examples**.

3) Click **Compute** to see reactions, the beam sketch, **V(x)** and **M(x)**, and equilibrium checks.

**Sign convention**
- Upward reactions positive. Downward loads positive inputs.
- Shear: positive upward on left face; Moment: sagging positive.
"""

# ---------- Small helpers for add/delete ----------
def _parse_row_list(s, n_rows):
    if not s:
        return []
    idxs = []
    for tok in str(s).replace(" ", "").split(","):
        if not tok:
            continue
        try:
            k = int(tok)
            if 1 <= k <= n_rows:
                idxs.append(k-1)
        except ValueError:
            pass
    return sorted(set(idxs), reverse=True)

def delete_point_rows(df_points, rows_to_delete, confirm):
    if not confirm:
        return df_points, "Deletion NOT performed (check 'Confirm delete')."
    if df_points is None or len(df_points) == 0:
        return pd.DataFrame(columns=["x (m)","P (N, down +)"]), "No point rows to delete."
    idxs = _parse_row_list(rows_to_delete, len(df_points))
    if not idxs:
        return df_points, "No valid point row indices provided."
    df = df_points.copy()
    for i in idxs:
        df = df.drop(df.index[i])
    df = df.reset_index(drop=True)
    return df, f"Deleted point rows: {', '.join(str(i+1) for i in idxs)}."

def delete_udl_rows(df_udl, rows_to_delete, confirm):
    if not confirm:
        return df_udl, "Deletion NOT performed (check 'Confirm delete')."
    if df_udl is None or len(df_udl) == 0:
        return pd.DataFrame(columns=["x1 (m)","x2 (m)","w (N/m, down +)"]), "No UDL rows to delete."
    idxs = _parse_row_list(rows_to_delete, len(df_udl))
    if not idxs:
        return df_udl, "No valid UDL row indices provided."
    df = df_udl.copy()
    for i in idxs:
        df = df.drop(df.index[i])
    df = df.reset_index(drop=True)
    return df, f"Deleted UDL rows: {', '.join(str(i+1) for i in idxs)}."

def clear_points():
    return pd.DataFrame(columns=["x (m)","P (N, down +)"]), "Cleared all point loads."

def clear_udls():
    return pd.DataFrame(columns=["x1 (m)","x2 (m)","w (N/m, down +)"]), "Cleared all UDLs."

def reset_examples_points():
    return EX_POINT.copy(), "Restored example point loads."

def reset_examples_udls():
    return EX_UDL.copy(), "Restored example UDLs."

# --- ADD with limit = 2 ---
MAX_POINTS = 2
MAX_UDLS = 2

def add_point_row(df_points, x, P):
    df = df_points.copy() if df_points is not None else pd.DataFrame(columns=["x (m)","P (N, down +)"])
    if len(df) >= MAX_POINTS:
        return df, f"Limit reached: max {MAX_POINTS} point loads."
    try:
        x = float(x); P = float(P)
    except Exception:
        return df, "Provide numeric x and P."
    if P < 0:
        return df, "Use P ≥ 0 (downward +)."
    df = pd.concat([df, pd.DataFrame([[x, P]], columns=df.columns)], ignore_index=True)
    return df, f"Added point load #{len(df)}."

def add_udl_row(df_udl, x1, x2, w):
    df = df_udl.copy() if df_udl is not None else pd.DataFrame(columns=["x1 (m)","x2 (m)","w (N/m, down +)"])
    if len(df) >= MAX_UDLS:
        return df, f"Limit reached: max {MAX_UDLS} UDLs."
    try:
        x1 = float(x1); x2 = float(x2); w = float(w)
    except Exception:
        return df, "Provide numeric x1, x2, w."
    if not (x1 < x2):
        return df, "Require x1 < x2."
    if w < 0:
        return df, "Use w ≥ 0 (downward +)."
    df = pd.concat([df, pd.DataFrame([[x1, x2, w]], columns=df.columns)], ignore_index=True)
    return df, f"Added UDL #{len(df)}."

# ---------- Core compute ----------
def run_once(L, df_points, df_udl, npts):
    try:
        L = float(L)
        n = int(npts)
        if L <= 0:
            raise ValueError("Beam length L must be > 0.")
        if n < 201:
            n = 201

        point_loads = []
        if df_points is not None and len(df_points) > 0:
            for _, row in df_points.iterrows():
                x = float(row[0]); P = float(row[1])
                if 0 <= x <= L and P >= 0:
                    point_loads.append((x, P))

        udls = []
        if df_udl is not None and len(df_udl) > 0:
            for _, row in df_udl.iterrows():
                x1 = float(row[0]); x2 = float(row[1]); w = float(row[2])
                if 0 <= x1 < x2 <= L and w >= 0:
                    udls.append((x1, x2, w))

        x, V, M, (RA, RB) = vm_diagrams(L, point_loads, udls, n=n)

        fig_beam = draw_beam_sketch(L, point_loads, udls, RA, RB)
        fig_V = plot_line(x, V, "Shear Force Diagram V(x)", "V [N]")
        fig_M = plot_line(x, M, "Bending Moment Diagram M(x)", "M [N·m]")

        total_down = sum(P for _, P in point_loads) + sum(w*(x2-x1) for (x1,x2,w) in udls)
        eq_sumF = RA + RB - total_down
        Mp = sum(P * x for (x, P) in point_loads)
        Mu = sum(w*(x2-x1)*( (x1+x2)/2 ) for (x1,x2,w) in udls)
        eq_MA = RB*L - (Mp + Mu)

        df = pd.DataFrame([{
            "RA [N]": round(RA, 3),
            "RB [N]": round(RB, 3),
            "ΣF (should ≈ 0) [N]": round(eq_sumF, 6),
            "ΣM_A (should ≈ 0) [N·m]": round(eq_MA, 6),
            "max|V| [N]": round(float(np.max(np.abs(V))), 3),
            "max|M| [N·m]": round(float(np.max(np.abs(M))), 3),
        }])

        return df, fig_beam, fig_V, fig_M, ""
    except Exception as e:
        return pd.DataFrame(), None, None, None, f"Input error:\n{e}"

# ---------- UI ----------
with gr.Blocks(title="Shear & Moment Diagrams — Simply Supported Beam") as demo:
    gr.Markdown("# Shear & Bending Moment Diagram Generator")
    gr.Markdown(HELP_MD)

    with gr.Row():
        with gr.Column():
            L = gr.Number(value=10.0, label="Span L [m]")
            npts = gr.Slider(minimum=201, maximum=5001, step=100, value=1201, label="Resolution (points)")

        with gr.Column():
            gr.Markdown("### Point loads (downward +) — max 2")
            df_points = gr.Dataframe(
                value=EX_POINT, headers=["x (m)","P (N, down +)"],
                datatype=["number","number"], row_count=(1, "dynamic")
            )
            with gr.Row():
                x_new = gr.Number(value=1.0, label="x (m)")
                P_new = gr.Number(value=5000.0, label="P (N, down +)")
                btn_add_pt = gr.Button("Add point load", variant="primary")
            with gr.Row():
                del_pts_rows = gr.Textbox(label="Delete point rows (e.g., 1,2)")
                confirm_pts = gr.Checkbox(value=False, label="Confirm delete")
            with gr.Row():
                btn_del_pts = gr.Button("Delete selected point rows", variant="secondary")
                btn_clr_pts = gr.Button("Clear ALL point loads", variant="stop")
                btn_rst_pts = gr.Button("Reset example points", variant="primary")

        with gr.Column():
            gr.Markdown("### Uniform distributed loads (downward +) — max 2")
            df_udl = gr.Dataframe(
                value=EX_UDL, headers=["x1 (m)","x2 (m)","w (N/m, down +)"],
                datatype=["number","number","number"], row_count=(1, "dynamic")
            )
            with gr.Row():
                x1_new = gr.Number(value=2.0, label="x1 (m)")
                x2_new = gr.Number(value=4.0, label="x2 (m)")
                w_new  = gr.Number(value=1000.0, label="w (N/m)")
                btn_add_udl = gr.Button("Add UDL", variant="primary")
            with gr.Row():
                del_udl_rows = gr.Textbox(label="Delete UDL rows (e.g., 1)")
                confirm_udl = gr.Checkbox(value=False, label="Confirm delete")
            with gr.Row():
                btn_del_udl = gr.Button("Delete selected UDL rows", variant="secondary")
                btn_clr_udl = gr.Button("Clear ALL UDLs", variant="stop")
                btn_rst_udl = gr.Button("Reset example UDLs", variant="primary")

    go_btn = gr.Button("Compute", variant="primary")

    gr.Markdown("### Reactions & Checks")
    out_tbl = gr.Dataframe(interactive=False)

    gr.Markdown("### Beam sketch")
    out_beam = gr.Plot()

    with gr.Row():
        out_V = gr.Plot()
        out_M = gr.Plot()

    status = gr.Textbox(label="Status / Errors", interactive=False)

    # --- Wire add actions (limit enforced) ---
    def _add_point(df, x, P):
        new_df, msg = add_point_row(df, x, P)
        return new_df, msg
    btn_add_pt.click(_add_point, inputs=[df_points, x_new, P_new], outputs=[df_points, status])

    def _add_udl(df, x1, x2, w):
        new_df, msg = add_udl_row(df, x1, x2, w)
        return new_df, msg
    btn_add_udl.click(_add_udl, inputs=[df_udl, x1_new, x2_new, w_new], outputs=[df_udl, status])

    # --- Wire delete/clear/reset actions ---
    def _delete_pts(df, rows_txt, confirm):
        new_df, msg = delete_point_rows(df, rows_txt, confirm)
        return new_df, msg
    btn_del_pts.click(_delete_pts, inputs=[df_points, del_pts_rows, confirm_pts],
                      outputs=[df_points, status])

    def _clear_pts():
        return clear_points()
    btn_clr_pts.click(_clear_pts, inputs=None, outputs=[df_points, status])

    def _reset_pts():
        return reset_examples_points()
    btn_rst_pts.click(_reset_pts, inputs=None, outputs=[df_points, status])

    def _delete_udl(df, rows_txt, confirm):
        new_df, msg = delete_udl_rows(df, rows_txt, confirm)
        return new_df, msg
    btn_del_udl.click(_delete_udl, inputs=[df_udl, del_udl_rows, confirm_udl],
                      outputs=[df_udl, status])

    def _clear_udl():
        return clear_udls()
    btn_clr_udl.click(_clear_udl, inputs=None, outputs=[df_udl, status])

    def _reset_udl():
        return reset_examples_udls()
    btn_rst_udl.click(_reset_udl, inputs=None, outputs=[df_udl, status])

    # --- Compute ---
    go_btn.click(
        fn=run_once,
        inputs=[L, df_points, df_udl, npts],
        outputs=[out_tbl, out_beam, out_V, out_M, status]
    )

if __name__ == "__main__":
    demo.launch(debug=False)