import math # For access to infinity

import gradio # For building the interface
import pandas # For working with tables

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline # For LLMS

# Instantiate the model that we'll be calling. This is a tiny one!
MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
pipe = pipeline(
    task="text-generation",
    model=AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
    ),
    tokenizer=tokenizer
)

# Create a function to do the tensile test calculations
def simple_tensile_test_calc(Force_N: float, Area_m2: float, YoungsModulus_GPa: float) -> dict:
    """
    Simple tensile test calculation.
    Calculates stress and strain.
    """

    # Convert Young's Modulus from GPa to Pa
    YoungsModulus_Pa = YoungsModulus_GPa * 1e9

    # Calculate Stress (Pa)
    Stress_Pa = Force_N / Area_m2
    Stress_MPa = Stress_Pa / 1e6

    # Calculate Strain (dimensionless)
    Strain = Stress_Pa / YoungsModulus_Pa

    return dict(
        results={
            "Stress_MPa": Stress_MPa,
            "Strain": Strain,
        }
    )

# This helper function applies a chat format to help the LLM understand what
# is going on
def _format_chat(system_prompt: str, user_prompt: str) -> str:
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    template = getattr(tokenizer, "chat_template", None)
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

# This functoin uses hte LLM to generate a response.
def _llm_generate(prompt: str, max_tokens: int) -> str:
    out = pipe(
        prompt,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=0.5,
        return_full_text=False,
    )
    return out[0]["generated_text"]

# This function generates an explanation of the results
def llm_explain(results: dict, inputs: list) -> str:
    Force_N, Area_m2, YoungsModulus_GPa = inputs
    r = results["results"]

    system_prompt = (
        "You explain engineering to a smart 5-year-old. "
        "Use food-based analogies to support the explanation."
        "You always return CONCISE responses, only one sentence."
    )

    user_prompt = (
        f"A force of {Force_N:g} N is applied to a material with a cross-sectional area of {Area_m2:g} m².\n"
        f"The material has a Young's Modulus of {YoungsModulus_GPa:g} GPa.\n"
        f"The calculated stress is {r['Stress_MPa']:.2f} MPa and the strain is {r['Strain']:.4f}.\n"
        "Explain what stress and strain mean in ONE friendly sentence for a non-expert"
        ""
    )

    formatted = _format_chat(system_prompt, user_prompt)
    return _llm_generate(formatted, max_tokens=128)


# This function ties everythign together (evaluation, LLM explanaation, output)
# And will be out main entry point for teh GUI
def run_once(Force_N, Area_m2, YoungsModulus_GPa):
    inputs = [Force_N, Area_m2, YoungsModulus_GPa]
    d = simple_tensile_test_calc(
        Force_N=float(Force_N),
        Area_m2=float(Area_m2),
        YoungsModulus_GPa=float(YoungsModulus_GPa),
    )

    df = pandas.DataFrame([{
        "Stress [MPa]": round(d["results"]["Stress_MPa"], 3),
        "Strain [-]": round(d["results"]["Strain"], 6),
    }])

    narrative = llm_explain(d, inputs).split("\n")[0]
    return df, narrative

# Last but not high_light, here's the UI!
with gradio.Blocks() as demo:

    # Let's start by adding a title and introduction
    gradio.Markdown(
        "# Run and Explain Tensile Test Calculations"
    )
    gradio.Markdown(
        "This app runs simple calculations for a tensile test and returns a natural language description of the results"
    )

    # This row contains all of the physical parameters
    with gradio.Row():
        Force_N = gradio.Number(value=1000.0, label="Applied Force [N]")
        Area_m2 = gradio.Number(value=0.0001, label="Cross-sectional Area [m²]")


    # This row contains the material properties
    with gradio.Row():
        YoungsModulus_GPa = gradio.Number(value=200.0, label="Young's Modulus [GPa]")


    # Add a button to click to run the interface
    run_btn = gradio.Button("Compute")

    # These are the outputs. We use both a dataframe (for tabular info) and a markdown box
    # for info from teh LLM
    results_df = gradio.Dataframe(label="Numerical results (deterministic)", interactive=False)
    explain_md = gradio.Markdown(label="Explanation")

    # Run the calculations when the button is clicked
    run_btn.click(fn=run_once, inputs=[Force_N, Area_m2, YoungsModulus_GPa], outputs=[results_df, explain_md])

    # Finally, add a few examples
    gradio.Examples(
        examples=[
            [1000.0, 0.0001, 200.0],
            [2000.0, 0.0002, 210.0],
            [500.0, 0.00005, 70.0],
        ],
        inputs=[Force_N, Area_m2, YoungsModulus_GPa],
        label="Representative cases",
        examples_per_page=3,
        cache_examples=False,
    )

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