import subprocess
import sys
import json
import math
import time
import datetime
import yaml

# 版本号
# 首位数字为大版本号
# 后两位取发布日志
VERSION = "1.5.27"

# 用于缓存pitch的文件，平地标定前读10次pitch，取平均值，越障标定时读取并减去缓存的patch偏差
pitch_file_in_device = "/tmp/log/4cm_point_tof_test_pitch.txt"
# 单点tof的文件
file_path = "/mnt/misc/single_dtof.flag"
# 当前的功率模式
power_mode = "low_power"

# unical-conf 配置的文件
CONST_UNICAL_CONF = "/data/log/dreame_factory_run_parameters.yaml"

# tof 下方空旷无遮挡距离
global_xtalk_pre_limit = 30000
#抬升后pitch阈值
climb_up_pitch_limit = [1200, 1400]
# tof下视离地高度,单位:mm
climb_tof_to_ground_distance = 23
#后小轮离地高度,单位:mm
climb_rear_wheel_to_ground_distance = 7
#后小轮离双轮轴间距,单位:mm
climb_rear_wheel_to_wheel_distance = 163.9
#tof下视离双轮轴间距,单位:mm
climb_tof_to_wheel_distance = 149.16

def call_adb_command(command):
    """Helper function to call adb command and capture its output."""
    try:
        print("RUN CMD:", command)
        result = subprocess.run(command, shell=True, check=True, text=True, capture_output=True, encoding='utf-8', errors='ignore')
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"Error executing command: {command}\n{e.stderr}")
        return ""

def check_file_exists(file_path):
    # 构造检查文件存在的ADB命令
    command = f'adb shell ls {file_path}'
    result = call_adb_command(command)
    
    if 'No such file or directory' in result:
        return False
    else:
        return True

def calib_xtalk(modulename):
    print(datetime.datetime.now(), " RESULT:")
    # 先读取一下点tof值，确认附近没有障碍物
    command1 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"get_distance\"}'"
    result_str = call_adb_command(command1)
    print(datetime.datetime.now(), " RESULT:", result_str)
    ret_data = json.loads(result_str)
    tmp_value = ret_data["value"]
    if (ret_data["value"] < global_xtalk_pre_limit):
        print(f"{modulename} ERROR: make there is no object neer tof during xtalk calib!{tmp_value}vs{global_xtalk_pre_limit}")
        sys.exit(0)
    command = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"cali_xtalk\"}'"
    result_str = call_adb_command(command)
    time.sleep(1)
    while 1:
        command1 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"get_cali_status\"}'"
        result_str = call_adb_command(command1)
        print(datetime.datetime.now(), "RESULT:", result_str)
        ret_data = json.loads(result_str)
        if (ret_data["value"] == 2):
            print(modulename, "OK")
            break
        elif ret_data["value"] == 3:
            print(modulename, "ERROR:", result_str)
            sys.exit(0)
        else:
            time.sleep(0.2)



def calib_ground_offset(modulename, depth):
    command2 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"cali_offset\",\"value\":" + str(climb_tof_to_ground_distance + depth) + "}'"
    result_str = call_adb_command(command2)
    print("RESULT:", result_str)
    time.sleep(1)
    while 1:
        command1 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"get_cali_status\"}'"
        result_str = call_adb_command(command1)
        ret_data = json.loads(result_str)
        if (ret_data["value"] == 2):
            print(modulename, "OK")
            break
        elif ret_data["value"] == 3:
            print(modulename, "ERROR:", result_str)
            sys.exit(0)
        else:
            time.sleep(0.2)

def calib_crossing(modulename, target):
    # 发送标定指令
    command2 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"cali_gradient\",\"value\":" + str(target) + "}'"
    result_str = call_adb_command(command2)
    print("RESULT:", result_str)
    time.sleep(1)
    while 1:
        command1 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"get_cali_status\"}'"
        result_str = call_adb_command(command1)
        ret_data = json.loads(result_str)
        if (ret_data["value"] == 2):
            print(modulename, "OK")
            break
        elif ret_data["value"] == 3:
            print(modulename, "ERROR:", result_str)
            sys.exit(0)
        else:
            time.sleep(0.2)

def change_range_mode(modulename, rangemode):
    # 170是低光功率模式，85是高光功率模式
    command1 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"set_range_mode\",\"mode\":" + str(rangemode) + "}'"
    result_str = call_adb_command(command1)

def check_ground_offset(modulename, depth):
    command2 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"get_distance\"}'"
    result_str = call_adb_command(command2)
    print(f"RESULT: {result_str}")
    ret_data = json.loads(result_str)
    print(modulename, power_mode, "VALUE:", float(ret_data["value"]), ", DIFF:", float(ret_data["value"]) * 0.01 - (climb_tof_to_ground_distance + depth), "mm")

def check_crossing(modulename, target):
    # 发送标定指令
    command2 = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"module\":\"" + modulename + r"\",\"cmd\":\"get_distance\"}'"
    # 取两次数据，把之前缓存的脏数扔掉
    result_str = call_adb_command(command2)
    print("RESULT:", result_str)

    ret_data = json.loads(result_str)
    print(modulename, power_mode, "VALUE:", float(ret_data["value"]), ", DIFF:", float(ret_data["value"]) * 0.01 - target, "mm")

# 返回roll pitch角度，单位0.01度
# 支持多次读roll和pitch取平均值,时间间隔0.1s
def get_roll_pitch(sample_count = 1):
    roll = 0
    pitch = 0
    for i in range(sample_count):
        command = r"adb shell avacmd clb '{\"type\":\"system\",\"cmd\":\"print_roll_pitch\"}'"
        result_str = call_adb_command(command)
        print("ROLL-PITCH:", result_str)
        ret_data = json.loads(result_str)
        roll += ret_data["roll"]
        pitch += ret_data["pitch"]
        time.sleep(0.1)
    print("get_roll_pitch return ROLL:", roll / sample_count, ", PITCH:", pitch / sample_count)
    return roll / sample_count, pitch / sample_count

def dump_global_config():
    print(f"global_xtalk_pre_limit:{global_xtalk_pre_limit}")
    print(f"climb_up_pitch_limit:{climb_up_pitch_limit}")
    print(f"climb_tof_to_ground_distance:{climb_tof_to_ground_distance}")
    print(f"climb_rear_wheel_to_ground_distance:{climb_rear_wheel_to_ground_distance}")
    print(f"climb_rear_wheel_to_wheel_distance:{climb_rear_wheel_to_wheel_distance}")
    print(f"climb_tof_to_wheel_distance:{climb_tof_to_wheel_distance}")

def init_global_config():
    if check_file_exists(CONST_UNICAL_CONF):
        global climb_up_pitch_limit
        global climb_tof_to_ground_distance
        global climb_rear_wheel_to_ground_distance
        global climb_rear_wheel_to_wheel_distance
        global climb_tof_to_wheel_distance
        yaml_string = call_adb_command(f"adb shell cat {CONST_UNICAL_CONF}")
        data = yaml.safe_load(yaml_string)
        # 访问解析后的数据
        climb_up_pitch_limit = data["climb_up_pitch_limit"]
        climb_tof_to_ground_distance = data["climb_tof_to_ground_distance"]
        climb_rear_wheel_to_ground_distance = data["climb_rear_wheel_to_ground_distance"]
        climb_rear_wheel_to_wheel_distance = data["climb_rear_wheel_to_wheel_distance"]
        climb_tof_to_wheel_distance = data["climb_tof_to_wheel_distance"]

        dump_global_config()
        return True
    else:
        print(f"ERROR: {CONST_UNICAL_CONF} not found, please do 3-in-1 unical calibration first!!!")
        return False

def main():
    global power_mode
    if not init_global_config():
        return
    if len(sys.argv) < 3:
        print("Usage: 4cm_point_tof_test.py calib_xtalk calib_ground <value> calib_crossing <value>")
        return
    module_str = sys.argv[2]
    module_list = module_str.split(",")
    print(module_list)
    single_dtof = check_file_exists(file_path)

    if sys.argv[1] == 'calib_xtalk':
        print("start calibrate xtalk")
        # 触发一次开tof,避免ERP模式影响测试
        command_open_tof = r"adb shell avacmd mcubase '{\"type\":\"tof\",\"cmd\":\"open_tof\"}'"
        ret = call_adb_command(command_open_tof)
        print(ret)
        # 这里增加5秒延时，确保打开成功后再执行后续命令
        # 从产线标定发现，完整开完左前右前两个tof模组，耗时约3秒
        time.sleep(5)
        for module in module_list:
            calib_xtalk(module)
            time.sleep(0.5)
        sys.exit(0)
    elif sys.argv[1] == 'calib_ground_offset':
        print("start calibrate ground_offset")

        # 校验一下机器姿态，如果不平，说明要不没有做IMU标定，要不放的不平
        roll, pitch = get_roll_pitch(10)
        if abs(roll) >= 100 or abs(pitch) >= 100:
            print(f"{module_str} ERROR: fail to calibrate, the robot pose is not level, PLEASE make sure the robot is put level, and calibrate IMU before [{roll} {pitch}]vs[100]")
            sys.exit(0)
        # 把pitch缓存到机器内文件
        tmp_store_cmd = f'adb shell "echo -n {pitch} > {pitch_file_in_device}'
        call_adb_command(tmp_store_cmd)
        # 先校验一下
        depth = 0  # 单位mm，默认为0，即平地标定时，机器驱动轮贴地平面，和点tof打的平面一样高，如果点tof打地平面更矮，则为正值
        if len(sys.argv) >=4:
            depth = float(sys.argv[3])
            print("DEPTH:", depth)
        for module in module_list:
            calib_ground_offset(module, depth)
            time.sleep(0.5)
        sys.exit(0)
    elif sys.argv[1] == 'calib_crossing':
        print("start calibrate crossing gradient")
        depth = 0  # 单位mm，默认为0，即平地标定时，机器驱动轮贴地平面，和点tof打的平面一样高，如果点tof打地平面更矮，则为正值
        if len(sys.argv) >=4:
            depth = float(sys.argv[3])
            print("DEPTH:", depth)
        # 这一步需要读取IMU数值,小轮支起的时候，roll应该在1度以内，pitch应该在12 - 14 度
        roll, pitch = get_roll_pitch()
        if abs(roll) >= 100 or  not (climb_up_pitch_limit[0] <= pitch  and pitch <= climb_up_pitch_limit[1]):
            print(f"{module_str} ERROR: the robot pose is not valid, PLEASE make sure the robot is stand in good way, pose:[{roll}vs100] [{pitch}vs{climb_up_pitch_limit}]")
            sys.exit(0)
        # 撑起标定时,从机器内读出缓存的pitch
        tmp_store_cmd = f'adb shell cat {pitch_file_in_device}'
        cached_pitch = call_adb_command(tmp_store_cmd)
        try:
            cached_pitch = float(cached_pitch)
        except Exception as e:
            print(e)
            cached_pitch = 0
        patch_org = pitch
        # 将撑起标定时读到的pitch值减去缓存的pitch值
        # pitch -= cached_pitch
        # 根据IMU读数计算理论的点tof距离，进行标定
        target_org = math.tan(patch_org * 0.01 * 3.14159 / 180) * (climb_rear_wheel_to_wheel_distance + climb_tof_to_wheel_distance) + (climb_tof_to_ground_distance - climb_rear_wheel_to_ground_distance) + depth / math.cos(patch_org * 0.01 * 3.14159 / 180)
        target = math.tan(pitch * 0.01 * 3.14159 / 180) * (climb_rear_wheel_to_wheel_distance + climb_tof_to_wheel_distance) + (climb_tof_to_ground_distance - climb_rear_wheel_to_ground_distance) + depth / math.cos(pitch * 0.01 * 3.14159 / 180)
        print("PITCH:", pitch, ", target distance:", target, ", patch_org:", patch_org, ", target_org:", target_org)
        for module in module_list:
            calib_crossing(module, target)
            time.sleep(0.5)
        sys.exit(0)
    elif sys.argv[1] == 'check_ground_offset':
        print("start check ground_offset")

        # 校验一下机器姿态，如果不平，说明要不没有做IMU标定，要不放的不平
        roll, pitch = get_roll_pitch()
        if abs(roll) >= 100 or abs(pitch) >= 100:
            print(f"{module_str} ERROR: fail to calibrate, the robot pose is not level, PLEASE make sure the robot is put level, and calibrate IMU before [{roll} {pitch}] vs [100]")
            sys.exit(0)
        # 先校验一下
        depth = 0  # 单位mm，默认为0，即平地标定时，机器驱动轮贴地平面，和点tof打的平面一样高，如果点tof打地平面更矮，则为正值
        if len(sys.argv) >=4:
            depth = float(sys.argv[3])
            print("DEPTH:", depth)
        for module in module_list:
            change_range_mode(module, 170)
        power_mode = "low_power"
        time.sleep(1)
        for module in module_list:
            check_ground_offset(module, depth)
        
        time.sleep(0.1)
        for module in module_list:
            change_range_mode(module, 85)
        power_mode = "high_power"
        time.sleep(1)
        for module in module_list:
            check_ground_offset(module, depth)
        sys.exit(0)
    elif sys.argv[1] == 'check_crossing':
        print("start check crossing distance")
        depth = 0  # 单位mm，默认为0，即平地标定时，机器驱动轮贴地平面，和点tof打的平面一样高，如果点tof打地平面更矮，则为正值
        if len(sys.argv) >=4:
            depth = float(sys.argv[3])
            print("DEPTH:", depth)
        # 这一步需要读取IMU数值,小轮支起的时候，roll应该在1度以内，pitch应该在12 - 14 度
        roll, pitch = get_roll_pitch()
        if abs(roll) >= 100 or  not (climb_up_pitch_limit[0] <= pitch  and pitch <= climb_up_pitch_limit[1]):
            print(f"{module_str} ERROR: the robot pose is not valid, PLEASE make sure the robot is stand in good way, pose:[{roll}vs100] [{pitch}vs{climb_up_pitch_limit}]")
            sys.exit(0)

        # 撑起标定时,从机器内读出缓存的pitch
        tmp_store_cmd = f'adb shell cat {pitch_file_in_device}'
        cached_pitch = call_adb_command(tmp_store_cmd)
        try:
            cached_pitch = float(cached_pitch)
        except Exception as e:
            print(e)
            cached_pitch = 0
        patch_org = pitch
        # 将撑起标定时读到的pitch值减去缓存的pitch值
        # pitch -= cached_pitch

        # 根据IMU读数计算理论的点tof距离，进行标定
        target_org = math.tan(patch_org * 0.01 * 3.14159 / 180) * (climb_rear_wheel_to_wheel_distance + climb_tof_to_wheel_distance) + (climb_tof_to_ground_distance - climb_rear_wheel_to_ground_distance) + depth / math.cos(patch_org * 0.01 * 3.14159 / 180)
        target = math.tan(pitch * 0.01 * 3.14159 / 180) * (climb_rear_wheel_to_wheel_distance + climb_tof_to_wheel_distance) + (climb_tof_to_ground_distance - climb_rear_wheel_to_ground_distance) + depth / math.cos(pitch * 0.01 * 3.14159 / 180)
        print("PITCH:", pitch, ", target distance:", target)

        for module in module_list:
            change_range_mode(module, 170)
        power_mode = "low_power"
        time.sleep(1)
        for module in module_list:
            check_crossing(module, target)
        
        time.sleep(0.1)
        for module in module_list:
            change_range_mode(module, 85)
        power_mode = "high_power"
        time.sleep(1)
        for module in module_list:
            check_crossing(module, target)
        sys.exit(0)
    else:
        print("ERROR: Invalid cmd", sys.argv)

if __name__ == '__main__':
    print(f"VERSION:{VERSION}")
    main()
