import os
import time
import argparse
import sys
import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
from torch.cuda import amp
from torch.utils.tensorboard import SummaryWriter
import torchvision
import numpy as np
import logging
import pprint

sys.path.append("../")
from data import SimulatorDataset

from spikingjelly.activation_based import neuron, encoding, functional, surrogate, layer
class ANN(nn.Module):
    def __init__(self, tau):
        super().__init__()

        self.layer = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28*28, 28*28 *2, bias=True),
            nn.ReLU(),
            nn.Linear(28*28 *2, 10, bias=False),
            )


    def forward(self, x: torch.Tensor):
        return self.layer(x)


class ANN2(nn.Module):
    def __init__(self, tau):
        super().__init__()

        self.layer = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28*28, 200, bias=True),
            nn.ReLU(),
            nn.Linear(200, 100, bias=False),
            nn.ReLU(),
            nn.Linear(100, 10, bias=False),
            )


    def forward(self, x: torch.Tensor):
        return self.layer(x)

class ANN3(nn.Module):
    def __init__(self, tau):
        super().__init__()

        self.layer = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28*28,10, bias=True)
            )


    def forward(self, x: torch.Tensor):
        return self.layer(x)

def create_logger(name,log_file, level=logging.INFO):

    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    logger = logging.getLogger(name)
    formatter = logging.Formatter(
        '[%(asctime)s][%(filename)15s][line:%(lineno)4d][%(levelname)8s] %(message)s')
    logger.setLevel(level = logging.INFO)
    
    handler = logging.FileHandler(log_file)
    handler.setLevel(logging.INFO)   
    handler.setFormatter(formatter)
    
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    console.setFormatter(formatter)
    
    logger.addHandler(handler)
    logger.addHandler(console)

    return logger

def main():

    
    parser = argparse.ArgumentParser(description='LIF MNIST Training')
    parser.add_argument('-T', default=10, type=int, help='simulating time-steps')
    parser.add_argument('-device', default='cuda:3', help='device')
    parser.add_argument('-b', default=6400, type=int, help='batch size')
    parser.add_argument('-epochs', default=100, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('-j', default=4, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
    parser.add_argument('-data-dir', type=str, help='root dir of MNIST dataset')
    parser.add_argument('-out-dir', type=str, default='../logs', help='root dir for saving logs and checkpoint')
    parser.add_argument('-resume', type=str, help='resume from the checkpoint path')
    parser.add_argument('-amp', action='store_true', help='automatic mixed precision training')
    parser.add_argument('-opt', type=str, choices=['sgd', 'adam'], default='adam', help='use which optimizer. SGD or Adam')
    parser.add_argument('-momentum', default=0.9, type=float, help='momentum for SGD')
    parser.add_argument('-lr', default=1e-3, type=float, help='learning rate')
    parser.add_argument('-tau', default=2.0, type=float, help='parameter tau of LIF neuron')
    parser.add_argument('-net_name', type=str, help='resume from the checkpoint path')


    args = parser.parse_args()

    out_dir = os.path.join(args.out_dir, f'{args.net_name}_CEloss_b{args.b}_{args.opt}_lr{args.lr}_epoch{args.epochs}')
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
        print(f'Mkdir {out_dir}.')

    logger = create_logger(__name__,os.path.join(out_dir, 'log.txt'))
    logger.info(f'---------------------- A NEW TESTING ----------------------')
    logger.info(f'config: {pprint.pformat(args)}')

    if args.net_name == 'ANN2':
        net = ANN2(tau=args.tau)
    elif args.net_name == 'ANN3':
        net = ANN3(tau=args.tau)
    else:
        net = ANN(tau=args.tau)

    logger.info(f'config: {pprint.pformat(net)}')

    net.to(args.device)

    test_dataset = torchvision.datasets.MNIST(
        root=args.data_dir,
        train=False,
        transform=torchvision.transforms.ToTensor(),
        download=False
    )
    test_data_loader = data.DataLoader(
        dataset=test_dataset,
        batch_size=args.b,
        shuffle=False,
        drop_last=False,
        num_workers=args.j,
        pin_memory=True
    )

    LossFunc = nn.CrossEntropyLoss()

    if args.resume:
        checkpoint = torch.load(args.resume, map_location='cpu')
        net.load_state_dict(checkpoint['net'])




    net.eval()
    test_loss = 0
    test_acc = 0
    test_samples = 0
    with torch.no_grad():
        for img, label in test_data_loader:
            img = img.to(args.device).float()
            label = label.to(args.device)
            label_onehot = F.one_hot(label, 10).float()
            out = net(img)
            out_softmax = F.softmax(out)
            loss = LossFunc(out, label_onehot)

            test_samples += label.numel()
            test_loss += loss.item() * label.numel()
            test_acc += (out_softmax.argmax(1) == label).float().sum().item()
    test_loss /= test_samples
    test_acc /= test_samples

    logger.info(f'test_loss ={test_loss: .4f}, test_acc ={test_acc: .4f}')
   
if __name__ == '__main__':
    main()