# Path: snntoolbox/bin/run.py
# def main(filepath=None):
#     """Entry point for running the toolbox.
# 
#     Note
#     ----
# 
#     There is no need to call this function directly, because python sets up an
#     executable during :ref:`installation` that can be called from terminal.
# 
#     """
#     from snntoolbox.bin.utils import update_setup, run_pipeline
# 
#     if filepath is not None:
#         config = update_setup(filepath)
#         run_pipeline(config)
#         return
# 
#     parser = argparse.ArgumentParser(
#         description='Run SNN toolbox to convert an analog neural network into '
#                     'a spiking neural network, and optionally simulate it.')
#     parser.add_argument('config_filepath', nargs='?',
#                         help='Path to configuration file.')
#     parser.add_argument('-t', '--terminal', action='store_true',
#                         help='Set this flag to run the toolbox from terminal. '
#                              'Omit this flag to open GUI.')
#     args = parser.parse_args()
# 
#     _filepath = os.path.abspath(args.config_filepath)
#     if _filepath is not None:
#         config = update_setup(_filepath)
# 
#         if args.terminal:
#             run_pipeline(config)
#         else:
#             from snntoolbox.bin.gui import gui
#             gui.main()
#     else:
#         if args.terminal:
#             parser.error("When using the SNN toolbox from terminal, a "
#                          "config_filepath argument must be provided.")
#             return
#         else:
#             from snntoolbox.bin.gui import gui
#             gui.main()
# 
# Path: snntoolbox/utils/utils.py
# def import_configparser():
#     try:
#         import configparser
#     except ImportError:
#         # noinspection PyPep8Naming
#         import ConfigParser as configparser
# 
#     return configparser
# 
# Path: tests/parsing/models/pytorch.py
# class Model(nn.Module):
#     def __init__(self):
#         super(Model, self).__init__()
# 
#         # The input_shape field is required by SNN toolbox.
#         self.input_shape = (1, 28, 28)
# 
#         layers_trunk = [
#             nn.Conv2d(1, 16, kernel_size=5, stride=2),
#             # BatchNorm doesn't work with Keras==2.3.1 because for some reason
#             # they put the batch-norm axis in a list.
#             # nn.BatchNorm2d(16),
#             nn.ReLU(),
#             nn.AvgPool2d(kernel_size=2, stride=2)]
#         layers_branch1 = [
#             nn.Conv2d(16, 32, kernel_size=3, padding=1),
#             nn.ReLU()]
#         layers_branch2 = [
#             nn.Conv2d(16, 8, kernel_size=1),
#             nn.ReLU()]
#         layers_head = [
#             nn.Conv2d(40, 8, kernel_size=1),
#             nn.ReLU()]
#         layers_classifier = [
#             nn.Dropout(1e-5),
#             nn.Linear(288, 10),
#             nn.Softmax(1)]
#         self.trunk = nn.Sequential(*layers_trunk)
#         self.branch1 = nn.Sequential(*layers_branch1)
#         self.branch2 = nn.Sequential(*layers_branch2)
#         self.head = nn.Sequential(*layers_head)
#         self.classifier = nn.Sequential(*layers_classifier)
# 
#     def forward(self, x):
#         x = self.trunk(x)
#         x1 = self.branch1(x)
#         x2 = self.branch2(x)
#         x = torch.cat([x1, x2], 1)
#         x = self.head(x)
#         x = x.view(-1, 288)  # Flatten
#         x = self.classifier(x)
#         return x


import os
import shutil
import inspect
import time
import numpy as np
import torch
import torch.nn as nn
from tensorflow.keras import backend
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
from snntoolbox.bin.run import main
from snntoolbox.utils.utils import import_configparser
from tests.parsing.models.pytorch import Model
class PytorchDataset(torch.utils.data.Dataset):
    def __init__(self, data, target, transform=None):
        self.data = torch.from_numpy(data).float()
        self.target = torch.from_numpy(target).long()
        self.transform = transform

    def __getitem__(self, index):
        x = self.data[index]

        if self.transform:
            x = self.transform(x)

        return x, self.target[index]

    def __len__(self):
        return len(self.data)


trainset = torch.utils.data.DataLoader(PytorchDataset(x_train, y_train),
                                       batch_size=64)
testset = torch.utils.data.DataLoader(PytorchDataset(x_test, y_test),
                                      batch_size=64)

# CREATE ANN #
##############

# This section creates a CNN using pytorch, and trains it with backpropagation.
# There are no spikes involved at this point.

# Create pytorch model from definition in separate script.

# Path: snntoolbox/bin/run.py
# def main(filepath=None):
#     """Entry point for running the toolbox.
# 
#     Note
#     ----
# 
#     There is no need to call this function directly, because python sets up an
#     executable during :ref:`installation` that can be called from terminal.
# 
#     """
#     from snntoolbox.bin.utils import update_setup, run_pipeline
# 
#     if filepath is not None:
#         config = update_setup(filepath)
#         run_pipeline(config)
#         return
# 
#     parser = argparse.ArgumentParser(
#         description='Run SNN toolbox to convert an analog neural network into '
#                     'a spiking neural network, and optionally simulate it.')
#     parser.add_argument('config_filepath', nargs='?',
#                         help='Path to configuration file.')
#     parser.add_argument('-t', '--terminal', action='store_true',
#                         help='Set this flag to run the toolbox from terminal. '
#                              'Omit this flag to open GUI.')
#     args = parser.parse_args()
# 
#     _filepath = os.path.abspath(args.config_filepath)
#     if _filepath is not None:
#         config = update_setup(_filepath)
# 
#         if args.terminal:
#             run_pipeline(config)
#         else:
#             from snntoolbox.bin.gui import gui
#             gui.main()
#     else:
#         if args.terminal:
#             parser.error("When using the SNN toolbox from terminal, a "
#                          "config_filepath argument must be provided.")
#             return
#         else:
#             from snntoolbox.bin.gui import gui
#             gui.main()
# 
# Path: snntoolbox/utils/utils.py
# def import_configparser():
#     try:
#         import configparser
#     except ImportError:
#         # noinspection PyPep8Naming
#         import ConfigParser as configparser
# 
#     return configparser
# 
# Path: tests/parsing/models/pytorch.py
# class Model(nn.Module):
#     def __init__(self):
#         super(Model, self).__init__()
# 
#         # The input_shape field is required by SNN toolbox.
#         self.input_shape = (1, 28, 28)
# 
#         layers_trunk = [
#             nn.Conv2d(1, 16, kernel_size=5, stride=2),
#             # BatchNorm doesn't work with Keras==2.3.1 because for some reason
#             # they put the batch-norm axis in a list.
#             # nn.BatchNorm2d(16),
#             nn.ReLU(),
#             nn.AvgPool2d(kernel_size=2, stride=2)]
#         layers_branch1 = [
#             nn.Conv2d(16, 32, kernel_size=3, padding=1),
#             nn.ReLU()]
#         layers_branch2 = [
#             nn.Conv2d(16, 8, kernel_size=1),
#             nn.ReLU()]
#         layers_head = [
#             nn.Conv2d(40, 8, kernel_size=1),
#             nn.ReLU()]
#         layers_classifier = [
#             nn.Dropout(1e-5),
#             nn.Linear(288, 10),
#             nn.Softmax(1)]
#         self.trunk = nn.Sequential(*layers_trunk)
#         self.branch1 = nn.Sequential(*layers_branch1)
#         self.branch2 = nn.Sequential(*layers_branch2)
#         self.head = nn.Sequential(*layers_head)
#         self.classifier = nn.Sequential(*layers_classifier)
# 
#     def forward(self, x):
#         x = self.trunk(x)
#         x1 = self.branch1(x)
#         x2 = self.branch2(x)
#         x = torch.cat([x1, x2], 1)
#         x = self.head(x)
#         x = x.view(-1, 288)  # Flatten
#         x = self.classifier(x)
#         return x


# Path: examples/mnist_pytorch_INI.py
import os
import shutil
import inspect
import time
import numpy as np
import torch
import torch.nn as nn
from tensorflow.keras import backend
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
from snntoolbox.bin.run import main
from snntoolbox.utils.utils import import_configparser
from tests.parsing.models.pytorch import Model

class PytorchDataset(torch.utils.data.Dataset):
    def __init__(self, data, target, transform=None):
        self.data = torch.from_numpy(data).float()
        self.target = torch.from_numpy(target).long()
        self.transform = transform

    def __getitem__(self, index):
        x = self.data[index]

        if self.transform:
            x = self.transform(x)

        return x, self.target[index]

    def __len__(self):
        return len(self.data)


trainset = torch.utils.data.DataLoader(PytorchDataset(x_train, y_train),
                                       batch_size=64)
testset = torch.utils.data.DataLoader(PytorchDataset(x_test, y_test),
                                      batch_size=64)

# CREATE ANN #
##############

# This section creates a CNN using pytorch, and trains it with backpropagation.
# There are no spikes involved at this point.

# Create pytorch model from definition in separate script.

model = Model()