"""Alpha research tools: IC, signal decay, quantile returns, breadth, transfer coefficient.

Implements the diagnostic toolkit from Grinold-Kahn's fundamental law:
IR = IC x sqrt(Breadth) x TC
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np
from pydantic import BaseModel, Field

EPSILON: float = 1e-15


class ICResult(BaseModel):
    ic: float = 0.0
    ic_t_stat: float = 0.0
    ic_p_value: float = 1.0
    hit_rate: float = 0.0
    method: str = "rank"
    n_observations: int = 0


class DecayResult(BaseModel):
    ic_by_horizon: Dict[int, float] = Field(default_factory=dict)
    half_life_days: Optional[float] = None
    optimal_horizon: int = 0


class QuantileResult(BaseModel):
    quantile_returns: Dict[int, float] = Field(default_factory=dict)
    long_short_spread: float = 0.0
    monotonicity_score: float = 0.0
    n_quantiles: int = 5


class TurnoverMetrics(BaseModel):
    autocorrelation_by_lag: Dict[int, float] = Field(default_factory=dict)
    average_turnover: float = 0.0
    implied_holding_period: float = 0.0


def information_coefficient(
    signals: Dict[str, float],
    forward_returns: Dict[str, float],
    *,
    method: str = "rank",
) -> ICResult:
    """Information Coefficient: correlation between signal and forward returns.

    Args:
        signals: ``{asset: signal_value}`` cross-sectional signals.
        forward_returns: ``{asset: forward_return}`` realised returns.
        method: ``"rank"`` (Spearman, default) or ``"pearson"``.
    """
    from scipy import stats as sp_stats

    common = sorted(set(signals.keys()) & set(forward_returns.keys()))
    n = len(common)
    if n < 3:
        return ICResult(n_observations=n, method=method)

    s = np.array([signals[a] for a in common], dtype=np.float64)
    r = np.array([forward_returns[a] for a in common], dtype=np.float64)

    if method == "rank":
        corr, pval = sp_stats.spearmanr(s, r)
    elif method == "pearson":
        corr, pval = sp_stats.pearsonr(s, r)
    else:
        raise ValueError(f"Unknown method '{method}'. Use: rank, pearson")

    t_stat = corr * np.sqrt((n - 2) / max(1.0 - corr ** 2, EPSILON))
    hit = np.mean(np.sign(s) == np.sign(r))

    return ICResult(
        ic=round(float(corr), 6),
        ic_t_stat=round(float(t_stat), 4),
        ic_p_value=round(float(pval), 6),
        hit_rate=round(float(hit), 4),
        method=method,
        n_observations=n,
    )


def rolling_ic(
    signals_series: List[Dict[str, float]],
    returns_series: List[Dict[str, float]],
    *,
    window: int = 60,
    method: str = "rank",
) -> List[float]:
    """Rolling IC over time.

    Args:
        signals_series: List of cross-sectional signal dicts (one per period).
        returns_series: List of forward return dicts (one per period).
        window: Rolling window size.
        method: IC method (``"rank"`` or ``"pearson"``).

    Returns:
        List of rolling IC values.
    """
    n = min(len(signals_series), len(returns_series))
    ics: List[float] = []

    for i in range(window, n):
        window_ics = []
        for j in range(i - window, i):
            result = information_coefficient(signals_series[j], returns_series[j], method=method)
            window_ics.append(result.ic)
        ics.append(round(float(np.mean(window_ics)), 6))

    return ics


def signal_decay(
    signals: Dict[str, float],
    returns_at_horizons: Dict[int, Dict[str, float]],
    *,
    method: str = "rank",
) -> DecayResult:
    """IC at multiple forward horizons to measure signal half-life.

    Args:
        signals: ``{asset: signal_value}`` cross-sectional signals.
        returns_at_horizons: ``{horizon_days: {asset: return}}`` — returns
            measured at different holding periods (e.g. 1, 5, 10, 20, 60 days).
        method: IC method.
    """
    ic_by_horizon: Dict[int, float] = {}
    for horizon, fwd_rets in sorted(returns_at_horizons.items()):
        result = information_coefficient(signals, fwd_rets, method=method)
        ic_by_horizon[horizon] = result.ic

    horizons = sorted(ic_by_horizon.keys())
    ics = [ic_by_horizon[h] for h in horizons]

    half_life = None
    if len(ics) > 1 and abs(ics[0]) > EPSILON:
        half_ic = abs(ics[0]) / 2.0
        for i in range(1, len(ics)):
            if abs(ics[i]) <= half_ic:
                prev_h = horizons[i - 1]
                curr_h = horizons[i]
                prev_ic = abs(ics[i - 1])
                curr_ic = abs(ics[i])
                if prev_ic - curr_ic > EPSILON:
                    half_life = prev_h + (half_ic - prev_ic) / (curr_ic - prev_ic) * (curr_h - prev_h)
                else:
                    half_life = float(curr_h)
                break

    optimal = horizons[np.argmax(np.abs(ics))] if ics else 0

    return DecayResult(
        ic_by_horizon=ic_by_horizon,
        half_life_days=round(half_life, 1) if half_life is not None else None,
        optimal_horizon=optimal,
    )


def quantile_returns(
    signals: Dict[str, float],
    forward_returns: Dict[str, float],
    *,
    n_quantiles: int = 5,
) -> QuantileResult:
    """Sort assets into quantiles by signal, measure average return per quantile.

    Args:
        signals: ``{asset: signal_value}``.
        forward_returns: ``{asset: forward_return}``.
        n_quantiles: Number of quantiles (default 5 = quintiles).
    """
    common = sorted(set(signals.keys()) & set(forward_returns.keys()))
    n = len(common)
    if n < n_quantiles:
        return QuantileResult(n_quantiles=n_quantiles)

    s = np.array([signals[a] for a in common], dtype=np.float64)
    r = np.array([forward_returns[a] for a in common], dtype=np.float64)

    sorted_idx = np.argsort(s)
    quantile_size = n // n_quantiles

    q_returns: Dict[int, float] = {}
    for q in range(n_quantiles):
        start = q * quantile_size
        end = start + quantile_size if q < n_quantiles - 1 else n
        indices = sorted_idx[start:end]
        q_returns[q + 1] = round(float(np.mean(r[indices])) * 100, 4)

    ls_spread = q_returns.get(n_quantiles, 0.0) - q_returns.get(1, 0.0)

    vals = list(q_returns.values())
    mono = 0
    for i in range(1, len(vals)):
        if vals[i] > vals[i - 1]:
            mono += 1
        elif vals[i] < vals[i - 1]:
            mono -= 1
    mono_score = mono / max(len(vals) - 1, 1)

    return QuantileResult(
        quantile_returns=q_returns,
        long_short_spread=round(ls_spread, 4),
        monotonicity_score=round(mono_score, 4),
        n_quantiles=n_quantiles,
    )


def signal_turnover(
    signals_series: List[Dict[str, float]],
    *,
    lags: Optional[List[int]] = None,
) -> TurnoverMetrics:
    """Signal autocorrelation and turnover metrics.

    Args:
        signals_series: List of cross-sectional signal dicts over time.
        lags: Lags to compute autocorrelation for (default [1,5,10,20]).
    """
    if lags is None:
        lags = [1, 5, 10, 20]

    n = len(signals_series)
    if n < 2:
        return TurnoverMetrics()

    all_assets = set()
    for s in signals_series:
        all_assets.update(s.keys())
    assets = sorted(all_assets)

    matrix = np.zeros((n, len(assets)))
    for i, s in enumerate(signals_series):
        for j, a in enumerate(assets):
            matrix[i, j] = s.get(a, 0.0)

    autocorr: Dict[int, float] = {}
    for lag in lags:
        if lag >= n:
            continue
        corrs = []
        for j in range(len(assets)):
            col = matrix[:, j]
            if np.std(col, ddof=1) < EPSILON:
                continue
            c = np.corrcoef(col[:-lag], col[lag:])[0, 1]
            if not np.isnan(c):
                corrs.append(c)
        autocorr[lag] = round(float(np.mean(corrs)), 6) if corrs else 0.0

    turnovers = []
    for i in range(1, n):
        diff = matrix[i] - matrix[i - 1]
        turnovers.append(float(np.sum(np.abs(diff))))

    avg_turnover = float(np.mean(turnovers)) if turnovers else 0.0
    holding_period = 1.0 / max(avg_turnover, EPSILON) if avg_turnover > EPSILON else float("inf")

    return TurnoverMetrics(
        autocorrelation_by_lag=autocorr,
        average_turnover=round(avg_turnover, 6),
        implied_holding_period=round(holding_period, 2),
    )


def breadth(
    signals: Dict[str, float],
    *,
    min_signal: float = 0.0,
) -> float:
    """Effective breadth: number of independent bets.

    Approximation: ``BR = N * (1 - avg_pairwise_signal_correlation)``.
    When signals are independent, BR = N.

    Args:
        signals: ``{asset: signal_value}``.
        min_signal: Minimum absolute signal to count as an active bet.
    """
    active = {k: v for k, v in signals.items() if abs(v) > min_signal}
    return float(len(active))


def transfer_coefficient(
    alphas: Dict[str, float],
    unconstrained_weights: Dict[str, float],
    constrained_weights: Dict[str, float],
) -> float:
    """Transfer coefficient: how much alpha survives constraints.

    TC = corr(alpha, w_c * alpha) / corr(alpha, w_u * alpha)

    TC = 1 means constraints don't hurt alpha capture.
    TC < 1 means constraints are eroding alpha.

    Args:
        alphas: ``{asset: alpha_estimate}``.
        unconstrained_weights: Optimal weights without constraints.
        constrained_weights: Optimal weights with constraints.
    """
    common = sorted(set(alphas.keys()) & set(unconstrained_weights.keys()) & set(constrained_weights.keys()))
    if len(common) < 2:
        return 1.0

    a = np.array([alphas[c] for c in common], dtype=np.float64)
    wu = np.array([unconstrained_weights[c] for c in common], dtype=np.float64)
    wc = np.array([constrained_weights[c] for c in common], dtype=np.float64)

    active_alpha_c = wc * a
    active_alpha_u = wu * a

    std_a = np.std(a, ddof=1)
    std_u = np.std(active_alpha_u, ddof=1)
    std_c = np.std(active_alpha_c, ddof=1)

    if std_a < EPSILON or std_u < EPSILON or std_c < EPSILON:
        return 1.0

    corr_c = float(np.corrcoef(a, active_alpha_c)[0, 1])
    corr_u = float(np.corrcoef(a, active_alpha_u)[0, 1])

    return round(corr_c / max(abs(corr_u), EPSILON), 4)
