from fluxion import *
import numpy as np

class Heatmap3D(ThreeDScene):
    def construct(self):
        # 1. SETUP
        self.set_camera_orientation(phi=75 * DEGREES, theta=-30 * DEGREES)
        
        title = Text("3D Surface Heatmap", font_size=40)
        self.add_fixed_in_frame_mobjects(title)
        title.to_edge(UP)
        
        # 2. CREATE AXES
        axes = ThreeDAxes(
            x_range=[-3, 3, 1],
            y_range=[-3, 3, 1],
            z_range=[-2, 2, 1]
        )
        
        self.play(Create(axes))
        
        # 3. CREATE SURFACE
        # Function: z = sin(x*y) * cos(y)
        # 3D equivalent of a complex heatmap
        
        def param_surface(u, v):
            x = u
            y = v
            z = np.sin(x) * np.cos(y) * 1.5 * np.exp(-(x**2 + y**2)/10)
            return np.array([x, y, z])
            
        resolution_fa = 24
        surface = Surface(
            param_surface,
            resolution=(resolution_fa, resolution_fa),
            v_range=[-3, 3],
            u_range=[-3, 3]
        )
        
        # Style the surface
        surface.set_style(fill_opacity=0.8, stroke_color=BLACK, stroke_width=0.5)
        surface.set_fill_by_checkerboard(BLUE, RED, opacity=0.8) # Fallback texture
        
        # Try to apply gradient if possible, else checkerboard looks cool
        # We can color based on Z-height manually if we want, but let's stick to simple first
        
        self.play(Create(surface), run_time=3)
        self.wait(1)
        
        # 4. CAMERA MOVEMENT
        self.begin_ambient_camera_rotation(rate=0.2)
        
        # 5. DYNAMIC COLOR/SHAPE CHANGE
        # We will transform the surface into a different function
        
        def param_surface_2(u, v):
            x = u
            y = v
            # Ripple effect
            d = np.sqrt(x*x + y*y)
            z = np.sin(d * 3) * 0.5
            return np.array([x, y, z])
            
        surface2 = Surface(
            param_surface_2,
            resolution=(resolution_fa, resolution_fa),
            v_range=[-3, 3],
            u_range=[-3, 3]
        )
        surface2.set_style(fill_opacity=0.8, stroke_color=WHITE, stroke_width=0.2)
        surface2.set_fill_by_checkerboard(TEAL, YELLOW, opacity=0.8)
        
        self.play(
            Transform(surface, surface2),
            run_time=4
        )
        
        self.wait(2)
        
        # 6. FLY AROUND
        self.move_camera(phi=45*DEGREES, theta=45*DEGREES, zoom=1.5, run_time=3)
        self.move_camera(phi=90*DEGREES, zoom=1, run_time=3)
        
        self.stop_ambient_camera_rotation()
        self.play(FadeOut(axes), FadeOut(surface), FadeOut(title))
