"""Mobjects representing objects from probability theory and statistics."""

from __future__ import annotations

__all__ = ["SampleSpace", "BarChart"]


from collections.abc import Iterable, MutableSequence, Sequence
from typing import Any

import numpy as np

from manim import config, logger
from manim.constants import *
from manim.mobject.geometry.polygram import Rectangle
from manim.mobject.graphing.coordinate_systems import Axes
from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject
from manim.mobject.svg.brace import Brace
from manim.mobject.text.tex_mobject import MathTex, Tex
from manim.mobject.types.vectorized_mobject import VGroup, VMobject
from manim.typing import Vector3D
from manim.utils.color import (
    BLUE_E,
    DARK_GREY,
    GREEN_E,
    LIGHT_GREY,
    MAROON_B,
    YELLOW,
    ParsableManimColor,
    color_gradient,
)
from manim.utils.iterables import tuplify

EPSILON = 0.0001


class SampleSpace(Rectangle):
    """A mobject representing a twodimensional rectangular
    sampling space.

    Examples
    --------
    .. manim:: ExampleSampleSpace
        :save_last_frame:

        class ExampleSampleSpace(Scene):
            def construct(self):
                poly1 = SampleSpace(stroke_width=15, fill_opacity=1)
                poly2 = SampleSpace(width=5, height=3, stroke_width=5, fill_opacity=0.5)
                poly3 = SampleSpace(width=2, height=2, stroke_width=5, fill_opacity=0.1)
                poly3.divide_vertically(p_list=np.array([0.37, 0.13, 0.5]), colors=[BLACK, WHITE, GRAY], vect=RIGHT)
                poly_group = VGroup(poly1, poly2, poly3).arrange()
                self.add(poly_group)
    """

    def __init__(
        self,
        height: float = 3,
        width: float = 3,
        fill_color: ParsableManimColor = DARK_GREY,
        fill_opacity: float = 1,
        stroke_width: float = 0.5,
        stroke_color: ParsableManimColor = LIGHT_GREY,
        default_label_scale_val: float = 1,
    ):
        super().__init__(
            height=height,
            width=width,
            fill_color=fill_color,
            fill_opacity=fill_opacity,
            stroke_width=stroke_width,
            stroke_color=stroke_color,
        )
        self.default_label_scale_val = default_label_scale_val

    def add_title(
        self, title: str = "Sample space", buff: float = MED_SMALL_BUFF
    ) -> None:
        # TODO, should this really exist in SampleSpaceScene
        title_mob = Tex(title)
        if title_mob.width > self.width:
            title_mob.width = self.width
        title_mob.next_to(self, UP, buff=buff)
        self.title = title_mob
        self.add(title_mob)

    def add_label(self, label: str) -> None:
        self.label = label

    def complete_p_list(self, p_list: float | Iterable[float]) -> list[float]:
        p_list_tuplified: tuple[float] = tuplify(p_list)
        new_p_list = list(p_list_tuplified)
        remainder = 1.0 - sum(new_p_list)
        if abs(remainder) > EPSILON:
            new_p_list.append(remainder)
        return new_p_list

    def get_division_along_dimension(
        self,
        p_list: float | Iterable[float],
        dim: int,
        colors: Sequence[ParsableManimColor],
        vect: Vector3D,
    ) -> VGroup:
        p_list_complete = self.complete_p_list(p_list)
        colors_in_gradient = color_gradient(colors, len(p_list_complete))

        last_point = self.get_edge_center(-vect)
        parts = VGroup()
        for factor, color in zip(p_list_complete, colors_in_gradient, strict=False):
            part = SampleSpace()
            part.set_fill(color, 1)
            part.replace(self, stretch=True)
            part.stretch(factor, dim)
            part.move_to(last_point, -vect)
            last_point = part.get_edge_center(vect)
            parts.add(part)
        return parts

    def get_horizontal_division(
        self,
        p_list: float | Iterable[float],
        colors: Sequence[ParsableManimColor] = [GREEN_E, BLUE_E],
        vect: Vector3D = DOWN,
    ) -> VGroup:
        return self.get_division_along_dimension(p_list, 1, colors, vect)

    def get_vertical_division(
        self,
        p_list: float | Iterable[float],
        colors: Sequence[ParsableManimColor] = [MAROON_B, YELLOW],
        vect: Vector3D = RIGHT,
    ) -> VGroup:
        return self.get_division_along_dimension(p_list, 0, colors, vect)

    def divide_horizontally(self, *args: Any, **kwargs: Any) -> None:
        self.horizontal_parts = self.get_horizontal_division(*args, **kwargs)
        self.add(self.horizontal_parts)

    def divide_vertically(self, *args: Any, **kwargs: Any) -> None:
        self.vertical_parts = self.get_vertical_division(*args, **kwargs)
        self.add(self.vertical_parts)

    def get_subdivision_braces_and_labels(
        self,
        parts: VGroup,
        labels: list[str | VMobject | OpenGLVMobject],
        direction: Vector3D,
        buff: float = SMALL_BUFF,
        min_num_quads: int = 1,
    ) -> VGroup:
        label_mobs = VGroup()
        braces = VGroup()
        for label, part in zip(labels, parts, strict=False):
            brace = Brace(part, direction, min_num_quads=min_num_quads, buff=buff)
            if isinstance(label, (VMobject, OpenGLVMobject)):
                label_mob = label
            else:
                label_mob = MathTex(label)
                label_mob.scale(self.default_label_scale_val)
            label_mob.next_to(brace, direction, buff)

            braces.add(brace)
            assert isinstance(label_mob, VMobject)
            label_mobs.add(label_mob)
        parts.braces = braces  # type: ignore[attr-defined]
        parts.labels = label_mobs  # type: ignore[attr-defined]
        parts.label_kwargs = {  # type: ignore[attr-defined]
            "labels": label_mobs.copy(),
            "direction": direction,
            "buff": buff,
        }
        return VGroup(parts.braces, parts.labels)

    def get_side_braces_and_labels(
        self,
        labels: list[str | VMobject | OpenGLVMobject],
        direction: Vector3D = LEFT,
        **kwargs: Any,
    ) -> VGroup:
        assert hasattr(self, "horizontal_parts")
        parts = self.horizontal_parts
        return self.get_subdivision_braces_and_labels(
            parts, labels, direction, **kwargs
        )

    def get_top_braces_and_labels(
        self, labels: list[str | VMobject | OpenGLVMobject], **kwargs: Any
    ) -> VGroup:
        assert hasattr(self, "vertical_parts")
        parts = self.vertical_parts
        return self.get_subdivision_braces_and_labels(parts, labels, UP, **kwargs)

    def get_bottom_braces_and_labels(
        self, labels: list[str | VMobject | OpenGLVMobject], **kwargs: Any
    ) -> VGroup:
        assert hasattr(self, "vertical_parts")
        parts = self.vertical_parts
        return self.get_subdivision_braces_and_labels(parts, labels, DOWN, **kwargs)

    def add_braces_and_labels(self) -> None:
        for attr in "horizontal_parts", "vertical_parts":
            if not hasattr(self, attr):
                continue
            parts = getattr(self, attr)
            for subattr in "braces", "labels":
                if hasattr(parts, subattr):
                    self.add(getattr(parts, subattr))

    def __getitem__(self, index: int) -> SampleSpace:
        if hasattr(self, "horizontal_parts"):
            val: SampleSpace = self.horizontal_parts[index]
            return val
        elif hasattr(self, "vertical_parts"):
            val = self.vertical_parts[index]
            return val
        return self.split()[index]


class BarChart(Axes):
    """Creates a bar chart. Inherits from :class:`~.Axes`, so it shares its methods
    and attributes. Each axis inherits from :class:`~.NumberLine`, so pass in ``x_axis_config``/``y_axis_config``
    to control their attributes.

    Parameters
    ----------
    values
        A sequence of values that determines the height of each bar. Accepts negative values.
    bar_names
        A sequence of names for each bar. Does not have to match the length of ``values``.
    y_range
        The y_axis range of values. If ``None``, the range will be calculated based on the
        min/max of ``values`` and the step will be calculated based on ``y_length``.
    x_length
        The length of the x-axis. If ``None``, it is automatically calculated based on
        the number of values and the width of the screen.
    y_length
        The length of the y-axis.
    bar_colors
        The color for the bars. Accepts a sequence of colors (can contain just one item).
        If the length of``bar_colors`` does not match that of ``values``,
        intermediate colors will be automatically determined.
    bar_width
        The length of a bar. Must be between 0 and 1.
    bar_fill_opacity
        The fill opacity of the bars.
    bar_stroke_width
        The stroke width of the bars.

    Examples
    --------
    .. manim:: BarChartExample
        :save_last_frame:

        class BarChartExample(Scene):
            def construct(self):
                chart = BarChart(
                    values=[-5, 40, -10, 20, -3],
                    bar_names=["one", "two", "three", "four", "five"],
                    y_range=[-20, 50, 10],
                    y_length=6,
                    x_length=10,
                    x_axis_config={"font_size": 36},
                )

                c_bar_lbls = chart.get_bar_labels(font_size=48)

                self.add(chart, c_bar_lbls)
    """

    def __init__(
        self,
        values: MutableSequence[float],
        bar_names: Sequence[str] | None = None,
        y_range: Sequence[float] | None = None,
        x_length: float | None = None,
        y_length: float | None = None,
        bar_colors: Iterable[str] = [
            "#003f5c",
            "#58508d",
            "#bc5090",
            "#ff6361",
            "#ffa600",
        ],
        bar_width: float = 0.6,
        bar_fill_opacity: float = 0.7,
        bar_stroke_width: float = 3,
        **kwargs: Any,
    ):
        if isinstance(bar_colors, str):
            logger.warning(
                "Passing a string to `bar_colors` has been deprecated since v0.15.2 and will be removed after v0.17.0, the parameter must be a list.  "
            )
            bar_colors = list(bar_colors)

        y_length = y_length if y_length is not None else config.frame_height - 4

        self.values = values
        self.bar_names = bar_names
        self.bar_colors = bar_colors
        self.bar_width = bar_width
        self.bar_fill_opacity = bar_fill_opacity
        self.bar_stroke_width = bar_stroke_width

        x_range = [0, len(self.values), 1]

        if y_range is None:
            y_range = [
                min(0, min(self.values)),
                max(0, max(self.values)),
                round(max(self.values) / y_length, 2),
            ]

        elif len(y_range) == 2:
            y_range = [*y_range, round(max(self.values) / y_length, 2)]

        if x_length is None:
            x_length = min(len(self.values), config.frame_width - 2)

        x_axis_config = {"font_size": 24, "label_constructor": Tex}
        self._update_default_configs(
            (x_axis_config,), (kwargs.pop("x_axis_config", None),)
        )

        self.bars: VGroup = VGroup()
        self.x_labels: VGroup | None = None
        self.bar_labels: VGroup | None = None

        super().__init__(
            x_range=x_range,
            y_range=y_range,
            x_length=x_length,
            y_length=y_length,
            x_axis_config=x_axis_config,
            tips=kwargs.pop("tips", False),
            **kwargs,
        )

        self._add_bars()

        if self.bar_names is not None:
            self._add_x_axis_labels()

        self.y_axis.add_numbers()

    def _update_colors(self) -> None:
        """Initialize the colors of the bars of the chart.

        Sets the color of ``self.bars`` via ``self.bar_colors``.

        Primarily used when the bars are initialized with ``self._add_bars``
        or updated via ``self.change_bar_values``.
        """
        self.bars.set_color_by_gradient(*self.bar_colors)

    def _add_x_axis_labels(self) -> None:
        """Essentially :meth`:~.NumberLine.add_labels`, but differs in that
        the direction of the label with respect to the x_axis changes to UP or DOWN
        depending on the value.

        UP for negative values and DOWN for positive values.
        """
        assert isinstance(self.bar_names, list)
        val_range = np.arange(
            0.5, len(self.bar_names), 1
        )  # 0.5 shifted so that labels are centered, not on ticks

        labels = VGroup()

        for i, (value, bar_name) in enumerate(
            zip(val_range, self.bar_names, strict=False)
        ):
            # to accommodate negative bars, the label may need to be
            # below or above the x_axis depending on the value of the bar
            direction = UP if self.values[i] < 0 else DOWN
            bar_name_label: MathTex = self.x_axis.label_constructor(bar_name)

            bar_name_label.font_size = self.x_axis.font_size
            bar_name_label.next_to(
                self.x_axis.number_to_point(value),
                direction=direction,
                buff=self.x_axis.line_to_number_buff,
            )

            labels.add(bar_name_label)

        self.x_axis.labels = labels
        self.x_axis.add(labels)

    def _create_bar(self, bar_number: int, value: float) -> Rectangle:
        """Creates a positioned bar on the chart.

        Parameters
        ----------
        bar_number
            Determines the x-position of the bar.
        value
            The value that determines the height of the bar.

        Returns
        -------
        Rectangle
            A positioned rectangle representing a bar on the chart.
        """
        # bar measurements relative to the axis

        # distance from between the y-axis and the top of the bar
        bar_h = abs(self.c2p(0, value)[1] - self.c2p(0, 0)[1])
        # width of the bar
        bar_w = self.c2p(self.bar_width, 0)[0] - self.c2p(0, 0)[0]

        bar = Rectangle(
            height=bar_h,
            width=bar_w,
            stroke_width=self.bar_stroke_width,
            fill_opacity=self.bar_fill_opacity,
        )

        pos = UP if (value >= 0) else DOWN
        bar.next_to(self.c2p(bar_number + 0.5, 0), pos, buff=0)
        return bar

    def _add_bars(self) -> None:
        for i, value in enumerate(self.values):
            tmp_bar = self._create_bar(bar_number=i, value=value)
            self.bars.add(tmp_bar)

        self._update_colors()
        self.add_to_back(self.bars)

    def get_bar_labels(
        self,
        color: ParsableManimColor | None = None,
        font_size: float = 24,
        buff: float = MED_SMALL_BUFF,
        label_constructor: type[MathTex] = Tex,
    ) -> VGroup:
        """Annotates each bar with its corresponding value. Use ``self.bar_labels`` to access the
        labels after creation.

        Parameters
        ----------
        color
            The color of each label. By default ``None`` and is based on the parent's bar color.
        font_size
            The font size of each label.
        buff
            The distance from each label to its bar. By default 0.4.
        label_constructor
            The Mobject class to construct the labels, by default :class:`~.Tex`.

        Examples
        --------
        .. manim:: GetBarLabelsExample
            :save_last_frame:

            class GetBarLabelsExample(Scene):
                def construct(self):
                    chart = BarChart(values=[10, 9, 8, 7, 6, 5, 4, 3, 2, 1], y_range=[0, 10, 1])

                    c_bar_lbls = chart.get_bar_labels(
                        color=WHITE, label_constructor=MathTex, font_size=36
                    )

                    self.add(chart, c_bar_lbls)
        """
        bar_labels = VGroup()
        for bar, value in zip(self.bars, self.values, strict=False):
            bar_lbl: MathTex = label_constructor(str(value))

            if color is None:
                bar_lbl.set_color(bar.get_fill_color())
            else:
                bar_lbl.set_color(color)

            bar_lbl.font_size = font_size

            pos = UP if (value >= 0) else DOWN
            bar_lbl.next_to(bar, pos, buff=buff)
            bar_labels.add(bar_lbl)

        return bar_labels

    def change_bar_values(
        self, values: Iterable[float], update_colors: bool = True
    ) -> None:
        """Updates the height of the bars of the chart.

        Parameters
        ----------
        values
            The values that will be used to update the height of the bars.
            Does not have to match the number of bars.
        update_colors
            Whether to re-initalize the colors of the bars based on ``self.bar_colors``.

        Examples
        --------
        .. manim:: ChangeBarValuesExample
            :save_last_frame:

            class ChangeBarValuesExample(Scene):
                def construct(self):
                    values=[-10, -8, -6, -4, -2, 0, 2, 4, 6, 8, 10]

                    chart = BarChart(
                        values,
                        y_range=[-10, 10, 2],
                        y_axis_config={"font_size": 24},
                    )
                    self.add(chart)

                    chart.change_bar_values(list(reversed(values)))
                    self.add(chart.get_bar_labels(font_size=24))
        """
        for i, (bar, value) in enumerate(zip(self.bars, values, strict=False)):
            chart_val = self.values[i]

            if chart_val > 0:
                bar_lim = bar.get_bottom()
                aligned_edge = DOWN
            else:
                bar_lim = bar.get_top()
                aligned_edge = UP

            # check if the bar has height
            if chart_val != 0:
                quotient = value / chart_val
                if quotient < 0:
                    aligned_edge = UP if chart_val > 0 else DOWN

                    # if the bar is already positive, then we now want to move it
                    # so that it is negative. So, we move the top edge of the bar
                    # to the location of the previous bottom

                    # if already negative, then we move the bottom edge of the bar
                    # to the location of the previous top

                bar.stretch_to_fit_height(abs(quotient) * bar.height)

            else:
                # create a new bar since the current one has a height of zero (doesn't exist)
                temp_bar = self._create_bar(i, value)
                self.bars.remove(bar)
                self.bars.insert(i, temp_bar)

            bar.move_to(bar_lim, aligned_edge)

        if update_colors:
            self._update_colors()

        self.values[: len(list(values))] = values
