"""
Multi-Asset Correlation — model cross-asset effects for portfolio risk.

Improves strategy selection by considering correlation with BTC and other assets.
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

import numpy as np
from numba import njit


@njit(cache=True)
def compute_rolling_correlation(
    returns_a: np.ndarray,
    returns_b: np.ndarray,
    window: int = 20,
) -> float:
    """
    Compute rolling Pearson correlation between two return series.
    """
    if len(returns_a) < window or len(returns_b) < window:
        return 0.0

    a = returns_a[-window:]
    b = returns_b[-window:]

    mean_a = np.mean(a)
    mean_b = np.mean(b)

    var_a = np.var(a)
    var_b = np.var(b)

    if var_a <= 0 or var_b <= 0:
        return 0.0

    cov = np.mean((a - mean_a) * (b - mean_b))
    return cov / math.sqrt(var_a * var_b)


@njit(cache=True)
def compute_correlation_regime(
    correlation: float,
    correlation_vol: float,
) -> float:
    """
    Compute correlation regime score (0-1).

    High correlation (>0.8) → regime = 1 (correlated)
    Low correlation (<0.2) → regime = 0 (uncorrelated)
    """
    # Sigmoid mapping
    return 1.0 / (1.0 + math.exp(-5.0 * (correlation - 0.5)))


class MultiAssetCorrelationModel:
    """
    Multi-asset correlation model for portfolio risk.

    Tracks correlations between assets and uses them for:
      - Portfolio risk management
      - Correlation-based strategy selection
      - Hedging decisions
    """

    def __init__(self) -> None:
        self._returns: Dict[str, list[float]] = {}
        self._correlations: Dict[Tuple[str, str], float] = {}

    def update_returns(self, symbol: str, ret: float) -> None:
        """Update return series for an asset."""
        if symbol not in self._returns:
            self._returns[symbol] = []
        self._returns[symbol].append(ret)
        if len(self._returns[symbol]) > 1000:
            self._returns[symbol] = self._returns[symbol][-500:]

    def compute_correlation(self, symbol_a: str, symbol_b: str, window: int = 20) -> float:
        """Compute correlation between two assets."""
        if symbol_a not in self._returns or symbol_b not in self._returns:
            return 0.0
        returns_a = np.array(self._returns[symbol_a], dtype=np.float64)
        returns_b = np.array(self._returns[symbol_b], dtype=np.float64)
        corr = compute_rolling_correlation(returns_a, returns_b, window)
        self._correlations[(symbol_a, symbol_b)] = corr
        self._correlations[(symbol_b, symbol_a)] = corr
        return corr

    def get_correlation(self, symbol_a: str, symbol_b: str) -> float:
        """Get cached correlation."""
        return self._correlations.get((symbol_a, symbol_b), 0.0)

    def get_btc_correlation(self, symbol: str) -> float:
        """Get correlation with BTC."""
        return self.get_correlation(symbol, "BTCUSDT")

    @property
    def asset_count(self) -> int:
        return len(self._returns)