# PyTorch implementation of
# https://github.com/tensorflow/privacy/blob/master/research/mi_lira_2021/train.py
#
# author: Chenxiang Zhang (orientino)

import argparse
import os
import time
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
import wandb
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import models, transforms
from torchvision.datasets import CIFAR10
from tqdm import tqdm
from opacus.validators import ModuleValidator
from opacus import PrivacyEngine
from opacus.utils.batch_memory_manager import BatchMemoryManager
import pyvacy
#from pyvacy import optim#, analysis, sampling

from wide_resnet import WideResNet

parser = argparse.ArgumentParser()
parser.add_argument("--lr", default=0.1, type=float)
parser.add_argument("--epochs", default=1, type=int)
parser.add_argument("--n_shadows", default=16, type=int)
parser.add_argument("--shadow_id", default=1, type=int)
parser.add_argument("--model", default="resnet18", type=str)
parser.add_argument("--pkeep", default=0.5, type=float)
parser.add_argument("--savedir", default="exp/cifar10", type=str)
parser.add_argument("--debug", action="store_true")
args = parser.parse_args()

DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps")
EPOCHS = args.epochs


def run():
    seed = np.random.randint(0, 1000000000)
    seed ^= int(time.time())
    pl.seed_everything(seed)

    args.debug = True
    wandb.init(project="lira", mode="disabled" if args.debug else "online")
    wandb.config.update(args)

    # Dataset
    train_transform = transforms.Compose(
        [
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(32, padding=4),
            transforms.ToTensor(),
            transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]),
        ]
    )
    test_transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]),
        ]
    )
    datadir = Path().home() / "opt/data/cifar"
    train_ds = CIFAR10(root=datadir, train=True, download=True, transform=train_transform)
    test_ds = CIFAR10(root=datadir, train=False, download=True, transform=test_transform)

    # Compute the IN / OUT subset:
    # If we run each experiment independently then even after a lot of trials
    # there will still probably be some examples that were always included
    # or always excluded. So instead, with experiment IDs, we guarantee that
    # after `args.n_shadows` are done, each example is seen exactly half
    # of the time in train, and half of the time not in train.

    size = len(train_ds)
    np.random.seed(seed)
    if args.n_shadows is not None:
        np.random.seed(0)
        keep = np.random.uniform(0, 1, size=(args.n_shadows, size))
        order = keep.argsort(0)
        keep = order < int(args.pkeep * args.n_shadows)
        keep = np.array(keep[args.shadow_id], dtype=bool)
        keep = keep.nonzero()[0]
    else:
        keep = np.random.choice(size, size=int(args.pkeep * size), replace=False)
        keep.sort()
    keep_bool = np.full((size), False)
    keep_bool[keep] = True

    train_ds = torch.utils.data.Subset(train_ds, keep)
    train_dl = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=4)
    test_dl = DataLoader(test_ds, batch_size=128, shuffle=False, num_workers=4)

    # Model
    if args.model == "wresnet28-2":
        m = WideResNet(28, 2, 0.0, 10)
        print("one")
    elif args.model == "wresnet28-10":
        m = WideResNet(28, 10, 0.3, 10)
        print("two")
    elif args.model == "resnet18":
        print("three")
        m = models.resnet18(weights=None, num_classes=10)
        m.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        m.maxpool = nn.Identity()
    else:
        raise NotImplementedError
    m = m.to(DEVICE)

    m = ModuleValidator.fix(m)
    ModuleValidator.validate(m, strict=True)

    optim = torch.optim.SGD(m.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
    #optim = pyvacy.DPSGD(
    #    params=m.parameters(),
    #    lr=args.lr,
    #    momentum=0.9,
    #    weight_decay=5e-4,
    #)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=args.epochs)

    privacy_engine = PrivacyEngine()
    m, optim, train_dl = privacy_engine.make_private_with_epsilon(
        module=m,
        optimizer=optim,
        data_loader=train_dl,
        epochs=args.epochs,
        target_epsilon=1,
        target_delta=1e-4,
        max_grad_norm=1.0,
        batch_first=True,
    )

    print(f"Device: {DEVICE}")
    accumulation_steps = 10

    # Train
    with BatchMemoryManager(
        data_loader=train_dl,
        max_physical_batch_size=1000,
        optimizer=optim
    ) as memory_safe_data_loader:
        for i in tqdm(range(args.epochs)):
            m.train()
            loss_total = 0
            pbar = tqdm(memory_safe_data_loader, leave=False)
            #pbar = tqdm(train_dl, leave=False)
            for itr, (x, y) in enumerate(pbar):
                x, y = x.to(DEVICE), y.to(DEVICE)
                if False:
                    loss = F.cross_entropy(m(x), y) / accumulation_steps
                    loss_norm = loss / accumulation_steps
                    loss_total += loss_norm
                    loss_norm.backward()
                    pbar.set_postfix_str(f"loss: {loss:.2f}")

                    if ((itr + 1) % accumulation_steps == 0) or (itr + 1 == len(memory_safe_data_loader)):
                        optim.step()
                        optim.zero_grad()
                else:
                    loss = F.cross_entropy(m(x), y)
                    loss_total += loss

                    pbar.set_postfix_str(f"loss: {loss:.2f}")
                    optim.zero_grad()
                    loss.backward()
                    optim.step()
            sched.step()

            wandb.log({"loss": loss_total / len(train_dl)})

    print(f"[test] acc_test: {get_acc(m, test_dl):.4f}")
    wandb.log({"acc_test": get_acc(m, test_dl)})

    savedir = os.path.join(args.savedir, str(args.shadow_id))
    os.makedirs(savedir, exist_ok=True)
    np.save(savedir + "/keep.npy", keep_bool)
    torch.save(m.state_dict(), savedir + "/model.pt")


@torch.no_grad()
def get_acc(model, dl):
    acc = []
    for x, y in dl:
        x, y = x.to(DEVICE), y.to(DEVICE)
        acc.append(torch.argmax(model(x), dim=1) == y)
    acc = torch.cat(acc)
    acc = torch.sum(acc) / len(acc)

    return acc.item()


if __name__ == "__main__":
    run()