from machine import I2C, Pin
import time

class INA3221:
    INA3221_ADDR = 0x40
    INA3221_CONFIG = 0x00
    INA3221_SHUNT_VOLTAGE_BASE = 0x01
    INA3221_BUS_VOLTAGE_BASE = 0x02

    def __init__(self, i2c):
        self.i2c = i2c
        self.i2c.writeto_mem(self.INA3221_ADDR, self.INA3221_CONFIG, bytes([0x71, 0x5C]))

    def read_shunt_voltage(self, channel):
        addr = self.INA3221_SHUNT_VOLTAGE_BASE + (channel * 2)
        raw_data = self.i2c.readfrom_mem(self.INA3221_ADDR, addr, 2)
        shunt_voltage = int.from_bytes(raw_data, 'big')
        if shunt_voltage & (1 << 15):  # check if the value is negative
            shunt_voltage = shunt_voltage - (1 << 16)
        return shunt_voltage * 0.005  # LSB: 5 uV

    def read_bus_voltage(self, channel):
        addr = self.INA3221_BUS_VOLTAGE_BASE + (channel * 2)
        raw_data = self.i2c.readfrom_mem(self.INA3221_ADDR, addr, 2)
        return (int.from_bytes(raw_data, 'big') >> 3) * 0.001  # LSB: 1 mV

i2c = I2C(scl=Pin(22), sda=Pin(21), freq=400000)
ina3221 = INA3221(i2c)

SHUNT_RESISTOR_VALUE = 0.1   # Shunt resistor value in OHMs

# Name the 3 channels that are measuring current.
pan_motor_channel = 1
tilt_motor_channel = 2
solar_panel_channel = 3


while True:
    pan_motor_current = ina3221.read_shunt_voltage(pan_motor_channel) / SHUNT_RESISTOR_VALUE
    pan_bus_voltage = ina3221.read_bus_voltage(pan_motor_channel)
    tilt_motor_current = ina3221.read_shunt_voltage(tilt_motor_channel) / SHUNT_RESISTOR_VALUE
    tilt_bus_voltage = ina3221.read_bus_voltage(tilt_motor_channel)
    solar_panel_current = ina3221.read_shunt_voltage(solar_panel_channel) / SHUNT_RESISTOR_VALUE
    solar_panel_bus_voltage = ina3221.read_bus_voltage(solar_panel_channel)
    
    print("  Pan Motor:  Bus Voltage = ", pan_bus_voltage, "V", "  Current = ", pan_motor_current, "A")
    print(" Tilt Motor:  Bus Voltage = ", tilt_bus_voltage, "V", "  Current = ", tilt_motor_current, "A")
    print("Solar Panel:  Bus Voltage = ", solar_panel_bus_voltage, "V", "  Current = ", solar_panel_current, "A")
    print("----------------------------------"
    
    time.sleep(1)