import torch
import gradio as gr
import torch.nn.functional as F
import pandas as pd

df = pd.read_csv("Cancer_Data.csv")
df = df.dropna(axis='columns')
df = df.drop(['id'],axis=1)

headers = ["diagnosis","radius_mean","texture_mean","perimeter_mean","area_mean","smoothness_mean","compactness_mean","concavity_mean","concave points_mean","symmetry_mean","fractal_dimension_mean","radius_se","texture_se","perimeter_se","area_se","smoothness_se","compactness_se","concavity_se","concave points_se","symmetry_se","fractal_dimension_se","radius_worst","texture_worst","perimeter_worst","area_worst","smoothness_worst","compactness_worst","concavity_worst","concave points_worst","symmetry_worst","fractal_dimension_worst"]
inputs = [gr.Dataframe(headers = headers, row_count = (2, "dynamic"), col_count=(31,"dynamic"), label="Input Data", interactive=False)]
outputs = [gr.Dataframe(row_count = (2, "dynamic"), col_count=(1, "fixed"), label="Predictions", headers=["results"])]

def classify_cell(df_input):
    # Dropping diagnosis
    cells = df_input.drop(['diagnosis'],axis=1).values
    
    # Classes
    cancer_classes = ['Benign','Malignant']
    # Loading model
    cell_model = torch.jit.load('cancer_classifier.ptl')
    # List I will pass into a dataframe as the return object
    cell_results = []

    # Converting to tensor and casting it as float32
    cell = torch.tensor(cells)
    cell = cell.to(torch.float)

    # Running through model and applying softmax to output
    cell_pred = cell_model(cell)
    cell_pred = F.softmax(cell_pred,dim=1)
    # Looping through model output to format string for each cell
    for i in range(len(cell_pred)):
        cell_prob = round(cell_pred[i][cell_pred[i].argmax()].item()*100,2)
        output_string = f"{cancer_classes[cell_pred[i].argmax()]} Cancer Cell : {cell_prob}% confident"
        # Appending formatted string for a cell to cell results list
        cell_results.append(output_string)
    return pd.DataFrame(cell_results,columns=["results"])

demo = gr.Interface(classify_cell,
                        inputs = inputs,
                        outputs = outputs,
                        title = "Classify Cancer Cell", 
                        description="Classifies Cancer Cell into <b>Malignant</b> or <b>Benign</b> based on its features. Click on the example to classify 5 cells! </br> First Column shows what the cell should be classified as",
                        examples=[df.iloc[17:22]],
                        cache_examples=False
                   )
                    
demo.launch(inline=False)