import argparse
import os
from pathlib import Path
import tqdm
import cv2
import tarfile
import zipfile
from functools import partial
import numpy as np
import pickle
import multiprocessing as mp


def make_pkl_for_one_person(id_, output_path, img_size=64):
    if id_.split(".")[-1] != "tar" or not os.path.exists(os.path.join(output_path, id_)):
        return
    with tarfile.TarFile(os.path.join(output_path, id_)) as f:
        f.extractall(output_path)
    os.remove(os.path.join(output_path, id_))
    id_path = id_.split(".")[0]
    input_path = os.path.join(output_path, "forTrain", id_path)
    base_pkl_path = os.path.join(output_path, "opengait", id_path)
    if not os.path.isdir(input_path):
        print("Path not found: "+input_path)
    for height in sorted(os.listdir(input_path)):
        height_path = os.path.join(input_path, height)
        for scene in sorted(os.listdir(height_path)):
            scene_path = os.path.join(height_path, scene)
            for type_ in sorted(os.listdir(scene_path)):
                type_path = os.path.join(scene_path, type_)
                for view in sorted(os.listdir(type_path)):
                    view_path = os.path.join(type_path, view)
                    for num in sorted(os.listdir(view_path)):
                        num_path = os.path.join(view_path, num)
                        imgs = []
                        for file_ in sorted(os.listdir(num_path)):
                            img = cv2.imread(os.path.join(
                                num_path, file_), cv2.IMREAD_GRAYSCALE)
                            if img_size != img.shape[0]:
                                img = cv2.resize(
                                    img, dsize=(img_size, img_size))
                            imgs.append(img)
                        if len(imgs) > 5:
                            pkl_path = os.path.join(
                                base_pkl_path, f"{height}-{scene}-{type_}-{num}", view)
                            os.makedirs(pkl_path, exist_ok=True)
                            pickle.dump(np.asarray(imgs), open(
                                os.path.join(pkl_path, f"{view}.pkl"), "wb"))
                        else:
                            print("No enough imgs: "+num_path)


def extractall(base_path: Path, output_path: Path, workers=1, img_size=64) -> None:
    """Extract all archives in base_path to output_path.

    Args:
        base_path (Path): Path to the directory containing the archives.
        output_path (Path): Path to the directory to extract the archives to.
    """

    os.makedirs(output_path, exist_ok=True)
    print("Unzipping train set...")
    with open(os.path.join(base_path, 'train001-500.zip'), 'rb') as f:
        z = zipfile.ZipFile(f)
        z.extractall(output_path)
    print("Unzipping validation set...")
    with open(os.path.join(base_path, 'val501-614.zip'), 'rb') as f:
        z = zipfile.ZipFile(f)
        z.extractall(output_path)
    print("Unzipping test set...")
    with open(os.path.join(base_path, 'test615-1014.zip'), 'rb') as f:
        z = zipfile.ZipFile(f)
        z.extractall(output_path)
    print("Extracting tar file...")
    os.makedirs(os.path.join(output_path,"forTrain"))
    os.makedirs(os.path.join(output_path,"opengait"))
    ids = os.listdir(os.path.join(output_path))
    progress = tqdm.tqdm(total=len(ids), desc='Pretreating', unit='id')

    with mp.Pool(workers) as pool:
        for _ in pool.imap_unordered(partial(make_pkl_for_one_person, output_path=output_path, img_size=img_size), ids):
            progress.update(1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='CASIA-E extractor')
    parser.add_argument('-b', '--input_path', type=str,
                        required=True, help='Base path to CASIA-E zip files')
    parser.add_argument('-o', '--output_path', type=str,
                        required=True, help='Output path for extracted files. The pickle files are generated in ${output_path}/opengait/')
    parser.add_argument('-s', '--img_size', default=64,
                        type=int, help='Image resizing size. Default 64')
    parser.add_argument('-n', '--num_workers',
                        type=int, default=1, help='Number of workers')
    args = parser.parse_args()

    extractall(Path(args.input_path), Path(args.output_path),
               args.num_workers, args.img_size)