struct Best end # best e for all combinations of a_k

# best e for all combinations of subset of a_k's at positions of ones in c
struct SubSet{T<:Integer}
    a::Vector{T}      # subset to optimize
end

struct SlidingWindow
    width :: Int
end

struct MabeyBest end # optimized version of Best

function search(::Best,s::Integer,n::Integer,::Vector=[])
    a = Array{Int}(1:s-1)
    a_to_n = Array{BigInt}(BigInt.(a).^n)
    lhs = BigInt(s)^n
    bc,best_error = best_combination(a_to_n, lhs)
    collect(OnePositions(bc)),best_error
end

function search(ss::SubSet{T}, s::T, n::Integer, a::Vector{T}) where {T<:Integer}
    @assert length(a)<s
    lhs =  BigInt(s)^n
    a_not_in_subset = setdiff(a,ss.a)
    for a_k in a_not_in_subset
        lhs-=BigInt(a_k)^n
    end
    bc_subset,best_error = best_combination([BigInt(x)^n for x in ss.a], lhs)
    best_a = a_not_in_subset
    sizehint!(best_a,length(best_a)+count_ones(bc_subset))
    for k in OnePositions(bc_subset)
        push!(best_a,ss.a[k])
    end
    best_a, best_error
end

function best_combination(a_to_n,
                          lhs::BigInt,
                          best_error::BigInt = BigInt(lhs)
                          )
    bc = 0
    for combination in 1:2^length(a_to_n)-1
        error = BigInt(lhs)
        for k in OnePositions(combination)
            error-=a_to_n[k]
        end
        if abs(error) < abs(best_error)
            best_error = error
            bc = combination
        end
    end
    bc,best_error
end

function search(sw::SlidingWindow, s::Int, n::Int, a::Vector{Int})
    window_max = s-1
    window_min = s-1-sw.width
    ss = SubSet([x for x in window_max:-1:window_min])
    e = BigInt(0)
    for i in 1:window_min
        a,e = search(ss,s,n,a)
        ss.a.-=1
    end
    a,e
end

function search(::MabeyBest, s::Int, n::Int, ::Vector{Int}=Vector{Int}([]))
    max_k = s-1
    all_a_to_n = [BigInt(x)^n for x in 1:max_k]
    lhs = BigInt(s)^n
    select_a = if max_k<65
        UInt64(0)
    elseif max_k<129
        UInt128(0)
    else
        BigInt(0)
    end
    best = Tracker(select_a,lhs)
    _mabey_best!(select_a, best, lhs, max_k, all_a_to_n)
    best_a, best_error = best()
    Int.(collect(OnePositions(best_a))), best_error
end
function _mabey_best!(select_a::T, best::Tracker, lhs::BigInt, max_k::Int, all_a_to_n) where T<:Integer
    # select_a and best are modified
    is_error_zero(best) && return nothing
    if max_k == 0
        best(select_a, lhs)
        return nothing
    end
    split,e = find_split(lhs, max_k, all_a_to_n)
    if split-1==max_k
        all_ones_to_split = select_a | T(2)^(split-1)-1
        best(all_ones_to_split, e)
        return nothing
    end
    upper_view_a_to_n = view(all_a_to_n, split:max_k)
    lower_view_a_to_n = view(all_a_to_n, 1:split-1)
    upper_combinations_range = 1:T(2)^(max_k-split)
    for upper_combinations in upper_combinations_range
        new_select_a = select_a ⊻ upper_combinations<<(split-1)
        new_lhs = lhs-sum(upper_view_a_to_n[OnePositions(upper_combinations)])
        if new_lhs>0
            _mabey_best!(new_select_a, best, new_lhs, split-1, lower_view_a_to_n)
        else
            best(new_select_a, new_lhs)
        end
    end
    nothing
end
function find_split(lhs::BigInt, max_k::Int, all_a_to_n)
    e = lhs
    for k in 1:max_k
        previous_error = e
        e-=all_a_to_n[k]
        e<=0 && return k,previous_error
    end
    return max_k,e
end
function true_positions(x::BitArray{1})
    positions = Vector{Int}()
    for i in eachindex(x)
        i[x] && push!(positions,i)
    end
    positions
end

function max_s_for_all_a(n)
    s = 1
    old_s_to_n = BigInt(1)
    e = old_s_to_n
    while true
        s+=1
        new_s_to_n = BigInt(s)^n
        e = e + new_s_to_n - 2*old_s_to_n
        e<0 && return s-1
        old_s_to_n = new_s_to_n
    end
end


function lfur_pv(X, c)
  #
  #
  q = size(X, 1);
  tmp = ((1:1:q)*pi).^2;
  cov_H0 = 1./tmp;
  chol_H0 = sqrt(cov_H0);
  chol_H0_inv = sqrt(1./cov_H0);
  X0 = chol_H0_inv.*X;
  cov_HA = Sigma_Compute2(0, c, 1, q)
  chol_HA_inv = chol(Hermitian(inv(cov_HA), :L));
  XA = chol_HA_inv*X;
  lfur = sum(X0.^2)/sum(XA.^2);
  # Generate 1000 draws under Ho for pvalue calculation
  nsim = 50000;
  Xsim = rand(Normal(0,1), q, nsim).*repmat(chol_H0, 1, nsim);
  X0sim = Xsim.*repmat(chol_H0_inv, 1, nsim);
  XAsim = chol_HA_inv*Xsim;
  nx = sum(X0sim.^2, 1);
  dx = sum(XAsim.^2, 1);
  lfur_sim = squeeze((nx./dx)',2);
  lfur_pvalue =  mean(lfur_sim.>=lfur)
  return lfur, lfur_pvalue
end


using Plots
using FFTW

function find_nyquist_frequency(filebase::String, res::Int64)
    arr = Array{Float64}(undef, res)
    for i = 1:res
        for j = 1:res
             arr[j] = sin(j*i*2*pi/res)
        end
        plot(arr, ylim = (-1,1))
        savefig(filebase * string(lpad(string(i-1), 4, string(0))) * ".png")

        plot(abs.(fft(arr)))
        savefig(filebase *"FFT"* string(lpad(string(i-1), 4, string(0)))
                         * ".png")
        println(i)
    end
end

find_nyquist_frequency("images/nyfreq", 40)


## Author: Victor H. Aguiar

## questions: slack

## Import packages, make sure to install them.
## Julia version: v"1.5.2"
import JuMP
import Ipopt

## define a model
robinson=JuMP.Model(Ipopt.Optimizer)
#Knitro
# define variables
#coconuts
@JuMP.variable(robinson,c>=0)
#labor
@JuMP.variable(robinson,l>=0)
#wages
@JuMP.variable(robinson,w>=0)

# define parameters
# you must change this values to fit the dataset given to you
α1=.4
α2=(1-α1)
β = .1
A=50
Lbar=100.0

# Solve model
# each term of the summation is an equation of the system of equations
tax=0
@JuMP.NLobjective(robinson,Min,(α1/(α2)*(Lbar-l)*w-c)^2+(A*l^β*(1-tax)-c)^2+(A*β*l^(β-1)*(1-tax)-w)^2)
JuMP.optimize!(robinson)

JuMP.value(c)
JuMP.value(l)
JuMP.value(w)

taxvec=collect(0:.01:.99)
revenuevec=zeros(length(taxvec))
for t in 1:length(taxvec)
    tax=taxvec[t]
    @JuMP.NLobjective(robinson,Min,(α1/(α2)*(Lbar-l)*w-c)^2+(A*l^β*(1-tax)-c)^2+(A*β*l^(β-1)*(1-tax)-w)^2)
    JuMP.optimize!(robinson)
    revenuevec[t]=A*JuMP.value(l)^β*tax
end

sol=findmax(revenuevec)
taxvec[sol[2]]


using Submodular, Convex
using FactCheck

Tol = 1e-3

facts("Prox") do

  context("prox of associated polyhedra") do
    S = SetVariable(4)
    p(z) = -0.5*z^2 + 4*z + 0.5 * z
    perm_func = compose(p, card(S))
    P = SubmodPoly(perm_func)

    x = Variable(4)
    w₁ = [1, 2, 5, 3]
    w₂ = [-2, -2, -5, -3]
    @fact prox(x in P, w₁) --> roughly([1, 2, 4, 3], Tol)
    @fact prox(x in P, w₂) --> roughly([-2, -2, -5, -3], Tol)
  end

  context("prox of convex functions") do
    w = rand(4)
    x = Variable(4)
    f = 0.5 * norm(x)^2
    prox(f, w)
    @fact evaluate(0.5 * w - x) --> roughly(zeros(4, 1), Tol)
  end

end


@doc raw"""
    SLIM(l1_ratio::Float64 = 0.5, λminratio::Float64 = 1e-4, k::Int = -1)

Sparse linear machine for recommendation, modified with Elastic Net loss. The prediction is made by
```math
\hat r_{ui} = \sum_{j\neq i} w_{ij} r_{uj}
```
where ``r_{ui}`` is the actual rating for item ``i`` by user ``u``, and ``\hat r_{ui}`` is the predicted value. ``w_{ij}`` is the model weght matrix. See Refs for algorithm details. SLIM uses [`Lasso.jl`](https://github.com/JuliaStats/Lasso.jl) for optimization.

# Constructor arguments
- `l1_ratio`: ratio of coefficients between ``L_1`` and ``L_2`` losses. `l1_ratio` ``\to 0`` means the Ridge regularization, while `l1_ratio` ``\to \infty`` the Lasso.
- `λminratio`: parameter which governs the strength of regularization. See the docs of `Lasso.jl`.
- `k`: the nearest neighborhood size, similar to `ItemkNN`. If `k` < 1, the neigoborhood size is infinity.

# References
- X. Ning and G. Karypis (2011), [SLIM: Sparse Linear Methods for Top-N Recommender Systems](http://glaros.dtc.umn.edu/gkhome/node/774)
- M. Levy (2013), [Efficient Top-N Recommendation by Linear Regression](https://www.slideshare.net/MarkLevy/efficient-slides)
"""
mutable struct SLIM <: AbstractRecommender
    l1_ratio::Float64
    λminratio::Float64
    k::Int

    W::Union{SparseMatrixCSC{Float64,Int},Nothing}

    user2uidx::Union{Dict,Nothing}
    item2iidx::Union{Dict,Nothing}
    iidx2item::Union{Dict,Nothing}
    R::Union{SparseMatrixCSC{Float64,Int},Nothing}

    SLIM(l1_ratio::Float64 = 0.5, λminratio::Float64 = 1e-4, k::Int = -1) =
        new(l1_ratio, λminratio, k, nothing, nothing, nothing, nothing, nothing)
end

function truncate_at_k!(w::SparseVector, k::Int)
    is, ws = findnz(w)
    if k >= length(ws) || k < 1
        return
    end
    arg_outof_k = sortperm(ws, rev = true)[(k+1):end]
    w[is[arg_outof_k]] .= 0
end

"""
    fit!(model::SLIM, table; col_user = :userid, col_item = :itemid, col_rating = :rating, cd_tol = 1e-7, nλ = 100)

Fit the SLIM model.

# Model-specific arguments
- `cd_tol`: tolerance paramerer for convergence, see `Lasso.jl`
- `nλ`: length of regularization path, see `Lasso.jl`
"""
function fit!(
    model::SLIM,
    table;
    col_user = :userid,
    col_item = :itemid,
    col_rating = :rating,
    cd_tol = 1e-7,
    nλ = 100,
    kwargs...,
)
    table, model.user2uidx, model.item2iidx, model.iidx2item =
        make_idmap(table, col_user = col_user, col_item = col_item)

    model.R = rows2sparse(
        table,
        col_user = col_user,
        col_item = col_item,
        col_rating = col_rating,
    )

    unique_items = collect(keys(model.iidx2item))
    n_item = length(unique_items)

    model.W = spzeros(n_item, n_item)
    # for multithreading
    tmp_ws = Vector{SparseVector{Float64,Int}}(undef, n_item)

    Threads.@threads for i = 1:n_item
        mask = spzeros(n_item, n_item)
        mask[i, i] = 1
        lasso_res = fit(
            LassoPath,
            model.R - model.R * mask,
            Vector(model.R[:, i]),
            α = model.l1_ratio,
            λminratio = model.λminratio,
            standardize = false,
            intercept = false,
            stopearly = true,
            cd_tol = cd_tol,
            nλ = nλ,
        )
        wi = lasso_res.coefs[:, end]
        truncate_at_k!(wi, model.k)
        dropzeros!(wi)
        tmp_ws[i] = wi
    end

    for i = 1:n_item
        model.W[:, i] = tmp_ws[i]
    end
end

@doc raw"""
    predict_u2i(model::SLIM, userid::Union{AbstractString,Int}, n::Int64; drop_history = false)

Make predictions by SLIM model.
"""
function predict_u2i(
    model::SLIM,
    userid::Union{AbstractString,Int},
    n::Int64;
    drop_history = false,
    kwargs...,
)
    if !(userid in keys(model.user2uidx))
        return []
    end

    uidx = model.user2uidx[userid]
    user_history = model.R[uidx, :]

    pred = (user_history' * model.W)'
    viewed_item, ratings = findnz(user_history)
    for (iidx, rui) in zip(viewed_item, ratings)
        pred[iidx] -= rui * model.W[iidx, iidx]
    end

    pred_iidx = sortperm(pred, rev = true)
    # this is very slow
    if drop_history
        filter!(p -> !(p in viewed_item), pred_iidx)
    end

    pred_items = [model.iidx2item[iidx] for iidx in pred_iidx]

    n = min(n, length(pred_items))
    return pred_items[1:n]
end


function diagnose_u(ψ, d::Domain)
    u = new_field(d)
    for I = CartesianIndices(d)
        u[I] = -∂y(I,ψ,d)
    end
    return u
end

function diagnose_v(ψ, d::Domain)
    v = new_field(d)
    for I = CartesianIndices(d)
        v[I] = ∂x(I,ψ,d)
    end
    return v
end

function diagnose_umag(ψ, d::Domain)
    umag = new_field(d)
    for I = CartesianIndices(d)
        umag[I] = sqrt(∂y(I,ψ,d)^2 + ∂x(I,ψ,d)^2)
    end
    return umag
end

function diagnose_ζ(ψ, d::Domain)
    ζ = new_field(d)
    for I = CartesianIndices(d)
        ζ[I] = ∂x²(I,ψ,d) + ∂y²(I,ψ,d)
    end
    return ζ
end

function diagnose_θ(ϕ, d::Domain)
    θ = new_field(d)
    for I = CartesianIndices(d)
        θ[I] = -∂z(I,ϕ,d)
    end
    return θ
end

function diagnose_penetration_depth(ψ, q, d::Domain)
    qmax, Imax = findmax(q)
    kmax = Imax[3]
    umag = diagnose_umag(ψ, d)
    ubot = umag[:,:,1]
    umax = umag[:,:,kmax]
    return maximum(ubot)/maximum(umax)
end


module TestAnalysis
using CustomTest
using Base.Test
using ControlSystems

@test_approx_eq stabregionPID("exp(-sqrt(s))", logspace(-3,3), kd=1)[2][1] -1.022356911142034

P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=true)
@test_approx_eq kp 0.82274734724854
@test_approx_eq ki -0.45472580052281375


using ReachabilityAnalysis, LinearAlgebra, SparseArrays

function tline(;η=3, R=1.00, Rd=10.0, L=1e-10, C=1e-13 * 4.00)
    A₁₁ = zeros(η, η)
    A₁₂ = Bidiagonal(fill(-1/C, η), fill(1/C, η-1), :U)
    A₂₁ = Bidiagonal(fill(1/L, η), fill(-1/L, η-1), :L)
    A₂₂ = Diagonal(vcat(-Rd/L, fill(-R/L, η-1)))
    A  = [A₁₁ A₁₂; A₂₁ A₂₂]
    B = sparse([η+1], [1], 1/L, 2η, 1)
    return A, B
end

using Plots

A, _ = tline(η=20)
spy(A, legend=nothing, markersize=2.0, title="Sparsity pattern of A", xlab="columns", ylab="rows")

function scale!(s, α=1.0)
    s.A .*= α
    s.B .*= α
    return s
end

η = 20 # order
n = 2η # state-space dimension
A, B = tline(η=η);

Uin_ss = Interval(-0.2, 0.2)
□(ϵ) = BallInf(zeros(n), ϵ)
X0 = -inv(Matrix(A)) * B * Uin_ss ⊕ □(0.001);

Uin = Interval(0.99, 1.01)
s = @system(x'= A*x + B*u, x ∈ Universe(n), u ∈ Uin)
α = 1e-9 # scaling factor
scale!(s, α);

T = 0.7 * 1e-11 # time horizon
P = InitialValueProblem(s, X0);

sol = solve(P, T=0.7, alg=BOX(δ=1e-3));

Uout_vs_t = @. (-1.0) * project(sol, η);

plot(Uout_vs_t, vars=(0, η), color=:blue, xlab="t", ylab="Uout", alpha=.5, lw=0.5)


"
solves AL * C = C * A in-place
"
function uniform_leftorth!(AL,CR, A; tol = Defaults.tolgauge, maxiter = Defaults.maxiter)
    #(_,CR[end]) = leftorth!(CR[end], alg = TensorKit.QRpos());
    normalize!(CR[end]);

    iteration = 1;
    delta = 2*tol;

    while iteration < maxiter && delta > tol

        if iteration>10 #when qr starts to fail, start using eigs - should be kw arg

            alg = Arnoldi(krylovdim = 30, tol = max(delta*delta,tol/10),maxiter=maxiter)

            (vals,vecs) = @closure eigsolve(CR[end], 1, :LM,alg) do x
                transfer_left(x,A,AL)
            end

            (_,CR[end]) = leftorth!(vecs[1],alg=TensorKit.QRpos())
        end

        cold = CR[end]

        for loc in 1:length(AL)
            AL[loc] = _transpose_front(CR[mod1(loc-1,end)]*_transpose_tail(A[loc]))
            AL[loc], CR[loc] = leftorth!(AL[loc], alg = QRpos())
            normalize!(CR[loc])
        end

        #update delta
        if domain(cold) == domain(CR[end]) && codomain(cold) == codomain(CR[end])
            delta = norm(cold-CR[end])
        end

        iteration += 1
    end

    delta>tol && @warn "leftorth failed to converge $(delta)"

    AL,CR
end

"
solves C * AR = C * A in-place
"
function uniform_rightorth!(AR,CR,A; tol = Defaults.tolgauge, maxiter = Defaults.maxiter)
    #(CR[end],_) = rightorth!(CR[end], alg = TensorKit.LQpos());
    normalize!(CR[end])

    iteration = 1;
    delta = 2*tol;
    while iteration<maxiter && delta>tol

        if iteration>10#when qr starts to fail, start using eigs
            alg = Arnoldi(krylovdim = 30, tol = max(delta*delta,tol/10),maxiter=maxiter)
            #Projection of the current guess onto its largest self consistent eigenvector + isolation of the unitary part

            (vals,vecs) = @closure eigsolve(CR[end], 1, :LM,alg) do x
                transfer_right(x,A,AR)
            end
            (CR[end],_) = rightorth!(vecs[1],alg=TensorKit.LQpos())
        end

        cold = CR[end]
        for loc in length(AR):-1:1
            AR[loc] = A[loc]*CR[loc]

            CR[mod1(loc-1,end)], temp = rightorth!(_transpose_tail(AR[loc]), alg=LQpos())
            AR[loc] = _transpose_front(temp)
            normalize!(CR[mod1(loc-1,end)])
        end

        #update counters and delta
        if domain(cold) == domain(CR[end]) && codomain(cold) == codomain(CR[end])
            delta = norm(cold-CR[end])
        end

        iteration += 1
  end

  delta>tol && @warn "rightorth failed to converge $(delta)"

  return AR, CR
end


using Distributions

include("ADVI.jl")

# Dirichlet
d = Dirichlet(collect(1:3)); x = collect(1:3) / sum(1:3)
@assert ADVI.compute_lpdf(d, x) == logpdf(d, x)

# Normal
d = Normal(2, 3); x = -5.0
@assert ADVI.compute_lpdf(d, x) == logpdf(d, x)

# Beta
d = Beta(2, 3); x = .6
@assert ADVI.compute_lpdf(d, x) == logpdf(d, x)

# Gamma
d = Gamma(2,4); x = 5.0
@assert ADVI.compute_lpdf(d, x) == logpdf(d, x)

# LogNormal
d = LogNormal(2, 4); x = 1.5
@assert ADVI.compute_lpdf(d, x) == logpdf(d, x)
# ADVI.compute_lpdf(d, x)
# logpdf(d, x)

# Uniform
d = Uniform(2,4); x = 3.5
@assert ADVI.compute_lpdf(d, x) == logpdf(d, x)


function inner_dual(thePortfolio::SparsePortfolioData, indices::Array{Int64})

    n=thePortfolio.f
    m=thePortfolio.m
    f=size(indices,1)

    innerQP = Model(solver=MosekSolver(MSK_DPAR_INTPNT_QO_TOL_PFEAS=1e-6, MSK_DPAR_INTPNT_QO_TOL_DFEAS=1e-6, MSK_IPAR_LOG=0, MSK_IPAR_MAX_NUM_WARNINGS=0)) #
    #

    @variable(innerQP, α[1:n])
    @variable(innerQP, λ)
    @variable(innerQP, w[1:f])
    @variable(innerQP, ρ[1:f]>=0.0)
    @variable(innerQP, βl[1:m]>=0.0)
    @variable(innerQP, βu[1:m]>=0.0)

    @constraint(innerQP, w.>=(thePortfolio.X[:, indices]'*α+thePortfolio.A[:, indices]'*(βl-βu)+λ*ones(f)+ρ-thePortfolio.d[indices]))
    @objective(innerQP, Max, (-0.5*α'*α-sum((thePortfolio.γ[indices]/2.0).*w.*w)+thePortfolio.Y'*α+λ+βl'*thePortfolio.l-βu'*thePortfolio.u+ρ'*thePortfolio.min_investment[indices])')

    status=solve(innerQP; suppress_warnings=true)
    return Dual(getvalue(α), getvalue(λ), getvalue(βl), getvalue(βu), getvalue(ρ), getvalue(w), getobjectivevalue(innerQP), status)
end

function inner_dual2(thePortfolio::SparsePortfolioData, indices::Array{Int64}, s::Array{Float64})

    n=thePortfolio.f
    m=thePortfolio.m
    f=size(indices,1)

    innerQP = Model(solver=MosekSolver(MSK_DPAR_INTPNT_QO_TOL_PFEAS=1e-6, MSK_DPAR_INTPNT_QO_TOL_DFEAS=1e-6, MSK_IPAR_LOG=0, MSK_IPAR_MAX_NUM_WARNINGS=0)) #
    #

    @variable(innerQP, α[1:n])
    @variable(innerQP, λ)
    @variable(innerQP, w[1:f])
    @variable(innerQP, ρ[1:f]>=0.0)
    @variable(innerQP, βl[1:m]>=0.0)
    @variable(innerQP, βu[1:m]>=0.0)

    @constraint(innerQP, w.>=(thePortfolio.X[:, indices]'*α+thePortfolio.A[:, indices]'*(βl-βu)+λ*ones(f)+ρ-thePortfolio.d[indices]))
    @objective(innerQP, Max, (-0.5*α'*α-sum((thePortfolio.γ[indices]/2.0).*s[indices].*w.*w)+thePortfolio.Y'*α+λ+βl'*thePortfolio.l-βu'*thePortfolio.u+ρ'*(s[indices].*thePortfolio.min_investment[indices]))')

    status=solve(innerQP; suppress_warnings=true)
    return Dual(getvalue(α), getvalue(λ), getvalue(βl), getvalue(βu), getvalue(ρ), getvalue(w), getobjectivevalue(innerQP), status)
end


#Reference https://math.stackexchange.com/questions/1435448/negacyclic-fft-multiplication

module MulFFT
    using FFTW
    struct TwistFFTPlan
        fftplan::FFTW.Plan
        ifftplan::FFTW.Plan
        twist::Array{ComplexF64,1}
    end
    function TwistFFTPlan(N::Int32)
        TwistFFTPlan(FFTW.plan_fft(Float64.(rand(-2^9:2^9,div(N,2)))+im*rand(-2^9:2^9,div(N,2))),FFTW.plan_ifft(Float64.(rand(-2^9:2^9,div(N,2)))+im*rand(-2^9:2^9,div(N,2))),TwistGen(N))
    end
    function TwistGen(N::Int32)
        return [exp(im*k*pi/N) for k in 0:div(N,2)-1]
    end

    function TwistFFT(a::Array{Int32,1},fftplan::FFTW.Plan,twist::Array{Complex{Float64}})
        return fftplan*(complex.(Float64.(a[1:div(length(a),2)]),Float64.(a[div(length(a),2)+1:length(a)])).* twist)
    end
    function TwistIFFT(a::Array{Complex{Float64},1},ifftplan::FFTW.Plan,twist::Array{Complex{Float64}})
        b = (ifftplan*a).* conj(twist)
        return append!(real(b),imag(b))
    end

    function PolyMul(a::Array{Int32,1},b::Array{Int32,1},TwistFFTPlan::TwistFFTPlan)
        return map(x->UInt32(mod(round(Int64,x),Int64(2)^32)),TwistIFFT(TwistFFT(a,TwistFFTPlan.fftplan,TwistFFTPlan.twist).*TwistFFT(b,TwistFFTPlan.fftplan,TwistFFTPlan.twist),TwistFFTPlan.ifftplan,TwistFFTPlan.twist))
    end
end


# SPDX-License-Identifier: MIT

"""
    abstract type Semiring <: Number end

Abstract base type for all semirings.
"""
abstract type Semiring <: Number end

"""
    abstract type Semifield <: Semiring end

Abstract base type for all semifields.
"""
abstract type Semifield <: Semiring end

Base.zero(x::T) where T<:Semiring = zero(T)
Base.one(x::T) where T<:Semiring = one(T)
Base.convert(::Type{T}, x::T) where T<:Semiring = x
Base.convert(T::Type{<:Number}, x::Semiring) = T(x.val)
Base.promote_rule(x::Type{T}, y::Type{<:Real}) where T <: Semiring = T
Base.show(io::IO, x::Semiring) = print(io, x.val)


######################################################################
# Mesh
#######################################################################

abstract type AbstractMesh{D} end

struct Mesh{D,T,V<:AbstractArray{SVector{D,T}}} <: AbstractMesh{D}   # D is dimension of parameter space
    vertices::V                         # Iterable vertex container with SVector{D,T} eltype
    adjmat::SparseMatrixCSC{Bool,Int}   # Undirected graph: both dest > src and dest < src
end

# const Mesh{D,T} = Mesh{D,T,Vector{SVector{D,T}},Vector{Tuple{Int,Vararg{Int,D}}}}

# Mesh{D,T}() where {D,T} = Mesh(SVector{D,T}[], sparse(Int[], Int[], Bool[]), NTuple{D+1,Int}[])

function Base.show(io::IO, mesh::Mesh{D}) where {D}
    i = get(io, :indent, "")
    print(io,
"$(i)Mesh{$D}: mesh of a $D-dimensional manifold
$i  Vertices   : $(nvertices(mesh))
$i  Edges      : $(nedges(mesh))")
end

nvertices(m::Mesh) = length(m.vertices)

nedges(m::Mesh) = div(nnz(m.adjmat), 2)

nsimplices(m::Mesh) = length(simplices(m))

vertices(m::Mesh) = m.vertices

edges(m::Mesh, src) = nzrange(m.adjmat, src)

edgedest(m::Mesh, edge) = rowvals(m.adjmat)[edge]

edgevertices(m::Mesh) =
    ((vsrc, m.vertices[edgedest(m, edge)]) for (i, vsrc) in enumerate(m.vertices) for edge in edges(m, i))

function minmax_edge_length(m::Mesh{D,T}) where {D,T<:Real}
    minlen2 = typemax(T)
    maxlen2 = zero(T)
    verts = vertices(m)
    for src in eachindex(verts), edge in edges(m, src)
        dest = edgedest(m, edge)
        dest > src || continue # Need only directed graph
        vec = verts[dest] - verts[src]
        norm2 = vec' * vec
        norm2 < minlen2 && (minlen2 = norm2)
        norm2 > maxlen2 && (maxlen2 = norm2)
    end
    return sqrt(minlen2), sqrt(maxlen2)
end

######################################################################
# Compute N-simplices (N = number of vertices)
######################################################################
function simplices(mesh::Mesh{D}, ::Val{N} = Val(D+1)) where {D,N}
    N > 0 || throw(ArgumentError("Need a positive number of vertices for simplices"))
    N == 1 && return Tuple.(1:nvertices(mesh))
    simps = NTuple{N,Int}[]
    buffer = (NTuple{N,Int}[], NTuple{N,Int}[], Int[])
    for src in eachindex(vertices(mesh))
        append!(simps, _simplices(buffer, mesh, src))
    end
    N > 2 && alignnormals!(simps, vertices(mesh))
    return simps
end

# Add (greater) neighbors to last vertex of partials that are also neighbors of scr, till N
function _simplices(buffer::Tuple{P,P,V}, mesh, src) where {N,P<:AbstractArray{<:NTuple{N}},V}
    partials, partials´, srcneighs = buffer
    resize!(srcneighs, 0)
    resize!(partials, 0)
    for edge in edges(mesh, src)
        srcneigh = edgedest(mesh, edge)
        srcneigh > src || continue # Directed graph, to avoid simplex duplicates
        push!(srcneighs, srcneigh)
        push!(partials, padright((src, srcneigh), 0, Val(N)))
    end
    for pass in 3:N
        resize!(partials´, 0)
        for partial in partials
            nextsrc = partial[pass - 1]
            for edge in edges(mesh, nextsrc)
                dest = edgedest(mesh, edge)
                dest > nextsrc || continue # If not directed, no need to check
                dest in srcneighs && push!(partials´, modifyat(partial, pass, dest))
            end
        end
        partials, partials´ = partials´, partials
    end
    return partials
end

modifyat(s::NTuple{N,T}, ind, el) where {N,T} = ntuple(i -> i === ind ? el : s[i], Val(N))

function alignnormals!(simplices, vertices)
    for (i, s) in enumerate(simplices)
        volume = elementvolume(vertices, s)
        volume < 0 && (simplices[i] = switchlast(s))
    end
    return simplices
end

# Project N-1 edges onto (N-1)-dimensional vectors to have a deterministic volume
elementvolume(verts, s::NTuple{N,Int}) where {N} =
    elementvolume(hcat(ntuple(i -> padright(SVector(verts[s[i+1]] - verts[s[1]]), Val(N-1)), Val(N-1))...))
elementvolume(mat::SMatrix{N,N}) where {N} = det(mat)

switchlast(s::NTuple{N,T}) where {N,T} = ntuple(i -> i < N - 1 ? s[i] : s[2N - i - 1] , Val(N))

######################################################################
# Special meshes
######################################################################
"""
    marchingmesh(npoints::Integer...; axes = 1.0 * I, shift = missing)

Creates a L-dimensional marching-tetrahedra `Mesh`. The mesh is confined to the box defined
by the rows of `axes`, shifted by `shift` (an `NTuple` or `Number`) if not `missing`, and
contains `npoints[i]` along each axis `i`.

    marchingmesh(ranges::AbstractRange...; axes = 1.0 * I, shift = missing)

The same as above, but allows to specify the points in axis `i` by a range `ranges[i]`, such
as e.g. `-1.0:0.1:1.0`.

Note that the size of `axes` should match the number `L` of elements in `npoints` or
`ranges`. The `eltype` of points is given by that of `ranges` or `axes`.

    marchingmesh(h::Hamiltonian{<:Lattice}; npoints = 13, shift = missing)

Equivalent to `marchingmesh(ntuple(_ -> range(-π, π; length = npoints), Val(L))...)` where
`L` is the dimension of the Hamiltonian's lattice.

# External links

- Marching tetrahedra (https://en.wikipedia.org/wiki/Marching_tetrahedra) in Wikipedia
"""
marchingmesh(ranges::Vararg{AbstractRange,L}; axes = 1.0 * I, shift = missing) where {L} =
    _marchingmesh(shiftranges(ranges, shift), SMatrix{L,L}(axes))

function marchingmesh(npoints::Vararg{Integer,L}; axes = 1.0 * I, shift = missing) where {L}
    ranges = meshranges(npoints, Val(L))
    ranges´ = shiftranges(ranges, shift)
    return _marchingmesh(ranges´, SMatrix{L,L}(axes))
end

function marchingmesh(h::Hamiltonian{<:Lattice,L}; npoints = 13, shift = missing) where {L}
    ranges = meshranges(npoints, Val(L))
    ranges´ = shiftranges(ranges, shift)
    return _marchingmesh(ranges´, SMatrix{L,L}(I))
end

meshranges(n::Int, ::Val{L}) where {L} = meshranges(filltuple(n, Val(L)), Val(L))
meshranges(ns::NTuple{L,Int}, ::Val{L}) where {L} = (n -> range(-π, π; length = n)).(ns)

shiftranges(rs, shift::Missing) = rs
shiftranges(rs::NTuple{L}, shift::Number) where {L} = shiftranges(rs, filltuple(shift, Val(L)))
shiftranges(rs::NTuple{L}, shifts::NTuple{L}) where {L} = ((r, s) -> r .+ s).(rs, shifts)

function _marchingmesh(ranges::NTuple{D,AbstractRange}, axes::SMatrix{D,D}) where {D,T<:AbstractFloat}
    npoints = length.(ranges)
    projection = axes' # ./ (SVector(npoints) - 1)' # Projects binary vector to m box with npoints
    cs = CartesianIndices(ntuple(n -> 1:npoints[n], Val(D)))
    ls = LinearIndices(cs)
    csinner = CartesianIndices(ntuple(n -> 1:npoints[n]-1, Val(D)))

    # edge vectors for marching tetrahedra in D-dimensions (skip zero vector [first])
    uedges = [c for c in CartesianIndices(ntuple(_ -> 0:1, Val(D)))][2:end]
    # tetrahedra built from the D unit-length uvecs added in any permutation
    perms = permutations(
            ntuple(i -> CartesianIndex(ntuple(j -> i == j ? 1 : 0, Val(D))), Val(D)))
    utets = [cumsum(pushfirst!(perm, zero(CartesianIndex{D}))) for perm in perms]

    # We don't use generators because their non-inferreble eltype causes problems later
    verts = [projection * SVector(getindex.(ranges, Tuple(c))) for c in cs]

    s = SparseMatrixBuilder{Bool}(length(cs), length(cs))
    for c in cs
        for u in uedges
            dest = c + u    # dest > src
            dest in cs && pushtocolumn!(s, ls[dest], true)
            dest = c - u    # dest < src
            dest in cs && pushtocolumn!(s, ls[dest], true)
        end
        finalizecolumn!(s)
    end
    adjmat = sparse(s)

    return Mesh(verts, adjmat)
end


"""
    SemidirectProductOperation(action::AbstractGroupAction)

Group operation of a semidirect product group. The operation consists of the operation
`opN` on a normal subgroup `N`, the operation `opH` on a subgroup `H`, and an automorphism
`action` of elements of `H` on `N`. Only the action is stored.
"""
struct SemidirectProductOperation{A<:AbstractGroupAction} <: AbstractGroupOperation
    action::A
end

function show(io::IO, op::SemidirectProductOperation)
    print(io, "SemidirectProductOperation($(op.action))")
end

const SemidirectProductGroup{N,H,A} =
    GroupManifold{ProductManifold{Tuple{N,H}},SemidirectProductOperation{A}}

@doc raw"""
    SemidirectProductGroup(N::GroupManifold, H::GroupManifold, A::AbstractGroupAction)

A group that is the semidirect product of a normal group $\mathcal{N}$ and a subgroup
$\mathcal{H}$, written $\mathcal{G} = \mathcal{N} ⋊_θ \mathcal{H}$, where
$θ: \mathcal{H} × \mathcal{N} → \mathcal{N}$ is an automorphism action of $\mathcal{H}$ on
$\mathcal{N}$. The group $\mathcal{G}$ has the composition rule

````math
g \circ g' = (n, h) \circ (n', h') = (n \circ θ_h(n'), h \circ h')
````

and the inverse

````math
g^{-1} = (n, h)^{-1} = (θ_{h^{-1}}(n^{-1}), h^{-1}).
````
"""
function SemidirectProductGroup(N::GroupManifold, H::GroupManifold, A::AbstractGroupAction)
    N === g_manifold(A) || error("Subgroup $(N) must be the G-manifold of action $(A)")
    H === base_group(A) || error("Subgroup $(H) must be the base group of action $(A)")
    op = SemidirectProductOperation(A)
    M = ProductManifold(N, H)
    return GroupManifold(M, op)
end

function show(io::IO, G::SemidirectProductGroup)
    M = base_manifold(G)
    N, H = M.manifolds
    A = G.op.action
    print(io, "SemidirectProductGroup($(N), $(H), $(A))")
end

submanifold(G::SemidirectProductGroup, i) = submanifold(base_manifold(G), i)

_padpoint!(G::SemidirectProductGroup, q) = q

_padvector!(G::SemidirectProductGroup, X) = X

inv(G::GT, e::Identity{GT}) where {GT<:SemidirectProductGroup} = e

function inv!(G::SemidirectProductGroup, q, p)
    M = base_manifold(G)
    N, H = M.manifolds
    A = G.op.action
    np, hp = submanifold_components(G, p)
    nq, hq = submanifold_components(G, q)
    inv!(H, hq, hp)
    npinv = inv(N, np)
    apply!(A, nq, hq, npinv)
    @inbounds _padpoint!(G, q)
    return q
end
inv!(G::AG, p, e::Identity{AG}) where {AG<:SemidirectProductGroup} = identity!(G, p, e)

identity(G::GT, e::Identity{GT}) where {GT<:SemidirectProductGroup} = e

function identity!(G::SemidirectProductGroup, q, p)
    M = base_manifold(G)
    N, H = M.manifolds
    np, hp = submanifold_components(G, p)
    nq, hq = submanifold_components(G, q)
    identity!(N, nq, np)
    identity!(H, hq, hp)
    @inbounds _padpoint!(G, q)
    return q
end
identity!(G::GT, e::E, ::E) where {GT<:SemidirectProductGroup,E<:Identity{GT}} = e

compose(G::GT, p, e::Identity{GT}) where {GT<:SemidirectProductGroup} = p
compose(G::GT, e::Identity{GT}, p) where {GT<:SemidirectProductGroup} = p
compose(G::GT, e::E, ::E) where {GT<:SemidirectProductGroup,E<:Identity{GT}} = e

function compose!(G::SemidirectProductGroup, x, p, q)
    M = base_manifold(G)
    N, H = M.manifolds
    A = G.op.action
    np, hp = submanifold_components(G, p)
    nq, hq = submanifold_components(G, q)
    nx, hx = submanifold_components(G, x)
    compose!(H, hx, hp, hq)
    nxtmp = apply(A, hp, nq)
    compose!(N, nx, np, nxtmp)
    @inbounds _padpoint!(G, x)
    return x
end
compose!(G::GT, x, ::Identity{GT}, q) where {GT<:SemidirectProductGroup} = copyto!(x, q)
compose!(G::GT, x, p, ::Identity{GT}) where {GT<:SemidirectProductGroup} = copyto!(x, p)
function compose!(G::GT, x, e::E, ::E) where {GT<:SemidirectProductGroup,E<:Identity{GT}}
    return identity!(G, x, e)
end

function translate_diff!(G::SemidirectProductGroup, Y, p, q, X, conX::LeftAction)
    M = base_manifold(G)
    N, H = M.manifolds
    A = G.op.action
    np, hp = submanifold_components(G, p)
    nq, hq = submanifold_components(G, q)
    nX, hX = submanifold_components(G, X)
    nY, hY = submanifold_components(G, Y)
    translate_diff!(H, hY, hp, hq, hX, conX)
    nZ = apply_diff(A, hp, nq, nX)
    nr = apply(A, hp, nq)
    translate_diff!(N, nY, np, nr, nZ, conX)
    @inbounds _padvector!(G, Y)
    return Y
end

function get_vector!(G::SemidirectProductGroup, Y, p, X, B::VeeOrthogonalBasis)
    M = base_manifold(G)
    N, H = M.manifolds
    dimN = manifold_dimension(N)
    dimH = manifold_dimension(H)
    @assert length(X) == dimN + dimH
    np, hp = submanifold_components(G, p)
    nY, hY = submanifold_components(G, Y)
    get_vector!(N, nY, np, view(X, 1:dimN), B)
    get_vector!(H, hY, hp, view(X, dimN+1:dimN+dimH), B)
    @inbounds _padvector!(G, Y)
    return Y
end
eval(quote
    @invoke_maker 1 Manifold get_vector!(
        M::SemidirectProductGroup,
        Xⁱ,
        e::Identity,
        X,
        B::VeeOrthogonalBasis,
    )
end)

function get_coordinates!(G::SemidirectProductGroup, Y, p, X, B::VeeOrthogonalBasis)
    M = base_manifold(G)
    N, H = M.manifolds
    dimN = manifold_dimension(N)
    dimH = manifold_dimension(H)
    @assert length(Y) == dimN + dimH
    np, hp = submanifold_components(G, p)
    nY, hY = submanifold_components(G, X)
    get_coordinates!(N, view(Y, 1:dimN), np, nY, B)
    get_coordinates!(H, view(Y, dimN+1:dimN+dimH), hp, hY, B)
    return Y
end
eval(quote
    @invoke_maker 1 Manifold get_coordinates!(
        M::SemidirectProductGroup,
        Y,
        e::Identity,
        X,
        B::VeeOrthogonalBasis,
    )
end)

function zero_tangent_vector(G::SemidirectProductGroup, p)
    X = allocate_result(G, zero_tangent_vector, p)
    zero_tangent_vector!(G, X, p)
    return X
end

function zero_tangent_vector!(G::SemidirectProductGroup, X, p)
    M = base_manifold(G)
    N, H = M.manifolds
    np, hp = submanifold_components(G, p)
    nX, hX = submanifold_components(G, X)
    zero_tangent_vector!(N, nX, np)
    zero_tangent_vector!(H, hX, hp)
    return X
end

function isapprox(G::SemidirectProductGroup, p, q; kwargs...)
    M = base_manifold(G)
    N, H = M.manifolds
    np, hp = submanifold_components(G, p)
    nq, hq = submanifold_components(G, q)
    return isapprox(N, np, nq; kwargs...) && isapprox(H, hp, hq; kwargs...)
end
function isapprox(G::SemidirectProductGroup, p, X, Y; kwargs...)
    M = base_manifold(G)
    N, H = M.manifolds
    np, hp = submanifold_components(G, p)
    nX, hX = submanifold_components(G, X)
    nY, hY = submanifold_components(G, Y)
    return isapprox(N, np, nX, nY; kwargs...) && isapprox(H, hp, hX, hY; kwargs...)
end
function isapprox(G::GT, p, e::Identity{GT}; kwargs...) where {GT<:SemidirectProductGroup}
    return isapprox(G, e, p; kwargs...)
end
function isapprox(G::GT, e::Identity{GT}, p; kwargs...) where {GT<:SemidirectProductGroup}
    return isapprox(G, identity(G, p), p; kwargs...)
end
function isapprox(
    ::GT,
    ::E,
    ::E;
    kwargs...,
) where {GT<:SemidirectProductGroup,E<:Identity{GT}}
    return true
end


using DataStructures

input = joinpath(@__DIR__, "input")
lines = readlines(input)

m = hcat(map(line -> map(n -> parse(Int, n), collect(line)), lines)...)

function expand(m)
    acc = [m]
    for i in 1:4
        push!(acc, map(x -> mod1(x, 9), m.+ i))
    end
    tmp = hcat(acc...)
    acc = [tmp]
    for i in 1:4
        n = tmp.+ 1
        push!(acc, map(x -> mod1(x, 9), tmp.+ i))
    end
    return vcat(acc...)
end

function neighbors(m, idx)
    x, y = Tuple(idx)
    return filter(c -> checkbounds(Bool, m, c), [
        CartesianIndex(x + 1, y),
        CartesianIndex(x - 1, y),
        CartesianIndex(x, y + 1),
        CartesianIndex(x, y - 1)
    ])
end

function path(m, origin, dest)
    unvisited = Set{}(collect(CartesianIndices(m)))
    candidates = PriorityQueue{CartesianIndex, Int}()
    node = origin
    # initial node is free
    cost = 0

    while dest in unvisited
        delete!(unvisited, node)
        for n in neighbors(m, node)
            if n in unvisited
                haskey(candidates, n) ?
                    candidates[n] = minimum([candidates[n], cost + m[n]]) :
                    candidates[n] = cost + m[n]
            end
        end
        node, cost = dequeue_pair!(candidates)
        node == dest && delete!(unvisited, node)
    end
    return cost
end

p1 = path(m, CartesianIndex(1, 1), CartesianIndex(size(m)))
m = expand(m)
p2 = path(m, CartesianIndex(1, 1), CartesianIndex(size(m)))

println("-----------------------------------------------------------------------")
println("chiton -- part one :: $p1")
println("chiton -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 790)
@assert(p2 == 2998)


using StaticArrays
using Plots
using BenchmarkTools

function f(u, p)
    α, β, γ, δ = p

    du1 = α*u[1] - β*u[1]*u[2]
    du2 = -γ*u[2] + δ*u[1]*u[2]

    @SVector [du1, du2]
end

function dopri5!(u_cache, u, p, t_span, f)
    # Values from the Dormand Prince Tableau
    a21  =  (1.0/5.0)
    a31  = (3.0/40.0)
    a32  = (9.0/40.0)
    a41  = (44.0/45.0)
    a42  = (-56.0/15.0)
    a43  = (32.0/9.0)
    a51  = (19372.0/6561.0)
    a52  = (-25360.0/2187.0)
    a53  = (64448.0/6561.0)
    a54  = (-212.0/729.0)
    a61  = (9017.0/3168.0)
    a62  = (-355.0/33.0)
    a63  = (46732.0/5247.0)
    a64  = (49.0/176.0)
    a65  = (-5103.0/18656.0)
    a71  = (35.0/384.0)
    a72  = (0.0)
    a73  = (500.0/1113.0)
    a74  = (125.0/192.0)
    a75  = (-2187.0/6784.0)
    a76  = (11.0/84.0)

    b1   = (35.0/384.0)
    b2   = (0.0)
    b3   = (500.0/1113.0)
    b4   = (125.0/192.0)
    b5   = (-2187.0/6784.0)
    b6   = (11.0/84.0)
    b7   = (0.0)

    dt = 0.25

    @inbounds for i in 2:Int64(t_span/dt)
        k1 = f(u, p)
        k2 = f(u + dt.*(a21.*k1), p)
        k3 = f(u + dt.*(a31.*k1 + a32.*k2), p)
        k4 = f(u + dt.*(a41.*k1 + a42.*k2 + a43.*k3), p)
        k5 = f(u + dt.*(a51.*k1 + a52.*k2 + a53.*k3 + a54.*k4), p)
        k6 = f(u + dt.*(a61.*k1 + a62.*k2 + a63.*k3 + a64.*k4 + a65.*k5), p)
        u = u + dt.*(b1.*k1 + b3.*k3 + b4.*k4 + b5.*k5 + b6.*k6)
        u_cache[:, i] = u
    end
end

p = (1.5, 1.0, 3.0, 1.0)
u₀ = [1.0, 1.0]
u_cache = Array{Float64}(undef, 2, 40)
u_cache[:, 1] = u₀
t_span = 10.0
u_cache

@btime dopri5!(u_cache, u₀, p, t_span, f)

dopri5!(u_cache, u₀, p, t_span, f)

u_cache

plot(0.25:0.25:t_span, u_cache[1, :])
plot!(0.25:0.25:t_span, u_cache[2, :])


#=
zwart_powell.jl
Translated by Harshit Nanda from ir_radon_zwart_powell.m in 2022.
Matlab author: Seongjin Yoon, Univ. of Michigan, Jan 2015
=#

export zwart_powell

const TOLERANCE = 1e-07

"""
    output = zwart_powell(r, ϕ)
Analytic 2D Radon transform value of Zwart-Powell box spline,
for radial distance `r` (normalized by pixel size) and angle `ϕ` (in radians).
"""
function zwart_powell(r::Real, ϕ::Real)
    s, c = sincos(ϕ)
    return zwart_powell(r, (c, s, c+s, c-s))
end

function zwart_powell(r::Real, z::NTuple{4,T}) where T <: Real
    N = sum(z -> abs(z) > TOLERANCE, z)
    return BoxSp4(r, z..., N) / factorial(N-1)
end

BoxSp0(y::RealU, N::Int) = y > 0 ? y^(N-1) : zero(y)

BoxSp1(y::RealU, z1::T, N::Int) where T <: Real =
    abs(z1) > TOLERANCE ?
        (BoxSp0(y + 0.5 * z1, N) -
         BoxSp0(y - 0.5 * z1, N)) / z1 :
        BoxSp0(y, N)

BoxSp2(y::RealU, z1::T, z2::T, N::Int) where T <: Real =
    abs(z2) > TOLERANCE ?
        (BoxSp1(y + 0.5 * z2, z1, N) -
         BoxSp1(y - 0.5 * z2, z1, N)) / z2 :
        BoxSp1(y, z1, N)

BoxSp3(y::RealU, z1::T, z2::T, z3::T, N::Int) where T <: Real =
    abs(z3) > TOLERANCE ?
        (BoxSp2(y + 0.5 * z3, z1, z2, N) -
         BoxSp2(y - 0.5 * z3, z1, z2, N)) / z3 :
        BoxSp2(y, z1, z2, N)

function BoxSp4(y::RealU, z1::T, z2::T, z3::T, z4::T, N::Int) where T <: Real
    return abs(z4) > TOLERANCE ?
        (BoxSp3(y + 0.5 * z4, z1, z2, z3, N) -
         BoxSp3(y - 0.5 * z4, z1, z2, z3, N)) / z4 :
        BoxSp3(y, z1, z2, z3, N)
end


# see tutorial https://turinglang.github.io/TuringCallbacks.jl/dev/
#turing test

module TuringCLIExample

using Turing, TuringCallbacks, TensorBoardLogger, ArgParse, StatsPlots, CSV, DataFrames

function simulate_and_estimate(;
    num_samples,
    num_adapts,
    target_acceptance_rate,
    s_prior_alpha,
    s_prior_theta,
    kwargs...,
)
    @model function demo(x; s_prior_alpha, s_prior_theta)
        s ~ InverseGamma(s_prior_alpha, s_prior_theta)
        m ~ Normal(0, √s)
        for i in eachindex(x)
            x[i] ~ Normal(m, √s)
        end
    end

    xs = randn(100) .+ 1
    model = demo(xs; s_prior_alpha, s_prior_theta)

    # Sampling
    println("Generating $num_samples samples")
    callback = TensorBoardCallback("tensorboard_logs/run")
    alg = NUTS(num_adapts, target_acceptance_rate)
    chain = sample(model, alg, num_samples; callback)

    println("Generating trace plot")
    trace_plot = plot(chain, seriestype=:traceplot)
    savefig(trace_plot,  joinpath(callback.logger.logdir, "traceplots.png"))

    println("Summarizing the chain")
    sum_stats = describe(chain)
    param_names = sum_stats[1][:,1]
    param_mean = sum_stats[1][:,2]
    param_sd = sum_stats[1][:,3]
    param_ess = sum_stats[1][:,6]
    param_rhat = sum_stats[1][:,7]
    param_ess_per_sec = sum_stats[1][:, 8]

    CSV.write(
        joinpath(callback.logger.logdir, "summary.csv"),
        DataFrame(
            parameter=param_names,
            mean=param_mean,
            sd=param_sd,
            ess=param_ess,
            rhat=param_rhat,
            ess_per_sec= param_ess_per_sec
        )
    )

    # Log the ESS/sec and rhat.  Nice to show as summary results from tensorboard
    for (i, name) = enumerate(param_names)
        TensorBoardLogger.log_value(
            callback.logger,
            "$(name)_ess_per_sec",
            param_ess_per_sec[i],
        )
        TensorBoardLogger.log_value(
            callback.logger,
            "$(name)_rhat",
            param_rhat[i],
        )
    end
end

# Entry for script
function main()
    d = parse_commandline()

    # Generic code to convert the dictionary to
    dictkeys = (collect(Symbol.(keys(d)))...,)
    dictvalues = (collect(values(d))...,)
    args = NamedTuple{dictkeys}(dictvalues)
    # parses converts all arguments to named tuple then splat into solution
    simulate_and_estimate(; args...)
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--num_samples"
        help = "samples to draw in chain"
        arg_type = Int64
        default = 10000
        "--num_adapts"
        help = "number of adaptations for NUTS"
        arg_type = Int64
        default = 100
        "--target_acceptance_rate"
        help = "Target acceptance rate for dual averaging."
        arg_type = Float64
        default = 0.65
        "--s_prior_alpha"
        help = "alpha in InverseGamma prior for s"
        arg_type = Float64
        default = 2.0
        "--s_prior_theta"
        help = "theta in InverseGamma prior for s"
        arg_type = Float64
        default = 3.0
    end

    return parse_args(s)
end

# Can experiment later, or can package compiler
# function _precompile_()
#     ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
#     Base.precompile(Tuple{typeof(main)})
# end

# # Only really intended for final versions rather than during interactive use
# # NOTE: Consider commenting out or fix if changing signature of simulate_and_estimte
# if Base.VERSION >= v"1.4.2"
#     precompile(Tuple{typeof(TuringCLIExample.main)})

#     # Is this necessary or does the main() above do it? Comment out if any problem.
#     precompile(Tuple{TuringCLIExample.var"#simulate_and_estimate##kw", NamedTuple{(:num_samples, :num_adapts, :target_acceptance_rate, :s_prior_theta, :s_prior_alpha), Tuple{Int64, Int64, Float64, Float64, Float64}}, typeof(TuringCLIExample.simulate_and_estimate)})

#     _precompile_()
# end
end #module


@testset "pwmeva.jl" begin

    θ = [0.0, 0.0, 0.1]

    pd = GeneralizedExtremeValue(θ[1], exp(θ[2]), θ[3])

    y = [0]

    fm = Extremes.pwmEVA{BlockMaxima, GeneralizedExtremeValue}(Extremes.BlockMaxima(Variable("y", y)), θ)

    @testset "getdistribution(fittedmodel)" begin
        @test Extremes.getdistribution(fm)[] == pd
    end

    @testset "quantile(fm, p)" begin
        # p outside of [0, 1] throws
        @test_throws AssertionError Extremes.quantile(fm, -1)

        # Test with known values
        @test quantile(fm, .99)[] ≈ quantile(pd,.99)

    end

    @testset "returnlevel(fm, returnPeriod)" begin

        # returnPeriod < 0 throws
        @test_throws AssertionError returnlevel(fm, -1)

        # Test with known values
        @test returnlevel(fm, 100).value[] ≈ quantile(pd, 1-1/100)
    end

    @testset "parametervar(fm, nboot)" begin
        # nboot < 0 throws
        @test_throws AssertionError Extremes.parametervar(fm, -1)

        # Tested in cint(fm)
    end

    @testset "quantilevar(fm, level, nboot)" begin
        # Tested in cint(fm)
    end

    @testset "cint(fm)" begin
        # returnPeriod < 0 throws
        @test_throws AssertionError cint(ReturnLevel(Extremes.BlockMaximaModel(fm), -1, [1.0]), 0.95)

        # confidencelevel not in [0, 1]
        @test_throws AssertionError cint(ReturnLevel(Extremes.BlockMaximaModel(fm), 1, [1.0]), -1)

        niter = 1000

        cover = falses(niter,3)

        Threads.@threads  for i in 1:niter
            y = rand(pd, 100)
            fm = gevfitpwm(y)
            ci = cint(fm, .95 ,1000)
            for j in eachindex(θ)
                cover[i,j] = ci[j][1] < θ[j] < ci[j][2]
            end
        end

        @test all(count.(eachcol(cover))/niter .> .85)
    end

    @testset "cint(fm, returnPeriod, confidencelevel)" begin
        # confidencelevel not in [0, 1]
        @test_throws AssertionError cint(ReturnLevel(Extremes.BlockMaximaModel(fm), 1, [1.0]), -1)

        # Test with known values

        θ = [100; log(5); .1]

        pd = GeneralizedExtremeValue(θ[1],exp(θ[2]),θ[3])

        q = quantile(pd, 1-1/100)

        niter = 1200

        cover = falses(niter)

        Threads.@threads  for i in 1:niter
            y = rand(pd, 200)
            fm = gevfitpwm(y)
            rl = returnlevel(fm, 100)
            ci = cint(rl, .95 , 500)
            cover[i] = ci[][1] < q < ci[][2]
        end

        @test count(cover)/niter > .85
    end

    @testset "showfittedEVA(io, obj, prefix)" begin
        # print does not throw
        buffer = IOBuffer()
        @test_logs Extremes.showfittedEVA(buffer, fm, prefix = "\t")

    end

end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using DelimitedFiles
Random.seed!(233)

function halo_exchange_two_d(u,fill_value,m,n)
    halo_exchange_two_d_ = load_op_and_grad("./build/libHaloExchangeTwoD","halo_exchange_two_d")
    u,fill_value,m,n = convert_to_tensor(Any[u,fill_value,m,n], [Float64,Float64,Int64,Int64])
    halo_exchange_two_d_(u,fill_value,m,n)
end

mpi_init()
U = reshape(1:24, 4, 6)'|>Array

m = 3
n = 2
fill_value = 1.0

M = mpi_rank()÷n+1
N = mpi_rank()%n+1
ulocal = U[(M-1)*2 + 1: M * 2, (N-1)*2+1:N*2]
# TODO: specify your input parameters
u = halo_exchange_two_d(ulocal,fill_value,m,n)
sess = Session(); init(sess)
Uval = run(sess, u)

sleep(mpi_rank())
println("=========================== rank = $(mpi_rank()) ======================")
writedlm(stdout, Uval)
mpi_finalize()


@enum Direction FFT_FORWARD=-1 FFT_BACKWARD=1
@enum Pow24 POW2=2 POW4=1

abstract type AbstractFFTType end

# Represents a Composite Cooley-Tukey FFT
struct CompositeFFT <: AbstractFFTType end

# Represents a Radix-2 Cooley-Tukey FFT
struct Pow2FFT <: AbstractFFTType end

# Represents a Radix-3 Cooley-Tukey FFT
struct Pow3FFT <: AbstractFFTType end

# Represents a Radix-4 Cooley-Tukey FFT
struct Pow4FFT <: AbstractFFTType end

# Represents an O(N²) DFT
struct DFT <: AbstractFFTType end

"""
$(TYPEDSIGNATURES)
Node of a call graph

# Arguments
`left`: Offset to the left child node
`right`: Offset to the right child node
`type`: Object representing the type of FFT
`sz`: Size of this FFT

"""
struct CallGraphNode
    left::Int
    right::Int
    type::AbstractFFTType
    sz::Int
    s_in::Int
    s_out::Int
end

"""
$(TYPEDSIGNATURES)
Object representing a graph of FFT Calls

# Arguments
`nodes`: Nodes keeping track of the graph
`workspace`: Preallocated Workspace

"""
struct CallGraph{T<:Complex}
    nodes::Vector{CallGraphNode}
    workspace::Vector{Vector{T}}
end

# Get the node in the graph at index i
Base.getindex(g::CallGraph{T}, i::Int) where {T} = g.nodes[i]

"""
$(TYPEDSIGNATURES)
Check if `N` is a power of `base`

"""
function _ispow(N, base)
    while N % base == 0
        N = N/base
    end
    return N == 1
end

"""
$(TYPEDSIGNATURES)
Check if `N` is a power of 2 or 4

"""
function _ispow24(N::Int)
    N < 1 && return nothing
    while N & 0b11 == 0
        N >>= 2
    end
    return N < 3 ? Pow24(N) : nothing
end

"""
$(TYPEDSIGNATURES)
Recursively instantiate a set of `CallGraphNode`s

# Arguments
`nodes`: A vector (which gets expanded) of `CallGraphNode`s
`N`: The size of the FFT
`workspace`: A vector (which gets expanded) of preallocated workspaces
`s_in`: The stride of the input
`s_out`: The stride of the output

"""
function CallGraphNode!(nodes::Vector{CallGraphNode}, N::Int, workspace::Vector{Vector{T}}, s_in::Int, s_out::Int)::Int where {T}
    if iseven(N)
        pow = _ispow24(N)
        if !isnothing(pow)
            push!(workspace, T[])
            push!(nodes, CallGraphNode(0, 0, pow == POW2 ? Pow2FFT() : Pow4FFT(), N, s_in, s_out))
            return 1
        end
    end
    if N % 3 == 0
        if _ispow(N, 3)
            push!(workspace, T[])
            push!(nodes, CallGraphNode(0, 0, Pow3FFT(), N, s_in, s_out))
            return 1
        end
    end
    if isprime(N)
        push!(workspace, T[])
        push!(nodes, CallGraphNode(0,0, DFT(),N, s_in, s_out))
        return 1
    end
    Ns = [first(x) for x in collect(factor(N)) for _ in 1:last(x)]
    if Ns[1] == 2
        N1 = prod(Ns[Ns .== 2])
    elseif Ns[1] == 3
        N1 = prod(Ns[Ns .== 3])
    else
        # Greedy search for closest factor of N to sqrt(N)
        Nsqrt = sqrt(N)
        N_cp = cumprod(Ns[end:-1:1])[end:-1:1]
        N_prox = abs.(N_cp .- Nsqrt)
        _,N1_idx = findmin(N_prox)
        N1 = N_cp[N1_idx]
    end
    N2 = N ÷ N1
    push!(nodes, CallGraphNode(0,0,DFT(),N,s_in,s_out))
    sz = length(nodes)
    push!(workspace, Vector{T}(undef, N))
    left_len = CallGraphNode!(nodes, N1, workspace, N2, N2*s_out)
    right_len = CallGraphNode!(nodes, N2, workspace, N1*s_in, 1)
    nodes[sz] = CallGraphNode(1, 1 + left_len, CompositeFFT(), N, s_in, s_out)
    return 1 + left_len + right_len
end

"""
$(TYPEDSIGNATURES)
Instantiate a CallGraph from a number `N`

"""
function CallGraph{T}(N::Int) where {T}
    nodes = CallGraphNode[]
    workspace = Vector{Vector{T}}()
    CallGraphNode!(nodes, N, workspace, 1, 1)
    CallGraph(nodes, workspace)
end


"""
    init_input_layer(res_size::Int, in_size::Int, sigma::Float64)

Return a weighted input layer matrix, with random non-zero elements drawn from \$ [-\\text{sigma}, \\text{sigma}] \$, as described in [1].

[1] Lu, Zhixin, et al. "Reservoir observers: Model-free inference of unmeasured variables in chaotic systems." Chaos: An Interdisciplinary Journal of Nonlinear Science 27.4 (2017): 041102.
"""
function init_input_layer(res_size::Int,
        in_size::Int,
        sigma::Float64)

    W_in = zeros(Float64, res_size, in_size)
    q = floor(Int, res_size/in_size) #need to fix the reservoir input size. Check the constructor
    for i=1:in_size
        W_in[(i-1)*q+1 : (i)*q, i] = (2*sigma).*(rand(Float64, 1, q).-0.5)
    end
    return W_in

end

"""
    init_dense_input_layer(res_size::Int, in_size::Int, sigma::Float64)

Return a fully connected input layer matrix, with random non-zero elements drawn from \$ [-sigma, sigma] \$.
"""
function init_dense_input_layer(res_size::Int,
        in_size::Int,
        sigma::Float64)

    W_in = rand(Float64, res_size, in_size)
    W_in = 2.0 .*(W_in.-0.5)
    W_in = sigma .*W_in
    return W_in
end

"""
    init_sparse_input_layer(res_size::Int, in_size::Int, sigma::Float64, sparsity::Float64)

Return a sparsely connected input layer matrix, with random non-zero elements drawn from \$ [-sigma, sigma] \$ and given sparsity.
"""
function init_sparse_input_layer(res_size::Int,
        in_size::Int,
        sigma::Float64,
        sparsity::Float64)

    W_in = Matrix(sprand(Float64, res_size, in_size, sparsity))
    W_in = 2.0 .*(W_in.-0.5)
    replace!(W_in, -1.0=>0.0)
    W_in = sigma .*W_in
    return W_in
end

#from "minimum complexity echo state network" Rodan
"""
    min_complex_input(res_size::Int, in_size::Int, weight::Float64)

Return a fully connected input layer matrix with the same weights and sign drawn from a Bernoulli distribution, as described in [1].

[1] Rodan, Ali, and Peter Tino. "Minimum complexity echo state network." IEEE transactions on neural networks 22.1 (2010): 131-144.
"""
function min_complex_input(res_size::Int,
        in_size::Int,
        weight::Float64)

    W_in = Array{Float64}(undef, res_size, in_size)
    for i=1:res_size
        for j=1:in_size
            if rand(Bernoulli()) == true
                W_in[i, j] = weight
            else
                W_in[i, j] = -weight
            end
        end
    end
    return W_in
end

#from "minimum complexity echo state network" Rodan
#and "simple deterministically constructed cycle reservoirs with regular jumps" by Rodan and Tino

"""
    irrational_sign_input(res_size::Int, in_size::Int , weight::Float64 [, start::Int, irrational::Irrational])

Return a fully connected input layer matrix with the same weights and sign decided by the values of an irrational number, as described in [1] and [2].

[1] Rodan, Ali, and Peter Tino. "Minimum complexity echo state network." IEEE transactions on neural networks 22.1 (2010): 131-144.
[2] Rodan, Ali, and Peter Tiňo. "Simple deterministically constructed cycle reservoirs with regular jumps." Neural computation 24.7 (2012): 1822-1852.
"""
function irrational_sign_input(res_size::Int,
        in_size::Int,
        weight::Float64;
        start::Int = 1,
        irrational::Irrational = pi)

    setprecision(BigFloat, Int(ceil(log2(10)*(res_size*in_size+start+1))))
    ir_string = string(BigFloat(irrational)) |> collect
    deleteat!(ir_string, findall(x->x=='.', ir_string))
    ir_array = Array{Int}(undef, length(ir_string))
    W_in = Array{Float64}(undef, res_size, in_size)

    for i =1:length(ir_string)
        ir_array[i] = parse(Int, ir_string[i])
    end

    counter = start

    for i=1:res_size
        for j=1:in_size
            if ir_array[counter] < 5
                W_in[i, j] = -weight
            else
                W_in[i, j] = weight
            end
            counter += 1
        end
    end
    return W_in
end


using Colors, ColorSchemes
using DataFrames
using StatsPlots
using Statistics
using GLM
using Measures
using LsqFit

pyplot(grid=false)
fnt = Plots.font("Arial")
default(titlefont=fnt, guidefont=fnt, tickfont=fnt, legendfont=fnt)
#  Plots.resetfontsizes()

noise_level = 0.003544

###############################################################################
### FIG S3A
###############################################################################

function get_pct_off(t_recruit, t_release, k)
    ks = k['s']
    ka = k['a']
    ki = k['i']
    y1 = (ks+ka+ki - sqrt((ks+ka+ki)^2 - 4*ks*ki))/2
    y2 = (ks+ka+ki + sqrt((ks+ka+ki)^2 - 4*ks*ki))/2

    function c_r(t_recruit)
        return ks/(y2-y1) * (exp(-y1*t_recruit) - exp(-y2*t_recruit))
    end

    function c_i(t_recruit)
        if (ki==0)
            return 0
        else
            term1 = (ks*ki)/(y1*(y2-y1)) * exp(-y1*t_recruit)
            term2 = (ks*ki)/(y2*(y2-y1)) * exp(-y2*t_recruit)
            return 1 - term1 + term2
        end
    end

    return c_r(t_recruit) * exp(-ka * t_release) + c_i(t_recruit)
end

# see https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=1023694
function get_single_cell_cc(t_recruit, t_release, k)
    p = 1 - get_pct_off(t_recruit, t_release, k) # error probability
    return log2(1 + (1-p)^(p/(1-p)))
end

function get_population_cc(t_recruit, t_release, k, variance=0.05)
    power = get_pct_off(t_recruit, t_release, k)
    return maximum([0, 0.5 * log2(1 + power/variance)])
end

###############################################################################
### FIG 5D
###############################################################################

t_recruits = collect(0:0.001:1)
t_release = 0
k_p_ccs = [get_population_cc(t_recruit, t_release, k_krab, noise_level) for t_recruit in t_recruits]
h_p_ccs = [get_population_cc(t_recruit, t_release, k_hdac, noise_level) for t_recruit in t_recruits]

p1 = plot(t_recruits, k_p_ccs,
         label="KRAB",
         color=color_map["KRAB"],
         linewidth=2)
p1 = plot!(t_recruits, h_p_ccs,
          label="HDAC4",
          color=color_map["HDAC4"],
          linewidth=2,
          xlabel="Days Recruited",
          ylabel="Channel Capacity (bits)", yticks=[0, 1, 2, 3, 4],
          xticks = [0, 0.5, 1],
          legend=:right,
          legendfont = font(10),
          size=(600, 300))

savefig(p1, "fig_s3a.pdf")

###############################################################################
### FIG 5E
###############################################################################

t_recruits = 5
t_release = 1
sigmas = collect(0.00001:0.00001:0.25)
k_p_ccs = [get_population_cc(t_recruit, t_release, k_krab, sigma) for sigma in sigmas]
h_p_ccs = [get_population_cc(t_recruit, t_release, k_hdac, sigma) for sigma in sigmas]

p1 = plot(sigmas, k_p_ccs,
         label="KRAB",
         color=color_map["KRAB"],
         linewidth=2)
p1 = plot!(sigmas, h_p_ccs,
          label="HDAC4",
          color=color_map["HDAC4"],
          linewidth=2,
          xlabel="Output Noise Variance",
          # xscale=:log10,
          # legend=:outerright,
          ylabel="Channel Capacity (bits)",
          ylims=(0,8.5),
          legendfont = font(10),
          size=(600, 300))

println("KRAB: $(get_population_cc(t_recruit, t_release, k_krab, noise_level))")
println("HDAC4: $(get_population_cc(t_recruit, t_release, k_hdac, noise_level))")

p1 = vline!(p1, [noise_level], color=:black, linestyle=:dash, label=false)
savefig(p1, "fig_s3b.pdf")

###############################################################################
### FIG 5F
###############################################################################

function get_population_cc_input_noise(t_recruit, t_release, k, input_variance = 0.05, output_variance=0.05)
    target_poff = get_pct_off(t_recruit, t_release, k)

    # input b/w zero and 1 (power=1), and noise is input variance
    input_variance_scaled = target_poff * input_variance

    overall_cc = 0.5*log2(1 + target_poff/(input_variance_scaled + output_variance))

    return overall_cc
end

k_hdac = Dict('a' => 0.76,
              's' => 8.6,
              'i' => 0)

k_krab = Dict('a' => 0.31,
              's' => 11,
              'i' => 0.13)

t_recruit = 5
t_release = 5

sigmas = collect(0.00001:0.00001:0.25)

input_noise_sigmas = [0, 0.001, 0.01, 0.1, 1]
noise_colors = reverse([colorant"purple4", colorant"purple1", colorant"mediumorchid2", colorant"violet", colorant"plum1", colorant"thistle1"])
noise_colors = reverse(noise_colors)
noise_colors = get(ColorSchemes.viridis, reverse(range(0.0, stop=0.9, length=length(noise_colors))))

krab_cc_list = [[get_population_cc_input_noise(1, t_release, k_krab, input_noise_sigmas[i], sigma) for sigma in sigmas] for i in 1:length(input_noise_sigmas)]
hdac_cc_list = [[get_population_cc_input_noise(1, t_release, k_hdac, input_noise_sigmas[i], sigma) for sigma in sigmas] for i in 1:length(input_noise_sigmas)]

# t1_p_ccs = [get_population_cc_input_noise(1, t_release, k_krab, sigma) for sigma in sigmas]
# t2_p_ccs = [get_population_cc_input_noise(2, t_release, k_krab, sigma) for sigma in sigmas]
# t3_p_ccs = [get_population_cc_input_noise(3, t_release, k_krab, sigma) for sigma in sigmas]
# t4_p_ccs = [get_population_cc_input_noise(4, t_release, k_krab, sigma) for sigma in sigmas]
# t5_p_ccs = [get_population_cc_input_noise(5, t_release, k_krab, sigma) for sigma in sigmas]
#
# cc_list = [t1_p_ccs, t2_p_ccs]
p1 = plot(xlabel="Output Noise Variance",
          ylabel="Channel Capacity (bits)",
          title="KRAB",
          legendtitle="Input Noise Variance",
          legend=false,
          legendfont = font(10),
          xscale=:log10,
          ylims=(-0.25,8.5),
          size=(600, 300))
for i in 1:length(krab_cc_list)
    p1 = plot!(sigmas, krab_cc_list[i], label="$(input_noise_sigmas[i])",
               color=noise_colors[i],
               linewidth=2)
end

p2 = plot(xlabel="Output Noise Variance",
          ylabel="Channel Capacity (bits)",
          title="HDAC4",
          legendtitle="Input Noise Variance",
          legendfont = font(10),
          xscale=:log10,
          legend=:outerright,
          ylims=(-0.25,8.5),
          size=(1200, 300))
for i in 1:length(hdac_cc_list)
    p2 = plot!(sigmas, hdac_cc_list[i], label="$(input_noise_sigmas[i])",
               color=noise_colors[i],
               linewidth=2)
end

p1 = vline!(p1, [noise_level], color=:black, linestyle=:dash, label=false)
p2 = vline!(p2, [noise_level], color=:black, linestyle=:dash, label=false)

lay = @layout[a  b]
p = plot(p1, p2, layout=lay)
savefig(p, "fig_s3c.pdf")


# Hutchinson Trace Estimator for Inverse of Matrix
# Based on Hutchinson's Stochastic Estimation Proof
# Values from Extrapolation Methods for Estimating Trace of the Matrix Inverse by P. Fika
# (This method works for SPD with low condition number)

export hutch, hutch!, HutchWorkspace

using LinearAlgebra

struct HutchWorkspace{T, TM <: AbstractMatrix{T}, F <: Function, TV <: AbstractVector{T}}
    A::TM
    randfunc::F
    x::TV
    y::TV
    N::Int64
    skipverify::Bool
end

"""
    HutchWorkspace(A; N = 30, skipverify = false)
    HutchWorkspace(A, randfunc::Function; N = 30, skipverify = false)

# Arguments
- `A` : Symmetric Positive Definite Matrix with Low Condtion Number (k < 500)
- `randfunc` : Function to generate random values for x distributed uniformly
                (Base: rand(-1.0:2.0:1.0, size(A,1))
                 Example: f(n) = rand(-1.0:2.0:1.0, n)
- `N` : Number of iterations (Default: 30)
- `skipverify` : If false, it will check isposdef(A) (Default: false)
"""
function HutchWorkspace(A, randfunc::Function; N = 30, skipverify = false)
    x = randfunc(size(A,1))
    y = similar(x)
    return HutchWorkspace(A, randfunc, x, y, N, skipverify)
end

function HutchWorkspace(A; N = 30, skipverify = false)
    randfunc(n) = rand(-1:2:1, n)
    x = rand(eltype(A) <: Integer ? Float64 : eltype(A), size(A,1))
    y = similar(x)
    return HutchWorkspace(A, randfunc, x, y, N, skipverify)
end

# Calculating the moments and extrapolating c-1
function ev(w::HutchWorkspace)
    # Create new random values for x
    copyto!(w.x, w.randfunc(size(w.A, 1)))

    # Find c(r) = x' * A^r * x = dot(x', A^r * x)
    c0 = dot(w.x, w.x)
    mul!(w.y, w.A, w.x)
    c1 = dot(w.y, w.x)
    c2 = dot(w.y, w.y)
    mul!(w.x, w.A, w.y)
    c3 = dot(w.y, w.x)

    # intermediate arguments for v0
    nVal = (c1^2)/(c0 * c2)
    dVal = (c1 * c3)/(c2^2)

    if (nVal > 0) && (dVal > 0)
        # Find p (Page 175)
        p = (c0 * c2)/(c1^2)
        # Find v0 (Page 179, Numerical Example) -> Corollary 4 in [16]
        v0 = log10(nVal) / log10(dVal)
        # Finally find ev (Page 175)
        ev0 = (c0^2)/(p^v0 * c1)
        return ev0
    else
        @warn "v0 cannot be calculated. Aitken's process may give inaccurate results!"
        w1 = HutchWorkspace(w.A, w.randfunc, w.x, w.y, w.N, w.skipverify)
        return hutch!(w1, aitken=true)
    end
end

# Aitken's Process to Predict the negative moment
# (Page 176, eq 4)
function gfun(w::HutchWorkspace)
    copyto!(w.x, w.randfunc(size(w.A, 1)))

    c0 = dot(w.x, w.x)
    mul!(w.y, w.A, w.x)
    c1 = dot(w.y, w.x)
    c2 = dot(w.y, w.y)
    mul!(w.x, w.A, w.y)
    c3 = dot(w.y, w.x)

    #g1 = c0 - (((c1 - c0)^2) / (c2 - c1))
    g2 = c1 - (((c2 - c0) * (c2 - c1)) / (c3 - c2))
    #g = (g1 + g2)/2
end

"""
    hutch(A; N = 30, skipverify = false)

Take a HutchWorkspace object as input, apply hutchinson estimation algorithm and solve for trace
of inverse of the matrix. (in-place version also available)
"""
# Hutchinson trace estimation using extrapolation technique (Page 177)
function hutch(A; N=30, skipverify = false, aitken = false)
    w = HutchWorkspace(A, N = N, skipverify = skipverify)
    return hutch!(w, aitken)
end

"""
    hutch!(w::HutchWorkspace; aitken = false)
    hutch!(A::AbstractArray; N = 30, skipverify = false, aitken = false)

Take a HutchWorkspace object as input, apply hutchinson estimation algorithm and solve for trace
of inverse of the matrix. (in-place version of hutch)
"""
function hutch!(w::HutchWorkspace; aitken = false)
    if w.skipverify == true || isposdef(w.A)
        tr = zero(eltype(w.A))
        sum = zero(eltype(w.A))
        for i in 1:w.N
            # if aitken == true => use aitken process to predict the terms via gfun
            sum = sum + (aitken ? gfun(w) : ev(w))
        end
        tr = sum/w.N
    end
end

function hutch!(A::AbstractArray{<:Any, 2}; N = 30, skipverify = false, aitken = false)
    w = HutchWorkspace(A, N, skipverify)
    return hutch!(w, aitken)
end


"""
    get_mesh_surfacenormals(points, elements, facets)

Computing outwards oriented surface normals on a connected mesh compartment.
"""
function get_mesh_surfacenormals(points, elements, facets)

    _, areas, facet_centers, normals = get_mesh_surface(points, facets)

    # Sizes
    nfacet = size(facets, 2)
    nelement = size(elements, 2)

    # Find elements to which the facets belong
    combs = [[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]
    ef = [SVector{3}(sort(elements[comb, i])) for comb ∈ combs for i = 1:nelement]
    sort!(ef)
    ef_labels = repeat(1:nelement, 4)

    f = SVector{3}.(sort.(eachcol(facets)))
    matchinds = fill(0, length(f))
    for i ∈ eachindex(f)
        j = findfirst(isequal(f[i]), ef)
        matchinds[i] = ef_labels[j]
    end

    elements_match = elements[:, matchinds]

    # Identify centers of matching elements
    x = points[:, elements_match]
    element_centers = reshape(mean(x, dims = 2), 3, nfacet)

    # Create outward directed vectors (from cell center to outer facet)
    outwards_vectors = facet_centers - element_centers

    # Determine orientation of normals (+ for out, - for in)
    orientations = [sign(outwards_vectors[:, i]' * normals[:, i]) for i = 1:nfacet]

    # Orient all normals outwards
    normals = orientations' .* normals

    areas, facet_centers, normals
end


# using Test
# using StaticArrays
# using TrajectoryOptimization
# using RobotZoo: Cartpole
# using ForwardDiff
# using BenchmarkTools
# using RobotDynamics
# using DiffResults
# using FiniteDiff
# using LinearAlgebra
import TrajectoryOptimization: stage_cost, CostFunction, Expansion
# const TO = TrajectoryOptimization

##
struct CartpoleCost{T} <: TO.CostFunction
    Q::Vector{T}
    R::Vector{T}
end
function TO.stage_cost(cost::CartpoleCost, x, u)
    return TO.stage_cost(cost, x) + 0.5 * cost.R[1] * u[1]^2
end
function TO.stage_cost(cost::CartpoleCost, x)
    y = x[1]
    θ = x[2]
    ydot = x[3]
    θdot = x[4]
    J = cost.Q[2] * cos(θ/2)
    J += 0.5* (cost.Q[1] * y^2 + cost.Q[3] * ydot^2 + cost.Q[4] * θdot^2)
    return J
end
TO.state_dim(::CartpoleCost) = 4
TO.control_dim(::CartpoleCost) = 1

## Initialize Model
@testset "Nonlinear Costs" begin
    model = Cartpole()
    n,m = size(model)
    x = zeros(n)
    u = zeros(m)
    cst = CartpoleCost([1,2,3,4.], [2.])
    @test TO.stage_cost(cst, x, u) ≈ 2.0
    x[2] = π
    @test TO.stage_cost(cst, x) ≈ 0 atol = 10*eps()
    x,u = rand(model)
    z = KnotPoint(x,u,0.1)
    zterm = KnotPoint(x,u,0.0)

    ## Test ForwardDiff
    TO.diffmethod(::CartpoleCost) = RobotDynamics.ForwardAD()
    E = Expansion{Float64}(n,m)
    TO.gradient!(E, cst, zterm)
    TO.hessian!(E, cst, zterm)
    TO.gradient!(E, cst, z)
    TO.hessian!(E, cst, z)
    @test (@allocated TO.gradient!(E, cst, zterm)) == 0
    @test (@allocated TO.hessian!(E, cst, zterm)) == 0
    @test (@allocated TO.gradient!(E, cst, z)) == 0
    @test (@allocated TO.hessian!(E, cst, z)) == 0
    @test E.xx ≈ Diagonal([1, -0.5*cos(x[2]/2), 3, 4])
    @test E.uu ≈ Diagonal([2])
    @test E.x ≈ [1*x[1], -sin(x[2]/2), 3*x[3], 4*x[4]]
    @test E.u ≈ 2*u

    TO.diffmethod(::CartpoleCost) = RobotDynamics.FiniteDifference()
    E2 = Expansion{Float64}(n,m)
    TO.gradient!(E2, cst, zterm)
    TO.hessian!(E2, cst, zterm)
    TO.gradient!(E2, cst, z)
    TO.hessian!(E2, cst, z)

    cache = TO.ExpansionCache(cst)
    TO.gradient!(E2, cst, zterm, cache)
    TO.hessian!(E2, cst, zterm, cache)
    TO.gradient!(E2, cst, z, cache)
    TO.hessian!(E2, cst, z, cache)

    @test (@allocated TO.gradient!(E2, cst, zterm, cache)) == 0
    @test (@allocated TO.hessian!(E2, cst, zterm, cache)) == 0
    @test (@allocated TO.gradient!(E2, cst, z, cache)) == 0
    @test (@allocated TO.hessian!(E2, cst, z, cache)) == 0

    @test norm(E2.hess- E.hess) < 1e-6
    @test norm(E2.grad- E.grad) < 1e-6

    ## Test expansion of trajectory
    N = 11
    obj = Objective(cst, N)
    X = [rand(n) for k = 1:N]
    U = [rand(m) for k = 1:N-1]
    Z = Traj(X,U,fill(0.1,N))

    TO.diffmethod(::CartpoleCost) = RobotDynamics.ForwardAD()
    E0 = [Expansion{Float64}(n,m) for k = 1:N]
    cache = TO.ExpansionCache(obj)
    TO.cost_expansion!(E0, obj, Z, cache, init=true, rezero=true)
    # @test (@allocated TO.cost_expansion!(E0, obj, Z, cache, init=true, rezero=true)) == 0

    TO.diffmethod(::CartpoleCost) = RobotDynamics.FiniteDifference()
    E1 = [Expansion{Float64}(n,m) for k = 1:N]
    cache = TO.ExpansionCache(obj)
    TO.cost_expansion!(E1, obj, Z, cache, init=true, rezero=true)
    # @test (@allocated TO.cost_expansion!(E1, obj, Z, cache, init=true, rezero=true)) == 0
end


using Plots
using Random
using DifferentialEquations
using DynamicalSystems

function make_data(system::DynamicalSystem;
                   train_time = 100, nspin=500, test_time = 15, n_test = 10, Δt=0.01)
    """make training and testing data from the dynamical system"""
    options = (alg = Vern9(), abstol = 1e-12,reltol=1e-12)
    D = length(system.u0)
    train_data = trajectory(system, train_time, system.u0.*rand(D).+rand(D);
                            Δt=Δt,
                            Ttr=nspin, diffeq=options)
    test_data = Matrix{Float64}[]
    for i in 1:n_test
        data = trajectory(system, test_time, system.u0.*rand(D).+rand(D);
                          Δt=Δt,
                          Ttr=nspin, diffeq=options)
        push!(test_data, collect(Matrix(data)'))
    end
    return collect(Matrix(train_data)'), test_data
end

function plot_prediction(time, upred, utrue)
    """Plotting function"""
    p = []
    D = size(upred)[1]
    for i in 1:D
        if i == 1
            px = plot(time, upred[1, :], label="RC")
            plot!(px, time, utrue[1, :], label="Truth")
        else
            px = plot(time, upred[i, :])
            plot!(px, time, utrue[i, :], legend=false)
        end
        ylabel!(px, "X$i")
        if i!=D xticks!(px, Int[]) end
        push!(p, px)
    end

    xlabel!(p[end], "Time")
    return plot(p..., layout=(D, 1), link=:x)
end


    # Test for GaussLobatto()
    tol = 1e-14
    n = 42
    x = gausslobatto(n)
    w = x[2]; x = x[1]
    @test ( length(x) == n) && ( length(w) == n )
    @test abs(x[37] - 0.922259214258616) < tol
    @test abs(w[37] - 0.029306411216166) < tol
    @test ( x[1] == -1 && x[n] == 1 )


abstract type AbstractMachine end

function filtrations end

function (m::AbstractMachine)(y₀)
    forward, backward = filtrations(m, y₀)
    return solve(y₀, m.W, m.σ, (forward, backward))
end

sum_dims(dims::Tuple) = prod(dims[1:end-2]) * sum(dims[end-1:end])

# glorot initialization, from https://github.com/FluxML/Flux.jl
glorot_normal(dims...) = randn(Float32, dims...) .* sqrt(2.0f0 / sum_dims(dims))
glorot_uniform(dims...) = (rand(Float32, dims...) .- 0.5f0) .* sqrt(24.0f0 / sum_dims(dims))

function consecutive_ranges(dims::Vector{Int})
    return map(cumsum(dims), dims) do cumdim, dim
        return (cumdim - dim + 1):cumdim
    end
end

function dense_filtrations(y, dims::Vector{Int})
    channels, minibatch = sum(dims), size(y, 2)
    shapes = map(rg -> (rg, 1:minibatch), consecutive_ranges(dims))
    input(_) = (1:channels, 1:minibatch)
    function project(input_shape, output_shape)
        return (first(input_shape), first(output_shape))
    end
    return (
        forward=Filtration(shapes, input, project),
        backward=Filtration(reverse(shapes), input, project)
    )
end

"""
    DenseMachine(W, σ, dims::Vector{Int})

Create a `DenseMachine` object from a square weight matrix `W`, pointwise nonlinearity `σ`
and filtration sequence `dims`.
The values of `dims` specify how to split the input space into a sequence of subspaces.
In particular, it is required that `size(W) == (sum(dims), sum(dims))`.
"""
struct DenseMachine{T, F} <: AbstractMachine
    W::T
    σ::F
    dims::Vector{Int}
end

@functor DenseMachine (W, σ)

DenseMachine(σ, dims::Vector{Int}; init=glorot_uniform) = DenseMachine(dims, σ; init)

"""
    DenseMachine(dims::Vector{Int}, σ; init=glorot_uniform)

Return a `DenseMachine(W, σ, dims)` object, where `W = init(sum(dims), sum(dims))`.
Default to Glorot uniform initialization.
"""
function DenseMachine(dims::Vector{Int}, σ; init=glorot_uniform)
    W = init(sum(dims), sum(dims))
    return DenseMachine(W, σ, dims)
end

filtrations(m::DenseMachine, y₀) = dense_filtrations(y₀, m.dims)

nest_pads(::Tuple{}) = ()

function nest_pads(pad::Tuple)
    a, b, tail... = pad
    return ((a, b), nest_pads(tail)...)
end

function conv_range(output_range, (pad0, pad1)::Tuple{Int,Int})
    i0, i1 = first(output_range), last(output_range)
    return (i0 - pad0):(i1 + pad1)
end

function conv_filtrations(y::AbstractArray{T, N}, dims::Vector{Int}; pad::NTuple{M, Int}) where {T, N, M}
    M == 2 * (N - 2) || throw(ArgumentError("padding must have twice the length of image dimensions"))
    pads = nest_pads(pad)

    ranges = consecutive_ranges(dims)
    channels, minibatch = sum(dims), size(y, N)
    shapes = [ntuple(n -> ifelse(n == N - 1, rg, UnitRange(axes(y, n))), N) for rg in ranges]

    function input(shape)
        data = ntuple(n -> shape[n], N - 2)
        zdata = map(conv_range, data, pads)
        return (zdata..., 1:channels, 1:minibatch)
    end

    function input′(shape)
        data = ntuple(n -> shape[n], N - 2)
        zdata = map(conv_range, data, map(reverse, pads))
        return (zdata..., 1:channels, 1:minibatch)
    end

    function project(input_shape, output_shape)
        return ntuple(N) do n
            n == N && return output_shape[end - 1]
            n == N - 1 && return input_shape[end - 1]
            p0, p1 = pads[n]
            return 1:(p0 + p1 + 1)
        end
    end

    return (
        forward=Filtration(shapes, input, project),
        backward=Filtration(reverse(shapes), input′, project)
    )
end

"""
    ConvMachine(W, σ, dims::Vector{Int}, pad::Dims)

Create a `ConvMachine` object from a weight array `W`, pointwise nonlinearity `σ`,
filtration sequence `dims` and padding `pad`.
The values of `dims` specify how to split the input space into a sequence of subspaces.
In particular, it is required that `size(W) == (kernelsize..., sum(dims), sum(dims))`,
where `kernelsize` is such that convolution by `W` with padding `pad` preserves
input dimension.
Padding `pad` has length twice the number of kernel dimensions (for example, it takes
four values for image convolutions).
"""
struct ConvMachine{T, F, M} <: AbstractMachine
    W::T
    σ::F
    dims::Vector{Int}
    pad::NTuple{M, Int}
end

@functor ConvMachine (W, σ)

ConvMachine(W, σ, dims; pad) = ConvMachine(W, σ, dims, pad)

ConvMachine(σ, dims::Vector{Int}; pad, init=glorot_uniform) = ConvMachine(dims, σ; pad, init)

"""
    ConvMachine(dims::Vector{Int}, σ; pad, init=glorot_uniform)

Return a `ConvMachine(W, σ, dims, pad)` object, where `W = init(kernelsize..., sum(dims), sum(dims))`.
Default to Glorot uniform initialization.
Here `kernelsize` is such that convolution with a kernel of size `kernelsize`
and padding `pad` preserves input dimension.
Padding `pad` has length twice the number of kernel dimensions (for example, it takes
four values for image convolutions).
"""
function ConvMachine(dims::Vector{Int}, σ; pad, init=glorot_uniform)
    kernelsize = map(((p0, p1),) -> p0 + p1 + 1, nest_pads(pad))
    W = init(kernelsize..., sum(dims), sum(dims))
    return ConvMachine(W, σ, dims; pad)
end

filtrations(m::ConvMachine, y₀) = conv_filtrations(y₀, m.dims; m.pad)

function recur_filtrations(y, dims::Vector{Int}; pad, timeblock)
    ranges = consecutive_ranges(dims)
    datalength, channels, minibatch = size(y, 1), sum(dims), size(y, 3)
    N = ceil(Int, datalength / timeblock)
    shapes = [(timeblock  * (n  - 1) + 1 : timeblock * n, range, 1:minibatch) for n in 1:N for range in ranges]

    input((timerange, _, mb)) =  (conv_range(timerange, (pad, 0)), 1:channels, mb)
    input′((timerange, _, mb)) =  (conv_range(timerange, (0, pad)), 1:channels, mb)
    function project(input_shape, output_shape)
        input_range, output_range = first(input_shape), first(output_shape)
        w0 = 1 + clamp(first(output_range) - last(input_range), 0:pad)
        w1 = 1 + clamp(last(output_range) - first(input_range), 0:pad)
        return (w0:w1, input_shape[end-1], output_shape[end-1])
    end
    flip = (pad+1):-1:1
    function project′(input_shape, output_shape)
        input_range, output_range = first(input_shape), first(output_shape)
        w0 = 1 + clamp(first(input_range) - last(output_range), 0:pad)
        w1 = 1 + clamp(last(input_range) - first(output_range), 0:pad)
        return (flip[w1]:flip[w0], input_shape[end-1], output_shape[end-1])
    end
    return (
        forward=Filtration(shapes, input, project),
        backward=Filtration(reverse(shapes), input′, project′)
    )
end

"""
    RecurMachine(W, σ, dims::Vector{Int}, pad::Int, timeblock::Int)

Create a `RecurMachine` object from a weight array `W`, pointwise nonlinearity `σ`,
filtration sequence `dims`, padding `pad` and time block `timeblock`.
The values of `dims` specify how to split the input space into a sequence of subspaces.
In particular, it is required that `size(W) == (pad + 1, sum(dims), sum(dims))`.
"""
struct RecurMachine{T, F} <: AbstractMachine
    W::T
    σ::F
    dims::Vector{Int}
    pad::Int
    timeblock::Int
end

@functor RecurMachine (W, σ)

RecurMachine(W, σ, dims; pad, timeblock) = RecurMachine(W, σ, dims, pad, timeblock)

RecurMachine(σ, dims::Vector{Int}; pad, timeblock, init=glorot_uniform) = RecurMachine(dims, σ; pad, timeblock, init)

"""
    RecurMachine(dims::Vector{Int}, σ; pad, timeblock, init=glorot_uniform)

Return a `RecurMachine(W, σ, dims, pad, timeblock)` object, where `W = init(pad + 1, sum(dims), sum(dims))`.
Default to Glorot uniform initialization.
"""
function RecurMachine(dims::Vector{Int}, σ; pad, timeblock, init=glorot_uniform)
    W = init(pad + 1, sum(dims), sum(dims))
    return RecurMachine(W, σ, dims; pad, timeblock)
end

filtrations(m::RecurMachine, y₀) = recur_filtrations(y₀, m.dims; m.pad, m.timeblock)

function ChainRulesCore.rrule(::typeof(filtrations), m::AbstractMachine, y)
    res = filtrations(m, y)
    function filtrations_pullback(_, _)
        return NoTangent(), NoTangent(), NoTangent()
    end
    return res, filtrations_pullback
end


"""
    x = solve_pcg(A,b; maxiter=5000,eps=1e-15,reset=50,verbose=false)

Solve the symmetric positive definite system of equations:
`A` as sparse left-hand side and `b` as a vector of right-hand side.
The Jacobi preconditioner (`diag(A)`) is used.
The maximum number of iterations is specified with `maxiter`.
Convergence criterion `eps` is `norm(b-A*x)/norm(b)`.
The residual vector will be recalculated every `reset` rounds.
The details will be shown with `verbose=true`.

```juliadoctest
julia> x = solve_pcg(A,b)
```
"""
function solve_pcg(A, b::Vector{Tv};
                   maxiter=5000,eps=1e-15,reset=50,verbose=false) where {Tv}
   n = size(A,1)
   x = zeros(Tv,n)
   r = zeros(Tv,n)
   p = zeros(Tv,n)
   z = zeros(Tv,n)
   w = zeros(Tv,n)
   Mi = Diagonal(1 ./ diag(A))
   r .= b
   oldtau = 1.0
   for i=1:maxiter
      z .= Mi*r
      tau = z'*r
      if i==1
         beta = 0
         p .= z
      else
         beta = tau/oldtau
         p .= z + beta*p
      end
      w .= A*p
      alpha = tau/(p'*w)
      x .= x + alpha*p
      if i % reset == 0
         r .= b - A*x
      else
         r .= r - alpha*w
      end
      conv = norm(r)^2/norm(b)^2
      if verbose
         print(@sprintf("round %5d: conv %8.3e\n",i,conv))
      end
      if conv<eps; break; end
      oldtau = tau
   end
   x
end

"""
   tr = approximate_trace_of_inverse(A)
"""
function approximate_trace_of_inverse(A; rounds=1000, msglev=0)
   if rounds<1
      throw( ArgumentError("negative rounds") )
   end
   tr = Threads.Atomic{Float64}(0.0)
   done = Threads.Atomic{Int}(0)
   n = size(A,1)
   @threads for iter=1:rounds
      atomic_add!(done, 1)
      z = [ifelse(rand()>0.5,1.0,-1.0) for i in 1:n]
      #tr = tr + z'*y
      if msglev>0
         @time y = solve_pcg(A,z)
         println("iter $(done.value) thread $(Threads.threadid())")
      else
         y = solve_pcg(A,z)
      end
      atomic_add!(tr, z'*y)
   end
   return tr.value/rounds
end


@testset verbose = true "statistics" begin
    @testset "analytic_histogram" begin
        # check against uniform distribution
        d = Distributions.Uniform(0, 1)
        G = LinRange(0.0, 1.0, 10)
        res = analytic_histogram(d, G)

        @test res[1] ≈ step(G) / 2
        for i in 2:(length(G) - 1)
            @test res[i] ≈ step(G)
        end
        @test res[end] ≈ step(G) / 2
    end
end


module EPExperiments

using Reexport
@reexport using DifferentialEquations
using DiffEqSensitivity
@reexport using EPT
using StatsFuns: logsumexp, logistic
using DataFrames
@reexport using Parameters
using Logging
using LoggingExtras
using LaTeXStrings
import LinearAlgebra

export sir_ode!,
    base_reproduction_rate,
    predict,
    neg_bin_params2,
    bayes_sir,
    geomspace,
    SIRParameters,
    plot_predictive,
    process_results,
    COST_FUNCTIONS,
    plot_is_samples,
    MHTransitionConfig,
    MHFixedTuple,
    MHThreeLevel,
    MHFunction,
    AnISConfig,
    ISConfig,
    AnISHMCConfig,
    AnISHMCGranularConfig,
    AnISMHConfig,
    SIRExperimentConfig,
    convergence_plot,
    ISGroundTruthConfig,
    get_anis_alg,
    radon_hierarchical,
    preprocess_radon_df,
    POST_PRED_COUNTY_IDX,
    POST_PRED_FLOOR,
    RADON_COST_FNS,
    radon_target_f,
    plot_intermediate_ess,
    plot_intermediate_weights,
    plot_intermediate_acceptance,
    get_median_and_quantiles,
    check_diagnostics,
    ess_batch,
    make_ess_plot,
    relative_squared_error,
    plot_joint_dist,
    plot_sir_samples,
    post_pred,
    POST_PRED_Y_OBSERVED,
    POST_PRED_GROUND_TRUTH,
    POST_PRED_DIMENSION,
    compute_ess,
    banana,
    banana_f

# Constants for the loss function for radon problem.
const POST_PRED_COUNTY_IDX = 10
const POST_PRED_FLOOR = 0

include("configs.jl")
include("cost_functions.jl")
include("check_diagnostics.jl")
include("post_pred_model.jl")

include("radon_model.jl")

function sir_ode!(du,u,p,t)
    (S,I,R,C) = u
    (β,γ) = p
    N = S+I+R
    infection = β*I/N*S
    recovery = γ*I
    @inbounds begin
        du[1] = -infection
        du[2] = infection - recovery
        du[3] = recovery
        du[4] = infection
    end
    nothing
end

function base_reproduction_rate(β, γ)
    return β / γ
end

function predict(y,chain)
    # Length of data
    l = length(y)
    # Length of chain
    m = length(chain)
    # Choose random
    idx = sample(1:m)
    i₀ = Array(chain[:i₀])[idx]
    β = Array(chain[:β])[idx]
    I = i₀*10.0
    u0=[1000.0-I,I,0.0,0.0]
    p=[β,0.25]
    tspan = (0.0,float(l))
    prob = ODEProblem(sir_ode!,
            u0,
            tspan,
            p)
    sol = solve(prob,
                Tsit5(),
                saveat = 1.0)
    out = Array(sol)
    sol_X = [0.0; out[4,2:end] - out[4,1:(end-1)]]
    return hcat(sol.t, out', sol_X)
end

neg_bin_r(mean, var) = mean^2 / (var - mean)
neg_bin_p(r, mean) = r / (r + mean)

function neg_bin_params(mean, var)
    r = neg_bin_r(mean, var)
    return r, neg_bin_p(r, mean)
end

neg_bin_params2(mean, phi) = neg_bin_params(mean, mean + mean^2 / phi)

@expectation function bayes_sir(y, total_population, γ, cost_fn)
    # Calculate number of timepoints
    l = length(y)
    i₀ ~ truncated(Normal(10, 10), 0, total_population/10)
    β ~ truncated(Normal(2, 1.5), 0, Inf)

    I = i₀ * 10
    u0 = [total_population-I, I, 0.0, 0.0]
    p = [β, γ]

    tspan = (0.0, float(l))
    prob = ODEProblem(sir_ode!, u0, tspan, p)
    sol = solve(prob, Tsit5(), saveat = 1.0)
    sol_C = Array(sol)[4,:] # Cumulative cases
    sol_X = sol_C[2:end] - sol_C[1:(end-1)]

    l = length(y)
    if any(sol_X .< 0) || length(sol_X) != l
        # Check if we have negative cumulative cases
        @warn "Negative cumulative cases" β i₀
        Turing.acclogp!(_varinfo, -Inf)
        return l
    end

    phi = 0.5
    variance = sol_X .+ sol_X.^2 ./ phi
    rs = neg_bin_r.(sol_X, variance)
    ps = neg_bin_p.(rs, sol_X)
    if !all(rs .> 0) || !all(0 .< ps .<= 1)
        Turing.acclogp!(_varinfo, -Inf)
        return l
    end

    y ~ arraydist(NegativeBinomial.(rs, ps))
    #y ~ arraydist([Poisson(x) for x in sol_X])
    return cost_fn(β, γ)
end

function geomspace(start, stop, length)
    logstart = log10(start)
    logstop = log10(stop)
    points = 10 .^ range(logstart, logstop; length=length)
    points[1] = start
    points[end] = stop
    return points
end

function relative_squared_error(mu_true, mu_hat)
    return (mu_true - mu_hat)^2 / mu_true^2
end

function compute_ess(log_weights)
    denominator = logsumexp(2 * log_weights)
    numerator = 2 * logsumexp(log_weights)
    return exp(numerator - denominator)
end

include("plots.jl")

function est_exp(chn, f)
    log_weights = Array(chn[:log_weight])
    i₀ = Array(chn[:i₀])
    β = Array(chn[:β])

    normalisation = exp(logsumexp(log_weights))
    weighted_sum = sum(zip(log_weights, i₀, β)) do (lw, i, b)
        exp(lw) * f(i, b)
    end

    return weighted_sum / normalisation
end

function process_results(
    result_folder,
    logger;
    ode_tabi,
    R0_estimate,
    tabi,
    X,
    Y,
    obstimes,
    i_0_true,
    β_true,
    cost_fn
)
    betas = tabi.Z2_alg.inference_algorithm.betas

    savefig(plot_intermediate_weights(
        ode_tabi[:Z2_info].info[:intermediate_log_weights], betas),
        joinpath(result_folder, "intermediate_weights.png")
    )
    savefig(plot_intermediate_weights(
        ode_tabi[:Z2_info].info[:intermediate_log_weights], betas; ixs=1:14),
        joinpath(result_folder, "intermediate_weights_first40.png")
    )
    savefig(plot_intermediate_ess(
        ode_tabi[:Z2_info].info[:intermediate_log_weights], betas),
        joinpath(result_folder, "intermediate_ess.png")
    )
    savefig(plot_intermediate_ess(
        ode_tabi[:Z2_info].info[:intermediate_log_weights], betas; ixs=1:14),
        joinpath(result_folder, "intermediate_ess_first40.png")
    )
    savefig(
        plot_predictive(obstimes, X, Y, ode_tabi[:Z2_info]),
        joinpath(result_folder, "anis_post_pred.png")
    )
    savefig(
        plot_joint_dist(ode_tabi[:Z2_info], β_true, i_0_true),
        joinpath(result_folder, "joint_samples.png")
    )

    anis_estimate = est_exp(ode_tabi[:Z2_info], cost_fn)

    with_logger(logger) do
        @info "Expectation estimate TABI: $(R0_estimate)"
        @info "Expectation estimate AnIS: $(anis_estimate)"

        @info "Estimation algorithm: $(tabi)"

        for (k, est_results) in pairs(ode_tabi)
            if isnothing(est_results)
                continue
            end
            msg = string([
                "$(string(k)):\n",
                "ESS: $(est_results.info[:ess])\n",
                "Log evidence: $(est_results.logevidence)\n",
                "Log weights: $(get(est_results, :log_weight)[:log_weight])\n"
            ]...)

            @info "$(msg)"
        end
    end
end

end


const Maybe{T} = Union{T, Nothing} where T

const Symbolic = Union{Expr, Symbol}

struct Bounds
    a :: Number
    b :: Number
end

struct Integral
    expr   :: Union{Symbolic, Number}
    diff   :: Symbolic
    bounds :: Maybe{Bounds}
end

createIntegral(expr, diff::Symbol, bounds = nothing) = Integral(expr, diff, bounds)
function createIntegral(expr, diff::Expr, bounds = nothing)
    # We need to validate that the differential contains ONE AND ONLY ONE variable
    vars = 0
    f(x) = nothing
    f(x::Symbol) = (vars += 1)
    walk(f, diff)

    if vars ≠ 1
        error("the differential must contain strictly one variable.")
    end

    Integral(expr, diff, bounds)
end

struct Thunk
    expr :: Expr
end

createThunk(x) = x

# Eagerly evaluate eval() calls
function createThunk(ex::Expr)
    new_ex = Context.eagerCalc(ex)
    root_eager = ex.head == :call && ex.args[1] ∈ [:eval, :eval!]
    root_eager ? new_ex : Thunk(new_ex)
end

const Value = Union{Symbolic, Number, Integral, Thunk}

struct Assignment
    lvalue :: Union{Symbolic, Integral}
    rvalue :: Value
end

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

# string conversion overload
using Base: string
Base.string(t::Thunk) = Base.string(t.expr)

Base.string(bounds::Bounds) = "_$(wrap(bounds.a))^$(wrap(bounds.b))"

function Base.string(int::Integral)
    bounds = int.bounds ≠ nothing ? Base.string(int.bounds) : ""
    "int$(bounds) $(wrap(int.expr))*d$(wrap(int.diff))"
end


@testset "Levenberg-Marquardt: NIST Dataset (Eckerle4)" begin
    print_head("Levenberg-Marquardt: NIST Dataset (Eckerle4)")

    obj = optimizer(nd=3, ny=35, method="lm")
    fun = Eckerle4()
    sol = Vector{Float64}(undef, 3)

    parμ = [1.0, 10.0, 460.0]
    βvec = zeros(Float64, 35)
    βmat = zeros(Float64, 35, 35)

    print_body("\033[1m\033[34mLikelihood Precision Matrix: Identity\033[0m")

    @simd for i in eachindex(βvec)
        @inbounds βvec[i] = 1.0
    end

    @simd for i in eachindex(βvec)
        @inbounds βmat[i,i] = 1.0
    end

    ans = [1.55438286, 4.08883291, 451.541218]

    print_body(string("Precision Const.: ", optimize!(sol, fun, parμ, 1.0, obj)))
    @inbounds for i in eachindex(ans)
        @test sol[i] ≈ ans[i] rtol=1e-7
    end

    print_body(string("Precision Vector: ", optimize!(sol, fun, parμ, βvec, obj)))
    @inbounds for i in eachindex(ans)
        @test sol[i] ≈ ans[i] rtol=1e-7
    end

    print_body(string("Precision Matrix: ", optimize!(sol, fun, parμ, βmat, obj)))
    @inbounds for i in eachindex(ans)
        @test sol[i] ≈ ans[i] rtol=1e-7
    end

    print_body("\033[1m\033[34mLikelihood Precision Matrix: Custom\033[0m")

    @simd for i in eachindex(βvec)
        @inbounds βvec[i] = 1.0 + 0.1 * (i-1)
    end

    @simd for i in eachindex(βvec)
        @inbounds βmat[i,i] = 1.0 + 0.1 * (i-1)
    end

    @inbounds ans[1], ans[2], ans[3] = 1.55104316, 4.07147586, 451.543588

    print_body(string("Precision Vector: ", optimize!(sol, fun, parμ, βvec, obj)))
    @inbounds for i in eachindex(ans)
        @test sol[i] ≈ ans[i] rtol=1e-7
    end

    print_body(string("Precision Matrix: ", optimize!(sol, fun, parμ, βmat, obj)))
    @inbounds for i in eachindex(ans)
        @test sol[i] ≈ ans[i] rtol=1e-7
    end
end