from pathlib import Path
import shutil
from loop_alphapose_training import run_cmd, Iteration
import logging
import coloredlogs

logger = logging.getLogger(__name__)
coloredlogs.install(level=logging.INFO)

def collate(iterations):
    path = Path(f'out/test_imgs_{iterations[1].nr}/vis')
    images = [x for x in path.iterdir() if x.is_file()]
    for img in images:
        logger.info(f"collate {img.name}")
        target_dir = Path(f'out/test_imgs_loops/{img.stem}/')
        target_dir.mkdir(parents=True, exist_ok=True)
        
        for iteration in iterations:
            try:
                src = Path(f'out/test_imgs_{iteration.nr}/vis/') / img.name
                target = target_dir / (iteration.nr + img.suffix)
                logger.info(f"\tcopy {src} to {target}")
                shutil.copy(src, target)
            except Exception as e:
                logger.exception(e)
        

# collate([Iteration(i) for i in range(2,7)])
# exit()

if __name__ == '__main__':

    i = 0
    iterations = []
    while True:
        i+=1
        iteration = Iteration(i)
        path = iteration.model_path()
        if not path.exists():
           logger.warning(f"Model for iteration {iteration.nr} doesn't exist")
           break
        iterations.append(iteration)
        cmd = [
            'python', 'scripts/demo_inference.py',
            '--cfg', str(iteration.training_config_path()),
            '--checkpoint', str(iteration.model_path()),
            '--gpus', '0',
            '--format', 'coco',
            '--indir', 'data/test_imgs',
            '--outdir', f'/out/test_imgs_{iteration.nr}',
            '--save_img'

        ]
        logger.info(f"Running {cmd}")
        run_cmd(cmd, in_docker=True)
        # break


    collate(iterations)