import os
import argparse
import numpy as np
from tqdm import tqdm
import pandas as pd
import joblib
from collections import OrderedDict
from datetime import datetime
import matplotlib.pyplot as plt

import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets

from utils import *
from scaled_mnist.dataset import ScaledMNIST
import scaled_mnist.archs as archs

import torch
from torch.autograd import Variable

arch_names = archs.__dict__.keys()


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--name', default=None,
                        help='model name: (default: arch+timestamp)')
    parser.add_argument('--arch', '-a', metavar='ARCH', default='ScaledMNISTNet',
                        choices=arch_names,
                        help='model architecture: ' +
                             ' | '.join(arch_names) +
                             ' (default: ScaledMNISTNet)')
    parser.add_argument('--deform', default=True, type=str2bool,
                        help='use deform conv')
    parser.add_argument('--modulation', default=True, type=str2bool,
                        help='use modulated deform conv')
    parser.add_argument('--min-deform-layer', default=3, type=int,
                        help='minimum number of layer using deform conv')
    parser.add_argument('--epochs', default=10, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('--optimizer', default='SGD',
                        choices=['Adam', 'SGD'],
                        help='loss: ' +
                             ' | '.join(['Adam', 'SGD']) +
                             ' (default: Adam)')
    parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float,
                        metavar='LR', help='initial learning rate')
    parser.add_argument('--momentum', default=0.5, type=float,
                        help='momentum')
    parser.add_argument('--weight-decay', default=1e-4, type=float,
                        help='weight decay')
    parser.add_argument('--nesterov', default=False, type=str2bool,
                        help='nesterov')

    args = parser.parse_args()

    return args


def normalize_output(img):
    img = img - img.min()
    img = img / img.max()
    return img


def train(args, train_loader, model, criterion, optimizer, epoch, scheduler=None):
    losses = AverageMeter()
    scores = AverageMeter()

    model.train()

    for i, (input, target) in tqdm(enumerate(train_loader), total=len(train_loader)):
        input = input.cuda()
        target = target.cuda()

        output = model(input)
        loss = criterion(output, target)

        acc = accuracy(output, target)[0]

        losses.update(loss.item(), input.size(0))
        scores.update(acc.item(), input.size(0))

        # compute gradient and do optimizing step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    log = OrderedDict([
        ('loss', losses.avg),
        ('acc', scores.avg),
    ])

    return log


def validate(args, val_loader, model, criterion):
    losses = AverageMeter()
    scores = AverageMeter()

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        for i, (input, target) in tqdm(enumerate(val_loader), total=len(val_loader)):
            input = input.cuda()
            target = target.cuda()

            output = model(input)
            loss = criterion(output, target)

            acc = accuracy(output, target)[0]

            losses.update(loss.item(), input.size(0))
            scores.update(acc.item(), input.size(0))

    log = OrderedDict([
        ('loss', losses.avg),
        ('acc', scores.avg),
    ])

    return log


def main():
    args = parse_args()

    if args.name is None:
        args.name = '%s' % args.arch
        if args.deform:
            args.name += '_wDCN'
            if args.modulation:
                args.name += 'v2'
            args.name += '_c%d-4' % args.min_deform_layer

    if not os.path.exists('models/%s' % args.name):
        os.makedirs('models/%s' % args.name)

    print('Config -----')
    for arg in vars(args):
        print('%s: %s' % (arg, getattr(args, arg)))
    print('------------')

    with open('models/%s/args.txt' % args.name, 'w') as f:
        for arg in vars(args):
            print('%s: %s' % (arg, getattr(args, arg)), file=f)

    joblib.dump(args, 'models/%s/args.pkl' % args.name)

    criterion = nn.CrossEntropyLoss().cuda()

    cudnn.benchmark = True
    ##########################################################################
    # data loading code
    # transform_train = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))
    # ])
    #
    # transform_test = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))
    # ])
    #
    # train_set = ScaledMNIST(
    #     train=True,
    #     transform=transform_train)
    # train_loader = torch.utils.data.DataLoader(
    #     train_set,
    #     batch_size=32,
    #     shuffle=True,
    #     num_workers=8)
    #
    # test_set = ScaledMNIST(
    #     train=False,
    #     transform=transform_train)
    # test_loader = torch.utils.data.DataLoader(
    #     test_set,
    #     batch_size=32,
    #     shuffle=False,
    #     num_workers=8)
    ##########################################################################

    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_dataset = datasets.CIFAR10('~/data', train=True, download=True, transform=transform)
    test_dataset = datasets.CIFAR10('~/data', train=False, download=True, transform=transform)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=0)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=0)

    num_classes = 10

    # create model
    model = archs.__dict__[args.arch](args, num_classes)
    model = model.cuda()

    print(model)

    if args.optimizer == 'Adam':
        optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr)
    elif args.optimizer == 'SGD':
        optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr,
                              momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov)

    log = pd.DataFrame(index=[], columns=[
        'epoch', 'lr', 'loss', 'acc', 'val_loss', 'val_acc'
    ])

    best_acc = 0
    for epoch in range(args.epochs):
        print('Epoch [%d/%d]' % (epoch, args.epochs))

        # train for one epoch
        train_log = train(args, train_loader, model, criterion, optimizer, epoch)
        # evaluate on validation set
        val_log = validate(args, test_loader, model, criterion)

        print('loss %.4f - acc %.4f - val_loss %.4f - val_acc %.4f'
              % (train_log['loss'], train_log['acc'], val_log['loss'], val_log['acc']))

        tmp = pd.Series([
            epoch,
            1e-1,
            train_log['loss'],
            train_log['acc'],
            val_log['loss'],
            val_log['acc'],
        ], index=['epoch', 'lr', 'loss', 'acc', 'val_loss', 'val_acc'])

        log = log.append(tmp, ignore_index=True)
        log.to_csv('models/%s/log.csv' % args.name, index=False)

        if val_log['acc'] > best_acc:
            torch.save(model.state_dict(), 'models/%s/model.pth' % args.name)
            best_acc = val_log['acc']
            print("=> saved best model")

    print("best val_acc: %f" % best_acc)


def viualization_featureMap():
    args = parse_args()
    model = archs.__dict__[args.arch](args, 10)

    # for mnist load data
    # transform_train = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))
    # ])
    # test_set = ScaledMNIST(
    #     train=True,
    #     transform=transform_train)
    # test_loader = torch.utils.data.DataLoader(
    #     test_set,
    #     batch_size=32,
    #     shuffle=True,
    #     num_workers=0)

    # for cifar10 load data
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_dataset = datasets.CIFAR10('~/data', train=True, download=True, transform=transform)
    test_dataset = datasets.CIFAR10('~/data', train=False, download=True, transform=transform)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=0)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=0)

    unloader = transforms.ToPILImage()
    for i, (input, target) in tqdm(enumerate(test_loader), total=len(test_loader)):
        # input = input.cuda()
        target = target.cuda()

        output = model(input)
        ################################################
        # Plot some images
        idx = torch.randint(0, output.size(0), ())
        pred = normalize_output(output)[idx, 0]
        img = input[idx, 0]

        fig, axarr = plt.subplots(1, 2)
        print('original image:', img.detach().cpu().numpy().shape)
        print('predict image:', pred.detach().cpu().numpy().shape)
        axarr[0].imshow(img.detach().cpu().numpy())
        axarr[1].imshow(pred.detach().cpu().numpy())
        plt.show()
        ################################################


if __name__ == '__main__':
    main()

    # 修改archs.py中的def forward()，return self.features(input)
    # viualization_featureMap()

import os
import argparse
import numpy as np
from tqdm import tqdm
import pandas as pd
import joblib
from collections import OrderedDict
from datetime import datetime
import matplotlib.pyplot as plt

import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets

from utils import *
from scaled_mnist.dataset import ScaledMNIST
import scaled_mnist.archs as archs

import torch
from torch.autograd import Variable

arch_names = archs.__dict__.keys()


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--name', default=None,
                        help='model name: (default: arch+timestamp)')
    parser.add_argument('--arch', '-a', metavar='ARCH', default='ScaledMNISTNet',
                        choices=arch_names,
                        help='model architecture: ' +
                             ' | '.join(arch_names) +
                             ' (default: ScaledMNISTNet)')
    parser.add_argument('--deform', default=True, type=str2bool,
                        help='use deform conv')
    parser.add_argument('--modulation', default=True, type=str2bool,
                        help='use modulated deform conv')
    parser.add_argument('--min-deform-layer', default=3, type=int,
                        help='minimum number of layer using deform conv')
    parser.add_argument('--epochs', default=10, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('--optimizer', default='SGD',
                        choices=['Adam', 'SGD'],
                        help='loss: ' +
                             ' | '.join(['Adam', 'SGD']) +
                             ' (default: Adam)')
    parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float,
                        metavar='LR', help='initial learning rate')
    parser.add_argument('--momentum', default=0.5, type=float,
                        help='momentum')
    parser.add_argument('--weight-decay', default=1e-4, type=float,
                        help='weight decay')
    parser.add_argument('--nesterov', default=False, type=str2bool,
                        help='nesterov')

    args = parser.parse_args()

    return args


def normalize_output(img):
    img = img - img.min()
    img = img / img.max()
    return img


def train(args, train_loader, model, criterion, optimizer, epoch, scheduler=None):
    losses = AverageMeter()
    scores = AverageMeter()

    model.train()

    for i, (input, target) in tqdm(enumerate(train_loader), total=len(train_loader)):
        input = input.cuda()
        target = target.cuda()

        output = model(input)
        loss = criterion(output, target)

        acc = accuracy(output, target)[0]

        losses.update(loss.item(), input.size(0))
        scores.update(acc.item(), input.size(0))

        # compute gradient and do optimizing step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    log = OrderedDict([
        ('loss', losses.avg),
        ('acc', scores.avg),
    ])

    return log


def validate(args, val_loader, model, criterion):
    losses = AverageMeter()
    scores = AverageMeter()

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        for i, (input, target) in tqdm(enumerate(val_loader), total=len(val_loader)):
            input = input.cuda()
            target = target.cuda()

            output = model(input)
            loss = criterion(output, target)

            acc = accuracy(output, target)[0]

            losses.update(loss.item(), input.size(0))
            scores.update(acc.item(), input.size(0))

    log = OrderedDict([
        ('loss', losses.avg),
        ('acc', scores.avg),
    ])

    return log


def main():
    args = parse_args()

    if args.name is None:
        args.name = '%s' % args.arch
        if args.deform:
            args.name += '_wDCN'
            if args.modulation:
                args.name += 'v2'
            args.name += '_c%d-4' % args.min_deform_layer

    if not os.path.exists('models/%s' % args.name):
        os.makedirs('models/%s' % args.name)

    print('Config -----')
    for arg in vars(args):
        print('%s: %s' % (arg, getattr(args, arg)))
    print('------------')

    with open('models/%s/args.txt' % args.name, 'w') as f:
        for arg in vars(args):
            print('%s: %s' % (arg, getattr(args, arg)), file=f)

    joblib.dump(args, 'models/%s/args.pkl' % args.name)

    criterion = nn.CrossEntropyLoss().cuda()

    cudnn.benchmark = True
    ##########################################################################
    # data loading code
    # transform_train = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))
    # ])
    #
    # transform_test = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))
    # ])
    #
    # train_set = ScaledMNIST(
    #     train=True,
    #     transform=transform_train)
    # train_loader = torch.utils.data.DataLoader(
    #     train_set,
    #     batch_size=32,
    #     shuffle=True,
    #     num_workers=8)
    #
    # test_set = ScaledMNIST(
    #     train=False,
    #     transform=transform_train)
    # test_loader = torch.utils.data.DataLoader(
    #     test_set,
    #     batch_size=32,
    #     shuffle=False,
    #     num_workers=8)
    ##########################################################################

    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_dataset = datasets.CIFAR10('~/data', train=True, download=True, transform=transform)
    test_dataset = datasets.CIFAR10('~/data', train=False, download=True, transform=transform)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=0)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=0)

    num_classes = 10

    # create model
    model = archs.__dict__[args.arch](args, num_classes)
    model = model.cuda()

    print(model)

    if args.optimizer == 'Adam':
        optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr)
    elif args.optimizer == 'SGD':
        optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr,
                              momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov)

    log = pd.DataFrame(index=[], columns=[
        'epoch', 'lr', 'loss', 'acc', 'val_loss', 'val_acc'
    ])

    best_acc = 0
    for epoch in range(args.epochs):
        print('Epoch [%d/%d]' % (epoch, args.epochs))

        # train for one epoch
        train_log = train(args, train_loader, model, criterion, optimizer, epoch)
        # evaluate on validation set
        val_log = validate(args, test_loader, model, criterion)

        print('loss %.4f - acc %.4f - val_loss %.4f - val_acc %.4f'
              % (train_log['loss'], train_log['acc'], val_log['loss'], val_log['acc']))

        tmp = pd.Series([
            epoch,
            1e-1,
            train_log['loss'],
            train_log['acc'],
            val_log['loss'],
            val_log['acc'],
        ], index=['epoch', 'lr', 'loss', 'acc', 'val_loss', 'val_acc'])

        log = log.append(tmp, ignore_index=True)
        log.to_csv('models/%s/log.csv' % args.name, index=False)

        if val_log['acc'] > best_acc:
            torch.save(model.state_dict(), 'models/%s/model.pth' % args.name)
            best_acc = val_log['acc']
            print("=> saved best model")

    print("best val_acc: %f" % best_acc)


def viualization_featureMap():
    args = parse_args()
    model = archs.__dict__[args.arch](args, 10)

    # for mnist load data
    # transform_train = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))
    # ])
    # test_set = ScaledMNIST(
    #     train=True,
    #     transform=transform_train)
    # test_loader = torch.utils.data.DataLoader(
    #     test_set,
    #     batch_size=32,
    #     shuffle=True,
    #     num_workers=0)

    # for cifar10 load data
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_dataset = datasets.CIFAR10('~/data', train=True, download=True, transform=transform)
    test_dataset = datasets.CIFAR10('~/data', train=False, download=True, transform=transform)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=0)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=0)

    unloader = transforms.ToPILImage()
    for i, (input, target) in tqdm(enumerate(test_loader), total=len(test_loader)):
        # input = input.cuda()
        target = target.cuda()

        output = model(input)
        ################################################
        # Plot some images
        idx = torch.randint(0, output.size(0), ())
        pred = normalize_output(output)[idx, 0]
        img = input[idx, 0]

        fig, axarr = plt.subplots(1, 2)
        print('original image:', img.detach().cpu().numpy().shape)
        print('predict image:', pred.detach().cpu().numpy().shape)
        axarr[0].imshow(img.detach().cpu().numpy())
        axarr[1].imshow(pred.detach().cpu().numpy())
        plt.show()
        ################################################


if __name__ == '__main__':
    main()

    # 修改archs.py中的def forward()，return self.features(input)
    # viualization_featureMap()

# compute gradient and do optimizing step # switch to evaluate mode ########################################################################## # data loading code # transform_train = transforms.Compose([ #     transforms.ToTensor(), #     transforms.Normalize((0.1307,), (0.3081,)) # ]) # # transform_test = transforms.Compose([ #     transforms.ToTensor(), #     transforms.Normalize((0.1307,), (0.3081,)) # ]) # # train_set = ScaledMNIST( #     train=True, #     transform=transform_train) # train_loader = torch.utils.data.DataLoader( #     train_set, #     batch_size=32, #     shuffle=True, #     num_workers=8) # # test_set = ScaledMNIST( #     train=False, #     transform=transform_train) # test_loader = torch.utils.data.DataLoader( #     test_set, #     batch_size=32, #     shuffle=False, #     num_workers=8) ########################################################################## # create model # train for one epoch # evaluate on validation set # for mnist load data # transform_train = transforms.Compose([ #     transforms.ToTensor(), #     transforms.Normalize((0.1307,), (0.3081,)) # ]) # test_set = ScaledMNIST( #     train=True, #     transform=transform_train) # test_loader = torch.utils.data.DataLoader( #     test_set, #     batch_size=32, #     shuffle=True, #     num_workers=0) # for cifar10 load data # input = input.cuda() ################################################ # Plot some images ################################################ # 修改archs.py中的def forward()，return self.features(input) # viualization_featureMap()