"""
validation/metrics.py — Shared statistical metrics for factor and regime validation.
"""

import numpy as np
import pandas as pd
from scipy import stats
from typing import Optional


def information_coefficient(factor: pd.Series, forward_returns: pd.Series) -> float:
    """Spearman rank IC between factor values and forward returns."""
    mask = factor.notna() & forward_returns.notna()
    if mask.sum() < 10:
        return 0.0
    ic, _ = stats.spearmanr(factor[mask], forward_returns[mask])
    return float(ic) if not np.isnan(ic) else 0.0


def icir(ic_series: pd.Series) -> float:
    """Information Coefficient IR = mean(IC) / std(IC)."""
    if len(ic_series) < 5 or ic_series.std() == 0:
        return 0.0
    return float(ic_series.mean() / ic_series.std())


def hit_ratio(factor: pd.Series, forward_returns: pd.Series) -> float:
    """Fraction of times factor direction matches return direction."""
    mask = factor.notna() & forward_returns.notna()
    if mask.sum() < 10:
        return 0.5
    # Compare sign of factor deviation from median vs sign of returns
    factor_median = factor[mask].median()
    factor_sign = np.sign(factor[mask] - factor_median)
    return_sign = np.sign(forward_returns[mask])
    return float((factor_sign == return_sign).mean())


def quantile_spread(factor: pd.Series, forward_returns: pd.Series,
                    n_quantiles: int = 5) -> float:
    """Top vs bottom quantile return spread (分层回测)."""
    mask = factor.notna() & forward_returns.notna()
    if mask.sum() < n_quantiles * 3:
        return 0.0
    f = factor[mask]
    r = forward_returns[mask]
    labels = pd.qcut(f, n_quantiles, labels=False, duplicates="drop")
    top_ret = r[labels == labels.max()].mean()
    bot_ret = r[labels == labels.min()].mean()
    return float(top_ret - bot_ret)


def lead_lag_ic(factor: pd.Series, returns: pd.Series,
                max_lag: int = 14) -> dict:
    """Find the best leading/trailing relationship by computing IC at each lag."""
    results = {}
    for lag in range(-max_lag, max_lag + 1):
        if lag < 0:
            shifted = factor.shift(abs(lag))
            ic = information_coefficient(shifted, returns)
            results[f"lead_{abs(lag)}d"] = ic
        elif lag > 0:
            shifted = returns.shift(lag)
            ic = information_coefficient(factor, shifted)
            results[f"lag_{lag}d"] = ic
        else:
            ic = information_coefficient(factor, returns)
            results["sync"] = ic

    # Find best lead period
    lead_ics = {k: v for k, v in results.items() if k.startswith("lead_")}
    best_lead = max(lead_ics, key=lead_ics.get) if lead_ics else "sync"
    best_ic = lead_ics.get(best_lead, results.get("sync", 0))

    return {
        "best_lead": best_lead,
        "best_ic": best_ic,
        "ic_curve": results,
        "is_leading": best_lead.startswith("lead_") and abs(best_ic) > 0.03,
        "lead_days": int(best_lead.split("_")[1].rstrip("d")) if best_lead.startswith("lead_") else 0,
    }


def mutual_information(factor: pd.Series, labels: pd.Series,
                       n_bins: int = 10) -> float:
    """Mutual information between factor (binned) and discrete regime labels."""
    mask = factor.notna() & labels.notna()
    if mask.sum() < 20:
        return 0.0
    f = factor[mask]
    l = labels[mask]
    try:
        f_binned = pd.qcut(f, n_bins, labels=False, duplicates="drop")
    except ValueError:
        f_binned = pd.cut(f, n_bins, labels=False)
    mi = 0.0
    for fi in range(n_bins):
        p_f = (f_binned == fi).mean()
        if p_f == 0:
            continue
        for li in l.unique():
            p_l = (l == li).mean()
            p_joint = ((f_binned == fi) & (l == li)).mean()
            if p_joint > 0:
                mi += p_joint * np.log(p_joint / (p_f * p_l))
    return float(mi)


def kl_divergence(factor: pd.Series, labels: pd.Series,
                  regime_a: str, regime_b: str, n_bins: int = 10) -> float:
    """KL divergence between factor distributions in two regimes."""
    mask_a = (labels == regime_a) & factor.notna()
    mask_b = (labels == regime_b) & factor.notna()
    if mask_a.sum() < 10 or mask_b.sum() < 10:
        return 0.0
    try:
        hist_a, edges = np.histogram(factor[mask_a], bins=n_bins, density=True)
        hist_b, _ = np.histogram(factor[mask_b], bins=edges, density=True)
    except ValueError:
        return 0.0
    hist_a = np.clip(hist_a, 1e-10, None)
    hist_b = np.clip(hist_b, 1e-10, None)
    return float((hist_a * np.log(hist_a / hist_b)).sum())


def anova_f_score(factor: pd.Series, labels: pd.Series) -> float:
    """ANOVA F-statistic: how well factor separates different regimes."""
    mask = factor.notna() & labels.notna()
    if mask.sum() < 20:
        return 0.0
    groups = [factor[mask][labels[mask] == lbl] for lbl in labels[mask].unique()]
    groups = [g for g in groups if len(g) > 1]
    if len(groups) < 2:
        return 0.0
    f_stat, _ = stats.f_oneway(*groups)
    return float(f_stat) if not np.isnan(f_stat) else 0.0


def transition_matrix(labels: pd.Series) -> pd.DataFrame:
    """Compute Markov transition matrix from regime sequence."""
    unique = sorted(labels.dropna().unique())
    n = len(unique)
    matrix = np.zeros((n, n))
    seq = labels.dropna().values
    for i in range(len(seq) - 1):
        from_idx = unique.index(seq[i])
        to_idx = unique.index(seq[i + 1])
        matrix[from_idx][to_idx] += 1

    # Row-normalize
    row_sums = matrix.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1
    matrix = matrix / row_sums

    return pd.DataFrame(matrix, index=unique, columns=unique)


def regime_duration_stats(labels: pd.Series) -> dict:
    """Compute average duration, flip rate, state entropy for regime sequence."""
    seq = labels.dropna().values
    if len(seq) < 2:
        return {"avg_duration": 0, "flip_rate": 0, "state_entropy": 0, "n_days": len(seq)}

    # Count durations
    durations = []
    current = seq[0]
    count = 1
    flips = 0
    for i in range(1, len(seq)):
        if seq[i] == current:
            count += 1
        else:
            durations.append(count)
            current = seq[i]
            count = 1
            flips += 1
    durations.append(count)

    avg_dur = float(np.mean(durations)) if durations else 0
    flip_rate = flips / len(seq)

    # State entropy
    _, counts = np.unique(seq, return_counts=True)
    probs = counts / counts.sum()
    entropy = float(-(probs * np.log2(probs + 1e-10)).sum())

    return {
        "avg_duration": round(avg_dur, 1),
        "flip_rate": round(flip_rate, 3),
        "state_entropy": round(entropy, 3),
        "n_days": len(seq),
    }