"""Consistency-Distillation / Consistency-Training generator wrapper for CIFAR10.

Wraps OpenAI's CIFAR10 consistency model (`consistency_models_cifar10`, package
`jcm`) so it can stand in as a generator with a decode-style API: build once,
then `decode(z)` maps a noise-space input to a 32x32x3 image in [0, 1].

Latents are flat vectors of shape `(batch, 3072)` (== 32*32*3), drawn from
N(0, sigma_max**2 * I) and reshaped internally.

Only the generator path is exercised - LPIPS / classifier heads / training-time
losses are not imported. `distiller_fn` math is inlined (KVE-SDE Karras et al.
2022 conditioning) so we don't depend on `jcm.sde_lib`.
"""

from pathlib import Path

import jax
import jax.numpy as jnp
import ml_collections

# Registers "ncsnpp" in jcm.models.utils._MODELS via decorator side-effect.
from jcm.models import ncsnpp  # noqa: F401
from jcm.models import utils as mutils

Array = jax.Array

IMAGE_SIZE = 32
NUM_CHANNELS = 3
NOISE_DIM = IMAGE_SIZE * IMAGE_SIZE * NUM_CHANNELS  # 3072
SIGMA_MAX = 80.0
SIGMA_MIN = 0.002
DATA_STD = 0.5


def _default_config() -> ml_collections.ConfigDict:
    """NCSN++ config used by the CT-LPIPS CIFAR10 checkpoint."""
    c = ml_collections.ConfigDict()
    c.data = ml_collections.ConfigDict(
        dict(
            dataset="CIFAR10",
            image_size=IMAGE_SIZE,
            num_channels=NUM_CHANNELS,
            random_flip=False,
            uniform_dequantization=False,
        )
    )
    c.model = ml_collections.ConfigDict(
        dict(
            name="ncsnpp",
            ema_rate=0.9999,
            normalization="GroupNorm",
            nonlinearity="swish",
            nf=128,
            ch_mult=(2, 2, 2),
            num_res_blocks=4,
            attn_resolutions=(16,),
            resamp_with_conv=True,
            conditional=True,
            fir=True,
            fir_kernel=[1, 3, 3, 1],
            skip_rescale=True,
            resblock_type="biggan",
            progressive="none",
            progressive_input="residual",
            progressive_combine="sum",
            attention_type="ddpm",
            embedding_type="fourier",
            init_scale=0.0,
            fourier_scale=16,
            conv_size=3,
            rho=7.0,
            data_std=DATA_STD,
            num_scales=18,
            dropout=0.0,
            sigma_min=0.02,
            sigma_max=100,
            beta_min=0.1,
            beta_max=20.0,
            t_min=SIGMA_MIN,
            t_max=SIGMA_MAX,
            double_heads=False,
        )
    )
    c.training = ml_collections.ConfigDict(dict(sde="kvesde"))
    c.sampling = ml_collections.ConfigDict(dict(method="onestep", std=SIGMA_MAX))
    return c


def _build_model_and_shapes(config: ml_collections.ConfigDict, seed: int):
    import functools

    model_def = functools.partial(mutils.get_model("ncsnpp"), config=config)
    model = model_def()
    rng = jax.random.PRNGKey(seed)
    p_rng, d_rng = jax.random.split(rng)
    fake_x = jnp.zeros((1, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS))
    fake_t = jnp.zeros((1,), dtype=jnp.float32)
    variables = model.init({"params": p_rng, "dropout": d_rng}, fake_x, fake_t)
    return model, variables["params"]


class CD:
    """Consistency-Distillation / Consistency-Training generator.

    Instantiate with `pretrained` naming a checkpoint kind and an explicit
    `checkpoint_path` pointing at the slim pickle bundle
    (`{"params", "meta"}`). Latents are flat vectors of shape `(batch, 3072)`
    for consistency with the target's expected `n_dim`; the wrapper reshapes
    internally. `decode(z)` returns 32x32x3 images in [0, 1].
    """

    z_dim: int = NOISE_DIM
    image_size: int = IMAGE_SIZE
    num_channels: int = NUM_CHANNELS

    def __init__(
        self,
        pretrained: str = "ct-lpips",
        checkpoint_path: str | Path | None = None,
        seed: int = 0,
    ) -> None:
        if pretrained != "ct-lpips":
            raise ValueError(
                f"only 'ct-lpips' is wired up; got {pretrained!r}. "
                "Add another config in `_default_config` to support others."
            )
        if checkpoint_path is None:
            raise ValueError("checkpoint_path is required (path to the slim .pkl bundle)")

        self.pretrained = pretrained
        self.sigma_max = SIGMA_MAX
        self.sigma_min = SIGMA_MIN
        self.data_std = DATA_STD

        self.config = _default_config()
        self.model, init_params = _build_model_and_shapes(self.config, seed)

        import pickle

        from flax import serialization

        with open(checkpoint_path, "rb") as f:
            bundle = pickle.load(f)
        self.params = serialization.from_state_dict(init_params, bundle["params"])

        def _decode(z_flat: Array) -> Array:
            batch = z_flat.shape[0]
            # Reshape flat noise vector to image; scale by sigma_max (prior std).
            x = z_flat.reshape(batch, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS) * SIGMA_MAX
            img = _distill_one_step(self.model, self.params, x, SIGMA_MAX)
            # Consistency models output pixels in ~[-1, 1]; map to [0, 1].
            return jnp.clip((img + 1.0) * 0.5, 0.0, 1.0)

        self._decode = jax.jit(_decode)

    def decode(self, z: Array) -> Array:
        return self._decode(z)

    def sample_prior(self, key: Array, n: int) -> Array:
        z = jax.random.normal(key, (n, self.z_dim))
        return self.decode(z)


def _distill_one_step(model, params, x: Array, sigma: float) -> Array:
    """Inline KVE-SDE consistency-model one-step distillation (Karras et al. 2022).

    x: (B, H, W, C) noise-scaled inputs (i.e. x ~ N(0, sigma**2 * I) reshaped).
    sigma: scalar noise level (== SIGMA_MAX for one-step from prior).
    """
    t = jnp.asarray(sigma, dtype=x.dtype)
    pred_t = SIGMA_MIN
    d2 = DATA_STD * DATA_STD
    in_scale = 1.0 / jnp.sqrt(t * t + d2)
    cond_t = 0.25 * jnp.log(t)
    cond_batch = jnp.full((x.shape[0],), cond_t, dtype=x.dtype)
    raw = model.apply({"params": params}, x * in_scale, cond_batch, train=False)
    out_scale = (t - pred_t) * DATA_STD / jnp.sqrt(t * t + d2)
    skip_scale = d2 / ((t - pred_t) ** 2 + d2)
    return skip_scale * x + out_scale * raw