#!/usr/bin/env python
"""
A script to split Food-101 dataset.

Usage; food-101.split --base_dir <Food 101 dataset directory>

__author__ = "Hide Inada"
__copyright__ = "Copyright 2019, Hide Inada"
__license__ = "The MIT License"
__email__ = "hideyuki@gmail.com"
"""

import os
import logging
import argparse
from pathlib import Path

log = logging.getLogger(__name__)
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))  # Change the 2nd arg to INFO to suppress debug logging


def split_dataset(base_dir_path, image_dir_path):
    """

    Parameters
    ----------
    base_dir_path: pathlib.Path
        Path to the directory where Food-101 dataset was extracted.
    image_dir_path: pathlib.Path
        Path to the Images directory containing all the images.
    """

    # Create validation and test directories
    validation_dir_path = base_dir_path / Path("validation_images")
    if validation_dir_path.exists() is False:
        validation_dir_path.mkdir(exist_ok=True, parents=True)

    test_dir_path = base_dir_path / Path("test_images")
    if test_dir_path.exists() is False:
        test_dir_path.mkdir(exist_ok=True, parents=True)

    # Open the test set list
    validation_category_dir_path = None
    test_category_dir_path = None
    test_set_list_path = base_dir_path / Path("meta") / Path("test.txt")
    with open(test_set_list_path, "r") as f:
        prev_category = ""
        image_in_category = 0
        for line in f:
            line = line.strip()
            category, stem = line.split("/")

            if prev_category != category:
                image_in_category = 1
                validation_category_dir_path = validation_dir_path / Path(category)
                if validation_category_dir_path.exists() is False:
                    validation_category_dir_path.mkdir(exist_ok=True, parents=True)

                test_category_dir_path = test_dir_path / Path(category)
                if test_category_dir_path.exists() is False:
                    test_category_dir_path.mkdir(exist_ok=True, parents=True)

                prev_category = category

            else:
                image_in_category += 1

            src_image_path = image_dir_path / Path(category) / Path(stem + ".jpg")

            if image_in_category <= 100:
                target_image_path = validation_category_dir_path / Path(stem + ".jpg")
            else:
                target_image_path = test_category_dir_path / Path(stem + ".jpg")

            src_image_path.rename(target_image_path)
            log.info("Src: %s, Target: %s" % (src_image_path, target_image_path))


def main():
    """Defines an application's main functionality"""

    parser = argparse.ArgumentParser(description='Food-101 base directory')
    parser.add_argument('--base_dir',
                        type=str,
                        default=None,
                        help="Directory where Food-101 dataset was extracted.")

    args = parser.parse_args()
    base_dir = args.base_dir
    if base_dir is None:
        raise ValueError("Specify the Food-101 directory with the --base_dir option")

    base_dir_path = Path(base_dir)
    if base_dir_path.exists() is False:
        raise ValueError("base_dir %s not found" % (base_dir))

    image_dir_path = base_dir_path / Path("images")
    if image_dir_path.exists() is False:
        raise ValueError("Images directory not found under %s" % (base_dir))

    split_dataset(base_dir_path, image_dir_path)


if __name__ == "__main__":
    main()
