import os
from fleet_beam_search_3.dataloader import VRP_Dataset


def build_and_save_dataset(save_path: str, dataset_size: int = 1000):
    """
    Builds a VRP dataset and saves it to the given path.

    Args:
        save_path (str): Destination file path to save dataset.
        dataset_size (int): Number of samples in the dataset.
    """
    dataset = VRP_Dataset(dataset_size=dataset_size)

    # Check if dataset has a save method or export function
    if hasattr(dataset, "save") and callable(dataset.save):
        dataset.save(save_path)
        print(f"Dataset saved to: {save_path}")
    else:
        raise NotImplementedError("Dataset object lacks a save() method.")

        
if __name__ == "__main__":
    output_path = "data/generated_vrp_dataset.pkl"
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    build_and_save_dataset(save_path=output_path, dataset_size=1000)