import itertools
from typing import Iterable, Sequence

from manim import ManimColor as Color
from manim import *


class ChartBars(VGroup):
    def __init__(
        self,
        axes,
        values: Sequence[float | int],
        xs: Sequence[float] | None = None,
        width_ratio: float = 1.0,
        offset: float = 0.5,
        fill_color: Color = BLUE,
        fill_opacity: float = 0.5,
        stroke_color: Color = WHITE,
        stroke_width: float = 1.0,
    ):
        """Create a bar chart from values.

        Args:
            axes: The axes to plot the bars on.
            values: The values to plot. There should be one value for each bar.
            xs: The x-coordinates of the bars. If None, the bars will be centered on
                each integer from 1 to the number of values.
            width_ratio: The width of each bar as a fraction of the distance between x-coordinates.
            offset: The distance of the bars from the x-coordinates as a fraction of the distance between x-coordinates.
            fill_color: The fill color of the bars.
            fill_opacity: The fill opacity of the bars.
            stroke_color: The stroke color of the bars.
            stroke_width: The stroke width of the bars.
        """
        xs = xs if xs is not None else np.arange(*axes.x_range)

        self.x_to_index = dict(zip(xs, itertools.count()))
        x_step = xs[1] - xs[0]
        x_unit = axes.x_axis.get_unit_size()
        y_unit = axes.y_axis.get_unit_size()

        width = width_ratio * x_unit * x_step

        # Create a list of rectangles arranged side by side,
        # one for each x value
        rects = []
        epsilon = 1e-8
        for x, y in zip(xs, values):
            rect = Rectangle(
                width=width,
                height=max(y * y_unit, epsilon),
                fill_color=fill_color,
                fill_opacity=fill_opacity,
                stroke_color=stroke_color,
                stroke_width=stroke_width,
            )
            rect.move_to(axes.c2p(x + offset * x_step, 0), DOWN)
            rects.append(rect)
        super().__init__(*rects)
        self.axes = axes
        self.xs = xs
        self.set_values(values)

    def set_values(self, values: Iterable[float | int]):
        y_unit = self.axes.y_axis.get_unit_size()
        for rect, value in zip(self, values):
            new_height = value * y_unit
            height = rect.get_top()[1] - rect.get_bottom()[1]
            rescale = new_height / height
            rect.stretch(
                rescale,
                dim=1,
                about_edge=DOWN,
            )
