# SPDX-License-Identifier: MIT
"""Check the GPU encodings in lcdata.py against direct NumPy definitions.

  python test_encodings.py
"""
import numpy as np
import torch

import lcdata

rng = np.random.default_rng(0)
L, m, q = 1024, 64, 8
x = np.cumsum(rng.normal(size=(3, L)), 1).astype(np.float32)
xt = torch.tensor(x)


def paa_np(v, m):
    return v.reshape(m, -1).mean(1)


def mm(v, lo=-1, hi=1):
    return lo + (hi - lo) * (v - v.min()) / (v.max() - v.min())


for k in range(3):
    z = mm(paa_np(x[k], m))
    phi = np.arccos(np.clip(z, -1, 1))
    ref = {
        "gasf": np.cos(phi[:, None] + phi[None, :]),
        "gadf": np.sin(phi[:, None] - phi[None, :]),
        "rp": 1 - 2 * np.abs(mm(paa_np(x[k], m), 0, 1)[:, None] - mm(paa_np(x[k], m), 0, 1)[None, :]),
    }
    for name, r in ref.items():
        got = lcdata.ENCODERS[name](xt[k:k + 1], m)[0].numpy()
        err = np.abs(got - r).max()
        assert err < 1e-4, (name, err)

    # MTF: full L x L field, then block-average to m x m, then min-max to [-1, 1]
    edges = np.quantile(x[k], np.linspace(0, 1, q + 1)[1:-1])
    b = np.searchsorted(edges, x[k])
    W = np.zeros((q, q))
    for i in range(L - 1):
        W[b[i], b[i + 1]] += 1
    W = W / np.maximum(W.sum(1, keepdims=True), 1)
    full = W[b[:, None], b[None, :]]
    blk = full.reshape(m, L // m, m, L // m).mean((1, 3))
    got = lcdata.mtf(xt[k:k + 1], m, q)[0].numpy()
    err = np.abs(got - mm(blk)).max()
    assert err < 1e-4, ("mtf", err)

img = lcdata.encode(xt, "stack", 224)
assert img.shape == (3, 3, 224, 224) and img.min() >= -1.0001 and img.max() <= 1.0001
for e in ("gasf", "gadf", "mtf", "rp", "line", "stack"):
    im = lcdata.encode(xt, e, 224)
    assert torch.isfinite(im).all() and im.shape == (3, 3, 224, 224), e
print("all encodings match their definitions")