import random
import math
import numpy as np
from PIL import Image
from skimage.draw import line
from skimage import morphology
import cv2

def line_crosses_cracks(start, end, img):
    rr, cc = line(start[0], start[1], end[0], end[1])
    # Exclude the starting point from the line coordinates
    if len(rr) > 1 and len(cc) > 1:
        return np.any(img[rr[1:], cc[1:]] == 255)
    return False

def random_walk(img_array, k=8, m=0.1, min_steps=50, max_steps=200, length=2, degree_range=30, seed=None):
    
    if seed is not None:
        random.seed(seed)
        np.random.seed(seed)

    
    img_array = cv2.ximgproc.thinning(img_array)

    rows, cols = img_array.shape
    # Find all white pixels (existing cracks)
    white_pixels = np.column_stack(np.where(img_array == 255))
    original_crack_count = len(white_pixels)  # Count of original crack pixels

    # Select k random starting points from the white pixels
    if white_pixels.size == 0:
        raise ValueError("No initial crack pixels found in the image.")
    if k > len(white_pixels):
        raise ValueError("k is greater than the number of existing crack pixels.")
    initial_points = white_pixels[random.sample(range(len(white_pixels)), k)]

    # Initialize step count for each initial point with a random value between min_steps and max_steps
    step_counts = {i: random.randint(min_steps, max_steps) for i in range(k)}
    # Initialize main direction for each initial point (0 to 360 degrees)
    main_angles = {i: random.uniform(0, 360) for i in range(k)}

    grown_crack_count = 0  # Count of newly grown crack pixels

    # Start the random walk for each initial point
    for idx, point in enumerate(initial_points):
        current_pos = tuple(point)
        current_steps = 0
        while current_steps < step_counts[idx]:
            # Check the crack ratio
            current_ratio = np.sum(img_array == 255) / (rows * cols)
            if current_ratio >= m:
                return img_array, {'original_crack_count': original_crack_count, 'grown_crack_count': grown_crack_count}

            # Generate a random direction within the fan-shaped area around the main angle
            main_angle = main_angles[idx]
            angle = math.radians(main_angle + random.uniform(-degree_range, degree_range))
            
            # Determine the next position with the specified length
            delta_row = length * math.sin(angle)
            delta_col = length * math.cos(angle)
            next_pos = (int(current_pos[0] + delta_row), int(current_pos[1] + delta_col))
            
            # Check if the line from the current to the next position crosses existing cracks
            if 0 <= next_pos[0] < rows and 0 <= next_pos[1] < cols and not line_crosses_cracks(current_pos, next_pos, img_array):
                # Draw a line from the current position to the next position
                rr, cc = line(current_pos[0], current_pos[1], next_pos[0], next_pos[1])
                img_array[rr, cc] = 255  # Set the pixels along the line to white
                grown_crack_count += len(rr)  # Update the count of grown crack pixels
                current_pos = next_pos
                current_steps += 1
            else:
                # If the line crosses existing cracks or the next position is outside the boundaries, stop the walk for this point
                break

    return img_array, {'original_crack_count': original_crack_count, 'grown_crack_count': grown_crack_count}

# The rest of the test code remains the same.
# You can use this function in your test code to generate the image and get the counts.


# test code 
if __name__ == "__main__":
    # Updated parameters
    k = 8  # Number of initial white pixels to start the random walk
    m = 0.1  # Maximum ratio of crack pixels
    min_steps = 50
    max_steps = 200
    img_path = '/data/leiqin/diffusion/huggingface_diffusers/crack_label_creator/random_walk/thindata_256/2.png'
    img = Image.open(img_path)
    img_array = np.array(img)
    length = 2

    # Perform the modified random walk
    result_img_array_mod, pixels_dict = random_walk(img_array.copy(), k, m, min_steps, max_steps, length)

    # Convert the result to an image
    result_img_mod = Image.fromarray(result_img_array_mod.astype('uint8'))

    # Save the resulting image
    result_img_path_mod = 'resutls.png'
    result_img_mod.save(result_img_path_mod)
    print(pixels_dict)