from manim import *

class ConvexFunction(Scene):
    def construct(self):
        # Axes
        axes = Axes(
            x_range=[-2, 4, 1],
            y_range=[-0.5, 4, 1],
            axis_config={"include_numbers": False},
        )

        # Labels for axes
        x_label = axes.get_x_axis_label("x", edge=RIGHT, direction=RIGHT, buff=0.3)
        y_label = axes.get_y_axis_label("f(x)", edge=UP, direction=UP, buff=0.3)

        
        # Convex function
        def convex_function(x):
            return 0.5 * (x - 1)**2 + 1

        graph = axes.plot(convex_function, x_range=[-2, 3.5], color=BLUE)
        
        # Initial positions of points a, b, and c
        a_val_start = -1.5
        b_val_start = 3.5
        c_val = 2.0

        # Final positions of points a and b
        a_val_end = 0.5
        b_val_end = 3.0

        # Create ValueTrackers for a and b to smoothly interpolate between their values
        a_tracker = ValueTracker(a_val_start)
        b_tracker = ValueTracker(b_val_start)

        
        # Create the moving dots for a and b, and static dot for c
        a_dot = always_redraw(lambda: Dot(axes.c2p(a_tracker.get_value(), convex_function(a_tracker.get_value())), color=RED))
        b_dot = always_redraw(lambda: Dot(axes.c2p(b_tracker.get_value(), convex_function(b_tracker.get_value())), color=RED))
        c_dot = Dot(axes.c2p(c_val, convex_function(c_val)), color=RED)

        # Labels
        a_label = always_redraw(lambda: MathTex("a").next_to(a_dot, DOWN))
        b_label = always_redraw(lambda: MathTex("b").next_to(b_dot, DOWN))
        c_label = MathTex("c").next_to(c_dot, DOWN)
        
        # Secant lines (chords) that update dynamically as points move
        secant_ac = always_redraw(lambda: Line(
            start=a_dot.get_center(),
            end=c_dot.get_center(),
            color=YELLOW
        ))
        secant_ab = always_redraw(lambda: Line(
            start=a_dot.get_center(),
            end=b_dot.get_center(),
            color=GREEN
        ))
        secant_cb = always_redraw(lambda: Line(
            start=c_dot.get_center(),
            end=b_dot.get_center(),
            color=ORANGE
        ))

        # Inequality text
        inequality_text = MathTex(
            r"f: \mathbb{R} \to \mathbb{R} \text{ convex}, a < c < b, \quad", r"\frac{f(c) - f(a)}{c - a}", r"\leq", r"\frac{f(b) - f(a)}{b - a}", r"\leq", r"\frac{f(b) - f(c)}{b - c}"
        ).scale(0.75).to_corner(DOWN)
        inequality_text.set_color_by_tex_to_color_map({
            r"\frac{f(c) - f(a)}{c - a}": YELLOW,
            r"\frac{f(b) - f(a)}{b - a}": GREEN,
            r"\frac{f(b) - f(c)}{b - c}": ORANGE,
        })

        # Show all objects (axes, graph, points, secants, inequality) at the same time
        self.add(axes, x_label, y_label, graph, inequality_text,
             a_dot, b_dot, c_dot, a_label, b_label, c_label,
             secant_ac, secant_ab, secant_cb)

        # Define the motion of the points along the curve
        def move_along_graph(mob, t):
            mob.move_to(axes.c2p(t, convex_function(t)))

        # Animate the movement of a and b along the graph of the convex function smoothly
        self.play(
            a_tracker.animate.set_value(a_val_end),
            b_tracker.animate.set_value(b_val_end),
            run_time=4,
            rate_func=smooth
        )

        # Hold the final state for a few seconds
        self.wait(2)