# Kookabus-01 driver and demo for MicroPython (Raspberry Pi Pico, ESP32, any port with machine.I2C).
# Copy to the board as main.py, or import Kookabus from it.
#
# The board's INA226 runs from the host's 3.3 V, so it resets whenever the host does:
# begin() must run after every host power-up (the constructor does it).
from machine import I2C, Pin
import time

R_SHUNT = 0.005     # ohm; trim this against a multimeter if the current reads off
I_LSB = 0.0005      # A per CURRENT bit
CAL = round(0.00512 / (I_LSB * R_SHUNT))   # 2048
CONFIG = 0x4527     # 16-sample average, 1.1 ms conversions, continuous

CELLS = 5           # 5S tool packs
R_PATH = 0.03 + 0.10  # ohm: board (fuse, FET, shunt, copper) + an assumed pack resistance
V_CUTOFF = 3.0 * CELLS

# Resting Li-ion cell voltage -> % (typical NMC; tune for your packs)
OCV = ((3.00, 0), (3.30, 5), (3.45, 10), (3.60, 20), (3.70, 35),
       (3.80, 50), (3.90, 65), (4.00, 78), (4.10, 90), (4.20, 100))


def soc(pack_v, cells=CELLS):
    v = pack_v / cells
    if v <= OCV[0][0]:
        return 0
    for (v0, p0), (v1, p1) in zip(OCV, OCV[1:]):
        if v <= v1:
            return p0 + (p1 - p0) * (v - v0) / (v1 - v0)
    return 100


class Kookabus:
    def __init__(self, i2c, addr=0x40):
        self.i2c, self.addr = i2c, addr
        self.begin()

    def _w(self, reg, val):
        self.i2c.writeto_mem(self.addr, reg, bytes((val >> 8, val & 0xFF)))

    def _r(self, reg, signed=False):
        b = self.i2c.readfrom_mem(self.addr, reg, 2)
        v = (b[0] << 8) | b[1]
        return v - 0x10000 if signed and v & 0x8000 else v

    def begin(self):
        if self._r(0xFE) != 0x5449:
            raise OSError("no INA226 at 0x%02x" % self.addr)
        self._w(0x00, CONFIG)
        self._w(0x05, CAL)

    def voltage(self):      # V at the LOAD connector
        return self._r(0x02) * 0.00125

    def current(self):      # A into the load; negative only if something back-feeds
        return self._r(0x04, True) * I_LSB

    def power(self):        # W
        return self._r(0x03) * 25 * I_LSB

    def pack_voltage(self):  # V, estimated resting voltage with the sag added back
        return self.voltage() + self.current() * R_PATH


if __name__ == "__main__":
    # PiicoDev / Qwiic on a Pico: I2C0 on GP8 (SDA) / GP9 (SCL). ESP32: I2C(0, sda=Pin(21), scl=Pin(22)).
    bus = Kookabus(I2C(0, sda=Pin(8), scl=Pin(9), freq=400_000))
    while True:
        v, i, p = bus.voltage(), bus.current(), bus.power()
        pv = bus.pack_voltage()
        print("%6.2f V  %6.3f A  %6.1f W  ~%3.0f %%%s" %
              (v, i, p, soc(pv), "  LOW - stop the load" if v < V_CUTOFF else ""))
        time.sleep(1)