#imports
import argparse
import sys
import json
from Bio.Seq import Seq
import ggCaller_cpp
from ggCaller.graph_traversal import *
from multiprocessing import Pool
from functools import partial


def get_options():
    description = 'Generates ORFs from a Bifrost graph.'
    parser = argparse.ArgumentParser(description=description,
                                     prog='ggcaller')

    IO = parser.add_argument_group('Input/output')
    IO.add_argument('--graph',
                    default=None,
                    help='Bifrost GFA file generated by Bifrost build. ')
    IO.add_argument('--colours',
                    default=None,
                    help='Bifrost colours file generated by Bifrost build.  ')
    IO.add_argument('--not-ref',
                    action="store_false",
                    default=True,
                    help='If using existing graph, was not graph built exclusively with assembled genomes.  '
                         '[Default = False] ')
    IO.add_argument('--refs',
                    default=None,
                    help='List of reference genomes (one file path per line). ')
    IO.add_argument('--reads',
                    default=None,
                    help='List of reference genomes (one file path per line). ')
    IO.add_argument('--codons',
                    default=None,
                    help='JSON file containing start and stop codon sequences. ')
    IO.add_argument('--kmer',
                        default=31,
                        help='K-mer size used in Bifrost build (bp). '
                        '[Default = 31] ')
    IO.add_argument('--path',
                    default=10000,
                    help='Maximum path length during traversal (bp). '
                    '[Default = 10000] ')
    IO.add_argument('--orf',
                    default=90,
                    help='Minimum ORF length to return (bp). '
                    '[Default = 90] ')
    IO.add_argument('--maxoverlap',
                    default=60,
                    help='Maximum overlap allowed between overlapping ORFs. '
                    '[Default = 60] ')
    IO.add_argument('--min-path-score',
                    default=100,
                    help='Minimum total score for a path of ORFs to be returned. '
                    '[Default = 100] ')
    IO.add_argument('--min-orf-score',
                    default=100,
                    help='Minimum individual score for an ORF to be returned. '
                    '[Default = 100] ')
    IO.add_argument('--no-filter',
                    action="store_true",
                    default=False,
                    help='Do not filter ORF calls using Balrog. Will return all ORF calls. '
                         '[Default = False] ')
    IO.add_argument('--no-write-idx',
                    action="store_false",
                    default=True,
                    help='Do not write FMIndexes to file. '
                         '[Default = False] ')
    IO.add_argument('--no-write-graph',
                    action="store_false",
                    default=True,
                    help='Do not write Bifrost GFA and colours to file. '
                         '[Default = False] ')
    IO.add_argument('--repeat',
                    action="store_true",
                    default=False,
                    help='Enable traversal of nodes mulitple times. '
                         '[Default = False] ')
    IO.add_argument('--threads',
                    default=1,
                    help='Number of threads for FMIndexing '
                         '[Default = 1] ')
    IO.add_argument('--out',
                    default='calls.fasta',
                    help='Output FASTA file containing ORF sequences. ')
    return parser.parse_args()

def main():
    options = get_options()

    # parse command line arguments
    graph_file = options.graph
    colours_file = options.colours
    is_ref = bool(options.not_ref)
    refs_file = options.refs
    reads_file = options.reads
    codons = options.codons
    ksize = int(options.kmer)
    max_path_length = int(options.path)
    min_ORF_length = int(options.orf)
    max_ORF_overlap = int(options.maxoverlap)
    minimum_path_score = int(options.min_path_score)
    minimum_ORF_score = int(options.min_orf_score)
    no_filter = bool(options.no_filter)
    write_idx = bool(options.no_write_idx)
    write_graph = bool(options.no_write_graph)
    repeat = bool(options.repeat)
    num_threads = int(options.threads)
    output = options.out

    # define start/stop codons
    if codons != None:
        with open(codons, "r") as json_file:
            try:
                data = json.load(json_file)
                start_codons = data["codons"]["start"]
                stop_codon_for = data["codons"]["stop"]
                stop_codon_rev = [str((Seq(i)).reverse_complement()) for i in stop_codon_for]
            except:
                print("Please specify codons in the format shown in codons.json.")
                sys.exit(1)
    else:
        start_codons = ["ATG", "GTG", "TTG"]
        stop_codon_for = ["TAA", "TGA", "TAG"]
        stop_codon_rev = ["TTA", "TCA", "CTA"]

    # if build graph specified, build graph and then call ORFs
    if (graph_file != None) and (colours_file != None) and (refs_file == None) and (reads_file == None):
        called_ORF_tuple = ggCaller_cpp.call_genes_existing(graph_file, colours_file, start_codons, stop_codon_for, stop_codon_rev, num_threads, is_ref, write_idx, repeat, max_path_length, min_ORF_length, max_ORF_overlap)
    # if refs file specified for building
    elif (graph_file == None) and (colours_file == None) and (refs_file != None) and (reads_file == None):
        called_ORF_tuple = ggCaller_cpp.call_genes_build(refs_file, ksize, start_codons, stop_codon_for, stop_codon_rev, num_threads, True, write_idx, repeat, write_graph, max_path_length, min_ORF_length, max_ORF_overlap)
    # if reads file specified for building
    elif (graph_file == None) and (colours_file == None) and (refs_file == None) and (reads_file != None):
        called_ORF_tuple = ggCaller_cpp.call_genes_build(reads_file, ksize, start_codons, stop_codon_for, stop_codon_rev, num_threads, False, write_idx, repeat, write_graph, max_path_length, min_ORF_length, max_ORF_overlap)
    # if both reads and refs file specified for building
    elif (graph_file == None) and (colours_file == None) and (refs_file != None) and (reads_file != None):
        called_ORF_tuple = ggCaller_cpp.call_genes_build(refs_file, ksize, start_codons, stop_codon_for, stop_codon_rev, num_threads, False, write_idx, repeat, write_graph, max_path_length, min_ORF_length, max_ORF_overlap, reads_file)
    else:
        print("Error: incorrect number of input files specified. Please only specify the below combinations:\n"
              "- Bifrost GFA and Bifrost colours file\n"
              "- List of reference files\n"
              "- List of read files\n"
              "- A list of reference files and a list of read files.")
        sys.exit(1)

    # unpack ORF pair into overlap dictionary and list for gene scoring
    ORF_overlap_dict, full_ORF_dict, ORF_colour_ID_map = called_ORF_tuple

    # create list for high scoring ORFs to return
    true_genes = {}

    # if no filter specified, generate ORF sequences and combine colours of identical ORFS
    if no_filter == True:
        for colour, gene_set in ORF_colour_ID_map.items():
            for ORF_ID, ORF_seq in gene_set.items():
                gene = ORF_seq[16:]
                if gene not in true_genes:
                    true_genes[gene] = colour
                else:
                    updated_colour = update_colour(true_genes[gene], colour)
                    true_genes[gene] = updated_colour
    else:
        # generate gene scores using Balrog
        full_ORF_dict = score_genes(full_ORF_dict, ORF_colour_ID_map, minimum_ORF_score, num_threads)

        print("Generating highest scoring gene paths...")

        # merge full_ORF_dict and ORF_overlap_dict for mulitprocessing and free up memory
        ORF_dict = {}
        for colour in full_ORF_dict.keys():
            ORF_dict[colour] = {}
            ORF_dict[colour]["ID"] = full_ORF_dict[colour]
            ORF_dict[colour]["overlap"] = ORF_overlap_dict[colour]

        # multithread per colour in ORF_dict
        # Turn gt threading on
        if gt.openmp_enabled():
            gt.openmp_set_num_threads(1)

        with Pool(processes=num_threads) as pool:
            for colour, high_scoring_ORFs in pool.map(
                    partial(call_true_genes, minimum_path_score=minimum_path_score),
                    ORF_dict.items()):
                # remove TIS, and merge any matching genes which had differing TIS but were called together
                for ORF_ID in high_scoring_ORFs:
                    ORF_seq = ORF_colour_ID_map[colour][ORF_ID]
                    gene = ORF_seq[16:]
                    if gene not in true_genes:
                        true_genes[gene] = colour
                    else:
                        updated_colour = update_colour(true_genes[gene], colour)
                        true_genes[gene] = updated_colour

    print("Generating fasta file of gene calls...")
    # print output to file
    ORF_count = 1
    with open(output, "w") as f:
        for gene, colour in true_genes.items():
            f.write(">" + str(ORF_count) + "_" + str(colour) + "\n" + gene + "\n")
            ORF_count += 1

    print("Finished.")

    sys.exit(0)

if __name__ == '__main__':
    main()
