app.py

import gradio as gr
import folium
from folium import plugins
import requests
from datetime import datetime, timedelta
from io import BytesIO
from PIL import Image
import base64
import numpy as np

# Domain bounds for different regions
DOMAIN_BOUNDS = {
    'conus': [[24.0, -125.0], [50.0, -66.0]],
    'alaska': [[51.0, -180.0], [72.0, -130.0]],
    'hawaii': [[18.0, -161.0], [23.0, -154.0]],
    'full': [[18.0, -180.0], [72.0, -66.0]]
}

# HRRR image bounds (approximate for composite reflectivity images)
HRRR_IMAGE_BOUNDS = {
    'conus': [[20.0, -130.0], [52.0, -60.0]],
    'alaska': [[48.0, -180.0], [75.0, -125.0]],
    'hawaii': [[17.0, -162.0], [24.0, -153.0]]
}

def get_available_runs():
    """Generate list of recent model run times"""
    runs = []
    now = datetime.utcnow()
    current_hour = now.replace(minute=0, second=0, microsecond=0)

    for i in range(48):  # Last 48 hours
        run_time = current_hour - timedelta(hours=i)
        runs.append(run_time.strftime("%Y-%m-%d %H:00 UTC"))

    return runs

def try_fetch_hrrr_image(run_time_str, forecast_hour, domain='conus'):
    """
    Try to fetch HRRR composite reflectivity image from various NOAA sources

    Returns:
        PIL Image object if found, None otherwise
    """
    dt = datetime.strptime(run_time_str, "%Y-%m-%d %H:%M UTC")
    run_str = dt.strftime("%Y%m%d%H")

    # Map domain to HRRR naming
    domain_map = {
        'conus': 'conus',
        'alaska': 'alaska',
        'hawaii': 'hawaii',
        'full': 'conus'  # Start with conus for full view
    }

    hrrr_domain = domain_map.get(domain, 'conus')

    # Try multiple URL patterns for HRRR composite reflectivity
    url_patterns = [
        # Pattern 1: Standard for_web structure
        f"https://rapidrefresh.noaa.gov/hrrr/HRRR/for_web/hrrr_ncep_jet/{run_str}/{hrrr_domain}/refc_sfc_f{forecast_hour:02d}.png",
        # Pattern 2: Alternative structure
        f"https://rapidrefresh.noaa.gov/hrrr/HRRR/for_web/hrrr_{hrrr_domain}/{run_str}/refc_sfc_f{forecast_hour:02d}.png",
        # Pattern 3: Simplified path
        f"https://rapidrefresh.noaa.gov/hrrr/for_web/{run_str}/{hrrr_domain}/refc_sfc_f{forecast_hour:02d}.png",
        # Pattern 4: Direct HRRR graphics
        f"https://rapidrefresh.noaa.gov/hrrr/HRRR/displayMapLocalDiskDateDomainZipTZModel.cgi?keys=hrrr_ncep_jet:&runtime={run_str}&plot_type=refc&fcst={forecast_hour:02d}",
    ]

    for url in url_patterns:
        try:
            response = requests.get(url, timeout=10)
            if response.status_code == 200 and len(response.content) > 1000:  # Valid image
                img = Image.open(BytesIO(response.content))
                return img, url
        except:
            continue

    return None, None

def apply_hue_shift(image, hue_shift=0.3):
    """
    Apply hue shift to image to differentiate from NEXRAD

    Args:
        image: PIL Image
        hue_shift: Hue shift amount (0-1), 0.3 = greenish-blue tint

    Returns:
        PIL Image with hue shifted
    """
    # Convert to RGBA if not already
    if image.mode != 'RGBA':
        image = image.convert('RGBA')

    # Convert to numpy array
    img_array = np.array(image)

    # Separate RGB and alpha channels
    rgb = img_array[:, :, :3].astype(float)
    alpha = img_array[:, :, 3]

    # Convert RGB to HSV
    rgb_normalized = rgb / 255.0

    # Simple hue shift by rotating RGB values
    # This gives a greenish-blue tint to HRRR data
    r, g, b = rgb_normalized[:,:,0], rgb_normalized[:,:,1], rgb_normalized[:,:,2]

    # Apply color tint - shift toward cyan/green for HRRR
    r_new = r * 0.6 + g * 0.2 + b * 0.2  # Reduce red
    g_new = r * 0.1 + g * 0.8 + b * 0.1  # Enhance green
    b_new = r * 0.1 + g * 0.3 + b * 0.6  # Moderate blue

    # Stack and convert back
    rgb_shifted = np.stack([r_new, g_new, b_new], axis=2)
    rgb_shifted = np.clip(rgb_shifted * 255, 0, 255).astype(np.uint8)

    # Recombine with alpha
    img_shifted = np.dstack([rgb_shifted, alpha])

    return Image.fromarray(img_shifted, 'RGBA')

def image_to_data_url(image):
    """Convert PIL Image to data URL for folium"""
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

def create_legends():
    """Create HTML legends for both radar types"""
    legend_html = '''
    <div style="position: fixed; top: 80px; right: 20px; z-index: 1000;
                background-color: white; padding: 10px; border-radius: 5px;
                border: 2px solid #333; font-family: Arial; font-size: 11px;
                max-width: 200px;">
        <b style="font-size: 12px;">Reflectivity (dBZ)</b><br>
        <div style="margin-top: 8px;">
            <div style="display: flex; gap: 8px; margin-bottom: 6px;">
                <div style="flex: 1;">
                    <b style="font-size: 10px; color: #d9534f;">NEXRAD</b><br>
                    <div style="margin-top: 3px; font-size: 9px;">Real-time<br>(Standard colors)</div>
                </div>
                <div style="flex: 1;">
                    <b style="font-size: 10px; color: #5cb85c;">HRRR</b><br>
                    <div style="margin-top: 3px; font-size: 9px;">Forecast<br>(Green-blue tint)</div>
                </div>
            </div>
        </div>
        <div style="display: flex; flex-direction: column; margin-top: 8px; gap: 2px; border-top: 1px solid #ccc; padding-top: 6px;">
            <div style="font-size: 10px;"><span style="background: #9854c6; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>60+ (Extreme)</div>
            <div style="font-size: 10px;"><span style="background: #f800fd; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>50-60 (Severe)</div>
            <div style="font-size: 10px;"><span style="background: #bc0000; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>40-50 (Heavy)</div>
            <div style="font-size: 10px;"><span style="background: #fd0000; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>30-40 (Moderate)</div>
            <div style="font-size: 10px;"><span style="background: #fd9500; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>25-30 (Light)</div>
            <div style="font-size: 10px;"><span style="background: #fdf802; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>20-25</div>
            <div style="font-size: 10px;"><span style="background: #02fd02; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>10-20 (Weak)</div>
            <div style="font-size: 10px;"><span style="background: #019ff4; width: 16px; height: 8px; display: inline-block; margin-right: 4px;"></span>&lt;10 (Trace)</div>
        </div>
    </div>
    '''
    return legend_html

def generate_map(run_time_str, forecast_hour, domain_selection,
                 show_nexrad, nexrad_opacity, show_hrrr, hrrr_opacity):
    """
    Generate Folium map with both NEXRAD and HRRR overlays for comparison
    """
    # Set map center and zoom based on domain
    domain_configs = {
        'full': {'location': [45.0, -100.0], 'zoom': 3},
        'conus': {'location': [39.0, -98.0], 'zoom': 4},
        'alaska': {'location': [64.0, -152.0], 'zoom': 4},
        'hawaii': {'location': [20.5, -157.0], 'zoom': 7}
    }

    config = domain_configs.get(domain_selection, domain_configs['conus'])

    # Create base map
    m = folium.Map(
        location=config['location'],
        zoom_start=config['zoom'],
        tiles='CartoDB positron'
    )

    # Add alternative tile layers
    folium.TileLayer('OpenStreetMap', name='Street Map').add_to(m)
    folium.TileLayer('CartoDB dark_matter', name='Dark Map').add_to(m)

    try:
        dt = datetime.strptime(run_time_str, "%Y-%m-%d %H:%M UTC")
        valid_time = dt + timedelta(hours=int(forecast_hour))

        data_status = []

        # Add HRRR forecast overlay if requested
        if show_hrrr:
            hrrr_img, hrrr_url = try_fetch_hrrr_image(run_time_str, int(forecast_hour), domain_selection)

            if hrrr_img:
                # Apply hue shift to HRRR data (greenish-blue tint)
                hrrr_shifted = apply_hue_shift(hrrr_img)
                hrrr_data_url = image_to_data_url(hrrr_shifted)

                # Get bounds for this domain
                bounds = HRRR_IMAGE_BOUNDS.get(domain_selection, HRRR_IMAGE_BOUNDS['conus'])

                folium.raster_layers.ImageOverlay(
                    image=hrrr_data_url,
                    bounds=bounds,
                    opacity=hrrr_opacity,
                    name='HRRR Forecast (Green-Blue Tint)',
                    overlay=True,
                    control=True
                ).add_to(m)

                data_status.append(f"✓ HRRR F{int(forecast_hour):03d} loaded")
            else:
                data_status.append(f"✗ HRRR F{int(forecast_hour):03d} not available")

        # Add NEXRAD real-time radar if requested
        if show_nexrad:
            wms_url = 'https://mapservices.weather.noaa.gov/eventdriven/services/radar/radar_base_reflectivity/MapServer/WMSServer'

            folium.raster_layers.WmsTileLayer(
                url=wms_url,
                layers='0',
                name='NEXRAD Real-Time (Standard Colors)',
                format='image/png',
                transparent=True,
                opacity=nexrad_opacity,
                attr='NOAA',
                overlay=True,
                control=True
            ).add_to(m)

            data_status.append("✓ NEXRAD Real-Time loaded")

        # Add comparison info box
        comparison_html = f"""
        <div style='position: fixed; bottom: 20px; left: 20px; z-index: 1000;
                    background-color: rgba(255, 255, 255, 0.95); padding: 12px;
                    border-radius: 5px; border: 2px solid #0066cc; font-family: Arial;
                    max-width: 340px; box-shadow: 0 2px 5px rgba(0,0,0,0.3);'>
            <b style='color: #0066cc; font-size: 14px;'>📊 Data Comparison View</b><br>
            <div style='margin-top: 8px; font-size: 11px;'>
                <b>Model Run:</b> {dt.strftime("%Y-%m-%d %H:00 UTC")}<br>
                <b>Forecast Hour:</b> F{int(forecast_hour):03d}<br>
                <b>Valid Time:</b> {valid_time.strftime("%Y-%m-%d %H:00 UTC")}<br>
                <b>Domain:</b> {domain_selection.upper()}<br>
            </div>
            <div style='margin-top: 8px; padding: 6px; background: #e8f4f8; border-radius: 3px; font-size: 10px;'>
                <b>Data Status:</b><br>
                {'<br>'.join(data_status) if data_status else 'No data layers selected'}
            </div>
            <div style='margin-top: 6px; padding: 6px; background: #fff3cd; border-radius: 3px; font-size: 9px;'>
                <b>💡 Comparison Tips:</b><br>
                • <span style="color: #d9534f;">Red/Purple</span> = NEXRAD standard colors<br>
                • <span style="color: #5cb85c;">Green/Cyan</span> = HRRR with color shift<br>
                • Overlapping areas appear <b>mixed</b><br>
                • Perfect alignment = good model performance<br>
                • Use opacity sliders to adjust visibility
            </div>
        </div>
        """
        m.get_root().html.add_child(folium.Element(comparison_html))

        # Add legend
        m.get_root().html.add_child(folium.Element(create_legends()))

        # Add domain boundary
        if domain_selection in DOMAIN_BOUNDS:
            bounds = DOMAIN_BOUNDS[domain_selection]
            folium.Rectangle(
                bounds=bounds,
                color='#3388ff',
                fill=False,
                weight=2,
                popup=f"{domain_selection.upper()} Domain"
            ).add_to(m)

    except Exception as e:
        error_html = f"""
        <div style='position: fixed; top: 50%; left: 50%; transform: translate(-50%, -50%);
                    z-index: 1001; background-color: #ffcccc; padding: 20px;
                    border-radius: 10px; border: 2px solid #ff0000; font-family: Arial;'>
            <h3 style='margin-top: 0;'>Error Loading Data</h3>
            <p>Error: {str(e)}</p>
        </div>
        """
        m.get_root().html.add_child(folium.Element(error_html))

    # Add layer control
    folium.LayerControl(position='topright', collapsed=False).add_to(m)

    # Add fullscreen and measure controls
    plugins.Fullscreen(position='topleft').add_to(m)
    plugins.MeasureControl(position='bottomright', primary_length_unit='miles').add_to(m)

    return m

def create_interface():
    """Create Gradio interface"""

    with gr.Blocks(title="HRRR vs NEXRAD Radar Comparison", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🌩️ HRRR vs NEXRAD Radar Comparison Viewer

        Compare NOAA HRRR forecast composite reflectivity with real-time NEXRAD radar data.
        **Both layers shown simultaneously with different colors for visual alignment checking.**

        **Data Sources:** NOAA NEXRAD Real-Time Radar (standard colors) + HRRR Model Forecast (green-blue tint)
        """)

        with gr.Row():
            with gr.Column(scale=1):
                run_time = gr.Dropdown(
                    choices=get_available_runs(),
                    value=get_available_runs()[0],
                    label="🕐 Model Run Time (UTC)",
                    info="HRRR initialization time"
                )

            with gr.Column(scale=1):
                forecast_hour = gr.Slider(
                    minimum=0,
                    maximum=18,
                    step=1,
                    value=0,
                    label="⏱️ Forecast Hour",
                    info="F000 = Analysis (best for comparison)"
                )

            with gr.Column(scale=1):
                domain = gr.Radio(
                    choices=['conus', 'alaska', 'hawaii'],
                    value='conus',
                    label="🗺️ Domain",
                    info="Geographic region"
                )

        gr.Markdown("### 🎨 Layer Controls - Adjust to Compare Alignment")

        with gr.Row():
            with gr.Column(scale=1):
                show_nexrad = gr.Checkbox(
                    value=True,
                    label="📡 Show NEXRAD Real-Time",
                    info="Standard red/purple colors"
                )
                nexrad_opacity = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.6,
                    step=0.1,
                    label="NEXRAD Opacity",
                    info="Lower to see HRRR underneath"
                )

            with gr.Column(scale=1):
                show_hrrr = gr.Checkbox(
                    value=True,
                    label="🛰️ Show HRRR Forecast",
                    info="Green/cyan color tint"
                )
                hrrr_opacity = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.5,
                    step=0.1,
                    label="HRRR Opacity",
                    info="Adjust for blending"
                )

        load_btn = gr.Button("🔄 Load Comparison View", variant="primary", size="lg")

        with gr.Row():
            map_output = gr.HTML(label="Comparison Map")

        def load_map(run_time, forecast_hour, domain, show_nexrad, nexrad_opacity, show_hrrr, hrrr_opacity):
            m = generate_map(run_time, int(forecast_hour), domain,
                           show_nexrad, nexrad_opacity, show_hrrr, hrrr_opacity)
            return m._repr_html_()

        load_btn.click(
            fn=load_map,
            inputs=[run_time, forecast_hour, domain, show_nexrad, nexrad_opacity, show_hrrr, hrrr_opacity],
            outputs=map_output
        )

        # Auto-load on startup
        demo.load(
            fn=load_map,
            inputs=[run_time, forecast_hour, domain, show_nexrad, nexrad_opacity, show_hrrr, hrrr_opacity],
            outputs=map_output
        )

        gr.Markdown("""
        ---
        ## 📊 How to Use This Comparison Tool

        ### Visual Alignment Check

        1. **Set Forecast Hour to 0** (F000 = HRRR analysis)
        2. **Enable both NEXRAD and HRRR** layers
        3. **Adjust opacity sliders** to see both layers clearly
        4. **Look for alignment:**
           - **Perfect overlap** = HRRR correctly assimilated radar data
           - **Offset/misalignment** = potential data issues or timing differences
           - **Different intensities** = model vs. observation differences

        ### Color Coding

        - **🔴 NEXRAD (Standard Colors):** Red, purple, yellow = Real-time radar observations
        - **🟢 HRRR (Green-Blue Tint):** Cyan, green = Model forecast with color shift
        - **Mixed Areas:** Where both overlap, you'll see blended colors

        ### Recommended Settings for Comparison

        | Purpose | NEXRAD Opacity | HRRR Opacity | Notes |
        |---------|---------------|--------------|-------|
        | Check alignment | 0.6 | 0.5 | Balanced visibility |
        | Focus on NEXRAD | 0.8 | 0.3 | HRRR as reference |
        | Focus on HRRR | 0.3 | 0.7 | NEXRAD as reference |
        | See differences | 0.5 | 0.5 | Equal blending |

        ### Understanding Forecast Hours

        - **F000**: HRRR analysis - should match NEXRAD closely (uses radar data assimilation)
        - **F001-F003**: Very short-term forecast - minor divergence expected
        - **F006-F012**: Short-term forecast - moderate divergence
        - **F012-F018**: Medium-range forecast - larger divergence from real-time

        ### Model Coverage

        - **CONUS HRRR**: Continental US at 3km resolution, updated hourly
        - **Alaska HRRR**: Alaska domain at 3km resolution
        - **Hawaii HRRR**: Hawaiian Islands at 3km resolution

        ### Data Availability

        **HRRR Data:**
        - Images may not always be available from rapidrefresh.noaa.gov
        - Try recent run times (last 6-12 hours) for best availability
        - GRIB2 data always available from NOMADS/AWS S3

        **NEXRAD Data:**
        - Real-time WMS service, always current
        - Updates every ~5 minutes
        - Covers CONUS, Alaska, Hawaii, Puerto Rico

        ### 🔗 References

        - [HRRR Information](https://rapidrefresh.noaa.gov/hrrr/)
        - [NEXRAD Documentation](https://www.ncei.noaa.gov/products/radar/next-generation-weather-radar)
        - [HRRR on AWS](https://registry.opendata.aws/noaa-hrrr/)

        ---

        <p style='text-align: center; color: #666; font-size: 11px;'>
        Data: NOAA | For research/educational purposes only | Not for operational use
        </p>
        """)

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()