# Copyright 2026 Toyota Research Institute.  All rights reserved.

import io

import numpy as np
import torch
from PIL import Image


def check_invalid_resolution(res, target):
    if target is None:
        return False
    return list(res.shape[-2:]) != target


def check_invalid_matrix(value):
    try:
        arr = np.asarray(value)
    except Exception:
        return False
    if arr.dtype == object:
        return False
    if not np.issubdtype(arr.dtype, np.number):
        return False
    return np.isnan(arr).any() or np.isinf(arr).any() or np.sum(arr) == 0


def rgb_to_byte_array(rgb, format="jpeg", quality=95):
    if isinstance(rgb, np.ndarray):
        rgb = Image.fromarray(np.uint8(rgb)).convert('RGB')        
    rgb_bytes = io.BytesIO()
    rgb.save(rgb_bytes, format=format, quality=quality)
    return rgb_bytes.getvalue()


def _mask_to_pil(mask):
    if isinstance(mask, Image.Image):
        return mask.convert("L")
    if isinstance(mask, torch.Tensor):
        mask = mask.detach().cpu().numpy()
    if not isinstance(mask, np.ndarray):
        mask = np.asarray(mask)
    if mask.ndim == 3 and mask.shape[0] == 1:
        mask = mask[0]
    if mask.dtype == np.bool_:
        mask = mask.astype(np.uint8) * 255
    elif np.issubdtype(mask.dtype, np.floating):
        max_val = float(np.max(mask)) if mask.size else 0.0
        if max_val <= 1.0:
            mask = np.clip(mask * 255.0, 0, 255).astype(np.uint8)
        else:
            mask = np.clip(mask, 0, 255).astype(np.uint8)
    else:
        mask = np.clip(mask, 0, 255).astype(np.uint8)
    return Image.fromarray(mask, mode="L")


def mask_to_byte_array(mask):
    mask_bytes = io.BytesIO()
    _mask_to_pil(mask).save(mask_bytes, format="PNG")
    return mask_bytes.getvalue()

## TODO: This is a temporary workaround for loading VAE and text encoder weights from S3. We should eventually replace this with a more robust solution that can handle correct inputs.
def _ensure_latent_tensor(val, label):
    if isinstance(val, torch.Tensor):
        tensor = val
    elif isinstance(val, Image.Image):
        arr = np.asarray(val)
        if arr.ndim == 2:
            arr = arr[None, :, :]
        else:
            arr = arr.transpose(2, 0, 1)
        tensor = torch.from_numpy(arr)
    else:
        arr = np.asarray(val)
        if arr.ndim == 2:
            arr = arr[None, :, :]
        elif arr.ndim == 3 and arr.shape[-1] in [1, 3]:
            arr = arr.transpose(2, 0, 1)
        tensor = torch.from_numpy(arr)

    if tensor.ndim != 3:
        raise ValueError(f"Expected 3D tensor-like input for {label}, got shape {tuple(tensor.shape)}")

    tensor = tensor.float()
    if label == 'rgb':
        if tensor.shape[0] == 1:
            tensor = tensor.repeat(3, 1, 1)
        if tensor.max() > 1.0:
            tensor = tensor / 255.0
    elif label == 'depth':
        if tensor.shape[0] > 1:
            tensor = tensor[:1]
        if tensor.max() > 1.0:
            tensor = tensor / 255.0
    return tensor


def encode_latents(data, time_st, time_fn, cams, label, vae, b, stride=1, device="cuda:0"):
    if vae is None:
        return {}

    labs = []
    used_cams = []
    for cam in cams:
        lab = data[label]
        lab = {
            key: _ensure_latent_tensor(val, label).to(device=device, dtype=torch.bfloat16, non_blocking=True)
            for key, val in lab.items()
            if time_st <= key[0] < time_fn and key[1] == cam and ((key[0] - time_st) % stride == 0)
        }
        lab_vals = [val for val in lab.values()]
        if len(lab_vals) == 0:
            continue
        while len(lab_vals) < b + 1:
            lab_vals.append(lab_vals[-1].clone())
        lab_clip = torch.stack(lab_vals, 1)
        labs.append(lab_clip)
        used_cams.append(cam)

    if len(labs) == 0:
        return {}

    lab = torch.stack(labs, 0)
    with torch.no_grad():
        if label == 'rgb':
            latents = vae.encode_single_video_rgb(lab)
        elif label == 'depth':
            latents = vae.encode_single_video_depth(lab)
        else:
            raise ValueError(f"Unsupported latent label {label}")

    latents = torch.split(latents, latents.shape[0] // len(labs), 0)

    latents_dict = {}
    for i in range(len(latents)):
        latents_i = latents[i]
        for j in range(latents_i.shape[0]):
            latents_dict[j + time_st, used_cams[i]] = latents_i[j]

    return latents_dict