"""Flat sheet: inversion of the map pixel -> mm and resampling of the original photo.

The inverse is solved by Newton's method on every `coarse`-th output pixel and interpolated
bilinearly in between.
"""

import numpy as np


def _upsample(a, k):
    """Bilinear upsampling by k with corner-aligned samples: out[i*k] = a[i]."""
    a = a.astype(np.float32)
    ny, nx = a.shape
    yi = np.arange((ny - 1) * k + 1) / k
    xi = np.arange((nx - 1) * k + 1) / k
    y0 = np.minimum(np.floor(yi).astype(int), ny - 2)
    x0 = np.minimum(np.floor(xi).astype(int), nx - 2)
    fy, fx = (yi - y0)[:, None], (xi - x0)[None, :]
    rows = a[y0] * (1 - fy) + a[y0 + 1] * fy
    return rows[:, x0] * (1 - fx) + rows[:, x0 + 1] * fx


def inverse_map(original_shape, factor, mesh, nodes, out_ppm, iterations=8, coarse=8):
    """map_x, map_y for cv2.remap: original-photo coordinates of every flat-sheet pixel at out_ppm px/mm
    (-1 where outside the photo or not converged). factor: original photo -> 5.5 px/mm scale."""
    height, width = int(round(original_shape[0] * factor)), int(round(original_shape[1] * factor))
    node_px = np.stack(np.meshgrid(np.arange(mesh.nx) * mesh.step, np.arange(mesh.ny) * mesh.step), -1).reshape(-1, 2)
    inside = (node_px[:, 0] <= width - 1 + mesh.step) & (node_px[:, 1] <= height - 1 + mesh.step)
    design = np.stack([nodes.real, nodes.imag, np.ones(mesh.count)], 1)[inside]
    affine, *_ = np.linalg.lstsq(design, node_px[inside], rcond=None)
    border = np.concatenate([np.stack([np.linspace(0, width - 1, 64), np.zeros(64)], 1),
                             np.stack([np.linspace(0, width - 1, 64), np.full(64, height - 1.0)], 1),
                             np.stack([np.zeros(64), np.linspace(0, height - 1, 64)], 1),
                             np.stack([np.full(64, width - 1.0), np.linspace(0, height - 1, 64)], 1)])
    paper = mesh.sample(nodes, border)
    n_u = int((paper.real.max() - paper.real.min()) * out_ppm)
    n_v = int((paper.imag.max() - paper.imag.min()) * out_ppm)
    u = paper.real.min() + np.arange(0, n_u + coarse, coarse) / out_ppm
    v = paper.imag.min() + np.arange(0, n_v + coarse, coarse) / out_ppm
    uu, vv = np.meshgrid(u, v)
    target = uu + 1j * vv
    p = np.stack([uu, vv, np.ones_like(uu)], -1) @ affine
    eps = 0.5
    for _ in range(iterations):
        f = mesh.sample(nodes, p)
        fx = (mesh.sample(nodes, p + [eps, 0.0]) - mesh.sample(nodes, p - [eps, 0.0])) / (2 * eps)
        fy = (mesh.sample(nodes, p + [0.0, eps]) - mesh.sample(nodes, p - [0.0, eps])) / (2 * eps)
        a, b, c, d = fx.real, fy.real, fx.imag, fy.imag
        det = np.where(np.abs(a * d - b * c) < 1e-12, 1e-12, a * d - b * c)
        r = target - f
        p = p + np.stack([np.clip((d * r.real - b * r.imag) / det, -20, 20),
                          np.clip((-c * r.real + a * r.imag) / det, -20, 20)], -1)
    residual = np.abs(mesh.sample(nodes, p) - target)
    valid = (residual < 0.05) & (p[..., 0] >= 0) & (p[..., 0] <= width - 1) & (p[..., 1] >= 0) & (p[..., 1] <= height - 1)

    def up(a):
        return _upsample(a, coarse)[:n_v, :n_u]
    map_x = up((p[..., 0] + 0.5) / factor - 0.5)
    map_y = up((p[..., 1] + 0.5) / factor - 0.5)
    ok = up(valid.astype(np.float32)) > 0.999
    return np.where(ok, map_x, -1).astype(np.float32), np.where(ok, map_y, -1).astype(np.float32)


def crop_to(image, keep, pad=8):
    """Crop to the bounding box of the bool mask keep (+ pad px)."""
    ys, xs = np.nonzero(keep)
    if len(xs) == 0:
        return image
    return image[max(0, ys.min() - pad):ys.max() + pad + 1, max(0, xs.min() - pad):xs.max() + pad + 1]