#!/usr/bin/env python3
import time
import os
import sys
import json
import csv
import serial
from enum import Enum
import threading

SERIAL_PORT = '/dev/ttyACM0'
BAUDRATE = 9600 
TIMEOUT = 1     

class OutputType(Enum):
    NoN=0
    Volt=1
    Curr=2

class SmuPower:
    
    def __init__(self):
        self.smu = None
        self.is_smu = False
        self.output = OutputType.NoN

    def send_and_read(self, command, wait_time=0.01):
        if self.smu is None:
            raise Exception("Not load SMU")
        full_command = command + '\n' # 添加换行符作为命令结束标志
        self.smu.write(full_command.encode()) # 发送命令
        # print(f"Sent: {command}")
        # time.sleep(wait_time)
        if command.strip().endswith('?'):
            response = self.smu.readline().decode().strip()
            # print(f"Received: {response}")
            return response
        else:
            time.sleep(wait_time)
            return None

    def load_smu(self):
        if os.path.exists(SERIAL_PORT):
            self.smu = serial.Serial(port=SERIAL_PORT, baudrate=BAUDRATE,
                            bytesize=8, parity='N', stopbits=1,
                            timeout=TIMEOUT)
            for _ in range(5):
                self.is_smu = self.is_keithley()
                if self.is_smu:
                    break
        else:
            raise Exception(f"Not found {SERIAL_PORT}")
    
    def is_keithley(self):
        t_msg = self.send_and_read("*IDN?")
        return "KEITHLEY" in t_msg

    def set_volt(self, volt) -> bool:
        if self.smu is None:
            return False
        self.send_and_read(f':SOUR:VOLT {volt}')
        self.send_and_read(':OUTP ON')
        time.sleep(0.05)

    def set_curr(self, curr) -> bool:
        if self.smu is None:
            return False
        self.send_and_read(f':SOUR:CURR {curr}')
        self.send_and_read(':OUTP ON')
        time.sleep(0.05)

    def close_power(self):
        if self.smu is not None:
            self.send_and_read(':OUTP OFF')

    def set_output_type(self, output: OutputType):
        self.send_and_read('*RST')
        self.output = output
        if output == OutputType.Volt:
            self.send_and_read(':SOUR:FUNC VOLT')
            self.send_and_read(':SENS:CURR:PROT 0.1') # 设置电流钳位为 0.1 A (100 mA)
            self.send_and_read(':SOUR:VOLT:RANG 10')    # 设定最大电压
        elif output == OutputType.Curr:
            self.send_and_read(':SOUR:FUNC CURR')
            self.send_and_read(':SENS:VOLT:PROT 10')     # 设置电压钳位为 10V
            self.send_and_read(':SOUR:CURR:RANG 0.02')    # 设定最大电流范围 20mA
            
    def close(self):
        if self.smu is not None:
            self.send_and_read(':OUTP OFF')
            self.smu.close()

class TestRelay:
    CMD_WRITE_SINGLE_COIL = 0x05
    CMD_READ_HOLDING_REGISTERS = 0x03
    
    def __init__(self, port: str, slave_id: int = 0x01, baudrate: int = 9600, timeout: float = 1.0):
        self.port = port
        self.slave_id = slave_id
        self.baudrate = baudrate
        self.timeout = timeout
        self.serial = None
        self.lock = threading.Lock()  # 串口操作锁，防止多线程冲突
        
    def connect(self) -> bool:
        try:
            self.serial = serial.Serial(
                port=self.port,
                baudrate=self.baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=self.timeout
            )
            self.serial.reset_input_buffer()
            self.serial.reset_output_buffer()
            print(f"成功连接到 {self.port} 波特率 {self.baudrate}")
            return True
        except serial.SerialException as e:
            print(f"串口打开失败: {e}")
            return False
    
    def disconnect(self):
        if self.serial and self.serial.is_open:
            self.serial.close()
    
    def is_connected(self) -> bool:
        return self.serial is not None and self.serial.is_open
    
    def _calculate_crc(self, data: bytes) -> bytes:
        crc = 0xFFFF
        for cur_byte in data:
            crc ^= cur_byte
            for _ in range(8):
                if crc & 0x0001:
                    crc = (crc >> 1) ^ 0xA001
                else:
                    crc >>= 1
        return bytes([crc & 0xFF, (crc >> 8) & 0xFF])
    
    def _send_command(self, cmd_bytes: bytes, expected_len: int = 8) -> bytes:
        with self.lock: 
            if not self.is_connected():
                raise ConnectionError("串口未连接，请先调用 connect()")
            
            crc = self._calculate_crc(cmd_bytes)
            full_cmd = cmd_bytes + crc
            
            self.serial.reset_input_buffer()
            self.serial.write(full_cmd)
            time.sleep(0.05)
            response = self.serial.read(expected_len)
            if len(response) >= 4:
                return response
            else:
                print(f"警告：响应长度异常 (期望 {expected_len}，实际 {len(response)})")
                return b''
    
    def set_relay(self, channel: int, state: bool) -> bool:

        cmd = bytearray()
        cmd.append(self.slave_id)
        cmd.append(self.CMD_WRITE_SINGLE_COIL)
        cmd.append((channel >> 8) & 0xFF)
        cmd.append(channel & 0xFF)
        if state:
            cmd.extend([0xFF, 0x00])
        else:
            cmd.extend([0x00, 0x00])
        
        response = self._send_command(cmd, expected_len=8)
        if len(response) == 8:
            if response[0] == self.slave_id and response[1] == self.CMD_WRITE_SINGLE_COIL:
                return True
        return False
    
    def set_all_relays(self, states: list) -> bool:
        success = True
        for idx, state in enumerate(states):
            if not self.set_relay(idx, state):
                print(f"设置第 {idx+1} 路继电器失败")
                success = False
        return success
    
    def read_relay_state(self, channel: int) -> bool:
        cmd = bytearray()
        cmd.append(self.slave_id)
        cmd.append(0x01) 
        cmd.append((channel >> 8) & 0xFF)
        cmd.append(channel & 0xFF)
        cmd.append(0x00)
        cmd.append(0x01)
        
        response = self._send_command(cmd, expected_len=6) 
        if len(response) >= 4:
            if response[0] == self.slave_id and response[1] == 0x01:
                data_byte = response[3]
                return (data_byte & 0x01) == 1
        return False

def read_channel_raw(channel):
    channel1 = channel - 1
    raw_file = f"/sys/bus/iio/devices/iio:device0/in_voltage{channel1}_raw"
    try:
        with open(raw_file, 'r', encoding='utf-8') as f:
            return int(f.read().strip())

    except Exception as e:
        print(str(e))
    return -1

def get_ad(channel):
    t_count = 5
    i = 0
    t_ad = []
    while True:
        if i < t_count:
            time.sleep(0.001)
            t_ad.append(read_channel_raw(channel))
            i += 1
        else:
            break
    t_max = max(t_ad)
    t_min = min(t_ad)
    t_ad.remove(t_max)
    t_ad.remove(t_min)
    return round(sum(t_ad)/len(t_ad))

def is_ad_failed(ad, expected):
    return abs(ad - expected) > 10

def load_compensation(data, index):
    t_ad = []
    t_ad.append(data[index][0] - data[index][1])
    for i in range(1,4):
        if (index - i) > 0:
            t_ad.append(data[index-i][0] - data[index-i][1])
        if (index + i) < len(data):
            t_ad.append(data[index+i][0] - data[index+i][1])

    t_max = max(t_ad)
    t_min = min(t_ad)
    t_ad.remove(t_max)
    t_ad.remove(t_min)
    return round(sum(t_ad)/len(t_ad))

def do_compensation(channel, t_data):
    adc_json = "/opt/ed-adc/compensation.json"
    calibration = None
    if os.path.exists(adc_json):
        with open(adc_json) as f:
            try:
                calibration = json.load(f)
            except Exception as e:
                print(f"[Error] Failed load compensation, {str(e)}")
    if calibration is None:
        calibration = {}
    
    if "channels" not in calibration:
        calibration["channels"] = {}
            
    t_len = len(t_data)
    i = 1
    channel_data = []
    i=0
    while True:
        if i < t_len:
            t_index_1 = len(channel_data)
            t_compensation = load_compensation(t_data, i)
            if t_index_1 == 0:
                channel_data.append({"max": t_data[i][0], "compensation": t_compensation})
            else:
                if channel_data[t_index_1-1]["max"] != t_data[i][0]:
                    if channel_data[t_index_1-1]["compensation"] == t_compensation:
                        channel_data[t_index_1-1]["max"] = t_data[i][0]
                    else:
                        channel_data.append({"max": t_data[i][0], "compensation": t_compensation})
            i += 1
        else:
            break

    calibration["channels"][str(channel)] = channel_data
    zero_data = {
        "1": get_ad(1),
        "2": get_ad(2),
        "3": get_ad(3),
        "4": get_ad(4)
    }
    calibration["zero"] = zero_data
    with open(adc_json, 'w', encoding='utf-8') as f:
        json.dump(calibration, f, ensure_ascii=False, indent=4)

def calibrate(smu: SmuPower, channel):
    if channel == 1 or channel == 2:
        smu.set_output_type(OutputType.Volt)
    elif channel == 3 or channel == 4:
        smu.set_output_type(OutputType.Curr)
    else:
        print("[Error] unknow channel.")
        return
    
    index = 0
    interval = 5
    data = []
    while True:
        if index > 1000:
            break
        if index > 0:
            if channel == 1 or channel == 2:
                smu.set_volt(index * 0.01)
            elif channel == 3 or channel == 4:
                smu.set_curr(index * 0.00002)
         
            t_ad = get_ad(channel)
            if is_ad_failed(t_ad, index):
                print(f"[Error] the AD value is invalid({index}-{t_ad}).")
                return
            data.append([index, t_ad])
        
        print(f"[Debug] Calibrating with value: {index/10}%")
        index += interval
    
    smu.close_power()
    with open(f'/opt/ed-adc/data-{channel}.csv', 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerows(data)
    do_compensation(channel, data)

if __name__ == "__main__":
    if os.geteuid() != 0:
        print("请以root用户运行此脚本")
        sys.exit(1)
    os.system("mkdir -p /opt/ed-adc")
    os.system("rm -rf /opt/ed-adc/*")
    relay = TestRelay(port='/dev/com2', slave_id=0x01, baudrate=38400)
    if not relay.connect():
        print("无法连接到继电器模块，请检查端口和硬件连接")
        exit(1)
    smu = SmuPower()
    smu.load_smu()
    relay.set_relay(0, False)
    relay.set_relay(1, False)
    relay.set_relay(2, False)
    relay.set_relay(3, False)
    try:
        channels = ["V0", "V1", "A0", "A1"]
        for i in range(1,5):
            if relay.set_relay(i-1, True):
                time.sleep(0.1)
                print(f"--------------------开始标定{channels[i-1]}-------------------------")
                t_startTime = time.time()
                calibrate(smu, i)
                t_time = time.time() - t_startTime
                print(f"[Info] {channels[i-1]} 标定时间: {t_time} seconds")
                if not relay.set_relay(i-1, False):
                    print(f"控制继电器第{i-1}失败")
                    break
                time.sleep(0.1)
            else:
                print(f"控制继电器第{i-1}失败")
                break
    finally:
        smu.close()
        relay.disconnect()