# Import Gradio and other libraries
import gradio as gr
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from PIL import Image
import pandas as pd
import joblib,os

from sklearn.preprocessing import MinMaxScaler,LabelEncoder

n_tab_features = 8
num_classes = 11
n_pain_cls = 4
n_hidden = 128

project = os.getcwd()
scaler = MinMaxScaler(feature_range=(0,1))
encoder = LabelEncoder()

fun_tran= lambda x: [int(y) if y!=1 else 0 for y in np.exp( 5*np.log( [k if k>0 else 1 for k in x.flatten()] ) ) ]

# Define the model class
class DosagePredictionModel(nn.Module):
    def __init__(self, n_tab_features, n_pain_cls):
        super(DosagePredictionModel, self).__init__()
        self.fc1 = nn.Linear(n_tab_features, n_hidden)  # n_tab_features input features
        self.fc2 = nn.Linear(n_hidden, n_hidden)
        self.fc3 = nn.Linear(n_hidden, n_hidden//2)
        self.fc4 = nn.Linear(n_hidden//2, n_hidden//2)
        self.bn5 = nn.BatchNorm1d(n_hidden//2)
        self.fc6 = nn.Linear(n_hidden//2, num_classes)  # Output for Medicine classification
        self.fc7 = nn.Linear(n_hidden//2, n_pain_cls)  # Output for Pain classification
        self.fc8 = nn.Linear(n_hidden//2, 1)  # Output for regression
        

    def forward(self, x, regression_targets=None,classification_targets=None, pain_cls_tgt=None ):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = self.bn5(x)

        regression_output = self.fc8(x)  # Regression output
        regression_y = regression_output
        
        # Medicine Classification Layer
        output_classification = self.fc6(x)
        output_classification = torch.softmax(output_classification,dim=1)
        pred_cls_y = torch.argmax(output_classification, dim = 1)
        
        # Pain Classification Layer
        pain_output_cls = self.fc7(x)
        pain_output_cls = torch.softmax(pain_output_cls,dim=1)
        pred_pain_cls_y = torch.argmax(pain_output_cls, dim = 1)
        
        if (classification_targets is None) or (regression_targets is None) or (pain_cls_tgt is None):
            #print('Inference Mode')
            loss, regression_loss, med_cls_loss, pain_cls_loss = None, None, None, None
        else:
            loss_fn = nn.SmoothL1Loss()  # mean square error
            regression_loss = loss_fn(regression_output, regression_targets.view(-1,1))
            #Classification
            med_cls_loss = nn.CrossEntropyLoss()(output_classification, classification_targets )
            #pain_class
            pain_cls_loss = nn.CrossEntropyLoss()(pain_output_cls, pain_cls_tgt )
            
            weight_dosage, weight_med_cls = 0.2, 0.4
            loss = weight_dosage * regression_loss + weight_med_cls * med_cls_loss + (1-(weight_dosage+weight_med_cls))*pain_cls_loss
        return loss,regression_loss,med_cls_loss,pain_cls_loss, regression_y, pred_cls_y, pred_pain_cls_y


# Create an instance of the model
model = DosagePredictionModel(n_tab_features, n_pain_cls)

# Load the model weights from a file (assuming it is saved as "model.pt")
model.load_state_dict(torch.load("med_dos_pain_1.pt"))

# Set the model to evaluation mode
model.eval()

# Define a list of possible medicines
medicines = ["Aspirin", "Lisinopril", "Metoprolol", "Hydrochlorothiazide"]
dense_features = ['Weight_in_Kgs','Heart_rate','pulse_rate','Systolic_BP','Diastolic_BP','BIS_Value','SPO2']
sparse_features = ['Pain_Position']

# Define a function to get the medicine name from the model output
def get_medicine_name(pred_medicine_class, pred_pain_cls):
    med_encoder = joblib.load(project+'/cat_encoder.joblib')
    pred_medicine_name = med_encoder.inverse_transform(pred_medicine_class)
    
    pain_level_enc= joblib.load(project+'/cat_dos_level_encoder.joblib')
    pred_dos_level_nm = pain_level_enc.inverse_transform(pred_pain_cls)
    
    # Return the corresponding medicine name
    return pred_medicine_name, pred_dos_level_nm

# Define a function to get the model input from the user input
def get_model_input(Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP, Diastolic_BP, BIS_Value, SPO2, Pain_Position):
    values = [[Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP, Diastolic_BP, BIS_Value, SPO2]]
    cat_values = [[Pain_Position]]
    all_values = [[Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP, Diastolic_BP, BIS_Value, SPO2,Pain_Position]]
    temp = pd.DataFrame(all_values, columns= dense_features + sparse_features)
#     print(temp.dtypes)
    # Normalize the dense feature values to be between 0 and 1
    scaler = joblib.load(project+'/scaler.joblib')
    temp[dense_features] = scaler.transform( temp[dense_features] )
    #print(temp)

    # Encode the categorical features 0, 1 etc.
    painpos_encoder = joblib.load(project+'/cat_pain_level_encoder.joblib')
    temp[sparse_features] = painpos_encoder.transform( cat_values[0] )

    # Create a tensor from the normalized values
    input = torch.tensor(temp.to_numpy()).float()
    #print(input)
    # Return the input tensor
    return input

# Define a function to get the prediction from the user input
def predict(Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP, Diastolic_BP, BIS_Value, SPO2, Pain_Position):
    # Get the model input from the user input
    X = get_model_input(Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP, Diastolic_BP, BIS_Value, SPO2, Pain_Position)
    #print(X.dtype)
    # Get the model output from the model input
    _, _, _, _, pred_dosage, pred_medicine_class, pred_pain_cls = model(X)
    # Get the medicine name from the model output
    medicine, dosage_level_nm = get_medicine_name(pred_medicine_class, pred_pain_cls)
    if Pain_Position =='No Pain':
        medicine, dosage_level_nm = 'NA', 'NA'
    # Return the predicted medicine name as a string
    return "The predicted medicine for you is: " + medicine[0], dosage_level_nm[0]

Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP = 75.3, 85.7, 84, 141
Diastolic_BP, BIS_Value, SPO2, Pain_Position       = 92, 80, 90, 'Lower Back'
predict(Weight_in_Kgs, Heart_rate, pulse_rate, Systolic_BP, Diastolic_BP, BIS_Value, SPO2, Pain_Position)
# Create a Gradio interface with two number inputs and a text output
interface = gr.Interface(
    fn=predict,
    inputs=[gr.inputs.Slider(minimum=40,maximum=120,step=1,label="Weight in Kgs"), 
            gr.inputs.Slider(minimum=40,maximum=105,step=1,label="Heart rate"), 
            gr.inputs.Slider(minimum=60,maximum=100,step=1,label="Pulse rate"),
            gr.inputs.Slider(minimum=115,maximum=142,step=1,label="Systolic BP"), 
            gr.inputs.Slider(minimum=70,maximum=93,step=1,label="Diastolic BP"),
            gr.inputs.Slider(minimum=30,maximum=100,step=1,label="BIS Value"),
            gr.inputs.Slider(minimum=71,maximum=100,step=1,label="SPO2"), 
            gr.inputs.Dropdown(['Lower Back', 'Shoulder ', 'Abdomine', 'Muscle', 'Neck',
       'Fybromialgia', 'Knee Joint', 'Joint pain', 'Shoulder', 'No Pain'], label="Pain Position", default="Lower Back") ],
    outputs=[gr.outputs.Textbox(label="Medicine"),gr.outputs.Textbox(label="Dosage Level")] )

# Launch the interface on Hugging Face Spaces (assuming you have an account)
interface.launch()