#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import pandas as pd
import numpy as np

import torch
import torch.nn as nn
import torchtuples as tt

from pycox.models import CoxPH
from pycox.evaluation import EvalSurv

from sklearn.model_selection import KFold
from sklearn.preprocessing import StandardScaler
from sklearn_pandas import DataFrameMapper

import matplotlib.pyplot as plt

import gradio as gr
import dill


# In[ ]:


import pathlib
temp = pathlib.WindowsPath
pathlib.WindowsPath = pathlib.PosixPath


# In[ ]:


class MLP(nn.Module):
        def __init__(self, input_size, output_size, n_layers, layer_sizes, activation, dropout_rate):
            super(MLP, self).__init__()
            self.layers = nn.ModuleList()
            self.activation = getattr(nn, activation)()
            self.dropout = nn.Dropout(p=dropout_rate)
    
            # Input layer
            self.layers.append(nn.Linear(input_size, layer_sizes[0]))
    
            # Hidden layers
            for i in range(n_layers - 2):
                self.layers.append(nn.Linear(layer_sizes[i], layer_sizes[i+1]))
    
            # Output layer
            self.layers.append(nn.Linear(layer_sizes[-1], output_size))
    
        def forward(self, x):
            for layer in self.layers[:-1]:
                x = self.activation(layer(x))
                x = self.dropout(x)
            x = self.layers[-1](x)
            return x


# In[ ]:


with open('model.pkl', 'rb') as file:
    model = dill.load(file)


# In[ ]:


def predict(n, thy_volume, thy_v30, thy_v50, thy_vs10, thy_vs40, thy_vs65, pit_max_dose):
    
    input_data = pd.DataFrame([[n, thy_volume, thy_v30, thy_v50, thy_vs10, thy_vs40, thy_vs65, pit_max_dose]], 
                              columns=["n", "thy_volume", "thy_v30", "thy_v50", "thy_vs10", 'thy_vs40', 'thy_vs65', 'pit_max_dose'])

    cols = input_data.columns
    cols = cols.tolist()
    cols_standardize = []
    cols_leave = cols
    
    standardize = [([col], StandardScaler()) for col in cols_standardize]
    leave = [(col, None) for col in cols_leave]
    
    x_mapper = DataFrameMapper(standardize + leave)
    x_test = x_mapper.fit_transform(input_data).astype('float32')
    
    surv = model.predict_surv_df(x_test)
    
    fig, ax = plt.subplots()
    ax.plot(surv.index, 1 - surv.iloc[:, 0])  
    ax.set_xlim(0, 60)                    
    ax.set_xticks(range(0, 61, 12))
    ax.set_ylim(0, 1)  
    
    ax.set_xlabel("Months")
    ax.set_ylabel("Risk")
    ax.set_title("RIHT Risk Curve")
    
    plt.close(fig)  # Close the plot to prevent it from displaying now

    specific_points = [12, 36, 60]
    selected_data = surv.loc[specific_points]  
    selected_df = pd.DataFrame({
        'Years': ['1-year', '3-year', '5-year'],
        'Risk': 1 - selected_data.iloc[:, 0].values 
    })

    selected_df['Risk'] = selected_df['Risk'].apply(lambda x: f"{x:.2%}")
    
    return fig, selected_df


# In[ ]:


iface = gr.Interface(fn=predict,
                     inputs=[
                         gr.Dropdown(value = 2, choices=[0, 1, 2, 3], label="N Stage"),
                         gr.Number(value=14, label="Thyroid Volume (cm³)"),
                         gr.Number(value=55, label="Thyroid V30 (%)"),
                         gr.Number(value=2, label="Thyroid V50 (%)"),
                         gr.Number(value=0, label="Thyroid VS10 (cm³)"),
                         gr.Number(value=11, label="Thyroid VS40 (cm³)"),
                         gr.Number(value=13, label="Thyroid VS65 (cm³)"),
                         gr.Number(value=51, label="Pituitary Max Dose (Gy)")
                     ],
                     outputs=[
                         gr.Plot(label="RIHT Risk Curve"),
                         gr.Dataframe(label = 'RIHT Risk by time')
                     ],
                     title="RIHT Prediction Web App",
                     description='''<strong>This app predicts the risk of Radiation-Induced Hypothyroidism (RIHT)</strong>
                     <p>For academic purposes only</p>
                     <p>Please enter the parameters and hit submit</p>
                     ''')


# In[ ]:


iface.launch()


# In[ ]: