import os
import click
import json
from tqdm import tqdm


@click.command()
@click.option("--swim_dir", type=str, default="datasets/swim_data")
@click.option("--output_dir", type=str, default="datasets/swim_data_cyclegan")
@click.option("--type", type=str, help="fog|rain|snow|night", required=True)
@click.option("--no_night", is_flag=True)
def build_cyclegan_dataset(swim_dir: str, output_dir: str, type: str, no_night: bool):
    # build the dataset with format
    # swim_data_cyclegan
    # ├── trainA
    # │   ├── 000000.jpg
    # │   ├── ...
    # ├── trainB
    # │   ├── 000000.jpg
    # │   ├── ...
    # |── testA
    # │   ├── 000000.jpg
    # │   ├── ...
    # |── testB
    # │   ├── 000000.jpg
    # │   ├── ...
    #
    # note:
    # trainA and trainB are extracted from swim_data train set
    # testA and testB are extracted from swim_data val set
    # A - clear
    # B - fog|rain|snow|night

    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, "trainA"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "trainB"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "testA"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "testB"), exist_ok=True)

    with open(os.path.join(swim_dir, "train", "labels.json"), "r") as f:
        train_labels = json.load(f)

    with open(os.path.join(swim_dir, "val", "labels.json"), "r") as f:
        val_labels = json.load(f)

    if type != "night":
        for label in tqdm(train_labels, desc="train"):
            if no_night and label["timeofday"] == "night":
                continue

            if label["weather"] == type:
                os.system(
                    f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainB', label['name'])}"
                )
            elif label["weather"] == "clear":
                os.system(
                    f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainA', label['name'])}"
                )

        for label in tqdm(val_labels, desc="val"):
            if no_night and label["timeofday"] == "night":
                continue

            if label["weather"] == type:
                os.system(
                    f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testB', label['name'])}"
                )
            elif label["weather"] == "clear":
                os.system(
                    f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testA', label['name'])}"
                )
    else:
        for label in tqdm(train_labels, desc="train"):
            if label["weather"] != "clear":
                continue

            if label["timeofday"] == "night":
                os.system(
                    f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainB', label['name'])}"
                )
            elif label["timeofday"] == "daytime":
                os.system(
                    f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainA', label['name'])}"
                )

        for label in tqdm(val_labels, desc="val"):
            if label["weather"] != "clear":
                continue

            if label["timeofday"] == "night":
                os.system(
                    f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testB', label['name'])}"
                )
            elif label["timeofday"] == "daytime":
                os.system(
                    f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testA', label['name'])}"
                )


if __name__ == "__main__":
    build_cyclegan_dataset()