using Roots
using Test
using JSON
using Printf

## testing/benchmarking derivative-free methods

mutable struct Func1
    name :: Symbol
    val :: Function
    x0 :: Function
    alpha # a value or values
    params :: Vector{Any}
end
function Base.show(io::IO, f::Func1)
    @printf io "Func(%s)" f.name
end

#Construct a function object,
macro Func1(name)
    @gensym f p x
    esc(quote
        $f = Func1($name, val, x0, alpha, params)
        push!(known_functions, $f)
        $f
    end)
end

known_functions = Func1[]

## This set of tests is useful for benchmarking the number of function
## calls, failures, and max errors for the various derivative_free methods.

# easy ones
func1 = let
    val = (_, x) -> cos(x) - x/2
    a,b, N = 0.0, pi/2, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 1.0298665293222589
    @Func1 :func1
end

func2 = let
    val = (_, x) -> exp(x) - x^4
    a, b, N = 5.0, 20.0, 11
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = find_zeros(x->val(0,x), -5.0, 20.0)
    @Func1 :func2
end

func3 = let
    val = (_, x) -> x^5 - x - 1
    a, b, N = 0.5, 2.0, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = .1673039782614187
    @Func1 :func3
end

# wider range
func4 = let
    val = (_,x) -> (1 + cos(x)) * (exp(x) - 2)
    a, b, N = 0.0, 1.7, 15
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.6931471805599453
    @Func1 :func4
end

func5 = let
    val = (_,x) -> 2*x - exp(-x)
    a, b, N = -5.0, 5.0, 15
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.35173371124919584
    @Func1 :func5
end

# once over hump, hard to get back to 0
func6 = let
    val = (_,x) -> x * exp(-x)
    a, b, N = -5.0, 5.0, 15
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.0
    @Func1 :func6
end

# once outside of extrema, points away from zero
func7 = let
    val = (_,x) ->  20.0 * x / ( 100.0 * x^2 + 1.0)
    a, b, N = -0.2, 0.2, 15
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.0
    @Func1 :func7
end

# should have some failures here

func8 = let
 val = (_,x) -> cos(x) - x/2
    a, b, N = 0.0, 3pi/4, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 1.0298665293222589
    @Func1 :func8
end

func9 = let
    val = (_,x) -> tanh(x) - tan(x)
    a, b, N = 5.0, 10.0, 8
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = find_zeros(x -> val(0,x), 5, 12)
    @Func1 :func9
end

func10 = let
    val = (_,x) -> asin(x^2 - 1) - x/2 + 1
    a, b, N = -0.5, 0.9, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = find_zeros(x -> val(0,x), -0.5, 1.0)
    @Func1 :func10
end

func11 = let
    val = (_,x) -> exp(-x) - cos(x)
    a, b, N = -2.0, 2.0, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = find_zeros(x -> exp(-x) - cos(x), -2,2)
    @Func1 :func11
end

func12 = let
    val = (_,x) -> sqrt(x) - 1/x - 3
    a, b, N = 5, 20, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 9.633595562832696
    @Func1 :func12
end

## multiplicity: http://ir.igsnrr.ac.cn/bitstream/311030/8840/1/%E4%BE%AF%E9%BA%9F%E7%A7%91(SCI)2.pdf
func13 = let
 val = (_,x) -> (x - sqrt(5))^4 / ((x-1)^2 + 2)
    a, b, N = 3.0, 4.5, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha =  2.236067977499790
    @Func1 :func13
end

func14 = let
 val = (_,x) -> (sin(x)^2 - 2x + 1)^5
    a, b, N = 3.0, 5.0, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.71483582544138924
    @Func1 :func14
end

func15 = let
 val = (_,x) -> (8x*exp(-x^2) -2x - 3)^8
    a, b, N = -2.0,-1.5, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = -1.7903531791589544
    @Func1 :func15
end

# this is 1e6 * (x-1)^7, has many "zeros" near 1
func16 = let
    val = (_,x) -> 1e6 * (x^7 -7x^6 +21x^5 -35x^4 +35x^3-21x^2+7x-1)
    a, b, N = -0.0, 2.0, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 1.0
    @Func1 :func16
end

func17 = let
 val = (_,x) -> (exp(-x^2 + x + 3) - x + 2)^9
    a, b, N = 2.4, 2.5, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha =  2.4905398276083051
    @Func1 :func17
end

func18 = let
 val = (_,x) -> (exp(-x) + 2sin(x))^4
    a, b, N = 3.0, 3.4, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 3.1627488709263654
    @Func1 :func18
end

## hard w/o aggressive bracketing
func19 = let
 val = (_,x) -> cbrt(x)
    a, b, N = -3.0, 3.0, 10
    x0(n) = range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.0
    @Func1 :func19
end

func20 = let
    val = (_,x) -> x < -1/4 ? 3/4*x-0.3125 : x < 1/4 ? 2x : 3/4*x+0.3125
    a,b,N = -1.0, 1.0, 10
    x0(n) =  range(a, stop=b, length=N)[n]
    params = 1:N
    alpha = 0.0
    @Func1 :func20
end

mutable struct MethodResults1
    name
    problems::Int
    evalcount :: Int
maxresidual :: Float64
maxdifference :: Float64
    failures :: Vector{Tuple{Func1, Int}}
end
MethodResults1() = MethodResults1(nothing, 0, 0, 0.0, Inf,  Tuple{Func1, Int}[])
Base.show(io::IO, results::MethodResults1) =
    print(io, "MethodResults($(results.name), evalcount=$(results.evalcount), numfailures=$(length(results.failures)), maxresidual=$(results.maxresidual))")

## Run a method on all known functions.
mindiff(a, alpha) =  minimum([a - i for i in alpha])
function run_tests(method; verbose=false, trace=false, name=nothing, abandon=false)
    results = MethodResults1()
    results.name = name
    results.problems = 0
    for f in known_functions
        for i in 1:length(f.params)
            p = f.params[i]
            evalcount = 0
            function feval(x)
                evalcount += 1
                result = f.val(p, x)
                trace && @printf "%s[%d]: %s ⇒ %s\n" f i x result
                result
            end
            result, residual = nothing, nothing
            try
                result = method(feval, f.x0(p))
                md = abs(mindiff(result, f.alpha) )
                isnan(result) && error("NaN")
                (md > 1) && error("ran away")

                results.problems += 1
                results.evalcount += evalcount
                residual = f.val(p, result)

                verbose && @printf "%s[%d] ⇒ %d / %s, residual %.5e\n" f.name i evalcount result residual

                ## Some functions might return non-real values on failures
                if isa(result, AbstractFloat) && isa(residual, AbstractFloat) && isfinite(residual)
                    results.maxresidual = max(results.maxresidual, abs(residual))
                    results.maxdifference = max(results.maxdifference, mindiff(result, f.alpha))
                end

                verbose && abs(residual) >1/10 && @printf "Large Residual [%s] %d/%d\n" f.name i p

            catch ex
                verbose && @printf "%s[%d] ⇒ FAILED: %s\n" f.name i ex
                push!(results.failures, (f, i))
                abandon && rethrow(ex)
            end
        end
    end
    results
end

avg(x) = sum(x)/length(x)
D(f, h=1e-4) = x -> (f(x+h) - f(x-h))/(2h)
D2(f, h=1e-4) = x -> (f(x+h) - 2f(x) + f(x-h))/h^2

if !isinteractive()

    @testset "derivative free methods" begin

        ## Test for failures, ideally all of these would be 0
        ## test for residual, ideally small
        ## test for evaluation counts, ideally not so low for these problems

        ## basic methods
        Ms = [Roots.Order1(), Roots.Order1B(), Roots.Order2(), Roots.Order2B(),
              Roots.Order5(), Roots.Order8(), Roots.Order16()]
        results = [run_tests((f,b) -> find_zero(f, b, M), name="$M") for M in Ms]

        failures = [length(result.failures) for result in results]
        residuals = [result.maxresidual for result in results]
        cnts = [result.evalcount/result.problems for result in results]

        @test maximum(failures) <= 60
        @test maximum(residuals) <= 5e-5
        @test avg(cnts) <= 40

        ## methods which fall back to bisection when bracket found
        Ms = [Roots.Order0()]
        results = [run_tests((f,b) -> find_zero(f, b, M), name="$M") for M in Ms]

        failures = [length(result.failures) for result in results]
        residuals = [result.maxresidual for result in results]
        cnts = [result.evalcount/result.problems  for result in results]

        @test maximum(failures) <= 40
        @test maximum(residuals) <= 1e-5
        @test avg(cnts) <= 30

        ## Newton and Halley
        Fs = [(f,b) -> find_zero((f, D(f)), b, Roots.Newton()),
              (f,b) -> find_zero((f, D(f), D2(f)), b, Roots.Halley()),
              (f, b) -> find_zero((f, D(f), D2(f)), b, Roots.Schroder())
              ]
        results = [run_tests(F) for F in Fs]

        failures = [length(result.failures) for result in results]
        residuals = [result.maxresidual for result in results]
        cnts = [result.evalcount/result.problems  for result in results]

        @test maximum(failures) <= 70
        @test maximum(residuals) <= 5e-4
        @test avg(cnts) <= 125

    end

    @testset "derivative free, non Float64" begin

        Ms = [Roots.Order0(), Roots.Order1(), Roots.Order1B(), Roots.Order2(), Roots.Order2B(),
              Roots.Order5(), Roots.Order8(), Roots.Order16()]
        Ts = [Float16, Float32, BigFloat]

        fn = x -> x^3 - 2x - 5
        alpha = 2.094551481542326591482386540579302963857306105628239180304128529045312189983499
        x0 = 2.0

        for T in Ts
            for M in Ms
                xstar = find_zero(fn, T(x0), M)
                @test xstar ≈ T(alpha) atol=max(sqrt(eps(T)), eps())
                @test isa(xstar, T)
            end
        end

        for T in Ts
            xstar = Roots.find_zero((fn,D(fn,sqrt(eps(T)))), T(x0), Roots.Newton())
            @test xstar ≈ T(alpha)  atol=max(sqrt(eps(T)), eps())
            @test isa(xstar, T)

            xstar = Roots.find_zero((fn,D(fn,sqrt(eps(T))),D2(fn,sqrt(eps(T)))), T(x0), Roots.Halley())
            @test xstar ≈ T(alpha) atol=max(sqrt(eps(T)), eps())
            @test isa(xstar, T)

            xstar = Roots.find_zero((fn,D(fn,sqrt(eps(T))),D2(fn,sqrt(eps(T)))), T(x0), Roots.Schroder())
            @test xstar ≈ T(alpha) atol=max(sqrt(eps(T)), eps())
            @test isa(xstar, T)
        end
    end
end


# ------------------------------------------ #
#                   Importing                #
# ------------------------------------------ #
using KissMCMC, LaTeXStrings

# ------------------------------------------ #
#           Loop over the two days           #
# ------------------------------------------ #

model_runs = Dict()
if !@isdefined n_iterations
    n_iterations = 5*10^4 # at least 10^6 for nice distribution in plot
end

if !@isdefined use_progress_meter
    use_progress_meter = true # shows nice progress bar; but hide it for Literate.jl
end

for date in ["0908", "2108"]

    out = Dict()

    # ------------------------------------------ #
    #            free-gradient model, MCMC       #
    # ------------------------------------------ #

    # initial point of walker
    if date == "0908"
        gradT0 = -5e-4;
    elseif date == "2108"
        gradT0 = -2e-3;
    end
    Sinit_0 = pmean(mid_309_265[date][:S][1]);
    Sinit_std = pstd(mid_309_265[date][:S][1]);
    theta0 = [gradT0, Sinit_0];

    # prior information
    dTmin, dTmax = -1e-2, -1e-6;
    # logarithmic posterior density function
    logposterior = R.make_logpdf(date, mid_309_265, [dTmin, dTmax]);

    # emcee MCMC sampler:
    thetas, accept_ratioe, logdensities, blobs = emcee(logposterior, make_theta0s(theta0, [1e-3, Sinit_std],
                                                                                logposterior, 6, hasblob=true),
                                                       niter=n_iterations, hasblob=true, use_progress_meter=use_progress_meter)
    thetas, accept_ratioe, logdensities, blobs = squash_walkers(thetas, accept_ratioe, logdensities, blobs); # puts all walkers into one

    out[:dTdz_MCMC] = [thetas[k][1] for k in 1:length(thetas)];
    # out[:Sinit] = [thetas[k][2] for k in 1:length(thetas)];

    out[:S_MCMC] = [blobs[k][1] for k in 1:length(blobs)];
    out[:dSdt_MCMC] = [blobs[k][2] for k in 1:length(blobs)];
    out[:dSdt_sensible_MCMC] = [blobs[k][3] for k in 1:length(blobs)];

    # ------------------------------------------ #
    #                   ct-model                 #
    # ------------------------------------------ #

    dp = mean(mid_309_265[date][:dp], dims=1) # mean pressure difference over the tracer experiments
    out[:dTdz_ct] = dp * ct / mid_309_265[date][:dz]
    out[:S_ct],
        out[:dSdt_ct],
        out[:dSdt_sensible_ct] = R.model_S(mid_309_265, date, out[:dTdz_ct], Sinit_0)

    # ------------------------------------------ #
    #                closure rate                #
    # ------------------------------------------ #

    out[:closure] = R.closure_rate(ctd309, ctd265, mid_309_265, date)

    # ------------------------------------------ #
    #            dT/dz from measurements         #
    # ------------------------------------------ #

    i = indices[date][1][1]:indices[date][4][end] # indices from appearance of first tracer signal until disappearance of last one
    out[:dTdz_measured] = (Particles.(n_partcl, Normal.(
                            mean(ctd265[date][:temp_PT][i]), e_T))
                        - Particles.(n_partcl, Normal.(
                            mean(ctd309[date][:temp_PT][i]), e_T))
                        ) / mid_309_265[date][:dz]

    # ------------------------------------------ #
    #              Save in dictionary            #
    # ------------------------------------------ #

    model_runs[date] = out

end

## Table S2
## print latex table of total opening rate, opening rate due to sensible heat and closure rate
table_s3 = ""
for (MCMC, c_t, description) in zip([:dSdt_MCMC, :dSdt_sensible_MCMC],
                                 [:dSdt_ct, :dSdt_sensible_ct],
                                 [L"Total opening rate ($\mathrm{m^2\,s^{-1}}$)", L"Opening rate due to sensible heat ($\mathrm{m^2\,s^{-1}}$)"])
    for date in ["0908", "2108"]
        if date == "0908"
            o = description
        else
            o = ""
        end
        values_ct = only(mean(model_runs[date][c_t], dims=1))
        values_MCMC = only(mean(
                    [ bootstrap(
                        Particles([model_runs[date][MCMC][i][k] for i in 1:length(model_runs[date][MCMC])]),
                        n_partcl)
                      for k in 1:length(model_runs[date][MCMC][1]) ]
                    , dims=1))
        if date == "0908"
            o = o * " & " * string(pmean(values_ct) ± pstd(values_ct)) * " & " * string(pmean(values_MCMC) ± pstd(values_MCMC)) * "\\\\ \n"
        else
            o = o * " & \\textbf{" * string(pmean(values_ct) ± pstd(values_ct)) * "} & \\textbf{" * string(pmean(values_MCMC) ± pstd(values_MCMC)) * "} \\\\ \n"
        end
        global table_s3 *= o
    end
end
table_s3 *= "\\hline \n"
closure_rate = only(mean(model_runs["0908"][:closure], dims=1))
o = L"Closure rate ($\mathrm{m^2\,s^{-1}}$)" * " & \\multicolumn{2}{" * string(pmean(closure_rate) ± pstd(closure_rate)) * "} \\\\ \n"
table_s3 *= o
closure_rate = only(mean(model_runs["2108"][:closure], dims=1))
o = " & \\multicolumn{2}{\\textbf{" * string(pmean(closure_rate) ± pstd(closure_rate)) * "}} \\\\ \n"
table_s3 *= o


###############################################################################
#
#  Efficiently computing the finite sum part of the Riemann theta function.
#
#  Original Authors
#  -------
#  * Chris Swierczewski (@cswiercz) - September 2012, July 2016
#  * Grady Williams (@gradyrw) - October 2012
#  * Jeremy Upsal (@jupsal) - July 2016
#
###############################################################################

###############################################################################
#  exppart
#  -------
#  A helper function for the finite sum functions. Computes
#             2π < (n-intshift), (1/2)X(n-intshift) + x >
###############################################################################
@inline function exppart(n::Vector{Float64}, X::Matrix{Float64},
                 x::Vector{Float64}, intshift::Vector{Float64})::Float64
    tmp1 = n - intshift
    tmp2 = 0.5 * X * tmp1 + x
    2π * dot(tmp1, tmp2)
end

###############################################################################
# normpart
# --------
# A helper function for the finite sum functions. Computes
#            -pi * || T*(n+fracshift) ||^2
###############################################################################
@inline function normpart(n::Vector{Float64}, T::AbstractMatrix{Float64},
                  fracshift::Vector{Float64})::Float64
    tmp1 = n + fracshift
    tmp2 = T * tmp1
    -π * dot(tmp2, tmp2)
end

"""
      deriv_prod(n::Vector{Float64},
                 intshift::Vector{Float64},
                 derivs::Vector{Vector{ComplexF64}})::ComplexF64

Compute the real and imaginary parts of the product
         ___
         | |    2π * I <d, n-intshift>    | |
           d in derivs
for a given n in ZZ^g.

Parameters
----------
- n : An integer vector in the finite sum ellipsoid.
- intshift : The integer part of Yinv*y.
- deriv_real, deriv_imag : The real and imaginary parts of the derivative directional vectors.

Returns
-------
- The complex "derivative product".
"""
function deriv_prod(n::Vector{Float64},
                    intshift::Vector{Float64},
                    derivs::Vector{Vector{ComplexF64}})::ComplexF64
    # compute n-intshift
    nmintshift = n - intshift

    #   Computes the dot product of each directional derivative and nmintshift.
    #   Then it computes the product of the resulting complex scalars.
    total = Complex(1., 0.)
    for der in derivs
        total *= sum(der .* nmintshift)
    end

    # Compute (2*pi*i)^(nderivs) * (total_real + total_imag*i)
    pi_mult = (2π) ^ length(derivs)

    # Determines what the result of i^nderivs is, and performs the correct
    #  multiplication afterwards.
    remain = length(derivs) % 4
    (remain == 0) && return pi_mult * total
    (remain == 1) && return pi_mult * total * Complex(0., 1.)
    (remain == 2) && return pi_mult * (-total)
    (remain == 3) && return pi_mult * total * Complex(0., -1.)
end

"""
    function finite_sum(X::Matrix{Float64},
                        Yinv::Matrix{Float64},
                        T::Matrix{Float64},
                        z::Vector{Vector{ComplexF64}},
                        S::Vector{Vector{Float64}},
                        derivs::Vector{Vector{ComplexF64}})

Computes the real and imaginary parts of the finite sum with derivatives.

Parameters
----------
X, Yinv, T : Row-major matrices such that the Riemann matrix, Omega is equal to (X +
  iY). T is the Cholesky decomposition of Y.
z : Input vectors.
S : The set of points in ZZ^g over which to compute the finite sum
derivs : the derivative directional vectors.

Returns
-------
  The finite sums for each z.
"""
function finite_sum(X::Matrix{Float64},
                    Yinv::Matrix{Float64},
                    T::AbstractMatrix{Float64},
                    z::Vector{Vector{ComplexF64}},
                    S::Vector{Vector{Float64}},
                    derivs::Vector{Vector{ComplexF64}})

    num_vectors, num_points, g = length(z), length(S), size(T,1)

    if num_vectors * num_points * g > 4000 # TODO : improve this heuristic ?
        finite_sum_large(X, Yinv, T, z, S, derivs)
    else
        finite_sum_small(X, Yinv, T, z, S, derivs)
    end
end

function finite_sum_small(X::Matrix{Float64},
                          Yinv::Matrix{Float64},
                          T::AbstractMatrix{Float64},
                          z::Vector{Vector{ComplexF64}},
                          S::Vector{Vector{Float64}},
                          derivs::Vector{Vector{ComplexF64}})

    num_vectors = length(z)
    values = zeros(ComplexF64, num_vectors)

    for kk in 1:num_vectors
        #   compute the shifted vectors: shift = Yinv*y as well as its integer and
        #   fractional parts
        shift = Yinv * imag(z[kk])
        intshift = round.(shift)
        fracshift = shift - intshift
        rz = real(z[kk])

        # compute the finite sum
        for s in S
            tmp1 = s - intshift
            tmp2 = 0.5 * X * tmp1 + rz
            ep = 2π * dot(tmp1, tmp2) # π omitted since sinpi / cospi is used

            tmp1 = s + fracshift
            tmp2 = T * tmp1
            np = -π * dot(tmp2, tmp2)

            pt = exp(np) * cis(ep)

            if length(derivs) == 0
                values[kk] += pt
            else
                values[kk] += deriv_prod(s, intshift, derivs) * pt
            end
        end
    end
    values
end

# use matrix operations as much as possible
function finite_sum_large(X::Matrix{Float64},
                          Yinv::Matrix{Float64},
                          T::AbstractMatrix{Float64},
                          z::Vector{Vector{ComplexF64}},
                          S::Vector{Vector{Float64}},
                          derivs::Vector{Vector{ComplexF64}})

    num_vectors, num_points, g = length(z), length(S), size(T,1)

    Z = reshape(vcat(z...), g, num_vectors)
    Smat = reshape(vcat(S...), g, num_points)

    RealZ = real.(Z)
    tmp = Yinv * imag.(Z)
    IS = round.(tmp)
    FS = tmp .- IS
    TF = T * FS
    TF2 = sum(TF .* TF, dims=1)

    ### exponential part
    XIS = X * IS
    XS  = X * Smat
    EP = Smat' * RealZ .- sum(RealZ .* IS, dims=1) .+
         0.5 * ( sum(Smat .* XS, dims=1)'  .+
                 sum(IS .* XIS, dims=1) .-
                 Smat' * XIS .- XS' * IS)
    EP = 2π .* EP

    ### normal part
    TS  = T * Smat
    TS2 = sum(TS .* TS, dims=1)'
    TSTF = TS' * TF
    NP = exp.( -π .* ( TS2 .+ TF2 .+ 2. * TSTF ) )

    if length(derivs) == 0
        values = vec(sum( NP .* cis.(EP), dims=1))
    else
        values = zeros(ComplexF64, num_vectors)
        PT = NP .* cis.(EP)
        for iv in 1:num_vectors
            for ip in 1:num_points
                dp = deriv_prod(S[ip], IS[:,iv], derivs)
                values[iv] += PT[ip,iv] * dp
            end
        end
    end
    values
end


module SignalDecomposition

export decompose, Decomposition

"Supertype of all decomposition methods."
abstract type Decomposition end

"""
    decompose([t, ] s, method::Decomposition) → x, r
Decompose an 1D input signal or timeseries `s(t)` into components, `x, r`,
using the given `method`. `t` defaults to `1:length(s)`.

What are `x` and `r` really depends on your point of view and your application.
They can be structure `x` and noise `r` (i.e. noise reduction). They can be
seasonal/periodic `x` and residual component `r`. They can even be multiplier `x` and
input `r`.
"""
decompose(s::AbstractVector, method::Decomposition; kwargs...) =
decompose(1:length(s), s, method; kwargs...)

using Statistics
include("utils.jl")
include("linear/fourier.jl")
include("linear/lpv.jl")
include("product/matrixinversion.jl")
include("nonlinear/extremelysimple.jl")
include("nonlinear/projection.jl")
include("misc/anomaly.jl")

end # module


using .TurbulenceClosures: ∂z_aaf, ▶z_aaf

const g_Earth = 9.80665

#=
Supported buoyancy types:

- Nothing
- SeawaterBuoyancy
=#

#####
##### Functions for buoyancy = nothing
#####

@inline buoyancy_perturbation(i, j, k, grid::AbstractGrid{T}, ::Nothing, C) where T = zero(T)
@inline buoyancy_frequency_squared(i, j, k, grid::AbstractGrid{T}, ::Nothing, C) where T = zero(T)

#####
##### Seawater buoyancy for buoyancy determined by temperature and salinity
#####

"""
    BuoyancyTracer <: AbstractBuoyancy{Nothing}

Type indicating that the tracer `T` represents buoyancy.
"""
struct BuoyancyTracer <: AbstractBuoyancy{Nothing} end

@inline buoyancy_perturbation(i, j, k, grid, ::BuoyancyTracer, C) = @inbounds C.T[i, j, k]
@inline buoyancy_frequency_squared(i, j, k, grid, ::BuoyancyTracer, C) = ∂z_aaf(i, j, k, grid, C.T)

"""
    SeawaterBuoyancy{G, EOS} <: AbstractBuoyancy{EOS}

Buoyancy model for temperature- and salt-stratified seawater.
"""
struct SeawaterBuoyancy{G, EOS} <: AbstractBuoyancy{EOS}
    gravitational_acceleration :: G
    equation_of_state :: EOS
end

"""
    SeawaterBuoyancy([T=Float64;] gravitational_acceleration = g_Earth,
                                  equation_of_state = LinearEquationOfState(T))

Returns parameters for a temperature- and salt-stratified seawater buoyancy model
with a `gravitational_acceleration` constant (typically called 'g'), and an
`equation_of_state` that related temperature and salinity (or conservative temperature
and absolute salinity) to density anomalies and buoyancy.
"""
function SeawaterBuoyancy(T=Float64;
                          gravitational_acceleration = g_Earth,
                          equation_of_state = LinearEquationOfState(T))
    return SeawaterBuoyancy{T, typeof(equation_of_state)}(gravitational_acceleration, equation_of_state)
end

"""
    buoyancy_frequency_squared(i, j, k, grid, b::SeawaterBuoyancy, C)

Returns the buoyancy frequency squared for temperature and salt-stratified water,

```math
N^2 = g \\left ( \\alpha \\partial_z T - \\beta \\partial_z S \\right ) \\, ,
```

where ``\$ g \$`` is gravitational acceleration, ``\$ \\alpha \$`` is the thermal expansion
coefficient, ``\$ \\beta \$`` is the haline contraction coefficient, ``\$ T \$`` is
temperature or conservative temperature, where applicable, and ``\$ S \$`` is the
salinity or absolute salinity, where applicable.
"""
@inline buoyancy_frequency_squared(i, j, k, grid, b::SeawaterBuoyancy, C) =
    b.gravitational_acceleration * (
           thermal_expansion(i, j, k, grid, b.equation_of_state, C) * ∂z_aaf(i, j, k, grid, C.T)
        - haline_contraction(i, j, k, grid, b.equation_of_state, C) * ∂z_aaf(i, j, k, grid, C.S) )

#####
##### Linear equation of state
#####

"""
    LinearEquationOfState{T} <: AbstractEquationOfState

Linear equation of state for seawater.
"""
struct LinearEquationOfState{T} <: AbstractEquationOfState
    α :: T
    β :: T
end

"""
    LinearEquationOfState([T=Float64;] α=1.67e-4, β=7.80e-4)

Returns parameters for a linear equation of state for seawater with
thermal expansion coefficient `α` [K⁻¹] and haline contraction coefficient
`β` [ppt⁻¹]. The buoyancy perturbation associated with a linear equation of state is

    `b = g * α * T - g * β * S`

Default constants are taken from Table 1.2 (page 33) of Vallis, "Atmospheric and Oceanic Fluid
Dynamics: Fundamentals and Large-Scale Circulation" (2ed, 2017).
"""
LinearEquationOfState(T=Float64; α=1.67e-4, β=7.80e-4) =
    LinearEquationOfState{T}(α, β)

const LinearSeawaterBuoyancy = SeawaterBuoyancy{FT, <:LinearEquationOfState} where FT

@inline buoyancy_perturbation(i, j, k, grid, b::LinearSeawaterBuoyancy, C) =
    return @inbounds b.gravitational_acceleration * (   b.equation_of_state.α * C.T[i, j, k]
                                                      - b.equation_of_state.β * C.S[i, j, k] )

@inline  thermal_expansion(i, j, k, grid, eos::LinearEquationOfState, C) = eos.α
@inline haline_contraction(i, j, k, grid, eos::LinearEquationOfState, C) = eos.β

#####
##### Nonlinear equations of state
#####

@inline buoyancy_perturbation(i, j, k, grid, b::AbstractBuoyancy{<:AbstractNonlinearEquationOfState}, C) =
    - b.gravitational_acceleration * ρ′(i, j, k, grid, b.equation_of_state, C) / b.ρ₀

#####
##### Roquet et al 2015 idealized nonlinear equations of state
#####

# Reference: Table 3 in Roquet et el, "Defining a Simplified yet 'Realistic' Equation of State for Seawater", (JPO, 2015)
optimized_roquet_coeffs = Dict(
                  :linear => (R₀₁₀ = -1.775e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = 0,         R₀₁₁ = 0,          R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
               :cabbeling => (R₀₁₀ = -0.844e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = -4.561e-3, R₀₁₁ = 0,          R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
:cabbeling_thermobaricity => (R₀₁₀ = -0.651e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = -5.027e-3, R₀₁₁ = -2.5681e-5, R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
                :freezing => (R₀₁₀ = -0.491e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = -5.027e-3, R₀₁₁ = -2.5681e-5, R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
            :second_order => (R₀₁₀ =  0.182e-1, R₁₀₀ = 8.078e-1, R₀₂₀ = -4.937e-3, R₀₁₁ = -2.4677e-5, R₂₀₀ = -1.115e-4, R₁₀₁ = -8.241e-6, R₁₁₀ = -2.446e-3)
)

"""
    RoquetIdealizedNonlinearEquationOfState{F, C, T} <: AbstractNonlinearEquationOfState

Parameters associated with the idealized nonlinear equation of state proposed by
Roquet et al., "Defining a Simplified yet 'Realistic' Equation of State for Seawater",
Journal of Physical Oceanography (2015).
"""
struct RoquetIdealizedNonlinearEquationOfState{F, C, T} <: AbstractNonlinearEquationOfState
                   ρ₀ :: T
    polynomial_coeffs :: C
end

type_convert_roquet_coeffs(T, coeffs) = NamedTuple{propertynames(coeffs)}(Tuple(T(R) for R in coeffs))

"""
    RoquetIdealizedNonlinearEquationOfState([T=Float64,] flavor, ρ₀=1024.6,
                                            polynomial_coeffs=optimized_roquet_coeffs[flavor])

Returns parameters for the idealized polynomial nonlinear equation of state with
reference density `ρ₀` and `polynomial_coeffs` proposed by Roquet et al., "Defining a
Simplified yet 'Realistic' Equation of State for Seawater", Journal of Physical
Oceanography (2015). The default reference density is `ρ₀ = 1024.6 kg m⁻³`, the average
surface density of seawater in the world ocean.

The `flavor` of the nonlinear equation of state is a symbol that selects a set of optimized
polynomial coefficients defined in Table 2 of Roquet et al., "Defining a Simplified yet
'Realistic' Equation of State for Seawater", Journal of Physical Oceanography (2015), and
further discussed in the text surrounding equations (12)--(15). The optimization minimizes
errors in estimated horizontal density gradient estiamted from climatological temperature
and salinity distributions between the 5 simplified forms chosen by Roquet et. al
and the full-fledged [TEOS-10](http://www.teos-10.org) equation of state.

The equations of state define the density anomaly `ρ′`, and have
the polynomial form

    `ρ′(T, S, D) = Σᵢⱼₐ Rᵢⱼₐ Tⁱ Sʲ Dᵃ`,

where `T` is conservative temperature, `S` is absolute salinity, and `D` is the
geopotential depth, currently just `D = -z`. The `Rᵢⱼₐ` are constant coefficients.

Flavors of idealized nonlinear equations of state
=================================================

    - `:linear`: a linear equation of state, `ρ′ = R₁₀₀ * T + R₀₁₀ * S`

    - `:cabbeling`: includes quadratic temperature term,
                    `ρ′ = R₁₀₀ * T + R₀₁₀ * S + R₀₂₀ * T^2`

    - `:cabbeling_thermobaricity`: includes 'thermobaricity' term,
                                   `ρ′ = R₁₀₀ * T + R₀₁₀ * S + R₀₂₀ * T^2 + R₀₁₁ * T * D`

    - `:freezing`: same as `:cabbeling_thermobaricity` with modified constants to increase
                   accuracy near freezing

    - `:second_order`: includes quadratic salinity, halibaricity, and thermohaline term,
                       `ρ′ = R₁₀₀ * T + R₀₁₀ * S + R₀₂₀ * T^2 + R₀₁₁ * T * D`
                             + R₂₀₀ * S^2 + R₁₀₁ * S * D + R₁₁₀ * S * T`

Example
=======

julia> using Oceananigans

julia> eos = Oceananigans.RoquetIdealizedNonlinearEquationOfState(:cabbeling);

julia> eos.polynomial_coeffs
(R₀₁₀ = -0.0844, R₁₀₀ = 0.7718, R₀₂₀ = -0.004561, R₀₁₁ = 0.0, R₂₀₀ = 0.0, R₁₀₁ = 0.0, R₁₁₀ = 0.0)

References
==========

    - Roquet et al., "Defining a Simplified yet 'Realistic' Equation of State for
      Seawater", Journal of Physical Oceanography (2015).

    - "Thermodynamic Equation of State for Seawater" (TEOS-10), http://www.teos-10.org
"""
function RoquetIdealizedNonlinearEquationOfState(T, flavor=:cabbeling_thermobaricity;
                                                 polynomial_coeffs=optimized_roquet_coeffs[flavor], ρ₀=1024.6)
    typed_coeffs = type_convert_roquet_coeffs(T, polynomial_coeffs)
    return RoquetIdealizedNonlinearEquationOfState{flavor, typeof(typed_coeffs), T}(ρ₀, typed_coeffs)
end

RoquetIdealizedNonlinearEquationOfState(flavor::Symbol=:cabbeling_thermobaricity; kwargs...) =
    RoquetIdealizedNonlinearEquationOfState(Float64, flavor; kwargs...)

""" Return the geopotential depth at `i, j, k` at cell centers. """
@inline D_aac(i, j, k, grid) = @inbounds -grid.zC[k]
const D = D_aac

""" Return the geopotential depth at `i, j, k` at cell z-interfaces. """
@inline D_aaf(i, j, k, grid) = @inbounds -grid.zF[k]

@inline ρ′(i, j, k, grid, eos::RoquetIdealizedNonlinearEquationOfState, C) =
    @inbounds (   eos.polynomial_coeffs.R₁₀₀ * C.S[i, j, k]
                + eos.polynomial_coeffs.R₀₁₀ * C.T[i, j, k]
                + eos.polynomial_coeffs.R₀₂₀ * C.T[i, j, k]^2
                + eos.polynomial_coeffs.R₀₁₁ * C.T[i, j, k] * D(i, j, k, grid)
                + eos.polynomial_coeffs.R₂₀₀ * C.S[i, j, k]^2
                + eos.polynomial_coeffs.R₁₀₁ * C.S[i, j, k] * D(i, j, k, grid)
                + eos.polynomial_coeffs.R₁₁₀ * C.S[i, j, k] * C.T[i, j, k] )

@inline thermal_expansion(i, j, k, grid, eos::RoquetIdealizedNonlinearEquationOfState, C) =
    @inbounds (   eos.polynomial_coeffs.R₀₁₀
                + 2 * eos.polynomial_coeffs.R₀₂₀ * ▶z_aaf(i, j, k, grid, C.T)
                + eos.polynomial_coeffs.R₀₁₁ * D_aaf(i, j, k, grid)
                + eos.polynomial_coeffs.R₁₁₀ * ▶z_aaf(i, j, k, grid, C.T) )

@inline haline_contraction(i, j, k, grid, eos::RoquetIdealizedNonlinearEquationOfState, C) =
    @inbounds (   eos.polynomial_coeffs.R₁₀₀
                + 2 * eos.polynomial_coeffs.R₂₀₀ * ▶z_aaf(i, j, k, grid, C.S)
                + eos.polynomial_coeffs.R₁₀₁ * D_aaf(i, j, k, grid)
                + eos.polynomial_coeffs.R₁₁₀ * ▶z_aaf(i, j, k, grid, C.S) )


importall Base
using Base.LinAlg.Givens

# This needs to be executed once
# Pkg.clone("https://github.com/jiahao/DoublelengthFloat.jl.git")
# using DoublelengthFloat # by Jiahao Chen or use other by Simon Byrne
# Pkg.clone("https://github.com/simonbyrne/DoubleDouble.jl.git")

using DoubleDouble # Jiahao's has error, this is Simon's version

# using DoublelengthFloat

# Define a Symmetric Arrow Matrix Type
immutable SymArrow{T<:Real} <: AbstractMatrix{T}
    D::Vector{T} # diagonal
    z::Vector{T} # 1st row[2:n]
    a::T # arrow top value
    i::Int # arrow top position
end # immutable

# Define its size

size(A::SymArrow, dim::Integer) = length(A.D)+1
size(A::SymArrow)= size(A,1), size(A,1)

# Index into a SymArrow
function getindex(A::SymArrow,i::Integer,j::Integer)
    n=size(A,1)
    if i==j<A.i; return A.D[i]
    elseif i==j>A.i; return A.D[i-1]
    elseif i==j==A.i; return A.a
    elseif i==A.i&&j<A.i; return A.z[j]
    elseif i==A.i&&j>A.i; return A.z[j-1]
    elseif j==A.i&&i<A.i; return A.z[i]
    elseif j==A.i&&i>A.i; return A.z[i-1]
    else return zero(typeof(A.a))
    end
end # getindex

# Dense version of SymArrow
full(A::SymArrow) =[A[i,j] for i=1:size(A,1), j=1:size(A,2)]

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

# Define a Symmetric Diagonal+Rank-One Type
immutable SymDPR1{T} <: AbstractMatrix{T}
    D::Vector{T} # diagonal
    u::Vector{T} # rank one, length n
    r::T # rho
end # immutable

# Define its size

size(A::SymDPR1, dim::Integer) = length(A.D)
size(A::SymDPR1)= size(A,1), size(A,1)

# Index into a SymDPR1
function getindex(A::SymDPR1,i::Integer,j::Integer)
    Aij=A.r*A.u[i]*A.u[j]
    return i==j ? A.D[i]+Aij : Aij
end # getindex

# Dense version of SymDPR1
full(A::SymDPR1) =[A[i,j] for i=1:size(A,1), j=1:size(A,2)]


insts = OrderedDict()

insts["test"] = [
    ((4, MatNegLog()),),
    ((4, MatNegLogDirect()),),
    ((6, MatNegRtdet()),),
    ((6, MatNegRtdetEFExp()),),
]

inversecovariance_insts(f::MatSpecExt, max_d::Int) = [((d, f),) for d in vcat(5, 5:max_d)]

insts["nat"] = inversecovariance_insts(MatNegLog(), 16)
insts["ext"] = inversecovariance_insts(MatNegLogDirect(), 16)

return (InverseCovariance, insts)


"""
    fpw(X, r)

ntrols the random number generator used for initialising the starting point of the optimisation.

# Example
```julia-repl
julia> X = spiraldata(300)
julia> mix = fpw(X, 0.05)
julia> x = rand(mix, 1000)'
julia> plot(X[:,1], X[:,2], "bo", label="data")
julia> plot(x[:,1], x[:,2], ".r", label="generated", alpha=0.7)
julia> legend()
```
"""
fpw(X, r, gamma = 1e-6) = fastparzenwindows(X, r, gamma)

function fastparzenwindows(X, r, gamma=1e-6)

  centres_ind = partition(X, r)

  Q, mu, C = softparzen(X, centres_ind, r, gamma)

  getmixturemodel(Q, mu, C)

end


using Optim

# Proxy function for optimization
function optFunc(x::Vector{CONST_TYPE}, dataset::Dataset{T}, baseline::T,
                 tree::Node, options::Options)::T where {T<:Real}
    setConstants(tree, x)
    return scoreFunc(dataset, baseline, tree, options)
end

# Use Nelder-Mead to optimize the constants in an equation
function optimizeConstants(dataset::Dataset{T},
                           baseline::T, member::PopMember,
                           options::Options)::PopMember where {T<:Real}

    nconst = countConstants(member.tree)
    if nconst == 0
        return member
    end
    x0 = getConstants(member.tree)
    f(x::Vector{CONST_TYPE})::T = optFunc(x, dataset, baseline, member.tree, options)
    if nconst == 1
        algorithm = Newton
    else
        algorithm = NelderMead
    end

    try
        result = optimize(f, x0, algorithm(), Optim.Options(iterations=100))
        # Try other initial conditions:
        for i=1:options.nrestarts
            new_start = x0 .* (convert(CONST_TYPE, 1) .+ convert(CONST_TYPE, 1//2)*randn(CONST_TYPE, size(x0)[1]))
            tmpresult = optimize(f, new_start, algorithm(), Optim.Options(iterations=100))

            if tmpresult.minimum < result.minimum
                result = tmpresult
            end
        end

        if Optim.converged(result)
            setConstants(member.tree, result.minimizer)
            member.score = convert(T, result.minimum)
            member.birth = getTime()
        else
            setConstants(member.tree, x0)
        end
    catch error
        # Fine if optimization encountered domain error, just return x0
        if isa(error, AssertionError)
            setConstants(member.tree, x0)
        else
            throw(error)
        end
    end
    return member
end


using OrdinaryDiffEq, Test, LinearAlgebra

function f(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end
function foop(u,p,t)
  dx = p[1]*u[1] - p[2]*u[1]*u[2]
  dy = -p[3]*u[2] + p[4]*u[1]*u[2]
  [dx,dy]
end

p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(f,[1.0;1.0],(0.0,10.0),p)
proboop = ODEProblem(foop,[1.0;1.0],(0.0,10.0),p)

SPECIAL_INTERPS = [Tsit5(),DP5(),SSPRK22(),OwrenZen3(),OwrenZen4(),OwrenZen5(),
                   BS5(),Vern6(),Vern7(),Vern8(),Vern9(),DP8(),Rosenbrock23(),
                   Rodas4(),Rodas5()]

y1 = zeros(2); y2 = zeros(2)
for alg in SPECIAL_INTERPS
  @show alg
  sol = solve(prob,alg,dt=0.0033,abstol=1e-14,reltol=1e-14)
  soloop = solve(proboop,alg,adaptive=false,tstops=sol.t,abstol=1e-14,reltol=1e-14)
  @test maximum(norm(soloop(t) - sol(t)) for t in 0:0.001:10) < 1e-10
  @test maximum(norm(soloop(y1,t) - sol(y2,t)) for t in 0:0.001:10) < 1e-10
end


struct PlanetaryConstants{T<:AbstractFloat} <: ConstantsCollection
    Ω::T  # Rotation rate of the planet [rad/s].
    f::T  # Nominal value for the Coriolis frequency [rad/s].
    g::T  # Standard acceleration due to gravity [m/s²].
end

PlanetaryConstants(; Ω=1.0, f=0.0, g=1.0) = PlanetaryConstants(Ω, f, g)

function choose_f(Ω, f, lat)
end

function Earth(; f=nothing, lat=nothing)
    Ω = 7.2921150e-5
    g = 9.80665

    if isnothing(f) && isnothing(lat)
        f′ = 1e-4  # Corresponds to a latitude of 43.29°N.
    elseif isnothing(lat)
        f′ = f
    elseif isnothing(f)
        f′ = 2*Ω*sind(lat)
    else
        throw(ArgumentError("Cannot specify both f and lat!"))
    end

    abs(f′) <= 2Ω || throw(ArgumentError("Coriolis parameter |f| cannot be larger than 2Ω!"))

    PlanetaryConstants(Ω, f′, g)
end

function EarthStationary()
    Ω = 0
    f = 0
    g = 9.80665
    PlanetaryConstants(Ω, f, g)
end

function Europa(; f=nothing, lat=nothing)
    # Values taken from Soderlund (Table 1, 2019) [arXiv:1901.04093].
    Ω = 2.1e-5
    g = 1.3

    if isnothing(f) && isnothing(lat)
        throw(ArgumentError("Must specify Coriolis parameter f or latitude!"))
    elseif isnothing(lat)
        f′ = f
    elseif isnothing(f)
        f′ = 2*Ω*sind(lat)
    else
        throw(ArgumentError("Cannot specify both f and lat!"))
    end

    abs(f′) <= 2Ω || throw(ArgumentError("Coriolis parameter |f| cannot be larger than 2Ω!"))

    PlanetaryConstants(Ω, f, g)
end

function Enceladus(; f=nothing, lat=nothing)
    # Values taken from Soderlund (Table 1, 2019) [arXiv:1901.04093].
    Ω = 5.3e-5
    g = 0.1

    if isnothing(f) && isnothing(lat)
        throw(ArgumentError("Must specify Coriolis parameter f or latitude!"))
    elseif isnothing(lat)
        f′ = f
    elseif isnothing(f)
        f′ = 2*Ω*sind(lat)
    else
        throw(ArgumentError("Cannot specify both f and lat!"))
    end

    abs(f′) <= 2Ω || throw(ArgumentError("Coriolis parameter |f| cannot be larger than 2Ω!"))

    PlanetaryConstants(Ω, f, g)
end


using Base.LinAlg.BLAS

N = 4000
srand(0)
A = rand(N,N)
B = rand(N,N)
result = zeros(N,N)

function substep{T}(a::Matrix{T}, beta::T, b::Matrix{T}, result::Matrix{T})
    for j=1:size(a,2)
        for i=1:size(b,1)
            @inbounds result[i,j] = a[i,j] + beta*b[i,j]
        end
    end
end

function f0{T}(A::Matrix{T}, alpha::T, B::Matrix{T}, result::Matrix{T}, repeat=1)
    for i=1:repeat
        substep(A, alpha, B, result)
    end
    return nothing
end

@time f0(A, 0.1, B, result)
@time f0(A, 0.1, B, result)

A_vec = reshape(A, length(A))
result_vec = reshape(result, length(result))
N = length(A_vec)
@time BLAS.axpy!(N, 0.1, A, 1, B, 1)
@time BLAS.axpy!(N, 0.1, A, 1, B, 1)

# @time BLAS.blascopy!(N, A_vec, 1, result_vec, 1)
# @time BLAS.blascopy!(N, A_vec, 1, result_vec, 1)
# @time BLAS.scale!(0.1, A)
# @time BLAS.scale!(0.1, A)


# a simple LP example
#
#   maximize 3x + 2y
#
#   s.t. 1 x + 1 y == 20
#        x \in [0, 10], y \in [0, 5]
#

using Xpress, Compat.Test

# @testset "Basics 1" begin
    model = Xpress.Model("lp_01", :minimize)

    # add variables
    add_cvar!(model, 3.0, 0., 10.)  # x
    add_cvar!(model, 2.0,  0., 5.)  # y

    # add constraints
    add_constr!(model, [1., 1.], '=', 20.)

    @test Xpress.get_obj(model) == [3, 2]
    @test Xpress.get_constrmatrix(model) == sparse([1 1])
    @test Xpress.get_rhs(model) == [20]
    @test Xpress.get_lb(model) == [0, 0]
    @test Xpress.get_ub(model) == [10, 5]
    @test Xpress.get_sense(model) == [:(==)]

    # perform optimization
    optimize(model)

    ans = get_optiminfo(model)
    @test ans.status_lp == :infeasible

    lrp = Float64[0.0]
    grp = Float64[0.1]
    lbp = Float64[0.0,0.0]
    ubp = Float64[0.0,0.0]

    @test 0 == Xpress.repairweightedinfeasibility(model, lrp, grp, lbp, ubp, phase2 = Cchar('d'))
    ans = get_optiminfo(model)
    @test ans.status_lp == :optimal

@show Xpress.get_objval(model)
@show Xpress.get_solution(model)
@show Xpress.get_dual(model)

    @test 0 == Xpress.repairweightedinfeasibility(model, lrp, grp, lbp, ubp, phase2 = Cchar('x'))
    ans = get_optiminfo(model)
    @test ans.status_lp == :optimal

@show Xpress.get_objval(model)
@show Xpress.get_solution(model)
@show Xpress.get_dual(model)

    @test 0 == Xpress.repairweightedinfeasibility(model, lrp, grp, lbp, ubp, phase2 = Cchar('f'))
    ans = get_optiminfo(model)
    @test ans.status_lp == :optimal

@show Xpress.get_objval(model)
@show Xpress.get_solution(model)
@show Xpress.get_dual(model)

    @test 0 == Xpress.repairweightedinfeasibility(model, lrp, grp, lbp, ubp, phase2 = Cchar('n'))
    ans = get_optiminfo(model)
    @test ans.status_lp == :optimal

@show Xpress.get_objval(model)
@show Xpress.get_solution(model)
@show Xpress.get_dual(model)
;


# ---
# title: 1385. Find the Distance Value Between Two Arrays
# id: problem1385
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/find-the-distance-value-between-two-arrays/description/>
# hidden: true
# ---
#
# Given two integer arrays `arr1` and `arr2`, and the integer `d`, _return the
# distance value between the two  arrays_.
#
# The distance value is defined as the number of elements `arr1[i]` such that
# there is not any element `arr2[j]` where `|arr1[i]-arr2[j]| <= d`.
#
#
#
# **Example 1:**
#
#
#
#     Input: arr1 = [4,5,8], arr2 = [10,9,1,8], d = 2
#     Output: 2
#     Explanation:
#     For arr1[0]=4 we have:
#     |4-10|=6 > d=2
#     |4-9|=5 > d=2
#     |4-1|=3 > d=2
#     |4-8|=4 > d=2
#     For arr1[1]=5 we have:
#     |5-10|=5 > d=2
#     |5-9|=4 > d=2
#     |5-1|=4 > d=2
#     |5-8|=3 > d=2
#     For arr1[2]=8 we have:
#     **|8-10|=2 <= d=2**
#     **|8-9|=1 <= d=2**
#     |8-1|=7 > d=2
#     **|8-8|=0 <= d=2**
#
#
# **Example 2:**
#
#
#
#     Input: arr1 = [1,4,2,3], arr2 = [-4,-3,6,10,20,30], d = 3
#     Output: 2
#
#
# **Example 3:**
#
#
#
#     Input: arr1 = [2,1,100,3], arr2 = [-5,-2,10,-3,7], d = 6
#     Output: 1
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr1.length, arr2.length <= 500`
#   * `-10^3 <= arr1[i], arr2[j] <= 10^3`
#   * `0 <= d <= 100`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


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

T = 100                                 # Number of periods in the data
P = 2                                   # Number of lags in the factor equation
L = 2                                   # Number of AR lags in the observation equation
N = 10                                  # Total number of series
K = 1                                   # Total number of factors (only the global factor)

H = [ones(N) ident(N) zeros(N)]
R = zeros(N, N)
A = zeros(N, 2)
Z = zeros(2, 2)

F = zeros(P + N, P + N)
F[1, 1] = 0.5
F[1, P+N] = 0.1
F[P+N, 1] = 1.0
F[2:(P+N-1), 2:(P+N-1)] = 0.25 * ident(N)

Q = zeros(P + N, P + N)
Q[1, 1] = 1.0
for i in 1:N
    Q[1+i, 1+i] = 2.0
end

μ = zeros(P + N)

num_obs = T

# Gather all SS parameters
ssmodel = SSModel(H, A, F, μ, R, Q, Z)

# Simulate common components model in state space form
data_y, data_z, data_β = simulateSSModel(num_obs, ssmodel)

samplerparams = DFMStruct(P, L, 1000, 50)            # Set model priors

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

results = PCA1LevelEstimator(data_y, samplerparams)

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

medians = Any[]
quant33 = Any[]
quant66 = Any[]
stds = Any[]

j = 1
for i in 1:size(results.F)[1]
    push!(stds, std(results.F[i, j, :]))
    push!(quant33, quantile(results.F[i, :], 0.33))
    push!(quant66, quantile(results.F[i, :], 0.66))
    push!(medians, median(results.F[i, :]))
end

plot(data_β[:,1])
plot!(results.means.F[:, j])
#plot(medians)
plot!(quant33)
plot!(quant66)


lines = readlines("input.txt")
field = zeros(Int, (1024,1024))
re = r"(\d+)"

for line = lines
    x1,y1,x2,y2 = 1 .+ parse.(Int, m.match for m = eachmatch(re, line))
    dx,dy = x1 < x2 ? 1 : -1, y1 < y2 ? 1 : -1
    if x1 != x2 && y1 != y2
        for ij = CartesianIndex.(zip(y1:dy:y2, x1:dx:x2))
            field[ij] += 1
        end
    else
        field[y1:dy:y2, x1:dx:x2] .+= 1
    end
end

println(sum(@. field > 1))


"""
This module reimplements models from the paper:

    Spoerer, C.J., McClure, P. and Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
module Spoerer2017
using Flux
using ...Splitters
using ...Mergers
using ...FeedbackChains
using ..LRNs
using ..Flatten

export spoerer_model_b, spoerer_model_bf, spoerer_model_bk,
       spoerer_model_bl, spoerer_model_bt, spoerer_model_blt

"""
    spoerer_model_fw(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)

Generate one of the forward models (B, B-K, B-F) from the paper:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
function spoerer_model_fw(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)
    return Chain(
        Conv(kernel, channels=>features, relu, pad=map(x -> x ÷ 2, kernel)),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool((2,2), stride=(2,2)),
        Conv(kernel, features=>features, relu, pad=map(x -> x ÷ 2, kernel)),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool(map(x -> x ÷ 2, inputsize)),
        flatten,
        Dense(features, classes, σ)
    )
end # function spoerer_model_fw

"""
    spoerer_model_b(T; channels=1, inputsize=(28, 28), classes=10)

Generate the bottom-up (B) convolutional neural network from:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
spoerer_model_b(T; channels=1, inputsize=(28,28), classes=10) =
    spoerer_model_fw(T, channels=channels, inputsize=inputsize, classes=classes)

"""
    spoerer_model_bk(T; channels=1, inputsize=(28, 28), classes=10)

Generate the convolutional neural network with increased kernel size (BK) from:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
spoerer_model_bk(T; channels=1, inputsize=(28,28), classes=10) =
    spoerer_model_fw(T, channels=channels, inputsize=inputsize, classes=classes, kernel=(5,5))

"""
    spoerer_model_bf(T; channels=1, inputsize=(28, 28), classes=10)

Generate the convolutional neural network with additional feature maps (BF) from:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
spoerer_model_bf(T; channels=1, inputsize=(28,28), classes=10) =
    spoerer_model_fw(T, channels=channels, inputsize=inputsize, classes=classes, features=64)

# TODO: currently, the forward, lateral and backward convolutions all have their
#       own biases. This should not make the model more powerful, as they are
#       combined additively before the non-linearity, but it wastes some resources.
"""
    spoerer_model_bl(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)

Generate the convolutional neural network with lateral recurrence (BL) from:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
function spoerer_model_bl(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)
    return FeedbackChain(
        Conv(kernel, channels=>features, pad=map(x -> x ÷ 2, kernel)),
        Merger("l1", ConvTranspose(kernel, features=>features, pad=map(x -> x ÷ 2, kernel)), +),
        x -> relu.(x),
        Splitter("l1"),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool((2,2), stride=(2,2)),
        Conv(kernel, features=>features, pad=map(x -> x ÷ 2, kernel)),
        Merger("l2", ConvTranspose(kernel, features=>features, pad=map(x -> x ÷ 2, kernel)), +),
        x -> relu.(x),
        Splitter("l2"),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool(map(x -> x ÷ 2, inputsize)),
        flatten,
        Dense(features, classes, σ)
    )
end # function spoerer_model_bl

"""
    spoerer_model_bt(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)

Generate the convolutional neural network with top-down recurrence (BT) from:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
function spoerer_model_bt(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)
    return FeedbackChain(
        Conv(kernel, channels=>features, pad=map(x -> x ÷ 2, kernel)),
        Merger("l2",  ConvTranspose((2,2), features=>features, stride=2), +),
        x -> relu.(x),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool((2,2), stride=(2,2)),
        Conv(kernel, features=>features, relu, pad=map(x -> x ÷ 2, kernel)),
        Splitter("l2"),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool(map(x -> x ÷ 2, inputsize)),
        flatten,
        Dense(features, classes, σ)
    )
end # function spoerer_model_bt

"""
    spoerer_model_blt(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)

Generate the convolutional neural network with lateral and top-down recurrence (BLT) from:

    Spoerer, C.J., McClure, P. & Kriegeskorte, N. (2017).
    Recurrent convolutional neural networks: a better model of biological object recognition.
    Frontiers in Psychology 8, 1551.
"""
function spoerer_model_blt(T; channels=1, inputsize=(28, 28), kernel=(3,3), features=32, classes=10)
    return FeedbackChain(
        Conv(kernel, channels=>features, pad=map(x -> x ÷ 2, kernel)),
        Merger("l1", ConvTranspose(kernel, features=>features, pad=map(x -> x ÷ 2, kernel)), +),
        Merger("l2", ConvTranspose((2,2), features=>features, stride=2), +),
        x -> relu.(x),
        Splitter("l1"),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool((2,2), stride=(2,2)),
        Conv(kernel, features=>features, pad=map(x -> x ÷ 2, kernel)),
        Merger("l2", ConvTranspose(kernel, features=>features, pad=map(x -> x ÷ 2, kernel)), +),
        x -> relu.(x),
        Splitter("l2"),
        LRN(T(1.0), T(0.0001), T(0.5), 5),
        MaxPool(map(x -> x ÷ 2, inputsize)),
        flatten,
        Dense(features, classes, σ)
    )
end # function spoerer_model_blt

end # module Spoerer2017


using Test
using ParticleInCell

@testset "ParticleMeshCoupling 2D" begin

    n_cells = 10 # Number of cells
    n_particles = 4 # Number of particles
    spline_degree = 3 # Spline degree
    xmin, xmax, nx = 0.0, 2.0, n_cells
    ymin, ymax, ny = 0.0, 1.0, n_cells

    mesh = TwoDGrid(xmin, xmax, nx, ymin, ymax, ny)

    volume = (xmax - xmin) * (ymax - ymin)

    x_vec = [0.1 0.65 0.7 1.5; 0.0 0.0 0.0 0.0]
    v_vec = [1.5 0.00 0.0 0.0; 0.0 0.5 0.0 0.0]

    pg = ParticleGroup{2,2}(n_particles, charge = 1.0, mass = 1.0, n_weights = 1)

    for i_part = 1:n_particles
        pg.array[1, i_part] = x_vec[1, i_part]
        pg.array[2, i_part] = x_vec[2, i_part]
        pg.array[3, i_part] = v_vec[1, i_part]
        pg.array[4, i_part] = v_vec[2, i_part]
        pg.array[5, i_part] = 1.0 / n_particles
    end

    # Initialize the kernel
    kernel = ParticleMeshCoupling2D(pg, mesh, spline_degree, :collocation)

    # Reference values of the shape factors
    index_grid = zeros(Int, (2, 4))
    index_grid[1, :] .= [-2, 1, 1, 5]
    index_grid[2, :] .= [-3, -3, -3, -3]

    v_grid = zeros(Float64, (4, 2, 4))
    v_grid[:, 1, 1] .= [
        2.0833333333333332E-002,
        0.47916666666666663,
        0.47916666666666663,
        2.0833333333333332E-002,
    ]
    v_grid[:, 1, 3] .= v_grid[:, 1, 1]
    v_grid[:, 1, 4] .= v_grid[:, 1, 1]
    v_grid[:, 1, 2] .= [
        7.0312500000000000E-002,
        0.61197916666666663,
        0.31510416666666663,
        2.6041666666666665E-003,
    ]
    v_grid[1, 2, :] .= 0.0
    v_grid[2, 2, :] .= 1.0 / 6.0
    v_grid[3, 2, :] .= 2.0 / 3.0
    v_grid[4, 2, :] .= 1.0 / 6.0

    ρ_dofs = zeros(Float64, 100)
    ρ_dofs1 = zeros(Float64, 100)
    ρ_dofs_ref = zeros(Float64, 100)
    ρ_dofs_pp = zeros(Float64, (16, 100))

    # Accumulate ρ
    for i_part = 1:n_particles
        x = pg.array[1, i_part]
        y = pg.array[2, i_part]
        w = pg.array[5, i_part]
        GEMPIC.add_charge!(ρ_dofs, kernel, x, y, w)
        GEMPIC.add_charge_pp!(ρ_dofs1, kernel, x, y, w)
    end

    ρ_dofs_ref[8:10] .= v_grid[1:3, 1, 1]
    ρ_dofs_ref[1] = v_grid[4, 1, 1]
    ρ_dofs_ref[1:4] .= ρ_dofs_ref[1:4] .+ v_grid[:, 1, 2] .+ v_grid[:, 1, 3]
    ρ_dofs_ref[5:8] .= ρ_dofs_ref[5:8] .+ v_grid[:, 1, 4]

    ρ_dofs_ref[71:80] .= ρ_dofs_ref[1:10] ./ 6.0
    ρ_dofs_ref[81:90] .= ρ_dofs_ref[1:10] .* 2.0 / 3.0
    ρ_dofs_ref[91:100] .= ρ_dofs_ref[1:10] ./ 6.0
    ρ_dofs_ref[1:10] .= 0.0

    ρ_dofs_ref .*= n_cells^2 / volume / n_particles

    @test ρ_dofs ≈ ρ_dofs_ref
    @test ρ_dofs1 ≈ ρ_dofs_ref

    GEMPIC.b_to_pp_2d!(ρ_dofs_pp, kernel.spline1, kernel.spline2, ρ_dofs)

    # For evaluation check
    particle_values = zeros(4)
    particle_values1 = zeros(4)
    particle_values_ref = zeros(4)

    # Test function evaluation
    for i_part = 1:n_particles
        xp = pg.array[1, i_part]
        yp = pg.array[2, i_part]
        particle_values[i_part] = evaluate(kernel, xp, yp, ρ_dofs)
        particle_values1[i_part] = evaluate_pp(kernel, xp, yp, ρ_dofs_pp)
    end

    particle_values_ref =
        [1.1560058593749998, 2.3149278428819446, 2.2656250000000000, 1.1512586805555554] ./
        volume

    fill!(particle_values_ref, 0.0)

    for i_part = 1:n_particles
        for i = 1:4
            i1 = mod(index_grid[1, i_part] + i - 2, n_cells)
            for j = 1:4
                i2 = mod(index_grid[2, i_part] + j - 2, n_cells)
                res =
                    v_grid[i, 1, i_part] *
                    v_grid[j, 2, i_part] *
                    ρ_dofs_ref[i1+i2*n_cells+1]
                particle_values_ref[i_part] += res
            end
        end
    end

    @test particle_values ≈ particle_values_ref
    @test particle_values1 ≈ particle_values_ref

end


export plasma

function plasma( mesh :: Mesh, nbpart :: Int64 )

    kx     = 0.5
    alpha  = 0.05

    nx, ny = mesh.nx, mesh.ny
    dx, dy = mesh.dx, mesh.dy
    dimx = mesh.xmax - mesh.xmin
    dimy = mesh.ymax - mesh.ymin

    weight = (dimx * dimy) / nbpart

    particles = Particles(nbpart, weight )

    k = 1
    while (k<=nbpart)

        xi   = rand() * dimx
        yi   = rand() * dimy
        zi   = (2.0 + alpha) * rand()
        temm = 1.0 + sin(yi) + alpha * cos(kx*xi)

        if (temm >= zi)
            particles.x[1,k] = xi
            particles.x[2,k] = yi
            k = k + 1
        end

    end

    k = 1
    while (k<=nbpart)

        xi   = (rand()-0.5)*10
        yi   = (rand()-0.5)*10
        zi   = rand()
        temm = ( exp(-((xi-2)^2 + yi^2)/2)
               + exp(-((xi+2)^2 + yi^2)/2))/2

        if (temm >= zi)
            particles.v[1,k] = xi
            particles.v[2,k] = yi
            k = k + 1
        end

    end

    particles

end


using ZeroInflatedDistributions
using Distributions
using QuadGK
using Random
using Test

@testset "ZeroInflatedDistributions.jl" begin
    @testset "Link functions" begin
        p1 = 0
        p2 = 1
        # Logit-log link tests
        ll = LogitLogLink()
        @test encprob(ll, p1, p2) == 1//2
        @test posrate(ll, p1, p2) == Float64(ℯ)
        # Poisson link tests
        pl = PoissonLink()
        @test isone(pl.offset)
        @test_throws DomainError PoissonLink(-1)
        @test encprob(pl, p1, p2) == 0.6321205588285577
        @test posrate(pl, p1, p2) == 4.300258535328371
        # Identity link tests
        il = IdentityLink()
        @test_throws DomainError encprob(il, -1)
        @test_throws DomainError encprob(il, 1.1)
        @test encprob(il, 0.5) == 0.5
        @test posrate(il, 1.0) == 1
    end

    @testset "Zero-inflated distribution constructors" begin
        p1 = 0
        p2 = 1
        disp = 0.5
        ll = LogitLogLink()
        pl = PoissonLink()
        # Test inner constructor
        zil0 = ZeroInflatedDistribution(Bernoulli(0.5), LogNormal(0.0, 1.0))
        @test logpdf(zil0, 0) == log(0.5)
        @test logpdf(zil0, 1) == logpdf(LogNormal(0.0, 1.0), 1) + log(0.5)

        # Without bias correction, posrate is the median of the log-normal
        zil1 = ZeroInflatedDistribution(ll, LogNormal, p1, p2, disp; biascorrect = false)
        @test median(zil1.posdist) == posrate(ll, p1, p2)

        # With bias correction (default) it is the mean
        zil2 = ZeroInflatedDistribution(ll, LogNormal, p1, p2, disp)
        @test mean(zil2.posdist) == posrate(ll, p1, p2)

        zil3 = ZeroInflatedDistribution(pl, Gamma, p1, p2, disp)
        @test mean(zil3.posdist) == posrate(pl, p1, p2)
        @test std(zil3.posdist) == disp

        zil4 = ZeroInflatedDistribution(pl, InverseGamma, p1, p2, disp)
        @test mean(zil4.posdist) == posrate(pl, p1, p2)
        @test std(zil4.posdist) == disp

        zil5 = ZeroInflatedDistribution(pl, InverseGaussian, p1, p2, disp)
        @test mean(zil5.posdist) == posrate(pl, p1, p2)
        @test shape(zil5.posdist) == disp
    end

    @testset "Zero-inflated log-pdfs" begin
        zil = ZeroInflatedDistribution(Bernoulli(0.5), LogNormal(0.0, 1.0))
        @test logpdf(zil, 0) == log(0.5)
        @test logpdf(zil, 1) == logpdf(LogNormal(0.0, 1.0), 1) + log(0.5)
    end

    @testset "Random zero-inflated data generation" begin
        n = 1_000
        zil = ZeroInflatedDistribution(Bernoulli(0.5), LogNormal(-1/2, 1.0))
        x = rand(zil, n)
        @test any(x .== 0)
        @test all(x .≥ 0)
    end

    @testset "Zero-inflated distribution statistics" begin
        zil = ZeroInflatedDistribution(Bernoulli(0.75), LogNormal(-1/2, 1.0))
        zil0 = ZeroInflatedDistribution(Bernoulli(0), LogNormal(-1/2, 1.0))
        zil1 = ZeroInflatedDistribution(Bernoulli(1), LogNormal(-1/2, 1.0))
        zilgam = ZeroInflatedDistribution(Bernoulli(0.85), Gamma(2, 4))

        @testset "Zero-inflated distribution supports" begin
            @test !insupport(zil, -1)
            @test all(minimum.([zil, zil0, zil1, zilgam]) .== 0)
            @test all(maximum.([zil, zil0, zil1, zilgam]) .== Inf)
        end

        @testset "Test means and variances via simulation" begin
            Random.seed!(13579)
            n = 10_000

            x = rand(zil, n)
            x0 = rand(zil0, n)
            x1 = rand(zil1, n)

            sd = std(zil) / sqrt(n)
            sd1 = std(zil1) / sqrt(n)

            @test isapprox(mean(x), 0.75, atol = 2 * sd)
            @test all(x0 .== 0)
            @test isapprox(mean(x1), 1, atol = 2 * sd1)

            # @test isapprox(var(zil), var(x))
            @test var(zil0) == 0
            @test var(zil1) == var(zil1.posdist)
        end

        @testset "Test means and variances via quadrature" begin
            # 𝔼[X]
            quadex = quadgk(x -> x * pdf(zil, x), 0, Inf)
            quadexgam = quadgk(x -> x * pdf(zilgam, x), 0, Inf)
            # 𝔼[X²]
            quadex2 = quadgk(x -> x^2 * pdf(zil, x), 0, Inf)
            quadexgam2 = quadgk(x -> x^2 * pdf(zilgam, x), 0, Inf)
            # 𝕍[X] = 𝔼[X^2] - 𝔼[X]^2
            quadvar = quadex2[1] - mean(zil)^2
            quadvargam = quadexgam2[1] - mean(zilgam)^2

            @test isapprox(quadex[1], mean(zil), atol = quadex[2])
            @test isapprox(quadvar, var(zil), atol = quadex2[2])
            @test isapprox(quadexgam[1], mean(zilgam), atol = quadexgam[2])
            @test isapprox(quadvargam, var(zilgam), atol = quadexgam2[2])
        end

        @testset "Test CDF function" begin
            @test cdf(zil, 0) == failprob(zil.encdist)
            @test cdf(zil0, 0) == 1
            @test cdf(zil1, 2) == cdf(zil1.posdist, 2)
            @test cdf(zilgam, 10) ==
                failprob(zilgam.encdist) +
                cdf(zilgam.posdist, 10) * succprob(zilgam.encdist)
        end

        @testset "Test quantile function" begin
            @test quantile(zil, 0.2) == 0
            @test quantile(zilgam, 0.1) == 0
            @test quantile(zil, 0.5) ==
                quantile(zil.posdist,
                         (0.5 - failprob(zil.encdist)) / succprob(zil.encdist))
            @test quantile(zil1, 0.9) == quantile(zil1.posdist, 0.9)
            @test quantile(zilgam, 0.75) ==
                quantile(zilgam.posdist,
                         (0.75 - failprob(zilgam.encdist)) /
                         succprob(zilgam.encdist))
        end

        @testset "Test modes function" begin
            @test all(modes(zil) .== [0, mode(LogNormal(-0.5, 1))])
            @test all(modes(zilgam) .== [0, mode(Gamma(2, 4))])
        end
    end
end


@testset "DMD" begin
    # Create some linear data
    A = [0.9 -0.2; 0.0 0.2]
    y = [[10.; -10.]]
    for i in 1:10
        push!(y, A*y[end])
    end
    X = hcat(y...)

    estimator = DMD(X)
    # TODO export
    #@test isstable(estimator)
    @test operator(estimator) ≈ A
    @test eigvals(estimator) ≈ eigvals(A)
    @test eigvecs(estimator) ≈ eigvecs(A)
    @test is_discrete(estimator)
    # TODO add Array{T,2}
    #@test estimator(X[:, 1]) ≈ X[:, 2]
    @test_nowarn update!(estimator, X[:, end-1], X[:,end])
    @test_throws AssertionError outputmap(estimator)

    @variables u[1:2]
    b = Basis(u, u)
    sys = EDMD(X[:, 1:3], b)
    @test_nowarn update!(sys, X[:, 3:end-1], X[:, 4:end], threshold = 1e-17)
    @test operator(sys) ≈ A

    # Test for a large system with reduced svd
    dd(u, p, t) = -0.9*ones(length(u))*u[1] - 0.05*u
    u0 = randn(100)*100.0
    u0[1] = 100.0
    prob = ODEProblem(dd, u0, (0.0, 100.0))
    sol = solve(prob, Tsit5(), saveat = 0.1)

    X = sol[:,:]
    DX = sol(sol.t, Val{1})[:,:]

    t = DMDSVD(0.01)
    d = gDMD(X, DX, alg = t)

    test = ODEProblem(d, u0, (0.0, 100.0))
    sol_ = solve(test, Tsit5(), saveat = 0.1)
    @test norm(sol-sol_, Inf) < 2.0

    d2 = gDMD(sol.t, X)
    d3 = gDMD(sol.t, X, dt = 0.2)
    @test generator(d) ≈ generator(d2) atol = 1e-1
    @test generator(d) ≈ generator(d3) atol = 1e-1

    K = -0.5*I + [0 0 -0.2; 0.1 0 -0.1; 0. -0.2 0]
    F = qr(randn(20, 3))
    Q = F.Q[:, 1:3]
    dudt(u, p, t) = K*u
    prob = ODEProblem(dudt, [10.0; 0.3; -5.0], (0.0, 10.0))
    sol_ = solve(prob, Tsit5(), saveat = 0.01)
    alg = TOTALDMD(3, DMDPINV())
    X = Q*sol_[:,:] + 1e-3*randn(20, 1001)
    DX = Q*sol_(sol_.t, Val{1})[:,:] + 1e-3*randn(20, 1001)
    approx = gDMD(X, DX, alg = alg)
    @test Q'*generator(approx)*Q ≈ K atol = 1e-1
    @test Q*K*Q' ≈ generator(approx) atol = 1e-1
    alg = TOTALDMD(0.01, DMDSVD())
    approx = gDMD(X, DX, alg = alg)
    @test Q'*generator(approx)*Q ≈ K atol = 1e-1
    @test Q*K*Q' ≈ generator(approx) atol = 1e-1

end

@testset "EDMD" begin
    # Test for linear system
    function linear_sys(u, p, t)
        x = -0.9*u[1]
        y = -0.3*u[2]
        return [x;y]
    end

    u0 = [π; 1.0]
    tspan = (0.0, 50.0)
    prob = DiscreteProblem(linear_sys, u0, tspan)
    sol = solve(prob,FunctionMap())

    @variables u[1:2]
    h = [u[1]; u[2]; sin(u[1]); cos(u[1]); u[1]*u[2]; u[2]^2]
    basis = Basis(h, u)

    estimator = EDMD(sol[:,:], basis, alg = DMDSVD())
    @test basis == estimator.basis
    basis_2 = reduce_basis(estimator, threshold = 1e-5)
    @test size(basis_2)[1] < size(basis)[1]

    estimator_2 = EDMD(sol[:,:], basis_2)
    p1 = DiscreteProblem(estimator, u0, tspan)
    s1 = solve(p1,FunctionMap())
    p2 = DiscreteProblem(estimator_2, u0, tspan)
    s2 = solve(p2,FunctionMap())
    # TODO add linear dynamics ?
    #p3 = DiscreteProblem(linear_dynamics(estimator_2), estimator_2(u0), tspan, [])
    #s3 = solve(p3,FunctionMap())

    @test sol[:,:] ≈ s1[:,:]
    @test sol[:,:] ≈ s2[:,:]
    @test eigvals(estimator_2) ≈ [-0.9; -0.3]
    @test isstable(estimator_2)

    # Test for nonlinear system
    function nonlinear_sys(du, u, p, t)
        du[1] = 0.9u[1] + 0.1u[2]^2
        du[2] = sin(u[1]) - 0.1u[1]
    end

    prob = DiscreteProblem(nonlinear_sys, u0, tspan)
    sol = solve(prob,FunctionMap())
    estimator = EDMD(sol[:,:], basis, alg = DMDPINV())
    sys = ODESystem(estimator)
    #@test isa(sys, ODESystem)
    dudt = ODEFunction(sys)
    p4 = DiscreteProblem(dudt, u0, tspan)

    s4 = solve(p4,FunctionMap())
    @test sol[:,:] ≈ s4[:,:]

    function slow_manifold(du, u, p, t)
      du[1] = p[1]*u[1]
      du[2] = p[2]*(u[2]-u[1]^2)
    end

    u0 = [3.0; -2.0]
    tspan = (0.0, 10.0)
    p = [-0.05, -0.9]

    problem = ODEProblem(slow_manifold, u0, tspan, p)
    sol_slow = solve(problem, Tsit5(), saveat = 0.2)

    X = Array(sol_slow)
    # This enforces more accurate results
    DX = similar(X)
    for (i,dx) in enumerate(eachcol(DX))
        slow_manifold(dx, X[:, i], p, sol_slow.t[i])
    end
    basis = Basis([u; u[1]^2], u)
    estimator2 = gEDMD(X, DX, basis)
    A_analytical = [p[1] 0 0; 0 p[2] -p[2]; 0 0 2*p[1]]
    outputmap(estimator2)
    generator(estimator2)
    @test generator(estimator2) ≈ [p[1] 0 0; 0 p[2] -p[2]; 0 0 2*p[1]] atol =  1e-3
    @test outputmap(estimator2) ≈ [1 0 0 ; 0 1 0] atol = 1e-3
    @test abs.(modes(estimator2)) ≈ abs.(eigvecs(A_analytical)) atol = 1e-3
    @test frequencies(estimator2) ≈ eigvals(A_analytical) atol = 1e-3

    function nonlinear_sys2(du, u, p, t)
        du[1] = u[2]
        du[2] = -0.9u[1]
    end

    prob_nl = ODEProblem(nonlinear_sys2, u0, (0.0, 10.0))
    sol_nl = solve(prob_nl, Tsit5())

    DX = sol_nl(sol_nl.t, Val{1})[:,:]
    X = Array(sol_nl)

    basis = Basis(u, u)
    estimator3 = gEDMD(X, DX, basis)
    @test generator(estimator3) ≈ [0  1.0; -0.9 0]
    @test outputmap(estimator3) ≈ I(2)

    d2 = gEDMD(sol_nl.t, X, basis)
    d3 = gEDMD(sol_nl.t, X, basis, dt = 0.1)
    @test generator(estimator3) ≈ generator(d2) atol = 1e-1
    @test generator(estimator3) ≈ generator(d3) atol = 1e-1

end

@testset "DMDc" begin
    # Define measurements from unstable system with known control input
    X = [4 2 1 0.5 0.25; 7 0.7 0.07 0.007 0.0007]
    U = [-4 -2 -1 -0.5]
    B = Float32[1; 0]

    # But with a little more knowledge
    sys = DMDc(X, U, B = B)
    @test operator(sys) ≈[1.5 0; 0 0.1]
    @test inputmap(sys) ≈ [1.0; 0.0]
    @test !isstable(sys)
    @test_nowarn eigen(sys)
    sys2 = DMDc(X, U, B = B, alg = DMDSVD())
    @test operator(sys2) ≈ operator(sys)
    sys3 = gDMDc(X[:, 1:end-1], X[:, 2:end], U, B = B)
    @test generator(sys3) ≈ [1.5 0; 0 0.1]

    A = [0.9 0.3; -0.2 0.9]
    B = [1.0 0;0 0.5]
    X = zeros(Float64, 2, 11)
    U = zeros(Float64, 2, 10)
    X[:, 1] = [1.0; -3.0]
    for i in 1:10
        U[:, i] .= [sin(2.0/5.0*i); cos(5.0/7.0*i)]
        X[:, i+1] .= A*X[:,i]+B*[sin(2.0/5.0*i); cos(5.0/7.0*i)]
    end

    sys4 = DMDc(X, U)
    @test operator(sys4) ≈ A
    @test inputmap(sys4) ≈ B
    sys5 = gDMDc(X[:, 1:end-1], X[:, 2:end], U[:, 1:end])
    @test_throws AssertionError operator(sys5)
    @test generator(sys5) ≈ A
    @test inputmap(sys5) ≈ B

    sys6 = gDMDc(collect(0.0:9.0), X[:, 1:end-1], U, B = B)
    @test exp(generator(sys6)) ≈ A atol = 1e-1
    @test inputmap(sys6) ≈ B

end


module DataStructuresAlgorithms

using Random

export Stack, DynamicArray
export linearsearch, binarysearch
export selectionsort!, insertionsort!, mergesort, quicksort!
export SLinkedList, DLinkedList, addlast!, addfirst!, add!, removelast!, removefirst!, remove!, peeklast, peekfirst
export Queue, enqueue!, dequeue!, peek, peekfirst, peeklast
export RMHeap, BNode, add!, remove!, meld!, MinPQ, add!, isheap, greater, getmin
export BST, BSTNode, printsort, isbst, maximum_node
export SCHash, SCLHash, OAHash, compress

import Base.push!, Base.pop!, Base.isempty, Base.peek, Base.show
import Base.getindex, Base.resize!, Base.setindex!, Base.:(==)
import Base.iterate, Base.eltype, Base.length, Base.size, Base.get!, Base.put!
import Base.delete!, Base.collect, Base.minimum, Base.maximum, Base.get, Base.hash

include("dynamicarray.jl")
include("stack.jl")
include("list.jl")
include("queue.jl")
include("sort.jl")
include("search.jl")
include("search.jl")
include("heap.jl")
include("tree.jl")
include("hash.jl")

end # module


"""
$(SIGNATURES)

This function is the main API for eQTL scans. EQTL scan process includes
    If no covariate:
    - filter maf if the threshold is greater than 0
    - standardizing phenotype matrix (Y) and genotype matrix (G)
    - calculate correlation (R) matrix.
    - computes log of odds (LOD) score matrix, or p-value if `lod_or_pval` is set to `lod`
    - calculate maximum LOD score if `export_matrix` is set to false.

    If covariates exists:
    - calculate px
    - computes Y hat and G hat by matrix multiplication.
    - substract Y hat and G hat from Y and G respectively.
    - filter maf if the threshold is greater than 0
    - standardizing phenotype matrix (Y) and genotype matrix (G)
    - calculate correlation (R) matrix.
    - computes log of odds (LOD) score matrix, or p-value if `lod_or_pval` is set to `lod`
    - calculate maximum LOD score if `export_matrix` is set to false.

# Arguments:
- `Y` : a matrix of phenotypes.
- `G` : a matrix of genotypes.
- `X` : a matrix of covariates. Default is `nothing`. If `nothing`, scan is run without covariates.
- `maf_threshold`: a floating point number to indicate the maf_threshold. Default is 0.05. Set to 0 if no maf filtering should be done.
- `export_matrix` : a boolean value that determines whether the result should be the maximum value of LOD score of each phenotype and its corresponding index, or the whole LOD score matrix.
- `usegpu` : a boolean value that indicates whether to run scan function on GPU or CPU. Default is false, which runs scan on CPU.
- `lod_or_pval`: a string value of either `lod` or `pval` to indicate the desired output.
- `timing_file`: a string that indicates the file location for the timing outputs. Default is nothing.
# Output:
returns LOD score or pval, in vector or matrix format depending on value of `export_matrix`.
"""
function scan(Y::AbstractArray{<:Real, 2}, G::AbstractArray{<:Real, 2},  X::Union{AbstractArray{<:Real, 2}, Nothing}=nothing; maf_threshold=0.05, export_matrix=false, usegpu=false, lod_or_pval="lod", timing_file="")
    if usegpu
        # TODO: return p-value functionality implementation in GPU.
        return LiteQTL.gpurun(Y, G, X, maf_threshold=maf_threshold, export_matrix=export_matrix, timing_file=timing_file)
    end

    return LiteQTL.cpurun(Y,G,X, maf_threshold=maf_threshold, export_matrix = export_matrix, lod_or_pval = lod_or_pval, timing_file=timing_file)
end


@kernel function _compute_vertically_integrated_lateral_areas!(∫ᶻ_A, grid)
    i, j, k = @index(Global, NTuple)

    @inbounds begin
        ∫ᶻ_A.xᶠᶜᶜ[i, j, 1] = 0
        ∫ᶻ_A.yᶜᶠᶜ[i, j, 1] = 0

        @unroll for k in 1:grid.Nz
            ∫ᶻ_A.xᶠᶜᶜ[i, j, 1] += Δyᶠᶜᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid)
            ∫ᶻ_A.yᶜᶠᶜ[i, j, 1] += Δxᶜᶠᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid)
        end
    end
end

function compute_vertically_integrated_lateral_areas!(∫ᶻ_A, grid, arch)

    event = launch!(arch,
                    grid,
                    :xyz,
                    _compute_vertically_integrated_lateral_areas!,
                    ∫ᶻ_A,
                    grid,
                    dependencies=Event(device(arch)))

    wait(device(arch), event)

    return
end


if VERSION <= v"1.1"
    const TH_EXP = 100
    const TH_INV = 1000
else
    const TH_EXP = 5000
    const TH_INV = 100_000
end

#========== exp! log! inv! ==========#

export exp0, exp_, exp!, exp!!
export log0, log_, log!, log!!
export inv0, inv_, inv!, inv!!

"""
    exp!(A)
    exp_(A) = exp!(similar(A), A)
    exp0(A) ≈ exp.(A)

Element-wise in-place exponential, and friends.
Multi-threaded when `length(A) >= $TH_EXP`.
Will be handled by `Yeppp` or `AppleAccelerate` or `IntelVectorMath` if you load one of them,
note that `Yeppp` may well be slower.
"""
function exp! end

exp0(A) = similar(A) .= exp.(A) # maps Adjoint -> Adjoint etc

@doc @doc(exp!)
exp_(A) = exp!(similar(A), A)
exp!(A) = exp!(A, A)
exp!!(A) = exp!(A) # differs in gradient

function exp!(B, A)
    @assert size(A)==size(B)
    if length(A) < TH_EXP
        for I in eachindex(A)
            @inbounds B[I] = exp1(A[I])
        end
    else
        Threads.@threads for I in eachindex(A)
            @inbounds B[I] = exp1(A[I])
        end
    end
    B
end

"""
    log!(A)
    log_(A) ≈ log!(similar(A), A)
    log0(A) = log.(A)

Element-wise in-place natural logarithm, and friends.
Multi-threaded when `length(A) >= $TH_EXP`.
Will be handled by `Yeppp` or `AppleAccelerate` or `IntelVectorMath` if you load one of them.
"""
function log! end

log0(A) = similar(A) .= log.(A)

@doc @doc(log!)
log_(A) = log!(similar(A), A)
log!(A) = log!(A, A)
log!!(A) = log!(A) # differs in gradient

function log!(B, A)
    @assert size(A)==size(B)
    if length(A) < TH_EXP
        for I in eachindex(A)
            @inbounds B[I] = log1(A[I])
        end
    else
        Threads.@threads for I in eachindex(A)
            @inbounds B[I] = log1(A[I])
        end
    end
    B
end

# These are a little faster than Julia's built-in functions?
exp1(x::Float64) = ccall(:exp, Cdouble, (Cdouble,), x)
exp1(x) = exp(x)
log1(x::Float64) = ccall(:log, Cdouble, (Cdouble,), x)
log1(x) = log(x)

# Versions which use cache
exp_(name::Symbol, A) = exp!(similar_(name, A), A)
log_(name::Symbol, A) = log!(similar_(name, A), A)

"""
    inv!(A) ≈ 1 ./ A
    inv!(A, b::Number) ≈ b ./ A

And `inv_(A)` which copies, and `inv0(A)` simple broadcasting.
Multi-threaded when `length(A) >= $TH_INV`.
Will be handled by `AppleAccelerate` if you load it.
"""
function inv! end

inv0(A::AbstractArray, b::Number=1) = similar(A) .= b ./ A # maps Adjoint -> Adjoint etc

@doc @doc(inv!)
inv_(A::AbstractArray, b::Number=1) = inv!(similar(A), A, b)
inv_(a::Number) = 1/a # for iscale_

inv!(A::AbstractArray, b::Number=1) = inv!(A, A, 1)
inv!(a::Number) = 1/a
function inv!(C::AbstractArray, A::AbstractArray, b::Number=1)
    @assert size(A)==size(C)
    if length(A) < TH_INV
        for I in eachindex(A)
            @inbounds C[I] = b / A[I]
        end
    else
        Threads.@threads for I in eachindex(A)
            @inbounds C[I] = b / A[I]
        end
    end
    C
end

inv_(name::Symbol, A::AbstractArray, b::Number=1) = inv!(similar_(name, A), A, b)
inv_(name::Symbol, a::Number=1) = 1/a # for iscale_

#========== scale! iscale! ==========#

export scale0, scale_, scale!, scale!!
export iscale0, iscale_, iscale!, iscale!!

using LinearAlgebra: Adjoint, Transpose

const ARVector = Union{Adjoint{<:Any, <:AbstractVector}, Transpose{<:Any, <:AbstractVector}}
const RVector = Union{Adjoint{<:Any, <:Vector}, Transpose{<:Any, <:Vector}}

"""
    scale!(A, b::Number) ≈ A .* b
    scale!(M, v::Vector) ≈ A .* v       # M::Matrix
    scale!(M, r::Adjoint) ≈ A .* r      # r::RowVector / Transpose etc.
    scale!(A, B) ≈ A .* B               # A,B same ndims

In-place scaling by a constant or (in the case of a matrix) by a row- or column-vector.
For each of these, there is also also `scale_(A, ...)` non-mutating but perhaps accellerated,
and `scale0(A, ...)` simple broadcasting.
"""
function scale! end

using LinearAlgebra

scale0(A::AbstractArray, b) = similar(A) .= A .* b

scale_(A::Array, b::Number) = rmul!(copy(A), b)
scale!(A::Array, b::Number) = rmul!(A, b)
scale!!(A::Array, b) = scale!(A,b) # differs in gradient

scale_(A::RVector, b::Number) = rmul!(copy(A), b) # scale_(::Abstract...) causes flux ambiguities
scale!(A::RVector, b::Number) = rmul!(A, b)
scale!!(A::RVector, b) = scale!(A,b)

@doc @doc(scale!)
scale_(A::Matrix, v::Vector) = lmul!(Diagonal(v), copy(A))
scale!(A::Matrix, v::Vector) = lmul!(Diagonal(v), A)

scale_(A::Matrix, r::RVector) = rmul!(copy(A), Diagonal(transpose(r)))
scale!(A::Matrix, r::RVector) = rmul!(A, Diagonal(transpose(r)))

# scale_(A::AbstractArray{T,N}, B::AbstractArray{T,N}) where {T,N} = similar(A) .= A .* B
# scale!(A::AbstractArray{T,N}, B::AbstractArray{T,N}) where {T,N} = A .= A .* B

function scale!(C::AbstractArray{T,N}, A::AbstractArray{TA,N}, B::AbstractArray{TB,N}) where {T,N, TA, TB}
    @assert size(A) == size(B) == size(C)
    for i in eachindex(A)
        @inbounds C[i] = A[i] * B[i]
    end
    C
end
scale_(A::AbstractArray{T,N}, B::AbstractArray{T,N}) where {T,N} = scale!(similar(A), A, B)
scale!(A::AbstractArray{T,N}, B::AbstractArray{T,N}) where {T,N} = scale!(A, A, B)

scale_(name::Symbol, A::Array, b::Number) = rmul!(copy_(name, A), b)
scale_(name::Symbol, A::Matrix, v::Vector) = lmul!(Diagonal(v), copy_(name, A))
scale_(name::Symbol, A::Matrix, r::RVector) = rmul!(copy_(name, A), Diagonal(transpose(r)))
scale_(name::Symbol, A::AbstractArray{T,N}, B::AbstractArray{T,N}) where {T,N} =
    scale!(similar_(name, A), A, B)

"""
    iscale!(A, b::Number) ≈ A ./ b
    iscale!(A, v::Vector) ≈ A ./ v      # A::Matrix
    iscale!(A, r::Adjoint) ≈ A ./ r     # r::RowVector / Transpose etc.
    iscale!(A, B) ≈ A ./ B

For each of these, there is also `iscale_(A, ...)` non-mutating but perhaps accellerated,
and `iscale0(A, ...)` simple broadcasting.
Finally there is `iscale!!(A, x)` which mutate both arguments, wihch may be a terrible idea.
"""
function iscale! end

iscale0(A::AbstractArray, b) = similar(A) .= A ./ b

@doc @doc(iscale!)
iscale_(A::AbstractArray, b) = scale_(A, inv_(b))
iscale!(A::AbstractArray, b) = scale!(A, inv_(b))
iscale!!(A::AbstractArray, b) = scale!(A, inv!(b))

function iscale!(C::Array{T,N}, A::Array{TA,N}, B::Array{TB,N}) where {T,N, TA, TB}
    @assert size(A) == size(B) == size(C)
    for i in eachindex(A)
        @inbounds C[i] = A[i] / B[i]
    end
    C
end
iscale_(A::Array{T,N}, B::Array{T,N}) where {T,N} = iscale!(similar(A), A, B)
iscale!(A::Array{T,N}, B::Array{T,N}) where {T,N} = iscale!(A, A, B)

# On square matrices this is a tie, but on (n,N) ./ N' it is faster
# The equivalent for iscale(Matrix, Vector) however is slower, wrong order
function iscale!(C::Matrix, A::Matrix, r::RVector)
    @assert size(A)==size(C)
    axes(A,2) == axes(r,2) || throw(DimensionMismatch("size disagreement in iscale?(Matrix,RowVector)"))
    @inbounds for j in axes(A,2)
        invr = inv(r[j])
        for i in axes(A,1)
            C[i,j] = A[i,j] * invr
        end
    end
    C
end
iscale_(A::Matrix, r::RVector) = iscale!(similar(A), A, r)
iscale!(A::Matrix, r::RVector) = iscale!(A, A, r)

iscale_(name::Symbol, A::AbstractArray, b) = scale_(name, A, inv_(name, b))
iscale_(name::Symbol, A::Matrix, r::RVector) = iscale!(similar_(name, A), A, r)
iscale_(name::Symbol, A::Array{T,N}, B::Array{T,N}) where {T,N} =
    iscale!(similar_(name, A), A, B)

#========== Accelerators ==========#

const CFloat = Union{Float64, Float32}
const CFloatArray{N} = Array{<:CFloat, N}
const CFloatMatrix = Matrix{<:CFloat}

IVERBOSE = false

VEC = ""
function load_note(str)
    global VEC
    if VEC == ""
        @info "ArrayAllez loaded code for $str"
        VEC = str
    else
        @warn "ArrayAllez loaded code for $str, perhaps overwriting $VEC"
        VEC *= " then $str"
    end
end

using Requires

@init @require Yeppp = "6310b701-1812-5374-a82f-9f6f2d54a40a" begin
    using .Yeppp

    exp!(B::CFloatArray, A::CFloatArray) = Yeppp.exp!(B, A)

    log!(B::CFloatArray, A::CFloatArray) = Yeppp.log!(B, A) # log_(A) calls log!(B,A)

    scale_(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} = Yeppp.multiply(A,B)
    scale!(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} = Yeppp.multiply!(A,A,B)

    IVERBOSE && load_note("Yeppp")
end

@init @require AppleAccelerate = "13e28ba4-7ad8-5781-acae-3021b1ed3924" begin
    using .AppleAccelerate

    exp!(B::CFloatArray, A::CFloatArray) = AppleAccelerate.exp!(B, A)

    log!(B::CFloatArray, A::CFloatArray) = AppleAccelerate.log!(B, A)

    inv!(A::CFloatArray) = AppleAccelerate.rec!(A, A)

    scale_(A::Vector{T}, B::Vector{T}) where {T<:CFloat} = AppleAccelerate.vmul(A,B)
    scale!(A::Vector{T}, B::Vector{T}) where {T<:CFloat} = AppleAccelerate.vmul!(A,A,B)

    scale_(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} =
        reshape(AppleAccelerate.vmul(vec(A),vec(B)), size(A))  # vmul is literally only vectors
    scale!(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} =
        begin AppleAccelerate.vmul!(vec(A),vec(A),vec(B)); A end

    iscale_(A::Vector{T}, B::Vector{T}) where {T<:CFloat} = AppleAccelerate.vdiv(A,B)
    iscale!(A::Vector{T}, B::Vector{T}) where {T<:CFloat} = AppleAccelerate.vdiv!(A,A,B)

    iscale_(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} =
        reshape(AppleAccelerate.vdiv(vec(A),vec(B)), size(A))
    iscale!(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} =
        begin AppleAccelerate.vdiv!(vec(A),vec(A),vec(B)); A end

    IVERBOSE && load_note("AppleAccelerate")
end

@init @require IntelVectorMath = "c8ce9da6-5d36-5c03-b118-5a70151be7bc" begin
    using .IntelVectorMath

    exp!(B::CFloatArray, A::CFloatArray) = IVM.exp!(B, A)

    log!(B::CFloatArray, A::CFloatArray) = IVM.log!(B, A) # log_(A) calls log!(B,A)

    iscale_(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} = IVM.divide(A,B)
    iscale!(A::Array{T,N}, B::Array{T,N}) where {T<:CFloat,N} = IVM.divide!(A,A,B)

    IVERBOSE && load_note("IntelVectorMath")
end

#========== The End ==========#


export
    Unconstrained,
    Polygonal,
    nextboundary

abstract type Domain end

"""
    Unconstrained

Signature of domains without boundaries.
"""
struct Unconstrained <: Domain end

"""
    Polygonal

Encapsulate a domain with polygonal boundaries. The boundaries are defined with
the normals to the facets and the intercepts of the corresponding hyperplanes.
"""
struct Polygonal <: Domain
    normals::Matrix{Real} # dim CxD where C=#constraints, D=#features
    intercepts::Vector{Real} # dim C
end

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

"""
    nextboundary(ud::Unconstrained, x, v)

Return (NaN, NaN) corresponding to the unconstrained case.
"""
function nextboundary(ud::Unconstrained,
                      x::Vector{<:Real}, v::Vector{<:Real})
    return (NaN, NaN)
end

"""
    nextboundary(pd::Polygonal, x, v)

Return the time of next boundary hit along the current trajectory when the
domain is polygonal and return the normal of the corresponding boundary.
The current point is `x` and the velocity `v`.
"""
function nextboundary(pd::Polygonal,
                      x::Vector{<:Real}, v::Vector{<:Real})
    # hitting time along trajectory (x+tv) for a boundary (normal,intercept) is
    # t = intercept/(normal dot v) - (x dot normal)/(normal dot v)
    nsv  = pd.normals * v
    hits = (pd.intercepts - pd.normals * x) ./ nsv
    # hard threshold times with nsv ~ 0 (near-parallel case),
    # remove negative times and times ~ 0 (for numerical stability)
    mask = map(|, abs.(nsv) .< 1e-10, hits .< 1e-10)
    for (i, e) ∈ enumerate(mask)
        e && (hits[i] = Inf)
    end
    # get time of hit + index of corresponding boundary
    (t_hit, j) = findmin(hits)
    # return time of hit + normal vector to boundary
    (t_hit, pd.normals[j, :])
end


"""
$(SIGNATURES)

Sample particle velocity from Maxwellian. Unconstraint methods use Box-Muller sampling,
while constraint one depends on Distributions.jl.
"""
function sample_maxwell(λ::Real, v = 0.0::Real)
    rd1 = rand()
    rd2 = rand()

    return sqrt(-log(rd1) / λ) * sin(2.0 * π * rd2) + v
end

"""
$(SIGNATURES)
"""
function sample_maxwell(prim::T) where {T<:AA{<:Real,1}}

    u = prim[2]
    if length(prim) == 3
        v = 0.0
        w = 0.0
    elseif length(prim) == 4
        v = prim[3]
        w = 0.0
    elseif length(prim) == 5
        v = prim[3]
        w = prim[4]
    end

    rd1 = rand(3)
    rd2 = rand(3)

    return @. sqrt(-log(rd1) / prim[end]) * sin(2.0 * π * rd2) + [u, v, w]

end

"""
$(SIGNATURES)

Truncated sampling with Distributions.jl
"""
function sample_maxwell(λ::Real, l, u)
    d = truncated(Normal(0.0, sqrt(1.0 / λ)), l, u)
    return rand(d)
end

"""
$(SIGNATURES)
"""
function sample_maxwell(λ::Real, v, l, u)
    d = truncated(Normal(0.0, sqrt(1.0 / λ)), l, u)
    return rand(d) + v
end

"""
$(SIGNATURES)
"""
function sample_maxwell(prim::T, l, u) where {T<:AA{<:Real,1}}
    u = prim[2]
    if length(prim) == 3
        v = 0.0
        w = 0.0
    elseif length(prim) == 4
        v = prim[3]
        w = 0.0
    elseif length(prim) == 5
        v = prim[3]
        w = prim[4]
    end

    d = truncated(Normal(0.0, sqrt(1.0 / prim[end])), l, u)
    return rand(d, 3) + [u, v, w]
end

"""
$(SIGNATURES)

Calculate the time interval until next collision
"""
next_collision_time(τ) = -τ * log(rand())

"""
$(SIGNATURES)

Calculate traveling time to edges
"""
function boundary_time(x, v::Real, xL, xR)
    if v < 0.0
        tb = (xL - x) / v
    elseif v > 0.0
        tb = (xR - x) / v
    else
        tb = 1e8
    end

    return tb
end

"""
$(SIGNATURES)
"""
function boundary_time(x, v::T, xL, xR) where {T<:AA{<:Real,1}}
    if v[1] < 0.0
        tb = (xL - x) / v[1]
    elseif v[1] > 0.0
        tb = (xR - x) / v[1]
    else
        tb = 1e8
    end

    return tb
end


module patch_test_2d_examples
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtools.MeshExportModule
import LinearAlgebra: cholesky

function q4_stress()
    println("Q4. Plane stress.")

    E = 1.0;
    nu = 1.0/3;
    alpha, beta, gamma, delta, eta, phi= 1.0/30, 1.0/34, -1.0/21, -1.0/51, -1.0/26, -1.0/35
    ux(x, y) = alpha + beta * x + gamma * y
    uy(x, y) = delta + eta * x + phi * y
    MR = DeforModelRed2DStress

    fens = FENodeSet([1.0 -0.3; 2.3 -0.3; 2.3 0.95; 1.0 0.95; 1.4 0.05; 1.9 -0.03; 1.7 0.5; 1.3 0.6])
    fes = FESetQ4([1 2 6 5; 6 2 3 7; 7 3 4 8; 8 4 1 5; 5 6 7 8])

    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz, 1), 2)) # displacement field

    # Apply prescribed displacements to exterior nodes
    for i in 1:4
    	setebc!(u, [i], 1, val=ux(fens.xyz[i, :]...))
    	setebc!(u, [i], 2, val=uy(fens.xyz[i, :]...))
    end

    applyebc!(u)
    numberdofs!(u)

    material = MatDeforElastIso(MR,  0.0, E, nu, 0.0)

    femm = FEMMDeforLinear(MR, IntegDomain(fes,  GaussRule(2, 2)),  material)

    F = nzebcloadsstiffness(femm,  geom,  u)
    K = stiffness(femm,  geom,  u)
    K = cholesky(K)
    U=  K\(F)
    scattersysvec!(u, U[:])

    for i in 5:8
    	uexact = [ux(fens.xyz[i, :]...), uy(fens.xyz[i, :]...)]
    	println("u.values[$i, :] = $(u.values[i, :]), uexact = [$(uexact)]")
    end

    File =  "a.vtk"
    vtkexportmesh(File,  fes.conn,  geom.values,
    FinEtools.MeshExportModule.Q4; vectors=[("u", u.values)])

    true

end # cookstress

function q4_stress_export()
    println("Q4. Plane stress.")

    E = 1.0;
    nu = 1.0/3;
    alpha, beta, gamma, delta, eta, phi= 1.0/30, 1.0/34, -1.0/21, -1.0/51, -1.0/26, -1.0/35
    ux(x, y) = alpha + beta * x + gamma * y
    uy(x, y) = delta + eta * x + phi * y
    MR = DeforModelRed2DStress

    fens = FENodeSet([1.0 -0.3; 2.3 -0.3; 2.3 0.95; 1.0 0.95; 1.4 0.05; 1.9 -0.03; 1.7 0.5; 1.3 0.6])
    fes = FESetQ4([1 2 6 5; 6 2 3 7; 7 3 4 8; 8 4 1 5; 5 6 7 8])

    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz, 1), 2)) # displacement field

    # Apply prescribed displacements to exterior nodes
    for i in 1:4
    	setebc!(u, [i], 1, val=ux(fens.xyz[i, :]...))
    	setebc!(u, [i], 2, val=uy(fens.xyz[i, :]...))
    end

    applyebc!(u)
    numberdofs!(u)

    material = MatDeforElastIso(MR,  0.0, E, nu, 0.0)

    femm = FEMMDeforLinear(MR, IntegDomain(fes,  GaussRule(2, 2)),  material)

    F = nzebcloadsstiffness(femm,  geom,  u)
    K = stiffness(femm,  geom,  u)
    K = cholesky(K)
    U=  K\(F)
    scattersysvec!(u, U[:])

    for i in 5:8
    	uexact = [ux(fens.xyz[i, :]...), uy(fens.xyz[i, :]...)]
    	println("u.values[$i, :] = $(u.values[i, :]), uexact = [$(uexact)]")
    end

    AE = AbaqusExporter("q4_stress_export");
    HEADING(AE, "q4_stress_export");
    COMMENT(AE, "");
    PART(AE, "part1");
    END_PART(AE);
    ASSEMBLY(AE, "ASSEM1");
    INSTANCE(AE, "INSTNC1", "PART1");
    NODE(AE, fens.xyz);
    ELEMENT(AE, "CPS4", "AllElements", connasarray(fes))
    NSET_NSET(AE, "clamped", 1:4)
    ORIENTATION(AE, "GlobalOrientation", vec([1. 0 0]), vec([0 1. 0]));
    SOLID_SECTION(AE, "elasticity", "GlobalOrientation", "AllElements", 1.0);
    END_INSTANCE(AE);
    END_ASSEMBLY(AE);
    MATERIAL(AE, "elasticity")
    ELASTIC(AE, E, nu)
    STEP_PERTURBATION_STATIC(AE)
    BOUNDARY(AE, "ASSEM1.INSTNC1", 1:4, fill(true, 4, 2), [[ux(fens.xyz[i, :]...) for i in 1:4] [uy(fens.xyz[i, :]...) for i in 1:4]])
    END_STEP(AE)
    close(AE)

    true

end # cookstress

function allrun()
    println("#####################################################")
    println("# q4_stress ")
    q4_stress()
    println("#####################################################")
    println("# q4_stress_export ")
    q4_stress_export()
    return true
end # function allrun

end # module patch_test_2d_examples


# Number differentiation

function gradient(f::Number)
  @inline function grad_f(x::Point)
    zero(return_type(outer,x,f))
  end
end

function ∇∇(f::Number)
  @inline function hess_f(x::Point)
    g = gradient(f)(x)
    gradient(g)(x)
  end
end

# Automatic differentiation of functions

function gradient(f::Function)
  @inline function _gradient(x)
    gradient(f,x)
  end
end

function symmetric_gradient(f::Function)
  @inline function _symmetric_gradient(x)
    symmetric_gradient(f,x)
  end
end

function divergence(f::Function)
  @inline function _divergence(x)
    divergence(f,x)
  end
end

function curl(f::Function)
  @inline function _curl(x)
    curl(f,x)
  end
end

function laplacian(f::Function)
  @inline function _laplacian(x)
    laplacian(f,x)
  end
end

# Specialization when x::Point

function gradient(f::Function,x::Point)
  gradient(f,x,return_value(f,x))
end

@inline function gradient(f::Function,x::Point,fx::Number)
  VectorValue(ForwardDiff.gradient(f,get_array(x)))
end

@inline function gradient(f::Function,x::Point,fx::VectorValue)
  TensorValue(transpose(ForwardDiff.jacobian(y->get_array(f(y)),get_array(x))))
end

function gradient(f::Function,x::Point,fx::MultiValue)
  @notimplemented
end

function symmetric_gradient(f::Function,x::Point)
  symmetric_part(gradient(f,x))
end

function divergence(f::Function,x::Point)
  divergence(f,x,return_value(f,x))
end

function divergence(f::Function,x::Point,fx)
  tr(gradient(f,x,fx))
end

function divergence(f::Function,x::Point,fx::TensorValue{2,2})
  g(x) = SVector(f(x).data)
  a = ForwardDiff.jacobian(g,get_array(x))
  VectorValue(
    a[1,1]+a[2,2],
    a[3,1]+a[4,2],
  )
end

function divergence(f::Function,x::Point,fx::TensorValue{3,3})
  g(x) = SVector(f(x).data)
  a = ForwardDiff.jacobian(g,get_array(x))
  VectorValue(
    a[1,1]+a[2,2]+a[3,3],
    a[4,1]+a[5,2]+a[6,3],
    a[7,1]+a[8,2]+a[9,3],
   )
end

function divergence(f::Function,x::Point,fx::SymTensorValue{2})
  g(x) = SVector(f(x).data)
  a = ForwardDiff.jacobian(g,get_array(x))
  VectorValue(
    a[1,1]+a[2,2],
    a[2,1]+a[3,2],
  )
end

function divergence(f::Function,x::Point,fx::SymTensorValue{3})
  g(x) = SVector(f(x).data)
  a = ForwardDiff.jacobian(g,get_array(x))
  VectorValue(
    a[1,1]+a[2,2]+a[3,3],
    a[2,1]+a[4,2]+a[5,3],
    a[3,1]+a[5,2]+a[6,3],
   )
end

function curl(f::Function,x::Point)
  grad2curl(gradient(f,x))
end

function laplacian(f::Function,x::Point)
  laplacian(f,x,return_value(f,x))
end

@inline function laplacian(f::Function,x::Point,fx::Number)
  tr(ForwardDiff.jacobian(y->ForwardDiff.gradient(f,y), get_array(x)))
end

@inline function laplacian(f::Function,x::Point,fx::VectorValue)
  A = length(x)
  B = length(fx)
  a = ForwardDiff.jacobian(y->transpose(ForwardDiff.jacobian(z->get_array(f(z)),y)), get_array(x))
  tr(ThirdOrderTensorValue{A,A,B}(Tuple(transpose(a))))
end

function laplacian(f::Function,x::Point,fx::MultiValue)
  @notimplemented
end

# Specialization when x::StaticVector

@inline gradient(f::Function,x::SVector) = gradient(f,Point(x))
@inline symmetric_gradient(f::Function,x::SVector) = symmetric_gradient(f,Point(x))
@inline divergence(f::Function,x::SVector) = divergence(f,Point(x))
@inline curl(f::Function,x::SVector) = curl(f,Point(x))
@inline laplacian(f::Function,x::SVector) = laplacian(f,Point(x))


using Test

using BEAST
using CompScienceMeshes

o, x, y = point(0,0,0), point(1,0,0), point(0,1,0)
G = Mesh([o,x,y,x+y],[index(1,2,3)])

@test numvertices(G) == 4
@test numcells(G) == 1

c = 1.0
K = BEAST.HH3DDoubleLayerTDBIO(speed_of_light=c)

X = lagrangecxd0(G)
@test numfunctions(refspace(X)) == 1

Y = duallagrangec0d1(G)
@test numfunctions(Y) == 1
@test numfunctions(refspace(Y)) == 3

# X = duallagrangecxd0(G, boundary(G))
# Y = lagrangec0d1(G, dirichlet=false)
# @test numfunctions(X) == 3
# @test numfunctions(Y) == 3

Δt = 20.0
Nt = 5
T = timebasiscxd0(Δt, Nt)
V = X ⊗ T
W = Y ⊗ T

fr2, store2 = BEAST.allocatestorage(K, V, V, Val{:densestorage}, BEAST.LongDelays{:ignore})
BEAST.assemble!(K, W, V, store2)
Z = fr2()
@test all(==(0), Z[:,:,1])

# import WiltonInts84
# trial_element = chart(G, first(cells(G)))
# x = neighborhood(trial_element,(1/3,1/3))
# x = neighborhood(trial_element,(0.0,0.0))
# r, R = 0.0, 20.0
# workspace = WiltonInts84.workspace(eltype(τ.vertices))
# G, vG, GG = WiltonInts84.wiltonints(
#     trial_element[1],
#     trial_element[2],
#     trial_element[3],
#     cartesian(x), r, R, Val{0},workspace)


include("../src/FiberNlse.jl")
using Plots;gr()
using FFTW;
using Peaks;
using SpecialFunctions;
U = FiberNlse; # Alias

begin

# CONSTANTS
    T = 50000*U.ps;
    Fs = 200*U.GHz;
# Simulation dimension
    Nₜ, Nₗ = (Int(round(Fs*T)),500);

# Fiber properties

    D = -7*U.ps/U.nm/U.km;
    α = 0*0.046/U.km;
    γ =  1.1/U.W/U.km;
    L = 4131;
    fib = U.Fiber(L, α, D, γ)

# Signal properties
λ1 = 1550.0*U.nm; # Wavelength
λ2 = λ1+1*U.pm; # Wavelength
λ0 = 0.5*(λ1+λ2);
Δλ = 0.5*abs(λ1-λ2)
sim, t = U.configure(Nₜ,Nₗ,D, γ, α, L, T, λ0);
Δw = 2pi*U.c*abs(λ1-λ2)/(λ1*λ2)
end

SP(ψ) = fftshift(abs.(fft(ψ)).^2)/sim.Nₜ^2
ν = fftshift(FFTW.fftfreq(sim.Nₜ, 1.0 / (sim.dt)));

function run(P)
    # Input construction
    Ψₒ = @. sqrt(0.5*P)*(exp.(1im*Δw.*t).+exp.(-1im*Δw.*t))
    U.inputSignal(sim,Ψₒ);
    U.simulate(sim, false);
    return (sim.Ψ[1,:],sim.Ψ[end,:])
end

ψi,ψf = run(0.5);
sp = SP(ψf)
idx = Peaks.argmaxima(sp)
idx = idx[findall(ν[idx] .< 4*Δw/(2pi))]
idx = idx[findall(ν[idx] .> 0)]
plot(SP(ψf))
scatter!(idx, sp[idx])

begin
Pl = (30:10:300)*U.mW
r = []
pf = []
for p in Pl
ψi,ψf = run(p);
sp = SP(ψf)
idx = Peaks.argmaxima(sp)
idx = idx[findall(ν[idx] .< 4*Δw/(2pi))]
idx = idx[findall(ν[idx] .> 0)]

if length(idx)>1
    append!(r, 10*log10(sp[idx[1]]/sp[idx[2]]))
    append!(pf, p)
end
end
end
pth = range(30, 300, 1000)*U.mW
a=0.82pi/maximum(pth)

rr(x) = @. (besselj0(0.5x)^2+besselj1(0.5x)^2)/(besselj1(0.5x)^2+besseljx(2,0.5x)^2)
scatter(pf,r)
plot!(pth, 10*log10.(rr.((a*pth))))

γexp = 0.5*a/L
δγ = abs(γexp-γ)/γ*100