"""Claim 1 audit: TS-Fingerprint compresses a variable-length MedTS into a fixed set of
k learnable Fingerprint Tokens via a learnable query set Q in R^{k x d}, with k << T.

Checks:
  1. the encoder output is exactly (B, k, d) for every input length T;
  2. Q is a genuine nn.Parameter that receives gradient (i.e. it is *learnable*),
     and the parameter count of the bottleneck is independent of T;
  3. k << T holds at the paper's settings for all 7 benchmark datasets of Table 5;
  4. the decoder sees *only* F' plus positional mask tokens, so the data-processing
     chain X -> F' -> X_hat is strict: perturbing a masked input patch can change
     X_hat only through F'.  We verify this by freezing F' and re-running the decoder.
"""

import json
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import TSFingerprint  # noqa: E402

# Table 5 of the paper: (#Timestamps T, #Channel C, patch size)
DATASETS = {
    "ADFTD": (256, 19, 8), "PTB": (300, 15, 10), "PTB-XL": (250, 12, 10),
    "APAVA": (256, 16, 8), "FLAAP": (100, 6, 10), "UCIHAR": (128, 9, 32),
    "Sleep-EDF": (3000, 1, 60),
}
K, D = 8, 128
rep = {}

torch.manual_seed(0)
m = TSFingerprint(c_in=19, patch_size=8, n_classes=3, d_model=D, k=K, max_patches=4096)

# --- 1. shape invariance to T -------------------------------------------------
shapes = {}
for T in (64, 128, 256, 512, 1024, 2048, 4096, 8192):
    with torch.no_grad():
        f = m.encoder(torch.randn(3, T, 19))
    shapes[T] = {"n_patches": T // 8, "F_shape": list(f.shape)}
rep["shape_invariance"] = shapes
rep["all_outputs_are_k_by_d"] = all(v["F_shape"] == [3, K, D] for v in shapes.values())

# --- 2. Q is learnable, and bottleneck size is T-independent -------------------
rep["Q_shape"] = list(m.encoder.Q.shape)
rep["Q_is_parameter"] = isinstance(m.encoder.Q, torch.nn.Parameter)
rep["Q_requires_grad"] = bool(m.encoder.Q.requires_grad)
m.zero_grad()
loss, _, _ = m.pretrain_step(torch.randn(4, 256, 19))
loss.backward()
g = m.encoder.Q.grad
rep["Q_grad_norm_after_backward"] = float(g.norm())
rep["Q_receives_gradient"] = bool(g is not None and g.norm() > 0)
rep["encoder_params_excluding_pos_embed"] = sum(
    p.numel() for n, p in m.encoder.named_parameters() if n != "pos")

# --- 3. k << T on every benchmark ---------------------------------------------
ratios = {}
for name, (T, C, ps) in DATASETS.items():
    n_patch = T // ps
    ratios[name] = {"T": T, "C": C, "patch": ps, "n_patches": n_patch, "k": K,
                    "k_over_T": round(K / T, 5), "k_over_n_patches": round(K / n_patch, 4),
                    "k_much_less_than_T": bool(K < T / 4)}
rep["compression_per_dataset"] = ratios
rep["k_ll_T_on_all_datasets"] = all(v["k_much_less_than_T"] for v in ratios.values())
# compression factor: raw floats in vs floats out
rep["compression_factor_TxC_to_kxd"] = {
    n: round((v["T"] * v["C"]) / (K * D), 3) for n, v in ratios.items()}

# --- 4. strict bottleneck: decoder output depends on x only through F' ---------
m.eval()
x = torch.randn(2, 256, 19)
with torch.no_grad():
    f = m.encoder(x)
    rec_a = m.decoder(f, 32)
    x2 = x.clone()
    x2[:, 100:140] += 5.0                       # large perturbation of the input
    f2 = m.encoder(x2)
    rec_b = m.decoder(f2, 32)                   # differs, because F' differs
    rec_c = m.decoder(f, 32)                    # same F' -> must be bit-identical
rep["decoder_conditioned_solely_on_F"] = {
    "delta_recon_when_F_changes": float((rec_a - rec_b).abs().max()),
    "delta_recon_when_F_held_fixed": float((rec_a - rec_c).abs().max()),
    "strict_chain_X_to_F_to_Xhat": bool((rec_a - rec_c).abs().max() == 0.0
                                        and (rec_a - rec_b).abs().max() > 0),
}
# the decoder signature takes only (F', n_patches) -- no path for raw patches
import inspect  # noqa: E402
rep["decoder_forward_signature"] = str(inspect.signature(m.decoder.forward))

print(json.dumps(rep, indent=2))
os.makedirs("results", exist_ok=True)
with open("results/claim1_audit.json", "w") as fh:
    json.dump(rep, fh, indent=2)

print("\n---- verdict ----")
print(f"encoder output is (B, {K}, {D}) for every T in "
      f"{sorted(shapes)}: {rep['all_outputs_are_k_by_d']}")
print(f"Q in R^(k x d) = {rep['Q_shape']}, learnable nn.Parameter, "
      f"grad norm {rep['Q_grad_norm_after_backward']:.4f}")
print(f"k << T on all 7 Table-5 datasets: {rep['k_ll_T_on_all_datasets']} "
      f"(k/T from {min(v['k_over_T'] for v in ratios.values()):.5f} to "
      f"{max(v['k_over_T'] for v in ratios.values()):.5f})")
print("decoder is strictly conditioned on F': "
      f"{rep['decoder_conditioned_solely_on_F']['strict_chain_X_to_F_to_Xhat']}")