import gradio as gr
import numpy as np
import joblib
from sklearn.preprocessing import StandardScaler
import os

# Define model paths
model_paths = {
    'Path': {
        '3 hours': 'path_to_3H_model.pkl',
        '6 hours': 'path_to_6H_model.pkl',
        '9 hours': 'path_to_9H_model.pkl'
    },
    'Speed': {
        '3 hours': 'path_to_3H_speed_model.pkl',
        '6 hours': 'path_to_6H_speed_model.pkl',
        '9 hours': 'path_to_9H_speed_model.pkl'
    },
    'Pressure': {
        '3 hours': 'path_to_3H_pressure_model.pkl',
        '6 hours': 'path_to_6H_pressure_model.pkl',
        '9 hours': 'path_to_9H_pressure_model.pkl'
    }
}

# Define scaler paths
scaler_paths = {
    'Path': {
        '3 hours': 'path_to_3H_scaler.pkl',
        '6 hours': 'path_to_6H_scaler.pkl',
        '9 hours': 'path_to_9H_scaler.pkl'
    },
    'Speed': {
        '3 hours': 'path_to_3H_speed_scaler.pkl',
        '6 hours': 'path_to_6H_speed_scaler.pkl',
        '9 hours': 'path_to_9H_speed_scaler.pkl'
    },
    'Pressure': {
        '3 hours': 'path_to_3H_pressure_scaler.pkl',
        '6 hours': 'path_to_6H_pressure_scaler.pkl',
        '9 hours': 'path_to_9H_pressure_scaler.pkl'
    }
}

def process_input(input_data, scaler):
    input_data = np.array(input_data).reshape(-1, 7)
    processed_data = input_data[:2].reshape(1, -1)
    processed_data = scaler.transform(processed_data)
    return processed_data

def load_model_and_predict(prediction_type, time_interval, input_data):
    try:
        # Load the model and scaler based on user selection
        model = joblib.load(model_paths[prediction_type][time_interval])
        scaler = joblib.load(scaler_paths[prediction_type][time_interval])
        
        # Process input and predict
        processed_data = process_input(input_data, scaler)
        prediction = model.predict(processed_data)
        
        if prediction_type == 'Path':
            return f"Predicted Latitude: {prediction[0][0]}, Predicted Longitude: {prediction[0][1]}"
        elif prediction_type == 'Speed':
            return f"Predicted Speed: {prediction[0]}"
        elif prediction_type == 'Pressure':
            return f"Predicted Pressure: {prediction[0]}"
    except Exception as e:
        return str(e)

# Gradio interface components
with gr.Blocks() as cyclone_predictor:
    gr.Markdown("# Cyclone Prediction App")
    
    # Dropdown for Prediction Type
    prediction_type = gr.Dropdown(
        choices=['Path', 'Speed', 'Pressure'],
        label="Select Prediction Type"
    )
    
    # Dropdown for Time Interval
    time_interval = gr.Dropdown(
        choices=['3 hours', '6 hours', '9 hours'],
        label="Select Time Interval"
    )
    
    # Input fields for user data
    input_data = gr.Textbox(
        placeholder="Enter cyclone data as list of lists, e.g., [[15.54,90.64,31,2024,10,23,0], [15.71,90.29,32,2024,10,23,3]]",
        label="Input Cyclone Data (2 rows required)"
    )
    
    # Output prediction
    prediction_output = gr.Textbox(label="Prediction Output")
    
    # Predict button
    predict_button = gr.Button("Predict")
    
    # Linking function to UI elements
    predict_button.click(
        load_model_and_predict,
        inputs=[prediction_type, time_interval, input_data],
        outputs=prediction_output
    )

cyclone_predictor.launch()