"""
Statistical drift detection for ML data and concept drift.
"""

from typing import Dict, List, Optional
from dataclasses import dataclass, field
from enum import Enum

import pandas as pd
import numpy as np
from scipy import stats
from sklearn.preprocessing import LabelEncoder
from sklearn.ensemble import RandomForestClassifier


class DriftType(Enum):
    """Types of drift detected."""

    DATA_DRIFT = "data_drift"
    CONCEPT_DRIFT = "concept_drift"
    COVARIATE_DRIFT = "covariate_drift"
    CONCEPT_SHIFT = "concept_shift"


@dataclass
class DriftReport:
    """Comprehensive drift detection report."""

    drift_detected: bool
    drift_type: List[DriftType]
    timestamp: str
    column_reports: Dict[str, Dict] = field(default_factory=dict)
    global_score: float = 0.0
    details: Dict = field(default_factory=dict)

    def to_dict(self) -> Dict:
        return {
            "drift_detected": self.drift_detected,
            "drift_type": [t.value for t in self.drift_type],
            "timestamp": self.timestamp,
            "global_score": self.global_score,
            "column_reports": self.column_reports,
            "details": self.details,
        }


class DriftDetector:
    """
    Multi-method drift detection for ML pipelines.
    """

    def __init__(
        self,
        numerical_method: str = "ks",
        categorical_method: str = "chi2",
        threshold: float = 0.05,
        psi_bins: int = 10,
    ):
        self.numerical_method = numerical_method
        self.categorical_method = categorical_method
        self.threshold = threshold
        self.psi_bins = psi_bins
        self._reference_data: Optional[pd.DataFrame] = None
        self._reference_target: Optional[pd.Series] = None

    def fit(self, reference_data: pd.DataFrame, target: Optional[pd.Series] = None):
        """Set reference distribution."""
        self._reference_data = reference_data.copy()
        if target is not None:
            self._reference_target = target.copy()
        else:
            self._reference_target = None
        return self

    def detect(
        self,
        current_data: pd.DataFrame,
        target: Optional[pd.Series] = None,
        columns: Optional[List[str]] = None,
    ) -> DriftReport:
        """
        Detect drift between reference and current data.
        """
        if self._reference_data is None:
            raise ValueError("Detector not fitted. Call .fit() first.")

        columns = columns or self._reference_data.columns.tolist()
        timestamp = pd.Timestamp.now().isoformat()

        column_reports = {}
        drifted_columns = []
        drift_types = []

        for col in columns:
            if col not in current_data.columns:
                continue

            report = self._check_column_drift(col, current_data[col])  # type: ignore
            column_reports[col] = report

            if report.get("drift_detected", False):
                drifted_columns.append(col)

        concept_drift_score = 0.0
        if target is not None and self._reference_target is not None:
            concept_report = self._check_concept_drift(current_data, target)
            concept_drift_score = concept_report.get("score", 0.0)
            if concept_report.get("drift_detected", False):
                drift_types.append(DriftType.CONCEPT_DRIFT)
                column_reports["__concept__"] = concept_report

        global_score = len(drifted_columns) / len(columns) if columns else 0

        drift_detected = len(drifted_columns) > len(columns) * 0.1 or concept_drift_score > 0.1

        if drift_detected and drifted_columns:
            drift_types.append(DriftType.DATA_DRIFT)

        return DriftReport(
            drift_detected=drift_detected,
            drift_type=list(set(drift_types)),
            timestamp=timestamp,
            column_reports=column_reports,
            global_score=global_score,
            details={
                "drifted_columns": drifted_columns,
                "total_columns": len(columns),
                "concept_drift_score": concept_drift_score,
            },
        )

    def _check_column_drift(self, col_name: str, current: pd.Series) -> Dict:
        """Check drift for single column."""
        reference = self._reference_data[col_name]

        is_numeric = pd.api.types.is_numeric_dtype(reference)

        if is_numeric:
            return self._numerical_drift_test(reference, current)
        else:
            return self._categorical_drift_test(reference, current)

    def _numerical_drift_test(self, reference: pd.Series, current: pd.Series) -> Dict:
        """Statistical test for numerical drift."""
        ref_clean = reference.dropna()
        curr_clean = current.dropna()

        if len(ref_clean) == 0 or len(curr_clean) == 0:
            return {"drift_detected": False, "method": "none", "reason": "empty_data"}

        if self.numerical_method == "ks":
            statistic, p_value = stats.ks_2samp(ref_clean, curr_clean)
            drift_detected = p_value < self.threshold

            return {
                "drift_detected": drift_detected,
                "method": "ks",
                "statistic": float(statistic),
                "p_value": float(p_value),
                "threshold": self.threshold,
            }

        elif self.numerical_method == "psi":
            psi_score = self._calculate_psi(ref_clean, curr_clean)
            drift_detected = psi_score > 0.2

            return {
                "drift_detected": drift_detected,
                "method": "psi",
                "psi_score": float(psi_score),
                "threshold": 0.2,
            }

        elif self.numerical_method == "wasserstein":
            from scipy.stats import wasserstein_distance

            distance = wasserstein_distance(ref_clean, curr_clean)
            normalized = distance / (ref_clean.std() + 1e-10)
            drift_detected = normalized > 0.1

            return {
                "drift_detected": drift_detected,
                "method": "wasserstein",
                "distance": float(distance),
                "normalized": float(normalized),
                "threshold": 0.1,
            }

        return {"drift_detected": False, "method": "unknown"}

    def _categorical_drift_test(self, reference: pd.Series, current: pd.Series) -> Dict:
        """Statistical test for categorical drift."""
        ref_counts = reference.value_counts(normalize=True)
        curr_counts = current.value_counts(normalize=True)

        all_cats = set(ref_counts.index) | set(curr_counts.index)

        if self.categorical_method == "chi2":
            ref_freq = [ref_counts.get(cat, 0) * len(reference) for cat in all_cats]
            curr_freq = [curr_counts.get(cat, 0) * len(current) for cat in all_cats]

            try:
                chi2, p_value, dof, expected = stats.chi2_contingency([ref_freq, curr_freq])
                drift_detected = p_value < self.threshold

                return {
                    "drift_detected": drift_detected,
                    "method": "chi2",
                    "chi2": float(chi2),
                    "p_value": float(p_value),
                    "threshold": self.threshold,
                    "new_categories": list(set(curr_counts.index) - set(ref_counts.index)),
                    "missing_categories": list(set(ref_counts.index) - set(curr_counts.index)),
                }
            except Exception as e:
                return {"drift_detected": False, "method": "chi2", "error": str(e)}

        elif self.categorical_method == "psi":
            ref_probs = [ref_counts.get(cat, 0) for cat in all_cats]
            curr_probs = [curr_counts.get(cat, 0) for cat in all_cats]

            psi = sum(
                (p - q) * np.log((p + 1e-10) / (q + 1e-10))
                for p, q in zip(ref_probs, curr_probs)
                if p > 0 or q > 0
            )

            return {
                "drift_detected": psi > 0.2,
                "method": "psi",
                "psi_score": float(psi),
                "new_categories": list(set(curr_counts.index) - set(ref_counts.index)),
            }

        return {"drift_detected": False, "method": "unknown"}

    def _calculate_psi(self, reference: pd.Series, current: pd.Series) -> float:
        """Calculate Population Stability Index."""
        min_val, max_val = reference.min(), reference.max()
        bins = np.linspace(min_val, max_val, self.psi_bins + 1)

        ref_percents = np.histogram(reference, bins=bins)[0] / len(reference)
        curr_percents = np.histogram(current, bins=bins)[0] / len(current)

        ref_percents = ref_percents + 1e-10
        curr_percents = curr_percents + 1e-10

        psi = np.sum((ref_percents - curr_percents) * np.log(ref_percents / curr_percents))
        return float(psi)

    def _check_concept_drift(self, current_data: pd.DataFrame, current_target: pd.Series) -> Dict:
        """
        Detect concept drift using model performance comparison.
        """
        if self._reference_target is None:
            return {"drift_detected": False, "reason": "no_reference_target"}

        try:
            ref_X = self._reference_data.select_dtypes(include=[np.number]).fillna(0)  # type: ignore
            curr_X = current_data[ref_X.columns].fillna(0)

            le = LabelEncoder()
            ref_y = le.fit_transform(self._reference_target)

            clf = RandomForestClassifier(n_estimators=50, max_depth=5, random_state=42)
            clf.fit(ref_X, ref_y)

            ref_score = clf.score(ref_X, ref_y)

            curr_y_transformed: List[int] = []
            for val in current_target:
                if val in le.classes_:
                    curr_y_transformed.append(le.transform([val])[0])
                else:
                    curr_y_transformed.append(-1)

            curr_y_transformed = np.array(curr_y_transformed)
            mask = curr_y_transformed != -1  # type: ignore

            if mask.sum() == 0:  # type: ignore
                return {
                    "drift_detected": True,
                    "reason": "all_unknown_classes",
                    "score": 1.0,
                }

            curr_score = clf.score(curr_X[mask], curr_y_transformed[mask])

            score_drop = ref_score - curr_score
            drift_detected = score_drop > 0.15

            return {
                "drift_detected": drift_detected,
                "method": "model_proxy",
                "reference_score": float(ref_score),
                "current_score": float(curr_score),
                "score_drop": float(score_drop),
                "score": float(score_drop),
            }

        except Exception as e:
            return {"drift_detected": False, "method": "model_proxy", "error": str(e)}

    def get_drift_importance(self, report: DriftReport) -> pd.DataFrame:
        """Rank features by drift importance."""
        rows = []
        for col, details in report.column_reports.items():
            if col.startswith("__"):
                continue

            rows.append(
                {
                    "column": col,
                    "drift_detected": details.get("drift_detected", False),
                    "score": details.get("statistic", details.get("psi_score", 0)),
                    "p_value": details.get("p_value", 1.0),
                    "method": details.get("method", "unknown"),
                }
            )

        df = pd.DataFrame(rows)
        if not df.empty:
            df = df.sort_values("score", ascending=False)
        return df
