import math
import numpy as np
import torch
import hydra
from scipy.spatial.transform import Rotation as R
import os
import xml.etree.ElementTree as ET
import gdown
import zipfile
import random
from PIL import Image
import importlib.util
import sys
from pathlib import Path


def load_policy_from_file(policy_path, config_overrides=None):
    policy_path = Path(policy_path)
    
    spec = importlib.util.spec_from_file_location("policy_module", policy_path)
    policy_module = importlib.util.module_from_spec(spec)
    sys.modules["policy_module"] = policy_module
    spec.loader.exec_module(policy_module)
    
    config = policy_module.get_config()
    
    if config_overrides:
        for key, value in config_overrides.items():
            config[key] = value
    
    policy_class = None
    for attr_name in dir(policy_module):
        attr = getattr(policy_module, attr_name)
        if isinstance(attr, type) and attr_name.endswith('Policy') and attr_name != 'BasePolicy':
            if hasattr(attr, '__module__') and attr.__module__ == 'policy_module':
                policy_class = attr
                break
    
    if policy_class is None:
        raise ValueError(f"No policy class found in {policy_path}")
    
    return policy_class(config=config)


def maybe_download():
    objects_path = os.path.join(os.path.dirname(__file__), "assets/objects")
    if not os.path.exists(objects_path):
        print("Downloading Objects...")
        download_and_extract_zip(
            gdrive_link="https://drive.google.com/file/d/1wDczBM6H3rQvFUGKe_vQINE5Bm2fBlU4",
            output_dir=os.path.join(os.path.dirname(__file__), "assets"),
        )
    return objects_path

def get_textures(root_folder=None):
    results = []
    if root_folder is None:
        root_folder = os.path.join(os.path.dirname(__file__), "assets/textures")

    for dirpath, _, filenames in os.walk(root_folder):
        for filename in filenames:
            if filename.endswith('.png'):
                full_path = os.path.join(dirpath, filename)
                rel_path = os.path.relpath(full_path, root_folder)
                results.append(rel_path)
    
    return results

def get_objects(root_folder=None):
    results = []
    if root_folder is None:
        root_folder = maybe_download()

    for dirpath, _, filenames in os.walk(root_folder):
        if "model.xml" in filenames:
            full_path = os.path.join(dirpath, "model.xml")
            parent_folder = os.path.basename(dirpath)
            results.append((parent_folder, full_path))

    results.append(("cabinet", "procedurally_generated"))
    results.append(("drawer", "procedurally_generated"))
    
    return results

def make_objects_manager(objects_set, np_random, shuffle=True):
    from egogym.managers import ObjectsManager
    objects = []
    available_objects = get_objects()
    available_objects = {obj[0]: obj[1] for obj in available_objects}
    for obj_name in objects_set:
        if obj_name in available_objects:
            objects.append((obj_name, available_objects[obj_name]))
    return ObjectsManager(objects, np_random, shuffle=shuffle)

def make_textures_manager(textures_set, np_random):
    from egogym.managers import TexturesManager
    textures = []
    available_textures = get_textures(root_folder=os.path.join(os.path.dirname(__file__), "assets/textures"))
    for tex_path in textures_set:
        for available_tex_name in available_textures:
            if tex_path in available_tex_name:
                textures.append(tex_path)
    return TexturesManager(textures, np_random)

def include_in_scene(scene_xml_content, include_file_path):
    root = ET.fromstring(scene_xml_content)
    new_include = ET.Element("include", file=include_file_path)
    includes = root.findall("include")
    if includes:
        index = list(root).index(includes[-1]) + 1
        root.insert(index, new_include)
    else:
        root.insert(0, new_include)
    return ET.tostring(root, encoding="unicode")

def download_and_extract_zip(gdrive_link, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    file_id = gdrive_link.split("/d/")[1].split("/")[0]
    download_url = f"https://drive.google.com/uc?id={file_id}"
    zip_path = os.path.join(output_dir, "temp_download.zip")

    gdown.download(download_url, zip_path, quiet=False)

    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(output_dir)

    os.remove(zip_path)

def euler_to_mat_xyz(rx, ry, rz):
    cx = torch.cos(rx)
    sx = torch.sin(rx)
    cy = torch.cos(ry)
    sy = torch.sin(ry)
    cz = torch.cos(rz)
    sz = torch.sin(rz)
    B = rx.shape[0]
    Rmats = torch.empty(B, 3, 3, device=rx.device)
    Rmats[:, 0, 0] = cy * cz
    Rmats[:, 0, 1] = -cy * sz
    Rmats[:, 0, 2] = sy
    Rmats[:, 1, 0] = cx * sz + cz * sx * sy
    Rmats[:, 1, 1] = cx * cz - sx * sy * sz
    Rmats[:, 1, 2] = -cy * sx
    Rmats[:, 2, 0] = sx * sz - cx * cz * sy
    Rmats[:, 2, 1] = cz * sx + cx * sy * sz
    Rmats[:, 2, 2] = cx * cy
    return Rmats

def make_pose(pos, rot):
    pose = np.eye(4)
    pose[:3, 3] = pos
    pose[:3, :3] = rot.as_matrix()
    return pose

def get_pos_rot(data, name, type):
    type_lower = type.lower()
    if type_lower == "body-control":
        pos = data.body(name).xpos.copy()
        rot_mat = data.body(name).xmat.copy().reshape((3, 3))
        return pos, rot_mat
    if type_lower == "body":
        pos = data.body(name).xipos.copy()
        rot_mat = data.body(name).xmat.copy().reshape((3, 3))
        return pos, rot_mat
    elif type_lower == "site":
        pos = data.site(name).xpos.copy()
        rot_mat = data.site(name).xmat.copy().reshape((3, 3))
        return pos, rot_mat
    elif type_lower == "geom":
        pos = data.geom(name).xpos.copy()
        rot_mat = data.geom(name).xmat.copy().reshape((3, 3))
        return pos, rot_mat
    elif type_lower == "camera":
        cam_id = data.model.camera(name).id
        pos = data.cam_xpos[cam_id].copy()
        rot_mat = data.cam_xmat[cam_id].copy().reshape((3, 3))
        return pos, rot_mat
    else:
        raise ValueError(
            f'type_ must be one of "geom", "body", or "site", got "{type}".'
        )

def get_pose(data, name, type):
    pos, rot_mat = get_pos_rot(data, name, type)
    pose = np.eye(4)
    pose[:3, 3] = pos
    pose[:3, :3] = rot_mat
    return pose

def position_sampler(
    np_random,
    num_samples,
    box_half_ranges,
    min_distance,
    max_total_tries=100,
    start_position=None,
    thickness=0.1,
):
    hx, hy = map(float, box_half_ranges)
    samples = []
    samples_arr = []
    cell = float(min_distance) / np.sqrt(2.0)
    ox, oy = -hx, -hy
    grid = {}

    def to_cell_idx(p):
        return (
            int(np.floor((p[0] - ox) / cell)),
            int(np.floor((p[1] - oy) / cell)),
        )

    def is_far_enough(p):
        cx, cy = to_cell_idx(p)
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                key = (cx + dx, cy + dy)
                j = grid.get(key)
                if j is None:
                    continue
                q = samples_arr[j]
                if (p[0] - q[0]) ** 2 + (p[1] - q[1]) ** 2 < min_distance**2:
                    return False
        return True

    tries = 0
    while len(samples) < num_samples and tries < max_total_tries:
        tries += 1
        px = np_random.uniform(-hx, hx)
        py = np_random.uniform(-hy, hy)
        p = (px, py)
        if not samples_arr:
            samples.append([px, py])
            samples_arr.append(np.array(p, dtype=float))
            grid[to_cell_idx(p)] = 0
            continue
        if is_far_enough(p):
            idx = len(samples)
            samples.append([px, py])
            samples_arr.append(np.array(p, dtype=float))
            grid[to_cell_idx(p)] = idx

    if start_position is None:
        return samples

    A = np.array(start_position, dtype=float)
    blocked = [False] * len(samples)
    eps = 1e-9

    S = np.array(samples, dtype=float)
    for i in range(len(samples)):
        B = S[i]
        AB = B - A
        ab2 = float(AB[0] * AB[0] + AB[1] * AB[1])
        if ab2 <= eps:
            blocked[i] = False
            continue
        for j in range(len(samples)):
            if j == i:
                continue
            P = S[j]
            t = float(np.dot(P - A, AB) / ab2)
            if t <= eps or t >= 1.0 - eps:
                continue
            C = A + t * AB
            if float(np.sum((P - C) ** 2)) <= thickness * thickness:
                blocked[i] = True
                break

    if blocked[0]:
        for i in range(1, len(blocked)):
            if not blocked[i]:
                samples[0], samples[i] = samples[i], samples[0]
                blocked[0], blocked[i] = blocked[i], blocked[0]
                break
    return samples


def add_xml_to_scene(original_xml_string, new_xml_string):
    """
    Merges new_xml_string into original_xml_string by adding elements from corresponding sections.
    Elements from new XML's worldbody go under original's worldbody, assets under assets, etc.
    """

    original_root = ET.fromstring(original_xml_string)
    new_root = ET.fromstring(new_xml_string)

    for new_section in new_root:
        section_name = new_section.tag

        original_section = original_root.find(section_name)

        if original_section is not None:
            for child in new_section:
                original_section.append(child)
        else:
            original_root.append(new_section)

    return ET.tostring(original_root, encoding="unicode")

def make_handle_xml(
    handle_type="pullbar",
    handle_name="handle",
    position=(0, 0, 0),
    euler_orientation=(0, 0, 0),
    handle_radius=0.009,
    handle_length=0.05,
    handle_material="M_hinge_metal",
    mass=0.1,
    friction_geom="0 0 0",
    hinge_side="left",
    np_random=None,
):
    pos_x, pos_y, pos_z = position
    euler_x, euler_y, euler_z = euler_orientation

    if handle_type == "knob":
        if handle_name == "handle":
            knob_geom_name = "handle"
            knob_geom_viz_name = "handle_viz"
        else:
            knob_geom_name = f"knob_geo_{hinge_side}"
            knob_geom_viz_name = f"knob_geom_viz_{hinge_side}"
        friction_geom = "5 5 5"
        handle_body_xml = f"""
        <body name="{handle_name}_body" pos="{pos_x + (-0.02 if hinge_side == "left" else 0.02)} {pos_y + 0.07} {pos_z}" euler="{euler_x} {euler_y} 0">
                    <body name="knob">
                    <geom name="{knob_geom_name}" mesh="knob" class="visual" material="{handle_material}"/>
                    <geom name="{knob_geom_viz_name}" mesh="knob_collision_0" rgba="0.12319198224928918 0.8185897489755214 0.9473380251850289 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_1" rgba="0.9952048472348035 0.9587846043226812 0.44813684566262046 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_2" rgba="0.5562189501844684 0.3964830683252689 0.24776822693559275 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_3" rgba="0.5665588588501549 0.4670123537790013 0.1620234740959262 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_4" rgba="0.6170081600309119 0.6175499636469116 0.07870231532672867 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_5" rgba="0.5008884031488108 0.5029776170199916 0.7335611007545255 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_6" rgba="0.6962939670098736 0.9990710674845018 0.1259585452380786 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_7" rgba="0.25804299152558585 0.3180969182215061 0.9745095357400368 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_8" rgba="0.012626418834853048 0.6472133449453128 0.6250494987727182 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    <geom mesh="knob_collision_9" rgba="0.5557254592689047 0.6025541199284764 0.902390111497001 1" class="collision" friction="{friction_geom}" mass="{mass}"/>
                    </body>
                  </body>"""
        scale = np_random.uniform(0.012, 0.0135)
        base_path = os.path.join(os.path.dirname(__file__), "assets/meshes")
        mesh_asset_xml = f"""
        <mesh file="{base_path}/knob/knob.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_0.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_1.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_2.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_3.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_4.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_5.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_6.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_7.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_8.obj" scale="{scale} {scale} {scale}"/>
        <mesh file="{base_path}/knob/knob_collision_9.obj" scale="{scale} {scale} {scale}"/>
        """

    elif handle_type == "pullbar":
        friction_geom = "5 5 5"
        handle_body_xml = f"""
        <body name="{handle_name}_body" pos="{pos_x} {pos_y} {pos_z}" euler="{euler_x} {euler_y} {euler_z}">
            <geom material="{handle_material}" pos="0 0 0" size="{handle_radius} {handle_radius} {handle_length}" type="box" />
            <geom name="{handle_name}" material="{handle_material}" pos="0 0 0" size="{handle_radius} {handle_radius} {handle_length}" type="box"/>
            <geom material="{handle_material}" pos="-0.04 0 {handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" mass="{mass / 10}"/>
            <geom material="{handle_material}" pos="-0.04 0 -{handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" mass="{mass / 10}"/>
            <geom material="{handle_material}" pos="-0.04 0 {handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" conaffinity="0" contype="0"/>
            <geom material="{handle_material}" pos="-0.04 0 -{handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" conaffinity="0" contype="0"/>
        </body>"""
        mesh_asset_xml = ""

    else:
        friction_geom = "5 5 5"
        handle_body_xml = f"""
        <body name="{handle_name}_body" pos="{pos_x} {pos_y} {pos_z}" euler="{euler_x} {euler_y} {euler_z}">
            <geom material="{handle_material}" pos="0 0 0" size="{handle_radius} {handle_length}" type="cylinder" friction="{friction_geom}"/>
            <geom name="{handle_name}" material="{handle_material}" pos="0 0 0" size="{handle_radius} {handle_length}" type="cylinder" mass="{mass}" friction="{friction_geom}"/>
            <geom material="{handle_material}" pos="-0.04 0 {handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" mass="{mass / 10}"/>
            <geom material="{handle_material}" pos="-0.04 0 -{handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" mass="{mass / 10}"/>
            <geom material="{handle_material}" pos="-0.04 0 {handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" conaffinity="0" contype="0"/>
            <geom material="{handle_material}" pos="-0.04 0 -{handle_length - 0.03}" euler="0 1.57 0" size="{handle_radius / 1.2} {0.04} {0.021}" type="cylinder" conaffinity="0" contype="0"/>
        </body>"""
        mesh_asset_xml = ""

    return handle_body_xml, mesh_asset_xml


def make_cabinet_xml(np_random=None):
    if np_random is None:
        np_random = np.random
    
    hinge_side = np_random.choice(["left", "right"])
    holder_type = np_random.choice(["handle", "pullbar"])
    double_doors = np_random.choice([0, 1])
    
    if double_doors:
        width = np_random.uniform(0.4, 0.8)
    else:
        width = np_random.uniform(0.2, 0.5)
    
    height = np_random.uniform(0.4, 0.8)
    depth = np_random.uniform(0.1, 0.3)
    handle_position = np_random.uniform(0.5, 0.8)
    handle_length = np_random.uniform(0.05, 0.1)
    handle_radius = np_random.uniform(0.004, 0.006)
    
    wall_thickness = 0.01
    door_gap = 0.00
    add_wall = 1
    handle_offset_y = 0.05
    handle_offset_z = 0.00
    door_mass = 0.001
    friction_geom = "3 3 3" if holder_type == "knob" else "0 0 0"
    door_gap_between = 0.005
    texrepeat_xy = (5, 5)
    paint_rgba = (1.0, 1.0, 1.0, 1.0)
    collision_rgba = (0.3, 0.3, 1.0, 0.5)

    base_path = os.path.join(os.path.dirname(__file__), "assets/textures")
    door_texture_png = np_random.choice([f"{base_path}/door/{i}.png" for i in range(2)])
    cabinet_texture_png = np_random.choice([f"{base_path}/door/{i}.png" for i in range(2)])
    handle_texture_png = np_random.choice([f"{base_path}/handle/{i}.png" for i in range(4)])
    floor_texture_png = np_random.choice([f"{base_path}/tiles/{i}.png" for i in range(3)])
    wall_texture_png = np_random.choice([f"{base_path}/wall/{i}.png" for i in range(6)])

    hx = width / 2.0
    hy = depth / 2.0
    hz = height / 2.0
    t = wall_thickness

    interior_width = width - 2 * t
    door_thickness_y = t

    if double_doors:
        door_half_x = (hx - t / 2.0) / 2.0
    else:
        door_half_x = hx - t / 2.0

    door_half_y = door_thickness_y / 2.0
    door_half_z = hz - t / 2.0

    side_x = hx - t / 2.0
    top_z = hz - t / 2.0
    bottom_z = -hz + t / 2.0
    back_y = hy - t / 2.0

    side_half_x = t / 2.0
    side_half_y = (depth - t) / 2.0
    side_half_z = hz

    tb_half_x = interior_width / 2.0
    tb_half_y = (depth - t) / 2.0
    tb_half_z = t / 2.0

    back_half_x = hx - t / 2.0
    back_half_y = t / 2.0
    back_half_z = hz - t / 2.0

    repx, repy = texrepeat_xy
    wood_texrepeat = f"{repx} {repy}"

    paint_r, paint_g, paint_b, paint_a = paint_rgba
    col_r, col_g, col_b, col_a = collision_rgba

    wall_thickness_back = 0.02
    wall_y_pos = hy + wall_thickness_back / 2.0

    wall_section = ""
    if add_wall:
        wall_section = f"""
        <!-- Large wall behind cabinet -->
        <body name="front_wall">    
                <geom name="wall" pos="0 {wall_y_pos} 0" size="0 0 0.05" euler="1.57 0 0" type="plane" material="M_wall"/>
          </body>
        """

    def build_holder(
        hinge_side,
        holder_type,
        handle_x,
        handle_y,
        handle_z,
        friction_geom,
        door_mass,
        handle_radius,
        handle_length,
        handle_name,
    ):
        # Use the new reusable handle function
        handle_body_xml, mesh_asset_xml = make_handle_xml(
            handle_type=holder_type,
            handle_name=handle_name,
            position=(handle_x, handle_y, handle_z),
            euler_orientation=(0, 0, -1.57),
            handle_radius=handle_radius,
            handle_length=handle_length,
            handle_material="M_hinge_metal",
            mass=0.1,
            friction_geom=friction_geom,
            hinge_side=hinge_side,
            np_random=np_random,
        )

        return handle_body_xml, mesh_asset_xml

    def make_single_door(
        name_suffix, x_offset, hinge_side, joint_name, handle_name, door_width_half
    ):
        hinge_sign = +1.0 if hinge_side.lower() == "right" else -1.0

        # For double doors, adjust hinge position so doors are flush with cabinet sides and add gap
        if double_doors:
            door_side_offset = hx - t / 2.0
            gap = door_gap_between / 2.0
            if hinge_side.lower() == "left":
                door_root_x = -door_side_offset - gap
            else:
                door_root_x = door_side_offset + gap
        else:
            door_root_x = hinge_sign * side_x

        door_root_z = 0.0
        door_center_x = -hinge_sign * door_width_half
        door_center_y = -(t / 2.0)

        if double_doors:
            handle_x = door_center_x - hinge_sign * (door_width_half - 0.06)
        else:
            handle_x = door_center_x - hinge_sign * (door_width_half - 0.05)
        handle_y = -door_half_y - handle_offset_y
        door_bottom = -door_half_z
        door_top = door_half_z
        handle_z = (
            door_bottom + handle_position * (door_top - door_bottom) + handle_offset_z
        )

        if hinge_side.lower() == "right":
            jmin, jmax = 0.0, 1.57
        else:
            jmin, jmax = -1.57, 0.0

        holder_xml, mesh_asset = build_holder(
            hinge_side,
            holder_type,
            handle_x,
            handle_y,
            handle_z,
            friction_geom,
            door_mass,
            handle_radius,
            handle_length,
            handle_name,
        )

        return (
            f"""
            <body name="hingedoor{name_suffix}" pos="{door_root_x} {-hy - door_gap} {door_root_z}">
                <joint axis="0 0 1" name="{joint_name}" range="{jmin} {jmax}"  damping="2" frictionloss="2" armature=".01"/>
                {holder_xml}
                <body name="door_body{name_suffix}">
                    <geom material="M_door_wood" pos="{door_center_x} {door_center_y} 0" size="{door_width_half} {door_half_z} {door_half_y}" euler="1.57 0 0" type="box"/>
                </body>
                <site type="sphere" name="boxdoor_site{name_suffix}" pos="{handle_x} {handle_y - 0.018} {handle_z}" size=".01" group="3" rgba="1 0 0 1"/>
            </body>
        """,
            mesh_asset,
        )

    def make_single_cabinet(name_suffix, x_offset, hinge_side, joint_name, handle_name):
        hinge_sign = +1.0 if hinge_side.lower() == "right" else -1.0
        door_root_x = hinge_sign * side_x
        door_root_z = 0.0
        door_center_x = -hinge_sign * (door_half_x)
        door_center_y = -(t / 2.0)

        handle_x = door_center_x - hinge_sign * (door_half_x - 0.05)
        handle_y = -door_half_y - handle_offset_y
        door_bottom = -door_half_z
        door_top = door_half_z
        handle_z = (
            door_bottom + handle_position * (door_top - door_bottom) + handle_offset_z
        )

        if hinge_side.lower() == "right":
            jmin, jmax = 0.0, 1.57
        else:
            jmin, jmax = -1.57, 0.0

        holder_xml, mesh_asset = build_holder(
            hinge_side,
            holder_type,
            handle_x,
            handle_y,
            handle_z,
            friction_geom,
            door_mass,
            handle_radius,
            handle_length,
            handle_name,
        )

        return (
            f"""
        <body name="cabinet_object{name_suffix}" childclass="boxcabinet" pos="{x_offset} 0 {hz}">
            <body name="back_panel{name_suffix}">
                <geom material="M_cabinet_paint" pos="-{side_x} {-t / 4.0} 0" size="{side_half_x} {side_half_y} {side_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="{side_x} {-t / 4.0} 0" size="{side_half_x} {side_half_y} {side_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="0 {-t / 4.0} {top_z}" size="{tb_half_x} {tb_half_y} {tb_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="0 {-t / 4.0} {bottom_z}" size="{tb_half_x} {tb_half_y} {tb_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="0 {back_y} 0" size="{back_half_x} {back_half_y} {back_half_z}" type="box" conaffinity="0" contype="0"/>
            </body>

            <body name="hingedoor{name_suffix}" pos="{door_root_x} {-hy - door_gap} {door_root_z}">
                <joint axis="0 0 1" name="{joint_name}" range="{jmin} {jmax}" damping="2" frictionloss="2" armature=".01"/>
                {holder_xml}
                <body name="door_body{name_suffix}">
                    <geom material="M_door_wood" pos="{door_center_x} {door_center_y} 0" size="{door_half_x} {door_half_z} {door_half_y}" euler="1.57 0 0" type="box"/>
                </body>
                <site type="sphere" name="boxdoor_site{name_suffix}" pos="{handle_x} {handle_y - 0.018} {handle_z}" size=".01" group="3" rgba="1 0 0 1"/>
            </body>
        </body>
        """,
            mesh_asset,
        )

    cabinets_xml, extra_meshes = "", ""

    if double_doors:
        # Randomly pick which door gets 'task_joint' and 'handle_1'
        if np_random.uniform() > 0.5:
            left_joint = "task_joint"
            right_joint = "task_joint_2"
        else:
            left_joint = "task_joint_2"
            right_joint = "task_joint"

        left_handle = "handle" if left_joint == "task_joint" else "handle_2"
        right_handle = "handle" if right_joint == "task_joint" else "handle_2"

        left_door_xml, mesh1 = make_single_door(
            "_L", -door_half_x, "left", left_joint, left_handle, door_half_x
        )
        right_door_xml, mesh2 = make_single_door(
            "_R", door_half_x, "right", right_joint, right_handle, door_half_x
        )

        cabinets_xml = f"""
        <body name="cabinet_object" childclass="boxcabinet" pos="0 0 {hz}">
            <body name="back_panel">
                <geom material="M_cabinet_paint" pos="-{side_x} {-t / 4.0} 0" size="{side_half_x} {side_half_y} {side_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="{side_x} {-t / 4.0} 0" size="{side_half_x} {side_half_y} {side_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="0 {-t / 4.0} {top_z}" size="{tb_half_x} {tb_half_y} {tb_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="0 {-t / 4.0} {bottom_z}" size="{tb_half_x} {tb_half_y} {tb_half_z}" type="box" conaffinity="0" contype="0"/>
                <geom material="M_cabinet_paint" pos="0 {back_y} 0" size="{back_half_x} {back_half_y} {back_half_z}" type="box" conaffinity="0" contype="0"/>
            </body>
            {left_door_xml}
            {right_door_xml}
        </body>
        """
        extra_meshes = mesh1 + mesh2

        exclude_doors_xml = """
        <contact> <exclude body1="hingedoor_L" body2="hingedoor_R"/> <exclude body1="door_body_L" body2="door_body_R"/>
        
        <exclude body1="left_collisions" body2="door_body_L"/>
        <exclude body1="left_collisions" body2="door_body_R"/>
        <exclude body1="right_collisions" body2="door_body_L"/>
        <exclude body1="right_collisions" body2="door_body_R"/>

        </contact>

        """
    else:
        cabinets_xml, extra_meshes = make_single_cabinet(
            "", 0, hinge_side, "task_joint", "handle"
        )
        exclude_doors_xml = """
        <contact>
        <exclude body1="left_collisions" body2="door_body"/>
        <exclude body1="right_collisions" body2="door_body"/>



        </contact>

        """

     
    xml = f"""
            <mujoco model="hinge cabinet">
            <compiler angle="radian"/>
                <asset>
                    {extra_meshes}
                    <texture name="T_door_wood" type="2d" file="{door_texture_png}"/>
                    <texture name="T_cabinet_wood" type="2d" file="{cabinet_texture_png}"/>
                    <texture name="T_hinge_metal" type="2d" file="{handle_texture_png}"/>
                    <texture name="T_floor" type="2d" file="{floor_texture_png}"/>
                    <texture name="T_wall" type="2d" file="{wall_texture_png}"/>
                    <material name="M_door_wood" texture="T_door_wood" texrepeat="5 5" reflectance="0.0" shininess=".1" texuniform="true"/>
                    <material name="M_cabinet_wood" texture="T_cabinet_wood" texrepeat="{wood_texrepeat}" reflectance="0.0" shininess=".1" texuniform="true"/>
                    <material name="M_hinge_metal" texture="T_hinge_metal" texrepeat="4 4" reflectance="0.6" shininess=".6" texuniform="true"/>
                    <material name="M_floor" texture="T_floor" texrepeat="3 3" reflectance="0.0" shininess=".1" texuniform="true"/>
                    <material name="M_wall" texture="T_wall" texrepeat="4 4" reflectance="0.0" shininess=".1" texuniform="true"/>
                    <material name="M_cabinet_paint" texture="T_cabinet_wood" rgba="{paint_r} {paint_g} {paint_b} {paint_a}" reflectance="0.0" shininess=".01" texuniform="true"/>
                    <material name="hinge_collision_blue" rgba="{col_r} {col_g} {col_b} {col_a}" shininess="0" specular="0"/>
                </asset>
                <default>
                    <default class="boxcabinet">
                        <joint damping="2" frictionloss="2" limited="true"/>
                        <geom group="1" material="M_cabinet_wood" type="mesh"/>
                        <default class="box_collision">
                            <geom conaffinity="1" condim="6" contype="1" group="4" material="hinge_collision_blue"/>
                        </default>
                    </default>
                </default>
                <worldbody>
                    {wall_section}
                    {cabinets_xml}
                
                    <body name="floor">
                        <geom name="floor" pos="0 0 0" size="0 0 0.05" euler="0 0 0" type="plane" material="M_floor"/>
                    </body>
                </worldbody>
                {exclude_doors_xml}
            </mujoco>
            """
    return xml

def make_drawer_xml(np_random=None):
    if np_random is None:
        np_random = np.random
    
    handle_type = np_random.choice(["handle", "pullbar"])
    height_per_drawer = np_random.uniform(0.05, 0.15)
    handle_radius = np_random.uniform(0.005, 0.01)
    width = np_random.uniform(0.2, 0.4)
    depth = np_random.uniform(0.15, 0.3)
    handle_length = np_random.uniform(0.07, max(0.05, width / 1.2))
    num_drawers = np_random.integers(1, 5)
    
    position = (0, 0, 0)
    rotation = (0, 0, -1.57)
    handle_protrusion = 0.05
    thickness = 0.005
    cover_color = (0.3, 0.2, 0.05, 1.0)
    drawer_color = (0.4, 0.25, 0.1, 1.0)
    handle_color = (0.2, 0.2, 0.2, 1.0)
    texrepeat_xy = (3, 3)
    add_wall = 1
    joint_damping = 2

    base_path = os.path.join(os.path.dirname(__file__), "assets/textures")
    cover_texture_png = np_random.choice([f"{base_path}/door/{i}.png" for i in range(3)])
    drawer_texture_png = np_random.choice([f"{base_path}/door/{i}.png" for i in range(3)])
    door_texture_png = np_random.choice([f"{base_path}/door/{i}.png" for i in range(3)])
    handle_texture_png = np_random.choice([f"{base_path}/handle/{i}.png" for i in range(4)])
    floor_texture_png = np_random.choice([f"{base_path}/tiles/{i}.png" for i in range(3)])
    wall_texture_png = np_random.choice([f"{base_path}/wall/{i}.png" for i in range(6)])

    half_total_height = num_drawers * height_per_drawer + thickness * (num_drawers + 1)
    total_height = half_total_height * 2

    adjusted_position = (position[0], position[1], thickness)

    repx, repy = texrepeat_xy
    texrepeat_str = f"{repx} {repy}"

    cover_r, cover_g, cover_b, cover_a = cover_color
    drawer_r, drawer_g, drawer_b, drawer_a = drawer_color
    handle_r, handle_g, handle_b, handle_a = handle_color

    base_template = f"""<mujoco model="drawer">
    <asset>
        <texture name="T_cover" type="2d" height="1" width="1" file="{cover_texture_png}"/>
        <texture name="T_drawer" type="2d" height="1" width="1" file="{drawer_texture_png}"/>
        <texture name="T_door" type="2d" height="1" width="1" file="{door_texture_png}"/>
        <texture name="T_handle" type="2d" height="1" width="1" file="{handle_texture_png}"/>
        <texture name="T_floor" type="2d" height="2" width="2" file="{floor_texture_png}"/>
        <texture name="T_wall" type="2d" height="2" width="1" file="{wall_texture_png}"/>
    
    <material name="M_cover" texture="T_cover" texrepeat="{texrepeat_str}" reflectance="0.0" shininess=".1" texuniform="false"/>
    <material name="M_drawer" texture="T_drawer" texrepeat="{texrepeat_str}" reflectance="0.0" shininess=".1" texuniform="false"/>
    <material name="M_door" texture="T_door" texrepeat="{texrepeat_str}" reflectance="0.0" shininess=".1" texuniform="false"/>
    <material name="M_handle" texture="T_handle" texrepeat="{texrepeat_str}" reflectance="0.0" shininess=".1" texuniform="false"/>
    <material name="M_floor" texture="T_floor" texrepeat="5 5" reflectance="0.0" shininess=".1" texuniform="false"/>
    <material name="M_wall" texture="T_wall" texrepeat="2 2" reflectance="0.0" shininess=".1" texuniform="false"/>
  </asset>

    <worldbody>
      
        <body name="drawer_object" pos="{adjusted_position[0]} {adjusted_position[1]} {adjusted_position[2]}" euler="{rotation[0]} {rotation[1]} {rotation[2]}">
            <body name="outer_frame">
      <geom size="{depth + thickness * 3} {thickness} {half_total_height}" pos="0 -{width + thickness * 2} {half_total_height}" type="box" material="M_cover" />
      <geom size="{depth + thickness * 3} {thickness} {half_total_height}" pos="0 {width + thickness * 2} {half_total_height}" type="box" material="M_cover" />
      <geom size="{depth + thickness * 3} {width + thickness * 3} {thickness}" pos="0 0 {thickness}" type="box" material="M_cover" />
      <geom size="{depth + thickness * 3} {width + thickness * 3} {thickness}" pos="0 0 {total_height}" type="box" material="M_cover" />
      <geom size="{thickness} {width + thickness * 3} {half_total_height}" pos="-{depth + thickness * 2} 0 {half_total_height}" type="box" material="M_cover" />

        </body>
      """

    drawer_bodies = ""

    task_drawer_index = np_random.integers(0, num_drawers)

    for i in range(num_drawers):
        drawer_name = f"drawer_{i}"

        if i == task_drawer_index:
            joint_name = "task_joint"
            handle_geom_name = "handle"
        else:
            joint_name = f"drawer_{i}_slide"
            if i == 1 and task_drawer_index != 1:
                handle_geom_name = "handle_1_normal"
            else:
                handle_geom_name = f"handle_{i}"

        z_pos = height_per_drawer * (2 * i + 1) + thickness * (2 * i + 3)
        collision_thickness = thickness * 0.8
        collision_size_adjust = thickness * 0.8

        drawer_body = f"""
      <body name="{drawer_name}" pos="{thickness * 3} 0 {z_pos}">
        <joint name="{joint_name}" type="slide" axis="1 0 0" limited="true" range="0 {
            depth * 2
        }"  damping="10" frictionloss="10" armature=".01"   />
        
        <geom size="{depth} {thickness} {height_per_drawer}" type="box" pos="0 {
            width - thickness
        } 0" material="M_drawer" conaffinity="0" contype="0"/>
        <geom size="{depth} {thickness} {height_per_drawer}" type="box" pos="0 -{
            width - thickness
        } 0" material="M_drawer" conaffinity="0" contype="0"/>
        <geom size="{height_per_drawer} {width} {thickness}" type="box" pos="{
            depth - thickness
        } 0 0" material="M_drawer" conaffinity="0" contype="0" euler="1.57 1.57 1.57"/>
        <geom size="{height_per_drawer} {width} {thickness}" type="box" pos="{
            depth - thickness
        } 0 0" material="M_drawer" conaffinity="0" contype="0" euler="1.57 1.57 1.57"/>
          <geom size="{height_per_drawer} {width} {thickness}" type="box" pos="{
            depth - thickness
        } 0 0" material="M_drawer" euler="1.57 1.57 1.57"/>
        
        <geom size="{depth * 0.9} {width * 0.9} {
            collision_thickness
        }" type="box" pos="0 0 -{height_per_drawer - thickness}" mass="0.0001"/>
        <geom size="{depth * 0.9} {collision_size_adjust} {
            height_per_drawer * 0.8
        }" type="box" pos="0 {width - thickness} 0" mass="0.0001"/>
        <geom size="{depth * 0.9} {collision_size_adjust} {
            height_per_drawer * 0.8
        }" type="box" pos="0 -{width - thickness} 0" mass="0.0001"/>
        <geom size="{collision_size_adjust} {width * 0.9} {
            height_per_drawer * 0.8
        }" type="box" pos="{depth - thickness} 0 0" mass="0.0001"/>
        <geom size="{collision_size_adjust} {width * 0.9} {
            height_per_drawer * 0.8
        }" type="box" pos="-{depth - thickness} 0 0" mass="0.0001"/>
        
        {
            make_handle_xml(
                handle_type=handle_type,
                handle_name=handle_geom_name,
                position=(depth + handle_protrusion, 0, 0),
                euler_orientation=(1.57, 0, 0),
                handle_radius=handle_radius,
                handle_length=handle_length,
                handle_material="M_handle",
                mass=0.1,
                friction_geom="0 0 0",
            )[0]
        }
      </body>"""

        drawer_bodies += drawer_body

    wall_y_pos = -(depth + 0.55)

    floor_wall_section = """
        <body name="floor">
            <geom name="floor" pos="0 0 0" size="0 0 0.05" euler="0 0 0" type="plane" material="M_floor" conaffinity="0" contype="0"/>
        </body>"""

    if add_wall:
        floor_wall_section += f"""
        <!-- Large wall behind drawer -->
        <body name="front_wall">    
                <geom name="wall" pos="0 {-wall_y_pos} 0" size="0 0 0.05" euler="1.57 0 0" type="plane" material="M_wall"/>
        </body>
       """

    # Dynamically generate contact exclusions based on num_drawers
    contact_excludes = ""
    for i in range(num_drawers):
        contact_excludes += (
            f'        <exclude body1="drawer_{i}" body2="outer_frame"/>\n'
        )

    closing_template = (
        """
    </body>"""
        + floor_wall_section
        + f"""
  </worldbody>

  <contact>
{contact_excludes}  </contact>
</mujoco>"""
    )

    drawer_xml_string = base_template + drawer_bodies + closing_template
    return drawer_xml_string

def resize_with_pad(images, height, width):
    if images.shape[-3:-1] == (height, width):
        return images

    original_shape = images.shape
    images = images.reshape(-1, *original_shape[-3:])
    resized = np.stack(
        [
            _resize_with_pad_pil(Image.fromarray(im), height, width, method=Image.BILINEAR)
            for im in images
        ]
    )
    return resized.reshape(*original_shape[:-3], *resized.shape[-3:])


def _resize_with_pad_pil(image, height, width, method):
    cur_width, cur_height = image.size
    if cur_width == width and cur_height == height:
        return image

    ratio = max(cur_width / width, cur_height / height)
    resized_height = int(cur_height / ratio)
    resized_width = int(cur_width / ratio)
    resized_image = image.resize((resized_width, resized_height), resample=method)

    zero_image = Image.new(resized_image.mode, (width, height), 0)
    pad_height = max(0, int((height - resized_height) / 2))
    pad_width = max(0, int((width - resized_width) / 2))
    zero_image.paste(resized_image, (pad_width, pad_height))
    assert zero_image.size == (width, height)
    return zero_image

def set_table_size(scene_xml_content, table_width, table_length):
    return scene_xml_content.replace("{TABLE_WIDTH}", str(table_width)).replace(
        "{TABLE_LENGTH}", str(table_length)
    )

def pixel_to_world(u, v, depth, model, data, cam_name, img_width, img_height):
    fovy = math.radians(model.camera(cam_name).fovy[0])
    f = img_height / (2 * np.tan(fovy / 2))
    cam_mat = np.array([[f, 0, img_width / 2], [0, f, img_height / 2], [0, 0, 1]])

    cx = cam_mat[0, 2]
    cy = cam_mat[1, 2]
    fx = cam_mat[0, 0]
    fy = cam_mat[1, 1]

    z_cam = -depth
    x_cam = -(u - cx) * z_cam / fx
    y_cam = -(cy - v) * z_cam / fy

    p_cam = np.array([x_cam, y_cam, z_cam])

    xmat = data.camera(cam_name).xmat
    R_wc = xmat.reshape(3, 3)
    pos = data.camera(cam_name).xpos
    p_world = R_wc @ p_cam + pos
    return p_world