"""
Data cleaning utilities.
Single Responsibility: Clean and preprocess hydrological/meteorological data.
"""
import pandas as pd
import numpy as np
from typing import Optional, List


class DataCleaner:
    """
    Cleans hydrological and meteorological datasets.
    Handles missing values, outliers, and data quality issues.
    """

    def __init__(
        self,
        missing_threshold: float = 0.3,
        outlier_method: str = "iqr"
    ):
        """
        Initialize data cleaner.

        Args:
            missing_threshold: Max fraction of missing values allowed per column
            outlier_method: Method for outlier detection ("iqr", "zscore", "none")
        """
        self.missing_threshold = missing_threshold
        self.outlier_method = outlier_method

    def clean(
        self,
        df: pd.DataFrame,
        date_col: str = "date",
        station_col: str = "station_id"
    ) -> pd.DataFrame:
        """
        Clean dataframe with multiple steps.

        Args:
            df: Input dataframe
            date_col: Name of date column
            station_col: Name of station ID column

        Returns:
            Cleaned dataframe
        """
        df = df.copy()

        # Remove duplicates
        df = self._remove_duplicates(df, date_col, station_col)

        # Handle missing values
        df = self._handle_missing(df, date_col, station_col)

        # Remove outliers if requested
        if self.outlier_method != "none":
            df = self._remove_outliers(df, date_col, station_col)

        # Sort by date and station
        if date_col in df.columns and station_col in df.columns:
            df = df.sort_values([station_col, date_col]).reset_index(drop=True)

        return df

    def _remove_duplicates(
        self,
        df: pd.DataFrame,
        date_col: str,
        station_col: str
    ) -> pd.DataFrame:
        """Remove duplicate rows based on date and station."""
        if date_col in df.columns and station_col in df.columns:
            df = df.drop_duplicates(subset=[date_col, station_col], keep="first")
        return df

    def _handle_missing(
        self,
        df: pd.DataFrame,
        date_col: str,
        station_col: str
    ) -> pd.DataFrame:
        """
        Handle missing values.
        - Drop columns with >threshold missing
        - Forward fill small gaps in time series
        """
        # Drop columns with too many missing values
        missing_frac = df.isnull().mean()
        cols_to_keep = missing_frac[missing_frac <= self.missing_threshold].index
        df = df[cols_to_keep]

        # Forward fill small gaps (max 3 days) within each station
        if station_col in df.columns:
            numeric_cols = df.select_dtypes(include=[np.number]).columns
            df[numeric_cols] = df.groupby(station_col)[numeric_cols].transform(
                lambda x: x.fillna(method='ffill', limit=3)
            )

        return df

    def _remove_outliers(
        self,
        df: pd.DataFrame,
        date_col: str,
        station_col: str
    ) -> pd.DataFrame:
        """
        Remove outliers using IQR or Z-score method.
        Only applies to numeric columns.
        """
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        numeric_cols = [c for c in numeric_cols if c != station_col]

        if self.outlier_method == "iqr":
            df = self._remove_outliers_iqr(df, numeric_cols, station_col)
        elif self.outlier_method == "zscore":
            df = self._remove_outliers_zscore(df, numeric_cols, station_col)

        return df

    def _remove_outliers_iqr(
        self,
        df: pd.DataFrame,
        cols: List[str],
        station_col: str
    ) -> pd.DataFrame:
        """Remove outliers using IQR method (per station)."""
        for col in cols:
            if col in df.columns:
                # Calculate IQR per station
                Q1 = df.groupby(station_col)[col].transform(lambda x: x.quantile(0.25))
                Q3 = df.groupby(station_col)[col].transform(lambda x: x.quantile(0.75))
                IQR = Q3 - Q1
                lower_bound = Q1 - 3 * IQR
                upper_bound = Q3 + 3 * IQR

                # Set outliers to NaN
                df.loc[(df[col] < lower_bound) | (df[col] > upper_bound), col] = np.nan

        return df

    def _remove_outliers_zscore(
        self,
        df: pd.DataFrame,
        cols: List[str],
        station_col: str,
        threshold: float = 4.0
    ) -> pd.DataFrame:
        """Remove outliers using Z-score method (per station)."""
        for col in cols:
            if col in df.columns:
                # Calculate Z-scores per station
                mean = df.groupby(station_col)[col].transform('mean')
                std = df.groupby(station_col)[col].transform('std')
                z_scores = np.abs((df[col] - mean) / std)

                # Set outliers to NaN
                df.loc[z_scores > threshold, col] = np.nan

        return df

    def get_cleaning_report(self, df_before: pd.DataFrame, df_after: pd.DataFrame) -> dict:
        """
        Generate report of cleaning operations.

        Args:
            df_before: DataFrame before cleaning
            df_after: DataFrame after cleaning

        Returns:
            Dictionary with cleaning statistics
        """
        return {
            "rows_before": len(df_before),
            "rows_after": len(df_after),
            "rows_removed": len(df_before) - len(df_after),
            "cols_before": len(df_before.columns),
            "cols_after": len(df_after.columns),
            "cols_removed": len(df_before.columns) - len(df_after.columns)
        }