"""Bias correction model for three-point estimates, including an AI covariate."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping, Sequence

import math
import numpy as np

from ..domain import DeveloperId, WorkItemEstimate


@dataclass(frozen=True)
class BiasObservation:
    """Historical data point for estimating estimation bias.

    Parameters correspond to the paper's log-ratio error:
        eta_ij = log(actual_ij / m_ij)

    Attributes
    ----------
    engineer_id:
        Developer who produced the estimate.
    estimate_most_likely:
        The historical "most-likely" estimate m_ij.
    actual:
        Realised effort.
    ai_intensity:
        AI usage intensity u_ij.
    """

    engineer_id: DeveloperId
    estimate_most_likely: float
    actual: float
    ai_intensity: float = 0.0

    def __post_init__(self) -> None:
        if self.estimate_most_likely <= 0 or self.actual <= 0:
            raise ValueError("estimate_most_likely and actual must be positive")
        if self.ai_intensity < 0:
            raise ValueError("ai_intensity must be non-negative")

    @property
    def log_ratio_error(self) -> float:
        return math.log(self.actual / self.estimate_most_likely)


@dataclass(frozen=True)
class BiasHyperparameters:
    """Priors and noise settings for the hierarchical bias model."""

    mu_prior_mean: float = 0.0
    mu_prior_std: float = 1.0
    beta_ai_prior_mean: float = 0.0
    beta_ai_prior_std: float = 1.0
    alpha_prior_std: float = 1.0
    noise_std: float | None = None


@dataclass(frozen=True)
class BiasModelParameters:
    """Parameters used for bias correction in forecasts."""

    mu_alpha: float
    beta_ai: float
    alpha_by_engineer: Mapping[DeveloperId, float]
    delta_b: float = 0.0

    def alpha_for(self, engineer_id: DeveloperId) -> float:
        return float(self.alpha_by_engineer.get(engineer_id, 0.0))

    def correction_log_factor(self, engineer_id: DeveloperId, item_ai_intensity: float) -> float:
        """Return log multiplicative correction applied to all triad points."""
        if item_ai_intensity < 0:
            raise ValueError("item_ai_intensity must be non-negative")
        return float(self.mu_alpha + self.alpha_for(engineer_id) - self.beta_ai * item_ai_intensity)

    def correction_factor(self, engineer_id: DeveloperId, item_ai_intensity: float) -> float:
        return math.exp(self.correction_log_factor(engineer_id, item_ai_intensity))

    def correct_triad(self, item: WorkItemEstimate) -> tuple[float, float, float]:
        """Return (a', m', b') after bias + AI correction."""
        factor = self.correction_factor(item.owner_id, item.ai_intensity)
        a_p = factor * item.optimistic
        m_p = factor * item.most_likely
        b_p = math.exp(self.delta_b) * factor * item.pessimistic
        return (a_p, m_p, b_p)


class HierarchicalBiasEstimator:
    """Estimate bias parameters with shrinkage (linear-Gaussian posterior mean).

    Model:
        eta_ij = mu_alpha + alpha_i - beta_ai * u_ij + eps
        eps ~ Normal(0, sigma^2)
        mu_alpha ~ Normal(m0, s0^2)
        beta_ai ~ Normal(b0, sb^2)
        alpha_i ~ Normal(0, tau^2)

    The posterior mean is obtained by solving a ridge-regularised least squares
    system (equivalent to the conjugate Gaussian posterior when sigma is known).
    """

    def __init__(self, hyper: BiasHyperparameters | None = None):
        self.hyper = hyper or BiasHyperparameters()

    def fit(
        self,
        observations: Sequence[BiasObservation],
        engineer_ids: Sequence[DeveloperId],
        *,
        delta_b: float = 0.0,
    ) -> BiasModelParameters:
        if len(engineer_ids) == 0:
            raise ValueError("engineer_ids must be non-empty")
        if len(observations) == 0:
            # With no data, return prior means and zero engineer effects.
            return BiasModelParameters(
                mu_alpha=self.hyper.mu_prior_mean,
                beta_ai=self.hyper.beta_ai_prior_mean,
                alpha_by_engineer={eid: 0.0 for eid in engineer_ids},
                delta_b=delta_b,
            )

        y = np.array([obs.log_ratio_error for obs in observations], dtype=float)
        u = np.array([obs.ai_intensity for obs in observations], dtype=float)

        engineer_index = {eid: idx for idx, eid in enumerate(engineer_ids)}
        n_eng = len(engineer_ids)
        n_obs = len(observations)

        # Parameter vector: [mu_alpha, beta_ai, alpha_0..alpha_{n_eng-1}]
        p = 2 + n_eng
        X = np.zeros((n_obs, p), dtype=float)
        X[:, 0] = 1.0
        X[:, 1] = -u
        for r, obs in enumerate(observations):
            idx = engineer_index.get(obs.engineer_id)
            if idx is None:
                raise KeyError(f"Unknown engineer_id in observations: {obs.engineer_id}")
            X[r, 2 + idx] = 1.0

        sigma = self.hyper.noise_std
        if sigma is None:
            # Conservative default: use sample std of y.
            sigma = float(np.std(y, ddof=1)) if n_obs > 1 else 1.0
            sigma = max(sigma, 1e-6)

        prior_mean = np.zeros(p, dtype=float)
        prior_mean[0] = self.hyper.mu_prior_mean
        prior_mean[1] = self.hyper.beta_ai_prior_mean

        prior_var = np.zeros(p, dtype=float)
        prior_var[0] = self.hyper.mu_prior_std**2
        prior_var[1] = self.hyper.beta_ai_prior_std**2
        prior_var[2:] = self.hyper.alpha_prior_std**2

        prior_prec = np.diag(1.0 / np.maximum(prior_var, 1e-12))

        XtX = (X.T @ X) / (sigma**2)
        Xty = (X.T @ y) / (sigma**2)

        A = XtX + prior_prec
        b = Xty + prior_prec @ prior_mean

        theta_hat = np.linalg.solve(A, b)

        mu_alpha = float(theta_hat[0])
        beta_ai = float(theta_hat[1])
        alpha_vec = theta_hat[2:]

        alpha_by_engineer = {eid: float(alpha_vec[i]) for i, eid in enumerate(engineer_ids)}

        return BiasModelParameters(
            mu_alpha=mu_alpha,
            beta_ai=beta_ai,
            alpha_by_engineer=alpha_by_engineer,
            delta_b=delta_b,
        )
