import os

import time

import numpy as np

import tensorflow as tf

import pickle

from multiprocessing.pool import ThreadPool

train_stats = (

    'Training statistics: \n'

    '\tLearning rate : {}\n'

    '\tBatch size    : {}\n'

    '\tEpoch number  : {}\n'

    '\tBackup every  : {}'

)

pool = ThreadPool()

def _save_ckpt(self, step, loss_profile):

    file = '{}-{}{}'

    model = self.meta['name']

    profile = file.format(model, step, '.profile')

    profile = os.path.join(self.FLAGS.backup, profile)

    with open(profile, 'wb') as profile_ckpt:

        pickle.dump(loss_profile, profile_ckpt)

    ckpt = file.format(model, step, '')

    ckpt = os.path.join(self.FLAGS.backup, ckpt)

    self.say('Checkpoint at step {}'.format(step))

    self.saver.save(self.sess, ckpt)

def train(self):

    loss_ph = self.framework.placeholders

    loss_mva = None;

    profile = list()

    batches = self.framework.shuffle()

    loss_op = self.framework.loss

    for i, (x_batch, datum) in enumerate(batches):

        if not i: self.say(train_stats.format(

            self.FLAGS.lr, self.FLAGS.batch,

            self.FLAGS.epoch, self.FLAGS.save

        ))

        feed_dict = {

            loss_ph[key]: datum[key]

            for key in loss_ph}

        feed_dict[self.inp] = x_batch

        feed_dict.update(self.feed)

        fetches = [self.train_op, loss_op]

        if self.FLAGS.summary:

            fetches.append(self.summary_op)

        fetched = self.sess.run(fetches, feed_dict)

        loss = fetched[1]

        if loss_mva is None: loss_mva = loss

        loss_mva = .9 * loss_mva + .1 * loss

        step_now = self.FLAGS.load + i + 1

        if self.FLAGS.summary:

            self.writer.add_summary(fetched[2], step_now)

        form = 'step {} - loss {} - moving ave loss {}'

        self.say(form.format(step_now, loss, loss_mva))

        profile += [(loss, loss_mva)]

        ckpt = (i + 1) % (self.FLAGS.save // self.FLAGS.batch)

        args = [step_now, profile]

        if not ckpt: _save_ckpt(self, *args)

    if ckpt: _save_ckpt(self, *args)

def return_predict(self, im):

    assert isinstance(im, np.ndarray), \

        'Image is not a np.ndarray'

    h, w, _ = im.shape

    im = self.framework.resize_input(im)

    this_inp = np.expand_dims(im, 0)

    feed_dict = {self.inp: this_inp}

    out = self.sess.run(self.out, feed_dict)[0]

    boxes = self.framework.findboxes(out)

    threshold = self.FLAGS.threshold

    boxesInfo = list()

    for box in boxes:

        tmpBox = self.framework.process_box(box, h, w, threshold)

        if tmpBox is None:

            continue

        boxesInfo.append({

            "label": tmpBox[4],

            "confidence": tmpBox[6],

            "topleft": {

                "x": tmpBox[0],

                "y": tmpBox[2]},

            "bottomright": {

                "x": tmpBox[1],

                "y": tmpBox[3]}

        })

    return boxesInfo

import math

def predict(self):

    inp_path = self.FLAGS.imgdir

    all_inps = os.listdir(inp_path)

    all_inps = [i for i in all_inps if self.framework.is_inp(i)]

    if not all_inps:

        msg = 'Failed to find any images in {} .'

        exit('Error: {}'.format(msg.format(inp_path)))

    batch = min(self.FLAGS.batch, len(all_inps))

    # predict in batches

    n_batch = int(math.ceil(len(all_inps) / batch))

    for j in range(n_batch):

        from_idx = j * batch

        to_idx = min(from_idx + batch, len(all_inps))

        # collect images input in the batch

        this_batch = all_inps[from_idx:to_idx]

        inp_feed = pool.map(lambda inp: (

            np.expand_dims(self.framework.preprocess(

                os.path.join(inp_path, inp)), 0)), this_batch)

        # Feed to the net

        feed_dict = {self.inp: np.concatenate(inp_feed, 0)}

        self.say('Forwarding {} inputs ...'.format(len(inp_feed)))

        start = time.time()

        out = self.sess.run(self.out, feed_dict)

        stop = time.time();

        last = stop - start

        self.say('Total time = {}s / {} inps = {} ips'.format(

            last, len(inp_feed), len(inp_feed) / last))

        # Post processing

        self.say('Post processing {} inputs ...'.format(len(inp_feed)))

        start = time.time()

        pool.map(lambda p: (lambda i, prediction:

                            self.framework.postprocess(

                                prediction, os.path.join(inp_path, this_batch[i])))(*p),

                 enumerate(out))

        stop = time.time();

        last = stop - start

        # Timing

        self.say('Total time = {}s / {} inps = {} ips'.format(

            last, len(inp_feed), len(inp_feed) / last))