"""Visualization of predicted and ground truth for a single batch."""

from typing import Any, Dict

import numpy as np
import torch

from siclib.geometry.perspective_fields import get_latitude_field
from siclib.models.utils.metrics import latitude_error, up_error
from siclib.utils.conversions import rad2deg
from siclib.utils.tensor import batch_to_device
from siclib.visualization.viz2d import (
    plot_confidences,
    plot_heatmaps,
    plot_image_grid,
    plot_latitudes,
    plot_vector_fields,
)


def make_up_figure(
    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
) -> Dict[str, Any]:
    """Get predicted and ground truth up fields and errors.

    Args:
        pred (Dict[str, torch.Tensor]): Predicted up field.
        data (Dict[str, torch.Tensor]): Ground truth up field.
        n_pairs (int): Number of pairs to visualize.

    Returns:
        Dict[str, Any]: Dictionary with figure.
    """
    pred = batch_to_device(pred, "cpu", detach=True)
    data = batch_to_device(data, "cpu", detach=True)

    n_pairs = min(n_pairs, len(data["image"]))

    if "up_field" not in pred.keys():
        return {}

    errors = up_error(pred["up_field"], data["up_field"])

    up_fields = []
    for i in range(n_pairs):
        row = [data["up_field"][i], pred["up_field"][i], errors[i]]
        titles = ["Up GT", "Up Pred", "Up Error"]

        if "up_confidence" in pred.keys():
            row += [pred["up_confidence"][i]]
            titles += ["Up Confidence"]

        row = [r.float().numpy() if isinstance(r, torch.Tensor) else r for r in row]
        up_fields.append(row)

    # create figure
    N, M = len(up_fields), len(up_fields[0]) + 1
    imgs = [[data["image"][i].permute(1, 2, 0).cpu().clip(0, 1)] * M for i in range(n_pairs)]
    fig, ax = plot_image_grid(imgs, titles=[["Image"] + titles] * N, return_fig=True, set_lim=True)
    ax = np.array(ax)

    for i in range(n_pairs):
        plot_vector_fields(up_fields[i][:2], axes=ax[i, [1, 2]])
        plot_heatmaps([up_fields[i][2]], cmap="turbo", colorbar=True, axes=ax[i, [3]])

        if "up_confidence" in pred.keys():
            plot_confidences([up_fields[i][3]], axes=ax[i, [4]])

    return {"up": fig}


def make_latitude_figure(
    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
) -> Dict[str, Any]:
    """Get predicted and ground truth latitude fields and errors.

    Args:
        pred (Dict[str, torch.Tensor]): Predicted latitude field.
        data (Dict[str, torch.Tensor]): Ground truth latitude field.
        n_pairs (int, optional): Number of pairs to visualize. Defaults to 2.

    Returns:
        Dict[str, Any]: Dictionary with figure.
    """
    pred = batch_to_device(pred, "cpu", detach=True)
    data = batch_to_device(data, "cpu", detach=True)

    n_pairs = min(n_pairs, len(data["image"]))
    latitude_fields = []

    if "latitude_field" not in pred.keys():
        return {}

    errors = latitude_error(pred["latitude_field"], data["latitude_field"])
    for i in range(n_pairs):
        row = [
            rad2deg(data["latitude_field"][i][0]),
            rad2deg(pred["latitude_field"][i][0]),
            errors[i],
        ]
        titles = ["Latitude GT", "Latitude Pred", "Latitude Error"]

        if "latitude_confidence" in pred.keys():
            row += [pred["latitude_confidence"][i]]
            titles += ["Latitude Confidence"]

        row = [r.float().numpy() if isinstance(r, torch.Tensor) else r for r in row]
        latitude_fields.append(row)

    # create figure
    N, M = len(latitude_fields), len(latitude_fields[0]) + 1
    imgs = [[data["image"][i].permute(1, 2, 0).cpu().clip(0, 1)] * M for i in range(n_pairs)]
    fig, ax = plot_image_grid(imgs, titles=[["Image"] + titles] * N, return_fig=True, set_lim=True)
    ax = np.array(ax)

    for i in range(n_pairs):
        plot_latitudes(latitude_fields[i][:2], is_radians=False, axes=ax[i, [1, 2]])
        plot_heatmaps([latitude_fields[i][2]], cmap="turbo", colorbar=True, axes=ax[i, [3]])

        if "latitude_confidence" in pred.keys():
            plot_confidences([latitude_fields[i][3]], axes=ax[i, [4]])

    return {"latitude": fig}


def make_camera_figure(
    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
) -> Dict[str, Any]:
    """Get predicted and ground truth camera parameters.

    Args:
        pred (Dict[str, torch.Tensor]): Predicted camera parameters.
        data (Dict[str, torch.Tensor]): Ground truth camera parameters.
        n_pairs (int, optional): Number of pairs to visualize. Defaults to 2.

    Returns:
        Dict[str, Any]: Dictionary with figure.
    """
    pred = batch_to_device(pred, "cpu", detach=True)
    data = batch_to_device(data, "cpu", detach=True)

    n_pairs = min(n_pairs, len(data["image"]))

    if "camera" not in pred.keys():
        return {}

    latitudes = []
    for i in range(n_pairs):
        titles = ["Cameras GT"]
        row = [get_latitude_field(data["camera"][i], data["gravity"][i])]

        if "camera" in pred.keys() and "gravity" in pred.keys():
            row += [get_latitude_field(pred["camera"][i], pred["gravity"][i])]
            titles += ["Cameras Pred"]

        row = [rad2deg(r).squeeze(-1).float().numpy()[0] for r in row]
        latitudes.append(row)

    # create figure
    N, M = len(latitudes), len(latitudes[0]) + 1
    imgs = [[data["image"][i].permute(1, 2, 0).cpu().clip(0, 1)] * M for i in range(n_pairs)]
    fig, ax = plot_image_grid(imgs, titles=[["Image"] + titles] * N, return_fig=True, set_lim=True)
    ax = np.array(ax)

    for i in range(n_pairs):
        plot_latitudes(latitudes[i], is_radians=False, axes=ax[i, 1:])

    return {"camera": fig}


def make_perspective_figures(
    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
) -> Dict[str, Any]:
    """Get predicted and ground truth perspective fields.

    Args:
        pred (Dict[str, torch.Tensor]): Predicted perspective fields.
        data (Dict[str, torch.Tensor]): Ground truth perspective fields.
        n_pairs (int, optional): Number of pairs to visualize. Defaults to 2.

    Returns:
        Dict[str, Any]: Dictionary with figure.
    """
    n_pairs = min(n_pairs, len(data["image"]))
    figures = make_up_figure(pred, data, n_pairs)
    figures |= make_latitude_figure(pred, data, n_pairs)
    figures |= make_camera_figure(pred, data, n_pairs)

    {f.tight_layout() for f in figures.values()}

    return figures