"""Correlation matrix operations for portfolio diversification.

Feature: autonomous-trading-engine

Pure computation module for price correlation coefficients between
tracked companies. Used by the Position Sizer to prevent
over-concentration in correlated positions.
"""

from __future__ import annotations

from services.trading.models import OpenPosition


class CorrelationMatrix:
    """In-memory correlation matrix for ticker pairs.

    Stores pairwise correlation coefficients and provides lookup
    methods for individual pairs and weighted portfolio averages.
    """

    def __init__(self) -> None:
        self._data: dict[tuple[str, str], float] = {}

    def load(self, data: dict[tuple[str, str], float]) -> None:
        """Load correlation data from a dict of (ticker_a, ticker_b) -> coefficient."""
        self._data = dict(data)

    def get_correlation(self, ticker_a: str, ticker_b: str) -> float:
        """Return the correlation coefficient for a ticker pair.

        Checks both orderings (a, b) and (b, a). Returns 0.0 if the
        pair is not in the matrix.
        """
        if ticker_a == ticker_b:
            return 1.0
        return self._data.get(
            (ticker_a, ticker_b),
            self._data.get((ticker_b, ticker_a), 0.0),
        )

    def get_portfolio_correlation(
        self,
        candidate: str,
        positions: list[OpenPosition],
    ) -> float:
        """Compute weighted average correlation between candidate and positions.

        Weights are based on each position's market_value. Returns 0.0
        if there are no positions or total weight is zero.
        """
        if not positions:
            return 0.0

        total_weight = 0.0
        weighted_corr = 0.0

        for pos in positions:
            corr = self.get_correlation(candidate, pos.ticker)
            weight = pos.market_value
            weighted_corr += corr * weight
            total_weight += weight

        if total_weight == 0.0:
            return 0.0

        return weighted_corr / total_weight