import tensorflow as tf
import os
from .. import dark
import numpy as np
from os.path import basename

tf = tf.compat.v1
tf.disable_v2_behavior()


class loader(object):
    """
    interface to work with both .weights and .ckpt files
    in loading / recollecting / resolving mode
    """
    VAR_LAYER = ['convolutional', 'connected', 'local',
                 'select', 'conv-select',
                 'extract', 'conv-extract']

    def __init__(self, *args):
        self.src_key = list()
        self.vals = list()
        self.load(*args)

    def __call__(self, key):
        for idx in range(len(key)):
            val = self.find(key, idx)
            if val is not None: return val
        return None

    def find(self, key, idx):
        up_to = min(len(self.src_key), 4)
        for i in range(up_to):
            key_b = self.src_key[i]
            if key_b[idx:] == key[idx:]:
                return self.yields(i)
        return None

    def yields(self, idx):
        del self.src_key[idx]
        temp = self.vals[idx]
        del self.vals[idx]
        return temp


class weights_loader(loader):
    """one who understands .weights files"""

    _W_ORDER = dict({  # order of param flattened into .weights file
        'convolutional': [
            'biases', 'gamma', 'moving_mean', 'moving_variance', 'kernel'
        ],
        'connected': ['biases', 'weights'],
        'local': ['biases', 'kernels']
    })

    def load(self, path, src_layers):
        self.src_layers = src_layers
        walker = weights_walker(path)

        for i, layer in enumerate(src_layers):
            if layer.type not in self.VAR_LAYER: continue
            self.src_key.append([layer])

            if walker.eof:
                new = None
            else:
                args = layer.signature
                new = dark.darknet.create_darkop(*args)
            self.vals.append(new)

            if new is None: continue
            order = self._W_ORDER[new.type]
            for par in order:
                if par not in new.wshape: continue
                val = walker.walk(new.wsize[par])
                new.w[par] = val
            new.finalize(walker.transpose)

        if walker.path is not None:
            assert walker.offset == walker.size, \
                'expect {} bytes, found {}'.format(
                    walker.offset, walker.size)
            print('Successfully identified {} bytes'.format(
                walker.offset))


class checkpoint_loader(loader):
    """
    one who understands .ckpt files, very much
    """

    def load(self, ckpt, ignore):
        meta = ckpt + '.meta'
        with tf.Graph().as_default() as graph:
            with tf.Session().as_default() as sess:
                saver = tf.train.import_meta_graph(meta)
                saver.restore(sess, ckpt)
                for var in tf.global_variables():
                    name = var.name.split(':')[0]
                    packet = [name, var.get_shape().as_list()]
                    self.src_key += [packet]
                    self.vals += [var.eval(sess)]


def create_loader(path, cfg=None):
    if path is None:
        load_type = weights_loader
    elif '.weights' in path:
        load_type = weights_loader
    else:
        load_type = checkpoint_loader

    return load_type(path, cfg)


class weights_walker(object):
    """incremental reader of float32 binary files"""

    def __init__(self, path):
        self.eof = False  # end of file
        self.path = path  # current pos
        if path is None:
            self.eof = True
            return
        else:
            self.size = os.path.getsize(path)  # save the path
            major, minor, revision, seen = np.memmap(path,
                                                     shape=(), mode='r', offset=0,
                                                     dtype='({})i4,'.format(4))
            self.transpose = major > 1000 or minor > 1000
            self.offset = 16

    def walk(self, size):
        if self.eof: return None
        end_point = self.offset + 4 * size
        assert end_point <= self.size, \
            'Over-read {}'.format(self.path)

        float32_1D_array = np.memmap(
            self.path, shape=(), mode='r',
            offset=self.offset,
            dtype='({})float32,'.format(size)
        )

        self.offset = end_point
        if end_point == self.size:
            self.eof = True
        return float32_1D_array


def model_name(file_path):
    file_name = basename(file_path)
    ext = str()
    if '.' in file_name:  # exclude extension
        file_name = file_name.split('.')
        ext = file_name[-1]
        file_name = '.'.join(file_name[:-1])
    if ext == str() or ext == 'meta':  # ckpt file
        file_name = file_name.split('-')
        num = int(file_name[-1])
        return '-'.join(file_name[:-1])
    if ext == 'weights':
        return file_name