from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

"""Language model agent.

Agent outputs code in a sequence just like a language model. Can be trained
as a language model or using RL, or a combination of the two.
"""

from collections import namedtuple
from math import exp, log
import time

import numpy as np
from six.moves import xrange
import tensorflow as tf

import cibi.rollout as rollout_lib  # brain coder
from cibi import utils
from cibi import bf
from cibi.codebase import make_dev_codebase, make_codebase_like, Codebase

import logging
logger = logging.getLogger(f'cibi.{__file__}')

# Experiments in the ICLR 2018 paper used reduce_sum instead of reduce_mean for
# some losses. We make all loses be batch_size independent, and multiply the
# changed losses by 64, which was the fixed batch_size when the experiments
# where run. The loss hyperparameters still match what is reported in the paper.
MAGIC_LOSS_MULTIPLIER = 64

def rshift_time(tensor_2d, fill):
  """Right shifts a 2D tensor along the time dimension (axis-1)."""
  dim_0 = tf.shape(tensor_2d)[0]
  fill_tensor = tf.fill([dim_0, 1], fill)
  return tf.concat([fill_tensor, tensor_2d[:, :-1]], axis=1)

ReflectionResult = namedtuple(
    'ReflectionResult',
    ['global_step', 'global_npe', 'summaries_list', 'gradients_dict'])

LMConfig = namedtuple(
  'NeuralParams',
  ['policy_lstm_sizes', 'obs_embedding_size', 'grad_clip_threshold',
   'param_init_factor', 'lr', 'pi_loss_hparam', 'entropy_beta', 'regularizer',
   'softmax_tr', 'optimizer', 'topk', 'topk_loss_hparam', 'topk_batch_size',
   'ema_baseline_decay', 'alpha', 'iw_normalize', 'batch_size', 'timestep_limit']
)

default_config = {
  'batch_size': 4,
  'policy_lstm_sizes': [50,50],
  'obs_embedding_size': 15,
  'grad_clip_threshold': 10.0,
  'param_init_factor':1.0,
  'lr':5e-5,
  'pi_loss_hparam':1.0,
  'entropy_beta':1e-2,
  'regularizer':0.0,
  'softmax_tr':1.0,  # Reciprocal temperature.
  'optimizer':'rmsprop',  # 'adam', 'sgd', 'rmsprop'
  'topk':5,  # Top-k unique codes will be stored.
  'topk_loss_hparam':0.5,  # off policy loss multiplier.
  # Uniformly sample this many episodes from topk buffer per batch.
  # If topk is 0, this has no effect.
  'topk_batch_size':1,
  'timestep_limit': 1024,
  # Exponential moving average baseline for REINFORCE.
  # If zero, A2C is used.
  # If non-zero, should be close to 1, like .99, .999, etc.
  'ema_baseline_decay':0.99,
  # Replay probability. 1 : always replay, 0 : always on policy.
  'alpha':0.75,
  # Whether to normalize importance weights in each minibatch.
  'iw_normalize':True
}

def list_(x):
  return [] if x is None else list(x)

def make_optimizer(kind, lr):
  if kind == 'sgd':
    return tf.train.GradientDescentOptimizer(lr)
  elif kind == 'adam':
    return tf.train.AdamOptimizer(lr)
  elif kind == 'rmsprop':
    return tf.train.RMSPropOptimizer(learning_rate=lr, decay=0.99)
  else:
    raise ValueError('Optimizer type "%s" not recognized.' % kind)

class LinearWrapper(tf.contrib.rnn.RNNCell):
  """RNNCell wrapper that adds a linear layer to the output."""

  def __init__(self, cell, output_size, dtype=tf.float32, suppress_index=None):
    self.cell = cell
    self._output_size = output_size
    self._dtype = dtype
    self._suppress_index = suppress_index
    self.smallest_float = -2.4e38

  def __call__(self, inputs, state, scope=None):
    with tf.variable_scope(type(self).__name__):
      outputs, state = self.cell(inputs, state, scope=scope)
      logits = tf.matmul(
          outputs,
          tf.get_variable('w_output',
                          [self.cell.output_size, self.output_size],
                          dtype=self._dtype))
      if self._suppress_index is not None:
        # Replace the target index with -inf, so that it never gets selected.
        batch_size = tf.shape(logits)[0]
        logits = tf.concat(
            [logits[:, :self._suppress_index],
             tf.fill([batch_size, 1], self.smallest_float),
             logits[:, self._suppress_index + 1:]],
            axis=1)

    return logits, state

  @property
  def output_size(self):
    return self._output_size

  @property
  def state_size(self):
    return self.cell.state_size

  def zero_state(self, batch_size, dtype):
    return self.cell.zero_state(batch_size, dtype)

class AttrDict(dict):
  """Dict with attributes as keys.

  https://stackoverflow.com/a/14620633
  """

  def __init__(self, *args, **kwargs):
    super(AttrDict, self).__init__(*args, **kwargs)
    self.__dict__ = self

class LanguageModel:
  """Language model agent."""
  def __init__(self, language, config={},
               global_best_reward_fn=None,
               program_count=None,
               do_iw_summaries=False,
               dtype=tf.float32,
               verbose_level=0,
               is_local=True):
    if type(config) == LMConfig:
      self.config = config
    elif type(config) == dict:
      self.config = LMConfig(**{**default_config, **config})
      config = None # prevent accidentaly using partial config
    else:
      raise ValueError('Unsupported configuration format. Should be a dict or LMConfigs')

    self.verbose_level = verbose_level
    self.global_best_reward_fn = global_best_reward_fn
    self.parent_scope_name = tf.get_variable_scope().name
    self.dtype = dtype
    self.is_local = is_local

    self.top_reward = 0.0
    self.embeddings_trainable = True

    self.action_space = len(language['alphabet'])
    self.observation_space = len(language['alphabet'])
    self.int_to_char = language['int_to_char']
    self.char_to_int = language['char_to_int']
    self.eos_char = language['eos_char']
    self.eos_int = language['eos_int']

    self.no_op = tf.no_op()

    self.learning_rate = tf.constant(
        self.config.lr, dtype=dtype, name='learning_rate')
    self.initializer = tf.contrib.layers.variance_scaling_initializer(
        factor=self.config.param_init_factor,
        mode='FAN_AVG',
        uniform=True,
        dtype=dtype)  # TF's default initializer.
    tf.get_variable_scope().set_initializer(self.initializer)

    logger.info('Using exponential moving average REINFORCE baselines.')
    self.ema_by_len = [0.0] * self.config.timestep_limit

    # Top-k
    topk_params = (self.config.topk, self.config.topk_loss_hparam, self.config.topk_batch_size)
    if any(p != 0 for p in topk_params):
      assert all(p > 0 for p in topk_params)

    # Experience replay.
    if self.config.alpha == 0:
      self.num_replay_per_batch = 0
    else:
      self.num_replay_per_batch = int(self.config.batch_size * self.config.alpha)
    self.num_on_policy_per_batch = (
        self.config.batch_size - self.num_replay_per_batch)
    self.replay_alpha = (
        self.num_replay_per_batch / float(self.config.batch_size))
    logger.info('num_replay_per_batch: %d', self.num_replay_per_batch)
    logger.info('num_on_policy_per_batch: %d', self.num_on_policy_per_batch)
    logger.info('replay_alpha: %s', self.replay_alpha)

    if program_count is not None:
      self.program_count = program_count
      self.program_count_add_ph = tf.placeholder(
          tf.int64, [], 'program_count_add_ph')
      self.program_count_add_op = self.program_count.assign_add(
          self.program_count_add_ph)

    ################################
    # RL policy network #
    ################################
    batch_size = self.config.batch_size
    logger.info('batch_size: %d', batch_size)

    self.policy_cell = LinearWrapper(
        tf.contrib.rnn.MultiRNNCell(
            [tf.contrib.rnn.BasicLSTMCell(cell_size)
             for cell_size in self.config.policy_lstm_sizes]),
        self.action_space,
        dtype=dtype,
        suppress_index=None)

    obs_embedding_scope = 'obs_embed'
    with tf.variable_scope(
        obs_embedding_scope,
        initializer=tf.random_uniform_initializer(minval=-1.0, maxval=1.0)):
      obs_embeddings = tf.get_variable(
          'embeddings',
          [self.observation_space, self.config.obs_embedding_size],
          dtype=dtype, trainable=self.embeddings_trainable)
      self.obs_embeddings = obs_embeddings

    ################################
    # RL policy and value networks #
    ################################

    initial_state = tf.fill([batch_size], self.eos_int)
    def loop_fn(loop_time, cell_output, cell_state, loop_state):
      """Function called by tf.nn.raw_rnn to instantiate body of the while_loop.

      See https://www.tensorflow.org/api_docs/python/tf/nn/raw_rnn for more
      information.

      When time is 0, and cell_output, cell_state, loop_state are all None,
      `loop_fn` will create the initial input, internal cell state, and loop
      state. When time > 0, `loop_fn` will operate on previous cell output,
      state, and loop state.

      Args:
        loop_time: A scalar tensor holding the current timestep (zero based
            counting).
        cell_output: Output of the raw_rnn cell at the current timestep.
        cell_state: Cell internal state at the current timestep.
        loop_state: Additional loop state. These tensors were returned by the
            previous call to `loop_fn`.

      Returns:
        elements_finished: Bool tensor of shape [batch_size] which marks each
            sequence in the batch as being finished or not finished.
        next_input: A tensor containing input to be fed into the cell at the
            next timestep.
        next_cell_state: Cell internal state to be fed into the cell at the
            next timestep.
        emit_output: Tensor to be added to the TensorArray returned by raw_rnn
            as output from the while_loop.
        next_loop_state: Additional loop state. These tensors will be fed back
            into the next call to `loop_fn` as `loop_state`.
      """
      if cell_output is None:  # 0th time step.
        next_cell_state = self.policy_cell.zero_state(batch_size, dtype)
        elements_finished = tf.zeros([batch_size], tf.bool)
        output_lengths = tf.ones([batch_size], dtype=tf.int32)
        next_input = tf.gather(obs_embeddings, initial_state)
        emit_output = None
        next_loop_state = (
            tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True),
            output_lengths,
            elements_finished
        )
      else:
        scaled_logits = cell_output * self.config.softmax_tr  # Scale temperature.
        prev_chosen, prev_output_lengths, prev_elements_finished = loop_state
        next_cell_state = cell_state
        chosen_outputs = tf.to_int32(tf.where(
            tf.logical_not(prev_elements_finished),
            tf.multinomial(logits=scaled_logits, num_samples=1)[:, 0],
            tf.zeros([batch_size], dtype=tf.int64)))
        elements_finished = tf.logical_or(
            tf.equal(chosen_outputs, self.eos_int),
            loop_time >= self.config.timestep_limit)
        output_lengths = tf.where(
            elements_finished,
            prev_output_lengths,
            # length includes EOS token. empty seq has len 1.
            tf.tile(tf.expand_dims(loop_time + 1, 0), [batch_size])
        )
        next_input = tf.gather(obs_embeddings, chosen_outputs)
        emit_output = scaled_logits
        next_loop_state = (prev_chosen.write(loop_time - 1, chosen_outputs),
                           output_lengths,
                           tf.logical_or(prev_elements_finished,
                                         elements_finished))
      return (elements_finished, next_input, next_cell_state, emit_output,
              next_loop_state)

    with tf.variable_scope('policy'):
      (decoder_outputs_ta,
       _,  # decoder_state
       (sampled_output_ta, output_lengths, _)) = tf.nn.raw_rnn(
           cell=self.policy_cell,
           loop_fn=loop_fn)
    policy_logits = tf.transpose(decoder_outputs_ta.stack(), (1, 0, 2),
                                 name='policy_logits')
    sampled_tokens = tf.transpose(sampled_output_ta.stack(), (1, 0),
                                  name='sampled_tokens')
    # for sampling actions from the agent, and which told tensors for doing
    # gradient updates on the agent.
    self.sampled_batch = AttrDict(
        logits=policy_logits,
        tokens=sampled_tokens,
        episode_lengths=output_lengths,
        probs=tf.nn.softmax(policy_logits),
        log_probs=tf.nn.log_softmax(policy_logits))

    # adjusted_lengths can be less than the full length of each episode.
    # Use this to train on only part of an episode (starting from t=0).
    self.adjusted_lengths = tf.placeholder(
        tf.int32, [None], name='adjusted_lengths')
    self.policy_multipliers = tf.placeholder(
        dtype,
        [None, None],
        name='policy_multipliers')
    # Empirical value, i.e. discounted sum of observed future rewards from each
    # time step in the episode.
    self.empirical_values = tf.placeholder(
        dtype,
        [None, None],
        name='empirical_values')

    # Off-policy training. Just add supervised loss to the RL loss.
    self.off_policy_targets = tf.placeholder(
        tf.int32,
        [None, None],
        name='off_policy_targets')
    self.off_policy_target_lengths = tf.placeholder(
        tf.int32, [None], name='off_policy_target_lengths')

    self.actions = tf.placeholder(tf.int32, [None, None], name='actions')
    # Add SOS to beginning of the sequence.
    inputs = rshift_time(self.actions, fill=self.eos_int)
    with tf.variable_scope('policy', reuse=True):
      logits, _ = tf.nn.dynamic_rnn(
          self.policy_cell, tf.gather(obs_embeddings, inputs),
          sequence_length=self.adjusted_lengths,
          dtype=dtype)

    self.given_batch = AttrDict(
        logits=logits,
        tokens=sampled_tokens,
        episode_lengths=self.adjusted_lengths,
        probs=tf.nn.softmax(logits),
        log_probs=tf.nn.log_softmax(logits))

    # Episode masks.
    max_episode_length = tf.shape(self.actions)[1]
    # range_row shape: [1, max_episode_length]
    range_row = tf.expand_dims(tf.range(max_episode_length), 0)
    episode_masks = tf.cast(
        tf.less(range_row, tf.expand_dims(self.given_batch.episode_lengths, 1)),
        dtype=dtype)
    episode_masks_3d = tf.expand_dims(episode_masks, 2)

    # Length adjusted episodes.
    self.a_probs = a_probs = self.given_batch.probs * episode_masks_3d
    self.a_log_probs = a_log_probs = (
        self.given_batch.log_probs * episode_masks_3d)
    self.a_policy_multipliers = a_policy_multipliers = (
        self.policy_multipliers * episode_masks)

    # pi_loss is scalar
    acs_onehot = tf.one_hot(self.actions, self.action_space, dtype=dtype)
    self.acs_onehot = acs_onehot
    chosen_masked_log_probs = acs_onehot * a_log_probs
    pi_target = tf.expand_dims(a_policy_multipliers, -1)
    pi_loss_per_step = chosen_masked_log_probs * pi_target  # Maximize.
    self.pi_loss = pi_loss = (
        -tf.reduce_mean(tf.reduce_sum(pi_loss_per_step, axis=[1, 2]), axis=0)
        * MAGIC_LOSS_MULTIPLIER)  # Minimize.
    assert len(self.pi_loss.shape) == 0  # pylint: disable=g-explicit-length-test

    # shape: [batch_size, time]
    self.chosen_log_probs = tf.reduce_sum(chosen_masked_log_probs, axis=2)
    self.chosen_probs = tf.reduce_sum(acs_onehot * a_probs, axis=2)

    # Maximize entropy regularizer
    self.entropy = entropy = (
        -tf.reduce_mean(
            tf.reduce_sum(a_probs * a_log_probs, axis=[1, 2]), axis=0)
        * MAGIC_LOSS_MULTIPLIER)  # Maximize
    self.negentropy = -entropy  # Minimize negentropy.
    assert len(self.negentropy.shape) == 0  # pylint: disable=g-explicit-length-test

    # off-policy loss
    self.offp_switch = tf.placeholder(dtype, [], name='offp_switch')
    if self.config.topk != 0:
      # Add SOS to beginning of the sequence.
      offp_inputs = tf.gather(obs_embeddings,
                              rshift_time(self.off_policy_targets,
                                          fill=self.eos_int))
      with tf.variable_scope('policy', reuse=True):
        offp_logits, _ = tf.nn.dynamic_rnn(
            self.policy_cell, offp_inputs, self.off_policy_target_lengths,
            dtype=dtype)  # shape: [batch_size, time, action_space]
      topk_loss_per_step = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=self.off_policy_targets,
          logits=offp_logits,
          name='topk_loss_per_logit')
      # Take mean over batch dimension so that the loss multiplier strength is
      # independent of batch size. Sum over time dimension.
      topk_loss = tf.reduce_mean(
          tf.reduce_sum(topk_loss_per_step, axis=1), axis=0)
      assert len(topk_loss.shape) == 0  # pylint: disable=g-explicit-length-test
      self.topk_loss = topk_loss * self.offp_switch
      logger.info('Including off policy loss.')
    else:
      self.topk_loss = topk_loss = 0.0

    self.entropy_hparam = tf.constant(
        self.config.entropy_beta, dtype=dtype, name='entropy_beta')

    self.pi_loss_term = pi_loss * self.config.pi_loss_hparam
    self.entropy_loss_term = self.negentropy * self.entropy_hparam
    self.topk_loss_term = self.config.topk_loss_hparam * topk_loss
    self.loss = (
        self.pi_loss_term
        + self.entropy_loss_term
        + self.topk_loss_term)

    params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                               tf.get_variable_scope().name)
    self.trainable_variables = params
    self.sync_variables = self.trainable_variables
    non_embedding_params = [p for p in params
                            if obs_embedding_scope not in p.name]
    self.non_embedding_params = non_embedding_params
    self.params = params

    if self.config.regularizer:
      logger.info('Adding L2 regularizer with scale %.2f.',
                   self.config.regularizer)
      self.regularizer = self.config.regularizer * sum(
          tf.nn.l2_loss(w) for w in non_embedding_params)
      self.loss += self.regularizer
    else:
      logger.info('Skipping regularizer.')
      self.regularizer = 0.0

    # Only build gradients graph for local model.
    if self.is_local:
      unclipped_grads = tf.gradients(self.loss, params)
      self.dense_unclipped_grads = [
          tf.convert_to_tensor(g) for g in unclipped_grads]
      self.grads, self.global_grad_norm = tf.clip_by_global_norm(
          unclipped_grads, self.config.grad_clip_threshold)
      self.gradients_dict = dict(zip(params, self.grads))
      self.optimizer = make_optimizer(self.config.optimizer, self.learning_rate)
      self.all_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                             tf.get_variable_scope().name)

    self.do_iw_summaries = do_iw_summaries
    if self.do_iw_summaries:
      b = None
      self.log_iw_replay_ph = tf.placeholder(tf.float32, [b],
                                             'log_iw_replay_ph')
      self.log_iw_policy_ph = tf.placeholder(tf.float32, [b],
                                             'log_iw_policy_ph')
      self.log_prob_replay_ph = tf.placeholder(tf.float32, [b],
                                               'log_prob_replay_ph')
      self.log_prob_policy_ph = tf.placeholder(tf.float32, [b],
                                               'log_prob_policy_ph')
      self.log_norm_replay_weights_ph = tf.placeholder(
          tf.float32, [b], 'log_norm_replay_weights_ph')
      self.iw_summary_op = tf.summary.merge([
          tf.summary.histogram('is/log_iw_replay', self.log_iw_replay_ph),
          tf.summary.histogram('is/log_iw_policy', self.log_iw_policy_ph),
          tf.summary.histogram('is/log_prob_replay', self.log_prob_replay_ph),
          tf.summary.histogram('is/log_prob_policy', self.log_prob_policy_ph),
          tf.summary.histogram(
              'is/log_norm_replay_weights', self.log_norm_replay_weights_ph),
      ])

  def _iw_summary(self, session, replay_iw, replay_log_probs,
                  norm_replay_weights, on_policy_iw,
                  on_policy_log_probs):
    """Compute summaries for importance weights at a given batch.

    Args:
      session: tf.Session instance.
      replay_iw: Importance weights for episodes from replay buffer.
      replay_log_probs: Total log probabilities of the replay episodes under the
          current policy.
      norm_replay_weights: Normalized replay weights, i.e. values in `replay_iw`
          divided by the total weight in the entire replay buffer. Note, this is
          also the probability of selecting each episode from the replay buffer
          (in a roulette wheel replay buffer).
      on_policy_iw: Importance weights for episodes sampled from the current
          policy.
      on_policy_log_probs: Total log probabilities of the on-policy episodes
          under the current policy.

    Returns:
      Serialized TF summaries. Use a summary writer to write these summaries to
      disk.
    """
    return session.run(
        self.iw_summary_op,
        {self.log_iw_replay_ph: np.log(replay_iw),
         self.log_iw_policy_ph: np.log(on_policy_iw),
         self.log_norm_replay_weights_ph: np.log(norm_replay_weights),
         self.log_prob_replay_ph: replay_log_probs,
         self.log_prob_policy_ph: on_policy_log_probs})

  def make_summary_ops(self):
    """Construct summary ops for the model."""
    # size = number of timesteps across entire batch. Number normalized by size
    # will not be affected by the amount of padding at the ends of sequences
    # in the batch.
    size = tf.cast(
        tf.reduce_sum(self.given_batch.episode_lengths), dtype=self.dtype)
    offp_size = tf.cast(tf.reduce_sum(self.off_policy_target_lengths),
                        dtype=self.dtype)
    scope_prefix = self.parent_scope_name

    def _remove_prefix(prefix, name):
      assert name.startswith(prefix)
      return name[len(prefix):]

    # RL summaries.
    self.rl_summary_op = tf.summary.merge(
        [tf.summary.scalar('model/policy_loss', self.pi_loss / size),
         tf.summary.scalar('model/topk_loss', self.topk_loss / offp_size),
         tf.summary.scalar('model/entropy', self.entropy / size),
         tf.summary.scalar('model/loss', self.loss / size),
         tf.summary.scalar('model/grad_norm',
                           tf.global_norm(self.grads)),
         tf.summary.scalar('model/unclipped_grad_norm', self.global_grad_norm),
         tf.summary.scalar('model/non_embedding_var_norm',
                           tf.global_norm(self.non_embedding_params)),
         tf.summary.scalar('hparams/entropy_beta', self.entropy_hparam),
         tf.summary.scalar('hparams/topk_loss_hparam', self.config.topk_loss_hparam),
         tf.summary.scalar('hparams/learning_rate', self.learning_rate),
         tf.summary.scalar('model/trainable_var_norm',
                           tf.global_norm(self.trainable_variables)),
         tf.summary.scalar('loss/loss', self.loss),
         tf.summary.scalar('loss/entropy', self.entropy_loss_term),
         tf.summary.scalar('loss/policy', self.pi_loss_term),
         tf.summary.scalar('loss/offp', self.topk_loss_term)] +
        [tf.summary.scalar(
            'param_norms/' + _remove_prefix(scope_prefix + '/', p.name),
            tf.norm(p))
         for p in self.params] +
        [tf.summary.scalar(
            'grad_norms/' + _remove_prefix(scope_prefix + '/', p.name),
            tf.norm(g))
         for p, g in zip(self.params, self.grads)] +
        [tf.summary.scalar(
            'unclipped_grad_norms/' + _remove_prefix(scope_prefix + '/',
                                                     p.name),
            tf.norm(g))
         for p, g in zip(self.params, self.dense_unclipped_grads)])

    self.text_summary_placeholder = tf.placeholder(tf.string, shape=[])
    self.rl_text_summary_op = tf.summary.text('rl',
                                              self.text_summary_placeholder)

  def _rl_text_summary(self, session, step, npe, tot_r, num_steps, code, result):
    """Logs summary about a single episode and creates a text_summary for TB.

    Args:
      session: tf.Session instance.
      step: Global training step.
      npe: Number of programs executed so far.
      tot_r: Total reward.
      num_steps: Number of timesteps in the episode (i.e. code length).
      code: String representation of the code.
      result: Result of program execution

    Returns:
      Serialized text summary data for tensorboard.
    """
    if not code:
      code = ' '
    text = (
        'Tot R: **%.2f**;  Len: **%d**;  Result: **%s**\n\n'
        '\n\nCode: **`%s`**'
        % (tot_r, num_steps, result, code))
    text_summary = session.run(self.rl_text_summary_op,
                               {self.text_summary_placeholder: text})
    logger.info(
        'Step %d.\t NPE: %d\t Result: %s.\t Tot R: %.2f.\t Length: %d. \tProgram: %s',
        step, npe, result, tot_r, num_steps, code)
    return text_summary

  def _rl_reward_summary(self, total_rewards):
    """Create summary ops that report on episode rewards.

    Creates summaries for average, median, max, and min rewards in the batch.

    Args:
      total_rewards: Tensor of shape [batch_size] containing the total reward
          from each episode in the batch.

    Returns:
      tf.Summary op.
    """
    tr = np.asarray(total_rewards)
    reward_summary = tf.Summary(value=[
        tf.Summary.Value(
            tag='reward/avg',
            simple_value=np.mean(tr)),
        tf.Summary.Value(
            tag='reward/med',
            simple_value=np.median(tr)),
        tf.Summary.Value(
            tag='reward/max',
            simple_value=np.max(tr)),
        tf.Summary.Value(
            tag='reward/min',
            simple_value=np.min(tr))])
    return reward_summary

  def write_programs(self, session, inspiration_branch):
    """Generate Brainfuck programs with this language model"""

    # Sample new programs from the policy.
    # Note: batch size is constant. A full batch will be sampled, but not all
    # programs will be executed and added to the replay buffer. Those which
    # are not executed will be discarded and not counted.

    length_limit = f'code.str.len() < {self.config.timestep_limit}'
    self.inspiration_branch = inspiration_branch.query(length_limit)

    batch_actions, episode_lengths, log_probs = session.run(
        [self.sampled_batch.tokens,
         self.sampled_batch.episode_lengths,
         self.sampled_batch.log_probs])

    if episode_lengths.size == 0:
      # This should not happen.
      logger.warn(
          'Shapes:\n'
          'batch_actions.shape: %s\n'
          'episode_lengths.shape: %s\n',
          batch_actions.shape, episode_lengths.shape)

    programs = make_dev_codebase()
    for actions, episode_length, action_log_probs in zip(
        batch_actions, episode_lengths, log_probs
      ):
      code = self.int_to_char(actions[:episode_length])
      if code[-1] == self.eos_char:
        code = code[:-1]

      logger.info(f'Wrote program: {code}')
      metrics = {'log_prob': np.choose(actions[:episode_length],
                                       action_log_probs[:episode_length].T).sum()}
      programs.commit(code, metrics=metrics)

    return programs

  def accept_feedback(self, session, feedback_branch, train_op, global_step_op, return_gradients=False):
    """Improve language generation based on rewards from the real world

    Args:
      session: tf.Session instance.
      reinforcement: a Reinforcement object that summarizes the results of
          recent program runs
      train_op: A TF op which will perform the gradient update. LMAgent does not
          own its training op, so that trainers can do distributed training
          and construct a specialized training op.
      global_step_op: A TF op which will return the current global step when
          run (should not increment it).
      return_gradients: If True, the gradients will be saved and returned from
          this method call. This is useful for testing.

    Returns:
      Results from the update step in a ReflectionResult namedtuple, including
      global step, global NPE, serialized summaries, and optionally gradients.
    """
    assert self.is_local
    assert len(feedback_branch) == self.config.batch_size

    if self.config.topk != 0:
      off_policy_branch = self.inspiration_branch.top_k('quality', self.config.topk)
      if len(off_policy_branch) > self.config.topk_batch_size:
        off_policy_branch = off_policy_branch.sample(self.config.topk_batch_size)
    else:
      off_policy_branch = make_codebase_like(self.inspiration_branch)

    if len(off_policy_branch):
      off_policy_target_lengths, off_policy_targets, _, _ = \
        self.process_episodes(off_policy_branch)

      offp_switch = 1
    else:
      off_policy_targets = [[0]]
      off_policy_target_lengths = [1]
      offp_switch = 0

    feedback_lengths, feedback_actions, feedback_targets, feedback_returns = \
      self.process_episodes(feedback_branch)

    # Do update for REINFORCE or REINFORCE + replay buffer.
    if self.num_replay_per_batch == 0:
      # Train with on-policy REINFORCE.
      num_programs_from_policy = self.config.batch_size

      # Process on-policy samples.

      batch_policy_multipliers = feedback_targets
      batch_emp_values = [[]]

      fetches = {
          'global_step': global_step_op,
          'program_count': self.program_count,
          'summaries': self.rl_summary_op,
          'train_op': train_op,
          'gradients': self.gradients_dict if return_gradients else self.no_op}

      fetched = session.run(
          fetches,
          {self.actions: feedback_actions,
          self.empirical_values: batch_emp_values,
          self.policy_multipliers: batch_policy_multipliers,
          self.adjusted_lengths: feedback_lengths,
          self.off_policy_targets: off_policy_targets,
          self.off_policy_target_lengths: off_policy_target_lengths,
          self.offp_switch: offp_switch})

      combined_adjusted_lengths = feedback_lengths
      combined_returns = feedback_returns
    else:
      # Train with REINFORCE + off-policy replay buffer by using importance
      # sampling.

      # Sample from experince replay buffer
      empty_replay_buffer = len(self.inspiration_branch) < self.num_replay_per_batch
      num_programs_from_replay_buff = (
          self.num_replay_per_batch if not empty_replay_buffer else 0)
      num_programs_from_policy = (
          self.config.batch_size - num_programs_from_replay_buff)
      if (not empty_replay_buffer) and num_programs_from_replay_buff:
        replay_branch = self.inspiration_branch.sample(num_programs_from_replay_buff)
        replay_lengths, replay_actions, replay_batch_targets, replay_returns = \
          self.process_episodes(replay_branch)

        replay_episode_actions = utils.stack_pad(replay_actions, pad_axes=0,
                                                 dtype=np.int32)

        # compute log probs for replay samples under current policy
        all_replay_log_probs, = session.run(
            [self.given_batch.log_probs],
            {self.actions: replay_episode_actions,
             self.adjusted_lengths: replay_lengths})

        replay_log_probs = [
            np.choose(replay_actions[i, :l], all_replay_log_probs[i, :l].T).sum()
            for i, l in enumerate(replay_lengths)]
        replay_branch['log_prob'] = replay_log_probs

        # Convert 2D array back into ragged 2D list.
        replay_policy_multipliers = [
            replay_batch_targets[i, :l]
            for i, l
            in enumerate(
                replay_lengths[:num_programs_from_replay_buff])]
      else:
        replay_lengths = None
        replay_actions = None
        replay_batch_targets = None
        replay_returns = None
        replay_policy_multipliers = None

      # Process on-policy samples.
      p = num_programs_from_policy

      on_policy_branch = feedback_branch.sample(p)
      adjusted_lengths, episode_actions, batch_targets, batch_returns = \
        self.process_episodes(on_policy_branch)
      batch_policy_multipliers = batch_targets
      batch_emp_values = [[]]
      on_policy_returns = batch_returns

      # On-policy episodes.
      if num_programs_from_policy:
        separate_actions = [
            episode_actions[i, :l]
            for i, l in enumerate(adjusted_lengths)]
        new_experiences = [
            (separate_actions[i],
            on_policy_branch['total_reward'][i],
            on_policy_branch['log_prob'][i], l)
            for i, l in enumerate(adjusted_lengths)]
        on_policy_policy_multipliers = [
            batch_policy_multipliers[i, :l]
            for i, l in enumerate(adjusted_lengths)]
        (on_policy_actions,
        _,  # rewards
        on_policy_log_probs,
        on_policy_adjusted_lengths) = zip(*new_experiences)
      else:
        new_experiences = []
        on_policy_policy_multipliers = []
        on_policy_actions = []
        on_policy_log_probs = []
        on_policy_adjusted_lengths = []

      if (not empty_replay_buffer) and num_programs_from_replay_buff:
        on_policy_branch['quality'] = 0
        # Look for new experiences in replay buffer. Assign weight if an episode
        # is in the buffer.
        on_policy_branch.replace(replay_branch)

      # Randomly select on-policy or off policy episodes to train on.
      combined_adjusted_lengths = list_(replay_lengths) + list_(on_policy_adjusted_lengths)
      combined_returns = utils.stack_pad(list_(replay_returns) + list_(on_policy_returns), pad_axes=0)
      combined_actions = utils.stack_pad(list_(replay_actions) + list_(on_policy_actions), pad_axes=0)
      combined_policy_multipliers = utils.stack_pad(list_(replay_policy_multipliers) + list_(on_policy_policy_multipliers),
                                                    pad_axes=0)

      # Importance adjustment. Naive formulation:
      # E_{x~p}[f(x)] ~= 1/N sum_{x~p}(f(x)) ~= 1/N sum_{x~q}(f(x) * p(x)/q(x)).
      # p(x) is the policy, and q(x) is the off-policy distribution, i.e. replay
      # buffer distribution. Importance weight w(x) = p(x) / q(x).

      # Instead of sampling from the replay buffer only, we sample from a
      # mixture distribution of the policy and replay buffer.
      # We are sampling from the mixture a*q(x) + (1-a)*p(x), where 0 <= a <= 1.
      # Thus the importance weight w(x) = p(x) / (a*q(x) + (1-a)*p(x))
      # = 1 / ((1-a) + a*q(x)/p(x)) where q(x) is 0 for x sampled from the
      #                             policy.
      # Note: a = self.replay_alpha
      if empty_replay_buffer:
        # The replay buffer is empty.
        # Do no gradient update this step. The replay buffer will have stuff in
        # it next time.
        combined_policy_multipliers *= 0
      elif not num_programs_from_replay_buff:
        combined_policy_multipliers = np.ones([len(combined_actions), 1],
                                              dtype=np.float32)
      else:
        # If a < 1 compute importance weights
        # importance weight
        # = 1 / [(1 - a) + a * exp(log(replay_weight / total_weight / p))]
        # = 1 / ((1-a) + a*q/p)
        combined_branch = Codebase(on_policy_branch.metrics,
                                   on_policy_branch.metadata,
                                   deduplication=False)
        combined_branch.merge(on_policy_branch)
        combined_branch.merge(replay_branch)

        importance_weights = compute_iw(combined_branch, self.replay_alpha)

        if self.config.iw_normalize:
          importance_weights *= (
              float(self.config.batch_size) / importance_weights.sum())
        combined_policy_multipliers *= importance_weights.reshape(-1, 1)

      # Train on replay batch, top-k MLE.
      assert self.program_count is not None
      fetches = {
          'global_step': global_step_op,
          'program_count': self.program_count,
          'summaries': self.rl_summary_op,
          'train_op': train_op,
          'gradients': self.gradients_dict if return_gradients else self.no_op}
      fetched = session.run(
          fetches,
          {self.actions: combined_actions,
          self.empirical_values: [[]],  # replay_emp_values,
          self.policy_multipliers: combined_policy_multipliers,
          self.adjusted_lengths: combined_adjusted_lengths,
          self.off_policy_targets: off_policy_targets,
          self.off_policy_target_lengths: off_policy_target_lengths,
          self.offp_switch: offp_switch})

    # Update program count.
    session.run(
        [self.program_count_add_op],
        {self.program_count_add_ph: num_programs_from_policy})

    # Update EMA baselines on the mini-batch which we just did traning on.
    for i, program in enumerate(feedback_branch['code']):
      episode_length = combined_adjusted_lengths[i]
      empirical_returns = combined_returns[i, :episode_length]
      for j in xrange(episode_length):
        # Update ema_baselines in place.
        self.ema_by_len[j] = (
            self.config.ema_baseline_decay * self.ema_by_len[j]
            + (1 - self.config.ema_baseline_decay) * empirical_returns[j])

    global_step = fetched['global_step']
    global_npe = fetched['program_count']
    core_summaries = fetched['summaries']
    summaries_list = [core_summaries]

    if num_programs_from_policy:
      s_i = 0
      text_summary = self._rl_text_summary(
          session,
          global_step,
          global_npe,
          feedback_branch['total_reward'][s_i],
          feedback_lengths[s_i],
          feedback_branch['code'][s_i],
          feedback_branch['result'][s_i])
      reward_summary = self._rl_reward_summary(feedback_branch['total_reward'])

      max_i = np.argmax(feedback_branch['total_reward'])
      max_tot_r = feedback_branch['total_reward'][max_i]
      if max_tot_r >= self.top_reward:
        if max_tot_r >= self.top_reward:
          self.top_reward = max_tot_r
        logger.info('Top code: r=%.2f, \t%s', max_tot_r, feedback_branch['code'][max_i])

      summaries_list += [text_summary, reward_summary]

      if self.do_iw_summaries and not empty_replay_buffer:
        # prob of replay samples under replay buffer sampling.
        total_weight = sum(replay_branch['quality'])
        norm_replay_weights = [
            w / total_weight
            for w in replay_branch['quality']]
        replay_iw = compute_iw(replay_branch, self.replay_alpha)
        on_policy_iw = compute_iw(on_policy_branch, self.replay_alpha)
        summaries_list.append(
            self._iw_summary(
                session, replay_iw, replay_log_probs, norm_replay_weights,
                on_policy_iw, on_policy_log_probs))

    return ReflectionResult(
        global_step=global_step,
        global_npe=global_npe,
        summaries_list=summaries_list,
        gradients_dict=fetched['gradients'])

  def process_episodes(self, reinforce_branch):
    """Compute REINFORCE targets.

    Treat a program as a reinforcement learning episode where a character
    is an action and 'quality' is the reward for the last character

    REINFORCE here takes the form:
    grad_t = grad[log(pi(a_t|c_t))*target_t]
    where c_t is context: i.e. RNN state or environment state (or both).

    Args:
      reinforce_branch: A Codebase with 'quality' metric
      baselines: Provide baselines for each timestep. This is a
          list (or indexable container) of length max_time. Note: baselines are
          shared across all episodes, which is why there is no batch dimension.
          It is up to the caller to update baselines accordingly.

    Returns:
      lengths: Length of every episode
      actions: List of actions for every episode
      targets: REINFORCE targets for each episode and timestep. A numpy
          array of shape [batch_size, max_sequence_length].
      returns: Returns computed for each episode and timestep. This is for
          reference, and is not used in the REINFORCE gradient update (but was
          used to compute the targets). A numpy array of shape
          [batch_size, max_sequence_length].
    """
    assert 'quality' in reinforce_branch.metrics, 'Reinforcement has to contain a quality metric'
    assert len(reinforce_branch)

    lengths = [len(code) for code in reinforce_branch['code']]
    baselines = self.ema_by_len

    # We only reward the last character
    # Everything else is a preparation for dat final punch
    action_rewards = utils.stack_pad(
      [[0] * (episode_length - 1) + [episode_reward]
      for episode_reward, episode_length in
      zip(reinforce_branch['quality'], lengths)],
      pad_axes=0
    )

    num_programs = len(action_rewards)

    batch_returns = [None] * num_programs
    batch_targets = [None] * num_programs
    for i in xrange(num_programs):
      episode_length = lengths[i]
      # Compute target for each timestep.
      #    target_t = R_t - baselines[t]
      #    where `baselines` are provided.
      # In practice we use a more generalized formulation of advantage. See docs
      # for `discounted_advantage_and_rewards`.

      # Compute return for each timestep. See section 3 of
      # https://arxiv.org/pdf/1602.01783.pdf
      assert baselines is not None
      empirical_returns = rollout_lib.discount(action_rewards[i], gamma=1.0)
      targets = [None] * episode_length
      for j in xrange(episode_length):
        targets[j] = empirical_returns[j] - baselines[j]
      batch_returns[i] = empirical_returns
      batch_targets[i] = targets

    actions = utils.stack_pad(
              [self.char_to_int(code)
                for code in reinforce_branch['code']],
                pad_axes=0)
    actions = np.array(actions, dtype=int)

    returns = utils.stack_pad(batch_returns, 0)
    if num_programs:
      targets = utils.stack_pad(batch_targets, 0)
    else:
      targets = np.array([], dtype=np.float32)

    return (lengths, actions, targets, returns)

def compute_iw(codebase, replay_alpha):
    """Compute importance weights for a batch of episodes.

    The codebase has to contain a 'quality' and 'log_prob' metrics

    Returns:
      Numpy array of shape [batch_size] containing the importance weight for
      each episode in the batch.
    """
    total_replay_weight = sum(codebase['quality'])
    if total_replay_weight == 0:
      # This happens when:
      # - codebase is empty
      # or
      # - all replay weights are zero due to floating point limitations
      # replay weight are exponents and sometimes exp(minus a lot) is rounded to 0
      n = len(codebase)
      return np.ones(n) / n
    else:
      try:
        log_total_replay_weight = log(total_replay_weight)

        # importance weight
        # = 1 / [(1 - a) + a * exp(log(replay_weight / total_weight / p))]
        # = 1 / ((1-a) + a*q/p)
        a = float(replay_alpha)
        a_com = 1.0 - a  # compliment of a

        importance_weights = np.asarray(
            [1.0 / (a_com
                    + a * exp((log(replay_weight) - log_total_replay_weight)
                              - log_p))
            if replay_weight > 0 else 1.0 / a_com
            for log_p, replay_weight
            in zip(codebase['log_prob'], codebase['quality'])])
        return importance_weights
      except OverflowError:
        # This Softmax is too close for the CPU to handle
        # So it's safe to turn it into just max
        importance_weights = np.zeros(len(codebase))
        importance_weights[np.argmax(codebase['quality'])] = 1
      return importance_weights


import tensorflow as tf
from tools.utils import helper
from models.yolonet import yoloconv, pureconv, yolonet
from tensorflow.contrib import slim
import os
from datetime import datetime
from tqdm import tqdm
import numpy as np
import sys
import argparse
import importlib

def restore_ckpt(sess: tf.Session, var_list: list, ckptdir: str):
    if ckptdir == '' or ckptdir == None:
        pass
    elif 'mobilenet' in ckptdir:
        variables_to_restore = slim.get_model_variables()
        loader = tf.train.Saver([var for var in variables_to_restore if 'MobilenetV1' in var.name])
        loader.restore(sess, ckptdir)
    else:
        ckpt = tf.train.get_checkpoint_state(ckptdir)
        loader = tf.train.Saver(var_list=var_list)
        loader.restore(sess, ckpt.model_checkpoint_path)

def main(train_list,
         pre_ckpt,
         model_def,
         is_augmenter,
         depth_multiplier,
         image_size,
         output_size,
         batch_size,
         rand_seed,
         max_nrof_epochs,
         init_learning_rate,
         learning_rate_decay_epochs,
         learning_rate_decay_factor,
         pb_name,
         log_dir):
    g = tf.get_default_graph()
    tf.set_random_seed(rand_seed)
    """ import network """
    network = eval(model_def)

    """ generate the dataset """
    fddb = helper(train_list, image_size, output_size)
    fddb.set_dataset(batch_size, rand_seed, is_training=is_augmenter)
    next_img, next_label = fddb.get_iter()

    """ define the model """
    batch_image = tf.placeholder_with_default(next_img, shape=[None, image_size[0], image_size[1], 3], name='Input_image')
    batch_label = tf.placeholder_with_default(next_label, shape=[None, output_size[0], output_size[1], 5], name='Input_label')
    training_control = tf.placeholder_with_default(True, shape=[], name='training_control')
    true_label = tf.identity(batch_label)
    nets, endpoints = network(batch_image, depth_multiplier, is_training=training_control)

    """ reshape the model output """
    pred_label = tf.identity(nets, name='predict')
    """ split the label """
    pred_xywh = pred_label[:, :, :, 0:4]
    pred_xywh = tf.nn.sigmoid(pred_xywh)
    pred_confidence = pred_label[:, :, :, 4:5]
    pred_confidence_sigmoid = tf.nn.sigmoid(pred_confidence)

    true_xywh = true_label[:, :, :, 0:4]
    true_confidence = true_label[:, :, :, 4:5]

    obj_mask = true_confidence[:, :, :, 0] > .7

    """ define loss """
    mse_loss = tf.losses.mean_squared_error(labels=tf.boolean_mask(true_xywh, obj_mask), predictions=tf.boolean_mask(pred_xywh, obj_mask))
    obj_loss = 2 * tf.reduce_sum(tf.boolean_mask(tf.nn.sigmoid_cross_entropy_with_logits(labels=true_confidence, logits=pred_confidence), obj_mask))/batch_size
    noobj_loss = tf.reduce_sum(tf.boolean_mask(tf.nn.sigmoid_cross_entropy_with_logits(labels=true_confidence, logits=pred_confidence), tf.logical_not(obj_mask)))/batch_size
    total_loss = obj_loss+noobj_loss+mse_loss

    """ define steps """
    global_steps = tf.train.create_global_step()

    """ define train optimizer """
    current_learning_rate = tf.train.exponential_decay(init_learning_rate, global_steps, fddb.epoch_step // learning_rate_decay_epochs,
                                                       learning_rate_decay_factor, staircase=False)
    """ define train_op """
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = tf.train.AdamOptimizer(current_learning_rate).minimize(total_loss, global_steps)
    # train_op = slim.learning.create_train_op(total_loss, tf.train.AdamOptimizer(current_learning_rate), global_steps)

    """ calc the accuracy """
    # pred_obj_num = tf.reduce_sum(tf.cast(pred_confidence_sigmoid > .7, tf.float32))
    # confidence_acc = tf.boolean_mask(tf.equal(true_confidence, tf.cast(pred_confidence_sigmoid > .7, tf.float32)), obj_mask)
    # confidence_acc = tf.reduce_mean(tf.cast(confidence_acc, tf.float32))
    confidence_acc, acc_op = tf.metrics.accuracy(tf.boolean_mask(true_confidence, obj_mask), tf.boolean_mask(tf.cast(pred_confidence_sigmoid > .7, tf.float32), obj_mask))
    test_confidence_acc, test_acc_op = tf.metrics.accuracy(tf.boolean_mask(true_confidence, obj_mask), tf.boolean_mask(tf.cast(pred_confidence_sigmoid > .7, tf.float32), obj_mask))
    with tf.Session() as sess:
        """ must save the bn paramter! """
        var_list = tf.global_variables()  # list(set(tf.trainable_variables() + [g for g in tf.global_variables() if 'moving_' in g.name]))
        saver = tf.train.Saver(var_list)

        # init the model and restore the pre-train weight
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())  # NOTE the accuracy must init local variable
        restore_ckpt(sess, var_list, pre_ckpt)
        # define the log and saver
        subdir = os.path.join(log_dir, datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S'))
        writer = tf.summary.FileWriter(subdir, graph=sess.graph)
        # write_arguments_to_file(os.path.join(subdir, 'arguments.txt'))
        tf.summary.scalar('total_loss', total_loss)
        tf.summary.scalar('obj_loss', obj_loss)
        tf.summary.scalar('noobj_loss', noobj_loss)
        tf.summary.scalar('mse_loss', mse_loss)
        # tf.summary.scalar('pred_obj_num', pred_obj_num)
        tf.summary.scalar('leraning rate', current_learning_rate)
        tf.summary.scalar('test acc', test_confidence_acc)
        tf.summary.scalar('train acc', confidence_acc)
        merged = tf.summary.merge_all()

        try:
            for i in range(max_nrof_epochs):
                with tqdm(total=fddb.epoch_step, bar_format='{n_fmt}/{total_fmt} |{bar}| {rate_fmt}{postfix}', unit=' batch', dynamic_ncols=True) as t:
                    for j in range(fddb.epoch_step):
                        if j % 40 == 0:
                            summary, test_con_acc, _, step_cnt = sess.run(
                                [merged,  test_confidence_acc, test_acc_op, global_steps], feed_dict={training_control: False})
                        else:
                            summary, _, total_l, con_acc, _,  lr, step_cnt = sess.run(
                                [merged, train_op, total_loss, confidence_acc, acc_op, current_learning_rate, global_steps])
                            t.set_postfix(loss='{:<5.3f}'.format(total_l), acc='{:<4.2f}%'.format(con_acc*100), lr='{:f}'.format(lr))
                        writer.add_summary(summary, step_cnt)
                        t.update()
            saver.save(sess, save_path=os.path.join(subdir, 'model.ckpt'), global_step=global_steps)
            """ save as pb """
            constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['predict'])
            with tf.gfile.GFile(pb_name, 'wb') as f:
                f.write(constant_graph.SerializeToString())
            print('save over')
        except KeyboardInterrupt as e:
            saver.save(sess, save_path=os.path.join(subdir, 'model.ckpt'), global_step=global_steps)
            """ save as pb """
            constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['Yolo/Final/conv2d/BiasAdd'])
            with tf.gfile.GFile(pb_name, 'wb') as f:
                f.write(constant_graph.SerializeToString())
            print('save over')

def parse_arguments(argv):
    parser = argparse.ArgumentParser()

    parser.add_argument('--train_list',                 type=str,   help='trian file lists',           default='data/train.list')
    parser.add_argument('--pre_ckpt',                   type=str,   help='pre-train ckpt dir',         default='mobilenet_v1_0.5_224/mobilenet_v1_0.5_224.ckpt')
    parser.add_argument('--model_def',                  type=str,   help='Model definition.',          choices=['yoloseparabe', 'yoloconv', 'pureconv'], default='yoloconv')
    parser.add_argument('--augmenter',                  type=bool,  help='use image augmenter',        default=True)
    parser.add_argument('--depth_multiplier',           type=float, help='mobilenet depth_multiplier', default=0.5)
    parser.add_argument('--image_size',                 type=int,   help='net work input image size',  default=(224, 320), nargs='+')
    parser.add_argument('--output_size',                type=int,   help='net work output image size', default=(7, 10), nargs='+')
    parser.add_argument('--batch_size',                 type=int,   help='batch size',                 default=32)
    parser.add_argument('--rand_seed',                  type=int,   help='random seed',                default=6)
    parser.add_argument('--max_nrof_epochs',            type=int,   help='epoch num',                  default=10)
    parser.add_argument('--init_learning_rate',         type=float, help='init learing rate',          default=0.0005)
    parser.add_argument('--learning_rate_decay_epochs', type=int,   help='learning rate decay epochs', default=10)
    parser.add_argument('--learning_rate_decay_factor', type=int,   help='learning rate decay factor', default=1.0)
    parser.add_argument('--pb_name',                    type=str,   help='pb name',                    default='Training_save.pb')
    parser.add_argument('--log_dir',                    type=str,   help='log dir',                    default='log')

    return parser.parse_args(argv)

if __name__ == "__main__":
    args = parse_arguments(sys.argv[1:])
    main(args.train_list,
         args.pre_ckpt,
         args.model_def,
         args.augmenter,
         args.depth_multiplier,
         args.image_size,
         args.output_size,
         args.batch_size,
         args.rand_seed,
         args.max_nrof_epochs,
         args.init_learning_rate,
         args.learning_rate_decay_epochs,
         args.learning_rate_decay_factor,
         args.pb_name,
         args.log_dir)


import logging
import os
import shutil
import sys
from argparse import (ArgumentParser)
from mob_suite.version import __version__
import pandas as pd
import scipy
import scipy.cluster.hierarchy as sch
from Bio import SeqIO
from scipy.cluster.hierarchy import fcluster
from mob_suite.blast import BlastRunner
from mob_suite.utils import \
    read_fasta_dict
from mob_suite.wrappers import mash

LOG_FORMAT = '%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'

def init_console_logger(lvl):
    logging_levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]
    report_lvl = logging_levels[lvl]

    logging.basicConfig(format=LOG_FORMAT, level=report_lvl)
    return logging

def parse_args():
    "Parse the input arguments, use '-h' for help"
    parser = ArgumentParser(description='Mob-Suite: Generate and update existing plasmid clusters')
    parser.add_argument('-m','--mode', type=str, required=True, help='Build: Create a new database from scratch, Update: Update an existing database with one or more sequences')
    parser.add_argument('-o','--outdir', type=str, required=True, help='Output Directory to put results')
    parser.add_argument('-i','--infile', type=str, required=True, help='Input fasta file of one or more closed plasmids to process')
    parser.add_argument('--ref_cluster_file', type=str, required=False, help='Reference mob-cluster file')
    parser.add_argument('--ref_fasta_file', type=str, required=False, help='Reference mob-cluster fasta file')
    parser.add_argument('--ref_mash_db', type=str, required=False, help='Reference mob-cluster mash sketch file')
    parser.add_argument('--num_threads', type=int, required=False, help='Number of threads to be used', default=1)
    parser.add_argument('-w','--overwrite',  required=False, help='Overwrite the MOB-suite databases with results', action='store_true')
    return parser.parse_args()

def read_cluster_assignments(file):
    if os.path.getsize(file) == 0:
        return dict()
    data = pd.read_csv(file, sep='\t', header=0,index_col=0)
    clusters = dict()
    for index, row in data.iterrows():
        record = list()
        for r in row:
            record.append(r)
        clusters[index] = record

    return clusters

def calcDistances(input_fasta,ref_sketch_db,output_file):
    m = mash()
    mash_results = dict()
    with open(output_file, 'w') as oh:
        m.run_mash(ref_sketch_db, input_fasta, oh)
        mash_lines = m.read_mash(output_file)
        for line in mash_lines:
            row = line.strip("\n").split("\t")
            ref_id = row[0]
            query_id = row[1]
            distance = float(row[2])
            if not query_id in mash_results:
                mash_results[query_id] = dict()
            mash_results[query_id][ref_id] = distance
    return mash_results

def write_clusters(file,cluster_assignments,header):
    with open(file, 'w') as outfile:
        for id in cluster_assignments:
            outfile.write(str(id) + "\t" + "\t".join(cluster_assignments[id]))
        outfile.close()

def build_cluster_db(distance_matrix_file,distances):
    data = pd.read_csv(distance_matrix_file, sep='\t', header=0,
                       index_col=0)
    distance_matrix = data.as_matrix()
    condensed_matrix = scipy.spatial.distance.squareform(distance_matrix)
    Z = scipy.cluster.hierarchy.linkage(condensed_matrix, method='complete')

    clust_assignments = dict()

    for dist in distances:
        index = 0
        clusters = fcluster(Z, dist, criterion='distance')
        for id in data.columns.values:
            if not id in clust_assignments:
                clust_assignments[id] = list()
            clust_assignments[id].append(str(clusters[index]))
            index += 1

    return clust_assignments

def add_new_record(fasta_file,ref_mashdb,mash_results_file,ref_cluster_file,distances_thresholds,num_threads):

    ref_clusters = read_cluster_assignments(ref_cluster_file)

    cluster_pointers = list()
    for id in ref_clusters:
        for i in range(0,len(ref_clusters[id])):
            if len(cluster_pointers) == 0:
                cluster_pointers = list(ref_clusters[id])
                continue
            if ref_clusters[id][i] > cluster_pointers[i]:
                cluster_pointers[i] = ref_clusters[id][i]

    for i in range(0,len(cluster_pointers)):
        cluster_pointers[i]+=1

    max_thresh = max(distances_thresholds)
    min_thresh = min(distances_thresholds)

    # Run mash to get distances to the reference
    out_dir = os.path.dirname(fasta_file)
    mashObj = mash()
    mashObj.mashsketch(fasta_file, output_path=fasta_file, sketch_ind=True, num_threads=1, kmer_size=21, sketch_size=1000)
    mashfile_handle = open(mash_results_file,'w')
    mashObj.run_mash(ref_mashdb, fasta_file + '.msh', mashfile_handle,table=False,num_threads=num_threads)
    mashfile_handle.close()

    mashfile_handle = open(mash_results_file, 'r')
    query_mash_distances = dict()

    for line in mashfile_handle:
        row = line.split('\t')
        query_id = row[1]
        ref_id = row[0].split('|')[0]
        distance = float(row[2])

        if distance > max_thresh:
            continue
        if not query_id in query_mash_distances:
            query_mash_distances[query_id] = dict()

        query_mash_distances[query_id][ref_id] = distance

    for query_id in query_mash_distances:

        min_distance = min(query_mash_distances[query_id].values())
        min_dist_key = min(query_mash_distances[query_id], key=query_mash_distances[query_id].get)
        clusters = list()
        if float(min_distance) > float(max_thresh):
            for i in range(0,len(distances_thresholds)):
                clusters.append(cluster_pointers[i])
                cluster_pointers[i]+=1

        elif float(min_distance) <= float(min_thresh):

            clusters = list(ref_clusters[min_dist_key])

        else:

            clusters = list(ref_clusters[min_dist_key])
            for i in range(0, len(distances_thresholds)):
                dist = distances_thresholds[i]
                if float(min_distance) <= float(dist):
                    continue
                clusters[i] = cluster_pointers[i]
                cluster_pointers[i] += 1
        ref_clusters[query_id] = clusters

    return ref_clusters

def writeClusterAssignments(output_file,header,cluster_assignmnets):
    with open(output_file, 'w') as out:
        out.write("\t".join(map(str,header)) + "\n")
        for id in cluster_assignmnets:
            out.write("{}\t{}".format(id,"\t".join(map(str,cluster_assignmnets[id]))) + "\n")
        out.close()

def updateFastaFile(in_fasta_file,out_fasta_file,cluster_assignments):
    out = open(out_fasta_file,'w')
    with open(in_fasta_file, "r") as handle:
        for record in SeqIO.parse(handle, "fasta"):
            row = str(record.id).split('|')
            id = row[0]

            if not id in cluster_assignments:
                continue

            out.write(">{}|{}\n{}\n".format(id,cluster_assignments[id],record.seq))
        handle.close()
    out.close()

def selectCluster(clust_assignments,column):
    out = dict()
    for id in clust_assignments:
        row = str(id).split('|')
        out[row[0]] = clust_assignments[id][column]
    return out

def update_existing(input_fasta,tmp_dir,ref_mash_db,tmp_cluster_file,header,tmp_ref_fasta_file,update_fasta,num_threads=1):
    sequences = read_fasta_dict(input_fasta)

    for id in sequences:
        seq = sequences[id]
        tmp_fasta = os.path.join(tmp_dir,id + '_tmp.fasta')
        with open(tmp_fasta , "w") as fh:
            fh.write("\n>{}\n{}\n".format(id,seq))
            fh.close()
        tmp_mash = os.path.join(tmp_dir,id + '_tmp.txt')
        clust_assignments = add_new_record(tmp_fasta, ref_mash_db ,tmp_mash, tmp_cluster_file,(0.05, 0.0001),num_threads)

        writeClusterAssignments(tmp_cluster_file , header, clust_assignments)
        with open(tmp_ref_fasta_file , "a") as fh:
            fh.write("\n>{}\n{}\n".format(id,seq))
            fh.close()

    clust_dict = selectCluster(clust_assignments, 1)

    updateFastaFile(tmp_ref_fasta_file ,update_fasta, clust_dict)

def main():
    args = parse_args()
    logging = init_console_logger(3)
    logging.info('Running Mob-Suite Clustering toolkit v. {}'.format(__version__))
    logging.info('Processing fasta file {}'.format(args.infile))
    logging.info('Analysis directory {}'.format(args.outdir))

    input_fasta = args.infile
    if not os.path.isfile(input_fasta):
        logging.error('Error, input fasta specified does not exist: {}'.format(input_fasta ))
        sys.exit()

    out_dir = args.outdir
    num_threads = args.num_threads
    if not os.path.isdir(out_dir):
        os.mkdir(out_dir, 0o755)
    tmp_dir = os.path.join(out_dir, '__tmp')
    if not os.path.isdir(tmp_dir):
        os.mkdir(tmp_dir, 0o755)

    mode = str(args.mode).lower()

    if mode not in ('update','build'):
        logging.error('Error you have not entered a valid mode of build or update, you entered: {}'.format(mode))
        print(('Error you have not entered a valid mode of build or update, you entered: {}'.format(mode)))
        sys.exit()

    header = ('id', 0.05, 0.0001)
    tmp_cluster_file = os.path.join(out_dir, 'clusters.txt')
    tmp_ref_fasta_file = os.path.join(tmp_dir, 'references_tmp.fasta')
    update_fasta = os.path.join(out_dir, 'references_updated.fasta')

    if mode == 'update':
        if args.ref_cluster_file is None:
            logging.error('Reference fasta file must be specified, please check help for parameter reference')
            sys.exit()

        ref_fasta = args.ref_fasta_file

        if not os.path.isfile(ref_fasta ):
            logging.error('Reference fasta file specified does not exist: {}'.format(ref_fasta))
            sys.exit()

        if args.ref_cluster_file is None:
            logging.error('Reference cluster file must be specified, please check help for parameter reference')
            sys.exit()

        ref_cluster_file = args.ref_cluster_file

        if not os.path.isfile(ref_cluster_file):
            logging.error('Reference cluster file specified does not exist: {}'.format(ref_cluster_file))
            sys.exit()

        if args.ref_mash_db is None:
            logging.error('Reference mash sketch file must be specified, please check help for parameter reference')
            sys.exit()

        ref_mash_db = args.ref_mash_db
        if not os.path.isfile(ref_mash_db):
            logging.error('Reference mash file specified does not exist: {}'.format(ref_mash_db))
            sys.exit()

        logging.info('Running mob-cluster in update mode with input file: {}'.format(input_fasta))
        logging.info('Running mob-cluster in update mode with output directory: {}'.format(out_dir))
        logging.info('Running mob-cluster in update mode on reference fasta file: {}'.format(ref_fasta))
        logging.info('Reading previous cluster reference assignments from : {}'.format(ref_cluster_file))

        shutil.copy(ref_cluster_file, tmp_cluster_file)
        shutil.copy(ref_fasta, tmp_ref_fasta_file)
        update_existing(input_fasta, tmp_dir, ref_mash_db, tmp_cluster_file, header, tmp_ref_fasta_file, update_fasta)

        if args.overwrite:
            shutil.move(update_fasta,ref_fasta)
            shutil.move(tmp_cluster_file,ref_cluster_file)
            mash_db_file = "{}.msh".format(input_fasta)
            mObj = mash()
            mObj.mashsketch(input_fasta, mash_db_file, num_threads=num_threads)
            blast_runner = BlastRunner(ref_fasta, '')
            blast_runner.makeblastdb(ref_fasta, 'nucl')
    else:
        mashObj = mash()
        mashObj.mashsketch(input_fasta,input_fasta+".msh",num_threads=num_threads)
        distance_matrix_file = os.path.join(tmp_dir,'mash_dist_matrix.txt')
        mashfile_handle = open(distance_matrix_file,'w')

        mashObj.run_mash(input_fasta+'.msh', input_fasta+'.msh', mashfile_handle,table=True,num_threads=num_threads)
        clust_assignments = build_cluster_db(distance_matrix_file, (0.05, 0.0001))
        writeClusterAssignments(tmp_cluster_file, header, clust_assignments)
        clust_dict = selectCluster(clust_assignments, 1)
        shutil.copy(input_fasta, tmp_ref_fasta_file)
        updateFastaFile(tmp_ref_fasta_file ,update_fasta, clust_dict)

# call main function
if __name__ == '__main__':
    main()


import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
import matplotlib.pyplot as plt
from wordcloud import WordCloud
from math import log, sqrt
import pandas as pd
import numpy as np
import re
import pickle
from sklearn.externals import joblib
#print (1000)
#classifier = joblib.load('model.pkl')
#print(1)
#predict = classifier.predict(msg_test)
#print (2)
#print (predict)
from trained_model2 import classified
a = classified("my name ios apple")
print(a)
print("hi this is me")


# import numpy as np
# import gym
# import torch
# from typing import Dict, Any
# from copy import deepcopy
# from malib.algorithm.common.model import MLPCritic
# from malib.algorithm.common.policy import Policy
#
#
# class QMIX(Policy):
#     def __init__(
#         self,
#         registered_name: str,
#         observation_space: gym.spaces.Space,
#         action_space: gym.spaces.Space,
#         model_config: Dict[str, Any] = None,
#         custom_config: Dict[str, Any] = None,
#     ):
#         super(QMIX, self).__init__(
#             registered_name=registered_name,
#             observation_space=observation_space,
#             action_space=action_space,
#             model_config=model_config,
#             custom_config=custom_config,
#         )
#         self.eps_min = (
#             1e-2 if custom_config is None else custom_config.get("eps_min", 1e-5)
#         )
#         self.eps_max = (
#             1.0 if custom_config is None else custom_config.get("eps_max", 1.0)
#         )
#         self.eps_decay = (
#             100 if custom_config is None else custom_config.get("eps_decay", 100)
#         )
#         self.polyak = (
#             0.99 if custom_config is None else custom_config.get("polyak", 0.99)
#         )
#         self.delta = (self.eps_max - self.eps_min) / self.eps_decay
#         self.step = 0
#
#         self.obs_dim = self.preprocessor.size
#         self.act_dim = action_space.n
#
#         self.q = MLPCritic(self.obs_dim, self.act_dim, model_config)
#         self.q_targ = deepcopy(self.q)
#
#     def compute_actions(self, observation, **kwargs):
#         pass
#
#     def _calc_eps(self):
#         return max(self.eps_min, self.eps_max - self.delta * self.step)
#
#     def compute_action(self, observation, **kwargs):
#         # self.step += 1
#         if np.random.random() < self._calc_eps():
#             actions = range(self.action_space.n)
#             if "legal_moves" in kwargs:
#                 actions = kwargs["legal_moves"]
#             elif "action_mask" in kwargs:
#                 actions = np.where(kwargs["action_mask"] == 1)[0]
#             action = np.random.choice(actions)
#             action_prob = torch.zeros(self.action_space.n)
#             action_prob[action] = 1.0
#             return action, None, {"action_probs": action_prob}
#         probs = torch.softmax(self.q(observation), dim=-1)
#         if "legal_moves" in kwargs:
#             mask = torch.zeros_like(probs)
#             mask[kwargs["legal_moves"]] = 1
#             probs = mask * probs
#         elif "action_mask" in kwargs:
#             mask = torch.FloatTensor(kwargs["action_mask"])
#             probs = mask * probs
#         # probs = probs / probs.sum()
#         # action = Categorical(probs=probs).sample()
#         action = probs.argmax().view(1)
#
#         extra_info = {"action_probs": probs.detach().numpy()}
#         return action.item(), None, extra_info
#
#     def compute_q(self, obs):
#         return self.q(obs)
#
#     def compute_target_q(self, obs):
#         return self.q_targ(obs)
#
#     def get_parameters(self):
#         return self.q.parameters()
#
#     def update_target(self):
#         # self.q_targ.load_state_dict(self.q.state_dict())
#         with torch.no_grad():
#             for p, p_targ in zip(self.q.parameters(), self.q_targ.parameters()):
#                 p_targ.data.mul_(self.polyak)
#                 p_targ.data.add_((1 - self.polyak) * p.data)
#
#     def state_dict(self):
#         return {
#             "q": self.q.state_dict(),
#             "q_target": self.q_targ.state_dict(),
#             "step": self.step,
#         }
#
#     def set_weights(self, parameters):
#         self.q.load_state_dict(parameters["q"])
#         self.q_targ.load_state_dict(parameters["q_target"])
#         self.step = parameters["step"]
#
#     def train(self):
#         pass
#
#     def eval(self):
#         pass


from flask import *
from keras.utils import np_utils
import numpy as np
from glob import glob

from keras.applications.resnet50 import ResNet50, preprocess_input, decode_predictions
from keras.preprocessing import image
from keras.backend import set_session, get_session

from keras import backend as K
from tqdm import tqdm
import pickle
import cv2
import glob
import tensorflow as tf
import pandas as pd
import numpy as np
import os
from keras.models import model_from_json, Sequential, Model, load_model

from keras.layers import Input, Conv2D, MaxPooling2D, GlobalAveragePooling2D, Flatten, Dense, Lambda, ELU, Cropping2D, Dropout, BatchNormalization
from keras.layers.convolutional import Convolution2D, Conv2D
from keras.layers.pooling import MaxPooling2D
from keras.optimizers import Adam
#import matplotlib.pyplot as plan
from random import randint

import json
dog_names=[]
with open('data/dog_names.json') as json_file:
    dog_names = json.load(json_file)

cat_names=[]
with open('data/cat_names.json') as json_file:
    cat_names = json.load(json_file)
dog_describe = []
with open('data/dog_des.json', encoding='UTF8') as json_file:
    dog_describe = json.load(json_file)
cat_describe = []
with open('data/cat_des.json', encoding='UTF8') as json_file:
   cat_describe = json.load(json_file)

global sess
global graph
sess = tf.Session()
graph = tf.get_default_graph()
cat_graph = tf.get_default_graph()

with graph.as_default():
    set_session(sess)
    ResNet50_model_for_dog_breed = ResNet50(weights='imagenet')
    graph = tf.get_default_graph()

with graph.as_default():
    set_session(sess)
    Res_model_for_adjusting_shape = ResNet50(weights='imagenet', include_top=False)
    graph = tf.get_default_graph()

with graph.as_default():
    set_session(sess)
    bottleneck_features = np.load('bottleneck_features/DogResnet50Data.npz')
    train_Resnet = bottleneck_features['train']
    valid_Resnet = bottleneck_features['valid']
    test_Resnet = bottleneck_features['test']
    Resnet_Model = Sequential()
    Resnet_Model.add(GlobalAveragePooling2D(input_shape=train_Resnet.shape[1:]))
    Resnet_Model.add(Dense(133,activation='softmax'))
    Resnet_Model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    Resnet_Model.load_weights('saved_models/weights.best.Resnet.hdf5')
    graph = tf.get_default_graph()

with cat_graph.as_default():
    set_session(sess)
    ResNet50_model_for_cat_breed = ResNet50(weights='imagenet')
    cat_graph = tf.get_default_graph()

with cat_graph.as_default():
    set_session(sess)
    Res_model_for_adjusting_shape2 = ResNet50(weights='imagenet', include_top=False)
    cat_graph = tf.get_default_graph()

with cat_graph.as_default():
    set_session(sess)
    bottleneck_features = np.load('bottleneck_features/CatResnet50Data.npz')
    train_Resnet = bottleneck_features['train']
    valid_Resnet = bottleneck_features['valid']
    test_Resnet = bottleneck_features['test']
    Resnet_Model2 = Sequential()
    Resnet_Model2.add(GlobalAveragePooling2D(input_shape=train_Resnet.shape[1:]))
    Resnet_Model2.add(Dense(30,activation='softmax'))
    Resnet_Model2.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    Resnet_Model2.load_weights('saved_models/new_cat_classifier.h5')
    cat_graph = tf.get_default_graph()

def extract_Resnet50(tensor):
	return ResNet50(weights='imagenet', include_top=False).predict(preprocess_input(tensor))

#define generic function for pre-processing images into 4d tensor as input for CNN
def path_to_tensor(img_path):
    # loads RGB image as PIL.Image.Image type
    img = image.load_img(img_path, target_size=(224, 224))
    # convert PIL.Image.Image type to 3D tensor with shape (224, 224, 3)
    x = image.img_to_array(img)
    # convert 3D tensor to 4D tensor with shape (1, 224, 224, 3) and return 4D tensor
    return np.expand_dims(x, axis=0)

def extract_Resnet50(tensor):
	from keras.applications.resnet50 import ResNet50, preprocess_input
	return ResNet50(weights='imagenet', include_top=False).predict(preprocess_input(tensor))

#predicts the dog breed based on the pretrained ResNet50 models with weights from imagenet
def ResNet50_predict_labels(img_path):
    with graph.as_default():
        # returns prediction vector for dog image located at img_path
        img = preprocess_input(path_to_tensor(img_path))
        return np.argmax(ResNet50_model_for_dog_breed.predict(img))

#predicts the cat breed based on the pretrained ResNet50 models with weights from imagenet
def ResNet50_predict_cat_labels(img_path):
    with cat_graph.as_default():
        # returns prediction vector for dog image located at img_path
        img = preprocess_input(path_to_tensor(img_path))
        return np.argmax(ResNet50_model_for_cat_breed.predict(img))

### returns "True" if a dog is detected in the dog image stored at img_path
def dog_detector(img_path):
    prediction = ResNet50_predict_labels(img_path)
    print(prediction)
    return ((prediction <= 268) & (prediction >= 151))

### returns "True" if a dog is detected in the cat image stored at img_path
def cat_detector(img_path):
    prediction = ResNet50_predict_cat_labels(img_path)
    print(prediction)
    return ((prediction <= 285) & (prediction >= 281))

# returns "True" if face is detected in image stored at img_path
def face_detector(img_path):
    img = cv2.imread(img_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_cascade = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_alt.xml')
    faces = face_cascade.detectMultiScale(gray)
    print(len(faces))
    return len(faces) > 0

def Resnet_predict_breed(img_path):
    # extract bottleneck features
    #bottleneck_feature = extract_Resnet50(path_to_tensor(img_path))
    y = path_to_tensor(img_path)
    print(y.shape)
    y = preprocess_input(y)
    print(y.shape)
    x = Res_model_for_adjusting_shape.predict(y)
    #bottleneck_feature = Res_model_for_adjusting_shape.predict(y)
    print(x.shape)
    # obtain predicted vector
    predicted_vector = Resnet_Model.predict(x)
    # return dog breed that is predicted by the model
    return dog_names[np.argmax(predicted_vector)],dog_describe[np.argmax(predicted_vector)]

def Resnet_predict_cat_breed(img_path):
    # extract bottleneck features
    #bottleneck_feature = extract_Resnet50(path_to_tensor(img_path))
    y = path_to_tensor(img_path)
    print(y.shape)
    y = preprocess_input(y)
    print(y.shape)
    x = Res_model_for_adjusting_shape2.predict(y)
    #bottleneck_feature = Res_model_for_adjusting_shape.predict(y)
    print(x.shape)
    # obtain predicted vector
    predicted_vector = Resnet_Model2.predict(x)
    # return dog breed that is predicted by the model
    return cat_names[np.argmax(predicted_vector)],cat_describe[np.argmax(predicted_vector)]

def predict_image(img_path):
    vowels=["a","e","i","o","u"]
    #if a dog is detected in the image, return the predicted breed.
    with graph.as_default():
        set_session(sess)
        if dog_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_breed(img_path)
            predicted_breed = predicted_breed.rsplit('.',1)[1].replace("_", " ")
            return "The predicted dog breed is " + str(predicted_breed) + ".",des
        #if a human is detected in the image, return the resembling dog breed.
        if face_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_breed(img_path)
            predicted_breed=predicted_breed.rsplit('.',1)[1].replace("_", " ")
            return "This photo looks like " + str(predicted_breed) + ".", des
        #if neither is detected in the image, provide output that indicates an error.
        else:
            des ="Detect Failed"
            return "Detect Failed, Upload another picture.", des

def predict_cat_image(img_path):
    vowels=["a","e","i","o","u"]
    #if a dog is detected in the image, return the predicted breed.
    with cat_graph.as_default():
        set_session(sess)
        if cat_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_cat_breed(img_path)
            predicted_breed=predicted_breed.rsplit('_',1)[1].replace("_", " ")
            return "The predicted cat breed is " + str(predicted_breed) + ".", des
        #if a human is detected in the image, return the resembling dog breed.
        if face_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_cat_breed(img_path)
            predicted_breed = predicted_breed.rsplit('_', 1)[1].replace("_", " ")
            return "This photo looks like " + str(predicted_breed) + ".", des
        #if neither is detected in the image, provide output that indicates an error.
        else:
            des ="Detect Failed"
            return "Detect Failed, Upload another picture.", des

def result_dog(img_path):
    vowels=["a","e","i","o","u"]
    #if a dog is detected in the image, return the predicted breed.
    with graph.as_default():
        set_session(sess)
        if dog_detector(img_path)==True:
            predicted_breed,des = Resnet_predict_breed(img_path)
            predicted_breed= predicted_breed.rsplit('.',1)[1].replace("_", " ")
            txt_dog = str(predicted_breed)
            path_dir = './static/dogSimilar'
            file_list = os.listdir(path_dir)
            for i in range(len(file_list)):
                if txt_dog == file_list[i]:
                    path_dir = path_dir + "/" + txt_dog
                    file_list = os.listdir(path_dir)
                    print(file_list[0])
                    txt_dog = path_dir + "/" + file_list[0]
                    txt_dog = txt_dog[1:]
                    break
            return str(txt_dog)
        #if a human is detected in the image, return the resembling dog breed.
        if face_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_breed(img_path)
            predicted_breed = predicted_breed.rsplit('.', 1)[1].replace("_", " ")
            txt_dog = str(predicted_breed)
            path_dir = './static/dogSimilar'
            file_list = os.listdir(path_dir)
            for i in range(len(file_list)):
                if txt_dog == file_list[i]:
                    path_dir = path_dir + "/" + txt_dog
                    file_list = os.listdir(path_dir)
                    print(file_list[0])
                    txt_dog = path_dir + "/" + file_list[0]
                    txt_dog = txt_dog[1:]
                    break
            return str(txt_dog)
        #if neither is detected in the image, provide output that indicates an error.
        else:
            return "Detect Failed, Upload another picture."

def result_cat(img_path):
    vowels=["a","e","i","o","u"]
    #if a dog is detected in the image, return the predicted breed.
    with cat_graph.as_default():
        set_session(sess)
        if cat_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_cat_breed(img_path)
            predicted_breed=predicted_breed.rsplit('_',1)[1].replace("_", " ")
            txt_cat = str(predicted_breed)
            path_dir = './static/catSimilar'
            file_list = os.listdir(path_dir)
            for i in range(len(file_list)) :
                if txt_cat == file_list[i] :
                    path_dir = path_dir + "/" + txt_cat
                    file_list = os.listdir(path_dir)
                    print(file_list[0])
                    txt_cat = path_dir + "/" + file_list[0]
                    txt_cat = txt_cat[1:]
                    print(txt_cat)
                    break
            return str(txt_cat)
        #if a human is detected in the image, return the resembling dog breed.
        if face_detector(img_path)==True:
            predicted_breed, des = Resnet_predict_cat_breed(img_path)
            predicted_breed = predicted_breed.rsplit('_', 1)[1].replace("_", " ")
            txt_cat = str(predicted_breed)
            path_dir = './static/catSimilar'
            file_list = os.listdir(path_dir)
            for i in range(len(file_list)):
                if txt_cat == file_list[i]:
                    path_dir = path_dir + "/" + txt_cat
                    file_list = os.listdir(path_dir)
                    print(file_list[0])
                    txt_cat = path_dir + "/" + file_list[0]
                    txt_cat = txt_cat[1:]
                    print(txt_cat)
                    break
            return str(txt_cat)
        #if neither is detected in the image, provide output that indicates an error.
        else:
            return "Detect Failed, Upload another picture."

IMAGE_FOLDER = 'static/'
#PROCESSED_FOLDER = 'processed/'

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = IMAGE_FOLDER
#app.config['PROCESSED_FOLDER'] = PROCESSED_FOLDER

@app.route('/')
def upload():
	return render_template("index.html")

@app.route('/success', methods = ['POST'])
def success():
	if request.method == 'POST':
		f = request.files['file']
		f.save(os.path.join(app.config['UPLOAD_FOLDER'], f.filename))
		full_filename = os.path.join(app.config['UPLOAD_FOLDER'], f.filename)
		image_ext = cv2.imread(full_filename)
		img_path = full_filename
		#print(image_ext.shape)
		with graph.as_default():
			set_session(sess)
			txt,des = predict_image(img_path)
			txt_dog = result_dog(img_path)
			#txt = result
		final_text = 'Results after Detecting Dog Breed in Input Image'
		return render_template("success.html", name = final_text, img = full_filename, out_1 = txt, out_2 = f.filename, out_3 = txt_dog, out_4 = des)

@app.route('/cat_success', methods = ['POST'])
def cat_success():
	if request.method == 'POST':
		f = request.files['file']
		f.save(os.path.join(app.config['UPLOAD_FOLDER'], f.filename))
		full_filename = os.path.join(app.config['UPLOAD_FOLDER'], f.filename)
		image_ext = cv2.imread(full_filename)
		img_path = full_filename
		with cat_graph.as_default():
			set_session(sess)
			txt,des = predict_cat_image(img_path)
			txt_cat = result_cat(img_path)
			#txt = result
		final_text = 'Results after Detecting Cat Breed in Input Image'
		return render_template("cat_success.html", name = final_text, img = full_filename, out_1 = txt, out_2 = f.filename,out_3 = txt_cat, out_4 = des)

@app.route('/')
def info():
	return render_template("index.html")

@app.route('/dog')
def dog():
	return render_template("dog.html")

@app.route('/cat')
def cat():
	return render_template("cat.html")

if __name__ == '__main__':
	app.run(host="0.0.0.0",port=8080,debug=True)


# ----------------------------------------
# Written by Yude Wang
# ----------------------------------------

from tqdm import tqdm
from lib.net.sync_batchnorm import SynchronizedBatchNorm2d
from lib.utils.imutils import img_denorm
from lib.utils.finalprocess import writelog
from lib.utils.configuration import Configuration
from lib.net.sync_batchnorm.replicate import patch_replication_callback
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from torch import optim
from lib.net.generateNet import generate_net
from lib.datasets.generateData import generate_dataset
from experiment.deeplabv3_voc.config import config_dict
import time
import os
from torch import nn
import torch
import numpy as np
import random
torch.manual_seed(1)  # cpu
torch.cuda.manual_seed_all(1)  # gpu
np.random.seed(1)  # numpy
random.seed(1)  # random and transforms
torch.backends.cudnn.deterministic = True  # cudnn

cfg = Configuration(config_dict)

def train_net():
    period = 'train'
    transform = 'weak'
    dataset = generate_dataset(cfg, period=period, transform=transform)

    def worker_init_fn(worker_id):
        np.random.seed(1 + worker_id)
    dataloader = DataLoader(dataset,
                            batch_size=cfg.TRAIN_BATCHES,
                            shuffle=cfg.TRAIN_SHUFFLE,
                            num_workers=cfg.DATA_WORKERS,
                            pin_memory=True,
                            drop_last=True,
                            worker_init_fn=worker_init_fn)

    if cfg.GPUS > 1:
        net = generate_net(cfg, batchnorm=SynchronizedBatchNorm2d, dilated=cfg.MODEL_BACKBONE_DILATED,
                           multi_grid=cfg.MODEL_BACKBONE_MULTIGRID, deep_base=cfg.MODEL_BACKBONE_DEEPBASE)
    else:
        net = generate_net(cfg, batchnorm=nn.BatchNorm2d, dilated=cfg.MODEL_BACKBONE_DILATED,
                           multi_grid=cfg.MODEL_BACKBONE_MULTIGRID, deep_base=cfg.MODEL_BACKBONE_DEEPBASE)
    if cfg.TRAIN_CKPT:
        net.load_state_dict(torch.load(cfg.TRAIN_CKPT), strict=True)
        print('load pretrained model')
    if cfg.TRAIN_TBLOG:
        from tensorboardX import SummaryWriter
        # Set the Tensorboard logger
        tblogger = SummaryWriter(cfg.LOG_DIR)

    print('Use %d GPU' % cfg.GPUS)
    device = torch.device(0)
    if cfg.GPUS > 1:
        net = nn.DataParallel(net)
        patch_replication_callback(net)
        parameter_source = net.module
    else:
        parameter_source = net
    net.to(device)
    criterion = nn.CrossEntropyLoss(ignore_index=255)
    optimizer = optim.SGD(
        params=[
            {'params': get_params(parameter_source, key='backbone'), 'lr': cfg.TRAIN_LR},
            {'params': get_params(parameter_source, key='cls'),      'lr': 10*cfg.TRAIN_LR},
            {'params': get_params(parameter_source, key='others'),   'lr': cfg.TRAIN_LR}
        ],
        momentum=cfg.TRAIN_MOMENTUM,
        weight_decay=cfg.TRAIN_WEIGHT_DECAY
    )
    itr = cfg.TRAIN_MINEPOCH * len(dataset)//(cfg.TRAIN_BATCHES)
    max_itr = cfg.TRAIN_ITERATION
    max_epoch = max_itr*(cfg.TRAIN_BATCHES)//len(dataset)+1
    tblogger = SummaryWriter(cfg.LOG_DIR)
    criterion = nn.CrossEntropyLoss(ignore_index=255)
    scaler = torch.cuda.amp.GradScaler()
    with tqdm(total=max_itr) as pbar:
        for epoch in range(cfg.TRAIN_MINEPOCH, max_epoch):
            for i_batch, sample in enumerate(dataloader):

                now_lr = adjust_lr(optimizer, itr, max_itr, cfg.TRAIN_LR, cfg.TRAIN_POWER)
                optimizer.zero_grad()

                inputs, seg_label = sample['image'], sample['segmentation']
                n, c, h, w = inputs.size()

                with torch.cuda.amp.autocast():
                    pred1 = net(inputs.to(0))
                    loss = criterion(pred1, seg_label.to(0))
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()

                pbar.set_description("loss=%g " % (loss.item()))
                pbar.update(1)
                time.sleep(0.001)
                # print('epoch:%d/%d\tbatch:%d/%d\titr:%d\tlr:%g\tloss:%g' %
                #	(epoch, max_epoch, i_batch, len(dataset)//(cfg.TRAIN_BATCHES),
                #	itr+1, now_lr, loss.item()))
                if cfg.TRAIN_TBLOG and itr % 100 == 0:
                    inputs1 = img_denorm(inputs[-1].cpu().numpy()).astype(np.uint8)
                    label1 = sample['segmentation'][-1].cpu().numpy()
                    label_color1 = dataset.label2colormap(label1).transpose((2, 0, 1))

                    n, c, h, w = inputs.size()
                    seg_vis1 = torch.argmax(pred1[-1], dim=0).detach().cpu().numpy()
                    seg_color1 = dataset.label2colormap(seg_vis1).transpose((2, 0, 1))

                    tblogger.add_scalar('loss', loss.item(), itr)
                    tblogger.add_scalar('lr', now_lr, itr)
                    tblogger.add_image('Input', inputs1, itr)
                    tblogger.add_image('Label', label_color1, itr)
                    tblogger.add_image('SEG1', seg_color1, itr)
                itr += 1
                if itr >= max_itr:
                    break
            save_path = os.path.join(cfg.MODEL_SAVE_DIR, '%s_%s_%s_epoch%d.pth' %
                                     (cfg.MODEL_NAME, cfg.MODEL_BACKBONE, cfg.DATA_NAME, epoch))
            torch.save(parameter_source.state_dict(), save_path)
            print('%s has been saved' % save_path)
            remove_path = os.path.join(cfg.MODEL_SAVE_DIR, '%s_%s_%s_epoch%d.pth' %
                                       (cfg.MODEL_NAME, cfg.MODEL_BACKBONE, cfg.DATA_NAME, epoch-1))
            if os.path.exists(remove_path):
                os.remove(remove_path)

    save_path = os.path.join(cfg.MODEL_SAVE_DIR, '%s_%s_%s_itr%d_all.pth' % (
        cfg.MODEL_NAME, cfg.MODEL_BACKBONE, cfg.DATA_NAME, cfg.TRAIN_ITERATION))
    torch.save(parameter_source.state_dict(), save_path)
    if cfg.TRAIN_TBLOG:
        tblogger.close()
    print('%s has been saved' % save_path)
    writelog(cfg, period)

def adjust_lr(optimizer, itr, max_itr, lr_init, power):
    now_lr = lr_init * (1 - itr/(max_itr+1)) ** power
    optimizer.param_groups[0]['lr'] = now_lr
    optimizer.param_groups[1]['lr'] = 10*now_lr
    optimizer.param_groups[2]['lr'] = now_lr
    return now_lr

def get_params(model, key):
    for m in model.named_modules():
        if key == 'backbone':
            if ('backbone' in m[0]) and isinstance(m[1], (nn.Conv2d, SynchronizedBatchNorm2d, nn.BatchNorm2d, nn.InstanceNorm2d)):
                for p in m[1].parameters():
                    yield p
        elif key == 'cls':
            if ('cls_conv' in m[0]) and isinstance(m[1], (nn.Conv2d, SynchronizedBatchNorm2d, nn.BatchNorm2d, nn.InstanceNorm2d)):
                for p in m[1].parameters():
                    yield p
        elif key == 'others':
            if ('backbone' not in m[0] and 'cls_conv' not in m[0]) and isinstance(m[1], (nn.Conv2d, SynchronizedBatchNorm2d, nn.BatchNorm2d, nn.InstanceNorm2d)):
                for p in m[1].parameters():
                    yield p

if __name__ == '__main__':
    train_net()


# -*- coding: utf-8 -*-
"""
Created on Thu Mar 23 18:24:57 2017

@author: 93137
"""
import numpy as np
print np.log10(102323)
print int(round(0))