#!/usr/bin/env python

"""
Script report_pnp
==========================
This script is used for counting PNPs and non-PNPs from data chunks.
"""

# standard
import warnings
from pathlib import Path
import logging
import re
import sys
from datetime import timedelta
from datetime import datetime
import pandas as pd
import argparse
from pandas import DataFrame
# chemoinformatics
import rdkit
# dev
import npfc
from npfc import load
from npfc import utils
import subprocess
# disable SettingWithCopyWarning warnings
pd.options.mode.chained_assignment = None  # default='warn'


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ FUNCTIONS ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

# todo:
# 1. iterate over each chunk
# 2. count PNPs, non-PNPs
# 3. draw n examples per chunk
# 4. plot results


def get_df_pnp_count(df):
    pass

def draw_pnp_examples(df, num_examples):
    pass

def draw_non_pnp_examples(df, num_examples, num_references):
    pass

def plot_pnp_ratio(df_pnp_count):
    pass


def main():

    # init
    d0 = datetime.now()
    description = """Script is used for gathering computational time in each step of the pipeline
    for a given chunk using all corresponding log files.

    I takes two inputs:

        - input directory
        - pattern for matching the corresponding chunk

    It creates one output:
        - a csv file with all computational times for each step as a single row.

    All chunk computational times can then easilsy gathered manually using pandas.

    It uses the installed npfc libary in your favorite env manager.

    Example:

        >>> report_time fc/04_synthetic/chembl/data/prep 'chembl_001*' fc/04_synthetic/chembl/data/prep/report/chembl_001_time.csv

    """

    # parameters CLI
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('input_dir', type=str, default=None, help="Input directory (i.e. fc/04_synthetic/chembl/data/prep)")
    parser.add_argument('pattern', type=str, default=None, help="Pattern of the file (i.e. prefix of a chunk) to investigate. All matching files in subdirectories will be processed.")
    parser.add_argument('output_csv', type=str, default=None, help="Output file for recording the computational times.")
    parser.add_argument('-p', '--prep', type=str, default=None, help="Subfolder name for prep step.")
    parser.add_argument('-n', '--natref', type=str, default=None, help="Subfolder name for natref step.")
    parser.add_argument('-f', '--frags', type=str, default=None, help="Subfolder name for frags step.")
    parser.add_argument('--log', type=str, default='INFO', help="Specify level of logging. Possible values are: CRITICAL, ERROR, WARNING, INFO, DEBUG.")
    args = parser.parse_args()

    # logging
    logger = utils._configure_logger(args.log)
    logger.info("RUNNING REPORT TIME")
    warnings.filterwarnings('ignore', category=pd.io.pytables.PerformanceWarning)  # if None is returned instead of a molecule, do not complain about mixed types
    pad = 40

    # parse arguments
    utils.check_arg_input_dir(args.input_dir)
    if args.input_dir.endswith('/'):
        args.input_dir = args.input_dir[0:-1]
    utils.check_arg_output_file(args.output_csv)
    out_format, out_compression = utils.get_file_format(args.output_csv)

    # display infos
    logger.info("LIBRARY VERSIONS:")
    logger.info("rdkit".ljust(pad) + f"{rdkit.__version__}")
    logger.info("pandas".ljust(pad) + f"{pd.__version__}")
    logger.info("npfc".ljust(pad) + f"{npfc.__version__}")
    logger.info("ARGUMENTS:")
    logger.info("INPUT_DIR".ljust(pad) + f"{args.input_dir}")
    logger.info("PATTERN".ljust(pad) + f"{args.pattern}")
    logger.info("OUTPUT_FILE".ljust(pad) + f"{args.output_csv}")
    logger.info("OUT_FORMAT".ljust(pad) + f"{out_format}")
    logger.info("OUT_COMPRESSION".ljust(pad) + f"{out_compression}")
    logger.info("LOG".ljust(pad) + f"{args.log}")
    d1 = datetime.now()

    # begin
    logger.info("BEGIN")

    # finding the input files
    input_files = find_input_files(args.input_dir, args.pattern, prep=args.prep, natref=args.natref, frags=args.frags)
    d = {}
    logger.info("STEP".ljust(15) + "TIME".ljust(15) + "INPUT_FILE".ljust(15))
    for input_file in input_files:
        step = Path(input_file).parent.parent.name
        comput_time = get_time(input_file)
        d["pattern"] = [args.pattern]
        d[f"{step}_time"] = [comput_time]
        logger.info(f"{step}".ljust(15) + f"{comput_time}".ljust(15) + f"{input_file}".ljust(15))

    # assemble data into a DF of one single row
    df = DataFrame(d)
    d2 = datetime.now()

    # export results
    if Path(args.output_csv).exists():
        df.to_csv(args.output_csv, mode='a', sep='|', header=False, index=False)
    else:
        df.to_csv(args.output_csv, mode='w+', sep='|', header=True, index=False)

    # end
    d3 = datetime.now()
    logger.info("SUMMARY")
    logger.info("COMPUTATIONAL TIME: CONFIGURING JOB".ljust(pad * 2) + f"{d1-d0}")
    logger.info("COMPUTATIONAL TIME: PARSING INPUT FILES".ljust(pad * 2) + f"{d2-d1}")
    logger.info("COMPUTATIONAL TIME: EXPORTING RESULTS".ljust(pad * 2) + f"{d3-d2}")
    logger.info("COMPUTATIONAL TIME: TOTAL".ljust(pad * 2) + f"{d3-d0}")
    logger.info("END")


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ MAIN ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #


if __name__ == '__main__':
    main()
    sys.exit(0)
