'''Training script generator for subject-adaptive classification.

'''

from os import makedirs

from os.path import join as pjoin

with open(pjoin('tasks_adapt', 'script.sh.template'), 'r') as f:

    sh_template = f.read()

with open(pjoin('tasks_adapt', 'train.pbs.template'), 'r') as f:

    pbs_template = f.read()

def list_to_str(l):

    return ' '.join(map(str, l))

command_template = "python train_adapt.py $datapath $modelpath ./result_adapt{0}/r{1} -scheme {0} -trfrate {1} -gpu {2} > $logpath/stdout.r{1}.out &"

for scheme in range(1, 6):

    tasks_dir = "tasks_adapt{}".format(scheme)

    result_dir = "result_adapt{}".format(scheme)

    makedirs(tasks_dir, exist_ok=True)

    for rate in range(10, 101, 10):

        makedirs(pjoin(result_dir, "r{}".format(rate)), exist_ok=True)

    commands = []

    for ind, rate in enumerate(range(10, 81, 10)):

        commands.append(command_template.format(scheme, rate, ind))

    commands.append('wait')

    content = sh_template.format(scheme) + '\n'.join(commands)

    ncpu = 40

    ngpu = 8

    with open(pjoin(tasks_dir, 'script.' + str(0) + '.sh'), 'w') as f:

        f.write(content)

    content = pbs_template.format(0, scheme, ncpu, ngpu)

    with open(pjoin(tasks_dir, 'train.' + str(0) + '.pbs'), 'w') as f:

        f.write(content)

tasks_dir = "tasks_adaptr"

command_template = "python train_adapt.py $datapath $modelpath ./result_adapt{0}/r{1} -scheme {0} -trfrate {1} -gpu {2} > $logpath/stdout.s{0}.r{1}.out &"

makedirs(tasks_dir, exist_ok=True)

commands = []

for scheme in range(1, 5):

    for ind, rate in enumerate(range(90, 101, 10)):

        igpu = (scheme - 1) * 2 + ind

        commands.append(command_template.format(scheme, rate, igpu))

commands.append('wait')

content = sh_template.format('r') + '\n'.join(commands)

ncpu = 40

ngpu = 8

with open(pjoin(tasks_dir, 'script.' + str(0) + '.sh'), 'w') as f:

    f.write(content)

content = pbs_template.format(0, 'r', ncpu, ngpu)

with open(pjoin(tasks_dir, 'train.' + str(0) + '.pbs'), 'w') as f:

    f.write(content)

commands = []

scheme = 5

for ind, rate in enumerate(range(90, 101, 10)):

    commands.append(command_template.format(scheme, rate, ind))

commands.append('wait')

content = sh_template.format('r') + '\n'.join(commands)

ncpu = 10

ngpu = 2

with open(pjoin(tasks_dir, 'script.' + str(1) + '.sh'), 'w') as f:

    f.write(content)

content = pbs_template.format(1, 'r', ncpu, ngpu)

with open(pjoin(tasks_dir, 'train.' + str(1) + '.pbs'), 'w') as f:

    f.write(content)