##############################################################################
# pymbar: A Python Library for MBAR
#
# Copyright 2010-2014 University of Virginia, Memorial Sloan-Kettering Cancer Center
#
# Authors: Michael Shirts, John Chodera
# Contributors: Kyle Beauchamp
#
# pymbar is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with pymbar. If not, see <http://www.gnu.org/licenses/>.
##############################################################################

"""
A module for extracting uncorrelated samples from correlated timeseries data.

This module provides various tools that allow one to examine the correlation functions and
integrated autocorrelation times in correlated timeseries data, compute statistical inefficiencies,
and automatically extract uncorrelated samples for data analysis.

Please reference the following if you use this code in your research:

[1] Shirts MR and Chodera JD. Statistically optimal analysis of samples from multiple equilibrium states.
J. Chem. Phys. 129:124105, 2008
http://dx.doi.org/10.1063/1.2978177

[2] J. D. Chodera, W. C. Swope, J. W. Pitera, C. Seok, and K. A. Dill. Use of the weighted
histogram analysis method for the analysis of simulated and parallel tempering simulations.
JCTC 3(1):26-41, 2007.

"""

#=============================================================================================
# COPYRIGHT NOTICE
#
# Written by John D. Chodera <jchodera@gmail.com> and Michael R. Shirts <mrshirts@gmail.com>.
#
# Copyright (c) 2007 The Regents of the University of California.  All Rights Reserved.
# Portions of this software are Copyright (c) 2007 Stanford University and Columbia University.
#
# This program is free software; you can redistribute it and/or modify it under the terms of
# the GNU General Public License as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
# without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with this program;
# if not, write to the Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
# Boston, MA  02110-1301, USA.
#=============================================================================================

#=============================================================================================
# TODO
# * Implement unit tests that generate timeseries with various levels of Gaussian correlation to test all methods.
# * Add Zwanzig procedure for estimating statistical uncertainties in correlation functions
# (by making Gaussian process assumptions).
#=============================================================================================


__authors__ = "Michael R. Shirts and John D. Chodera."
__license__ = "GPL 2.0"

#=============================================================================================
# IMPORTS
#=============================================================================================
import math
import numpy
import numpy.linalg
from pymbar.utils import ParameterError

#=============================================================================================
# Issue warning on import.
#=============================================================================================

LongWarning = "Warning on use of the timeseries module: If the inherent timescales of the system are long compared to those being analyzed, this statistical inefficiency may be an underestimate.  The estimate presumes the use of many statistically independent samples.  Tests should be performed to assess whether this condition is satisfied.   Be cautious in the interpretation of the data."

#sys.stderr.write(LongWarning + '\n')

#=============================================================================================
# METHODS
#=============================================================================================

#=============================================================================================


def statisticalInefficiency(A_n, B_n=None, fast=False, mintime=3):
    """Compute the (cross) statistical inefficiency of (two) timeseries.

    Parameters
    ----------
    u_kln : np.ndarray, float, shape=(K, L, N_max)
        u_kln[k,l,n] is the reduced potential energy of uncorrelated
        configuration n sampled from state k, evaluated at state l.  K >= L             
    A_n : np.ndarray, float
        A_n[n] is nth value of timeseries A.  Length is deduced from vector.
    B_n : np.ndarray, float, optional
        B_n[n] is nth value of timeseries B.  Length is deduced from vector.
        If supplied, the cross-correlation of timeseries A and B will be estimated instead of the
        autocorrelation of timeseries A.  
    fast : bool, optional, default=False
        f True, will use faster (but less accurate) method to estimate correlation
        time, described in Ref. [1] (default: False)
    mintime : int, optional, default=3
        minimum amount of correlation function to compute (default: 3)
        The algorithm terminates after computing the correlation time out to mintime when the
        correlation function furst goes negative.  Note that this time may need to be increased
        if there is a strong initial negative peak in the correlation function.

    Returns
    -------
    g : np.ndarray,
        g is the estimated statistical inefficiency (equal to 1 + 2 tau, where tau is the correlation time).
        We enforce g >= 1.0.

    Notes
    -----
    The same timeseries can be used for both A_n and B_n to get the autocorrelation statistical inefficiency.
    The fast method described in Ref [1] is used to compute g.

    References
    ----------
    [1] J. D. Chodera, W. C. Swope, J. W. Pitera, C. Seok, and K. A. Dill. Use of the weighted
        histogram analysis method for the analysis of simulated and parallel tempering simulations.
        JCTC 3(1):26-41, 2007.

    Examples
    --------

    Compute statistical inefficiency of timeseries data with known correlation time.  

    >>> from pymbar.testsystems import correlated_timeseries_example
    >>> A_n = correlated_timeseries_example(N=100000, tau=5.0)
    >>> g = statisticalInefficiency(A_n, fast=True)

    """

    # Create numpy copies of input arguments.
    A_n = numpy.array(A_n)
    if B_n is not None:
        B_n = numpy.array(B_n)
    else:
        B_n = numpy.array(A_n)

    # Get the length of the timeseries.
    N = A_n.size

    # Be sure A_n and B_n have the same dimensions.
    if(A_n.shape != B_n.shape):
        raise ParameterError('A_n and B_n must have same dimensions.')

    # Initialize statistical inefficiency estimate with uncorrelated value.
    g = 1.0

    # Compute mean of each timeseries.
    mu_A = A_n.mean()
    mu_B = B_n.mean()

    # Make temporary copies of fluctuation from mean.
    dA_n = A_n.astype(numpy.float64) - mu_A
    dB_n = B_n.astype(numpy.float64) - mu_B

    # Compute estimator of covariance of (A,B) using estimator that will ensure C(0) = 1.
    sigma2_AB = (dA_n * dB_n).mean()  # standard estimator to ensure C(0) = 1

    # Trap the case where this covariance is zero, and we cannot proceed.
    if(sigma2_AB == 0):
        raise ParameterError('Sample covariance sigma_AB^2 = 0 -- cannot compute statistical inefficiency')

    # Accumulate the integrated correlation time by computing the normalized correlation time at
    # increasing values of t.  Stop accumulating if the correlation function goes negative, since
    # this is unlikely to occur unless the correlation function has decayed to the point where it
    # is dominated by noise and indistinguishable from zero.
    t = 1
    increment = 1
    while (t < N - 1):

        # compute normalized fluctuation correlation function at time t
        C = numpy.sum(dA_n[0:(N - t)] * dB_n[t:N] + dB_n[0:(N - t)] * dA_n[t:N]) / (2.0 * float(N - t) * sigma2_AB)
        # Terminate if the correlation function has crossed zero and we've computed the correlation
        # function at least out to 'mintime'.
        if (C <= 0.0) and (t > mintime):
            break

        # Accumulate contribution to the statistical inefficiency.
        g += 2.0 * C * (1.0 - float(t) / float(N)) * float(increment)

        # Increment t and the amount by which we increment t.
        t += increment

        # Increase the interval if "fast mode" is on.
        if fast:
            increment += 1

    # g must be at least unity
    if (g < 1.0):
        g = 1.0

    # Return the computed statistical inefficiency.
    return g
#=============================================================================================


def statisticalInefficiencyMultiple(A_kn, fast=False, return_correlation_function=False):
    """Estimate the statistical inefficiency from multiple stationary timeseries (of potentially differing lengths).

    Parameters
    ----------
    A_kn : list of np.ndarrays
        A_kn[k] is the kth timeseries, and A_kn[k][n] is nth value of timeseries k.  Length is deduced from arrays.

    fast : bool, optional, default=False
        f True, will use faster (but less accurate) method to estimate correlation
        time, described in Ref. [1] (default: False)
    return_correlation_function : bool, optional, default=False
        if True, will also return estimates of normalized fluctuation correlation function that were computed (default: False)

    Returns
    -------
    g : np.ndarray,
        g is the estimated statistical inefficiency (equal to 1 + 2 tau, where tau is the correlation time).
        We enforce g >= 1.0.
    Ct : list (of tuples)
        Ct[n] = (t, C) with time t and normalized correlation function estimate C is returned as well if return_correlation_function is set to True

    Notes
    -----
    The autocorrelation of the timeseries is used to compute the statistical inefficiency.
    The normalized fluctuation autocorrelation function is computed by averaging the unnormalized raw correlation functions.
    The fast method described in Ref [1] is used to compute g.

    References
    ----------
    [1] J. D. Chodera, W. C. Swope, J. W. Pitera, C. Seok, and K. A. Dill. Use of the weighted
        histogram analysis method for the analysis of simulated and parallel tempering simulations.
        JCTC 3(1):26-41, 2007.

    Examples
    --------

    Estimate statistical efficiency from multiple timeseries of different lengths.

    >>> from pymbar import testsystems
    >>> N_k = [1000, 2000, 3000, 4000, 5000]
    >>> tau = 5.0 # exponential relaxation time
    >>> A_kn = [ testsystems.correlated_timeseries_example(N=N, tau=tau) for N in N_k ]
    >>> g = statisticalInefficiencyMultiple(A_kn)

    Also return the values of the normalized fluctuation autocorrelation function that were computed.

    >>> [g, Ct] = statisticalInefficiencyMultiple(A_kn, return_correlation_function=True)

    """

    # Convert A_kn into a list of arrays if it is not in this form already.
    if (type(A_kn) == numpy.ndarray):
        A_kn_list = list()
        if A_kn.ndim == 1:
            A_kn_list.append(A_kn.copy())
        else:
            [K, N] = A_kn.shape
            for k in range(K):
                A_kn_list.append(A_kn[k, :].copy())
        A_kn = A_kn_list

    # Determine number of timeseries.
    K = len(A_kn)

    # Get the length of each timeseries.
    N_k = numpy.zeros([K], numpy.int32)
    for k in range(K):
        N_k[k] = A_kn[k].size

    # Compute average timeseries length.
    Navg = numpy.array(N_k, numpy.float64).mean()

    # Determine total number of samples.
    N = numpy.sum(N_k)

    # Initialize statistical inefficiency estimate with uncorrelated value.
    g = 1.0

    # Compute sample mean.
    mu = 0.0
    for k in range(K):
        mu += numpy.sum(A_kn[k])
    mu /= float(N)

    # Construct and store fluctuation timeseries.
    dA_kn = list()
    for k in range(K):
        dA_n = A_kn[k] - mu
        dA_kn.append(dA_n.copy())

    # Compute sample variance from mean of squared fluctuations, to ensure that C(0) = 1.
    sigma2 = 0.0
    for k in range(K):
        sigma2 += numpy.sum(dA_kn[k] ** 2)
    sigma2 /= float(N)

    # Initialize statistical inefficiency estimate with uncorrelated value.
    g = 1.0

    # Initialize storage for correlation function.
    Ct = list()  # Ct[n] is a tuple (t, C) of the time lag t and estimate of normalized fluctuation correlation function C

    # Accumulate the integrated correlation time by computing the normalized correlation time at
    # increasing values of t.  Stop accumulating if the correlation function goes negative, since
    # this is unlikely to occur unless the correlation function has decayed to the point where it
    # is dominated by noise and indistinguishable from zero.
    t = 1
    increment = 1
    while (t < N_k.max() - 1):
        # compute unnormalized correlation function
        numerator = 0.0
        denominator = 0.0
        for k in range(K):
            if (t >= N_k[k]):
                continue  # skip trajectory if lag time t is greater than its length
            dA_n = dA_kn[k]  # retrieve trajectory
            x = dA_n[0:(N_k[k] - t)] * dA_n[t:N_k[k]]
            numerator += numpy.sum(x)  # accumulate contribution from trajectory k
            denominator += float(x.size)  # count how many overlapping time segments we've included

        C = numerator / denominator

        # compute normalized fluctuation correlation function at time t
        C = C / sigma2
        # print "C[%5d] = %16f (%16f / %16f)" % (t, C, numerator, denominator)

        # Store estimate of correlation function.
        Ct.append((t, C))

        # Terminate if the correlation function has crossed zero.
        # Note that we've added a hack (t > 10) condition to avoid terminating too early in correlation functions that have a strong negative peak at
        if (C <= 0.0) and (t > 10):
            break

        # Accumulate contribution to the statistical inefficiency.
        g += 2.0 * C * (1.0 - float(t) / Navg) * float(increment)

        # Increment t and the amount by which we increment t.
        t += increment

        # Increase the interval if "fast mode" is on.
        if fast:
            increment += 1

    # g must be at least unity
    if (g < 1.0):
        g = 1.0

    # Return statistical inefficency and correlation function estimate, if requested.
    if return_correlation_function:
        return (g, Ct)

    # Return the computed statistical inefficiency.
    return g
#=============================================================================================


def integratedAutocorrelationTime(A_n, B_n=None, fast=False, mintime=3):
    """Estimate the integrated autocorrelation time."""

    g = statisticalInefficiency(A_n, B_n, fast, mintime)
    tau = (g - 1.0) / 2.0
    return tau
#=============================================================================================


def integratedAutocorrelationTimeMultiple(A_kn, fast=False):
    """Estimate the integrated autocorrelation time from multiple timeseries."""

    g = statisticalInefficiencyMultiple(A_kn, fast, False)
    tau = (g - 1.0) / 2.0
    return tau
#=============================================================================================


def normalizedFluctuationCorrelationFunction(A_n, B_n=None, N_max=None):
    """Compute the normalized fluctuation (cross) correlation function of (two) stationary timeseries.

    C(t) = (<A(t) B(t)> - <A><B>) / (<AB> - <A><B>)

    This may be useful in diagnosing odd time-correlations in timeseries data.

    Parameters
    ----------
    A_n : np.ndarray
        A_n[n] is nth value of timeseries A.  Length is deduced from vector.
    B_n : np.ndarray
        B_n[n] is nth value of timeseries B.  Length is deduced from vector.
    N_max : int, default=None
        if specified, will only compute correlation function out to time lag of N_max

    Returns
    -------
    C_n : np.ndarray
        C_n[n] is the normalized fluctuation auto- or cross-correlation function for timeseries A(t) and B(t).

    Notes
    -----
    The same timeseries can be used for both A_n and B_n to get the autocorrelation statistical inefficiency.
    This procedure may be slow.
    The statistical error in C_n[n] will grow with increasing n.  No effort is made here to estimate the uncertainty.

    References
    ----------
    [1] J. D. Chodera, W. C. Swope, J. W. Pitera, C. Seok, and K. A. Dill. Use of the weighted
    histogram analysis method for the analysis of simulated and parallel tempering simulations.
    JCTC 3(1):26-41, 2007.

    Examples
    --------

    Estimate normalized fluctuation correlation function.

    >>> from pymbar import testsystems
    >>> A_t = testsystems.correlated_timeseries_example(N=10000, tau=5.0)
    >>> C_t = normalizedFluctuationCorrelationFunction(A_t, N_max=25)

    """

    # If B_n is not specified, set it to be identical to A_n.
    if B_n is None:
        B_n = A_n

    # Create numpy copies of input arguments.
    A_n = numpy.array(A_n)
    B_n = numpy.array(B_n)

    # Get the length of the timeseries.
    N = A_n.size

    # Set maximum time to compute correlation functon for.
    if (not N_max) or (N_max > N - 1):
        N_max = N - 1

    # Be sure A_n and B_n have the same dimensions.
    if(A_n.shape != B_n.shape):
        raise ParameterError('A_n and B_n must have same dimensions.')

    # Initialize statistical inefficiency estimate with uncorrelated value.
    g = 1.0

    # Compute means and variance.
    mu_A = A_n.mean()
    mu_B = B_n.mean()

    # Make temporary copies at high precision with means subtracted off.
    dA_n = A_n.astype(numpy.float64) - mu_A
    dB_n = B_n.astype(numpy.float64) - mu_B

    # sigma2_AB = sum((A_n-mu_A) * (B_n-mu_B)) / (float(N)-1.0) # unbiased estimator
    sigma2_AB = (dA_n * dB_n).mean()  # standard estimator to ensure C(0) = 1
    if(sigma2_AB == 0):
        raise ParameterError('Sample covariance sigma_AB^2 = 0 -- cannot compute statistical inefficiency')

    # allocate storage for normalized fluctuation correlation function
    C_n = numpy.zeros([N_max + 1], numpy.float64)

    # Compute normalized correlation funtion.
    t = 0
    for t in range(0, N_max + 1):
        # compute normalized fluctuation correlation function at time t
        C_n[t] = numpy.sum(dA_n[0:(N - t)] * dB_n[t:N] + dB_n[0:(N - t)] * dA_n[t:N]) / (2.0 * float(N - t) * sigma2_AB)

    # Return the computed correlation function
    return C_n
#=============================================================================================


def normalizedFluctuationCorrelationFunctionMultiple(A_kn, B_kn=None, N_max=None):
    """Compute the normalized fluctuation (cross) correlation function of (two) timeseries from multiple timeseries samples.

    C(t) = (<A(t) B(t)> - <A><B>) / (<AB> - <A><B>)
    This may be useful in diagnosing odd time-correlations in timeseries data.

    Parameters
    ----------
    A_kn : Python list of numpy arrays
        A_kn[k] is the kth timeseries, and A_kn[k][n] is nth value of timeseries k.  Length is deduced from arrays.
    B_kn : Python list of numpy arrays
        B_kn[k] is the kth timeseries, and B_kn[k][n] is nth value of timeseries k.  B_kn[k] must have same length as A_kn[k]
    N_max : int, optional, default=None
        if specified, will only compute correlation function out to time lag of N_max

    Returns
    -------
    C_n[n] : np.ndarray
        The normalized fluctuation auto- or cross-correlation function for timeseries A(t) and B(t).

    Notes
    -----
    The same timeseries can be used for both A_n and B_n to get the autocorrelation statistical inefficiency.
    This procedure may be slow.
    The statistical error in C_n[n] will grow with increasing n.  No effort is made here to estimate the uncertainty.

    References
    ----------
    [1] J. D. Chodera, W. C. Swope, J. W. Pitera, C. Seok, and K. A. Dill. Use of the weighted
    histogram analysis method for the analysis of simulated and parallel tempering simulations.
    JCTC 3(1):26-41, 2007.

    Examples
    --------

    Estimate a portion of the normalized fluctuation autocorrelation function from multiple timeseries of different length.

    >>> from pymbar import testsystems
    >>> N_k = [1000, 2000, 3000, 4000, 5000]
    >>> tau = 5.0 # exponential relaxation time
    >>> A_kn = [ testsystems.correlated_timeseries_example(N=N, tau=tau) for N in N_k ]
    >>> C_n = normalizedFluctuationCorrelationFunctionMultiple(A_kn, N_max=25)

    """

    # If B_kn is not specified, define it to be identical with A_kn.
    if B_kn is None:
        B_kn = A_kn

    # TODO: Change this to support other iterable types, like sets.
    # Make sure A_kn and B_kn are both lists
    if (type(A_kn) is not list) or (type(B_kn) is not list):
        raise ParameterError("A_kn and B_kn must each be a list of numpy arrays.")

    # Ensure the same number of timeseries are stored in A_kn and B_kn.
    if (len(A_kn) != len(B_kn)):
        raise ParameterError("A_kn and B_kn must contain corresponding timeseries -- different numbers of timeseries detected in each.")

    # Determine number of timeseries stored.
    K = len(A_kn)

    # Ensure both observable trajectories in each timeseries are of the same length.
    for k in range(K):
        A_n = A_kn[k]
        B_n = B_kn[k]
        if A_n.size != B_n.size:
            raise ParameterError("A_kn and B_kn must contain corresponding timeseries -- lack of correspondence in timeseries lenghts detected.")

    # Get the length of each timeseries.
    N_k = numpy.zeros([K], numpy.int32)
    for k in range(K):
        N_k[k] = A_kn[k].size

    # Determine total number of samples.
    N = numpy.sum(N_k)

    # Set maximum time to compute correlation functon for.
    if (not N_max) or (N_max > max(N_k) - 1):
        N_max = max(N_k) - 1

    # Compute means.
    mu_A = 0.0
    mu_B = 0.0
    for k in range(K):
        mu_A += numpy.sum(A_kn[k])
        mu_B += numpy.sum(B_kn[k])
    mu_A /= float(N)
    mu_B /= float(N)

    # Compute fluctuation timeseries.
    dA_kn = list()
    dB_kn = list()
    for k in range(K):
        dA_n = A_kn[k] - mu_A
        dB_n = B_kn[k] - mu_B
        dA_kn.append(dA_n)
        dB_kn.append(dB_n)

    # Compute covariance.
    sigma2_AB = 0.0
    for k in range(K):
        sigma2_AB += numpy.sum(dA_kn[k] * dB_kn[k])
    sigma2_AB /= float(N)

    # allocate storage for normalized fluctuation correlation function
    C_n = numpy.zeros([N_max + 1], numpy.float64)

    # Accumulate the integrated correlation time by computing the normalized correlation time at
    # increasing values of t.  Stop accumulating if the correlation function goes negative, since
    # this is unlikely to occur unless the correlation function has decayed to the point where it
    # is dominated by noise and indistinguishable from zero.
    t = 0
    for t in range(0, N_max + 1):
        # compute unnormalized correlation function
        numerator = 0.0
        denominator = 0.0
        for k in range(K):
            if (t >= N_k[k]):
                continue  # skip this trajectory if t is longer than the timeseries
            numerator += numpy.sum(dA_kn[k][0:(N_k[k] - t)] * dB_kn[k][t:N_k[k]])
            denominator += float(N_k[k] - t)
        C = numerator / denominator

        # compute normalized fluctuation correlation function at time t
        C /= sigma2_AB

        # Store correlation function.
        C_n[t] = C

    # Return the computed fluctuation correlation function.
    return C_n
#=============================================================================================


def subsampleCorrelatedData(A_t, g=None, fast=False, conservative=False, verbose=False):
    """Determine the indices of an uncorrelated subsample of the data.

    Parameters
    ----------
    A_t : np.ndarray
        A_t[t] is the t-th value of timeseries A(t).  Length is deduced from vector.
    g : float, optional
        if provided, the statistical inefficiency g is used to subsample the timeseries -- otherwise it will be computed (default: None)
    fast : bool, optional, default=False
        fast can be set to True to give a less accurate but very quick estimate (default: False)
    conservative : bool, optional, default=False
        if set to True, uniformly-spaced indices are chosen with interval ceil(g), where
        g is the statistical inefficiency.  Otherwise, indices are chosen non-uniformly with interval of
        approximately g in order to end up with approximately T/g total indices
    verbose : bool, optional, default=False
        if True, some output is printed

    Returns
    -------
    indices : list of int
        the indices of an uncorrelated subsample of the data

    Notes
    -----
    The statistical inefficiency is computed with the function computeStatisticalInefficiency().

    ToDo
    ----
    Instead of using regular stride, use irregular stride so more data can be fit in when g is non-integral.

    Examples
    --------

    Subsample a correlated timeseries to extract an effectively uncorrelated dataset.

    >>> from pymbar import testsystems
    >>> A_t = testsystems.correlated_timeseries_example(N=10000, tau=5.0) # generate a test correlated timeseries
    >>> indices = subsampleCorrelatedData(A_t) # compute indices of uncorrelated timeseries
    >>> A_n = A_t[indices] # extract uncorrelated samples

    Extract uncorrelated samples from multiple timeseries data from the same process.

    >>> # Generate multiple correlated timeseries data of different lengths.
    >>> T_k = [1000, 2000, 3000, 4000, 5000]
    >>> K = len(T_k) # number of timeseries
    >>> tau = 5.0 # exponential relaxation time
    >>> A_kt = [ testsystems.correlated_timeseries_example(N=T, tau=tau) for T in T_k ] # A_kt[k] is correlated timeseries k
    >>> # Estimate statistical inefficiency from all timeseries data.
    >>> g = statisticalInefficiencyMultiple(A_kt)
    >>> # Count number of uncorrelated samples in each timeseries.
    >>> N_k = numpy.array([ len(subsampleCorrelatedData(A_t, g=g)) for A_t in A_kt ]) # N_k[k] is the number of uncorrelated samples in timeseries k
    >>> N = N_k.sum() # total number of uncorrelated samples
    >>> # Subsample all trajectories to produce uncorrelated samples
    >>> A_kn = [ A_t[subsampleCorrelatedData(A_t, g=g)] for A_t in A_kt ] # A_kn[k] is uncorrelated subset of trajectory A_kt[t]
    >>> # Concatenate data into one timeseries.
    >>> A_n = numpy.zeros([N], numpy.float32) # A_n[n] is nth sample in concatenated set of uncorrelated samples
    >>> A_n[0:N_k[0]] = A_kn[0]
    >>> for k in range(1,K): A_n[N_k[0:k].sum():N_k[0:k+1].sum()] = A_kn[k]

    """

    # Create numpy copy of arrays.
    A_t = numpy.array(A_t)

    # Get the length of the timeseries.
    T = A_t.size

    # Compute the statistical inefficiency for the timeseries.
    if not g:
        if verbose:
            print "Computing statistical inefficiency..."
        g = statisticalInefficiency(A_t, A_t, fast=fast)
        if verbose:
            print "g = %f" % g

    if conservative:
        # Round g up to determine the stride we can use to pick out regularly-spaced uncorrelated samples.
        stride = int(math.ceil(g))
        if verbose:
            print "conservative subsampling: using stride of %d" % stride

        # Assemble list of indices of uncorrelated snapshots.
        indices = range(0, T, stride)
    else:
        # Choose indices as floor(n*g), with n = 0,1,2,..., until we run out of data.
        indices = []
        n = 0
        while int(round(n * g)) < T:
            t = int(round(n * g))
            # ensure we don't sample the same point twice
            if (n == 0) or (t != indices[n - 1]):
                indices.append(t)
            n += 1
        if verbose:
            print "standard subsampling: using average stride of %f" % g

    # Number of samples in subsampled timeseries.
    N = len(indices)

    if verbose:
        print "The resulting subsampled set has %d samples (original timeseries had %d)." % (N, T)

    # Return the list of indices of uncorrelated snapshots.
    return indices


def detectEquilibration(A_t, fast=True, nskip=1):
    """Automatically detect equilibrated region of a dataset using a heuristic that maximizes number of effectively uncorrelated samples.

    Parameters
    ----------
    A_t : np.ndarray 
        timeseries
    nskip : int, optional, default=1
        number of samples to sparsify data by in order to speed equilibration detection

    Returns
    -------
    t : int
        start of equilibrated data
    g : float
        statistical inefficiency of equilibrated data
    Neff_max : float
        number of uncorrelated samples   

    ToDo
    ----
    Consider implementing a binary search for Neff_max.

    Examples
    --------

    Determine start of equilibrated data for a correlated timeseries.

    >>> from pymbar import testsystems
    >>> A_t = testsystems.correlated_timeseries_example(N=1000, tau=5.0) # generate a test correlated timeseries
    >>> [t, g, Neff_max] = detectEquilibration(A_t) # compute indices of uncorrelated timeseries

    Determine start of equilibrated data for a correlated timeseries with a shift.

    >>> from pymbar import testsystems
    >>> A_t = testsystems.correlated_timeseries_example(N=1000, tau=5.0) + 2.0 # generate a test correlated timeseries
    >>> B_t = testsystems.correlated_timeseries_example(N=10000, tau=5.0) # generate a test correlated timeseries
    >>> C_t = numpy.concatenate([A_t, B_t])
    >>> [t, g, Neff_max] = detectEquilibration(C_t, nskip=50) # compute indices of uncorrelated timeseries

    """
    T = A_t.size

    # Special case if timeseries is constant.
    if A_t.std() == 0.0:
        return (0, 1, T)

    g_t = numpy.ones([T - 1], numpy.float32)
    Neff_t = numpy.ones([T - 1], numpy.float32)
    for t in range(0, T - 1, nskip):
        g_t[t] = statisticalInefficiency(A_t[t:T], fast=fast)
        Neff_t[t] = (T - t + 1) / g_t[t]

    Neff_max = Neff_t.max()
    t = Neff_t.argmax()
    g = g_t[t]

    return (t, g, Neff_max)
