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

# Load your model from Hugging Face Hub
MODEL_ID = "Muhammadidrees/my-gpt-oss"

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
    try:
        height_m = height / 100  # cm → m
        bmi = round(weight / (height_m ** 2), 2)
    except Exception:
        bmi = "N/A"

    # Improved system prompt with clearer instructions
    system_prompt = """You are a professional AI Medical Assistant analyzing patient biomarkers.

CRITICAL: Generate a COMPLETE report following this EXACT structure. Do not stop mid-sentence.

=== REQUIRED OUTPUT FORMAT ===

1. Executive Summary
   - Top Priority Issues: [List 2-3 main concerns]
   - Key Strengths: [List 2-3 positive findings]

2. System-Specific Analysis
   - Blood Health (MCV, RDW, Lymphocytes, WBC)
   - Protein & Liver Health (Albumin, ALP)
   - Kidney Health (Creatinine)
   - Metabolic Health (Glucose, CRP)

3. Personalized Action Plan
   - Medical: [Recommended tests/consultations]
   - Nutrition: [Dietary recommendations and supplements]
   - Lifestyle: [Exercise, hydration, sleep guidance]
   - Testing: [Follow-up labs needed]

4. Interaction Alerts
   [Explain how biomarkers interact and influence each other]

6. Tabular Mapping
   | Biomarker | Value | Status | AI Insight | Client Message |
   |-----------|-------|--------|------------|----------------|
   [Complete table for all biomarkers]

7. Enhanced AI Insights
   - Subclinical Nutrient Analysis (Iron, B12, Folate status)
   - ALP Interpretation (bone vs liver origin)
   - Immune System Assessment (WBC & lymphocyte trends)
   - Long-term Health Considerations

=== END FORMAT ===

Now analyze the following patient data and provide a COMPLETE report:"""

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

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

Generate complete analysis now:"""

    prompt = system_prompt + "\n" + patient_input

    try:
        # Increased max_new_tokens significantly for complete output
        result = pipe(
            prompt,
            max_new_tokens=2500,  # INCREASED from 1000
            do_sample=True,
            temperature=0.7,  # INCREASED from 0.3 for better generation
            top_p=0.92,
            top_k=50,
            repetition_penalty=1.1,  # Prevent repetition
            return_full_text=False,
            pad_token_id=tokenizer.eos_token_id
        )

        output_text = result[0]["generated_text"].strip()
        
        # Clean up output - remove any prompt leakage
        if "Executive Summary" in output_text:
            idx = output_text.find("Executive Summary")
            output_text = output_text[idx:]
        elif "1. Executive Summary" in output_text:
            idx = output_text.find("1. Executive Summary")
            output_text = output_text[idx:]
        
        # If output seems incomplete, add a note
        if len(output_text) < 500:
            output_text += "\n\n⚠️ Note: Output may be incomplete. Consider re-running the analysis."
        
        return output_text
    
    except Exception as e:
        return f"Error during analysis: {str(e)}\n\nPlease check your input values and try again."


# Build Gradio UI with improved layout
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🧪 Wellness Insights AI
    ### Enter Patient Profile & Lab Values for Comprehensive Analysis
    """)

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 👤 Demographics")
            age = gr.Number(label="Age (years)", value=30)
            gender = gr.Dropdown(choices=["Male", "Female"], label="Gender", value="Male")
            height = gr.Number(label="Height (cm)", value=175)
            weight = gr.Number(label="Weight (kg)", value=70)
        
        with gr.Column():
            gr.Markdown("### 🩸 Blood Health Markers")
            wbc = gr.Number(label="White Blood Cell Count (K/uL)", value=7.0)
            lymph = gr.Number(label="Lymphocyte Percentage (%)", value=30)
            mcv = gr.Number(label="Mean Cell Volume (fL)", value=90)
            rdw = gr.Number(label="Red Cell Distribution Width (%)", value=13)

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🫀 Metabolic Markers")
            glucose = gr.Number(label="Glucose (mg/dL)", value=95)
            crp = gr.Number(label="C-Reactive Protein (mg/L)", value=1.5)
        
        with gr.Column():
            gr.Markdown("### 🧬 Organ Function Markers")
            albumin = gr.Number(label="Albumin (g/dL)", value=4.2)
            creatinine = gr.Number(label="Creatinine (mg/dL)", value=1.0)
            alp = gr.Number(label="Alkaline Phosphatase (U/L)", value=70)

    analyze_btn = gr.Button("🔎 Generate Comprehensive Analysis", variant="primary", size="lg")
    
    gr.Markdown("### 📋 Analysis Report")
    output = gr.Textbox(
        label="AI-Generated Lab Report",
        lines=25,
        max_lines=50,
        show_copy_button=True
    )

    analyze_btn.click(
        fn=analyze,
        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
                wbc, lymph, age, gender, height, weight],
        outputs=output
    )
    
    gr.Markdown("""
    ---
    **Note:** This tool provides educational insights based on biomarker analysis. 
    Always consult healthcare professionals for medical advice.
    """)

demo.launch(share=False)