from data_stat import Cpu, Format

import argparse
import glob
import os
import subprocess
import random

parser = argparse.ArgumentParser()
parser.add_argument('cpu', choices=[x.name.lower() for x in Cpu])
parser.add_argument('output_dir')
parser.add_argument('matrix_dir')
parser.add_argument('format', type=str,
        choices=[fmt.name.lower() for fmt in Format])
parser.add_argument('base_iterations', type=int)
parser.add_argument('min_time_s', type=int)
parser.add_argument('baseline_time_s', type=int)
parser.add_argument('baseline_delay_s', type=int)
parser.add_argument('--perf', action='store_const', const='--perf')
parser.add_argument('--power', action='store_const', const='--power')
parser.add_argument('--distribute', action='store_true')
args = parser.parse_args()
args.cpu = Cpu[args.cpu.upper()]
args.format = Format[args.format.upper()]

srun_args = {
        Cpu.ALTRA: [
            '--account', 'oasis',
            '--partition', 'oasis',
            '--qos', 'oasis-exempt',
            #'--cpus-per-task 160',
            '--cpus-per-task', '160',
            #'--mem 28114',
            '--mem', '16G',
            '--ntasks-per-node', '1'#,
            #'--exclusive',
            #'--output', '/dev/null',
            #'--error', '/dev/null'
        ],
        Cpu.EPYC_7313P: [
            '--account', 'nexus',
            '--partition', 'tron',
            '--qos', 'high',
            '--cpus-per-task', '16',
            '--ntasks-per-node', '1',
            '--prefer', 'EPYC-7313P'
        ],
        Cpu.XEON_4216: [
            '--account', 'nexus',
            '--partition', 'tron',
            '--qos', 'tron-exempt',
            '--cpus-per-task', '32',
            '--ntasks-per-node', '1',
            '--prefer', 'Xeon,4216'
        ]
}
python = {
        Cpu.ALTRA: 'python3',
        Cpu.EPYC_7313P: 'python3.11',
        Cpu.XEON_4216: 'python3.11'
}

def srun(srun_args_list: list, run_args, matrix_file: str) -> list:
    run_args_list = [
            args.cpu.name.lower(),
            matrix_file,
            args.format.name.lower(),
            str(args.base_iterations),
            str(args.min_time_s),
            str(args.baseline_time_s),
            str(args.baseline_delay_s)]
    if args.perf is not None:
        run_args_list += [args.perf]
    if args.power is not None:
        run_args_list += [args.power]
    return ['srun'] + srun_args_list + [python[args.cpu], 'run.py'] + run_args_list

processes = list()

for i, matrix in enumerate(glob.glob(f'{args.matrix_dir.rstrip("/")}/*.mtx')):
    if args.distribute:
        if args.cpu == Cpu.ALTRA:
            i = i % 40
            srun_args_temp = srun_args[args.cpu] + ['--nodelist', f'oasis{i:02}']
        elif args.cpu == Cpu.EPYC_7313P:
            srun_args_temp = srun_args[args.cpu]
    else:
        srun_args_temp = srun_args[args.cpu]

    output_filename = '_'.join([
            args.cpu.name.lower(),
            str(args.min_time_s),
            str(args.baseline_time_s),
            str(args.baseline_delay_s),
            os.path.splitext(os.path.basename(matrix))[0]])

    json_filepath = f'{args.output_dir.rstrip("/")}/{output_filename}.json'
    raw_filepath = f'{args.output_dir.rstrip("/")}/{output_filename}.output'
    with open(json_filepath, 'w') as json_file, open(raw_filepath, 'w') as raw_file:
        print(srun(srun_args_temp, args, matrix))
        print(json_filepath)
        print(raw_filepath)

        processes.append(subprocess.Popen(
                srun(srun_args_temp, args, matrix),
                stdout=json_file,
                stderr=raw_file))

    # Wait on every 10 jobs to avoid socket timeout.
    if i % 10 == 9:
        print("Waiting on 10 jobs")
        for process in processes:
            process.wait()

        processes = list()

for process in processes:
    process.wait()