# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
Visualization utility functions for colorization and color bars.
"""

import dataclasses
from typing import Optional, Tuple

import numpy as np
import torch
import cv2
import matplotlib.cm as cm


@dataclasses.dataclass
class CameraState:
    """Camera state for rendering."""
    fov: float
    aspect: float
    c2w: np.ndarray

    def get_K(self, img_wh: Tuple[int, int]) -> np.ndarray:
        """Get camera intrinsic matrix from FOV and image size."""
        W, H = img_wh
        focal_length = H / 2.0 / np.tan(self.fov / 2.0)
        K = np.array([
            [focal_length, 0.0, W / 2.0],
            [0.0, focal_length, H / 2.0],
            [0.0, 0.0, 1.0],
        ])
        return K


def get_vertical_colorbar(
    h: int,
    vmin: float,
    vmax: float,
    cmap_name: str = "jet",
    label: Optional[str] = None,
    cbar_precision: int = 2
) -> np.ndarray:
    """
    Create a vertical colorbar image.

    Args:
        h: Height in pixels
        vmin: Minimum value
        vmax: Maximum value
        cmap_name: Colormap name
        label: Optional label for the colorbar
        cbar_precision: Decimal precision for tick labels

    Returns:
        Colorbar image as numpy array (H, W, 3)
    """
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_agg import FigureCanvasAgg
    import matplotlib as mpl

    fig = Figure(figsize=(2, 8), dpi=100)
    fig.subplots_adjust(right=1.5)
    canvas = FigureCanvasAgg(fig)

    ax = fig.add_subplot(111)
    cmap = cm.get_cmap(cmap_name)
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)

    tick_cnt = 6
    tick_loc = np.linspace(vmin, vmax, tick_cnt)
    cb1 = mpl.colorbar.ColorbarBase(
        ax, cmap=cmap, norm=norm, ticks=tick_loc, orientation="vertical"
    )

    tick_label = [str(np.round(x, cbar_precision)) for x in tick_loc]
    if cbar_precision == 0:
        tick_label = [x[:-2] for x in tick_label]

    cb1.set_ticklabels(tick_label)
    cb1.ax.tick_params(labelsize=18, rotation=0)
    if label is not None:
        cb1.set_label(label)

    canvas.draw()
    s, (width, height) = canvas.print_to_buffer()

    im = np.frombuffer(s, np.uint8).reshape((height, width, 4))
    im = im[:, :, :3].astype(np.float32) / 255.0

    if h != im.shape[0]:
        w = int(im.shape[1] / im.shape[0] * h)
        im = cv2.resize(im, (w, h), interpolation=cv2.INTER_AREA)

    return im


def colorize_np(
    x: np.ndarray,
    cmap_name: str = "jet",
    mask: Optional[np.ndarray] = None,
    range: Optional[Tuple[float, float]] = None,
    append_cbar: bool = False,
    cbar_in_image: bool = False,
    cbar_precision: int = 2,
) -> np.ndarray:
    """
    Turn a grayscale image into a color image.

    Args:
        x: Input grayscale image [H, W]
        cmap_name: Colormap name
        mask: Optional mask image [H, W]
        range: Value range for scaling [min, max], automatic if None
        append_cbar: Whether to append colorbar
        cbar_in_image: Put colorbar inside image
        cbar_precision: Colorbar tick precision

    Returns:
        Colorized image [H, W, 3]
    """
    if range is not None:
        vmin, vmax = range
    elif mask is not None:
        vmin = np.min(x[mask][np.nonzero(x[mask])])
        vmax = np.max(x[mask])
        x[np.logical_not(mask)] = vmin
    else:
        vmin, vmax = np.percentile(x, (1, 100))
        vmax += 1e-6

    x = np.clip(x, vmin, vmax)
    x = (x - vmin) / (vmax - vmin)

    cmap = cm.get_cmap(cmap_name)
    x_new = cmap(x)[:, :, :3]

    if mask is not None:
        mask = np.float32(mask[:, :, np.newaxis])
        x_new = x_new * mask + np.ones_like(x_new) * (1.0 - mask)

    cbar = get_vertical_colorbar(
        h=x.shape[0],
        vmin=vmin,
        vmax=vmax,
        cmap_name=cmap_name,
        cbar_precision=cbar_precision,
    )

    if append_cbar:
        if cbar_in_image:
            x_new[:, -cbar.shape[1]:, :] = cbar
        else:
            x_new = np.concatenate(
                (x_new, np.zeros_like(x_new[:, :5, :]), cbar), axis=1
            )
        return x_new
    else:
        return x_new


def colorize(
    x: torch.Tensor,
    cmap_name: str = "jet",
    mask: Optional[torch.Tensor] = None,
    range: Optional[Tuple[float, float]] = None,
    append_cbar: bool = False,
    cbar_in_image: bool = False
) -> torch.Tensor:
    """
    Turn a grayscale image into a color image (PyTorch tensor version).

    Args:
        x: Grayscale image tensor [H, W] or [B, H, W]
        cmap_name: Colormap name
        mask: Optional mask tensor [H, W] or [B, H, W]
        range: Value range for scaling
        append_cbar: Whether to append colorbar
        cbar_in_image: Put colorbar inside image

    Returns:
        Colorized tensor
    """
    device = x.device
    x = x.cpu().numpy()
    if mask is not None:
        mask = mask.cpu().numpy() > 0.99
        kernel = np.ones((3, 3), np.uint8)

    if x.ndim == 2:
        x = x[None]
        if mask is not None:
            mask = mask[None]

    out = []
    for x_ in x:
        if mask is not None:
            mask = cv2.erode(mask.astype(np.uint8), kernel, iterations=1).astype(bool)

        x_ = colorize_np(x_, cmap_name, mask, range, append_cbar, cbar_in_image)
        out.append(torch.from_numpy(x_).to(device).float())
    out = torch.stack(out).squeeze(0)
    return out