import itertools as it
import collections as co
import json
import autograd.numpy as np
import numpy as raw_np
from cached_property import cached_property
import scipy
import scipy.sparse
import gzip
import os
import logging
from .compressed_counts import _hashed2config, _config2hashable
from .compressed_counts import CompressedAlleleCounts
from .configurations import ConfigList
from .configurations import _ConfigList_Subset
from ..util import memoize_instance


def site_freq_spectrum(sampled_pops, freqs_by_locus, length=None):
    """Create :class:`Sfs` from a list of dicts of counts.

    ``freqs_by_locus[l][config]`` gives the frequency of ``config`` \
    on locus ``l``.

    A ``config`` is a ``(n_pops,2)``-array, where ``config[p][a]`` \
    is the count of allele ``a`` in population ``p``. ``a=0`` represents \
    the ancestral allele while ``a=1`` represents the derived allele.

    :param list sampled_pops: list of population names
    :param list freqs_by_locus: list of dict
    :param float,None length: Number of bases in the dataset. Set to ``None`` if unknown.
    :rtype: :class:`Sfs`
    """
    index2loc = []
    index2count = []
    loci_counters = []

    def chained_sequences():
        for loc, locus_iter in enumerate(freqs_by_locus):
            loci_counters.append(co.Counter())
            try:
                locus_iter = locus_iter.items()
            except AttributeError:
                locus_iter = zip(locus_iter, iter(lambda: 1, 0))
            for config, count in locus_iter:
                index2loc.append(loc)
                index2count.append(count)
                yield config

    compressed_counts = CompressedAlleleCounts.from_iter(
        chained_sequences(), len(sampled_pops))
    config_array = compressed_counts.config_array
    index2uniq = compressed_counts.index2uniq

    assert len(index2loc) == len(index2count) and len(
        index2count) == len(index2uniq)
    for loc, count, uniq in zip(index2loc, index2count, index2uniq):
        loci_counters[loc][uniq] += count

    configs = ConfigList(sampled_pops, config_array, None, None)
    return Sfs(loci_counters, configs, folded=False, length=length)



class Sfs(object):
    """
    Class for representing observed site frequnecy spectrum data.

    To create a :class:`Sfs` object, use \
    :meth:`SnpAlleleCounts.extract_sfs`, \
    :meth:`Sfs.load`, or :func:`site_freq_spectrum`. \
    Do NOT call the class constructor directly, it is for internal use only.
    """
    @classmethod
    def from_matrix(cls, mat, configs, *args, **kwargs):
        # convert to csc
        mat = scipy.sparse.csc_matrix(mat)
        indptr = mat.indptr
        loci = []
        for i in range(mat.shape[1]):
            loci.append(np.array([
                mat.indices[indptr[i]:indptr[i+1]],
                mat.data[indptr[i]:indptr[i+1]]]))

        return cls(loci, configs, *args, **kwargs)

    @classmethod
    def load(cls, f):
        """Load :class:`Sfs` from file created by :meth:`Sfs.dump` or ``python -m momi.extract_sfs``

        :param str,file f: file object or file name to read in
        :rtype: :class:`Sfs`
        """
        if isinstance(f, str):
            fname = os.path.expanduser(f)
            if fname.endswith(".gz"):
                with gzip.open(fname, "rt") as f:
                    return cls.load(f)
            else:
                with open(fname) as f:
                    return cls.load(f)

        logging.getLogger(__name__).info("Reading json...")
        info = json.load(f)
        logging.getLogger(__name__).info("Finished reading json...")

        logging.getLogger(__name__).info("Constructing configs...")
        configs = info.pop("configs")
        logging.getLogger(__name__).info("{} unique configs detected".format(len(configs)))
        # don't use autograd here for performance
        configs = raw_np.array(configs, dtype=int)
        logging.getLogger(__name__).info("Converted to numpy array...")
        configs = ConfigList(info.pop("sampled_pops"), configs)
        logging.getLogger(__name__).info("Finished constructing configs")

        logging.getLogger(__name__).info("Constructing SFS...")
        loci = []
        for locus, locus_rows in it.groupby(
                info.pop("(locus,config_id,count)"), lambda x: x[0]):
            loci.append({config_id: count
                        for _, config_id, count in locus_rows})

        ret = cls(loci, configs, **info)
        logging.getLogger(__name__).info("Finished constructing SFS")

        return ret

    def __init__(self, loci, configs, folded, length):
        self.folded = folded
        self._length = length

        self.configs = configs

        self.loc_idxs, self.loc_counts = [], []
        for loc in loci:
            if len(loc) == 0:
                self.loc_idxs.append(np.array([], dtype=int))
                self.loc_counts.append(np.array([], dtype=float))
            else:
                try:
                    loc.items()
                except:
                    loc = np.array(loc)
                    if len(loc.shape) == 2:
                        assert loc.shape[0] == 2
                        idxs, cnts = loc[0, :], loc[1, :]
                    else:
                        idxs, cnts = np.unique(loc, return_counts=True)
                else:
                    idxs, cnts = zip(*loc.items())
                self.loc_idxs.append(np.array(idxs, dtype=int))
                self.loc_counts.append(np.array(cnts, dtype=float))

        if len(self.loc_idxs) > 1:
            self._total_freqs = self.freqs_matrix.dot(np.ones(self.n_loci))
            assert self._total_freqs.shape == (self.freqs_matrix.shape[0],)
        else:
            # avoid costly building of frequency matrix, when there are many
            # Sfs's of a single locus (e.g. in many stochastic minibatches)
            idxs, = self.loc_idxs
            cnts, = self.loc_counts
            self._total_freqs = np.zeros(len(self.configs))
            self._total_freqs[idxs] = cnts

        assert not np.any(self._total_freqs == 0)

    def dump(self, f):
        """Write Sfs to file

        :param str,file f: Filename or object. If name ends with ".gz" gzip it.
        """
        if isinstance(f, str):
            fname = os.path.expanduser(f)
            if fname.endswith(".gz"):
                with gzip.open(fname, "wt") as f:
                    self.dump(f)
            else:
                with open(fname, 'w') as f:
                    self.dump(f)
            return

        print("{", file=f)
        print('\t"sampled_pops": {},'.format(
            json.dumps(list(self.sampled_pops))), file=f)
        print('\t"folded": {},'.format(
            json.dumps(self.folded)), file=f)
        print('\t"length": {},'.format(
            json.dumps(self._length)), file=f)
        print('\t"configs": [', file=f)
        n_configs = len(self.configs.value)
        for i, c in enumerate(self.configs.value.tolist()):
            if i != n_configs - 1:
                print("\t\t{},".format(c), file=f)
            else:
                # no trailing comma
                print("\t\t{}".format(c), file=f)
        print("\t],", file=f)
        print('\t"(locus,config_id,count)": [', file=f)
        for loc in range(self.n_loci):
            n_loc_rows = len(self.loc_idxs[loc])
            assert n_loc_rows == len(self.loc_counts[loc])
            for i in range(n_loc_rows):
                config_id = int(self.loc_idxs[loc][i])
                count = float(self.loc_counts[loc][i])
                if loc == self.n_loci - 1 and i == n_loc_rows - 1:
                    # no trailing comma
                    print("\t\t{}".format(json.dumps(
                        (loc, config_id, count))), file=f)
                else:
                    print("\t\t{},".format(json.dumps(
                        (loc, config_id, count))), file=f)
        print("\t]", file=f)
        print("}", file=f)

    @property
    def populations(self):
        return self.sampled_pops

    @memoize_instance
    def combine_loci(self):
        # return copy with all loci combined
        return self.from_matrix(self.freqs_matrix.sum(axis=1),
                                self.configs, self.folded,
                                self._length)

    @property
    def freqs_matrix(self):
        """Sparse matrix representing the frequencies at each locus.

        The ``(i,j)``-th entry gives the frequency of the ``i``-th config in :attr:`Sfs.config_array` at locus ``j``

        :rtype: :class:`scipy.sparse.csr_matrix`
        """
        return self.csr_freqs_matrix

    @cached_property
    def csr_freqs_matrix(self):
        return _csr_freq_matrix_from_counters(
            self.loc_idxs, self.loc_counts, len(self.configs))

    @cached_property
    def avg_pairwise_hets(self):
        # avg number of hets per ind per pop (assuming Hardy-Weinberg)
        n_nonmissing = np.sum(self.configs.value, axis=2)
        # for denominator, assume 1 allele is drawn from whole sample, and 1
        # allele is drawn only from nomissing alleles
        denoms = np.maximum(n_nonmissing * (self.sampled_n - 1), 1.0)
        p_het = 2 * self.configs.value[:, :, 0] * \
            self.configs.value[:, :, 1] / denoms

        return self.freqs_matrix.T.dot(p_het)

    def resample(self):
        """Create a new SFS by resampling blocks with replacement.

        Note the resampled SFS is assumed to have the same length in base pairs \
        as the original SFS, which may be a poor assumption if the blocks are not of equal length.

        :returns: Resampled SFS
        :rtype: :class:`Sfs`
        """
        loci = np.random.choice(
            np.arange(self.n_loci), size=self.n_loci, replace=True)
        mat = self.freqs_matrix[:, loci]
        to_keep = np.asarray(mat.sum(axis=1) > 0).squeeze()
        to_keep = np.arange(len(self.configs))[to_keep]
        mat = mat[to_keep, :]
        configs = _ConfigList_Subset(self.configs, to_keep)

        return self.from_matrix(mat, configs, self.folded, self.length)

    @property
    def sampled_n(self):
        """Number of samples per population

        ``sampled_n[i]`` is the number of haploids \
        in the ``i``-th population of :attr:`Sfs.sampled_pops`

        :returns: 1-d array of integers
        :rtype: :class:`numpy.ndarray`
        """
        return self.configs.sampled_n

    @property
    def ascertainment_pop(self):
        return self.configs.ascertainment_pop

    @property
    def sampled_pops(self):
        """Names of sampled populations

        :returns: 1-d array of strings
        :rtype: :class:`numpy.ndarray`
        """
        return self.configs.sampled_pops

    @property
    def n_loci(self):
        """The number of loci in the dataset

        :rtype: int
        """
        return len(self.loc_idxs)

    @property
    def n_nonzero_entries(self):
        return len(self.configs)

    @property
    def config_array(self):
        """Array of unique configs.

        This is a 3-d array with shape ``(n_configs, n_pops, 2)``. \
        ``config_array[i, p, a]`` gives the count of allele ``a`` \
        in population ``p`` in configuration ``i``. \
        Allele ``a=0`` is ancestral, ``a=1`` is derived. \
        :attr:`Sfs.sampled_pops` gives the ordering of the populations.

        :returns: array of configs
        :rtype: :class:`numpy.ndarray`
        """
        return self.configs.config_array

    @property
    def length(self):
        """Length of data in bases. ``None`` if unknown.

        :rtype: float
        """
        return self._length

    @memoize_instance
    def n_snps(self, vector=False):
        """Number of SNPs, either per-locus or overall.

        :param bool vector: Whether to return a single total count, or an array of counts per-locus
        """
        if vector:
            return raw_np.array(self.freqs_matrix.sum(axis=0)).reshape((-1,))
        else:
            return np.sum(self._total_freqs)

    def __eq__(self, other):
        loci_dicts = self._get_dict(vector=True)
        try:
            return loci_dicts == other._get_dict(vector=True)
        except AttributeError:
            return False

    def _get_dict(self, vector=False, locus=None):
        if vector:
            assert locus is None
            return [self._get_dict(locus=loc) for loc in range(self.n_loci)]
        elif locus is None:
            return dict(zip(map(_config2hashable, self.configs),
                            self._total_freqs))
        return dict(zip(
            (_config2hashable(self.configs[i]) for i in self.loc_idxs[locus]),
            self.loc_counts[locus]))

    def to_dict(self, vector=False):
        if not vector:
            return {_hashed2config(k): v for k, v in self._get_dict().items()}
        else:
            return [{_hashed2config(k): v for k, v in d.items()}
                    for d in self._get_dict(vector=True)]

    # FIXME: This method returns the _negative_ entropy, and should be renamed.
    @cached_property
    def _entropy(self):
        counts = self._total_freqs
        n_snps = float(self.n_snps())
        p = counts / n_snps
        # return np.sum(p * np.log(p))
        ret = np.sum(p * np.log(p))

        # correct for missing data
        sampled_n = np.sum(self.configs.value, axis=2)
        sampled_n_counts = co.Counter()
        assert len(counts) == len(sampled_n)
        for c, n in zip(counts, sampled_n):
            n = tuple(n)
            sampled_n_counts[n] += c
        sampled_n_counts = np.array(
            list(sampled_n_counts.values()), dtype=float)

        ret = ret + np.sum(sampled_n_counts / n_snps *
                           np.log(n_snps / sampled_n_counts))
        assert not np.isnan(ret)
        return ret

    def _get_muts_poisson_entropy(self, use_pairwise_diffs):
        if use_pairwise_diffs:
            return self._pairwise_muts_poisson_entropy
        else:
            return self._total_muts_poisson_entropy

    @cached_property
    def _total_muts_poisson_entropy(self):
        lambd = self.n_snps(vector=True)
        lambd = lambd[lambd > 0]
        return np.sum(-lambd + lambd * np.log(lambd)
                      - scipy.special.gammaln(lambd + 1))

    @cached_property
    def _pairwise_muts_poisson_entropy(self):
        lambd = self.avg_pairwise_hets
        ret = -lambd + lambd * np.log(lambd) - scipy.special.gammaln(lambd + 1)
        ret[lambd <= 1e-16] = 0
        ret = ret * self.sampled_n / float(np.sum(
            self.sampled_n[self.ascertainment_pop]))
        return np.sum(ret[:, self.ascertainment_pop])

    def fold(self):
        """
        :returns: A copy of the SFS, but with folded entries.
        :rtype: :class:`Sfs`
        """
        def get_folded(config):
            if tuple(config[:, 0]) < tuple(config[:, 1]):
                return config[:, ::-1]
            else:
                return config

        compressed_folded = CompressedAlleleCounts.from_iter(
            map(get_folded, self.configs),
            len(self.sampled_pops), sort=False)

        return self.from_matrix(
            compressed_folded.index2uniq_mat @ self.freqs_matrix,
            ConfigList(self.sampled_pops, compressed_folded.config_array,
                        sampled_n=self.sampled_n,
                        ascertainment_pop=self.ascertainment_pop),
            folded=True, length=self._length)

    def _copy(self, sampled_n=None):
        """
        See also: ConfigList._copy()
        """
        if sampled_n is None:
            sampled_n = self.sampled_n
        return self.from_matrix(
            self.csr_freqs_matrix,
            ConfigList(self.sampled_pops, self.configs.value,
                        sampled_n=sampled_n,
                        ascertainment_pop=self.ascertainment_pop),
            self.folded, self._length)

    def _integrate_sfs(self, weights, vector=False, locus=None):
        if vector:
            assert locus is None
            return np.array([self._integrate_sfs(weights, locus=loc)
                             for loc in range(self.n_loci)])
        if locus is None:
            idxs, counts = slice(None), self._total_freqs
        else:
            idxs, counts = self.loc_idxs[locus], self.loc_counts[locus]
        return np.sum(weights[idxs] * counts)

    def subset_populations(self, populations, non_ascertained_pops=None):
        """Returns lower-dimensional SFS restricted to a subset of populations.

        :param list populations: list of populations to subset to
        :param list non_ascertained_pops: list of populations to treat as non-ascertained (SNPs must be polymorphic with respect to the ascertained populations)
        :returns: Lower-dimensional SFS restricted to a subset of populations.
        :rtype: :class:`Sfs`
        """
        if non_ascertained_pops is None:
            non_ascertained_pops = []
            for pop, asc in zip(self.sampled_pops, self.ascertainment_pop):
                if not asc:
                    non_ascertained_pops.append(pop)
        return self._subset_populations(tuple(populations),
                                        tuple(non_ascertained_pops))

    @memoize_instance
    def _subset_populations(self, populations, non_ascertained_pops):
        # old indexes of the new population ordering
        old_sampled_pops = list(self.sampled_pops)
        old_pop_idx = np.array([
            old_sampled_pops.index(p) for p in populations], dtype=int)

        # old ascertainment
        ascertained = dict(zip(self.sampled_pops,
                               self.ascertainment_pop))
        # restrict to new populations
        ascertained = {p: ascertained[p] for p in populations}
        # previously non-ascertained pops should remain so
        for pop, is_asc in ascertained.items():
            if not is_asc:
                assert pop in non_ascertained_pops
        # update ascertainment
        for pop in non_ascertained_pops:
            ascertained[pop] = False
        # convert from dict to numpy array
        ascertained = np.array([ascertained[p] for p in populations])

        # keep only polymorphic configs
        asc_only = self.configs[:, old_pop_idx[ascertained], :]
        asc_is_poly = (asc_only.sum(axis=1) != 0).all(axis=1)
        asc_is_poly = np.arange(len(asc_is_poly))[asc_is_poly]
        sub_sfs = self._subset_configs(asc_is_poly)

        # get the new configs
        new_configs = CompressedAlleleCounts.from_iter(
            sub_sfs.configs[:, old_pop_idx, :],
            len(populations), sort=False)

        return self.from_matrix(
            new_configs.index2uniq_mat @ sub_sfs.freqs_matrix,
            ConfigList(populations, new_configs.config_array,
                        ascertainment_pop=ascertained),
            self.folded, self._length)

    def _subset_configs(self, idxs):
        return self.from_matrix(
            self.csr_freqs_matrix[idxs, :],
            _ConfigList_Subset(self.configs, idxs),
            self.folded, self._length)

    @property
    def sfs(self):
        return self

    @cached_property
    def p_missing(self):
        """Estimate of probability that a random allele from each population is missing.

        Missingness is estimated as follows: \
        from each SNP remove a random allele; \
        if the resulting config is monomorphic, then ignore. \
        If polymorphic, then count whether the removed allele \
        is missing or not.

        This avoids bias from fact that we don't observe \
        some polymorphic configs that appear monomorphic \
        after removing the missing alleles.

        :returns: 1-d array of missingness per population
        :rtype: :class:`numpy.ndarray`
        """
        counts = self._total_freqs
        sampled_n = self.sampled_n
        n_pops = len(self.sampled_pops)

        config_arr = self.configs.value
        # augment config_arr to contain the missing counts
        n_miss = sampled_n - np.sum(config_arr, axis=2)
        config_arr = np.concatenate((config_arr, np.reshape(
            n_miss, list(n_miss.shape)+[1])), axis=2)

        ret = []
        for i in range(n_pops):
            n_valid = []
            for allele in (0, 1, -1):
                # configs with removed allele
                removed = np.array(config_arr)
                removed[:, i, allele] -= 1
                # is the resulting config polymorphic?
                valid_removed = (removed[:, i, allele] >= 0) & np.all(
                    np.sum(removed[:, :, :2], axis=1) > 0, axis=1)

                # sum up the valid configs
                n_valid.append(np.sum(
                    (counts * config_arr[:, i, allele])[valid_removed]))
            # fraction of valid configs with missing additional allele
            ret.append(n_valid[-1] / float(sum(n_valid)))
        return np.array(ret)


def _csr_freq_matrix_from_counters(idxs_by_loc, cnts_by_loc,
                                   n_configs):
    data = []
    indices = []
    indptr = []
    n_loc = 0
    for idxs, cnts in zip(idxs_by_loc, cnts_by_loc):
        indptr.append(len(data))
        data.extend(cnts)
        indices.extend(idxs)
        n_loc += 1
    indptr.append(len(data))
    return scipy.sparse.csc_matrix((data, indices, indptr),
                                   shape=(n_configs, n_loc)).tocsr()


def _get_subsample_counts(configs, n):
    subconfigs, weights = [], []
    for pop_comb in it.combinations_with_replacement(configs.sampled_pops, n):
        subsample_n = co.Counter(pop_comb)
        subsample_n = np.array([subsample_n[pop]
                                for pop in configs.sampled_pops], dtype=int)
        if np.any(subsample_n > configs.sampled_n):
            continue

        for sfs_entry in it.product(*(range(sub_n + 1)
                                      for sub_n in subsample_n)):
            sfs_entry = np.array(sfs_entry, dtype=int)
            if np.all(sfs_entry == 0) or np.all(sfs_entry == subsample_n):
                # monomorphic
                continue

            sfs_entry = np.transpose([subsample_n - sfs_entry, sfs_entry])
            cnt_vec = configs.subsample_probs(sfs_entry)
            if not np.all(cnt_vec == 0):
                subconfigs.append(sfs_entry)
                weights.append(cnt_vec)

    return np.array(subconfigs), np.array(weights)


class SimpleNamespace(object):
    pass
