import pyvista as pv
import numpy as np

class AIDog:
    def __init__(self):
        """初始化渲染器，启用高质量渲染"""
        self.plotter = pv.Plotter()
        # 启用抗锯齿和高质量渲染
        self.plotter.enable_anti_aliasing()
        self.plotter.enable_depth_peeling(10)
        
        # 添加快捷键
        self.plotter.add_key_event('s', self.save_screenshot)
        self.plotter.add_key_event('q', self.plotter.close)
        
        # 添加快捷键提示
        self.plotter.add_text(
            "Keyboard shortcuts:\n"
            "S: Save screenshot\n"
            "Q: Quit",
            position='upper_left',
            font_size=12,
            color='black'
        )
        
        # 整体下移的偏移量
        self.z_offset = -0.1

    def create_head(self):
        """创建狗头"""
        parts = []
        
        # 主体 - 使用圆柱体和球体组合
        head_base = pv.Cylinder(
            center=(0.25, 0, 0.33),
            direction=(1, 0, 0),
            radius=0.06,
            height=0.1,
            resolution=32
        )
        parts.append(head_base)
        
        # 头部圆角
        head_front = pv.Sphere(
            center=(0.3, 0, 0.33),
            radius=0.06,
            phi_resolution=32,
            theta_resolution=32
        )
        head_back = pv.Sphere(
            center=(0.2, 0, 0.33),
            radius=0.06,
            phi_resolution=32,
            theta_resolution=32
        )
        parts.extend([head_front, head_back])
        
        # 眼睛外壳 - 凸起的圆柱体
        left_eye_shell = pv.Cylinder(
            center=(0.32, 0.04, 0.34),
            direction=(1, 0.1, 0),
            radius=0.015,
            height=0.01,
            resolution=32
        )
        right_eye_shell = pv.Cylinder(
            center=(0.32, -0.04, 0.34),
            direction=(1, -0.1, 0),
            radius=0.015,
            height=0.01,
            resolution=32
        )
        parts.extend([left_eye_shell, right_eye_shell])
        
        # 眼睛镜片 - 稍微凸出的深色圆柱体
        left_eye_lens = pv.Cylinder(
            center=(0.325, 0.04, 0.34),
            direction=(1, 0.1, 0),
            radius=0.01,
            height=0.005,
            resolution=32
        )
        right_eye_lens = pv.Cylinder(
            center=(0.325, -0.04, 0.34),
            direction=(1, -0.1, 0),
            radius=0.01,
            height=0.005,
            resolution=32
        )
        parts.extend([left_eye_lens, right_eye_lens])
        
        # 鼻子部分 - 传感器样式
        nose = pv.Cylinder(
            center=(0.35, 0, 0.32),
            direction=(1, 0, 0),
            radius=0.015,
            height=0.01,
            resolution=32
        )
        nose_sensor = pv.Cylinder(
            center=(0.355, 0, 0.32),
            direction=(1, 0, 0),
            radius=0.01,
            height=0.005,
            resolution=32
        )
        parts.extend([nose, nose_sensor])
        
        return parts, [left_eye_lens, right_eye_lens, nose_sensor]  # 返回两组部件，第二组是深色部件

    def create_body(self):
        """创建机器狗躯干"""
        parts = []
        
        # 主体 - 使用圆柱体组合
        body_main = pv.Cylinder(
            center=(0, 0, 0.3),
            direction=(1, 0, 0),
            radius=0.08,
            height=0.3,
            resolution=32
        )
        parts.append(body_main)
        
        # 前后圆角
        front_cap = pv.Sphere(
            center=(0.15, 0, 0.3),
            radius=0.08,
            phi_resolution=32,
            theta_resolution=32
        )
        back_cap = pv.Sphere(
            center=(-0.15, 0, 0.3),
            radius=0.08,
            phi_resolution=32,
            theta_resolution=32
        )
        parts.extend([front_cap, back_cap])
        
        # 顶部装甲 - 更薄的圆柱体
        top_armor = pv.Cylinder(
            center=(0, 0, 0.36),
            direction=(1, 0, 0),
            radius=0.07,
            height=0.25,
            resolution=32
        )
        parts.append(top_armor)
        
        return parts

    def create_leg(self, position, side='right'):
        """创建单条腿"""
        # 髋关节
        hip = pv.Sphere(radius=0.03, center=position)
        
        # 大腿
        thigh_length = 0.12
        thigh_pos = (position[0], 
                    position[1], 
                    position[2] - thigh_length/2)
        thigh = pv.Cylinder(
            center=thigh_pos,
            direction=(0, 0, 1),
            radius=0.015,
            height=thigh_length,
            resolution=32
        )
        
        # 膝关节
        knee_pos = (position[0], 
                   position[1], 
                   position[2] - thigh_length)
        knee = pv.Sphere(radius=0.02, center=knee_pos)
        
        # 小腿
        shin_length = 0.15
        shin_pos = (knee_pos[0],
                   knee_pos[1],
                   knee_pos[2] - shin_length/2)
        shin = pv.Cylinder(
            center=shin_pos,
            direction=(0, 0, 1),
            radius=0.012,
            height=shin_length,
            resolution=32
        )
        
        # 脚部
        foot_pos = (knee_pos[0],
                   knee_pos[1],
                   knee_pos[2] - shin_length)
        foot = pv.Box(bounds=(
            foot_pos[0]-0.02, foot_pos[0]+0.02,
            foot_pos[1]-0.015, foot_pos[1]+0.015,
            foot_pos[2]-0.01, foot_pos[2]+0.01
        ))
        
        return [hip, thigh, knee, shin, foot]

    def create_model(self):
        """创建完整模型"""
        # 材质颜色定义 - 更亮的金属色
        colors = {
            'head': '#BDC3C7',  # 更亮的金属灰
            'body': '#BDC3C7',  # 更亮的金属灰
            'limbs': '#95A5A6', # 金属灰
            'sensors': '#34495E' # 深色传感器
        }
        
        # 统一的金属材质参数 - 增加光泽
        metal_props = {
            'metallic': 0.9,
            'roughness': 0.1,  # 降低粗糙度
            'specular': 1.0,   # 增加反光
            'specular_power': 20
        }
        
        # 传感器材质参数
        sensor_props = {
            'metallic': 0.9,
            'roughness': 0.1,
            'specular': 1.0,
            'specular_power': 20
        }
        
        # 添加头部
        head_parts, sensor_parts = self.create_head()
        for part in head_parts:
            self.plotter.add_mesh(part, 
                                color=colors['head'],
                                **metal_props)
        # 添加深色传感器部件
        for part in sensor_parts:
            self.plotter.add_mesh(part,
                                color=colors['sensors'],
                                **sensor_props)
        
        # 添加躯干
        for part in self.create_body():
            self.plotter.add_mesh(part,
                                color=colors['body'],
                                **metal_props)
            
        # 添加四条腿
        leg_positions = [
            (0.12, 0.08, 0.3),   # 右前
            (0.12, -0.08, 0.3),  # 左前
            (-0.12, 0.08, 0.3),  # 右后
            (-0.12, -0.08, 0.3)  # 左后
        ]
        
        for pos in leg_positions:
            for part in self.create_leg(pos):
                self.plotter.add_mesh(part,
                                    color=colors['limbs'],
                                    **metal_props)

        # 渲染设置 - 使用更亮的背景
        self.plotter.set_background('#CCCCCC')  # 更亮的灰色背景
        self.plotter.camera_position = 'iso'
        self.plotter.camera.zoom(1.5)
        
    def show(self):
        """显示模型"""
        self.plotter.show()

    def save_screenshot(self, filename=None):
        """保存截图并显示提示"""
        if filename is None:
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f'aidog_{timestamp}.png'
        
        self.plotter.screenshot(filename)
        print(f"Screenshot saved as: {filename}")
        # 添加临时提示文字
        self.plotter.add_text(
            f"Saved: {filename}",
            position='upper_right',
            font_size=12,
            color='black'
        )

if __name__ == '__main__':
    dog = AIDog()
    dog.create_model()
    dog.show()