# Copyright 2026 Toyota Research Institute.  All rights reserved.

import cv2
import torch
import numpy as np
from matplotlib.pyplot import get_cmap

from anydata.utils.decorators import iterate1
from anydata.utils.types import is_tensor, is_list
from anydata.geometry.depth import depth2inv

try:
    import flow_vis  #@IgnoreException
except:
    pass

def draw_points_on_image(frame, pixel_coords, palette="Purples", dot_size=5):
    """
    Draw colored dots on an image at given pixel coordinates.

    Args:
        frame: (H, W, C) numpy array image
        pixel_coords: (N, 2) numpy array of pixel coordinates (x, y)
        palette: Matplotlib colormap name for coloring dots
        dot_size: Radius of drawn circles

    Returns:
        (H, W, C) numpy array with dots drawn
    """
    frame = frame.astype(np.uint8).copy()
    if isinstance(pixel_coords, tuple):
        pixel_coords = [pixel_coords]

    # Get color palette
    color_palette = get_cmap(palette)
    color_palette = color_palette(np.linspace(0, 1, len(pixel_coords)))
    color_palette = (color_palette[:, :3] * 255).astype(np.uint8)
    color_palette = color_palette.tolist()

    for i, pixel_val in enumerate(pixel_coords):
        try:
            x, y = int(pixel_val[0]), int(pixel_val[1])

            # Skip extreme coordinates (likely overflow)
            if abs(x) > 1e6 or abs(y) > 1e6:
                continue

            # Skip coordinates far outside image bounds
            threshold = 1e-3
            if x < -threshold or x > frame.shape[1] + threshold or y < -threshold or y > frame.shape[0] + threshold:
                continue

            frame = cv2.circle(frame, (x, y), dot_size, color_palette[i], -1)
        except Exception as e:
            print(f"Error drawing point {pixel_val} on image: {e}")
            pass

    return frame


def flow_to_color(flow_uv, clip_flow=None):
    """Calculate color from optical flow"""
    # Clip if requested
    if clip_flow is not None:
        flow_uv = np.clip(flow_uv, -clip_flow, clip_flow)
    # Get optical flow channels
    u = flow_uv[:, :, 0]
    v = flow_uv[:, :, 1]
    # Calculate maximum radian
    rad_max = np.sqrt(2) * clip_flow if clip_flow is not None else \
        np.max(np.sqrt(np.square(u) + np.square(v)))
    # Normalize optical flow channels
    epsilon = 1e-5
    u = u / (rad_max + epsilon)
    v = v / (rad_max + epsilon)
    # Return colormap [0,1]
    return flow_vis.flow_uv_to_colors(u, v, convert_to_bgr=False) / 255


@iterate1
@iterate1
def viz_inv_depth(inv_depth, normalizer=None, percentile=95,
                  colormap='plasma', filter_zeros=False):
    """Converts an inverse depth map to a colormap for visualization"""
    if is_list(inv_depth):
        return [viz_inv_depth(
            inv[0], normalizer, percentile, colormap, filter_zeros)
            for inv in inv_depth]
    # If a tensor is provided, convert to numpy
    if is_tensor(inv_depth):
        # If it has a batch size, use first one
        if len(inv_depth.shape) == 4:
            inv_depth = inv_depth[0]
        # Squeeze if depth channel exists
        if len(inv_depth.shape) == 3:
            inv_depth = inv_depth.squeeze(0)
        inv_depth = inv_depth.detach().cpu().numpy()
    cm = get_cmap(colormap)
    if normalizer is None:
        if (inv_depth > 0).sum() == 0:
            normalizer = 1.0
        else:
            normalizer = np.percentile(
                inv_depth[inv_depth > 0] if filter_zeros else inv_depth, percentile)
    inv_depth = inv_depth / (normalizer + 1e-6)
    colormap = cm(np.clip(inv_depth, 0., 1.0))[:, :, :3]
    colormap[inv_depth == 0] = 0
    return colormap


@iterate1
@iterate1
def viz_depth(depth, *args, **kwargs):
    """Same as viz_inv_depth, but takes depth as input instead"""
    return viz_inv_depth(depth2inv(depth), *args, **kwargs)


@iterate1
@iterate1
def viz_normals(normals):
    """Converts normals map to a colormap for visualization"""
    # If a tensor is provided, convert to numpy
    if is_tensor(normals):
        normals = normals.permute(1, 2, 0).detach().cpu().numpy()
    return (normals + 1) / 2


@iterate1
@iterate1
def viz_optflow(optflow, clip_value=100., multiplier=None, normalizer=None):
    """Returns a colorized version of an optical flow map"""
    # If a tensor is provided, convert to numpy
    if is_list(optflow):
        return [viz_optflow(opt[0], multiplier, normalizer) for opt in optflow]
    if is_tensor(optflow):
        if len(optflow.shape) == 4:
            optflow = optflow[0]
        optflow = optflow.permute(1, 2, 0).detach().cpu().numpy()
    if multiplier is not None:
        optflow = optflow * multiplier
    if normalizer is not None:
        optflow = (optflow / np.abs(optflow).max()) * normalizer
    # Return colorized optical flow
    return flow_to_color(optflow, clip_flow=clip_value)


@iterate1
@iterate1
def viz_photo(photo, colormap='viridis', normalize=False):
    """Return a colorized version of photometric loss"""
    if is_tensor(photo):
        if len(photo.shape) == 4:
            photo = photo[0]
        if len(photo.shape) == 3:
            photo = photo.squeeze(0)
        photo = photo.detach().cpu().numpy()
    cm = get_cmap(colormap)
    if normalize:
        photo -= photo.min()
        photo /= (photo.max() + 1e-6)
    colormap = cm(np.clip(photo, 0., 1.0))[:, :, :3]
    colormap[photo == 0] = 0
    return colormap


@iterate1
@iterate1
def viz_semantic(semantic, ontology):
    """Return colorized version of a semantic map (given ontology)"""
    # If it is a tensor, cast to numpy
    if is_tensor(semantic):
        if semantic.dim() == 3:
            semantic = semantic.squeeze(0)
        semantic = semantic.detach().cpu().numpy()
    color = np.zeros([semantic.shape[0], semantic.shape[1], 3])
    for key, val in ontology.items():
        color[semantic == int(key)] = np.array(val['color'])
    return color / 255            


@iterate1
@iterate1
def viz_camera(camera, filter_zeros=False):
    """Return colorized version of camera rays"""
    if is_tensor(camera):
        # If it's a tensor, reshape it
        rays = camera[-3:].permute(1, 2, 0).detach().cpu().numpy()
    else:
        # If it's a camera, get viewing rays
        rays = camera.no_translation().get_viewdirs(normalize=True, flatten=False, to_world=True)
        rays = rays[0].permute(1, 2, 0).detach().cpu().numpy()

    if filter_zeros:
        # Only normalize valid rays
        valid = rays.sum(2) != 0.0
        valid = valid.repeat(1, 3, 1, 1)
        rays[valid] = (rays[valid] + 1) / 2
    else:
        # Normalize all rays
        rays = (rays + 1) / 2

    return rays

@iterate1
@iterate1
def viz_stddev(stddev, normalizer=None):
    """Return colorized version of standard deviation"""
    return viz_inv_depth(stddev, colormap='jet', normalizer=normalizer)


@iterate1
@iterate1
def viz_scene_flow(scnflow, clip_value=10):
    # If a tensor is provided, convert to numpy
    if is_list(scnflow):
        return [viz_scene_flow(scn[0], clip_value) for scn in scnflow]
    # If a tensor is provided, convert to numpy
    if is_tensor(scnflow):
        if len(scnflow.shape) == 4:
            scnflow = scnflow[0]
        scnflow = scnflow.permute(1, 2, 0).detach().cpu().numpy()
    return (np.clip(scnflow / clip_value, -1, 1) + 1) / 2


@iterate1
@iterate1
def viz_features(features):

    c, h, w = features.shape
    features = features.permute(1, 2, 0).reshape(-1, c).float()

    with torch.amp.autocast(features.device.type, enabled=False):
        (U, S, V) = torch.pca_lowrank(features, q=None, center=True, niter=2)

    k = 3
    pca = torch.matmul(features, V[:, :k])

    low =  pca.min(dim=0).values
    high = pca.max(dim=0).values
    pca = (pca - low) / (high - low)

    pca = pca.view(h, w, k).detach().cpu().numpy()
    return pca