import sys
import cv2
import os
import yaml
import numpy as np
from multiprocessing import Process
from tqdm import tqdm

def run_test(param_list, wid, skip):
    for cfg_file, test_folder, ckpt_file, out_dir in param_list:
        if os.path.exists(out_dir):
            if skip:
                print(f'{out_dir} exists, skip it.')
                continue
            else:
                print(f'{out_dir} exists, force to remove.')
                os.system(f'rm -rf {out_dir}')
        cmd = f'python3 test.py --cfg_file {cfg_file} --test_folder {test_folder} --ckpt {ckpt_file} --overwrite_output_dir {out_dir}'
        print(f'Working process {wid} runs: {cmd}')
        os.system(cmd)

def load_image(input_folder, output_folder, ckpt, fn, postfix, resize_to):
    if ckpt == 'blank':
        return np.zeros([resize_to[1], resize_to[0], 3], dtype=np.uint8)
    if ckpt == 'input':
        img = cv2.imread(os.path.join(input_folder, fn.replace(postfix, '')))
    else:
        img = cv2.imread(os.path.join(output_folder, ckpt, fn))
    if not (img.shape[0] == resize_to[1] and img.shape[1] == resize_to[0]):
        img = cv2.resize(img, resize_to)
    return cv2.putText(img, ckpt, (10, 30), cv2.FONT_HERSHEY_COMPLEX, 1.0, (0, 0, 255), 2)

if __name__ == '__main__':
    assert len(sys.argv) == 2, 'Usage: python3 test_multi_epochs.py [path to config].'

    # read config
    with open(sys.argv[1], 'rt') as f:
        config = yaml.safe_load(f)
    assert type(config) is dict, 'Loading config file failed.'

    # check config
    gpu_ids = config['gpu_ids']
    skip_existed = config['skip_existed']
    cfg_file = config['cfg_file']
    ckpt_folder = config['ckpt_folder']
    ckpt_layout = config['ckpt_layout']
    test_folder = config['test_folder']
    output_folder = config['output_folder']
    resize_to = config['resize_to']

    with open(cfg_file, 'rt') as f:
        temp_config = yaml.safe_load(f)
    assert type(temp_config) is dict, 'Loading config file failed.'

    postfix = temp_config['testing']['visual_names']
    assert len(postfix) == 1, 'Count of visual names should be 1.'
    postfix = '_' + postfix[0]

    temp_config['testing']['image_format'] = 'input'

    # run test
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    gpu_count = len(gpu_ids)
    param_lists = [ [] for _ in range(gpu_count) ]
    cnt = 0
    for line in ckpt_layout:
        assert len(line) == len(ckpt_layout[0]), 'Each row should have the same count of items.'
        for ckpt in line:
            if not ckpt in ['input', 'blank']:
                ckpt_path = os.path.join(ckpt_folder, ckpt + '.pth')
                assert os.path.exists(ckpt_path), f'Checkpoint not found: {ckpt_path}'
                new_cfg_file = os.path.join(output_folder, ckpt + '.yaml')
                out_dir = os.path.join(output_folder, ckpt)
                with open(new_cfg_file, 'wt') as f:
                    temp_config['common']['gpu_ids'] = [ gpu_ids[cnt % gpu_count] ]
                    yaml.safe_dump(temp_config, f)
                param_lists[cnt % gpu_count].append((new_cfg_file, test_folder, ckpt_path, out_dir))
                cnt += 1
    pool = []
    for i in range(gpu_count):
        p = Process(target=run_test, args=(param_lists[i], i, skip_existed))
        p.start()
        pool.append(p)
    for p in pool:
        p.join()

    # concat image
    print('Begin to concat.')
    out_dir = os.path.join(output_folder, 'concat')
    if os.path.exists(out_dir):
        print(f'{out_dir} exists, force to remove.')
        os.system(f'rm -rf {out_dir}')
    os.makedirs(out_dir)
    files = os.listdir(param_lists[0][0][3])
    for fn in tqdm(files):
        if not os.path.splitext(fn)[1].lower() in ['.jpg', '.jpeg', '.png']:
            continue
        try:
            result = np.vstack([
                np.hstack([
                    load_image(test_folder, output_folder, ckpt, fn, postfix, resize_to) for ckpt in line
                ]) for line in ckpt_layout
            ])
            cv2.imwrite(os.path.join(out_dir, fn.replace(postfix, '')), result)
        except Exception as e:
            print(fn, e)