#!/usr/bin/env python3
import smbus2

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

CURRENT_LSB = 0.00032
VOLTAGE_LSB = 0.00125
SHUNT_LSB = 0.0000025

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

class INA226:
    def __init__(self, bus: smbus2.SMBus, addr: int, cal: int = 0):
        self.addr = addr
        self.cal = cal
        self.bus = bus
        self.data = PowerData()

        # Initial calibration
        self.calibrate(cal)

    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, cal: int) -> None:
        self.__writeRegWord(CAL_REG_ADDR, cal)

    # 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)) * CURRENT_LSB * 25.0
        self.data.current = float(self.__readRegWord(CURRENT_REG_ADDR)) * CURRENT_LSB
        if (self.data.current > 20.96):
            self.data.current = 0.0
        return self.data