import argparse
import os
import sys
import torch

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

SAMPLES = {
    "cmapss_FD001": torch.full((1, 30, 14), 0.5),
    "cmapss_FD002": torch.full((1, 30, 14), 0.5),
    "cmapss_FD003": torch.full((1, 30, 14), 0.5),
    "cmapss_FD004": torch.full((1, 30, 14), 0.5),
    "cwru": torch.linspace(-1, 1, 2048).reshape(1, 2048, 1),
}
VERIFICATION = {
    "cmapss_FD001": {"value": 0.390998, "tol": 0.001},
    "cmapss_FD002": {"value": 0.038257, "tol": 0.001},
    "cmapss_FD003": {"value": 0.788864, "tol": 0.001},
    "cmapss_FD004": {"value": 0.414077, "tol": 0.001},
    "cwru": {"value": 0, "tol": 0.001},
}
EXPECTED_SHAPE = {
    "cmapss_FD001": (1,),
    "cmapss_FD002": (1,),
    "cmapss_FD003": (1,),
    "cmapss_FD004": (1,),
    "cwru": (1, 10),
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--variant", default="cmapss_FD001", choices=list(SAMPLES))
    ap.add_argument("--weights", default=os.environ.get("FELA_PDM_WEIGHTS", "."))
    args = ap.parse_args()
    model = load_model(args.weights, variant=args.variant)
    x = SAMPLES[args.variant]
    with torch.no_grad():
        out = model(x)
    exp = EXPECTED_SHAPE[args.variant]
    if tuple(out.shape) != exp:
        print(f"Fail: output shape {tuple(out.shape)} != expected {exp}")
        sys.exit(1)
    print(f"Shape OK: {tuple(out.shape)}")
    g = VERIFICATION[args.variant]
    if g["value"] is None:
        if args.variant.startswith("cmapss"):
            captured = float(out.reshape(-1)[0])
        else:
            captured = int(out.argmax(-1).item())
        print(f"Captured output: {captured}")
        print(
            "Verification value is a placeholder. Paste this captured value into VERIFICATION and re-run to enable the check. Shape check passed."
        )
        return
    if args.variant.startswith("cmapss"):
        got = float(out.reshape(-1)[0])
        if abs(got - g["value"]) > g["tol"]:
            print(f"Fail: RUL {got} vs verification {g['value']} (tol {g['tol']})")
            sys.exit(1)
    else:
        got = int(out.argmax(-1).item())
        if got != g["value"]:
            print(f"Fail: class {got} vs verification {g['value']}")
            sys.exit(1)
    print("Verification check OK")


if __name__ == "__main__":
    main()