import time
import sys

sys.path.extend(['../module'])

import pywfs

import itertools
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation as anim
import time

LOCAL_ADDR = "192.168.1.204"
nodes = [ "192.168.1.180" ]

# ---------- CONNECT TO AND PREPARE SAMPLING SYSTEM -----------

# enable logging onto STDOUT
pywfs.Logger.start()

# discover nodes on the network
b = pywfs.ServerBeacon()
b.setInterfaceAddr(LOCAL_ADDR)
b.singleScan()
#nodes = b.getNodesOnNetwork()

# query acquisition format (assume all nodes are identical in this aspect)
f = b.queryFromNode(nodes[0], "snd acqformat mr")
acqf = pywfs.AcquisitionFormat(f)

# create MSR
msr = pywfs.MultiStreamReceiver(nodes, f)
ch_n = msr.getChannelCount()

# save samples to file
# mstf = pywfs.MultiStreamToFile.create("test")
# msr.listen(mstf)

# create oscilloscope data feeder object
osc = pywfs.MultiStreamOscilloscope.create()

#index = 2 * nodes.index("192.168.1.180")

# create a slope trigger
edge_trigger = pywfs.EdgeTrigger()
edge_trigger.level = 0.1
edge_trigger.edge = pywfs.FALLING
#edge_trigger.ch = 0 + index
edge_trigger.ch = 0

osc.trigger = edge_trigger
osc.setScreenPeriod(int(10e+06))

# start listening for samples
msr.listen(osc)

# enable network sample transmission from every device
#b.execCmdOnAllNodes("snd connect 10.42.0.1 20220")
for node in nodes:
    b.execCmdOnNode(node, "snd connect " + LOCAL_ADDR)


# # data collection
# osc.armTrigger()  # arm trigger
# samples = osc.capture()  # wait for captured data
#
# # data collection
# osc.armTrigger()  # arm trigger
# samples = osc.capture()  # wait for captured data
#
# exit(0)

# ------------------ DISPLAY COLLECTED SAMPLES ----------------

def data_gen():
    real_time_last = time.time()
    index_last = 0

    screen_lim = osc.getScreenTimeLimits()
    T = [screen_lim[0] * 1e-09, screen_lim[1] * 1e-09]
    Y = [0, 0]

    for i in itertools.count():
        # FPS measurement
        real_time = time.time()
        if (real_time - real_time_last) > 1.0:
            print(i - index_last, " FPS")
            real_time_last = time.time()
            index_last = i

        # data collection
        osc.armTrigger()  # arm trigger
        channels = osc.capture(200)  # wait for captured data

        # if we have incoming data
        if (len(channels) > 0):
            T = list()
            Y = list()

            for ch in channels:
                t = list()
                y_t = list()

                for sp in ch:
                    t.append(sp.t * 1E-09)
                    y_t.append(sp.y)

                T.append(t)
                Y.append(y_t)

        yield T, Y


def init():
    ax.set_ylim(-1.1, 1.1)
    #ax.set_xlim(-0.01, 0.06)

    xlims = osc.getScreenTimeLimits()
    ax.set_xlim(xlims[0] * 1E-09, xlims[1] * 1E-09)



fig, ax = plt.subplots()
lines = list()

# trigger mark
screen_period = osc.getScreenPeriod() * 1E-09
trig_mark_halflen = 0.02
trig_t = [-trig_mark_halflen * screen_period, trig_mark_halflen * screen_period ]
trig_y = [ edge_trigger.level, edge_trigger.level ]

for k in range(0, ch_n):
    lines.append(ax.plot([], [], lw=2))
    t, y = list(), list()

ax.plot(trig_t, trig_y, color = 'r')

def run(data):
    T, Y = data
    for k in range(0, ch_n):
        lines[k][0].set_data(T[k], Y[k])

    return lines[0], lines[1],


ani = anim.FuncAnimation(fig, run, data_gen, interval=20, init_func=init)
plt.grid()
plt.show()

# disconnect nodes
b.execCmdOnAllNodes("snd disconnect")
msr.close()