import argparse
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
# from torchvision import models
from model.wide_res_net import WideResNet
from model.smooth_cross_entropy import smooth_crossentropy
from data.cifar import Cifar
from utility.log import Log
from utility.initialize import initialize
from utility.step_lr import StepLR
import sys; sys.path.append("..")
from sam import SAM

from numpy.random import RandomState
import numpy as np
import torch
import torch.optim as optim
from torch.utils.data import Subset
import torchvision

def train(model, classifier, device, train_loader, optimizer, epoch, display=True):
    classifier.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = classifier(model(data))
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
    if display:
        print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            epoch, batch_idx * len(data), len(train_loader.dataset),
            100. * batch_idx / len(train_loader), loss.item()))

def test(model, classifier, device, test_loader):
    classifier.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = classifier(model(data))
            test_loss += F.cross_entropy(output, target, size_average=False).item() # sum up batch loss
            pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    return 100. * correct / len(test_loader.dataset)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", default=128, type=int, help="Batch size used in the training and validation loop.")
    parser.add_argument("--depth", default=16, type=int, help="Number of layers.")
    parser.add_argument("--dropout", default=0.0, type=float, help="Dropout rate.")
    parser.add_argument("--epochs", default=200, type=int, help="Total number of epochs.")
    parser.add_argument("--label_smoothing", default=0.1, type=float, help="Use 0.0 for no label smoothing.")
    parser.add_argument("--learning_rate", default=0.1, type=float, help="Base learning rate at the start of the training.")
    parser.add_argument("--momentum", default=0.9, type=float, help="SGD Momentum.")
    parser.add_argument("--threads", default=2, type=int, help="Number of CPU threads for dataloaders.")
    parser.add_argument("--rho", default=0.05, type=int, help="Rho parameter for SAM.")
    parser.add_argument("--weight_decay", default=0.0005, type=float, help="L2 weight decay.")
    parser.add_argument("--width_factor", default=8, type=int, help="How many times wider compared to normal ResNet.")
    args = parser.parse_args()

    initialize(args, seed=42)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    dataset = Cifar(args.batch_size, args.threads)
    log = Log(log_each=10)
    model = WideResNet(args.depth, args.width_factor, args.dropout, in_channels=3, labels=95).to(device)

    base_optimizer = torch.optim.SGD
    optimizer = SAM(model.parameters(), base_optimizer, rho=args.rho, lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
    scheduler = StepLR(optimizer, args.learning_rate, args.epochs)

    for epoch in range(args.epochs):
        model.train()
        log.train(len_dataset=len(dataset.train))

        for inputs, targets in dataset.train:
            for i, t in enumerate(targets):
                for label in dataset.labels_to_remove:
                    if t > label:
                        targets[i] -= 1

            inputs, targets = inputs.to(device), targets.to(device)

            # first forward-backward step
            predictions = model(inputs)
            loss = smooth_crossentropy(predictions, targets)
            loss.mean().backward()
            optimizer.first_step(zero_grad=True)

            # second forward-backward step
            smooth_crossentropy(model(inputs), targets).mean().backward()
            optimizer.second_step(zero_grad=True)

            with torch.no_grad():
                correct = torch.argmax(predictions.data, 1) == targets
                log(model, loss.cpu(), correct.cpu(), scheduler.lr())
                scheduler(epoch)

        model.eval()
        log.eval(len_dataset=len(dataset.test))

        with torch.no_grad():
            for inputs, targets in dataset.test:
                for i, t in enumerate(targets):
                    for label in dataset.labels_to_remove:
                        if t > label:
                            targets[i] -= 1

                inputs, targets = inputs.to(device), targets.to(device)

                predictions = model(inputs)
                loss = smooth_crossentropy(predictions, targets)
                correct = torch.argmax(predictions, 1) == targets
                log(model, loss.cpu(), correct.cpu())

    log.flush()

    ### persistent model
    torch.save(model.state_dict(), "./model_params.sam")


    ### transfer to cifar10
    model.f[-1] = nn.Identity()
    model.eval()
    
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                      std=[0.229, 0.224, 0.225])

    transform_val = transforms.Compose([transforms.ToTensor(), normalize]) #careful to keep this one same
    transform_train = transforms.Compose([transforms.ToTensor(), normalize]) 
    ##### Cifar Data
    cifar_data = torchvision.datasets.CIFAR10(root='.',train=True, transform=transform_train, download=True)
    cifar_data_val = torchvision.datasets.CIFAR10(root='.',train=True, transform=transform_val, download=True)

    accs = []
    for seed in [i for i in range(10)]:
        prng = RandomState(seed)
        random_permute = prng.permutation(np.arange(0, 5000))
        indx_train = np.concatenate([np.where(np.array(cifar_data.targets) == classe)[0][random_permute[0:10]] for classe in range(0, 10)])
        indx_val = np.concatenate([np.where(np.array(cifar_data.targets) == classe)[0][random_permute[10:210]] for classe in range(0, 10)])
        
        train_data = Subset(cifar_data, indx_train)
        val_data = Subset(cifar_data_val, indx_val)
        train_loader = torch.utils.data.DataLoader(train_data, batch_size=128, shuffle=True)
        val_loader = torch.utils.data.DataLoader(val_data, batch_size=128, shuffle=False)

        classifier = nn.Linear(4 * 16 * 8, 10).to(device)
        optimizer = torch.optim.Adam(classifier.parameters(), lr=3e-4, weight_decay=0.01, eps=1e-4)

        for epoch in range(25):
            train(model, classifier, device, train_loader, optimizer, epoch, (epoch + 1) % 5 == 0)
        accs.append(test(model, classifier, device, val_loader))

    accs = np.array(accs)
    print('Acc over 2 instances: %.2f +- %.2f'%(accs.mean(),accs.std()))
import argparse
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
# from torchvision import models
from model.wide_res_net import WideResNet
from model.smooth_cross_entropy import smooth_crossentropy
from data.cifar import Cifar
from utility.log import Log
from utility.initialize import initialize
from utility.step_lr import StepLR
import sys; sys.path.append("..")
from sam import SAM

from numpy.random import RandomState
import numpy as np
import torch
import torch.optim as optim
from torch.utils.data import Subset
import torchvision

def train(model, classifier, device, train_loader, optimizer, epoch, display=True):
    classifier.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = classifier(model(data))
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
    if display:
        print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            epoch, batch_idx * len(data), len(train_loader.dataset),
            100. * batch_idx / len(train_loader), loss.item()))

def test(model, classifier, device, test_loader):
    classifier.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = classifier(model(data))
            test_loss += F.cross_entropy(output, target, size_average=False).item() # sum up batch loss
            pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    return 100. * correct / len(test_loader.dataset)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", default=128, type=int, help="Batch size used in the training and validation loop.")
    parser.add_argument("--depth", default=16, type=int, help="Number of layers.")
    parser.add_argument("--dropout", default=0.0, type=float, help="Dropout rate.")
    parser.add_argument("--epochs", default=200, type=int, help="Total number of epochs.")
    parser.add_argument("--label_smoothing", default=0.1, type=float, help="Use 0.0 for no label smoothing.")
    parser.add_argument("--learning_rate", default=0.1, type=float, help="Base learning rate at the start of the training.")
    parser.add_argument("--momentum", default=0.9, type=float, help="SGD Momentum.")
    parser.add_argument("--threads", default=2, type=int, help="Number of CPU threads for dataloaders.")
    parser.add_argument("--rho", default=0.05, type=int, help="Rho parameter for SAM.")
    parser.add_argument("--weight_decay", default=0.0005, type=float, help="L2 weight decay.")
    parser.add_argument("--width_factor", default=8, type=int, help="How many times wider compared to normal ResNet.")
    args = parser.parse_args()

    initialize(args, seed=42)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    dataset = Cifar(args.batch_size, args.threads)
    log = Log(log_each=10)
    model = WideResNet(args.depth, args.width_factor, args.dropout, in_channels=3, labels=95).to(device)

    base_optimizer = torch.optim.SGD
    optimizer = SAM(model.parameters(), base_optimizer, rho=args.rho, lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
    scheduler = StepLR(optimizer, args.learning_rate, args.epochs)

    for epoch in range(args.epochs):
        model.train()
        log.train(len_dataset=len(dataset.train))

        for inputs, targets in dataset.train:
            for i, t in enumerate(targets):
                for label in dataset.labels_to_remove:
                    if t > label:
                        targets[i] -= 1

            inputs, targets = inputs.to(device), targets.to(device)

            # first forward-backward step
            predictions = model(inputs)
            loss = smooth_crossentropy(predictions, targets)
            loss.mean().backward()
            optimizer.first_step(zero_grad=True)

            # second forward-backward step
            smooth_crossentropy(model(inputs), targets).mean().backward()
            optimizer.second_step(zero_grad=True)

            with torch.no_grad():
                correct = torch.argmax(predictions.data, 1) == targets
                log(model, loss.cpu(), correct.cpu(), scheduler.lr())
                scheduler(epoch)

        model.eval()
        log.eval(len_dataset=len(dataset.test))

        with torch.no_grad():
            for inputs, targets in dataset.test:
                for i, t in enumerate(targets):
                    for label in dataset.labels_to_remove:
                        if t > label:
                            targets[i] -= 1

                inputs, targets = inputs.to(device), targets.to(device)

                predictions = model(inputs)
                loss = smooth_crossentropy(predictions, targets)
                correct = torch.argmax(predictions, 1) == targets
                log(model, loss.cpu(), correct.cpu())

    log.flush()

    ### persistent model
    torch.save(model.state_dict(), "./model_params.sam")


    ### transfer to cifar10
    model.f[-1] = nn.Identity()
    model.eval()
    
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                      std=[0.229, 0.224, 0.225])

    transform_val = transforms.Compose([transforms.ToTensor(), normalize]) #careful to keep this one same
    transform_train = transforms.Compose([transforms.ToTensor(), normalize]) 
    ##### Cifar Data
    cifar_data = torchvision.datasets.CIFAR10(root='.',train=True, transform=transform_train, download=True)
    cifar_data_val = torchvision.datasets.CIFAR10(root='.',train=True, transform=transform_val, download=True)

    accs = []
    for seed in [i for i in range(10)]:
        prng = RandomState(seed)
        random_permute = prng.permutation(np.arange(0, 5000))
        indx_train = np.concatenate([np.where(np.array(cifar_data.targets) == classe)[0][random_permute[0:10]] for classe in range(0, 10)])
        indx_val = np.concatenate([np.where(np.array(cifar_data.targets) == classe)[0][random_permute[10:210]] for classe in range(0, 10)])
        
        train_data = Subset(cifar_data, indx_train)
        val_data = Subset(cifar_data_val, indx_val)
        train_loader = torch.utils.data.DataLoader(train_data, batch_size=128, shuffle=True)
        val_loader = torch.utils.data.DataLoader(val_data, batch_size=128, shuffle=False)

        classifier = nn.Linear(4 * 16 * 8, 10).to(device)
        optimizer = torch.optim.Adam(classifier.parameters(), lr=3e-4, weight_decay=0.01, eps=1e-4)

        for epoch in range(25):
            train(model, classifier, device, train_loader, optimizer, epoch, (epoch + 1) % 5 == 0)
        accs.append(test(model, classifier, device, val_loader))

    accs = np.array(accs)
    print('Acc over 2 instances: %.2f +- %.2f'%(accs.mean(),accs.std()))
# from torchvision import models # sum up batch loss # get the index of the max log-probability # first forward-backward step # second forward-backward step ### persistent model ### transfer to cifar10 #careful to keep this one same ##### Cifar Data