import math
import gradio
import pandas
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# 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 Young's Modulus calculation
def youngs_modulus_calc(F_N: float, A_m2: float, L0_m: float, delta_L_m: float) -> float:
    """
    Calculates Young's Modulus (E) in GPa.
    F_N: Force in Newtons
    A_m2: Area in square meters
    L0_m: Original length in meters
    delta_L_m: Change in length in meters
    """
    if A_m2 == 0 or L0_m == 0 or delta_L_m == 0:
        return math.inf  # Avoid division by zero
    stress_Pa = F_N / A_m2
    strain = delta_L_m / L0_m
    E_Pa = stress_Pa / strain
    E_GPa = E_Pa / 1e9
    return E_GPa

# 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(E_GPa: float, inputs: list) -> str:
    system_prompt = (
        "Explain engineering concepts to a smart 5-year-old using food analogies. "
        "Keep responses concise, only one sentence."
    )

    F_N, A_m2, L0_m, delta_L_m = inputs

    aluminum_E_GPa = 70.0
    comparison = "more flexible" if E_GPa < aluminum_E_GPa else "stiffer" if E_GPa > aluminum_E_GPa else "about as stiff"


    user_prompt = (
        f"An unknown material was tested with a force of {F_N:g} N, area of {A_m2:g} m², original length of {L0_m:g} m, and change in length of {delta_L_m:g} m.\n"
        f"The unknown material's Young's Modulus is {E_GPa:.2f} GPa.\n"
        f"Aluminum's Young's Modulus is {aluminum_E_GPa} GPa.\n"
        f"Explain if this material is {comparison} than aluminum."
        "Explain the Young's modulus of the unknown material and the compare it to aluminum 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(F_N, A_m2, L0_m, delta_L_m):
    # Young's Modulus calculation part
    youngs_modulus_inputs = [F_N, A_m2, L0_m, delta_L_m]
    youngs_modulus_result = youngs_modulus_calc(
        F_N=float(F_N),
        A_m2=float(A_m2),
        L0_m=float(L0_m),
        delta_L_m=float(delta_L_m)
    )

    aluminum_E_GPa = 70.0
    stiffness_comparison = "stiff" if youngs_modulus_result > aluminum_E_GPa else "flexible" if youngs_modulus_result < aluminum_E_GPa else "similar stiffness"

    youngs_modulus_df = pandas.DataFrame([{
        "Young's Modulus [GPa]": round(youngs_modulus_result, 3),
        "Compared to Aluminum": stiffness_comparison
        }])

    youngs_modulus_narrative = llm_explain(youngs_modulus_result, youngs_modulus_inputs).split("\n")[0]


    return youngs_modulus_df, youngs_modulus_narrative

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

    # Let's start by adding a title and introduction
    gradio.Markdown(
        "# Run and Explain Young's Modulus Calculation"
    )
    gradio.Markdown(
        "This app calculates Young's Modulus and provides a natural language description of the result, comparing it to aluminum."
    )

    # This row contains all of the physical parameters for Young's Modulus
    with gradio.Row():
        F_N = gradio.Number(value=1000.0, label="Force [N]")
        A_m2 = gradio.Number(value=0.01, label="Area [m²]")
        L0_m = gradio.Number(value=1.0, label="Original Length [m]")
        delta_L_m = gradio.Number(value=0.001, label="Change in Length [m]")

    # Add a button to click to run the interface
    youngs_modulus_run_btn = gradio.Button("Compute Young's Modulus")

    # These are the outputs for Young's Modulus
    youngs_modulus_results_df = gradio.Dataframe(label="Numerical results (deterministic)", interactive=False)
    youngs_modulus_explain_md = gradio.Markdown(label="Explanation")

    # Run the calculations when the button is clicked
    youngs_modulus_run_btn.click(
        fn=run_once,
        inputs=[F_N, A_m2, L0_m, delta_L_m],
        outputs=[youngs_modulus_results_df, youngs_modulus_explain_md]
    )

    # Finally, add a few examples
    gradio.Examples(
        examples=[
            [1000000.0, 0.001, 1.0, 0.01], # Example 1 (should be around 100 GPa - stiffer than aluminum)
            [50000.0, 0.01, 1.0, 0.0002], # Example 2 (should be around 25 GPa - more flexible than aluminum)
            [700000.0, 0.01, 1.0, 0.001], # Example 3 (should be around 70 GPa - about as stiff as aluminum)
        ],
        inputs=[F_N, A_m2, L0_m, delta_L_m],
        label="Representative Young's Modulus Cases",
        examples_per_page=3,
        cache_examples=False,
    )


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