cd("/Volumes/SSD Hans/Github/MacroFall2020/AdvMacro/Assignment5")
# mkpath("Figures")
using Plots
# using LateXStrings # Pkg.add("LaTeXStrings") # https://github.com/stevengj/LaTeXStrings.jl
using Dierckx # Pkg.add("Dierckx") # https://github.com/kbarbary/Dierckx.jl
using Interpolations # Pkg.add("Interpolations") # https://github.com/JuliaMath/Interpolations.jl
using ForwardDiff # Pkg.add("ForwardDiff") # https://github.com/JuliaDiff/ForwardDiff.jl
using Optim # Pkg.add("Optim") # https://julianlsolvers.github.io/Optim.jl/stable/
    using Optim: converged, maximum, maximizer, minimizer, iterations
using Roots # Pkg.add("Roots") # https://github.com/JuliaMath/Roots.jl
using Parameters # Pkg.add("Parameters") # https://github.com/mauro3/Parameters.jl
using Distributions #Pkg.add("Distributions")
using QuadGK # Pkg.add("QuadGK") # https://juliamath.github.io/QuadGK.jl/latest/
using LinearAlgebra
using Random
using Statistics
using Latexify
Pkg.add("StatsPlots")
using StatsPlots
# Call Scaled Interpolation Functions
    include("../Scaled_Interpolation_Functions.jl")

# Set random seed
Random.seed!(0021);

# Paramters
    # Generate structure for parameters using Parameters module
    # Set default values for parameters
    @with_kw struct Par
        # Model Parameters
        ρ::Float64 = 0.9    ; # AR(1) persistence
        σ::Float64 = 2.0    ; # Standard deviation of inovations η
        n_draw::Int64 = 10000 ; # Number of draws for simulated markov chain
    end
# Allocate parameters to object p for future calling
p = Par()

# Function to generate markov chain z using monte carlo simulation
function Simulate_MC(p::Par)
    @unpack ρ,σ,n_draw = p
    # Generate inovations η
    η = rand(Normal(0.0,σ),n_draw)
    # initialize process for z
    z = zeros(n_draw)
    # Compute z according to AR(1) process
    z[1] = η[1]
    for i in 2:n_draw
        z[i] = ρ*z[i-1] + η[i]
    end
    return z
end

# Function to distretize AR(1) markov process with Tauchen (1986)
function Tauchen86(N,p::Par,Ω::Any=3)
    @unpack ρ,σ = p
    # INPUTS:
        # ρ: persistence of unerlying AR(1) process where z' = ρz+η
        # σ: Std dev of inovation η in AR(1) process where η∼N(0,σ^2)
        # N: Size of grid for discrete process
        # Ω: Denotes boundaries of the grid as the number of standard deviations below and above the mean (lb=-Ωσ,ub=Ωσ)
    # OUTPUT:
        # z: All possible values of discretized AR(1) process, equally spaced grid of size N
        # Π: Matrix of transition probabilities
        # PDF_z: Stationary PDF of z
    #--------------------------------------------------------------------------
     z = collect(range(-Ω*σ,Ω*σ;length=N))
     Δz= 2*Ω*σ/(N-1)
     # initialize and compute probability transition matrix
     Π = zeros(N,N)
     F(x) = cdf.(Normal(),x)
     function prob_ij(zzp,zz,z::Vector)
         if zzp == z[1]
             return F((zzp+Δz/2-ρ*zz)/σ)
         elseif zzp == z[N]
             return 1-F((zzp-Δz/2-ρ*zz)/σ)
         else
             return F((zzp+Δz/2-ρ*zz)/σ)-F((zzp-Δz/2-ρ*zz)/σ)
         end
     end
     for i in 1:N
         for j in 1:N
            Π[i,j] = prob_ij(z[j],z[i],z)
        end
     end
     # Stationary distribution
     PDF_z = real(eigvecs(Π')[:,end]); PDF_z = PDF_z/sum(PDF_z)
     return (z,Π,PDF_z)
end

# Function to distretize AR(1) markov process with Rouwenhorst (1995)
function Rouwenhorst95(N,p::Par)
    @unpack ρ,σ=p
    # INPUTS:
        # ρ: persistence of unerlying AR(1) process where z' = ρz+η
        # σ: Std dev of inovation η in AR(1) process where η∼N(0,σ^2)
        # N: Size of grid for discrete process
    # OUTPUT:
        # z: All possible values of discretized AR(1) process, equally spaced grid of size N
        # Π: Matrix of transition probabilities
        # PDF_z: Stationary PDF of z
    #---------------------------------------------------------------------------
    Π = zeros(N,N)
    Π_Nm = zeros(N-1,N-1)
    P = (1+ρ)/2
    ϕ = σ*(sqrt((N-1)/(1-ρ^2)))
    z = range(-ϕ,ϕ;length=N)
    if N==2
        Π = [P 1-P;1-P P]
    else
        Π_Nm = Rouwenhorst95(N-1,p)[2]
        o = zeros(N-1)
        Π = P*[Π_Nm o; o' 0] + (1-P)*[o Π_Nm; 0 o'] + (1-P)*[o' 0; Π_Nm o] + P*[0 o';o Π_Nm]
        Π = Π./repeat(sum(Π,dims=2),1,N)
    end
    PDF_z = pdf.(Binomial(N-1,0.5),(0:N-1))
    return (z,Π,PDF_z)
end

# Function to simulate a discrete markov process
function Simulate_discrete(N,z,PDF_s,Π,p::Par)
    @unpack ρ,σ,n_draw = p
    # INPUT:
        # n_draw: Number of draws for the simulation
        # N     : Number of possible values for the discrete process
        # z     : Vector of possible values for the discrete process
        # PDF_s : Stationary PDF (used to draw initial point)
        # Π     : Transition probability matrix
    # OUTPUT:
        # z_MC  : Simulated draws
    # Initialization
    z_ind = zeros(Int64,n_draw) # index values of z
    z_MC = zeros(n_draw)
    # Initial value
    z_ind[1] = rand(Categorical(PDF_s))
    z_MC[1] = z[z_ind[1]]
    # Simulation
    for i in 2:n_draw
        z_ind[i] = rand(Categorical(Π[z_ind[i-1],:]))
        z_MC[i] = z[z_ind[i]]
    end
    return z_MC
end

# Function that calculate numerical moments out of a given array
function num_moments(z)
    # First four moments
    num_mean = mean(z)
    num_var = var(z)
    num_skew = skewness(z)
    num_kurt = kurtosis(z)
    # Autocorrelation going for t-1 to t-4
    autocorr = zeros(4)
    for i in 1:4
        autocorr[i] = cor(z[(i+1):end],z[1:(end-i)])
    end
    return (num_mean,num_var,num_skew,num_kurt,autocorr)
end

                    ### a) ###
# Simulate markov chain
z = Simulate_MC(p)

                    ### b) ###
N = [5,15]
# Tauchen (1986)
ztauc_MC = zeros(p.n_draw,2)
for i in 1:2
    (ztauc,Π_tauc,PDF_ztauc) = Tauchen86(N[i],p)
    ztauc_MC[:,i] = Simulate_discrete(N[i],ztauc,PDF_ztauc,Π_tauc,p)
end
# Rouwenhorst (1995)
zrouw_MC = zeros(p.n_draw,2)
for i in 1:2
    (zrouw,Π_rouw,PDF_zrouw) = Rouwenhorst95(N[i],p)
    zrouw_MC[:,i] = Simulate_discrete(N[i],zrouw,PDF_zrouw,Π_rouw,p)
end

                    ### c) ###
# Report moments for simulations, Ns=10000
tmean,rmean,tvar,rvar,tskew,rskew,tkur,rkur = [zeros(2) for i in 1:8]
mcmean,mcvar,mcskew,mckur = [zeros(1) for i in 1:4]
tautocorr,rautocorr = [zeros(4,2) for i in 1:2]
mcautocorr = zeros(4)
for i in 1:2
    (tmean[i],tvar[i],tskew[i],tkur[i],tautocorr[:,i]) = num_moments(ztauc_MC[:,i])
    (rmean[i],rvar[i],rskew[i],rkur[i],rautocorr[:,i]) = num_moments(zrouw_MC[:,i])
end
(mcmean,mcvar,mcskew,mckur,mcautocorr) = num_moments(z)
r(x) = round(x,digits=3)
Moments_Mat = ["" "Data" "Tauchen" "Tauchen" "" "Rouwenhorst" "Rouwenhorst";
               "" " " "N=5" "N=15" "" "N=5" "N=15" ;
               "mean" r(mcmean) r(tmean[1]) r(tmean[2]) "" r(rmean[1]) r(rmean[2]);
               "var"  r(mcvar)  r(tvar[1]) r(tvar[2]) "" r(rvar[1]) r(rvar[2]);
               "skew"  r(mcskew)  r(tskew[1]) r(tskew[2]) "" r(rskew[1]) r(rskew[2]);
               "kurt" r(mckur) r(tkur[1]) r(tkur[2]) "" r(rkur[1]) r(rkur[2]);
               "acorr" r(mcautocorr[1]) r(tautocorr[1,1]) r(tautocorr[1,2]) "" r(rautocorr[1,1]) r(rautocorr[1,2])]

println("\n Moments from simulated paths: \n")
display(Moments_Mat)

copy_to_clipboard(true)
mdtable(Moments_Mat)

#-------------------------------------------------------
# Plot paths for simulations, Ns=100
gr()
plot(title="Tauchen Simulation Paths",legend=:outerright,foreground_color_legend = nothing,background_color_legend = nothing)
plot!(1:100,z[1:100], linewidth=3,label="Simulation")
plot!(1:100,ztauc_MC[1:100,1] ,linewidth=3,label="N=5",linestyle=(:dot) )
plot!(1:100,ztauc_MC[1:100,2],linewidth=3,label="N=15",linestyle=(:dash))
xlabel!("Time")
savefig("./Figures/Tauchen")

gr()
plot(title="Rouwenhorst Simulation Paths",legend=:outerright,foreground_color_legend = nothing,background_color_legend = nothing)
plot!(1:100,R_s_5[1:100] ,linewidth=3,label="N=5" )
plot!(1:100,R_s_11[1:100],linewidth=3,label="N=11",linestyle=(:dash))
plot!(1:100,R_s_21[1:100],linewidth=3,label="N=21",linestyle=(:dot))
xlabel!("Time")
savefig("./Figures/Rouwenhorst")])

density(z, label="Simulation", linewidth=3,)
density!(ztauc_MC[:,1], label="N=5", linewidth=3,)
density!(ztauc_MC[:,2], label="N=15", linewidth=3,)
title!("Tauchen")
savefig("./Figures/Tauchen")

density(z, label="Simulation", linewidth=3,)
density!(zrouw_MC[:,1], label="N=5", linewidth=3,)
density!(zrouw_MC[:,2], label="N=15", linewidth=3,)
title!("Rouwenhorst")
savefig("./Figures/Rouwen")


# This file is a part of Julia. License is MIT: https://julialang.org/license

using Test

using Base.LinAlg: BlasComplex, BlasFloat, BlasReal, QRPivoted

n = 10

# Split n into 2 parts for tests needing two matrices
n1 = div(n, 2)
n2 = 2*n1

srand(1234321)

areal = randn(n,n)/2
aimg  = randn(n,n)/2

@testset for eltya in (Float32, Float64, Complex64, Complex128, Int)
    aa = eltya == Int ? rand(1:7, n, n) : convert(Matrix{eltya}, eltya <: Complex ? complex.(areal, aimg) : areal)
    asym = aa'+aa                  # symmetric indefinite
    apd  = aa'*aa                 # symmetric positive-definite
    for (a, asym, apd) in ((aa, asym, apd),
                           (view(aa, 1:n, 1:n),
                            view(asym, 1:n, 1:n),
                            view(apd, 1:n, 1:n)))
        ε = εa = eps(abs(float(one(eltya))))

        α = rand(eltya)
        β = rand(eltya)
        eab = eig(α,β)
        @test eab[1] == eigvals(fill(α,1,1),fill(β,1,1))
        @test eab[2] == eigvecs(fill(α,1,1),fill(β,1,1))

        @testset "non-symmetric eigen decomposition" begin
            d, v = eig(a)
            for i in 1:size(a,2)
                @test a*v[:,i] ≈ d[i]*v[:,i]
            end
            f = eigfact(a)
            @test det(a) ≈ det(f)
            @test inv(a) ≈ inv(f)
            @test isposdef(a) == isposdef(f)
            @test eigvals(f) === f[:values]
            @test eigvecs(f) === f[:vectors]
            @test Array(f) ≈ a

            num_fact = eigfact(one(eltya))
            @test num_fact.values[1] == one(eltya)
            h = asym
            @test minimum(eigvals(h)) ≈ eigmin(h)
            @test maximum(eigvals(h)) ≈ eigmax(h)
            @test_throws DomainError eigmin(a - a')
            @test_throws DomainError eigmax(a - a')
        end
        @testset "symmetric generalized eigenproblem" begin
            if isa(a, Array)
                asym_sg = asym[1:n1, 1:n1]
                a_sg = a[:,n1+1:n2]
            else
                asym_sg = view(asym, 1:n1, 1:n1)
                a_sg = view(a, 1:n, n1+1:n2)
            end
            f = eigfact(asym_sg, a_sg'a_sg)
            @test asym_sg*f[:vectors] ≈ (a_sg'a_sg*f[:vectors]) * Diagonal(f[:values])
            @test f[:values] ≈ eigvals(asym_sg, a_sg'a_sg)
            @test prod(f[:values]) ≈ prod(eigvals(asym_sg/(a_sg'a_sg))) atol=200ε
            @test eigvecs(asym_sg, a_sg'a_sg) == f[:vectors]
            @test eigvals(f) === f[:values]
            @test eigvecs(f) === f[:vectors]
            @test_throws KeyError f[:Z]

            d,v = eig(asym_sg, a_sg'a_sg)
            @test d == f[:values]
            @test v == f[:vectors]
        end
        @testset "Non-symmetric generalized eigenproblem" begin
            if isa(a, Array)
                a1_nsg = a[1:n1, 1:n1]
                a2_nsg = a[n1+1:n2, n1+1:n2]
            else
                a1_nsg = view(a, 1:n1, 1:n1)
                a2_nsg = view(a, n1+1:n2, n1+1:n2)
            end
            f = eigfact(a1_nsg, a2_nsg)
            @test a1_nsg*f[:vectors] ≈ (a2_nsg*f[:vectors]) * Diagonal(f[:values])
            @test f[:values] ≈ eigvals(a1_nsg, a2_nsg)
            @test prod(f[:values]) ≈ prod(eigvals(a1_nsg/a2_nsg)) atol=50000ε
            @test eigvecs(a1_nsg, a2_nsg) == f[:vectors]
            @test_throws KeyError f[:Z]

            d,v = eig(a1_nsg, a2_nsg)
            @test d == f[:values]
            @test v == f[:vectors]
        end
    end
end

@testset "eigenvalue computations with NaNs" begin
    for eltya in (NaN16, NaN32, NaN)
        @test_throws(ArgumentError, eig(fill(eltya, 1, 1)))
        @test_throws(ArgumentError, eig(fill(eltya, 2, 2)))
        test_matrix = rand(typeof(eltya),3,3)
        test_matrix[2,2] = eltya
        @test_throws(ArgumentError, eig(test_matrix))
    end
end

# test a matrix larger than 140-by-140 for #14174
let aa = rand(200, 200)
    for a in (aa, view(aa, 1:n, 1:n))
        f = eigfact(a)
        @test a ≈ f[:vectors] * Diagonal(f[:values]) / f[:vectors]
    end
end


@testset "Charts tests:" begin

    umin   = 0.0
    umax   = 2.0
    unodes = 16

    ucoord = GaussLobatto("u", umin, umax, unodes)
    u, = Jecco.cheb(umin, umax, unodes)
    @test ucoord[:] ≈ u

    xmin   = -10.0
    xmax   =  10.0
    xnodes =  20

    xcoord = Cartesian{2}("x", xmin, xmax, xnodes, endpoint=false)
    @test xcoord[:] == collect(-10.0:1.0:9.0)

    ymin   = -20.0
    ymax   =  20.0
    ynodes =  20

    ycoord = Cartesian{3}("y", ymin, ymax, ynodes, endpoint=false)

    chart = Chart(ucoord, xcoord, ycoord)

    u, x, y = chart[:]

    @test x == collect(-10.0:1.0:9.0)
    @test y == collect(-20.0:2.0:19.0)

end


using Test
using TestSetExtensions
using LinearAlgebra
using Qaintessent

##==----------------------------------------------------------------------------------------------------------------------

# adapted from https://github.com/FluxML/Zygote.jl/blob/master/test/gradcheck.jl
function ngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in 1:length(x)
        δ = sqrt(eps())
        tmp = x[i]
        x[i] = tmp - δ/2
        y1 = f(xs...)
        x[i] = tmp + δ/2
        y2 = f(xs...)
        x[i] = tmp
        Δ[i] = (y2-y1)/δ
        if eltype(x) <: Complex
            # derivative with respect to imaginary part
            x[i] = tmp - im*δ/2
            y1 = f(xs...)
            x[i] = tmp + im*δ/2
            y2 = f(xs...)
            x[i] = tmp
            Δ[i] += im*(y2-y1)/δ
        end
    end
    return grads
end

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "statevector circuit gradients" begin

    # construct parametrized circuit
    N = 4
    A = rand(ComplexF64, 2 ,2)
    U, R = qr(A)
    U = Array(U)
    i = rand(1:N)

    cgc(θ, ϕ, χ, κ, ωn, η, ι, λ) = CircuitGate[
        circuit_gate(3, HadamardGate()),
        circuit_gate(2, RzGate(θ), (1, 4)),
        circuit_gate(2, TGate(), 4),
        circuit_gate(2, 3, SwapGate()),
        circuit_gate(2, SdagGate(), 1),
        circuit_gate(3, PhaseShiftGate(ϕ)),
        circuit_gate(1, RyGate(χ)),
        circuit_gate(4, HadamardGate(), 2),
        circuit_gate(3, X, 1),
        circuit_gate(2, Z, 4),
        circuit_gate(4, Y, 1),
        circuit_gate(3, 1, EntanglementYYGate(κ)),
        circuit_gate(3, RotationGate(ωn)),
        circuit_gate(3, RxGate(η), (1, 4)),
        circuit_gate(2, 3, EntanglementXXGate(ι)),
        circuit_gate(4, 1, EntanglementZZGate(λ)),
        circuit_gate(4, SGate()),
        circuit_gate(2, TdagGate(), 3),
        circuit_gate(i, MatrixGate(U)),
        circuit_gate(2, Z),
        circuit_gate(3, X),
        circuit_gate(2, Y, 4),
        circuit_gate(1, SGate()),
        circuit_gate(2, 3, SwapGate(), 4),
        circuit_gate(3, SdagGate()),
        circuit_gate(2, TGate()),
    ]

    # measurement operators
    meas(M) = MeasurementOperator.([Matrix{Float64}(I, 2^N, 2^N), Hermitian(M)], (Tuple(1:N),))

    # parameter values
    θ = 1.5π
    ϕ = 0.3
    χ = √2
    κ = exp(-0.4)
    η = √3
    ι = 3.33π
    λ = exp(-0.11)
    n = randn(Float64, 3)
    n /= norm(n)
    ωn = 0.2π * n
    M = randn(ComplexF64, 2^N, 2^N)
    M = 0.5*(M + adjoint(M))

    c = Circuit{N}(cgc(θ, ϕ, χ, κ, ωn, η, ι, λ), meas(M))
    # input quantum state
    ψ = Statevector(randn(ComplexF64, 2^N))
    ψl = deepcopy(ψ)
    # fictitious gradients of cost function with respect to circuit output
    Δ = [0.3, -1.2]

    dc = Qaintessent.gradients(c, ψ, Δ)[1]

    f(rθ, rϕ, rχ, rκ, ωn, rη, rι, rλ, M) = dot(Δ, apply!(copy!(ψl, ψ), Circuit{N}(cgc(rθ[], rϕ[], rχ[], rκ[], ωn, rη[], rι[], rλ[]), meas(M))))
    # numeric gradients
    ngrad = ngradient(f, [θ], [ϕ], [χ], [κ], ωn, [η], [ι], [λ], M)
    # symmetrize gradient with respect to measurement operator

    ngrad[end][:] = 0.5*(ngrad[end] + adjoint(ngrad[end]))
    @test all(isapprox.(ngrad,
        (dc.moments[1][2].gate.U.θ,
         dc.moments[4][2].gate.ϕ,
         dc.moments[5][1].gate.θ,
         dc.moments[8][1].gate.θ,
         dc.moments[9][1].gate.nθ,
         dc.moments[10][1].gate.U.θ,
         dc.moments[11][1].gate.θ,
         dc.moments[11][2].gate.θ,
         sparse_matrix(dc.meas[2])), rtol=1e-5, atol=1e-5))
end

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "circuit gradients with moments" begin

   # construct parametrized circuit
   N = 4
    cgc(θ, ϕ, χ, ωn) = Moment[
        Moment([
            circuit_gate(3, HadamardGate()),
            circuit_gate(2, RzGate(θ), (1, 4)),
        ]),
        Moment(circuit_gate(2, 3, SwapGate())),
        Moment(circuit_gate(3, PhaseShiftGate(ϕ))),
        Moment([
            circuit_gate(3, RotationGate(ωn)),
            circuit_gate(1, RyGate(χ)),
        ]),
    ]

   # measurement operators
   meas(M) = MeasurementOperator.([Matrix{Float64}(I, 2^N, 2^N), Hermitian(M)], (Tuple(1:N),))

   # parameter values
   θ = 1.5π
   ϕ = 0.3
   χ = √2
   n = randn(Float64, 3)
   n /= norm(n)
   ωn = 0.2π * n
   M = randn(ComplexF64, 2^N, 2^N)
   M = 0.5*(M + adjoint(M))

   c = Circuit{N}(cgc(θ, ϕ, χ, ωn), meas(M))

   # input quantum state
   ψ = Statevector(randn(ComplexF64, 2^N))
   ψl = deepcopy(ψ)
   # fictitious gradients of cost function with respect to circuit output
   Δ = [0.3, -1.2]

   dc = Qaintessent.gradients(c, ψ, Δ)[1]

   f(rθ, rϕ, rχ, ωn, M) = dot(Δ, apply!(copy!(ψl, ψ), Circuit{N}(cgc(rθ[], rϕ[], rχ[], ωn), meas(M))))
   # numeric gradients
   ngrad = ngradient(f, [θ], [ϕ], [χ], ωn, M)
   # symmetrize gradient with respect to measurement operator
   ngrad[end][:] = 0.5*(ngrad[end] + adjoint(ngrad[end]))

   @test all(isapprox.(ngrad,
       (dc.moments[1][2].gate.U.θ,
        dc.moments[3][1].gate.ϕ,
        dc.moments[4][2].gate.θ,
        dc.moments[4][1].gate.nθ,
        sparse_matrix(dc.meas[2])), rtol=1e-5, atol=1e-5))
end


## ASCII display

function irange(u, W, H)
    x, y = u.x.val, u.y.val
    rx = (1 + trunc(Integer, x.lo)):(1 + trunc(Integer, x.hi) - 1)
    ry = (1 + H - trunc(Integer, y.hi)):(1 + H - trunc(Integer, y.lo) - 1)
    rx, ry
end

"""

Function to graph, using ascii characters, a predicate related to a
two-dimensional equation, such as `f == 0` or `f < g`, where `f` and
`g` are functions, e.g., `f(x,y) = y - sqrt(x)`. The positional
arguments are `L`, `R`, `B`, `T` which indicate the x-y range of the
graph of the equation. The keywork arguments `W` and `H` indicate the number
of pixels to use.

Set `offset=0` to see algorithm

A pixel is important, as the graph will color a pixel

- "white" (e.g., " ") if no solution exists in the region indicated by the pixel
- "black" (e.g., "x") if a solution is known to exist in the region
- "red" (e.g. ".") if a solution _might_ exist

"""
function asciigraph(r, L=-5, R=5, B=-5, T=5; W=2^4, H=2^4)

    cols = Dict()
    cols[:red]="."; cols[:white]=" ";cols[:black]="x"

    red, black, white = GRAPH(r, L, R, B, T, W, H)

    graph = repmat([:red], H, W)
    for u in white
        r1, r2 = irange(u, W, H)
        graph[r2, r1] = :white
    end
    for u in black
        r1, r2 = irange(u, W, H)
        graph[r2, r1] = :black
    end

    mat = graph
    m,n = size(mat)
    for i in 1:m
        for j in 1:n
            print(cols[mat[i,j]])
        end
        print("\n")
    end
end


using ControlSystems
using LinearAlgebra

##
# system of the form
# 1 | 2 3
# ———————
# 4 | 5 6
# 7 | 8 9
sys1 = ControlSystems.PartionedStateSpace(ss(1.0, [2.0 3.0], [4.0; 7.0], [5.0 6.0; 8.0 9.0]), 1, 1)

matrix(x::Number) = fill(x, 1, 1)

sys1.A == matrix(1.0)
sys1.B1 == matrix(2.0)
sys1.B2 == matrix(3.0)
sys1.C1 == matrix(4.0)
sys1.C2 == matrix(7.0)
sys1.D11 == matrix(5.0)
sys1.D12 == matrix(6.0)
sys1.D21 == matrix(8.0)
sys1.D22 == matrix(9.0)

##
sys2 = ControlSystems.PartionedStateSpace(ss(fill(1.0, 2, 2), fill(2.0, 2, 5), fill(3.0, 7, 2), fill(4.0, 7, 5)), 2, 3)

@test sys2.A == fill(1.0, 2, 2)
@test sys2.B1 == fill(2.0, 2, 2)
@test sys2.B2 == fill(2.0, 2, 3)
@test sys2.C1 == fill(3.0, 3, 2)
@test sys2.C2 == fill(3.0, 4, 2)
@test sys2.D11 == fill(4.0, 3, 2)
@test sys2.D12 == fill(4.0, 3, 3)
@test sys2.D21 == fill(4.0, 4, 2)
@test sys2.D22 == fill(4.0, 4, 3)

# TODO: Add some tests for interconnections, implicitly tested through delay system implementations though
@test (sys1 + sys1).P[1, 1] == (sys1.P[1,1] + sys1.P[1,1])


# polynomial drift
include("common.jl")
import SDELab2
#
T=1 #  final time
x0=1. # initial data
alp=2. # parameter
sig=1 #
quad0=SDELab2.delta_measure(x0)
###########################################
function drift_ff(u,alp)
  x=u[1]; mf=u[2]; mf_sq=u[3]
  return alp*x+mf-x*mf_sq
end
##########################################
function diff_gg(u,sig)
  x=u[1]
  return sig*x
end
#########################################
function mf_r(x)
  return [x; x^2]
end
#########################################
fcn=SDELab2.set_fcn_mf(u->drift_ff(u,alp),
          u->diff_gg(u,sig),
          x->mf_r(x))
######################################
function odef(t, r)
    # Extract the coordinates from the r vector
    (x, y) = r
    dx_dt = (alp+1)*x- x*y
    dy_dt = (2*alp+1)*y + 2*x*x -2*y*y
    # Return the derivatives as a vector
    return [dx_dt; dy_dt]
end
####################################
using ODE
const odedt = 0.1
tvec = 0:odedt:T
tic()
tvec, pos = ode45(odef, [x0;x0], tvec,abstol=1e-9,reltol=1e-9)
exact=pos[end];
println("ode soln ",exact,"; cpu time ",toq()," secs.")
#############################################

dt=1e-2
solver="MFEM0"
params=SDELab2.set_opt(dt,solver)
params["cutoff"]=40 # range R for Alg 3.1 cutoff
params["const"]=1 # consts for no_Gauss_points
params["length_scale"]=10 # for no_Gauss_points
params["skiphalf"]=true
params["periodic"]=false
params["no_partition"]=1
dtpts=[0.05,0.025,0.01,0.005,0.0025]#,0.001]
errors=zeros(length(dtpts),2)
errors_no=deepcopy(errors);errors_nd=deepcopy(errors)
cpu_times=zeros(length(dtpts));
cpu_times_no=deepcopy(cpu_times);cpu_times_nd=deepcopy(cpu_times)
phi=x->[x;x.^2]
params["Solver"]="MFEM0"
dtpts_no=dtpts/4
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts_no[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_no[i,:],cpu_times_no[i]=helper(qoi, exact, cpu_time, "normal")
end
########################
params["Solver"]="MFEM0"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,true)
    errors[i,:],cpu_times[i]=helper(qoi, exact, cpu_time, "normal")
end
################
params["Solver"]="MFEM2"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_nd[i,:],cpu_times_nd[i]=helper(qoi, exact, cpu_time, "normal")
end
#########################################
str1="--" # extrap
str2="k-." # 2nd-order
figure(1);
# first functional
i=1
time_step_err(dtpts,errors[:,i],
  dtpts_no,errors_no[:,i],
  dtpts,errors_nd[:,i],[1e0,1e0],utf8("mom1.pdf"))
err_cpu(errors[:,i],cpu_times[:],
  errors_no[:,i],cpu_times_no[:],
  errors_nd[:,i],cpu_times_nd[:],[1e0,1e0], utf8("mom1_.pdf"))
# second functional
i=2
time_step_err(dtpts,errors[:,i],
      dtpts_no,errors_no[:,i],
      dtpts,errors_nd[:,i],[1e0,1],utf8("mom2.pdf"))
err_cpu(errors[:,i], cpu_times[:],
      errors_no[:,i],cpu_times_no[:],
      errors_nd[:,i],cpu_times_nd[:],[1e0,1], utf8("mom2_.pdf"))
#########################################
println("finished preparing plots")


function Sphere1(totalstep)
    function integrand(config)
        X = config.var[1]
        if (X[1]^2 + X[2]^2 < 1.0)
            return 1.0
        else
            return 0.0
        end
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable += weight / abs(weight) * factor
    end

    T = MCIntegration.Tau(1.0, 1.0 / 2.0)
    dof = [[2,],] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (T,), dof, 0.0)
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end

function Sphere2(totalstep)
    function integrand(config)
        X = config.var[1]
        if (X[1][1]^2 + X[1][2]^2 < 1.0)
            return 1.0
        else
            return 0.0
        end
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable += weight / abs(weight) * factor
    end

    T = MCIntegration.TauPair(1.0, 1.0 / 2.0)
    dof = [[1,],] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (T,), dof, 0.0)
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end

function Sphere3(totalstep; offset = 0)
    function integrand(config)
        @assert config.curr == 1 || config.curr == 2
        X = config.var[1]
        if config.curr == 1
            if (X[1+offset]^2 + X[2+offset]^2 < 1.0)
                return 1.0
            else
                return 0.0
            end
        else
            if (X[1+offset]^2 + X[2+offset]^2 + X[3+offset]^2 < 1.0)
                return 1.0
            else
                return 0.0
            end
        end
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable[config.curr] += weight / abs(weight) * factor
    end

    T = MCIntegration.Continuous([0.0, 1.0], 1.0 / 2.0, offset = offset)
    dof = [[2,], [3,]] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (T,), dof, [0.0, 0.0])
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end

function Exponential1(totalstep)
    function integrand(config)
        X = config.var[1]
        return exp(-X[1])
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable += weight / abs(weight) * factor
    end

    K = MCIntegration.RadialFermiK(1.0, 0.01)
    dof = [[1,],] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (K,), dof, 0.0)
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end
const LorentzN = 1000

function Lorentz1(totalstep)
    function integrand(config)
        X = config.var[1]
        return 1.0 / ((X[1] - 1)^2 + (1.0 / LorentzN)^2)
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable += weight / abs(weight) * factor
    end

    K = MCIntegration.RadialFermiK(1.0, 0.001)
    dof = [[1,],] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (K,), dof, 0.0)
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end

function Exponential2(totalstep)
    function integrand(config)
        @assert config.curr == 1 || config.curr == 2
        X = config.var[1]
        if config.curr == 1
            return exp(-X[1])
        else
            return exp(-X[1] - X[2])
        end
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable[config.curr] += weight / abs(weight) * factor
    end

    K = MCIntegration.RadialFermiK(1.0, 0.01)
    dof = [[1,], [2,]] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (K,), dof, [0.0, 0.0])
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end

function Lorentz2(totalstep)
    function integrand(config)
        @assert config.curr == 1 || config.curr == 2
        X = config.var[1]
        if config.curr == 1
            return 1.0 / ((X[1] - 1)^2 + (1.0 / LorentzN)^2)
        else
            return 1.0 / ((X[1] - 1)^2 + (1.0 / LorentzN)^2) * exp(-X[2])
        end
    end

    function measure(config)
        factor = 1.0 / config.reweight[config.curr]
        weight = integrand(config)
        config.observable[config.curr] += weight / abs(weight) * factor
    end

    K = MCIntegration.RadialFermiK(1.0, 0.01)
    dof = [[1,], [2,]] # number of T variable for the normalization and the integrand
    config = MCIntegration.Configuration(totalstep, (K,), dof, [0.0, 0.0])
    avg, err = MCIntegration.sample(config, integrand, measure; Nblock = 64, print = -1)
    # avg, err = MonteCarlo.sample(totalstep, (T,), dof, [0.0, ], integrand, measure; Nblock=64, print=-1)

    return avg, err
end

@testset "MonteCarlo Sampler" begin
    totalStep = 1000_00

    avg, err = Sphere1(totalStep)
    println("MC integration 1: $avg ± $err (exact: $(π / 4.0))")
    @test abs(avg - π / 4.0) < 5.0 * err
    # @test abs(avg[1] - π / 4.0) < 5.0 * err[1]

    avg, err = Sphere2(totalStep)
    println("MC integration 2: $avg ± $err (exact: $(π / 4.0))")
    @test abs(avg - π / 4.0) < 5.0 * err
    # @test abs(avg[1] - π / 4.0) < 5.0 * err[1]

    avg, err = Sphere3(totalStep)
    println("MC integration 3: $(avg[1]) ± $(err[1]) (exact: $(π / 4.0))")
    println("MC integration 3: $(avg[2]) ± $(err[2]) (exact: $(4.0 * π / 3.0 / 8))")
    @test abs(avg[1] - π / 4.0) < 5.0 * err[1]
    @test abs(avg[2] - π / 6.0) < 5.0 * err[2]

    avg, err = Sphere3(totalStep, offset = 2)
    println("MC integration 3 with offset: $(avg[1]) ± $(err[1]) (exact: $(π / 4.0))")
    println("MC integration 3 with offset: $(avg[2]) ± $(err[2]) (exact: $(4.0 * π / 3.0 / 8))")
    @test abs(avg[1] - π / 4.0) < 5.0 * err[1]
    @test abs(avg[2] - π / 6.0) < 5.0 * err[2]

    avg, err = Exponential1(totalStep)
    println("MC integration 4: $avg ± $err (exact: $(1.0))")
    @test abs(avg - 1.0) < 5.0 * err

    avg, err = Exponential2(totalStep)
    println("MC integration 5: $(avg[1]) ± $(err[1]) (exact: $(1.0))")
    println("MC integration 5: $(avg[2]) ± $(err[2]) (exact: $(1.0))")
    @test abs(avg[1] - 1.0) < 5.0 * err[1]
    @test abs(avg[2] - 1.0) < 5.0 * err[2]

    avg, err = Lorentz1(totalStep)
    println("MC integration 6: $avg ± $err (exact: $((LorentzN / 2 * π + LorentzN * atan(LorentzN)))")
    @test abs(avg - (LorentzN / 2 * π + LorentzN * atan(LorentzN))) < 5.0 * err

    avg, err = Lorentz2(totalStep)
    println("MC integration 7: $(avg[1]) ± $(err[1]) (exact: $((LorentzN / 2 * π + LorentzN * atan(LorentzN)))")
    println("MC integration 7: $(avg[2]) ± $(err[2]) (exact: $((LorentzN / 2 * π + LorentzN * atan(LorentzN)))")
    @test abs(avg[1] - ((LorentzN / 2 * π + LorentzN * atan(LorentzN)))) < 5.0 * err[1]
    @test abs(avg[2] - ((LorentzN / 2 * π + LorentzN * atan(LorentzN)))) < 5.0 * err[2]

end


    # Test for GaussRadau()
    tol = 1e-14
    n = 42
    ntests = 4
    x = gaussradau(n)
    w = x[2]; x = x[1]
    @test ( length(x) == n) && ( length(w) == n )
    @test abs(x[37] - 0.908847278001044) < tol
    @test abs(w[37] - 0.031190846817016) < tol
    @test ( x[1] == -1 )


#!/usr/bin/julia
#
# Definitions of useful potential functions for use in simulations.
#
using UnitfulAtomic;
export gs_polarisation, radial_gaussian, radial_gaussian_harmonic_freq,
       tweezer_potential_depth, simple_harmonic, translate_trajectory, merge_radial_gaussian_harmonic_freq;

"""
    gs_polarisation(atom::String, λ::Number)

Find the frequency dependent ground state polarisability for Rubidium and
Caesium at the given wavelength.
"""
function gs_polarisation(atom::String, λ::Number)
  # Apply the relevant data.
  # Data from M. S. Safronova et al., Phys. Rev. A 73, 022505 (2006).
  # https://doi-org.ezphost.dur.ac.uk/10.1103/PhysRevA.73.022505
  if atom == "Cs" || atom == "cs"
    dip12 = 4.4890;
    dip32 = 6.3238;
    δE12 = 0.050932;
    δE32 = 0.053456;
    A = 17.35;
  elseif atom == "Rb" || atom == "rb"
    dip12 = 4.231;
    dip32 = 5.977;
    δE12 = 0.057314;
    δE32 = 0.058396;
    A = 10.54;
  else throw(ArgumentError("only data for Rb and Cs has been implemented")); end

  # Calculate frequency in atomic units.
  ω = ustrip(auconvert(2π*Unitful.c0/λ));

  # Return the ground state dynamic polarisability.
  α = 1/3 * ((δE12 * dip12^2)/(δE12^2 - ω^2)
             + (δE32 * dip32^2)/(δE32^2 - ω^2))
      + A;

  # Apply the atomic units.
  return α * unit(1Unitful.u"a0_au^3");
end

"""
    radial_gaussian(ω₀::Number, U₀::Number, r::AbstractVector{T}) where T<:Number

Returns the numerical values for the radial slice of a Gaussian function with
the parameters given defined over the values contained in `r`.
"""
function radial_gaussian(ω₀::Number, U₀::Number, r::AbstractVector{T}) where T<:Number
  return U₀ * (1 .- exp.(-2r.^2 / ω₀^2));
end

"""
    radial_gaussian_harmonic_freq(ω₀::Number, U₀::Number, m::Number)

Calculate the radial harmonic frequency for the Gaussian potential with the
given parameters
"""
function radial_gaussian_harmonic_freq(ω₀::Number, U₀::Number, m::Number)
  (4U₀/(m * ω₀^2))^(1/2)
end

"""
    tweezer_potential_depth(P::Number, α::Number, ω₀::Number)

Calculate the trap depth, U₀, for an optical tweezer generated by a Gaussian
laser beam with power `P`, frequency-dependent atomic ground state
polarisability of `α` (in a₀³) and beam waist `ω₀`.

Uses the SI unit conversion of `α(C m²/V) = 4π ϵ₀ a₀³ α(a₀³)` with the potential
depth given by `U₀ = (P α)/(π c ϵ₀ ω₀²)`.
"""
function tweezer_potential_depth(P::Number, α::Number, ω₀::Number)
  return (4*ustrip(α)*Unitful.uconvert(Unitful.u"m", 1Unitful.u"a0_au")^3*P)/
         (Unitful.c0 * ω₀^2);
end

"""
    simple_harmonic(m::Number, ω::Number, x::AbstractVector{T}) where T<:Number

Returns the numerical values for the simple harmonic function with the
parameters given defined over the values contained in `x`.
"""
function simple_harmonic(m::Number, ω::Number, x::AbstractVector{T}) where T<:Number
  return 0.5 * m * ω^2 * x.^2;
end

"""
    translate_trajectory(r, b, Ua, ω₀a, Ub, ω₀b)

Translate between the intended trajectory for the minimum of the radial
component of two merging Gaussian potentials and the location of the potential.
"""
function translate_trajectory(r, b, Ua, ω₀a, Ub, ω₀b)
  c = Ub/ω₀b^2 * (r-b) * exp(-2(r-b)^2/ω₀b^2);
  arg = -(2c * ω₀a/Ua)^2;
  #println("Argument: ", arg);
  if (abs(arg) > 1.0/exp(1))
    println("WARNING: outside of radius of convergence of series expansion!");
    println("Argument was ", arg);
  end
  #println(arg);

  lambert = 0.0;
  for i ∈ 1:10
    lambert += (-i)^(i-1)/factorial(i) * arg^i;
  end
  #println("Correction: ", 0.5ω₀a*sqrt(-lambert));
  return r + 0.5ω₀a*sqrt(-lambert);
end

"""
    merge_radial_gaussian_harmonic_freq(m, a, Ua, wa, b, Ub, wb)

Get the equivalent harmonic frequency of the trap in the radial direction given
the locations of the merging Gaussian potentials.
"""
function merge_radial_gaussian_harmonic_freq(m, a, Ua, wa, b, Ub, wb)
  # Imaginary values for repulsive potential, real for attractive.
  ω = 2.0*sqrt(complex(1/m*(Ub/wb^2*exp(-2(a-b)^2/wb^2)*(4(a-b)^2/wb^2 - 1) - Ua/wa^2)));
  return abs(ω);
end

function desired_trap_frequency2(γ, t, tfinal, ω0)
  s = t/tfinal;
  #println("s = ", s);
  ρ = 1 + 10*(γ-1)*s^3 - 15*(γ-1)*s^4 + 6*(γ-1)*s^5;
  #println("ρ = ", ρ);
  ρdd = (60*(γ-1)*s - 180*(γ-1)*s^2 + 120*(γ-1)*s^3)/Unitful.uconvert(unit(ω0)^-1, tfinal)^2;
  #println("ρdd = ", ρdd);
  return (ω0^2/ρ^4 - ρdd/ρ);
end

function tweezer_power_from_potential(U, w, α)
  return ((Unitful.c0 * w^2 * -U)/(4*ustrip(α)*Unitful.uconvert(Unitful.u"m", 1Unitful.u"a0_au")^3));
end

function Ua_from_ω2(ω2, m, Ub, wa, wb, a, b)
  return (wa^2*(Ub/wb^2 * exp(-2*(a-b)^2/wb^2)*((4*(a-b)^2)/wb^2 - 1) - m*ω2/4));
end

function Ub_from_ω2(ω2, m, Ua, wa, wb, a, b)
  return ((m*ω2/4 + Ua/wa^2) * wb^2/(exp(-2*(a-b)^2/wb^2)*((4*(a-b)^2)/wb^2 - 1)));
end


#=
This file contains code for calculating intervals for plotting purposes.
These functions should, at minimum, take in some form of data input and
return a tuple (min, max) of the limits corresponding to the function
=#

using Statistics: mean, median, std

"""
    zscale(input::AbstractArray,
        nsamples=1000;
        contrast=0.25,
        max_reject=0.5,
        min_npixels=5,
        k_rej=2.5,
        max_iterations=5)

Implementation of the `zscale` IRAF function for finding
colorbar limits of `input`, which showcase data near the median.
This is useful for data with extreme outliers, such as astronomical images.

## Keyword arguments
* `nsamples` - The number of samples to use from `input`. If fewer than `nsamples` are present, will use the full input
* `contrast` - The desired contrast
* `k_rej` - The number of standard deviations above which data is rejected
* `max_iteration` - The number of iterations used for fitting samples
* `max_reject` - The maximum number of pixels to reject during the iterative fitting
* `min_npixels` - The minimum number of pixels to calculate the limits after the iterative fitting

See the extended help (in REPL, `??zscale`) for technical details.

## Examples
```jldoctest
julia> img = 0:9999

julia> zscale(img)
(0, 9990)
```

## Extended help
## Description

The zscale algorithm is designed to display the image values
near the median image value, without the time consuming process
of computing a full image histogram.

This is particularly useful for astronomical images, which generally
have a very peaked histogram corresponding to the background sky
in direct imaging, or the continuum in a two dimensional spectrum.

A subset of the image is examined, and approximately `nsamples`
pixels are sampled evenly over the image. The number of lines is a
user parameter, `nsample_lines`.

The pixels are ranked in brightness to form the function `I(i)`,
where `i` is the rank of the pixel, and `I` is its value. Generally,
the midpoint of this function (the median) is very near the peak
of the image histogram.  There is a well defined slope about the
midpoint, which is related to the width of the histogram.

At the ends of the `I(i)` function, there are a few very bright and
dark pixels due to objects and defects in the field. To determine
the slope, a linear function is fit with iterative rejection:

```
I(i) = intercept + slope * (i - midpoint)
```

If more than half of the points are rejected, then there is no well
defined slope, and the full range of the sample defines `z1` and
`z2`. Otherwise, the endpoints of the linear function are used
(provided they are within the original range of the sample):

```
z1 = I(midpoint) + (slope / contrast) * (1 - midpoint)
z2 = I(midpoint) + (slope / contrast) * (npoints - midpoint)
```
"""
function zscale(
    input::AbstractArray,
    nsamples::Int = 1000;
    contrast = 0.25,
    max_reject = 0.5,
    min_npixels = 5,
    k_rej = 2.5,
    max_iterations = 5,
)

    # get samples from finite values of input
    values = float(filter(isfinite, input))
    stride = max(1, round(Int, length(values) / nsamples))
    samples = @view values[1:stride:end][1:min(nsamples, end)]
    sort!(samples)

    N = length(samples)
    vmin = first(samples)
    vmax = last(samples)

    # fit a line to the sorted samples
    min_pix = max(min_npixels, round(Int, N * max_reject))
    x = 0:(N - 1)

    ngood = N
    last_good = N + 1

    # good pixel mask (array bool faster than bitarray)
    mask = ones(Bool, N)

    # get number of neighbors to mask if a pixel is bad
    ngrow = max(1, round(Int, N / 100) ÷ 2)

    local β # line slope
    # iteratively fit samples and reject sigma-clipped outliers
    for _ in 1:max_iterations
        (ngood ≥ last_good || ngood < min_pix) && break

        # linear fit using mask
        x_ = @view x[mask]
        y = @view samples[mask]
        α, β = fit_line(x_, y)
        flat = @. samples - (α + β * x)

        # k-sigma rejection threshold
        threshold = k_rej * std(flat[mask])

        # detect and reject outliers based on threshold
        @. mask[!(-threshold ≤ flat ≤ threshold)] = false

        # dilate mask
        mask = dilate_mask(mask, ngrow)

        last_good = ngood
        ngood = count(mask)
    end

    if ngood ≥ min_pix
        slope = contrast > 0 ? β / contrast : β
        center = (N - 1) ÷ 2
        m = median(samples)
        vmin = max(vmin, m - (center - 1) * slope)
        vmax = min(vmax, m + (N - center) * slope)
    end

    return vmin, vmax
end

"""
    dilate_mask(mask, ngrow)

Takes a mask and dilates each `false` with `ngrow` `false`s on
either side.  This is equivalent to boolean "convolution".
"""
function dilate_mask(mask, ngrow)
    out = similar(mask)
    idxs = CartesianIndices(mask)
    mindx = idxs[1].I[1]
    maxdx = idxs[end].I[1]
    @inbounds for idx in idxs
        lower = max(mindx, idx.I[1] - ngrow)
        upper = min(maxdx, idx.I[1] + ngrow)
        # output will only be true if there are no falses in the input section
        out[idx] = all(mask[lower:upper])
    end
    return out
end

"""
    fit_line(x, y)

Simple linear regression, returns `(intercept, slope)`.
https://en.wikipedia.org/wiki/Simple_linear_regression
"""
function fit_line(x, y)
    mx = mean(x)
    my = mean(y)
    SSxy = sum(zip(x, y)) do (xi, yi)
        (xi - mx) * (yi - my)
    end
    SSx = sum(x) do xi
        (xi - mx)^2
    end
    β = SSxy / SSx
    α = my - β * mx
    return α, β
end


export AtomicMeasure, WeightedDiracMeasure

"""
    struct WeightedDiracMeasure{T}
        center::Vector{T}
        weight::T
    end

Represents a weighted dirac measure with centered at `center` and with weight `weight`.
"""
struct WeightedDiracMeasure{T, AT<:AbstractVector{T}} # Cannot subtype AbstractMeasureLike, since it does not contain the variables corresponding to the coefficients of the center
    center::AT
    weight::T
end
function WeightedDiracMeasure(center::AbstractVector{S}, weight::T) where {S, T}
    U = promote_type(S, T)
    return WeightedDiracMeasure(convert(AbstractVector{U}, center),
                                convert(U, weight))
end

"""
    struct AtomicMeasure{T, AT, V} <: AbstractMeasureLike{T}
        variables::V                           # Vector/Tuple of variables
        atoms::Vector{WeightedDiracMeasure{T, AT}} # Atoms of the measure
    end

An measure is said to be *atomic* if it is a sum of weighted dirac measures.
For instance, ``\\eta = 2 \\delta_{(1, 0)} + 3 \\delta_{(1/2, 1/2)}`` is an atomic measure since it is a sum of the diracs centered at ``(1, 0)`` and ``(1/2, 1/2)`` and weighted respectively by 2 and 3.
That is, ``\\mathbb{E}_{\\eta}[p] = 2p(1, 0) + 3p(1/2, 1/2)``.

The `AtomicMeasure` struct stores an atomic measure where `variables` contains the variables and `atoms` contains atoms of the measure.
"""
struct AtomicMeasure{T, AT, V} <: AbstractMeasureLike{T}
    variables::V                           # Vector/Tuple of variables
    atoms::Vector{WeightedDiracMeasure{T, AT}} # Atoms of the measure
end

function Base.show(io::IO, η::AtomicMeasure)
    print(io, "Atomic measure on the variables $(join(η.variables, ", ")) with $(length(η.atoms)) atoms:")
    for atom in η.atoms
        println(io)
        print(io, " at $(atom.center) with weight $(atom.weight)")
    end
end

measure(η::AtomicMeasure, x) = Measure(η, x)
function Measure(η::AtomicMeasure{T}, x::AbstractVector{TT}) where {T, TT}
    Measure{T, MB.MonomialBasis{monomialtype(TT), monovectype(x)}}(η, x)
end
function Measure{T, BT}(η::AtomicMeasure{T}, x) where {T, BT}
    X = monovec(x)
    sum(atom.weight * dirac(X, η.variables => atom.center) for atom in η.atoms)
end

expectation(η::AtomicMeasure, p::APL) = sum(δ -> δ.weight * p(η.variables => δ.center), η.atoms)
expectation(p::APL, η::AtomicMeasure) = expectation(η, p)

function permcomp(f, m)
    picked = BitSet()
    for i in 1:m
        k = 0
        for j in 1:m
            if !(j in picked) && f(i, j)
                k = j
                break
            end
        end
        if k == 0
            return false
        end
        push!(picked, k)
    end
    true
end

Base.isapprox(η1::WeightedDiracMeasure, η2::WeightedDiracMeasure; kws...) = isapprox(η1.weight, η2.weight; kws...) && isapprox(η1.center, η2.center; kws...)
function Base.isapprox(η1::AtomicMeasure, η2::AtomicMeasure; kws...)
    m = length(η1.atoms)
    length(η2.atoms) == m && permcomp((i, j) -> isapprox(η1.atoms[i], η2.atoms[j]; kws...), m)
end


using JSON, PiecewiseDeterministicMarkovProcesses, LinearAlgebra, Random

const p0  = convert(Dict{AbstractString,Float64}, JSON.parsefile("../examples/ml.json")["type II"])
const p1  = ( JSON.parsefile("../examples/ml.json"))
include("morris_lecar_variables.jl")
const p_ml = ml(p0)

function F_ml!(xcdot, xc, xd, parms, t::Float64)
	# vector field used for the continuous variable
	#compute the current, v = xc[1]
	xcdot[1] = xd[2] / p_ml.N * (p_ml.g_Na * (p_ml.v_Na - xc[1])) + xd[4] / p_ml.M * (p_ml.g_K  * (p_ml.v_K  - xc[1]))  + (p_ml.g_L  * (p_ml.v_L  - xc[1])) + p_ml.I_app
	nothing
end

function R_ml!(rate, xc, xd, parms, t, issum::Bool)
	if issum == false
		rate[1] = p_ml.beta_na * exp(4.0 * p_ml.gamma_na * xc[1] + 4.0 * p_ml.k_na) * xd[1]
		rate[2] = p_ml.beta_na * xd[2]
		rate[3] = p_ml.beta_k * exp(p_ml.gamma_k * xc[1] + p_ml.k_k) * xd[3]
		rate[4] = p_ml.beta_k * exp(-p_ml.gamma_k * xc[1]  -p_ml.k_k) * xd[4]
		return 0.
	else
		return (p_ml.beta_na * exp(4.0 * p_ml.gamma_na * xc[1] + 4.0 * p_ml.k_na) * xd[1] +
		p_ml.beta_na * xd[2] +
		p_ml.beta_k * exp( p_ml.gamma_k * xc[1] + p_ml.k_k) * xd[3] +
		p_ml.beta_k * exp(-p_ml.gamma_k * xc[1] - p_ml.k_k) * xd[4])
	end
end

xc0 = vec([p1["v(0)"]])
xd0 = vec([Int(p0["N"]),	#Na closed
	0,			   			#Na opened
	Int(p0["M"]),			#K closed
	0])						#K opened

nu_ml = [[-1 1 0 0];[1 -1 0 1];[0 0 -1 1];[0 0 1 -1]]
parms = vec([0.])

tf = p1["t_end"]
tf = 350.

Random.seed!(123)
println("--> chv_optim - call")
pb = PDMP.PDMPProblem(F_ml!, R_ml!, nu_ml, xc0, xd0, parms, (0.0, tf))
# result =		PDMP.pdmp!(xc0,xd0, F_ml!, R_ml!, nu_ml, parms,0.0,tf,algo=:chv_optim,n_jumps = 6)
# result =  @time PDMP.pdmp!(xc0,xd0, F_ml!, R_ml!, nu_ml, parms,0.0,tf,algo=:chv_optim,n_jumps = 4500) #cpp = 100ms/2200 jumps
res = @time PDMP.solve(pb, CHV(Tsit5()), n_jumps = 2200, save_positions = (false, true))


"""
    `CalcVSDiffusionCoefs(boundary, alphas, dt, nu)`

Calculates the RBF coefficients for the diffusion step.

# ARGUMENTS
* `boundary::Boundary`     : Boundary object
* `etas::Array{T}`         : RBF coefficients for rho
* `dt::Real`               : Time step
* `nu::Real`               : Kinematic viscosity
* `sigma = 0.2`            : Gaussian width
where {T<:Real}

# RETURNS
* `betas::Array{Float64}`  : Array of RBF coefficients
"""
function CalcVSDiffusionCoefs(boundary::Boundary,
                              alphas::Array{T},
                              dt::Real,
                              nu::Real;
                              sigma=0.2) where {T<:Real}

    # === Extract geometry from boundary ===
    X = [[point[1], point[2], 0.0] for point in boundary.bodyPTS]
    n_hats = [[n_hat[1], n_hat[2], 0.0] for n_hat in boundary.nHats]
    t_hats = [[t_hat[1], t_hat[2], 0.0] for t_hat in boundary.tHats]
    NPTS = boundary.NPTS_BODY

    # === Constants ===
    CONST1 = -1 / (nu * dt)

    # === Build matrix ===
    A = zeros(NPTS, NPTS)
    for k in 1:NPTS
        for i in 1:NPTS
            if k != i
                R = X[k] - X[i]
                r = norm(R)
                r_hat = R ./ r
                n_hat = n_hats[k]
                A[k,i] = dot(RBF_gauss(r, deriv = 1, sigma=sigma) .* r_hat, n_hat)
            else
                R = [0, 0, 0]
                r = norm(R)
                r_hat = [0, 0, 0]
                n_hat = n_hats[k]
                A[k,i] = dot(RBF_gauss(r, deriv = 1, sigma=sigma) .* r_hat, n_hat)
            end
        end
    end

    # === Build RHS ===
    b = zeros(NPTS)
    for k in 1:NPTS
        b[k] = CalcVS(boundary, alphas, X[k]) * CONST1
    end

    # === Solve linear system ===
    betas = A\b
    return betas

end

"""
    `CalcDiffusion(boundary, beta, X_eval)`

Evaluates the RBF interpolation function: omega(X) = sum(beta_i * phi(X - X_i)).

# ARGUMENTS
* `boundary::Boundary` : Boundary object
* `beta::Array{T}`     : Array containing RBF coefficients
* `X_eval::Array{T}`   : Point to evaluate function at
* `sigma = 0.2`        : Gaussian width
where {T<:Real}

# RETURNS
* `omega::Float64`     : Evaluated function value
"""
function CalcDiffusion(boundary::Boundary,
                       beta::Array{T},
                       X_eval::Array{T};
                       sigma=0.2) where {T<:Real}

    # === Unpack geometry ===
    XP = [[point[1], point[2], 0.0] for point in boundary.bodyPTS]
    NPTS = boundary.NPTS_BODY

    # === Summation over all points for gamma ===
    omega = sum( [beta[i] * RBF_gauss(norm(X_eval - XP[i]), sigma=sigma) for i in 1:NPTS ] )
    return omega
end


using JuMP
using Ipopt
using PowerModels
using StochasticPowerModels

# constants
const PM = PowerModels
const SPM = StochasticPowerModels

# solvers
ipopt_solver = Ipopt.Optimizer

# input
deg  = 2
aux  = true
red  = false
case = "case5_spm.m"

# data
file  = joinpath(BASE_DIR, "test/data/matpower", case)

#-----------------------------------
# run the convenience functions for stochastic OPF for IVR and ACR
result_ivr = run_sopf_iv(file, PM.IVRPowerModel, ipopt_solver, aux=aux, deg=deg, red=red)
result_acr = run_sopf_acr(file, PM.ACRPowerModel, ipopt_solver, aux=aux, deg=deg)

@assert result_ivr["termination_status"] == PM.LOCALLY_SOLVED
@assert result_acr["termination_status"] == PM.LOCALLY_SOLVED
#the optimal objective values (expectation) are
obj_ivr = result_ivr["objective"]
obj_acr = result_acr["objective"]
@assert obj_ivr ≈ obj_acr

# print variables for all polynomial indices k
SPM.print_summary(result_ivr["solution"])

# print variables for a specific index k
k=1
SPM.print_summary(result_ivr["solution"]["nw"]["$k"])

# get polynomial chaos coefficients for specific component
pg_coeff = pce_coeff(result_ivr, "gen", 1, "pg")

# obtain 10 samples of the generator active power output variable
pg_sample = sample(result_ivr, "gen", 1, "pg"; sample_size=10)

# obtain an kernel density estimate of the generator active power output variable
pg_density = density(result_ivr, "gen", 1, "pg"; sample_size=10)

#-----------------------------------
# alternatively, you can first read in PowerModels dict,
# from a file with stochastic data extensions:
data  = PM.parse_file(file)

result_ivr2 = run_sopf_iv(data, PM.IVRPowerModel, ipopt_solver; aux=aux, deg=deg)
@assert result_ivr2["termination_status"] == PM.LOCALLY_SOLVED
obj_ivr2 = result_ivr2["objective"]

#-----------------------------------
# finally, we can also build the multinetwork dictionary here.
# we run the stochastic replicate function, introduces polynomial index k
sdata = SPM.build_stochastic_data(data, deg)

# run the reduced IVR with auxiliary variables
result_ivr3 = PM.run_model(sdata, PM.IVRPowerModel, ipopt_solver, SPM.build_sopf_iv_with_aux; multinetwork=true, solution_processors=[PM.sol_data_model!])
@assert result_ivr3["termination_status"] == PM.LOCALLY_SOLVED
obj_ivr3 = result_ivr3["objective"]

@assert obj_ivr ≈ obj_ivr2 ≈ obj_ivr3


module DivConformingFESpacesTests

using Test
using Gridap.Geometry
using Gridap.ReferenceFEs
using Gridap.FESpaces
using Gridap.CellData
using Gridap.TensorValues
using Gridap.Fields
using Gridap.Io

function test_div_v_q_equiv(U,V,P,Q,Ω)
  v=get_fe_basis(V)
  u=get_trial_fe_basis(U)

  q=get_fe_basis(Q)
  p=get_trial_fe_basis(P)

  dΩ=Measure(Ω,1)
  dΩᵣ=Measure(Ω,1,ReferenceDomain())

  a1(p,v)=∫(divergence(v)*p)dΩ
  a2(p,v)=∫(DIV(v)*p)dΩᵣ

  tol=1.0e-12
  assem = SparseMatrixAssembler(P,V)
  data = collect_cell_matrix(P,V,a1(p,v))
  A1 = assemble_matrix(assem,data)
  data = collect_cell_matrix(P,V,a2(p,v))
  A2 = assemble_matrix(assem,data)
  @test norm(A1-A2) < tol

  a3(u,q)=∫(q*divergence(u))dΩ
  a4(u,q)=∫(q*DIV(u))dΩᵣ
  assem = SparseMatrixAssembler(U,Q)
  data = collect_cell_matrix(U,Q,a3(u,q))
  A3 = assemble_matrix(assem,data)
  data = collect_cell_matrix(U,Q,a4(u,q))
  A4 = assemble_matrix(assem,data)
  @test norm(A3-A4) < tol

  uh=FEFunction(U,rand(num_free_dofs(U)))
  l1(q)=∫(q*divergence(uh))dΩ
  l2(q)=∫(q*DIV(uh))dΩᵣ
  v1=assemble_vector(l1,Q)
  v2=assemble_vector(l2,Q)
  @test norm(v1-v2) < tol
end

domain =(0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)

order = 1

u(x) = x

reffe = ReferenceFE(raviart_thomas,order)

V = TestFESpace(model,reffe,dirichlet_tags = [1,6])
test_single_field_fe_space(V)
U = TrialFESpace(V,u)

reffe = ReferenceFE(lagrangian,Float64,order)
Q = TestFESpace(model,reffe,conformity=:L2)
P = TrialFESpace(Q)

uh = interpolate(u,U)

e = u - uh

Ω = Triangulation(model)
dΩ = Measure(Ω,2*order)

el2 = sqrt(sum( ∫( e⋅e )*dΩ ))
@test el2 < 1.0e-10

test_div_v_q_equiv(U,V,P,Q,Ω)

#using Gridap.Visualization
#
#writevtk(Ω,"trian",nsubcells=10,cellfields=["uh"=>uh])

order = 1

reffe = ReferenceFE(TET,raviart_thomas,order)

domain =(0,1,0,1,0,1)
partition = (3,3,3)
model = simplexify(CartesianDiscreteModel(domain,partition))

labels = get_face_labeling(model)
dir_tags = Array{Integer}(undef,0)

V = FESpace(model,reffe,conformity=DivConformity())
U = TrialFESpace(V,u)

reffe = ReferenceFE(lagrangian,Float64,order)
Q = TestFESpace(model,reffe,conformity=:L2)
P = TrialFESpace(Q)

v(x) = VectorValue(-0.5*x[1]+1.0,-0.5*x[2],-0.5*x[3])
vh = interpolate(v,V)

e = v - vh

Ω = Triangulation(model)
dΩ = Measure(Ω,2*order)

el2 = sqrt(sum( ∫( e⋅e )*dΩ ))
@test el2 < 1.0e-10

test_div_v_q_equiv(U,V,P,Q,Ω)

#using Gridap.Visualization
#
#writevtk(trian,"test",order=3,cellfields=["vh"=>vh])

# Tests on manifold

# Create domain
domain = (0,1,0,1,0,1)
cells  = (2,2,2)
model  = CartesianDiscreteModel(domain,cells)

# Restrict model to cube surface (using new BoundaryDiscreteModel)
labels = get_face_labeling(model)
bgface_to_mask = get_face_mask(labels,"boundary",2)
Γface_to_bgface = findall(bgface_to_mask)
Dc2Dp3model = BoundaryDiscreteModel(Polytope{2},model,Γface_to_bgface)

order  = 0
degree = 1

reffe_rt = ReferenceFE(raviart_thomas,Float64,order)
V  = FESpace(Dc2Dp3model, reffe_rt ; conformity=:HDiv)
U = TrialFESpace(V,u)
reffe = ReferenceFE(lagrangian,Float64,order)
Q = TestFESpace(Dc2Dp3model,reffe,conformity=:L2)
P = TrialFESpace(Q)
uh = FEFunction(V,rand(num_free_dofs(V)))
vh = interpolate_everywhere(uh,V)

Ω = Triangulation(Dc2Dp3model)
dΩ = Measure(Ω,2*order)
e=sqrt(sum(∫((uh-vh)⋅(uh-vh))dΩ))
@test e < 1.0e-12

test_div_v_q_equiv(U,V,P,Q,Ω)

end # module


## Factor over Z[x]
## big prime

# Algo 15.2
# f is square free
function factor_Zx_big_prime_squarefree(f)
    n = degree(f)
    n == 1 && return f

    A = maxnorm(f)
    b = value(lead(f))

    # we use BigInt for p, B if we need to
    lambda = log(typemax(Int)/(4*A*b))/(1/2 + log(2))
    if n < -1 #lambda
        p = 0
        B = floor(Int, sqrt(n+1)*big(2)^n*big(A)*b)
    else
        p = big(0)
        B = floor(BigInt, sqrt(n+1)*big(2)^n*big(A)*b)
    end
    while true
        p = Primes.nextprime(rand(2B:4B))
        g = gcd(as_poly_Zp(f, p, "x"), as_poly_Zp(derivative(f), p, "x"))
        isone(g) && break
    end

    fZp = as_poly_Zp(f, p, "x")
    fZp = fZp * inv(lead(fZp))

    d = factor_Zp_squarefree(fZp, p, variable(fZp))
    gs = collect(keys(d))

    R,x = ZZ["x"] # big?
    Gs = as_poly.(poly_coeffs.(gs), x)

    _factor_combinations(f, Gs, p, 1, x, b, B)

end

##################################################

# algo 15.10
# g, h rel prime
# lifts f,g,h,s,t over F_m to values in F_m^2
function hensel_step(f, g, h, s, t, m)
    # f, g, h,s, t are in Z[x]

    isone(lead(h)) || error("h must be monic")
    degree(f) == degree(g) + degree(h) || error("degree(f) != degree(g) + degree(h)")
    degree(s) < degree(h) && degree(t) < degree(g) || error("degree(s) !< degree(h) or degree(t) !< degree(g)")

    e = as_poly_modp(as_poly(f) - as_poly(g) * as_poly(h), m^2)
    f, g, h, s, t = as_poly_modp.((f,g,h,s,t), m^2)
    q, r = divrem(s*e, h)

    gstar = g + t * e + q * g
    hstar = h + r

    b = s*gstar + t * hstar - one(t)
    c, d = divrem(s * b, hstar)

    sstar = s - d
    tstar = t - t*b - c*gstar

    iszero(f - gstar * hstar) || error("f != g^* * h^* mod m^2")
#    isone(sstar * tstar + tstar * hstar) || error("st + th != 1 mod m^2")

    as_poly.((gstar, hstar, sstar, tstar))
end

# collect factors into a tree for apply HenselStep
abstract type AbstractFactorTree end

mutable struct FactorTree <: AbstractFactorTree
    fg
    children
    s
    t
    FactorTree(fg) = new(fg, Any[])
end

mutable struct FactorTree_over_Zp <: AbstractFactorTree
    fg
    children
    s
    t
    FactorTree_over_Zp(fg, p) = new(fg, Any[])
end

function make_factor_tree_over_Zp(f, fs, p) # fs factors over Zp
    N = length(fs)
    n = ceil(Int, log2(N))
    tau = FactorTree(f)
    N == 1 && return tau

    k = 2^(n-1)
    fls = fs[1:k]
    fl = prod(as_poly.(fls))
    frs = fs[(k+1):end]
    fr = prod(as_poly.(frs))

    l, r = tau.children = Any[make_factor_tree_over_Zp(fl, fls, p),
                              make_factor_tree_over_Zp(fr, frs, p)]
    g, s, t = gcdx(as_poly_Zp(l.fg,p), as_poly_Zp(r.fg,p))

    gi = invmod(value(coeff(g, 0)), p)
    tau.s = as_poly(gi*s); tau.t = as_poly(gi*t)
    tau
end

function hensel_step_update_factor_tree!(tau, p)
    !has_children(tau) && return
    l,r = tau.children
    f, g,h,s,t = tau.fg, l.fg, r.fg, tau.s, tau.t

    g1,h1,s1,t1 = hensel_step(f, g, h, s, t, p)

    tau.s, tau.t = s1, t1
    l.fg, r.fg = g1, h1

    hensel_step_update_factor_tree!(l, p)
    hensel_step_update_factor_tree!(r, p)
end

has_children(tau::AbstractFactorTree) = length(tau.children) == 2

function all_children(tau::AbstractFactorTree)
   has_children(tau) ? vcat(all_children(tau.children[1]), all_children(tau.children[2])) : [tau.fg]
end

"""

Algo 15.17 multifactor Hensel lifting

"""
function hensel_lift(f, facs, m::T, a0, l) where {T}

    tau = make_factor_tree_over_Zp(f, facs, m)

    d = ceil(Int, log2(l))
    for j = 1:d
        a0 = mod(2*a0 - lead(f) * a0^2, m^2^j)
        tau.fg = a0 * f
        hensel_step_update_factor_tree!(tau, m^2^(j-1))
    end

    tau
end

# factor square free poly in Z[x] using hensel lifting techique
# can resolve factors of p^l using brute force or LLL algorith (which is
# not as competitive here)
# Could rewrite latter to use floating point numbers (http://perso.ens-lyon.fr/damien.stehle/downloads/LLL25.pdf)
function factor_Zx_prime_power_squarefree(f,lll=false)

    n = degree(f)

    n == 1 && return (f,)

    A = big(maxnorm(f))
    b = abs(value(lead(f)))

    B = sqrt(n+1)*big(2)^n * A * b
    C = big(n+1)^(2n) * A^(2n-1)

    gamma = 2*log2(C) + 1

    # hack to try and factor_Zp over small prime if possible
    #
    gb = 2 * gamma*log(gamma)
    if gb < (typemax(Int))^(1/10)
        gamma_bound = floor(Int, gb)
    else
        gamma_bound = floor(BigInt, gb)
    end

    gamma_bound_lower = floor(BigInt, sqrt(gb))

    fs = poly_coeffs(f)

    ## find p to factor over.
    ## for lower degree polynomials we choose more than 1 p, and select
    ## that with the fewest factors over Zp.
    ## This trades of more time factoring for less time resolving the factors

    p, l = zero(gamma_bound), 0
    fbar = f # not type stable, but can't be, as cycle through GF(p)["x"]
    hs = nothing

    P, M = 0, Inf
    for k in 1:(max(0,3-div(n,15)) + 1)

        while true
            P = Primes.prevprime(rand(gamma_bound_lower:gamma_bound))
            iszero(mod(b, P)) && continue
            fbar = as_poly_Zp(fs, P, "x")
            isone(gcd(fbar, derivative(fbar))) && break
        end

        L = floor(Int, log(P, 2B+1))
        d = factor_Zp(fbar, P)

        # modular factorization
        ds = collect(keys(d))
        nfacs = length(ds)

        if nfacs < M
            p = P
            l = L
            hs = ds
            M = nfacs
        end

    end

    # hensel lifting
    a = invmod(b, p)
    # set  up factor tree

    tau = make_factor_tree_over_Zp(f, hs, p)
    # use 15.17 to lift f = b * g1 ... gr mod p^l, gi = hi mod p
    d = ceil(Int, log2(1+l))
    for j = 1:d
        a = mod(2*a - lead(f) * a^2, p^2^j)
        tau.fg = a * f
        hensel_step_update_factor_tree!(tau, p^2^(j-1))
    end

    hs = all_children(tau)

    if lll
        # slower, and not quite right
        identify_factors_lll(f, hs, p, 2^d, b, B)
    else
        x = variable(hs[1])
        _factor_combinations(f, hs, p, 2^d, x, b, B)
    end

end

"""
Division algorithm for Z[x]

returns q,r with

* `a = q*b + r`
* `degree(r) < degree(b)`

If no such q and r can be found, then both `q` and `r` are 0.

"""
function exact_divrem(a, b)
    f,g = a, b
    x = variable(g)
    q, r = zero(f), f

    while degree(r) >= degree(g)
        u, v = divrem(lead(r), lead(g))
        v != 0 && return zero(a), zero(a)
        term = u * x^(degree(r) - degree(g))
        q = q + term
        r = r - term * g
    end
    (q, r)
end

# a = b^k c, return (c, k)
function deflate(a, b)
    k = 0
    while degree(a) >= degree(b)
        c, r = exact_divrem(a, b)
        iszero(c) && return(a, k)
        !iszero(r) && return (a, k)
        k += 1
        a = c
    end
    return (a, k)
end

## f in Z[x]
## find square free
## return fsq, g where fsq * g = f
function square_free(f)
    degree(f) <= 1 && return f

    g = gcd(f, derivative(f))
    d = degree(g)

    if  d > 0
        fsq = divexact(f, g) #faster than exact_divrem if r=0 is known
    else
        fsq = f
    end
    fsq, g
end

##################################################
### Resolve factors over F_p^l

# Brute force method
# f has factors G over Fp^l
function _factor_combinations(f, Gs, p, l, x, b, B)
    r = length(Gs)
    q = p^l

    r == 1 && return (as_poly(as_poly_modp(b*Gs[1],q,x)), )

    for s in 1:(r ÷ 2)
        for inds in Combinatorics.combinations(1:r, s)
            _inds = setdiff(1:r, inds)
            _gstar, _hstar = one(x), one(x)
            for i in inds
                _gstar *= Gs[i]
            end
            for j in _inds
                _hstar *= Gs[j]
            end
            gstar, hstar = as_poly.(as_poly_modp.((b*_gstar, b*_hstar), q))

            if onenorm(gstar) * onenorm(hstar) <= B
                Gs = Gs[_inds]
                f = as_poly(primpart(hstar))
                b = value(lead(f))
                return (primpart(gstar), _factor_combinations(f, Gs, p, l, x, b, B)...)
            end
        end
    end

    return (as_poly(as_poly_modp(b*prod(Gs),p)), )

end

### Interfaces

function poly_factor(f::AbstractAlgebra.Generic.Poly{Rational{T}}) where {T <: Integer}
    x, qs = variable(f), poly_coeffs(f)
    common_denom = lcm(denominator.(qs))
    fs = numerator.(common_denom * qs)
    fz = as_poly(fs)

    U = poly_factor(fz)
    V = Dict{typeof(f), Int}()
    for (k,v) in U
        j = as_poly(poly_coeffs(k), x)
        V[j] = v
    end
    a = 1//common_denom * one(f)
    if !isone(a)
        d = filter(kv -> degree(kv[1]) == 0, V)
        if length(d) == 0
            V[a] = 1
        else
            k,v = first(d)
            V[k*a] = 1
            delete!(V, k)
        end
    end
    V
end

function poly_factor(f::AbstractAlgebra.Generic.Poly{T}) where {T <: Integer}

    # want content free poly with positive leading coefficient
    degree(f) < 0 && return Dict(f=>1)
    c, pp = content(f), primpart(f)
    if sign(lead(f)) < 0
        c,f,pp = -c, -f, -pp
    end

    if degree(f) == 1
        isone(c) && return Dict(f=>1)
        return Dict(c * one(f)=>1, pp=>1)
    end

    fsq, g = square_free(pp)
    fs = factor_Zx_prime_power_squarefree(fsq)

    # get factors
    U = Dict(u=>1 for u in fs)
    if !isone(c)
        U[c*one(f)] = 1
    end

    # get multiplicities
    for u in fs
        g, k = deflate(g, u)
        if k > 0
            U[u] += k
        end
        degree(g) <= 0 && break
    end

    U

end

# an iterable yielding a0,a1, ..., an
# ? How wide a type signature do we want here
function poly_factor(f::Function)
    z = f(0)
    if isa(z, Rational)
        R,x = QQ["x"]
    elseif isa(z, Integer)
        R,x = ZZ["x"]
    else
        error("f(x) is not in ZZ[x] or QQ[x]")
    end

    poly_factor(f(x))
end
poly_factor(as) = poly_factor(as_poly(as))


######### Stan ODE example  ###########
#
# This is very experimental.
#
# Running this example with a single chain often returns a completed simulation.
# As more chains are added, you will find the ODE solver and/or Stan HMC might not finish,
# If this happens, in the REPL, ^C will terminate the run. Trying it a few more time
# will typically lead to a successful simulation. On my system completion takes < 1 minute.
#
# Bob Carpenter suggested to lower the stepsize and increase acceptance rate. This can
# be done as shown below (commented out lines 97 and 97). Limited testing has not shown
# guaranteed improvements in successful completion of the simulation.
#

using Stan, Distributions

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  odemodel = "
  functions {
         real[] sho(real t,
                    real[] y,
                    real[] theta,
                    real[] x_r,
                    int[] x_i) {
           real dydt[2];
           dydt[1] <- y[2];
           dydt[2] <- -y[1] - theta[1] * y[2];
           return dydt;
          }
  }

  data {
    int<lower=1> T;
    int<lower=1> M;
    real t0;
    #real y0[M];
    real ts[T];
    real y[T,M];
  }

  transformed data {
    real x_r[0];
    int x_i[0];
  }

  parameters {
    real y0[2];
    vector<lower=0>[2] sigma;
    real theta[1];
  }

  model {
    real y_hat[T,2];
    sigma ~ cauchy(0,2.5);
    theta ~ normal(0,1);
    y0 ~ normal(0,1);
    y_hat <- integrate_ode(sho, y0, t0, ts, theta, x_r, x_i);
    for (t in 1:T)
     y[t] ~ normal(y_hat[t], sigma);
  }
  "

  #Just store date in ho.csv, drop the header
  ho_data = readdlm("ho.csv.data", header=true)[1];

  #Extract initial values in row 1 of ho_data
  t0 = ho_data[1, 1]
  y0 = reshape(ho_data[1, 2:3], 2)

  # Now we can drop the first row of ho_data
  ts = ho_data[2:end, 1]
  ho_data = ho_data[2:end, 2:3]

  #Get the proper sizes to pass to CmdStan
  T = size(ho_data, 1)
  M = size(ho_data, 2)

  #Prepare to add the noise
  sigma = 0.15
  noise = [rand(Normal(0, sigma)) for i in 1:T, j in 1:M]
  y = ho_data + noise

  odedict =
    Dict(
      "T" => T,
      "M" => M,
      "t0" => t0,
      "y0" => y0,
      "ts" => ts,
      "y" => y)

  global stanmodel, rc, sim1
  stanmodel = Stanmodel(name="ode", model=odemodel, nchains=2);
  #stanmodel.method.algorithm.stepsize = 0.7 # default stepsize=1.0
  #stanmodel.method.adapt.delta = 0.9 # default delta=0.8

  println("\nStanmodel that will be used:")
  stanmodel |> display
  println()

  rc, sim1 = stan(stanmodel, [odedict], ProjDir, diagnostics=false,
    CmdStanDir=CMDSTAN_HOME);

end # cd


## Aircraft stability analysis example
using AeroMDAO
using LinearAlgebra # For norm()

# Lifting surfaces
wing  = WingSection(span       = 4.0,
                    dihedral   = 5.0,
                    sweep_LE   = 15.0,
                    taper      = 0.4,
                    root_chord = 2.0,
                    root_twist = 0.0,
                    tip_twist  = -2.0,
                    root_foil  = naca4((2,4,1,2)),
                    tip_foil   = naca4((2,4,1,2)))
wing_mac  = mean_aerodynamic_center(wing)
wing_pos  = [0., 0., 0.]
wing_plan = plot_wing(wing;
                      position = wing_pos)

print_info(wing, "Wing")

htail = WingSection(span       = 1.0,
                    dihedral   = 0.0,
                    sweep_LE   = 15.0,
                    taper      = 0.6,
                    root_chord = 0.8,
                    root_twist = 0.0,
                    tip_twist  = 0.0,
                    root_foil  = naca4((0,0,1,2)),
                    tip_foil   = naca4((0,0,0,9)));
htail_mac  = mean_aerodynamic_center(htail)
htail_pos  = [5., 0., 0.]
α_h_i      = 0.
htail_plan = plot_wing(htail;
                       position = htail_pos)

print_info(htail, "Horizontal Tail")

vtail = HalfWingSection(span       = 0.8,
                        dihedral   = 0.0,
                        sweep_LE   = 8.0,
                        taper      = 0.6,
                        root_chord = 0.8,
                        root_twist = 0.0,
                        tip_twist  = 0.,
                        root_foil  = naca4((0,0,0,9)),
                        tip_foil   = naca4((0,0,0,9)))
vtail_mac  = mean_aerodynamic_center(vtail) # NEEDS FIXING FOR ROTATION
vtail_pos  = [5., 0., 0.]
vtail_plan = plot_wing(vtail;
                       position = vtail_pos,
                       angle    = π/2)

print_info(vtail, "Vertical Tail")

## Static stability
S, b, c = projected_area(wing), span(wing), mean_aerodynamic_chord(wing)
x_w = wing_pos + [ wing_mac[1], 0, 0 ]

# Horizontal tail volume coefficient
S_h = projected_area(htail)
x_h = htail_pos + [ htail_mac[1], 0, 0 ]
l_h = norm(x_h - x_w)
V_h = S_h * l_h / (S * c)
println("Horizontal TVC      V_h: $V_h")

# Vertical tail volume coefficient
S_v = projected_area(vtail)
x_v = vtail_pos + [ vtail_mac[1], 0, 0 ]
l_v = norm(x_v - x_w)
V_v = S_v * l_v / (S * b)
println("Vertical TVC        V_v: $V_v")

## Panelling and assembly
wing_panels  = 	panel_wing(wing, [20], 10;
                           position = wing_pos
                          )
htail_panels =	panel_wing(htail, [10], 5;
                           position = htail_pos,
                           angle    = deg2rad(α_h_i),
                           axis     = [0., 1., 0.]
                          )
vtail_panels = 	panel_wing(vtail, [10], 5;
                           position = vtail_pos,
                           angle    = π/2
                          )

aircraft = Dict("Wing"            => wing_panels,
                "Horizontal Tail" => htail_panels,
                "Vertical Tail"   => vtail_panels);

## Evaluate case
ac_name = "My Aircraft"
ρ 		= 1.225
ref     = x_w
V, α, β = 1.0, 3.0, 0.0
Ω 		= [0.0, 0.0, 0.0]
fs 	    = Freestream(V, α, β, Ω)

@time dv_data =
solve_stability_case(aircraft, fs;
                     rho_ref     = ρ,           # Reference density
                     r_ref       = ref,         # Reference point for moments
                     area_ref    = S,           # Reference area
                     span_ref    = b,           # Reference span
                     chord_ref   = c,           # Reference chord
                     name        = ac_name,     # Aircraft name
                     print       = true,        # Prints the results for the entire aircraft
                     print_components = true,   # Prints the results for each component
                    );

## Process data
labels = (collect ∘ keys)(dv_data) # Get aircraft component names from analysis
comp = labels[1]                   # Pick your component
nf, ff, dvs = dv_data[comp];       # Get the nearfield, farfield, and stablity derivative coefficients
print_case(dv_data, comp)          # Pretty-print the results

## Aerodynamic quantities of aircraft
nf_plane, ff_plane, dvs_plane = dv_data[labels[1]]

# Center of pressure
Cm   = nf_plane[5]	# Moment coefficient
CL   = nf_plane[3]	# Lift coefficient
cp   = -c * Cm / CL

# Neutral point
Cm_α = dvs_plane[5,1]
CL_α = dvs_plane[3,1]
np   = -c * Cm_α / CL_α

# Spiral stability
Cl_β = dvs_plane[4,2]
Cl_r = dvs_plane[4,5]
Cn_r = dvs_plane[6,5]
Cn_β = dvs_plane[6,2]

γ = Cl_β * Cn_r / (Cl_r * Cn_β) # Check degree-radian issue

# Locations
x_np = [ ref[1] .+ np; zeros(2) ]	# Neutral point
x_cp = [ ref[1] .+ cp; zeros(2) ]	# Center of pressure

println("Aerodynamic Center x_ac: $(x_w[1]) m")
println("Neutral Point      x_np: $(x_np[1]) m")
println("Center of Pressure x_cp: $(x_cp[1]) m")
println("Spiral Stability      γ: $γ")

## Plotting everything
using Plots
gr(dpi = 300)

#
z_limit = b
aircraft_plot =
plot(xaxis = "x", yaxis = "y", zaxis = "z",
     aspect_ratio = 1,
     camera = (30, 60),
     xlim = (0, z_limit),
    #  ylim = (-z_limit/2, z_limit/2),
     zlim = (-z_limit/2, z_limit/2),
     size = (1280, 720)
    )

plot!(wing_plan, label = "Wing")
plot!(htail_plan, label = "Horizontal Tail")
plot!(vtail_plan, label = "Vertical Tail")

scatter!(tuple(x_w...), color = :blue, label = "Wing MAC")
scatter!(tuple(x_h...), color = :red, label = "Horizontal Tail MAC")
scatter!(tuple(x_v...), color = :green, label = "Vertical Tail MAC")

scatter!(tuple(x_np...), color = :orange, label = "Neutral Point")
scatter!(tuple(x_cp...), color = :brown, label = "Center of Pressure")
# savefig(aircraft_plot, "Aircraft.png")
plot!()


using OrdinaryDiffEq
using Trixi

"""
    initial_condition_jeans_instability(x, t,
                                        equations::Union{CompressibleEulerEquations2D,
                                                         HyperbolicDiffusionEquations2D})

The classical Jeans instability taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
- Dominik Derigs, Andrew R. Winters, Gregor J. Gassner, Stefanie Walch (2016)
  A Novel High-Order, Entropy Stable, 3D AMR MHD Solver with Guaranteed Positive Pressure
  [arXiv: 1605.03572](https://arxiv.org/abs/1605.03572)
- Flash manual https://flash.uchicago.edu/site/flashcode/user_support/flash_ug_devel.pdf
in CGS (centimeter, gram, second) units.
"""
function initial_condition_jeans_instability(x, t,
                                            equations::CompressibleEulerEquations2D)
  # Jeans gravitational instability test case
  # see Derigs et al. https://arxiv.org/abs/1605.03572; Sec. 4.6
  # OBS! this uses cgs (centimeter, gram, second) units
  # periodic boundaries
  # domain size [0,L]^2 depends on the wave number chosen for the perturbation
  # OBS! Be very careful here L must be chosen such that problem is periodic
  # typical final time is T = 5
  # gamma = 5/3
  dens0  = 1.5e7 # g/cm^3
  pres0  = 1.5e7 # dyn/cm^2
  delta0 = 1e-3
  # set wave vector values for pertubation (units 1/cm)
  # see FLASH manual: https://flash.uchicago.edu/site/flashcode/user_support/flash_ug_devel.pdf
  kx = 2.0*pi/0.5 # 2π/λ_x, λ_x = 0.5
  ky = 0.0   # 2π/λ_y, λ_y = 1e10
  k_dot_x = kx*x[1] + ky*x[2]
  # perturb density and pressure away from reference states ρ_0 and p_0
  dens = dens0*(1.0 + delta0*cos(k_dot_x))                 # g/cm^3
  pres = pres0*(1.0 + equations.gamma*delta0*cos(k_dot_x)) # dyn/cm^2
  # flow starts as stationary
  velx = 0.0 # cm/s
  vely = 0.0 # cm/s
  return prim2cons((dens, velx, vely, pres), equations)
end

function initial_condition_jeans_instability(x, t,
                                             equations::HyperbolicDiffusionEquations2D)
  # gravity equation: -Δϕ = -4πGρ
  # Constants taken from the FLASH manual
  # https://flash.uchicago.edu/site/flashcode/user_support/flash_ug_devel.pdf
  rho0   = 1.5e7
  delta0 = 1e-3

  phi = rho0*delta0 # constant background pertubation magnitude
  q1  = 0.0
  q2  = 0.0
  return (phi, q1, q2)
end

initial_condition = initial_condition_jeans_instability

###############################################################################
# semidiscretization of the compressible Euler equations
gamma = 5/3
equations_euler = CompressibleEulerEquations2D(gamma)

polydeg = 3
solver_euler = DGSEM(polydeg, flux_hll)

coordinates_min = (0, 0)
coordinates_max = (1, 1)
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=4,
                n_cells_max=10_000)

semi_euler = SemidiscretizationHyperbolic(mesh, equations_euler, initial_condition, solver_euler)

###############################################################################
# semidiscretization of the hyperbolic diffusion equations
equations_gravity = HyperbolicDiffusionEquations2D()

solver_gravity = DGSEM(polydeg, flux_lax_friedrichs)

semi_gravity = SemidiscretizationHyperbolic(mesh, equations_gravity, initial_condition, solver_gravity,
                                            source_terms=source_terms_harmonic)

###############################################################################
# combining both semidiscretizations for Euler + self-gravity
parameters = ParametersEulerGravity(background_density=1.5e7, # aka rho0
                                    gravitational_constant=6.674e-8, # aka G
                                    cfl=1.6,
                                    resid_tol=1.0e-4,
                                    n_iterations_max=1000,
                                    timestep_gravity=timestep_gravity_carpenter_kennedy_erk54_2N!)

semi = SemidiscretizationEulerGravity(semi_euler, semi_gravity, parameters)

###############################################################################
# ODE solvers, callbacks etc.
tspan = (0.0, 5.0)
ode = semidiscretize(semi, tspan);

summary_callback = SummaryCallback()

stepsize_callback = StepsizeCallback(cfl=1.0)

save_solution = SaveSolutionCallback(interval=10,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

save_restart = SaveRestartCallback(interval=100,
                                   save_final_restart=true)

analysis_interval = 100
alive_callback = AliveCallback(analysis_interval=analysis_interval)

Trixi.pretty_form_utf(::Val{:energy_potential}) = "∑e_potential"
Trixi.pretty_form_ascii(::Val{:energy_potential}) = "e_potential"

function Trixi.analyze(::Val{:energy_potential}, du, u_euler, t, semi::SemidiscretizationEulerGravity)

  u_gravity = Trixi.wrap_array(semi.cache.u_ode, semi.semi_gravity)

  mesh, equations_euler, dg, cache = Trixi.mesh_equations_solver_cache(semi.semi_euler)
  _, equations_gravity, _, _ = Trixi.mesh_equations_solver_cache(semi.semi_gravity)

  e_potential = Trixi.integrate_via_indices(u_euler, mesh, equations_euler, dg, cache, equations_gravity, u_gravity) do u, i, j, element, equations_euler, dg, equations_gravity, u_gravity
    u_euler_local   = Trixi.get_node_vars(u_euler,   equations_euler,   dg, i, j, element)
    u_gravity_local = Trixi.get_node_vars(u_gravity, equations_gravity, dg, i, j, element)
    # OBS! subtraction is specific to Jeans instability test where rho0 = 1.5e7
    return (u_euler_local[1] - 1.5e7) * u_gravity_local[1]
  end
  return e_potential
end

analysis_callback = AnalysisCallback(semi_euler, interval=analysis_interval,
                                     save_analysis=true,
                                     extra_analysis_integrals=(energy_total, energy_kinetic, energy_internal, Val(:energy_potential)))

callbacks = CallbackSet(summary_callback, stepsize_callback,
                        save_restart, save_solution,
                        analysis_callback, alive_callback)

###############################################################################
# run the simulation
sol = solve(ode, CarpenterKennedy2N54(williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary
println("Number of gravity subcycles: ", semi.gravity_counter.ncalls_since_readout)


using MUMPS
if VERSION >= v"0.5.0-dev+7720"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end
include("getDivGrad.jl");
Ns   = (32,48,64)
nrhs = (1,10,100,500)
rhsDensity = 0.2
MUMPStimeSparse =zeros(length(nrhs))
MUMPStimeDense  =zeros(length(nrhs))

for i=1:length(Ns)
	println(@sprintf("Factoring div-grad matrix on %d x %d x %d grid",Ns[i],Ns[i],Ns[i]))
	A = getDivGrad(Ns[i],Ns[i],Ns[i]);
	n = size(A,1);

	#Factor the matrix
	Afac1 = factorMUMPS(A,1)
	Afac2 = factorMUMPS(A,1)

	#rhsInit = spdiagm((ones(n),ones(n-Ns[i]),ones(n-2*Ns[i])),[0,Ns[i],2*Ns[i]],n,n)
	rhsInit = sprandn(n,nrhs[end],rhsDensity)
	for j = 1:length(nrhs)
	  println(@sprintf("Solve with %d rhs",nrhs[j]));
	  #rhs = sprandn(n,nrhs[j],rhsDensity)
	  #rhs = spdiagm((0.8*ones(n-1),1.25*ones(n)),[-1,0],n,n)
	  rhs = rhsInit[:,1:nrhs[j]]

	  # solve using mumps with sparse RHS
	  tic()
	  x1 = applyMUMPS(Afac1,rhs)
	  MUMPStimeSparse[j]= toc();

	  # solve using mumps with dense rhs
	  rhs = full(rhs)
	  tic()
	  x2 = applyMUMPS(Afac2,rhs)
	  MUMPStimeDense[j]= toc();
	  for k = 1:nrhs[j]
	    @test norm(x1[:,k]-x2[:,k])/norm(x1[:,k]) < 1e-9
	  end
        end
	destroyMUMPS(Afac1)
	destroyMUMPS(Afac2)

        println(@sprintf("| %4s | %8s | %8s | %8s |","nrhs","sparseRHS","denseRHS","speedup"))

        for j=1:length(nrhs)
	  println(@sprintf("| %d | %1.5f | %1.5f | %1.5f |",nrhs[j],MUMPStimeSparse[j],MUMPStimeDense[j],MUMPStimeDense[j]/MUMPStimeSparse[j]))
        end
        println(" ")
end


function water_flood(core_props, fluids, rel_perms, core_flood)
    fw, dfw = fractional_flow_function(rel_perms, fluids)
    sw_init = core_flood.initial_water_saturation
    sw_shock = tangent_line_saturation(rel_perms, fluids, (sw_init, fw(sw_init)))
    # println("sw_shock = $sw_shock")
    t_D_BT = 1/dfw(sw_shock) # breakthrough (BT) time [#PV]
    # println("breakthrough time = $t_D_BT")

    # tracer
    sw_tracer = tangent_line_saturation(rel_perms, fluids, (0.0, 0.0))
    t_D_tracer = 1/dfw(sw_tracer)
    xt_tracer_shock = dfw(sw_tracer)
    # construct the recovery factor curve versus the # of PV
    # R = zeros(1)
    # pv_R = zeros(1)
    # at breakthrough
    # push!(R, (1-fw(sw_init))*t_D_BT/(1-sw_init)) # recovery at BT
    # push!(pv_R, t_D_BT) # BT time
    # after breakthrough
    pv_inj = max(core_flood.injected_pore_volume, 2.0) # at least inject 2 pv
    # sw_max = outlet_saturation(pv_inj, sw_shock, dfw, rel_perms)
    # println(sw_max)
    # sw_tmp = linspace(sw_shock, sw_max, 100)
    # t_D_tmp = 1./dfw.(sw_tmp)

    # s_av_tmp = sw_tmp-(fw.(sw_tmp)-1).*t_D_tmp
    # R_tmp = (s_av_tmp-sw_init)/(1-sw_init)
    # append!(R, R_tmp[2:end])
    # append!(pv_R, t_D_tmp[2:end])

    # saturation profile
    sor = rel_perms.sor
    sw_inj = core_flood.injected_water_saturation
    phi = core_props.porosity
    ut = core_flood.injection_velocity
    L = core_props.length
    pv_to_t = phi*L/ut
    s1 = collect(linspace(min(sw_inj, 1-sor-eps()), sw_shock, 500))
    xt_s1 = dfw.(s1)
    xt_shock = dfw(sw_shock)
    xt_prf=[xt_s1; xt_shock+eps(); 2*xt_shock]
    xt_prf[1]=0.0
    sw_prf=[s1; sw_init; sw_init]

    xt_tracer = [0.0, xt_tracer_shock, xt_tracer_shock+eps(), xt_prf[end]]
    c_tracer = [1.0, 1.0, 0.0, 0.0]

    # calculate the pressure drop
    k = core_props.permeability
    krw, kro, dkrw, dkro = rel_perm_functions(rel_perms)
    muo, muw = fluids.oil_viscosity, fluids.water_viscosity

    xt = copy(xt_prf)
    sw = copy(sw_prf)
    i=1
    while(true)
        if i==length(xt)
            break
        elseif xt[i]>=xt[i+1]
            deleteat!(xt, i)
            deleteat!(sw, i)
        else
            i+=1
        end
    end

    x = collect(linspace(0,L,200))
    sw_int = Spline1D(xt, sw, k=1, bc="nearest")
    # print(sw_int)
    # println(ut/phi.*xt)
    # println(sw)
    t_inj=pv_inj*pv_to_t
    t = collect(linspace(0.0,t_inj, 200)) # [s] time
    # println(x/t[2])

    # println(sw_int.(x/t[2]))
    p_inj = zeros(length(t))
    R_int = zeros(length(t))
    WC_int = zeros(length(t))
    p_inj[1]=L*ut./(k*(kro.(sw_init)/muo+krw.(sw_init)/muw))
    WC_int[1] = (krw(sw_init)/muw)/(kro(sw_init)/muo+krw(sw_init)/muw)
    for i in 2:length(t)
        # try
        xt_real = x/t[i]/(ut/phi)
        p_inj[i] = trapz(x, ut./(k*(kro.(sw_int(xt_real))/muo+krw.(sw_int(xt_real))/muw)))
        R_int[i] = (trapz(x, sw_int(xt_real))/L-sw_init)/(1-sw_init)
        WC_int[i] = (krw(sw_int(xt_real[end]))/muw)/(kro(sw_int(xt_real[end]))/muo+krw(sw_int(xt_real[end]))/muw)
        # catch
            # println(x/t[i])
            # error("interpolator at $i !!!")
        # end
    end

    return FracFlowResults([fw], [Line([sw_init, fw(sw_init)], [sw_shock, fw(sw_shock)])],
                            [t/pv_to_t R_int], [t R_int], [xt sw], [xt_tracer c_tracer],
                            [t/pv_to_t p_inj], [t p_inj], [t/pv_to_t WC_int], [t WC_int])
    # return pv_R, R, xt_prf, sw_prf # for the time being to test the code

end


module IsingModelTest

    export IsingModel, Ensemble, Plot2D, Gif2D, GetUpSpins, CentralMoments, EnsProbUpSpins

    using Colors, Plots, StatsBase, Parameters, Distributions, StatsBase, LinearAlgebra;

    """
    Define the Ising model struct.
    """
    @with_kw struct IsingModel
        N::Int64 = 100
        h::Float64
        J::Float64
        p0::Float64 = 0.5 # initial prob of spin up
        periodic::Int64 = 1
        mf::Int64 = 0
        fully_connected::Int64 = 0
    end

    """
    Define the struct for an ensemble of sims.
    """
    @with_kw struct Ensemble
        IM::IsingModel
        e_size::Int64 = 1 # size of ensemble
        method::String = "Glauber" # or "SSA"
        tₘ::Real = 10^5 # no units of time
        Δt::Real = 10^2 # no units of time
        ts::Real = 1.0 # timescale: real time τ = ts*t
        sims = SimRun(method,e_size,tₘ,Δt,ts,IM)
    end

    """
    Choice of algorithm
    """
    function SimRun(alg::String,e_size::Int64,tₘ::Real,Δt::Real,ts::Real,IM::IsingModel)
        if alg == "Glauber"
            return Glauber(e_size,tₘ,Δt,ts,IM)
        elseif alg == "SSA"
            return SSA(e_size,tₘ,Δt,ts,IM)
        else
            error(alg, " is not an appropriate choice of algorithm.
            Please choose either 'Glauber' or 'SSA'.")
        end
    end

    """
    Make the initial NxN matrix of spins.
    """
    function S_make(p0::Float64, N::Int64)
        if p0 > 1 || p0 < 0
            error("p must be bounded between 0 and 1")
        else
            return ones(N,N)-2*sample([1,0],Weights([1-p0,p0]),(N,N))
        end
    end

    """
    Define the effective field at each spin.
    """
    function hₑ(i::Int64, j::Int64, S::Matrix{Float64}, IM::IsingModel, periodic::Int64)
        @unpack J, h, N, fully_connected = IM;
        if fully_connected == 1
            return 2*h + 2*(J/N)*(sum(S)-S[i,j]) # J scaled by N.
        elseif fully_connected == 0 && periodic == 0
            if i==1 && j==1
                return 2*h + 2*J*(S[i+1,j]+S[i,j+1])
            elseif i==N && j==N
                return 2*h + 2*J*(S[i-1,j]+S[i,j-1])
            elseif i==1 && j==N
                return 2*h + 2*J*(S[i+1,j]+S[i,j-1])
            elseif i==N && j==1
                return 2*h + 2*J*(S[i-1,j]+S[i,j+1])
            elseif i==1
                return 2*h + 2*J*(S[i+1,j]+S[i,j+1]+S[i,j-1])
            elseif j==1
                return 2*h + 2*J*(S[i+1,j]+S[i-1,j]+S[i,j+1])
            elseif i==N
                return 2*h + 2*J*(S[i-1,j]+S[i,j+1]+S[i,j-1])
            elseif j==N
                return 2*h + 2*J*(S[i+1,j]+S[i-1,j]+S[i,j-1])
            else
                return 2*h + 2*J*(S[i+1,j]+S[i-1,j]+S[i,j+1]+S[i,j-1])
            end
        else # is periodic BC
            return 2*h + 2*J*(S[mod1(i+1,N),j]+S[mod1(i-1,N),j]+S[i,mod1(j+1,N)]+S[i,mod1(j-1,N)])
        end
    end

    """
    Effective field in MF setting.
    """
    function hₘ(i::Int64, j::Int64, S::Matrix{Float64}, IM::IsingModel, periodic::Int64)
        @unpack J, h, N, fully_connected = IM;
        m = mean(S);
        Jm = J*m;
        if fully_connected == 1
            return 2*h + 2*J*m # J scaled by N.
        elseif fully_connected == 0 && periodic == 0
            if i==1 && j==1
                return 2*h + 2*2*Jm
            elseif i==N && j==N
                return 2*h + 2*2*Jm
            elseif i==1 && j==N
                return 2*h + 2*2*Jm
            elseif i==N && j==1
                return 2*h + 2*2*Jm
            elseif i==1
                return 2*h + 2*3*Jm
            elseif j==1
                return 2*h + 2*3*Jm
            elseif i==N
                return 2*h + 2*3*Jm
            elseif j==N
                return 2*h + 2*3*Jm
            else
                return 2*h + 2*4*Jm
            end
        else # is periodic BC
            return 2*h + 2*4*Jm
        end
    end

    """
    Energy difference calculation
    """
    function ΔH(i::Int64, j::Int64, S::Matrix{Float64}, IM::IsingModel, periodic::Int64=1, mf::Int64=0)
        if mf == 0
            Δe = hₑ(i, j, S, IM, periodic)*S[i,j]
        elseif mf == 1
            Δe = hₘ(i, j, S, IM, periodic)*S[i,j]
        else
            error("mf must be 1 or 0.")
        end
        return Δe
    end

    """
    Logistic rule.
    """
    function LogisticRule(E₀::Float64)
        return (1+exp(E₀))^-1
    end

    """
    Glauber simulation.
    """
    function Glauber(ens_its::Int64,tₘ::Int64,Δt::Int64,μ₁::Float64,IM::IsingModel)
        @unpack p0, N, mf, periodic = IM;
        ens_store = Array{Vector{Matrix{Float64}}}(undef, ens_its);
        for it in 1:ens_its
            S = S_make(p0, N); # make initial config
            S_temp = copy(S); # use S_temp for updating.
            store = Array{Matrix{Float64}}(undef, tₘ÷Δt + 1);
            store[1] = copy(S_temp);
            for n in 1:tₘ+1
                i,j = Int.(map(round,rand(2).*(N-1))) .+ [1,1]; # agent chosen at random
                E₀ = ΔH(i,j,S_temp,IM,periodic,mf) # periodic BC
                pr = LogisticRule(E₀);
                weights = [pr,1-pr]; # flip, don't flip
                flip = sample([1,0],Weights(weights));
                if flip == 1 # if chosen to flip perform
                    S_temp[i,j] = -S_temp[i,j]
                end
                if n%Δt == 0 # store every Δt steps
                    store[(n÷Δt)] = copy(S_temp)
                end
            end
            ens_store[it] = store
        end
        return ens_store
    end

    """
    Propensity for SSA
    """
    function propensity(S::Matrix{Float64},IM::IsingModel,γ::Float64)
        @unpack N, mf, periodic = IM;
        γ = γ/(N^2); # normalise the timescale to match Glauber.
        props = zeros(size(S))
        for i in 1:size(S)[1]
            for j in 1:size(S)[2]
                E₀ = ΔH(i,j,S,IM,periodic,mf); # periodic BC
                pr = LogisticRule(E₀);
                props[i,j] = γ*pr;
            end
        end
        return props
    end

    """
    SSA simulations
    """
    function SSA(ens_its::Int64,τₘ::Real,dτ::Real,γ::Real,IM::IsingModel)
        @unpack p0, N, mf, periodic = IM;
        ens_store = Array{Vector{Matrix{Float64}}}(undef, ens_its);
        for it in 1:ens_its
            S = S_make(p0, N); # make initial config
            S_temp = copy(S); # use S_temp for updating.
            store = Array{Matrix{Float64}}(undef, Int(τₘ÷dτ)+1);
            τ = 0.0;
            sim = 0;
            while τ < τₘ
                props = propensity(S_temp,IM,γ);
                props_flat = vec(transpose(props)); # concatted rows
                w0 = sum(props);
                r1, r2 = rand(2);
                u = (1/w0)*log(1/r1);
                k = findfirst(x -> x>=r2*w0,cumsum(props_flat));
                n = 0;
                while τ+n*dτ < τ+u
                    if τ+n*dτ > τₘ
                        store[Int(τₘ÷dτ)+1] = copy(S_temp)
                        break
                    else
                        store[Int((τ+n*dτ)÷dτ)+1] = copy(S_temp)
                        n+=1;
                    end
                end
                new_i = ((k-1)÷N)+1;
                new_j = k-(new_i-1)*N;
                S_temp[new_i,new_j] = -S_temp[new_i,new_j];
                τ += u;
                if τ > τₘ
                    store[Int(τₘ÷dτ)+1] = copy(S_temp)
                end
                sim += 1;
                if sim%1000==0
                    println(τ)
                end
            end
            ens_store[it] = store;
        end
        return ens_store
    end

    """
    Plot slice at some time.
    """
    function Plot2D(t::Int64, m::Int64, ens::Ensemble) # m = sim set
        @unpack method, Δt, tₘ, ts, sims = ens;
        if method == "Glauber" # then t is timeslice
            plot(Gray.(sims[m][t]))
        elseif method == "SSA" # then t is the actual time
            n = Int((t/ts)÷Δt)
            if n <= tₘ
                plot(Gray.(sims[m][n]))
            else
                error("Plot time must be within the simulated times.")
            end
        else
            error("Simulation method chosen is invalid.")
        end
    end

    """
    IM gif maker.
    """
    function Gif2D(m::Int64, ens::Ensemble, fname::String, frames::Real=40)
        @unpack sims = ens;
        anim = @animate for i ∈ 1:length(sims[m])
            plot(Gray.(sims[m][i]))
        end
        return gif(anim, join([fname,".gif"]), fps = frames)
    end

    """
    Take the trajectories of up spins from each set of sims in an ensemble.
    """
    function GetUpSpins(ens::Ensemble, norm::Int64=0)
        @unpack IM, sims = ens
        @unpack N = IM
        store = Array{Float64,2}(undef, length(sims), length(sims[1]));#Array{Array{Float64,1}}
        for m in 1:length(sims)
            for t in 1:length(sims[1])
                store[m,t] = count(j->(j==1.0),sims[m][t])
            end
        end
        if norm == 0
            return store
        elseif norm == 1
            return store/N^2
        else
            error("norm must take a value of 0 or 1.")
        end
    end

    """
    Get the central moments from the ensemble
    """
    function CentralMoments(ens::Ensemble,n::Int64,norm::Int64=0)
        upspins = GetUpSpins(ens, norm)
        store = Array{Float64,1}(undef, length(upspins[1,:]))
        if n == 1
            for t in 1:length(upspins[1,:])
                store[t] = mean(upspins[:,t])
            end
        else
            for t in 1:length(upspins[1,:])
                store[t] = moment(upspins[:,t],n)
            end
        end
        # need to sort out times
        @unpack Δt, ts = ens;
        no_unit_t = LinRange(0:Δt:length(upspins[1,:])-1);
        times = no_unit_t * ts;
        return (times, store)
    end

    """
    Get prob distribution at all times for an ensemble
    """
    function EnsProbUpSpins(ens::Ensemble, T::Int64)
        UpSpins = GetUpSpins(ens)
        mod_bins = LinRange(minimum(UpSpins[:,T])-0.5:1:maximum(UpSpins[:,T])+0.5);
        mid_pts = LinRange(minimum(UpSpins[:,T]):1:maximum(UpSpins[:,T]));
        bin_vals = normalize(fit(Histogram, UpSpins[:,T], mod_bins), mode=:probability).weights;
        return (mid_pts, bin_vals)
    end

end # module end


const _default_θn_bias = Float64(0.35)
const _default_θe_bias = Float64(0.35)
const _default_θc_bias = Float64(0.35)

struct Summary
    label::Float64
    means::Vector{Float64}
    sd::Vector{Float64}
    members::Int
end

function summarize(data)
    # assumes dimensions by samples format
    d, n = size(data)

    # assumes last row is class label
    k = size(unique(data[end, :]), 1)

    res = Vector{Summary}(undef, k)

    for i in 1:k
        inds = findall(x -> x == Float64(i), data[end, :])
        temp = data[:, inds]

        sd = Vector{Float64}(undef, d-1)
        means = Vector{Float64}(undef, d-1)

        for j in 1:d-1
            means[j] = sum(temp[j, :])/size(temp, 2)
            sd[j] = sqrt(sum((temp[j, :] .- means[j]).^2)/n)
        end

        res[i] = Summary(Float64(i), means, sd, size(temp, 2))
    end

    return res
end

"""
    iso_param(data, summaries) -> Float64, Float64, Float64

    Given data and the summarized data (based on class label),
    attempt to generate ISODATA parameters based on a bias
    percentage.
"""
function iso_param(
	data,
	summaries::Array{Summary};
	_θn_bias::Float64=_default_θn_bias,
	_θe_bias::Float64=_default_θe_bias,
	_θc_bias::Float64=_default_θc_bias
)
	d, n = size(data)
	m = size(summaries, 1)

	θn, θe, θc = 0.0, 0.0, 0.0

	# find θn, minimum number of cluster members
	min_members = 0
	total_members = 0
	for i = 1:m
		if min_members < summaries[i].members
			min_members = summaries[i].members
		end
		total_members += summaries[i].members
	end

	# θn is _θn_bias% lower than min_members/total_members
	θn = (min_members/total_members)*(1-_θn_bias)

	# find θe, maximum standard deviation allowed for a cluster
	max_sd = 0.0
	for i = 1:m
		for j = 1:d
			if summaries[i].sd[j] > max_sd
				max_sd = summaries[i].sd[j]
			end
		end
	end

	# θe is _θe_bias% higher than max_sd
	θe = max_sd*(1+_θe_bias)

	# find θc, minimum distance between clusters
	means = Array{Float64}(undef, d, m)
	for i = 1:m
		means[:, i] = summaries[i].means
	end

	dists = pairwise(Euclidean(), means, means)
	min_dist = maximum(dists)

	for i = 1:m
		for j = 1:m
			if (dists[i, j] < min_dist && i != j)
				min_dist = dists[i, j]
			end
		end
	end

	# θc is _θc_bias% lower than min_dist
	θc = min_dist*(1-_θc_bias)

	return θn, θe, θc
end


# pseudospectral radius and abscissa
using Pseudospectra, Test

@testset "PS radius & abscissa" begin
    A = Pseudospectra.grcar(100)

    opts = Dict{Symbol,Any}()

    # also check that this gets the eigenvalues
    ps_data = new_matrix(A,opts)
    topeig = maximum(imag(ps_data.ps_dict[:ews]))
    @test isapprox(topeig, 2.2617668, atol=1e-6)

    # reference values from Guglielmi & Overton (2012)
    r,zr = psa_radius(A,1e-4,ps_data.ps_dict[:ews])
    @test isapprox(r, 2.85216, atol=1e-5)
    α,za = psa_abscissa(A,1e-4,ps_data.ps_dict[:ews])
    @test isapprox(α, 2.41276, atol=1e-5)

    # Note: if we ever interface to a Hamiltonian eigensolver, we could use
    # the ϵ=1e-8 case in S&P.
end