"""
ML — Anomaly detection: IQR, Z-score, Isolation Forest, and LOF.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np
import pandas as pd
from sklearn.ensemble import IsolationForest
from sklearn.neighbors import LocalOutlierFactor
from sklearn.preprocessing import StandardScaler

from bizanalyticsai.core.dataframe import BizFrame


@dataclass
class AnomalyResult:
    df:          BizFrame      # original df + anomaly_score, is_anomaly columns
    n_anomalies: int
    anomaly_pct: float
    method:      str

    def anomalies(self) -> BizFrame:
        return BizFrame(self.df[self.df["is_anomaly"]])

    def summary(self) -> str:
        return (
            f"Method: {self.method} | Anomalies: {self.n_anomalies} "
            f"({self.anomaly_pct:.2f}% of {len(self.df)} rows)"
        )


def detect_anomalies(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    method: Literal["iqr", "zscore", "isolation_forest", "lof"] = "isolation_forest",
    contamination: float = 0.05,
    **kwargs,
) -> AnomalyResult:
    """
    Detect anomalies in a DataFrame.

    Methods:
    - iqr:               IQR-based (fast, univariate)
    - zscore:            Z-score (assumes normality, univariate)
    - isolation_forest:  Tree-based (multivariate, handles high-dimensional data)
    - lof:               Local Outlier Factor (density-based, multivariate)

    Args:
        df:            Input DataFrame.
        cols:          Columns to analyze. Defaults to all numeric.
        method:        Detection method.
        contamination: Expected proportion of anomalies (0.0–0.5).
        **kwargs:      Extra args passed to sklearn estimator.

    Returns:
        AnomalyResult with is_anomaly and anomaly_score columns.
    """
    num_cols = cols or list(df.select_dtypes(include="number").columns)
    out = df.copy()
    X = out[num_cols].fillna(out[num_cols].median()).values

    if method == "iqr":
        scores = np.zeros(len(out))
        flags  = np.zeros(len(out), dtype=bool)
        for i, col in enumerate(num_cols):
            s  = out[col].fillna(out[col].median())
            q1, q3 = s.quantile(0.25), s.quantile(0.75)
            iqr = q3 - q1
            lo, hi = q1 - 3 * iqr, q3 + 3 * iqr
            col_flag = (s < lo) | (s > hi)
            flags |= col_flag
            scores += np.where(col_flag, np.abs((s - s.mean()) / (s.std() + 1e-9)), 0)
        out["anomaly_score"] = scores
        out["is_anomaly"]    = flags

    elif method == "zscore":
        Xs = StandardScaler().fit_transform(X)
        z_max = np.abs(Xs).max(axis=1)
        threshold = kwargs.get("threshold", 3.0)
        out["anomaly_score"] = z_max
        out["is_anomaly"]    = z_max > threshold

    elif method == "isolation_forest":
        model = IsolationForest(contamination=contamination, random_state=42, **kwargs)
        preds = model.fit_predict(X)
        out["anomaly_score"] = -model.score_samples(X)  # higher = more anomalous
        out["is_anomaly"]    = preds == -1

    elif method == "lof":
        n_neighbors = kwargs.pop("n_neighbors", 20)
        model = LocalOutlierFactor(n_neighbors=n_neighbors, contamination=contamination, **kwargs)
        preds = model.fit_predict(X)
        out["anomaly_score"] = -model.negative_outlier_factor_
        out["is_anomaly"]    = preds == -1

    else:
        raise ValueError(f"Unknown method: {method}")

    n_anomalies = int(out["is_anomaly"].sum())
    return AnomalyResult(
        df=BizFrame(out),
        n_anomalies=n_anomalies,
        anomaly_pct=round(n_anomalies / max(len(out), 1) * 100, 2),
        method=method,
    )


def time_series_anomalies(
    series: pd.Series,
    window: int = 7,
    threshold_std: float = 3.0,
) -> BizFrame:
    """
    Detect anomalies in a time series using rolling Z-score.

    Args:
        series:        Numeric time series (index should be datetime).
        window:        Rolling window size.
        threshold_std: Anomaly if |z-score| > threshold.

    Returns:
        BizFrame with original values, rolling mean, rolling std, z_score, is_anomaly.
    """
    roll_mean = series.rolling(window=window, center=True, min_periods=1).mean()
    roll_std  = series.rolling(window=window, center=True, min_periods=1).std().fillna(1)
    z_score   = (series - roll_mean) / roll_std

    result = pd.DataFrame({
        "value":       series,
        "rolling_mean": roll_mean,
        "rolling_std":  roll_std,
        "z_score":      z_score,
        "is_anomaly":   z_score.abs() > threshold_std,
    })
    return BizFrame(result)
