"""Focus/productivity multiplier model, including AI covariates."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping, Sequence

import math
import numpy as np

from ..domain import DeveloperId


@dataclass(frozen=True)
class FocusObservation:
    """Historical observation of a developer-day focus multiplier."""

    engineer_id: DeveloperId
    observed_multiplier: float
    ai_intensity: float = 0.0

    def __post_init__(self) -> None:
        if self.observed_multiplier <= 0.0:
            raise ValueError("observed_multiplier must be positive")
        if self.ai_intensity < 0.0:
            raise ValueError("ai_intensity must be non-negative")

    @property
    def log_value(self) -> float:
        return math.log(self.observed_multiplier)


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

    m_theta_prior_mean: float = 0.0
    m_theta_prior_std: float = 1.0
    beta_ai_prior_mean: float = 0.0
    beta_ai_prior_std: float = 1.0
    delta_prior_std: float = 1.0
    noise_std: float | None = None


@dataclass(frozen=True)
class FocusModelParameters:
    """Parameters for the lognormal focus model."""

    base_log_mu_by_engineer: Mapping[DeveloperId, float]
    log_sigma: float
    beta_ai: float

    def base_mu_for(self, engineer_id: DeveloperId) -> float:
        return float(self.base_log_mu_by_engineer.get(engineer_id, 0.0))

    def effective_log_mu(self, engineer_id: DeveloperId, ai_intensity: float) -> float:
        if ai_intensity < 0.0:
            raise ValueError("ai_intensity must be non-negative")
        return float(self.base_mu_for(engineer_id) + self.beta_ai * ai_intensity)


class HierarchicalFocusEstimator:
    """Estimate focus parameters with shrinkage.

    Model on the log scale:
        y_{i,t} = m_theta + delta_i + beta_ai * u_{i,t} + eps
        eps ~ Normal(0, sigma^2)
        m_theta ~ Normal(m0, s0^2)
        beta_ai ~ Normal(b0, sb^2)
        delta_i ~ Normal(0, tau^2)

    We return posterior means for m_theta + delta_i as per-engineer base mu.
    The noise std is estimated if not supplied.
    """

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

    def fit(
        self,
        observations: Sequence[FocusObservation],
        engineer_ids: Sequence[DeveloperId],
    ) -> FocusModelParameters:
        if len(engineer_ids) == 0:
            raise ValueError("engineer_ids must be non-empty")

        if len(observations) == 0:
            base = {
                eid: self.hyper.m_theta_prior_mean
                for eid in engineer_ids
            }
            return FocusModelParameters(
                base_log_mu_by_engineer=base,
                log_sigma=float(self.hyper.noise_std or 1.0),
                beta_ai=self.hyper.beta_ai_prior_mean,
            )

        y = np.array([obs.log_value 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: [m_theta, beta_ai, delta_0..delta_{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:
            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.m_theta_prior_mean
        prior_mean[1] = self.hyper.beta_ai_prior_mean

        prior_var = np.zeros(p, dtype=float)
        prior_var[0] = self.hyper.m_theta_prior_std**2
        prior_var[1] = self.hyper.beta_ai_prior_std**2
        prior_var[2:] = self.hyper.delta_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)

        m_theta = float(theta_hat[0])
        beta_ai = float(theta_hat[1])
        delta = theta_hat[2:]

        base_mu_by_engineer = {
            eid: float(m_theta + delta[i]) for i, eid in enumerate(engineer_ids)
        }

        # Estimate sigma from residuals using posterior mean parameters.
        residuals = y - X @ theta_hat
        sigma_hat = float(np.sqrt(np.mean(residuals**2)))
        sigma_hat = max(sigma_hat, 1e-6)

        return FocusModelParameters(
            base_log_mu_by_engineer=base_mu_by_engineer,
            log_sigma=sigma_hat,
            beta_ai=beta_ai,
        )


class LogNormalFocusSampler:
    """Strategy for sampling F_{i,t} given FocusModelParameters."""

    def __init__(self, params: FocusModelParameters):
        self.params = params

    def sample(
        self,
        engineer_id: DeveloperId,
        ai_intensity: float,
        rng: np.random.Generator,
        size: int,
        *,
        base_log_mu: float | None = None,
        log_sigma: float | None = None,
    ) -> np.ndarray:
        """Sample focus multipliers.

        Parameters
        ----------
        base_log_mu:
            Optional override. If provided, this is treated as mu^{(F)}_{i,t}.
        log_sigma:
            Optional override for the log-scale std.
        """
        if size <= 0:
            raise ValueError("size must be positive")
        if ai_intensity < 0.0:
            raise ValueError("ai_intensity must be non-negative")

        mu0 = self.params.base_mu_for(engineer_id) if base_log_mu is None else base_log_mu
        sig = self.params.log_sigma if log_sigma is None else log_sigma
        mu_eff = float(mu0 + self.params.beta_ai * ai_intensity)
        return np.exp(rng.normal(loc=mu_eff, scale=sig, size=size))
