#!/usr/bin/env python3
"""rssi_fit.py - calibrate a Bluetooth RSSI-to-distance model from your own readings.

Fits the log-distance path-loss model

    RSSI = TxPower - 10 * n * log10(distance_m)

by ordinary least squares, and prints TxPower (the RSSI you would read at 1 m),
n (the path-loss exponent) and the RMS residual in dB.

Input: one reading per line, "distance_m rssi_dbm", e.g.

    1 -58
    1 -61
    2 -66
    # comments and blank lines are ignored

Usage:
    python3 rssi_fit.py readings.txt
    python3 rssi_fit.py < readings.txt
    python3 rssi_fit.py --selftest      # fit seeded synthetic data with a known answer

Standard library only. Python 3.6+. From https://dankdev.com/Blog/bluetooth-rssi-to-distance
"""
import math
import random
import statistics
import sys


def fit(readings):
    """Least-squares fit of RSSI = TxPower - 10*n*log10(d).

    readings: list of (distance_m, rssi_dbm). Needs at least two distinct distances.
    Returns (tx_power, n, rms_residual_db).
    """
    if len(readings) < 3:
        raise ValueError("need at least 3 readings")
    xs = [math.log10(d) for d, _ in readings]
    ys = [float(r) for _, r in readings]
    mx = sum(xs) / len(xs)
    my = sum(ys) / len(ys)
    sxx = sum((x - mx) ** 2 for x in xs)
    if sxx == 0:
        raise ValueError("all readings are at one distance; measure at two or more")
    slope = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sxx
    tx_power = my - slope * mx
    n = -slope / 10.0
    rms = math.sqrt(sum((y - (tx_power + slope * x)) ** 2 for x, y in zip(xs, ys)) / len(xs))
    return tx_power, n, rms


def rssi_to_distance(rssi, tx_power, n):
    """Distance in metres implied by an RSSI under the fitted model."""
    return 10 ** ((tx_power - rssi) / (10.0 * n))


def parse(lines):
    readings = []
    for lineno, line in enumerate(lines, 1):
        line = line.split("#", 1)[0].strip()
        if not line:
            continue
        parts = line.replace(",", " ").split()
        if len(parts) != 2:
            raise ValueError("line %d: expected 'distance_m rssi_dbm', got %r" % (lineno, line))
        d, r = float(parts[0]), float(parts[1])
        if d <= 0:
            raise ValueError("line %d: distance must be > 0" % lineno)
        readings.append((d, r))
    return readings


def report(readings):
    tx_power, n, rms = fit(readings)
    dists = sorted(set(d for d, _ in readings))
    print("readings: %d at %d distances (%s m)"
          % (len(readings), len(dists), ", ".join("%g" % d for d in dists)))
    print("TxPower = %.1f dBm   (fitted RSSI at 1 m)" % tx_power)
    print("n       = %.2f       (path-loss exponent)" % n)
    print("rms residual = %.1f dB" % rms)
    for d in dists:
        rs = [r for dd, r in readings if dd == d]
        print("  at %g m: median RSSI %g dBm -> model says %.2f m"
              % (d, statistics.median(rs), rssi_to_distance(statistics.median(rs), tx_power, n)))
    return tx_power, n, rms


def selftest():
    true_tx, true_n, noise_sd, seed = -62.0, 2.7, 3.0, 20250219
    rng = random.Random(seed)
    readings = []
    for d in (1, 2, 4, 8):
        for _ in range(10):
            clean = true_tx - 10 * true_n * math.log10(d)
            readings.append((d, round(clean + rng.gauss(0, noise_sd))))
    print("selftest: true TxPower = %.1f dBm, true n = %.2f, noise sd = %.1f dB, seed %d"
          % (true_tx, true_n, noise_sd, seed))
    print("          10 whole-dBm readings at each of 1, 2, 4 and 8 m")
    tx_power, n, rms = report(readings)
    ok = abs(tx_power - true_tx) <= 1.5 and abs(n - true_n) <= 0.3
    print("selftest: %s (|TxPower error| = %.1f dB, |n error| = %.2f)"
          % ("PASS" if ok else "FAIL", abs(tx_power - true_tx), abs(n - true_n)))
    return 0 if ok else 1


def main(argv):
    if "--selftest" in argv:
        return selftest()
    if "-h" in argv or "--help" in argv:
        print(__doc__)
        return 0
    paths = [a for a in argv if not a.startswith("-")]
    if paths:
        with open(paths[0]) as fh:
            readings = parse(fh)
    else:
        readings = parse(sys.stdin)
    report(readings)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except ValueError as exc:
        print("rssi_fit: %s" % exc, file=sys.stderr)
        sys.exit(2)