import struct
import pickle


class RamPickleWrite(object):
    def __init__(self, sequence):
        self.sequence = sequence

    def pickle(self, fd):
        base_offset = fd.tell()
        index = []
        fd.write(struct.pack('L', 0))
        i = 0
        while True:
            batch = self.sequence[i:i + 100]
            if not batch:
                break
            index.append(fd.tell())
            pickle.dump(batch, fd)
            i += 100
        index_offset = fd.tell()
        pickle.dump(index, fd)
        fd.seek(base_offset)
        fd.write(struct.pack('L', index_offset))


class RamPickleRead(object):
    def __init__(self, fd):
        self.fd = fd
        buf = fd.read(struct.calcsize('L'))
        index_offset, = struct.unpack('L', buf)
        fd.seek(index_offset)
        self.index = pickle.load(fd)
        self.batches = {}

    def load_batch(self, index):
        page = index // 100
        if len(self.index) <= page:
            return
        if page not in self.batches:
            self.fd.seek(self.index[page])
            self.batches[page] = pickle.load(self.fd)
        return self.batches[page]

    def __getitem__(self, index):
        if isinstance(index, (long, int)):
            batch = self.load_batch(index)
            if not batch:
                raise IndexError(index)
            offset = index % 100
            if len(batch) <= offset:
                raise IndexError(index)
            return batch[offset]
        elif isinstance(index, slice):
            l = []
            i = index.start or 0
            while index.stop is None or i < index.stop:
                try:
                    l.append(self[i])
                except IndexError:
                    break
                i += index.step or 1
            return l
        else:
            raise TypeError(index)
