"""标定板和标靶创建工具

需要在 Isaac Sim 初始化后导入。
"""

import math

import numpy as np
from omni.isaac.core.objects import FixedCuboid
from omni.isaac.core.utils.rotations import euler_angles_to_quat

from usd_utils import create_textured_board, create_textured_top_strip


def add_down_camera_intrinsic_target(world, stage, args, material):
    """添加下视相机内参标定台"""
    length = args.down_camera_target_length
    width = args.down_camera_target_width
    center_x = args.down_camera_target_x
    center_y = args.down_camera_target_y
    z_low = 0.012
    z_high = args.down_camera_target_max_height
    x_start = center_x - length / 2.0
    x_low_end = x_start + length * 0.22
    x_ramp_end = x_start + length * 0.74
    x_end = center_x + length / 2.0
    y_min = center_y - width / 2.0
    y_max = center_y + width / 2.0
    base_thickness = 0.006

    world.scene.add(FixedCuboid(
        prim_path="/World/Workshop/DownCameraIntrinsicTarget/Base",
        name="down_camera_intrinsic_target_base",
        position=np.array([center_x, center_y, base_thickness / 2.0]),
        scale=np.array([length + 0.08, width + 0.08, base_thickness]),
        color=np.array([0.045, 0.050, 0.052]),
    ))

    create_textured_top_strip(
        stage,
        "/World/Workshop/DownCameraIntrinsicTarget/CharucoRampSurface",
        [x_start, x_low_end, x_ramp_end, x_end],
        y_min,
        y_max,
        [z_low, z_low, z_high, z_high],
        material,
    )

    specs = [{
        "id": "down_camera_3d_charuco_ramp",
        "target_type": "down_camera_intrinsic",
        "pattern": "charuco",
        "purpose": "intrinsic_calibration_with_depth_and_pose_gradient",
        "center": [center_x, center_y, (z_low + z_high) / 2.0],
        "length_m": length,
        "width_m": width,
        "squares_x": args.down_camera_target_squares_x,
        "squares_y": args.down_camera_target_squares_y,
        "square_size_m": width / args.down_camera_target_squares_y,
        "min_height_m": z_low,
        "max_height_m": z_high,
        "recommended_drive_axis": "+x",
        "recommended_speed_mps": 0.10,
        "surfaces": [
            {
                "id": "low_flat_charuco",
                "type": "flat",
                "x_range_m": [x_start, x_low_end],
                "z_range_m": [z_low, z_low],
            },
            {
                "id": "continuous_slope_charuco",
                "type": "continuous_slope",
                "x_range_m": [x_low_end, x_ramp_end],
                "z_range_m": [z_low, z_high],
            },
            {
                "id": "high_flat_charuco",
                "type": "flat",
                "x_range_m": [x_ramp_end, x_end],
                "z_range_m": [z_high, z_high],
            },
        ],
    }]
    return specs


def lidar_2d_target_wall_slots(args):
    """计算 2D LiDAR 靶标墙面插槽位置"""
    if args.disable_2d_lidar_targets:
        return {}

    width = args.lidar_2d_target_width
    height = args.lidar_2d_target_height
    gap = 0.12
    edge_margin = 0.18
    bay_margin = 0.07
    bay_width = max(1.90, 2.0 * width + gap + 2.0 * bay_margin)
    bay_z_min = 0.0
    bay_z_max = args.room_height
    center_spacing = width + gap

    def group_offsets(span, side):
        usable_min = -span / 2.0 + edge_margin
        usable_max = span / 2.0 - edge_margin
        usable_width = max(0.0, usable_max - usable_min)
        effective_bay_width = min(bay_width, usable_width) if usable_width > 0.0 else bay_width
        if usable_width <= effective_bay_width:
            axis_min = usable_min
            axis_max = usable_max
        elif side == "max":
            axis_max = usable_max
            axis_min = axis_max - effective_bay_width
        else:
            axis_min = usable_min
            axis_max = axis_min + effective_bay_width
        group_center = (axis_min + axis_max) / 2.0
        vertical_offset = group_center - center_spacing / 2.0
        slope_offset = group_center + center_spacing / 2.0
        return vertical_offset, slope_offset, axis_min, axis_max

    front_vertical, front_slope, front_min, front_max = group_offsets(args.room_width, "min")
    _, _, right_min, right_max = group_offsets(args.room_length, "max")
    right_vertical = (right_min + right_max) / 2.0

    slots = {
        "front": {
            "vertical_offset": front_vertical,
            "slope_offset": front_slope,
            "reserved_axis_min": front_min,
            "reserved_axis_max": front_max,
            "reserved_z_min": bay_z_min,
            "reserved_z_max": bay_z_max,
        },
        "right": {
            "vertical_offset": right_vertical,
            "reserved_axis_min": right_min,
            "reserved_axis_max": right_max,
            "reserved_z_min": bay_z_min,
            "reserved_z_max": bay_z_max,
        },
    }
    return slots


def lidar_2d_checkerboard_reserved_zones(args):
    """获取 2D LiDAR 标靶预留区域"""
    slots = lidar_2d_target_wall_slots(args)
    reserved_zones = {}
    for wall, slot in slots.items():
        reserved_zones[wall] = [{
            "axis_min": slot["reserved_axis_min"],
            "axis_max": slot["reserved_axis_max"],
            "z_min": slot["reserved_z_min"],
            "z_max": slot["reserved_z_max"],
            "reason": "2d_lidar_extrinsic_corner_bay",
        }]
    return reserved_zones


def calibration_board_layout(args, board_width, board_height):
    """计算标定板布局"""
    wall_standoff = 0.015
    front_x = args.room_length / 2.0 - wall_standoff
    back_x = -args.room_length / 2.0 + wall_standoff
    left_y = args.room_width / 2.0 - wall_standoff
    right_y = -args.room_width / 2.0 + wall_standoff
    horizontal_gap = 0.18
    vertical_gap = 0.16
    side_margin = 0.35
    bottom_margin = 0.32
    top_margin = 0.28

    def axis_positions(span, item_size, margin, gap):
        available = span - 2.0 * margin
        count = max(1, int((available + gap) // (item_size + gap)))
        if count == 1:
            return [0.0]
        used = count * item_size + (count - 1) * gap
        start = -used / 2.0 + item_size / 2.0
        return [start + index * (item_size + gap) for index in range(count)]

    def z_positions():
        available = args.room_height - bottom_margin - top_margin
        count = max(1, int((available + vertical_gap) // (board_height + vertical_gap)))
        if count == 1:
            return [bottom_margin + board_height / 2.0]
        used = count * board_height + (count - 1) * vertical_gap
        start = bottom_margin + board_height / 2.0 + max(0.0, available - used) / 2.0
        return [start + index * (board_height + vertical_gap) for index in range(count)]

    zs = z_positions()
    front_back_offsets = axis_positions(args.room_width, board_width, side_margin, horizontal_gap)
    side_offsets = axis_positions(args.room_length, board_width, side_margin, horizontal_gap)
    reserved_zones = lidar_2d_checkerboard_reserved_zones(args)

    board_specs = []

    def overlaps_reserved_zone(wall, offset, z):
        board_axis_min = offset - board_width / 2.0
        board_axis_max = offset + board_width / 2.0
        board_z_min = z - board_height / 2.0
        board_z_max = z + board_height / 2.0
        for zone in reserved_zones.get(wall, []):
            axis_overlaps = board_axis_min < zone["axis_max"] and board_axis_max > zone["axis_min"]
            z_overlaps = board_z_min < zone["z_max"] and board_z_max > zone["z_min"]
            if axis_overlaps and z_overlaps:
                return True
        return False

    def add_wall_grid(wall, fixed_value, offsets, rotation_deg):
        for row, z in enumerate(zs):
            for col, offset in enumerate(offsets):
                if overlaps_reserved_zone(wall, offset, z):
                    continue
                if wall == "front":
                    position = [fixed_value, offset, z]
                elif wall == "back":
                    position = [fixed_value, offset, z]
                elif wall == "left":
                    position = [offset, fixed_value, z]
                else:
                    position = [offset, fixed_value, z]

                board_specs.append({
                    "id": f"{wall}_wall_r{row:02d}_c{col:02d}",
                    "wall": wall,
                    "mount": "flush",
                    "purpose": "wall_checkerboard_array",
                    "position": position,
                    "rotation_deg": rotation_deg,
                })

    add_wall_grid("front", front_x, front_back_offsets, [90, 0, 90])
    add_wall_grid("back", back_x, front_back_offsets, [90, 0, -90])
    add_wall_grid("left", left_y, side_offsets, [90, 0, 0])
    add_wall_grid("right", right_y, side_offsets, [90, 0, 180])

    return board_specs


def add_calibration_boards(stage, args, board_width, board_height, material):
    """添加墙面标定板"""
    board_specs = calibration_board_layout(args, board_width, board_height)
    for spec in board_specs:
        create_textured_board(
            stage,
            f"/World/Workshop/CalibrationBoards/{spec['id']}",
            board_width,
            board_height,
            spec["position"],
            spec["rotation_deg"],
            material,
        )
    return board_specs


def add_lidar_2d_panel(world, prim_path, name, position, scale, color, rotation_deg=None):
    """添加 2D LiDAR 标定面板"""
    orientation = None
    if rotation_deg is not None:
        orientation = np.array(euler_angles_to_quat(np.array(rotation_deg), degrees=True))
    world.scene.add(FixedCuboid(
        prim_path=prim_path,
        name=name,
        position=np.array(position),
        orientation=orientation,
        scale=np.array(scale),
        color=np.array(color),
    ))


def add_2d_lidar_calibration_targets(world, args):
    """添加 2D LiDAR 标定靶标"""
    thickness = args.lidar_2d_target_thickness
    width = args.lidar_2d_target_width
    height = args.lidar_2d_target_height
    bottom_z = args.lidar_2d_target_bottom_z
    center_z = bottom_z + height / 2.0
    angle_deg = args.lidar_2d_target_slope_angle_deg
    angle_rad = math.radians(angle_deg)
    sloped_length = height / max(math.cos(angle_rad), 1e-3)
    standoff = 0.06
    slots = lidar_2d_target_wall_slots(args)

    front_wall_x = args.room_length / 2.0
    right_wall_y = -args.room_width / 2.0

    front_base_x = front_wall_x - standoff - thickness / 2.0
    side_right_base_y = right_wall_y + standoff + thickness / 2.0

    front_slope_center_x = front_wall_x - standoff - height / 2.0 - thickness

    front_vertical_y = slots["front"]["vertical_offset"]
    front_slope_y = slots["front"]["slope_offset"]
    right_vertical_x = slots["right"]["vertical_offset"]

    vertical_color = [0.92, 0.82, 0.18]
    slope_color = [0.12, 0.65, 0.95]
    specs = []

    def add_spec(spec):
        specs.append(spec)
        return spec

    front_vertical = add_spec({
        "id": "front_vertical_reference_panel",
        "wall": "front",
        "type": "vertical_reference",
        "position": [front_base_x, front_vertical_y, center_z],
        "scale": [thickness, width, height],
        "rotation_deg": [0.0, 0.0, 0.0],
        "nominal_plane": "x = room_length/2 - standoff - thickness",
    })
    add_lidar_2d_panel(
        world,
        "/World/Workshop/Lidar2DCalibrationTargets/FrontVerticalReference",
        "front_vertical_reference_panel",
        front_vertical["position"],
        front_vertical["scale"],
        vertical_color,
    )

    front_slope = add_spec({
        "id": "front_45deg_height_encoding_panel",
        "wall": "front",
        "type": "height_encoding_slope",
        "position": [front_slope_center_x, front_slope_y, center_z],
        "scale": [thickness, width, sloped_length],
        "rotation_deg": [0.0, -angle_deg, 0.0],
        "slope_angle_deg": angle_deg,
        "height_to_range_sign": "higher_scan_plane_farther_from_front_wall",
    })
    add_lidar_2d_panel(
        world,
        "/World/Workshop/Lidar2DCalibrationTargets/FrontSlope45",
        "front_45deg_height_encoding_panel",
        front_slope["position"],
        front_slope["scale"],
        slope_color,
        front_slope["rotation_deg"],
    )

    right_vertical = add_spec({
        "id": "right_vertical_reference_panel",
        "wall": "right",
        "type": "vertical_reference",
        "position": [right_vertical_x, side_right_base_y, center_z],
        "scale": [width, thickness, height],
        "rotation_deg": [0.0, 0.0, 0.0],
        "nominal_plane": "y = -room_width/2 + standoff + thickness",
    })
    add_lidar_2d_panel(
        world,
        "/World/Workshop/Lidar2DCalibrationTargets/RightVerticalReference",
        "right_vertical_reference_panel",
        right_vertical["position"],
        right_vertical["scale"],
        vertical_color,
    )

    return specs


def add_floor_marker(world, prim_path, name, position, scale, color):
    """添加地面标记"""
    world.scene.add(FixedCuboid(
        prim_path=prim_path,
        name=name,
        position=np.array(position),
        scale=np.array(scale),
        color=np.array(color),
    ))


def add_calibration_floor_fixtures(world, args):
    """添加地面标定设施"""
    z = 0.012
    thickness = 0.012
    length = args.straight_track_length
    width = args.straight_track_width

    add_floor_marker(
        world,
        "/World/Workshop/CalibrationFixtures/StraightTrack",
        "straight_track",
        [0.0, 0.0, z],
        [length, width, thickness],
        [0.08, 0.12, 0.10],
    )
    add_floor_marker(
        world,
        "/World/Workshop/CalibrationFixtures/StraightCenterLine",
        "straight_center_line",
        [0.0, 0.0, z + thickness],
        [length, 0.035, thickness],
        [0.95, 0.95, 0.15],
    )
    for side, y in (("Left", width / 2.0), ("Right", -width / 2.0)):
        add_floor_marker(
            world,
            f"/World/Workshop/CalibrationFixtures/StraightBoundary{side}",
            f"straight_boundary_{side.lower()}",
            [0.0, y, z + thickness],
            [length, 0.025, thickness],
            [0.95, 0.78, 0.08],
        )

    stop_x = min(length / 2.0 - 0.35, args.room_length / 2.0 - 1.0)
    add_floor_marker(
        world,
        "/World/Workshop/CalibrationFixtures/StopAccuracyTarget",
        "stop_accuracy_target",
        [stop_x, 0.0, z + 2.0 * thickness],
        [0.6, 0.9, thickness],
        [0.88, 0.10, 0.10],
    )

    lateral_x = -min(length / 2.0 - 0.8, args.room_length / 2.0 - 1.2)
    add_floor_marker(
        world,
        "/World/Workshop/CalibrationFixtures/LateralMotionPad",
        "lateral_motion_pad",
        [lateral_x, 0.0, z + thickness],
        [0.75, 2.0, thickness],
        [0.10, 0.34, 0.88],
    )

    arc_center = np.array([0.0, -args.room_width * 0.20])
    arc_radius = min(args.arc_track_radius, args.room_width * 0.32)
    marker_count = 25
    for index in range(marker_count):
        theta = math.radians(20.0 + 140.0 * index / (marker_count - 1))
        x = arc_center[0] + arc_radius * math.cos(theta)
        y = arc_center[1] + arc_radius * math.sin(theta)
        add_floor_marker(
            world,
            f"/World/Workshop/CalibrationFixtures/ArcMarker_{index:02d}",
            f"arc_marker_{index:02d}",
            [x, y, z + 3.0 * thickness],
            [0.10, 0.10, thickness],
            [0.12, 0.75, 0.35],
        )

    add_floor_marker(
        world,
        "/World/Workshop/CalibrationFixtures/SensorCapturePad",
        "sensor_capture_pad",
        [0.0, args.room_width * 0.28, z + thickness],
        [1.2, 0.9, thickness],
        [0.42, 0.18, 0.78],
    )
