import cv2
import numpy as np

def identity_transform(self, imgs: np.ndarray) -> np.ndarray:
    """Identity transform - returns input as is."""
    return imgs


def yolo_transform(self, imgs: np.ndarray) -> np.ndarray:
    """YOLO transform - normalize to [0, 1] range."""
    return imgs / 255.0


def v5_pose_transform(self, imgs: np.ndarray) -> np.ndarray:
    """
    YOLOv5 pose transform - reshape and normalize.
    Input: (N, 320, 320, 3)
    Output: (N, 160, 160, 12)
    """
    # 確保輸入是 4D (N, H, W, C)
    if imgs.ndim == 3:
        imgs = np.expand_dims(imgs, 0)
    
    # Reshape: Nx320x320x3 -> (N, 160, 2, 160, 2, 3)
    h, w = imgs.shape[1:3]
    imgs = imgs.reshape(-1, h//2, 2, w//2, 2, 3)
    
    # Transpose & Reshape to (N, H/2, W/2, 12)
    imgs = np.transpose(imgs, (0, 1, 3, 2, 4, 5)) 
    imgs = imgs.reshape(-1, h//2, w//2, 12)
    imgs = imgs.astype(np.float32) / 255.0
    
    return imgs


def ir_transform(self, imgs: np.ndarray) -> np.ndarray:
    """
    IR (face recognition) transform.
    Resize to 112x112, convert BGR to RGB, normalize to [-1, 1].
    """
    if imgs.ndim == 3:
        imgs = np.expand_dims(imgs, 0)

    imgs_out = np.empty((imgs.shape[0], 112, 112, 3), dtype=np.float32)
    
    for idx in range(imgs.shape[0]):
        img = imgs[idx]
        
        # Resize
        img = cv2.resize(img, (112, 112), interpolation=cv2.INTER_AREA)
        
        # BGR -> RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        imgs_out[idx] = img.astype(np.float32)
    
    # Normalize to [-1, 1]
    imgs_out /= 255.0
    imgs_out = (imgs_out - 0.5) / 0.5
    
    return imgs_out


def reid_transform(self, imgs: np.ndarray) -> np.ndarray:
    """
    ReID transform - ImageNet normalization.
    Input: (N, H, W, 3), value range [0, 255]
    Output: Normalized with ImageNet mean/std
    """
    if imgs.ndim == 3:
        imgs = np.expand_dims(imgs, 0)
        
    imgs = imgs.astype(np.float32) / 255.0
    
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    
    imgs = (imgs - mean) / std  # Broadcasting
    return imgs


def par_transform(self, imgs: np.ndarray) -> np.ndarray:
    """
    PAR (person attribute recognition) transform.
    Convert BGR to RGB, normalize with CLIP mean/std, convert to CHW format.
    """
    if imgs.ndim == 3:
        imgs = np.expand_dims(imgs, 0)

    clip_mean = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
    clip_std = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
    
    imgs_out = np.empty(imgs.shape, dtype=np.float32)
    
    for idx in range(imgs.shape[0]):
        img = imgs[idx]
        # BGR -> RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        imgs_out[idx] = img.astype(np.float32) / 255.0
    
    # Normalize
    imgs_out = (imgs_out - clip_mean) / clip_std
    
    return imgs_out.astype(np.float32)