import gradio as gr
import pandas as pd
import plotly.graph_objects as go

data = {
    'years_after_diagnosis': [5, 8, 7, 7, 8, 7, 5, 6, 8, 5],
    'age_at_diagnosis': [54, 41, 65, 85, 60, 66, 61, 55, 54, 48],
    'stage level': [6, 6, 5, 5, 10, 8, 6, 3, 8, 6],
    'status': [1, 1, 0, 1, 1, 1, 0, 0, 0, 0]
}

df = pd.DataFrame(data)

def plot_probability_chart(years_after_diagnosis, age_at_diagnosis, stage_level, chemotherapy, brachtherapy,
                           chemoradiation, radiotherapy, radiation, menopause, MENO_post, HISTOLOGY, CM_1, CM_2, CM_3):
    match = df.loc[
        (df['years_after_diagnosis'] == years_after_diagnosis) &
        (df['age_at_diagnosis'] == age_at_diagnosis) &
        (df['stage level'] == stage_level)
    ]

    if len(match) == 0:
        probability_alive = 0.5
        probability_dead = 0.5
        label = ["No matching record", "No matching record"]
    else:   
        status = match.iloc[0]['status']
        label = ["Alive", "Dead"]
        
        if status == 0:
            probability_alive = 0.2
            probability_dead = 0.8
        else:
            probability_alive = 0.7
            probability_dead = 0.3

    fig = go.Figure(
        go.Bar(
            x=label,
            y=[probability_alive, probability_dead]
        )
    )
    fig.update_layout(
        title="Survival Probability Chart",
        xaxis_title="Status",
        yaxis_title="Probability",
    )
    
    return f"Patient Status: {label[status]}, Probability Score: {probability_alive if status == 1 else probability_dead}", fig

inputs = [
    gr.inputs.Slider(minimum=0, maximum=10, step=1, default=5, label='years_after_diagnosis'),
    gr.inputs.Slider(minimum=0, maximum=100, step=1, default=54, label='age_at_diagnosis'),
    gr.inputs.Slider(minimum=0, maximum=10, step=1, default=6, label='stage level'),
    gr.inputs.Checkbox(label='chemotherapy'),
    gr.inputs.Checkbox(label='brachtherapy'),
    gr.inputs.Checkbox(label='chemoradiation'),
    gr.inputs.Checkbox(label='radiotherapy'),
    gr.inputs.Checkbox(label='radiation'),
    gr.inputs.Slider(minimum=0, maximum=30, step=1, default=0, label='menopause'),
    gr.inputs.Slider(minimum=0, maximum=30, step=1, default=0, label='MENO_post'),
    gr.inputs.Slider(minimum=0, maximum=2, step=1, default=2, label='HISTOLOGY'),
    gr.inputs.Checkbox(label='CM_1'),
    gr.inputs.Checkbox(label='CM_2'),
    gr.inputs.Checkbox(label='CM_3')
]

title = "Cervical Cancer Survival Predictor"
subtitle = "Ojie, Deborah Voke (PG/2021/274546)"
description = f"<small><center>{subtitle}</center></small>"

interface = gr.Interface(fn=plot_probability_chart, inputs=inputs, outputs=["markdown", "plot"], 
                         title=title, description=description, examples=None)
interface.launch()