import gradio as gr
import pandas as pd
import io
import zipfile
import matplotlib.pyplot as plt
import seaborn as sns
from logic.care_gap_engine import evaluate_care_gaps
from logic.financial_model import estimate_financial_recovery
from logic.nlp_report import generate_patient_summary, summarize_data_quality
from logic.exporter import export_to_docx_and_zip

CONFIG = {
    "base_rate": 9500,
    "care_gap_rules": {
        "Breast Cancer Screening": {"gender": "F", "min_age": 50, "max_age": 74, "interval_days": 730},
        "Colorectal Cancer Screening": {"min_age": 50, "max_age": 75, "interval_days": 365*10},
        "Blood Pressure Control": {"bp_threshold": 140},
        "Follow-Up Care": {},
        "Readmission Risk": {}
    },
    "sdoh_modifiers": [
        "Housing_Instability", "Food_Insecurity", "Transportation_Access", "Health_Literacy"
    ]
}

import os

def run_analysis(file, tone, gap_filter, custom_prompt):
    # If no file is uploaded, use default CSV from data folder
    if file is None:
        default_path = os.path.join(os.path.dirname(__file__), "data", "Enriched_Patient_Dataset.csv")
        df = pd.read_csv(default_path)
    else:
        df = pd.read_csv(file.name)
    df = df.convert_dtypes()  # safer type conversion

    quality_issues_df = summarize_data_quality(df)

    care_gaps = evaluate_care_gaps(df, CONFIG)
    financials = estimate_financial_recovery(care_gaps, df, CONFIG['base_rate'], CONFIG['sdoh_modifiers'])
    merged = care_gaps.merge(financials, on="patient_id")

    if gap_filter:
        merged = merged[merged['care_gaps'].apply(lambda gaps: gap_filter in gaps)]

    summaries = {}
    for _, row in merged.iterrows():
        summaries[row['patient_id']] = generate_patient_summary(row, tone, custom_prompt)

    zip_bytes, file_names = export_to_docx_and_zip(merged, summaries)

    # Risk score graph
    plt.figure(figsize=(8, 4))
    plt.hist(merged['risk_score'], bins=10)
    plt.title("Risk Score Distribution")
    plt.xlabel("Risk Score")
    plt.ylabel("Number of Patients")
    plt.tight_layout()
    plt.savefig("risk_score_plot.png")

    # Gap frequency bar chart
    gap_freq = {}
    for gaps in merged['care_gaps']:
        for g in gaps:
            gap_freq[g] = gap_freq.get(g, 0) + 1
    gap_df = pd.DataFrame(list(gap_freq.items()), columns=["Gap", "Count"])
    plt.figure(figsize=(10, 5))
    sns.barplot(x="Count", y="Gap", data=gap_df.sort_values(by="Count", ascending=False))
    plt.title("Care Gap Frequency")
    plt.tight_layout()
    plt.savefig("care_gap_freq.png")

    # Top 10 revenue opportunities
    top10 = merged.sort_values(by="projected_gain", ascending=False).head(10)
    top10_chart_path = "top10_revenue.png"
    plt.figure(figsize=(10, 5))
    sns.barplot(x="projected_gain", y="patient_id", data=top10)
    plt.title("Top 10 Patients by Projected Gain")
    plt.xlabel("Projected Revenue Recovery ($)")
    plt.ylabel("Patient ID")
    plt.tight_layout()
    plt.savefig(top10_chart_path)

    return quality_issues_df, merged, zip_bytes, "risk_score_plot.png", "care_gap_freq.png", top10_chart_path

iface = gr.Interface(
    fn=run_analysis,
    inputs=[
        gr.File(label="Upload Patient CSV", file_count="single", type="filepath"),
        gr.Dropdown(choices=["clinical", "executive", "casual"], label="Report Tone", value="executive"),
        gr.Textbox(label="Filter by Care Gap (e.g. 'Breast Cancer Screening')", value=""),
        gr.Textbox(label="Custom Prompt Override (optional)", value="")
    ],
    outputs=[
        gr.Dataframe(label="⚠️ Data Quality Issues"),
        gr.Dataframe(label="Care Gap & Financial Report"),
        gr.File(label="Download All Reports as ZIP"),
        gr.Image(label="Risk Score Distribution"),
        gr.Image(label="Care Gap Frequency"),
        gr.Image(label="Top 10 Revenue Opportunities")
    ],
    title="AI Medicare Advantage Analyzer",
    description="Upload a patient dataset, or leave blank to use the sample data. Evaluate CMS care gaps and SDOH risk, forecast financial recovery, generate NLP summaries, and validate data quality."
)

if __name__ == "__main__":
    iface.launch(share=True)