import numpy as np
import argparse
import time
import yaml
import os
import datetime

import torch
import torch.nn as nn
import torchvision.utils
from loader import create_loader
from timm.models import create_model, safe_model_name, resume_checkpoint, load_checkpoint
from timm.utils import *

from transforms_factory import contrastive_learning_transforms as DataTransforms
import model
import barlowtwins

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler


def cleanup():
    dist.destroy_process_group()

def reduce_tensor(tensor, world_size):
    if world_size > 1:
        rt = tensor.clone()
        dist.all_reduce(rt, op=dist.ReduceOp.SUM)
        rt /= world_size
        return rt
    else:
        return tensor


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

parser = argparse.ArgumentParser(description='PyTorch ')
parser.add_argument('--num-classes', type=int, default=1000, metavar='N',
                    help='number of label classes (Model default if None)')
parser.add_argument('--lr', default=1e-3, type=float,
                    help='learning rate')
parser.add_argument('--dataset', '-d', metavar='NAME', default='dataset',
                    help='dataset type (default: ImageFolder/ImageTar if empty)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='Resume full model and optimizer state from checkpoint (default: none)')
parser.add_argument('--eval_proportion', default='0.01', type=float,
                    help='proportion of test dataset for evaluation')
parser.add_argument('--val_epochs', default='60', type=int,
                    help='For 1% labeled data we fine-tune for 60 epochs, and for 10% labeled data we fine-tune for 30 \
                         epochs')
parser.add_argument('--local-rank', default=0, type=int, help='Local rank for DistributedDataParallel')


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


def train(args, loader, model, criterion, optimizer, world_size):
    loss_epoch = 0
    accuracy_epoch_top1 = 0
    accuracy_epoch_top5 = 0
    for step, (x, y) in enumerate(loader):
        optimizer.zero_grad()

        x = x.to(args.device)
        y = y.to(args.device)

        output, _ = model(x)
        loss = criterion(output, y)

        # top-1
        predicted = output.argmax(1)
        acc = (predicted == y).sum().item() / y.size(0)
        accuracy_epoch_top1 += acc
        # top-5
        _, pred = output.topk(5, 1, True, True)
        pred = pred.t()
        acc_5 = pred.eq(y[None])
        acc_5 = acc_5.flatten().sum(dtype=torch.float32) / y.size(0)
        accuracy_epoch_top5 += acc_5

        loss.backward()
        optimizer.step()

        loss_epoch += loss.item()
        # if step % 100 == 0:
        #     print(
        #         f"Step [{step}/{len(loader)}]\t Loss: {loss.item()}\t Accuracy: {acc}"
        #     )

    # return loss_epoch, accuracy_epoch_top1, accuracy_epoch_top5
    loss_epoch = reduce_tensor(torch.tensor(loss_epoch, device=args.device), world_size)
    accuracy_epoch_top1 = reduce_tensor(torch.tensor(accuracy_epoch_top1, device=args.device), world_size)
    accuracy_epoch_top5 = reduce_tensor(torch.tensor(accuracy_epoch_top5, device=args.device), world_size)

    return loss_epoch.item(), accuracy_epoch_top1.item(), accuracy_epoch_top5.item()


def test(args, loader, model, criterion):
    loss_epoch = 0
    accuracy_epoch_top1 = 0
    accuracy_epoch_top5 = 0
    model.eval()
    for step, (x, y) in enumerate(loader):
        model.zero_grad()

        x = x.to(args.device)
        y = y.to(args.device)

        with torch.no_grad():
            output, _ = model(x)

        loss = criterion(output, y)

        # top-1
        predicted = output.argmax(1)
        acc = (predicted == y).sum().item() / y.size(0)
        accuracy_epoch_top1 += acc
        # top-5
        _, pred = output.topk(5, 1, True, True)
        pred = pred.t()
        acc_5 = pred.eq(y[None])
        acc_5 = acc_5.flatten().sum(dtype=torch.float32) / y.size(0)
        accuracy_epoch_top5 += acc_5

        loss_epoch += loss.item()

    return loss_epoch, accuracy_epoch_top1, accuracy_epoch_top5


def main(rank, world_size, args):
    # init DDP
    args.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
    if world_size > 1:
        dist.init_process_group("nccl", init_method="env://", rank=rank, world_size=world_size, timeout=datetime.timedelta(seconds=120))
        torch.cuda.set_device(rank)

    if rank == 0:
        print("Now start loading dataset")

    if 'imagenet' in args.dataset:
        dataset_train = torchvision.datasets.ImageFolder(
            os.path.join(args.dataset, 'train'),
            transform=DataTransforms(size=args.img_size).train_evaluation_transform,
        )
        test_dataset = torchvision.datasets.ImageFolder(
            os.path.join(args.dataset, 'val'),
            transform=DataTransforms(size=args.img_size).test_transform,
        )
    else:
        raise NotImplementedError

    # Create a 1%/10%/100% data subset using the random indices
    num_samples = int(len(dataset_train) * args.eval_proportion)
    indices = torch.randperm(len(dataset_train))[:num_samples]
    train_subset = torch.utils.data.Subset(dataset_train, indices)

    if world_size > 1:
        sampler = DistributedSampler(train_subset, num_replicas=world_size, rank=rank, shuffle=True)
    else:
        sampler = None

    train_loader = torch.utils.data.DataLoader(
        train_subset,
        batch_size=args.val_batch_size,
        shuffle=(sampler is None),
        drop_last=True,
        num_workers=4,
        sampler=sampler
    )

    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=args.val_batch_size,
        shuffle=False,
        drop_last=True,
        num_workers=4,
    )

    # load pre-trained model from checkpoint
    model = create_model(
        'spikformer',
        pretrained=False,
        drop_rate=0.,
        drop_path_rate=0.,
        drop_block_rate=None,
        img_size_h=args.img_size, img_size_w=args.img_size,
        patch_size=args.patch_size, embed_dims=args.dim, num_heads=args.num_heads, mlp_ratios=args.mlp_ratio,
        in_channels=3, num_classes=args.num_classes, qkv_bias=False,
        depths=args.layer, sr_ratios=1,
        T=args.time_step
    )

    # Load weights from pre-trained file. Map the keys from the source model to the target model
    model_dict = model.state_dict()
    # pretrained_dict = torch.load(args.resume, map_location=args.device.type)
    pretrained_dict = torch.load(args.resume, map_location=args.device.type, weights_only=False)  # 或 weights_only=True 如果可以
    # print(checkpoint.keys())

    new_dict = {}
    for k, v in pretrained_dict['state_dict'].items():
        if "backbone." in k:
            # Remove the 'backbone.' prefix and use the rest as the key for the target model
            key = k.replace("backbone.", "")
            new_dict[key] = v

    model.load_state_dict(new_dict, strict=False)
    model = model.to(args.device)

    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=3e-4)
    criterion = torch.nn.CrossEntropyLoss()

    for epoch in range(60):
        if world_size > 1:  # shuffle
            sampler.set_epoch(epoch)
        loss_epoch, accuracy_epoch_top1, accuracy_epoch_top5 = train(
            args, train_loader, model, criterion, optimizer, world_size
        )
        if epoch % 20 == 0:
            print(
                f"Epoch [{epoch}/{60}]\t Loss: {loss_epoch / len(train_loader)}\t Accuracy top-1: {accuracy_epoch_top1 / len(train_loader)}\t Accuracy top-5: {accuracy_epoch_top5 / len(train_loader)}"
            )
    # final testing
    loss_epoch, accuracy_top1, accuracy_top5 = test(
        args, test_loader, model, criterion
    )
    print(
        f"[FINAL]\t Loss: {loss_epoch / len(test_loader)}\t Accuracy top-1: {accuracy_top1 / len(test_loader)}\t Accuracy top-5: {accuracy_top5 / len(test_loader)}"
    )


if __name__ == "__main__":
    args, args_text = _parse_args()
    args.lr = float(args.lr)
    print(vars(args))


    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "1234"

    world_size = torch.cuda.device_count()
    if world_size > 1:
        torch.multiprocessing.spawn(main, args=(world_size, args), nprocs=world_size, join=True)
    else:
        main(0, 1, args)