"""VaR-based risk budgeting across strategies."""

from __future__ import annotations

import logging
import math
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger("horizon.fund.var_budget")


@dataclass(frozen=True)
class VaRAllocation:
    """VaR budget allocation for a single strategy."""
    strategy: str
    budget_pct: float
    budget_amount: float
    used_amount: float
    usage_pct: float
    available: float
    at_limit: bool


@dataclass(frozen=True)
class VaRBudgetSnapshot:
    """Point-in-time snapshot of fund VaR budget."""
    timestamp: float
    total_budget: float
    used_budget: float
    usage_pct: float
    reserve_pct: float
    reserve_amount: float
    allocations: list[VaRAllocation]
    fund_var_95: float
    fund_cvar_95: float


class VaRBudgetManager:
    """Allocates and tracks a total VaR budget across strategies.

    Each strategy gets a fraction of the total fund VaR budget.
    As strategies take positions, they consume their allocated budget.
    When a strategy hits its limit, it cannot increase exposure.

    VaR is estimated using historical simulation (percentile of
    rolling return distribution).
    """

    def __init__(
        self,
        total_var_budget: float = 50_000.0,
        reserve_pct: float = 0.25,
        confidence: float = 0.95,
        lookback: int = 100,
    ) -> None:
        self._total_budget = total_var_budget
        self._reserve_pct = reserve_pct
        self._confidence = confidence
        self._lookback = lookback
        self._allocations: dict[str, float] = {}  # strategy -> fraction of budget
        self._returns: dict[str, deque[float]] = {}
        self._last_nav: dict[str, float] = {}
        self._lock = threading.Lock()

    def set_allocation(self, strategy: str, fraction: float) -> None:
        """Set budget allocation for a strategy (as fraction of total, 0-1)."""
        with self._lock:
            self._allocations[strategy] = max(0.0, min(1.0 - self._reserve_pct, fraction))

    def set_allocations(self, allocations: dict[str, float]) -> None:
        """Batch set allocations. Normalizes so total + reserve = 1.0."""
        total = sum(allocations.values())
        available = 1.0 - self._reserve_pct
        with self._lock:
            if total > 0:
                for name, frac in allocations.items():
                    self._allocations[name] = (frac / total) * available
            else:
                for name in allocations:
                    self._allocations[name] = 0.0

    def remove_strategy(self, strategy: str) -> None:
        with self._lock:
            self._allocations.pop(strategy, None)
            self._returns.pop(strategy, None)
            self._last_nav.pop(strategy, None)

    def update_navs(self, strategy_navs: dict[str, float]) -> None:
        """Update with current strategy NAVs to track returns."""
        with self._lock:
            for name, nav in strategy_navs.items():
                if name not in self._returns:
                    self._returns[name] = deque(maxlen=self._lookback)
                    self._last_nav[name] = nav
                    continue

                prev = self._last_nav[name]
                if prev > 0:
                    ret = (nav - prev) / prev
                    self._returns[name].append(ret)
                self._last_nav[name] = nav

    def compute_strategy_var(self, strategy: str) -> float:
        """Compute VaR for a single strategy using historical simulation."""
        with self._lock:
            returns = list(self._returns.get(strategy, []))

        if len(returns) < 5:
            return 0.0

        nav = self._last_nav.get(strategy, 0.0)
        return self._historical_var(returns, nav, self._confidence)

    def snapshot(self) -> VaRBudgetSnapshot:
        """Compute current VaR budget snapshot."""
        now = time.time()
        allocations: list[VaRAllocation] = []
        total_used = 0.0

        with self._lock:
            all_returns: list[float] = []

            for name in self._allocations:
                strat_returns = list(self._returns.get(name, []))
                nav = self._last_nav.get(name, 0.0)
                budget_pct = self._allocations.get(name, 0.0)
                budget_amount = budget_pct * self._total_budget

                var = self._historical_var(strat_returns, nav, self._confidence) if strat_returns else 0.0
                available = max(0.0, budget_amount - var)

                allocations.append(VaRAllocation(
                    strategy=name,
                    budget_pct=budget_pct,
                    budget_amount=budget_amount,
                    used_amount=var,
                    usage_pct=(var / budget_amount * 100.0) if budget_amount > 0 else 0.0,
                    available=available,
                    at_limit=var >= budget_amount,
                ))
                total_used += var

                # Aggregate returns for fund-level VaR
                for r in strat_returns:
                    if len(all_returns) < len(strat_returns):
                        all_returns.extend([0.0] * (len(strat_returns) - len(all_returns)))
                    # Simple sum (ignores correlation for aggregation)

            # Fund-level VaR (sum of individual VaRs - conservative)
            fund_nav = sum(self._last_nav.values())
            fund_var = self._historical_var(
                [sum(r) for r in zip(*[list(self._returns.get(n, [])) for n in self._allocations]) if r]
                if len(self._allocations) > 1 else list(next(iter(self._returns.values()), [])),
                fund_nav,
                self._confidence,
            ) if self._returns else 0.0

            fund_cvar = self._historical_cvar(
                [sum(r) for r in zip(*[list(self._returns.get(n, [])) for n in self._allocations]) if r]
                if len(self._allocations) > 1 else list(next(iter(self._returns.values()), [])),
                fund_nav,
                self._confidence,
            ) if self._returns else 0.0

        reserve_amount = self._reserve_pct * self._total_budget
        usage_pct = (total_used / self._total_budget * 100.0) if self._total_budget > 0 else 0.0

        return VaRBudgetSnapshot(
            timestamp=now,
            total_budget=self._total_budget,
            used_budget=total_used,
            usage_pct=usage_pct,
            reserve_pct=self._reserve_pct,
            reserve_amount=reserve_amount,
            allocations=allocations,
            fund_var_95=fund_var,
            fund_cvar_95=fund_cvar,
        )

    def can_increase_exposure(self, strategy: str) -> bool:
        """Check if a strategy has remaining VaR budget."""
        with self._lock:
            budget_pct = self._allocations.get(strategy, 0.0)
            budget_amount = budget_pct * self._total_budget
            returns = list(self._returns.get(strategy, []))
            nav = self._last_nav.get(strategy, 0.0)

        var = self._historical_var(returns, nav, self._confidence) if returns else 0.0
        return var < budget_amount

    def marginal_var(self, strategy: str, additional_return: float) -> float:
        """Estimate how much VaR would increase with an additional position return."""
        with self._lock:
            returns = list(self._returns.get(strategy, []))
            nav = self._last_nav.get(strategy, 0.0)

        if not returns:
            return 0.0

        current_var = self._historical_var(returns, nav, self._confidence)
        augmented = returns + [additional_return]
        new_var = self._historical_var(augmented, nav, self._confidence)
        return new_var - current_var

    @staticmethod
    def _historical_var(returns: list[float], nav: float, confidence: float) -> float:
        """VaR using historical simulation (percentile method)."""
        if not returns or nav <= 0:
            return 0.0
        sorted_returns = sorted(returns)
        idx = int((1 - confidence) * len(sorted_returns))
        idx = max(0, min(idx, len(sorted_returns) - 1))
        return abs(sorted_returns[idx] * nav)

    @staticmethod
    def _historical_cvar(returns: list[float], nav: float, confidence: float) -> float:
        """CVaR (Expected Shortfall) using historical simulation."""
        if not returns or nav <= 0:
            return 0.0
        sorted_returns = sorted(returns)
        idx = int((1 - confidence) * len(sorted_returns))
        idx = max(1, min(idx, len(sorted_returns)))
        tail = sorted_returns[:idx]
        if not tail:
            return 0.0
        return abs(sum(tail) / len(tail) * nav)
