import argparse

import os
import sys
import torch
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from safetensors.torch import load_file

from litespecformer import (
    LiteSpecFormerModel,
    LiteSpecFormerConfig,
    LiteSpecFormerPipeline,
)

parser = argparse.ArgumentParser(description="Zero-Shot Prediction with LiteSpecFormer")
parser.add_argument(
    "--dataset_name_or_path", type=str, help="The name or path of the dataset to use."
)
parser.add_argument(
    "--seq_length", type=int, default=336, help="The length of the input context."
)
parser.add_argument(
    "--prediction_length",
    type=int,
    default=96,
    help="The length of the prediction horizon.",
)
parser.add_argument(
    "--checkpoint_path", type=str, help="The path to the model checkpoint."
)
parser.add_argument(
    "--batch_size", type=int, default=256, help="The batch size for prediction."
)
parser.add_argument(
    "--output_dir", type=str, default=None, help="The directory to save the results."
)

args = parser.parse_args()

params = load_file(args.checkpoint_path)
model = LiteSpecFormerModel(LiteSpecFormerConfig())
model.load_state_dict(params, strict=False)

# Create the pipeline for zero-shot prediction
pipeline = LiteSpecFormerPipeline(model=model)

# Evaluate the model on the spectrum library
pipeline.evaluate_spectrum_library(
    dataset_name_or_path=args.dataset_name_or_path,
    seq_length=args.seq_length,
    prediction_length=args.prediction_length,
    batch_size=args.batch_size,
)
