from manim import *
import numpy as np


def color_to_numpy(color):
    color = color.lstrip("#")
    return np.array([int(color[i : i + 2], 16) for i in (0, 2, 4)])


def numpy_to_color(arr):
    r, g, b = arr.round().astype(int)
    return "#%02x%02x%02x" % (r, g, b)


class ScaleExpDistribution(Scene):
    def get_rectangle_corners(self, bottom_left, top_right):
        return [
            (top_right[0], top_right[1]),
            (bottom_left[0], top_right[1]),
            (bottom_left[0], bottom_left[1]),
            (top_right[0], bottom_left[1]),
        ]

    def construct(self):
        FONT_SIZE = 96
        START_VALUE = 1
        END_VALUE = 2
        START_COLOR = color_to_numpy(YELLOW)
        END_COLOR = color_to_numpy(RED)

        t = ValueTracker(START_VALUE)

        axes = Axes(x_range=(0, 3.2, 0.5), y_range=(0, 1.1, 0.1))
        axes.add_coordinates()
        x_label = axes.get_x_axis_label("x")
        x_label.shift(0.7 * DOWN)
        y_label = axes.get_y_axis_label("f_\eta (x)")
        y_label.shift(0.8 * LEFT)

        def get_graph():
            t_value = t.get_value()
            alpha = (t_value - START_VALUE) / (END_VALUE - START_VALUE)
            color = numpy_to_color(START_COLOR + (alpha * (END_COLOR - START_COLOR)))
            xi_text = MathTex(r"\xi \sim Exp(1)", font_size=FONT_SIZE)

            number = DecimalNumber().set_color(color).set_value(t_value)
            number.font_size = FONT_SIZE
            eta_text = VGroup(
                MathTex(r"\eta = ", font_size=FONT_SIZE),
                number,
                MathTex(r"\xi", font_size=FONT_SIZE),
            ).arrange()

            dist_text = VGroup(xi_text, eta_text).arrange(DOWN)
            dist_text.to_corner(UP + RIGHT)

            rate = 1 / t_value
            exp_pdf = lambda x: rate * np.exp(-rate * x)
            graph = axes.plot(exp_pdf, x_range=(0, 3.1), color=color)

            starting_points = [0.25, 0.5, 0.75, 1, 1.25, 1.5]
            polygons = []
            for starting_point in starting_points:
                x_value = starting_point * (t_value / START_VALUE)
                polygon = Polygon(
                    *[
                        axes.c2p(*i)
                        for i in self.get_rectangle_corners(
                            (0, 0), (x_value, exp_pdf(x_value))
                        )
                    ]
                )
                polygon.stroke_width = 1
                polygon.set_stroke(WHITE)
                polygons.append(polygon)

            return VGroup(graph, *polygons, dist_text)

        graph = always_redraw(get_graph)
        self.add(axes, graph, x_label, y_label)

        self.wait()
        self.play(t.animate.set_value(END_VALUE), run_time=3, rate_func=linear)
        self.wait()