from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from pandas.api.types import is_bool_dtype, is_numeric_dtype

from featureleak.report import LeakIssue

try:
    from sklearn.feature_selection import mutual_info_classif
    from sklearn.metrics import accuracy_score
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import LabelEncoder
    from sklearn.tree import DecisionTreeClassifier

    SKLEARN_AVAILABLE = True
except Exception:  # pragma: no cover
    SKLEARN_AVAILABLE = False


TARGET_NAME_TOKENS = (
    "target",
    "label",
    "outcome",
    "response",
    "sales",
    "revenue",
    "profit",
    "default",
    "churn",
    "fraud",
    "grade",
    "class",
    "category",
    "pass",
    "fail",
)

FUTURE_NAME_TOKENS = (
    "future",
    "next",
    "lead",
    "ahead",
    "t+1",
    "t_plus_1",
    "tomorrow",
)

EXACT_COPY_NAME_TOKENS = (
    "actual",
    "observed",
    "realized",
    "final",
)


def _safe_abs_corr(feature: pd.Series, target: pd.Series) -> float | None:
    try:
        x = pd.to_numeric(feature, errors="coerce")
        y = pd.to_numeric(target, errors="coerce")
        mask = x.notna() & y.notna()
        if int(mask.sum()) < 2:
            return None
        corr = x[mask].corr(y[mask])
        if pd.isna(corr):
            return None
        return float(abs(corr))
    except Exception:
        return None


def _equality_ratio(feature: pd.Series, target: pd.Series) -> float:
    try:
        lhs = feature.astype("object").where(feature.notna(), "__NaN__").astype(str)
        rhs = target.astype("object").where(target.notna(), "__NaN__").astype(str)
        return float((lhs == rhs).mean())
    except Exception:
        return 0.0


def _shift_match_ratio(feature: pd.Series, target: pd.Series, shift: int) -> float:
    try:
        shifted = target.shift(shift)
        lhs = feature.astype("object").where(feature.notna(), "__NaN__").astype(str)
        rhs = shifted.astype("object").where(shifted.notna(), "__NaN__").astype(str)
        valid = shifted.notna()
        if int(valid.sum()) == 0:
            return 0.0
        return float((lhs[valid] == rhs[valid]).mean())
    except Exception:
        return 0.0


def _suspicious_name(column_name: str, target_name: str) -> dict[str, Any]:
    col = str(column_name).lower()
    target_lower = str(target_name).lower()

    matched: list[str] = []

    if target_lower and target_lower in col and target_lower != col:
        matched.append("target_name")

    matched.extend([tok for tok in TARGET_NAME_TOKENS if tok in col])
    matched.extend([tok for tok in FUTURE_NAME_TOKENS if tok in col])
    matched.extend([tok for tok in EXACT_COPY_NAME_TOKENS if tok in col])

    return {
        "matched_tokens": sorted(set(matched)),
        "has_target_semantics": any(tok in col for tok in TARGET_NAME_TOKENS)
        or (target_lower in col and target_lower != col),
        "has_future_semantics": any(tok in col for tok in FUTURE_NAME_TOKENS),
        "has_exact_copy_semantics": any(tok in col for tok in EXACT_COPY_NAME_TOKENS),
    }


def _should_run_expensive_checks(
    feature: pd.Series,
    target: pd.Series,
    col: str,
    target_name: str,
    *,
    unique_ratio_gate: float = 0.999,
) -> bool:
    """
    Cheap screening stage before more expensive statistical / ML checks.
    """
    try:
        name_info = _suspicious_name(col, target_name)
        if name_info["matched_tokens"]:
            return True

        non_null = feature.dropna()
        if non_null.empty:
            return False

        unique_ratio = non_null.nunique(dropna=False) / max(len(non_null), 1)
        if unique_ratio >= unique_ratio_gate and not is_numeric_dtype(feature):
            # Extremely high-cardinality non-numeric columns are usually IDs.
            return False

        if is_numeric_dtype(feature) or is_bool_dtype(feature):
            return True

        if non_null.nunique(dropna=False) <= 1000:
            return True
    except Exception:
        return True

    return False


def _encode_categorical_series(series: pd.Series) -> tuple[np.ndarray, dict[str, int]] | tuple[None, None]:
    try:
        values = (
            series.astype("object")
            .where(series.notna(), "__MISSING__")
            .astype(str)
        )
        categories = pd.Index(sorted(values.unique()))
        mapping = {cat: idx for idx, cat in enumerate(categories)}
        encoded = values.map(mapping).to_numpy(dtype=int)
        return encoded, mapping
    except Exception:
        return None, None


def _safe_mutual_info_classif(feature: pd.Series, target: pd.Series) -> float | None:
    """
    Mutual information between a single feature and a categorical target.
    Supports numeric or categorical feature.
    """
    if not SKLEARN_AVAILABLE:
        return None

    try:
        mask = feature.notna() & target.notna()
        if int(mask.sum()) < 20:
            return None

        x = feature[mask]
        y = target[mask]

        y_encoded, _ = _encode_categorical_series(y)
        if y_encoded is None or len(np.unique(y_encoded)) < 2:
            return None

        if is_numeric_dtype(x) and not is_bool_dtype(x):
            x_num = pd.to_numeric(x, errors="coerce")
            valid = x_num.notna()
            if int(valid.sum()) < 20:
                return None
            X = x_num[valid].to_numpy(dtype=float).reshape(-1, 1)
            y_final = y_encoded[valid.to_numpy()]
            mi = mutual_info_classif(
                X,
                y_final,
                discrete_features=False,
                random_state=42,
            )[0]
            return float(mi)

        x_encoded, _ = _encode_categorical_series(x)
        if x_encoded is None:
            return None

        X = x_encoded.reshape(-1, 1)
        mi = mutual_info_classif(
            X,
            y_encoded,
            discrete_features=True,
            random_state=42,
        )[0]
        return float(mi)
    except Exception:
        return None


def _safe_cramers_v(feature: pd.Series, target: pd.Series) -> float | None:
    """
    Bias-corrected Cramér's V for categorical-categorical association.
    """
    try:
        lhs = feature.astype("object").where(feature.notna(), "__MISSING__").astype(str)
        rhs = target.astype("object").where(target.notna(), "__MISSING__").astype(str)

        if lhs.nunique(dropna=False) < 2 or rhs.nunique(dropna=False) < 2:
            return None

        table = pd.crosstab(lhs, rhs)
        if table.empty:
            return None

        observed = table.to_numpy(dtype=float)
        n = observed.sum()
        if n <= 0:
            return None

        row_sum = observed.sum(axis=1, keepdims=True)
        col_sum = observed.sum(axis=0, keepdims=True)
        expected = row_sum @ col_sum / n

        with np.errstate(divide="ignore", invalid="ignore"):
            chi2 = np.nansum((observed - expected) ** 2 / expected)

        phi2 = chi2 / n
        r, k = observed.shape

        phi2corr = max(0.0, phi2 - ((k - 1) * (r - 1)) / max(n - 1, 1))
        rcorr = r - ((r - 1) ** 2) / max(n - 1, 1)
        kcorr = k - ((k - 1) ** 2) / max(n - 1, 1)

        denom = min(kcorr - 1, rcorr - 1)
        if denom <= 0:
            return None

        return float(np.sqrt(phi2corr / denom))
    except Exception:
        return None


def _safe_tree_probe_accuracy(feature: pd.Series, target: pd.Series) -> float | None:
    """
    Tiny single-feature decision tree to detect near-deterministic mappings,
    e.g. final_grade -> grade_category or pass_fail -> grade_category.
    """
    if not SKLEARN_AVAILABLE:
        return None

    try:
        mask = feature.notna() & target.notna()
        if int(mask.sum()) < 40:
            return None

        x = feature[mask]
        y = target[mask]

        if y.nunique(dropna=True) < 2:
            return None

        y_enc, _ = _encode_categorical_series(y)
        if y_enc is None:
            return None

        if is_numeric_dtype(x) and not is_bool_dtype(x):
            x_num = pd.to_numeric(x, errors="coerce")
            valid = x_num.notna()
            if int(valid.sum()) < 40:
                return None
            X = x_num[valid].to_numpy(dtype=float).reshape(-1, 1)
            y_final = y_enc[valid.to_numpy()]
        else:
            x_enc, _ = _encode_categorical_series(x)
            if x_enc is None:
                return None
            X = x_enc.reshape(-1, 1).astype(float)
            y_final = y_enc

        if len(np.unique(y_final)) < 2:
            return None

        X_train, X_test, y_train, y_test = train_test_split(
            X,
            y_final,
            test_size=0.30,
            random_state=42,
            stratify=y_final,
        )

        clf = DecisionTreeClassifier(
            max_depth=3,
            min_samples_leaf=5,
            random_state=42,
        )
        clf.fit(X_train, y_train)
        preds = clf.predict(X_test)
        return float(accuracy_score(y_test, preds))
    except Exception:
        return None


def detect_target_leakage(
    df: pd.DataFrame,
    target: str,
    corr_threshold: float = 0.98,
    equality_threshold: float = 0.98,
    shift_threshold: float = 0.98,
    mutual_info_threshold: float = 0.50,
    cramers_v_threshold: float = 0.90,
    tree_accuracy_threshold: float = 0.95,
    max_expensive_columns: int = 100,
) -> list[LeakIssue]:
    """
    Detect direct leakage from target-derived features.

    This version handles:
    - numeric target <-> numeric feature via correlation
    - categorical target <- numeric feature via mutual information / tree probe
    - categorical target <-> categorical feature via Cramér's V / mutual information
    - exact and shifted copies for all dtypes

    Performance-aware behavior:
    - cheap screening on all columns
    - expensive checks only on promising columns
    """
    issues: list[LeakIssue] = []

    if df.empty or target not in df.columns:
        return issues
    
    target_series = df[target]
    target_is_numeric = is_numeric_dtype(target_series) and not is_bool_dtype(target_series)

    candidate_columns: list[str] = []
    precomputed_name_info: dict[str, dict[str, Any]] = {}

    for col in df.columns:
        if col == target:
            continue

        name_info = _suspicious_name(col, target)
        precomputed_name_info[col] = name_info

        if _should_run_expensive_checks(df[col], target_series, col, target):
            candidate_columns.append(col)

    candidate_columns = candidate_columns[:max_expensive_columns]

    for col in candidate_columns:
        feature = df[col]
        details: dict[str, Any] = {}
        reasons: list[str] = []
        severity: str | None = None

        name_info = precomputed_name_info[col]
        details["matched_name_tokens"] = name_info["matched_tokens"]

        if name_info["has_target_semantics"]:
            reasons.append("target_semantic_name")
            severity = "medium"

        if name_info["has_future_semantics"]:
            reasons.append("future_semantic_name")
            severity = "medium" if severity is None else severity

        if name_info["has_exact_copy_semantics"]:
            reasons.append("exact_copy_semantic_name")
            severity = "medium" if severity is None else severity

        eq_ratio = _equality_ratio(feature, target_series)
        details["equality_ratio"] = round(eq_ratio, 6)

        if eq_ratio >= equality_threshold:
            reasons.append("direct_target_copy")
            severity = "critical"

        shift_minus_1 = _shift_match_ratio(feature, target_series, -1)
        shift_plus_1 = _shift_match_ratio(feature, target_series, 1)
        details["shift_match_ratio_minus_1"] = round(shift_minus_1, 6)
        details["shift_match_ratio_plus_1"] = round(shift_plus_1, 6)

        if max(shift_minus_1, shift_plus_1) >= shift_threshold:
            reasons.append("shifted_target_copy")
            severity = "critical"

        corr = None
        if target_is_numeric and is_numeric_dtype(feature) and not is_bool_dtype(feature):
            corr = _safe_abs_corr(feature, target_series)
            details["abs_target_corr"] = None if corr is None else round(corr, 6)

            if corr is not None and corr >= corr_threshold:
                reasons.append("near_perfect_target_correlation")
                if severity != "critical":
                    severity = "high"

        mi = _safe_mutual_info_classif(feature, target_series)
        details["mutual_information"] = None if mi is None else round(mi, 6)

        if mi is not None and mi >= mutual_info_threshold:
            reasons.append("strong_target_dependence_mutual_information")
            if severity not in ("critical", "high"):
                severity = "high"

        cramers_v = None
        feature_is_effectively_categorical = (
            (not is_numeric_dtype(feature))
            or is_bool_dtype(feature)
            or feature.dropna().nunique() <= 20
        )
        target_is_effectively_categorical = not target_is_numeric

        if feature_is_effectively_categorical and target_is_effectively_categorical:
            cramers_v = _safe_cramers_v(feature, target_series)
            details["cramers_v"] = None if cramers_v is None else round(cramers_v, 6)

            if cramers_v is not None and cramers_v >= cramers_v_threshold:
                reasons.append("strong_categorical_target_association")
                if severity not in ("critical", "high"):
                    severity = "high"

        tree_acc = None
        if target_is_effectively_categorical:
            tree_acc = _safe_tree_probe_accuracy(feature, target_series)
            details["tree_probe_accuracy"] = None if tree_acc is None else round(tree_acc, 6)

            if tree_acc is not None and tree_acc >= tree_accuracy_threshold:
                reasons.append("single_feature_near_deterministic_prediction")
                if severity == "medium":
                    severity = "high"
                elif severity is None:
                    severity = "high"

        # Escalation rules for suspicious names with moderately strong evidence.
        if severity == "medium":
            if eq_ratio >= 0.90:
                severity = "high"
            elif corr is not None and corr >= 0.90:
                severity = "high"
            elif max(shift_minus_1, shift_plus_1) >= 0.90:
                severity = "high"
            elif mi is not None and mi >= 0.20:
                severity = "high"
            elif cramers_v is not None and cramers_v >= 0.75:
                severity = "high"
            elif tree_acc is not None and tree_acc >= 0.85:
                severity = "high"

        if severity is None:
            continue

        details["reasons"] = sorted(set(reasons))

        issues.append(
            LeakIssue(
                severity=severity,
                check="target_leakage",
                title="Potential target leakage detected",
                summary=(
                    f"Column `{col}` appears strongly tied to target `{target}` and may "
                    f"contain direct or derived label information."
                ),
                affected_columns=[col],
                details=details,
                suggested_fix=(
                    "Drop this feature or recompute it using only information available at "
                    "prediction time, ideally inside the training fold only."
                ),
            )
        )

    return issues