import pandas as pd
import numpy as np
import gradio as gr
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score, classification_report
import matplotlib.pyplot as plt

# Loading the dataset
df = pd.read_csv('assignment-2-k2461469.csv')

# Splitting the data into features and target variable
X = df[["dirty", "wait", "lastyear", "usa"]]
y = df["good"]

# Splitting the dataset into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Creating and fitting the logistic regression model
logreg_model = LogisticRegression()
logreg_model.fit(X_train, y_train)

# Creating and fitting the neural network model
nn_model = MLPClassifier(random_state=42)
nn_model.fit(X_train, y_train)

# Function to make predictions and display them on a graph (Logistic Regression)
def predict_and_plot_logreg(dirty, wait, lastyear, usa):
    input_data = np.array([[dirty, wait, lastyear, usa]])
    predicted_value = logreg_model.predict(input_data)[0]

    y_pred = logreg_model.predict(X_test)

    plt.figure(figsize=(12, 6))

    # Dirty vs Good plot
    plt.subplot(1, 2, 1)
    plt.scatter(X_test.iloc[:, 0], y_test, color='blue', alpha=0.6, label='Actual Good Values')  # X_test.iloc[:, 0] corresponds to 'dirty'
    plt.scatter(dirty, predicted_value, color='red', label=f'Input: Dirty={dirty}, Predicted Good={predicted_value:.2f}', zorder=5)
    plt.title('Dirty vs Good')
    plt.xlabel('Dirty')
    plt.ylabel('Good')
    plt.legend()
    plt.grid(True)

    # Wait vs Good plot
    plt.subplot(1, 2, 2)
    plt.scatter(X_test.iloc[:, 1], y_test, color='blue', alpha=0.6, label='Actual Good Values')
    plt.scatter(wait, predicted_value, color='red', label=f'Input: Wait={wait}, Predicted Good={predicted_value:.2f}', zorder=5)
    plt.title('Wait vs Good')
    plt.xlabel('Wait')
    plt.ylabel('Good')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.savefig('output_plot_logreg.png')
    plt.close()

    return predicted_value, 'output_plot_logreg.png'

# Function to make predictions and display them on a graph (Neural Network)
def predict_and_plot_nn(dirty, wait, lastyear, usa):
    input_data = np.array([[dirty, wait, lastyear, usa]])
    predicted_value = nn_model.predict(input_data)[0]

    y_pred = nn_model.predict(X_test)

    plt.figure(figsize=(12, 6))

    # Dirty vs Good plot
    plt.subplot(1, 2, 1)
    plt.scatter(X_test.iloc[:, 0], y_test, color='blue', alpha=0.6, label='Actual Good Values')  # X_test.iloc[:, 0] corresponds to 'dirty'
    plt.scatter(dirty, predicted_value, color='red', label=f'Input: Dirty={dirty}, Predicted Good={predicted_value:.2f}', zorder=5)
    plt.title('Dirty vs Good (NN)')
    plt.xlabel('Dirty')
    plt.ylabel('Good')
    plt.legend()
    plt.grid(True)

    # Wait vs Good plot
    plt.subplot(1, 2, 2)
    plt.scatter(X_test.iloc[:, 1], y_test, color='blue', alpha=0.6, label='Actual Good Values')
    plt.scatter(wait, predicted_value, color='red', label=f'Input: Wait={wait}, Predicted Good={predicted_value:.2f}', zorder=5)
    plt.title('Wait vs Good (NN)')
    plt.xlabel('Wait')
    plt.ylabel('Good')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.savefig('output_plot_nn.png')
    plt.close()

    return predicted_value, 'output_plot_nn.png'

# Creating Gradio UI with two tabs for Logistic Regression and Neural Network
with gr.Blocks() as demo:
    gr.Markdown("# Prediction Models")

    with gr.Tab("Logistic Regression"):
        gr.Markdown("### Logistic Regression Prediction")

        with gr.Row():
            dirty_slider = gr.Slider(minimum=0, maximum=6, step=0.01, label="Dirty")
            wait_slider = gr.Slider(minimum=0, maximum=5.3, step=0.01, label="Wait")
            lastyear_slider = gr.Slider(minimum=0, maximum=70, step=0.01, label="Last Year")
            usa_slider = gr.Slider(minimum=0, maximum=1, step=0.01, label="USA")

        predict_button_logreg = gr.Button("Predict")

        predicted_value_output_logreg = gr.Textbox(label="Predicted Value (Logistic Regression)")
        plot_output_logreg = gr.Image(label="Actual vs Predicted Graph (Logistic Regression)")

        predict_button_logreg.click(
            fn=predict_and_plot_logreg,
            inputs=[dirty_slider, wait_slider, lastyear_slider, usa_slider],
            outputs=[predicted_value_output_logreg, plot_output_logreg]
        )

    with gr.Tab("Neural Network"):
        gr.Markdown("### Neural Network Prediction")

        with gr.Row():
            dirty_slider_nn = gr.Slider(minimum=0, maximum=6, step=0.01, label="Dirty")
            wait_slider_nn = gr.Slider(minimum=0, maximum=5.3, step=0.01, label="Wait")
            lastyear_slider_nn = gr.Slider(minimum=0, maximum=70, step=0.01, label="Last Year")
            usa_slider_nn = gr.Slider(minimum=0, maximum=1, step=0.01, label="USA")

        predict_button_nn = gr.Button("Predict")

        predicted_value_output_nn = gr.Textbox(label="Predicted Value (Neural Network)")
        plot_output_nn = gr.Image(label="Actual vs Predicted Graph (Neural Network)")

        predict_button_nn.click(
            fn=predict_and_plot_nn,
            inputs=[dirty_slider_nn, wait_slider_nn, lastyear_slider_nn, usa_slider_nn],
            outputs=[predicted_value_output_nn, plot_output_nn]
        )

demo.launch()