"""
preprocessing.py - Data Preprocessing for LAI Prediction Pipeline

This module provides routines for preprocessing LAI input data before model training,
including:

- Missing value imputation (simple and KNN-based)
- Log transformation and scaling of targets/features
- One-hot encoding of categorical land cover classes
- Data reshaping and conversion to PyTorch tensors
- Preparation of training, validation, and test DataLoaders
- Loading and applying persistent water masks to exclude water pixels

These functions ensure the input data is clean, consistent, and formatted for deep learning workflows.
"""

import torch
import os
import zarr
import logging
from torch.utils.data import DataLoader
import numpy as np
from sklearn.impute import SimpleImputer, KNNImputer
from modeling import LAIDataset
from utils import scale_data, plot_histogram

def apply_imputation(X_train, X_val, X_test, strategy="mean"):
    """
    Applies missing value imputation to 3D input data using the specified strategy.

    Args:
        X_train (np.ndarray): Training data (N x T x F).
        X_val (np.ndarray): Validation data (N x T x F).
        X_test (np.ndarray): Test data (N x T x F).
        strategy (str): Imputation method ('mean', 'median', 'most_frequent', or 'knn').

    Returns:
        tuple: (X_train, X_val, X_test) with missing values filled.
    """
    if strategy in ['mean', 'median', 'most_frequent']:
        imputer = SimpleImputer(strategy=strategy)
    else:
        imputer = KNNImputer(n_neighbors=5)

    # Reshape from (N, T, F) → (N*T, F) for imputation (imputation for all location-year pairs a the same time)
    X_train_2d = X_train.reshape(-1, X_train.shape[-1])
    X_val_2d = X_val.reshape(-1, X_val.shape[-1])
    X_test_2d = X_test.reshape(-1, X_test.shape[-1])

    X_train_2d = imputer.fit_transform(X_train_2d)
    X_val_2d = imputer.transform(X_val_2d)
    X_test_2d = imputer.transform(X_test_2d)

    X_train = X_train_2d.reshape(X_train.shape)
    X_val = X_val_2d.reshape(X_val.shape)
    X_test = X_test_2d.reshape(X_test.shape)

    return X_train, X_val, X_test

def apply_imputation_yearly(X_train, X_val, X_test, numeric_strategy="mean", categorical_strategy="most_frequent"):
    """
    Separately imputes numeric and categorical variables in yearly feature data.

    Args:
        X_train, X_val, X_test (np.ndarray): Input arrays (N x T x F), where the last feature is categorical.
        numeric_strategy (str): Strategy for numeric features ('mean', 'median', etc.).
        categorical_strategy (str): Strategy for categorical feature ('most_frequent' recommended).

    Returns:
        tuple: Imputed arrays (X_train, X_val, X_test) with final feature recombined.
    """
    logger = logging.getLogger(__name__)
    
    # Split into numeric and categorical
    X_train_num, X_train_cat = X_train[:, :, :-1], X_train[:, :, -1:]
    X_val_num, X_val_cat = X_val[:, :, :-1], X_val[:, :, -1:]
    X_test_num, X_test_cat = X_test[:, :, :-1], X_test[:, :, -1:]

    logger.info(f"[LAI + LUH2 IMPUTATION] Missing values before imputation - Train: {np.isnan(X_train_num).sum()}, Val: {np.isnan(X_val_num).sum()}, Test: {np.isnan(X_test_num).sum()} (Strategy: {numeric_strategy})")

    # Impute numeric part
    X_train_num, X_val_num, X_test_num = apply_imputation(
        X_train_num, X_val_num, X_test_num, strategy=numeric_strategy
    )

    logger.info(f"[HILDA+ IMPUTATION] Missing values before imputation - Train: {np.isnan(X_train_cat).sum()}, Val: {np.isnan(X_val_cat).sum()}, Test: {np.isnan(X_test_cat).sum()} (Strategy: {categorical_strategy})")
    
    # Reshape and impute categorical (reshape: from (N, T, 1) to (N*T, 1), again for all location-year pairs)
    shape = X_train_cat.shape
    imp_cat = SimpleImputer(strategy=categorical_strategy)
    X_train_cat_flat = imp_cat.fit_transform(X_train_cat.reshape(-1, 1)).reshape(shape)
    X_val_cat_flat = imp_cat.transform(X_val_cat.reshape(-1, 1)).reshape(X_val_cat.shape)
    X_test_cat_flat = imp_cat.transform(X_test_cat.reshape(-1, 1)).reshape(X_test_cat.shape)

    # Recombine LUH2 numerical features and HILDA+ after imputing sepparately
    X_train_combined = np.concatenate([X_train_num, X_train_cat_flat], axis=-1)
    X_val_combined = np.concatenate([X_val_num, X_val_cat_flat], axis=-1)
    X_test_combined = np.concatenate([X_test_num, X_test_cat_flat], axis=-1)

    return X_train_combined, X_val_combined, X_test_combined


def one_hot_encode_yearly(X):
    """
    One-hot encodes the last feature (HILDA+) of a 3D array containing land cover class labels.

    Assumes the last feature contains HILDA+ class codes (e.g., 11, 22, ..., 99).
    Unknown or unmapped values are converted to 99 before encoding.

    Args:
        X (np.ndarray): Input array (N x T x F), where the last feature is categorical.

    Returns:
        np.ndarray: Transformed array with one-hot encoded land cover (N x T x F+num_classes-1).
    """
    # Full set of land cover classes including unknown (99)
    landcover_classes = [11, 22, 33, 40, 41, 42, 43, 44, 45, 55, 66, 77, 99]
    class_map = {v: i for i, v in enumerate(landcover_classes)}

    numeric = X[:, :, :-1]  # (N, T, F-1)
    categorical = X[:, :, -1].astype(int)  # (N, T)

    logging.info(f"Original unique HILDA+ values before mapping: {np.unique(categorical)}")

    # Replace unknown values with 99
    known_set = set(landcover_classes)
    categorical[~np.isin(categorical, list(known_set))] = 99

    logging.info(f"Unique HILDA+ values after remapping unknowns to 99: {np.unique(categorical)}")

    # Now safe to encode
    mapped = np.vectorize(lambda x: class_map[x])(categorical)
    
    # After encoding: one_hot → (N, T, num_classes)
    one_hot = np.eye(len(landcover_classes))[mapped]

    # Final concat → (N, T, F-1 + num_classes)
    return np.concatenate([numeric, one_hot], axis=-1)


def preprocess_lai_data(
    X_train_monthly, X_train_yearly, y_train,
    X_val_monthly, X_val_yearly, y_val,
    X_test_monthly, X_test_yearly, y_test,
    scale_gt=False, log_transform_gt=False,
    lower_bound=0, upper_bound=7,
    plot=False, save_data_path=".", region="unknown",
    timestamp="0000", imputation_method="mean"
):
    """
    Applies full preprocessing pipeline to LAI input features and targets.

    Includes:
    - Imputation
    - Optional log transform
    - Optional scaling to [0,1]
    - One-hot encoding of land cover
    - Plotting histograms of transformations
    - Conversion to PyTorch tensors

    Args:
        X_train_monthly, X_val_monthly, X_test_monthly (np.ndarray): Monthly LAI features.
        X_train_yearly, X_val_yearly, X_test_yearly (np.ndarray): Yearly static features.
        y_train, y_val, y_test (np.ndarray): LAI targets.
        scale_gt (bool): Whether to scale targets and features.
        log_transform_gt (bool): Whether to apply log1p transform.
        lower_bound (float): Minimum value for scaling.
        upper_bound (float): Maximum value for scaling.
        plot (bool): Whether to generate histogram plots.
        save_data_path (str): Path to save plots.
        region (str): Region identifier for filenames.
        timestamp (str): Timestamp string for filenames.
        imputation_method (str): Imputation strategy.

    Returns:
        tuple: Preprocessed tensors (X_train_monthly, X_train_yearly, y_train, ...)
               along with unnormalized target copies and scaling bounds.
    """

    # Impute missing values
    X_train_monthly, X_val_monthly, X_test_monthly = apply_imputation(
        X_train_monthly, X_val_monthly, X_test_monthly, strategy=imputation_method
    )
    X_train_yearly, X_val_yearly, X_test_yearly = apply_imputation_yearly(
        X_train_yearly, X_val_yearly, X_test_yearly
    )

    # Initialize scaling min and max values
    scaling_min = lower_bound
    scaling_max = upper_bound
    
    # Plot raw histograms of train and validation data
    if plot:
        plot_histogram(
            y_train.reshape(-1), y_val.reshape(-1),
            title="Raw LAI Train vs Val",
            filename=os.path.join(save_data_path, f"lai_hist_raw_{region}_{timestamp}.png")
        ) 

    # Log transform
    if log_transform_gt:
        y_train = np.log1p(y_train)
        y_val = np.log1p(y_val)
        y_test = np.log1p(y_test)

        X_train_monthly = np.log1p(X_train_monthly)
        X_val_monthly = np.log1p(X_val_monthly)
        X_test_monthly = np.log1p(X_test_monthly)

        if plot:
            plot_histogram(
                y_train.reshape(-1), y_val.reshape(-1),
                title="Log-transformed LAI Train vs Val",
                filename=os.path.join(save_data_path, f"lai_hist_log_{region}_{timestamp}.png"),
                xlabel="LAI value"
            )

    # Scaling to [0,1] range
    if scale_gt:
        if log_transform_gt:
            # Compute dynamic min/max from log-transformed training data
            scaling_min = min(y_train.min(), X_train_monthly.min())
            scaling_max = max(y_train.max(), X_train_monthly.max())

            # Scale valdiation and test with train values to avoid data leakage
            y_train = scale_data(y_train, scaling_min, scaling_max)
            y_val = scale_data(y_val, scaling_min, scaling_max)
            y_test = scale_data(y_test, scaling_min, scaling_max)

            X_train_monthly = scale_data(X_train_monthly, scaling_min, scaling_max)
            X_val_monthly = scale_data(X_val_monthly, scaling_min, scaling_max)
            X_test_monthly = scale_data(X_test_monthly, scaling_min, scaling_max)
            
            if plot:
                plot_histogram(
                    y_train.reshape(-1), y_val.reshape(-1),
                    title="Log-transformed and scaled (range [0,1]) LAI Train vs Val",
                    filename=os.path.join(save_data_path, f"lai_hist_log_scaled_{region}_{timestamp}.png"),
                    xlabel="Log-transformed and scaled LAI value"
                )
        else:
            # If no log-transformation scale with the default bounds [0,7]
            y_train = scale_data(y_train, lower_bound, upper_bound)
            y_val = scale_data(y_val, lower_bound, upper_bound)
            y_test = scale_data(y_test, lower_bound, upper_bound)

            X_train_monthly = scale_data(X_train_monthly, lower_bound, upper_bound)
            X_val_monthly = scale_data(X_val_monthly, lower_bound, upper_bound)
            X_test_monthly = scale_data(X_test_monthly, lower_bound, upper_bound)

            if plot:
                plot_histogram(
                    y_train.reshape(-1), y_val.reshape(-1),
                    title="Scaled (range [0,1]) LAI Train vs Val",
                    filename=os.path.join(save_data_path, f"lai_hist_scaled_{region}_{timestamp}.png"),
                    xlabel="Scaled LAI value"
                )

    # Save "post" versions (for inverse-transforming predictions later)
    y_train_post = y_train.copy()
    y_val_post = y_val.copy()
    y_test_post = y_test.copy()

    # One-hot encode the last yearly variable (variable 15) containing 12 classes
    logging.info(f"One-hot encoding for TRAIN:")
    X_train_yearly = one_hot_encode_yearly(X_train_yearly) # After one-hot: X_train_yearly → (N, T, 27)
    logging.info(f"One-hot encoding for VALIDATION:")
    X_val_yearly = one_hot_encode_yearly(X_val_yearly)
    logging.info(f"One-hot encoding for TEST:")
    X_test_yearly = one_hot_encode_yearly(X_test_yearly)
    
    # Convert to torch tensors
    X_train_monthly = torch.tensor(X_train_monthly, dtype=torch.float32)
    X_train_yearly = torch.tensor(X_train_yearly, dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.float32)

    X_val_monthly = torch.tensor(X_val_monthly, dtype=torch.float32)
    X_val_yearly = torch.tensor(X_val_yearly, dtype=torch.float32)
    y_val = torch.tensor(y_val, dtype=torch.float32)

    X_test_monthly = torch.tensor(X_test_monthly, dtype=torch.float32)
    X_test_yearly = torch.tensor(X_test_yearly, dtype=torch.float32)
    y_test = torch.tensor(y_test, dtype=torch.float32)

    return (
        X_train_monthly, X_train_yearly, y_train,
        X_val_monthly, X_val_yearly, y_val,
        X_test_monthly, X_test_yearly, y_test,
        y_train_post, y_val_post, y_test_post,
        scaling_min, scaling_max,
    )


def prepare_dataloaders(
    X_train_monthly, X_train_yearly, y_train,
    X_val_monthly, X_val_yearly, y_val,
    X_test_monthly, X_test_yearly, y_test,
    batch_size=16384, num_workers=32, num_yearly_variables=27
):
    """
    Prepares PyTorch DataLoaders for training, validation, and test sets.

    - Flattens data across spatial dimensions
    - Flips monthly sequences (Dec→Jan)
    - Wraps inputs into LAIDataset
    - Builds DataLoaders for each set

    Args:
        X_train_monthly, X_val_monthly, X_test_monthly (Tensor): Monthly inputs (B x 12).
        X_train_yearly, X_val_yearly, X_test_yearly (Tensor): Yearly inputs (B x F).
        y_train, y_val, y_test (Tensor): LAI targets (B x 12).
        batch_size (int): DataLoader batch size.
        num_workers (int): Number of parallel workers for data loading.
        num_yearly_variables (int): Number of static variables per location.

    Returns:
        tuple: (train_loader, val_loader, test_loader, train_dataset, val_dataset, test_dataset)
    """


    # Reshape from (N, T, 12) → (N*T, 12), then reverse time: Dec → Jan
    X_train_monthly_flat = torch.flip(X_train_monthly.reshape(-1, 12), dims=[1]) # (N*T, 12)
    X_train_yearly_flat = X_train_yearly.reshape(-1, num_yearly_variables) # (N*T, 27)
    y_train_flat = torch.flip(y_train.reshape(-1, 12), dims=[1]) # (N*T, 12)

    X_val_monthly_flat = torch.flip(X_val_monthly.reshape(-1, 12), dims=[1])
    X_val_yearly_flat = X_val_yearly.reshape(-1, num_yearly_variables)
    y_val_flat = torch.flip(y_val.reshape(-1, 12), dims=[1])

    X_test_monthly_flat = torch.flip(X_test_monthly.reshape(-1, 12), dims=[1])
    X_test_yearly_flat = X_test_yearly.reshape(-1, num_yearly_variables)
    y_test_flat = torch.flip(y_test.reshape(-1, 12), dims=[1])

    # Build datasets
    train_dataset = LAIDataset(X_train_monthly_flat, X_train_yearly_flat, y_train_flat)
    val_dataset = LAIDataset(X_val_monthly_flat, X_val_yearly_flat, y_val_flat)
    test_dataset = LAIDataset(X_test_monthly_flat, X_test_yearly_flat, y_test_flat)

    # Build data loaders
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)

    return train_loader, val_loader, test_loader, train_dataset, val_dataset, test_dataset


def load_combined_water_mask(years, region, save_data_path, logger=None):
    """
    Loads and combines annual water masks into a persistent land mask.

    Applies a logical AND across all years to retain only permanently non-water locations.

    Args:
        years (list of int): List of years to include.
        region (str): Region identifier.
        save_data_path (str): Directory containing water mask Zarr files.
        logger (logging.Logger, optional): Logger for status messages.

    Returns:
        np.ndarray: Flattened boolean array where True indicates land (not water).

    Raises:
        RuntimeError: If no valid water masks could be loaded.
    """
    
    logger = logger or logging.getLogger(__name__)
    combined = None

    for year in years:
        path = os.path.join(save_data_path, f"water_mask_{region}_{year}.zarr")
        if not os.path.exists(path):
            logger.warning(f"Water mask not found for {year}: {path}")
            continue
        
        # Each mask: zarr.load → shape (lat * lon,), e.g., (15052800,)
        mask = zarr.load(path).astype(bool).flatten(order="C")
        logger.info(f"Loaded water mask for {year}, shape: {mask.shape}")

        # Combined: (15052800,) boolean array where True = land
        if combined is None:
            combined = mask
        else:
            combined &= mask  # keep only persistent land

    if combined is None:
        raise RuntimeError("No valid water masks loaded.")

    logger.info(f"Water mask applied — retained {np.sum(combined):,} land locations.")
    return combined
