import json
import sys
from pathlib import Path

import numpy as np
from PIL import Image

folder = Path(sys.argv[1] if len(sys.argv) > 1 else ".")
checks = []
for angle in [0, 4, 12]:
    y, x = np.mgrid[:384, :384]
    u, v = (x + 0.5) / 384 - 0.5, (y + 0.5) / 384 - 0.5
    theta = np.deg2rad(angle)
    first = (1 + np.cos(2 * np.pi * 38 * u)) / 2
    second = (1 + np.cos(2 * np.pi * 38 * (u * np.cos(theta) + v * np.sin(theta)))) / 2
    expected = np.floor(255 * first * second + 0.5).astype(np.uint8)
    actual = np.asarray(Image.open(folder / f"moire-angle-{angle}.png").convert("RGBA"))
    assert np.array_equal(actual[:, :, 0], expected)
    assert np.array_equal(actual[:, :, 1], expected) and np.array_equal(actual[:, :, 2], expected)
    assert np.all(actual[:, :, 3] == 255)
    record = json.loads((folder / f"moire-angle-{angle}.json").read_text())
    beat = 2 * 38 * np.sin(theta / 2)
    assert abs(beat - record["beatCyclesPerWidth"]) < 1e-12
    checks.append({"name": f"overlap {angle}", "pixels": 384 * 384, "exactRGBA": True})
for name, gain in [("point", 1), ("prefiltered", np.exp(-2 * np.pi**2 * 0.75**2 * 0.58**2))]:
    values = np.floor(
        255 * (0.5 + 0.5 * gain * np.cos(2 * np.pi * 0.58 * np.arange(384))) + 0.5
    ).astype(np.uint8)
    actual = np.asarray(Image.open(folder / f"moire-sampling-{name}.png").convert("RGBA"))
    assert np.array_equal(actual[:, :, :3], np.broadcast_to(values[None, :, None], (384, 384, 3)))
    assert np.all(actual[:, :, 3] == 255)
    checks.append({"name": name, "pixels": 384 * 384, "exactRGBA": True})
root = json.loads((folder / "prompt2rig-root-review.json").read_text())
expected = [[0.2, 1, 0], [0.4, 1, 0.3], [0.8, 1, -0.2]]
assert np.allclose([r["local"] for r in root["frames"]], expected, atol=1e-12, rtol=0)
assert [r["time"] for r in root["frames"]] == [0, 0.5, 1.5]
proof = {
    "scope": "Independent saved numerical output checks not an image restoration or motion capture benchmark",
    "checks": checks,
    "rootPoints": len(expected),
}
(folder / "moire-rig-validation.json").write_text(json.dumps(proof, indent=2) + "\n")
print(json.dumps(proof))