"""Control-map builders for the structured-noise phase slot (``rgb{v}_edge_control``).

Sits next to ``custom/wan/structured_noise.py``: that module owns the K-subspace math and the
per-step source PICK (``sn_canonical_mix`` / ``sn_pick_modality``); this module owns turning a raw
anydata modality into the pseudo-RGB [0,1] control video that fills the phase slot —
  - RGB  -> Canny edge map      (``compute_canny_edge_control``)
  - depth / colorized semantic  (``compute_depth_control``, a generic pseudo-RGB->[0,1] normalizer)
plus the dataloader entry point ``populate_edge_control`` that picks a source and fills the slot.

Extracted verbatim from ``custom/dataloader/unified_anydrive.py`` so ``anydata_to_any4d`` doesn't
keep growing (and to shrink the merge surface on that hot file). Behavior is unchanged.

The last picked source is published via ``set_last_sn_pick`` and read back by the model/val pipe
(``get_last_sn_pick``) to select the matching per-modality projection P_m — done on the same
rank/thread within one forward, so it's reliable and availability-proof.
"""

import numpy as np
import torch

# Last structured-noise source modality picked for the phase slot (per rank/process).
_LAST_SN_PICK = {'modality': None, 'phase': None, 'train_step': -1, 'val_step': -1}


def set_last_sn_pick(modality, phase=None, train_step=-1, val_step=-1):
    _LAST_SN_PICK.update(modality=modality, phase=phase,
                         train_step=int(train_step), val_step=int(val_step))


def get_last_sn_pick():
    """Return the most recent structured-noise source modality picked (canonical name) or None."""
    return _LAST_SN_PICK.get('modality')


def compute_canny_edge_control(anydata_rgb, low_range=(50.0, 130.0), high_range=(150.0, 280.0)):
    """Derive a Canny edge-map control video from RGB frames (Wan Fun-InP+canny).

    Returns a modality dict with the SAME ``(time, camera) -> (B, 3, H, W)`` structure
    as ``anydata_rgb``, where each frame is the Canny edge map (replicated to 3 channels,
    range [0, 1]). A random (low, high) threshold pair is drawn per batch element and
    shared across that element's frames/cameras, so the edge style stays temporally
    consistent within a clip while varying across samples (data augmentation).

    :param anydata_rgb: dict mapping (t, cam) -> (B, 3, H, W) float RGB in [0, 1].
    :return: dict mapping (t, cam) -> (B, 3, H, W) float canny pseudo-RGB in [0, 1].
    """
    import random as _random

    import cv2

    keys = list(anydata_rgb.keys())
    if not keys:
        return None
    B = anydata_rgb[keys[0]].shape[0]
    # One random Canny threshold pair per batch element (consistent across frames/cams).
    thresholds = [
        (_random.uniform(*low_range), _random.uniform(*high_range)) for _ in range(B)
    ]
    edge_control = {}
    for key, frame in anydata_rgb.items():
        dev, dt = frame.device, frame.dtype
        f = frame.detach().float().clamp(0.0, 1.0).cpu().numpy()  # (B, 3, H, W)
        out = np.zeros_like(f)
        for b in range(B):
            lo, hi = thresholds[b]
            img = (f[b].transpose(1, 2, 0) * 255.0).astype(np.uint8)  # (H, W, 3) RGB
            gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
            edge = cv2.Canny(gray, lo, hi).astype(np.float32) / 255.0  # (H, W) in {0, 1}
            out[b] = np.repeat(edge[None], 3, axis=0)  # (3, H, W) pseudo-RGB
        edge_control[key] = torch.from_numpy(out).to(device=dev, dtype=dt)
    return edge_control


def compute_depth_control(anydata_depth, anydata_rgb):
    """Route a precomputed grayscale DEPTH video modality into the structured-noise
    phase-source slot (rgb{v}_edge_control), mirroring compute_canny_edge_control's output
    contract: dict (t, cam) -> (B, 3, H, W) float in [0, 1] on the rgb device/dtype.

    AnyData loads the depth mp4 frame-aligned to rgb (shared indexer) and resized to the rgb
    resolution, but channel-last (B, H, W, 3) float in ~[0, 255]. We move it to channel-first,
    normalize to [0, 1], replicate to 3 channels if needed, and (defensively) resize to the
    rgb spatial size. The depth is consumed ONLY as the structured-noise phase source (same as
    canny): the DiT never sees it as an explicit input, and it is excluded from the loss.
    """
    out = {}
    for key, rgb in anydata_rgb.items():
        d = anydata_depth.get(key)
        if d is None:
            continue
        dev, dt = rgb.device, rgb.dtype
        H, W = rgb.shape[-2], rgb.shape[-1]
        d = torch.as_tensor(d).float()
        if d.ndim == 4 and d.shape[-1] in (1, 3) and d.shape[1] not in (1, 3):
            d = d.permute(0, 3, 1, 2)            # (B, H, W, C) -> (B, C, H, W)
        elif d.ndim == 3:
            d = d.unsqueeze(1)                   # (B, H, W) -> (B, 1, H, W)
        if d.shape[1] == 1:
            d = d.repeat(1, 3, 1, 1)
        if float(d.max()) > 1.5:
            d = d / 255.0
        if tuple(d.shape[-2:]) != (H, W):
            d = torch.nn.functional.interpolate(d, size=(H, W), mode='bilinear', align_corners=False)
        out[key] = d.clamp(0.0, 1.0).to(device=dev, dtype=dt)
    return out


def populate_edge_control(anydata_dict, config, phase, train_step=-1, val_step=-1):
    """Fill the structured-noise phase-source slot ``rgb_edge_control`` when the experiment declares
    edge_control entries but no precomputed control modality is present, and return the resulting
    control dict (also stored into ``anydata_dict['rgb_edge_control']``); returns None / the
    pre-existing value otherwise.

    Priority: multi-source ``structured_noise_source_mix`` (per-step pick) > DEPTH (depth-phase)
    > SEMANTIC (seg-phase, a colorized label-map pseudo-RGB) > derive Canny from the RGB
    (canny-phase). ``compute_depth_control`` is a generic pseudo-RGB->[0,1] normalizer reused for
    the colorized semantic modality.
    """
    anydata_rgb_edge_control = anydata_dict.get('rgb_edge_control', None)
    if (anydata_rgb_edge_control is None
            and any('edge_control' in e for e in config.all_highdim_entries)):
        anydata_rgb = anydata_dict.get('rgb', None)
        anydata_depth = anydata_dict.get('depth', None)
        anydata_semantic = anydata_dict.get('semantic', None)
        # Multi-source phase mixture: fill the single phase slot from ONE source chosen per
        # step (TRAIN: uniform random keyed on train_step; VAL: round-robin keyed on val_step
        # so curves stay comparable). See Any4DConfig.structured_noise_source_mix.
        source_mix = getattr(config, 'structured_noise_source_mix', None)
        if source_mix:
            from custom.wan.structured_noise import sn_canonical_mix, sn_pick_modality
            # name -> zero-arg builder, only for sources actually available this batch.
            builders = {}
            if anydata_depth is not None:
                builders['depth'] = lambda: compute_depth_control(anydata_depth, anydata_rgb)
            if anydata_semantic is not None:
                builders['seg'] = lambda: compute_depth_control(anydata_semantic, anydata_rgb)
            if anydata_rgb is not None:
                builders['canny'] = lambda: compute_canny_edge_control(anydata_rgb)
                # 'rgb' == image-phase: route the clean RGB ([0,1], 3ch) into the slot so it
                # VAE-encodes to ~the rgb latent (the DiT still never sees it as an input).
                builders['rgb'] = lambda: {
                    k: v.detach().float().clamp(0.0, 1.0) for k, v in anydata_rgb.items()}
            wanted = sn_canonical_mix(source_mix)
            avail = [s for s in wanted if s in builders]
            assert avail, (
                f'structured_noise_source_mix={source_mix} but none available '
                f'(have {sorted(builders)}); check dataset labels.')
            pick = sn_pick_modality(avail, train_step=(int(train_step) if phase == 'train' else -1),
                                    val_step=(int(val_step) if phase != 'train' else -1))
            # Publish the chosen source so the model/pipe selects the matching per-modality P_m.
            # Same rank/thread as the upcoming compute_loss within this forward, so it's reliable
            # and availability-proof (no re-derivation that could desync when a modality is absent).
            set_last_sn_pick(pick, phase=phase, train_step=int(train_step), val_step=int(val_step))
            anydata_rgb_edge_control = builders[pick]()
        elif anydata_depth is not None:
            anydata_rgb_edge_control = compute_depth_control(anydata_depth, anydata_rgb)
        elif anydata_semantic is not None:
            anydata_rgb_edge_control = compute_depth_control(anydata_semantic, anydata_rgb)
        elif anydata_rgb is not None:
            anydata_rgb_edge_control = compute_canny_edge_control(anydata_rgb)
        if anydata_rgb_edge_control is not None:
            anydata_dict['rgb_edge_control'] = anydata_rgb_edge_control
    return anydata_rgb_edge_control