from __future__ import annotations
import numpy as np
import torch

def resolve_source_padding_vq_lock(model_config, override=None, *, has_edit_image):
    if not isinstance(has_edit_image, bool):
        raise TypeError('has_edit_image must be boolean')
    if override is None:
        resolved = getattr(model_config, 'lock_source_padding_vq', False)
        if resolved is None:
            resolved = False
        if not isinstance(resolved, bool):
            raise TypeError('config.lock_source_padding_vq must be boolean or null')
        if resolved and (not has_edit_image):
            return False
        return resolved
    if not isinstance(override, bool):
        raise TypeError('lock_source_padding_vq override must be boolean or null')
    if override and (not has_edit_image):
        raise ValueError('lock_source_padding_vq requires an edit_image')
    return override

def _resolve_source_native_size(value, *, fallback):
    resolved = fallback if value is None else value
    if not isinstance(resolved, (tuple, list)) or len(resolved) != 2 or any((isinstance(item, bool) or not isinstance(item, (int, np.integer)) or int(item) < 1 for item in resolved)):
        raise ValueError('edit_source_native_size must contain positive integer width and height')
    return (int(resolved[0]), int(resolved[1]))

def _source_padding_vq_mask(source_native_size, raw_grid_shape, *, device):
    (width, height) = _resolve_source_native_size(source_native_size, fallback=source_native_size)
    if not isinstance(raw_grid_shape, (tuple, list)) or len(raw_grid_shape) != 2 or any((isinstance(item, bool) or not isinstance(item, (int, np.integer)) or int(item) < 1 for item in raw_grid_shape)):
        raise ValueError('raw_grid_shape must contain positive integer height and width')
    (grid_height, grid_width) = (int(value) for value in raw_grid_shape)
    if grid_height != grid_width:
        raise ValueError('source-padding VQ locking requires a square raw VQ grid')
    side = max(width, height)
    left = round((side - width) * 0.5)
    top = round((side - height) * 0.5)
    right = left + width
    bottom = top + height
    columns = torch.arange(grid_width, dtype=torch.long, device=device)
    rows = torch.arange(grid_height, dtype=torch.long, device=device)
    horizontally_outside = ((columns + 1) * side <= left * grid_width) | (columns * side >= right * grid_width)
    vertically_outside = ((rows + 1) * side <= top * grid_height) | (rows * side >= bottom * grid_height)
    return (vertically_outside[:, None] | horizontally_outside[None, :]).reshape(1, -1)

def _augment_source_padding_vq_mask(state, raw_padding_mask, content_positions):
    if not torch.is_tensor(state) or state.ndim != 2:
        raise ValueError('generation state must have shape [batch, tokens]')
    if not torch.is_tensor(raw_padding_mask) or raw_padding_mask.dtype is not torch.bool or raw_padding_mask.ndim != 2 or (raw_padding_mask.shape[0] != state.shape[0]) or (raw_padding_mask.device != state.device):
        raise TypeError('raw_padding_mask must be device-local bool with the state batch size')
    if not torch.is_tensor(content_positions) or content_positions.ndim != 1 or content_positions.dtype is not torch.long or (content_positions.device != state.device) or (content_positions.numel() != raw_padding_mask.shape[1]):
        raise TypeError('content_positions must be a device-local long vector covering raw cells')
    if bool((content_positions < 0).any()) or bool((content_positions >= state.shape[1]).any()):
        raise ValueError('content_positions contains an out-of-range state index')
    locked_mask = torch.zeros_like(state, dtype=torch.bool)
    locked_mask[:, content_positions] = raw_padding_mask
    return locked_mask

def _prefill_source_padding_vq(state, source_codes, raw_padding_mask, content_positions):
    if not torch.is_tensor(state) or state.ndim != 2:
        raise ValueError('generation state must have shape [batch, tokens]')
    if not torch.is_tensor(source_codes) or source_codes.ndim != 2 or source_codes.dtype != state.dtype or (source_codes.device != state.device):
        raise TypeError('source_codes must match the state integer dtype and device')
    if source_codes.shape != raw_padding_mask.shape:
        raise TypeError('raw_padding_mask must be bool and match source_codes')
    locked_mask = _augment_source_padding_vq_mask(state, raw_padding_mask, content_positions)
    locked_codes = torch.zeros_like(state)
    locked_codes[:, content_positions] = source_codes
    prefilled = state.clone()
    prefilled[locked_mask] = locked_codes[locked_mask]
    return (prefilled, locked_mask, locked_codes)

def _restore_source_padding_vq(state, locked_mask, locked_codes):
    if locked_mask is None:
        return state
    if not torch.is_tensor(state) or state.ndim != 2 or (not torch.is_tensor(locked_mask)) or (locked_mask.dtype is not torch.bool) or (locked_mask.shape != state.shape) or (locked_mask.device != state.device) or (not torch.is_tensor(locked_codes)) or (locked_codes.shape != state.shape) or (locked_codes.dtype != state.dtype) or (locked_codes.device != state.device):
        raise TypeError('locked source-padding tensors must align with sampler state')
    restored = state.clone()
    restored[locked_mask] = locked_codes[locked_mask]
    return restored

def _assert_source_padding_vq_locked(state, locked_mask, locked_codes):
    if locked_mask is None:
        return
    if not torch.is_tensor(locked_mask) or locked_mask.dtype is not torch.bool or locked_mask.shape != state.shape or (not torch.is_tensor(locked_codes)) or (locked_codes.shape != state.shape):
        raise TypeError('locked source-padding tensors must align with sampler state')
    if not torch.equal(state[locked_mask], locked_codes[locked_mask]):
        raise RuntimeError('a locked source-padding VQ code changed during sampling')