"""Utilities for computing and displaying summary statistics on data sets."""

import logging
import multiprocessing as mp
import os
import os.path as pth
import sys
from collections import namedtuple
from functools import partial

import allel
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tqdm

from .datasets import DNADataset
from .simulation import Simulation
from .utils.config import Config, ConfigError, ConfigMixIn
from .utils.decorators import cached_classproperty
from .utils.misc import stdio_redirect_tqdm, zero_pad_format
from .utils.plugins import load_plugins


log = logging.getLogger(__name__)


def plot_sfs(df_sfs, groupby='scenario', window=1, ax=None, legend=True):
    """
    Function to plot sfs

    Parameters
    ----------
    df_sfs : `~pandas.DataFrame`
        `~pandas.DataFrame` generated by the `SummaryStatistics.sfs` function.
    groupby : str
        Column on which to group the data.  By default, ``'scenario'``, but it
        could be ``'label'``, or a another generated column.
    step : int
        If you want to subsample your sfs, plot every ``step`` value instead of
        all.
    ax : `~matplotlib.axes.Axes`
        Axes to plot in a subplot.
    legend: bool
        Whether to plot the legend.
    """

    uniq_id = df_sfs[groupby].sort_values().unique()
    nsamp = df_sfs.indiv_idx.max()
    if len(uniq_id) > 1:
        dic_color = {j: plt.cm.viridis(int(i * 255 / (len(uniq_id) - 1)))
                     for i, j in enumerate(uniq_id)}
    else:
        dic_color = {uniq_id[0]: plt.cm.viridis(128)}

    if ax is None:
        fig, ax = plt.subplots(1, 1)

    for g in df_sfs.groupby(groupby):
        color = dic_color[g[0]]
        yerr = g[1].groupby("indiv_idx")[["freq_indiv", "i_xi_sem_norm"]].mean()

        if window > 1:
            yerr = yerr.rolling(window).i_xi_sem_norm.mean()
            by_n_indiv = g[1].groupby("indiv_idx")[["freq_indiv", "i_xi_norm"]]
            by_n_indiv.mean().rolling(window).mean().dropna().plot(
                    x="freq_indiv",
                    y="i_xi_norm",
                    yerr=yerr,
                    ax=ax,
                    label=g[0],
                    color=color)

        else:
            # It groupby on n_indiv even though it supposed to have 1 value
            # but it depends on the outer group defined by "by=". If it is
            # something else than "scenario", it can have more than 1 value for
            # a given n_indiv category.
            yerr = g[1].groupby("indiv_idx")[["freq_indiv", "i_xi_sem_norm"]].mean()
            g[1].groupby("indiv_idx")[["freq_indiv", "i_xi_norm"]].mean().plot(
                    x="freq_indiv",
                    y="i_xi_norm",
                    yerr=yerr,
                    ax=ax,
                    color=color,
                    label=g[0])

    ax.axhline(1 / (nsamp),
               color="0.5",
               zorder=10,
               linestyle="--")
    if legend:
        ax.legend(loc=6, bbox_to_anchor=(1, 0.5))
    else:
        ax.legend_.set_visible(False)


def plot_ld(df_ld, groupby='scenario', ax=None, legend=True):
    """Function to plot the Linkage Desiquilibrium.

    Parameters
    ----------
    df_ld : `~pandas.DataFrame`
        `~pandas.DataFrame` generated by the `SummaryStatistics.ld` function.
    groupby : str
        Column on which to group the data.  By default, ``'scenario'``, but it
        could be ``'label'``, or a another generated column.
    ax : `~matplotlib.axes.Axes`
        Axes to plot in a subplot.
    legend: bool
        Whether to plot the legend.
    """

    uniq_id = df_ld[groupby].sort_values().unique()
    if len(uniq_id) > 1:
        dic_color = {j: plt.cm.viridis(int(i * 255 / (len(uniq_id) - 1)))
                     for i, j in enumerate(uniq_id)}
    else:
        dic_color = {uniq_id[0]: plt.cm.viridis(128)}

    if ax is None:
        fig, ax = plt.subplots(1, 1)

    for g in df_ld.groupby(groupby):
        color = dic_color[g[0]]
        g[1].groupby("dist_group").mean().plot(x="mean_dist",
                                               y="mean_r2",
                                               yerr=g[1].groupby("dist_group").mean_r2.sem(),
                                               kind="scatter",
                                               label=g[0],
                                               ax=ax,
                                               color=color,
                                               legend=True)
    ax.set_xscale("log")
    ax.set_xlim(10,
                10**round(np.log10(df_ld.mean_dist.max())))
    if legend:
        ax.legend(loc=6, bbox_to_anchor=(1, 0.5))
    else:
        ax.legend_.set_visible(False)


def plot_sel(df_sel, groupby='scenario', which="all", ax=None, legend=True):
    """Function to plot selection sumstats.

    Parameters
    ----------
    df_sel : `~pandas.DataFrame`
        `~pandas.DataFrame` generated by the selection statistics functions
        in `SummaryStatistics.statistics`, namely `SummaryStatistics.tajimasD`,
        `SummaryStatistics.ihs` and `SummaryStatistics.nsl`.
    groupby : str
        Column on which to group the data.  By default, ``'scenario'``, but it
        could be ``'label'``, or a another generated column.
    which : str
        Which one to plot: ``'tajimasd'``, ``'nsl'``, ``'ihs'``.
    ax : `~matplotlib.axes.Axes`
        Axes to plot in a subplot.
    legend: bool
    """

    uniq_id = df_sel[groupby].sort_values().unique()
    if len(uniq_id) > 1:
        dic_color = {j: plt.cm.viridis(int(i * 255 / (len(uniq_id) - 1)))
                     for i, j in enumerate(uniq_id)}
    else:
        dic_color = {uniq_id[0]: plt.cm.viridis(128)}

    if ax is None:
        if which == "all":
            fig, ax = plt.subplots(3, 1)
        else:
            fig, ax = plt.subplots(1, 1)
    else:
        assert len(ax) == 3 and which == "all"

    for g in df_sel.groupby(groupby):
        color = dic_color[g[0]]
        g[1].groupby("dist_group").mean().plot(x="mean_dist",
                                               y="mean_r2",
                                               yerr=g[1].groupby("dist_group").mean_r2.sem(),
                                               kind="scatter",
                                               label=g[0],
                                               ax=ax,
                                               color=color,
                                               legend=True)
    ax.set_xscale("log")
    ax.set_xlim(10,
                10**round(np.log10(df_sel.mean_dist.max())))
    if legend:
        ax.legend(loc=6, bbox_to_anchor=(1, 0.5))
    else:
        ax.legend_.set_visible(False)


class statistic:
    """
    Decorator which marks a method of the `SummaryStatistics` class or a
    subclass thereof as a statistic that should be computed when computing
    summary statistics for a dataset.
    """
    def __init__(self, *, name=None, description=None):
        self.description = description
        self.name = name.lower() if name else name

    def __call__(self, func):
        """
        Decorator for methods that attach a `statistic` object to them so they
        are marked as statistic functions.
        """

        # Support wrapping staticmethods and classmethods
        if hasattr(func, '__func__'):
            wrapped_func = func
            func = func.__func__
        else:
            wrapped_func = func

        if self.name is None:
            self.name = func.__name__.lower()

        func.statistic = self
        return wrapped_func


class SumStats(namedtuple('SumStats', ('sfs', 'ld', 'sel'))):
    @classmethod
    def from_csv(cls, filename_format, format_fields, **kwargs):
        """
        Counterpart to `SumStats.to_csv`.

        Arguments are the same as `SumStats.to_csv`, except that ``**kwargs``
        are passed to `pandas.read_csv`.
        """

        # Specify dtypes for standard columns in each table; the dtypes of
        # other columns will be inferred using the standard rules from Pandas
        # This seems to be the magic combination of parameters to get floating
        # point columns containing NaNs to be read as floating point (with NaN)
        # values, while empty string columns are read as string columns
        read_csv_kwargs = {
            'dtype': {'scenario_idx': int, 'model': str, 'label': str},
            'na_values': ['nan', 'NaN', 'NAN'],
            'keep_default_na': False
        }
        read_csv_kwargs.update(kwargs)

        stats = []
        for stat_type in cls._fields:
            filename = filename_format.format(type=stat_type, **format_fields)
            stats.append(pd.read_csv(filename, **read_csv_kwargs))

        return cls(*stats)

    def to_csv(self, filename_format, format_fields={}, overwrite=False,
               **kwargs):
        """
        Write summary statistics tables to CSV files.

        Parameters
        ----------
        filename_format : str
            The filename format should be a template string containing a
            ``{type}`` field, where ``type`` is the statistic type, currently
            either ``'sfs'``, ``'ld'``, and ``'sel'``.
        format_fields : dict, optional
            Additional values to fill into fields in the ``filename_format``
            template, if any.  E.g. if ``filename_format`` is
            ``'{dataset_name}_{type}.csv'`` then it is also necessary to pass
            ``format_fields={'dataset_name': 'my_simulation'}``.
        overwrite : bool, optional
            If `True`, overwrite any existing file of the same name, otherwise
            if the file already exists append new statistics to its end instead
            of overwriting it (default: `False`).
        **kwargs
            Additional keyword arguments are passed through to
            `pandas.DataFrame.to_csv`.
        """

        def write_csv(type, df):
            filename = filename_format.format(type=type, **format_fields)
            dirname = pth.dirname(filename)

            if not os.path.isdir(dirname):
                os.makedirs(dirname, exist_ok=True)

            mode = 'w' if overwrite else 'a'
            with open(filename, mode) as fobj:
                header = (fobj.tell() == 0) or overwrite
                to_csv_kwargs = {
                    'sep': ',',
                    'na_rep': 'nan',
                    'index': False,
                    'header': header
                }
                to_csv_kwargs.update(kwargs)
                df.to_csv(filename, **to_csv_kwargs)

        for type_ in self._fields:
            write_csv(type_, getattr(self, type_))


class SummaryStatistics(ConfigMixIn):
    """Computes summary statistics for a dataset."""

    config_schema = 'summary-statistics'
    sumstats_cls = SumStats

    def __init__(self, config={}, validate=True, simulation=None, plugins=None):
        # Summary statistics configuration can be loaded either from a
        # stand-alone summary statistics config which references a simulation
        # config, or from a simulation config containing an embedded summary
        # statistics config; see dnadna/schemas/summary-statistics.yml for more
        # detail
        if 'simulation' in config:
            simulation_config = config['simulation']
            summary_statistics_config = config
        else:
            try:
                # It may also be a dataset config
                filename = getattr(config, 'filename', None)
                config = Config(config, filename=filename, schema='dataset',
                                validate=True)
            except ConfigError:
                raise ConfigError(config,
                    "summary statistics config must either contain a "
                    "'simulation' property, or must be a simulation "
                    "config containing a 'summary_statistics' property")
            else:
                simulation_config = config
                summary_statistics_config = config.pop('summary_statistics', {})
                summary_statistics_config['simulation'] = simulation_config

        # load any plugins listed in the config; plugins can use
        # SummaryStatistics.register_statistic to register any additional
        # selection statistics we wish to compute
        # It is also possible to specify additional plugins not included
        # in the config via the plugins argument--if so these are loaded
        # after the plugins from the config (and thus may override them)
        all_plugins = summary_statistics_config.get('plugins', [])
        if plugins is not None:
            all_plugins += list(plugins)

        load_plugins(all_plugins)

        super().__init__(config=summary_statistics_config, validate=True)

        # Already validated by super().__init__ so use validate=False
        if simulation is None:
            self.simulation = Simulation(simulation_config, validate=False)
        else:
            # Override the simulation from the config file with an existing
            # Simulation instance
            self.simulation = simulation

    @classmethod
    def from_config_file(cls, filename, validate=True, simulation=None,
                         plugins=None, **kwargs):
        config = Config.from_file(filename, validate=False, **kwargs)
        return cls(config=config, validate=validate, simulation=simulation,
                   plugins=plugins)

    @classmethod
    def register_statistic(cls, func, name=None, description=None,
                           override=False):
        """
        Registers a new statistic function on this class.

        New statistics can be registered either by subclassing
        `SummaryStatistics` and adding additional `statistic` decorators, or by
        calling this method on an existing class.

        Use `SummaryStatistics.unregister_statistic` with the function name
        (all lower-case) to remove the statistic; this does *not* remove
        statistics that were registered with the `statistic` decorator.

        Examples
        --------

        >>> from dnadna.summary_statistics import SummaryStatistics
        >>> def my_stat(*args, **kwargs):
        ...     pass
        >>> SummaryStatistics.register_statistic(my_stat,
        ...     description='My Statistic')
        >>> SummaryStatistics.statistics
        {... 'my_stat': (<function my_stat at 0x...>, 'My Statistic')...}
        >>> SummaryStatistics.unregister_statistic('my_stat')
        >>> 'my_stat' in SummaryStatistics.statistics
        False

        Built-in statistics can also be overridden with different
        implementations:

        >>> SummaryStatistics.register_statistic(my_stat,
        ...     name='ld', description='my LD implementation', override=True)
        >>> SummaryStatistics.statistics['ld']
        (<function my_stat at 0x...>, 'my LD implementation')
        >>> SummaryStatistics.unregister_statistic('ld')
        >>> SummaryStatistics.statistics['ld']
        (<function SummaryStatistics.ld at 0x...>, 'LD')

        Invalidate the cached `SummaryStatistics.statistics` completely so that
        it's reset for future tests (the unregister_statistic call already
        reset it, but could result in re-ordering of dict keys).  This cannot
        simply be done with ``del`` since ``statistics`` is a
        `~dnadna.utils.decorators.classproperty`:

        >>> SummaryStatistics.reset_statistics()
        """

        if name is None:
            name = func.__name__.lower()

        if description is None:
            description = func.__name__.title()

        if name in cls.statistics and not override:
            raise ValueError(
                f"a statistic named {name} is already registered on "
                f"{cls}; use override=True to force overriding it")

        func.statistic = statistic(name=name, description=description)

        cls.statistics[name] = (func, description)

    @classmethod
    def unregister_statistic(cls, name):
        """See `SummaryStatistics.register_statistic`."""

        cls_func = getattr(cls, name, None)
        cls_statistic = getattr(cls_func, 'statistic', None)
        cache_key = name.lower()
        cached_func = cls.statistics.get(cache_key, (None,))[0]

        if (isinstance(cls_statistic, statistic) and
                cached_func.statistic is cls_statistic):
            # This may not be the class if the built-in statistic of the same
            # name was overridden by SummaryStatistics.register_statistic
            raise ValueError(
                "statistics registered as class-level methods cannot be "
                "unregistered")

        # Allow KeyErrors to bubble up if an invalid name is unregistered
        del cls.statistics[cache_key]

        # Re-register the original class-level statistic if it was overridden
        if cls_statistic is not None:
            cls.statistics[cache_key] = (cls_func, cls_statistic.description)

    @classmethod
    def reset_statistics(cls):
        """
        Resets the class's registered statistics to the defaults.

        Do not use ``del SummaryStatistics.register_statistics``: due to the
        behavior of data-descriptors, calling ``del`` on a class attribute does
        not go through the ``__delete__`` method if the attribute is a
        descriptor; modifying this behavior would require a special metaclass.

        See `SummaryStatistics.register_statistic`.
        """

        cls.__dict__['statistics'].invalidate(cls)

    # TODO: This is not being propertly picked up by Sphinx autodoc, probably
    # due to being a specialized subclass of 'property'; I've had problems like
    # this before with autodoc--will need to investigate.
    @cached_classproperty
    def statistics(cls):
        """
        Dict of all statistic functions that are computed as part of the
        summary statistics.

        Maps a normalized short name of the function to a pair consisting of
        the function (or bound method) itself, and a display name for the
        statistic.

        Examples
        --------

        >>> from dnadna.summary_statistics import SummaryStatistics
        >>> SummaryStatistics.statistics
        {'sfs': (<function SummaryStatistics.sfs at 0x...>, 'SFS'),
         'ld': (<function SummaryStatistics.ld at 0x..., 'LD'),
         'tajimasd': (<function SummaryStatistics.tajimasD at 0x...>, "Tajima's D"),
         'ihs': (<function SummaryStatistics.ihs at 0x...>, 'iHS'),
         'nsl': (<function SummaryStatistics.nsl at 0x...>, 'nSL')}
        """

        statistics = {}
        for k, f in vars(cls).items():
            if hasattr(f, '__func__'):
                # Handle functions wrapped in class/staticmethods
                f = f.__func__

            if isinstance(getattr(f, 'statistic', None), statistic):
                s = f.statistic
                statistics[s.name] = (f, s.description)

        return statistics

    @statistic(description='SFS')
    @staticmethod
    def sfs(haplotype, ac, n_indiv=None, folded=False):
        """Compute SFS for SNP matrix"""
        # TODO: Make improvements to this function's implementation
        # A lot of it appears to be redundant
        if n_indiv is None:
            n_indiv = haplotype.shape[1]
        tmp_df = pd.DataFrame({"indiv_idx": range(1, n_indiv)})
        if folded:
            df_sfs = pd.DataFrame(allel.sfs_folded(ac), columns=["n_snp"])
            df_sfs["i_xi"] = allel.sfs_folded_scaled(ac)
            df_sfs.index.name = "indiv_idx"
            df_sfs.reset_index(inplace=True)
            df_sfs = df_sfs.merge(tmp_df, on="indiv_idx", how="right")
            df_sfs = df_sfs.fillna(0).astype(int)
        else:
            df_sfs = pd.DataFrame(allel.sfs(ac.T[1]), columns=["n_snp"])
            df_sfs["i_xi"] = allel.sfs_scaled(ac.T[1])
            df_sfs.index.name = "indiv_idx"
            df_sfs.reset_index(inplace=True)
            df_sfs = df_sfs.merge(tmp_df, on="indiv_idx", how="right")
            df_sfs = df_sfs.fillna(0).astype(int)

        df_sfs["freq_indiv"] = df_sfs.indiv_idx / n_indiv
        df_sfs['n_indiv'] = n_indiv
        return df_sfs

    @statistic(description='LD')
    @staticmethod
    def ld(haplotype, pos_vec, chromosome_size,
           circular=False, distance_bins=None, gaps_type="short",
           min_snp_pairs=300):
        """
        Compute LD for a subset of SNPs drawn with different gap sizes in
        between them.

        Gap sizes follow power 2 distribution.
        The LD is then computed and averaged over different bin
        (``distance_bins``) sizes.

        Parameters
        ----------
        haplotype : 2-D `~numpy.ndarray` or `allel.HaplotypeArray`
            SNP matrix where in the first dimension are the SNP (rows) and in
            the second dimension (columns) are the samples.
        pos_vec : 1-D `~numpy.ndarray`
            array of absolute positions in [0, ``chromosome_size``].
        chromosome_size : int
            Size of the chromosome.
        circular : bool
            Whether to consider the chromosome circular or not.  If circular,
            the maximum distance between 2 SNPs is thus half the chromosome.
        distance_bins : int or list
            LD will be averaged by bins of distances e.g. if ``distance_bins``
            = [0, 100, 1000, 10000], LD will be averaged for the groups
            [0,100), [100, 1000), and [1000, 10000). If ``distance_bins`` is an
            int, it defines the number of bins of distances for which to
            compute the LD The bins are created in a logspace if
            ``distance_bins`` is a list, they will be used instead.
        gaps_type: str
            Pairs of SNP considered are separated by a given number (gap) of
            columns. Not all pairs are considered.  By default (``'short'``),
            gaps are power of 2 up to the closest power of 2 of the number of
            SNP.
            Meaning that most of the comparisons will be done on close SNPs
            (short distance).
            If one wants to sample more at large distance (to test for
            circularity for instance), use ``'long'`` instead of ``'short'``
            Using ``'long'`` will add gaps like: ``n_snp - gaps``. It will take
            more time to run.
        min_snp_pairs: int
            Minimum number of pairs of SNP to consider for a given gap size.
            If the gap size is big enough such that there is less than
            ``min_snp_pairs`` possible pairs, then all pairs are considered.

        Returns
        -------
        `~pandas.DataFrame`
            Table with the ``distance_bins`` as index, and the mean values
            over each bin.
        """
        # TODO: More of these options should be explicitly documented in the
        # summary statistics config schema as well.
        # TODO: Make improvements to this function's implementation

        if distance_bins is None or isinstance(distance_bins, int):
            if isinstance(distance_bins, int):
                n_bins = distance_bins - 1
            else:
                # TODO: Why 19???
                n_bins = 19

            # TODO: This can be done more simply, without using np.insert
            if circular:
                distance_bins = np.logspace(2, np.log10(chromosome_size // 2),
                                            n_bins)
                distance_bins = np.insert(distance_bins, 0, [0])
            else:
                distance_bins = np.logspace(2, np.log10(chromosome_size),
                                            n_bins)
                distance_bins = np.insert(distance_bins, 0, [0])

        n_snp, n_samples = haplotype.shape

        # gaps are distance between SNPs in term of position in the snp matrix (not in bp)
        # log2 scales of intervals
        gap_intervals = (2 ** np.arange(0, np.log2(n_snp), 1)).astype(int)
        if gaps_type.lower() == "long":
            # TODO: I'm sure this can be done better
            gap_intervals_shifted = np.array(list(n_snp - gap_intervals)[::-1])
            gap_intervals = np.unique(np.concatenate([
                gap_intervals, gap_intervals_shifted])).astype(int)

        else:
            # TODO: This should be an exception instead
            if gaps_type.lower() != "short":
                log.warning(
                    f"gaps should be either `short` or `long`. Using short "
                    f"instead of {gaps_type}")

        def snp_pairs_for_gap(gap):
            if circular:
                max_value = n_snp
            else:
                max_value = n_snp - gap

            # TODO: I'm not so sure it's necessary to generate so many pairs
            # right away in memory; this could be a generator.  However,
            # doing it in memory might be faster, so need to profile which is
            # better.

            # min_snp_pairs : min number of SNP pairs to consider.
            if max_value < min_snp_pairs:
                # if not many possible pairs possible, just take them all
                # directly, instead of reaching that number after many more
                # random trials
                snps = np.arange(0, n_snp, gap)
                snp_pairs = np.unique([
                    ((snps[i] + i) % n_snp, (snps[i + 1] + i) % n_snp)
                    for i in range(len(snps) - 1)], axis=0)
                return np.concatenate([
                    (snp_pairs + i) % n_snp for i in range(max_value)], axis=0)
            else:
                # TODO: As this uses randomization, we should also have an option
                # to set the PRNG seed
                # adding a random start (+1, bc 2nd bound in randint is exlusive)
                snps = (np.arange(0, n_snp, gap) +
                        np.random.randint(0, (n_snp - 1) % gap + 1))
                # non overlapping contiguous pairs
                # snps=[ 196, 1220, 2244] becomes
                # snp_pairs=[(196, 1220), (1221, 2245)]
                snp_pairs = np.unique([
                    ((snps[i] + i) % n_snp, (snps[i + 1] + i) % n_snp)
                    for i in range(len(snps) - 1)], axis=0)

                # If we don't have enough pairs (typically when gap is large),
                # we add a random rotation until we have at least 300)

                if not circular:
                    # remove pairs that are over the edges
                    snp_pairs = snp_pairs[snp_pairs[:, 0] < snp_pairs[:, 1]]
                last_pair = snp_pairs[-1]

                while len(snp_pairs) < min(min_snp_pairs, max_value):
                    shift = np.random.randint(1, n_snp) % n_snp
                    new_pair = (last_pair + shift) % n_snp
                    snp_pairs = np.unique(np.concatenate([
                        snp_pairs, new_pair.reshape(1, 2)]), axis=0)
                    last_pair = new_pair

                    if not circular:
                        snp_pairs = snp_pairs[snp_pairs[:, 0] < snp_pairs[:, 1]]

                return snp_pairs

        # Functions to aggregate the values within each distance bin
        agg_bins = {c: tuple(c.rsplit('_', 1)) for c in [
            'snp_dist_mean', 'r2_mean', 'r2_count', 'r2_sem'
        ]}

        lds = []
        generate_pairs = (snp_pairs_for_gap(gap) for gap in gap_intervals)

        for i, snps_pos in enumerate(generate_pairs):
            if circular:
                # %chromosome_size/2 because max distance btw 2 SNP is
                # chromosome_size/2
                snp_dist = (((np.diff(pos_vec[snps_pos]) % chromosome_size)) %
                            (chromosome_size // 2))
                sd = pd.DataFrame(snp_dist, columns=["snp_dist"])
            else:
                sd = pd.DataFrame((np.diff(pos_vec[snps_pos])),
                                  columns=["snp_dist"])

            rh = allel.rogers_huff_r

            def memoized_rh(pair, memo={}):
                k = tuple(pair)
                try:
                    return memo[k]
                except KeyError:
                    return memo.setdefault(k, rh(haplotype[pair]) ** 2)

            sd['dist_group'] = pd.cut(sd.snp_dist, bins=distance_bins)
            sd['r2'] = [memoized_rh(pair) for pair in snps_pos]
            sd['gap_id'] = i
            lds.append(sd)

        ld = pd.concat(lds)
        return ld.dropna().groupby("dist_group").agg(**agg_bins)

    @statistic(description="Tajima's D")
    @staticmethod
    def tajimasD(haplotype, pos_vec=None, window=None):
        """
        Given a snp_mat, return tajima's D
        window: Number of window of equal size to slice the position vector.

        if windowed stat, provide pos_vec too.
        """
        # TODO: Improve this method's docstring.
        # TODO: This function is mostly ok as-is, but can probably use a little
        # improvement in its implementation
        if not window:
            allele_count = haplotype.count_alleles()
            return allel.tajima_d(allele_count)
        else:
            pos = pd.DataFrame(pos_vec, columns=["pos"])
            pos["pos_cat"] = pd.cut(pos.pos, window,
                                    labels=range(1, window + 1))
            pos.index.name = "SNP"
            snp_per = pos.reset_index().groupby("pos_cat").SNP.unique()
            all_taj_d = np.full(len(snp_per), np.nan)
            for idx, per in enumerate(snp_per):
                if len(per):
                    allele_count = haplotype[per].count_alleles()
                    all_taj_d[idx] = allel.tajima_d(allele_count)
            return all_taj_d

    @statistic(description='iHS')
    @staticmethod
    def ihs(haplotype, pos_vec, window=None):
        """Compute the standardize integrated haplotype score"""
        # TODO: This function is mostly ok as-is, but can probably use a little
        # improvement in its implementation
        # TODO: Should min_maf be a customizable argument?
        ihs = allel.ihs(haplotype, pos_vec, min_maf=0.01,
                        include_edges=True)
        aac = haplotype.count_alleles().T[1]
        ihs_stand = _standardize_by_allele_count_safe(ihs, aac)

        if window:
            di = pd.DataFrame(ihs_stand, columns=["iHS"])
            di["pos_cat"] = pd.cut(pos_vec, window, labels=range(1, window + 1))
            dig = di.groupby("pos_cat").iHS.mean()
            return dig
        else:
            return ihs_stand

    @statistic(description='nSL')
    @staticmethod
    def nsl(haplotype, pos_vec=None, window=None):
        """
        Compute the standardize number of segregating sites by length (nSl)
        for each variant, comparing the reference and alternate alleles,
        after Ferrer-Admetlla et al. (2014)

        If windowed stat, provide pos_vec too.
        """
        # TODO: This function is mostly ok as-is, but can probably use a little
        # improvement in its implementation

        nsl = allel.nsl(haplotype)
        aac = haplotype.count_alleles().T[1]
        nsl_stand = _standardize_by_allele_count_safe(nsl, aac)

        if window:
            dn = pd.DataFrame(nsl_stand, columns=["nSL"])
            dn["pos_cat"] = pd.cut(pos_vec, window, labels=range(1, window + 1))
            dng = dn.groupby("pos_cat").nSL.mean()
            return dng
        else:
            return nsl_stand

    def run(self, n_replicates=None, label='', n_cpus=None,
            ignore_missing=False, raise_exc=False, overwrite=False,
            progress_bar=False):
        """
        Compute summary statistics for all scenarios in a dataset, and save the
        resulting tables as CSV files.

        The summary statistics filenames are paths are determined from the
        ``filename_format`` property in the summary statistics config.  By
        default this is relative to the ``data_root`` for the simulation, and
        has the format
        ``sumstats/scenario_{scenario}/{dataset_name}_{scenario}_{type}.csv``
        where ``{type}`` is each of "ld", "sfs", and "sel".

        If connected to a TTY and ``progress_bar=True``, a progress bar is
        displayed while computing the statistics.
        """

        n_scenarios = len(self.simulation.scenario_params.index)
        n_cpus = n_cpus or mp.cpu_count()

        log.info(
            f'Starting sumstats with {n_cpus} CPU(s) on {n_scenarios} '
            f'scenario(s)')

        def write_csv(scenario_idx, sumstats):
            filename_format = zero_pad_format(self.filename_format,
                    scenario=n_scenarios)
            filename_format = pth.join(self.simulation.data_root,
                                       filename_format)
            sumstats.to_csv(filename_format,
                    format_fields={
                        'scenario': scenario_idx,
                        'dataset_name': self.simulation.dataset_name
                    }, overwrite=overwrite)

        progress_bar = progress_bar and sys.stderr.isatty()

        it = self.compute_scenarios(n_replicates=n_replicates, label=label,
                n_cpus=n_cpus, ignore_missing=ignore_missing,
                raise_exc=raise_exc)

        with stdio_redirect_tqdm() as orig_stdio:
            with tqdm.tqdm(it, total=n_scenarios, unit='simulation',
                           file=orig_stdio.stdout, dynamic_ncols=True,
                           disable=(not progress_bar)) as bar:
                for scenario_idx, sumstats in bar:
                    write_csv(scenario_idx, sumstats)

        log.info("Done")

    def compute_scenarios(self, n_replicates=None, label='', n_cpus=None,
                          ignore_missing=False, raise_exc=False):
        """
        Returns a generator yielding the scenario_idx, and computed summary
        statistics `DataFrames <pandas.DataFrame>` for each scenario.
        """

        n_cpus = n_cpus or mp.cpu_count()
        scenarios = self.simulation.scenario_params.index
        use_mp = (n_cpus != 1)

        if use_mp:
            compute_scenario = self._compute_scenario_wrapper
        else:
            compute_scenario = self.compute_scenario

        compute_scenario = partial(compute_scenario,
                n_replicates=n_replicates, label=label,
                ignore_missing=ignore_missing, raise_exc=raise_exc)

        if use_mp:
            def iter_statistics():
                with mp.Pool(n_cpus) as pool:
                    it = pool.imap_unordered(compute_scenario, scenarios)
                    for res in it:
                        yield res
        else:
            def iter_statistics():
                for scenario_idx in scenarios:
                    yield scenario_idx, compute_scenario(scenario_idx)

        for scenario_idx, sumstats in iter_statistics():
            yield scenario_idx, sumstats

    def compute_scenario(self, scenario_idx, n_replicates=None, label='',
                         ignore_missing=False, raise_exc=False):
        # Create a DNADataset just for this scenario
        dataset = DNADataset({'ignore_missing': ignore_missing},
                             validate=False,
                             source=self.simulation.source,
                             scenario_params=self.simulation.scenario_params,
                             scenario_set=scenario_idx)

        all_sfs = []
        all_ld = []
        all_sel = []

        for idx, (_, _, sample, _) in enumerate(dataset):
            if n_replicates and idx >= n_replicates:
                break

            if sample is None:
                # Returned in case of a missing file and ignore_missing=True
                continue

            sfs, ld, sel = self._compute_single_simulation(
                    scenario_idx, sample, raise_exc=raise_exc)

            # TODO: As mentioned in the previous comment, fix this
            all_sfs.append(sfs)
            all_ld.append(ld)
            all_sel.append(sel)

        all_sfs = pd.concat(all_sfs)
        all_ld = pd.concat(all_ld)
        all_sel = pd.concat(all_sel)

        def add_standard_columns(df):
            df.insert(0, 'scenario_idx', scenario_idx)
            df.insert(1, 'model', self.simulation.dataset_name)
            df.insert(2, 'label', label)
            df.reset_index(inplace=True)

        # TODO: Clean this up as well
        # It looks like there is some duplicate work fo assigning
        # model/scenario/label columns that is also done in
        # _compute_single_simulation; maybe it would make more sense to
        # limit that action to this part of the code instead

        # TODO: These have individual routines for gathering all the results
        # for each statistic; perhaps this could be done in a more extensible
        # manner as well, because right now each statistic is hard-coded here.
        all_sfs2 = all_sfs.groupby("indiv_idx").mean()
        all_sfs2["i_xi_norm"] = all_sfs2.i_xi / all_sfs2.i_xi.mean()
        # TODO: I think this might be broken; the old code assumed all
        # simulations in the dataset had the same n_indiv, but if that's not
        # true then this is broken, so for now we add an n_indiv column to the
        # sfs() output, but need to confirm that this is correct...
        all_sfs2["freq_indiv"] = all_sfs2.index / all_sfs2.n_indiv
        all_sfs2["i_xi_sem_norm"] = (all_sfs.groupby("indiv_idx").i_xi.sem() /
                                     all_sfs2.i_xi.mean())
        all_sfs2["i_xi_std_norm"] = (all_sfs.groupby("indiv_idx").i_xi.std() /
                                     all_sfs2.i_xi.mean())
        add_standard_columns(all_sfs2)

        all_ld2 = all_ld.groupby("dist_group").mean()
        add_standard_columns(all_ld2)

        all_sel2 = all_sel.groupby("position_percent").mean()
        for c in all_sel2.columns:
            all_sel2[f"{c}_std"] = all_sel.groupby("position_percent")[c].std()
            all_sel2[f"{c}_sem"] = all_sel.groupby("position_percent")[c].sem()
        add_standard_columns(all_sel2)

        return self.sumstats_cls(all_sfs2, all_ld2, all_sel2)

    def load(self, scenario_idx, ignore_missing=False):
        """
        Load existing summary statistics for a single scenario, as output with
        the current configuration.
        """

        # Specify dtypes for standard columns in each table; the dtypes of
        # other columns will be inferred using the standard rules from Pandas
        # This seems to be the magic combination of parameters to get floating
        # point columns containing NaNs to be read as floating point (with NaN)
        # values, while empty string columns are read as string columns
        n_scenarios = len(self.simulation.scenario_params.index)
        filename_format = zero_pad_format(self.filename_format,
                scenario=n_scenarios)
        filename_format = pth.join(self.simulation.data_root, filename_format)
        name = self.simulation.dataset_name

        try:
            return self.sumstats_cls.from_csv(filename_format, format_fields={
                'scenario': scenario_idx,
                'dataset_name': name
            })
        except FileNotFoundError as exc:
            log.warning(
                f'missing summary statistics file "{exc.filename}"')
            if not ignore_missing:
                raise

    def load_all(self, ignore_missing=False):
        """
        Returns an iterator of ``(scenario_idx, (sfs, ld, sel))`` tuples for
        all scenarios in a dataset for which summary statistics have been
        computed.

        Parameters
        ----------
        ignore_missing : bool, optional
            If `True`, simply output a warning for missing summary stats files,
            otherwise an exception is raised (default: `False`).
        """

        scenarios = self.simulation.scenario_params.index
        for scenario_idx in scenarios:
            yield (scenario_idx, self.load(scenario_idx,
                ignore_missing=ignore_missing))

    def _compute_scenario_wrapper(self, scenario_idx, **kwargs):
        """
        Same as `SummaryStatistics.compute_scenario` but also returns the
        ``scenario_idx``, for us with multiprocessing.
        """

        stats = self.compute_scenario(scenario_idx, **kwargs)
        return scenario_idx, stats

    def _compute_single_simulation(self, scenario_idx, sample, raise_exc=False):
        # convert to absolute positions
        # TODO: Fix this once we incorporate better handling of position
        # normalization
        pos = sample.pos.numpy()
        if pos.max() <= 1:
            pos = (pos * self.chromosome_size).round().astype(int)

        haplotype = allel.HaplotypeArray(sample.snp.numpy().T)
        allele_count = haplotype.count_alleles()

        # TODO: concat-ing into an empty DataFrame (as we do from the method
        # that calls this one) is rather inefficient; do this better
        done_stats = set()

        def compute_statistic(stat, *args, **kwargs):
            try:
                ret = self._compute_statistic(stat, *args, **kwargs)
            except Exception as exc:
                self._handle_error_for(stat, scenario_idx, exc,
                                       raise_exc=raise_exc)
                ret = pd.DataFrame()

            done_stats.add(stat)

            return ret

        sfs = compute_statistic('sfs', haplotype, allele_count,
                                **self.sfs_options)
        ld = compute_statistic('ld', haplotype, pos, self.chromosome_size,
                               **self.ld_options)

        # Remaining stats all take the same options
        sel_size = self.sel_options.get('window', sample.n_snp)
        sel_stats = []

        for stat in self.statistics:
            if stat in done_stats:
                continue

            val = compute_statistic(stat, haplotype, pos, **self.sel_options)

            # Some of the different stats have different types of return
            # values (either a DataFrame, an ndarray, or a scalar value)
            # so we perform some normalization here; in the case of a scalar
            # (like Tajima's D with window=None) it is expanded to an
            # array/series filled with the value to the length of the number of
            # SNPs.  I'm not totally happy with this inconsistency, but in fact
            # it makes sense that different statistics would compute different
            # things.

            if not isinstance(val, (np.ndarray, pd.DataFrame, pd.Series)):
                val = np.full(sel_size, val)

            if isinstance(val, np.ndarray):
                # Wrap in a named series
                val = pd.Series(val, name=stat)

            sel_stats.append(val)

        sel = pd.concat(sel_stats, axis=1)
        # TODO: Need to undertand why this is being done...
        sel.index.name = "position_percent"
        sel.reset_index(inplace=True)

        return self.sumstats_cls(sfs, ld, sel)

    def _compute_statistic(self, stat, *args, **kwargs):
        """Compute a statistic from the `SummaryStatistics.statistics` table."""

        try:
            stat_func, _ = self.statistics[stat]
        except KeyError:
            raise RuntimeError(f'unknown statistic: "{stat}"')

        return stat_func(*args, **kwargs)

    def _handle_error_for(self, stat, scenario_idx, exc, raise_exc=False):
        descr = self.statistics[stat][1]
        log.error(
            f'while computing {descr} for scenario {scenario_idx}\n'
            f'>>> Error: {exc}')
        if raise_exc:
            raise exc


def _standardize_by_allele_count_safe(score, aac):
    # NOTE: if aac is all NaN, standardize_by_allele_count crashes because it
    # doesn't properly handle the case where there are zero bins, so we just
    # don't call it
    nan_score = np.isnan(score)

    if nan_score.all():
        score_standardized = score
    else:
        # NOTE: standardize_by_allele_count sometimes divides by zero when
        # computing the standardized score, if the standard deviation in a bin
        # is zero; for now it is simply ignored, but perhaps the warning should
        # still be logged?  Normally this won't happen except with very small
        # SNP matrices (e.g. in testing)
        #
        # Also, if the number of alleles for non-NaN scores counted in aac is
        # fewer than the default number of bins used by
        # standardize_by_allele_count (which
        # happens to be `np.max(aac) // 2`; otherwise
        # standardize_by_allele_count crashes in this case too.
        n_bins = min(np.max(aac) // 2, sum(~nan_score))

        with np.errstate(divide='ignore', invalid='ignore'):
            score_standardized, _ = allel.standardize_by_allele_count(
                    score, aac, n_bins=n_bins, diagnostics=False)

    return score_standardized
