import os
HF_TOKEN = os.getenv("HF_TOKEN")

import numpy as np
import pandas as pd
import sklearn
import sklearn.metrics
from math import sqrt
from scipy import stats as st
from matplotlib import pyplot as plt

from sklearn.utils import resample
from sklearn.calibration import CalibratedClassifierCV

import shap
import gradio as gr
import random
import re
import textwrap
from datasets import load_dataset


#Read data.

x1 = load_dataset("mertkarabacak/DMVO-mRS", data_files="gradio_data_x.csv", use_auth_token = HF_TOKEN)
x1 = pd.DataFrame(x1['train'])
x1 = x1.iloc[:, 1:]

print(x1.columns)
print(x1.shape)

#Define feature names.
f1_names = ['Age', 'Current or Former Smoker', 'Diabetes', 'History of Malignancy', 'Prior DVT or PE', 'Admission BMI', 'Admission NIHSS Score', 'Premorbid mRS Score', 'Occlusion Laterality', 'Hyperdense MCA', 'MT - Type of Thrombectomy', 'MT - mTICI Score', 'Hypertension', 'Antiplatelet Use', 'Admission Hemoglobin', 'Intravenous Thrombolysis']

#Prepare training data for the outcome 1.
y1 = x1.pop('OUTCOME')

#Training model.
from tabpfn import TabPFNClassifier
tabpfn = TabPFNClassifier(device='cuda', N_ensemble_configurations=8)

y1_calib_model = CalibratedClassifierCV(tabpfn, method='isotonic', cv=5)
y1_calib_model = y1_calib_model.fit(x1, y1)

y1_explainer = shap.Explainer(y1_calib_model.predict, x1)


output_y1 = (
    """
        <br/>
        <center><h3 style='font-size: 24px;'>The probability of 90-day mRS score of 3-6:</center>
        <br/>
        <center><h1 style='font-size: 36px;'>{}%</h1></center>
    """
)


#Define predict for y1.
def y1_predict(*args):
    df1 = pd.DataFrame([args], columns=x1.columns)
    pos_pred = y1_calib_model.predict_proba(df1)
    prob = pos_pred[0][1]
    prob_percent = round(prob * 100)  # Round to nearest whole number
    output = output_y1.format(prob_percent)
    return output


#Define function for wrapping feature labels.
def wrap_labels(ax, width, break_long_words=False):
    labels = []
    for label in ax.get_yticklabels():
        text = label.get_text()
        labels.append(textwrap.fill(text, width=width, break_long_words=break_long_words))
    ax.set_yticklabels(labels, rotation=0)
    

#Define interpret for y1.
def y1_interpret(*args):
    df1 = pd.DataFrame([args], columns=x1.columns)
    shap_values1 = y1_explainer(df1).values
    shap_values1 = np.abs(shap_values1)
    shap.bar_plot(shap_values1[0], max_display = 16, show = False, feature_names = f1_names)
    fig = plt.gcf()
    ax = plt.gca()
    wrap_labels(ax, 50)
    ax.figure
    plt.tight_layout()
    fig.set_figheight(9)
    fig.set_figwidth(9)
    plt.xlabel("SHAP value (impact on model output)", fontsize =12, fontweight = 'heavy', labelpad = 8)
    plt.tick_params(axis="y",direction="out", labelsize = 12)
    plt.tick_params(axis="x",direction="out", labelsize = 12)
    return fig


with gr.Blocks(title = "DMVO-mRS") as demo:
        
    gr.Markdown(
        """
    <br/>
    <center><h2>NOT FOR CLINICAL USE</h2><center>    
    <br/>    
    <center><h1>DMVO 90-Day mRS Score</h1></center>
    <center><h2>Prediction Tool</h2></center>
    <br/>
    <center><h3>This web application should not be used to guide any clinical decisions.</h3><center>
    <br/>
    <center><i>The publication describing the details of this prediction tool will be posted here upon the acceptance of publication.</i><center>
        """
    )

    gr.Markdown(
        """
        <center><h3>Model Performance</h3></center>
        <table style="margin-left: auto; margin-right: auto;">
        <tr>
            <th>Precision</th>
            <td>0.711 (0.634 - 0.765)</td>
        </tr>
        <tr>
            <th>Recall</th>
            <td>0.628 (0.553 - 0.694)</td>
        </tr>
        <tr>
            <th>F1 Score</th>
            <td>0.656 (0.585 - 0.708)</td>
        </tr>
        <tr>
            <th>Accuracy</th>
            <td>0.724 (0.696 - 0.752)</td>
        </tr>
        <tr>
            <th>Matthew's Correlation Coefficient</th>
            <td>0.450 (0.390 - 0.503)</td>
        </tr>
        <tr>
            <th>AUROC</th>
            <td>0.815 (0.790 - 0.841)</td>
        </tr>
        <tr>
            <th>AUPRC</th>
            <td>0.808 (0.781 - 0.832)</td>
        </tr>
        <tr>
            <th>Brier Score</th>
            <td>0.190 (0.177 - 0.202)</td>             
        </tr>
        </table>
        """
    )
   

    with gr.Row():

        with gr.Column():

            Age = gr.Slider(label="Age", minimum = 18, maximum = 99, step = 1, value = 55)

            Current_or_Former_Smoker = gr.Dropdown(label = "Current or Former Smoker", choices = ['No', 'Yes'], type = 'index', value = 'No')
            
            Diabetes = gr.Dropdown(label = "Diabetes", choices = ['No', 'Yes'], type = 'index', value = 'No')

            Hypertension = gr.Dropdown(label = "Hypertension", choices = ['No', 'Yes'], type = 'index', value = 'No')

            History_of_Malignancy = gr.Dropdown(label = "History of Malignancy", choices = ['No', 'Yes'], type = 'index', value = 'No')
            
            DVT_or_PE = gr.Dropdown(label = "Prior DVT or PE", choices = ['No', 'Yes'], type = 'index', value = 'No')

            Antiplatelet_Use = gr.Dropdown(label = "Antiplatelet Use", choices = ['No', 'Yes'], type = 'index', value = 'No')
            
            Admission_BMI = gr.Slider(label="Admission BMI", minimum = 15, maximum = 60, step = 0.1, value = 25)

            Admission_Hemoglobin = gr.Slider(label="Admission Hemoglobin", minimum = 5, maximum = 25, step = 0.1, value = 15)

            Admission_NIHSS = gr.Slider(label="Admission NIHSS Score", minimum = 0, maximum = 42, step = 1, value = 1)

            Premorbid_mRS = gr.Slider(label="Premorbid mRS Score", minimum = 0, maximum = 5, step = 1, value = 0)
                                      
            Occlusion_Laterality = gr.Dropdown(label = "Occlusion Laterality", choices = ['Left', 'Right'], type = 'index', value = 'Left')

            Hyperdense_MCA = gr.Dropdown(label = "Hyperdense MCA", choices = ['No', 'Yes'], type = 'index', value = 'No')

            IVTPA = gr.Dropdown(label = "Intravenous Thrombolysis", choices = ['No', 'Yes'], type = 'index', value = 'No')
            
            Type_of_Thrombectomy = gr.Dropdown(label = "MT - Type of Thrombectomy", choices = ['MT not attempted', 'Direct aspiration', 'Stent retriever', 'Combined'], type = 'index', value = 'Stent retriever')
            
            mTICI = gr.Dropdown(label = 'MT - mTICI Score', choices = ['MT not attempted' '0', '1', '2a', '2b', '2c', '3'], type = 'index', value = '3')
            
        with gr.Column():
            
            with gr.Box():
                
                with gr.Row():
                    y1_predict_btn = gr.Button(value="Predict")
                
                gr.Markdown(
                    """
                    <br/>
                    """
                    )
                
                label1 = gr.Markdown()
                
                gr.Markdown(
                    """
                    <br/>
                    """
                    )
                
                with gr.Row():
                    y1_interpret_btn = gr.Button(value="Explain")
                
                gr.Markdown(
                    """
                    <br/>
                    """
                    )
                
                plot1 = gr.Plot()
                
                gr.Markdown(
                    """
                    <br/>
                    """
                    )
                
                y1_predict_btn.click(
                    y1_predict,
                    inputs = [Age, Current_or_Former_Smoker, Diabetes, History_of_Malignancy, DVT_or_PE, Admission_BMI, Admission_NIHSS, Premorbid_mRS, Occlusion_Laterality, Hyperdense_MCA, Type_of_Thrombectomy, mTICI, Hypertension, Antiplatelet_Use, Admission_Hemoglobin, IVTPA],
                    outputs = [label1]
                )               

                y1_interpret_btn.click(
                    y1_interpret,
                    inputs = [Age, Current_or_Former_Smoker, Diabetes, History_of_Malignancy, DVT_or_PE, Admission_BMI, Admission_NIHSS, Premorbid_mRS, Occlusion_Laterality, Hyperdense_MCA, Type_of_Thrombectomy, mTICI, Hypertension, Antiplatelet_Use, Admission_Hemoglobin, IVTPA],
                    outputs = [plot1],
                )         
                
    gr.Markdown(
                """    
                <center><h2>Disclaimer</h2>
                <center> 
                This predictive tool, available on this webpage, is designed to provide general health information only and is not a substitute for professional medical advice, diagnosis, or treatment. It is strongly recommended that users consult with their own healthcare provider for any health-related concerns or issues. The authors make no warranties or representations, express or implied, regarding the accuracy, timeliness, relevance, or utility of the information contained in this tool. The health information in the prediction tool is subject to change and can be affected by various confounders, therefore it may be outdated, incomplete, or incorrect. No doctor-patient relationship is created by using this prediction tool and the authors have not validated its content. The authors do not record any specific user information or initiate contact with users. Before making any healthcare decisions or taking or refraining from any action based on the information in this prediction tool, it is advisable to seek professional advice from a healthcare provider. By using the prediction tool, users acknowledge and agree that neither the authors nor any other party will be liable for any decisions made, actions taken or not taken as a result of the information provided herein.
                <br/>
                <h4>By using this tool, you accept all of the above terms.<h4/>
                </center>
                """
    )                
                
demo.launch()