#!/usr/bin/python3
"""Euclidean distance transform.

This is my misinterpretation of the non-subpixel part of Steven
Wittens’s <https://acko.net/blog/subpixel-distance-transform/>.

Run tests with pytest euclideandt.py.  They’re broken.

"""


def hordt(xs, inf=255):
    "Compute the squared Euclidean distance within a row."
    distsq = delta = inf
    result = []
    for x in xs:
        if x < distsq:
            distsq, delta = x, 1
        result.append(distsq)
        distsq += delta
        delta += 2

    distsq = delta = inf
    for i in range(len(xs))[::-1]:
        x = xs[i]
        if x < distsq:
            distsq, delta = x, 1
        result[i] = min(distsq, result[i])
        distsq += delta
        delta += 2

    return result


def test_hordt():
    assert (hordt([255,   0, 255, 255, 255, 255, 255, 0, 255]) ==
                  [  1,   0,   1,   4,   9,   4,   1, 0,   1])
    assert (hordt([255, 255, 255, 255, 255, 255, 255, 0, 255]) ==
                  [ 49,  36,  25,    16,   9,   4, 1, 0,  1])
    assert (hordt([255, 255, 255, 255, 255, 255, 255]) ==
                  [255, 255, 255, 255, 255, 255, 255])


class Column:
    def __init__(self, matrix, j):
        self.matrix = matrix
        self.j = j

    def __len__(self):
        return len(self.matrix)

    def __getitem__(self, i):
        if type(i) is slice:
            return [self[n] for n in slice_to_range(i, 0, len(self.matrix))]

        else:
            return self.matrix[i][self.j]

    def __setitem__(self, i, x):
        if type(i) is slice:
            indices = slice_to_range(i, 0, len(self.matrix))
            assert len(x) == len(indices)
            for index, xi in zip(indices, x):
                self[index] = xi

        else:
            self.matrix[i][self.j] = x


def slice_to_range(s, start, stop):
    start = (start if s.start is None or s.start <= start else
             s.start if s.start >= 0 else
             stop + s.start)
    stop = (stop if s.stop is None or s.stop >= stop else
            s.stop if s.stop >= 0 else
            stop + s.stop)
    step = (1 if s.step is None else
            s.step)
    return range(start, stop, step)


def test_column():
    m = [[1, 2], [3, 4]]
    assert Column(m, 0)[1] == 3
    Column(m, 0)[1] = 53
    assert m == [[1, 2], [53, 4]]
    Column(m, 1)[:] = [5353, 37]
    assert m == [[1, 5353], [53, 37]]


def dtsq(matrix):
    result = [hordt(matrix[i]) for i in range(len(matrix))]
    for j in range(len(result[0])):
        col = Column(result, j)
        col[:] = hordt(col)

    return result


def test_dt():
    assert (dtsq([[255, 255, 255, 255],
                  [255, 255, 255, 255],
                  [255, 255,   0, 255],
                  [255, 255, 255, 255],
                  [255, 255, 255, 255]]) ==
                 [[  8,   5,   4,   5],
                  [  5,   2,   1,   2],
                  [  4,   1,   0,   1],
                  [  5,   2,   1,   2],
                  [  8,   5,   4,   5]])

    assert (dtsq([[255, 255, 255, 255],
                  [  0, 255, 255, 255],
                  [255, 255,   0, 255],
                  [255, 255, 255, 255],
                  [255, 255, 255, 255]]) ==
                 [[  1,   2,   4,   5],
                  [  0,   1,   1,   2],
                  [  1,   1,   0,   1],
                  [  4,   2,   1,   2],
                  [  8,   5,   4,   5]])  # XXX fails