using Flux
println("NNPDE_tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
using GalacticOptim
using Optim

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

## Example 1, 1D ode
@parameters t θ
@variables u(..)
@derivatives Dt'~t

# 1D ODE
eq = Dt(u(t,θ)) ~ t^3 + 2*t + (t^2)*((1+3*(t^2))/(1+t+(t^3))) - u(t,θ)*(t + ((1+3*(t^2))/(1+t+t^3)))

# Initial and boundary conditions
bcs = [u(0.,θ) ~ 1.0]

# Space and time domains
domains = [t ∈ IntervalDomain(0.0,1.0)]
# Discretization
dt = 0.1
# Neural network
chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,1))

discretization = NeuralPDE.PhysicsInformedNN(dt,
                                             chain,
                                             init_params = nothing;
                                             phi = nothing,
                                             autodiff=false,
                                             derivative = nothing,
                                             strategy = NeuralPDE.GridTraining())

pde_system = PDESystem(eq,bcs,domains,[t],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.1), progress = false; cb = cb, maxiters=1500)
phi = discretization.phi

analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2
ts = [domain.domain.lower:dt/10:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(t) for t in ts]
u_predict  = [first(phi(t,res.minimizer)) for t in ts]

@test u_predict ≈ u_real atol = 0.2

# using Plots
# t_plot = collect(ts)
# plot(t_plot ,u_real)
# plot!(t_plot ,u_predict)

## Example 2, 2D Poisson equation
@parameters x y θ
@variables u(..)
@derivatives Dxx''~x
@derivatives Dyy''~y

# 2D PDE
eq  = Dxx(u(x,y,θ)) + Dyy(u(x,y,θ)) ~ -sin(pi*x)*sin(pi*y)

# Initial and boundary conditions
bcs = [u(0,y,θ) ~ 0.f0, u(1,y,θ) ~ -sin(pi*1)*sin(pi*y),
       u(x,0,θ) ~ 0.f0, u(x,1,θ) ~ -sin(pi*x)*sin(pi*1)]
# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0),
           y ∈ IntervalDomain(0.0,1.0)]
# Discretization
dx = 0.05

# Neural network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))

discretization = NeuralPDE.PhysicsInformedNN(dx,
                                             chain,
                                             strategy = NeuralPDE.StochasticTraining(include_frac = 0.4))

pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.1), progress = false; cb = cb, maxiters=700)
phi = discretization.phi

xs,ys = [domain.domain.lower:dx/10:domain.domain.upper for domain in domains]
analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2)

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@test u_predict ≈ u_real atol = 3.0

# p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
# p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 3, 3D PDE
@parameters x y t θ
@variables u(..)
@derivatives Dxx''~x
@derivatives Dyy''~y
@derivatives Dt'~t

# 3D PDE
eq  = Dt(u(x,y,t,θ)) ~ Dxx(u(x,y,t,θ)) + Dyy(u(x,y,t,θ))
# Initial and boundary conditions
bcs = [u(x,y,0,θ) ~ exp(x+y)*cos(x+y) ,
      # u(x,y,2,θ) ~ exp(x+y)*cos(x+y+4*2) ,
      u(0,y,t,θ) ~ exp(y)*cos(y+4t),
      u(2,y,t,θ) ~ exp(2+y)*cos(2+y+4t) ,
      u(x,0,t,θ) ~ exp(x)*cos(x+4t),
      u(x,2,t,θ) ~ exp(x+2)*cos(x+2+4t)]
# Space and time domains
domains = [x ∈ IntervalDomain(0.0,2.0),
          y ∈ IntervalDomain(0.0,2.0),
          t ∈ IntervalDomain(0.0,2.0)]

# Discretization
dx = 0.25; dy= 0.25; dt = 0.25
# Neural network
chain = FastChain(FastDense(3,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))

discretization = NeuralPDE.PhysicsInformedNN([dx,dy,dt],
                                             chain,
                                             strategy = NeuralPDE.StochasticTraining(include_frac=0.9))
pde_system = PDESystem(eq,bcs,domains,[x,y,t],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.1), progress = false; cb = cb, maxiters=1000)
phi = discretization.phi

xs,ys,ts = [domain.domain.lower:dx:domain.domain.upper for (dx,domain) in zip([dx,dy,dt],domains)]
analytic_sol_func(x,y,t) = exp(x+y)*cos(x+y+4t)
u_real = [reshape([analytic_sol_func(x,y,t) for x in xs  for y in ys], (length(xs),length(ys)))  for t in ts ]
u_predict = [reshape([first(phi([x,y,t],res.minimizer)) for x in xs  for y in ys], (length(xs),length(ys)))  for t in ts ]
@test u_predict ≈ u_real atol = 200.0

# p1 =plot(xs, ys, u_predict, st=:surface);
# p2 = plot(xs, ys, u_real, st=:surface);
# plot(p1,p2)

## Example 4, high-order ode
@parameters x θ
@variables u(..)
@derivatives Dxxx'''~x
@derivatives Dx'~x

# ODE
eq = Dxxx(u(x,θ)) ~ cos(pi*x)

# Initial and boundary conditions
bcs = [u(0.,θ) ~ 0.0,
       u(1.,θ) ~ cos(pi),
       Dx(u(1.,θ)) ~ 1.0]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0)]

# Discretization
dx = 0.05

# Neural network
chain = FastChain(FastDense(1,8,Flux.σ),FastDense(8,1))

discretization = NeuralPDE.PhysicsInformedNN(dx,chain,strategy= NeuralPDE.StochasticTraining(include_frac=0.5))
pde_system = PDESystem(eq,bcs,domains,[x],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.01), progress = false; cb = cb, maxiters=2000)
phi = discretization.phi

analytic_sol_func(x) = (π*x*(-x+(π^2)*(2*x-3)+1)-sin(π*x))/(π^3)

xs = [domain.domain.lower:dx/10:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

@test u_predict ≈ u_real atol = 1.0

# x_plot = collect(xs)
# plot(x_plot ,u_real)
# plot!(x_plot ,u_predict)

## Example 5, system of pde
@parameters x, y, θ
@variables u1(..), u2(..)
@derivatives Dx'~x
@derivatives Dy'~y

# System of pde
eqs = [Dx(u1(x,y,θ)) + 4*Dy(u2(x,y,θ)) ~ 0,
      Dx(u2(x,y,θ)) + 9*Dy(u1(x,y,θ)) ~ 0,
      3*u1(x,0,θ) ~ 2*u2(x,0,θ)]

# Initial and boundary conditions
bcs = [[u1(x,0,θ) ~ 2x, u2(x,0,θ) ~ 3x]]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0), y ∈ IntervalDomain(0.0,1.0)]

# Discretization
dx = 0.1

# Neural network
chain = FastChain(FastDense(2,8,Flux.σ),FastDense(8,2))
discretization = NeuralPDE.PhysicsInformedNN(dx,chain)
pde_system = PDESystem(eqs,bcs,domains,[x,y],[u1,u2])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,Optim.BFGS(); cb = cb, maxiters=300)
phi = discretization.phi

analytic_sol_func(x,y) =[1/3*(6x - y), 1/2*(6x - y)]
xs,ys = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
u_real  = [[analytic_sol_func(x,y)[i] for x in xs  for y in ys] for i in 1:2]
u_predict  = [[phi([x,y],res.minimizer)[i] for x in xs  for y in ys] for i in 1:2]

@test u_predict ≈ u_real atol = 10.0

# p1 =plot(xs, ys, u_predict, st=:surface);
# p2 = plot(xs, ys, u_real, st=:surface);
# plot(p1,p2)

## Example 6, 2d wave equation, neumann boundary condition
#here we use low level api for build solution
@parameters x, t, θ
@variables u(..)
@derivatives Dxx''~x
@derivatives Dtt''~t
@derivatives Dt'~t

#2D PDE
C=1
eq  = Dtt(u(x,t,θ)) ~ C^2*Dxx(u(x,t,θ))

# Initial and boundary conditions
bcs = [u(0,t,θ) ~ 0.,# for all t > 0
       u(1,t,θ) ~ 0.,# for all t > 0
       u(x,0,θ) ~ x*(1. - x), #for all 0 < x < 1
       Dt(u(x,0,θ)) ~ 0. ] #for all  0 < x < 1]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0),
           t ∈ IntervalDomain(0.0,1.0)]
# Discretization
dx = 0.1

# Neural network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))
strategy=NeuralPDE.GridTraining()
discretization = NeuralPDE.PhysicsInformedNN(dx,chain,strategy=strategy)

phi = discretization.phi
derivative = discretization.derivative
initθ = discretization.initθ

indvars = [x,t]
depvars = [u]
dim = length(domains)

expr_pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars)
expr_bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars) for bc in bcs]

train_sets = NeuralPDE.generate_training_sets(domains,dx,bcs,indvars,depvars)

train_domain_set,train_bound_set,train_set= train_sets

pde_loss_function = NeuralPDE.get_loss_function(eval(expr_pde_loss_function),
                                      train_domain_set,
                                      phi,
                                      derivative,
                                      strategy)
bc_loss_function = NeuralPDE.get_loss_function(eval.(expr_bc_loss_functions),
                                     train_bound_set,
                                     phi,
                                     derivative,
                                     strategy)

function loss_function(θ,p)
    return pde_loss_function(θ) + bc_loss_function(θ)
end

f = OptimizationFunction(loss_function, initθ, GalacticOptim.AutoZygote())
prob = GalacticOptim.OptimizationProblem(f, initθ)

res = GalacticOptim.solve(prob,Optim.BFGS(); cb = cb, maxiters=400)

xs,ts = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
analytic_sol_func(x,t) =  sum([(8/(k^3*pi^3)) * sin(k*pi*x)*cos(C*k*pi*t) for k in 1:2:50000])

u_predict = reshape([first(phi([x,t],res.minimizer)) for x in xs for t in ts],(length(xs),length(ts)))
u_real = reshape([analytic_sol_func(x,t) for x in xs for t in ts], (length(xs),length(ts)))

@test u_predict ≈ u_real atol = 10.0

# diff_u = abs.(u_predict .- u_real)
# p1 = plot(xs, ts, u_real, linetype=:contourf,title = "analytic");
# p2 =plot(xs, ts, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ts, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)


using Statistics
using NCDatasets

using ArgParse
using JSON

include("LinearRegression.jl")

correlation = (x1, x2) -> x1' * x2 / (sum(x1.^2)*sum(x2.^2)).^0.5

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file"
            help = "Ocean data file. New variable will be appended."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

        "--SST"
            help = "Variable name of SST."
            arg_type = String
            required = true

        "--beg-year"
            help = "Year of begin."
            arg_type = Int64
            required = true

        "--end-year"
            help = "Year of end."
            arg_type = Int64
            required = true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

Dataset(parsed["data-file"], "r") do ds

    dims = size(ds[parsed["SST"]])

    beg_t = (parsed["beg-year"] - 1) * 12 + 1
    end_t = (parsed["end-year"] - 1) * 12 + 12

    if length(dims) == 3 # 2D + time
        rng = (:,:,beg_t:end_t)
    elseif length(dims) == 4 # 2D + time
        rng = (:,:,1,beg_t:end_t)
    else
        ErrorException("Unknown dimension") |> throw
    end

    global SST  = replace(ds[parsed["SST"]][rng...], missing=>NaN)

    global (Nx, Ny, Nt) = size(SST)

    if mod(Nt, 12) != 0

        SST = SST[:, :, 1:12*floor(Integer, Nt/12.0)]
        (Nx, Ny, Nt) = size(SST)

        #ErrorException("Time record is not multiple of 12") |> throw
    end

    global nyears = Int64(Nt / 12)
end

Dataset(parsed["domain-file"], "r") do ds
    global mask = replace(ds["mask"], missing=>NaN)
end

# MM   = Monthly Mean
# A    = Anomaly
# YYC  = Year-to-Year Correlation

SSTMM   = zeros(Float64, size(SST)[1:2]..., 12)
SSTA    = zeros(Float64, size(SST)...)
SSTAYYC = zeros(Float64, size(SST)[1:2]..., 12)
SSTAVAR = zeros(Float64, size(SST)[1:2]..., 12)

N = size(SST)[3]
x = collect(Float64, 1:N)
for i=1:Nx, j=1:Ny

    d = detrend(x, view(SST, i, j, :))

    if mask[i, j] != 0
        SSTMM[i, j, :]  .= NaN
        SSTA[i, j, :]   .= NaN
        SSTAYYC[i, j, :] .= NaN
        continue
    end

    SSTMM[i, j, :] = mean( reshape(d, 12, :), dims=2 )[:, 1]
    SSTA[i, j, :]  = d - repeat( SSTMM[i, j, :], outer=nyears)

    for m = 1:12
        d_yy = view(SSTA, i, j, m:12:(m+nyears*12-1))
        SSTAYYC[i, j, m] = correlation(d_yy[1:end-1], d_yy[2:end])
        SSTAVAR[i, j, m] = std(d_yy)
    end

end

Dataset(parsed["data-file"], "a") do ds
    if ! haskey(ds.dim, "months")
        defDim(ds, "months", 12)
    end
end

Dataset(parsed["data-file"], "a") do ds
    for (varname, vardata, vardim, attrib) in [
        ("SSTMM",  SSTMM, ("Nx", "Ny", "months"), Dict()),
        ("SSTA",    SSTA, ("Nx", "Ny", "time"),   Dict()),
        ("SSTAYYC",  SSTAYYC, ("Nx", "Ny", "months"), Dict()),
        ("SSTAVAR",  SSTAVAR, ("Nx", "Ny", "months"), Dict()),
    ]

        if ! haskey(ds, varname)
            var = defVar(ds, varname, Float64, vardim)
            var.attrib["_FillValue"] = 1e20
        end

        var = ds[varname]

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end
end


"""
    create_RTLinvR(cache::BasicILMCache[;scale=1.0])

Using the provided cache `cache`, construct the square matrix ``-R^T L^{-1}R``, which maps data of type `ScalarData`
to data of the same type. The operators `R^T` and `R` correspond to [`interpolate!`](@ref) and
[`regularize!`](@ref) and `L^{-1}` is the inverse of the grid Laplacian. The optional keyword `scale` multiplies the
matrix by the designated value.
"""
function create_RTLinvR(cache::BasicILMCache{N};scale=1.0) where {N}
    @unpack L, sdata_cache, gdata_cache = cache

    len = length(sdata_cache)
    A = Matrix{eltype(sdata_cache)}(undef,len,len)
    fill!(sdata_cache,0.0)

    for col in 1:len
        sdata_cache[col] = 1.0
        fill!(gdata_cache,0.0)

        regularize!(gdata_cache,sdata_cache,cache)
        inverse_laplacian!(gdata_cache,cache)
        interpolate!(sdata_cache,gdata_cache,cache)

        A[:,col] = -scale*sdata_cache
        fill!(sdata_cache,0.0)
    end

    return A

end

"""
    create_CLinvCT(cache::BasicILMCache[;scale=1.0])

Using the provided cache `cache`, construct the square matrix ``-C_s L^{-1}C_s^T``, which maps
data of the primary point data type of the cache to data of the same type. The operators `C_s` and `C_s^T` correspond to [`surface_curl!`](@ref)
and `L` is the grid Laplacian. The optional keyword `scale` multiplies the
matrix by the designated value.
"""
function create_CLinvCT(cache::BasicILMCache{N};scale=1.0) where {N}
    @unpack L, sdata_cache, gcurl_cache = cache

    len = length(sdata_cache)
    A = Matrix{eltype(sdata_cache)}(undef,len,len)
    fill!(sdata_cache,0.0)

    for col in 1:len
        sdata_cache[col] = 1.0
        fill!(gcurl_cache,0.0)
        surface_curl!(gcurl_cache,sdata_cache,cache)

        inverse_laplacian!(gcurl_cache,cache)
        surface_curl!(sdata_cache,gcurl_cache,cache)

        A[:,col] = -scale*sdata_cache
        fill!(sdata_cache,0.0)
    end

    return A

end

"""
    create_CL2invCT(cache::BasicILMCache[;scale=1.0])

Using the provided cache `cache`, construct the square matrix ``-C_s L^{-2} C_s^T``, which maps
data of the primary point data type of the cache to data of the same type. The operators `C_s` and `C_s^T` correspond to [`surface_curl!`](@ref)
and `L` is the grid Laplacian. The optional keyword `scale` multiplies the
matrix by the designated value.
"""
function create_CL2invCT(cache::BasicILMCache{N};scale=1.0) where {N}
    @unpack L, sdata_cache, gdata_cache, gcurl_cache = cache

    len = length(sdata_cache)
    A = Matrix{eltype(sdata_cache)}(undef,len,len)
    fill!(sdata_cache,0.0)

    for col in 1:len
        sdata_cache[col] = 1.0
        fill!(gdata_cache,0.0)
        fill!(gcurl_cache,0.0)

        surface_curl!(gcurl_cache,sdata_cache,cache)
        inverse_laplacian!(gcurl_cache,cache)
        inverse_laplacian!(gcurl_cache,cache)
        surface_curl!(sdata_cache,gcurl_cache,cache)

        A[:,col] = -scale*sdata_cache
        fill!(sdata_cache,0.0)
    end

    return A

end

"""
    create_GLinvD(cache::BasicILMCache[;scale=1.0])

Using the provided cache `cache`, construct the square matrix ``-G_s L^{-1}D_s``, which maps data of type `ScalarData`
to data of the same type. The operators `G_s` and `D_s` correspond to [`surface_grad!`](@ref)
and [`surface_divergence!`](@ref), and `L` is the grid Laplacian. The optional keyword `scale` multiplies the
matrix by the designated value.
"""
function create_GLinvD(cache::BasicILMCache{N};scale=1.0) where {N}
    @unpack L, sdata_cache, gdata_cache = cache

    len = length(sdata_cache)
    A = Matrix{eltype(sdata_cache)}(undef,len,len)
    fill!(sdata_cache,0.0)

    for col in 1:len
        sdata_cache[col] = 1.0
        fill!(gdata_cache,0.0)
        surface_divergence!(gdata_cache,sdata_cache,cache)
        inverse_laplacian!(gdata_cache,cache)
        surface_grad!(sdata_cache,gdata_cache,cache)

        A[:,col] .= -scale*sdata_cache
        fill!(sdata_cache,0.0)
    end

    return A

end

"""
    create_GLinvD_cross(cache::BasicILMCache[;scale=1.0])

Using the provided cache `cache`, construct the square matrix ``-\\hat{G}_s L^{-1}\\hat{D}_s``, which maps data of type `ScalarData`
to data of the same type. The operators `G_s` and `D_s` correspond to [`surface_grad_cross!`](@ref)
and [`surface_divergence_cross!`](@ref), and `L` is the grid Laplacian. The optional keyword `scale` multiplies the
matrix by the designated value.
"""
function create_GLinvD_cross(cache::BasicILMCache{N};scale=1.0) where {N}
    @unpack L, sdata_cache, gdata_cache = cache

    len = length(sdata_cache)
    A = Matrix{eltype(sdata_cache)}(undef,len,len)
    fill!(sdata_cache,0.0)

    for col in 1:len
        sdata_cache[col] = 1.0
        fill!(gdata_cache,0.0)
        surface_divergence_cross!(gdata_cache,sdata_cache,cache)
        inverse_laplacian!(gdata_cache,cache)
        surface_grad_cross!(sdata_cache,gdata_cache,cache)

        A[:,col] .= -scale*sdata_cache
        fill!(sdata_cache,0.0)
    end

    return A

end

"""
    create_nRTRn(cache::BasicILMCache[;scale=1.0])

Using the provided cache `cache`, construct the square matrix ``n\\cdot R_f^T R_f n \\circ``, which maps data of type `ScalarData`
to data of the same type. The operators `R_f^T` and `R_f` correspond to the interpolation
and regularization matrices. The optional keyword `scale` multiplies the
matrix by the designated value.
"""
function create_nRTRn(cache::BasicILMCache{N};scale=1.0) where {N}
    @unpack sdata_cache, gsnorm_cache = cache

    len = length(sdata_cache)
    A = Matrix{eltype(sdata_cache)}(undef,len,len)
    fill!(sdata_cache,0.0)

    for col in 1:len
        sdata_cache[col] = 1.0
        fill!(gsnorm_cache,0.0)
        regularize_normal!(gsnorm_cache,sdata_cache,cache)
        normal_interpolate!(sdata_cache,gsnorm_cache,cache)

        A[:,col] .= scale*sdata_cache
        fill!(sdata_cache,0.0)
    end

    return A

end

"""
    create_surface_filter(cache::BasicILMCache)

Create a surface filtering matrix operator ``\\tilde{R}^T R``, where ``\\tilde{R}^T``
represents a modified version of the
interpolation operator. The resulting matrix can be applied
to surface data to filter out high-frequency components.
"""
function create_surface_filter(cache::BasicILMCache{N,SCA}) where {N,SCA}
    @unpack R, gdata_cache, sdata_cache, g = cache
    regfilt = _get_regularization(points(cache),areas(cache),g,
                                  _ddf_type(cache),SCA,filter=true)
    Ef = InterpolationMatrix(regfilt,gdata_cache,sdata_cache)

    len = length(sdata_cache)
    C = Matrix{eltype(sdata_cache)}(undef,len,len)
    return mul!(C,Ef,R)

end


# day 01 ---------------------------------------------
input_path = joinpath(pwd(), "inputs", "input_01.txt")
x = parse.(Int, readlines(input_path))

# 01a ------------------------------------------------
sum(x[i] > x[i-1] for i in 2:lastindex(x))
# or
sum(diff(x) .> 0)

# 01b ------------------------------------------------
sum(4:lastindex(x)) do i
   sum(x[(i-3):(i-1)]) < sum(x[i-2:i])
end

# alternative solution
# sum([b, c, d]) > sum([a, b, c])
# d + sum([b, c]) > a + sum([b, c])
# d > a
sum(4:lastindex(x)) do i
   x[i] > x[(i-3)]
end


## In LinearAlgebra/src/qr.jl
## the QR factorization stores
## * Q in a compressed form using Householder transformations
## * R as an upper triangular matrix
##
## QRFactorization stores
## * Q in a compressed form using Givens rotations, either descending or twisted
## * R in a compressed form (RankOne, Pencil, or Identity) or as a full matrix that is upper triangular
## * this factorization gets updated by the AMRVW_algorithm!, and hence `eigvals`
##
abstract type AbstractQRRotatorFactorization  end

##  We  have  two types, though could consolidate
##
struct QRFactorization{QF, RF} <: AbstractQRRotatorFactorization
  QF::QF
  RF::RF
  end

# use identity R when not specified
function QRFactorization(QF::AbstractQFactorization{T,S}) where {T,S}
    RF = RFactorizationIdentity{T,S}()
    QRFactorization(QF, RF)
end

function Base.size(state::AbstractQRRotatorFactorization)
    m1, n1 = size(state.QF)
    m2, n2 = size(state.RF)
    return (max(m1,m2), max(n1, n2))
end

Base.copy(F::QRFactorization) = QRFactorization(copy(F.QF), copy(F.RF))
Base.length(state::QRFactorization) = length(state.QF)+1
Base.eltype(state::QRFactorization) = eltype(state.QF)

# return A = QR
function Base.Matrix(state::QRFactorization)

    Q = Matrix(state.QF)
    R = Matrix(state.RF)

    ## May have Q smaller than  R, so will pad in that case
    if R != I
        m,n =  size(Q)[1], size(R)[1]
        n  < m && error("R is  too  small")
        return Q * R[1:m, 1:m]
    else
        return Q
    end

end

"""

    qr_factorization(H::Matrix; unitary=false)

For a Hessenberg matrix `H` return a factorization,`Q⋅R`, where `Q` is a `QFactorization` object of descending rotators and `R` is either a `AMRVW.RFactorizationUnitaryDiagonal` object (if `unitary=true`) or a `RFactorizationUpperTriangular` object.

`H` may be a full matrix or the `.H` compontent  of a `hessenberg` factorization.

"""
function qr_factorization(H::AbstractMatrix; unitary=false)
    R = copy(H)
    qr_factorization!(R; unitary=unitary)
end

function qr_factorization!(R::AbstractMatrix{S}; unitary=false) where {S}

    n = size(R, 1)
    Qs = Vector{Rotator{real(S),S}}(undef, n-1)
    for i in 1:n-1
        c,s,r = givensrot(R[i,i], R[i+1,i])
        Ui =  Rotator(c,s,i)
        Qs[i] = Ui'
        j = i + 1
        R[i,i] = r
        R[j,i] = zero(S)
        for k in j:n
            rik, rjk =  R[i,k],  R[j,k]
            R[i,k] = c * rik + s * rjk
            R[j,k] = -conj(s) * rik + conj(c) * rjk
        end
    end

    _qr_factorization(DescendingChain(Qs), R, Val(unitary))
end

function _qr_factorization(Des::DescendingChain{T,S,V}, R, unitary::Val{true}) where {T, S, V}

    QF = QFactorization(Des)
    RF = RFactorizationUnitaryDiagonal(sign.(diag(R)))

    QRFactorization(QF, RF)
end

function _qr_factorization(Des, R, unitary::Val{false})

    QF = QFactorization(Des)
    RF = RFactorizationUpperTriangular(R)

    F = QRFactorization(QF, RF)
end


function getrk!(tempk,press,rk)

#=**************************************************************************

 This is a function in an ozone NOx formaldehyde mechanism, adapted from Fortran
 by Obin Sturm (UC Davis) in October 2019.
 The original comments and file description are below.

 created by: Mike Kleeman
             ECI 298 Air Quality Modeling
             April 1999

 The purpose of this subroutine is to calculate the reaction rate constants
 for the ozone NOx formaldehyde system.  The order of equations is assumed
 to be:
         1) NO2 + HV = NO + O
         2) O + O2 = O3
         3) O3 + NO = NO2 + O2
         4) HCHO + HV = 2 HO2. + CO
         5) HCHO + HV = H2 + CO
         6) HCHO + HO. = HO2. + CO + H2O
         7) HO2. + NO = HO. + NO2
         8) HO. + NO2 = HNO3
         9) HO2H + HV = 2 HO.
        10) HO2H + HO. = HO2. + H2O

 It is assumed that the photolytic constants are set at a constant value
 appropriate for mid-day in Southern California.

 Inputs:
  tempk - atmospheric temperature in degrees Kelvin
  press - atmospheric pressure in atmospheres

 Outputs:
  rk    - reaction rate constants in ppm min units

****************************************************************************=#

# function definition to convert from cm3 molec-1 s-1 to ppm min-1
    rkppmmin(rk1,press,tempk) = rk1 * 60 *  7.34e+15 * press / tempk

# begin by setting the rate constants for the photolytic reactions
    rk[1] = 0.5
    rk[4] = 0.015
    rk[5] = 0.022
    rk[9] = 0.0003

# reaction 2
    rk[2] = 6.0e-34 * ( tempk/300. )^(-2.3)
    rm = 7.34e+21 * press / tempk
    rk[2] = rk[2] * rm
    rk[2] = rkppmmin(rk[2],press,tempk)

# reaction 3
    rk[3] = 2.e-12 * exp(-1400/tempk)
    rk[3] = rkppmmin(rk[3],press,tempk)

# reaction 6
    rk[6] = 1.2e-14 * tempk^1 * exp(+287/tempk)
    rk[6] = rkppmmin(rk[6],press,tempk)

# reaction 7
    rk[7] = 3.7e-12 * exp(+250/tempk)
    rk[7] = rkppmmin(rk[7],press,tempk)

# reaction 8
# note: the following 2 lines are adapted from comments from the o.g. Fortran module
#      rk0 = 2.6e-30 * ( tempk/300 )^-3.2
#      rki = 2.4e-11 * ( tempk/300 )^-1.3
    rk0 = 2.0e-30 * ( tempk/300 )^(-3.0)
    rki = 2.5e-11 * ( tempk/300 )^0
    rm = 7.34e+21 * press / tempk
    rb = 0.6^(1.0 / ( 1.0 + (log10(rk0*rm/rki))^2.0 ))
    rk[8] = rk0 * rm * rb / (1.0 + rk0*rm/rki)
    rk[8] = rkppmmin(rk[8],press,tempk)

# reaction 10
    rk[10] = 3.3e-12 * exp(-200/tempk)
    rk[10] = rkppmmin(rk[10],press,tempk)

# return to the calling subroutine
    return rk
end


#
# Use this script to plot the evolution of the boundary layer depth
#

## Plotting mixed layer depth
using Plots, PyPlot

using OceanTurbulenceParameterEstimation
using Oceananigans.Fields: interior

# les analysis
# derivative function
function δ(z, Φ)
    m = length(Φ)-1
    Φz = ones(m)
    for i in 1:m
        Φz[i] = (Φ[i+1]-Φ[i])/(z[i+1]-z[i])
    end
    return Φz
end

"""
Approximates the mixed layer depth for pure convection.
"""
function approximate_mixed_layer_depth(model_time_series, data::OneDimensionalTimeSeries, targets)

        data.constants[:αg]
        Qᵇ = data.boundary_conditions.Qᵇ
        N² = data.boundary_conditions.dbdz_bottom

        Nz = data.grid.Nz
        z = data.grid.zC
        Lz = data.grid.Lz
        Nt = length(data.t)

        # For the LES solution
        h2_les = randn(Nt)
        for i in targets

            b = interior(data.b[i]) # remove halos

            Bz = δ([z...], [b...])
            mBz = maximum(Bz)
            tt = (2*N² + mBz)/3
            bools = Bz .> tt
            zA = (z[1:(Nz-1)] .+ z[2:Nz] ) ./ 2

            h2_les[i] = any(bools) ? -minimum(zA[bools]) : Lz
        end

        # For the model solution
        h2_model = randn(Nt)
        for i in 1:Nt-1
            B = interior(model_time_series.b[i])
            Bz = δ([z...], [B...])
            mBz = maximum(Bz)
            tt = (2*N² + mBz)/3
            bools = Bz .> tt
            zA = (z[1:(Nz-1)] + z[2:Nz] )./2

            h2_model[i] = any(bools) ? -minimum(zA[bools]) : Lz # mixed layer reached the bottom
        end

    return [h2_les, h2_model]
end

LESdata = GeneralStrat

ParametersToOptimize = CATKEParametersRiDependent
RelevantParameters = ParametersToOptimize

params = Parameters(RelevantParameters = RelevantParameters,
               ParametersToOptimize = ParametersToOptimize)

function make_all_the_plots(params)

    directory = pwd() * "/Results/plotGeneralStrat_default_parameters/$(RelevantParameters)/"
    isdir(directory) || mkpath(directory)

    for (i, LEScase) in enumerate(values(LESdata))

        td = OneDimensionalTimeSeries(LEScase.filename; grid_type=OneDimensionalEnsembleGrid, Nz=32)

        relative_weights = Dict(:b => 1.0, :u => 1.0, :v => 1.0, :e => 1.0)

        # Single simulation
        case_loss, default_parameters = get_loss(LEScase, td, params, relative_weights; Δt=10.0, fields=relevant_fields(LEScase),
                                                parameter_specific_kwargs[params.RelevantParameters]...)

        targets = case_loss.loss.targets

        ℱ = model_time_series(case_loss.simulation, case_loss.data, default_parameters)
        h2_les, h2_model = approximate_mixed_layer_depth(ℱ, td, targets)

        days = td.t[targets] ./ 86400
        first = LEScase.first
        last = isnothing(LEScase.last) ? length(td.t) : LEScase.last

        Plots.plot(days, h2_les, label = "LES mixed layer depth", linewidth = 3 , color = :red, grid = true, gridstyle = :dash, gridalpha = 0.25, framestyle = :box , title=td.name, ylabel = "Depth [meters]", xlabel = "days", legend = :topleft)
        p = plot!(days, h2_model, color = :purple, label = "TKE mixed layer depth", linewidth = 3 )
        Plots.savefig(p, directory*td.name*"_mixed_layer_depth.pdf")

        Plots.plot(days, h2_les, label = "LES mixed layer depth", linewidth = 3 , color = :red, grid = true, gridstyle = :dash, gridalpha = 0.25, framestyle = :box , title=td.name, ylabel = "Depth [meters]", xlabel = "days", yscale = :log10, xscale = :log10, legend = :topleft)
        p = plot!(days, h2_model, color = :purple, label = "TKE mixed layer depth", linewidth = 3 )
        Plots.savefig(p, directory*td.name*"_mixed_layer_depth_log10.pdf")

        # p = visualize!(md, td, 1:180:length(td), best_parameters)
        # PyPlot.savefig(directory*td.name*".png")
    end
end

make_all_the_plots(params)


module SimplexQuad

using LinearAlgebra

export simplexquad

"""
    simplexquad([::Type{T},] N::Int, n::Int)
    simplexquad(N::Int, vert::AbstractMatrix)

Construct Gauss points and weights for an `n`-dimensional simplex
domain with vertices specified by the `n*(n-1)` matrix `vert`, where
each row contains the coordinates `(x1,...,xn)` for a vertex. The
order of the quadrature scheme in each dimension must be the same in
this implementation.

# Arguments
- `N`: Number of quadrature points
- `T`: Return type (defaults to Float64)
- `n`: Number of dimensions
- `vert`: matrix of simplex vertex coordinates

# Result
- `X`: Matrix of quadrature points, where each row contains the
  coordinates for a point
- `W`: Array of quadrature weights

# Sample Usage:

    X, W = simplexquad(n, vert)     # Specify the vertices
    X, W = simplexquad(T, n, dim)   # Specify the dimension and use unit simplex

`X` will be a matrix for which the `j`th column are the grid points in
each coordinate `xj`.

To integrate a function `f`, use e.g.

    sum(W[i] * f(X[i,:]) for i in 1:length(W))

Note: The standard `n`-dimensional simplex has vertices specified
`vert = diagm(n+1, n, ones(n))
"""
simplexquad

function simplexquad(::Type{T}, N::Int, n::Int) where {T}
    m = n + 1
    vert = diagm(m, n, ones(T, n))
    return simplexquad(N, vert)
end
simplexquad(N::Int, n::Int) = simplexquad(float(Int), N, n)

function simplexquad(N::Int, vert::AbstractMatrix{T}) where {T}
    N <= 0 && @error "Number of Gauss points must be a natural number"

    m, n = size(vert)
    m != n + 1 &&
        @error "The matrix of vertices must have n+1 rows and n columns"

    Nn = N^n
    if n == 1
        # The 1-D simplex is only an interval
        q, w = rquad(T, N, 0)
        len = diff(vert; dims = 1)
        X = vert[1] .+ len[1] * q
        W = abs.(len[1]) * w
        X = reshape(X, :, 1)
    else
        # Find quadrature rules for higher dimensional domains
        q = Array{Vector}(undef, n)
        w = Array{Vector}(undef, n)
        for k in 1:n
            q[k], w[k] = rquad(T, N, n - k)
        end
        Q = ndgrid(q...)
        q = reshape(cat(Q...; dims = n), Nn, n)
        W = ndgrid(w...)
        w = reshape(cat(W...; dims = n), Nn, n)
        map = zeros(T, m, m) + I
        map[2:m, 1] .= -1
        c = map * vert
        W = abs(det(c[2:m, :])) * prod(w; dims = 2)
        qp = cumprod(q; dims = 2)
        e = ones(T, Nn, 1)
        X = [e [(1 .- q[:, 1:(n - 1)]) e] .* [e qp[:, 1:(n - 2)] qp[:, n]]] * c
        @assert size(W, 2) == 1
        W = reshape(W, :)
    end

    X::Matrix{T}
    @assert size(X, 2) == n
    W::Vector{T}
    @assert length(W) == size(X, 1)
    return X, W
end

function rquad(::Type{T}, N::Int, k::Int) where {T}
    k1 = k + 1
    k2 = k + 2

    n = collect(1:N)'
    nnk = 2 * n .+ k
    A = [T(k) / k2 repeat([T(k^2)], 1, N) ./ (nnk .* (nnk .+ 2))]

    n = collect(2:N)'
    nnk = nnk[n]
    B1 = T(4) * k1 / (k2 * k2 * (k + 3))

    nk = n .+ k
    nnk2 = nnk .* nnk
    B = T(4) * (n .* nk) .^ 2 ./ (nnk2 .* nnk2 - nnk2)

    ab = [A' [T(2^k1) / k1; B1; B']]
    s = sqrt.(ab[2:N, 2])
    X, V = eigen(SymTridiagonal(ab[1:N, 1], s))
    I = sortperm(X)
    X = X[I]
    x = (X .+ 1) / 2
    w = (T(1) / 2)^k1 * ab[1, 2] * V[1, I] .^ 2

    x::Vector{T}
    @assert length(x) == N
    w::Vector{T}
    @assert length(w) == N
    return x, w
end

# These functions were formerly a part of Julia. License is MIT:
# https://julialang.org/license

ndgrid(v::AbstractVector) = copy(v)

function ndgrid(v1::AbstractVector{T}, v2::AbstractVector{T}) where {T}
    m, n = length(v1), length(v2)
    v1 = reshape(v1, m, 1)
    v2 = reshape(v2, 1, n)
    return (repeat(v1, 1, n), repeat(v2, m, 1))
end

function ndgrid_fill(a, v, s, snext)
    return for j in 1:length(a)
        a[j] = v[div(rem(j - 1, snext), s) + 1]
    end
end

function ndgrid(vs::AbstractVector{T}...) where {T}
    n = length(vs)
    sz = map(length, vs)
    out = ntuple(i -> Array{T}(undef, sz), n)
    s = 1
    for i in 1:n
        a = out[i]::Array
        v = vs[i]
        snext = s * size(a, i)
        ndgrid_fill(a, v, s, snext)
        s = snext
    end
    return out
end

integrate(f, X, W) = sum(W[i] * f(X[i, :]) for i in 1:length(W))

end


using Distributions

"Boostrapping for degnerate U-statistic."
bootstrap_degenerate(samples::AbstractArray, f) = begin
    # bootstrap sample
    res = 0.0

    n = size(samples, 2)

    # sample bootstrap weights
    sample_p = Multinomial(n, n)
    w = (rand(sample_p) / n) .- (1.0 / n)

    for i = 1:n
        for j = 1:n
            if i == j
                continue
            end
            r = f(samples[:, i], samples[:, j])
            res += ((w[i] * w[j]) * r)
        end
    end

    return res
end


p = BigFloat(pi)
e = exp(BigFloat(1))
table = [
    ["PI", p],
    ["E", e],
    ["LOG2_E", log2(e)],
    ["LN_10", log(big"10")],
    ["LN_2", log(big"2")],
    ["LOG10_2", log10(big"2")],
    ["LOG10_E", log10(e)],
    ["LOG2_10", log2(big"10")],
    ["FRAC_1_PI", 1 / p],
    ["FRAC_1_SQRT_2", 1 / sqrt(big"2")],
    ["FRAC_2_PI", 2 / p],
    ["FRAC_2_SQRT_PI", 2 / sqrt(p)],
    ["FRAC_PI_2", p / 2],
    ["FRAC_PI_3", p / 3],
    ["FRAC_PI_4", p / 4],
    ["FRAC_PI_6", p / 6],
    ["FRAC_PI_8", p / 8],
    ["SQRT_2", sqrt(big"2")],
    ["TAU", p * 2],
]

for (name, value) in table
    println("#[inline]")
    hi = Float64(value)
    lo = Float64(value - hi)
    println("fn $name() -> Self { Self { hi: $hi, lo: $lo } }")
    println("")
end


#=
Fast Type-I Discrete Cosine Transform and Chebyshev polynomial interpolation
=#
module Fct

#fast type I DCT:
#(fct(x_0,...,x_N))_n == sum_{0<=m<=N}x_m*cos(pi*m*n/N)
#with 1/2 weight given to first and last summands.
#WARNING: assumes N is power of 2.
function fct(x::AbstractVector,T::Type=Complex{Float64})::AbstractVector{T}
    fct!(x,fct_alloc(x,T)...)
end
function fct_alloc(x::AbstractVector,T::Type=Complex{Float64})
    N = length(x)-1
    twoN = 2N
    Vector{T}(undef,twoN), Vector{T}(undef,twoN), Vector{T}(undef,twoN)
end
function fct!(dct_in::AbstractVector, #input
              ifft_in::AbstractVector, #workspace: ifft input
              ifft_out::AbstractVector, #output is view(ifft_out,1:N+1)
              ifft_work::AbstractVector, #workspace: ifft workspace
              N=length(dct_in)-1)
    ifft_in[1] = dct_in[1]
    i,j = 2,2N
    while i<=N
        z = dct_in[i]
        ifft_in[i] = z
        ifft_in[j] = z
        i+=1
        j-=1
    end
    ifft_in[N+1] = dct_in[N+1]
    ifft!(ifft_in,ifft_out,ifft_work,2N)
    dct_out = view(ifft_out,1:N+1)
    dct_out .*= 0.5
end

#inverse of fct and dct
function idct(x::AbstractVector,T::Type=Complex{Float64})::AbstractVector{T}
    idct!(x,fct_alloc(x,T)...)
end
function idct!(idct_in::AbstractVector, #input
               ifft_in::AbstractVector, #workspace: ifft input
               ifft_out::AbstractVector, #output is view(ifft_out,1:N+1)
               ifft_work::AbstractVector, #workspace: ifft workspace
               n=length(idct_in)-1)
    idct_out = fct!(idct_in,ifft_in,ifft_out,ifft_work,n)
    T = eltype(ifft_work)
    r = real(T)(2)/n
    idct_out .*= r
end

ifft_cache=Dict{Pair{Int64,Type},Vector}()

#fast inverse DFT without 1/N factor
#WARNING: assumes N is power of 2.
#(ifft(x_0,...,x_{N-1}))_n = sum_{0<=m<N}x_m*exp(2*pi*m*n/N)
function ifft(x::AbstractVector,T::Type=Complex{Float64})
    N=length(x)
    y=Vector{T}(undef,N) #output
    z=Vector{T}(undef,N) #work space
    ifft!(x,y,z)
end
function ifft!(input::AbstractVector,output::AbstractVector,workspace::AbstractVector,
               N=length(input))
    T=eltype(workspace)
    key = N=>T
    if key in keys(ifft_cache)
        exps = ifft_cache[key]
    else
        exps = ifft_cache[key] = exp_populate(N,T)
    end
    ifft_helper!(input,output,workspace,exps,0,1,N)
end

#x is input; y is output; z is work space
#assumes length(x)==length(y)==length(exps)<=length(z)
#input/output indices: 1+offset,1+offset+stride,...,1+offset+(len-1)*stride
function ifft_helper!(x::AbstractVector,
                      y::AbstractVector,
                      z::AbstractVector,
                      exps::AbstractVector,
                      offset::Integer,
                      stride::Integer,
                      len::Integer)
    if len==1
        y[1+offset]=x[1+offset]
        return y
    end
    doublestride=stride<<1
    halflen=len>>1
    ifft_helper!(x,y,z,exps,offset,doublestride,halflen)
    ifft_helper!(x,y,z,exps,offset+stride,doublestride,halflen)
    i=1
    j=1+offset
    k=1
    while k<=halflen
        z[k]=y[j]+y[j+stride]*exps[i]
        i+=stride
        j+=doublestride
        k+=1
    end
    j=1+offset
    while k<=len
        z[k]=y[j]+y[j+stride]*exps[i]
        i+=stride
        j+=doublestride
        k+=1
    end
    #println()
    i=1+offset
    k=1
    while k<=len
        y[i]=z[k]
        i+=stride
        k+=1
    end
    y
end

#output: exps[1+m]==exp(2im*pi*m/N) for 0<=m<N
#N can be any positive integer.
function exp_populate(N::Integer,T::Type=Complex{Float64})
    exps=Vector{T}(undef,N)
    w = 2im*T(pi)/N
    for m in 0:N-1
        exps[m+1] = exp(w*m)
    end
    exps
end

#=
Chebyshev polynomial interpolation of f on [-1,1] intended for Clenshaw-Curtis quadrature.
On input y_0,...,y_N, outputs c_0,...,c_N such that y_n == sum_m c_m cos(m*n*pi/N).
y_m should be f(cos(m*pi/N)).
WARNING: N must be a power of 2.
=#
function chebcoeffs(samples::AbstractVector,T::Type=promote_type(Complex{Bool},eltype(samples)))
    chebcoeffs!(samples,fct_alloc(samples,T)...)
end
function chebcoeffs!(samples::AbstractVector,
                     ifft_in::AbstractVector,
                     ifft_out::AbstractVector,#output is view(ifft_out,1:length(samples))
                     ifft_work::AbstractVector,
                     n=length(samples)-1)
    coeffs = idct!(samples,ifft_in,ifft_out,ifft_work,n)
    coeffs[1] *= 0.5
    coeffs[end] *= 0.5
    coeffs
end

function chebsample(f,n::Integer,base=0,stride=1;T::Type=Complex{Float64})
    samples = Vector{T}(undef,length(base:stride:n))
    chebsample!(f,samples,n,base,stride)
end
function chebsample!(f,samples,n::Integer=length(samples)-1,base=0,stride=1)
    p = real(eltype(samples))(pi)
    i=1
    for j in base:stride:n
        samples[i] = f(cos(p*j/n))
        i += 1
    end
    samples
end

function shiftedchebsample(center,radius,f1,f2,n::Integer,base=0,stride=1;T::Type=Complex{Float64})
    samples = Vector{T}(undef,length(base:stride:n))
    shiftedchebsample!(center,radius,f1,f2,samples,n,base,stride)
end
function shiftedchebsample!(center,radius,f1,f2,samples,n::Integer=length(samples)-1,base=0,stride=1)
    p = real(eltype(samples))(pi)
    i=1
    for j in base:stride:n
        x = center + radius*cos(p*j/n)
        samples[i] = f1(x)*f2(x)
        i += 1
    end
    samples
end

function tonechebsample(cfreq,center,radius,f1,f2,n::Integer,base=0,stride=1;T::Type=Complex{Float64})
    samples = Vector{T}(undef,length(base:stride:n))
    tonechebsample!(cfreq,center,radius,f1,f2,samples,n,base,stride)
end
function tonechebsample!(cfreq,center,radius,f1,f2,samples,n::Integer=length(samples)-1,base=0,stride=1)
    p = real(eltype(samples))(pi)
    i=1
    for j in base:stride:n
        x = cos(p*j/n)
        y = center + radius*x
        z = x * -cfreq
        s,c = sincos(z)
        samples[i] = f1(y)*f2(y)*complex(c,s)
        i += 1
    end
    samples
end

function chirpchebsample(cfreq,chirp,center,radius,f1,f2,n::Integer,base=0,stride=1;T::Type=Complex{Float64})
    samples = Vector{T}(undef,length(base:stride:n))
    chirpchebsample!(cfreq,chirp,center,radius,f1,f2,samples,n,base,stride)
end
function chirpchebsample!(cfreq,chirp,center,radius,f1,f2,samples,n::Integer=length(samples)-1,base=0,stride=1)
    p = real(eltype(samples))(pi)
    i=1
    for j in base:stride:n
        x = cos(p*j/n)
        y = center + radius*x
        z = -x*(cfreq + x*chirp)
        s,c = sincos(z)
        samples[i] = f1(y)*f2(y)*complex(c,s)
        i += 1
    end
    samples
end

#double Chebyshev degree
function doublesample(f,oldsamples)
    N = length(oldsamples)-1
    T = eltype(oldsamples)
    newsamples = chebsample(f,2N,1,2;T=T)
    samples = Vector{T}(undef,2N+1)
    samples[1] = oldsamples[1]
    for n in 1:N
        samples[2n] = newsamples[n]
        samples[2n+1] = oldsamples[n+1]
    end
    samples
end

function shifteddoublesample(center,radius,f1,f2,oldsamples)
    N = length(oldsamples)-1
    T = eltype(oldsamples)
    newsamples = shiftedchebsample(center,radius,f1,f2,2N,1,2;T=T)
    samples = Vector{T}(undef,2N+1)
    samples[1] = oldsamples[1]
    for n in 1:N
        samples[2n] = newsamples[n]
        samples[2n+1] = oldsamples[n+1]
    end
    samples
end

function chirpdoublesample(cfreq,chirp,center,radius,f1,f2,oldsamples)
    N = length(oldsamples)-1
    T = eltype(oldsamples)
    newsamples = chirpchebsample(cfreq,chirp,center,radius,f1,f2,2N,1,2;T=T)
    samples = Vector{T}(undef,2N+1)
    samples[1] = oldsamples[1]
    for n in 1:N
        samples[2n] = newsamples[n]
        samples[2n+1] = oldsamples[n+1]
    end
    samples
end

function doublesample!(f,samples,N)
    for n in N:-1:1
        samples[2n+1] = samples[n+1]
    end
    chebsample!(f,view(samples,2:2:2N),2N,1,2)
    samples
end

function shifteddoublesample!(center,radius,f1,f2,samples,N)
    for n in N:-1:1
        samples[2n+1] = samples[n+1]
    end
    shiftedchebsample!(center,radius,f1,f2,view(samples,2:2:2N),2N,1,2)
    samples
end

function tonedoublesample!(cfreq,center,radius,f1,f2,samples,N)
    for n in N:-1:1
        samples[2n+1] = samples[n+1]
    end
    tonechebsample!(cfreq,center,radius,f1,f2,view(samples,2:2:2N),2N,1,2)
    samples
end

end #module


using Pumas, MCMCChains, StatsPlots, Random
plotly()

# Define Population
population = Population(
    map(i -> Subject(
        id = i),
        1:8))

# Build model
tumour_growth_model = @model begin
    @param   begin
        # mean_initial_volume ~ Constrained(
        #     Normal(0.2, 1), lower=0.001, upper=Inf, init=0.2)
        # mean_critical_volume ~ Constrained(
        #     Normal(0.2, 1), lower=0.001, upper=Inf, init=0.2)
        # mean_growth_rate ~ Constrained(
        #     Normal(0.2, 1), lower=0.001, upper=Inf, init=0.2)
        # std ~ Constrained(
        #     MvNormal([1 , 1, 1]),
        #     lower=[0, 0, 0],
        #     upper=[Inf, Inf, Inf],
        #     init=[0.09,0.09, 0.09])
        # σ_base ~ Constrained(
        #     Normal(0.2, 1), lower=0.001, upper=Inf, init=0.2)
        # σ_rel ~ Constrained(
        #     Normal(0.2, 1), lower=0.001, upper=Inf, init=0.2)
        mean_initial_volume ∈ RealDomain(lower=0, init = 0.2)
        mean_critical_volume ∈ RealDomain(lower=0, init = 1)
        mean_growth_rate ∈ RealDomain(lower = 0, init= 0.1)
        std ∈ PDiagDomain(init=[0.09,0.09, 0.09])
        σ_base ∈ RealDomain(lower=0, init=0.04)
        σ_rel ∈ RealDomain(lower=0, init=0.04)
    end

    @random begin
        η ~ MvNormal(std)
    end

    @pre begin
        initial_volume = mean_initial_volume * exp(η[1])
        critical_volume = mean_critical_volume * exp(η[2])
        growth_rate = mean_growth_rate * exp(η[3])
    end

    @init begin
        volume = initial_volume
    end

    @dynamics begin
       volume' = (growth_rate * critical_volume * volume) / (volume + critical_volume)
    end

  @derived begin
      σ_tot := σ_base .+ volume * σ_rel
      Measurement ~ @. Normal(volume, σ_tot)
    end
end

# Simulate Population
param = init_param(tumour_growth_model)
obs = simobs(tumour_growth_model, population, param, obstimes=0:1:30)
scatter(obs)

# Convert to dataset
simdf = DataFrame(obs)
data = read_pumas(
    simdf, observations=[:Measurement], event_data=false)

# Infer parameters
param = init_param(tumour_growth_model)
result = fit(
    tumour_growth_model, data, param, Pumas.BayesMCMC();
    nsamples=2000, nadapts=1000)

# Visualise population parameters
chains = Pumas.Chains(result)
plot(chains[1000:end])

tumour_growth_model.prob


# a!/b!
factorial_ratio(a::I, b::I) where {I<:Integer} = gamma(a+1)/gamma(b+1) #binomial(a,b)*factorial(a-b)

function coulomb_analytical(k::I, γ::R, ℓ::I, r̃::R) where {I<:Integer, R<:Real}
    e⁻ᵞʳ = exp(-γ*r̃)

    # S = -gamma(k+ℓ+1)/r̃^(ℓ+1)*(e⁻ᵞʳ-1)/γ^(k+ℓ+1)
    S = 2gamma(k+ℓ+1)*exp(-γ*r̃/2 -(ℓ+1)*log(r̃))*sinh(γ*r̃/2)/γ^(k+ℓ+1)

    for i ∈ 0:k+ℓ-1
        S -= factorial_ratio(k+ℓ,k+ℓ-i)*r̃^(k-i-1)*e⁻ᵞʳ/γ^(i+1)
    end

    for i ∈ 0:k-ℓ-1
        S += factorial_ratio(k-ℓ-1,k-ℓ-i-1)*r̃^(k-i-1)*e⁻ᵞʳ/γ^(i+1)
    end

    S
end


#############################################################################
# Computes the bases that maxize a linear function over the base polyhedron
# using the greedy algorithm.
#############################################################################

export greedy, greedy_rand

TOL = 1e-2

# Computes one solution
function greedy(F::Function, x::AbstractArray,
                S = collect(1:length(x))::AbstractArray)
  @assert F([])[1] == 0 # f should be 0 at the empty set
  n = length(x)
  @assert length(S) == n
  i = sortperm(x, rev = true)
  V = sort(S)
  w = zeros(n)
  for ii = 1:n
    w[i[ii]] = F(V[i[1: ii]])[1] - F(V[i[1: ii-1]])[1]
  end
  return w
end

greedy(F::SubmodFunc, x::AbstractArray) = greedy(x -> evaluate(F, x), x)

greedy(f::LovaszExtAtom, x::AbstractArray) = greedy(f.func, x)

function greedy(f::LovaszExtAbsAtom, x::AbstractArray)
  w = greedy(f.func, abs.(x))
  for i = 1:length(w)
    if x[i] < 0
      w[i] = -w[i]
    end
  end
  return w
end

# Compute a solution with indexed on the same level set permutated
function greedy_rand(F::Function, x::AbstractArray, tol = 1e-3 ::Number,
                    S = collect(1:length(x))::AbstractArray)
  @assert F([])[1] == 0       # f should be 0 at the empty set
  n = length(x)
  @assert length(S) == n
  ordering = sortperm(x, rev = true)
  pernum = 1
  V = sort(S)
  ordering = V[ordering]      # the indexed reordered in a descending order
  # Mark level sets
  index = 1
  level_sets = zeros(2, n)    # the beginning and ending of each level set
  level_sets[1, 1] = 1
  for i = 2:n
    if x[ordering[i - 1]] - x[ordering[i]] >= tol
    level_sets[2, index] = i-1
    index += 1
    level_sets[1, index] = i
    end
  end
  level_sets[2, index] = n
  new_ordering = zeros(n)
  for i = 1:index
    new_ordering[Int(level_sets[1, i]) : Int(level_sets[2, i])] = shuffle(ordering[Int(level_sets[1, i]) : Int(level_sets[2, i])])
    pernum *= factorial(level_sets[2, i] - level_sets[1, i] + 1)
  end
  new_ordering = Int.(new_ordering)
  w = zeros(n)
  for ii = 1:n
    w[new_ordering[ii]] = F(V[new_ordering[1: ii]])[1] - F(V[new_ordering[1: ii-1]])[1]
  end
  return w, pernum, new_ordering
end

greedy_rand(F::SubmodFunc, x::AbstractArray, tol::Number) = greedy_rand(x -> evaluate(F, x), x, tol)


using DiffEqSensitivity
using OrdinaryDiffEq, Calculus, Test
using Zygote

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

function condition(u,t,integrator) # Event when event_f(u,t) == 0
  u[1]
end

function affect!(integrator)
  @show integrator.t
  println("bounced.")
  integrator.u[2] = -integrator.p[2]*integrator.u[2]
end

cb = ContinuousCallback(condition, affect!)
p = [9.8, 0.8]
prob = ODEProblem(f,eltype(p).([1.0,0.0]),eltype(p).((0.0,1.0)),copy(p))

function test_f(p)
  _prob = remake(prob, p=p)
  solve(_prob,Tsit5(),abstol=1e-14,reltol=1e-14,callback=cb,save_everystep=false)[end]
end
findiff = Calculus.finite_difference_jacobian(test_f,p)
findiff

using ForwardDiff
ad = ForwardDiff.jacobian(test_f,p)
ad

@test ad ≈ findiff

function test_f2(p, sensealg=ForwardDiffSensitivity(), controller=nothing, alg=Tsit5())
  _prob = remake(prob, p=p)
  u = solve(_prob,alg,sensealg=sensealg,controller=controller,
    abstol=1e-14,reltol=1e-14,callback=cb,save_everystep=false)
  u[end][end]
end

@test test_f2(p) == test_f(p)[end]

g1 = Zygote.gradient(θ->test_f2(θ,ForwardDiffSensitivity()), p)
g2 = Zygote.gradient(θ->test_f2(θ,ReverseDiffAdjoint()), p)
g3 = Zygote.gradient(θ->test_f2(θ,ReverseDiffAdjoint(), IController()), p)
g4 = Zygote.gradient(θ->test_f2(θ,ReverseDiffAdjoint(), PIController(7//50, 2//25)), p)
@test_broken g5 = Zygote.gradient(θ->test_f2(θ,ReverseDiffAdjoint(), PIDController(1/18. , 1/9., 1/18.)), p)
g6 = Zygote.gradient(θ->test_f2(θ,ForwardDiffSensitivity(),
  OrdinaryDiffEq.PredictiveController(), TRBDF2()), p)
@test_broken g7 = Zygote.gradient(θ->test_f2(θ,ReverseDiffAdjoint(),
    OrdinaryDiffEq.PredictiveController(), TRBDF2()), p)

@test g1[1] ≈ findiff[2,1:2]
@test g2[1] ≈ findiff[2,1:2]
@test g3[1] ≈ findiff[2,1:2]
@test g4[1] ≈ findiff[2,1:2]
@test_broken g5[1] ≈ findiff[2,1:2]
@test g6[1] ≈ findiff[2,1:2]
@test_broken g7[1] ≈ findiff[2,1:2]


## Assumptions
## http://docs.sympy.org/0.7.2/modules/assumptions/index.html

"""
     refine

Simplify an expression using assumptions; [refine](https://docs.sympy.org/dev/modules/assumptions/refine.html).
"""
refine(ex, assumpts...) = sympy.refine(ex, assumpts...)

"""
    ask(query)

Returns `true`, `false`, or `nothing`; [ask](https://docs.sympy.org/dev/modules/assumptions/ask.html)

Example:

```jldoctest ask
julia> using SymPyCall

julia> @syms x::integer y::integer
(x, y)

julia> ask(𝑄.integer(x*y), And(𝑄.integer(x), 𝑄.integer(y)))
ERROR: UndefVarError: And not defined
Stacktrace:
 [1] top-level scope
   @ none:1

julia> ## really slow isprime:
       filter(x -> ask(𝑄.prime(x)), 1:10)
4-element Vector{Int64}:
 2
 3
 5
 7
```

"""
function ask(x::Sym, args...)
    val = sympy.ask(x, args...)
    Bool(getpy(val) == pybuiltins.None) && return nothing
    Bool(getpy(val) == pybuiltins.True) && return true
    Bool(getpy(val) == pybuiltins.False) && return false
    throw(DomainError("Argument not none, true, or false"))
end

ask(x::Bool, args...) = x
ask(x::Nothing, args...) = x

# ## should we support & and | for (sym,sym) pairs? Not sure
# ## dependso on what happens to | (x==0) and ex == x usage for ==
# ## for now, we can combine terms logically with And, Or, Not...
# ## these are in logic module

# ## We make a module Q to hold the assumptions
# ## this follows this page http://docs.sympy.org/0.7.5/_modules/sympy/assumptions/ask.html
# """
#     𝑄 [\itQ]
#     SymPy.Q

# Documentation for the `SymPy.Q` module, exported as `𝑄`.

# SymPy allows for
# [assumptions](https://docs.sympy.org/latest/modules/assumptions/index.html)
# on variables. These may be placed on free sympols at construction.

# For example, the following creates a real value variable `x` and a postive, real variable `y`:

# ```jldoctest 𝑄
# julia> using SymPy

# julia> @syms x::real
# (x,)

# julia> @syms y::(real,positive)
# (y,)
# ```

# The `Q` module exposes a means to *q*uery the assumptions on a
# variable. For example,

# ```jldoctest 𝑄
# julia> ask(𝑄.positive(y))  # true
# true

# julia> ask(𝑄.negative(y))  # false
# false

# julia> ask(SymPyCall.Q.positive(x))  # `nothing`

# julia> ask(SymPyCall..Q.positive(x^2)) # `nothing` -- might be 0

# julia> ask(SymPyCall.Q.positive(1 + x^2)) # true  -- must be postive now.
# true
# ```

# The ask function uses tri-state logic, returning one of 3 values:
# `true`; `false`; or `nothing`, when the query is indeterminate.

# The construction of predicates is done through `Q` methods. These can
# be combined logically. For example, this will be `true`:

# ```jldoctest 𝑄
# julia> ask(𝑄.positive(y) & 𝑄.negative(-x^2 - 1))

# ```

# The above use `&` as an infix operation for the binary operator
# `And`. Values can also be combined with `Or`, `Not`, `Xor`, `Nand`,
# `Nor`, `Implies`, `Equivalent`, and `satisfiable`.

# !!! note "Matrix predicates"
#     As `SymPy.jl` converts symbolic matrices into Julia's `Array`
# type and not as matrices within Python, the predicate functions from SymPy for
# matrices are not used, though a replacement is given.
# """
# module Q
# import SymPyCall: sympy
# import PythonCall
# import LinearAlgebra: det, norm

# ##http://docs.sympy.org/dev/_modules/sympy/assumptions/ask.html#ask
# Q_predicates = (:antihermitian,
#                 :bounded, :finite, # bounded deprecated
#                 :commutative,
#                 :complex,
#                 :composite,
#                 :even,
#                 :extended_real,
#                 :hermitian,
#                 :imaginary,
#                 :infinitesimal,
#                 :infinity, :infinite, # :infinity deprecated
#                 :integer,
#                 :irrational,
#                 :rational,
#                 :algebraic,
#                 :transcendental,
#                 :negative,
#                 :nonzero, :zero,
#                 :positive,
#                 :prime,
#                 :real,
#                 :odd,
#                 :is_true,
#                 :nonpositive,
#                 :nonnegative
# #                :symmetric,
# #                :invertible,
# #                :singular,
# #                :orthogonal,
# #                :unitary,
# #                :normal,
# #                :positive_definite,
# #                :upper_triangular,
# #                :lower_triangular,
# #                :diagonal,
# #                :triangular,
# #                :unit_triangular,
# #                :fullrank,
# #                :square,
# #                :real_elements,
# #                :complex_elements,
# #                :integer_elements
# )

# for meth in Q_predicates
#    nm = string(meth)
#       @eval begin
#           ($meth)(x) = sympy.Q.$(nm)(x)
#               #PyCall.pycall(SymPy.sympy.Q.$nm, SymPy.Sym, x)::SymPy.Sym
#    end
# end

# symmetric(M::Array{T,2}) where {T <: SymPyCall.Sym} = SymPyCall.issymmetric(M)
# function invertible(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     d = det(M)
#     pos = SymPyCall.ask(positive(d))
#     if pos == nothing
#         return nothing
#     elseif pos == true
#         return true
#     end
#     neg = SymPyCall.ask(negative(d))
#     if neg == nothing
#         return nothing
#     end
#     z = SymPyCall.ask(zero(d))
#     if z == nothing
#         return nothing
#     elseif z == true
#         return false
#     end

#     return true
# end

# function singular(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     !invertible(M)
# end

# function orthogonal(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     vals = SymPyCall.simplify.(SymPyCall.simplify.(M*transpose(M)) .== one(T))
#     no_nothing = 0
#     for val in vals
#         a = SymPyCall.ask(zero(val))
#         if a == nothing
#             no_nothing += 1
#         elseif a == false
#             return false
#         end
#     end

#     no_nothing > 0 && return nothing
#     return true
# end

# function unitary(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     vals = SymPyCall.simplify.(SymPyCall.simplify.(M*ctranspose(M)) .== one(T))
#     no_nothing = 0
#     for val in vals
#         a = SymPyCall.ask(zero(val))
#         if a == nothing
#             no_nothing += 1
#         elseif a == false
#             return false
#         end
#     end

#     no_nothing > 0 && return nothing
#     return true
# end

# function normal(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     lhs = ctranspose(M) * M
#     rhs = M * ctranspose(M)
#     vals = zero.(SymPyCall.simplify.(lhs - rhs))
#     no_nothing = 0
#     for val in vals
#         a = SymPyCall.ask(val)
#         if a == nothing
#             no_nothing += 1
#         elseif a == false
#             return false
#         end
#     end

#     no_nothing > 0 && return nothing
#     return true
# end

# # Use [Sylvester's](https://en.wikipedia.org/wiki/Sylvester%27s_criterion) Criteria
# function positive_definite(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     !SymPyCall.ask(square(M)) && return false
#     !SymPyCall.ask(symmetric(M))  && return false
#     m, n = size(M)
#     no_false = 0
#     no_nothing = 0
#     for i in 1:m
#         a = SymPyCall.ask(Q.positive(det(M[1:i, 1:i])))
#         if a == nothing no_nothing += 1 end
#         if a == false no_false += 1 end
#     end
#     if no_false > 0
#         return false
#     elseif no_nothing > 0
#         return nothing
#     else
#         return true
#     end
# end

# upper_triangular(M::Array{T,2}) where {T <: SymPyCall.Sym} = SymPyCall.istriu(M)
# lower_triangular(M::Array{T,2}) where {T <: SymPyCall.Sym} = SymPyCall.istril(M)
# diagonal(M::Array{T,2}) where {T <: SymPyCall.Sym} = upper_triangular(M) && lower_triangular(M)
# triangular(M::Array{T,2}) where {T <: SymPyCall.Sym} = upper_triangular(M) || lower_triangular(M)

# ## This is likely not the best way as it is a bit fidgety due
# ## to the call to rref.
# function full_rank(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     m,n = size(M)
#     m <= n || return full_rank(transpose(M))

#     rr, p = SymPyCall.rref(M)
#     lr = rr[end, :] # is this zero?
#     no_nothing = 0
#     no_nonzero = 0
#     for val in lr
#         a = SymPyCall.ask(nonzero(val))
#         if a == nothing
#             no_nothing += 1
#         end
#         if a == true
#             no_nonzero += 1
#         end
#     end
#     if no_nothing > 0
#         return nothing
#     elseif no_nonzero == 0
#         return false
#     else
#         return true
#     end

# end

# function square(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     m,n = SymPyCall.size(M)
#     m == n
# end

# function real_elements(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     vals = real.(M)
#     for val in vals
#         a = SymPyCall.ask(real(val))
#         (a == nothing || a == false) && return false
#     end
#     return true
# end

# function complex_elements(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     vals = real.(M)
#     for val in vals
#         a = SymPyCall.ask(SymPyCall.sympy."Q".complex(val))
#         (a == nothing || a == false) && return false
#     end
#     return true

# end

# function integer_elements(M::Array{T,2}) where {T <: SymPyCall.Sym}
#     vals = real.(M)
#     for val in vals
#         a = SymPyCall.ask(integer(val))
#         (a == nothing || a == false) && return false
#     end
#     return true
# end

# end

# ## Issue  #354; request to *not*  export  Q
# ## export
# #export Q

# const 𝑄 = Q

# """
#     𝑄

# Exported  symbol for  [`SymPy.Q`](@ref), a  Julia  module implementing `sympy.Q`. "Questions" can be asked through the patterns
# `𝑄.query(value)` (𝑄 is entered as  [slash]itQ[tab]) or `SymPy.Q.query(value)` *but  not* as `sympy.Q.query(value)`

# !!! note
#     At one time, the symbol `Q` was exported for this. To avoid namespace clutter, the unicode alternative is now used. Legacy code would need a definition like `const Q = SymPy.Q`  to work.

# """
# 𝑄
# export  𝑄


#! /usr/bin/julia

# Rosetta Code, Spiral matrix

immutable Spiral
    m::Int
    n::Int
    cmax::Int
    dir::Array{Array{Int,1},1}
    bdelta::Array{Array{Int,1},1}
end

function Spiral(m::Int, n::Int)
    cmax = m*n
    dir = Array{Int,1}[[0,1], [1,0], [0,-1], [-1,0]]
    bdelta = Array{Int,1}[[0,0,0,1], [-1,0,0,0],
                          [0,-1,0,0], [0,0,1,0]]
    Spiral(m, n, cmax, dir, bdelta)
end

function spiral(m::Int, n::Int)
    0<m&&0<n || error("The matrix dimensions must be positive.")
    Spiral(m, n)
end
spiral(n::Int) = spiral(n, n)

type SpState
    cnt::Int
    dirdex::Int
    cell::Array{Int,1}
    bounds::Array{Int,1}
end

Base.length(sp::Spiral) = sp.cmax
Base.start(sp::Spiral) = SpState(1, 1, [1,1], [sp.n,sp.m,1,1])
Base.done(sp::Spiral, sps::SpState) = sps.cnt > sp.cmax

function Base.next(sp::Spiral, sps::SpState)
    s = sub2ind((sp.m, sp.n), sps.cell[1], sps.cell[2])
    if sps.cell[rem1(sps.dirdex+1, 2)] == sps.bounds[sps.dirdex]
        sps.bounds += sp.bdelta[sps.dirdex]
        sps.dirdex = rem1(sps.dirdex+1, 4)
    end
    sps.cell += sp.dir[sps.dirdex]
    sps.cnt += 1
    return (s, sps)
end

using Formatting

function width{T<:Integer}(n::T)
    w = ndigits(n)
    n < 0 || return w
    return w + 1
end

function pretty{T<:Integer}(a::Array{T,2}, indent::Int=4)
    lo, hi = extrema(a)
    w = max(width(lo), width(hi))
    id = " "^indent
    fe = FormatExpr(@sprintf(" {:%dd}", w))
    s = id
    nrow = size(a)[1]
    for i in 1:nrow
        for j in a[i,:]
            s *= format(fe, j)
        end
        i != nrow || continue
        s *= "\n"*id
    end
    return s
end

n = 5
println("The n = ", n, " spiral matrix:")
a = zeros(Int, (n, n))
for (i, s) in enumerate(spiral(n))
    a[s] = i-1
end
println(pretty(a))

m = 3
println()
println("Generalize to a non-square matrix (", m, "x", n, "):")
a = zeros(Int, (m, n))
for (i, s) in enumerate(spiral(m, n))
    a[s] = i-1
end
println(pretty(a))

p = primes(10^3)
n = 7
println()
println("An n = ", n, " prime spiral matrix:")
a = zeros(Int, (n, n))
for (i, s) in enumerate(spiral(n))
    a[s] = p[i]
end
println(pretty(a))


using TheAlgorithms
using Test

using LinearAlgebra

@testset "TheAlgorithms" begin

include("conversions.jl")
include("data_structures.jl")
include("knapsack.jl")
include("math.jl")
include("matrix.jl")
include("project-rosalind.jl")
include("scheduling.jl")
include("searches.jl")
include("sorts.jl")
include("statistics.jl")
include("strings.jl")

end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ f80fd13d-207a-458a-b326-77256fe946e0
using Plots

# ╔═╡ bda0f8d4-5ab0-4c65-9d79-71ccbfe8dd3a
using SpecialFunctions

# ╔═╡ b608f792-dd91-4bd2-97ab-baac7b7d2f9c
using Dierckx

# ╔═╡ 240e07ba-1508-485c-8fe1-d33e8fa52804
using SpecialFunctions

# ╔═╡ 4e16dbbe-b0fc-11eb-2349-ad7ccd0ff56b
md"""
# Simple questions and simple answers about the Boys function
Rick Muller

A large part of electronic structure theory lies in computing the repulsive
interaction between electrons, and a large part of electron repulsion
is captured in the Boys function that computes the repulsion between
partially overlapping electronic wave functions.

The Boys function is given by [^boys50]

$Fm(T)=\int_0^1 u^{2m}\exp(-Tu^2)du$

and was introduced by Boys in 1950. Since that time, a surprising amount
of work has gone into efficiently computing [^gjp91] this function. The
best of these methods [^gjp91] divides the integration over $u$ into multiple slices
and approximates the integrand in this region using Taylor or Chebyshev
polynomials. Efficient programs [^libint] to evaluate these terms have detailed
tables with nodes and weights for the interpolation and quadrature schemes.

One says "surprising" because we very much know what electron repulsion
looks like in crude terms: at long distances, the interaction should look like Coulombic $1/r$ repulsion, and at shorter distances the interaction rolls
over to a constant value with zero slope. This behavior was approximated by
Mataga and Nishimoto [^mn57] as

$F_{MN}(r) \sim \frac{1}{a+r}$

which we plot below.

"""

# ╔═╡ 8a736133-17f2-446d-9ab5-4bc46bdd556b
fmn(r,a=1) = 1/(a+r);

# ╔═╡ bf014265-210c-412f-9af4-5a654f557de9
begin
	rs = 0:100
	mnvals =
	plot(rs,[fmn(r,1) for r in rs],label="fmn(r,a=1)",
		title="Behavior of the Mataga-Nishimoto approximation",
		xlabel="R (au)")
end

# ╔═╡ d8061039-367c-42ee-8a65-a32dca37fc67
md"""
We can use the following simple definition for the Boys function taken from
[^libint]:
"""

# ╔═╡ bc53663d-1f9f-43a6-8140-e8eeac1618bb
"Reference Fm from libint"
function fboys(m,T,eps = 1e-10)
    denom = (m + 0.5)
    term = exp(-T) /2denom
    old_term = 0.0
    sum = term
    while true
        denom += 1
        old_term = term
        term = old_term * T / denom
        sum += term
        term > sum*eps || old_term < term || break
    end
    return sum
end;

# ╔═╡ 1a0bffb2-8656-44db-b3b3-6cf6b2494b7c
md"This is the version that uses the gamma function:"

# ╔═╡ b6b757ad-7fca-46c5-b42f-20f94ae4cbd7
"gammainc returns the lower incomplete gamma function"
gammainc(a,x) = gamma(a)*gamma_inc(a,x)[1]

# ╔═╡ 145a0007-b8d7-46a1-9cb8-ee063057527c
"Boys Fgamma function, using the lower incomplete gamma function."
function Fgamma(m,x,SMALL=1e-18,Tcrit=20.0)
    # Note, most programs use a much larger value for Tcrit (117)
    mhalf = m+0.5
    x = max(x,SMALL) # Evidently needs underflow protection
    if x>Tcrit
        retval = sqrt(pi/2)*factorial2(2m-1)/(2x)^mhalf
    else
        retval = 0.5*x^-mhalf*gammainc(mhalf,x)
    end
    return retval
end

# ╔═╡ f6fa9869-b454-458d-95e3-818a9112cc52
md"""
The asymptotic behavior of the Boys function for large $T$ is known
to be

$F_\mathrm{asymp}(T,m) = \sqrt{\frac{\pi}{2}}(2m-1)!!/(2T)^{m+0.5}$

which leads to a slightly more complicated approximation form of

$F_\mathrm{approx}(T)=\sqrt{\frac{\pi}{2}}\frac{1}{a\sqrt{\pi/2}+\sqrt{2T}}$
"""

# ╔═╡ e7454913-e406-4b4a-856f-47103a7bf23d
fasymp(m,T) = sqrt(0.5*pi)*prod(2m-1:-2:1)*(2T)^(-m-0.5);

# ╔═╡ cdc2c6c9-e567-4ee6-a002-8fea660e35fa
function fapprox(T,a=1)
	pre = sqrt(0.5*pi)
	return pre/(a*pre + (2T)^(0.5))
end;

# ╔═╡ b5f42a47-eb6c-4b59-9827-a07b81618018
md"""
This approximation leads to good agreement with the Boys function.
"""

# ╔═╡ 105717a4-f481-479d-bfbb-dc35cc3c0997
begin
	plot(rs,[fboys(0,r) for r in rs],label="fboys(0,r)")
	plot!(rs,[fapprox(r,1) for r in rs],label="fapprox(r,a=1)")
	#title!("Parameter-free approximation of Boys function")
end

# ╔═╡ 0b513de7-cca5-45ca-a4fa-2f2604eb7da9
md"""
We note that although the behavior of `fapprox` plotted above is not exact,
it is essentially parameter-free when `a=1`, and thus represents a very efficient starting point for further analysis.
"""

# ╔═╡ 4dd0b1e6-61f0-4daa-80f1-625dc793b33a
md"""
As an aside, we certainly don't mean to denigrate the excellent work that has
gone into the development of the approximations methods [^gjp91] or their
efficient implementation [^libint]. These are remarkably accurate and fast
methods that have been derived with great thought and applied repeatedly to
chemical calculations with excellent results. Our intent here is to merely
understand whether a more intuitive approach is possible to this admittedly
solved problem.
"""

# ╔═╡ fdd2691f-f26f-4c26-8cc6-0d8edaa5278a
md"""
Instead of

$F_\mathrm{approx}(T)=\sqrt{\frac{\pi}{2}}\frac{1}{a\sqrt{\pi/2}+\sqrt{2T}}$

we can expand the denominator in a term that will go to zero at large r:

$F_\mathrm{approx2}(T)=\sqrt{\frac{\pi}{2}}\frac{1}{a\exp(-bT)\sqrt{\pi/2}+\sqrt{2T}}$
"""

# ╔═╡ a0a09c9f-c6d2-422f-b32a-8e3f8131b3c3
function fapprox2(T,a=1,b=1)
	pre = sqrt(0.5*pi)
	return pre/(a*pre*exp(-b*T) + (2T)^(0.5))
end;

# ╔═╡ ae9658a3-9241-423d-9daa-d6644e8a206a
begin
	rs2=0:0.1:15
	plot(rs2,[fboys(0,r) for r in rs2],label="fboys(0,r)",ylim=(0,1))
	plot!(rs2,[fasymp(0,r) for r in rs2],label="fasymp(0,r)")
	plot!(rs2,[fapprox(r,1) for r in rs2],label="fapprox(r,1)")
	plot!(rs2,[fapprox2(r,1,1.5) for r in rs2],label="fapprox2(r,a=1,b=1.5)")
	#title!("Parameter-free approximation of Boys function")
end

# ╔═╡ 8c8adf9b-5eff-4942-9a49-95a4aa1e2002
md"""
With very few parameters (essentially just estimating the `b=1.5` fit) it's possible to get a very good reproduction of the Boys function.

If rather than $exp(-bT)$ we assume there is some function of $T$ that multiplies
the decay of the constant term, what would that function look like?
"""

# ╔═╡ 9428670f-5cf7-4c9e-bb93-dd0747a34182
boysdenom(T) = 1/fboys(0,T) - sqrt(2T)/sqrt(pi/2);

# ╔═╡ 771d4eeb-72e6-4863-a114-ff203a42455b
begin
	rs3 = 0:0.1:3
	plot(rs3,boysdenom.(rs3),label="Boys denom decay")
	plot!(rs3,[exp(-T-0.56) for T in rs3],label="-T-0.56")
	plot!(rs3,[exp(-4T) for T in rs3],label="-4T")
	#plot!(rs3,[exp(-4T)/2+exp(-T-0.56)/2 for T in rs3],label="avg")
end

# ╔═╡ 997e6018-bb0e-4933-8b13-b904e96fc819
md"Plot this as `fapprox3`"

# ╔═╡ 894135da-1ef9-4170-aa49-1b6307e398eb
function fapprox3(T,a=1)
	pre = sqrt(0.5*pi)
	return pre/(a*pre*exp(-T-0.56) + (2T)^(0.5))
end;

# ╔═╡ 9dfb09bb-22d8-44f2-9c4b-98777e43de64
begin
	plot(rs3,[fboys(0,r) for r in rs3],label="fboys(0,r)")
	plot!(rs3,[fapprox2(r) for r in rs3],label="fapprox2(r)")
	plot!(rs3,[fapprox3(r) for r in rs3],label="fapprox3(r)")
end

# ╔═╡ d526caa1-ff7e-4ec5-ab59-2a10e9ad9c5e
md"""
## Plot of multiple Fm(T) values

Plot the family of m values on a semilog-y plot. Very regular behavior:
"""

# ╔═╡ 3923126f-93be-48f2-bd37-3ef50588ff47
begin
	rs4 = 0:0.1:10
	plot(rs4,fboys.(0,rs4),yaxis=:log,legend=false)
	plot!(rs4,Fgamma.(0,rs4))
	plot!(rs4,fboys.(1,rs4))
	plot!(rs4,fboys.(2,rs4))
	plot!(rs4,fboys.(3,rs4))
	plot!(rs4,fboys.(4,rs4))
	plot!(rs4,fboys.(5,rs4))
	plot!(rs4,fboys.(6,rs4))
end

# ╔═╡ eb05a798-8eb2-4ff1-b88f-b58babe64891
denom(T,m) = 1/fboys(m,T);

# ╔═╡ 84df8774-a1f6-49f7-b84f-c1482b51dd9f
denomg(T,m) = 1/Fgamma(m,T);

# ╔═╡ 08c5dc9b-0cb1-4bea-bd97-4872f46e63bf
begin
	let rs4 = 0:0.1:5
	plot(rs4,denom.(0,rs4),yaxis=:log,legend=false)
	plot!(rs4,denom.(1,rs4))
	plot!(rs4,denom.(2,rs4))
	plot!(rs4,denom.(3,rs4))
	plot!(rs4,denom.(4,rs4))
	plot!(rs4,denom.(5,rs4))
	plot!(rs4,denom.(6,rs4))
	end
end

# ╔═╡ 347ca2e4-b932-4451-a12e-15ca7e23d052
md"""
## Interpolation
The other option that is appealing is to simply tabulate the values of the function
for the m-values that matter [0,12] on a relatively fine mesh and then interpolate
them either linearly or using Interpolate.jl.

Tasks:
- Express Fm(T) on a mesh
	- Decide on linear or log
	- Might have to be linear, since we definitely want values at T=0.
- Write a function to compute the maximum error with that discretization.
"""

# ╔═╡ b15df2d5-f407-4d2d-959f-18b21c8a3b60
ts = 0:0.01:100 # Intentially too coarse a mesh

# ╔═╡ ae817c36-d9b8-4d14-8a0e-91f83185d0ad
table_0 = fboys.(0,ts)

# ╔═╡ c615c015-34f6-441f-a285-0c7b8926660b
f0 = Spline1D(ts,fboys.(0,ts))

# ╔═╡ 85999318-1639-4677-bd9c-81723b4c809f
f0(0.5)

# ╔═╡ 61e753cb-f6c5-4301-8248-da8c21ba9186
fboys(0,0.5)

# ╔═╡ 24273947-c334-4508-a8e3-eab88100a745
function compute_error(steps)
	errs = []
	for step in steps
		mesh = 0:step:100
		f = Spline1D(mesh,fboys.(0,mesh))
		finemesh = 0:0.01:100
		error(x) = abs(fboys(0,x)-f(x))
		push!(errs,maximum(error(x) for x in finemesh))
	end
	return errs
end

# ╔═╡ 31c80276-35e5-4dc8-9a26-9600805a32da
steps = [0.05,0.1,0.2,0.5,1.0]

# ╔═╡ e3d4d6c7-c54b-408e-8575-331c7db3fec3
errs = compute_error(steps)

# ╔═╡ ca6f1d12-53c9-4df1-9194-aedb9480ec7c
md"Plot the dependence of the maximum error as a function of the 1d mesh spacing"

# ╔═╡ d73c0af0-559b-450a-b429-6d18991c839a
plot(steps,errs,yaxis=:log)

# ╔═╡ e0954971-c7a3-4847-a45a-ff22adc71490
md"Preliminary timing suggests that this is slow"

# ╔═╡ 88d466c8-1f3f-4c96-9507-18dc8cc123d1
@time f0(0.123)

# ╔═╡ 53659193-15dd-4bdc-820b-43422887b580
@time fboys(0,0.123)

# ╔═╡ 4161d94c-d545-4105-84e0-c92fd1344b05
md"## Implementing computation of a series of m-values for Fm(T)"

# ╔═╡ 44107e8f-116c-475f-b6d4-2b5aecc2f1e8
md"""
## References
[^boys50]: S.F. Boys. "Electronic Wave Functions. I. A General Method of Calculation for the Stationary States of Any Molecular System." *Proc. R. Soc. Lond.*, A200, 542-554 (1950).
[^gjp91]: P.M.W. Gill, B.G. Johnson, "Two-Electron Integrals Over Gaussian *s* Functions". J.A. Pople. *IJQC*. 40, 745-752, 1991.
[^libint]: E.F. Valeev. "Libint: A library for the evaluation of molecular integrals of many-body operators over Gaussian functions." [http://libint.valeyev.net/](http://libint.valeyev.net/)
[^mn57]: K. Nishimoto, N. Mataga. "Electronic Structure and Spectra of Some Nitrogen Heterocycles." *Z. Phys.*, 12, 335-338 (1957).
"""

# ╔═╡ Cell order:
# ╠═f80fd13d-207a-458a-b326-77256fe946e0
# ╠═bda0f8d4-5ab0-4c65-9d79-71ccbfe8dd3a
# ╟─4e16dbbe-b0fc-11eb-2349-ad7ccd0ff56b
# ╠═8a736133-17f2-446d-9ab5-4bc46bdd556b
# ╟─bf014265-210c-412f-9af4-5a654f557de9
# ╟─d8061039-367c-42ee-8a65-a32dca37fc67
# ╠═bc53663d-1f9f-43a6-8140-e8eeac1618bb
# ╟─1a0bffb2-8656-44db-b3b3-6cf6b2494b7c
# ╠═145a0007-b8d7-46a1-9cb8-ee063057527c
# ╠═b6b757ad-7fca-46c5-b42f-20f94ae4cbd7
# ╟─f6fa9869-b454-458d-95e3-818a9112cc52
# ╠═e7454913-e406-4b4a-856f-47103a7bf23d
# ╠═cdc2c6c9-e567-4ee6-a002-8fea660e35fa
# ╟─b5f42a47-eb6c-4b59-9827-a07b81618018
# ╟─105717a4-f481-479d-bfbb-dc35cc3c0997
# ╟─0b513de7-cca5-45ca-a4fa-2f2604eb7da9
# ╟─4dd0b1e6-61f0-4daa-80f1-625dc793b33a
# ╟─fdd2691f-f26f-4c26-8cc6-0d8edaa5278a
# ╠═a0a09c9f-c6d2-422f-b32a-8e3f8131b3c3
# ╠═ae9658a3-9241-423d-9daa-d6644e8a206a
# ╟─8c8adf9b-5eff-4942-9a49-95a4aa1e2002
# ╠═9428670f-5cf7-4c9e-bb93-dd0747a34182
# ╠═771d4eeb-72e6-4863-a114-ff203a42455b
# ╟─997e6018-bb0e-4933-8b13-b904e96fc819
# ╠═894135da-1ef9-4170-aa49-1b6307e398eb
# ╠═9dfb09bb-22d8-44f2-9c4b-98777e43de64
# ╠═d526caa1-ff7e-4ec5-ab59-2a10e9ad9c5e
# ╠═3923126f-93be-48f2-bd37-3ef50588ff47
# ╠═eb05a798-8eb2-4ff1-b88f-b58babe64891
# ╠═84df8774-a1f6-49f7-b84f-c1482b51dd9f
# ╠═08c5dc9b-0cb1-4bea-bd97-4872f46e63bf
# ╟─347ca2e4-b932-4451-a12e-15ca7e23d052
# ╠═b608f792-dd91-4bd2-97ab-baac7b7d2f9c
# ╠═b15df2d5-f407-4d2d-959f-18b21c8a3b60
# ╠═ae817c36-d9b8-4d14-8a0e-91f83185d0ad
# ╠═c615c015-34f6-441f-a285-0c7b8926660b
# ╠═85999318-1639-4677-bd9c-81723b4c809f
# ╠═61e753cb-f6c5-4301-8248-da8c21ba9186
# ╠═24273947-c334-4508-a8e3-eab88100a745
# ╠═31c80276-35e5-4dc8-9a26-9600805a32da
# ╠═e3d4d6c7-c54b-408e-8575-331c7db3fec3
# ╟─ca6f1d12-53c9-4df1-9194-aedb9480ec7c
# ╠═d73c0af0-559b-450a-b429-6d18991c839a
# ╟─e0954971-c7a3-4847-a45a-ff22adc71490
# ╠═88d466c8-1f3f-4c96-9507-18dc8cc123d1
# ╠═53659193-15dd-4bdc-820b-43422887b580
# ╠═4161d94c-d545-4105-84e0-c92fd1344b05
# ╠═44107e8f-116c-475f-b6d4-2b5aecc2f1e8


# ---
# title: 1420. Build Array Where You Can Find The Maximum Exactly K Comparisons
# id: problem1420
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/build-array-where-you-can-find-the-maximum-exactly-k-comparisons/description/>
# hidden: true
# ---
#
# Given three integers `n`, `m` and `k`. Consider the following algorithm to
# find the maximum element of an array of positive integers:
#
# ![](https://assets.leetcode.com/uploads/2020/04/02/e.png)
#
# You should build the array arr which has the following properties:
#
#   * `arr` has exactly `n` integers.
#   * `1 <= arr[i] <= m` where `(0 <= i < n)`.
#   * After applying the mentioned algorithm to `arr`, the value `search_cost` is equal to `k`.
#
# Return _the number of ways_ to build the array `arr` under the mentioned
# conditions. As the answer may grow large, the answer  **must be**  computed
# modulo `10^9 + 7`.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 2, m = 3, k = 1
#     Output: 6
#     Explanation: The possible arrays are [1, 1], [2, 1], [2, 2], [3, 1], [3, 2] [3, 3]
#
#
# **Example 2:**
#
#
#
#     Input: n = 5, m = 2, k = 3
#     Output: 0
#     Explanation: There are no possible arrays that satisify the mentioned conditions.
#
#
# **Example 3:**
#
#
#
#     Input: n = 9, m = 1, k = 1
#     Output: 1
#     Explanation: The only possible array is [1, 1, 1, 1, 1, 1, 1, 1, 1]
#
#
# **Example 4:**
#
#
#
#     Input: n = 50, m = 100, k = 25
#     Output: 34549172
#     Explanation: Don't forget to compute the answer modulo 1000000007
#
#
# **Example 5:**
#
#
#
#     Input: n = 37, m = 17, k = 7
#     Output: 418930126
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 50`
#   * `1 <= m <= 100`
#   * `0 <= k <= n`
#
#
## @lc code=start
using LeetCode

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


# code for Video 5.1:
julia> using RDatasets
julia> iris = dataset("datasets", "iris")
julia> features = convert(Array, iris[:, 1:4]) # measurements or independent variables
# 1) sepal_length, 2) sepal_width, 3) petal_length, 4) petal_width
julia> labels = convert(Array, iris[:, 5])     # targets or dependent variable, here the species name

# general application cycle of an algorithm:
train = # Array of length nrow(iris) and random true / false values
# then the training data set can be constructed as:
features[train]
labels[train]
# and the testing data set is then:
test = !train
features[test]
labels[test]

train  = Array(Bool, nrow(iris)) # this has the right dimension, but contains only false values

julia> using Distributions
julia> rand(Bernoulli(0.75), nrow(iris))
150-element Array{Int64,1}:
 1
 1
 0
 1
 1
 1
 1
 1
 1
 1
 1
 1
 ⋮

julia> train = rand(Bernoulli(0.75), nrow(iris)) .== 1
150-element BitArray{1}:
  true
  true
  true
  true
  true
  true
  true
 false
  true
 false
  true

# training and testing a model:
model = build_model(labels[train], features[train, :])
model = prune_model(model, 0.7)
print_tree(model))
predictions = apply_model(model, features[test,:])

### MLBase:
julia> Pkg.add("MLBase")
julia> using MLBase
## Data preprocessing:
julia> species = ["setosa", "virginica", "versicolor"]
3-element Array{ASCIIString,1}:
 "setosa"
 "virginica"
 "versicolor"

julia> arrs = repeach(species, [2,1,5])
8-element Array{ASCIIString,1}:
 "setosa"
 "setosa"
 "virginica"
 "versicolor"
 "versicolor"
 "versicolor"
 "versicolor"
 "versicolor"

julia> lm = labelmap(species)
LabelMap (with 3 labels):
[1] setosa
[2] virginica
[3] versicolor

julia> labelencode(lm, arrs)
8-element Array{Int64,1}:
 1
 1
 2
 3
 3
 3
 3
 3

julia> groupindices(lm, arrs)
3-element Array{Array{Int64,1},1}:
 [1,2]
 [3]
 [4,5,6,7,8]

# Performance evaluation:
julia> gt = ["setosa", "versicolor", "virginica", "setosa", "virginica"]
5-element Array{ASCIIString,1}:
 "setosa"
 "versicolor"
 "virginica"
 "setosa"
 "virginica"

julia> pred = ["setosa", "virginica", "setosa", "setosa", "virginica"]
5-element Array{ASCIIString,1}:
 "setosa"       # correct predition
 "virginica"    # error predition
 "setosa"       # error predition
 "setosa"       # correct predition
 "virginica"    # correct predition

julia> gtl = labelencode(lm, gt)
5-element Array{Int64,1}:
 1
 3
 2
 1
 2

julia> predl = labelencode(lm, pred)
5-element Array{Int64,1}:
 1
 2
 1
 1
 2

julia> correctrate(gtl, predl)
0.6

julia> errorrate(gtl, predl)
0.4

julia> confusmat(3, gtl, predl)
3x3 Array{Int64,2}:
 2  0  0
 1  1  0
 0  1  0

 ### Receiver Operating Characteristics (ROC)
julia> r = roc(gtl, predl)
MLBase.ROCNums{Int64}
  p = 5     # positive in ground-truth
  n = 0     # negative in ground-truth
  tp = 3    # true positives
  tn = 0    # true negatives
  fp = 0    # false positives
  fn = 0    # false negatives

true_positive_rate(r) # 0.6
true_negative_rate(r) # NaN
precision(r)          # 1.0
recall(r)             # 0.6

# for a real example: see section 2 and code file 5_2.jl

### Cross validation schemes
# for real examples: see section 2 and code file 5_2.j


# ---
# title: 368. Largest Divisible Subset
# id: problem368
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Dynamic Programming
# link: <https://leetcode.com/problems/largest-divisible-subset/description/>
# hidden: true
# ---
#
# Given a set of **distinct** positive integers, find the largest subset such
# that every pair (Si, Sj) of elements in this subset satisfies:
#
# Si % Sj = 0 or Sj % Si = 0.
#
# If there are multiple solutions, return any subset is fine.
#
# **Example 1:**
#
#
#
#     Input: [1,2,3]
#     Output: [1,2] (of course, [1,3] will also be ok)
#
#
# **Example 2:**
#
#
#
#     Input: [1,2,4,8]
#     Output: [1,2,4,8]
#
#
#
## @lc code=start
using LeetCode

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


## Generic and Fallbacks
"""
    discount(yc, to)
    discount(yc, from,to)

The discount factor for the yield curve `yc` for times `from` through `to`.
"""
discount(yc::T, from, to) where {T<:AbstractYieldCurve}= discount(yc, to) / discount(yc, from)

"""
    forward(yc, from, to, CompoundingFrequency=Periodic(1))

The forward `Rate` implied by the yield curve `yc` between times `from` and `to`.
"""
function forward(yc::T, from, to) where {T<:AbstractYieldCurve}
    return forward(yc, from, to, DEFAULT_COMPOUNDING)
end

function forward(yc::T, from, to, cf::CompoundingFrequency) where {T<:AbstractYieldCurve}
    r = Periodic((accumulation(yc, to) / accumulation(yc, from))^(1 / (to - from)) - 1, 1)
    return convert(cf, r)
end

function forward(yc::T, from) where {T<:AbstractYieldCurve}
    to = from + 1
    return forward(yc, from, to)
end

function CompoundingFrequency(curve::T) where {T<:AbstractYieldCurve}
    return DEFAULT_COMPOUNDING
end

"""
    par(curve,time;frequency=2)

Calculate the par yield for maturity `time` for the given `curve` and `frequency`. Returns a `Rate` object with periodicity corresponding to the `frequency`. The exception to this is if `time` is less than what the payments allowed by frequency (e.g. a time `0.5` but with frequency `1`) will effectively assume frequency equal to 1 over `time`.

# Examples

julia> c = Yields.Constant(0.04);

julia> Yields.par(c,4)
Yields.Rate{Float64, Yields.Periodic}(0.03960780543711406, Yields.Periodic(2))

julia> Yields.par(c,4;frequency=1)
Yields.Rate{Float64, Yields.Periodic}(0.040000000000000036, Yields.Periodic(1))

julia> Yields.par(c,0.6;frequency=4)
Yields.Rate{Float64, Yields.Periodic}(0.039413626195875295, Yields.Periodic(4))

julia> Yields.par(c,0.2;frequency=4)
Yields.Rate{Float64, Yields.Periodic}(0.039374942589460726, Yields.Periodic(5))

julia> Yields.par(c,2.5)
Yields.Rate{Float64, Yields.Periodic}(0.03960780543711406, Yields.Periodic(2))

"""
function par(curve, time; frequency=2)
    mat_disc = discount(curve, time)
    coup_times = coupon_times(time,frequency)
    coupon_pv = sum(discount(curve,t) for t in coup_times)
    Δt = step(coup_times)
    r = (1-mat_disc) / coupon_pv
    cfs = [t == last(coup_times) ? 1+r : r for t in coup_times]
    # `sign(r)`` is used instead of `1` because there are times when the coupons are negative so we want to flip the sign
    cfs = [-1;cfs]
    r = internal_rate_of_return(cfs,[0;coup_times])
    frequency_inner = min(1/Δt,max(1 / Δt, frequency))
    r = convert(Periodic(frequency_inner),r)
    return r
end

"""
    zero(curve,time)
    zero(curve,time,CompoundingFrequency)

Return the zero rate for the curve at the given time.
"""
function Base.zero(c::YC, time) where {YC<:AbstractYieldCurve}
     zero(c, time, CompoundingFrequency(c))
end

function Base.zero(c::YC, time, cf::C) where {YC<:AbstractYieldCurve,C<:CompoundingFrequency}
    df = discount(c, time)
    r = -log(df)/time
    return convert(cf, Continuous(r)) # c.zero is a curve of continuous rates represented as floats. explicitly wrap in continuous before converting
end

"""
    accumulation(yc, from, to)

The accumulation factor for the yield curve `yc` for times `from` through `to`.
"""
function accumulation(yc::AbstractYieldCurve, time)
    return 1 ./ discount(yc, time)
end

function accumulation(yc::AbstractYieldCurve, from, to)
    return 1 ./ discount(yc, from, to)
end

"""
    Par(rates, maturities=eachindex(rates)
    Par(p::YieldCurveFitParameters, rates, maturities=eachindex(rates)

Construct a curve given a set of bond equivalent yields and the corresponding maturities. Assumes that maturities <= 1 year do not pay coupons and that after one year, pays coupons with frequency equal to the CompoundingFrequency of the corresponding rate (normally the default for a `Rate` is `1`, but when constructed via `Par` the default compounding Frequency is `2`).

See [`bootstrap`](@ref) for more on the `interpolation` parameter, which is set to `QuadraticSpline()` by default.

# Examples

```julia-repl

julia> par = [6.,8.,9.5,10.5,11.0,11.25,11.38,11.44,11.48,11.5] ./ 100
julia> maturities = [t for t in 1:10]
julia> curve = Par(par,maturities);
julia> zero(curve,1)
Rate(0.06000000000000005, Periodic(1))

```
"""
function Yields.Par(rates,maturities=eachindex(rates))
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))
    return Yields.Par(Bootstrap(),rates,maturities)
end

"""
    Forward(rates,maturities=eachindex(rates))
    Forward(p::YieldCurveFitParameters, rates,maturities=eachindex(rates))

Takes a vector of 1-period forward rates and constructs a discount curve. The method of fitting the curve to the data is determined by the [`YieldCurveFitParameters`](@ref) object `p`, which is a `Boostrap(QuadraticSpline())` by default.

If `rates` is a vector of floating point number instead of a vector `Rate`s, see the [`YieldCurveFitParameters`](@ref) for how the rate will be interpreted.

# Examples

```julia-repl
julia> Yields.Forward( [0.01,0.02,0.03] );

julia> Yields.Forward( Yields.Continuous.([0.01,0.02,0.03]) );

```
"""
function Yields.Forward(rates,maturities=eachindex(rates))
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))
    return Yields.Forward(Bootstrap(),rates,maturities)
end

"""
    Zero(rates, maturities=eachindex(rates))
    Zero(p::YieldCurveFitParameters,rates, maturities=eachindex(rates))

Construct a yield curve with given zero-coupon spot `rates` at the given `maturities`. The method of fitting the curve to the data is determined by the [`YieldCurveFitParameters`](@ref) object `p`, which is a `Boostrap(QuadraticSpline())` by default.

If `rates` is a vector of floating point number instead of a vector `Rate`s, see the [`YieldCurveFitParameters`](@ref) for how the rate will be interpreted.

# Examples

```julia-repl
julia> Yields.Zero([0.01,0.02,0.04,0.05],[1,2,5,10])

              ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀Yield Curve (Yields.BootstrapCurve)
              ┌────────────────────────────────────────────────────────────┐
         0.05 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠤⠒⠒⠒⠒⠒⠤⠤⠤⢄⣀⣀⣀⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ Zero rates
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠖⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠉⠉⠉⠒⠒⠒⠢⠤⠤⠤⣄⣀⣀⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠉⠉⠑⠒⠒⠒⠦⠤⠤⠤⣀⣀⣀⡀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⡔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠉│
              │⠀⠀⠀⠀⠀⠀⠀⢠⠎⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⢠⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⢠⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
   Continuous │⠀⠀⠀⠀⢀⠇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⡜⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⡸⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⢠⠇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⡜⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠒⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
            0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              └────────────────────────────────────────────────────────────┘
              ⠀0⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀time⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀30

"""
function Yields.Zero(rates,maturities=eachindex(rates))
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))
    return Yields.Zero(Bootstrap(),rates,maturities)
end

"""
    Yields.CMT(rates, maturities; interpolation=QuadraticSpline())
    Yields.CMT(p::YieldCurveFitParameters,rates, maturities)

Takes constant maturity (treasury) yields (bond equivalent), and assumes that instruments <= one year maturity pay no coupons and that the rest pay semi-annual.

The method of fitting the curve to the data is determined by the [`YieldCurveFitParameters`](@ref) object `p`, which is a `Boostrap(QuadraticSpline())` by default.

# Examples

```
# 2021-03-31 rates from Treasury.gov
rates =[0.01, 0.01, 0.03, 0.05, 0.07, 0.16, 0.35, 0.92, 1.40, 1.74, 2.31, 2.41] ./ 100
mats = [1/12, 2/12, 3/12, 6/12, 1, 2, 3, 5, 7, 10, 20, 30]

Yields.CMT(rates,mats)
```
"""
function Yields.CMT(rates,maturities=eachindex(rates))
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))
    return Yields.CMT(Bootstrap(),rates,maturities)
end

"""
    OIS(rates, maturities)
    OIS(p::YieldCurveFitParameters, rates, maturities)

Takes Overnight Index Swap rates, and assumes that instruments <= one year maturity are settled once and other agreements are settled quarterly with a corresponding CompoundingFrequency.

The method of fitting the curve to the data is determined by the [`YieldCurveFitParameters`](@ref) object `p`, which is a `Boostrap(QuadraticSpline())` by default.

# Examples
```
julia> ois = [1.8, 2.0, 2.2, 2.5, 3.0, 4.0] ./ 100;
julia> mats = [1 / 12, 1 / 4, 1 / 2, 1, 2, 5];
julia> curve = Yields.OIS(ois, mats)

              ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀Yield Curve (Yields.BootstrapCurve)
              ┌────────────────────────────────────────────────────────────┐
          0.1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ Zero rates
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
   Continuous │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⠀⠀⠀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⠀⠀⣀⠔⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠀⡠⠎⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              │⠜⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
         0.01 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
              └────────────────────────────────────────────────────────────┘
              ⠀0⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀time⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀30
```
"""
function Yields.OIS(rates,maturities=eachindex(rates))
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))
    return Yields.OIS(Bootstrap(),rates,maturities)
end


using StructuralCausalModels, Test

ProjDir = @__DIR__
cd(ProjDir) #do

d = OrderedDict(
  :w => :s,
  :d => [:a, :w, :m],
  :m => [:a, :s],
  :a => :s
);

dag = DAG("test_open_paths_01", d);
fname = joinpath(ProjDir, "test_open_paths_01.dot")
to_graphviz(dag, fname)
#Sys.isapple() && run(`open -a GraphViz.app $(fname)`)

to_dagitty(d)

bs = basis_set(dag)
ap  = all_paths(dag, :w, :d)
bp = backdoor_paths(dag, ap, :w)

cs = Vector{Symbol}[]
syms = syms_in_paths(bp, :w, :d)

function adjsets(d::DAG, paths, syms)
  lsyms = deepcopy(syms)
  adjustment_sets = Vector{Symbol}[]
  for s in lsyms
    if sym_in_all_paths(bp, s)
      if length(open_paths(dag, bp, [s])) == 0
        #println("$s closes all paths")
        append!(adjustment_sets, [[s]])
        setdiff!(lsyms, [s])
      end
    end
  end
  len = 2
  local csyms
  while length(lsyms) >= len
    csyms = collect(combinations(lsyms, len))
    for s in csyms
      if length(open_paths(dag, bp, s)) == 0
        append!(adjustment_sets, [s])
        setdiff!(lsyms, s)
      end
    end
    len += 1
  end
  adjustment_sets
end

@testset "Open_path_01" begin

  adjs = adjsets(dag, bp, syms)
  @test adjs == adjustment_sets(dag, :w, :d)

end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: -all
#     formats: jl:light,notebooks//ipynb
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Julia 1.5.1
#     language: julia
#     name: julia-1.5
# ---

# # Example Julia
#
# ## Documentation de fonction

"""
    bspline(p, j, x)

Return the value at x in [0,1[ of the B-spline with integer nodes of degree p with support starting at j.
Implemented recursively using the [De Boor's Algorithm](https://en.wikipedia.org/wiki/De_Boor%27s_algorithm)

```math
B_{i,0}(x) := \\left\\{
\\begin{matrix}
1 & \\mathrm{if}  \\quad t_i ≤ x < t_{i+1} \\\\
0 & \\mathrm{otherwise}
\\end{matrix}
\\right.
```

```math
B_{i,p}(x) := \\frac{x - t_i}{t_{i+p} - t_i} B_{i,p-1}(x)
+ \\frac{t_{i+p+1} - x}{t_{i+p+1} - t_{i+1}} B_{i+1,p-1}(x).
```
"""
function bspline(p::Int, j::Int, x::Float64)
   if p == 0
       if j == 0
           return 1.0
       else
           return 0.0
       end
   else
       w = (x - j) / p
       w1 = (x - j - 1) / p
   end
   return (w * bspline(p - 1, j, x) + (1 - w1) * bspline(p - 1, j + 1, x))
end

# Switch the next cell from `Markdown` to `Code` format

# ?bspline

# +
using Random, Plots

function generate_data( n = 2000, seed = 1234 )
    seuil = 0.25
    rng = MersenneTwister(seed)
    X1 = rand( rng, n)
    X2 = rand( rng, n)
    U  = rand( rng, n)
    Y  = zeros(Int,n)
    Y[(X1 .<= 0.25) .& (U  .<= seuil)] .= 1
    Y[(X1 .>  0.25) .& (X2 .>= 0.75) .& (U .<= seuil)] .= 1
    Y[(X1 .>  0.25) .& (X2 .<  0.75) .& (U .>  seuil)] .= 1
    return X1, X2, Y
end

X1, X2, Y = generate_data()
scatter(X1,X2, marker_z = Y)
# -
using DataFrames, StatsPlots

data = DataFrame( X1=X1, X2=X2, Y=Y)
head(data)

@df data scatter(:X1,:X2, zcolor= :Y, xaxis = "X1", yaxis="X2", lab="Y")

"""
[x1,x2,x1^2,x1x2,x2^2……x2^6]
"""
function map_features(X1,X2)
    degree = 6
    out = ones(size(X1[:,1]))
    for i=1:6
        for j=0:i
            out = hcat(out,(X1.^(i-j)).*(X2.^j))
        end
    end
    return out
end

X = map_features(X1,X2)