"""
rtk_modes.py: standalone, RTK float, and RTK fixed from the same walk.

EECE 5554 Robotics Sensing and Navigation, demo for Lab 2.

This script simulates what a receiver like the Lab 2 RTK mosaic-X5 reports in
its three modes. The numbers are made up but realistic in size:

    standalone   meter level      no corrections
    RTK float    decimeter level  corrections, carrier ambiguities not yet integers
    RTK fixed    centimeter level corrections, ambiguities fixed to integers

GNSS errors are not independent from one epoch to the next. Each mode gets a
slowly wandering error (a first-order Gauss-Markov process) plus a bias, which
is what multipath and atmosphere do. That is why a cloud of points can be
tight (precise) and still sit in the wrong place (inaccurate).

    precision = how tightly the points cluster around their own center
    accuracy  = how far that center is from the true position

Run it from a terminal:
    python3 rtk_modes.py            print the statistics for every case
"""
import json
import math
import sys

import numpy as np

# per-axis 1-sigma wander (m), its correlation time (s), and bias size (m), for each mode and scene
MODES = {
    "open": {
        "standalone": {"sigma": 0.80, "tau": 40.0, "bias": 0.70, "avail": 1.00},
        "float":      {"sigma": 0.15, "tau": 30.0, "bias": 0.12, "avail": 1.00},
        "fixed":      {"sigma": 0.008, "tau": 10.0, "bias": 0.006, "avail": 1.00},
    },
    "occluded": {   # trees and buildings: more multipath, and the fix comes and goes
        "standalone": {"sigma": 2.20, "tau": 60.0, "bias": 2.50, "avail": 1.00},
        "float":      {"sigma": 0.60, "tau": 45.0, "bias": 0.45, "avail": 0.85},
        "fixed":      {"sigma": 0.015, "tau": 10.0, "bias": 0.012, "avail": 0.35},
    },
}
LABELS = {"standalone": "Standalone", "float": "RTK float", "fixed": "RTK fixed"}


def gauss_markov(n, sigma, tau, dt, rng):
    """First-order Gauss-Markov noise: x[k] = a x[k-1] + w, with steady-state std sigma."""
    a = math.exp(-dt / tau)
    w = rng.normal(0.0, sigma * math.sqrt(1 - a * a), size=n)
    x = np.empty(n)
    x[0] = rng.normal(0.0, sigma)
    for k in range(1, n):
        x[k] = a * x[k - 1] + w[k]
    return x


def availability_mask(n, avail, rng, seg=20):
    """Which epochs have this solution: on and off in runs of about `seg` seconds."""
    if avail >= 1.0:
        return np.ones(n, dtype=bool)
    m = np.zeros(n, dtype=bool)
    k = 0
    on = rng.random() < avail
    while k < n:
        mean = seg * (avail / (1 - avail) if on else 1.0)
        run = max(3, int(rng.exponential(mean)))
        m[k:k + run] = on
        k += run
        on = not on
    return m


def truth_path(motion, n, dt=1.0):
    """Stand still at the origin, or walk a 120 m by 75 m rectangle at about 1.3 m/s."""
    if motion == "still":
        return np.zeros((n, 2))
    w, h = 120.0, 75.0
    per = 2 * (w + h)
    s = (np.arange(n) * dt * 1.3) % per
    e = np.where(s < w, s, np.where(s < w + h, w, np.where(s < 2 * w + h, w - (s - w - h), 0.0)))
    nn = np.where(s < w, 0.0, np.where(s < w + h, s - w, np.where(s < 2 * w + h, h, h - (s - 2 * w - h))))
    return np.column_stack([e, nn])


def simulate(scene="open", motion="still", seconds=300, seed=5):
    rng = np.random.default_rng(seed)
    n = int(seconds)
    truth = truth_path(motion, n)
    out = {}
    for mode, prm in MODES[scene].items():
        err = np.column_stack([gauss_markov(n, prm["sigma"], prm["tau"], 1.0, rng) for _ in range(2)])
        ang = rng.uniform(0, 2 * math.pi)
        bias = prm["bias"] * np.array([math.cos(ang), math.sin(ang)])
        mask = availability_mask(n, prm["avail"], rng)
        pos = truth + err + bias
        out[mode] = {"pos": pos[mask], "truth": truth[mask], "mask": mask}
    return truth, out


def stats(pos, truth):
    """Precision, accuracy, and total error, all in meters."""
    if len(pos) == 0:
        return None
    err = pos - truth                       # error at each epoch
    center = err.mean(axis=0)               # average error = offset of the cloud
    spread = err - center
    return {
        "n": int(len(pos)),
        "precision_m": float(math.sqrt((spread ** 2).sum(axis=1).mean())),   # RMS distance from own center
        "accuracy_m": float(math.hypot(*center)),                             # center to truth
        "rms_error_m": float(math.sqrt((err ** 2).sum(axis=1).mean())),       # total, vs truth
        "center_offset": [float(center[0]), float(center[1])],
    }


def demo_api(params_json):
    """Web page entry point. Input: {"scene": "open"|"occluded", "motion": "still"|"walk", "seed": int}."""
    p = json.loads(params_json or "{}")
    scene = p.get("scene", "open")
    motion = p.get("motion", "still")
    truth, sims = simulate(scene, motion, int(p.get("seconds", 300)), int(p.get("seed", 5)))
    res = {"scene": scene, "motion": motion, "truth": truth.round(3).tolist(), "modes": {}}
    for mode, d in sims.items():
        s = stats(d["pos"], d["truth"])
        step = max(1, len(d["pos"]) // 400)
        res["modes"][mode] = {
            "label": LABELS[mode],
            "availability": float(d["mask"].mean()),
            "stats": s,
            "points": d["pos"][::step].round(4).tolist(),
        }
    return json.dumps(res)


def main(argv):
    for scene in ("open", "occluded"):
        for motion in ("still", "walk"):
            _, sims = simulate(scene, motion)
            print("{} sky, {}".format(scene, motion))
            for mode, d in sims.items():
                s = stats(d["pos"], d["truth"])
                print("  {:11s} avail {:4.0%}  precision {:7.3f} m  accuracy {:7.3f} m  RMS error {:7.3f} m".format(
                    LABELS[mode], d["mask"].mean(), s["precision_m"], s["accuracy_m"], s["rms_error_m"]))


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