from __future__ import annotations
import gradio as gr
from src.config import APP_TITLE, APP_MODE
from src.styles import APP_CSS
from src.inference import get_predictor
from src.imaging import nifti_preview
from src.ui_helpers import summary_html, interpretation_html, survival_figure

predictor = get_predictor()

HERO='''<section class="hero"><div class="eyebrow">AI-BEACON · Multimodal oncology research</div><h1>Individualized breast cancer prognosis</h1><p>Combine longitudinal MRI, segmentation, radiology text, and structured clinical variables in one transparent prognostic workflow.</p><div class="badge-row"><span class="badge">Research use only</span><span class="badge">Multimodal model</span><span class="badge">Mode: %s</span></div></section>''' % APP_MODE.title()
FLOW='''<div class="flow"><div class="flow-item"><b>01 · Eligibility</b>Confirm the intended cohort</div><div class="flow-item"><b>02 · Imaging</b>Upload MRI and mask data</div><div class="flow-item"><b>03 · Clinical profile</b>Review disease variables</div><div class="flow-item"><b>04 · Prognosis</b>Generate and interpret output</div></div>'''

def section(title, subtitle):
    return f'<div class="section-title"><h3>{title}</h3><p>{subtitle}</p></div>'

def run_analysis(pre_mri, post_mri, mask, report, age, pre_stage, post_stage, family_history,
                 er, pr, her2, tumor_type, treatment, distant_metastasis):
    try:
        result=predictor.predict(pre_mri=pre_mri,post_mri=post_mri,mask=mask,report=report,age=age,
            pre_stage=pre_stage,post_stage=post_stage,family_history=family_history,er=er,pr=pr,her2=her2,
            tumor_type=tumor_type,treatment=treatment,distant_metastasis=distant_metastasis)
        pre=nifti_preview(pre_mri,mask); post=nifti_preview(post_mri,mask)
        return summary_html(result),survival_figure(result),pre,post,interpretation_html(result),"Analysis completed."
    except Exception as exc:
        raise gr.Error(str(exc))

def load_example(kind):
    examples={
      "Lower-risk example":["No","A small enhancing lesion with favorable treatment response and no suspicious nodal progression.",46,"I","0","No","Yes","Yes","No","Invasive ductal carcinoma","Neoadjuvant therapy completed"],
      "Intermediate-risk example":["No","Residual enhancement and limited nodal disease are described after treatment.",58,"II","II","Yes","Yes","No","Yes","Invasive ductal carcinoma","Neoadjuvant therapy completed"],
      "Higher-risk example":["No","Persistent extensive residual disease with nodal involvement and imaging features concerning for progression.",67,"III","III","No","No","No","No","Triple-negative breast cancer","Partial systemic treatment"],
    }
    return examples[kind]

with gr.Blocks(title=APP_TITLE, css=APP_CSS, theme=gr.themes.Base(primary_hue="slate",secondary_hue="blue",neutral_hue="slate",radius_size="lg",spacing_size="md")) as demo:
    gr.HTML(HERO); gr.HTML(FLOW)
    with gr.Row(equal_height=False):
      with gr.Column(scale=5):
        with gr.Group(elem_classes=["section-card"]):
          gr.HTML(section("Eligibility","This workflow is designed for non-metastatic baseline assessment."))
          distant_metastasis=gr.Radio(["No","Yes"],value="No",label="Distant metastasis at baseline")
        with gr.Group(elem_classes=["section-card"]):
          gr.HTML(section("Imaging and radiology report","Upload NIfTI volumes. A central-slice preview is generated after analysis."))
          with gr.Row():
            pre_mri=gr.File(label="Pre-treatment MRI",file_types=[".nii",".gz"],type="filepath",elem_classes=["upload-card"])
            post_mri=gr.File(label="Post-treatment MRI",file_types=[".nii",".gz"],type="filepath",elem_classes=["upload-card"])
            mask=gr.File(label="Segmentation mask",file_types=[".nii",".gz"],type="filepath",elem_classes=["upload-card"])
          report=gr.Textbox(label="Radiology report",lines=7,placeholder="Paste the original radiology report...")
      with gr.Column(scale=4):
        with gr.Group(elem_classes=["section-card"]):
          gr.HTML(section("Clinical profile","Provide the variables used by the prognostic model."))
          with gr.Row(): age=gr.Number(label="Age",value=55,minimum=18,maximum=110); family_history=gr.Radio(["No","Yes"],value="No",label="Family history")
          with gr.Row(): pre_stage=gr.Dropdown(["0","I","II","III","IV","Unknown"],value="II",label="Pre-treatment stage"); post_stage=gr.Dropdown(["0","I","II","III","IV","Unknown"],value="II",label="Post-treatment stage")
          with gr.Row(): er=gr.Radio(["Yes","No"],value="Yes",label="ER positive"); pr=gr.Radio(["Yes","No"],value="Yes",label="PR positive"); her2=gr.Radio(["Yes","No"],value="No",label="HER2 positive")
          tumor_type=gr.Dropdown(["Invasive ductal carcinoma","Invasive lobular carcinoma","Triple-negative breast cancer","Other"],value="Invasive ductal carcinoma",label="Tumor type")
          treatment=gr.Dropdown(["Neoadjuvant therapy completed","Partial systemic treatment","Surgery first","Other"],value="Neoadjuvant therapy completed",label="Treatment")
        with gr.Group(elem_classes=["section-card"]):
          gr.HTML(section("Example profiles","Load clinical values for a fast interface test. Imaging files remain empty."))
          example=gr.Dropdown(["Lower-risk example","Intermediate-risk example","Higher-risk example"],value="Intermediate-risk example",label="Example case")
          load_btn=gr.Button("Load example",elem_classes=["secondary-action"])
        run_btn=gr.Button("Generate prognostic assessment →",variant="primary",elem_classes=["primary-action"])
        status=gr.Textbox(label="Status",interactive=False,value="Ready.")
    gr.HTML('<div style="height:12px"></div>')
    result_summary=gr.HTML()
    with gr.Row():
      with gr.Column(scale=6): curve=gr.Plot(label="Survival trajectory")
      with gr.Column(scale=4):
        with gr.Row(): pre_preview=gr.Image(label="Pre-treatment MRI",height=245); post_preview=gr.Image(label="Post-treatment MRI",height=245)
    interpretation=gr.HTML()
    gr.HTML('<div class="partner-strip"><span>Developed for translational research collaboration</span><div class="partner-names"><span>AI-BEACON</span><span>RUMC</span><span>NKI</span><span>BIG</span></div></div>')

    example_outputs=[distant_metastasis,report,age,pre_stage,post_stage,family_history,er,pr,her2,tumor_type,treatment]
    load_btn.click(load_example,inputs=example,outputs=example_outputs)
    inputs=[pre_mri,post_mri,mask,report,age,pre_stage,post_stage,family_history,er,pr,her2,tumor_type,treatment,distant_metastasis]
    run_btn.click(run_analysis,inputs=inputs,outputs=[result_summary,curve,pre_preview,post_preview,interpretation,status])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch()