import numpy as np

from pynet import nn
from .Net import Net

__all__ = ['MyAlexNet']

model_urls = {
    'myalexnet': ''
}

class MyAlexNet(Net):
    def __init__(self, dropout=0.5, weight_scale=0.01, cut_percent=0.9, cut_weight=False, cut_weight_mode='same',
                 cut_grad=False):
        super(MyAlexNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 3, 3, 96, padding=1, weight_scale=weight_scale, cut_weight=cut_weight,
                               cut_percent=cut_percent, cut_weight_mode=cut_weight_mode)
        self.conv2 = nn.Conv2d(96, 3, 3, 256, padding=1, weight_scale=weight_scale, cut_weight=cut_weight,
                               cut_percent=cut_percent, cut_weight_mode=cut_weight_mode)
        self.conv3 = nn.Conv2d(256, 3, 3, 384, padding=1, weight_scale=weight_scale, cut_weight=cut_weight,
                               cut_percent=cut_percent, cut_weight_mode=cut_weight_mode)
        self.conv4 = nn.Conv2d(384, 3, 3, 384, padding=1, weight_scale=weight_scale, cut_weight=cut_weight,
                               cut_percent=cut_percent, cut_weight_mode=cut_weight_mode)
        self.conv5 = nn.Conv2d(384, 3, 3, 256, padding=1, weight_scale=weight_scale, cut_weight=cut_weight,
                               cut_percent=cut_percent, cut_weight_mode=cut_weight_mode)

        self.maxPool1 = nn.MaxPool(2, 2, 96, stride=2)
        self.maxPool2 = nn.MaxPool(2, 2, 256, stride=2)
        self.maxPool3 = nn.MaxPool(2, 2, 256, stride=2)

        self.fc1 = nn.FC(4096, 2048, weight_loc=0., weight_scale=weight_scale, bias_value=0.)
        self.fc2 = nn.FC(2048, 2048, weight_loc=0., weight_scale=weight_scale, bias_value=0.)
        self.fc3 = nn.FC(2048, 10, weight_loc=0., weight_scale=weight_scale, bias_value=0.)

        self.flat = nn.Flat((256, 4, 4), 4096)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout()
        self.bn1 = nn.BN()
        self.bn2 = nn.BN()
        self.prune = nn.Prune(cut_grad=cut_grad, cut_percent=cut_percent)

        self.other_params = {'dropout_param': {'mode': 'train', 'p': dropout},
                             'bn1': {'mode': 'train'},
                             'bn2': {'mode': 'train'},
                             }
        self.params = self._get_params()

        self.layers = []
        self._add_layers([(self.prune, None, None),
                          (self.conv1, ('W1', 'b1'), None),
                          (self.relu, None, None),
                          (self.maxPool1, None, None),
                          (self.prune, None, None),
                          (self.conv2, ('W2', 'b2'), None),
                          (self.relu, None, None),
                          (self.maxPool2, None, None),
                          (self.prune, None, None),
                          (self.conv3, ('W3', 'b3'), None),
                          (self.relu, None, None),
                          (self.prune, None, None),
                          (self.conv4, ('W4', 'b4'), None),
                          (self.relu, None, None),
                          (self.prune, None, None),
                          (self.conv5, ('W5', 'b5'), None),
                          (self.relu, None, None),
                          (self.maxPool3, None, None),
                          (self.prune, None, None),
                          (self.flat, None, None),
                          (self.dropout, None, 'dropout_param'),
                          (self.fc1, ('W6', 'b6'), None),
                          (self.prune, None, None),
                          (self.bn1, ('gamma1', 'beta1'), 'bn1'),
                          (self.relu, None, None),
                          (self.dropout, None, 'dropout_param'),
                          (self.fc2, ('W7', 'b7'), None),
                          (self.prune, None, None),
                          (self.bn2, ('gamma2', 'beta2'), 'bn2'),
                          (self.relu, None, None),
                          (self.prune, None, None),
                          (self.fc3, ('W8', 'b8'), None)
                          ])

    def __call__(self, inputs):
        return self.forward(inputs)

    def forward(self, inputs):
        # inputs.shape = [N, 3, 32, 32]
        assert len(inputs.shape) == 4
        assert inputs.shape[1:] == (3, 32, 32)

        feature = inputs
        for layer in self.layers:
            # print(layer['grad_keys'])
            params = None
            if layer['params_key'] is not None:
                params = [self.params[k] for k in layer['params_key']]

            other = None
            if layer['other_key'] is not None:
                other = self.other_params[layer['other_key']]

            feature, layer['cache'] = layer['layer'](feature, params, other)

        return feature

    def backward(self, grad_out):
        # grad_out.shape = [N, 10]
        assert len(grad_out.shape) == 2
        assert grad_out.shape[1] == 10

        grads = dict()
        feature_grad = grad_out
        for layer in reversed(self.layers):
            backward_out = layer['layer'].backward(feature_grad, layer['cache'])
            if layer['params_key'] is not None:
                feature_grad = backward_out[0]
                assert len(layer['params_key']) == len(backward_out) - 1
                for i in range(len(layer['params_key'])):
                    grads[layer['params_key'][i]] = backward_out[i + 1]
            else:
                feature_grad = backward_out

        return grads

    def _get_params(self):
        params = dict()
        params['W1'], params['b1'] = self.conv1.get_params()
        params['W2'], params['b2'] = self.conv2.get_params()
        params['W3'], params['b3'] = self.conv3.get_params()
        params['W4'], params['b4'] = self.conv4.get_params()
        params['W5'], params['b5'] = self.conv5.get_params()

        params['W6'], params['b6'] = self.fc1.get_params()
        params['W7'], params['b7'] = self.fc2.get_params()
        params['W8'], params['b8'] = self.fc3.get_params()

        params['gamma1'], params['beta1'] = self.bn1.get_params(2048)
        params['gamma2'], params['beta2'] = self.bn2.get_params(2048)

        return params

    def _add_layers(self, layer_list):
        for layer, params_key, other_key in layer_list:
            self.layers.append({
                'layer': layer,
                'params_key': params_key,
                'cache': None,
                'other_key': other_key
            })

    def train(self):
        self.conv1.train()
        self.conv2.train()
        self.conv3.train()
        self.conv4.train()
        self.conv5.train()
        for v in self.other_params.values():
            v['mode'] = 'train'

    def eval(self):
        self.conv1.eval()
        self.conv2.eval()
        self.conv3.eval()
        self.conv4.eval()
        self.conv5.eval()
        for v in self.other_params.values():
            v['mode'] = 'test'


"""
Module for handling atomic elements.
"""
from numpy import vectorize

class Element:
    """
    Class for storing and accessing atomic elements.

    Args:
        input_value: The atomic number, symbol, or name of the element
    """

    def __init__(self, input_value):
        self.input = input_value

        # list with atomic number z, short name, full name, valence,
        # valence electrons, covalent radius, vdW radius, metallic radius
        self.elements_list = [
            (1, "H", "Hydrogen", 1.0, 1, 0.31, 1.20, None),
            (2, "He", "Helium", 0.5, 2, 0.28, 1.40, None),
            (3, "Li", "Lithium", 1.0, 1, 1.28, 1.82, 1.52),
            (4, "Be", "Beryllium", 2.0, 2, 0.96, 1.53, 1.12),
            (5, "B", "Boron", 3.0, 3, 0.84, 1.92, None),
            (6, "C", "Carbon", 4.0, 4, 0.78, 1.70, None),
            (7, "N", "Nitrogen", 3.0, 5, 0.78, 1.55, None),
            (8, "O", "Oxygen", 2.0, 6, 0.70, 1.52, None),
            (9, "F", "Fluorine", 1.0, 7, 0.57, 1.47, None),
            (10, "Ne", "Neon", 0.5, 8, 0.58, 1.54, None),
            (11, "Na", "Sodium", 1.0, 1, 1.66, 2.27, 1.86),
            (12, "Mg", "Magnesium", 2.0, 2, 1.41, 1.73, 1.60),
            (13, "Al", "Aluminium", 3.0, 3, 1.21, 1.84, 1.43),
            (14, "Si", "Silicon", 4.0, 4, 1.11, 2.10, None),
            (15, "P", "Phosphorus", 3.0, 5, 1.07, 1.80, None),
            (16, "S", "Sulfur", 2.0, 6, 1.25, 1.80, None),
            (17, "Cl", "Chlorine", 1.0, 7, 1.02, 1.75, None),
            (18, "Ar", "Argon", 0.5, 8, 1.06, 1.88, None),
            (19, "K", "Potassium", 1.0, 1, 2.03, 2.75, 2.27),
            (20, "Ca", "Calcium", 2.0, 2, 1.76, 2.31, 1.97),
            (21, "Sc", "Scandium", 3.0, 3, 1.70, 2.11, 1.62),
            (22, "Ti", "Titanium", 4.0, 4, 1.60, 2.00, 1.47),
            (23, "V", "Vanadium", 4.0, 5, 1.53, 2.00, 1.34),
            (24, "Cr", "Chromium", 3.0, 6, 1.39, 2.00, 1.28),
            (25, "Mn", "Manganese", 4.0, 5, 1.39, 2.00, 1.27),
            (26, "Fe", "Iron", 3.0, 3, 1.32, 2.00, 1.26),
            (27, "Co", "Cobalt", 3.0, 3, 1.26, 2.00, 1.25),
            (28, "Ni", "Nickel", 2.0, 3, 1.24, 1.63, 1.24),
            (29, "Cu", "Copper", 2.0, 2, 1.32, 1.40, 1.28),
            (30, "Zn", "Zinc", 2.0, 2, 1.22, 1.39, 1.34),
            (31, "Ga", "Gallium", 3.0, 3, 1.22, 1.87, 1.35),
            (32, "Ge", "Germanium", 4.0, 4, 1.20, 2.11, None),
            (33, "As", "Arsenic", 3.0, 5, 1.19, 1.85, None),
            (34, "Se", "Selenium", 2.0, 6, 1.20, 1.90, None),
            (35, "Br", "Bromine", 1.0, 7, 1.20, 1.85, None),
            (36, "Kr", "Krypton", 0.5, 8, 1.16, 2.02, None),
            (37, "Rb", "Rubidium", 1.0, 1, 2.20, 3.03, 2.48),
            (38, "Sr", "Strontium", 2.0, 2, 1.95, 2.49, 2.15),
            (39, "Y", "Yttrium", 3.0, 3, 1.90, 2.00, 1.80),
            (40, "Zr", "Zirconium", 4.0, 4, 1.75, 2.00, 1.60),
            (41, "Nb", "Niobium", 5.0, 5, 1.64, 2.00, 1.46),
            (42, "Mo", "Molybdenum", 4.0, 6, 1.54, 2.00, 1.39),
            (43, "Tc", "Technetium", 4.0, 5, 1.47, 2.00, 1.36),
            (44, "Ru", "Ruthenium", 4.0, 3, 1.46, 2.00, 1.34),
            (45, "Rh", "Rhodium", 4.0, 3, 1.42, 1.63, 1.34),
            (46, "Pd", "Palladium", 4.0, 3, 1.39, 1.72, 1.37),
            (47, "Ag", "Silver", 1.0, 2, 1.45, 1.58, 1.44),
            (48, "Cd", "Cadmium", 2.0, 2, 1.44, 1.93, 1.51),
            (49, "In", "Indium", 3.0, 3, 1.42, 2.17, 1.67),
            (50, "Sn", "Tin", 4.0, 4, 1.39, 2.06, None),
            (51, "Sb", "Antimony", 3.0, 5, 1.39, 2.06, None),
            (52, "Te", "Tellurium", 2.0, 6, 1.38, 2.06, None),
            (53, "I", "Iodine", 1.0, 7, 1.39, 1.98, None),
            (54, "Xe", "Xenon", 0.5, 8, 1.40, 2.16, None),
            (55, "Cs", "Caesium", 1.0, 1, 2.44, 3.43, 2.65),
            (56, "Ba", "Barium", 2.0, 2, 2.15, 2.68, 2.22),
            (57, "La", "Lanthanum", 3.0, 3, 2.07, 2.10, 1.87),
            (58, "Ce", "Cerium", 4.0, 3, 2.04, 2.10, 1.818),
            (59, "Pr", "Praseodymium", 3.0, 3, 2.03, 2.10, 1.824),
            (60, "Nd", "Neodymium", 3.0, 3, 2.01, 2.10, 1.814),
            (61, "Pm", "Promethium", 3.0, 3, 1.99, 2.10, 1.834),
            (62, "Sm", "Samarium", 3.0, 3, 1.98, 2.10, 1.804),
            (63, "Eu", "Europium", 3.0, 3, 1.98, 2.10, 1.804),
            (64, "Gd", "Gadolinium", 3.0, 3, 1.96, 2.10, 1.804),
            (65, "Tb", "Terbium", 3.0, 3, 1.94, 2.10, 1.773),
            (66, "Dy", "Dysprosium", 3.0, 3, 1.92, 2.10, 1.781),
            (67, "Ho", "Holmium", 3.0, 3, 1.92, 2.10, 1.762),
            (68, "Er", "Erbium", 3.0, 3, 1.89, 2.10, 1.761),
            (69, "Tm", "Thulium", 3.0, 3, 1.90, 2.10, 1.759),
            (70, "Yb", "Ytterbium", 3.0, 3, 1.87, 2.10, 1.76),
            (71, "Lu", "Lutetium", 3.0, 3, 1.87, 2.10, 1.738),
            (72, "Hf", "Hafnium", 4.0, 3, 1.75, 2.10, 1.59),
            (73, "Ta", "Tantalum", 5.0, 3, 1.70, 2.10, 1.46),
            (74, "W", "Tungsten", 4.0, 3, 1.62, 2.10, 1.39),
            (75, "Re", "Rhenium", 4.0, 3, 1.51, 2.10, 1.37),
            (76, "Os", "Osmium", 4.0, 3, 1.44, 2.10, 1.35),
            (77, "Ir", "Iridium", 4.0, 3, 1.41, 2.10, 1.355),
            (78, "Pt", "Platinum", 4.0, 3, 1.36, 1.75, 1.385),
            (79, "Au", "Gold", 1.0, 3, 1.36, 1.66, 1.44),
            (80, "Hg", "Mercury", 2.0, 3, 1.32, 1.55, 1.51),
            (81, "Tl", "Thallium", 3.0, 3, 1.45, 1.96, 1.70),
            (82, "Pb", "Lead", 4.0, 4, 1.46, 2.02, None),
            (83, "Bi", "Bismuth", 3.0, 5, 1.48, 2.07, None),
            (84, "Po", "Polonium", 2.0, 6, 1.40, 1.97, None),
            (85, "At", "Astatine", 1.0, 7, 1.50, 2.02, None),
            (86, "Rn", "Radon", 0.5, 8, 1.50, 2.20, None),
            (87, "Fr", "Francium", 1.0, 1, 2.60, 3.48, None),
            (88, "Ra", "Radium", 2.0, 2, 2.21, 2.83, None),
            (89, "Ac", "Actinium", 3.0, 3, 2.15, 2.20, None),
            (90, "Th", "Thorium", 4.0, 3, 2.06, 2.20, 1.79),
            (91, "Pa", "Protactinium", 4.0, 3, 2.00, 2.20, 1.63),
            (92, "U", "Uranium", 4.0, 3, 1.96, 2.20, 1.56),
            (93, "Np", "Neptunium", 4.0, 3, 1.90, 2.20, 1.55),
            (94, "Pu", "Plutonium", 4.0, 3, 1.87, 2.20, 1.59),
            (95, "Am", "Americium", 4.0, 3, 1.80, 2.20, 1.73),
            (96, "Cm", "Curium", 4.0, 3, 1.69, 2.20, 1.74),
            (97, "Bk", "Berkelium", 4.0, 3, None, None, 1.70),
            (98, "Cf", "Californium", 4.0, 3, None, None, 1.86),
            (99, "Es", "Einsteinium", 4.0, 3, None, None, 1.86),
            (100, "Fm", "Fermium", 4.0, 3, None, None, None),
            (101, "Md", "Mendelevium", 4.0, 3, None, None, None),
            (102, "No", "Nobelium", 4.0, 3, None, None, None),
            (103, "Lr", "Lawrencium", 4.0, 3, None, None, None),
            (104, "Rf", "Rutherfordium", 4.0, 3, None, None, None),
            (105, "Db", "Dubnium", 2.0, 3, None, None, None),
        ]
        """A list of atomic numbers, symbols, names, and other information, up
        to atomic number 105"""

        # scatter factor
        self.sf = [
            [0.493, 0.323, 0.140, 0.041, 10.511, 26.126, 3.142, 57.800, 0.003],
            [0.873, 0.631, 0.311, 0.178, 9.104, 3.357, 22.928, 0.982, 0.006],
            [1.128, 0.751, 0.618, 0.465, 3.955, 1.052, 85.391, 168.261, 0.038],
            [1.592, 1.128, 0.539, 0.703, 43.643, 1.862, 103.483, 0.542, 0.038],
            [2.055, 1.333, 1.098, 0.707, 23.219, 1.021, 60.350, 0.140, -0.193],
            [2.310, 1.020, 1.589, 0.865, 20.844, 10.208, 0.569, 51.651, 0.216],
            [12.213, 3.132, 2.013, 1.166, 0.006, 9.893, 28.997, 0.583, -11.529],
            [3.049, 2.287, 1.546, 0.867, 13.277, 5.701, 0.324, 32.909, 0.251],
            [3.539, 2.641, 1.517, 1.024, 10.283, 4.294, 0.262, 26.148, 0.278],
            [3.955, 3.112, 1.455, 1.125, 8.404, 3.426, 0.231, 21.718, 0.352],
            [4.763, 3.174, 1.267, 1.113, 3.285, 8.842, 0.314, 129.424, 0.676],
            [5.420, 2.174, 1.227, 2.307, 2.828, 79.261, 0.381, 7.194, 0.858],
            [6.420, 1.900, 1.594, 1.965, 3.039, 0.743, 31.547, 85.089, 1.115],
            [6.292, 3.035, 1.989, 1.541, 2.439, 32.334, 0.678, 81.694, 1.141],
            [6.435, 4.179, 1.780, 1.491, 1.907, 27.157, 0.526, 68.164, 1.115],
            [6.905, 5.203, 1.438, 1.586, 1.468, 22.215, 0.254, 56.172, 0.867],
            [11.460, 7.196, 6.256, 1.645, 0.010, 1.166, 18.519, 47.778, -9.557],
            [7.484, 6.772, 0.654, 1.644, 0.907, 14.841, 43.898, 33.393, 1.444],
            [8.219, 7.440, 1.052, 0.866, 12.795, 0.775, 213.187, 41.684, 1.423],
            [8.627, 7.387, 1.590, 1.021, 10.442, 0.660, 85.748, 178.437, 1.375],
            [9.189, 7.368, 1.641, 1.468, 9.021, 0.573, 136.108, 51.353, 1.333],
            [9.759, 7.356, 1.699, 1.902, 7.851, 0.500, 35.634, 116.105, 1.281],
            [10.297, 7.351, 2.070, 2.057, 6.866, 0.438, 26.894, 102.478, 1.220],
            [10.641, 7.354, 3.324, 1.492, 6.104, 0.392, 20.263, 98.740, 1.183],
            [11.282, 7.357, 3.019, 2.244, 5.341, 0.343, 17.867, 83.754, 1.090],
            [11.769, 7.357, 3.522, 2.305, 4.761, 0.307, 15.354, 76.881, 1.037],
            [12.284, 7.341, 4.003, 2.349, 4.279, 0.278, 13.536, 71.169, 1.012],
            [12.838, 7.292, 4.444, 2.380, 3.878, 0.257, 12.176, 66.342, 1.034],
            [13.338, 7.168, 5.616, 1.673, 3.583, 0.247, 11.397, 64.831, 1.191],
            [14.074, 7.032, 5.165, 2.410, 3.266, 0.233, 10.316, 58.710, 1.304],
            [15.235, 6.701, 4.359, 2.962, 3.067, 0.241, 10.781, 61.414, 1.719],
            [16.082, 6.375, 3.707, 3.683, 2.851, 0.252, 11.447, 54.763, 2.131],
            [16.672, 6.070, 3.431, 4.278, 2.635, 0.265, 12.948, 47.797, 2.531],
            [17.001, 5.820, 3.973, 4.354, 2.410, 0.273, 15.237, 43.816, 2.841],
            [17.179, 5.236, 5.638, 3.985, 2.172, 16.580, 0.261, 41.433, 2.956],
            [17.355, 6.729, 5.549, 3.537, 1.938, 16.562, 0.226, 39.397, 2.825],
            [17.178, 9.644, 5.140, 1.529, 1.789, 17.315, 0.275, 164.934, 3.487],
            [17.566, 9.818, 5.422, 2.669, 1.556, 14.099, 0.166, 132.376, 2.506],
            [17.776, 10.295, 5.726, 3.266, 1.403, 12.801, 0.261, 104.354, 1.912],
            [17.876, 10.948, 5.417, 3.657, 1.276, 11.916, 0.118, 87.663, 2.069],
            [17.614, 12.014, 4.042, 3.533, 1.189, 11.766, 0.205, 69.796, 3.756],
            [3.703, 17.236, 12.888, 3.743, 0.277, 1.096, 11.004, 61.658, 4.387],
            [19.130, 11.095, 4.649, 2.713, 0.864, 8.145, 21.571, 86.847, 5.404],
            [19.267, 12.918, 4.863, 1.568, 0.809, 8.435, 24.800, 94.293, 5.379],
            [19.296, 14.350, 4.734, 1.289, 0.752, 8.218, 25.875, 98.606, 5.328],
            [19.332, 15.502, 5.295, 0.606, 0.699, 7.989, 25.205, 76.899, 5.266],
            [19.281, 16.688, 4.805, 1.046, 0.645, 7.473, 24.660, 99.816, 5.179],
            [19.221, 17.644, 4.461, 1.603, 0.595, 6.909, 24.701, 87.482, 5.069],
            [19.162, 18.560, 4.295, 2.040, 0.548, 6.378, 25.850, 92.803, 4.939],
            [19.189, 19.101, 4.458, 2.466, 5.830, 0.503, 26.891, 83.957, 4.782],
            [19.642, 19.045, 5.037, 2.683, 5.303, 0.461, 27.907, 75.283, 4.591],
            [19.964, 19.014, 6.145, 2.524, 4.817, 0.421, 28.528, 70.840, 4.352],
            [20.147, 18.995, 7.514, 2.273, 4.347, 0.381, 27.766, 66.878, 4.071],
            [20.293, 19.030, 8.977, 1.990, 3.928, 0.344, 26.466, 64.266, 3.712],
            [20.389, 19.106, 10.662, 1.495, 3.569, 0.311, 24.388, 213.904, 3.335],
            [20.336, 19.297, 10.888, 2.696, 3.216, 0.276, 20.207, 167.202, 2.773],
            [20.578, 19.599, 11.373, 3.287, 2.948, 0.244, 18.773, 133.124, 2.147],
            [21.167, 19.770, 11.851, 3.330, 2.812, 0.227, 17.608, 127.113, 1.863],
            [22.044, 19.670, 12.386, 2.824, 2.774, 0.222, 16.767, 143.644, 2.058],
            [22.684, 19.685, 12.774, 2.851, 2.662, 0.211, 15.885, 137.903, 1.985],
            [23.340, 19.610, 13.123, 2.875, 2.563, 0.202, 15.101, 132.721, 2.029],
            [24.004, 19.426, 13.440, 2.896, 2.473, 0.196, 14.400, 128.007, 2.210],
            [24.627, 19.089, 13.760, 2.293, 2.388, 0.194, 13.755, 123.174, 2.575],
            [25.071, 19.080, 13.852, 3.545, 2.253, 0.182, 12.933, 101.398, 2.420],
            [25.898, 18.219, 14.317, 2.954, 2.243, 0.196, 12.665, 115.362, 3.583],
            [26.507, 17.638, 14.560, 2.966, 2.180, 0.202, 12.190, 111.874, 4.297],
            [26.905, 17.294, 14.558, 3.638, 2.071, 0.198, 11.441, 92.657, 4.568],
            [27.656, 16.428, 14.978, 2.982, 2.074, 0.224, 11.360, 105.703, 5.920],
            [28.182, 15.885, 15.154, 2.987, 2.029, 0.239, 10.998, 102.961, 6.756],
            [28.664, 15.434, 15.309, 2.990, 1.989, 0.257, 10.665, 100.417, 7.567],
            [28.948, 15.221, 15.100, 3.716, 1.902, 9.985, 0.261, 84.330, 7.976],
            [29.144, 15.173, 14.759, 4.300, 1.833, 9.600, 0.275, 72.029, 8.582],
            [29.202, 15.229, 14.514, 4.765, 1.773, 9.370, 0.296, 63.364, 9.244],
            [0.000, 0.000, 0.000, 0.000, 1.722, 9.231, 0.323, 57.725, 9.858],
            [28.762, 15.719, 14.556, 5.442, 1.672, 9.092, 0.350, 52.086, 10.472],
            [28.189, 16.155, 14.931, 5.676, 1.629, 8.979, 0.383, 48.165, 11.000],
            [27.305, 16.730, 15.611, 5.834, 1.593, 8.866, 0.418, 45.001, 11.472],
            [27.006, 17.764, 15.713, 5.784, 1.513, 8.812, 0.425, 38.610, 11.688],
            [16.882, 18.591, 25.558, 5.860, 0.461, 8.622, 1.483, 36.396, 12.066],
            [20.681, 19.042, 21.657, 5.968, 0.545, 8.448, 1.573, 38.325, 12.609],
            [27.545, 19.158, 15.538, 5.526, 0.655, 8.708, 1.963, 45.815, 13.175],
            [31.062, 13.064, 18.442, 5.970, 0.690, 2.358, 8.618, 47.258, 13.412],
            [33.369, 12.951, 16.588, 6.469, 0.704, 2.924, 8.794, 48.009, 13.578],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [35.660, 23.103, 12.598, 4.087, 0.589, 3.652, 18.599, 117.020, 13.527],
            [35.564, 23.422, 12.747, 4.807, 0.563, 3.462, 17.831, 99.172, 13.431],
            [35.885, 23.295, 14.189, 4.173, 0.548, 3.415, 16.924, 105.251, 13.429],
            [0.000, 0.000, 0.000, 0.000, 0.530, 3.335, 16.143, 101.371, 13.393],
            [36.187, 23.596, 15.640, 4.186, 0.512, 3.254, 15.362, 97.491, 13.357],
            [36.526, 23.808, 16.771, 3.479, 0.499, 3.264, 14.946, 105.980, 13.381],
        ]
        """A list of scatter factors for the elements"""

        self.z = None
        """atomic number"""
        self.short_name = None
        """atomic symbol"""
        self.long_name = None
        """atomic name"""
        self.valence = None
        """valence value"""
        self.valence_electrons = None
        """number of valence electrons"""
        self.covalent_radius = None
        """atomic radius used for distance checking within crystals"""
        self.vdw_radius = None
        """atomic radius used for volume estimation within crystals"""
        self.metallic_radius = None
        """atomic radius used for distance checking within metallic crystals"""

        pos = None

        try:
            int(self.input)
            self.z = self.input

            for i, el in enumerate(self.elements_list):
                if el[0] == self.z:
                    pos = i
                    self.short_name = el[1]
                    self.long_name = el[2]
                    break
        except ValueError:
            self.short_name = self.input
            for i, el in enumerate(self.elements_list):
                if el[1] == self.short_name:
                    pos = i
                    self.z = el[0]
                    self.long_name = el[2]
                    break

            if not self.z:
                self.short_name = None
                self.long_name = self.input
                for i, el in enumerate(self.elements_list):
                    if el[2] == self.long_name:
                        pos = i
                        self.z = el[0]
                        self.short_name = el[1]
                        break
                if not self.z:
                    self.long_name = None

        if pos is not None:
            self.valence = self.elements_list[pos][3]
            self.valence_electrons = self.elements_list[pos][4]
            self.covalent_radius = self.elements_list[pos][5]
            self.vdw_radius = self.elements_list[pos][6]
            self.metallic_radius = self.elements_list[pos][7]
            self.scatter = self.sf[pos]

    def get_all(self, pos):
        """
        Return all [pos] elements in the full element list

        Args:
            pos: the index of the elements to retrieve

        Returns:
            a list containing only the [pos] elements of self.elements_list
        """
        els = []
        for el in self.elements_list:
            els.append(el[pos])
        return els

    def get_sf(self, pos):
        """
        Get the scattering factor for an element.

        Args:
            pos: the atomic number of the element

        Returns:
            the scattering factor of the element
        """
        els = []
        for el in self.sf:
            els.append(el[pos])
        return els

    # TODO: add docstrings for miscellaneous functions
    def all_z(self):
        return self.get_all(0)

    def all_short_names(self):
        return self.get_all(1)

    def all_long_names(self):
        return self.get_all(2)

    def all_valences(self):
        return self.get_all(3)

    def all_valence_electrons(self):
        return self.get_all(4)

    def all_covalent_radii(self):
        return self.get_all(5)

    def all_vdw_radii(self):
        return self.get_all(6)

    def all_metallic_radii(self):
        return self.get_all(7)

    def get_sf(self):
        return self.get_sf()

    def number_from_specie(specie):
        if type(specie) == int or type(specie) == float:
            if specie <= 105 and specie >= 1:
                index = int(specie)
            else:
                print(specie)
                print("Error: Atomic number must be between 1 and 105.")
                return
        elif type(specie) == str:
            try:
                el = Element(specie)
                index = el.z
            except:
                print("Error: Invalid atomic symbol, name, or number.")
                return
        elif type(specie) == Element:
            try:
                index = specie.z
            except:
                print("Error: Element object has no atomic number 'z'.")
                return
        else:
            try:
                el = Element(specie.number)
                index = el.z
            except:
                print("Error: Invalid atomic symbol, name, or number.")
                return
        return index


# Copyright (c) 2019 Eric Steinberger 2020 Vsevolod Kompantsev

import ctypes
import os
from os.path import join as ospj

import itertools
import numpy as np

from PokerRL._.CppWrapper import CppWrapper
from PokerRL.game._.rl_env.game_rules import HoldemRules

class CppHandeval(CppWrapper):

    def __init__(self):
        super().__init__(path_to_dll=ospj(os.path.dirname(os.path.realpath(__file__)),
                                          "lib_hand_eval." + self.CPP_LIB_FILE_ENDING))
        self._clib.get_hand_rank_52_holdem.argtypes = [
            self.ARR_2D_ARG_TYPE,
            self.ARR_2D_ARG_TYPE
        ]
        self._clib.get_hand_rank_52_holdem.restype = ctypes.c_int32

        self._clib.get_hand_rank_all_hands_on_given_boards_52_holdem.argtypes = [
            self.ARR_2D_ARG_TYPE,
            self.ARR_2D_ARG_TYPE,
            ctypes.c_int32,
            self.ARR_2D_ARG_TYPE,
            self.ARR_2D_ARG_TYPE
        ]
        self._clib.get_hand_rank_all_hands_on_given_boards_52_holdem.restype = None

    def get_hand_rank_52_holdem(self, hand_2d, board_2d):
        """
        Args:
            hand_2d (np.ndarray(shape=[5,2], dtype=int8)):      [rank, suit], [rank, suit]]
            board_2d (np.ndarray(shape=[5,2], dtype=int8)):     [rank, suit], [rank, suit], ...]

        Returns:
            int: integer representing strength of the strongest 5card hand in the 7 cards. higher is better.
        """
        return self._clib.get_hand_rank_52_holdem(self.np_2d_arr_to_c(hand_2d), self.np_2d_arr_to_c(board_2d))

    def get_hand_rank_52_plo(self, hand_2d, board_2d):
        """
        Args:
            hand_2d (np.ndarray(shape=[5,2], dtype=int8)):      [rank, suit], [rank, suit]]
            board_2d (np.ndarray(shape=[5,2], dtype=int8)):     [rank, suit], [rank, suit], ...]

        Returns:
            int: integer representing strength of the strongest 5card hand
             in the 4 player + 5 board card, given that player have to use any 2. higher is better.
        Works:
            computes hand strength for all 6 possible combs of given Omaha hand, returns highest
            ---NEXT recreate this shit in pure C with fast 7-card evaluate algo(now its naive algo ~50k evals/sec)
            now it is totaling for a PLO hand at 7k hands/sec, while it can be around N mil hands/sec
        """
        hand_2c_combs = np.empty([6, 2, 2], dtype=np.int8)
        idxs = np.array([0, 1, 2, 3])
        dt = np.dtype([('', idxs.dtype)] * 2)
        # fill array right from combinations() func iteratively - its fast!
        idxs_combs = np.fromiter(itertools.combinations(idxs, 2), dt)
        for n, i in enumerate(idxs_combs):
            hand_2c_combs[n, 0] = hand_2d[i[0]]
            hand_2c_combs[n, 1] = hand_2d[i[1]]
        maxres = -1
        for hand in hand_2c_combs:
            res = self._clib.get_hand_rank_52_holdem(self.np_2d_arr_to_c(hand), self.np_2d_arr_to_c(board_2d))
            if res > maxres:
                maxres = res
        return maxres

    def get_hand_rank_all_hands_on_given_boards_52_holdem(self, boards_1d, lut_holder):
        """
        Args:
            boards_1d (np.ndarray(shape=[N, 5], dtype=int8)):   [[c1, c2, c3, c4, c5], [c1, c2, .., c5], ...}

        Returns:
            np.ndarray(shape=[N, RANGE_SIZE], dtype=int32):     hand_rank for each possible hand; -1 for
                                                                blocked on each of the given boards

        """
        assert len(boards_1d.shape) == 2
        assert boards_1d.shape[1] == 5
        hand_ranks = np.full(shape=(boards_1d.shape[0], HoldemRules.RANGE_SIZE), fill_value=-1, dtype=np.int32)
        self._clib.get_hand_rank_all_hands_on_given_boards_52_holdem(
            self.np_2d_arr_to_c(hand_ranks),  # int32**
            self.np_2d_arr_to_c(boards_1d),  # int8**
            boards_1d.shape[0],  # int (number of boards)
            self.np_2d_arr_to_c(lut_holder.LUT_IDX_2_HOLE_CARDS),  # int8**
            self.np_2d_arr_to_c(lut_holder.LUT_1DCARD_2_2DCARD)  # int8**
        )
        return hand_ranks


import string
import argparse

import sys
import torch
import torch.backends.cudnn as cudnn
import torch.utils.data

from utils import CTCLabelConverter, AttnLabelConverter, BertConverter, TransformerConverter
from dataset import FileAlignCollate, AlignCollate
from model import Model
from new_model import NewModel
from PIL import Image
from BertAttnModel import BertAttnModel, Identity
from torch.utils.data import Dataset, DataLoader
import numpy as np
import cv2
import math

from model_cfg import RecognitionConfig, g_BERT_Attn_config, g_BERT_Attn_cpx_config, g_BERT_CTC_config, g_BERT_CTC_cpx_config

BERT = False

class Recognition():

    def __init__(self, recog_config = None, gpu_id = None ):

        """
            gpu_id : int
        """
        if recog_config :
            self.opt = recog_config
        else :
            self.opt = RecognitionConfig(
                saved_model = 'saved_models/BERT_CTC/best_norm_ED.pth',
                batch_max_length = 35,
                imgH = 32,
                imgW = 256,
                character = 'new_dict.txt',
                minC = 5000,
                PAD = True,
                Transformation = 'TPS',
                FeatureExtraction = 'SEResNet',
                SequenceModeling = 'BERT',
                Prediction = 'CTC',
                num_fiducial = 20,
                input_channel = 1,
                output_channel = 512,
                hidden_size = 256
            )

        # tmp
        # parser = argparse.ArgumentParser()
        # parser.add_argument('--saved_model', type=str,default='saved_models/BERT_Attn_cpx/best_norm_ED.pth' ,help="path to saved_model to evaluation")
        # """ Data processing """
        # parser.add_argument('--batch_max_length', type=int, default=35, help='maximum-label-length')
        # parser.add_argument('--imgH', type=int, default=32, help='the height of the input image')
        # parser.add_argument('--imgW', type=int, default=256, help='the width of the input image')
        # # parser.add_argument('--rgb', action='store_true', help='use rgb input')
        # parser.add_argument('--rgb', default=False, help='use rgb input')
        # parser.add_argument('--character', type=str, default='new_dict.txt', help='character label')
        # parser.add_argument('--minC', type=int, default=10000, help='min classes of dict ')
        # # parser.add_argument('--sensitive', action='store_true', help='for sensitive character mode')
        # parser.add_argument('--PAD', type=bool, default= True, help='whether to keep ratio then pad for image resize')
        # """ Model Architecture """
        # parser.add_argument('--Transformation', type=str, default='TPS', help='Transformation stage. None|TPS')
        # parser.add_argument('--FeatureExtraction', type=str, default='SEResNet', help='FeatureExtraction stage. VGG|RCNN|ResNet')
        # parser.add_argument('--SequenceModeling', type=str, default='BERT', help='SequenceModeling stage. None|BiLSTM')
        # parser.add_argument('--Prediction', type=str, default='Attn', help='Prediction stage. CTC|Attn')
        # parser.add_argument('--num_fiducial', type=int, default=20, help='number of fiducial points of TPS-STN')
        # parser.add_argument('--input_channel', type=int, default=1, help='the number of input channel of Feature extractor')
        # parser.add_argument('--output_channel', type=int, default=512,
        #                     help='the number of output channel of Feature extractor')
        # parser.add_argument('--hidden_size', type=int, default=256, help='the size of the LSTM hidden state')

        c_dict=open(self.opt.character , encoding='utf-8')
        c_dict=c_dict.read().split()
        self.opt.character=''.join(c_dict)
        self.gpu_id = gpu_id

        self.model= self.buildModel()

        self.AlignCollate_demo = FileAlignCollate(imgH=self.opt.imgH, imgW=self.opt.imgW, keep_ratio_with_pad=self.opt.PAD, c=self.opt.input_channel)

    def _cutRects(self, np_image, np_bboxes):
        """
            np_image: 整张图片 numpy格式, 需要RGB顺序
            np_bboxes: 这张图片的文字区域的bbox, numpy array (n,4,2), 点可以乱序,需要(x,y)格式
        """
        def order_points_old(pts):
            #TODO: parallel code

            rect = np.zeros_like(pts)

            s = pts.sum(axis=1)
            rect[0] = pts[np.argmin(s)]
            rect[2] = pts[np.argmax(s)]

            diff = np.diff(pts, axis=1)
            rect[1] = pts[np.argmin(diff)]
            rect[3] = pts[np.argmax(diff)]

            return rect

        rects = np.zeros_like(np_bboxes)
        for i in range(rects.shape[0]):
            rects[i] = order_points_old(np_bboxes[i])

        pil_results=[]
        for i, rect in enumerate(rects):

            w = int(np.linalg.norm(rect[0]-rect[1]))
            h = int(np.linalg.norm(rect[0]-rect[-1]))

            four_points=np.array(((0,0),(w-1,0),(w-1,h-1),(0,h-1)),dtype=np.float32)
            target_points = rect.astype(np.float32)

            try:
                M = cv2.getPerspectiveTransform(target_points, four_points)
                Rotated= cv2.warpPerspective(np_image, M, (w, h))
                pil_pic = Image.fromarray(Rotated)

                if self.opt.input_channel == 1 :
                    pil_pic = pil_pic.convert('L')
                else :
                    pil_pic = pil_pic.convert('RGB')
                pil_results.append(pil_pic)

            except:
                print('cut failed ')
                pass

        return pil_results

    def bboxPicRecognition(self ,np_image, np_bboxes):
        """
            np_image: 整张图片 numpy格式, 需要RGB顺序
            np_bboxes: 这张图片的文字区域的bbox, numpy array (n,4,2), 点可以乱序,需要(x,y)格式
        """
        if len(np_bboxes) == 0 :
            print('empty !')
            image_list = []
        else :
            image_list = self._cutRects(np_image, np_bboxes)

        print('preds str len{}-{}'.format(len(np_bboxes),len(image_list)))
        result = self.imageBatchRecognition(image_list)

        return result

    def buildModel(self):
        opt=self.opt

        """ model configuration """

        if len(opt.character) < opt.minC:
            opt.character += ' '*(opt.minC-len(opt.character))

        if BERT:
            print('bert converter !')
            converter = BertConverter(opt.character)
        elif 'CTC' in opt.Prediction:
            converter = CTCLabelConverter(opt.character)
        elif opt.Prediction == 'Transformer':
            converter = TransformerConverter(opt.character)
        else:
            converter = AttnLabelConverter(opt.character)

        self.converter = converter

        opt.num_class = len(converter.character)

        if BERT:
            print(opt.num_class)
            model = BertAttnModel(num_class=opt.num_class)
        else :
            model = Model(opt)
        # model=NewModel(opt)

        if self.gpu_id :
            self.gpu_id = [self.gpu_id]

        model = torch.nn.DataParallel(model, device_ids = self.gpu_id)
        if torch.cuda.is_available():
            model = model.cuda()

        # load model
        if opt.saved_model:
            print('start load model')
            print('loading  model from {}'.format(opt.saved_model))
            model.load_state_dict(torch.load(opt.saved_model))
            print('success load model ')
        return model

    @torch.no_grad()
    def _recognition(self,images):
        model=self.model
        opt=self.opt
        model.eval()
        with torch.no_grad():
            images = images.cuda()
            length_for_pred = torch.cuda.IntTensor([opt.batch_max_length] * len(images))
            text_for_pred = torch.cuda.LongTensor(len(images), opt.batch_max_length + 1).fill_(0)

        if BERT:
            scores = model(images)
            _, preds = scores.max(-1)
            # preds_str = self.converter.decode(preds,decode_type=0)
            preds_str = self.converter.decode(preds)

        elif 'CTC' in opt.Prediction:
            preds=model(images,text_for_pred)
            preds_size=torch.IntTensor([preds.size(1)]*len(images))
            _, preds_index = preds.permute(1,0,2).max(2)

            preds_index = preds_index.transpose(1,0).contiguous().view(-1)
            preds_str= self.converter.decode(preds_index.data,preds_size.data)

        elif opt.Prediction == 'Transformer':
            preds_index = model(images, None, is_train= False)
            preds_str = self.converter.decode(preds_index)

        else :
            preds = model(images, text_for_pred, is_train=False)
            _, preds_index = preds.max(2)
            preds_str = self.converter.decode(preds_index, length_for_pred)

        result=[]

        for pred in preds_str:
            if 'Attn' in opt.Prediction:
                pred = pred[:pred.find('[s]')]  # prune after "end of sentence" token ([s])
            result.append(pred)

        return result

    def fileRecognition(self,file_name):

        opt=self.opt
        model=self.model
        model.eval()
        if opt.input_channel == 1 :
            pil_image = Image.open(file_name).convert('L')
        else :
            pil_image = Image.open(file_name).convert('RGB')

        image_tensors=self.AlignCollate_demo([pil_image])

        # predict
        # model.eval()
        # with torch.no_grad():
        #     image = image_tensors.cuda()
        #     length_for_pred = torch.cuda.IntTensor([opt.batch_max_length] * 1)
        #     text_for_pred = torch.cuda.LongTensor(1, opt.batch_max_length + 1).fill_(0)

        # if 'CTC' in opt.Prediction:
        #     # preds=model(image,text_for_pred).log_softmax(2)

        #     preds=model(image,text_for_pred)
        #     preds_size=torch.IntTensor([preds.size(1)])
        #     _, preds_index = preds.permute(1,0,2).max(2)
        #     preds_index = preds_index.transpose(1,0).contiguous().view(-1)
        #     preds_str= self.converter.decode(preds_index.data,preds_size.data)

        # else :
        #     preds = model(image, text_for_pred, is_train=False)
        #     _, preds_index = preds.max(2)
        #     preds_str = self.converter.decode(preds_index, length_for_pred)

        result=self._recognition(image_tensors)

        return result[0]

    def imageBatchRecognition(self, image_list, batch_size=16):
        """
            image_batch : list of PIL Image
            return : list of labels
        """

        index=0
        result=[]

        while index < len(image_list):
            print(index,end='\r')

            images=self.AlignCollate_demo(image_list[index:index+batch_size])
            part_result=self._recognition(images)
            result+=part_result
            index+=batch_size

        return result

g_Recognizer = Recognition(g_BERT_CTC_cpx_config)

if __name__ =='__main__':
    import time
    import os
    import sys
    #model=Recognition(g_BERT_CTC_cpx_config)
    model = g_Recognizer
    start_time=time.time()
    # g_root_dir='cut_test'
    # g_root_dir = '/ai/local/menglc/v3_dataset/raw_img/validation/ReCTS'
    # g_root_dir = '/ai/local/menglc/ocrdata/test_cut'
    # g_root_dir = '/home/menglc/ocrdata/cut_img'
    # g_root_dir = '/ai/local/menglc/v3_dataset/raw_img/validation/ctw'
    g_root_dir = '/home/menglc/ocrdata/TextRecognitionDataGenerator-master/TextRecognitionDataGenerator/poster_simple'
    count=1e-8
    correct_count=0
    mode=4

    if mode==0:
        # from collections import OrderedDict
        # model_dict = model.model.module.SequenceModeling.state_dict()

        # from torchvision.models import resnet34

        # net = resnet34()
        # net.load_state_dict(torch.load('pretrain/resnet34-333f7ec4.pth'))
        # print(net)

        # net.avgpool= Identity()
        # net.fc = Identity()
        # new_dict= OrderedDict()
        # for k,v in model_dict.items():
        #     new_dict[k[7:]]=v

        from transformer_decoder import TransFormerDecoder
        from modules.prediction import Attention

        new_decoder = Attention(256, 256, 10002)
        # new_decoder = TransFormerDecoder(50,5004)

        pre_decoder = model.model.module.Prediction

        # new_decoder = torch.nn.Linear(256, 10001)

        # print('pre decoder weight:{}, bias:{} .'.format(pre_decoder.weight.shape, pre_decoder.bias.shape))

        # torch.nn.init.constant_(new_decoder.bias, 0.0)
        # torch.nn.init.kaiming_normal_(new_decoder.weight)

        # new_decoder.weight[:5001,:] = pre_decoder.weight
        # new_decoder.bias[:5001] = pre_decoder.bias
        # new_decoder.weight[5001:,:] = 0.0
        # new_decoder.bias[5001:] = 0.0

        model.model.module.Prediction =  new_decoder

        torch.save(model.model.state_dict(),'pretrain/pretrain_bigger.pth')

        import sys
        sys.exit()

        # torch.save(net.state_dict(),'pretrain/my_resnet34.pth')
        # torch.save(model.model.module.Transformation.state_dict(),'pretrain/HV_CTC_tps.pth')
        # torch.save(model.model.module.FeatureExtraction.state_dict(),'pretrain/HV_CTC_seresnet.pth')

    elif mode ==1:
        log_file=open('Cut_result.txt','w',encoding='utf-8')
        log_file.write('file_name\tresult\tpred\tgt\n')
        dir_list=os.listdir(g_root_dir)
        dir_list.sort()
        pil_images=[]
        gts=[]
        for i,name in enumerate(dir_list):
            if name.endswith('jpg'):
                print(i,end='\r')
                pil_image = Image.open(os.path.join(g_root_dir,name)).convert('L')
                # pil_image = Image.open(os.path.join(g_root_dir,name)).convert('RGB')
                pil_images.append(pil_image)
                # result=model.fileRecognition(os.path.join(g_root_dir,name))
                try:
                    gt = open(os.path.join(g_root_dir,name[:-4]+'.txt')).read()
                    gt=gt.split()[0]
                except:
                    gt ='###'
                gts.append(gt)

        results = model.imageBatchRecognition(pil_images)
        print('len results: {}, len gts: {}, '.format(len(results),len(gts)))
        for result, gt in zip(results,gts):

            log_file.write('{}\t{}\t{}\t{}\n'.format(name,(result==gt),result,gt))
            if result==gt :
                correct_count+=1
            # log_file.write('{}\t{}\n'.format(name,result))
            count+=1

    elif mode==2:
        log_file=open('Cut_result.txt','w',encoding='utf-8')
        log_file.write('file_name\tresult\tpred\tgt\n')
        dir_list=os.listdir(g_root_dir)
        dir_list.sort()
        for i,name in enumerate(dir_list):
            if name.endswith('jpg'):
                print(count,end='\r')
                result=model.fileRecognition(os.path.join(g_root_dir,name))
                try:
                    gt = open(os.path.join(g_root_dir,name[:-4]+'.txt')).read()
                    gt=gt.split()[0]
                except:
                    gt ='###'

                log_file.write('{}\t{}\t{}\t{}\n'.format(name,(result==gt),result,gt))
                if result==gt :
                    correct_count+=1
                # log_file.write('{}\t{}\n'.format(name,result))
                count+=1

    elif mode==3:
        g_root_dir='normal_test/results_feiweijin'
        dir_list=os.listdir(g_root_dir)
        dir_list.sort()
        for i,name in enumerate(dir_list):
            # if name.endswith('jpg'):
            # print(i,end='\r')
            result=model.fileRecognition(os.path.join(g_root_dir,name))
            # gt = open(os.path.join(g_root_dir,name[:-4]+'.txt')).read()
            # log_file.write('{}\t{}\t\n'.format((result==gt),result))
            # log_file.write('{}\t{}\n'.format(name,result))
            print('{}\t{}'.format(name,result))
            count+=1

    elif mode ==4:
        count=1
        result=model.fileRecognition('test_2.png')
        print('result ',result)

    elif mode==5:
        log_file=open('Cut_result.txt','w',encoding='utf-8')
        log_file.write('file_name\tresult\tpred\tgt\n')
        dir_list=os.listdir(g_root_dir)
        dir_list.sort()
        #gts = open(os.path.join(g_root_dir,'result.txt')).readlines()
        gts = open('/ai/local/menglc/v3_dataset/raw_img/ctwReCTS_validate.txt').readlines()[:200]
        for i, data in enumerate(gts):
            img_path, gt = data.split('\t')
            gt = gt.split('\n')[0]
            print(count,end='\r')
            result=model.fileRecognition(img_path)
            name = img_path.split('/')[-1]
            log_file.write('{}\t{}\t{}\t{}\n'.format(name,(result==gt),result,gt))
            if result==gt :
                correct_count+=1
            # log_file.write('{}\t{}\n'.format(name,result))
            count+=1
    elif mode ==6 :
        log_file=open('Cut_result.txt','w',encoding='utf-8')
        log_file.write('file_name\tresult\tpred\tgt\n')
        #gt_file = open('/ai/local/menglc/v3_dataset/raw_img/ctwReCTS_validate.txt').readlines()
        # gt_file = open('/home/menglc/ocrdata/poster_simple/result.txt').readlines()
        gt_file = open('/home/menglc/ocrdata/cut_img/result.txt').readlines()
        # gt_file = open('/ai/local/menglc/LSVT/total_data/cut_validate/result.txt').readlines()
        gts=[]
        pil_images=[]
        for i, data in enumerate(gt_file):
            name, gt = data.split('\t')
            gt = gt.split('\n')[0]
            pil_image = Image.open(os.path.join(g_root_dir,name)).convert('L')
            pil_images.append(pil_image)
            gts.append(gt)
            print(i,end='\r')

        count=0
        results = model.imageBatchRecognition(pil_images)
        print('len results: {}, len gts: {}, '.format(len(results),len(gts)))
        for result, gt in zip(results,gts):

            log_file.write('{}\t{}\t{}\t{}\n'.format('name',(result==gt),result,gt))
            if result==gt :
                correct_count+=1
            # log_file.write('{}\t{}\n'.format(name,result))
            count+=1

    elif mode == 7:
        AlignCollate_valid= AlignCollate(imgH=32, imgW=256, keep_ratio_with_pad=True ,c=1 )
        test_dataset = TestDataset('/ai/local/menglc/v3_dataset/raw_img/ctwReCTS_validate.txt')
        valid_loader = DataLoader(
            test_dataset, batch_size=32,
            shuffle=True,
            collate_fn=AlignCollate_valid, pin_memory=True
        )

        count =0
        for i, (images, labels) in enumerate(valid_loader):
            print(i, end='\r')
            results = model._recognition(images)

            for result, label in zip(results, labels):
                if result == label:
                    correct_count +=1

                count+=1

    elif mode == 8 :
        import os
        import os.path as osp

        log_file=open('Cut_result.txt','w',encoding='utf-8')
        # image_root = '/home/menglc/ocrdata/SceneTextGenerator/poster_simple'
        image_root = '/ai/local/menglc/LSVT/total_data/validate'
        image_name = os.listdir(image_root)
        for i, name in enumerate(image_name) :
            if name.endswith('.jpg'):
                print(i//2, end='\r')
                image = cv2.imread(osp.join(image_root, name))
                # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                ps_str = open(osp.join(image_root,name.split('.')[0]+'.txt')).read().split('\n')
                pps = []
                labels = []
                for line in ps_str :
                    if line == '':
                        continue
                    line = line.split(',')

                    ps = [int(float(x)) for x in line[:8]]
                    pps.append((ps[0:2],ps[2:4],ps[4:6],ps[6:8]))
                    labels.append(line[8])

                pps = np.array(pps)
                # print(type(pps), pps.shape)

                results = model.bboxPicRecognition(image, pps)

                for result, gt in zip(results, labels):
                    # print('pred\t:{}\ngt:\t{}'.format(result, gt))
                    log_file.write('{}\t{}\t{}\t{}\n'.format(name,(result==gt),result,gt))
                    if result==gt :
                        correct_count+=1
                        # log_file.write('{}\t{}\n'.format(name,result))
                    count+=1

    total_time=time.time()-start_time
    log_file.write('infer time:{}, avg:{}, acc:{}'.format(total_time,total_time/count,correct_count/count))
    print('infer time:{}, avg:{}, acc:{}'.format(total_time,total_time/count,correct_count/count))


"""
This module consists database class to handle different different data split, loading data into memory, dividing long
audio file into segments and tokenize these segments
Note: 32bit is sufficient for DL models.
Note on terminology on audio length: frames -> segments -> chunk/clip -> file
"""
import logging
import os
from typing import List

import h5py
import numpy as np
import pandas as pd

class SedDoaDatabase:
    """
    Class to handle different extracted features for SED or DOA separately.
    """
    def __init__(self,
                 feature_root_dir: str = '/media/tho_nguyen/disk2/new_seld/dcase2020/features/'
                                         'logmel_norm/24000fs_1024nfft_300nhop_128nmels_Falsestd',
                 gt_meta_root_dir: str = '/media/tho_nguyen/disk1/audio_datasets/dcase2020/task3',
                 audio_format: str = 'foa', n_classes: int = 14, fs: int = 24000, n_fft: int = 1024, hop_len: int = 300,
                 label_rate: float = 10, train_chunk_len_s: float = 4.0, train_chunk_hop_len_s: float = 0.5,
                 test_chunk_len_s: float = 4.0, test_chunk_hop_len_s: float = 2.0, scaler_type: str = 'vector'):
        """
        :param feature_root_dir: Feature directory. can be SED or DOA feature.
        The data are organized in the following format:
            |__feature_root_dir/
                |__foa_dev/
                |__foa_eval/
                |__mic_dev/
                |__mic_eval/
                |__foa_feature_scaler.h5
                |__mic_feature_scaler.h5
        :param gt_meta_root_dir: Directory that contains groundtruth meta data.
        The data are orgamized in the following format:
            |__gt_meta_dir/
                |__/metadata_dev/
                |__/metadata_eval/
                |__metadata_eval_info.csv
        """
        self.feature_root_dir = feature_root_dir
        self.gt_meta_root_dir = gt_meta_root_dir
        self.audio_format = audio_format
        self.n_classes = n_classes
        self.fs = fs
        self.n_fft = n_fft
        self.hop_len = hop_len
        self.label_rate = label_rate
        self.train_chunk_len = self.second2frame(train_chunk_len_s)
        self.train_chunk_hop_len = self.second2frame(train_chunk_hop_len_s)
        self.test_chunk_len = self.second2frame(test_chunk_len_s)
        self.test_chunk_hop_len = self.second2frame(test_chunk_hop_len_s)
        self.scaler_type = scaler_type

        assert audio_format in ['foa', 'mic'], 'Incorrect value for audio format {}'.format(audio_format)
        assert os.path.isdir(os.path.join(self.feature_root_dir, self.audio_format + '_dev')), \
            '"dev" folder is not found'

        self.chunk_len = None
        self.chunk_hop_len = None
        self.n_frames = int(np.floor((self.fs * 60 - (self.n_fft - self.hop_len)) / self.hop_len)) + 2  #+ 2 because of padding
        self.feature_rate = self.fs/self.hop_len  # Frame rate per second
        self.label_upsample_ratio = int(self.feature_rate / self.label_rate)
        self.mean, self.std = self.load_scaler()

        logger = logging.getLogger('lightning')
        logger.info('Load feature database from {}.'.format(self.feature_root_dir))
        logger.info('train_chunk_len = {}, train_chunk_hop_len = {}'.format(
            self.train_chunk_len, self.train_chunk_hop_len))
        logger.info('test_chunk_len = {}, test_chunk_hop_len = {}'.format(
            self.test_chunk_len, self.test_chunk_hop_len))

    def get_split(self, split: str, split_meta_dir: str = '/meta/original', doa_format: str = 'xyz'):
        """
        Function to load all data of a split into memory, divide long audio clip/file into smaller chunks, and assign
        labels for clips and chunks. List of SED labels:

        :param split: Split of data, choices:
            'train', 'val', 'test', 'eval': load chunk of data
        :param split_meta_dir: Directory where meta of split is stored.
        :param doa_format: Choices are 'xyz' or 'polar'.
        :return:
        """
        assert doa_format in ['xyz', 'polar'], 'Incorrect value for doa format {}'.format(doa_format)
        # Get feature dir, filename list, and gt_meta_dir
        if split == 'eval':
            split_feature_dir = os.path.join(self.feature_root_dir, self.audio_format + '_eval')
            csv_filename = os.path.join(os.path.split(split_meta_dir)[0], 'eval.csv')
            gt_meta_dir = os.path.join(self.gt_meta_root_dir, 'metadata_eval')
        else:
            split_feature_dir = os.path.join(self.feature_root_dir, self.audio_format + '_dev')
            csv_filename = os.path.join(split_meta_dir, split + '.csv')
            gt_meta_dir = os.path.join(self.gt_meta_root_dir, 'metadata_dev')
        meta_df = pd.read_csv(csv_filename)
        split_filenames = meta_df['filename'].tolist()
        # Get chunk len and chunk hop len
        if split in ['train', 'trainval', 'val']:
            self.chunk_len = self.train_chunk_len
            self.chunk_hop_len = self.train_chunk_hop_len
        elif split in ['test', 'eval']:
            self.chunk_len = self.test_chunk_len
            self.chunk_hop_len = self.test_chunk_hop_len
        else:
            raise NotImplementedError('chunk len is not assigned for split {}'.format(split))

        # Load and crop data
        features, sed_targets, doa_targets, chunk_idxes, filename_list, test_batch_size = self.load_chunk_data(
            split_filenames=split_filenames, split_feature_dir=split_feature_dir, gt_meta_dir=gt_meta_dir,
            doa_format=doa_format, split=split)
        # pack data
        db_data = {
            'features': features,
            'sed_targets': sed_targets,
            'doa_targets': doa_targets,
            'chunk_idxes': chunk_idxes,
            'filename_list': filename_list,
            'test_batch_size': test_batch_size
        }

        return db_data

    def second2frame(self, second):
        """
        Convert seconds to frame unit.
        """
        sample = int(second * self.fs)
        frame = int(round(sample/self.hop_len))
        return frame

    def load_scaler(self):
        scaler_fn = os.path.join(self.feature_root_dir, self.audio_format + '_feature_scaler.h5')
        if self.scaler_type == 'vector':
            with h5py.File(scaler_fn, 'r') as hf:
                mean = hf['mean'][:]
                std = hf['std'][:]
        elif self.scaler_type == 'scalar':
            with h5py.File(scaler_fn, 'r') as hf:
                mean = hf['scalar_mean'][:]
                std = hf['scalar_std'][:]
        else:
            mean = 0
            std = 1
        return mean, std

    def load_chunk_data(self, split_filenames: List, split_feature_dir: str, gt_meta_dir: str,
                        doa_format: str = 'xyz', split: str = 'train'):
        """
        Load feature, crop data and assign labels.
        :param split_filenames: List of filename in the split.
        :param split_feature_dir: Feature directory of the split
        :param gt_meta_dir: Ground truth meta directory of the split.
        :param doa_format: Choices are 'xyz' or 'polar'.
        :param split: Name of split, can be 'train', 'trainval', 'val', 'test', 'eval'.
        :return: features, targets, chunk_idxes, filename_list
        """
        pointer = 0
        features_list = []
        filename_list = []
        sed_targets_list = []
        doa_targets_list = []
        idxes_list = []
        for filename in split_filenames:
            feature_fn = os.path.join(split_feature_dir, filename + '.h5')
            # Load feature
            with h5py.File(feature_fn, 'r') as hf:
                feature = hf['feature'][:]
            # Normalize feature
            feature = (feature - self.mean) / self.std
            n_frames = feature.shape[1]
            # Load gt info from metadata
            gt_meta_fn = os.path.join(gt_meta_dir, filename + '.csv')
            df = pd.read_csv(gt_meta_fn, header=None,
                             names=['frame_number', 'sound_class_idx', 'track_number', 'azimuth', 'elevation'])
            frame_number = df['frame_number'].values
            sound_class_idx = df['sound_class_idx'].values
            track_number = df['track_number'].values
            azimuth = df['azimuth'].values
            elevation = df['elevation'].values
            # Generate target data
            sed_target = np.zeros((n_frames, self.n_classes), dtype=np.float32)
            azi_target = np.zeros((n_frames, self.n_classes), dtype=np.float32)
            ele_target = np.zeros((n_frames, self.n_classes), dtype=np.float32)
            nsources_target = np.zeros((n_frames, 3), dtype=np.float32)
            count_sources_target = np.zeros((n_frames,), dtype=np.float32)
            for itrack in np.arange(5):
                track_idx = track_number == itrack
                frame_number_1 = frame_number[track_idx]
                sound_class_idx_1 = sound_class_idx[track_idx]
                azimuth_1 = azimuth[track_idx]
                elevation_1 = elevation[track_idx]
                for idx, iframe in enumerate(frame_number_1):
                    start_idx = int(iframe * self.label_upsample_ratio - self.label_upsample_ratio//2)
                    start_idx = np.max((0, start_idx))
                    end_idx = int(start_idx + self.label_upsample_ratio)
                    end_idx = np.min((end_idx, n_frames))
                    class_idx = int(sound_class_idx_1[idx])
                    sed_target[start_idx:end_idx, class_idx] = 1.0
                    azi_target[start_idx:end_idx, class_idx] = azimuth_1[idx] * np.pi / 180.0  # Radian unit
                    ele_target[start_idx:end_idx, class_idx] = elevation_1[idx] * np.pi / 180.0  # Radian unit
                    count_sources_target[start_idx:end_idx] += 1
            # Convert nsources to one-hot encoding
            for i in np.arange(3):
                idx = count_sources_target == i
                nsources_target[idx, i] = 1.0
            # Doa target
            if doa_format == 'polar':
                doa_target = np.concatenate((azi_target, ele_target), axis=-1)
            elif doa_format == 'xyz':
                x = np.cos(azi_target) * np.cos(ele_target)
                y = np.sin(azi_target) * np.cos(ele_target)
                z = np.sin(ele_target)
                doa_target = np.concatenate((x, y, z), axis=-1)
            # Get segment indices
            n_crop_frames = n_frames
            assert self.chunk_len <= n_crop_frames, 'Number of cropped frame is less than chunk len'
            idxes = np.arange(pointer, pointer + n_crop_frames - self.chunk_len + 1, self.chunk_hop_len).tolist()
            # Include the leftover of the cropped data
            if (n_crop_frames - self.chunk_len) % self.chunk_hop_len != 0:
                idxes.append(pointer + n_crop_frames - self.chunk_len)
            pointer += n_crop_frames
            # Append data
            features_list.append(feature)
            filename_list.extend([filename] * len(idxes))
            sed_targets_list.append(sed_target)
            doa_targets_list.append(doa_target)
            idxes_list.append(idxes)

        if len(features_list) > 0:
            features = np.concatenate(features_list, axis=1)
            sed_targets = np.concatenate(sed_targets_list, axis=0)
            doa_targets = np.concatenate(doa_targets_list, axis=0)
            chunk_idxes = np.concatenate(idxes_list, axis=0)
            test_batch_size = len(idxes)  # to load all chunks of the same file
            return features, sed_targets, doa_targets, chunk_idxes, filename_list, test_batch_size
        else:
            return None, None, None, None, None

class SeldDatabase():
    """
    Database class to handle two input streams, one for SED, one for DOA. Use this database to train alignment module
    """
    pass

if __name__ == '__main__':
    tp_db = SedDoaDatabase()
    db_data = tp_db.get_split(split='val', split_meta_dir='meta/original')
    print(db_data['features'].shape)
    print(db_data['sed_targets'].shape)
    print(len(db_data['chunk_idxes']))
    print(len(db_data['filename_list']))


import numpy as np
import os
import tensorflow as tf
import glob
import time
import random
import argparse

from functools import partial
from tqdm import tqdm
from collections import namedtuple
from typing import Sequence, Tuple, List

## (shiyao) TODO : adjust these imports
from feedback_synthesizer.hparams import hparams
from deep_speaker.resnet import ResNet
from feedback_synthesizer.models.embedding.Resnet import resnet_hparams

trainline = namedtuple(
    'trainline',
    [
        'audio',
        'mel',
        'embed',
        'nsample',
        'nframe',
        'text'
    ]
)

def build_and_restore_resnet(sess : tf.Session, ckpt_path : str) -> Tuple[tf.placeholder, tf.placeholder] :
    """
    build and restore resnet in a given session, return placehoders for input and output values
    :param sess: tensorflow session in which the graph state is to be restored
    :param ckpt_path: path to saved model chekcpoint to load state from

    :return: a tuple of tf.Placeholder s, `mel_input` and `gv`, which are the input and output placehoders
    for the resnet encoder
    """

    with tf.variable_scope('resnet', reuse = tf.AUTO_REUSE) as resnet_scopt :
        mel_input = tf.placeholder(tf.float32, (None, None, hparams.num_mels), name='mel_input')
        labels    = tf.placeholder(tf.int32, (None), name='labels')

        net       = ResNet(resnet_hparams, mel_input, labels, 'eval')
        net.build_graph()

    ckpt_state = tf.train.get_checkpoint_state(ckpt_path)
    saver      = tf.train.Saver()

    saver.restore(sess, ckpt_state.model_checkpoint_path)

    return mel_input, net.gv

def parse_metafile(filename : str) -> List[trainline] :
    with open(filename) as f :
        return [
            trainline(*line.strip().split('|'))
            for line in f
        ]

def audio_path_from_metaline(l : trainline) -> str :
    return os.path.join(
        'audio',
        l.spkid,
        f'{l.sid}.wav'
    )

def _embedding_from_line(sess : tf.Session, nodes : Tuple[tf.placeholder, tf.placeholder], mel_dir : str, l : trainline) -> np.ndarray :

    mel_node, gv_node = nodes

    # load the mel-spectrogram
    mel = np.load(os.path.join(mel_dir, l.mel))

    gvector = sess.run(
        gv_node,
        {
            mel_node : mel[None]
        }
    )

    return gvector[0]

def _embedding_from_line_direct(sess : tf.Session, nodes : Tuple[tf.placeholder, tf.placeholder],p_mel : str) -> np.ndarray :

    mel_node, gv_node = nodes

    # load the mel-spectrogram
    # mel = np.load(p_mel).T  # when used for generated mel
    mel = np.load(p_mel)    # when used for extracted mel
    # the difference being axis order

    gvector = sess.run(
        gv_node,
        {
            mel_node : mel[None]
        }
    )

    return gvector[0]

def main() -> None :

    parser = argparse.ArgumentParser()
    parser.add_argument('dataset_dir', help='path to dataset dir, containing train.txt and mel-spectrograms', type=str)
    parser.add_argument('--gvec_ckpt', help='path ot resnet checkpoint', type=str, default='synthesizer_with_encoder/models/embedding/save_models')
    args = parser.parse_args()

    metafile = parse_metafile(os.path.join(args.dataset_dir, 'train.txt'))

    embed_dir = os.path.join(args.dataset_dir, 'embeds')
    mel_dir   = os.path.join(args.dataset_dir, 'mels')
    os.makedirs(embed_dir, exist_ok=True)

    # load network and parameters
    with tf.Session() as sess :
        node = build_and_restore_resnet(sess, args.gvec_ckpt)
        embedding_from_line = partial(_embedding_from_line, sess, node, mel_dir)

        for tl in tqdm(metafile) :
            embedding = embedding_from_line(tl)
            np.save(os.path.join(embed_dir, tl.embed), embedding)

        print(f'n of metafile : {len(metafile)}')
        print(f'e.g. {random.sample(metafile, 5)}')

def main_tsne() :

    gvec_ckpt = 'synthesizer_with_encoder/models/embedding/save_models'

    mel_dir = 't_sne_data/trial_4/dep_mels'
    out_dir = 't_sne_data/trial_4/dep_embeds'

    os.makedirs(out_dir, exist_ok=True)

    with tf.Session() as sess :
        node = build_and_restore_resnet(sess, gvec_ckpt)
        embedding_from_line = partial(_embedding_from_line_direct, sess, node)

        for mel_path in tqdm(glob.glob(os.path.join(mel_dir, '*.npy'))) :
            embedding = embedding_from_line(mel_path)
            np.save(os.path.join(out_dir, os.path.basename(mel_path)), embedding)

if __name__ == '__main__' :
    main()
    # main_tsne()


"""Utilities for coordinate-based meta-analysis estimators."""
import logging
import os

import nibabel as nib
import numpy as np
import numpy.linalg as npl
from scipy import ndimage

from .. import references
from ..due import due
from ..extract import download_peaks2maps_model
from ..utils import _determine_chunk_size

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
LGR = logging.getLogger(__name__)

def model_fn(features, labels, mode, params):
    """Run model function used internally by peaks2maps.

    .. deprecated:: 0.0.11
        `model_fn` will be removed in NiMARE 0.0.13.

    .. versionadded:: 0.0.4

    """
    import tensorflow as tf
    from tensorflow.python.estimator.export.export_output import PredictOutput

    ngf = 64
    layers = []

    training_flag = mode == tf.estimator.ModeKeys.TRAIN

    input_images_placeholder = tf.expand_dims(features, -1)

    conv_args = {
        "strides": 2,
        "kernel_size": 4,
        "padding": "valid",
        "activation": tf.nn.leaky_relu,
        "kernel_initializer": tf.random_normal_initializer(0, 0.02),
        "name": "conv",
        "use_bias": False,
    }

    deconv_args = conv_args.copy()
    deconv_args["padding"] = "same"
    deconv_args["name"] = "deconv"

    batchnorm_args = {
        "scale": True,
        "gamma_initializer": tf.random_normal_initializer(1.0, 0.02),
        "center": True,
        "beta_initializer": tf.zeros_initializer(),
        "name": "batchnorm",
        "training": training_flag,
    }

    def pad_and_conv(input, out_channels, conv_args):
        padded_input = tf.pad(input, [[0, 0], [1, 1], [1, 1], [1, 1], [0, 0]], mode="CONSTANT")
        convolved = tf.compat.v1.layers.conv3d(padded_input, out_channels, **conv_args)
        return convolved

    # encoder_1: [batch, 256, 256, in_channels] => [batch, 128, 128, ngf]
    with tf.compat.v1.variable_scope("encoder_1"):
        this_args = conv_args.copy()
        output = pad_and_conv(input_images_placeholder, ngf, this_args)
        layers.append(output)

    layer_specs = [
        (ngf * 2, 0.2),
        # encoder_3: [batch, 64, 64, ngf * 2] => [batch, 32, 32, ngf * 4]
        (ngf * 2, 0.2),
        # encoder_3: [batch, 64, 64, ngf * 2] => [batch, 32, 32, ngf * 4]
        (ngf * 4, 0.2),
        # encoder_4: [batch, 32, 32, ngf * 4] => [batch, 16, 16, ngf * 8]
        (ngf * 8, 0.2),
        # encoder_5: [batch, 16, 16, ngf * 8] => [batch, 8, 8, ngf * 8]
        # ngf * 8,
        # # encoder_6: [batch, 8, 8, ngf * 8] => [batch, 4, 4, ngf * 8]
    ]

    for out_channels, dropout in layer_specs:
        with tf.compat.v1.variable_scope("encoder_%d" % (len(layers) + 1)):
            # [batch, in_height, in_width, in_channels] => [batch, in_height/2, in_width/2,
            # out_channels]
            convolved = pad_and_conv(layers[-1], out_channels, conv_args)
            output = tf.compat.v1.layers.batch_normalization(convolved, **batchnorm_args)
            if dropout > 0.0:
                output = tf.compat.v1.layers.dropout(output, rate=dropout, training=training_flag)
            layers.append(output)

    layer_specs = [
        # (ngf * 8, 0.5),
        # # decoder_6: [batch, 4, 4, ngf * 8 * 2] => [batch, 8, 8, ngf * 8 * 2]
        (ngf * 8, 0.2),
        # decoder_5: [batch, 8, 8, ngf * 8 * 2] => [batch, 16, 16, ngf * 8 * 2]
        (ngf * 4, 0.2),
        # decoder_4: [batch, 16, 16, ngf * 8 * 2] => [batch, 32, 32, ngf * 4 * 2]
        (ngf * 2, 0.2),
        # decoder_3: [batch, 32, 32, ngf * 4 * 2] => [batch, 64, 64, ngf * 2 * 2]
        (ngf * 2, 0.2),
        # decoder_2: [batch, 64, 64, ngf * 2 * 2] => [batch, 128, 128, ngf * 2]
    ]

    num_encoder_layers = len(layers)
    for decoder_layer, (out_channels, dropout) in enumerate(layer_specs):
        skip_layer = num_encoder_layers - decoder_layer - 1
        with tf.compat.v1.variable_scope("decoder_%d" % (skip_layer + 1)):
            if decoder_layer == 0:
                # first decoder layer doesn't have skip connections
                # since it is directly connected to the skip_layer
                input = layers[-1]
            else:
                input = tf.concat([layers[-1], layers[skip_layer]], axis=4)

            output = tf.compat.v1.layers.conv3d_transpose(input, out_channels, **deconv_args)
            output = tf.compat.v1.layers.batch_normalization(output, **batchnorm_args)

            if dropout > 0.0:
                output = tf.compat.v1.layers.dropout(output, rate=dropout, training=training_flag)

            layers.append(output)

    # decoder_1: [batch, 128, 128, ngf * 2] => [batch, 256, 256, generator_outputs_channels]
    with tf.compat.v1.variable_scope("decoder_1"):
        input = tf.concat([layers[-1], layers[0]], axis=4)
        this_args = deconv_args.copy()
        this_args["activation"] = None
        output = tf.compat.v1.layers.conv3d_transpose(input, 1, **this_args)
        layers.append(output)

    predictions = tf.squeeze(layers[-1], -1)

    if mode == tf.estimator.ModeKeys.PREDICT:
        temp = tf.compat.v1.saved_model.signature_constants
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            export_outputs={temp.DEFAULT_SERVING_SIGNATURE_DEF_KEY: PredictOutput(predictions)},
        )
    else:
        labels, filenames = labels
        loss = tf.losses.mean_squared_error(labels, predictions)

        # Add a scalar summary for the snapshot loss.
        # Create the gradient descent optimizer with the given learning rate.
        optimizer = tf.train.AdamOptimizer(params.learning_rate)
        # Use the optimizer to apply the gradients that minimize the loss
        # (and also increment the global step counter) as a single training step.
        extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(extra_update_ops):
            train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            loss=loss,
            train_op=train_op,
            export_outputs={"output": predictions},
        )

def _get_resize_arg(target_shape):
    """Get resizing arguments, as used by peaks2maps.

    .. deprecated:: 0.0.11
        `_get_resize_arg` will be removed in NiMARE 0.0.13.

    .. versionadded:: 0.0.1

    """
    mni_shape_mm = np.array([148.0, 184.0, 156.0])
    target_resolution_mm = np.ceil(mni_shape_mm / np.array(target_shape)).astype(np.int32)
    target_affine = np.array(
        [
            [4.0, 0.0, 0.0, -75.0],
            [0.0, 4.0, 0.0, -105.0],
            [0.0, 0.0, 4.0, -70.0],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    target_affine[0, 0] = target_resolution_mm[0]
    target_affine[1, 1] = target_resolution_mm[1]
    target_affine[2, 2] = target_resolution_mm[2]
    return target_affine, list(target_shape)

def _get_generator(contrasts_coordinates, target_shape, affine, skip_out_of_bounds=False):
    """Get generator, as used by peaks2maps."""

    def generator():
        for contrast in contrasts_coordinates:
            encoded_coords = np.zeros(list(target_shape))
            for real_pt in contrast:
                vox_pt = np.rint(nib.affines.apply_affine(npl.inv(affine), real_pt)).astype(int)
                if skip_out_of_bounds and (vox_pt[0] >= 32 or vox_pt[1] >= 32 or vox_pt[2] >= 32):
                    continue
                encoded_coords[vox_pt[0], vox_pt[1], vox_pt[2]] = 1
            yield (encoded_coords, encoded_coords)

    return generator

@due.dcite(
    references.PEAKS2MAPS,
    description="Transforms coordinates of peaks to unthresholded maps using a deep "
    "convolutional neural net.",
)
def compute_p2m_ma(
    contrasts_coordinates, skip_out_of_bounds=True, tf_verbosity_level=None, model_dir="auto"
):
    """Generate modeled activation (MA) maps using deep ConvNet model peaks2maps.

    .. deprecated:: 0.0.11
        `compute_p2m_ma` will be removed in NiMARE 0.0.13.

    Parameters
    ----------
    contrasts_coordinates : list of lists that are len == 3
        List of contrasts and their coordinates
    skip_out_of_bounds : bool, optional
        Remove coordinates outside of the bounding box of the peaks2maps model
    tf_verbosity_level : int
        Tensorflow verbosity logging level
    model_dir : str, optional
        Location of peaks2maps model. Default is "auto".

    Returns
    -------
    ma_values : array-like
        1d array of modeled activation values.
    """
    try:
        import tensorflow as tf
    except ImportError as e:
        if "No module named 'tensorflow'" in str(e):
            raise Exception(
                "tensorflow not installed - see https://www.tensorflow.org/install/ "
                "for instructions"
            )
        else:
            raise

    if tf_verbosity_level is None:
        tf_verbosity_level = tf.compat.v1.logging.FATAL
    target_shape = (32, 32, 32)
    affine, _ = _get_resize_arg(target_shape)
    tf.compat.v1.logging.set_verbosity(tf_verbosity_level)

    def generate_input_fn():
        dataset = tf.compat.v1.data.Dataset.from_generator(
            _get_generator(
                contrasts_coordinates, target_shape, affine, skip_out_of_bounds=skip_out_of_bounds
            ),
            (tf.float32, tf.float32),
            (tf.TensorShape(target_shape), tf.TensorShape(target_shape)),
        )
        dataset = dataset.batch(1)
        iterator = dataset.make_one_shot_iterator()
        return iterator.get_next()

    # download_peaks2maps_model expects None for "auto"
    model_dir = None if model_dir == "auto" else model_dir
    model_dir = download_peaks2maps_model(data_dir=model_dir)
    model = tf.estimator.Estimator(model_fn, model_dir=model_dir)

    results = model.predict(generate_input_fn)
    results = [result for result in results]
    assert len(results) == len(contrasts_coordinates), "returned %d" % len(results)

    niis = [nib.Nifti1Image(np.squeeze(result), affine) for result in results]
    return niis

def compute_kda_ma(
    shape,
    vox_dims,
    ijks,
    r,
    value=1.0,
    exp_idx=None,
    sum_overlap=False,
    memory_limit=None,
    memmap_filename=None,
):
    """Compute (M)KDA modeled activation (MA) map.

    .. versionchanged:: 0.0.8

        * [ENH] Add *memmap_filename* parameter for memory mapping arrays.

    .. versionadded:: 0.0.4

    Replaces the values around each focus in ijk with binary sphere.

    Parameters
    ----------
    shape : :obj:`tuple`
        Shape of brain image + buffer. Typically (91, 109, 91).
    vox_dims : array_like
        Size (in mm) of each dimension of a voxel.
    ijks : array-like
        Indices of foci. Each row is a coordinate, with the three columns
        corresponding to index in each of three dimensions.
    r : :obj:`int`
        Sphere radius, in mm.
    value : :obj:`int`
        Value for sphere.
    exp_idx : array_like
        Optional indices of experiments. If passed, must be of same length as
        ijks. Each unique value identifies all coordinates in ijk that come from
        the same experiment. If None passed, it is assumed that all coordinates
        come from the same experiment.
    sum_overlap : :obj:`bool`
        Whether to sum voxel values in overlapping spheres.
    memory_limit : :obj:`str` or None, optional
        Memory limit to apply to data. If None, no memory management will be applied.
        Otherwise, the memory limit will be used to (1) assign memory-mapped files and
        (2) restrict memory during array creation to the limit.
        Default is None.
    memmap_filename : :obj:`str`, optional
        If passed, use this file for memory mapping arrays

    Returns
    -------
    kernel_data : :obj:`numpy.array`
        3d or 4d array. If `exp_idx` is none, a 3d array in the same shape as
        the `shape` argument is returned. If `exp_idx` is passed, a 4d array
        is returned, where the first dimension has size equal to the number of
        unique experiments, and the remaining 3 dimensions are equal to `shape`.
    """
    squeeze = exp_idx is None
    if exp_idx is None:
        exp_idx = np.ones(len(ijks))

    uniq, exp_idx = np.unique(exp_idx, return_inverse=True)
    n_studies = len(uniq)

    kernel_shape = (n_studies,) + shape
    if memmap_filename:
        # Use a memmapped 4D array
        kernel_data = np.memmap(memmap_filename, dtype=type(value), mode="w+", shape=kernel_shape)
    else:
        kernel_data = np.zeros(kernel_shape, dtype=type(value))

    n_dim = ijks.shape[1]
    xx, yy, zz = [slice(-r // vox_dims[i], r // vox_dims[i] + 0.01, 1) for i in range(n_dim)]
    cube = np.vstack([row.ravel() for row in np.mgrid[xx, yy, zz]])
    kernel = cube[:, np.sum(np.dot(np.diag(vox_dims), cube) ** 2, 0) ** 0.5 <= r]

    if memory_limit:
        chunk_size = _determine_chunk_size(limit=memory_limit, arr=ijks[0])

    for i, peak in enumerate(ijks):
        sphere = np.round(kernel.T + peak)
        idx = (np.min(sphere, 1) >= 0) & (np.max(np.subtract(sphere, shape), 1) <= -1)
        sphere = sphere[idx, :].astype(int)
        exp = exp_idx[i]
        if sum_overlap:
            kernel_data[exp][tuple(sphere.T)] += value
        else:
            kernel_data[exp][tuple(sphere.T)] = value

        if memmap_filename and i % chunk_size == 0:
            # Write changes to disk
            kernel_data.flush()

    if squeeze:
        kernel_data = np.squeeze(kernel_data, axis=0)

    return kernel_data

def compute_ale_ma(shape, ijk, kernel):
    """Generate ALE modeled activation (MA) maps.

    Replaces the values around each focus in ijk with the contrast-specific
    kernel. Takes the element-wise maximum when looping through foci, which
    accounts for foci which are near to one another and may have overlapping
    kernels.

    Parameters
    ----------
    shape : tuple
        Shape of brain image + buffer. Typically (91, 109, 91) + (30, 30, 30).
    ijk : array-like
        Indices of foci. Each row is a coordinate, with the three columns
        corresponding to index in each of three dimensions.
    kernel : array-like
        3D array of smoothing kernel. Typically of shape (30, 30, 30).

    Returns
    -------
    ma_values : array-like
        1d array of modeled activation values.
    """
    ma_values = np.zeros(shape)
    mid = int(np.floor(kernel.shape[0] / 2.0))
    mid1 = mid + 1
    for j_peak in range(ijk.shape[0]):
        i, j, k = ijk[j_peak, :]
        xl = max(i - mid, 0)
        xh = min(i + mid1, ma_values.shape[0])
        yl = max(j - mid, 0)
        yh = min(j + mid1, ma_values.shape[1])
        zl = max(k - mid, 0)
        zh = min(k + mid1, ma_values.shape[2])
        xlk = mid - (i - xl)
        xhk = mid - (i - xh)
        ylk = mid - (j - yl)
        yhk = mid - (j - yh)
        zlk = mid - (k - zl)
        zhk = mid - (k - zh)

        if (
            (xl >= 0)
            & (xh >= 0)
            & (yl >= 0)
            & (yh >= 0)
            & (zl >= 0)
            & (zh >= 0)
            & (xlk >= 0)
            & (xhk >= 0)
            & (ylk >= 0)
            & (yhk >= 0)
            & (zlk >= 0)
            & (zhk >= 0)
        ):
            ma_values[xl:xh, yl:yh, zl:zh] = np.maximum(
                ma_values[xl:xh, yl:yh, zl:zh], kernel[xlk:xhk, ylk:yhk, zlk:zhk]
            )
    return ma_values

@due.dcite(references.ALE_KERNEL, description="Introduces sample size-dependent kernels to ALE.")
def get_ale_kernel(img, sample_size=None, fwhm=None):
    """Estimate 3D Gaussian and sigma (in voxels) for ALE kernel given sample size or fwhm."""
    if sample_size is not None and fwhm is not None:
        raise ValueError('Only one of "sample_size" and "fwhm" may be specified')
    elif sample_size is None and fwhm is None:
        raise ValueError('Either "sample_size" or "fwhm" must be provided')
    elif sample_size is not None:
        uncertain_templates = (
            5.7 / (2.0 * np.sqrt(2.0 / np.pi)) * np.sqrt(8.0 * np.log(2.0))
        )  # pylint: disable=no-member
        # Assuming 11.6 mm ED between matching points
        uncertain_subjects = (11.6 / (2 * np.sqrt(2 / np.pi)) * np.sqrt(8 * np.log(2))) / np.sqrt(
            sample_size
        )  # pylint: disable=no-member
        fwhm = np.sqrt(uncertain_subjects ** 2 + uncertain_templates ** 2)

    fwhm_vox = fwhm / np.sqrt(np.prod(img.header.get_zooms()))
    sigma_vox = (
        fwhm_vox * np.sqrt(2.0) / (np.sqrt(2.0 * np.log(2.0)) * 2.0)
    )  # pylint: disable=no-member

    data = np.zeros((31, 31, 31))
    mid = int(np.floor(data.shape[0] / 2.0))
    data[mid, mid, mid] = 1.0
    kernel = ndimage.filters.gaussian_filter(data, sigma_vox, mode="constant")

    # Crop kernel to drop surrounding zeros
    mn = np.min(np.where(kernel > np.spacing(1))[0])
    mx = np.max(np.where(kernel > np.spacing(1))[0])
    kernel = kernel[mn : mx + 1, mn : mx + 1, mn : mx + 1]
    mid = int(np.floor(data.shape[0] / 2.0))
    return sigma_vox, kernel