"""Multi-Strategy Book -- run and track multiple strategies with isolated risk budgets.

Manages a portfolio of independent strategy engines, each with its own risk
budget, and provides aggregated analytics, correlation tracking, rebalancing,
and intervention alerts.

Usage:
    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    book.add_strategy("mean_revert", engine_b, risk_budget=0.6)

    # Each cycle:
    book.update()
    report = book.report()
    print(f"Portfolio Sharpe: {report.portfolio_sharpe:.2f}")

    # Check for problems:
    troubled = book.needs_intervention(max_drawdown=0.10)

    # Rebalance:
    new_budgets = book.rebalance_capital(method="risk_parity")

    # Pipeline integration:
    pipeline = [model, book.strategy_pipeline("momentum"), quoter]
"""

from __future__ import annotations

import math
from collections import deque
from dataclasses import dataclass, field
from typing import Any, Callable


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass
class StrategySlot:
    """A single strategy within the multi-strategy book.

    Attributes:
        name: Unique identifier for this strategy.
        engine: Reference to the strategy's Engine instance.
        risk_budget: Fraction of total capital allocated (0.0 to 1.0).
        pnl: Current unrealized + realized PnL.
        sharpe: Rolling Sharpe ratio.
        drawdown: Current drawdown as a fraction of peak equity.
        is_active: Whether the strategy is currently active.
    """

    name: str
    engine: Any  # Engine reference -- not owned
    risk_budget: float
    pnl: float = 0.0
    sharpe: float = 0.0
    drawdown: float = 0.0
    is_active: bool = True


@dataclass
class StrategyCorrelation:
    """Pairwise correlation between two strategy PnL streams.

    Attributes:
        strategy_a: Name of the first strategy.
        strategy_b: Name of the second strategy.
        correlation: Pearson correlation of PnL changes.
        n_observations: Number of overlapping observations used.
    """

    strategy_a: str
    strategy_b: str
    correlation: float
    n_observations: int


@dataclass
class MultiStrategyReport:
    """Aggregated report across all strategies in the book.

    Attributes:
        strategies: Snapshot of each strategy slot.
        total_pnl: Sum of all strategy PnLs.
        total_drawdown: Portfolio-level drawdown.
        portfolio_sharpe: Sharpe ratio of the combined equity curve.
        correlations: Pairwise correlations between strategies.
        capital_utilization: Fraction of total capital in use.
    """

    strategies: list[StrategySlot]
    total_pnl: float
    total_drawdown: float
    portfolio_sharpe: float
    correlations: list[StrategyCorrelation]
    capital_utilization: float


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _pearson_correlation(xs: list[float], ys: list[float]) -> float:
    """Compute Pearson correlation between two series of equal length.

    Returns 0.0 if insufficient data or zero variance.
    """
    n = min(len(xs), len(ys))
    if n < 2:
        return 0.0

    mean_x = sum(xs[:n]) / n
    mean_y = sum(ys[:n]) / n

    cov = 0.0
    var_x = 0.0
    var_y = 0.0
    for i in range(n):
        dx = xs[i] - mean_x
        dy = ys[i] - mean_y
        cov += dx * dy
        var_x += dx * dx
        var_y += dy * dy

    denom = math.sqrt(var_x * var_y)
    if denom < 1e-15:
        return 0.0

    r = cov / denom
    return max(-1.0, min(1.0, r))


def _rolling_sharpe(returns: list[float]) -> float:
    """Compute annualized Sharpe ratio from a list of period returns.

    Assumes returns are per-cycle (high frequency). Uses 252 * 6.5 * 60
    periods per year as a rough approximation for minute-level data, but
    the ratio is still meaningful relative to other strategies in the book.
    """
    n = len(returns)
    if n < 2:
        return 0.0

    mean_r = sum(returns) / n
    variance = sum((r - mean_r) ** 2 for r in returns) / (n - 1)
    std_r = math.sqrt(max(variance, 0.0))

    if std_r < 1e-15:
        return 0.0

    # Annualize with sqrt(n) since we don't know the exact frequency.
    # This gives a scale-free comparison across strategies.
    return (mean_r / std_r) * math.sqrt(min(n, 252))


def _drawdown_from_equity(equity: list[float]) -> float:
    """Compute current drawdown as a fraction of peak equity.

    Returns a non-negative value (0.0 = at peak, 0.1 = 10% below peak).
    """
    if not equity:
        return 0.0

    peak = equity[0]
    for val in equity:
        if val > peak:
            peak = val

    if peak <= 0:
        return 0.0

    current = equity[-1]
    dd = (peak - current) / peak
    return max(0.0, dd)


# ---------------------------------------------------------------------------
# StrategyBook
# ---------------------------------------------------------------------------


_MAX_HISTORY = 10000


class StrategyBook:
    """Multi-strategy book with isolated risk budgets and aggregate analytics.

    Manages multiple Engine references, tracks per-strategy PnL histories,
    computes correlations, and provides rebalancing recommendations.

    The book does NOT own engines -- it holds references and queries them
    via ``engine.status()``. Each strategy is responsible for its own
    execution; the book provides oversight and analytics.

    Thread-safe: no shared mutable state between strategies. Each strategy's
    PnL history is stored in its own deque.
    """

    def __init__(self, total_capital: float = 100000.0) -> None:
        from horizon._horizon import auth_require_pro
        auth_require_pro()

        self._total_capital = total_capital
        self._strategies: dict[str, StrategySlot] = {}
        self._pnl_histories: dict[str, deque[float]] = {}
        self._equity_curves: dict[str, deque[float]] = {}
        self._portfolio_equity: deque[float] = deque(maxlen=_MAX_HISTORY)

    # -- Strategy management ------------------------------------------------

    def add_strategy(
        self,
        name: str,
        engine: Any,
        risk_budget: float | None = None,
    ) -> None:
        """Add a strategy to the book.

        If *risk_budget* is None, redistributes budgets equally across all
        strategies (including the new one).

        Args:
            name: Unique strategy name.
            engine: Reference to the strategy's Engine instance.
            risk_budget: Fraction of total capital to allocate (0.0 to 1.0).
                If None, budgets are distributed equally.

        Raises:
            ValueError: If name already exists or total budget would exceed 1.0.
        """
        if name in self._strategies:
            raise ValueError(f"Strategy '{name}' already exists in the book")

        if risk_budget is None:
            # Redistribute equally
            n = len(self._strategies) + 1
            equal_budget = 1.0 / n
            for slot in self._strategies.values():
                slot.risk_budget = equal_budget
            risk_budget = equal_budget
        else:
            # Validate total budget
            current_total = sum(s.risk_budget for s in self._strategies.values())
            if current_total + risk_budget > 1.0 + 1e-10:
                raise ValueError(
                    f"Total risk budget would be {current_total + risk_budget:.4f}, "
                    f"exceeding 1.0"
                )

        self._strategies[name] = StrategySlot(
            name=name,
            engine=engine,
            risk_budget=risk_budget,
        )
        self._pnl_histories[name] = deque(maxlen=_MAX_HISTORY)
        self._equity_curves[name] = deque(maxlen=_MAX_HISTORY)

    def remove_strategy(self, name: str) -> None:
        """Remove a strategy from the book.

        Args:
            name: Strategy name to remove.
        """
        self._strategies.pop(name, None)
        self._pnl_histories.pop(name, None)
        self._equity_curves.pop(name, None)

    # -- Update cycle -------------------------------------------------------

    def update(self) -> None:
        """Record PnL snapshots for each strategy.

        Call this once per cycle (or at regular intervals) to build up the
        equity curves used for correlation and Sharpe computation.
        """
        total_equity = 0.0

        for name, slot in self._strategies.items():
            if not slot.is_active:
                continue

            # Query engine status for current PnL
            pnl = 0.0
            try:
                status = slot.engine.status()
                total_pnl = getattr(status, "total_pnl", None)
                if callable(total_pnl):
                    pnl = total_pnl()
                elif total_pnl is not None:
                    pnl = float(total_pnl)
            except Exception:
                pass

            slot.pnl = pnl

            # Record PnL for correlation tracking
            self._pnl_histories[name].append(pnl)

            # Build equity curve (capital * budget + pnl)
            equity = self._total_capital * slot.risk_budget + pnl
            self._equity_curves[name].append(equity)
            total_equity += equity

            # Update rolling metrics
            history = list(self._pnl_histories[name])
            if len(history) >= 2:
                # Compute returns as PnL differences
                returns = [
                    history[i] - history[i - 1]
                    for i in range(1, len(history))
                ]
                slot.sharpe = _rolling_sharpe(returns)

            equity_list = list(self._equity_curves[name])
            slot.drawdown = _drawdown_from_equity(equity_list)

        self._portfolio_equity.append(total_equity)

    # -- Reporting ----------------------------------------------------------

    def report(self) -> MultiStrategyReport:
        """Generate a multi-strategy report.

        Computes per-strategy metrics, pairwise PnL correlations, and
        aggregate portfolio-level analytics.

        Returns:
            :class:`MultiStrategyReport` with all current metrics.
        """
        # Strategy snapshots
        strategy_snapshots = [
            StrategySlot(
                name=s.name,
                engine=s.engine,
                risk_budget=s.risk_budget,
                pnl=s.pnl,
                sharpe=s.sharpe,
                drawdown=s.drawdown,
                is_active=s.is_active,
            )
            for s in self._strategies.values()
        ]

        # Total PnL
        total_pnl = sum(s.pnl for s in self._strategies.values())

        # Portfolio drawdown
        portfolio_equity = list(self._portfolio_equity)
        total_drawdown = _drawdown_from_equity(portfolio_equity)

        # Portfolio Sharpe (from combined equity curve)
        portfolio_sharpe = 0.0
        if len(portfolio_equity) >= 2:
            returns = [
                portfolio_equity[i] - portfolio_equity[i - 1]
                for i in range(1, len(portfolio_equity))
            ]
            portfolio_sharpe = _rolling_sharpe(returns)

        # Pairwise correlations
        correlations = self._compute_correlations()

        # Capital utilization
        active_budget = sum(
            s.risk_budget for s in self._strategies.values() if s.is_active
        )
        capital_utilization = min(1.0, active_budget)

        return MultiStrategyReport(
            strategies=strategy_snapshots,
            total_pnl=total_pnl,
            total_drawdown=total_drawdown,
            portfolio_sharpe=portfolio_sharpe,
            correlations=correlations,
            capital_utilization=capital_utilization,
        )

    def _compute_correlations(self) -> list[StrategyCorrelation]:
        """Compute pairwise Pearson correlations of PnL changes."""
        names = [n for n, s in self._strategies.items() if s.is_active]
        correlations: list[StrategyCorrelation] = []

        for i in range(len(names)):
            for j in range(i + 1, len(names)):
                name_a = names[i]
                name_b = names[j]

                hist_a = list(self._pnl_histories.get(name_a, []))
                hist_b = list(self._pnl_histories.get(name_b, []))

                # Compute changes (first differences) for correlation
                n_obs = min(len(hist_a), len(hist_b))
                if n_obs < 3:
                    correlations.append(
                        StrategyCorrelation(
                            strategy_a=name_a,
                            strategy_b=name_b,
                            correlation=0.0,
                            n_observations=n_obs,
                        )
                    )
                    continue

                changes_a = [hist_a[k] - hist_a[k - 1] for k in range(1, n_obs)]
                changes_b = [hist_b[k] - hist_b[k - 1] for k in range(1, n_obs)]

                corr = _pearson_correlation(changes_a, changes_b)

                correlations.append(
                    StrategyCorrelation(
                        strategy_a=name_a,
                        strategy_b=name_b,
                        correlation=corr,
                        n_observations=len(changes_a),
                    )
                )

        return correlations

    # -- Rebalancing --------------------------------------------------------

    def rebalance_capital(self, method: str = "equal") -> dict[str, float]:
        """Compute target risk budgets for each strategy.

        Args:
            method: Rebalancing method. One of:

                - ``"equal"``: Equal allocation across active strategies.
                - ``"risk_parity"``: Allocate inversely to volatility
                  (lower vol gets more capital).
                - ``"performance"``: Allocate proportionally to Sharpe ratio
                  (higher Sharpe gets more capital).

        Returns:
            Dictionary mapping strategy name to target risk_budget (summing
            to 1.0 across active strategies).

        Raises:
            ValueError: If *method* is not recognized.
        """
        active = {
            name: slot
            for name, slot in self._strategies.items()
            if slot.is_active
        }

        if not active:
            return {}

        if method == "equal":
            budget = 1.0 / len(active)
            return {name: budget for name in active}

        elif method == "risk_parity":
            # Inverse volatility weighting
            inv_vols: dict[str, float] = {}
            for name in active:
                history = list(self._pnl_histories.get(name, []))
                if len(history) >= 2:
                    changes = [
                        history[i] - history[i - 1]
                        for i in range(1, len(history))
                    ]
                    mean_c = sum(changes) / len(changes)
                    variance = sum((c - mean_c) ** 2 for c in changes) / max(len(changes) - 1, 1)
                    vol = math.sqrt(max(variance, 0.0))
                    if vol > 1e-15:
                        inv_vols[name] = 1.0 / vol
                    else:
                        inv_vols[name] = 1.0
                else:
                    inv_vols[name] = 1.0

            total_inv = sum(inv_vols.values())
            if total_inv < 1e-15:
                budget = 1.0 / len(active)
                return {name: budget for name in active}

            return {name: iv / total_inv for name, iv in inv_vols.items()}

        elif method == "performance":
            # Allocate proportionally to Sharpe ratio
            # Shift Sharpe values so all are positive (use softmax-like approach)
            sharpes: dict[str, float] = {}
            for name, slot in active.items():
                sharpes[name] = slot.sharpe

            # Shift: add |min| + small epsilon so all values are positive
            min_sharpe = min(sharpes.values()) if sharpes else 0.0
            shift = abs(min_sharpe) + 0.01

            shifted: dict[str, float] = {
                name: s + shift for name, s in sharpes.items()
            }

            total_s = sum(shifted.values())
            if total_s < 1e-15:
                budget = 1.0 / len(active)
                return {name: budget for name in active}

            return {name: s / total_s for name, s in shifted.items()}

        else:
            raise ValueError(
                f"Unknown rebalancing method '{method}'. "
                f"Use 'equal', 'risk_parity', or 'performance'."
            )

    # -- Intervention -------------------------------------------------------

    def needs_intervention(
        self,
        max_drawdown: float = 0.1,
        min_sharpe: float = 0.0,
    ) -> list[str]:
        """Return names of strategies that breach risk thresholds.

        A strategy needs intervention if:
        - Its drawdown exceeds *max_drawdown*, OR
        - Its Sharpe ratio falls below *min_sharpe* (and it has enough
          history for a meaningful Sharpe calculation).

        Args:
            max_drawdown: Maximum allowed drawdown (e.g. 0.1 = 10%).
            min_sharpe: Minimum acceptable Sharpe ratio.

        Returns:
            List of strategy names that need attention.
        """
        troubled: list[str] = []

        for name, slot in self._strategies.items():
            if not slot.is_active:
                continue

            if slot.drawdown > max_drawdown:
                troubled.append(name)
                continue

            # Only check Sharpe if we have enough data
            history = self._pnl_histories.get(name, deque())
            if len(history) >= 10 and slot.sharpe < min_sharpe:
                troubled.append(name)

        return troubled

    # -- Pipeline integration -----------------------------------------------

    def strategy_pipeline(self, strategy_name: str) -> Callable:
        """Return a pipeline function that records metrics for the named strategy.

        Injects the following into ``ctx.params``:

        - ``"strategy_pnl"``: Current PnL for this strategy.
        - ``"strategy_drawdown"``: Current drawdown fraction.
        - ``"strategy_sharpe"``: Rolling Sharpe ratio.

        Args:
            strategy_name: Name of the strategy (must already be in the book).

        Returns:
            Pipeline function compatible with ``hz.run()``.
        """
        book = self

        def _pipeline(ctx: Any, signal: Any = None) -> Any:
            slot = book._strategies.get(strategy_name)
            params = getattr(ctx, "params", {})

            if slot is None:
                params["strategy_pnl"] = 0.0
                params["strategy_drawdown"] = 0.0
                params["strategy_sharpe"] = 0.0
                return signal

            # Query fresh status from engine
            pnl = slot.pnl
            try:
                status = slot.engine.status()
                total_pnl = getattr(status, "total_pnl", None)
                if callable(total_pnl):
                    pnl = total_pnl()
                elif total_pnl is not None:
                    pnl = float(total_pnl)
                slot.pnl = pnl
            except Exception:
                pass

            params["strategy_pnl"] = slot.pnl
            params["strategy_drawdown"] = slot.drawdown
            params["strategy_sharpe"] = slot.sharpe

            return signal

        _pipeline.__name__ = f"strategy_pipeline_{strategy_name}"
        return _pipeline

    # -- Dunder methods -----------------------------------------------------

    def __repr__(self) -> str:
        n = len(self._strategies)
        active = sum(1 for s in self._strategies.values() if s.is_active)
        return (
            f"StrategyBook(capital={self._total_capital}, "
            f"strategies={n}, active={active})"
        )

    def __len__(self) -> int:
        return len(self._strategies)
