import subprocess
import os

# Change the working directory to where the training command needs to be run
BASE_PATH = os.environ.get("BASE_PATH","")
if BASE_PATH != "":
    os.chdir(BASE_PATH)
# Run the command and capture the output

def run_command(command, loss_upper_bound, test_name):
    #output = subprocess.check_output(command, shell=True, text=True)
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    output, _ = process.communicate()
    # Parse the output to find excess_Xentropy_loss value
    lines = output.split('\n')[::-1]
    excess_Xentropy_loss = None
    for line in lines:
        
        if "val/excess_Xentropy_loss=" in line:
            tokens = line.split(' ')
            for token in tokens:
                if "val/excess_Xentropy_loss=" in token:
                    excess_Xentropy_loss = float(token.split('=')[1][:-1])
            
            break

    # Check if excess_Xentropy_loss is less than 0.005
    if excess_Xentropy_loss is not None and excess_Xentropy_loss < loss_upper_bound:
        print(f"Test {test_name} PASSED! Loss is {excess_Xentropy_loss} and upper bound is {loss_upper_bound}")
    else:
        if excess_Xentropy_loss is not None:
            print(f"Test {test_name} FAILED: excess_Xentropy_loss is {excess_Xentropy_loss}, which is not less than {loss_upper_bound}")
        else:
            print(f"Test {test_name} FAILED: excess_Xentropy_loss value not found in the output")

commands_small = [
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_SMALL ++dataset.num_seq=100 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=1 ++trainer.max_epochs=20",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_SMALL ++dataset.num_seq=100 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=1 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_SMALL ++dataset.num_seq=100 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=1 ++trainer.max_epochs=100",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_SMALL ++dataset.num_seq=100 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=1 ++trainer.max_epochs=20",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_SMALL ++dataset.num_seq=100 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=1 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_SMALL ++dataset.num_seq=100 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=1 ++trainer.max_epochs=100",
    ]
loss_upper_bound_small = [0.03,0.007,0.03, 0.01, 0.004, 0.04]  #updated nov 12
test_name_small = [
    "linear, 20 epochs, 100 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 40 epochs, 100 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 100 epochs, 100 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 20 epochs, 100 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 40 epochs, 100 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 100 epochs, 100 sequences, 5 batch size, 1 pool size, veasy",
    ]
assert len(commands_small) == len(loss_upper_bound_small) == len(test_name_small)
def run_small():
    for i in range(len(commands_small)):
        run_command(commands_small[i], loss_upper_bound_small[i], test_name_small[i])
#run_small()

commands_medium = [
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=20",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=100",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=20",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=100",
    ]

loss_upper_bound_medium = [0.011,0.007,0.013, 0.0013, 0.006, 0.006]
test_name_medium = [
    "linear, 20 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 40 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 100 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 20 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 40 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 100 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    ]


def run_medium():
    for i in range(len(commands_medium)):
        run_command(commands_medium[i], loss_upper_bound_medium[i], test_name_medium[i])

commands_large = [
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_LARGE ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=20",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_LARGE ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_LARGE ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=100",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=20",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=100",
    ]
loss_upper_bound_large = [0.0025,0.002,0.002, 0.002, 0.002, 0.0004]
test_name_large = [
    "linear, 20 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 100 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 20 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 100 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    ]

def run_large():
    for i in range(len(commands_large)):
        run_command(commands_large[i], loss_upper_bound_large[i], test_name_large[i])


commands_large_2 = [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_easy ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_medium ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_hard ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_vhard ++dataset.num_seq=10000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40",

    ]
loss_upper_bound_large_2 = [0.0020,0.003,0.003, 0.0035, 0.007]
test_name_large_2 = [
    "long conv, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, easy",
    "long conv, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, medium",
    "long conv, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, hard",
    "long conv, 40 epochs, 10000 sequences, 5 batch size, 1 pool size, vhard",

    ]

def run_large_2():
    for i in range(len(commands_large_2)):
        run_command(commands_large_2[i], loss_upper_bound_large_2[i], test_name_large_2[i])


###TESTING THE SET ARCHITECTURE:
commands_medium_set= [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_veasy_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_easy_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_medium_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_hard_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_vhard_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_veasy_pooled ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_easy_pooled ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_medium_pooled ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_hard_pooled ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM_vhard_pooled ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    # The last five models are not set models, there is a pool dependency in the loans that is ignored
    # The set model aims to create a low dimensional feature to model the dependencies between the loans.
    ]
loss_upper_bound_medium_set = [0.017,0.09,0.03, 0.08, 0.08,0.018,0.08,0.03, 0.08, 0.07]
test_name_medium_set = [
    "long conv set, pool 10, 40 epochs, 100 sequences, 1 batch size, veasy",
    "long conv set, pool 10, 40 epochs, 100 sequences, 1 batch size, easy",
    "long conv set, pool 10, 40 epochs, 100 sequences, 1 batch size, medium",
    "long conv set, pool 10, 40 epochs, 100 sequences, 1 batch size, hard",
    "long conv set, pool 10, 40 epochs, 100 sequences, 1 batch size, vhard",
    "long conv, pool 10, 40 epochs, 100 sequences, 1 batch size, veasy",
    "long conv, pool 10, 40 epochs, 100 sequences, 1 batch size, easy",
    "long conv, pool 10, 40 epochs, 100 sequences, 1 batch size, medium",
    "long conv, pool 10, 40 epochs, 100 sequences, 1 batch size, hard",
    "long conv, pool 10, 40 epochs, 100 sequences, 1 batch size, vhard",

    ]

def run_medium_set():
    for i in range(len(commands_medium_set)):
        run_command(commands_medium_set[i], loss_upper_bound_medium_set[i], test_name_medium_set[i])



commands_large_set= [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy_set ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_easy_set ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_medium_set ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_hard_set ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_vhard_set ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy_pooled ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_easy_pooled ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_medium_pooled ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_hard_pooled ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_vhard_pooled ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    # The last five models are not set models, there is a pool dependency in the loans that is ignored
    # The set model aims to create a low dimensional feature to model the dependencies between the loans.
    ]
loss_upper_bound_large_set = [0.003,0.004,0.0035, 0.004, 0.006,0.0035,0.0045,0.0045,0.0045, 0.006]
test_name_large_set = [
    "long conv set, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, easy",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, medium",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, hard",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, vhard",
    "long conv, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, veasy",
    "long conv, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, easy",
    "long conv, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, medium",
    "long conv, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, hard",
    "long conv, pool 10, 40 epochs, 1000 sequences, 1 batch size, 10 pool size, vhard",
    ]

def run_large_set():
    for i in range(len(commands_large_set)):
        run_command(commands_large_set[i], loss_upper_bound_large_set[i], test_name_large_set[i])



commands_xlarge_set= [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_veasy_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_easy_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_medium_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_veasy_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_easy_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_medium_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    # The last five models are not set models, there is a pool dependency in the loans that is ignored
    # The set model aims to create a low dimensional feature to model the dependencies between the loans.
    ]
loss_upper_bound_xlarge_set = [0.0006,0.0004,0.0006, 0.0015, 0.0009,0.0035,0.0035,0.0035,0.0015, 0.0015]
test_name_xlarge_set = [
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, easy",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, medium",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, hard",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, vhard",
    "long conv, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, veasy",
    "long conv, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, easy",
    "long conv, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, medium",
    "long conv, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, hard",
    "long conv, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, vhard",
   ]

def run_xlarge_set():
   for i in range(len(commands_xlarge_set)):
        run_command(commands_xlarge_set[i], loss_upper_bound_xlarge_set[i], test_name_xlarge_set[i])



commands_medium_supereasy = [
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=100",
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=supereasy ++trainer.devices=8 ++trainer.max_epochs=100",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=100",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_MEDIUM ++dataset.num_seq=1000 ++loader.batch_size=5 ++dataset.loan_pool_size=1 ++dataset.generator.level=supereasy ++trainer.devices=8 ++trainer.max_epochs=100",
    ]

loss_upper_bound_medium_supereasy = [0.008,0.003, 0.005, 0.003]
test_name_medium_supereasy = [
    "linear, 100 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "linear, 100 epochs, 1000 sequences, 5 batch size, 1 pool size, supereasy",
    "long conv, 100 epochs, 1000 sequences, 5 batch size, 1 pool size, veasy",
    "long conv, 100 epochs, 1000 sequences, 5 batch size, 1 pool size, supereasy",
    ]


def run_medium_supereasy():
    for i in range(len(commands_medium_supereasy)):
        run_command(commands_medium_supereasy[i], loss_upper_bound_medium_supereasy[i], test_name_medium_supereasy[i])

###Experiment to add: Many loans with no pool dimension. 

commands_xlarge_stack_pool= [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_veasy_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_easy_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_medium_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    ]
loss_upper_bound_xlarge_stack_pool = [0.003,0.03,0.1, 0.08, 0.1]
test_name_xlarge_stack_pool = [
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, easy",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, medium",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, hard",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, vhard",

   ]
def run_xlarge_stack_pool():
   for i in range(len(commands_xlarge_stack_pool)):
        run_command(commands_xlarge_stack_pool[i], loss_upper_bound_xlarge_stack_pool[i], test_name_xlarge_stack_pool[i])




commands_xlarge_vhard= [
   
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_stack_pool ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++decoder.loan_pool_size=10",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_pooled ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    ]
loss_upper_bound_xlarge_vhard = [0.06,0.09,0.0014, 0.0009, 0.0014, 0.0015]
test_name_xlarge_vhard = [
    "long conv set stack, pool 10, 40 epochs, 10000 sequences, 1 batch size, hard",
    "long conv set stack, pool 10, 40 epochs, 10000 sequences, 1 batch size, vhard",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, hard",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, vhard",
    "long conv set pooled, pool 10, 40 epochs, 10000 sequences, 1 batch size, hard",
    "long conv set pooled, pool 10, 40 epochs, 10000 sequences, 1 batch size, vhard",

   ]
def run_xlarge_vhard():
    for i in range(len(commands_xlarge_vhard)):
        run_command(commands_xlarge_vhard[i], loss_upper_bound_xlarge_vhard[i], test_name_xlarge_vhard[i])




commands_xlarge_set_lbz= [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_veasy_set ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_easy_set ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_medium_set ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_set ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_set ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_veasy_pooled ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_easy_pooled ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_medium_pooled ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_pooled ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_vhard_pooled ++dataset.num_seq=10000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=vhard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.debug=True ++decoder.loan_pool_size=1",
    # The last five models are not set models, there is a pool dependency in the loans that is ignored
    # The set model aims to create a low dimensional feature to model the dependencies between the loans.
    ]


loss_upper_bound_xlarge_set_lbz = [0.0006,0.0004,0.0006, 0.0015, 0.0009,0.0035,0.0035,0.0035,0.0015, 0.0015]
test_name_xlarge_set_lbz = [
    "long conv set, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, easy",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, medium",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, hard",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, vhard",
    "long conv, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, veasy",
    "long conv, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, easy",
    "long conv, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, medium",
    "long conv, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, hard",
    "long conv, pool 10, 40 epochs, 10000 sequences, 20 batch size, 10 pool size, vhard",
   ]

def run_xlarge_set_lbz():
    for i in range(len(commands_xlarge_set_lbz)):
        run_command(commands_xlarge_set_lbz[i], loss_upper_bound_xlarge_set_lbz[i],test_name_xlarge_set_lbz[i])



test_diff_pool = [
    "long conv set, pool 10, 40 epochs, 1000 sequences, 20 batch size, 10 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 4 batch size, 50 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 1 batch size, 200 pool size, veasy",
    "long conv set, pool 10, 40 epochs, 1000 sequences, 40 batch size, 5 pool size, veasy"]
loss_diff_pool = [0.0015,0.0015,0.0015, 0.0015]
### Test different number of pool dimensions (keeping to total batch size fixed)
commands_different_pool= [
"python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy_set_pool_10 ++dataset.num_seq=1000 ++loader.batch_size=20 ++dataset.loan_pool_size=10 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
"python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy_set_pool_50 ++dataset.num_seq=1000 ++loader.batch_size=4 ++dataset.loan_pool_size=50 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
"python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy_set_pool_200 ++dataset.num_seq=1000 ++loader.batch_size=1 ++dataset.loan_pool_size=200 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1",
"python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_LARGE_veasy_set_pool_5 ++dataset.num_seq=1000 ++loader.batch_size=40 ++dataset.loan_pool_size=5 ++dataset.generator.level=veasy ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1"
]

def run_large_pool():
    for i in range(len(commands_different_pool)):
        run_command(commands_different_pool[i], loss_diff_pool[i], test_diff_pool[i])



#run_medium_supereasy()
#run_small()
#run_medium()
#run_large() #Needs update
#run_large_2() #Needs update
#run_medium_set()
#run_large_set() #Needs update
#run_xlarge_set() # Needs update
#run_xlarge_stack_pool() # Needs update
#run_xlarge_vhard() # Needs update
#run_xlarge_set_lbz() # Needs update
#run_large_pool() # Needs update
#
# "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_2d_path_dep_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=2d_path_dependency ++trainer.devices=1 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.common_pool_embedding_dim=2 ++decoder.loan_pool_size=1"
# "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_2d_path_dep_set ++dataset.num_seq=10000 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=2d_path_dependency ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++encoder.common_pool_embedding_dim=2 ++decoder.loan_pool_size=1"

# "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_2d_path_dep_set ++dataset.num_seq=100 ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=2d_path_dependency ++trainer.devices=1 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1
        


# Experiments needed for the papers:
# 1. Test the performance of the set sequence model in isolation.
        

nr_seq = 10000
commands_large_set_sequence= [
   
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_seasy1d_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_1d ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_seasy2d_no_loan_feature_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_2d_no_loan_specific_feature ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_seasy2d_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_2d ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_easy2d_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=easy_2d ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_medium_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=medium ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_hard_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=hard ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    # The set model aims to create a low dimensional feature to model the dependencies between the loans.
    ]


loss_upper_bound_large_set_sequence = [0.00212,0.011,0.0192, 0.056, 0.0549,0.0259]
test_name_large_set_sequence = [
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_1d",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_2d_no_loan_specific_feature",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_2d",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, easy_2d",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, medium",
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, hard",
   ]

def run_large_set_sequence():
    for i in range(len(commands_large_set_sequence)):
        run_command(commands_large_set_sequence[i], loss_upper_bound_large_set_sequence[i],test_name_large_set_sequence[i])

#run_large_set_sequence()

# 2. Compare the performance of the set sequence model with other baselines. 
#Models to compare against
# 1. Void set model
# 2. Stacking in the sequence dimension
# 3. One head per loan
# 4. Linear model
# 5. Attention
# 6. Feedforward NN
nr_seq = 10000
commands_compare_set_sequence_baseline= [
    "python -m train experiment=timeseries/ts_lc.yaml ++wandb.name=ts_lc_TEST_xLARGE_seasy2d_no_loan_feature_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_2d_no_loan_specific_feature ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_id.yaml ++wandb.name=ts_id_TEST_xLARGE_seasy2d_no_loan_feature_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_2d_no_loan_specific_feature ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_ff.yaml ++wandb.name=ts_ff_TEST_xLARGE_seasy2d_no_loan_feature_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_2d_no_loan_specific_feature ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    "python -m train experiment=timeseries/ts_mha.yaml ++wandb.name=ts_mha_TEST_xLARGE_seasy2d_no_loan_feature_set ++dataset.num_seq={} ++loader.batch_size=1 ++dataset.loan_pool_size=10 ++dataset.generator.level=supereasy_2d_no_loan_specific_feature ++trainer.devices=8 ++trainer.max_epochs=40 ++encoder._name_=set_encoder ++decoder.loan_pool_size=1".format(nr_seq),
    ]


loss_upper_bound_compare_set_sequence_baseline = [0.01, 0.01, 0.01, 0.01]
test_name_compare_set_sequence_baseline = [
    "long conv set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_2d_no_loan_specific_feature",
    "linear set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_2d_no_loan_specific_feature",
    "fully connected set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_2d_no_loan_specific_feature",
    "multi-head attention set, pool 10, 40 epochs, 10000 sequences, 1 batch size, 10 pool size, supereasy_2d_no_loan_specific_feature",
   ]

def run_compare_set_sequence_to_baseline():
    for i in range(len(commands_compare_set_sequence_baseline)):
        run_command(commands_compare_set_sequence_baseline[i], loss_upper_bound_compare_set_sequence_baseline[i],test_name_compare_set_sequence_baseline[i])

run_compare_set_sequence_to_baseline()
