"""Runnable smoke test for the non-intrusive DS reproduction (1-channel).

Verifies, on synthetic 16 kHz audio, that:
  1. BSRNN SE forward works and exposes hidden H of shape (B, N, T, K).
  2. The decoupled DS module and DS4BSRNN produce coefficients S in [0,1]
     with shape (B, K, T) and a refined complex spectrogram.
  3. Algorithm 1 interpolation invariants hold:
       S == 1  ->  X_hat == X        (use only original)
       S == 0  ->  X_hat == X_tilde  (use only enhanced)
     and full-band == sub-band when S is constant across bands.
  4. Non-intrusive training: a single optimisation step runs, the loss is
     finite, gradients reach ONLY the DS parameters (SE + ASR stay frozen),
     and the scheduled-DS-coefficient bias is frozen during warmup then
     unfreezes afterwards.

Run:
    python scripts/smoke_test.py
"""

import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ds4se import BSRNN, DS4BSRNN, DSModule, ds_interpolate  # noqa: E402
from ds4se.stft import STFT  # noqa: E402
from ds4se.train_ds import build_pipeline, train_step  # noqa: E402


def section(msg):
    print("\n" + "=" * 70 + f"\n{msg}\n" + "=" * 70)


def main():
    torch.manual_seed(2025)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device = {device}")

    B, sr, dur = 2, 16000, 1.0
    n_fft, hop = 512, 128
    samples = int(sr * dur)
    noisy = torch.randn(B, samples, device=device) * 0.1

    stft = STFT(n_fft=n_fft, hop_length=hop, win_length=n_fft)
    X = stft.stft(noisy)  # (B, T, F, 2)
    Bn, T, F, two = X.shape
    print(f"noisy {tuple(noisy.shape)} -> STFT {tuple(X.shape)}")
    assert two == 2 and F == n_fft // 2 + 1

    # ---- 1. BSRNN SE + hidden H -------------------------------------------
    section("1. BSRNN SE forward + hidden representation H")
    se = BSRNN(n_fft=n_fft, num_channel=32, num_layer=2).to(device).eval()
    with torch.no_grad():
        enhanced, H = se(X, return_hidden=True)
    K = se.num_bands
    print(f"enhanced {tuple(enhanced.shape)}, H {tuple(H.shape)}, num_bands K={K}")
    assert enhanced.shape == X.shape
    assert H.shape == (B, 32, T, K)

    # ---- 2. DS module + DS4BSRNN coefficients -----------------------------
    section("2. DS module + DS4BSRNN coefficient shapes / range")
    ds = DSModule(n_fft=n_fft, subbands=se.subbands, mode="sub-band",
                  warmup_steps=10).to(device)
    X_hat, S = ds(X, enhanced)
    print(f"[decoupled] X_hat {tuple(X_hat.shape)} complex={torch.is_complex(X_hat)}, "
          f"S {tuple(S.shape)} in [{S.min():.3f}, {S.max():.3f}]")
    assert torch.is_complex(X_hat) and X_hat.shape == (B, T, F)
    assert S.shape == (B, K, T)
    assert S.min() >= 0.0 and S.max() <= 1.0

    ds4 = DS4BSRNN(bsrnn_num_channel=32, subbands=se.subbands, mode="sub-band",
                   warmup_steps=10).to(device)
    X_hat4, S4 = ds4(X, enhanced, H)
    print(f"[coupled]   X_hat {tuple(X_hat4.shape)}, S {tuple(S4.shape)} "
          f"in [{S4.min():.3f}, {S4.max():.3f}]")
    assert X_hat4.shape == (B, T, F) and S4.shape == (B, K, T)
    assert S4.min() >= 0.0 and S4.max() <= 1.0

    # ---- 3. Algorithm 1 interpolation invariants --------------------------
    section("3. Algorithm 1 interpolation invariants")
    Xc = torch.view_as_complex(X.contiguous())
    Ec = torch.view_as_complex(enhanced.contiguous())
    ones = torch.ones(B, K, T, device=device)
    zeros = torch.zeros(B, K, T, device=device)
    x_hat_one = ds_interpolate(X, enhanced, ones, se.subbands, mode="sub-band")
    x_hat_zero = ds_interpolate(X, enhanced, zeros, se.subbands, mode="sub-band")
    err_one = (x_hat_one - Xc).abs().max().item()
    err_zero = (x_hat_zero - Ec).abs().max().item()
    print(f"S=1 -> ||X_hat - X||_inf      = {err_one:.2e}  (expect ~0)")
    print(f"S=0 -> ||X_hat - X_tilde||_inf = {err_zero:.2e}  (expect ~0)")
    assert err_one < 1e-5 and err_zero < 1e-5

    # full-band == sub-band when S is constant across bands
    s_const = torch.rand(B, 1, T, device=device).expand(B, K, T).contiguous()
    fb = ds_interpolate(X, enhanced, s_const, se.subbands, mode="full-band")
    sb = ds_interpolate(X, enhanced, s_const, se.subbands, mode="sub-band")
    err_fb = (fb - sb).abs().max().item()
    print(f"full-band vs sub-band (const S) = {err_fb:.2e}  (expect ~0)")
    assert err_fb < 1e-5

    # ---- 4. Non-intrusive training step -----------------------------------
    for coupled in (False, True):
        tag = "DS4BSRNN (coupled)" if coupled else "DS module (decoupled)"
        section(f"4. Non-intrusive training step: {tag}")
        pipe = build_pipeline(
            coupled=coupled, n_fft=n_fft, hop_length=hop,
            se_num_channel=32, se_num_layer=2, warmup_steps=2,
        )
        pipe.se.to(device)
        pipe.ds.to(device)
        pipe.asr.to(device)

        # frozen modules must have no trainable params
        n_se = sum(p.requires_grad for p in pipe.se.parameters())
        n_asr = sum(p.requires_grad for p in pipe.asr.parameters())
        n_ds = sum(p.requires_grad for p in pipe.ds.parameters())
        print(f"trainable params: SE={n_se}, ASR={n_asr}, DS={n_ds}")
        assert n_se == 0 and n_asr == 0 and n_ds > 0

        opt = torch.optim.Adam(
            [p for p in pipe.ds.parameters() if p.requires_grad], lr=1e-2
        )

        # synthetic CTC targets
        L = 5
        targets = torch.randint(1, pipe.asr.vocab_size, (B, L), device=device)
        target_lengths = torch.full((B,), L, dtype=torch.long, device=device)

        # bias should start frozen (scheduled DS coefficients)
        assert pipe.ds.coef_head._bias_frozen, "bias must be frozen at start"
        bias_before = pipe.ds.coef_head.linear.bias.detach().clone()

        losses = []
        for step in range(4):
            noisy_b = torch.randn(B, samples, device=device) * 0.1
            loss = train_step(pipe, opt, noisy_b, targets, target_lengths, step)
            losses.append(loss)
            assert torch.isfinite(torch.tensor(loss)), "loss not finite"

        # after warmup_steps=2, bias must be trainable
        assert not pipe.ds.coef_head._bias_frozen, "bias must unfreeze after warmup"

        # SE params must be unchanged / receive no grad
        se_grads = [p.grad for p in pipe.se.parameters() if p.grad is not None]
        asr_grads = [p.grad for p in pipe.asr.parameters() if p.grad is not None]
        ds_grads = [p.grad for p in pipe.ds.parameters()
                    if p.requires_grad and p.grad is not None]
        print(f"grad tensors: SE={len(se_grads)}, ASR={len(asr_grads)}, "
              f"DS={len(ds_grads)}")
        assert len(se_grads) == 0 and len(asr_grads) == 0
        assert len(ds_grads) > 0
        # while frozen, bias did not move on step 0; overall it is now trainable
        print(f"losses: {[round(x, 4) for x in losses]}")
        print(f"bias frozen->trainable ok; bias_init={bias_before.item():.3f}")

    section("ALL SMOKE TESTS PASSED")


if __name__ == "__main__":
    main()