# %%
import numpy as np
from typing import NamedTuple
class Partition2d(NamedTuple):
    weights: np.ndarray
    row_range: tuple[int, int]
    column_range: tuple[int, int]
# %%

def switch_axis(axis: int):
    return int(not bool(axis))

def bisection_cut(partitions, axis: int):
    out = []
    for p in partitions:
        weight_along_axis = p.weights.sum(axis=switch_axis(axis))
        
        # Calculate the cumulative weights and find the cut point
        cut_weight = weight_along_axis.sum() / 2
        cumulative_weight = weight_along_axis.cumsum()
        cut_index = int(np.searchsorted(cumulative_weight, cut_weight)) + 1
        
        # Ensure cut_index is valid
        if cut_index == len(weight_along_axis):
            cut_index -= 1
        elif cut_index == 0:
            cut_index = 1
        
        if axis == 0:  # Cut along rows
            p0 = Partition2d(
                p.weights[:cut_index, :], 
                (p.row_range[0], p.row_range[0] + cut_index), 
                p.column_range
            )
            p1 = Partition2d(
                p.weights[cut_index:, :], 
                (p.row_range[0] + cut_index, p.row_range[1]), 
                p.column_range
            )
        else:  # Cut along columns
            p0 = Partition2d(
                p.weights[:, :cut_index], 
                p.row_range, 
                (p.column_range[0], p.column_range[0] + cut_index)
            )
            p1 = Partition2d(
                p.weights[:, cut_index:], 
                p.row_range, 
                (p.column_range[0] + cut_index, p.column_range[1])
            )
        
        out.extend([p0, p1])
    
    return out

def weightkey(p: Partition2d):
    return p.weights.sum()

def ratiokey(p: Partition2d):
    # Probably a pretty bad measure
    nrow = p.row_range[1] - p.row_range[0]
    ncol = p.column_range[1] - p.column_range[0]
    return abs(nrow - ncol)

def recursive_bisection_cut(weights: np.ndarray, n_partition: int, partition_sortkey):
    nrow, ncol = weights.shape
    partitions = [Partition2d(weights, (0, nrow), (0, ncol))]
    axis = 0
    while len(partitions) < n_partition:
        remainder = n_partition - len(partitions)
        # Each cut results in +1 partitions
        n_cuts = min(len(partitions), remainder)
        partitions = bisection_cut(partitions[:n_cuts], axis) + partitions[n_cuts:]
        # Sort partitions by some property, e.g. sum of weights
        partitions = sorted(partitions, key=partition_sortkey, reverse=True)
        axis=switch_axis(axis)
    return partitions

def paint(weights, partitions):
    labels = np.zeros_like(weights)
    for i, p in enumerate(partitions):
        labels[p.row_range[0]: p.row_range[1], p.column_range[0]: p.column_range[1]] = i
    return labels
# %%
weights = np.random.rand(10, 10)
# %%
partitions = recursive_bisection_cut(weights, 13, weightkey)
labels = paint(weights, partitions)
import matplotlib.pyplot as plt
plt.imshow(labels)
# %%