"""Functions to generate batches for the reinforcement learning part.
Mainly intended for training, though during the playing phase, the same
functions are used."""
from __future__ import print_function, division

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

import models as models_reinforced

import cv2
from config import Config
import imgaug as ia
from lib.util import to_variable, to_cuda, to_numpy
from lib import util
from lib import actions as actionslib
from lib import replay_memory
import numpy as np
from scipy import misc
import multiprocessing
import threading
import random
import time

if sys.version_info[0] == 2:
    import cPickle as pickle
    from Queue import Full as QueueFull
elif sys.version_info[0] == 3:
    import pickle
    from queue import Full as QueueFull
    xrange = range

GPU = Config.GPU
NB_REWARD_BINS = 101

class BatchData(object):
    """Method encapsulating the data of a single batch.

    TODO some of the functions are named like properties, rename
    """

    def __init__(self, curr_idx, images_by_timestep, images_prev_by_timestep, multiactions, rewards, speeds, is_reverse, steering_wheel, steering_wheel_raw, previous_states_distances):
        self.curr_idx = curr_idx
        self.images_by_timestep = images_by_timestep
        self.images_prev_by_timestep = images_prev_by_timestep
        self.multiactions = multiactions
        self.rewards = rewards
        self.speeds = speeds
        self.is_reverse = is_reverse
        self.steering_wheel = steering_wheel
        self.steering_wheel_raw = steering_wheel_raw
        self.previous_states_distances = previous_states_distances

    @property
    def batch_size(self):
        return self.images_by_timestep.shape[1]

    @property
    def nb_future(self):
        return self.images_prev_by_timestep.shape[0] - 1

    @property
    def nb_prev_per_image(self):
        return self.images_prev_by_timestep.shape[2]

    def reward_bin_idx(self, timestep, inbatch_idx):
        timestep = self.curr_idx + timestep
        reward = self.rewards[timestep, inbatch_idx]
        reward_norm = (reward - Config.MIN_REWARD) / (Config.MAX_REWARD - Config.MIN_REWARD)
        reward_norm = 1 - reward_norm # top to bottom
        rewbin = np.clip(int(reward_norm * NB_REWARD_BINS), 0, NB_REWARD_BINS-1) # clip here, because MAX_REWARD ends up giving bin NB_REWARD_BINS, which is 1 too high
        return rewbin

    def rewards_bins(self, timestep):
        timestep = self.curr_idx + timestep
        T, B = self.rewards.shape
        result = np.zeros((B, NB_REWARD_BINS), dtype=np.float32)
        for b in xrange(B):
            rewbin = self.reward_bin_idx(timestep-self.curr_idx, b)
            result[b, rewbin] = 1
        return result

    def rewards_bins_all(self):
        T, B = self.rewards.shape
        bins_over_time = [self.rewards_bins(t) for t in xrange(-self.curr_idx, T-self.curr_idx)]
        return np.array(bins_over_time, dtype=np.float32)

    def inputs_supervised(self, volatile=False, requires_grad=True, gpu=GPU):
        images = to_cuda(to_variable(self.images_by_timestep[0], volatile=volatile, requires_grad=requires_grad), gpu)
        images_prev = to_cuda(to_variable(self.images_prev_by_timestep[0], volatile=volatile, requires_grad=requires_grad), gpu)
        return images, images_prev

    def inputs_reinforced_add_numpy(self, timestep=0):
        timestep = self.curr_idx + timestep
        B = self.batch_size

        prev_indices_exclusive = [timestep - d for d in self.previous_states_distances]
        prev_indices_inclusive = [timestep] + prev_indices_exclusive

        ma_vecs = np.zeros((self.batch_size, len(prev_indices_exclusive), 9), dtype=np.float32)
        for i, idx in enumerate(prev_indices_exclusive):
            mas = self.multiactions[idx]
            for b, ma in enumerate(mas):
                ma_vecs[b, i, :] = actionslib.ACTIONS_TO_MULTIVEC[ma]
        ma_vecs = ma_vecs.reshape(self.batch_size, -1) # (B, P*9) with P=number of previous images

        speeds = self.speeds[prev_indices_inclusive, :]
        steering_wheel = (self.steering_wheel[prev_indices_inclusive, :] - Config.STEERING_WHEEL_CNN_MIN) / (Config.STEERING_WHEEL_CNN_MAX - Config.STEERING_WHEEL_CNN_MIN)
        steering_wheel_raw = (self.steering_wheel_raw[prev_indices_inclusive, :] - Config.STEERING_WHEEL_RAW_CNN_MIN) / (Config.STEERING_WHEEL_RAW_CNN_MAX - Config.STEERING_WHEEL_RAW_CNN_MIN)
        vals = {
            "speeds": np.squeeze(np.clip(speeds / Config.MAX_SPEED, 0, 1)),
            "is_reverse": np.squeeze(self.is_reverse[prev_indices_inclusive, :]),
            "steering_wheel": np.squeeze(steering_wheel*2 - 1),
            "steering_wheel_raw": np.squeeze(steering_wheel_raw*2 - 1),
            "multiactions_vecs": ma_vecs
        }
        if B == 1:
            vals["speeds"] = vals["speeds"][:, np.newaxis]
            vals["is_reverse"] = vals["is_reverse"][:, np.newaxis]
            vals["steering_wheel"] = vals["steering_wheel"][:, np.newaxis]
            vals["steering_wheel_raw"] = vals["steering_wheel_raw"][:, np.newaxis]
        vals["speeds"] = vals["speeds"].transpose((1, 0)) # (P, B) => (B, P) with P=number of previous images
        vals["is_reverse"] = vals["is_reverse"].transpose((1, 0)) # (P, B) => (B, P) with P=number of previous images
        vals["steering_wheel"] = vals["steering_wheel"].transpose((1, 0)) # (P, B) => (B, P) with P=number of previous images
        vals["steering_wheel_raw"] = vals["steering_wheel_raw"].transpose((1, 0)) # (P, B) => (B, P) with P=number of previous images

        return vals

    def inputs_reinforced_add(self, volatile=False, requires_grad=True, gpu=GPU):
        return to_cuda(to_variable(self.inputs_reinforced_add_numpy(), volatile=volatile, requires_grad=requires_grad), gpu)

    def future_inputs_supervised(self, volatile=False, requires_grad=True, gpu=GPU):
        images = to_cuda(to_variable(self.images_by_timestep[1:], volatile=volatile, requires_grad=requires_grad), gpu)
        images_prev = to_cuda(to_variable(self.images_prev_by_timestep[1:], volatile=volatile, requires_grad=requires_grad), gpu)
        return images, images_prev

    def future_reinforced_add(self, volatile=False, requires_grad=True, gpu=GPU):
        vals = {
            "speeds": [],
            "is_reverse": [],
            "steering_wheel": [],
            "steering_wheel_raw": [],
            "multiactions_vecs": []
        }
        for timestep in xrange(1, self.nb_future+1):
            inputs_ts = self.inputs_reinforced_add_numpy(timestep=timestep)
            vals["speeds"].append(inputs_ts["speeds"])
            vals["is_reverse"].append(inputs_ts["is_reverse"])
            vals["steering_wheel"].append(inputs_ts["steering_wheel"])
            vals["steering_wheel_raw"].append(inputs_ts["steering_wheel_raw"])
            vals["multiactions_vecs"].append(inputs_ts["multiactions_vecs"])
        vals["speeds"] = np.array(vals["speeds"], dtype=np.float32)
        vals["is_reverse"] = np.array(vals["is_reverse"], dtype=np.float32)
        vals["steering_wheel"] = np.array(vals["steering_wheel"], dtype=np.float32)
        vals["steering_wheel_raw"] = np.array(vals["steering_wheel_raw"], dtype=np.float32)
        vals["multiactions_vecs"] = np.array(vals["multiactions_vecs"], dtype=np.float32)

        T, B, _ = vals["speeds"].shape
        vals_flat = {
            "speeds": vals["speeds"].reshape((T*B, -1)),
            "is_reverse": vals["is_reverse"].reshape((T*B, -1)),
            "steering_wheel": vals["steering_wheel"].reshape((T*B, -1)),
            "steering_wheel_raw": vals["steering_wheel_raw"].reshape((T*B, -1)),
            "multiactions_vecs": vals["multiactions_vecs"].reshape((T*B, -1))
        }

        return to_cuda(to_variable(vals_flat, volatile=volatile, requires_grad=requires_grad), gpu)

    def inputs_successor_multiactions_vecs(self, volatile=False, requires_grad=True, gpu=GPU):
        # the successor gets in actions a and has to predict the next
        # state, i.e. for tuples (s, a, r, s') it gets a and predicts s',
        # hence the future actions here start at curr_idx (current state index)
        # and end at -1
        arr = models_reinforced.SuccessorPredictor.multiactions_to_vecs(self.multiactions[self.curr_idx:-1])
        assert arr.shape == (self.nb_future, self.batch_size, 9)
        return to_cuda(to_variable(arr, volatile=volatile, requires_grad=requires_grad), gpu)

    def direct_rewards_values(self, volatile=False, requires_grad=True, gpu=GPU):
        rews = self.rewards[self.curr_idx, :][:,np.newaxis]
        rews = np.tile(rews, (1, 9))
        return to_cuda(to_variable(rews, volatile=volatile, requires_grad=requires_grad), gpu)

    def future_direct_rewards_values(self, volatile=False, requires_grad=True, gpu=GPU):
        rews = self.rewards[self.curr_idx+1:, :][:, :, np.newaxis]
        rews = np.tile(rews, (1, 1, 9))
        return to_cuda(to_variable(rews, volatile=volatile, requires_grad=requires_grad), gpu)

    def outputs_dr_gt(self, volatile=False, requires_grad=True, gpu=GPU):
        # for a tuple (s, a, r, s'), the reward r is ought to be predicted
        # that is, the reward for the previous action, which is dependent
        # on the new state s' that it created
        # it is saved at the previous timestep, i.e. at state s, hence
        # here -1
        bins = self.rewards_bins(-1)
        return to_cuda(to_variable(bins, volatile=volatile, requires_grad=requires_grad), gpu)

    def outputs_dr_future_gt(self, volatile=False, requires_grad=True, gpu=GPU):
        # starting at curr_idx and ending at -1 here for the same reason
        # as above
        bins = self.rewards_bins_all()
        bins = bins[self.curr_idx:-1]
        return to_cuda(to_variable(bins, volatile=volatile, requires_grad=requires_grad), gpu)

    def outputs_ae_gt(self, volatile=False, requires_grad=True, gpu=GPU):
        imgs = self.images_by_timestep[0, ...]
        imgs = np.clip(imgs*255, 0, 255).astype(np.uint8).transpose((0, 2, 3, 1))
        imgs_rs = ia.imresize_many_images(imgs, (45, 80))
        imgs_rs = (imgs_rs / 255.0).astype(np.float32).transpose((0, 3, 1, 2))
        return to_cuda(to_variable(imgs_rs, volatile=volatile, requires_grad=requires_grad), gpu)

    def chosen_action_indices(self):
        mas_timestep = self.multiactions[self.curr_idx]
        indices = [np.argmax(actionslib.ACTIONS_TO_MULTIVEC[ma]) for ma in mas_timestep]
        return indices

    def chosen_action_indices_future(self):
        indices_by_timestep = []
        for t_idx in xrange(self.nb_future):
            mas_timestep = self.multiactions[t_idx]
            indices = [np.argmax(actionslib.ACTIONS_TO_MULTIVEC[ma]) for ma in mas_timestep]
            indices_by_timestep.append(indices)
        return indices_by_timestep

    def draw(self, timestep=0, inbatch_idx=0):
        timestep = self.curr_idx + timestep
        img = self.images_by_timestep[timestep-self.curr_idx, inbatch_idx, :, :, :]
        img = (img.transpose((1, 2, 0))*255).astype(np.uint8)
        imgs_prev = self.images_prev_by_timestep[timestep-self.curr_idx, inbatch_idx, :, :, :]
        imgs_prev = (imgs_prev.transpose((1, 2, 0))*255).astype(np.uint8)

        h, w = img.shape[0:2]
        imgs_viz = [img] + [np.tile(imgs_prev[..., i][:, :, np.newaxis], (1, 1, 3)) for i in xrange(imgs_prev.shape[2])]
        imgs_viz = [ia.imresize_single_image(im, (h, w), interpolation="cubic") for im in imgs_viz]
        imgs_viz = np.hstack(imgs_viz)

        rewards_bins = self.rewards_bins_all()
        mas = [self.multiactions[i][inbatch_idx] for i in xrange(timestep-self.nb_prev_per_image, timestep)]
        pos = [timestep] + [timestep-d for d in self.previous_states_distances]
        reinforced_add = self.inputs_reinforced_add_numpy(timestep=timestep-self.curr_idx)
        outputs_dr_gt = self.outputs_dr_gt()[inbatch_idx]
        texts = [
            "pos: " + " ".join([str(i) for i in pos]),
            "Rewards:      " + " ".join(["%.2f" % (self.rewards[i, inbatch_idx],) for i in pos]),
            "Rewards bins: " + " ".join(["%d" % (np.argmax(rewards_bins[i, inbatch_idx]),) for i in pos]),
            "Speeds:       " + " ".join(["%.2f" % (self.speeds[i, inbatch_idx],) for i in pos]),
            "Multiactions: " + " ".join(["%s%s" % (ma[0], ma[1]) for ma in mas]),
            "Speeds RA:    " + " ".join(["%.3f" % (reinforced_add["speeds"][inbatch_idx, i],) for i in xrange(reinforced_add["speeds"].shape[1])]),
            "outputs_dr_gt[t=-1]: " + "%d" % (np.argmax(to_numpy(outputs_dr_gt)),)
        ]
        texts = "\n".join(texts)

        result = np.zeros((imgs_viz.shape[0]*3, imgs_viz.shape[1], 3), dtype=np.uint8)
        util.draw_image(result, x=0, y=0, other_img=imgs_viz, copy=False)
        result = util.draw_text(result, x=0, y=imgs_viz.shape[0]+4, text=texts, size=9)
        return result

def states_to_batch(previous_states_list, states_list, augseq, previous_states_distances, model_height, model_width, model_prev_height, model_prev_width):
    """Convert multiple chains of states into a batch.

    Parameters
    ----------
    previous_states_list : list of list of State
        Per chain of states a list of the previous states.
        First index of the list is the batch index,
        second index is the timestep. The oldest states come first.
    states_list : list of list of State
        Per chain of states a list of states that contain the "current"
        state at the start, followed by future states.
        First index is batch index, second timestep.
    augseq : Augmenter
        Sequence of augmenters to apply to each image. Use Noop() to make
        no changes.
    previous_states_distances : list of int
        List of distances relative to the current state. Each distance
        refers to one previous state to add to the model input.
        E.g. [2, 1] adds the state 200ms and 100ms before the current "state".
    model_height : int
        Height of the model input images (current state).
    model_width : int
        Width of the model input images (current state).
    model_prev_height : int
        Height of the model input images (previous states).
    model_prev_width : int
        Width of the model input images (previous states).

    Returns
    ----------
    List of BatchData
    """
    assert isinstance(previous_states_list, list)
    assert isinstance(states_list, list)
    assert isinstance(previous_states_list[0], list)
    assert isinstance(states_list[0], list)
    assert len(previous_states_list) == len(states_list)

    B = len(states_list)
    H, W = model_height, model_width
    Hp, Wp = model_prev_height, model_prev_width

    nb_prev_load = max(previous_states_distances)
    nb_future_states = len(states_list[0]) - 1
    nb_timesteps = nb_prev_load + 1 + nb_future_states
    #images = np.zeros((nb_timesteps, B, H, W, 3), dtype=np.uint8)
    #images_gray = np.zeros((nb_timesteps, B, Hp, Wp), dtype=np.float32)
    images_by_timestep = np.zeros((1+nb_future_states, B, H, W, 3), dtype=np.float32)
    images_gray = np.zeros((nb_timesteps, B, Hp, Wp), dtype=np.float32)
    multiactions = [[] for i in xrange(nb_timesteps)]
    rewards = np.zeros((nb_timesteps, B), dtype=np.float32)
    speeds = np.zeros((nb_timesteps, B), dtype=np.float32)
    is_reverse = np.zeros((nb_timesteps, B), dtype=np.float32)
    steering_wheel = np.zeros((nb_timesteps, B), dtype=np.float32)
    steering_wheel_raw = np.zeros((nb_timesteps, B), dtype=np.float32)

    augseqs_det = [augseq.to_deterministic() for _ in xrange(len(states_list))]

    for b, (previous_states, states) in enumerate(zip(previous_states_list, states_list)):
        augseq_det = augseqs_det[b]

        all_states = previous_states + states
        for t, state in enumerate(all_states):
            imgy = cv2.cvtColor(state.screenshot_rs, cv2.COLOR_RGB2GRAY)
            imgy_rs = downscale(imgy, Hp, Wp)
            imgy_rs_aug = augseq_det.augment_image(imgy_rs)
            images_gray[t, b, ...] = imgy_rs

            multiactions[t].append(state.multiaction)
            rewards[t, b] = state.reward
            if state.speed is not None:
                speeds[t, b] = state.speed
            if state.is_reverse is not None:
                is_reverse[t, b] = int(state.is_reverse)
            if state.steering_wheel_cnn is not None:
                steering_wheel[t, b] = state.steering_wheel_cnn
            if state.steering_wheel_raw_cnn is not None:
                steering_wheel_raw[t, b] = state.steering_wheel_raw_cnn
    images_gray = images_gray[..., np.newaxis]

    for b, states in enumerate(states_list):
        augseq_det = augseqs_det[b]

        for i, state in enumerate(states):
            state = states[i]
            images_by_timestep[i, b, ...] = augseq_det.augment_image(downscale(state.screenshot_rs, H, W))

    nb_prev_per_img = len(previous_states_distances)
    images_prev_by_timestep = np.zeros((1+nb_future_states, B, Hp, Wp, nb_prev_per_img), dtype=np.float32)
    for t in xrange(1 + nb_future_states):
        indices = [nb_prev_load+t-d for d in previous_states_distances]
        prev = images_gray[indices]
        prev = prev.transpose((1, 2, 3, 4, 0)).reshape((B, Hp, Wp, nb_prev_per_img))
        images_prev_by_timestep[t] = prev
    images_by_timestep = (images_by_timestep.astype(np.float32) / 255.0).transpose((0, 1, 4, 2, 3))
    images_prev_by_timestep = (images_prev_by_timestep.astype(np.float32) / 255.0).transpose((0, 1, 4, 2, 3))

    return BatchData(nb_prev_load, images_by_timestep, images_prev_by_timestep, multiactions, rewards, speeds, is_reverse, steering_wheel, steering_wheel_raw, previous_states_distances)

def downscale(im, h, w):
    if im.ndim == 2:
        im = im[:, :, np.newaxis]
        return np.squeeze(ia.imresize_single_image(im, (h, w), interpolation="cubic"))
    else:
        return ia.imresize_single_image(im, (h, w), interpolation="cubic")

class BatchLoader(object):
    """Class to load batches from the replay memory."""

    def __init__(self, val, batch_size, augseq, previous_states_distances, nb_future_states, model_height, model_width, model_prev_height, model_prev_width):
        self.val = val
        self.batch_size = batch_size
        self.augseq = augseq.deepcopy()
        self.augseq.reseed(random.randint(0, 10**6))
        self.previous_states_distances = previous_states_distances
        self.nb_future_states = nb_future_states
        self.model_height = model_height
        self.model_width = model_width
        self.model_prev_height = model_prev_height
        self.model_prev_width = model_prev_width
        self._memory = None

    def load_random_batch(self):
        if self._memory is None:
            self._memory = replay_memory.ReplayMemory.create_instance_reinforced(val=self.val)
            self._memory.update_caches()
            print("Connected memory to %s, idmin=%d, idmax=%d" % ("val" if self.val else "train", self._memory.id_min, self._memory.id_max))
        memory = self._memory

        nb_prev = max(self.previous_states_distances)
        nb_timesteps = nb_prev + 1 + self.nb_future_states

        previous_states_list = []
        states_list = []
        for b in xrange(self.batch_size):
            statechain = memory.get_random_state_chain(nb_timesteps)
            previous_states_list.append(statechain[:nb_prev])
            states_list.append(statechain[nb_prev:])

        return states_to_batch(previous_states_list, states_list, self.augseq, self.previous_states_distances, self.model_height, self.model_width, self.model_prev_height, self.model_prev_width)

class BackgroundBatchLoader(object):
    """Class that takes a BatchLoader and executes it many times in background
    processes."""

    def __init__(self, batch_loader, queue_size, nb_workers, threaded=False):
        self.queue = multiprocessing.Queue(queue_size)
        self.workers = []
        self.exit_signal = multiprocessing.Event()
        for i in range(nb_workers):
            seed = random.randint(1, 10**6)
            if threaded:
                worker = threading.Thread(target=self._load_batches, args=(batch_loader, self.queue, self.exit_signal, None))
            else:
                worker = multiprocessing.Process(target=self._load_batches, args=(batch_loader, self.queue, self.exit_signal, seed))
            worker.daemon = True
            worker.start()
            self.workers.append(worker)

    def get_batch(self):
        return pickle.loads(self.queue.get())

    def _load_batches(self, batch_loader, queue, exit_signal, seed=None):
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)
            batch_loader.augseq.reseed(seed)
            ia.seed(seed)

        while not exit_signal.is_set():
            batch = batch_loader.load_random_batch()

            start_time = time.time()
            batch_str = pickle.dumps(batch, protocol=-1)
            added_to_queue = False # without this, it will add the batch countless times to the queue
            while not added_to_queue and not exit_signal.is_set():
                try:
                    queue.put(batch_str, timeout=1)
                    added_to_queue = True
                except QueueFull as e:
                    pass
            end_time = time.time()
        batch_loader._memory.close()

    def join(self):
        self.exit_signal.set()
        time.sleep(5)

        while not self.queue.empty():
            _ = self.queue.get()
        #self.queue.join()

        for worker in self.workers:
            #worker.join()
            worker.terminate()

if __name__ == "__main__":
    from scipy import misc
    from imgaug import augmenters as iaa

    MODEL_HEIGHT = 90
    MODEL_WIDTH = 160
    MODEL_PREV_HEIGHT = 45
    MODEL_PREV_WIDTH = 80

    loader = BatchLoader(
        val=False, batch_size=8, augseq=iaa.Noop(),
        previous_states_distances=[2, 4, 6, 8, 10],
        nb_future_states=10,
        model_height=MODEL_HEIGHT, model_width=MODEL_WIDTH,
        model_prev_height=MODEL_PREV_HEIGHT, model_prev_width=MODEL_PREV_HEIGHT
    )
    for _ in xrange(1000):
        for t in xrange(3):
            imgs = []
            for b in xrange(3):
                print(t, b)
                batch = loader.load_random_batch()
                imgs.append(batch.draw(timestep=t, inbatch_idx=b))
            misc.imshow(np.vstack(imgs))
