# -*- coding: utf-8 -*-
from os import path, listdir, remove
from contextlib import contextmanager
from multiprocessing import Pool, cpu_count
from distutils.dir_util import copy_tree

from Bio import SeqIO
from Bio.SeqFeature import SeqFeature, FeatureLocation
from Bio.Alphabet import generic_dna
import gffutils

__author__ = u"Mikołaj Dziurzyński"


SIMILARITY_THRESHOLD = 50

PROGRAM_PRIORITY = {"genemarks": 1, "prodigal": 2}


def _are_overlapping(mp_feature, ap_feature):
    mp_start = mp_feature.location.start.position
    mp_end = mp_feature.location.end.position
    ap_start = ap_feature.location.start.position
    ap_end = ap_feature.location.end.position

    if (mp_start >= ap_start and mp_start <= ap_end) or \
       (mp_end >= ap_start and mp_end <= ap_end) or \
       (mp_start <= ap_start and mp_end >= ap_end):
        return True
    else:
        return False


def _merge_gbk_runner(name, mp_name, mp_path, ap_name, ap_path):

    mp_file_path = path.join(mp_path, name + ".gbk")
    ap_file_path = path.join(ap_path, name + ".gbk")

    main_gbk = SeqIO.read(mp_file_path, 'genbank', generic_dna)
    ap_gbk = SeqIO.read(ap_file_path, 'genbank', generic_dna)

    new_features_list = []

    matched_index = set()

    for mp_feature in main_gbk.features:
        for i, ap_feature in enumerate(ap_gbk.features):

            if ap_feature.location.start.position > mp_feature.location.end.position:
                break

            if mp_feature.strand != ap_feature.strand or \
               ap_feature.location.end.position < mp_feature.location.start.position or \
               mp_feature.type != ap_feature.type:
                continue

            if mp_feature.strand == 1:
                mp_frame = mp_feature.location.start.position % 3
                ap_frame = ap_feature.location.start.position % 3
            elif mp_feature.strand == -1:
                mp_frame = (len(main_gbk) - mp_feature.location.end.position - 1) % 3
                ap_frame = (len(ap_gbk) - ap_feature.location.end.position - 1) % 3
            else:
                raise Exception("Wrong strand value! {0}".format(main_gbk.name))

            if mp_frame != ap_frame:
                continue

            are_overlapping = _are_overlapping(mp_feature, ap_feature)

            if not are_overlapping:
                continue
            else:
                # add qualifier to original feature
                if ap_feature.location.start.position != mp_feature.location.start.position or \
                   ap_feature.location.end.position != mp_feature.location.end.position:
                    mp_feature.qualifiers.update({
                        "maisen_coords_{0}".format(ap_name): "{0}:{1}".format(
                            ap_feature.location.start.position+1,
                            ap_feature.location.end.position)
                    })
                matched_index = matched_index | {i}

        new_features_list.append(mp_feature)

    for i, feature in enumerate(ap_gbk.features):
        if i not in matched_index:
            new_features_list.append(feature)

    new_features_list.sort(key=lambda x: x.location.start.position)
    main_gbk.features = new_features_list

    SeqIO.write(main_gbk, mp_file_path, 'genbank')


def merge_gbk(main_program, additional_program):
    """
    Based on SIMILARITY_THRESHOLD add new genes from other gbks into the main one.

    NOTE: check if genes are in the same reading frame
    """

    # add features from 'additional_program' to 'main_program'
    # gbk and save in main_program loc

    mp_path = path.join("tmp/intermediate_gbk", main_program)
    ap_path = path.join("tmp/intermediate_gbk", additional_program)

    mp_gbk_list = {f.split('.')[0] for f in listdir(mp_path) if path.isfile(path.join(mp_path, f))}
    ap_gbk_list = {f.split('.')[0] for f in listdir(ap_path) if path.isfile(path.join(ap_path, f))}

    not_in_mp = ap_gbk_list - mp_gbk_list
    not_in_ap = mp_gbk_list - ap_gbk_list
    if not_in_ap or not_in_mp:
        print("The following gbk files are not \
            present at:\n-{0}: {1}\n-{2}: {3}".format(main_program, not_in_mp, additional_program, not_in_ap))

    present_in_both = mp_gbk_list & ap_gbk_list

    numer_of_workers = cpu_count() - 1

    with Pool(processes=numer_of_workers) as pool:
        multiple_results = [pool.apply_async(
            _merge_gbk_runner, (name, main_program, mp_path, additional_program, ap_path))
            for name in present_in_both]

        [res.get() for res in multiple_results]


def _gff_gbr_runner(name, gff_dir, fasta_dir, program):
    record = SeqIO.read(path.join(fasta_dir, name + ".fasta"), "fasta", generic_dna)
    with gff_db_manager(path.join(gff_dir, name+".gff")) as db:
        if db:
            for feature in db.all_features():
                f_strand = feature.strand == '+' and 1 or feature.strand == '-' and -1 or 0
                f_qualifiers = {
                    "maisen": True,
                    "maisen_coords_{0}".format(program): "{0}:{1}".format(feature.start, feature.end),
                }
                new_feature = SeqFeature(
                    location=FeatureLocation(feature.start - 1, feature.end),
                    type=feature.featuretype,
                    strand=f_strand,
                    qualifiers=f_qualifiers
                    )
                # if feature.featuretype.upper() == 'CDS':
                #     new_feature.qualifiers.update({
                #         'translation': str(new_feature.extract(record.seq).translate(stop_symbol=''))
                #     })

                record.features.append(new_feature)

    output_gbk = path.join("tmp/intermediate_gbk", program, name + ".gbk")
    SeqIO.write(record, output_gbk, "genbank")


@contextmanager
def gff_db_manager(gff_path):
    db_name = "{0}_db".format(path.basename(gff_path).split(".")[0])
    db_path = path.join("tmp/gff_dbs", db_name)
    try:
        db = gffutils.create_db(gff_path, db_path)
    # empty gff?
    except ValueError:
        db = []
    yield db
    del db
    remove(db_path)


def parse_gff_to_gbk(program, gff_dir, fasta_dir):
    """
    Create gbk file for each gff in program directory (tmp/output/*program*/)
    """

    # get gff_dir files
    gff_names = [f.split('.gff')[0] for f in listdir(gff_dir) if path.isfile(path.join(gff_dir, f))]

    fasta_names = [f.split('.fasta')[0] for f in listdir(fasta_dir) if path.isfile(path.join(fasta_dir, f))]

    # for each single_fasta check if gff exists
    run_names = []
    for fasta in fasta_names:
        if fasta in gff_names:
            run_names.append(fasta)
        else:
            print("No matching gff for {0}.fasta from program {1}".format(fasta, program))

    number_of_workers = cpu_count() - 1

    with Pool(processes=number_of_workers) as pool:
        multiple_results = [pool.apply_async(
            _gff_gbr_runner, (name, gff_dir, fasta_dir, program))
            for name in run_names]

        [res.get() for res in multiple_results]

    return True


def create_gbk_output(program_list, output_path, cwd):
    """
    Create gbk output from computed gff files. GFF from GeneMarkerS
    creates basic gbk, other programs output is added to it.
    """

    # create gbk from every fasta+gff pair per program

    gff_output_dir = path.join(cwd, "tmp/outputs")
    input_fastas_dir = path.join(cwd, "tmp/single_fastas")

    for program in program_list:
        parse_gff_to_gbk(program, path.join(gff_output_dir, program), input_fastas_dir)

    # merge all gbks into one using PROGRAM_PRIORITY list
    main_program_number = min([PROGRAM_PRIORITY[program] for program in program_list])
    rev_priority_dict = {}
    for k in PROGRAM_PRIORITY:
        rev_priority_dict[PROGRAM_PRIORITY[k]] = k
    main_program = rev_priority_dict[main_program_number]
    remaining_programs = [program for program in program_list if program != main_program]

    for program in remaining_programs:
        merge_gbk(main_program, program)

    # save_at_output(main_program, output_path)
    mp_path = path.join("tmp/intermediate_gbk", main_program)
    copy_tree(mp_path, output_path)

    return True
