import pandas as pd
import gradio as gr


def risk_stratify(file):
    try:
        df = pd.read_csv(file.name)
    except Exception as e:
        return f"❌ Error reading file: {e}"

    required_cols = [
        "member_id",
        "age",
        "chronic_conditions",
        "ed_visits_past_year",
        "inpatient_admissions",
        "total_cost",
        "behavioral_health_flag",
        "pharmacy_rx_count",
        "social_risk_score"
    ]

    missing = [c for c in required_cols if c not in df.columns]
    if missing:
        return f"❌ Missing required columns: {missing}"

    # Risk scoring
    df["risk_score"] = 0
    df["risk_score"] += (df["age"] >= 65).astype(int)
    df["risk_score"] += (df["chronic_conditions"] >= 2).astype(int) * 2
    df["risk_score"] += (df["ed_visits_past_year"] >= 3).astype(int) * 2
    df["risk_score"] += (df["inpatient_admissions"] >= 1).astype(int) * 2
    df["risk_score"] += (df["total_cost"] >= 20000).astype(int) * 2
    df["risk_score"] += (df["behavioral_health_flag"] == 1).astype(int) * 2
    df["risk_score"] += (df["pharmacy_rx_count"] >= 10).astype(int)
    df["risk_score"] += (df["social_risk_score"] >= 5).astype(int)

    # Categorization
    def categorize(score):
        if score >= 6:
            return "High Risk"
        elif score >= 3:
            return "Moderate Risk"
        else:
            return "Low Risk"

    df["risk_category"] = df["risk_score"].apply(categorize)

    # Summary output
    counts = df["risk_category"].value_counts().to_dict()

    summary = f"""
# Medicaid Risk Stratification Results

### **Member Count by Risk Group:**
- High Risk: **{counts.get("High Risk", 0)}**
- Moderate Risk: **{counts.get("Moderate Risk", 0)}**
- Low Risk: **{counts.get("Low Risk", 0)}**

---

### **Method Summary**
This risk stratification uses:
- Chronic disease counts  
- ED and inpatient utilization  
- Annual cost  
- Pharmacy use  
- Behavioral health indicator  
- Social risk  
- Member age  

A risk score ≥ 6 indicates **High Need** for care management outreach.

Upload another file to re-run the analysis.
"""

    return summary


with gr.Blocks() as demo:
    gr.Markdown("""
# Medicaid Risk Stratification Tool  
Upload a Medicaid member dataset to calculate risk scores.

### Required CSV Columns:
`member_id, age, chronic_conditions, ed_visits_past_year, inpatient_admissions, total_cost, behavioral_health_flag, pharmacy_rx_count, social_risk_score`
""")

    file = gr.File(label="Upload Medicaid CSV")
    output = gr.Markdown()

    file.upload(risk_stratify, file, output)

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