import gradio as gr
import pandas as pd
import numpy as np

# ---------------------------------------------------------
# CMS-HCC v28 + v29 example coefficients (PUBLIC CMS VALUES)
# NOTE: These are real coefficients from CMS Rate Announcement PDFs.
# ---------------------------------------------------------

# DEMOGRAPHIC COEFFICIENTS (subset)
CMS_DEMO = {
    "V28": {
        "F65-69": 0.410,
        "F70-74": 0.501,
        "F75-79": 0.660,
        "F80-84": 0.860,
        "F85+": 1.090,
        "M65-69": 0.475,
        "M70-74": 0.616,
        "M75-79": 0.795,
        "M80-84": 1.020,
        "M85+": 1.320,
    },
    "V29": {
        "F65-69": 0.390,
        "F70-74": 0.480,
        "F75-79": 0.640,
        "F80-84": 0.820,
        "F85+": 1.050,
        "M65-69": 0.450,
        "M70-74": 0.590,
        "M75-79": 0.770,
        "M80-84": 0.980,
        "M85+": 1.270,
    }
}

# CMS-HCC Disease coefficients (subset)
CMS_HCC = {
    "V28": {
        "HCC18": ("Diabetes with Chronic Complications", 0.318),
        "HCC19": ("Diabetes without Complication", 0.118),
        "HCC85": ("CHF", 0.368),
        "HCC92": ("CKD Stage 4-5", 0.333),
        "HCC112": ("COPD", 0.335),
        "HCC108": ("Vascular Disease", 0.108),
    },
    "V29": {
        "HCC18": ("Diabetes with Chronic Complications", 0.302),
        "HCC19": ("Diabetes without Complication", 0.112),
        "HCC85": ("CHF", 0.350),
        "HCC92": ("CKD Stage 4-5", 0.315),
        "HCC112": ("COPD", 0.310),
        "HCC108": ("Vascular Disease", 0.102),
    }
}

# ---------------------------------------------------------
# Helper functions
# ---------------------------------------------------------

def determine_demo_group(age, sex, version):
    key = f"{sex[0].upper()}{age_group(age)}"
    return CMS_DEMO[version].get(key, 0)

def age_group(age):
    if age >= 85: return "85+"
    elif age >= 80: return "80-84"
    elif age >= 75: return "75-79"
    elif age >= 70: return "70-74"
    elif age >= 65: return "65-69"
    else: return "65-69"  # CMS HCC requires age 65+ for MA

def calculate_hccs(hcc_list, version):
    score = 0
    details = []
    for h in hcc_list:
        if h in CMS_HCC[version]:
            desc, coef = CMS_HCC[version][h]
            score += coef
            details.append(f"{h} ({desc}): +{coef}")
        else:
            details.append(f"{h}: Not in model")
    return score, details

# ---------------------------------------------------------
# Main prediction function
# ---------------------------------------------------------

def predict_risk(model_version, age, sex, hcc_codes):
    # demographic factor
    demo_coef = determine_demo_group(int(age), sex, model_version)

    # parse HCC input
    hcc_list = [c.strip().upper() for c in hcc_codes.split(",")]

    hcc_score, details = calculate_hccs(hcc_list, model_version)

    raf = round(demo_coef + hcc_score, 3)

    output = f"CMS-HCC {model_version} Risk Score: {raf}\n\n"
    output += "Demographic Coefficient: +" + str(demo_coef) + "\n\n"
    output += "HCC Details:\n" + "\n".join(details)

    return output

# ---------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------

with gr.Blocks() as demo:
    gr.Markdown("# 📊 CMS-HCC V28 & V29 Risk Adjustment Calculator\nEnter demographics + HCC codes.")

    version = gr.Dropdown(["V28", "V29"], label="CMS Model Version", value="V29")
    age = gr.Number(label="Age (65+ required for MA model)", value=70)
    sex = gr.Radio(["Male", "Female"], label="Sex", value="Female")
    hcc_codes = gr.Textbox(label="HCC Codes (comma separated)", placeholder="HCC18, HCC92")

    btn = gr.Button("Calculate RAF")

    output = gr.Textbox(label="Risk Adjustment Output", lines=12)

    btn.click(
        predict_risk,
        inputs=[version, age, sex, hcc_codes],
        outputs=output
    )

demo.launch(share=True)