"""Stage B: dependence-aware joint path head (Gaussian copula with low-rank + diagonal correlation over the H future
minutes, Student-t marginals from Stage A). Training: copula log-likelihood of the probability-integral transforms of
the observed targets with the marginals frozen. Sampling: z = L eps + sqrt(psi) eta, standardised, u = Phi(z),
y_h = F_t^{-1}(u_h; mu_h, s_h, nu_h). Missing targets are marginalised exactly by replacing their rows/columns of the
correlation matrix with the identity. All linear algebra in FP32 on CPU when the device lacks Cholesky support."""
from __future__ import annotations

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy import stats

SQRT2 = math.sqrt(2.0)


class LowRankCopulaHead(nn.Module):
    def __init__(self, d_model: int, horizon: int, rank: int = 4):
        super().__init__()
        self.rank, self.H = rank, horizon
        self.load = nn.Linear(d_model, rank)
        self.logpsi = nn.Linear(d_model, 1)
        nn.init.normal_(self.load.weight, std=0.02); nn.init.zeros_(self.load.bias)
        nn.init.zeros_(self.logpsi.weight); nn.init.zeros_(self.logpsi.bias)

    def correlation(self, h: torch.Tensor) -> torch.Tensor:
        """h: [B, H, d] -> correlation matrices [B, H, H] (FP32)."""
        lam = self.load(h.float())                       # [B, H, k]
        psi = F.softplus(self.logpsi(h.float())).squeeze(-1) + 1e-3   # [B, H]
        cov = lam @ lam.transpose(1, 2) + torch.diag_embed(psi)
        d = torch.rsqrt(torch.diagonal(cov, dim1=1, dim2=2))
        return cov * d.unsqueeze(1) * d.unsqueeze(2)

    @staticmethod
    def _masked_corr(R: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        m = mask.float()
        return R * m.unsqueeze(1) * m.unsqueeze(2) + torch.diag_embed(1.0 - m)

    def nll(self, h: torch.Tensor, z: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """Gaussian-copula negative log density (per observed dimension): z = Phi^{-1}(u) [B, H], mask [B, H]."""
        R = self._masked_corr(self.correlation(h), mask)
        dev = R.device
        Rc = R.cpu() if dev.type == "mps" else R
        zc = (z.float() * mask.float()).unsqueeze(-1)
        zc = zc.cpu() if dev.type == "mps" else zc
        L, info = torch.linalg.cholesky_ex(Rc)
        logdet = 2.0 * torch.log(torch.diagonal(L, dim1=1, dim2=2)).sum(-1)
        sol = torch.cholesky_solve(zc, L)
        quad = (zc * sol).sum(dim=(1, 2))
        # copula density: -0.5 logdet - 0.5 z'(R^-1 - I) z
        ll = -0.5 * logdet - 0.5 * (quad - (zc * zc).sum(dim=(1, 2)))
        n = mask.float().sum(-1).clamp_min(1.0)
        n = n.cpu() if dev.type == "mps" else n
        return -(ll / n).mean().to(dev)

    def sample_z(self, h: torch.Tensor, n: int, gen: torch.Generator | None = None) -> torch.Tensor:
        """Correlated standard normals [n, B, H] from the low-rank structure (no Cholesky needed)."""
        lam = self.load(h.float()); psi = F.softplus(self.logpsi(h.float())).squeeze(-1) + 1e-3
        B, H, k = lam.shape
        eps = torch.randn((n, B, k), generator=gen, device=lam.device)
        eta = torch.randn((n, B, H), generator=gen, device=lam.device)
        z = torch.einsum("bhk,nbk->nbh", lam, eps) + torch.sqrt(psi).unsqueeze(0) * eta
        sd = torch.sqrt((lam * lam).sum(-1) + psi).unsqueeze(0)
        return z / sd


def pit_student_t(y: np.ndarray, mu: np.ndarray, scale: np.ndarray, nu: np.ndarray) -> np.ndarray:
    y, mu, scale, nu = (np.asarray(a, dtype=np.float64) for a in (y, mu, scale, nu))
    return stats.t.cdf((y - mu) / scale, nu)


def z_from_pit(u: np.ndarray) -> np.ndarray:
    u = np.clip(u, 1e-6, 1 - 1e-6)
    return stats.norm.ppf(u)


def samples_from_copula(z: torch.Tensor, mu: torch.Tensor, scale: torch.Tensor, nu: torch.Tensor) -> torch.Tensor:
    """z [n, B, H] correlated normals -> Student-t marginal samples [n, B, H] via the PIT (CPU numpy)."""
    # float64 throughout: scipy's t.ppf returns NaN for float32 inputs near 0.5 (observed with scipy 1.17)
    u = (0.5 * (1.0 + torch.erf(z.detach().cpu().double() / SQRT2))).clamp(1e-6, 1 - 1e-6).numpy()
    mu_, sc_, nu_ = (t.detach().cpu().double().numpy() for t in (mu, scale, nu))
    y = stats.t.ppf(u, nu_[None]) * sc_[None] + mu_[None]
    return torch.from_numpy(np.asarray(y, dtype=np.float32))