import gradio as gr
import numpy as np
import pickle
from tensorflow.keras.models import load_model
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image
import json
import pandas as pd

model = load_model('models/new_merged_weather_financial_q_l4q_l24q_general.keras')
scaler = MinMaxScaler()

with open('data/2024_11_7_total_return_sample.pkl', 'rb') as f:
    data = pickle.load(f)

columns_to_drop = ['Date', 'CBSA_Name', 'Id', 'iname', 'type', 'Cluster', 'region', 'division', 'state', 'msa', 'tret',
                   'iret', 'aret', 'treturn', 'tot_index', 'inc_index', 'app_index', 'count', 'emv', 'bmv', 'income',
                   'psales', 'capimp', 'ireturn', 'areturn', 'Latitude_x', 'Longitude_x', 'Latitude_y', 'Longitude_y',
                   'tmin', 'tmean', 'tmax', 'tdmean', 'ppt', 'vpdmin', 'vpdmax', 'tmin_low', 'tmax_high',
                   'num_of_positive_headlines', 'num_of_negative_headlines', 'remote_work_shift']

features = data.drop(columns=columns_to_drop)
scaler.fit(features)


def create_sequences(data_param, target_param, input_steps=12, forecast_steps=4):
    X, y = [], []
    for i in range(len(data_param) - input_steps - forecast_steps):
        X.append(data_param[i:(i + input_steps)])
        y.append(target_param[(i + input_steps):(i + input_steps + forecast_steps)].mean())
    return np.array(X), np.array(y)


def predict_and_plot(cbsa_name):
    cbsa_data = data[data['CBSA_Name'] == cbsa_name]

    print(cbsa_data['YYYYQ'].sort_values())

    cbsa_features = cbsa_data.drop(columns=columns_to_drop)
    cbsa_features = cbsa_features[features.columns]

    cbsa_target = cbsa_data['tret']

    cbsa_scaled_features = scaler.transform(cbsa_features)

    X_cbsa, y_cbsa = create_sequences(cbsa_scaled_features, cbsa_target)

    predictions = model.predict(X_cbsa)
    predictions = np.squeeze(predictions)

    shift_steps = 5
    predictions = np.roll(predictions, -shift_steps)

    future_quarters = [f"{year}-Q{quarter}" for year in range(2024, 2026) for quarter in range(1, 5)]

    num_future_steps = len(future_quarters)

    future_predictions = []
    current_input = X_cbsa[-1]

    for _ in range(num_future_steps):
        next_prediction = model.predict(current_input.reshape(1, -1, X_cbsa.shape[2]))

        future_predictions.append(next_prediction.squeeze())

        current_input = np.roll(current_input, -1, axis=0)

        current_input[-1] = next_prediction.squeeze()

    predictions = np.concatenate((predictions, np.array(future_predictions)))

    time_index = (
        cbsa_data['YYYYQ'].iloc[:len(y_cbsa)]
        .apply(lambda x: f"{str(x)[:4]}-Q{str(x)[4]}")
        .sort_values()
    )

    future_time_index = pd.Series(future_quarters)
    full_time_index = pd.concat([time_index, future_time_index]).reset_index(drop=True)

    actual_index = full_time_index[:len(y_cbsa)]
    predicted_index = full_time_index[:len(predictions)]

    predicted_json = [
        {
            "year": int(year_quarter.split("-")[0]),
            "quarter": int(year_quarter.split("-")[1][1]),
            "total_return": round(float(pred), 4)
        }
        for year_quarter, pred in zip(full_time_index[-num_future_steps:], future_predictions)
    ]

    json_output = {
        "CBSA": cbsa_name,
        "Label": "Total Return",
        "Predicted": predicted_json
    }

    plt.figure(figsize=(20, 6))
    plt.plot(actual_index, y_cbsa, label='Actual Total Return', color='green', linestyle='-')
    plt.plot(predicted_index, predictions, label='Predicted Total Return', color='red', linestyle='-')
    plt.title(f'Model Predictions vs Actual for {cbsa_name}')
    plt.xlabel('Time (YYYYQ)')
    plt.ylabel('Total Return')

    plt.xticks(rotation=90, fontsize=8)
    plt.gca().tick_params(axis='x', pad=15)

    plt.legend()
    plt.tight_layout()

    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = Image.open(buf)
    img_array = np.array(img)

    return img_array, json.dumps(json_output, indent=2)


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            cbsa_dropdown = gr.Dropdown(choices=data['CBSA_Name'].unique().tolist(), label="Select CBSA Area")
            predict_button = gr.Button("Submit")
        with gr.Column():
            gr.Markdown("""
            **Total Return**: Total return is a measure of the performance of an asset or investment over a specific period, defined as the sum of **Income Return** and **Asset Return**. 
            - **Income Return**: The net income generated by a property, calculated as rental income minus operating and capital expenditures.
            - **Asset Return**: The appreciation in the market value of a property from purchase to sale.
            **CBSA (Core-Based Statistical Area)**: Represents a geographical area defined by the Office of Management and Budget, typically used for statistical purposes in the U.S. It consists of counties and county equivalents centered around an urban center with a high degree of social and economic integration.
            """)
    with gr.Row():
        output_image = gr.Image(type="numpy", label="Actual vs Predicted Total Return (2015-2025)")

    with gr.Row():
        json_display = gr.JSON(label="Prediction JSON Output")

    predict_button.click(fn=predict_and_plot, inputs=cbsa_dropdown, outputs=[output_image, json_display])

demo.launch()