#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
设备管理模块
负责摄像头、IMU传感器和压力传感器的连接和数据采集
"""

import cv2
import numpy as np
import time
import threading
import json
from datetime import datetime
from typing import Dict, List, Optional, Any
import logging

logger = logging.getLogger(__name__)

class DeviceManager:
    """设备管理器"""
    
    def __init__(self):
        self.camera = None
        self.imu_device = None
        self.pressure_device = None
        self.device_status = {
            'camera': False,
            'imu': False,
            'pressure': False
        }
        self.calibration_data = {}
        self.data_lock = threading.Lock()
        self.latest_data = {}
        
        # 初始化设备
        self._init_devices()
    
    def _init_devices(self):
        """初始化所有设备"""
        try:
            self._init_camera()
            self._init_imu()
            self._init_pressure_sensor()
            logger.info('设备初始化完成')
        except Exception as e:
            logger.error(f'设备初始化失败: {e}')
    
    def _init_camera(self):
        """初始化摄像头"""
        try:
            # 尝试连接默认摄像头
            self.camera = cv2.VideoCapture(0)
            if self.camera.isOpened():
                # 设置摄像头参数
                self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
                self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
                self.camera.set(cv2.CAP_PROP_FPS, 30)
                
                self.device_status['camera'] = True
                logger.info('摄像头初始化成功')
            else:
                logger.warning('摄像头连接失败')
                self.camera = None
        except Exception as e:
            logger.error(f'摄像头初始化异常: {e}')
            self.camera = None
    
    def _init_imu(self):
        """初始化IMU传感器"""
        try:
            # 这里应该连接实际的IMU设备
            # 目前使用模拟数据
            self.imu_device = MockIMUDevice()
            self.device_status['imu'] = True
            logger.info('IMU传感器初始化成功（模拟）')
        except Exception as e:
            logger.error(f'IMU传感器初始化失败: {e}')
            self.imu_device = None
    
    def _init_pressure_sensor(self):
        """初始化压力传感器"""
        try:
            # 这里应该连接实际的压力传感器
            # 目前使用模拟数据
            self.pressure_device = MockPressureDevice()
            self.device_status['pressure'] = True
            logger.info('压力传感器初始化成功（模拟）')
        except Exception as e:
            logger.error(f'压力传感器初始化失败: {e}')
            self.pressure_device = None
    
    def get_device_status(self) -> Dict[str, bool]:
        """获取设备状态"""
        return self.device_status.copy()
    
    def get_connected_devices(self) -> List[str]:
        """获取已连接的设备列表"""
        return [device for device, status in self.device_status.items() if status]
    
    def refresh_devices(self):
        """刷新设备连接"""
        logger.info('刷新设备连接...')
        
        # 重新初始化所有设备
        if self.camera:
            self.camera.release()
        
        self._init_devices()
    
    def calibrate_devices(self) -> Dict[str, Any]:
        """校准设备"""
        calibration_result = {}
        
        try:
            # 摄像头校准
            if self.device_status['camera']:
                camera_calibration = self._calibrate_camera()
                calibration_result['camera'] = camera_calibration
            
            # IMU校准
            if self.device_status['imu']:
                imu_calibration = self._calibrate_imu()
                calibration_result['imu'] = imu_calibration
            
            # 压力传感器校准
            if self.device_status['pressure']:
                pressure_calibration = self._calibrate_pressure()
                calibration_result['pressure'] = pressure_calibration
            
            self.calibration_data = calibration_result
            logger.info('设备校准完成')
            
        except Exception as e:
            logger.error(f'设备校准失败: {e}')
            raise
        
        return calibration_result
    
    def _calibrate_camera(self) -> Dict[str, Any]:
        """校准摄像头"""
        if not self.camera or not self.camera.isOpened():
            return {'status': 'failed', 'error': '摄像头未连接'}
        
        try:
            # 获取几帧图像进行校准
            frames = []
            for _ in range(10):
                ret, frame = self.camera.read()
                if ret:
                    frames.append(frame)
                time.sleep(0.1)
            
            if not frames:
                return {'status': 'failed', 'error': '无法获取图像'}
            
            # 计算平均亮度和对比度
            avg_brightness = np.mean([np.mean(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)) for frame in frames])
            
            calibration = {
                'status': 'success',
                'brightness': float(avg_brightness),
                'resolution': (int(self.camera.get(cv2.CAP_PROP_FRAME_WIDTH)), 
                              int(self.camera.get(cv2.CAP_PROP_FRAME_HEIGHT))),
                'fps': float(self.camera.get(cv2.CAP_PROP_FPS)),
                'timestamp': datetime.now().isoformat()
            }
            
            return calibration
            
        except Exception as e:
            return {'status': 'failed', 'error': str(e)}
    
    def _calibrate_imu(self) -> Dict[str, Any]:
        """校准IMU传感器"""
        if not self.imu_device:
            return {'status': 'failed', 'error': 'IMU设备未连接'}
        
        try:
            # 收集静态数据进行零点校准
            samples = []
            for _ in range(100):
                data = self.imu_device.read_data()
                samples.append(data)
                time.sleep(0.01)
            
            # 计算零点偏移
            accel_offset = {
                'x': np.mean([s['accel']['x'] for s in samples]),
                'y': np.mean([s['accel']['y'] for s in samples]),
                'z': np.mean([s['accel']['z'] for s in samples]) - 9.8  # 重力补偿
            }
            
            gyro_offset = {
                'x': np.mean([s['gyro']['x'] for s in samples]),
                'y': np.mean([s['gyro']['y'] for s in samples]),
                'z': np.mean([s['gyro']['z'] for s in samples])
            }
            
            calibration = {
                'status': 'success',
                'accel_offset': accel_offset,
                'gyro_offset': gyro_offset,
                'timestamp': datetime.now().isoformat()
            }
            
            return calibration
            
        except Exception as e:
            return {'status': 'failed', 'error': str(e)}
    
    def _calibrate_pressure(self) -> Dict[str, Any]:
        """校准压力传感器"""
        if not self.pressure_device:
            return {'status': 'failed', 'error': '压力传感器未连接'}
        
        try:
            # 收集零压力数据
            samples = []
            for _ in range(50):
                data = self.pressure_device.read_data()
                samples.append(data)
                time.sleep(0.02)
            
            # 计算零点偏移
            zero_offset = {
                'left_foot': np.mean([s['left_foot'] for s in samples]),
                'right_foot': np.mean([s['right_foot'] for s in samples])
            }
            
            calibration = {
                'status': 'success',
                'zero_offset': zero_offset,
                'timestamp': datetime.now().isoformat()
            }
            
            return calibration
            
        except Exception as e:
            return {'status': 'failed', 'error': str(e)}
    
    def collect_data(self) -> Dict[str, Any]:
        """采集所有设备数据"""
        data = {}
        timestamp = datetime.now().isoformat()
        
        try:
            # 采集摄像头数据
            if self.device_status['camera']:
                camera_data = self._collect_camera_data()
                if camera_data:
                    data['camera'] = camera_data
            
            # 采集IMU数据
            if self.device_status['imu']:
                imu_data = self._collect_imu_data()
                if imu_data:
                    data['imu'] = imu_data
            
            # 采集压力传感器数据
            if self.device_status['pressure']:
                pressure_data = self._collect_pressure_data()
                if pressure_data:
                    data['pressure'] = pressure_data
            
            # 添加时间戳
            data['timestamp'] = timestamp
            
            # 更新最新数据
            with self.data_lock:
                self.latest_data = data.copy()
            
        except Exception as e:
            logger.error(f'数据采集失败: {e}')
        
        return data
    
    def _collect_camera_data(self) -> Optional[Dict[str, Any]]:
        """采集摄像头数据"""
        if not self.camera or not self.camera.isOpened():
            return None
        
        try:
            ret, frame = self.camera.read()
            if not ret:
                return None
            
            # 进行姿态检测（这里应该集成实际的姿态检测算法）
            pose_data = self._detect_pose(frame)
            
            return {
                'frame_available': True,
                'frame_size': frame.shape,
                'pose_data': pose_data,
                'timestamp': datetime.now().isoformat()
            }
            
        except Exception as e:
            logger.error(f'摄像头数据采集失败: {e}')
            return None
    
    def _detect_pose(self, frame: np.ndarray) -> Dict[str, Any]:
        """姿态检测（模拟实现）"""
        # 这里应该集成实际的姿态检测算法，如MediaPipe或OpenPose
        # 目前返回模拟数据
        return {
            'center_of_gravity': {
                'x': np.random.normal(0, 5),
                'y': np.random.normal(0, 5)
            },
            'body_angle': {
                'pitch': np.random.normal(0, 2),
                'roll': np.random.normal(0, 2),
                'yaw': np.random.normal(0, 1)
            },
            'confidence': np.random.uniform(0.8, 1.0)
        }
    
    def _collect_imu_data(self) -> Optional[Dict[str, Any]]:
        """采集IMU数据"""
        if not self.imu_device:
            return None
        
        try:
            raw_data = self.imu_device.read_data()
            
            # 应用校准偏移
            if 'imu' in self.calibration_data:
                calibration = self.calibration_data['imu']
                if calibration.get('status') == 'success':
                    accel_offset = calibration.get('accel_offset', {})
                    gyro_offset = calibration.get('gyro_offset', {})
                    
                    # 校正加速度数据
                    raw_data['accel']['x'] -= accel_offset.get('x', 0)
                    raw_data['accel']['y'] -= accel_offset.get('y', 0)
                    raw_data['accel']['z'] -= accel_offset.get('z', 0)
                    
                    # 校正陀螺仪数据
                    raw_data['gyro']['x'] -= gyro_offset.get('x', 0)
                    raw_data['gyro']['y'] -= gyro_offset.get('y', 0)
                    raw_data['gyro']['z'] -= gyro_offset.get('z', 0)
            
            return raw_data
            
        except Exception as e:
            logger.error(f'IMU数据采集失败: {e}')
            return None
    
    def _collect_pressure_data(self) -> Optional[Dict[str, Any]]:
        """采集压力传感器数据"""
        if not self.pressure_device:
            return None
        
        try:
            raw_data = self.pressure_device.read_data()
            
            # 应用校准偏移
            if 'pressure' in self.calibration_data:
                calibration = self.calibration_data['pressure']
                if calibration.get('status') == 'success':
                    zero_offset = calibration.get('zero_offset', {})
                    
                    raw_data['left_foot'] -= zero_offset.get('left_foot', 0)
                    raw_data['right_foot'] -= zero_offset.get('right_foot', 0)
            
            # 计算重心位置
            total_pressure = raw_data['left_foot'] + raw_data['right_foot']
            if total_pressure > 0:
                center_of_pressure = {
                    'x': (raw_data['right_foot'] - raw_data['left_foot']) / total_pressure * 100,
                    'y': 0  # 简化模型
                }
            else:
                center_of_pressure = {'x': 0, 'y': 0}
            
            raw_data['center_of_pressure'] = center_of_pressure
            raw_data['total_pressure'] = total_pressure
            
            return raw_data
            
        except Exception as e:
            logger.error(f'压力传感器数据采集失败: {e}')
            return None
    
    def get_latest_data(self) -> Dict[str, Any]:
        """获取最新采集的数据"""
        with self.data_lock:
            return self.latest_data.copy()
    
    def start_video_recording(self, output_path: str) -> bool:
        """开始视频录制"""
        if not self.camera or not self.camera.isOpened():
            return False
        
        try:
            # 获取摄像头参数
            width = int(self.camera.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(self.camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fps = int(self.camera.get(cv2.CAP_PROP_FPS))
            
            # 创建视频写入器
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            self.video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            
            if self.video_writer.isOpened():
                logger.info(f'开始视频录制: {output_path}')
                return True
            else:
                logger.error('视频写入器创建失败')
                return False
                
        except Exception as e:
            logger.error(f'开始视频录制失败: {e}')
            return False
    
    def stop_video_recording(self):
        """停止视频录制"""
        if hasattr(self, 'video_writer') and self.video_writer:
            self.video_writer.release()
            self.video_writer = None
            logger.info('视频录制已停止')
    
    def cleanup(self):
        """清理资源"""
        try:
            if self.camera:
                self.camera.release()
            
            if hasattr(self, 'video_writer') and self.video_writer:
                self.video_writer.release()
            
            logger.info('设备资源已清理')
            
        except Exception as e:
            logger.error(f'清理设备资源失败: {e}')


class MockIMUDevice:
    """模拟IMU设备"""
    
    def __init__(self):
        self.noise_level = 0.1
    
    def read_data(self) -> Dict[str, Any]:
        """读取IMU数据"""
        return {
            'accel': {
                'x': np.random.normal(0, self.noise_level),
                'y': np.random.normal(0, self.noise_level),
                'z': np.random.normal(9.8, self.noise_level)  # 重力加速度
            },
            'gyro': {
                'x': np.random.normal(0, self.noise_level),
                'y': np.random.normal(0, self.noise_level),
                'z': np.random.normal(0, self.noise_level)
            },
            'temperature': np.random.normal(25, 2),
            'timestamp': datetime.now().isoformat()
        }


class MockPressureDevice:
    """模拟压力传感器设备"""
    
    def __init__(self):
        self.base_pressure = 500  # 基础压力值
        self.noise_level = 10
    
    def read_data(self) -> Dict[str, Any]:
        """读取压力数据"""
        # 模拟轻微的左右脚压力差异
        left_pressure = self.base_pressure + np.random.normal(0, self.noise_level)
        right_pressure = self.base_pressure + np.random.normal(0, self.noise_level)
        
        return {
            'left_foot': max(0, left_pressure),
            'right_foot': max(0, right_pressure),
            'timestamp': datetime.now().isoformat()
        }