"""
Data normalization strategies.
Single Responsibility: Normalize hydrological/meteorological features.
Open/Closed Principle: Easy to add new normalization methods.
"""
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple
from abc import ABC, abstractmethod
import pickle


class BaseNormalizer(ABC):
    """
    Abstract base class for normalization strategies.
    Dependency Inversion: Depend on abstraction, not concrete normalizers.
    """

    def __init__(self):
        """Initialize normalizer."""
        self.stats: Dict[str, Dict] = {}
        self.is_fitted = False

    @abstractmethod
    def fit(self, df: pd.DataFrame, columns: Optional[List[str]] = None) -> 'BaseNormalizer':
        """
        Fit normalizer to data (compute statistics).

        Args:
            df: Input dataframe
            columns: Columns to normalize (None = all numeric)

        Returns:
            self (for chaining)
        """
        pass

    @abstractmethod
    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Transform data using fitted statistics.

        Args:
            df: Input dataframe

        Returns:
            Normalized dataframe
        """
        pass

    def fit_transform(self, df: pd.DataFrame, columns: Optional[List[str]] = None) -> pd.DataFrame:
        """
        Fit and transform in one step.

        Args:
            df: Input dataframe
            columns: Columns to normalize

        Returns:
            Normalized dataframe
        """
        self.fit(df, columns)
        return self.transform(df)

    def inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Reverse normalization.

        Args:
            df: Normalized dataframe

        Returns:
            Original scale dataframe
        """
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before inverse transform")
        return df.copy()

    def save(self, filepath: str) -> None:
        """Save normalizer statistics to file."""
        with open(filepath, 'wb') as f:
            pickle.dump(self.stats, f)

    def load(self, filepath: str) -> None:
        """Load normalizer statistics from file."""
        with open(filepath, 'rb') as f:
            self.stats = pickle.load(f)
        self.is_fitted = True


class StandardNormalizer(BaseNormalizer):
    """
    Z-score normalization: (x - mean) / std
    Best for normally distributed data.
    """

    def fit(self, df: pd.DataFrame, columns: Optional[List[str]] = None) -> 'StandardNormalizer':
        """Fit by computing mean and std for each column."""
        if columns is None:
            columns = df.select_dtypes(include=[np.number]).columns.tolist()

        for col in columns:
            if col in df.columns:
                self.stats[col] = {
                    'mean': df[col].mean(),
                    'std': df[col].std()
                }

        self.is_fitted = True
        return self

    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Apply z-score normalization."""
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before transform")

        df_norm = df.copy()
        for col, stats in self.stats.items():
            if col in df_norm.columns:
                df_norm[col] = (df_norm[col] - stats['mean']) / stats['std']

        return df_norm

    def inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Reverse z-score normalization."""
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before inverse transform")

        df_orig = df.copy()
        for col, stats in self.stats.items():
            if col in df_orig.columns:
                df_orig[col] = df_orig[col] * stats['std'] + stats['mean']

        return df_orig


class MinMaxNormalizer(BaseNormalizer):
    """
    Min-max normalization: (x - min) / (max - min)
    Scales to [0, 1] range.
    """

    def fit(self, df: pd.DataFrame, columns: Optional[List[str]] = None) -> 'MinMaxNormalizer':
        """Fit by computing min and max for each column."""
        if columns is None:
            columns = df.select_dtypes(include=[np.number]).columns.tolist()

        for col in columns:
            if col in df.columns:
                self.stats[col] = {
                    'min': df[col].min(),
                    'max': df[col].max()
                }

        self.is_fitted = True
        return self

    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Apply min-max normalization."""
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before transform")

        df_norm = df.copy()
        for col, stats in self.stats.items():
            if col in df_norm.columns:
                denominator = stats['max'] - stats['min']
                if denominator == 0:
                    df_norm[col] = 0  # Constant column
                else:
                    df_norm[col] = (df_norm[col] - stats['min']) / denominator

        return df_norm

    def inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Reverse min-max normalization."""
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before inverse transform")

        df_orig = df.copy()
        for col, stats in self.stats.items():
            if col in df_orig.columns:
                df_orig[col] = df_orig[col] * (stats['max'] - stats['min']) + stats['min']

        return df_orig


class RobustNormalizer(BaseNormalizer):
    """
    Robust normalization using median and IQR: (x - median) / IQR
    Best for data with outliers (recommended for hydrological data).
    """

    def fit(self, df: pd.DataFrame, columns: Optional[List[str]] = None) -> 'RobustNormalizer':
        """Fit by computing median and IQR for each column."""
        if columns is None:
            columns = df.select_dtypes(include=[np.number]).columns.tolist()

        for col in columns:
            if col in df.columns:
                q1 = df[col].quantile(0.25)
                q3 = df[col].quantile(0.75)
                self.stats[col] = {
                    'median': df[col].median(),
                    'iqr': q3 - q1
                }

        self.is_fitted = True
        return self

    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Apply robust normalization."""
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before transform")

        df_norm = df.copy()
        for col, stats in self.stats.items():
            if col in df_norm.columns:
                if stats['iqr'] == 0:
                    df_norm[col] = 0  # Constant column
                else:
                    df_norm[col] = (df_norm[col] - stats['median']) / stats['iqr']

        return df_norm

    def inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Reverse robust normalization."""
        if not self.is_fitted:
            raise RuntimeError("Normalizer must be fitted before inverse transform")

        df_orig = df.copy()
        for col, stats in self.stats.items():
            if col in df_orig.columns:
                df_orig[col] = df_orig[col] * stats['iqr'] + stats['median']

        return df_orig


def get_normalizer(method: str = "robust") -> BaseNormalizer:
    """
    Factory function to get normalizer by name.
    Open/Closed Principle: Easy to extend with new normalizers.

    Args:
        method: Normalization method ("standard", "minmax", "robust")

    Returns:
        Normalizer instance
    """
    normalizers = {
        "standard": StandardNormalizer,
        "minmax": MinMaxNormalizer,
        "robust": RobustNormalizer
    }

    if method not in normalizers:
        raise ValueError(f"Unknown normalization method: {method}. "
                        f"Choose from {list(normalizers.keys())}")

    return normalizers[method]()