export classic_sta_lta, classic_sta_lta!, delayed_sta_lta, delayed_sta_lta!
export recursive_sta_lta, recursive_sta_lta!, carl_sta_trig, carl_sta_trig!
export z_detect, z_detect!, trigger_onset

@doc """
  classic_sta_lta!(S, sta, lta)

Computes classic STA/LTA.

The length of the STA and LTA are given by `sta` and `lta` in seconds, respectively.

!!! note

    The in-place version (classic_sta_lta!) overwrites the seismic trace with the STA/LTA
    characteristic function.

# Arguments
- `S::GphysData`: GphysData (SeisData, SeisChannel, SeisIO.Nodal.NodalChannel, SeisIO.Quake.EventChannel).
- `sta:Int`: Length of short time average window in seconds
- `lta:Int`: Length of long time average window in seconds
""" classic_sta_lta!
function classic_sta_lta!(S::GphysData, sta::Real, lta::Real)
    for ii in 1:S.n
        classic_sta_lta!(S[ii], sta, lta)
    end
    return nothing
end
function classic_sta_lta!(C::GphysChannel, sta::Real, lta::Real)
    nsta = round(Int64, C.fs * sta)
    nlta = round(Int64, C.fs * lta)
    C.x .= classic_sta_lta(C.x, nsta, nlta)
    return nothing
end

@doc classic_sta_lta!
function classic_sta_lta(S::GphysData, sta::Real, lta::Real)
    U = deepcopy(S)
    classic_sta_lta!(U, sta, lta)
    return U
end
function classic_sta_lta(C::GphysChannel, sta::Real, lta::Real)
    U = deepcopy(C)
    classic_sta_lta!(U, sta, lta)
    return U
end
function classic_sta_lta(A::AbstractArray, nsta::Real, nlta::Real)
    nsta > nlta ? throw(DomainError((nsta, nlta), "Short-term average must be less than long-term average")) : nothing
    nsta < 0 ? throw(DomainError(nsta, "Short-term average must be less greater than zero")) : nothing
    nlta < 0 ? throw(DomainError(nlta, "Long-term average must be greater than zero")) : nothing
    N = size(A,1)
    T = eltype(A)
    nsta > N ? throw(DomainError(nsta, "Short-term average must be less than length of time series")) : nothing
    nlta > N ? throw(DomainError(nlta, "Long-term average must be less than length of time series")) : nothing

    # cumulative sum can be exploited to calculate a moving average
    sta = cumsum(A .^ 2)
    lta = deepcopy(sta)

    # compute the STA and LTA
    sta[nsta + 1:end] = sta[nsta + 1:end] .- sta[1:end-nsta]
    sta ./= nsta
    lta[nlta + 1:end] = lta[nlta + 1:end] .- lta[1:end-nlta]
    lta ./= nlta

    # pad zeros
    sta[1:nlta] .= 0.0

    # avoid division by zero by setting zero values to tiny float
    tinyfloat = eps(T)
    ind = findall(lta .< tinyfloat)
    lta[ind] .= tinyfloat
    return sta ./ lta
end

@doc """
  delayed_sta_lta!(S, sta, lta)

Delayed STA/LTA from Withers, 1998.

!!! note

    The in-place version (delayed_sta_lta!) overwrites the seismic trace with the STA/LTA
    characteristic function.

# Arguments
- `S::GphysData`: GphysData (SeisData, SeisChannel, SeisIO.Nodal.NodalChannel, SeisIO.Quake.EventChannel).
- `sta::Real`: Length of short time average window in seconds
- `lta::Real`: Length of long time average window in seconds
""" delayed_sta_lta!
function delayed_sta_lta!(S::SeisData, sta::Real, lta::Real)
    for ii in 1:S.n
        delayed_sta_lta!(S[ii], sta, lta)
    end
    return nothing
end
function delayed_sta_lta!(C::GphysChannel, sta::Real, lta::Real)
    nsta = round(Int64, C.fs * sta)
    nlta = round(Int64, C.fs * lta)
    C.x .= delayed_sta_lta(C.x, nsta, nlta)
    return nothing
end

@doc delayed_sta_lta!
function delayed_sta_lta(S::GphysData, sta::Real, lta::Real)
    U = deepcopy(S)
    delayed_sta_lta!(U, sta, lta)
    return U
end
function delayed_sta_lta(C::GphysChannel, sta::Real, lta::Real)
   U = deepcopy(C)
   delayed_sta_lta!(U, sta, lta)
   return U
end
function delayed_sta_lta(A::AbstractArray, nsta::Int, nlta::Int)
    nsta > nlta ? throw(DomainError((nsta, nlta), "Short-term average must be less than long-term average")) : nothing
    nsta < 0 ? throw(DomainError(nsta, "Short-term average must be less greater than zero")) : nothing
    nlta < 0 ? throw(DomainError(nlta, "Long-term average must be greater than zero")) : nothing
    N = size(A, 1)
    T = eltype(A)
    nsta > N ? throw(DomainError(nsta, "Short-term average must be less than length of time series")) : nothing
    nlta > N ? throw(DomainError(nlta, "Long-term average must be less than length of time series")) : nothing

    sta = zeros(T, N)
    lta = zeros(T, N)
    sta[1] = A[1] ^ 2 / nsta
    for ii in 2:N
        sta[ii] = (A[ii] ^ 2 + A[max(1, ii - nsta)] ^ 2) / nsta + sta[ii - 1]
        lta[ii] = (A[max(1, ii - nsta - 1)] ^ 2 + A[max(1, ii - nsta - nlta -1)] ^ 2) / nlta + lta[max(1, ii -1)]
    end
    sta[1:min(N,nlta + nsta + 50)] .= 0.0
    lta[1:min(N,nlta + nsta + 50)] .= 1.0 # avoid division by zero
    return sta ./ lta
end

@doc """
  recursive_sta_lta!(S, sta, lta)

Computes recursive STA/LTA.

The length of the STA and LTA are given by `sta` and `sta` in seconds, respectively.

!!! note

    The in-place version (recursive_sta_lta!) overwrites the seismic trace with the STA/LTA
    characteristic function.

# Arguments
- `S::GphysData`: GphysData (SeisData, SeisChannel, SeisIO.Nodal.NodalChannel, SeisIO.Quake.EventChannel).
- `sta:Int`: Length of short time average window in seconds
- `lta:Int`: Length of long time average window in seconds

See also: [`classic_sta_lta`](@ref), [`carl_sta_trig`](@ref), [`delayed_sta_lta`](@ref) [`z_detect`](@ref)
""" recursive_sta_lta!
function recursive_sta_lta!(S::GphysData, sta::Real, lta::Real)
    for ii in 1:S.n
        recursive_sta_lta!(S[ii], sta, lta)
    end
    return nothing
end
function recursive_sta_lta!(C::GphysChannel, sta::Real, lta::Real)
    nsta = round(Int64, C.fs * sta)
    nlta = round(Int64, C.fs * lta)
    C.x .= recursive_sta_lta(C.x, nsta, nlta)
    return nothing
end

@doc recursive_sta_lta!
function recursive_sta_lta(S::GphysData, sta::Real, lta::Real)
    U = deepcopy(S)
    recursive_sta_lta!(U, sta, lta)
    return U
end
function recursive_sta_lta(C::GphysChannel, sta::Real, lta::Real)
    U = deepcopy(C)
    recursive_sta_lta!(U, sta, lta)
    return U
end
function recursive_sta_lta(A::AbstractArray, nsta::Int, nlta::Int)
    nsta > nlta ? throw(DomainError((nsta, nlta), "Short-term average must be less than long-term average")) : nothing
    nsta < 0 ? throw(DomainError(nsta, "Short-term average must be less greater than zero")) : nothing
    nlta < 0 ? throw(DomainError(nlta, "Long-term average must be greater than zero")) : nothing
    N = size(A,1)
    T = eltype(A)
    nsta > N ? throw(DomainError(nsta, "Short-term average must be less than length of time series")) : nothing
    nlta > N ? throw(DomainError(nlta, "Long-term average must be less than length of time series")) : nothing

    # compute short time average (STA) and long term average (LTA)
    # given by Evans and Allen
    csta = 1.0 / nsta
    clta = 1.0 / nlta
    sta = 0.0
    lta = 1e-99 # avoid zero division
    charfct = zeros(eltype(A), size(A))
    icsta = 1.0 - csta
    iclta = 1.0 - clta

    for ii in 2:N
        sq = A[ii] ^ 2
        sta = csta * sq + icsta * sta
        lta = clta * sq + iclta * lta
        charfct[ii] = sta / lta

        if ii < nlta
            charfct[ii] = 0.0
        end
    end
    return charfct
end

@doc """
  carl_sta_trig!(S, sta, lta, ratio, quiet)

Computes the CarlSTAtrig characteristic function.

eta = star - (ratio * ltar) - abs(sta - lta) - quiet

!!! note

    The in-place version (carl_sta_trig!) overwrites the seismic trace with the STA/LTA
    characteristic function.

# Arguments
- `S::GphysData`: GphysData (SeisData, SeisChannel, SeisIO.Nodal.NodalChannel, SeisIO.Quake.EventChannel).
- `sta::Real`: Length of short time average window in seconds
- `lta::Real`: Length of long time average window in seconds
- `ratio::AbstractFloat`: as `ratio` gets smaller, carl_sta_trig gets more sensitive
- `Quiet::AbstractFloat`: as `quiet` gets smaller, carl_sta_trig gets more sensitive
""" carl_sta_trig!
function carl_sta_trig!(S::GphysData, sta::Real, lta::Real, ratio::Real, quiet::Real)
    for ii in 1:S.n
        carl_sta_trig!(S[ii],sta, lta, ratio, quiet)
    end
    return nothing
end
function carl_sta_trig!(C::GphysChannel, sta::Real, lta::Real, ratio::Real, quiet::Real)
    nsta = round(Int64, C.fs * sta)
    nlta = round(Int64, C.fs * lta)
    C.x .= carl_sta_trig(C.x, nsta, nlta, ratio, quiet)
    return nothing
end

@doc carl_sta_trig!
function carl_sta_trig(S::GphysData, sta::Real, lta::Real, ratio::Real, quiet::Real)
    U = deepcopy(S)
    carl_sta_trig!(U, sta, lta, ratio, quiet)
    return U
end
function carl_sta_trig(C::GphysChannel, sta::Real, lta::Real, ratio::Real, quiet::Real)
    U = deepcopy(C)
    carl_sta_trig!(U, sta, lta, ratio, quiet)
    return U
end
function carl_sta_trig(A::AbstractArray, nsta::Int, nlta::Int, ratio::Real, quiet::Real)
    nsta > nlta ? throw(DomainError((nsta,nlta), "Short-term average must be less than long-term average")) : nothing
    nsta < 0 ? throw(DomainError(nsta, "Short-term average must be less greater than zero")) : nothing
    nlta < 0 ? throw(DomainError(nlta, "Long-term average must be greater than zero")) : nothing
    N = size(A,1)
    T = eltype(A)
    nsta > N ? throw(DomainError(nsta, "Short-term average must be less than length of time series")) : nothing
    nlta > N ? throw(DomainError(nlta, "Long-term average must be less than length of time series")) : nothing
    sta = zeros(T,N)
    lta = deepcopy(sta)
    star = deepcopy(sta)
    ltar = deepcopy(sta)
    pad_sta = zeros(T, nsta)
    pad_lta = zeros(T, nlta)

    # compute the short time average (STA)
    for ii in 1:nsta
        sta .+= vcat(pad_sta, A[ii:N - nsta - 1 + ii])
    end
    sta ./= nsta

    # compute the long time average (LTA)
    for ii in 1:nlta
        lta .+= vcat(pad_lta, sta[ii:N - nlta - 1 + ii])
    end
    lta ./= nlta
    lta = vcat(T(0), lta[1:end-1])

    # compute star, average of abs diff between trace and lta
    for ii in 1:nsta
        star .+= vcat(pad_sta, abs.(A[ii:N - nsta - 1 + ii] .- lta[ii:N - nsta - 1 + ii]))
    end
    star ./= nsta

    # compute ltar, long term average over star
    for ii in 1:nlta
        ltar .+= vcat(pad_lta, star[ii:N - nlta - 1 + ii])
    end
    ltar ./= nlta

    eta = star .- (T(ratio) .* ltar) .- abs.(sta .- lta) .- T(quiet)
    eta[1:nlta] .= T(-1.0)
    return eta
end

@doc """
  z_detect!(S, sta)

Z-detector from Withers, 1998.

!!! note

    The in-place version (delayed_sta_lta!) overwrites the seismic trace with the STA/LTA
    characteristic function.

# Arguments
- `S::GphysData`: GphysData (SeisData, SeisChannel, SeisIO.Nodal.NodalChannel, SeisIO.Quake.EventChannel).
- `sta::Real`: Length of short time average window in seconds
""" z_detect!
function z_detect!(S::GphysData, sta::Real)
    for ii in 1:S.n
        z_detect!(S[ii], sta)
    end
    return nothing
end
function z_detect!(C::GphysChannel, sta::Real)
    nsta = round(Int64, C.fs * sta)
    C.x .= z_detect(C.x, nsta)
    return nothing
end

@doc z_detect!
function z_detect(S::GphysData, sta::Real)
    U = deepcopy(S)
    z_detect!(U, sta)
    return U
end
function z_detect(C::GphysChannel, sta::Real)
   U = deepcopy(C)
   z_detect!(U, sta)
   return U
end
function z_detect(A::AbstractArray, nsta::Int)
    nsta < 0 ? throw(DomainError(nsta, "Short-term average must be less greater than zero")) : nothing
    N = size(A, 1)
    T = eltype(A)
    nsta > N ? throw(DomainError(nsta, "Short-term average must be less than length of time series")) : nothing

    # Z-detector given by Swindell and Snell, 1977
    sta = zeros(T, N)
    A2 = A .^ 2

    # replace this with convolution
    for ii in 1:nsta
        sta[ii:N - nsta - 1 + ii] .+= A2[ii:N - nsta - 1 + ii]
    end
    return (sta .- mean(sta)) ./ std(sta)
end

@doc """
  trigger_onset(charfct, thresh1, thresh2; max_len=1e99, max_len_delete=false)

Calculate trigger on and off times.

Given `thresh1` and `thresh2` calculate trigger on and off times from
characteristic function.

# Arguments
- `charfct::AbstractArray`: Characteristic function of e.g. STA/LTA trigger.
- `thresh1::Real`: Value above which trigger (of characteristic function)
    is activated (higher threshold).
- `thresh2::Real`: Value below which trigger (of characteristic function)
    is deactivated (lower threshold).
# Keywords
- `max_len::Int`: Maximum length of triggered event in samples. A new
    event will be triggered as soon as the signal reaches again above thresh1.
- `max_len_delete::Bool`: Do not write events longer than max_len into report file.
""" trigger_onset
function trigger_onset(charfct::AbstractArray, thresh1::Real, thresh2::Real;
    max_len::Int=10^10, max_len_delete::Bool=false,
)
    thresh2 > thresh1 ? throw(DomainError(thresh1, "thresh1 $thresh1 must be greater than thresh2 $thresh2")) : nothing
    ind1 = findall(charfct .> thresh1)
    if length(ind1) == 0
        return Array{Int64}(undef,0,0)
    end
    ind2 = findall(charfct .> thresh2)

    on = [ind1[1]]
    off = [-1]

    # determine the indices where charfct falls below off-threshold
    ind2_ = Array{Bool}(undef,length(ind2))
    ind2_[1:end-1] .= diff(ind2) .> 1
    # last occurence is missed by diff, add it manually
    ind2_[end] = true
    append!(off,ind2[ind2_])
    append!(on,ind1[findall(diff(ind1) .> 1) .+ 1])
    # include last pick if trigger is on or drop it
    if max_len_delete
        # drop it
        append!(off,max_len)
        append!(on,on[end])
    else
        # include it
        append!(off,ind2[end])
    end

    pick = []
    while on[end] > off[1]
        while on[1] <= off[1]
            deleteat!(on,1)
        end
        while off[1] < on[1]
            deleteat!(off,1)
        end
        if off[1] - on[1] > max_len
            if max_len_delete
                deleteat!(on,1)
                continue
            end
            prepend!(off,on[1] + max_len)
        end
        push!(pick,[on[1],off[1]])
    end
    return permutedims(hcat(pick...))
end


function algpoly(plist::Vector{T}, f, n) where {T}
    @debug "Algorithm [algpoly]" plist f
    R = parent(n)
    # Construct polynomials qj that arise from treating the shift operator
	# as the difference operator plus the identity operator.
    r = length(plist) - 1
    qlist = [sum(binomial(i, j) * plist[i+1] for i in j:r) for j in 0:r]

    # Find all candidates for the degree bound on output polynomial.
    b = maximum([(Nemo.degree(q) - (j - 1)) for (j, q) in enumerate(qlist)])
    alpha = sum(lead(q)*fallfactorial(n, j-1) for (j, q) in enumerate(qlist) if Nemo.degree(q) - (j-1) == b)
    roots = [-Nemo.coeff(f, 0) for (f, _) in Nemo.factor(R(alpha))]
    roots = [r for r in roots if r >= 0 && isone(denominator(r))]

    first = Nemo.degree(f) - b
    second = -b - 1
    third = maximum([roots; 0])
    d = convert(Int, numerator(max(first, second, third)))
    @debug "Degree bound" first second third

    # Use method of undetermined coefficients to find the output polynomial.
    R, varlist = PolynomialRing(base_ring(n), ["x$i" for i in 1:d+1])
    arg = change_base_ring(R, n)
    genpoly = sum(v*arg^(i-1) for (i,v) in enumerate(varlist))
    _plist = map(x->change_base_ring(R, x), plist)
    @debug "Generic polynomial" genpoly varlist

    poly = sum(genpoly(arg+(i - 1)) * p for (i, p) in enumerate(_plist))
    poly -= change_base_ring(R, f)
    @debug "Poly" poly

    coefficients = [Nemo.coeff(poly, i) for i in 0:length(poly)]
    A = [Nemo.coeff(cf, v) for cf in coefficients, v in varlist]
    A = MatrixSpace(base_ring(R), size(A)...)(A)
    rank, ker = kernel(A)
    solutions = typeof(n)[]
    S = base_ring(n)
    for j in 1:rank
        col = Array(ker[:, j])
        s = sum(Nemo.coeff(genpoly, i)(col...)*n^i for i in 0:length(genpoly))
        push!(solutions, s)
    end
    @debug "Algorithm [algpoly] return" solutions
    return solutions
end

function monic_factors(p::PolyElem)
    fac = Nemo.factor(p)
    u = Nemo.coeff(parent(p)(unit(fac)), 0)
    isempty(fac) && return u, [one(p)]
    fs = [(lead(f), div(f, parent(p)(lead(f)))) for (f, m) in fac for i in 1:m]
    return u*prod(map(first, fs)), map(last, fs)
end

function all_monic_factors(p::PolyElem)
    unit, fs = monic_factors(p)
    fs = [one(p); fs]
    unique(prod(s) for s in Combinatorics.powerset(fs, 1))
end

_prod(itr) = reduce(*, itr, init=1)

function alghyper(plist::Vector{T}, n; all_solutions::Bool=false) where {T}
    @debug "Algorithm [alghyper]" plist
    d = length(plist)

    alist = all_monic_factors(plist[1](n))
    blist = all_monic_factors(plist[end](n-(d-2)))

    @debug "All monic factors" alist blist

    solutions = FracElem{typeof(n)}[]
    for a in alist
        for b in blist
            @debug "Monic factors" a b

            ps = [plist[i+1] * _prod(a(n+j) for j in 0:i-1) * _prod(b(n+j) for j in i:d-2) for i in 0:d-1]

            m = maximum(Nemo.degree(p) for p in ps)
            alphas = map(x->Nemo.coeff(x, m), ps)
            zpoly = sum(c*n^(i-1) for (i, c) in enumerate(alphas))
            zs = roots(zpoly)
            zs = [z for z in zs if !iszero(z)]
            @debug "Nonzero roots" alphas zpoly zs
            for z in zs
                _ps = [z^(i-1)*p for (i, p) in enumerate(ps)]
                @debug "Coefficients auxiliary recurrence" _ps

                _sols = algpoly(_ps, zero(n), n)
                @debug "Solutions of algpoly" _sols

                for c in _sols
                    s = z * (a(n) // b(n)) * (c(n+1) // c(n))
                    !all_solutions && return s
                    push!(solutions, s)
                end
            end
        end
    end
    @debug "Algorithm [alghyper] return" unique(solutions) solutions
    !all_solutions && return nothing
    return unique(solutions)
end

function gosper(t::Seq{T}, arg) where {T <: FieldElem}
    @debug "[Gosper]" t arg
    r = div(t(arg+1), t(arg))
    @assert isrational(r)
    r = first(coeffs(r))
    c1, fac1 = monic_factors(numerator(r))
    c2, fac2 = monic_factors(denominator(r))
    if isempty(fac1) && isempty(fac2)
        @debug "Rational coefficients" c1 c2
        return parent(t)(c1//c2*arg)
    end
    R = parent(arg)
    f = R(_prod(fac1))
    g = R(_prod(fac2))
    @debug "Gosper start" t r f g
    @assert gcd(f, g) == 1

    _R, _h = PolynomialRing(base_ring(R), "h")
    _arg = change_base_ring(_R, arg)
    _res = resultant(f(_arg), g(_arg+_h))
    @debug "Resultant" _res

    hs = roots(_res)
    hs = [numerator(r) for r in hs if r >= 0 && isone(denominator(r))]
    @debug "Non-negative integer roots of resultant" hs

    sz, type = length(hs), typeof(arg)
    ss = Array{type}(undef, sz)
    ps = Array{type}(undef, sz+1)
    qs = Array{type}(undef, sz+1)
    ps[1] = f
    qs[1] = g
    for j in 1:length(hs)
        ss[j] = gcd(ps[j](arg), qs[j](arg+hs[j]))
        ps[j+1] = div(ps[j], ss[j])
        qs[j+1] = div(qs[j], ss[j](arg-hs[j]))
    end
    a = c1//c2 * ps[end]
    b = qs[end]
    c = R(_prod(ss[i](arg-j) for i in 1:length(hs) for j in 1:Int(hs[i])))
    @debug "Coefficients for auxiliary equation" a b c ss

    if degree(a) != degree(b) || lead(a) != lead(b)
        D = [degree(c) - max(degree(a), degree(b))]
    else
        d = degree(a)
        if iszero(d)
            D = [degree(c) - d + 1]
        else
            A = Nemo.coeff(a, d-1)
            B = Nemo.coeff(b(arg-1), d-1)
            dd = (B-A)/lead(a)
            if dd >= 0 && isone(denominator(dd))
                D = [degree(c) - d + 1, convert(Int, numerator(dd))]
            else
                D = [degree(c) - d + 1]
            end
        end
    end
    filter!(x->x>=0, D)
    @debug "Degree set" D
    isempty(D) && return nothing
    d = maximum(D)

    # Use method of undetermined coefficients
    _S, _vs = PolynomialRing(base_ring(R), ["v$i" for i = 0:d])
    _arg = change_base_ring(_S, arg)
    _a = change_base_ring(_S, a)
    _b = change_base_ring(_S, b)
    aux = sum(_vs[i+1]*_arg^i for i = 0:d)
    poly = _a(_arg)*aux(_arg+1) - _b(_arg-1)*aux(_arg)
    @debug "Solution template" aux poly

    deg = degree(poly)
    cs = [Nemo.coeff(poly, i) for i in 0:deg]
    mA = [Nemo.coeff(c, v) for c in cs, v in _vs]
    mB = [Nemo.coeff(c, i) for i in 0:deg]
    @debug "Solve" mA mB

    mM = hcat(mA, mB)
    mM = MatrixSpace(base_ring(arg), size(mM)...)(mM)
    r, mM = rref(mM)

    xs = zeros(base_ring(arg), d+1)
    nc = ncols(mM)
    for i in 1:r
        I = findfirst(isone, Array(mM[i, 1:nc-1]))
        xs[I[2]] = mM[i, nc]
    end
    @debug "Method of undetermined coefficients" r mM xs
    iszero(xs) && return nothing

    x = sum(xs[i+1]*arg^i for i = 0:d)
    z = parent(t)(b(arg-1)*x(arg)//c(arg))*t
    @debug "Solution from gosper" z x b c
    @assert z(arg+1) - z(arg) == t
    return z
end

issolution(s::Seq, cfs::Vector{<:FracElem}, arg::PolyElem) = iszero(sum(c * s(arg + i - 1) for (i, c) in enumerate(cfs)))

function petkovsek(cfs::Vector{<:FracElem}, arg::PolyElem, level::Int=0)
    @debug "[Petkovsek]" cfs arg
    # Get rid of denominator
    m = reduce(lcm, map(denominator, cfs))
    polys = [numerator(c * m) for c in cfs]

    ls = map(monic_factors, polys)
    cs = map(first, ls)
    g = reduce(gcd, cs)
    cs = map(x->div(x, g), cs)
    polys = [parent(arg)(c*p) for (c, p) in zip(cs, map(_prod∘last, ls))]

    g = reduce(gcd, polys)
    polys = map(x->div(x, g), polys)
    S = SequenceRing(parent(arg))
    sz = length(polys)
    if sz == 2
        # Recurrence is hypergeometric
        return [S((-polys[1]//polys[2])(arg), arg+1)(arg-1)]
    end

    # Get first hypergeometric solution
    hgterm = alghyper(polys, arg)
    if isnothing(hgterm)
        @debug "No solution from Hyper"
        return nothing
    end
    s_n = S(hgterm(arg), arg+1)(arg-1)
    @debug "Hyper" hgterm s_n
    @assert issolution(s_n, cfs, arg)

    # Take y(n) = z(n)s(n) for unknown z(n), and plug into original recurrence.
    # Then reduce order by taking u(n) = z(n+1) - z(n)
    R, zs = PolynomialRing(S, ["z$i" for i in 0:sz-1])
    poly = sum(s_n(arg+(i-1))*S(cfs[i])*zs[i] for i in 1:sz)

    bs = [Nemo.coeff(poly, z) for z in zs]
    for i in reverse(2:length(bs)-1)
        bs[i] = bs[i] + bs[i+1]
    end
    bs = bs[2:end]
    g = reduce(gcd, bs)
    bs = map(x->div(x, g), bs)
    @debug "Coefficients for order-reduced recurrence" g bs map(isrational, bs) terms(bs[end])
    @assert all(isrational(b) for b in bs)
    # Recursively solve for u(n)
    u_basis = petkovsek(map(first∘coeffs, bs), arg, level+1)
    @assert all(iszero(sum(c * u(arg + i - 1) for (i, c) in enumerate(bs))) for u in u_basis)
    # Use Gosper to resolve antidifference, i.e. find z(n) satisfying u(n) = z(n+1) - z(n)
    z_basis = [gosper(u, arg) for u in u_basis]
    if any(map(isnothing, z_basis))
        @debug "No solution from Gosper"
        return nothing
    end
    @debug "Solutions from Gosper" z_basis
    ss = [s_n * z for z in z_basis]
    @assert all(issolution(s, cfs, arg) for s in ss)
    [s_n; ss]
end


struct MEM{T,S}
    Ω::T
    Φ::S
    fdspec::Matrix{Float64}
    D::Matrix{Float64}
    function MEM(S, Ω, Φ)
        D = zeros(length(Φ), length(Ω)) # transpose on return because should be frequency direction
        fdspec = similar(D)
        for (i, s) ∈ enumerate(S)
            a1, b1, a2, b2 = freqdir_fouriercoef(s)
            c1 = a1 + 1im * b1
            c2 = a2 + 1im * b2
            F₁ = (c1 - c2 * conj(c1)) / (1 - abs(c1)^2)
            F₂ = c2 - c1 * F₁

            @. D[:, i] = real(
                (1 / (2pi)) * (1 - F₁ * conj(c1) - F₂ * conj(c2)) /
                (abs(1 - F₁ * exp(-1im * Φ) - F₂ * exp(-2im * Φ))^2),
            )

            for j = 1:length(Φ)
                fdspec[j, i] = D[j, i] * real.(s[1, 1])
            end
        end
        new{typeof(Ω), typeof(Φ)}(Ω, Φ, fdspec', D')
    end
end
function freqdir_fouriercoef(S)
    denom = sqrt(abs((S[2, 2] + S[3, 3]) * S[1, 1]))
    a1 = imag(S[2,1]) / denom
    b1 = imag(S[3,1]) / denom
    a2 = real((S[2, 2] - S[3, 3]) / (S[2, 2] + S[3, 3]))
    b2 = 2 * real((S[2, 3]) / (S[2, 2] + S[3, 3]))
    return a1, b1, a2, b2
end


using Omega
using UnicodePlots

weight = β(2.0,2.0)
beta_samples = rand(weight, 10000)

UnicodePlots.histogram(beta_samples)

nflips = 4
coinflips_ = [bernoulli(weight, Bool) for i = 1:nflips]


using Combinatorics

# Part 2
checksum = 0;
for line in eachline(STDIN)
  input = parse.(Int, split(line, ' '))

  for (x, y) in combinations(input, 2)
    if x % y == 0
      checksum += Int(x / y)
      break
    end

    if y % x == 0
      checksum += Int(y / x)
      break
    end
  end
end

println("Part 2: $checksum")


# Poower-law distributions of infections in a simple metapopulation model
# this corresponds to Fig.3 in
# Blasius B. (2020) Power-law distribution in the number of confirmed COVID-19 cases. Chaos
# https://doi.org/10.1063/5.0013031

module C

using Dates, LsqFit, PyPlot, Random, StatsBase, DifferentialEquations, Polynomials

include("tools.jl")

function SIRD(du,u,p,t)  # standard SIR model + dead class
  S,I,R,D = u
  b,g,m = p
  N = S+I+R+D
  du[1] = dS = -b*S*I/N       # susceptible
  du[2] = dI = b*S*I/N  - g*I # infected
  du[3] = dR = (1-m)*g*I      # recovered
  du[4] = dD = m*g*I          # dead
end

function SIRDmeta(du,u,p,t)  # SIRD metapopulation model (uncoupled countries)
  b,g,m,N = p
  j = 1
  for i=1:N
      S,I,R,D = u[j:j+3]
      Nc = S+I+R+D
      du[j] = dS = -b*S*I/Nc       # susceptible
      du[j+1] = dI = b*S*I/Nc - g*I # infected
      du[j+2] = dR = (1-m)*g*I      # recovered
      du[j+3] = dD = m*g*I          # dead
      j = j+4
  end
end

function nlys_discrete_SIR()
    Ti = 6.0         # infectious period
    gamma = 1. / Ti  # recovery rate
    m = 0.01         # case fatality rate
    b = 0.4          # contact rate

    Nc  = 50 * 1e6  # population size

    I0 = 1
    S0 = Nc - I0
    R0 = 0
    D0 = 0

    ndays = 75    # total runtime

    ndays = 80    # total runtime
    nB = 62   # runtime until states are saved
    nB1 = 80   # runtime until states are saved
    nB2 = 65   # runtime until states are saved
    nB3 = 50   # runtime until states are saved
    nB4 = 35   # runtime until states are saved
    nB5 = 20   # runtime until states are saved
    nB6 = 5   # runtime until states are saved

    N = 200       # max number of countries
    p = 0.0005    # invasion probability
    p = 0.0006    # invasion probability
    firsts = [0]

    # local growth rate
    r = b - gamma
    C_max = exp(r*(ndays+1))
    println("Cmax: ",C_max)

    alphaI = zeros(ndays)
    alphaD = zeros(ndays)
    C = zeros(0)
    D = zeros(0)
    N1 = 1
    u0 = [S0;I0;R0;D0]   # initial values
    tspan = (0.0,1.0)    # run the SIR-model for bursts of 1 day
    for i=1:ndays
       par = (b,gamma,m,N1)
       prob = ODEProblem(SIRDmeta,u0,tspan,par)
       sol = solve(prob)
       u0 = sol[:,end]

       for j=N1+1:N  # infections of other countries
           if minimum(rand(N1)) < p
               push!(u0,S0)
               push!(u0,I0)
               push!(u0,R0)
               push!(u0,D0)
               N1 += 1
               push!(firsts,i)  # invasion at time i
           end
       end
       # calculate time-dependent critical exponent by log-likelihood
       tmp = reshape(u0,(4,:)) # vector of infected
       Iv = tmp[2,:] # vector of infected
       Rv = tmp[3,:] # vector of recovered
       Dv = tmp[4,:] # vector of dead
       Cv = Iv .+ Rv .+ Dv   # all cases
       if i == nB
          C = Cv
          D = Dv
       end
       if i == nB1
           println("here")
          C1 = Cv
          D1 = Dv
       elseif i == nB2
          C2 = Cv
          D2 = Dv
       elseif i == nB3
          C3 = Cv
          D3 = Dv
       elseif i == nB4
          C4 = Cv
          D4 = Dv
       elseif i == nB5
          C5 = Cv
          D5 = Dv
       elseif i == nB6
          C6 = Cv
          D6 = Dv
       end
       ind = findall(x->x>=1,Dv)
       Dv = Dv[ind]
       ind = findall(x->x>=1,Cv)
       Cv = Cv[ind]
       if length(Cv) >= 5
          mu,st = estimate(Cv)
          alphaI[i] = mu
       end
       if length(Dv) >= 5
          mu,st = estimate(Dv)
          alphaD[i] = mu
       end
    end

   # ************ distribution   ********************************
   ind = findall(x->x>=0.5, C)
   C = sort(C[ind], rev=true)
   nC = length(C)
   println("maximum C: ",maximum(C))
   ex1,ey1 = log2histogram(C)
   x1 = log.(ex1)
   y1 = log.(ey1)
   pol = Polynomials.fit(x1,y1,1)
   muC=-coeffs(pol)[2]
   println("linear fit muC:",muC)
   pol = Polynomials.fit(x1,y1,1)
   ey1fit = exp.(pol.(x1))
   #alpha = 1 + length(C) / sum(log.(C)) # calculate exponent by log-likelihood
   #println("alpha ", alphaI[62])

   println("length D ", length( findall(x->x>=1,D )))
   ind = findall(x->x>=0.5, D)
   D = sort(D[ind], rev=true)
   nD = length(D)
   ex1D,ey1D = log2histogram(D)
   x1D = log.(ex1D)
   y1D = log.(ey1D)
   pol = Polynomials.fit(x1D,y1D,1)
   muD=-coeffs(pol)[2]
   println("linear fit muD: ",muD)
   pol = Polynomials.fit(x1D,y1D,1)
   ey1fitD = exp.(pol.(x1D))

    C1 = sort(C1[findall(x->x>0.5, C1)],rev=true)
    exC1,eyC1 = log2histogram(C1)
    nC1 = length(C1)

    C2 = sort(C2[findall(x->x>0.5, C2)],rev=true)
    exC2,eyC2 = log2histogram(C2)
    nC2 = length(C2)

    C3 = sort(C3[findall(x->x>0.5, C3)],rev=true)
    exC3,eyC3 = log2histogram(C3)
    nC3 = length(C3)

    C4 = sort(C4[findall(x->x>0.5, C4)],rev=true)
    exC4,eyC4 = log2histogram(C4)
    nC4 = length(C4)

    C5 = sort(C5[findall(x->x>0.5, C5)],rev=true)
    exC5,eyC5 = log2histogram(C5)
    nC5 = length(C5)

    #  *********** Arrival time histogram  ***********************
    println(N1, " countries infected")
    h = StatsBase.fit(Histogram, firsts,  nbins=100, closed=:left)
    h_weight = h.weights
    h_bins = h.edges[1][1:end-1]

    @. model(x, p) = p[1]*exp(x*p[2])
    p0 = [0.5, 0.5]
    expfit = curve_fit(model, collect(h_bins[1:62]), h_weight[1:62], p0)
    println("Spreading rate: ",coef(expfit))
    cf = coef(expfit)

   ax = PyPlot.axes([fig_x2, fig_y1, fig_width_x, fig_width_y])
       loglog(exC1,eyC1,"-", color="k")
       scatter(exC1,eyC1,s=50, color="k", label="arrival time: $nB1 days")

       loglog(exC2,eyC2,"-", color="r")
       scatter(exC2,eyC2,s=50, color="r", label="arrival time: $nB2 days")

       loglog(exC3,eyC3,"-", color="b")
       scatter(exC3,eyC3,s=50, color="b", label="arrival time: $nB3 days")
       loglog(exC4,eyC4,"-", color="g")
       scatter(exC4,eyC4,s=50, color="g", label="arrival time: $nB4 days")
       loglog(exC5,eyC5,"-", color="m")
       scatter(exC5,eyC5,s=50, color="m", label="arrival time: $nB5 days")
       legend(loc="lower left",fancybox="true", fontsize=12)
       xticks(fontsize=14)
       yticks([1e-8,1e-7,1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=14)
       xlim(0.8,1e8)
       ylim(5e-9,1)
       xlabel("cases, n",fontsize=16)
       ax1 = ax.transAxes
       text(-0.14, 0.95, "c)", fontsize=20, fontweight="bold", transform = ax1)

   ax = PyPlot.axes([fig_Ix3, fig_y1+fig_Iy, fig_Iwidth_x, fig_Iwidth_y]) # Inset: cumulative distribution
       loglog(C1,(1:nC1),"-", color="k")
       loglog(C2,(1:nC2),"-", color="r")
       loglog(C3,(1:nC3),"-", color="b")
       loglog(C4,(1:nC4),"-", color="g")
       loglog(C5,(1:nC5),"-", color="m")
       ylim(1,200)
       xlim(1,1e8)
       xticks(fontsize=12)
       yticks(fontsize=12)
       xlabel("cases, n", fontsize=14)
       ylabel(L"$N \cdot C(n)$", fontsize=14, labelpad=0)

   ax = PyPlot.axes([fig_x2, fig_y2, fig_width_x, fig_width_y])
       bar(h_bins,h_weight)
       plot(h_bins[1:63], cf[1]*exp.(cf[2]*h_bins[1:63]),"k", linewidth=2, linestyle="dashed")
       plot([62,62],[0,14],"-r",linewidth=2)
       ylabel("countries",fontsize=16, labelpad=5)
       xticks(fontsize=14)
       yticks(0:2:12,fontsize=14)
       xlim(0,ndays)
       ylim(0,14)
       xlabel("arrival time (days)",fontsize=16)
       ax1 = ax.transAxes
       text(-0.14, 0.95, "d)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix2, fig_y2+fig_Iy, fig_Iwidth_x1, fig_Iwidth_y1]) # Inset: cumulative distribution
       ind = findfirst(x->x>0, alphaI)
       plot(ind:ndays,alphaI[ind:ndays], linewidth=2, color="black")
       ind = findfirst(x->x>0, alphaD)
       plot(ind:ndays,alphaD[ind:ndays], linewidth=2, color="blue")
       plot([62,62],[1,2],"r-")
       xticks([0,20,40,60,80],fontsize=12)
       yticks([1,1.2,1.4,1.6,1.8,2.0],fontsize=12)
       xlim(10,ndays)
       ylim(1,1.8)
       ax.yaxis.grid()
       xlabel("time (days)", fontsize=14)
       ylabel(L"exponent, $\mu$", fontsize=14, labelpad=5)

   # **************************************************************************
   # plot distribution
   ax = PyPlot.axes([fig_x1, fig_y1, fig_width_x, fig_width_y])
       loglog(ex1,ey1,"-",linewidth=2)
       scatter(ex1,ey1,s=50,label="total detected cases")
       loglog(ex1,ey1fit,"k", label=L"regression line, $\mu=1.18$")
       xlim(0.8,1e7)
       xlabel("cases, n",fontsize=16)
       ylabel(L"frequency, $P_C(n)$",fontsize=16,labelpad=10)
       xticks(fontsize=14)
       yticks([1e-8,1e-7,1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=14)
       legend(bbox_to_anchor=(0.01, 0.25), loc="upper left", fancybox="true", fontsize=14)
       ax1 = ax.transAxes
       text(-0.2, 0.95, "a)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix1, fig_y1+fig_Iy, fig_Iwidth_x1, fig_Iwidth_y1]) # Inset: cumulative distribution
       cc = 10 .^ range(0,7, length=50)
       loglog(cc, cumulpdf.(cc,muC,5e6), color="black", linewidth=2)
       scatter(C,(1:nC)./nC,s=20)
       xticks(fontsize=12)
       yticks(fontsize=12)
       xlim(1,1e7)
       ylim(5e-3,1)
       xlabel("cases, n",fontsize=14)
       ylabel("C(n)",fontsize=14)

   ax = PyPlot.axes([fig_x1, fig_y2, fig_width_x, fig_width_y])
       loglog(ex1D,ey1D,"-",linewidth=2)
       scatter(ex1D,ey1D,s=50,label="deaths")
       loglog(ex1D,ey1fitD,"k", label=L"regression line, $\mu=1.31$")
       xticks(fontsize=14)
       yticks([1e-6,1e-5,1e-4,1e-3,1e-2,1e-1],fontsize=14)
       xlim(0.8,5e4)
       xlabel("cases, n",fontsize=16)
       ylabel(L"frequency, $P_D(n)$",fontsize=16,labelpad=10)
       legend(bbox_to_anchor=(0.01, 0.25), loc="upper left", fancybox="true", fontsize=14)
       ax1 = ax.transAxes
       text(-0.2, 0.95, "b)", fontsize=20, fontweight="bold", transform = ax1)
   ax = PyPlot.axes([fig_Ix1, fig_y2+fig_Iy, fig_Iwidth_x1, fig_Iwidth_y1]) # Inset: cumulative distribution
       cc = 10 .^ range(0,4.2, length=50)
       loglog(cc, cumulpdf.(cc,muD,5e4), color="black", linewidth=2)
       scatter(D,(1:nD)/nD,s=20)
       xticks(fontsize=12)
       yticks(fontsize=12)
       xlim(1,6e4)
       ylim(1e-2,1)
       xlabel("cases, n",fontsize=14)
       ylabel("C(n)",fontsize=14)

end

Random.seed!(12)

fig_Iy = 0.25
fig_Ix1 = 0.34
fig_Ix3 = 0.83
fig_Ix2 = 0.63
fig_Iwidth_y  = 0.12
fig_Iwidth_x  = 0.13
fig_Iwidth_y1 = 0.13
fig_Iwidth_x1 = 0.15

fig_sep_y   = 0.06
fig_sep_y   = 0.08
fig_width_y = 0.4
fig_width_x = 0.4
fig_x1 = 0.1
fig_x2 = 0.57
fig_y2=0.08
fig_y1=fig_y2 + fig_width_y + fig_sep_y

figure("distribution_meta1", figsize=(12,8))
clf()
nlys_discrete_SIR()
#savefig("../figures/meta_all.png")

end # module


module TensorAlgebra

using LinearAlgebra

export Tensor, Covector, VectorSpace, DualSpace, ProductSpace, TensorSpace, TensorProduct
export spaces, tensors, label, dimensions, field, degree, dual, domain, ×, ⊗

abstract type AbstractSpace{K,N} end

struct VectorSpace{K,L,D} <: AbstractSpace{K,1} end

struct DualSpace{K,L,D} <: AbstractSpace{K,1} end

abstract type AbstractProductSpace{K,N,S} <: AbstractSpace{K,N} end

struct ProductSpace{K,N,S} <: AbstractProductSpace{K,N,S} end

struct TensorSpace{K,N,S} <: AbstractProductSpace{K,N,S} end

abstract type AbstractTensor{K,N} <: AbstractArray{K,N} end

struct Tensor{K,N,D} <: AbstractTensor{K,N}
    array::Array{K,N}
end

const Covector{K,L,D} = Tensor{K,1,VectorSpace{K,L,D}}

struct TensorProduct{K,N,S} <: AbstractTensor{K,N}
    scalar::K
    tensors::S
end

Base.collect(t::Tensor) = t.array

field(::AbstractSpace{K}) where {K} = K

field(::AbstractTensor{K}) where {K} = K

field(::Vector{K}) where {K} = K

degree(::AbstractSpace{K,N}) where {K,N} = N

degree(::AbstractTensor{K,N}) where {K,N} = N

spaces(as::AbstractSpace{K,1}) where {K} = (as,)

spaces(::AbstractProductSpace{K,N,S}) where {K,N,S} = S

label(::Union{VectorSpace{K,L},DualSpace{K,L}}) where {K,L} = L

dimensions(::Union{VectorSpace{K,L,D},DualSpace{K,L,D}}) where {K,L,D} = D

dual(::VectorSpace{K,L,D}) where {K,L,D} = DualSpace{K,L,D}()

dual(::DualSpace{K,L,D}) where {K,L,D} = VectorSpace{K,L,D}()

function Covector(vs::VectorSpace{K,L,D}, a) where {K,L,D}
    size(a)[1] === D || throw(DimensionMismatch("Tried to create a covector of dimension $(size(a)[1]) in the vector space $(label(vs)) of dimension $(D)."))
    Tensor{K,1,TensorSpace{K,1,(VectorSpace{K,L,D}(),)}}(a)
end

function Vector(vs::VectorSpace{K,L,D}, a) where {K,L,D}
    size(a)[1] === D || throw(DimensionMismatch("Tried to create a vector of dimension $(size(a)[1]) in the vector space $(label(vs)) of dimension $(D)."))
    Tensor{K,1,TensorSpace{K,1,(DualSpace{K,L,D}(),)}}(a)
end

VectorSpace(L::Symbol,::Type{K},dim::Int) where {K,D} = VectorSpace{K,L,dim}()

ProductSpace(args::AbstractSpace{K,1}...) where {K} = ProductSpace{K,length(args),(args...,)}()

TensorSpace(args::Vararg{<:AbstractSpace{K,1}}) where {K} = TensorSpace{K,length(args),(args...,)}()

TensorSpace(::ProductSpace{K,N,S}) where {K,N,S} = TensorSpace{K,N,S}()

TensorSpace(ts::TensorSpace{K,N,S}) where {K,N,S} = ts

(::Type{PS})(args...) where {PS <: ProductSpace} = ProductSpace(args...)

(::Type{TS})(args...) where {TS <: TensorSpace} = TensorSpace(args...)

dual(ps::PS) where {PS <: AbstractProductSpace} = PS(dual.(spaces(ps))...)

Tensor(dom::D,a) where {K,N,D <: AbstractSpace{K,N}} = Tensor{K,N,typeof(TensorSpace(dom))}(a)

Tensor(::D,a) where {K,N,D <: TensorSpace{K,N}} = Tensor{K,N,D}(a)

TensorProduct(args::Tensor{K}...) where {K} = TensorProduct{K,sum(degree.(args)),typeof((args...,))}(one(K), (args...,))

TensorProduct(t::Tensor) = t

domain(::Tensor{K,N,D}) where {K,N,D} = D()

domain(tp::TensorProduct{K,N,S}) where {K,N,S} = TensorSpace{K,N,((spaces.(domain.(tensors(tp)))...)...,)}()

dimensions(at::AbstractTensor) = ((dimensions.(spaces(domain(at)))...)...,)

scalar(::Tensor{K}) where {K} = one(K)

scalar(tp::TensorProduct) = tp.scalar

tensors(t::Tensor) = t

tensors(tp::TensorProduct) = tp.tensors

function (f::AbstractTensor{K,2})(::typeof(-), x::Tensor{K,1}) where {K}
    dom = spaces(domain(f))
    codom = spaces(domain(x))
    dual(dom[2]) === codom[1] || throw(DomainError(x,"Domain mismatch: Expected $(dual(dom[2])). Got $(codom[1])."))
    Tensor(dom[1], collect(f) * collect(x))
end

function (f::AbstractTensor{K,2})(x::Tensor{K,1}, ::typeof(-)) where {K}
    dom = spaces(domain(f))
    codom = spaces(domain(x))
    dual(dom[1]) === codom[1] || throw(DomainError(x,"Domain mismatch: Expected $(dual(dom[1])). Got $(codom[1])."))
    Tensor(dom[2], (collect(x)' * collect(f)).parent)
end

function (f::AbstractTensor{K,N})(x::AbstractTensor{K,N}) where {K,N}
    dual(domain(f)) === domain(x) || throw(DomainError(x,"Domain mismatch: Expected $(dual(domain(f))). Got $(domain(x))."))
    dot(f, x)
end

(f::Tensor{K,N})(::Vararg{typeof(-),N}) where {K,N} = f

function (f::Tensor{K,1})(x::Tensor{K,1}) where {K}
    dual(domain(f)) === domain(x) || throw(DomainError(x,"Domain mismatch: Expected $(dual(domain(f))). Got $(domain(x))."))
    dot(f, x)
end

function (f::Tensor{K,N})(xs::Vararg{Tensor{K,1},N}) where {K,N}
    x = TensorProduct(xs...)
    dual(domain(f)) === domain(x) || throw(DomainError(x,"Domain mismatch: Expected $(dual(domain(f))). Got $(domain(x))."))
    dot(f, x)
end

function (f::AbstractTensor{K,N})(xs::Vararg{Union{typeof(-),AbstractTensor{K,1}},N}) where {K,N}
    isarg = xs .!== -
    indims = findall(isarg)
    x = TensorProduct(xs[indims]...)
    any(.!isarg) || return dot(f, x)
    outdims = findall(.!isarg)
    a = reshape(mapslices(f, dims=indims) do slice
        dot(x, slice)
    end,size(f)[outdims]...)
    d = TensorSpace(spaces(domain(f))[outdims]...)
    Tensor(d, a)
end

Base.size(t::Tensor) = size(t.array)

Base.size(t::TensorProduct) = ((size.(t.tensors)...)...,)

Base.getindex(t::Tensor,ix::Vararg{Int}) = getindex(t.array, ix...)

function Base.getindex(tp::TensorProduct, ix::Vararg{Int})
    value = one(field(tp))
    offset = 0
    for (i, tensor) in enumerate(tp.tensors)
        it = degree(tensor) === 1 ? ix[offset + 1] : ix[offset + 1:offset + degree(tensor)]
        value *= getindex(tensor, it...)
        offset += degree(tensor)
    end
    value * scalar(tp)
end

Base.:^(v::AbstractSpace,::typeof(*)) = dual(v)

Base.first(ts::AbstractProductSpace) = first(spaces(ts))

Base.last(ts::AbstractProductSpace) = last(spaces(ts))

Base.in(t::AbstractTensor{K}, vs::AbstractSpace{K}) where {K} = dual(domain(t)) === TensorSpace(vs)

function Base.:+(at1::AbstractTensor{K,N}, at2::AbstractTensor{K,N}) where {K,N}
    domain(at1) === domain(at2) || throw(DomainError(at2,"Domain mismatch: Expected $(domain(at1)). Got $(domain(at2))."))
    Tensor{K,N,typeof(domain(at1))}(collect(at1) + collect(at2))
end

function Base.:-(at1::AbstractTensor{K,N}, at2::AbstractTensor{K,N}) where {K,N}
    domain(at1) === domain(at2) || throw(DomainError(at2,"Domain mismatch: Expected $(domain(at1)). Got $(domain(at2))."))
    Tensor{K,N,typeof(domain(at1))}(collect(at1) - collect(at2))
end

function Base.:-(t::Tensor{K,N,D}) where {K,N,D}
    Tensor{K,N,D}(-t.array)
end

function Base.:-(tp::TensorProduct{K,N,S}) where {K,N,S}
    TensorProduct{K,N,S}(-tp.scalar, tp.tensors)
end

Base.:*(x::Number,t::Tensor{K,N,D}) where {K,N,D} = Tensor{K,N,D}(x * t.array)

Base.:*(t::Tensor{K,N,D},x::Number) where {K,N,D} = Tensor{K,N,D}(x * t.array)

Base.:*(t1::Tensor{K,1,D},t2::Tensor{K,1,D}) where {K,D} = Tensor{K,1,D}(t1.array .* t2.array)

Base.:*(x::Number,tp::TensorProduct{K,N,S}) where {K,N,S} = TensorProduct{K,N,S}(x * scalar(tp), tensors(tp))

Base.:*(tp::TensorProduct{K,N,S},x::Number) where {K,N,S} = x * tp

×(vs1::AbstractSpace{K,1},vs2::AbstractSpace{K,1}) where {K} = ProductSpace(vs1, vs2)

×(vs::AbstractSpace{K,1},ts::ProductSpace{K}) where {K} = ProductSpace(vs, spaces(ts)...)

×(ts::ProductSpace{K},vs::AbstractSpace{K,1}) where {K} = ProductSpace(spaces(ts)..., vs)

×(ts1::ProductSpace{K},ts2::ProductSpace{K}) where {K} = ProductSpace((spaces(ts1)..., spaces(ts2)...))

⊗(vs1::AbstractSpace{K,1},vs2::AbstractSpace{K,1}) where {K} = TensorSpace(vs1, vs2)

⊗(vs::AbstractSpace{K,1},ts::TensorSpace{K}) where {K} = TensorSpace(vs, spaces(ts)...)

⊗(ts::TensorSpace{K},vs::AbstractSpace{K,1}) where {K} = TensorSpace(spaces(ts)..., vs)

⊗(ts1::TensorSpace{K},ts2::TensorSpace{K}) where {K} = TensorSpace((spaces(ts1)..., spaces(ts2)...))

⊗(t1::Tensor{K,R1},t2::Tensor{K,R2}) where {K,R1,R2} = TensorProduct{K,R1 + R2,typeof((t1, t2))}(one(K), (t1, t2))

⊗(tp::TensorProduct{K,R1},t::Tensor{K,R2}) where {K,R1,R2} = TensorProduct{K,R1 + R2,typeof((tensors(tp)..., t))}(scalar(tp), (tensors(tp)..., t))

⊗(t::Tensor{K,R1},tp::TensorProduct{K,R2}) where {K,R1,R2} = TensorProduct{K,R1 + R2,typeof((t, tensors(tp)...))}(scalar(tp), (t, tensors(tp)...))

⊗(t1::TensorProduct{K,R1},t2::TensorProduct{K,R2}) where {K,R1,R2} = TensorProduct{K,R1 + R2,typeof((tensors(t1)..., tensors(t2)...))}(scalar(t1) * scalar(t2), (tensors(t1)..., tensors(t2)...))

⊗(t1::AbstractTensor,t2::AbstractTensor,ts::Vararg{<:AbstractTensor}) = TensorProduct(scalar(t1) * scalar(t2), t1 ⊗ t2, ts...)

Base.show(io::IO, ::VectorSpace{K,L}) where {K,L} = print(io, L)

# Base.show(io::IO, ::Type{VectorSpace{K,L}}) where {K,L} = print(io, L)

Base.show(io::IO, ::DualSpace{K,L}) where {K,L} = print(io, L, "^*")

# Base.show(io::IO, ::Type{DualSpace{K,L}}) where {K,L} = print(io, L, "^*")

# Base.show(io::IO, ::ProductSpace{K,N,S}) where {K,N,S} = print(io, join(S, " × "))

# Base.show(io::IO, ::Type{ProductSpace{K,N,S}}) where {K,N,S} = print(io, join(S, " × "))

# Base.show(io::IO, ::TensorSpace{K,N,S}) where {K,N,S} = print(io, join(S, " ⊗ "))

# Base.show(io::IO, ::Type{TensorSpace{K,N,S}}) where {K,N,S} = print(io, join(S, " ⊗ "))

end # module


const δθ = π / 18
const π_third = π / 3
const π_fourth = π / 4
const π_sixth = π / 6

"""
    η_B(θ_B::Real, θ_crit::Real)

Calculate B angle dependent efficiency component following Pais et. al. 2018, eq. 7, https://arxiv.org/pdf/1805.00128.pdf
"""
@inline function η_B(θ_B::Real, θ_crit::Real)
    (tanh((θ_crit - θ_B) / δθ) + 1) / 2
end

"""
    ηB_acc_p(θ_B::Real)

Magnetic field geometry dependent efficiency for protons at initial acceleration.
"""
function ηB_acc_p(θ_B::Real)
    η_B(θ_B, π_fourth)
end

"""
    ηB_reacc_p(θ_B::Real)

Magnetic field geometry dependent efficiency for protons at reacceleration.
"""
function ηB_reacc_p(θ_B::Real)
    η_B(θ_B, π_third)
end

"""
    ηB_acc_e(θ_B::Real)

Magnetic field geometry dependent efficiency for electrons at initial acceleration.
"""
function ηB_acc_e(θ_B::Real)
    1 - η_B(θ_B, π_fourth)
end

"""
    ηB_reacc_e(θ_B::Real)

Magnetic field geometry dependent efficiency for electrons at reacceleration.
"""
function ηB_reacc_e(θ_B::Real)
    1 - η_B(θ_B, π_sixth)
end


"""
**AtiyahBott** is a module containing an implementation of the Atiyah-Bott residue formula in the Julia language.
The theory and the algorithm behind the package is described in the paper
"Effective computations of the Atiyah-Bott formula" by Giosuè Muratore and
Csaba Schneider (https://arxiv.org/pdf/2105.11183.pdf).

The colorations (useful to speed up the code) are up to projective spaces of dimension 29.
The full list is here: https://github.com/mgemath/Colorations/.
"""
module AtiyahBott

using Combinatorics
using LightGraphs
using ProgressMeter
using Downloads
using CodecZlib
using Nemo

const current_dir = @__DIR__

include("Marked.jl")
include("GraphFunctions.jl")
include("EquivariantClasses.jl")
include("Checks.jl")
include("Main.jl")

end # module


"""
    contact_point
Find approximate contact point. Current available methods (cp) are `threshold' (i.e. the contact point is identified by fixing a force threshold) or `hertz' (i.e. the contact point is identified by fitting a hertz model to the curve).
"""
function contact_point(self::RheoTimeData, interface::Interface, cp::String, param::NamedTuple; savelog = true,  sec = :)

    contactmodels = Dict("threshold" => contact_threshold,
                         "hertz" => contact_hertz);

    data = self[interface];

    cp_index = contactmodels[cp](data[:f], data[:d], param);

    if cp_index != nothing
        t = self.t[cp_index:end] .- self.t[cp_index];
        d = data[:d][cp_index:end] .- data[:d][cp_index];
        f = data[:f][cp_index:end] .- data[:f][cp_index];
    else
        t = self.t[:]
        d = data[:d][:]
        f = data[:f][:]
    end

    ϵ,σ = interface.to_ϵσ(d, f);

    log = self.log == nothing ? nothing : [self.log; RheoLogItem( (type=:process, funct=:contact_point, params=(cp = cp, param = param ), keywords=(sec = sec, ) ),
    (comment="Contact point detection and data shift",) ) ]

 return RheoTimeData(σ, ϵ, t, log)

end

"""
    contact_threshold(f::Array{Float64,1}, d::Array{Float64,1}; _param::NamedTuple)

Applies force threshold to find approximate contact point.
"""
function contact_threshold(force::Array{RheoFloat}, disp::Array{RheoFloat}, param::NamedTuple)

    # get index of contact
    cp_index = findfirst(force .>= param.threshold);

    # return
    return cp_index

end

"""
    contact_hertz(f::Array{Float64,1}, δ::Array{Float64,1}; _param::Float64 = NaN64)

    Apply spherical Hertz contact model to find approximate contact point.
"""
function contact_hertz(f::Array{RheoFloat}, δ::Array{RheoFloat}, param::NamedTuple)

    # regularise displacement
    δ = δ .- minimum(δ);

    # radius of sphere
    R = param.R;

    # poisson ratio - assume incompressible
    ν = 0.5;

    # section of the curve to be used to identify the first approximation
    s = param.s

    # get approximate δ0 and YM values
    (init_δ₀, init_YM, tilt_approx, offset_approx) =  hertz_approx(f, δ, R, ν, s);

    # least squares fit
    opt = Opt(:LN_SBPLX, 4); # 4 parameters to fit

    bounds_param = 0.25 # +/- limits allowance from initial contact point guess
    lower_bounds!(opt, [init_δ₀*(1 - bounds_param), 0.0, -5.0, -5.0]);
    upper_bounds!(opt, [init_δ₀*(1 + bounds_param), 1e10, 5.0, 5.0]);

    xtol_rel!(opt,1e-5);

    min_objective!(opt, (params, grad) -> hertz_cost(f, δ, R, ν, params[1], params[2], params[3], params[4]));

    params_init = [init_δ₀, init_YM, tilt_approx, offset_approx];

    (minf, minx, ret) = optimize(opt, params_init);

    # return index of contact point
    δ₀_index = argmin(abs.(δ .- minx[1]));

    return (δ₀_index)

end

###########################################
# Convenience functions for contact_hertz #
###########################################
"""
Generate approximate Hertz fit a, b, E, cp to use as initial conditions in proper fit.
"""

function hertz_approx(f::Array{RheoFloat,1}, δ::Array{RheoFloat,1}, R::RheoFloat, ν::RheoFloat, s::RheoFloat)

    # approximate contact point, assume linear near end of approach section
    # and trace line down to axis
    deriv = RHEOS.derivCD(RHEOS.derivCD(f, δ),δ);
    value_max, index_max = findmax(deriv);

    # #coeff = (1-index_max/length(f))
    last_section = round(Int64, s*length(f)):length(f)

    avg_gradient = sum(RHEOS.derivBD(f[last_section], δ[last_section]))/length(last_section);

    c = f[end] - avg_gradient*δ[end];

    approx_δ₀ = -c/avg_gradient;

    # transform to find approx YM
    δᵦ = (heaviside(δ .- approx_δ₀).*δ).^(3/2);

    prefactor = (4/3)*sqrt(R)/(1 - ν^2);

    avg_gradient_transformed = sum(RHEOS.derivBD(f[last_section], δᵦ[last_section]))/length(last_section);

    approx_YM = avg_gradient_transformed/prefactor;

    # estimate offset and tilt
    first_section = 1:round(Int64, 0.5*length(f));

    df_dh = RHEOS.derivBD(f, δ);

    tilt_approx = sum(df_dh[first_section])/length(first_section);

    offset_approx = sum(f[first_section])/length(first_section);

    if approx_δ₀<0
        approx_δ₀ = 0
    end

    if approx_YM<0
        approx_YM = 50
    end

    # return
    approx_δ₀, approx_YM, tilt_approx, offset_approx;
end

"""
Heaviside function
"""
function heaviside(x)
   0.5*(sign.(x) .+ 1);
end

"""
Hertz spherical model
"""
function hertz_model(δ::Array{RheoFloat,1}, R::RheoFloat, ν::RheoFloat, δ₀::RheoFloat, E::RheoFloat, tilt::RheoFloat, offset::RheoFloat)
    # define prefactor for Hertz sphere model
    prefactor = (4/3)*sqrt(R)/(1 - ν^2);

    δ_shift = δ - δ₀;

    # force
    F = tilt*δ + offset + prefactor*E*((δ_shift.*heaviside(δ_shift)).^(3/2));
end

"""
Hertz cost function
"""
function hertz_cost(f::Array{RheoFloat,1}, δ::Array{RheoFloat,1}, R::RheoFloat, ν::RheoFloat, δ₀::RheoFloat, E::RheoFloat, tilt::RheoFloat, offset::RheoFloat)
    # difference
    Δ = f - hertz_model(δ, R, ν, δ₀, E, tilt, offset);

    # regularisation parameter for offset and tilt
    λ = 0.1;

    # sum of squares
    cost = sum(Δ.^2); # Maybe not needed + λ*(tilt + offset)^2
end


import Base.Iterators:product

function read_puzzle(filename)

    open(filename) do f
         return parse.(Int, readlines(f))
    end

end

function ex1( puzzle )
    for item in product(puzzle,puzzle)
        if length(unique(item)) == 2
            sum(item) == 2020  && return prod(item)
        end
    end
end

function ex2( puzzle )
    for item in product(puzzle, puzzle, puzzle)
        if length(unique(item)) == 3
            sum(item) == 2020 && return prod(item)
        end
    end
end

puzzle = read_puzzle("input01.txt")
println(ex1(puzzle))
println(ex2(puzzle))


read_data() = readlines("day08.txt")

# read_data() = split("""
# nop +0
# acc +1
# jmp +4
# acc +3
# jmp -3
# acc -99
# acc +1
# jmp -4
# acc +6""", "\n")

function parse_instructions()
    lines = read_data()
    instructions = split.(lines, " ")
    return map(x -> (Symbol(x[1]), parse(Int, x[2])), instructions)
end

function execute(instructions)
    val = 0
    len = length(instructions)
    visited = falses(len)
    i = 1
    @inbounds while true
        # needed by Part2
        i > len && return (val, :terminated)

        # Next instruction
        instr = instructions[i]
        visited[i] && return (val, :infinite)

        # Remember this instruction
        # push!(visited, i)
        visited[i] = true

        # Execute!
        if instr[1] === :acc
            val += instr[2]
        elseif instr[1] === :jmp
            i += instr[2] - 1
        end
        i += 1
    end
end

part1() = execute(parse_instructions())

# part 2 brute force

function part2()
    prog = parse_instructions()
    sz = length(prog)
    jmps = findall(x -> x[1] == :jmp, prog)
    for j in jmps
        np = modified_program(prog, j, :nop)
        result = execute(np)
        result[2] == :terminated && return result[1]
    end
    # NOTE:
    # I didn't have to complete the code by replacing nop with jmp because I have
    # already reached an answer. Otherwise, it would be simialr code as follows.
    # nops = [i for i in 1:sz if prog[i][1] == :nop && prog[i][2] != 0]
    return "not found"
end

function modified_program(prog, i, chg)
    prog = copy(prog)
    prog[i] = (chg, prog[i][2])
    return prog
end


# Methods for checking matrix sizes etc.
import LinearAlgebra: issymmetric, eigmin

issquare(x::Matrix) = size(x, 1) == size(x, 2)
ispossemidef(x::Matrix) = issymmetric(x) && (eigmin(x) >= 0)

function confirm_matrix_sizes(F, B, V, G, W)

    nx = size(F, 1)
    nu = size(B, 2)
    ny = size(G, 1)

    @assert size(F) == (nx, nx)
    @assert size(B) == (nx, nu)
    @assert size(V) == (nx, nx)

    @assert size(G) == (ny, nx)
    @assert size(W) == (ny, ny)

    return nx, ny, nu

end


"""
    step!(new, old, ::DistributionStrategy)

Perform one iteration of Conway's Game of Life.
"""
step!

"""
    prepare!(grid)

Prepare the grid to be used as the old/current world stage for `step!`.
Assuming that only inner cells, `grid[2:end-1,2:end-1]`, contain actual
world information, mirror the edges of that grid.
"""
function prepare!(grid)
    m, n = size(grid)
    @views begin
        grid[1,2:n-1] .= grid[m-1,2:n-1]
        grid[m,2:n-1] .= grid[2,2:n-1]
        # Now we are good to include the corners as well:
        grid[1:m,1] .= grid[1:m,n-1]
        grid[1:m,n] .= grid[1:m,2]
    end
    nothing
end

function life_rule!(new, old, newindices, oldindices)
    neighbours = [CartesianIndex(x,y) for x in -1:1, y in -1:1 if x!=0 || y!=0]
    for (i,j) in zip(CartesianIndices(oldindices), CartesianIndices(newindices))
        nc = sum(old[i+n] for n in neighbours)
        if old[i]
            new[j] = 2 <= nc <= 3
        else
            new[j] = nc == 3
        end
    end
    nothing
end

interior(r::OrdinalRange) = r[2:end-1]
interior(t::Tuple) = map(interior, t)
interior(m::AbstractMatrix) = @view m[interior(axes(m))...]


"""
    EconometricsModel <: RegressionModel

Abstract type provided by Econometrics.jl
"""
abstract type EconometricsModel <: RegressionModel end
implicit_intercept(::Type{<:EconometricsModel}) = true
# Estimators
"""
    ModelEstimator

Abstract type for model estimators.
"""
abstract type ModelEstimator end
"""
    LinearModelEstimators

Abstract type for linear model estimators.
"""
abstract type LinearModelEstimators <: ModelEstimator end
"""
    ContinuousResponse(groups::Vector{Vector{Vector{Int}}}) <: LinearModelEstimators

Continuous response estimator with potential features absorption.
"""
struct ContinuousResponse <: LinearModelEstimators
    groups::Vector{Vector{Vector{Int}}}
end
show(io::IO, estimator::ContinuousResponse) = println(io, "Continuous Response Model")
"""
    BetweenEstimator(effect::Symbol,
                     groups::Vector{Vector{Int}}) <: LinearModelEstimators

Continuous response estimator collapsing a dimension in a longitudinal setting.
"""
struct BetweenEstimator <: LinearModelEstimators
    effect::Symbol
    groups::Vector{Vector{Int}}
end
function show(io::IO, estimator::BetweenEstimator)
    @unpack effect, groups = estimator
    println(io, "Between Estimator")
    println(io, "$effect with $(sum(length, groups)) groups")
    T = length.(groups)
    if !isempty(T)
        E = extrema(T)
        if E[1] ≈ E[2]
            println(io, "Balanced groups with size $(E[1])")
        else
            T̄ = round(harmmean(T), digits = 2)
            println(io, "Groups of sizes [$(E[1]), $(E[2])] with harmonnic mean of $(T̄)")
        end
    end
end
"""
    RandomEffectsEstimator(pid::Tuple{Symbol,Vector{Vector{Int}}},
                           tid::Tuple{Symbol,Vector{Vector{Int}}},
                           idiosyncratic::Float64,
                           individual::Float64,
                           θ::Vector{Float64}) <: LinearModelEstimators

Swamy-Arora estimator.
"""
struct RandomEffectsEstimator <: LinearModelEstimators
    pid::Tuple{Symbol,Vector{Vector{Int}}}
    tid::Tuple{Symbol,Vector{Vector{Int}}}
    idiosyncratic::Float64
    individual::Float64
    θ::Vector{Float64}
    RandomEffectsEstimator(pid, tid) =
        new((pid, Vector{Vector{Int}}()),
            (tid, Vector{Vector{Int}}()), NaN, NaN, zeros(0))
    function RandomEffectsEstimator(pid, tid, X, y, z, Z, wts)
        Xbe, ybe, βbe, Ψbe, ŷbe, wtsbe, pivbe =
            solve(BetweenEstimator(pid[1], pid[2]), X, y, z, Z, wts)
        Xfe, yfe, βfe, Ψfe, ŷfe, wtsfe, pivfe =
            solve(ContinuousResponse([pid[2]]), X, y, z, Z, wts)
        ivk = size(Z, 2)
        T = length.(pid[2])
        T̄ = harmmean(T)
        σₑ² = sum(wᵢ * (yᵢ - ŷᵢ)^2 for (wᵢ, yᵢ, ŷᵢ) ∈ zip(wtsfe, yfe, ŷfe)) /
        	(sum(wtsfe) - length(βfe) - ivk - length(pid[2]) + 1)
        σᵤ² = max(0, sum((yᵢ - ŷᵢ)^2 for (yᵢ, ŷᵢ) ∈ zip(ybe, ŷbe)) /
                     (length(ybe) - length(βbe) - ivk) - σₑ² / T̄)
        θ = 1 .- sqrt.(σₑ² ./ (T .* σᵤ² .+ σₑ²))
        new(pid, tid, √σₑ², √σᵤ², θ)
    end
end
function show(io::IO, estimator::RandomEffectsEstimator)
    pid, D = estimator.pid
    tid, T = estimator.tid
    @unpack idiosyncratic, individual = estimator
    println(io, "One-way Random Effect Model")
    println(io, "Longitudinal dataset: $pid, $tid")
    L = length.(D)
    if !isempty(L)
        E = extrema(L)
        if E[1] ≈ E[2]
            println(io, "Balanced dataset with $(length(L)) panels of length $(E[1])")
        else
            T̄ = round(harmmean(L), digits = 2)
            println(io, "Unbalanced dataset with $(length(L)) panels")
            println(io, "Panel lengths [$(E[1]), $(E[2])] with harmonnic mean of $(T̄)")
        end
    end
    println(io, "individual error component: $(round(individual, digits = 4))")
    println(io, "idiosyncratic error component: $(round(idiosyncratic, digits = 4))")
    println(io, "ρ: $(round(individual^2 / (individual^2 + idiosyncratic^2), digits = 4))")
end
"""
    NominalResponse(categories::ContrastsMatrix) <: ModelEstimator

Multinomial logistic regression.
"""
struct NominalResponse{T} <: ModelEstimator
    categories::Vector{T}
    function NominalResponse(obj::ContrastsMatrix)
        categories = isnothing(obj.contrasts.base) ?
            obj.levels :
            union(vcat(obj.contrasts.base, obj.levels))
        new{eltype(categories)}(categories)
    end
end
function show(io::IO, estimator::NominalResponse)
    println(io, "Probability Model for Nominal Response")
    println(io, "Categories: $(join(estimator.categories, ", "))")
end
"""
    OrdinalResponse(categories::ContrastsMatrix) <: ModelEstimator

Proportional odds logistic regression.
"""
struct OrdinalResponse{T} <: ModelEstimator
    categories::Vector{T}
    function OrdinalResponse(obj::ContrastsMatrix)
        categories = isnothing(obj.contrasts.levels) ?
            obj.levels :
            obj.contrasts.levels
        new{eltype(categories)}(categories)
    end
end
function show(io::IO, estimator::OrdinalResponse)
    println(io, "Probability Model for Ordinal Response")
    println(io, "Categories: $(join(estimator.categories, " < "))")
end
"""
    VCE

Variance-covariance estimators.

- Observed Information Matrix (OIM)
- Heteroscedasticity Consistent: HC0, HC1, HC2, HC3, HC4
"""
@enum VCE begin
    OIM = -1
    HC0 = 0
    HC1 = 1
    HC2 = 2
    HC3 = 3
    HC4 = 4
end


#
# SubGradient Plan
#
export SubGradientMethodOptions, SubGradientProblem
export getCost, getSubGradient
export getStepsize
#
# Problem
@doc doc"""
    SubGradientProblem <: Problem
A structure to store information about a subgradient based optimization problem

# Fields
* `M` – a [`Manifold`](@ref)
* `costFunction` – the function $F$ to be minimized
* `subGradient` – a function returning a subgradient $\partial F$ of $F$
"""
mutable struct SubGradientProblem{mT <: Manifold} <: Problem
    M::mT
    costFunction::Function
    subGradient::Function
end
"""
    getSubGradient(p,x)

evaluate the gradient of a [`SubGradientProblem`](@ref)` p` at the [`MPoint`](@ref) `x`.
"""
getSubGradient(p::P,x::MP) where {P <: SubGradientProblem{M} where M <: Manifold, MP <: MPoint} = p.subGradient(x)
"""
    SubGradientMethodOptions <: Options
stories option values for a [`subGradientMethod`](@ref) solver

# Fields
* `retraction` – the retration to use within
* `stepsize` – a [`Stepsize`](@ref)
* `stop` – a [`StoppingCriterion`](@ref)
* `x` – (initial or current) value the algorithm is at
* `optimalX` – optimal value
"""
mutable struct SubGradientMethodOptions{P,T} <: Options where {P <: MPoint, T <: TVector}
    retraction::Function
    stepsize::Stepsize
    stop::StoppingCriterion
    x::P
    xOptimal::P
    ∂::T
    SubGradientMethodOptions{P,T}(x::P,sC::StoppingCriterion,s::Stepsize,retr::Function=exp) where {P <: MPoint, T <: TVector} = (
        o = new{P,T}(); o.x = x; o.xOptimal = x;
        o.stepsize = s; o.retraction = retr;
        o.stop = sC;
        return o
    )
end
SubGradientMethodOptions(x::P,sC::StoppingCriterion,s::Stepsize,retr::Function=exp) where {P <: MPoint} = SubGradientMethodOptions{P,typeofTVector(x)}(x,sC,s,retr)


export @nview, apy2, nrm2, swap!, scal!, axpy!, xpby!, axpby!

function mach_acc() # @code_warntype ✓
    sfmin = nextfloat(0.0)
    small = inv(prevfloat(Inf))
    if small ≥ sfmin
        # Use "small" plus 1 bit, to avoid the possibility of rounding
        # causing overflow when computing 1/sfmin.
        sfmin = small * (1.0 + eps())
    end
    return sfmin
end

const SAFMIN = mach_acc() / eps()

macro nview(arr::Symbol, N::Int)
    a = Vector{Expr}(undef, N)
    @inbounds for i in 1:N
        a[i] = Expr(:(=), Symbol("k", i), :(view($arr, :, $i)))
    end
    return Expr(:escape, Expr(:block, a...))
end

macro nview(arr::Symbol, vars::Symbol...)
    N = length(vars)
    a = Vector{Expr}(undef, N)
    @inbounds for i in 1:N
        a[i] = Expr(:(=), vars[i], :(view($arr, :, $i)))
    end
    return Expr(:escape, Expr(:block, a...))
end

function apy2(x::Real, y::Real) # @code_warntype ✓
    isnan(x) && return x
    isnan(y) && return y
    # general case
    xabs = abs(x)
    yabs = abs(y)
    w = max(xabs, yabs)
    z = min(xabs, yabs)
    iszero(z) && return w
    return w * sqrt(1.0 + abs2(z / w))
end

function swap!(v::VecI, i::Int, j::Int) # @code_warntype ✓
    @inbounds temp = v[i]
    @inbounds v[i] = v[j]
    @inbounds v[j] = temp
    return nothing
end

function swap!(v::MatI, i1::Int, j1::Int, i2::Int, j2::Int) # @code_warntype ✓
    @inbounds temp     = v[i1,j1]
    @inbounds v[i1,j1] = v[i2,j2]
    @inbounds v[i2,j2] = temp
    return nothing
end

function nrm2(x::VecI{Tx}, y::VecI{Ty}, b::VecB{Tb}) where {Tx<:Real,Ty<:Real,Tb<:Real} # @code_warntype ✓
    @simd for i in eachindex(b)
        @inbounds b[i] = abs2(x[i] - y[i])
    end
    return sqrt(sum(b))
end

function scal!(a::Real, x::VecO{Tx}) where Tx<:Real # @code_warntype ✓
    isone(a) && return nothing
    if iszero(a)
        zeroTx = zero(Tx)
        @simd for i in eachindex(x)
            @inbounds x[i] = zeroTx
        end
        return nothing
    end
    @simd for i in eachindex(x)
        @inbounds x[i] *= a
    end
end

# a * x + y → y, @code_warntype ✓
function axpy!(a::Real, x::VecI{Tx}, y::VecO{Ty}) where {Tx<:Real,Ty<:Real}
    iszero(a) && return nothing
    if isone(a)
        @simd for i in eachindex(y)
            @inbounds y[i] += x[i]
        end
        return nothing
    end
    @simd for i in eachindex(y)
        @inbounds y[i] += a * x[i]
    end
end

# x + b * y → y, @code_warntype ✓
function xpby!(x::VecI{Tx}, b::Real, y::VecO{Ty}) where {Tx<:Real,Ty<:Real}
    if iszero(b)
        @simd for i in eachindex(y)
            @inbounds y[i] = x[i]
        end
        return nothing
    end
    if isone(b)
        @simd for i in eachindex(y)
            @inbounds y[i] += x[i]
        end
        return nothing
    end
    @simd for i in eachindex(y)
        @inbounds y[i] = x[i] + b * y[i]
    end
end

# a * x + b * y → y, @code_warntype ✓
function axpby!(a::Real, x::VecI{Tx}, b::Real, y::VecO{Ty}) where {Tx<:Real,Ty<:Real}
    isone(b) && return axpy!(a, x, y)
    if iszero(b)
        iszero(a) && return scal!(0, y)
        if isone(a)
            @simd for i in eachindex(y)
                @inbounds y[i] = x[i]
            end
        else
            @simd for i in eachindex(y)
                @inbounds y[i] = a * x[i]
            end
        end
        return nothing
    end
    iszero(a) && return scal!(b, y)
    if isone(a)
        @simd for i in eachindex(y)
            @inbounds y[i] = x[i] + b * y[i]
        end
    else
        @simd for i in eachindex(y)
            @inbounds y[i] = a * x[i] + b * y[i]
        end
    end
end


"""
x = fnnls(AtA, Atb; ...)

Returns x that solves A*x = b in the least-squares sense, subject to x >=0. The
inputs are the cross-products AtA = A'*A and Atb = A'*b. Uses the modified
active set method of Bro and De Jong (1997).

Optional arguments:
    tol: tolerance for nonnegativity constraints
    max_iter: maximum number of iterations (counts inner loop iterations)

References:
    Bro R, De Jong S. A fast non-negativitity-constrained least squares
    algorithm. Journal of Chemometrics. 11, 393–401 (1997)
"""
function fnnls(AtA, Atb::AbstractVector{T}; kwargs...) where {T}
    work = fnnls_work(AtA, Atb)
    fnnls!(work, AtA, Atb; kwargs...)
end

function fnnls_work(AtA, Atb::AbstractVector{T}) where {T}
    n    = size(AtA,1)
    x    = zeros(T, n)
    s    = zeros(T, n)
    w    = similar(x)
    P    = BitArray(undef, n)
    nP   = BitArray(undef, n)
    iP   = BitArray(undef, n)
    AtAP = similar(AtA)
    AtbP = similar(Atb)
    return (x=x, s=s, w=w, P=P, nP=nP, iP=iP, AtAP=AtAP, AtbP=AtbP)
end

function fnnls!(
        work,
        AtA,
        Atb::AbstractVector{T};
        tol::T = sqrt(eps(T)),
        max_iter::Int = 30 * size(AtA, 2)
    ) where {T}

    x, s, w, P, nP, iP, AtAP, AtbP = work
    n = size(AtA,1)
    x .= 0
    s .= 0

    # P is a bool array storing positive elements of x
    # i.e., x[P] > 0 and x[~P] == 0
    P  .= x .> tol
    nP .= .!(P)

    mul!(w, AtA, x)
    w .= Atb .- w; #w = Atb - AtA*x

    function anyP(f, x, P)
        @inbounds for i in eachindex(x,P)
            if P[i] && f(x[i])
                return true
            end
        end
        return false
    end

    function argmaxP(x, P)
        α = T(-Inf)
        out = 1
        @inbounds for i in eachindex(x, P)
            if P[i] && (x[i] > α)
                α = max(α, x[i])
                out = i # i = argmax(w .* nP)
            end
        end
        return out
    end

    function assigntoP!(dst, srcP, P)
        j = 0
        @inbounds for i in eachindex(dst, P)
            if P[i]
                j += 1
                dst[i] = srcP[j]
            end
        end
        return dst
    end

    function copytocorner!(dst, src, P)
        j1, j2 = 0, 0
        @inbounds for i2 in 1:size(src, 2)
            if P[i2]
                j2 += 1
                j1  = 0
                for i1 in 1:size(src, 1)
                    if P[i1]
                        j1 += 1
                        dst[j1,j2] = src[i1,i2]
                    end
                end
            end
        end
        return j1, j2
    end

    function copytostart!(dst, src, P)
        j1, j2, idx = 0, 0, 0
        @inbounds for i2 in 1:size(src, 2)
            if P[i2]
                j2 += 1
                j1  = 0
                for i1 in 1:size(src, 1)
                    if P[i1]
                        j1  += 1
                        idx += 1
                        dst[idx] = src[i1,i2]
                    end
                end
            end
        end
        return j1, j2
    end

    function copytotop!(dst, src, P)
        j1 = 0
        @inbounds for i1 in 1:size(src, 1)
            if P[i1]
                j1 += 1
                dst[j1] = src[i1]
            end
        end
        return j1
    end

    function alpha(x, s, P)
        α = oftype(Inf, one(eltype(x))/one(eltype(s)))
        @inbounds for i in eachindex(x, s, P)
            if P[i]
                α = min(α, x[i]/(x[i] - s[i]))
            end
        end
        return α
    end

    function solveP!(y, Atmp, ytmp, A, x, P) # s[P] = AtA[P,P] \ Atb[P]
        # (1)
        # y[P] .= A[P,P] \ x[P] # 644 - 721 us
        # (2)
        # assigntoP!(y, Symmetric(A[P,P]) \ x[P], P) # 723-836 us
        assigntoP!(y, A[P,P] \ x[P], P) # 611-709 us
        # (3)
        # invA = cholesky!(Symmetric(A[P,P]); check = false) # 650 us
        # invA = cholesky!(A[P,P]; check = false) # 650-850 ua
        # assigntoP!(y, ldiv!(invA, x[P]), P)
        # (4)
        # m_, n_ = copytocorner!(Atmp, A, P)
        # p_ = copytotop!(ytmp, x, P)
        # invA = cholesky!(Symmetric(view(Atmp, 1:m_, 1:n_)); check = false)
        # outP = ldiv!(invA, view(ytmp, 1:p_)) # 630-720 us
        # assigntoP!(y, outP, P)
        # (5)
        # m_, n_ = copytostart!(Atmp, A, P)
        # p_ = copytotop!(ytmp, x, P)
        # Atmp_view = reshape(view(Atmp, 1:m_*n_), m_, n_)
        # # invA = cholesky!(Atmp_view; check = false) # 660 - 730 us
        # invA = cholesky!(Symmetric(Atmp_view); check = false) # 650 - 715 us
        # outP = ldiv!(invA, view(ytmp, 1:p_))
        # assigntoP!(y, outP, P)
    end

    # We have reached an optimum when either:
    #   (a) all elements of x are positive (no nonneg constraints activated)
    #   (b) ∂f/∂x = A' * (b - A*x) > 0 for all nonpositive elements of x
    iter = 0
    while sum(P) < n && anyP(>(tol), w, nP) && iter < max_iter

        # find i that maximizes w, restricting i to indices not in P
        # Note: the while loop condition guarantees at least one w[~P]>0
        i = argmaxP(w, nP) # i = argmax(w .* nP)

        # Move i to P
        P[i]  = true
        nP[i] = false

        # Solve least-squares problem, with zeros for columns/elements not in P
        solveP!(s, AtAP, AtbP, AtA, Atb, P) # s[P] = AtA[P,P] \ Atb[P]
        @inbounds for i in eachindex(s, nP)
            nP[i] && (s[i] = 0) # s[nP] .= zero(eltype(s)) # zero out elements not in P
        end

        # Inner loop: deal with negative elements of s
        while anyP(<=(tol), s, P)
            iter += 1

            # find indices in P where s is negative
            @inbounds iP .= (s .<= tol) .& P

            # calculate step size, α, to prevent any xᵢ from going negative
            α = alpha(x, s, iP) # minimum(x[iP] ./ (x[iP] .- s[iP]))

            # update solution (pushes some xᵢ to zero)
            @inbounds x .+= α .* (s .- x)

            # Remove all i in P where x[i] == 0
            @inbounds for i = 1:n
                if P[i] && abs(x[i]) < tol
                    P[i]  = false # remove i from P
                    nP[i] = true
                end
            end

            # Solve least-squares problem again, zeroing nonpositive columns
            solveP!(s, AtAP, AtbP, AtA, Atb, P) # s[P] = AtA[P,P] \ Atb[P]
            @inbounds for i in eachindex(s, nP)
                nP[i] && (s[i] = 0) # s[nP] .= zero(eltype(s)) # zero out elements not in P
            end
        end

        # update solution
        x .= s
        mul!(w, AtA, x)
        w .= Atb .- w; #w = Atb - AtA*x
    end

    return x
end

function fnnls(A,
               B::AbstractMatrix;
               gram::Bool = false,
               use_parallel::Bool = true,
               kwargs...)

    n = size(A,2)
    k = size(B,2)

    if gram
        # A,B are actually Gram matrices
        AtA = A
        AtB = B
    else
        # cache matrix computations
        AtA = A'*A
        AtB = A'*B
    end

    if use_parallel && nprocs()>1
        X = @distributed (hcat) for i=1:k
            fnnls(AtA, AtB[:,i]; kwargs...)
        end
    else
        X = Array{eltype(B)}(undef,n,k)
        for i = 1:k
            X[:,i] = fnnls(AtA, AtB[:,i]; kwargs...)
        end
    end

    return X
end


#=
   Binomial Trial Count.

   From infer.net test/Tests/BlogTests.cs
   """
   // Example of inferring the size of a population.
   // Reference: A. Raftery, "Inference for the binomial N parameter: A hierarchical Bayes approach", Biometrika 1988
   // http://pluto.huji.ac.il/~galelidan/52558/Material/Raftery.pdf

   ....
   ExpectationPropagation
   impala
   theta = Beta(80.24,91.17)[mean=0.4681]
   N mode = 44, median = 45
   waterbuck
   theta = Beta(238,233.8)[mean=0.5044]
   N mode = 124, median = 125
   impala
   N mode = 37, median = 65
   waterbuck
   N mode = 122, median = 208

   VariationalMessagePassing
   impala
   theta = Beta(106,370.9)[mean=0.2223]
   N mode = 94, median = 94
   waterbuck
   theta = Beta(316,393)[mean=0.4457]
   N mode = 141, median = 141
   impala
   N mode = 37, median = 65
   waterbuck
   N mode = 122, median = 208
   """

   Model 1: Impala
   Summary Statistics
   parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

     beta[1]   43.9012   26.9240     0.2692    1.6380   232.9437    1.0011      110.3999
     beta[2]   59.2836   25.3251     0.2533    1.2813   328.1459    1.0000      155.5194
       theta    0.4091    0.1625     0.0016    0.0090   272.1842    1.0010      128.9972
           n   73.4049   87.6808     0.8768    4.6531   335.0041    1.0000      158.7697

   Model 2: Waterbuck
   Summary Statistics
   parameters       mean        std   naive_se      mcse        ess      rhat   ess_per_sec
      Symbol    Float64    Float64    Float64   Float64    Float64   Float64       Float64

     beta[1]    42.1176    25.6791     0.2568    1.7376   152.0795    1.0034       90.3085
     beta[2]    63.5964    24.1648     0.2416    1.7421   106.0723    1.0363       62.9883
       theta     0.3738     0.1557     0.0016    0.0104   164.0217    1.0003       97.4000
           n   215.9076   134.5262     1.3453    9.2679   136.5959    1.0000       81.1139

   Cf ~/webppl/binomial_trial_count.wppl

=#

using Turing
using Plots, StatsPlots
include("jl_utils.jl")

@model function binomial_trial_count(data)
    maxN = 1000
    beta ~ filldist(Uniform(0.1,100),2)
    theta ~ Beta(beta[1],beta[2])

    n ~ Categorical(simplex([1/i for i in 1:maxN]))
    data ~ filldist(Binomial(n, theta), length(data))

end

impala = [15,20,21,23,26]
waterbuck = [53,57,66,67,72]
println("Model 1: Impala")
model = binomial_trial_count(impala)

num_chains = 4
# chs = sample(model, Prior(), 1000)
# chs = sample(model, MH(), 10_000)
# chs = sample(model, PG(5), 10_000)
# chs = sample(model, IS(), 10_000)
chs = sample(model, SMC(), 10_000)
# chs = sample(model, SMC(), MCMCThreads(), 10_000, num_chains)

# chs = sample(model, SGLD(), 10_000)
# chs = sample(model,NUTS(), 10_000)
# chs = sample(model,HMC(0.01,5), 10_000)

display(chs)
# display(plot(chs))

println("\n\nModel 2: Waterbuck")
model = binomial_trial_count(waterbuck)

num_chains = 4
# chs = sample(model, Prior(), 1000)
# chs = sample(model, MH(), 10_000)
# chs = sample(model, PG(5), 10_000)
# chs = sample(model, IS(), 10_000)
chs = sample(model, SMC(), 10_000)
# chs = sample(model, SMC(), MCMCThreads(), 10_000, num_chains)

# chs = sample(model, SGLD(), 10_000)
# chs = sample(model,NUTS(), 10_000)
# chs = sample(model,HMC(0.01,5), 10_000)

display(chs)
# display(plot(chs))


using Base.Test
using EGR

println("TestM1P1")

a = MinusPlusOneVector([1,0,-1],Set([-1]))
b = MinusPlusOneVector([-1,-1,1])
@test a.field == b.field
@test a.numPlus == b.numPlus
@test_throws ArgumentError MinusPlusOneVector([1,1,-2])
@test_throws ArgumentError MinusPlusOneVector([1.0,1,-2])

println("TestM1P1 successful!")
println()


brng_type = VSL_BRNG_ARS5
info("Testing discrete distributions with $brng_type...")
brng = BasicRandomNumberGenerator(brng_type, 200701281)

stdout_ = STDOUT
pwd_ = pwd()
cd(joinpath(Pkg.dir("VSL"), "test/"))
mkpath("./actual")
outfile = open("./actual/discrete.out", "w")
redirect_stdout(outfile)

t_uniform = UniformDiscrete(brng, Cint(1), Cint(10))
@test_rand(rand(t_uniform))
@test_rand(rand(t_uniform, 2, 3))

t_uniform_bits = UniformBits(brng)
@test_rand(rand(t_uniform_bits))
@test_rand(rand(t_uniform_bits, 2, 3))

t_uniform_bits32 = UniformBits32(brng)
@test_rand(rand(t_uniform_bits32))
@test_rand(rand(t_uniform_bits32, 2, 3))

t_uniform_bits64 = UniformBits64(brng)
@test_rand(rand(t_uniform_bits64))
@test_rand(rand(t_uniform_bits64, 2, 3))

t_bernoulli = Bernoulli(brng, 0.3)
@test_rand(rand(t_bernoulli))
@test_rand(rand(t_bernoulli, 4, 5))

t_geometric = Geometric(brng, 0.6)
@test_rand(rand(t_geometric))
@test_rand(rand(t_geometric, 6, 3))

t_binomial = Binomial(brng, Cint(10), 0.6)
@test_rand(rand(t_binomial))
@test_rand(rand(t_binomial, 6, 3))

t_hypergeometric = Hypergeometric(brng, Cint(10), Cint(5), Cint(6))
@test_rand(rand(t_hypergeometric))
@test_rand(rand(t_hypergeometric, 3, 6))

t_poisson = Poisson(brng, 3.0)
@test_rand(rand(t_poisson))
@test_rand(rand(t_poisson, 3, 6))

t_poisson_v = PoissonV(brng, [1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8])
@test_rand(rand(t_poisson_v))
@test_rand(rand(t_poisson_v, 3, 3))

t_neg_binomial = NegBinomial(brng, 3.0, 0.7)
@test_rand(rand(t_neg_binomial))
@test_rand(rand(t_neg_binomial, 3, 4))

close(outfile)
redirect_stdout(stdout_)

@test success(`diff ./expected/discrete.out ./actual/discrete.out`)
cd(pwd_)


#=
1D advection-diffusion using FV
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("FVaddiff1d");

useLog("FVaddiff1dlog", level=3)

# Configuration setup
domain(1)
solverType(FV)
timeStepper(RK4)

# Mesh
n = 40 # number of elements
mesh(LINEMESH, elsperdim=n, bids=2)

# Variables and BCs
u = variable("u", location=CELL)
v = variable("v", location=CELL)
boundary(u, 1, FLUX, 0)
boundary(u, 2, NO_BC)
boundary(v, 1, FLUX, 0)
boundary(v, 2, NO_BC)

# Time interval and initial condition
T = 0.5;
timeInterval(T)
initial(u, "x<0.1 ? 1 : 0")
initial(v, "x<0.1 ? 1 : 0")

# The flux and source terms of the conservation equation
# F and S in the following equation:
# Dt(int(u dx)) = int(S dx) - int(F.n ds)
coefficient("a", 1) # advection velocity
# The "upwind" function applies upwinding to the term (a.n)*u with flow velocity a.
# The optional third parameter is for tuning. Default upwind = 0, central = 1. Choose something between these.
coefficient("D", 0.01) # Diffusion rate
flux([u, v], ["upwind(a,u) - D*dot(grad(u),normal())", "upwind(a,v)"])
# Note that there is no source() for this problem.

#@exportCode("fvad1dcode") # uncomment to export generated code to a file
#@importCode("fvad1dcode") # uncomment to import code from a file

solve([u,v])

finalize_finch()

##### Uncomment below to plot

# # The initial condition
# init = zeros(n);
# x = Finch.fv_info.cellCenters[:]
# for i=1:n
#     if x[i] < 0.1
#         init[i] = 1
#     end
# end

# using Plots
# pyplot();
# display(plot([x x x], [init u.values[:] v.values[:]], markershape=:circle, label=["initial" "upwind advection, central diffusion" "advection only"]))


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

using ACE, BenchmarkTools
using ACE.JacobiPolys: Jacobi
using JuLIP: evaluate

x = 2*rand() - 1
α, β = rand(), rand()
N = 30
P = zeros(N+1)
dP = zeros(N+1)
J = Jacobi(α, β, N)

evaluate!(P, nothing, J, x)
evaluate_d!(P, dP, nothing, J, x)
@info("Timing for evaluate!")
@btime evaluate!($P, nothing, $J, $x)
@btime evaluate!($P, nothing, $J, $x)
@info("Timing for evaluate_d!")
@btime evaluate_d!($P, $dP, nothing, $J, $x)
@btime evaluate_d!($P, $dP, nothing, $J, $x)

# # Julia Bug? => not a bug but a current limitation
# @info("Looking at that strange allocation?")
# using BenchmarkTools
# f(A, B) = A, B
# g(A, B) = A
# A = rand(5)
# B = rand(5)
# @btime f($A, $B);
# @btime f(1, 2);
# @btime g($A, $B);
#
# function runn(P, dP, J, x, N)
#    for n = 1:1000
#       evaluate_d!(P, dP, J, rand(), N)
#    end
#    return nothing
# end
#
# @info("Try again inside a function - no allocation!")
# @btime runn($P, $dP, $J, $x, 15)


module TestApproxRidge
    using StreamStats
    using Distributions
    using Base.Test

    for n in rand(100_000:1_000_000, 10)
        p = rand(1:100)
        β₀ = randn()
        β = randn(p)
        xs = randn(p, n)
        ys = Float64[β₀ + dot(β, xs[:, i]) + randn() for i in 1:n]
        stat = StreamStats.ApproxRidge(p, 1.0)
        for i in 1:n
            update!(stat, xs[:, i], ys[i])
        end
        @test cor(vcat(β₀, β), state(stat)) > 0.99
        # TODO: Add real tests
        # l2_err = norm(vcat(β₀, β) - state(stat))
        # @test l2_err < 0.05 * p
        # @test nobs(stat) == n
    end
end


const β_min, β_step = 0.01, 0.01

function x_upper(β)
    q = -Inf
    x = 0.01
    while cdf(GeneralizedKolmogorov(β), x; precomputed_interpolation = false) < 0.9999
        x *= 2
    end
    return x
end

function make_cdfs()
    β_grid = β_min:β_step:1.0
    δ = 0.001

    cdf_inter_results = []
    @showprogress for β in β_grid
        upper_value = x_upper(β)
        x_grid = 0:δ:upper_value
        cdf_vals = vcat([cdf(GeneralizedKolmogorov(β), x; precomputed_interpolation = false) for x in x_grid], 1.0)
        x_grid = union(x_grid, 2*upper_value)
        push!(cdf_inter_results, LinearInterpolation(x_grid, cdf_vals, extrapolation_bc=Flat()))
    end
    return cdf_inter_results
end


using Distributions

# Uncertain parameters defined in "Projections and Uncertainties about Climate
# Change in an Era of Minimal Climate Policies" (Nordhaus, 2018) EViews file
# provided by the AEJ : https://www.aeaweb.org/articles?id=10.1257/pol.20170046

dice_default_mcs = @defsim begin
    climatedynamics.t2xco2 = LogNormal(1.106, 0.2646)
	totalfactorproductivity.ga0 = Normal(0.076, 0.056)
	emissions.gsigma1 = Normal(-0.0152, 0.0032)
	co2cycle.mueq = LogNormal(5.851, 0.2649)
	damages.a2 = Normal(0.00227, 0.001135)

    # save(damages.DAMAGES)
	# save(grosseconomy.YGROSS)
	# save(emissions.E)
	# save(co2cycle.MU)
	# save(damages.DAMFRAC)
end

function getmcs()
    return deepcopy(dice_default_mcs)
end


using SpecialFunctions
using QuadGK

function ourquadgk(f, a, b)
    (value, _) = quadgk(f, a, b, )
    return value
end

"""
K(ρ, β, x, r) computes the integral original function, gives a more convinent way to calculste the cumulative sum of the inner function.
"""

"""
wright(ρ, β, z) computes the Wright function of z.

While in the paper, q₁ is a fixed small number which 0<q₁<1 and q₂ is a fixed big number,
Here we set q₁ = 0.5 and q₂ = 100
"""
q₁ = 2
q₂ = 100

function wright(ρ, β, z)
    ϵ = 0.2
    x=abs(z)
    if abs(z) ≤ q₁ #HACK: This part is working fine
        if ρ > 0
            k₀ = max((1-β)/ρ, log(ϵ)/log(abs(z))-1)
        elseif -1 < ρ < 0
            med = (1-β)/(1+ρ)
            k₀ = max(-β/ρ, med, log(π*(1+ρ)*ϵ*abs(z)^(-med))/(log(abs(z))-1-ρ)+med)
        end
        s = zero(z)
        k₀ = ceil(k₀)
        for k=0:k₀
            s += (z^k)/(factorial(k)*gamma(β + ρ*k))#+μz #NOTE: factorial might led to overflow when k is bigger than 21
        end
        return s

    elseif q₁ < abs(z) < q₂ #The most important in applications

        #TODO: What is x stand for?? The real part of z or just a defination??
        med1 = 1-z*cos(ρ*π)
        med2 = 1/ρ+1-β/ρ
        r₀=0
        if (-1 < ρ < -1/2) && (1+ρ ≤ β)
            r₀ = max(med1^(1/(1+ρ)), (-log(-π*ρ*ϵ))^(-1/ρ))
        elseif (-1 < ρ < -1/2) && (β<1+ρ)
            B = -1/(ρ*π)*(abs(med2)+2)*(-2*med2)^(-med2)
            r₀ = max((med1), (abs(med2)+1)^(-1/ρ), (-2*log(ϵ/B))^(-1/ρ))
        elseif (-1/2 <= ρ < 0) & (β ≥ 0)
            r₀ = -log(ϵ*π)
        elseif (-1 < ρ < -1/2) & (β<0)
            r₀ = max(abs(β)+1, -2*log(ϵ*π/(4(abs(β)+2)*β^2)))
        elseif (ρ > 0) & (β ≥ 0)
            r₀ = max(abs(x*cos(ρ*π))^(1/ρ), -log(ϵ*π)+1)
        elseif (ρ > 0) & (β < 0)
            r₀ = max(abs(x*cos(ϵ*π))^(1/ρ), abs(β)+1, -2*log(ϵ*π/(4*(abs(β)+2)*β^2))+2)
        end
        function K(r)
            return exp(-r-x*r^(-ρ)*cos(ρ*π))*r^(-β)*sin(-x*r^(-ρ)*sin(ρ*π)+π*β)
        end
        return ourquadgk(K, 1, r₀) # The improper integral require the lower limit ∈ {0, 1} and r₀≥1, we set 0 as the lower limit.

    elseif abs(z)>q₂
    end

    #FIXME: The third part abs(z) > q₂ is not implement yet

end


for N in [Float64, Float32, Rational{Int}]

    # interval in union of intervals
    X = Interval(N(0), N(1))

    # included
    YU = Interval(N(1//2), N(2)) ∪ Interval(N(0), N(2//3))
    YU_hr = convert(Hyperrectangle, YU.X) ∪ convert(Hyperrectangle, YU.Y)
    Yarr = UnionSetArray([YU.X, YU.Y])
    for Y in [YU, YU_hr, Yarr]
        res, w = ⊆(X, Y, true)
        @test X ⊆ Y && res && w == N[]
    end

    # not included
    YU = Interval(N(1//2), N(5//6)) ∪ Interval(N(0), N(1//3))
    YU_hr = convert(Hyperrectangle, YU.X) ∪ convert(Hyperrectangle, YU.Y)
    Yarr = UnionSetArray([YU.X, YU.Y])
    for Y in [YU, YU_hr, Yarr]
        res, w = ⊆(X, Y, true)
        @test !(X ⊆ Y) && !res && w ∈ X && w ∉ Y
    end

    # using IA types
    X = IA.interval(N(0), N(1)) # IA
    Y = Interval(N(-1), N(2))
    @test X ⊆ Y
    @test !(Y ⊆ X)

    X = IntervalBox(IA.interval(N(0), N(1)), IA.interval(N(0), N(1)))
    Y = Hyperrectangle(low=[N(-1), N(-1)], high=[N(2), N(2)])
    @test X ⊆ Y
    @test !(Y ⊆ X)

end


@inline isnearint(v, tol) = abs(v-round(v)) < tol

# TODO (cg 2021/10/22 17:04): The cutoffs chosen here are actually to some
# degree chosen as a balance between pointwise accuracy AND AD derivative
# accuracy. For example, the series gets worse faster for derivatives than it
# does for raw evals. Not entirely sure what to make of that, but here we are.
#
# TODO (cg 2021/10/27 10:28): the weak point is clearly with |z| between, say, 8
# and 15. We only get atols on the order of 1e-12 for that, which is not quite
# good enough to declare total victory. I can only manage to match the speed of
# AMOS in that part of the domain, too.
function _besselk(v, x, maxit=100, tol=1e-12, order=6)
  @assert x >= zero(x) DomainError("x needs to be non-negative.")
  iszero(x) && return Inf
  (real(v) < 0) && return _besselk(-v, x, maxit, tol, order)
  # TODO (cg 2021/11/16 16:44):  this is not the right way to test if you're
  # doing AD. But I don't want to hard-bake the ForwardDiff.Dual type in here.
  is_ad = !(v isa AbstractFloat)
  # Special cases, now just half-integers:
  #
  # TODO (cg 2021/12/16 12:41): for the moment, specifically at half
  # integer values the second derivatives are more accurate using the direct
  # series, even at values in which for direct evaluations v is large enough
  # that the series isn't so great. It isn't earth-shattering and the temme one
  # is much more expensive (300 ns vs 100 ns on my machine), but in the interest
  # of maximum safety I'm switching to this behavior.
  #
  # What would really be nice is some way if checking if (v isa Dual{T,V,N}
  # where T<: Dual). But again, I'm worried about getting too stuck with
  # ForwardDiff.
  if isinteger(v-1/2) && (x < 8.5)
    if is_ad
      return _besselk_ser(v, x, maxit, tol, false)
    else
      # Probably don't need the is_ad correction here.
      return _besselk_as(v, x, Int(ceil(v)), false)
    end
  end
  #
  # General cases:
  #
  # TODO (cg 2021/11/01 18:03): These branches are not perfect. If you go into
  # ./testing/accuracy.jl and track down the largest rtols between this code and
  # AMOS, you will be able to fiddle around with what version you use and get a
  # better rtol/atol. But I'm at the point where whenever I tweak something like
  # that, something else gets worse and makes it a wash. I think for the time
  # being I have to stop playing with things.
  if abs(x) < 8.5 # (x < 9)
    if (v > 2.85) || isnearint(v, 0.01) || ((x > 4) && is_ad)
      return _besselk_temme(v,x,maxit,tol,false) # direct series.
    else
      return _besselk_ser(v,x,maxit,tol,false) # direct series.
    end
  elseif abs(x) < 15.0
    return _besselk_asv(v,x,12) # uniform large order expn.
  elseif abs(x) < 30.0
    return _besselk_asv(v,x,8)
  else
    if abs(v) > 1.5 || is_ad
      return _besselk_asv(v,x,6)
    else
      return _besselk_as(v,x,order)
    end
  end
end

# A very simple wrapper that will use AMOS when possible, but fall back to this
# code. So now you can use AD on this but still get AMOS for direct evals.
#
# TODO (cg 2021/11/05 16:57): what's the most sensible naming thing here?
# Calling it besselk and not exporting it seems reasonable enough, but users
# will obviously want to import it. So not obvious what's best to do here.
function adbesselk(v, x, maxit=100, tol=1e-12, order=5)
  if v isa AbstractFloat
    SpecialFunctions.besselk(v, x)
  else
    _besselk(v, x, maxit, tol)
  end
end

# Not exactly a taylor series, but accurate enough.
#
# TODO (cg 2021/11/10 18:33): an enhancement here would be something that also
# worked for integers. But that seems hard. I did the whole Temme thing because
# integers are hard. But as it turns out, we really need it, because the Temme
# recursion depends on (x^v)*besselk(v,x) ->_{z->0} some finite number. But for
# v=0, which is needed to compute for _all_ integer v, that doesn't hold! An
# expansion like this that is valid for all v, including integer v, but is ALSO
# AD-compatible would take care of that entirely, because the K_{v+1}(x) term in
# the Temme series doesn't need f0.
@inline function besselkxv_t0(v, x)
  gv  = gamma(v)
  _2v = 2^(v-1)
  cof = (gv*_2v, zero(v), (_2v/4)*gv/(v-1), zero(v), (_2v/16)*gv/(v*v - 3*v + 2))
  evalpoly(x, cof)
end

# Note that the special series cutofs are low compared to the above function. In
# general, the adbesselk* functions that are exported really should be pretty
# near machine precision here or should just fall back to AMOS when possible. It
# turns out that the *xv modifications in the code are really only helpful when
# the argument is pretty small.
function adbesselkxv(v, x, maxit=100, tol=1e-12, order=5)
  (iszero(v) && iszero(x)) && return Inf
  is_ad = !(v isa AbstractFloat)
  xcut  = is_ad ? 6.0 : 2.0
  if !isinteger(v) && (abs(x) <= 1e-8) # use Taylor at zero.
    return besselkxv_t0(v, x)
  elseif (x < xcut) && (v < 5.75) && !isnearint(v, 0.01)
    return _besselk_ser(v, x, maxit, tol, true)
  elseif (x < xcut)
    return _besselk_temme(v, x, maxit, tol, true)
  elseif is_ad && (x > xcut) && (x < 15.0)
    return _besselk_asv(v, x, 12, true)
  else
    return adbesselk(v, x, maxit, tol, order)*(x^v)
  end
end

# Unlike adbesselkxv, this function is pure BesselK.jl, including the cases in
# which special branches for (x^v)*besselk(v,x) don't come up. This is primarily
# used for testing.
function _besselkxv(v, x, maxit=100, tol=1e-12, order=5)
  (iszero(v) && iszero(x)) && return Inf
  is_ad = !(v isa AbstractFloat)
  xcut  = is_ad ? 6.0 : 2.0
  if !isinteger(v) && (abs(x) <= 1e-8) # use Taylor at zero.
    return besselkxv_t0(v, x)
  elseif (x < xcut) && (v < 5.75) && !isnearint(v, 0.01)
    return _besselk_ser(v, x, maxit, tol, true)
  elseif (x < xcut)
    return _besselk_temme(v, x, maxit, tol, true)
  elseif is_ad && (x > xcut) && (x < 15.0)
    return _besselk_asv(v, x, 12, true)
  else
    return _besselk(v, x, maxit, tol, order)*(x^v)
  end
end