# 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 parse_crop_center(orig, shape):
    h = (orig[1] - shape[0]) / 2
    w = (orig[0] - shape[1]) / 2 
    l, r = int(np.floor(w)), int(np.ceil(w))
    u, d = int(np.floor(h)), int(np.ceil(h))
    return (l, u, r, d)


def parse_crop_random(orig, shape):
    l = int(random.random() * (orig[0] - shape[1]))
    u = int(random.random() * (orig[1] - shape[0]))
    r = orig[0] - (l + shape[1])
    d = orig[1] - (u + shape[0])
    return (l, u, r, d)


def parse_crop_stride(orig, shape):
    h = apply_stride(orig[1], shape[0], capped=True)
    w = apply_stride(orig[0], shape[0], capped=True)
    lr, ud = orig[0] - w, orig[1] - h
    l, r = int(np.floor(lr / 2)), int(np.ceil(lr / 2))
    u, d = int(np.floor(ud / 2)), int(np.ceil(ud / 2))
    return (l, u, r, d)


def parse_crop_params(orig, shape, mode):
    if 'random' in mode:
        borders = {key[1]: parse_crop_random(val, shape) 
            for key, val in orig.items() if key[0] == 0}
        borders = {key: borders[key[1]] for key in orig.keys()}
    elif 'center' in mode:
        borders = {key[1]: parse_crop_center(val, shape) 
            for key, val in orig.items() if key[0] == 0}
        borders = {key: borders[key[1]] for key in orig.keys()}
    elif 'stride' in mode:
        borders = {key[1]: parse_crop_stride(val, shape)
            for key, val in orig.items() if key[0] == 0}
        borders = {key: borders[key[1]] for key in orig.keys()}
    else:
        raise ValueError(f'Invalid cropping function {mode}')
    return borders


@iterate12
def crop_mask(mask, borders):
    if isinstance(mask, np.ndarray):
        return crop_npy(mask, borders)
    elif isinstance(mask, Image.Image):
        return crop_pil(mask, borders)
    else:
        raise ValueError(f'Invalid mask type {type(mask)}')


@iterate12
def crop_pil(image, borders):
    """Crops PIL image"""
    return image.crop((
        borders[0], 
        borders[1], 
        image.size[0] - borders[2], 
        image.size[1] - borders[3],
    ))


@iterate12
def crop_intrinsics(intrinsics, borders):
    """Crop camera intrinsics matrix to match a target resolution"""
    intrinsics = np.copy(intrinsics)

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

    intrinsics[0, 2] -= borders[0]
    intrinsics[1, 2] -= borders[1]

    return intrinsics


@iterate12
def crop_npy(data, borders):
    """Crops numpy array"""
    u, d = borders[1], None if borders[3] == 0 else - borders[3]
    l, r = borders[0], None if borders[2] == 0 else - borders[2]
    data = data[:, u:d, l:r]
    return np.ascontiguousarray(data)


@iterate12
def crop_bbox2d(bbox2d, borders):
    """Crop 2D bounding boxes to match a target resolution"""
    bbox2d = np.copy(bbox2d)
    bbox2d[:, 0] -= borders[0]
    bbox2d[:, 1] -= borders[1]
    bbox2d[:, 2] -= borders[0]
    bbox2d[:, 3] -= borders[1]
    return bbox2d


def crop_sample(sample, shape, mode):
    """Crops 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()}
    borders = parse_crop_params(orig, shape, mode)

    ### INTRINSICS
    for key in keys_with(sample, 'intrinsics'):
        sample[key] = crop_intrinsics(sample[key], borders)
    ### BBOX2D
    for key in keys_with(sample, 'bbox2d'):
        sample[key] = crop_bbox2d(sample[key], borders)
    ### RGB
    for key in keys_with(sample, 'rgb'):
        if not key.startswith('lat'):
            sample[key] = crop_pil(sample[key], borders)
    ### DEPTH
    for key in keys_with(sample, 'depth'):
        if 'mask' in key: continue
        sample[key] = crop_npy(sample[key], borders)
    ### NORMALS
    for key in keys_with(sample, 'normals'):
        sample[key] = crop_npy(sample[key], borders)
    ### MASKS
    for key in keys_with(sample, 'mask'):
        sample[key] = crop_mask(sample[key], borders)
    ### SEMANTIC
    for key in keys_with(sample, 'semantic'):
        sample[key] = crop_npy(sample[key], borders)
    ### OPTICAL FLOW
    for key in keys_with(sample, 'optflow'):
        sample[key] = crop_npy(sample[key], borders)
    ### SCENE FLOW
    for key in keys_with(sample, 'scnflow'):
        sample[key] = crop_npy(sample[key], borders)

    # Switch from borders to resolution
    resolution = {key: [
        orig[key][1] - (val[1] + val[3]), 
        orig[key][0] - (val[0] + val[2]),
    ] for key, val in borders.items()}

    # Return resized sample
    return sample, resolution