# Copyright 2026 Toyota Research Institute.  All rights reserved.

import cv2
import numpy as np
from functools import partial
from PIL import Image
import random

import torch
import torchvision.transforms as transforms
from torchvision.transforms import InterpolationMode

from anydata.utils.data import keys_with, first_value
from anydata.utils.decorators import iterate1, iterate12, iterate123
from anydata.utils.types import is_seq, is_int, is_list, is_numpy, is_pil, is_seq, is_dict
from anydata.augmentations.misc import apply_stride 


def get_mask_fn(data):
    """Retrun function used to resize mask (numpy or pil)"""
    if is_dict(data):
        return get_mask_fn(first_value(data))
    elif is_seq(data):
        return get_mask_fn(data[0])
    else:
        return resize_npy if is_numpy(data) else \
            partial(resize_pil, interpolation=InterpolationMode.NEAREST)


def parse_resize_params(orig, shape):
    """Parse resize parameters to generate new width and height"""
    # If it's a dict, parse each dimension separately
    if is_dict(orig):
        shape = {key[1]: parse_resize_params(val, shape)
            for key, val in orig.items() if key[0] == 0}
        return {key: shape[key[1]] for key in orig.keys()}

    # Get original resize ratios
    ratio_hw = orig[1] / orig[0]
    ratio_wh = orig[0] / orig[1]

    # If each dimension is a list, randomly select a value within that interval
    if is_seq(shape[0]):
        height = 0 if len(shape[0]) == 0 else shape[0][0] if len(shape[0]) == 1 else \
            shape[0][0] + random.random() * (shape[0][1] - shape[0][0])
        width = 0 if len(shape[1]) == 0 else shape[1][0] if len(shape[1]) == 1 else \
            shape[1][0] + random.random() * (shape[1][1] - shape[1][0])
        new_shape = [height, width]
    elif 0 < shape[1] < 10: # Assume a value 0 < x < 10 is a ratio, not a direct value
        new_shape = [o * shape[1] for o in orig]
        if shape[0] < 0: # Apply stride if requested
            stride = - shape[0]
            new_shape[0] = apply_stride(new_shape[0], stride, capped=True)
            new_shape[1] = apply_stride(new_shape[1], stride, capped=True)
        new_shape = new_shape[::-1]
    # Nothing fancy, just use those values
    else:
        new_shape = shape[:2]

    # If one dimension is zero, preserve ratio
    if new_shape[0] == 0 and new_shape[1] > 0:
        new_shape[0] = np.round(orig[1] / orig[0] * new_shape[1])
    # If another dimension is zero, preserve ratio
    elif new_shape[1] == 0 and new_shape[0] > 0:
        new_shape[1] = np.round(orig[0] / orig[1] * new_shape[0])
    # If first dimension is negative, use valid dimension as the largest dimension
    elif new_shape[0] < 0:
        stride = - new_shape[0]
        target = new_shape[1]
        # Check with dimension is larger
        if orig[0] > orig[1]:
            new_shape[1] = target
            new_shape[0] = np.round(orig[1] / orig[0] * new_shape[1])
        else:
            new_shape[0] = target
            new_shape[1] = np.round(orig[0] / orig[1] * new_shape[0])
        # Apply stride resizing on both
        new_shape[0] = apply_stride(new_shape[0], stride, capped=True)
        new_shape[1] = apply_stride(new_shape[1], stride, capped=True)
    elif new_shape[0] == 0 and new_shape[1] == 0:
        new_shape = list(orig[::-1])

    # Additional resizing options (third value in the list)
    if len(shape) > 2:
        for s in shape[2:]:
            # Shorthand for stride (sX)
            if isinstance(s, str):
                if s.startswith('s'):
                    s = ['stride',int(s[1:])]
            # There is a minimum dimension value
            if s[0] == 'min':
                min_dim_h = s[1][0] if is_list(s[1]) else s[1]
                min_dim_w = s[1][1] if is_list(s[1]) else s[1]
                if new_shape[0] < min_dim_h:
                    new_shape[0], new_shape[1] = min_dim_h, int(orig[0] / orig[1] * min_dim_h)
                if new_shape[1] < min_dim_w:
                    new_shape[1], new_shape[0] = min_dim_w, int(orig[1] / orig[0] * min_dim_w)
            # There is a maximum dimension value
            elif s[0] == 'max':
                max_dim_h = s[1][0] if is_list(s[1]) else s[1]
                max_dim_w = s[1][1] if is_list(s[1]) else s[1]
                if new_shape[0] >= max_dim_h:
                    new_shape[0], new_shape[1] = max_dim_h, int(orig[0] / orig[1] * max_dim_h)
                if new_shape[1] >= max_dim_w:
                    new_shape[1], new_shape[0] = max_dim_w, int(orig[1] / orig[0] * max_dim_w)
            # Needs to respect stride value
            elif s[0] == 'stride':
                stride = s[1]
                new_shape[0] = apply_stride(new_shape[0], stride, capped=True)
                new_shape[1] = apply_stride(new_shape[1], stride, capped=True)
            # Needs to fit in a certain dimensions
            elif s[0] == 'fit':
                if new_shape[1] > s[2]:
                    new_shape = list([ratio_hw * s[2], s[2]])
                if new_shape[0] > s[1]:
                    new_shape = list([s[1], ratio_wh * s[1]])
                if new_shape[1] > orig[0]:
                    new_shape[1] = ratio_wh * new_shape[0]
                if new_shape[0] > orig[1]:
                    new_shape[0] = ratio_hw * new_shape[1]
    # Return rounded values
    return tuple([int(np.round(s)) for s in new_shape])


@iterate12
def resize_pil(image, shape, interpolation=InterpolationMode.LANCZOS):
    """Resizes PIL image"""
    # If a single number is provided, use resize ratio
    if not is_seq(shape):
        shape = tuple(int(s * shape) for s in image.size[::-1])
    transform = transforms.Resize(shape, interpolation=interpolation)
    return transform(image)


@iterate12
def resize_npy(depth, shape):
    """Resizes numpy array"""
    # If a single number is provided, use resize ratio
    if not is_seq(shape):
        shape = tuple(int(s * shape) for s in depth.shape[:2])
    # Resize depth map
    depth = np.transpose(depth, (1, 2, 0))
    depth = cv2.resize(depth, dsize=tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST)
    # Return resized depth map
    if len(depth.shape) == 2:
        depth = np.expand_dims(depth, axis=0)
    return depth


@iterate12
def resize_npy_preserve(depth, shape):
    """Resizes depth map preserving all valid depth pixels (numpy)"""
    # If a single number is provided, use resize ratio
    if not is_seq(shape):
        shape = tuple(int(s * shape) for s in depth.shape[:2])
    # Store dimensions and reshapes to single column
    depth = np.squeeze(depth)
    h, w = depth.shape
    x = depth.reshape(-1)
    # Create coordinate grid
    uv = np.mgrid[:h, :w].transpose(1, 2, 0).reshape(-1, 2)
    # Filters valid points
    idx = x > 0
    crd, val = uv[idx], x[idx]
    # Downsamples coordinates
    crd[:, 0] = (crd[:, 0] * (shape[0] / h)).astype(np.int32)
    crd[:, 1] = (crd[:, 1] * (shape[1] / w)).astype(np.int32)
    # Filters points inside image
    idx = (crd[:, 0] < shape[0]) & (crd[:, 1] < shape[1])
    crd, val = crd[idx], val[idx]
    # Creates downsampled depth image and assigns points
    depth = np.zeros(shape)
    depth[crd[:, 0], crd[:, 1]] = val
    # Return resized depth map
    if len(depth.shape) == 2:
        depth = np.expand_dims(depth, axis=0)
    return depth


@iterate12
def resize_torch_preserve(depth, shape):
    """Resizes depth map preserving all valid depth pixels (torch)"""
    # If a single number is provided, use resize ratio
    if not is_seq(shape):
        shape = tuple(int(s * shape) for s in depth.shape)
    # Store dimensions and reshapes to single column
    c, h, w = depth.shape
    # depth = np.squeeze(depth)
    # h, w = depth.shape
    x = depth.reshape(-1)
    # Create coordinate grid
    uv = np.mgrid[:h, :w].transpose(1, 2, 0).reshape(-1, 2)
    # Filters valid points
    idx = x > 0
    crd, val = uv[idx], x[idx]
    # Downsamples coordinates
    crd[:, 0] = (crd[:, 0] * (shape[0] / h)).astype(np.int32)
    crd[:, 1] = (crd[:, 1] * (shape[1] / w)).astype(np.int32)
    # Filters points inside image
    idx = (crd[:, 0] < shape[0]) & (crd[:, 1] < shape[1])
    crd, val = crd[idx], val[idx]
    # Creates downsampled depth image and assigns points
    depth = torch.zeros(shape, device=depth.device, dtype=depth.dtype)
    depth[crd[:, 0], crd[:, 1]] = val
    # Return resized depth map
    return depth.unsqueeze(0)


@iterate12
@iterate1
def resize_npy_multiply(data, shape):
    """Resizes flow map, where we need to multiply by the ratio"""
    if data is None:
        return data
    ratio_w = shape[0] / data.shape[0]
    ratio_h = shape[1] / data.shape[1]
    out = resize_npy(data, shape)
    out[..., 0] *= ratio_h
    out[..., 1] *= ratio_w
    return out


@iterate123
def resize_intrinsics(intrinsics, original, resized):
    """Resize camera intrinsics matrix to match a target resolution"""
    resized = resized[::-1] if is_seq(resized) else resized

    # If a single number is provided, use resize ratio
    if not is_seq(resized):
        resized = tuple(int(o * resized) for o in original)

    intrinsics = np.copy(intrinsics)

    ratio_w = resized[0] / original[0]
    ratio_h = resized[1] / original[1]

    if len(intrinsics.shape) == 1:
        if len(intrinsics) in [4, 8, 9, 16]:  # flat pinhole or distorted model
            intrinsics[0] *= ratio_w
            intrinsics[1] *= ratio_h
            intrinsics[2] *= ratio_w
            intrinsics[3] *= ratio_h
        elif len(intrinsics) == 14:  # ftheta model
            # [sx, sy, cx, cy, fw0, fw1, fw2, fw3, fw4, bw0, bw1, bw2, bw3, bw4]
            intrinsics[0] *= ratio_w
            intrinsics[1] *= ratio_h
        else:
            print('### CAREFUL, intrinsics are not being resized!!')
        return intrinsics

    intrinsics[0, 0] *= ratio_w
    intrinsics[1, 1] *= ratio_h

    intrinsics[0, 2] = intrinsics[0, 2] * ratio_w
    intrinsics[1, 2] = intrinsics[1, 2] * ratio_h
    #     intrinsics[0, 2] = (intrinsics[0, 2] - 0.5) * ratio_w + 0.5
    #     intrinsics[1, 2] = (intrinsics[1, 2] - 0.5) * ratio_h + 0.5

    return intrinsics


@iterate123
def resize_bbox2d(bbox2d, original, resized):
    """Resize 2D bounding boxes to match a target resolution"""
    if bbox2d.shape[0] == 0:
        return bbox2d  # Do nothing if there are no bboxes

    resized = resized[::-1] if is_seq(resized) else resized

    # If a single number is provided, use resize ratio
    if not is_seq(resized):
        resized = tuple(int(o * resized) for o in original)

    bbox2d = np.copy(bbox2d)

    ratio_w = resized[0] / original[0]
    ratio_h = resized[1] / original[1]

    bbox2d[:, 0] = bbox2d[:, 0] * ratio_w
    bbox2d[:, 1] = bbox2d[:, 1] * ratio_h
    bbox2d[:, 2] = bbox2d[:, 2] * ratio_w
    bbox2d[:, 3] = bbox2d[:, 3] * ratio_h

    return bbox2d


def resize_sample(sample, shape, preserve_depth=False, pil_interpolation=InterpolationMode.LANCZOS):
    """Resizes a sample, including all relevant labels."""
    # Cast RGB back to PIL if needed
    for key, val in sample['rgb'].items():
        if isinstance(val, np.ndarray):
            sample['rgb'][key] = Image.fromarray(np.uint8(val)).convert('RGB')

    # Get shapes
    orig = {key: val.size for key, val in sample['rgb'].items()}
    shape = parse_resize_params(orig, shape)

    ### INTRINSICS
    for key in keys_with(sample, 'intrinsics'):
        original = {k: v.size for k, v in sample['rgb'].items()}
        sample[key] = resize_intrinsics(sample[key], original, shape)
    ### BBOX2D
    for key in keys_with(sample, 'bbox2d'):
        original = {k: v.size for k, v in sample['rgb'].items()}
        sample[key] = resize_bbox2d(sample[key], original, shape)
    ### RGB
    for key in keys_with(sample, 'rgb'):
        if not key.startswith('lat'):
            sample[key] = resize_pil(sample[key], shape, interpolation=pil_interpolation)
    ### DEPTH
    for key in keys_with(sample, 'depth'):
        if 'mask' in key:
            continue
        resize_npy_depth = resize_npy_preserve if preserve_depth else resize_npy
        sample[key] = resize_npy_depth(sample[key], shape)
    ### NORMALS
    for key in keys_with(sample, 'normals'):
        sample[key] = resize_npy(sample[key], shape)
    ### MASKS
    for key in keys_with(sample, 'mask'):
        resize_fn = get_mask_fn(sample[key])
        sample[key] = resize_fn(sample[key], shape)
    ### SEMANTIC
    for key in keys_with(sample, 'semantic'):
        sample[key] = resize_npy(sample[key], shape)
    ### OPTICAL FLOW
    for key in keys_with(sample, 'optflow'):
        sample[key] = resize_npy_multiply(sample[key], shape)
    ### SCENE FLOW
    for key in keys_with(sample, 'scnflow'):
        sample[key] = resize_npy(sample[key], shape)

    # Return resized sample
    return sample, shape