__all__ = [
    're_instantiate_sampler'
]
from array import array
from typing import Sequence
from collections import deque


def re_instantiate_sampler(sampler, new_sampler_class=None):
    all_attributes = vars(sampler)
    if new_sampler_class is not None:
        return new_sampler_class(**all_attributes)
    return type(sampler)(**all_attributes)


def create_array(length, fill_value) -> array:
    """
    根据长度自动创建 array ，超过 4294967295 需要使用 'L', 否则使用 'I'

    :param length:
    :param fill_value:
    :return:
    """
    if not isinstance(fill_value, Sequence):
        fill_value = [fill_value]*length

    if length > 4294967295:
        _index_lst = array("L", fill_value)
    else:
        _index_lst = array("I", fill_value)
    return _index_lst


class NumConsumedSamplesArray:
    def __init__(self, buffer_size=2000, num_consumed_samples=0):
        """
        保留 buffer_size 个 num_consumed_samples 数据，可以索引得到某个 index 下的 num_consumed_samples 多少

        Example::

            array = NumConsumedSamplesArray(buffer_size=3)
            for i in range(10):
                array.push(i)

            array[9]  # 输出为9，表示这个位置真实的 num_consumed_samples 是多少。
            array[6]  # 报错，因为只保留了3个最近的数据，6超过了最大buffer的记录了，即 [7, 8, 9]

        暂时由于 sampler 的 batch 都是规整的，先保留

        :param buffer_size: 报错多少个历史。
        :param num_consumed_samples: 第一个 num_consumed_samples 是多少。
        """
        self.count = 0
        self.deque = deque(maxlen=buffer_size)
        if num_consumed_samples is not None:
            self.push(num_consumed_samples)
        self.buffer_size = buffer_size

    def __getitem__(self, item):
        if len(self.deque) == 0:  # 如果没有任何缓存的内容，说明还没有写入，直接返回0
            return 0
        assert isinstance(item, int), "Only int index allowed."
        assert self.count-len(self.deque)<=item<self.count, f"Only keep {len(self.deque)} history index."
        index = len(self.deque) - (self.count - item)
        return self.deque[index]

    def push(self, num_consumed_samples):
        self.deque.append(num_consumed_samples)
        self.count += 1