"""Module for calculating the complexity of a DNASequence.

Complexity rules borrowed from IDT.

* less than 25% GC content
* higher than 75% GC content
* homopolymeric runs of 10 or more A/Ts
* homopolymeric runs of 6 or more G/Cs
* repeats greater than 14bp
* hairpins
* GC scew?
* windowed scew?
"""
import hashlib
import json
from copy import copy as do_copy
from copy import deepcopy as do_deepcopy
from functools import lru_cache

import numpy as np


class DNAStatsMeta:
    def __init__(self):
        self.fwd_signatures = None
        self.rev_signatures = None
        self.rolling_stats = None
        self.repeat_signatures = None


# TODO:' compute cost for extremes of GC content
# TODO: this could be a separate library
class DNAStats:
    """
    Class for computing statistics for DNA Sequences.

    ::

        {
            "n_repeats",
            "n_hairpins",
            "window_cost",
            "gc_cost"
        }
    """

    BASES = "AGCT"
    np.random.seed(1)
    BASE_SIGNATURES = np.random.uniform(0.0, 1.0, size=len(BASES)).reshape(-1, 4)

    DEFAULT_MODE = ("window", "hairpin", "repeat")
    ONLY_WINDOW = ("window",)
    ONLY_HAIRPIN = ("hairpin",)
    ONLY_REPEAT = "repeat"

    __slicable__ = [
        "fwd_signatures",
        "rev_signatures",
        "repeat_signatures",
        "seq",
        "seq_onehot",
        "rolling_stats",
    ]

    __copyable__ = [
        "bases",
        "repeat_window",
        "stats_window",
        "hairpin_window",
        "gc_content_threshold",
        "at_content_threshold",
        "base_percentage_threshold",
        "conv_seed_repeat",
        "mode",
        "conv_seed_hairpin",
    ]

    __slots__ = __copyable__ + __slicable__

    def __init__(
        self,
        seq,
        repeat_window,
        stats_window,
        hairpin_window,
        base_percentage_threshold: float = 0.75,
        gc_content_threshold: float = 0.85,
        at_content_threshold: float = 0.90,
        conv_seed_repeat: np.array = None,
        conv_seed_hairpin: np.array = None,
        mode: str = None,
    ):
        """

        :param seq:
        :param repeat_window: the length of the k-mers to search for repeats.
        :param stats_window: the width of the stats window
        :param hairpin_window: length of hte k-mers to search for hairpins.
        :param base_percentage_threshold: if the per base percentage threshold
            to count in slinding window calculatiions. For example,
            for a sliding window of 20, if the threshold is 0.75, if 15 bases
            or of a single base pair, this will get counted in the `window_cost`
            method.
        :param gc_content_threshold: if the gc percentage threshold
            to count in sliding window calculatiions. For example,
            for a sliding window of 20, if the threshold is 0.75, if 15 bases
            of any window of 20 were gc, this will get counted in the `window_cost`
            method.
        :param conv_seed_repeat: (optional) the array to convolve the one-hot
            encoded DNA sequence with. If not provided, will be generated by
            `np.random.uniform(0.0, 100.0, size=repeat_window)`
        :param conv_seed_hairpin: (optional) the array to convolve the one-hot
            encoded DNA sequence with to find hairpin signatures.
            If not provided, will be generated by
            `np.random.uniform(0.0, 100.0, size=hairpin_window)`
        :param mode: (optional) Set to DNAStats.ONLY_HAIRPIN to only
            calculate hairpin signatures, DNAStats.ONLY_REPEATS to only calculate
            repeat signatures. DNAStats.ONLY_WINDOW to perform only
            sliding window calculations.
        """
        if not isinstance(seq, str):
            raise ValueError("Expected a str not a '{}'".format(type(seq)))
        seq = seq.upper()
        self.seq = seq
        arr = np.array(list(self.seq))
        self.bases = self.BASES.upper()
        self.seq_onehot = self.one_hot(arr, self.bases)
        self.repeat_window = repeat_window
        self.stats_window = stats_window
        self.hairpin_window = hairpin_window
        self.gc_content_threshold = gc_content_threshold
        self.at_content_threshold = at_content_threshold
        self.base_percentage_threshold = base_percentage_threshold
        if not mode:
            mode = self.DEFAULT_MODE
        self.mode = mode
        if conv_seed_repeat is not None:
            self.conv_seed_repeat = conv_seed_repeat
        else:
            self.conv_seed_repeat = np.random.uniform(0.0, 100.0, size=repeat_window)
        if conv_seed_hairpin is not None:
            self.conv_seed_hairpin = conv_seed_hairpin
        else:
            self.conv_seed_hairpin = np.random.uniform(0.0, 100.0, size=hairpin_window)

        if "window" in self.mode:
            rolling_stats = self.get_base_stats(stats_window)
            gc_content = (
                rolling_stats[self.BASES.index("G"), :]
                + rolling_stats[self.BASES.index("C"), :]
            )
            at_content = (
                rolling_stats[self.BASES.index("A"), :]
                + rolling_stats[self.BASES.index("T"), :]
            )
            self.rolling_stats = np.vstack((rolling_stats, gc_content, at_content))
        else:
            self.rolling_stats = None

        ###
        # get repeat signatures
        ###
        seq_signatures = np.sum(
            np.multiply(self.seq_onehot, self.BASE_SIGNATURES.T), axis=0
        )
        if "repeat" in self.mode:
            mv = np.convolve(seq_signatures, self.conv_seed_repeat, mode="valid")
            self.repeat_signatures = np.concatenate(
                ([np.NaN for _ in range(repeat_window - 1)], mv)
            )
        else:
            self.repeat_signatures = None

        ###
        # get hairpin signatures
        ###
        if "hairpin" in self.mode:
            revcomp_signatures = np.sum(
                np.multiply(self.seq_onehot[:, ::-1], self.BASE_SIGNATURES[:, ::-1].T),
                axis=0,
            )
            self.fwd_signatures = np.convolve(
                seq_signatures, self.conv_seed_hairpin, mode="valid"
            )
            self.rev_signatures = np.convolve(
                revcomp_signatures, self.conv_seed_hairpin, mode="valid"
            )[::-1]
        else:
            self.fwd_signatures = None
            self.rev_signatures = None

    def _view(self, index: slice = None, copy: bool = False, deepcopy: bool = False):
        x = object.__new__(self.__class__)
        for k in self.__slots__:
            v = getattr(self, k)
            if k in self.__slicable__:
                if copy or deepcopy:
                    if isinstance(v, np.ndarray):
                        v = v.copy()
                    else:
                        v = do_copy(v)
                if index:
                    if isinstance(v, np.ndarray):
                        if v.ndim == 1:
                            v = v[index]
                        else:
                            v = v[:, index]
                    else:
                        v = v[index]
            elif k in self.__copyable__:
                if deepcopy:
                    v = do_deepcopy(getattr(self, k))
                else:
                    v = do_copy(getattr(self, k))
            setattr(x, k, v)
        return x

    def copy(self, index: slice = None):
        return self._view(index=index, copy=True)

    def deepcopy(self, index: slice = None):
        return self._view(index=index, deepcopy=True)

    def view(self, index: slice = None):
        return self._view(index=index)

    @property
    def config(self):
        return dict(
            repeat_window=self.repeat_window,
            stats_window=self.stats_window,
            hairpin_window=self.hairpin_window,
            base_percentage_threshold=self.base_percentage_threshold,
            gc_content_threshold=self.gc_content_threshold,
            at_content_threshold=self.at_content_threshold,
            mode=self.mode,
        )

    def copy_with_new_seq(self, seq: str):
        return self.__class__(
            seq,
            repeat_window=self.repeat_window,
            stats_window=self.stats_window,
            hairpin_window=self.hairpin_window,
            base_percentage_threshold=self.base_percentage_threshold,
            gc_content_threshold=self.gc_content_threshold,
            at_content_threshold=self.at_content_threshold,
            conv_seed_repeat=self.conv_seed_repeat,
            conv_seed_hairpin=self.conv_seed_hairpin,
            mode=self.mode,
        )

    @staticmethod
    def one_hot(sequence, categories):
        arrs = []
        for i, c in enumerate(categories):
            a = sequence == c
            a = a.astype(int)
            arrs.append(a)
        return np.vstack(arrs)

    @staticmethod
    def rolling_window(a, window):
        shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
        strides = a.strides + (a.strides[-1],)
        return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

    @staticmethod
    def rolling_average(x, n):
        mv = np.convolve(x, np.ones(n) / n, mode="valid")
        return np.concatenate(([np.NaN for k in range(n - 1)], mv))

    def get_base_stats(self, window):
        return np.mean(self.rolling_window(self.seq_onehot, window), axis=2)

    def get_repeat_signatures(self, window):
        seq_signatures = np.sum(
            np.multiply(self.seq_onehot, self.BASE_SIGNATURES.T), axis=0
        )
        x = np.random.uniform(0.0, 100.0, size=window)
        mv = np.convolve(seq_signatures, x, mode="valid")
        signatures = np.concatenate(([np.NaN for _ in range(window - 1)], mv))
        return signatures

    def get_hairpin_signatures(self, window):
        seq_signatures = np.sum(
            np.multiply(self.seq_onehot, self.BASE_SIGNATURES.T), axis=0
        )
        revcomp_signatures = np.sum(
            np.multiply(self.seq_onehot[:, ::-1], self.BASE_SIGNATURES[:, ::-1].T),
            axis=0,
        )
        x = np.random.uniform(0.0, 100.0, size=window)
        mv1 = np.convolve(seq_signatures, x, mode="valid")
        mv2 = np.convolve(revcomp_signatures, x, mode="valid")
        return mv1, mv2[::-1]

    @staticmethod
    def count_signatures(a):
        return len(np.where(np.unique(a, return_counts=True)[1] > 1)[0])

    def slice_repeats(self, i, j):
        return self.count_signatures(self.repeat_signatures[i:j])

    def slice_hairpins(self, i, j):
        f = self.fwd_signatures[i:j]
        r = self.rev_signatures[i:j]
        d = np.concatenate((f, r))
        num_hairpins = self.count_signatures(d) - 2 * self.count_signatures(f)
        return num_hairpins

    def window_cost(
        self,
        i,
        j,
        base_threshold_perc: float = None,
        gc_content_threshold: float = None,
        at_content_threshold: float = None,
    ):

        if base_threshold_perc is None:
            base_threshold_perc = self.base_percentage_threshold
        if gc_content_threshold is None:
            gc_content_threshold = self.gc_content_threshold

        if at_content_threshold is None:
            at_content_threshold = self.at_content_threshold
        d = self.rolling_stats[:4, i:j]

        # keep each base below 0.8
        a = np.sum(d > base_threshold_perc, axis=1)

        # keep gc content below extremes 0.7
        b = np.sum(self.rolling_stats[4, i:j] > gc_content_threshold)

        # keep gc content below extremes 0.7
        c = np.sum(self.rolling_stats[5, i:j] > at_content_threshold)
        return a, b, c

    # def get_GC_content_complexity(seq):
    #     gc = get_GC_content(seq)
    #     return abs(gc * 100.0 - 50) * 17 / 25.0

    @staticmethod
    def _optimize_partition_helper(
        signatures: np.ndarray, step: int, i: int = None, j: int = None
    ):
        """Optimize partition by minimizing the number of signatures in the
        given array.

        :param signatures: array of signatures
        :param step: step size
        :param i:
        :param j:
        :return:
        """
        d = []

        if i is None:
            i = 0
        if j is None:
            j = signatures.shape[1]

        for x in range(i, j, step):
            m1 = np.empty(signatures.shape[1])
            m2 = m1.copy()
            m1.fill(np.nan)
            m2.fill(np.nan)

            m1[:x] = np.random.uniform(1, 10)
            m2[x:] = np.random.uniform(1, 10)

            d += [m1, m2]
        d = np.vstack(d)
        z = np.tile(d, signatures.shape[0]) * signatures.flatten()

        partition_index = np.repeat(
            np.arange(0, signatures.shape[1], step),
            signatures.shape[0] * signatures.shape[1] * 2,
        )

        a, b, c = np.unique(z, return_counts=True, return_index=True)
        i = b[np.where(c > 1)]
        a, c = np.unique(partition_index[i], return_counts=True)
        if len(c):
            arg = c.argmin()
            return a[arg], c[arg]

    def count_repeats_from_slice(self, i, j):
        """Counts the number of repeats found in slices [0, i) and [j, None) of
        all kmers in the sequence found within the [i, j) slice.

        .. note::
            The hairpin_window_length of the DNAStats instance is used as the kmer
            length.

        :param i: start of the slice (inclusive)
        :param j: end of the slice (exclusive)
        :return: tuple of repeat counts found in left and right slices respectively.
        """
        if i < 0:
            raise IndexError("i={} cannot be < 0")
        if j > len(self.seq):
            raise IndexError("j={} cannot be >= length {}".format(j, len(self.seq)))
        kmer_length = self.hairpin_window

        # we ust signatures (src) of kmers inside i:j...
        src_indices = np.arange(i, j - kmer_length + 1)

        # too look for signatures (dest) in the left and right flanks
        left_flank_indices = np.arange(0, i)
        right_flank_indices = np.arange(j, len(self.seq) - kmer_length)
        dest_indices = np.concatenate([left_flank_indices, right_flank_indices])

        src = self.fwd_signatures[src_indices]

        # TODO: unsure why the rev indices need to be shifted...
        src_rc = self.rev_signatures[src_indices]
        dest = self.fwd_signatures[dest_indices]

        def _mispriming_counts(src_sig, src_sig_rc, dest_sig):
            # 5' end, top repeats to self
            counts_trim_top = self.count_signatures(src_sig)

            # 5' end, calculate priming to top strand
            x, counts = np.unique(src_sig, return_counts=True)
            top_counts = counts[np.where(np.isin(x, dest_sig))].sum()

            # 5' end, calculate priming to bottom strand
            x, counts = np.unique(src_sig_rc, return_counts=True)
            bottom_counts = (
                counts[np.where(np.isin(x, dest_sig))].sum() + (counts - 1).sum()
            )

            return top_counts + bottom_counts + counts_trim_top

        return _mispriming_counts(src, src_rc, dest)

    def cost(self, i=None, j=None):
        d = self(i, j)
        w1 = d["window_cost"][0]
        w2 = d["window_cost"][1]
        total = (
            d["n_repeats"] + d["n_hairpins"] + np.sum(w1) + np.sum(w2) + d["gc_cost"]
        )
        return total

    @staticmethod
    def _hash(string: str):
        return hashlib.sha1(string.encode("utf-8")).hexdigest()

    def __hash__(self):
        config = json.dumps(self.config)
        s = self.seq.upper() + config
        return int(hashlib.sha1(s.encode("utf-8")).hexdigest(), 16)

    def gc_cost(self, i, j):
        mn = np.mean(self.seq_onehot[:, i:j], axis=1)
        gi = self.bases.index("G")
        ci = self.bases.index("C")
        gc = mn[gi] + mn[ci]
        gccost = abs(gc * 100.0 - 50) * 17 / 25.0
        return gccost

    def __call__(self, i=None, j=None):
        return {
            "n_repeats": self.slice_repeats(i, j),
            "n_hairpins": self.slice_hairpins(i, j),
            "window_cost": self.window_cost(i, j),
            "gc_cost": self.gc_cost(i, j),
        }

    def __len__(self):
        return len(self.seq)

    def __getitem__(self, item):
        return self.view(item)


@lru_cache(512)
def cached_stats(seq: str, hairpin_window: int):
    return DNAStats(seq, 1, 1, hairpin_window, mode=DNAStats.ONLY_HAIRPIN)


@lru_cache(512)
def count_misprimings_in_amplicon(
    seq,
    i,
    j,
    min_primer_anneal: int = 12,
    max_primer_anneal: int = 30,
    cyclic: bool = False,
    max_length: int = 3000,
):
    """Counts the estimated number of misprimings in an amplicon.

    :param seq: The template sequence
    :param i: amplicon start (inclusive)
    :param j: amplicon end (exclusive)
    :param min_primer_anneal: The minimum number of bases to look for misprimings
        (exact matches only).
    :param max_primer_anneal: The maximum size of the annealing sequence of the primers.
        Internally, this takes a window slice from the ends of the amplicon to look
        for repeats in the template.
    :param cyclic: If True, assumes the sequence is cyclic.
    :param max_length: The maximum length to consider a mispriming. Generally,
        this should be about the bp distinguishable in your fragment purification
        (e.g. gel)
    :return:
    """
    if j == i:
        return 0
    if i < 0:
        raise IndexError("i cannot be < 0")
    if j < 0:
        raise IndexError("j cannot be < 0")
    elif j < i and not cyclic:
        raise IndexError(
            "Invalid indices provided. "
            "Indices [{i},j) indicates a cyclic sequence, but "
            "cyclic was set to False."
        )
    elif cyclic:
        seq = seq[i:] + seq[:i]
        if j < i:
            length = len(seq) - i + j
            i = 0
            j = length
        else:
            j = j - i
            i = 0

        # TODO: trim extra long sequences
        # trim extra long sequences
        # if len(seq) - j > 2 * max_length:
        #     seq = seq[:j] + seq[j:][max_length] + seq[-max_length:]
    # else:
    #     # trim extra long sequences
    #     if i > max_length:
    #         delta = max_length - i
    #     else:
    #         delta = 0
    #     seq = seq[delta:i] + seq[i:j] + seq[j:][:max_length]
    #     if i > max_length:
    #         delta = max_length - i
    #         i -= delta
    #         j -= delta
    stats = cached_stats(seq, min_primer_anneal)
    n1 = stats.count_repeats_from_slice(i, min(i + max_primer_anneal, len(seq) - 1))
    k = j - max_primer_anneal
    if k < 0:
        n2 = 0
    else:
        n2 = stats.count_repeats_from_slice(max(j - max_primer_anneal, 0), j)
    return n1 + n2
