import socket
import struct
import uart_monitor_config as c
from networkx import Graph, DiGraph, simple_cycles, connected_components
import networkx as nx
import sys
import time
import numpy as np
from functools import partial

MCAST_GRP = c.udp_grp
MCAST_PORT = c.udp_raw_port

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(('', MCAST_PORT))

mreq = struct.pack("4sl", socket.inet_aton(MCAST_GRP), socket.INADDR_ANY)
sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)


samplerate = c.samplerate
channel_idx = 2
baud_rates_test = c.baud_rates_test


do_graph_analysis = False
recv_size = 1024*256

sock_send = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
sock_send.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, c.MULTICAST_TTL)

sock_send_funcs = dict()

def send_packet(data, flags = None):
    print(data, flags)
    sock_send.sendto(data, flags)

for b_rate, udp_port in c.baud_rate_map.items():
#    sock_send_funcs[b_rate] = partial(sock_send.sendto, flags = (MCAST_GRP,MCAST_PORT))
    sock_send_funcs[b_rate] = partial(send_packet, flags = (MCAST_GRP,udp_port))
    

# %%
    


bcands = dict()
bcands_zero = dict()
pre_bits = {x:np.zeros(0,dtype=np.uint8) for x in baud_rates_test}
while True:
    dat = sock.recv(recv_size)
    bits = np.unpackbits(np.frombuffer(dat,dtype=np.uint8)).reshape((-1,8))[:,::-1]
    bit_stream_pre = bits[:,channel_idx]

    
    for baud_rate in baud_rates_test:
        bit_stream = np.concatenate([ pre_bits[baud_rate]  ,bit_stream_pre])
        all_to_zero_total = np.where(np.diff(bit_stream) == 255)[0] #bit_stream is uint8, so transition f
        bit_period = 1.0 / float(baud_rate)
        brhz=1/samplerate
        bspace = bit_period/brhz
        to_keep = all_to_zero_total < (len(bit_stream) - 10*bspace)
        to_not_keep = np.logical_not(to_keep)
        all_to_zero = all_to_zero_total [to_keep]
        frames = np.arange(0,8*bspace, bspace, dtype=np.int64) + int(bspace/2 + bspace)
        frame_length = 9*bspace
        idx_scan = all_to_zero[:,None] + frames[None,:]
        possible_arrays = bit_stream[idx_scan]
        if len(possible_arrays) > 0:
            packed_bits = np.packbits(possible_arrays,axis=1,bitorder='little')
        else:
            packed_bits = []

        diff_this = np.insert(all_to_zero,0,0)
        sum_ending = np.sum(all_to_zero_total > (len(bit_stream) - 11*bspace))
        big_spaces = np.where(np.diff(diff_this) / (bspace/2) > 10.25)[0]

        diff_this_total =  np.insert(all_to_zero_total,0,0)
        big_spaces_total = np.where(np.diff(diff_this_total)/ (bspace/2) > 10)[0]

        if len(big_spaces) > 0:
            relevant_bits = packed_bits[big_spaces].flatten()
            print(baud_rate, list(relevant_bits))
            sys.stdout.flush()

            for x in relevant_bits:
                print('Sending')
                sock_send_funcs[baud_rate](x)
        else:
            relevant_bits = []

 
        set_diff = set(big_spaces_total) - set(big_spaces)
        if len(set_diff)>0:
            space_start = max(set_diff)
            m1 = np.ones((int(11*bspace)), dtype=np.uint8)
            m2 = bit_stream[int(diff_this_total[space_start+1] - bspace):]
            pre_bits[baud_rate]=np.concatenate((m1,m2))

        elif np.sum(to_not_keep) > 0:
            m1 = np.ones((int(11*bspace)), dtype=np.uint8)
            m2 = bit_stream[all_to_zero_total[np.argmax(to_not_keep)]:]
            pre_bits[baud_rate]=np.concatenate((m1,m2))
        else:
            pre_bits[baud_rate] = np.zeros(0)    

        pre_bits[baud_rate] = bit_stream[-int(frame_length):]
        
        if do_graph_analysis and len(all_to_zero)>0:
            atz = all_to_zero
            diff_mat = atz[:,None] - atz[None,:]
            #id_trav = np.logical_or(diff_mat > frame_length,diff_mat == 0)
            id_trav = diff_mat > frame_length
            conn = np.argmax(id_trav,axis=0)
            con_dict = dict()

            for idf, cc in zip(range(len(conn)),conn):
                if cc == 0:
                    leaves.append(str(idf))
                    continue

                con_dict[str(idf)] = [str(cc)]


            G = DiGraph(con_dict)
            roots = [i for i,d in G.in_degree() if d == 0]
            leaves = [i for i,d in G.out_degree() if d == 0]            


            all_ps = list()
            for root in roots:
                paths = nx.all_simple_paths(G, root, leaves)
                all_ps.extend(paths)