forecasting/inference.py

"""
Inference Script for Solar Flare Prediction Models
=================================================

Runs a trained ViTLocal checkpoint over a folder of AIA data. Computes soft
X-ray (SXR) predictions, and saves attention weights, flux contributions, and
final outputs.

The workflow includes:
- Loading configuration parameters from a YAML file.
- Resolving dynamic variables in the config.
- Loading the model checkpoint and preparing it for inference.
- Performing batched evaluation over AIA/GOES datasets.
- Saving predicted fluxes, ground truth, and visualization-ready artifacts.

"""

import argparse
import re
import sys
import gc
import pandas as pd
import torch
import numpy as np
from torch.utils.data import DataLoader
from pathlib import Path
import os
import yaml
import torch.nn.functional as F
from torch.nn import HuberLoss
from tqdm import tqdm

# Add project root to Python path
PROJECT_ROOT = Path(__file__).parent.parent.absolute()
sys.path.insert(0, str(PROJECT_ROOT))

from forecasting.dataset import AIAGOESDataset
from forecasting.model import ViTLocal


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')




def evaluate_model_on_dataset(model, dataset, batch_size=16, times=None, config_data=None,
                              save_weights=True, input_size=512, patch_size=16, save_flux=False):
    """
    Run batched inference on the dataset and yield predictions, attention maps, and flux data.
    
    Memory optimization: Processes attention weights immediately and moves to CPU to reduce GPU memory usage.

    Parameters
    ----------
    model : torch.nn.Module
        Loaded solar flare prediction model.
    dataset : torch.utils.data.Dataset
        Dataset containing AIA images and corresponding SXR values.
    batch_size : int, default=16
        Number of samples per batch.
    times : list, optional
        List of timestamps corresponding to each sample.
    config_data : dict, optional
        YAML configuration dictionary.
    save_weights : bool, default=True
        Whether to save attention weights for visualization.
    input_size : int, default=512
        Input image resolution.
    patch_size : int, default=16
        Patch size for ViT-based models.
    save_flux : bool, default=False
        Whether to save flux contributions.

    Yields
    ------
    tuple
        (predictions, ground_truth, attention_map, flux_map, global_index)
    """
    model.eval()
    data_device = next(model.parameters()).device
    
    # Optimize DataLoader settings
    num_workers = config_data.get('num_workers', 4) if config_data else 4
    pin_memory = config_data.get('pin_memory', True) if config_data else True
    
    n_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
    use_multi_gpu = str(config_data.get('multi_gpu', False) if config_data else False).lower() == 'true'
    # Unwrapped model used as fallback for batches smaller than n_gpus
    base_model = model.module if isinstance(model, torch.nn.DataParallel) else model
    loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers,
                       pin_memory=pin_memory, shuffle=False,
                       multiprocessing_context='spawn' if num_workers > 0 else None)
    
    # Load SXR normalization only if path is provided and not empty
    sxr_norm = None
    if (config_data and 'data' in config_data and 'sxr_norm_path' in config_data['data'] 
        and config_data['data']['sxr_norm_path'] and config_data['data']['sxr_norm_path'].strip()):
        try:
            sxr_norm = np.load(config_data['data']['sxr_norm_path'])
        except FileNotFoundError:
            sxr_norm = None

    # All models are ViTLocal with localized attention (no CLS token)
    grid_h, grid_w = input_size // patch_size, input_size // patch_size
    
    use_amp = (data_device.type == 'cuda'
               and config_data.get('use_amp', False) if config_data else False)
    try:
      with torch.no_grad():
        for batch_idx, batch in enumerate(loader):
            aia_imgs = batch[0]
            sxr = batch[1]
            aia_imgs = aia_imgs.to(data_device, non_blocking=True)

            # Call model — optionally use AMP (set use_amp: true in config).
            # FP16 on V100 can spike peak memory due to FP32 fallbacks in attention.
            # CRITICAL: ViTLocal defaults to return_attention=True, which uses massive memory
            # Only compute attention weights if we're saving them (save_weights=True)
            # Fall back to single GPU for batches smaller than n_gpus to avoid
            # DataParallel crashing when some replicas receive empty inputs.
            active_model = (base_model
                            if isinstance(model, torch.nn.DataParallel) and aia_imgs.shape[0] < n_gpus
                            else model)
            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
                if save_weights:
                    pred = active_model(aia_imgs, return_attention=True)
                else:
                    pred = active_model(aia_imgs, return_attention=False)

            # Extract outputs (ViTLocal returns tuple of (predictions, attention_weights, flux_contributions) when return_attention=True)
            if isinstance(pred, tuple) and len(pred) >= 3:
                predictions = pred[0]
                weights = pred[1] if save_weights else None
                flux_contributions = pred[2] if save_flux else None
            elif isinstance(pred, tuple) and len(pred) == 2:
                # When return_attention=False, returns (predictions, flux_contributions)
                predictions = pred[0]
                weights = None
                flux_contributions = pred[1] if save_flux else None
            else:
                predictions = pred
                weights = None
                flux_contributions = None

            current_batch_size = predictions.shape[0]
            
            # Process each sample in the batch to reduce memory footprint
            for i in range(current_batch_size):
                global_idx = batch_idx * batch_size + i
                
                # Process attention weights immediately and move to CPU to free GPU memory
                weight_data = None
                if save_weights and weights is not None:
                    try:
                        # Extract attention for this sample only (localized model - no CLS token)
                        last_layer_attention = weights[-1][i].detach()  # Detach to allow garbage collection
                        
                        if last_layer_attention is not None:
                            avg_attention = last_layer_attention.mean(dim=0)
                            
                            if not torch.isnan(avg_attention).any():
                                # Move to CPU immediately to free GPU memory
                                # Localized attention: average over heads, then over patches
                                patch_attention = avg_attention.mean(dim=0).cpu()
                                
                                attention_map = patch_attention.reshape(grid_h, grid_w)
                                weight_data = attention_map.numpy()
                                
                                # Clear GPU tensors
                                del last_layer_attention, avg_attention, patch_attention
                                
                                # Save weight if needed
                                if config_data and 'weight_path' in config_data and global_idx < len(times):
                                    weight_dir = Path(config_data['weight_path'])
                                    weight_dir.mkdir(parents=True, exist_ok=True)
                                    weight_file = os.path.join(config_data['weight_path'], f"{times[global_idx]}.npy")
                                    np.save(weight_file, weight_data)
                    except Exception as e:
                        weight_data = np.zeros((grid_h, grid_w))
                
                # Process flux contributions
                flux_data = None
                if save_flux and flux_contributions is not None:
                    try:
                        flux_contrib = flux_contributions[i].detach().cpu()
                        flux_contrib_map = flux_contrib.reshape(grid_h, grid_w)
                        flux_data = flux_contrib_map.numpy()
                        
                        # Save flux if needed
                        if config_data and 'flux_path' in config_data and global_idx < len(times):
                            flux_dir = Path(config_data['flux_path'])
                            flux_dir.mkdir(parents=True, exist_ok=True)
                            flux_file = os.path.join(config_data['flux_path'], f"{times[global_idx]}.npy")
                            np.save(flux_file, flux_data)
                        
                        del flux_contrib, flux_contrib_map
                    except Exception:
                        flux_data = None
                
                # Yield result immediately - don't accumulate in memory
                yield (predictions[i].cpu().numpy(), sxr[i].cpu().numpy(), weight_data, flux_data, global_idx)
            
            # Clear all batch-level tensors from GPU memory immediately
            # CRITICAL: Delete weights first as they're the biggest memory consumer
            if weights is not None:
                # Delete each layer's attention weights explicitly to free memory
                for layer_idx, layer_weights in enumerate(weights):
                    if isinstance(layer_weights, torch.Tensor):
                        del layer_weights
                    elif isinstance(layer_weights, (list, tuple)):
                        for w in layer_weights:
                            if isinstance(w, torch.Tensor):
                                del w
                del weights  # Delete the list itself
                weights = None  # Ensure reference is cleared
            
            del predictions, aia_imgs
            if flux_contributions is not None:
                del flux_contributions
                flux_contributions = None
            
            # Only clear cache occasionally — doing it every batch stalls the GPU pipeline
            #if batch_idx % 2 == 0:
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    finally:
        # Explicitly shut down DataLoader workers so they don't linger into the
        # next condition and exhaust shared memory / semaphore limits.
        del loader
        gc.collect()


def load_model_from_config(config_data):
    """
    Load the model from checkpoint based on configuration data.

    Parameters
    ----------
    config_data : dict
        Configuration dictionary from YAML file.

    Returns
    -------
    torch.nn.Module
        Loaded model ready for inference.
    """
    checkpoint_path = config_data['data']['checkpoint_path']

    print("Loading ViTLocal model...")

    # Use GPU(s) if available
    if torch.cuda.is_available():
        load_device = torch.device('cuda:0')
        print(f"Using GPU(s) for inference")
    else:
        load_device = torch.device('cpu')
        print("Using CPU for inference")

    if ".ckpt" in checkpoint_path:
        # Lightning checkpoint format
        model = ViTLocal.load_from_checkpoint(checkpoint_path, map_location=load_device, weights_only=False)
    else:
        state = torch.load(checkpoint_path, map_location=load_device, weights_only=False)
        model = state['model']
        model = model.to(load_device)

    model.eval()

    raw_multi_gpu = config_data.get('multi_gpu', False) if config_data else False
    use_multi_gpu = (str(raw_multi_gpu).lower() == 'true'
                     and torch.cuda.is_available()
                     and torch.cuda.device_count() > 1)
    if use_multi_gpu:
        n_gpus = torch.cuda.device_count()
        model = torch.nn.DataParallel(model)
        print(f"Using DataParallel across {n_gpus} GPUs — ensure batch_size >= {n_gpus}")

    return model


def main():
    """
    Main function to execute solar flare model inference pipeline.

    Steps
    -----
    1. Parse YAML configuration and resolve ${variable} placeholders.
    2. Load pretrained model and dataset.
    3. Run batched inference and optionally save attention/flux maps.
    4. Save predictions and ground truth results to CSV.
    """
    def resolve_config_variables(config_dict):
        """Recursively resolve ${variable} references within config."""
        variables = {}
        for key, value in config_dict.items():
            if isinstance(value, str) and not value.startswith('${'):
                variables[key] = value

        def substitute_value(value, variables):
            if isinstance(value, str):
                pattern = r'\$\{([^}]+)\}'
                for match in re.finditer(pattern, value):
                    var_name = match.group(1)
                    if var_name in variables:
                        value = value.replace(f'${{{var_name}}}', variables[var_name])
            return value

        def recursive_substitute(obj, variables):
            if isinstance(obj, dict):
                return {k: recursive_substitute(v, variables) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [recursive_substitute(item, variables) for item in obj]
            else:
                return substitute_value(obj, variables)

        return recursive_substitute(config_dict, variables)

    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='inference_config.yaml', required=True,
                        help='Path to the inference configuration YAML file.')
    args = parser.parse_args()

    with open(args.config, 'r') as stream:
        config_data = yaml.load(stream, Loader=yaml.SafeLoader)

    config_data: dict = resolve_config_variables(config_data)

    model_params = config_data.get('model_params', {})
    input_size = model_params.get('input_size', 512)
    patch_size = model_params.get('patch_size', 16)
    batch_size = model_params.get('batch_size', 10)
    no_weights = model_params.get('no_weights', False)
    no_flux = model_params.get('no_flux', False)
    
    print(f"Using parameters from config:")
    print(f"  Input size: {input_size}\n  Patch size: {patch_size}\n  Batch size: {batch_size}\n  Skip weights: {no_weights}\n  Skip flux: {no_flux}")

    model = load_model_from_config(config_data)

    save_weights = not no_weights
    if no_weights:
        print("Skipping attention weight saving (no_weights=true).")
        print("  Note: This saves ~3GB per batch by not computing attention weights.")
    else:
        print("Will save attention weights during inference.")
        print("\n Memory note:")
        print("   - Attention weights from all layers use significant GPU memory")
        print("   - For ViT with 8 layers, 8 heads, 4096 patches: ~3GB+ per batch with attention!")
        print("   - If you get OOM errors, set no_weights=true to skip attention saving\n")

    save_flux = config_data and 'flux_path' in config_data and not no_flux
    if no_flux:
        print("Skipping flux contribution saving (no_flux=true).")
    elif config_data and 'flux_path' in config_data:
        print("Will save flux contributions during inference.")
    else:
        print("No flux path specified.")

    torch.backends.cudnn.benchmark = True

    print("Loading dataset...")
    # Check if running in prediction-only mode
    prediction_only = config_data.get('prediction_only', 'false').lower() == 'true'

    dataset = AIAGOESDataset(aia_dir=config_data['data']['aia_dir'],
                              sxr_dir=config_data['data'].get('sxr_dir') if not prediction_only else None,
                              wavelengths=config_data['wavelengths'],
                              only_prediction=prediction_only)

    times = dataset.samples
    
    # Load SXR normalization only if not in prediction-only mode
    if prediction_only:
        print("Running in prediction-only mode - skipping SXR normalization loading")
        sxr_norm = None
    else:
        try:
            sxr_norm = np.load(config_data['data']['sxr_norm_path'])
            print("SXR normalization loaded successfully")
        except FileNotFoundError:
            print(f"Warning: SXR normalization file not found: {config_data['data']['sxr_norm_path']}")
            print("Continuing without normalization...")
            sxr_norm = None

    timestamp, predictions, ground = [], [], []
    total_samples = len(times)
    print(f"Processing {total_samples} samples with batch size {batch_size}...")

    print("Running inference...")
    
    # Create progress bar
    pbar = tqdm(
        evaluate_model_on_dataset(
            model, dataset, batch_size, times, config_data,
            save_weights, input_size, patch_size, save_flux
        ),
        total=total_samples,
        desc="Inference",
        unit="sample",
        ncols=100
    )
    
    for prediction, sxr, weight, flux_data, idx in pbar:
        # ViTLocal models already return unnormalized predictions, so no need to unnormalize
        pred = prediction

        predictions.append(pred.item() if hasattr(pred, 'item') else float(pred))
        
        # Only collect ground truth if not in prediction-only mode
        if not prediction_only:
            ground.append(sxr.item() if hasattr(sxr, 'item') else float(sxr))
        else:
            ground.append(0.0)  # Dummy value for prediction-only mode
        
        timestamp.append(str(times[idx]))
        
        # Update progress bar
        pbar.set_postfix({'sample': idx + 1, 'total': total_samples})

    output_df = pd.DataFrame({'timestamp': timestamp, 'predictions': predictions, 'groundtruth': ground})
    output_dir = Path(config_data['output_path']).parent
    output_dir.mkdir(parents=True, exist_ok=True)
    output_df.to_csv(config_data['output_path'], index=False)
    
    if prediction_only:
        print(f"Predictions saved to {config_data['output_path']} (prediction-only mode)")
    else:
        print(f"Predictions saved to {config_data['output_path']}")


if __name__ == '__main__':
    main()