import numpy as np
from dnachisel.biotools import reverse_complement
from Bio import Restriction
import proglog

from .tools import (
    enzymes_to_dna_pattern,
    get_enzymes_ATGC_sequences,
    find_patterns_matches,
    reverse_complement,
)


def sites_difference(site1, site2):
    """Return minimal sequence of nucleotides that should be added at the end
    of site1 to make site2 appear."""
    for i in range(len(site2), -1, -1):
        if site2[:i] == site1[-i:]:
            return site2[i:]
    return site2


def _enzymes_names_to_distances_graph(enzymes_names):
    enzymes_names = enzymes_names[:]
    np.random.shuffle(enzymes_names)
    enzymes_sites = get_enzymes_ATGC_sequences(enzymes_names)
    patterns = enzymes_to_dna_pattern(enzymes_names)
    core_enzymes = {}
    for e1 in list(enzymes_names):
        site1 = enzymes_sites[e1]
        rev_site1 = reverse_complement(site1)
        for e2 in list(enzymes_names):
            if e1 == e2:
                continue
            site2 = enzymes_sites[e2]
            pattern2 = patterns[e2]
            if any(
                [
                    site1 == site2,
                    site2 in site1,
                    rev_site1 == site2,
                    site2 in rev_site1,
                    len(pattern2.find_matches(site1)),
                    len(pattern2.find_matches(rev_site1)),
                ]
            ):
                if e1 not in core_enzymes:
                    core_enzymes[e1] = []
                if e2 not in core_enzymes[e1]:
                    core_enzymes[e1].append(e2)
    graph = {}
    for e1 in enzymes_names:
        site1 = enzymes_sites[e1]
        for e2 in enzymes_names:
            if e1 == e2:
                continue
            site2 = enzymes_sites[e2]
            diff = sites_difference(site1, site2)
            graph[(e1, e2)] = dict(diff=diff, dist=len(diff))
    return graph, enzymes_sites, core_enzymes


def _enzyme_path_to_sequence(path, graph, enzymes_sites):
    """Convert a path of successive enzymes into a sequence."""
    return "".join(
        [enzymes_sites[path[0]]]
        + [graph[(n1, n2)]["diff"] for n1, n2 in zip(path, path[1:])]
    )


class NoSuitableSequenceFound(Exception):
    pass


def stacked_sites_array(
    enzymes_names,
    forbidden_enzymes=(),
    unique_sites=True,
    tries=10,
    logger="bar",
    success_condition=None,
):
    """Generate a sequence.

    Parameters
    ----------
    enzymes_names: list
        Names of enzyme sites which the algorithm should try to include in the
        sequence.

    forbidden_enzymes: list
        List of enzymes names whose sites should not be in the sequence.

    unique_sites: bool
        If True, all the enzymes in enzymes_names will have no more than one site
        in the final sequence.

    tries: int
        Number of tries. The sequence returned is the one with the best score,
        i.e., with the least enzymes left over, or the shortest sequence
        in case of equality.

    success_condition: function
        A function evaluated at the end of each try to validate the designed sequence.
        Should be of the form ``f(seq, inseq, notinseq) => True/False`` where ``seq``
        is the final sequence, ``inseq`` the set of enzymes in the sequence,
        and ``notinseq`` the set of enzymes not in the sequence.

    logger: {'bar', None} or Proglog progress logger, optional
        Either "bar" for a progress bar, None for no progress logger,
        or any Proglog progress logger.

    Returns
    -------
    tuple
        A triplet (sequence, enzymes_in_sequence, enzymes_not_in_sequence).
    """

    enzymes_names = sorted(set(enzymes_names))
    logger = proglog.default_bar_logger(logger, min_time_interval=0.2)
    patterns = enzymes_to_dna_pattern(list(enzymes_names) + list(forbidden_enzymes))
    graph, enzymes_sites, core_enzymes = _enzymes_names_to_distances_graph(
        enzymes_names
    )

    def one_try():
        pickable_enzymes_not_in_seq = set(enzymes_names)
        enzymes_not_in_seq = set(enzymes_names)
        enzymes_in_seq = set()

        def add_enzyme(e, as_core=True):
            if e not in enzymes_in_seq:
                if e in pickable_enzymes_not_in_seq:
                    pickable_enzymes_not_in_seq.remove(e)
                enzymes_not_in_seq.remove(e)
                enzymes_in_seq.add(e)
                if as_core:
                    for enz in core_enzymes.get(e, []):
                        add_enzyme(enz, as_core=False)

        list_pickable_enzymes_not_in_seq = list(pickable_enzymes_not_in_seq)
        path = [
            list_pickable_enzymes_not_in_seq[
                np.random.randint(len(list_pickable_enzymes_not_in_seq))
            ]
        ]
        seq = _enzyme_path_to_sequence(path, graph, enzymes_sites)
        add_enzyme(path[-1])
        while len(pickable_enzymes_not_in_seq):
            nodes_distances = sorted(
                [
                    (graph[(path[-1], n)]["dist"], n)
                    for n in pickable_enzymes_not_in_seq
                ],
                key=lambda dn: (dn[0], np.random.rand()),
            )
            choices = [n for d, n in nodes_distances]
            for new_node in choices:
                new_path = path + [new_node]
                new_seq = _enzyme_path_to_sequence(new_path, graph, enzymes_sites)
                new_diff = len(new_seq) - len(seq)
                end_seq = new_seq[-(6 + new_diff) :]
                matches = find_patterns_matches(end_seq, patterns)
                if unique_sites:
                    new_matches = {
                        e: [m for m in matches[e] if m.end > len(end_seq) - new_diff]
                        for e in matches
                    }
                    forbidden = list(enzymes_in_seq) + list(forbidden_enzymes)
                    if any([len(new_matches[e]) for e in forbidden]):
                        continue
                    if any([len(new_matches[e]) > 1 for e in enzymes_not_in_seq]):
                        continue
                seq = new_seq
                path = new_path
                add_enzyme(new_node)
                for e in list(enzymes_not_in_seq):
                    if len(matches[e]) > 0:
                        add_enzyme(e)
                break
            else:
                return seq, enzymes_in_seq, enzymes_not_in_seq
        return seq, enzymes_in_seq, enzymes_not_in_seq

    failure_score = (10000, 10000)

    def score(seq, in_seq, leftover):
        if unique_sites or len(forbidden_enzymes):
            matches = find_patterns_matches(seq, patterns)
            if unique_sites and any([len(matches[e]) != 1 for e in in_seq]):
                return failure_score
            if any([len(matches[e]) for e in forbidden_enzymes]):
                return failure_score
        if success_condition is not None:
            if not success_condition(seq, in_seq, leftover):
                return failure_score
        return (len(leftover), len(seq))

    current_best = seq, in_seq, leftover = one_try()
    current_score = score(seq, in_seq, leftover)
    for _ in logger.iter_bar(tries=range(tries)):
        new_result = seq, in_seq, leftover = one_try()
        new_score = score(seq, in_seq, leftover)
        if new_score < current_score:
            current_score = new_score
            current_best = new_result
    if current_score == failure_score:
        raise NoSuitableSequenceFound("Could not find a suitable sequence.")
    return current_best
