import torch
import inspect

import torch.nn.functional as F

from typing import TypeAlias, Iterable, Callable, _CallableGenericAlias, _GenericAlias
from types import GenericAlias, UnionType
from dataclasses import dataclass
from enum import Enum
from math import sqrt, log, pi
from torch import Tensor
from scipy.special import erfinv


#               #
# === TYPES === #
#               #

Filter: TypeAlias = Callable[[Tensor], Tensor]


#                             #
# === BOUNDING BOX (BBOX) === #
#                             #

# BBox stores the widths of segments i.e. left is the width of the segment to 
# the left of the box, width is the width of the box, and right is the width of
# the segment to the right of the box.

@dataclass
class BBox:
    left: int
    width: int
    right: int
    up: int
    height: int
    down: int

    @property
    def total_width(self):
        return self.left + self.width + self.right

    @property
    def total_height(self):
        return self.up + self.height + self.down

    @property
    def down_offset(self):
        return self.up + self.height

    @property
    def right_offset(self):
        return self.left + self.width

    @property
    def valid(self) -> bool:
        return \
            self.left   >= 0 and \
            self.width  >= 0 and \
            self.right  >= 0 and \
            self.up     >= 0 and \
            self.height >= 0 and \
            self.down   >= 0


def bbox_from_2d_binary(tensor: Tensor) -> BBox:
    _check([
        "tensor.dim() == 2",
        "tensor.dtype == torch.bool",
    ])
    u, h, d = _h_min_max(tensor)
    l, w, r = _h_min_max(tensor.transpose(0, 1))
    bbox = BBox(l, w, r, u, h, d)
    return bbox


def _h_min_max(tensor: Tensor) -> tuple[int, int, int]:
    th, _ = tensor.shape
    h_index = torch.arange(0, th)
    h_index_reversed = (th - 1) - h_index
    crushed = torch.sum(tensor, 1, dtype=torch.bool)
    down_offset = int(torch.amax(crushed * h_index)) + 1
    up_offset = int(torch.amax(crushed * h_index_reversed)) + 1
    down = th - down_offset
    up = th - up_offset
    height = th - (down + up)
    return (up, height, down)


def bbox_expand_area(bbox: BBox, area_multiplier: float) -> BBox:
    _check([
        "bbox.valid",
        "area_multiplier >= 0",
    ])
    side_multiplier = sqrt(area_multiplier)
    _, _, _, am = _axis_expand(bbox.left, bbox.width, bbox.right, side_multiplier)
    u, h, d, am = _axis_expand(bbox.up, bbox.height, bbox.down, area_multiplier / am)
    l, w, r, __ = _axis_expand(bbox.left, bbox.width, bbox.right, area_multiplier / am)
    expanded = BBox(l, w, r, u, h, d)
    return expanded


def bbox_outer_square(bbox: BBox) -> BBox:
    _check([
        "bbox.valid",
    ])
    w_mul = max(bbox.height / bbox.width, 1)
    h_mul = max(bbox.width / bbox.height, 1)
    nl, nw, nr, _ = _axis_expand(bbox.left, bbox.width, bbox.right, w_mul) 
    nu, nh, nd, _ = _axis_expand(bbox.up, bbox.height, bbox.down, h_mul)
    new_bbox = BBox(nl, nw, nr, nu, nh, nd)
    return new_bbox


def _axis_expand(a: int, b: int, c: int, m: float) -> tuple[int, int, int, float]:
    total = a + b + c
    new_b = min(total, round(b * m))
    actual_m = new_b / b
    delta = new_b - b
    new_a = max(0, a - delta // 2)
    new_c = max(0, total - (new_a + new_b))
    new_a = total - (new_b + new_c)
    return (new_a, new_b, new_c, actual_m)


#                    #
# === COLORSPACE === #
#                    #

# Reference: http://www.ericbrasseur.org/gamma.html?i=1#formulas
# NB: Colorspace conversion should generally be done with 64-bit floats and
# conversions back to gamma should be quantized to the number of colors that
# the final encoding supports (e.g. quantized to 256 values if the final
# encoding is 8-bits per channel).

def colorspace_srgb_linear_from_gamma(tensor: Tensor) -> Tensor:
    _check([
        "tensor.is_floating_point()"
    ])
    piece_lo = tensor / 12.92
    A = 0.055
    piece_hi = torch.pow((tensor + A) / (1 + A), 2.4)
    linear = torch.where(tensor <= 0.04045, piece_lo, piece_hi)
    return linear


def colorspace_srgb_gamma_from_linear(tensor: Tensor) -> Tensor:
    _check([
        "tensor.is_floating_point()"
    ])
    piece_lo = tensor * 12.92
    A = 0.055
    piece_hi = (1 + A) * torch.pow(tensor, 1/2.4) - A
    gamma = torch.where(tensor <= 0.0031308, piece_lo, piece_hi)
    return gamma


def colorspace_linear_from_gamma(tensor: Tensor, gamma: float = 2.2) -> Tensor:
    _check([
        "tensor.is_floating_point()"
    ])
    linear = torch.pow(tensor, gamma)
    return linear


def colorspace_gamma_from_linear(tensor: Tensor, gamma: float = 2.2) -> Tensor:
    _check([
        "tensor.is_floating_point()"
    ])
    gamma = torch.pow(tensor, 1 / gamma)
    return gamma


def colorspace_log2_from_linear(tensor: Tensor, eps: float = 0.00001) -> Tensor:
    _check([
        "tensor.is_floating_point()"
    ])
    log = torch.log2(tensor + eps)
    return log


def colorspace_linear_from_log2(tensor: Tensor, eps: float = 0.00001) -> Tensor:
    _check([
        "tensor.is_floating_point()"
    ])
    linear = torch.exp2(tensor) - eps
    return linear


#                    #
# === CONVERSION === #
#                    #

def convert_luminance_from_linear_srgb(tensor: Tensor, channel_dim: int) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "_valid_dim(tensor, channel_dim)",
        "tensor.shape[channel_dim] == 3",
    ])
    coef_shape = [1] * tensor.dim()
    coef_shape[channel_dim] = 3
    luminance_coef = torch.tensor((0.2126, 0.7152, 0.0722)).reshape(*coef_shape)
    luminance = (tensor * luminance_coef).sum(channel_dim, keepdims=True)
    return luminance


def convert_component_coloring_from_2d_binary(tensor: Tensor, diagonals: bool = False) -> Tensor:
    _check([
        "tensor.dim() == 2",
        "tensor.dtype == torch.bool",
    ])
    h, w = tensor.shape
    integer = tensor.to(torch.int)

    # 1D component coloring
    prepend = torch.zeros(h, 1, dtype=torch.int)
    diffed = torch.diff(integer, dim=1, prepend=prepend).clamp(0, 1)
    cumsum = torch.cumsum(diffed.flatten(), 0).reshape(h, w)
    coloring_1d = cumsum * integer
    coloring_2d = coloring_1d.clone()
    max_unique = 0

    # Find equivalences between strips
    adjacent_middle = torch.logical_and(tensor[:-1, :], tensor[1:, :])
    equ_m_u = coloring_1d[:-1, :][adjacent_middle]
    equ_m_d = coloring_1d[1:, :][adjacent_middle]
    equivalences = torch.stack((equ_m_u, equ_m_d), dim=1)
    if diagonals:  # 8-way equivalance
        adjacent_left = torch.logical_and(tensor[:-1, :-1], tensor[1:, 1:])
        adjacent_right = torch.logical_and(tensor[:-1, 1:], tensor[1:, :-1])
        equ_l_u = coloring_1d[:-1, :-1][adjacent_left]
        equ_r_u = coloring_1d[:-1, 1:][adjacent_right]
        equ_l_d = coloring_1d[1:, 1:][adjacent_left]
        equ_r_d = coloring_1d[1:, :-1][adjacent_right]
        equ_l = torch.stack((equ_l_u, equ_l_d), dim=1)
        equ_r = torch.stack((equ_r_u, equ_r_d), dim=1)
        equivalences = torch.cat((equivalences, equ_l, equ_r), dim=0)

    # Find all of the strip coloring identifiers
    unique = torch.unique_consecutive(cumsum).tolist()
    if unique[0] == 0:
        unique = unique[1:]
    if len(unique) == 0:
        coloring_2d[:, :] = 0
        return (coloring_2d, max_unique)

    # Find all of the sets of 1D colorings that should be tha same color
    umax = int(unique[-1])
    parent = list(range(umax + 1))
    size = [1] * (umax + 1)
    for pair in equivalences.tolist():
        p0, p1 = pair
        _disjoint_set_union(p0, p1, parent, size)
    for x in range(umax + 1):
        parent[x] = _disjoint_set_find(x, parent)

    # Relabel the previously found groups so that they're sequential integers
    parent = Tensor(parent)
    parent_unique = torch.unique(parent)
    max_unique = max(max_unique, int(parent_unique.shape[0]) - 1)
    index_map = torch.zeros((umax + 1, ), dtype=torch.long)
    index_map[parent_unique] = torch.arange(0, parent_unique.shape[0])
    parent_relabel = index_map.index_select(0, parent)

    # Recolor the 1D colorings to their grouped and relabeled 2D coloring
    coloring_2d.reshape(-1)[:] = parent_relabel.index_select(0, coloring_2d.reshape(-1))
    return (coloring_2d, max_unique)


def _disjoint_set_union(x, y, parents, size):
    px = _disjoint_set_find(x, parents)
    py = _disjoint_set_find(y, parents)
    if px != py:
        if size[px] < size[py]:
            (px, py) = (py, px)
        parents[py] = px
        size[px] += size[py]


def _disjoint_set_find(x, parents):
    while parents[x] != x:
        parents[x] = parents[parents[x]]
        x = parents[x]
    return x


def convert_gradient_suppression_mask(gtf_r: Tensor, gtf_theta: Tensor) -> Tensor:
    padded = transform_pad_dim_reflect(transform_pad_dim_reflect(gtf_r, 3, (1, 1)), 2, (1, 1))
    views = []
    for i in range(9):
        x, y = i % 3, i // 3
        nx, ny = _ztn(-(2-x)), _ztn(-(2-y))
        views += [padded[:, :, y:ny, x:nx]]
    comparisons = []
    for i in range(4):
        lower = views[4] >= views[i]
        upper = views[4] >= views[8-i]
        comparisons += [torch.logical_and(lower, upper)]
    angle_index = (((gtf_theta + (9 * pi / 8)) * (4 / pi)).to(torch.int) + 3) % 4
    odd = (angle_index % 2).to(torch.bool)
    l = torch.where(odd, comparisons[1], comparisons[0])
    u = torch.where(odd, comparisons[3], comparisons[2])
    mask = torch.where((angle_index // 2).to(torch.bool), u, l)
    return mask


def _ztn(value: int) -> int | None:
    return None if value == 0 else value


# https://en.wikipedia.org/wiki/Otsu%27s_method
def convert_otsus_method(gtf: Tensor, bins: int) -> Tensor:
    b, c, h, w = gtf.shape
    binned = (util_range_normalize(gtf, (2, 3)) * (bins - 1)).to(torch.int)
    batches = []
    for bi in range(b):
        channels = []
        for ci in range(c):
            n = torch.bincount(binned[bi, ci].flatten(), minlength=bins)
            n_cumsum = torch.cumsum(n, 0)
            N = n_cumsum[-1]
            omega = n_cumsum / N
            mu = torch.cumsum(n_cumsum * torch.arange(bins), 0) / N
            mu_T = mu[-1]
            sigma_2_B = (mu_T * omega - mu)**2 / (omega * (1 - omega))
            threshold = torch.argmax(sigma_2_B.nan_to_num()) / (bins - 1)
            if threshold.dim() == 1:
                return threshold[0]
            channels += [threshold.reshape(1, 1)]
        batches += [torch.stack(channels)]
    thresholds = torch.stack(batches)
    return thresholds


#                   #
# === FILTERING === #
#                   #

def filter_convolve_2d(tensor: Tensor, kernel: Tensor, dims: tuple[int, int]) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_ascending(dims)",
        "kernel.dim() == 2",
        "kernel.shape[0] % 2 == 1 and kernel.shape[1] % 2 == 1",
    ])
    kh, kw = kernel.shape
    kernel_reshaped = kernel.reshape(1, 1, kh, kw)
    ph, pw = kh // 2, kw // 2
    permuted = utils_dims_to_end(tensor, dims)
    padded = transform_pad_dim2_reflect(permuted, (tensor.dim() - 2, tensor.dim() - 1), ((ph, ph), (pw, pw)))
    flattened = padded.flatten(0,-3).unsqueeze(0)
    convolved = F.conv2d(flattened, kernel_reshaped, groups=flattened.shape[0])
    reshaped = convolved.reshape(*permuted.shape)
    unpermuted = utils_dims_from_end(reshaped, dims)
    return unpermuted


def filter_convolve_2d_separable(tensor: Tensor, kernels: tuple[Tensor, Tensor], dims: tuple[int, int]) -> Tensor:
    _check([
        "kernels[0].dim() == 1 and kernels[1].dim() == 1",
        "kernels[0].shape[0] % 2 == 1 and kernels[1].shape[0] % 2 == 1",
        "_ascending(dims)",
        "_valid_dims(tensor, dims)",
    ])
    c0 = tensor
    c1 = filter_convolve_2d(c0, kernels[0].reshape(-1, 1), dims)
    c2 = filter_convolve_2d(c1, kernels[1].reshape(1, -1), dims)
    return c2


def filter_hysteresis_threshold(gtf_weak: Tensor, gtf_strong: Tensor) -> Tensor:
    b, c, h, w = gtf_weak.shape
    coloring, max_unique = TF.component_coloring(gtf_weak, True)
    strong_components = gtf_strong.to(torch.bool) * coloring
    batches = []
    for bi in range(b):
        channels = []
        for ci in range(c):
            strong_select = torch.zeros(max_unique + 1)
            strong_select[strong_components[bi, ci].flatten()] = 1
            strong_select[0] = 0
            result = strong_select.index_select(0, coloring[bi, ci].flatten()).reshape(h, w)
            channels += [result]
        batches += [torch.stack(channels)]
    thresholded = torch.stack(batches)
    return thresholded


def filter_patch_min(gtf: Tensor, radius: int) -> Tensor:
    k = 2 * radius + 1
    padded = transform_pad_dim2_reflect(gtf, (2, 3), (radius, radius))
    minimum = -F.max_pool2d(-padded, k, stride=1)
    return minimum


def filter_patch_max(gtf: Tensor, radius: int) -> Tensor:
    k = 2 * radius + 1
    padded = transform_pad_dim2_reflect(gtf, (2, 3), (radius, radius))
    maximum = F.max_pool2d(padded, k, stride=1)
    return maximum


def filter_patch_median(gtf: Tensor, radius: int) -> Tensor:
    def unfold(gtf: Tensor, kh: int, kw: int) -> Tensor:
        unfolded = gtf.unfold(3 , kw, 1).unfold(2, kh, 1)
        return unfolded
    b, c, h, w = gtf.shape
    k = 2 * radius + 1
    padded = transform_pad_dim2_reflect(gtf, (2, 3), (radius, radius))
    unfolded = unfold(padded, k, k).flatten(-2)
    median = unfolded.median(dim=-1, keepdim=False).values
    return median


def filter_patch_range_normalize(gtf: Tensor, radius: int) -> Tensor: 
    r = radius 
    k = 2 * r + 1
    minimum = patch_min(gtf, radius)
    ln = gtf - minimum
    maximum = patch_max(ln, radius)
    normalized = ln / maximum
    return normalized


# MORPHOLOGICAL

def filter_morpho_dilate(tensor: Tensor, radius: int, dims: tuple[int, int]) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "radius > 0",
    ])
    kernel_size = 2 * radius - 1
    padding = radius - 1
    permuted = utils_dims_to_end(tensor, dims)
    reshaped = permuted.flatten(0, -3).unsqueeze(0)
    dilated = F.max_pool2d(reshaped, kernel_size, stride=1, padding=padding)
    unpermuted = utils_dims_from_end(dilated.reshape(*permuted.shape), dims)
    return unpermuted


# def filter_morpho_erode(tensor: Tensor, radius: int) -> Tensor:
#     eroded = _invert(dilate(_invert(tensor), radius))
#     return eroded
# 
# 
# def filter_morpho_close(tensor: Tensor, radius: int) -> Tensor:
#     closed = erode(dilate(tensor, radius), radius)
#     return closed
# 
# 
# def filter_morpho_open(tensor: Tensor, radius: int) -> Tensor:
#     opened = dilate(erode(tensor, radius), radius)
#     return opened
# 
# 
# def _invert(tensor: Tensor) -> Tensor:
#     inverted = (1.0 - tensor)
#     return inverted


# OTHER

class RoundingMode(Enum):
    ROUND = torch.round
    FLOOR = torch.floor
    CEILING = torch.ceil


def filter_quantize(tensor: Tensor, steps: int, mode: RoundingMode) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "steps > 1",
    ])
    max_val = steps - 1
    quantized = mode.value(tensor * max_val) / max_val
    return quantized


#                    #
# === RESAMPLING === #
#                    #

# NEAREST NEIGHBOR

def resample_nearest_neighbor_2d(
    tensor: Tensor, 
    new_lens: tuple[int, int], 
    dims: tuple[int, int]
) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_positive(new_lens)",
    ])
    D_0 = _nearest_neighbor_1d(tensor, new_lens[0], dims[0])
    D_1 = _nearest_neighbor_1d(D_0, new_lens[1], dims[1])
    return D_1


def _nearest_neighbor_indices(
    I_t: Tensor,
    l: int,
    L: int
) -> Tensor:
    i_snn = ((2 * I_t + 1) * l) // (2 * L)
    return i_snn


def _nearest_neighbor_1d(
    d_s: Tensor,
    L: int,
    dim: int = 0
) -> Tensor:
    L = L
    l = int(d_s.shape[dim])
    I_t = torch.arange(L)
    indices = _nearest_neighbor_indices(I_t, l, L)
    resampled = d_s.index_select(dim, indices)
    return resampled


# FILTER BASED

def resample_filter_2d_separable(
    tensor: Tensor,
    new_lens: tuple[int, int],
    radius: int,
    filters: tuple[Filter, Filter],
    dims: tuple[int, int],
) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_positive(new_lens)",
        "radius > 0",
    ])
    D_0 = _filter_1d(tensor, new_lens[0], radius, filters[0], dims[0])
    D_1 = _filter_1d(D_0, new_lens[1], radius, filters[1], dims[1])
    resampled = D_1.to(tensor.dtype)
    return resampled


def resample_filter_2d(
    tensor: Tensor,
    new_lens: tuple[int, int],
    radius: int,
    filter: Filter,
    dims: tuple[int, int],
) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_positive(new_lens)",
        "radius > 0",
    ])
    resampled = _filter_2d(tensor, new_lens, radius, filter, dims).to(tensor.dtype)
    return resampled


def _window_lengths(
    l: int,
    L: int,
    radius: int,
) -> int:
    if radius == 0:
        w, W = 1, 1
    elif L <= l:
        W = 2 * radius
        w = -((-l * W) // L)
    else:
        w = 2 * radius
        W = -((-L * w) // l)
    return (w, W)


def _padding_1d(
    i_snn: Tensor,
    I_t: Tensor,
    l: int,
    L: int,
    W: int,
    w: int,
) -> Tensor:
    if w == 1 and W == 1:
        p = i_snn * 0
    elif L <= l:
        p = (W * l + L * (2 * i_snn + 1) - l * (2 * I_t + 1)) // (2 * L)
    else:
        c = l * (2 * I_t + 1) > L * (2 * i_snn + 1)
        p = (w // 2) - c.to(torch.int)
    return p


def _x_values_1d(
    i_snn: Tensor,
    i_w: Tensor,
    I_t: Tensor,
    l: int,
    L: int,
    p: Tensor
) -> Tensor:
    X_w_n = _outer_sum(
        (2 * (i_snn - p) + 1) * L - (2 * I_t + 1) * l, (2 * L) * i_w
    )
    if L <= l:
        x_values = X_w_n / (2 * l)
    else:
        # breaks notation convention
        x_values = X_w_n / (2 * L)
    return x_values


def _data_values_1d(
    d_p: Tensor,
    i_snn: Tensor,
    i_w: Tensor,
    p: Tensor,
    p0: int,
    L: int,
    w: int,
    dim: int = 0
) -> Tensor:
    o = i_snn - p + p0
    i_2d = _outer_sum(o, i_w)
    i_1d = i_2d.flatten()
    d_1d = d_p.index_select(dim, i_1d)
    shape = list(d_p.shape)
    shape = shape[:dim] + [L, w] + shape[dim+1:]
    d_2d = d_1d.reshape(*shape)
    return d_2d


def _filter_1d(
    d_s: Tensor,
    L: int,
    radius: int,
    filter: Filter,
    dim: int = 0,
) -> Tensor:
    l = int(d_s.shape[dim])
    (w, W) = _window_lengths(l, L, radius)
    I_t = torch.arange(L)
    i_snn = _nearest_neighbor_indices(I_t, l, L)
    p = _padding_1d(i_snn, I_t, l, L, W, w)
    i_w = torch.arange(w)
    x_values = _x_values_1d(i_snn, i_w, I_t, l, L, p)
    p_0 = int(p[0])
    d_p = transform_pad_dim_reflect(d_s, dim, (p_0, p_0 + 1))
    d_2d = _data_values_1d(d_p, i_snn, i_w, p, p_0, L, w, dim)
    f = filter(x_values)
    f_n = util_sum_normalize(f, (1,))
    shape = [1] * d_2d.dim()
    shape[dim] = L
    shape[dim+1] = w
    d_f = d_2d * f_n.reshape(*shape)
    d_r = d_f.sum(dim+1)
    return d_r


def _x_values_2d(
    X_w: tuple[Tensor, Tensor],
    L: tuple[int, int],
    w: tuple[int, int],
) -> Tensor:
    X_w_sqr = (X_w[0] * X_w[0], X_w[1] * X_w[1])
    X_w2_sqr = X_w_sqr[0].reshape(L[0], 1, w[0], 1) + \
        X_w_sqr[1].reshape(1, L[1], 1, w[1])
    X_w2 = torch.sqrt(X_w2_sqr)
    return X_w2


def _data_values_2d(
    d_p: Tensor,
    i_snn: tuple[Tensor, Tensor],
    p: tuple[Tensor, Tensor],
    i_w: tuple[Tensor, Tensor],
    p0: tuple[int, int],
    w: tuple[int, int],
    L: tuple[int, int],
    l_p: tuple[int, int],
    dim: tuple[int, int],
) -> Tensor:
    o = (i_snn[0] - p[0] + p0[0], i_snn[1] - p[1] + p0[1])
    i_2d = (_outer_sum(o[0], i_w[0]), _outer_sum(o[1], i_w[1]))
    i_4d = (i_2d[0] * l_p[1]).reshape(L[0], 1, w[0], 1) + \
        i_2d[1].reshape(1, L[1], 1, w[1])
    i_1d = i_4d.flatten()
    d_p = d_p.transpose(dim[1], -1)
    d_p = d_p.transpose(dim[0], -2)
    d_1d = d_p.flatten(start_dim=-2, end_dim=-1).index_select(-1, i_1d)
    shape = list(d_1d.shape)[:-1]
    d_4d = d_1d.reshape(*shape, L[0], L[1], w[0], w[1])
    d_4d = d_4d.transpose(dim[0], -4)
    d_4d = d_4d.transpose(dim[1], -3)
    return d_4d


def _filter_2d(
    d_s: Tensor,
    L: tuple[int, int],
    radius: int,
    filter: Filter,
    dim: tuple[int, int],
) -> Tensor:
    l = (int(d_s.shape[dim[0]]), int(d_s.shape[dim[1]]))
    (w, W) = tuple(zip(
        _window_lengths(l[0], L[0], radius),
        _window_lengths(l[1], L[1], radius)
    ))
    I_t = (torch.arange(L[0]), torch.arange(L[1]))
    i_snn = (
        _nearest_neighbor_indices(I_t[0], l[0], L[0]),
        _nearest_neighbor_indices(I_t[1], l[1], L[1])
    )
    p = (
        _padding_1d(i_snn[0], I_t[0], l[0], L[0], W[0], w[0]),
        _padding_1d(i_snn[1], I_t[1], l[1], L[1], W[1], w[0])
    )
    i_w = (torch.arange(w[0]), torch.arange(w[1]))
    X_w = (
        _x_values_1d(i_snn[0], i_w[0], I_t[0], l[0], L[0], p[0]),
        _x_values_1d(i_snn[1], i_w[1], I_t[1], l[1], L[1], p[1])
    )
    X_w2 = _x_values_2d(X_w, L, w)
    p0 = (int(p[0][0]), int(p[1][0]))
    d_p = transform_pad_dim_reflect(d_s, dim[0], (p0[0], p0[0] + 1))
    d_p = transform_pad_dim_reflect(d_p, dim[1], (p0[1], p0[1] + 1))
    l_p = (int(d_p.shape[dim[0]]), int(d_p.shape[dim[1]]))
    d_4d = _data_values_2d(d_p, i_snn, p, i_w, p0, w, L, l_p, dim)
    f = filter(X_w2)
    f_n = util_sum_normalize(f, (2, 3))
    shape = [1] * d_4d.dim()
    shape[dim[0]] = L[0]
    shape[dim[1]] = L[1]
    shape[-2] = w[0]
    shape[-1] = w[1]
    d_f = d_4d * f_n.reshape(*shape)
    d_r = d_f.sum(-1).sum(-1)
    return d_r


# HELPERS

def _outer_sum(lhs: Tensor, rhs: Tensor) -> Tensor:
    ret = lhs.unsqueeze(1) + rhs.unsqueeze(0)
    return ret


#                           #
# === SPECIAL FUNCTIONS === #
#                           #

# FUNCTIONS

def special_jinc(x: Tensor) -> Tensor:
    _check([
        "x.is_floating_point()",
    ])
    p0: Tensor = (2 / pi) * torch.special.bessel_j1(pi * x) / x
    j = p0.where(x != 0, 1)
    return j


def special_gaussian(x: Tensor, sigma: float) -> Tensor:
    _check([
        "x.is_floating_point()",
        "sigma > 0",
    ])
    coef = 1 / (sigma * sqrt(2 * pi))
    g = coef * torch.exp(x**2 / (-2 * sigma**2))
    return g


def special_derivative_of_gaussian(x: Tensor, sigma: float) -> Tensor:
    _check([
        "x.is_floating_point()",
        "sigma > 0",
    ])
    coef = -1 / (sigma**3 * sqrt(2 * pi))
    dog = coef * x * torch.exp(x**2 / (-2 * sigma**2))
    return dog


# HELPERS

def special_gaussian_area_radius(area: float, sigma: float) -> float:
    _check([
        "sigma > 0",
        "0 < area < 1",
    ])
    radius = sqrt(2) * sigma * erfinv(area)
    return radius


def special_derivative_of_gaussian_area_radius(area: float, sigma: float) -> float:
    _check([
        "sigma > 0",
        "0 < area < 1",
    ])
    radius = sqrt(2) * sigma * sqrt(log(1 / (1 - area)))
    return radius


#                     #
# === TONEMAPPING === #
#                     #

# All based on https://64.github.io/tonemapping/

def tonemap_reinhard(tensor: Tensor) -> Tensor:
    _check([
        "tensor.is_floating_point()",
    ])
    tonemapped = tensor / (1 + tensor)
    return tonemapped


def tonemap_reinhard_luminance(tensor: Tensor, luminance: Tensor) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "luminance.is_floating_point()",
        "_is_broadcastable_to(luminance, tensor)",
    ])
    tonemapped = tensor / (1 + luminance)
    return tonemapped


def tonemap_reinhard_extended(tensor: Tensor, whitepoint: Tensor | float) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "type(whitepoint) == float or whitepoint.is_floating_point()",
        "_is_broadcastable_to(whitepoint, tensor)",
    ])
    tonemapped = (tensor * (1 + (tensor / (whitepoint ** 2)))) / (1 + tensor)
    return tonemapped


def tonemap_reinhard_extended_luminance(tensor: Tensor, luminance: Tensor, whitepoint: Tensor | float) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "luminance.is_floating_point()",
        "_is_broadcastable_to(luminance, tensor)",
        "type(whitepoint) == float or whitepoint.is_floating_point()",
        "_is_broadcastable_to(whitepoint, tensor)",
    ])
    tonemapped = (tensor * (1 + (luminance / (whitepoint ** 2)))) / (1 + luminance)
    return tonemapped


def tonemap_reinhard_jodie(tensor: Tensor, luminance: Tensor) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "luminance.is_floating_point()",
        "_is_broadcastable_to(luminance, tensor)",
    ])
    t_reinhard = tonemap_reinhard(tensor)
    t_reinhard_luminance = tonemap_reinhard_luminance(tensor, luminance)
    lerped = torch.lerp(t_reinhard_luminance, t_reinhard, t_reinhard)
    return lerped


def tonemap_reinhard_jodie_extended(tensor: Tensor, luminance: Tensor, whitepoint: Tensor | float) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "luminance.is_floating_point()",
        "_is_broadcastable_to(luminance, tensor)",
        "type(whitepoint) == float or whitepoint.is_floating_point()",
        "_is_broadcastable_to(whitepoint, tensor)",
    ])
    t_reinhard = tonemap_reinhard_extended(tensor, whitepoint)
    t_reinhard_luminance = tonemap_reinhard_extended_luminance(tensor, luminance, whitepoint)
    lerped = torch.lerp(t_reinhard_luminance, t_reinhard, t_reinhard)
    return lerped


def tonemap_uncharted_2(tensor: Tensor) -> Tensor:
    _check([
        "tensor.is_floating_point()",
    ])
    def tonemap_partial(t: Tensor | float) -> Tensor | float:
        A, B, C, D, E, F = 0.15, 0.5, 0.1, 0.2, 0.02, 0.3
        partial = ((t*(A*t+C*B)+D*E)/(t*(A*t+B)+D*F))-E/F
        return partial
    EXPOSURE_BIAS = 2.0
    W = 11.2
    partial = tonemap_partial(tensor * EXPOSURE_BIAS)
    white_scale = 1 / tonemap_partial(W)
    tonemapped = partial * white_scale
    return tonemapped


def tonemap_aces(tensor: Tensor, channel_dim: int) -> Tensor:
    _check([
        "tensor.is_floating_point()",
        "_valid_dim(tensor, channel_dim)",
        "tensor.shape[channel_dim] == 3",
    ])
    INPUT_MAT = torch.tensor((
        (0.59719, 0.35458, 0.04823),
        (0.07600, 0.90834, 0.01566),
        (0.02840, 0.13383, 0.83777),
    ))
    OUTPUT_MAT = torch.tensor((
        (+1.60475, -0.53108, -0.07367),
        (-0.10208, +1.10813, -0.00605),
        (-0.00327, -0.07276, +1.07602),
    ))

    permuted = utils_dims_to_end(tensor, (channel_dim, ))
    input = (permuted.unsqueeze(-2) * INPUT_MAT).sum(-1)

    a = input * (input + 0.0245786) - 0.000090537
    b = input * (0.983729 * input + 0.4329510) + 0.238081
    rtt_odt_fit = a / b

    tonemapped = (rtt_odt_fit.unsqueeze(-2) * OUTPUT_MAT).sum(-1).clamp(0, 1)
    unpermuted = utils_dims_from_end(tonemapped, (channel_dim, ))
    return unpermuted


#                    #
# === TRANSFORMS === #
#                    #

def transform_crop_dim(tensor: Tensor, dim: int, crop: tuple[int, int]) -> Tensor:
    _check([
        "_valid_dim(tensor, dim)",
        "_nonnegative(crop)",
    ])
    lower = crop[0]
    upper = crop[1] if crop[1] != 0 else None
    cropped = tensor.clone()[util_slice_dim(dim, lower, upper)]
    return cropped


def transform_pad_dim_reflect(tensor: Tensor, dim: int, pad: tuple[int, int]) -> Tensor:
    _check([
        "_valid_dim(tensor, dim)",
        "_nonnegative(pad)",
        "tensor.shape[dim] >= max(pad) + 1",
    ])
    length = tensor.shape[dim]
    start = tensor[util_slice_dim(dim, 1, pad[0] + 1)].flip(dim)
    end = tensor[util_slice_dim(dim, length - (pad[1] + 1), -1)].flip(dim)
    padded = torch.cat((start, tensor, end), dim)
    return padded


def transform_pad_dim2_reflect(tensor: Tensor, dims: tuple[int, int], pad: tuple[tuple[int, int], tuple[int, int]]) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_nonnegative(pad[0]) and _nonnegative(pad[1])",
    ])
    p0 = tensor
    p1 = transform_pad_dim_reflect(p0, dims[0], pad[0])
    p2 = transform_pad_dim_reflect(p1, dims[1], pad[1])
    return p2


def transform_pad_dim_zero(tensor: Tensor, dim: int, pad: tuple[int, int]) -> Tensor:
    _check([
        "_valid_dim(tensor, dim)",
        "_nonnegative(pad)",
    ])
    dims = tensor.dim()
    pad_list = [0, 0] * (dims - dim - 1) + list(pad)
    padded = F.pad(tensor, pad_list)
    return padded


def transform_pad_dim2_zero(tensor: Tensor, dims: tuple[int, int], pad: tuple[tuple[int, int], tuple[int, int]]) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_nonnegative(pad[0]) and _nonnegative(pad[1])",
    ])
    p0 = tensor
    p1 = transform_pad_dim_zero(p0, dims[0], pad[0])
    p2 = transform_pad_dim_zero(p1, dims[1], pad[1])
    return p2


def transform_uncrop_from_bbox(tensor: Tensor, bbox: BBox, dim_h: int, dim_w: int) -> Tensor:
    _check([
        "bbox.valid",
        "_valid_dims(tensor, (dim_h, dim_w))",
        "tensor.shape[dim_h] == bbox.height",
        "tensor.shape[dim_w] == bbox.width",
    ])
    uncropped_h = transform_pad_dim_zero(tensor, dim_h, (bbox.up, bbox.down))
    uncropped = transform_pad_dim_zero(uncropped_h, dim_w, (bbox.left, bbox.right))
    return uncropped


def transform_crop_to_bbox(tensor: Tensor, bbox: BBox, dim_h: int, dim_w: int) -> Tensor:
    _check([
        "bbox.valid",
        "_valid_dims(tensor, (dim_h, dim_w))",
        "tensor.shape[dim_h] == bbox.total_height",
        "tensor.shape[dim_w] == bbox.total_width",
    ])
    c0 = tensor
    c1 = c0[util_slice_dim(dim_h, bbox.up, bbox.down_offset)]
    c2 = c1[util_slice_dim(dim_w, bbox.left, bbox.right_offset)]
    return c2


def transform_connected_components_from_2d_binary(tensor: Tensor, diagonals: bool = False) -> list[Tensor]:
    (coloring, max_unique) = convert_component_coloring(tensor, diagonals)
    if max_unique == 0:
        return [coloring]
    colorings = []
    for i in range(1, max_unique + 1):
        colorings += [(coloring == i)]
    return colorings


#                   #
# === WINDOWING === #
#                   #

# FUNCTIONS

def window_boxcar(x: Tensor, radius: int) -> Tensor:
    _check([
        "radius >= 0",
    ])
    windowed = torch.abs(x) <= radius
    return windowed


def window_triangle(x: Tensor, radius: int) -> Tensor:
    _check([
        "radius >= 0",
    ])
    windowed = ((radius - torch.abs(x)) / radius).clamp(0, 1)
    return windowed


def window_lanczos(x: Tensor, radius: int) -> Tensor:
    _check([
        "radius > 0",
    ])
    windowed = torch.sinc(x / radius) * (torch.abs(x) <= radius)
    return windowed


def window_mitchell_netravali(x: Tensor, b: float, c: float) -> Tensor:
    _check([
        "0 <= b <= 1",
        "0 <= c <= 1",
    ])
    x_abs = torch.abs(x)
    x_abs2 = x_abs * x_abs
    x_abs3 = x_abs2 * x_abs
    p0 = (12 + -9*b + -6*c)*x_abs3 + (-18 + 12*b + 6*c)*x_abs2 + (6 + -2*b)
    p1 = (-b + -6*c)*x_abs3 + (6*b + 30*c)*x_abs2 + (-12*b + -48*c)*x_abs + (8*b + 24*c)
    p0_mask = x_abs < 1
    p1_mask = torch.logical_and(torch.logical_not(p0_mask), x_abs < 2)
    p = p0 * p0_mask + p1 * p1_mask
    mn = p / 6
    return mn


def window_area(x: Tensor, old_len: int, new_len: int) -> Tensor:
    _check([
        "old_len > 0",
        "new_len > 0",
    ])
    l, L = (old_len, new_len) if new_len <= old_len else (new_len, old_len)
    windowed = (((L + l) - (2 * l) * torch.abs(x)) / (2 * L)).clamp(0, 1)
    return windowed


# HELPERS

def window_mitchell_netravali_radius() -> int:
    return 2


def window_area_radius() -> int:
    return 1


#                   #
# === UTILITIES === #
#                   #

def util_round_up_to_mult_of(number: int, mult_of: int) -> int:
    aligned = ((number + mult_of - 1) // mult_of) * mult_of
    return aligned


def util_delta_split(delta: int, t: float) -> tuple[int, int]:
    upper = int(delta * t)
    lower = delta - upper
    deltas = (lower, upper)
    return deltas


def util_tensor_min(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    curr = tensor
    for dim in dims:
        curr = torch.amin(curr, dim, True)
    return curr


def util_tensor_max(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    curr = tensor
    for dim in dims:
        curr = torch.amax(curr, dim, True)
    return curr


def util_tensor_mean(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    curr = tensor
    for dim in dims:
        curr = torch.mean(curr, dim, True)
    return curr


def util_slice_dim(
    dim: int,
    start: int | None,
    stop: int | None,
    step: int | None = 1
) -> list[slice]:
    _check([
        "dim >= 0",
    ])
    slices = [slice(None) for _ in range(dim)] + [slice(start, stop, step)]
    return slices


def util_sum_normalize(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    dims = sorted(dims)
    denominator = tensor
    for dim in reversed(dims):
        denominator = denominator.sum(dim)
    for dim in dims:
        denominator = denominator.unsqueeze(dim)
    normalized = tensor / denominator
    return normalized


def util_range_normalize(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    cur_min = gtf_min(tensor, dims)
    minned = tensor - cur_min
    cur_max = gtf_max(minned, dims)
    maxxed = minned / cur_max
    clamped = maxxed.clamp(0, 1)
    return clamped


def util_polar_from_cartesian(tensor_x: Tensor, tensor_y: Tensor) -> tuple[Tensor, Tensor]:
    r = torch.hypot(tensor_x, tensor_y)
    theta = torch.atan2(tensor_y, tensor_x).nan_to_num()
    return (r, theta)


def util_cartesian_from_polar(tensor_r: Tensor, tensor_theta: Tensor) -> tuple[Tensor, Tensor]:
    x = tensor_r * torch.cos(tensor_theta)
    y = tensor_r * torch.sin(tensor_theta)
    return (x, y)


def utils_dims_to_end(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_ascending(dims)",
    ])
    permutation = [x for x in range(tensor.dim()) if x not in dims] + list(dims)
    permuted = tensor.permute(*permutation)
    return permuted


def utils_dims_from_end(tensor: Tensor, dims: Iterable[int]) -> Tensor:
    _check([
        "_valid_dims(tensor, dims)",
        "_ascending(dims)",
    ])
    unpermutation = [x for x in range(tensor.dim() - len(dims))]
    dim_indices = [x for x in range(tensor.dim() - len(dims), tensor.dim())]
    for dim, index in zip(reversed(dims), reversed(dim_indices)):
        unpermutation = unpermutation[:dim] + [index] + unpermutation[dim:]
    unpermuted = tensor.permute(unpermutation)
    return unpermuted


def util_shift_mean_terms(tensor: Tensor, new_mean: float, dims: Iterable[int], domain: tuple[float, float]) -> tuple[Tensor, Tensor]:
    _check([
        "_valid_dims(tensor, dims)",
        "domain[0] <= new_mean <= domain[1]",
        "domain[0] <= torch.min(tensor)",
        "domain[1] >= torch.max(tensor)",
    ])
    dmin = domain[0]
    dmax = domain[1]
    mean = util_tensor_mean(tensor, dims)
    minimum = util_tensor_min(tensor, dims)
    maximum = util_tensor_max(tensor, dims)
    delta = new_mean - mean 
    scale = torch.ones(*mean.shape)
    max_scale = (dmax - new_mean) / (maximum - mean)
    min_scale = (dmin - new_mean) / (minimum - mean)
    max_delta = dmax - max_scale * maximum
    min_delta = dmin - min_scale * minimum
    need_max = maximum + delta > dmax
    need_min = minimum + delta < dmin
    scale = torch.where(need_max, max_scale, scale)
    scale = torch.where(need_min, min_scale, scale)
    delta = torch.where(need_max, max_delta, delta)
    delta = torch.where(need_min, min_delta, delta)
    return (scale, delta)


def util_shift_mean(tensor: Tensor, terms: tuple[Tensor, Tensor]) -> Tensor:
    shifted = tensor * terms[0] + terms[1]
    return shifted


def util_unshift_mean(tensor: Tensor, terms: tuple[Tensor, Tensor]) -> Tensor:
    unshifted = (tensor - terms[1]) / terms[0]
    return unshifted


#                             #
# === CHECKING (INTERNAL) === #
#                             #

def _check(preconditions):
    caller_frame = inspect.currentframe().f_back
    locals_ = caller_frame.f_locals
    globals_ = caller_frame.f_globals
    function = globals_[caller_frame.f_code.co_name]

    # Check types
    for arg_name, arg_type in function.__annotations__.items():
        if arg_name in locals_:
            if not _is_type(locals_[arg_name], arg_type):
                raise TypeError(f"Violated: type({arg_name}) == {arg_type}")

    # Check preconditions
    for precondition in preconditions:
        if eval(precondition, globals_, locals_) == False:
            raise ValueError(f"Violated: {precondition}")


def _is_type(value, expected_type) -> bool:
    if isinstance(expected_type, GenericAlias):
        origin_type = expected_type.__origin__
        if not isinstance(value, origin_type):
            return False
        ga_args = expected_type.__args__
        if origin_type == tuple:
            if len(value) != len(ga_args):
                return False
            for v, t in zip(value, ga_args):
                if not _is_type(v, t):
                    return False
        elif origin_type == list:
            if len(ga_args) != 1:
                return False
            for v in value:
                if not _is_type(v, ga_args[0]):
                    return False
        else:
            raise NotImplementedError(f"Type checking not yet implemented for subscripted {vtype}.")
        return True 
    elif isinstance(expected_type, _CallableGenericAlias):
        # We don't do param/return type checking
        origin_type = expected_type.__origin__
        return isinstance(value, origin_type)
    elif isinstance(expected_type, _GenericAlias):
        if expected_type.__origin__ == Iterable.__origin__:
            if not isinstance(value, Iterable):
                return False
            if len(expected_type.__args__) != 1:
                raise NotImplementedError()
            for v in value:
                if not _is_type(v, expected_type.__args__[0]):
                    return False
            return True
        else:
            raise NotImplementedError()
    elif isinstance(expected_type, UnionType):
        ut_args = expected_type.__args__
        for t in ut_args:
            if _is_type(value, t):
                return True
        return False
    else:
        return isinstance(value, expected_type)


def _valid_dim(tensor: Tensor, dim: int) -> bool:
    if dim >= 0:
        return 0 <= dim < tensor.dim()
    else:
        return -tensor.dim() <= dim < 0


def _valid_dims(tensor: Tensor, dims: Iterable[int]) -> bool:
    for dim in dims:
        if not _valid_dim(tensor, dim):
            return False
    return True


def _nonnegative(iter: Iterable) -> bool:
    for value in iter:
        if value < 0:
            return False
    return True


def _positive(iter: Iterable) -> bool:
    for value in iter:
        if value <= 0:
            return False
    return True


def _is_broadcastable_to(from_: Tensor | int | float, to_: Tensor) -> bool:
    if type(from_) in (int, float) or from_.dim() == 0:
        return True
    if from_.dim() > to_.dim():
        return False
    for f, t in zip(reversed(from_.shape), reversed(to_.shape)):
        if f != 1 and f != t:
            return False
    return True


def _ascending(iter: Iterable) -> bool:
    for i in range(1, len(iter)):
        if iter[i] <= iter[i - 1]:
            return False
    return True