local Rx = require 'Rx'

local minRunsRequiredByStatFunctions = 5

--------------------------------------------------------------------------------

local benchmarkPrototype = {

    _testName = false,
        -- (string) name of the test

    _minRuns = minRunsRequiredByStatFunctions,
        -- (number) minimum number of times to run a test, must be >= 5

    _maxRuns = 1000,
        -- (number) maximum number of times to run a test

    _numGCCyclesToRun = 3,
        -- (number) number of times to run garbagecollect before starting a timing

    _overheadSamples = {},
        -- (table) filled in with timings of test runs, reused for each testSize

    _oldSamples = {},
        -- (table) filled in with timings of test runs, reused for each testSize

    _newSamples = {},
        -- (table) filled in with timings of test runs, reused for each testSize

    _overheadFn = false,
        -- (function) a function that contains overhead time to subtract from tests

    _oldFn = false,
        -- (function) the function being compared against

    _newFn = false,
        -- (function) the new version of the function to test

    _rangeStart = 1,
        -- (number)

    _rangeEnd = 1,
        -- (number)

    _rangeMultiplier = 2,
        -- (number) when using a range this value is multiplied for each iteration

    _results = false,
        -- (table) filled with the results after the benchmark has been run

    _showResults = true,
        -- (boolean) show results of tests

    _showDetails = false,
        -- (boolean) shows detailed stats

    _printFn = false,
        -- (function) to use when printing results

    class = {},
        -- class methods
}

--------------------------------------------------------------------------------

function benchmarkPrototype:__init(testName, minRuns, maxRuns, numGCCyclesToRun)

    self._testName = testName or "Unnamed"

    if minRuns and minRuns >= minRunsRequiredByStatFunctions then
        self._minRuns = minRuns
    end

    self._maxRuns = maxRuns or self._maxRuns
    if self._maxRuns < self._minRuns then
        self._maxRuns = self._minRuns
    end

    self._numGCCyclesToRun = numGCCyclesToRun or self._numGCCyclesToRun
    self._printFn = print

end

--------------------------------------------------------------------------------

local stat = {

    min = function(samples)

        local numSamples = #samples

        if numSamples > 0 then

            local min = samples[1]
            for _, v in ipairs(samples) do
                if v < min then
                    min = v
                end
            end

            return min

        else
            return 0
        end

    end,

    ----------------------------------------------------------------------------

    max = function(samples)

        local numSamples = #samples

        if numSamples > 0 then

            local max = samples[1]
            for _, v in ipairs(samples) do
                if v > max then
                    max = v
                end
            end

            return max

        else
            return 0
        end

    end,

    ----------------------------------------------------------------------------

    mean = function(samples)

        local numSamples = #samples

        if numSamples > 0 then

            local sum = 0
            for _, v in ipairs(samples) do
                sum = sum + v
            end

            return sum / numSamples

        else
            return 0
        end

    end,

    ----------------------------------------------------------------------------

    variance = function(samples, mean)

        local numSamples = #samples

        if numSamples > 1 then

            local sum = 0
            for _, v in ipairs(samples) do
                sum = sum + (v - mean)^2
            end

            return sum / (numSamples-1)

        else
            return 0
        end

    end

}

--------------------------------------------------------------------------------

-- T-Distribution two-tailed critical values for 95% confidence

local tTable = {
    12.706, 4.303, 3.182, 2.776, 2.571, 2.447, 2.365, 2.306, 2.262, 2.228,
    02.201, 2.179, 2.160, 2.145, 2.131, 2.120, 2.110, 2.101, 2.093, 2.086,
    02.080, 2.074, 2.069, 2.064, 2.060, 2.056, 2.052, 2.048, 2.045, 2.042,
    infinity = 1.960
}

--------------------------------------------------------------------------------

-- Critical Mann-Whitney U-values for 95% confidence

local uTable = {
    [  5 ] = { 0, 1, 2 },
    [  6 ] = { 1, 2, 3, 5 },
    [  7 ] = { 1, 3, 5, 6, 8 },
    [  8 ] = { 2, 4, 6, 8, 10, 13 },
    [  9 ] = { 2, 4, 7, 10, 12, 15, 17 },
    [ 10 ] = { 3, 5, 8, 11, 14, 17, 20, 23 },
    [ 11 ] = { 3, 6, 9, 13, 16, 19, 23, 26, 30 },
    [ 12 ] = { 4, 7, 11, 14, 18, 22, 26, 29, 33, 37 },
    [ 13 ] = { 4, 8, 12, 16, 20, 24, 28, 33, 37, 41, 45 },
    [ 14 ] = { 5, 9, 13, 17, 22, 26, 31, 36, 40, 45, 50, 55 },
    [ 15 ] = { 5, 10, 14, 19, 24, 29, 34, 39, 44, 49, 54, 59, 64 },
    [ 16 ] = { 6, 11, 15, 21, 26, 31, 37, 42, 47, 53, 59, 64, 70, 75 },
    [ 17 ] = { 6, 11, 17, 22, 28, 34, 39, 45, 51, 57, 63, 67, 75, 81, 87 },
    [ 18 ] = { 7, 12, 18, 24, 30, 36, 42, 48, 55, 61, 67, 74, 80, 86, 93, 99 },
    [ 19 ] = { 7, 13, 19, 25, 32, 38, 45, 52, 58, 65, 72, 78, 85, 92, 99, 106, 113 },
    [ 20 ] = { 8, 14, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90, 98, 105, 112, 119, 127 },
    [ 21 ] = { 8, 15, 22, 29, 36, 43, 50, 58, 65, 73, 80, 88, 96, 103, 111, 119, 126, 134, 142 },
    [ 22 ] = { 9, 16, 23, 30, 38, 45, 53, 61, 69, 77, 85, 93, 101, 109, 117, 125, 133, 141, 150, 158 },
    [ 23 ] = { 9, 17, 24, 32, 40, 48, 56, 64, 73, 81, 89, 98, 106, 115, 123, 132, 140, 149, 157, 166, 175 },
    [ 24 ] = { 10, 17, 25, 33, 42, 50, 59, 67, 76, 85, 94, 102, 111, 120, 129, 138, 147, 156, 165, 174, 183, 192 },
    [ 25 ] = { 10, 18, 27, 35, 44, 53, 62, 71, 80, 89, 98, 107, 117, 126, 135, 145, 154, 163, 173, 182, 192, 201, 211 },
    [ 26 ] = { 11, 19, 28, 37, 46, 55, 64, 74, 83, 93, 102, 112, 122, 132, 141, 151, 161, 171, 181, 191, 200, 210, 220, 230 },
    [ 27 ] = { 11, 20, 29, 38, 48, 57, 67, 77, 87, 97, 107, 118, 125, 138, 147, 158, 168, 178, 188, 199, 209, 219, 230, 240, 250 },
    [ 28 ] = { 12, 21, 30, 40, 50, 60, 70, 80, 90, 101, 111, 122, 132, 143, 154, 164, 175, 186, 196, 207, 218, 228, 239, 250, 261, 272 },
    [ 29 ] = { 13, 22, 32, 42, 52, 62, 73, 83, 94, 105, 116, 127, 138, 149, 160, 171, 182, 193, 204, 215, 226, 238, 249, 260, 271, 282, 294 },
    [ 30 ] = { 13, 23, 33, 43, 54, 65, 76, 87, 98, 109, 120, 131, 143, 154, 166, 177, 189, 200, 212, 223, 235, 247, 258, 270, 282, 293, 305, 317 }
}

--------------------------------------------------------------------------------

local function calcStats(sample)

    local mean = stat.mean(sample)
    local variance = stat.variance(sample, mean)
    local standardDeviation = math.sqrt(variance)
    local standardError = standardDeviation / math.sqrt(#sample)

    local degreesOfFreedom = #sample - 1
    if degreesOfFreedom < 1 then
        degreesOfFreedom = 1
    end

    local criticalValue = tTable[degreesOfFreedom] or tTable['infinity']

    local marginOfError = standardError * criticalValue

    local relativeMarginOfError = 0
    if mean > 0 then
        relativeMarginOfError = (marginOfError / mean) * 100
    end

    return {
        numSamples = #sample,
        min = stat.min(sample),
        max = stat.max(sample),
        mean = mean,
        variance = variance,
        standardDeviation = standardDeviation,
        standardError = standardError,
        degreesOfFreedom = degreesOfFreedom,
        criticalValue = criticalValue,
        marginOfError = marginOfError,
        relativeMarginOfError = relativeMarginOfError
    }

end

--------------------------------------------------------------------------------

local function calcRankSum(sampleA, sampleB)

    local function getScore(xA, sampleB)

        local score = 0

        for i = 1, #sampleB do

            local xB = sampleB[i]

            if xB > xA then
                score = score + 1
            elseif xB == xA then
                score = score + 0.5
            end

        end

        return score

    end

    ----------------------------------------------------------------------------

    local u = 0

    for i = 1, #sampleA do

        local xA = sampleA[i]

        u = u + getScore(xA, sampleB)

    end

    return u

end

--------------------------------------------------------------------------------

local function calcZStat(u, n1, n2)

    return (u - ((n1 * n2) / 2)) / math.sqrt((n1 * n2 * (n1 + n2 + 1)) / 12)

end

--------------------------------------------------------------------------------

local function compareSamples(sample1, sample2)

-- Returns:
--   -1 if sample1 slower than sample2.
--    0 if indeterminate.
--    1 if sample1 faster than sample2.

    local size1 = #sample1
    local size2 = #sample2
    local u1 = calcRankSum(sample1, sample2)
    local u2 = calcRankSum(sample2, sample1)
    local u = math.min(u1, u2)

    -- reject the null hyphothesis the two samples come from the
    -- same population (i.e. have the same median) if...

    local nullHypothesisRejected = false

    if size1 + size2 > 30 then

        -- ...the z-stat is less than -1.96
        -- http://www.statisticslectures.com/topics/mannwhitneyu/

        local zStat = calcZStat(u, size1, size2)

        nullHypothesisRejected = zStat <= -1.96

    else

        -- ...the U value is less than or equal the critical U value
        -- http://www.geoib.com/mann-whitney-u-test.html

        local maxSize = math.max(size1, size2)
        local minSize = math.min(size1, size2)

        local critical = 0

        if maxSize >= minRunsRequiredByStatFunctions and minSize >= 3 then

            critical = uTable[maxSize][minSize - 2]

            nullHypothesisRejected = u <= critical

        end

    end

    local result = 0

    if nullHypothesisRejected then

        if u1 > u2 then
            result = 1
        else
            result = -1
        end

    end

    return result, u1, u2

end

--------------------------------------------------------------------------------

function benchmarkPrototype:printStats(title, stats, showDetails)

    if stats.numSamples > 0 then

        if showDetails then

            self._printFn("Stats for " .. title .. ":")

            for k, v in pairs(stats) do
                self._printFn(" ", k, v)
            end

        end

        self._printFn(
            string.format("%s x %.2f ops/sec +/- %.2f%% (%d runs sampled)",
                title,
                1 / stats.mean,
                stats.relativeMarginOfError,
                stats.numSamples))

    end

end

--------------------------------------------------------------------------------

function benchmarkPrototype:_timeFunc(fn, testSize)

    local samples = {}

    local marginOfError = 100

    local i = 0

    while (i < self._minRuns or marginOfError > 1.0) and (i < self._maxRuns) do

        for gcX = 1, self._numGCCyclesToRun do
            collectgarbage()
        end

        local startTime = os.clock()

        fn(testSize)

        samples[#samples + 1] = os.clock() - startTime

        i = i + 1

        if i > self._minRuns then

            local stats = calcStats(samples)
            marginOfError = stats.relativeMarginOfError

        end

    end

    return samples

end

--------------------------------------------------------------------------------

function benchmarkPrototype:addOverhead(overheadFn)
    self._overheadFn = overheadFn
end

--------------------------------------------------------------------------------

function benchmarkPrototype:addOld(oldFn)
    self._oldFn = oldFn
end

--------------------------------------------------------------------------------

function benchmarkPrototype:addNew(newFn)
    self._newFn = newFn
end

--------------------------------------------------------------------------------

function benchmarkPrototype:testRange(rangeStart, rangeEnd, rangeMultiplier)

    self._rangeStart = rangeStart
    self._rangeEnd = rangeEnd

    self._rangeMultiplier = rangeMultiplier
    if not self._rangeMultiplier or self._rangeMultiplier <= 1 then
        self._rangeMultiplier = 2
    end

end

--------------------------------------------------------------------------------

local function timesFaster(new, old, overhead)

    local newDuration = new - overhead
    local oldDuration = old - overhead

    return oldDuration / newDuration

end

--------------------------------------------------------------------------------

local function printResults(self, testSize, results, showDetails)

    local testName = self._testName
    if testSize and self._rangeEnd > 1 then
        testName = testName .. "(" .. tostring(testSize) .. ")"
    end

    self._printFn()

    local comparisonToName = {
        [-1] = "old",
        [ 0] = "undetermined",
        [ 1] = "new"
    }

    self._printFn(
        string.format("%s - fastest is %s.",
            testName,
            comparisonToName[results.compareResult]))

    local overheadStats = results.overheadStats
    local oldStats = results.oldStats
    local newStats = results.newStats

    self:printStats(testName .. " - overhead", overheadStats, showDetails)
    self:printStats(testName .. " - old", oldStats, showDetails)
    self:printStats(testName .. " - new", newStats, showDetails)

    local newDuration = newStats.min - overheadStats.min
    local oldDuration = oldStats.min - overheadStats.min

    self._printFn(
        string.format("%s - Fastest time - new vs old: %f vs %f  (%.2f times faster)",
            testName,
            newDuration,
            oldDuration,
            oldDuration / newDuration))

    newDuration = newStats.max - overheadStats.max
    oldDuration = oldStats.max - overheadStats.max

    self._printFn(
        string.format("%s - Slowest time - new vs old: %f vs %f  (%.2f times faster)",
            testName,
            newDuration,
            oldDuration,
            oldDuration / newDuration))

    newDuration = newStats.mean - overheadStats.mean
    oldDuration = oldStats.mean - overheadStats.mean

    self._printFn(
        string.format("%s - Mean time - new vs old: %f vs %f  (%.2f times faster)",
            testName,
            newDuration,
            oldDuration,
            oldDuration / newDuration))

end

--------------------------------------------------------------------------------

function benchmarkPrototype:run()

    if not (self._oldFn and self._newFn) then
        error("must call addOld and addNew before running benchmarks")
    end

    if self._results then
        return self._results
    end

    local results = {}

    local testSize = self._rangeStart

    while testSize <= self._rangeEnd do

        if self._overheadFn then
            self._overheadSamples = self:_timeFunc(self._overheadFn, testSize)
        end

        self._oldSamples = self:_timeFunc(self._oldFn, testSize)

        self._newSamples = self:_timeFunc(self._newFn, testSize)

        local compareResult = compareSamples(self._newSamples, self._oldSamples)

        local overheadStats = calcStats(self._overheadSamples)
        local oldStats = calcStats(self._oldSamples)
        local newStats = calcStats(self._newSamples)

        results[testSize] = {
            compareResult = compareResult,
            overheadStats = overheadStats,
            oldStats = oldStats,
            newStats = newStats,
        }

        if self._showResults then
            printResults(self, testSize, results[testSize], self._showDetails)
        end

        testSize = testSize * self._rangeMultiplier

    end

    self._results = results

    return self._results

end

--------------------------------------------------------------------------------

function benchmarkPrototype:timesFaster(testSize)

    local results = self:run()

    if testSize == nil then
        testSize = 1
    end

    local v = results[testSize]

    if v then
        return timesFaster(v.newStats.mean, v.oldStats.mean, v.overheadStats.mean)
    else
        return 0
    end

end

--------------------------------------------------------------------------------

function benchmarkPrototype:newIsFaster(minImprovement)

    -- Return true if at least one test from the test range is faster
    -- by the specified percentage and none are slower.

    if minImprovement == nil then
        minImprovement = 1.0
    end

    local results = self:run()

    local someAreFaster = false
    local someAreSlower = false

    for k, v in pairs(results) do

        if v.compareResult == 1 then

            local minXFaster = timesFaster(v.newStats.min, v.oldStats.min, v.overheadStats.min)
            local meanXFaster = timesFaster(v.newStats.mean, v.oldStats.mean, v.overheadStats.mean)

            if minXFaster >= minImprovement and meanXFaster >= minImprovement then
                someAreFaster = true
            end

        end

        if v.compareResult == -1 then
            someAreSlower = true
        end

    end

    return someAreFaster and not someAreSlower

end

--------------------------------------------------------------------------------

function benchmarkPrototype:newIsSlower()

    -- Return true if at least one test from the test range is slower.

    local results = self:run()

    local someAreSlower = false

    for k, v in pairs(results) do

        if v.compareResult == -1 then
            someAreSlower = true
        end

    end

    return someAreSlower

end

--------------------------------------------------------------------------------

function benchmarkPrototype:setPrintMode(showResults, showDetails, printFn)

    self._showResults = showResults

    if showDetails then
        self._showDetails = showDetails
    end

    if printFn then
        self._printFn = printFn
    end

end

--------------------------------------------------------------------------------

function benchmarkPrototype.class.compareSamples(_, sample1, sample2)
    return compareSamples(sample1, sample2)
end

--------------------------------------------------------------------------------

function benchmarkPrototype.class.calcStats(_, sample)
    return calcStats(sample)
end

--------------------------------------------------------------------------------

return Rx.Class('Rx.internals.Benchmark', benchmarkPrototype)
