import Primes

@testset "Constants" begin
    @test Groebner.FIRST_COMPUTE_PRIME == 2^31-1
    @test Groebner.FIRST_CHECK_PRIME == 2^30 + 3
end

@testset "Lucky primes" begin
    v = [[BigInt(1), BigInt(2), BigInt(7)*123, BigInt(9)*(2^31-1)], [BigInt(11)*4]]
    vc = deepcopy(v)

    @test Groebner.isluckyprime(v, 5)
    @test !Groebner.isluckyprime(v, 11)
    @test !Groebner.isluckyprime(v, 2)
    @test !Groebner.isluckyprime(v, 7)
    @test !Groebner.isluckyprime(v, 2^31-1)
    @test Groebner.isluckyprime(v, 13)

    @test Groebner.nextluckyprime([BigInt[]]) == Groebner.FIRST_COMPUTE_PRIME
    @test Groebner.nextluckyprime(v) == Groebner.nextluckyprime(v, 2^31-1) == Primes.nextprime(2^31)
    @test Groebner.nextluckyprime([[Primes.nextprime(BigInt(2^31))]]) == 2^31-1
    @test Groebner.nextluckyprime([[Primes.nextprime(BigInt(2^31))]], 2^31-1) == Primes.nextprime(1+Primes.nextprime(BigInt(2^31)))

    @test v == vc
end

@testset "Good primes" begin
    v = [[BigInt(1), BigInt(2), BigInt(7)*123, BigInt(9)*(2^31-1)], [BigInt(11)*4]]
    vc = deepcopy(v)

    moduli = [17]
    @test Groebner.nextgoodprime(v, moduli, 16) == 19
    @test Groebner.nextgoodprime(v, [17, 19], 16) == 23
    @test Groebner.nextgoodprime(v, Int[], 17) == 19
    @test Groebner.nextgoodprime(v, Int[]) == Groebner.FIRST_CHECK_PRIME

    @test v == vc
end

@testset "Scaling coefficients" begin
    v = [BigInt(1) // 3, BigInt(2) // 6, BigInt(-5) // 12, BigInt(-3)//12]
    vv = Groebner.scale_denominators(v)
    @test vv == BigInt[4, 4, -5, -3]

    v = [
            [BigInt(1)//5, BigInt(2)//3, BigInt(7)//11],
            [BigInt(11), 1],
            [BigInt(-2)//1, BigInt(-1)//2],
            [BigInt(-1)//7, BigInt(3)//14, BigInt(-12)//7],
            [BigInt(1)//2],
            [BigInt(4115)//4115226, 1]
    ]
    vc = deepcopy(v)
    vv = Groebner.scale_denominators(v)

    @test vv == [
            [BigInt(33), BigInt(55)*2, BigInt(15)*7],
            [BigInt(11), 1],
            [BigInt(-4), BigInt(-1)],
            [BigInt(-2), BigInt(3), BigInt(-24)],
            [BigInt(1)],
            [BigInt(4115), 4115226]
    ]
    @test vc == v
end

@testset "Reconstructing coefficients modulo" begin
    v = [[BigInt(3), BigInt(4)], [BigInt(5), BigInt(6)]]
    qq = Vector{Vector{Rational{BigInt}}}(undef, 0)
    modulo = BigInt(11)
    vc = deepcopy(v)

    Groebner.reconstruct_modulo!(qq, v, modulo)
    @test v == vc
    @test qq == [[-2//3, 1//3], [-1//2, 1//2]]

    Groebner.reconstruct_modulo!(qq, v, modulo)
    @test v == vc
    @test qq == [[-2//3, 1//3], [-1//2, 1//2]]

    v = [[BigInt(9), BigInt(10)], [BigInt(11), BigInt(12)]]
    modulo = BigInt(37)
    vc = deepcopy(v)
    Groebner.reconstruct_modulo!(qq, v, modulo)
    @test v == vc
    @test qq == [[-1//4, 3//4], [-4//3, -1//3]]

end

@testset "Reconstructing coefficients crt" begin
    modulo = BigInt(1)
    ring_ff = Groebner.PolyRing(2, 3, :deglex, 5, :input)
    coeffs_ff = [[UInt64(1), UInt64(2)], [UInt64(4), UInt64(3)]]
    gbcoeffs_accum = Vector{Vector{BigInt}}(undef, 0)
    coeffs_ff_c = deepcopy(coeffs_ff)

    Groebner.reconstruct_crt!(gbcoeffs_accum, modulo, coeffs_ff, ring_ff)

    @test coeffs_ff_c == coeffs_ff
    @test gbcoeffs_accum == [[BigInt(1), 2], [BigInt(4), 3]]
    @test modulo == 5

    ring_ff = Groebner.PolyRing(2, 3, :deglex, 7, :input)
    coeffs_ff = [[UInt64(6), UInt64(1)], [UInt64(5), UInt64(3)]]
    coeffs_ff_c = deepcopy(coeffs_ff)

    Groebner.reconstruct_crt!(gbcoeffs_accum, modulo, coeffs_ff, ring_ff)
    @test coeffs_ff_c == coeffs_ff
    @test gbcoeffs_accum == [[BigInt(6), 22], [BigInt(19), 3]]
    @test modulo == 35
end

@testset "Reducing coefficients" begin
    coeffs = [
        BigInt(0), BigInt(100), BigInt(99), BigInt(6),
        BigInt(3), BigInt(4), BigInt(5), BigInt(-1),
        BigInt(-2), BigInt(-3), BigInt(-4), BigInt(-5),
        BigInt(-6), BigInt(-7), BigInt(-8), BigInt(-9),
        BigInt(-10), BigInt(-11), BigInt(-12), BigInt(-13),
        BigInt(-14), BigInt(-15), BigInt(-16), BigInt(-17)
    ]
    coeffs_v = deepcopy(coeffs)
    coeffs_ff = Vector{UInt64}(undef, length(coeffs))
    prime = 5

    Groebner.reduce_modulo!(coeffs, prime, coeffs_ff)
    @test coeffs == coeffs_v
    @test coeffs_ff == UInt64[
        0, 0, 4, 1,
        3, 4, 0, 4,
        3, 2, 1, 0,
        4, 3, 2, 1,
        0, 4, 3, 2,
        1, 0, 4, 3
    ]

    ring, coeffs_ff = Groebner.reduce_modulo([coeffs], Groebner.PolyRing(2, 3, :a, 4, :b), prime)
    @test coeffs == coeffs_v
    @test ring.ch == 5
    @test coeffs_ff == [UInt64[
        0, 0, 4, 1,
        3, 4, 0, 4,
        3, 2, 1, 0,
        4, 3, 2, 1,
        0, 4, 3, 2,
        1, 0, 4, 3
    ]]

end


#=
 Copyright 2013 - 2015 Marco Nehmeier (nehmeier@informatik.uni-wuerzburg.de)
 Copyright 2015 Oliver Heimlich (oheim@posteo.de)

 Original author: Marco Nehmeier (unit tests in libieeep1788)
 Converted into portable ITL format by Oliver Heimlich with minor corrections.

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
=#
#Language imports

#Test library imports
if VERSION < v"0.7.0-DEV.2004"
    using Base.Test
else
    using Test
end

#Arithmetic library imports
using IntervalArithmetic

#Preamble
setprecision(53)
setprecision(Interval, Float64)
setrounding(Interval, :tight)
# Set full format, and show decorations
@format full

@testset "minimal_inf_test" begin
    @test inf(∅) == Inf
    @test inf(Interval(-Inf, Inf)) == -Inf
    @test inf(Interval(1.0, 2.0)) == 1.0
    @test inf(Interval(-3.0, -2.0)) == -3.0
    @test inf(Interval(-Inf, 2.0)) == -Inf
    @test inf(Interval(-Inf, 0.0)) == -Inf
    @test inf(Interval(-Inf, -0.0)) == -Inf
    @test inf(Interval(-2.0, Inf)) == -2.0
    @test inf(Interval(0.0, Inf)) == -0.0
    @test inf(Interval(-0.0, Inf)) == -0.0
    @test inf(Interval(-0.0, 0.0)) == -0.0
    @test inf(Interval(0.0, -0.0)) == -0.0
    @test inf(Interval(0.0, 0.0)) == -0.0
    @test inf(Interval(-0.0, -0.0)) == -0.0
end

@testset "minimal_inf_dec_test" begin
    @test inf(DecoratedInterval(∅, trv)) == Inf
    @test inf(DecoratedInterval(Interval(-Inf, Inf), def)) == -Inf
    @test inf(DecoratedInterval(Interval(1.0, 2.0), com)) == 1.0
    @test inf(DecoratedInterval(Interval(-3.0, -2.0), trv)) == -3.0
    @test inf(DecoratedInterval(Interval(-Inf, 2.0), dac)) == -Inf
    @test inf(DecoratedInterval(Interval(-Inf, 0.0), def)) == -Inf
    @test inf(DecoratedInterval(Interval(-Inf, -0.0), trv)) == -Inf
    @test inf(DecoratedInterval(Interval(-2.0, Inf), trv)) == -2.0
    @test inf(DecoratedInterval(Interval(0.0, Inf), def)) == -0.0
    @test inf(DecoratedInterval(Interval(-0.0, Inf), trv)) == -0.0
    @test inf(DecoratedInterval(Interval(-0.0, 0.0), dac)) == -0.0
    @test inf(DecoratedInterval(Interval(0.0, -0.0), trv)) == -0.0
    @test inf(DecoratedInterval(Interval(0.0, 0.0), trv)) == -0.0
    @test inf(DecoratedInterval(Interval(-0.0, -0.0), trv)) == -0.0
end

@testset "minimal_sup_test" begin
    @test sup(∅) == -Inf
    @test sup(Interval(-Inf, Inf)) == Inf
    @test sup(Interval(1.0, 2.0)) == 2.0
    @test sup(Interval(-3.0, -2.0)) == -2.0
    @test sup(Interval(-Inf, 2.0)) == 2.0
    @test sup(Interval(-Inf, 0.0)) == 0.0
    @test sup(Interval(-Inf, -0.0)) == 0.0
    @test sup(Interval(-2.0, Inf)) == Inf
    @test sup(Interval(0.0, Inf)) == Inf
    @test sup(Interval(-0.0, Inf)) == Inf
    @test sup(Interval(-0.0, 0.0)) == 0.0
    @test sup(Interval(0.0, -0.0)) == 0.0
    @test sup(Interval(0.0, 0.0)) == 0.0
    @test sup(Interval(-0.0, -0.0)) == 0.0
end

@testset "minimal_sup_dec_test" begin
    @test sup(DecoratedInterval(∅, trv)) == -Inf
    @test sup(DecoratedInterval(Interval(-Inf, Inf), def)) == Inf
    @test sup(DecoratedInterval(Interval(1.0, 2.0), com)) == 2.0
    @test sup(DecoratedInterval(Interval(-3.0, -2.0), trv)) == -2.0
    @test sup(DecoratedInterval(Interval(-Inf, 2.0), dac)) == 2.0
    @test sup(DecoratedInterval(Interval(-Inf, 0.0), def)) == 0.0
    @test sup(DecoratedInterval(Interval(-Inf, -0.0), trv)) == 0.0
    @test sup(DecoratedInterval(Interval(-2.0, Inf), trv)) == Inf
    @test sup(DecoratedInterval(Interval(0.0, Inf), def)) == Inf
    @test sup(DecoratedInterval(Interval(-0.0, Inf), trv)) == Inf
    @test sup(DecoratedInterval(Interval(-0.0, 0.0), dac)) == +0.0
    @test sup(DecoratedInterval(Interval(0.0, -0.0), trv)) == +0.0
    @test sup(DecoratedInterval(Interval(0.0, 0.0), trv)) == +0.0
    @test sup(DecoratedInterval(Interval(-0.0, -0.0), trv)) == +0.0
end

@testset "minimal_mid_test" begin
    @test mid(Interval(-Inf, Inf)) == 0.0
    @test mid(Interval(-0x1.fffffffffffffp1023, +0x1.fffffffffffffp1023)) == 0.0
    @test mid(Interval(0.0, 2.0)) == 1.0
    @test mid(Interval(2.0, 2.0)) == 2.0
    @test mid(Interval(-2.0, 2.0)) == 0.0
    @test mid(Interval(-0x0.0000000000002p-1022, 0x0.0000000000001p-1022)) == 0.0
    @test mid(Interval(-0x0.0000000000001p-1022, 0x0.0000000000002p-1022)) == 0.0
end

@testset "minimal_mid_dec_test" begin
    @test mid(DecoratedInterval(Interval(-Inf, Inf), def)) == 0.0
    @test mid(DecoratedInterval(Interval(-0x1.fffffffffffffp1023, +0x1.fffffffffffffp1023), trv)) == 0.0
    @test mid(DecoratedInterval(Interval(0.0, 2.0), com)) == 1.0
    @test mid(DecoratedInterval(Interval(2.0, 2.0), dac)) == 2.0
    @test mid(DecoratedInterval(Interval(-2.0, 2.0), trv)) == 0.0
    @test mid(DecoratedInterval(Interval(-0x0.0000000000002p-1022, 0x0.0000000000001p-1022), trv)) == 0.0
    @test mid(DecoratedInterval(Interval(-0x0.0000000000001p-1022, 0x0.0000000000002p-1022), trv)) == 0.0
end

@testset "minimal_rad_test" begin
    @test radius(Interval(0.0, 2.0)) == 1.0
    @test radius(Interval(2.0, 2.0)) == 0.0
    @test radius(Interval(-Inf, Inf)) == Inf
    @test radius(Interval(0.0, Inf)) == Inf
    @test radius(Interval(-Inf, 1.2)) == Inf
end

@testset "minimal_rad_dec_test" begin
    @test radius(DecoratedInterval(Interval(0.0, 2.0), trv)) == 1.0
    @test radius(DecoratedInterval(Interval(2.0, 2.0), com)) == 0.0
    @test radius(DecoratedInterval(Interval(-Inf, Inf), trv)) == Inf
    @test radius(DecoratedInterval(Interval(0.0, Inf), def)) == Inf
    @test radius(DecoratedInterval(Interval(-Inf, 1.2), trv)) == Inf
end

@testset "minimal_wid_test" begin
    @test diam(Interval(2.0, 2.0)) == 0.0
    @test diam(Interval(1.0, 2.0)) == 1.0
    @test diam(Interval(1.0, Inf)) == Inf
    @test diam(Interval(-Inf, 2.0)) == Inf
    @test diam(Interval(-Inf, Inf)) == Inf
end

@testset "minimal_wid_dec_test" begin
    @test diam(DecoratedInterval(Interval(2.0, 2.0), com)) == 0.0
    @test diam(DecoratedInterval(Interval(1.0, 2.0), trv)) == 1.0
    @test diam(DecoratedInterval(Interval(1.0, Inf), trv)) == Inf
    @test diam(DecoratedInterval(Interval(-Inf, 2.0), def)) == Inf
    @test diam(DecoratedInterval(Interval(-Inf, Inf), trv)) == Inf
end

@testset "minimal_mag_test" begin
    @test mag(Interval(1.0, 2.0)) == 2.0
    @test mag(Interval(-4.0, 2.0)) == 4.0
    @test mag(Interval(-Inf, 2.0)) == Inf
    @test mag(Interval(1.0, Inf)) == Inf
    @test mag(Interval(-Inf, Inf)) == Inf
    @test mag(Interval(-0.0, 0.0)) == 0.0
    @test mag(Interval(-0.0, -0.0)) == 0.0
end

@testset "minimal_mag_dec_test" begin
    @test mag(DecoratedInterval(Interval(1.0, 2.0), com)) == 2.0
    @test mag(DecoratedInterval(Interval(-4.0, 2.0), trv)) == 4.0
    @test mag(DecoratedInterval(Interval(-Inf, 2.0), trv)) == Inf
    @test mag(DecoratedInterval(Interval(1.0, Inf), def)) == Inf
    @test mag(DecoratedInterval(Interval(-Inf, Inf), trv)) == Inf
    @test mag(DecoratedInterval(Interval(-0.0, 0.0), trv)) == 0.0
    @test mag(DecoratedInterval(Interval(-0.0, -0.0), trv)) == 0.0
end

@testset "minimal_mig_test" begin
    @test mig(Interval(1.0, 2.0)) == 1.0
    @test mig(Interval(-4.0, 2.0)) == 0.0
    @test mig(Interval(-4.0, -2.0)) == 2.0
    @test mig(Interval(-Inf, 2.0)) == 0.0
    @test mig(Interval(-Inf, -2.0)) == 2.0
    @test mig(Interval(-1.0, Inf)) == 0.0
    @test mig(Interval(1.0, Inf)) == 1.0
    @test mig(Interval(-Inf, Inf)) == 0.0
    @test mig(Interval(-0.0, 0.0)) == 0.0
    @test mig(Interval(-0.0, -0.0)) == 0.0
end

@testset "minimal_mig_dec_test" begin
    @test mig(DecoratedInterval(Interval(1.0, 2.0), com)) == 1.0
    @test mig(DecoratedInterval(Interval(-4.0, 2.0), trv)) == 0.0
    @test mig(DecoratedInterval(Interval(-4.0, -2.0), trv)) == 2.0
    @test mig(DecoratedInterval(Interval(-Inf, 2.0), def)) == 0.0
    @test mig(DecoratedInterval(Interval(-Inf, -2.0), trv)) == 2.0
    @test mig(DecoratedInterval(Interval(-1.0, Inf), trv)) == 0.0
    @test mig(DecoratedInterval(Interval(1.0, Inf), trv)) == 1.0
    @test mig(DecoratedInterval(Interval(-Inf, Inf), trv)) == 0.0
    @test mig(DecoratedInterval(Interval(-0.0, 0.0), trv)) == 0.0
    @test mig(DecoratedInterval(Interval(-0.0, -0.0), trv)) == 0.0
end
# FactCheck.exitstatus()


@enum MatrixShape FullSymbolic UpperTriangular UnitUpperTriangular Companion UserSpecific

_FullSymbolic(s::Int)        = [mkpoly(mkvar("b$i$j")) for i in 1:s, j in 1:s]
_UpperTriangular(s::Int)     = [j>=i ? mkpoly(mkvar("b$i$j")) : mkpoly(0) for i in 1:s, j in 1:s]
_UnitUpperTriangular(s::Int) = [j>i ? mkpoly(mkvar("b$i$j")) : i==j ? mkpoly(1) : mkpoly(0) for i in 1:s, j in 1:s]
_Companion(s::Int)           = [i==s ? mkpoly(mkvar("b$i$j")) : i+1==j ? mkpoly(1) : mkpoly(0) for i in 1:s, j in 1:s]
_UserSpecific(s::Int)        = error("Should not be called")

function _add_const_one(M::Matrix)
    s = size(M, 1) + 1
    col = [mkpoly(mkvar("b$i$s")) for i in 1:s-1]
    _add_row_one(hcat(M, col))
end

function _add_row_one(M::Matrix)
    T = eltype(M)
    M = vcat(M, zeros(T, 1, size(M, 2)))
    M[end,end] = one(T)
    M
end

function initmatrix(vars::Vector{Symbol}, params::Vector{SymOrNum})
    rows, cols = length(vars), length(params)
    A = fill(mkpoly(1), (rows,cols))

    for i in 1:rows, j in 1:cols
        u, v = vars[i], params[j]
        if u == v
            A[i,j] = j == findfirst(x->x==u, params) ? 1 : 0
        else
            A[i,j] = findfirst(x->x==u, params) === nothing ? mkvar("a$i$j") : 0
        end
    end
    A
end

function bodymatrix(s::Int, shape::MatrixShape)
    f = Meta.parse(string("_", shape))
    eval(:($f($s)))
end

# ------------------------------------------------------------------------------

struct RecTemplate
    vars::Vector{Symbol}
    params::Vector{SymOrNum}
    body::Matrix{<:Poly}
    init::Matrix{<:Poly}
end

function RecTemplate(vars::Vector{Symbol}; shape::MatrixShape=UnitUpperTriangular, constant_one::Bool=true, params::Vector{Symbol}=Symbol[])
    size = length(vars)
    params = SymOrNum[params; 1]

    B = bodymatrix(size, shape)
    A = initmatrix(vars, params)

    if constant_one
        push!(vars, CONST_ONE_SYM)
        B = _add_const_one(B)
        A = _add_row_one(A)
    end

    RecTemplate(vars, params, B, A)
end

variables(r::RecTemplate) = r.vars

has_constant_one(r::RecTemplate) = r.vars[end] == CONST_ONE_SYM

function Base.summary(io::IO, r::RecTemplate)
    compact = get(io, :compact, false)
    vars = [v == CONST_ONE_SYM ? 1 : v for v in r.vars]
    if compact
        print(io, "$(typeof(r)) for some permutation of ")
        print(io, vars)
    else
        print(io, "$(typeof(r)) for some permutation σ of ")
        print(io, vars)
    end
end

function Base.show(io::IO, r::RecTemplate)
    summary(io, r)
    compact = get(io, :compact, false)
    if !compact
        println(io, ":")
        _show(io, r)
    end
end

function _show(io::IO, r::RecTemplate)
    init = r.init * map(mkpoly, r.params)
    vars = [v == CONST_ONE_SYM ? "σ(1)" : "σ($v)" for v in r.vars]
    _print_recsystem(io, vars, r.body, init)
    # @time _print_rectemplate(io, vars, r.body, init)
end

lpar2(h::Int, d = "") = h == 1 ? "($(d)" : ["⎛$(d)"; fill("⎜$(d)", h - 2); "⎝$(d)"]
rpar2(h::Int, d = "") = h == 1 ? "$(d))" : ["$(d)⎞"; fill("$(d)⎟", h - 2); "$(d)⎠"]

space2(h::Int, sp = " ") = fill(sp, h)

function symstr2(h::Int, symbol::String)
    a = fill(" ", h)
    a[Int(ceil(h/2))] = "$(symbol)"
    a
end

function _print_rectemplate(io::IO, vars, body, init)
    h = size(body, 1)
    lstr, rstr = lpar2(h), rpar2(h, "")
    eq = symstr2(h, "=")

    # lc = "\u2099"
    # lp, rp, plus = "", "", "\u208A"
    # zero, one = "\u2080", "\u2081"

    lc = "n"
    lp, rp, plus = "(", ")", "+"
    zero, one = "0", "1"

    vars0 = ["$v$lp$zero$rp" for v in vars]
    vars1 = ["$v$lp$lc$rp" for v in vars]
    vars2 = ["$v$lp$lc$plus$one$rp" for v in vars]

    body = map(prettify_number, body)
    init = map(prettify_number, init)
    # add extra column such that all rows have equal length
    # body = hcat(body, fill(nothing, size(body, 1)))
    # init = hcat(init, fill(nothing, size(init, 1)))
    # body = sprint(Base.print_matrix, body)
    # init = sprint(Base.print_matrix, init)
    # body = Base.replace(body, "  $nothing" => "")
    # init = Base.replace(init, "  $nothing" => "")

    x = hcat(space2(h, " "), lstr, vars2, rstr, eq, lstr, body, rstr, lstr, vars1, rstr, space2(h, "    "), lstr, vars0, rstr, eq, lstr, init, rstr, space2(h, ""))
    str = sprint((args...)->Base.print_matrix(args...; ), x)
    str = Base.replace(str, "\""=>"")
    str = Base.replace(str, " ⎛ "=>"⎛")
    str = Base.replace(str, " ⎜ "=>"⎜")
    str = Base.replace(str, " ⎝ "=>"⎝")
    str = Base.replace(str, " ⎞ "=>"⎞")
    str = Base.replace(str, " ⎟ "=>"⎟")
    str = Base.replace(str, " ⎠ "=>"⎠")

    print(io, str)
end


# samplepath.jl

"""
    abstract type AbstractSamplePath <: Any

subtype: SamplePath.

Its subtype implements a continuous samplepath.
"""
abstract type AbstractSamplePath end

"""
    isincreasing(x::AbstractVector{T})::Bool where T <: Number
    isincreasing(x::Range{T})::Bool where T <: Number

Tests whether a vector of numbers is strictly increasing. Is internal to
NonparametricBayesForDiffusions.
"""
function isincreasing(x::AbstractVector{T}) where T <: Number
  for i in 1:length(x)-1
    x[i+1] <= x[i] && return false
  end
  return true
end

isincreasing(x::Range{T}) where T <: Number = step(x) > 0

"""
    SamplePath(timeinterval::S, samplevalues::T) where
        {S<:AbstractVector{Float64}, T<:AbstractVector{Float64}}

 Implements a continuous samplepath where the timeinterval is not necessarily
 equally spaced. Sample value samplevalues[k] is the value of the process at
 time timeinterval[k]. timeinterval is an increasing vector. timeinterval and
 samplevalues are of equal length.

 Constructors:
 SamplePath(timeinterval, samplevalues) (as above)
 SamplePath(timeinterval, f::Function) =
    SamplePath(timeinterval, f.(timeinterval))

 # Examples

```julia
using PyPlot
clf()
t = 0.0:0.1:2.0
v = sinpi.(t)
X = SamplePath(t, v)
plot(X)
```
"""
struct SamplePath{S<:AbstractVector{Float64},
        T<:AbstractVector{Float64}} <: AbstractSamplePath
  timeinterval::S
  samplevalues::T

  function SamplePath(timeinterval::S, samplevalues::T) where
      {S<:AbstractVector{Float64}, T<:AbstractVector{Float64}}
    length(timeinterval) == length(samplevalues) ||
      error("Length of timeinterval should be equal to the length of the samplevalues vector")
    isincreasing(timeinterval) ||
      error("Timeinterval should be increasing")
    new{S,T}(timeinterval, samplevalues)
  end
end

#constructor
SamplePath(timeinterval, f::Function) = SamplePath(timeinterval, f.(timeinterval))

step(X::SamplePath{S}) where S<:Range{<:Number} = step(X.timeinterval)

# Returns the length of timeinterval == samplevalues vectors, not the endtime!
# TO DO: Answers to examples
"""
    length(X::SamplePath)

Returns the length of the vector timeinterval == length vector samplevalues.

# Examples
```julia
X = SamplePath([0.,1.,2.], [3.,5., -1.])
length(X)
```

```julia
X = SamplePath(0.0:0.1:2π, sin)
length(X) == length(0.0:0.1:2π)
```
"""
length(X::SamplePath)=length(X.timeinterval)


using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges
const MOIBC = MathOptInterface.Bridges.Constraint

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()))
config = MOIT.TestConfig(duals = false)

@testset "RSOCtoNonConvexQuad" begin

    @test MOIBC.RSOCtoNonConvexQuadBridge{Float64} == MOIBC.concrete_bridge_type(
        MOIBC.RSOCtoNonConvexQuadBridge{Float64},
        MOI.VectorOfVariables,
        MOI.RotatedSecondOrderCone)
    @test MOI.supports_constraint(MOIBC.RSOCtoNonConvexQuadBridge{Float64},
                                  MOI.VectorOfVariables,
                                  MOI.RotatedSecondOrderCone)
    @test !MOI.supports_constraint(MOIBC.RSOCtoNonConvexQuadBridge{Float64},
                                   MOI.ScalarAffineFunction{Float64},
                                   MOI.RotatedSecondOrderCone)

    bridged_mock = MOIB.Constraint.RSOCtoNonConvexQuad{Float64}(mock)

    MOIT.basic_constraint_tests(
        bridged_mock, config,
        include = [(F, S)
                   for F in [MOI.VectorOfVariables]
                   for S in [MOI.RotatedSecondOrderCone]])

    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [0.5, 1.0, 1/√2, 1/√2])
    MOIT.rotatedsoc1vtest(bridged_mock, config)

    ci = first(MOI.get(bridged_mock,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables,
                                    MOI.RotatedSecondOrderCone}()))

    test_delete_bridge(bridged_mock, ci, 4,
                    (
                        (MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}, 0),
                        (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}, 0),
                    ))
end

@testset "SOCtoNonConvexQuad" begin

    @test MOIBC.SOCtoNonConvexQuadBridge{Float64} == MOIBC.concrete_bridge_type(
        MOIBC.SOCtoNonConvexQuadBridge{Float64},
        MOI.VectorOfVariables,
        MOI.SecondOrderCone)
    @test MOI.supports_constraint(MOIBC.SOCtoNonConvexQuadBridge{Float64},
                                  MOI.VectorOfVariables,
                                  MOI.SecondOrderCone)
    @test !MOI.supports_constraint(MOIBC.SOCtoNonConvexQuadBridge{Float64},
                                   MOI.ScalarAffineFunction{Float64},
                                   MOI.SecondOrderCone)

    bridged_mock = MOIB.Constraint.SOCtoNonConvexQuad{Float64}(mock)

    MOIT.basic_constraint_tests(
        bridged_mock, config,
        include = [(F, S)
                   for F in [MOI.VectorOfVariables]
                   for S in [MOI.SecondOrderCone]])

    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0, 1/√2, 1/√2])
    MOIT.soc1vtest(bridged_mock, config)

    ci = first(MOI.get(bridged_mock,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables,
                                    MOI.SecondOrderCone}()))

    test_delete_bridge(bridged_mock, ci, 3,
                       (
                           (MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}, 0),
                           (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}, 0),
                       ))
end


function myfunction(n)
    for i = 1:n
        A = randn(100,100,20)
        m = maximum(A)
        Am = mapslices(sum, A; dims=2)
        B = A[:,:,5]
        Bsort = mapslices(sort, B; dims=1)
        b = rand(100)
        C = B.*b
    end
end

myfunction(1)

@profiler myfunction(10)

# Short for:
# Profile.clear()
# @profile myfunction(10)
# Juno.profiler()


"""
extract a fully sampled calibration area of size `M` by `N` from the Matrix `A`

# Arguments
* `A::Array{Int64,2}` - Array to be sampled
* `M::Int64`          - size of A along the first dimension
* `N::Int64`          - size of N along the second dimension
* `redFac::Float64`   - reduction factor of the target sampling pattern (used for a sanity check)
* `calsize::Int64`    - size of the calibration area
"""
function setcalibrationarea(A::Array{Int64,2},M::Int64,N::Int64,redFac::Float64,calsize::Int64)
  cal_x, cal_y = min(calsize,M), min(calsize,N)
  ChosenSamples = 0
  NumberSamples = floor(Int64,M*N/redFac)
  if calsize > 0
    if cal_x*cal_y > NumberSamples
      @error "Number of calibration samples exceeds number of required total samples"
    end
    ChosenSamples = cal_x*cal_y
    A[floor(Int64,end/2)-floor(Int64,cal_x/2)+1:floor(Int64,end/2)-floor(Int64,cal_x/2)+cal_x,floor(Int64,end/2)-floor(Int64,cal_y/2)+1:floor(Int64,end/2)-floor(Int64,cal_y/2)+cal_y].=1
  end

  return A, ChosenSamples

end


using CairoMakie

"""
plot_mcmc_densities!(chain_X; n_columns=3)

Plots marginal density distributions for each parameter in the
constituent parameter vectors of `chain_X`, each representing a sample
generated by MCMC.
"""
function plot_mcmc_densities2!(fig, axes, chain_X, parameter_names;
                                n_columns = 3,
                                directory = pwd(),
                                filename = "mcmc_densities.png",
                                label = false,
                                color = (:blue, 0.5),
                                type = "hist",
                                show_means = false,
                                kwargs...)
    @assert type in ["hist", "density"]

    sample_means = []
    for (i, param_name) in enumerate(parameter_names)
        samples = getindex.(chain_X, i)

        ax = axes[i+1]
        if type == "hist"
            hist!(ax, samples; label, bins = 50, color, normalization = :pdf, kwargs...)
        elseif type == "density"
            density!(ax, samples; label, color, bandwidth=sqrt(var(samples))/15, kwargs...)
        end

        push!(sample_means, mean(samples))
        show_means && vlines!(ax,  [mean(samples)]; color=(color[1], 1.0), label="Sample mean", linestyle=nothing, linewidth=5)
    end

    fig[1,1] = Legend(fig, axes[end], nothing; framevisible=true,
                                                            tellheight = false,
                                                            tellwidth = false,
                                                            nbanks = 2)
    save(joinpath(directory, filename), fig)

    return sample_means
end

function plot_mcmc_densities2(chain_X, parameter_names;
                                n_columns = 3, kwargs...)

    N_axes = length(first(chain_X)) + 1
    n_rows = Int(ceil(N_axes / n_columns))

    fig = Figure(resolution = (500n_columns, 200n_rows))

    ax_coords = [(i, j) for i = 1:n_rows, j = 1:n_columns]

    # Reserve position [1,1] for the legend
    ax1 = Axis(fig[1,1])
    hidedecorations!(ax1)
    hidespines!(ax1)
    axes = [ax1]
    for (i, param_name) in enumerate(parameter_names)
        coords = ax_coords[i+1]
        ax = Axis(fig[coords...],
            xlabel = string(param_name),
            ylabel = "Density")
        ax.xlabelsize = 20
        push!(axes, ax)
    end

    sample_means = plot_mcmc_densities2!(fig, axes, chain_X, parameter_names; n_columns, kwargs...)

    return fig, axes
end


include("../src/sir.jl")
using PyPlot
K = 50
Np = 2500
σ_o = 0.1
σ_d = 0.0
Δ = 5
Nth = 500
obsfun = obs1

x_trj, w_trj, y, x_true = assimilate(K, Np,
							σ_o, σ_d,
							Δ, Nth,
							obsfun)
x_trj = permutedims(x_trj,[3,2,1])
trj_len = K*Δ
y = permutedims(y,[2,1])
w_trj = permutedims(w_trj,[2,1])
x_true = permutedims(x_true,[2,1])

function plot_trj()
	#n_t = 10
	#trj_len = size(x_true)[1] - n_t + 1
	time_arr = LinRange(0, trj_len -1, trj_len)
	time_arr = dt*time_arr

	for i = 1:d
		fig, ax = subplots()
		ax.grid(true)
		ax.plot(time_arr, x_trj[:,:,i],"b",lw=0.9)
		ax.plot(time_arr, x_trj[:,end,i],"b",lw=0.9,label="\$ x^{($i)}\$ orbit of particle")
		if i==d
		    ax.plot(time_arr[1:Δ:end], y[:,1],"k",lw=3.0,label="Observed \$ x^{($i)}\$")
		end
		ax.legend(fontsize=28)
		ax.xaxis.set_tick_params(labelsize=28)
		ax.yaxis.set_tick_params(labelsize=28)
		ax.set_xlabel("Time",fontsize=28)
		ax.set_title("# particles = $(Np), partial observation", fontsize=28)
	end
end
function plot_rmse()
    filter_mean = reshape(sum(x_trj,dims=2)/Np, trj_len, d)
    norm2(x,dim) = sum(x -> x^2, x, dims=dim)
    rmse = norm2(x_true .- filter_mean, 2)
	rmse .= sqrt.(rmse)/d
	time_arr = LinRange(0, trj_len -1, trj_len)
	time_arr = dt*time_arr
	fig, ax = subplots()
	ax.grid(true)
	ax.plot(time_arr, rmse,"r",lw=3.0,label="RMSE")
	ax.legend(fontsize=28)
	ax.xaxis.set_tick_params(labelsize=28)
	ax.yaxis.set_tick_params(labelsize=28)
	ax.set_xlabel("Time",fontsize=28)
	ax.set_title("# particles = $(Np), partial observation", fontsize=28)
end
function plot_weight_stats()
	fig, ax = subplots()
	ax.grid(true)
	ax.plot(w_trj,"b",lw=0.2)
	ax.plot(w_trj[:,end],"b",lw=1.0,label="particle weight orbit")
	ax.legend(fontsize=28)
	ax.xaxis.set_tick_params(labelsize=28)
	ax.yaxis.set_tick_params(labelsize=28)
end


path = dirname(@__FILE__)

# Initialize arguments to function
@load "$path/reference/kalman_filter_args.jld2" y T R C Q Z D E z0 P0

file  = h5open("$path/reference/kalman_filter_out.h5", "r")
pred  = read(file, "pred")
vpred = read(file, "vpred")
filt  = read(file, "filt")
vfilt = read(file, "vfilt")
close(file)

# Run smoothers
states = Dict{Symbol, Matrix{Float64}}()
shocks = Dict{Symbol, Matrix{Float64}}()

states[:hamilton], shocks[:hamilton] = hamilton_smoother(y, T, R, C, Q, Z, D, E, z0, P0)
states[:koopman], shocks[:koopman] = koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0, pred, vpred)
states[:carter_kohn], shocks[:carter_kohn] = carter_kohn_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = false)
states[:durbin_koopman], shocks[:durbin_koopman] = durbin_koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = false)

# Check that last-period smoothed states equal last-period filtered states
@testset "Ensure last smoothed state = the last filtered state" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        @test filt[:, end] ≈ states[smoother][:, end]
    end
end

# Compare to expected output
exp_states, exp_shocks = h5open("$path/reference/smoothers_out.h5", "r") do file
    read(file, "exp_states"), read(file, "exp_shocks")
end

@testset "Check expected output for all smoothers" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        @test exp_states ≈ states[smoother]
        @test exp_shocks ≈ shocks[smoother]
    end
end

# Make sure that simulation smoothers run with `draw_states` on
carter_kohn_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = true)
durbin_koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = true)

nothing


"""
   DifferentialExtinction
   Original Python code (solar_diffex.py) by Andrea Lin
   Ported to Julia by Eric Ford
   Based on Collier-Cameron et al. (2019) MNRAS 487, 1082.  doi:10.1093/mnras/stz1215
   See also Astronomical Algorithms (Meeus, 1991) p151 & p178
"""
__precompile__() # this module is safe to precompile
module DifferentialExtinction
using PyCall
using Dates
using DataFrames, CSV

const JplHorizons = PyNULL()

function __init__()
	copy!(JplHorizons , pyimport("astroquery.jplhorizons") )
end

import ..NeidSolarScripts
import ..NeidSolarScripts.SolarRotation

AU = 149597870.7  # km
R_sun = 695700.0  # km # should we use better, but non-standard value, e.g, 695660?
valid_obs = [:WIYN, :HARPSN]

""" get_obs_loc(obs::Symbol)
 Returns a Dict with long & lat (degrees) and elevation (km)
 Warning: Currently only has info for :WIYN and :HARPSN.
"""
function get_obs_loc(obs::Symbol)
	@assert obs ∈ valid_obs
	if obs == :HARPSN
		loc = Dict("lon"=> -17.88905, "lat"=> 28.754, "elevation"=> 2.3872)
	elseif obs == :WIYN
		loc = Dict("lon"=> -111.600562, "lat"=> 31.958092, "elevation"=> 2.091)
	else
		@error("Don't have coordinates for obs = " * string(obs))
	end
	return loc
end

f_z(alt) = π/2-alt

function f_airmass(z)
	secz = sec(z)
	return secz * (1.0 - 0.0012*(secz^2 - 1.0))
end

"""fractional brightness diff (from sun center)"""
function f_epsilon(z_center, dist_sun) #
        z_upper = z_center - asin(R_sun/dist_sun)
        z_lower = z_center + asin(R_sun/dist_sun)
        delta_airmass = f_airmass(z_upper) - f_airmass(z_lower)
        k_ex = 0.157482  # WARNING: Hardwired for WIYN.  TODO: Should recalc for each obs?
		#=
		 X = [ones(length(eph.columns["airmass"])) eph.columns["airmass"]]
		 y =  eph.columns["magextinct"]
		 coeff = X \y
		 k_ex = coeff[2]
		 =#
		return k_ex * delta_airmass/2.0 * log(10.)/2.5
end

"parallactic angle"
function f_q(H_sun, z_center, lat)
	#lat = deg2rad(wiyn["lat"]) #* u.deg # observer's latitude # no longer hard-wired
	return asin(sin(H_sun) * cos(lat)/sin(z_center))
end

f_phi(pole_angle, q) = (pole_angle - q)

function f_pole_angle_ecl(alpha_sun)
	e_tilt = deg2rad(23.4) # ecliptic tilt
    asin(sin(e_tilt) * cos(alpha_sun))
end

f_psi(pole_angle_ecl, q) = (pole_angle_ecl - q)

 """
 f_B0(JD)
  see Astronomical Algorithms (Meeus, 1991) p151 & p178
 """
 function f_B0(JD)

    T = (JD - 2451545.) / 36525.
    L0 = deg2rad(280.46645 + 36000.76983*T + 0.0003032*T^2)
    MA = deg2rad(357.52910 + 35999.05030*T - 0.0001559*T^2 - 0.00000048*T^3)
    sun_center = (1.914600 - 0.004817*T - 0.000014*T^2) * sin(MA) +
    			 (0.019993 - 0.000101*T) * sin(2*MA) +
    			 0.000290 * sin(3*MA)
	sun_center = deg2rad(sun_center)
    sun_true_lon = L0 + sun_center
    lam = sun_true_lon - deg2rad(0.00569) - deg2rad(0.00478 * sin(deg2rad(125.04 - 1934.136*T)))
    kappa = deg2rad(73.6667 + 1.3958333 * (JD - 2396758.) / 36525.)
    return asin(sin(lam - kappa) * sin(deg2rad(7.25)))
end

function diff_extinct_corr(JD, alt, HA, RA, pole_angle, dist_sun, lat)
    z = f_z(deg2rad(alt))
    epsilon = f_epsilon(z, dist_sun)
    q = f_q(deg2rad(HA*15), z, deg2rad(lat))

    phi = (f_phi(deg2rad(pole_angle), q))
    psi = f_psi(f_pole_angle_ecl(deg2rad(RA)), q)

    inc_sol_axis = f_B0(JD)

	ld_coeff = 0.6 # limb darkening coeff (linear)

	# sidereal diff rotation (Snodgrass & Ulrich 1990)
	sr_A = 2.972e-6 #* u.rad / u.s
	sr_B = -0.484e-6 #* u.rad / u.s
	sr_C = -0.361e-6 #* u.rad / u.s

	omega_orb = (2*π) / (365.25 * 24. * 3600.) #* u.rad / u.s # omega of Earth around Sun

    diff_rot_fac = sr_A*(7*ld_coeff - 15.) / (20*(3-ld_coeff)) +
      sr_B*(19 *ld_coeff - 35 + cos(inc_sol_axis)^2 * (3*ld_coeff - 35)) / (280*(3-ld_coeff)) +
      sr_C*(187*ld_coeff - 315 - cos(inc_sol_axis)^2 * (630 - 118*ld_coeff) -
      105 *cos(inc_sol_axis)^4 * (ld_coeff-1.)) / (6720*(3-ld_coeff))

    delta_vr = (epsilon * sin(phi) * R_sun * sin(inc_sol_axis) * diff_rot_fac -
       (7*ld_coeff - 15) * epsilon * sin(psi) / (20*(3-ld_coeff)) * R_sun * omega_orb)
	delta_vr *= 1e3 # km/s -> m/s.  WARN: Andrea's python code returns cm/s

    return delta_vr
end

"""
   calc_Δv_diff_extinction(bjd; obs, ...)
   calc_Δv_diff_extinction(start, stop; obs, dt, dt_unit, ...)
Estimate apparent Δv (m/s) in solar observations due to differential extinction.
"""
function calc_Δv_diff_extinction end

function query_horizons(;epochs::Union{Vector{Float64},Dict}, obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	if typeof(epochs) <: Dict
		@assert haskey(epochs,"start")
		@assert haskey(epochs,"stop")
		@assert haskey(epochs,"step")
	elseif typeof(epochs) <: Vector{Float64}
		@assert 1 <= length(epochs) <= 10000
	end
	@assert 1 < max_airmass <= 10.0
	loc = get_obs_loc(obs)
	# Sometimes this fails.  Why?
	obj = JplHorizons.Horizons(id="10", location=loc, epochs=epochs, id_type="majorbody")
	eph = obj.ephemerides(quantities="1,4,8,17,20,42", airmass_lessthan=max_airmass, refraction=true)
	if verbose  println(eph)  end
	return eph
end

function calc_Δv_diff_extinction_horizons(;epochs::Union{Vector{Float64},Dict}, obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	loc = get_obs_loc(obs)
	eph = query_horizons(epochs=epochs,obs=obs,max_airmass=max_airmass, verbose=verbose)
	delta_vr = diff_extinct_corr.(eph.columns["datetime_jd"], eph.columns["EL"],
	 	eph.columns["hour_angle"], eph.columns["RA"],
		eph.columns["NPole_ang"], eph.columns["delta"].*AU, loc["lat"])
	return (jd=eph.columns["datetime_jd"], Δv=delta_vr)
end

function cache_horizons_output(year::Int64, obs::Symbol)
	fn = "horizons_" * string(obs) * "_" * string(year) * ".csv"
	filename = joinpath(dirname(pathof(NeidSolarScripts)),"..","data","horizons",fn)
	if isfile(filename)
		return
	end
	df = DataFrame(:jd=>Float64[], :sol_npole_ang=>Float64[], :sol_dist_au=>Float64[])
	samples_per_day = 4
	epoch_start = Dates.datetime2julian(DateTime(year,1,1))
	for i in 1:samples_per_day
		epoch_stop = epoch_start + ceil(Int64,366/samples_per_day)
		epochs = collect(range(epoch_start,stop=epoch_stop,step=1.0/samples_per_day) )
		eph = query_horizons(epochs=epochs, obs=obs, max_airmass=10.0)
		df_tmp = DataFrame(:jd=>eph.columns["datetime_jd"], :sol_npole_ang=>eph.columns["NPole_ang"], :sol_dist_au=>eph.columns["delta"])
		append!(df, df_tmp)
		epoch_start = epoch_stop
	end
	CSV.write(filename,df)
	return df
end

function read_horizons_output(jd::Union{Float64,Vector{Float64}}, obs::Symbol)
	@assert obs ∈ valid_obs
	years = unique(year.(Dates.julian2datetime.(jd)))
	@assert length(years) >= 1
	cache = Dict{Int64,Any}()
	for year in years
		fn = "horizons_" * string(obs) * "_" * string(year) * ".csv"
		filename = joinpath(dirname(pathof(NeidSolarScripts)),"..","data","horizons",fn)
		if isfile(filename)
			cache[year] = CSV.read(filename, DataFrame)
		else
			cache[year] = cache_horizons_output(year, obs)
		end
		cache[year].sol_npole_ang[cache[year].sol_npole_ang .> 180] .-= 360
	end
	if length(years) == 1
		return cache[first(keys(cache))]
	else
		sorted_keys = sort(collect(keys(cache)))
		output = cache[first(sorted_keys)]
		for k in sorted_keys[2:end]
			append!(output,cache[k])
		end
		return output
	end
end

function get_horizons_output(jd::Float64, obs::Symbol)
  df = read_horizons_output(jd, obs)
  idxlo = searchsortedlast(df.jd, jd)
  @assert 1 <= idxlo < size(df,1)
  idxhi = idxlo + 1
  result = Dict{Symbol,Float64}()
  num = jd-df.jd[idxlo]
  denom = df.jd[idxhi]-df.jd[idxlo]
  whi = num/denom
  wlo = 1-whi
  for k in Symbol.(names(df))
	if k == :jd continue end
  	result[k] = df[idxlo,k] * wlo +  df[idxhi,k] * whi
  end
  result[:jd] = jd
  return result
end

function get_horizons_output(jd::Vector{Float64}, obs::Symbol)
  df = read_horizons_output(jd, obs)
  idxlo = map(t->searchsortedlast(df.jd, t),jd)
  @assert all(1 .<= idxlo .< size(df,1))
  result = DataFrame(:jd=>jd)
  whi = (jd.-df.jd[idxlo]) ./ (df.jd[idxlo.+1].-df.jd[idxlo])
  #wlo = 1.-whi
  for k in  Symbol.(names(df))
	if k == :jd continue end
  	result[!,k] = df[idxlo,k] .* (1.0.-whi) .+  df[idxlo.+1,k] .* whi
  end
  return result
end

function calc_Δv_diff_extinction_astropy(;epochs::Vector{Float64}, obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	#=if typeof(epochs) <: Dict
		@assert haskey(epochs,"start")
		@assert haskey(epochs,"stop")
		@assert haskey(epochs,"step")
	else=#
	if typeof(epochs) <: Vector{Float64}
		@assert 1 <= length(epochs) <= 10000
	end
	@assert 1 < max_airmass <= 10.0
	sol = SolarRotation.get_solar_info(epochs, obs=obs)
	ha = sol.lst .- sol.ra./15
	idx_visible = 0 .< sol.airmass .<= max_airmass
	horizons = get_horizons_output(epochs[idx_visible],obs)
	#delta = SolarRotation.get_earth_sun_dist.(epochs[idx_visible], obs=obs)
	delta = sun.sol_dist_au
	loc = get_obs_loc(obs)
	#println(" jd = ", epochs, " delta_h = ", horizons.sol_dist_au, " delta_ap = ", delta)
	#println(" diff = ", horizons.sol_dist_au.-delta)
	delta_vr = diff_extinct_corr.(epochs[idx_visible], sol.alt[idx_visible], ha[idx_visible], sol.ra[idx_visible], horizons.sol_npole_ang, horizons.sol_dist_au.*AU, loc["lat"]) #eph["NPole_ang"], eph["delta"])
	return delta_vr
end

function calc_Δv_diff_extinction_astropy(epoch::Float64; obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	@assert 1 < max_airmass <= 10.0
	sol = SolarRotation.get_solar_info(epoch, obs=obs)
	ha = sol[:lst] .- sol[:ra]./15
	idx_visible = 0 < sol[:airmass] <= max_airmass
	if !idx_visible return (jd=epoch, Δv = 0.0) end
	horizons = get_horizons_output(epoch,obs)
	#delta = SolarRotation.get_earth_sun_dist(epoch, obs=obs)
	loc = get_obs_loc(obs)
	delta_vr = diff_extinct_corr(epoch, sol[:alt], ha, sol[:ra], horizons[:sol_npole_ang], sol[:sol_dist_au].*AU, loc["lat"]) #eph["NPole_ang"], eph["delta"])
	return delta_vr
end

function calc_Δv_diff_extinction(epoch::Float64; obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	calc_Δv_diff_extinction_astropy(epoch; obs=obs, max_airmass=max_airmass, verbose=verbose)
end

function compare_Horizons_Astropy(;epochs::Vector{Float64}, obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	eph = query_horizons(epochs=epochs,obs=obs,max_airmass=max_airmass, verbose=verbose)
	delta_vr = diff_extinct_corr.(eph.columns["datetime_jd"], eph.columns["EL"],
	 	eph.columns["hour_angle"], eph.columns["RA"],
		eph.columns["NPole_ang"], eph.columns["delta"].*AU, loc["lat"])

	sol = SolarRotation.get_solar_info(epochs, obs=obs)
	ha = sol.lst .- sol.ra./15
	idx_visible = 0 .< sol.airmass .<= max_airmass
	delta = SolarRotation.get_earth_sun_dist.(epochs[idx_visible], obs=obs)

	println("Horizons airmass =", eph.columns["airmass"])
	println("Astropy  airmass =", sol.airmass[idx_visible])
	println("Δjd = ", eph.columns["datetime_jd"], " vs ", epochs[idx_visible], " = ", eph.columns["datetime_jd"].-epochs[idx_visible])
	println("Δalt = ", eph.columns["EL"].- sol.alt[idx_visible])
	println("Δha = ", eph.columns["hour_angle"].- ha[idx_visible])
	println("ΔRA = ", eph.columns["RA"].- sol.ra[idx_visible])

	horizons = get_horizons_output(epochs[idx_visible],obs)
	delta_vr_new = diff_extinct_corr.(epochs[idx_visible], sol.alt[idx_visible], ha[idx_visible], sol.ra[idx_visible], zeros(length(epochs[idx_visible])), fill(AU,length(epochs[idx_visible])), loc["lat"]) #eph["NPole_ang"], eph["delta"])

	println("delta_vr_old = ", delta_vr, "\ndelta_vr_new = ", delta_vr_new)
	println("delta = ", delta_vr.- delta_vr_new)
	return (eph=eph, sol=sol, delta_vr_old=delta_vr, delta_vr_new=delta_vr_nw)
end

function calc_Δv_diff_extinction(bjd::Real; obs::Symbol, max_airmass::Float64 = 10.0, verbose::Bool = false)
	epochs = [bjd]
	(bjd_out, Δv) = calc_Δv_diff_extinction(epochs=epochs, obs=obs, max_airmass=max_airmass, verbose=verbose)
	return first(Δv)
end

function calc_Δv_diff_extinction(start_date::String, stop_date::String; obs::Symbol, dt::Integer = 1, dt_unit::String = "m", max_airmass::Float64 = 10.0, verbose::Bool = false)
	@assert 1 <= dt <= 24*60
	@assert dt_unit ∈ ["m", "h", "d"]
	@assert 1 < max_airmass <= 10.0
	epochs = Dict("start"=>start_date, "stop"=>stop_date, "step"=>string(dt) * dt_unit)
	(bjd_out, Δv) = calc_Δv_diff_extinction(epochs=epochs, obs=obs, max_airmass=max_airmass, verbose=verbose)
	return (bjd=bjd_out, Δv=Δv)
end

function calc_Δv_diff_extinction(start_bjd::Real, stop_bjd::Real; obs::Symbol, dt::Integer = 1, dt_unit::String = "m", max_airmass::Float64 = 10.0, verbose::Bool=false)
	start_date = string(Date(Dates.julian2datetime(start_bjd)))
	stop_date = string(Date(Dates.julian2datetime(stop_bjd)))
	calc_Δv_diff_extinction(start_date, stop_date, obs=obs, dt=dt, dt_unit=dt_unit, max_airmass=max_airmass, verbose=verbose)
end

function calc_Δv_diff_extinction_demo(; verbose::Bool = false)
	start_date = "2021-01-01"
	stop_date = "2021-02-11"
	calc_Δv_diff_extinction(start_date, stop_date, obs=:WIYN, verbose=verbose)
end

export calc_Δv_diff_extinction

end # module DifferentialExtinction


import Base: copy

# -------------
# Abstract type
# -------------

@doc """
# Interface

where Q is a subtype of AbstractQFunction

* policy(Q) -> a vector of the (currently) optimal policy for all states
* policy(Q, s) -> the (currently) optimal action to take at state s
* value(Q) -> a vector of the currently optimal values for all states
* value(Q, s) -> the value expected from state s given (currently) optimal actions are taken
* setvalue!(Q, v, s, a) -> set the value for state action
* num_states(Q) -> the number of states

## Optional

* getvalue(Q, s, a) -> the value expected from state s given action a is taken
* num_actions(Q) -> the number of actions

""" ->
abstract AbstractQFunction

# ----------
# Array type
# ----------

type ArrayQFunction{T<:Real} <: AbstractQFunction
    array::Array{T,2}
end

ArrayQFunction{T}(::Type{T}, states, actions) =
    ArrayQFunction(zeros(T, states, actions))

ArrayQFunction(states, actions) = ArrayQFunction(Float64, states, actions)

QFunction(A::Matrix) = ArrayQFunction(A)

# Base methods
# ------------

==(a::ArrayQFunction, b::ArrayQFunction) = a.array == b.array

isequal(a::ArrayQFunction, b::ArrayQFunction) = isequal(a.array, b.array)

copy(Q::ArrayQFunction) = ArrayQFunction(copy(Q.array))

# Size methods
# ------------

num_actions(Q::ArrayQFunction) = size(Q.array, 2)

num_states(Q::ArrayQFunction) = size(Q.array, 1)

# Value methods
# -------------

value(Q::ArrayQFunction, state) = maximum(Q.array[state, :])

function value!(A::Vector, Q::ArrayQFunction)
    for s = 1:num_states(Q)
        A[s] = value(Q, s)
    end
    return A
end

value{T}(Q::ArrayQFunction{T}) = value!(Array(T, num_states(Q)), Q)

getvalue(Q::ArrayQFunction, state, action) = getindex(Q.array, state, action)

setvalue!(Q::ArrayQFunction, v, state, action) =
    (setindex!(Q.array, v, state, action); Q)

valuetype{T}(Q::ArrayQFunction{T}) = T

# Policy methods
# --------------

policy(Q::ArrayQFunction, state) = indmax(Q.array[state, :])

function policy!(A::Vector, Q::ArrayQFunction)
    for s = 1:num_states(Q)
        A[s] = policy(Q, s)
    end
    return A
end

policy{T}(::Type{T}, Q::ArrayQFunction) = policy!(Array(T, num_states(Q)), Q)

policy(Q::ArrayQFunction) = policy(Int, Q)

# -----------
# Vector type
# -----------

type VectorQFunction{V<:Real,A<:Integer} <: AbstractQFunction
    value::Vector{V}
    policy::Vector{A}

    function VectorQFunction(value::Vector{V}, policy::Vector{A})
        length(value) == length(policy) || throw(
            DimensionMismatch("The lengths of value and policy must match.")
        )
        new(value, policy)
    end
end

VectorQFunction{V,A}(value::Vector{V}, policy::Vector{A}) =
    VectorQFunction{V,A}(value, policy)

VectorQFunction{V,A}(::Type{V}, ::Type{A}, states) =
    VectorQFunction(zeros(V, states), zeros(A, states))

VectorQFunction(states) = VectorQFunction(Float64, Int, states)

QFunction(A::Vector, B::Vector) = VectorQFunction(A, B)

# Base methods
# ------------

==(a::VectorQFunction, b::VectorQFunction) =
    (a.value == b.value) & (a.policy == b.policy)

isequal(a::ArrayQFunction, b::ArrayQFunction) =
    isequal(a.value, b.value) & isequal(a.policy, b.policy)

copy(Q::VectorQFunction) = VectorQFunction(copy(Q.value), copy(Q.policy))

# Size methods
# ------------

# VectorQFunction can not calculate the number of actions so cannot define
# num_actions(Q::VectorQFunction)

num_states(Q::VectorQFunction) = length(Q.value)

# Value methods
# -------------

value!(A::Vector, Q::VectorQFunction) = copy!(A, Q.value)

value(Q::VectorQFunction) = Q.value

value(Q::VectorQFunction, state) = getindex(Q.value, state)

@doc """
Returns negative infinity if `a` is not the optimal action.
""" ->
getvalue{V,A}(Q::VectorQFunction{V,A}, s, a) = Q.policy[s] == a ? Q.value[s] : convert(V, -Inf)

function setvalue!(Q::VectorQFunction, v, state, action)
    if v > value(Q, state)
        setindex!(Q.value, v, state)
        setindex!(Q.policy, action, state)
    end
    return Q
end

valuetype{V,A}(Q::VectorQFunction{V,A}) = V

# Policy methods
# --------------

policy!(A::Vector, Q::VectorQFunction) = copy!(A, Q.policy)

policy(Q::VectorQFunction) = Q.policy

policy(Q::VectorQFunction, state) = getindex(Q.policy, state)


# Testing discrete univariate distributions

using Distributions
import JSON
using Base.Test

function verify_and_test_drive(jsonfile, n_tsamples::Int)
    R = JSON.parsefile(jsonfile)
    for (ex, dct) in R
        println("    testing $(ex)")

        # check type
        dtype = eval(symbol(dct["dtype"]))
        d = eval(parse(ex))
        if dtype == TruncatedNormal
            @test isa(d, Truncated{Normal})
        else
            @assert isa(dtype, Type) && dtype <: UnivariateDistribution
            @test isa(d, dtype)
        end

        # verification and testing
        verify_and_test(d, dct, n_tsamples)
    end
end

_parse_x(d::DiscreteUnivariateDistribution, x) = int(x)
_parse_x(d::ContinuousUnivariateDistribution, x) = float64(x)

function verify_and_test(d::UnivariateDistribution, dct::Dict, n_tsamples::Int)
    # verify parameters
    pdct = dct["params"]
    for (fname, val) in pdct
        f = eval(symbol(fname))
        @assert isa(f, Function)
        Base.Test.test_approx_eq(f(d), val, "$fname(d)", "val")
    end

    # verify stats
    r_min = float64(dct["minimum"])
    r_max = float64(dct["maximum"])
    @assert r_min < r_max

    @test_approx_eq minimum(d) r_min
    @test_approx_eq maximum(d) r_max
    @test_approx_eq_eps mean(d) float64(dct["mean"]) 1.0e-8
    @test_approx_eq_eps var(d) float64(dct["var"]) 1.0e-8
    @test_approx_eq_eps median(d) dct["median"] 1.0

    if applicable(entropy, d)
        @test_approx_eq_eps entropy(d) dct["entropy"] 1.0e-7
    end

    # verify quantiles
    @test_approx_eq_eps quantile(d, 0.10) dct["q10"] 1.0e-8
    @test_approx_eq_eps quantile(d, 0.25) dct["q25"] 1.0e-8
    @test_approx_eq_eps quantile(d, 0.50) dct["q50"] 1.0e-8
    @test_approx_eq_eps quantile(d, 0.75) dct["q75"] 1.0e-8
    @test_approx_eq_eps quantile(d, 0.90) dct["q90"] 1.0e-8

    # verify logpdf and cdf at certain points
    pts = dct["points"]
    for pt in pts
        x = _parse_x(d, pt["x"])
        lp = float64(pt["logpdf"])
        cf = float64(pt["cdf"])
        Base.Test.test_approx_eq(logpdf(d, x), lp, "logpdf(d, $x)", "lp")
        Base.Test.test_approx_eq(cdf(d, x), cf, "cdf(d, $x)", "cf")
    end

    # generic testing
    test_distr(d, n_tsamples)
end

## main

for c in ["discrete",
          "continuous"]

    title = string(uppercase(c[1]), c[2:end])
    println("    [$title]")
    println("    ------------")
    jsonfile = joinpath(dirname(@__FILE__), "$(c)_test.json")
    verify_and_test_drive(jsonfile, 10^6)
    println()
end


"""
$(TYPEDEF)

Represents a difference operator.

# Fields
$(FIELDS)

# Examples

```jldoctest
julia> using Symbolics

julia> @variables t;

julia> D = Difference(t; dt=0.01)
Difference(t; dt=0.01)
```
"""
struct Difference <: Function
    """Fixed Difference"""
    t
    dt
    Difference(t; dt) = new(value(t), dt)
end
(D::Difference)(t) = Term{symtype(t)}(D, [t])
(D::Difference)(t::Num) = Num(D(value(t)))
SymbolicUtils.promote_symtype(::Difference, t) = t

Base.show(io::IO, D::Difference) = print(io, "Difference(", D.t, "; dt=", D.dt, ")")

Base.:(==)(D1::Difference, D2::Difference) = isequal(D1.t, D2.t) && isequal(D1.dt, D2.dt)
Base.isequal(D1::Difference, D2::Difference) = isequal(D1.t, D2.t) && isequal(D1.dt, D2.dt)
Base.hash(D::Difference, u::UInt) = hash(D.dt, hash(D.t, xor(u, 0x055640d6d952f101)))


function assemble_shortest_paths(g::GenericGraph, s, dist_key; distance::Bool=false)

    nodes = [k.index => k for k = vertices(g)]

    Source_nodes = Int64[]
    Predecessors = [k.index => Int64[] for k = vertices(g)] # vymenene {} za []
    sigma = [vertex.index => 0.0 for vertex = vertices(g)]
    Distances = Dict{Int64, Float64}()
    sigma[s] = 1.0
    visited = Dict(s => 0)
    c = 0
    Q = (Array{Any, 1})[]
    Collections.heappush!(Q, [0, c += 1 , s, s],Base.Order.ord((x,y)->isless(x[1], y[1]), identity, false))

    while length(Q) > 0

        heap_element = Collections.heappop!(Q, Base.Order.ord((x,y)->isless(x[1], y[1]), identity, false))
        dist, _, pred, vertex = heap_element[1] , heap_element[2] , heap_element[3] , heap_element[4]

        if haskey(Distances, vertex)
            continue
        end

        sigma[vertex] += sigma[pred]
        push!(Source_nodes,vertex)
        Distances[vertex] = dist

        for edge in out_edges(nodes[vertex], g)

            vw_dist = dist + edge.attributes[dist_key]

            if (haskey(Distances, edge.target.index) == false) && ( haskey(visited, edge.target.index) == false || vw_dist < visited[edge.target.index] )

                visited[edge.target.index] = vw_dist
                Collections.heappush!(Q, [vw_dist, c += 1, vertex, edge.target.index],Base.Order.ord((x,y)->isless(x[1], y[1]), identity, false))
                sigma[edge.target.index] = 0.0
                Predecessors[edge.target.index] = [vertex]

            elseif vw_dist == visited[edge.target.index]

                sigma[edge.target.index] += sigma[vertex]
                push!(Predecessors[edge.target.index],vertex)

            end

        end

    end

    if distance
      return Source_nodes, Predecessors, Distances
    elseif !(distance)
      return Source_nodes, Predecessors, sigma
    end

end


using SparseArrays
using Random

# split from runtests.jl for docs generating script (gen_imgs.jl)

# sprand or sprandn is not stable across versions (e.g. 1.0 vs 1.6)
function _stable_sprand(r, m::Integer, n::Integer, density::AbstractFloat)
    I = Int[]
    J = Int[]
    for li in randsubseq(r, 1:(m*n), density)
        j, i = divrem(li - 1, m)
        push!(I, i + 1)
        push!(J, j + 1)
    end
    V = rand(r, length(I))
    sparse(I, J, V)
end


# interval.jl

iscompact(d::TypedEndpointsInterval{:closed,:closed}) = true
iscompact(d::TypedEndpointsInterval) = false

##################
### An interval
##################

approx_indomain(x, d::AbstractInterval, tolerance) =
    (x <= rightendpoint(d)+tolerance) && (x >= leftendpoint(d)-tolerance)

function point_in_domain(d::AbstractInterval)
    isempty(d) && throw(BoundsError())
    mean(d)
end

isapprox(d1::AbstractInterval, d2::AbstractInterval; kwds...) =
    isapprox(leftendpoint(d1), leftendpoint(d2); kwds...) &&
    isapprox(rightendpoint(d1), rightendpoint(d2); kwds...)

boundary(d::TypedEndpointsInterval{:closed,:closed}) = Point(leftendpoint(d)) ∪ Point(rightendpoint(d))
boundary(d::TypedEndpointsInterval{:closed,:open}) = Point(leftendpoint(d))
boundary(d::TypedEndpointsInterval{:open,:closed}) = Point(rightendpoint(d))
boundary(d::TypedEndpointsInterval{:open,:open,T}) where T = EmptySpace{T}()
boundary(d::AbstractInterval) = boundary(Interval(d))

# We extend some functionality of intervals to mapped intervals
const MappedInterval{D <: AbstractInterval,T} = MappedDomain{D,T}

endpoints(d::MappedInterval) = forward_map(d) .* endpoints(source(d))
for op in (:leftendpoint, :rightendpoint)
    @eval $op(d::MappedInterval) = forward_map(d) * $op(source(d))
end

for op in (:isleftopen, :isrightopen, :isleftclosed, :isrightclosed)
    @eval $op(d::MappedInterval) = $op(source(d))
end

## Some special intervals follow, e.g.:
# - the unit interval [0,1]
# - the 'Chebyshev' interval [-1,1]
# - ...
# Unlike a generic interval, these specific intervals have no data, and only one
# type parameter (T).

"""
The abstract type `FixedInterval` is the supertype of intervals with endpoints
determined by the type, rather than field values. Examples include `UnitInterval`
and `ChebyshevInterval`.
"""
abstract type FixedInterval{L,R,T} <: TypedEndpointsInterval{L,R,T} end
const ClosedFixedInterval{T} = FixedInterval{:closed,:closed,T}

"""
Return an interval that is similar to the given interval, but with endpoints
`a` and `b` instead.
"""# Assume a closed interval by default
similar_interval(d::ClosedFixedInterval{T}, a, b) where {T} = ClosedInterval{float(T)}(a, b)

"The closed unit interval [0,1]."
struct UnitInterval{T} <: ClosedFixedInterval{T} end

UnitInterval() = UnitInterval{Float64}()

endpoints(d::UnitInterval{T}) where {T} = (zero(T), one(T))

"The closed interval [-1,1]."
struct ChebyshevInterval{T} <: ClosedFixedInterval{T}
end

ChebyshevInterval() = ChebyshevInterval{Float64}()

endpoints(d::ChebyshevInterval{T}) where {T} = (-one(T),one(T))

"The half-open positive halfline `[0,∞)`."
struct HalfLine{T} <: FixedInterval{:closed,:open,T} end
HalfLine() = HalfLine{Float64}()

endpoints(d::HalfLine{T}) where {T} = (zero(T), T(Inf))

indomain(x, d::HalfLine) = x >= 0

function similar_interval(d::HalfLine, a, b)
    @assert a == 0
    @assert isinf(b) && b > 0
    d
end

point_in_domain(d::HalfLine) = zero(eltype(d))

"The open negative halfline `(-∞,0)`."
struct NegativeHalfLine{T} <: FixedInterval{:open,:open,T} end
NegativeHalfLine() = NegativeHalfLine{Float64}()

endpoints(d::NegativeHalfLine{T}) where {T} = (-T(Inf), zero(T))

# Open at both endpoints

indomain(x, d::NegativeHalfLine) = x < 0

function similar_interval(d::NegativeHalfLine, a, b)
    @assert isinf(a) && a < 0
    @assert b == 0
    d
end

point_in_domain(d::NegativeHalfLine) = -one(eltype(d))

similar_interval(d::Interval{L,R,T}, a, b) where {L,R,T} =
    Interval{L,R,float(T)}(a, b)

#########################################
# A few set operations with known result
#########################################

# We define an exhaustive list of combinations of the four fixed intervals
# combined above in the routines 'intersect', 'union' and 'setdiff' where
# the output is known explicitly.

# Since fixed intervals are fully determined by their type,
# the result of intersect, union or setdiff is always known for two
# domains of the same type.
intersect(d1::D, d2::D) where {D <: FixedInterval} = d1
union(d1::D, d2::D) where {D <: FixedInterval} = d1
setdiff(d1::D, d2::D) where {D <: FixedInterval} = EmptySpace{eltype(D)}()

# [0,1] ∩ [-1,1] = [0,1]
intersect(d1::UnitInterval{T}, d2::ChebyshevInterval{T}) where {T} = UnitInterval{T}()
intersect(d1::ChebyshevInterval{T}, d2::UnitInterval{T}) where {T} = UnitInterval{T}()
# [0,1] ∩ [0,∞) = [0,1]
intersect(d1::UnitInterval{T}, d2::HalfLine{T}) where {T} = UnitInterval{T}()
intersect(d1::HalfLine{T}, d2::UnitInterval{T}) where {T} = UnitInterval{T}()
# [0,1] ∩ (-∞,0) = {}
intersect(d1::UnitInterval{T}, d2::NegativeHalfLine{T}) where {T} = EmptySpace{T}()
intersect(d1::NegativeHalfLine{T}, d2::UnitInterval{T}) where {T} = EmptySpace{T}()
# [-1,1] ∩ [0,∞) = [0,1]
intersect(d1::ChebyshevInterval{T}, d2::HalfLine{T}) where {T} = UnitInterval{T}()
intersect(d1::HalfLine{T}, d2::ChebyshevInterval{T}) where {T} = UnitInterval{T}()
# [0,∞) ∩ (-∞,0) = {}
intersect(d1::HalfLine{T}, d2::NegativeHalfLine{T}) where {T} = EmptySpace{T}()
intersect(d1::NegativeHalfLine{T}, d2::HalfLine{T}) where {T} = EmptySpace{T}()

# [0,1] ∪ [-1,1] = [-1,1]
union(d1::UnitInterval{T}, d2::ChebyshevInterval{T}) where {T} = ChebyshevInterval{T}()
union(d1::ChebyshevInterval{T}, d2::UnitInterval{T}) where {T} = ChebyshevInterval{T}()
# [0,1] ∪ [0,∞) = [0,∞)
union(d1::UnitInterval{T}, d2::HalfLine{T}) where {T} = HalfLine{T}()
union(d1::HalfLine{T}, d2::UnitInterval{T}) where {T} = HalfLine{T}()

# (-∞,0) ∪ [0,∞) = (-∞,∞)
# Note: T<:real to ensure that FullSpace{T} is not larger than intended.
union(d1::NegativeHalfLine{T}, d2::HalfLine{T}) where {T<:Real} = FullSpace{T}()
union(d1::HalfLine{T}, d2::NegativeHalfLine{T}) where {T<:Real} = FullSpace{T}()

# [0,1] ∖ [-1,1] = {}
setdiff(d1::UnitInterval{T}, d2::ChebyshevInterval{T}) where {T} = EmptySpace{T}()
# [0,1] ∖ [0,∞) = {}
setdiff(d1::UnitInterval{T}, d2::HalfLine{T}) where {T} = EmptySpace{T}()
# [0,1] ∖ (-∞,0) = [0,1]
setdiff(d1::UnitInterval{T}, d2::NegativeHalfLine{T}) where {T} = UnitInterval{T}()
# [-1,1] ∖ (-∞,0) = [0,1]
setdiff(d1::ChebyshevInterval{T}, d2::NegativeHalfLine{T}) where {T} = UnitInterval{T}()
# [0,∞) ∖ (-∞,0) = [0,∞)
setdiff(d1::HalfLine{T}, d2::NegativeHalfLine{T}) where {T} = HalfLine{T}()
# (-∞,0) ∖ [0,1] = (-∞,0)
setdiff(d1::NegativeHalfLine{T}, d2::UnitInterval{T}) where {T} = NegativeHalfLine{T}()
# (-∞,0) ∖ [0,∞) = (-∞,0)
setdiff(d1::NegativeHalfLine{T}, d2::HalfLine{T}) where {T} = NegativeHalfLine{T}()

#################################
# Conversions between intervals
#################################
# Avoid depcrecated warning: Warning: Constructors no longer fall back to `convert`.
# example: A constructor `AbstractInterval{Float64}(::IntervalSets.ClosedInterval{Float64})` should be defined instead.

function convert(::Type{UnitInterval{T}}, d::AbstractInterval) where {T}
    endpoints(d) == (0,1) || throw(InexactError(:convert,UnitInterval,d))
    UnitInterval{T}()
end

function convert(::Type{ChebyshevInterval{T}}, d::AbstractInterval) where {T}
    endpoints(d) == (-1,1)|| throw(InexactError(:convert,ChebyshevInterval,d))
    ChebyshevInterval{T}()
end

UnitInterval{T}(d::AbstractInterval) where T = convert(UnitInterval{T}, d)
ChebyshevInterval{T}(d::AbstractInterval) where T = convert(ChebyshevInterval{T}, d)

########################
# Arithmetic operations
########################

# Some computations with intervals simplify without having to use a mapped domain.
# This is only the case for Interval{L,R,T}, and not for any of the FixedIntervals
# because the endpoints of the latter are, well, fixed.

-(d::ChebyshevInterval) = d
-(d::AbstractInterval) = similar_interval(d, -rightendpoint(d), -leftendpoint(d))

for op in (:+, :-), Inter in (:AbstractInterval, :ClosedInterval)
    @eval $op(d::$Inter, x::Real) = similar_interval(d, $op(leftendpoint(d),x), $op(rightendpoint(d),x))
end

for Inter in (:AbstractInterval, :ClosedInterval)
    @eval begin
        +(x::Real, d::$Inter) = similar_interval(d, x+leftendpoint(d), x+rightendpoint(d))
        -(x::Real, d::$Inter) = similar_interval(d, x-rightendpoint(d), x-leftendpoint(d))
    end
end

for op in (:*, :/)
    @eval function $op(d::AbstractInterval, x::Real)
        if x ≥ 0 # -{x : 0 ≤ x ≤ 1} should be {x : -1 ≤ x ≤ 0}, not empty set {x : 0 ≤ x ≤ -1}
            similar_interval(d, $op(leftendpoint(d),x), $op(rightendpoint(d),x))
        else
            similar_interval(d, $op(rightendpoint(d),x), $op(leftendpoint(d),x))
        end
    end
end

for op in (:*, :\), Inter in (:AbstractInterval, :ClosedInterval)
    @eval function $op(x::Real, d::$Inter)
        if x ≥ 0 # -{x : 0 ≤ x ≤ 1} should be {x : -1 ≤ x ≤ 0}, not empty set {x : 0 ≤ x ≤ -1}
            similar_interval(d, $op(x,leftendpoint(d)), $op(x,rightendpoint(d)))
        else
            similar_interval(d, $op(x,rightendpoint(d)), $op(x,leftendpoint(d)))
        end
    end
end

function show(io::IO, d::ChebyshevInterval)
    print(io, Interval(d))
    print(io, " (Chebyshev)")
end
function show(io::IO, d::UnitInterval)
    print(io, Interval(d))
    print(io, " (Unit)")
end
function show(io::IO, d::HalfLine)
    print(io, Interval(d))
    print(io, " (HalfLine)")
end

function setdiff(d1::AbstractInterval{T}, d2::AbstractInterval{T}) where T
    a1 = leftendpoint(d1)
    b1 = rightendpoint(d1)
    a2 = leftendpoint(d2)
    b2 = rightendpoint(d2)

    isempty(d1) && return d1
    isempty(d2) && return d1
    b1 < a2 && return d1
    a1 < a2 ≤ b1 ≤ b2 && return (a1 .. a2)
    a1 < a2 ≤ b2 < b1 && return UnionDomain(a1 .. a2) ∪ UnionDomain(b2 .. b1)
    a2 ≤ a1 < b2 < b1 && return (b2 .. b1)
    a2 ≤ a1 ≤ b1 ≤ b2 && return EmptySpace{T}()

    @assert b2 ≤ a1
    d1
end

setdiff(d1::AbstractInterval, d2::AbstractInterval) = setdiff(promote(d1,d2)...)

function *(map::AffineMap, domain::AbstractInterval)
    le = map*leftendpoint(domain)
    re = map*rightendpoint(domain)
    if le<re
        similar_interval(domain,le,re)
    else
        similar_interval(domain,re,le)
    end
end


module CellularAutomata

import Base.show

include("1dim.jl")
include("2dim.jl")

export CellularAutomaton,
       CA2d,
       rule,
       show

#Print CA to screen
function show(io::IO, ca::CellularAutomaton)
    io2 = IOBuffer()

    arr = ca.cells
    h, w = size(arr)

    print(io2, "$(w)x$(h) Cellular Automaton")
    print(io2, "\n")

    #TODO: print rule

    for i = 1:h
        print(io2,"\t")
        for j = 1:w
            if arr[i,j] == 0
                print(io2," ")
            elseif ca.k == 2 && arr[i,j] == 1
                print(io2,"#")
            elseif ca.k == 3 && arr[i,j] == 1
                print(io2,"X")
            elseif arr[i,j] == 2
                print(io2, "#")
            end
        end
        print(io2,"\n")
    end

    print(io, bytestring(io2))
end

#Print CA to screen
function show(io::IO, ca::CA2d)
    io2 = IOBuffer()

    arr = ca.cells
    h, w, gen = size(arr)

    print(io2, "$(w)x$(h)x$(gen) Cellular Automaton")
    print(io2, "\n")

    #TODO: print rule

    for g in 1:gen
        for i = 1:h
            print(io2,"\t")
            for j = 1:w
                if arr[i,j, g] == 0
                    print(io2," ")
                elseif arr[i,j, g] == 1
                    print(io2, "#")
                end
            end
            print(io2,"\n")
        end
        print(io2,"\n\n")
    end

    print(io, bytestring(io2))
end

#end of module
end


function MixtureModel(hmm::HMM)
    sdists = hmm.A[1, :]
    MixtureModel([hmm.B...], sdists)
end

function HMM(m::MixtureModel)
    K = ncomponents(m)
    a = probs(m)
    A = repeat(permutedims(m.prior.p), K, 1)
    B = m.components
    HMM(a, A, B)
end
#
# function PeriodicHMM(vec_mix::Vector{MixtureModel})
#     T = length(vec_mix)
#     K = ncomponents(vec_mix[1])
#     a = probs(vec_mix[1])
#     A = repeat.(permutedims.(probs.(vec_mix)), K, 1)
#     B = components.(vec_mix)
#     PeriodicHMM(a, [A[t][k,l] for k in 1:K, l in 1:K, t in 1:T], [B[t][k] for k in 1:K, t in 1:T])
# end

function copy(mix::MixtureModel)
    MixtureModel(components(mix), probs(mix))
end

function fit_mle(mix::MixtureModel, observations; init = :none, kwargs...)
    mix = copy(mix)

    if init == :kmeans
        kmeans_init!(mix, observations, display = get(kwargs, :display, :none))
    end
    fit_em(mix, observations; kwargs...)
end

function update_a!(
    a::AbstractVector,
    A::AbstractMatrix
)
    a[:] = A[1, :]
end

function update_A!(
    A::AbstractMatrix,
    γ::AbstractMatrix
)
    K = size(A, 1)
    A[:, :] = repeat(mean(γ, dims = 1), K, 1)
end

function posteriors!(
    γ::AbstractMatrix,
    c::AbstractVector,
    A::AbstractMatrix,
    LL::AbstractMatrix
)
    K = size(A, 1)
    a = similar(LL)
    α = log.(A[1, :])
    [a[:, k] = LL[:, k] .+ α[k] for k = 1:K]
    c[:] = logsumexp(a, dims = 2)
    γ[:, :] = exp.(a .- c)
end

function fit_em(mix::MixtureModel, y; display = :none, maxiter = 100, tol = 1e-3, robust = false)

    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K = length(y), ncomponents(mix)
    history = EMHistory(false, 0, [])

    # Allocate memory for in-place updates

    LL = zeros(N, K)
    γ = similar(LL)
    c = zeros(N)
    types = typeof.(components(mix))

    # Initial parameters
    α = copy(probs(mix))
    dists = copy(components(mix))

    # E-step
    # evaluate likelihood for each type k
    [LL[:, k] = log(α[k]) .+ logpdf.(dists[k], y) for k = 1:K]
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    # get posterior of each category
    c[:] = logsumexp(LL, dims = 2)
    γ[:, :] = exp.(LL .- c)

    # Loglikelihood
    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter

        # M-step
        # with γ in hand, maximize (update) the parameters
        α[:] = mean(γ, dims = 1)
        dists[:] = [fit_mle(types[k], y, γ[:, k]) for k = 1:K]

        # E-step
        # evaluate likelihood for each type k
        [LL[:, k] = log(α[k]) .+ logpdf.(dists[k], y) for k = 1:K]
        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
        # get posterior of each category
        c[:] = logsumexp(LL, dims = 2)
        γ[:, :] = exp.(LL .- c)

        # Loglikelihood
        logtotp = sum(c)
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    return MixtureModel(dists, α), history
end

# function fit_em_GE(mix::MixtureModel, y, M1; display = :none, maxiter = 100, tol = 1e-3, robust = false)
#
# 	@argcheck display in [:none, :iter, :final]
#     @argcheck maxiter >= 0
#
# 	N, K = length(y), ncomponents(mix)
# 	history = EMHistory(false, 0, [])
#
# 	# Allocate memory for in-place updates
#
#     LL = zeros(N,K)
#     γ = similar(LL)
# 	c = zeros(N)
# 	types = typeof.(components(mix))
# 	α = copy(probs(mix))
# 	dists = copy(components(mix))
#
# 	# evaluate likelihood for each type k
# 	[LL[:,k] = log(α[k]).+logpdf.(dists[k], y) for k in 1:K]
# 	robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
#
# 	# get posterior of each type
# 	c[:] = logsumexp(LL, dims = 2)
# 	γ[:,:] =  exp.(LL .- c)
#
# 	logtot = sum(c)
# 	(display == :iter) && println("Iteration 0: logtot = $logtot")
#
# 	for it = 1:maxiter
#
# 		# with p in hand, update
# 		α[:] = mean(γ, dims=1)
# 		θ = [mean(y, Weights(γ[:,k])) for k in 1:K]
# 		alpha = (M1 - θ[2]α[2])/(α[1]θ[1])
# 		dists[1] = Gamma(alpha, θ[1])
# 		dists[2] = Exponential(θ[2])
#
#         # evaluate likelihood for each type k
# 		[LL[:,k] = log(α[k]) .+ logpdf.(dists[k], y) for k in 1:K]
# 		robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
#
# 		# get posterior of each type
#         c[:] = logsumexp(LL, dims = 2)
#         γ[:,:] =  exp.(LL .- c)
#
#         # Log likelihood
# 		logtotp = sum(c)
# 		(display == :iter) && println("Iteration $it: logtot = $logtotp")
#
# 		push!(history.logtots, logtotp)
#         history.iterations += 1
#
# 		if abs(logtotp - logtot) < tol
#             (display in [:iter, :final]) &&
#                 println("EM converged in $it iterations, logtot = $logtotp")
#             history.converged = true
#             break
#         end
#
#         logtot = logtotp
#     end
#
#     return MixtureModel(dists, α), history
# end

# same as fit_em but a tad slower (howeve)
function fit_mle_mixture!(
    hmm::HMM,
    observations,
    ;
    display = :none,
    maxiter = 100,
    tol = 1e-3,
    robust = false,
    estimator = fit_mle
)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K = size(observations, 1), size(hmm, 1)
    history = EMHistory(false, 0, [])

    # Allocate memory for in-place updates

    c = zeros(N)
    γ = zeros(N, K)
    LL = zeros(N, K)

    loglikelihoods!(LL, hmm, observations)
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

    posteriors!(γ, c, hmm.A, LL)

    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter

        update_A!(hmm.A, γ)
        update_B!(hmm.B, γ, observations, estimator)

        # Ensure the "connected-ness" of the states,
        # this prevents case where there is no transitions
        # between two extremely likely observations.
        robust && (hmm.A .+= eps())

        @check istransmat(hmm.A)

        loglikelihoods!(LL, hmm, observations)
        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

        posteriors!(γ, c, hmm.A, LL)

        logtotp = sum(c)
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end
    update_a!(hmm.a, hmm.A)

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end
    history
end

## * Parametric model: Trigonometric function

function model_for_dist(dist::Exponential, θ_Y::AbstractArray, γ::AbstractVector, y, n2t, T::Int; silence = true)
    d = (size(θ_Y, 1) - 1) ÷ 2
    @argcheck length(y) == length(n2t)
    K = size(γ, 2)
    N = length(n2t)
    model = Model(Ipopt.Optimizer)
    silence && set_silent(model)
    f = 2π / T

    cos_nj = [cos(f * j * n2t[n]) for n = 1:N, j = 1:d]
    sin_nj = [sin(f * j * n2t[n]) for n = 1:N, j = 1:d]

    trig = [[1; interleave2(cos_nj[n, :], sin_nj[n, :])] for n = 1:N]

    @variable(model, θ_jump[j = 1:(2d+1)])
    # Polynomial P
    @NLexpression(model, Pθ[n = 1:N], sum(trig[n][j] * θ_jump[j] for j = 1:(2d+1)))

    @NLparameter(model, πk[n = 1:N] == γ[n])
    @NLparameter(model, R[n = 1:N] == y[n])

    @NLobjective(
        model, Max,
        -sum(πk[n] * (Pθ[n] + R[n] * exp(-Pθ[n])) for n = 1:N)
    )

    # I don't know if it is the best but https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:πk] = πk
    model[:R] = R

    return model
end

function model_for_dist(dist::Gamma, θ_Y::AbstractArray, γ::AbstractVector, y, n2t, T::Int; silence = true)
    d = (size(θ_Y, 2) - 1) ÷ 2
    K = size(γ, 2)
    N = length(n2t)
    model = Model(Ipopt.Optimizer)
    set_optimizer_attributes(model, "max_iter" => 20)
    silence && set_silent(model)
    f = 2π / T
    cos_nj = [cos(f * j * n2t[n]) for n = 1:N, j = 1:d]
    sin_nj = [sin(f * j * n2t[n]) for n = 1:N, j = 1:d]

    trig = [[1; interleave2(cos_nj[n, :], sin_nj[n, :])] for n = 1:N]

    @variable(model, θ_jump[j = 1:(2d+1)])
    @variable(model, a_jump[j = 1:(2d+1)])
    # set_lower_bound(θ_jump[1], 0)
    # set_upper_bound.(θ_jump, 6)
    # set_lower_bound(a_jump[1], 0)
    # set_upper_bound(a_jump[1], 3)

    # Polynomial P
    @NLexpression(model, Pθ[n = 1:N], sum(trig[n][j] * θ_jump[j] for j = 1:length(trig[n])))
    @NLexpression(model, Pa[n = 1:N], sum(trig[n][j] * a_jump[j] for j = 1:length(trig[n])))

    @NLparameter(model, πk[n = 1:N] == γ[n])
    @NLparameter(model, R[n = 1:N] == y[n])
    register(model, :loggamma, 1, loggamma; autodiff = true) # cf the JuMP warning that prevent loggamma with more than one argument
    @NLobjective(
        model, Max,
        # -sum(πk[n]*(Pθ[n]+R[n]*exp(-Pθ[n])) for n in 1:N)
        # -sum(πk[n]*(exp(Pa[n])*Pθ[n]+R[n]*exp(-Pθ[n])+log(gamma(exp(Pa[n])))-expm1(Pa[n])*log(R[n])) for n in 1:N)
        -sum(πk[n] * (exp(Pa[n]) * Pθ[n] + R[n] * exp(-Pθ[n]) + loggamma(exp(Pa[n])) - expm1(Pa[n]) * log(R[n])) for n = 1:N)
    )
    # @NLexpression(model, mle,
    # sum(πk[n]*(-exp(Pa[n])*Pθ[n]-R[n]*exp(-Pθ[n])-loggamma(exp(Pa[n]))+expm1(Pa[n])*log(R[n])) for n in N)
    # )
    # I don't know if it is the best but https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:πk] = πk
    model[:R] = R
    return model
end

function model_for_α(θ_α::AbstractMatrix, γ::AbstractMatrix, n2t, T::Int; silence = true)
    N, K = size(γ)
    d = (size(θ_α, 2) - 1) ÷ 2
    model = Model(Ipopt.Optimizer)
    silence && set_silent(model)
    f = 2π / T
    cos_nj = [cos(f * j * n2t[n]) for n = 1:N, j = 1:d]
    sin_nj = [sin(f * j * n2t[n]) for n = 1:N, j = 1:d]

    trig = [[1; interleave2(cos_nj[n, :], sin_nj[n, :])] for n = 1:N]
    @variable(model, θ_jump[k = 1:K-1, j = 1:(2d+1)])
    # Polynomial P
    @NLexpression(model, Pn[n = 1:N, k = 1:K-1], sum(trig[n][j] * θ_jump[k, j] for j = 1:length(trig[n])))

    @NLparameter(model, πk[n = 1:N, k = 1:K] == γ[n, k])
    @NLobjective(
        model,
        Max,
        sum(sum(πk[n, k] * Pn[n, k] for k = 1:K-1) - log1p(sum(exp(Pn[n, k]) for k = 1:K-1)) for n = 1:N) #Σ_k πk[n,k] = 1
    )
    # I don't know if it is the best but https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:πk] = πk
    return model
end

function update_α!(
    α::AbstractArray where {T},
    θ::AbstractArray where {T},
    γ::AbstractMatrix,
    model::Model;
    warm_start = true
)
    N, K = size(γ)
    T = size(α, 2)

    # ξ are the filtering probablies
    θ_jump = model[:θ_jump]
    πk = model[:πk]
    ## Update the smoothing parameters in the JuMP model
    [set_value(πk[n, k], γ[n, k]) for n = 1:N, k = 1:K]
    warm_start && set_start_value.(θ_jump, θ)
    # Optimize the updated model
    set_silent(model)

    optimize!(model)
    # Obtained the new parameters
    θ[:, :] = value.(θ_jump)
    [α[k, t] = exp(polynomial_trigo(t, θ[k, :], T = T)) for k = 1:K-1, t = 1:T]
    [α[K, t] = 1 for t = 1:T] # last colum is 1/normalization (one could do otherwise)
    normalization_polynomial = [1 + sum(α[k, t] for k = 1:K-1) for t = 1:T]
    [α[k, t] /= normalization_polynomial[t] for k = 1:K, t = 1:T]
end

function update_dist(dist::AbstractVector{F} where {F<:Exponential}, θ_Y::AbstractArray, γ::AbstractVector, observations, model; warm_start = true)
    @argcheck size(γ, 1) == size(observations, 1)
    dist = copy(dist)
    θ_Y = copy(θ_Y)

    N = size(γ, 1)
    T = size(dist, 1)
    θ_jump = model[:θ_jump]
    πk = model[:πk]
    R = model[:R]

    ## Update the smoothing parameters in the JuMP model
    for n = 1:N
        set_value(πk[n], γ[n])
        set_value(R[n], observations[n])
    end

    warm_start && set_start_value.(θ_jump, θ_Y[:])
    set_silent(model)

    optimize!(model)
    θ_Y[:] = value.(θ_jump)

    p = [exp(polynomial_trigo(t, θ_Y[:], T = T)) for t = 1:T]
    dist[:] = Exponential.(p)
    return dist, θ_Y
end

function update_dist(dist::AbstractVector{F} where {F<:Gamma}, θ_Y::AbstractArray, γ::AbstractVector, observations, model; warm_start = true)
    @argcheck size(γ, 1) == size(observations, 1)
    dist = copy(dist)
    θ_Y = copy(θ_Y)

    N = size(γ, 1)
    T = size(dist, 1)
    θ_jump = model[:θ_jump]
    a_jump = model[:a_jump]
    πk = model[:πk]
    R = model[:R]

    ## Update the smoothing parameters in the JuMP model
    for n = 1:N
        set_value(πk[n], γ[n])
        set_value(R[n], observations[n])
    end

    warm_start && set_start_value.(a_jump, θ_Y[1, :]) # a parameter of Gamma
    warm_start && set_start_value.(θ_jump, θ_Y[2, :]) # θ parameter of Gamma

    optimize!(model)
    θ_Y[1, :] = value.(a_jump)
    θ_Y[2, :] = value.(θ_jump)

    a = [exp(polynomial_trigo(t, θ_Y[1, :], T = T)) for t = 1:T]
    p = [exp(polynomial_trigo(t, θ_Y[2, :], T = T)) for t = 1:T]
    dist[:] = [Gamma(a[t], p[t]) for t = 1:T]
    return dist, θ_Y

end

# * Improved and distributed compatible version

function model_for_dist_d(dist::Exponential, θ_Y::AbstractArray, T::Int; silence = true)
    d = (size(θ_Y, 1) - 1) ÷ 2

    model = Model(Ipopt.Optimizer)
    silence && set_silent(model)
    f = 2π / T

    cos_nj = [cos(f * j * t) for t = 1:T, j = 1:d]
    sin_nj = [sin(f * j * t) for t = 1:T, j = 1:d]
    trig = [[1; interleave2(cos_nj[t, :], sin_nj[t, :])] for t = 1:T]

    @variable(model, θ_jump[j = 1:(2d+1)])
    # Polynomial P
    @NLexpression(model, Pθ[t = 1:T], sum(trig[t][j] * θ_jump[j] for j = 1:(2d+1)))

    @NLparameter(model, s_π[t = 1:T] == 1)
    @NLparameter(model, s_π_R[t = 1:T] == 1) # same initialization but will be changed later

    """
        -sum(πk[n]*(Pθ[n]+R[n]*exp(-Pθ[n])) for n in 1:N)
    """
    @NLobjective(
        model, Max,
        -sum(s_π[t] * Pθ[t] + s_π_R[t] * exp(-Pθ[t]) for t = 1:T)
    )

    # I don't know if it is the best but https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:s_π] = s_π
    model[:s_π_R] = s_π_R

    return model
end

function update_dist_d(dist::AbstractVector{F} where {F<:Exponential}, θ_Y::AbstractArray, γ::AbstractVector, observations::AbstractArray, n_in_t, model::Model; warm_start = true)
    @argcheck size(γ, 1) == size(observations, 1)
    dist = copy(dist)
    θ_Y = copy(θ_Y)

    N = size(γ, 1)
    T = size(dist, 1)
    θ_jump = model[:θ_jump]
    s_π = model[:s_π]
    s_π_R = model[:s_π_R]

    ## Update the smoothing parameters in the JuMP model
    for t = 1:T
        set_value(s_π[t], sum(γ[n] for n in n_in_t[t]; init = 0))
        set_value(s_π_R[t], sum(γ[n] * observations[n] for n in n_in_t[t]; init = 0))
    end

    warm_start && set_start_value.(θ_jump, θ_Y[:])

    optimize!(model)
    θ_Y[:] = value.(θ_jump)

    p = [exp(polynomial_trigo(t, θ_Y[:], T = T)) for t = 1:T]
    dist[:] = Exponential.(p)
    return dist, θ_Y
end

function model_for_dist_d(dist::Gamma, θ_Y::AbstractArray, T::Int; silence = true)
    d = (size(θ_Y, 2) - 1) ÷ 2
    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "max_cpu_time", 60.0)
    set_optimizer_attribute(model, "max_iter", 100)
    silence && set_silent(model)
    f = 2π / T
    cos_nj = [cos(f * j * t) for t = 1:T, j = 1:d]
    sin_nj = [sin(f * j * t) for t = 1:T, j = 1:d]
    trig = [[1; interleave2(cos_nj[t, :], sin_nj[t, :])] for t = 1:T]

    @variable(model, θ_jump[j = 1:(2d+1)])
    @variable(model, a_jump[j = 1:(2d+1)])
    # set_lower_bound(θ_jump[1], 0)
    # set_upper_bound.(θ_jump, 6)
    # set_lower_bound(a_jump[1], 0)
    # set_upper_bound(a_jump[1], 3)

    # Polynomial P
    @NLexpression(model, Pθ[t = 1:T], sum(trig[t][j] * θ_jump[j] for j = 1:length(trig[t])))
    @NLexpression(model, Pa[t = 1:T], sum(trig[t][j] * a_jump[j] for j = 1:length(trig[t])))

    @NLparameter(model, s_π[t = 1:T] == 1)
    @NLparameter(model, s_π_R[t = 1:T] == 1)
    @NLparameter(model, s_π_logR[t = 1:T] == 1)

    register(model, :loggamma, 1, loggamma; autodiff = true) # cf the JuMP warning that prevent loggamma with more than one argument
    """
        - sum(πk[n]*(exp(Pa[n])*Pθ[n]+R[n]*exp(-Pθ[n])+loggamma(exp(Pa[n]))-expm1(Pa[n])*log(R[n])) for n in 1:N)
    """
    @NLobjective(
        model, Max,
        -sum(s_π[t] * exp(Pa[t]) * Pθ[t] + s_π_R[t] * exp(-Pθ[t])
             + s_π[t] * loggamma(exp(Pa[t])) - s_π_logR[t] * expm1(Pa[t]) for t = 1:T)
    )

    model[:s_π] = s_π
    model[:s_π_R] = s_π_R
    model[:s_π_logR] = s_π_logR

    return model
end

function update_dist_d(dist::AbstractVector{F} where {F<:Gamma}, θ_Y::AbstractArray, γ::AbstractVector, observations::AbstractArray, n_in_t, model::Model; warm_start = true)
    @argcheck size(γ, 1) == size(observations, 1)
    dist = copy(dist)
    θ_Y = copy(θ_Y)

    N = size(γ, 1)
    T = size(dist, 1)
    θ_jump = model[:θ_jump]
    a_jump = model[:a_jump]
    s_π = model[:s_π]
    s_π_R = model[:s_π_R]
    s_π_logR = model[:s_π_logR]

    ## Update the smoothing parameters in the JuMP model
    for t = 1:T
        set_value(s_π[t], sum(γ[n] for n in n_in_t[t]; init = 0))
        set_value(s_π_R[t], sum(γ[n] * observations[n] for n in n_in_t[t]; init = 0))
        set_value(s_π_logR[t], sum(γ[n] * log(observations[n]) for n in n_in_t[t]; init = 0))
    end

    warm_start && set_start_value.(a_jump, θ_Y[1, :]) # a parameter of Gamma
    warm_start && set_start_value.(θ_jump, θ_Y[2, :]) # θ parameter of Gamma

    optimize!(model)
    θ_Y[1, :] = value.(a_jump)
    θ_Y[2, :] = value.(θ_jump)

    a = [exp(polynomial_trigo(t, θ_Y[1, :], T = T)) for t = 1:T]
    p = [exp(polynomial_trigo(t, θ_Y[2, :], T = T)) for t = 1:T]
    dist[:] = [Gamma(a[t], p[t]) for t = 1:T]
    return dist, θ_Y

end

function s_γα!(s_ξ, ξ, n_in_t)
    T, K = size(s_ξ, 1), size(s_ξ, 2)
    for t = 1:T
        for k = 1:K
            s_ξ[t, k] = sum(ξ[n, k] for n in n_in_t[t]; init = 0)
        end
    end
end

function model_for_α_d(θ_α::AbstractMatrix, s_γ::AbstractMatrix; silence = true)
    T, K = size(s_γ)
    d = (size(θ_α, 2) - 1) ÷ 2
    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "max_cpu_time", 60.0)
    set_optimizer_attribute(model, "max_iter", 100)
    silence && set_silent(model)
    f = 2π / T
    cos_nj = [cos(f * j * t) for t = 1:T, j = 1:d]
    sin_nj = [sin(f * j * t) for t = 1:T, j = 1:d]
    trig = [[1; interleave2(cos_nj[t, :], sin_nj[t, :])] for t = 1:T]

    @variable(model, θ_jump[k = 1:K-1, j = 1:(2d+1)])
    # Polynomial P
    @NLexpression(model, P[t = 1:T, k = 1:K-1], sum(trig[t][j] * θ_jump[k, j] for j = 1:(2d+1)))

    @NLparameter(model, s_π[t = 1:T, k = 1:K] == s_γ[t, k])
    @NLobjective(
        model,
        Max,
        sum(sum(s_π[t, k] * P[t, k] for k = 1:K-1) - sum(s_π[t, k] for k = 1:K) * log1p(sum(exp(P[t, k]) for k = 1:K-1)) for t = 1:T) #Σ_k πk[n,k] = 1
    )
    # I don't know if it is the best but https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:s_π] = s_π
    return model
end

function update_α_d!(
    α::AbstractArray where {T},
    θ::AbstractArray where {T},
    s_γ::AbstractMatrix,
    model::Model;
    warm_start = true
)
    T, K = size(s_γ)

    # ξ are the filtering probablies
    θ_jump = model[:θ_jump]
    s_π = model[:s_π]
    ## Update the smoothing parameters in the JuMP model
    for t = 1:T, k = 1:K
        set_value(s_π[t, k], s_γ[t, k])
    end
    warm_start && set_start_value.(θ_jump, θ)
    # Optimize the updated model
    optimize!(model)
    # Obtained the new parameters
    θ[:, :] = value.(θ_jump)
    [α[k, t] = exp(polynomial_trigo(t, θ[k, :], T = T)) for k = 1:K-1, t = 1:T]
    [α[K, t] = 1 for t = 1:T] # last colum is 1/normalization (one could do otherwise)
    normalization_polynomial = [1 + sum(α[k, t] for k = 1:K-1) for t = 1:T]
    [α[k, t] /= normalization_polynomial[t] for k = 1:K, t = 1:T]
end

# TODO use the in place version (but it was bugy)

function fit_mle(mix::Vector{F} where {F<:MixtureModel}, θ_α::AbstractArray, θ_Y::AbstractArray, y, n2t;
    display = :none, maxiter = 100, tol = 1e-3, robust = false, silence = true, warm_start = true)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K, T = length(y), ncomponents(mix[1]), size(mix, 1)
    # TODO deal with K=1 case
    @argcheck K > 1
    d = (length(θ_α[1, :]) - 1) ÷ 2 #
    history = EMHistory(false, 0, [])
    # Allocate memory for in-place updates

    LL = zeros(N, K)
    γ = similar(LL)
    s_γ = zeros(T, K)

    c = zeros(N)

    # Initial parameters
    θ_α = copy(θ_α)
    θ_Y = copy(θ_Y)
    n_in_t = [findall(n2t .== t) for t = 1:T]

    α = copy(hcat(probs.(mix)...))
    dists = copy(hcat(components.(mix)...))
    types = typeof.(dists[:, 1])
    dists = [convert(Vector{types[k, 1]}, dists[k, :]) for k = 1:K] # I have to separate the array to preserve types

    model_α = model_for_α_d(θ_α, s_γ; silence = silence)
    model_dist = [model_for_dist_d(dists[k][1], θ_Y[k], T; silence = silence) for k = 1:K]

    # E-step
    # evaluate likelihood for each type k
    for k = 1:K, n = 1:N
        LL[n, k] = log(α[k, n2t[n]]) + logpdf(dists[k][n2t[n]], y[n])
    end
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    # get posterior of each category
    c[:] = logsumexp(LL, dims = 2)
    γ[:, :] = exp.(LL .- c)

    # Loglikelihood
    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter

        # M-step
        # with γ in hand, maximize (update) the parameters
        s_γα!(s_γ, γ, n_in_t)
        update_α_d!(α, θ_α, s_γ, model_α; warm_start = warm_start)
        for k = 1:K
            dists[k][:], θ_Y[k] = update_dist_d(dists[k], θ_Y[k], γ[:, k], y, n_in_t, model_dist[k], warm_start = warm_start)
        end
        # E-step
        # evaluate likelihood for each type k
        for k = 1:K, n = 1:N
            LL[n, k] = log(α[k, n2t[n]]) + logpdf(dists[k][n2t[n]], y[n])
        end
        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
        # get posterior of each category
        c[:] = logsumexp(LL, dims = 2)
        γ[:, :] = exp.(LL .- c)

        # Loglikelihood
        logtotp = sum(c)
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    return [MixtureModel([dists[k][t] for k = 1:K], [α[k, t] for k = 1:K]) for t = 1:T], θ_α, θ_Y, history
end

function fit_em_GE(mix::Vector{F} where {F<:MixtureModel}, θ_α::AbstractArray, θ_Y::AbstractArray, y, n2t;
    display = :none, maxiter = 100, tol = 1e-3, robust = false, silence = true, warm_start = true)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K, T = length(y), ncomponents(mix[1]), size(mix, 1)
    # TODO deal with K=1 case
    @argcheck K > 1
    d = (length(θ_α[1, :]) - 1) ÷ 2 #
    history = EMHistory(false, 0, [])
    # Allocate memory for in-place updates

    LL = zeros(N, K)
    γ = similar(LL)
    c = zeros(N)

    # Initial parameters
    θ_α = copy(θ_α)
    θ_Y = copy(θ_Y)

    α = copy(hcat(probs.(mix)...))
    dists = copy(hcat(components.(mix)...))
    types = typeof.(dists[:, 1])
    dists = [convert(Vector{types[k, 1]}, dists[k, :]) for k = 1:K] # I have to separate the array to preserve types

    model_α = model_for_α(θ_α, γ, n2t, T; silence = silence)
    model_dist = [model_for_dist(dists[k][1], θ_Y[k], γ[:, k], y, n2t, T; silence = silence) for k = 1:K]

    # E-step
    # evaluate likelihood for each type k
    for k = 1:K, n = 1:N
        LL[n, k] = log(α[k, n2t[n]]) + logpdf(dists[k][n2t[n]], y[n])
    end
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    # get posterior of each category
    c[:] = logsumexp(LL, dims = 2)
    γ[:, :] = exp.(LL .- c)

    # Loglikelihood
    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter

        # M-step
        # with γ in hand, maximize (update) the parameters
        update_α!(α, θ_α, γ, model_α; warm_start = warm_start)
        for k = 1:K
            dists[k][:], θ_Y[k] = update_dist(dists[k], θ_Y[k], γ[:, k], y, model_dist[k], warm_start = warm_start)
            # update_dist!(dists[k], θ_Y[k], γ[:,k], y, model_dist[k], warm_start = warm_start)
        end
        # E-step
        # evaluate likelihood for each type k
        for k = 1:K, n = 1:N
            LL[n, k] = log(α[k, n2t[n]]) + logpdf(dists[k][n2t[n]], y[n])
        end
        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
        # get posterior of each category
        c[:] = logsumexp(LL, dims = 2)
        γ[:, :] = exp.(LL .- c)

        # Loglikelihood
        logtotp = sum(c)
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    return [MixtureModel([dists[k][t] for k = 1:K], [α[k, t] for k = 1:K]) for t = 1:T], θ_α, θ_Y, history
end


# ---
# title: 901. Online Stock Span
# id: problem901
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Stack
# link: <https://leetcode.com/problems/online-stock-span/description/>
# hidden: true
# ---
#
# Write a class `StockSpanner` which collects daily price quotes for some stock,
# and returns the _span_  of that stock's price for the current day.
#
# The span of the stock's price today is defined as the maximum number of
# consecutive days (starting from today and going backwards) for which the price
# of the stock was less than or equal to today's price.
#
# For example, if the price of a stock over the next 7 days were `[100, 80, 60,
# 70, 60, 75, 85]`, then the stock spans would be `[1, 1, 1, 2, 1, 4, 6]`.
#
#
#
# **Example 1:**
#
#
#
#     Input: ["StockSpanner","next","next","next","next","next","next","next"], [[],[100],[80],[60],[70],[60],[75],[85]]
#     Output: [null,1,1,1,2,1,4,6]
#     Explanation:
#     First, S = StockSpanner() is initialized.  Then:
#     S.next(100) is called and returns 1,
#     S.next(80) is called and returns 1,
#     S.next(60) is called and returns 1,
#     S.next(70) is called and returns 2,
#     S.next(60) is called and returns 1,
#     S.next(75) is called and returns 4,
#     S.next(85) is called and returns 6.
#
#     Note that (for example) S.next(75) returned 4, because the last 4 prices
#     (including today's price of 75) were less than or equal to today's price.
#
#
#
#
# **Note:**
#
#   1. Calls to `StockSpanner.next(int price)` will have `1 <= price <= 10^5`.
#   2. There will be at most `10000` calls to `StockSpanner.next` per test case.
#   3. There will be at most `150000` calls to `StockSpanner.next` across all test cases.
#   4. The total time limit for this problem has been reduced by 75% for C++, and 50% for all other languages.
#
#
## @lc code=start
using LeetCode

Base.@kwdef struct StockSpanner
    history::Vector{Pair{Int,Int}} = []
end

function next!(s::StockSpanner, price::Int)
    span = 1
    while !isempty(s.history) && s.history[end][1] <= price
        span += s.history[end][2]
        pop!(s.history)
    end
    push!(s.history, price => span)
    return span
end
## @lc code=end


## Dimensions and descriptions

desc = doc"
# Describing `RLEVector` Objects
For an RLEVector `x = RLEVector([4,5,6],[3,6,9])`
 * `length(x)` The full length of the vector, uncompressed
 * `size(x)` Same as `length`, as for any other vector
 * `size(x,dim)` Returns `(length(x),1) for dim == 1`
 * `rstart(x)` The index of the beginning of each run
 * `rwidth(x)` The width of each run
 * `rstop(x)` The index of the end of each run
 * `rlast(x)` The data value for each run
 * `isempty(x)` Returns boolean, as for any other vector
 * `nrun(x)` Returns the number of runs represented in the array
 * `endtype(x)` Returns the element type of the run ends

## See also
length, size, rstart, rwidth, rstop, rvalue, isempty, nrun, endtype
"

@doc desc ->
function ndims(x::RLEVector)
    return(1)
end

@doc desc ->
function nrun(x::RLEVector)
  length(x.runends)
end

@doc desc ->
function length{T1,T2<:Integer}(x::RLEVector{T1,T2})
  re = x.runends
  ind = endof(re)
  if (ind > 0)
    @inbounds len = re[ind]
  else
    len = zero(T2)
  end
  return(len)
end

@doc desc ->
function size(x::RLEVector)
  (length(x),)
end

@doc desc ->
function size(x::RLEVector, dim::Integer)
  len = length(x)
  if dim == 1
    return(len)
  else
    return( (len,1) )
  end
end

@doc desc ->
function isempty(x::RLEVector)
  isempty(x.runends)
end

### Getters
@doc desc ->
function rfirst(x::RLEVector)
  re = x.runends
  rval = similar(re)
  prev = zero(eltype(re))
    @inbounds for i in eachindex(re)
        rval[i] = prev + 1
        prev = re[i]
  end
  return(rval)
end

@doc desc ->
function rfirst(x::RLEVector, run::Integer)
    num_one = one(eltype(x.runends))
    run == 1 ? num_one : x.runends[run-1] + num_one
end

@doc desc ->
function rwidth(x::RLEVector)
  re = x.runends
  rval = similar(re)
  prev = zero(eltype(re))
  @inbounds for i in eachindex(re)
      rei = re[i]
      rval[i] = rei - prev
      prev = rei
  end
  return(rval)
end

function rwidth(x::RLEVector, run::Integer)
  run == 1 ? x.runends[1] : x.runends[run] - x.runends[run-1]
end

@doc desc ->
function rlast(x::RLEVector)
  x.runends
end

@doc desc ->
function rvalue(x::RLEVector)
  x.runvalues
end

@doc desc ->
endtype(x::RLEVector) = eltype(rlast(x))


#!/usr/bin/env julia

function fact(n::Int)::Int64
    result = 1
    for i in 2:n
        result *= i
    end
    return result
end

for n in 0:5
    println("fact($n) = $(fact(n))")
end


x = LinRange(0,2,100)
plot(x, [x.^α for α in 0:0.2:10], c = :plasma, colorbar = true,
    linez = collect(0:0.2:10)', legend = false, ylim = (-0.1,10),
    xlab = L"x", size = (600,300))
annotate!(0.5, 7.5, text(L"f(x) = x^\alpha", :purple))
annotate!(2.15, 10.2, text(L"\alpha", :purple))
# because colorbar_title is currently not implemented for GR
savefig("figpow.pdf")