#!/usr/bin/env python3

import subprocess
import argparse
import shutil
import glob
import sys
import os
import re

from logzero import logger as log

#----------------------------
class data:
    jobid   = None
    nfits   = None
    job_dir = os.environ['JOBDIR']
    sandbox = f'{job_dir}/extractor'
    vers    = None
    upgrade = None
    sft_path= None
    sft_dir = '/publicfs/lhcb/user/campoverde/SFT'
    sft_name= 'RK_TOY'
#----------------------------
def get_args():
    parser = argparse.ArgumentParser(description='Used to run local tests of toy fits')
    parser.add_argument('-j', '--job_id' , type=int, help='Index of job'                         , required=True)
    parser.add_argument('-n', '--nfits'  , type=int, help='Number of fits per job'               , required=True)
    parser.add_argument('-v', '--vers'   , type=str, help='Version of output'                    , required=True)
    parser.add_argument('-s', '--sndbx'  , type=str, help='Directory for output, optional'       , default=data.sandbox)
    parser.add_argument('-u', '--upgrade', type=str, help='Will upgrade before running'          , default='0', choices=['0', '1'])
    parser.add_argument('-d', '--dirname', type=str, help='Name of directory with code, optional', default=data.sft_name)
    args = parser.parse_args()

    data.jobid   = args.job_id
    data.nfits   = args.nfits
    data.vers    = args.vers
    data.sandbox = f'{args.sndbx}/{str(data.jobid).rjust(3, "0")}_{str(data.nfits).rjust(3, "0")}'
    data.upgrade = args.upgrade
    data.sft_name= args.dirname
    data.sft_path= f'{data.sft_dir}/{data.sft_name}'

    try:
        os.makedirs(data.sandbox, exist_ok=True)
    except:
        log.error(f'Cannot make sandbox: {data.sandbox}')
        raise
#----------------------------
def print_args():
    log.info('-' * 40)
    log.info(f'Args for {os.path.basename(__file__)}:')
    log.info('-' * 40)
    log.info(f'{"Sandbox":<20}{data.sandbox:<20}')
    log.info(f'{"JOBID":<20}{data.jobid:<20}')
    log.info(f'{"NFits":<20}{data.nfits:<20}')
    log.info(f'{"CFG Version":<20}{data.vers:<20}')
    log.info(f'{"SFT path":<20}{data.sft_path:<20}')
    log.info('-' * 40)
#----------------------------
def copy_files():
    l_script = ['rxe_toys', 'rxe_run_toys']
    l_file   = l_script 
    for file_path in l_file:
        file_name = os.path.basename(file_path)
        shutil.copyfile(file_path, f'{data.sandbox}/{file_name}')

    for script in l_script:
        subprocess.run(['chmod', '+x', f'{data.sandbox}/{script}'] )
#----------------------------
def run():
    l_command = ['./rxe_run_toys', data.vers, data.upgrade]

    log.info(f'Running: {l_command}, with RK_TOY={data.sft_path}')
    stat = subprocess.run(l_command, env={'RK_TOY' : data.sft_path})
    if stat.returncode != 0:
        log.error(f'Process returned exit status: {stat.returncode}')
        raise
#----------------------------
def make_seeds():
    with open(f'{data.jobid}.sd', 'w') as ofile:
        for i_fit in range(data.nfits):
            seed = i_fit + 1000 * data.jobid
            ofile.write(f'{seed}\n')
#----------------------------
def main():
    get_args()
    print_args()
    copy_files()
    os.chdir(data.sandbox)
    make_seeds()
    run()
#----------------------------
if __name__ == '__main__':
    main()

