"""
postprocessing.py - Postprocessing Utilities for LAI Prediction Pipeline

This module handles model output processing after inference, including:

- Applying inverse transformations (log and scaling)
- Applying persistent water masks to exclude non-land areas
- Computing bias and performing bias correction
- Clipping predictions to valid LAI range
- Generating post-inference plots (residuals, histograms)
- Calculating evaluation metrics (MSE, RMSE, MAE, R²)

These utilities are typically used after model evaluation on the test set.
"""

import os
import torch
import numpy as np
import zarr
import logging
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from scipy.stats import gaussian_kde
from utils import inverse_scale_data, log_stats, plot_postprocessing_outputs
from preprocessing import load_combined_water_mask 

def apply_water_mask(y_pred, y_pred_corrected, test_years, region, save_data_path, logger=None):
    """
    Applies a spatial mask to exclude water-covered areas from predictions.

    Args:
        y_pred (np.ndarray): Raw model predictions (shape: N x 12).
        y_pred_corrected (np.ndarray): Bias-corrected predictions (shape: N x 12).
        test_years (list of int): Years corresponding to the test data.
        region (str): Region identifier.
        save_data_path (str): Path to directory containing water mask files.
        logger (logging.Logger, optional): Logger for tracking.

    Returns:
        tuple: (masked_y_pred, masked_y_pred_corrected, list of water masks).
    """
    logger = logger or logging.getLogger(__name__)

    water_masks = []
    for yr in test_years:
        mask_path = os.path.join(save_data_path, f'water_mask_{region}_{yr}.zarr')
        if os.path.exists(mask_path):
            water_mask = zarr.load(mask_path)
            logger.info(f"Water mask shape for year {yr}: {water_mask.shape}")
            water_masks.append(water_mask.flatten(order='C'))
        else:
            logger.warning(f"Water mask for year {yr} not found.")

    if not water_masks:
        logger.warning("No valid water masks found. Predictions were not masked.")
        return y_pred, y_pred_corrected, []

    combined_mask = np.concatenate(water_masks, axis=0)
    logger.info(f"Combined water mask shape: {combined_mask.shape}")

    # Repeat the mask across months to match the shape of the prediction array (N x 12)
    combined_mask = np.tile(combined_mask[:, np.newaxis], (1, y_pred.shape[1]))
    logger.info(f"Tiled water mask shape: {combined_mask.shape}")

    if combined_mask.shape != y_pred.shape:
        raise ValueError("Mismatch between water mask and predictions dimensions.")

    # Apply mask: zero out predictions over water
    y_pred_masked = y_pred * combined_mask
    y_pred_corrected_masked = y_pred_corrected * combined_mask
    logger.info("Applied water mask to predictions.")

    return y_pred_masked, y_pred_corrected_masked, water_masks

def plot_monthly_histograms(y_true, y_pred, region, timestamp, save_path, bins=50):
    """
    Plots overlaid histograms for each month comparing ground truth and predictions.

    Args:
        y_true (np.ndarray): Ground truth LAI values (N x 12).
        y_pred (np.ndarray): Model predictions (N x 12).
        region (str): Region name for labeling outputs.
        timestamp (str): Unique identifier for filenames.
        save_path (str): Directory to save plots.
        bins (int): Number of histogram bins.
    """
    os.makedirs(save_path, exist_ok=True)

    # Overlaid histograms to visualize distribution shifts per month
    for month in range(12):
        gt = y_true[:, month]
        pred = y_pred[:, month]

        plt.figure(figsize=(10, 6))
        plt.hist(gt, bins=bins, alpha=0.5, label='Ground Truth')
        plt.hist(pred, bins=bins, alpha=0.5, label='Prediction')
        plt.title(f'LAI Distribution - Month {month + 1:02d}')
        plt.xlabel('LAI')
        plt.ylabel('Frequency')
        plt.xlim(0, 7)

        plt.legend(loc='upper left')

        # Stats
        gt_stats = f"GT: μ={gt.mean():.2f}, med={np.median(gt):.2f}, σ={gt.std():.2f}"
        pred_stats = f"Pred: μ={pred.mean():.2f}, med={np.median(pred):.2f}, σ={pred.std():.2f}"
        stats_text = f"{gt_stats}\n{pred_stats}"

        # Top center placement (axes coordinates)
        plt.gca().text(
            0.5, 0.95, stats_text,
            transform=plt.gca().transAxes,
            fontsize=10,
            verticalalignment='top',
            horizontalalignment='center',
            bbox=dict(boxstyle="round,pad=0.4", facecolor='white', alpha=0.85)
        )

        out_path = os.path.join(save_path, f"histogram_{region}_month_{month+1:02d}_{timestamp}.png")
        plt.tight_layout()
        plt.savefig(out_path)
        plt.close()
        
def plot_monthly_residual_curves(y_true, y_pred, region, timestamp, save_path, bandwidth=0.1):
    """
    Plots kernel density estimates (KDE) of monthly residuals (y_true - y_pred).

    Args:
        y_true (np.ndarray): Ground truth LAI values (N x 12).
        y_pred (np.ndarray): Predicted LAI values (N x 12).
        region (str): Region name for filenames.
        timestamp (str): Identifier for output files.
        save_path (str): Directory where plots are saved.
        bandwidth (float): Smoothing bandwidth for KDE.
    """
    os.makedirs(save_path, exist_ok=True)

    # Estimate and plot residual distribution using Gaussian KDE
    for month in range(12):
        residuals = y_true[:, month] - y_pred[:, month]

        # Remove NaNs/Infs
        residuals = residuals[np.isfinite(residuals)]

        # Kernel density estimate
        kde = gaussian_kde(residuals, bw_method=bandwidth)
        x_vals = np.linspace(residuals.min(), residuals.max(), 1000)
        y_vals = kde(x_vals)

        plt.figure(figsize=(10, 6))
        plt.plot(x_vals, y_vals, label=f"Residuals (GT - Pred)", color='crimson')
        plt.axvline(0, color='black', linestyle='--', label='Zero Error')

        plt.title(f"Residual Distribution - Month {month + 1:02d}")
        plt.xlabel("Residual (GT - Prediction)")
        plt.ylabel("Density")
        plt.legend(loc='upper right')
        plt.xlim(-1, 1)

        # Add stats box
        mean = residuals.mean()
        median = np.median(residuals)
        std = residuals.std()
        stats_text = f"μ={mean:.3f}, med={median:.3f}, σ={std:.3f}"

        plt.gca().text(
            0.5, 0.95, stats_text,
            transform=plt.gca().transAxes,
            fontsize=10,
            verticalalignment='top',
            horizontalalignment='center',
            bbox=dict(boxstyle="round,pad=0.4", facecolor='white', alpha=0.85)
        )

        out_path = os.path.join(save_path, f"residual_curve_{region}_month_{month+1:02d}_{timestamp}.png")
        plt.tight_layout()
        plt.savefig(out_path)
        plt.close()
        
def evaluate_and_postprocess(
    model, test_loader, y_train_post, y_test_post,
    scaling_min, scaling_max, ground_truth_scale, ground_truth_log_transform,
    region, timestamp, model_name, save_data_path,
    test_years, lat_dim, lon_dim, lower_bound=0, upper_bound=7,
    post_plots=False, logger=logging
):
    """
    Performs postprocessing on model predictions after evaluation.

    Includes:
    - Forward pass to get predictions
    - Inverse scaling and log transform (if used)
    - Water mask application to exclude water bodies
    - Bias correction and clipping to valid LAI range
    - Logging of prediction statistics
    - Plot generation (optional)
    - Metric computation (MSE, RMSE, MAE, R²)

    Args:
        model (torch.nn.Module): Trained LAI prediction model.
        test_loader (DataLoader): DataLoader for test dataset.
        y_train_post (np.ndarray): Unscaled/logged training targets.
        y_test_post (np.ndarray): Unscaled/logged test targets.
        scaling_min (float): Minimum value used during scaling.
        scaling_max (float): Maximum value used during scaling.
        ground_truth_scale (bool): Whether scaling was applied.
        ground_truth_log_transform (bool): Whether log transform was applied.
        region (str): Region name.
        timestamp (str): Time identifier for output files.
        model_name (str): Model identifier for output labeling.
        save_data_path (str): Path to save plots and outputs.
        test_years (list): Years used in the test set.
        lat_dim (int): Latitude dimension of the spatial grid.
        lon_dim (int): Longitude dimension of the spatial grid.
        lower_bound (float): Minimum valid LAI value.
        upper_bound (float): Maximum valid LAI value.
        post_plots (bool): Whether to generate postprocessing plots.
        logger (logging.Logger): Logger for output tracking.

    Returns:
        dict: A dictionary containing:
            - y_pred: Full grid predictions after clipping
            - y_pred_land: Land-only raw predictions
            - y_pred_land_clipped: Land-only clipped predictions
            - y_pred_land_corrected: Land-only bias-corrected and clipped predictions
            - metrics: Dictionary of evaluation metrics (raw, clipped, corrected, bias)
    """

    device = next(model.parameters()).device

    logger.info("Starting Evaluation step on test set")
    logger.info(f"Model architecture:\n{model}")

    # Make predictions on the test set
    model.eval()
    y_pred = []
    
    # Forward pass without gradient computation
    with torch.no_grad():
        for lai_next_year, yearly_vars, _ in test_loader:
            lai_next_year = lai_next_year.to(device)
            yearly_vars = yearly_vars.to(device)
            output = model(lai_next_year, yearly_vars)
            y_pred.append(output.cpu().numpy())

    # Concatenate predictions
    y_pred = np.concatenate(y_pred, axis=0)  # Shape: (num_locations * num_test_years, 12)

    # Reverse time dimension (Dec → Jan) to original order
    y_pred = np.flip(y_pred, axis=1)

    # Flatten y_test_post to match y_pred
    y_train = y_train_post.reshape(-1, 12)
    y_test = y_test_post.reshape(-1, 12)

    logger.info(f"Predictions shape: {y_pred.shape}")
    logger.info(f"Ground Truth shape: {y_test.shape}")

    assert y_pred.shape == y_test.shape, "Shape mismatch between predictions and ground truth"

    # Postprocessing
    if ground_truth_scale:
        logger.info("Applying inverse scaling...")
        y_pred = inverse_scale_data(y_pred, scaling_min, scaling_max)
        y_test = inverse_scale_data(y_test, scaling_min, scaling_max)

        logger.info(f"After inverse scaling — y_pred min: {np.nanmin(y_pred):.4f}, max: {np.nanmax(y_pred):.4f}, mean: {np.nanmean(y_pred):.4f}")
        logger.info(f"After inverse scaling — y_test min: {np.nanmin(y_test):.4f}, max: {np.nanmax(y_test):.4f}, mean: {np.nanmean(y_test):.4f}")
        logger.info(f"NaNs in y_pred after inverse scaling: {np.isnan(y_pred).sum()}, Infs: {np.isinf(y_pred).sum()}")
        logger.info(f"NaNs in y_test after inverse scaling: {np.isnan(y_test).sum()}, Infs: {np.isinf(y_test).sum()}")

    if ground_truth_log_transform:
        logger.info("Applying inverse log transform...")
        y_pred = np.expm1(y_pred)
        y_test = np.expm1(y_test)

        logger.info(f"After inverse log — y_pred min: {np.nanmin(y_pred):.4f}, max: {np.nanmax(y_pred):.4f}")
        logger.info(f"After inverse log — y_test min: {np.nanmin(y_test):.4f}, max: {np.nanmax(y_test):.4f}")
        logger.info(f"NaNs in y_pred after inverse log: {np.isnan(y_pred).sum()}, Infs: {np.isinf(y_pred).sum()}")
        logger.info(f"NaNs in y_test after inverse log: {np.isnan(y_test).sum()}, Infs: {np.isinf(y_test).sum()}")

    # Get water mask for test set (same dimensions as full test region)
    water_mask_test = load_combined_water_mask(test_years, region, save_data_path, logger)

    # Flatten the 2D spatial mask to apply to flattened prediction/gt arrays
    water_mask_flat = water_mask_test.flatten()

    # Apply mask to y_pred and y_test
    y_pred_land = y_pred[water_mask_flat]
    y_test_land = y_test[water_mask_flat]

    logger.info(f"Predictions shape (land-only): {y_pred_land.shape}")
    logger.info(f"Ground Truth shape (land-only): {y_test_land.shape}")

    # Log number of predictions outside valid LAI range (0-7)
    num_below_0 = np.sum(y_pred_land < lower_bound)
    num_above_7 = np.sum(y_pred_land > upper_bound)
    logger.info(f"Predictions outside valid LAI range: below 0 = {num_below_0}, above 7 = {num_above_7}")

    # Compute bias based on land-only residuals
    bias = np.mean(y_test_land - y_pred_land)
    logger.info(f"Bias Correction (land-only): {abs(bias):.6f}")

    # Apply clipping alone and bias correction and clipping to land-only and full-grid predictions
    y_pred_land_clipped = np.clip(y_pred_land, lower_bound, upper_bound)
    y_pred_corrected = np.clip(y_pred + abs(bias), lower_bound, upper_bound)
    y_pred_land_corrected = np.clip(y_pred_land + abs(bias), lower_bound, upper_bound)

    # Apply clipping on y_pred (full grid) and zero-out water locations to have final product (assuming no bias correction)
    y_pred = np.clip(y_pred, lower_bound, upper_bound)
    y_pred, y_pred_corrected, _ = apply_water_mask(
        y_pred, y_pred_corrected, test_years=test_years, region=region, save_data_path=save_data_path, logger=logger
    )

    # Summary stats
    logger.info("Bias and clipping ([0, 7]) applied to land-only and full predictions.")
    log_stats("Ground Truth", y_test_land)
    log_stats("Raw Predictions", y_pred_land)
    log_stats("Clipped Predictions", y_pred_land_clipped)
    log_stats("Corrected Predictions", y_pred_land_corrected)

    # Generate post processing plots 
    if post_plots:
        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, upper_bound)
        plot_monthly_histograms(y_test_land, y_pred_land_clipped, region, timestamp, save_data_path)
        plot_monthly_residual_curves(y_test_land, y_pred_land, region, timestamp, save_data_path)
        logger.info("Monthly histograms (ground truth vs prediction) saved")

    # Raw metrics
    metrics = lambda yt, yp: (
        mean_squared_error(yt, yp),
        np.sqrt(mean_squared_error(yt, yp)),
        mean_absolute_error(yt, yp),
        r2_score(yt, yp)
    )

    raw_metrics = metrics(y_test_land, y_pred_land)
    clipped_metrics = metrics(y_test_land, y_pred_land_clipped)
    corrected_metrics = metrics(y_test_land, y_pred_land_corrected)

    logger.info(f'Raw RNN Metrics (land-only) — MSE: {raw_metrics[0]:.4f}, RMSE: {raw_metrics[1]:.4f}, MAE: {raw_metrics[2]:.4f}, R²: {raw_metrics[3]:.4f}')
    logger.info(f'Clipped RNN Metrics (land-only) — MSE: {clipped_metrics[0]:.4f}, RMSE: {clipped_metrics[1]:.4f}, MAE: {clipped_metrics[2]:.4f}, R²: {clipped_metrics[3]:.4f}')
    logger.info(f'Bias Corrected and Clipped RNN Metrics (land-only) — MSE: {corrected_metrics[0]:.4f}, RMSE: {corrected_metrics[1]:.4f}, MAE: {corrected_metrics[2]:.4f}, R²: {corrected_metrics[3]:.4f}')

    # Return dictionary of predictions and evaluation metrics
    # - Full-grid outputs: y_pred (clipped), y_pred_corrected
    # - Land-only slices: y_pred_land, clipped, corrected
    return {
        "y_pred": y_pred,
        "y_pred_land": y_pred_land,
        "y_pred_land_clipped": y_pred_land_clipped,
        "y_pred_land_corrected": y_pred_land_corrected,
        "metrics": {
            "raw": raw_metrics,
            "clipped": clipped_metrics,
            "corrected": corrected_metrics,
            "bias": bias
        }
    }
