"""
utils.py - Utility Functions for LAI Prediction Pipeline

This module contains helper functions used across different stages of the LAI prediction pipeline, including:

- Plotting utilities (e.g., histograms, residuals, spatial maps, training curves)
- Data loading from Zarr format
- Data scaling and inverse transformations
- Model evaluation and logging
- JSON serialization and metadata handling
- General-purpose functions for reshaping, logging system info, and statistics

These utilities are shared across preprocessing, training, and evaluation scripts.
"""

# ===========================
# Standard Imports
# ===========================
import os
import sys
import json
import logging
import numpy as np
import matplotlib.pyplot as plt
import zarr
import torch
import platform
import psutil

# ===========================
# Plotting Utilities
# ===========================

def plot_data(data, title, filename, vmin=None, vmax=None):
    """
    Visualizes 2D spatial data (e.g., LAI maps) using imshow and saves the plot.

    Args:
        data (np.ndarray): 2D spatial data array to visualize.
        title (str): Title of the plot.
        filename (str): File path to save the output figure.
        vmin (float, optional): Minimum color scale value.
        vmax (float, optional): Maximum color scale value.
    """
    plt.figure(figsize=(10, 6))
    im = plt.imshow(np.flipud(data), cmap='viridis', origin='upper', vmin=vmin, vmax=vmax)
    plt.colorbar(im, label='Value')
    plt.title(title)
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.savefig(filename)
    plt.close()

def plot_histogram(data1, data2, title, filename, xlabel="LAI value", ylabel="Frequency"):
    """
    Plots and compares histograms of two datasets (e.g., train and validation LAI distributions).

    Args:
        data1 (np.ndarray): First dataset (e.g., training).
        data2 (np.ndarray): Second dataset (e.g., validation).
        title (str): Plot title.
        filename (str): File path to save the plot.
        xlabel (str): Label for the x-axis.
        ylabel (str): Label for the y-axis.
    """

    data1 = data1.reshape(-1)
    data2 = data2.reshape(-1)
    plt.figure(figsize=(10, 6))
    plt.hist(data1, bins=50, alpha=0.7, label='Training Data')
    plt.hist(data2, bins=50, alpha=0.7, label='Validation Data')
    plt.legend()
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(filename)
    plt.close()
    logging.info(f"{title} plot to: {filename}") 

def plot_residuals(y_true, residuals, title, filename, xlabel="True Values", ylabel="Residuals"):
    """
    Plots residuals versus ground-truth values.

    Args:
        y_true (np.ndarray): True LAI values.
        residuals (np.ndarray): Prediction residuals (e.g., y_true - y_pred).
        title (str): Plot title.
        filename (str): Output file path.
        xlabel (str): X-axis label.
        ylabel (str): Y-axis label.
    """
    y_true = y_true.reshape(-1)
    residuals = residuals.reshape(-1)
    plt.figure(figsize=(10, 6))
    plt.scatter(y_true, residuals, alpha=0.5)
    plt.axhline(0, color='red', linestyle='--')
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(filename)
    plt.close()

def plot_training_curves(train_losses, val_losses, filename="training_curve.png"):
    """
    Plots training and validation loss curves across epochs.

    Args:
        train_losses (list): Training loss per epoch.
        val_losses (list): Validation loss per epoch.
        filename (str): Path to save the loss curve plot.
    """
    epochs = range(1, len(train_losses) + 1)
    plt.figure(figsize=(10, 6))
    plt.plot(epochs, train_losses, label='Training Loss')
    plt.plot(epochs, val_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training vs Validation Loss')
    plt.legend()
    plt.grid(True)
    plt.savefig(filename)
    plt.close()

def plot_postprocessing_outputs(
    y_train, y_test, y_pred, y_pred_corrected,
    save_data_path, region, timestamp, model_name,
    lat_dim, lon_dim, lower_bound=0, upper_bound=7
):
    """
    Generates a suite of post-processing plots including histograms, residuals, and spatial predictions.

    Args:
        y_train (np.ndarray): Training target values.
        y_test (np.ndarray): Test target values.
        y_pred (np.ndarray): Model predictions.
        y_pred_corrected (np.ndarray): Bias-corrected predictions.
        save_data_path (str): Directory to save plots.
        region (str): Region name.
        timestamp (str): Timestamp string for output file naming.
        model_name (str): Name of the model used.
        lat_dim (int): Latitude dimension of the spatial grid.
        lon_dim (int): Longitude dimension of the spatial grid.
        lower_bound (float): Minimum value for spatial plot color scale.
        upper_bound (float): Maximum value for spatial plot color scale.

    Returns:
        dict: Mapping of plot types to saved file paths.
    """

    # File paths
    target_dist_path = os.path.join(save_data_path, f'target_distribution_{region}_{timestamp}.png')
    residuals_plot_path = os.path.join(save_data_path, f'residuals_plot_{region}_{timestamp}.png')
    corrected_histogram_path = os.path.join(save_data_path, f'histogram_corrected_predictions_{region}_{timestamp}.png')
    corrected_residuals_plot_path = os.path.join(save_data_path, f'residuals_corrected_predictions_{region}_{timestamp}.png')

    # 1. Target distribution
    plot_histogram(y_train, y_test, "Target Distribution", target_dist_path)
    logging.info(f"Target distribution plot saved at: {target_dist_path}")

    # 2. Raw residuals
    plot_residuals(y_test, y_test - y_pred, "Residuals Plot", residuals_plot_path)
    logging.info(f"Residuals plot saved at: {residuals_plot_path}")

    # 3. Histogram - Corrected
    plot_histogram(y_pred_corrected, y_test, "Histogram - Corrected Predictions vs True Values", corrected_histogram_path)
    logging.info(f"Corrected predictions histogram saved at: {corrected_histogram_path}")

    # 4. Residuals - Corrected
    corrected_residuals = y_test - y_pred_corrected
    plot_residuals(y_test, corrected_residuals, "Residuals Plot - Corrected Predictions", corrected_residuals_plot_path)
    logging.info(f"Residuals plot for corrected predictions saved at: {corrected_residuals_plot_path}")

    # 5. Spatial plots for each month
    for month in range(y_pred.shape[1]):
        predicted_data = reshape_with_padding(y_pred[:, month], (lat_dim, lon_dim))
        true_data = reshape_with_padding(y_test[:, month], (lat_dim, lon_dim))
        corrected_data = reshape_with_padding(y_pred_corrected[:, month], (lat_dim, lon_dim))

        pred_path = os.path.join(save_data_path, f'prediction_plot_{model_name}_month{month+1}_{region}_{timestamp}.png')
        true_path = os.path.join(save_data_path, f'true_plot_{model_name}_month{month+1}_{region}_{timestamp}.png')
        corrected_path = os.path.join(save_data_path, f'prediction_plot_corrected_{model_name}_month{month+1}_{region}_{timestamp}.png')

        plot_data(predicted_data, f"Predicted LAI Month {month+1}", pred_path, vmin=lower_bound, vmax=upper_bound)
        plot_data(true_data, f"True LAI Month {month+1}", true_path, vmin=lower_bound, vmax=upper_bound)
        plot_data(corrected_data, f"Corrected LAI Month {month+1}", corrected_path, vmin=lower_bound, vmax=upper_bound)

        logging.info(f"Saved plots for Month {month+1} — Pred: {pred_path}, True: {true_path}, Corrected: {corrected_path}")

    return {
        "target_distribution": target_dist_path,
        "residuals_raw": residuals_plot_path,
        "residuals_corrected": corrected_residuals_plot_path,
        "histogram_corrected_predictions": corrected_histogram_path
    }

# ===========================
# Data Loading
# ===========================

def load_data_for_years(years, region, data_path):
    """
    Loads feature and LAI data across specified years and region.

    Args:
        years (list of int): Years to load data for.
        region (str): Region identifier.
        data_path (str): Base path to the Zarr data files.

    Returns:
        tuple: (features_list, LAI_list) where each is a list of np.ndarray.
    """
    features_list = []
    LAI_list = []

    for year in years:
        features_path = os.path.join(data_path, f'features_training_flat_{year}_{region}.zarr')
        lai_path = os.path.join(data_path, f'LAI_{year}_{region}.zarr')

        if not os.path.exists(features_path):
            logging.error(f"Feature file not found: {features_path}")
            sys.exit(1)
        if not os.path.exists(lai_path):
            logging.error(f"LAI file not found: {lai_path}")
            sys.exit(1)

        try:
            features = zarr.load(features_path)
            lai = zarr.load(lai_path)
        except Exception as e:
            logging.error(f"Error loading data for year {year}: {e}")
            sys.exit(1)

        if features.size == 0 or lai.size == 0:
            logging.error(f"Empty data found for year {year}. Exiting.")
            sys.exit(1)

        logging.info(f"Year {year} - features shape: {features.shape}, LAI shape: {lai.shape}")
        features_list.append(features)
        LAI_list.append(lai)

    logging.info(f"Data loaded as a sequence. Number of years: {len(features_list)}")
    return features_list, LAI_list

# ===========================
# Scaling Functions
# ===========================

def scale_data(data, lower_bound, upper_bound):
    """
    Scales data to [0, 1] using specified bounds.

    Args:
        data (np.ndarray): Raw data.
        lower_bound (float): Minimum expected value.
        upper_bound (float): Maximum expected value.

    Returns:
        np.ndarray: Scaled data.
    """

    return (data - lower_bound) / (upper_bound - lower_bound)

def inverse_scale_data(data, lower_bound, upper_bound):
    """
    Converts scaled data back to original range.

    Args:
        data (np.ndarray): Scaled data in [0, 1].
        lower_bound (float): Original minimum value.
        upper_bound (float): Original maximum value.

    Returns:
        np.ndarray: Rescaled data.
    """

    return data * (upper_bound - lower_bound) + lower_bound

# ===========================
# Model Evaluation 
# ===========================

def evaluate(model, loader, device, criterion):
    """
    Evaluates a PyTorch model over a DataLoader and returns average loss.

    Args:
        model (torch.nn.Module): Trained model to evaluate.
        loader (DataLoader): Evaluation DataLoader.
        device (torch.device): Device for computation (CPU/GPU).
        criterion (loss function): PyTorch loss function.

    Returns:
        float: Average loss over dataset.
    """
    model.eval()  # Set model to evaluation mode
    total_loss = 0
    with torch.no_grad():  # No need to compute gradients during evaluation
        for lai_next_year, yearly_vars, lai_target in loader:
            lai_next_year = lai_next_year.to(device)
            yearly_vars = yearly_vars.to(device)
            lai_target = lai_target.to(device)

            output = model(lai_next_year, yearly_vars)
            loss = criterion(output, lai_target)  # Compute loss
            total_loss += loss.item()  # Accumulate loss

    return total_loss / len(loader)  # Return the average loss

# ===========================
# JSON & Serialization Helpers
# ===========================

def load_metadata(save_data_path, start_year, end_year, region):
    """
    Loads model training metadata from JSON.

    Args:
        save_data_path (str): Path to directory containing metadata.
        start_year (int): Start year of training and test data.
        end_year (int): End year of training and test data.
        region (str): Region identifier.

    Returns:
        dict: Loaded metadata.
    """
    metadata_path = os.path.join(save_data_path, f'metadata_training_{start_year}_{end_year}_{region}.json')

    if not os.path.exists(metadata_path):
        raise FileNotFoundError(f"Metadata file not found: {metadata_path}")

    with open(metadata_path, 'r') as f:
        metadata = json.load(f)

    logging.info(f"Loaded metadata: {metadata}")

    return metadata

def sanitize_params(params):
    """
    Sanitizes a parameter dictionary for JSON serialization.

    Args:
        params (dict): Dictionary of parameters.

    Returns:
        dict: JSON-serializable dictionary.
    """
    sanitized = {}
    for key, value in params.items():
        try:
            json.dumps(value)
            sanitized[key] = value
        except TypeError:
            sanitized[key] = str(value)
    return sanitized


def convert_numpy_types(obj):
    """
    Recursively converts NumPy types to native Python types for JSON serialization.

    Args:
        obj: Any object potentially containing NumPy types.

    Returns:
        obj: JSON-safe object.
    """
    if isinstance(obj, (np.integer, int)):
        return int(obj)
    elif isinstance(obj, (np.floating, float)):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {str(k): convert_numpy_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_numpy_types(v) for v in obj]
    elif isinstance(obj, tuple):
        return tuple(convert_numpy_types(v) for v in obj)
    else:
        return obj

# ===========================
# General Utilities
# ===========================

def configure_logging(start_year, end_year, region, timestamp):
    """
    Sets up logging to a file and stdout for a training session.

    Args:
        start_year (int): Start year of training and test period.
        end_year (int): End year of training and test period.
        region (str): Region name.
        timestamp (str): Timestamp string for filename uniqueness.

    Returns:
        str: Path to the log file.
    """
    current_path = os.getcwd()
    logs_path = os.path.join(current_path, "Logs")
    os.makedirs(logs_path, exist_ok=True)

    log_file_name = f"train_log_{start_year}_{end_year}_{region}_{timestamp}.log"
    log_file_path = os.path.join(logs_path, log_file_name)

    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file_path),
            logging.StreamHandler()
        ]
    )

    return log_file_path

def reshape_with_padding(arr, target_shape):
    """
    Reshapes a 1D array to 2D shape, padding with zeros if needed.

    Args:
        arr (np.ndarray): Input 1D array.
        target_shape (tuple): Desired 2D shape (rows, columns).

    Returns:
        np.ndarray: Reshaped 2D array.
    """
    target_size = target_shape[0] * target_shape[1]
    current_size = arr.size
    if current_size < target_size:
        padded_arr = np.pad(arr, (0, target_size - current_size), mode='constant')
        return padded_arr.reshape(target_shape, order='C')
    return arr[:target_size].reshape(target_shape, order='C')

def log_system_info():
    """
    Logs system information including CPU, RAM, GPU (if available), OS, and Python environment.
    """
    # Device type
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logging.info(f"Using device: {device}")

    # CPU info
    num_cores = os.cpu_count()
    cpu_freq = psutil.cpu_freq()
    logging.info(f"CPU: {platform.processor()} with {num_cores} cores")
    logging.info(f"CPU Frequency: {cpu_freq.current:.2f} MHz (min: {cpu_freq.min:.2f}, max: {cpu_freq.max:.2f})")

    # RAM info
    mem = psutil.virtual_memory()
    logging.info(f"Total RAM: {mem.total / (1024 ** 3):.2f} GB")
    logging.info(f"Available RAM: {mem.available / (1024 ** 3):.2f} GB")

    # GPU info
    if device.type == "cuda":
        gpu_name = torch.cuda.get_device_name(0)
        gpu_props = torch.cuda.get_device_properties(0)
        gpu_memory = gpu_props.total_memory / (1024 ** 3)
        logging.info(f"GPU: {gpu_name}")
        logging.info(f"GPU Memory: {gpu_memory:.2f} GB")
        logging.info(f"CUDA Capability: {gpu_props.major}.{gpu_props.minor}")
    else:
        logging.info("CUDA not available — running on CPU.")

    # Python & OS info
    logging.info(f"Python Version: {platform.python_version()}")
    logging.info(f"OS: {platform.system()} {platform.release()} ({platform.version()})")

def log_tensor_stats(y_train, y_val, y_test, X_train_monthly, X_val_monthly, X_test_monthly, message_prefix=""):
    """
    Logs statistical summaries for training/validation/test tensors.

    Args:
        y_train, y_val, y_test (array-like): Target arrays.
        X_train_monthly, X_val_monthly, X_test_monthly (array-like): Monthly feature arrays.
        message_prefix (str): Optional string to prefix all log lines.
    """
    def stats(tensor):
        return (
            tensor.min().item(), 
            tensor.max().item(), 
            tensor.mean().item(), 
            tensor.median().item(), 
            tensor.var().item()
        )

    y_train_min, y_train_max, y_train_mean, y_train_median, y_train_var = stats(torch.tensor(y_train) if not torch.is_tensor(y_train) else y_train.clone().detach())
    y_val_min, y_val_max, y_val_mean, y_val_median, y_val_var = stats(torch.tensor(y_val) if not torch.is_tensor(y_val) else y_val.clone().detach())
    y_test_min, y_test_max, y_test_mean, y_test_median, y_test_var = stats(torch.tensor(y_test) if not torch.is_tensor(y_test) else y_test.clone().detach())

    x_train_min, x_train_max, x_train_mean, x_train_median, x_train_var = stats(torch.tensor(X_train_monthly) if not torch.is_tensor(X_train_monthly) else X_train_monthly.clone().detach())
    x_val_min, x_val_max, x_val_mean, x_val_median, x_val_var = stats(torch.tensor(X_val_monthly) if not torch.is_tensor(X_val_monthly) else X_val_monthly.clone().detach())
    x_test_min, x_test_max, x_test_mean, x_test_median, x_test_var = stats(torch.tensor(X_test_monthly) if not torch.is_tensor(X_test_monthly) else X_test_monthly.clone().detach())

    logging.info(f"{message_prefix}")
    logging.info(f"y_train — Min: {y_train_min:.4f}, Max: {y_train_max:.4f}, Mean: {y_train_mean:.4f}, Median: {y_train_median:.4f}, Var: {y_train_var:.4f}")
    logging.info(f"y_val   — Min: {y_val_min:.4f}, Max: {y_val_max:.4f}, Mean: {y_val_mean:.4f}, Median: {y_val_median:.4f}, Var: {y_val_var:.4f}")
    logging.info(f"y_test  — Min: {y_test_min:.4f}, Max: {y_test_max:.4f}, Mean: {y_test_mean:.4f}, Median: {y_test_median:.4f}, Var: {y_test_var:.4f}")
    logging.info(f"X_train_monthly — Min: {x_train_min:.4f}, Max: {x_train_max:.4f}, Mean: {x_train_mean:.4f}, Median: {x_train_median:.4f}, Var: {x_train_var:.4f}")
    logging.info(f"X_val_monthly   — Min: {x_val_min:.4f}, Max: {x_val_max:.4f}, Mean: {x_val_mean:.4f}, Median: {x_val_median:.4f}, Var: {x_val_var:.4f}")
    logging.info(f"X_test_monthly  — Min: {x_test_min:.4f}, Max: {x_test_max:.4f}, Mean: {x_test_mean:.4f}, Median: {x_test_median:.4f}, Var: {x_test_var:.4f}")

def log_stats(name, data):
    """
    Logs summary statistics for a given NumPy array.

    Args:
        name (str): Label for the dataset.
        data (np.ndarray): Data to summarize.
    """
    logging.info(f"{name} — Min: {np.min(data):.4f}, Max: {np.max(data):.4f}, Mean: {np.mean(data):.4f}, Median: {np.median(data):.4f}, Var: {np.var(data):.4f}")

def build_experiment_log(
    args, model_name, model_path, final_epoch, best_epoch, training_duration, training_curve_path,
    num_model_params, num_train_years, val_years_count, num_test_years, num_locations,
    num_monthly_variables, num_yearly_variables, metrics, bias, train_losses, val_losses
):
    """
    Compiles all experiment configuration and results into a structured dictionary.

    Args:
        args (Namespace): Training arguments.
        model_name (str): Name of the model.
        model_path (str): Path to saved model.
        final_epoch (int): Final training epoch.
        best_epoch (int): Epoch with best validation performance.
        training_duration (float): Total training time in seconds.
        training_curve_path (str): Path to training curve image.
        num_model_params (int): Number of model parameters.
        num_train_years (int): Number of training years.
        val_years_count (int): Number of validation years.
        num_test_years (int): Number of testing years.
        num_locations (int): Spatial locations.
        num_monthly_variables (int): Monthly input features.
        num_yearly_variables (int): Yearly input features.
        metrics (dict): Dictionary of raw, clipped, and corrected evaluation metrics.
        bias (float): Prediction bias value.
        train_losses (list): List of training loss values.
        val_losses (list): List of validation loss values.

    Returns:
        dict: Structured log of the experiment.
    """

    return {
        "timestamp": args.model_timestamp,
        "model_name": model_name,
        "model_path": model_path,
        "region": args.region,
        "start_year": args.start_year,
        "end_year": args.end_year,
        "train_years": "-".join(map(str, args.train_years)),
        "test_years": "-".join(map(str, args.test_years)),
        "model_type": "rnn",
        "cell_type": args.rnn_cell_type,
        "attention": args.use_attention,
        "dropout": args.dropout,
        "weight_decay": args.weight_decay,
        "hidden_size": args.hidden_size,
        "num_epochs_requested": args.num_epochs,
        "final_epoch": final_epoch,
        "best_epoch": best_epoch,
        "early_stopping": args.early_stopping,
        "early_stopping_patience": args.early_stopping_patience,
        "teacher_forcing_ratio": args.teacher_forcing_ratio,
        "batch_size": args.batch_size,
        "num_train_years": num_train_years,
        "num_val_years": val_years_count,
        "num_test_years": num_test_years,
        "num_locations": num_locations,
        "num_monthly_features": num_monthly_variables,
        "num_yearly_features": num_yearly_variables,
        "device": str(torch.device("cuda" if torch.cuda.is_available() else "cpu")),
        "training_duration_sec": round(training_duration, 2),

        # Software/hardware
        "python_version": platform.python_version(),
        "torch_version": torch.__version__,
        "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "None",
        "gpu_memory_gb": round(torch.cuda.get_device_properties(0).total_memory / (1024 ** 3), 2) if torch.cuda.is_available() else 0,

        # Preprocessing
        "imputation_method": args.imputation,
        "ground_truth_log_transform": args.ground_truth_log_transform,
        "ground_truth_scale": args.ground_truth_scale,
        "scaling_range": "[0,7]",
        "water_mask_applied": args.water_filter,

        # Resources
        "num_workers": 32,
        "cpu_cores": os.cpu_count(),
        "ram_gb": round(psutil.virtual_memory().total / (1024 ** 3), 2),

        # Evaluation metrics
        "mse": round(metrics["raw"][0], 5),
        "rmse": round(metrics["raw"][1], 5),
        "mae": round(metrics["raw"][2], 5),
        "r2": round(metrics["raw"][3], 5),

        "mse_clipped": round(metrics["clipped"][0], 5),
        "rmse_clipped": round(metrics["clipped"][1], 5),
        "mae_clipped": round(metrics["clipped"][2], 5),
        "r2_clipped": round(metrics["clipped"][3], 5),

        "mse_corrected": round(metrics["corrected"][0], 5),
        "rmse_corrected": round(metrics["corrected"][1], 5),
        "mae_corrected": round(metrics["corrected"][2], 5),
        "r2_corrected": round(metrics["corrected"][3], 5),

        "bias_correction": round(abs(bias), 5),

        # Training info
        "train_losses": ",".join([f"{v:.5f}" for v in train_losses]),
        "val_losses": ",".join([f"{v:.5f}" for v in val_losses]),
        "num_model_parameters": num_model_params,
        "training_curve_path": training_curve_path
    }
