#!/usr/bin/env python3

def cli():
    import os,shutil,argparse
    cli = argparse.ArgumentParser(description='CLARA scaling test')
    cli.add_argument('-y','--yaml',   metavar='YAML',help='path to YAML file',required=True)
    cli.add_argument('-c','--clara',  metavar='DIR',help='CLARA_HOME path (default=$CLARA_HOME)',default=os.getenv('CLARA_HOME',None))
    cli.add_argument('-t','--threads',metavar='#',help='threads (default=8,16,24,34,44)',default='8,16,24,34,44')
    cli.add_argument('-e','--events', metavar='#',help='events per thread (default=555)',default=555,type=int)
    cli.add_argument('-N','--numa',   metavar='#',help='restrict to a NUMA socket (choices=[0,1])',choices=[0,1])
    cli.add_argument('-C','--cpus',   metavar='#',help='restrict to a CPU list (e.g. 0,1,2,3)')
    cli.add_argument('-T','--cputhr', metavar='#',help='restrict to same number of CPUs as threads')
    cli.add_argument('datafile', help='input EVIO/HIPO data file')
    cfg = cli.parse_args()
    cfg.threads = cfg.threads.split(',')
    if not cfg.clara:
        cli.error('cannot find CLARA installation via -c or $CLARA_HOME')
    elif shutil.which('run-clara'):
        cfg.run_clara = shutil.which('run-clara')
    elif os.path.exists(clara_home + '/plugins/clas12/bin/run-clara'):
        cfg.run_clara = clara_home + '/plugins/clas12/bin/run-clara'
    else:
        cli.error('cannot find run-clara in $PATH or $CLARA_HOME')
    if cfg.cpus:
        try:
            cfg.cpus = ','.join([str(int(i)) for i in cfg.cpus.split(',')])
        except Exception as e:
            cli.error('invalid --cpus argument:  {cfg.cpus}')
    elif cfg.numa or cfg.cputhr:
        if not cfg.numa:
            cfg.numa = 0
        if not shutil.which('numactl'):
            cli.error('numactl is not in $PATH, --numa option not supported')
        for x in run(['numactl','-H']):
            if x.startswith(f'node {cfg.numa} cpus:'):
                cfg.cpus = ','.join([str(int(i)) for i in x.strip().split()[3:]])
                break
        if not cfg.cpus:
            cli.error('error determining cpu list for -C or -T')
    return cfg

def run(cmd):
    import subprocess
    print('run >>> '+' '.join(cmd))
    p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,universal_newlines=True,encoding='latin-1')
    for line in iter(p.stdout.readline, ''):
        line = line.strip()
        if len(line) > 0:
            yield line
    p.wait()

def benchmark(cfg, threads, log):
    import collections
    cmd = []
    # use taskset for thread affinity:
    if cfg.cpus:
        if cfg.cputhr:
            cmd.extend(['taskset','-c',','.join(cfg.cpus.split(',')[:threads])])
        else:
            cmd.extend(['taskset','-c',cfg.cpus])
    # add the run-clara command:
    cmd.extend([cfg.run_clara,
            '-c',cfg.clara,
            '-n',str(cfg.events*int(threads)),
            '-t',str(threads),
            '-l',
            '-y',cfg.yaml,
            '-o',f'tmp-scaling-{threads}',
            cfg.datafile])
    exiting,benchmarks = False,collections.OrderedDict()
    for line in run(cmd):
        cols = line.split()
        print(line)
        log.write(line+'\n')
        try:
            if line.find('Benchmark results:') >= 0:
                exiting = True
            elif line.find('Processing is complete') >= 0:
                exiting = False
            elif len(cols) > 20:
                if line.find('Processed') >= 0:
                    benchmarks['Event'] = float(cols[12])
            elif exiting:
                # catch-all for services:
                if len(cols) > 14:
                    if 'Services' not in benchmarks:
                        benchmarks['Services'] = collections.OrderedDict()
                    benchmarks['Services'][cols[2]] = float(cols[14])
                # FIXME:  what are these, why don't they add up?
                elif line.find('Average processing time') >= 0:
                    benchmarks['Avg'] = float(cols[6])
                elif line.find('Total processing time') >= 0:
                    benchmarks['Total'] = float(cols[6])
                elif line.find('Total orchestrator time') >= 0:
                    benchmarks['Orch'] = float(cols[6])
        except ValueError:
            pass
    return benchmarks

def table(benchmarks):
    table = []
    header = [ 'Threads' ]
    b = benchmarks[0][1]
    header.extend([x for x in b if x != 'Services' and x != 'Event'])
    if 'Services' in b:
        header.extend(b['Services'].keys())
    table.append(header)
    for b in benchmarks:
        threads,benchmark = b[0],b[1]
        row = [threads]
        for k in ['Avg','Total','Orch','Services']:
            if k in benchmark:
                if k == 'Services':
                    row.extend(benchmark[k].values())
                else:
                    row.append(benchmark[k])
        table.append(row)
    return table

def show(benchmarks):
    for row in table(benchmarks):
        print(' '.join([str(x) for x in row])) 

def save(benchmarks):
    with open('scaling.txt','w') as f:
        for row in table(benchmarks):
            f.write(' '.join([str(x) for x in row])+'\n') 

def plot():
    import os,shutil,subprocess
    if shutil.which('gnuplot'):
        g = os.path.dirname(os.path.abspath(__file__))+'/scaling.gpl'
        r = subprocess.run(['gnuplot',g], capture_output=True, text=True)
        if r.returncode == 0:
            with open('scaling.svg','w') as f:
                f.write(r.stdout)
                print('Info:  created scaling.svg from scaling.txt and scaling.gpl!')
        else:
            print(r.stdout)
            print(r.stderr)
            print('Error:  gnuplot failed!')
    else:
        print('Warning:  gnuplot is not in $PATH.')

if __name__ == '__main__':
    cfg = cli()
    import os
    benchmarks = []
    for threads in cfg.threads:
        os.makedirs('tmp-scaling-'+threads)
        with open(f'tmp-scaling-{threads}/run-clara.log','w') as log:
            benchmarks.append([threads, benchmark(cfg, threads, log)])
    show(benchmarks)
    save(benchmarks)
    plot()

