"""Global correction of the network map by the numbered grid nodes, and the self-check.

The network map F places node (i, j) at F(node); a cubic polynomial T is fitted so that
T(F(node)) lies on the 5 mm lattice, and the corrected map is T(F). A second candidate is a
homography fitted to the nodes directly.

Self-check: nodes are split like a chessboard by the parity of i + j; the correction is fitted on
the even nodes and measured on the odd ones after the best similarity to the lattice. It measures
the SHAPE of the lattice (the global scale is fitted out) and is blind to detector errors.
"""

import cv2
import numpy as np

OUTLIER_MM = 1.0


def design(z):
    """Cubic polynomial columns of (x, y) = (Re z, Im z)."""
    x, y = z.real, z.imag
    return np.stack([np.ones_like(x), x, y, x * x, x * y, y * y, x ** 3, x * x * y, x * y * y, y ** 3], 1)


def fit_correction(mapped, ij):
    """Fit T: mapped node positions (complex mm) -> 5 * (i + 1j * j), refit without > OUTLIER_MM outliers.
    The lattice may come mirrored (j up), so both orientations are tried; reflections are rejected.
    -> function applying T to map values."""
    target = 5.0 * (ij[:, 0] + 1j * ij[:, 1])
    centre = mapped.mean()
    z = mapped - centre
    scale = np.abs(z).max()
    keep = np.ones(len(z), bool)
    for _ in range(2):
        best = None
        for q in (target, np.conj(target)):
            m = design(z[keep] / scale)
            cu, *_ = np.linalg.lstsq(m, q[keep].real, rcond=None)
            cv, *_ = np.linalg.lstsq(m, q[keep].imag, rcond=None)
            if cu[1] * cv[2] - cu[2] * cv[1] <= 0:
                continue
            residual = np.abs(design(z / scale) @ cu + 1j * (design(z / scale) @ cv) - q)
            if best is None or np.median(residual) < np.median(best[2]):
                best = (cu, cv, residual)
        cu, cv, residual = best
        keep = residual < OUTLIER_MM

    def apply(values):
        m = design((values - centre) / scale)
        return m @ cu + 1j * (m @ cv)
    return apply


def _fit_similarity(p, q):
    """Residuals |p - a*q - b| of the best similarity (complex a, b)."""
    pc, qc = p - p.mean(), q - q.mean()
    a = (np.conj(qc) * pc).sum() / (np.abs(qc) ** 2).sum()
    return np.abs(pc - a * qc)


def score(p, ij):
    """Residuals (mm) of points p against the 5 mm lattice ij after the best similarity (mirror allowed)."""
    q = 5.0 * (ij[:, 0] + 1j * ij[:, 1])
    r1, r2 = _fit_similarity(p, q), _fit_similarity(p, np.conj(q))
    r = r1 if np.sqrt((r1 ** 2).mean()) <= np.sqrt((r2 ** 2).mean()) else r2
    trimmed = np.sort(r)[:int(0.98 * len(r))]
    return {"rms_mm": float(np.sqrt((r ** 2).mean())), "rms98_mm": float(np.sqrt((trimmed ** 2).mean())),
            "median_mm": float(np.median(r)), "p95_mm": float(np.percentile(r, 95)), "points": int(len(r))}


def halves(ij):
    """Chessboard split: even i + j to fit, odd to check."""
    return (ij.sum(1) % 2) == 0, (ij.sum(1) % 2) == 1


def holdout_network(mesh, nodes, pts, ij):
    fit, test = halves(ij)
    apply = fit_correction(mesh.sample(nodes, pts[fit]), ij[fit])
    return score(apply(mesh.sample(nodes, pts[test])), ij[test])


def homography_nodes(mesh, pts, ij):
    """Homography map from the nodes, evaluated at the mesh nodes; None if it fails."""
    h, _ = cv2.findHomography(pts.astype(np.float64), 5.0 * ij.astype(np.float64), cv2.RANSAC, 5.0)
    if h is None:
        return None
    gx, gy = np.meshgrid(np.arange(mesh.nx) * mesh.step, np.arange(mesh.ny) * mesh.step)
    q = cv2.perspectiveTransform(np.stack([gx.ravel(), gy.ravel()], 1).reshape(-1, 1, 2).astype(np.float64), h)
    return q[:, 0, 0] + 1j * q[:, 0, 1]


def holdout_homography(mesh, pts, ij):
    fit, test = halves(ij)
    h_nodes = homography_nodes(mesh, pts[fit], ij[fit])
    if h_nodes is None:
        return None
    return score(mesh.sample(h_nodes, pts[test]), ij[test])