################################################################################
# Name:                 Ingestion Program
# Author:             Zhengying Liu, Isabelle Guyon, Adrien Pavao, Zhen Xu
# Update time:    13 Aug 2019
# Usage: python ingestion.py --dataset_dir=<dataset_dir> --output_dir=<prediction_dir> --ingestion_program_dir=<ingestion_program_dir> --code_dir=<code_dir> --score_dir=<score_dir>

# AS A PARTICIPANT, DO NOT MODIFY THIS CODE.

VERSION = "v20191204"
DESCRIPTION = """This is the "ingestion program" written by the organizers. It takes the
code written by participants (with `workflow.py`) and one dataset as input,
run the code on the dataset and produce predictions on test set. For more
information on the code/directory structure, please see comments in this
code (ingestion.py) and the README file of the starting kit.
Previous updates:
20191204: [ZY] Add timer and separate workflow initialization from train/predict
                             process, : now workflow initilization doesn't consume time budget
                             quota (but can only use 20min)
20190820: [ZY] Mark the beginning of ingestion right before workflow.py to reduce
                             variance
20190708: [ZY] Integrate Julien's parallel data loader
20190516: [ZY] Change time budget to 20 minutes.
20190508: [ZY] Add time_budget to 'start.txt'
20190507: [ZY] Write timestamps to 'start.txt'
20190505: [ZY] Use argparse to parse directories AND time budget;
                             Rename input_dir to dataset_dir;
                             Rename submission_dir to code_dir;
20190504: [ZY] Check if workflow.py has attribute done_training and use it to
                             determinate whether ingestion has ended;
                             Use module-specific logger instead of logging (with root logger);
                             At beginning, write start.txt with ingestion_pid and start_time;
                             In the end, write end.txt with end_time and ingestion_success;
20190429: [ZY] Remove useless code block; better code layout.
20190425: [ZY] Check prediction shape.
20190424: [ZY] Use logging instead of logger; remove start.txt checking;
20190419: [ZY] Try-except clause for training process;
                    always terminates successfully.
"""
# The dataset directory dataset_dir (e.g. AutoDL_sample_data/) contains one dataset
# folder (e.g. adult.data/) with the training set (train/)    and test set (test/),
# each containing an some tfrecords data with a `metadata.textproto` file of
# metadata on the dataset. So one AutoDL dataset will look like
#
#     adult.data
#     ├── test
#     │   ├── metadata.textproto
#     │   └── sample-adult-test.tfrecord
#     └── train
#             ├── metadata.textproto
#             └── sample-adult-train.tfrecord
#
# The output directory output_dir (e.g. AutoDL_sample_result_submission/)
# will first have a start.txt file written by ingestion then receive
# all predictions made during the whole train/predict process
# (thus this directory is updated when a new prediction is made):
# 	adult.predict_0
# 	adult.predict_1
# 	adult.predict_2
#                ...
# after ingestion has finished, a file end.txt will be written, containing
# info on the duration ingestion used. This file is also used as a signal
# for scoring program showing that ingestion has terminated.
#
# The code directory submission_program_dir (e.g. AutoDL_sample_code_submission/)
# should contain your code submission workflow.py (and possibly other functions
# it depends upon).
#
# We implemented several classes:
# 1) DATA LOADING:
#        ------------
# dataset.py
# dataset.AutoDLMetadata: Read metadata in metadata.textproto
# dataset.AutoDLDataset: Read data and give tf.data.Dataset
# 2) LEARNING MACHINE:
#        ----------------
# workflow.py
# workflow.workflow.train
# workflow.workflow.test
#
# ALL INFORMATION, SOFTWARE, DOCUMENTATION, AND DATA ARE PROVIDED "AS-IS".
# UNIVERSITE PARIS SUD, CHALEARN, AND/OR OTHER ORGANIZERS OR CODE AUTHORS DISCLAIM
# ANY EXPRESSED OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE, AND THE
# WARRANTY OF NON-INFRIGEMENT OF ANY THIRD PARTY'S INTELLECTUAL PROPERTY RIGHTS.
# IN NO EVENT SHALL UNIVERSITE PARIS SUD AND/OR OTHER ORGANIZERS BE LIABLE FOR ANY SPECIAL,
# INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF SOFTWARE, DOCUMENTS, MATERIALS,
# PUBLICATIONS, OR INFORMATION MADE AVAILABLE FOR THE CHALLENGE.
#
# Main contributors: Isabelle Guyon and Zhengying Liu

# =========================== BEGIN OPTIONS ==============================

# Verbosity level of logging:
##############
# Can be: NOTSET, DEBUG, INFO, WARNING, ERROR, CRITICAL
verbosity_level = "INFO"

import warnings

warnings.filterwarnings("ignore", category=FutureWarning)

# Some common useful packages
from contextlib import contextmanager
from os import getcwd as pwd
from os.path import join
from sys import argv, path
import argparse
import datetime
import glob
import logging
import math
import numpy as np
import os
import sys
import signal
import time
import tensorflow as tf


from ...task import Task
from ...data import DataBundle
from ...data import data_io
from ..logging import get_logger

logger = get_logger(verbosity_level)


def _HERE(*args):
    """Helper function for getting the current directory of this script."""
    h = os.path.dirname(os.path.realpath(__file__))
    return os.path.abspath(os.path.join(h, *args))


def write_start_file(output_dir, start_time=None, time_budget=None, task_name=None):
    """Create start file 'start.txt' in `output_dir` with ingestion's pid and
    start time.

    The content of this file will be similar to:
            ingestion_pid: 1
            task_name: beatriz
            time_budget: 7200
            start_time: 1557923830.3012087
            0: 1557923854.504741
            1: 1557923860.091236
            2: 1557923865.9630117
            3: 1557923872.3627956
            <more timestamps of predictions>
    """
    ingestion_pid = os.getpid()
    start_filename = "start.txt"
    start_filepath = os.path.join(output_dir, start_filename)
    with open(start_filepath, "w") as f:
        f.write("ingestion_pid: {}\n".format(ingestion_pid))
        f.write("task_name: {}\n".format(task_name))
        f.write("time_budget: {}\n".format(time_budget))
        f.write("start_time: {}\n".format(start_time))
    logger.debug("Finished writing 'start.txt' file.")


def write_timestamp(output_dir, predict_idx, timestamp):
    start_filename = "start.txt"
    start_filepath = os.path.join(output_dir, start_filename)
    with open(start_filepath, "a") as f:
        f.write("{}: {}\n".format(predict_idx, timestamp))
    logger.debug(
        "Wrote timestamp {} to 'start.txt' for predition {}.".format(
            timestamp, predict_idx
        )
    )


class WorkflowApiError(Exception):
    pass


class TimeoutException(Exception):
    pass


class Timer:
    def __init__(self):
        self.duration = 0
        self.total = None
        self.remain = None
        self.exec = None

    def set(self, time_budget):
        self.total = time_budget
        self.remain = time_budget
        self.exec = 0

    @contextmanager
    def time_limit(self, pname):
        def signal_handler(signum, frame):
            raise TimeoutException("Timed out!")

        signal.signal(signal.SIGALRM, signal_handler)
        signal.alarm(int(math.ceil(self.remain)))
        start_time = time.time()

        try:
            yield
        finally:
            exec_time = time.time() - start_time
            signal.alarm(0)
            self.exec += exec_time
            self.duration += exec_time
            self.remain = self.total - self.exec

            logger.info("{} success, time spent so far {} sec".format(pname, self.exec))

            if self.remain <= 0:
                raise TimeoutException("Timed out for the process: {}!".format(pname))


# =========================== BEGIN PROGRAM ================================

if __name__ == "__main__":

    #### Check whether everything went well
    ingestion_success = True

    # Parse directories from input arguments
    root_dir = _HERE(os.pardir)
    default_dataset_dir = join(root_dir, "AutoDL_sample_data")
    default_output_dir = join(root_dir, "AutoDL_sample_result_submission")
    default_ingestion_program_dir = join(root_dir, "AutoDL_ingestion_program")
    default_code_dir = join(root_dir, "AutoDL_sample_code_submission")
    default_score_dir = join(root_dir, "AutoDL_scoring_output")
    default_time_budget = 1200
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset_dir",
        type=str,
        default=default_dataset_dir,
        help="Directory storing the dataset (containing " + "e.g. adult.data/)",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=default_output_dir,
        help="Directory storing the predictions. It will "
        + "contain e.g. [start.txt, adult.predict_0, "
        + "adult.predict_1, ..., end.txt] when ingestion "
        + "terminates.",
    )
    parser.add_argument(
        "--ingestion_program_dir",
        type=str,
        default=default_ingestion_program_dir,
        help="Directory storing the ingestion program "
        + "`ingestion.py` and other necessary packages.",
    )
    parser.add_argument(
        "--code_dir",
        type=str,
        default=default_code_dir,
        help="Directory storing the submission code "
        + "`workflow.py` and other necessary packages.",
    )
    parser.add_argument(
        "--score_dir",
        type=str,
        default=default_score_dir,
        help="Directory storing the scoring output "
        + "e.g. `scores.txt` and `detailed_results.html`.",
    )
    parser.add_argument(
        "--time_budget",
        type=float,
        default=default_time_budget,
        help="Time budget for running ingestion program.",
    )
    args = parser.parse_args()
    logger.debug("Parsed args are: " + str(args))
    logger.debug("-" * 50)
    dataset_dir = args.dataset_dir
    output_dir = args.output_dir
    ingestion_program_dir = args.ingestion_program_dir
    code_dir = args.code_dir
    score_dir = args.score_dir
    time_budget = args.time_budget
    if dataset_dir.endswith("run/input") and code_dir.endswith("run/program"):
        logger.debug(
            "Since dataset_dir ends with 'run/input' and code_dir "
            "ends with 'run/program', suppose running on "
            + "CodaLab platform. Modify dataset_dir to 'run/input_data' "
            "and code_dir to 'run/submission'. "
            + "Directory parsing should be more flexible in the code of "
            + "compute worker: we need explicit directories for "
            + "dataset_dir and code_dir."
        )
        dataset_dir = dataset_dir.replace("run/input", "run/input_data")
        code_dir = code_dir.replace("run/program", "run/submission")

    # Show directories for debugging
    logger.debug("sys.argv = " + str(sys.argv))
    logger.debug("Using dataset_dir: " + dataset_dir)
    logger.debug("Using output_dir: " + output_dir)
    logger.debug("Using ingestion_program_dir: " + ingestion_program_dir)
    logger.debug("Using code_dir: " + code_dir)

    # Our libraries
    path.append(ingestion_program_dir)
    path.append(code_dir)
    # IG: to allow submitting the starting kit as sample submission
    path.append(code_dir + "/AutoDL_sample_code_submission")

    data_io.mkdir(output_dir)

    #### INVENTORY DATA (and sort dataset names alphabetically)
    datanames = data_io.inventory_data(dataset_dir)
    #### Delete zip files and metadata file
    datanames = [x for x in datanames if x.endswith(".data")]

    if len(datanames) != 1:
        raise ValueError(
            "{} datasets found in dataset_dir={}!\n".format(len(datanames), dataset_dir)
            + "Please put only ONE dataset under dataset_dir."
        )

    basename = datanames[0]

    logger.info("************************************************")
    logger.info("******** Processing dataset " + basename[:-5].capitalize() + " ********")
    logger.info("************************************************")
    logger.debug("Version: {}. Description: {}".format(VERSION, DESCRIPTION))

    ##### Begin creating training set and test set #####
    logger.info("Reading training set and test set...")
    data_bundle = DataBundle(dataset_dir=dataset_dir)
    ##### End creating training set and test set #####

    # Creating Task
    task = Task(data_bundle=data_bundle, metric=None, resources=None)

    # 20 min for participants to initializing and install other packages
    try:
        init_time_budget = 20 * 60  # time budget for initilization.
        timer = Timer()
        timer.set(init_time_budget)
        with timer.time_limit("Initialization"):
            ##### Begin creating workflow #####
            logger.info("Creating workflow...this process should not exceed 20min.")
            from workflow import Workflow  # in participants' workflow.py

            # The metadata of D_train and D_test only differ in sample_count
            workflow = Workflow(task.get_train_databundle().get_metadata())
            ###### End creating workflow ######
    except TimeoutException as e:
        logger.info(
            "[-] Initialization phase exceeded time budget. Move to train/predict phase"
        )
    except Exception as e:
        logger.error("Failed to initializing workflow.")
        logger.error("Encountered exception:\n" + str(e), exc_info=True)

    # Mark starting time of ingestion
    start = time.time()
    logger.info(
        "=" * 5
        + " Start core part of ingestion program. "
        + "Version: {} ".format(VERSION)
        + "=" * 5
    )

    write_start_file(
        output_dir,
        start_time=start,
        time_budget=time_budget,
        task_name=basename.split(".")[0],
    )

    try:
        # Check if the workflow has methods `train` and `test`.
        for attr in ["train", "test"]:
            if not hasattr(workflow, attr):
                raise WorkflowApiError(
                    "Your workflow object doesn't have the method "
                    + "`{}`. Please implement it in workflow.py."
                )

        # Check if workflow.py uses new done_training API instead of marking
        # stopping by returning None
        use_done_training_api = hasattr(workflow, "done_training")
        if not use_done_training_api:
            logger.warning(
                "Your workflow object doesn't have an attribute "
                + "`done_training`. But this is necessary for ingestion "
                + "program to know whether the workflow has done training "
                + "and to decide whether to proceed more training. "
                + "Please add this attribute to your workflow."
            )

        workflow.run(task, start, output_dir, basename)

    except Exception as e:
        ingestion_success = False
        logger.info("Failed to run ingestion.")
        logger.error("Encountered exception:\n" + str(e), exc_info=True)

    # Finishing ingestion program
    end_time = time.time()
    overall_time_spent = end_time - start

    # Write overall_time_spent to a end.txt file
    end_filename = "end.txt"
    with open(os.path.join(output_dir, end_filename), "w") as f:
        f.write("ingestion_duration: " + str(overall_time_spent) + "\n")
        f.write("ingestion_success: " + str(int(ingestion_success)) + "\n")
        f.write("end_time: " + str(end_time) + "\n")
        logger.info(
            "Wrote the file {} marking the end of ingestion.".format(end_filename)
        )
        if ingestion_success:
            logger.info("[+] Done. Ingestion program successfully terminated.")
            logger.info("[+] Overall time spent %5.2f sec " % overall_time_spent)
        else:
            logger.info("[-] Done, but encountered some errors during ingestion.")
            logger.info("[-] Overall time spent %5.2f sec " % overall_time_spent)

    # Copy all files in output_dir to score_dir
    logger.info(f"Copy {output_dir} to {score_dir}")
    # os.system("cp -R {} {}".format(os.path.join(output_dir, "*"), score_dir))
    os.system("cp -R {} {}".format(output_dir, score_dir))
    logger.debug("Copied all ingestion output to scoring output directory.")

    logger.info("[Ingestion terminated]")
