"""
Data Preprocessing Module

Handles normalization, missing value interpolation, outlier detection,
and data validation for climate data. All preprocessing is designed
to be deterministic and reproducible.

Design Decisions:
- Statistics computed once on training data, applied to all splits
- Missing values handled by spatial/temporal interpolation
- Outliers are flagged, not removed (climate extremes are real)
- All transformations are invertible for output interpretation

Why Z-Score Normalization?
- Neural networks train better with zero-mean, unit-variance inputs
- Gradient flow is more stable across layers
- Prevents any single variable from dominating

Time Complexity: O(n) for n data points
Space Complexity: O(n) for data + O(1) for statistics
"""

import numpy as np
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
from scipy.ndimage import binary_dilation
from scipy.interpolate import griddata
import warnings


class Preprocessor:
    """
    Climate data preprocessor with normalization and quality assurance.
    
    Designed to be fitted once on training data and applied to all splits.
    Statistics are persistable for inference on new data.
    
    Attributes:
        config: Configuration dictionary
        statistics: Dict of per-variable normalization statistics
        is_fitted: Whether the preprocessor has been fitted
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        Initialize the preprocessor.
        
        Args:
            config: Configuration dictionary with preprocessing settings
        """
        self.config = config
        
        # Extract settings
        preproc_config = config.get("preprocessing", {})
        self.normalize = preproc_config.get("normalize", True)
        self.normalize_method = preproc_config.get("normalize_method", "zscore")
        self.handle_missing = preproc_config.get("handle_missing", "interpolate")
        self.outlier_method = preproc_config.get("outlier_method", "zscore")
        self.outlier_threshold = preproc_config.get("outlier_threshold", 3.0)
        self.clip_outliers = preproc_config.get("clip_outliers", False)
        
        # Per-variable statistics (computed during fit)
        self.statistics: Dict[str, Dict[str, float]] = {}
        self.is_fitted = False
        
        # Variables that require log-transformation (skewed distributions)
        self.log_vars = {"tp", "precip", "total_precipitation", "precipitation"}
    
    def _is_log_var(self, variable: str) -> bool:
        """Check if variable requires log transformation."""
        return variable.lower() in self.log_vars

    def fit(self, data: np.ndarray, variable: str) -> 'Preprocessor':
        """
        Compute normalization statistics from training data.
        
        Only call this on training data to avoid data leakage.
        
        Args:
            data: Array of shape (time, lat, lon)
            variable: Variable name for statistics storage
            
        Returns:
            self for method chaining
        """
        # Handle missing values first
        clean_data = self._handle_missing_values(data)
        
        # Apply log transform if needed
        if self._is_log_var(variable):
            # log1p(x) = log(x + 1) handles zeros gracefully
            # ensure non-negative
            clean_data = np.log1p(np.maximum(clean_data, 0))
        
        # Compute statistics (ignoring NaN if any remain)
        if self.normalize_method == "zscore":
            mean = float(np.nanmean(clean_data))
            std = float(np.nanstd(clean_data))
            # Prevent division by zero
            if std < 1e-8:
                std = 1.0
                warnings.warn(f"Variable {variable} has near-zero std, using 1.0")
            
            self.statistics[variable] = {
                "mean": mean,
                "std": std,
                "min": float(np.nanmin(clean_data)),
                "max": float(np.nanmax(clean_data)),
            }
        
        elif self.normalize_method == "minmax":
            min_val = float(np.nanmin(clean_data))
            max_val = float(np.nanmax(clean_data))
            # Prevent division by zero
            if max_val - min_val < 1e-8:
                max_val = min_val + 1.0
                warnings.warn(f"Variable {variable} has near-zero range")
            
            self.statistics[variable] = {
                "min": min_val,
                "max": max_val,
                "mean": float(np.nanmean(clean_data)),
                "std": float(np.nanstd(clean_data)),
            }
        
        self.is_fitted = True
        return self
    
    def transform(
        self,
        data: np.ndarray,
        variable: str
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply preprocessing transformations.
        
        Steps:
        1. Handle missing values
        2. Detect outliers
        3. Apply normalization
        
        Args:
            data: Array of shape (time, lat, lon)
            variable: Variable name for looking up statistics
            
        Returns:
            Tuple of (transformed_data, outlier_mask)
        """
        if not self.is_fitted:
            raise RuntimeError("Preprocessor must be fitted before transform")
        
        if variable not in self.statistics:
            raise KeyError(f"No statistics for variable '{variable}'. Fit first.")
        
        # Step 1: Handle missing values
        processed = self._handle_missing_values(data.copy())
        
        # Apply log transform if needed (before outlier/norm)
        if self._is_log_var(variable):
            processed = np.log1p(np.maximum(processed, 0))
        
        # Step 2: Detect outliers
        outlier_mask = self._detect_outliers(processed, variable)
        
        # Optionally clip outliers
        if self.clip_outliers:
            processed = self._clip_outliers(processed, variable)
        
        # Step 3: Normalize
        if self.normalize:
            processed = self._normalize(processed, variable)
        
        return processed.astype(np.float32), outlier_mask

    def fit_transform(
        self,
        data: np.ndarray,
        variable: str
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Fit and transform in one step (for training data).
        
        Args:
            data: Array of shape (time, lat, lon)
            variable: Variable name
            
        Returns:
            Tuple of (transformed_data, outlier_mask)
        """
        self.fit(data, variable)
        return self.transform(data, variable)

    def inverse_transform(self, data: np.ndarray, variable: str) -> np.ndarray:
        """
        Reverse the normalization transformation.
        
        Used to convert model outputs back to physical units.
        
        Args:
            data: Normalized data array
            variable: Variable name
            
        Returns:
            Data in original physical units
        """
        if not self.is_fitted:
            raise RuntimeError("Preprocessor must be fitted first")
        
        stats = self.statistics[variable]
        
        # 1. Denormalize
        if self.normalize_method == "zscore":
            denorm = data * stats["std"] + stats["mean"]
        elif self.normalize_method == "minmax":
            denorm = data * (stats["max"] - stats["min"]) + stats["min"]
        else:
            denorm = data
            
        # 2. Inverse log (expm1) if needed
        if self._is_log_var(variable):
            denorm = np.expm1(denorm)
            # Clip negative values that might result from numerical noise
            denorm = np.maximum(denorm, 0)
        
        return denorm
    
    def _handle_missing_values(self, data: np.ndarray) -> np.ndarray:
        """
        Handle missing values in climate data.
        
        Strategies:
        - interpolate: Spatial/temporal interpolation
        - mask: Keep NaN and let training handle it
        - drop: Not recommended, raises warning
        
        Args:
            data: Input array (may contain NaN)
            
        Returns:
            Array with missing values handled
        """
        # Identify missing values
        missing_mask = np.isnan(data) | np.isinf(data)
        
        if not missing_mask.any():
            return data
        
        n_missing = missing_mask.sum()
        total = data.size
        missing_pct = 100 * n_missing / total
        
        if missing_pct > 10:
            warnings.warn(
                f"High missing data percentage: {missing_pct:.1f}%. "
                "Consider data quality review."
            )
        
        if self.handle_missing == "mask":
            # Keep NaN - training will use masked loss
            return data
        
        elif self.handle_missing == "drop":
            warnings.warn("'drop' strategy removes data. Use 'interpolate' instead.")
            return data
        
        elif self.handle_missing == "interpolate":
            return self._interpolate_missing(data, missing_mask)
        
        return data
    
    def _interpolate_missing(
        self,
        data: np.ndarray,
        missing_mask: np.ndarray
    ) -> np.ndarray:
        """
        Interpolate missing values using spatial then temporal interpolation.
        
        Strategy:
        1. Try spatial interpolation within each timestep
        2. Fall back to temporal interpolation for remaining gaps
        3. Use mean for any remaining values
        
        Args:
            data: Data array with missing values
            missing_mask: Boolean mask of missing locations
            
        Returns:
            Interpolated data array
        """
        result = data.copy()
        
        # Process each timestep
        for t in range(data.shape[0]):
            frame = result[t]
            mask = missing_mask[t]
            
            if not mask.any():
                continue
            
            # Get coordinates of valid and missing points
            valid_points = np.argwhere(~mask)
            missing_points = np.argwhere(mask)
            
            if len(valid_points) < 4:
                # Not enough valid points for interpolation
                # Use previous/next timestep if available
                if t > 0:
                    result[t][mask] = result[t-1][mask]
                elif t < data.shape[0] - 1:
                    result[t][mask] = data[t+1][mask]
                continue
            
            # Spatial interpolation
            valid_values = frame[~mask]
            
            try:
                interpolated = griddata(
                    valid_points,
                    valid_values,
                    missing_points,
                    method='linear',
                    fill_value=np.nanmean(valid_values)
                )
                
                # Fill in interpolated values
                for i, point in enumerate(missing_points):
                    result[t, point[0], point[1]] = interpolated[i]
            except Exception:
                # Fall back to mean fill
                result[t][mask] = np.nanmean(frame)
        
        # Any remaining NaN gets filled with global mean
        remaining_nan = np.isnan(result)
        if remaining_nan.any():
            result[remaining_nan] = np.nanmean(result)
        
        return result
    
    def _detect_outliers(
        self,
        data: np.ndarray,
        variable: str
    ) -> np.ndarray:
        """
        Detect outliers using configured method.
        
        Note: Outliers are flagged but not removed by default.
        Climate extremes (heat waves, heavy rain) are real events.
        
        Args:
            data: Data array
            variable: Variable name
            
        Returns:
            Boolean mask where True indicates outlier
        """
        if self.outlier_method == "none":
            return np.zeros_like(data, dtype=bool)
        
        stats = self.statistics[variable]
        
        if self.outlier_method == "zscore":
            z_scores = np.abs((data - stats["mean"]) / stats["std"])
            return z_scores > self.outlier_threshold
        
        elif self.outlier_method == "iqr":
            # Compute quartiles from stored statistics
            # This is an approximation using normal distribution assumption
            q1 = stats["mean"] - 0.675 * stats["std"]
            q3 = stats["mean"] + 0.675 * stats["std"]
            iqr = q3 - q1
            lower = q1 - self.outlier_threshold * iqr
            upper = q3 + self.outlier_threshold * iqr
            return (data < lower) | (data > upper)
        
        return np.zeros_like(data, dtype=bool)
    
    def _clip_outliers(self, data: np.ndarray, variable: str) -> np.ndarray:
        """
        Clip outliers to threshold boundaries.
        
        Args:
            data: Data array
            variable: Variable name
            
        Returns:
            Clipped data array
        """
        stats = self.statistics[variable]
        
        if self.outlier_method == "zscore":
            lower = stats["mean"] - self.outlier_threshold * stats["std"]
            upper = stats["mean"] + self.outlier_threshold * stats["std"]
        else:
            lower = stats["min"]
            upper = stats["max"]
        
        return np.clip(data, lower, upper)
    
    def _normalize(self, data: np.ndarray, variable: str) -> np.ndarray:
        """
        Apply normalization transformation.
        
        Args:
            data: Data array
            variable: Variable name
            
        Returns:
            Normalized data array
        """
        stats = self.statistics[variable]
        
        if self.normalize_method == "zscore":
            return (data - stats["mean"]) / stats["std"]
        
        elif self.normalize_method == "minmax":
            return (data - stats["min"]) / (stats["max"] - stats["min"])
        
        return data
    
    def save_statistics(self, path: str) -> None:
        """
        Save fitted statistics to disk for later use.
        
        Args:
            path: Path to save statistics (JSON-like format via NumPy)
        """
        save_path = Path(path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        
        np.savez(
            save_path,
            statistics=np.array([self.statistics], dtype=object),
            normalize_method=self.normalize_method,
            outlier_method=self.outlier_method,
            outlier_threshold=self.outlier_threshold,
        )
    
    def load_statistics(self, path: str) -> None:
        """
        Load previously saved statistics.
        
        Args:
            path: Path to statistics file
        """
        loaded = np.load(path, allow_pickle=True)
        self.statistics = loaded["statistics"].item()
        self.normalize_method = str(loaded["normalize_method"])
        self.outlier_method = str(loaded["outlier_method"])
        self.outlier_threshold = float(loaded["outlier_threshold"])
        self.is_fitted = True
    
    def get_report(self) -> Dict[str, Any]:
        """
        Generate a preprocessing report.
        
        Returns:
            Dictionary with preprocessing summary
        """
        return {
            "normalize_method": self.normalize_method,
            "outlier_method": self.outlier_method,
            "outlier_threshold": self.outlier_threshold,
            "handle_missing": self.handle_missing,
            "variables": list(self.statistics.keys()),
            "statistics": self.statistics,
            "is_fitted": self.is_fitted,
        }