from random import shuffle
from random import seed
from threading import Thread
import numpy as np
import time
import tensorflow as tf
  import Queue as Q  # ver. < 3.0
except ImportError:
  import queue as Q
from sklearn.preprocessing import normalize

PriorityQueue = Q.PriorityQueue

# Custom prioriy queue that is able to clear the queue once it is full
# and cut it to half. Therefore, everytime size of buffer is dqn_replay_buffer_size,
# we will keep half of the most valuable states and remove the rest to provide
# space for new experiences
class CustomQueue(PriorityQueue):
  A custom queue subclass that provides a :meth:`clear` method.
  def __init__(self, size):
    PriorityQueue.__init__(self, size)

  def clear(self):
    Clears all items from the queue.
    with self.mutex:
      unfinished = self.unfinished_tasks - len(self.queue)
      if unfinished <= 0:
        if unfinished < 0:
          raise ValueError('task_done() called too many times')
      self.queue = self.queue[0:len(self.queue)/2]
      self.unfinished_tasks = unfinished + len(self.queue)

  def isempty(self):
    with self.mutex:
      return len(self.queue) == 0

  def isfull(self):
    with self.mutex:
      return len(self.queue) == self.maxsize

class Transition(object):
  A class for holding the experiences collected from seq2seq model
  def __init__(self, state, action, state_prime, action_prime, reward, q_value, done):
        state: current decoder output state
        action: the greedy action selected from current decoder output
        state_prime: next decoder output state
        reward: reward of the greedy action selected
        q_value: Q-value of the greedy action selected
        done: whether we reached End-Of-Sequence or not
    self.state = state # size: dqn_input_feature_len
    self.action = action # size: 1
    self.state_prime = state_prime # size: dqn_input_feature_len
    self.action_prime = action_prime
    self.reward = reward # size: vocab_size
    self.q_value = q_value # size: vocab_size
    self.done = done # true/false

  def __cmp__(self, item):
    """ PriorityQueue uses this functino to sort the rewards
        We sort the queue such that items with higher rewards are in the head of max-heap
    return cmp(item.reward, self.reward) # bigger numbers have more priority

class ReplayBatch(object):
  """ A class for creating batches required for training DDQN. """

  def __init__(self, hps, example_list, dqn_batch_size, use_state_prime = False, max_art_oovs = 0):
       hps: seq2seq model parameters
       example_list: list of experiences
       dqn_batch_size: DDQN batch size
       use_state_prime: whether to use the next decoder state to make the batch or the current one
       max_art_oovs: number of OOV tokens in current batch

        _x: The input to DDQN model for training, this is basically the decoder output (dqn_batch_size, dqn_input_feature_len)
        _y: The Q-estimation (dqn_batch_size, vocab_size)
        _y_extended: The Q-estimation (dqn_batch_size, vocab_size + max_art_oovs)
    self._x = np.zeros((dqn_batch_size, hps.dqn_input_feature_len))
    self._y = np.zeros((dqn_batch_size, hps.vocab_size))
    self._y_extended = np.zeros((dqn_batch_size, hps.vocab_size + max_art_oovs))
    for i,e in enumerate(example_list):
      if use_state_prime:
        self._y[i,:]=normalize(e.q_value[0:hps.vocab_size], axis=1, norm='l1')
      if max_art_oovs == 0:
        self._y_extended[i,:] = normalize(e.q_value[0:hps.vocab_size], axis=1, norm='l1')
        self._y_extended[i,:] = e.q_value

class ReplayBuffer(object):
  """ A class for implementing the priority experience buffer. """

  BATCH_QUEUE_MAX = 100 # max number of batches the batch_queue can hold

  def __init__(self, hps):
    self._hps = hps
    self._buffer = CustomQueue(self._hps.dqn_replay_buffer_size)

    self._batch_queue = Q.Queue(self.BATCH_QUEUE_MAX)
    self._example_queue = Q.Queue(self.BATCH_QUEUE_MAX * self._hps.dqn_batch_size)
    self._num_example_q_threads = 1 # num threads to fill example queue
    self._num_batch_q_threads = 1  # num threads to fill batch queue
    self._bucketing_cache_size = 100 # how many batches-worth of examples to load into cache before bucketing

    # Start the threads that load the queues
    self._example_q_threads = []
    for _ in range(self._num_example_q_threads):
      self._example_q_threads[-1].daemon = True
    self._batch_q_threads = []
    for _ in range(self._num_batch_q_threads):
      self._batch_q_threads[-1].daemon = True

    # Start a thread that watches the other threads and restarts them if they're dead
    self._watch_thread = Thread(target=self.watch_threads)
    self._watch_thread.daemon = True

  def next_batch(self):
    """Return a Batch from the batch queue.

    If mode='decode' then each batch contains a single example repeated beam_size-many times; this is necessary for beam search.

      batch: a Batch object, or None if we're in single_pass mode and we've exhausted the dataset.
    # If the batch queue is empty, print a warning
    if self._batch_queue.qsize() == 0:
      tf.logging.warning('Bucket input queue is empty when calling next_batch. Bucket queue size: %i, Input queue size: %i', self._batch_queue.qsize(), self._example_queue.qsize())
      return None

    batch = self._batch_queue.get() # get the next Batch
    return batch

  def create_batch(_hps, batch, batch_size, use_state_prime=False, max_art_oovs=0):
    """ Create a DDQN-compatible batch from the input transitions

        _hps: seq2seq model parameters
        batch: a list of Transitions
        dqn_batch_size: DDQN batch size
        use_state_prime: whether to use the next decoder state to make the batch or the current one
        max_art_oovs: number of OOV tokens in current batch

        An object of ReplayBatch class

    return ReplayBatch(_hps, batch, batch_size, use_state_prime, max_art_oovs)

  def fill_example_queue(self):
    """Reads data from file and processes into Examples which are then placed into the example queue."""
    while True:
        input_gen = self._example_generator().next()
      except StopIteration: # if there are no more examples:"The example generator for this example queue filling thread has exhausted data.")
        raise Exception("single_pass mode is off but the example generator is out of data; error.")
      self._example_queue.put(input_gen) # place the pair in the example queue.

  def fill_batch_queue(self):
    """Takes Examples out of example queue, sorts them by encoder sequence length, processes into Batches and places them in the batch queue."""
    while True:
      # Get bucketing_cache_size-many batches of Examples into a list, then sort
      inputs = []
      for _ in range(self._hps.dqn_batch_size * self._bucketing_cache_size):

      # feed back all the samples to the buffer

      # Group the sorted Examples into batches, optionally shuffle the batches, and place in the batch queue.
      batches = []
      for i in range(0, len(inputs), self._hps.dqn_batch_size):
        batches.append(inputs[i:i + self._hps.dqn_batch_size])
      for b in batches:  # each b is a list of Example objects
        self._batch_queue.put(ReplayBatch(self._hps, b, self._hps.dqn_batch_size))

  def watch_threads(self):
    """Watch example queue and batch queue threads and restart if dead."""
    while True:
      for idx,t in enumerate(self._example_q_threads):
        if not t.is_alive(): # if the thread is dead
          tf.logging.error('Found example queue thread dead. Restarting.')
          new_t = Thread(target=self.fill_example_queue)
          self._example_q_threads[idx] = new_t
          new_t.daemon = True
      for idx,t in enumerate(self._batch_q_threads):
        if not t.is_alive(): # if the thread is dead
          tf.logging.error('Found batch queue thread dead. Restarting.')
          new_t = Thread(target=self.fill_batch_queue)
          self._batch_q_threads[idx] = new_t
          new_t.daemon = True

  def add(self, items):
    """ Adding a list of experiences to the buffer. When buffer is full,
      we get rid of half of the least important experiences and keep the rest.

        items: A list of experiences of size (batch_size, k, max_dec_steps, hidden_dim)
    for item in items:
      if not self._buffer.isfull():
        print('Replay Buffer is full, getting rid of unimportant transitions...')
    print('ReplayBatch size: {}'.format(self._buffer.qsize()))
    print('ReplayBatch example queue size: {}'.format(self._example_queue.qsize()))
    print('ReplayBatch batch queue size: {}'.format(self._batch_queue.qsize()))

  def _buffer_len(self):
    return self._buffer.qsize()

  def _example_generator(self):
    while True:
      if not self._buffer.isempty():
        item = self._buffer.get_nowait()
        yield item