import argparse
import os
import sys
import torch

sys.path.insert(0, os.path.dirname(__file__))
from modeling import load_model

SHAPE = (1, 8, 96, 96)
VERIFICATION = 0.566771
TOL = 0.001


def fixed_input():
    torch.manual_seed(0)
    return torch.randn(*SHAPE)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default=".")
    args = ap.parse_args()
    model = load_model(args.weights)
    x = fixed_input()
    with torch.no_grad():
        out = model(x)
    if tuple(out.shape) != (1, 1, 96, 96):
        print(
            f"Fail: unexpected output shape {tuple(out.shape)}, expected (1, 1, 96, 96)"
        )
        sys.exit(1)
    center = out[0, 0, 48, 48].item()
    print(f"Captured center value: {center:.6f}")
    if abs(center - VERIFICATION) > TOL:
        print(
            f"Fail: center {center:.6f} differs from verification {VERIFICATION:.6f} by more than {TOL}"
        )
        sys.exit(1)
    print(f"Verification check OK (center within {TOL} of {VERIFICATION:.6f})")
    sys.exit(0)


if __name__ == "__main__":
    main()