"""
数据清洗工具模块

提供通用的数据清洗和标准化功能
"""

import logging
from typing import List, Literal, Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class DataCleaner:
    """
    数据清洗工具类

    提供数据清洗、去重、类型转换、异常值检测等功能
    """

    def __init__(self):
        """初始化数据清洗工具"""
        logger.info("DataCleaner initialized")

    def remove_nulls(
        self,
        df: pd.DataFrame,
        columns: Optional[List[str]] = None,
        how: Literal["any", "all"] = "any",
    ) -> pd.DataFrame:
        """
        去除空值

        Args:
            df: 输入的DataFrame
            columns: 需要检查的列名列表，None表示检查所有列
            how: 'any'表示只要有一个空值就删除行，'all'表示所有值都为空才删除行

        Returns:
            清理后的DataFrame
        """
        logger.info(f"Removing null values from dataframe, shape before: {df.shape}")

        if columns:
            result = df.dropna(subset=columns, how=how)
        else:
            result = df.dropna(how=how)

        logger.info(f"Shape after removing nulls: {result.shape}")
        return result

    def remove_duplicates(
        self,
        df: pd.DataFrame,
        columns: Optional[List[str]] = None,
        keep: Literal["first", "last", False] = "first",
    ) -> pd.DataFrame:
        """
        去除重复数据

        Args:
            df: 输入的DataFrame
            columns: 用于判断重复的列名列表，None表示使用所有列
            keep: 'first'保留第一个，'last'保留最后一个，False删除所有重复

        Returns:
            去重后的DataFrame
        """
        logger.info(f"Removing duplicates from dataframe, shape before: {df.shape}")

        result = df.drop_duplicates(subset=columns, keep=keep)

        logger.info(f"Shape after removing duplicates: {result.shape}")
        return result

    def convert_types(self, df: pd.DataFrame, type_mapping: dict) -> pd.DataFrame:
        """
        数据类型转换

        Args:
            df: 输入的DataFrame
            type_mapping: 列名到目标类型的映射，例如 {'age': int, 'price': float}

        Returns:
            类型转换后的DataFrame
        """
        logger.info(f"Converting data types for columns: {list(type_mapping.keys())}")

        result = df.copy()

        for col, dtype in type_mapping.items():
            if col in result.columns:
                try:
                    result[col] = result[col].astype(dtype)
                    logger.info(f"Column '{col}' converted to {dtype}")
                except Exception as e:
                    logger.error(
                        f"Failed to convert column '{col}' to {dtype}: {str(e)}"
                    )
                    raise
            else:
                logger.warning(f"Column '{col}' not found in dataframe")

        return result

    def detect_outliers(
        self, df: pd.DataFrame, column: str, method: str = "iqr", threshold: float = 1.5
    ) -> pd.Series:
        """
        异常值检测

        Args:
            df: 输入的DataFrame
            column: 需要检测的列名
            method: 检测方法，'iqr'（四分位距）或'zscore'（标准分数）
            threshold: 阈值，IQR方法默认1.5，Z-score方法默认3

        Returns:
            布尔Series，True表示异常值
        """
        logger.info(f"Detecting outliers in column '{column}' using {method} method")

        if column not in df.columns:
            raise ValueError(f"Column '{column}' not found in dataframe")

        data = df[column]

        if method == "iqr":
            # 使用四分位距方法
            Q1 = data.quantile(0.25)
            Q3 = data.quantile(0.75)
            IQR = Q3 - Q1

            lower_bound = Q1 - threshold * IQR
            upper_bound = Q3 + threshold * IQR

            outliers = (data < lower_bound) | (data > upper_bound)

        elif method == "zscore":
            # 使用Z-score方法
            if threshold == 1.5:  # 如果使用默认IQR阈值，改为Z-score默认阈值
                threshold = 3

            mean = data.mean()
            std = data.std()

            if std == 0:
                logger.warning(f"Standard deviation is 0 for column '{column}'")
                return pd.Series([False] * len(data), index=data.index)

            z_scores = np.abs((data - mean) / std)
            outliers = z_scores > threshold

        else:
            raise ValueError(f"Unknown method: {method}. Use 'iqr' or 'zscore'")

        outlier_count = outliers.sum()
        logger.info(f"Found {outlier_count} outliers in column '{column}'")

        return outliers

    def remove_outliers(
        self, df: pd.DataFrame, column: str, method: str = "iqr", threshold: float = 1.5
    ) -> pd.DataFrame:
        """
        移除异常值

        Args:
            df: 输入的DataFrame
            column: 需要处理的列名
            method: 检测方法，'iqr'或'zscore'
            threshold: 阈值

        Returns:
            移除异常值后的DataFrame
        """
        outliers = self.detect_outliers(df, column, method, threshold)
        result: pd.DataFrame = df[~outliers].copy()  # type: ignore[assignment]

        logger.info(f"Removed {outliers.sum()} outliers from dataframe")
        return result

    def clean_data(
        self,
        df: pd.DataFrame,
        remove_nulls: bool = True,
        remove_duplicates: bool = True,
        type_mapping: Optional[dict] = None,
        outlier_columns: Optional[List[str]] = None,
    ) -> pd.DataFrame:
        """
        一键数据清洗

        Args:
            df: 输入的DataFrame
            remove_nulls: 是否去除空值
            remove_duplicates: 是否去重
            type_mapping: 类型转换映射
            outlier_columns: 需要检测异常值的列名列表

        Returns:
            清洗后的DataFrame
        """
        logger.info(f"Starting data cleaning, input shape: {df.shape}")

        result = df.copy()

        # 去除空值
        if remove_nulls:
            result = self.remove_nulls(result)

        # 去重
        if remove_duplicates:
            result = self.remove_duplicates(result)

        # 类型转换
        if type_mapping:
            result = self.convert_types(result, type_mapping)

        # 异常值处理
        if outlier_columns:
            for col in outlier_columns:
                if col in result.columns:
                    result = self.remove_outliers(result, col)

        logger.info(f"Data cleaning completed, output shape: {result.shape}")
        return result


# 便捷函数
def clean_data(df: pd.DataFrame, **kwargs) -> pd.DataFrame:
    """
    便捷的数据清洗函数

    Args:
        df: 输入的DataFrame
        **kwargs: 传递给DataCleaner.clean_data的参数

    Returns:
        清洗后的DataFrame
    """
    cleaner = DataCleaner()
    return cleaner.clean_data(df, **kwargs)
