#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*----------------------------------------------------------------------
#* Author: Paula B. Staudt
#* $Id$
#*--------------------------------------------------------------------
type tank_simplified
	tank_simplified()=begin
		new(
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"Valve Constant",
				:Unit=>"m^2.5/h",
				:Default=>4
			)),
			area(Dict{Symbol,Any}(
				:Brief=>"Tank area",
				:Default=>2
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Tank level"
			)),
			flow_vol(Dict{Symbol,Any}(
				:Brief=>"Input flow",
				:PosX=>0.3037,
				:PosY=>0
			)),
			flow_vol(Dict{Symbol,Any}(
				:Brief=>"Output flow",
				:PosX=>1,
				:PosY=>1
			)),
			[
				:(diff(A*Level) = Fin - Fout),
				:(Fout = k*sqrt(Level)),
			],
			[
				"Mass balance","Valve equation",
			],
			[:k,:A,],
			[:Level,:Fin,:Fout,]
		)
	end
	k::DanaReal
	A::area
	Level::length
	Fin::flow_vol
	Fout::flow_vol
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export tank_simplified
function setEquationFlow(in::tank_simplified)
	addEquation(1)
	addEquation(2)
end
function atributes(in::tank_simplified,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=true
	fields[:Icon]="icon/Tank"
	fields[:Brief]="Model of a simplified tank."
	fields[:Info]="== Specify ==
* the Inlet flow rate;

== Initial Conditions ==
* the tank initial level (Level);
"
	drive!(fields,_)
	return fields
end
tank_simplified(_::Dict{Symbol,Any})=begin
	newModel=tank_simplified()
	newModel.attributes=atributes(newModel,_)
	newModel
end


module GlobalSensitivity

using Statistics, RecursiveArrayTools, LinearAlgebra, Random
using QuasiMonteCarlo, ForwardDiff, KernelDensity, Trapz
using Parameters: @unpack
using FFTW, Distributions, StatsBase

abstract type GSAMethod end

include("morris_sensitivity.jl")
include("sobol_sensitivity.jl")
include("regression_sensitivity.jl")
include("DGSM_sensitivity.jl")
include("eFAST_sensitivity.jl")
include("delta_sensitivity.jl")
include("easi_sensitivity.jl")
include("rbd-fast_sensitivity.jl")
include("fractional_factorial_sensitivity.jl")

export gsa

export Sobol, Morris, RegressionGSA, DGSM, eFAST, DeltaMoment, EASI, FractionalFactorial

end # module


## This script produces the time series we use in the tests

using DifferentialEquations
using Random
using DelimitedFiles

# For comparison reasons using Travis CI we carry out the integration on a UNIX
# OS and save the resulting time series

# Lorenz
lo = Systems.lorenz([1.0, 1.0, 50.0])
tr = trajectory(lo, 100; dt = 0.01, Ttr = 10)
x = tr[:, 1] # x-component of time series
y = tr[:, 2] #y-component of time series
writedlm("./test/timeseries/lorenz_pecora_uni_x.csv", x)
writedlm("./test/timeseries/lorenz_pecora_uni_y.csv", y)
writedlm("./test/timeseries/lorenz_pecora_multi.csv", tr)

# Lorenz2
lo = Systems.lorenz([1.0, 1.0, 50.0]; ρ=60)
tr = trajectory(lo, 100; dt = 0.01, Ttr = 10)
x = tr[:, 1] # x-component of time series
y = tr[:, 2] #y-component of time series
writedlm("./test/timeseries/lorenz2_pecora_uni_x.csv", x)
writedlm("./test/timeseries/lorenz2_pecora_uni_y.csv", y)
writedlm("./test/timeseries/lorenz2_pecora_multi.csv", tr)


import Base.Math: @horner

const FORWARD  =  true
const BACKWARD = false

const sqrtpi = 1.772453850905516027298
const edivsqrt2pi = 1.084437551419227546612

"""
Compute a typed 0.5.
"""
half(x::Number) = oftype(x,0.5)
half(x::Integer) = half(float(x))
half(::Type{T}) where {T<:Number} = convert(T,0.5)
half(::Type{T}) where {T<:Integer} = half(AbstractFloat)

"""
Compute a typed 2.
"""
two(x::Number) = oftype(x,2)
two(::Type{T}) where {T<:Number} = convert(T,2)

"""
The Kronecker ``\\delta`` function:

```math
\\delta_{k,j} = \\left\\{\\begin{array}{ccc} 1 & {\\rm for} & k = j,\\\\ 0 & {\\rm for} & k \\ne j.\\end{array}\\right.
```
"""
δ(k::Integer,j::Integer) = k == j ? 1 : 0

"""
Pochhammer symbol ``(x)_n = \\frac{\\Gamma(x+n)}{\\Gamma(x)}`` for the rising factorial.
"""
function pochhammer(x::Number,n::Integer)
    ret = one(x)
    if n≥0
        for i=0:n-1
            ret *= x+i
        end
    else
        ret /= pochhammer(x+n,-n)
    end
    ret
end

pochhammer(x::Number,n::Number) = isinteger(n) ? pochhammer(x,Int(n)) : ogamma(x)/ogamma(x+n)

function pochhammer(x::Number,n::UnitRange{T}) where T<:Real
    ret = Vector{promote_type(typeof(x),T)}(length(n))
    ret[1] = pochhammer(x,first(n))
    for i=2:length(n)
        ret[i] = (x+n[i]-1)*ret[i-1]
    end
    ret
end

ogamma(x::Number) = (isinteger(x) && x<0) ? zero(float(x)) : inv(gamma(x))

"""
Stirling's asymptotic series for ``\\Gamma(z)``.
"""
stirlingseries(z) = gamma(z)*sqrt((z/π)/2)*exp(z)/z^z

function stirlingseries(z::Float64)
    if z ≥ 3274.12075200175       # N =  4
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273)
    elseif z ≥ 590.1021805526798  # N =  5
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917)
    elseif z ≥ 195.81733962412835 # N =  6
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666)
    elseif z ≥ 91.4692823071966   # N =  7
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5)
    elseif z ≥ 52.70218954633605  # N =  8
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939)
    elseif z ≥ 34.84031591198865  # N =  9
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5)
    elseif z ≥ 25.3173982783047   # N = 10
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873)
    elseif z ≥ 19.685015283078513 # N = 11
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5)
    elseif z ≥ 16.088669099569266 # N = 12
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5,-0.0019144384985654776)
    elseif z ≥ 13.655055978888104 # N = 13
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5,-0.0019144384985654776,-0.00016251626278391583)
    elseif z ≥ 11.93238782087875  # N = 14
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5,-0.0019144384985654776,-0.00016251626278391583,0.00640336283380807)
    elseif z ≥ 10.668852439197263 # N = 15
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5,-0.0019144384985654776,-0.00016251626278391583,0.00640336283380807,0.0005401647678926045)
    elseif z ≥ 9.715358216638403  # N = 16
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5,-0.0019144384985654776,-0.00016251626278391583,0.00640336283380807,0.0005401647678926045,-0.02952788094569912)
    elseif z ≥ 8.979120323411497  # N = 17
        @horner(inv(z),1.0,0.08333333333333333,0.003472222222222222,-0.0026813271604938273,-0.00022947209362139917,0.0007840392217200666,6.972813758365857e-5,-0.0005921664373536939,-5.171790908260592e-5,0.0008394987206720873,7.204895416020011e-5,-0.0019144384985654776,-0.00016251626278391583,0.00640336283380807,0.0005401647678926045,-0.02952788094569912,-0.002481743600264998)
    else
        gamma(z)*sqrt(z/2π)*exp(z)/z^z
    end
end

stirlingremainder(z::Number,N::Int) = (1+zeta(N))*gamma(N)/((2π)^(N+1)*z^N)/stirlingseries(z)

Aratio(n::Int,α::Float64,β::Float64) = exp((n/2+α+1/4)*log1p(-β/(n+α+β+1))+(n/2+β+1/4)*log1p(-α/(n+α+β+1))+(n/2+1/4)*log1p(α/(n+1))+(n/2+1/4)*log1p(β/(n+1)))
Aratio(n::Number,α::Number,β::Number) = (1+(α+1)/n)^(n+α+1/2)*(1+(β+1)/n)^(n+β+1/2)/(1+(α+β+1)/n)^(n+α+β+1/2)/(1+(zero(α)+zero(β))/n)^(n+1/2)

Cratio(n::Int,α::Float64,β::Float64) = exp((n+α+1/2)*log1p((α-β)/(2n+α+β+2))+(n+β+1/2)*log1p((β-α)/(2n+α+β+2))-log1p((α+β+2)/2n)/2)/sqrt(n)
Cratio(n::Number,α::Number,β::Number) = n^(-1/2)*(1+(α+1)/n)^(n+α+1/2)*(1+(β+1)/n)^(n+β+1/2)/(1+(α+β+2)/2n)^(2n+α+β+3/2)

Anαβ(n::Number,α::Number,β::Number) = 2^(α+β+1)/(2n+α+β+1)*exp(lgamma(n+α+1)-lgamma(n+α+β+1)+lgamma(n+β+1)-lgamma(n+1))
function Anαβ(n::Integer,α::Number,β::Number)
    if n==0
        2^(α+β+1)*beta(α+1,β+1)
    else
        val = Anαβ(0,α,β)
        for i=1:n
            val *= (i+α)*(i+β)/(i+α+β)/i*(2i+α+β-1)/(2i+α+β+1)
        end
        val
    end
end

function Anαβ(n::Integer,α::Float64,β::Float64)
    if n+min(α,β,α+β,0) ≥ 7.979120323411497
        2 ^ (α+β+1)/(2n+α+β+1)*stirlingseries(n+α+1)*Aratio(n,α,β)/stirlingseries(n+α+β+1)*stirlingseries(n+β+1)/stirlingseries(n+one(Float64))
    else
        (n+1)*(n+α+β+1)/(n+α+1)/(n+β+1)*Anαβ(n+1,α,β)*((2n+α+β+3)/(2n+α+β+1))
    end
end

"""
The Lambda function ``\\Lambda(z) = \\frac{\\Gamma(z+\\frac{1}{2})}{\\Gamma(z+1)}`` for the ratio of gamma functions.
"""
Λ(z::Number) = exp(lgamma(z+half(z))-lgamma(z+one(z)))
"""
For 64-bit floating-point arithmetic, the Lambda function uses the asymptotic series for ``\\tau`` in Appendix B of

I. Bogaert and B. Michiels and J. Fostier, 𝒪(1) computation of Legendre polynomials and Gauss–Legendre nodes and weights for parallel computing, *SIAM J. Sci. Comput.*, **34**:C83–C101, 2012.
"""
function Λ(x::Float64)
    if x > 9.84475
        xp = x+0.25
        @horner(inv(xp^2),1.0,-1.5625e-02,2.5634765625e-03,-1.2798309326171875e-03,1.343511044979095458984375e-03,-2.432896639220416545867919921875e-03,6.7542375336415716446936130523681640625e-03)/sqrt(xp)
    else
        (x+1.0)*Λ(x+1.0)/(x+0.5)
    end
end

"""
The Lambda function ``\\Lambda(z,λ₁,λ₂) = \\frac{\\Gamma(z+\\lambda_1)}{Γ(z+\\lambda_2)}`` for the ratio of gamma functions.
"""
Λ(z::Number,λ₁::Number,λ₂::Number) = exp(lgamma(z+λ₁)-lgamma(z+λ₂))
function Λ(x::Float64,λ₁::Float64,λ₂::Float64)
    if min(x+λ₁,x+λ₂) ≥ 8.979120323411497
        exp(λ₂-λ₁+(x-.5)*log1p((λ₁-λ₂)/(x+λ₂)))*(x+λ₁)^λ₁/(x+λ₂)^λ₂*stirlingseries(x+λ₁)/stirlingseries(x+λ₂)
    else
        (x+λ₂)/(x+λ₁)*Λ(x+1.,λ₁,λ₂)
    end
end

## TODO: deprecate when Lambert-W is supported in a mainstream repository such as SpecialFunctions.jl
"""
The principal branch of the Lambert-W function, defined by ``x = W_0(x) e^{W_0(x)}``, computed using Halley's method for ``x \\in [-e^{-1},\\infty)``.
"""
function lambertw(x::AbstractFloat)
    if x < -exp(-one(x))
        return throw(DomainError())
    elseif x == -exp(-one(x))
        return -one(x)
    elseif x < 0
        w0 = ℯ*x/(1+inv(inv(sqrt(2*ℯ*x+2))+inv(ℯ-1)-inv(sqrt(2))))
    else
        log1px = log1p(x)
        w0 = log1px*(1-log1p(log1px)/(2+log1px))
    end
    expw0 = exp(w0)
    w1 = w0 - (w0*expw0 - x)/((w0 + 1)*expw0 -
        (w0 + 2) * (w0*expw0 - x)/(2w0 + 2))
    while abs(w1/w0 - 1) > 2eps(typeof(x))
        w0 = w1
        expw0 = exp(w0)
        w1 = w0 - (w0*expw0 - x)/((w0 + 1)*expw0 -
            (w0 + 2) * (w0*expw0 - x)/(2w0 + 2))
    end
    return w1
end
lambertw(x::Real) = lambertw(float(x))

Cnλ(n::Integer,λ::Float64) = 2^λ/sqrtpi*Λ(n+λ)
Cnλ(n::Integer,λ::Number) = 2^λ/sqrt(oftype(λ,π))*Λ(n+λ)
function Cnλ(n::UnitRange{T},λ::Number) where T<:Integer
    ret = Vector{typeof(λ)}(undef, length(n))
    ret[1] = Cnλ(first(n),λ)
    for i=2:length(n)
        ret[i] = (n[i]+λ-half(λ))/(n[i]+λ)*ret[i-1]
    end
    ret
end

function Cnmλ(n::Integer,m::Integer,λ::Number)
    if m == 0
        Cnλ(n,λ)
    else
        (λ+m-1)/2/m*(m-λ)/(n+λ+m)*Cnmλ(n,m-1,λ)
    end
end

function Cnαβ(n::Integer,α::Number,β::Number)
    if n==0
        2^(α+β+1)*beta(α+1,β+1)/π
    else
        val = Cnαβ(0,α,β)
        for i=1:n
            val *= (i+α)*(i+β)/(i+(α+β+1)/2)/(i+(α+β)/2)
        end
        val
    end
end

function Cnαβ(n::Integer,α::Float64,β::Float64)
    if n+min(α,β) ≥ 7.979120323411497
        stirlingseries(n+α+1)/sqrtpi/stirlingseries(2n+α+β+2)*Cratio(n,α,β)*stirlingseries(n+β+1)
    else
        (n+(α+β+3)/2)/(n+β+1)*(n+(α+β+2)/2)/(n+α+1)*Cnαβ(n+1,α,β)
    end
end

function Cnmαβ(n::Integer,m::Integer,α::Number,β::Number)
    if m == 0
        Cnαβ(n,α,β)
    else
        Cnmαβ(n,m-1,α,β)/2(2n+α+β+m+1)
    end
end

function Cnmαβ(n::Integer,m::Integer,α::AbstractArray{T},β::AbstractArray{T}) where T<:Number
    shp = promote_shape(size(α),size(β))
    reshape([ Cnmαβ(n,m,α[i],β[i]) for i in eachindex(α,β) ], shp)
end

function absf(α::Number,β::Number,m::Int,θ::Number)
    ret = zero(θ)
    for l=0:m
        ret += pochhammer(half(α)+α,l)*pochhammer(half(α)-α,l)*pochhammer(half(β)+β,m-l)*pochhammer(half(β)-β,m-l)/factorial(l)/factorial(m-l)/sinpi(θ/2)^(l+α+half(α))/cospi(θ/2)^(m-l+β+half(β))
    end
    ret
end

function absf(α::AbstractArray{T},β::AbstractArray{T},m::Int,θ::Number) where T<:Number
    shp = promote_shape(size(α),size(β))
    reshape([ absf(α[i],β[i],m,θ) for i in eachindex(α,β) ], shp)
end

function absf(α::Number,β::Number,m::Int,θ::AbstractArray{T,1}) where T<:Number
    ret = zero(θ)
    cfs = zeros(T,m+1)
    for l=0:m
        @inbounds cfs[l+1] = pochhammer(half(α)+α,l)*pochhammer(half(α)-α,l)*pochhammer(half(β)+β,m-l)*pochhammer(half(β)-β,m-l)/factorial(l)/factorial(m-l)
    end
    @inbounds for i=1:length(θ),l=0:m
        ret[i] += cfs[l+1]/sinpi(θ[i]/2)^(l+α+half(α))/cospi(θ[i]/2)^(m-l+β+half(β))
    end
    ret
end
absf(α::Number,β::Number,m::Int,θ::AbstractArray{T,2}) where {T<:Number} = [ absf(α,β,m,θ[i,j]) for i=1:size(θ,1), j=1:size(θ,2) ]
absf(α::Number,β::Number,m::Int,θ::AbstractArray{T}) where {T<:Number} = reshape([ absf(α,β,m,θ[i]) for i in eachindex(θ) ], size(θ))

function compute_absf!(ret::Vector{T},cfs::Matrix{T},α::T,β::T,tempcos::Vector{T},tempsin::Vector{T},tempcosβsinα::Vector{T},m::Int) where T<:AbstractFloat
    @inbounds for i=1:length(ret)
        temp = inv(tempcos[i]^m*tempcosβsinα[i])
        ret[i] = cfs[m+1,1]*temp
        for l=1:m
            temp *= tempcos[i]/tempsin[i]
            ret[i] += cfs[m+1,l+1]*temp
        end
    end
    ret
end

function compute_absf!(ret::Vector{T},tempsin::Vector{T},tempsinλ::Vector{T},m::Int) where T<:AbstractFloat
    for i=1:length(ret)
        @inbounds ret[i] = inv(tempsin[i]^m*tempsinλ[i])
    end
    ret
end

function compute_umvm!(um::Vector{T},vm::Vector{T},cfs::Matrix{T},α::T,β::T,tempcos::Vector{T},tempsin::Vector{T},tempcosβsinα::Vector{T},m::Int,θ::Vector{T},ir::UnitRange{Int}) where T<:AbstractFloat
    @inbounds for i in ir
        temp = inv(tempcos[i]^m*tempcosβsinα[i])
        ϑ = (α+half(α))/2-(α+β+m+1)*θ[i]/2
        um[i] = cfs[m+1,1]*cospi(ϑ)*temp
        vm[i] = cfs[m+1,1]*sinpi(ϑ)*temp
        @inbounds for l=1:m
            temp *= tempcos[i]/tempsin[i]
            ϑ = (α+l+half(α))/2-(α+β+m+1)*θ[i]/2
            um[i] += cfs[m+1,l+1]*cospi(ϑ)*temp
            vm[i] += cfs[m+1,l+1]*sinpi(ϑ)*temp
        end
    end
end

function compute_umvm!(um::Vector{T},vm::Vector{T},λ::T,tempsinλm::Vector{T},m::Int,θ::Vector{T},ir::UnitRange{Int}) where T<:AbstractFloat
    @inbounds @simd for i in ir
        temp = inv(tempsinλm[i])
        ϑ = (m+λ)*(half(T)-θ[i])
        um[i] = cospi(ϑ)*temp
        vm[i] = sinpi(ϑ)*temp
    end
end

function findmindices!(Rαβjm::Vector{T},cfs::Matrix{T},α::T,β::T,j::Int,m::Int,tempcos::Vector{T},tempsin::Vector{T},tempcosβsinα::Vector{T}) where T<:AbstractFloat
    compute_absf!(Rαβjm,cfs,α,β,tempcos,tempsin,tempcosβsinα,m)
    rmul!(Rαβjm,Cnmαβ(j,m,α,β))
    rmin,imin = findmin(Rαβjm)
    if rmin < eps(T)
        i₁ = imin-1
        while i₁ ≥ 3
            if Rαβjm[i₁] < eps(T)
                i₁-=1
            else
                break
            end
        end
        i₂ = imin+1
        while i₂ ≤ length(Rαβjm)-2
            if Rαβjm[i₂] < eps(T)
                i₂+=1
            else
                break
            end
        end
        return i₁,i₂
    else
        return 1,0# error("There are no indices such that the interior asymptotic formula is valid.")# but adding an error is type-unstable :(.
    end
end

function findmindices!(Rαβjm::Vector{T},λ::T,j::Int,m::Int,tempsin::Vector{T},tempsinλ::Vector{T}) where T<:AbstractFloat
    compute_absf!(Rαβjm,tempsin,tempsinλ,m)
    rmul!(Rαβjm,Cnmλ(j,m,λ))
    rmin,imin = findmin(Rαβjm)
    if rmin < eps(T)
        i₁ = imin-1
        while i₁ ≥ 3
            if Rαβjm[i₁] < eps(T)
                i₁-=1
            else
                break
            end
        end
        i₂ = imin+1
        while i₂ ≤ length(Rαβjm)-2
            if Rαβjm[i₂] < eps(T)
                i₂+=1
            else
                break
            end
        end
        return i₁,i₂
    else
        return 1,0# error("There are no indices such that the interior asymptotic formula is valid.")# but adding an error is type-unstable :(.
    end
end

# initialization methods

function init_cfs(α::T,β::T,M::Int) where T<:AbstractFloat
    cfs = zeros(T,M+1,M+1)
    @inbounds for m=0:M,l=0:m
        cfs[m+1,l+1] = pochhammer(half(α)+α,l)*pochhammer(half(α)-α,l)*pochhammer(half(β)+β,m-l)*pochhammer(half(β)-β,m-l)/factorial(l)/factorial(m-l)
    end
    cfs
end

function init_c₁c₂!(c₁::Vector,c₂::Vector,a::Vector,b::Vector,j₁::Int,j₂::Int)
    @inbounds for j=1:j₁-1 c₁[j] = 0 end
    @inbounds for j=j₁:j₂ c₁[j] = a[j]*b[j] end
    @inbounds for j=j₂+1:length(c₁) c₁[j] = 0 end
    copyto!(c₂,c₁)
end

function init_c₁c₂!(c₁::Vector,c₂::Vector,u::Vector,v::Vector,c::Vector,i₁::Int,i₂::Int)
    @inbounds for i=1:i₁-1
        c₁[i] = 0
        c₂[i] = 0
    end
    @inbounds for i=i₁:i₂
        c₁[i] = u[i]*c[i]
        c₂[i] = v[i]*c[i]
    end
    @inbounds for i=i₂+1:length(c₁)
        c₁[i] = 0
        c₂[i] = 0
    end
end

"""
Modified Chebyshev moments of the first kind:

```math
    \\int_{-1}^{+1} T_n(x) {\\rm\\,d}x.
```
"""
function chebyshevmoments1(::Type{T}, N::Int) where T
    μ = zeros(T, N)
    for i = 0:2:N-1
        @inbounds μ[i+1] = two(T)/T(1-i^2)
    end
    μ
end

"""
Modified Chebyshev moments of the first kind with respect to the Jacobi weight:

```math
    \\int_{-1}^{+1} T_n(x) (1-x)^\\alpha(1+x)^\\beta{\\rm\\,d}x.
```
"""
function chebyshevjacobimoments1(::Type{T}, N::Int, α, β) where T
    μ = zeros(T, N)
    N > 0 && (μ[1] = 2 .^ (α+β+1)*beta(α+1,β+1))
    if N > 1
        μ[2] = μ[1]*(β-α)/(α+β+2)
        for i=1:N-2
            @inbounds μ[i+2] = (2(β-α)*μ[i+1]-(α+β-i+2)*μ[i])/(α+β+i+2)
        end
    end
    μ
end

"""
Modified Chebyshev moments of the first kind with respect to the logarithmic weight:

```math
    \\int_{-1}^{+1} T_n(x) \\log\\left(\\frac{1-x}{2}\\right){\\rm\\,d}x.
```
"""
function chebyshevlogmoments1(::Type{T}, N::Int) where T
    μ = zeros(T, N)
    N > 0 && (μ[1] = -two(T))
    if N > 1
        μ[2] = -one(T)
        for i=1:N-2
            cst = isodd(i) ? T(4)/T(i^2-4) : T(4)/T(i^2-1)
            @inbounds μ[i+2] = ((i-2)*μ[i]+cst)/(i+2)
        end
    end
    μ
end

"""
Modified Chebyshev moments of the second kind:

```math
    \\int_{-1}^{+1} U_n(x) {\\rm\\,d}x.
```
"""
function chebyshevmoments2(::Type{T}, N::Int) where T
    μ = zeros(T, N)
    for i = 0:2:N-1
        @inbounds μ[i+1] = two(T)/T(i+1)
    end
    μ
end

"""
Modified Chebyshev moments of the second kind with respect to the Jacobi weight:

```math
    \\int_{-1}^{+1} U_n(x) (1-x)^\\alpha(1+x)^\\beta{\\rm\\,d}x.
```
"""
function chebyshevjacobimoments2(::Type{T}, N::Int, α, β) where T
    μ = zeros(T, N)
    N > 0 && (μ[1] = 2 .^ (α+β+1)*beta(α+1,β+1))
    if N > 1
        μ[2] = 2μ[1]*(β-α)/(α+β+2)
        for i=1:N-2
            @inbounds μ[i+2] = (2(β-α)*μ[i+1]-(α+β-i)*μ[i])/(α+β+i+2)
        end
    end
    μ
end

"""
Modified Chebyshev moments of the second kind with respect to the logarithmic weight:

```math
    \\int_{-1}^{+1} U_n(x) \\log\\left(\\frac{1-x}{2}\\right){\\rm\\,d}x.
```
"""
function chebyshevlogmoments2(::Type{T}, N::Int) where T
    μ = chebyshevlogmoments1(T, N)
    if N > 1
        μ[2] *= two(T)
        for i=1:N-2
            @inbounds μ[i+2] = 2μ[i+2] + μ[i]
        end
    end
    μ
end

"""
Compute Jacobi expansion coefficients in ``P_n^{(\\alpha+1,\\beta)}(x)`` given Jacobi expansion coefficients in ``P_n^{(\\alpha,\\beta)}(x)`` in-place.
"""
function incrementα!(c::AbstractVector,α,β)
    αβ,N = α+β,length(c)
    N > 1 && (c[1] -= (β+1)/(αβ+3)*c[2])
    @inbounds for i=2:N-1 c[i] = (αβ+i)/(αβ+2i-1)*c[i] - (β+i)/(αβ+2i+1)*c[i+1] end
    N > 1 && (c[N] *= (αβ+N)/(αβ+2N-1))
    c
end

"""
Compute Jacobi expansion coefficients in ``P_n^{(\\alpha,\\beta+1)}(x)`` given Jacobi expansion coefficients in ``P_n^{(\\alpha,\\beta)}(x)`` in-place.
"""
function incrementβ!(c::AbstractVector,α,β)
    αβ,N = α+β,length(c)
    N > 1 && (c[1] += (α+1)/(αβ+3)*c[2])
    @inbounds for i=2:N-1 c[i] = (αβ+i)/(αβ+2i-1)*c[i] + (α+i)/(αβ+2i+1)*c[i+1] end
    N > 1 && (c[N] *= (αβ+N)/(αβ+2N-1))
    c
end

"""
Compute Jacobi expansion coefficients in ``P_n^{(\\alpha+1,\\alpha+1)}(x)`` given Jacobi expansion coefficients in ``P_n^{(\\alpha,\\alpha)}(x)`` in-place.
"""
function incrementαβ!(c::AbstractVector,α,β)
    @assert α == β
    N = length(c)
    if N == 2
        c[2] *= (2α+2)/(2α+4)
    elseif N > 2
        c[1] -= (α+2)/(4α+10)*c[3]
        @inbounds for i=2:N-2 c[i] = (2α+i)*(2α+i+1)/(2α+2i-1)/(2α+2i)*c[i] - (α+i+1)/(4α+4i+6)*c[i+2] end
        c[N-1] *= (2α+N-1)*(2α+N)/(2α+2N-3)/(2α+2N-2)
        c[N] *= (2α+N)*(2α+N+1)/(2α+2N-1)/(2α+2N)
    end
    c
end

"""
Compute Jacobi expansion coefficients in ``P_n^{(\\alpha-1,\\beta)}(x)`` given Jacobi expansion coefficients in ``P_n^{(\\alpha,\\beta)}(x)`` in-place.
"""
function decrementα!(c::AbstractVector,α,β)
    αβ,N = α+β,length(c)
    N > 1 && (c[N] *= (αβ+2N-2)/(αβ+N-1))
    @inbounds for i=N-1:-1:2 c[i] = (αβ+2i-2)/(αβ+i-1)*c[i] + (αβ+2i-2)/(αβ+2i)*(β+i)/(αβ+i-1)*c[i+1] end
    N > 1 && (c[1] += (β+1)/(αβ+2)*c[2])
    c
end

"""
Compute Jacobi expansion coefficients in ``P_n^{(\\alpha,\\beta-1)}(x)`` given Jacobi expansion coefficients in ``P_n^{(\\alpha,\\beta)}(x)`` in-place.
"""
function decrementβ!(c::AbstractVector,α,β)
    αβ,N = α+β,length(c)
    N > 1 && (c[N] *= (αβ+2N-2)/(αβ+N-1))
    @inbounds for i=N-1:-1:2 c[i] = (αβ+2i-2)/(αβ+i-1)*c[i] - (αβ+2i-2)/(αβ+2i)*(α+i)/(αβ+i-1)*c[i+1] end
    N > 1 && (c[1] -= (α+1)/(αβ+2)*c[2])
    c
end

"""
Compute Jacobi expansion coefficients in ``P_n^{(\\alpha-1,\\alpha-1)}(x)`` given Jacobi expansion coefficients in ``P_n^{(\\alpha,\\alpha)}(x)`` in-place.
"""
function decrementαβ!(c::AbstractVector,α,β)
    @assert α == β
    N = length(c)
    if N == 2
        c[2] *= (2α+2)/(2α)
    elseif N > 2
        c[N] *= (2α+2N-3)*(2α+2N-2)/(2α+N-2)/(2α+N-1)
        c[N-1] *= (2α+2N-5)*(2α+2N-4)/(2α+N-3)/(2α+N-2)
        @inbounds for i=N-2:-1:2 c[i] = (2α+2i-3)*(2α+2i-2)/(2α+i-2)/(2α+i-1)*(c[i] + (α+i)/(4α+4i+2)*c[i+2]) end
        c[1] += (α+1)/(4α+6)*c[3]
    end
    c
end

function modαβ(α)
    if -0.5 < α ≤ 0.5
        a = α
    else
        a = α%1
        a > 0.5 && (a-=1)
        a ≤ -0.5 && (a+=1)
    end
    a
end

function modλ(λ)
    if 0 ≤ λ < 1
        l = λ
    else
        l = λ%1
        l < 0 && (l+=1)
    end
    l
end

function tosquare!(ret::AbstractVector,α,β)
    a,b = modαβ(α),modαβ(β)
    A,B = α-a,β-b
    if α ≤ -0.5 && β ≤ -0.5
        incrementα!(ret,α,β)
        incrementβ!(ret,a,β)
    elseif α ≤ -0.5
        incrementα!(ret,α,β)
        for j=B:-1:1
            decrementβ!(ret,a,j+b)
        end
    elseif β ≤ -0.5
        incrementβ!(ret,α,β)
        for i=A:-1:1
            decrementα!(ret,i+a,b)
        end
    else
        for i=A:-1:1
            decrementα!(ret,i+a,β)
        end
        for j=B:-1:1
            decrementβ!(ret,a,j+b)
        end
    end
    ret
end

function fromsquare!(ret::AbstractVector,α,β)
    a,b = modαβ(α),modαβ(β)
    A,B = α-a,β-b
    if α ≤ -0.5 && β ≤ -0.5
        decrementα!(ret,a,b)
        decrementβ!(ret,α,b)
    elseif α ≤ -0.5
        decrementα!(ret,a,b)
        for j=0:B-1
            incrementβ!(ret,α,j+b)
        end
    elseif β ≤ -0.5
        decrementβ!(ret,a,b)
        for i=0:A-1
            incrementα!(ret,i+a,β)
        end
    else
        for i=0:A-1
            incrementα!(ret,i+a,b)
        end
        for j=0:B-1
            incrementβ!(ret,α,j+b)
        end
    end
    ret
end

function toline!(ret::AbstractVector,α,β)
    @assert α == β
    a,b = modαβ(α),modαβ(β)
    A,B = α-a,β-b
    if α ≤ -0.5
        incrementαβ!(ret,α,β)
    else
        for i=A:-1:1
            decrementαβ!(ret,i+a,i+a)
        end
    end
    ret
end

function fromline!(ret::AbstractVector,α,β)
    @assert α == β
    a,b = modαβ(α),modαβ(β)
    A,B = α-a,β-b
    if α ≤ -0.5
        decrementαβ!(ret,a,b)
    else
        for i=0:A-1
            incrementαβ!(ret,i+a,i+a)
        end
    end
    ret
end


# example: julia opf.jl 66200 /path/to/JuMPSupplement/acpower
# JuMPSupplement is at https://github.com/mlubin/JuMPSupplement

#include("compare_jump.jl")

using JuMP
using Compat
using Ipopt
using ReverseDiffSparse2

type Bus
    bustype::Int
    name::ASCIIString
    voltage0::Float64
    angle0::Float64
    p_gen::Float64
    q_gen::Float64
    q_min::Float64
    q_max::Float64
    p_load::Float64
    q_load::Float64
    g_shunt::Float64
    b_shunt0::Float64
    b_shunt_min::Float64
    b_shunt_max::Float64
    b_dispatch::Float64
    area::Float64
end

type Branch
    from::Int
    to::Int
    branchtype::Int
    r::Float64
    x::Float64
    c::Float64
    tap0::Float64
    tap_min0::Float64
    tap_max0::Float64
    def0::Float64
    def_min::Float64
    def_max::Float64
    g::Float64
    b::Float64
end

function opf(bus, branch,
    bus_voltage_min, bus_voltage_max,
    branch_tap_min, branch_tap_max,
    p_gen_upper, p_gen_lower)

    mod = Model()

    nbus = length(bus)
    nbranch = length(branch)
    @defVar(mod, bus_voltage_min[bus[i].bustype] <= bus_voltage[i=1:nbus] <= bus_voltage_max[bus[i].bustype], start=1)
    @defVar(mod, bus[i].b_shunt_min <= bus_b_shunt[i=1:nbus] <= bus[i].b_shunt_max, start=bus[i].b_shunt0)
    @defVar(mod, bus_angle[1:nbus], start=0)

    @defVar(mod, branch_tap_min <= branch_tap[1:nbranch] <= branch_tap_max, start=1)
    @defVar(mod, branch[i].def_min <= branch_def[i=1:nbranch] <= branch[i].def_max, start = branch[i].def0)

    in_lines = [Int[] for i in 1:nbus] # indices of incoming branches
    out_lines = [Int[] for i in 1:nbus] # indices of outgoing branches
    @show nbranch
    @show nbus
    for i in 1:nbranch
        b = branch[i]
        push!(out_lines[b.from],i)
        push!(in_lines[b.to],i)
        @assert 1 <= b.to <= nbus
    end

    @defNLExpr(mod, Gself[k=1:nbus],
        bus[k].g_shunt +
        sum{branch[i].g*branch_tap[i]^2, i=out_lines[k]} +
        sum{ branch[i].g, i=in_lines[k]})

    @defNLExpr(mod, Gout[i=1:nbranch],
        (-branch[i].g*cos(branch_def[i])+branch[i].b*sin(branch_def[i]))*branch_tap[i])
    @defNLExpr(mod, Gin[i=1:nbranch],
        (-branch[i].g*cos(branch_def[i])-branch[i].b*sin(branch_def[i]))*branch_tap[i])

    @defNLExpr(mod, Bself[k=1:nbus],
        bus_b_shunt[k] +
        sum{branch[i].b*branch_tap[i]^2 + branch[i].c/2, i=out_lines[k]} +
        sum{ branch[i].b + branch[i].c/2, i=in_lines[k]})

    @defNLExpr(mod, Bin[i=1:nbranch],
        (branch[i].g*sin(branch_def[i])-branch[i].b*cos(branch_def[i]))*branch_tap[i])
    @defNLExpr(mod, Bout[i=1:nbranch],
        (-branch[i].g*sin(branch_def[i])-branch[i].b*cos(branch_def[i]))*branch_tap[i])

    # Minimize active power

    @setNLObjective(mod, Min, sum{ (bus[k].p_load +
        sum{ bus_voltage[k] * bus_voltage[branch[i].from] *
            (Gin[i] * cos(bus_angle[k] - bus_angle[branch[i].from]) +
             Bin[i] * sin(bus_angle[k] - bus_angle[branch[i].from])), i=in_lines[k] } +
        sum{ bus_voltage[k] * bus_voltage[branch[i].to] *
            (Gout[i] * cos(bus_angle[k] - bus_angle[branch[i].to]) +
             Bout[i] * sin(bus_angle[k] - bus_angle[branch[i].to])), i=out_lines[k] }
        + bus_voltage[k]^2*Gself[k] )^2,
        k in 1:nbus; bus[k].bustype == 2 || bus[k].bustype == 3})
    for k in 1:nbus
        bus[k].bustype == 0 || continue
        # p_load
        @addNLConstraint(mod, bus[k].p_gen - bus[k].p_load -
            sum{ bus_voltage[k] * bus_voltage[branch[i].from] *
                (Gin[i] * cos(bus_angle[k] - bus_angle[branch[i].from]) +
                 Bin[i] * sin(bus_angle[k] - bus_angle[branch[i].from])), i=in_lines[k] } -
            sum{ bus_voltage[k] * bus_voltage[branch[i].to] *
                (Gout[i] * cos(bus_angle[k] - bus_angle[branch[i].to]) +
                 Bout[i] * sin(bus_angle[k] - bus_angle[branch[i].to])), i=out_lines[k] }
            - bus_voltage[k]^2*Gself[k] == 0)
        # q_load
        @addNLConstraint(mod, bus[k].q_gen - bus[k].q_load -
            sum{ bus_voltage[k] * bus_voltage[branch[i].from] *
                (Gin[i] * sin(bus_angle[k] - bus_angle[branch[i].from]) -
                 Bin[i] * cos(bus_angle[k] - bus_angle[branch[i].from])), i=in_lines[k] } -
            sum{ bus_voltage[k] * bus_voltage[branch[i].to] *
                (Gout[i] * sin(bus_angle[k] - bus_angle[branch[i].to]) -
                 Bout[i] * cos(bus_angle[k] - bus_angle[branch[i].to])), i=out_lines[k] }
             + bus_voltage[k]^2*Bself[k] == 0)
    end
    for k in 1:nbus
        (bus[k].bustype == 2 || bus[k].bustype == 3) || continue
        # q_inj
        @addNLConstraint(mod, bus[k].q_min <= bus[k].q_load +
            sum{ bus_voltage[k] * bus_voltage[branch[i].from] *
                (Gin[i] * sin(bus_angle[k] - bus_angle[branch[i].from]) -
                 Bin[i] * cos(bus_angle[k] - bus_angle[branch[i].from])), i=in_lines[k] } +
            sum{ bus_voltage[k] * bus_voltage[branch[i].to] *
                (Gout[i] * sin(bus_angle[k] - bus_angle[branch[i].to]) -
                 Bout[i] * cos(bus_angle[k] - bus_angle[branch[i].to])), i=out_lines[k] } +
            - bus_voltage[k]^2*Bself[k] <= bus[k].q_max)
        # p_inj
        @addNLConstraint(mod, 0 <= bus[k].p_load +
            sum{ bus_voltage[k] * bus_voltage[branch[i].from] *
                (Gin[i] * cos(bus_angle[k] - bus_angle[branch[i].from]) +
                 Bin[i] * sin(bus_angle[k] - bus_angle[branch[i].from])), i=in_lines[k] } +
            sum{ bus_voltage[k] * bus_voltage[branch[i].to] *
                (Gout[i] * cos(bus_angle[k] - bus_angle[branch[i].to]) +
                 Bout[i] * sin(bus_angle[k] - bus_angle[branch[i].to])), i=out_lines[k] } +
            + bus_voltage[k]^2*Gself[k] <= p_gen_upper*bus[k].p_gen)
    end

    for i in 1:nbus
        #---- FREEZE THE REFERENCE BUS ANGLE TO ZERO.
        if bus[i].bustype == 3
            setLower(bus_angle[i], 0.0)
            setUpper(bus_angle[i], 0.0)
        end
        #---- FREEZE ANY DISPATCHABLE SHUNTS.
        if bus[i].b_dispatch == 0
            setLower(bus_b_shunt[i],bus[i].b_shunt0)
            setUpper(bus_b_shunt[i],bus[i].b_shunt0)
        end
    end

    for i in 1:nbranch
        #---- FREEZE ANY BRANCH TAPS.
        if branch[i].branchtype == 0 || branch[i].branchtype == 3
            setLower(branch_tap[i],1)
            setUpper(branch_tap[i],1)
        end
        #---- FREEZE CERTAIN PHASE SHIFTERS.
        if branch[i].branchtype != 4
            setLower(branch_def[i],getValue(branch_def[i]))
            setUpper(branch_def[i],getValue(branch_def[i]))
        end
    end

    return mod

end

function prepdata(bus, branch)

    buses = Array(Bus,size(bus,1))
    refbus = 0
    shunt = 0
    busmap = Dict()
    for i in 1:size(bus,1)
        @assert !haskey(busmap, bus[i,1])
        busmap[bus[i,1]] = i
        buses[i] = Bus(bus[i,2:end]...)
        # rescale
        buses[i].p_gen /= 100
        buses[i].q_gen /= 100
        buses[i].q_min /= 100
        buses[i].q_max /= 100
        buses[i].p_load /= 100
        buses[i].q_load /= 100
        if buses[i].bustype == 3
            refbus += 1
        end
        if buses[i].b_dispatch == 0
            shunt += 1
        end
    end
    @show refbus
    @show shunt

    branches = Array(Branch,size(branch,1))
    branchtap = 0
    phaseshift = 0
    for i in 1:size(branch,1)
        branches[i] = Branch(branch[i,2:end]...,0,0)
        @assert branch[i,1] == i
        branches[i].to = busmap[branches[i].to]
        branches[i].from = busmap[branches[i].from]
        branches[i].g = branches[i].r /(branches[i].r^2 + branches[i].x^2)
        branches[i].b = -branches[i].x /(branches[i].r^2 + branches[i].x^2)
        # rescale
        branches[i].def_min *= 3.14159/180
        branches[i].def_max *= 3.14159/180
        branches[i].def0 *= -3.14159/180
        if branches[i].branchtype == 0 || branches[i].branchtype == 3
            branchtap += 1
        end
        if branches[i].branchtype != 4
            phaseshift += 1
        end
    end
    @show branchtap
    @show phaseshift
    return branches, buses
end

numbuses = ARGS[1]
path = ARGS[2]
branches, buses = prepdata(readdlm("$path/IEEE$numbuses.bus"),readdlm("$path/IEEE$numbuses.branch"))

bus_voltage_min = @compat Dict(0 => 0.85, 1 => 0.85, 2 => 0.92, 3 => 0.99)
bus_voltage_max = @compat Dict(0 => 1.15, 1 => 1.15, 2 => 1.08, 3 => 1.01)
branch_tap_min = 0.85
branch_tap_max = 1.15

p_gen_upper = 1.10
p_gen_lower = 0.90
mod = opf(buses, branches,
    bus_voltage_min, bus_voltage_max,
    branch_tap_min, branch_tap_max,
    p_gen_upper, p_gen_lower)

setSolver(mod,IpoptSolver(max_iter=100))
#setSolver(mod,RDSSolver(IpoptSolver()))
println("solving")
solve(mod)
#compare_jump(mod)


using Distributions
using DistributionsAD
using ForwardDiff
using Gen
struct DistributionsForwardBacked{T} <: Gen.Distribution{T}
    to_dist
    has_arg_grads :: Vector{Bool}
    has_out_grad :: Bool
end

Gen.random(d::DistributionsForwardBacked, args...) = rand(d.to_dist(args...))
Gen.logpdf(d::DistributionsForwardBacked{T}, v::T, args...) where T  = Distributions.logpdf(d.to_dist(args...), v)
Gen.logpdf_grad(d::DistributionsForwardBacked{T}, v::T,
            args...) where T = ForwardDiff.gradient((v, args...) -> Distributions.logpdf(d.to_dist(args...), v), v, args...)
(d::DistributionsForwardBacked)(args...) = Gen.random(d, args...)
Gen.has_output_grad(d::DistributionsForwardBacked) = d.has_out_grad
Gen.has_argument_grads(d::DistributionsForwardBacked) = d.has_arg_grads


using ProgressMeter
using Dates

function main()

    start_t = now()
    out_dir  = "result"
    PARAMDAT = "PARAMDAT.json"
    fwrite   = "write"

    nval = 4
    Rd   = 287.0
    R    = 8.314

    xmax, ymax, nodes, vecAx, vecAy = read_allgrid()
    out_file_front, out_ext, restartnum, restart_file, init_small, norm_ok,
    time_integ, nt, dt, every_outnum, in_nt, dtau, cfl,
    init_rho, init_u, init_v, init_p, init_T, specific_heat_ratio, Rd, bdcon = input_para(PARAMDAT)

    init_p = init_rho*Rd*init_T

    Qbase, cellxmax, cellymax, restartnum = set_initQbase(xmax, ymax, restart_file, init_rho, init_u, init_v, init_p, init_T,
                                                        specific_heat_ratio, out_file_front, out_ext, out_dir, restartnum, Rd, nval)

    # init Delta_Qcon_hat
    volume = set_volume(nodes, cellxmax, cellymax)
    dx, dy = set_dx_lts(nodes, cellxmax, cellymax)
    reset_write(fwrite)

    # main loop
    print("threads num : ")
    println(Threads.nthreads())

    # check bd
    check_bd(bdcon)

    #throw(UndefVarError(:x))
    # main loop
    prog = Progress(nt,1)
    @time for t in 1:nt
        next!(prog)

        evalnum = t+restartnum
        if time_integ == "1"
            # exlicit scheme

            Qbase    = set_boundary(Qbase, cellxmax, cellymax, vecAx, vecAy, bdcon, Rd, specific_heat_ratio, nval)
            Qcon     = base_to_conservative(Qbase, cellxmax, cellymax, specific_heat_ratio, nval)
            Qcon_hat = setup_Qcon_hat(Qcon, cellxmax, cellymax, volume, nval)

            # initial_setup
            mu     = set_mu(Qbase, cellxmax, cellymax, specific_heat_ratio, Rd)
            lambda = set_lambda(Qbase, cellxmax, cellymax, mu, specific_heat_ratio, Rd)

            # RHS
            # advection_term
            #E_adv_hat, F_adv_hat = AUSM(Qbase,Qcon,cellxmax,cellymax,vecAx,vecAy,specific_heat_ratio)
            E_adv_hat, F_adv_hat = AUSM_plus(Qbase, Qcon, cellxmax, cellymax, vecAx, vecAy, specific_heat_ratio, volume, nval)

            # viscos_term
            E_vis_hat, F_vis_hat = central_diff(Qbase, Qcon, cellxmax, cellymax, mu, lambda,
                                                vecAx, vecAy, specific_heat_ratio, volume, Rd, nval)

            RHS = setup_RHS(cellxmax, cellymax, E_adv_hat, F_adv_hat, E_vis_hat, F_vis_hat, nval, volume)

            # time integral
            Qcon_hat = time_integration_explicit(dt, Qcon_hat, RHS, cellxmax, cellymax, nval)

            Qcon  = Qhat_to_Q(Qcon_hat, cellxmax, cellymax, volume, nval)
            Qbase = conservative_to_base(Qcon, cellxmax, cellymax, specific_heat_ratio, nval)

        elseif time_integ == "2"
            output_physicaltime(fwrite, t, dt)
            # Qcn
            Qbasen    = copy(Qbase)
            Qconn     = base_to_conservative(Qbasen, cellxmax, cellymax, specific_heat_ratio, nval)
            Qconn_hat = setup_Qcon_hat(Qconn, cellxmax, cellymax, volume, nval)

            Qbasem = copy(Qbase)

            for tau in 1:in_nt
                # LHS (A_adv_hat=jacobian)
                Qbasem = set_boundary(Qbasem, cellxmax, cellymax, vecAx, vecAy, bdcon, Rd, specific_heat_ratio, nval)

                Qcon     = base_to_conservative(Qbasem, cellxmax, cellymax, specific_heat_ratio, nval)
                Qcon_hat = setup_Qcon_hat(Qcon, cellxmax, cellymax, volume, nval)

                # initial_setup
                mu     = set_mu(Qbasem, cellxmax, cellymax, specific_heat_ratio, Rd)
                lambda = set_lambda(Qbasem, cellxmax, cellymax, mu, specific_heat_ratio, Rd)
                dtau   = set_lts(Qbasem, cellxmax, cellymax, mu, dx, dy, vecAx, vecAy, volume, specific_heat_ratio, cfl)

                # RHS
                #advection_term
                E_adv_hat, F_adv_hat = AUSM_plus(Qbasem, Qcon, cellxmax, cellymax, vecAx, vecAy, specific_heat_ratio, volume, nval)

                # viscos_term
                E_vis_hat, F_vis_hat = central_diff(Qbasem, Qcon, cellxmax, cellymax, mu, lambda,
                                                vecAx, vecAy, specific_heat_ratio, volume, Rd, nval)

                RHS = setup_RHS(cellxmax, cellymax, E_adv_hat, F_adv_hat, E_vis_hat, F_vis_hat, nval, volume)

                # lusgs_advection_term
                A_adv_hat_p, A_adv_hat_m, B_adv_hat_p, B_adv_hat_m, A_beta_shig, B_beta_shig = one_wave(Qbasem, Qcon, cellxmax, cellymax, vecAx, vecAy, specific_heat_ratio, volume, nval)
                # lusgs_viscos_term
                jalphaP, jbetaP = central_diff_jacobian(Qbasem, Qcon, cellxmax, cellymax, mu, lambda, vecAx, vecAy, specific_heat_ratio, volume, nval)

                # LUSGS
                delta_Q = zeros(cellxmax, cellymax, nval)

                ite = 0
                norm2 = zeros(nval)
                while true
                    delta_Q_temp = copy(delta_Q)

                    delta_Q = lusgs(dt, dtau, Qcon_hat, Qconn_hat, delta_Q, A_adv_hat_p,  A_adv_hat_m,  B_adv_hat_p,  B_adv_hat_m,  A_beta_shig,  B_beta_shig, jalphaP,  jbetaP, RHS, cellxmax, cellymax, volume, nval)

                    res   = set_res(delta_Q, delta_Q_temp, cellxmax, cellymax, nval)
                    norm2 = check_converge(res, RHS, cellxmax, cellymax, init_small, nval)

                    ite += 1
                    if norm2[1] < norm_ok && norm2[2] < norm_ok && norm2[3] < norm_ok && norm2[4] < norm_ok
                        break
                    end
                    if ite % 100 ==0
                        println(" now cal norm2 ")
                        println(norm2)
                    end
                end
                output_innertime(fwrite, tau, norm2, nval)

                for l in 1:nval
                    for j in 2:cellymax-1
                        for i in 2:cellxmax-1
                            Qcon_hat[i,j,l] = Qcon_hat[i,j,l] + delta_Q[i,j,l]
                        end
                    end
                end

                Qcon = Qhat_to_Q(Qcon_hat, cellxmax, cellymax, volume, nval)
                Qbasem = conservative_to_base(Qcon, cellxmax, cellymax, specific_heat_ratio, nval)

            end
            Qbase = copy(Qbasem)
        end

        if round(evalnum) % every_outnum == 0
            println("\n")
            println("nt_______________________________"*string(round(evalnum)))
            output_result(evalnum, Qbase, cellxmax, cellymax, specific_heat_ratio, out_file_front, out_ext, out_dir, Rd, nval)
        end

        check_divrege(Qbase, cellxmax, cellymax, Rd, fwrite)
    end

    end_t = now()
    output_fin(fwrite, start_t, end_t, nt, dt, in_nt, cellxmax, cellymax)
end

# -- main --
main()


using Oceananigans.Advection

#####
##### Momentum advection operators
#####

"""
    u∇u(i, j, k, grid, U)

Calculate the advection of momentum in the x-direction U·∇u

    1/Vᵘ * [δxᶠᵃᵃ(ℑxᶜᵃᵃ(Ax * u) * ℑxᶜᵃᵃ(u)) + δy_fca(ℑxᶠᵃᵃ(Ay * v) * ℑyᵃᶠᵃ(u)) + δz_fac(ℑxᶠᵃᵃ(Az * w) * ℑzᵃᵃᶠ(u))]

which will end up at the location `fcc`.
"""
@inline function div_ũu(i, j, k, grid, advection, U)
    return 1/Vᵃᵃᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, momentum_flux_uu, advection, U.u)    +
                                    δyᵃᶜᵃ(i, j, k, grid, momentum_flux_uv, advection, U.u, U.v) +
                                    δzᵃᵃᶜ(i, j, k, grid, momentum_flux_uw, advection, U.u, U.w))
end

"""
    u∇v(i, j, k, grid, U)

Calculates the advection of momentum in the y-direction U·∇v

    1/Vʸ * [δx_cfa(ℑyᵃᶠᵃ(Ax * u) * ℑxᶠᵃᵃ(v)) + δyᵃᶠᵃ(ℑyᵃᶜᵃ(Ay * v) * ℑyᵃᶜᵃ(v)) + δz_afc(ℑxᶠᵃᵃ(Az * w) * ℑzᵃᵃᶠ(w))]

which will end up at the location `cfc`.
"""
@inline function div_ũv(i, j, k, grid, advection, U)
    return 1/Vᵃᵃᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_vu, advection, U.u, U.v) +
                                    δyᵃᶠᵃ(i, j, k, grid, momentum_flux_vv, advection, U.v)    +
                                    δzᵃᵃᶜ(i, j, k, grid, momentum_flux_vw, advection, U.v, U.w))
end

"""
    u∇w(i, j, k, grid, U)

Calculates the advection of momentum in the z-direction U·∇w

    1/Vʷ * [δx_caf(ℑzᵃᵃᶠ(Ax * u) * ℑxᶠᵃᵃ(w)) + δy_acf(ℑzᵃᵃᶠ(Ay * v) * ℑyᵃᶠᵃ(w)) + δzᵃᵃᶠ(ℑzᵃᵃᶜ(Az * w) * ℑzᵃᵃᶜ(w))]

which will end up at the location `ccf`.
"""
@inline function div_ũw(i, j, k, grid, advection, U)
    return 1/Vᵃᵃᶠ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_wu, advection, U.u, U.w) +
                                    δyᵃᶜᵃ(i, j, k, grid, momentum_flux_wv, advection, U.v, U.w) +
                                    δzᵃᵃᶠ(i, j, k, grid, momentum_flux_ww, advection, U.w))
end


function Mass2DMatrix(mesh)
    return assembleScalar(mesh,localMass2D!,0.0)
end
function Mass2DMatrix(mesh,scalarFunc)
    return assembleScalar(mesh,localMassWithScalar2D!,scalarFunc)
end

function Laplace2DMatrix(mesh,farr,param::T) where T<:AbstractDummyParameter
    D = assembleScalar(mesh,localLaplace2D!,0.0)
    F = WeakScalar2D(mesh,farr)
    return D,F
end

function Darcy2DMatrix(mesh,farr,param::T) where T<:AbstractConstantParameter
    D = assembleScalar(mesh,localLaplace2D!,0.0)
    D = param.α*D
    F = WeakScalar2D(mesh,farr)
    return D,F
end
function Darcy2DMatrix(mesh,farr,param::T) where T<:AbstractVariableParameter
    D = assembleScalar(mesh,localDarcy2D!,param.α)
    F = WeakScalar2D(mesh,farr)
    return D,F
end
function Darcy2DMatrix_2b(mesh,farr,param::T,prob) where T<:AbstractVariableParameter
    D, F = assembleScalarandRHS_2b(mesh,localDarcy2D!,param.α,prob)
    return D, F
end

function AdvDiff2DMatrix(mesh,farr,param::T) where T<:AbstractConstantParameter
    D = assembleScalar(mesh,localLaplace2D!,0.0)
    A = assembleScalar(mesh,localAdvDiff2D!,param)
    F = WeakScalar2D(mesh,farr)
    Stiff = A + D/param.Pe
    return Stiff, F
end
function AdvDiff2DMatrix(mesh,farr,param::T) where T<:AbstractVariableParameter
    D = assembleScalar(mesh,localDarcy2D!,param.Pe)
    A = assembleScalar(mesh,localAdvDiff2D!,param)
    F = WeakScalar2D(mesh,farr)
    Stiff = A + D
    return Stiff, F
end

function Stokes2DMatrix(mesh::FluidMesh,prob,param::T) where T<:AbstractConstantParameter
    Stiff = assembleHalfFluid(mesh,localStokesConst2D!,param.μ)
    Fx    = WeakScalar2D(mesh,prob.bcval[:forcingX])
    Fy    = WeakScalar2D(mesh,prob.bcval[:forcingY])
    G     = zeros(length(mesh.xyp))
    F     = vcat(Fx,Fy,G)

    return Stiff, F
end
function Stokes2DMatrix(mesh::FluidMesh,prob,param::T) where T<:AbstractVariableParameter
    error("Stokes 2D variable viscosity not working right now")

    # old original -- THIS ONE IS WORKING (? need to test this)
    Stiff = assembleFullFluid_WRONG(mesh,localStokesVar2D_WRONG!,param.μ)
    Fx    = WeakScalar2D(mesh,prob.bcval[:forcingX])
    Fy    = WeakScalar2D(mesh,prob.bcval[:forcingY])
    G     = zeros(length(mesh.xyp))
    F     = vcat(Fx,Fy,G)
    return Stiff, F
end

function Brinkman2DMatrix(mesh::FluidMesh,prob,param::T) where T<:AbstractConstantParameter
    Stiff = assembleHalfFluid(mesh,localBrinkmanConst2D!,param.μ,param.α)
    Fx    = WeakScalar2D(mesh,prob.bcval[:forcingX])
    Fy    = WeakScalar2D(mesh,prob.bcval[:forcingY])
    G     = zeros(length(mesh.xyp))
    F     = vcat(Fx,Fy,G)

    return Stiff, F
end
function Brinkman2DMatrix(mesh::FluidMesh,prob,param::T) where T<:AbstractVariableParameter
    Stiff = assembleFullFluid(mesh,localBrinkmanVar2D!,param.μ,param.α)
    Fx    = WeakScalar2D(mesh,prob.bcval[:forcingX])
    Fy    = WeakScalar2D(mesh,prob.bcval[:forcingY])
    G     = zeros(length(mesh.xyp))
    F     = vcat(Fx,Fy,G)

    return Stiff, F
end

"""
Generates matrix operator to be used to solve brinkman-like equations
  -α1*lap(u) + α2*u + α3*grad(p) = f
"""
function BrinkmanMP2DMatrix(mesh::FluidMesh,prob,param::T) where T<:AbstractConstantParameter
	Stiff = assembleHalfFluid(mesh,localBrinkmanMPConst2D!,param.α1,param.α2,param.α3)
	Fx    = WeakScalar2D(mesh,prob.bcval[:forcingX])
	Fy    = WeakScalar2D(mesh,prob.bcval[:forcingY])
	G     = zeros(length(mesh.xyp))
	F     = vcat(Fx,Fy,G)

	return Stiff, F
end
function BrinkmanMP2DMatrix(mesh::FluidMesh,prob,param::T) where T<:AbstractVariableParameter
    Stiff = assembleMPB(mesh,localMPBVar2D!,param.α1,param.α2,param.α3)
	Fx    = WeakScalar2D(mesh,prob.bcval[:forcingX])
	Fy    = WeakScalar2D(mesh,prob.bcval[:forcingY])
	G     = zeros(length(mesh.xyp))
	F     = vcat(Fx,Fy,G)

	return Stiff, F
end

function LaplaceASMatrix(mesh,farr,param)
    D     = assembleScalar(mesh,localLaplaceConstAS!,param.κ)
    Stiff = D
    F     = WeakScalarAS(mesh, farr)
    return Stiff, F
end

function AdvDiffASMatrix(mesh,farr,param::T) where T<:AbstractVariableParameter
    D     = assembleScalar(mesh,localLaplaceVarAS!,param.Pe)
    A     = assembleScalar(mesh,localAdvDiffAS!,param)
    Stiff = A + D
    F     = WeakScalarAS(mesh,farr)
    return Stiff, F
end

function StokesASMatrix(mesh,prob,param::T) where T<:AbstractConstantParameter
    Stiff = assembleHalfFluid(mesh,localStokesConstAS!,param.μ)
    Fx    = WeakScalarAS(mesh,prob.bcval[:forcingX])
    Fy    = WeakScalarAS(mesh,prob.bcval[:forcingY])
    G     = zeros(Float64,length(mesh.xyp))
    F     = vcat(Fx,Fy,G)
    return Stiff, F
end
function StokesASMatrix(mesh,prob,param::T) where T<:AbstractVariableParameter
    Stiff = assembleFullFluid_WRONG(mesh,localStokesVarAS!,param.μ)
    Fx    = WeakScalarAS(mesh,prob.bcval[:forcingX])
    Fy    = WeakScalarAS(mesh,prob.bcval[:forcingY])
    G     = zeros(Float64,length(mesh.xyp))
    F     = vcat(Fx,Fy,G)
    return Stiff, F
end


@testset "Durou" begin
    p, q = synthetic_gradient(SynthSphere(50))
    Z = Durou()(p, q)
    @test Z[38,38] ≈ 1.5742826661339246
    @test Z[1,1] ≈ 0.44970037692631165
    @test Z[76,76] ≈ 47.228212708897075
end


function wk3odes(t::Float64,y::Vector{Float64},dy::Vector{Float64},mparams::WK3.ModelParams,k::Float64)
    # ventricular elastance and pressure

    E,dE = elastancefn(t*ts,mparams,k);
    Pv = E.*(y[2].*Vs.-mparams.V0)./Ps;

    # derivatives
    # println("ODEs, first loop:")
    if y[3] != 0
        dy[1] = ts/Ps*((1+mparams.Zc/mparams.R)*mparams.C)^-1*((1+mparams.Zc/mparams.R)*Qs*y[3]
            .+ mparams.C*mparams.Zc/mparams.R*(-E*Qs*y[3].+dE*(Vs*y[2]-mparams.V0)) .- Ps*y[1]/mparams.R);
    else
        dy[1] = ts/Ps*((1+mparams.Zc/mparams.R)*mparams.C)^-1.*(-Ps.*y[1]./mparams.R);
    end
    # println("ODEs, second loop:")
    if Pv > y[1] || (Pv < y[1] && y[3] > 0)
        dy[2] = -ts/Vs*Qs*y[3];
        dy[3] = ts/(Qs*(mparams.Rav+mparams.Zc))*(-E*Qs*y[3] .+ dE*(Vs*y[2].-mparams.V0) .-
            (1+mparams.Zc/mparams.R)*Qs/mparams.C*y[3] .+ Ps/(mparams.R*mparams.C)*y[1]);
    elseif Pv < y[1] && y[3] <= 0
        dy[2] = 0.;
        dy[3] = 0.;
    end

    # for i = 1:length(dy)
    #     if dy[i] > 1e6
    #         warn("dy/dt above tolerance at t = $t. dE/dt = $dE.")
    #     end
    # end
end


function constraint_converter_losses(pm::AbstractLPACModel, n::Int, cnd::Int, i::Int, a, b, c, plmax)
    pconv_ac = PowerModels.var(pm, n, cnd, :pconv_ac, i)
    pconv_dc = PowerModels.var(pm, n, cnd, :pconv_dc, i)
    iconv = PowerModels.var(pm, n, cnd, :iconv_ac, i)

    PowerModels.con(pm, n, cnd, :conv_loss)[i] = @constraint(pm.model, pconv_ac + pconv_dc == a + b*iconv)
end

function constraint_conv_transformer(pm::AbstractLPACModel, n::Int, cnd::Int, i::Int, rtf, xtf, acbus, tm, transformer)
    ptf_fr = PowerModels.var(pm, n, cnd, :pconv_tf_fr, i)
    qtf_fr = PowerModels.var(pm, n, cnd, :qconv_tf_fr, i)
    ptf_to = PowerModels.var(pm, n, cnd, :pconv_tf_to, i)
    qtf_to = PowerModels.var(pm, n, cnd, :qconv_tf_to, i)

    phi = PowerModels.var(pm, n, cnd, :phi, acbus)
    phi_vmf = PowerModels.var(pm, n, cnd, :phi_vmf, i)
    va = PowerModels.var(pm, n, cnd, :va, acbus)
    vaf = PowerModels.var(pm, n, cnd, :vaf, i)
    cs = PowerModels.var(pm, n, cnd, :cs_vaf,i)

    ztf = rtf + im*xtf
    if transformer
        ytf = 1/(rtf + im*xtf)
        gtf = real(ytf)
        btf = imag(ytf)
        c1, c2, c3, c4 = lpac_power_flow_constraints(pm.model, gtf, btf, phi, phi_vmf, va, vaf, ptf_fr, ptf_to, qtf_fr, qtf_to, tm, cs)
        c5 = constraint_cos_angle_diff_PWL(pm, n, cnd, cs, va, vaf)
        PowerModels.con(pm, n, cnd, :conv_tf_p_fr)[i] = c1
        PowerModels.con(pm, n, cnd, :conv_tf_q_fr)[i] = c2
        PowerModels.con(pm, n, cnd, :conv_tf_p_to)[i] = c3
        PowerModels.con(pm, n, cnd, :conv_tf_q_to)[i] = c4
    else
        PowerModels.con(pm, n, cnd, :conv_tf_p_fr)[i] = @constraint(pm.model, ptf_fr + ptf_to == 0)
        PowerModels.con(pm, n, cnd, :conv_tf_q_fr)[i] = @constraint(pm.model, qtf_fr + qtf_to == 0)
        @constraint(pm.model, va == vaf)
        @constraint(pm.model, (1+phi) == (1+phi_vmf))
    end
end

"constraints for a voltage magnitude transformer + series impedance"

function lpac_power_flow_constraints(model, g, b, phi_fr, phi_to, va_fr, va_to, p_fr, p_to, q_fr, q_to, tm, cs)

    c1 = @constraint(model, p_fr ==  g/(tm^2)*(1.0 + 2*phi_fr) + (-g/tm)*(cs + phi_fr + phi_to) + (-b/tm)*(va_fr-va_to))
    c2 = @constraint(model, q_fr == -b/(tm^2)*(1.0 + 2*phi_fr) - (-b/tm)*(cs + phi_fr + phi_to) + (-g/tm)*(va_fr-va_to))
    c3 = @constraint(model, p_to ==  g*(1.0 + 2*phi_to) + (-g/tm)*(cs + phi_fr + phi_to) + (-b/tm)*-(va_fr-va_to))
    c4 = @constraint(model, q_to == -b*(1.0 + 2*phi_to) - (-b/tm)*(cs + phi_fr + phi_to) + (-g/tm)*-(va_fr-va_to))
    return c1, c2, c3, c4
end

function constraint_conv_reactor(pm::AbstractLPACModel, n::Int, cnd::Int, i::Int, rc, xc, reactor)
    ppr_fr = PowerModels.var(pm, n, cnd, :pconv_pr_fr, i)
    qpr_fr = PowerModels.var(pm, n, cnd, :qconv_pr_fr, i)
    ppr_to = -PowerModels.var(pm, n, cnd, :pconv_ac, i)
    qpr_to = -PowerModels.var(pm, n, cnd, :qconv_ac, i)
    phi_vmc = PowerModels.var(pm, n, cnd, :phi_vmc, i)
    phi_vmf = PowerModels.var(pm, n, cnd, :phi_vmf, i)
    vac = PowerModels.var(pm, n, cnd, :vac, i)
    vaf = PowerModels.var(pm, n, cnd, :vaf, i)
    cs = PowerModels.var(pm, n, cnd, :cs_vac,i)

    phi_vmc_ub = JuMP.upper_bound(phi_vmc)
    ppr_to_ub = JuMP.upper_bound(PowerModels.var(pm, n, cnd)[:pconv_ac][i])
    qpr_to_ub = JuMP.upper_bound(PowerModels.var(pm, n, cnd)[:qconv_ac][i])
    Smax = sqrt(ppr_to_ub^2 + qpr_to_ub^2)
    zc = rc + im*xc
    if reactor
        yc = 1/(zc)
        gc = real(yc)
        bc = imag(yc)
        c1, c2, c3, c4 = lpac_power_flow_constraints(pm.model, gc, bc, phi_vmf, phi_vmc, vaf, vac, ppr_fr, ppr_to, qpr_fr, qpr_to, 1, cs)
        c5 = constraint_cos_angle_diff_PWL(pm, n, cnd, cs, vaf, vac)
        PowerModels.con(pm, n, cnd, :conv_pr_p)[i] = c3
        PowerModels.con(pm, n, cnd, :conv_pr_q)[i] = c4
        c6 = constraint_conv_capacity_PWL(pm, n, cnd, ppr_to, qpr_to, ppr_to_ub, qpr_to_ub, Smax)
   else
        PowerModels.con(pm, n, cnd, :conv_pr_p)[i] = @constraint(pm.model, ppr_fr + ppr_to == 0)
        PowerModels.con(pm, n, cnd, :conv_pr_q)[i] = @constraint(pm.model, qpr_fr + qpr_to == 0)
        @constraint(pm.model, vac == vaf)
        @constraint(pm.model, (1+phi_vmf) == (1+phi_vmc))

    end
end

function constraint_conv_filter(pm::AbstractLPACModel, n::Int, cnd::Int, i::Int, bv, filter)
    ppr_fr = PowerModels.var(pm, n, cnd, :pconv_pr_fr, i)
    qpr_fr = PowerModels.var(pm, n, cnd, :qconv_pr_fr, i)
    ptf_to = PowerModels.var(pm, n, cnd, :pconv_tf_to, i)
    qtf_to = PowerModels.var(pm, n, cnd, :qconv_tf_to, i)
    phi_vmf = PowerModels.var(pm, n, cnd, :phi_vmf, i)

    PowerModels.con(pm, n, cnd, :conv_kcl_p)[i] = @constraint(pm.model, ppr_fr + ptf_to == 0 )
    PowerModels.con(pm, n, cnd, :conv_kcl_q)[i] = @constraint(pm.model, qpr_fr + qtf_to + -bv*filter*(1+2*phi_vmf) == 0)
end

function constraint_converter_current(pm::AbstractLPACModel, n::Int, cnd::Int, i::Int, Umax, Imax)
    phi_vmc = PowerModels.var(pm, n, cnd, :phi_vmc, i)
    pconv_ac = PowerModels.var(pm, n, cnd, :pconv_ac, i)
    qconv_ac = PowerModels.var(pm, n, cnd, :qconv_ac, i)
    iconv = PowerModels.var(pm, n, cnd, :iconv_ac, i)

    # PowerModels.con(pm, n, cnd, :conv_i_sqrt)[i] = @NLconstraint(pm.model, pconv_ac^2 + qconv_ac^2 <= (1.0+phi_vmc)^2 * (iconv)^2)
    @constraint(pm.model, iconv <= Imax)
end

function variable_converter_filter_voltage(pm::AbstractLPACModel; kwargs...)
    variable_converter_filter_voltage_magnitude(pm; kwargs...)
    variable_converter_filter_voltage_angle_cs(pm; kwargs...)
    variable_converter_filter_voltage_angle(pm; kwargs...)
end

function variable_converter_internal_voltage(pm::AbstractLPACModel; kwargs...)
    variable_converter_internal_voltage_magnitude(pm; kwargs...)
    variable_converter_internal_voltage_angle_cs(pm; kwargs...)
    variable_converter_internal_voltage_angle(pm; kwargs...)
end

function variable_converter_filter_voltage_angle_cs(pm::AbstractLPACModel; nw::Int=pm.cnw, cnd::Int=pm.ccnd, bounded = true)
    PowerModels.var(pm, nw, cnd)[:cs_vaf] = @variable(pm.model,
    [i in PowerModels.ids(pm, nw, :convdc)], base_name="$(nw)_$(cnd)_cs_vaf",
    lower_bound =  0, # using the maximum bound allowed in cosine linearization. can be changed to bus angle +/- 10 degree?
    upper_bound =  1,
    start = 0
    )
end

function variable_converter_internal_voltage_angle_cs(pm::AbstractLPACModel; nw::Int=pm.cnw, cnd::Int=pm.ccnd, bounded = true)
    PowerModels.var(pm, nw, cnd)[:cs_vac] = @variable(pm.model,
    [i in PowerModels.ids(pm, nw, :convdc)], base_name="$(nw)_$(cnd)_cs_vac",
    lower_bound =  0, # using the maximum bound allowed in cosine linearization. can be changed to bus angle +/- 10 degree?
    upper_bound =  1,
    start = 0
    )
end

function variable_converter_filter_voltage_magnitude(pm::AbstractLPACModel; nw::Int=pm.cnw, cnd::Int=pm.ccnd, bounded = true)
    PowerModels.var(pm, nw, cnd)[:phi_vmf] = JuMP.@variable(pm.model,
            [i in PowerModels.ids(pm, nw, :convdc)], base_name="$(nw)_$(cnd)_phi_vmf",
            lower_bound = PowerModels.ref(pm, nw, :convdc, i, "Vmmin", cnd) - 1.0,
            upper_bound = PowerModels.ref(pm, nw, :convdc, i, "Vmmax", cnd) - 1.0,
            start = comp_start_value(ref(pm, nw, :convdc, i), "phi_start", cnd)
        )
end

function variable_converter_internal_voltage_magnitude(pm::AbstractLPACModel; nw::Int=pm.cnw, cnd::Int=pm.ccnd, bounded = true)
    PowerModels.var(pm, nw, cnd)[:phi_vmc] = JuMP.@variable(pm.model,
        [i in PowerModels.ids(pm, nw, :convdc)], base_name="$(nw)_$(cnd)_phi_vmc",
        lower_bound = PowerModels.ref(pm, nw, :convdc, i, "Vmmin", cnd) - 1.0,
        upper_bound = PowerModels.ref(pm, nw, :convdc, i, "Vmmax", cnd) - 1.0,
        start = comp_start_value(ref(pm, nw, :convdc, i), "phi_start", cnd)
        )
end

function constraint_conv_capacity_PWL(pm::AbstractLPACModel, n::Int, cnd::Int, ppr_to, qpr_to, Umax, Imax, Smax)
    np = 20 #no. of segments, can be passed as an argument later
    l = 0
    for i = 1:np
        a= Smax*sin(l)
        b = Smax*cos(l)
        c6 = @constraint(pm.model, a*ppr_to + b*qpr_to <= Smax^2) #current and voltage bounds to be proper to use Umax*Imax because Umax*Imax == Smax
        l = l + 2*pi/np
    end
end

function constraint_cos_angle_diff_PWL(pm::AbstractLPACModel, n::Int, cnd::Int, cs, va_fr, va_to)
    nb = 20 #no. of segments, can be passed as an argument later
    l = -pi/6
    h = pi/6
    inc = (h-l)/(nb+1)
    a = l + inc
    diff = va_fr - va_to
    for i = 1:nb
        c5 = @constraint(pm.model, cs <= -sin(a)*(diff-a) + cos(a))
        a = a + inc
    end
end

function add_dcconverter_voltage_setpoint(sol, pm::AbstractLPACModel)
    PowerModels.add_setpoint!(sol, pm, "convdc", "vmconv", :phi_vmc, status_name="islcc", inactive_status_value = 4, scale = (x,item,cnd) -> 1.0+x)
    PowerModels.add_setpoint!(sol, pm, "convdc", "vmfilt", :phi_vmf, status_name="islcc", inactive_status_value = 4, scale = (x,item,cnd) -> 1.0+x)
    PowerModels.add_setpoint!(sol, pm, "convdc", "vaconv", :vac, status_name="islcc", inactive_status_value = 4)
    PowerModels.add_setpoint!(sol, pm, "convdc", "vafilt", :vaf, status_name="islcc", inactive_status_value = 4)
end


"""
The first part of make transform matrix, usually used by the `learn_paths` function to
save time and computing resources.
"""
function make_transform_fac end

"""
Using Cholesky decomposition to calculate the transformation matrix from S to C
or from C to S.
"""
function make_transform_matrix end

"""
    make_transform_fac(X::SparseMatrixCSC)

Calculate the first step of Cholesky decomposition for sparse matrices.
"""
function make_transform_fac(
    X::SparseMatrixCSC;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX, shift = shift)
    else
        # convert value to Float64
        # otherwise multiplier would raise error
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    fac
end

"""
    make_transform_fac(X::Matrix)

Calculate the first step of Cholesky decomposition for dense matrices.
"""
function make_transform_fac(
    X::Matrix;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX + shift * I)
    else
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    fac
end

"""
    make_transform_matrix(fac::Union{LinearAlgebra.Cholesky, SuiteSparse.CHOLMOD.Factor}, X::Union{SparseMatrixCSC, Matrix}, Y::Union{SparseMatrixCSC, Matrix})

Second step in calculating the Cholesky decomposition for the transformation matrix.
"""
function make_transform_matrix(
    fac::Union{LinearAlgebra.Cholesky,SuiteSparse.CHOLMOD.Factor},
    X::Union{SparseMatrixCSC,Matrix},
    Y::Union{SparseMatrixCSC,Matrix};
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    make_transform_matrix(X::SparseMatrixCSC, Y::Matrix)

Use Cholesky decomposition to calculate the transformation matrix from X to Y,
where X is a sparse matrix and Y is a dense matrix.

# Obligatory Arguments
- `X::SparseMatrixCSC`: the X matrix, where X is a sparse matrix
- `Y::Matrix`: the Y matrix, where Y is a dense matrix

# Optional Arguments
- `method::Symbol = :additive`: whether :additive or :multiplicative decomposition is required
- `shift::Float64 = 0.02`: shift value for :additive decomposition
- `multiplier::Float64 = 1.01`: multiplier value for :multiplicative decomposition
- `output_format::Symbol = :auto`: to force output format to dense(:dense) or sparse(:sparse), make it auto(:auto) to determined by the program
- `sparse_ratio::Float64 = 0.05`: the ratio to decide whether a matrix is sparse
- `verbose::Bool = false`: if true, more information will be printed out

# Examples
```julia
# additive mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :additive,
    shift = 0.02,
    verbose = false)

# multiplicative mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :multiplicative,
    multiplier = 1.01,
    verbose = false)

# further control of sparsity ratio
JudiLing.make_transform_matrix(
  ...
    output_format = :auto,
    sparse_ratio = 0.05,
  ...)
```
"""
function make_transform_matrix(
    X::SparseMatrixCSC,
    Y::Matrix;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX, shift = shift)
    else
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    # M is in sparse format
    # but sometimes it is actually a dense matrix
    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    make_transform_matrix(X::Matrix, Y::Union{SparseMatrixCSC, Matrix})

Use the Cholesky decomposition to calculate the transformation matrix from X to Y,
where X is a dense matrix and Y is either a dense matrix or a sparse matrix.

# Obligatory Arguments
- `X::Matrix`: the X matrix, where X is a dense matrix
- `Y::Union{SparseMatrixCSC, Matrix}`: the Y matrix, where Y is either a sparse or a dense matrix

# Optional Arguments
- `method::Symbol = :additive`: whether :additive or :multiplicative decomposition is required
- `shift::Float64 = 0.02`: shift value for :additive decomposition
- `multiplier::Float64 = 1.01`: multiplier value for :multiplicative decomposition
- `output_format::Symbol = :auto`: to force output format to dense(:dense) or sparse(:sparse), make it auto(:auto) to determined by the program
- `sparse_ratio::Float64 = 0.05`: the ratio to decide whether a matrix is sparse
- `verbose::Bool = false`: if true, more information will be printed out

# Examples
```julia
# additive mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :additive,
    shift = 0.02,
    verbose = false)

# multiplicative mode
JudiLing.make_transform_matrix(
    C,
    S,
    method=:multiplicative,
    multiplier = 1.01,
    verbose = false)

# further control of sparsity ratio
JudiLing.make_transform_matrix(
    ...
    output_format = :auto,
    sparse_ratio = 0.05,
    ...)
```
"""
function make_transform_matrix(
    X::Matrix,
    Y::Union{SparseMatrixCSC,Matrix};
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX + shift * I)
    else
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    # M is in sparse format
    # but sometimes it is actually a dense matrix
    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    make_transform_matrix(X::SparseMatrixCSC, Y::SparseMatrixCSC)

Use the Cholesky decomposition to calculate the transformation matrix from X to Y,
where X is a sparse matrix and Y is a sparse matrix.

# Obligatory Arguments
- `X::SparseMatrixCSC`: the X matrix, where X is a sparse matrix
- `Y::SparseMatrixCSC`: the Y matrix, where Y is a sparse matrix

# Optional Arguments
- `method::Symbol = :additive`: whether :additive or :multiplicative decomposition is required
- `shift::Float64 = 0.02`: shift value for :additive decomposition
- `multiplier::Float64 = 1.01`: multiplier value for :multiplicative decomposition
- `output_format::Symbol = :auto`: to force output format to dense(:dense) or sparse(:sparse), make it auto(:auto) to determined by the program
- `sparse_ratio::Float64 = 0.05`: the ratio to decide whether a matrix is sparse
- `verbose::Bool = false`: if true, more information will be printed out

# Examples
```julia
# additive mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :additive,
    shift = 0.02,
    verbose = false)

# multiplicative mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :multiplicative,
    multiplier = 1.01,
    verbose = false)

# further control of sparsity ratio
JudiLing.make_transform_matrix(
    ...
    output_format = :auto,
    sparse_ratio = 0.05,
    ...)
```
"""
function make_transform_matrix(
    X::SparseMatrixCSC,
    Y::SparseMatrixCSC;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX, shift = shift)
    else
        # convert value to Float64
        # otherwise multiplier would raise error
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    # M is in sparse format
    # but sometimes it is actually a dense matrix
    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    format_matrix(M::Union{SparseMatrixCSC, Matrix}, output_format=:auto)

Convert output matrix format to either a dense matrix or a sparse matrix.
"""
function format_matrix(
    M::Union{SparseMatrixCSC,Matrix},
    output_format = :auto;
    sparse_ratio = 0.05,
    verbose = false,
)

    if output_format == :dense
        verbose && println("Returning a dense matrix format")
        Array(M)
    elseif output_format == :sparse
        verbose && println("Returning a sparse matrix format")
        sparse(M)
    else
        verbose && print("Auto mode: ")
        if is_truly_sparse(M, threshold = sparse_ratio, verbose = verbose)
            verbose && println("Returning a sparse matrix format")
            return M
        else
            verbose && println("Returning a dense matrix format")
            return Array(M)
        end
    end
end


@inline near(I::CartesianIndex,a=0) = (2I-2oneunit(I)):(2I-oneunit(I)-δ(a,I))

function restrictML(b::Array{Float64,m}) where m
    N = ntuple(i-> i==m ? m-1 : 1+size(b,i)÷2, m)
    a = zeros(N)
    restrictL!(a,b)
    Poisson(a)
end
@fastmath function restrictL!(a::Array{Float64,m},b::Array{Float64,m}) where m
    @inbounds for i ∈ 1:m-1, I ∈ inside(size(a)[1:m-1])
        a[I,i] = 0.5sum(@inbounds(b[J,i]) for J ∈ near(I,i))
    end
end

@fastmath restrict!(a::Array{Float64},b::Array{Float64}) = @inbounds @simd for I ∈ inside(a)
    a[I] = sum(@inbounds(b[J]) for J ∈ near(I))
end

prolongate!(a::Array{Float64},b::Array{Float64}) = @inbounds for I ∈ inside(b)
    @simd for J ∈ near(I)
        a[J] = b[I]
end;end

@inline divisible(N) = mod(N,2)==0 && N>4

struct MultiLevelPoisson{N,M} <: AbstractPoisson{N,M}
    levels :: Vector{Poisson{N,M}}
    n :: Vector{Int16}
    function MultiLevelPoisson(L::Array{Float64,n}) where n
        levels = [Poisson(L)]
        while all(size(levels[end].x) .|> divisible)
            push!(levels,restrictML(levels[end].L))
        end
        text = "MultiLevelPoisson requires size=a2ⁿ, where a<31, n>2"
        @assert (length(levels)>2 && all(size(levels[end].x).<31)) text
        new{n-1,n}(levels,[])
    end
end
function update!(ml::MultiLevelPoisson,L::Array)
    update!(ml.levels[1],L)
    for l ∈ 2:length(ml.levels)
        restrictL!(ml.levels[l].L,ml.levels[l-1].L)
        set_diag!(ml.levels[l])
    end
end

function Vcycle!(ml::MultiLevelPoisson;l=1)
    fine,coarse = ml.levels[l],ml.levels[l+1]
    # set up coarse level
    GS!(fine,it=0)
    restrict!(coarse.r,fine.r)
    fill!(coarse.x,0.)
    # solve coarse (with recursion if possible)
    l+1<length(ml.levels) && Vcycle!(ml,l=l+1)
    GS!(coarse,it=2)
    # correct fine
    prolongate!(fine.ϵ,coarse.x)
    increment!(fine)
end

mult(ml::MultiLevelPoisson,x) = mult(ml.levels[1],x)

function solver!(x::Array{Float64,n},ml::MultiLevelPoisson{n},b::Array{Float64,n};log=false,tol=1e-3,itmx=32) where n
    p = ml.levels[1]
    p.x .= x
    residual!(p,b); r₂ = L₂(p.r)
    log && (res = [r₂])
    nᵖ=0
    while r₂>tol && nᵖ<itmx
        Vcycle!(ml)
        GS!(p,it=2); r₂ = L₂(p.r)
        5tol>r₂>tol && (GS!(p,it=2); r₂ = L₂(p.r))
        log && push!(res,r₂)
        nᵖ+=1
    end
    x .= p.x
    push!(ml.n,nᵖ)
    log && return res
end


@defcomp SulphateForcing begin
    region = Index()

    se0_sulphateemissionsbase = Parameter(index=[region], unit="TgS/year")
    pse_sulphatevsbase = Parameter(index=[time, region], unit="%")
    se_sulphateemissions = Variable(index=[time, region], unit="TgS/year")
    area = Parameter(index=[region], unit ="km^2")

    sfx_sulphateflux = Variable(index=[time, region], unit="TgS/km^2/yr")

    d_sulphateforcingbase = Parameter(unit="W/m2", default=-0.46666666666666673)
    ind_slopeSEforcing_indirect = Parameter(unit="W/m2", default=-0.4000000000000001)
    nf_naturalsfx = Parameter(index=[region], unit="TgS/km^2/yr")

    fs_sulphateforcing = Variable(index=[time, region], unit="W/m2")

    function run_timestep(p, v, d, tt)
        bigSFX0 = p.se0_sulphateemissionsbase ./ p.area

        for rr in d.region
            v.se_sulphateemissions[tt, rr] = p.se0_sulphateemissionsbase[rr] * p.pse_sulphatevsbase[tt, rr] / 100

            # Eq.17 from Hope (2006) - sulfate flux
            v.sfx_sulphateflux[tt,rr] = v.se_sulphateemissions[tt,rr] / p.area[rr]
            # Update for Eq. 18 from Hope (2009) - sulfate radiative forcing effect
            bigSFD0 = p.d_sulphateforcingbase * bigSFX0[rr] / (sum(bigSFX0 .* p.area) / sum(p.area))
            fsd_term = bigSFD0 * v.sfx_sulphateflux[tt,rr] / bigSFX0[rr]
            fsi_term = p.ind_slopeSEforcing_indirect/log(2) * log((p.nf_naturalsfx[rr] + v.sfx_sulphateflux[tt, rr]) / p.nf_naturalsfx[rr])

            v.fs_sulphateforcing[tt, rr] = fsd_term + fsi_term
        end
    end
end


function gaussian(windowLength::Integer, sigma::Float64 = 1.0)
  samplingStep = 6/(windowLength-1)
  return [exp(-((x/sigma)^2)) for x in -3:samplingStep:3]
end


# Experimental threads example. WIP!

using CmdStan
#using Distributed
using Statistics
using StatsBase: sample

ProjDir = mktempdir()
cd(ProjDir) #do

bernoullimodel = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";
n = 10;
observeddata = Dict("N" => n, "y" => sample([0,1],n))

sm = Stanmodel(name="bernoulli", model=bernoullimodel;
  output_format=:namedtuple);

println("\nThreads loop\n")
p1 = 16 # p1 is the number of models to fit
estimates = Vector(undef, p1)
Threads.@threads for i in 1:p1
    pdir = pwd()
    while ispath(pdir)
        pdir = tempname()
    end
    println(pdir)

    new_model= deepcopy(sm)
    new_model.pdir = pdir
    new_model.tmpdir = joinpath(splitpath(pdir)...,"tmp")

    mkpath(new_model.tmpdir)

    CmdStan.update_model_file(joinpath(new_model.tmpdir, "$(new_model.name).stan"), strip(new_model.model))

    rc, samples, cnames = stan(new_model, observeddata, new_model.pdir;
    #  summary=false
    );

    if rc == 0
      estimates[i] = [mean(reshape(samples.theta, 4000)), std(reshape(samples.theta, 4000))]
    end

    #rm(pdir; force=true, recursive=true)
end

estimates |> display

#end


# SPDX-License-Identifier: X11
# 2020-11-20
# 1% (200pt)

const interest = big(101)//100

f(i, b, t) = b >= t ? i : f(i + 1, floor(b * interest), t)

(println ∘ f)(0, 100, parse(Int, readline()))


using LaTeXStrings
import AlgebraicMultigrid
import DPFEHM
import IterativeSolvers
import JLD
import LinearAlgebra
import PyPlot
import Statistics

function solve_cg_amg(A, b)
	Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.smoothed_aggregation(A))
	x, ch = IterativeSolvers.cg(A, b; maxiter=10000, reltol=1e-12, log=true, Pl=Pl)
	return x, ch.data[:resnorm]
end

function solve_cholesky(A, b)
	Af = LinearAlgebra.cholesky(A)
	x = Af \ b
	return x, [sqrt(sum((A * x .- b) .^ 2) / sum(b .^ 2))]
end

function setup_problem(n)
	mins = [0, 0, 0]#meters
	maxs = [100, 100, 100]#meters
	ns = [n, n, n]
	coords, neighbors, areasoverlengths, _ = DPFEHM.regulargrid3d(mins, maxs, ns)
	Qs = zeros(size(coords, 2))
	dirichletnodes = Int[]
	dirichleths = zeros(size(coords, 2))
	for i = 1:size(coords, 2)
		if coords[1, i] == mins[1]
			push!(dirichletnodes, i)
			dirichleths[i] = 1.0
		elseif coords[1, i] == maxs[1]
			push!(dirichletnodes, i)
			dirichleths[i] = 0.0
		end
	end
	Ks = ones(length(neighbors))
	args = (zeros(length(Qs) - length(dirichletnodes)), Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, ones(length(Qs)), ones(length(Qs)))
	b = -DPFEHM.groundwater_residuals(args...)
	A = DPFEHM.groundwater_h(args...)
	return A, b, (neighbors, areasoverlengths, Ks, Qs, dirichletnodes, dirichleths)
end

for i = 1:10#this loop gets compilation out of the way
	A, b, _ = setup_problem(2 ^ 2)
	solve_cholesky(A, b)
	solve_cg_amg(A, b)
end
if !isfile("porous_media_results.jld")
	ns = 2 .^ (2:7)
	num_solves = 10
	cholesky_times = zeros(length(ns), num_solves)
	cg_times = zeros(length(ns), num_solves)
	for (j, n) in enumerate(ns)
		A, b = setup_problem(n)
		these_cholesky_times = Float64[]
		for i = 1:num_solves
			@show n, i
			t = @elapsed solve_cholesky(A, b)
			@show t
			cholesky_times[j, i] = t
			t = @elapsed solve_cg_amg(A, b)
			@show t
			cg_times[j, i] = t
		end
	end
	JLD.save("porous_media_results.jld", "cholesky_times", cholesky_times, "cg_times", cg_times, "ns", ns)
end

function plotit(cholesky_times, cg_times, ns, ax)
	ax.loglog(ns, Statistics.median(cholesky_times; dims=2), ".", label="Cholesky", color="C0", ms=10, alpha=0.75)
	ax.loglog(ns, Statistics.median(cg_times; dims=2), ".", label="CG+AMG", color="C1", ms=10, alpha=0.75)
	for i = 1:size(cholesky_times, 1)
		ax.plot([ns[i], ns[i]], sort(cholesky_times[i, :])[[2, end - 1]], color="C0", alpha=0.75)
		ax.plot([ns[i], ns[i]], sort(cg_times[i, :])[[2, end - 1]], color="C1", alpha=0.75)
	end
	ax.set(xlabel="Degrees of freedom", ylabel="Time [s]")
	ax.legend()
end

cholesky_times, cg_times, ns = JLD.load("porous_media_results.jld", "cholesky_times", "cg_times", "ns")
fig, axs = PyPlot.subplots(1, 3; figsize=(12, 4))
plotit(cholesky_times, cg_times, ns .^ 3, axs[3])

function setup_problem_2d(n)
	mins = [0, 0]#meters
	maxs = [100, 100]#meters
	ns = [n, n]
	coords, neighbors, areasoverlengths, _ = DPFEHM.regulargrid2d(mins, maxs, ns, 1.0)
	Qs = zeros(size(coords, 2))
	dirichletnodes = Int[]
	dirichleths = zeros(size(coords, 2))
	for i = 1:size(coords, 2)
		if coords[1, i] == mins[1]
			push!(dirichletnodes, i)
			dirichleths[i] = 1.0
		elseif coords[1, i] == maxs[1]
			push!(dirichletnodes, i)
			dirichleths[i] = 0.0
		end
	end
	Ks = ones(length(neighbors))
	args = (zeros(length(Qs) - length(dirichletnodes)), Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, ones(length(Qs)), ones(length(Qs)))
	b = -DPFEHM.groundwater_residuals(args...)
	A = DPFEHM.groundwater_h(args...)
	return A, b, (neighbors, areasoverlengths, Ks, Qs, dirichletnodes, dirichleths)
end

for i = 1:10#this loop gets compilation out of the way
	A, b, _ = setup_problem_2d(2 ^ 2)
	solve_cholesky(A, b)
	solve_cg_amg(A, b)
end
if !isfile("porous_media_results_2d.jld")
	ns = 2 .^ (3:13)
	num_solves = 10
	cholesky_times = zeros(length(ns), num_solves)
	cg_times = zeros(length(ns), num_solves)
	for (j, n) in enumerate(ns)
		A, b = setup_problem_2d(n)
		these_cholesky_times = Float64[]
		for i = 1:num_solves
			@show n, i
			t = @elapsed solve_cholesky(A, b)
			@show t
			cholesky_times[j, i] = t
			t = @elapsed solve_cg_amg(A, b)
			@show t
			cg_times[j, i] = t
		end
	end
	JLD.save("porous_media_results_2d.jld", "cholesky_times", cholesky_times, "cg_times", cg_times, "ns", ns)
end

cholesky_times, cg_times, ns = JLD.load("porous_media_results_2d.jld", "cholesky_times", "cg_times", "ns")
plotit(cholesky_times, cg_times, ns .^ 2, axs[2])

function setup_problem_1d(n)
	mins = [0]#meters
	maxs = [100]#meters
	ns = [n]
	coords, neighbors, areasoverlengths, _ = DPFEHM.regulargrid1d(mins, maxs, ns, 1.0, 1.0)
	Qs = zeros(size(coords, 2))
	dirichletnodes = Int[]
	dirichleths = zeros(size(coords, 2))
	for i = 1:size(coords, 2)
		if coords[1, i] == mins[1]
			push!(dirichletnodes, i)
			dirichleths[i] = 1.0
		elseif coords[1, i] == maxs[1]
			push!(dirichletnodes, i)
			dirichleths[i] = 0.0
		end
	end
	Ks = ones(length(neighbors))
	args = (zeros(length(Qs) - length(dirichletnodes)), Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, ones(length(Qs)), ones(length(Qs)))
	b = -DPFEHM.groundwater_residuals(args...)
	A = DPFEHM.groundwater_h(args...)
	return A, b, (neighbors, areasoverlengths, Ks, Qs, dirichletnodes, dirichleths)
end

for i = 1:10#this loop gets compilation out of the way
	A, b, _ = setup_problem_1d(2 ^ 2)
	solve_cholesky(A, b)
	solve_cg_amg(A, b)
end
if !isfile("porous_media_results_1d.jld")
	ns = 10 .^ (1:8)
	num_solves = 10
	cholesky_times = zeros(length(ns), num_solves)
	cg_times = zeros(length(ns), num_solves)
	for (j, n) in enumerate(ns)
		A, b = setup_problem_1d(n)
		these_cholesky_times = Float64[]
		for i = 1:num_solves
			@show n, i
			t = @elapsed solve_cholesky(A, b)
			@show t
			cholesky_times[j, i] = t
			t = @elapsed solve_cg_amg(A, b)
			@show t
			cg_times[j, i] = t
		end
	end
	JLD.save("porous_media_results_1d.jld", "cholesky_times", cholesky_times, "cg_times", cg_times, "ns", ns)
end

cholesky_times, cg_times, ns = JLD.load("porous_media_results_1d.jld", "cholesky_times", "cg_times", "ns")
plotit(cholesky_times, cg_times, ns, axs[1])
for i = 1:3
	axs[i].set_title("$(i)D")
end
fig.tight_layout()
display(fig)
println()
fig.savefig("porous_media_comparison.pdf")
PyPlot.close(fig)


using JuMP
using ParameterJuMP
using Polyhedra

mutable struct Conitope{T, VT <: AbstractVector{T}, D<:Polyhedra.FullDim}
    d::D
    points::Vector{VT}
    factory::Union{Nothing, JuMP.OptimizerFactory}
    model::Union{Nothing, JuMP.Model}
    z::Union{Nothing, Vector{ParameterJuMP.ParameterRef}}
    t_0::Union{Nothing, JuMP.VariableRef}
    function Conitope{T, VT, D}(
        d::Polyhedra.FullDim, points::Polyhedra.PointIt,
        factory::Union{Nothing, JuMP.OptimizerFactory}=nothing) where {T, VT, D}
        new{T, VT, D}(Polyhedra.FullDim_convert(D, d),
                      Polyhedra.lazy_collect(points), factory, nothing, nothing, nothing)
    end
end
function Conitope(d::Polyhedra.FullDim, points::Polyhedra.PointIt, args...)
    return Conitope{Polyhedra.coefficient_type(eltype(points)), eltype(points),
                    typeof(d)}(d, points, args...)
end

Polyhedra.FullDim(v::Conitope) = v.d

function _build_model(brp::Conitope)
    n = length(brp.points)
    brp.model = ParameterJuMP.ModelWithParams(brp.factory)
    @variable(brp.model, t[1:n] ≥ 0)
    @constraint(brp.model, sum(t) ≤ 1)
    return sum(t[i] .* brp.points[i] for i in 1:n)
end
function _parametrized_model(brp::Conitope, v)
    hull = _build_model(brp)
    brp.z = ParameterJuMP.add_parameters(brp.model, collect(v))
    @constraint(brp.model, brp.z .≤ hull)
end
function _ratio_model(brp::Conitope, v)
    hull = _build_model(brp)
    brp.t_0 = @variable(brp.model)
    @constraint(brp.model, brp.t_0 .* v .≤ hull)
end

function _ei(i::Integer, d::Integer, T::Type)
    e = zeros(T, d)
    e[i] = one(T)
    return e
end
function Polyhedra.polyhedron(
    c::Conitope{T},
    lib::Polyhedra.Library=Polyhedra.default_library(Polyhedra.FullDim(c),
                                                     T)) where T
    n = fulldim(c)
    p = polyhedron(vrep([Ray(-_ei(i, n, T)) for i in 1:n]) + vrep(c.points),
                   lib)
    return p ∩ (Matrix(-LinearAlgebra.I, n, n) * p)
end


include("../src/HssMatrices.jl")
using .HssMatrices
using LinearAlgebra
using BenchmarkTools
using Random
Random.seed!(123)

### run benchmarks on Cauchy matrix
K(x,y) = (x-y) != 0 ? 1/(x-y) : 10000.
A = [ K(x,y) for x=-1:0.001:1, y=-1:0.001:1];
b = randn(size(A,2), 5);

# test the simple implementation of cluster trees
m, n = size(A)
#lsz = 64;
lsz = 64;
rcl = bisection_cluster(1:m, leafsize=lsz)
ccl = bisection_cluster(1:n, leafsize=lsz)

hssA = compress(A, rcl, ccl);

# time conversion to dense
println("Benchmarking full...")
@btime full(hssA);

# time access
println("Benchmarking getindex...")
ii = randperm(100); jj = randperm(100)
@btime Aij = hssA[ii,jj];

# time compression
println("Benchmarking compression...")
@btime hssA = compress(A, rcl, ccl);

# time compression
println("Benchmarking randomized compression...")
@btime hssA = randcompress_adaptive(A, rcl, ccl);

println("Benchmarking re-compression...")
hssB = copy(hssA)
@btime hssA = recompress!(hssB; atol=1e-3, rtol=1e-3);

println("Benchmarking proper...")
hssB = hssA + hssA
@btime orthonormalize_generators!(hssB)

println("Benchmarking addition...")
@btime hssC = hssA + hssA

# time matvec
println("Benchmarking matvec...")
x = randn(size(A,2), 10);
@btime y = hssA*x

# time ulvfactsolve
println("Benchmarking ulvfactsolve...")
b = randn(size(A,2), 10);
@btime x = ulvfactsolve(hssA, b);

# time hssldivide
println("Benchmarking hssldivide...")
hssX = compress(1.0*Matrix(I, n, n), ccl, ccl, atol=1e-3, rtol=1e-3); # this should probably be one constructor
@btime hssC = ldiv!(copy(hssA), hssX);


using CellListMap
using StaticArrays
import Random

#
# In this test we compute the average displacement of the x coordinates of the particles
#
function average_displacement(;N=100_000,parallel=true,x=nothing)

    # Number of particles, sides and cutoff
    sides = @SVector [250,250,250]
    cutoff = 10
    box = Box(sides, cutoff)

    # Particle positions
    Random.seed!(321)
    if x === nothing
        x = [ sides .* rand(SVector{3,Float64}) for i in 1:N ]
    end

    # Initialize auxiliary linked lists
    cl = CellList(x, box, parallel=parallel)

    # Function to be evalulated for each pair: sum of displacements on x
    f(x, y, avg_dx) = avg_dx + abs(x[1] - y[1])

    avg_dx = (N / (N * (N - 1) / 2)) * map_pairwise(
        (x, y, i, j, d2, avg_dx) -> f(x, y, avg_dx),
        0.,box,cl,
        parallel=parallel
    )

    return avg_dx
end


using CommutativeRings: hensel_lift, partsums, subset, remove_subset!, allgcdx

let x = monom(ZZ{Int}[:x]), u = x^8 + x^6 -3x^4 -3x^3 +8x^2 + 2x - 5

@testset "integer polynomials over $T" for T in (Int64, BigInt)
    x = monom(ZZ{T}[:x])
    b = x^2 + 5x + 1
    a = x + 2
    c = a * b
    @test factor(c) == [a =>1; b=>1]
end

@testset "coeffbounds" begin
    @test_throws ArgumentError coeffbounds(zero(u), 0)
    @test_throws ArgumentError coeffbounds(10one(u), 1)
    @test coeffbounds(10one(u), 0) == [10]
    @test coeffbounds(u, 0) == [1]
    @test coeffbounds(u, 1) == [5, 1]
    @test coeffbounds(u, 2) <= [5, 12, 1]
    @test coeffbounds(u, 3) <= [5, 21, 13, 1]
    @test coeffbounds(u, 4) <= [5, 26, 36, 14, 1]
end

@testset "hensel_lifting" for (p, q) in [(13, 13)]
    X = varname(x)

    r = gcd(p, q)
    Pq = (ZZ/q)[X]
    Pqr = (ZZ/(q*r))[X]
    v = first.(factor(Pq(u)))
    a = allgcdx(v)
    V, qr = hensel_lift(u, v, a)
    @test Pqr(u) == Pqr(prod(V))
    @test all(Pq.(V) .== v)

    #=
    p, q = p, q * r
    r = gcd(p, q)
    Pq = (ZZ/q)[X]
    Pqr = (ZZ/(q*r))[X]
    v = V
    V, qr = hensel_lift(u, v, a)
    @test Pqr(u) == prod(V)
    @test all(Pq.(V) .== v)
    =#
end

@testset "partsums" begin
    s = [2, 3, 10]
    @test partsums(s) == ( 0xb42d, [1,0,1,1,0,1,0,0], [[0], [], [1], [2], [], [3], [], []])

end

@testset "subsets" begin
    vv = [10, 20, 30, 40]
    @test subset(vv, 0) == Int[]
    @test subset(vv, 6) == [20, 30]
    @test subset(vv, -1) == vv

    @test remove_subset!(copy(vv), -1) |> isempty
    @test remove_subset!(copy(vv), 5) == [20, 40]
    remove_subset!(vv, 8)
    @test vv == [10, 20, 30]
end

end


mutable struct SingleLayer{F, FP, A}
    σ::F
    σ′::FP
    W::A
end

SingleLayer(in::Integer, out::Integer, σ, σ′; init=(dims...)->zeros(Float32, dims...)) =
    SingleLayer(σ, σ′, init(out, in))

(layer::SingleLayer)(x) = layer.σ.(layer.W*x)
deriv(layer::SingleLayer, x) = layer.σ′.(layer.W*x)

# Sparse Updates

Linear(in::Integer, out::Integer) =
    SingleLayer(in, out, identity, (x)->1.0)

# SparseLayer(num_weights, args...; kwargs...) =
#     SingleLayer(num_weights, out, args...; kwargs...)

# sigmoid = Flux.sigmoid
sigmoid′(x) = sigmoid(x)*(1.0-sigmoid(x))
identity′(x) = 1.0
tanh′(x) = 1 - (tanh(x))^2

mutable struct TabularLayer{A}
    W::A
end
Tabular(dims::Integer...) = TabularLayer(zeros(dims...))

@forward TabularLayer.W Base.getindex, Base.setindex!

(layer::TabularLayer)(x) = layer[x]
(layer::TabularLayer)(x::Array{T, 1}) where {T<:Integer} = layer[x...]

mutable struct SparseLayer{F, FP, A, P}
    σ::F
    σ′::FP
    W::A
    ϕ::P
end

SparseLayer(num_weights::Integer, out::Integer, σ=identity, σ′=identity′; init=(dims...)->zeros(Float32, dims...)) =
    SparseLayer(σ, σ′, init(out, num_weights), zeros(Float32, out, num_weights))

(layer::SparseLayer)(x::Array{Int64, 1}) = layer.σ.(sum(layer.W[x]))
deriv(layer::SparseLayer, x::Array{Int64, 1}) = layer.σ′.(sum(layer.W[x]))

(layer::SparseLayer)(x::Array{CartesianIndex{1}, 1}) = layer.σ.(sum(layer.W[x]))
deriv(layer::SparseLayer, x::Array{CartesianIndex{1}, 1}) = layer.σ′.(sum(layer.W[x]))


# using Revise

using DistributedFactorGraphs
using IncrementalInference

using Test

@testset "test CSM runaway on upsolve, (issue 427)" begin

## parameters
lm_prior_noise = 0.1
meas_noise = 0.25
odom_noise = 0.1
n_samples = 100

# initialize mean landmark locations
l1 = 50.0
l2 = -50.0
l3 = 80.0

# "Ground-truth" robot poses
x1 = 0.0
x2 = 0.0
x3 = 0.0

## Initialize empty factor graph
fg = initfg()

addVariable!(fg, Symbol("l1"), ContinuousScalar, N=n_samples)
addFactor!(fg, [:l1], Prior(Normal(l1, lm_prior_noise)))

addVariable!(fg, Symbol("l2"), ContinuousScalar, N=n_samples)
addFactor!(fg, [:l2], Prior(Normal(l2, lm_prior_noise)))

addVariable!(fg, Symbol("l1_0"), ContinuousScalar, N=n_samples)
addVariable!(fg, Symbol("l2_0"), ContinuousScalar, N=n_samples)

# Add first pose
addVariable!(fg, :x1, ContinuousScalar, N=n_samples)

addFactor!(fg, [:x1; :l1; :l1_0], LinearRelative(Normal(40., meas_noise)), multihypo=[1.0; 1.0/2.0; 1.0/2.0])

# Add second pose
addVariable!(fg, :x2, ContinuousScalar, N=n_samples)

# Gaussian transition model
addFactor!(fg, [:x1; :x2], LinearRelative(Normal(0., odom_noise)))

# Make second "door" measurement
# addFactor!(fg, [:x1; :l1], LinearRelative(Normal(0, meas_noise)) )
addFactor!(fg, [:x2; :l2; :l2_0], LinearRelative(Normal(-40., meas_noise)), multihypo=[1.0; 1.0/2.0; 1.0/2.0])

# drawGraph(fg)

# ensureAllInitialized!(fg)

approxConv(fg, :x1l1l1_0f1, :l1_0)
approxConv(fg, :x1l1l1_0f1, :x1)
# doautoinit!(fg, :l1_0)
# doautoinit!(fg, :x1)

## Run solver
getSolverParams(fg).limititers = 30 # previous runaway CSM issue due to excessive limits on autoinit.
# getSolverParams(fg).dbg = false
# getSolverParams(fg).async = false
# getSolverParams(fg).drawtree = false
# getSolverParams(fg).showtree = false

tree, smt, hist = solveTree!(fg, recordcliqs=ls(fg))

# drawGraph(fg)
# fetchAssignTaskHistoryAll!(tree, smt)
# printCliqHistorySummary(tree, :l1)
# getTreeCliqsSolverHistories(fg, tree)
#
# csmAnimate(fg, tree, [:l1;:l2])
#
# # Base.rm("/tmp/caesar/csmCompound/out.ogv")
# run(`ffmpeg -r 10 -i /tmp/caesar/csmCompound/csm_%d.png -c:v libtheora -vf fps=25 -pix_fmt yuv420p -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" -q 10 /tmp/caesar/csmCompound/out.ogv`)
# @async run(`totem /tmp/caesar/csmCompound/out.ogv`)

# using RoMEPlotting
# plotKDE(fg, ls(fg))

end

#


# type PartitionFunc <: ScalarFunc3d{Float64}
#     at::Atom
#     atoms::Vector{Atom}
# end

# function Base.call(f::PartitionFunc, x::Real, y::Real, z::Real)
#     f.at.h(x,y,z)/sum([μ.h(x,y,z) for μ in f.atoms])
# end

# Base.call(f::PartitionFunc, p::Point3{Float64}) = f(p[1],p[2],p[3])

# @vectorize2 Base.call PartitionFunc Point3{Float64} Float64

function delleypartition!(atoms::Atoms)
    for at in atoms
        r = at[:g].r
        at[:h] = RadialFunc(r,at[:rho].u(r)./r.^2,R=at.R,rmax=at[:rho].rmax,npts=length(r))
    end
    for at in atoms
        at[:p] = r -> at[:h](r)./sum([μ[:h](r) for μ in atoms])
    end
    return
end

# function delleypartition!(atoms::Atoms)
#     for at in atoms
#         r,rho = at[:rho].u.x,at[:rho].u.y
#         at[:h] = RadialFunc(r,rho./r.^2,R=at[:R],rmax=at[:rho].rmax,npts=length(r))
#     end
#     for at in atoms
#         at[:p] = r -> at[:h](r)./sum([μ[:h](r) for μ in atoms])
#     end
#     return
# end

function delleypartition(atoms::Atoms)
    let atoms = copy(atoms)
        delleypartition!(atoms)
        return atoms
    end
end


using Test

using HydrophoneCalibrations
using HydrophoneCalibrations: enforce_hertz_unit, interp
using HydrophoneCalibrations.Unitful

tolerance_v2p = 1e4u"Pa/V";

@testset "HydrophoneCalibrations.jl" begin
    @testset "precision acoustics hydrophone" begin
        v2p = volt_to_pressure(1e6, :PA_3197)
        sensitivity_at_1mhz = 34u"mV/MPa"
        conversion_factor = uconvert(u"Pa/V", 1/sensitivity_at_1mhz)
        @test isapprox(v2p, conversion_factor, atol=tolerance_v2p)
    end

    @testset "onda hydrophone" begin
        v2p = volt_to_pressure(1e6, :Onda_HGL0200_2322)
        sensitivity_at_1mhz = 4.11e-8u"V/Pa"
        conversion_factor = uconvert(u"Pa/V", 1/sensitivity_at_1mhz)
        @test isapprox(v2p, conversion_factor, atol=tolerance_v2p)
    end

    @testset "onda hydrophone with preamplifier (linear calculations)" begin
        v2p = volt_to_pressure(1e6, :Onda_HGL0200_2322, :Onda_AH2020_1238_20dB)
        sensitivity_at_1mhz = 4.11e-8u"V/Pa"
        hydrophone_capacitance_at_1mhz = 9.963u"pF"
        preamp_capacitance_at_1mhz = 7.08u"pF"
        preamp_gain_at_1mhz = 19.92u"dB"
        capacitance_factor = hydrophone_capacitance_at_1mhz/(hydrophone_capacitance_at_1mhz+preamp_capacitance_at_1mhz)
        system_sensitivity = uconvertrp(NoUnits, preamp_gain_at_1mhz) * sensitivity_at_1mhz * capacitance_factor
        conversion_factor = uconvert(u"Pa/V", 1/system_sensitivity)
        @test isapprox(v2p, conversion_factor, atol=tolerance_v2p)
    end

    @testset "onda hydrophone with preamplifier (logarithmic calculations)" begin
        v2p = volt_to_pressure(1e6, :Onda_HGL0200_2322, :Onda_AH2020_1238_20dB)
        sensitivity_at_1mhz = 4.11e-8u"V/Pa"
        hydrophone_capacitance_at_1mhz = 9.963u"pF"
        preamp_capacitance_at_1mhz = 7.08u"pF"
        preamp_gain_at_1mhz = 19.92u"dB"
        capacitance_factor = hydrophone_capacitance_at_1mhz/(hydrophone_capacitance_at_1mhz+preamp_capacitance_at_1mhz)
        system_sensitivity = preamp_gain_at_1mhz + ustrip(uconvertrp(u"dBV/Pa", sensitivity_at_1mhz))*u"dBV" + uconvertrp(u"dB", capacitance_factor)
        system_sensitivity = uconvertrp(u"V", system_sensitivity) / u"Pa"
        conversion_factor = uconvert(u"Pa/V", 1/system_sensitivity)
        @test isapprox(v2p, conversion_factor, atol=tolerance_v2p)
    end

    @testset "onda hydrophone with unknown preamplifier" begin
        @test_throws MethodError volt_to_pressure(1e6, :Onda_HGL0200_2322, :my_dummy)
    end

    @testset "volt_to_pressure_and_phase() with unknown preamp" begin
        @test_throws MethodError volt_to_pressure_and_phase(1e6, :Onda_HGL0200_2322, :my_dummy)
    end

    @testset "enforce_hertz_unit()" begin
        @test enforce_hertz_unit(2.0u"mHz") == 0.002u"Hz"
        @test enforce_hertz_unit(2u"Hz")    == 2u"Hz"
        @test enforce_hertz_unit(2.2u"MHz") == 2.2e6u"Hz"
        @test_throws ErrorException enforce_hertz_unit(2.2u"m")
        @test enforce_hertz_unit(2)     == 2u"Hz"
        @test enforce_hertz_unit(2.2e6) == 2.2e6u"Hz"
    end

    @testset "interp()" begin
        f_axis = collect(0:1:10)
        values = f_axis

        f0_0 = 0.01
        @test interp(f_axis, values, f0_0) == 0.01
        f0_1 = 0.5
        @test interp(f_axis, values, f0_1) == 0.5
        f0_low = -0.1
        @test_throws "f0 must be between $f_axis[1]-$f_axis[end]" interp(f_axis, values, f0_low)
        f0_high = 10.1
        @test_throws "f0 must be between $f_axis[1]-$f_axis[end]" interp(f_axis, values, f0_high)
    end
end


@testset "dirclustercoeffs" begin
    (A,xy,labels) = load_matrix_network_metadata("celegans")
    (cc, cccyc, ccmid, ccin, ccout, nf) = dirclustercoeffs(A, true, true)
    @test length(cc) == 202
    (maxval, maxind) = findmax(cc)
    @test maxind == 113
    (cc, cccyc, ccmid, ccin, ccout, nf) = dirclustercoeffs(A)
    @test length(cc) == 202
    (maxval, maxind) = findmax(cc)
    @test maxind == 113
    (cc, cccyc, ccmid, ccin, ccout, nf) = dirclustercoeffs(A, true)
    @test length(cc) == 202
    (maxval, maxind) = findmax(cc)
    @test maxind == 113
    # negative weights
    @test_throws Exception dirclustercoeffs(-sparse(1.0I,5,5))
end