import torch
import torch.nn as nn
import torchvision
import numpy as np
import argparse
import torch.backends.cudnn as cudnn
import os

# distributed training
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DataParallel
from torch.nn.parallel import DistributedDataParallel as DDP


from utils import yaml_config_hook
from model import load_optimizer, save_model
from modules.transformations import DataTransforms

from timm.models import create_model
from transforms_factory import contrastive_learning_transforms
from loss import BarlowTwinsLoss

import model
import barlowtwins


config_parser = parser = argparse.ArgumentParser(description='Training Config', add_help=False)
parser.add_argument('-c', '--config', default='imagenet.yml', type=str, metavar='FILE',
                    help='YAML config file specifying default arguments')

def train(args, train_loader, model, criterion, optimizer, writer):
    loss_epoch = 0
    for step, ((x_i, x_j), _) in enumerate(train_loader):
        optimizer.zero_grad()
        x_i = x_i.cuda(non_blocking=True)
        x_j = x_j.cuda(non_blocking=True)

        # positive pair, with encoding
        if args.spiking:
            h_i, h_j, z_i, z_j, z_i_temporal, z_j_temporal = model(x_i, x_j)
        else:
            h_i, h_j, z_i, z_j = model(x_i, x_j)

        if args.temporal_loss and args.spiking:
            loss = criterion(z_i_temporal, z_j_temporal)
        else:
            loss = criterion(z_i, z_j)

        loss.backward()
        optimizer.step()

        # use all_reduce to store loss correctly
        if dist.is_available() and dist.is_initialized():
            loss = loss.data.clone()
            dist.all_reduce(loss.div_(dist.get_world_size()))

        if args.rank == 0 and step % 100 == 0:
            print(f"Step [{step}/{len(train_loader)}]\t Loss: {loss.item()}")

        if args.rank == 0:
            writer.add_scalar("Loss/train_epoch", loss.item(), args.global_step)
            args.global_step += 1

        loss_epoch += loss.item()
    return loss_epoch

def main(gpu, args):
    args.rank = args.nr * args.gpus + gpu
    args.device = torch.device(f"cuda:{gpu}" if torch.cuda.is_available() else "cpu")

    if args.world_size > 1:
        dist.init_process_group("nccl", rank=args.rank, world_size=args.world_size)
        torch.cuda.set_device(gpu)

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    if args.dataset == "CIFAR10":
        train_dataset = torchvision.datasets.CIFAR10(
            args.dataset_dir,
            train=True,
            download=True,
            transform=DataTransforms(size=32),
        )
    elif args.dataset == "CIFAR100":
        train_dataset = torchvision.datasets.CIFAR100(
            args.dataset_dir,
            train=True,
            download=True,
            transform=DataTransforms(size=32),
        )
    elif args.dataset == "ImageNet":
        train_dataset = torchvision.datasets.ImageFolder(
            os.path.join(args.dataset_dir, 'train'),
            transform=DataTransforms(size=args.image_size),
        )
    else:
        raise NotImplementedError

    if args.world_size > 1:
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_dataset, num_replicas=args.world_size, rank=args.rank , shuffle=True
        )
    else:
        train_sampler = None

    dataloader_train_ssl = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=(train_sampler is None),
        drop_last=True,
        num_workers=args.num_workers,
        sampler=train_sampler
    )

    # initialize ResNet, encoder is resnet/resnet_snn
    if args.spiking:
        if "resnet" in args.model:
            backbone = get_resnet_spiking(args.model, args.timestep, args.sync_norm, args.n_classes)
            n_features = backbone.fc.in_features  # get dimensions of fc layer
            backbone.fc = nn.Identity()
            # convert to cifar-fitted structure
            if "CIFAR" in args.dataset:
                backbone = modify_resnet_model(backbone)

        elif "vgg" in args.model:
            backbone = get_vgg_spiking(args.model, args.timestep, args.sync_norm, args.n_classes)
            n_features = backbone.fc[0].in_features
            backbone.fc = nn.Identity()
        model = BartonTwinsSpiking(backbone, in_dim=n_features, out_dim=args.projection_dim, timestep=args.timestep)
    else:
        if "resnet" in args.model:
            backbone = get_resnet(args.model, args.n_classes)
            n_features = backbone.fc.in_features
            backbone.fc = nn.Identity()
            # convert to cifar-fitted structure
            if "CIFAR" in args.dataset:
                backbone = modify_resnet_model(backbone)

        elif "vgg" in args.model:
            backbone = get_vgg(args.model, args.n_classes)
            n_features = backbone.fc[0].in_features
            backbone.fc = nn.Identity()
        model = BartonTwins(backbone, in_dim=n_features, out_dim=args.projection_dim)

    if args.reload:
        model_fp = os.path.join(
            args.model_path, "checkpoint_{}.tar".format(args.epoch_num)
        )
        model.load_state_dict(torch.load(model_fp, map_location=args.device.type))
    model = model.to(args.device)

    # optimizer / loss
    optimizer, scheduler = load_optimizer(args, model)
    if args.temporal_loss and args.spiking:
        criterion = BarlowTwinsTemporalLoss(args.device, args.world_size, args.timestep, args.cross_temporal, args.simplified_loss)
    else:
        criterion = BarlowTwinsLoss(device=args.device, world_size=args.world_size)

    # DDP
    if args.world_size > 1:
        if not args.spiking:
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
        else:
            model = SynctdBatchNorm.convert_sync_tdBatchNorm(model)
        model = DDP(model, device_ids=[gpu])

    model = model.to(args.device)
    writer = None
    if args.rank == 0:
        writer = SummaryWriter()

    args.global_step = 0
    args.current_epoch = args.start_epoch
    for epoch in range(args.start_epoch, args.epochs):
        if args.rank == 0:
            print('===================================')

        if train_sampler is not None:
            train_sampler.set_epoch(epoch)

        lr = optimizer.param_groups[0]["lr"]
        loss_epoch = train(args, dataloader_train_ssl, model, criterion, optimizer, writer)

        if scheduler:
            scheduler.step()

        if args.rank == 0 and epoch % 10 == 0:
            save_model(args, model, optimizer)

        if args.rank == 0:
            writer.add_scalar("Loss/train", loss_epoch / len(dataloader_train_ssl), epoch)
            writer.add_scalar("Misc/learning_rate", lr, epoch)
            print(
                f"Epoch [{epoch}/{args.epochs}]\t Loss: {loss_epoch / len(dataloader_train_ssl)}\t lr: {round(lr, 7)}"
            )
            args.current_epoch += 1

    ## end training
    if args.rank == 0:
        save_model(args, model, optimizer)


def _parse_args():
    # Do we have a config file to parse?
    args_config, remaining = config_parser.parse_known_args()
    if args_config.config:
        with open(args_config.config, 'r') as f:
            cfg = yaml.safe_load(f)
            parser.set_defaults(**cfg)

    # The main arg parser parses the rest of the args, the usual
    # defaults will have been overridden if config file specified.
    args = parser.parse_args(remaining)

    # Cache the args as a text string to save them in the output dir later
    args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)
    return args, args_text



if __name__ == "__main__":
    args, args_text = _parse_args()
    parser = argparse.ArgumentParser()
    config = yaml_config_hook("./config/config.yaml")
    for k, v in config.items():
        parser.add_argument(f"--{k}", default=v, type=type(v))

    args = parser.parse_args()

    # Master address for distributed data parallel
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "8090"

    if not os.path.exists(args.model_path):
        os.makedirs(args.model_path)

    args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    args.world_size = args.gpus * args.nodes
    args.lr = float(args.lr)
    print(vars(args))

    cudnn.benchmark = True

    if args.world_size > 1:
        print(
            f"Training with {args.world_size} , waiting until all world_size join before starting training"
        )
        mp.spawn(main, args=(args,), nprocs=args.gpus, join=True)
    else:
        main(0, args)
