"""
features_ransac.py: fit a homography to feature matches when some of them are wrong.

EECE 5554 Robotics Sensing and Navigation, demo for Lecture 13.

Two photos of a flat scene are related by one homography H (Lecture 12). A
feature matcher (ORB or SIFT plus the ratio test) gives point pairs, but some
pairs are always wrong. This script makes such a set of matches, then fits H
two ways:

    least squares   the DLT on every match: one wrong match can ruin it
    RANSAC          sample 4 pairs, fit, count the inliers, repeat, keep the
                    best, and refit on its inliers

It also computes how many RANSAC tries are needed for a given share of wrong
matches e and sample size s, with probability p of at least one clean sample:

    N = log(1 - p) / log(1 - (1 - e)^s)

Images are 1616 x 1240 pixels, the size of the course's Blackfly S camera.

Run it from a terminal:
    python3 features_ransac.py            compare the two fits at 30 percent wrong matches
    python3 features_ransac.py 0.6        the same at 60 percent
"""
import json
import math
import sys

import numpy as np

W, H = 1616, 1240
CORNERS = np.array([[0, 0], [W, 0], [W, H], [0, H]], dtype=float)


def true_homography():
    """The same flat scene seen from two camera poses: a fixed, realistic H."""
    t = 0.12
    return np.array([[0.98 * math.cos(t), -0.12, 110.0],
                     [0.10, 0.93, 50.0],
                     [0.00008, -0.00003, 1.0]])


def apply_h(Hm, pts):
    """Map an (n, 2) array of points through a 3 x 3 homography."""
    p = np.column_stack([pts, np.ones(len(pts))]) @ Hm.T
    return p[:, :2] / p[:, 2:3]


def normalizer(pts):
    """Hartley normalization: move to the centroid and scale to an average distance of sqrt(2)."""
    c = pts.mean(axis=0)
    d = np.sqrt(((pts - c) ** 2).sum(axis=1)).mean()
    s = math.sqrt(2) / max(d, 1e-12)
    return np.array([[s, 0, -s * c[0]], [0, s, -s * c[1]], [0, 0, 1.0]])


def dlt(src, dst):
    """H with dst ~ H src, from 4 or more point pairs (the Lecture 12 DLT)."""
    Ts, Td = normalizer(src), normalizer(dst)
    s, d = apply_h(Ts, src), apply_h(Td, dst)
    rows = []
    for (x, y), (u, v) in zip(s, d):
        rows.append([-x, -y, -1, 0, 0, 0, u * x, u * y, u])
        rows.append([0, 0, 0, -x, -y, -1, v * x, v * y, v])
    _, _, vt = np.linalg.svd(np.array(rows))
    Hn = vt[-1].reshape(3, 3)
    Hm = np.linalg.inv(Td) @ Hn @ Ts
    return Hm / Hm[2, 2]


def make_matches(n=60, wrong=0.3, noise=1.0, seed=7):
    """n matches between two images: right ones follow H plus pixel noise, wrong ones land anywhere."""
    rng = np.random.default_rng(seed)
    Ht = true_homography()
    src = np.column_stack([120 + rng.random(n) * 1380, 120 + rng.random(n) * 1000])
    dst = apply_h(Ht, src) + rng.normal(0, noise, (n, 2))
    bad = rng.random(n) < wrong
    dst[bad] = np.column_stack([100 + rng.random(bad.sum()) * 1416, 100 + rng.random(bad.sum()) * 1040])
    return src, dst, bad, Ht


def tries_needed(e, s, p=0.99):
    """RANSAC tries for at least one sample with no wrong points, with probability p."""
    clean = (1 - e) ** s
    if clean <= 0:
        return math.inf
    if clean >= 1:
        return 1
    return math.ceil(math.log(1 - p) / math.log(1 - clean))


def residuals(Hm, src, dst):
    return np.sqrt(((apply_h(Hm, src) - dst) ** 2).sum(axis=1))


def ransac_homography(src, dst, thresh=3.0, p=0.99, max_tries=2000, seed=0):
    """RANSAC with 4-point samples. Updates the number of tries as it finds more inliers."""
    rng = np.random.default_rng(seed)
    n = len(src)
    best = np.zeros(n, dtype=bool)
    tries, need = 0, max_tries
    while tries < min(need, max_tries):
        tries += 1
        idx = rng.choice(n, 4, replace=False)
        try:
            Hm = dlt(src[idx], dst[idx])
        except np.linalg.LinAlgError:
            continue
        if not np.all(np.isfinite(Hm)):
            continue
        inl = residuals(Hm, src, dst) < thresh
        if inl.sum() > best.sum():
            best = inl
            need = tries_needed(1 - best.mean(), 4, p)
    if best.sum() < 4:
        return None, best, tries
    Hm = dlt(src[best], dst[best])                       # refit on all inliers,
    again = residuals(Hm, src, dst) < thresh             # then count once more
    if again.sum() >= best.sum():
        best = again
        Hm = dlt(src[best], dst[best])
    return Hm, best, tries


def corner_error(Hm, Ht):
    """Average distance, in pixels, between where Hm and the true H put the four image corners."""
    return float(np.sqrt(((apply_h(Hm, CORNERS) - apply_h(Ht, CORNERS)) ** 2).sum(axis=1)).mean())


def demo_api(params_json):
    """Web page entry point. Input: {"wrong": 0..0.9, "noise": px, "seed": int, "n": int}."""
    p = json.loads(params_json or "{}")
    wrong, noise = float(p.get("wrong", 0.3)), float(p.get("noise", 1.0))
    seed, n = int(p.get("seed", 7)), int(p.get("n", 60))
    src, dst, bad, Ht = make_matches(n, wrong, noise, seed)
    H_ls = dlt(src, dst)
    H_r, inl, tries = ransac_homography(src, dst, thresh=3.0 + 2.0 * noise, seed=seed * 31 + 5)
    out = {
        "W": W, "H": H, "src": src.round(2).tolist(), "dst": dst.round(2).tolist(), "wrong": bad.tolist(),
        "n_wrong": int(bad.sum()),
        "corners_true": apply_h(Ht, CORNERS).round(1).tolist(),
        "corners_ls": apply_h(H_ls, CORNERS).round(1).tolist(),
        "err_ls": corner_error(H_ls, Ht),
        "tries_formula": tries_needed(bad.mean(), 4),
        "ransac": None,
        "table": [{"s": s, "name": nm, "tries": [tries_needed(e, s) for e in (0.1, 0.3, 0.5, 0.7)]}
                  for s, nm in ((2, "line"), (4, "homography"), (8, "fundamental, 8 points"))],
    }
    if H_r is not None:
        out["ransac"] = {"corners": apply_h(H_r, CORNERS).round(1).tolist(), "err": corner_error(H_r, Ht),
                         "tries": tries, "inliers": inl.tolist(), "n_inliers": int(inl.sum()),
                         "missed_good": int((~inl & ~bad).sum()), "kept_bad": int((inl & bad).sum())}
    if out["tries_formula"] == math.inf:
        out["tries_formula"] = None
    return json.dumps(out)


def main(argv):
    wrong = float(argv[1]) if len(argv) > 1 else 0.3
    r = json.loads(demo_api(json.dumps({"wrong": wrong})))
    print("{} of 60 matches are wrong".format(r["n_wrong"]))
    print("least squares: corners off by {:.1f} px".format(r["err_ls"]))
    if r["ransac"]:
        q = r["ransac"]
        print("RANSAC: {} inliers after {} tries, corners off by {:.1f} px".format(q["n_inliers"], q["tries"], q["err"]))
    print("tries the formula asks for:", r["tries_formula"])


if __name__ == "__main__":
    main(sys.argv)