wave1 = [0.0, 50.0, 100.0]
wave2 = [0.5, 20.5, 40.5, 60.5, 80.5, 100.5, 120.5, 140.5]
wave3 = [1.0, 11.0, 21.0, 31.0, 41.0, 51.0, 61.0, 71.0, 81.0, 91.0, 101.0, 111.0, 121.0, 131.0, 141.0]

worst = []
i = 0
n = 6                           # Number of worst combination we are looking for (the worst of the worst)

freqs = [20, 30, 40, 50, 60, 80, 100, 125, 250, 500]

for subset in itertools.combinations(freqs, 5):
    stim_overlap = optimal_delay(optimization_method, t0, tf, list(subset), shifts)

    if i < n:
        worst.append((subset, stim_overlap[2]))
        i += 1
        continue
    elif True in [stim_overlap[2] > worst[n][1] for n in range(len(worst))]:
        id = [stim_overlap[2] > worst[n][1] for n in range(len(worst))].index(True)
        worst[id] = (subset, stim_overlap[2])
        i += 1
        continue
    else:
        i +=1
        continue

((40, 60, 80, 250, 500), 1.0673737373737373)
((50, 60, 80, 250, 500), 1.0248737373737373)
((60, 80, 125, 250, 500), 0.961540404040404)
((40, 80, 125, 250, 500), 0.785)
((60, 80, 100, 250, 500), 0.7357070707070708)
((60, 80, 100, 125, 500), 0.6732070707070708)

# -*- coding: utf-8 -*-

import numpy as np
import itertools

"""-----------------------------------
------------- Pulse shape ------------
-----------------------------------"""

# Pulse width
global width
width = 0.3

"""-----------------------------------
------------- Parameters -------------
-----------------------------------"""

t0 = 0
tf = 200                                        # Pulses can't start after 200-0.3 ms.
shifts = np.arange(0, 2.05, 0.05)               # shifts range

"""-----------------------------------
------------- Functions --------------
-----------------------------------"""

def time_builder(f, t0=0, tf=300):
    """
    Function building a list of time points at which a pulse is triggered.
    """
    return list(np.round(np.arange(t0, tf, 1/f*1000),3))

def duo_stim_overlap(t1, t2):
    """
    Function taking 2 timelines build by time_builder function in input
    and returning the ids of overlapping pulses between the 2.
    len(t1) < len(t2)
    """
    pulse_id_t1 = [x for x in range(len(t1)) for y in range(len(t2)) if abs(t1[x] - t2[y]) < width]
    pulse_id_t2 = [x for x in range(len(t2)) for y in range(len(t1)) if abs(t2[x] - t1[y]) < width]

    return pulse_id_t1, pulse_id_t2

def optimal_delay(t0, tf, frequences, shifts):
    """
    Function computing the optimal shift with the optimization method chosen for n waveforms.
    - t0 and tf in ms
    - frequences: a list of the frequencies to use
    - shifts: 1D array like with the shifts to test
    """

    optimal = []                            # Result tuple of the form (shift, dict{f_id:  [overlapping pulses ids]}, optim parameter)
    First = True                            # First shift iteration condition

    for s in shifts:
        delay = 0                           # delay between signals, i.e shift
        overlap = dict()
        timelines = list()

        for i in range(len(frequences)):
            timelines.append(time_builder(frequences[i], t0+delay, tf))
            overlap[i] = list()
            delay += s

        for subset in itertools.combinations(timelines, 2):
            p1_stim, p2_stim = duo_stim_overlap(subset[0], subset[1])
            overlap[timelines.index(subset[0])] += p1_stim
            overlap[timelines.index(subset[1])] += p2_stim

        optim_param = 0
        for key, items in overlap.items():
            optim_param += (len(list(set(items))) / len(timelines[key]))

        if First or optim_param < optimal[2]:
            optimal = (s, overlap, optim_param)
            First = False
        else:
            continue

    return optimal

"""-----------------------------------
------------- Program ----------------
-----------------------------------"""

worst = []
i = 0
n = 6                           # Number of worst combination we are looking for (the worst of the worst)

freqs = [20, 30, 40, 50, 60, 80, 100, 125, 250, 500]

for subset in itertools.combinations(freqs, 5):
    stim_overlap = optimal_delay(t0, tf, list(subset), shifts)

    if i < n:
        worst.append((subset, stim_overlap[2]))
        i += 1
        continue
    elif True in [stim_overlap[2] > worst[n][1] for n in range(len(worst))]:
        id = [stim_overlap[2] > worst[n][1] for n in range(len(worst))].index(True)
        worst[id] = (subset, stim_overlap[2])
        i += 1
        continue
    else:
        i +=1
        continue