"""

WARNING: spaghetti code.

"""

import numpy as np

import pickle

import os

def parser(model):

    """

    Read the .cfg file to extract layers into `layers`

    as well as model-specific parameters into `meta`

    """

    def _parse(l, i=1):

        return l.split('=')[i].strip()

    with open(model, 'rb') as f:

        lines = f.readlines()

    lines = [line.decode() for line in lines]

    meta = dict();

    layers = list()  # will contains layers' info

    h, w, c = [int()] * 3;

    layer = dict()

    for line in lines:

        line = line.strip()

        line = line.split('#')[0]

        if '[' in line:

            if layer != dict():

                if layer['type'] == '[net]':

                    h = layer['height']

                    w = layer['width']

                    c = layer['channels']

                    meta['net'] = layer

                else:

                    if layer['type'] == '[crop]':

                        h = layer['crop_height']

                        w = layer['crop_width']

                    layers += [layer]

            layer = {'type': line}

        else:

            try:

                i = float(_parse(line))

                if i == int(i): i = int(i)

                layer[line.split('=')[0].strip()] = i

            except:

                try:

                    key = _parse(line, 0)

                    val = _parse(line, 1)

                    layer[key] = val

                except:

                    'banana ninja yadayada'

    meta.update(layer)  # last layer contains meta info

    if 'anchors' in meta:

        splits = meta['anchors'].split(',')

        anchors = [float(x.strip()) for x in splits]

        meta['anchors'] = anchors

    meta['model'] = model  # path to cfg, not model name

    meta['inp_size'] = [h, w, c]

    return layers, meta

def cfg_yielder(model, binary):

    """

    yielding each layer information to initialize `layer`

    """

    layers, meta = parser(model);

    yield meta;

    h, w, c = meta['inp_size'];

    l = w * h * c

    # Start yielding

    flat = False  # flag for 1st dense layer

    conv = '.conv.' in model

    for i, d in enumerate(layers):

        # -----------------------------------------------------

        if d['type'] == '[crop]':

            yield ['crop', i]

        # -----------------------------------------------------

        elif d['type'] == '[local]':

            n = d.get('filters', 1)

            size = d.get('size', 1)

            stride = d.get('stride', 1)

            pad = d.get('pad', 0)

            activation = d.get('activation', 'logistic')

            w_ = (w - 1 - (1 - pad) * (size - 1)) // stride + 1

            h_ = (h - 1 - (1 - pad) * (size - 1)) // stride + 1

            yield ['local', i, size, c, n, stride,

                   pad, w_, h_, activation]

            if activation != 'linear': yield [activation, i]

            w, h, c = w_, h_, n

            l = w * h * c

        # -----------------------------------------------------

        elif d['type'] == '[convolutional]':

            n = d.get('filters', 1)

            size = d.get('size', 1)

            stride = d.get('stride', 1)

            pad = d.get('pad', 0)

            padding = d.get('padding', 0)

            if pad: padding = size // 2

            activation = d.get('activation', 'logistic')

            batch_norm = d.get('batch_normalize', 0) or conv

            yield ['convolutional', i, size, c, n,

                   stride, padding, batch_norm,

                   activation]

            if activation != 'linear': yield [activation, i]

            w_ = (w + 2 * padding - size) // stride + 1

            h_ = (h + 2 * padding - size) // stride + 1

            w, h, c = w_, h_, n

            l = w * h * c

        # -----------------------------------------------------

        elif d['type'] == '[maxpool]':

            stride = d.get('stride', 1)

            size = d.get('size', stride)

            padding = d.get('padding', (size - 1) // 2)

            yield ['maxpool', i, size, stride, padding]

            w_ = (w + 2 * padding) // d['stride']

            h_ = (h + 2 * padding) // d['stride']

            w, h = w_, h_

            l = w * h * c

        # -----------------------------------------------------

        elif d['type'] == '[avgpool]':

            flat = True;

            l = c

            yield ['avgpool', i]

        # -----------------------------------------------------

        elif d['type'] == '[softmax]':

            yield ['softmax', i, d['groups']]

        # -----------------------------------------------------

        elif d['type'] == '[connected]':

            if not flat:

                yield ['flatten', i]

                flat = True

            activation = d.get('activation', 'logistic')

            yield ['connected', i, l, d['output'], activation]

            if activation != 'linear': yield [activation, i]

            l = d['output']

        # -----------------------------------------------------

        elif d['type'] == '[dropout]':

            yield ['dropout', i, d['probability']]

        # -----------------------------------------------------

        elif d['type'] == '[select]':

            if not flat:

                yield ['flatten', i]

                flat = True

            inp = d.get('input', None)

            if type(inp) is str:

                file = inp.split(',')[0]

                layer_num = int(inp.split(',')[1])

                with open(file, 'rb') as f:

                    profiles = pickle.load(f, encoding='latin1')[0]

                layer = profiles[layer_num]

            else:

                layer = inp

            activation = d.get('activation', 'logistic')

            d['keep'] = d['keep'].split('/')

            classes = int(d['keep'][-1])

            keep = [int(c) for c in d['keep'][0].split(',')]

            keep_n = len(keep)

            train_from = classes * d['bins']

            for count in range(d['bins'] - 1):

                for num in keep[-keep_n:]:

                    keep += [num + classes]

            k = 1

            while layers[i - k]['type'] not in ['[connected]', '[extract]']:

                k += 1

                if i - k < 0:

                    break

            if i - k < 0:

                l_ = l

            elif layers[i - k]['type'] == 'connected':

                l_ = layers[i - k]['output']

            else:

                l_ = layers[i - k].get('old', [l])[-1]

            yield ['select', i, l_, d['old_output'],

                   activation, layer, d['output'],

                   keep, train_from]

            if activation != 'linear': yield [activation, i]

            l = d['output']

        # -----------------------------------------------------

        elif d['type'] == '[conv-select]':

            n = d.get('filters', 1)

            size = d.get('size', 1)

            stride = d.get('stride', 1)

            pad = d.get('pad', 0)

            padding = d.get('padding', 0)

            if pad: padding = size // 2

            activation = d.get('activation', 'logistic')

            batch_norm = d.get('batch_normalize', 0) or conv

            d['keep'] = d['keep'].split('/')

            classes = int(d['keep'][-1])

            keep = [int(x) for x in d['keep'][0].split(',')]

            segment = classes + 5

            assert n % segment == 0, \

                'conv-select: segment failed'

            bins = n // segment

            keep_idx = list()

            for j in range(bins):

                offset = j * segment

                for k in range(5):

                    keep_idx += [offset + k]

                for k in keep:

                    keep_idx += [offset + 5 + k]

            w_ = (w + 2 * padding - size) // stride + 1

            h_ = (h + 2 * padding - size) // stride + 1

            c_ = len(keep_idx)

            yield ['conv-select', i, size, c, n,

                   stride, padding, batch_norm,

                   activation, keep_idx, c_]

            w, h, c = w_, h_, c_

            l = w * h * c

        # -----------------------------------------------------

        elif d['type'] == '[conv-extract]':

            file = d['profile']

            with open(file, 'rb') as f:

                profiles = pickle.load(f, encoding='latin1')[0]

            inp_layer = None

            inp = d['input']

            out = d['output']

            inp_layer = None

            if inp >= 0:

                inp_layer = profiles[inp]

            if inp_layer is not None:

                assert len(inp_layer) == c, \

                    'Conv-extract does not match input dimension'

            out_layer = profiles[out]

            n = d.get('filters', 1)

            size = d.get('size', 1)

            stride = d.get('stride', 1)

            pad = d.get('pad', 0)

            padding = d.get('padding', 0)

            if pad: padding = size // 2

            activation = d.get('activation', 'logistic')

            batch_norm = d.get('batch_normalize', 0) or conv

            k = 1

            find = ['[convolutional]', '[conv-extract]']

            while layers[i - k]['type'] not in find:

                k += 1

                if i - k < 0: break

            if i - k >= 0:

                previous_layer = layers[i - k]

                c_ = previous_layer['filters']

            else:

                c_ = c

            yield ['conv-extract', i, size, c_, n,

                   stride, padding, batch_norm,

                   activation, inp_layer, out_layer]

            if activation != 'linear': yield [activation, i]

            w_ = (w + 2 * padding - size) // stride + 1

            h_ = (h + 2 * padding - size) // stride + 1

            w, h, c = w_, h_, len(out_layer)

            l = w * h * c

        # -----------------------------------------------------

        elif d['type'] == '[extract]':

            if not flat:

                yield ['flatten', i]

                flat = True

            activation = d.get('activation', 'logistic')

            file = d['profile']

            with open(file, 'rb') as f:

                profiles = pickle.load(f, encoding='latin1')[0]

            inp_layer = None

            inp = d['input']

            out = d['output']

            if inp >= 0:

                inp_layer = profiles[inp]

            out_layer = profiles[out]

            old = d['old']

            old = [int(x) for x in old.split(',')]

            if inp_layer is not None:

                if len(old) > 2:

                    h_, w_, c_, n_ = old

                    new_inp = list()

                    for p in range(c_):

                        for q in range(h_):

                            for r in range(w_):

                                if p not in inp_layer:

                                    continue

                                new_inp += [r + w * (q + h * p)]

                    inp_layer = new_inp

                    old = [h_ * w_ * c_, n_]

                assert len(inp_layer) == l, \

                    'Extract does not match input dimension'

            d['old'] = old

            yield ['extract', i] + old + [activation] + [inp_layer, out_layer]

            if activation != 'linear': yield [activation, i]

            l = len(out_layer)

        # -----------------------------------------------------

        elif d['type'] == '[route]':  # add new layer here

            routes = d['layers']

            if type(routes) is int:

                routes = [routes]

            else:

                routes = [int(x.strip()) for x in routes.split(',')]

            routes = [i + x if x < 0 else x for x in routes]

            for j, x in enumerate(routes):

                lx = layers[x];

                xtype = lx['type']

                _size = lx['_size'][:3]

                if j == 0:

                    h, w, c = _size

                else:

                    h_, w_, c_ = _size

                    assert w_ == w and h_ == h, \

                        'Routing incompatible conv sizes'

                    c += c_

            yield ['route', i, routes]

            l = w * h * c

        # -----------------------------------------------------

        elif d['type'] == '[reorg]':

            stride = d.get('stride', 1)

            yield ['reorg', i, stride]

            w = w // stride;

            h = h // stride;

            c = c * (stride ** 2)

            l = w * h * c

        # -----------------------------------------------------

        else:

            exit('Layer {} not implemented'.format(d['type']))

        d['_size'] = list([h, w, c, l, flat])

    if not flat:

        meta['out_size'] = [h, w, c]

    else:

        meta['out_size'] = l