import tensorflow as tf

import time

from . import help

from . import flow

from .ops import op_create, identity

from .ops import HEADER, LINE

from .framework import create_framework

from ..dark.darknet import Darknet

import json

import os

tf = tf.compat.v1

tf.disable_v2_behavior()

class TFNet(object):

    _TRAINER = dict({

        'rmsprop': tf.train.RMSPropOptimizer,

        'adadelta': tf.train.AdadeltaOptimizer,

        'adagrad': tf.train.AdagradOptimizer,

        'adagradDA': tf.train.AdagradDAOptimizer,

        'momentum': tf.train.MomentumOptimizer,

        'adam': tf.train.AdamOptimizer,

        'ftrl': tf.train.FtrlOptimizer,

        'sgd': tf.train.GradientDescentOptimizer

    })

    # imported methods

    _get_fps = help._get_fps

    say = help.say

    train = flow.train

    camera = help.camera

    predict = flow.predict

    return_predict = flow.return_predict

    to_darknet = help.to_darknet

    build_train_op = help.build_train_op

    load_from_ckpt = help.load_from_ckpt

    def __init__(self, FLAGS, darknet=None):

        self.ntrain = 0

        if isinstance(FLAGS, dict):

            from ..defaults import argHandler

            newFLAGS = argHandler()

            newFLAGS.setDefaults()

            newFLAGS.update(FLAGS)

            FLAGS = newFLAGS

        self.FLAGS = FLAGS

        if self.FLAGS.pbLoad and self.FLAGS.metaLoad:

            self.say('\nLoading from .pb and .meta')

            self.graph = tf.Graph()

            device_name = FLAGS.gpuName \

                if FLAGS.gpu > 0.0 else None

            with tf.device(device_name):

                with self.graph.as_default() as g:

                    self.build_from_pb()

            return

        if darknet is None:

            darknet = Darknet(FLAGS)

            self.ntrain = len(darknet.layers)

        self.darknet = darknet

        args = [darknet.meta, FLAGS]

        self.num_layer = len(darknet.layers)

        self.framework = create_framework(*args)

        self.meta = darknet.meta

        self.say('\nBuilding net ...')

        start = time.time()

        self.graph = tf.Graph()

        device_name = FLAGS.gpuName \

            if FLAGS.gpu > 0.0 else None

        with tf.device(device_name):

            with self.graph.as_default() as g:

                self.build_forward()

                self.setup_meta_ops()

        self.say('Finished in {}s\n'.format(

            time.time() - start))

    def build_from_pb(self):

        with tf.gfile.FastGFile(self.FLAGS.pbLoad, "rb") as f:

            graph_def = tf.GraphDef()

            graph_def.ParseFromString(f.read())

        tf.import_graph_def(

            graph_def,

            name=""

        )

        with open(self.FLAGS.metaLoad, 'r') as fp:

            self.meta = json.load(fp)

        self.framework = create_framework(self.meta, self.FLAGS)

        # Placeholders

        self.inp = tf.get_default_graph().get_tensor_by_name('input:0')

        self.feed = dict()  # other placeholders

        self.out = tf.get_default_graph().get_tensor_by_name('output:0')

        self.setup_meta_ops()

    def build_forward(self):

        verbalise = self.FLAGS.verbalise

        # Placeholders

        inp_size = [None] + self.meta['inp_size']

        self.inp = tf.placeholder(tf.float32, inp_size, 'input')

        self.feed = dict()  # other placeholders

        # Build the forward pass

        state = identity(self.inp)

        roof = self.num_layer - self.ntrain

        self.say(HEADER, LINE)

        for i, layer in enumerate(self.darknet.layers):

            scope = '{}-{}'.format(str(i), layer.type)

            args = [layer, state, i, roof, self.feed]

            state = op_create(*args)

            mess = state.verbalise()

            self.say(mess)

        self.say(LINE)

        self.top = state

        self.out = tf.identity(state.out, name='output')

    def setup_meta_ops(self):

        cfg = dict({

            'allow_soft_placement': False,

            'log_device_placement': False

        })

        utility = min(self.FLAGS.gpu, 1.)

        if utility > 0.0:

            self.say('GPU mode with {} usage'.format(utility))

            cfg['gpu_options'] = tf.GPUOptions(

                per_process_gpu_memory_fraction=utility)

            cfg['allow_soft_placement'] = True

        else:

            self.say('Running entirely on CPU')

            cfg['device_count'] = {'GPU': 0}

        if self.FLAGS.train: self.build_train_op()

        if self.FLAGS.summary:

            self.summary_op = tf.summary.merge_all()

            self.writer = tf.summary.FileWriter(self.FLAGS.summary + 'train')

        self.sess = tf.Session(config=tf.ConfigProto(**cfg))

        self.sess.run(tf.global_variables_initializer())

        if not self.ntrain: return

        self.saver = tf.train.Saver(tf.global_variables(),

                                    max_to_keep=self.FLAGS.keep)

        if self.FLAGS.load != 0: self.load_from_ckpt()

        if self.FLAGS.summary:

            self.writer.add_graph(self.sess.graph)

    def savepb(self):

        """

        Create a standalone const graph def that

        C++ can load and run.

        """

        darknet_pb = self.to_darknet()

        flags_pb = self.FLAGS

        flags_pb.verbalise = False

        flags_pb.train = False

        # rebuild another tfnet. all const.

        tfnet_pb = TFNet(flags_pb, darknet_pb)

        tfnet_pb.sess = tf.Session(graph=tfnet_pb.graph)

        # tfnet_pb.predict() # uncomment for unit testing

        name = 'built_graph/{}.pb'.format(self.meta['name'])

        os.makedirs(os.path.dirname(name), exist_ok=True)

        # Save dump of everything in meta

        with open('built_graph/{}.meta'.format(self.meta['name']), 'w') as fp:

            json.dump(self.meta, fp)

        self.say('Saving const graph def to {}'.format(name))

        graph_def = tfnet_pb.sess.graph_def

        tf.train.write_graph(graph_def, './', name, False)