#

using LinearAlgebra
using Gadfly
Gadfly.set_default_plot_size(35cm,20cm)

p1 = 0.7 + 0.7*im

b1 = 0.4 + 0.2*im
b2 = 0.3 + 0.9*im
b3 = 0.8 + 0.1*im

r1 = 100*norm(b1-p1)
r2 = 100*norm(b2-p1)
r3 = 100*norm(b3-p1)

N = 50000

# random range samples
Cc1 = (rand(Complex{Float64},N) .- b1)*100
Cc2 = (rand(Complex{Float64},N) .- b2)*100
Cc3 = (rand(Complex{Float64},N) .- b3)*100

A = zeros(100,100)

kr1 = 1 ./ (abs.(r1 .- norm.(Cc1)) )
kr2 = 1 ./ (abs.(r2 .- norm.(Cc2)) )
kr3 = 1 ./ (abs.(r3 .- norm.(Cc3)) )

kr1 .= exp.(kr1) .- 1
kr2 .= exp.(kr2) .- 1
kr3 .= exp.(kr3) .- 1

mask1 = rand(N) .< kr1
mask2 = rand(N) .< kr2
mask3 = rand(N) .< kr3

X1 = Cc1[mask1] .|> real .|> x->round(Int,x)
Y1 = Cc1[mask1] .|> imag .|> x->round(Int,x)
X1 .+= 100*real(b1)
Y1 .+= 100*imag(b1)
X1[X1 .== 0] .= 1
Y1[Y1 .== 0] .= 1
for i in 1:sum(mask1)
  A[X1[i],Y1[i]] += 1
end

X2 = Cc2[mask2] .|> real .|> x->round(Int,x)
Y2 = Cc2[mask2] .|> imag .|> x->round(Int,x)
X2 .+= 100*real(b2)
Y2 .+= 100*imag(b2)
X2[X2 .== 0] .= 1
Y2[Y2 .== 0] .= 1
for i in 1:sum(mask2)
  A[X2[i],Y2[i]] += 1
end

X3 = Cc3[mask3] .|> real .|> x->round(Int,x)
Y3 = Cc3[mask3] .|> imag .|> x->round(Int,x)
X3 .+= 100*real(b3)
Y3 .+= 100*imag(b3)
X3[X3 .== 0] .= 1
Y3[Y3 .== 0] .= 1
for i in 1:sum(mask3)
  A[X3[i],Y3[i]] += 1
end

pl = spy(A)

#


#ENV["CPLEX_STUDIO_BINARIES"]="/home/software/cplex/128/cplex/bin/x86-64_linux"
using DelimitedFiles, LinearAlgebra
include("src/core.jl")

function getProblemDataBeasley(pathName)
    testData=readdlm(pathName)
    n=testData[1,1]
    μ=testData[2:n+1, 1]
    μ=convert(Array{Float64}, μ)
    σ=testData[2:n+1, 2]
    Σ=zeros(n,n)
    for iter=n+2:size(testData,1)
        # This loop inserts diagonal elements twice. This is ok as benign, and a "fix" would require more lines of code.
        Σ[testData[iter, 1], testData[iter, 2]]=σ[testData[iter, 1]]*σ[testData[iter, 2]]*testData[iter, 3]
        Σ[testData[iter, 2], testData[iter, 1]]=σ[testData[iter, 1]]*σ[testData[iter, 2]]*testData[iter, 3]
    end
    X=sqrt(Σ)
    Y=(X*X')\(X*μ)
    d=X'*Y-μ
    d=round.(d, digits=8) #avoids round-off errors when Σ is full rank
    A=zeros(1,size(X,2))
    l=zeros(size(A,1))
    u=zeros(size(A,1))
    min_investment=zeros(n)
    f=size(X, 2)
    m=size(A, 1)
    k=5
    γ=(100/sqrt(n))*ones(n)

    return SparsePortfolioData(μ, Y, d, X, A, l, u, k, n, m, f, min_investment, γ)
end

for portfolioNumber=1:5
    pathName = "/Users/Ryancw2/Dropbox (MIT)/Sparse Portfolios Project/OptimalMarkowitzPortfolios.jl/data/port"
    pathName=pathName*string(portfolioNumber)*".txt"
    thePortfolioData=getProblemDataBeasley(pathName)

    sparse_Portfolio=cutting_planes_portfolios(thePortfolioData; useSOCPLB=false, useHeuristic=false);
    @show @test(abs.(sum(sparse_Portfolio.w)-1.0)<1e-4)
    println("Variance is:", sparse_Portfolio.portfolioVariance)
    println("Return is:", sparse_Portfolio.expectedReturn)
    println("Reg term is:", 0.5*sum(sparse_Portfolio.w[i]^2/thePortfolioData.γ[sparse_Portfolio.indices[i]] for i=1:size(sparse_Portfolio.indices,1)))
    open("results/Table3ResultsOutput.csv","a") do fp
    println(fp, "Instance, cardinality, regularizer, method, solvetime, nodecount, iters")
    println(fp, "beasley"*string(portfolioNumber), ",", thePortfolioData.k, ",",
                     "100/sqrt(n)", ",", "OA", ",",  sparse_Portfolio.solveTime, ",", sparse_Portfolio.nodeCount, ",", sparse_Portfolio.numIters)
    end

    # # Test against Big-M for correctness
    sparse_Portfolio_Cplex_bigM=cplex_raw_bigM(thePortfolioData, 300.0, 1e-8, 1)
    @show @test(abs.(sum(sparse_Portfolio_Cplex_bigM.w)-1.0)<1e-4)
    println("Variance is:", sparse_Portfolio_Cplex_bigM.portfolioVariance)
    println("Return is:", sparse_Portfolio_Cplex_bigM.expectedReturn)
    println("Reg term is:", 0.5*sum(sparse_Portfolio_Cplex_bigM.w[i]^2/thePortfolioData.γ[sparse_Portfolio_Cplex_bigM.indices[i]] for i=1:size(sparse_Portfolio_Cplex_bigM.indices,1)))
    open("results/Table3ResultsOutput.csv","a") do fp
    println(fp, "Instance, cardinality, regularizer, method, solvetime, nodecount")
    println(fp, "beasley"*string(portfolioNumber), ",", thePortfolioData.k, ",",
                     "100/sqrt(n)", ",", "CPLEX Big-M", ",",  sparse_Portfolio_Cplex_bigM.solveTime, ",", sparse_Portfolio_Cplex_bigM.nodeCount)
    end

    # # Test against MISOCP for correctness
    sparse_Portfolio_Cplex=cplex_raw_MISOCP(thePortfolioData, 300.0, 1e-8, 1)
    @show @test(abs.(sum(sparse_Portfolio_Cplex.w)-1.0)<1e-4)
    println("Variance is:", sparse_Portfolio_Cplex.portfolioVariance)
    println("Return is:", sparse_Portfolio_Cplex.expectedReturn)
    println("Reg term is:", 0.5*sum(sparse_Portfolio_Cplex.w[i]^2/thePortfolioData.γ[sparse_Portfolio_Cplex.indices[i]] for i=1:size(sparse_Portfolio_Cplex.indices,1)))
    open("results/Table3ResultsOutput.csv","a") do fp
    println(fp, "Instance, cardinality, regularizer, method, solvetime, nodecount")
    println(fp, "beasley"*string(portfolioNumber), ",", thePortfolioData.k, ",",
                    "100/sqrt(n)", ",", "CPLEX MISOCP", ",",  sparse_Portfolio_Cplex.solveTime, ",", sparse_Portfolio_Cplex.nodeCount)
    end
end


let
	# if maxlayers is redefined, must remake the calibration cache or run without it (set usecache=false)
	global const maxlayers = 50			# 15-100 (model gets unstable if too high relative to time step, see stabilitycheck())

	secondsperyear	= 365*24*60*60		# [s/year]
	k				= 1.5/1e4			# thermal diffusivity of deep ocean [m^2/s]  (default 1.5 cm^2/s)
	w				= 4/secondsperyear	# mean ocean upwelling rate [m/s]  (default 4 m/yr)
	polarwaterratio	= 0.2				# ratio of Δtemp of downwelling polar water to mean Δtemp of surface ocean [unitless]
	mixeddepth		= 70.0				# depth of mixed layer [m]
	rho				= 1026				# density of seawater [kg/m3]
	c				= 3996				# specific heat capacity of seawater [J/kgK]
	oceanarea		= 0.71				# fraction of earth's surface covered by ocean [unitless]
	a_L_O_exch		= 0.31
	b_L_O_exch		= 1.59
	marinewarming 	= 1.3				# marine surface air warming enhancement [unitless]
	oceandepth		= 3688				# average depth of the ocean [m]

	zstep			= (oceandepth - mixeddepth)/maxlayers	# vertical length of each discrete layer below surface layer [m]
	midlayers 		= 2:maxlayers-1

	function derivative(T, step, position)
		if position == 1
			# This one-sided three point approximation is more accurate than (T[2] - T[1])/step.
			# It is O(step^2) instead of O(step).
			(-3*T[1] + 4*T[2] - T[3]) / (2*step)
		elseif position == length(T)
			# ditto
			(3*T[end] - 4*T[end-1] + T[end-2]) / (2*step)
		else
			# This central difference is actually also a three point approximation, and also O(step^2).
			(T[position+1] - T[position-1]) / (2*step)
		end
	end

	function secondderivative(T, step, position)
		# The position check is quite expensive, so let's comment it out and make sure we only call this in the interior.
		# If at the extremes, calculate from position one step inward.
		#position = clamp.(position, 2, length(T)-1)

		# This central three point approximation is also correct for the extreme left or right (if shifted inward one step).
		# It is O(step^2) in the interior and O(step) at the extremes.
		(T[position-1] - 2*T[position] + T[position+1]) / step^2
	end

	# Calculate ocean temperature using an upwelling-diffusion energy balance model.
	# Solve the advection-diffusion differential equation using explicit Euler finite differences.
	# The discretization is forward-in-time, central-in-space (FTCS).
	# The function does only one iteration of the time step but calculates temperature in all ocean layers (vectorized).
	global function temperatures!(s::ClimateState, p::ClimateParams)
		@unpack TotalRadiativeForcing, Temp, Temp_land, Temp_global = s
		@unpack timestep, lambda = p

		lambda_L = lambda
		lambda_O = lambda/marinewarming
		k_landocean = 1 # + b_L_O_exch - a_L_O_exch*lambda_L		#  land-ocean heat exchange coefficient [(W/m^2)/K]

		# We need some spatial derivatives of the temperature profile
		dz_Temp_1 = derivative(Temp, zstep, 1)
		dz_Temp = derivative.(Ref(Temp), zstep, midlayers)
		dz2_Temp = secondderivative.(Ref(Temp), zstep, midlayers)

		# First boundary condition (type Dirichlet):
		# Temp_mixed [°C]: temperature (change since preindustrial times) of mixed layer (i.e. surface ocean)
		Temp_mixed = Temp[1]

		# Temp_land [°C]: land temperature (change since preindustrial times)
		Temp_land = (TotalRadiativeForcing + k_landocean/(1-oceanarea)*marinewarming*Temp_mixed) /
						(1/lambda_L + k_landocean/(1-oceanarea))

		# heat transfer from land to ocean [W/m^2]
		Heatflux_land2ocean = k_landocean * (Temp_land - marinewarming*Temp_mixed)

		# heat transfer from mixed layer to deep ocean [W/m^2]
		Heatflux_mixed2deep = -rho*c*k * dz_Temp_1 - rho*c*w * (1-polarwaterratio) * Temp_mixed
		#Heatflux_mixed2deep = -rho*c*k * dz_Temp_1 - rho*c*w * (Temp[2] - polarwaterratio*Temp_mixed)

		# time derivative of temperature [°C/year]
		# timestep unit is [per year], so all time derivatives are also [per year]
		dt_Temp = Vector{Float64}(undef, maxlayers)

		# Energy balance for the mixed layer
		dt_Temp[1] = secondsperyear / (rho*mixeddepth*c) *
			(TotalRadiativeForcing - Temp_mixed/lambda_O - Heatflux_mixed2deep + Heatflux_land2ocean/oceanarea)

		# Main advection-diffusion differential equation:
		# We could speed things up a tiny bit by assigning Temp[midlayers+1] and Temp[midlayers-1] to variables, but meh.
		dt_Temp[midlayers] = secondsperyear * (k * dz2_Temp + w * dz_Temp)

		# Temp [°C]: temperature change since preindustrial times in each ocean layer
		Temp[1:maxlayers-1] += timestep * dt_Temp[1:maxlayers-1]

		# Second boundary condition (type Robin): approximate the derivative to get an equation for Temp[max]
		dz_Temp_max = w / k * (polarwaterratio*Temp_mixed - Temp[maxlayers])
		Temp[maxlayers] = Temp[maxlayers-1] + zstep * dz_Temp_max
		# more accurate approximation below, but accuracy is not important at ocean bottom
		#Temp[maxlayers] = -4*Temp[maxlayers-1] + Temp[maxlayers-2] + 2*3*zstep * dz_Temp_maxlayers

		# Temp_global [°C]: average land & ocean temperature change since preindustrial times
		Temp_global = marinewarming*oceanarea*Temp[1] + (1-oceanarea)*Temp_land

		@pack! s = Temp, Temp_land, Temp_global
	end

	global function init_temperatures!(s::ClimateState)
		Temp = zeros(maxlayers)
		Temp_global = 0.0
		Temp_land = 0.0
		@pack! s = Temp, Temp_global, Temp_land
	end

	global function stabilitycheck(timestep)
		# Courant number, see p56 of Kajishima (2017), Computational Fluid Dynamics
		courant = w * timestep * secondsperyear / zstep

		# cell Reynolds number, see Kajishima p59-60
		reynolds = w * zstep / k

		# stability relation
		@assert courant <= min(reynolds/2, 2/reynolds) "Instable discretization of temperature differential equation, " *
			"try reducing timestep or increasing maxlayers."
	end
end


"""
    Network

Tools for creating axial (i.e., 1d) thermoacoustic network models.
The formulation is Riemann-based:
p=F exp(+iω*l/c)+G exp(-iω*l/c)
Au=A/ρc[ F exp(+iω*l/c)-G exp(-iω*l/c)]
"""
module Network
using WavesAndEigenvalues.NLEVP
export discretize
include("NLEVP_exports.jl")

"""
    out=duct(l,c,A,ρ=1.4*101325/c^2)

Create local matrix coefficients for duct element.

# Arguments
- `l`: length of the duct
- `c`: speed of sound
- `A` cross-sectional area of the duct
- `ρ` density
"""
function duct(l,c,A,ρ=1.4*101325/c^2)
    M=[-1      -1;        #p_previous-p_current=0
        -A/(ρ*c) A/(ρ*c); #A u_previous-A u_current=0
        0          0;     #p_current-p_next=0
        0          0]     #A u_current-A u_next=0

    M31=[0        0;
        0         0;
        1         0;
        0         0]

    M32= [0        0;
        0         0;
        0         1;
        0         0]
    M41= [0        0;
        0        0;
        0         0;
        1         0]
    M42= [0        0;
        0         0;
        0         0;
        0         -1]
    out=[[M,(),()],
         [M31, (generate_exp_az(1im*l/c),), ((:ω,),),],
         [M32, (generate_exp_az(-1im*l/c),), ((:ω,),),],
         [M41*A/(ρ*c), (generate_exp_az(1im*l/c),), ((:ω,),),],
         [M42*A/(ρ*c), (generate_exp_az(-1im*l/c),), ((:ω,),),],
         ]

    return out
end

"""
    out=terminal(R,c,A,ρ=1.4*101325/c^2;init=true)

Create local matrix coefficients for terminal element.

# Arguments
- `R`: reflection coefficient
- `c`: speed of sound
- `A`: cross-sectional area of the terminal
- `ρ`: density
- `init`: (optional) if true its the left end else it is the right one.
"""
function terminal(R,c,A,ρ=1.4*101325/c^2;init=true)
    #if hasmethod(R,(ComplexF64,Int))
    if typeof(R)<:Number
    else
        println("Error, unknown impedance type $(typeof(Z))")
    end

    if init
         M=[+R      -1;
           +1       +1;
           1*A/(ρ*c)  -1*A/(ρ*c)]

    else
        M=[ -1       -1;
            -1*A/(ρ*c)  +1*A/(ρ*c);
          -1     +R]
    end
    return [[M,(),()],]
end

"""
    out=flame(c1,c2,A,ρ=1.4*101325/c1^2)

Create local matrix coefficients for flame element with n-tau dynamics.

# Arguments
- `c1`: speed of sound on the unburnt side
- `c2`: speed of sound on the burnt side
- `A`: cross-sectional area of the terminal
- `ρ`: density
"""
function flame(c1,c2,A,ρ=1.4*101325/c1^2)
    M= [0        0;
        0        0;
        0         0;
        1         -1]

    out=duct(0,c1,A,ρ)
    push!(out,[M*(c2^2/c1^2-1)*A/(ρ*c1), (pow1,exp_delay), ((:n,),(:ω,:τ),),])
    return out
end

"""
    out=helmholtz(V,l_n,d_n,c,A,ρ=1.4*101325/c1^2)

Create local matrix coefficients for a sidewall Helmholtz damper element.

# Arguments
- `V`: Volume of the damper
- `l_n`: length of the damper neck
- `d_n`: diameter of the damper neck
- `c`: speed of sound
- `A`: cross-sectional area of the terminal
- `ρ`: density

# Notes:

The model is build on  the transfer matrix from [1]

# References

[1] F. P. Mechel, Formulas of Acoustic, Springer, 2004, p. 728
"""
function helmholtz(V,l_n,d_n,c,A,ρ=1.4*101325/c^2)
    #=====

    p_u=p_d
    u_u=1/Z p_d+u_d <==> A u_u=A/Z p_d+A u_d>

    TODO:
    <==> Z*A*u_u=A*p_d+Z*Au_d (roe easier representation)

    where

    Z(ω)=ρ[ω^2/(π*c)*[2-r_n/r_u]+0.425*M*c*/S_n+1im[ω*l/S_n-c^2/(ω*V)]]

    where
    ρ: density
    c: speed of sound
    r_n: neck radius of the Helmholtz damper
    r_u: radius of the main duct
    M: Mach number
    S_n = π*r_n2 : crossectional area of the neck
    V = : volume of the helmholtz damper
    and
    l=l_n+t_w+0.85 r_n*(2-r_n/r_u)

    with tw: wall thickness?

    ## adaption to Riemann invariants

    p_u-A_+*exp(im*ω*l/c)+A_-*exp(-im*ω*l/c)=0
    u_u-A_+/ρc*exp(im*ω*l/c)+A_+/ρc*exp(-im*ω*l/c)=0
    The last equation is to be multiplied by A.

    ## notes
    this element does not resolve what is between the usptream and the downstream
    site. It directly describes the jump!
    =====#

    r_n=d_n/2
    r_u=sqrt(A/pi)
    S_n=pi*r_n^2
    t_w=0
    l=l_n+t_w+0.85*r_n*(2-r_n/r_u)
    M=0
    M21= [0        0;
        -1       -1;
        0         0;
        0         0]
    #M22= [0        0;
    #    0         -1;
    #    0         0;
    #    0         0]

    function impedance(ω::ComplexF64,k::Int)::ComplexF64
        let ρ=ρ, r_n=r_n, r_u=r_u, M=M, c=c, l=l, V=V ,S_n=S_n
            return ρ*(pow2(ω,k)/(π*c)*(2-r_n/r_u)+pow0(ω,k)*0.425*M*c/S_n
            +1.0im*pow1(ω,k)*l/S_n-1.0im*c^2/V*pow(ω,k,-1))
        end
    end

    function admittance(ω::ComplexF64,k::Int)::ComplexF64
        let Z=impedance
            if k==0
                return 1/Z(ω,0)
            elseif k==1
                return -Z(ω,1)/Z(ω,0)^2 #TODO: implement arbitray order chain rule
            else
                return NaN+NaN*im
            end
        end
    end
    function admittance(ω::Symbol)::String
        return "1/Z($ω)"
    end

    out=duct(0,c,A,ρ)
    #push!(out,[M21, (admittance,), ((:ω,),),])
    push!(out,[-M21/ρ, (admittance,), ((:ω,),),])
    return out
end

"""
    sidewallimp(imp,c,A,ρ=1.4*101325/c^2)

Use function predefined  function `imp` to prescripe a frequency-dependent
side wall impedance.
"""
function sidewallimp(imp,c,A,ρ=1.4*101325/c1^2)
    M21= [0        0;
        -1       -1;
        0         0;
        0         0]
    function admittance(ω::ComplexF64,k::Int)::ComplexF64
        let Z=imp
            if k==0
                return 1/Z(ω,0)
            elseif k==1
                return -Z(ω,1)/Z(ω,0)^2 #TODO: implement arbitray order chain rule
            else
                return NaN+NaN*im
            end
        end
    end
    function admittance(ω::Symbol)::String
        return "1/Z($ω)"
    end

    out=duct(0,c,A,ρ)
    #push!(out,[M21*A*(ρ*c), (admittance,), ((:ω,),),])
    push!(out,[M21, (admittance,), ((:ω,),),])
end
"""
    out=lhr(V,l_n,d_n,c,A,ρ=1.4*101325/c^2)

Create linear Helmholtz model according to [1].

#Reference
[1] Nonlinear effects in acoustic metamaterial based on a cylindrical pipe with
    ordered Helmholtz resonators,J. Lan, Y. Li, H. Yu, B. Li, and X. Liu, Phys.
    Rev. Lett. A., 2017, [doi:10.1016/j.physleta.2017.01.036](https://doi.org/10.1016/j.physleta.2017.01.036)
"""
function lhr(V,l_n,d_n,c,A,ρ=1.4*101325/c^2)
    r_n=d_n/2
    S_n=pi*r_n^2
    B0=ρ*c^2
    η=1.5E-5 #dynamic viscosity or air in m^2/s
    R_vis=ρ*l_n/r_n*sqrt(η/2)*S_n# *ω^(1/2) #!!!! wrong in Lan it must be divided not miultiplied by r_n
    R_rad=1/4*ρ*r_n^2/c*S_n#*ω^2
    l=l_n+1.7*r_n #length correction
    Cm=V/(ρ*c^2*S_n^2)
    Mm=ρ*l*S_n
    #δ=Rm/Mm
    ω0=1/(sqrt(Cm*Mm))
    println("M: $Mm, C: $Cm, freq: $ω0")
    C=B0*S_n/(1im*ω0^2*V)/(S_n)  #last division to convert between impedance and flow impedance
    function impedance(ω::ComplexF64,k::Int)::ComplexF64
        let C=C, R_vis=R_vis, R_rad=R_rad, ω0=ω0
            return C*pow1(ω,k)-C*ω0^2*pow(ω,k,-1)-C*1im*R_vis/Mm*pow(ω,k,1/2)-C*1im*R_rad/Mm*pow2(ω,k)
        end
    end

    return sidewallimp(impedance,c,A,ρ)
end

"""
    L=discretize(network)

Discretize network model.

# Arguments
- `network`

# Returns
- `L::LinearOperatorFamily`

# The network is specified of a list of network elements.
elements together with there options. Available elements are.
- `:duct`: duct element with options `(l,c,A,ρ)`
- `:flame`: n-tau-flame element with options `(c1,c2,A,ρ)`
- `:helmholtz`:Helmholtz damper element with options `(V,l_n,d_n,c,A,ρ)`
- `:unode`: sound hard boundary, i.e. a velocity node with options `(c,A,ρ)`
- `:pnode`: sound soft boundary, i.e., a pressure node with options `(c,A,ρ)`
- `:anechoic`: anechoic boundary condition with option `(c,A,ρ)`

For all elements the density `ρ` is an optional parameter which is computed
as `ρ=1.4*101325/c^2` if unspecified (air at atmospheric pressure).

# Example
This is an example for the discretization of a Rijke tube, with a Helmholtz
damper, a velocity node at the inlet and a pressure node at the outlet:
    network=[(:unode,(347.0, 0.01), ),
            (:duct,(0.25, 347.0, 0.01)),
            (:flame,(347,347*2,0.01)),
            (:duct,(0.25, 347.0*2, 0.01)),
            (:helmholtz,(0.02^3,0.01,.005,347*2,0.01)),
            (:duct,(0.25, 347.0*2, 0.01)),
            (:pnode,(347.0*2, 0.01), )
            ]

    L=discretize(network)
    L.params[n]=1
    L.params[τ]=0.001
"""
function discretize(network)
    i,j=1,1
    L=LinearOperatorFamily((:ω,))
    N=length(network)
    A=zeros(2*N,2*N)
    for (idx,(elmt, data)) in enumerate(network)
        if elmt==:duct
            terms=duct(data...)
        elseif elmt==:unode
            R=1
            if idx==1
                init=true
            elseif idx==length(network)
                init=false
            else
                println("Error, terminal element at intermediate position $idx!")
            end
            #terms=velocity_node(data...)
            terms=terminal(R,data...,init=init)
        elseif elmt==:pnode
            R=-1
            if idx==1
                init=true
            elseif idx==length(network)
                init=false
            else
                println("Error, terminal element at intermediate position $idx!")
            end
            terms=terminal(R,data...,init=init)
        elseif elmt==:anechoic
            R=0
            if idx==1
                init=true
            elseif idx==length(network)
                init=false
            else
                println("Error, terminal element at intermediate position $idx!")
            end
            terms=terminal(R,data...,init=init)
        elseif elmt==:pnodeend
            terms=pressure_node_end(data...)
        elseif elmt==:vnodestart
            terms=velocity_node_start(data...)
        elseif elmt==:flame
            terms=flame(data...)
        elseif elmt==:helmholtz
            terms=helmholtz(data...)
        elseif elmt==:helmholtz2
            terms=lhr(data...)
        else
            println("Warning: unknown element $elmt")
            continue
        end
        I,J=size(terms[1][1])

        for (coeff,func,arg) in terms
            M=deepcopy(A)
            M[i:i+I-1,j:j+J-1]=coeff
            push!(L,Term(M,func,arg,"M"))
        end
        i+=I-2
        j+=2
    end
    return L
end
end #module network


function quaternion_map(ω, timestep)
    return Quaternion(sqrt(4 / timestep^2 - dot(ω, ω)), ω)
end

function quaternion_map_jacobian(ω::SVector{3}, timestep)
    msq = -sqrt(4 / timestep^2 - dot(ω, ω))
    return [ω' / msq; I]
end

function cayley(ω)
    Quaternion(1.0 / sqrt(1.0 + ω' * ω) * [1.0; ω])
end

function cayley_jacobian(ω)
    ω₁, ω₂, ω₃ = ω
    a = (sqrt(1.0 + ω₁^2 + ω₂^2 + ω₃^2)^-3)
    b = sqrt(1.0 + ω₁^2 + ω₂^2 + ω₃^2)^-1
    SMatrix{4,3}([
                 -ω₁*a -ω₂*a -ω₃*a;
                 (b - (ω₁^2)*a) (-ω₁ * ω₂ * a) (-ω₁ * ω₃ * a);
                 (-ω₁ * ω₂ * a) (b - (ω₂^2)*a) (-ω₂ * ω₃ * a);
                 (-ω₁ * ω₃ * a) (-ω₂ * ω₃ * a) (b - (ω₃^2)*a);
                 ])
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

# flip arguments so that points always come first
evaluate(d::PreMetric, g::Geometry, p::Point) = evaluate(d, p, g)

"""
    evaluate(Euclidean(), point, line)

Evaluate the Euclidean distance between `point` and `line`.
"""
function evaluate(::Euclidean, p::Point, l::Line)
  a, b = points(l)
  u = p - a
  v = b - a
  α = (u ⋅ v) / (v ⋅ v)
  norm(u - α*v)
end

"""
    evaluate(::PreMetric, point1, point2)

Evaluate pre-metric between coordinates of `point2` and `point2`.
"""
evaluate(d::PreMetric, p1::Point, p2::Point) =
  evaluate(d, coordinates(p1), coordinates(p2))


using QuantumLab
using Test
using Dates
using LinearAlgebra

INFO(str) = @info "$(Dates.format(now(),dateformat"Y/m/d HH:MM:SS,sss"))  $str"

if (!@isdefined(indent))
  indent = ""
end

INFO("$(indent)READING:   h2o.xyz -> h2o::Geometry ")
h2o = readGeometryXYZ("h2o.xyz")

INFO("$(indent)OBTAINING:   BasisSetExchange -> sto3g::BasisSet")
sto3g = BasisSet("sto-3g")

INFO("$(indent)COMPUTING:   sto3g,h2o -> bas")
bas = computeBasis(sto3g,h2o)

INFO("$(indent)COMPUTING:   bas,h2o -> matrixOverlap, matrixKinetic, matrixNuclearAttraction, ERIs")
matrixOverlap = computeMatrixOverlap(bas)
matrixKinetic = computeMatrixKinetic(bas)
matrixNuclearAttraction = computeMatrixNuclearAttraction(bas,h2o)
ERIs = computeTensorElectronRepulsionIntegrals(bas)

INFO("$(indent)COMPUTING:   h2o -> matrixSADguess")
matrixSADguess = computeDensityGuessSAD("HF","STO-3G",h2o)

INFO("$(indent)COMPUTING:   sto3g, h2o -> shells, shells_native")
shells = computeBasisShellsLibInt2(sto3g,h2o)
shells_native = computeBasisShells(sto3g,h2o)

INFO("$(indent)COMPUTING:   (HartreeFock)  shells, h2o, matrixSADguess -> density")
density = evaluateSCF(shells,h2o,sum(matrixSADguess)/2,info=false,detailedinfo=false)[3]

INFO("$(indent)COMPUTING:   density, matrixKinetic, matrixNuclearAttraction, ERIs -> matrixFock")
matrixFock = computeMatrixFock(density,matrixKinetic,matrixNuclearAttraction,ERIs)

densityvirt = inv(matrixOverlap) - density

INFO("$(indent)COMPUTING:   shells, density -> S, T, J, K, Fock")
S = computeMatrixOverlap(shells)
T = computeMatrixKinetic(shells)
J = computeMatrixCoulomb(shells,density)
K = computeMatrixExchange(shells,density)
V = computeMatrixNuclearAttraction(shells,h2o)
fock = computeMatrixFock(T,V,J,K)
moenergies = eigvals(Symmetric(fock),Symmetric(S))


using CSV, DataFrames, Statistics
using AxisKeys, MonteCarloMeasurements, DimensionalData
using StatsBase, NamedTupleTools
using StanSample

df = CSV.read(joinpath(@__DIR__, "..", "..", "..", "data", "WaffleDivorce.csv"), DataFrame);
df.D = zscore(df.Divorce)
df.M = zscore(df.Marriage)
df.A = zscore(df.MedianAgeMarriage)
data = (N=size(df, 1), D=df.D, A=df.A, M= df.M)

stan5_1 = "
data {
    int < lower = 1 > N; // Sample size
    vector[N] D; // Outcome
    vector[N] A; // Predictor
}
parameters {
    real a; // Intercept
    real bA; // Slope (regression coefficients)
    real < lower = 0 > sigma;    // Error SD
}
transformed parameters {
    vector[N] mu;               // mu is a vector
    for (i in 1:N)
        mu[i] = a + bA * A[i];
}
model {
    a ~ normal(0, 0.2);         //Priors
    bA ~ normal(0, 0.5);
    sigma ~ exponential(1);
    D ~ normal(mu , sigma);     // Likelihood
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(D[i] | mu[i], sigma);
}
";

m5_1s = SampleModel("m5.1s", stan5_1)
rc5_1s = stan_sample(m5_1s; data)

if success(rc5_1s)

    nt5_1s = read_samples(m5_1s, :particles)
    NamedTupleTools.select(nt5_1s, (:a, :bA, :sigma)) |> display

    ka_chns = read_samples(m5_1s, :keyedarray)
    ka_chns |> display

    da_chns = read_samples(m5_1s, :dimarray)
    da_chns |> display
end


export  PlaneStressIncompressibleRivlinSaunders, getStress

@doc """
Pascon, João Paulo.
"Large deformation analysis of plane-stress hyperelastic problems via triangular membrane finite elements."
International Journal of Advanced Structural Engineering (2019): 1-20.
"""->
mutable struct PlaneStressIncompressibleRivlinSaunders
    ρ::Float64 # density
    C1::Float64 # Young's modulus
    C2::Float64 # Poisson's ratio

    # The classical incompressible PlaneStressIncompressibleRivlinSaunders strain energy function
    # I3 = 0
    # W = C1(I1 - 3) + C2(I2 - 2)

    σ0::Array{Float64} # stress at last time step
    σ0_::Array{Float64} # σ0 to be updated in `commitHistory`
    ε0::Array{Float64} # stress at last time step
    ε0_::Array{Float64} # σ0 to be updated in `commitHistory`
end

function PlaneStressIncompressibleRivlinSaundersJacobian(E11::T, E22::T, G12::T, C1::S, C2::S) where {S<:Number, T<:Number}
    f = x->PlaneStressIncompressibleRivlinSaundersStress(x[1],x[2],x[3], C1, C2)
    J = ForwardDiff.jacobian(f, [E11;E22;G12])
    σ = f([E11;E22;G12])
    σ, J
end

function PlaneStressIncompressibleRivlinSaundersStress(E11::T, E22::T, G12::T, C1::S, C2::S) where {S<:Number, T<:Number}
    C11, C22, C12 = 2*E11+1.0, 2*E22+1.0, G12
    det22C = C11*C22 - C12*C12
    C33 = 1/det22C

    I1 = C11 + C22 + C33

    C = [ C11    C12      0
          C12    C22      0
          0       0     C33]

    Cinv = [ C22/det22C    -C12/det22C        0
            -C12/det22C     C11/det22C        0
                  0             0         1.0/C33]
    svol = Cinv
    Siso = 2.0*C1*UniformScaling(1.0) + 2.0*C2*(C + UniformScaling(I1))

    p = -C33*Siso[3,3]

    σ = zeros(T, 3)
    σ[1] = Siso[1,1] + p*svol[1,1]
    σ[2] = Siso[2,2] + p*svol[2,2]
    σ[3] = Siso[1,2] + p*svol[1,2]
    return σ
end

function PlaneStressIncompressibleRivlinSaunders(prop::Dict{String, Any})
    C1 = prop["C1"]; C2 = prop["C2"]; ρ = prop["rho"];
    σ0 = zeros(3); σ0_ = zeros(3); ε0 = zeros(3); ε0_ = zeros(3)
    PlaneStressIncompressibleRivlinSaunders(ρ, C1, C2, σ0, σ0_, ε0, ε0_)
end

function getStress(self::PlaneStressIncompressibleRivlinSaunders,  strain::Array{Float64},  Dstrain::Array{Float64}, Δt::Float64 = 0.0)
    local dΔσdΔε
    C1, C2 = self.C1, self.C2
    E11, E22, G12 = strain[1], strain[2], strain[3]
    σ, dΔσdΔε = PlaneStressIncompressibleRivlinSaundersJacobian(E11, E22, G12, C1, C2)

    self.σ0_ = σ[:]
    self.ε0_ = strain

    #@show σ, dΔσdΔε
    return σ, dΔσdΔε
end

function getTangent(self::PlaneStressIncompressibleRivlinSaunders)
    error("Not implemented")
end

function commitHistory(self::PlaneStressIncompressibleRivlinSaunders)
    self.σ0 = self.σ0_
    self.ε0 = self.ε0_
end

# # W = C1(I1 - 3) + C2(I2 - 2)
# #
# # E11, E22, G12 = strain[1], strain[2], strain[3]
# # E12 = G12/2.0
# #
# #      E11 E12 0
# # E =  E21 E22 0
# #      0   0   E33
# # incompressibility
# # det(2E + 1) = det(C) = det(F^T F) = 1
# # E33 = (1.0/(1 + 2*E11 + G12 + 4*E11*E22 - G12*G12) - 1.0) /2.0
# # C = 2E + 1
# # I1(C) = lambda1 + lambda2 + lambda3 = 3 + 2(E11 + E22 + E33)
# # I2(C) = lambda1*lambda2 + lambda2*lambda3 + lambda3*lambda1
# #       = (1 + 2E11)*(1 + 2E22) + (1 + 2E22)*(1 + 2E33) + (1 + 2E33)*(1 + 2E11) - G12 * G21

# using SymPy
# E11, E22, G12, C1, C2 = @vars E11 E22 G12 C1 C2
# E33 = (1.0/(1 + 2*E11 + G12 + 4*E11*E22 - G12*G12) - 1.0) /2.0
# I1 = 3 + 2*(E11 + E22 + E33)
# I2 = 3 + 4*(E11 + E22 + E33) + 4*E11*E22 + 4*E11*E22 + 4*E11*E33 - G12*G12
# W = C1*(I1 - 3) + C2*(I2 - 3)
# σ = Array{Sym}(undef,3)
# σ[1] = simplify(diff(W, E11))
# σ[2] = simplify(diff(W, E22))
# σ[3] = simplify(diff(W, G12))
# J = Matrix{Sym}(undef, 3, 3)
# for i = 1:3
#     J[i,1] = simplify(diff(σ[i], E11))
#     J[i,2] = simplify(diff(σ[i], E22))
#     J[i,3] = simplify(diff(σ[i], G12))
# end
# s = "J = zeros(3,3);\n"
# for i = 1:3
#     for j = 1:3
#         global s *= "J[$i,$j]="*sympy.julia_code(J[i,j])*"\n"
#     end
# end
# t = "σ = zeros(3)\n"
# for i = 1:3
#     global t *= "σ[$i]="*sympy.julia_code(σ[i])*"\n"
# end


using CSV, StateSpaceModels, Plots, Statistics, Dates

# Load the AirPassengers dataset
AP = CSV.read("AirPassengers.csv")

# Take the log of the series
logAP = log.(Vector{Float64}(AP.Passengers))

# Plot the data
p1 = plot(AP.Date, logAP, label = "Log-airline passengers", legend = :topleft, color = :black)

# Specify the state-space model
model = structural(logAP, 12)

# Estimate the state-space model
ss = statespace(model)

# Analyze its decomposition in trend, slope and seasonal components
p2 = plot(AP.Date, [ss.smoother.alpha[:, 1] ss.smoother.alpha[:, 2] ss.smoother.alpha[:, 3]], layout = (3, 1),
            label = ["Trend" "Slope" "Seasonal"], legend = :topleft)

# Forecast 24 months ahead
N = 24
pred, dist = forecast(ss, N)

# Define forecasting dates
firstdate = AP.Date[end] + Month(1)
newdates = collect(firstdate:Month(1):firstdate + Month(N - 1))

p3 = plot!(p1, newdates, pred, label = "Forecast")

# 10 scenarios for 24 months ahead
sim = simulate(ss, N, 10)
p4 = plot!(p1, newdates, sim[:, 1, :], label = "", color = "Grey", width = 0.2)


abstract type SSM end

struct LGSSM{T <: Real} <: SSM
    A::AbstractArray{T} # State matrix
    B::AbstractArray{T} # Measurement matrix
    R::AbstractArray{T} # State Cholesky
    S::LowerTriangular{T, Array{T,2}} # Measurement Cholesky
    Σᵥ::Matrix{T} # Initial state covariance
    # μᵥ::Vector{T} # Initial state mean
end


function construct_model(t_train, t_test, u0_2, u0_3; hidden_dim = 10, test_model = true)
    ann = Chain(Dense(4, hidden_dim, relu), Dense(hidden_dim, 1))
    p1, re = Flux.destructure(ann)
    p2 = Float64[0.5, 0.03]
    p3 = [p1; p2]

    ps = Flux.params(p3)

    u0 = Float64[25.0e6, u0_2, u0_3, 10]

    function QSIR(du, u, p, t)
        β = abs(p[length(p1) + 1])
        γ = 0.007627
        du[1]=  -β * u[1] * u[2] / u0[1]
        du[2] = β * u[1] * u[2] / u0[1] - γ * u[2] - abs(re(p[1:length(p1)])(u)[1]) * u[2] / u0[1]
        du[3] = γ * u[2]
        du[4] = abs(re(p[1:length(p1)])(u)[1] * u[2] / u0[1])
    end

    tspan = (1.0, maximum(t_test))
    prob = ODEProblem(QSIR, u0, tspan, p3)

    if test_model
        sol = Array(concrete_solve(prob, Rosenbrock23(autodiff = false), u0, p3, saveat=t_train))
    end

    return ps, re, u0, prob, p3, length(p1)
end


using JuMP
using DisjunctiveProgramming

m = Model()
@variable(m, -5 ≤ x ≤ 10)
@disjunction(
    m,
    0 ≤ x ≤ 3,
    5 ≤ x ≤ 9,
    reformulation=:big_m,
    name=:y
)
@proposition(m, y[1] ∨ y[2]) #this is a redundant proposition

print(m)

# ┌ Warning: disj_y[1] : x in [0.0, 3.0] uses the `MOI.Interval` set. Each instance of the interval set has been split into two constraints, one for each bound.
# ┌ Warning: disj_y[2] : x in [5.0, 9.0] uses the `MOI.Interval` set. Each instance of the interval set has been split into two constraints, one for each bound.
# Feasibility
# Subject to
#  XOR(disj_y) : y[1] + y[2] == 1.0         <- XOR constraint
#  y[1] ∨ y[2] : y[1] + y[2] >= 1.0         <- reformulated logical proposition (name is the proposition)
#  disj_y[1,lb] : -x + 5 y[1] <= 5.0       <- left-side of constraint in 1st disjunct (name is assigned to disj_y[1][lb])
#  disj_y[1,ub] : x + 7 y[1] <= 10.0       <- right-side of constraint in 1st disjunct (name is assigned to disj_y[1][ub])
#  disj_y[2,lb] : -x + 10 y[2] <= 5.0      <- left-side of constraint in 2nd disjunct (name is assigned to disj_y[2][lb])
#  disj_y[2,ub] : x + y[2] <= 10.0         <- right-side of constraint in 2nd disjunct (name is assigned to disj_y[2][ub])
#  x >= -5.0                                <- variable lower bound
#  x <= 10.0                                <- variable upper bound
#  y[1] >= 0.0                              <- lower bound on binary
#  y[2] >= 0.0                              <- lower bound on binary
#  y[1] <= 1.0                              <- upper bound on binary
#  y[2] <= 1.0                              <- upper bound on binary
#  y[1] binary                              <- indicator variable (1st disjunct) is binary
#  y[2] binary                              <- indicator variable (2nd disjunct) is binary


function ricker(f₀::Float64, Δt::Float64)
    @assert f₀ > 0 "Peak frequency f₀ must be positive!"
    @assert Δt > 0 "Time interval Δt must be positive!"

    # signal length
    length = 2.0 / f₀ / Δt

    # number of points in each direction
    n = ceil(Int, length / 2)

    # time axis
    t = Δt * (-n:n)

    # amplitude
    x = (π * f₀ * t) .^ 2
    A = (1 .- 2x) .* exp.(-x)
end

function d2t(lyr::Dict{Symbol,Vector{Float64}})
    vₚ = lyr[:vp]
    d = lyr[:dh]
    t = cumsum(2d ./ vₚ)
    pushfirst!(t,0)
end

function d2t(lyr::Vector{Dict{Symbol,Vector{Float64}}})
    stacked = stack(lyr)
    d2t(stacked)
end

function seqcounts(labels::Vector)
    # find sequence boundaries
    indices = findall(diff(labels) .!= 0)

    # get sequence of unique labels
    sequence = labels[indices]

    # get label counts
    push!(sequence, labels[end])
    pushfirst!(indices, 0)
    push!(indices, length(labels))
    counts = diff(indices)

    sequence, Int64.(counts)
end


## Creates a simple Euler--Maruyama discretization of a diffusion in which the
## drift coefficient corresponds to the Lorenz96 model and the diffusion
## coefficient is constant.

module Lorenz96

using SequentialMonteCarlo
using RNGPool
using StaticArrays
import SMCExamples.Particles.MVFloat64Particle
using Random

struct LorenzTheta
  σ::Float64
  F::Float64
  δ::Float64
  Δ::Float64
  steps::Int64
end

struct LorenzScratch{d}
  t1::MVector{d, Float64}
  t2::MVector{d, Float64}
end
LorenzScratch{d}() where d = LorenzScratch{d}(MVector{d, Float64}(undef),
  MVector{d, Float64}(undef))

@inline function lorenz(F::Float64, x::MVector{d, Float64},
  z::MVector{d, Float64}) where d
  @inbounds z[1] = (x[2] - x[d-1]) * x[d]
  @inbounds z[2] = (x[3] - x[d]) * x[1]
  for i = 3:d-1
    @inbounds z[i] = (x[i+1] - x[i-2]) * x[i-1]
  end
  @inbounds z[d] = (x[1] - x[d-2]) * x[d-1]
  z .+= F
end

function makeLorenzModel(theta::LorenzTheta,
  ys::Vector{SVector{d, Float64}}) where d
  n::Int64 = length(ys)
  σ::Float64 = theta.σ
  F::Float64 = theta.F
  invδ²over2 = 0.5 / (theta.δ * theta.δ)
  logncG = - 0.5 * d * log(2 * π * theta.δ * theta.δ)
  steps::Float64 = theta.steps
  Δ::Float64 = theta.Δ
  h::Float64 = Δ/steps
  sqrth::Float64 = sqrt(Δ/steps)

  @inline function lG(p::Int64, particle::MVFloat64Particle{d}, ::LorenzScratch{d})
    r::Float64 = 0.0
    for i = 1:d
      @inbounds v::Float64 = particle.x[i] - ys[p][i]
      r += logncG - v * invδ²over2 * v
    end
    return r
  end
  ## Euler--Maruyama
  @inline function M!(newParticle::MVFloat64Particle{d}, rng::RNG, p::Int64,
    particle::MVFloat64Particle{d}, scratch::LorenzScratch{d})
    if p == 1
      randn!(rng, newParticle.x)
      newParticle.x .*= σ
    else
      scratch.t1 .= particle.x
      for i = 1:steps
        lorenz(F, scratch.t1, scratch.t2)
        scratch.t2 .*= h
        randn!(rng, scratch.t1)
        scratch.t1 .*= σ * sqrth
        scratch.t1 .+= scratch.t2
      end
      newParticle.x .= scratch.t1
    end
  end
  return SMCModel(M!, lG, length(ys), MVFloat64Particle{d}, LorenzScratch{d})
end

function simulateLorenzModel(theta::LorenzTheta, d::Int64, n::Int64)
  model = makeLorenzModel(theta, Vector{SVector{d, Float64}}(undef, 0))
  ys = Vector{SVector{d, Float64}}(undef, n)
  xParticle = MVFloat64Particle{d}()
  xScratch = LorenzScratch{d}()
  rng = getRNG()
  for p in 1:n
    model.M!(xParticle, rng, p, xParticle, xScratch)
    ys[p] = xParticle.x + theta.δ * randn(rng, d)
  end
  return ys
end

function defaultLorenzModel(d::Int64, n::Int64)
  theta = LorenzTheta(0.5, 8.0, 0.7, 0.05, 20)
  ys = simulateLorenzModel(theta, d, n)
  model = makeLorenzModel(theta, ys)
  return model, theta, ys
end

end


@testset "$(rpad("Runge-Kutta Tableaus",80))" begin

    # explicit midpoint
    a = Rational{Int64}[[0    0]
                        [1//2 0]]
    b = Rational{Int64}[0, 1   ]
    c = Rational{Int64}[0, 1//2]
    o = 2

    @test TableauERK(:explicit_midpoint, o, a, b, c) == getTableauExplicitMidpoint()

    tab_explicit_midpoint1 = @test_nowarn TableauERK(:explicit_midpoint, o, a, b, c)
    tab_explicit_midpoint2 = @test_logs (:warn, r"Initializing TableauDIRK with explicit tableau explicit_midpoint.*") TableauDIRK(:explicit_midpoint, o, a, b, c)
    tab_explicit_midpoint3 = @test_logs (:warn, r"Initializing TableauFIRK with explicit tableau explicit_midpoint.*") TableauFIRK(:explicit_midpoint, o, a, b, c)

    @test tab_explicit_midpoint1 == tab_explicit_midpoint2
    @test tab_explicit_midpoint2 == tab_explicit_midpoint3
    @test tab_explicit_midpoint3 == tab_explicit_midpoint1

    @test !isequal(tab_explicit_midpoint1, tab_explicit_midpoint2)
    @test !isequal(tab_explicit_midpoint2, tab_explicit_midpoint3)
    @test !isequal(tab_explicit_midpoint3, tab_explicit_midpoint1)

    tab_explicit_midpoint1 = TableauERK(:explicit_midpoint, o, a, b, c)
    tab_explicit_midpoint2 = TableauERK(:explicit_midpoint, o, a, b, c)
    @test tab_explicit_midpoint1 == tab_explicit_midpoint2
    @test isequal(tab_explicit_midpoint1, tab_explicit_midpoint2)

    tmp = mktempdir()
    tab_explicit_midpoint1 = TableauERK(:explicit_midpoint, o, a, b, c)
    writeTableauToFile(tmp, tab_explicit_midpoint1)
    tab_explicit_midpoint2 = readTableauERKFromFile(tmp, "explicit_midpoint")
    rm(tmp, recursive=true)
    @test tab_explicit_midpoint1 == tab_explicit_midpoint2
    @test !isequal(tab_explicit_midpoint1, tab_explicit_midpoint2) # fails because data arrays in tab_explicit_midpoint1 are of type Rational{Int64} while data arrays in tab_explicit_midpoint2 are of type Float

    # Heun
    a = [[0.0 0.0]
         [1.0 0.0]]
    b = [0.5, 0.5]
    c = [0.0, 1.0]
    o = 2

    @test_nowarn tab_explicit_heun = TableauERK(:heun, o, a, b, c)
    @test_logs (:warn, r"Initializing TableauDIRK with explicit tableau heun.*") TableauDIRK(:heun, o, a, b, c)
    @test_logs (:warn, r"Initializing TableauFIRK with explicit tableau heun.*") TableauFIRK(:heun, o, a, b, c)

    @test TableauERK(:heun, o, a, b, c) == getTableauHeun()

    tmp = mktempdir()
    tab_explicit_heun1 = TableauERK(:heun, o, a, b, c)
    writeTableauToFile(tmp, tab_explicit_heun1)
    tab_explicit_heun2 = readTableauERKFromFile(tmp, "heun")
    rm(tmp, recursive=true)
    @test tab_explicit_heun1 == tab_explicit_heun2
    @test isequal(tab_explicit_heun1, tab_explicit_heun2)

    # Crouzeix
    fac = 0.5/√3
    a = [[ 0.5+fac 0.0    ]
         [-2.0*fac 0.5+fac]]
    b = [0.5,     0.5    ]
    c = [0.5+fac, 0.5-fac]
    o = 2

    @test TableauDIRK(:crouzeix, o, a, b, c) == getTableauCrouzeix()

    @test_throws AssertionError tab_explicit_crouzeix = TableauERK(:crouzeix, o, a, b, c)
    @test_nowarn tab_explicit_crouzeix = TableauDIRK(:crouzeix, o, a, b, c)
    @test_logs (:warn, r"Initializing TableauFIRK with diagonally implicit tableau crouzeix.*") tab_explicit_crouzeix = TableauFIRK(:crouzeix, o, a, b, c)

    # implicit midpoint
    a = ones(Rational{Int64}, 1, 1) // 2
    b = Rational{Int64}[1   ]
    c = Rational{Int64}[1//2]
    o = 2

    @test TableauFIRK(:implicit_midpoint, o, a, b, c) == getTableauImplicitMidpoint()
    @test TableauFIRK(:implicit_midpoint, o, a, b, c) == getTableauGLRK(1)

    @test_throws AssertionError tab_implicit_midpoint = TableauERK(:implicit_midpoint, o, a, b, c)
    @test_throws AssertionError tab_implicit_midpoint = TableauDIRK(:implicit_midpoint, o, a, b, c)
    tab_implicit_midpoint = TableauFIRK(:implicit_midpoint, o, a, b, c)

    # Gauss-Legendre
    a = Array{Float64}(@dec128 [
         [0.25       0.25-√3/6]
         [0.25+√3/6  0.25     ]
        ])
    b = Array{Float64}(@dec128 [0.5,      0.5     ])
    c = Array{Float64}(@dec128 [0.5-√3/6, 0.5+√3/6])
    o = 4

    @test TableauFIRK(:glrk2, o, a, b, c) == getTableauGLRK(2)

    @test_throws AssertionError tab_explicit_glrk2 = TableauERK(:glrk2, o, a, b, c)
    @test_throws AssertionError tab_explicit_glrk2 = TableauDIRK(:glrk2, o, a, b, c)
    tab_explicit_glrk2 = TableauFIRK(:glrk2, o, a, b, c)

    # instantiate all explicit tableaus
    @test typeof(getTableauExplicitEuler()) <: TableauERK
    @test typeof(getTableauExplicitMidpoint()) <: TableauERK
    @test typeof(getTableauHeun()) <: TableauERK
    @test typeof(getTableauKutta()) <: TableauERK
    @test typeof(getTableauERK4()) <: TableauERK
    @test typeof(getTableauERK438()) <: TableauERK

    # instantiate all diagonally implicit tableaus
    @test typeof(getTableauCrouzeix()) <: TableauDIRK

    # instantiate all fully implicit tableaus
    @test typeof(getTableauImplicitEuler()) <: TableauFIRK
    @test typeof(getTableauImplicitMidpoint()) <: TableauFIRK
    @test typeof(getTableauGLRK(1)) <: TableauFIRK
    @test typeof(getTableauGLRK(2)) <: TableauFIRK

    # instatiate all partitioned tableaus
    @test typeof(getTableauSymplecticEulerA()) <: TableauEPRK
    @test typeof(getTableauSymplecticEulerB()) <: TableauEPRK

    # test instatiation of partioned tableau by composition of two RK tableaus
    @test typeof(TableauEPRK(:PERK4, 4, getTableauERK4().q, getTableauERK4().q)) <: TableauEPRK
    @test TableauEPRK(:PERK4, 4, getTableauERK4().q, getTableauERK4().q) == TableauEPRK(:PERK4, 4, getTableauERK4().q)

    # TODO Add tests for TableauIPRK, TableauSARK, TableauSPARK and TableauGLM.

end


#Load some packages and add workers
using Images
using Distributed

using LinearAlgebra
using Statistics

#Load the package on all workers
addprocs(4)
@everywhere using DPMMSubClusters

#Load image, taken from SINTEL http://sintel.is.tue.mpg.de/
img = load("frame_0001.png")

#Change to channel view
img_channels = channelview(img)
z,x,y = size(img_channels)

#Create input
input_arr = zeros(5,x*y)
for i=1:x
    for j=1:y
        input_arr[:,(i-1)*y+j]=[img_channels[1,i,j],img_channels[2,i,j],img_channels[3,i,j],i,j]
    end
end

#Create HyperParams
#The rgb,xy multiplier allows us to both play with the weight of the xy/rgb
rgb_prior_multiplier = 1
xy_prior_multiplier = 0.1

data_cov = cov(input_arr')
data_cov[4:5,1:3] .= 0
data_cov[1:3,4:5] .= 0

data_cov[1:3,1:3] .*= rgb_prior_multiplier
data_cov[4:5,4:5] .*= xy_prior_multiplier

data_mean = mean(input_arr,dims = 2)[:]

hyper_params = DPMMSubClusters.niw_hyperparams(1.0,
           data_mean,
           8,
           data_cov)

#Run the model
labels,clusters,weights = DPMMSubClusters.fit(input_arr,hyper_params,50000.0,iters = 200, verbose = true)

#Get the cluster color means
color_means = [x.μ[1:3] for x in clusters]

segemnated_image = zeros(3,x,y)
for i=1:x
    for j=1:y
        segemnated_image[:,i,j] = color_means[labels[(i-1)*y+j]]
    end
end
segemnated_image = colorview(RGB,segemnated_image)


"""
    curl(dx::AbstractVector, A::AbstractArray)

Calculate 2nd order cell-centered ∇×A where `A` is a 4D array of size (3, nx, ny, nz) and
`dx` is a vector of grid intervals in each dimension.
"""
function curl(dx::AbstractVector, A::AbstractArray{T,N}) where {T,N}
   @assert N == 4 && length(dx) == 3 "Input vector shall be indexed in 3D!"

   @views Ax, Ay, Az = A[1,:,:,:], A[2,:,:,:], A[3,:,:,:]

   B = zeros(T, size(A))
   @views Bx, By, Bz = B[1,:,:,:], B[2,:,:,:], B[3,:,:,:]

   invdx = inv.(dx)

   if any(==(1), size(A)) # 2D
      if size(A,4) == 1
         @inbounds for j in 2:size(A,3)-1, i in 2:size(A,2)-1
            ∂Ay∂x = (Ay[i-1,j,1] - 2Ay[i,j,1] + Ay[i+1,j,1]) * invdx[1]
            ∂Az∂x = (Az[i-1,j,1] - 2Az[i,j,1] + Az[i+1,j,1]) * invdx[1]
            ∂Ax∂y = (Ax[i-1,j,1] - 2Ax[i,j,1] + Ax[i+1,j,1]) * invdx[2]
            ∂Az∂y = (Az[i-1,j,1] - 2Az[i,j,1] + Az[i+1,j,1]) * invdx[2]
            ∂Ax∂z = (Ax[i-1,j,1] - 2Ax[i,j,1] + Ax[i+1,j,1]) * invdx[3]
            ∂Ay∂z = (Ay[i-1,j,1] - 2Ay[i,j,1] + Ay[i+1,j,1]) * invdx[3]

            Bx[i,j,1] = ∂Az∂y - ∂Ay∂z
            By[i,j,1] = ∂Ax∂z - ∂Az∂x
            Bz[i,j,1] = ∂Ay∂x - ∂Ax∂y
         end
      elseif size(A,3) == 1
         @inbounds for k in 2:size(A,4)-1, i in 2:size(A,2)-1
            ∂Ay∂x = (Ay[i-1,1,k] - 2Ay[i,1,k] + Ay[i+1,1,k]) * invdx[1]
            ∂Az∂x = (Az[i-1,1,k] - 2Az[i,1,k] + Az[i+1,1,k]) * invdx[1]
            ∂Ax∂y = (Ax[i-1,1,k] - 2Ax[i,1,k] + Ax[i+1,1,k]) * invdx[2]
            ∂Az∂y = (Az[i-1,1,k] - 2Az[i,1,k] + Az[i+1,1,k]) * invdx[2]
            ∂Ax∂z = (Ax[i-1,1,k] - 2Ax[i,1,k] + Ax[i+1,1,k]) * invdx[3]
            ∂Ay∂z = (Ay[i-1,1,k] - 2Ay[i,1,k] + Ay[i+1,1,k]) * invdx[3]

            Bx[i,1,k] = ∂Az∂y - ∂Ay∂z
            By[i,1,k] = ∂Ax∂z - ∂Az∂x
            Bz[i,1,k] = ∂Ay∂x - ∂Ax∂y
         end
      end
   else # 3D
      @inbounds for k in 2:size(A,4)-1, j in 2:size(A,3)-1, i in 2:size(A,2)-1
         ∂Ay∂x = (Ay[i-1,j,k] - 2Ay[i,j,k] + Ay[i+1,j,k]) * invdx[1]
         ∂Az∂x = (Az[i-1,j,k] - 2Az[i,j,k] + Az[i+1,j,k]) * invdx[1]
         ∂Ax∂y = (Ax[i-1,j,k] - 2Ax[i,j,k] + Ax[i+1,j,k]) * invdx[2]
         ∂Az∂y = (Az[i-1,j,k] - 2Az[i,j,k] + Az[i+1,j,k]) * invdx[2]
         ∂Ax∂z = (Ax[i-1,j,k] - 2Ax[i,j,k] + Ax[i+1,j,k]) * invdx[3]
         ∂Ay∂z = (Ay[i-1,j,k] - 2Ay[i,j,k] + Ay[i+1,j,k]) * invdx[3]

         Bx[i,j,k] = ∂Az∂y - ∂Ay∂z
         By[i,j,k] = ∂Ax∂z - ∂Az∂x
         Bz[i,j,k] = ∂Ay∂x - ∂Ax∂y
      end
   end
   B
end


immutable HyperboloidIntercept{T} <: Model
    β₀::T
    a::T
    α::T
    μ::T
end

function Base.show(io::IO, ::Type{HyperboloidIntercept})
    return print(io, "Intercept Hyperboloid")
end

function Base.show(io::IO, m::HyperboloidIntercept)
    @printf(
        io,
        "Intercept Hyperboloid (β₀ = %.3f, a = %.3f, α = %.3f, μ = %.3f)",
        m.β₀,
        m.a,
        m.α,
        m.μ,
    )
end

function initial_parameters(::Type{HyperboloidIntercept})
    return Float64[0.0, log(1.0), log(1.0), log(1.0)]
end

canonical_parameters(m::HyperboloidIntercept) = (m.a, m.α, m.μ, m.β₀)

function HyperboloidIntercept(parameters::Vector)
    β₀ = parameters[1]
    a = exp(parameters[2])
    α = exp(parameters[3])
    μ = exp(parameters[4])
    return HyperboloidIntercept(β₀, a, α, μ)
end

@inline function predict(
    m::HyperboloidIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, α, μ = m.β₀, m.a, m.α, m.μ
    u1 = x1 * (1 + α * t1)^(-μ)
    u2 = x2 * (1 + α * t2)^(-μ)
    z = β₀ + a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::HyperboloidIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, α, μ = m.β₀, m.a, m.α, m.μ
    u1 = x1 * (1 + α * t1)^(-μ)
    u2 = x2 * (1 + α * t2)^(-μ)
    gr = (
        1.0,
        (u2 - u1) * a,
        a * (x2 * -μ * (1 + α * t2)^(-μ - 1) * t2 - x1 * -μ * (1 + α * t1)^(-μ - 1) * t1) * α,
        a * (-x2 * (1 + α * t2)^(-μ) * log(1 + α * t2) + x1 * (1 + α * t1)^(-μ) * log(1 + α * t1)) * μ,
    )
    return gr
end


mutable struct MultiLogitProblem{Ty, TX, Tp} <: NUTSGibbs.AbstractProblem
    y::Ty
    X::TX
    prior::Tp
end

function (problem::MultiLogitProblem)(theta)
    @unpack y, X, prior = problem
    @unpack xi = theta

    N, p = size(X)
    G = size(xi, 1)
    ll = 0.0

    for i in 1:N
        tmp = zeros(eltype(xi), G+1)
        ##tmp = 1.0

        for g in 1:G
            xx = dot(X[i,:], xi[g,:])
            tmp[g] = xx
            ##tmp += exp(xx)
            ll += y[i,g]*xx
        end

        ll -= logsumexp(tmp)
        ##ll -= log(tmp)
    end

    for g in 1:G
        ll += sum(logpdf.(prior, xi[g,:]))
    end

    isfinite(ll) || return -Inf

    ll
end

function mlogit_transformation(p::MultiLogitProblem)
    D = size(p.X,2)
    G = size(p.y,2)
    as((xi = as(Array, G-1, D),))
end

function nuts_mlogit_init(y, X, prior, warmup)
    problem = MultiLogitProblem(y, X, prior)
    NUTSGibbs.init(problem, mlogit_transformation, warmup)
end

function nuts_mlogit_step!(x::NUTSGibbsSampler, y, xi)
    x.problem.y = y
    return NUTSGibbs.step(x, (xi=xi,))
end


import LinearAlgebra.logdet

struct LogDetAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int,Int}

    function LogDetAtom(x::AbstractExpr)
        children = (x,)
        return new(:logdet, hash(children), children, (1, 1))
    end
end

function sign(x::LogDetAtom)
    return NoSign()
end

function monotonicity(x::LogDetAtom)
    return (NoMonotonicity(),)
end

function curvature(x::LogDetAtom)
    return ConcaveVexity()
end

function evaluate(x::LogDetAtom)
    return log(det(evaluate(x.children[1])))
end

function conic_form!(x::LogDetAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        A = x.children[1]
        D = Variable(size(A)) # diagonal matrix
        U = Variable(size(A)) # upper triangular matrix

        # objective given by the sum of the log of diagonal matrix D
        objective = conic_form!(sum(log(diag(D))), unique_conic_forms)

        # enforce D and U to be diagonal and upper triangular
        for i in 1:A.size[1]
            for j in 1:A.size[2]
                if i != j
                    conic_form!(D[i, j] == 0, unique_conic_forms)
                end
                if i > j
                    conic_form!(U[i, j] == 0, unique_conic_forms)
                end
            end
        end

        # diagonals of D and U are the same
        conic_form!(diag(D) == diag(U), unique_conic_forms)

        # A and [D U; U' A] need to be positive semidefinite
        conic_form!(isposdef(A), unique_conic_forms)
        conic_form!(isposdef([D U; U' A]), unique_conic_forms)

        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

logdet(x::AbstractExpr) = LogDetAtom(x)


module _AttnEn
export AttnEn
using StatsBase: Histogram, fit
    """
        Av4, (Hxx,Hnn,Hxn,Hnx) = AttnEn(Sig)

    Returns the attention entropy (`Av4`) calculated as the average of the
    sub-entropies (`Hxx`,`Hxn`,`Hnn`,`Hnx`) estimated from the data sequence
    (`Sig`) using a base-2 logarithm.

        Av4, (Hxx, Hnn, Hxn, Hnx) = AttnEn(Sig::AbstractArray{T,1} where T<:Real; Logx::Real=2)

    Returns the attention entropy (`Av4`) and the sub-entropies (`Hxx`,`Hnn`,`Hxn`,`Hnx`)
    from the data sequence (`Sig`) where,
    Hxx:    entropy of local-maxima intervals
    Hnn:    entropy of local minima intervals
    Hxn:    entropy of intervals between local maxima and subsequent minima
    Hnx:    entropy of intervals between local minima and subsequent maxima

    # Arguments:
    `Logx`  - Logarithm base, a positive scalar
              (Enter 0 for natural logarithm)

    See also `EnofEn`, `SpecEn`, `XSpecEn`, `PermEn`, `MSEn`

    # References:
        [1] Jiawei Yang, et al.,
            "Classification of Interbeat Interval Time-series Using
            Attention Entropy."
            IEEE Transactions on Affective Computing
            (2020)

    """
    function AttnEn(Sig::AbstractArray{T,1} where T<:Real; Logx::Real=2)

    (Logx == 0) ? Logx = exp(1) : nothing
    N = size(Sig,1)
    (N > 10) ? nothing : error("Sig:   must be a numeric vector")
    (Logx>0) ? nothing : error("Logx:  must be a positive number > 0")

    Xmax = PkFind(Sig)
    Xmin = PkFind(-Sig)
    Txx = diff(Xmax)
    Tnn = diff(Xmin)
    Temp = diff(sort(vcat(Xmax, Xmin)))

    if isempty(Xmax)
        error("No local maxima found!")
    elseif isempty(Xmin)
        error("No local minima found!")
    end

    (Xmax[1]<Xmin[1]) ? (Txn = Temp[1:2:end]; Tnx = Temp[2:2:end]) :
            (Txn = Temp[2:2:end]; Tnx = Temp[1:2:end])

    Edges = -0.5:N
    Pnx = fit(Histogram,Tnx,Edges).weights
    Pnn = fit(Histogram,Tnn,Edges).weights
    Pxx = fit(Histogram,Txx,Edges).weights
    Pxn = fit(Histogram,Txn,Edges).weights

    Pnx = Pnx[Pnx.!=0]/size(Tnx,1)
    Pxn = Pxn[Pxn.!=0]/size(Txn,1)
    Pnn = Pnn[Pnn.!=0]/size(Tnn,1)
    Pxx = Pxx[Pxx.!=0]/size(Txx,1)

    Hxx = -sum(Pxx.*(log.(Logx,Pxx)))
    Hxn = -sum(Pxn.*(log.(Logx,Pxn)))
    Hnx = -sum(Pnx.*(log.(Logx,Pnx)))
    Hnn = -sum(Pnn.*(log.(Logx,Pnn)))
    Av4 = (Hnn + Hxx + Hxn + Hnx)/4

    return Av4, (Hxx,Hnn,Hxn,Hnx)
    end

    function PkFind(X)
        Nx = size(X,1)
        Indx = zeros(Int,Nx);
        for n = 2:Nx-1
            if X[n-1]< X[n] > X[n+1]
                Indx[n] = n

            elseif X[n-1] < X[n] == X[n+1]
                k = 1
                while (n+k)<Nx && X[n] == X[n+k]
                    k +=1
                end
                if X[n] > X[n+k]
                    Indx[n] = n + floor((k-1)/2)
                end
            end
        end
        Indx = Indx[Indx.!==0]
    return Indx
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

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

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

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

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


##################### negative logarithm of density ##############################
function nld(G::Univariate, x::Number)
    ((x-mean(G))^2/var(G) + log(var(G)) + log(2π))/2
end

function nld(G::Multivariate, x::AbstractVector)
    x = difference(x, mean(G))
	nld(cov(G), x)
end

function nld(Σ::AbstractMatOrFacOrUni, x::AbstractVector)
    Σ = factorize(Σ) # since factorize on a Inverse is a no-op, saves unnecessary factorization in laplace's method
    Λ = inverse(Σ)
    (dot(x, Λ, x) + logdet(Σ) + length(x) * log(2π))/2
end

# probability density function
pdf(G::Gaussian, x) = exp(-nld(G, x))
pdf(x::Real) = exp(-x^2/2) / sqrt(2π)
function pdf(G::Uni, x::Real)
	σ = sqrt(G.Σ)
	pdf((x - G.μ) / σ) / σ
end

using SpecialFunctions: erf
cdf(x::Real) = (1 + erf(x / sqrt(2))) / 2
cdf(G::Uni, x::Real) = cdf((x - G.μ) / G.Σ)

# lazy
nld(G::Gaussian) = x -> nld(G, x)
pdf(G::Gaussian) = x -> pdf(G, x)
cdf(G::Uni) = x -> cdf(G, x)

entropy(G::Univariate) = log(2π*ℯ*cov(G)) / 2
entropy(G::Multivariate) = (length(G) * log(2π*ℯ) + logdet(cov(G))) / 2

# for Gaussian process
function nld(G::Gaussian, x::AbstractVector, y::AbstractVector)
	nld(marginal(G, x), y)
end
# negative log marginal likelihood
# with noiseless observations y
nlml(K::AbstractMatOrFac, y::AbstractVector) = nld(K, y)
nlml(k, x::AbstractVector, y::AbstractVector) = nlml(gramian(k, x), y)
nlml(G::Gaussian, x::AbstractVector, y::AbstractVector) = nlml(cov(G), x, y)

# with noisy observations y
# if first argument is not Gaussian, assume it's a kernel
function nlml(k, x::AbstractVector, Gy::Multivariate; tol::Real = tol)
	G = gramian(k, x)
	F = factorize(G, tol = tol) # this should be a pivoted cholesky or low rank
	y, Σ = mean(Gy), cov(Gy)
	K = Woodbury(Σ, F)
	nlml(K, y) # reduces to "noiseless" call
end
function nlml(G::Gaussian, x::AbstractVector, Gy::Multivariate; tol::Real = tol)
	if mean(G) != zero_mean # subtract prior mean, if applicable
		mean(Gy) .-= mean(G, x)
	end
	nlml(cov(G), x, Gy)
end
# convenience for passing noise covariance separately
function nlml(k, x::AbstractVector, y::AbstractVector, Σ::Union{Real, AbstractMatOrFacOrUni})
	nlml(k, x, 𝒩(y, Σ))
end

# value and gradient of nlml with noiseless observations
function pushforward(::typeof(nlml), k, x::AbstractVector, y::AbstractVector)
	K = gramian(k, x)
	pushforward(nlml, K, y)
end
# value and gradient of nlml with noisy observations
function pushforward(::typeof(nlml), k, x::AbstractVector, Gy::Multivariate; tol::Real = tol)
	K = gramian(k, x)
	K = factorize(K, tol = tol) # ideally, this is a pivoted cholesky or low rank
	y, Σ = mean(Gy), cov(Gy)
	K = Woodbury(Σ, K)
	pushforward(nlml, K, y)
end
function pushforward(::typeof(nlml), K::AbstractMatOrFac, y::AbstractVector)
	K = factorize(K)
    α = K \ y # coefficients α are the same as in conditional mean, could pre-allocate
	val = (dot(y, α) + logdet(K) + length(y) * log(2π))/2
	K⁻¹ = Matrix(inv(K))
	@. K⁻¹ -= α * α'
	val, (dK::AbstractMatOrFac) -> dot(K⁻¹, dK)/2 # this is O(n^2) per evaluation
end

# could be useful if certain hyper-parameter combinations lead to rank-deficiency
# function safe_nlml(Σ, y)
#     try nlml(Σ, y) catch e;
#         println(e)
#         Inf
#     end
# end

# # TODO: stochastic gradient via stochastic trace estimator and logdet
# # derivative for Multivariate nld
# function pushforward(::typeof(nlml), Σ::AbstractMatOrFac, y::AbstractVector,
# 										dense_inverse::Val{false} = Val(false))
# 	Σ = factorize(Σ)
#     α = Σ \ y # could pre-allocate here for Σ\∇,
# 	val = (dot(y, α) + logdet(Σ) + length(y) * log(2π))/2
#     val, function (dΣ)
# 				(tr(inverse(Σ)*dΣ) - dot(α, dΣ, α))/2 # this is O(n^3) per evaluation
# 			end
# end


using jInv.Mesh;
using ForwardHelmholtz
using ForwardHelmholtz.PointSourceADR

plotting = true;

if plotting
	using PyPlot;
	# close("all");
end

include("getModels.jl");

NeumannAtFirstDim = true;
Sommerfeld = true;
gamma_val = 0.00001*2*pi;

# model = "const";
# model = "linear";
# model = "linear squared slowness";
# model = "Gaussian"
# model = "Wave-guide"
model = "Wedge"
domain = [0.0,1.0,0.0,1.0];
	n_cells = [512,512];
src = div(n_cells,2)+1;
src[2] = 1;

if model=="const"
	NeumannAtFirstDim = false;
	Sommerfeld = true;
	f = 40.0;
	src = div(n_cells,2)+1;
	m_coarse = getConstModel(n_cells+1);
	pad_cells = [20;20];
elseif model=="linear"
	f = 40.0
	m_coarse = getLinearModel(1.6,3.5,n_cells+1);m_coarse = 1./m_coarse.^2
	pad_cells = [20;20];
elseif model=="linear squared slowness"
	f = 70.0
	m_coarse = getInvLinearModel(1.6,3.5,n_cells+1);
	pad_cells = [20;20];
elseif model=="Gaussian"
	f = 40.0
	pad_cells = [50;50];
	Minv = getRegularMesh(domain,n_cells);
	m_coarse = getSmoothGaussianMedium(n_cells+1,Minv.h)[1];
elseif model=="Wave-guide"
	pad_cells = [64;64];
	f = 20.0;
	m_coarse = waveGuideModel(n_cells+1)[1];
elseif model=="Wedge"
	f = 20.0;
	pad_cells = [36;36];
	n_cells = [512,512];  (m_coarse,Minv) = getWedgeModel(n_cells+1);domain = Minv.domain;
else
	error("unknown model");
end
Minv = getRegularMesh(domain,n_cells);

# m_coarse = getLinearModel(1.6,3.5,n_cells+1);
# m_coarse = 1./m_coarse.^2
# m_coarse = getLinearModel(1./((1.6)^2),1./(3.5)^2,n_cells+1);
# m_coarse = getSmoothGaussianMedium(n_cells+1,Minv.h)[1];
#

w = 2*pi*f

figure()
imshow(m_coarse'); colorbar();
title(string("The ",model," model"))
# error("Here")

n_tup = tuple((n_cells+1)...);
q_coarse = zeros(Complex,n_tup)
q_coarse[src[1],src[2]] = 1/(Minv.h[2]^2);

# src2 = div(3*n_cells,4)+1;
# q_coarse[src2[1],src2[2]] = 1/(Minv.h[2]^2);
# src = src2;

println("omega*h:");
println(w*Minv.h*sqrt(maximum(m_coarse)));

# m = readdlm("SEGmodel2Dsalt.dat"); m = m'; m = m*1e-3;
maxOmega = getMaximalFrequency(m_coarse,Minv);
ABLamp = maxOmega;

### GENERATE FINE REFERENCE ##############################################################################
factor = 4;
n_cells_fine = factor*n_cells;
pad_cells_fine = factor*pad_cells;

Minv_fine = getRegularMesh(domain,n_cells_fine);
src_fine = div(n_cells_fine,2)+1;
if src[2]==1
	src_fine[2] = 1;
end
n_tup = tuple((n_cells_fine+1)...);
q = zeros(Complex,n_tup)
q[src_fine[1],src_fine[2]] = 1/(Minv_fine.h[2]^2);

# src2 = div(3*n_cells_fine,4)+1;
# q[src2[1],src2[2]] = 1/(Minv_fine.h[2]^2);

m=0;
if model=="const"
	m = getConstModel(n_cells_fine+1);
elseif model=="linear"
	m = getLinearModel(1.6,3.5,n_cells_fine+1);m = 1./m.^2
elseif model=="linear squared slowness"
	m = getInvLinearModel(1.6,3.5,n_cells_fine+1);
elseif model=="Gaussian"
	m = getSmoothGaussianMedium(n_cells_fine+1,Minv_fine.h)[1];
elseif model=="Wave-guide"
	m = waveGuideModel(n_cells_fine+1)[1];
elseif model=="Wedge"
	m = getWedgeModel(n_cells_fine+1)[1];
end
#
#

H, = GetHelmholtzOperator(Minv_fine, m, w, gamma_val*ones(Float32,size(m)),NeumannAtFirstDim,pad_cells_fine+1,2.0*ABLamp,Sommerfeld);
xh_fine = (H\q[:]);
xh_fine = reshape(xh_fine,n_tup);

n_tup = tuple(n_cells+1...);
xh_ref = zeros(Complex128,n_tup);
xh_ref = xh_fine[1:factor:end,1:factor:end];
Iref = xh_ref
figure()
imshow(real(xh_ref)');colorbar();
### FINE REFERENCE DONE ##############################################################################

m = m_coarse;

gamma = gamma_val*ones(Float32,size(m)) + getABL(Minv,NeumannAtFirstDim,pad_cells+1,ABLamp);
(ADR_long,ADR_short,T) = getHelmholtzADR(true,Minv, m, w, gamma,NeumannAtFirstDim,src,Sommerfeld);
(ADR_longUp,ADR_short,T) = getHelmholtzADR(false,Minv, m, w, gamma,NeumannAtFirstDim,src,Sommerfeld);

figure();
imshow(real(exp(1im*w*reshape(T,n_tup)))');

H, = GetHelmholtzOperator(Minv, m, w, gamma_val*ones(Float32,size(m)),NeumannAtFirstDim,pad_cells+1,ABLamp,Sommerfeld);
S =  GetHelmholtzShiftOP(m, w,0.2);

q = q_coarse;
xh = (H\q[:])
q = q[:].*exp(1im*w*T);
xl = (ADR_long\q[:]).*exp(-1im*w*T);
xs = ((ADR_short )\q[:]).*exp(-1im*w*T);
us = (H+S)\q[:];
alu = (ADR_longUp\q[:]);
xlu = alu.*exp(-1im*w*T);

I1 = reshape(xh,n_tup);
I2 = reshape(xl,n_tup);
I3 = reshape(xlu,n_tup);
I4 = reshape(xs,n_tup);
I5 = reshape(us,n_tup);

b = maximum(real([I1[:];I2[:];I3[:]]));
a = minimum(real([I1[:];I2[:];I3[:]]));
a_err = 1.0;
N = abs(Iref);
N = N + 0.1*mean(N);
# N[:]=1.0;

E1 = abs(Iref-I1)./N;
E2 = abs(Iref-I2)./N;
E3 = abs(Iref-I3)./N;
E4 = abs(Iref-I4)./N;
# E1[src[1],src[2]] = 0.0;
# E2[src[1],src[2]] = 0.0;
# E3[src[1],src[2]] = 0.0;
# E4[src[1],src[2]] = 0.0;

if plotting
figure();
subplot(1,4,1);
imshow(real(I2'));
subplot(1,4,2);
imshow(reshape(T,n_tup)');
subplot(1,4,3);
imshow(reshape(real(alu),n_tup)');
subplot(1,4,4);
imshow(reshape(real(exp(-1im*w*T)),n_tup)');

figure();
subplot(1,3,1);
imshow(real(I1')); title("Full Helmholtz solution")
subplot(1,3,2);
imshow(reshape(real(us),n_tup)'); title("Shifted Helmholtz solution")
subplot(1,3,3);
imshow(reshape(real(xs),n_tup)'); title("First order upwind ADR solution")

# figure()
# subplot(2,4,1);
# imshow(real(I1'),clim = [a,b]); colorbar();title("Standard discretization")
# subplot(2,4,2);
# imshow(real(I2'),clim = [a,b]);colorbar();title("ADR discretization - central")
# subplot(2,4,3);
# imshow(real(I3'),clim = [a,b]);colorbar();title("ADR discretization - 2nd upwind")
# subplot(2,4,4);
# imshow(real(I4'),clim = [a,b]);colorbar();title("ADR discretization - 1st upwind")

# subplot(2,4,5);
# imshow(E1',clim = [0.0,a_err]);colorbar();title("Relative error - standard")
# subplot(2,4,6);
# imshow(E2',clim = [0.0,a_err]);colorbar();title("Relative error - ADR central")
# subplot(2,4,7);
# imshow(E3',clim = [0.0,a_err]);colorbar();title("Relative error - ADR upwind")
# subplot(2,4,8);
# imshow(E4',clim = [0.0,a_err]);colorbar();

figure()
subplot(2,3,1);
imshow(real(I1'),clim = [a,b]); colorbar();title("Standard discretization");
subplot(2,3,2);
imshow(real(I2'),clim = [a,b]);colorbar();title("ADR discretization - central")
subplot(2,3,3);
imshow(real(I3'),clim = [a,b]);colorbar();title("ADR discretization - 2nd order upwind")

subplot(2,3,4);
imshow(E1',clim = [0.0,a_err]);colorbar();title("Relative error - standard")
subplot(2,3,5);
imshow(E2',clim = [0.0,a_err]);colorbar();title("Relative error - ADR central")
subplot(2,3,6);
imshow(E3',clim = [0.0,a_err]);colorbar();title("Relative error - ADR upwind")

# figure()
# subplot(2,3,1);
# imshow(reshape(T,n[1],n[2])');colorbar();
# subplot(2,3,2);
# imshow(reshape(real(exp(1im*w*T)),n[1],n[2])'); colorbar();
# subplot(2,3,3);
# imshow(reshape(G1,n[1],n[2])');colorbar();
# subplot(2,3,4);
# imshow(reshape(G2,n[1],n[2])');colorbar();
# subplot(2,3,5);
# imshow(reshape(G1.^2 + G2.^2,n[1],n[2])');colorbar();
# subplot(2,3,6);
# imshow(reshape(LT,n[1],n[2])');colorbar();

end

error("ET")


using GeometricBase.Config
using GeometricBase.Utils
using GeometricIntegrators.Integrators
using GeometricIntegrators.Integrators.VPRK
using GeometricIntegrators.Tableaus
using GeometricProblems.LotkaVolterra2d
using SimpleSolvers
using Test

SimpleSolvers.set_config(:nls_atol, 8eps())
SimpleSolvers.set_config(:nls_rtol, 2eps())

const Δt = 0.01
const nt = 10
const q₀ = [1.0, 1.0]
const parameters = (a₁=1.0, a₂=1.0, b₁=-1.0, b₂=-2.0)

ode  = lotka_volterra_2d_ode(q₀; params=parameters)
iode = lotka_volterra_2d_iode(q₀; params=parameters)
lode = lotka_volterra_2d_lode(q₀; params=parameters)
ldae = lotka_volterra_2d_ldae(q₀; params=parameters)

int  = IntegratorFIRK(ode, TableauGauss(8), Δt)
sol  = integrate(ode, int, nt)
refx = sol.q[end]

@testset "$(rpad("VPRK integrators",80))" begin

    sol = integrate(iode, TableauVPRK(:pglrk, 2, TableauGauss(1), -1), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 2E-6

    sol = integrate(iode, TableauVPRK(:pglrk, 4, TableauGauss(2), +1), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 8E-7

    sol = integrate(iode, TableauVPRK(:pglrk, 6, TableauGauss(3), -1), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-12

    sol = integrate(iode, TableauVPLobattoIIIA(2), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-6

    sol = integrate(iode, TableauVPLobattoIIIA(3), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 8E-7

    sol = integrate(iode, TableauVPLobattoIIIA(4), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 3E-11

    sol = integrate(iode, TableauVPLobattoIIIB(2), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 2E-6

    sol = integrate(iode, TableauVPLobattoIIIB(3), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 8E-7

    sol = integrate(iode, TableauVPLobattoIIIB(4), Δt, nt)
    @test relative_maximum_error(sol.q, refx) < 2E-11

end

@testset "$(rpad("VPRK integrators with standard projection",80))" begin

    int = IntegratorVPRKpStandard(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-6

    int = IntegratorVPRKpStandard(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-11

    int = IntegratorVPRKpStandard(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-15

end

@testset "$(rpad("VPRK integrators with symplectic projection",80))" begin

    int = IntegratorVPRKpSymplectic(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-6

    int = IntegratorVPRKpSymplectic(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-11

    int = IntegratorVPRKpSymplectic(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 8E-12

end

@testset "$(rpad("VPRK integrators with symmetric projection",80))" begin

    int = IntegratorVPRKpSymmetric(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-6

    int = IntegratorVPRKpSymmetric(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-11

    int = IntegratorVPRKpSymmetric(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-16

end

@testset "$(rpad("VPRK integrators with midpoint projection",80))" begin

    int = IntegratorVPRKpMidpoint(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-6

    int = IntegratorVPRKpMidpoint(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-11

    int = IntegratorVPRKpMidpoint(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-16

end

@testset "$(rpad("VPRK integrators with internal projection",80))" begin

    int = IntegratorVPRKpInternal(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 2E-6

    int = IntegratorVPRKpInternal(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-11

    int = IntegratorVPRKpInternal(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-12

    int = IntegratorVPRKpInternal(iode, TableauVPGLRK(4), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 8E-16

end

@testset "$(rpad("VPRK integrators with projection on secondary constraint",80))" begin

    # TODO: reactivate

    # int = IntegratorVPRKpSecondary(ldae, TableauVPGLRK(1), Δt)
    # sol = integrate(ldae, int, nt)
    # @test relative_maximum_error(sol.q, refx) < 2E-6

    # int = IntegratorVPRKpSecondary(ldae, TableauVPGLRK(2), Δt)
    # sol = integrate(ldae, int, nt)
    # @test relative_maximum_error(sol.q, refx) < 8E-7

    # int = IntegratorVPRKpSecondary(ldae, TableauVPGLRK(3), Δt)
    # sol = integrate(ldae, int, nt)
    # @test relative_maximum_error(sol.q, refx) < 4E-12

end

@testset "$(rpad("VPRK integrators with variational projection",80))" begin

    intV1 = IntegratorVPRKpVariational(iode, TableauVPGLRK(1), Δt)
    solV1 = integrate(iode, intV1, nt)
    @test relative_maximum_error(solV1.q, refx) < 8E-7

    intV2 = IntegratorVPRKpVariational(iode, TableauVPGLRK(2), Δt)
    solV2 = integrate(iode, intV2, nt)
    @test relative_maximum_error(solV2.q, refx) < 8E-8

    intV3 = IntegratorVPRKpVariational(iode, TableauVPGLRK(3), Δt)
    solV3 = integrate(iode, intV3, nt)
    @test relative_maximum_error(solV3.q, refx) < 1E-11

    intQ1 = IntegratorVPRKpVariationalQ(iode, TableauVPGLRK(1), Δt)
    solQ1 = integrate(iode, intQ1, nt)
    @test relative_maximum_error(solQ1.q, refx) < 4E-5

    intQ2 = IntegratorVPRKpVariationalQ(iode, TableauVPGLRK(2), Δt)
    solQ2 = integrate(iode, intQ2, nt)
    @test relative_maximum_error(solQ2.q, refx) < 2E-4

    intQ3 = IntegratorVPRKpVariationalQ(iode, TableauVPGLRK(3), Δt)
    solQ3 = integrate(iode, intQ3, nt)
    @test relative_maximum_error(solQ3.q, refx) < 1E-8

    intP1 = IntegratorVPRKpVariationalP(iode, TableauVPGLRK(1), Δt)
    solP1 = integrate(iode, intP1, nt)
    @test relative_maximum_error(solP1.q, refx) < 8E-7

    intP2 = IntegratorVPRKpVariationalP(iode, TableauVPGLRK(2), Δt)
    solP2 = integrate(iode, intP2, nt)
    @test relative_maximum_error(solP2.q, refx) < 8E-8

    intP3 = IntegratorVPRKpVariationalP(iode, TableauVPGLRK(3), Δt)
    solP3 = integrate(iode, intP3, nt)
    @test relative_maximum_error(solP3.q, refx) < 1E-11

    @test relative_maximum_error(solV1.q, solP1.q[end]) == 0
    @test relative_maximum_error(solV2.q, solP2.q[end]) == 0
    @test relative_maximum_error(solV3.q, solP3.q[end]) == 0

end

@testset "$(rpad("VPRK integrators with projection on Runge-Kutta tableau",80))" begin

    int = IntegratorVPRKpTableau(iode, CoefficientsPGLRK(5), Δt*5)
    sol = integrate(iode, int, div(nt,5))
    @test relative_maximum_error(sol.q, refx) < 8E-12

    int = IntegratorVPRKpTableau(iode, CoefficientsPGLRK(6), Δt*5)
    sol = integrate(iode, int, div(nt,5))
    @test relative_maximum_error(sol.q, refx) < 4E-14

end

@testset "$(rpad("Degenerate symplectic partitioned Runge-Kutta methods",80))" begin

    int = IntegratorVPRKdegenerate(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 2E-5

    int = IntegratorVPRKdegenerate(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 4E-7

    int = IntegratorVPRKdegenerate(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 2E-10

end

@testset "$(rpad("VSPRK integrators with Legendre projection",80))" begin

    int = IntegratorVPRKpLegendre(iode, TableauVPGLRK(1), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-6

    int = IntegratorVPRKpLegendre(iode, TableauVPGLRK(2), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 1E-11

    int = IntegratorVPRKpLegendre(iode, TableauVPGLRK(3), Δt)
    sol = integrate(iode, int, nt)
    @test relative_maximum_error(sol.q, refx) < 8E-16

end


function regulargrid2d(mins, maxs, ns, dz)
	linearindex = (i1, i2)->i2 + ns[2] * (i1 - 1)
	coords = Array{Float64}(undef, 2, prod(ns))
	xs = range(mins[1]; stop=maxs[1], length=ns[1])
	ys = range(mins[2]; stop=maxs[2], length=ns[2])
	dx = xs[2] - xs[1]
	dy = ys[2] - ys[1]
	j = 1
	neighbors = Array{Pair{Int, Int}}(undef, 2 * prod(ns) - ns[1] - ns[2])
	areasoverlengths = Array{Float64}(undef, 2 * prod(ns) - ns[1] - ns[2])
	volumes = Array{Float64}(undef, 0)
	for i1 = 1:ns[1]
		for i2 = 1:ns[2]
			push!(volumes, dx * dy * dz)
			coords[1, linearindex(i1, i2)] = xs[i1]
			coords[2, linearindex(i1, i2)] = ys[i2]
			if i1 < ns[1]
				neighbors[j] = linearindex(i1, i2)=>linearindex(i1 + 1, i2)
				areasoverlengths[j] = dy * dz / dx
				j += 1
			end
			if i2 < ns[2]
				neighbors[j] = linearindex(i1, i2)=>linearindex(i1, i2 + 1)
				areasoverlengths[j] = dx * dz / dy
				j += 1
			end
		end
	end
	return coords, neighbors, areasoverlengths, volumes
end

function regulargrid3d(mins, maxs, ns)
	linearindex = (i1, i2, i3)->i3 + ns[3] * (i2 - 1) + ns[3] * ns[2] * (i1 - 1)
	coords = Array{Float64}(undef, 3, prod(ns))
	xs = range(mins[1]; stop=maxs[1], length=ns[1])
	ys = range(mins[2]; stop=maxs[2], length=ns[2])
	zs = range(mins[3]; stop=maxs[3], length=ns[3])
	dx = xs[2] - xs[1]
	dy = ys[2] - ys[1]
	dz = zs[2] - zs[1]
	j = 1
	neighbors = Array{Pair{Int, Int}}(undef, 3 * prod(ns) - ns[1] * ns[2] - ns[1] * ns[3] - ns[2] * ns[3])
	areasoverlengths = Array{Float64}(undef, 3 * prod(ns) - ns[1] * ns[2] - ns[1] * ns[3] - ns[2] * ns[3])
	volumes = fill(dx * dy * dz, size(coords, 2))
	for i1 = 1:ns[1]
		for i2 = 1:ns[2]
			for i3 = 1:ns[3]
				coords[1, linearindex(i1, i2, i3)] = xs[i1]
				coords[2, linearindex(i1, i2, i3)] = ys[i2]
				coords[3, linearindex(i1, i2, i3)] = zs[i3]
				if i1 < ns[1]
					neighbors[j] = linearindex(i1, i2, i3)=>linearindex(i1 + 1, i2, i3)
					areasoverlengths[j] = dy * dz / dx
					j += 1
				end
				if i2 < ns[2]
					neighbors[j] = linearindex(i1, i2, i3)=>linearindex(i1, i2 + 1, i3)
					areasoverlengths[j] = dx * dz / dy
					j += 1
				end
				if i3 < ns[3]
					neighbors[j] = linearindex(i1, i2, i3)=>linearindex(i1, i2, i3 + 1)
					areasoverlengths[j] = dx * dy / dz
					j += 1
				end
			end
		end
	end
	return coords, neighbors, areasoverlengths, volumes
end

function darcy_velocity(h, Ks, coords, ns)
	@assert length(ns) == 2
	@show size(Ks)
	@show size(ns)
	hr = reshape(h, ns...)
	vr = zeros(2, ns...)
	@show size(vr)
	@show ns
	coordsr = reshape(coords, 2, reverse(ns)...)
	@show size(coordsr)
	@show size(hr)
	for i = 2:ns[1] - 1
		for j = 2:ns[2] - 1
			#=
			if coordsr[1, j, i + 1, j] - coordsr[1, i - 1, j] == 0
				@show i, j
				@show coordsr[1, i + 1, j], coordsr[1, i - 1, j]
			end
			=#
			vr[1, i, j] = Ks[i, j] * (hr[i + 1, j] - hr[i - 1, j]) / (coordsr[1, j, i + 1] - coordsr[1, j, i - 1])
			vr[2, i, j] = Ks[i, j] * (hr[i, j + 1] - hr[i, j - 1]) / (coordsr[2, j + 1, i] - coordsr[2, j - 1, i])
		end
	end
	return vr
end


using Distributions:Random
using AdaptiveABC
using SequentialLOB
using Random
using Distributions
using Optim
using NLopt
using JLD2
using Plots
using DelimitedFiles

# Set Simulator Parameters and Seed
parameter_names = ["D", "σ", "nu", "μ", "α"]
plotting_parameter_names = ["D", "σ", "ν", "μ", "α"]
n_replications = 10
Random.seed!(1717617000301)
# Constants
num_paths = 1; M = 400 ; T = 2299 ; p₀ = 238.75 ; L = 200 ; λ = 1.0

# Read in Observed Market Data

observed_price_path = round.(readdlm("$(dirname(@__FILE__))/Original_Price_Bars_2300.csv", ',', Float64, '\n')[:,1], digits=2)
observed_log_returns = diff(log.(observed_price_path[:,1]))
observed_summary_stats = get_summary_stats(observed_log_returns)

# plt = plot(
#     1:T + 1, observed_price_path[:,1], legend=false,
#     xlab="Time", ylab="Mid-Price",
# )

# Define Summary function used in AdaptiveABC. Takes in a vector of parameters and outputs a success boolean and the
# simulated summart statistics
function summary_fn(parameters, n_summary_stats, n_replications)
    try
        D, σ, nu, μ, α_slob = parameters
        model = SLOB(num_paths, T, p₀, M, L, D, σ, nu, α_slob, SourceTerm(λ, μ))
        summary_stats = Array{Float64,2}(undef, n_summary_stats, n_replications)
        for i in 1:n_replications
            sim_price_path = model()
            sim_obs = diff(log.(sim_price_path[:, 1]))
            summary_stats[:, i] = get_summary_stats(observed_log_returns, sim_obs)
        end
        return true, summary_stats
    catch e
        return false, zeros(n_summary_stats, n_replications)
    end
end

# Setup for Calibration of Parameters through AdaptiveABC
# Define Prior distribution using Prior struct
prior = Prior([
    Uniform(0.5, 3.0), # D
    Uniform(0.1, 3.0), # σ
    Uniform(0.05, 0.95), # nu
    Uniform(0.1, 3.0), # μ
    Uniform(50.0, 200.0) # α_slob
])

#####
# Calibration Technique: SLOB ABC-PMC BBWM
#####

# Distance Function to measure distance between observed summary stats and simulated summary stats
boot_weight_matrix = BlockBootstrapWeightMatrix(9201292, vcat([0], observed_log_returns), get_summary_stats, 100, 10_000)
weighted_bootstrap = WeightedBootstrap(vcat([0], observed_log_returns), get_summary_stats, boot_weight_matrix)

abc_input_wb = ABCInput(
    prior,
    parameter_names,
    observed_summary_stats,
    summary_fn,
    17,
    weighted_bootstrap
)
abc_pmc_out_wb = JLD2.load_object("$(dirname(@__FILE__))/abc_pmc_out_wb.jld2")
# abc_pmc_out_wb = ABC_PMC(
#     abc_input_wb,
#     100,
#     n_replications,
#     0.25,
#     10_000,
#     10;
#     parallel=true,
#     batch_size=250,
#     seed=18291002012543
# )
posterior_means = AdaptiveABC.parameter_means(abc_pmc_out_wb)

println()
println("ABC-PMC BBWM: D, σ, nu, μ, α_slob")
print(round.(posterior_means[:,end], digits=3))
println()

# JLD2.save_object("$(dirname(@__FILE__))/abc_pmc_out_wb.jld2", abc_pmc_out_wb)

# ABC-PMC BBWM: D, σ, nu, μ, α_slob
# [2.087, 0.421, 0.603, 1.202, 106.504]

param_inds = [1,2,3,4,5]
plt = plot(abc_pmc_out_wb, iteration_colours=cgrad(:blues, 5, categorical=true),
    iterations=[1, abc_pmc_out_wb.n_iterations ÷ 2, abc_pmc_out_wb.n_iterations],
    prior_dists=abc_input_wb.prior.distribution[param_inds],
    param_inds=param_inds, param_names=plotting_parameter_names[param_inds])
plt
savefig(plt, "$(dirname(@__FILE__))/abc_pmc_wb.pdf")

#####
# Calibration Technique: ABC-PMC MADWE
#####
abc_input_we = ABCInput(
    prior,
    parameter_names,
    observed_summary_stats,
    summary_fn,
    17,
    WeightedEuclidean(observed_log_returns, get_summary_stats, ones(17), "MAD")
)

abc_pmc_out_we = JLD2.load_object("$(dirname(@__FILE__))/abc_pmc_out_we.jld2")
# abc_pmc_out_we = ABC_PMC(
#     abc_input_we,
#     100,
#     n_replications,
#     0.15,
#     10_000,
#     10;
#     parallel=true,
#     batch_size=250,
#     seed=77261511
# )
posterior_means = AdaptiveABC.parameter_means(abc_pmc_out_we)
println("")
println("ABC-PMC MADWE: D, σ, nu, μ, α")
print(round.(posterior_means[:,end], digits=3))
println("")
# ABC-PMC MADWE: D, σ, nu, μ, α
# [1.661, 1.127, 0.667, 1.987, 127.368]

# JLD2.save_object("$(dirname(@__FILE__))/abc_pmc_out_we.jld2", abc_pmc_out_we)

param_inds = [1,2,3,4,5]
plt = plot(abc_pmc_out_we, iteration_colours=cgrad(:blues, 5, categorical=true),
    iterations=[1,abc_pmc_out_we.n_iterations],
    prior_dists=abc_input_we.prior.distribution[param_inds],
    param_inds=param_inds, param_names=plotting_parameter_names[param_inds])
savefig(plt, "$(dirname(@__FILE__))/abc_pmc_we.pdf")
plt


using LinearAlgebra: givensAlgorithm
using Base: @propagate_inbounds

import LinearAlgebra: lmul!, rmul!
import Base: Matrix

@propagate_inbounds is_offdiagonal_small(H::AbstractMatrix{T}, i::Int, tol = eps(real(T))) where {T} =
    abs(H[i+1,i]) ≤ tol*(abs(H[i,i]) + abs(H[i+1,i+1]))

abstract type SmallRotation end

"""
Given's rotation acting on rows i:i+1
"""
struct Rotation2{Tc,Ts} <: SmallRotation
    c::Tc
    s::Ts
    i::Int
end

"""
Two Given's rotations acting on rows i:i+2. This could also be implemented as one Householder
reflector!
"""
struct Rotation3{Tc,Ts} <: SmallRotation
    c₁::Tc
    s₁::Ts
    c₂::Tc
    s₂::Ts
    i::Int
end

# Some utility to materialize a rotation to a matrix.
function Matrix(r::Rotation2{Tc,Ts}, n::Int) where {Tc,Ts}
    r.i < n || throw(ArgumentError("Matrix should have order $(r.i+1) or larger"))
    G = Matrix{promote_type(Tc,Ts)}(I, n, n)
    G[r.i+0,r.i+0] = r.c
    G[r.i+1,r.i+0] = -conj(r.s)
    G[r.i+0,r.i+1] = r.s
    G[r.i+1,r.i+1] = r.c
    return G
end

function Matrix(r::Rotation3{Tc,Ts}, n::Int) where {Tc,Ts}
    G₁ = Matrix(Rotation2(r.c₁, r.s₁, r.i + 1), n)
    G₂ = Matrix(Rotation2(r.c₂, r.s₂, r.i), n)
    return G₂ * G₁
end

"""
Get a rotation that maps [p₁, p₂] to a multiple of [1, 0]
"""
function get_rotation(p₁, p₂, i::Int)
    c, s, nrm = givensAlgorithm(p₁, p₂)
    Rotation2(c, s, i), nrm
end

"""
Get a rotation that maps [p₁, p₂, p₃] to a multiple of [1, 0, 0]
"""
function get_rotation(p₁, p₂, p₃, i::Int)
    c₁, s₁, nrm₁ = givensAlgorithm(p₂, p₃)
    c₂, s₂, nrm₂ = givensAlgorithm(p₁, nrm₁)
    Rotation3(c₁, s₁, c₂, s₂, i), nrm₂
end

"""
Passed into the Schur factorization function if you do not wish to have the Schur vectors.
"""
struct NotWanted end

@inline lmul!(G::SmallRotation, A::AbstractMatrix) = lmul!(G, A, 1, size(A, 2))
@inline rmul!(A::AbstractMatrix, G::SmallRotation) = rmul!(A, G, 1, size(A, 1))

lmul!(::SmallRotation, ::NotWanted, args...) = nothing
rmul!(::NotWanted, ::SmallRotation, args...) = nothing

@inline function lmul!(G::Rotation3, A::AbstractMatrix, from::Int, to::Int)
    @inbounds for j = from:to
        a₁ = A[G.i+0,j]
        a₂ = A[G.i+1,j]
        a₃ = A[G.i+2,j]

        a₂′ = G.c₁ * a₂ + G.s₁ * a₃
        a₃′ = -G.s₁' * a₂ + G.c₁ * a₃

        a₁′′ = G.c₂ * a₁ + G.s₂ * a₂′
        a₂′′ = -G.s₂' * a₁ + G.c₂ * a₂′

        A[G.i+0,j] = a₁′′
        A[G.i+1,j] = a₂′′
        A[G.i+2,j] = a₃′
    end

    A
end

@inline function rmul!(A::AbstractMatrix, G::Rotation3, from::Int, to::Int)
    @inbounds for j = from:to
        a₁ = A[j,G.i+0]
        a₂ = A[j,G.i+1]
        a₃ = A[j,G.i+2]

        a₂′ = a₂ * G.c₁ + a₃ * G.s₁'
        a₃′ = a₂ * -G.s₁ + a₃ * G.c₁

        a₁′′ = a₁ * G.c₂ + a₂′ * G.s₂'
        a₂′′ = a₁ * -G.s₂ + a₂′ * G.c₂

        A[j,G.i+0] = a₁′′
        A[j,G.i+1] = a₂′′
        A[j,G.i+2] = a₃′
    end
    A
end

@inline function lmul!(G::Rotation2, A::AbstractMatrix, from::Int, to::Int)
    @inbounds for j = from:to
        a₁ = A[G.i+0,j]
        a₂ = A[G.i+1,j]

        a₁′ = G.c * a₁ + G.s * a₂
        a₂′ = -G.s' * a₁ + G.c * a₂

        A[G.i+0,j] = a₁′
        A[G.i+1,j] = a₂′
    end

    A
end

@inline function rmul!(A::AbstractMatrix, G::Rotation2, from::Int, to::Int)
    @inbounds for j = from:to
        a₁ = A[j,G.i+0]
        a₂ = A[j,G.i+1]

        a₁′ = a₁ * G.c + a₂ * G.s'
        a₂′ = a₁ * -G.s + a₂ * G.c

        A[j,G.i+0] = a₁′
        A[j,G.i+1] = a₂′
    end
    A
end

function double_shift_schur!(H::AbstractMatrix{Tv}, from::Int, to::Int,
                             μ::Complex, Q = NotWanted()) where {Tv<:Real}
    m, n = size(H)

    # Compute the nonzero entries of p = (H - μI)(H - μ'I)e₁.
    # Because of the Hessenberg structure we only need H[min:min+2,min:min+1] to form p.
    @inbounds H₁₁ = H[from+0,from+0]
    @inbounds H₂₁ = H[from+1,from+0]

    @inbounds H₁₂ = H[from+0,from+1]
    @inbounds H₂₂ = H[from+1,from+1]
    @inbounds H₃₂ = H[from+2,from+1]

    p₁ = abs2(μ) - 2real(μ) * H₁₁ + H₁₁ * H₁₁ + H₁₂ * H₂₁
    p₂ = -2real(μ) * H₂₁ + H₂₁ * H₁₁ + H₂₂ * H₂₁
    p₃ = H₃₂ * H₂₁

    # Map that column to a mulitiple of e₁ via two Given's rotations
    G₁, nrm = get_rotation(p₁, p₂, p₃, from)

    # Apply the Given's rotations
    lmul!(G₁, H, from, n)
    rmul!(H, G₁, 1, min(from + 3, m))
    rmul!(Q, G₁)

    # Bulge chasing. First step of the for-loop below looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x + + + x x x
    # i → x x x x x x x     + + + + + + +     x + + + x x x
    #     x x x x x x x     o + + + + + +       + + + x x x
    #     x x x x x x x  ⇒  o + + + + + +  ⇒    + + + x x x
    #       |   x x x x           x x x x       + + + x x x
    #       |     x x x             x x x             x x x
    #       |       x x               x x               x x
    #       ↑
    #       i
    #
    # Last iterations looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x x x x + + +
    #     x x x x x x x     x x x x x x x     x x x x + + +
    #       x x x x x x       x x x x x x       x x x + + +
    #         x x x x x  ⇒    x x x x x x  ⇒      x x + + +
    # i → ----- x x x x           + + + +           x + + +
    #           x x x x           o + + +             + + +
    #           x x x x           o + + +             + + +
    #             ↑
    #             i

    @inbounds for i = from + 1 : to - 2
        p₁ = H[i+0,i-1]
        p₂ = H[i+1,i-1]
        p₃ = H[i+2,i-1]

        G, nrm = get_rotation(p₁, p₂, p₃, i)

        # First column is done by hand
        H[i+0,i-1] = nrm
        H[i+1,i-1] = zero(Tv)
        H[i+2,i-1] = zero(Tv)

        # Rotate remaining columns
        lmul!(G, H, i, n)

        # Create a new bulge
        rmul!(H, G, 1, min(i + 3, m))
        rmul!(Q, G)
    end

    # Last bulge is just one Given's rotation
    #     from          to
    #       ↓           ↓
    #from → x x x x x x x    x x x x x x x    x x x x x + +
    #       x x x x x x x    x x x x x x x    x x x x x + +
    #         x x x x x x      x x x x x x      x x x x + +
    #           x x x x x  ⇒     x x x x x  ⇒     x x x + +
    #             x x x x          x x x x          x x + +
    #               x x x            + + +            x + +
    # to  → ------- x x x            o + +              + +

    @inbounds Gₙ, nrm = get_rotation(H[to-1,to-2], H[to,to-2], to-1)
    @inbounds H[to-1,to-2] = nrm
    @inbounds H[to,to-2] = zero(Tv)

    lmul!(Gₙ, H, to-1, n)
    rmul!(H, Gₙ, 1, to)
    rmul!(Q, Gₙ)

    H
end

function single_shift_schur!(H::AbstractMatrix{Tv}, from::Int, to::Int,
                             μ::Number, Q = NotWanted()) where {Tv<:Number}
    m, n = size(H)

    # Compute the nonzero entries of p = (H - μI)e₁.
    @inbounds H₁₁ = H[from+0,from+0]
    @inbounds H₂₁ = H[from+1,from+0]

    p₁ = H₁₁ - μ
    p₂ = H₂₁

    # Map that column to a mulitiple of e₁ via two Given's rotations
    G₁, nrm = get_rotation(p₁, p₂, from)

    # Apply the Given's rotations
    lmul!(G₁, H, from, n)
    rmul!(H, G₁, 1, min(from + 2, m))
    rmul!(Q, G₁)

    # Bulge chasing. First step of the for-loop below looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x + + x x x x
    # i → x x x x x x x     + + + + + + +     x + + x x x x
    #     x x x x x x x     o + + + + + +       + + x x x x
    #         x x x x x  ⇒      x x x x x  ⇒    + + x x x x
    #       |   x x x x           x x x x           x x x x
    #       |     x x x             x x x             x x x
    #       |       x x               x x               x x
    #       ↑
    #       i
    #
    # Last iterations looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x x x x x + +
    #     x x x x x x x     x x x x x x x     x x x x x + +
    #       x x x x x x       x x x x x x       x x x x + +
    #         x x x x x  ⇒      x x x x x  ⇒      x x x + +
    #           x x x x           x x x x           x x + +
    # i → ------- x x x             + + +             x + +
    #             x x x             0 + +               + +
    #               ↑
    #               i

    @inbounds for i = from + 1 : to - 1
        p₁ = H[i+0,i-1]
        p₂ = H[i+1,i-1]

        G, nrm = get_rotation(p₁, p₂, i)

        # First column is done by hand
        H[i+0,i-1] = nrm
        H[i+1,i-1] = zero(Tv)

        # Rotate remaining columns
        lmul!(G, H, i, n)

        # Create a new bulge
        rmul!(H, G, 1, min(i + 2, m))
        rmul!(Q, G)
    end

    H
end

###
### Real arithmetic
###
function local_schurfact!(H::AbstractMatrix{T}, start::Int, to::Int,
                          Q = NotWanted(), tol = eps(T),
                          maxiter = 100*size(H, 1)) where {T<:Real}
    # iteration count
    iter = 0

    @inbounds while true
        iter += 1

        if iter > maxiter
            throw("QR algorithm did not converge")
        end

        # Indexing
        # `to` points to the column where the off-diagonal value was last zero.
        # while `from` points to the smallest index such that there is no small off-diagonal
        # value in columns from:end-1. Sometimes `from` is just 1. Cartoon of a split
        # with from != 1:
        #
        #  + + | | | | + +
        #  + + | | | | + +
        #    o X X X X = =
        #      X X X X = =
        #      . X X X = =
        #      .   X X = =
        #      .     o + +
        #      .     . + +
        #      ^     ^
        #   from   to
        # The X's form the unreduced Hessenberg matrix we are applying QR iterations to,
        # the | and = values get updated by Given's rotations, while the + values remain
        # untouched! The o's are zeros -- or numerically considered zeros.

        # We keep `from` one column past the zero off-diagonal value, so we check whether
        # the `from - 1` column has a small off-diagonal value.
        from = to
        while from > start && !is_offdiagonal_small(H, from - 1, tol)
            from -= 1
        end

        if from == to
            # This just means H[to, to-1] == 0, so one eigenvalue converged at the end
            H[from,from-1] = zero(T)
            to -= 1
        else
            # Now we are sure we can work with a 2×2 block H[to-1:to,to-1:to]
            # We check if this block has a conjugate eigenpair, which might mean we have
            # converged w.r.t. this block if from + 1 == to.
            # Otherwise, if from + 1 < to, we do either a single or double shift, based on
            # whether the H[to-1:to,to-1:to] part has real eigenvalues or a conjugate pair.

            H₁₁, H₁₂ = H[to-1,to-1], H[to-1,to]
            H₂₁, H₂₂ = H[to  ,to-1], H[to  ,to]

            # Scaling to avoid losing precision in the case where we have nearly
            # repeated eigenvalues.
            scale = abs(H₁₁) + abs(H₁₂) + abs(H₂₁) + abs(H₂₂)
            H₁₁ /= scale; H₁₂ /= scale; H₂₁ /= scale; H₂₂ /= scale

            # Trace and discriminant of small eigenvalue problem.
            t = (H₁₁ + H₂₂) / 2
            d = (H₁₁ - t) * (H₂₂ - t) - H₁₂ * H₂₁
            sqrt_discr = sqrt(abs(d))

            # Very important to have a strict comparison here!
            if d < zero(T)
                # Real eigenvalues.
                # Note that if from + 1 == to in this case, then just one additional
                # iteration is necessary, since the Wilkinson shift will do an exact shift.

                # Determine the Wilkinson shift -- the closest eigenvalue of the 2x2 block
                # near H[to,to]

                λ₁ = t + sqrt_discr
                λ₂ = t - sqrt_discr
                λ = abs(H₂₂ - λ₁) < abs(H₂₂ - λ₂) ? λ₁ : λ₂
                λ *= scale

                # Run a bulge chase
                single_shift_schur!(H, from, to, λ, Q)
            else
                # Conjugate pair
                if from + 1 == to
                    # A conjugate pair has converged apparently!
                    if from != 1
                        H[from,from-1] = zero(T)
                    end
                    to -= 2
                else
                    # Otherwise we do a double shift!
                    complex_shift = scale * (t + sqrt_discr * im)
                    double_shift_schur!(H, from, to, complex_shift, Q)
                end
            end
        end

        # Converged!
        to ≤ start && break
    end

    return true
end

###
### Generic implementation
###
function local_schurfact!(H::AbstractMatrix{T}, start::Int, to::Int,
                          Q = NotWanted(), tol = eps(real(T)),
                          maxiter = 100*size(H, 1)) where {T}
    # iteration count
    iter = 0

    @inbounds while true
        iter += 1
        iter > maxiter && return false

        # Indexing, see the real arithmetic version!

        from = to
        while from > start && !is_offdiagonal_small(H, from - 1, tol)
            from -= 1
        end

        if from == to
            # This just means H[to, to-1] == 0, so one eigenvalue converged at the end
            H[from,from-1] = zero(T)
            to -= 1
        else
            # Compute Wilkinson shift
            H₁₁, H₁₂ = H[to-1,to-1], H[to-1,to]
            H₂₁, H₂₂ = H[to  ,to-1], H[to  ,to]
            d = H₁₁ * H₂₂ - H₂₁ * H₁₂
            t = H₁₁ + H₂₂
            sqr = sqrt(t * t - 4d)
            λ₁ = (t + sqr) / 2
            λ₂ = (t - sqr) / 2
            λ = abs(H₂₂ - λ₁) < abs(H₂₂ - λ₂) ? λ₁ : λ₂

            # Run a bulge chase
            single_shift_schur!(H, from, to, λ, Q)
        end

        # Converged!
        to ≤ start && break
    end

    return true
end

local_schurfact!(H::AbstractMatrix{T}, Q = NotWanted(), tol = eps(real(T)), maxiter = 100*size(H, 1)) where {T} =
    local_schurfact!(H, 1, size(H, 2), Q, tol, maxiter)


using Underscores
using Utils

function parse_file(f)
    @_ map(parse(Int64, _), eachline(f))
end

function load()
    open(joinpath(@__DIR__, "09_encrypted.txt"), "r") do f
        parse_file(f)
    end
end

function search_summands(target, candidates)
    for i in candidates, j in candidates
        if i + j == target
            return i, j
        end
    end
end

function ex1(code, lookback)
    for i in lookback+1:length(code)
        if isnothing(search_summands(code[i], view(code, i-lookback:i-1)))
            return code[i]
        end
    end
end

function ex2(code, nonsum)
    l = length(code)
    for i in eachindex(code)
        s = code[i]
        for j in i+1:l
            s += code[j]
            if s == nonsum
                mi, ma = extrema(view(code, i:j))
                return mi + ma
            elseif s > nonsum
                break
            end
        end
    end
end

code = load()

nonsum = ex1(code, 25)
println("The first number not a sum is: ", nonsum)
println("The code breaking number is ", ex2(code, nonsum))

using Test

@testset "Encoding Error" begin
    example = [
        35,
        20,
        15,
        25,
        47,
        40,
        62,
        55,
        65,
        95,
        102,
        117,
        150,
        182,
        127,
        219,
        299,
        277,
        309,
        576,
    ]

    @testset "example 1" begin
        @test ex1(example, 5) == 127
    end

    @testset "example 2" begin
        @test ex2(example, 127) == 62
    end

    @testset "results" begin
        @test ex1(code, 25) == 31161678
        @test ex2(code, 31161678) == 5453868
    end
end