"""In this example, we learn how to remove events on the device containing
no information to save USB bandwidth."""

from matplotlib import pyplot as plt
import sys

from Swabian import TimeTagger


def find_random_signal(tt, threshold_rate):
    # This function scans the trigger level in the same way as
    # it is shown in example 2-B.
    channels = tt.getChannelList(TimeTagger.ChannelEdge.Rising)
    level = max(-0.01, tt.getTriggerLevelRange(channels[0])[0])
    # If using the Time Tagger X, the input hysteresis value should be set to the lowest possible value for all input channels.
    if tagger.getModel() == 'Time Tagger X':
        [tt.setInputHysteresis(channel,  1) for channel in channels]
    count = TimeTagger.Countrate(tagger=tt, channels=channels)
    print("")
    while level < 0.01:
        sys.stdout.write("Searching level: {:5.1f} mV\r".format(level * 1000))
        for channel in channels:
            tt.setTriggerLevel(channel, level)
        count.startFor(int(1E8))
        count.waitUntilFinished()
        data = count.getData()
        if max(data) > threshold_rate:
            return channels[data.argmax()], level, True
        level += 0.0001
    return 0, .0, False


# Create a TimeTagger instance to control your hardware
tagger = TimeTagger.createTimeTagger()

# Set a threshold for the random counts according to the Time Tagger model
min_random_rate = 3E7 if tagger.getModel() == "Time Tagger 20" else 1E8

# Try to find a channel that exceeds the count rate threshold and set the
# trigger level accordingly
random_channel, random_level, random_found = find_random_signal(
    tagger, min_random_rate)
if random_found:
    print("\nFound a random signal on channel {} at {:2.3} mV.".format(
        random_channel, 1000 * random_level))
    tagger.setTriggerLevel(channel=random_channel, voltage=random_level)

    # The default dead-time is defined by the internal clock of the Time Tagger
    # TT X: 750 MHz, TT Ultra: 500 MHz, TT 20: 166.6 MHz
    default_deadtime = tagger.getDeadtime(random_channel)

    # To measure the count rate that exceeds the USB transfer limit without overflows,
    # we use the EventDivider. With divider=100, we discard 99 out of 100 events and
    # let only one pass every 100th event, which is within the USB bandwidth
    tagger.setEventDivider(channel=random_channel, divider=100)
    count = TimeTagger.Countrate(tagger, [random_channel])
    count.startFor(int(1E12))
    count.waitUntilFinished()

    # To obtain the real count rate, we have to multiply by 100
    countrate = count.getData()[0] * 100
    specified_usb_rate = 7.5E6 if tagger.getModel() == "Time Tagger 20" else 82E6
    print(
        "\nThe countrate of the random channel is {:.1f} Mtags/s".format(countrate / 1E6))
    print("This is {:.1f} % of the specified USB data rate and {:.1f} % of the maximal time-to-digital conversion rate.".format(
        100*countrate/specified_usb_rate, 100*countrate/(1E12/default_deadtime)
    ))

    # Switch off the EventDivider = set it to 1
    tagger.setEventDivider(channel=random_channel, divider=1)

    # Take a second input channel which is NOT the random channel and let it
    # measure the built-in test signal, i.e. a periodic signal
    periodic_channel = 2 if random_channel == 1 else 1
    tagger.setTestSignal(periodic_channel, True)

    # We want to measure only the very first random click after a periodic click.
    # This is achieved by the ConditionalFilter. It opens a gate for any event
    # on the 'trigger' channels and lets only one event of each 'filtered' channel
    # pass. After that, the gate is closed until the next trigger event.
    tagger.setConditionalFilter(
        trigger=[periodic_channel], filtered=[random_channel])

    # We measure the correlation between the periodic and the FILTERED random channels
    corr = TimeTagger.Correlation(tagger=tagger,
                                  channel_1=random_channel,
                                  channel_2=periodic_channel,
                                  binwidth=default_deadtime//100,
                                  n_bins=3000)
    plt.figure(figsize=(9, 6))
    print("\nMeasuring with different dead-times:")
    for i in range(1, 6):
        # We repeat this measurement for five different dead-times
        tagger.setDeadtime(random_channel, i * default_deadtime)
        deadtime = tagger.getDeadtime(random_channel)
        print("Default dead-time x {} = {}".format(i, deadtime))
        corr.startFor(int(1E12))
        corr.waitUntilFinished()
        plt.plot(corr.getIndex()/default_deadtime,
                 corr.getData(), label="{} ps".format(deadtime))
    plt.title("Correlation for different dead-times")
    plt.xlabel("Delay/(default_deadtime = {} ps)".format(default_deadtime))
    plt.ylabel("Counts/bin")
    plt.annotate("""Within the dead-time,
the Correlation is more
or less constant, i.e
randomly distributed.
For larger delays, the
Correlation drops with
the probability of finding
no preceding time-tag.""",
                 (2, max(corr.getData())),
                 xytext=(-150, 0),
                 textcoords='offset pixels',
                 ha='right',
                 arrowprops={'arrowstyle': '->'})
    plt.legend()
    plt.show()

TimeTagger.freeTimeTagger(tagger)
