"""Toy pose-graph SLAM along a corridor: robot goes 20 m to the lab and back.
Compares: odometry only, naive (trust every loop closure), DCS, and a reference
that knows which loop closure is false. Everything is linear, so each solve is one
weighted least-squares problem and the whole script runs in about a second.
"""
import sys
import numpy as np

N = 41                                   # poses: 0 = start, 20 = lab door, 40 = back at start
SIG_ODO, SIG_LOOP = 0.02, 0.10           # std dev of odometry and loop-closure noise (m)
DRIFT = np.array([0.010, 0.004])         # odometry bias per step (m), does NOT cancel on the way back
TRUE = np.array([[min(i, 40 - i) * 1.0, 0.0] for i in range(N)])


def make_data(seed):
    rng = np.random.default_rng(seed)
    odom = [(i, i + 1, TRUE[i + 1] - TRUE[i] + DRIFT + rng.normal(0, SIG_ODO, 2)) for i in range(N - 1)]
    good = [(i, j, rng.normal(0, SIG_LOOP, 2)) for i, j in [(10, 30), (0, 40), (5, 35)]]  # true: same place
    bad = [(10, 26, np.zeros(2))]            # false: poses 4 m apart that look identical
    return odom, good, bad


def solve(odom, loops, scale=None):
    A, b, w = [], [], []
    def add(i, j, z, om):
        r = np.zeros(N); r[j], r[i] = 1, -1
        A.append(r); b.append(z); w.append(om)
    for i, j, z in odom: add(i, j, z, 1 / SIG_ODO**2)
    for k, (i, j, z) in enumerate(loops):
        s = 1.0 if scale is None else scale[k]
        add(i, j, z, s * s / SIG_LOOP**2)    # DCS scales the information by s^2
    r = np.zeros(N); r[0] = 1                # anchor the first pose
    A.append(r); b.append(TRUE[0]); w.append(1e6)
    A, b, W = np.array(A), np.array(b), np.diag(w)
    return np.linalg.solve(A.T @ W @ A, A.T @ W @ b)


def chi2(x, loops):
    return np.array([np.sum((x[j] - x[i] - z) ** 2) / SIG_LOOP**2 for i, j, z in loops])


def dcs(odom, loops, phi=10.0, iters=30):
    x, s = solve(odom, []), np.ones(len(loops))
    for _ in range(iters):
        c = chi2(x, loops)
        s = np.minimum(1.0, 2 * phi / (phi + c))   # the DCS scale
        x = solve(odom, loops, s)
    return x, s, chi2(x, loops)


rmse = lambda x: float(np.sqrt(np.mean(np.sum((x - TRUE) ** 2, 1))))
final = lambda x: float(np.linalg.norm(x[-1] - TRUE[-1]))

if __name__ == "__main__":
    seed = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    odom, good, bad = make_data(seed)
    loops = good + bad
    runs = {"odometry only": solve(odom, []), "naive (trust all loops)": solve(odom, loops),
            "DCS (phi=10)": dcs(odom, loops)[0], "reference (correct loops only)": solve(odom, good)}
    for k, x in runs.items():
        print(f"{k:34s} RMSE {rmse(x):.3f} m   error at return {final(x):.3f} m")
    _, s, c = dcs(odom, loops)
    print("loop-closure scales s:", np.round(s, 4), " chi2 after DCS:", np.round(c, 1))