# (c) 2024 LAiSR-SK

# This code is licensed under the MIT license (see LICENSE.md).

"""Train Adversarially Robust Models with Feature Scattering"""

import os
from collections import namedtuple
from os import PathLike

import torch
import torch.backends.cudnn as cudnn
import torchvision
import torchvision.transforms as transforms
from torch import optim
from tqdm import tqdm

from airtk.defense.feature_scatter.fs_attack_methods import (
    Attack_FeaScatter,
)
from airtk.model import WideResNet


class FeatureScatterTraining:
    def __init__(
        self,
        dataset: str,
        n_classes: int = 10,
        image_size: int = 32,
        epsilon: float = 8.0 / 512,
        lr: float = 0.1,
        num_steps: int = 1,
        step_size: int = 8.0 / 512,
        ls_factor: float = 0.5,
        epochs: int = 200,
        batch_size: int = 128,
        momentum: float = 0.9,
        save_epochs: int = 100,
        weight_decay: float = 2e-4,
        decay_epoch1: int = 60,
        decay_epoch2: int = 90,
        decay_rate: float = 0.1,
        log_step: int = 10,
        model_dir: PathLike = "data/model",
    ) -> None:
        ArgPrototype = namedtuple(
            "ArgPrototype",
            [
                "resume",
                "lr",
                "adv_mode",
                "model_dir",
                "init_model_pass",
                "max_epoch",
                "save_epochs",
                "decay_epoch1",
                "decay_epoch2",
                "decay_rate",
                "batch_size_train",
                "momentum",
                "weight_decay",
                "log_step",
                "num_classes",
                "image_size",
                "dataset",
            ],
        )

        if dataset == "cifar10":
            n_classes = 10

            image_size = 32

        elif dataset == "cifar100":
            n_classes = 100

            image_size = 32

        self._args = ArgPrototype(
            False,
            lr,
            "feature_scatter",
            model_dir,
            -1,
            epochs,
            save_epochs,
            decay_epoch1,
            decay_epoch2,
            decay_rate,
            batch_size,
            momentum,
            weight_decay,
            log_step,
            n_classes,
            image_size,
            dataset,
        )
        self._ds_name = dataset

        self._config = {
            "train": True,
            "epsilon": epsilon,
            "num_steps": num_steps,
            "step_size": step_size,
            "random_start": True,
            "ls_factor": ls_factor,
        }

    def __call__(self) -> None:
        main_featurescatter(self._ds_name, self._args, self._config)


# parser = argparse.ArgumentParser(description="Feature Scatterring Training")


# add type keyword to registries

# parser.register("type", "bool", str2bool)


# parser.add_argument(

# "--resume", "-r", action="store_true", help="resume from checkpoint"

# )

# parser.add_argument("--lr", default=0.1, type=float, help="learning rate")

# parser.add_argument(

# "--adv_mode",

# default="feature_scatter",

# type=str,

# help="adv_mode (feature_scatter)",

# )

# parser.add_argument(

# "--model_dir", default="\data\model", type=str, help="model path"

# )

# parser.add_argument(

# "--init_model_pass",

# default="-1",

# type=str,

# help="init model pass (-1: from scratch; K: checkpoint-K)",

# )

# parser.add_argument(

# "--max_epoch", default=200, type=int, help="max number of epochs"

# )

# parser.add_argument("--save_epochs", default=100, type=int, help="save period")

# parser.add_argument(

# "--decay_epoch1",

# default=60,

# type=int,

# help="learning rate decay epoch one",

# )

# parser.add_argument(

# "--decay_epoch2",

# default=90,

# type=int,

# help="learning rate decay point two",

# )

# parser.add_argument(

# "--decay_rate", default=0.1, type=float, help="learning rate decay rate"

# )

# parser.add_argument(

# "--batch_size_train", default=128, type=int, help="batch size for training"

# )

# parser.add_argument(

# "--momentum", default=0.9, type=float, help="momentum (1-tf.momentum)"

# )

# parser.add_argument(

# "--weight_decay", default=2e-4, type=float, help="weight decay"

# )

# parser.add_argument("--log_step", default=10, type=int, help="log_step")

#

# number of classes and image size will be updated below based on the dataset

# parser.add_argument("--num_classes", default=10, type=int, help="num classes")

# parser.add_argument("--image_size", default=32, type=int, help="image size")

# parser.add_argument(

# "--dataset", default="cifar10", type=str, help="dataset"

# )  # concat cascade


# args = parser.parse_args()

# device = "cuda" if torch.cuda.is_available() else "cpu"


def print_para(net):
    for name, param in net.named_parameters():
        if param.requires_grad:
            print(name)
            print(param.data)

        break


def zero_gradients(x):
    """Manually zeroes the gradients when pytorch is unable"""

    if x.grad is not None:
        x.grad.zero_()


def train_fun(epoch, net, trainloader, optimizer, args):
    print("\nEpoch: %d" % epoch)
    net.train()

    # update learning rate

    if epoch < args.decay_epoch1:
        lr = args.lr

    elif epoch < args.decay_epoch2:
        lr = args.lr * args.decay_rate

    else:
        lr = args.lr * args.decay_rate * args.decay_rate

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    def get_acc(outputs, targets):
        _, predicted = outputs.max(1)

        total = targets.size(0)

        correct = predicted.eq(targets).sum().item()

        acc = 1.0 * correct / total
        return acc

    iterator = tqdm(trainloader, ncols=0, leave=False)

    device = "cuda" if torch.cuda.is_available() else "cpu"

    for batch_idx, (inputs, targets) in enumerate(iterator):
        inputs, targets = inputs.to(device), targets.to(device)

        adv_acc = 0

        optimizer.zero_grad()

        # forward
        outputs, loss_fs = net(inputs.detach(), targets)

        optimizer.zero_grad()
        loss = loss_fs.mean()

        loss.backward()

        optimizer.step()

        if batch_idx % args.log_step == 0:
            if adv_acc == 0:
                adv_acc = get_acc(outputs, targets)

            iterator.set_description(str(adv_acc))

            nat_outputs, _ = net(inputs, targets, attack=False)

            """print(

                "epoch %d, step %d, lr %.4f, duration %.2f, training nat acc %.2f, training adv acc %.2f, training adv loss %.4f"

                % (epoch, batch_idx, lr, duration, 100 * nat_acc,

                   100 * adv_acc, train_loss))"""

    if epoch % args.save_epochs == 0 or epoch >= args.max_epoch - 2:
        print("Saving..")

        f_path = os.path.join(args.model_dir, ("checkpoint-%s" % epoch))

        state = {
            "net": net.state_dict(),
            # 'optimizer': optimizer.state_dict()
        }

        if not os.path.isdir(args.model_dir):
            os.mkdir(args.model_dir)

        torch.save(state, f_path)

    if epoch >= 0:
        print("Saving latest @ epoch %s.." % (epoch))

        f_path = os.path.join(args.model_dir, "latest")

        state = {
            "net": net.state_dict(),
            "epoch": epoch,
            "optimizer": optimizer.state_dict(),
        }

        if not os.path.isdir(args.model_dir):
            os.mkdir(args.model_dir)

        torch.save(state, f_path)

        torch.save(
            net.state_dict(),
            f"data/model/feature-attack-epoch{epoch}.pt",
        )

        print("Epoch " + str(epoch) + " saved")


def main_featurescatter(ds_name, args, config_feature_scatter):
    start_epoch = 0

    transform_train = transforms.Compose(
        [
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # [-1 1]
        ]
    )

    if args.dataset == "cifar10":
        trainset = torchvision.datasets.CIFAR10(
            root="./data", train=True, download=True, transform=transform_train
        )

    elif args.dataset == "cifar100":
        trainset = torchvision.datasets.CIFAR100(
            root="./data", train=True, download=True, transform=transform_train
        )

    else:
        raise NotImplementedError()

    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=2
    )

    print("---wide resenet-----")

    device = "cuda" if torch.cuda.is_available() else "cpu"

    basic_net = WideResNet(
        depth=34, num_classes=args.num_classes, widen_factor=10
    ).to(device)

    if args.adv_mode.lower() == "feature_scatter":
        print("-----Feature Scatter mode -----")

        net = Attack_FeaScatter(basic_net, config_feature_scatter)

    else:
        print("-----OTHER_ALGO mode -----")

        raise NotImplementedError("Please implement this algorithm first!")

    if device == "cuda":
        net = torch.nn.DataParallel(net)

        cudnn.benchmark = True

    optimizer = optim.SGD(
        net.parameters(),
        lr=args.lr,
        momentum=args.momentum,
        weight_decay=args.weight_decay,
    )

    if args.resume and args.init_model_pass != "-1":
        # Load checkpoint.

        print("==> Resuming from checkpoint..")

        f_path_latest = os.path.join(args.model_dir, "latest")

        f_path = os.path.join(
            args.model_dir, ("checkpoint-%s" % args.init_model_pass)
        )

        if not os.path.isdir(args.model_dir):
            print("train from scratch: no checkpoint directory or file found")
        elif args.init_model_pass == "latest" and os.path.isfile(
            f_path_latest
        ):
            checkpoint = torch.load(f_path_latest)

            net.load_state_dict(checkpoint["net"])

            start_epoch = checkpoint["epoch"] + 1

            print("resuming from epoch %s in latest" % start_epoch)

        elif os.path.isfile(f_path):
            checkpoint = torch.load(f_path)

            net.load_state_dict(checkpoint["net"])

            start_epoch = checkpoint["epoch"] + 1

            print("resuming from epoch %s" % (start_epoch - 1))

        elif not os.path.isfile(f_path) or not os.path.isfile(f_path_latest):
            print("train from scratch: no checkpoint directory or file found")

    for epoch in range(1, args.max_epoch):
        train_fun(epoch, net, trainloader, optimizer, args)


if __name__ == "__main__":
    main_featurescatter("cifar10")
