!/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_dfs_pnp(WD: Path) -> DataFrame:
    """Get a DF summarizing the results of the pnp step.

    At the moment only one subset is recognized (i.e. 'subset' subfolder in WD).


    :param WD: the main directory of the dataset data (i.e. 'natural/dnp/data')
    :return: a DF summarizing results of the murcko subset step
    """
    logger.info("PNP -- COMPUTING RESULTS FOR PNP")
    if not isinstance(WD, Path):
        WD = Path(WD)
    # parse results before fragment search
    WD_PNP = [str(x) for x in list(WD.glob("*_pnp"))][0]    # get latest step
    pattern = ".*_list_pnp.csv.gz$"
    chunks = _get_chunks(f"{WD_PNP}/log", pattern)
    chunks = [c for c in chunks if c.endswith('.csv.gz')]

    num_pnp = 0
    num_non_pnp = 0
    num_total = 0

    pool = Pool()
    results = pool.map(parse_chunk_pnp, chunks)
    pool.close()
    pool.join()

    # sum of tuples
    num_pnp = sum([x['pnp'] for x in results])
    num_non_pnp = sum([x['non_pnp'] for x in results])
    # dfs
    dfs_pnp_nnprefperfcg = [x['df_pnp_nnprefperfcg'] for x in results]

    # ratio PNP/non-PNP
    num_total = num_pnp + num_non_pnp
    logger.info(f"pnp: {num_pnp:,} + non-pnp: {num_non_pnp:,} = tot: {num_total:,}")
    # create a dataframe with counts
    df_pnp_ratio = pd.DataFrame({'Category': ['PNP', 'Non-PNP'], 'Count': [num_pnp, num_non_pnp]})
    df_pnp_ratio['Perc_Mols'] = df_pnp_ratio['Count'].map(lambda x: f"{x/num_total:.2%}")
    logger.info(f"PNP -- RESULTS FOR LABELLING PNPs IN {num_total:,} MOLECULES:\n\n{df_pnp_ratio}\n")

    # Number of NP refs per FCG
    df_pnp_nnprefperfcg = pd.concat(dfs_pnp_nnprefperfcg)
    df_pnp_nnprefperfcg = df_pnp_nnprefperfcg[['NumNPRef', 'Count']].groupby('NumNPRef').sum().reset_index()
    tot = df_pnp_nnprefperfcg['Count'].sum()
    df_pnp_nnprefperfcg['Perc'] = df_pnp_nnprefperfcg['Count'].map(lambda x: f"{x / tot:.2%}")
    logger.info(f"PNP -- RESULTS NUMBER OF NP REFERENCES PER FCG:\n\n{df_pnp_nnprefperfcg}\n")

    return {'df_pnp_ratio': df_pnp_ratio, 'df_pnp_nnprefperfcg': df_pnp_nnprefperfcg}

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:
        num_pnp = 0
        num_non_pnp = 0
        num_total = 0
        df = load.file(input_file, decode=True)[['idm', 'idfcg', '_pnp_ref', 'mol', '_colormap']]
        df_pnp = df[df['pnp_mol']].groupby('idm')
        df_non_pnp = df[~df['pnp_mol']].groupby('idm')





        # number of PNP references per FCG
        df = df[['idm', '_pnp_ref']]
        df['_pnp_ref'] = df['_pnp_ref'].map(lambda x: len(utils.decode_object(x)))
        df = df.rename({'_pnp_ref': 'NumNPRef', 'idm': 'Count'}, axis=1).groupby('NumNPRef').count().reset_index()


    # 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)
