import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import gradio as gr

model = nn.Sequential(
            nn.Linear(11, 20),
            nn.ReLU(),
            nn.Linear(20, 5, bias=True))

PATH = "wine_model.pth"

model.load_state_dict(torch.load(PATH, weights_only=False))

def forward(model, input):
  preds = model(input)
  predicted_class = torch.argmax(preds, dim=-1) + 4
  return predicted_class

def process_data(input_dataframe):
    # Perform operations on the input_dataframe
    if isinstance(input_dataframe, pd.DataFrame):
      wineq_np = input_dataframe.to_numpy(dtype=np.float32)
      wineq_t = torch.from_numpy(wineq_np)
      return forward(model, wineq_t)
    return "Invalid input type"

columns = ['fixed acidity',
 'volatile acidity',
 'citric acid',
 'residual sugar',
 'chlorides',
 'free sulfur dioxide',
 'total sulfur dioxide',
 'density',
 'pH',
 'sulphates',
 'alcohol']


with gr.Blocks() as demo:
    gr.Markdown("Enter your wine data below:")
    input_df = gr.Dataframe(
        row_count=(1, "dynamic"),  # Allows adding/removing rows
        col_count=(11, "dynamic"),  # Allows adding/removing columns
        headers=columns,
        label="Input Data",
        interactive=True,
        type="pandas" # Specify the desired input type for your function
    )

    submit_button = gr.Button("Process Data")
    output_text = gr.Textbox(label="Processed Output")

    submit_button.click(
        fn=process_data, 
        inputs=input_df, 
        outputs=output_text
    )

    submit_button.click(fn=process_data, inputs=input_df, outputs=output_text)

demo.launch()