#%%
import numpy as np
import cv2

def merge_intervals(intervals):
    # Merge overlapping 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):
    # Compute forbidden intervals for this channel from neighbor values.
    forbidden = []
    for c in neighbor_values:
        low = max(0, c - delta)
        high = min(1, c + delta)
        forbidden.append((low, high))
    # Merge overlapping forbidden intervals.
    merged = merge_intervals(forbidden)
    allowed = []
    # Allowed region before the first forbidden interval.
    if merged:
        if merged[0][0] > 0:
            allowed.append((0, merged[0][0]))
        # Allowed regions between forbidden intervals.
        for i in range(1, len(merged)):
            allowed.append((merged[i-1][1], merged[i][0]))
        # Allowed region after the last forbidden interval.
        if merged[-1][1] < 1:
            allowed.append((merged[-1][1], 1))
    else:
        allowed.append((0, 1))
    return allowed

def sample_from_allowed(allowed):
    # Choose a random value from the allowed intervals with probability proportional to interval length.
    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):
    # Compute allowed intervals and sample a value.
    allowed = get_allowed_intervals(neighbor_values, delta)
    if not allowed:
        return None
    return sample_from_allowed(allowed)

def clip_allowed_to_min(allowed, min_value):
    clipped = []
    for start, end in allowed:
        start = max(start, min_value)
        if end > start:
            clipped.append((start, end))
    return clipped


def generate_valid_gray(neighbor_grays, r, black_threshold):
    delta = np.sqrt(r)
    min_gray = min_gray_normalized(black_threshold)
    allowed = clip_allowed_to_min(get_allowed_intervals(neighbor_grays, delta), min_gray)
    if not allowed:
        return np.random.uniform(min_gray, 1.0)
    sampled = sample_from_allowed(allowed)
    if sampled is None:
        return np.random.uniform(min_gray, 1.0)
    return enforce_min_gray(sampled, black_threshold)


def rgb_to_gray(color):
    # ITU-R BT.601 luma, color channels in [0, 1].
    r, g, b = color
    return 0.299 * r + 0.587 * g + 0.114 * b


def min_gray_normalized(black_threshold):
    # Strictly above black_threshold in uint8 space.
    return (float(black_threshold) + 1.0) / 255.0


def gray_to_rgb(gray):
    gray = float(gray)
    return np.array([gray, gray, gray], dtype=np.float32)


def enforce_min_gray(gray, black_threshold):
    return max(float(gray), min_gray_normalized(black_threshold))


def generate_valid_color(neighbors, r):
    """
    Generate a random color that satisfies, for each neighbor,
    |R - R_neighbor| >= sqrt(r), and similarly for G and B.
    neighbors: list of neighbor colors (each is an array of shape (3,))
    r: threshold for squared difference.
    """
    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 the allowed intervals are empty, use a fallback.
        if sampled is None:
            sampled = np.random.rand()
        color[channel] = sampled
    return color

def fill_pattern_remainder(pattern, h, w, d, black_threshold):
    """Fill pixels not covered by full dxd patches (image size may not be divisible by d)."""
    min_gray = min_gray_normalized(black_threshold)
    filled_h = (h // d) * d
    filled_w = (w // d) * d

    if filled_h < h and filled_w > 0:
        remainder_h = h - filled_h
        source = pattern[filled_h - d:filled_h, :filled_w]
        pattern[filled_h:h, :filled_w] = np.tile(
            source, ((remainder_h + d - 1) // d, 1)
        )[:remainder_h, :]

    if filled_w < w and filled_h > 0:
        remainder_w = w - filled_w
        source = pattern[:filled_h, filled_w - d:filled_w]
        pattern[:filled_h, filled_w:w] = np.tile(source, (1, (remainder_w + d - 1) // d))[:, :remainder_w]

    if filled_h < h and filled_w < w:
        remainder_h = h - filled_h
        remainder_w = w - filled_w
        source = pattern[filled_h - d:filled_h, filled_w - d:filled_w]
        pattern[filled_h:h, filled_w:w] = np.tile(
            source,
            ((remainder_h + d - 1) // d, (remainder_w + d - 1) // d),
        )[:remainder_h, :remainder_w]

    pattern[:] = np.maximum(pattern, min_gray)
    return pattern


def generate_color_pattern(h, w, d, r, black_threshold):
    """
    Generate a dense grayscale pattern.
    h, w: height and width of the output image in pixels.
    d: patch size in pixels (each patch is d x d).
    r: randomness regulation factor. For each neighbor the condition is:
       (|channel_new - channel_neighbor|)^2 >= r.
    black_threshold: minimum grayscale value in [0, 255]; patches cannot be black.
    """
    pattern = np.zeros((h, w), dtype=np.float32)
    num_patches_x = h // d
    num_patches_y = w // d
    min_gray = min_gray_normalized(black_threshold)

    for i in range(num_patches_x):
        for j in range(num_patches_y):
            # Collect neighbor patch grays as RGB triples for the existing sampler.
            neighbors = []
            neighbor_grays = []
            if i > 0:
                patch_above = pattern[(i-1)*d:i*d, j*d:(j+1)*d]
                gray_above = float(np.mean(patch_above))
                neighbor_grays.append(gray_above)
                neighbors.append(gray_to_rgb(gray_above))
            if j > 0:
                patch_left = pattern[i*d:(i+1)*d, (j-1)*d:j*d]
                gray_left = float(np.mean(patch_left))
                neighbor_grays.append(gray_left)
                neighbors.append(gray_to_rgb(gray_left))

            if len(neighbors) == 0:
                # First patch: sample directly above the black threshold.
                gray = np.random.uniform(min_gray, 1.0)
            else:
                new_color = generate_valid_color(neighbors, r)
                gray = rgb_to_gray(new_color)
                if gray < min_gray:
                    gray = generate_valid_gray(neighbor_grays, r, black_threshold)
                else:
                    gray = enforce_min_gray(gray, black_threshold)

            pattern[i*d:(i+1)*d, j*d:(j+1)*d] = gray
            print('Patch ({}, {}) gray={:.3f}'.format(i, j, gray))

    return fill_pattern_remainder(pattern, h, w, d, black_threshold)

def main():
    # Parameters for the pattern.
    # For example, using 100x100 pixels image with patches of 3x3 pixels.
    # height = 2750
    # width = 1950
    height = 1340
    width = 1360
    patch_size = 12
    black_threshold = 125
    # Set r to a value that gives sufficient difference.
    # Note: r is the threshold for squared difference per channel, and the condition is
    # |c - c_neighbor| >= sqrt(r). Very high r may make it impossible to satisfy.
    r = 1000

    pattern = generate_color_pattern(height, width, patch_size, r, black_threshold)
    # Scale to 0-255 for display and saving.
    pattern_display = (pattern * 255).astype(np.uint8)

    if pattern_display.min() <= black_threshold:
        raise RuntimeError(
            f"Generated pattern contains values <= black_threshold ({black_threshold}); "
            f"min={int(pattern_display.min())}"
        )

    # Display the pattern with OpenCV.
    cv2.imshow('Pattern', pattern_display)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # Save the pattern as an image file.
    cv2.imwrite('gray_pattern.png', pattern_display)

if __name__ == '__main__':
    main()