#!/usr/bin/env python

import csv
import sys
import json
import logging
import multiprocessing
import pandas

import ndasynapse

pandas.options.display.max_rows = None
pandas.options.display.max_columns = None
pandas.options.display.max_colwidth = 1000

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

def get_guid(auth, args):
    guids = ndasynapse.nda.get_guid(auth, args.guid)

def get_collection_submissions(auth, args):
    logger.debug("collectionids = {collection_id}".format(collection_id=args.collection_id))
    submissions = ndasynapse.nda.get_submissions(auth=auth, 
                                                 collectionid=[str(x) for x in args.collection_id])
    submissions_processed = ndasynapse.nda.process_submissions(submissions)

    submissions_processed.to_csv(sys.stdout, index=False, quoting=csv.QUOTE_NONNUMERIC,
                                 encoding='utf-8')

def get_submission(auth, args):
    submission = ndasynapse.nda.NDASubmission(auth=auth, submission_id=args.submission_id)
    if args.json:
        sys.stdout.write(json.dumps(submission.submission, indent=2))
    else:
        submission.processed_submission.to_csv(sys.stdout, index=False, 
                                               quoting=csv.QUOTE_NONNUMERIC,
                                               encoding='utf-8')

def get_submission_files(auth, args):
    submission = ndasynapse.nda.get_submission_files(auth=auth,
                                                     submissionid=args.submission_id)
    submissions_processed = ndasynapse.nda.process_submission_files(submission)
    submissions_processed.to_csv(sys.stdout, index=False,
                                 quoting=csv.QUOTE_NONNUMERIC,
                                 encoding='utf-8')

def get_collection_submission_files(auth, args):
    nda_collection = ndasynapse.nda.NDACollection(auth=auth,
                                                  collection_id=args.collection_id) # pylint: disable:line-too-long
    
    sub_files = [sub.submission_files['processed_files'] for sub in nda_collection.submissions] # pylint: disable:line-too-long
    submission_files_processed = pandas.concat(sub_files)
    submission_files_processed.to_csv(sys.stdout, index=False,
                                      quoting=csv.QUOTE_NONNUMERIC,
                                      encoding='utf-8')


def get_experiments(auth, args):
    data = ndasynapse.nda.get_experiments(auth, args.experiment_id)

    if args.json:
        sys.stdout.write(json.dumps(data, indent=2))
    else:
        data = ndasynapse.nda.process_experiments(data)
        data = data.drop_duplicates()
        data.to_csv(sys.stdout, index=False,
                    quoting=csv.QUOTE_NONNUMERIC,
                    encoding='utf-8')


def get_samples(auth, args):
    data = ndasynapse.nda.get_samples(auth=auth, guid=args.guid)

    if args.json:
        sys.stdout.write(json.dumps(data, indent=2))
    else:
        data_df = ndasynapse.nda.process_guid_data(data)
        data_df = ndasynapse.nda.process_samples(data_df)
        data_df.to_csv(sys.stdout, index=False, 
                       quoting=csv.QUOTE_NONNUMERIC,
                       encoding='utf-8')


def get_subjects(auth, args):
    data = ndasynapse.nda.get_subjects(auth=auth, guid=args.guid)

    if args.json:
        sys.stdout.write(json.dumps(data, indent=2))
    else:
        data_df = ndasynapse.nda.subjects_to_df(data)
        data_df = ndasynapse.nda.process_subjects(data_df)
        data_df.to_csv(sys.stdout, index=False,
                       quoting=csv.QUOTE_NONNUMERIC,
                       encoding='utf-8')

_guid_fieldnames = ['submission_id', 'guid']


def get_collection_guids(auth, args):
    nda_collection = ndasynapse.nda.NDACollection(
        auth=auth, collection_id=args.collection_id)
    guids = [str(guid) for guid in nda_collection.guids]
    sys.stdout.write("\n".join(guids))


def get_collection_manifests(auth, args):
    """Get all original manifests submitted with each submission in a collection.

    NDA does not update these files if metadata change requests are made.
    They only update metadata in their database, accessible through the GUID API.
    Records obtained here should be used for historical purposes only.

    Args:
        auth: a requests.auth.HTTPBasicAuth object authenticating to NDA.
        args: arpgarse arguments
    Returns:
        Output to stdout in CSV format of all submission manifests concatenated together.
    """
    
    all_data = []

    for collection_id in args.collection_id:
        nda_collection = ndasynapse.nda.NDACollection(
            auth=auth, collection_id=collection_id)

        manifest_data = nda_collection.get_collection_manifests(
            manifest_type=args.manifest_type)
        all_data.append(manifest_data)

    if all_data:
        all_data_df = pandas.concat(all_data, axis=0, 
                                    ignore_index=True, sort=False)
        all_data_df.to_csv(sys.stdout, index=False,
                           quoting=csv.QUOTE_NONNUMERIC,
                           encoding='utf-8')


def get_guid_collection_manifests(auth, args):
    pool = multiprocessing.dummy.Pool(args.parallel)

    guid_worker = lambda guid: ndasynapse.nda.get_guid_data(
        auth=auth, subjectkey=guid,
        short_name=args.manifest_type)

    collection_worker = lambda coll_id: ndasynapse.nda.NDACollection(auth=auth,
                                                                     collection_id=coll_id)

    collections = pool.map(collection_worker, args.collection_id)

    all_collections_data = [] #pandas.DataFrame()

    for nda_collection in collections:

        all_guids_data = []
        coll_id = nda_collection.collection_id
        guid_data_list = pool.map(guid_worker, nda_collection.guids)

        for (guid, guid_data) in zip(nda_collection.guids, guid_data_list):
            # It is possible for there to be no data for the specified
            # manifest type. If this is the case, the GUID API will return an
            # OK status (status_code = 200) and an empty data structure, which
            # will cause the code to crash further down, so check to make sure
            # that the data structure is not empty before continuing.
            if guid_data is None or not guid_data["age"]:
                logger.warn(f"No data for guid {guid}")
                continue

            data = ndasynapse.nda.process_guid_data(guid_data, collection_ids=[int(coll_id)],
                                                    drop_duplicates=True)
            all_guids_data.append(data)
        
        all_guids_df = pandas.concat(all_guids_data, axis=0, 
                                     ignore_index=True, sort=False)

        all_collections_data.append(all_guids_df)

    all_collections_df = pandas.concat(all_collections_data, axis=0, 
                                       ignore_index=True, sort=False)

    all_collections_df.to_csv(sys.stdout, index=False,
                              quoting=csv.QUOTE_NONNUMERIC,
                              encoding='utf-8')


def main():

    import argparse
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--verbose", action="store_true", default=False)
    parser.add_argument("--version", action="store_true", 
                        default=False, help="Print version.")
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--json", action="store_true", default=False,
                        help="Output in JSON format, if possible. Default is to output in CSV format.")
    parser.add_argument("--parallel", type=int, default=4,
                        help="Run in parallel threads, if enabled.")

    subparsers = parser.add_subparsers(help='sub-command help')

    parser_get_samples = subparsers.add_parser('get-samples',
                                               help='Get samples from NDA.')
    parser_get_samples.add_argument('--guid', type=str, 
                                    help='NDA subject key (GUID).')
    parser_get_samples.set_defaults(func=get_samples)

    parser_get_subjects = subparsers.add_parser('get-subjects',
                                                help='Get subjects from NDA.')
    parser_get_subjects.add_argument('--guid', type=str,
                                     help='NDA subject key (GUID).')
    parser_get_subjects.set_defaults(func=get_subjects)

    parser_get_experiments = subparsers.add_parser('get-experiments', help='Get experiments from NDA.')
    parser_get_experiments.add_argument('--experiment_id', type=int, nargs="+", help='NDA experiment IDs.')
    parser_get_experiments.set_defaults(func=get_experiments)

    parser_get_submissions = subparsers.add_parser('get-collection-submissions', help='Get submissions in NDA collections.')
    parser_get_submissions.add_argument('--collection_id', type=int, nargs="+", help='NDA collection IDs.')
    parser_get_submissions.set_defaults(func=get_collection_submissions)

    parser_get_submission = subparsers.add_parser('get-submission', help='Get an NDA submission.')
    parser_get_submission.add_argument('--submission_id', type=int, help='NDA submission ID.')
    parser_get_submission.set_defaults(func=get_submission)

    parser_get_submission_files = subparsers.add_parser('get-submission-files', help='Get files for an NDA submission.')
    parser_get_submission_files.add_argument('--submission_id', type=int, help='NDA submission ID.')
    parser_get_submission_files.set_defaults(func=get_submission_files)

    parser_get_collection_submission_files = \
        subparsers.add_parser('get-collection-submission-files', 
                              help='Get all submission files in an NDA collection.')
    parser_get_collection_submission_files.add_argument('--collection_id', 
                                                        type=int, 
                                                        help='NDA collection IDs.')
    parser_get_collection_submission_files.set_defaults(func=get_collection_submission_files)

    parser_get_collection_guids = \
        subparsers.add_parser('get-collection-guids', 
                              help='Get all GUIDS used in an NDA collection.')
    parser_get_collection_guids.add_argument('--collection_id', 
                                             type=int, 
                                             help='NDA collection ID.')
    parser_get_collection_guids.set_defaults(func=get_collection_guids)

    parser_get_collection_manifest = subparsers.add_parser(
        'get-collection-manifests',
        help='Get manifest files that were submitted in an NDA collection.')
    parser_get_collection_manifest.add_argument('--collection_id', type=int, 
                                                nargs="+", help='NDA collection ID.')
    parser_get_collection_manifest.add_argument('--manifest_type', type=str,
                                                help='Manifest type to retrieve.',
                                                choices=["genomics_sample", 
                                                         "genomics_subject",
                                                         "nichd_btb"])
    parser_get_collection_manifest.set_defaults(func=get_collection_manifests)

    parser_get_guid_collection_manifests = subparsers.add_parser(
        'get-guid-collection-manifests',
        help='Get manifest data for for completed submissions using the NDA GUID service.')
    parser_get_guid_collection_manifests.add_argument('--collection_id', type=int, 
                                                     nargs="+", help='NDA collection ID.')
    parser_get_guid_collection_manifests.add_argument('--manifest_type', type=str,
                                                    help='Manifest type to retrieve.',
                                                    choices=["genomics_sample03", 
                                                             "genomics_subject02",
                                                             "nichd_btb02"])
    parser_get_guid_collection_manifests.set_defaults(func=get_guid_collection_manifests)

    args = parser.parse_args()

    if args.version:
        print(ndasynapse.__version__)
        sys.exit()
    
    config = json.load(open(args.config))
    auth = ndasynapse.nda.authenticate(config)
    logger.info(auth)
    
    args.func(auth, args)


if __name__ == "__main__":
    main()
