"""Transparent industry/style factor risk reference implementation."""

from __future__ import annotations

import math
import statistics
from dataclasses import dataclass
from typing import Mapping, Sequence

from .portable_core import normalize_symbol
from .research import DeterministicRidge


class RiskModelError(ValueError):
    pass


def _finite(value: float, name: str) -> float:
    numeric = float(value)
    if not math.isfinite(numeric):
        raise RiskModelError(f"{name} must be finite")
    return numeric


@dataclass(frozen=True, slots=True)
class FactorExposureSnapshot:
    as_of: str
    exposures: dict[str, dict[str, float]]

    def __post_init__(self) -> None:
        if not str(self.as_of).strip():
            raise RiskModelError("exposure as_of is required")
        if not self.exposures:
            raise RiskModelError("factor exposures are empty")
        normalized: dict[str, dict[str, float]] = {}
        factor_set = None
        for raw_symbol, raw_factors in self.exposures.items():
            symbol = normalize_symbol(raw_symbol)
            factors = {
                str(name): _finite(value, f"{symbol}.{name}")
                for name, value in raw_factors.items()
            }
            if not factors:
                raise RiskModelError(f"factor exposures are empty for {symbol}")
            if factor_set is None:
                factor_set = set(factors)
            elif set(factors) != factor_set:
                raise RiskModelError("factor exposure columns are inconsistent")
            normalized[symbol] = dict(sorted(factors.items()))
        object.__setattr__(self, "exposures", dict(sorted(normalized.items())))

    @property
    def factor_names(self) -> tuple[str, ...]:
        first = next(iter(self.exposures.values()))
        return tuple(sorted(first))


@dataclass(frozen=True, slots=True)
class CrossSectionalFactorFit:
    factor_returns: dict[str, float]
    residuals: dict[str, float]
    observation_count: int


def estimate_factor_returns(
    *,
    returns: Mapping[str, float],
    exposure: FactorExposureSnapshot,
    ridge_alpha: float = 1e-4,
) -> CrossSectionalFactorFit:
    normalized_returns = {
        normalize_symbol(symbol): _finite(value, f"return.{symbol}")
        for symbol, value in returns.items()
    }
    symbols = sorted(set(normalized_returns) & set(exposure.exposures))
    factors = exposure.factor_names
    if len(symbols) <= len(factors):
        raise RiskModelError(
            "cross-sectional regression needs more securities than factors"
        )
    matrix = [
        [exposure.exposures[symbol][factor] for factor in factors]
        for symbol in symbols
    ]
    target = [normalized_returns[symbol] for symbol in symbols]
    model = DeterministicRidge(
        alpha=ridge_alpha,
        fit_intercept=False,
    ).fit(matrix, target)
    predictions = model.predict(matrix)
    assert model.coef_ is not None
    return CrossSectionalFactorFit(
        factor_returns={
            factor: model.coef_[index]
            for index, factor in enumerate(factors)
        },
        residuals={
            symbol: observed - predicted
            for symbol, observed, predicted in zip(
                symbols,
                target,
                predictions,
            )
        },
        observation_count=len(symbols),
    )


@dataclass(frozen=True, slots=True)
class FactorRiskForecast:
    factor_names: tuple[str, ...]
    factor_covariance: dict[str, dict[str, float]]
    specific_variance: dict[str, float]
    half_life: float
    shrinkage: float
    variance_floor: float


def _ewma_weights(count: int, half_life: float) -> list[float]:
    if count <= 0:
        raise RiskModelError("EWMA history is empty")
    if not math.isfinite(float(half_life)) or half_life <= 0:
        raise RiskModelError("half_life must be finite and positive")
    decay = math.exp(math.log(0.5) / float(half_life))
    raw = [decay ** (count - 1 - index) for index in range(count)]
    total = sum(raw)
    return [value / total for value in raw]


def build_factor_risk_forecast(
    *,
    factor_return_history: Sequence[Mapping[str, float]],
    residual_history: Mapping[str, Sequence[float]],
    half_life: float = 20.0,
    shrinkage: float = 0.20,
    variance_floor: float = 1e-8,
) -> FactorRiskForecast:
    if not factor_return_history:
        raise RiskModelError("factor return history is empty")
    if not math.isfinite(float(shrinkage)) or not (0 <= shrinkage <= 1):
        raise RiskModelError("shrinkage must lie in [0, 1]")
    if not math.isfinite(float(variance_floor)) or variance_floor <= 0:
        raise RiskModelError("variance_floor must be finite and positive")
    factors = tuple(sorted(factor_return_history[0]))
    if not factors:
        raise RiskModelError("factor return columns are empty")
    rows = []
    for raw_row in factor_return_history:
        if set(raw_row) != set(factors):
            raise RiskModelError("factor return columns are inconsistent")
        rows.append(
            {
                factor: _finite(raw_row[factor], f"factor_return.{factor}")
                for factor in factors
            }
        )
    weights = _ewma_weights(len(rows), half_life)
    means = {
        factor: sum(
            weight * row[factor] for weight, row in zip(weights, rows)
        )
        for factor in factors
    }
    covariance: dict[str, dict[str, float]] = {}
    for left in factors:
        covariance[left] = {}
        for right in factors:
            sample = sum(
                weight
                * (row[left] - means[left])
                * (row[right] - means[right])
                for weight, row in zip(weights, rows)
            )
            if left == right:
                value = max(variance_floor, sample)
            else:
                value = (1.0 - shrinkage) * sample
            covariance[left][right] = value

    specific: dict[str, float] = {}
    for raw_symbol, raw_values in residual_history.items():
        symbol = normalize_symbol(raw_symbol)
        values = [
            _finite(value, f"specific_residual.{symbol}")
            for value in raw_values
        ]
        if len(values) < 2:
            specific[symbol] = variance_floor
        else:
            specific[symbol] = max(
                variance_floor,
                statistics.variance(values),
            )
    if not specific:
        raise RiskModelError("specific residual history is empty")
    return FactorRiskForecast(
        factor_names=factors,
        factor_covariance=covariance,
        specific_variance=dict(sorted(specific.items())),
        half_life=float(half_life),
        shrinkage=float(shrinkage),
        variance_floor=float(variance_floor),
    )


def portfolio_factor_variance(
    *,
    portfolio_weights: Mapping[str, float],
    exposure: FactorExposureSnapshot,
    forecast: FactorRiskForecast,
) -> float:
    weights = {
        normalize_symbol(symbol): _finite(value, f"weight.{symbol}")
        for symbol, value in portfolio_weights.items()
    }
    if any(value < 0 for value in weights.values()):
        raise RiskModelError("portfolio weights must be non-negative")
    missing_exposure = set(weights) - set(exposure.exposures)
    missing_specific = set(weights) - set(forecast.specific_variance)
    if missing_exposure or missing_specific:
        raise RiskModelError(
            "risk inputs missing for "
            f"exposure={sorted(missing_exposure)}, "
            f"specific={sorted(missing_specific)}"
        )
    if set(exposure.factor_names) != set(forecast.factor_names):
        raise RiskModelError("exposure/forecast factors differ")
    factor_weight = {
        factor: sum(
            weight * exposure.exposures[symbol][factor]
            for symbol, weight in weights.items()
        )
        for factor in forecast.factor_names
    }
    factor_variance = sum(
        factor_weight[left]
        * forecast.factor_covariance[left][right]
        * factor_weight[right]
        for left in forecast.factor_names
        for right in forecast.factor_names
    )
    specific_variance = sum(
        weight * weight * forecast.specific_variance[symbol]
        for symbol, weight in weights.items()
    )
    result = factor_variance + specific_variance
    if not math.isfinite(result) or result < -1e-12:
        raise RiskModelError("portfolio variance is invalid")
    return max(0.0, result)


def factor_stress_loss(
    *,
    portfolio_weights: Mapping[str, float],
    exposure: FactorExposureSnapshot,
    factor_shocks: Mapping[str, float],
) -> float:
    shocks = {
        str(factor): _finite(value, f"shock.{factor}")
        for factor, value in factor_shocks.items()
    }
    unknown = set(shocks) - set(exposure.factor_names)
    if unknown:
        raise RiskModelError(f"unknown stress factors: {sorted(unknown)}")
    loss = 0.0
    for raw_symbol, raw_weight in portfolio_weights.items():
        symbol = normalize_symbol(raw_symbol)
        if symbol not in exposure.exposures:
            raise RiskModelError(f"missing exposure for {symbol}")
        weight = _finite(raw_weight, f"weight.{symbol}")
        return_shock = sum(
            exposure.exposures[symbol][factor] * shock
            for factor, shock in shocks.items()
        )
        loss -= weight * return_shock
    return loss