#%%
import numpy as np
import cv2

# Flat pattern scale before warping to a dome: 1 mm == 100 pixels.
PX_PER_MM = 100.0


def mm_to_px(mm, scale_px_per_mm=PX_PER_MM):
    """Convert millimetres to pixels (default 1 mm = 100 px)."""
    return mm * scale_px_per_mm


def px_to_mm(px, scale_px_per_mm=PX_PER_MM):
    """Convert pixels to millimetres (default 1 mm = 100 px)."""
    return px / scale_px_per_mm


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 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 _easing_progress(t, easing, gamma):
    """
    Map normalized position t in [0, 1] along an axis to a progress u in [0, 1]
    for patch-size interpolation: size = start + (end - start) * u.
    """
    t = float(np.clip(t, 0.0, 1.0))
    e = (easing or "linear").lower()
    if e == "linear":
        return t
    if e == "power":
        g = float(gamma)
        return t**g
    if e == "cosine":
        return 0.5 * (1.0 - np.cos(np.pi * t))
    if e in ("dome", "dome_sin2"):
        # Shallow rise from left (t=0) to right (t=1); slows change near the left
        # (useful when the left is farther on a dome and needs larger flat patches).
        return np.sin(0.5 * np.pi * t) ** 2
    raise ValueError(
        f"Unknown easing {e!r}; use 'linear', 'power', 'cosine', or 'dome'"
    )


def patch_sizes_nonlinear_along_length(
    total_len_px,
    patch_px_start,
    patch_px_end,
    *,
    easing="power",
    gamma=2.0,
):
    """
    Integer segment lengths along an axis of length ``total_len_px``, with **nonlinear**
    nominal patch size from ``patch_px_start`` (at the beginning of the axis, e.g. left)
    to ``patch_px_end`` (at the end, e.g. right).

    At cumulative pixel position ``pos`` (start of a segment), let ``t = pos / total_len_px``.
    Nominal size is::

        patch_px_start + (patch_px_end - patch_px_start) * easing(t)

    Segments are at least 1 pixel wide; the sum of segment lengths equals ``total_len_px``.

    Parameters
    ----------
    total_len_px : int
        Full extent in pixels (e.g. image width after ``mm_to_px(length_mm)``).
    patch_px_start, patch_px_end : int
        Patch size in pixels at the start vs end of the axis.
    easing : str
        - ``'linear'``: same as uniform interpolation in ``t``
        - ``'power'``: ``t ** gamma`` (gamma > 1 keeps sizes nearer ``patch_px_start`` longer)
        - ``'cosine'``: smooth cosine half-period
        - ``'dome'``: ``sin(pi/2 * t) ** 2`` (gentle dome-style progression)
    gamma : float
        Exponent for ``easing='power'`` (ignored otherwise).

    See Also
    --------
    patch_sizes_from_mm : same logic with inputs in millimetres.
    """
    if total_len_px <= 0:
        return []
    a = max(1, int(round(float(patch_px_start))))
    b = max(1, int(round(float(patch_px_end))))
    tl = float(total_len_px)
    sizes = []
    pos = 0
    while pos < total_len_px:
        t = pos / tl
        u = _easing_progress(t, easing, gamma)
        nominal = a + (b - a) * u
        seg = int(max(1, round(float(nominal))))
        seg = min(seg, total_len_px - pos)
        sizes.append(seg)
        pos += seg
    diff = total_len_px - sum(sizes)
    if diff != 0 and sizes:
        sizes[-1] += diff
        sizes[-1] = max(1, sizes[-1])
    return sizes


def patch_sizes_from_mm(
    length_mm,
    patch_mm_start,
    patch_mm_end,
    *,
    px_per_mm=PX_PER_MM,
    **kwargs,
):
    """
    Same as ``patch_sizes_nonlinear_along_length`` but axis length and patch bounds
    are given in **millimetres**; converted with ``1 mm = px_per_mm`` pixels (default 100).
    """
    L = int(round(float(length_mm) * float(px_per_mm)))
    a = max(1, int(round(float(patch_mm_start) * float(px_per_mm))))
    b = max(1, int(round(float(patch_mm_end) * float(px_per_mm))))
    return patch_sizes_nonlinear_along_length(L, a, b, **kwargs)


def patch_sizes_along_length(total_len, size_at_start, size_at_end):
    """
    Integer segment lengths along an axis of length `total_len`, varying **linearly**
    from `size_at_start` (at the beginning of the axis) to `size_at_end` (at the end).

    Equivalent to ``patch_sizes_nonlinear_along_length(..., easing='linear')``.
    """
    return patch_sizes_nonlinear_along_length(
        total_len, size_at_start, size_at_end, easing="linear", gamma=1.0
    )


def generate_color_pattern_varying_horizontal(
    h,
    w,
    patch_w_max_left,
    patch_w_min_right,
    r,
    *,
    easing="linear",
    gamma=2.0,
):
    """
    Like ``generate_color_pattern``, but patch **size** varies along x (largest on the
    left, smallest on the right). Each patch is **square**: side length equals that
    column's width, and the same step is used vertically within the column.

    Parameters
    ----------
    h, w : int
        Output image size.
    patch_w_max_left : int
        Square side length at the left edge (maximum).
    patch_w_min_right : int
        Square side length toward the right edge (minimum).
    r : float
        Same neighbor color constraint as ``generate_color_pattern``.
    easing : str
        Passed to ``patch_sizes_nonlinear_along_length`` (e.g. ``'dome'`` for dome warp).
    gamma : float
        Exponent when ``easing='power'``.

    Note
    ----
    If ``h`` is not a multiple of a column's side length ``s``, the last row in that
    column is shorter than ``s`` pixels tall while the width stays ``s`` (a thin strip).
    """
    pattern = np.zeros((h, w, 3), dtype=np.float32)

    col_widths = patch_sizes_nonlinear_along_length(
        w, patch_w_max_left, patch_w_min_right, easing=easing, gamma=gamma
    )
    x_edges = np.concatenate([[0], np.cumsum(col_widths, dtype=np.int64)])
    num_cols = len(col_widths)

    for j in range(num_cols):
        s = int(col_widths[j])
        x0 = int(x_edges[j])
        x1 = int(x_edges[j + 1])
        y = 0
        while y < h:
            y1 = min(y + s, h)
            neighbors = []
            if y >= s:
                patch_above = pattern[y - s : y, x0:x1]
                if patch_above.size:
                    neighbors.append(np.mean(patch_above, axis=(0, 1)))
            if j > 0:
                xl0 = int(x_edges[j - 1])
                patch_left = pattern[y:y1, xl0:x0]
                if patch_left.size:
                    neighbors.append(np.mean(patch_left, axis=(0, 1)))

            if len(neighbors) == 0:
                new_color = np.random.rand(3)
            else:
                new_color = generate_valid_color(neighbors, r)

            pattern[y:y1, x0:x1] = new_color
            y = y1

    return pattern


def generate_color_pattern(h, w, d, r):
    """
    Generate a dense color 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.
    """
    pattern = np.zeros((h, w, 3), dtype=np.float32)
    num_patches_x = h // d
    num_patches_y = w // d

    for i in range(num_patches_x):
        for j in range(num_patches_y):
            # Collect neighbor patch colors (using mean because patches have uniform color).
            neighbors = []
            if i > 0:
                patch_above = pattern[(i-1)*d:i*d, j*d:(j+1)*d]
                neighbors.append(np.mean(patch_above, axis=(0, 1)))
            if j > 0:
                patch_left = pattern[i*d:(i+1)*d, (j-1)*d:j*d]
                neighbors.append(np.mean(patch_left, axis=(0, 1)))
                
            if len(neighbors) == 0:
                # For the first patch use a random color.
                new_color = np.random.rand(3)
            else:
                new_color = generate_valid_color(neighbors, r)
                
            pattern[i*d:(i+1)*d, j*d:(j+1)*d] = new_color
            print('Patch ({}, {})'.format(i, j))
    return pattern


def draw_center_cross_bgr(
    img,
    *,
    color=(0, 0, 0),
    arm_half_len=None,
    thickness=3,
):
    """
    Draw a black (+) cross centered on the image for print / camera alignment.
    `img` is modified in place (BGR uint8).
    """
    h, w = img.shape[:2]
    cx, cy = w // 2, h // 2
    if arm_half_len is None:
        arm_half_len = max(8, min(h, w) // 40)
    half = int(arm_half_len)
    cv2.line(
        img,
        (cx - half, cy),
        (cx + half, cy),
        color,
        thickness,
        lineType=cv2.LINE_AA,
    )
    cv2.line(
        img,
        (cx, cy - half),
        (cx, cy + half),
        color,
        thickness,
        lineType=cv2.LINE_AA,
    )


def main():
    # Parameters for the pattern.
    # For example, using 100x100 pixels image with patches of 3x3 pixels.
    # height = 2750
    # width = 1950
    height = 1502
    width = 1738
    patch_size = 12
    # 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

    # 1 mm = 100 px: e.g. width_mm = px_to_mm(width); patch sizes in mm -> px with mm_to_px.
    # Column widths only (nonlinear for dome — left farther): patch_sizes_from_mm(
    #     length_mm=px_to_mm(width), patch_mm_start=0.30, patch_mm_end=0.06, easing="dome"
    # )
    pattern = generate_color_pattern_varying_horizontal(
        height, width, patch_w_max_left=30, patch_w_min_right=6, r=r, easing="linear"
    )
    # pattern = generate_color_pattern(height, width, patch_size, r)
    # Scale to 0-255 for display and saving.
    pattern_display = (pattern * 255).astype(np.uint8)
    draw_center_cross_bgr(pattern_display)

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

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

if __name__ == '__main__':
    main()