app.py

"""Gradio app: underdamped step response calculator with an LLM explanation panel."""
from html import escape

import gradio as gr
import numpy as np
import pandas as pd
from matplotlib.figure import Figure

import calculator as C
import explainer as E
import theme as T

MODES = {"Normalised (ωn, ζ)": "normalised", "Mass-spring-damper (m, c, k)": "mass-spring-damper"}
TABLE_COLUMNS = ["Quantity", "Value", "Unit", "Formula or method"]
STALE_NOTE = "_Select **Explain these results** to generate a plain-language explanation of the numbers above._"


def plot_response(record):
    """Response curve with final value, settling band, envelope and key markers."""
    cv = record["_curve"]
    t = np.linspace(0.0, cv["t_end"], 1500)
    c = C.normalised_response(t, cv["zeta"], cv["wn"])
    y, yf, band = cv["y_final"] * c, cv["y_final"], cv["band"]
    root = (1 - cv["zeta"] ** 2) ** 0.5
    env = np.exp(-cv["zeta"] * cv["wn"] * t) / root  # decay envelope of the oscillation

    fig = Figure(figsize=(7.2, 3.6), dpi=110)
    fig.patch.set_facecolor(T.P["surface"])
    ax = fig.add_subplot(111)
    ax.set_facecolor(T.P["surface"])
    ax.axhspan(yf * (1 - band), yf * (1 + band), color=T.P["data"], alpha=0.32, lw=0,
               label=f"{record['settling_band']} band")
    ax.plot(t, yf * (1 + env), color=T.P["muted"], lw=0.9, ls=":", label="Envelope")
    ax.plot(t, yf * (1 - env), color=T.P["muted"], lw=0.9, ls=":")
    ax.axhline(yf, color=T.P["ink"], lw=0.8, ls="--", label="Final value")
    ax.plot(t, y, color=T.P["accent"], lw=2.2, label="Response")

    m = record["metrics"]
    tp, yp, ts = m["peak_time"]["value"], m["peak_value"]["value"], m["ts_exact"]["value"]
    ax.plot([tp], [yp], "o", color=T.P["accent"], ms=6)
    ax.annotate(f"Peak {C.fmt(yp)} at {C.fmt(tp)} s", (tp, yp), xytext=(8, 6),
                textcoords="offset points", fontsize=8.5, color=T.P["ink"])
    ax.axvline(ts, color=T.P["good"], lw=1.1)
    ax.annotate(f"Settled {C.fmt(ts)} s", (ts, yf * (1 - band)), xytext=(4, -14),
                textcoords="offset points", fontsize=8.5, color=T.P["good"])

    unit = f" ({record['output_unit']})" if record["output_unit"] else ""
    ax.set_xlabel("Time (s)", color=T.P["ink"])
    ax.set_ylabel(f"{record['output_name'].capitalize()}{unit}", color=T.P["ink"])
    lo = min(float(y.min()), 0.0, yf * (1 - band))    # works for negative steps too
    hi = max(float(y.max()), 0.0, yf * (1 + band))
    pad = 0.06 * (hi - lo)
    ax.set_ylim(lo - pad, hi + pad)
    ax.set_xlim(0, cv["t_end"])
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(T.P["line"])
    ax.tick_params(colors=T.P["muted"], labelsize=8.5)
    ax.legend(loc="lower right", fontsize=8, frameon=False, labelcolor=T.P["ink"])
    fig.tight_layout()
    return fig


def summary_card(record):
    """Headline numbers plus one badge per design check."""
    m = record["metrics"]
    unit = f" {record['output_unit']}" if record["output_unit"] else ""
    title = f"{C.fmt(m['overshoot']['value'])} % overshoot, settled in {C.fmt(m['ts_exact']['value'])} s"
    sub = (f"The {escape(record['output_name'])} peaks at {C.fmt(m['peak_value']['value'])}{escape(unit)} "
           f"after {C.fmt(m['peak_time']['value'])} s and settles at {C.fmt(m['final']['value'])}{escape(unit)}.")
    badges = []
    for ch in record["checks"]:
        word = "passes" if ch["result"] == "PASS" else "fails"
        badges.append(T.badge(f"{ch['name']} {word}", "good" if ch["result"] == "PASS" else "bad"))
    body = " ".join(badges) if badges else T.badge("No design limits set", "warn")
    foot = " ".join(escape(w) for w in record["warnings"]) or \
        "Closed-form values from the standard second-order model; settling time found numerically."
    return T.verdict_html(title, sub, f"<p>{body}</p>", foot)


def run_calculation(mode_label, wn, zeta, amplitude, gain, mass, damping, stiffness, force,
                    band, os_on, os_limit, ts_on, ts_limit):
    """Validate, compute, and render every numerical output. Errors clear stale results."""
    try:
        record = C.compute(
            MODES.get(mode_label, ""), wn=wn, zeta=zeta, amplitude=amplitude, gain=gain,
            mass=mass, damping=damping, stiffness=stiffness, force=force, band_name=band,
            os_limit=os_limit if os_on else None, ts_limit=ts_limit if ts_on else None,
        )
    except (C.InputError, TypeError, ValueError) as exc:
        card = T.verdict_html("Check the inputs", escape(str(exc)), empty=True)
        return (card, None, pd.DataFrame(columns=TABLE_COLUMNS), None, "", {},
                "_Fix the inputs above to see results and an explanation._", "")
    public = {k: v for k, v in record.items() if not k.startswith("_")}  # hide plotting internals
    table = pd.DataFrame(C.to_rows(record), columns=TABLE_COLUMNS)
    return (summary_card(record), plot_response(record), table, record, C.to_text(record),
            public, STALE_NOTE, "")


def run_explanation(record):
    """Ask the LLM for prose, then audit it against the record."""
    if record is None:
        return "_There are no valid results to explain yet._", ""
    text, source = E.explain(record)
    report = E.grounding_check(text, record)
    if report["grounded"]:
        status = T.badge(f"Grounded: all {report['numbers_found']} numbers match the record", "good")
    else:
        issues = [f"not in record: {', '.join(report['unsupported'])}"] if report["unsupported"] else []
        issues += report["contradictions"]
        status = T.badge("Check this explanation against the table", "bad") + \
            f'<p class="expl-note">{escape("; ".join(issues))}</p>'
    status += f'<p class="expl-note">Written by {escape(source)}. The table above is the source of truth.</p>'
    return text, status


def toggle_mode(mode_label):
    """Show only the input group that matches the chosen mode."""
    physical = MODES.get(mode_label) == "mass-spring-damper"
    return gr.update(visible=not physical), gr.update(visible=physical)


EXTRA_CSS = f"""
#explain-panel {{ border-top: 1px solid {T.P['line']}; margin-top: 8px; padding-top: 8px; }}
#explain-text {{ font-size: 16px; line-height: 1.6; max-width: 75ch; }}
.expl-note {{ color: {T.P['muted']}; font-size: 13px; margin: 6px 0 0; }}
"""

with gr.Blocks(title="Step response explainer") as demo:
    gr.HTML(T.header_html(
        "Step response explainer",
        "Enter an underdamped second-order system, see its step response computed from first "
        "principles, then get a plain-language explanation checked against the numbers."))

    with gr.Row(equal_height=False):
        with gr.Column(scale=5, min_width=300):
            mode = gr.Radio(list(MODES), value="Normalised (ωn, ζ)", label="Describe the system by",
                            info="Physical values are converted to ωn = √(k/m) and ζ = c / (2√(km)).")
            with gr.Column(visible=True) as norm_group:
                wn = gr.Number(2.0, label="Natural frequency ωn (rad/s)", minimum=0.1, maximum=500,
                               info="Valid range 0.1 to 500 rad/s.")
                zeta = gr.Slider(0.05, 0.95, value=0.5, step=0.01, label="Damping ratio ζ",
                                 info="Underdamped only. Below 0.05 or at 1 and above is out of scope.")
                with gr.Row():
                    amplitude = gr.Number(1.0, label="Step size A", info="Nonzero.")
                    gain = gr.Number(1.0, label="DC gain K", info="Final value is K·A.")
            with gr.Column(visible=False) as phys_group:
                with gr.Row():
                    mass = gr.Number(1.0, label="Mass m (kg)", minimum=0.01, maximum=10000,
                                     info="0.01 to 10,000 kg.")
                    stiffness = gr.Number(100.0, label="Spring stiffness k (N/m)", minimum=1,
                                          maximum=1e7, info="1 to 10⁷ N/m.")
                with gr.Row():
                    damping = gr.Number(4.0, label="Damping c (N·s/m)", minimum=0,
                                        info="Must give 0.05 ≤ ζ ≤ 0.95.")
                    force = gr.Number(10.0, label="Step force F (N)", info="Nonzero. Output shown in mm.")
            band = gr.Radio(list(C.BANDS), value="2%", label="Settling band",
                            info="Settled means staying within this fraction of the final value.")
            with gr.Accordion("Design limits (optional)", open=True):
                with gr.Row():
                    os_on = gr.Checkbox(True, label="Check overshoot")
                    os_limit = gr.Number(10.0, label="Max overshoot (%)", minimum=0.1, maximum=99)
                with gr.Row():
                    ts_on = gr.Checkbox(True, label="Check settling time")
                    ts_limit = gr.Number(3.0, label="Max settling time (s)", minimum=0.0001)

        with gr.Column(scale=6, min_width=320):
            card = gr.HTML()
            plot = gr.Plot(label="Step response", show_label=False)
            table = gr.Dataframe(headers=TABLE_COLUMNS, label="Numerical results",
                                 interactive=False, wrap=True)

    with gr.Column(elem_id="explain-panel"):
        gr.HTML(T.section_html("Explain the results", "A small language model reads the structured "
                               "record and describes it in plain words. It does not do any math."))
        explain_btn = gr.Button("Explain these results", variant="primary")
        explanation = gr.Markdown(STALE_NOTE, elem_id="explain-text")
        grounding = gr.HTML()
        with gr.Accordion("What the language model receives", open=False):
            llm_input = gr.Textbox(label="Structured record as text", lines=18, interactive=False)
        with gr.Accordion("Full structured record (JSON)", open=False):
            record_json = gr.JSON(label="Record")

    record_state = gr.State(None)
    inputs = [mode, wn, zeta, amplitude, gain, mass, damping, stiffness, force,
              band, os_on, os_limit, ts_on, ts_limit]
    mode.change(toggle_mode, mode, [norm_group, phys_group], api_name=False)
    # With no explicit triggers, gr.on runs on load and whenever an input changes
    gr.on(fn=run_calculation, inputs=inputs,
          outputs=[card, plot, table, record_state, llm_input, record_json, explanation, grounding],
          trigger_mode="always_last", concurrency_limit=4, api_name="calculate")
    explain_btn.click(run_explanation, record_state, [explanation, grounding],
                      concurrency_limit=1, api_name="explain")

    gr.Examples(
        examples=[
            ["Normalised (ωn, ζ)", 2.0, 0.5, 1.0, 1.0, 1.0, 4.0, 100.0, 10.0, "2%", True, 10.0, True, 3.0],
            ["Normalised (ωn, ζ)", 5.0, 0.7, 1.0, 1.0, 1.0, 4.0, 100.0, 10.0, "5%", False, 10.0, False, 3.0],
            ["Mass-spring-damper (m, c, k)", 2.0, 0.5, 1.0, 1.0, 250.0, 1000.0, 16000.0, 2500.0, "2%", True, 25.0, True, 1.5],
            ["Mass-spring-damper (m, c, k)", 2.0, 0.5, 1.0, 1.0, 0.5, 4.0, 2000.0, 5.0, "2%", True, 20.0, True, 0.5],
            ["Mass-spring-damper (m, c, k)", 2.0, 0.5, 1.0, 1.0, 5.0, 60.0, 50.0, 20.0, "2%", False, 10.0, False, 3.0],
        ],
        example_labels=[
            "Textbook case, fails both limits",
            "ζ = 0.7, settles before its peak",
            "Quarter-car suspension, 2.5 kN load",
            "Lightly damped sensor mount",
            "Door closer, overdamped (out of scope)",
        ],
        inputs=inputs, label="Try an example",
    )

    gr.HTML(T.footer_html([
        "Assumes a linear second-order system with no zeros, an ideal step at t = 0, and zero initial "
        "conditions. Overdamped and critically damped systems are outside this calculator's scope.",
        f'Explanations by <a href="https://huggingface.co/{E.MODEL_ID}">{E.MODEL_ID}</a> at a pinned '
        "revision, run on CPU. 24-679 Homework 3 by Jack Stevens.",
    ]))

if __name__ == "__main__":
    demo.launch(theme=T.build_theme(), css=T.CSS + EXTRA_CSS, footer_links=[])