import argparse

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

from n3ml.model import DynamicModel_STBP_SNN


def validate(val_loader, model, encoder, criterion, opt):
    model.eval()

    total_images = 0
    num_corrects = 0
    total_loss = 0

    with torch.no_grad():
        for step, (images, labels) in enumerate(val_loader):
            images = images.cuda()
            labels = labels.cuda()

            preds = model(encoder, images, opt.num_steps)
            labels_ = torch.zeros(torch.numel(labels), 10, device=labels.device)
            labels_ = labels_.scatter_(1, labels.view(-1, 1), 1)

            loss = criterion(preds, labels_)

            num_corrects += torch.argmax(preds, dim=1).eq(labels).sum(dim=0)
            total_loss += loss.cpu().detach().numpy() * images.size(0)
            total_images += images.size(0)

    val_acc = num_corrects.float() / total_images
    val_loss = total_loss / total_images

    return val_acc, val_loss


def app(opt):
    print(opt)

    val_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST(
            opt.data,
            train=False,
            download=True,
            transform=torchvision.transforms.Compose([transforms.ToTensor()])),
        batch_size=opt.batch_size)

    state_dict = torch.load(opt.pretrained)

    model = DynamicModel_STBP_SNN(batch_size=opt.batch_size)
    for m in state_dict['arch']:
        model.add_module(m[0], m[1])

    if torch.cuda.is_available():
        model.cuda()

    encoder = lambda x: (x > torch.rand(x.size(), device=x.device)).float()

    criterion = nn.MSELoss()

    acc, loss = validate(val_loader, model, encoder, criterion, opt)
    print("In test, loss: {} - acc: {}".format(loss, acc))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--data', default='data')
    parser.add_argument('--batch_size', default=100, type=int)
    parser.add_argument('--num_steps', default=15, type=int)
    parser.add_argument('--pretrained', default='pretrained/stbp_dynamic_acc_9897.pt')

    app(parser.parse_args())

[((50, 17, 50, 43), 'torch.load', 'torch.load', ({(50, 28, 50, 42): 'opt.pretrained'}, {}), '(opt.pretrained)', False, 'import torch\n'), ((52, 12, 52, 60), 'n3ml.model.DynamicModel_STBP_SNN', 'DynamicModel_STBP_SNN', (), '', False, 'from n3ml.model import DynamicModel_STBP_SNN\n'), ((56, 7, 56, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((61, 16, 61, 28), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((68, 13, 68, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((18, 9, 18, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((24, 34, 24, 53), 'torch.numel', 'torch.numel', ({(24, 46, 24, 52): 'labels'}, {}), '(labels)', False, 'import torch\n'), ((47, 54, 47, 75), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', True, 'import torchvision.transforms as transforms\n'), ((29, 28, 29, 54), 'torch.argmax', 'torch.argmax', (), '', False, 'import torch\n')]