import json
import pandas as pd
import gradio as gr
from transformers import pipeline, set_seed

SCOPE = """
**Scope:** DC, steady-state Ohm's Law.
**Assumptions:** Ideal components; temperature effects ignored.
**Units:** V (volts), I (amps), R (ohms).
**Valid ranges:** 0 ≤ V ≤ 1000, 0 ≤ I ≤ 100, 0.1 ≤ R ≤ 1e6.
"""

set_seed(42)
try:
    explain_llm = pipeline("text2text-generation", model="google/flan-t5-small")
except Exception as e:
    print("LLM pipeline init failed; using fallback text only.", e)
    explain_llm = None

def validate(target, V, I, R):
    errs = []
    if target != "Voltage (V)"   and not (0 <= V <= 1000): errs.append("V must be 0–1000 V.")
    if target != "Current (I)"   and not (0 <= I <= 100):  errs.append("I must be 0–100 A.")
    if target != "Resistance (R)"and not (0.1 <= R <= 1e6):errs.append("R must be 0.1–1e6 Ω.")
    if target == "Current (I)"   and R == 0: errs.append("R cannot be 0 for I = V/R.")
    if target == "Resistance (R)"and I == 0: errs.append("I cannot be 0 for R = V/I.")
    return errs

def compute_structured(target, V, I, R):
    if target == "Voltage (V)":
        value = I * R
        steps = [f"Use V = I·R", f"V = {I} A × {R} Ω = {value} V"]
        out = {"quantity":"V","value":value,"units":"V"}
    elif target == "Current (I)":
        value = V / R
        steps = [f"Use I = V/R", f"I = {V} V ÷ {R} Ω = {value} A"]
        out = {"quantity":"I","value":value,"units":"A"}
    else:
        value = V / I
        steps = [f"Use R = V/I", f"R = {V} V ÷ {I} A = {value} Ω"]
        out = {"quantity":"R","value":value,"units":"Ω"}

    return {
        "calculation":"Ohm's Law",
        "equations":{"V":"V=I·R","I":"I=V/R","R":"R=V/I"},
        "target":target,
        "inputs":{"V":V,"I":I,"R":R,"units":{"V":"V","I":"A","R":"Ω"}},
        "steps":steps,
        "output":out
    }

def explain(_data: dict) -> str:
    """
    Always ask the LLM to explain what Ohm's Law is; fallback if unavailable.
    """
    prompt = (
        "Explain Ohm’s Law for a non-expert in <=80 words. "
        "Use symbols ×, ·, and Ω (no Unicode escapes). "
        "Include the three forms (V = I·R, I = V/R, R = V/I) and list the units."
    )
    if explain_llm is not None:
        try:
            txt = explain_llm(prompt, do_sample=False, num_beams=1, max_new_tokens=120)[0]["generated_text"].strip()
            if txt:
                return (txt.replace("\u00b7","·")
                          .replace("\u00D7","×").replace("\u00d7","×")
                          .replace("\u03a9","Ω").replace("\u03A9","Ω"))
        except Exception as e:
            print("LLM explanation failed:", e)
    return ("Ohm’s Law relates voltage (V), current (I), and resistance (R): "
            "V = I·R, I = V/R, R = V/I. Units: V (volts), A (amps), Ω (ohms).")

def run(target, V, I, R):
    errs = validate(target, V, I, R)
    if errs: raise gr.Error("\n".join(errs))
    data = compute_structured(target, V, I, R)
    df = pd.DataFrame([{
        "Target": data["target"],
        "V (V)": data["inputs"]["V"],
        "I (A)": data["inputs"]["I"],
        "R (Ω)": data["inputs"]["R"],
        "Result": f'{data["output"]["quantity"]} = {data["output"]["value"]} {data["output"]["units"]}'
    }])
    return df, json.dumps(data, indent=2), explain(data)

with gr.Blocks(title="Ohm's Law — see the math") as demo:
    gr.Markdown("# 🔌 Ohm's Law — Deterministic Calculator")
    gr.Markdown(SCOPE)
    with gr.Row():
        with gr.Column():
            target = gr.Radio(["Voltage (V)","Current (I)","Resistance (R)"], value="Voltage (V)", label="Compute")
            V = gr.Number(value=12.0, label="Voltage V (V) [ignored if computing V]")
            I = gr.Number(value=2.0,  label="Current I (A) [ignored if computing I]")
            R = gr.Number(value=6.0,  label="Resistance R (Ω) [ignored if computing R]")
            btn = gr.Button("Compute", variant="primary")
        with gr.Column():
            with gr.Tabs():
                with gr.Tab("Results"):
                    out_table = gr.Dataframe(label="Results", interactive=False)
                with gr.Tab("Explanation"):
                    out_expl = gr.Markdown(value="_Explanation will appear here after you click **Compute**._")
                with gr.Tab("Structured JSON"):
                    out_json = gr.Code(label="Structured output (JSON)", language="json")
    inputs = [target, V, I, R]
    btn.click(run, inputs=inputs, outputs=[out_table, out_json, out_expl])
    gr.Examples(
        inputs=inputs,
        examples=[
            ["Voltage (V)", 0.0, 2.0, 50.0],
            ["Current (I)", 12.0, 0.0, 6.0],
            ["Resistance (R)", 5.0, 0.01, 0.0]
        ],
        label="Examples"
    )
if __name__ == "__main__":
    demo.launch()