from manim import *
import numpy as np

class FullVideo(Scene):
    def construct(self):
        # Part 1: Introduction
        title = Text("Gradients, Optimization, and Bayesian Updating").scale(0.8).to_edge(UP)
        intro_text = Text(
            "Exploring how gradients guide optimization\nand how beliefs evolve with evidence.",
            font_size=24
        ).next_to(title, DOWN)
        self.play(Write(title), Write(intro_text))
        self.wait(2)
        self.play(FadeOut(intro_text))

        # Transition to gradients
        self.play(FadeOut(title))
        self.wait(0.5)

        # Part 2: Understanding Gradients
        axes = Axes(
            x_range=[0, 10, 1],
            y_range=[0, 25, 5],
            axis_config={"include_numbers": True},
        )
        graph = axes.plot(lambda x: (x - 5)**2, color=BLUE)
        func_label = MathTex("f(x) = (x - 5)^2").next_to(axes, UP)
        self.play(Create(axes), Write(func_label))
        self.play(Create(graph))
        self.wait(1)

        # Gradient descent animation
        dot_min = Dot(axes.coords_to_point(7, (7 - 5)**2), color=RED)
        self.play(FadeIn(dot_min))
        for _ in range(5):
            new_x = dot_min.get_center()[0] - 0.5
            new_y = (new_x - 5)**2
            new_dot = Dot(axes.coords_to_point(new_x, new_y), color=RED)
            self.play(Transform(dot_min, new_dot), run_time=0.5)
        self.wait(1)

        # Zoom effect on the minimum point
        self.play(
            axes.animate.scale(0.8).shift(LEFT * 2),
            dot_min.animate.scale(1.5),
            run_time=1.5
        )
        self.wait(1)
        self.play(FadeOut(axes), FadeOut(dot_min), FadeOut(func_label))

        # Part 3: Comfort Score Function
        axes_3d = ThreeDAxes(
            x_range=[60, 80, 5],
            y_range=[30, 50, 5],
            z_range=[0, 100, 20],
            x_length=8,
            y_length=8,
            z_length=6
        )
        def comfort_score(t, h):
            return 72 - (t - 70)**2 - 2 * (h - 40)**2

        surface = Surface(
            lambda u, v: axes_3d.c2p(u, v, comfort_score(u, v)),
            u_range=[60, 80],
            v_range=[30, 50],
            resolution=(20, 20),
            fill_opacity=0.7
        )
        surface.set_fill_by_value(
            axes=axes_3d,
            colors=[(RED, 0), (YELLOW, 50), (GREEN, 100)]
        )

        # Set camera orientation and animate rotation
        self.set_camera_orientation(phi=75 * DEGREES, theta=-45 * DEGREES)
        self.add(axes_3d, surface)
        self.begin_ambient_camera_rotation(rate=0.2)
        self.wait(5)
        self.stop_ambient_camera_rotation()

        # Zoom into the peak of the surface
        self.move_camera(phi=90 * DEGREES, theta=-90 * DEGREES, zoom=1.5, run_time=2)
        self.wait(1)
        self.play(FadeOut(axes_3d), FadeOut(surface))

        # Part 4: Bayesian Updating
        prior = [0.3, 0.7]
        bar_chart = BarChart(
            prior,
            max_value=1,
            bar_names=["Rain", "No Rain"],
            bar_colors=[BLUE, YELLOW]
        )
        prior_label = Text("Prior Probabilities").next_to(bar_chart, DOWN)
        self.play(Create(bar_chart), Write(prior_label))
        self.wait(2)

        posterior = [0.6, 0.4]
        updated_bar_chart = BarChart(
            posterior,
            max_value=1,
            bar_names=["Rain", "No Rain"],
            bar_colors=[BLUE, YELLOW]
        )
        posterior_label = Text("Posterior Probabilities").next_to(updated_bar_chart, DOWN)
        self.play(Transform(bar_chart, updated_bar_chart), Transform(prior_label, posterior_label))
        self.wait(2)

        bayes_formula = MathTex(r"P(H|E) = \frac{P(E|H)P(H)}{P(E)}").next_to(bar_chart, DOWN)
        self.play(Write(bayes_formula))
        self.wait(2)

        # Zoom out for conclusion
        self.play(
            bar_chart.animate.scale(0.8).shift(LEFT * 2),
            bayes_formula.animate.scale(0.8).shift(RIGHT * 2),
            run_time=1.5
        )
        self.wait(1)

        # Part 5: Connecting the Dots
        self.clear()
        conclusion = Text(
            "Gradients and Bayesian updating both rely\non iterative refinement to achieve their goals.",
            font_size=24
        ).to_edge(UP)
        self.play(Write(conclusion))
        self.wait(3)

        # Part 6: Call to Action
        call_to_action = Text(
            "Explore more about optimization and Bayesian methods!",
            font_size=24
        ).next_to(conclusion, DOWN)
        self.play(Write(call_to_action))
        self.wait(3)