import json
import numpy as np
import torch
from laplacian_attention import FirstSpikeLaplacianAttention


def main():
    torch.manual_seed(7)
    dtype = torch.float64
    b, h, nq, nk, c, d = 2, 3, 5, 7, 4, 6
    q = torch.rand(b, h, nq, c, dtype=dtype, requires_grad=True)
    k = torch.rand(b, h, nk, c, dtype=dtype, requires_grad=True)
    v = torch.rand(b, h, nk, d, dtype=dtype, requires_grad=True)
    mod = FirstSpikeLaplacianAttention(h, init_sigma=.7).to(dtype)
    out, a, _ = mod(q, k, v, return_attention=True)
    out.square().mean().backward()

    # 13 independent valid cases: near distance 0, far distance delta >= 0.
    # The expected probability is the exact two-key normalization after L1.
    delta = torch.linspace(0, 3, 13, dtype=dtype)
    q2 = torch.zeros(13, 1, 1, 1, dtype=dtype)
    k2 = torch.stack([torch.zeros_like(delta), delta], dim=-1).reshape(13, 1, 2, 1)
    v2 = torch.eye(2, dtype=dtype).reshape(1, 1, 2, 2).expand(13, -1, -1, -1)
    sigma = .7
    mod2 = FirstSpikeLaplacianAttention(1, init_sigma=sigma).to(dtype)
    _, a2, _ = mod2(q2, k2, v2, return_attention=True)
    observed_prob = a2[:, 0, 0, 0].detach().numpy()
    expected_prob = (1 / (1 + torch.exp(-delta / sigma))).numpy()

    q3 = torch.zeros(1, 1, 1, 1, dtype=dtype)
    k3 = torch.tensor([[[[.8], [2.4]]]], dtype=dtype)
    _, a3, _ = mod2(q3, k3, torch.ones(1,1,2,1,dtype=dtype), return_attention=True)
    observed_log_ratio = float(torch.log(a3[0,0,0,0] / a3[0,0,0,1]))
    predicted_log_ratio = (2.4-.8)/sigma
    result = {
      'shape': list(out.shape),
      'row_sum_error': float((a.sum(-1)-1).abs().max()),
      'nonnegative': bool((a >= 0).all()),
      'positive_bandwidths': mod.bandwidth().detach().tolist(),
      'gradient_finite': bool(all(x.grad is not None and torch.isfinite(x.grad).all() for x in [q,k,v,mod.theta])),
      'two_key_max_error': float(np.max(np.abs(observed_prob-expected_prob))),
      'two_key_crossing_delta': float(delta[np.argmin(np.abs(observed_prob-.5))]),
      'log_ratio_observed': observed_log_ratio,
      'log_ratio_predicted': predicted_log_ratio,
      'log_ratio_error': abs(observed_log_ratio-predicted_log_ratio),
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()