from fluxion import *
import numpy as np

class HeatmapShowcase(Scene):
    def construct(self):
        # Setup
        self.camera.background_color = "#111111"
        
        title = Text("Heatmap Visualization", font_size=40).to_edge(UP)
        self.play(Write(title))

        # 1. CREATE GRID
        # Create a grid of squares to represent heatmap cells
        rows, cols = 10, 16
        cell_size = 0.5
        
        grid = VGroup()
        for i in range(rows):
            for j in range(cols):
                square = Square(side_length=cell_size)
                square.set_stroke(BLACK, width=1)
                square.move_to(
                    (j - cols/2 + 0.5) * cell_size * RIGHT + 
                    (i - rows/2 + 0.5) * cell_size * UP
                )
                grid.add(square)
        
        self.play(Create(grid), run_time=2)
        
        # 2. INITIALIZE HEATMAP
        # Function to generate heat values
        def get_heat_values(t):
            values = []
            center_x = 3 * np.sin(t)
            center_y = 2 * np.cos(t * 1.5)
            
            for i in range(rows):
                for j in range(cols):
                    # Coordinates relative to center
                    x = (j - cols/2 + 0.5) * cell_size
                    y = (i - rows/2 + 0.5) * cell_size
                    
                    # Distance from moving heat source
                    dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                    
                    # Heat formula (gaussian-like)
                    heat = np.exp(-dist**2 / 4)
                    values.append(heat)
            return values

        # Initial color set
        initial_heat = get_heat_values(0)
        
        # Color mapping helper (simple gradient interpolation)
        colors = [BLUE, TEAL, GREEN, YELLOW, RED]
        def get_color_from_value(val):
            # val is 0 to 1
            idx = val * (len(colors) - 1)
            idx_floor = int(idx)
            idx_ceil = min(idx_floor + 1, len(colors) - 1)
            alpha = idx - idx_floor
            return interpolate_color(colors[idx_floor], colors[idx_ceil], alpha)
            
        for square, val in zip(grid, initial_heat):
            square.set_fill(get_color_from_value(val), opacity=1)
            
        self.play(grid.animate.set_opacity(1), run_time=1)
        
        # 3. DYNAMIC UPDATE
        self.play(
            Transform(title, Text("Dynamic Diffusion Simulation", font_size=40).to_edge(UP))
        )
        
        # Add updater to simulate changing heat
        def update_heatmap(mob, dt):
            mob.time_tracker += dt
            current_vals = get_heat_values(mob.time_tracker)
            
            for square, val in zip(mob, current_vals):
                square.set_fill(get_color_from_value(val), opacity=1)
                
        grid.time_tracker = 0
        grid.add_updater(update_heatmap)
        
        self.wait(8)
        
        grid.remove_updater(update_heatmap)
        
        # 4. TRANSITION TO 3D BARS (Pseudo-3D effect)
        self.play(
            Transform(title, Text("3D Bar Chart View", font_size=40).to_edge(UP))
        )
        
        # Since this is a 2D scene, we simulate 3D via scaling and opacity or just scale
        # Scale squares based on their last heat value
        # (simulating height map)
        
        last_vals = get_heat_values(grid.time_tracker)
        anims = []
        for square, val in zip(grid, last_vals):
            # Scale factor based on heat
            scale_factor = 0.2 + val * 0.8
            anims.append(square.animate.scale(scale_factor))
            
        self.play(*anims, run_time=2)
        self.wait(2)
        
        # Finale
        self.play(FadeOut(grid), FadeOut(title))
