import numpy as np
import copy
from pybullet_utils.logger import Logger
import inspect as inspect
from pybullet_envs.deep_mimic.env.env import Env
import pybullet_utils.math_util as MathUtil

class ReplayBuffer(object):
    TERMINATE_KEY = 'terminate'
    PATH_START_KEY = 'path_start'
    PATH_END_KEY = 'path_end'

    def __init__(self, buffer_size):
        assert buffer_size > 0

        self.buffer_size = buffer_size
        self.total_count = 0
        self.buffer_head = 0
        self.buffer_tail = MathUtil.INVALID_IDX
        self.num_paths = 0
        self._sample_buffers = dict()
        self.buffers = None

        self.clear()
        return

    def sample(self, n):
        curr_size = self.get_current_size()
        assert curr_size > 0

        idx = np.empty(n, dtype=int)
        # makes sure that the end states are not sampled
        for i in range(n):
            while True:
                curr_idx = np.random.randint(0, curr_size, size=1)[0]
                curr_idx += self.buffer_tail
                curr_idx = np.mod(curr_idx, self.buffer_size)

                if not self.is_path_end(curr_idx):
                    break
            idx[i] = curr_idx

        return idx

    def sample_filtered(self, n, key):
        assert key in self._sample_buffers
        curr_buffer = self._sample_buffers[key]
        idx = curr_buffer.sample(n)
        return idx

    def count_filtered(self, key):
        curr_buffer = self._sample_buffers[key]
        return curr_buffer.count

    def get(self, key, idx):
        return self.buffers[key][idx]

    def get_all(self, key):
        return self.buffers[key]

    def get_idx_filtered(self, key):
        assert key in self._sample_buffers
        curr_buffer = self._sample_buffers[key]
        idx = curr_buffer.slot_to_idx[:curr_buffer.count]
        return idx
    
    def get_path_start(self, idx):
        return self.buffers[self.PATH_START_KEY][idx]

    def get_path_end(self, idx):
        return self.buffers[self.PATH_END_KEY][idx]

    def get_pathlen(self, idx):
        is_array = isinstance(idx, np.ndarray) or isinstance(idx, list)
        if not is_array:
            idx = [idx]

        n = len(idx)
        start_idx = self.get_path_start(idx)
        end_idx = self.get_path_end(idx)
        pathlen = np.empty(n, dtype=int)

        for i in range(n):
            curr_start = start_idx[i]
            curr_end = end_idx[i]
            if curr_start < curr_end:
                curr_len = curr_end - curr_start
            else:
                curr_len = self.buffer_size - curr_start + curr_end
            pathlen[i] = curr_len

        if not is_array:
            pathlen = pathlen[0]

        return pathlen

    def is_valid_path(self, idx):
        start_idx = self.get_path_start(idx)
        valid = start_idx != MathUtil.INVALID_IDX
        return valid

    def store(self, path):
        start_idx = MathUtil.INVALID_IDX
        n = path.pathlength()
        
        if (n > 0):
            assert path.is_valid()

            if path.check_vals():
                if self.buffers is None:
                    self._init_buffers(path)

                idx = self._request_idx(n + 1)
                self._store_path(path, idx)
                self._add_sample_buffers(idx)

                self.num_paths += 1
                self.total_count += n + 1
                start_idx = idx[0]
            else:
                Logger.print2('Invalid path data value detected')
        
        return start_idx

    def clear(self):
        self.buffer_head = 0
        self.buffer_tail = MathUtil.INVALID_IDX
        self.num_paths = 0

        for key in self._sample_buffers:
            self._sample_buffers[key].clear()
        return

    def get_next_idx(self, idx):
        next_idx = np.mod(idx + 1, self.buffer_size)
        return next_idx

    def is_terminal_state(self, idx):
        terminate_flags = self.buffers[self.TERMINATE_KEY][idx]
        terminate = terminate_flags != Env.Terminate.Null.value
        is_end = self.is_path_end(idx)
        terminal_state = np.logical_and(terminate, is_end)
        return terminal_state

    def check_terminal_flag(self, idx, flag):
        terminate_flags = self.buffers[self.TERMINATE_KEY][idx]
        terminate = terminate_flags == flag.value
        return terminate

    def is_path_end(self, idx):
        is_end = self.buffers[self.PATH_END_KEY][idx] == idx
        return is_end

    def add_filter_key(self, key):
        assert self.get_current_size() == 0
        if key not in self._sample_buffers:
            self._sample_buffers[key] = SampleBuffer(self.buffer_size)
        return

    def get_current_size(self):
        if self.buffer_tail == MathUtil.INVALID_IDX:
            return 0
        elif self.buffer_tail < self.buffer_head:
            return self.buffer_head - self.buffer_tail
        else:
            return self.buffer_size - self.buffer_tail + self.buffer_head

    def _check_flags(self, key, flags):
        return (flags & key) == key

    def _add_sample_buffers(self, idx):
        flags = self.buffers['flags']
        for key in self._sample_buffers:
            curr_buffer = self._sample_buffers[key]
            filter_idx = [i for i in idx if (self._check_flags(key, flags[i]) and not self.is_path_end(i))]
            curr_buffer.add(filter_idx)
        return

    def _free_sample_buffers(self, idx):
        for key in self._sample_buffers:
            curr_buffer = self._sample_buffers[key]
            curr_buffer.free(idx)    
        return

    def _init_buffers(self, path):
        self.buffers = dict()
        self.buffers[self.PATH_START_KEY] = MathUtil.INVALID_IDX * np.ones(self.buffer_size, dtype=int);
        self.buffers[self.PATH_END_KEY] = MathUtil.INVALID_IDX * np.ones(self.buffer_size, dtype=int);

        for key in dir(path):
            val = getattr(path, key)
            if not key.startswith('__') and not inspect.ismethod(val):
                if key == self.TERMINATE_KEY:
                    self.buffers[self.TERMINATE_KEY] = np.zeros(shape=[self.buffer_size], dtype=int)
                else:
                    val_type = type(val[0])
                    is_array = val_type == np.ndarray
                    if is_array:
                        shape = [self.buffer_size, val[0].shape[0]]
                        dtype = val[0].dtype
                    else:
                        shape = [self.buffer_size]
                        dtype = val_type
                    
                    self.buffers[key] = np.zeros(shape, dtype=dtype)
        return

    def _request_idx(self, n):
        assert n + 1 < self.buffer_size # bad things can happen if path is too long

        remainder = n
        idx = []

        start_idx = self.buffer_head
        while remainder > 0:
            end_idx = np.minimum(start_idx + remainder, self.buffer_size)
            remainder -= (end_idx - start_idx)

            free_idx = list(range(start_idx, end_idx))
            self._free_idx(free_idx)
            idx += free_idx
            start_idx = 0

        self.buffer_head = (self.buffer_head + n) % self.buffer_size
        return idx

    def _free_idx(self, idx):
        assert(idx[0] <= idx[-1])
        n = len(idx)
        if self.buffer_tail != MathUtil.INVALID_IDX:
            update_tail = idx[0] <= idx[-1] and idx[0] <= self.buffer_tail and idx[-1] >= self.buffer_tail
            update_tail |= idx[0] > idx[-1] and (idx[0] <= self.buffer_tail or idx[-1] >= self.buffer_tail)
            
            if update_tail:
                i = 0
                while i < n:
                    curr_idx = idx[i]
                    if self.is_valid_path(curr_idx):
                        start_idx = self.get_path_start(curr_idx)
                        end_idx = self.get_path_end(curr_idx)
                        pathlen = self.get_pathlen(curr_idx)

                        if start_idx < end_idx:
                            self.buffers[self.PATH_START_KEY][start_idx:end_idx + 1] = MathUtil.INVALID_IDX
                            self._free_sample_buffers(list(range(start_idx, end_idx + 1)))
                        else:
                            self.buffers[self.PATH_START_KEY][start_idx:self.buffer_size] = MathUtil.INVALID_IDX
                            self.buffers[self.PATH_START_KEY][0:end_idx + 1] = MathUtil.INVALID_IDX
                            self._free_sample_buffers(list(range(start_idx, self.buffer_size)))
                            self._free_sample_buffers(list(range(0, end_idx + 1)))
                        
                        self.num_paths -= 1
                        i += pathlen + 1
                        self.buffer_tail = (end_idx + 1) % self.buffer_size;
                    else:
                        i += 1
        else:
            self.buffer_tail = idx[0]
        return

    def _store_path(self, path, idx):
        n = path.pathlength()
        for key, data in self.buffers.items():
            if key != self.PATH_START_KEY and key != self.PATH_END_KEY and key != self.TERMINATE_KEY:
                val = getattr(path, key)
                val_len = len(val)
                assert val_len == n or val_len == n + 1
                data[idx[:val_len]] = val

        self.buffers[self.TERMINATE_KEY][idx] = path.terminate.value
        self.buffers[self.PATH_START_KEY][idx] = idx[0]
        self.buffers[self.PATH_END_KEY][idx] = idx[-1]
        return

class SampleBuffer(object):
    def __init__(self, size):
        self.idx_to_slot = np.empty(shape=[size], dtype=int)
        self.slot_to_idx = np.empty(shape=[size], dtype=int)
        self.count = 0
        self.clear()
        return
    
    def clear(self):
        self.idx_to_slot.fill(MathUtil.INVALID_IDX)
        self.slot_to_idx.fill(MathUtil.INVALID_IDX)
        self.count = 0
        return

    def is_valid(self, idx):
        return self.idx_to_slot[idx] != MathUtil.INVALID_IDX

    def get_size(self):
        return self.idx_to_slot.shape[0]

    def add(self, idx):
        for i in idx:
            if not self.is_valid(i):
                new_slot = self.count
                assert new_slot >= 0

                self.idx_to_slot[i] = new_slot
                self.slot_to_idx[new_slot] = i
                self.count += 1
        return

    def free(self, idx):
        for i in idx:
            if self.is_valid(i):
                slot = self.idx_to_slot[i]
                last_slot = self.count - 1
                last_idx = self.slot_to_idx[last_slot]

                self.idx_to_slot[last_idx] = slot
                self.slot_to_idx[slot] = last_idx
                self.idx_to_slot[i] = MathUtil.INVALID_IDX
                self.slot_to_idx[last_slot] = MathUtil.INVALID_IDX
                self.count -= 1
        return

    def sample(self, n):
        if self.count > 0:
            slots = np.random.randint(0, self.count, size=n)
            idx = self.slot_to_idx[slots]
        else:
            idx = np.empty(shape=[0], dtype=int)
        return idx

    def check_consistency(self):
        valid = True
        if self.count < 0:
            valid = False

        if valid:
            for i in range(self.get_size()):
                if self.is_valid(i):
                    s = self.idx_to_slot[i]
                    if self.slot_to_idx[s] != i:
                        valid = False
                        break

                s2i = self.slot_to_idx[i] 
                if s2i != MathUtil.INVALID_IDX:
                    i2s = self.idx_to_slot[s2i]
                    if i2s != i:
                        valid = False
                        break

        count0 = np.sum(self.idx_to_slot == MathUtil.INVALID_IDX)
        count1 = np.sum(self.slot_to_idx == MathUtil.INVALID_IDX)
        valid &= count0 == count1
        return valid
