import bpy
import bmesh
import numpy as np
from mathutils import Vector
import os

# ----------------------------------------------------------------------
# Scanline polygon fill for Blender Image pixels
# ----------------------------------------------------------------------
def fill_polygon(image, pixels, width, height, polygon, color):
    poly = [(int(round(x)), int(round(y))) for x, y in polygon]
    if len(poly) < 3:
        return

    min_x = max(0, min(p[0] for p in poly))
    max_x = min(width-1, max(p[0] for p in poly))
    min_y = max(0, min(p[1] for p in poly))
    max_y = min(height-1, max(p[1] for p in poly))

    for y in range(min_y, max_y + 1):
        intersections = []
        for i in range(len(poly)):
            x1, y1 = poly[i]
            x2, y2 = poly[(i+1) % len(poly)]
            if y1 == y2:
                continue
            if (y1 <= y < y2) or (y2 <= y < y1):
                t = (y - y1) / (y2 - y1)
                x = x1 + t * (x2 - x1)
                intersections.append(x)
        intersections.sort()
        for k in range(0, len(intersections), 2):
            if k+1 >= len(intersections):
                break
            x_start = int(round(intersections[k]))
            x_end = int(round(intersections[k+1]))
            x_start = max(0, min(width-1, x_start))
            x_end = max(0, min(width-1, x_end))
            for x in range(x_start, x_end + 1):
                idx = (y * width + x) * 4
                pixels[idx:idx+4] = color

# ----------------------------------------------------------------------
# Get interpolated UV from face and hit location (fixed version)
# ----------------------------------------------------------------------
def get_uv_from_face(obj, face_index, location):
    me = obj.data
    local_loc = obj.matrix_world.inverted() @ location

    poly = me.polygons[face_index]
    vert_indices = poly.vertices
    loop_indices = poly.loop_indices

    if len(vert_indices) < 3:
        return None

    v0 = me.vertices[vert_indices[0]].co
    v1 = me.vertices[vert_indices[1]].co
    v2 = me.vertices[vert_indices[2]].co

    uv_layer = me.uv_layers.active.data
    uv0 = uv_layer[loop_indices[0]].uv
    uv1 = uv_layer[loop_indices[1]].uv
    uv2 = uv_layer[loop_indices[2]].uv

    def tri_area(a, b, c):
        return 0.5 * ((b - a).cross(c - a)).length

    total_area = tri_area(v0, v1, v2)
    if total_area == 0:
        return None

    w0 = tri_area(local_loc, v1, v2) / total_area
    w1 = tri_area(local_loc, v2, v0) / total_area
    w2 = 1.0 - w0 - w1

    uv = w0 * uv0 + w1 * uv1 + w2 * uv2
    return Vector((uv.x, uv.y))

# ----------------------------------------------------------------------
# Color generation functions (unchanged)
# ----------------------------------------------------------------------
def merge_intervals(intervals):
    if not intervals:
        return []
    intervals.sort(key=lambda x: x[0])
    merged = [intervals[0]]
    for current in intervals[1:]:
        last = merged[-1]
        if current[0] <= last[1]:
            merged[-1] = (last[0], max(last[1], current[1]))
        else:
            merged.append(current)
    return merged

def get_allowed_intervals(neighbor_values, delta):
    forbidden = []
    for c in neighbor_values:
        low = max(0, c - delta)
        high = min(1, c + delta)
        forbidden.append((low, high))
    merged = merge_intervals(forbidden)
    allowed = []
    if merged:
        if merged[0][0] > 0:
            allowed.append((0, merged[0][0]))
        for i in range(1, len(merged)):
            allowed.append((merged[i-1][1], merged[i][0]))
        if merged[-1][1] < 1:
            allowed.append((merged[-1][1], 1))
    else:
        allowed.append((0, 1))
    return allowed

def sample_from_allowed(allowed):
    lengths = [end - start for start, end in allowed]
    total_length = sum(lengths)
    if total_length <= 0:
        return None
    point = np.random.rand() * total_length
    cumulative = 0
    for (start, end), length in zip(allowed, lengths):
        cumulative += length
        if point <= cumulative:
            return np.random.uniform(start, end)
    return allowed[-1][0]

def sample_channel(neighbor_values, delta):
    allowed = get_allowed_intervals(neighbor_values, delta)
    if not allowed:
        return None
    return sample_from_allowed(allowed)

def generate_valid_color(neighbors, r):
    delta = np.sqrt(r)
    color = np.zeros(3)
    for channel in range(3):
        neighbor_channel = [n[channel] for n in neighbors]
        sampled = sample_channel(neighbor_channel, delta)
        if sampled is None:
            sampled = np.random.rand()
        color[channel] = sampled
    return color

# ----------------------------------------------------------------------
# Main function
# ----------------------------------------------------------------------
def generate_fingerprint_texture(
    obj_name=None,
    camera_name="Camera",
    nx=40,
    ny=30,
    r=0.05,
    texture_width=2048,
    texture_height=2048,
    output_path="//finger_texture.png"
):
    scene = bpy.context.scene
    depsgraph = bpy.context.evaluated_depsgraph_get()
    
    if obj_name:
        obj = bpy.data.objects[obj_name]
    else:
        obj = bpy.context.active_object
    if obj is None or obj.type != 'MESH':
        raise ValueError("No valid mesh object selected")
    
    cam = bpy.data.objects[camera_name]
    if cam is None or cam.type != 'CAMERA':
        raise ValueError("Camera not found")
    
    if not obj.data.uv_layers.active:
        obj.data.uv_layers.new(name="UVMap")
    
    render = scene.render
    res_x = render.resolution_x
    res_y = render.resolution_y
    
    cam_data = cam.data
    if cam_data.sensor_fit == 'HORIZONTAL':
        sensor_width = cam_data.sensor_width
        sensor_height = sensor_width / (res_x / res_y)
    else:
        sensor_height = cam_data.sensor_height
        sensor_width = sensor_height * (res_x / res_y)
    focal = cam_data.lens
    frame_width = sensor_width / focal
    frame_height = sensor_height / focal
    
    patch_uv_polygons = []
    patch_colors = np.zeros((ny, nx, 3), dtype=np.float32)
    
    for j in range(ny):
        v_start = j / ny
        v_end = (j+1) / ny
        for i in range(nx):
            u_start = i / nx
            u_end = (i+1) / nx
            
            corners_ndc = [
                (u_start, v_start),
                (u_end,   v_start),
                (u_end,   v_end),
                (u_start, v_end)
            ]
            
            uv_poly = []
            for (u_ndc, v_ndc) in corners_ndc:
                sx = (u_ndc - 0.5) * frame_width
                sy = (v_ndc - 0.5) * frame_height
                direction_cam = Vector((sx, sy, -1.0))
                direction_world = cam.matrix_world.to_quaternion() @ direction_cam
                direction_world.normalize()
                origin = cam.matrix_world.translation
                
                hit, location, normal, face_index, hit_obj, matrix = scene.ray_cast(depsgraph, origin, direction_world)
                if hit and hit_obj == obj:
                    uv = get_uv_from_face(obj, face_index, location)
                    if uv is not None:
                        uv_poly.append(uv)
            
            if len(uv_poly) < 3:
                patch_uv_polygons.append(None)
                continue
            patch_uv_polygons.append(uv_poly)
            
            neighbors = []
            if j > 0:
                neighbors.append(patch_colors[j-1, i])
            if i > 0:
                neighbors.append(patch_colors[j, i-1])
            
            if len(neighbors) == 0:
                color = np.random.rand(3)
            else:
                color = generate_valid_color(neighbors, r)
            patch_colors[j, i] = color
    
    img_name = "FingerprintTexture"
    if img_name in bpy.data.images:
        bpy.data.images.remove(bpy.data.images[img_name])
    image = bpy.data.images.new(img_name, width=texture_width, height=texture_height, alpha=True)
    pixels = list(image.pixels)
    
    idx = 0
    for j in range(ny):
        for i in range(nx):
            uv_poly = patch_uv_polygons[idx]
            idx += 1
            if uv_poly is None:
                continue
            pixel_poly = [(u * (texture_width - 1), v * (texture_height - 1)) for (u, v) in uv_poly]
            color = patch_colors[j, i]
            rgba = (color[0], color[1], color[2], 1.0)
            fill_polygon(image, pixels, texture_width, texture_height, pixel_poly, rgba)
    
    image.pixels[:] = pixels
    image.update()
    
    abs_path = bpy.path.abspath(output_path)
    image.filepath_raw = abs_path
    image.file_format = 'PNG'
    image.save()
    print(f"Texture saved to {abs_path}")
    return image

# ----------------------------------------------------------------------
# Run
# ----------------------------------------------------------------------
if __name__ == "__main__":
    generate_fingerprint_texture(
        obj_name="FingerMesh",   # Change to your mesh name
        nx=50, ny=40,
        r=0.05,
        texture_width=4096,
        texture_height=4096,
        output_path="//finger_texture.png"
    )