# -*- coding: utf-8 -*-
"""natural_language_gradio.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/135ewhMpX2YZA9ysE_QNN0yFQVWYaeW34
"""

import json, pandas as pd
from dataclasses import dataclass, asdict

import math
from dataclasses import dataclass, asdict
from typing import Dict, Any

@dataclass
class PIDInputs:
    tau_s: float            # plant time constant [s], tau > 0
    Kp: float               # proportional gain
    Ki: float               # integral gain (>0 for step tracking)
    Kd: float               # derivative gain
    step_amplitude: float = 1.0   # unit step default
    settling_pct: float = 0.02    # 2% criterion for settling time

def validate_inputs(x: PIDInputs) -> Dict[str, Any]:
    issues = []
    # Reasonable ranges for a compact, safe demo. Adjust as needed.
    if not (0.01 <= x.tau_s <= 10.0):
        issues.append(f"tau must be in [0.01, 10] s, got {x.tau_s:.4g}")
    if not (-0.9 <= x.Kp <= 200.0):
        issues.append(f"Kp should be in [-0.9, 200], got {x.Kp:.4g}")
    if not (1e-6 <= x.Ki <= 1e4):
        issues.append(f"Ki should be in [1e-6, 1e4], got {x.Ki:.4g}")
    if not (-0.009 <= x.Kd <= 100.0):
        issues.append(f"Kd should be in [-0.009, 100], got {x.Kd:.4g}")
    if x.tau_s + x.Kd <= 0:
        issues.append("tau + Kd must be > 0 for a proper 2nd-order form.")
    if x.step_amplitude == 0:
        issues.append("step amplitude should be non-zero for meaningful metrics.")
    if not (0.005 <= x.settling_pct <= 0.1):
        issues.append("settling_pct should be within [0.005, 0.1] (i.e., 0.5% to 10%).")
    return {"ok": len(issues) == 0, "issues": issues}

def compute_pid(x: PIDInputs) -> Dict[str, Any]:
    val = validate_inputs(x)
    status = "ok" if val["ok"] else "invalid"

    wn = None
    zeta = None
    if (x.tau_s + x.Kd) > 0 and x.Ki > 0:
        wn = math.sqrt(x.Ki / (x.tau_s + x.Kd))
        denom = 2.0 * math.sqrt(x.Ki * (x.tau_s + x.Kd))
        zeta = (x.Kp + 1.0) / denom

    # --- NEW: poles & damping classification ---
    poles = None
    damping_class = None
    if wn is not None and wn > 0 and zeta is not None:
        # Standard 2nd-order characteristic: s^2 + 2ζωn s + ωn^2 = 0
        # Poles: s = -ζωn ± ωn*sqrt(ζ^2 - 1)
        re = -zeta * wn
        disc = zeta**2 - 1.0
        if disc < 0:
            # complex-conjugate poles
            im = wn * math.sqrt(1.0 - zeta**2)
            poles = [complex(re,  im), complex(re, -im)]
            damping_class = "underdamped (ζ<1): complex-conjugate poles"
        elif abs(disc) < 1e-12:
            poles = [complex(re, 0.0), complex(re, 0.0)]
            damping_class = "critically damped (ζ≈1): repeated real pole"
        else:
            # distinct real poles
            root = wn * math.sqrt(disc)
            poles = [complex(re + root, 0.0), complex(re - root, 0.0)]
            damping_class = "overdamped (ζ>1): two distinct real poles"

    metrics = {}
    if wn is not None and zeta is not None and wn > 0 and zeta > 0:
        if zeta < 1.0:
            wd = wn * math.sqrt(1.0 - zeta**2)
            Tp = math.pi / wd
            Mp = math.exp(-math.pi * zeta / math.sqrt(1.0 - zeta**2))  # ratio
        else:
            wd = None
            Tp = None
            Mp = 0.0
        Ts = 4.0 / (zeta * wn) * (0.02 / x.settling_pct)
        if zeta < 1.0:
            theta = math.acos(zeta)
            Tr = (math.pi - theta) / (wn * math.sqrt(1.0 - zeta**2))
        else:
            Tr = 2.0 / wn
        ess = 0.0

        metrics = {
            "wn_rad_s": wn,
            "zeta": zeta,
            "wd_rad_s": wd,
            "Mp_pct": 100.0 * Mp,
            "Tp_s": Tp,
            "Ts_s": Ts,
            "Tr_s": Tr,
            "ess": ess,
        }

    structured = {
        "meta": {
            "model": "PID_on_1stOrder_v1",
            "deterministic": True,
            "assumptions": [
                "Unity feedback.",
                "1st-order plant G(s) = 1/(tau s + 1).",
                "Linear time-invariant dynamics.",
                "PID controller C(s) = Kp + Ki/s + Kd s.",
                "Small-signal step response analysis."
            ],
            "units": {
                "tau_s": "s",
                "wn_rad_s": "rad/s",
                "wd_rad_s": "rad/s",
                "Tp_s": "s",
                "Ts_s": "s",
                "Tr_s": "s",
                "Mp_pct": "%"
            },
            "valid_ranges": {
                "tau_s": "[0.01, 10] s",
                "Kp": "[-0.9, 200]",
                "Ki": "[1e-6, 1e4]",
                "Kd": "[-0.009, 100]",
                "tau+Kd": "> 0",
                "Ki_positive": "> 0",
                "settling_pct": "[0.005, 0.1]"
            }
        },
        "inputs": asdict(x),
        "validation": val,
        "normalized_second_order": {
            "a2": x.tau_s + x.Kd,
            "a1": 1.0 + x.Kp,
            "a0": x.Ki,
            "wn": wn,
            "zeta": zeta
        },
        # --- NEW: add poles & classification in the payload ---
        "poles": [complex(p).real if abs(p.imag) < 1e-15 else p for p in (poles or [])],
        "damping_class": damping_class,
        "metrics": metrics,
        "status": status
    }
    return structured

import gradio as gr
import pandas as pd
from transformers import pipeline
from typing import Dict, Any

# from core import PIDInputs, compute_pid

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
explainer = pipeline(task="text-generation", model=_model, tokenizer=_tokenizer)

def explain_structured(d: dict) -> str:
    """
    Explain what the OUTPUT means (stability class, ωn, ζ, poles, overshoot, Tr/Tp/Ts, ess).
    Uses the SmolLM explainer with deterministic decoding, then falls back to a
    deterministic Markdown explanation if the model returns too little text.
    """
    meta  = d.get("meta", {})
    m     = d.get("metrics", {})
    norm  = d.get("normalized_second_order", {})
    poles = d.get("poles", [])
    dampc = d.get("damping_class", None)
    val   = d.get("validation", {})
    status = d.get("status")
    issues = val.get("issues", [])

    # ---------- helpers ----------
    def r(v, n=4, na="N/A"):
        try:
            return f"{float(v):.{n}g}"
        except Exception:
            return na if v is None else str(v)

    def pstr(p):
        try:
            # p may already be complex or a float
            if isinstance(p, complex) or (hasattr(p, "imag") and p.imag != 0):
                return f"{p.real:+.4g} {'+' if p.imag>=0 else '-'} j{abs(p.imag):.4g}"
            return f"{float(p):+.4g}"
        except Exception:
            return str(p)

    def dedup_lines(md: str) -> str:
        seen, out = set(), []
        for line in md.splitlines():
            key = line.strip()
            # never dedup headers; only de-dup plain bullet/paragraph lines
            if key and not key.startswith("#") and key in seen:
                continue
            seen.add(key)
            out.append(line)
        return "\n".join(out)

    # ---------- invalid → deterministic, no LLM ----------
    if status != "ok" or issues:
        bullets = "\n".join([f"- {iss}" for iss in issues]) if issues else "- Check inputs."
        return f"""# Results Explanation

**Status:** ❌ Invalid inputs

Fix these first:
{bullets}

**Why it matters**
- τ+Kd must be > 0 to form a valid 2nd-order model.
- Ki > 0 (type-1) gives zero steady-state error to a step.
"""

    # ---------- numeric snapshot for prompt & fallback ----------
    wn   = norm.get("wn")
    zeta = norm.get("zeta")
    Mp   = m.get("Mp_pct")
    Tp   = m.get("Tp_s")
    Ts   = m.get("Ts_s")
    Tr   = m.get("Tr_s")
    ess  = m.get("ess")
    poles_text = ", ".join(pstr(p) for p in poles) if poles else "N/A"

    snapshot = (
        f"- ωₙ (natural frequency): {r(wn)} rad/s\n"
        f"- ζ (damping ratio): {r(zeta)}  → {dampc or 'N/A'}\n"
        f"- Poles: {poles_text}\n"
        f"- Overshoot: ≈ {r(Mp,3)} %\n"
        f"- Rise time Tr: ≈ {r(Tr)} s\n"
        f"- Peak time Tp: ≈ {r(Tp)} s\n"
        f"- Settling time Ts: ≈ {r(Ts)} s\n"
        f"- Steady-state error (step): {r(ess)}"
    )

    # ---------- LLM prompt (deterministic, stability-focused) ----------
    prompt = (
        "You are a controls engineer. Explain what the OUTPUT VALUES MEAN.\n"
        "Write CLEAR MARKDOWN with short, specific bullets. No repetition.\n\n"
        "## Stability classification (what ζ and the poles tell you)\n"
        "- State whether the system is underdamped, critically damped, or overdamped based on ζ and the pole pattern.\n"
        "- Explain what complex vs real poles imply for oscillations and smoothness.\n\n"
        "## What ωₙ means (speed)\n"
        "- Explain that ωₙ sets the overall speed scale of the response (larger ωₙ → shorter Tr and Ts).\n\n"
        "## What ζ means (smoothness vs overshoot)\n"
        "- Interpret ζ ranges (<1, ≈1, >1) in terms of oscillation and overshoot.\n\n"
        "## What each time/percent metric means\n"
        "- Overshoot: how much the peak exceeds final value.\n"
        "- Tr: time to go from low to near-final (e.g., 10–90%).\n"
        "- Tp: time to first peak.\n"
        "- Ts: time to settle within the chosen band.\n"
        "- ess: final error for a step; with Ki>0 it is 0.\n\n"
        "## How the poles relate to that behavior\n"
        "- Connect pole real part to decay speed; imaginary part to oscillation frequency.\n\n"
        "## Numeric snapshot\n"
        f"{snapshot}\n"
    )

    # ---------- deterministic generation with anti-repetition ----------
    gen = explainer(
        prompt,
        max_new_tokens=220,
        do_sample=False,
        temperature=0.0,
        top_p=1.0,
        top_k=0,
        repetition_penalty=1.15,
        no_repeat_ngram_size=4,
        eos_token_id=_tokenizer.eos_token_id,
        pad_token_id=_tokenizer.eos_token_id,
        return_full_text=False
    )[0]["generated_text"]

    # ---------- SHORT-OUTPUT FALLBACK (your requested addition) ----------
    MIN_WORDS = 30
    if not gen or len(gen.split()) < MIN_WORDS:
        gen = f"""## Stability classification
- ζ = {r(zeta)} → {dampc or 'N/A'}.

## Meaning of ωₙ and ζ
- ωₙ = {r(wn)} rad/s sets the speed scale (larger ωₙ → faster rise/settle).
- ζ controls smoothness/overshoot: ζ<1 underdamped; ζ≈1 critically damped; ζ>1 overdamped.

## Poles and behavior
- Poles: {poles_text}
- More negative real part → faster decay; nonzero imaginary part → oscillations.

## Time-domain metrics
- Overshoot ≈ {r(Mp,3)} %  |  Tr ≈ {r(Tr)} s  |  Tp ≈ {r(Tp)} s  |  Ts ≈ {r(Ts)} s  |  ess = {r(ess)}

## Tuning tip
- Raise Ki to increase ωₙ (speed). If overshoot or oscillation appears (ζ too low), add Kd or increase Kp to raise damping.
"""

    return dedup_lines(gen)

def run_calc(tau_s, Kp, Ki, Kd, step_amp, settling_pct):
    x = PIDInputs(
        tau_s=float(tau_s), Kp=float(Kp), Ki=float(Ki), Kd=float(Kd),
        step_amplitude=float(step_amp), settling_pct=float(settling_pct) / 100.0  # slider in %, convert to fraction
    )
    structured = compute_pid(x)

    # Display normalized form + metrics in a compact table
    rows = []
    for k, v in structured.get("normalized_second_order", {}).items():
        rows.append(["2nd-order", k, v])
    for k, v in structured.get("metrics", {}).items():
        rows.append(["metrics", k, v])
    df = pd.DataFrame(rows, columns=["section", "key", "value"])

    explanation = explain_structured(structured)
    return df, explanation, structured

with gr.Blocks(title="PID Controls Calculator (1st-Order Plant)", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# PID Feedback Controls — Deterministic Calculator")
    gr.Markdown(
        "Unity-feedback PID on a first-order plant G(s)=1/(τs+1). "
        "We derive the equivalent 2nd-order parameters (ωₙ, ζ) and step-response metrics (overshoot, rise, peak, settling)."
    )

    with gr.Row():
        with gr.Column():
            tau_s = gr.Slider(0.01, 10.0, value=0.5, step=0.01, label="Plant time constant τ [s]")
            Kp = gr.Slider(-0.9, 200.0, value=1.0, step=0.1, label="Kp")
            Ki = gr.Slider(1e-6, 1e4, value=1.0, step=0.1, label="Ki")
            Kd = gr.Slider(-0.009, 100.0, value=0.0, step=0.001, label="Kd")
            step_amp = gr.Slider(0.1, 10.0, value=1.0, step=0.1, label="Step amplitude")
            settling_pct = gr.Slider(0.5, 10.0, value=2.0, step=0.1, label="Settling band [%]")
            go = gr.Button("Compute", variant="primary")

        with gr.Column():
            gr.Markdown("### Numerical Results")
            table = gr.Dataframe(headers=["section", "key", "value"], interactive=False)
            gr.Markdown("### Explain the Results")
            explanation = gr.Markdown()
            gr.Markdown("### Raw Structured Output")
            json_out = gr.JSON(label="Structured JSON")

    go.click(run_calc, inputs=[tau_s, Kp, Ki, Kd, step_amp, settling_pct],
             outputs=[table, explanation, json_out])

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