import numpy as np

from .flow import CHUNK


class SharedFrames:
    """Two input slots shared by extraction and generation."""

    slots = 2

    def __init__(self, context, size):
        width, height = size
        self.pose_shape = (CHUNK, height, width, 3)
        self.face_shape = (CHUNK, 512, 512, 3)
        self.pose_bytes = int(np.prod(self.pose_shape))
        self.slot_bytes = self.pose_bytes + int(np.prod(self.face_shape))
        self.storage = context.RawArray("B", self.slots * self.slot_bytes)

    def read(self, slot, count):
        if not 0 <= slot < self.slots or count not in (9, CHUNK):
            raise ValueError("Invalid shared input slot or frame count.")
        offset = slot * self.slot_bytes
        pose = np.ndarray(self.pose_shape, np.uint8, self.storage, offset=offset)
        face = np.ndarray(self.face_shape, np.uint8, self.storage, offset=offset + self.pose_bytes)
        return pose[:count], face[:count]

    def write(self, slot, pose, face):
        targets = self.read(slot, len(pose))
        for target, value in zip(targets, (pose, face)):
            if value.shape != target.shape or value.dtype != np.uint8:
                raise ValueError("Invalid extracted frame shape or dtype.")
        for target, value in zip(targets, (pose, face)):
            np.copyto(target, value)