import gradio as gr
import xarray as xr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from datetime import datetime, timedelta
import warnings
import logging
import traceback

warnings.filterwarnings('ignore')

# Set up detailed logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Catalog configuration - using latest.zarr URLs as per dynamical-org/notebooks
CATALOG = {
    "NOAA GFS Analysis (Hourly)": {
        "url": "https://data.dynamical.org/noaa/gfs/analysis-hourly/latest.zarr?email=treesixtyweather@gmail.com",
        "type": "analysis",
        "variables": ["temperature_2m", "precipitation", "wind_u_10m", "wind_v_10m", "mean_sea_level_pressure"]
    },
    "NOAA GFS Forecast": {
        "url": "https://data.dynamical.org/noaa/gfs/forecast/latest.zarr?email=treesixtyweather@gmail.com",
        "type": "forecast",
        "variables": ["temperature_2m", "precipitation", "wind_u_10m", "wind_v_10m", "mean_sea_level_pressure"]
    },
    "NOAA GEFS Analysis": {
        "url": "https://data.dynamical.org/noaa/gefs/analysis/latest.zarr?email=treesixtyweather@gmail.com",
        "type": "analysis",
        "variables": ["temperature_2m", "precipitation", "wind_u_10m", "wind_v_10m"]
    },
    "NOAA GEFS Forecast (35-day)": {
        "url": "https://data.dynamical.org/noaa/gefs/forecast-35-day/latest.zarr?email=treesixtyweather@gmail.com",
        "type": "forecast",
        "variables": ["temperature_2m", "precipitation", "wind_u_10m", "wind_v_10m"]
    },
    "NOAA HRRR Forecast (48-hour)": {
        "url": "https://data.dynamical.org/noaa/hrrr/forecast-48-hour/latest.zarr?email=treesixtyweather@gmail.com",
        "type": "forecast",
        "variables": ["temperature_2m", "precipitation", "wind_u_10m", "wind_v_10m"]
    }
}

# Cache for loaded datasets
dataset_cache = {}

def load_dataset(dataset_name, use_cache=True):
    """Load a dataset from the Dynamical catalog"""
    logger.info(f"=== Loading dataset: {dataset_name} ===")

    if use_cache and dataset_name in dataset_cache:
        logger.info(f"Dataset found in cache: {dataset_name}")
        return dataset_cache[dataset_name], None

    try:
        url = CATALOG[dataset_name]["url"]
        logger.info(f"Opening zarr store at: {url}")

        # Open zarr store - using approach from dynamical-org/notebooks
        logger.info("Opening zarr store with chunks=None")
        ds = xr.open_zarr(url, chunks=None)
        logger.info(f"Successfully opened zarr store")
        logger.info(f"Dataset dimensions: {dict(ds.dims)}")
        logger.info(f"Dataset variables: {list(ds.data_vars)}")
        logger.info(f"Dataset coordinates: {list(ds.coords)}")

        if use_cache:
            dataset_cache[dataset_name] = ds
            logger.info(f"Dataset cached: {dataset_name}")

        return ds, None
    except Exception as e:
        error_msg = f"Error loading dataset: {str(e)}"
        logger.error(f"=== ERROR loading {dataset_name} ===")
        logger.error(f"URL: {CATALOG[dataset_name]['url']}")
        logger.error(f"Exception type: {type(e).__name__}")
        logger.error(f"Exception message: {str(e)}")
        logger.error(f"Traceback:\n{traceback.format_exc()}")
        return None, error_msg

def create_map_visualization(dataset_name, variable, time_index=0):
    """Create an interactive map visualization of the selected variable"""
    logger.info(f"=== Creating map visualization ===")
    logger.info(f"Dataset: {dataset_name}, Variable: {variable}, Time index: {time_index}")

    try:
        ds, error = load_dataset(dataset_name)
        if ds is None:
            logger.error(f"Dataset loading returned None: {error}")
            return None, f"Error loading dataset: {dataset_name}\n{error}"

        logger.info(f"Dataset loaded successfully")

        # Check if variable exists
        if variable not in ds.variables:
            available_vars = list(ds.data_vars)
            logger.error(f"Variable '{variable}' not found. Available: {available_vars}")
            return None, f"Variable '{variable}' not found. Available: {available_vars}"

        logger.info(f"Variable '{variable}' found in dataset")

        # Get the data
        data_var = ds[variable]
        logger.info(f"Variable shape: {data_var.shape}, dims: {data_var.dims}")

        # Handle time dimension
        if 'time' in data_var.dims:
            logger.info(f"Time dimension found, length: {len(ds.time)}")
            if time_index >= len(ds.time):
                time_index = 0
            data_var = data_var.isel(time=time_index)
            logger.info(f"Selected time index: {time_index}")

        # Handle ensemble dimension if present
        if 'ensemble' in data_var.dims:
            logger.info(f"Ensemble dimension found, selecting ensemble 0")
            data_var = data_var.isel(ensemble=0)

        logger.info(f"Data variable shape after slicing: {data_var.shape}")

        # Load data into memory (subsample for performance)
        step = max(1, len(ds.latitude) // 200)  # Limit to ~200 points per dimension
        logger.info(f"Subsampling with step: {step}")
        data_var = data_var.isel(latitude=slice(None, None, step), longitude=slice(None, None, step))
        logger.info(f"Computing data values...")
        data_values = data_var.compute().values
        logger.info(f"Data values shape: {data_values.shape}, min: {data_values.min()}, max: {data_values.max()}")

        # Get coordinates
        lats = ds.latitude.isel(latitude=slice(None, None, step)).values
        lons = ds.longitude.isel(longitude=slice(None, None, step)).values
        logger.info(f"Lat range: [{lats.min()}, {lats.max()}], Lon range: [{lons.min()}, {lons.max()}]")

        # Create plotly figure
        fig = go.Figure(data=go.Heatmap(
            z=data_values,
            x=lons,
            y=lats,
            colorscale='RdBu_r',
            hovertemplate='Lat: %{y:.2f}<br>Lon: %{x:.2f}<br>Value: %{z:.2f}<extra></extra>'
        ))

        time_str = ""
        if 'time' in ds[variable].dims:
            time_val = pd.to_datetime(ds.time.isel(time=time_index).values)
            time_str = f" - {time_val.strftime('%Y-%m-%d %H:%M UTC')}"

        fig.update_layout(
            title=f"{dataset_name}: {variable}{time_str}",
            xaxis_title="Longitude",
            yaxis_title="Latitude",
            height=600,
            hovermode='closest'
        )

        logger.info(f"Map visualization created successfully")
        return fig, f"Successfully loaded {dataset_name}"

    except Exception as e:
        error_msg = f"Error creating visualization: {str(e)}"
        logger.error(f"=== ERROR creating visualization ===")
        logger.error(f"Exception type: {type(e).__name__}")
        logger.error(f"Exception message: {str(e)}")
        logger.error(f"Traceback:\n{traceback.format_exc()}")
        return None, error_msg

def get_point_forecast(dataset_name, lat, lon, variable):
    """Get forecast data for a specific point"""
    logger.info(f"=== Getting point forecast ===")
    logger.info(f"Dataset: {dataset_name}, Lat: {lat}, Lon: {lon}, Variable: {variable}")

    try:
        ds, error = load_dataset(dataset_name)
        if ds is None:
            logger.error(f"Dataset loading failed: {error}")
            return None, f"Error loading dataset: {error}"

        if variable not in ds.variables:
            logger.error(f"Variable '{variable}' not found in dataset")
            return None, f"Variable '{variable}' not found in dataset"

        logger.info(f"Selecting nearest point to ({lat}, {lon})")

        # Find nearest point
        data_var = ds[variable].sel(latitude=lat, longitude=lon, method='nearest')

        # Handle ensemble dimension
        if 'ensemble' in data_var.dims:
            logger.info(f"Handling ensemble dimension")
            data_var = data_var.isel(ensemble=0)

        logger.info(f"Point data shape: {data_var.shape}, dims: {data_var.dims}")

        # Load data
        logger.info(f"Computing point data values...")
        data_values = data_var.compute().values
        logger.info(f"Point data computed, shape: {data_values.shape}")

        # Create time series plot
        if 'time' in ds[variable].dims:
            times = pd.to_datetime(ds.time.values)
            logger.info(f"Creating time series plot with {len(times)} time steps")

            fig = go.Figure()
            fig.add_trace(go.Scatter(
                x=times,
                y=data_values,
                mode='lines+markers',
                name=variable
            ))

            fig.update_layout(
                title=f"Point Forecast: {variable} at ({lat:.2f}, {lon:.2f})",
                xaxis_title="Time (UTC)",
                yaxis_title=variable,
                height=400,
                hovermode='x unified'
            )

            # Create data table
            df = pd.DataFrame({
                'Time (UTC)': times,
                variable: data_values
            })

            logger.info(f"Point forecast created successfully")
            return fig, df.to_html(index=False)
        else:
            logger.warning(f"No time dimension found for {variable}")
            return None, f"No time dimension found for {variable}"

    except Exception as e:
        error_msg = f"Error getting point forecast: {str(e)}"
        logger.error(f"=== ERROR getting point forecast ===")
        logger.error(f"Exception type: {type(e).__name__}")
        logger.error(f"Exception message: {str(e)}")
        logger.error(f"Traceback:\n{traceback.format_exc()}")
        return None, error_msg

def update_available_variables(dataset_name):
    """Update the variable dropdown based on selected dataset"""
    logger.info(f"=== Updating available variables for {dataset_name} ===")

    try:
        ds, error = load_dataset(dataset_name, use_cache=False)
        if ds is None:
            logger.warning(f"Could not load dataset, using default variables: {error}")
            return gr.Dropdown(choices=CATALOG[dataset_name]["variables"], value=CATALOG[dataset_name]["variables"][0])

        available_vars = list(ds.data_vars)
        logger.info(f"Available variables: {available_vars}")
        return gr.Dropdown(choices=available_vars, value=available_vars[0] if available_vars else None)
    except Exception as e:
        logger.error(f"Error updating variables: {str(e)}")
        logger.error(f"Traceback:\n{traceback.format_exc()}")
        return gr.Dropdown(choices=CATALOG[dataset_name]["variables"], value=CATALOG[dataset_name]["variables"][0])

# Create Gradio interface
with gr.Blocks(title="Dynamical Weather Catalog Viewer") as app:
    gr.Markdown("""
    # 🌍 Dynamical Weather Catalog Viewer

    Explore weather analysis and forecast data from the [Dynamical.org catalog](https://dynamical.org/catalog/).

    **Features:**
    - Visualize global weather data on interactive maps
    - Click a location to get point forecasts
    - Browse multiple datasets: NOAA GFS, GEFS, and HRRR
    """)

    with gr.Row():
        with gr.Column(scale=1):
            dataset_dropdown = gr.Dropdown(
                choices=list(CATALOG.keys()),
                value=list(CATALOG.keys())[0],
                label="Select Dataset"
            )
            variable_dropdown = gr.Dropdown(
                choices=CATALOG[list(CATALOG.keys())[0]]["variables"],
                value=CATALOG[list(CATALOG.keys())[0]]["variables"][0],
                label="Select Variable"
            )
            time_slider = gr.Slider(
                minimum=0,
                maximum=10,
                step=1,
                value=0,
                label="Time Index"
            )
            load_btn = gr.Button("Load Map", variant="primary")
            status_text = gr.Textbox(label="Status", interactive=False)

        with gr.Column(scale=2):
            map_plot = gr.Plot(label="Map Visualization")

    gr.Markdown("## 📍 Point Forecast")
    gr.Markdown("Enter coordinates to get a time series forecast for a specific location")

    with gr.Row():
        lat_input = gr.Number(value=40.7, label="Latitude", precision=2)
        lon_input = gr.Number(value=-74.0, label="Longitude", precision=2)
        forecast_btn = gr.Button("Get Point Forecast", variant="secondary")

    with gr.Row():
        forecast_plot = gr.Plot(label="Time Series Forecast")

    forecast_table = gr.HTML(label="Forecast Data")

    # Event handlers
    dataset_dropdown.change(
        fn=update_available_variables,
        inputs=[dataset_dropdown],
        outputs=[variable_dropdown]
    )

    load_btn.click(
        fn=create_map_visualization,
        inputs=[dataset_dropdown, variable_dropdown, time_slider],
        outputs=[map_plot, status_text]
    )

    forecast_btn.click(
        fn=get_point_forecast,
        inputs=[dataset_dropdown, lat_input, lon_input, variable_dropdown],
        outputs=[forecast_plot, forecast_table]
    )

if __name__ == "__main__":
    app.launch()