import cv2
import numpy as np
from PIL import Image


from .ops import letterbox


# =============================================================================
# Allowed and default steps for preprocessing
# =============================================================================
PREPROCESS_ALLOWED_STEPS = {
    'detect'  : ['list2arr', 'letterbox', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous'],
    'segment' : ['list2arr', 'letterbox', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous'],
    'pose'    : ['list2arr', 'letterbox', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous'],
    'classify': ['list2arr', 'classify', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous']
}

PREPROCESS_DEFAULT_STEPS = {
    'detect'  : ['letterbox', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous'],
    'segment' : ['letterbox', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous'],
    'pose'    : ['letterbox', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous'],
    'classify': ['classify', 'bgr2rgb', 'hwc2chw', 'normalize', 'contiguous']
}


# ===== Operators =====

def letterbox_batch(imgs, input_shape):
    """Apply letterbox resize on batch."""
    return np.stack([letterbox(img, input_shape) for img in imgs])


def list2arr(imgs):
    """Convert list to stack."""
    return np.stack(imgs, axis=0)


def bgr2rgb(arr):
    """Convert BGR to RGB (batch)."""
    return arr[..., ::-1] if arr.shape[-1] == 3 else arr


def hwc2chw(arr):
    """HWC -> CHW for batch."""
    return arr.transpose(0, 3, 1, 2)


def normalize(arr):
    """Normalize [0,255] -> [0,1]."""
    return arr.astype(np.float32) / 255.0


def contiguous(arr):
    """Ensure memory is contiguous."""
    return np.ascontiguousarray(arr)


def letterbox_classify_batch(imgs, transforms):
    """Preprocess images for classification."""
    img_stack = np.stack(
        [transforms(Image.fromarray(im)) for im in imgs]
    )
    return img_stack