import pyrealsense2 as rs
import numpy as np
import cv2
import os
import time
import paramiko
from datetime import datetime
import socket
import urx
import math

# 服务器配置 - 请根据实际情况修改
SERVER_CONFIG = {
    'host': '10.103.69.239',
    'port': 27277,
    'username': 'chubin',
    'password': 'ZQSg7YOAeCqGPRky=Aa0#',
    'remote_dir': '/mnt/disk_3/chubin/img/'
}

REMOTE_HOST = "localhost"  # 控制命令远程主机
REMOTE_PORT = 9999        # 控制端口
UR_ROBOT_IP = "192.168.10.99"  # 机械臂IP
instruction = input("请输入任务描述：").strip()
# 配置摄像头管道
def configure_camera():
    pipeline = rs.pipeline()
    config = rs.config()
    
    config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
    
    profile = pipeline.start(config)
    
    depth_sensor = profile.get_device().first_depth_sensor()
    depth_scale = depth_sensor.get_depth_scale()
    print(f"深度比例系数为: {depth_scale}")
    
    align_to = rs.stream.color
    align = rs.align(align_to)
    
    return pipeline, align

# 创建视频写入器
def create_video_writer(width, height, fps=30):
    if not os.path.exists("videos"):
        os.makedirs("videos")
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    color_filename = f"videos/color_{timestamp}.avi"
    depth_filename = f"videos/depth_{timestamp}.avi"
    
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    color_out = cv2.VideoWriter(color_filename, fourcc, fps, (width, height))
    depth_out = cv2.VideoWriter(depth_filename, fourcc, fps, (width, height))
    
    return color_out, depth_out

# 上传文件到服务器
def upload_to_server(local_path, remote_path, server_config):
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(
            server_config['host'],
            port=server_config['port'],
            username=server_config['username'],
            password=server_config['password']
        )
        
        sftp = ssh.open_sftp()
        sftp.put(local_path, remote_path)
        print(f"文件 {local_path} 已上传到 {remote_path}")
        
        sftp.close()
        ssh.close()
        return True
    except Exception as e:
        print(f"上传失败: {e}")
        return False

# 获取机械臂控制动作
def get_action_from_remote(instruction: str) -> np.ndarray:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.connect((REMOTE_HOST, REMOTE_PORT))
        sock.sendall(instruction.encode())
        data = sock.recv(4096)
        action_str = data.decode()
        action = np.array([float(x) for x in action_str.split(",") if x.strip() != ""])
        return action[:-1]

# 机械臂控制
def move_robot(action, ur):
    x, y, z, rx, ry, rz = action
    #x = x * 1000  # 转换为毫米
    #y = y * 1000
    #z = (-0.4 + z) * 1000
    rx = 2.578 + 1.752 + rx
    ry = -0.771 + ry
    rz = 4.48 + rz
    ur.movej((x, y, z, rx, ry, rz), 0.5, 0.05, relative=False, wait=False)
    print("机械臂动作已执行:", action)

# 主程序
def main():
    # 连接到机械臂
    ur = urx.Robot(UR_ROBOT_IP)
    
    # 配置摄像头和视频写入器
    pipeline, align = configure_camera()
    color_out, depth_out = create_video_writer(640, 480)

    print("正在等待控制指令... 按 'q' 退出")
    
    recording = True
    last_capture_time = 0
    capture_interval = 5  # 每2秒截图一次

    while True:
        frames = pipeline.wait_for_frames()
        aligned_frames = align.process(frames)
        
        aligned_depth_frame = aligned_frames.get_depth_frame()
        color_frame = aligned_frames.get_color_frame()
        
        if not aligned_depth_frame or not color_frame:
            continue
        
        depth_image = np.asanyarray(aligned_depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())
        
        # 显示录制状态
        if recording:
            cv2.putText(color_image, "REC", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
        color_out.write(color_image)
        depth_out.write(cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET))
        
        # 每隔一段时间截图并上传
        current_time = time.time()
        if current_time - last_capture_time >= capture_interval:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            color_capture_path = f"captures/color_capture_{timestamp}.jpg"
            depth_capture_path = f"captures/depth_capture_{timestamp}.jpg"
            
            # 保存截图
            cv2.imwrite(color_capture_path, color_image)
            cv2.imwrite(depth_capture_path, depth_image)
            
            # 上传文件到服务器
            remote_color_path = os.path.join(SERVER_CONFIG['remote_dir'], f"color_capture_{timestamp}.jpg")
            remote_depth_path = os.path.join(SERVER_CONFIG['remote_dir'], f"depth_capture_{timestamp}.jpg")
            
            upload_to_server(color_capture_path, remote_color_path, SERVER_CONFIG)
            upload_to_server(depth_capture_path, remote_depth_path, SERVER_CONFIG)
            
            last_capture_time = current_time
        
        # 控制机械臂
       
        if instruction:
            print(f"发送任务到远程模型服务器...")
            action = get_action_from_remote(instruction)
            print(f"接收到动作: {action}")
            move_robot(action, ur)
        
        # 退出控制
        key = cv2.waitKey(1)
        if key & 0xFF == ord('q'):
            break
        elif key & 0xFF == ord('s'):  # 按s暂停/继续录制
            recording = not recording
            print("录制已暂停" if not recording else "录制继续")
    
    # 释放资源
    cv2.destroyAllWindows()
    pipeline.stop()
    color_out.release()
    depth_out.release()
    ur.close()

if __name__ == "__main__":
    main()