#!/usr/bin/env python
# -*- coding: utf-8
"""Returns sequences for a given list of gene caller ids"""

import sys

import anvio
import anvio.utils as utils
import anvio.terminal as terminal
import anvio.filesnpaths as filesnpaths
import anvio.genomestorage as genomestorage

from anvio.errors import ConfigError, FilesNPathsError
from anvio.dbops import ContigsSuperclass


__author__ = "Developers of anvi'o (see AUTHORS.txt)"
__copyright__ = "Copyleft 2015-2018, the Meren Lab (http://merenlab.org/)"
__credits__ = []
__license__ = "GPL 3.0"
__version__ = anvio.__version__
__maintainer__ = "A. Murat Eren"
__email__ = "a.murat.eren@gmail.com"
__requires__ = ['contigs-db',]
__provides__ = ['genes-fasta',]
__description__ = "A script to get back sequences for gene calls"


run = terminal.Run()
progress = terminal.Progress()


def export_from_contigs(args):
    c = ContigsSuperclass(args)

    output_file_path = args.output_file if args.output_file else 'sequences_for_gene_calls.txt'
    filesnpaths.is_output_file_writable(output_file_path)

    gene_caller_ids = list(utils.get_gene_caller_ids_from_args(args.gene_caller_ids, args.delimiter))

    func_kwargs = dict(
        gene_caller_ids_list=gene_caller_ids,
        output_file_path=args.output_file,
        simple_headers=not args.report_extended_deflines,
        wrap=args.wrap
    )

    if args.export_gff3:
        if args.get_aa_sequences:
            raise ConfigError("AA sequences can only be reported in FASTA format, please remove the --export-gff3 flag to continue.")

        c.gen_GFF3_file_of_sequences_for_gene_caller_ids(**func_kwargs)
    else:
        c.gen_FASTA_file_of_sequences_for_gene_caller_ids(**func_kwargs, report_aa_sequences=args.get_aa_sequences)


def export_from_genomes_storage(genomes_storage_db_path, output_file_path):
    g = genomestorage.GenomeStorage(genomes_storage_db_path)
    g.gen_combined_aa_sequences_FASTA(output_file_path)


def main(args):
    A = lambda x: args.__dict__[x] if x in args.__dict__ else None
    contigs_db_path = A('contigs_db')
    genomes_storage_db_path = A('genomes_storage')
    output_file_path = A('output_file')

    if not contigs_db_path and not genomes_storage_db_path:
        raise ConfigError("You must give this program either a contigs or a genomes storage database so it can like "
                          "export sequences for your genes? :/")

    if contigs_db_path and genomes_storage_db_path:
        raise ConfigError("You can either ask for sequences in a contigs database, or a genomes storage. But not both "
                          "(obviously).")

    if contigs_db_path:
        export_from_contigs(args)
    elif genomes_storage_db_path:
        export_from_genomes_storage(genomes_storage_db_path, output_file_path)
    else:
        raise ConfigError("o_O")


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description=__description__)

    groupA = parser.add_argument_group('OPTION #1: EXPORT FROM CONTIGS DB')
    groupA.add_argument(*anvio.A('contigs-db'), **anvio.K('contigs-db', {'required': False}))
    groupA.add_argument(*anvio.A('gene-caller-ids'), **anvio.K('gene-caller-ids'))
    groupA.add_argument(*anvio.A('delimiter'), **anvio.K('delimiter'))
    groupA.add_argument(*anvio.A('report-extended-deflines'), **anvio.K('report-extended-deflines'))
    groupA.add_argument(*anvio.A('wrap'), **anvio.K('wrap'))
    groupA.add_argument(*anvio.A('export-gff3'), **anvio.K('export-gff3'))
    groupA.add_argument(*anvio.A('get-aa-sequences'), **anvio.K('get-aa-sequences'))

    groupB = parser.add_argument_group('OPTION #2: EXPORT FROM A GENOMES STORAGE')
    groupB.add_argument(*anvio.A('genomes-storage'), **anvio.K('genomes-storage', {'required': False}))
    groupB.add_argument(*anvio.A('genomes-names'), **anvio.K('genomes-names'))

    groupC = parser.add_argument_group('OPTIONS COMMON TO ALL INPUTS')
    groupC.add_argument(*anvio.A('output-file'), **anvio.K('output-file', {'required': True}))

    args = anvio.get_args(parser)

    try:
        main(args)
    except ConfigError as e:
        print(e)
        sys.exit(-1)
    except FilesNPathsError as e:
        print(e)
        sys.exit(-2)
