app.py

import warnings
warnings.filterwarnings("ignore")

import sys
import argparse
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import seaborn as sns
import gradio as gr
import torch
from matplotlib.patches import Rectangle
from matplotlib.gridspec import GridSpec

# Configure matplotlib rcParams for consistent font sizes
plt.rcParams.update({
    'axes.titlesize': 14,        # Title font size for plots
    'axes.labelsize': 14,        # Axis label font size
    'xtick.labelsize': 14,       # X-axis tick label font size
    'ytick.labelsize': 14,       # Y-axis tick label font size
    'legend.fontsize': 14,       # Legend font size
    'figure.titlesize': 16,      # Figure title font size
    'font.size': 14,             # Base font size
})

# Define specific font sizes for different plot types
PLOT_TITLE_SIZE = plt.rcParams['axes.titlesize']      # 14
IMPORTANCE_TITLE_SIZE = 12                            # Smaller for importance plots

# Parse command line arguments before Gradio processes them
def parse_args():
    parser = argparse.ArgumentParser(description='TeleViT XAI Viewer', add_help=False)
    parser.add_argument('mode', nargs='?', choices=['local'], help='Use local data instead of remote')
    
    # Parse known args to avoid conflicts with Gradio
    args, unknown = parser.parse_known_args()
    
    # Remove our custom arguments from sys.argv so Gradio doesn't see them
    if args.mode == 'local':
        sys.argv = [sys.argv[0]] + unknown
    
    return args.mode == 'local'

# Parse arguments early
USE_LOCAL_DATA = parse_args()

# --- Constants and configuration (copied/kept consistent with main.py) ---
COMMON_VARS = {
    "input_global_shape":  [14, 1, 360, 180],
    "input_local_shape":   [14, 1, 80, 80],
    "input_oci_shape":     [10, 10],
    "patch_global_shape":  [14, 1, 30, 30],
    "patch_local_shape":   [14, 1, 16, 16],
    "patch_oci_shape":     [1, 1],
    "local_vars": ['lst_day', 'mslp', 'ndvi', 'pop_dens', 'ssrd', 'sst', 'swvl1', 't2m_mean', 'tp', 'vpd'],
    "global_vars": ['lst_day', 'mslp', 'ndvi', 'pop_dens', 'ssrd', 'sst', 'swvl1', 't2m_mean', 'tp', 'vpd'],
    "positional_vars": ['cos_lat', 'sin_lat', 'cos_lon', 'sin_lon'],
    "oci_vars": ['oci_censo', 'oci_ea', 'oci_epo', 'oci_ao', 'oci_nao', 'oci_nina34_anom', 'oci_pdo', 'oci_pna', 'oci_soi', 'oci_wp'],
    "var_2_label":  {'oci_censo' : 'Bivariate ENSO',
        'oci_ea' : 'Eastern Atlantic',
        'oci_epo' : 'East Pacific',
        'oci_ao' : 'Arctic',
        'oci_nao' : 'North Atlantic',
        'oci_nina34_anom' : 'NINA 3.4',
        'oci_pdo' : 'Pacific Decadal',
        'oci_pna' : 'Pacific/North Amer.',
        'oci_soi' : 'Southern',
        'oci_wp' : 'Western Pacific',
        'argmax_var': 'Most Important Variable',
         "lst_day": "Land Surface Temperature",
         "mslp": "Mean Sea Level Pressure",
         "ndvi": "Normalized Difference Vegetation Index",
         "pop_dens": "Population Density",
         "ssrd": "Surface Solar Radiation Downwards",
         "sst": "Sea Surface Temperature",
         "swvl1": "Soil Wetness Layer 1",
         "t2m_mean": "2m Temperature",
         "tp": "Total Precipitation",
         "vpd": "Vapour Pressure Deficit"
     },
     'multiplier': {'local': 80*80, 'global': 360*180, 'oci': 10}
}

# Set DATA_PATH based on argument
if USE_LOCAL_DATA:
    DATA_PATH = './data'
else:
    DATA_PATH = 'https://huggingface.co/datasets/iprapas/televit-xai-data/resolve/main'

LEAD_TIMES = [0, 1, 2, 4, 8, 16]
TIME_RANGE = (0, 45)

# Mapping between preset places and (latitude_idx, longitude_idx)
PLACE_TO_IDX: dict[str, tuple[int, int]] = {
    'north_america': (2, 3),
    'mediterranean': (2, 9),
    'north_africa': (3, 10),
    'amazon': (4, 5),
    'south_africa': (5, 10),
    'australia': (5, 15),
    'southeast_asia': (4, 14),
}

# Lazily initialized datasets
_ds_dict: dict | None = None


def _load_datasets_once() -> dict:
    global _ds_dict
    if _ds_dict is not None:
        return _ds_dict

    path_dict: dict[str, str] = {}
    path_dict['attentions'] = f'{DATA_PATH}/attentions-televit_ig-2019.zip'
    path_dict['predictions'] = f'{DATA_PATH}/predictions-televit_ig-2019.zip'
    path_dict['seasfire'] = f'{DATA_PATH}/seasfire_2019_v0.4.zip'
    path_dict['seasfire_1deg'] = f'{DATA_PATH}/seasfire_2019_1deg_v0.4.zip'
    for input_type in ['local', 'global', 'oci']:
        path_dict[f'xai-{input_type}'] = f'{DATA_PATH}/intgrad-televit_ig-{input_type}-2019.zip'

    ds_dict: dict[str, xr.Dataset] = {}
    for key, value in path_dict.items():
        # Open remote zipped zarr stores
        ds_dict[key] = xr.open_zarr(f'zip://::{value}', consolidated=True)

    _ds_dict = ds_dict
    return _ds_dict


def get_available_vars(dataset_type: str) -> list[str]:
    return ['argmax_var'] + COMMON_VARS[f"{dataset_type}_vars"]


def get_code_to_label_map(dataset_type: str) -> dict[str, str]:
    codes = get_available_vars(dataset_type)
    return {code: COMMON_VARS['var_2_label'].get(code, code) for code in codes}


def get_label_to_code_map(dataset_type: str) -> dict[str, str]:
    m = get_code_to_label_map(dataset_type)
    return {label: code for code, label in m.items()}


def get_var_labels(dataset_type: str) -> list[str]:
    return list(get_code_to_label_map(dataset_type).values())


def render_plot(var: str, time: int, lead_time: int, aggregation: str, dataset_type: str):
    ds_dict = _load_datasets_once()

    agg = None if (aggregation is None or str(aggregation).lower() == 'none') else aggregation

    # Select the dataset and the requested variable slice
    ds_key = f"xai-{dataset_type}"
    ds = ds_dict[ds_key].sel(lead_time=lead_time)[var]
    if agg == 'median':
        data = ds.median(dim="time")
        time_label = "median"
    else:
        data = ds.isel(time=int(time))
        # format time for title
        tval = str(data.time.values)
        time_label = tval.split("T")[0]

    # Prepare colormap / normalization
    if var == 'argmax_var':
        labels = COMMON_VARS[f"{dataset_type}_vars"]
        num_classes = len(labels)
        palette = sns.color_palette("muted", n_colors=num_classes)
        cmap = mcolors.ListedColormap(palette)
        boundaries = np.arange(num_classes + 1) - 0.5
        norm = mcolors.BoundaryNorm(boundaries=boundaries, ncolors=num_classes)
    else:
        cmap = 'viridis'
        norm = mcolors.Normalize()

    # Projection and figure
    include_secondary = (dataset_type == 'oci' and var != 'argmax_var')
    proj = ccrs.Robinson()
    if include_secondary:
        fig = plt.figure(figsize=(12, 9))
        ax = fig.add_subplot(2, 1, 1, projection=proj)
    else:
        fig, ax = plt.subplots(figsize=(12, 6), subplot_kw={'projection': proj})

    ax.add_feature(cfeature.LAND, facecolor='lightgray', zorder=0)
    ax.add_feature(cfeature.COASTLINE, linewidth=0.6)
    ax.set_global()
    ax.gridlines(draw_labels=False, linestyle='--', linewidth=0.5, alpha=0.7)

    # Plot
    im = ax.pcolormesh(data['longitude'], data['latitude'], data*COMMON_VARS['multiplier'][dataset_type] if var != 'argmax_var' else data,
                        transform=ccrs.PlateCarree(), cmap=cmap, norm=norm, alpha=0.9)

    # Colorbar
    if var == 'argmax_var':
        labels = COMMON_VARS[f"{dataset_type}_vars"]
        num_classes = len(labels)
        cbar = plt.colorbar(im, ax=ax, orientation='vertical', ticks=range(num_classes), pad=0.05, shrink=0.7)
        cbar.ax.set_yticklabels([COMMON_VARS["var_2_label"].get(x, x) for x in labels])
        cbar.set_label("Variables")
    else:
        cbar = plt.colorbar(im, ax=ax, orientation='vertical', pad=0.05, shrink=0.7)
        cbar.set_label(COMMON_VARS["var_2_label"].get(var, var))

    cbar.ax.tick_params(left=False, right=False, labelleft=False, labelbottom=False)

    # Title
    if var == 'argmax_var':
        ax.set_title(f"Most important {dataset_type} variable - time={time_label} - lead time={lead_time}x8d", fontsize=IMPORTANCE_TITLE_SIZE)
    else:
        ax.set_title(f"{COMMON_VARS['var_2_label'].get(var, var)} importance - time={time_label} - lead time={lead_time}x8d", fontsize=IMPORTANCE_TITLE_SIZE)

    # Secondary subplot for OCI raw variable when applicable
    if include_secondary:
        try:
            ax2 = fig.add_subplot(2, 1, 2)
            ds_dict['seasfire'][var].plot(ax=ax2)
            ax2.set_title(f"{COMMON_VARS['var_2_label'].get(var, var)} (seasfire)")
            # Mark selected time when applicable
            if agg is None:
                try:
                    selected_time = data['time'].values
                    ax2.axvline(selected_time, color='red', linestyle='--', linewidth=1.5)
                except Exception:
                    pass
        except Exception:
            ax2 = fig.add_subplot(2, 1, 2)
            ax2.axis('off')
            ax2.text(0.5, 0.5, 'Unable to plot seasfire variable', ha='center', va='center')

    fig.tight_layout()

    return fig


# --- Gradio UI ---

def _update_var_choices(dataset_type: str):
    labels = get_var_labels(dataset_type)
    default_label = COMMON_VARS['var_2_label'].get('argmax_var', 'argmax_var')
    return gr.update(choices=labels, value=default_label)


def _update_time_interactive(aggregation: str):
    agg = None if (aggregation is None or str(aggregation).lower() == 'none') else aggregation
    return gr.update(interactive=(agg is None))


def _predict(dataset_type: str, var_label: str, lead_time: int, aggregation: str, time: int):
    try:
        # Map human-friendly label back to code
        code = get_label_to_code_map(dataset_type).get(var_label, var_label)
        fig = render_plot(var=code, time=time, lead_time=int(lead_time), aggregation=aggregation, dataset_type=dataset_type)
        return fig
    except Exception as e:
        raise gr.Error(str(e))


def _predict_fixed(dataset_type: str):
    def _fn(var_label: str, lead_time: int, aggregation: str, time: int):
        return _predict(dataset_type, var_label, lead_time, aggregation, time)
    return _fn


def _build_tab_for(dataset_type: str):
    with gr.Tab(dataset_type.capitalize() + " importance per patch"):
        gr.Markdown(f"""
        Visualize importance maps for the {dataset_type} inputs.
        - Each map shows the summed integrated gradients per local patch for the selected variable.
        - "Most Important Variable" colors each patch by the variable with the maximum summed attribution.
        - Select a lead time (8-day steps) and a time index; or use median to aggregate across time.
        """)
        var = gr.Dropdown(label="Variable", choices=get_var_labels(dataset_type), value=COMMON_VARS['var_2_label'].get('argmax_var', 'argmax_var'))
        lead_time = gr.Dropdown(label="Lead Time (x8d)", choices=LEAD_TIMES, value=0)
        aggregation = gr.Dropdown(label="Aggregation", choices=["none", "median"], value="none")
        time = gr.Slider(label="Time Index", minimum=TIME_RANGE[0], maximum=TIME_RANGE[1], step=1, value=0, interactive=True)

        submit = gr.Button("Render")
        out = gr.Plot(label="Map")

        aggregation.change(_update_time_interactive, inputs=aggregation, outputs=time)
        submit.click(_predict_fixed(dataset_type), inputs=[var, lead_time, aggregation, time], outputs=out)


def _get_attentions(ds_attn):
    input_global_shape = [10, 1, 180, 360]
    input_local_shape = [10, 1, 80, 80]
    input_oci_shape = [10, 10]
    patch_global_shape = [10, 1, 30, 30]
    patch_local_shape = [10, 1, 16, 16]
    patch_oci_shape = [1, 1]

    token_global_shape = [x//y for x,y in zip(input_global_shape, patch_global_shape)]  # [1,1,6,12] actually for 180x360
    num_global_tokens = np.prod(token_global_shape)
    token_local_shape = [x//y for x,y in zip(input_local_shape, patch_local_shape)]  # [1,1,5,5]
    num_local_tokens = np.prod(token_local_shape)
    token_oci_shape = [x//y for x,y in zip(input_oci_shape, patch_oci_shape)]  # [10,10]
    num_oci_tokens = np.prod(token_oci_shape)

    attns = ds_attn
    attns_t = attns

    local_to_others_attn = attns_t[:num_local_tokens].mean(dim=0)
    local_to_local_attn = local_to_others_attn[:num_local_tokens].reshape(*token_local_shape).mean(dim=(0,1))

    start_idx = num_local_tokens
    end_idx = start_idx + num_global_tokens
    local_to_global_attn = local_to_others_attn[start_idx:end_idx].reshape(*token_global_shape).mean(dim=(0,1))

    start_idx = end_idx
    end_idx = start_idx + num_oci_tokens
    local_to_oci_attn = local_to_others_attn[start_idx:end_idx].reshape(*token_oci_shape)

    return local_to_local_attn, local_to_global_attn, local_to_oci_attn


def _predict_global_comparison(lead_time: int, time: int):
    ds_dict = _load_datasets_once()
    
    ds = ds_dict['seasfire'] if 'seasfire' in ds_dict else None
    preds_ds = ds_dict['predictions'] if 'predictions' in ds_dict else None
    
    plots = []
    
    # Create Target plot (global)
    fig_target, ax_target = plt.subplots(figsize=(15, 8), subplot_kw={'projection': ccrs.PlateCarree()})
    
    if ds is not None:
        target_data = ds.isel(time=time)
        
        # Add cartopy features
        ax_target.add_feature(cfeature.COASTLINE, linewidth=0.5)
        ax_target.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3)
        
        # Plot target data - using the mask as specified: (gwis_ba > 0).where(gwis_ba > 0)
        mask = target_data['gwis_ba']
        mask_filtered = np.log(mask + 1).where(mask > 0)
        im_target = ax_target.pcolormesh(target_data['longitude'], target_data['latitude'], mask_filtered,
                                       transform=ccrs.PlateCarree(), alpha=0.9, cmap='Reds')
        
        # Set global extent
        ax_target.set_global()
        ax_target.gridlines(draw_labels=False, linestyle='--', alpha=0.5,
                          xlocs=np.arange(-180, 181, 30), ylocs=np.arange(-90, 91, 30))
        
        # Format time for title
        tval = str(target_data.time.values)
        time_label = tval.split("T")[0]
        ax_target.set_title(f'Global Targets (Burned Area > 0) - time={time_label}')
        
        cbar_target = fig_target.colorbar(im_target, ax=ax_target, fraction=0.046, pad=0.04)
        cbar_target.set_label('Fire Occurrence')
    else:
        ax_target.axis('off')
        ax_target.text(0.5, 0.5, 'Target dataset not available', ha='center', va='center', transform=ax_target.transAxes)
    
    fig_target.tight_layout()
    plots.append(fig_target)

    # Create Prediction plot (global)
    fig_pred, ax_pred = plt.subplots(figsize=(15, 8), subplot_kw={'projection': ccrs.PlateCarree()})
    
    try:
        if preds_ds is not None:
            preds = preds_ds[f'predictions_{lead_time}'].isel(time=time)
            
            # Add cartopy features
            ax_pred.add_feature(cfeature.COASTLINE, linewidth=0.5)
            ax_pred.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3)
            
            # Plot predictions
            im_pred = ax_pred.pcolormesh(preds['longitude'], preds['latitude'], 
                                       np.ma.masked_where(preds.values < 0.01, preds.values),
                                       transform=ccrs.PlateCarree(), alpha=0.9, cmap='Spectral_r', vmin=0, vmax=1)
            
            # Set global extent
            ax_pred.set_global()
            ax_pred.gridlines(draw_labels=False, linestyle='--', alpha=0.5,
                            xlocs=np.arange(-180, 181, 30), ylocs=np.arange(-90, 91, 30))
            
            # Format time for title
            tval = str(preds.time.values)
            time_label = tval.split("T")[0]
            ax_pred.set_title(f'Global Predictions - time={time_label} - lead time={lead_time}x8d')
            
            cbar_pred = fig_pred.colorbar(im_pred, ax=ax_pred, fraction=0.046, pad=0.04)
            cbar_pred.set_label('Prediction Confidence')
        else:
            ax_pred.axis('off')
            ax_pred.text(0.5, 0.5, 'Predictions dataset not available', ha='center', va='center', transform=ax_pred.transAxes)
    except Exception as e:
        ax_pred.axis('off')
        ax_pred.text(0.5, 0.5, f'Unable to plot predictions: {str(e)}', ha='center', va='center', transform=ax_pred.transAxes)
    
    fig_pred.tight_layout()
    plots.append(fig_pred)
    
    return plots


def _predict_attentions(place: str, latitude_idx: int, longitude_idx: int, lead_time: int, time: int):
    ds_dict = _load_datasets_once()

    # Resolve place presets
    if place and place in PLACE_TO_IDX:
        latitude_idx, longitude_idx = PLACE_TO_IDX[place]

    ds_attns_tmp = ds_dict[f'attentions'][f'attentions_{lead_time}'].isel(time=time, latitude=latitude_idx, longitude=longitude_idx)
    longitude = ds_attns_tmp.longitude
    latitude = ds_attns_tmp.latitude
    local_to_local, local_to_global_attn, local_to_oci_attn = _get_attentions(torch.from_numpy(ds_attns_tmp.to_numpy()))

    ds = ds_dict['seasfire'] if 'seasfire' in ds_dict else None
    ds_1deg = ds_dict['seasfire_1deg'] if 'seasfire_1deg' in ds_dict else None
    preds_ds = ds_dict['predictions'] if 'predictions' in ds_dict else None

    plots = []
    
    # Create Target plot
    fig_target, ax_target = plt.subplots(figsize=(10, 8), subplot_kw={'projection': ccrs.PlateCarree()})
    
    if ds is not None:
        target_data = ds.sel(longitude=slice(longitude-10, longitude+10), latitude=slice(latitude+10, latitude-10)).isel(time=time)
                
        import numpy as np
        import matplotlib.colors as mcolors

        # Create binary mask (1 where burned area, 0 elsewhere)
        mask = (target_data['gwis_ba'] > 0).astype(int)

        # Create categorical colormap: index 0 = transparent, index 1 = red
        cmap = mcolors.ListedColormap([(1, 0, 0, 0),  # RGBA for transparent
                                    (1, 0, 0, 1)]) # RGBA for red
        bounds = [-0.5, 0.5, 1.5]
        norm = mcolors.BoundaryNorm(bounds, cmap.N)

                # Add cartopy features
        ax_target.add_feature(cfeature.COASTLINE, linewidth=0.5)
        ax_target.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3)

        # Plot binary mask
        im_target = ax_target.pcolormesh(
            target_data['longitude'], target_data['latitude'], mask,
            transform=ccrs.PlateCarree(), cmap=cmap, norm=norm
        )

        # Set extent based on data
        ax_target.set_extent(
            [target_data['longitude'].min(), target_data['longitude'].max(),
            target_data['latitude'].min(), target_data['latitude'].max()],
            ccrs.PlateCarree()
        )
        ax_target.gridlines(draw_labels=False, linestyle='--', alpha=0.5, 
                            xlocs=np.arange(-180, 181, 4), ylocs=np.arange(-90, 91, 4))
        ax_target.set_title('Burned Area Mask')
    
    fig_target.tight_layout()
    plots.append(fig_target)

    # Create Prediction plot
    fig_pred, ax_pred = plt.subplots(figsize=(10, 8), subplot_kw={'projection': ccrs.PlateCarree()})
    
    try:
        if preds_ds is not None:
            preds = preds_ds[f'predictions_{lead_time}'].isel(time=time)
            preds_sel = preds.sel(longitude=slice(longitude-10, longitude+10), latitude=slice(latitude+10, latitude-10))
            
            # Add cartopy features
            ax_pred.add_feature(cfeature.COASTLINE, linewidth=0.5)
            ax_pred.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3)
            
            # Plot predictions
            im_pred = ax_pred.pcolormesh(preds_sel['longitude'], preds_sel['latitude'], 
                                       preds_sel.values, 
                                       transform=ccrs.PlateCarree(), alpha=0.9, cmap='Spectral_r')
            
            # Set extent and styling
            ax_pred.set_extent([preds_sel['longitude'].min(), preds_sel['longitude'].max(),
                              preds_sel['latitude'].min(), preds_sel['latitude'].max()], ccrs.PlateCarree())
            ax_pred.gridlines(draw_labels=False, linestyle='--', alpha=0.5,
                            xlocs=np.arange(-180, 181, 4), ylocs=np.arange(-90, 91, 4))
            ax_pred.set_title('Prediction')
            cbar_pred = fig_pred.colorbar(im_pred, ax=ax_pred, fraction=0.046, pad=0.04)
            cbar_pred.set_label('Confidence')
        else:
            ax_pred.axis('off')
            ax_pred.text(0.5, 0.5, 'Predictions dataset not available', ha='center', va='center', transform=ax_pred.transAxes)
    except Exception:
        ax_pred.axis('off')
        ax_pred.text(0.5, 0.5, 'Unable to plot predictions', ha='center', va='center', transform=ax_pred.transAxes)
    
    fig_pred.tight_layout()
    plots.append(fig_pred)

    # Create Local to Local Attention plot
    fig_local, ax_local = plt.subplots(figsize=(10, 8), subplot_kw={'projection': ccrs.PlateCarree()})
    
    # Add cartopy features
    ax_local.add_feature(cfeature.COASTLINE, linewidth=0.5)
    ax_local.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3)
    
    if ds is not None:
        target_data = ds.sel(longitude=slice(longitude-10, longitude+10), latitude=slice(latitude+10, latitude-10)).isel(time=time)
        
        # Create coordinate grids for the attention data
        lon_coords = target_data['longitude']
        lat_coords = target_data['latitude']
        
        # Interpolate attention data to match the spatial resolution
        attention_interp = torch.nn.functional.interpolate(
            local_to_local.unsqueeze(0).unsqueeze(0), 
            size=(len(lat_coords), len(lon_coords)), 
            mode='bilinear'
        ).squeeze().detach().numpy()
        
        # Plot attention data
        im_local = ax_local.pcolormesh(lon_coords, lat_coords, attention_interp,
                                     transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9)
        
        # Set extent and styling
        ax_local.set_extent([lon_coords.min(), lon_coords.max(),
                           lat_coords.min(), lat_coords.max()], ccrs.PlateCarree())
    else:
        # Fallback: create simple coordinate grid around the selected location
        lon_range = np.linspace(longitude-10, longitude+10, 20)
        lat_range = np.linspace(latitude-10, latitude+10, 20)
        attention_interp = torch.nn.functional.interpolate(
            local_to_local.unsqueeze(0).unsqueeze(0), 
            size=(20, 20), 
            mode='bilinear'
        ).squeeze().detach().numpy()
        
        im_local = ax_local.pcolormesh(lon_range, lat_range, attention_interp,
                                     transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9)
        ax_local.set_extent([lon_range.min(), lon_range.max(),
                           lat_range.min(), lat_range.max()], ccrs.PlateCarree())
    
    ax_local.gridlines(draw_labels=False, linestyle='--', alpha=0.5,
                     xlocs=np.arange(-180, 181, 4), ylocs=np.arange(-90, 91, 4))
    ax_local.set_title('Local to Local Attention')
    cbar_local = fig_local.colorbar(im_local, ax=ax_local, fraction=0.046, pad=0.04)
    cbar_local.set_label('Attention weight')
    
    fig_local.tight_layout()
    plots.append(fig_local)

    # Create Local to OCI Attention plot
    fig_oci, ax_oci = plt.subplots(figsize=(10, 8))
    
    im_oci = ax_oci.imshow(local_to_oci_attn, cmap='viridis', alpha=0.9)
    ax_oci.set_title('Local to OCI Attention')
    ax_oci.set_yticks(np.arange(10))
    ax_oci.set_yticklabels([COMMON_VARS["var_2_label"].get(x, x) for x in COMMON_VARS["oci_vars"]])
    ax_oci.set_xticks(np.arange(10))
    ax_oci.set_xticklabels(np.arange(10, 0, -1))
    ax_oci.set_ylabel('OCI Variables')
    ax_oci.set_xlabel('Months Before Prediction')
    cbar_oci = fig_oci.colorbar(im_oci, ax=ax_oci, fraction=0.046, pad=0.04)
    cbar_oci.set_label('Attention weight')
    
    fig_oci.tight_layout()
    plots.append(fig_oci)

    # Create Local to Global Attention plot
    fig_global, ax_global = plt.subplots(figsize=(15, 8), subplot_kw={'projection': ccrs.PlateCarree()})
    
    # Add cartopy features
    ax_global.add_feature(cfeature.COASTLINE, linewidth=0.5)
    ax_global.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3)
    
    if ds_1deg is not None:
        # Create coordinate grids for the global attention data
        global_lons = ds_1deg['longitude']
        global_lats = ds_1deg['latitude']
        
        # Interpolate attention data to match the global grid
        attention_global_interp = torch.nn.functional.interpolate(
            local_to_global_attn.unsqueeze(0).unsqueeze(0), 
            size=(len(global_lats), len(global_lons)), 
            mode='bilinear'
        ).squeeze().detach().numpy()
        
        # Plot global attention
        im_global = ax_global.pcolormesh(global_lons, global_lats, attention_global_interp,
                                       transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9)
        
        # Draw rectangle for the corresponding local patch on the global grid
        try:
            lat_min = latitude - 10
            lat_max = latitude + 10
            lon_min = longitude - 10
            lon_max = longitude + 10

            lat_min_v = float(np.asarray(getattr(lat_min, 'values', lat_min)))
            lat_max_v = float(np.asarray(getattr(lat_max, 'values', lat_max)))
            lon_min_v = float(np.asarray(getattr(lon_min, 'values', lon_min)))
            lon_max_v = float(np.asarray(getattr(lon_max, 'values', lon_max)))

            # Draw rectangle in geographic coordinates
            from matplotlib.patches import Rectangle as GeoRectangle
            rect = GeoRectangle((lon_min_v, lat_min_v), 
                              lon_max_v - lon_min_v, lat_max_v - lat_min_v,
                              linewidth=2.5, edgecolor='cyan', facecolor='none', alpha=0.9,
                              transform=ccrs.PlateCarree())
            ax_global.add_patch(rect)
        except Exception:
            pass
            
    else:
        # Fallback: create global coordinate grid
        global_lons = np.linspace(-180, 180, 360)
        global_lats = np.linspace(-90, 90, 180)
        
        attention_global_interp = torch.nn.functional.interpolate(
            local_to_global_attn.unsqueeze(0).unsqueeze(0), 
            size=(180, 360), 
            mode='bilinear'
        ).squeeze().detach().numpy()
        
        im_global = ax_global.pcolormesh(global_lons, global_lats, attention_global_interp,
                                       transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9)

    # Set global extent and styling
    ax_global.set_global()
    ax_global.gridlines(draw_labels=False, linestyle='--', alpha=0.5,
                      xlocs=np.arange(-180, 181, 30), ylocs=np.arange(-90, 91, 30))
    ax_global.set_title('Local to Global Attention')
    cbar_global = fig_global.colorbar(im_global, ax=ax_global, fraction=0.046, pad=0.04)
    cbar_global.set_label('Attention weight')
    
    fig_global.tight_layout()
    plots.append(fig_global)

    return plots


def _build_attentions_tab():
    with gr.Tab("Attentions"):
        gr.Markdown(
            """
        Explore attention maps at a specific location and time.
        - Local→Local shows how tokens within the local patch attend to each other (upsampled to the 80×80 input grid).
        - Local→Global shows attention from local tokens to global tokens (upsampled to the 180×360 grid).
        - Local→OCI shows attention from local tokens to the 10×10 monthly climate indices grid.
        - The attention overlays are drawn on top of their originating input grids (land/LSM or target), to provide spatial context.
        """
        )
        place = gr.Dropdown(label="Place", choices=[""] + list(PLACE_TO_IDX.keys()), value="mediterranean")
        latitude_idx = gr.Slider(label="Latitude Index", minimum=0, maximum=9, step=1, value=2)
        longitude_idx = gr.Slider(label="Longitude Index", minimum=0, maximum=18, step=1, value=9)
        lead_time = gr.Dropdown(label="Lead Time (x8d)", choices=LEAD_TIMES, value=0)
        time = gr.Slider(label="Time Index", minimum=TIME_RANGE[0], maximum=TIME_RANGE[1], step=1, value=0)

        # Hidden state to indicate when lat/lon updates originate from a place selection
        syncing_from_place = gr.State(False)

        submit = gr.Button("Render")
        target_plot = gr.Plot(label="Target")
        pred_plot = gr.Plot(label="Prediction") 
        local_plot = gr.Plot(label="Local to Local Attention")
        oci_plot = gr.Plot(label="Local to OCI Attention")
        global_plot = gr.Plot(label="Local to Global Attention")
        
        outputs = [target_plot, pred_plot, local_plot, oci_plot, global_plot]

        # When place changes, update lat/lon and set syncing flag so slider change handlers don't clear place
        def _on_place_change(selected_place: str, _flag: bool):
            if selected_place and selected_place in PLACE_TO_IDX:
                lat, lon = PLACE_TO_IDX[selected_place]
                return gr.update(value=lat), gr.update(value=lon), gr.update(value=True)
            return gr.update(), gr.update(), gr.update(value=False)

        place.change(_on_place_change, inputs=[place, syncing_from_place], outputs=[latitude_idx, longitude_idx, syncing_from_place])

        # When user manually changes lat/lon, update place if coordinates match a preset, otherwise clear it
        def _on_index_change(lat: int, lon: int, flag: bool):
            if flag:
                # Reset the flag, keep the current place
                return gr.update(), gr.update(value=False)
            
            # Manual change: check if the coordinates match any preset place
            for place_name, (place_lat, place_lon) in PLACE_TO_IDX.items():
                if lat == place_lat and lon == place_lon:
                    # Coordinates match a preset place, update to that place
                    return gr.update(value=place_name), gr.update(value=False)
            
            # Coordinates don't match any preset, clear place
            return gr.update(value=""), gr.update(value=False)

        latitude_idx.change(_on_index_change, inputs=[latitude_idx, longitude_idx, syncing_from_place], outputs=[place, syncing_from_place])
        longitude_idx.change(_on_index_change, inputs=[latitude_idx, longitude_idx, syncing_from_place], outputs=[place, syncing_from_place])

        submit.click(_predict_attentions, inputs=[place, latitude_idx, longitude_idx, lead_time, time], outputs=outputs)


def _build_global_predictions_tab():
    with gr.Tab("Global Predictions vs Targets"):
        gr.Markdown(
            """
        Visualize global predictions vs targets for wildfire occurrence.
        - **Targets**: Shows fire occurrence (burned area > 0) from the GWIS dataset.
        - **Predictions**: Shows model predictions for the selected lead time.
        - Select lead time (8-day steps) and time index to compare model predictions with actual fire occurrence.
        - Both plots are shown on a global scale to understand spatial patterns of fire activity and model performance.
        """
        )
        lead_time = gr.Dropdown(label="Lead Time (x8d)", choices=LEAD_TIMES, value=0)
        time = gr.Slider(label="Time Index", minimum=TIME_RANGE[0], maximum=TIME_RANGE[1], step=1, value=0)

        submit = gr.Button("Render")
        target_plot = gr.Plot(label="Global Targets")
        pred_plot = gr.Plot(label="Global Predictions") 
        
        outputs = [target_plot, pred_plot]

        submit.click(_predict_global_comparison, inputs=[lead_time, time], outputs=outputs)


def _warmup():
    # Load all datasets once on first page load
    _load_datasets_once()
    return None


def build_interface() -> gr.Blocks:
    data_source_note = f"**Data Source:** {'Local data (`./data`)' if USE_LOCAL_DATA else 'Remote data (HuggingFace)'}"
    loading_note = "Note: The first render may take longer as datasets are loaded from " + ("local storage." if USE_LOCAL_DATA else "remote storage.")
    
    with gr.Blocks(title="TeleViT XAI Viewer") as demo:
        gr.Markdown(f"""
        ### TeleViT XAI Viewer
        {data_source_note}
        
        Select dataset type tab, then choose variable, time, lead time, and aggregation to visualize the corresponding map.
        
        {loading_note}
        
        Scientific context:
        - For local, global, and oci inputs, the maps show the sum of integrated gradients within each local patch for the selected variable. This shows the importance of the variable in the local patch. Choosing "Most Important Variable" renders, for each patch, the variable with the highest summed attribution (argmax over variables).
        - Lead time indexes the prediction horizon in steps of 8 days; time selects the date of the forecast. "median" aggregates attributions across time before rendering.
        - Attention maps visualize transformer attention from local tokens to: Local (within-patch), Global (global tokens), and OCI (10×10 monthly climate indices). These attentions are upsampled to their native input grids (80×80 local; 180×360 global; 10×10 OCI) and overlaid on the corresponding backgrounds to provide spatial context. In the Attentions tab you can choose a preset place or indices, plus lead time and time, to render these views.
        """)

        with gr.Tabs():
            _build_attentions_tab()
            _build_tab_for("local")
            _build_tab_for("global")
            _build_tab_for("oci")
            _build_global_predictions_tab()

        # Trigger dataset loading on first page load (must be inside Blocks context)
        demo.load(_warmup)

    return demo


demo = build_interface()


def main():
    # Print data source information
    if USE_LOCAL_DATA:
        print("Using local data from:", DATA_PATH)
    else:
        print("Using remote data from:", DATA_PATH)
    
    # Ensure datasets are lazily loaded on first use
    demo.queue()
    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    main()