"""Inference: multi-condition classifier-free guidance + ODE (docs/04 "Çok koşullu CFG").

Three passes per step, batched together:

* ``v∅``  — unconditional: text dropped, emotion/intensity null (the prompt policy is a
  parameter: ``drop_prompt_uncond``);
* ``vT``  — text + tags, emotion null;
* ``vTE`` — text + tags + emotion.

``v = v∅ + w_text · (vT − v∅) + w_emo · (vTE − vT)``. With ``w_text = w_emo = 1`` this is
exactly ``vTE``. ``w_emo`` is the user-facing "emotion intensity" knob.
"""

from __future__ import annotations

from dataclasses import replace
from typing import Literal

import torch
from torch import Tensor

from nida.model.dit import Conditions, DiT
from nida.model.flow import sample_ode

__all__ = ["guided_velocity", "synthesize_latents"]


def _null_emotion(model: DiT, c: Conditions) -> Conditions:
    batch = c.emotion.shape[0]
    return replace(
        c,
        emotion=torch.full(
            (batch,), model.cfg.n_emotions, dtype=c.emotion.dtype, device=c.emotion.device
        ),
        intensity_known=torch.zeros(batch, dtype=torch.bool, device=c.emotion.device),
    )


def _concat(parts: list[Conditions]) -> Conditions:
    """Concatenate along the batch; an optional field must be set in all parts or none."""
    values: dict[str, Tensor | None] = {}
    for name in Conditions.__dataclass_fields__:
        tensors = [getattr(p, name) for p in parts]
        if all(x is None for x in tensors):
            values[name] = None
        elif any(x is None for x in tensors):
            raise ValueError(f"Conditions field {name!r} is set in some parts only")
        else:
            values[name] = torch.cat(tensors)
    return Conditions(**values)  # type: ignore[arg-type]  # required fields are never None


def guided_velocity(
    model: DiT,
    x_t: Tensor,
    t: Tensor,
    c: Conditions,
    *,
    w_text: float,
    w_emo: float,
    drop_prompt_uncond: bool = False,
) -> Tensor:
    """Guided velocity ``(B, T, C)`` from one batched forward of ``3·B`` items."""
    batch = x_t.shape[0]
    c_te = c
    c_t = _null_emotion(model, c)
    c_u = replace(c_t, drop_text=torch.ones(batch, dtype=torch.bool, device=x_t.device))
    if drop_prompt_uncond:
        c_u = replace(c_u, cond_latent=torch.zeros_like(c.cond_latent))
    v = model(torch.cat([x_t] * 3), torch.cat([t] * 3), _concat([c_u, c_t, c_te]))
    v_u, v_t, v_te = v.split(batch)
    guided: Tensor = v_u + w_text * (v_t - v_u) + w_emo * (v_te - v_t)
    return guided


@torch.no_grad()
def synthesize_latents(
    model: DiT,
    c: Conditions,
    prompt_mask: Tensor,
    *,
    steps: int,
    w_text: float,
    w_emo: float,
    seed: int,
    method: Literal["euler", "midpoint"] = "euler",
    sway: float = -1.0,
    drop_prompt_uncond: bool = False,
) -> Tensor:
    """Generate latents ``(B, T, C)``.

    ``prompt_mask`` ``(B, T)`` bool, True = prompt frame: copied from ``cond_latent`` in the
    output (they are context, not generated). The noise comes from a CPU generator seeded
    with ``seed`` (callers always record it).
    """
    batch, length = c.frame_mask.shape
    if prompt_mask.shape != (batch, length) or prompt_mask.dtype != torch.bool:
        raise ValueError(f"prompt_mask must be bool ({batch}, {length})")
    device = c.cond_latent.device
    generator = torch.Generator(device="cpu").manual_seed(seed)
    x0 = torch.randn(batch, length, model.cfg.latent_dim, generator=generator)
    x0 = x0.to(device=device, dtype=c.cond_latent.dtype)
    model.eval()
    out = sample_ode(
        lambda x, t: guided_velocity(
            model, x, t, c, w_text=w_text, w_emo=w_emo, drop_prompt_uncond=drop_prompt_uncond
        ),
        x0,
        steps=steps,
        method=method,
        sway=sway,
    )
    result: Tensor = torch.where(prompt_mask[..., None], c.cond_latent, out)
    result = result * c.frame_mask[..., None]
    return result