import os
import json
import math

# ===== 输入路径 =====
scene_dir = '/home/sig/sig/qianluo/qianluo/3DGS_VLN_Benchmark/test/839873'
map_file = os.path.join(scene_dir, '2D_Semantic_Map_839873_Complete_v3.json')
traj_file = os.path.join(scene_dir, 'trajectories_0033_839873_recollected_gvln.json')

# ===== 读取 2D Map 获取边界 =====
with open(map_file) as f:
    map_data = json.load(f)
all_y = [float(y) for inst in map_data for y, x in inst['mask_coords_m']]
all_x = [float(x) for inst in map_data for y, x in inst['mask_coords_m']]
min_y, max_y = min(all_y), max(all_y)
min_x, max_x = min(all_x), max(all_x)
print(f"Map bounds X:[{min_x}, {max_x}], Y:[{min_y}, {max_y}]")

# ===== 镜像翻转函数 =====
def flip_position(px, py, flip_x=True, flip_y=True, negate=False):
    """先做局部镜像，再可选对X、Y整体取负号"""
    if flip_x:
        px = (min_x + max_x) - px
    if flip_y:
        py = (min_y + max_y) - py
    if negate:
        px = -px
        py = -py
    return px, py

def yaw_from_quaternion(qx, qy, qz, qw):
    """从四元数取出绕Z的旋转角"""
    siny_cosp = 2 * (qw * qz + qx * qy)
    cosy_cosp = 1 - 2 * (qy * qy + qz * qz)
    return math.atan2(siny_cosp, cosy_cosp)

def quaternion_from_yaw(yaw):
    # 原函数按Z轴生成，现在其实只给我们临时计算用
    qz = math.sin(yaw / 2.0)
    qw = math.cos(yaw / 2.0)
    return 0.0, 0.0, qz, qw

# ===== 读取原轨迹 =====
with open(traj_file) as f:
    traj_data = json.load(f)

# ===== 开关控制 =====
flip_x = True   # 左右镜像
flip_y = True   # 上下镜像
negate_xy = True  # 最后整体取负号

# ===== 对每个点进行处理 =====
for scene in traj_data["scenes"]:
    for sample in scene["samples"]:
        points = sample["points"]
        for idx, pt in enumerate(points):
            pos = pt["position"]
            rot = pt["rotation"]
            px_old, py_old, pz = pos

            # 翻转+取负
            px_new, py_new = flip_position(px_old, py_old,
                                           flip_x=flip_x,
                                           flip_y=flip_y,
                                           negate=negate_xy)
            pt["position"] = [px_new, py_new, pz]

            # 旋转部分处理
            if idx == len(points) - 1:
                # 最后一个点直接恒等旋转
                pt["rotation"] = [0.0, 0.0, 0.0, 1.0]
            else:
                # 原来的 yaw 计算保持，用于得出标准 qz,qw
                yaw = yaw_from_quaternion(*rot)
                yaw_new = yaw + math.pi
                if yaw_new > math.pi:
                    yaw_new -= 2 * math.pi
                qx_tmp, qy_tmp, qz_tmp, qw_tmp = quaternion_from_yaw(yaw_new)

                # 将原qz放到x，并且取反，其它分量按你的规则设置
                qx = -qz_tmp
                qy = 0.0
                qz = 0.0
                qw = qw_tmp

                pt["rotation"] = [qx, qy, qz, qw]

# ===== 输出新json =====
out_file = traj_file.replace('.json', '_trans.json')
with open(out_file, 'w') as f:
    json.dump(traj_data, f, indent=2)

print(f"已保存轨迹到: {out_file}")