"""PNI RM3100 geomagnetic sensor driver.

Register numbers and sequences follow the RM3100 & RM2100 Sensor Suite User
Manual (Doc 1017252 R07), section 5.

Knows nothing about USB: the bus object need only provide write(addr, data)
and read(addr, count).
"""

import time
from collections import namedtuple

# Register addresses (manual Table 5-1).
REG_POLL = 0x00  # single measurement trigger
REG_CMM = 0x01  # continuous measurement mode
REG_CCX = 0x04  # cycle counts, 6 bytes: CCX, CCY, CCZ as uint16 big-endian
REG_TMRC = 0x0B  # continuous mode update rate
REG_MX = 0x24  # measurements, 9 bytes: X, Y, Z as int24 big-endian
REG_BIST = 0x33
REG_STATUS = 0x34  # bit 7 = DRDY
REG_HSHAKE = 0x35
REG_REVID = 0x36

# The manual's own examples (sections 5.7.2 and 5.8.3) use 0x79 for "measure all
# three axes, raise DRDY once the whole sequence is done", even though Table 5-1
# describes bit 3 as reserved-zero. Follow the examples.
CMM_ALL_AXES = 0x79
CMM_OFF = 0x00

POLL_ALL_AXES = 0x70

# HSHAKE with DRC1=1, DRC0=0: DRDY is cleared by reading the measurement
# registers, but *not* by an arbitrary register write. The 0x1B default has
# DRC0=1, which would mean the pointer write needed to read STATUS clears the
# very flag we are about to sample, so polling could never observe it set.
HSHAKE_DRDY_ON_READ_ONLY = 0x0A

# TMRC values (manual Table 5-4), mapped to their approximate rates in Hz.
TMRC_RATES = {
    0x92: 600.0, 0x93: 300.0, 0x94: 150.0, 0x95: 75.0,
    0x96: 37.0, 0x97: 18.0, 0x98: 9.0, 0x99: 4.5,
    0x9A: 2.3, 0x9B: 1.2, 0x9C: 0.6, 0x9D: 0.3,
}

STATUS_DRDY = 0x80

EXPECTED_REVID = 0x22

# 50 buys bandwidth almost for free. Against cycle count 200 it costs 4.4% in
# post-filter noise (89% duty against 97%) and nothing else, while giving 3.7x
# the spectrum -- 267 Hz of Nyquist against 73 Hz. That matters because aliased
# interference cannot be filtered out afterwards at any cycle count, so seeing
# it is worth more than a few percent of noise.
#
# The one thing to verify rather than assume is dither: 50 sits at 0.58 LSB of
# intrinsic noise, which simulation puts safely in the region where averaging
# still recovers sub-LSB resolution (it fails below ~0.2). characterize.py
# prints sd/LSB, which answers it from the first capture. Fall back to 100 or
# 200 if that comes back low.
DEFAULT_CYCLE_COUNT = 50

# The fastest rate register. Used by default so the cycle count, not TMRC, sets
# the rate -- TMRC offers only factor-of-two steps and has no effect at all once
# the cycle count governs.
TMRC_FASTEST = 0x92


# Table 3-1 quotes gain in LSB per microtesla; these are the coefficients of a
# linear fit to its 20, 38 and 75 LSB/uT at cycle counts 50, 100 and 200, which
# reproduces all three to within a count. Working units in this codebase are
# tesla, so the fit is scaled by 1e6 on the way out.
_GAIN_SLOPE_LSB_PER_UT = 0.3671
_GAIN_OFFSET_LSB_PER_UT = 1.5
UT_PER_TESLA = 1e6

# Displays and plots use nanotesla; no computation is done in it.
NT_PER_TESLA = 1e9


def gain_lsb_per_tesla(cycle_count):
    """Sensitivity in LSB per tesla for a given cycle count."""
    return (_GAIN_SLOPE_LSB_PER_UT * cycle_count
            + _GAIN_OFFSET_LSB_PER_UT) * UT_PER_TESLA


def tesla_per_count(cycle_count):
    """Calibration constant: multiply a raw count by this to get tesla.

    Stored in capture headers, so a reader needs no knowledge of the gain fit --
    and it is a multiply at the point of use rather than a divide.
    """
    return 1.0 / gain_lsb_per_tesla(cycle_count)


# Timing model for a three-axis measurement:
#
#     per-axis time = cycle_count / COUNTS_PER_SECOND + AXIS_OVERHEAD_S
#
# COUNTS_PER_SECOND comes from the specification, not from a fit: Table 3-1
# gives a 180 kHz circuit oscillation frequency, and section 4.1 measures each
# cycle count in *both* bias directions, so one count costs two oscillations.
# Rates measured here across cycle counts 229..29769 agree with that figure to
# within 0.8%, which confirms the spec value rather than improving on it -- the
# residual is this particular unit's oscillator sitting inside ordinary
# component tolerance, and another part would sit somewhere else.
#
# The overhead has no specified value and must be measured; 68.7 us reproduces
# the observed rates to better than 1% when the divisor is held at spec. It is
# why the naive "rate x cycle_count" constant is not constant, drifting from
# 84,429 at cycle count 100 to 89,191 at 1200.
COUNTS_PER_SECOND = 90000.0        # 180 kHz (Table 3-1) / 2 bias directions
AXIS_OVERHEAD_S = 68.7e-6
AXES = 3

# Rates predicted from the model are good to roughly this much on a given unit,
# and no better across units, since the manual quotes no tolerance on the
# oscillator. Anything needing the real number measures it: logger.py calibrates
# the period against the host clock before recording.
RATE_TOLERANCE = 0.02

# Table 3-1 quotes 30/20/15 nT at cycle counts 50/100/200, which fits
# K/sqrt(cycle_count). Extrapolation past ~400 is unverified: the manual calls
# that its useful upper limit for noise and gives no data beyond it.
_NOISE_K_NT = 208.0


def sample_period(cycle_count, axes=AXES):
    """Seconds between measurements when the cycle count governs the rate."""
    return axes * (cycle_count / COUNTS_PER_SECOND + AXIS_OVERHEAD_S)


def integration_time(cycle_count, axes=AXES):
    """Seconds per sample actually spent integrating, excluding fixed overhead.

    Only this part reduces noise. Against the sample period it gives the duty
    cycle: how much of the wall clock the sensor is doing useful work rather
    than idling between measurements or paying per-axis overhead.
    """
    return axes * cycle_count / COUNTS_PER_SECOND


def cycle_count_for_rate(rate_hz, axes=AXES):
    """Cycle count that makes the sensor free-run at about rate_hz.

    Inverse of sample_period(). Use with a TMRC faster than the target so the
    cycle count governs -- then this is a continuous rate knob, where TMRC only
    offers factor-of-two steps, and the duty cycle is ~100% by construction.

    The achieved rate will sit within about RATE_TOLERANCE of the target, set by
    oscillator tolerance rather than by this calculation. Measure it if it
    matters.
    """
    cc = round((1.0 / (rate_hz * axes) - AXIS_OVERHEAD_S) * COUNTS_PER_SECOND)
    return max(1, min(0xFFFF, cc))


# Section 5.1: "quantization issues generally dictate working above a cycle
# count value of ~30". Below that the LSB grows faster than the sensor's own
# noise, so the quantiser stops being dithered and averaging stalls.
MIN_CYCLE_COUNT = 30
MAX_CYCLE_COUNT = 0xFFFF
MIN_RATE_BY_CYCLE_COUNT = 1.0 / (AXES * (MAX_CYCLE_COUNT / COUNTS_PER_SECOND
                                         + AXIS_OVERHEAD_S))

Plan = namedtuple("Plan", "cycle_count tmrc predicted_hz governed_by duty notes")


def plan(rate_hz=None, cycle_count=None, tmrc=None):
    """Resolve a full configuration from whichever knob the caller specified.

    Two ceilings compete and **the slower one wins** (manual section 5.2.1): the
    cycle count sets how long a measurement takes, TMRC sets how often one is
    started. Leaving TMRC faster than the cycle-count ceiling means the sensor
    free-runs at ~100% duty; setting it slower makes the sensor idle, which
    costs sensitivity for nothing unless low power is the goal.

    So the default in every branch is TMRC_FASTEST, and the cycle count is the
    rate knob -- it is continuous where TMRC offers only factor-of-two steps.
    """
    notes = []
    if cycle_count is None:
        if rate_hz is None:
            cycle_count = DEFAULT_CYCLE_COUNT
        elif rate_hz < MIN_RATE_BY_CYCLE_COUNT:
            # The register is 16 bits, so below ~0.46 Hz the cycle count runs
            # out of range and TMRC is the only way to go slower. Max the cycle
            # count anyway: it costs nothing and buys resolution.
            cycle_count = MAX_CYCLE_COUNT
            notes.append(
                f"{rate_hz:g} Hz is below the {MIN_RATE_BY_CYCLE_COUNT:.3f} Hz "
                f"floor of a {MAX_CYCLE_COUNT:,}-count measurement, so TMRC "
                "must set the cadence and the sensor will idle")
            if tmrc is None:
                tmrc = min(TMRC_RATES, key=lambda t: abs(TMRC_RATES[t] - rate_hz))
        else:
            cycle_count = cycle_count_for_rate(rate_hz)

    if cycle_count < MIN_CYCLE_COUNT:
        notes.append(f"cycle count raised to the {MIN_CYCLE_COUNT} the manual "
                     "advises as a quantisation floor (section 5.1)")
    cycle_count = max(MIN_CYCLE_COUNT, min(MAX_CYCLE_COUNT, cycle_count))
    ceiling = 1.0 / sample_period(cycle_count)
    if tmrc is None:
        tmrc = TMRC_FASTEST
    requested = TMRC_RATES[tmrc]

    if requested <= ceiling:
        predicted, governed_by = requested, "TMRC"
        idle = 1.0 - requested / ceiling
        if idle > 0.05:
            notes.append(f"TMRC leaves the sensor idle {idle:.0%} of each "
                         "period, which costs sensitivity")
    else:
        predicted, governed_by = ceiling, "cycle count"

    if rate_hz is not None and predicted < rate_hz * 0.98:
        notes.append(f"{rate_hz:g} Hz is faster than this configuration can "
                     f"reach; {predicted:.1f} Hz is the ceiling")
    return Plan(cycle_count, tmrc, predicted, governed_by,
                integration_time(cycle_count) * predicted, notes)


def expected_noise_nt(cycle_count):
    """Per-sample noise from the Table 3-1 fit, in nanotesla."""
    return _NOISE_K_NT / cycle_count ** 0.5


def decode_measurements(data):
    """Decode 9 bytes from REG_MX into (x, y, z) signed counts.

    Each axis is 24-bit two's complement, most significant byte first.
    """
    if len(data) != 9:
        raise ValueError(f"Expected 9 measurement bytes, got {len(data)}")
    return tuple(
        int.from_bytes(data[i:i + 3], "big", signed=True) for i in (0, 3, 6)
    )


class RM3100:
    """An RM3100 on an I2C bus."""

    # The top 5 bits of the address are fixed at 0b01000; SA1/SA0 are strapped
    # on the module, so any of these four is possible (manual section 4.5).
    ADDRESSES = range(0x20, 0x24)

    def __init__(self, bus, address):
        self.bus = bus
        self.address = address
        self.cycle_count = DEFAULT_CYCLE_COUNT

    def read_reg(self, reg, count=1):
        """Read count bytes starting at reg, using the sensor's auto-increment.

        Returns (data, mono, wall): the bus stamps the clocks when the reply
        lands, and the stamp travels with its own data.

        Prefers a single combined transaction (repeated START) where the bus
        offers one, since the host round trip is what limits the sample rate.
        Falls back to the manual's STOP-then-START form (sections 4.5.2, 5.8.4)
        for a bus that cannot do it.
        """
        combined = getattr(self.bus, "write_read", None)
        if combined is not None:
            return combined(self.address, [reg], count)
        self.bus.write(self.address, [reg])
        return self.bus.read(self.address, count)

    def write_reg(self, reg, data):
        self.bus.write(self.address, bytes([reg]) + bytes(data))

    def revid(self):
        return self.read_reg(REG_REVID)[0][0]

    def set_cycle_counts(self, count):
        """Set all three axes to the same cycle count."""
        if not 0 <= count <= 0xFFFF:
            raise ValueError(f"Cycle count {count} outside 0..65535")
        self.write_reg(REG_CCX, count.to_bytes(2, "big") * 3)
        self.cycle_count = count

    def get_cycle_counts(self):
        """Read back (ccx, ccy, ccz)."""
        data, _, _ = self.read_reg(REG_CCX, 6)
        return tuple(
            int.from_bytes(data[i:i + 2], "big") for i in (0, 2, 4)
        )

    def set_rate(self, tmrc):
        if tmrc not in TMRC_RATES:
            raise ValueError(
                f"TMRC 0x{tmrc:02x} not one of "
                f"{', '.join(f'0x{v:02x}' for v in TMRC_RATES)}"
            )
        self.write_reg(REG_TMRC, [tmrc])

    def configure(self):
        """Put DRDY into a state where polling STATUS actually works.

        Verified rather than assumed: exactly-once sampling relies on DRC1=1
        clearing DRDY when the measurement registers are read, which is what
        makes a second read of the same measurement impossible.
        """
        self.write_reg(REG_HSHAKE, [HSHAKE_DRDY_ON_READ_ONLY])
        # Bits 4-6 are read-only NACK status, so compare only the writable ones.
        readback = self.read_reg(REG_HSHAKE)[0][0] & 0x0F
        if readback != HSHAKE_DRDY_ON_READ_ONLY & 0x0F:
            raise IOError(
                f"HSHAKE did not take: wrote 0x{HSHAKE_DRDY_ON_READ_ONLY:02x}, "
                f"read back 0x{readback:02x}. Exactly-once sampling cannot be "
                "guaranteed without DRC1=1."
            )

    def start_cmm(self):
        self.write_reg(REG_CMM, [CMM_ALL_AXES])

    def stop_cmm(self):
        self.write_reg(REG_CMM, [CMM_OFF])

    def data_ready(self):
        """True if a measurement is waiting."""
        return self.poll_ready()[0]

    def poll_ready(self):
        """Like data_ready(), but returns (ready, mono, wall).

        The stamp is when the STATUS reply landed, which is the tightest bound
        available on when DRDY actually went high.
        """
        data, mono, wall = self.read_reg(REG_STATUS)
        return bool(data[0] & STATUS_DRDY), mono, wall

    def wait_for_data(self, timeout=2.0, interval=0.001):
        """Block until DRDY is set. Returns False if timeout elapses first."""
        deadline = time.monotonic() + timeout
        while True:
            if self.data_ready():
                return True
            if time.monotonic() >= deadline:
                return False
            time.sleep(interval)

    def read_raw(self):
        """Return ((x, y, z) counts, mono, wall) -- the fast path.

        Callers logging at high rates should use this and defer the microtesla
        conversion, so the sampling loop does I2C and nothing else.
        """
        data, mono, wall = self.read_reg(REG_MX, 9)
        return decode_measurements(data), mono, wall

    def read_measurements(self):
        """Return ((x, y, z) counts, (x, y, z) tesla)."""
        counts, _, _ = self.read_raw()
        gain = gain_lsb_per_tesla(self.cycle_count)
        return counts, tuple(c / gain for c in counts)