import argparse
import os
import sys
import torch

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

SHAPES = {"solar": (1, 6, 20), "wind": (1, 12, 15)}
VERIFICATION = {"solar": 1e-06, "wind": 0.446117}
TOL = 0.001


def fixed_input(track):
    torch.manual_seed(0)
    return torch.randn(*SHAPES[track])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--track", choices=["solar", "wind"], required=True)
    ap.add_argument("--weights", default=".")
    args = ap.parse_args()
    model = load_model(args.weights, track=args.track)
    x = fixed_input(args.track)
    with torch.no_grad():
        out = model(x)
    if out.dim() != 2 or out.shape[0] != 1 or out.shape[-1] != 99:
        print(f"Fail: unexpected output shape {tuple(out.shape)}, expected (1, 99)")
        sys.exit(1)
    p50 = out[0, 49].item()
    print(f"Captured first-hour P50 for {args.track}: {p50:.6f}")
    ref = VERIFICATION[args.track]
    if abs(p50 - ref) > TOL:
        print(
            f"Fail: P50 {p50:.6f} differs from verification {ref:.6f} by more than {TOL}"
        )
        sys.exit(1)
    print(f"Verification check OK (P50 within {TOL} of {ref:.6f})")
    sys.exit(0)


if __name__ == "__main__":
    main()