import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
import gradio as gr

# Define the mapping between user-friendly names and NASA variable names
data_mapping = {
    "Temperature": {
        "Earth Skin Temperature": "TS",
        "Temperature at 2 Meters": "T2M",
        "Wet Bulb Temperature at 2 Meters": "T2MWET"
    },
    "Moisture & Precipitation": {
        "Precipitation Average": "PRECTOTCORR",
        "Profile Soil Moisture (surface to bedrock)": "GWETPROF",
        "Root Zone Soil Wetness (surface to 100 cm below)": "GWETROOT",
        "Surface Soil Wetness (surface to 5 cm below)": "GWETTOP"
    },
    "Air & Pressure": {
        "Specific Humidity at 2 Meters": "QV2M",
        "Surface Pressure": "PS"
    },
    "Wind": {
        "Wind Direction at 10 Meters": "WD10M",
        "Wind Direction at 2 Meters": "WD2M",
        "Wind Speed at 10 Meters": "WS10M",
        "Wind Speed at 2 Meters": "WS2M"
    }
}

def load_and_prepare_data():
    data = pd.read_csv('data/preprocessed_data.csv')
    nasa_features = [var for category in data_mapping.values() for var in category.values()]
    target = 'Close_^FTSE'

    X = data[nasa_features]
    y = data[target]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, shuffle=False)

    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    return X_train_scaled, X_test_scaled, y_train, y_test, scaler, nasa_features, data

def train_model(X_train_scaled, y_train):
    model = LinearRegression()
    model.fit(X_train_scaled, y_train)
    return model

def create_prediction_interface(model, scaler, features_mapping, data):
    def predict_func(*args):
        input_values = pd.DataFrame([args], columns=nasa_features)
        input_values_scaled = scaler.transform(input_values)
        prediction = model.predict(input_values_scaled)
        return f"${round(prediction[0], 2)} USD"

    with gr.Blocks(theme=gr.themes.Default()) as demo:
        gr.Markdown("### Stock Price Prediction")
        gr.Markdown("Adjust NASA POWER DAS variables to see predicted stock prices.")
        
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("#### Adjust Inputs")
                inputs = []
                for category, variables in features_mapping.items():
                    with gr.Accordion(category, open=False):
                        for csv_name, das_var in variables.items():
                            min_val = data[das_var].min()
                            max_val = data[das_var].max()
                            slider = gr.Slider(
                                label=csv_name,
                                minimum=min_val,
                                maximum=max_val,
                                value=(min_val + max_val) / 2,
                                step=(max_val - min_val) / 100
                            )
                            inputs.append(slider)
                predict_btn = gr.Button("Predict", variant="primary")
            
            with gr.Column(scale=1):
                gr.Markdown("#### Predicted Result", elem_id="result-header")
                output = gr.Textbox(label="Predicted Close_^FTSE", interactive=False, value="$0.00 USD")
                predict_btn.click(fn=predict_func, inputs=inputs, outputs=output)
    
    return demo

if __name__ == "__main__":
    X_train_scaled, X_test_scaled, y_train, y_test, scaler, nasa_features, data = load_and_prepare_data()
    model = train_model(X_train_scaled, y_train)
    demo = create_prediction_interface(model, scaler, data_mapping, data)
    demo.launch()