import hashlib
import importlib.util
import json
import sys
from pathlib import Path

import numpy as np
import torch

source = Path(sys.argv[1])
out = Path(sys.argv[2])
audit = json.loads((out / "ptd-source-audit.json").read_text())
expected = next(d["sha256"] for d in audit["documents"] if d["name"] == "minimax_h3_ptd.py")
assert hashlib.sha256(source.read_bytes()).hexdigest() == expected
spec = importlib.util.spec_from_file_location("pinned_ptd", source)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
cases = []
for hops in [4, 8]:
    record = json.loads((out / f"ptd-{hops}-schedule.json").read_text())
    for shift, key in [(12, "video"), (3, "audio")]:
        grid = module.pdd_time_grid(shift, 32)
        sigma = np.linspace(1, 0, 33)
        independent = 1 - shift * sigma / (1 + (shift - 1) * sigma)
        assert np.allclose(grid.numpy(), independent, rtol=0, atol=1e-12)
        assert np.allclose(grid.numpy(), record[key + "Grid"], rtol=0, atol=1e-12)
        for hop in range(hops):
            start = hop * (32 // hops)
            block = 32 // hops
            plan = module.pdd_sampling_plan(grid.diff(), start, block)[0].numpy()
            assert np.allclose(plan, record["rows"][hop][key + "Plan"], rtol=0, atol=1e-12)
            assert abs(plan.sum() - 1) < 1e-12
        cases.append(
            {"hops": hops, "shift": shift, "intervalPlans": hops, "gridAndPlanMatched": True}
        )
toys = []
for hops in [4, 8]:
    base = torch.nn.Linear(3, 2, bias=True, dtype=torch.float64)
    with torch.no_grad():
        base.weight.copy_(torch.tensor([[1.0, 0.0, 0.5], [0.0, 1.0, -0.5]], dtype=torch.float64))
        base.bias.copy_(torch.tensor([0.1, -0.2], dtype=torch.float64))
    layer = module.PerHopLoRALinear(base, hops, rank=1, alpha=1)
    x = torch.tensor([[2.0, 3.0, 1.0]], dtype=torch.float64)
    assert torch.equal(layer(x), base(x))
    with torch.no_grad():
        for hop in range(hops):
            layer.lora_down[hop].copy_(torch.tensor([[1.0, 0.0, 0.0]], dtype=torch.float64))
            layer.lora_up[hop].copy_(
                torch.tensor([[hop + 1.0], [-(hop + 1.0)]], dtype=torch.float64)
            )
    for hop in range(hops):
        layer.set_hop(hop)
        actual = layer(x).detach().numpy()
        independent = np.array([[2.6, 2.3]]) + np.array([[2 * (hop + 1), -2 * (hop + 1)]])
        assert np.allclose(actual, independent, rtol=0, atol=1e-12)
        toys.append(
            {"hops": hops, "activeHop": hop, "output": actual.tolist(), "independentMatched": True}
        )
    layer.set_hop(-1)
    assert layer.hop == 0
    layer.set_hop(999)
    assert layer.hop == hops - 1
proof = {
    "scope": "Executed pinned source schedule and per hop linear helpers on CPU with original tiny tensors No full model weights apply_ptd_lora pipeline ComfyUI or GPU generation executed",
    "revision": audit["revision"],
    "sourceSHA256": expected,
    "torch": torch.__version__,
    "numpy": np.__version__,
    "scheduleCases": cases,
    "toyLinearCases": toys,
    "unloadedAdapterNoOp": True,
    "hopClampingChecked": True,
    "tolerance": 1e-12,
}
(out / "ptd-source-validation.json").write_text(json.dumps(proof, indent=2) + "\n")
print("Matched", len(cases), "schedule cases and", len(toys), "tiny per hop linear cases")