LogicLibrary = {}

local tonumber = tonumber

local NN = function(sequence, default)
    return sequence or Sequence:newWithScalar(default)
end

local map = function(func)
    return function(sequence)
        return NN(sequence):map(func)
    end
end

local numberMap = function(func)
    return map(function(value)
        value = tonumber(value)
        if value then
            return func(value)
        end
    end)
end

local stringMap = function(func)
    return map(function(value)
        value = todisplaystring(value)
        if value then
            return func(value)
        end
    end)
end

local dateMap = function(func)
    return map(function(value)
        if value and type(value) == 'userdata' and value.year then
            return func(value)
        end
    end)
end

local multiMap = function(nargs, func)
    return function(s1, ...)
        local sequences = {...}
        for index = 1, nargs - 1 do
            sequences[index] = NN(sequences[index])
        end
        return NN(s1):multiMap(func, unpack(sequences))
    end
end

local map2 = function(func)
    return function(s1, s2)
        return NN(s1):map2(func, NN(s2))
    end
end

local numberMap2 = function(func)
    return map2(function(value1, value2)
        value1, value2 = tonumber(value1), tonumber(value2)
        if value1 and value2 then
            return func(value1, value2)
        end
    end)
end

local formatMap = function(func)
    return function(sequence, ...)
        local formatter = func()
        return NN(sequence):multiMap(function(value, ...)
            return formatter:format(value, ...)
        end, ...)
    end
end

local function _meanAndVariance(sequence)
    local s0, s1, s2 = 0, 0, 0
    for _, value in sequence:iter() do
        value = tonumber(value)
        if value then
            s0 = s0 + 1
            s1 = s1 + value
            s2 = s2 + value ^ 2
        end
    end
    return s1 / s0, (s0 * s2 - s1 ^ 2) / (s0 * math.max(1, s0 - 1))
end

local function _meanAndVariance2(sequence1, sequence2)
    local s0, s11, s12, s21, s22 = 0, 0, 0, 0, 0
    for _, value1, value2 in sequence1:pairIter(sequence2) do
        value1, value2 = tonumber(value1), tonumber(value2)
        if value1 and value2 then
            s0 = s0 + 1
            s11 = s11 + value1
            s12 = s12 + value1 ^ 2
            s21 = s21 + value2
            s22 = s22 + value2 ^ 2
        end
    end
    local mean1 = s11 / s0
    local mean2 = s21 / s0
    local var1 = (s0 * s12 - s11 ^ 2) / (s0 * math.max(1, s0 - 1))
    local var2 = (s0 * s22 - s21 ^ 2) / (s0 * math.max(1, s0 - 1))
    return mean1, mean2, var1, var2
end

local function _covariance(sequence1, sequence2, mean1, mean2)
    local count, sum = 0, 0
    for _, value1, value2 in sequence1:pairIter(sequence2) do
        value1, value2 = tonumber(value1), tonumber(value2)
        if value1 and value2 then
            count = count + 1
            sum = sum + (value1 - mean1) * (value2 - mean2)
        end
    end
    return sum / (count - 1)
end

local function _mapReduce(identity, func, ...)
    local sequences = {...}
    local scalar = true
    local length = 0
    for index = 1, #sequences do
        local sequence = sequences[index]
        if not sequence:isScalar() then
            scalar = false
            length = math.max(length, sequence:length())
        end
    end
    if scalar then
        length = math.huge
    end
    return Sequence:new{
        getter = function(index)
            local result = identity
            for sindex = 1, #sequences do
                local sequence = sequences[sindex]
                local value = sequence:getValue(index)
                result = func(result, value)
            end
            return result
        end,
        length = function()
            return length
        end,
    }
end

LogicLibrary.average = function(sequence, rangeMin, rangeMax)
    if not rangeMin or not rangeMax then
        rangeMin = Sequence:newWithScalar(1)
        rangeMax = Sequence:newWithScalar(math.huge)
    end
    local function average(rangeMin, rangeMax)
        rangeMin, rangeMax = tonumber(rangeMin), tonumber(rangeMax)
        if rangeMin and rangeMax then
            local count = 0
            local sum = 0
            local function stepFunc(value)
                value = tonumber(value)
                if value then
                    count = count + 1
                    sum = sum + value
                end
            end
            Sequence.reduceRange(sequence, stepFunc, rangeMin, rangeMax)
            return sum / count
        end
    end
    if rangeMin:isScalar() and rangeMax:isScalar() then
        return Sequence:newWithScalar(average(rangeMin:getValue(1), rangeMax:getValue(1)))
    else
        return Sequence:new{
            getter = function(index)
                return average(rangeMin:getValue(index), rangeMax:getValue(index))
            end,
            length = function()
                return math.min(rangeMin:length(), rangeMax:length())
            end,
        }
    end
end

LogicLibrary.conjunction = function(...)
    -- if any false then false, elseif any truthy then true, else nil
    return _mapReduce(nil, function(conjunction, value)
        if value ~= nil then
            conjunction = value and (conjunction ~= false)
        end
        return conjunction
    end, ...)
end

LogicLibrary.count = function(sequence, rangeMin, rangeMax)
    if not rangeMin or not rangeMax then
        rangeMin = Sequence:newWithScalar(1)
        rangeMax = Sequence:newWithScalar(math.huge)
    end
    local function count(rangeMin, rangeMax)
        rangeMin, rangeMax = tonumber(rangeMin), tonumber(rangeMax)
        if rangeMin and rangeMax then
            local count = 0
            local function stepFunc(value)
                if value then
                    count = count + 1
                end
            end
            Sequence.reduceRange(sequence, stepFunc, rangeMin, rangeMax)
            return count
        end
    end
    if rangeMin:isScalar() and rangeMax:isScalar() then
        return Sequence:newWithScalar(count(rangeMin:getValue(1), rangeMax:getValue(1)))
    else
        return Sequence:new{
            getter = function(index)
                return count(rangeMin:getValue(index), rangeMax:getValue(index))
            end,
            length = function()
                return math.min(rangeMin:length(), rangeMax:length())
            end,
        }
    end
end

LogicLibrary.correl = function(sequence1, sequence2)
    local result
    if sequence1 and sequence2 then
        local mean1, mean2, var1, var2 = _meanAndVariance2(sequence1, sequence2)
        local cov = _covariance(sequence1, sequence2, mean1, mean2)
        result = cov / math.sqrt(var1 * var2)
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.cov = function(sequence1, sequence2)
    local result
    if sequence1 and sequence2 then
        local mean1, mean2 = _meanAndVariance2(sequence1, sequence2)
        local cov = _covariance(sequence1, sequence2, mean1, mean2)
        result = cov
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.disjunction = function(...)
    -- if any truthy then true, elseif any false then false, else nil
    return _mapReduce(nil, function(disjunction, value)
        if value ~= nil then
            disjunction = disjunction or (value ~= false)
        end
        return disjunction
    end, ...)
end

LogicLibrary.geomean = function(sequence, rangeMin, rangeMax)
    if not rangeMin or not rangeMax then
        rangeMin = Sequence:newWithScalar(1)
        rangeMax = Sequence:newWithScalar(math.huge)
    end
    local function geomean(rangeMin, rangeMax)
        rangeMin, rangeMax = tonumber(rangeMin), tonumber(rangeMax)
        if rangeMin and rangeMax then
            local count = 0
            local logSum = 0
            local function stepFunc(value)
                value = tonumber(value)
                if value then
                    count = count + 1
                    logSum = logSum + math.log(value)
                end
            end
            Sequence.reduceRange(sequence, stepFunc, rangeMin, rangeMax)
            return math.exp(logSum / count)
        end
    end
    if rangeMin:isScalar() and rangeMax:isScalar() then
        return Sequence:newWithScalar(geomean(rangeMin:getValue(1), rangeMax:getValue(1)))
    else
        return Sequence:new{
            getter = function(index)
                return geomean(rangeMin:getValue(index), rangeMax:getValue(index))
            end,
            length = function()
                return math.min(rangeMin:length(), rangeMax:length())
            end,
        }
    end
end

LogicLibrary.max = function(sequence)
    local result
    if sequence then
        result = sequence:reduce(-math.huge, function(max, value)
            value = tonumber(value)
            if value and value > max then
                max = value
            end
            return max
        end)
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.min = function(sequence)
    local result
    if sequence then
        result = sequence:reduce(math.huge, function(min, value)
            value = tonumber(value)
            if value and value < min then
                min = value
            end
            return min
        end)
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.median = function(sequence)
    local result
    if sequence then
        local values = {}
        for _, value in sequence:iter() do
            value = tonumber(value)
            if value then
                values[#values + 1] = value
            end
        end
        if #values > 0 then
            table.sort(values)
            local index = (#values + 1) / 2
            result = (values[math.floor(index)] + values[math.ceil(index)]) / 2
        end
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.mode = function(sequence)
    local result
    if sequence then
        local counts = {}
        for _, value in sequence:iter() do
            if value ~= nil then
                counts[value] = (counts[value] or 0) + 1
            end
        end
        local maxValue
        local maxCount = 0
        for value, count in pairs(counts) do
            if count > maxCount then
                maxCount = count
                maxValue = value
            end
        end
        result = maxValue
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.product = function(sequence, rangeMin, rangeMax)
    if not rangeMin or not rangeMax then
        rangeMin = Sequence:newWithScalar(1)
        rangeMax = Sequence:newWithScalar(math.huge)
    end
    local function product(rangeMin, rangeMax)
        rangeMin, rangeMax = tonumber(rangeMin), tonumber(rangeMax)
        if rangeMin and rangeMax then
            local product = 1
            local function stepFunc(value)
                value = tonumber(value)
                if value then
                    product = product * value
                end
            end
            Sequence.reduceRange(sequence, stepFunc, rangeMin, rangeMax)
            return product
        end
    end
    if rangeMin:isScalar() and rangeMax:isScalar() then
        return Sequence:newWithScalar(product(rangeMin:getValue(1), rangeMax:getValue(1)))
    else
        return Sequence:new{
            getter = function(index)
                return product(rangeMin:getValue(index), rangeMax:getValue(index))
            end,
            length = function()
                return math.min(rangeMin:length(), rangeMax:length())
            end,
        }
    end
end

LogicLibrary.std = function(sequence)
    local result
    if sequence then
        local mean, var = _meanAndVariance(sequence)
        result = math.sqrt(var)
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.sum = function(sequence, rangeMin, rangeMax)
    if not rangeMin or not rangeMax then
        rangeMin = Sequence:newWithScalar(1)
        rangeMax = Sequence:newWithScalar(math.huge)
    end
    local function sum(rangeMin, rangeMax)
        rangeMin, rangeMax = tonumber(rangeMin), tonumber(rangeMax)
        if rangeMin and rangeMax then
            local sum = 0
            local function stepFunc(value)
                value = tonumber(value)
                if value then
                    sum = sum + value
                end
            end
            Sequence.reduceRange(sequence, stepFunc, rangeMin, rangeMax)
            return sum
        end
    end
    if rangeMin:isScalar() and rangeMax:isScalar() then
        return Sequence:newWithScalar(sum(rangeMin:getValue(1), rangeMax:getValue(1)))
    else
        return Sequence:new{
            getter = function(index)
                return sum(rangeMin:getValue(index), rangeMax:getValue(index))
            end,
            length = function()
                return math.min(rangeMin:length(), rangeMax:length())
            end,
        }
    end
end

LogicLibrary.var = function(sequence)
    local result
    if sequence then
        local mean, var = _meanAndVariance(sequence)
        result = var
    end
    return Sequence:newWithScalar(result)
end

LogicLibrary.functions = {
    ABS = numberMap(math.abs),
    ACOS = numberMap(math.acos),
    AND = LogicLibrary.conjunction,
    ASIN = numberMap(math.asin),
    ATAN = numberMap(math.atan),
    ATAN2 = numberMap2(math.atan2),
    AVERAGE = LogicLibrary.average,
    CEILING = multiMap(2, function(value, precision)
        value, precision = tonumber(value), tonumber(precision)
        if value then
            precision = (precision and precision ~= 0 and math.abs(precision)) or 1
            return math.ceil(value / precision) * precision
        end
    end),
    CLASSIFY = function(bins, tests, counts)
        if bins and tests then
            counts = NN(counts, 1)
            local result = {}
            local binValue = {}
            for index, bin in bins:iter() do
                result[index] = 0
                binValue[index] = bin
            end
            for index, test in tests:iter() do
                for binIndex = 1, #binValue do
                    if test == binValue[binIndex] then
                        local count = tonumber(counts:getValue(index)) or 0
                        result[binIndex] = result[binIndex] + count
                    end
                end
            end
            return Sequence:newWithArray(result)
        end
    end,
    COEFF = formatMap(function() return NumberStringFormatter:coefficient() end),
    COMPACT = function(sequence) return NN(sequence):compact() end,
    CONTAINS = multiMap(3, function(haystack, needle, start)
        haystack, needle, start = todisplaystring(haystack), todisplaystring(needle), tonumber(start)
        if haystack and needle and string.find(haystack, needle, start, true) then
            return true
        end
        return false
    end),
    CORREL = LogicLibrary.correl,
    COS = numberMap(math.cos),
    COSH = numberMap(math.cosh),
    COUNT = LogicLibrary.count,
    COV = LogicLibrary.cov,
    CURRENCY = formatMap(function() return NumberStringFormatter:currency() end),
    DATE = multiMap(6, function(year, month, day, hour, minute, second) return Date:get(year, month, day, hour, minute, second) end),
    DATEISO = dateMap(function(date) return date:iso() end),
    DATESTRING = formatMap(function() return DateStringFormatter:boolean(true, false) end),
    DATETIMESTRING = formatMap(function() return DateStringFormatter:boolean(true, true) end),
    DAY = dateMap(function(date) return date:day() end),
    DECIMAL = formatMap(function() return NumberStringFormatter:decimal() end),
    DEGREES = numberMap(math.deg),
    EXP = numberMap(math.exp),
    FILTER = function(sequence, filter) return NN(sequence):filter(NN(filter)) end,
    FIND = multiMap(3, function(haystack, needle, start)
        haystack, needle, start = todisplaystring(haystack), todisplaystring(needle), tonumber(start)
        if haystack and needle then
            return string.find(haystack, needle, start, true)
        end
    end),
    FIRST = function(sequence) return Sequence:newWithScalar(NN(sequence):first()) end,
    FLOOR = multiMap(2, function(value, precision)
        value, precision = tonumber(value), tonumber(precision)
        if value then
            precision = (precision and precision ~= 0 and math.abs(precision)) or 1
            return math.floor(value / precision) * precision
        end
    end),
    GEOMEAN = LogicLibrary.geomean,
    HOUR = dateMap(function(date) return date:hour() end),
    HYPOT = numberMap2(function(x, y) return math.sqrt(x ^ 2 + y ^ 2) end),
    IF = function(condition, valueIfTrue, valueIfFalse) return NN(condition):ifMap(NN(valueIfTrue), NN(valueIfFalse)) end,
    ISODATE = map(function(input) return Date:iso(input) end),
    ISVALUE = map(function(value) return value ~= nil end),
    LAST = function(sequence) return Sequence:newWithScalar(NN(sequence):last()) end,
    LENGTH = stringMap(string.len),
    LN = numberMap(math.log),
    LOG = numberMap2(function(number, base) return math.log(number) / math.log(base) end),
    LOG10 = numberMap(function(number) return math.log(number) / math.log(10) end),
    LOWER = stringMap(string.lower),
    MAX = LogicLibrary.max,
    MEDIAN = LogicLibrary.median,
    METRIC = formatMap(function() return NumberStringFormatter:metric() end),
    MIN = LogicLibrary.min,
    MINUTE = dateMap(function(date) return date:minute() end),
    MOD = numberMap2(math.fmod),
    MODE = LogicLibrary.mode,
    MONTH = dateMap(function(date) return date:month() end),
    NEXT = function(sequence) return NN(sequence):shift(1) end,
    NOT = map(function(a) if a ~= nil then return not a end end),
    NTH = function(sequence, index) return NN(sequence):nth(NN(index)) end,
    OR = LogicLibrary.disjunction,
    PERCENT = formatMap(function() return NumberStringFormatter:percent() end),
    POW = numberMap2(math.pow),
    PREV = function(sequence) return NN(sequence):shift(-1) end,
    PRODUCT = LogicLibrary.product,
    QUANTILE = function(sequence, groups)
        local length = NN(sequence):length()
        return NN(sequence):rank():map2(function(rank, groups)
            return 1 + math.floor(groups * rank / (length + 1))
        end, NN(groups, 10))
    end,
    RADIANS = numberMap(math.rad),
    RANGE = function(min, max)
        min = math.ceil(tonumber(NN(min):first() or nil) or math.huge)
        max = math.floor(tonumber(NN(max):first() or nil) or -math.huge)
        return Sequence:new{
            getter = function(index)
                if min + index - 1 <= max then
                    return min + index - 1
                end
            end,
            length = function()
                return math.max(max - min + 1, 0)
            end,
        }
    end,
    RANK = function(sequence) return NN(sequence):rank() end,
    REPLACE = multiMap(3, safegsub),
    REVERSE = function(sequence) return NN(sequence):reverse() end,
    ROUND = multiMap(2, function(value, precision)
        value, precision = tonumber(value), tonumber(precision)
        if value then
            precision = (precision and precision ~= 0 and math.abs(precision)) or 1
            return math.floor(0.5 + value / precision) * precision
        end
    end),
    SCALE = formatMap(function() return NumberStringFormatter:scale() end),
    SCIENTIFIC = formatMap(function() return NumberStringFormatter:scientific() end),
    SECOND = dateMap(function(date) return date:second() end),
    SEQUENCE = function(...)
        local value = Sequence:newWithArray({})
        local values = {...}
        for i = 1, #values do
            value = value:join(values[i])
        end
        return value
    end,
    SIN = numberMap(math.sin),
    SINH = numberMap(math.sinh),
    SORT = function(sequence, order) return NN(sequence):sort(order or NN(sequence)) end,
    SPAN = function(min, max, steps)
        min = tonumber(NN(min):first() or nil) or math.huge
        max = tonumber(NN(max):first() or nil) or -math.huge
        steps = math.floor(tonumber(NN(steps):first() or nil) or -1)
        if min > max or steps <= 0 or steps == math.huge then
            steps = -1
        end
        return Sequence:new{
            getter = function(index)
                if index <= steps + 1 then
                    return min * (1 - (index - 1) / steps) + max * ((index - 1) / steps)
                end
            end,
            length = function()
                return steps + 1
            end,
        }
    end,
    SQRT = numberMap(math.sqrt),
    STD = LogicLibrary.std,
    SUB = multiMap(3, function(input, start, last)
        input, start, last = tostring(input), tonumber(start) or 1, tonumber(last) or -1
        if input then
            return string.sub(input, start, last)
        end
    end),
    SUM = LogicLibrary.sum,
    TAN = numberMap(math.tan),
    TANH = numberMap(math.tanh),
    TIME = multiMap(3, function(hour, minute, second) return Date:getTime(hour, minute, second) end),
    TIMESTRING = formatMap(function() return DateStringFormatter:boolean(false, true) end),
    UNIQUE = function(sequence) return NN(sequence):unique() end,
    UPPER = stringMap(string.upper),
    VAR = LogicLibrary.var,
    WEEKDAY = dateMap(function(date) return date:weekday() end),
    YEAR = dateMap(function(date) return date:year() end),
    YEARF = dateMap(function(date)
        local year = date:year()
        local rangeStart = Date:get(year):numeric()
        local rangeEnd = Date:get(year + 1):numeric()
        return year + (date:numeric() - rangeStart) / (rangeEnd - rangeStart)
    end),
}

LogicLibrary.constants = {
    NIL = Sequence:newWithScalar(nil),
    E = Sequence:newWithScalar(math.exp(1)),
    PI = Sequence:newWithScalar(math.pi),
    INFINITY = Sequence:newWithScalar(math.huge),
    FALSE = Sequence:newWithScalar(false),
    TRUE = Sequence:newWithScalar(true),
}

return LogicLibrary
