#! /usr/bin/env python3

# This script takes piano roll image files and converts them to MIDI
# It was intended for the P909 dataset
# It will create a directory of images for each MIDI file, where each image is a frame of the MIDI file

import os
import sys
import pretty_midi
import argparse
from multiprocessing import Pool, cpu_count
from tqdm import tqdm
from control_toys.data import fast_scandir
from functools import partial
from control_toys.pianoroll import img_file_2_midi_file

if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__,
                            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # use --no-onsets to disable requirment of onsets
    p.add_argument('--diff', default='', help='diff against this background image, new notes go in instrument 2')
    p.add_argument('--onsets', action=argparse.BooleanOptionalAction, default=True)
    p.add_argument('--separators', default=0, type=int, help='draw separators every this many pixels. (0=none)')
    p.add_argument("img_dirs", nargs='+', help="directories containing image files")
    p.add_argument("output_dir", help="output directory")
    args = p.parse_args()
    #print('args = ',args)

    img_dirs = args.img_dirs
    output_dir = args.output_dir


    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    img_files = []
    if os.path.isfile(img_dirs[0]):
        img_files = img_dirs
    else:
        for imgdir in img_dirs:
            i_subdirs, imfile = fast_scandir(imgdir, ['png', 'jpg', 'jpeg'])
            if imfile != []: img_files = img_files + imfile
    if len(img_files) == 0:
        print(f"No image files found in {img_dirs}")
        sys.exit(1)

    process_one = partial(img_file_2_midi_file, output_dir=output_dir, require_onsets=args.onsets, 
                          separators=args.separators, diff_img_file=args.diff)
    cpus = cpu_count()
    with Pool(cpus) as p:
        list(tqdm(p.imap(process_one, img_files), total=len(img_files), desc='Processing image files'))