"""
Nests images named in the Labeled Face in the Wild format into subdirectories with the
identity's name as the name of the subdirectory.
"""

import argparse
import os
import shutil

from pathlib import Path

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--source_dir", dest="source_dir", required=True, 
                        type=str, help="Source directory of images in LFW name format")
    parser.add_argument("-o", "--output_dir", dest="output_dir", required=True, 
                        type=str, help="Output directory to put LFW formatted images " 
                                       "in subdirectories of the name of the identities")
    args = parser.parse_args()
    
    source_dir = Path(args.source_dir)
    output_dir = Path(args.output_dir)
    
    for filepath in source_dir.iterdir():
        filename = os.path.basename(filepath)
        identity_dir = filename.rsplit("_", 1)[0]
        (output_dir / identity_dir).mkdir(parents=True, exist_ok=True)
        shutil.move(filepath, output_dir / identity_dir / filename)
