from __future__ import annotations

__all__ = ["OpenGLPMobject", "OpenGLPGroup", "OpenGLPMPoint"]

from typing import TYPE_CHECKING

import moderngl
import numpy as np

from manim.constants import *
from manim.mobject.opengl.opengl_mobject import OpenGLMobject
from manim.utils.bezier import interpolate
from manim.utils.color import (
    BLACK,
    WHITE,
    YELLOW,
    ParsableManimColor,
    color_gradient,
    color_to_rgba,
)
from manim.utils.config_ops import _Uniforms
from manim.utils.iterables import resize_with_interpolation

if TYPE_CHECKING:
    from typing import Self

    from manim.typing import (
        FloatRGBA_Array,
        FloatRGBALike_Array,
        Point3D_Array,
        Point3DLike_Array,
    )

__all__ = ["OpenGLPMobject", "OpenGLPGroup", "OpenGLPMPoint"]


class OpenGLPMobject(OpenGLMobject):
    shader_folder = "true_dot"
    # Scale for consistency with cairo units
    OPENGL_POINT_RADIUS_SCALE_FACTOR = 0.01
    shader_dtype = [
        ("point", np.float32, (3,)),
        ("color", np.float32, (4,)),
    ]

    point_radius = _Uniforms()

    def __init__(
        self,
        stroke_width: float = 2.0,
        color: ParsableManimColor = YELLOW,
        render_primitive: int = moderngl.POINTS,
        **kwargs,
    ):
        self.stroke_width = stroke_width
        super().__init__(color=color, render_primitive=render_primitive, **kwargs)
        self.point_radius = (
            self.stroke_width * OpenGLPMobject.OPENGL_POINT_RADIUS_SCALE_FACTOR
        )

    def reset_points(self) -> Self:
        self.rgbas: FloatRGBA_Array = np.zeros((1, 4))
        self.points: Point3D_Array = np.zeros((0, 3))
        return self

    def get_array_attrs(self):
        return ["points", "rgbas"]

    def add_points(
        self,
        points: Point3DLike_Array,
        rgbas: FloatRGBALike_Array | None = None,
        color: ParsableManimColor | None = None,
        opacity: float | None = None,
    ) -> Self:
        """Add points.

        Points must be a Nx3 numpy array.
        Rgbas must be a Nx4 numpy array if it is not None.
        """
        if rgbas is None and color is None:
            color = YELLOW
        self.append_points(points)
        # rgbas array will have been resized with points
        if color is not None:
            if opacity is None:
                opacity = self.rgbas[-1, 3]
            new_rgbas = np.repeat([color_to_rgba(color, opacity)], len(points), axis=0)
        elif rgbas is not None:
            new_rgbas = rgbas
        elif len(rgbas) != len(points):
            raise ValueError("points and rgbas must have same length")
        self.rgbas = np.append(self.rgbas, new_rgbas, axis=0)
        return self

    def thin_out(self, factor=5):
        """Removes all but every nth point for n = factor"""
        for mob in self.family_members_with_points():
            num_points = mob.get_num_points()

            def thin_func(num_points=num_points):
                return np.arange(0, num_points, factor)

            if len(mob.points) == len(mob.rgbas):
                mob.set_rgba_array_direct(mob.rgbas[thin_func()])
            mob.set_points(mob.points[thin_func()])

        return self

    def set_color_by_gradient(self, *colors):
        self.rgbas = np.array(
            list(map(color_to_rgba, color_gradient(*colors, self.get_num_points()))),
        )
        return self

    def set_colors_by_radial_gradient(
        self,
        center=None,
        radius=1,
        inner_color=WHITE,
        outer_color=BLACK,
    ):
        start_rgba, end_rgba = list(map(color_to_rgba, [inner_color, outer_color]))
        if center is None:
            center = self.get_center()
        for mob in self.family_members_with_points():
            distances = np.abs(self.points - center)
            alphas = np.linalg.norm(distances, axis=1) / radius

            mob.rgbas = np.array(
                np.array(
                    [interpolate(start_rgba, end_rgba, alpha) for alpha in alphas],
                ),
            )
        return self

    def match_colors(self, pmobject):
        self.rgbas[:] = resize_with_interpolation(pmobject.rgbas, self.get_num_points())
        return self

    def fade_to(self, color, alpha, family=True):
        rgbas = interpolate(self.rgbas, color_to_rgba(color), alpha)
        for mob in self.submobjects:
            mob.fade_to(color, alpha, family)
        self.set_rgba_array_direct(rgbas)
        return self

    def filter_out(self, condition):
        for mob in self.family_members_with_points():
            to_keep = ~np.apply_along_axis(condition, 1, mob.points)
            for key in mob.data:
                mob.data[key] = mob.data[key][to_keep]
        return self

    def sort_points(self, function=lambda p: p[0]):
        """function is any map from R^3 to R"""
        for mob in self.family_members_with_points():
            indices = np.argsort(np.apply_along_axis(function, 1, mob.points))
            for key in mob.data:
                mob.data[key] = mob.data[key][indices]
        return self

    def ingest_submobjects(self):
        for key in self.data:
            self.data[key] = np.vstack([sm.data[key] for sm in self.get_family()])
        return self

    def point_from_proportion(self, alpha):
        index = alpha * (self.get_num_points() - 1)
        return self.points[int(index)]

    def pointwise_become_partial(self, pmobject, a, b):
        lower_index = int(a * pmobject.get_num_points())
        upper_index = int(b * pmobject.get_num_points())
        for key in self.data:
            self.data[key] = pmobject.data[key][lower_index:upper_index]
        return self

    def get_shader_data(self):
        shader_data = np.zeros(len(self.points), dtype=self.shader_dtype)
        self.read_data_to_shader(shader_data, "point", "points")
        self.read_data_to_shader(shader_data, "color", "rgbas")
        return shader_data

    @staticmethod
    def get_mobject_type_class():
        return OpenGLPMobject


class OpenGLPGroup(OpenGLPMobject):
    def __init__(self, *pmobs, **kwargs):
        if not all(isinstance(m, OpenGLPMobject) for m in pmobs):
            raise Exception("All submobjects must be of type OpenglPMObject")
        super().__init__(**kwargs)
        self.add(*pmobs)

    def fade_to(self, color, alpha, family=True):
        if family:
            for mob in self.submobjects:
                mob.fade_to(color, alpha, family)


class OpenGLPMPoint(OpenGLPMobject):
    def __init__(self, location=ORIGIN, stroke_width=4.0, **kwargs):
        self.location = location
        super().__init__(stroke_width=stroke_width, **kwargs)

    def init_points(self):
        self.points = np.array([self.location], dtype=np.float32)
