from fluxion import *
import numpy as np

class MindBender(ThreeDScene):
    def construct(self):
        # 0. CONFIG 
        self.camera.background_color = "#000000"
        
        # ---------------------------------------------------------
        # SEGMENT 1: HYPNOTIC INTERFERENCE
        # ---------------------------------------------------------
        title = Text("FLUXION", font_size=96, weight=BOLD).set_color_by_gradient(RED, ORANGE, YELLOW, GREEN, BLUE, PURPLE)
        self.play(Write(title), run_time=3)
        self.play(title.animate.scale(3).set_opacity(0), run_time=2)
        
        # Circles
        circles = VGroup(*[
            Circle(radius=i*0.2, stroke_width=2).set_stroke(opacity=0.5)
            for i in range(20)
        ])
        circles.set_color_by_gradient(TEAL, PURPLE) # Replaced CYAN with TEAL
        
        circles2 = circles.copy()
        circles2.set_color_by_gradient(YELLOW, PINK)
        
        self.add(circles, circles2)
        
        def circle_updater(m, dt):
            m.rotate(dt * 0.5)
            new_width = 2 + np.sin(self.renderer.time * 5)
            m.set_stroke(width=new_width)
            
        def circle_updater2(m, dt):
            m.rotate(-dt * 0.7)
            m.shift(RIGHT * 0.5 * np.sin(self.renderer.time) * dt)

        circles.add_updater(circle_updater)
        circles2.add_updater(circle_updater2)
        
        self.wait(10)
        
        circles.remove_updater(circle_updater)
        circles2.remove_updater(circle_updater2)
        
        self.play(
            FadeOut(circles), FadeOut(circles2),
            run_time=2
        )

        # ---------------------------------------------------------
        # SEGMENT 2: 3D MATRIX FLIGHT
        # ---------------------------------------------------------
        self.set_camera_orientation(phi=75 * DEGREES, theta=30 * DEGREES)
        
        # Manually creating cubes since Cube might be missing
        # A simple 3D dot cloud is safer and looks cool too
        dots = VGroup()
        for x in range(-2, 3):
            for y in range(-2, 3):
                for z in range(-2, 3):
                    dot = Dot3D(point=[x, y, z], radius=0.1, color=BLUE)
                    dot.set_opacity(0.5)
                    dots.add(dot)
        
        self.play(Create(dots), run_time=4)
        
        self.begin_ambient_camera_rotation(rate=0.2)
        
        self.play(
            dots.animate.scale(1.5).rotate(PI/4, axis=UP),
            run_time=5
        )
        
        # Wave animation
        self.play(
            dots.animate.apply_function(
                lambda p: p + np.array([0, 0, 0.5 * np.sin(p[0] + p[1])])
            ),
            run_time=5
        )
        
        self.move_camera(phi=90*DEGREES, theta=-90*DEGREES, zoom=2, run_time=5)
        self.move_camera(phi=45*DEGREES, theta=45*DEGREES, zoom=0.5, run_time=5)
        
        self.stop_ambient_camera_rotation()
        self.play(FadeOut(dots), run_time=2)
        
        self.set_camera_orientation(phi=0, theta=-90*DEGREES)

        # ---------------------------------------------------------
        # SEGMENT 3: RECURSIVE FRACTAL CHAOS
        # ---------------------------------------------------------
        
        def get_fractal_layer(radius, depth, max_depth):
            if depth == max_depth:
                # Use Dot instead of Circle for filled look and safety
                return Dot(radius=radius, color=random_bright_color()).set_opacity(0.5)
            
            group = VGroup()
            for i in range(3):
                angle = i * (2 * PI / 3)
                sub = get_fractal_layer(radius/2, depth+1, max_depth)
                sub.shift(radius/2 * np.array([np.cos(angle), np.sin(angle), 0]))
                group.add(sub)
            return group

        colors = [RED, ORANGE, YELLOW, GREEN, BLUE, PURPLE]
        def random_bright_color():
            return colors[np.random.randint(0, len(colors))]

        fractal = get_fractal_layer(3, 0, 4)
        
        self.play(GrowFromCenter(fractal), run_time=4) # Replaced SpinInFromNothing
        
        self.play(
            Rotate(fractal, angle=2*PI, run_time=10, rate_func=linear),
            fractal.animate.set_color_by_gradient(TEAL, PINK, GOLD),
            run_time=10
        )
        
        def kaleidoscope_update(m, dt):
            m.rotate(dt)
            m.scale(1 + 0.05 * np.sin(self.renderer.time * 2))
            
        fractal.add_updater(kaleidoscope_update)
        self.wait(10)
        fractal.remove_updater(kaleidoscope_update)
        
        self.play(ShrinkToCenter(fractal), run_time=2)

        # ---------------------------------------------------------
        # SEGMENT 4: POLAR GRAPH MORPHING
        # ---------------------------------------------------------
        
        # Use simple lines for axes
        axes = VGroup(
            Line(LEFT*4, RIGHT*4),
            Line(DOWN*3, UP*3)
        ).set_color(GRAY)
        self.play(Create(axes))
        
        # Parametric functions
        def rose_curve(theta):
            k = 4
            r = 3 * np.cos(k * theta)
            return [r * np.cos(theta), r * np.sin(theta), 0]
            
        rose = ParametricFunction(rose_curve, t_range=[0, 2*PI], color=RED)
        
        self.play(Create(rose), run_time=4)
        
        def butterfly(theta):
            r = np.exp(np.sin(theta)) - 2*np.cos(4*theta) + np.sin((2*theta - PI)/24)**5
            return [2 * r * np.cos(theta), 2 * r * np.sin(theta), 0]
            
        butt = ParametricFunction(butterfly, t_range=[0, 10*PI], color=BLUE)
        
        # Use ReplacementTransform instead of direct Transform for parametric functions
        # to match structure nicely
        self.play(ReplacementTransform(rose, butt), run_time=6)
        
        self.play(Rotate(butt, angle=2*PI), run_time=5, rate_func=smooth)
        
        self.play(FadeOut(butt), FadeOut(axes), run_time=2)

        # ---------------------------------------------------------
        # SEGMENT 5: FINAL TEXTURE
        # ---------------------------------------------------------
        
        final_text = Text("Thanks for watching", font_size=40)
        
        particles = VGroup(*[
            Dot(radius=0.05, color=random_bright_color())
            for _ in range(100)
        ])
        particles.arrange_in_grid()
        
        self.play(Create(particles))
        self.play(
            particles.animate.scale(5).arrange_in_grid(buff=0.5).rotate(PI),
            run_time=3
        )
        self.play(ReplacementTransform(particles, final_text))
        self.wait(2)
        self.play(FadeOut(final_text))
