/**Pearson, Spearman and Kendall correlations, covariance.
 *
 * Author:  David Simcha*/
 /*
 * License:
 * Boost Software License - Version 1.0 - August 17th, 2003
 *
 * Permission is hereby granted, free of charge, to any person or organization
 * obtaining a copy of the software and accompanying documentation covered by
 * this license (the "Software") to use, reproduce, display, distribute,
 * execute, and transmit the Software, and to prepare derivative works of the
 * Software, and to permit third-parties to whom the Software is furnished to
 * do so, all subject to the following:
 *
 * The copyright notices in the Software and this entire statement, including
 * the above license grant, this restriction and the following disclaimer,
 * must be included in all copies of the Software, in whole or in part, and
 * all derivative works of the Software, unless such copies or derivative
 * works are solely in the form of machine-executable object code generated by
 * a source language processor.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
 * SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
 * FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

module dstats.cor;

import std.conv, std.range, std.typecons, std.exception, std.math,
    std.traits, std.typetuple, std.algorithm, std.parallelism, std.numeric;

import dstats.sort, dstats.base, dstats.alloc, dstats.summary;

version(unittest) {
    import std.stdio, dstats.random;

    Random gen;

    shared static this()
    {
        gen.seed(unpredictableSeed);
    }
}

/**Convenience function for calculating Pearson correlation.
 * When the term correlation is used unqualified, it is
 * usually referring to this quantity.  This is a parametric correlation
 * metric and should not be used with extremely ill-behaved data.
 * This function works with any pair of input ranges.
 *
 * Note:  The PearsonCor struct returned by this function is alias this'd to the
 * correlation coefficient.  Therefore, the result from this function can
 * be treated simply as a floating point number.
 *
 * References: Computing Higher-Order Moments Online.
 * http://people.xiph.org/~tterribe/notes/homs.html
 */
PearsonCor pearsonCor(T, U)(T input1, U input2)
if(doubleInput!(T) && doubleInput!(U)) {
    PearsonCor corCalc;

    static if(isRandomAccessRange!T && isRandomAccessRange!U &&
        hasLength!T && hasLength!U) {

        // ILP parallelization optimization.  Sharing a k between a bunch
        // of implicit PearsonCor structs cuts down on the amount of divisions
        // necessary.  Using nILP of them instead of one improves CPU pipeline
        // performance by reducing data dependency.  When the stack is
        // properly aligned, this can result in about 2x speedups compared
        // to simply submitting everything to a single PearsonCor struct.
        dstatsEnforce(input1.length == input2.length,
            "Ranges must be same length for Pearson correlation.");

        enum nILP = 8;
        size_t i = 0;
        if(input1.length > 2 * nILP) {

            double _k = 0;
            double[nILP] _mean1 = 0, _mean2 = 0, _var1 = 0, _var2 = 0, _cov = 0;

            for(; i + nILP < input1.length; i += nILP) {
                immutable kMinus1 = _k;
                immutable kNeg1 = 1 / ++_k;

                foreach(j; StaticIota!nILP) {
                    immutable double delta1 = input1[i + j] - _mean1[j];
                    immutable double delta2 = input2[i + j] - _mean2[j];
                    immutable delta1N = delta1 * kNeg1;
                    immutable delta2N = delta2 * kNeg1;

                    _mean1[j] += delta1N;
                    _var1[j]  += kMinus1 * delta1N * delta1;
                    _cov[j]   += kMinus1 * delta1N * delta2;
                    _var2[j]  += kMinus1 * delta2N * delta2;
                    _mean2[j] += delta2N;
                }
            }

            corCalc._k = _k;
            corCalc._mean1 = _mean1[0];
            corCalc._mean2 = _mean2[0];
            corCalc._var1 = _var1[0];
            corCalc._var2 = _var2[0];
            corCalc._cov = _cov[0];

            foreach(j; 1..nILP) {
                auto tmp = PearsonCor(_k, _mean1[j], _mean2[j],
                        _var1[j], _var2[j], _cov[j]);
                corCalc.put(tmp);
            }
        }

        // Handle remainder.
        for(; i < input1.length; i++) {
            corCalc.put(input1[i], input2[i]);
        }

    } else {
        while(!input1.empty && !input2.empty) {
            corCalc.put(input1.front, input2.front);
            input1.popFront;
            input2.popFront;
        }

        dstatsEnforce(input1.empty && input2.empty,
            "Ranges must be same length for Pearson correlation.");
    }

    return corCalc;
}

unittest {
    assert(approxEqual(pearsonCor([1,2,3,4,5][], [1,2,3,4,5][]).cor, 1));
    assert(approxEqual(pearsonCor([1,2,3,4,5][], [10.0, 8.0, 6.0, 4.0, 2.0][]).cor, -1));
    assert(approxEqual(pearsonCor([2, 4, 1, 6, 19][], [4, 5, 1, 3, 2][]).cor, -.2382314));

        // Make sure everything works with lowest common denominator range type.
    static struct Count {
        uint num;
        uint upTo;
        @property size_t front() {
            return num;
        }
        void popFront() {
            num++;
        }
        @property bool empty() {
            return num >= upTo;
        }
    }

    Count a, b;
    a.upTo = 100;
    b.upTo = 100;
    assert(approxEqual(pearsonCor(a, b).cor, 1));

    PearsonCor cor1 = pearsonCor([1,2,4][], [2,3,5][]);
    PearsonCor cor2 = pearsonCor([4,2,9][], [2,8,7][]);
    PearsonCor combined = pearsonCor([1,2,4,4,2,9][], [2,3,5,2,8,7][]);

    cor1.put(cor2);

    foreach(ti, elem; cor1.tupleof) {
        assert(approxEqual(elem, combined.tupleof[ti]));
    }

    assert(approxEqual(pearsonCor([1,2,3,4,5,6,7,8,9,10][],
        [8,6,7,5,3,0,9,3,6,2][]).cor, -0.4190758));

    foreach(iter; 0..1000) {
        // Make sure results for the ILP-optimized and non-optimized versions
        // agree.
        auto foo = randArray!(rNormal)(uniform(5, 100), 0, 1);
        auto bar = randArray!(rNormal)(foo.length, 0, 1);
        auto res1 = pearsonCor(foo, bar);
        PearsonCor res2;
        foreach(i; 0..foo.length) {
            res2.put(foo[i], bar[i]);
        }

        foreach(ti, elem; res1.tupleof) {
            assert(approxEqual(elem, res2.tupleof[ti]));
        }

        PearsonCor resCornerCase;  // Test where one N is zero.
        resCornerCase.put(res1);
        PearsonCor dummy;
        resCornerCase.put(dummy);
        foreach(ti, elem; res1.tupleof) {
            assert(isIdentical(resCornerCase.tupleof[ti], elem));
        }
    }
}

/**Allows computation of mean, stdev, variance, covariance, Pearson correlation online.
 * Getters for stdev, var, cov, cor cost floating point division ops.  Getters
 * for means cost a single branch to check for N == 0.  This struct uses O(1)
 * space.
 *
 * PearsonCor.cor is alias this'd, so if this struct is used as a float, it will
 * be converted to a simple correlation coefficient automatically.
 *
 * Bugs:  Alias this disabled due to compiler bugs.
 *
 * References: Computing Higher-Order Moments Online.
 * http://people.xiph.org/~tterribe/notes/homs.html
 */
struct PearsonCor {
package:
    double _k = 0, _mean1 = 0, _mean2 = 0, _var1 = 0, _var2 = 0, _cov = 0;

public:
    alias cor this;

    ///
    void put(double elem1, double elem2) nothrow @safe {
        immutable kMinus1 = _k;
        immutable kNeg1 = 1 / ++_k;
        immutable delta1 = elem1 - _mean1;
        immutable delta2 = elem2 - _mean2;
        immutable delta1N = delta1 * kNeg1;
        immutable delta2N = delta2 * kNeg1;

        _mean1 += delta1N;
        _var1  += kMinus1 * delta1N * delta1;
        _cov   += kMinus1 * delta1N * delta2;
        _var2  += kMinus1 * delta2N * delta2;
        _mean2 += delta2N;
    }

    /// Combine two PearsonCor's.
    void put(const ref typeof(this) rhs) nothrow @safe {
        if(_k == 0) {
            foreach(ti, elem; rhs.tupleof) {
                this.tupleof[ti] = elem;
            }
            return;
        } else if(rhs._k == 0) {
            return;
        }

        immutable totalN = _k + rhs._k;
        immutable delta1 = rhs.mean1 - mean1;
        immutable delta2 = rhs.mean2 - mean2;

        _mean1 = _mean1 * (_k / totalN) + rhs._mean1 * (rhs._k / totalN);
        _mean2 = _mean2 * (_k / totalN) + rhs._mean2 * (rhs._k / totalN);

        _var1 = _var1 + rhs._var1 + (_k / totalN * rhs._k * delta1 * delta1 );
        _var2 = _var2 + rhs._var2 + (_k / totalN * rhs._k * delta2 * delta2 );
        _cov  =  _cov + rhs._cov  + (_k / totalN * rhs._k * delta1 * delta2 );
        _k = totalN;
    }

    const pure nothrow @property @safe {

        ///
        double var1() {
            return (_k < 2) ? double.nan : _var1 / (_k - 1);
        }

        ///
        double var2() {
            return (_k < 2) ? double.nan : _var2 / (_k - 1);
        }

        ///
        double stdev1() {
            return sqrt(var1);
        }

        ///
        double stdev2() {
            return sqrt(var2);
        }

        ///
        double cor() {
            return cov / stdev1 / stdev2;
        }

        ///
        double cov() {
            return (_k < 2) ? double.nan : _cov / (_k - 1);
        }

        ///
        double mean1() {
            return (_k == 0) ? double.nan : _mean1;
        }

        ///
        double mean2() {
            return (_k == 0) ? double.nan : _mean2;
        }

        ///
        double N() {
            return _k;
        }

    }
}

///
double covariance(T, U)(T input1, U input2)
if(doubleInput!(T) && doubleInput!(U)) {
    return pearsonCor(input1, input2).cov;
}

unittest {
    assert(approxEqual(covariance([1,4,2,6,3].dup, [3,1,2,6,2].dup), 2.05));
}

/**Spearman's rank correlation.  Non-parametric.  This is essentially the
 * Pearson correlation of the ranks of the data, with ties dealt with by
 * averaging.*/
double spearmanCor(R, S)(R input1, S input2)
if(isInputRange!(R) && isInputRange!(S) &&
is(typeof(input1.front < input1.front) == bool) &&
is(typeof(input2.front < input2.front) == bool)) {

    static if(hasLength!S && hasLength!R) {
        if(input1.length < 2) {
            return double.nan;
        }
    }

    auto alloc = newRegionAllocator();

    double[] spearmanCorRank(T)(T someRange) {
        static if(hasLength!(T) && isRandomAccessRange!(T)) {
            double[] ret = alloc.uninitializedArray!(double[])(someRange.length);
            rank(someRange, ret);
        } else {
            auto iDup = alloc.array(someRange);
            double[] ret = alloc.uninitializedArray!(double[])(iDup.length);
            rankSort(iDup, ret);
        }
        return ret;
    }

    try {
        auto ranks1 = spearmanCorRank(input1);
        auto ranks2 = spearmanCorRank(input2);
        dstatsEnforce(ranks1.length == ranks2.length,
            "Ranges must be same length for Spearman correlation.");

        return pearsonCor(ranks1, ranks2).cor;
    } catch(SortException) {
        return double.nan;
    }
}

unittest {
    //Test against a few known values.
    assert(approxEqual(spearmanCor([1,2,3,4,5,6].dup, [3,1,2,5,4,6].dup), 0.77143));
    assert(approxEqual(spearmanCor([3,1,2,5,4,6].dup, [1,2,3,4,5,6].dup ), 0.77143));
    assert(approxEqual(spearmanCor([3,6,7,35,75].dup, [1,63,53,67,3].dup), 0.3));
    assert(approxEqual(spearmanCor([1,63,53,67,3].dup, [3,6,7,35,75].dup), 0.3));
    assert(approxEqual(spearmanCor([1.5,6.3,7.8,4.2,1.5].dup, [1,63,53,67,3].dup), .56429));
    assert(approxEqual(spearmanCor([1,63,53,67,3].dup, [1.5,6.3,7.8,4.2,1.5].dup), .56429));
    assert(approxEqual(spearmanCor([1.5,6.3,7.8,7.8,1.5].dup, [1,63,53,67,3].dup), .79057));
    assert(approxEqual(spearmanCor([1,63,53,67,3].dup, [1.5,6.3,7.8,7.8,1.5].dup), .79057));
    assert(approxEqual(spearmanCor([1.5,6.3,7.8,6.3,1.5].dup, [1,63,53,67,3].dup), .63246));
    assert(approxEqual(spearmanCor([1,63,53,67,3].dup, [1.5,6.3,7.8,6.3,1.5].dup), .63246));
    assert(approxEqual(spearmanCor([3,4,1,5,2,1,6,4].dup, [1,3,2,6,4,2,6,7].dup), .6829268));
    assert(approxEqual(spearmanCor([1,3,2,6,4,2,6,7].dup, [3,4,1,5,2,1,6,4].dup), .6829268));
    uint[] one = new uint[1000], two = new uint[1000];
    foreach(i; 0..100) {  //Further sanity checks for things like commutativity.
        size_t lowerBound = uniform(0, one.length);
        size_t upperBound = uniform(0, one.length);
        if(lowerBound > upperBound) swap(lowerBound, upperBound);
        foreach(ref o; one) {
            o = uniform(1, 10);  //Generate lots of ties.
        }
        foreach(ref o; two) {
             o = uniform(1, 10);  //Generate lots of ties.
        }
        double sOne =
             spearmanCor(one[lowerBound..upperBound], two[lowerBound..upperBound]);
        double sTwo =
             spearmanCor(two[lowerBound..upperBound], one[lowerBound..upperBound]);
        foreach(ref o; one)
            o*=-1;
        double sThree =
             -spearmanCor(one[lowerBound..upperBound], two[lowerBound..upperBound]);
        double sFour =
             -spearmanCor(two[lowerBound..upperBound], one[lowerBound..upperBound]);
        foreach(ref o; two) o*=-1;
        one[lowerBound..upperBound].reverse();
        two[lowerBound..upperBound].reverse();
        double sFive =
             spearmanCor(one[lowerBound..upperBound], two[lowerBound..upperBound]);
        assert(approxEqual(sOne, sTwo) || (isNaN(sOne) && isNaN(sTwo)));
        assert(approxEqual(sTwo, sThree) || (isNaN(sThree) && isNaN(sTwo)));
        assert(approxEqual(sThree, sFour) || (isNaN(sThree) && isNaN(sFour)));
        assert(approxEqual(sFour, sFive) || (isNaN(sFour) && isNaN(sFive)));
    }

    // Test input ranges.
    static struct Count {
        uint num;
        uint upTo;
        @property size_t front() {
            return num;
        }
        void popFront() {
            num++;
        }
        @property bool empty() {
            return num >= upTo;
        }
    }

    Count a, b;
    a.upTo = 100;
    b.upTo = 100;
    assert(approxEqual(spearmanCor(a, b), 1));
}

version(unittest) {
    // Make sure when we call kendallCor, the large N version always executes.
    private enum kendallSmallN = 1;
} else {
    private enum kendallSmallN = 15;
}

package template isDefaultSorted(R) {
    static if(!is(typeof(R.init.release()))) {
        enum isDefaultSorted = false;
    } else {
        enum isDefaultSorted =
            is(R == SortedRange!(typeof(R.init.release()), "a < b"));
    }
}

unittest {
    import std.algorithm;
    auto foo = sort([1, 2, 3]);
    auto bar = sort!"a > b"([1, 2, 3]);
    static assert(isDefaultSorted!(typeof(foo)));
    static assert(!isDefaultSorted!(typeof(bar)));
    static assert(!isDefaultSorted!(typeof([1, 2, 3])));
}

/**
Kendall's Tau-b, O(N log N) version.  This is a non-parametric measure
of monotonic association and can be defined in terms of the
bubble sort distance, or the number of swaps that would be needed in a
bubble sort to sort input2 into the same order as input1.

Since a copy of the inputs is made anyhow because they need to be sorted,
this function can work with any input range.  However, the ranges must
have the same length.

Note:

As an optimization, when a range is a SortedRange with predicate "a < b",
it is assumed already sorted and not sorted a second time by this function.
This is useful when applying this function multiple times with one of the
arguments the same every time:

---
auto lhs = randArray!rNormal(1_000, 0, 1);
auto indices = new size_t[1_000];
makeIndex(lhs, indices);

foreach(i; 0..1_000) {
    auto rhs = randArray!rNormal(1_000, 0, 1);
    auto lhsSorted = assumeSorted(
        indexed(lhs, indices)
    );

    // Rearrange rhs according to the sorting permutation of lhs.
    // kendallCor(lhsSorted, rhsRearranged) will be much faster than
    // kendallCor(lhs, rhs).
    auto rhsRearranged = indexed(rhs, indices);
    assert(kendallCor(lhsSorted, rhsRearranged) == kendallCor(lhs, rhs));
}
---

References:
A Computer Method for Calculating Kendall's Tau with Ungrouped Data,
William R. Knight, Journal of the American Statistical Association, Vol.
61, No. 314, Part 1 (Jun., 1966), pp. 436-439

The Variance of Tau When Both Rankings Contain Ties.  M.G. Kendall.
Biometrika, Vol 34, No. 3/4 (Dec., 1947), pp. 297-298
 */
double kendallCor(T, U)(T input1, U input2)
if(isInputRange!(T) && isInputRange!(U)) {
    static if(isArray!(T) && isArray!(U)) {
        dstatsEnforce(input1.length == input2.length,
            "Ranges must be same length for Kendall correlation.");
        if(input1.length <= kendallSmallN) {
            return kendallCorSmallN(input1, input2);
        }
    }

    auto alloc = newRegionAllocator();

    auto prepare(V)(V range) {
        static if(isDefaultSorted!V) {
            return range;
        } else {
            return prepareForSorting!compFun(alloc.array(range));
        }
    }

    try {
        auto i1d = prepare(input1);
        auto i2d = prepare(input2);

        dstatsEnforce(i1d.length == i2d.length,
            "Ranges must be same length for Kendall correlation.");

        if(i1d.length <= kendallSmallN) {
            return kendallCorSmallN(i1d, i2d);
        } else {
            return kendallCorDestructive(i1d, i2d);
        }
    } catch(SortException) {
        return double.nan;
    }
}

/**
Kendall's Tau-b O(N log N), overwrites input arrays with undefined data but
uses only O(log N) stack space for sorting, not O(N) space to duplicate
input.  R1 and R2 must be either SortedRange structs with the default predicate
or arrays.
 */
double kendallCorDestructive(R1, R2)(R1 input1, R2 input2) {
    dstatsEnforce(input1.length == input2.length,
        "Ranges must be same length for Kendall correlation.");
    try {
        return kendallCorDestructiveLowLevel(input1, input2, false).tau;
    } catch(SortException) {
        return double.nan;
    }
}

//bool compFun(T)(T lhs, T rhs) { return lhs < rhs; }
private enum compFun = "a < b";

// Make sure arguments are in the right order, etc. to simplify implementation
// an dallow buffer recycling.
auto kendallCorDestructiveLowLevel(R1, R2)
(R1 input1, R2 input2, bool needTies) {
    static if(isDefaultSorted!R1) {
        return kendallCorDestructiveLowLevelImpl(input1, input2, needTies);
    } else static if(isDefaultSorted!R2) {
        return kendallCorDestructiveLowLevelImpl(input2, input1, needTies);
    } else {
        static assert(isArray!R1);
        static assert(isArray!R2);
        alias typeof(R1.init[0]) T;
        alias typeof(R2.init[0]) U;

        static if(T.sizeof > U.sizeof) {
            return kendallCorDestructiveLowLevelImpl(input1, input2, needTies);
        } else {
            return kendallCorDestructiveLowLevelImpl(input2, input1, needTies);
        }
    }
}

struct KendallLowLevel {
    double tau;
    long s;

    // Notation as in Kendall, 1947, Biometrika

    ulong tieCorrectT1;  // sum{t(t - 1)(2t + 5)}
    ulong tieCorrectT2;  // sum{t(t - 1)(t - 2)}
    ulong tieCorrectT3;  // sum{t(t - 1)}

    ulong tieCorrectU1;  // sum{u(u - 1)(2u + 5)}
    ulong tieCorrectU2;  // sum{u(u - 1)(u - 2)}
    ulong tieCorrectU3;  // sum{u(u - 1)}
}

package KendallLowLevel kendallCorDestructiveLowLevelImpl
(R1, R2)(R1 input1, R2 input2, bool needTies)
in {
    assert(input1.length == input2.length);
} body {
    static ulong getMs(V)(V data) {  //Assumes data is sorted.
        ulong Ms = 0, tieCount = 0;
        foreach(i; 1..data.length) {
            if(data[i] == data[i - 1]) {
                tieCount++;
            } else if(tieCount) {
                Ms += (tieCount * (tieCount + 1)) / 2;
                tieCount = 0;
            }
        }
        if(tieCount) {
            Ms += (tieCount * (tieCount + 1)) / 2;
        }
        return Ms;
    }

    void computeTies(V)
    (V arr, ref ulong tie1, ref ulong tie2, ref ulong tie3) {
        if(!needTies) {
            return;  // If only computing correlation, this is a waste of time.
        }

        ulong tieCount = 1;
        foreach(i; 1..arr.length) {
            if(arr[i] == arr[i - 1]) {
                tieCount++;
            } else if(tieCount > 1) {
                tie1 += tieCount * (tieCount - 1) * (2 * tieCount + 5);
                tie2 += tieCount * (tieCount - 1) * (tieCount - 2);
                tie3 += tieCount * (tieCount - 1);
                tieCount = 1;
            }
        }

        // Handle last run.
         if(tieCount > 1) {
            tie1 += tieCount * (tieCount - 1) * (2 * tieCount + 5);
            tie2 += tieCount * (tieCount - 1) * (tieCount - 2);
            tie3 += tieCount * (tieCount - 1);
        }
    }

    ulong m1 = 0;
    ulong nPair = (cast(ulong) input1.length *
                  ( cast(ulong) input1.length - 1UL)) / 2UL;
    KendallLowLevel ret;
    ret.s = to!long(nPair);

    static if(!isDefaultSorted!R1) {
        qsort!(compFun)(input1, input2);
    }

    uint tieCount = 0;
    foreach(i; 1..input1.length) {
        if(input1[i] == input1[i - 1]) {
            tieCount++;
        } else if(tieCount > 0) {
            static if(!isDefaultSorted!R2) {
                qsort!(compFun)(input2[i - tieCount - 1..i]);
            }
            m1 += tieCount * (tieCount + 1UL) / 2UL;
            ret.s += getMs(input2[i - tieCount - 1..i]);
            tieCount = 0;
        }
    }
    if(tieCount > 0) {
        static if(!isDefaultSorted!R2) {
            qsort!(compFun)(input2[input1.length - tieCount - 1..input1.length]);
        }
        m1 += tieCount * (tieCount + 1UL) / 2UL;
        ret.s += getMs(input2[input1.length - tieCount - 1..input1.length]);
    }

    computeTies(input1, ret.tieCorrectT1, ret.tieCorrectT2, ret.tieCorrectT3);

    static if(isArray!R2) {
        ulong swapCount = 0;
        static if(isArray!R1) {
            // We've already guaranteed that input1 is the bigger array by
            // bytes and we own these arrays and won't use input1 again, so
            // this is safe.
            alias typeof(input2[0]) U;
            U[] input1Temp = (cast(U*) input1.ptr)[0..input2.length];
            mergeSortTemp!(compFun)(input2, input1Temp, &swapCount);
        } else {
            mergeSort!(compFun)(input2, &swapCount);
        }
    } else {
        enum ulong swapCount = 0;  // If they're both sorted then tau == 1.
    }

    immutable m2 = getMs(input2);
    computeTies(input2, ret.tieCorrectU1, ret.tieCorrectU2, ret.tieCorrectU3);

    ret.s -= (m1 + m2) + 2 * swapCount;
    immutable double denominator1 = nPair - m1;
    immutable double denominator2 = nPair - m2;
    ret.tau = ret.s / sqrt(denominator1) / sqrt(denominator2);
    return ret;
}

/* Kendall's Tau correlation, O(N^2) version.  This is faster than the
 * more asymptotically efficient version for N <= about 15, and is also useful
 * for testing.  Yes, the sorts for the large N impl fall back on insertion
 * sorting for moderately small N, but due to additive constants and O(N) terms
 * this algorithm is still faster for very small N.  (Besides, I can't
 * delete it anyhow because I need it for testing.)
 */
private double kendallCorSmallN(R1, R2)(R1 input1, R2 input2)
in {
    assert(input1.length == input2.length);

    // This function should never be used for any inputs even close to this
    // large because it's a small-N optimization and a more efficient
    // implementation exists in this module for large N, but when N gets this
    // large it's not even correct due to overflow errors.
    assert(input1.length < 1 << 15);
} body {
    int m1 = 0, m2 = 0;
    int s = 0;

    foreach(i; 0..input2.length) {
        foreach (j; i + 1..input2.length) {
            if(input2[i] > input2[j]) {
                if (input1[i] > input1[j]) {
                    s++;
                } else if(input1[i] < input1[j]) {
                    s--;
                } else if(input1[i] == input1[j]) {
                    m1++;
                } else {
                    return double.nan;
                }
            } else if(input2[i] < input2[j]) {
                if (input1[i] > input1[j]) {
                    s--;
                } else if(input1[i] < input1[j]) {
                    s++;
                } else if(input1[i] == input1[j]) {
                    m1++;
                } else {
                    return double.nan;
                }
            } else if(input2[i] == input2[j]) {
                m2++;

                if(input1[i] < input1[j]) {
                } else if(input1[i] > input1[j]) {
                } else if(input1[i] == input1[j]) {
                    m1++;
                } else {
                    return double.nan;
                }

            } else {
                return double.nan;
            }
        }
    }

    immutable nCombination = input2.length * (input2.length - 1) / 2;
    immutable double denominator1 = nCombination - m1;
    immutable double denominator2 = nCombination - m2;
    return s / sqrt(denominator1) / sqrt(denominator2);
}


unittest {
    //Test against known values.
    assert(approxEqual(kendallCor([1,2,3,4,5].dup, [3,1,7,4,3].dup), 0.1054093));
    assert(approxEqual(kendallCor([3,6,7,35,75].dup,[1,63,53,67,3].dup), 0.2));
    assert(approxEqual(kendallCor([1.5,6.3,7.8,4.2,1.5].dup, [1,63,53,67,3].dup), .3162287));

    static void doKendallTest(T)() {
        T[] one = new T[1000], two = new T[1000];
        // Test complex, fast implementation against straightforward,
        // slow implementation.
        foreach(i; 0..100) {
            size_t lowerBound = uniform(0, 1000);
            size_t upperBound = uniform(0, 1000);
            if(lowerBound > upperBound) swap(lowerBound, upperBound);
            foreach(ref o; one) {
                o = uniform(cast(T) -10, cast(T) 10);
            }
            foreach(ref o; two) {
                 o = uniform(cast(T) -10, cast(T) 10);
            }
            double kOne =
                 kendallCor(one[lowerBound..upperBound], two[lowerBound..upperBound]);
            double kTwo =
                 kendallCorSmallN(one[lowerBound..upperBound], two[lowerBound..upperBound]);
            assert(approxEqual(kOne, kTwo) || (isNaN(kOne) && isNaN(kTwo)));
        }
    }

    doKendallTest!int();
    doKendallTest!float();
    doKendallTest!double();

    // Make sure everything works with lowest common denominator range type.
    static struct Count {
        uint num;
        uint upTo;
        @property size_t front() {
            return num;
        }
        void popFront() {
            num++;
        }
        @property bool empty() {
            return num >= upTo;
        }
    }

    Count a, b;
    a.upTo = 100;
    b.upTo = 100;
    assert(approxEqual(kendallCor(a, b), 1));

    // This test will fail if there are overflow bugs, especially in tie
    // handling.
    auto rng = chain(repeat(0, 100_000), repeat(1, 100_000));
    assert(approxEqual(kendallCor(rng, rng), 1));

    // Test the case where we have one range sorted already.
    assert(kendallCor(iota(5), [3, 1, 2, 5, 4]) ==
        kendallCor(assumeSorted(iota(5)), [3, 1, 2, 5, 4])
    );

    assert(kendallCor(iota(5), [3, 1, 2, 5, 4]) ==
        kendallCor([3, 1, 2, 5, 4], assumeSorted(iota(5)))
    );

    assert(approxEqual(
        kendallCor(assumeSorted(iota(5)), assumeSorted(iota(5))), 1
    ));

    auto lhs = randArray!rNormal(1_000, 0, 1);
    auto indices = new size_t[1_000];
    import std.algorithm;
    makeIndex(lhs, indices);

    foreach(i; 0..1_000) {
        auto rhs = randArray!rNormal(1_000, 0, 1);
        auto lhsSorted = assumeSorted(
            indexed(lhs, indices)
        );

        // Rearrange rhs according to the sorting permutation of lhs.
        // kendallCor(lhsSorted, rhsRearranged) will be much faster than
        // kendallCor(lhs, rhs).
        auto rhsRearranged = indexed(rhs, indices);
        assert(kendallCor(lhsSorted, rhsRearranged) == kendallCor(lhs, rhs));
    }
}

// Alias to old correlation function names, but don't document them.  These will
// eventually be deprecated.
alias PearsonCor Pcor;
alias pearsonCor pcor;
alias spearmanCor scor;
alias kendallCor kcor;
alias kendallCorDestructive kcorDestructive;

/**Computes the partial correlation between vec1, vec2 given
 * conditions.  conditions can be either a tuple of ranges, a range of ranges,
 * or (for a single condition) a single range.
 *
 * cor is the correlation metric to use.  It can be either pearsonCor,
 * spearmanCor, kendallCor, or any custom correlation metric you can come up
 * with.
 *
 * Examples:
 * ---
 * uint[] stock1Price = [8, 6, 7, 5, 3, 0, 9];
 * uint[] stock2Price = [3, 1, 4, 1, 5, 9, 2];
 * uint[] economicHealth = [2, 7, 1, 8, 2, 8, 1];
 * uint[] consumerFear = [1, 2, 3, 4, 5, 6, 7];
 *
 * // See whether the prices of stock 1 and stock 2 are correlated even
 * // after adjusting for the overall condition of the economy and consumer
 * // fear.
 * double partialCor =
 *   partial!pearsonCor(stock1Price, stock2Price, economicHealth, consumerFear);
 * ---
 */
double partial(alias cor, T, U, V...)(T vec1, U vec2, V conditionsIn)
if(isInputRange!T && isInputRange!U && allSatisfy!(isInputRange, V)) {
    auto alloc = newRegionAllocator();
    static if(V.length == 1 && isInputRange!(ElementType!(V[0]))) {
        // Range of ranges.
        static if(isArray!(V[0])) {
            alias conditionsIn[0] cond;
        } else {
            auto cond = tempdup(cond[0]);
        }
    } else {
        alias conditionsIn cond;
    }

    auto corMatrix = doubleMatrix(cond.length + 2, cond.length + 2, alloc);
    foreach(i; 0..corMatrix.rows) corMatrix[i, i] = 1;

    corMatrix[0, 1] = corMatrix[1, 0] = cast(double) cor(vec1, vec2);
    foreach(i, condition; cond) {
        immutable conditionIndex = i + 2;
        corMatrix[0, conditionIndex] = cast(double) cor(vec1, condition);
        corMatrix[conditionIndex, 0] =  corMatrix[0, conditionIndex];
        corMatrix[1, conditionIndex] = cast(double) cor(vec2, condition);
        corMatrix[conditionIndex, 1] = corMatrix[1, conditionIndex];
    }

    foreach(i, condition1; cond) {
        foreach(j, condition2; cond[i + 1..$]) {
            immutable index1 = i + 2;
            immutable index2 = index1 + j + 1;
            corMatrix[index1, index2] = cast(double) cor(condition1, condition2);
            corMatrix[index2, index1] = corMatrix[index1, index2];
        }
    }

    auto invMatrix = doubleMatrix(cond.length + 2, cond.length + 2, alloc);
    invert(corMatrix, invMatrix);

    // This code is written so verbosely to work around a compiler segfault
    // that I have no idea how to reduce.
    immutable denom = sqrt(invMatrix[0, 0] * invMatrix[1, 1]);
    immutable negNumer = invMatrix[0, 1];
    return -negNumer / denom;
}

unittest {
    // values from Matlab.
    uint[] stock1Price = [8, 6, 7, 5, 3, 0, 9];
    uint[] stock2Price = [3, 1, 4, 1, 5, 9, 2];
    uint[] economicHealth = [2, 7, 1, 8, 2, 8, 1];
    uint[] consumerFear = [1, 2, 3, 4, 5, 6, 7];
    double partialCor =
    partial!pearsonCor(stock1Price, stock2Price, [economicHealth, consumerFear][]);
    assert(approxEqual(partialCor, -0.857818));

    double spearmanPartial =
    partial!spearmanCor(stock1Price, stock2Price, economicHealth, consumerFear);
    assert(approxEqual(spearmanPartial, -0.7252));
}

private __gshared TaskPool emptyPool;
shared static this() {
    emptyPool = new TaskPool(0);
}

private struct RorToMatrix(RoR) {
    RoR* ror;

    auto ref opIndex(size_t i, size_t j) {
        return (*ror)[i][j];
    }

    void opIndexAssign(typeof((*ror)[0][0]) val, size_t i, size_t j) {
        (*ror)[i][j] = val;
    }

    size_t rows() @property {
        return (*ror).length;
    }
}

private auto makeMatrixLike(RoR)(ref RoR ror) {
    return RorToMatrix!RoR(&ror);
}

private template isMatrixLike(T) {
    enum isMatrixLike = is(typeof({
        T t;
        size_t r = t.rows;
        t[0, 0] = 2.0;
    }));
}

version(scid) {
import scid.matrix;

/**
These functions allow efficient calculation of the Pearson, Spearman and
Kendall correlation matrices and the covariance matrix respectively.  They
are optimized to avoid computing certain values multiple times when computing
correlations of one vector to several others.

Note:  These functions are only available when SciD is installed and
       dstats is compiled with version=scid.

Params:

mat = A range of ranges to be treated as a set of vectors.  This must be
      a finite range, and must be rectangular, i.e. all elements must be
      the same length.  For Pearson correlation and covariance, the ranges
      must also have elements implicitly convertible to double.
      SciD matrices can be treated as ranges of ranges and can be used with
      these functions.

taskPool = A std.parallelism.TaskPool used to parallelize the computation.
           This is useful for very large matrices.  If none is provided,
           the computation will not be parallelized.


Examples:
---
auto input = [[8.0, 6, 7, 5],
              [3.0, 0, 9, 3],
              [1.0, 4, 1, 5]];
auto pearson = pearsonMatrix(input);
assert(approxEqual(pearson[0, 0], 1));
---
*/
SymmetricMatrix!double pearsonMatrix(RoR)(RoR mat, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) &&
   is(ElementType!(ElementType!RoR) : double) &&
   !isInfinite!RoR
) {
    SymmetricMatrix!double ret;
    pearsonSpearmanCov!true(mat, pool, CorCovType.pearson, ret);
    return ret;
}

/// Ditto
SymmetricMatrix!double spearmanMatrix(RoR)(RoR mat, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) && !isInfinite!RoR) {
    SymmetricMatrix!double ret;
    pearsonSpearmanCov!true(mat, pool, CorCovType.spearman, ret);
    return ret;
}

/// Ditto
SymmetricMatrix!double kendallMatrix(RoR)(RoR mat, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) && !isInfinite!RoR) {
    SymmetricMatrix!double ret;
    kendallMatrixImpl!true(mat, ret, pool);
    return ret;
}

/// Ditto
SymmetricMatrix!double covarianceMatrix(RoR)(RoR mat, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) &&
   is(ElementType!(ElementType!RoR) : double) &&
   !isInfinite!RoR
) {
    SymmetricMatrix!double ret;
    pearsonSpearmanCov!true(mat, pool, CorCovType.covariance, ret);
    return ret;
}
}

/**
These overloads allow for correlation and covariance matrices to be computed
with the results being stored in a pre-allocated variable, ans.  ans must
be either a SciD matrix or a random-access range of ranges with assignable
elements of a floating point type.  It must have the same number of rows
as the number of vectors in mat and must have at least enough columns in
each row to support storing the lower triangle.  If ans is a full rectangular
matrix/range of ranges, only the lower triangle results will be stored.

Note:  These functions can be used without SciD because they don't return
       SciD types.

Examples:
---
auto pearsonRoR = [[0.0], [0.0, 0.0], [0.0, 0.0, 0.0]];
pearsonMatrix(input, pearsonRoR);
assert(approxEqual(pearsonRoR[1][1], 1));
---
*/
void pearsonMatrix(RoR, Ret)(RoR mat, ref Ret ans, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) &&
   is(ElementType!(ElementType!RoR) : double) &&
   !isInfinite!RoR &&
   (is(typeof(ans[0, 0] = 2.0)) || is(typeof(ans[0][0] = 2.0)))
) {
    static if(isMatrixLike!Ret) {
        alias ans ret;
    } else {
        auto ret = makeMatrixLike(ans);
    }

    pearsonSpearmanCov!false(mat, pool, CorCovType.pearson, ret);
}

/// Ditto
void spearmanMatrix(RoR, Ret)(RoR mat, ref Ret ans, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) && !isInfinite!RoR &&
    (is(typeof(ans[0, 0] = 2.0)) || is(typeof(ans[0][0] = 2.0)))
) {
    static if(isMatrixLike!Ret) {
        alias ans ret;
    } else {
        auto ret = makeMatrixLike(ans);
    }

    pearsonSpearmanCov!false(mat, pool, CorCovType.spearman, ret);
}

/// Ditto
void kendallMatrix(RoR, Ret)(RoR mat, ref Ret ans, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) && !isInfinite!RoR &&
   (is(typeof(ans[0, 0] = 2.0)) || is(typeof(ans[0][0] = 2.0)))
) {
    static if(isMatrixLike!Ret) {
        alias ans ret;
    } else {
        auto ret = makeMatrixLike(ans);
    }

    kendallMatrixImpl!false(mat, ret, pool);
}

/// Ditto
void covarianceMatrix(RoR, Ret)(RoR mat, ref Ret ans, TaskPool pool = null)
if(isInputRange!RoR && isInputRange!(ElementType!RoR) &&
   is(ElementType!(ElementType!RoR) : double) &&
   !isInfinite!RoR &&
   (is(typeof(ans[0, 0] = 2.0)) || is(typeof(ans[0][0] = 2.0)))
) {
    static if(isMatrixLike!Ret) {
        alias ans ret;
    } else {
        auto ret = makeMatrixLike(ans);
    }

    pearsonSpearmanCov!false(mat, pool, CorCovType.covariance, ret);
}

private void kendallMatrixImpl(bool makeNewMatrix, RoR, Matrix)
(RoR mat, ref Matrix ret, TaskPool pool = null) {
    if(pool is null) {
        pool = emptyPool;
    }

    auto alloc = newRegionAllocator();
    alias ElementType!RoR R;
    alias ElementType!R E;

    static if(!isRandomAccessRange!R || !isForwardRange!RoR) {
        auto randomMat = alloc.array(mat
                .map!(r => prepareForSorting(alloc.array(r))));
    } else {
        alias mat randomMat;
    }

    if(randomMat.empty) return;
    immutable nElems = randomMat.front.length;

    foreach(row; randomMat.save) {
        dstatsEnforce(row.length == nElems,
            "Range of ranges must be rectangular for kendallMatrix."
        );
    }

    static if(makeNewMatrix) {
        static assert(is(typeof(ret) == SymmetricMatrix!double));
        ret = SymmetricMatrix!double(walkLength(randomMat));
    }

    // HACK:  Before the multithreaded portion of this algorithm
    // starts, make sure that there's no need to unshare ret if it's
    // using ref-counted COW semantics.
    ret[0, 0] = 0;

    foreach(i, row1; pool.parallel(randomMat.save, 1)) {
        auto alloc2 = newRegionAllocator();
        auto indices = alloc2.uninitializedArray!(size_t[])(row1.length);
        foreach(ii, ref elem; indices) elem = ii;

        bool less(size_t a, size_t b) {
            return row1[a] < row1[b];
        }

        qsort!less(indices);
        auto row1Sorted = assumeSorted(indexed(row1, indices));

        size_t j = 0;
        foreach(row2; randomMat.save) {
            scope(exit) j++;
            if(i == j) {
                ret[i, i] = 1;
                break;
            }

            auto row2Rearranged = indexed(row2, indices);
            ret[i, j] = kendallCor(row1Sorted, row2Rearranged);
        }
    }
}

private enum CorCovType {
    pearson,
    spearman,
    covariance
}

private void pearsonSpearmanCov(bool makeNewMatrix, RoR, Matrix)
(RoR mat, TaskPool pool, CorCovType type, ref Matrix ret) {
    import dstats.summary : mean;
    if(pool is null) {
        pool = emptyPool;
    }

    auto alloc = newRegionAllocator();

    auto normalized = alloc.array(mat
            .map!(r => alloc.array(r.map!(to!double))));

    foreach(row; normalized[1..$]) {
        dstatsEnforce(row.length == normalized[0].length,
            "Matrix must be rectangular for pearsonMatrix.");
    }

    immutable double nCols = (normalized.length) ? (normalized[0].length) : 0;

    final switch(type) {
        case CorCovType.pearson:
            foreach(row; pool.parallel(normalized)) {
                immutable msd = meanStdev(row);
                row[] = (row[] - msd.mean) / sqrt(msd.mse * nCols);
            }

            break;
        case CorCovType.covariance:
            immutable divisor = sqrt(nCols - 1.0);
            foreach(row; pool.parallel(normalized)) {
                immutable mu = mean(row).mean;
                row[] -= mu;
                row[] /= divisor;
            }
            break;
        case CorCovType.spearman:
            foreach(row; pool.parallel(normalized)) {
                auto alloc = newRegionAllocator();
                auto buf = alloc.uninitializedArray!(double[])(row.length);
                rank(row, buf);

                // Need to find mean, stdev separately for every row b/c
                // of ties.
                immutable msd = meanStdev(buf);
                row[] = (buf[] - msd.mean) / sqrt(msd.mse * nCols);
            }

            break;
    }

    static if(makeNewMatrix) {
        static assert(is(typeof(ret) == SymmetricMatrix!double));
        ret = SymmetricMatrix!double(normalized.length);
    }

    dotMatrix(normalized, ret, pool);
}

// This uses an array-of-arrays to avoid heap fragmentation issues with large
// matrices and because the loss of efficiency from poitner chasing is
// negligible given that we always access them one row at a time.
private void dotMatrix(Matrix)(
    double[][] rows,
    ref Matrix ret,
    TaskPool pool
) in {
    foreach(row; rows[1..$]) {
        assert(row.length == rows[0].length);
    }

    assert(ret.rows == rows.length);
} body {
    // HACK:  Before the multithreaded portion of this algorithm
    // starts, make sure that there's no need to unshare ret if it's
    // using ref-counted COW semantics.
    ret[0, 0] = 0;

    foreach(i; pool.parallel(iota(0, rows.length), 1)) {
        auto row1 = rows[i];

        foreach(j; 0..i + 1) {
            ret[i, j] = dotProduct(row1, rows[j]);
        }
    }
}

unittest {
    auto input = [[8.0, 6, 7, 5],
                  [3.0, 0, 9, 3],
                  [1.0, 4, 1, 5]];


    static double[][] makeRoR() {
        return [[0.0], [0.0, 0.0], [0.0, 0.0, 0.0]];
    }

    auto pearsonRoR = makeRoR();
    pearsonMatrix(input, pearsonRoR);

    auto spearmanRoR = makeRoR();
    spearmanMatrix(input, spearmanRoR);

    auto kendallRoR = makeRoR();
    kendallMatrix(input, kendallRoR);

    auto covRoR = makeRoR();
    covarianceMatrix(input, covRoR);

    // Values from R.

    alias approxEqual ae; // Save typing.
    assert(ae(pearsonRoR[0][0], 1));
    assert(ae(pearsonRoR[1][1], 1));
    assert(ae(pearsonRoR[2][2], 1));
    assert(ae(pearsonRoR[1][0], 0.3077935));
    assert(ae(pearsonRoR[2][0], -0.9393364));
    assert(ae(pearsonRoR[2][1], -0.6103679));

    assert(ae(spearmanRoR[0][0], 1));
    assert(ae(spearmanRoR[1][1], 1));
    assert(ae(spearmanRoR[2][2], 1));
    assert(ae(spearmanRoR[1][0], 0.3162278));
    assert(ae(spearmanRoR[2][0], -0.9486833));
    assert(ae(spearmanRoR[2][1], -0.5));

    assert(ae(kendallRoR[0][0], 1));
    assert(ae(kendallRoR[1][1], 1));
    assert(ae(kendallRoR[2][2], 1));
    assert(ae(kendallRoR[1][0], 0.1825742));
    assert(ae(kendallRoR[2][0], -0.9128709));
    assert(ae(kendallRoR[2][1], -0.4));

    assert(ae(covRoR[0][0], 1.66666));
    assert(ae(covRoR[1][1], 14.25));
    assert(ae(covRoR[2][2], 4.25));
    assert(ae(covRoR[1][0], 1.5));
    assert(ae(covRoR[2][0], -2.5));
    assert(ae(covRoR[2][1], -4.75));

    version(scid) {
    static bool test(double[][] a, SymmetricMatrix!double b) {
        foreach(i; 0..3) foreach(j; 0..i + 1) {
            if(!ae(a[i][j], b[i, j])) return false;
        }

        return true;
    }

    auto pearson = pearsonMatrix(input, taskPool);
    auto spearman = spearmanMatrix(input, taskPool);
    auto kendall = kendallMatrix(input, taskPool);
    auto cov = covarianceMatrix(input, taskPool);

    test(pearsonRoR, pearson);
    test(spearmanRoR, spearman);
    test(kendallRoR, kendall);
    test(covRoR, cov);
    }
}