import pygame
import math
import numpy as np
import random
import cv2
import os

# 初始化Pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 1200, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.DOUBLEBUF | pygame.HWSURFACE)
pygame.display.set_caption("🚀 究极炫酷3D立方体宇宙 🌌")
clock = pygame.time.Clock()

# 颜色和渐变
class ColorPalette:
    def __init__(self):
        self.palettes = [
            [(255, 50, 50), (50, 255, 50), (50, 50, 255)],  # 经典RGB
            [(255, 100, 0), (255, 200, 0), (200, 0, 255)],  # 霓虹
            [(0, 255, 255), (255, 0, 255), (255, 255, 0)],  # 赛博朋克
            [(100, 0, 255), (0, 255, 100), (255, 0, 100)],  # 科技感
        ]
        self.current_palette = random.choice(self.palettes)
        self.transition_speed = 0.02
        self.color_index = 0
        self.next_color_index = (self.color_index + 1) % len(self.current_palette)
        self.color_progress = 0

    def get_interpolated_color(self):
        start_color = self.current_palette[self.color_index]
        end_color = self.current_palette[self.next_color_index]
        
        r = start_color[0] + (end_color[0] - start_color[0]) * self.color_progress
        g = start_color[1] + (end_color[1] - start_color[1]) * self.color_progress
        b = start_color[2] + (end_color[2] - start_color[2]) * self.color_progress
        
        self.color_progress += self.transition_speed
        if self.color_progress >= 1:
            self.color_progress = 0
            self.color_index = self.next_color_index
            self.next_color_index = (self.color_index + 1) % len(self.current_palette)
        
        return (int(r), int(g), int(b))

# 背景星空特效
class StarField:
    def __init__(self, num_stars=200):
        self.stars = []
        for _ in range(num_stars):
            x = random.randint(0, WIDTH)
            y = random.randint(0, HEIGHT)
            z = random.uniform(0, 1)
            self.stars.append([x, y, z])

    def move_stars(self):
        for star in self.stars:
            star[2] -= 0.05
            if star[2] <= 0:
                star[0] = random.randint(0, WIDTH)
                star[1] = random.randint(0, HEIGHT)
                star[2] = 1

    def draw(self, screen):
        for star in self.stars:
            size = int((1 - star[2]) * 3)
            brightness = int(255 * (1 - star[2]))
            color = (brightness, brightness, brightness)
            pygame.draw.circle(screen, color, (int(star[0]), int(star[1])), size)

# 3D变换矩阵（增加更多变换方式）
def create_transformation_matrices():
    matrices = {
        'rotate_x': lambda angle: np.array([
            [1, 0, 0],
            [0, math.cos(angle), -math.sin(angle)],
            [0, math.sin(angle), math.cos(angle)]
        ]),
        'rotate_y': lambda angle: np.array([
            [math.cos(angle), 0, math.sin(angle)],
            [0, 1, 0],
            [-math.sin(angle), 0, math.cos(angle)]
        ]),
        'rotate_z': lambda angle: np.array([
            [math.cos(angle), -math.sin(angle), 0],
            [math.sin(angle), math.cos(angle), 0],
            [0, 0, 1]
        ]),
        'shear_x': lambda s: np.array([
            [1, s, 0],
            [0, 1, 0],
            [0, 0, 1]
        ]),
        'shear_y': lambda s: np.array([
            [1, 0, 0],
            [s, 1, 0],
            [0, 0, 1]
        ])
    }
    return matrices

# 立方体生成器（增加随机性和复杂性）
def create_dynamic_cube(size=1, complexity=1):
    half = size / 2
    base_vertices = np.array([
        [-half, -half, -half],
        [half, -half, -half],
        [half, half, -half],
        [-half, half, -half],
        [-half, -half, half],
        [half, -half, half],
        [half, half, half],
        [-half, half, half]
    ])
    
    # 随机扰动顶点
    if complexity > 0:
        noise = np.random.uniform(-complexity * 0.1, complexity * 0.1, base_vertices.shape)
        base_vertices += noise
    
    return base_vertices

# 投影函数（改进的透视投影）
def advanced_project(point, distance=5, width=WIDTH, height=HEIGHT):
    x, y, z = point
    # 非线性深度缩放
    scale_factor = 1 / (1 + math.exp(-z))
    perspective_factor = distance / (distance + z + 0.1)
    
    x_proj = x * perspective_factor * width * 0.3
    y_proj = y * perspective_factor * height * 0.3
    
    return (
        int(x_proj + width / 2),
        int(y_proj + height / 2),
        scale_factor
    )

# 主动画类（增加更多动态效果）
class AdvancedCubeAnimation:
    def __init__(self):
        self.matrices = create_transformation_matrices()
        self.cube = create_dynamic_cube(complexity=1.5)
        self.color_palette = ColorPalette()
        self.starfield = StarField()
        
        # 多样的动画参数
        self.rotation_params = {
            'x': {'angle': 0, 'speed': 0.02, 'amplitude': 1},
            'y': {'angle': 0, 'speed': 0.015, 'amplitude': 1.2},
            'z': {'angle': 0, 'speed': 0.01, 'amplitude': 0.8}
        }
        
        # 额外的变形参数
        self.morph_params = {
            'shear_x': 0,
            'shear_y': 0,
            'morph_speed_x': 0.05,
            'morph_speed_y': 0.03
        }
        
        # 炫酷特效参数
        self.effect_params = {
            'pulse_scale': 1,
            'pulse_speed': 0.02,
            'trail_mode': False,
            'glitch_mode': False
        }

    def update_transformations(self):
        # 更新旋转
        for axis, params in self.rotation_params.items():
            params['angle'] += params['speed'] * params['amplitude']
        
        # 更新形变
        self.morph_params['shear_x'] = math.sin(pygame.time.get_ticks() * 0.001) * 0.5
        self.morph_params['shear_y'] = math.cos(pygame.time.get_ticks() * 0.002) * 0.5
        
        # 脉冲效果
        self.effect_params['pulse_scale'] = 1 + math.sin(pygame.time.get_ticks() * self.effect_params['pulse_speed']) * 0.2

    def apply_transformations(self):
        # 组合变换矩阵
        Rx = self.matrices['rotate_x'](self.rotation_params['x']['angle'])
        Ry = self.matrices['rotate_y'](self.rotation_params['y']['angle'])
        Rz = self.matrices['rotate_z'](self.rotation_params['z']['angle'])
        Sx = self.matrices['shear_x'](self.morph_params['shear_x'])
        Sy = self.matrices['shear_y'](self.morph_params['shear_y'])
        
        # 组合旋转和形变
        transformation = Rz @ Ry @ Rx @ Sy @ Sx
        return self.cube @ transformation.T

    def draw_cube(self, screen):
        # 应用变换
        transformed_cube = self.apply_transformations()
        
        # 立方体的边
        edges = [
            (0, 1), (1, 2), (2, 3), (3, 0),  # 底面
            (4, 5), (5, 6), (6, 7), (7, 4),  # 顶面
            (0, 4), (1, 5), (2, 6), (3, 7)   # 连接边
        ]
        
        # 获取动态颜色
        base_color = self.color_palette.get_interpolated_color()
        
        # 绘制边
        for edge in edges:
            start = transformed_cube[edge[0]]
            end = transformed_cube[edge[1]]
            
            start_2d = advanced_project(start)
            end_2d = advanced_project(end)
            
            # 根据深度调整线条
            depth_factor = start_2d[2]
            line_width = max(1, int(4 * depth_factor))
            
            # 应用脉冲和颜色变化
            pulse_color = tuple(int(c * (0.7 + depth_factor * 0.3)) for c in base_color)
            
            pygame.draw.line(screen, pulse_color, start_2d[:2], end_2d[:2], line_width)

def main():
    # 初始化
    animation = AdvancedCubeAnimation()
    
    # 视频导出准备
    os.makedirs('output', exist_ok=True)
    
    # OpenCV视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter('output/3d_cube_animation.mp4', fourcc, 60.0, (WIDTH, HEIGHT))
    
    # 创建背景
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    
    # 总帧数和持续时间
    frame_count = 0
    max_frames = 600  # 10秒 * 60帧
    
    while frame_count < max_frames:
        # 背景处理（拖尾效果）
        background.set_alpha(50)
        screen.blit(background, (0, 0))
        
        # 绘制星空背景
        animation.starfield.move_stars()
        animation.starfield.draw(screen)
        
        # 绘制立方体
        animation.update_transformations()
        animation.draw_cube(screen)
        
        # 更新显示
        pygame.display.flip()
        
        # 捕获帧
        frame_surface = screen.copy()
        frame_data = pygame.surfarray.array3d(frame_surface)
        frame_data = frame_data.transpose([1, 0, 2])
        
        # 转换为 OpenCV 格式 (BGR)
        frame_bgr = cv2.cvtColor(frame_data, cv2.COLOR_RGB2BGR)
        out.write(frame_bgr)
        
        frame_count += 1
        clock.tick(60)
    
    # 释放资源
    out.release()
    pygame.quit()
    
    print("视频导出完成！保存在 output/3d_cube_animation.mp4")

if __name__ == "__main__":
    main()