from typing import Any, Dict
import argparse

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

FC1_DIM = 1024
FC2_DIM = 128


class MLP(nn.Module):
    """Simple MLP suitable for recognizing single characters."""

    def __init__(
        self,
        data_config: Dict[str, Any],
        args: argparse.Namespace = None,
    ) -> None:
        super().__init__()
        self.args = vars(args) if args is not None else {}

        input_dim = np.prod(data_config["input_dims"])
        num_classes = len(data_config["mapping"])

        self.dropout = nn.Dropout(0.5)
        layers = self.args.get("layers", FC1_DIM)
        self.layers = layers
        if layers:
            fcn = (int(FC1_DIM - x * ((FC1_DIM - FC2_DIM)//(layers-1))) for x in range(layers))
            fcl = input_dim
            fcv = []
            for fci in fcn:
                fcv.append(nn.Linear(fcl, fci))
                fcl = fci
            fcv.append(nn.Linear(fcl, num_classes))
            self.fcv = nn.Sequential(*fcv)
        else:
            fc1_dim = self.args.get("fc1", FC1_DIM)
            fc2_dim = self.args.get("fc2", FC2_DIM)

            self.fc1 = nn.Linear(input_dim, fc1_dim)
            self.fc2 = nn.Linear(fc1_dim, fc2_dim)
            self.fc3 = nn.Linear(fc2_dim, num_classes)

    def forward(self, x):
        x = torch.flatten(x, 1)
        if self.layers:
            for fci in self.fcv[:-1]:
                x = fci(x)
                x = F.relu(x)
                x = self.dropout(x)
            x = self.fcv[-1](x)
        else:
            x = self.fc1(x)
            x = F.relu(x)
            x = self.dropout(x)
            x = self.fc2(x)
            x = F.relu(x)
            x = self.dropout(x)
            x = self.fc3(x)
        return x

    @staticmethod
    def add_to_argparse(parser):
        parser.add_argument("--layers", type=int, default=None, choices=range(2, 20))
        parser.add_argument("--fc1", type=int, default=1024)
        parser.add_argument("--fc2", type=int, default=128)
        return parser

[((24, 20, 24, 54), 'numpy.prod', 'np.prod', ({(24, 28, 24, 53): "data_config['input_dims']"}, {}), "(data_config['input_dims'])", True, 'import numpy as np\n'), ((27, 23, 27, 38), 'torch.nn.Dropout', 'nn.Dropout', ({(27, 34, 27, 37): '0.5'}, {}), '(0.5)', True, 'import torch.nn as nn\n'), ((48, 12, 48, 31), 'torch.flatten', 'torch.flatten', ({(48, 26, 48, 27): 'x', (48, 29, 48, 30): '1'}, {}), '(x, 1)', False, 'import torch\n'), ((38, 23, 38, 42), 'torch.nn.Sequential', 'nn.Sequential', ({(38, 37, 38, 41): '*fcv'}, {}), '(*fcv)', True, 'import torch.nn as nn\n'), ((43, 23, 43, 52), 'torch.nn.Linear', 'nn.Linear', ({(43, 33, 43, 42): 'input_dim', (43, 44, 43, 51): 'fc1_dim'}, {}), '(input_dim, fc1_dim)', True, 'import torch.nn as nn\n'), ((44, 23, 44, 50), 'torch.nn.Linear', 'nn.Linear', ({(44, 33, 44, 40): 'fc1_dim', (44, 42, 44, 49): 'fc2_dim'}, {}), '(fc1_dim, fc2_dim)', True, 'import torch.nn as nn\n'), ((45, 23, 45, 54), 'torch.nn.Linear', 'nn.Linear', ({(45, 33, 45, 40): 'fc2_dim', (45, 42, 45, 53): 'num_classes'}, {}), '(fc2_dim, num_classes)', True, 'import torch.nn as nn\n'), ((57, 16, 57, 25), 'torch.nn.functional.relu', 'F.relu', ({(57, 23, 57, 24): 'x'}, {}), '(x)', True, 'import torch.nn.functional as F\n'), ((60, 16, 60, 25), 'torch.nn.functional.relu', 'F.relu', ({(60, 23, 60, 24): 'x'}, {}), '(x)', True, 'import torch.nn.functional as F\n'), ((37, 23, 37, 50), 'torch.nn.Linear', 'nn.Linear', ({(37, 33, 37, 36): 'fcl', (37, 38, 37, 49): 'num_classes'}, {}), '(fcl, num_classes)', True, 'import torch.nn as nn\n'), ((52, 20, 52, 29), 'torch.nn.functional.relu', 'F.relu', ({(52, 27, 52, 28): 'x'}, {}), '(x)', True, 'import torch.nn.functional as F\n'), ((35, 27, 35, 46), 'torch.nn.Linear', 'nn.Linear', ({(35, 37, 35, 40): 'fcl', (35, 42, 35, 45): 'fci'}, {}), '(fcl, fci)', True, 'import torch.nn as nn\n')]