import argparse
import os
import sys
import torch

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--track", choices=["solar", "wind"], default="solar")
    ap.add_argument(
        "--weights",
        default=None,
        help="path to the track's safetensors (defaults to <track>.safetensors)",
    )
    args = ap.parse_args()
    weights = args.weights or f"{args.track}.safetensors"
    steps, feats = modeling.expected_shape(args.track)
    raw_window = torch.randn(steps, feats)
    nwp = modeling.preprocess_nwp(raw_window, track=args.track)
    model = modeling.load_model(weights, track=args.track)
    with torch.no_grad():
        quantiles = model(nwp)
    p10 = quantiles[0, 9].item()
    p50 = quantiles[0, 49].item()
    p90 = quantiles[0, 89].item()
    print(f"Track: {args.track}")
    print(
        f"Output shape: {tuple(quantiles.shape)}  (batch, 99 quantiles for the center hour)"
    )
    print("Center forecast hour, power as fraction of site capacity:")
    print(f"  P10 (low):    {p10:.4f}")
    print(f"  P50 (median): {p50:.4f}")
    print(f"  P90 (high):   {p90:.4f}")
    print(f"  P10..P90 uncertainty band width: {p90 - p10:.4f}")


if __name__ == "__main__":
    main()