"""INT8 DCNv4 gate (Level 1: int8 storage + fp32 arithmetic).

Validates ext.dcnv4_int8_forward against the fp32 op on dequantized inputs:
the int8 path must be bit-equivalent to dequantize -> fp32 DCNv4 -> quantize,
so the only allowed error is output rounding (<= 0.5 quantization steps,
plus tiny fp accumulation noise).

Checks:
  1. correctness vs fp32 op for square (3,3), strip (1,9)/(9,1), SOFIA
     Stage-1 width (C=192, G=12);
  2. padding channels of offset_mask are ignored;
  3. DCNv4Strip.forward_int8 == dequant->fp forward->quant (within 1 step);
  4. benchmark int8 vs fp16 op (memory-bound shapes) — informational.

PASS gate: max |int8_out - quantize(fp32_out)| <= 1 LSB; >=99% exact match.

Run on the training server after rebuilding the extension:
    cd DCNv4_op && pip install -e . && python scripts/test_dcnv4_int8.py
"""

from __future__ import annotations

import time

import torch

from DCNv4 import DCNv4Strip, ext
from DCNv4.functions import DCNv4Function, dcnv4_int8_forward


def quantize(x: torch.Tensor, scale: float) -> torch.Tensor:
    return torch.clamp(torch.round(x / scale), -128, 127).to(torch.int8)


def run_fp32(value_f32, om_f32, kh, kw, pad_h, pad_w, group, gc, offset_scale=1.0):
    return DCNv4Function.apply(
        value_f32, om_f32, kh, kw, 1, 1, pad_h, pad_w, 1, 1,
        group, gc, offset_scale, 256, 0,
    )


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

    value_f = torch.randn(n, h, w, c, device="cuda") * 2.0
    v_scale = value_f.abs().max().item() / 127.0
    v_i8 = quantize(value_f, v_scale)
    v_deq = v_i8.float() * v_scale                      # exact int8 grid

    om = torch.zeros(n, h, w, p, device="cuda")
    om[..., : group * k * 3] = torch.randn(n, h, w, group * k * 3,
                                           device="cuda") * 0.7
    om_h = om.half()  # int8 path consumes fp16 offsets — give the fp32
    #                   reference the SAME rounded values for a clean compare

    ref = run_fp32(v_deq, om_h.float(), kh, kw, pad_h, pad_w, group, gc,
                   offset_scale)
    o_scale = ref.abs().max().item() / 127.0
    ref_q = quantize(ref, o_scale)

    out_i8 = dcnv4_int8_forward(
        v_i8, om_h, kh, kw, pad_h, pad_w, group, gc,
        offset_scale, v_scale, o_scale,
    ).reshape(ref_q.shape)

    diff = (out_i8.int() - ref_q.int()).abs()
    max_lsb = diff.max().item()
    exact = (diff == 0).float().mean().item()
    # fp16 offsets in the int8 path vs fp32 in the reference cause sub-LSB
    # coordinate differences; allow 1 LSB and require >=99% exact agreement.
    assert max_lsb <= 1, f"({kh}x{kw}) int8 deviates by {max_lsb} LSB"
    assert exact >= 0.99, f"({kh}x{kw}) only {exact:.4f} exact"
    print(f"  OK ({kh}x{kw}, C={c}, G={group}, os={offset_scale}): "
          f"max {max_lsb} LSB, exact {exact:.4f}")

    # Padding channels must be ignored.
    om_p = om.half().clone()
    om_p[..., group * k * 3:] = 30000.0
    out_p = dcnv4_int8_forward(
        v_i8, om_p, kh, kw, pad_h, pad_w, group, gc,
        offset_scale, v_scale, o_scale,
    )
    assert torch.equal(out_p, out_i8.reshape(out_p.shape)), \
        f"({kh}x{kw}) padding channels are NOT ignored"


def check_module() -> None:
    torch.manual_seed(0)
    m = DCNv4Strip(channels=192, k=9, orientation="h", group=12).cuda().half()
    torch.nn.init.normal_(m.offset_mask.weight, std=0.02)
    torch.nn.init.normal_(m.offset_mask.bias, std=0.02)

    x_f = torch.randn(2, 64 * 64, 192, device="cuda") * 1.5
    v_scale = x_f.abs().max().item() / 127.0
    x_i8 = quantize(x_f, v_scale)
    x_deq = (x_i8.half() * v_scale)

    ref = m(x_deq, shape=(64, 64)).float()
    o_scale = max(ref.abs().max().item() / 127.0, 1e-8)
    ref_q = quantize(ref, o_scale)

    out = m.forward_int8(x_i8, (64, 64), v_scale, o_scale)
    diff = (out.int() - ref_q.int()).abs()
    assert diff.max().item() <= 1, f"module int8 deviates {diff.max().item()} LSB"
    print(f"  OK DCNv4Strip.forward_int8: max {diff.max().item()} LSB, "
          f"exact {(diff == 0).float().mean().item():.4f}")


def bench(n=8, hw=64, c=192, group=12, k=9, iters=100) -> None:
    gc = c // group
    p = ((group * k * 3 + 7) // 8) * 8
    v_f16 = torch.randn(n, hw, hw, c, device="cuda", dtype=torch.float16)
    v_i8 = quantize(v_f16.float(), 0.05)
    om16 = torch.randn(n, hw, hw, p, device="cuda", dtype=torch.float16)

    def timeit(fn):
        for _ in range(20):
            fn()
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(iters):
            fn()
        torch.cuda.synchronize()
        return (time.perf_counter() - t0) / iters * 1e3

    t_fp16 = timeit(lambda: DCNv4Function.apply(
        v_f16, om16, 1, k, 1, 1, 0, (k - 1) // 2, 1, 1,
        group, gc, 1.0, 256, 0))
    t_int8 = timeit(lambda: ext.dcnv4_int8_forward(
        v_i8, om16, 1, k, 1, 1, 0, (k - 1) // 2, 1, 1, group, gc, 1.0, 1.0))
    print(f"  strip(1,{k}) [{n},{hw}²,{c}] fwd: fp16 {t_fp16:.3f} ms, "
          f"int8 {t_int8:.3f} ms ({t_fp16 / t_int8:.2f}x)")
    if t_int8 > t_fp16:
        print("  NOTE: v1 int8 kernel is correctness-first (no shm/vector "
              "tuning); optimize only after the accuracy gate passes.")


def main() -> None:
    assert torch.cuda.is_available(), "CUDA required"
    print("== correctness (int8 vs dequant->fp32->quant) ==")
    check_case(3, 3)
    check_case(1, 9)
    check_case(9, 1)
    check_case(1, 9, c=192, group=12)
    check_case(1, 9, c=192, group=12, offset_scale=0.5)
    print("== module ==")
    check_module()
    print("== benchmark (informational) ==")
    bench()
    print("\nALL INT8 GATES PASSED")


if __name__ == "__main__":
    main()