"""USD 材质和几何体创建工具

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

from pxr import Gf, Sdf, UsdGeom, UsdShade, Vt


def create_raw_usd_material(stage, mat_path, tex_path):
    """创建 USD PBR 材质"""
    material = UsdShade.Material.Define(stage, mat_path)
    pbr_shader = UsdShade.Shader.Define(stage, f"{mat_path}/PBRShader")
    pbr_shader.CreateIdAttr("UsdPreviewSurface")
    pbr_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0)
    pbr_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)

    tex_sampler = UsdShade.Shader.Define(stage, f"{mat_path}/diffuseTexture")
    tex_sampler.CreateIdAttr("UsdUVTexture")
    tex_sampler.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(Sdf.AssetPath(tex_path))
    tex_sampler.CreateInput("magFilter", Sdf.ValueTypeNames.Token).Set("nearest")
    tex_sampler.CreateInput("minFilter", Sdf.ValueTypeNames.Token).Set("nearest")

    st_reader = UsdShade.Shader.Define(stage, f"{mat_path}/stReader")
    st_reader.CreateIdAttr("UsdPrimvarReader_float2")
    st_reader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("st")

    tex_sampler.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(st_reader.ConnectableAPI(), "result")
    pbr_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(tex_sampler.ConnectableAPI(), "rgb")
    material.CreateSurfaceOutput().ConnectToSource(pbr_shader.ConnectableAPI(), "surface")
    return material


def create_textured_board(stage, prim_path, width, height, center, euler_rot_deg, usd_material):
    """创建带纹理的标定板"""
    mesh = UsdGeom.Mesh.Define(stage, prim_path)
    half_width, half_height = width / 2.0, height / 2.0
    mesh.GetPointsAttr().Set(Vt.Vec3fArray([
        Gf.Vec3f(-half_width, -half_height, 0),
        Gf.Vec3f(half_width, -half_height, 0),
        Gf.Vec3f(half_width, half_height, 0),
        Gf.Vec3f(-half_width, half_height, 0),
    ]))
    mesh.GetFaceVertexCountsAttr().Set([4])
    mesh.GetFaceVertexIndicesAttr().Set([0, 1, 2, 3])
    mesh.GetNormalsAttr().Set([Gf.Vec3f(0, 0, 1)] * 4)
    mesh.SetNormalsInterpolation(UsdGeom.Tokens.vertex)

    primvars_api = UsdGeom.PrimvarsAPI(mesh)
    st_primvar = primvars_api.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.vertex)
    st_primvar.Set([Gf.Vec2f(0, 0), Gf.Vec2f(1, 0), Gf.Vec2f(1, 1), Gf.Vec2f(0, 1)])
    mesh.GetExtentAttr().Set([Gf.Vec3f(-half_width, -half_height, -0.01), Gf.Vec3f(half_width, half_height, 0.01)])

    xform = UsdGeom.Xformable(mesh)
    xform.AddTranslateOp().Set(Gf.Vec3d(*center))
    xform.AddRotateXYZOp().Set(Gf.Vec3f(*euler_rot_deg))
    UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim()).Bind(usd_material)
    return mesh


def create_textured_top_strip(stage, prim_path, x_edges, y_min, y_max, z_values, usd_material):
    """创建顶部纹理条带（用于下视相机标定台）"""
    mesh = UsdGeom.Mesh.Define(stage, prim_path)
    x0 = x_edges[0]
    x1 = x_edges[-1]
    total_length = max(x1 - x0, 1e-6)

    points = []
    st_values = []
    for x, z in zip(x_edges, z_values):
        u = (x - x0) / total_length
        points.append(Gf.Vec3f(x, y_min, z))
        points.append(Gf.Vec3f(x, y_max, z))
        st_values.append(Gf.Vec2f(u, 0.0))
        st_values.append(Gf.Vec2f(u, 1.0))

    face_vertex_counts = []
    face_vertex_indices = []
    for index in range(len(x_edges) - 1):
        face_vertex_counts.append(4)
        face_vertex_indices.extend([2 * index, 2 * (index + 1), 2 * (index + 1) + 1, 2 * index + 1])

    mesh.GetPointsAttr().Set(Vt.Vec3fArray(points))
    mesh.GetFaceVertexCountsAttr().Set(face_vertex_counts)
    mesh.GetFaceVertexIndicesAttr().Set(face_vertex_indices)
    mesh.GetExtentAttr().Set([
        Gf.Vec3f(min(x_edges), y_min, min(z_values) - 0.005),
        Gf.Vec3f(max(x_edges), y_max, max(z_values) + 0.005),
    ])

    primvars_api = UsdGeom.PrimvarsAPI(mesh)
    st_primvar = primvars_api.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.vertex)
    st_primvar.Set(st_values)
    UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim()).Bind(usd_material)
    return mesh
