import hopsworks
import os
import joblib
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

os.environ["HOPSWORKS_API_KEY"] = "HLDlE8LYYNS4zTdX.GQvAhHVCPMDVOrXcbGK3IKrsUv2APbwSUzexJDgjhXHPRoC3aAjhOMKgkaM95UHt"

project = hopsworks.login()
fs = project.get_feature_store()
mr = project.get_model_registry()

mlp_model_meta = mr.get_model('mlp_airline_delay', version=1)
mlp_model_dir = mlp_model_meta.download()
mlp_model = joblib.load(f"{mlp_model_dir}/model.pkl")

# Read data from Hopsworks feature store
feature_group = fs.get_feature_group(name="new_flight_data", version=1)
df_lagged = feature_group.read()

# Infer min and max values for sliders
slider_ranges = {}
continuous_features = [
    'arr_flights', 'arr_diverted', 'temperature_2m', 'wind_speed_10m',
    'arr_delay_lag1', 'arr_flights_lag2', 'weather_delay', 'security_delay',
    'carrier_delay', 'arr_delay_lag1_lag2', 'arr_delay_lag1_lag3'
]

for col in continuous_features:
    slider_ranges[col] = {
        "min": round(df_lagged[col].min(), 1),
        "max": round(df_lagged[col].max(), 1)
    }

carrier_mapping = {
    "American Airlines": "carrier_aa",
    "Delta Air Lines": "carrier_dl",
    "United Airlines": "carrier_ua",
    "Southwest Airlines": "carrier_wn",
    "Alaska Airlines": "carrier_as",
    "JetBlue Airways": "carrier_b6",
    "Spirit Airlines": "carrier_nk",
    "Frontier Airlines": "carrier_f9",
    "Envoy Air": "carrier_mq",
    "SkyWest Airlines": "carrier_oo",
    "Mesa Airlines": "carrier_oh",
    "Republic Airways": "carrier_yx",
    "Horizon Air": "carrier_qx",
    "Allegiant Air": "carrier_ax",
    "US Airways": "carrier_us",
    "AirTran Airways": "carrier_fl"
}

def encode_month(month):
    month_sin = np.sin(2 * np.pi * month / 12)
    month_cos = np.cos(2 * np.pi * month / 12)
    return month_sin, month_cos

def predict_delay(carrier, arr_flights, arr_diverted, temperature, wind_speed, month,
                  arr_delay_lag1, arr_flights_lag2, weather_delay, security_delay,
                  carrier_delay, arr_delay_lag1_lag2, arr_delay_lag1_lag3):
    # Encode month for cyclical features
    month_sin, month_cos = encode_month(month)

    # Create input DataFrame
    data = {
        'arr_flights': [arr_flights],
        'arr_diverted': [arr_diverted],
        'temperature_2m': [temperature],
        'wind_speed_10m': [wind_speed],
        'month_sin': [month_sin],
        'month_cos': [month_cos],
        'arr_delay_lag1': [arr_delay_lag1],
        'arr_flights_lag2': [arr_flights_lag2],
        'weather_delay': [weather_delay],
        'security_delay': [security_delay],
        'carrier_delay': [carrier_delay],
        'arr_delay_lag1_lag2': [arr_delay_lag1_lag2],
        'arr_delay_lag1_lag3': [arr_delay_lag1_lag3]
    }
    input_df = pd.DataFrame(data)

    # One-hot encode carrier
    carrier_code = carrier_mapping.get(carrier)
    if carrier_code:
        input_df[carrier_code] = 1

    # Ensure all carrier columns exist and fill missing ones with 0
    for col in carrier_mapping.values():
        if col not in input_df:
            input_df[col] = 0

    # Reorder columns to match model's expected input
    input_df = input_df.reindex(columns=mlp_model.feature_names_in_, fill_value=0)

    # Make prediction
    prediction = mlp_model.predict(input_df)
    return round(prediction[0], 2)

import gradio as gr

def create_terminal_crowd(delay_minutes, n_passengers=200):
    # Normalize delay to determine crowd behavior
    # Assume delay range from -1 to 5 minutes
    normalized_delay = min(max((delay_minutes + 1) / 6, 0), 1)  # 0 to 1 scale
    
    # Terminal layout dimensions
    terminal_width = 100
    terminal_height = 60
    
    # Calculate crowd parameters based on delay
    crowd_density = normalized_delay * 0.8 + 0.2  # 0.2 to 1.0
    crowd_clustering = normalized_delay * 0.7 + 0.3  # 0.3 to 1.0
    
    # Generate passenger positions
    n_actual_passengers = int(n_passengers * crowd_density)
    
    # Create different passenger groups
    # 1. Seated passengers (at gates)
    # 2. Moving passengers (walking through terminal)
    # 3. Clustered passengers (at info desks, security, etc.)
    
    # Define gate areas (x, y, width, height)
    gates = [
        [10, 5, 20, 10],   # Gate A
        [40, 5, 20, 10],   # Gate B
        [70, 5, 20, 10],   # Gate C
        [10, 45, 20, 10],  # Gate D
        [40, 45, 20, 10],  # Gate E
        [70, 45, 20, 10]   # Gate F
    ]
    
    # Define info desk/security areas (x, y, radius)
    service_points = [
        [50, 30, 6],       # Main info desk
        [25, 30, 4],       # Security point 1
        [75, 30, 4]        # Security point 2
    ]
    
    # Allocate passengers to different areas based on delay
    seated_ratio = max(0.6 - normalized_delay * 0.4, 0.2)  # More delay = fewer seated
    service_ratio = normalized_delay * 0.5  # More delay = more at service points
    walking_ratio = 1 - seated_ratio - service_ratio
    
    n_seated = int(n_actual_passengers * seated_ratio)
    n_service = int(n_actual_passengers * service_ratio)
    n_walking = n_actual_passengers - n_seated - n_service
    
    # Generate positions
    passenger_x = []
    passenger_y = []
    passenger_colors = []
    
    # 1. Seated passengers (distributed across gates)
    gates_fullness = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8]  # Different fullness per gate
    for i in range(n_seated):
        # Select a gate, with higher delay making some gates more full than others
        gate_idx = min(int(i / n_seated * len(gates)), len(gates) - 1)
        gate = gates[gate_idx]
        
        # Add random position within gate area
        fullness = gates_fullness[gate_idx] * (1 + normalized_delay * 0.5)  # Adjusted for delay
        gate_crowd_factor = min(fullness, 1.0)  # Cap at 1.0
        
        # Generate position with more clustering as gate gets fuller
        if gate_crowd_factor < 0.5:
            # Spread out
            x = gate[0] + np.random.uniform(0, gate[2])
            y = gate[1] + np.random.uniform(0, gate[3])
        else:
            # Cluster more
            x = gate[0] + gate[2]/2 + np.random.normal(0, gate[2]/4)
            y = gate[1] + gate[3]/2 + np.random.normal(0, gate[3]/4)
            # Keep within bounds
            x = max(min(x, gate[0] + gate[2]), gate[0])
            y = max(min(y, gate[1] + gate[3]), gate[1])
        
        passenger_x.append(x)
        passenger_y.append(y)
        passenger_colors.append('blue')  # Seated passengers are blue
    
    # 2. Service point passengers (clustered at info desks/security)
    for i in range(n_service):
        # More delay means more concentration at the main info desk
        if normalized_delay > 0.7 and np.random.random() < 0.7:
            point_idx = 0  # Main info desk
        else:
            point_idx = min(int(i / n_service * len(service_points)), len(service_points) - 1)
        
        point = service_points[point_idx]
        
        # Cluster around the service point with degree based on delay
        cluster_tightness = 0.5 + normalized_delay * 0.5  # 0.5 to 1.0
        distance = np.random.exponential(point[2] / cluster_tightness)
        angle = np.random.uniform(0, 2 * np.pi)
        
        # Polar to cartesian conversion for clustered distribution
        x = point[0] + distance * np.cos(angle)
        y = point[1] + distance * np.sin(angle)
        
        # Keep within terminal bounds
        x = max(min(x, terminal_width), 0)
        y = max(min(y, terminal_height), 0)
        
        passenger_x.append(x)
        passenger_y.append(y)
        passenger_colors.append('red')  # Service point passengers are red (frustrated)
    
    # 3. Walking passengers (random paths in terminal)
    central_path_y = terminal_height / 2
    path_width = 20
    
    for i in range(n_walking):
        # With more delay, people tend to wander more randomly
        if normalized_delay < 0.4:
            # Mostly following the central path
            x = np.random.uniform(0, terminal_width)
            y = central_path_y + np.random.normal(0, path_width * 0.2)
        else:
            # More random movements
            x = np.random.uniform(0, terminal_width)
            y = np.random.uniform(0, terminal_height)
        
        # Keep within terminal bounds
        x = max(min(x, terminal_width), 0)
        y = max(min(y, terminal_height), 0)
        
        passenger_x.append(x)
        passenger_y.append(y)
        
        # Color walking passengers by delay frustration level
        if normalized_delay < 0.3:
            passenger_colors.append('green')  # Happy/Normal
        elif normalized_delay < 0.7:
            passenger_colors.append('orange')  # Mild frustration
        else:
            passenger_colors.append('red')  # High frustration
    
    return passenger_x, passenger_y, passenger_colors, normalized_delay

def predict_delay_interface(carrier, arr_flights, arr_diverted, temperature, wind_speed, month,
                            arr_delay_lag1, arr_flights_lag2, weather_delay, security_delay,
                            carrier_delay, arr_delay_lag1_lag2, arr_delay_lag1_lag3):
    # Get prediction
    delay_minutes = predict_delay(
        carrier, arr_flights, arr_diverted, temperature, wind_speed, month,
        arr_delay_lag1, arr_flights_lag2, weather_delay, security_delay,
        carrier_delay, arr_delay_lag1_lag2, arr_delay_lag1_lag3
    )
    
    # Create airport terminal crowd simulation
    passenger_x, passenger_y, passenger_colors, normalized_delay = create_terminal_crowd(delay_minutes)
    
    # Create terminal layout for visualization
    fig = make_subplots(rows=1, cols=1)
    
    # Add terminal background
    terminal_width = 100
    terminal_height = 60
    
    # Draw terminal structure
    # Main walking area
    fig.add_shape(
        type="rect",
        x0=0, y0=0, x1=terminal_width, y1=terminal_height,
        line=dict(width=2, color="gray"),
        fillcolor="lightgray",
        opacity=0.3
    )
    
    # Gates
    gates = [
        [10, 5, 20, 10, "Gate A"],
        [40, 5, 20, 10, "Gate B"],
        [70, 5, 20, 10, "Gate C"],
        [10, 45, 20, 10, "Gate D"],
        [40, 45, 20, 10, "Gate E"],
        [70, 45, 20, 10, "Gate F"]
    ]
    
    for gate in gates:
        fig.add_shape(
            type="rect",
            x0=gate[0], y0=gate[1], x1=gate[0]+gate[2], y1=gate[1]+gate[3],
            line=dict(width=1, color="blue"),
            fillcolor="lightblue",
            opacity=0.4
        )
        # Add gate label
        fig.add_annotation(
            x=gate[0]+gate[2]/2, y=gate[1]+gate[3]/2,
            text=gate[4],
            showarrow=False,
            font=dict(size=10)
        )
    
    # Info desk and security points
    service_points = [
        [50, 30, 6, "Information"],
        [25, 30, 4, "Security 1"],
        [75, 30, 4, "Security 2"]
    ]
    
    for point in service_points:
        fig.add_shape(
            type="circle",
            x0=point[0]-point[2], y0=point[1]-point[2], 
            x1=point[0]+point[2], y1=point[1]+point[2],
            line=dict(width=1, color="darkgray"),
            fillcolor="darkgray",
            opacity=0.4
        )
        # Add label
        fig.add_annotation(
            x=point[0], y=point[1],
            text=point[3],
            showarrow=False,
            font=dict(size=10, color="black")
        )
    
    # Add passengers as scatter points
    fig.add_trace(
        go.Scatter(
            x=passenger_x, y=passenger_y,
            mode='markers',
            marker=dict(
                size=7,
                color=passenger_colors,
                opacity=0.7
            ),
            name='Passengers'
        )
    )
    
    # Add delay information
    # Create color for delay text based on severity
    if delay_minutes < 0:
        delay_color = "green"
        delay_text = "Early by"
    elif delay_minutes < 1:
        delay_color = "green"
        delay_text = "On time (+"
    elif delay_minutes < 3:
        delay_color = "orange"
        delay_text = "Delayed by"
    else:
        delay_color = "red"
        delay_text = "Significantly delayed by"
    
    # Add delay annotation
    fig.add_annotation(
        xref="paper", yref="paper",
        x=0.5, y=0.97,
        text=f"<b>{delay_text} {abs(delay_minutes):.1f} minutes</b>",
        showarrow=False,
        font=dict(size=16, color=delay_color),
        align="center",
        bgcolor="rgba(255, 255, 255, 0.8)",
        bordercolor=delay_color,
        borderwidth=2,
        borderpad=4
    )
    
    # Add passenger count and frustration level
    total_passengers = len(passenger_x)
    frustrated_passengers = passenger_colors.count('red') + passenger_colors.count('orange')
    frustration_percent = round(frustrated_passengers / total_passengers * 100)
    
    fig.add_annotation(
        xref="paper", yref="paper",
        x=0.5, y=0.03,
        text=f"Terminal occupancy: {total_passengers} passengers | Frustration level: {frustration_percent}%",
        showarrow=False,
        font=dict(size=12),
        align="center",
        bgcolor="rgba(255, 255, 255, 0.8)",
        borderpad=4
    )
    
    # Customize layout
    fig.update_layout(
        title=f"Airport Terminal Crowd Simulation - {carrier}",
        showlegend=False,
        xaxis=dict(
            range=[-5, terminal_width+5],
            showgrid=False,
            zeroline=False,
            visible=False
        ),
        yaxis=dict(
            range=[-5, terminal_height+5],
            showgrid=False,
            zeroline=False,
            visible=False,
            scaleanchor="x",
            scaleratio=1
        ),
        margin=dict(l=20, r=20, t=50, b=20),
        plot_bgcolor='white'
    )
    
    return fig

# Define input components
inputs = [
    gr.Dropdown(choices=list(carrier_mapping.keys()), label="Carrier"),
    gr.Slider(minimum=slider_ranges['arr_flights']['min'], maximum=slider_ranges['arr_flights']['max'], step=0.1, label="Arriving Flights Previous Month (Scaled)"),
    gr.Slider(minimum=slider_ranges['arr_diverted']['min'], maximum=slider_ranges['arr_diverted']['max'], step=0.1, label="Diverted Flights Previous Month (Scaled)"),
    gr.Slider(minimum=slider_ranges['temperature_2m']['min'], maximum=slider_ranges['temperature_2m']['max'], step=0.1, label="Temperature (Scaled)"),
    gr.Slider(minimum=slider_ranges['wind_speed_10m']['min'], maximum=slider_ranges['wind_speed_10m']['max'], step=0.1, label="Wind Speed (Scaled)"),
    gr.Slider(minimum=1, maximum=12, step=1, label="Month"),
    gr.Slider(minimum=slider_ranges['arr_delay_lag1']['min'], maximum=slider_ranges['arr_delay_lag1']['max'], step=0.1, label="Arrival Delay Previous Month (Scaled)"),
    gr.Slider(minimum=slider_ranges['arr_flights_lag2']['min'], maximum=slider_ranges['arr_flights_lag2']['max'], step=0.1, label="Arriving Flights 2 Months Ago (Scaled)"),
    gr.Slider(minimum=slider_ranges['weather_delay']['min'], maximum=slider_ranges['weather_delay']['max'], step=0.1, label="Weather Delay (Scaled)"),
    gr.Slider(minimum=slider_ranges['security_delay']['min'], maximum=slider_ranges['security_delay']['max'], step=0.1, label="Security Delay (Scaled)"),
    gr.Slider(minimum=slider_ranges['carrier_delay']['min'], maximum=slider_ranges['carrier_delay']['max'], step=0.1, label="Carrier Delay (Scaled)"),
    gr.Slider(minimum=slider_ranges['arr_delay_lag1_lag2']['min'], maximum=slider_ranges['arr_delay_lag1_lag2']['max'], step=0.1, label="Arrival Delay 2 Months Ago (Scaled)"),
    gr.Slider(minimum=slider_ranges['arr_delay_lag1_lag3']['min'], maximum=slider_ranges['arr_delay_lag1_lag3']['max'], step=0.1, label="Arrival Delay 3 Months Ago (Scaled)")
]

outputs = gr.Plot(label="Airport Terminal Crowd Simulation")

gr.Interface(
    fn=predict_delay_interface,
    inputs=inputs,
    outputs=outputs,
    title="Flight Delay Prediction with Terminal Crowd Simulation",
    description="Predict flight delays and visualize how they affect the airport terminal crowd. As delays increase, passengers shift from gates to information desks and security points, with increasing frustration levels (blue = calm seated passengers not experiencing any delay, green = calm passengers who aren't frustrated by the delay, orange = mild frustration, red = high frustration)."
).launch()