import trimesh
import numpy as np
import mesh2sdf


def mesh_to_sdf_glsl(path, resolution=32):
    # Load mesh
    mesh = trimesh.load(path)
    print(f"\nInitial mesh:")
    print(f"  Vertices: {len(mesh.vertices)}")
    print(f"  Faces: {len(mesh.faces)}")
    print(f"  Bounds: {mesh.bounds}")
    print(f"  Extents: {mesh.extents}")

    # Center the mesh
    bounds = mesh.bounds
    center = (bounds[0] + bounds[1]) / 2.0
    mesh.vertices -= center
    print(f"\nAfter centering:")
    print(f"  Center: {center}")
    print(f"  Bounds: {mesh.bounds}")
    print(f"  Extents: {mesh.extents}")

    # Scale the mesh to fit in a reasonable size for SDF computation
    max_extent = np.max(mesh.extents)
    scale = 1.0 / max_extent
    mesh.vertices *= scale
    print(f"\nAfter scaling:")
    print(f"  Scale factor: {scale}")
    print(f"  Bounds: {mesh.bounds}")
    print(f"  Extents: {mesh.extents}")
    print(f"  Vertex range: [{np.min(mesh.vertices)}, {np.max(mesh.vertices)}]")

    # Convert vertices and faces to the correct types
    vertices = mesh.vertices.astype(np.float32)
    faces = mesh.faces.astype(np.uint32)
    print(f"\nVertex/Face types:")
    print(f"  Vertices dtype: {vertices.dtype}")
    print(f"  Faces dtype: {faces.dtype}")

    # Calculate signed distances using mesh2sdf
    print(f"\nComputing SDF with resolution {resolution}...")
    distances = mesh2sdf.compute(vertices, faces, resolution)

    # Debug output
    print(f"\nSDF computation results:")
    print(f"  Distance range: [{np.min(distances)}, {np.max(distances)}]")
    print(f"  Distance shape: {distances.shape}")
    print(f"  Mean distance: {np.mean(distances)}")
    print(f"  Std distance: {np.std(distances)}")

    # Check for any NaN or infinite values
    print(f"  NaN values: {np.isnan(distances).sum()}")
    print(f"  Inf values: {np.isinf(distances).sum()}")

    # Normalize distances to [-1, 1] range
    # This preserves the sign information which is important for SDF
    max_dist = np.max(np.abs(distances))
    distances = distances / max_dist
    print(f"\nAfter normalization:")
    print(f"  Max distance: {max_dist}")
    print(f"  New range: [{np.min(distances)}, {np.max(distances)}]")
    print(f"  Mean distance: {np.mean(distances)}")
    print(f"  Std distance: {np.std(distances)}")

    # Split the data into chunks of 1024 elements (GLSL array size limit)
    chunk_size = 1024
    chunks = []
    flat_distances = distances.ravel()
    print(f"\nChunking:")
    print(f"  Total elements: {len(flat_distances)}")
    print(f"  Number of chunks: {len(flat_distances) // chunk_size + 1}")

    for i in range(0, len(flat_distances), chunk_size):
        chunk = flat_distances[i : i + chunk_size]
        chunks.append(", ".join(f"{float(v):.4f}" for v in chunk))

    # Generate GLSL code with multiple arrays
    glsl_code = f"""// Auto-generated SDF GLSL from mesh
// Resolution: {resolution}x{resolution}x{resolution}

// Split into chunks to avoid GLSL array size limits
float sdfData0[{len(chunks[0].split(','))}] = float[](
    {chunks[0]}
);
"""

    # Add additional chunks if needed
    for i in range(1, len(chunks)):
        glsl_code += f"""
float sdfData{i}[{len(chunks[i].split(','))}] = float[](
    {chunks[i]}
);
"""

    # Add the SDF function that combines chunks
    glsl_code += f"""
float SDF(vec3 p) {{
    // Map from [-1,1] to [0,resolution-1]
    vec3 dim = vec3({resolution}.0);
    vec3 uv = (p + 1.0) * 0.5 * (dim - 1.0);
    
    // Add a small offset to avoid boundary issues
    uv = clamp(uv, vec3(0.0), dim - 1.0);
    
    // Use smooth interpolation between voxels
    ivec3 idx = ivec3(floor(uv));
    vec3 frac = uv - vec3(idx);
    
    // Convert 3D index to 1D array index
    int i = idx.x + idx.y * {resolution} + idx.z * {resolution * resolution};
    
    // Select the appropriate chunk and index within it
    int chunk = i / {chunk_size};
    int local_idx = i % {chunk_size};
    
    // Return the appropriate value based on chunk
    switch(chunk) {{
"""

    # Add switch cases for each chunk
    for i in range(len(chunks)):
        glsl_code += f"        case {i}: return sdfData{i}[local_idx];\n"

    glsl_code += """        default: return 1.0; // Outside the volume
    }
}
"""
    return glsl_code