"""
primer creation and evaluation
"""

# BUILTIN
import itertools
import re
import multiprocessing
import functools

# LIBS
from Bio.Seq import MutableSeq
from Bio import SeqIO
import primer3 as p3

# varVAMP
from varvamp.scripts import config


def calc_gc(seq):
    """
    calculate the gc of a sequence
    """
    return 100*(seq.count("g")+seq.count("c"))/len(seq)


def calc_temp(seq):
    """
    calculate the melting temperature
    """
    return p3.calc_tm(
            seq.upper(),
            mv_conc=config.PCR_MV_CONC,
            dv_conc=config.PCR_DV_CONC,
            dntp_conc=config.PCR_DNTP_CONC,
            dna_conc=config.PCR_DNA_CONC
        )


def calc_hairpin(seq):
    """
    calculates hairpins
    """
    return p3.calc_hairpin(
            seq.upper(),
            mv_conc=config.PCR_MV_CONC,
            dv_conc=config.PCR_DV_CONC,
            dntp_conc=config.PCR_DNTP_CONC,
            dna_conc=config.PCR_DNA_CONC
        )


def calc_dimer(seq1, seq2, structure=False):
    """
    Calculate the hetero-dimerization thermodynamics of two DNA sequences.
    Return primer3 thermo object.
    """
    return p3.calc_heterodimer(
        seq1.upper(),
        seq2.upper(),
        mv_conc=config.PCR_MV_CONC,
        dv_conc=config.PCR_DV_CONC,
        dna_conc=config.PCR_DNA_CONC,
        dntp_conc=config.PCR_DNTP_CONC,
        output_structure=structure
    )


def has_end_overlap(dimer_result):
    """
    checks if two oligos overlap at their ends
    Example:
        xxxxxxxxtagc-------
        --------atcgxxxxxxx
    """
    if dimer_result.structure_found:
        # clean structure
        structure = [x[4:] for x in dimer_result.ascii_structure_lines]
        # check if we have an overlap that is large enough
        overlap = len(structure[1].replace(" ", ""))
        if overlap <= config.END_OVERLAP:
            return False
        # not more than one conseq. internal mismatch
        if '  ' in structure[1].lstrip(" "):
            return False
        # The alignment length of the ACII structure is equal to the first part of the structure
        # and the maximum possible alignment length is the cumulative length of both primers (-> no overlap at all)
        alignment_length = len(structure[0])
        maximum_alignment_length = len(re.findall("[ATCG]", "".join(structure)))
        # this means that for a perfect end overlap the alignment length is equal to:
        # len(primer1) + len(primer2) - overlap.
        if alignment_length == maximum_alignment_length - overlap:
            return True

    return False


def is_dimer(seq1, seq2):
    """
    check if two sequences dimerize above threshold or are overlapping at their ends
    """
    dimer_result = calc_dimer(seq1, seq2, structure=True)
    # check both the temperature and the deltaG
    if dimer_result.tm > config.PRIMER_MAX_DIMER_TMP or dimer_result.dg < config.PRIMER_MAX_DIMER_DELTAG:
        return True
    # check for perfect end overlaps (this can result in primer extensions even though the tm/dg are okay)
    if has_end_overlap(dimer_result):
        return True

    return False


def calc_max_polyx(seq):
    """
    calculate maximum polyx of a seq
    """
    previous_nuc, counter, max_polyx = seq[0], 1, 1

    for nuc in seq[1:]:
        if nuc == previous_nuc:
            counter += 1
        else:
            counter = 1
            previous_nuc = nuc
        if counter > max_polyx:
            max_polyx = counter

    return max_polyx


def calc_max_dinuc_repeats(seq):
    """
    calculate maximum polyxy of a seq
    """
    for s in [seq, seq[1:]]:
        previous_dinuc = s[0:2]
        max_dinuc = 1
        counter = 1
        for i in range(2, len(s), 2):
            if s[i:i+2] == previous_dinuc:
                counter += 1
            else:
                if counter > max_dinuc:
                    max_dinuc = counter
                counter = 1
                previous_dinuc = s[i:i+2]

    return max_dinuc


def calc_end_gc(seq):
    """
    check how many gc nucleotides
    are within the last 5 bases of
    the 3' end
    """
    return seq[-5:].count('g') + seq[-5:].count('c')


def is_three_prime_ambiguous(amb_seq):
    """
    determine if a sequence contains an ambiguous char at the 3'prime
    """
    len_3_prime, is_amb = config.PRIMER_MIN_3_WITHOUT_AMB, False

    if len_3_prime != 0:
        for nuc in amb_seq[len(amb_seq)-len_3_prime:]:
            if nuc not in config.NUCS:
                is_amb = True
                break

    return is_amb


def rev_complement(seq):
    """
    reverse complement a sequence
    """
    return str(MutableSeq(seq).reverse_complement(inplace=True))


def calc_permutation_penalty(amb_seq):
    """
    get all permutations of a primer with ambiguous
    nucleotides and multiply with permutation penalty
    """
    permutations = 0

    for nuc in amb_seq:
        if nuc in config.AMBIG_NUCS:
            n = len(config.AMBIG_NUCS[nuc])
            if permutations != 0:
                permutations = permutations*n
            else:
                permutations = n

    return permutations*config.PRIMER_PERMUTATION_PENALTY


def calc_base_penalty(seq, primer_temps, gc_range, primer_sizes):
    """
    Calculate intrinsic primer penalty.
    """
    penalty = 0

    tm = calc_temp(seq)
    gc = calc_gc(seq)
    size = len(seq)

    # TEMP penalty
    if tm > primer_temps[2]:
        penalty += config.PRIMER_TM_PENALTY*(
            tm - primer_temps[2]
            )
    if tm < primer_temps[2]:
        penalty += config.PRIMER_TM_PENALTY*(
            primer_temps[2] - tm
            )
    # GC penalty
    if gc > gc_range[2]:
        penalty += config.PRIMER_GC_PENALTY*(
            gc - gc_range[2]
            )
    if gc < gc_range[2]:
        penalty += config.PRIMER_GC_PENALTY*(
            gc_range[2] - gc
        )
    # SIZE penalty
    if size > primer_sizes[2]:
        penalty += config.PRIMER_SIZE_PENALTY*(
            size - primer_sizes[2]
        )
    if size < primer_sizes[2]:
        penalty += config.PRIMER_SIZE_PENALTY * (
            primer_sizes[2] - size
        )

    return penalty


def calc_per_base_mismatches(kmer, alignment, ambiguous_consensus):
    """
    calculate for a given kmer with [seq, start, stop]
    the percent mismatch per kmer pos with the alignment.
    considers if kmer or aln sequences have an amb nuc. returns
    a list of percent mismatches for each kmer position.
    """
    # ini list
    mismatches = len(kmer[0])*[0]
    # get kmer with ambiguous nucs
    amb_kmer = ambiguous_consensus[kmer[1]:kmer[2]]
    # test it against all sequences in the alignment
    for sequence in alignment:
        # slice each sequence of the alignment for the kmer
        # start and stop positions
        seq_slice = sequence[1][kmer[1]:kmer[2]]
        for idx, slice_nuc in enumerate(seq_slice):
            # find the respective nuc to that of the slice
            current_kmer_pos = amb_kmer[idx]
            if slice_nuc == current_kmer_pos:
                continue
            # check if the slice nucleotide is an amb pos
            if slice_nuc in config.AMBIG_NUCS:
                # check if the kmer has an amb pos
                if current_kmer_pos in config.AMBIG_NUCS:
                    slice_nuc_set = set(config.AMBIG_NUCS[slice_nuc])
                    pri_set = set(config.AMBIG_NUCS[current_kmer_pos])
                    # check if these sets have no overlap
                    # -> mismatch
                    if len(slice_nuc_set.intersection(pri_set)) == 0:
                        mismatches[idx] += 1
                # if no amb pos is in kmer then check if kmer nuc
                # is part of the amb slice nuc
                elif current_kmer_pos not in config.AMBIG_NUCS[slice_nuc]:
                    mismatches[idx] += 1
            # check if kmer has an amb pos but the current
            # slice_nuc is not part of this amb nucleotide
            elif current_kmer_pos in config.AMBIG_NUCS:
                if slice_nuc not in config.AMBIG_NUCS[current_kmer_pos]:
                    mismatches[idx] += 1
            # mismatch
            else:
                mismatches[idx] += 1

    # gives a percent mismatch over all positions of the kmer from 5' to 3'
    mismatches = [round(x/len(alignment), 2) for x in mismatches]

    return mismatches


def calc_3_prime_penalty(direction, mismatches):
    """
    calculate the penalty for mismatches at the 3' end.
    the more mismatches are closer to the 3' end of the kmer,
    the higher the penalty. uses the previously calculated
    mismatch list.
    """

    penalty = 0

    if config.PRIMER_3_PENALTY:
        if direction == "-":
            penalty = sum([m * p for m, p in zip(mismatches[0:len(config.PRIMER_3_PENALTY)], config.PRIMER_3_PENALTY)])
        elif direction == "+":
            penalty = sum([m * p for m, p in zip(mismatches[::-1][0:len(config.PRIMER_3_PENALTY)], config.PRIMER_3_PENALTY)])

    return penalty


def filter_kmer_direction_independent(seq, primer_temps=config.PRIMER_TMP, gc_range=config.PRIMER_GC_RANGE, primer_sizes=config.PRIMER_SIZES):
    """
    filter kmer for temperature, gc content,
    poly x, dinucleotide repeats and homodimerization
    """

    return(
        (primer_temps[0] <= calc_temp(seq) <= primer_temps[1])
        and (gc_range[0] <= calc_gc(seq) <= gc_range[1])
        and (calc_max_polyx(seq) <= config.PRIMER_MAX_POLYX)
        and (calc_max_dinuc_repeats(seq) <= config.PRIMER_MAX_DINUC_REPEATS)
        and (calc_base_penalty(seq, primer_temps, gc_range, primer_sizes) <= config.PRIMER_MAX_BASE_PENALTY)
        and not is_dimer(seq, seq)
    )


def filter_kmer_direction_dependend(direction, kmer, ambiguous_consensus):
    """
    filter for 3'ambiguous, hairpin temp and end GC content.
    this differs depending on the direction of the kmer.
    """
    # get the correct amb kmer to test
    if direction == "+":
        kmer_seq = kmer[0]
        amb_kmer_seq = ambiguous_consensus[kmer[1]:kmer[2]]
    elif direction == "-":
        kmer_seq = rev_complement(kmer[0])
        amb_kmer_seq = rev_complement(ambiguous_consensus[kmer[1]:kmer[2]])
    # filter kmer
    return (
        (calc_hairpin(kmer_seq).tm <= config.PRIMER_HAIRPIN)
        and (config.PRIMER_GC_END[0] <= calc_end_gc(kmer_seq) <= config.PRIMER_GC_END[1])
        and not is_three_prime_ambiguous(amb_kmer_seq)
    )


def _process_kmer_batch(ambiguous_consensus, alignment, kmers):
    """
    Helper function for multiprocessing: process a batch of kmers.
    Returns (left_primers, right_primers) tuples.
    """
    left_primers = []
    right_primers = []

    for kmer in kmers:
        if not filter_kmer_direction_independent(kmer[0]):
            continue
        # calc penalties
        base_penalty = calc_base_penalty(kmer[0], config.PRIMER_TMP, config.PRIMER_GC_RANGE, config.PRIMER_SIZES)
        per_base_mismatches = calc_per_base_mismatches(kmer, alignment, ambiguous_consensus)
        permutation_penalty = calc_permutation_penalty(ambiguous_consensus[kmer[1]:kmer[2]])
        # some filters depend on the direction of each primer
        for direction in ["+", "-"]:
            if not filter_kmer_direction_dependend(direction, kmer, ambiguous_consensus):
                continue
            # calc penalties
            three_prime_penalty = calc_3_prime_penalty(direction, per_base_mismatches)
            primer_penalty = base_penalty + permutation_penalty + three_prime_penalty
            # add to lists depending on their direction
            if direction == "+":
                left_primers.append([kmer[0], kmer[1], kmer[2], primer_penalty, per_base_mismatches])
            else:
                right_primers.append([rev_complement(kmer[0]), kmer[1], kmer[2], primer_penalty, per_base_mismatches])

    return left_primers, right_primers


def find_primers(kmers, ambiguous_consensus, alignment, num_processes):
    """
    Filter kmers direction specific and append penalties --> potential primers.
    Uses multiprocessing to process kmers in parallel.
    """
    if not kmers:
        return [], []

    # Convert kmers set to list for slicing
    kmers = list(kmers)
    batch_size = max(1, int(len(kmers)/num_processes))

    # Split kmers into batches
    batches = [kmers[i:i + batch_size] for i in range(0, len(kmers), batch_size)]
    callable_f = functools.partial(
        _process_kmer_batch,
        ambiguous_consensus, alignment
    )

    # Solve dimers in parallel
    with multiprocessing.Pool(processes=num_processes) as pool:
        results = pool.map(callable_f, batches)

    # Aggregate results
    left_primer_candidates = []
    right_primer_candidates = []
    for left_primers, right_primers in results:
        left_primer_candidates.extend(left_primers)
        right_primer_candidates.extend(right_primers)

    return left_primer_candidates, right_primer_candidates


def create_primer_dictionary(primer_candidates, direction):
    """
    creates a primer dictionary from primer list
    """
    primer_dict = {}
    primer_idx = 0

    for primer in primer_candidates:
        if direction == "+":
            direction_name = "LEFT"
        else:
            direction_name = "RIGHT"
        primer_name = f"{direction_name}_{primer_idx}"
        primer_dict[primer_name] = primer
        primer_idx += 1

    return primer_dict


def find_best_primers(left_primer_candidates, right_primer_candidates, high_conservation:bool=False):
    """
    Primer candidates are likely overlapping. Here, the list of primers
    is sorted for the lowest to highest penalty. Then, the next lowest
    is retained if it does not have any nucleotides that have already
    been covered by the middle third of a better primer candidate.
    This significantly reduces the amount of primers while retaining
    the ones with the lowest penalties.

    Example:
    -------- (penalty 1) 1
        ------------- (penalty 1) 2
            ------------ (penalty 0.8) 3
                      ----------(penalty 0.9) 4

    --> primer 3 would be retained and primer 2 excluded, primer 4 and 1
    will however be considered in the next set of overlapping primers.

    """
    all_primers = {}

    for direction, primer_candidates in [("+", left_primer_candidates), ("-", right_primer_candidates)]:
        # sort the primers by penalty, and if same penalty by start
        primer_candidates.sort(key=lambda x: (x[3], x[1]))
        # ini everything with the primer with the lowest penalty
        to_retain = [primer_candidates[0]]
        primer_set = set(range(primer_candidates[0][1], primer_candidates[0][2]))

        for primer in primer_candidates[1:]:
            # for highly conserved alignments exclude everything that overlaps with the best primer
            # this reduces graph complexity by quite a large margin
            if high_conservation:
                primer_positions =set(range(primer[1], primer[2]))
            # get the thirds of the primer, only consider the middle
            else:
                thirds_len = int((primer[2] - primer[1])/3)
                primer_positions = set(range(primer[1] + thirds_len, primer[2] - thirds_len))
            # check if none of the nucleotides of the next primer
            # are already covered by a better primer
            if primer_set.isdisjoint(primer_positions):
                # update the primer set
                primer_set.update(primer_positions)
                # append this primer as it has a low penalty and is not overlapping
                # with another already retained primer
                to_retain.append(primer)

        # sort by start
        to_retain.sort(key=lambda x: x[1])
        # create dict
        all_primers[direction] = create_primer_dictionary(to_retain, direction)

    # and create a dict
    return all_primers


def get_permutations(seq):
    """
    get all permutations of an ambiguous sequence.
    """
    splits = [config.AMBIG_NUCS.get(nuc, [nuc]) for nuc in seq]

    return[''.join(p) for p in itertools.product(*splits)]


def parse_primer_fasta(fasta_path):
    """
    Parse a primer FASTA file and return a list of sequences using BioPython.
    """

    sequences = []

    for record in SeqIO.parse(fasta_path, "fasta"):
        seq = str(record.seq).lower()
        # Only include primers up to 40 nucleotides
        if len(seq) <= 40:
            sequences += get_permutations(seq)

    return list(set(sequences))  # deduplication


def check_primer_against_externals(external_sequences, primer):
    """
    Worker function to check a single primer against all external sequences.
    Returns the primer if it passes, None otherwise.
    Handles both list format and dict format (name, data) tuples.
    """

    # Extract sequence based on input format
    if isinstance(primer, tuple):
        name, data = primer
        seq = data[0]
    else:
        seq = primer[0]

    for ext_seq in external_sequences:
        if is_dimer(seq, ext_seq):
            return None

    return primer


def filter_non_dimer_candidates(primer_candidates, external_sequences, n_processes):
    """
    Filter out primer candidates that form dimers with external sequences.
    Uses multiprocessing to speed up checks.
    """
    is_dict = isinstance(primer_candidates, dict)

    callable_f = functools.partial(
        check_primer_against_externals,
        external_sequences
    )

    with multiprocessing.Pool(processes=n_processes) as pool:
        # Prepare arguments based on input type
        # qpcr probes are stored in dictionaries --> result in tuples when unpacked
        if is_dict:
            results = pool.map(callable_f, primer_candidates.items())
        else:
            results = pool.map(callable_f, primer_candidates)

    # Filter and restore original format
    if is_dict:
        filtered_results = [result for result in results if result is not None]
        return {name: data for name, data in filtered_results}
    else:
        return [primer for primer in results if primer is not None]
