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 modules import get_resnet, get_resnet_spiking, get_vgg, get_vgg_spiking, SynctdBatchNorm, modify_resnet_model
from loss import BarlowTwinsLoss, BarlowTwinsTemporalLoss

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

import modules.vit
from timm.models import create_model
from spikingjelly.clock_driven import functional

# TensorBoard
from torch.utils.tensorboard import SummaryWriter


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()

        # reset base.MemoryModule
        functional.reset_net(model)

        # 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),
        )
    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 vit
    backbone = create_model(
        'spikformer',
        pretrained=False,
        drop_rate=0.,
        drop_path_rate=0.,
        drop_block_rate=None,
        img_size_h=32, img_size_w=32,
        patch_size=128, embed_dims=384, num_heads=12, mlp_ratios=4,
        in_channels=3, num_classes=10, qkv_bias=False,
        depths=4, sr_ratios=1,
        T=args.timestep
    )
    n_features = 384  # get dimensions of fc layer
    backbone.head = nn.Identity()
    model = BartonTwinsSpiking(backbone, in_dim=n_features, out_dim=args.projection_dim, timestep=args.timestep)

    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)


if __name__ == "__main__":
    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)
