import gradio as gr
import logging
import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import io
import tempfile
from PIL import Image
from pathlib import Path
from matplotlib.colors import Normalize
import re
import dataclasses
import functools
import time

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

# Constants
MODEL_DIR = Path("./model_data")
MODEL_DIR.mkdir(exist_ok=True)

# Import GraphCast components
try:
    logger.info("Importing GraphCast components...")
    import jax
    import haiku as hk
    import xarray as xr
    from google.cloud import storage
    from graphcast import autoregressive, casting, checkpoint, data_utils, graphcast, normalization
    
    logger.info(f"JAX version: {jax.__version__}")
    logger.info(f"JAX devices: {jax.devices()}")
    HAS_GRAPHCAST = True
    logger.info("GraphCast components imported successfully")
except ImportError as e:
    logger.error(f"Error importing GraphCast: {e}")
    HAS_GRAPHCAST = False

# ==================== GRAPHCAST FIX FUNCTIONS ====================

def download_full_model_and_stats():
    """
    Download the full GraphCast model (not small) and normalization statistics.
    
    Returns:
        Tuple of (model_path, stats_paths)
    """
    logger.info("Setting up Google Cloud Storage client...")
    gcs_client = storage.Client.create_anonymous_client()
    gcs_bucket = gcs_client.get_bucket("dm_graphcast")
    dir_prefix = "graphcast/"
    
    # Find full GraphCast model (not small)
    logger.info("Finding full GraphCast model...")
    params_files = [
        name for blob in gcs_bucket.list_blobs(prefix=dir_prefix+"params/")
        if (name := blob.name.removeprefix(dir_prefix+"params/"))
    ]
    
    # Filter for the main GraphCast model, not the small version
    full_models = [name for name in params_files if "small" not in name.lower() and "operational" not in name.lower()]
    if not full_models:
        logger.warning("No full GraphCast model found, trying any model")
        full_models = params_files
    
    if not full_models:
        raise ValueError("No GraphCast models found in the bucket")
    
    model_name = full_models[0]
    model_path = MODEL_DIR / model_name
    
    # Download model if needed
    if not model_path.exists():
        logger.info(f"Downloading model: {model_name}")
        with open(model_path, 'wb') as f:
            model_blob = gcs_bucket.blob(f"{dir_prefix}params/{model_name}")
            logger.info(f"Model size: {model_blob.size / (1024*1024):.2f} MB")
            model_blob.download_to_file(f)
    else:
        logger.info(f"Using existing model: {model_name}")
    
    # Download normalization stats
    stats_files = ['diffs_stddev_by_level.nc', 'mean_by_level.nc', 'stddev_by_level.nc']
    stats_paths = []
    
    for stat_file in stats_files:
        stat_path = MODEL_DIR / stat_file
        stats_paths.append(stat_path)
        
        if not stat_path.exists():
            logger.info(f"Downloading stats: {stat_file}")
            with open(stat_path, 'wb') as f:
                gcs_bucket.blob(f"{dir_prefix}stats/{stat_file}").download_to_file(f)
        else:
            logger.info(f"Using existing stats: {stat_file}")
    
    return model_path, stats_paths

def download_compatible_dataset():
    """
    Download a dataset compatible with full GraphCast model.
    
    Returns:
        Path to the downloaded dataset
    """
    logger.info("Setting up Google Cloud Storage client...")
    gcs_client = storage.Client.create_anonymous_client()
    gcs_bucket = gcs_client.get_bucket("dm_graphcast")
    dir_prefix = "graphcast/"
    
    # List available datasets
    logger.info("Listing available datasets...")
    dataset_files = [
        name for blob in gcs_bucket.list_blobs(prefix=dir_prefix+"dataset/")
        if (name := blob.name.removeprefix(dir_prefix+"dataset/"))
    ]
    
    # For full GraphCast we need a 0.25 degree resolution dataset with 37 levels
    target_datasets = [
        name for name in dataset_files 
        if "res-0.25" in name and "levels-37" in name and "source-era5" in name
    ]
    
    if not target_datasets:
        logger.warning("No exact matching datasets found, trying with other criteria")
        target_datasets = [name for name in dataset_files if "levels-37" in name]
    
    if not target_datasets:
        raise ValueError("No compatible datasets found")
    
    # Sort by number of steps to get the smallest dataset
    target_datasets.sort(key=lambda x: int(re.search(r"steps-(\d+)", x).group(1)) if re.search(r"steps-(\d+)", x) else 99)
    
    dataset_name = target_datasets[0]
    dataset_path = MODEL_DIR / dataset_name
    
    # Download dataset if needed
    if not dataset_path.exists():
        logger.info(f"Downloading dataset: {dataset_name}")
        with open(dataset_path, 'wb') as f:
            dataset_blob = gcs_bucket.blob(f"{dir_prefix}dataset/{dataset_name}")
            logger.info(f"Dataset size: {dataset_blob.size / (1024*1024):.2f} MB")
            dataset_blob.download_to_file(f)
    else:
        logger.info(f"Using existing dataset: {dataset_name}")
    
    return dataset_path

def fix_encoder_shape_mismatch(model_config, task_config, params):
    """
    Fix the encoder shape mismatch by adjusting the encoder's weight matrices.
    
    Args:
        model_config: Original model configuration
        task_config: Task configuration
        params: Model parameters
        
    Returns:
        Tuple of (adjusted_params, new_model_config)
    """
    logger.info("Fixing encoder shape mismatch...")
    
    # Create a new model config with adjusted mesh size
    # Original is typically mesh size 5 or 6, we'll try with 2
    new_model_config = dataclasses.replace(
        model_config,
        mesh_size=2  # Use mesh size 2 (~162 nodes)
    )
    
    logger.info(f"Original mesh size: {model_config.mesh_size}")
    logger.info(f"New mesh size: {new_model_config.mesh_size}")
    
    # Find and fix the problematic encoder weights
    # The error is in 'grid2mesh_gnn/~_networks_builder/encoder_nodes_grid_nodes_mlp/~/linear_0/w'
    if 'grid2mesh_gnn' in params:
        encoder_path = 'grid2mesh_gnn/~_networks_builder/encoder_nodes_grid_nodes_mlp/~/linear_0'
        if encoder_path + '/w' in params:
            # Save original shape for logging
            orig_shape = params[encoder_path + '/w'].shape
            logger.info(f"Original encoder weight shape: {orig_shape}")
            
            # The second dimension (512) should be preserved
            latent_dim = orig_shape[1]
            
            # We'll use a smaller fixed size for the first dimension
            target_nodes = 98  # Target number of nodes from error message
            
            # Resize the weight matrix
            new_weight = np.zeros((target_nodes, latent_dim), dtype=params[encoder_path + '/w'].dtype)
            if orig_shape[0] > target_nodes:
                # If original has more rows, take a subset
                new_weight = params[encoder_path + '/w'][:target_nodes, :]
            else:
                # If original has fewer rows, pad with zeros
                new_weight[:orig_shape[0], :] = params[encoder_path + '/w']
                
            params[encoder_path + '/w'] = new_weight
            logger.info(f"Adjusted encoder weight shape: {params[encoder_path + '/w'].shape}")
            
            # Adjust bias if present
            if encoder_path + '/b' in params:
                # Most biases should remain the same size (usually matches latent_dim)
                # But check just in case
                orig_bias_shape = params[encoder_path + '/b'].shape
                logger.info(f"Original encoder bias shape: {orig_bias_shape}")
                
                # Usually bias has shape (latent_dim,)
                if orig_bias_shape[0] != latent_dim:
                    new_bias = np.zeros((latent_dim,), dtype=params[encoder_path + '/b'].dtype)
                    if orig_bias_shape[0] <= latent_dim:
                        new_bias[:orig_bias_shape[0]] = params[encoder_path + '/b']
                    else:
                        new_bias = params[encoder_path + '/b'][:latent_dim]
                        
                    params[encoder_path + '/b'] = new_bias
                    logger.info(f"Adjusted encoder bias shape: {params[encoder_path + '/b'].shape}")
    
    return params, new_model_config

def adjust_input_grid(dataset, expected_lat=721, expected_lon=1440):
    """
    Adjust the input grid to match the expected dimensions.
    
    Args:
        dataset: Input xarray dataset
        expected_lat: Expected number of latitude points
        expected_lon: Expected number of longitude points
        
    Returns:
        Adjusted dataset
    """
    if 'lat' in dataset.dims and 'lon' in dataset.dims:
        lat_size = dataset.sizes['lat']
        lon_size = dataset.sizes['lon']
        
        logger.info(f"Dataset dimensions: lat={lat_size}, lon={lon_size}")
        
        if lat_size != expected_lat or lon_size != expected_lon:
            logger.info(f"Adjusting grid to {expected_lat}x{expected_lon}...")
            
            # Create target grid
            target_lat = np.linspace(-90, 90, expected_lat)
            target_lon = np.linspace(-180, 180, expected_lon)
            
            # Interpolate to target grid
            adjusted_dataset = dataset.interp(lat=target_lat, lon=target_lon, method='linear')
            
            logger.info(f"Adjusted dimensions: lat={adjusted_dataset.sizes['lat']}, lon={adjusted_dataset.sizes['lon']}")
            return adjusted_dataset
    
    return dataset

def run_graphcast_prediction(model_config, task_config, params, inputs, targets, forcings,
                            diffs_stddev_by_level, mean_by_level, stddev_by_level):
    """
    Run GraphCast prediction with fixed model configuration.
    
    Args:
        model_config: Model configuration
        task_config: Task configuration
        params: Model parameters
        inputs: Input data
        targets: Target data template
        forcings: Forcing data
        diffs_stddev_by_level: Normalization statistics
        mean_by_level: Normalization statistics
        stddev_by_level: Normalization statistics
        
    Returns:
        Prediction output
    """
    logger.info("Setting up prediction function...")
    
    @hk.transform_with_state
    def run_forward(model_config, task_config, inputs, targets_template, forcings):
        predictor = graphcast.GraphCast(model_config, task_config)
        predictor = casting.Bfloat16Cast(predictor)
        predictor = normalization.InputsAndResiduals(
            predictor,
            diffs_stddev_by_level=diffs_stddev_by_level,
            mean_by_level=mean_by_level,
            stddev_by_level=stddev_by_level
        )
        predictor = autoregressive.Predictor(predictor, gradient_checkpointing=True)
        return predictor(inputs, targets_template=targets_template, forcings=forcings)
    
    # Initialize state
    state = {}
    
    # Create JIT-compiled function
    apply_fn = jax.jit(
        functools.partial(
            run_forward.apply,
            params=params,
            state=state,
            model_config=model_config,
            task_config=task_config
        )
    )
    
    # Run prediction
    logger.info("Running prediction...")
    predictions, _ = apply_fn(
        rng=jax.random.PRNGKey(0),
        inputs=inputs,
        targets_template=targets * np.nan,
        forcings=forcings
    )
    
    return predictions

def solve_graphcast_shape_mismatch():
    """
    Main function to solve the GraphCast shape mismatch issue.
    
    Returns:
        Tuple of (success, message, targets, predictions)
    """
    try:
        # Step 1: Download model and stats
        model_path, stats_paths = download_full_model_and_stats()
        
        # Step 2: Load model
        logger.info("Loading model...")
        with open(model_path, 'rb') as f:
            ckpt = checkpoint.load(f, graphcast.CheckPoint)
        
        original_model_config = ckpt.model_config
        task_config = ckpt.task_config
        
        logger.info(f"Model mesh size: {original_model_config.mesh_size}")
        logger.info(f"Model resolution: {original_model_config.resolution}")
        
        # Step 3: Fix encoder shape mismatch
        fixed_params, fixed_model_config = fix_encoder_shape_mismatch(
            original_model_config, task_config, ckpt.params
        )
        
        # Step 4: Download dataset
        dataset_path = download_compatible_dataset()
        
        # Step 5: Load dataset
        logger.info("Loading dataset...")
        with open(dataset_path, 'rb') as f:
            example_batch = xr.load_dataset(f, decode_timedelta=True).compute()
        
        # Step 6: Adjust input grid if needed
        example_batch = adjust_input_grid(example_batch)
        
        # Step 7: Extract inputs, targets, and forcings
        logger.info("Extracting inputs, targets, and forcings...")
        num_steps = 2  # Keep it small for faster prediction
        inputs, targets, forcings = data_utils.extract_inputs_targets_forcings(
            example_batch,
            target_lead_times=slice("6h", f"{num_steps*6}h"),
            **dataclasses.asdict(task_config)
        )
        
        # Step 8: Load normalization data
        logger.info("Loading normalization data...")
        with open(stats_paths[0], 'rb') as f:
            diffs_stddev_by_level = xr.load_dataset(f, decode_timedelta=True).compute()
        with open(stats_paths[1], 'rb') as f:
            mean_by_level = xr.load_dataset(f, decode_timedelta=True).compute()
        with open(stats_paths[2], 'rb') as f:
            stddev_by_level = xr.load_dataset(f, decode_timedelta=True).compute()
        
        # Step 9: Run prediction
        predictions = run_graphcast_prediction(
            fixed_model_config, task_config, fixed_params, inputs, targets, forcings,
            diffs_stddev_by_level, mean_by_level, stddev_by_level
        )
        
        logger.info("Prediction successful!")
        return True, "Prediction successful", targets, predictions
        
    except Exception as e:
        import traceback
        logger.error(f"Error: {e}")
        logger.error(traceback.format_exc())
        return False, f"Error: {str(e)}", None, None

# ==================== VISUALIZATION FUNCTIONS ====================

def select(data, variable, level=None, max_steps=None):
    """Select data based on variable, level, and maximum steps."""
    data = data[variable]
    if "batch" in data.dims:
        data = data.isel(batch=0)
    if max_steps is not None and "time" in data.sizes and max_steps < data.sizes["time"]:
        data = data.isel(time=range(0, max_steps))
    if level is not None and "level" in data.coords:
        data = data.sel(level=level)
    return data

def scale(data, center=None, robust=False):
    """Scale data for visualization."""
    vmin = np.nanpercentile(data, (2 if robust else 0))
    vmax = np.nanpercentile(data, (98 if robust else 100))
    if center is not None:
        diff = max(vmax - center, center - vmin)
        vmin = center - diff
        vmax = center + diff
    return (data, Normalize(vmin, vmax), ("RdBu_r" if center is not None else "viridis"))

def create_visualization(targets, predictions, variable='2m_temperature', level=None):
    """Create visualization of targets, predictions, and differences."""
    logger.info(f"Creating visualization for {variable}" + (f" at level {level}" if level else ""))
    
    # Create figure with subplots
    plt.figure(figsize=(18, 6))
    
    # Prepare data for visualization
    data_dict = {
        "Target": scale(select(targets, variable, level), robust=True),
        "Prediction": scale(select(predictions, variable, level), robust=True),
        "Difference": scale(
            (select(predictions, variable, level) - select(targets, variable, level)),
            center=0, robust=True
        )
    }
    
    # Create subplots
    for i, (title, (plot_data, norm, cmap)) in enumerate(data_dict.items()):
        ax = plt.subplot(1, 3, i+1)
        ax.set_title(title)
        im = ax.imshow(
            plot_data.isel(time=0, missing_dims="ignore"),
            norm=norm,
            origin="lower",
            cmap=cmap
        )
        plt.colorbar(im, ax=ax, shrink=0.7)
    
    # Set title
    title = f"{variable}"
    if level is not None:
        title += f" at {level} hPa"
    plt.suptitle(title, fontsize=16)
    
    # Save to temporary file
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
        plt.savefig(tmp.name, bbox_inches='tight', dpi=100)
        plt.close()
        return tmp.name

def create_sample_visualization():
    """Create a sample visualization when GraphCast is not available."""
    logger.info("Creating sample visualization")
    
    # Create grid for global data
    lats = np.linspace(-90, 90, 181)
    lons = np.linspace(-180, 180, 361)
    lon_grid, lat_grid = np.meshgrid(lons, lats)
    
    # Create sample data
    temp_target = 25 * np.cos(np.radians(lat_grid)) - 5 * np.cos(np.radians(lon_grid/2)) + 15
    temp_pred = 25 * np.cos(np.radians(lat_grid)) - 5 * np.cos(np.radians((lon_grid-20)/2)) + 15
    temp_diff = temp_pred - temp_target
    
    # Create figure with subplots
    plt.figure(figsize=(18, 6))
    
    # Plot target
    ax1 = plt.subplot(1, 3, 1)
    ax1.set_title("Target")
    im1 = ax1.imshow(temp_target, cmap='coolwarm', origin='lower', vmin=-10, vmax=35)
    plt.colorbar(im1, ax=ax1, shrink=0.7)
    
    # Plot prediction
    ax2 = plt.subplot(1, 3, 2)
    ax2.set_title("Prediction")
    im2 = ax2.imshow(temp_pred, cmap='coolwarm', origin='lower', vmin=-10, vmax=35)
    plt.colorbar(im2, ax=ax2, shrink=0.7)
    
    # Plot difference
    ax3 = plt.subplot(1, 3, 3)
    ax3.set_title("Difference")
    im3 = ax3.imshow(temp_diff, cmap='RdBu_r', origin='lower', vmin=-10, vmax=10)
    plt.colorbar(im3, ax=ax3, shrink=0.7)
    
    plt.suptitle("Sample 2m Temperature Visualization", fontsize=16)
    
    # Save to temporary file
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
        plt.savefig(tmp.name, bbox_inches='tight', dpi=100)
        plt.close()
        return tmp.name

# ==================== GRADIO APP FUNCTIONS ====================

def run_model_with_fix(progress=gr.Progress()):
    """Run GraphCast model with shape mismatch fix."""
    if not HAS_GRAPHCAST:
        return create_sample_visualization(), "GraphCast not available. Showing sample visualization."
    
    try:
        progress(0.1, "Setting up GraphCast...")
        progress(0.2, "Downloading full-size model (this may take a while)...")
        progress(0.4, "Downloading compatible dataset...")
        progress(0.5, "Fixing shape mismatch...")
        
        # Run the shape mismatch fix
        success, message, targets, predictions = solve_graphcast_shape_mismatch()
        
        if not success or predictions is None:
            return create_sample_visualization(), f"Error: {message}. Showing sample visualization."
        
        progress(0.8, "Creating visualization...")
        
        # Find a suitable variable to visualize
        if '2m_temperature' in predictions:
            variable = '2m_temperature'
            level = None
        elif 'temperature' in predictions and 'level' in predictions['temperature'].coords:
            variable = 'temperature'
            level = 500  # Mid-troposphere
        else:
            variable = list(predictions.data_vars.keys())[0]
            level = None if 'level' not in predictions[variable].coords else 500
        
        # Create visualization
        viz_path = create_visualization(targets, predictions, variable, level)
        
        progress(1.0, "Done!")
        return viz_path, f"Successfully ran full-size GraphCast with shape fix! Showing {variable}" + (f" at {level} hPa" if level else "")
        
    except Exception as e:
        import traceback
        logger.error(f"Error running model: {e}")
        logger.error(traceback.format_exc())
        return create_sample_visualization(), f"Error: {str(e)}. Showing sample visualization."

# ==================== GRADIO INTERFACE ====================

# Create Gradio interface
with gr.Blocks(title="GraphCast with Shape Fix") as app:
    gr.Markdown("# Full-Size GraphCast Weather Forecasting")
    gr.Markdown("""
    This application runs the full-size GraphCast model (not the small version) with real data.
    It includes a fix for the shape mismatch issue to enable successful execution.
    
    ## About GraphCast
    
    GraphCast is a state-of-the-art machine learning model for global weather forecasting developed by DeepMind.
    It uses graph neural networks to model Earth's atmosphere and can generate forecasts that are competitive 
    with traditional numerical weather prediction systems.
    
    ## The Shape Mismatch Fix
    
    This app implements a solution that:
    
    1. Downloads the full-size GraphCast model (~200MB)
    2. Fixes the encoder shape mismatch by adjusting the mesh size and weight matrices
    3. Ensures the input grid matches the expected dimensions
    4. Runs the model with real data (not simulations)
    
    Click the button below to run the model. Note that the first run may take a few minutes to download 
    the model and dataset, and to compile the JAX code.
    """)
    
    with gr.Row():
        run_btn = gr.Button("Run Full-Size GraphCast")
    
    with gr.Row():
        output_image = gr.Image(label="Weather Forecast")
    
    status_text = gr.Textbox(label="Status", value="Ready to run prediction", interactive=False)
    
    # Connect button to function
    run_btn.click(
        fn=run_model_with_fix,
        inputs=[],
        outputs=[output_image, status_text]
    )
    
    # Show a sample visualization on load
    output_image.value = create_sample_visualization()
    status_text.value = "Ready to run full-size GraphCast with shape fix"

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