import gradio as gr
import pandas as pd
import folium
import requests
from datetime import datetime, timedelta
import io
import os
import xarray as xr
import cfgrib # For reading GRIB2 files
import numpy as np # For numerical operations
import s3fs # For accessing S3
import fsspec # Needed for advanced caching

# --- Configuration ---
# NOAA_API_TOKEN is not needed for public GFS S3 access, but keep for consistency if using other NOAA APIs
NOAA_API_TOKEN = os.getenv("NOAA_API_TOKEN") 

# Base URL/Path for GFS data on AWS S3
GFS_S3_BASE = "s3://noaa-gfs-bdp-pds/"

# --- Helper Functions for GFS Data Retrieval ---

def fetch_gfs_wind_data(forecast_date: datetime, forecast_hour_utc: int = 0):
    """
    Fetches 10m wind data (u and v components) from GFS on AWS S3 for a given date and UTC hour.
    Returns a pandas DataFrame.
    """
    date_str = forecast_date.strftime("%Y%m%d")
    hour_str = f"{forecast_hour_utc:02d}" # e.g., "00", "06"
    
    # We'll fetch the analysis file (f000) for simplicity.
    # For forecasts, change fFFF accordingly.
    grib_file_path = f"{GFS_S3_BASE}{date_str}/{hour_str}/atmos/gfs.t{hour_str}z.pgrb2.0p25.f000.grib2"
    
    print(f"Attempting to open GFS file: {grib_file_path}")

    try:
        # **Crucial Change: Use fsspec's simplecache to specify a local cache directory**
        # This tells cfgrib to download the file (and create its index) in /tmp
        # /tmp is typically writable in containerized environments like Hugging Face Spaces
        cached_grib_file_path = f"simplecache::{grib_file_path}"
        
        # Configure the cache for s3fs (anon=True for public buckets)
        # and specify the cache directory for simplecache
        storage_options = {
            "simplecache": {
                "cache_storage": "/tmp/gfs_cache" # Use /tmp or another suitable writable path
            },
            "s3": {
                "anon": True # Anonymous access for public NOAA S3 bucket
            }
        }

        ds = xr.open_dataset(
            cached_grib_file_path, 
            engine="cfgrib", 
            backend_kwargs={'filter_by_keys': {'typeOfLevel': 'heightAboveGround', 'level': 10, 'shortName': ['u', 'v']}},
            # Pass storage_options to open_dataset
            chunks='auto', # Enable dask for out-of-core processing if files are large
            storage_options=storage_options
        )
        
        # Ensure the variable names are correct. GFS 0.25 typically uses 'u' and 'v'
        # if the filter_by_keys works correctly.
        # If not, you might see 'ugrd10m' and 'vgrd10m' or similar.
        u_var = [v for v in ds.data_vars if 'u' in v.lower() and '10m' in v.lower() or 'u' == v][0]
        v_var = [v for v in ds.data_vars if 'v' in v.lower() and '10m' in v.lower() or 'v' == v][0]

        ds['wind_speed'] = np.sqrt(ds[u_var]**2 + ds[v_var]**2)
        # Wind direction is typically measured clockwise from North.
        # atan2(U, V) gives math angle (counter-clockwise from East).
        # To convert to meteorological direction (from North, clockwise):
        # Angle = (270 - np.degrees(np.arctan2(U, V))) % 360
        ds['wind_direction'] = (270 - np.degrees(np.arctan2(ds[u_var], ds[v_var]))) % 360
        
        # Convert to pandas DataFrame for easier plotting with Folium
        df = ds[['wind_speed', 'wind_direction']].to_dataframe().reset_index()
        
        # Ensure latitude/longitude columns are named consistently for Folium
        # xarray might use 'latitude' and 'longitude' or 'lat' and 'lon'
        if 'latitude' not in df.columns and 'lat' in df.columns:
            df = df.rename(columns={'lat': 'latitude'})
        if 'longitude' not in df.columns and 'lon' in df.columns:
            df = df.rename(columns={'lon': 'longitude'})

        # Filter out NaN values that might result from calculations
        df = df.dropna(subset=['wind_speed', 'wind_direction', 'latitude', 'longitude'])
            
        return df

    except Exception as e:
        print(f"Error fetching GFS data: {e}")
        # Print current working directory and /tmp contents for debugging on HF Spaces
        print(f"Current directory: {os.getcwd()}")
        if os.path.exists('/tmp'):
            print(f"/tmp contents: {os.listdir('/tmp')}")
        return pd.DataFrame()


# --- Gradio Interface Function (remains mostly the same) ---

def visualize_global_wind(gfs_run_hour: str, days_ago: int):
    """
    Fetches global GFS wind data and generates a Folium map.
    """
    # No need for NOAA_API_TOKEN check here as GFS S3 access is anonymous
    
    # Get current time in Bozeman, MT (MDT is -6 UTC in summer, -7 UTC in winter)
    # The current time is 3:06:25 PM MDT on July 2, 2025.
    # MDT is UTC-6. So 3:06 PM MDT is 21:06 UTC.
    # GFS runs are 00Z, 06Z, 12Z, 18Z. The closest previous run to 21:06 UTC is 18Z.
    
    current_utc_time = datetime.utcnow()
    # Calculate the exact run date based on days_ago and selected hour
    # This logic ensures we pick the correct GFS run file path
    
    target_utc_date = current_utc_time - timedelta(days=days_ago)
    
    # GFS data is usually available for the 00, 06, 12, 18 UTC runs.
    # We want the *latest available* run for `days_ago`.
    # Let's say `gfs_run_hour` is '12'. We need the GFS file for `target_utc_date` at `12Z`.

    # The `fetch_gfs_wind_data` function is designed to fetch the analysis (f000) for a specific run.
    # So if you select "0 days ago" and "12" hour, it will attempt to get today's 12Z run.
    # This is correct.
    
    df_wind = fetch_gfs_wind_data(target_utc_date, forecast_hour_utc=int(gfs_run_hour))

    if df_wind.empty:
        return "<h3>Could not retrieve global wind data for the selected parameters. This might happen if the GFS data for this specific run is not yet available on AWS S3 or if there was an error processing it. Try a different run hour or 'Days Ago'.</h3>"

    # Get approximate center for the map (global view)
    # Use min/max lat/lon from the data if available to zoom to the data extent
    # Otherwise, keep a global view
    if not df_wind.empty and 'latitude' in df_wind.columns and 'longitude' in df_wind.columns:
        center_lat = df_wind['latitude'].mean()
        center_lon = df_wind['longitude'].mean()
        zoom_start = 2 # Global view
    else:
        center_lat = 0
        center_lon = 0
        zoom_start = 2

    m = folium.Map(location=[center_lat, center_lon], zoom_start=zoom_start)

    # Limit number of points for visualization performance
    # GFS 0.25 degree has ~259,200 grid points globally. Sampling is crucial.
    sample_df = df_wind.sample(n=min(10000, len(df_wind)), random_state=42) # Sample up to 10k points for better density

    for _, row in sample_df.iterrows():
        # Convert m/s to knots for display (1 m/s = ~1.944 knots)
        wind_speed_knots = row['wind_speed'] * 1.944

        folium.CircleMarker(
            location=[row['latitude'], row['longitude']],
            radius=3 + (wind_speed_knots / 8), # Adjust radius based on speed (knots)
            color='blue',
            fill=True,
            fill_color='darkblue',
            fill_opacity=0.6,
            tooltip=f"Speed: {wind_speed_knots:.1f} knots, Dir: {row['wind_direction']:.0f}°"
        ).add_to(m)

    # Convert the Folium map to HTML
    map_html = m._repr_html_()
    return map_html

# --- Gradio Interface ---

iface = gr.Interface(
    fn=visualize_global_wind,
    inputs=[
        gr.Dropdown(choices=["00", "06", "12", "18"], label="GFS Run Hour (UTC)", value="12"),
        gr.Slider(minimum=0, maximum=3, value=0, step=1, label="Days Ago (for GFS Run Date)")
    ],
    outputs=gr.HTML(label="Interactive Global Wind Map (GFS)"),
    title="Global Wind Data Visualizer (Gradio + Folium)",
    description=(
        "Visualize recent global wind data from NOAA's GFS model on an interactive map. "
        "Select the GFS model run hour (UTC) and how many days ago.<br>"
        "**Note:** Data is from AWS S3 (public access). Visualization samples points for performance."
    )
)

if __name__ == "__main__":
    # Add share=True for a public link
    iface.launch(share=True)