#!/usr/bin/env python3
# test_power_with_init.py

from maix import i2c, time
import sys

# INA226 register addresses
INA226_ADDR = 0x40
REG_CONFIGURATION = 0x00
REG_BUS_VOLTAGE = 0x02
REG_CURRENT = 0x04
REG_CALIBRATION = 0x05

# Configuration values
CONFIG_VALUE = 0x4527      # Configuration: 16 averages, 1.1ms conversion time, continuous mode
CALIBRATION_VALUE = 0x1400  # Calibration value

def write_register(bus, reg, value):
    """Write to INA226 register"""
    data = [(value >> 8) & 0xFF, value & 0xFF]
    bus.writeto_mem(INA226_ADDR, reg, bytes(data))

def read_register(bus, reg):
    """Read from INA226 register"""
    data = bus.readfrom_mem(INA226_ADDR, reg, 2)
    return (data[0] << 8) | data[1]

def init_ina226(bus):
    """Initialize INA226 chip"""
    try:
        # Write configuration register
        write_register(bus, REG_CONFIGURATION, CONFIG_VALUE)
        time.sleep_ms(10)
        
        # Write calibration register
        write_register(bus, REG_CALIBRATION, CALIBRATION_VALUE)
        time.sleep_ms(10)
        
        # Verify configuration by reading it back
        config_read = read_register(bus, REG_CONFIGURATION)
        if config_read != CONFIG_VALUE:
            print(f"  Warning: Config readback mismatch: 0x{config_read:04X} != 0x{CONFIG_VALUE:04X}")
        
        return True
    except Exception as e:
        print(f"  Init failed: {e}")
        return False

def read_voltage(bus):
    """Read bus voltage"""
    raw = read_register(bus, REG_BUS_VOLTAGE)
    voltage = raw * 1.25 / 1000
    return voltage

def read_current(bus):
    """Read current"""
    raw = read_register(bus, REG_CURRENT)
    # Handle signed value
    if raw & 0x8000:
        raw = raw - 0x10000
    current_lsb = 0.001 * CALIBRATION_VALUE / 4096
    current = raw * current_lsb * 1000  # mA
    return current

def test_i2c_bus(bus_num):
    """Test a single I2C bus with full initialization"""
    print(f"\n{'='*60}")
    print(f"Testing I2C Bus {bus_num}")
    print(f"{'='*60}")
    
    try:
        # Step 1: Initialize I2C bus
        print(f"  1. Initializing I2C bus...")
        bus = i2c.I2C(bus_num, i2c.Mode.MASTER)
        print(f"     OK")
        
        # Step 2: Initialize INA226
        print(f"  2. Initializing INA226...")
        if not init_ina226(bus):
            print(f"     FAILED")
            return False
        print(f"     OK")
        
        # Step 3: Read voltage multiple times
        print(f"  3. Reading voltage...")
        for i in range(5):
            try:
                voltage = read_voltage(bus)
                current = read_current(bus)
                print(f"     Read {i+1}: {voltage:.3f}V, {current:.1f}mA")
                time.sleep_ms(100)
            except Exception as e:
                print(f"     Read {i+1} failed: {e}")
        
        print(f"  SUCCESS")
        return True
        
    except Exception as e:
        print(f"  FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """Test all I2C buses"""
    print("INA226 Test with Proper Initialization")
    print("=" * 60)
    
    # Test buses in order of likelihood
    test_order = [5, 1, 3, 4, 0, 2]
    
    success_buses = []
    
    for bus_num in test_order:
        if test_i2c_bus(bus_num):
            success_buses.append(bus_num)
            # If we found a working bus, stop testing others
            break
    
    print(f"\n{'='*60}")
    print(f"Summary:")
    print(f"  Working buses: {success_buses}")
    if not success_buses:
        print(f"  ERROR: No working I2C bus found!")
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())