from pathlib import Path
from argparse import ArgumentParser
from traceback import print_exc
from concurrent.futures import ThreadPoolExecutor
import os

from arip2 import arip_main

def process_file(in_fp, output, ref, dx, disable_print):
    try:
        if not disable_print:
            print(f'>> process: {in_fp}, ref: {ref}')
        arip_main(in_fp, output, ref, dx, disable_print)
    except SystemExit:
        print(f'Error processing file: {in_fp}')
    except:
        print_exc()

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('-R', '--ref', default='data/s5000.xyz', type=Path, help='*.xyz coord file')
    parser.add_argument('input', type=Path, help='input *.pdb file or folder path')
    parser.add_argument('-O', '--output', default='out', type=Path, help='output folder path')
    parser.add_argument('--dx', default=0.2, type=float, help='sample delta')
    parser.add_argument('--h', action='store_true', help='use 15092 dots and 0.1 interval as surface and volume')
    args = parser.parse_args()

    if args.h:
        args.ref = 'data/s15092.xyz'
        args.dx = 0.1

    input: Path = args.input
    assert input.is_file() or input.is_dir(), f'{input} is not a valid file or directory'

    in_fps = [input] if input.is_file() else list(input.iterdir())
    disable_print = False if input.is_file() else True

    with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
        executor.map(process_file, in_fps, [args.output]*len(in_fps), [args.ref]*len(in_fps), [args.dx]*len(in_fps), [disable_print]*len(in_fps))