import mayavi.mlab as mlab
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch

from e2edet.utils.det3d.box_ops import center_to_corner_box2d, boxes_to_corners_3d

box_colormap = ["black", "peru", "red", "green", "purple"]


def visualize_pts(
    pts,
    fig=None,
    bgcolor=(0, 0, 0),
    fgcolor=(1.0, 1.0, 1.0),
    show_intensity=False,
    size=(600, 600),
    draw_origin=True,
):
    if isinstance(pts, torch.Tensor):
        pts = pts.cpu().numpy()

    if fig is None:
        fig = mlab.figure(
            figure=None, bgcolor=bgcolor, fgcolor=fgcolor, engine=None, size=size
        )

    if show_intensity:
        G = mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            pts[:, 3],
            mode="point",
            colormap="gnuplot",
            scale_factor=1,
            figure=fig,
        )
    else:
        G = mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            mode="point",
            colormap="gnuplot",
            scale_factor=1,
            figure=fig,
        )
    if draw_origin:
        mlab.points3d(0, 0, 0, color=(1, 1, 1), mode="cube", scale_factor=0.2)
        mlab.plot3d([0, 3], [0, 0], [0, 0], color=(0, 0, 1), tube_radius=0.1)
        mlab.plot3d([0, 0], [0, 3], [0, 0], color=(0, 1, 0), tube_radius=0.1)
        mlab.plot3d([0, 0], [0, 0], [0, 3], color=(1, 0, 0), tube_radius=0.1)

    return fig


def draw_sphere_pts(
    pts, color=(0, 1, 0), fig=None, bgcolor=(0, 0, 0), scale_factor=0.2
):
    if isinstance(pts, torch.Tensor):
        pts = pts.cpu().numpy()

    if fig is None:
        fig = mlab.figure(
            figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(600, 600)
        )

    if isinstance(color, np.ndarray) and color.shape[0] == 1:
        color = color[0]
        color = (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0)

    if isinstance(color, np.ndarray):
        pts_color = np.zeros((pts.__len__(), 4), dtype=np.uint8)
        pts_color[:, 0:3] = color
        pts_color[:, 3] = 255
        G = mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            np.arange(0, pts_color.__len__()),
            mode="sphere",
            scale_factor=scale_factor,
            figure=fig,
        )
        G.glyph.color_mode = "color_by_scalar"
        G.glyph.scale_mode = "scale_by_vector"
        G.module_manager.scalar_lut_manager.lut.table = pts_color
    else:
        mlab.points3d(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            mode="sphere",
            color=color,
            colormap="gnuplot",
            scale_factor=scale_factor,
            figure=fig,
        )

    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode="cube", scale_factor=0.2)
    mlab.plot3d(
        [0, 3],
        [0, 0],
        [0, 0],
        color=(0, 0, 1),
        line_width=3,
        tube_radius=None,
        figure=fig,
    )
    mlab.plot3d(
        [0, 0],
        [0, 3],
        [0, 0],
        color=(0, 1, 0),
        line_width=3,
        tube_radius=None,
        figure=fig,
    )
    mlab.plot3d(
        [0, 0],
        [0, 0],
        [0, 3],
        color=(1, 0, 0),
        line_width=3,
        tube_radius=None,
        figure=fig,
    )

    return fig


def draw_grid(x1, y1, x2, y2, fig, tube_radius=None, color=(0.5, 0.5, 0.5)):
    mlab.plot3d(
        [x1, x1],
        [y1, y2],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x2, x2],
        [y1, y2],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x1, x2],
        [y1, y1],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x1, x2],
        [y2, y2],
        [0, 0],
        color=color,
        tube_radius=tube_radius,
        line_width=1,
        figure=fig,
    )
    return fig


def draw_multi_grid_range(fig, grid_size=20, bv_range=(-60, -60, 60, 60)):
    for x in range(bv_range[0], bv_range[2], grid_size):
        for y in range(bv_range[1], bv_range[3], grid_size):
            fig = draw_grid(x, y, x + grid_size, y + grid_size, fig)

    return fig


def draw_scenes(
    points, gt_boxes=None, ref_boxes=None, ref_scores=None, ref_labels=None
):
    if isinstance(points, torch.Tensor):
        points = points.cpu().numpy()

    if ref_boxes is not None and not isinstance(ref_boxes, np.ndarray):
        ref_boxes = ref_boxes.cpu().numpy()
    if gt_boxes is not None and not isinstance(gt_boxes, np.ndarray):
        gt_boxes = gt_boxes.cpu().numpy()
    if ref_scores is not None and not isinstance(ref_scores, np.ndarray):
        ref_scores = ref_scores.cpu().numpy()
    if ref_labels is not None and not isinstance(ref_labels, np.ndarray):
        ref_labels = ref_labels.cpu().numpy()

    fig = visualize_pts(points)
    fig = draw_multi_grid_range(fig, bv_range=(0, -40, 80, 40))
    if gt_boxes is not None:
        corners3d = boxes_to_corners_3d(gt_boxes)
        fig = draw_corners3d(corners3d, fig=fig, color=(0, 0, 1), max_num=100)

    if ref_boxes is not None and len(ref_boxes) > 0:
        ref_corners3d = boxes_to_corners_3d(ref_boxes)
        if ref_labels is None:
            fig = draw_corners3d(
                ref_corners3d, fig=fig, color=(0, 1, 0), cls=ref_scores, max_num=100
            )
        else:
            for k in range(ref_labels.min(), ref_labels.max() + 1):
                cur_color = tuple(box_colormap[k % len(box_colormap)])
                mask = ref_labels == k
                fig = draw_corners3d(
                    ref_corners3d[mask],
                    fig=fig,
                    color=cur_color,
                    cls=ref_scores[mask],
                    max_num=100,
                )
    mlab.view(azimuth=-179, elevation=54.0, distance=104.0, roll=90.0)
    return fig


def draw_corners3d(
    corners3d,
    fig,
    color=(1, 1, 1),
    line_width=2,
    cls=None,
    tag="",
    max_num=500,
    tube_radius=None,
):
    """
    :param corners3d: (N, 8, 3)
    :param fig:
    :param color:
    :param line_width:
    :param cls:
    :param tag:
    :param max_num:
    :return:
    """
    import mayavi.mlab as mlab

    num = min(max_num, len(corners3d))
    for n in range(num):
        b = corners3d[n]  # (8, 3)

        if cls is not None:
            if isinstance(cls, np.ndarray):
                mlab.text3d(
                    b[6, 0],
                    b[6, 1],
                    b[6, 2],
                    "%.2f" % cls[n],
                    scale=(0.3, 0.3, 0.3),
                    color=color,
                    figure=fig,
                )
            else:
                mlab.text3d(
                    b[6, 0],
                    b[6, 1],
                    b[6, 2],
                    "%s" % cls[n],
                    scale=(0.3, 0.3, 0.3),
                    color=color,
                    figure=fig,
                )

        for k in range(0, 4):
            i, j = k, (k + 1) % 4
            mlab.plot3d(
                [b[i, 0], b[j, 0]],
                [b[i, 1], b[j, 1]],
                [b[i, 2], b[j, 2]],
                color=color,
                tube_radius=tube_radius,
                line_width=line_width,
                figure=fig,
            )

            i, j = k + 4, (k + 1) % 4 + 4
            mlab.plot3d(
                [b[i, 0], b[j, 0]],
                [b[i, 1], b[j, 1]],
                [b[i, 2], b[j, 2]],
                color=color,
                tube_radius=tube_radius,
                line_width=line_width,
                figure=fig,
            )

            i, j = k, k + 4
            mlab.plot3d(
                [b[i, 0], b[j, 0]],
                [b[i, 1], b[j, 1]],
                [b[i, 2], b[j, 2]],
                color=color,
                tube_radius=tube_radius,
                line_width=line_width,
                figure=fig,
            )

        i, j = 0, 5
        mlab.plot3d(
            [b[i, 0], b[j, 0]],
            [b[i, 1], b[j, 1]],
            [b[i, 2], b[j, 2]],
            color=color,
            tube_radius=tube_radius,
            line_width=line_width,
            figure=fig,
        )
        i, j = 1, 4
        mlab.plot3d(
            [b[i, 0], b[j, 0]],
            [b[i, 1], b[j, 1]],
            [b[i, 2], b[j, 2]],
            color=color,
            tube_radius=tube_radius,
            line_width=line_width,
            figure=fig,
        )

    return fig


def plot_bev(pts, boxes=None, labels=None, show_intensity=False, figure_num=None):
    if figure_num is not None:
        plt.figure(figure_num)
    else:
        fig, axs = plt.subplots(1, 1, figsize=(25, 25))

    plt.tick_params(labelbottom="off", labelleft="off")

    if show_intensity:
        plt.scatter(pts[:, 0], pts[:, 1], s=0.1, c=pts[3])
    else:
        plt.scatter(pts[:, 0], pts[:, 1], s=0.1)

    if boxes is not None:
        boxes = center_to_corner_box2d(boxes[:, :2], boxes[:, 3:5], boxes[:, -1])

        N = boxes.shape[0]
        for i in range(N):
            coord = boxes[i]
            x = coord[[0, 1, 2, 3, 0], 0].tolist()
            y = coord[[0, 1, 2, 3, 0], 1].tolist()

            plt.plot(x, y, color=box_colormap[labels[i]], linewidth=2.5)
    plt.show(block=False)
