import gradio as gr
import joblib
import pandas as pd
import os
from groq import Groq

from reportlab.platypus import (
    SimpleDocTemplate, Paragraph, Spacer,
    Table, TableStyle, Image
)
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib import colors
from reportlab.lib.units import inch

# ===============================
# LOAD MODEL
# ===============================
model = joblib.load("typhoidguard_model_reduced2.pkl")
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

last_llm_report = ""
last_basic_data = {}

LOGO_PATH = "logo1.png"   # 🔹 Place your logo file in HF repo

# ===============================
# LLM FUNCTION
# ===============================
def generate_llm_report(patient_info, prob_pct, risk_label, recommendation, factors):

    prompt = f"""
You are a responsible medical AI assistant helping rural healthcare units in Pakistan.

Patient Details:
{patient_info}

Predicted Risk: {prob_pct}%
Risk Category: {risk_label}

Clinical Risk Factors:
{", ".join(factors) if factors else "None"}

Agent Recommendation:
{recommendation}

Generate structured sections:
Clinical Interpretation
Possible Complications
Immediate Precautions
Preventative Measures
Follow-up Recommendations

Keep headings clear.
End with this line: "Remember, this advice is generated based on the provided information and should not replace the judgment of a qualified doctor. 
It is essential to consult a qualified doctor for personalized advice and treatment reminder to consult qualified doctor."
"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.4,
    )

    return response.choices[0].message.content


# ===============================
# BASIC PREDICTION
# ===============================
def predict_basic(name, age, gender, platelets, hb, duration,
                  blood_bacteria, severity, medication):

    global last_basic_data, last_llm_report
    last_llm_report = ""

    data = pd.DataFrame([{
        "Age": age,
        "Platelet Count": platelets,
        "Hemoglobin (g/dL)": hb,
        "Treatment Duration": duration,
        "Gender": gender,
        "Blood Culture Bacteria": blood_bacteria,
        "Symptoms Severity": severity,
        "Current Medication": medication
    }])

    prob = model.predict_proba(data)[0][1]
    prob_pct = round(prob * 100, 2)

    low_platelets = platelets < 150000
    low_hb = hb < 10

    if low_platelets or low_hb:
        prob_pct = 90.0
        risk_label = "HIGH RISK"
        color = "#b00020"
        recommendation = "Immediate referral to secondary care facility recommended."
    elif prob >= 0.65:
        risk_label = "HIGH RISK"
        color = "#b00020"
        recommendation = "Urgent clinical evaluation required."
    elif prob >= 0.35:
        risk_label = "MODERATE RISK"
        color = "#f57c00"
        recommendation = "Close monitoring advised."
    else:
        risk_label = "LOW RISK"
        color = "#2e7d32"
        recommendation = "Continue standard treatment protocol."

    factors = []
    if low_platelets:
        factors.append("Low Platelet Count")
    if low_hb:
        factors.append("Low Hemoglobin")

    last_basic_data = {
        "name": name,
        "age": age,
        "gender": gender,
        "platelets": platelets,
        "hb": hb,
        "duration": duration,
        "bacteria": blood_bacteria,
        "severity": severity,
        "medication": medication,
        "prob_pct": prob_pct,
        "risk_label": risk_label,
        "recommendation": recommendation,
        "factors": factors,
        "color": color
    }

    urdu = f"""
    <div style='background:#eef3fb;padding:10px;border-radius:8px;'>
    <b>اردو خلاصہ:</b><br>
    علاج ناکامی کا خطرہ <b>{prob_pct}%</b> ہے۔ درجہ: <b>{risk_label}</b>۔<br>
    ڈاکٹر سے مشورہ ضرور کریں۔
    </div>
    """

    return f"""
    <div style="font-family:Arial;padding:15px;">
    <h2 style="color:{color};">Patient: {name}</h2>

    <h3>
    Risk of Treatment Failure:
    <span style="background:{color};color:white;padding:6px 12px;border-radius:6px;">
    {prob_pct}% - {risk_label}
    </span>
    </h3>

    <h4>Clinical Risk Factors:</h4>
    <ul>
    {''.join([f"<li>{f}</li>" for f in factors]) if factors else "<li>None</li>"}
    </ul>

    <h4>Agent Recommendation:</h4>
    <p><b>{recommendation}</b></p>

    {urdu}
    </div>
    """


# ===============================
# DETAILED BUTTON
# ===============================
def generate_detailed(*inputs):
    global last_basic_data, last_llm_report

    basic_html = predict_basic(*inputs)
    d = last_basic_data

    patient_info = f"""
Age: {d['age']}
Gender: {d['gender']}
Platelets: {d['platelets']}
Hemoglobin: {d['hb']}
Duration: {d['duration']}
Bacteria: {d['bacteria']}
Severity: {d['severity']}
Medication: {d['medication']}
"""

    last_llm_report = generate_llm_report(
        patient_info,
        d["prob_pct"],
        d["risk_label"],
        d["recommendation"],
        d["factors"]
    )

    lines = last_llm_report.split("\n")

    formatted_parts = []
    disclaimer_html = ""

    for line in lines:
        line = line.strip()
        if not line:
            continue

        # Detect disclaimer line
        if "consult a qualified doctor" in line.lower():
            disclaimer_html = f"""
            <div style="
                margin-top:15px;
                padding:12px;
                background:#fff3cd;
                border-left:5px solid #ff9800;
                border-radius:6px;
                font-size:14px;">
                <b>Disclaimer:</b><br>
                {line}
            </div>
            """
            continue

        # Detect headings starting with ##
        if line.startswith("##"):
            clean_heading = line.replace("##", "").strip()
            formatted_parts.append(
                f"<br><b style='font-size:16px;color:#003366;'>{clean_heading}:</b><br>"
            )
        else:
            formatted_parts.append(f"{line}<br>")

    formatted = "".join(formatted_parts)

    return basic_html + f"""
    <hr>
    <h3 style="color:#003366;">Detailed AI Report & Recommendations</h3>
    <div style="background:#f9f9f9;padding:15px;border-radius:8px;">
    {formatted}
    {disclaimer_html}
    </div>
    """


# ===============================
# PDF GENERATION
# ===============================
def generate_pdf(*inputs):

    global last_basic_data, last_llm_report

    d = last_basic_data
    file_path = f"{d['name']}_Typhoid_Report.pdf"
    doc = SimpleDocTemplate(file_path)
    elements = []
    styles = getSampleStyleSheet()

    # ---- Custom Styles ----
    heading_style = ParagraphStyle(
        name="CustomHeading",
        parent=styles["Heading2"],
        fontSize=14,
        spaceAfter=6,
        textColor=colors.black
    )

    bullet_style = ParagraphStyle(
        name="BulletStyle",
        parent=styles["Normal"],
        leftIndent=15,
        bulletIndent=5,
        spaceAfter=4
    )

    disclaimer_heading_style = ParagraphStyle(
        name="DisclaimerHeading",
        parent=styles["Heading2"],
        textColor=colors.red,
        fontSize=13,
        spaceBefore=15
    )

    disclaimer_text_style = ParagraphStyle(
        name="DisclaimerText",
        parent=styles["Normal"],
        fontSize=10,
        textColor=colors.grey
    )

    # ---- LOGO ----
    if os.path.exists(LOGO_PATH):
        elements.append(Image(LOGO_PATH, width=1.6*inch, height=1.2*inch))
        elements.append(Spacer(1,0.2*inch))

    elements.append(Paragraph("<b>TyphoidGuard AI Clinical Report</b>", styles["Title"]))
    elements.append(Spacer(1, 0.3 * inch))

    # ---- RISK BADGE (Improved Height & Padding) ----
    risk_color = colors.HexColor(d["color"])
    badge_style = ParagraphStyle(
        name="Badge",
        parent=styles["Normal"],
        backColor=risk_color,
        textColor=colors.white,
        fontSize=13,
        leading=18,
        leftIndent=10,
        rightIndent=10,
        spaceBefore=10,
        spaceAfter=15
    )

    elements.append(Paragraph(
        f"<b>Risk: {d['prob_pct']}% - {d['risk_label']}</b>",
        badge_style
    ))
    elements.append(Spacer(1,0.3*inch))

    # ---- PATIENT TABLE ----
    table_data = [
        ["Patient Name", d["name"]],
        ["Age", d["age"]],
        ["Gender", d["gender"]],
        ["Platelet Count", d["platelets"]],
        ["Hemoglobin", d["hb"]],
        ["Treatment Duration", d["duration"]],
        ["Bacteria", d["bacteria"]],
        ["Severity", d["severity"]],
        ["Medication", d["medication"]],
    ]

    table = Table(table_data, colWidths=[2.5*inch, 3*inch])
    table.setStyle(TableStyle([
        ('GRID',(0,0),(-1,-1),1,colors.grey),
        ('FONTNAME',(0,0),(-1,-1),'Helvetica')
    ]))

    elements.append(table)
    elements.append(Spacer(1,0.4*inch))

    # ---- AGENT RECOMMENDATION ----
    elements.append(Paragraph("<b>Agent Recommendation</b>", heading_style))
    elements.append(Paragraph(d["recommendation"], styles["Normal"]))
    elements.append(Spacer(1,0.4*inch))

    # ---- DETAILED AI REPORT ----
    if last_llm_report:
        elements.append(Paragraph("<b>Detailed AI Report</b>", heading_style))
        elements.append(Spacer(1,0.2*inch))

        lines = last_llm_report.split("\n")

        for line in lines:
            line = line.strip()

            if not line:
                continue

            if "qualified doctor" in line.lower():
                continue

            if line.startswith("##"):
                clean_heading = line.replace("##", "").strip()
                elements.append(Spacer(1,0.2*inch))
                elements.append(
                    Paragraph(f"<b>{clean_heading}:</b>", heading_style)
                )
            else:
                elements.append(Paragraph(line, styles["Normal"]))

        elements.append(Spacer(1,0.5*inch))

    # ---- DISCLAIMER ----
    elements.append(Paragraph("Disclaimer", disclaimer_heading_style))
    elements.append(Spacer(1,0.1*inch))
    elements.append(Paragraph(
        "Remember to consult a qualified doctor for personalized advice and treatment. "
        "They will be able to provide guidance tailored to the patient's specific needs and circumstances.",
        disclaimer_text_style
    ))

    elements.append(Spacer(1,0.5*inch))

    # ---- TEAM ----
    elements.append(Paragraph("<b>Team Chinar AI (AJK)</b>", styles["Normal"]))
    elements.append(Paragraph(
        "Anees Qumar Abbasi (Team Lead) | Sharafat Hussain | Salma Asghar | "
        "Kaleem Hussain | Munazza Zahra | Kokub Khhurishid",
        styles["Normal"]
    ))

    doc.build(elements)

    # ✅ SUCCESS POPUP MESSAGE (Added Feature)
    
    gr.Info("Report Generated Successfully! Please download it below.")

    return file_path
# ===============================
# INTERFACE
# ===============================
# ===============================
# PREMIUM MEDICAL INTERFACE (UPDATED)
# ===============================
with gr.Blocks(
    theme=gr.themes.Soft(),
    css="""
    body {
        background: linear-gradient(135deg, #e3f2fd, #f4f9ff);
    }

    .center-header {
        text-align: center;
    }
    
    #header-center {
    text-align: center;
    }

    .instruction-text {
        text-align: center;
        font-size: 16px;
        color: #37474f;
        margin-bottom: 20px;
    }

    .gr-button {
        border-radius: 10px !important;
        font-weight: 600 !important;
        padding: 10px 18px !important;
        font-size: 15px !important;
    }

    .gr-textbox, .gr-number, .gr-dropdown {
        border-radius: 8px !important;
    }

    .footer-team {
        background: #ffffff;
        padding: 15px;
        border-radius: 10px;
        margin-top: 20px;
        text-align: center;
        font-size: 14px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.05);
    }
    """
) as demo:

    # ---- LOGO (Guaranteed to Display in HF) ----
    # ---- LOGO (Centered) ----
    '''gr.Markdown('<div style="text-align:center;">')
    gr.Image(LOGO_PATH, width=300, show_label=False, container=False)
    gr.Markdown('</div>')
    # ---- CENTERED HEADER ----
    gr.Markdown("""
    <div class="center-header">
        <h1 style="color:#0d47a1;">🏥 TyphoidGuard AI</h1>
        <h4 style="color:#37474f;">HEC Generative AI Hackathon – Pakistan</h4>
        <h4 style="color:#00695c;">Team Chinar AI (AJK)</h4>
    </div>
    """)'''
   
    # ---- CENTERED LOGO + HEADER ----
    with gr.Row():
        with gr.Column(scale=1, elem_id="header-center"):
            gr.Image(LOGO_PATH, width=300, show_label=False)
            gr.Markdown("""
            <div style="text-align:center;">
                <h1 style="color:#0d47a1; margin-top:15px;">🏥 TyphoidGuard AI</h1>
                <h4 style="color:#37474f; margin:5px 0;">HEC Generative AI Hackathon – Pakistan</h4>
                <h4 style="color:#00695c; margin:5px 0;">Team Chinar AI (AJK)</h4>
            </div>
            """)

    # ---- INSTRUCTION TEXT ----
    gr.Markdown("""
    <div class="instruction-text">
        Please enter the patient clinical values below to assess 
        the risk of typhoid treatment failure and receive AI-based recommendations.
    </div>
    """)

    # ---- INPUT FIELDS ----
    with gr.Row():
        name = gr.Textbox(label="Patient Name")
        age = gr.Number(label="Age")
        gender = gr.Radio(["Male", "Female"], label="Gender")

    with gr.Row():
        platelets = gr.Number(label="Platelet Count")
        hb = gr.Number(label="Hemoglobin (g/dL)")
        duration = gr.Number(label="Treatment Duration (Days)")

    with gr.Row():
        blood_bacteria = gr.Dropdown(
            ["Escherichia coli","Salmonella typhi","Unknown","Staphylococcus"],
            label="Blood Culture Bacteria"
        )
        severity = gr.Dropdown(
            ["Low","Moderate","High","Unknown"],
            label="Symptoms Severity"
        )
        medication = gr.Dropdown(
            ["Amoxicillin","Ceftriaxone","Azithromycin"],
            label="Current Medication"
        )

    # ---- BUTTONS ----
    with gr.Row():
        predict_btn = gr.Button("Basic Risk Assessment")
        detailed_btn = gr.Button("Detailed Report & Recommendations")
        pdf_btn = gr.Button("Generate PDF Report")

    output = gr.HTML()

    predict_btn.click(predict_basic,
        inputs=[name,age,gender,platelets,hb,duration,
                blood_bacteria,severity,medication],
        outputs=output)

    detailed_btn.click(generate_detailed,
        inputs=[name,age,gender,platelets,hb,duration,
                blood_bacteria,severity,medication],
        outputs=output)

    pdf_btn.click(generate_pdf,
        inputs=[name,age,gender,platelets,hb,duration,
                blood_bacteria,severity,medication],
        outputs=gr.File())

    # ---- TEAM FOOTER ----
    gr.Markdown("""
    <div class="footer-team">
        <b>Team Chinar AI (AJK)</b><br>
        Anees Qumar Abbasi (Team Lead) | Sharafat Hussain | Salma Asghar |
        Kaleem Hussain | Munazza Zahra | Kokub Khurishid
    </div>
    """)

demo.launch()