#!/usr/bin/env python3
import smbus2

VOLTAGE_REG_ADDR = 0x02
POWER_REG_ADDR = 0x03
CURRENT_REG_ADDR = 0x04
CAL_REG_ADDR = 0x05

VOLTAGE_LSB = 0.00125
SHUNT_LSB = 0.0000025

class PowerData:
    voltage = 0.0
    power = 0.0
    current = 0.0

class INA226:
    # I2C Address, Shunt Resistor, Current LSB, Calibration Multiplier
    def __init__(self, bus: smbus2.SMBus, addr: int, rshunt: int, clsb: float, cal: float = 1):
        self.addr = addr
        self.rshunt = rshunt,
        self.clsb = clsb,
        self.cal = cal
        self.bus = bus
        self.data = PowerData()

        # Initial calibration
        self.calibrate()

    def __readRegWord(self, reg: int) -> int:
        regWord = self.bus.read_word_data(self.addr, reg) & 0xFFFF
        regValue = ((regWord << 8) & 0xFF00) + (regWord >> 8)
        return regValue
    
    def __writeRegWord(self, reg: int, val: int) -> None:
        regBytes = [(val >> 8) & 0xFF,val & 0xFF]
        self.bus.write_i2c_block_data(self.addr, reg, regBytes)

    # Calibrate current sensor
    def calibrate(self) -> None:
        calReg = int((0.00512 / (self.clsb * self.rshunt)) * self.cal)
        self.__writeRegWord(CAL_REG_ADDR, calReg)

    # Retrieve power data from sensor
    def getData(self) -> PowerData:
        self.data.voltage = float(self.__readRegWord(VOLTAGE_REG_ADDR)) * VOLTAGE_LSB
        self.data.power = float(self.__readRegWord(POWER_REG_ADDR)) * self.clsb * 25.0
        self.data.current = float(self.__readRegWord(CURRENT_REG_ADDR)) * self.clsb
        if (self.data.current > 20.96):
            self.data.current = 0.0
        return self.data