from __future__ import annotations
import torch
from torch import nn
IMAGE_LOCAL_AR_CHILD_COUNT = 4

class PackedLinearImageCodec(nn.Module):

    def __init__(self, channels: int, *, norm_eps: float=1e-06, device: torch.device | str | None=None, dtype: torch.dtype | None=None) -> None:
        super().__init__()
        if isinstance(channels, bool) or not isinstance(channels, int) or channels < 1:
            raise ValueError('channels must be a positive integer')
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.channels = channels
        self.down_norm = nn.RMSNorm(channels, eps=norm_eps, elementwise_affine=True, **factory_kwargs)
        self.down_projection = nn.Linear(4 * channels, channels, bias=False, **factory_kwargs)
        self.up_norm = nn.RMSNorm(channels, eps=norm_eps, elementwise_affine=True, **factory_kwargs)
        self.up_projection = nn.Linear(channels, 4 * channels, bias=False, **factory_kwargs)

    @staticmethod
    def _shape(shape: tuple[int, int] | list[int]) -> tuple[int, int]:
        if not isinstance(shape, (tuple, list)) or len(shape) != 2:
            raise ValueError('image token shape must contain height and width')
        (height, width) = shape
        if any((isinstance(value, bool) or not isinstance(value, int) for value in shape)):
            raise TypeError('image token height and width must be integers')
        if height < 1 or width < 1 or height % 2 or width % 2:
            raise ValueError('packed-linear image grids must have positive even dimensions')
        return (height, width)

    @staticmethod
    def _patches(grid: torch.Tensor) -> torch.Tensor:
        return torch.stack((grid[..., 0::2, 0::2, :], grid[..., 0::2, 1::2, :], grid[..., 1::2, 0::2, :], grid[..., 1::2, 1::2, :]), dim=-2)

    @classmethod
    def pack_raster_values(cls, child_values: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
        (height, width) = cls._shape(shape)
        if not isinstance(child_values, torch.Tensor):
            raise TypeError('child_values must be a tensor')
        if child_values.ndim < 2 or child_values.shape[1] != height * width:
            raise ValueError(f'child_values must have shape [B, {height * width}, ...]')
        grid = child_values.reshape(child_values.shape[0], height, width, *child_values.shape[2:])
        patches = torch.stack((grid[:, 0::2, 0::2, ...], grid[:, 0::2, 1::2, ...], grid[:, 1::2, 0::2, ...], grid[:, 1::2, 1::2, ...]), dim=3)
        return patches.flatten(1, 2)

    @classmethod
    def unpack_patch_values(cls, patch_values: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
        (height, width) = cls._shape(shape)
        (parent_height, parent_width) = (height // 2, width // 2)
        parent_count = parent_height * parent_width
        if not isinstance(patch_values, torch.Tensor):
            raise TypeError('patch_values must be a tensor')
        if patch_values.ndim < 3 or patch_values.shape[1] != parent_count or patch_values.shape[2] != IMAGE_LOCAL_AR_CHILD_COUNT:
            raise ValueError(f'patch_values must have shape [B, {parent_count}, {IMAGE_LOCAL_AR_CHILD_COUNT}, ...]')
        patches = patch_values.reshape(patch_values.shape[0], parent_height, parent_width, IMAGE_LOCAL_AR_CHILD_COUNT, *patch_values.shape[3:])
        grid = patch_values.new_empty(patch_values.shape[0], height, width, *patch_values.shape[3:])
        grid[:, 0::2, 0::2, ...] = patches[:, :, :, 0, ...]
        grid[:, 0::2, 1::2, ...] = patches[:, :, :, 1, ...]
        grid[:, 1::2, 0::2, ...] = patches[:, :, :, 2, ...]
        grid[:, 1::2, 1::2, ...] = patches[:, :, :, 3, ...]
        return grid.flatten(1, 2)

    def compress_fixed_grid(self, child_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
        (height, width) = self._shape(shape)
        if child_states.ndim != 3 or child_states.shape[1:] != (height * width, self.channels):
            raise ValueError(f'child_states must have shape [B, {height * width}, {self.channels}]')
        grid = child_states.reshape(child_states.shape[0], height, width, self.channels)
        patches = self._patches(grid)
        normalized = self.down_norm(patches)
        parents = self.down_projection(normalized.reshape(child_states.shape[0], height // 2, width // 2, 4 * self.channels))
        return parents.flatten(1, 2)

    def expand_fixed_grid(self, parent_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
        children = self.expand_fixed_grid_patches(parent_states, shape)
        return self.unpack_patch_values(children, shape)

    def expand_fixed_grid_patches(self, parent_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
        (height, width) = self._shape(shape)
        parent_count = height // 2 * (width // 2)
        if parent_states.ndim != 3 or parent_states.shape[1:] != (parent_count, self.channels):
            raise ValueError(f'parent_states must have shape [B, {parent_count}, {self.channels}]')
        return self.up_projection(self.up_norm(parent_states)).reshape(parent_states.shape[0], parent_count, 4, self.channels)