import datetime
import random
from argparse import ArgumentParser
from pathlib import Path

import numpy as np
import torch
import yaml
from tqdm import tqdm

from datasets.video_dataset import AtlasDataset
from models.video_model import VideoModel
from util.atlas_loss import AtlasLoss
from util.util import get_optimizer
from util.video_logger import DataLogger


def train_model(config):
    # set seed
    seed = config["seed"]
    if seed == -1:
        seed = np.random.randint(2 ** 32)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    print(f"running with seed: {seed}.")

    dataset = AtlasDataset(config)
    model = VideoModel(config)
    criterion = AtlasLoss(config)
    optimizer = get_optimizer(config, model.parameters())

    logger = DataLogger(config, dataset)
    with tqdm(range(1, config["n_epochs"] + 1)) as tepoch:
        for epoch in tepoch:
            inputs = dataset[0]
            optimizer.zero_grad()
            outputs = model(inputs)
            losses = criterion(outputs, inputs)

            loss = 0.
            if config["finetune_foreground"]:
                loss += losses["foreground"]["loss"]
            elif config["finetune_background"]:
                loss += losses["background"]["loss"]

            lr = optimizer.param_groups[0]["lr"]
            log_data = logger.log_data(epoch, lr, losses, model, dataset)

            loss.backward()
            optimizer.step()
            optimizer.param_groups[0]["lr"] = max(config["min_lr"], config["gamma"] * optimizer.param_groups[0]["lr"])

            if config["use_wandb"]:
                wandb.log(log_data)
            else:
                if epoch % config["log_images_freq"] == 0:
                    logger.save_locally(log_data)

            tepoch.set_description(f"Epoch {epoch}")
            tepoch.set_postfix(loss=loss.item())


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--config",
        default="./configs/video_config.yaml",
        help="Config path",
    )
    parser.add_argument(
        "--example_config",
        default="car-turn_winter.yaml",
        help="Example config name",
    )
    args = parser.parse_args()
    config_path = args.config

    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    with open(f"./configs/video_example_configs/{args.example_config}", "r") as f:
        example_config = yaml.safe_load(f)
    config["example_config"] = args.example_config
    config.update(example_config)

    run_name = f"-{config['checkpoint_path'].split('/')[-2]}"
    if config["use_wandb"]:
        import wandb

        wandb.init(project=config["wandb_project"], entity=config["wandb_entity"], config=config, name=run_name)
        wandb.run.name = str(wandb.run.id) + wandb.run.name
        config = dict(wandb.config)
    else:
        now = datetime.datetime.now()
        run_name = f"{now.strftime('%Y-%m-%d_%H-%M-%S')}{run_name}"
        path = Path(f"{config['results_folder']}/{run_name}")
        path.mkdir(parents=True, exist_ok=True)
        with open(path / "config.yaml", "w") as f:
            yaml.dump(config, f)
        config["results_folder"] = str(path)

    train_model(config)
    if config["use_wandb"]:
        wandb.finish()