import argparse
import os
import sys
import torch

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--variant",
        default="cmapss_FD001",
        choices=[
            "cmapss_FD001",
            "cmapss_FD002",
            "cmapss_FD003",
            "cmapss_FD004",
            "cwru",
        ],
    )
    ap.add_argument(
        "--weights",
        default=os.environ.get("FELA_PDM_WEIGHTS", "."),
        help="directory with <variant>.safetensors + config.json, or a .pt checkpoint path",
    )
    args = ap.parse_args()
    src = args.weights
    variant_file = os.path.join(src, f"{args.variant}.safetensors")
    if os.path.isdir(src) and os.path.isfile(variant_file):
        model = load_model(src, variant=args.variant)
    elif os.path.isfile(src):
        model = load_model(src, variant=args.variant)
    else:
        raise SystemExit(
            f"No weights at {src}. Set FELA_PDM_WEIGHTS to a directory holding {args.variant}.safetensors and config.json (or pass a .pt checkpoint path). Weights are in lowdown-labs/FELA-pdm."
        )
    if args.variant.startswith("cmapss"):
        window = torch.randn(1, 30, 14)
        validate_window(window, model.cfg)
        rul = model.predict(window, task="rul")
        print(f"Variant: {args.variant}")
        print(f"Input shape: {tuple(window.shape)} (30 cycles, 14 sensors)")
        print(f"Estimated remaining useful life: {rul:.1f} cycles (capped at 125)")
    else:
        window = torch.randn(1, 2048, 1)
        validate_window(window, model.cfg)
        idx, prob = model.predict(window, task="cls")
        print(f"Variant: {args.variant}")
        print(f"Input shape: {tuple(window.shape)} (2048 vibration samples, 1 channel)")
        print(f"Predicted fault class index: {idx}  (probability {prob:.4f})")


if __name__ == "__main__":
    main()