import traceback

import datasets
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score


def process_data(vendor, model):
    data = datasets.load_dataset('anakib1/mango-ria', '13.08.2024')['train']
    # Handle cases where vendor or model is None or empty string
    if vendor:
        vendor = vendor.strip().lower()
    else:
        vendor = ''
    if model:
        model = model.strip().lower()
    else:
        model = ''

    rows = data.filter(lambda x: vendor in x['Title'].lower() and model in x['Title'].lower())

    dots = []
    for row in rows:
        # row[2] is the 'Title' field
        try:
            price = float(row['Price'])
            mileage = float(row['Mileage'].split()[0])
            dots.append((price, mileage))
        except:
            print(f"Could not parse row {row}. Ex = {traceback.format_exc()}")

    if not dots:
        return "No data found for the specified vendor and model.", None, None

    price, mileage = list(zip(*dots))

    # First plot: Histogram of prices
    fig1, ax1 = plt.subplots()
    ax1.hist(price)
    ax1.set_title('Histogram of Prices')
    ax1.set_xlabel('Price')
    ax1.set_ylabel('Frequency')

    # Second plot: Scatter plot with regression line
    fig2, ax2 = plt.subplots()
    model_lr = LinearRegression()
    model_lr.fit(np.array(mileage).reshape(-1, 1), price)
    y_hat = model_lr.predict(np.array(mileage).reshape(-1, 1))
    ax2.scatter(mileage, price)
    ax2.plot(mileage, y_hat, color='r',
             label='y = {:.2f} * x + {:.2f}. R2 = {:.2f}'.format(model_lr.coef_[0], model_lr.intercept_,
                                                                 r2_score(y_true=price, y_pred=y_hat)))
    ax2.legend()
    ax2.set_xlabel('Mileage')
    ax2.set_ylabel('Price')
    ax2.set_title('Price vs Mileage with Regression Line')

    # Return the figures
    return None, fig1, fig2


with gr.Blocks() as demo:
    gr.Markdown("# Car Data Analysis")
    with gr.Row():
        vendor_input = gr.Textbox(lines=1, label="Vendor", placeholder="Enter vendor, e.g., infiniti")
        model_input = gr.Textbox(lines=1, label="Model", placeholder="Enter model, e.g., q50")
    submit_btn = gr.Button("Submit")
    message_output = gr.Textbox(label="Message", interactive=False)
    plot_output1 = gr.Plot(label="Histogram of Prices")
    plot_output2 = gr.Plot(label="Price vs Mileage with Regression Line")

    submit_btn.click(process_data, inputs=[vendor_input, model_input],
                     outputs=[message_output, plot_output1, plot_output2])

demo.launch()