"""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

# 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

DEFAULT_CYCLE_COUNT = 200


def gain_lsb_per_ut(cycle_count):
    """Sensitivity in LSB per microtesla for a given cycle count.

    Linear fit to manual Table 3-1, which quotes 20, 38 and 75 LSB/uT at cycle
    counts of 50, 100 and 200; this reproduces all three to within a count.
    """
    return 0.3671 * cycle_count + 1.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.

        The pointer write is a separate transaction terminated by STOP rather
        than a repeated START, which is exactly what the manual's I2C read
        diagrams (sections 4.5.2 and 5.8.4) specify.
        """
        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]

    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."""
        self.write_reg(REG_HSHAKE, [HSHAKE_DRDY_ON_READ_ONLY])

    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):
        return bool(self.read_reg(REG_STATUS)[0] & STATUS_DRDY)

    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) as signed counts -- 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.
        """
        return decode_measurements(self.read_reg(REG_MX, 9))

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