"""d1_lune_sampler.py — SD15-Lune rectified-flow sampler + 225-token CLIP encoder.

PORTED VERBATIM (read-only law) from E:\\coworkers\\qwen_runner_1\\colab\\
sd15_lune_before_after.py (the sampler of record — repos/geolip-aleph-diffusion.md):
  encode_clip_225 — 3 x 75-token chunked CLIP -> [B, 227, 768]. ALL judging in
                    this line conditions through THIS, never a bare 77-token encode.
  flow_sample     — Euler ODE sigma 1->0 on the SHIFT-warped grid
                    sigma = (shift*u)/(1+(shift-1)*u), v = noise - x0, batched CFG.
                    NOTE: pod/v35_expP1_sd15.py used an UNSHIFTED grid — that
                    discrepancy is ledgered in R0a; THIS grid matches the trainer.
Defaults from the source bed: STEPS 30, SHIFT 2.5, VAE_SCALE 0.18215, 64x64
latents (512px). Guidance differs across Phil's beds (6.0 before/after bed,
4.0 expP1) — beds pass it explicitly and ledger it.

CPU cert: python pod2/d1_lune_sampler.py  (sigma-grid math + dummy-UNet
integration shape check only; real sampling is pod work).
"""
from __future__ import annotations

import torch

VAE_SCALE = 0.18215
SHIFT = 2.5
STEPS = 30


def encode_clip_225(prompts, tokenizer, text_encoder, device):
    """Encode text as 3 x 75-token chunks -> one 227-length hidden sequence."""
    chunk_len = tokenizer.model_max_length          # 77
    body_len = chunk_len - 2                        # 75
    n_chunks = 3
    ids = tokenizer(prompts, padding="max_length",
                    max_length=body_len * n_chunks + 2,
                    truncation=True, return_tensors="pt").input_ids   # [B, 227]
    bos, eos = ids[:, :1], ids[:, -1:]
    chunks = []
    for k in range(n_chunks):
        s = 1 + k * body_len
        chunks.append(torch.cat([bos, ids[:, s:s + body_len], eos], dim=1))
    ids = torch.stack(chunks, dim=1).reshape(-1, chunk_len).to(device)
    hs = text_encoder(ids)[0]                                   # [B*3, 77, 768]
    hs = hs.reshape(len(prompts), n_chunks * chunk_len, -1)
    out = [hs[:, :1]]
    for k in range(n_chunks):
        s = k * chunk_len + 1
        out.append(hs[:, s:s + body_len])
    out.append(hs[:, -1:])
    return torch.cat(out, dim=1)                                # [B, 227, 768]


def shifted_sigmas(n_steps: int, shift: float = SHIFT,
                   device="cpu") -> torch.Tensor:
    u = torch.linspace(1.0, 0.0, n_steps + 1, device=device)
    return (shift * u) / (1 + (shift - 1) * u)


@torch.no_grad()
def flow_sample(unet, ehs_cond, *, n_steps=STEPS, guidance=6.0, shift=SHIFT,
                seed=1234, device="cuda", latent_hw=64,
                encoder_attention_mask=None):
    """Integrate the rectified-flow ODE sigma 1->0; unet predicts v = noise - x0.
    encoder_attention_mask (optional, [2B, L]): the Tier-A masked-append toggle
    path — passed through to the UNet when provided."""
    B = ehs_cond.shape[0]
    g = torch.Generator(device=device).manual_seed(seed)
    x = torch.randn(B, 4, latent_hw, latent_hw, generator=g, device=device,
                    dtype=torch.float32).to(ehs_cond.dtype)   # noise drawn fp32
    # for seed comparability across dtypes, then cast to the UNet's dtype
    ehs_in = torch.cat([ehs_cond, torch.zeros_like(ehs_cond)], dim=0)
    sigmas = shifted_sigmas(n_steps, shift, device)
    kw = {}
    if encoder_attention_mask is not None:
        kw["encoder_attention_mask"] = encoder_attention_mask
    for i in range(n_steps):
        s, s_next = float(sigmas[i]), float(sigmas[i + 1])
        t = torch.full((2 * B,), s * 1000.0, device=device)
        v = unet(torch.cat([x, x], dim=0), t, ehs_in, return_dict=False,
                 **kw)[0]
        v_cond, v_uncond = v.chunk(2)
        v = v_uncond + guidance * (v_cond - v_uncond)
        x = x + (s_next - s) * v                                # Euler (sigma down)
    return x


@torch.no_grad()
def decode(vae, latents):
    img = vae.decode(latents / VAE_SCALE).sample
    img = (img / 2 + 0.5).clamp(0, 1)
    return img.permute(0, 2, 3, 1).float().cpu().numpy()


def _smoke():
    s = shifted_sigmas(30)
    assert abs(float(s[0]) - 1.0) < 1e-6 and abs(float(s[-1])) < 1e-6, \
        "shifted grid must span 1 -> 0 exactly"
    assert (s[:-1] > s[1:]).all(), "sigmas must be strictly decreasing"
    mid_shifted = float(shifted_sigmas(2)[1])       # u=0.5 -> 2.5*.5/1.75
    assert abs(mid_shifted - (1.25 / 1.75)) < 1e-6, "SHIFT=2.5 warp wrong"

    class DummyUNet:
        def __call__(self, x, t, ehs, return_dict=False, **kw):
            return (torch.zeros_like(x),)

    out = flow_sample(DummyUNet(), torch.zeros(2, 227, 768), n_steps=4,
                      guidance=1.0, seed=0, device="cpu")
    assert out.shape == (2, 4, 64, 64)
    print("d1_lune_sampler smoke PASSED (shifted grid endpoints + monotone, "
          "warp value, dummy integration shape)")


if __name__ == "__main__":
    _smoke()