import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# Load your model from Hugging Face Hub
MODEL_ID = "Muhammadidrees/MedicalInsights"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Function to build structured input and query the LLM
def analyze(
    albumin, creatinine, glucose, crp, mcv, rdw, alp,
    wbc, lymph, age, gender, height, weight
):
    # Calculate BMI (hidden from user, only passed to LLM)
    try:
        height_m = height / 100  # cm → m
        bmi = round(weight / (height_m ** 2), 2)
    except Exception:
        bmi = "N/A"

    # System-style instruction (non-medical, fixed headings)
    system_prompt = (
        """Output MUST strictly follow this structured format:

    1. Executive Summary
       - Longevity Vitality Score (out of 100)
       - Top Priority Issues
       - Key Strengths

    2. System-Specific Analysis
       - Blood Health (MCV, RDW, Lymphocytes, WBC)
       - Protein & Liver Health (Albumin, ALP)
       - Kidney Health (Creatinine µmol/L)
       - Metabolic Health (Glucose mmol/L, lnCRP)
       - Other relevant systems

    3. Personalized Action Plan
       - Medical (tests/consults)
       - Nutrition (diet & supplements)
       - Lifestyle (hydration, exercise, sleep)
       - Testing (follow-up labs: ferritin, Vitamin D, GGT)

    4. Interaction Alerts
       - How biomarkers interact (e.g., anemia ↔ infection cycle, ALP with bone/liver origin)

    5. Longevity Metrics
       - Metabolic Health Score
       - Cardiovascular & Cognitive risk trajectory

    . Enhanced AI Insights & Longitudinal Risk

"""
    )

    # Construct profile input
    patient_input = f"""
Patient Profile:
- Age: {age}
- Gender: {gender}
- Height: {height} cm
- Weight: {weight} kg
- BMI: {bmi}

Lab Values:
- Albumin: {albumin} g/dL
- Creatinine: {creatinine} mg/dL
- Glucose: {glucose} mg/dL
- C-Reactive Protein: {crp} mg/L
- Mean Cell Volume: {mcv} fL
- Red Cell Distribution Width: {rdw} %
- Alkaline Phosphatase: {alp} U/L
- White Blood Cell Count: {wbc} K/uL
- Lymphocyte Percentage: {lymph} %
"""

    prompt = system_prompt + "\n" + patient_input

    # Call LLM
    result = pipe(
        prompt,
        max_new_tokens=5000,
        do_sample=True,
        temperature=0.3,
        top_p=0.9,
        return_full_text=False
    )

    # Force output to start from "Executive Summary"
    output_text = result[0]["generated_text"].strip()
    if "Executive Summary" in output_text:
        output_text = output_text.split("Executive Summary", 1)[-1]
        output_text = "Executive Summary" + output_text

    return output_text

# Build Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Wellness Insights AI — Enter Profile Data")

    with gr.Row():
        albumin = gr.Number(label="Albumin (g/dL)")
        wbc = gr.Number(label="White Blood Cell Count (K/uL)")

    with gr.Row():
        creatinine = gr.Number(label="Creatinine (mg/dL)")
        lymph = gr.Number(label="Lymphocyte Percentage (%)")

    with gr.Row():
        glucose = gr.Number(label="Glucose (mg/dL)")
        age = gr.Number(label="Age (years)")

    with gr.Row():
        crp = gr.Number(label="C-Reactive Protein (mg/L)")
        gender = gr.Dropdown(choices=["Male", "Female"], label="Gender")

    with gr.Row():
        mcv = gr.Number(label="Mean Cell Volume (fL)")
        height = gr.Number(label="Height (cm)")

    with gr.Row():
        rdw = gr.Number(label="Red Cell Distribution Width (%)")
        weight = gr.Number(label="Weight (kg)")

    with gr.Row():
        alp = gr.Number(label="Alkaline Phosphatase (U/L)")

    analyze_btn = gr.Button("🔎 Analyze")
    output = gr.Textbox(label="AI Wellness Assessment", lines=12)

    analyze_btn.click(
        fn=analyze,
        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
                wbc, lymph, age, gender, height, weight],
        outputs=output
    )

demo.launch()