import matplotlib.pyplot as plt
import numpy as np
from scipy import signal
import time

format_string=(
    b"""packet_rtc 1 28 uint32_t t 0 4 RTC_TimeTypeDef sTime 4 20 RTC_DateTypeDef sDate 24 4
packet_vbatt 2 8 uint32_t t 0 4 uint16_t vbatt_cnts 4 2
packet_imu 9 12 uint32_t t 0 4 uint16_t readings_cnts[4] 4 2
packet_ekg 3 208 uint32_t t 0 4 uint8_t index 4 1 int32_t readings_cnts[50] 8 4
packet_strain 4 28 uint32_t t 0 4 uint8_t index 4 1 int32_t readings_cnts[5] 8 4
packet_outsideT 5 28 uint32_t t 0 4 uint8_t index 4 1 int32_t readings_cnts[5] 8 4
packet_insideT 6 28 uint32_t t 0 4 uint8_t index 4 1 int32_t readings_cnts[5] 8 4
packet_button 7 8 uint32_t t 0 4 uint8_t button_vec 4 1
packet_spo2 8 304 uint32_t t 0 4 uint32_t green_cnts[25] 4 4 uint32_t red_cnts[25] 104 4 uint32_t ir_cnts[25] 204 4""")

types = []

def arr_sizes(s):
    if s[-1:] != b']' or b'[' not in s:
        return 1
    v = int(s[s.rindex(b'[') + 1:-1])
    return v

for line in format_string.split(b"\n"):
    line = line.strip(b'\n').strip(b'\r')
    if line != b'':
        L = line.split(b" ")
        types.append({'type_name' : L[0],
                      'type_code' : int(L[1]),
                      'size' : int(L[2]),
                      'elements' : []
                      })
        i = 3
        while i < len(L):
            types[-1]['elements'].append({'type_name' : L[i], 'name' : L[i + 1], 'offset' : int(L[i + 2]), 'n_elements' : arr_sizes(L[i + 1]), 'size' : int(L[i + 3]) * arr_sizes(L[i + 1]), 'readings' : []})
            i += 4

import asyncio
from bleak import BleakScanner, BleakClient

async def scan():
    return await BleakScanner.find_device_by_name("XX-STM32")

async def connect(device):
    async with BleakClient(device) as client:
        value = bytes([0x01])
        #await client.write_gatt_char(TX_UUID, value, response=True)
        await client.start_notify(RX_UUID, cb)
        await asyncio.sleep(10000)

RX_UUID = "00000001-8E22-4541-9D4C-21EDAE82ED19"
TX_UUID = "00000000-8E22-4541-9D4C-21EDAE82ED19"

queue = asyncio.Queue()

def cb(sender, data):
    queue.put_nowait(data)

def my_filter(arr):
    inds = np.where(np.abs(np.diff(arr)) > 0.005)[0]
    for ind in inds:
        prev_ind = ind - 1
        while prev_ind in inds:
            prev_ind -= 1
        next_ind = ind + 1
        while next_ind in inds:
            next_ind += 1
        arr[ind] = 0.5 * arr[prev_ind] + 0.5 * arr[next_ind]
    return arr

# I think this is missing a lot of data (70Hz vs should have 250Hz)

async def update_with_data():

    fig, axs = plt.subplots(2)

    adcs = []
    accs = []
    gyros = []
    
    cons = bytes()
    start = time.time()
    while(True):
        data = await queue.get()
        cons = cons + data
        if len(cons) >= 1024:
            assert(len(cons) == 1024)

            index = 0
            while (index < 1024):
                packet_type = (cons[index + 1]<<8) + cons[index]
                print(packet_type, index)#, len(cons))
                ind = [i for i,t in enumerate(types) if t['type_code'] == packet_type]
                if len(ind) != 1:
                    if (index == 0):
                        cons = cons[128:]
                    else:
                        cons = bytes()
                    break
                t = types[ind[0]]
                d = cons[index + 2 : index + 2 + t['size']]
                for e in t['elements']:
                    block = d[e['offset']:e['offset'] + e['size']]
                    element_size = int(len(block) / e['n_elements'])
                    if t['type_name'] == b'packet_ekg' and e['name'] == b'readings_cnts[50]':
                        chunked = [(2.4 / (1<<24)) * int.from_bytes(block[i:i + element_size], byteorder='little', signed = True) for i in range(0, len(block), element_size)]
                        adcs += chunked
                        if len(adcs) > 500:
                            dat = np.array(adcs[-2048:])
                            filtered_dat = my_filter(dat)
                            dd = np.sort(dat)
                            center = 0.5 * (dd[-100] + dd[100])
                            range_ = dd[-100] - dd[100]
                            if len(dat) == 2048:
                                print("Freq:", 2048 / (time.time() - start))
                                axs[0].cla()
                                axs[1].cla()
                                axs[0].set_title("ADC {:3.3f}mV".format(1000 * (range_)))
                                               
                                #axs[0].plot(dat, 'r.', linestyle = '--')
                                axs[0].plot(filtered_dat, 'k.', linestyle = '--')
                                axs[0].set_ylim(center - range_, center + range_)
                                _fft = np.log(np.abs(np.fft.fft(filtered_dat)[1:]))
                                axs[1].plot(250 * np.fft.fftfreq(dat.shape[-1])[1:], _fft, 'k.')
                                axs[1].axvline(x = 60)
                                plt.savefig("image.png")

                    if t['type_name'] == b'packet_imu' and e['name'] == b'readings_cnts[4]':
                        imu_reading_type = block[0] >> 3
                        imu_reading_tag_cnt = (block[0] >> 1) & 3
                        reading_xyz = [int.from_bytes(block[2:4], byteorder = 'big', signed = True),
                                       int.from_bytes(block[4:6], byteorder = 'big', signed = True),
                                       int.from_bytes(block[6:8], byteorder = 'big', signed = True)]
                        if imu_reading_type == 1:
                            #print("{:02x} {:d}".format(imu_reading_type, imu_reading_tag_cnt))
                            gyros.append([250 * e / (1<<16) for e in reading_xyz])
                        elif imu_reading_type == 2:
                            accs.append([4 * e / (1<<16) for e in reading_xyz])
                        else:
                            print("ERR")
                index += 2 + t['size']

async def main():
    device = await scan()
    if not device:
        print("Device not found")
        return

    consumer_task = asyncio.create_task(update_with_data())

    await connect(device)

asyncio.run(main())