import time
import socket
import struct
import numpy as np
import atexit
import threading
from typing import Optional, Tuple
from matplotlib import pyplot as plt
import csv
from datetime import datetime

class XjcSensor:
    def __init__(self, host: str, port: int, rate: int = 250):
        """
        初始化 TCP 连接的力传感器
        :param host: 设备 IP 地址（如 "192.168.1.100"）
        :param port: 设备端口（如 502）
        :param rate: 数据采样率（100/250/500Hz）
        """
        self.host = host
        self.port = port
        self.rate = rate
        self.slave_address = 0x01  # 默认从机地址
        self.sock = None  # TCP socket 对象
        self.crc16_table = self.generate_crc16_table()
        # 后台读取相关的属性
        self._background_thread = None
        self._stop_background = False
        self._last_reading = None
        self._reading_lock = threading.Lock()
        # 建立 TCP 连接
        self.connect()
        
        # 注册退出时的清理函数
        atexit.register(self.disconnect)

    def __new__(cls, *args, **kwargs):
        """单例模式"""
        if not hasattr(cls, 'instance'):
            cls.instance = super(XjcSensor, cls).__new__(cls)
        return cls.instance

    def connect(self):
        """建立 TCP 连接"""
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.settimeout(1.0)  # 设置超时时间
            self.sock.connect((self.host, self.port))
            print(f"Connected to {self.host}:{self.port}")
        except Exception as e:
            print(f"TCP connection failed: {e}")
            raise

    def send_and_receive(self, data: bytes, expected_response_len: int) -> Optional[bytes]:
        """
        发送数据并接收响应
        :param data: 要发送的字节数据
        :param expected_response_len: 预期响应的长度
        :return: 收到的字节数据（失败返回 None）
        """
        try:
            self.sock.sendall(data)
            response = self.sock.recv(expected_response_len)
            return response
        except socket.timeout:
            print("Timeout while waiting for response")
            return None
        except Exception as e:
            print(f"TCP communication error: {e}")
            return None

    def set_zero(self) -> int:
        """传感器置零（TCP 版本）"""
        if self.slave_address == 0x01:
            zero_cmd = bytes.fromhex('01 10 46 1C 00 02 04 00 00 00 00 E8 95')
        elif self.slave_address == 0x09:
            zero_cmd = bytes.fromhex('09 10 46 1C 00 02 04 00 00 00 00 C2 F5')

        response = self.send_and_receive(zero_cmd, 8)
        if not response:
            print("set zero fail")
            return -1

        # CRC 校验（假设协议要求）
        hi_check, lo_check = self.crc16(response[:-2])
        if (response[0] != self.slave_address or 
            response[1] != 0x10 or 
            response[-2] != lo_check or 
            response[-1] != hi_check):
            print("Set zero failed: CRC check error")
            return -1

        print("Set zero success!")
        return 0

    def read_data_f32(self) -> np.ndarray:
        """
        TCP版本 - 读取六维力传感器数据 (float32格式)
        返回: 
            np.ndarray(6,) - 六维力数据 [Fx,Fy,Fz,Mx,My,Mz]
            读取失败返回 -1
        """
        # 构造读取命令

        if self.slave_address == 0x01:
            command = bytes.fromhex('01 04 00 00 00 0C F0 0F')
        elif self.slave_address == 0x09:
            command = bytes.fromhex('09 04 00 54 00 0C B0 97')
        
        try:
            # 清空接收缓冲区
            self._clear_tcp_buffer()
            
            # 接收完整响应
            response = self.send_and_receive(command, 29)
            
            # CRC校验
            Hi_check, Lo_check = self.crc16(response[:-2])
            if (response[0] != self.slave_address or 
                response[1] != 0x04 or 
                response[2] != 24 or  # 数据长度字节
                response[-2] != Lo_check or 
                response[-1] != Hi_check):
                print(f"Protocol error! Received: {response.hex(' ')}")
                print(f"Expected slave:{self.slave_address}, "
                    f"func:0x04, len:24, "
                    f"CRC:{Lo_check:02X}{Hi_check:02X}")
                return -1
            
            # 解析6个float32数据 (大端序)
            sensor_data = struct.unpack('>ffffff', response[3:27])
            return np.array(sensor_data)
            
        except socket.timeout:
            print("Timeout while waiting for sensor response")
            return -1
        except ConnectionError as e:
            print(f"TCP connection error: {e}")
            return -1
        except Exception as e:
            print(f"Unexpected error in read_data_f32: {str(e)}")
            return -1

    def _clear_tcp_buffer(self):
        """清空TCP接收缓冲区"""
        self.sock.settimeout(0.1)  # 短暂超时，避免阻塞
        try:
            while True:
                data = self.sock.recv(1024)
                if not data:  # 连接关闭或无数据
                    break
                print(f"Cleared TCP buffer: {data.hex(' ')}")  # 调试输出
        except socket.timeout:
            pass  # 预期内的超时，表示缓冲区已空
        finally:
            self.sock.settimeout(2.0)  # 恢复默认超时

    def enable_active_transmission(self) -> int:
        """启用主动传输模式（TCP 版本）"""
        if self.rate == 100:
            cmd = bytes.fromhex('01 10 01 9A 00 01 02 00 00 AB 6A') if self.slave_address == 0x01 else \
                  bytes.fromhex('09 10 01 9A 00 01 02 00 00 CC AA')
        elif self.rate == 250:
            cmd = bytes.fromhex('01 10 01 9A 00 01 02 00 01 6A AA') if self.slave_address == 0x01 else \
                  bytes.fromhex('09 10 01 9A 00 01 02 00 01 0D 6A')
        elif self.rate == 500:
            cmd = bytes.fromhex('01 10 01 9A 00 01 02 00 02 2A AB') if self.slave_address == 0x01 else \
                  bytes.fromhex('09 10 01 9A 00 01 02 00 02 4D 6B')
        else:
            print("Unsupported rate")
            return -1

        response = self.send_and_receive(cmd, 8)
        if response and response[1] == 0x10:  # 检查响应功能码
            print(f"Active transmission enabled at {self.rate}Hz")
            return 0
        else:
            print("Failed to enable active transmission")
            return -1

    def disconnect(self):
        """关闭 TCP 连接"""
        if self.sock:
            try:
                self.sock.close()
                print("TCP connection closed")
            except Exception as e:
                print(f"Error while closing socket: {e}")
            self.sock = None
    def start_background_reading(self):
        """
        启动后台读取任务，每秒读取一次传感器数据
        """
        if self._background_thread is not None and self._background_thread.is_alive():
            print("Background reading is already running")
            return False
        
        self._stop_background = False
        self._background_thread = threading.Thread(target=self._background_reading_task)
        self._background_thread.daemon = True  # 设置为守护线程
        self._background_thread.start()
        return True

    def stop_background_reading(self):
        """
        停止后台读取任务
        """
        if self._background_thread is None or not self._background_thread.is_alive():
            print("No background reading is running")
            return False
        
        self._stop_background = True
        self._background_thread.join(timeout=2.0)  # 等待线程结束，最多等待2秒
        self._background_thread = None
        return True

    def _background_reading_task(self):
        """
        后台读取任务的实现
        """
        while not self._stop_background:
            try:
                data = self.read_data_f32()
                if isinstance(data, np.ndarray):
                    with self._reading_lock:
                        self._last_reading = data
                else:
                    self.connect()
            except Exception as e:
                print(f"Error in background reading: {e}")
            time.sleep(1.0)  # 每秒读取一次

    def get_last_reading(self):
        """
        获取最近一次的读数
        """
        with self._reading_lock:
            return self._last_reading

    # ------------------------- 以下方法保持不变 -------------------------
    @staticmethod
    def generate_crc16_table():
        """CRC16 表生成（与原代码相同）"""
        table = []
        polynomial = 0xA001
        for i in range(256):
            crc = i
            for _ in range(8):
                if crc & 0x0001:
                    crc = (crc >> 1) ^ polynomial
                else:
                    crc >>= 1
            table.append(crc)
        return table

    def crc16(self, data: bytes):
        """CRC16 计算（与原代码相同）"""
        crc = 0xFFFF
        for byte in data:
            crc = (crc >> 8) ^ self.crc16_table[(crc ^ byte) & 0xFF]
        return (crc >> 8) & 0xFF, crc & 0xFF

    def __del__(self):
        """析构时自动断开连接"""
        self.disconnect()

    def disable_active_transmission(self) -> int:
        """
        禁用传感器的主动传输模式（TCP 版本）
        Returns:
            0 表示成功，-1 表示失败
        """
        try:
            # 构造禁用命令（原代码中的 FF x11 字节流）
            disable_cmd = bytes.fromhex('FF FF FF FF FF FF FF FF FF FF FF')
            print("Disabling active transmission mode...")

            # 检查 TCP 连接是否有效
            if not self.sock:
                raise ConnectionError("TCP connection is not established")

            # 发送禁用命令
            self.sock.sendall(disable_cmd)
            print("Active transmission disabled successfully")
            return 0

        except ConnectionError as e:
            print(f"Connection error: {e}")
            return -1
        except socket.timeout:
            print("Timeout while sending disable command")
            return -1
        except Exception as e:
            print(f"Unexpected error: {e}")
            return -1
        
    def read(self):
        """
        TCP Version - Read the sensor's data in passive mode.
        Returns:
            np.ndarray(6,) - Six-axis force/torque data [Fx, Fy, Fz, Mx, My, Mz]
            None if reading fails
        """
        try:
            # Clear any existing data in the TCP buffer first
            # self._clear_tcp_buffer()
            
            # Search for frame header (0x20 0x4E)
            header_found = False
            start_time = time.time()
            timeout = 1.0  # 1 second timeout
            
            while not header_found and (time.time() - start_time < timeout):
                # Read one byte at a time looking for header
                byte1 = self.sock.recv(1)
                if not byte1:
                    continue  # No data available yet
                
                if byte1 == b'\x20':
                    byte2 = self.sock.recv(1)
                    if byte2 == b'\x4E':
                        header_found = True
            
            if not header_found:
                print("Frame header not found within timeout period")
                return None
            
            # Now read the remaining 14 bytes of the frame
            response = bytearray([0x20, 0x4E])  # Start with the header we found
            remaining_bytes = 14
            bytes_received = 0
            start_time = time.time()
            
            while bytes_received < remaining_bytes and (time.time() - start_time < timeout):
                chunk = self.sock.recv(remaining_bytes - bytes_received)
                if chunk:
                    response.extend(chunk)
                    bytes_received += len(chunk)
            
            if bytes_received < remaining_bytes:
                print(f"Incomplete frame received. Got {len(response)} bytes, expected 16")
                return None
            
            # Verify CRC checksum
            Hi_check, Lo_check = self.crc16(response[:-2])
            if response[-1] != Hi_check or response[-2] != Lo_check:
                print("CRC check failed!")
                print(f"Received CRC: {response[-2]:02X}{response[-1]:02X}")
                print(f"Calculated CRC: {Lo_check:02X}{Hi_check:02X}")
                return None
            
            # Parse the sensor data
            sensor_data = self.parse_data_passive(response)
            return sensor_data
            
        except socket.timeout:
            print("Timeout while waiting for sensor data")
            return None
        except ConnectionError as e:
            print(f"TCP connection error: {e}")
            return None
        except Exception as e:
            print(f"Unexpected error in TCP read(): {str(e)}")
            return None

    def parse_data_passive(self, buffer):
        values = [
                int.from_bytes(buffer[i:i+2], byteorder='little', signed=True)
                for i in range(2, 14, 2)
                ]
        Fx, Fy, Fz = np.array(values[:3]) / 10.0
        Mx, My, Mz = np.array(values[3:]) / 1000.0
        return np.array([Fx, Fy, Fz, Mx, My, Mz])


# 外部
def test_sensor_frequency(sensor, mode='active', duration=5.0):
    """
    测试传感器在不同模式下的实际数据获取频率
    
    参数:
        sensor: XjcSensor实例
        mode: 'active'(主动模式) 或 'polling'(查询模式)
        duration: 测试持续时间(秒)
    
    返回:
        dict: 包含测试结果的字典
    """
    # 准备测试环境
    if mode == 'active':
        print("\n=== 测试主动传输模式 ===")
        sensor.enable_active_transmission()
        read_func = sensor.read
    else:
        print("\n=== 测试查询模式 ===")
        sensor.disable_active_transmission()  # 确保不在主动模式
        time.sleep(0.5)
        sensor.disable_active_transmission()  # 确保不在主动模式
        time.sleep(0.5)
        sensor.disable_active_transmission()  # 确保不在主动模式
        time.sleep(0.5)
        read_func = sensor.read_data_f32
    
    # 初始化测试变量
    timestamps = []
    data_count = 0
    start_time = time.perf_counter()
    end_time = start_time + duration
    
    print(f"开始测试，持续时间 {duration} 秒...")
    
    # 主测试循环
    while time.perf_counter() < end_time:
        data = read_func()
        print(data)
        if data is not None:
            timestamps.append(time.perf_counter())
            data_count += 1
    # print(timestamps)
    # 计算统计结果
    if data_count < 2:
        print("警告: 采集到的数据点不足")
        return None
    
    intervals = np.diff(timestamps)
    save_intervals_to_csv(timestamps)
    avg_interval = np.mean(intervals)
    min_interval = np.min(intervals)
    max_interval = np.max(intervals)
    std_interval = np.std(intervals)
    avg_freq = 1.0 / avg_interval
    
    # 打印结果
    print("\n测试结果:")
    print(f"总数据点数: {data_count}")
    print(f"平均频率: {avg_freq:.2f} Hz")
    print(f"最小间隔: {min_interval*1000:.3f} ms")
    print(f"最大间隔: {max_interval*1000:.3f} ms")
    print(f"间隔标准差: {std_interval*1000:.3f} ms")
    
    # 返回结果字典
    return {
        'mode': mode,
        'duration': duration,
        'data_count': data_count,
        'avg_freq': avg_freq,
        'min_interval': min_interval,
        'max_interval': max_interval,
        'std_interval': std_interval,
        'timestamps': timestamps,
        'intervals': intervals
    }

def plot_test_results(active_results, polling_results):
    """绘制两种模式的测试结果对比图"""
    plt.figure(figsize=(12, 8))
    
    # 间隔时间分布图
    plt.subplot(2, 1, 1)
    plt.hist(active_results['intervals']*1000, bins=50, alpha=0.7, label='主动模式')
    plt.hist(polling_results['intervals']*1000, bins=50, alpha=0.7, label='查询模式')
    plt.xlabel('间隔时间 (ms)')
    plt.ylabel('出现次数')
    plt.title('数据间隔时间分布')
    plt.legend()
    plt.grid(True)
    
    # 间隔时间序列图
    plt.subplot(2, 1, 2)
    plt.plot(np.arange(len(active_results['intervals'])), 
             active_results['intervals']*1000, 'b.', label='主动模式')
    plt.plot(np.arange(len(polling_results['intervals'])), 
             polling_results['intervals']*1000, 'r.', label='查询模式')
    plt.xlabel('数据点序号')
    plt.ylabel('间隔时间 (ms)')
    plt.title('数据间隔时间序列')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

def compare_modes(sensor, duration=5.0):
    """比较两种工作模式的性能"""
    # 测试主动模式
    active_results = test_sensor_frequency(sensor, 'active', duration)
    time.sleep(1)  # 模式切换间隔
    
    # 测试查询模式
    polling_results = test_sensor_frequency(sensor, 'polling', duration)
    time.sleep(1)
    
    # 打印对比报告
    print("\n=== 模式对比报告 ===")
    print(f"{'指标':<15} {'主动模式':>15} {'查询模式':>15}")
    print(f"{'平均频率(Hz)':<15} {active_results['avg_freq']:>15.2f} {polling_results['avg_freq']:>15.2f}")
    print(f"{'最小间隔(ms)':<15} {active_results['min_interval']*1000:>15.3f} {polling_results['min_interval']*1000:>15.3f}")
    print(f"{'最大间隔(ms)':<15} {active_results['max_interval']*1000:>15.3f} {polling_results['max_interval']*1000:>15.3f}")
    print(f"{'间隔标准差(ms)':<15} {active_results['std_interval']*1000:>15.3f} {polling_results['std_interval']*1000:>15.3f}")
    
    # 绘制对比图表
    plot_test_results(active_results, polling_results)
    
    return active_results, polling_results

def save_intervals_to_csv(timestamps, filename="sensor_intervals.csv"):
    """
    Save timestamp intervals to CSV file with metadata
    
    Args:
        timestamps: List of timestamp values (from time.perf_counter())
        filename: Output CSV filename
    """
    if len(timestamps) < 2:
        print("Not enough timestamps to calculate intervals")
        return
    
    # Calculate intervals in milliseconds
    intervals_ms = np.diff(timestamps) * 1000  
    
    # Prepare data rows
    rows = []
    for i, interval in enumerate(intervals_ms):
        rows.append({
            "reading_number": i+1,
            "interval_ms": f"{interval:.4f}",
            "timestamp": timestamps[i],
            "expected_interval_ms": (1000/250) if i==0 else "",  # For 250Hz
            "deviation_ms": f"{interval - (1000/250):.4f}" if i>0 else ""
        })
    
    # CSV header
    fieldnames = ["reading_number", "interval_ms", "timestamp", 
                 "expected_interval_ms", "deviation_ms"]
    
    # Write to file
    with open(filename, 'w', newline='') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)
    
    print(f"Saved {len(intervals_ms)} intervals to {filename}")

if __name__ == "__main__":
    # 替换为你的传感器 IP 和端口
    sensor = XjcSensor("192.168.5.1", 60000)
    
    # 示例操作
    
    sensor.disable_active_transmission()
    time.sleep(0.5)
    sensor.disable_active_transmission()
    time.sleep(0.5)
    sensor.disable_active_transmission()
    time.sleep(0.5)
    sensor.set_zero()

    # sensor.set_zero()
    # sensor.disable_active_transmission()
    
    # while True:
    #     data = sensor.read_data_f32()
    #     if data is not None:
    #         print(f"Force data: {data}")
    # sensor.enable_active_transmission()
    # while True:
    #     sensor_data = sensor.read()
    #     if sensor_data is None:
    #        print('failed to get force sensor data!')
    #     print(sensor_data)
    test_sensor_frequency(sensor)
    # compare_modes(sensor)