#!/usr/bin/env python
import argparse
import pkg_resources
import immunedb.common.config as config
import immunedb.exporting as exporting


if __name__ == '__main__':
    main_parser = config.get_base_arg_parser('Exports ImmuneDB data in '
                                             'various formats',
                                             multiproc=False)
    subparsers = main_parser.add_subparsers(dest='fmt', help='The format')
    subparsers.required = True

    # Genbank
    parser = subparsers.add_parser(
        'genbank', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--inference', default='ImmuneDB:{}'.format(
                        pkg_resources.get_distribution('immunedb').version),
                        help='Inference tool used for gene assignment')
    parser.add_argument('--gene-db', default='IMGT/GENE-DB',
                        help='Reference V/J gene database')
    parser.add_argument('--species', default='Homo sapiens',
                        help='Species of sequences')
    parser.add_argument('--mol-type', default='DNA', help='Molecule type')

    # Raw sequences
    parser = subparsers.add_parser(
        'sequences', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--sample-ids', type=int, nargs='+', help='''If
                        specified, limits the export to certain samples.''')
    parser.add_argument('--format', default='changeo', dest='out_format',
                        choices=['changeo', 'airr'], required=True)
    parser.add_argument('--clones-only', action='store_true', default=False,
                        help='Includes only sequences assigned to a clone')
    parser.add_argument('--min-subject-copies', default=None,
                        help='''The minimum copies in the subject required to
                        be exported.  Set to 1 for all subject-unique
                        sequences, or higher for more stringent filtering''')

    # Clone selection pressure
    parser = subparsers.add_parser(
        'selection', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--filter', default='both', dest='filter_type',
                        choices=['both', 'samples', 'overall'],
                        help='''What selection pressure to export.  "overall"
                        outputs one row per clone of its overall selection
                        pressure, "samples" outputs one row per sample per
                        clone indicating its selection pressure in those
                        samples, and "both" outputs both of the above.''')

    # Sample metadata
    parser = subparsers.add_parser(
        'samples', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--for-update', action='store_true', help='''When set,
        formats the metadata for use in `immunedb_update update-metadata` to
        update metadata.''')

    # Pooled analysis
    parser = subparsers.add_parser(
        'clones', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--pool-on', nargs='+', default=('sample',), help='''
        The metadata field(s) on which to pool.''')
    parser.add_argument('--format', default='immunedb', dest='out_format',
                        choices=['immunedb', 'vdjtools'],
                        help='''The format in which to export''')
    parser.add_argument('--sample-ids', type=int, nargs='+', help='''If
                        specified, limits the export to certain samples.''')

    # Overlap
    parser = subparsers.add_parser(
        'overlap',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('--pool-on', nargs='+', default='sample', help='''
        The metadata field(s) on which to pool.''')
    parser.add_argument('--sim-func', default='cosine',
                        choices=['cosine', 'jaccard'])
    parser.add_argument('--size-metric', default='copies',
                        choices=['copies', 'instances'])
    parser.add_argument('--agg-func', default='median',
                        choices=['mean', 'median'])
    parser.add_argument('--sample-ids', type=int, nargs='+')

    args = main_parser.parse_args()
    session = config.init_db(args.db_config)

    fmts = {
        'genbank': exporting.write_genbank,
        'sequences': exporting.write_sequences,
        'selection': exporting.write_selection,
        'samples': exporting.write_samples,
        'clones': exporting.write_pooled_clones,
        'overlap': exporting.write_clone_overlap,
    }

    fmts[args.fmt](session, **vars(args))
