# Import necessary libraries
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.linear_model import ElasticNet
from sklearn.model_selection import train_test_split
import gradio as gr

# Define the safety function
def safety(freedom):
    return 1 - 0.7374 * freedom**2

# Load and process data for feature importance
def load_data():
    # Replace this path with the appropriate dataset URL or local file path for Hugging Face deployment
    data = pd.read_csv("data_ml.csv")  # Ensure the dataset is uploaded to the Hugging Face repo
    X = data.drop(columns=["Freedom"], errors="ignore")
    y = data["Freedom"] if "Freedom" in data.columns else None
    return X, y

# Calculate top features affecting safety
def get_top_features(X, y):
    if y is None:
        return pd.DataFrame({"Features": [], "Importance": []})
    
    # Split data and train ElasticNet model
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)
    model = ElasticNet(alpha=1.0, l1_ratio=0.5, random_state=0).fit(X_train, y_train)
    
    # Get feature importance
    feature_importance = pd.DataFrame({
        "Features": X.columns,
        "Importance": np.abs(model.coef_)
    }).nlargest(7, "Importance").reset_index(drop=True)
    return feature_importance

# Generate PPF curve data
freedom = np.arange(0, 1.01, 0.01)
safety_values = safety(freedom)

# Initialize data and top features
X, y = load_data()
top_features = get_top_features(X, y)

# Function to update outputs when slider changes
def update_prediction(freedom_value):
    # Calculate safety
    safety_score = safety(freedom_value)
    
    # Create PPF plot
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.plot(freedom, safety_values, label='Safety = 1 - 0.7374 * Freedom²', color='blue', linestyle='--')
    ax.scatter(freedom_value, safety_score, color='red', zorder=5)
    ax.text(
        freedom_value,
        safety_score,
        f"Selected Point ({freedom_value:.2f}, {safety_score:.2f})",
        fontsize=9,
        verticalalignment='bottom'
    )
    ax.set_title("Safety vs Freedom Relationship")
    ax.set_xlabel("Freedom")
    ax.set_ylabel("Safety")
    ax.legend()
    ax.grid(False)
    
    return (
        f"Predicted Safety: {safety_score:.4f}",
        fig,
        top_features
    )

# Create Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# Country Safety Predictor")
    
    with gr.Row():
        # Freedom slider
        freedom_slider = gr.Slider(
            label="Freedom Score (0 = Least Free, 1 = Most Free)",
            minimum=0,
            maximum=1,
            value=0.5,
            step=0.01
        )
        
    with gr.Row():
        # Text box for safety prediction
        safety_text = gr.Textbox(label="Predicted Safety:", value="")
        
    with gr.Row():
        # Plot output for safety vs freedom
        plot_output = gr.Plot(label="Safety vs Freedom Relationship")
        
    with gr.Row():
        # Dataframe display for top features
        features_output = gr.Dataframe(label="Top Features Affecting Safety")
    
    # Set up the interface
    freedom_slider.change(
        fn=update_prediction,
        inputs=freedom_slider,
        outputs=[safety_text, plot_output, features_output]
    )

# Run the app
if __name__ == "__main__":
    app.queue().launch(debug=False, share=True)  # Enable sharing for Hugging Face deployment