"""Performance attribution module for jbqlab.

Analyze where returns came from and decompose performance.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

if TYPE_CHECKING:
    pass

__all__ = [
    "AttributionResult",
    "return_attribution",
    "factor_attribution",
    "timing_attribution",
    "holding_period_analysis",
]


@dataclass
class AttributionResult:
    """Results from performance attribution."""

    total_return: float
    market_contribution: float
    selection_contribution: float
    timing_contribution: float
    residual: float
    details: dict

    def __str__(self) -> str:
        """Human-readable summary."""
        lines = [
            "Performance Attribution",
            "=" * 40,
            f"Total Return:       {self.total_return:>10.2%}",
            "",
            "Decomposition:",
            f"  Market:           {self.market_contribution:>10.2%}",
            f"  Selection:        {self.selection_contribution:>10.2%}",
            f"  Timing:           {self.timing_contribution:>10.2%}",
            f"  Residual:         {self.residual:>10.2%}",
        ]
        return "\n".join(lines)


def return_attribution(
    strategy_returns: pd.Series,
    benchmark_returns: pd.Series,
    signal: pd.Series,
) -> AttributionResult:
    """Decompose returns into timing and selection.

    Args:
        strategy_returns: Strategy return series.
        benchmark_returns: Benchmark return series.
        signal: Signal series (1 = invested, 0 = cash).

    Returns:
        AttributionResult with decomposition.
    """
    # Align all series
    combined = pd.concat([strategy_returns, benchmark_returns, signal], axis=1).dropna()
    combined.columns = ["strategy", "benchmark", "signal"]

    strat = combined["strategy"]
    bench = combined["benchmark"]
    sig = combined["signal"].shift(1).fillna(0)  # Use previous signal for next return

    # Total return
    total_return = (1 + strat).prod() - 1
    bench_return = (1 + bench).prod() - 1

    # Market contribution (if fully invested)
    market_contribution = bench_return

    # Selection: difference when invested
    invested_strat = strat[sig == 1]
    invested_bench = bench[sig == 1]
    if len(invested_strat) > 0:
        selection = (1 + invested_strat).prod() / (1 + invested_bench).prod() - 1
    else:
        selection = 0

    # Timing: value from not being invested during down markets
    not_invested_bench = bench[sig == 0]
    if len(not_invested_bench) > 0:
        avoided_loss = -(not_invested_bench.sum())  # Loss we avoided
        timing = avoided_loss
    else:
        timing = 0

    # Residual
    residual = total_return - market_contribution - selection - timing

    return AttributionResult(
        total_return=total_return,
        market_contribution=market_contribution,
        selection_contribution=selection,
        timing_contribution=timing,
        residual=residual,
        details={
            "periods_invested": int(sig.sum()),
            "total_periods": len(sig),
            "time_in_market": float(sig.mean()),
        },
    )


def factor_attribution(
    strategy_returns: pd.Series,
    factor_returns: pd.DataFrame,
) -> dict:
    """Attribute returns to factors using regression.

    Args:
        strategy_returns: Strategy return series.
        factor_returns: DataFrame of factor returns.

    Returns:
        Dictionary with factor exposures and attribution.
    """
    # Align series
    combined = pd.concat([strategy_returns, factor_returns], axis=1).dropna()
    y = combined.iloc[:, 0].values
    X = combined.iloc[:, 1:].values

    # Add intercept
    X_with_const = np.column_stack([np.ones(len(X)), X])

    # Simple OLS
    try:
        betas = np.linalg.lstsq(X_with_const, y, rcond=None)[0]
    except np.linalg.LinAlgError:
        return {"error": "Regression failed"}

    alpha = betas[0]
    factor_betas = betas[1:]

    # Factor contributions
    factor_names = factor_returns.columns.tolist()
    factor_contrib = {}
    total_factor = 0

    for i, name in enumerate(factor_names):
        contrib = factor_betas[i] * factor_returns[name].mean() * 252  # Annualized
        factor_contrib[name] = {
            "beta": float(factor_betas[i]),
            "contribution": float(contrib),
        }
        total_factor += contrib

    # Residual
    predicted = X_with_const @ betas
    residuals = y - predicted
    r_squared = 1 - (residuals.var() / y.var())

    return {
        "alpha_annualized": float(alpha * 252),
        "factors": factor_contrib,
        "total_factor_contribution": float(total_factor),
        "r_squared": float(r_squared),
        "residual_return": float(strategy_returns.mean() * 252 - total_factor),
    }


def timing_attribution(
    signal: pd.Series,
    market_returns: pd.Series,
) -> dict:
    """Analyze timing skill.

    Measures how well signal predicts market direction.

    Args:
        signal: Signal series.
        market_returns: Market return series.

    Returns:
        Timing analysis results.
    """
    combined = pd.concat([signal, market_returns], axis=1).dropna()
    combined.columns = ["signal", "market"]

    sig = combined["signal"].shift(1).fillna(0)  # Use signal for next period
    mkt = combined["market"]

    # Correct predictions
    up_market = mkt > 0
    down_market = mkt < 0

    # When invested
    correct_up = ((sig == 1) & up_market).sum()
    wrong_up = ((sig == 1) & down_market).sum()

    # When out
    correct_down = ((sig == 0) & down_market).sum()
    wrong_down = ((sig == 0) & up_market).sum()

    total_up = up_market.sum()
    total_down = down_market.sum()

    # Hit rates
    up_hit_rate = correct_up / total_up if total_up > 0 else 0
    down_hit_rate = correct_down / total_down if total_down > 0 else 0

    # Henriksson-Merton timing measure
    # A random timer would have up_hit + down_hit = 1
    timing_measure = up_hit_rate + down_hit_rate - 1

    # Value of timing
    avg_up_return = mkt[up_market].mean()
    avg_down_return = mkt[down_market].mean()
    timing_value = (
        correct_up * avg_up_return +
        correct_down * (-avg_down_return) +
        wrong_up * avg_down_return +
        wrong_down * (-avg_up_return)
    ) / len(mkt)

    return {
        "up_market_hit_rate": float(up_hit_rate),
        "down_market_hit_rate": float(down_hit_rate),
        "timing_measure": float(timing_measure),  # > 0 indicates skill
        "timing_value_per_period": float(timing_value),
        "total_up_markets": int(total_up),
        "total_down_markets": int(total_down),
        "correct_bullish_calls": int(correct_up),
        "correct_bearish_calls": int(correct_down),
    }


def holding_period_analysis(
    signal: pd.Series,
    returns: pd.Series,
) -> dict:
    """Analyze returns by holding period.

    Args:
        signal: Signal series.
        returns: Return series.

    Returns:
        Holding period statistics.
    """
    # Find trades
    trades = []
    in_trade = False
    trade_start = None
    trade_returns = []

    for i, (_date, sig) in enumerate(signal.items()):
        if sig == 1 and not in_trade:
            in_trade = True
            trade_start = i
            trade_returns = []
        elif sig == 0 and in_trade:
            in_trade = False
            trades.append({
                "start_idx": trade_start,
                "end_idx": i,
                "duration": i - trade_start,
                "returns": trade_returns.copy(),
            })
            trade_returns = []

        if in_trade and i > 0:
            trade_returns.append(returns.iloc[i])

    # Handle open trade
    if in_trade:
        trades.append({
            "start_idx": trade_start,
            "end_idx": len(signal),
            "duration": len(signal) - trade_start,
            "returns": trade_returns.copy(),
        })

    if not trades:
        return {"error": "No trades found"}

    # Statistics
    durations = [t["duration"] for t in trades]
    trade_total_returns = [(1 + pd.Series(t["returns"])).prod() - 1 for t in trades if t["returns"]]

    profitable = [r for r in trade_total_returns if r > 0]
    losing = [r for r in trade_total_returns if r < 0]

    return {
        "n_trades": len(trades),
        "avg_duration": float(np.mean(durations)),
        "median_duration": float(np.median(durations)),
        "max_duration": int(max(durations)),
        "min_duration": int(min(durations)),
        "win_rate": len(profitable) / len(trade_total_returns) if trade_total_returns else 0,
        "avg_win": float(np.mean(profitable)) if profitable else 0,
        "avg_loss": float(np.mean(losing)) if losing else 0,
        "best_trade": float(max(trade_total_returns)) if trade_total_returns else 0,
        "worst_trade": float(min(trade_total_returns)) if trade_total_returns else 0,
        "expectancy": float(np.mean(trade_total_returns)) if trade_total_returns else 0,
    }
