using SDC
import Plots
using LaTeXStrings
using PrettyTables
"""
Currently a demonstration of SDC solving
"""
function main()

    # 5-node quadrature
    n = 5
    # IMEX Euler, half explicit, half implicit (dalhquist equation)
    fE(t, u) = 0.5*u
    fI(t, u) = 0.5*u
    # Solve implicit euler analytically
    implicit_solver(α, t, b) = b/(1+α)

    # Create SDC timestepper with provided parameters (should change code
    # structure to avoid this)
    my_SDC = make_SDC(n, fE, fI, implicit_solver)

    # Analytic solution
    u(t) = exp(t)
    # Start time
    t = 0.0
    # Initial value
    u_n = u(t)
    # Timesteps
    Δtimes = [1.0, 1.0, 1.0, 1.0, 1.0]
    # Non-start times, derived from timesteps
    times = cumsum(Δtimes) .+ t

    # Perform timestepping
    u_actual = []
    u_approx = []
    for i in eachindex(Δtimes)
        t += Δtimes[i]
        u_n = my_SDC(u_n, t, Δtimes[i])
        push!(u_actual, u(t))
        push!(u_approx, u_n)
    end

    # Calculate errors
    u_diff_abs = abs.(u_actual - u_approx)
    u_diff_rel = [u_diff_abs[i]/u_actual[i] for i in eachindex(u_diff_abs)]

    # Produce table of results
    table_data = hcat(times, u_actual, u_approx, u_diff_abs, u_diff_rel)
    header = ["t", "u(t)", "est u(t)", "Abs Err", "Rel Err"]
    pretty_table(table_data; header=header)

    # Produce graph of results
    p1 = Plots.plot(times, u_diff_abs, yaxis=:log, title = "Absolute Error of e^t, t_0=0, variable timesteps",
                   legend=false)
    Plots.xlabel!(p1, "Time")
    Plots.ylabel!(p1, "Absolute Error")
    p2 = Plots.plot(times, u_diff_rel, yaxis=:log, title = "Relative Error of e^t, t_0=0, variable timesteps",
                   legend=false)
    Plots.xlabel!(p2, "Time")
    Plots.ylabel!(p2, "Relative Error")
    plt = Plots.plot(p1, p2, layout=(2,1))
    Plots.savefig(plt, "./dahlquist_test.png")
    return plt
end

display(main())


convert_point(p::Vector, m::Matrix) = (([p, 1]*m)/m[3,3])[1:2]

M(s) = Diagonal([1, 1, 1/s]);

V(Θ) = [cos(Θ)  -sin(Θ) 0;
        sin(Θ)  cos(Θ)  0;
        0       0       1]

P(Δx, Δy) =[1   0   Δx;
            0   1   Δy;
            0   0   1]

size = 300;
models = Set{Matrix}();

function create_image(mvp::Matrix)
    mvp[3, 3] > size && return
    for u in [M(1/2)*V(π/6)*P(0.2, 0.2)]
        create_image(mvp*u)
    end
    push!(models, mvp)
end

function draw()
    for m in models

    end
end


function Jtriple(x::Vector{Float64},beta::Vector{Float64},rho::Float64)
    J = zeros(8,8);

    J[1,1] = 1.;
    J[2,1] = rho.*x[1]./(x[2].^2);
    J[3,1] = 1.;
    J[4,1] = 1.;
    J[5,1] = 1./x[2];

    J[2,2] = (0.5.*beta[1].*x[2].^-0.5 .- rho.*x[1].^2.*x[2].^-3);
    J[3,2] = J[2,2];
    J[4,2] = J[2,2];
    J[5,2] = -x[1].*x[2].^-2 .+ sqrt.(0.5.*beta[1]./rho).*x[2].^-0.75;

    J[1,3] = -1.;
    J[2,3] = -rho.*x[3]./(x[4].^2);
    J[6,3] = 1./x[4];

    J[2,4] = (-0.5.*beta[2].*x[4].^-0.5 .+ rho.*x[3].^2.*x[4].^-3);
    J[6,4] = -x[3].*x[4].^-2 .- sqrt.(0.5.*beta[2]./rho).*x[4].^-0.75;

    J[1,5] = -1.;
    J[3,5] = -rho.*x[5]./(x[6].^2);
    J[7,5] = 1./x[6];

    J[3,6] = (-0.5.*beta[3].*x[6].^-0.5 .+ rho.*x[5].^2.*x[6].^-3);
    J[7,6] = -x[5].*x[6].^-2 .- sqrt.(0.5.*beta[3]./rho).*x[6].^-0.75;

    J[1,7] = -1.;
    J[4,7] = -rho.*x[7]./(x[8].^2);
    J[8,7] = 1./x[8];

    J[4,8] = (-0.5.*beta[4].*x[8].^-0.5 .+ rho.*x[7].^2.*x[8].^-3);
    J[8,8] = -x[7].*x[8].^-2 .- sqrt.(0.5.*beta[4]./rho).*x[8].^-0.75;

    return J
end


#=
spatial:
- Julia version:
- Author: Dan
- Date: 2021-10-05
=#
export ramachandranangles, dihedralangle, phiangle, phiangles, psiangle, psiangles, rmsd, distance

ramachandranangles(kern::KernelInterface) = begin
    phiangles(kern), psiangles(kern)
end

ramachandranangles(kern::DataFrameSystem) = begin
    phiangles(kern), psiangles(kern)
end

"""
    dihedralangle(atom_a, atom_b, atom_c, atom_d)
    dihedralangle(vec_ab, vec_bc, vec_cd)
Calculate the dihedral angle in radians defined by four `AbstractAtom`s or
three vectors.
The angle between the planes defined by atoms (A, B, C) and (B, C, D) is
returned in the range -π to π.
Credit to this function goes to [BioStructures](https://doi.org/10.1093/bioinformatics/btaa502)
"""
dihedralangle(at_a::Union{DataFrameRow,AbstractAtom}, at_b::Union{DataFrameRow,AbstractAtom},
                at_c::Union{DataFrameRow,AbstractAtom}, at_d::Union{DataFrameRow,AbstractAtom}) =
    dihedralangle(
        coords(at_b) - coords(at_a),
        coords(at_c) - coords(at_b),
        coords(at_d) - coords(at_c)
    )

dihedralangle(vec_a::SVector{3,Float64}, vec_b::SVector{3,Float64}, vec_c::SVector{3,Float64}) =
    return atan(
        dot(cross(cross(vec_a, vec_b), cross(vec_b, vec_c)), vec_b / norm(vec_b)),
        dot(cross(vec_a, vec_b), cross(vec_b, vec_c)))

phiangle(prev_res::Union{T, DataFrameRow}, res::Union{T, DataFrameRow}) where T<: AbstractResidue = begin

    if !sequentialResidues(prev_res, res)
        return NaN
    end

    children = getChildren(res)
    child_names_res::Vector{String} = [x.name_ for x in children]
    child_names_prev_res::Vector{String} = [x.name_ for x in getChildren(prev_res)]
    if !in("C", child_names_prev_res) && !all(in(x, child_names_res) for x in ["C", "CA","N"])
        return NaN
    else
        c_atom_prev::Union{AbstractAtom, DataFrameRow} = getChildren(prev_res)[findfirst(x -> x == "C", child_names_prev_res)]
        n_atom::Union{AbstractAtom, DataFrameRow} =  children[findfirst(x -> x == "N", child_names_res)]
        ca_atom::Union{AbstractAtom, DataFrameRow} = children[findfirst(x -> x == "CA", child_names_res)]
        c_atom::Union{AbstractAtom, DataFrameRow} =  children[findfirst(x -> x == "C", child_names_res)]
        return dihedralangle(c_atom_prev, n_atom, ca_atom, c_atom)
    end
end

phiangles(kern::KernelInterface) = begin
    res_list = collectResidues(kern)
    if length(res_list) < 2
        throw(ArgumentError("At least 2 residues required to calculate dihedral angles"))
    end

    phi_angles = Float64[NaN]

    for i in 2:length(res_list)
        res = res_list[i]
        res_prev = res_list[i - 1]
        push!(phi_angles, phiangle(res_prev, res))
    end

    return phi_angles
end

phiangles(kern::DataFrameSystem) = begin
    res_list = eachResidue(kern)
    if length(res_list) < 2
        throw(ArgumentError("At least 2 residues required to calculate dihedral angles"))
    end

    phi_angles = Float64[NaN]

    for i in 2:length(res_list)
        res = res_list[i]
        res_prev = res_list[i - 1]
        push!(phi_angles, phiangle(res_prev, res))
    end

    return phi_angles
end

psiangle(res::Union{T, DataFrameRow}, res_next::Union{T, DataFrameRow}) where T<: AbstractResidue = begin

    if !sequentialResidues(res, res_next)
        return NaN
    end

    children = getChildren(res)
    child_names_res::Vector{String} = [x.name_ for x in children]
    child_names_next_res::Vector{String} = [x.name_ for x in getChildren(res_next)]
    if !in("N", child_names_next_res) && !all(in(x, child_names_res) for x in ["C", "CA","N"])
        return NaN
    else
        n_atom::Union{AbstractAtom, DataFrameRow} = children[findfirst(x -> x == "N", child_names_res)]
        n_atom_next::Union{AbstractAtom, DataFrameRow} = getChildren(res_next)[findfirst(x -> x == "N", child_names_next_res)]
        ca_atom::Union{AbstractAtom, DataFrameRow} = children[findfirst(x -> x == "CA", child_names_res)]
        c_atom::Union{AbstractAtom, DataFrameRow} = children[findfirst(x -> x == "C", child_names_res)]
        return dihedralangle(n_atom, ca_atom, c_atom, n_atom_next)
    end
end

psiangles(kern::KernelInterface) = begin
    res_list = collectResidues(kern)
    if length(res_list) < 2
        throw(ArgumentError("At least 2 residues required to calculate dihedral angles"))
    end

    psi_angles = Float64[]

    for i in 1:length(res_list) -1
        res = res_list[i]
        res_next = res_list[i + 1]
        push!(psi_angles, psiangle(res, res_next))
    end

    push!(psi_angles, NaN)
    return psi_angles
end

rmsd(dfs1::DataFrameSystem, dfs2::DataFrameSystem; calphaonly=true) = begin
    if calphaonly
        atoms1 = filter(x -> getName(x) == "CA", dfs1.atoms, view=true)
        atoms2 = filter(x -> getName(x) == "CA", dfs2.atoms, view=true)
    else
        atoms1 = dfs1.atoms
        atoms2 = dfs2.atoms
    end
    if size(atoms1)[1] != size(atoms2)[1]
        throw(ArgumentError("Inputs are of size $(length(atoms1)) and " *
                            "$(length(atoms2)) but must be the same for RMSD"))
    end
    diff = getPosition(atoms1) - getPosition(atoms2)
    return sqrt.(dot(diff,diff) / length(diff))
end

rmsd(comp1::AbstractComposite, comp2::AbstractComposite,calphaonly=true) = begin
    if calphaonly
        atoms1 = filter(x -> getName(x) == "CA", eachAtom(comp1))
        atoms2 = filter(x -> getName(x) == "CA", eachAtom(comp2))
    else
        atoms1 = eachAtom(comp1)
        atoms2 = eachAtom(comp2)
    end
    if length(atoms1) != length(atoms2)
        throw(ArgumentError("Inputs are of size $(length(atoms1)) and " *
                            "$(length(atoms2)) but must be the same for RMSD"))
    end
    diff = getPosition(atoms1) - getPosition(atoms2)
    return sqrt.(dot(diff,diff) / length(diff))
end

distance(el1::Union{KernelInterface,AbstractDataFrame, DataFrameRow},
    el2::Union{KernelInterface,AbstractDataFrame, DataFrameRow}) = begin
    return sqrt(squaredDistance(el1,el2))
end

squaredDistance(comp1::Union{KernelInterface,AbstractDataFrame, DataFrameRow},
    comp2::Union{KernelInterface,AbstractDataFrame, DataFrameRow}) = begin
    atoms1 = collect(eachAtom(comp1))
    atoms2 = collect(eachAtom(comp2))
    min_dist = Inf
    for i = 1:length(atoms1)
        for j = i:length(atoms2)
            cur_dist::Float64 = sum( ( getPosition(atoms1[i]) - getPosition(atoms2[j]) ).^2)
            min_dist = min(min_dist, cur_dist)
        end
    end
    return min_dist
end


@inline symmetric(X::AbstractMatrix) = Symmetric(X)
@inline symmetric(X::Diagonal) = X

diag_Xt_invA_X(A::Cholesky, X::AbstractVecOrMat) = AbstractGPs.diag_At_A(A.U' \ X)

Xt_invA_X(A::Cholesky, x::AbstractVector) = sum(abs2, A.U' \ x)

function Xt_invA_X(A::Cholesky, X::AbstractMatrix)
    V = A.U' \ X
    return Symmetric(V'V)
end

function diag_At_B(A::AbstractVecOrMat, B::AbstractVecOrMat)
    @assert size(A) == size(B)
    return vec(sum(A .* B; dims=1))
end


function compensated_summation(x::T, δ::T, ε::T) where {T}
    a = x
    ε = ε + δ
    x = a + ε
    ε = ε + (a - x)

    return (x, ε)
end

function compensated_summation(x::VT, δ::VT, ε::VT) where {T, VT <: AbstractVector{T}}
    a  = copy(x)
    ε  = ε .+ δ
    x̄  = a .+ ε
    ε̄ = ε .+ (a .- x̄)

    return (x̄, ε̄)
end

function compensated_summation!(x::VT, δ::VT, ε::VT) where {T, VT <: AbstractVector{T}}
    a  = copy(x)
    ε  = ε .+ δ
    x .= a .+ ε
    ε .= ε .+ (a .- x)

    return (x, ε)
end


# plot the distribution at one single point x0
using PyPlot
function plot_distribution_single(distribution_z0, x0, z0_true)
    pdf0 = distribution_z0.pdf
    mean = distribution_z0.mean
    # median = distribution_z0.median
    # CI = distribution_z0.CI
    # CI_left = CI[2]
    # CI_right = CI[3]

    # determine the range of plotting
    kmax = 20
    tol = 1e-3
    b0 = 5.0
    b, iter = zero_finding(pdf0, kmax, tol, b0)

    # plot pdf
    z_grid = range(1e-5, stop = b, step = 0.01)
    p_grid = pdf0.(z_grid);
    PyPlot.plot(z_grid, p_grid, label = "probability density function")
    # plot mean median and true value
    PyPlot.vlines(mean, 0, pdf0(mean), label = "mean", colors = "k")
    # PyPlot.vlines(median, 0, pdf0(median), label = "median",  colors = "b")
    PyPlot.vlines(z0_true, 0, pdf0(z0_true), label = "true value",  colors = "r")
    # plot CI
    # CI_x_range = range(CI_left, stop = CI_right, step = 0.01)
    # CI_y_range = pdf0.(CI_x_range)
    # PyPlot.fill_between(CI_x_range, 0, CI_y_range, alpha = 0.3, label = "95% confidence interval")
    PyPlot.legend()
    PyPlot.grid()
    PyPlot.title("Predicted Distribution at x = $x0")
#     PyPlot.savefig("Figure/Prediction")
end


function _bioclim_score(x)
    if isnothing(x)
        return nothing
    end
    @assert 0. <= x <= 1.
    if x >= 0.5
        return 2(1-x)
    else
        return 2x
    end
end

function bioclim(layer::T, records::GBIFRecords) where {T <: SimpleSDMLayer}
    bioclim_prediction = convert(SimpleSDMResponse, copy(layer))
    qf = StatsBase.ecdf(Float32.(filter(!isnothing, layer[records])))
    qfgrid = replace(x -> isnothing(x) ? nothing : qf(x), layer.grid)
    bioclim_prediction.grid = _bioclim_score.(qfgrid)
    return bioclim_prediction
end


__precompile__()
module eps_InAs_ptype
using QuadGK
export epsFCL_ptype,Gamma_ptype,m_star_ptype
#from integration import Integrate
#from scipy.integrate import quad
 #but I can't use QuadGk for larger values of eV, can only functionine function Integrate, tested E = 0.0037eV, get same eps of 10.240628772514803 + 0.03160036047883907im

const e = 1.602*^(10,-19)
const eV = 1.60218*^(10,-19)
const eps_0 = 8.8542*^(10,-12)
const c = 3*^(10,8)
const hbar = 1.05457*^(10,-34) #m2kg/s
const hbar_eV = 6.5821*^(10,-16) #eV s
const m_e =9.109*^(10,-31)#kg
const m_0 = 0.9109*^(10,-30)
const m_0_eV = 0.51099906*^(10,6) #eV

const E0_InAs = 0.42
const dt0_InAs = 0.28*0.001
const b0_InAs = 93
const A_InAs = 14*10^(-9) #for n-type  InAs
const B_InAs = 1.97*10^(-7)
const C_InAs = 57.9*10^(-12) #N in m^(-3)

const P_square = 11.9   #eV
const E_g_InAs= 0.416 #eV

const umin =20.0
const umax = 530.0
const Nref = 1.1*10^17*10^6 #in m^(-3) since N_base in m^(-3)
const phi = 0.46
const t1 = 2.3
const t2 = 3.0

@inline function m_star_ptype(n) # change m_e according to MB shift
    term = 8.0*P_square/eV*hbar^2*(3.0*pi^2*n)^(2/3)/(3.0*m_e*E_g_InAs^2)
    return (1.0+(4.0*P_square/(3.0*E_g_InAs))*(1.0+term)^(-1/2))^(-1)*m_0
end
precompile(m_star_ptype, (Float64,))

@inline function Gamma_ptype(N_Base,T)
    mu_franc =  umin + (umax*(300.0/T)^t1-umin)/(1+(N_Base/(Nref*(T/300.0)^t2))^phi)#from Francoeur
    return e/(m_star_ptype(N_Base)*mu_franc*10000.0) #s^(-1) change mu from cm^(-2) to m^(-2)
end
precompile(Gamma_ptype, (Float64,Float64))

const eps_inf_InAs = 12.3
const g = 2.89*10.0^11  #s^(-1), gamma
const o_TO = 4.12*10.0^(13) #s^(-1)
const o_LO = 4.58*10.0^(13)
@inline function epsFCL_ptype(omega,mstar,gamma,N_base)
    o_p_square=N_base*e^2/(12.3*eps_0*mstar) #s^(-2)
    return eps_inf_InAs*(1.0+(o_LO^2-o_TO^2)/(o_TO^2-omega^2-im*omega*g)-o_p_square/(omega*(omega+ im*gamma)))
end
precompile(epsFCL_ptype,(Float64,Float64,Float64,Float64))
end
#@inline function epsIB(omega,N0,T,E0_T_InAs_value) #use the function from ntype
#    if(hbar_eV*omega<=E0_T_InAs_value)
#        alpha_lb_2=0.0
#    else
#        alpha_lb_2= (3*10^8*(hbar_eV*omega-E0_T_InAs_value))^(1/2)*100
#    end
#    m_IB_pp = alpha_lb_2*c/(2*omega)

#    integratedfn(omega_p) = alpha_lb_2(omega_p,N0,T)/(omega^2-omega_p^2)
#    if(hbar_eV*omega<E0_T_InAs_value)
#        m_IB_p_value =1.0
#    else
     #   result  =1.0 + (c/pi)*quadgk(integratedfn,100.0,10.0^25,rtol=1e-8)[1] # result  =(1.0 + (c/pi)*Integrate(integratedfn, 10000,100.0,10.0^25))
#        m_IB_p_value = (1.0 + 0.00026172525)
#    end
#    return (m_IB_p_value^2-m_IB_pp^2) + im*2*m_IB_p_value*m_IB_pp
#end
#precompile(epsIB,(Float64,Float64,Float64,Float64))

#function eps_InAs_ntype(E_photon::Float64,N_base::Float64)::ComplexF64 #N_base in m^(-3)
    #E_photon= eVtoOmega(omega)
#    omega = eVtoOmega(E_photon)
#    T=300.0 #o_angular = 5.38*10^14   #hbar*5.38*10^14/eV = 0.35eV
#    if(E_photon>0.35)
#        return epsIB(omega,N_base,T) +  epsFCL(omega,N_base,T)
#    else
#        return epsFCL(omega,N_base,T)
#    end
#end

#omegalist= 2.0*10.0^14:5.75*10.0^13:2.5*10.0^15#
#eVlist=  0.2:0.02:1.5
#epslist = []
#for ind = 1:length(eVlist)
#    eps = eps_InAs_ntype(eVlist[ind],10.0^23)
#    println("eps_InAs",eps)
#    push!(epslist, imag(eps))
#end
#plot(eVlist,epslist)
#savefig("test_InAs_real.png")


function partitioned_matrices(n)

  reps = parse(Int64, ENV["LAMP_REPS"])

  m = Int(floor(n / 2))
  A1 = randn(m, m)
  A2 = randn(m, m)
  B = randn(2*m, 2*m)

  Benchmarker.add_data(csv, "compact", Benchmarker.measure(reps, blocked_solve_naive, A1, A2, B))
  Benchmarker.add_data(csv, "blocked", Benchmarker.measure(reps, blocked_solve_recommended, A1, A2, B))

  @test isapprox(blocked_solve_naive(A1, A2, B)[2], blocked_solve_recommended(A1, A2, B)[2])
end

function blocked_solve_naive(A1::DenseArray{Float64, 2}, A2::DenseArray{Float64, 2},B::DenseArray{Float64, 2})
  result = zeros(size(B))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = [A1 zeros(size(A1,1), size(A1, 2)); zeros(size(A2,1), size(A2, 2)) A2]\B

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function blocked_solve_recommended(A1::DenseArray{Float64, 2}, A2::DenseArray{Float64, 2},B::DenseArray{Float64, 2})
  result = zeros(size(B))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  b1 = @view B[1:size(A1, 1), 1:end] # do not use view with parenthesis
  b2 = @view B[size(A1, 1)+1:end, 1:end]
  result = [A1\b1; A2\b2]

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end


let
    n = 50
    p = 3
    t = (1:n)/n
    Ut,Vt = egrss.generators(t,p);

    Wt = egrss.potrf(Ut,Vt)
    L = LowerTriangular(tril(Ut'*Wt))
    @test ( isapprox(egrss.full_tril(Ut,Wt),L) )

    Wt,c = egrss.potrf(Ut,Vt,1e-3*ones(n))
    L = LowerTriangular(tril(Ut'*Wt,-1)+Diagonal(c[:]))
    @test ( isapprox(egrss.full_tril(Ut,Wt,c),L) )
end


# Separable sum, using slices of an array as variables

immutable SlicedSeparableSum{S <: AbstractArray, T <: AbstractArray} <: ProximableFunction
	fs::S
	is::T
	dim::Integer

function SlicedSeparableSum(a::S, b::T, dim::Integer)
		if size(a) != size(b)
			error("size(fs) must coincide with size(is)")
		else
			new(a, b, dim)
		end
	end
end

SlicedSeparableSum{S <: AbstractArray, T <: AbstractArray}(a::S, b::T, dim::Integer=1) =
	SlicedSeparableSum{S, T}(a, b, dim)

function SlicedSeparableSum{T <: Tuple}(ps::AbstractArray{T}, dim::Integer=1)
	a = Array{ProximableFunction}(length(ps))
	b = Array{AbstractArray}(length(ps))
	for i in eachindex(ps)
		a[i] = ps[i][1]
		b[i] = ps[i][2]
	end
	SlicedSeparableSum{typeof(a),typeof(b)}(a, b, dim)
end

function SlicedSeparableSum{P <: Pair}(p::AbstractArray{P}, dim::Integer = 1)
	a = Vector{ProximableFunction}(length(p))
	b = Vector{AbstractArray}(length(p))
	for i in eachindex(p)
		a[i] = p[i].first
		b[i] = p[i].second
	end
	return SlicedSeparableSum(a, b, dim)
end

function (f::SlicedSeparableSum{S}){S <: AbstractArray, T <: AbstractArray}(x::T)
	sum = 0.0
  for k in eachindex(f.fs)
		sum += f.fs[k](x[is[k]])
	end
	return sum
end

function prox!{T <: RealOrComplex}(f::SlicedSeparableSum, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  v = 0.0
  nd = ndims(x)

  for i in eachindex(f.fs)
	  z = slicedim(y,f.dim,f.is[i])
	  g = prox!(f.fs[i],slicedim(x,f.dim,f.is[i]),z,gamma)
	  y[[ n==f.dim ? f.is[i] : indices(y,n) for n in 1:nd ]...] = z
	  v += g
  end
  return v

end

fun_name(f::SlicedSeparableSum) = "sliced separable sum"
function fun_dom(f::SlicedSeparableSum)
	# s = ""
	# for k in eachindex(f.fs)
	# 	s = string(s, fun_dom(f.fs[k]), " × ")
	# end
	# return s
	return "n/a" # for now
end
fun_expr(f::SlicedSeparableSum) = "hard to explain"
function fun_params(f::SlicedSeparableSum)
	# s = "r = $(size(f.fs)), "
	# for k in eachindex(f.fs)
	# 	s = string(s, "f[$(k)] = $(typeof(f.fs[k])), ")
	# end
	# return s
	return "n/a" # for now
end


using Gridap
import Gridap: ∇

# Define manufactured functions
ufun(x) = x[1] + x[2]
ufun_grad(x) = VectorValue(1.0,1.0,0.0)
∇(::typeof(ufun)) = ufun_grad
bfun(x) = 0.0

# Define forms of the problem
a(v,u) = inner(∇(v), ∇(u))

# Define norms to measure the error
l2(u) = inner(u,u)
h1(u) = a(u,u) + l2(u)

function main(n)

  # Construct the discrete model
  println("model")
  @time model = CartesianDiscreteModel(partition=(n,n,n))

  # Construct the FEspace
  order = 1
  diritag = "boundary"
  println("fespace")
  @time fespace = ConformingFESpace(Float64,model,order,diritag)

  # Define test and trial spaces
  println("V")
  @time V = TestFESpace(fespace)

  println("U")
  @time U = TrialFESpace(fespace,ufun)

  # Define integration mesh and quadrature
  println("trian")
  @time trian = Triangulation(model)

  println("quad")
  @time quad = CellQuadrature(trian,degree=2)

  # Define the source term
  println("bfield")
  @time bfield = CellField(trian,bfun)

  # Define Assembler
  println("assem")
  @time assem = SparseMatrixAssembler(V,U)

  # Define the FEOperator
  b(v) = inner(v,bfield)
  println("op")
  @time op = LinearFEOperator(a,b,V,U,assem,trian,quad)

  # Define the FESolver
  ls = LUSolver()
  solver = LinearFESolver(ls)

  # Solve!
  println("uh")
  @time uh = solve(solver,op)

  # Define exact solution and error
  println("u")
  @time u = CellField(trian,ufun)

  println("e")
  @time e = u - uh

  # Compute errors
  println("el2")
  @time el2 = sqrt(sum( integrate(l2(e),trian,quad) ))

  println("eh1")
  @time eh1 = sqrt(sum( integrate(h1(e),trian,quad) ))

  @assert el2 < 1.e-8
  @assert eh1 < 1.e-8

end


# SPDX-License-Identifier: MIT

struct NormalModel{T} <: AbstractModel
    pθ::T
end

#function loglikelihood(::NormalModel, x::AbstractVector, uTθ)
#    Λμ, μᵀΛμ, Λ, logdetΛ = uTθ
#    D = length(x)
#    xᵀΛx = dot(x, Λ * x)
#    μᵀΛx = dot(x, Λμ)
#    -(1/2)*(xᵀΛx + (μᵀΛμ - logdetΛ + D*log(2π))) + μᵀΛx
#end

function loglikelihood(::NormalModel, X::AbstractMatrix, uTθ)
    Λμ, μᵀΛμ, Λ, logdetΛ = uTθ
    D = size(X, 1)
    XᵀΛX = sum(X .* (Λ * X), dims=1)
    μᵀΛX = Λμ' * X
    -(1/2)*(XᵀΛX .+ (μᵀΛμ - logdetΛ + D*log(2π))) .+ μᵀΛX
end

#######################################################################
# GSM prior/posterior

#getstats(model::NormalModel, pθ::NormalWishart) = μ(pθ)
#unpack(model::NormalModel, pθ::NormalWishart, Tθ) = unpack(pθ, Tθ)
#
#function newposterior(::NormalModel, qθ::NormalWishart, ∇μ::AbstractVector;
#                      lrate=1)
#    T = typeof(qθ)
#    ηq = η(qθ)
#    ∇̃ξ = ForwardDiff.derivative(t -> ξ(qθ, ηq + t*∇μ), 0)
#    ξᵗ⁺¹ = ξ(qθ, ηq) + lrate*∇̃ξ
#    NormalWishart(unpack(qθ, ξᵗ⁺¹)...)
#end


## findmin/findmax/minumum/maximum

A = sparse([1.0 3.0 6.0;
            5.0 2.0 4.0])
for (tup, rval, rind) in [((1,), [1.0 2.0 4.0], [1 4 6]),
                          ((2,), reshape([1.0,2.0], 2, 1), reshape([1,4], 2, 1)),
                          ((1,2), fill(1.0,1,1),fill(1,1,1))]
    @test findmin(A, tup) == (rval, rind)
    @test findmin!(similar(rval), similar(rind), A) == (rval, rind)
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((1,), [5.0 3.0 6.0], [2 3 5]),
                          ((2,), reshape([6.0,5.0], 2, 1), reshape([5,2], 2, 1)),
                          ((1,2), fill(6.0,1,1),fill(5,1,1))]
    @test findmax(A, tup) == (rval, rind)
    @test findmax!(similar(rval), similar(rind), A) == (rval, rind)
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

#issue 23209

A = sparse([1.0 3.0 6.0;
            NaN 2.0 4.0])
for (tup, rval, rind) in [((1,), [NaN 2.0 4.0], [2 4 6]),
                          ((2,), reshape([1.0, NaN], 2, 1), reshape([1,2], 2, 1)),
                          ((1,2), fill(NaN,1,1),fill(2,1,1))]
    @test string(findmin(A, tup)) == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((1,), [NaN 3.0 6.0], [2 3 5]),
                          ((2,), reshape([6.0, NaN], 2, 1), reshape([5,2], 2, 1)),
                          ((1,2), fill(NaN,1,1),fill(2,1,1))]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

A = sparse([1.0 NaN 6.0;
            NaN 2.0 4.0])
for (tup, rval, rind) in [((1,), [NaN NaN 4.0], [2 3 6]),
                          ((2,), reshape([NaN, NaN], 2, 1), reshape([3,2], 2, 1)),
                          ((1,2), fill(NaN,1,1),fill(2,1,1))]
    @test string(findmin(A, tup))  == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((1,), [NaN NaN 6.0], [2 3 5]),
                          ((2,), reshape([NaN, NaN], 2, 1), reshape([3,2], 2, 1)),
                          ((1,2), fill(NaN,1,1),fill(2,1,1))]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

A = sparse([Inf -Inf Inf  -Inf;
            Inf  Inf -Inf -Inf])
for (tup, rval, rind) in [((1,), [Inf -Inf -Inf -Inf], [1 3 6 7]),
                          ((2,), reshape([-Inf -Inf], 2, 1), reshape([3,6], 2, 1)),
                          ((1,2), fill(-Inf,1,1),fill(3,1,1))]
    @test string(findmin(A, tup))  == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((1,), [Inf Inf Inf -Inf], [1 4 5 7]),
                          ((2,), reshape([Inf Inf], 2, 1), reshape([1,2], 2, 1)),
                          ((1,2), fill(Inf,1,1),fill(1,1,1))]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

A = sparse([BigInt(10)])
for (tup, rval, rind) in [((2,), [BigInt(10)], [1])]
    @test string(findmin(A, tup))  == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((2,), [BigInt(10)], [1])]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

A = sparse([BigInt(-10)])
for (tup, rval, rind) in [((2,), [BigInt(-10)], [1])]
    @test string(findmin(A, tup))  == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((2,), [BigInt(-10)], [1])]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

A = sparse([BigInt(10) BigInt(-10)])
for (tup, rval, rind) in [((2,), [BigInt(-10)], [2])]
    @test string(findmin(A, tup))  == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((2,), [BigInt(10)], [1])]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end

A = sparse(["a", "b"])
for (tup, rval, rind) in [((1,), ["a"], [1])]
    @test string(findmin(A, tup))  == string((rval, rind))
    @test string(findmin!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(minimum(A, tup)) == string(rval)
    @test string(minimum!(similar(rval), A)) == string(rval)
    @test string(minimum!(copy(rval), A, init=false)) == string(rval)
end

for (tup, rval, rind) in [((1,), ["b"], [2])]
    @test string(findmax(A, tup)) == string((rval, rind))
    @test string(findmax!(similar(rval), similar(rind), A)) == string((rval, rind))
    @test string(maximum(A, tup)) == string(rval)
    @test string(maximum!(similar(rval), A)) == string(rval)
    @test string(maximum!(copy(rval), A, init=false)) == string(rval)
end


media_impact_fun(x,y) = exp(-x*y)


import Polynomials
using Polynomials

function polymult(z::AbstractVector{T}, l::AbstractVector{Int}) where {T<:Number}
#
# construct monic polynomial from its roots
# use sorting and a special order of processing to increase accuray.
# special implementation of poly()
# returns Real polynomal, if all conjugate roots with equal multiplicities.
#
#   INPUT:
#            z vector of distinct roots of polynomial
#            l vector of corresponding multiplicities
#   OUTPUT:
#            ff Poly representing the product (x - z[k]) ^ l[k]
#
    # println("polymult($(z), $(l)")
    # z = PolyZeros(z)
    products(z, l)
end

# Calculate the coefficient operator of the pejorative manifold
# PolyRoots is used to leverage calculation of (x-z[j]) - products
#
function polyG(z::AbstractVector{<:Number}, mult::AbstractVector{<:Integer})
    polyG(products(z, mult))
end
function polyG(p::AbstractArray{<:Number})
    p[2:end]
end

# The derivative of G at point z
function polyDG(zz::AbstractVector{S}, ll) where S<:Number
    m = length(zz)
    s = products(zz, ll .- 1)
    n0 = length(s)
    n = sum(ll)
    Df = zeros(S, n, m)

    for j = 1:m
        Df[1:n0,j] .= s
        ix = n0
        for k = 1:m
            if k != j && ll[j] > 0
                ix = product1!(ix, view(Df,1:n,j), 1, -zz[k])
            end
        end
        Df[1:n,j] .*= -ll[j]
    end
    isreal(Df) ? real(Df) : Df
end

function polyDG(zz::AbstractVector{S}, ll::AbstractVector{<:Integer}, p::AbstractArray{T}, zs::AbstractVector{<:Integer}) where {S<:Number,T<:Number}

    w = [1/max(abs(pp), 1) for pp in p]
    m = length(zz)
    n = sum(ll)
    Df = zeros(T, n, m)
    for j = 1:m
        zsj = zs[j]
        zsj <= 0 || zsj > j || continue
        q, r = adiv(p, -zz[j], w)
        q .*= -ll[j]
        if zsj <= 0
            Df[:,j] .= q
        else
            if zsj > j
                Df[:,j] .= real(q)
                Df[:,zsj] .= imag(q)
            end
        end
    end
    Df
end

function realtocoeff(z::AbstractVector{<:Real}, S::Type{<:Complex}, zs::AbstractVector)

    r = similar(z, S)
    all( zs .<= 0 ) && return r
    for j = 1:length(z)
        zsj = zs[j]
        if zsj <= 0
            r[j] = z[j]
        elseif zsj > j
            r[j] = z[j]/2 - im*z[zsj]/2
            r[zsj] = conj(r[j])
        end
    end
    r
end
function realtocoeff(z::AbstractVector, ::Type, zs::AbstractVector)
    z
end

"""

Calculate polynomial product of `p`and `q` given as vectors.
"""
product

function product!(r::AbstractVector{U}, p::AbstractVector{S}, q::AbstractVector{T}) where {U<:Number,S<:Number,T<:Number}

  np = length(p)
  nq = length(q)
  nr = length(r)
  nr >= np + nq - 1 || throw(ArgumentError("size of result vector must be exactly (length(a) + length(b) - 1 but $np + $nq - 1 != $nr"))

  nq > np && return product(q, p)
  U == promote_type(S, T, U) || throw(ArgumentError("result type not compatible with input"))

  if nq == 0
    return S[]
  end
  E = one(U)
  copy!(r, p)
  fill!(view(r, np+1:nr), zero(U))
  q1 = U(q[1])
  if q1 != E
    for j = 1:np
      r[j] *= q1
    end
  end
  for k = 2:nq
    q1 = U(q[k])
    for j = 1:np
      jk = j + k - 1
      r[jk] = muladd(q1, p[j], r[jk])
    end
  end
  r
end

function product(p::AbstractVector{S}, q::AbstractVector{T}) where {S<:Number,T<:Number}
  U = promote_type(S,T)
  product!(zeros(U, length(p) + length(q) - 1), p, q)
end

"""
    products(z::Vector{T}, ll::Vector{Int})

Calculate coefficients of polynomial given by its roots with multiplicities.
The z should be ordered by increasing absolute value for better accuracy.
"""
function products(z::AbstractVector{T}, ll::AbstractVector{Int}) where T <:Number
    m = length(z)
    m == length(ll) || throw(ArgumentError("root vector and muliplicities have differnet length"))
    all(ll .>= 0) || throw(ArgumentError("multiplicities must not be negative"))
    n = sum(ll)
    S, = zerostructure(z, ll)
    r = zeros(S, n+1)
    r[1] = 1
    Z = zero(S)
    ix = 1
    jj = sortperm(z, by=abs2, rev=false)
    z = view(z, jj)
    ll =  ll[jj]
    while sum(ll) > 0
        p = minpos(ll)
        jm = findall(x-> x >= p, ll)
        s = singprod(z[jm], S)
        s = power(s, p)
        ix = convolute!(ix, r, s)
        ll[jm] .-= p
    end
    r
end

"""
    singprod(z::Vector, S::Type))

Return coefficients of polynomial `prod((x - zz) for zz in z)`

"""
function singprod(z::AbstractVector{T}, ::Type{S}) where {T<:Number,S<:Number}
    m = length(z)
    r = Vector{S}(undef, m+1)
    r[1] = 1
    ix = 1
    if S<:Real && T<:Complex
        for k = 1:m
            zk = z[k]
            rr = -real(zk)
            ri = imag(zk)
            if iszero(ri)
                ix = product1!(ix, r, 1, rr)
            elseif ri > 0
                ix = product2!(ix, r, 1, 2*rr, abs2(zk))
            end
        end
    else
        for k = 1:m
            ix = product1!(ix, r, 1, -z[k])
        end
    end
    r
end

"""
    minpos(v::Vector{<:Integer})

Return the minimal positive integer in v
"""
function minpos(x::AbstractVector{T}) where T<:Integer
    y = typemax(T)
    minimum(x .+ y) - y
end

"""
    convolute!(ix, r, p)

If `r` is a vector representing a polynomial in `r[1:ix]` and p is a vector
containing another polynomial, modify `r` to contain the product.
Assume `ix >= 0` and `length(r) >= ix + length(p) - 1`.
"""
function convolute!(ix::Integer, r::AbstractVector, p::AbstractVector)
    ix <= 0 && return ix
    n = length(p) - 1
    ri = r[ix]
    for j = 0:n
        r[ix+j] = ri * p[j+1]
    end
    for i = ix-1:-1:1
        ri = r[i]
        for j = n:-1:1
            r[i+j] += ri * p[j+1]
        end
        r[i] = ri * p[1]
    end
    ix + n
end

"""
    convolute(p, q)

Return the convolution (polynomial product) of vectors p and q.
The coefficient order does not matter. No checks for zero coefficients.
"""
function convolute(p::AbstractVector{T}, q::AbstractVector{S}) where {T,S}
    R = promote_type(T,S)
    m = length(p)
    n = length(q)
    (n == 0 || m == 0) && return R[]
    r = similar(p, R, n + m - 1)
    copyto!(r, p)
    convolute!(m, r, q)
    r
end

"""
    product1!(ix, r, ll, p1)

multiply by (x + p1)^ll. ix is the current used size of r.
"""
function product1!(ix::Int, r::AbstractVector{T}, ll::Int, zk::T) where T<:Number
    for j = 1:ll
        r[ix+1] = r[ix] * zk
        for i = ix:-1:2
            r[i] = muladd(r[i-1], zk, r[i])
        end
        ix += 1
    end
    ix
end

"""
    product2!(ix, r, ll, p1, p2)

multiply by (x^2 + p1*x + p2)^ll. ix is the current used size of r.
"""
function product2!(ix::Int, r::AbstractVector{T}, ll::Int, p1::T, p2::T) where T<:Number
    for j = 1:ll
        r[ix+1] = r[ix+2] = 0
        for i = ix:-1:1
            rr = muladd(r[i+1], p1, r[i+2])
            r[i+2] = muladd(r[i], p2, rr)
        end
        r[2] = muladd(r[1], p1, r[2])
        ix += 2
    end
    ix
end

"""
    adiv(c::Vector, a::Number, weight::Vector) -> b::Vector

Divide Polynomial `x^n * c[1] + ... + c[n+1]` by `x + a` giving
`x^(n-1) * b[1] + ... + b[n]` under the condition that the division remainder is zero.

`weight` is a positive weight vector with the same size as `c` used to solve the
least squares problem `minimize norm(w .* (A * b - c)) where `A` is the `n+1 x n`
Cauchy matrix describing the multiplication with `x + a`.
The results tend to be much more accurate than that of the usual long division,
especially in the case of hughe coefficients, which are damped by `weight = 1 ./ abs.(c)`.
"""
function adiv(c::AbstractVector{<:Number}, a::S, w::AbstractVector{<:Real}) where S<:Number
    T = real(S)
    n = length(c) - 1
    aa2 = abs2(a)
    aa = sqrt(aa2)
    ac = conj(a)
    r = Vector{T}(undef, n)
    d = Vector{T}(undef, n)
    b = Vector{S}(undef, n)
    b[1] = 1
    b[n] = c[n+1] / a
    w1 = w[2]
    u1 = (c[2] - a) * w1
    for i = 2:n-2
        w2 = w[i+1]
        u2 = c[i+1] * w2
        z = d[i] = hypot(w1, w2*aa)
        w1z = w1 / z
        w2z = w2 / z
        r[i] = w2z * w2
        b[i] = w1z * u1 + w2z * u2 * ac
        u1 = -w2z * u1 * a + w1z * u2
        w1 *= w2z
    end
        w2 = w[n]
        u2 = (c[n] - b[n]) * w2
        z = d[n-1] = hypot(w1, w2*aa)
        w1z = w1 / z
        w2z = w2 / z
        r[n-1] = w2z * w2
        b[n-1] = w1z * u1 + w2z * u2 * ac
        u1 = -w2z * u1 * a + w1z * u2
    bi = b[n-1] = b[n-1] / d[n-1]
    for i = n-2:-1:2
        bi = b[i] = (b[i] - bi * r[i] * ac) / d[i]
    end
    b, u1
end

function power(p::AbstractArray, n::Integer)
    (Poly(p)^n).a
end

function Polynomials.roots(pc::AbstractVector{T}) where T<:Number
# redirect to implementation in Polynomials
	p = Poly(reverse(pc))
	roots(p)
end

function deriv(p::AbstractVector{T}) where T<:Number

    n = length(p) - 1
    map(k -> (p[k] * (n-k+1)), 1:n)
end


using Enzyme, BenchmarkTools, FiniteDiff, LinearAlgebra

f(x) = x^2

f(x) = x[1]*x[2]

inp =  [2.0 1.0
        0.0 1.0]
dinp = [0.0 0.0
        0.0 0.0]

f2(x) = inv(x)

@benchmark autodiff(det, Duplicated(inp, dinp))

@benchmark FiniteDi


using Revise
using ADCME
using ADSeismic
using PyPlot
using DelimitedFiles

##=======================
reset_default_graph()

scale = 100
param = AcousticPropagatorParams(NX=scale, NY=scale, NSTEP=3, DELTAT=1e-4,  DELTAX=1.0, DELTAY=1.0)

rc = Gauss(param, 1/pi, 100, 1e6)
srci = [div(param.NX,2)]
srcj = [div(param.NY,2)]
srcv = reshape(rc, :, 1)
src = AcousticSource(srci, srcj, srcv)

C = 3300*ones(param.NX+2, param.NY+2)
model = AcousticPropagatorSolver(param, src, C)

sess = Session(); init(sess)

run_profile(sess, model.u)
save_profile("acoustic-profile-$(has_gpu() ? "gpu" : "cpu").json")

##=======================
reset_default_graph()

scale = 100
param = ElasticPropagatorParams(NX=scale, NY=scale, NSTEP=2, DELTAT=1e-4,  DELTAX=1.0, DELTAY=1.0)

source = Gauss(param, 1/pi, 100, 1e6)
srci = [div(param.NX,2)]
srcj = [div(param.NY,2)]
srctype = [0]
srcv = reshape(source, :, 1)
src = ElasticSource(srci, srcj, srctype, srcv)

vp = 3300.
vs = 3300. / 1.732
rho = 2800.
λ, μ, ρ = compute_lame_parameters(param.NX, param.NY, vp, vs, rho)
model = ElasticPropagatorSolver(param, src, ρ, λ, μ)

sess = Session(); init(sess)

run_profile(sess, model.vx)
save_profile("elastic-profile-$(has_gpu() ? "gpu" : "cpu").json")


module Voronoi

using MiniQhull
using LinearAlgebra

import ChainRulesCore: rrule, NoTangent

∧(x,y) = x[1,:].*y[2,:] .- x[2,:].*y[1,:]

"""
    boundary(d)

Returns the boundary vertices of a unit cube in `d` dimensions.
"""
function boundary(d)
    d == 2 && return Float32.([[-1;-1] [-1;+1] [+1;+1] [+1;-1]])

    b = zeros(Float32, d, 2^d)
    for i in 0:(2^d-1)
        for n in 0:(d-1)
            b[n+1,i+1] = (((i >> n) & 1) == 1) ? +1 : -1
        end
    end
    return b
end

const NB = 4

"""
    tessellation(q)

Construct a voronoi tessellation from generating points `q`.
Return _just_ vertices of construction.
"""
function tessellation(q)
    # q = hcat(boundary, x)
    triangulation = delaunay(q)

    # order triangulation
    areas = (q[:,triangulation[2,:]] .- q[:,triangulation[1,:]]) ∧ (q[:,triangulation[3,:]] .- q[:,triangulation[2,:]])
    orientation = sign.(areas)
    for i in findall(orientation .== -1)
        triangulation[2,i], triangulation[3,i] = triangulation[3,i], triangulation[2,i]
    end

    # compute voronoi vertices
    r = zeros(eltype(q), size(q,1), size(triangulation,2))
    for t in 1:size(triangulation,2)
        for i in 1:3
            j = (i % 3) + 1
            k = (j % 3) + 1

            α,β,γ = triangulation[i,t], triangulation[j,t], triangulation[k,t]

            qα² = sum((q[:,α]).^2)
            r[1,t] += qα²*(q[2,γ] - q[2,β])
            r[2,t] -= qα²*(q[1,γ] - q[1,β])
        end
    end

    r[1,:] = r[1,:] ./ abs.(areas)
    r[2,:] = r[2,:] ./ abs.(areas)

    triangulation, r
end

"""
    areas(x)

Compute the areas of all 2-dimensional simplices generated by the delaunay construction for pointcloud `x`.
Assumes `x` is sized ``N \\times 2``
"""
function areas(x)
    q = hcat(boundary(2), x)
    triangulation = delaunay(q)

    a = 0.5*[
        let
            q[1,t[1]]*(q[2,t[2]]-q[2,t[3]]) +
            q[1,t[2]]*(q[2,t[3]]-q[2,t[1]]) +
            q[1,t[3]]*(q[2,t[1]]-q[2,t[2]])
        end for t in eachcol(triangulation)
    ]
    s = sign.(a)

    return s.*a
end

function rrule(::typeof(areas), x)
    q = hcat(boundary(2), x)
    triangulation = delaunay(q)

    a = 0.5*[
        let
            q[1,t[1]]*(q[2,t[2]]-q[2,t[3]]) +
            q[1,t[2]]*(q[2,t[3]]-q[2,t[1]]) +
            q[1,t[3]]*(q[2,t[1]]-q[2,t[2]])
        end for t in eachcol(triangulation)
    ]
    s = sign.(a)

    return s.*a, (∂a) -> let
        ∂x = zeros(size(x))

        for (i,t) in enumerate(eachcol(triangulation))
            if t[1] > NB
                ∂x[1,t[1]-NB] += (q[2,t[2]]-q[2,t[3]])*∂a[i]*s[i]
                ∂x[2,t[1]-NB] -= (q[1,t[2]]-q[1,t[3]])*∂a[i]*s[i]
            end
            if t[2] > 4
                ∂x[1,t[2]-NB] += (q[2,t[3]]-q[2,t[1]])*∂a[i]*s[i]
                ∂x[2,t[2]-NB] -= (q[1,t[3]]-q[1,t[1]])*∂a[i]*s[i]
            end
            if t[3] > 4
                ∂x[1,t[3]-NB] += (q[2,t[1]]-q[2,t[2]])*∂a[i]*s[i]
                ∂x[2,t[3]-NB] -= (q[1,t[1]]-q[1,t[2]])*∂a[i]*s[i]
            end
        end

        (NoTangent(), 0.5*∂x)
    end

end

# derivative of det given by adjugate (which is just rescaled inverse)
# see https://en.wikipedia.org/wiki/Adjugate_matrix

affine(simplex) = vcat(simplex, ones(Float32, 1, size(simplex,2)))
volume(simplex) = det(affine(simplex))

"""
    volumes(x)

Compute the volume of all d-dimensional simplices generated by the delaunay construction for pointcloud `x`.
Assumes `x` is sized ``N \\times d``
"""
function volumes(x)
    d = size(x,1)
    q = hcat(boundary(d), x)

    simplices = delaunay(q)

    Ω = [ volume(hcat((q[:,i] for i in simplex)...)) for simplex in eachcol(simplices) ]
    s = sign.(Ω)

    return s.*Ω
end

function rrule(::typeof(volumes), x)
    d = size(x,1)
    b = boundary(d)
    q = hcat(b, x)

    v₀ = size(b,2)
    simplices = delaunay(q)

    Ω = [ volume(hcat((q[:,i] for i in simplex)...)) for simplex in eachcol(simplices) ]
    s = sign.(Ω)

    Z = factorial(d)
    return (s.*Ω) ./ Z, function(∂Ω)
        ∂x = zeros(Float32,size(x))

        for (n,simplex) in enumerate(eachcol(simplices))
            ω = s[n]*Ω[n]*∂Ω[n]
            try
                Q = inv(affine(hcat((q[:,i] for i in simplex)...)))

                for (i,v) in enumerate(simplex)
                    if v > v₀
                        for j in 1:d
                            ∂x[j,v-v₀] += ω*Q[i,j]
                        end
                    end
                end
            catch error
                # XXX: do we need to do something more intelligent here?
                if isa(error, LinearAlgebra.SingularException)
                    continue
                else
                    throw(error)
                end
            end
        end

        return (NoTangent(), ∂x ./ Z)
    end
end

end


using FiniteVolumes
using Printf
using Plots

mesh = PeriodicCartesianMesh(40, 40)
model = LinearAdvectionFlux([1.0, 1.0])

is_in_square(x, side=0.5) = all(0.5-side/2 .<= x .<= 0.5+side/2)

w₀ = [is_in_square(x) ? 1.0 : 0.0 for x in cell_centers(mesh)]

anim = Animation()
function plot_callback(i, t, w)
    if i % 2 == 0
        title = @sprintf "time = %.2f" t
        plot(mesh, w, 1, title=title, clims=(0, 1))
        frame(anim)
    end
end

t, w = FiniteVolumes.run(directional_splitting(model), mesh, w₀,
                         time_step=FixedCourant(0.2),
                         nb_time_steps=200,
                         callback=plot_callback)
gif(anim, "advection.gif")


# Load file for needed functions
include("../Packages/ReciprocalVariety.jl")

# Basic setup. For ternary quartics: r=2, s=3
    r=2
    s=3
    m=binomial(r+s-1,s-1)
    N=binomial(2r+s-1,2r)
    M=binomial(m + 1, 2)
    @var x[0:2r, 0:2r, 0:2r],y[1:M]

# Define ideal equations
    cat=CatalecticantIndices(r,s,m)
    adj=adjugate(Expression.(vec_to_sym(y,m)));
    eq=CatalecticantSpace(cat,adj,m);

# Set up system
    F = System(eq; variables=y);

# Compute point on the variety
    x₀ = randn(ComplexF64,N)
    @var x[0:2r, 0:2r, 0:2r]
    X = map(ijk -> x[(ijk .+ 1)...], cat)
    y₀ = sym_to_vec(inv(X(variables(X) => x₀)),m)

# Put affine subspace through this point
    V₀ = let
       A₀ = randn(ComplexF64, N, M)
       b₀ = A₀ * y₀
       LinearSubspace(A₀, b₀)
    end

# Perform monodromy
    mres = monodromy_solve(
       F,
       y₀,
       V₀,
       parameter_sampler=_ -> LinearSubspace(
           randn(ComplexF64, N, M),
           randn(ComplexF64, N),
       ),
   )


# Interval Library
import Base.==
using Base.Test

type Interval
    left::Float64
    right::Float64
end

type EmptyInterval
end

function ==(a::Interval, b::Interval)
    return a.left == b.left && a.right == b.right
end

function ==(a::EmptyInterval, b::EmptyInterval)
    return true
end

function ==(a::Interval, b::EmptyInterval)
    return false
end

function ==(a::EmptyInterval, b::Interval)
    return false
end

function createInterval(a, b)
    left = convert(Float64,a)
    right = convert(Float64,b)
    if left <= right
        return Interval(left,right)
    else
        error("left > right this not define an interval.")
    end
end

function intersection(X::Interval,Y::Interval)
    if X.right <= Y.left || Y.right <= X.left
        return EmptyInterval()
    else
        a = max(X.left, Y.left)
        b = min(X.right, Y.right)
        return Interval(a, b)
    end
end

function intersection(X::EmptyInterval, Y::Interval)
    return EmptyInterval()
end

function intersection(X::Interval, Y::EmptyInterval)
    return EmptyInterval()
end

function intersection(X::EmptyInterval, Y::EmptyInterval)
  return EmptyInterval()
end

#Definition the interval hull
function intervalHull(X::Interval, Y::Interval)
    a = min(X.left, Y.left)
    b = max(X.right, Y.right)
    return Interval(a,b)
end

function intervalHull(X::EmptyInterval, Y::Interval)
    return Interval(Y.left, Y.right)
end

function intervalHull(X::Interval, Y::EmptyInterval)
    return Interval(X.left, X.right)
end

function intervalHull(X::EmptyInterval, Y::EmptyInterval)
  return EmptyInterval()
end

function ∪(X::Interval, Y::Interval)
  return Interval(X.left,Y.right)
end

# Testing

@test createInterval(-1, 10) == Interval(-1., 10.)


# fft version
using FFTW

# *******************************************
#  Configuration
# *******************************************
"""
    ConfigFFT{T <: Union{Float64,Complex{Float64}},N}

Assembly all basic configurations of the Fourier-spectral method.
T represents the type of values of functions in 'real' space,
and N represents the dimension. The followings are all members:

    ngrids::NTuple{N,Int}
... Each element means the number of grids in each axis

    xranges::NTuple{N,NTuple{2,Float64}}
... Tuple of ((xmin, xmax), (ymin, ymax), ...)

    Xcoords::NTuple{N,Array{Float64,N}}
... Coordinates of 'real' space in each axis

    Kcoords::NTuple{N,Array{Complex{Float64},N}}
... Wavenumbers in each axis. It is used to calculate derivatives.

    P_fft, P_ifft, P_fftpad, P_ifftpad
... Operators used to accelerate fft and ifft

    cut_zigzag_mode::Bool
... The highest mode is suppressed or not for axes whose grids number is even.
If you use the easy constructor (explained below), it is set to be 'true' by default.
This flag is used in 'X_K' and 'K_X' function, the fourier transformation. If you
try to accelerate as much as possible and can manage the highest wavenumber mode,
set 'cut_zigzag_mode=false'.

For normal use, the easy constructor below is recommended for initialization.

    ConfigFFT(
        ngrids::NTuple{N,Int},
        xranges::NTuple{N,NTuple{2,Float64}};
        use_complex::Bool=false,
        cut_zigzag_mode::Bool=true)

The function in 'real' space is set to be real-valued (Float64) by default.
If you want to use complex-valued function, set 'use_complex=true'.

# Examples
1) 1-dimensional case. The number of grid is 64, and the range is 0 ≦ x < 10.

    julia> ngrids = (64,);

    julia> xranges = ((0., 10.),);

    julia> config = ConfigFFT(ngrids, xranges);

2) 2-dimensional case. The numbers of grids are 128 in x-axis and 256 in y-axis.
   The ranges in each axis are 0 ≦ x < 2π and 0 ≦ y < 4π.

    julia> ngrids = (128, 256);

    julia> xranges = ((0., 2π), (0., 4π));

    julia> config = ConfigFFT(ngrids, xranges);

3) 1-dimensional case. The number of grid is 64, and the range is 0 ≦ x < 30.
   The complex-valued functions in 'real' space is used.

   julia> ngrids = (64,);

   julia> xranges = ((0., 30.),);

   julia> config = ConfigFFT(ngrids, xranges, use_complex=true);
"""
struct ConfigFFT{T<:Union{Float64,Complex{Float64}},N}

    ngrids::NTuple{N,Int} # Number Of Grids

    xranges::NTuple{N,NTuple{2,Float64}} # tuple of (min, max) s
    Xcoords::NTuple{N,Array{Float64,N}}
    Kcoords::NTuple{N,Array{Complex{Float64},N}}

    P_fft::FFTW.cFFTWPlan{T1} where T1
    P_ifft::AbstractFFTs.ScaledPlan{T2} where T2
    P_fftpad::FFTW.cFFTWPlan{T3} where T3
    P_ifftpad::AbstractFFTs.ScaledPlan{T4} where T4

    cut_zigzag_mode::Bool

    # CONSTRUCTOR
    function ConfigFFT{T,N}(
            ngrids, xranges, Xcoords, Kcoords,
            P_fft, P_ifft, P_fftpad, P_ifftpad;
            cut_zigzag_mode=true
        ) where {T,N}

        warn_if_ngrid_is_not_power_of_2(ngrids)

        new(ngrids, xranges, Xcoords, Kcoords,
            P_fft, P_ifft, P_fftpad, P_ifftpad,
            cut_zigzag_mode)
    end

    # EASY CONSTRUCTOR
    function ConfigFFT(
            ngrids::NTuple{N,Int},
            xranges::NTuple{N,NTuple{2,Float64}};
            use_complex::Bool=false,
            cut_zigzag_mode::Bool=true
        ) where N

        T = use_complex ? Complex{Float64} : Float64

        carts = CartesianIndices(ngrids)
        Xcoords = Xcoordsgen(ngrids, xranges)
        Kcoords = Kcoordsgen(ngrids, xranges)

        P_fft = plan_fft(Xcoords[1])
        P_ifft = plan_ifft(Xcoords[1])

        pad_ngrids = @. ngrids ÷ 2 * 3
        P_fftpad = plan_fft(zeros(pad_ngrids))
        P_ifftpad = plan_ifft(zeros(pad_ngrids))

        return ConfigFFT{T,N}(
            ngrids, xranges, Xcoords, Kcoords,
            P_fft, P_ifft, P_fftpad, P_ifftpad,
            cut_zigzag_mode=cut_zigzag_mode)
    end

end

# --- helper function for constuctor ---
function Xcoordsgen(
        ngrids::NTuple{N,Int},
        xranges::NTuple{N,NTuple{2,Float64}}
    ) where N

    _Xcoordgen(axis) = Xcoordgen(axis, ngrids, xranges)
    return ntuple(_Xcoordgen, N)

end

function Xcoordgen(
        axis::Int,
        ngrids::NTuple{N,Int},
        xranges::NTuple{N,NTuple{2,Float64}}
    ) where N

    ngrid = ngrids[axis]
    xrange = xranges[axis]
    _Xcoordgen(indices) = (
        (  (indices[axis] - 1)*xrange[2]
         + (ngrid - indices[axis] + 1)*xrange[1] ) / ngrid
    )
    return _Xcoordgen.(CartesianIndices(ngrids))

end

function Kcoordsgen(
        ngrids::NTuple{N,Int},
        xranges::NTuple{N,NTuple{2,Float64}}
    ) where N

    _Kcoordgen(axis) = Kcoordgen(axis, ngrids, xranges)
    return ntuple(_Kcoordgen, N)

end

function Kcoordgen(
        axis::Int,
        ngrids::NTuple{N,Int},
        xranges::NTuple{N,NTuple{2,Float64}}
    ) where N

    ngrid = ngrids[axis]
    xrange = xranges[axis]
    _Kcoordgen(indices) = (
        kval(indices[axis], ngrid, xrange)
    )

    return _Kcoordgen.(CartesianIndices(ngrids))

end

function kval(index, ngrid, xrange)::Complex{Float64}

    index0 = index - 1  # start by 0
    if 2*index0 < ngrid
        return index0
    elseif 2*index0 == ngrid
        return 0
    elseif 2*index0 > ngrid
        return index0 - ngrid
    end

end

# *******************************************
#  XFunc
# *******************************************
"""
    XFunc{T,N} <: AbstractArray{T,N}

A struct of a function on 'real' space.
It consists of the following:

    vals::Array{T,N}
... Values of the function on grids

    config::ConfigFFT{T,N}
... Configuration of Fourier Transformation

For normal use, the easy constructor below is recommended for initialization.

    XFunc(
        vals::Array{Tc,N},
        config::ConfigFFT{T,N}
        ) where {N, Tc<:Number, T}

The type of elements of input 'vals' is converted to T.
Instead, you may generate a XFunc object, whose values are undef as

    XFunc(undef, config)

XFunc type object is transformed by 'K_X' to KFunc type object,
which represents a function in 'wavenumber' space.
See also the explanation of KFunc.

# Examples
Suppose 'config' be ConfigFFT{Float64,2} type object,
and c.ngrids = (2, 2). Then, for example,

    v = [1. 2.; 3. 4.]
    X_func = XFunc(v, config)

generates a XFunc object whose 'vals' is v.
If 'K_X' is applied to this function, say,

    K_func = K_X(X_func)

a new object 'K_func' is returned.
It is KFunc type, and its 'vals' is Fourier transform of v.
"""
mutable struct XFunc{T,N} <: AbstractArray{T,N}

    vals::Array{T,N}
    config::ConfigFFT{T,N}

    # CONSTRUCTOR
    function XFunc{T,N}(
            vals::Array{T,N},
            config::ConfigFFT{T,N}
        ) where {T,N}

        check_size_consistency(vals, config)
        return new(vals, config)

    end

    # EASY CONSTRUCTOR
    function XFunc(
            vals::Array{Tv,N},
            config::ConfigFFT{Tc,N}
        ) where {N, Tv<:Number, Tc}

        if Tv <: Complex && Tc <: Real
            return XFunc{Tc,N}(Tc.(real(vals)), config)
        else
            return XFunc{Tc,N}(Tc.(vals), config)
        end

    end

    # UNDEF CONSTRUCTOR
    function XFunc(
            undef::UndefInitializer,
            config::ConfigFFT{T,N}
        ) where {T,N}

        f_undef = Array{T,N}(undef, config.ngrids)
        return XFunc{T,N}(f_undef, config)

    end

end

Base.:size(f::XFunc) = size(f.vals)

Base.:getindex(f::XFunc, i::Int) = getindex(f.vals, i)

function Base.:getindex(
        f::XFunc{T,N}, I::Vararg{Int,N}
    ) where {T,N}

    getindex(f.vals, I...)
end

Base.:setindex!(f::XFunc, v, i::Int) = setindex!(f.vals, v, i)

function Base.:setindex!(
        f::XFunc{T,N}, v, I::Vararg{Int,N}
    ) where {T,N}

    setindex!(f.vals, v, I...)
end

Base.:copy(f::XFunc) = XFunc(copy(f.vals), f.config)

# *******************************************
#  KFunc
# *******************************************
"""
    KFunc{T,N} <: AbstractArray{T,N}

A struct of a function on 'wavenumber' space.
It consists of the following:

    vals::Array{Complex{Float64},N}
... Values of the function on grids

    config::ConfigFFT{T,N}
... Configuration of Fourier Transformation

For normal use, the easy constructor below is recommended for initialization.

    KFunc(
        vals::Array{Tc,N},
        config::ConfigFFT{T,N}
    ) where {N, Tc<:Number, T}

The type of elements of input 'vals' is converted to T.
Instead, you may generate a KFunc object, whose values are undef as

    KFunc(undef, config)

KFunc type object is transformed by 'X_K' to XFunc type object,
which represents a function in 'real' space.
See also the explanation of XFunc.

# Examples
Suppose 'config' be ConfigFFT{Float64,2} type object,
and c.ngrids = (2, 2). Then, for example,

    v = [1. 2.; 3. 4.]
    K_func = KFunc(v, config)

generates a XFunc object whose 'vals' is v.
If 'X_K' is applied to this function, say,

    X_func = X_K(K_func)

a new object 'X_func' is returned.
It is XFunc type, and its 'vals' is inverse-Fourier transform of v.
"""
mutable struct KFunc{T,N} <: AbstractArray{Complex{Float64},N}

    vals::Array{Complex{Float64},N}
    config::ConfigFFT{T,N}

    # CONSTRUCTOR
    function KFunc{T,N}(
            vals::Array{Complex{Float64},N},
            config::ConfigFFT{T,N}
        ) where {T,N}

        check_size_consistency(vals, config)
        return new(vals, config)

    end

    # EASY CONSTRUCTOR
    function KFunc(
            vals::Array{Tv,N},
            config::ConfigFFT{Tc,N}
        ) where {N, Tv<:Number, Tc}

        KFunc{Tc,N}(Complex{Float64}.(vals), config)
    end

    # UNDEF CONSTRUCTOR
    function KFunc(
            undef::UndefInitializer,
            config::ConfigFFT{T,N}
        ) where {T,N}

        f_undef = Array{Complex{Float64},N}(undef, config.ngrids)
        return KFunc{T,N}(f_undef, config)

    end

end

Base.:size(f::KFunc) = size(f.vals)

Base.:getindex(f::KFunc, i::Int) = getindex(f.vals, i)

function Base.:getindex(
        f::KFunc{T,N}, I::Vararg{Int,N}
    ) where {T,N}

    getindex(f.vals, I...)
end

Base.:setindex!(f::KFunc, v, i::Int) = setindex!(f.vals, v, i)

function Base.:setindex!(
        f::KFunc{T,N}, v, I::Vararg{Int,N}
    ) where {T,N}

    setindex!(f.vals, v, I...)
end

Base.:copy(f::KFunc) = KFunc(copy(f.vals), f.config)

# *******************************************
#  Binomial Operators
# *******************************************
const BINOP = (
    (:+, :.+), (:-, :.-), (:*, :.*),
    (:/, :./), (:\, :.\), (:^, :.^)
)

# +++++ XFunc +++++
Base.:+(f::XFunc) = f
Base.:-(f::XFunc) = XFunc(-f.vals, f.config)
Base.:^(f::XFunc, n::Integer) = XFunc(f.vals .^ n, f.config)
Base.:inv(f::XFunc) = f \ 1.0

for (op, opd) = BINOP
    @eval begin
        function Base.$op(f::XFunc, g::XFunc)
            check_config_consistency(f.config, g.config)
            return XFunc($opd(f.vals, g.vals), f.config)
        end

        Base.$op(f::XFunc, a::Number) = (
            XFunc($opd(f.vals, a), f.config)
        )
        Base.$op(a::Number, f::XFunc) = (
               XFunc($opd(a, f.vals), f.config)
        )
    end
end

# +++++ KFunc +++++
Base.:+(f::KFunc) = f
Base.:-(f::KFunc) = KFunc(-f.vals, f.config)
Base.:^(f::KFunc, n::Integer) = KFunc(f.vals .^ n, f.config)
Base.:inv(f::KFunc) = f \ 1.0

for (op, opd) = BINOP
    @eval begin
        function Base.$op(f::KFunc, g::KFunc)
            check_config_consistency(f.config, g.config)
            return KFunc($opd(f.vals, g.vals), f.config)
        end

        Base.$op(f::KFunc, a::Number) = (
            KFunc($opd(f.vals, a), f.config)
        )
        Base.$op(a::Number, f::KFunc) = (
            KFunc($opd(a, f.vals), f.config)
        )
    end
end

# *******************************************
#  Operators for Complex Numbers
# *******************************************
const OPERATOR = (
    :real, :imag, :conj
)

for op = OPERATOR
    @eval begin
        Base.$op(f::XFunc) = XFunc($op(f.vals), f.config)
        Base.$op(f::KFunc) = KFunc($op(f.vals), f.config)
    end
end

Base.:reim(f::XFunc) = (real(f), imag(f))
Base.:reim(f::KFunc) = (real(f), imag(f))

# *******************************************
#  Low/High-pass Filter
# *******************************************
# +++++ general pass filter +++++
function pass_K!(
        f::KFunc{T,N},
        slices::NTuple{N,UnitRange{Int}}
    ) where {T,N}

    vals = copy(f.vals)
    vals[slices...] .= 0.0 + 0.0im
    f.vals -= vals

end

# +++++ high-pass filter +++++
function highpass_K!(
        f::KFunc{T,N},
        min_nwaves::NTuple{N,Int}
    ) where {T,N}

    ngrids = f.config.ngrids
    if any(@. min_nwaves > ngrids ÷ 2)
        println("WARNING: all waves are suppressed")
        f.vals .= 0.
    else
        floors = tuple(ones(Int, N)...)
        ceils = ngrids
        min_indices = @. max(floors, floors + min_nwaves)
        max_indices = @. min(ceils, ceils + 1 - min_nwaves)

        slices = (:).(min_indices, max_indices)
        pass_K!(f, slices)
    end
end

function K_highpass_K(
        f::KFunc{T,N},
        min_nwaves::NTuple{N,Int}
    ) where {T,N}

    g = copy(f)
    highpass_K!(g, min_nwaves)
    return g

end

# +++++ low-pass filter +++++
function lowpass_K!(
        f::KFunc{T,N},
        max_nwaves::NTuple{N,Int}
    ) where {T,N}

    ngrids = f.config.ngrids
    f.vals .= fftshift(f.vals)

    center_indices = @. ngrids ÷ 2 + 1
    if any(max_nwaves .< 0)
        println("WARNING: all waves are suppressed")
        f.vals .= 0.
    else
        floors = tuple(ones(Int, N)...)
        ceils = ngrids
        min_indices = (
            @. max(floors, center_indices - max_nwaves)
        )
        max_indices = (
            @. min(ceils, center_indices + max_nwaves)
        )

        slices = (:).(min_indices, max_indices)
        pass_K!(f, slices)
    end

    f.vals .= ifftshift(f.vals)

end

function K_lowpass_K(
        f::KFunc{T,N},
        max_nwaves::NTuple{N,Int}
    ) where {T,N}

    g = copy(f)
    lowpass_K!(g, max_nwaves)
    return g

end

# *******************************************
#  Fourier Transformation
# *******************************************
function K_X(f::XFunc{T,N}) where {T,N}

    P = f.config.P_fft
    g = KFunc(P * f.vals, f.config)

    if f.config.cut_zigzag_mode
        ngrids = f.config.ngrids
        max_nwaves = @. ngrids ÷ 2 - 1
        lowpass_K!(g, max_nwaves)
    end

    return g
end

function X_K(f::KFunc{T,N}) where {T,N}

    P = f.config.P_ifft
    if T <: Real
        XFunc(real(P * f.vals), f.config)
    else
        XFunc(P * f.vals, f.config)
    end
end

# *******************************************
#  Differentiation
# *******************************************
# +++++ destructive +++++
function ∂Xaxis_K!(f::KFunc, axis::Int)

    config = f.config
    xrange = config.xranges[axis]
    xlen = xrange[2] - xrange[1]
    Kcoord = config.Kcoords[axis]

    f.vals .*= (2π*im/xlen)*Kcoord

end

# +++++ non-destructive +++++
function K_∂Xaxis_K(f::KFunc, axis::Int)
    g = copy(f)
    ∂Xaxis_K!(g, axis)
    return g
end

function X_∂Xaxis_X(f::XFunc, axis::Int)
    X_K(K_∂Xaxis_K(K_X(f), axis))
end

function K_laplacian_K(f::KFunc{T,N}) where {T,N}

    return sum(
        K_∂Xaxis_K(K_∂Xaxis_K(f, axis), axis)
        for axis = 1:N
    )

end

K_Δ_K = K_laplacian_K
X_laplacian_X = X_K ∘ K_laplacian_K ∘ K_X
X_Δ_X = X_laplacian_X

"""
K_laplainv_K returns g satisfying Δg = f
"""
function K_laplainv_K(f::KFunc{T,N}) where {T,N}

    c = f.config
    xlens = xlens_xranges(c.xranges)

    lap = - 4π^2 * sum(
        (x -> x .^ 2).(c.Kcoords) ./ (xlens .^ 2)
    )
    gvals = f.vals ./ lap

    zero_indiceses = (
        map(c.ngrids) do x
            if x % 2 == 0
                return (1, x÷2 + 1)
            else
                return (1,)
            end
        end
    )
    for indices in Iterators.product(zero_indiceses...)
        gvals[indices...] = 0.
    end

    return KFunc(gvals, c)

end

K_Δ⁻¹_K = K_laplainv_K
X_laplainv_X = X_K ∘ K_laplainv_K ∘ K_X
X_Δ⁻¹_X = X_laplainv_X

# *******************************************
#  De-aliased Products
# *******************************************
# 3/2-rule ... zero-padding
# 2/3-rule ... truncating

# de-aliased product by 3/2-rule (zero padding)
function K_prod_32_K_K(
        f::KFunc{T,N}, g::KFunc{T,N}
    ) where {T,N}

    check_config_consistency(f.config, g.config)

    fvals_pad = padding(f)
    gvals_pad = padding(g)

    P_fft = f.config.P_fftpad
    P_ifft = f.config.P_ifftpad
    if T <: Real
        fgvals_pad = P_fft * (
            real(P_ifft * fvals_pad) .* real(P_ifft * gvals_pad)
        )
    else
        fgvals_pad = P_fft * (
            (P_ifft * fvals_pad) .* (P_ifft * gvals_pad)
        )
    end
    fgvals = truncate(fgvals_pad, f.config)

    return KFunc(fgvals, f.config)

end

# dot-product
function K_dot_32_Ks_Ks(
        f::Array{KFunc{T,N},1},
        g::Array{KFunc{T,N},1}
    ) where {T,N}

    check_length_consistency(f, g)
    ndim = length(f)
    for dim = 1:ndim
        check_config_consistency(
            f[dim].config, g[dim].config
        )
        if dim < ndim
            check_config_consistency(
                f[dim].config, f[dim+1].config
            )
        end
    end

    config = f[1].config

    fvals_pad = padding.(f)
    gvals_pad = padding.(g)

    P_fft = config.P_fftpad
    P_ifft = config.P_ifftpad

    fgvals_pad = P_fft * sum(
        (P_ifft * fvals_pad[dim]) .* (P_ifft * gvals_pad[dim])
        for dim = 1:ndim
    )

    fgvals = truncate(fgvals_pad, config)
    return KFunc(fgvals, config)

end

# cross-product
function K3_cross_32_K3_K3(
        f::Array{KFunc{T,N},1},
        g::Array{KFunc{T,N},1}
    ) where {T,N}

    if !(length(f) == length(g) == 3)
        errmsg = (
            "lengths of arrays must be 3"
        )
        throw(ErrorException(errmsg))
    end

    config = f[1].config

    fvals_pad = padding.(f)
    gvals_pad = padding.(g)

    P_fft = config.P_fftpad
    P_ifft = config.P_ifftpad

    xfvals_pad, yfvals_pad, zfvals_pad = [
        P_ifft * fvals_pad[dim] for dim = 1:3
    ]
    xgvals_pad, ygvals_pad, zgvals_pad = [
        P_ifft * gvals_pad[dim] for dim = 1:3
    ]

    fgvals = [
        truncate( P_fft * (
            yfvals_pad .* zgvals_pad .- zfvals_pad .* ygvals_pad
        ), config),
        truncate( P_fft * (
            zfvals_pad .* xgvals_pad .- xfvals_pad .* zgvals_pad
        ), config),
        truncate( P_fft * (
            xfvals_pad .* ygvals_pad .- yfvals_pad .* xgvals_pad
        ), config),
    ]

    return [
        KFunc(fgvals[dim], config) for dim = 1:3
    ]

end

# padding helper functions
function padding(f::KFunc)
    config = f.config

    ngrids = config.ngrids
    pad_ngrids = to_pad_ngrids(ngrids)

    slices = slices_padded_core(ngrids, pad_ngrids)

    vals_shift = fftshift(f.vals)
    padded = zeros(Complex{Float64}, pad_ngrids)
    padded[slices...] = vals_shift
    padded .= ifftshift(padded)

    return padded
end

function truncate(
        padded::Array{Complex{Float64},N},
        config::ConfigFFT{T,N}
    ) where {T,N}

    ngrids = config.ngrids
    pad_ngrids = to_pad_ngrids(ngrids)

    slices = slices_padded_core(ngrids, pad_ngrids)
    vals = ifftshift(fftshift(padded)[slices...])

    N_origin = prod(ngrids)
    N_padded = prod(pad_ngrids)
    vals *= N_padded / N_origin

    return vals

end

# helper function
function slices_padded_core(
        ngrids::NTuple{N1,Int},
        pad_ngrids::NTuple{N2,Int}
    ) where {N1,N2}

    nshifts = @. pad_ngrids÷2 - ngrids÷2
    min_nwaves = @. nshifts + 1
    max_nwaves = @. nshifts + ngrids
    return (:).(min_nwaves, max_nwaves)

end

function to_pad_ngrids(ngrids::NTuple{N,Int}) where N
    @. ngrids ÷ 2 * 3
end

# de-aliased product by 2/3-rule (truncation)
function K_prod_23_K_K(f::KFunc, g::KFunc)

    check_config_consistency(f.config, g.config)
    config = f.config
    ngrids = config.ngrids
    max_nwaves = @. ngrids ÷ 3
    f_trunc = K_lowpass_K(f, max_nwaves)
    g_trunc = K_lowpass_K(g, max_nwaves)
    return K_X(X_K(f_trunc) * X_K(g_trunc))

end

# dot-product
function K_dot_23_Ks_Ks(
        f::Array{KFunc{T,N},1},
        g::Array{KFunc{T,N},1}
    ) where {T,N}

    check_length_consistency(f, g)
    ndim = length(f)
    for dim = 1:ndim
        check_config_consistency(
            f[dim].config, g[dim].config
        )
        if dim < ndim
            check_config_consistency(
                f[dim].config, f[dim+1].config
            )
        end
    end

    config = f[1].config
    ngrids = config.ngrids
    max_nwaves = ngrids .÷ 3
    f_trunc = [
        K_lowpass_K(f[dim], max_nwaves)
        for dim = 1:ndim
    ]
    g_trunc = [
        K_lowpass_K(g[dim], max_nwaves)
        for dim = 1:ndim
    ]
    return K_X(
        sum( X_K.(f_trunc) .* X_K.(g_trunc) )
    )

end

# cross-product
function K3_cross_23_K3_K3(
        f::Array{KFunc{T,N},1},
        g::Array{KFunc{T,N},1}
    ) where {T,N}

    if !(length(f) == length(g) == 3)
        errmsg = (
            "lengths of arrays must be 3"
        )
        throw(ErrorException(errmsg))
    end

    config = f[1].config
    ngrids = config.ngrids
    max_nwaves = ngrids .÷ 3
    X_xf, X_yf, X_zf = X_K.([
        K_lowpass_K(f[dim], max_nwaves) for dim = 1:3
    ])
    X_xg, X_yg, X_zg = X_K.([
        K_lowpass_K(g[dim], max_nwaves) for dim = 1:3
    ])

    return K_X.([
        X_yf * X_zg - X_zf * X_yg,
        X_zf * X_xg - X_xf * X_zg,
        X_xf * X_yg - X_yf * X_xg
    ])

end

# +++++ aliases +++++
# --- 3/2-rule ---
# \ast
∗(f, g) = K_prod_32_K_K(f, g)
# \cdot
⋅(f, g) = K_dot_32_Ks_Ks(f, g)
# \times
×(f, g) = K3_cross_32_K3_K3(f, g)

# --- 2/3-rule ---
# \circledast
⊛(f, g) = K_prod_23_K_K(f, g)
# \odot
⊙(f, g) = K_dot_23_Ks_Ks(f, g)
# \otimes
⊗(f, g) = K3_cross_23_K3_K3(f, g)


using ImageBase
using ImageBase: varmult
using Statistics
using Test

@testset "Reductions" begin
    _abs(x::Colorant) = mapreducec(abs, +, 0, x)

    @testset "sumfinite, meanfinite, varfinite" begin
        for T in generate_test_types([N0f8, Float32], [Gray, RGB])
            A = rand(T, 5, 5)
            s12 = sum(A, dims=(1,2))
            @test eltype(s12) <: Union{T, float(T), float64(T)}

            @test sumfinite(A) ≈ sum(A)
            @test sumfinite(A, dims=1) ≈ sum(A, dims=1)
            @test sumfinite(A, dims=(1, 2)) ≈ sum(A, dims=(1, 2))

            @test meanfinite(A) ≈ mean(A)
            @test meanfinite(A, dims=1) ≈ mean(A, dims=1)
            @test meanfinite(A, dims=(1, 2)) ≈ mean(A, dims=(1, 2))

            @test varfinite(A) ≈ varmult(⋅, A)
            @test varfinite(A, dims=1) ≈ varmult(⋅, A, dims=1)
            @test varfinite(A, dims=(1, 2)) ≈ varmult(⋅, A, dims=(1, 2))

            # test NaN/Inf
            if eltype(T) != N0f8
                A = rand(T, 5, 5) .- 0.5 .* oneunit(T)
                A[1] = Inf
                @test sum(A) ≈ A[1]
                @test sum(abs, A) ≈ A[1]
                @test sumfinite(A) ≈ sum(A[2:end])
                @test sumfinite(abs, A) ≈ sum(abs, A[2:end])
                A[1] = NaN
                @test isnan(sum(A))
                @test isnan(sum(abs, A))
                @test sumfinite(A) ≈ sum(A[2:end])
                @test sumfinite(abs, A) ≈ sum(abs, A[2:end])

                A = rand(T, 5, 5) .- 0.5 .* oneunit(T)
                A[1] = Inf
                @test mean(A) ≈ A[1]
                @test mean(abs, A) ≈ A[1]
                @test meanfinite(A) ≈ mean(A[2:end])
                @test meanfinite(abs, A) ≈ mean(abs, A[2:end])
                A[1] = NaN
                @test isnan(mean(A))
                @test isnan(mean(abs, A))
                @test meanfinite(A) ≈ mean(A[2:end])
                @test meanfinite(abs, A) ≈ mean(abs, A[2:end])

                A = rand(T, 5, 5)
                A[1] = Inf
                @test isnan(varmult(⋅, A))
                @test varfinite(A) ≈ varmult(⋅, A[2:end])
                A[1] = NaN
                @test isnan(varmult(⋅, A))
                @test varfinite(A) ≈ varmult(⋅, A[2:end])
            end
        end

        A = [NaN, 1, 2, 3]
        @test meanfinite(A, dims=1) ≈ [2]
        @test varfinite(A, dims=1) ≈ [1]

        A = [NaN NaN 1;
            1 2 3]
        vf = varfinite(A, dims=2)
        @test isnan(vf[1])

        A = [NaN 1 2 3;
            NaN 6 5 4]
        mf = meanfinite(A, dims=1)
        vf = varfinite(A, dims=1)
        @test isnan(mf[1])
        @test mf[2:end] ≈ [3.5,3.5,3.5]
        @test isnan(vf[1])
        @test vf[2:end] ≈ [12.5,4.5,0.5]

        @test meanfinite(A, dims=2) ≈ reshape([2, 5], 2, 1)
        @test varfinite(A, dims=2) ≈ reshape([1, 1], 2, 1)

        @test meanfinite(A, dims=(1,2)) ≈ [3.5]
        @test varfinite(A, dims=(1,2)) ≈ [3.5]

        # Ensure we're consistant with our decision to `abs2` in ColorVectorSpace
        # See also: https://github.com/JuliaGraphics/ColorVectorSpace.jl/blob/master/README.md#abs-and-abs2
        A = rand(Gray, 4, 4)
        @test varfinite(A) ≈ varfinite(RGB.(A))
        A[1] = Inf
        @test varfinite(A) ≈ varfinite(RGB.(A))
        A[1] = NaN
        @test varfinite(A) ≈ varfinite(RGB.(A))
    end

    @testset "minfinite, maxfinite, maxabsfinite" begin
        for T in generate_test_types([N0f8, Float32], [Gray, ])
            A = rand(T, 5, 5)
            @test @inferred(minimum_finite(A)) == minimum(A)
            @test @inferred(maximum_finite(A)) == maximum(A)
            @test minimum_finite(A; dims=1) == minimum(A; dims=1)
            @test maximum_finite(A; dims=1) == maximum(A; dims=1)
            @test_broken @inferred maximum_finite(A; dims=1)
            @test_broken @inferred minimum_finite(A; dims=1)

            @test maximum_finite(abs2, A) == maximum(abs2, A)
            @test_broken @inferred maximum(abs2, A)
            @test_broken @inferred minimum(abs2, A)

            if eltype(T) != N0f8
                A = rand(T, 5, 5) .- 0.5 * rand(T, 5, 5)
                A[1] = Inf

                @test @inferred(minimum_finite(A)) == minimum(A[2:end])
                @test @inferred(maximum_finite(A)) == maximum(A[2:end])
                @test minimum_finite(abs, A) == minimum(abs, A[2:end])
                @test maximum_finite(abs2, A) == maximum(abs2, A[2:end])
            end
        end

        # minimum_finite and maximum_finite for RGB are processed per channel
        A = rand(RGB{Float32}, 5, 5)
        @test minimum_finite(A) == RGB(minimum(channelview(A), dims=(2, 3))...)
        @test maximum_finite(A) == RGB(maximum(channelview(A), dims=(2, 3))...)

        # Container of abstract type
        A = Any[1, 2, Gray(0.3)]
        @test minimum_finite(A) == 0.3
        @test maximum_finite(A) == 2.0
        @test Base.return_types(minimum_finite, (typeof(A),)) == Base.return_types(minimum, (typeof(A), ))
    end

end


export searchMatrix

"""
[240] Medium [Search a 2D Matrix II](https://leetcode.com/problems/search-a-2d-matrix-ii/description/)

Write an efficient algorithm that searches for a value in an *m* x *n* matrix. This matrix has the following properties:

+ Integers in each row are sorted in ascending from left to right.
+ Integers in each column are sorted in ascending from top to bottom.

## Example:

Consider the following matrix:

```
[
    [1,   4,  7, 11, 15],
    [2,   5,  8, 12, 19],
    [3,   6,  9, 16, 22],
    [10, 13, 14, 17, 24],
    [18, 21, 23, 26, 30]
]
```

Given target = `5`, return `true`.

Given target = `20`, return `false`.
"""
function searchMatrix(matrix::Matrix{Int}, target::Int)::Bool
    row, col = size(matrix)
    i, j = row, 1

    while (i > 0 && j <= col)
        last = matrix[i, j]

        if (last == target)
            return true
        elseif last > target
            i -= 1
        else
            j += 1
        end
    end

    false
end


using Winston
p = plot(abs(sin(0:.1:10)))
setattr(p, "ylog", true)
setattr(p, "yrange", (5e-6,0.2))


export LinearVApproximator, LinearQApproximator

using LinearAlgebra: dot

"""
    LinearVApproximator(weights::Array{Float64, N}) -> LinearVApproximator{N}
Use the weighted sum to represent the estimation of a state.
The state is expected to have the same length with `weights`.
See also [`LinearQApproximator`](@ref)
"""
struct LinearVApproximator{N} <: AbstractApproximator
    weights::Array{Float64,N}
end

RLCore.ApproximatorStyle(::LinearVApproximator) = V_APPROXIMATOR

# TODO: support Vector
(V::LinearVApproximator)(s) = dot(s, V.weights)

function RLBase.update!(V::LinearVApproximator, correction::Pair)
    s, e = correction
    V.weights .+= s .* e
end

"""
    LinearQApproximator(weights::Vector{Float64}, feature_func::F, actions::Vector{Int}) -> LinearQApproximator{F}
Use weighted sum to represent the estimation given a state and an action.
# Fields
- `weights::Vector{Float64}`: the weight of each feature.
- `feature_func::Function`: decide how to generate a feature vector of `length(weights)` given a state and an action as parameters.
- `actions::Vector{Int}`: all possible actions.
See also [`LinearVApproximator`](@ref).
"""
Base.@kwdef struct LinearQApproximator{F} <: AbstractApproximator
    weights::Vector{Float64}
    feature_func::F
    actions::Vector{Int}
end

RLCore.ApproximatorStyle(::LinearQApproximator) = Q_APPROXIMATOR

(Q::LinearQApproximator)(s, a::Int) = dot(Q.weights, Q.feature_func(s, a))

(Q::LinearQApproximator)(s) = [Q(s, a) for a in Q.actions]

function RLBase.update!(Q::LinearQApproximator, correction::Pair)
    (s, a), e = correction
    xs = Q.feature_func(s, a)
    Q.weights .+= xs .* e
end


# Advent of Code 2020 - Day 24

test_str = "sesenwnenenewseeswwswswwnenewsewsw
neeenesenwnwwswnenewnwwsewnenwseswesw
seswneswswsenwwnwse
nwnwneseeswswnenewneswwnewseswneseene
swweswneswnenwsewnwneneseenw
eesenwseswswnenwswnwnwsewwnwsene
sewnenenenesenwsewnenwwwse
wenwwweseeeweswwwnwwe
wsweesenenewnwwnwsenewsenwwsesesenwne
neeswseenwwswnwswswnw
nenwswwsewswnenenewsenwsenwnesesenew
enewnwewneswsewnwswenweswnenwsenwsw
sweneswneswneneenwnewenewwneswswnese
swwesenesewenwneswnwwneseswwne
enesenwswwswneneswsenwnewswseenwsese
wnwnesenesenenwwnenwsewesewsesesew
nenewswnwewswnenesenwnesewesw
eneswnwswnwsenenwnwnwwseeswneewsenese
neswnwewnwnwseenwseesewsenwsweewe
wseweeenwnesenwwwswnew"

test_list = split(test_str, '\n')

list = readlines("day24_input.txt")

function get_blacks(list)
    blacks = Set()
    for i = 1:length(list)
        steps = list[i]
        for subs in ("ne" => "o", "nw" => "m", "se" => "t", "sw" => "r")
            steps = replace(steps, subs)
        end
        position = [count(==('e'), steps) - count(==('w'), steps) +
                        (count(==('o'), steps) + count(==('t'), steps))//2 -
                        (count(==('m'), steps) + count(==('r'), steps))//2
                    (count(==('o'), steps) + count(==('m'), steps))//2 -
                        (count(==('r'), steps) + count(==('t'), steps))//2]
        if position ∈ blacks
            delete!(blacks, position)
        else
            push!(blacks, position)
        end
    end
    return blacks
end

function get_blacks_int(list)
    blacks = Set()
    for i = 1:length(list)
        steps = replace(replace(replace(replace(list[i], "ne" => "o"), "nw" => "m"), "se" => "t"), "sw" => "r")
        position = [2*(count(==('e'), steps) - count(==('w'), steps)) + count(==('o'), steps) + count(==('t'), steps) - count(==('m'), steps) - count(==('r'), steps)
                    count(==('o'), steps) + count(==('m'), steps) - count(==('r'), steps) - count(==('t'), steps)]
        if position ∈ blacks
            delete!(blacks, position)
        else
            push!(blacks, position)
        end
    end
    return blacks
end

@show length(get_blacks(test_list)) == 10
@show length(get_blacks(list)) == 356
@show length(get_blacks_int(test_list)) == 10
@show length(get_blacks_int(list)) == 356

#=
# A slight gain with an integer lattice
```julia
julia> @btime length(count_blacks(list));
  1.628 ms (22294 allocations: 1.37 MiB)

julia> @btime length(count_blacks_int(list));
  1.489 ms (22260 allocations: 1.36 MiB)
```
 =#

function count_blacks_moving(list, moves)
    blacks = Set()
    function get_neighbors(position)
        displacements = [[1,0], [1//2,1//2], [1//2,-1//2], [-1,0],[-1//2,-1//2],[-1//2,1//2]]
        return [position .+ displacement for displacement in displacements]
    end
    blacks = get_blacks(list)
    for n = 1:moves
        new_blacks = Set()
        for position in blacks
            neighbors = get_neighbors(position)
            if length(neighbors ∩ blacks) ∈ (1,2)
                push!(new_blacks, position)
            end
            for neighbor in neighbors
                if neighbor ∉ blacks && length(get_neighbors(neighbor) ∩ blacks) == 2
                    push!(new_blacks, neighbor)
                end
            end
        end
        blacks = new_blacks
    end
    return length(blacks)
end

function count_blacks_moving_int(list, moves)
    blacks = Set()
    function get_neighbors(position)
        displacements = [[2,0], [1,1], [1,-1], [-2,0],[-1,-1],[-1,1]]
        return [position .+ displacement for displacement in displacements]
    end
    blacks = get_blacks_int(list)
    for n = 1:moves
        new_blacks = Set()
        for position in blacks
            neighbors = get_neighbors(position)
            if length(neighbors ∩ blacks) ∈ (1,2)
                push!(new_blacks, position)
            end
            for neighbor in neighbors
                if neighbor ∉ blacks && length(get_neighbors(neighbor) ∩ blacks) == 2
                    push!(new_blacks, neighbor)
                end
            end
        end
        blacks = new_blacks
    end
    return length(blacks)
end

@show count_blacks_moving(test_list, 100) == 2208
@show count_blacks_moving(list, 100) == 3887
@show count_blacks_moving_int(test_list, 100) == 2208
@show count_blacks_moving_int(list, 100) == 3887

nothing
#=
```julia
julia> @btime count_blacks_moving(list, 100)
  9.949 s (19853343 allocations: 1.87 GiB)

julia> @btime count_blacks_moving_int(list, 100)
  6.001 s (17462036 allocations: 1.54 GiB)
```
 =#


struct GLMTrait{D<:ResponseDistribution, L<:InverseLinkFunction}
formula::String
mu::Vector{Float64}
dist:: D
link:: L
end

function GLMTrait(mu::Number, df, dist::D, link::L) where {D, L}
    return(GLMTrait{D, L}(string(mu), repeat([mu], size(df, 1)), dist, link))
end

function GLMTrait(formula::String, df, dist::D, link::L) where {D, L}
    mu = mean_formula(formula, df)
    return(GLMTrait{D, L}(formula, mu, dist, link))
end

function Multiple_GLMTraits(formulas, df, dist::ResponseDistribution, link::InverseLinkFunction)
  vec = [GLMTrait(formulas[i], df, dist, link) for i in 1:length(formulas)] #vector of GLMTrait objects
  return(vec)
end

# we put type of the dist vector as Any since we want to allow for any ResponseType{Poisson(), LogLink()}, ResponseType{Normal(), IdentityLink()}
function Multiple_GLMTraits(formulas::Vector{String}, df::DataFrame, dist::Vector, link::Vector)
  vec = [GLMTrait(formulas[i], df, dist[i], link[i]) for i in 1:length(formulas)]
  return(vec)
end

# lmm: multiple traits (MVN)

struct LMMTrait{T}
formulas::Vector{String}
mu::Matrix{Float64}
vc::T
  function LMMTrait(formulas, df, vc::T) where T
    n_traits = length(formulas)
    n_people = size(df)[1]
    mu = zeros(n_people, n_traits)
    for i in 1:n_traits
      #calculate the mean vector
      mu[:, i] += mean_formula(formulas[i], df)
    end
    return(new{T}(formulas, mu, vc))
  end
end


@testset "minimal_cancel_plus_test" begin
    @test cancelplus(interval(-Inf, -1.0), emptyinterval()) == entireinterval()
    @test cancelplus(interval(-1.0, Inf), emptyinterval()) == entireinterval()
    @test cancelplus(entireinterval(), emptyinterval()) == entireinterval()
    @test cancelplus(interval(-Inf, -1.0), interval(-5.0,1.0)) == entireinterval()
    @test cancelplus(interval(-1.0, Inf), interval(-5.0,1.0)) == entireinterval()
    @test cancelplus(entireinterval(), interval(-5.0,1.0)) == entireinterval()
    @test cancelplus(interval(-Inf, -1.0), entireinterval()) == entireinterval()
    @test cancelplus(interval(-1.0, Inf), entireinterval()) == entireinterval()
    @test cancelplus(emptyinterval(), interval(1.0, Inf)) == entireinterval()
    @test cancelplus(emptyinterval(), interval(-Inf,1.0)) == entireinterval()
    @test cancelplus(emptyinterval(), entireinterval()) == entireinterval()
    @test cancelplus(interval(-1.0,5.0), interval(1.0,Inf)) == entireinterval()
    @test cancelplus(interval(-1.0,5.0), interval(-Inf,1.0)) == entireinterval()
    @test cancelplus(interval(-1.0,5.0), entireinterval()) == entireinterval()
    @test cancelplus(entireinterval(), interval(1.0,Inf)) == entireinterval()
    @test cancelplus(entireinterval(), interval(-Inf,1.0)) == entireinterval()
    @test cancelplus(entireinterval(), entireinterval()) == entireinterval()
    @test cancelplus(interval(-5.0, -1.0), interval(1.0,5.1)) == entireinterval()
    @test cancelplus(interval(-5.0, -1.0), interval(0.9,5.0)) == entireinterval()
    @test cancelplus(interval(-5.0, -1.0), interval(0.9,5.1)) == entireinterval()
    @test cancelplus(interval(-10.0, 5.0), interval(-5.0,10.1)) == entireinterval()
    @test cancelplus(interval(-10.0, 5.0), interval(-5.1,10.0)) == entireinterval()
    @test cancelplus(interval(-10.0, 5.0), interval(-5.1,10.1)) == entireinterval()
    @test cancelplus(interval(1.0, 5.0), interval(-5.0,-0.9)) == entireinterval()
    @test cancelplus(interval(1.0, 5.0), interval(-5.1,-1.0)) == entireinterval()
    @test cancelplus(interval(1.0, 5.0), interval(-5.1,-0.9)) == entireinterval()
    @test cancelplus(interval(-10.0, -1.0), emptyinterval()) == entireinterval()
    @test cancelplus(interval(-10.0, 5.0), emptyinterval()) == entireinterval()
    @test cancelplus(interval(1.0, 5.0), emptyinterval()) == entireinterval()
    @test cancelplus(emptyinterval(), emptyinterval()) == emptyinterval()
    @test cancelplus(emptyinterval(), interval(1.0,10.0)) == emptyinterval()
    @test cancelplus(emptyinterval(), interval(-5.0,10.0)) == emptyinterval()
    @test cancelplus(emptyinterval(), interval(-5.0,-1.0)) == emptyinterval()
    @test cancelplus(interval(-5.1,-0.0), interval(0.0,5.0)) == interval(-0x1.999999999998P-4,0.0)
    @test cancelplus(interval(-5.1,-1.0), interval(1.0,5.0)) == interval(-0x1.999999999998P-4,0.0)
    @test cancelplus(interval(-5.0,-0.9), interval(1.0,5.0)) == interval(0.0, 0x1.9999999999998P-4)
    @test cancelplus(interval(-5.1,-0.9), interval(1.0,5.0)) == interval(-0x1.999999999998P-4,0x1.9999999999998P-4)
    @test cancelplus(interval(-5.0,-1.0), interval(1.0,5.0)) == interval(0.0,0.0)
    @test cancelplus(interval(-10.1, 5.0), interval(-5.0,10.0)) == interval(-0x1.999999999998P-4,0.0)
    @test cancelplus(interval(-10.0, 5.1), interval(-5.0,10.0)) == interval(0.0,0x1.999999999998P-4)
    @test cancelplus(interval(-10.1, 5.1), interval(-5.0,10.0)) == interval(-0x1.999999999998P-4,0x1.999999999998P-4)
    @test cancelplus(interval(-10.0, 5.0), interval(-5.0,10.0)) == interval(0.0,0.0)
    @test cancelplus(interval(0.9, 5.0), interval(-5.0,-1.0)) == interval(-0x1.9999999999998P-4,0.0)
    @test cancelplus(interval(1.0, 5.1), interval(-5.0,-1.0)) == interval(0.0,0x1.999999999998P-4)
    @test cancelplus(interval(0.0, 5.1), interval(-5.0,-0.0)) == interval(0.0,0x1.999999999998P-4)
    @test cancelplus(interval(0.9, 5.1), interval(-5.0,-1.0)) == interval(-0x1.9999999999998P-4,0x1.999999999998P-4)
    @test cancelplus(interval(1.0, 5.0), interval(-5.0,-1.0)) == interval(0.0,0.0)
    @test cancelplus(interval(0.0, 5.0), interval(-5.0,-0.0)) == interval(0.0,0.0)
    @test cancelplus(interval(0x1.FFFFFFFFFFFFP+0,0x1.FFFFFFFFFFFFP+0), interval(-0x1.999999999999AP-4,-0x1.999999999999AP-4)) == interval(0x1.E666666666656P+0,0x1.E666666666657P+0)
    @test cancelplus(interval(-0x1.999999999999AP-4,0x1.FFFFFFFFFFFFP+0), interval(-0x1.999999999999AP-4,0.01)) == interval(-0x1.70A3D70A3D70BP-4,0x1.E666666666657P+0)
    @test cancelplus(interval(0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == interval(0x1.FFFFFFFFFFFFFp1023,Inf)
    @test cancelplus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == interval(0.0,0.0)
    @test cancelplus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFEP+1023,0x1.FFFFFFFFFFFFFp1023)) == interval(0.0,0x1P+971)
    @test cancelplus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFEP+1023)) == interval(-0x1P+971,0.0)
    @test cancelplus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFEP+1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == entireinterval()
    @test cancelplus(interval(-0x1.FFFFFFFFFFFFEP+1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == entireinterval()
    @test cancelplus(interval(-0x1P+0,0x1.FFFFFFFFFFFFFP-53), interval(-0x1P+0,0x1.FFFFFFFFFFFFEP-53)) == interval(-0x1.FFFFFFFFFFFFFP-1,-0x1.FFFFFFFFFFFFEP-1)
    @test cancelplus(interval(-0x1P+0,0x1.FFFFFFFFFFFFEP-53), interval(-0x1P+0,0x1.FFFFFFFFFFFFFP-53)) == entireinterval()
end

@testset "minimal_cancel_minus_test" begin
    @test cancelminus(interval(-Inf, -1.0), emptyinterval()) == entireinterval()
    @test cancelminus(interval(-1.0, Inf), emptyinterval()) == entireinterval()
    @test cancelminus(entireinterval(), emptyinterval()) == entireinterval()
    @test cancelminus(interval(-Inf, -1.0), interval(-1.0,5.0)) == entireinterval()
    @test cancelminus(interval(-1.0, Inf), interval(-1.0,5.0)) == entireinterval()
    @test cancelminus(entireinterval(), interval(-1.0,5.0)) == entireinterval()
    @test cancelminus(interval(-Inf, -1.0), entireinterval()) == entireinterval()
    @test cancelminus(interval(-1.0, Inf), entireinterval()) == entireinterval()
    @test cancelminus(emptyinterval(), interval(-Inf, -1.0)) == entireinterval()
    @test cancelminus(emptyinterval(), interval(-1.0, Inf)) == entireinterval()
    @test cancelminus(emptyinterval(), entireinterval()) == entireinterval()
    @test cancelminus(interval(-1.0,5.0), interval(-Inf, -1.0)) == entireinterval()
    @test cancelminus(interval(-1.0,5.0), interval(-1.0, Inf)) == entireinterval()
    @test cancelminus(interval(-1.0,5.0), entireinterval()) == entireinterval()
    @test cancelminus(entireinterval(), interval(-Inf, -1.0)) == entireinterval()
    @test cancelminus(entireinterval(), interval(-1.0, Inf)) == entireinterval()
    @test cancelminus(entireinterval(), entireinterval()) == entireinterval()
    @test cancelminus(interval(-5.0, -1.0), interval(-5.1,-1.0)) == entireinterval()
    @test cancelminus(interval(-5.0, -1.0), interval(-5.0,-0.9)) == entireinterval()
    @test cancelminus(interval(-5.0, -1.0), interval(-5.1,-0.9)) == entireinterval()
    @test cancelminus(interval(-10.0, 5.0), interval(-10.1, 5.0)) == entireinterval()
    @test cancelminus(interval(-10.0, 5.0), interval(-10.0, 5.1)) == entireinterval()
    @test cancelminus(interval(-10.0, 5.0), interval(-10.1, 5.1)) == entireinterval()
    @test cancelminus(interval(1.0, 5.0), interval(0.9, 5.0)) == entireinterval()
    @test cancelminus(interval(1.0, 5.0), interval(1.0, 5.1)) == entireinterval()
    @test cancelminus(interval(1.0, 5.0), interval(0.9, 5.1)) == entireinterval()
    @test cancelminus(interval(-10.0, -1.0), emptyinterval()) == entireinterval()
    @test cancelminus(interval(-10.0, 5.0), emptyinterval()) == entireinterval()
    @test cancelminus(interval(1.0, 5.0), emptyinterval()) == entireinterval()
    @test cancelminus(emptyinterval(), emptyinterval()) == emptyinterval()
    @test cancelminus(emptyinterval(), interval(-10.0, -1.0)) == emptyinterval()
    @test cancelminus(emptyinterval(), interval(-10.0, 5.0)) == emptyinterval()
    @test cancelminus(emptyinterval(), interval(1.0, 5.0)) == emptyinterval()
    @test cancelminus(interval(-5.1,-0.0), interval(-5.0, 0.0)) == interval(-0x1.999999999998P-4,0.0)
    @test cancelminus(interval(-5.1,-1.0), interval(-5.0, -1.0)) == interval(-0x1.999999999998P-4,0.0)
    @test cancelminus(interval(-5.0,-0.9), interval(-5.0, -1.0)) == interval(0.0, 0x1.9999999999998P-4)
    @test cancelminus(interval(-5.1,-0.9), interval(-5.0, -1.0)) == interval(-0x1.999999999998P-4,0x1.9999999999998P-4)
    @test cancelminus(interval(-5.0,-1.0), interval(-5.0, -1.0)) == interval(0.0,0.0)
    @test cancelminus(interval(-10.1, 5.0), interval(-10.0, 5.0)) == interval(-0x1.999999999998P-4,0.0)
    @test cancelminus(interval(-10.0, 5.1), interval(-10.0, 5.0)) == interval(0.0,0x1.999999999998P-4)
    @test cancelminus(interval(-10.1, 5.1), interval(-10.0, 5.0)) == interval(-0x1.999999999998P-4,0x1.999999999998P-4)
    @test cancelminus(interval(-10.0, 5.0), interval(-10.0, 5.0)) == interval(0.0,0.0)
    @test cancelminus(interval(0.9, 5.0), interval(1.0, 5.0)) == interval(-0x1.9999999999998P-4,0.0)
    @test cancelminus(interval(-0.0, 5.1), interval(0.0, 5.0)) == interval(0.0,0x1.999999999998P-4)
    @test cancelminus(interval(1.0, 5.1), interval(1.0, 5.0)) == interval(0.0,0x1.999999999998P-4)
    @test cancelminus(interval(0.9, 5.1), interval(1.0, 5.0)) == interval(-0x1.9999999999998P-4,0x1.999999999998P-4)
    @test cancelminus(interval(1.0, 5.0), interval(1.0, 5.0)) == interval(0.0,0.0)
    @test cancelminus(interval(-5.0, 1.0), interval(-1.0, 5.0)) == interval(-4.0,-4.0)
    @test cancelminus(interval(-5.0, 0.0), interval(-0.0, 5.0)) == interval(-5.0,-5.0)
    @test cancelminus(interval(0x1.FFFFFFFFFFFFP+0,0x1.FFFFFFFFFFFFP+0), interval(0x1.999999999999AP-4,0x1.999999999999AP-4)) == interval(0x1.E666666666656P+0,0x1.E666666666657P+0)
    @test cancelminus(interval(-0x1.999999999999AP-4,0x1.FFFFFFFFFFFFP+0), interval(-0.01,0x1.999999999999AP-4)) == interval(-0x1.70A3D70A3D70BP-4,0x1.E666666666657P+0)
    @test cancelminus(interval(0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,-0x1.FFFFFFFFFFFFFp1023)) == interval(0x1.FFFFFFFFFFFFFp1023,Inf)
    @test cancelminus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == interval(0.0,0.0)
    @test cancelminus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFEP+1023)) == interval(0.0,0x1P+971)
    @test cancelminus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFEP+1023,0x1.FFFFFFFFFFFFFp1023)) == interval(-0x1P+971,0.0)
    @test cancelminus(interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFEP+1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == entireinterval()
    @test cancelminus(interval(-0x1.FFFFFFFFFFFFEP+1023,0x1.FFFFFFFFFFFFFp1023), interval(-0x1.FFFFFFFFFFFFFp1023,0x1.FFFFFFFFFFFFFp1023)) == entireinterval()
    @test cancelminus(interval(0x0.0000000000001p-1022,0x0.0000000000001p-1022), interval(0x0.0000000000001p-1022,0x0.0000000000001p-1022)) == interval(0.0,0.0)
    @test cancelminus(interval(0x0.0000000000001p-1022,0x0.0000000000001p-1022), interval(-0x0.0000000000001p-1022,-0x0.0000000000001p-1022)) == interval(0x0.0000000000002p-1022,0x0.0000000000002p-1022)
    @test cancelminus(interval(0x1P-1022,0x1.0000000000002P-1022), interval(0x1P-1022,0x1.0000000000001P-1022)) == interval(0.0,0x0.0000000000001P-1022)
    @test cancelminus(interval(0x1P-1022,0x1.0000000000001P-1022), interval(0x1P-1022,0x1.0000000000002P-1022)) == entireinterval()
    @test cancelminus(interval(-0x1P+0,0x1.FFFFFFFFFFFFFP-53), interval(-0x1.FFFFFFFFFFFFEP-53,0x1P+0)) == interval(-0x1.FFFFFFFFFFFFFP-1,-0x1.FFFFFFFFFFFFEP-1)
    @test cancelminus(interval(-0x1P+0,0x1.FFFFFFFFFFFFEP-53), interval(-0x1.FFFFFFFFFFFFFP-53,0x1P+0)) == entireinterval()
end


# day 04 ---------------------------------------------
using DelimitedFiles

input_path = joinpath(@__DIR__, "inputs", "input_04.txt")

draws = parse.(Int, split(readline(input_path), ","))
boards =
   reshape(
      readdlm(input_path, skipstart = 1, Int) |> permutedims,
      (5, 5, 100)
   )

# functions -----------------------------------------
function marknumber(checkboard, board, number::Int)
   for i = 1:5, j = 1:5
      if checkboard[i,j] ≠ 1
         checkboard[i,j] = board[i,j] == number
      end
   end
   return checkboard
end

function checkbingo(checkboard)
   h_sum = count(checkboard, dims=1)
   v_sum = count(checkboard, dims=2)
   any(x -> x == 5, h_sum) || any(x -> x == 5, v_sum)
end

getscore(checkboard, board, draw) = sum(boards[.!checkboard]) * draw

# 04a ------------------------------------------------
function sol_04a(boards, draws)
   checkboards = falses(size(boards))
   for draw in draws
      for i = 1:size(checkboards, 3)
         checkboards[:,:,i] = marknumber(checkboards[:,:,i], boards[:,:,i], draw)
         checkbingo(checkboards[:,:,i]) && return getscore(checkboards[:,:,i], boards[:,:,i], draw)
      end
   end
end

sol_04a(boards, draws)

# 04b ------------------------------------------------
function sol_04b(boards, draws)

   checkboards = falses(size(boards))
   nboards = size(checkboards, 3)
   bingocheck = falses(nboards)

   for nb in draws
      for i = 1:nboards
         if bingocheck[i] == false
            checkboards[:,:,i] = marknumber(checkboards[:,:,i], boards[:,:,i], nb)
            if checkbingo(checkboards[:,:,i])
               bingocheck[i] = true
               all(bingocheck) && return getscore(checkboards, boards, i, nb)
            end
         end
      end
   end
end

sol_04b(boards, draws)