"""
Per-parameter reparameterization bijectors for TPD bounded parameters.

Used by BoundedMechanismFlow (tpd_model.py) to map physically-bounded
parameters (e.g. theta_0 in [0, 1]) onto an unbounded space R before they
are modeled by the normalizing flow. This eliminates two failure modes of
the baseline architecture:

  1. Posterior samples that escape the physical domain (e.g. theta_0 < 0).
  2. Posterior collapse / under-coverage near the domain boundary because
     the spline coupling has to learn a hard constraint at +-tail_bound.

Each bijector exposes the same interface:

    forward(x)  -> (u,    log|du/dx|)         # physical -> unbounded
    inverse(u)  -> (x,    log|dx/du|)         # unbounded -> physical

`log|...|` has shape `[batch]`, summed over the (single) parameter dim
that the bijector acts on. Bijectors operate elementwise on the last
axis; composition across a parameter vector is handled by
`apply_param_bijectors_forward / _inverse` below.

EC code paths do NOT import this module. It is TPD-only.
"""

import math
import torch
import torch.nn as nn


# ----------------------------------------------------------------------------
# Base
# ----------------------------------------------------------------------------

class _Bijector(nn.Module):
    """Scalar elementwise bijector R -> R (or interval -> R).

    Subclasses implement `_forward` and `_inverse`, each returning
    `(value, log_abs_det)` where `log_abs_det` has shape `[batch]`
    (already summed across this parameter's contribution).
    """

    is_identity: bool = False

    def forward(self, x: torch.Tensor):
        return self._forward(x)

    def inverse(self, u: torch.Tensor):
        return self._inverse(u)

    def _forward(self, x):
        raise NotImplementedError

    def _inverse(self, u):
        raise NotImplementedError


# ----------------------------------------------------------------------------
# Identity (default for unbounded params)
# ----------------------------------------------------------------------------

class Identity(_Bijector):
    is_identity = True

    def _forward(self, x):
        return x, torch.zeros(x.shape[0], device=x.device, dtype=x.dtype)

    def _inverse(self, u):
        return u, torch.zeros(u.shape[0], device=u.device, dtype=u.dtype)


# ----------------------------------------------------------------------------
# Logit on [0, 1]
# ----------------------------------------------------------------------------

class Logit(_Bijector):
    """Maps x in (0, 1) -> u = log(x / (1-x)) in R.

    log|du/dx|   = -log(x) - log(1-x)
    log|dx/du|   = -u - 2 * softplus(-u)        (numerically stable form)

    Inputs are clamped to [eps, 1-eps] before forward to avoid +/-inf
    on training data that grazes the boundary (e.g. theta_0 == 1.0 exact).
    """

    def __init__(self, eps: float = 1e-5):
        super().__init__()
        self.eps = eps

    def _forward(self, x):
        x = x.clamp(self.eps, 1.0 - self.eps)
        u = torch.log(x) - torch.log1p(-x)
        # log|du/dx| = -log(x) - log(1-x)
        log_det = -torch.log(x) - torch.log1p(-x)
        return u, log_det

    def _inverse(self, u):
        x = torch.sigmoid(u)
        # log|dx/du| = log(sigmoid(u)) + log(1 - sigmoid(u))
        #            = -u - 2*softplus(-u)   (stable for large |u|)
        log_det = -u - 2.0 * nn.functional.softplus(-u)
        return x, log_det


# ----------------------------------------------------------------------------
# Affine + Logit on [low, high]
# ----------------------------------------------------------------------------

class AffineLogit(_Bijector):
    """Maps x in (low, high) -> u in R.

    Pipeline:  x -> y = (x - low) / (high - low)  -> u = logit(y)

    log|du/dx| = -log(high - low) - log(y) - log(1 - y)
    log|dx/du| =  log(high - low) + log(sigmoid(u)) + log(1 - sigmoid(u))
    """

    def __init__(self, low: float, high: float, eps: float = 1e-5):
        super().__init__()
        assert high > low, f"AffineLogit: high must exceed low, got [{low}, {high}]"
        self.low = float(low)
        self.high = float(high)
        self.scale = float(high - low)
        self.log_scale = math.log(self.scale)
        self.eps = eps

    def _forward(self, x):
        y = (x - self.low) / self.scale
        y = y.clamp(self.eps, 1.0 - self.eps)
        u = torch.log(y) - torch.log1p(-y)
        log_det = (-self.log_scale - torch.log(y) - torch.log1p(-y))
        return u, log_det

    def _inverse(self, u):
        y = torch.sigmoid(u)
        x = self.low + self.scale * y
        log_det = self.log_scale + (-u - 2.0 * nn.functional.softplus(-u))
        return x, log_det


# ----------------------------------------------------------------------------
# Composition over a parameter vector
# ----------------------------------------------------------------------------

def apply_param_bijectors_forward(theta: torch.Tensor, bijectors):
    """Apply per-parameter bijectors elementwise to theta.

    Args:
        theta: [B, D] physical parameters.
        bijectors: list of D bijector instances (one per parameter).

    Returns:
        u:        [B, D] unbounded representation.
        log_det:  [B]    sum of log|du_d/dtheta_d| across d.
    """
    assert theta.shape[1] == len(bijectors), (
        f"theta has {theta.shape[1]} dims but {len(bijectors)} bijectors provided"
    )
    cols = []
    log_det = torch.zeros(theta.shape[0], device=theta.device, dtype=theta.dtype)
    for d, bij in enumerate(bijectors):
        u_d, ld_d = bij.forward(theta[:, d])
        cols.append(u_d)
        log_det = log_det + ld_d
    return torch.stack(cols, dim=1), log_det


def apply_param_bijectors_inverse(u: torch.Tensor, bijectors):
    """Inverse of `apply_param_bijectors_forward`.

    Returns:
        theta:    [B, D] physical parameters.
        log_det:  [B]    sum of log|dtheta_d/du_d| across d.
    """
    assert u.shape[1] == len(bijectors)
    cols = []
    log_det = torch.zeros(u.shape[0], device=u.device, dtype=u.dtype)
    for d, bij in enumerate(bijectors):
        x_d, ld_d = bij.inverse(u[:, d])
        cols.append(x_d)
        log_det = log_det + ld_d
    return torch.stack(cols, dim=1), log_det


def all_identity(bijectors) -> bool:
    """True iff every bijector in the list is an Identity (no-op stack)."""
    return all(getattr(b, 'is_identity', False) for b in bijectors)


# ----------------------------------------------------------------------------
# Smoke test
# ----------------------------------------------------------------------------

if __name__ == "__main__":
    torch.manual_seed(0)

    print("=== Logit roundtrip ===")
    bij = Logit()
    x = torch.rand(8).clamp(0.05, 0.95)
    u, ldf = bij.forward(x)
    x2, ldi = bij.inverse(u)
    print(f"  max |x - x2| = {(x - x2).abs().max().item():.2e}")
    print(f"  ldf + ldi    = {(ldf + ldi).abs().max().item():.2e} (should be 0)")

    print("\n=== AffineLogit on [1, 10] roundtrip ===")
    bij = AffineLogit(1.0, 10.0)
    x = torch.empty(8).uniform_(1.5, 8.0)
    u, ldf = bij.forward(x)
    x2, ldi = bij.inverse(u)
    print(f"  max |x - x2| = {(x - x2).abs().max().item():.2e}")
    print(f"  ldf + ldi    = {(ldf + ldi).abs().max().item():.2e}")

    print("\n=== Per-parameter forward/inverse ===")
    bijs = [Identity(), Logit(), AffineLogit(1.0, 10.0)]
    theta = torch.stack([
        torch.randn(8),
        torch.rand(8).clamp(0.05, 0.95),
        torch.empty(8).uniform_(1.5, 8.0),
    ], dim=1)  # [8, 3]
    u, ldf = apply_param_bijectors_forward(theta, bijs)
    theta2, ldi = apply_param_bijectors_inverse(u, bijs)
    print(f"  max |theta - theta2| = {(theta - theta2).abs().max().item():.2e}")
    print(f"  ldf + ldi (per-row max) = {(ldf + ldi).abs().max().item():.2e}")