import gradio as gr
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
import joblib
import numpy as np


# Define the model architecture
model = joblib.load('Keras.pkl')


manual_encoding = {
    "Gender": {"Male": 1, "Female": 0, "-99": -99},
    "Race (Reported)": {"White": 0, "Other Mixed Race": 1, "African-American": 2, "Black": 3, "Hispanic": 4, "Japanese": 5, "Han Chinese": 6, "other": 7, "Asian": 8, "Korean": 9, "Caucasian": 10, "Black or African American": 11, "Other": 12, "Intermediate": 13, "Malay": 14, "Chinese": 15, "Indian": 16, "Black Caribbean": 17, "Other (Hungarian)": 18, "Other (Black British)": 19, "Black African": 20, "Black other": 21},
    "Age": {"90+": 6, "80 - 89": 5, "70 - 79": 4, "60 - 69": 3, "50 - 59": 2, "40 - 49": 1, "30 - 39": 0, "20 - 29": 7, "10 - 19": 8},
    "Diabetes": {1: 1, 0: 0},
    "Simvastatin (Zocor)": {1: 1, 0: 0},
    "Amiodarone (Cordarone)": {1: 1, 0: 0},
    "Target INR": {3.5: 5, 3.25: 5, 3.0: 5, 2.8: 5, 2.7: 5, 2.6: 5, 2.5: 5, 2.3: 5, 2.2: 5, 2.0: 5, 1.75: 5, 1.3: 5},
    "Cyp2C9 genotypes": {"*1/*1": 0, "*1/*3": 1, "*1/*2": 2, "*2/*2": 3, "*2/*3": 4, "*3/*3": 5, "*1/*5": 6, "*1/*13": 7, "*1/*14": 8, "*1/*11": 9, "*1/*6": 10},
    "VKORC1 genotype: -1639 G>A (3673); chr16:31015190; rs9923231; C/T": {"A/A": 0, "A/G": 1, "G/G": 2}
}

def apply_encoding(input_data):
    encoded_data = {}
    for key, value in input_data.items():
        if key in manual_encoding:
            encoded_data[key] = manual_encoding[key][value]
        else:
            encoded_data[key] = value
    return encoded_data

def predict_warfarin_dose(gender, race, age_group, height, weight, diabetes, simvastatin, amiodarone, target_inr, inr_on_reported_dose, cyp2c9, vkorc1):
    # Encode input data
    encoded_data = apply_encoding({
        "Gender": gender,
        "Race (Reported)": race,
        "Age": age_group,
        "Height": height,
        "Weight": weight,
        "Diabetes": diabetes,
        "Simvastatin (Zocor)": simvastatin,
        "Amiodarone (Cordarone)": amiodarone,
        "Target INR": target_inr,
        "INR on Reported Therapeutic Dose of Warfarin": inr_on_reported_dose,
        "Cyp2C9 genotypes": cyp2c9,
        "VKORC1 genotype: -1639 G>A (3673); chr16:31015190; rs9923231; C/T": vkorc1
    })

    # Convert encoded data to a format compatible with the model
    input_data = [encoded_data[key] for key in encoded_data]
    input_data = np.array(input_data).reshape(1, -1)

    # Make prediction using the trained model
    predicted_dose = model.predict(input_data)[0][0]

    return str(-1*predicted_dose)

iface = gr.Interface(
    fn=predict_warfarin_dose,
    inputs=[
        gr.Dropdown(["Male", "Female", "-99"], label="Gender"),
        gr.Dropdown(list(manual_encoding["Race (Reported)"].keys()), label="Race"),
        gr.Dropdown(list(manual_encoding["Age"].keys()), label="Age Group"),
        gr.Number(label="Height (cm)"),
        gr.Number(label="Weight (kg)"),
        gr.Dropdown([1, 0], label="Diabetes"),
        gr.Dropdown([1, 0], label="Simvastatin"),
        gr.Dropdown([1, 0], label="Amiodarone"),
        gr.Dropdown(list(manual_encoding["Target INR"].keys()), label="Target INR"),
        gr.Number(label="INR on Reported Therapeutic Dose of Warfarin"),
        gr.Dropdown(list(manual_encoding["Cyp2C9 genotypes"].keys()), label="CYP2C9 Genotypes"),
        gr.Dropdown(list(manual_encoding["VKORC1 genotype: -1639 G>A (3673); chr16:31015190; rs9923231; C/T"].keys()), label="VKORC1 Genotype"),
    ],
    outputs=gr.Textbox(label="Predicted Therapeutic Dose of Warfarin (mg/week)"),
    title="Warfarin Dosing Prediction",
    description="Enter patient information to predict the therapeutic dose of warfarin.",
)

iface.launch('share=True')