import numpy as np

from gempy.core.data import GeoModel
from gempy.core.data.orientations import OrientationsTable
from gempy.core.data.surface_points import SurfacePointsTable
from matplotlib.colors import ListedColormap

from gempy_viewer.modules.plot_2d.plot_2d_utils import get_geo_model_cmap
from gempy_viewer.modules.plot_3d.vista import GemPyToVista
from gempy_viewer.optional_dependencies import require_pyvista


def plot_data(gempy_vista: GemPyToVista,
              model: GeoModel,
              arrows_factor: float,
              transformed_data: bool = False,
              **kwargs):
    if transformed_data:
        surface_points_copy = model.surface_points_copy_transformed
        orientations_copy = model.orientations_copy_transformed
    else:
        surface_points_copy = model.surface_points_copy
        orientations_copy = model.orientations_copy

    plot_surface_points(
        gempy_vista=gempy_vista,
        surface_points=surface_points_copy,
        element_colors=model.structural_frame.elements_colors
    )

    plot_orientations(
        gempy_vista=gempy_vista,
        orientations=orientations_copy,
        surface_points=surface_points_copy,
        arrows_factor=arrows_factor
    )


def plot_surface_points(
        gempy_vista: GemPyToVista,
        surface_points: SurfacePointsTable,
        render_points_as_spheres=True,
        element_colors=None,
        point_size=10
):
    # Selecting the surfaces to plot
    xyz = surface_points.xyz
    if transfromed_data := False:  # TODO: Expose this to user
        xyz = surface_points.model_transform.apply(xyz)

    pv = require_pyvista()
    poly = pv.PolyData(xyz)

    ids = surface_points.ids
    if ids.shape[0] == 0:
        return
    vectorize_ids = _vectorize_ids(ids, ids)
    poly['id'] = vectorize_ids

    custom_cmap = ListedColormap(element_colors)
    
    gempy_vista.surface_points_mesh = poly
    gempy_vista.surface_points_actor = gempy_vista.p.add_mesh(
        mesh=poly,
        scalars='id',
        render_points_as_spheres=render_points_as_spheres,
        point_size=point_size,
        show_scalar_bar=False,
        cmap=custom_cmap,
        clim=(0, np.unique(vectorize_ids).shape[0])
    )


def plot_orientations(
        gempy_vista: GemPyToVista,
        orientations: OrientationsTable,
        surface_points: SurfacePointsTable,
        arrows_factor: float,
):
    orientations_xyz = orientations.xyz
    orientations_grads = orientations.grads

    if orientations_xyz.shape[0] == 0:
        return

    pv = require_pyvista()
    poly = pv.PolyData(orientations_xyz)

    poly['id'] = _vectorize_ids(
        mapping_ids=surface_points.ids,
        ids_to_map=orientations.ids
    )
    poly['vectors'] = orientations_grads

    arrows = poly.glyph(
        orient='vectors',
        scale=False,
        factor=arrows_factor,
    )

    gempy_vista.orientations_actor = gempy_vista.p.add_mesh(
        mesh=arrows,
        scalars='id',
        show_scalar_bar=False
    )
    gempy_vista.orientations_mesh = arrows


def _vectorize_ids(mapping_ids, ids_to_map):
    def _mapping_dict(ids):
        unique_values, first_indices = np.unique(ids, return_index=True)  # Find the unique elements and their first indices
        unique_values_order = unique_values[np.argsort(first_indices)]  # Sort the unique values by their first appearance in `a`
        # Flip order to please pyvista vertical scalarbar
        unique_values_order = unique_values_order[::-1]
        mapping_dict = {value: i + 1 for i, value in enumerate(unique_values_order)}  # Use a dictionary to map the original numbers to new values
        return mapping_dict

    mapping_dict = _mapping_dict(mapping_ids)
    mapped_array = np.vectorize(mapping_dict.get)(ids_to_map)  # Map the original array to the new values
    return mapped_array
