app.py

# --- deps (minimal) ---
# !pip -q install "gradio>=4.44.0" "transformers>=4.42.0"

# --- deterministic backend ---
from math import pi
import re
import gradio as gr
from transformers import pipeline

def validate_and_compute(
    driver_teeth:int,
    driven_teeth:int,
    input_torque_Nm:float,
    input_speed_rpm:float,
    efficiency_pct:float,
):
    """
    Deterministic, first-principles 2-gear calculator.

    SCOPE & ASSUMPTIONS
    - Two rigid spur gears (driver -> driven), single mesh, no slippage.
    - Quasi-steady operation; efficiency is mesh efficiency only.
    - No shaft/bearing losses; no strength checks.

    INPUTS & VALID RANGES
    - driver_teeth:  6..400 (int)
    - driven_teeth:  6..400 (int)
    - input_torque_Nm:  0..10_000  [N·m]
    - input_speed_rpm:  0..30_000  [rpm]
    - efficiency_pct:   50..100    [%]
    """
    problems = []
    if not isinstance(driver_teeth, int) or not (6 <= driver_teeth <= 400):
        problems.append("driver_teeth must be an integer in [6, 400].")
    if not isinstance(driven_teeth, int) or not (6 <= driven_teeth <= 400):
        problems.append("driven_teeth must be an integer in [6, 400].")
    if not (0 <= input_torque_Nm <= 10_000):
        problems.append("input_torque_Nm must be in [0, 10_000] N·m.")
    if not (0 <= input_speed_rpm <= 30_000):
        problems.append("input_speed_rpm must be in [0, 30_000] rpm.")
    if not (50 <= efficiency_pct <= 100):
        problems.append("efficiency_pct must be in [50, 100] %.")

    if problems:
        return {"ok": False, "errors": problems}

    gear_ratio = driven_teeth / driver_teeth               # >1 → torque↑ speed↓
    eta = efficiency_pct / 100.0

    # kinematics
    output_speed_rpm = input_speed_rpm / gear_ratio

    # torque (power*eta conserved across mesh)
    output_torque_Nm = input_torque_Nm * gear_ratio * eta

    # power check
    omega_in  = input_speed_rpm  * (2*pi/60.0)  # rad/s
    omega_out = output_speed_rpm * (2*pi/60.0)  # rad/s
    pin_W  = input_torque_Nm  * omega_in
    pout_W = output_torque_Nm * omega_out
    expected_pout_W = pin_W * eta
    power_balance_error_pct = (0 if expected_pout_W == 0 else
                               100.0 * (pout_W - expected_pout_W) / max(1e-12, expected_pout_W))

    speed_change = "decrease" if gear_ratio > 1 else ("increase" if gear_ratio < 1 else "no change")
    torque_change = "increase" if gear_ratio > 1 else ("decrease" if gear_ratio < 1 else "no change")

    return {
        "ok": True,
        "inputs": {
            "driver_teeth": driver_teeth,
            "driven_teeth": driven_teeth,
            "input_torque_Nm": float(input_torque_Nm),
            "input_speed_rpm": float(input_speed_rpm),
            "efficiency_pct": float(efficiency_pct)
        },
        "intermediate": {
            "gear_ratio": float(gear_ratio),
            "eta": float(eta)
        },
        "outputs": {
            "output_speed_rpm": float(output_speed_rpm),
            "output_torque_Nm": float(output_torque_Nm),
            "input_power_W": float(pin_W),
            "output_power_W": float(pout_W),
            "expected_output_power_W": float(expected_pout_W),
            "power_balance_error_percent": float(power_balance_error_pct),
            "qualitative": {
                "speed_change": speed_change,
                "torque_change": torque_change
            }
        }
    }

# --- compact, deterministic LLM explainer ---
explainer = pipeline("text2text-generation", model="google/flan-t5-small")

def sanitize_to_plain_sentences(text: str, max_sentences: int = 6) -> str:
    """Drop code/JSON-y lines and keep up to N plain sentences."""
    # Strip code blocks/backticks
    text = re.sub(r"```.*?```", "", text, flags=re.S)
    text = text.replace("`", "")
    # Keep only part after the last 'Explanation:' if present
    if "Explanation:" in text:
        text = text.split("Explanation:")[-1]
    # Remove obvious code-ish lines
    cleaned_lines = []
    for line in text.splitlines():
        l = line.strip()
        if not l:
            continue
        if re.search(r"\binput\(", l) or re.search(r"\bprint\(", l):
            continue
        if re.match(r"^\s*[A-Za-z_]\w*\s*=", l):  # assignments
            continue
        if l.startswith("{") or l.startswith("[") or l.endswith("}"):
            continue
        cleaned_lines.append(l)
    text = " ".join(cleaned_lines)

    # Split into sentences conservatively
    # (avoid splitting inside units like N·m, rpm—these don’t have periods)
    sentences = re.split(r"(?<=[.!?])\s+", text)
    # Trim and keep sentences that look like natural language
    plain = []
    for s in sentences:
        s = s.strip()
        if not s:
            continue
        # Heuristics to avoid code/fragments
        if re.search(r"[{};]|==|!=|<=|>=|\breturn\b|\bdef\b|\bclass\b", s):
            continue
        plain.append(s)
        if len(plain) >= max_sentences:
            break
    return " ".join(plain).strip()

def fallback_explanation(payload: dict) -> str:
    """Deterministic template if the LLM output gets sanitized to nothing/too short."""
    inp = payload["inputs"]; mid = payload["intermediate"]; out = payload["outputs"]
    sc = out["qualitative"]["speed_change"]; tc = out["qualitative"]["torque_change"]
    return (
        f"The driven gear has {inp['driven_teeth']} teeth and the driver has {inp['driver_teeth']}, "
        f"so the gear ratio is {mid['gear_ratio']:.3f} (driven/driver). "
        f"With an input torque of {inp['input_torque_Nm']:.3g} N·m at {inp['input_speed_rpm']:.3g} rpm, "
        f"the output speed becomes {out['output_speed_rpm']:.3g} rpm (a {sc} in speed) and the output torque is "
        f"{out['output_torque_Nm']:.3g} N·m (a {tc} in torque). "
        f"Input power is {out['input_power_W']:.3g} W; with an efficiency of {inp['efficiency_pct']:.3g}%, "
        f"the expected output power is {out['expected_output_power_W']:.3g} W, which matches the computed "
        f"{out['output_power_W']:.3g} W. "
        f"This calculation considers mesh efficiency only; check tooth strength, shafts, and bearings separately."
    )

def render_explanation(payload: dict) -> str:
    if not payload.get("ok"):
        return "There were validation errors; please fix inputs and try again."

    inp = payload["inputs"]; mid = payload["intermediate"]; out = payload["outputs"]

    # Only pass essentials so the model can’t riff on raw JSON.
    facts = (
        f"Driver teeth: {inp['driver_teeth']}; Driven teeth: {inp['driven_teeth']}. "
        f"Input torque: {inp['input_torque_Nm']:.4g} N·m; Input speed: {inp['input_speed_rpm']:.4g} rpm. "
        f"Efficiency: {inp['efficiency_pct']:.4g}%. "
        f"Gear ratio (driven/driver): {mid['gear_ratio']:.6g}. "
        f"Output speed: {out['output_speed_rpm']:.6g} rpm; Output torque: {out['output_torque_Nm']:.6g} N·m. "
        f"Input power: {out['input_power_W']:.6g} W; Output power: {out['output_power_W']:.6g} W; "
        f"Expected Pin×η: {out['expected_output_power_W']:.6g} W."
    )

    prompt = (
        "You are a careful engineering assistant. Using only the facts below, write 4–6 clear sentences that "
        "explain the gear calculation to a non-expert. Mention the gear ratio’s tradeoff (speed vs torque) and "
        "comment that output power is approximately input power times efficiency. "
        "Do NOT include code, equations, lists, JSON, or symbols like '=', '{', '}'. Plain sentences only.\n\n"
        f"{facts}\n\nExplanation:"
    )

    gen = explainer(
        prompt,
        max_new_tokens=140,
        do_sample=False,        # deterministic
        num_beams=4,
        early_stopping=True
    )[0]["generated_text"]

    cleaned = sanitize_to_plain_sentences(gen)
    # If the model still produced garbage, fall back to a deterministic template.
    if len(cleaned.split()) < 25:  # too short / likely sanitized away
        cleaned = fallback_explanation(payload)
    return cleaned

# --- Gradio app (no JSON panel; with clickable examples) ---
def app_fn(driver_teeth, driven_teeth, input_torque, input_speed, efficiency):
    result = validate_and_compute(
        int(driver_teeth),
        int(driven_teeth),
        float(input_torque),
        float(input_speed),
        float(efficiency),
    )

    if not result["ok"]:
        errors_md = "### Validation Errors\n" + "\n".join([f"- {e}" for e in result["errors"]])
        return errors_md, "—"

    r = result
    def fmt(x, digits=6):  # compact pretty printer
        return f"{x:.{digits}f}".rstrip('0').rstrip('.') if isinstance(x, float) else str(x)

    md = f"""
### Numerical Results
- **Gear ratio (driven/driver):** {fmt(r['intermediate']['gear_ratio'], 6)}
- **Output speed:** {fmt(r['outputs']['output_speed_rpm'], 6)} rpm  _(speed {r['outputs']['qualitative']['speed_change']})_
- **Output torque:** {fmt(r['outputs']['output_torque_Nm'], 6)} N·m  _(torque {r['outputs']['qualitative']['torque_change']})_
- **Input power:** {fmt(r['outputs']['input_power_W'], 6)} W  
- **Output power:** {fmt(r['outputs']['output_power_W'], 6)} W  
- **Expected output power (Pin×η):** {fmt(r['outputs']['expected_output_power_W'], 6)} W  
- **Power balance error:** {fmt(r['outputs']['power_balance_error_percent'], 6)} %
"""
    explanation = render_explanation(result)
    return md, explanation

with gr.Blocks(title="Deterministic Gear Calculator (See the Math)") as demo:
    gr.Markdown("# Deterministic Gear Calculator — with Natural-Language Explanation")
    gr.Markdown(
        "Two-gear train (driver → driven). Deterministic math with validation and a plain-English explanation."
    )
    with gr.Row():
        with gr.Column():
            in_driver = gr.Number(label="Driver gear teeth", value=20, precision=0)
            in_driven = gr.Number(label="Driven gear teeth", value=65, precision=0)
            in_torque = gr.Number(label="Input torque [N·m]", value=12.0)
            in_speed  = gr.Number(label="Input speed [rpm]", value=1800.0)
            in_eta    = gr.Slider(50, 100, value=97.0, step=0.1, label="Mesh efficiency [%]")
            run_btn   = gr.Button("Calculate", variant="primary")
        with gr.Column():
            results_md     = gr.Markdown(label="Numerical Results")
            explanation_md = gr.Markdown(label="Explain the results")

    gr.Markdown("### Examples (click to fill)")
    examples = [
        [20, 60, 10.0, 1800.0, 97.0],   # 3:1 reduction
        [18, 36, 5.0, 1500.0, 96.0],    # 2:1 reduction
        [32, 24, 8.0, 1200.0, 95.0],    # overdrive
        [25, 75, 15.0, 3600.0, 98.0],   # 3:1 with higher speed
    ]
    gr.Examples(
        examples=examples,
        inputs=[in_driver, in_driven, in_torque, in_speed, in_eta],
        label="Common scenarios"
    )

    run_btn.click(
        app_fn,
        [in_driver, in_driven, in_torque, in_speed, in_eta],
        [results_md, explanation_md]
    )

demo.launch()