#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations
import os
import numpy as np

import vedo.vtkclasses as vtki

from vedo.pointcloud import Points
from vedo.mesh import Mesh
from vedo.volume import Volume

import vedo
from vedo import settings
from vedo import utils

__doc__ = """Submodule to delegate jupyter notebook rendering"""

__all__ = []


############################################################################################
def get_notebook_backend(actors2show=()):
    """Return the appropriate notebook viewer"""

    backend = settings.default_backend

    if backend == "2d":
        return start_2d()
    if backend == "k3d":
        return start_k3d(actors2show)
    if backend.startswith("trame"):
        return start_trame()
    if backend.startswith("ipyvtk"):
        return start_ipyvtklink()
    if backend.startswith("panel"):
        return start_panel()

    vedo.logger.error(f"Unknown jupyter backend: {settings.default_backend}")
    return None


#####################################################################################
def start_2d():
    """Start a 2D display in the notebook"""
    try:
        import PIL.Image
        # import IPython
    except ImportError:
        print("PIL or IPython not available")
        return

    plt = vedo.current_plotter()
    if not plt:
        vedo.logger.error("No active Plotter found for the 2d backend.")
        return None

    if hasattr(plt, "window") and plt.window:
        try:
            nn = vedo.file_io.screenshot(asarray=True, scale=1)
            pil_img = PIL.Image.fromarray(nn)
        except ValueError:
            return

        # IPython.display.display(pil_img)
        vedo.set_current_notebook_plotter(pil_img)
        if settings.backend_autoclose and plt.renderer == plt.renderers[-1]:
            plt.close()
        return pil_img

#####################################################################################
def start_panel():
    """Start a panel display in the notebook"""
    try:
        import panel as pn # type: ignore
        pn.extension('vtk', design='material', sizing_mode='stretch_width', template='material')
        # pn.state.template.config.raw_css.append("""
        # #main {
        # padding: 0;
        # }""")
    except ImportError:
        print("panel is not installed, try:\n> conda install panel")
        return None

    print("panel backend NOT YET FUNCTIONAL")
    plt = vedo.current_plotter()
    if not plt:
        vedo.logger.error("No active Plotter found for the panel backend.")
        return None

    if hasattr(plt, "window") and plt.window:
        plt.renderer.ResetCamera()
        vtkpan = pn.pane.VTK(
            plt.window,
            margin=0, sizing_mode='stretch_both',
            min_height=600,
            orientation_widget=True,
            enable_keybindings=True,
        )
        vedo.set_current_notebook_plotter(vtkpan)
        return vedo.current_notebook_plotter()

####################################################################################
def start_k3d(actors2show):
    """Start a k3d display in the notebook"""
    try:
        # https://github.com/K3D-tools/K3D-jupyter
        import k3d
    except ModuleNotFoundError:
        print("\nCannot find k3d, install with:  pip install k3d")
        return None

    plt = vedo.current_plotter()
    if not plt:
        vedo.logger.error("No active Plotter found for the k3d backend.")
        return None

    def _setup_scalar_metadata(polydata, mapper):
        """Build scalar metadata for k3d color mapping."""
        vtkdata = polydata.GetPointData()
        vtkscals = vtkdata.GetScalars()

        if vtkscals is None:
            vtkdata = polydata.GetCellData()
            vtkscals = vtkdata.GetScalars()
            if vtkscals is not None:
                c2p = vtki.new("CellDataToPointData")
                c2p.SetInputData(polydata)
                c2p.Update()
                polydata = c2p.GetOutput()
                vtkdata = polydata.GetPointData()
                vtkscals = vtkdata.GetScalars()

        if vtkscals is None:
            return polydata, None, None, None, None

        if not vtkscals.GetName():
            vtkscals.SetName("scalars")
        scals_min, scals_max = mapper.GetScalarRange()
        color_attribute = (vtkscals.GetName(), scals_min, scals_max)
        lut = mapper.GetLookupTable()
        lut.Build()
        kcmap = []
        nlut = lut.GetNumberOfTableValues()
        for i in range(nlut):
            r, g, b, _ = lut.GetTableValue(i)
            kcmap += [i / (nlut - 1), r, g, b]
        return polydata, vtkscals, color_attribute, kcmap, (scals_min, scals_max)

    already_has_axes = False
    actors2show2 = []
    for ia in actors2show:
        if not ia:
            continue

        try:
            if ia.name == "Axes":
                already_has_axes = True
        except AttributeError:
            pass

        if isinstance(ia, vedo.Assembly):  # unpack assemblies
            actors2show2 += ia.recursive_unpack()
        else:
            actors2show2.append(ia)

    nbplot = k3d.plot(
        axes=["x", "y", "z"],
        menu_visibility=settings.k3d_menu_visibility,
        height=settings.k3d_plot_height,
        antialias=settings.k3d_antialias,
        background_color=_rgb2int(vedo.get_color(plt.backgrcol)),
        camera_fov=30.0,  # deg (this is the vtk default)
        lighting=settings.k3d_lighting,
        grid_color=_rgb2int(vedo.get_color(settings.k3d_axes_color)),
        label_color=_rgb2int(vedo.get_color(settings.k3d_axes_color)),
        axes_helper=settings.k3d_axes_helper,
    )

    # set k3d camera
    vedo.set_current_notebook_plotter(nbplot)
    nbplot.camera_auto_fit = settings.k3d_camera_autofit
    nbplot.axes_helper = settings.k3d_axes_helper
    nbplot.grid_auto_fit = settings.k3d_grid_autofit

    if already_has_axes:
        nbplot.grid_visible = False
    if settings.k3d_grid_visible is not None: # override if set
        nbplot.grid_visible = settings.k3d_grid_visible

    if plt.camera:
        nbplot.camera = utils.vtkCameraToK3D(plt.camera)

    for ia in actors2show2:

        if isinstance(ia, (vtki.vtkAssembly, vtki.vtkActor2D)):
            continue

        if hasattr(ia, "actor") and isinstance(
            ia.actor, (vtki.vtkAssembly, vtki.vtkActor2D)
        ):
            continue

        iacloned = ia

        kobj = None
        kcmap = None
        color_attribute = None
        vtkscals = None
        name = None
        if hasattr(ia, "filename"):
            if ia.filename:
                name = os.path.basename(ia.filename)
            if ia.name:
                name = os.path.basename(ia.name)

        ################################################################## scalars
        # work out scalars first, Points Lines are also Mesh objs
        if isinstance(ia, Points):
            # print('scalars', ia.name, ia.npoints)
            iap = ia.properties

            if ia.dataset.GetNumberOfPolys():
                iacloned = ia.clone()
                iapoly = iacloned.clean().triangulate().compute_normals().dataset
            else:
                iapoly = ia.dataset

            if ia.mapper.GetScalarVisibility() and ia.mapper.GetColorMode() > 0:
                iapoly, vtkscals, color_attribute, kcmap, scal_range = _setup_scalar_metadata(
                    iapoly, ia.mapper
                )
                if scal_range is not None:
                    scals_min, scals_max = scal_range

            else:
                color_attribute = ia.color()

        #####################################################################Volume
        if isinstance(ia, Volume):
            # print('Volume', ia.name, ia.dimensions())
            kx, ky, _ = ia.dimensions()
            arr = ia.pointdata[0]
            kimage = arr.reshape(-1, ky, kx)

            color_transfer_function = ia.properties.GetRGBTransferFunction()
            kcmap = []
            for i in range(128):
                r, g, b = color_transfer_function.GetColor(i / 127)
                kcmap += [i / 127, r, g, b]

            kbounds = np.array(ia.dataset.GetBounds()) + np.repeat(
                np.array(ia.dataset.GetSpacing()) / 2.0, 2
            ) * np.array([-1, 1] * 3)

            kobj = k3d.volume(
                kimage.astype(np.float32),
                color_map=kcmap,
                # color_range=ia.dataset.GetScalarRange(),
                alpha_coef=10,
                bounds=kbounds,
                name=name,
            )
            nbplot += kobj

        ################################################################ Text2D
        elif isinstance(ia, vedo.Text2D):
            # print('Text2D', ia.GetPosition())
            pos = (ia.GetPosition()[0], 1.0 - ia.GetPosition()[1])

            kobj = k3d.text2d(
                ia.text(),
                position=pos,
                color=_rgb2int(vedo.get_color(ia.c())),
                is_html=True,
                size=ia.properties.GetFontSize() / 22.5 * 1.5,
                label_box=bool(ia.properties.GetFrame()),
                # reference_point='bl',
            )
            nbplot += kobj

        ################################################################# Lines
        elif (
            hasattr(ia, "lines")
            and ia.dataset.GetNumberOfLines()
            and ia.dataset.GetNumberOfPolys() == 0
        ):

            for i, ln_idx in enumerate(ia.lines):

                if i > 200:
                    vedo.logger.warning("in k3d, nr. of lines is limited to 200.")
                    break

                pts = ia.coordinates[ln_idx]

                aves = ia.diagonal_size() * iap.GetLineWidth() / 100

                kobj = k3d.line(
                    pts.astype(np.float32),
                    color=_rgb2int(iap.GetColor()),
                    opacity=iap.GetOpacity(),
                    shader=settings.k3d_line_shader,
                    width=aves.astype(float),
                    name=name,
                )
                nbplot += kobj

        ################################################################## Mesh
        elif isinstance(ia, Mesh) and ia.npoints and ia.dataset.GetNumberOfPolys():
            # print('Mesh', ia.name, ia.npoints, len(ia.cells))

            if not vtkscals:
                color_attribute = None

            cols = []
            if ia.mapper.GetColorMode() == 2:  # direct RGB colors

                vcols = ia.dataset.GetPointData().GetScalars()

                if not vcols:
                    iacloned.map_cells_to_points()
                    vcols = iacloned.dataset.GetPointData().GetScalars()

                if vcols and vcols.GetNumberOfComponents() in (3, 4):
                    # vedo.logger.info("found RGB direct colors in Mesh")
                    cols = utils.vtk2numpy(vcols).astype(np.uint32)
                    cols = 65536 * cols[:, 0] + 256 * cols[:, 1] + cols[:, 2]

                    kobj = k3d.mesh(
                        iacloned.coordinates,
                        iacloned.cells,
                        colors=cols,
                        name=name,
                        opacity=iap.GetOpacity(),
                        side="double",
                        wireframe=(iap.GetRepresentation() == 1),
                    )

                else:
                    vedo.logger.warning("could not find RGB direct colors in Mesh")
                    kobj = k3d.mesh(
                        iacloned.coordinates,
                        iacloned.cells,
                        color=_rgb2int(iap.GetColor()),
                        name=name,
                        opacity=iap.GetOpacity(),
                        side="double",
                        wireframe=(iap.GetRepresentation() == 1),
                    )

            else:

                kobj = k3d.vtk_poly_data(
                    iapoly,
                    name=name,
                    color=_rgb2int(iap.GetColor()),
                    color_attribute=color_attribute,
                    color_map=kcmap,
                    opacity=iap.GetOpacity(),
                    side="double",
                    wireframe=(iap.GetRepresentation() == 1),
                )

            if iap.GetInterpolation() == 0:
                kobj.flat_shading = True

            nbplot += kobj

        #####################################################################Points
        elif isinstance(ia, Points):
            # print('Points', ia.name, ia.npoints)
            kcols = []
            if kcmap is not None and vtkscals:
                scals = utils.vtk2numpy(vtkscals)
                kcols = k3d.helpers.map_colors(
                    scals, kcmap, [scals_min, scals_max]
                ).astype(np.uint32)

            aves = ia.average_size() * iap.GetPointSize() / 200

            kobj = k3d.points(
                ia.coordinates.astype(np.float32),
                color=_rgb2int(iap.GetColor()),
                colors=kcols,
                opacity=iap.GetOpacity(),
                shader=settings.k3d_point_shader,
                point_size=aves.astype(float),
                name=name,
            )
            nbplot += kobj

        #####################################################################
        elif isinstance(ia, vedo.Image):
            vedo.logger.error("Sorry Image objects are not supported in k3d.")

    if plt and settings.backend_autoclose:
        plt.close()
    return nbplot


#####################################################################################
def start_trame():
    """Start a trame display in the notebook"""
    try:
        from trame.app import get_server, jupyter # type: ignore
        from trame.ui.vuetify import VAppLayout # type: ignore
        from trame.widgets import vtk as t_vtk, vuetify # type: ignore
    except ImportError:
        print("trame is not installed, try:\n> pip install trame==2.5.2")
        return

    plt = vedo.current_plotter()
    if not plt:
        vedo.logger.error("No active Plotter found for the trame backend.")
        return None
    if hasattr(plt, "window") and plt.window:
        plt.renderer.ResetCamera()
        server = get_server("jupyter-1")
        state, ctrl = server.state, server.controller
        plt.server = server
        plt.controller = ctrl
        plt.state = state

        with VAppLayout(server) as layout:

            with layout.root:

                with vuetify.VContainer(fluid=True, classes="pa-0 fill-height"):
                    plt.reset_camera()
                    view = t_vtk.VtkLocalView(plt.window)
                    ctrl.view_update = view.update
                    ctrl.view_reset_camera = view.reset_camera

        ctrl.on_server_exited.add(lambda **_: print("trame server exited"))
        jupyter.show(server)
        return
    vedo.logger.error("No window present for the trame backend.")
    return


#####################################################################################
def start_ipyvtklink():
    try:
        from ipyvtklink.viewer import ViewInteractiveWidget # type: ignore
    except ImportError:
        print("ipyvtklink is not installed, try:\n> pip install ipyvtklink")
        return None

    plt = vedo.current_plotter()
    if not plt:
        vedo.logger.error("No active Plotter found for the ipyvtklink backend.")
        return None
    if hasattr(plt, "window") and plt.window:
        plt.renderer.ResetCamera()
        nbplot = ViewInteractiveWidget(
            plt.window, allow_wheel=True, quality=100, quick_quality=50
        )
        vedo.set_current_notebook_plotter(nbplot)
        return nbplot
    vedo.logger.error("No window present for the ipyvtklink backend.")
    return None


#####################################################################################
def _rgb2int(rgb_tuple):
    # Return the int number of a color from (r,g,b), with 0<r<1 etc.
    rgb = (int(rgb_tuple[0] * 255), int(rgb_tuple[1] * 255), int(rgb_tuple[2] * 255))
    return 65536 * rgb[0] + 256 * rgb[1] + rgb[2]
