import math, json
import gradio as gr
import pandas as pd
from typing import Dict, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

LLM_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(LLM_ID)
llm = pipeline(
    task="text-generation",
    model=AutoModelForCausalLM.from_pretrained(LLM_ID),
    tokenizer=tokenizer,
    device_map="auto",
)

def projectile_calc(v0_mps: float, theta_deg: float, y0_m: float, g: float, weight_kg: float) -> Dict[str, Any]:
    errors = []
    if not (0 < v0_mps <= 500): errors.append("Initial speed must be in (0, 500] m/s.")
    if not (-10 <= theta_deg <= 90): errors.append("Launch angle must be between -10 and 90 degrees.")
    if not (0 <= y0_m <= 1000): errors.append("Initial height must be in [0, 1000] m.")
    if not (1 <= g <= 50): errors.append("Gravity must be in [1, 50] m/s^2.")
    if not (0.05 <= weight_kg <= 50): errors.append("Ball weight must be in [0.05, 50] kg.")
    if errors:
        return {"ok": False, "errors": errors}

    theta_rad = math.radians(theta_deg)
    v0x = v0_mps * math.cos(theta_rad)
    v0y = v0_mps * math.sin(theta_rad)

    # Time of flight (positive root)
    disc = v0y**2 + 2.0 * g * y0_m
    t_flight = (v0y + math.sqrt(max(disc, 0.0))) / g
    t_flight = max(t_flight, 0.0)

    # Range, apex height, impact speed
    range_m = v0x * t_flight
    t_apex = max(v0y / g, 0.0)
    y_apex = y0_m + v0y * t_apex - 0.5 * g * t_apex**2
    vy_final = -math.sqrt(max(v0y**2 + 2.0 * g * y0_m, 0.0))
    v_impact = math.sqrt(v0x**2 + vy_final**2)

    return {
        "ok": True,
        "inputs": {
            "v0_mps": v0_mps,
            "theta_deg": theta_deg,
            "y0_m": y0_m,
            "g_mps2": g,
            "weight_kg": weight_kg,
        },
        "derived": {
            "t_apex_s": t_apex,
        },
        "outputs": {
            "time_of_flight_s": t_flight,
            "range_m": range_m,
            "max_height_m": y_apex,
            "impact_speed_mps": v_impact,
        },
        "notes": {
            "assumptions": ["No air resistance.", "Flat landing at y=0.", "Constant gravity g."],
        }
    }

SYSTEM_MSG = (
    "You are a careful technical writer. "
    "Write EXACTLY ONE sentence using ONLY numbers provided in the JSON. "
    "Do not invent or rename fields; keep units as given; use ~2–3 sig figs."
)

def llm_one_sentence(structured: Dict[str, Any]) -> str:
    i = structured["inputs"]
    d = structured["derived"]
    o = structured["outputs"]

    payload = {
        "inputs": {
            "v0_mps": round(i["v0_mps"], 3),
            "theta_deg": round(i["theta_deg"], 3),
            "y0_m": round(i["y0_m"], 3),
            "g_mps2": round(i["g_mps2"], 3),
            "weight_kg": round(i["weight_kg"], 3),
        },
        "derived": {
            "t_apex_s": round(d["t_apex_s"], 4),
        },
        "outputs": {
            "time_of_flight_s": round(o["time_of_flight_s"], 4),
            "range_m": round(o["range_m"], 3),
            "max_height_m": round(o["max_height_m"], 3),
            "impact_speed_mps": round(o["impact_speed_mps"], 3),
        }
    }

    # Exact sentence shape we want.
    instruction = """Use only these keys: inputs, derived, outputs.
Return exactly one sentence in this form (fill in the braces with the JSON values):

If you threw a ball that weighed {weight_kg} kg at v0={v0_mps} m/s, theta={theta_deg} deg, from y0={y0_m} m under g={g_mps2} m/s^2,
it would stay in the air for {time_of_flight_s} s, reach {max_height_m} m, travel {range_m} m, and impact at {impact_speed_mps} m/s.
"""

    user_msg = "\n".join([
        "JSON:",
        json.dumps(payload, indent=2),
        "",
        "Produce the single sentence as specified."
    ])

    messages = [
        {"role": "system", "content": SYSTEM_MSG},
        {"role": "user", "content": instruction},
        {"role": "user", "content": user_msg},
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    out = llm(
        prompt,
        max_new_tokens=96,
        do_sample=False,
        temperature=0.0,
        return_full_text=False,
    )
    text = out[0]["generated_text"].strip()

    # Fallback if the model strays
    if ("If you threw a ball" not in text) or (len(text.split()) < 6):
        p = payload
        text = (
            f"If you threw a ball that weighed {p['inputs']['weight_kg']} kg at v0={p['inputs']['v0_mps']} m/s, "
            f"theta={p['inputs']['theta_deg']} deg, from y0={p['inputs']['y0_m']} m under g={p['inputs']['g_mps2']} m/s^2, "
            f"it would stay in the air for {p['outputs']['time_of_flight_s']} s, "
            f"reach {p['outputs']['max_height_m']} m, travel {p['outputs']['range_m']} m, "
            f"and impact at {p['outputs']['impact_speed_mps']} m/s."
        )

    return text

def run_once(v0_mps, theta_deg, y0_m, g_mps2, weight_kg):
    rec = projectile_calc(float(v0_mps), float(theta_deg), float(y0_m), float(g_mps2), float(weight_kg))
    if not rec.get("ok", False):
        errs = rec.get("errors", ["Invalid input."])
        df = pd.DataFrame([{"Error": "; ".join(errs)}])
        return df, "Please adjust inputs."
    i, d, o = rec["inputs"], rec["derived"], rec["outputs"]
    df = pd.DataFrame([{
        "v0_mps [m/s]": round(i["v0_mps"], 3),
        "theta_deg [deg]": round(i["theta_deg"], 3),
        "y0_m [m]": round(i["y0_m"], 3),
        "g_mps2 [m/s^2]": round(i["g_mps2"], 3),
        "weight_kg [kg]": round(i["weight_kg"], 3),
        "t_apex_s [s]": round(d["t_apex_s"], 4),
        "max_height_m [m]": round(o["max_height_m"], 3),
        "time_of_flight_s [s]": round(o["time_of_flight_s"], 4),
        "range_m [m]": round(o["range_m"], 3),
        "impact_speed_mps [m/s]": round(o["impact_speed_mps"], 3),
    }])
    sentence = llm_one_sentence(rec)
    return df, sentence

with gr.Blocks(title="Projectile Motion — Deterministic + One-Sentence LLM") as demo:
    gr.Markdown("# Projectile Motion — Deterministic Calculator + One-Sentence Explanation")
    gr.Markdown("Deterministic projectile motion (no air resistance). See all inputs and derived values, plus a single, grounded sentence.")

    with gr.Row():
        v0 = gr.Slider(1, 200, value=30.0, step=0.5, label="Initial speed v0 [m/s]")
        theta = gr.Slider(-10, 90, value=45.0, step=0.5, label="Launch angle theta [deg]")
        y0 = gr.Slider(0, 20, value=1.5, step=0.1, label="Initial height y0 [m]")
        g = gr.Slider(5, 20, value=9.81, step=0.01, label="Gravity g [m/s^2]")
        w = gr.Slider(0.05, 10.0, value=0.43, step=0.01, label="Ball weight [kg]")

    run_btn = gr.Button("Compute")

    results_df = gr.Dataframe(
        headers=[
            "v0_mps [m/s]", "theta_deg [deg]", "y0_m [m]", "g_mps2 [m/s^2]", "weight_kg [kg]",
            "t_apex_s [s]", "max_height_m [m]", "time_of_flight_s [s]", "range_m [m]", "impact_speed_mps [m/s]"
        ],
        label="All values (inputs + derived)",
        interactive=False
    )
    explain_md = gr.Markdown(label="One-sentence explanation")

    run_btn.click(run_once, inputs=[v0, theta, y0, g, w], outputs=[results_df, explain_md])

    gr.Examples(
        examples=[
            [30.0, 45.0, 1.5, 9.81, 0.43],
            [20.0, 30.0, 0.0, 9.81, 0.145],
            [40.0, 60.0, 0.0, 9.81, 0.45],
        ],
        inputs=[v0, theta, y0, g, w],
        label="Representative cases",
        examples_per_page=3,
        cache_examples=False,
    )

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