#==============================================================================

          Solving the HJB and KFE
            And then use the KFE to find the equilibrium interest rate

               Translated from matlab code on Ben Moll's website:
                              http://www.princeton.edu/~moll/HACTproject.htm

==============================================================================#

using LinearAlgebra, SparseArrays, Plots

γ = 2 # parameter from CRRA utility
ρ = 0.05 # the discount rate

z1=.1
z2 =.2
z= [z1 z2]

λ1 = 1.2
λ2 = 1.2

λ = [λ1 λ2]

r0=0.03 #inital guess
rmin= 0.01
rmax=0.04

H = 500 # number of points in the grid space
amin= -0.15
amax= 5

a = LinRange(amin,amax,H)
da = (amax-amin)/(H-1)

aa = [a a]
zz = ones(H,1)*z

maxit = 100 # the maximum number of iterations we allow the finite differencing algorithm
crit = 10^(-6) # our critical value
Δ = 40

dVf=zeros(H,2)
dVb = zeros(H,2)
c= zeros(H,2)

# Create the matrix that captures the systems dynamic from λ process
Aswitch = [-sparse(I,H,H)*λ[1] sparse(I,H,H)*λ[1]; sparse(I,H,H)*λ[2] -sparse(I,H,H)*λ[2]]

# Set up parameters for outer loop
Ir = 40 # number of iterations, values to check
crit_S=10^(-50)
r_grid =LinRange(rmin,rmax,Ir)

# Inital Guesses
r = r0

v0=zeros(H,2)
v0[:,1] = (z[1] .+ max(r,0.01).*a).^(1-γ)/(1-γ)/ρ
v0[:,2] = (z[2] .+ max(r,0.01).*a).^(1-γ)/(1-γ)/ρ

global v0=v0

dist=[]
V_n=[]
g_r=[]
dV_r=[]
V_r=[]
c_r=[]
adot=[]
r_r=[]
rmin_r =[]
rmax_r=[]

S =zeros(Ir,1)

# Loop that finds the equilibrium interest rate
for ir in 1:Ir
    push!(r_r,r)
    push!(rmin_r,rmin)
    push!(rmax_r,rmax)

    if ir>1
        global v0 = V_r[ir-1]
    end

    global v = v0

    # The finite differeing loop
    for n in 1:maxit
        global V=v
        push!(V_n,V)
        # Forward differencing
        dVf[1:H-1,:]=(V[2:H,:]-V[1:H-1,:])/da
        dVf[H,:] = (z .+ r.*amax).^(-γ) # impose state constraint at the max, just in case
        # Backward differencing
        dVb[2:H,:]=(V[2:H,:]-V[1:H-1,:])/da
        dVb[1,:]= (z .+r.*amin).^(-γ)

        # Find consumption and savings with the forward difference
        cf = dVf.^(-1/γ)
        sf = zz +r.*aa -cf

        # Find consumption and savings with the backward difference
        cb= dVb.^(-1/γ)
        sb = zz +r.*aa -cb

        # Find consumption and savings at steady state
        c0 = zz + r.*aa
        dV0 = c0.^(-γ)

        # Now implement the upwind scheme in order to select the best differencing method
        If = sf .> 0
        Ib = sb .< 0
        I0 = (1 .- If - Ib)

        # State constraint at amin is automatically implemented

        global dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0
        global c = dV_Upwind.^(-1/γ)

        u = c.^(1-γ)/(1-γ)
        # Construct matrix for the evolution of the system
        X = -min.(sb,0)/da
        Y= -max.(sf,0)/da + min.(sb,0)/da
        Z = max.(sf,0)/da

        A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
        A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
        AA = [A1 spzeros(H,H); spzeros(H,H) A2]

        global A = AA + Aswitch

        B = (ρ +1/Δ)*sparse(I,2*H,2*H) -A
        u_stacked = [u[:,1]; u[:,2]]
        V_stacked = [V[:,1]; V[:,2]]

        b = u_stacked + V_stacked/Δ
        V_stacked = B\b # Solves the system of equations

        V = [V_stacked[1:H] V_stacked[H+1:2*H]]

        V_change = V - v

        global v = V

        push!(dist, findmax(abs.(V_change))[1])

        if dist[n] < crit
            println("Value Function converged, Iteration=")
            println(n)
            break
        end
    end

    #====================================================
            Fokker-Planck Equation
    ====================================================#

    AT = A' # we need the transpose of our A matrix
    b =zeros(2*H,1)

    # We need to fix one value to prevent the matrix from being singular
    i_fix =1
    b[i_fix] = .1
    row = [zeros(1,i_fix-1) 1 zeros(1,2*H-i_fix)]
    AT[i_fix,:] = row

    # Using this solve the linear system
    gg = AT\b
    g_sum = gg'*ones(2*H,1)*da
    gg = gg./g_sum # we need gg to sum to one

    g = [gg[1:H] gg[H+1:2*H]]

    # Graphs
    push!(g_r, g)
    push!(adot, zz+r.*aa-c)
    push!(V_r, V)
    push!(c_r, c)
    push!(dV_r, dV_Upwind)

    S[ir] = g[:,1]'*a*da + g[:,2]'*a*da

    if S[ir]>crit_S
        println("Excess Supply")
        global rmax =r
        global r = 0.5*(r+rmin)
    elseif S[ir]<-crit_S
        println("Excess Demand")
        global rmin=r
        global r = 0.5*(r+rmax)
    elseif abs(S[ir])<crit_S
        println("Equilibrium Found")
        println(r)
        break
    end

#end of outer loop
end

# Graphs
amax1 = 0.6
amin1 = amin-0.03

plot(a,adot[end][:,1], label="\$ s_1(a)\$", xlabel="Wealth", ylabel="\$s_i(a)\$", xlims=(amin1,amax1))
plot!(a,adot[end][:,2], label="\$ s_2(a)\$")
plot!(LinRange(amin1,amax1,H),zeros(H,1), label="", line=:dash, color=:black)

plot(a,g_r[end][:,1], label="\$ g_1(a)\$", xlabel="Wealth", ylabel="\$g_i(a)\$", xlims=(amin1,amax1), ylims=(0,3))
plot!(a,g_r[end][:,2], label="\$ g_2(a)\$")


export expected_n,
       expected_pq

"""
    expected_n(basis::Vector{Int},
               phi::Vector{T},
               psi::Vector{T},
               p::Vector{Int}) where {T<:Number}

Compute the expectation value <phi|n_p[1]*n_p[2]*...*n_p[n]|psi>.
"""
function expected_n(basis::Vector{Int},
                    phi::Vector{T},
                    psi::Vector{T},
                    p::Vector{Int}) where {T<:Number}
    o = zero(Float64)
    for i in eachindex(basis)
        bs = basis[i]
        has_all_n = true
        for pi in p
            (bs>>(pi-1))&1 == 0 && (has_all_n = false)
        end
        if has_all_n
            o += conj(phi[i])*psi[i]
        end
    end
    return o
end

"""
    expected_n(basis::Vector{Int},
               psi::Vector{T},
               p::Vector{Int}) where {T<:Number}

Compute the expectation value <psi|n_p[1]*n_p[2]*...*n_p[n]|psi>.
"""
function expected_n(basis::Vector{Int},
                    psi::Vector{T},
                    p::Vector{Int}) where {T<:Number}
    return expected_n(basis, psi, psi, p)
end

expected_n(basis, psi, p::Int) = expected_n(basis, psi, Int[p])
expected_n(basis, phi, psi, p::Int) = expected_n(basis, phi, psi, Int[p])

"""
    expected_pq(basis::Vector{Int},
                phi::Vector{T},
                psi::Vector{T},
                p::Vector{Int},
                q::Vector{Int}) where {T<:Number}

Compute <phi|b^+_p[1]*...*b^+_p[n]*b_q[1]*...*b_q[m]|psi>.
"""
function expected_pq(basis::Vector{Int},
                     phi::Vector{T},
                     psi::Vector{T},
                     p::Vector{Int},
                     q::Vector{Int}) where {T<:Number}

    p = deepcopy(p)
    q = deepcopy(q)
    (allunique(p) && allunique(q)) || return zero(T)
    # Number operators.
    n_i = p ∩ q
    # Remove those indices where there are number ops from p and q.
    setdiff!(p, n_i)
    setdiff!(q, n_i)

    o = zero(T)
    for i in eachindex(basis)
        bs = basis[i]
        has_all_pq = true
        # See if b^+_p*b_q acting on bs is not zero.
        for pi in p
            (bs>>(pi-1))&1 != 0 && (has_all_pq = false)
        end
        for qi in q
            (bs>>(qi-1))&1 != 1 && (has_all_pq = false)
        end
        for ni in n_i
            (bs>>(ni-1))&1 != 1 && (has_all_pq = false)
        end
        # If b^+_q*b_q|bs> not zero, compute the resulting state.
        if has_all_pq
            rs = bs
            for pi in p
                rs += 1<<(pi-1)
            end
            for qi in q
                rs -= 1<<(qi-1)
            end
            j = searchsortedfirst(basis, rs)
            o += psi[i]*conj(phi[j])
        end
    end
    return o
end

"""
    expected_pq(basis::Vector{Int},
                psi::Vector{<:Number},
                phi::Vector{<:Number},
                p::Int,
                q::Int)

Compute <phi|b^+_p*b_q|psi>.
"""
function expected_pq(basis::Vector{Int},
                     psi::Vector{T},
                     phi::Vector{T},
                     p::Int,
                     q::Int) where {T<:Number}

    p == q && return expected_n(basis, psi, phi, Int[p])

    o = zero(T)
    for i in eachindex(basis)
        bs = basis[i]
        if (bs>>(p-1))&1 == 0 && (bs>>(q-1))&1 == 1
            rs = bs + 1<<(p-1) - 1<<(q-1)
            j = searchsortedfirst(basis, rs)
            o += psi[i]*conj(phi[j])
        end
    end
    return o
end

expected_pq(basis, psi, p::Int, q::Int) = expected_pq(basis, psi, psi, p, q)
expected_pq(basis, psi, p::Vector{Int}, q::Vector{Int}) = expected_pq(basis, psi, psi, p, q)


using LinearAlgebra:
        Diagonal,
        copyto!,
        diag,
        I

struct CoupledResonators
    yr::AbstractMatrix
    yi::AbstractMatrix
    wr::AbstractMatrix
    wi::AbstractMatrix
    γin::AbstractMatrix
    γex::AbstractMatrix
end

function CoupledResonators()
    return CoupledResonators(
        zeros(0, 0),
        zeros(0, 0),
        zeros(0, 0),
        zeros(0, 0),
        Diagonal{Real}([]),
        Diagonal{Real}([]),
     )
end

modes(cr::CoupledResonators) = size(cr.yr, 1)

function addModes(cr::CoupledResonators, l::Int)
    if l == 0
        return cr
    end
    return CoupledResonators(
    Matrix(cr.yr ⊕ zeros(l, l)),
    Matrix(cr.yi ⊕ zeros(l, l)),
    Matrix(cr.wr ⊕ zeros(l, l)),
    Matrix(cr.wi ⊕ zeros(l, l)),
    Diagonal( [ diag(cr.γin)..., zeros(l)... ] ),
    Diagonal( [ diag(cr.γex)..., ones(l)... ] )
    )
end

# add g a⁺ⱼaₖ + conj(g) aⱼa⁺ₖ
function addPassive(cr::CoupledResonators,
    g::Number, j::Int, k::Int)
    m = size(cr.yr, 1)
    cr = addModes(cr, max(j - m, k - m, 0 ))
    yr = cr.yr
    yi = cr.yi
    wi = cr.wi
    wr = cr.wr
    γin = cr.γin
    γex = cr.γex
    yr[j, k] = real(g)
    yr[k, j] = real(g)
    yi[j, k] = imag(g)
    yi[k, j] = -imag(g)
    return CoupledResonators(yr, yi, wr, wi, γin, γex)
end

# add g a⁺ⱼa⁺ₖ + conj(g) aⱼaₖ
function addActive(cr::CoupledResonators,
    g::Number, j::Int, k::Int)
    m = size(cr.yr, 1)
    cr = addModes(cr, max(j - m, k - m, 0 ))
    yr = cr.yr
    yi = cr.yi
    wi = cr.wi
    wr = cr.wr
    γin = cr.γin
    γex = cr.γex
    wr[j, k] = real(g) / 2
    wr[k, j] = real(g) / 2
    wi[j, k] = imag(g) / 2
    wi[k, j] = imag(g) / 2
    return CoupledResonators(yr, yi, wr, wi, γin, γex)
end

function addGammaIn(cr::CoupledResonators,
    γ::Number, l::Int)
    m = size(cr.yr, 1)
    cr = addModes(cr, max(l - m, 0 ))
    yr = cr.yr
    yi = cr.yi
    wi = cr.wi
    wr = cr.wr
    γin = cr.γin
    γex = cr.γex
    γin[l, l] = γ
    return CoupledResonators(yr, yi, wr, wi, γin, γex)
end

function addGammaEx(cr::CoupledResonators,
    γ::Number, l::Int)
    m = size(cr.yr, 1)
    cr = addModes(cr, max(l - m, 0 ))
    yr = cr.yr
    yi = cr.yi
    wi = cr.wi
    wr = cr.wr
    γin = cr.γin
    γex = cr.γex
    γex[l, l] = γ
    return CoupledResonators(yr, yi, wr, wi, γin, γex)
end

const e0 = I(4)

const e1 = [0 0 0 1;
            0 0 1 0;
            0 1 0 0;
            1 0 0 0]

const e2 = [0 0 1 0;
            0 0 0 -1;
            1 0 0 0;
            0 -1 0 0]

const e3 = [1 0 0 0;
            0 1 0 0;
            0 0 -1 0;
            0 0 0 -1]

function scatteringMatrix(ω, yr, yi, wr, wi, γin, γex)
    m = size(yr, 1)
    Θ = ω * inv(γex + γin)
    Yr = inv(sqrt.(γex)) * yr * inv(sqrt.(γex))
    Yi = inv(sqrt.(γex)) * yi * inv(sqrt.(γex))
    Wr = inv(sqrt.(γex)) * wr * inv(sqrt.(γex))
    Wi = inv(sqrt.(γex)) * wi * inv(sqrt.(γex))
    Γ = inv(sqrt.(γex)) * (γex + γin) * inv(sqrt.(γex)) / 2
    S = e0 ⊗ I(m) - inv(
          e0 ⊗ (Yi + Γ) -
          e1 ⊗ Wr +
          e2 ⊗ Wi -
          (e1 * e2) ⊗ Yr +
          (e1 * e2 * e3) ⊗ Θ
    )
    order = vcat([ [i, i + m] for i in 1:m]...,  [ [i + 2*m, i + 3*m] for i in 1:m]...)
    return S[order, order]
end
function scatteringMatrix(ω::Number, cr::CoupledResonators)
    yr = cr.yr
    yi = cr.yi
    wr = cr.wr
    wi = cr.wi
    γin = cr.γin
    γex = cr.γex
    return scatteringMatrix(ω, yr, yi, wr, wi, γin, γex)
end


# Definition of all routes of the route tree.

function fade(x, y, speed, δt)
    return x, y + speed * δt
end

function post(x, y, speed, δt; cut = 10.0, angle = 45.0)
    if y < cut
        return x, y + speed *δt
    else
        return x + speed * δt * cosd(angle), y + speed * δt * sind(angle)
    end
end


function prime_factors(n)
    found = false

    v = 2

    i = 2

    while i < n - 1 && !found
        if mod(n, i) == 0
            found = true

            v = i
        end

        i += 1
    end

    if found
        factors = vcat([v], prime_factors(div(n, v)))
    else
        factors = [n]
    end

    factors
end

println(maximum(prime_factors(600851475143)))


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using QHull

function GLLar(points::Array{Float64,2}, cells::Array{Array{Int64, 1}, 1})

	points = convert(Array{Float64,2}, points')
	vertices=Vector{Float32}()
	normals =Vector{Float32}()
	for cell in cells

		ch = QHull.chull(points[cell,:])
		verts = ch.vertices
		vdict = Dict(zip(verts, 1:length(verts)))
		hfaces = [[vdict[u],vdict[v],vdict[w]] for (u,v,w) in ch.simplices]
		hdict = Dict(zip(1:length(cell), cell))
		trias = [[hdict[t1],hdict[t2],hdict[t3]] for (t1,t2,t3) in hfaces]
		pts = points[verts,:]

		for tria in trias
			p3,p2,p1 = points[tria[1],:],points[tria[2],:],points[tria[3],:]
			p1 = convert(Point3d, p1)
			p2 = convert(Point3d, p2)
			p3 = convert(Point3d, p3)
			n=0.5*computeNormal(p1::Point3d,p2::Point3d,p3::Point3d)

			append!(vertices,p1); append!(normals,n)
			append!(vertices,p2); append!(normals,n)
			append!(vertices,p3); append!(normals,n)
		end
	end

	ret=GLMesh(GL_TRIANGLES)
	ret.vertices = GLVertexBuffer(vertices)
	ret.normals  = GLVertexBuffer(normals)
	return ret
end

V,CV = Lar.cuboidGrid([20,10,1])
VIEW([
      #GLCuboid(Box3d(Point3d(0,0,0),Point3d(1,1,1)))
      GLLar(V,CV)
      GLAxis(Point3d(0,0,0),Point3d(1,1,1))
])


using Random
import StatsBase

"""
An empirical distribution, defined by a histogram which assigns probabilities to
a set of axis-coordinates in a (possibly high-dimensional) grid.
"""
struct EmpiricalDistribution{Dims, T}
    prob::Array{Float64, Dims}
    axes::SVector{Dims, Vector{T}}
end

Base.eltype(::Type{<:EmpiricalDistribution{Dims, T}}) where {Dims, T} = Vector{T}

function empirical_dist(prob::AbstractArray{Float64,N}, axes::Vararg{<:AbstractVector,N}) where {N}
    EmpiricalDistribution{N, eltype(axes[1])}(prob ./ sum(prob), (axes...,))
end

function Random.rand(rng::AbstractRNG, d::Random.SamplerTrivial{<:EmpiricalDistribution})
    dist = d[]
	x = StatsBase.sample(rng, StatsBase.weights(vec(dist.prob)))
	index = CartesianIndices(size(dist.prob))
	i = index[x]
	map(j->dist.axes[j][i[j]], 1:length(i))
end

function Distributions.logpdf(dist::EmpiricalDistribution, v::AbstractVector)
	indices = [findfirst(==(v[i]), dist.axes[i]) for i=eachindex(dist.axes)]
	if nothing ∈ indices
		-Inf64
	else
		log(dist.prob[indices...])
	end
end
Distributions.logpdf(dist::EmpiricalDistribution, v) = logpdf(dist, [v])

function Distributions.pdf(dist::EmpiricalDistribution, v::AbstractVector)
	indices = [findfirst(==(v[i]), dist.axes[i]) for i=eachindex(dist.axes)]
	if nothing ∈ indices
		0.0
	else
		dist.prob[indices...]
	end
end
Distributions.pdf(dist::EmpiricalDistribution, v) = pdf(dist, [v])

function marginalize(dist::EmpiricalDistribution, axis_index::Integer)
	new_prob = sum(dist.prob, dims=axis_index)
	new_prob = dropdims(new_prob, dims=axis_index)
	new_axes = deleteat(dist.axes, axis_index)
	EmpiricalDistribution(new_prob, new_axes)
end


"""
    batch_quadrature(
        fs::AbstractVector,
        ms::AbstractVector{<:Real},
        σs::AbstractVector{<:Real},
        num_points::Integer,
    )

Approximate the integrals
```julia
∫ fs[n](x) pdf(Normal(ms[n], σs[n]), x) dx
```
for all `n` in `eachindex(fs)` using Gauss-Hermite quadrature with `num_points`.
"""
function batch_quadrature(
    fs::AbstractVector,
    ms::AbstractVector{<:Real},
    σs::AbstractVector{<:Real},
    num_points::Integer,
)
    # Check that as many bounds are provided as we have functions to integrate.
    length(fs) == length(ms) || throw(error("length(fs) != length(ms)"))
    length(fs) == length(σs) || throw(error("length(fs) != length(σs)"))

    # Construct the quadrature points.
    xs, ws = gausshermite(num_points)

    # Compute the integral.
    return map((f, m, σ) -> _gauss_hermite_quadrature(f, m, σ, xs, ws), fs, ms, σs)
end

Zygote.@nograd gausshermite

# Internal method. Assumes that the gradient w.r.t. xs and ws is never needed, so avoids
# computing it and returns nothing. This is potentially not what you want in general.
function _gauss_hermite_quadrature(f, m::Real, σ::Real, xs, ws)
    t(x, m, σ) = m + sqrt(2) * σ * x
    I = ws[1] * f(t(xs[1], m, σ))
    for j in 2:length(xs)
        I += ws[j] * f(t(xs[j], m, σ))
    end
    return I / sqrt(π)
end

function Zygote._pullback(
    ctx::Zygote.AContext, ::typeof(_gauss_hermite_quadrature), f, m::Real, σ::Real, xs, ws,
)
    function _gauss_hermite_quadrature_pullback(Δ::Real)
        g(f, x, w, m, σ) = w * f(m + sqrt(2) * σ * x)

        _, pb = Zygote._pullback(ctx, g, f, xs[1], ws[1], m, σ)
        _, Δf, _, _, Δm, Δσ = pb(Δ / sqrt(π))
        for j in 2:length(xs)
            _, pb = Zygote._pullback(ctx, g, f, xs[j], ws[j], m, σ)
            _, Δf_, _, _, Δm_, Δσ_ = pb(Δ / sqrt(π))
            Δf = Zygote.accum(Δf, Δf_)
            Δm = Zygote.accum(Δm, Δm_)
            Δσ = Zygote.accum(Δσ, Δσ_)
        end
        return nothing, Δf, Δm, Δσ, nothing, nothing
    end
    return _gauss_hermite_quadrature(f, m, σ, xs, ws), _gauss_hermite_quadrature_pullback
end


module DiscretisedFluidQueues

import Jacobi, LinearAlgebra, SparseArrays, StaticArrays
import Base: *, +, size, show, getindex, setindex!, sum

# Types
export Model, DiscretisedFluidQueue, BoundedFluidQueue, Phase, PhaseSet # Queues are <:Model
export AbstractMatrixExponential, ConcentratedMatrixExponential, MatrixExponential
export Mesh, DGMesh, FRAPMesh, FVMesh # are <:Mesh
export Generator, FullGenerator, LazyGenerator # are <:Generator
export SFMDistribution, AbstractIntegrationMethod, AutoQuadrature, Quadrature, TrapezoidRule
export ExplicitRungeKuttaScheme, ForwardEuler, Heuns, StableRK3, StableRK4 # Euler, RungeKutta4 <: TimeIntegrationScheme
export Simulation

# Functions
export augment_model, membership, N₋, N₊, n_phases, phases, rates # Model methods
export cell_nodes, Δ, n_bases_per_cell, n_bases_per_phase, n_intervals, total_n_bases # Mesh methods
export interior_point_mass, left_point_mass, right_point_mass, integrate_time # SFMDistribution methods
export simulate, fixed_time, n_jumps, first_exit_x # Simulation methods
export build_lazy_generator, build_full_generator, static_generator
export cme_params, pdf, ccdf, cdf, build_me, cell_probs
export normalised_closing_operator_cdf, normalised_closing_operator_pdf
export naive_normalised_closing_operator_cdf, naive_normalised_closing_operator_pdf
export unnormalised_closing_operator_cdf, unnormalised_closing_operator_pdf
export limit, Limiter, NoLimiter, GeneralisedMUSCL

# model
include("SFM.jl")

include("abstract_mesh.jl") # things which apply to all meshes
include("discretised_fluid_queue.jl")
include("lazy_generators.jl")
include("full_generators.jl")

# auxillary functions
include("ME_tools.jl") # used in FRAPApproximation.jl
include("polynomials.jl") # used in discontinuous_Galerkin.jl

include("discontinuous_Galerkin.jl")
include("finite_volume_method.jl")
include("FRAP_approximation.jl")

include("distributions.jl")
include("extend_Base_operations.jl")
include("limiters.jl")

include("time_integration.jl")

include("simulate.jl")

end


# Homology: type and methods

abstract type AbstractHomology end
grouptype(::Type{AbstractHomology}) = Nothing
group(h::AbstractHomology, dim::Int; kw...) = throw(MethodError(group,(typeof(h),Int)))

"""
Homology group iterator for an abstract complex
"""
struct Homology{C<:AbstractComplex, PID} <: AbstractHomology
    complex::C
end
homology(::Type{PID}, c::C) where {C<:AbstractComplex, PID} = Homology{C, PID}(c)

show(io::IO, h::Homology{C, PID}) where {C<:AbstractComplex, PID} = print(io, "Homology($PID)[$(h.complex)]")

#
# Interface methods
#

grouptype(::Type{Homology{C, PID}}) where {C<:AbstractComplex, PID} = PID

function group(h::Homology{C, PID}, p::Int; Dₚ::Int=0) where {C<:AbstractComplex, PID}
    cdim = dim(h.complex)
    @assert cdim >= p "Cannot define $p-th homology group for $cdim-dimensional complex"

    M = boundary(h.complex, p+1, PID)
    F = smith(M)
    D = diagm(F)

    nₚ, nₚ₊₁ = size(D)
    Dₚ₊₁ = trivial = 0

    # empty boundary
    nₚ₊₁ == 0 && return (0, 0, (F.S, F.Sinv, F.T, F.Tinv, D, Dₚ₊₁))

    # Calculate rank B_p = rank D_{p+1}
    for i in 1:nₚ₊₁
        D[i,i] == zero(PID) && break
        Dₚ₊₁ += 1
    end

    # rank of torsion-free subgroup
    for i in 1:min(nₚ, nₚ₊₁)
        if D[i,i] == one(PID) || D[i,i] == -one(PID)
           trivial += 1
        end
    end

    # β_p = rank H p = rank Z_p - rank B_p = n_p − rank D_p - rank D_{p+1}
    βₚ = nₚ - Dₚ - Dₚ₊₁
    τₚ = Dₚ₊₁ - trivial

    return (βₚ, τₚ, (F.S, F.Sinv, F.T, F.Tinv, D, Dₚ₊₁))
end

#
# Iterator methods
#
"""Return homology group type: dimension, Betti & torsion numbers."""
eltype(h::Homology) = Tuple{Int, Int, Int}

length(h::Homology) = dim(h.complex)

function iterate(h::Homology{C, PID}, state=nothing) where {C<:AbstractComplex, PID}
    if state === nothing
        Z = zeros(PID,0,0)
        snfstate = (Z,Z,Z,Z,Z,0)
        return iterate(h, (0, snfstate))
    end

    p = state[1]
    p >= dim(h.complex) && return nothing

    βₚ, τₚ, snfstate = group(h, p, Dₚ = state[2][end])
    return ((p, βₚ, τₚ), (p+1, snfstate))
end

"""
Homology generator iterator
"""
struct WithGenerators{H <: AbstractHomology}
    homology::H
end

"""
    withgenerators(hom)

An iterator that yields homology generators from a homology iterator `hom`.

Returns homology group parameters from iterator `hom` and generators as pair of `Chain` **x** and coefficient **k** that compose boundary **kx**.
When **k** is zero, then **x** is a boundary without any coefficient.
"""
withgenerators(h::H) where {H <: AbstractHomology} = WithGenerators{H}(h)

#
# Iterator methods
#
length(g::WithGenerators) = length(g.homology)
eltype(g::WithGenerators) = Tuple{Int,Int,Int,Dict{Chain, grouptype(typeof(g.homology))}}

function iterate(g::WithGenerators, state=nothing)
    h = g.homology

    # calculate betti, torsion & SNF
    result = iterate(h, state)
    result === nothing && return nothing

    GT = grouptype(typeof(h))
    chains = Dict{AbstractChain, GT}()

    p, βₚ, τₚ = result[1]
    U, Uinv, V, Vinv, D, t = result[2][2]
    cplx = g.homology.complex
    trivial = t - τₚ
    cls = cells(h.complex, p)

    A = view(U, 1:size(U,1), t+1:size(U,2))
    B = view(Uinv, t+1:size(Uinv,1), 1:size(Uinv,2))
    C = if p == 0
        n = size(cplx, p)
        SparseMatrixCSC{GT}(I, n, n)
    else
        VinvPrev = state[2][4]
        tPrev = state[2][6]
        view(VinvPrev, 1:size(VinvPrev,1), tPrev+1:size(VinvPrev,2))
    end

    P = A * (B * C)
    F = smith(P)
    G = F.S * diagm(F)

    # betti generator
    for j in 1:size(G,2)
        ii, V = findnz(G[:,j])
        length(ii) == 0 && continue
        ch = Chain(p, GT)
        for (i,v) in zip(ii,V)
            push!(ch, hash(cls[i])=>v)
        end
        chains[ch] = 0
    end

    # torsion generator
    for j in 1:τₚ
        ch = Chain(p, GT)
        I, V = findnz(D[:,trivial+j])
        for (i,v) in zip(I,V)
            push!(ch, hash(cls[i])=>v)
        end
        chains[ch] = D[trivial+j, trivial+j]
    end

    return (p, βₚ, τₚ, chains), result[2]
end

#
# Auxiliary  methods
#

"""Calculate Betti number"""
betti(g::AbstractHomology) = tuple((gst[2] for gst in g)...)

"""Calculate Euler characteristic"""
euler(btn::Tuple) = sum(isodd(i) ? v : -v for (i,v) in enumerate(btn))

"""Return linearly independent generators"""
function generators(g::WithGenerators{H}) where {H <: AbstractHomology}
    chains = Dict{Int,Vector{Chain}}()
    for (d,b,t,gens) in g
        chains[d] = collect(keys(gens))
    end
    return chains
end


# compute an ensemble of ϕ(u) scalar observable values
function run_obs_ensemble(u0data, Tmax, ϕ)

    n = length(u0data);
    samples = zeros(n);

    Threads.@threads for i in 1:n
        soln = solve_ode(u0data[i], Tmax);
        samples[i] = ϕ(soln)
    end

    return samples
end

# compute an ensemble of the paths at the indicated values of
function run_path_ensemble(u0data, Tmax, t_vals)

    n = length(u0data);
    u1_samples = zeros(n, length(t_vals));
    u2_samples = zeros(n, length(t_vals));

    Threads.@threads for i in 1:n
        soln = solve_ode(u0data[i], Tmax);
        soln_at_t = soln(t_vals);
        u1_samples[i,:] .= [u[1] for u in soln_at_t];
        u2_samples[i,:] .= [u[2] for u in soln_at_t];
    end

    return u1_samples,u2_samples
end


using MonteCarlo, Distributions, PyPlot, DataFrames, JLD

Tdist = Normal(MonteCarlo.IsingTc, .64)
n_Ts = 2^8 # 2^8
Ts = sort!(rand(Tdist, n_Ts))
Ts = Ts[Ts.>=1.2]
Ts = Ts[Ts.<=3.8]
therm = 10^4 # 10^4
sweeps = 10^3 # 10^3

df = DataFrame(L=Int[], T=Float64[], M=Float64[], χ=Float64[], E=Float64[], C_V=Float64[])

for L in 2 .^ [3, 4, 5, 6] #[3, 4, 5, 6]
	println("L = ", L)
	for (i, T) in enumerate(Ts)
		println("\t T = ", T)
		beta = 1/T
		model = IsingModel(dims=2, L=L)
		mc = MC(model, beta=beta)
		# mc.p.global_moves = true # enable Wolff cluster
		# mc.p.global_rate = 1
		run!(mc, sweeps=sweeps, thermalization=therm, verbose=false)
		meas = MonteCarlo.measurements(mc)[:ME]
		push!(df, [L, T, mean(meas[:Magn].m), mean(meas[:Magn].chi), mean(meas[:Energy].e), mean(meas[:Energy].C)])
	end
	flush(stdout)
end

sort!(df, [:L, :T])
@save "ising2d.jld" df

grps = groupby(df, :L)
# plot results individually
for g in grps
	L = g[:L][1]
	fig, ax = subplots(2,2, figsize=(12,8))
	ax[1].plot(g[:T], g[:E], "o", color="darkred", markeredgecolor="black")
	ax[1].set_ylabel("Energy")
	ax[1].set_xlabel("Temperature")

	ax[2].plot(g[:T], g[:C_V], "o", color="darkred", markeredgecolor="black")
	ax[2].set_ylabel("Specific heat")
	ax[2].set_xlabel("Temperature")

	ax[3].plot(g[:T], g[:M], "o", color="C0", markeredgecolor="black")
	ax[3].axvline(x=MonteCarlo.IsingTc, color="black", linewidth=2.0, label="\$ T_c \$")
	# ax[3].legend(loc="best")
	ax[3].set_ylabel("Magnetization")
	ax[3].set_xlabel("Temperature")

	ax[4].plot(g[:T], g[:χ], "o", color="C0", markeredgecolor="black")
	ax[4].axvline(x=MonteCarlo.IsingTc, color="black", linewidth=2.0, label="\$ T_c \$")
	# ax[4].legend(loc="best")
	ax[4].set_ylabel("Susceptibility χ")
	ax[4].set_xlabel("Temperature")
	tight_layout()
	savefig("ising2d_L_$(L).pdf")
end

# plot results together
fig, ax = subplots(2,2, figsize=(12,8))

for g in grps
	L = g[:L][1]
	ax[1].plot(g[:T], g[:E], "o", markeredgecolor="black", label="L=$L")
	ax[2].plot(g[:T], g[:C_V], "o", markeredgecolor="black", label="L=$L")
	ax[3].plot(g[:T], g[:M], "o", markeredgecolor="black", label="L=$L")
	ax[4].plot(g[:T], g[:χ], "o", markeredgecolor="black", label="L=$L")
end
ax[1].legend(loc="best")
ax[1].set_ylabel("Energy")
ax[1].set_xlabel("Temperature")

ax[2].set_ylabel("Specific heat")
ax[2].set_xlabel("Temperature")
ax[2].axvline(x=MonteCarlo.IsingTc, color="black", linestyle="dashed", label="\$ T_c \$")
ax[2].legend(loc="best")

ax[3].set_ylabel("Magnetization")
ax[3].set_xlabel("Temperature")
# ax[3].axvline(x=MonteCarlo.IsingTc, linewidth=2.0, color="black", label="\$ T_c \$")
x = range(1.2, stop=MonteCarlo.IsingTc, length=100)
y = (1 .- sinh.(2.0 ./ (x)).^(-4)).^(1/8)
ax[3].plot(x,y, "k--", label="exact")
ax[3].plot(range(MonteCarlo.IsingTc, stop=3.8, length=100), zeros(100), "k--")
ax[3].legend(loc="best")

ax[4].set_ylabel("Susceptibility χ")
ax[4].set_xlabel("Temperature")
ax[4].axvline(x=MonteCarlo.IsingTc, color="black", linestyle="dashed", label="\$ T_c \$")
ax[4].legend(loc="best")
tight_layout()
savefig("ising2d.pdf")
savefig("ising2d.svg")


"""
    randompartition(h::Hypergraph, n::Int)::Vector{Set{Int}}

Generates a random partition for vertices of a hypergraph `h` into `n` subsets.
"""
randompartition(h::Hypergraph, n::Int) = randompartition(nhv(h), n)

"""
    randompartition(N::Int, n::Int)::Vector{Set{Int}}

Generates a random partition for graph having `N` vertices into `n` subsets.
"""
function randompartition(N::Int, n::Int)
    res = [Set{Int}() for i in 1:n]
    for i in 1:N
        push!(res[rand(1:n)], i)
    end
    res
end

"""
    HypergraphAggs(h::Hypergraph)

Precomputes vertex and edge basic stats for a hypergraph.
The stats are being used for efficiency reasons by community search algorithms.
"""
struct HypergraphAggs
    hes::Vector{Int}
    max_hes::Int
    Ed::Vector{Int}
    deg_vs::Vector{Int}
    volV::Int
    HypergraphAggs(h::Hypergraph) = begin
        hes = [length(h.he2v[i]) for i in 1:nhe(h)]
        max_hes = maximum(hes)
        Ed = zeros(Int, max_hes)
        for ii in 1:length(hes)
            if hes[ii] >= 1
                Ed[hes[ii]]+=1
            end
        end
        deg_vs = [length(h.v2he[i]) for i in 1:nhv(h)]
        volV = sum(deg_vs)
        new(hes,max_hes,Ed,deg_vs,volV)
    end
end

"""
    LightGraphs.modularity(h::Hypergraph, partition::Vector{Set{Int}},
ha::HypergraphAggs=HypergraphAggs(h))

Calculates the strict modularity of a hypergraph `h` for a given `partition` using
the precomputed aggregates `ha`.
"""
@inline function LightGraphs.modularity(h::Hypergraph, partition::Vector{Set{Int}},
        ha::HypergraphAggs=HypergraphAggs(h))

    @boundscheck sum(length.(partition)) == nhv(h)
    @boundscheck union(partition...) == Set(1:nhv(h))
    volP_volV = [sum(ha.deg_vs[i] for i in p)/ha.volV for p in partition]
    eP = [count(i-> ha.hes[i]>0 && (keys(h.he2v[i]) ⊆ p), 1:nhe(h)) for p in partition]
    (sum(eP) - sum( ha.Ed[d]*sum(volP_volV.^d) for d in 1:ha.max_hes)) / nhe(h)
end

"""
The base type for all algorithms representing various community search patterns.
"""
abstract type AbstractCommunityFinder end

"""
    CFModularityRandom(n::Int, reps::Int) <: AbstractCommunityFinder

Represents a random search over the hypergraph `h` that finds
a partition into `n` communities (subsets) having the maximum modularity value.
During the search `reps` random `n`-partitions will be evaluated.
If there are many partitions having the same value the first
one that was randomly found will be returned.
"""
struct CFModularityRandom <: AbstractCommunityFinder
    n::Int
    reps::Int
end

"""
    findcommunities(h::Hypergraph, method::CFModularityRandom)

Makes a random search over the hypergraph `h` and finds
a partition into `method.n` communities (subsets) having the maximum modularity value.
During the search `method.reps` random `n`-partitions will be evaluated.
If there are many partitions having the same value the first
one that was randomly found will be returned.

Returns a `NamedTuple` where the field `bp` contains partition
and the field `bm` contains the modularity value for that partition.
"""
function findcommunities(h::Hypergraph, method::CFModularityRandom)
    bp = [Int[]]
    bm = -Inf
    ha = HypergraphAggs(h)
    for i in 1:method.reps
        p = Vector{Vector{Int}}([[]])
        while  minimum(length.(p)) == 0
            p = randompartition(h, method.n)
        end
        m = modularity(h, p, ha)
        if m > bm
            bm = m
            bp = p
        end
    end
    (bp=bp, bm=bm)
end

function find_first(c::Array{Set{Int}}, vals)
    for i in 1:length(c)
        for v in vals
            v in c[i] && return i
        end
    end
    throw("None of values in $vals found")
end

"""
    CFModularityCNMLike(n::Int, reps::Int) <: AbstractCommunityFinder

Represents a CNM-Like algorithm for finding communities.
In the algorithm we start with a partition where each node is in its own part.
Then in each step, we randomly select a hyperedge.
Subsequently, we consider merging each set of that parts it touches.
We actually merge the parts if the new best modularity is at least as high
as the modularity from the previous step.
The algortithm iterates through `reps` of repetitions.

For more information see `Algorithm 1` at:
Clustering via Hypergraph Modularity (submitted to Plos ONE), auhtors:
Bogumił Kamiński, Valerie Poulin, Paweł Prałat, Przemysław Szufel, Francois Theberge

"""
struct CFModularityCNMLike <: AbstractCommunityFinder
    reps::Int
end

"""
    findcommunities(h::Hypergraph, method::CFModularityCNMLike)

Iterates a CNM-Like algorithm for finding communities.
In the algorithm we start with a partition where each node is in its own part.
Then in each step, we randomly select a hyperedge.
Subsequently, we consider merging each set of that parts it touches.
We actually merge the parts if the new best modularity is at least as high
as the modularity from the previous step.

Returns a `NamedTuple` where the field `bp` contains partition
and the field `bm` contains the modularity value for that partition,
finally, the fiel `mod_history` represents modularities achieved
in subsequent steps of the algorithm.

For more information see `Algorithm 1` at:
Clustering via Hypergraph Modularity (submitted to Plos ONE), authors:
Bogumił Kamiński, Valerie Poulin, Paweł Prałat, Przemysław Szufel, Francois Theberge

"""
function findcommunities(h::Hypergraph, method::CFModularityCNMLike)
    ha = HypergraphAggs(h)
    best_modularity = 0
    comms = [Set(i) for i in 1:nhv(h)]
    mod_history = Vector{Float64}(undef, method.reps)
    for rep in 1:method.reps
        he = rand(1:nhe(h))
        vers = collect(keys(getvertices(h, he)))
        if length(vers) == 0
            continue
        end;
        c = deepcopy(comms)
        i0 = find_first(c, vers)
        max_i = length(c)
        i_cur = i0
        while i_cur < max_i
            i_cur += 1
            if length(intersect(c[i_cur],vers)) > 0
                union!(c[i0], c[i_cur])
                c[i_cur]=c[max_i]
                max_i += -1
            end
        end
        resize!(c,max_i)
        m = modularity(h, c, ha)
        if m > best_modularity
            best_modularity = m
            comms = c
        end
        mod_history[rep] = best_modularity
    end
    return (bm=best_modularity, bp=comms, mod_history=mod_history)
end


println("################################################################################")
# https://docs.julialang.org/en/v1/stdlib/Test/#Basic-Unit-Tests
using Test, Collatz

# Test function collatz_function(n::Integer; P::Integer=2, a::Integer=3, b::Integer=1)
@testset verbose = true "collatz_function" begin
    # Default/Any (P,a,b); 0 trap
    @test Collatz.collatz_function(0) == 0
    # Default/Any (P,a,b); 1 cycle; positives
    @test Collatz.collatz_function(1) == 4
    @test Collatz.collatz_function(4) == 2
    @test Collatz.collatz_function(2) == 1
    # Default/Any (P,a,b); -1 cycle; negatives
    @test Collatz.collatz_function(-1) == -2
    @test Collatz.collatz_function(-2) == -1
    # Test a wider modulo sweep by upping P to 5, a to 2, and b to 3.
    @test Collatz.collatz_function(1, P=5, a=2, b=3) == 5
    @test Collatz.collatz_function(2, P=5, a=2, b=3) == 7
    @test Collatz.collatz_function(3, P=5, a=2, b=3) == 9
    @test Collatz.collatz_function(4, P=5, a=2, b=3) == 11
    @test Collatz.collatz_function(5, P=5, a=2, b=3) == 1
    # Test negative P, a and b. %, used in the function, is "floor" in python
    # rather than the more reasonable euclidean, but we only use it's (0 mod P)
    # conjugacy class to determine functionality, so the flooring for negative P
    # doesn't cause any issue.
    @test Collatz.collatz_function(1, P=-3, a=-2, b=-5) == -7
    @test Collatz.collatz_function(2, P=-3, a=-2, b=-5) == -9
    @test Collatz.collatz_function(3, P=-3, a=-2, b=-5) == -1
    # Set P and a to 0 to @test on __assert_sane_parameterisation
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.collatz_function(1, P=0, a=2, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.collatz_function(1, P=0, a=0, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_A) Collatz.collatz_function(1, P=1, a=0, b=3)
end

# Test function reverse_collatz_function(n::Integer; P::Integer=2, a::Integer=3, b::Integer=1)
@testset verbose = true "reverse_collatz_function" begin
    # Default (P,a,b); 0 trap [as b is not a multiple of a]
    @test Collatz.reverse_collatz_function(0) == [0]
    # Default (P,a,b); 1 cycle; positives
    @test Collatz.reverse_collatz_function(1) == [2]
    @test Collatz.reverse_collatz_function(4) == [1, 8]
    @test Collatz.reverse_collatz_function(2) == [4]
    # Default (P,a,b); -1 cycle; negatives
    @test Collatz.reverse_collatz_function(-1) == [-2]
    @test Collatz.reverse_collatz_function(-2) == [-4, -1]
    # Test a wider modulo sweep by upping P to 5, a to 2, and b to 3.
    @test Collatz.reverse_collatz_function(1, P=5, a=2, b=3) == [-1, 5]
    @test Collatz.reverse_collatz_function(2, P=5, a=2, b=3) == [10]
    @test Collatz.reverse_collatz_function(3, P=5, a=2, b=3) == [15]  # also tests !0
    @test Collatz.reverse_collatz_function(4, P=5, a=2, b=3) == [20]
    @test Collatz.reverse_collatz_function(5, P=5, a=2, b=3) == [1, 25]
    # Test negative P, a and b. %, used in the function, is "floor" in python
    # rather than the more reasonable euclidean, but we only use it's (0 mod P)
    # conjugacy class to determine functionality, so the flooring for negative P
    # doesn't cause any issue.
    @test Collatz.reverse_collatz_function(1, P=-3, a=-2, b=-5) == [-3]  # != [-3, -3]
    @test Collatz.reverse_collatz_function(2, P=-3, a=-2, b=-5) == [-6]
    @test Collatz.reverse_collatz_function(3, P=-3, a=-2, b=-5) == [-9, -4]
    # Set P and a to 0 to @test on __assert_sane_parameterisation
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.reverse_collatz_function(1, P=0, a=2, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.reverse_collatz_function(1, P=0, a=0, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_A) Collatz.reverse_collatz_function(1, P=1, a=0, b=3)
    # If b is a multiple of a, but not of Pa, then 0 can have a reverse.
    @test Collatz.reverse_collatz_function(0, P=17, a=2, b=-6) == [0, 3]
    @test Collatz.reverse_collatz_function(0, P=17, a=2, b=102) == [0]
end

# Test function hailstone_sequence(initial_value::Integer; P::Integer=2, a::Integer=3, b::Integer=1, max_total_stopping_time::Integer=1000, total_stopping_time::Bool=true, verbose::Bool=true)
@testset verbose = true "hailstone_sequence" begin
    # Test 0's immediated termination.
    @test Collatz.hailstone_sequence(0) == [[_CC.ZERO_STOP, 0]]
    @test Collatz.hailstone_sequence(0, verbose=false) == [0]
    # The cycle containing 1 wont yield a cycle termination, as 1 is considered
    # the "total stop" that is the special case termination.
    @test Collatz.hailstone_sequence(1) == [[_CC.TOTAL_STOPPING_TIME, 0]]
    @test Collatz.hailstone_sequence(1, verbose=false) == [1]
    # Test the 3 known default parameter's cycles (ignoring [1,4,2])
    # If not verbose, then the result will be the cycle plus the final value
    # being the start of the cycle. If verbose, then the output will be the
    # values not in the cycle, a CC flag, then the cycle and another CC flag.
    for kc in [kc for kc in _KNOWN_CYCLES if !(1 in kc)]
        _temp_kc = deepcopy(kc); push!(_temp_kc, _temp_kc[1])
        @test Collatz.hailstone_sequence(kc[1], verbose=false) == _temp_kc
        @test Collatz.hailstone_sequence(kc[1]) == [_CC.CYCLE_INIT, kc, [_CC.CYCLE_LENGTH, length(kc)]]
    end
    # Test the lead into a cycle by entering two of the cycles.
    seq = [kc for kc in _KNOWN_CYCLES if -5 in kc][1]
    seq = append!(append!([seq[2]*4, seq[2]*2], seq[2:length(seq)]), [seq[1]])
    @test Collatz.hailstone_sequence(-56, verbose=false) == push!(deepcopy(seq), seq[3])
    @test Collatz.hailstone_sequence(-56) == [seq[1], seq[2], _CC.CYCLE_INIT, seq[3:length(seq)], [_CC.CYCLE_LENGTH, length(seq)-2]]
    seq = [kc for kc in _KNOWN_CYCLES if -17 in kc][1]
    seq = append!(append!([seq[2]*4, seq[2]*2], seq[2:length(seq)]), [seq[1]])
    @test Collatz.hailstone_sequence(-200, verbose=false) == push!(deepcopy(seq), seq[3])
    @test Collatz.hailstone_sequence(-200) == [seq[1], seq[2], _CC.CYCLE_INIT, seq[3:length(seq)], [_CC.CYCLE_LENGTH, length(seq)-2]]
    # 1's cycle wont yield a description of it being a "cycle" as far as the
    # hailstones are concerned, which is to be expected, so..
    @test Collatz.hailstone_sequence(4, verbose=false) == [4, 2, 1]
    @test Collatz.hailstone_sequence(4) == [4, 2, 1, [_CC.TOTAL_STOPPING_TIME, 2]]
    @test Collatz.hailstone_sequence(16, verbose=false) == [16, 8, 4, 2, 1]
    @test Collatz.hailstone_sequence(16) == [16, 8, 4, 2, 1, [_CC.TOTAL_STOPPING_TIME, 4]]
    # Test the regular stopping time check.
    @test Collatz.hailstone_sequence(4, total_stopping_time=false) == [4, 2, [_CC.STOPPING_TIME, 1]]
    @test Collatz.hailstone_sequence(5, total_stopping_time=false) == [5, 16, 8, 4, [_CC.STOPPING_TIME, 3]]
    # Test small max_total_stopping_time: (minimum internal value is one)
    @test Collatz.hailstone_sequence(4, max_total_stopping_time=-100) == [4, 2, [_CC.MAX_STOP_OOB, 1]]
    # Test the zero stop mid hailing. This wont happen with default params tho.
    @test Collatz.hailstone_sequence(3, P=2, a=3, b=-9) == [3, 0, [_CC.ZERO_STOP, -1]]
    # Lastly, while the function wont let you use a P value of 0, 1 and -1 are
    # still allowed, although they will generate immediate 1 or 2 length cycles
    # respectively, so confirm the behaviour of each of these hailstones.
    @test Collatz.hailstone_sequence(3, P=1, verbose=false) == [3, 3]
    @test Collatz.hailstone_sequence(3, P=-1, verbose=false) == [3, -3, 3]
    # Set P and a to 0 to @test on __assert_sane_parameterisation
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.hailstone_sequence(1, P=0, a=2, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.hailstone_sequence(1, P=0, a=0, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_A) Collatz.hailstone_sequence(1, P=1, a=0, b=3)
end

# Test function stopping_time(initial_value::Integer; P::Integer=2, a::Integer=3, b::Integer=1, max_stopping_time::Integer=1000, total_stopping_time::Bool=false)
@testset verbose = true "stopping_time" begin
    # Test 0's immediated termination.
    @test Collatz.stopping_time(0) == 0
    # The cycle containing 1 wont yield a cycle termination, as 1 is considered
    # the "total stop" that is the special case termination.
    @test Collatz.stopping_time(1) == 0
    # Test the 3 known default parameter's cycles (ignoring [1,4,2])
    # If not verbose, then the result will be the cycle plus the final value
    # being the start of the cycle. If verbose, then the output will be the
    # values not in the cycle, a CC flag, then the cycle and another CC flag.
    for c in [kc for kc in _KNOWN_CYCLES if !(1 in kc)]
        @test Collatz.stopping_time(c[1], total_stopping_time=true) == Base.Inf
    end
    # Test the lead into a cycle by entering two of the cycles. -56;-5, -200;-17
    @test Collatz.stopping_time(-56, total_stopping_time=true) == Base.Inf
    @test Collatz.stopping_time(-200, total_stopping_time=true) == Base.Inf
    # 1's cycle wont yield a description of it being a "cycle" as far as the
    # hailstones are concerned, which is to be expected, so..
    @test Collatz.stopping_time(4, total_stopping_time=true) == 2
    @test Collatz.stopping_time(16, total_stopping_time=true) == 4
    # Test the regular stopping time check.
    @test Collatz.stopping_time(4) == 1
    @test Collatz.stopping_time(5) == 3
    # Test small max_total_stopping_time: (minimum internal value is one)
    @test Collatz.stopping_time(5, max_stopping_time=-100) == nothing
    # Test the zero stop mid hailing. This wont happen with default params tho.
    @test Collatz.stopping_time(3, P=2, a=3, b=-9) == -1
    # Lastly, while the function wont let you use a P value of 0, 1 and -1 are
    # still allowed, although they will generate immediate 1 or 2 length cycles
    # respectively, so confirm the behaviour of each of these stopping times.
    @test Collatz.stopping_time(3, P=1) == Base.Inf
    @test Collatz.stopping_time(3, P=-1) == Base.Inf
    # One last one for the fun of it..
    @test Collatz.stopping_time(27, total_stopping_time=true) == 111
    # And for a bit more fun, common trajectories on
    for x in 0:0  #TODO: Extend this up to 0:4 when hailstone is made arbitrary integer safe.
        @test Collatz.stopping_time(27+576460752303423488*x) == 96
    end
    # Set P and a to 0 to @test on __assert_sane_parameterisation
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.stopping_time(1, P=0, a=2, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.stopping_time(1, P=0, a=0, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_A) Collatz.stopping_time(1, P=1, a=0, b=3)
end

# Test function tree_graph(initial_value::Integer, max_orbit_distance::Integer; P::Integer=2, a::Integer=3, b::Integer=1, __cycle_prevention::Union{Set{Integer},Nothing}=nothing)
@testset verbose = true "tree_graph" begin
    C = _CC.CYCLE_INIT  # Shorthand the cycle terminus
    D = Dict()  # Just to colourise the below in the editor..
    # The default zero trap
    @test Collatz.tree_graph(0, 0) == Dict(0=>D)
    @test Collatz.tree_graph(0, 1) == Dict(0=>Dict(C=>0))
    @test Collatz.tree_graph(0, 2) == Dict(0=>Dict(C=>0))
    # The roundings of the 1 cycle.
    @test Collatz.tree_graph(1, 1) == Dict(1=>Dict(2=>D))
    @test Collatz.tree_graph(1, 0) == Dict(1=>D)
    @test Collatz.tree_graph(1, 1) == Dict(1=>Dict(2=>D))
    @test Collatz.tree_graph(1, 2) == Dict(1=>Dict(2=>Dict(4=>D)))
    @test Collatz.tree_graph(1, 3) == Dict(1=>Dict(2=>Dict(4=>Dict(C=>1,8=>D))))
    @test Collatz.tree_graph(2, 3) == Dict(2=>Dict(4=>Dict(1=>Dict(C=>2),8=>Dict(16=>D))))
    @test Collatz.tree_graph(4, 3) == Dict(4=>Dict(1=>Dict(2=>Dict(C=>4)),8=>Dict(16=>Dict(5=>D,32=>D))))
    # The roundings of the -1 cycle
    @test Collatz.tree_graph(-1, 1) == Dict(-1=>Dict(-2=>D))
    @test Collatz.tree_graph(-1, 2) == Dict(-1=>Dict(-2=>Dict(-4=>D,C=>-1)))
    # Test a wider modulo sweep by upping P to 5, a to 2, and b to 3.
    T = function (x,y) Dict(1=>Dict(-1=>x,5=>y)) end
    orb_1 = T(D,D)
    orb_2 = T(Dict(-5=>D,-2=>D),Dict(C=>1,25=>D))
    T = function (x,y,z) Dict(1=>Dict(-1=>Dict(-5=>x,-2=>y),5=>Dict(C=>1,25=>z))) end
    orb_3 = T(Dict(-25=>D,-4=>D),Dict(-10=>D),Dict(11=>D,125=>D))
    @test Collatz.tree_graph(1, 1, P=5, a=2, b=3) == orb_1
    @test Collatz.tree_graph(1, 2, P=5, a=2, b=3) == orb_2
    @test Collatz.tree_graph(1, 3, P=5, a=2, b=3) == orb_3
    # Test negative P, a and b.
    orb_1 = Dict(1=>Dict(-3=>D))
    T = function (x,y) Dict(1=>Dict(-3=>Dict(-1=>x,9=>y))) end
    orb_2 = T(D,D)
    orb_3 = T(Dict(-2=>D,3=>D),Dict(-27=>D,-7=>D))
    @test Collatz.tree_graph(1, 1, P=-3, a=-2, b=-5) == orb_1
    @test Collatz.tree_graph(1, 2, P=-3, a=-2, b=-5) == orb_2
    @test Collatz.tree_graph(1, 3, P=-3, a=-2, b=-5) == orb_3
    # Set P and a to 0 to @test on __assert_sane_parameterisation
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.tree_graph(1, 1, P=0, a=2, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_P) Collatz.tree_graph(1, 1, P=0, a=0, b=3)
    @test_throws AssertionError(_ErrMsg.SANE_PARAMS_A) Collatz.tree_graph(1, 1, P=1, a=0, b=3)
    # If b is a multiple of a, but not of Pa, then 0 can have a reverse.
    @test Collatz.tree_graph(0, 1, P=17, a=2, b=-6) == Dict(0=>Dict(C=>0,3=>D))
    @test Collatz.tree_graph(0, 1, P=17, a=2, b=102) == Dict(0=>Dict(C=>0))
end

println("################################################################################")


# Run tests on Beyns contour integral method

using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using Random
using LinearAlgebra

@testset "Beyn contour" begin

    nep=nep_gallery("dep0")
    @bench @testset "disk at origin" begin
        Random.seed!(0);

        λ,v=contour_beyn(nep,logger=displaylevel,radius=1,neigs=1,sanity_check=false)

        for i = 1:size(λ,1)
            @info "$i: $(λ[i])"
            M=compute_Mder(nep,λ[i])
            minimum(svdvals(M))
            @test minimum(svdvals(M)) < eps()*1000
            @test norm(compute_Mlincomb(nep,λ[i],v[:,i]))/norm(v[:,i]) < eps()*500
        end

        λ,v=contour_beyn(nep,logger=displaylevel,radius=0.4,k=2,N=100,sanity_check=false)
        M=compute_Mder(nep,λ[1])
        minimum(svdvals(M))
        @test minimum(svdvals(M))<eps()*1000

    end
    @bench @testset "shifted disk" begin

        λ,v=contour_beyn(nep,logger=displaylevel,σ=0.2,radius=1.0,
                         neigs=4,sanity_check=false)

        @test size(λ,1)==3
        for i = 1:3
            @info "$i: $(λ[i])"
            M=compute_Mder(nep,λ[i])
            minimum(svdvals(M))
            @test minimum(svdvals(M)) < eps()*10000
            @test norm(compute_Mlincomb(nep,λ[i],v[:,i]))/norm(v[:,i]) < eps()*10000
        end

    end

end


using Gadfly, HilbertSpaceFillingCurve

x,y,z = Int[],Int[], Int[]
for d = 0:1:2^12-1

    p = hilbert(d, 3, 16)
    push!(x,p[1])
    push!(y,p[2])
    push!(z,p[3])
end
#
p = plot(
    x=x+z/40,y=y+z/40,Geom.line(preserve_order=true),
    Guide.xlabel(""),Guide.ylabel(""),
    Coord.cartesian(xmax=15,ymax=15),
    Guide.xticks(ticks=nothing), Guide.yticks(ticks=nothing)
)
Gadfly.draw(PNG(joinpath(Pkg.dir(),"HilbertSpaceFillingCurve","data","figure.png"),12cm,12cm),p)
p


using Distributions
using LinearAlgebra
using Random
using Test

using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions
using EnsembleKalmanProcesses.Localizers
import EnsembleKalmanProcesses: construct_mean, construct_cov, construct_sigma_ensemble
const EKP = EnsembleKalmanProcesses

@testset "Localization" begin

    # Seed for pseudo-random number generator
    rng_seed = 42
    rng = Random.MersenneTwister(rng_seed)

    # Linear problem with d == p >> N_ens - Section 6.1 of Tong and Morzfeld (2022)
    G(u) = u
    N_ens = 10
    p = 50
    N_iter = 20
    # Generate random truth
    y = 10.0 * rand(p)
    Γ = 1.0 * I

    #### Define prior information on parameters
    priors = ParameterDistribution[] #empty PD-array
    for i in 1:p
        push!(priors, ParameterDistribution(Parameterized(Normal(0.0, 0.5)), no_constraint(), string("u", i)))
    end
    prior = combine_distributions(priors)
    #prior_distns = repeat([Parameterized(Normal(0.0, 0.5))], p)
    #constraints = repeat([[no_constraint()]], p)
    #prior_names = [string("u_", i) for i in 1:p]
    #prior = ParameterDistribution(prior_distns, constraints, prior_names)

    initial_ensemble = EKP.construct_initial_ensemble(rng, prior, N_ens)

    # Solve problem without localization
    ekiobj_vanilla = EKP.EnsembleKalmanProcess(initial_ensemble, y, Γ, Inversion(); rng = rng)
    for i in 1:N_iter
        g_ens_vanilla = G(get_u_final(ekiobj_vanilla))
        EKP.update_ensemble!(ekiobj_vanilla, g_ens_vanilla)
    end
    nonlocalized_error = get_error(ekiobj_vanilla)[end]

    # Test different localizers
    loc_methods = [Delta(), RBF(1.0), RBF(0.1), BernoulliDropout(0.1)]

    for loc_method in loc_methods
        ekiobj =
            EKP.EnsembleKalmanProcess(initial_ensemble, y, Γ, Inversion(); rng = rng, localization_method = loc_method)
        @test isa(ekiobj.localizer, Localizer)

        for i in 1:N_iter
            g_ens = G(get_u_final(ekiobj))
            EKP.update_ensemble!(ekiobj, g_ens, deterministic_forward_map = true)
        end

        # Test that localized version does better in the setting p >> N_ens
        @test get_error(ekiobj)[end] < nonlocalized_error

        # Test Schur product theorem
        u_final = get_u_final(ekiobj)
        g_final = get_g_final(ekiobj)
        cov_ug = cov(u_final, g_final, dims = 2, corrected = false)
        kernel = ekiobj.localizer.kernel
        @test rank(cov_ug) < rank(kernel .* cov_ug)
    end

end


using jInvSeismic.BasicFWI
using jInvSeismic.Utils
using jInv.Mesh
using jInv.Utils
using jInv.ForwardShare
using Distributed

using DelimitedFiles

@everywhere begin
	using jInvSeismic.BasicFWI
	using jInv.Mesh
	using jInv.Utils
	using jInv.ForwardShare
end

# setup model and attenuation function
nx = 164; nz = 80
domain = [0.0,10.0,0.0,4.0]
Mr = getRegularMesh(domain,[nx,nz])

# the slowness model
m = 3.0*ones(nx,nz);
m[30:80,20:40] .= 5.0; # velocity...
m = 1.0./(m.^2); # convert to slowness squared

# m = readdlm("SEGmodel2Dsalt.dat");
# m = m*1e-3;
# m = m';
# m = 1.0./m^2

# attenuation for BC
padx = 20; padz = 20
a    = 2.0;
xc = getCellCenteredGrid(Mr)
gamma = getHelmholtzABL(Mr,true,[padx;padz],a);

# parameters for the Helmholtz (units in km)
h = Mr.h;
n = [nx; nz;]
omega = 2*pi*[0.5;1.0;2.0;3.0]
nfreq = length(omega)

# generate sources
q = zeros(tuple(n.+1...)); q[padx+1:4:end-padx-1,1] .= 1e4
Q = sdiag(vec(q))
Q = Q[:,(LinearIndices(sum(Q,dims=2) .!= 0))[findall(sum(Q,dims=2) .!= 0)]]
nsrc = size(Q,2)

# receivers
p = zeros(tuple(n.+1...)); p[padx+1:end-padx-1,1] .= 1
P = sdiag(vec(p))
P = P[:,(LinearIndices(sum(P,dims=2) .!= 0))[findall(sum(P,dims=2) .!= 0)]]
nrec = size(P,2)

pFor  = getBasicFWIparam(omega,gamma,Q,P,Mr)
pForp = getBasicFWIparam(omega,gamma,Q,P,Mr,true)

# inversion mesh and forward mesh are the same here
M2Mp = ones(length(pForp))


using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
using EntropyStableEuler
using EntropyStableEuler.Fluxes1D

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,(rhou,),E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end
unorm(U) = sum(map((x->x.^2),U))
function pfun_nd(rho, rhoU, E)
    rhoUnorm2 = unorm(rhoU)./rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,U,p)
    rhoU = rho.*U
    Unorm = unorm(U)
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rhoU,E)
end

function euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
    rhoR,uR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))

    unorm = (@. uL*uR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. f4aux*uavg)

    return (FxS1,FxS2,FxS3)
end

function euler_fluxes(rhoL,uL,betaL,rhoR,uR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
                           rhoR,uR,betaR,rhologR,betalogR)
end

function dopri45_coeffs()
    rk4a = [0.0             0.0             0.0             0.0             0.0             0.0         0.0
            0.2             0.0             0.0             0.0             0.0             0.0         0.0
            3.0/40.0        9.0/40.0        0.0             0.0             0.0             0.0         0.0
            44.0/45.0      -56.0/15.0       32.0/9.0        0.0             0.0             0.0         0.0
            19372.0/6561.0 -25360.0/2187.0  64448.0/6561.0  -212.0/729.0    0.0             0.0         0.0
            9017.0/3168.0  -355.0/33.0      46732.0/5247.0  49.0/176.0      -5103.0/18656.0 0.0         0.0
            35.0/384.0      0.0             500.0/1113.0    125.0/192.0     -2187.0/6784.0  11.0/84.0   0.0 ]

    rk4c = vec([0.0 0.2 0.3 0.8 8.0/9.0 1.0 1.0 ])

    # coefficients to evolve error estimator = b1-b2
    rk4E = vec([71.0/57600.0  0.0 -71.0/16695.0 71.0/1920.0 -17253.0/339200.0 22.0/525.0 -1.0/40.0 ])

    return rk4a,rk4E,rk4c
end

const TOL = 1e-16
"Approximation parameters"
N = 4 # The order of approximation
K = 2^7
T = 0.1

# Becker viscous shocktube
const γ = 1.4
const M_0 = 3.0
const mu = 0.1
const lambda = 2/3*mu
const Pr = 3/4
const cp = γ/(γ-1)
const cv = 1/(γ-1)
const kappa = mu*cp/Pr

const v_inf = 0.2
const rho_0 = 1.0
const v_0 = 1.0
const m_0 = rho_0*v_0
const v_1 = (γ-1+2/M_0^2)/(γ+1)
const v_01 = sqrt(v_0*v_1)

# const uL = v_0+v_inf
# const uR = v_1+v_inf
# const rhoL = m_0/v_0
# const rhoR = m_0/v_1
# const eL = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_0^2)
# const eR = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_1^2)
# const pL = (γ-1)*rhoL*eL
# const pR = (γ-1)*rhoR*eR
# const EL = pL/(γ-1)+0.5*rhoL*uL^2
# const ER = pR/(γ-1)+0.5*rhoR*uR^2

const XL = -2.0#-1.0
const XR = 2.0
gr(size=(300,300),ylims=(0,5.0),legend=false,markerstrokewidth=1,markersize=2)
plot()

"Mesh related variables"
VX = LinRange(XL,XR,K+1)
EToV = transpose(reshape(sort([1:K; 2:K+1]),2,K))
rd = init_reference_interval(N)

# High order mesh
@unpack r,M,Dr,Vq,Pq,Vf,LIFT,wq,VDM = rd
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1])
x = V1*VX[transpose(EToV)]
xf = Vf*x
mapM = reshape(1:2*K,2,K)
mapP = copy(mapM)
mapP[1,2:end] .= mapM[2,1:end-1]
mapP[2,1:end-1] .= mapM[1,2:end]

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nrJ = [-1;1]
wf  = [1;1]
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

# construct hybridized SBP operators
Qr = Pq'*M*Dr*Pq
Ef = Vf*Pq
Br = diagm(wf.*nrJ)
Qrh = .5*[Qr-Qr' Ef'*Br;
         -Br*Ef  Br]

Vh = [Vq;Vf]
Ph = M\transpose(Vh)
VhP = Vh*Pq

# make sparse skew symmetric versions of the operators"
# precompute union of sparse ids for Qr, Qs
Qrh_skew = .5*(Qrh-transpose(Qrh))

"""Initial condition"""
function bisection_solve_velocity(x,max_iter,tol)
    v_L = v_1
    v_R = v_0
    num_iter = 0

    L_k = kappa/m_0/cv
    f(v) = -x + 2*L_k/(γ+1)*log((v_0-v)^((v_0/(v_0-v_1)))*(v-v_1)^(-(v_1)/(v_0-v_1)))

    v_new = (v_L+v_R)/2
    while num_iter < max_iter
        v_new = (v_L+v_R)/2

        if abs(f(v_new)) < tol
            return v_new
        elseif sign(f(v_L)) == sign(f(v_new))
            v_L = v_new
        else
            v_R = v_new
        end
        num_iter += 1
    end

    return v_new
end

const max_iter = 10000
const tol = 1e-16

function exact_sol_viscous_shocktube(x,t)
    u   = bisection_solve_velocity(x-v_inf*t,max_iter,tol)
    rho = m_0/u
    e   = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-u^2)
    return rho, rho*(v_inf+u), rho*(e+1/2*(v_inf+u)^2)
end

U = exact_sol_viscous_shocktube.(x,0.0)
U = ([x[1] for x in U], [x[2] for x in U], [x[3] for x in U])

function flux_differencing!(UF,Uh,Qrh_skew,rxJ,Nh,Nq,Nc,K)
    for k = 1:K
        for j = 1:Nh # col idx
            for i = j:Nh # row idx
                if i <= Nq || j <= Nq # Skip lower right block
                    Fx = euler_fluxes(Uh[1][i,k],Uh[2][i,k],Uh[3][i,k],Uh[1][j,k],Uh[2][j,k],Uh[3][j,k])
                    for c = 1:Nc
                        update_val = 2*rxJ*Qrh_skew[i,j]*Fx[c]
                        UF[c][i,k] += update_val
                        UF[c][j,k] -= update_val
                    end
                end
            end
        end
    end
end

function rhs_inviscid(U,K,VhP,Vq,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,rhoL,uL,betaL,rhoR,uR,betaR)
    J = (XR-XL)/K/2 # assume uniform interval
    Nc = 3
    Nh,Nq = size(VhP)

    # Preallocate flux differencing array
    UF = [zeros(Float64,Nh,K) for i = 1:Nc]

    # entropy var projection
    VU = v_ufun((x->Vq*x).(U)...)
    VU = (x->VhP*x).(VU)
    Uh = u_vfun(VU...)

    # convert to rho,u,v,beta vars
    (rho,rhou,E) = Uh
    beta = betafun(rho,rhou,E)
    Ubh = (rho, rhou./rho, beta) # redefine Q = (rho,u,v,β)

    # compute face values
    UM = (x->x[Nq+1:Nh,:]).(Ubh)
    UP = (x->x[mapP]).(UM)
    UP[1][1] = rhoL
    UP[2][1] = uL
    UP[3][1] = betaL
    UP[1][end] = rhoR
    UP[2][end] = uR
    UP[3][end] = betaR

    # simple lax friedrichs dissipation
    Uf =  (x->x[Nq+1:Nh,:]).(Uh)
    (rhoM,rhouM,EM) = Uf
    rhoUM_n = @. (rhouM*nxJ)
    lam = abs.(wavespeed_1D(rhoM,rhoUM_n,EM))
    LFc = .25*max.(lam,lam[mapP])

    fSx = euler_fluxes(UM[1],UM[2],UM[3],UP[1],UP[2],UP[3])
    normal_flux(fx,u) = fx.*nxJ - LFc.*(u[mapP]-u)
    flux = normal_flux.(fSx,Uf)
    rhsU = (x->LIFT*x).(flux)

    flux_differencing!(UF,Ubh,Qrh_skew,rxJ,Nh,Nq,Nc,K)
    rhsU = (x->Ph*x).(UF) .+ rhsU
    rhsU = (x->-x./J).(rhsU)

    return rhsU
end

function rhs_viscous(U,K,N,Vq,Vf,Pq,Dr,LIFT,mapP,nxJ,v1L,v2L,v3L,v1R,v2R,v3R)
    J = (XR-XL)/K/2 # assume uniform interval
    Nc = 3

    # entropy var projection
    VU = v_ufun((x->Vq*x).(U)...)
    VU = (x->Pq*x).(VU)

    # compute and interpolate to quadrature
    VUf = (x->Vf*x).(VU)
    VUP = (x->x[mapP]).(VUf)
    VUP[1][1] = v1L
    VUP[2][1] = v2L
    VUP[3][1] = v3L
    VUP[1][end] = v1R
    VUP[2][end] = v2R
    VUP[3][end] = v3R

    surfx(uP,uf) = LIFT*(@. .5*(uP-uf)*nxJ)
    # Strong form, dv/dx
    VUx = (x->Dr*x).(VU)
    VUx = VUx .+ surfx.(VUP,VUf)
    VUx = VUx./J

    VUx = (x->Vq*x).(VUx)
    VUq = (x->Vq*x).(VU)

    # σ = K dv/dx
    sigma = zero.(VUx)
    for k = 1:K
        for i = 1:size(Vq,1)
            Kx = zeros(3,3)
            v1 = VUq[1][i,k]
            v2 = VUq[2][i,k]
            v4 = VUq[3][i,k]
            Kx[2,2] = -(2*mu-lambda)/v4
            Kx[2,3] = (2*mu-lambda)*v2/v4^2
            Kx[3,2] = Kx[2,3]
            Kx[3,3] = -(2*mu-lambda)*v2^2/v4^3+kappa/cv/v4^2

            sigma[2][i,k] += Kx[2,2]*VUx[2][i,k] + Kx[2,3]*VUx[3][i,k]
            sigma[3][i,k] += Kx[3,2]*VUx[2][i,k] + Kx[3,3]*VUx[3][i,k]
        end
    end
    sigma = (x->Pq*x).(sigma)

    sxf = (x->Vf*x).(sigma)
    sxP = (x->x[mapP]).(sxf)
    sxP[1][1] = sxf[1][1]
    sxP[2][1] = sxf[2][1]
    sxP[3][1] = sxf[3][1]
    sxP[1][end] = sxf[1][end]
    sxP[2][end] = sxf[2][end]
    sxP[3][end] = sxf[3][end]

    # strong form, dσ/dx
    penalization(uP,uf) = LIFT*(@. -.5*(uP-uf))
    sigmax = (x->Dr*x).(sigma)
    sigmax = sigmax .+ surfx.(sxP,sxf)
    sigmax = sigmax./J

    return sigmax
end

function rhs_ESDG(U,K,VhP,Vq,Vf,Pq,Dr,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,t)

    rhoL,rhouL,EL = exact_sol_viscous_shocktube(XL,t)
    rhoR,rhouR,ER = exact_sol_viscous_shocktube(XR,t)
    uL = rhouL/rhoL
    uR = rhouR/rhoR
    pL = (γ-1)*(EL-.5*rhoL*uL^2)
    pR = (γ-1)*(ER-.5*rhoR*uR^2)
    betaL = rhoL/(2*pL)
    betaR = rhoR/(2*pR)
    v1L,v2L,v3L = v_ufun(rhoL,rhoL*uL,EL)
    v1R,v2R,v3R = v_ufun(rhoR,rhoR*uR,ER)

    rhsI = rhs_inviscid(U,K,VhP,Vq,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,rhoL,uL,betaL,rhoR,uR,betaR)
    rhsV = rhs_viscous(U,K,N,Vq,Vf,Pq,Dr,LIFT,mapP,nxJ,v1L,v2L,v3L,v1R,v2R,v3R)
    return rhsI .+ rhsV
end

# Time stepping
"Time integration"
t = 0.0
U = collect(U)
resU = [zeros(size(x)),zeros(size(x)),zeros(size(x))]
resW = [zeros(size(x)),zeros(size(x)),zeros(size(x))]
resZ = [zeros(size(x)),zeros(size(x)),zeros(size(x))]

Vp = vandermonde_1D(N,LinRange(-1,1,10))/VDM

anim = Animation()
const ptL = XL+(XR-XL)/K/(N+1)/2
const ptR = XR-(XR-XL)/K/(N+1)/2
const hplot = (XR-XL)/K/(N+1)
i = 1

while t < T

    # SSPRK(3,3)
    # dt = min(1e-4,T-t)
    CN = (N+1)*(N+2)/2  # estimated trace constant
    dt = min(2/CN/K^2,T-t)
    rhsU = rhs_ESDG(U,K,VhP,Vq,Vf,Pq,Dr,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,t)
    @. resW = U + dt*rhsU
    rhsU = rhs_ESDG(resW,K,VhP,Vq,Vf,Pq,Dr,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,t)
    @. resZ = resW+dt*rhsU
    @. resW = 3/4*U+1/4*resZ
    rhsU = rhs_ESDG(resW,K,VhP,Vq,Vf,Pq,Dr,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,t)
    @. resZ = resW+dt*rhsU
    @. U = 1/3*U+2/3*resZ

    global t = t + dt
    global i = i + 1
    println("Current time $t with time step size $dt, and final time $T")
    if mod(i,1000) == 1
        plot(x,U[1])
        # for k = 1:K
        #     plot!(ptL+(k-1)*hplot*(N+1):hplot:ptL+k*hplot*(N+1), 1 .-L_plot[:,k],st=:bar,alpha=0.2)
        # end
        frame(anim)
    end

    # if i == 10
    #     break
    # end
end

# rka,rkE,rkc = dopri45_coeffs()
# bcopy!(x,y) = x .= y

# # DOPRI storage
# Utmp = similar.(U)
# rhsUrk = ntuple(x->zero.(U),length(rkE))

# errEst = 0.0
# prevErrEst = 0.0

# t = 0.0
# CN = (N+1)*(N+2)/2  # estimated trace constant
# dt = min(2/CN/K^2,T-t)
# dt0 = dt
# i = 0
# interval = 5

# dthist = Float64[dt]
# thist = Float64[0.0]
# errhist = Float64[0.0]
# wsJ = diagm(wf)

# rhsU = rhs_ESDG(U,K,VhP,Vq,Vf,Pq,Dr,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,t)
# bcopy!.(rhsUrk[1],rhsU) # initialize DOPRI rhs (FSAL property)

# while t < T
#     # DOPRI step and
#     rhstest = 0.0
#     for INTRK = 2:7
#         k = zero.(Utmp)
#         for s = 1:INTRK-1
#             bcopy!.(k, @. k + rka[INTRK,s]*rhsUrk[s])
#         end
#         bcopy!.(Utmp, @. U + dt*k)
#         rhsU = rhs_ESDG(Utmp,K,VhP,Vq,Vf,Pq,Dr,LIFT,Ph,Qrh_skew,rxJ,mapP,nxJ,t)#rhsRK(Qtmp,rd,md,ops,euler_fluxes)
#         bcopy!.(rhsUrk[INTRK],rhsU)
#     end
#     errEstVec = zero.(Utmp)
#     for s = 1:7
#         bcopy!.(errEstVec, @. errEstVec + rkE[s]*rhsUrk[s])
#     end

#     errTol = 1e-5
#     errEst = 0.0
#     for field = 1:length(Utmp)
#         errEstScale = @. abs(errEstVec[field]) / (errTol*(1+abs(U[field])))
#         errEst += sum(errEstScale.^2) # hairer seminorm
#     end
#     errEst = sqrt(errEst/(length(U[1])*4))
#     if errEst < 1.0 # if err small, accept step and update
#             bcopy!.(U, Utmp)
#             global t += dt
#             bcopy!.(rhsUrk[1], rhsUrk[7]) # use FSAL property
#     end
#     order = 5
#     dtnew = .8*dt*(.9/errEst)^(.4/(order+1)) # P controller
#     if i > 0 # use PI controller if prevErrEst available
#             dtnew *= (prevErrEst/max(1e-14,errEst))^(.3/(order+1))
#     end
#     global dt = max(min(10*dt0,dtnew),1e-9) # max/min dt
#     global prevErrEst = errEst

#     push!(dthist,dt)
#     push!(thist,t)

#     global i = i + 1  # number of total steps attempted
#     if i%interval==0
#         println("i = $i, t = $t, dt = $dtnew, errEst = $errEst, rhstest = $rhstest")
#     end
# end

# plot(x,U[1])
# gif(anim,"~/Desktop/tmp.gif",fps=15)

xq = Vq*x
exact_U = @. exact_sol_viscous_shocktube.(xq,T)
exact_rho = [x[1] for x in exact_U]
exact_rhou = [x[2] for x in exact_U]
exact_u = exact_rhou./exact_rho
exact_E = [x[3] for x in exact_U]

rho = Vq*U[1]
u = Vq*(U[2]./U[1])
rhou = Vq*U[2]
E = Vq*U[3]
p = Vq*pfun_nd.(U[1],U[2],U[3])
J = (XR-XL)/K/2

# Linferr = maximum(abs.(exact_rho-rho))/maximum(abs.(exact_rho)) +
#           maximum(abs.(exact_u-u))/maximum(abs.(exact_u)) +
#           maximum(abs.(exact_E-E))/maximum(abs.(exact_E))

# L1err = sum(J*wq.*abs.(exact_rho-rho))/sum(J*wq.*abs.(rho)) +
#         sum(J*wq.*abs.(exact_u-u))/sum(J*wq.*abs.(u)) +
#         sum(J*wq.*abs.(exact_E-E))/sum(J*wq.*abs.(E))

Linferr = maximum(abs.(exact_rho-rho))/maximum(abs.(exact_rho)) +
          maximum(abs.(exact_rhou-rhou))/maximum(abs.(exact_rhou)) +
          maximum(abs.(exact_E-E))/maximum(abs.(exact_E))

L1err = sum(J*wq.*abs.(exact_rho-rho))/sum(J*wq.*abs.(rho)) +
        sum(J*wq.*abs.(exact_rhou-rhou))/sum(J*wq.*abs.(rhou)) +
        sum(J*wq.*abs.(exact_E-E))/sum(J*wq.*abs.(E))

L2err = sqrt(sum(J*wq.*abs.(exact_rho-rho).^2))/sqrt(sum(J*wq.*abs.(rho).^2)) +
        sqrt(sum(J*wq.*abs.(exact_rhou-rhou).^2))/sqrt(sum(J*wq.*abs.(rhou).^2)) +
        sqrt(sum(J*wq.*abs.(exact_E-E).^2))/sqrt(sum(J*wq.*abs.(E).^2))

println("N = $N, K = $K")
println("L1 error is $L1err")
println("L2 error is $L2err")
println("Linf error is $Linferr")


using Gnuplot
# archimedean spiral
a, m, z₀ = 1, 2.1, 0
φ = LinRange(0,20π,500)
r = a*φ
x, y, z = r .* cos.(φ), r .* sin.(φ), m .* r .+ z₀
function theme(bgcolor = "#212946"; borderc ="white",
        ticsc = "white", ylabc = "white", gridst = 4,
        xlabc = "white", gridc = "#2A3459", labc = "white")
    bgcp1 = "set object rectangle from screen 0,0 to screen 1,1"
    bgcp2 = " behind fc '$(bgcolor)' fs solid noborder"
    bgcolor = bgcp1*bgcp2
    @gp :- "set border lw 1 lc '$(borderc)'"
    @gp :- "set tics textcolor rgb '$(ticsc)'"
    @gp :- bgcolor
    @gp :- "set ylabel textcolor '$(ylabc)'"
    @gp :- "set xlabel textcolor '$(xlabc)'"
    @gp :- "set grid ls 1 lc '$(gridc)' dt $(gridst)"
    @gp :- "set key t r textcolor '$(labc)'"
end

@gsp "set xyplane 0" "set view equal xyz" "set auto fix" :-
theme("black")
@gsp :- xlab = "x" ylab = "y"
@gsp :- xr=extrema(x) yr=extrema(y) zr=extrema(z)
@gsp :- "i=2"
# line and point
@gsp :- x y z "every ::1::i w l notit lc palette" palette(:Wistia)
@gsp :- x y z "every ::i::i w points notit lc 'red'"
@gsp :- x y 0*z z "every ::1::i w l notit lc palette lt 2"
# animation
for frame in 2:5:length(x)
  @gsp :- frame "i=$frame" "replot" :-
end
@gsp # show animation
save(term="gif animate size 900,600 delay 1", output="archimedean.gif") # hide


@inline iszero{T<:Real}(x::T) = (x == zero(T))
@inline isone{T<:Real}(x::T) = (x == one(T))

@inline function isfinite{T<:Rvl}(ivl::T)
    isfinite(ivl.lo) & isfinite(ivl.hi)
end
@inline function issubnormal{T<:Rvl}(ivl::T)
    issubnormal(ivl.lo) | issubnormal(ivl.hi)
end
@inline function isnormal{T<:Rvl}(ivl::T)
    isfinite(ivl) & !issubnormal(ivl)
end
@inline function isinf{T<:Rvl}(ivl::T)
    isinf(ivl.lo) | isinf(ivl.hi)
end
@inline function isnan{T<:Rvl}(ivl::T)
    isnan(ivl.lo) | isnan(ivl.hi)
end

@inline function isPoint{T<:Rvl}(ivl::T)
    (ivl.lo == ivl.hi) && isfinite(ivl.lo)
end
@inline function isNotPoint{T<:Rvl}(ivl::T)
    (ivl.lo != ivl.hi) || !isfinite(ivl.lo)
end
@inline function isBoundry{T<:Rvl}(ivl::T)
    (ivl.lo != ivl.hi) && isnan(ivl)
end
@inline function isNotBoundry{T<:Rvl}(ivl::T)
    (ivl.lo == ivl.hi) || !isnan(ivl)
end
@inline function isExact{T<:Rvl}(ivl::T)
    (isExact(ivl.lo) & isExact(ivl.hi)) && (isfinite(ivl.lo) & isfinite(ivl.hi))
end
@inline function isNotExact{T<:Rvl}(ivl::T)
    (isNotExact(ivl.lo) | isNotExact(ivl.hi)) || (!isfinite(ivl.lo) | !isfinite(ivl.hi))
end

@inline function isExactPoint{T<:Rvl}(ivl::T)
    isPoint(ivl) && isExact(ivl)
end
@inline function isNotExactPoint{T<:Rvl}(ivl::T)
    isNotPoint(ivl) || isNotExact(ivl)
end
@inline function isInexactPoint{T<:Rvl}(ivl::T)
    isPoint(ivl) && isNotExact(ivl)
end
@inline function isNotInexactPoint{T<:Rvl}(ivl::T)
    isNotPoint(ivl) || isExact(ivl)
end

@inline function isExactBoundry{T<:Rvl}(ivl::T)
    isBoundry(ivl) && isExact(ivl)
end
@inline function isNotExactBoundry{T<:Rvl}(ivl::T)
    isNotBoundry(ivl) || isNotExact(ivl)
end
@inline function isInexactBoundry{T<:Rvl}(ivl::T)
    isBoundry(ivl) && isNotExact(ivl)
end
@inline function isNotInexactBoundry{T<:Rvl}(ivl::T)
    isNotBoundry(ivl) || isExact(ivl)
end

@inline function isClosedOpen{T<:Rvl}(ivl::T)
    isExact(ivl.lo) & isInexact(ivl.hi)
end
@inline function isOpenClosed{T<:Rvl}(ivl::T)
    isInexact(ivl.lo) & isExact(ivl.hi)
end
@inline function isOpenOpen{T<:Rvl}(ivl::T)
    isInexact(ivl.lo) & isInexact(ivl.hi)
end
@inline function isClosedClosed{T<:Rvl}(ivl::T)
    isExact(ivl.lo) & isExact(ivl.hi)
end


pairs(n::Int64) = ((i, j) for i = 1:n for j = 1:i)
rpairs(n::Int64) = ((i,j) for i in 1:n for j in 1:n) # rectangular option to old pairs
spairs(n::Int64) = ((i, j) for i = 1:n for j = 1:(i-1)) # subdiagonal option to old pairs

triangle(i::Int64) = div(i*(i+1),2)
triangle(i::Int64,j::Int64) = i<j ? triangle(j-1)+i : triangle(i-1)+j

iiterator(n::Int64) = ((i,j,k,l) for (i,j) in pairs(n) for (k,l) in pairs(n) if triangle(i,j) <= triangle(k,l))

iindex(i::Int64,j::Int64,k::Int64,l::Int64) = triangle(triangle(i,j),triangle(k,l))
trace2(A,B) = sum(A.*B)

#=
function all_1e_ints( bfs::BasisSet, atoms::Atoms )
    n = length(bfs)
    S = Array{Float64}(undef,n,n)
    T = Array{Float64}(undef,n,n)
    V = Array{Float64}(undef,n,n)
    for j = 1:n
        for i = j:n
            S[i,j] = overlap(bfs[i], bfs[j])
            S[j,i] = S[i,j]
            T[i,j] = kinetic(bfs[i], bfs[j])
            T[j,i] = T[i,j]
            V[i,j] = nuclear_attraction(bfs[i], bfs[j], atoms)
            V[j,i] = V[i,j]
        end
    end
    return S,T,V
end
=#

function all_1e_ints(bfs::BasisSet, atoms::Atoms)
    n = length(bfs)
    S = Array{Float64}(undef,n,n)
    T = Array{Float64}(undef,n,n)
    V = Array{Float64}(undef,n,n)
    for (i,j) in pairs(n)
        a,b = bfs[i], bfs[j]
        S[i,j] = S[j,i] = overlap(a,b)
        T[i,j] = T[j,i] = kinetic(a,b)
        V[i,j] = V[j,i] = nuclear_attraction(a,b,atoms)
    end
    return S,T,V
end

#=
function all_twoe_ints(bfs,ERI=coulomb)
    n = length(bfs)
    totlen = div(n*(n+1)*(n*n+n+2),8)
    ints2e = zeros(totlen)  # using zeros instead of using Array constructor
    for i = 1:n
    for j = 1:n
        ij = triangle(i-1,j)
        for k = 1:n
        for l = 1:n
            kl = triangle(k-1,l)
            if ij <= kl
                ints2e[iindex(i,j,k,l)] = ERI(bfs[i], bfs[j], bfs[k], bfs[l])
            end
        end
        end
    end
    end
    return ints2e
end
=#

function all_twoe_ints(bfs,ERI=coulomb)
    n = length(bfs)
    totlen = div(n*(n+1)*(n*n+n+2),8)
    ints2e = Array{Float64}(undef,totlen)
    for (i,j,k,l) in iiterator(n)
        ints2e[iindex(i,j,k,l)] = ERI(bfs[i], bfs[j], bfs[k], bfs[l])
    end
    return ints2e
end

#=
function make2JmK(D::Array{Float64,2},Ints::Array{Float64,1})
    n = size(D,1)
    G = Array{Float64}(undef,n,n)
    D1 = reshape(D,n*n)
    temp = Array{Float64}(undef,n*n)
    for j = 1:n
    for i = j:n
        kl = 1
        for k = 1:n
            start = 1
            for l = start:n
                temp[kl] = 2*Ints[iindex(i,j,k,l)]-Ints[iindex(i,k,j,l)]
                kl += 1
            end
        end
        G[i,j] = dot(D1,temp)
        G[j,i] = G[i,j]
    end
    end
    return G
end
=#

function make2JmK(D::Array{Float64,2},Ints::Array{Float64,1})
    n = size(D,1)
    G = Array{Float64}(undef,n,n)
    D1 = reshape(D,n*n)
    temp = Array{Float64}(undef,n*n)
    for (i,j) in pairs(n)
        kl = 1
        for (k,l) in rpairs(n)
            temp[kl] = 2*Ints[iindex(i,j,k,l)]-Ints[iindex(i,k,j,l)]
            kl += 1
        end
        G[i,j] = G[j,i] = dot(D1,temp)
    end
    return G
end

dmat(U::Array{Float64,2},nocc::Int64) = U[:,1:nocc]*U[:,1:nocc]'

function RHF( atoms::Atoms, MaxIter::Int64=20, verbose::Bool=true)
    @printf("Starting RHF calculation:\n")

    Nelectrons = get_Nelectrons(atoms)
    nclosed,nopen = divrem( Int64(Nelectrons), 2 )
    @printf("Nel = %d, Nclosed = %d, nopen = %d\n", Nelectrons, nclosed, nopen)

    @printf("Building basis ...\n")
    bfs = build_basis_new(atoms)

    @printf("1e integrals ...\n")
    S,T,V = all_1e_ints(bfs,atoms)

    @printf("2e integrals ...\n")
    Ints = all_twoe_ints(bfs)

    h = T+V
    E,U = eigen(h,S)
    println("E = ", E)
    Enuke = nuclear_repulsion(atoms)
    Eold = 0.0
    Energy = 0.0

    IS_CONVERGED = false
    dEtot = 1.0
    for iter in 1:MaxIter
        D = dmat(U,nclosed)
        G = make2JmK(D,Ints)
        H = h + G
        E,U = eigen(H,S)
        Eone = trace2(D,h)
        Etwo = trace2(D,H)
        Energy = Enuke + Eone + Etwo
        dEtot = abs(Energy - Eold)
        #println("HF: $iter  $Energy : $Enuke  $Eone  $Etwo $dEtot")
        @printf("HF: %4d %18.10f %18.10f %18.10f: %18.10f %18.10e\n",
                iter, Enuke, Eone, Etwo, Energy, dEtot)
        if isapprox(Energy,Eold)
            IS_CONVERGED = true
            break
        end
        Eold  = Energy
    end

    if !IS_CONVERGED
        println("WARNING: RHF is not converged")
    end

    return Energy,E,U
end


@testset "add" begin
    x1 = [1,2,3]
    x2 = Var([1,2,3])
    y = x1 + x2
    @test y.data == [2,4,6]

    x = randn(3, 3)
    y = randn(3, 3)
    f = x -> x + y
    @test gradcheck(f, x)

    x = Var(randn(3, 3))
    y = randn(3)
    @test gradcheck(f, x)
end

@testset "mul" begin
    x1 = [1, 2, 3]
    x2 = Var([1, 2, 3])
    y = x1 .* x2
    @test y.data == [1, 4, 9]

    x = randn(3, 3)
    y = randn(3, 3)
    f = x -> x .* y
    @test gradcheck(f, x)

    y = randn(3, 1)
    @test gradcheck(f, x)

    y = randn(3)
    @test gradcheck(f, x)
    f = y -> x .* y
    @test gradcheck(f, x)
end

@testset "div" begin
    x1 = [1, 2, 3]
    x2 = Var([1, 2, 3])
    y = x1 / x2
    @test y.data == [1, 1, 1]

    x = randn(3, 3)
    y = randn(3, 3)
    f = x -> x ./ y
    @test gradcheck(f, x)
    y = randn(3)
    @test gradcheck(f, x)
end

@testset "pow" begin
    x = Var(2)
    @test (x^10).data == [1024]
    y = x^2
    gradient!(y)
    @test x.grad.data == [4.0]
end


using OrdinaryDiffEq, Zygote
using DiffEqSensitivity, Test, ForwardDiff

abstol = 1e-12
reltol = 1e-12
savingtimes = 0.5

# see https://diffeq.sciml.ai/stable/features/callback_functions/#VectorContinuousCallback-Example
function test_vector_continuous_callback(cb,g)

  function f(du, u, p, t)
    du[1] = u[2]
    du[2] = -p[1]
    du[3] = u[4]
    du[4] = 0.0
  end

  u0 = [50.0, 0.0, 0.0, 2.0]
  tspan = (0.0, 10.0)
  p = [9.8, 0.9]
  prob = ODEProblem(f,u0,tspan,p)
  sol = solve(prob, Tsit5(), callback=cb, abstol=abstol, reltol=reltol, saveat=savingtimes)

  du01, dp1 = @time Zygote.gradient(
      (u0, p) -> g(solve(prob, Tsit5(), u0=u0, p=p, callback=cb, abstol=abstol, reltol=reltol, saveat=savingtimes, sensealg=BacksolveAdjoint())),
      u0, p)

  dstuff = @time ForwardDiff.gradient(
      (θ) -> g(solve(prob, Tsit5(), u0=θ[1:4], p=θ[5:6], callback=cb, abstol=abstol, reltol=reltol, saveat=savingtimes)),
      [u0; p])

  @test du01 ≈ dstuff[1:4]
  @test dp1 ≈ dstuff[5:6]
end

@testset "VectorContinuous callbacks" begin
  @testset "MSE loss function bouncing-ball like" begin
    g(u) = sum((1.0.-u).^2)./2
    function condition(out, u, t, integrator) # Event when event_f(u,t) == 0
      out[1] = u[1]
      out[2] = (u[3]-10.0)u[3]
    end
    @testset "callback with linear affect" begin
      function affect!(integrator, idx)
        if idx == 1
            integrator.u[2] = -integrator.p[2]*integrator.u[2]
        elseif idx == 2
            integrator.u[4] = -integrator.p[2]*integrator.u[4]
        end
      end
      cb = VectorContinuousCallback(condition, affect!, 2)
      test_vector_continuous_callback(cb, g)
    end
  end
end


interc = readdlm("/home/alkan/c/nese/jtb/Adj_matrix_yeast.dat",' ',Int64);

n = length(interc[:,1])                 # number of nodes
selfact = diag(interc)                  # yeast specific
interc -= diagm(selfact)                # interaction matrix

# --- generate the symmetric adjacency matrix with elements 0,1 ---
function symadj(m)
    n = length(m[:,1])
    ad = zeros(Int64,n,n)
    for (i=1:n) for (j=i+1:n)       # adjacency matrix - without diagonals
        ad[i,j] = m[i,j]==0 && m[j,i]==0 ? 0 : 1
        ad[j,i] = ad[i,j]
    end end
    return ad
end

## ad = symadj(interc)

## --- check connectivity -----------
function isconnected (adj)
    n = length(adj[:,1])                # number of nodes
    nsq = n*n
    dum = adj
    (sumprev,sumdum) = (0,0)
    for (i=1:n)
        dum = sign(dum + dum*adj)
        sumdum = sum(dum)
        if (sumdum == sumprev)
            break
        end
        sumprev = sumdum
    end
    return (sumdum==nsq)
end

## --- node update function (integrator) -----------
function upd(s,intr,slf,n)
    sss = similar(s)
    inp = intr*s
    for (i=1:n)
        sss[i] = (inp[i]>0)?1:(inp[i]<0)?0:(slf[i]==0)?0:s[i]
    end
    return sss
end

## --- check fixed point ---
function isfixed (s,intr,slf,n)
    snew = upd(s,intr,slf,n)
    return (snew==s)
end

## ---- find point attractors of a small network ---
function getpointattr(intr,slf)
    n = length(slf)
    attrlist = zeros(Int64,n)                 # s=0 is always an attractor
    for (i=1:2^n-1)                           # visit each state except s=0
        s = digits(parseint(bits(i)))
        lens = length(s)
        if (lens < n)
            s = [s;zeros(Int64,n-lens)]
        end
        if (isfixed(s,intr,slf,n))
            attrlist = [attrlist s]
        end
    end
    return attrlist
end

## ---- find ALL attractors of a small network ---
## ---- by exhaustive search ---
## ---- output is the attractor list + periods ---
function getcycles(intr,slf)
    n = length(slf)
    nm1 = n-1
    attrlist = zeros(Int64,n)           # s=0 is always an attractor
    keyzeros = join(zeros(Int64,n))
    visdict = [keyzeros=>keyzeros] # add the zero fixed point to visited list
    attdict = [keyzeros=>1]    # values are the basin sizes
    for (i=1:2^n-1)            # visit each state except s=0
        skey = bits(i)[end-nm1:end]
        s = [int(skey[x])-48 for x = 1:n]
        if (!haskey(visdict,skey))      # if not seen this state before
            pathdict = Dict()
            while (!haskey(visdict,skey))   # generate a path starting here
                pathdict[skey] = "-"
                visdict[skey] = "-"
                s = upd(s,intr,slf,n)
                skey=join(s)         # join(s) == bits(i)[end-nm1:end]
            end
            attkey=visdict[skey]
            if (haskey(attdict,attkey)) # reached a known cycle or fp
                for (k in keys(pathdict))
                    visdict[k]=attkey
                    attdict[attkey] += length(keys(pathdict)) # basin
                end
            else                        # this is a new cycle of fp
                for (k in keys(pathdict)) visdict[k] = skey end
                attdict[skey]=length(keys(pathdict)) # basin
                attrlist = [attrlist s]
            end
        end
    end
    basinlist = [attdict[join(attrlist[:,i])]/(2^n-1) for i in 1:length(attrlist[1,:])]
    return attrlist,basinlist
end

## ---- output is the attractor list + basin sizes ---
## ---- calculated from a subset of initial conditions ---
function getcyclessample(intr,slf,ninit)

    n = length(slf)
    nm1 = n-1
    attrlist = zeros(Int64,n)           # s=0 is always an attractor
    keyzeros = join(zeros(Int64,n))
    attdict = [keyzeros=>0]    # values are the basin sizes
    cyclen = [keyzeros=>1]    # values are the cycle lengths
    for (i=1:ninit)       # calculate basins using ninit initial states
        s = rand(0:1,n)
        skey = join(s)
        pathdict = Dict()
        pathlen = 1
        while (!haskey(attdict,skey) && !haskey(pathdict,skey))
            pathdict[skey] = pathlen
            pathlen += 1
            s = upd(s,intr,slf,n)
            skey=join(s)         # join(s) == bits(i)[end-nm1:end]
        end
        if (haskey(attdict,skey)) attdict[skey] += 1
        else                        # this is a new cycle of fp
            attdict[skey] = 1
            cyclen[skey] = pathlen - pathdict[skey]
            attrlist = [attrlist s]
        end
    end
    basinlist = [attdict[join(attrlist[:,i])]/ninit for i in 1:length(attrlist[1,:])]
    cyclelengthlist = [cyclen[join(attrlist[:,i])] for i in 1:length(attrlist[1,:])]
    return attrlist,basinlist,cyclelengthlist
end