import json
import numpy as np
import matplotlib.pyplot as plt


def map2d_to_pose3d_with_flip(x2d, y2d, yaw2d=0.0, z_robot=0.4, action="",
                              x_min=None, x_max=None, y_min=None, y_max=None):
    '''
    x2d, y2d : 2D Map物理世界坐标 (单位米，带flip且单位和2D map一致)
    yaw2d    : 2D Map上的朝向（单位弧度，右手系，0为x正向）
    z_robot  : 机器人合理高度，默认0.4米
    action   : 可自定义字符串动作
    (x_min,x_max,y_min,y_max): 从Occupancy.json读取的地图世界边界
    '''
    assert x_min is not None and x_max is not None and y_min is not None and y_max is not None
    # 逆flip回原3D物理世界空间
    x3d = x_min + x_max - x2d
    y3d = y_min + y_max - y2d

    # yaw也要反变换：2D yaw在flip映射回3D时相对于world系按轴翻转
    # 这里示例：同时左右、上下flip后的yaw
    # x轴flip: yaw' = pi - yaw
    # y轴flip: yaw'' = -yaw'
    yaw3d = -(np.pi - yaw2d)

    return {
        "x": float(x3d),
        "y": float(y3d),
        "z": float(z_robot),
        "yaw": float(yaw3d),
        "action": action
    }


# Occupancy区间
with open('test/839920/occupancy.json') as f:
    meta = json.load(f)

x_min, y_min = meta['min'][:2]
x_max, y_max = meta['max'][:2]

# 示例 2D Map上的导航数据:
nav_path_2d = [
    # (x2d, y2d, yaw2d, action)  # 注意都是2D Map物理世界坐标，不是像素
    (1.23, -3.45, 0.0, "start"),
    (2.23, -3.45, 0.0, "move_forward"),
    (2.23, -2.45, np.pi/2, "turn_left"),
    (2.23, -1.45, np.pi/2, "move_forward"),
]

poses3d = [
    map2d_to_pose3d_with_flip(
        x2d, y2d, yaw2d, z_robot=0.4, action=act,
        x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max
    )
    for (x2d, y2d, yaw2d, act) in nav_path_2d
]

for pose in poses3d:
    print(pose)