chessbot

import smbus
from time import sleep, time
from i2cdev import I2C

import paho.mqtt.client as mqtt
import struct

bus_number = 1


COLUMNS = 8
ROWS = 8

DEBUG = False
PCA_OFFSET = 0  # 0 for normal use, 0-7 for debugging a single board that is not yet daisy chained
PCA_OG_ADDR = 0x70
PCA_ADDR = PCA_OG_ADDR + PCA_OFFSET
PCA_TO_PHYSICAL_WIRING_ORDER = [0,1,2,3,4,5,6,7]

current_power_state = [0x00] * 8 

#ban_addrs = banned_addresses_that_are_being_mean_to_me = []

# these addresses are tmaga1s (+- 40mT) whereas all the others are tmaga2s (+= 133mT)
tmag1s = [0x36]

#ban_adr = [((i - 0x20) // COLUMNS, PCA_TO_PHYSICAL_WIRING_ORDER.index((i-0x20) % COLUMNS)) for i in banned_addresses_that_are_being_mean_to_me]


mqttc = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)

def power_on(upto):
    upto += 1
    a=b''

    a += b'\xff' * (upto // 8)
    if len(a) < 8: a += bytes([0xff >> (8-(upto % 8))])
    if len(a) < 8: a += b'\x00' * ((8- upto // 8  - 1) )

    c = b''
    for i in range(64):
        c+= b'\x01\xff\x01' if i < upto else b'\x00\x00\x00'

    mqttc.publish("/magpower", a, retain=True)
    mqttc.publish("/led", c)




def power_on_single(addr):
   
    global current_power_state

    index = addr - 0x20      # 0–63 index
    row = index // 8
    col = index % 8

    # set that bit on
    current_power_state[row] |= (1 << col)

    # Publish updated power state
    a = bytes(current_power_state)
    mqttc.publish("/magpower", a, retain=True)

    #print(f"Turned ON TMAG at addr {hex(addr)} (row {row}, col {col})")






def power_off():
    a=b'\x00'*8
    mqttc.publish("/magpower", a, retain=True)


def power_off_single(addr):
    index = addr - 0x20          # map address to 0–63 index
    row = index // 8
    col = index % 8

    power_state = [0xFF] * 8     # start with all on
    power_state[row] &= ~(1 << col)

    a = bytes(power_state)
    mqttc.publish("/magpower", a, retain=True)

   # print(f"Turned off TMAG at addr {hex(addr)} (row {row}, col {col})")



#class PCAArray:
#    def __init__(self):
#        self.bus = smbus.SMBus(bus_number)
#    def softDisableAll(self):
#        for i in range(0, ROWS):
#            try: 
#                pigpio.i2c_write_byte(self.pcas[i], 0)
#                #if .read_byte(PCA_ADDR + i) != 0: print("ERR WRITING")
#                # TODO
#            except: pass
#     
#    # chip from 0 to 7
#    def enableOnly(self, chip, line):
#        self.softDisableAll()
#        # enable
#        self.bus.write_byte(PCA_ADDR + chip, 1 << line)
#    def unsafeEnableRow(self, chip):
#        on = 0xff
#        for chp, sq in ban_adr:
#            if chip == chp:
#                on &= ~(1 << sq)
#        #print(bin(on), chip)
#        self.bus.write_byte(PCA_ADDR + chip, on)
#    def unsafeDisableRow(self, chip):
#        self.bus.write_byte(PCA_ADDR + chip, 0x00)
#    def enableAll(self):
#        for i in range(0, ROWS):
#            try:
#                self.bus.write_byte(PCA_ADDR + i, enabled_cells[i])
#            except: pass
#switches = PCAArray()

tmagcache = {}
class TMAG5273:
    def __init__(self, addr):
        self.bus = smbus.SMBus(bus_number)
        self.device_address = addr
        if addr in tmag1s:
            self.RANGE = 80
        else:
            self.RANGE = 266

    @classmethod
    def init(self, new_addr, old_addr = 0x78):
        conf = self.config()
        conf[0xC] |= new_addr << 1
        smbus.SMBus(bus_number).write_i2c_block_data(old_addr, 0, conf)

    @classmethod
    def config(self):
        thing = [0 for _ in range(0, 0xD)]
        thing[0] |= 0 << 7; #NO TODO IT DOESN"T LIKE IT # Enable CRC
        # ignore magtempcompensation
        thing[0] |= 0x5 << 2; # 32x averaging = 400 samples / sec
        thing[0] |= 0x1; # 16 bit, 1-byte fast read command

        # cfg 2
        # ignore threshold stuff
        thing[1] |= 1 << 4; # low noise mode
        # leave glitch filter default
        # leave trigger mode default through i2c
        thing[1] |= 0x2; # continuous measure mode, TODO try wake up and sleep mode, 0x3

        # sensor cfg 1
        thing[2] |= 0x7 << 4; # enable xyz channels
        # leave w&s timing to the default of 1ms

        # sensor cfg 2
        # ignore threshold, angle calculation magic
        thing[3] |= 0x1 << 1; # XY axes range is the bigger one (40->80 or 133->266)
        thing[3] |= 0x1;      # z axis is bigger one

        # ignore threshold cfg registers
        # thing 4 to 6 = 0

        # enable temp acquisition cuz why not
        # no need for temp threshold
        thing[7] = 1; # 0th bit, enable data acquisition of temp
                      # interrupt pin
                      # most functions are disable by default
        thing[8] = 1; # 0th bit, Mask interrupt (it's connected to ground, so ignore the pin)

        # ignore mag gain/angle stuff registers
        # thing 9, A, B = 0

        # C 1..7 = I2C address will be put in by parent
        thing[0xC] = 1; # enable updating the address at the end of this transaction
        return thing


    def read_magnetism(self):
        # Command to start a measurement and read the data
        #self.bus.write_byte(self.device_address, 0x10)
        # Wait for the measurement to complete (adjust the delay according to the datasheet)
        #sleep(0.1)
        # Read 6 bytes of data (X, Y, Z magnetic field values)

        #if self.device_address in ban_addrs:
        #    return 10,10,-32
        if self.device_address not in tmagcache:
            tmagcache[self.device_address] = I2C(self.device_address, bus_number)
        i2c = tmagcache[self.device_address]
        #data = self.bus.read_i2c_block_data(self.device_address, 0x12, 6)
        data1 = list(bytearray(i2c.read(8)))         # read 1 byte
        data1 = data1[2:]
        #i2c.close()                 # close connection

        #print(data == data1, data, data1)
        #print(len(data))
        #print(data)
        data = data1
        x = twos_complement_int(data[1], data[0]) * self.RANGE  / 2**15
        y = twos_complement_int(data[3], data[2]) * self.RANGE  / 2**15
        z = twos_complement_int(data[5], data[4]) * self.RANGE  / 2**15
        return x, y, z


# CHATGPT AF
def twos_complement_int(LSB, MSB):
    # Concatenate the two bytes into a single 16-bit value
    value = (MSB << 8) | LSB
    
    # If the highest bit (bit 15) is set, the number is negative
    if value & 0x8000:
        # Convert to signed integer by subtracting 65536
        return -((~value & 0xFFFF) + 1)
    else:
        return value

def init_a_bunch():
    power_off()
    for i in range(ROWS):
        for j in range(COLUMNS):
            n = 0x20 + i * COLUMNS + (8-j -1)
            try:
                print (f'initalizing {i}, {j} to {hex(n)}')
                power_on(i*COLUMNS + j)
                sleep(.1)
                if DEBUG: input('Press enter to config it: ')
                try:
                    TMAG5273.init(n)
                except OSError as e: print(f'not at default addy', e)
 
            except Exception as e:
                print('could not initialize', e)
        #switches.enableAll(); input("row boundary") # for debug, enable all after each row
    #switches.enableAll()
    #try: print(switches.bus.read_byte(PCA_ADDR))
    #except: pass
    #print(switches.bus.read_byte(0x71))

def get_a_bunch(start, stop):
    ret = []
    for i in range(start, stop):
        #y,x = (i-0x10) // COLUMNS, (i-0x10) % COLUMNS
        #if x==0:
        #    switches.unsafeEnableRow(y)

        #if DEBUG: input("enabling " + str((y)))
        tmag5273 = TMAG5273(i)
        try:
            ret.append(tmag5273.read_magnetism())
        except OSError as e:
            print(f"error reading {hex(i)}: {e}. Attempting reinit...")

            try:
                # turn tmag on and off
                power_off_single(i)
                sleep(0.05)
                power_on_single(i)
                sleep(0.1)

                TMAG5273.init(i)

                tmag = TMAG5273(i)
                ret.append(tmag.read_magnetism())

                print(f"Recovered {hex(i)} successfully")

            except Exception as e2:
                print(f"Failed to recover {hex(i)}: {e2}")
                ret.append((0, 0, 0))  # placeholder if still failed

        except Exception as e:
            print(f"Unexpected error with {hex(i)}: {e}")
            ret.append((0, 0, 0))

        #if x==7:
        #    switches.unsafeDisableRow(y)
    return ret



def plotty_plot():
    import numpy as np
    import matplotlib.pyplot as plt
    
    # Enable interactive mode
    plt.ion()
    
    # Define the chessboard size
    rows, cols = 8, 8
    
    # Create a chessboard pattern
    chessboard = np.zeros((rows, cols))
    chessboard[1::2, ::2] = 1
    chessboard[::2, 1::2] = 1
    
    # Create an empty plot
    fig, ax = plt.subplots()
    im = ax.imshow(np.zeros((rows, cols)), cmap='bwr', interpolation='nearest', clim=(-5,5))
    #ax.imshow(chessboard, cmap='binary', alpha=0.3)
    colorbar = plt.colorbar(im, label='Score' )
    #colorbar.set_clim(-13, 13)
    plt.title('Chessboard with Randomly Updating Scores')
    plt.xticks(np.arange(cols), labels=[chr(97+i) for i in range(cols)])  # Add x-axis labels (a-h)
    plt.yticks(np.arange(rows), labels=[str(i+1) for i in range(rows)])  # Add y-axis labels (1-8)
    plt.gca().invert_yaxis()  # Invert y-axis to match chessboard orientation
    
    # Start a loop to continuously update the plot
    while True:
        # Generate random scores for each square
        scores = np.zeros((rows, cols))
        ans = get_a_bunch(0x20, 0x20 + ROWS*COLUMNS)
        ans = [i[2] for i in ans] # z-axis only
        for i in range(ROWS):
            for j in range(COLUMNS):
                scores[i][j] = ans[i * COLUMNS + j]
        
        print(scores)
        # Update the image data
        im.set_array(scores)
        
        # Redraw the plot
        fig.canvas.draw()
        fig.canvas.flush_events()

def on_connect(client, userdata, flags, reason_code, properties):
    if reason_code.is_failure:
        print(f"Failed to connect: {reason_code}. loop_forever() will retry connection")
    else:
        print(f"Successfully connected")
        # we should always subscribe from on_connect callback to be sure
        # our subscribed is persisted across reconnections.
        #client.subscribe("$SYS/#")
        #client.subscribe("/led")


if __name__ == "__main__":
    #plotty_plot()
    #exit(1)

    print(mqttc.connect("localhost"))
    mqttc.loop_start()
    #init_a_bunch()
    while True:
        t = time()
        ans = get_a_bunch(0x20, 0x20 + ROWS*COLUMNS)
        if DEBUG: print([round(i[2], 1) for i in ans])
        flat_list = [item for sublist in ans for item in sublist]
        byte_array = b''.join(struct.pack('f', num) for num in flat_list)
        p = mqttc.publish("/magnets", byte_array, retain=False, qos=0)
        #if DEBUG: print(p)
        #print(time(), 0.03 - ( time() - t))
        #print("Magnetic Field (X,Y,Z):", round(x, 1), round(y, 1), round(z, 1))
        sleep(max(0.03 - ( time() - t), 0 ) ) # approx 50 hz like Vex pid
    mqttc.loop_stop()
    # Wait for some time before the next reading