import argparse

import os

import shutil

from pathlib import Path

from typing import Tuple

from tqdm import tqdm

TOTAL_SUBJECTS = 10307

def sanitize(name: str) -> Tuple[str, str]:

    return name.split('_')

def rearrange(input_path: Path, output_path: Path) -> None:

    os.makedirs(output_path, exist_ok=True)

    progress = tqdm(total=TOTAL_SUBJECTS)

    for folder in input_path.iterdir():

        subject = folder.name

        for sid in folder.iterdir():

            view, seq = sanitize(sid.name)

            src = os.path.join(input_path, subject,sid.name)

            dst = os.path.join(output_path, subject, seq, view)

            os.makedirs(dst, exist_ok=True)

            for subfile in os.listdir(src):

                if subfile not in os.listdir(dst) and subfile.endswith('.json'):

                    os.symlink(os.path.join(src, subfile),

                               os.path.join(dst, subfile))

        progress.update(1)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='OUMVLP rearrange tool')

    parser.add_argument('-i', '--input_path', required=True, type=str,

                        help='Root path of raw dataset.')

    parser.add_argument('-o', '--output_path', default='OUMVLP_rearranged', type=str,

                        help='Root path for output.')

    args = parser.parse_args()

    input_path = Path(args.input_path).resolve()

    output_path = Path(args.output_path).resolve()

    rearrange(input_path, output_path)