import gradio as gr
import numpy as np
import pandas as pd
import torch
import pickle
from datetime import datetime, timedelta, timezone
import os
import requests
import sys
from huggingface_hub import hf_hub_download
from importlib.util import spec_from_file_location, module_from_spec

# ============================================================================
# CONFIGURATION
# ============================================================================

MODEL_REPO_ID = "assix-research/gcc-weather-forecast-transformer"
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"🚀 Using device: {DEVICE}")

# ============================================================================
# 📥 DOWNLOAD ARTIFACTS
# ============================================================================
print("\n📦 DOWNLOADING MODEL ARTIFACTS...")

def load_artifact(filename):
    try:
        path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=filename)
        return path
    except Exception as e:
        print(f"❌ Failed to download {filename}: {e}")
        return None

scaler_path = load_artifact("feature_scaler.pkl")
metadata_path = load_artifact("metadata.pkl")
model_weights_path = load_artifact("weather_transformer.pt")
model_code_path = load_artifact("model.py")

# ============================================================================
# LOAD RESOURCES
# ============================================================================

FEATURES = ['Temperature', 'Humidity', 'Pressure', 'Wind'] # Fallback
SEQUENCE_LENGTH = 72

if scaler_path:
    with open(scaler_path, 'rb') as f:
        feature_scaler = pickle.load(f)
        if hasattr(feature_scaler, 'feature_names_in_'):
            FEATURES = list(feature_scaler.feature_names_in_)
            print(f"✅ TRUSTING SCALER. True Training Order: {FEATURES}")
        else:
            print("⚠️ Scaler has no internal names, falling back to metadata.")

if metadata_path:
    with open(metadata_path, 'rb') as f:
        metadata = pickle.load(f)
        SEQUENCE_LENGTH = metadata.get('sequence_length', 72)

model = None
if model_code_path and model_weights_path:
    try:
        spec = spec_from_file_location("weather_model", model_code_path)
        weather_model = module_from_spec(spec)
        spec.loader.exec_module(weather_model)
        
        model = weather_model.WeatherTransformer(
            input_dim=len(FEATURES), 
            d_model=256, n_heads=8, n_layers=4, d_ff=512, dropout=0.1, seq_len=SEQUENCE_LENGTH
        ).to(DEVICE)
        
        state_dict = torch.load(model_weights_path, map_location=DEVICE)
        model.load_state_dict(state_dict)
        model.eval()
        print("✅ Model successfully loaded!")
    except Exception as e:
        print(f"❌ Error initializing model: {e}")

# ============================================================================
# DATA HELPERS
# ============================================================================

def format_local_datetime(dt_utc):
    if hasattr(dt_utc, 'to_pydatetime'): dt_utc = dt_utc.to_pydatetime()
    if dt_utc.tzinfo is None: dt_utc = dt_utc.replace(tzinfo=timezone.utc)
    dt_local = dt_utc.astimezone()
    return dt_local.strftime("%H:%M") 

def fetch_weather_context(city_name):
    locations = {
        'Dubai, UAE': {'lat': 25.2048, 'lon': 55.2708},
        'Riyadh, Saudi Arabia': {'lat': 24.7136, 'lon': 46.6753}
    }
    
    if city_name not in locations: return None
    coords = locations[city_name]
    
    end_date = datetime.now(timezone.utc)
    start_date = end_date - timedelta(hours=100)
    
    url = "https://archive-api.open-meteo.com/v1/archive"
    params = {
        'latitude': coords['lat'], 'longitude': coords['lon'],
        'start_date': start_date.strftime("%Y-%m-%d"), 'end_date': end_date.strftime("%Y-%m-%d"),
        'hourly': 'temperature_2m,relative_humidity_2m,pressure_msl,wind_speed_10m',
        'timezone': 'UTC'
    }
    
    try:
        r = requests.get(url, params=params)
        r.raise_for_status()
        data = r.json()
        
        df = pd.DataFrame({
            'time': pd.to_datetime(data['hourly']['time']),
            'temperature_2m': data['hourly']['temperature_2m'],
            'relative_humidity_2m': data['hourly']['relative_humidity_2m'],
            'pressure_msl': data['hourly']['pressure_msl'],
            'wind_speed_10m': data['hourly']['wind_speed_10m']
        })
        if len(df) >= SEQUENCE_LENGTH:
            return df.tail(SEQUENCE_LENGTH)
        return None
    except Exception:
        return None

# ============================================================================
# CORE PREDICTION LOGIC
# ============================================================================

def predict_weather(city, hours_ahead):
    if model is None or feature_scaler is None:
        return "System initializing...", None, None, gr.update(visible=True), "Temperature"
    
    # 1. Fetch History
    df = fetch_weather_context(city)
    if df is None:
        return f"❌ Data fetch failed for {city}", None, None, gr.update(), "Temperature"
    
    # --- STEP A: DYNAMIC COLUMN MAPPING ---
    col_map = {
        'Temperature': 'temperature_2m', 'Temp': 'temperature_2m', 'temperature': 'temperature_2m', 'temperature_2m': 'temperature_2m',
        'Humidity': 'relative_humidity_2m', 'relative_humidity': 'relative_humidity_2m',
        'Pressure': 'pressure_msl', 'pressure': 'pressure_msl',
        'Wind': 'wind_speed_10m', 'Wind Speed': 'wind_speed_10m', 'wind_speed': 'wind_speed_10m'
    }
    
    ordered_columns = []
    target_idx = 0 
    
    for idx, feat in enumerate(FEATURES):
        feat_lower = feat.lower()
        matched_col = None
        for k, v in col_map.items():
            if k.lower() in feat_lower:
                matched_col = v
                break
        
        if matched_col:
            ordered_columns.append(matched_col)
            if 'temp' in feat_lower:
                target_idx = idx
        else:
            ordered_columns.append('temperature_2m')
            
    # 2. Extract Data
    raw_data = df[ordered_columns].values
    X_scaled = feature_scaler.transform(raw_data)
    current_input = torch.FloatTensor(X_scaled).unsqueeze(0).to(DEVICE)
    
    last_time = df.iloc[-1]['time']
    current_vals = df.iloc[-1][ordered_columns].values
    
    future_timestamps = []
    future_predictions_actual = []
    
    # 3. Autoregressive Loop
    with torch.no_grad():
        for i in range(int(hours_ahead)):
            y_next_scaled = model(current_input) 
            y_np = y_next_scaled.cpu().numpy()
            
            last_input_step = current_input[0, -1, :].cpu().numpy()
            new_step_scaled = last_input_step.copy()
            
            if y_np.size == 1:
                new_step_scaled[target_idx] = y_np.item()
            else:
                new_step_scaled = y_np.reshape(len(FEATURES))
            
            new_step_scaled = new_step_scaled.reshape(1, len(FEATURES))
            pred_actual = feature_scaler.inverse_transform(new_step_scaled)[0]
            
            future_predictions_actual.append(pred_actual)
            
            next_step_tensor = torch.FloatTensor(new_step_scaled).unsqueeze(0).to(DEVICE)
            current_input = torch.cat((current_input[:, 1:, :], next_step_tensor), dim=1)
            
            last_time += timedelta(hours=1)
            future_timestamps.append(format_local_datetime(last_time))

    # 4. Post-Process
    predictions_actual = np.array(future_predictions_actual)
    
    # --- LOGIC FIX: Select Correct Prediction for Display ---
    if int(hours_ahead) > 1:
        # For long forecasts, show the LAST predicted value (e.g., value at hour 12)
        display_pred = predictions_actual[-1]
        time_label = f"+{int(hours_ahead)} Hours"
        display_feats = ['Temperature']
        grid_cols = "1fr" 
    else:
        # For 1h, show the IMMEDIATE next value
        display_pred = predictions_actual[0]
        time_label = "Next Hour"
        display_feats = ['Temperature', 'Humidity', 'Pressure', 'Wind']
        grid_cols = "1fr 1fr 1fr 1fr" 

    # --- DISPLAY CARDS LOGIC ---
    pred_dict = {feat: val for feat, val in zip(FEATURES, display_pred)}
    curr_dict = {feat: val for feat, val in zip(FEATURES, current_vals)}
    
    units = {'Temperature': '°C', 'Humidity': '%', 'Pressure': 'hPa', 'Wind': 'km/h'}
    colors = {'Temperature': '#ff6b6b', 'Humidity': '#4dabf7', 'Pressure': '#ffd43b', 'Wind': '#51cf66'}
    icons = {'Temperature': '🌡️', 'Humidity': '💧', 'Pressure': '🧭', 'Wind': '💨'}
    
    cards_html = f"""
    <div style="font-family: sans-serif; margin-bottom: 20px;">
        <div style="display: grid; grid-template-columns: {grid_cols}; gap: 10px;">
    """
    
    for display_name in display_feats:
        # Find key
        key = next((k for k in pred_dict if k.lower() == display_name.lower()), None)
        if not key: key = next((k for k in pred_dict if display_name.lower() in k.lower()), None)
        if not key and display_name == 'Wind': key = next((k for k in pred_dict if 'wind' in k.lower()), None)

        val = pred_dict.get(key, 0.0) if key else 0.0
        curr = curr_dict.get(key, 0.0) if key else 0.0
        
        # --- HIDE 0.0 PRESSURE FIX ---
        if display_name == 'Pressure' and val < 1.0:
            continue
        
        delta = val - curr
        arrow = "⬆" if delta > 0 else "⬇"
        
        if abs(delta) < 0.01:
                val_display = f"{curr:.1f}"
                change_text = "Persisted"
                change_color = "#888"
        else:
                val_display = f"{val:.1f}"
                change_text = f"{arrow} {abs(delta):.1f} change"
                change_color = colors[display_name]

        cards_html += f"""
        <div style="background: #2b2b2b; color: white; padding: 15px; border-radius: 8px; border-top: 4px solid {colors[display_name]}; text-align: center;">
            <div style="font-size: 20px; margin-bottom: 5px;">
                {icons[display_name]} <span style="font-size: 0.8em; margin-left: 5px; opacity: 0.9;">{display_name}</span>
            </div>
            <div style="font-size: 12px; opacity: 0.7;">{time_label}</div>
            <div style="font-size: 22px; font-weight: bold;">{val_display} <span style="font-size:12px">{units[display_name]}</span></div>
            <div style="font-size: 11px; color: {change_color}; margin-top: 5px;">{change_text}</div>
        </div>
        """
    cards_html += "</div></div>"

    # --- TREND LOGIC ---
    df_out = pd.DataFrame(predictions_actual, columns=FEATURES)
    df_out['Time'] = future_timestamps
    
    plot_y = next((f for f in FEATURES if 'temp' in f.lower()), FEATURES[0])
    
    temp_start = predictions_actual[0][target_idx] # Compare start of forecast
    temp_end = predictions_actual[-1][target_idx] # To end of forecast
    trend_diff = temp_end - temp_start
    trend_word = "Rise" if trend_diff > 0 else "Drop"
    
    if int(hours_ahead) <= 1:
        plot_visibility = gr.update(visible=False) 
        trend_msg = f"""
        ### ⏱️ Short-Term Outlook
        The model predicts a **{abs(trend_diff):.1f}°C {trend_word}** in temperature over the next hour.
        *Values for Humidity/Wind are persisted from current conditions as the model focuses on thermal trends.*
        """
    else:
        plot_visibility = gr.update(visible=True)
        trend_msg = f"""
        ### 📈 Trend Analysis
        Over the next **{hours_ahead} hours**, the model expects the temperature to **{trend_word}** by **{abs(trend_diff):.1f}°C**.
        """
    
    return cards_html, df_out, trend_msg, plot_visibility, plot_y

# ============================================================================
# GRADIO UI
# ============================================================================

with gr.Blocks() as app:
    
    gr.HTML("""
    <div style="text-align:center; padding: 20px;">
        <h1 style="margin:0;">🌦️ GCC Weather Transformer</h1>
        <p style="opacity: 0.6; margin-top: 5px;">Real-time AI Forecasting for Dubai & Riyadh</p>
        <p style="font-size: 0.85em;">
            <a href="https://huggingface.co/assix-research/gcc-weather-forecast-transformer" target="_blank" style="color:#2563eb; text-decoration:none;">
            🔗 View Model & Source Code on Hugging Face
            </a>
        </p>
    </div>
    """)
    
    with gr.Row():
        city_input = gr.Dropdown(["Dubai, UAE", "Riyadh, Saudi Arabia"], label="📍 Select City", value="Dubai, UAE")
    
    gr.Markdown("### ⏱️ Select Forecast Horizon")
    
    with gr.Row():
        btn_1h = gr.Button("🎯 1 Hour (High Confidence)", variant="primary")
        btn_12h = gr.Button("🌅 12 Hours")
        btn_24h = gr.Button("📅 24 Hours")
    
    with gr.Accordion("Or select custom duration", open=False):
        hours_slider = gr.Slider(minimum=1, maximum=24, step=1, value=1, label="Custom Hours (1-24)")
        btn_custom = gr.Button("Run Custom Forecast")

    gr.Markdown("### ⚡ Forecast Results")
    output_cards = gr.HTML()

    gr.Markdown("---")
    
    with gr.Group():
        trend_text = gr.Markdown("### 📈 Temperature Trend Projection")
        plot = gr.LinePlot(visible=False)

    output_table = gr.Dataframe(visible=False)
    plot_y_col = gr.State("Temperature") 
    
    def update_plot(df, y_col):
        return gr.LinePlot(
            value=df, x="Time", y=y_col, title="Temperature Trend (°C)", tooltip=["Time", y_col], visible=True
        )

    btn_1h.click(lambda: 1, outputs=hours_slider).then(
        predict_weather, inputs=[city_input, hours_slider], outputs=[output_cards, output_table, trend_text, plot, plot_y_col]
    )

    btn_12h.click(lambda: 12, outputs=hours_slider).then(
        predict_weather, inputs=[city_input, hours_slider], outputs=[output_cards, output_table, trend_text, plot, plot_y_col]
    ).success(update_plot, inputs=[output_table, plot_y_col], outputs=plot)

    btn_24h.click(lambda: 24, outputs=hours_slider).then(
        predict_weather, inputs=[city_input, hours_slider], outputs=[output_cards, output_table, trend_text, plot, plot_y_col]
    ).success(update_plot, inputs=[output_table, plot_y_col], outputs=plot)

    btn_custom.click(
        predict_weather, inputs=[city_input, hours_slider], outputs=[output_cards, output_table, trend_text, plot, plot_y_col]
    ).success(update_plot, inputs=[output_table, plot_y_col], outputs=plot)

    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; border-top: 1px solid #eee; padding-top: 20px; color: #888; font-size: 0.8em;">
        developed by assix research 2026 | Powered by Open-Meteo & PyTorch
    </div>
    """)

if __name__ == "__main__":
    app.launch(theme=gr.themes.Soft())