import os as _os
HERE = _os.path.dirname(_os.path.abspath(__file__))
ROOT = _os.path.dirname(HERE)
import sys; sys.path.insert(0, HERE)
import harness, torch, math, time
pack, nodes = harness.load_pack()
from ogg_pack.ogg import codec as C
torch.manual_seed(0)
m = C.FishCodec(dtype=torch.float32, device="cpu", operations=harness.ops.manual_cast)
for p in m.parameters(): p.data.normal_(0, 0.05)
print("codec params M:", round(sum(p.numel() for p in m.parameters()) / 1e6))
m.eval()
t = torch.arange(C.SAMPLE_RATE) / C.SAMPLE_RATE
wav = (0.3 * torch.sin(2 * math.pi * 220 * t) + 0.1 * torch.randn_like(t))[None]
with torch.inference_mode():
    t0 = time.time(); codes = m.encode(wav); print("encode", tuple(wav.shape), "->", tuple(codes.shape), codes.dtype, f"{time.time()-t0:.1f}s")
    assert codes.shape[1] == 10 and codes.shape[2] == math.ceil(C.SAMPLE_RATE / C.FRAME_LENGTH), codes.shape
    assert int(codes[:, 0].max()) < 4096 and int(codes[:, 1:].max()) < 1024
    t0 = time.time(); audio = m.decode(codes, torch.float32); print("decode ->", tuple(audio.shape), f"{time.time()-t0:.1f}s")
    assert audio.shape[-1] == codes.shape[-1] * C.FRAME_LENGTH and torch.isfinite(audio).all()
print("CODEC OK: [10, frames] in, frames*2048 samples out")