import math
def fin_array_calc(N: int, T_inf_C: float, T_b_C: float, L_m: float, w_m: float, t_m: float, h_W_m2K: float, k_W_mK: float) -> dict:
    """
    Straight rectangular fin array, adiabatic tip.
    Returns fin geometry, fin parameter, heat transfer per fin, and total heat transfer.
    """

    values = [N, T_inf_C, T_b_C, L_m, w_m, t_m, h_W_m2K, k_W_mK]
    if not all(math.isfinite(v) for v in values):
        raise ValueError("All inputs must be finite.")
    if N <= 0 or min(L_m, w_m, t_m, h_W_m2K, k_W_mK) <= 0:
        raise ValueError("Number of fins, dimensions, convection coefficient, and conductivity must be positive.")
    if T_b_C <= T_inf_C:
        raise ValueError("Base temperature must be greater than ambient temperature.")
    if N != int(N):
        raise ValueError("Number of fins must be an integer.")

    # Calculate fin geometry
    A_c = w_m * t_m
    P = 2.0 * (w_m + t_m)

    # Base excess temperature
    theta_b = T_b_C - T_inf_C

    # Fin parameter
    m = math.sqrt(h_W_m2K * P / (k_W_mK * A_c))
    M = math.sqrt(h_W_m2K * P * k_W_mK * A_c) * theta_b

    # Heat transfer for adiabatic tip
    q_fin_W = M * math.tanh(m * L_m)
    q_total_W = N * q_fin_W

    return dict(
        results={
            "A_c_m2": A_c,
            "P_m": P,
            "theta_b_C": theta_b,
            "m_1_per_m": m,
            "mL": m * L_m,
            "q_fin_W": q_fin_W,
            "q_total_W": q_total_W,
        },
        assumptions={
            "adiabatic_tip": True,
            "steady_state": True,
            "one_dimensional": True,
            "constant_h": True,
            "constant_T_inf": True,
            "constant_T_b": True,
            "constant_k": True,
            "uniform_cross_section": True,
            "identical_fins": True,
            "perfect_base_contact": True,
            "negligible_radiation": True,
            "no_fin_interaction": True,
        },
    )

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

MODEL_ID = "Qwen/Qwen3-0.6B"
MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, revision=MODEL_REVISION, dtype=torch.float32
).eval()  # CPU: works without a GPU runtime or a separate serving package.

def generate(messages, max_tokens=96):
    inputs = tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        enable_thinking=False, return_dict=True, return_tensors="pt",
    )
    config = GenerationConfig(
        max_new_tokens=max_tokens, do_sample=False,
        eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.eos_token_id,
    )
    model.generation_config = config  # Keep sampling defaults from overriding this classroom configuration.
    with torch.inference_mode():
        output = model.generate(**inputs, generation_config=config)
    # Decode only the new tokens, not the prompt.
    new_tokens = output[0, inputs["input_ids"].shape[-1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    if new_tokens[-1].item() != tokenizer.eos_token_id:
        text += "\n[Generation reached its length limit; this explanation may be incomplete.]"
    return text

import pandas as pd
import gradio as gr

TABLE_COLUMNS = ["Quantity", "Calculated", "Units"]

# Two examples teach the model how to describe the fin-array calculation.
FEW_SHOT_MESSAGES = [
    {"role": "user", "content": "Heat transfer per fin: 8.000 W. Number of fins: 10. Total heat transfer: 80.000 W. Base temperature: 100.000 C. Ambient temperature: 25.000 C."},
    {"role": "assistant", "content": "Each fin transfers 8 W from the base to the surroundings. With 10 identical fins, the total heat transfer is 80 W."},
    {"role": "user", "content": "Heat transfer per fin: 4.000 W. Number of fins: 20. Total heat transfer: 80.000 W. Base temperature: 80.000 C. Ambient temperature: 30.000 C."},
    {"role": "assistant", "content": "Each fin transfers 4 W from the base to the surroundings. With 20 identical fins, the total heat transfer is 80 W."},
]

def run_once(num_fins, ambient, base, length, width, thickness, h, conductivity):
    try:
        vals = [float(x) for x in [num_fins, ambient, base, length, width, thickness, h, conductivity]]
        result = fin_array_calc(int(vals[0]), vals[1], vals[2], vals[3], vals[4], vals[5], vals[6], vals[7])
    except (TypeError, ValueError, ArithmeticError) as exc:
        # Clear stale evidence and prose together when an input is invalid.
        return pd.DataFrame(columns=TABLE_COLUMNS), "Correct the inputs: " + str(exc)
    r = result["results"]
    table = pd.DataFrame([
        ["Cross-sectional area", f"{r['A_c_m2']:.6g}", "m²"],
        ["Fin perimeter", f"{r['P_m']:.6g}", "m"],
        ["Base excess temperature", f"{r['theta_b_C']:.3f}", "°C"],
        ["Fin parameter, m", f"{r['m_1_per_m']:.3f}", "1/m"],
        ["mL", f"{r['mL']:.3f}", "—"],
        ["Heat transfer per fin", f"{r['q_fin_W']:.3f}", "W"],
        ["Total array heat transfer", f"{r['q_total_W']:.3f}", "W"],
    ], columns=TABLE_COLUMNS)
    facts = (
        f"Heat transfer per fin: {r['q_fin_W']:.3f} W. Number of fins: {int(vals[0])}. Total heat transfer: {r['q_total_W']:.3f} W. "
        f"Base temperature: {vals[2]:.3f} C. Ambient temperature: {vals[1]:.3f} C. Fin parameter m: {r['m_1_per_m']:.3f} 1/m. mL: {r['mL']:.3f}."
    )
    explanation = generate([
        {"role": "system", "content": "Write exactly two short sentences, at most 35 words total. Sentence 1: state heat transfer per fin. Sentence 2: state total array heat transfer. Copy the supplied numbers exactly. Do not repeat values, recalculate, invent causes, or recommend design changes."},
        *FEW_SHOT_MESSAGES,
        {"role": "user", "content": facts},
    ])
    return table, explanation

with gr.Blocks() as demo:
    gr.Markdown("# Fin array calculation and explanation\nIdealized rectangular fins with adiabatic tips. Compare the generated description against the calculated values.")
    gr.Markdown(
      """
    **Assumptions**
    - Steady-state, one-dimensional conduction
    - Adiabatic fin tips
    - Constant `h`, `k`, `T∞`, and `Tb`
    - Uniform rectangular cross-section
    - Identical fins with perfect base contact
    - Negligible radiation and fin-to-fin interaction
    """
    )
    gr.Markdown(
    """
    **Valid inputs**
    - Number of fins: positive integer
    - Fin length, width, and thickness: greater than 0 m
    - Convection coefficient `h`: greater than 0 W/m²-K
    - Thermal conductivity `k`: greater than 0 W/m-K
    - Base temperature `Tb` must be greater than ambient temperature `T∞`
    """
    )
    with gr.Row():
        num_fins = gr.Number(value=10, precision=0, label="Number of fins")
        ambient = gr.Number(value=25, label="Ambient temperature (°C)")
        base = gr.Number(value=100, label="Base temperature (°C)")
        h = gr.Number(value=25, label="Convection coefficient (W/m²-K)")
    with gr.Row():
        length = gr.Number(value=0.05, label="Fin length (m)")
        width = gr.Number(value=0.03, label="Fin width (m)")
        thickness = gr.Number(value=0.002, label="Fin thickness (m)")
        conductivity = gr.Number(value=205, label="Thermal conductivity (W/m-K)")
    table = gr.Dataframe(headers=TABLE_COLUMNS, datatype="str", label="Calculated evidence — source of the results", interactive=False)
    prose = gr.Textbox(label="Generated explanation — verify against the table", lines=4, interactive=False)
    gr.Button("Calculate and explain").click(
        run_once, [num_fins, ambient, base, length, width, thickness, h, conductivity], [table, prose],
        concurrency_limit=1, api_name="run_once",
    )

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