"""Strip-DCN feasibility gate (Phase A3, PLAN_code_stripDCN_stripMix).

Validates that the stock compiled extension runs rectangular kernels:
a (1, 9) / (9, 1) strip hits the existing K=9 template instantiation
(switch-K in dcnv4_im2col_cuda.cuh covers {9, 25, 49}; the kernel body is
generic over kernel_h/kernel_w at runtime).

Checks:
  1. fwd correctness vs a pure-PyTorch bilinear reference for (1,9), (9,1)
     and (3,3) sanity, fp32;
  2. bwd correctness (grads wrt value and offset_mask) vs reference autograd;
  3. padding channels of offset_mask (beyond G*K*3) are ignored;
  4. fp16 fwd: finite, close to fp32 reference;
  5. benchmark: strip (1,9) vs square 7 (K=49) vs square 3 (K=9) at the
     SOFIA Stage-1 shape [8, 64*64, 192];
  6. DCNv4Strip module smoke (zero-init => zero output; random init finite).

PASS gate: max|d| < 1e-4 fp32 (fwd and grads), fp16 finite, strip fwd+bwd
speedup >= 2.5x vs square 7.

Run on the training server (CUDA build of the extension required):
    python scripts/test_dcnv4_strip.py
"""

from __future__ import annotations

import time

import torch

from DCNv4 import DCNv4Strip
from DCNv4.functions import DCNv4Function

ATOL_F32 = 1e-4
RTOL_F16 = 3e-2


def ref_dcnv4(
    value: torch.Tensor,        # [N, H, W, G*D]
    offset_mask: torch.Tensor,  # [N, H, W, P] (P >= G*K*3, tail = padding)
    kh: int, kw: int,
    pad_h: int, pad_w: int,
    group: int,
    offset_scale: float = 1.0,
) -> torch.Tensor:
    """Pure-PyTorch reference mirroring dcnv4_im2col_cuda.cuh semantics.

    Point order: outer loop over i in [0, kw), inner over j in [0, kh)
    (m = i*kh + j). Per group g the slab offset_mask[..., g*K*3:(g+1)*K*3]
    holds K interleaved (dx, dy) pairs first, then K mask values (no softmax).
    Sampling location for output pixel (hi, wi):
        x = wi - pad_w + (i + dx) * offset_scale_adj,  analogous for y,
    which for stride=1, dilation=1 reduces to x = wi - pad_w + i + dx
    when offset_scale == 1 (matches p0_w_/p0_h_ algebra in the kernel).
    Bilinear sampling, zero outside the image. Differentiable.
    """
    n, h, w, c = value.shape
    d = c // group
    k = kh * kw
    dev = value.dtype

    # Point grid in CUDA order: m = i*kh + j.
    ii = torch.arange(kw, device=value.device).repeat_interleave(kh)   # [K]
    jj = torch.arange(kh, device=value.device).repeat(kw)              # [K]

    base_y = torch.arange(h, device=value.device).view(1, h, 1, 1)
    base_x = torch.arange(w, device=value.device).view(1, 1, w, 1)

    out = value.new_zeros(n, h, w, group, d)
    half_w = (kw - 1) // 2
    half_h = (kh - 1) // 2
    for g in range(group):
        slab = offset_mask[..., g * k * 3: (g + 1) * k * 3]
        offs = slab[..., : k * 2].reshape(n, h, w, k, 2)
        mask = slab[..., k * 2: k * 3]                                  # [N,H,W,K]
        dx, dy = offs[..., 0], offs[..., 1]                             # [N,H,W,K]

        # p0 - centre*scale + (i*dil + dx)*scale  (stride=1, dil=1)
        x = (base_x - pad_w + half_w).to(dev) \
            + ((ii.view(1, 1, 1, k) - half_w).to(dev) + dx) * offset_scale
        y = (base_y - pad_h + half_h).to(dev) \
            + ((jj.view(1, 1, 1, k) - half_h).to(dev) + dy) * offset_scale

        x0 = torch.floor(x); y0 = torch.floor(y)
        wx1 = x - x0; wy1 = y - y0
        wx0 = 1.0 - wx1; wy0 = 1.0 - wy1

        vg = value[..., g * d: (g + 1) * d]                             # [N,H,W,D]
        acc = value.new_zeros(n, h, w, k, d)
        for oy, wy_ in ((y0, wy0), (y0 + 1, wy1)):
            for ox, wx_ in ((x0, wx0), (x0 + 1, wx1)):
                inside = (oy >= 0) & (oy <= h - 1) & (ox >= 0) & (ox <= w - 1)
                oy_c = oy.clamp(0, h - 1).long()
                ox_c = ox.clamp(0, w - 1).long()
                # Gather vg at [n, oy_c, ox_c] for every (h, w, k).
                idx = (oy_c * w + ox_c).reshape(n, -1)                  # [N, H*W*K]
                flat = vg.reshape(n, h * w, d)
                samp = torch.gather(
                    flat, 1, idx.unsqueeze(-1).expand(-1, -1, d),
                ).reshape(n, h, w, k, d)
                wgt = (wy_ * wx_ * inside.to(dev)).unsqueeze(-1)
                acc = acc + samp * wgt
        out[..., g, :] = (acc * mask.unsqueeze(-1)).sum(dim=3)

    return out.reshape(n, h * w, c)


def run_op(value_seq, offset_mask, kh, kw, pad_h, pad_w, group, gc):
    return DCNv4Function.apply(
        value_seq.view(value_seq.shape[0], int(value_seq.shape[1] ** 0.5),
                       int(value_seq.shape[1] ** 0.5), -1),
        offset_mask,
        kh, kw, 1, 1, pad_h, pad_w, 1, 1,
        group, gc, 1.0, 256, 0,
    )


def check_case(kh: int, kw: int, n=2, h=16, w=16, c=64, group=4) -> None:
    torch.manual_seed(42)
    gc = c // group
    k = kh * kw
    p = ((group * k * 3 + 7) // 8) * 8
    pad_h, pad_w = (kh - 1) // 2, (kw - 1) // 2

    value = torch.randn(n, h, w, c, device="cuda", dtype=torch.float32)
    om = torch.zeros(n, h, w, p, device="cuda", dtype=torch.float32)
    om[..., : group * k * 3] = torch.randn(n, h, w, group * k * 3, device="cuda") * 0.7

    v1 = value.clone().requires_grad_(True)
    o1 = om.clone().requires_grad_(True)
    v2 = value.clone().requires_grad_(True)
    o2 = om.clone().requires_grad_(True)

    out_op = DCNv4Function.apply(
        v1, o1, kh, kw, 1, 1, pad_h, pad_w, 1, 1, group, gc, 1.0, 256, 0,
    ).reshape(n, h * w, c)
    out_ref = ref_dcnv4(v2, o2, kh, kw, pad_h, pad_w, group)

    d_fwd = (out_op - out_ref).abs().max().item()
    assert d_fwd < ATOL_F32, f"({kh}x{kw}) fwd diverges: {d_fwd:.2e}"

    g = torch.randn_like(out_op)
    out_op.backward(g)
    out_ref.backward(g)
    d_gv = (v1.grad - v2.grad).abs().max().item()
    d_go = (o1.grad[..., : group * k * 3] - o2.grad[..., : group * k * 3]).abs().max().item()
    assert d_gv < ATOL_F32, f"({kh}x{kw}) grad_value diverges: {d_gv:.2e}"
    assert d_go < ATOL_F32, f"({kh}x{kw}) grad_offset diverges: {d_go:.2e}"

    # Padding channels must be ignored: poison them, output must not change.
    om_poison = om.clone()
    om_poison[..., group * k * 3:] = 1e6
    out_poison = DCNv4Function.apply(
        value, om_poison, kh, kw, 1, 1, pad_h, pad_w, 1, 1, group, gc, 1.0, 256, 0,
    ).reshape(n, h * w, c)
    d_pad = (out_poison - out_op.detach()).abs().max().item()
    assert d_pad == 0.0, f"({kh}x{kw}) padding channels are NOT ignored: {d_pad:.2e}"

    # fp16 forward: finite + close to fp32 reference.
    out_h = DCNv4Function.apply(
        value.half(), om.half(), kh, kw, 1, 1, pad_h, pad_w, 1, 1,
        group, gc, 1.0, 256, 0,
    ).reshape(n, h * w, c)
    assert torch.isfinite(out_h).all(), f"({kh}x{kw}) fp16 produced non-finite"
    rel = ((out_h.float() - out_ref).abs() / (out_ref.abs() + 1.0)).max().item()
    assert rel < RTOL_F16, f"({kh}x{kw}) fp16 rel err {rel:.2e}"

    print(f"  OK ({kh}x{kw}): fwd {d_fwd:.1e}, dval {d_gv:.1e}, "
          f"doff {d_go:.1e}, pad ignored, fp16 rel {rel:.1e}")


def bench(kh: int, kw: int, label: str, n=8, hw=64, c=192, group=12,
          iters=50) -> float:
    gc = c // group
    k = kh * kw
    p = ((group * k * 3 + 7) // 8) * 8
    value = torch.randn(n, hw, hw, c, device="cuda", requires_grad=True)
    om = torch.randn(n, hw, hw, p, device="cuda", requires_grad=True)
    args = (kh, kw, 1, 1, (kh - 1) // 2, (kw - 1) // 2, 1, 1, group, gc, 1.0, 256, 0)

    for _ in range(10):  # warmup
        out = DCNv4Function.apply(value, om, *args)
        out.sum().backward()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        out = DCNv4Function.apply(value, om, *args)
        out.sum().backward()
    torch.cuda.synchronize()
    ms = (time.perf_counter() - t0) / iters * 1e3
    print(f"  {label:14s} K={k:2d}: {ms:7.3f} ms/iter (fwd+bwd, "
          f"[{n},{hw}x{hw},{c}], G={group})")
    return ms


def main() -> None:
    assert torch.cuda.is_available(), "CUDA required"
    print("== correctness ==")
    check_case(3, 3)            # sanity: known-good square path (K=9)
    check_case(1, 9)            # strip-H (K=9, reuses square-3 instantiation)
    check_case(9, 1)            # strip-V
    check_case(1, 9, c=192, group=12)  # SOFIA Stage-1 width

    print("== DCNv4Strip module ==")
    m = DCNv4Strip(channels=192, k=9, orientation="h", group=12).cuda()
    assert (m.kernel_h, m.kernel_w) == (1, 9), "orientation='h' must give (1,9)"
    mv = DCNv4Strip(channels=192, k=9, orientation="v", group=12).cuda()
    assert (mv.kernel_h, mv.kernel_w) == (9, 1), "orientation='v' must give (9,1)"
    x = torch.randn(2, 64 * 64, 192, device="cuda")
    y = m(x, shape=(64, 64))
    assert y.shape == x.shape
    assert y.abs().max().item() < 1e-6, "zero-init must give ~zero output"
    torch.nn.init.normal_(m.offset_mask.weight, std=0.02)
    torch.nn.init.normal_(m.offset_mask.bias, std=0.02)
    x = x.requires_grad_(True)
    y = m(x, shape=(64, 64))
    assert torch.isfinite(y).all() and y.abs().max().item() > 0
    y.sum().backward()
    assert x.grad is not None and torch.isfinite(x.grad).all()
    print("  OK: module smoke (orientations, zero-init ~0, random init finite, bwd)")

    print("== benchmark (SOFIA Stage-1 shape) ==")
    t_strip = bench(1, 9, "strip (1,9)")
    t_sq3 = bench(3, 3, "square 3")
    t_sq7 = bench(7, 7, "square 7")
    speedup = t_sq7 / t_strip
    print(f"  strip vs square7 fwd+bwd speedup: {speedup:.2f}x "
          f"(target >= 2.5x, hard gate >= 1.8x), vs square3: {t_sq3 / t_strip:.2f}x")
    if speedup < 2.5:
        print(f"  WARNING: speedup {speedup:.2f}x below 2.5x target "
              f"(noisy benchmark? rerun on idle GPU)")
    assert speedup >= 1.8, "GATE FAILED: strip speedup below hard gate 1.8x"

    print("\nALL GATES PASSED — strip-DCN is viable on the stock binary.")


if __name__ == "__main__":
    main()