
import json
import numpy as np
import os
from subprocess import Popen
import time

tuples = ['runway street airplane']
#'initial-transfer', 'initial-tune', 'combined-transfer-ptune','added-transfer-ptune', 'removed-transfer-ptune', 'fs-tune-ptune'
modes = ['initial-tune', 'combined-transfer-ptune', 'fs-tune-ptune']
trials = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
num_gpus = 3

# Run each tuple in a batch
for t in tuples:
    info = t.split(' ')
    label1 = info[0]
    label2 = info[1]
    spurious = info[2]
    
    # Generate all of the configurations we want to run
    configs = []
    for mode in modes:
        for trial in trials:
            config = {}
            config['label1'] = label1
            config['label2'] = label2
            config['spurious'] = spurious
            config['mode'] = mode
            config['trial'] = trial
            configs.append(config)
        
    # Divide the configs among the workers
    configs_worker = [[] for i in range(num_gpus)]
    next_worker = 0
    for config in configs:
        configs_worker[next_worker].append(config)
        next_worker = (next_worker + 1) % num_gpus

    # Save the assignments
    for i in range(num_gpus):
        with open('./Models/{}.json'.format(i), 'w') as f:
            json.dump(configs_worker[i], f)

    # Launch the workers
    commands = []
    for i in range(num_gpus):
        command = 'CUDA_VISIBLE_DEVICES={} python Worker.py {}'.format(i, i)
        commands.append(command)

    procs = []
    for i in commands:
        procs.append(Popen(i, shell = True))
        time.sleep(np.random.uniform(4, 6))

    for p in procs:
       p.wait()
