"""Checkpoint surgery to enable SCAIL-2 in-context mask channels on a trained LTX-2 model.

Widening ``LTXModel.mask_conditioning_channels`` from 0 to ``C`` grows the video
``patchify_proj`` weight from ``[inner, in]`` to ``[inner, in + C]``. This helper
appends ``C`` **zero** input columns to that weight so a converted checkpoint is
numerically identical to the original until the new columns are finetuned -- the
extra channels contribute nothing at load time.

Usage::

    sd = load_state_dict(path)
    widen_patchify_proj_for_mask_channels(sd, mask_channels=56)
    model = LTXModel(..., mask_conditioning_channels=56)
    model.load_state_dict(sd)
"""

from __future__ import annotations

import torch

_VIDEO_PROJ_SUFFIX = "patchify_proj.weight"
_AUDIO_PROJ_SUFFIX = "audio_patchify_proj.weight"


def widen_module_patchify_proj_for_mask_channels(model: torch.nn.Module, mask_channels: int) -> torch.nn.Module:
    """In place, widen a live ``LTXModel``'s video ``patchify_proj`` to accept ``mask_channels`` extra inputs.

    Replaces ``model.patchify_proj`` with a wider ``Linear`` whose original input columns are copied and
    whose ``mask_channels`` new columns are zero (so the model reproduces its pre-widening output until
    those columns are trained). Also sets ``model.mask_conditioning_channels`` for metadata/consistency.
    Idempotent guard: raises if the model is already widened to a different value.
    """
    if mask_channels < 0:
        raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")

    proj = getattr(model, "patchify_proj", None)
    if not isinstance(proj, torch.nn.Linear):
        raise AttributeError("model has no linear 'patchify_proj' to widen")

    already = int(getattr(model, "mask_conditioning_channels", 0))
    if mask_channels in (0, already):
        model.mask_conditioning_channels = mask_channels
        return model
    if already != 0:
        raise ValueError(f"patchify_proj already widened for {already} mask channels, refusing to re-widen")

    new_in = proj.in_features + mask_channels
    wider = torch.nn.Linear(new_in, proj.out_features, bias=proj.bias is not None)
    wider = wider.to(device=proj.weight.device, dtype=proj.weight.dtype)
    with torch.no_grad():
        wider.weight.zero_()
        wider.weight[:, : proj.in_features].copy_(proj.weight)
        if proj.bias is not None:
            wider.bias.copy_(proj.bias)
    model.patchify_proj = wider
    model.mask_conditioning_channels = mask_channels
    return model


def widen_patchify_proj_for_mask_channels(
    state_dict: dict[str, torch.Tensor],
    mask_channels: int,
) -> dict[str, torch.Tensor]:
    """In place, append ``mask_channels`` zero input columns to the video ``patchify_proj`` weight.

    Only the video projection is touched (audio ``audio_patchify_proj`` is left alone). The bias,
    if present, is unchanged. Returns the same dict for convenience. Idempotency is the caller's
    responsibility -- calling twice widens twice.
    """
    if mask_channels < 0:
        raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")
    if mask_channels == 0:
        return state_dict

    target_keys = [
        key
        for key in state_dict
        if key.endswith(_VIDEO_PROJ_SUFFIX) and not key.endswith(_AUDIO_PROJ_SUFFIX)
    ]
    if not target_keys:
        raise KeyError(
            f"No '{_VIDEO_PROJ_SUFFIX}' weight found in the state dict; cannot widen for mask channels."
        )

    for key in target_keys:
        weight = state_dict[key]  # [inner_dim, in_features]
        if weight.ndim != 2:
            raise ValueError(f"Expected 2-D weight for '{key}', got shape {tuple(weight.shape)}")
        pad = weight.new_zeros(weight.shape[0], mask_channels)
        state_dict[key] = torch.cat([weight, pad], dim=1)

    return state_dict
