import os
import pytorch_lightning as pl
from argparse import ArgumentParser
from rich.console import Console
from torch_geometric.loader import DataLoader

from dev.datasets.scalable_dataset import MultiDataset, WaymoTargetBuilder
from dev.utils.func import load_config_act, Logging
from dev.model.smart import SMART

CONSOLE = Console(width=120)


if __name__ == '__main__':
    pl.seed_everything(2024, workers=True)
    parser = ArgumentParser()
    parser.add_argument('--seed', type=int, default=2024)
    parser.add_argument('--config', type=str, default='configs/train/train_scalable_with_state.yaml')
    parser.add_argument('--ckpt_path', type=str, default="")
    parser.add_argument('--insert_agent', action='store_true')
    parser.add_argument('--t', type=str, default=2)
    parser.add_argument('--save_path', type=str, default=None)
    args = parser.parse_args()
    pl.seed_everything(args.seed, workers=True)
    config = load_config_act(args.config)
    logger = Logging().log(level='DEBUG')

    data_config = config.Dataset
    val_dataset = MultiDataset(split='val',
                               raw_dir=data_config.val_raw_dir,
                               token_size=data_config.token_size,
                               transform=WaymoTargetBuilder(
                                   config.Model.num_historical_steps,
                                   config.Model.decoder.num_future_steps,
                                   max_num=data_config.max_num,
                                   training=False),
                               tfrecord_dir=data_config.val_tfrecords_splitted,
                               predict_motion=config.Model.predict_motion,
                               predict_state=config.Model.predict_state,
                               predict_map=config.Model.predict_map,
                               buffer_size=config.Model.buffer_size,
                               logger=logger,
                )
    dataloader = DataLoader(val_dataset,
                            shuffle=False,
                            num_workers=data_config.num_workers,
                            pin_memory=data_config.pin_memory,
                            persistent_workers=True if data_config.num_workers > 0 else False
                )

    if args.save_path is not None:
        save_path = args.save_path
    else:
        assert args.ckpt_path != "" and os.path.exists(args.ckpt_path), f"Path {args.ckpt_path} not exist!"
        save_path = os.path.join(os.path.dirname(args.ckpt_path), 'val')
    CONSOLE.log(f"Results will be saved to [yellow]{save_path}[/]")
    os.makedirs(save_path, exist_ok=True)

    model = SMART(config.Model, save_path=save_path, logger=logger, insert_agent=args.insert_agent, t=args.t)
    CONSOLE.log(f"Loaded model from [yellow]{args.ckpt_path}[/]")

    trainer_config = config.Trainer
    trainer = pl.Trainer(accelerator=trainer_config.accelerator,
                         devices=trainer_config.devices,
                         strategy='ddp', num_sanity_val_steps=0)
    trainer.validate(model, dataloader, ckpt_path=args.ckpt_path)

    CONSOLE.log(f"Validation done!")