"""
    fit_echo_linear(
        phas::AbstractArray{<:AbstractFloat, 4},
        W::AbstractArray{<:AbstractFloat, 4},
        TEs::NTuple{N, Real};
        phase_offset::Bool = true
    ) -> NTuple{1 or 2, typeof(similar(phas)){3}}

Weighted least squares for multi-echo data.

### Arguments
- `phas::AbstractArray{<:AbstractFloat, 4}`: unwrapped multi-echo phase
- `W::AbstractArray{<:AbstractFloat, 4}`: reciprocal of error variance of voxel
- `TEs::NTuple{N, Real}`: echo times

### Keywords
- `phase_offset::Bool = true`: model phase offset (`true`)

### Returns
- `typeof(similar(phas)){3}`: weighted least-squares estimate for phase
- [`typeof(similar(phas)){3}`]: weighted least-squares estimate for phase offset
    if `phase_offset = true`
"""
function fit_echo_linear(
    phas::AbstractArray{<:AbstractFloat, 4},
    W::AbstractArray{<:AbstractFloat, 4},
    TEs::NTuple{N, Real};
    phase_offset::Bool = true
) where {N}
    p = tzero(phas, size(phas)[1:3])
    if !phase_offset
        return fit_echo_linear!(p, phas, W, TEs)
    else
        return fit_echo_linear!(p, tzero(p), phas, W, TEs)
    end
end

"""
    fit_echo_linear!(
        p::AbstractArray{<:AbstractFloat, 3},
        phas::AbstractArray{<:AbstractFloat, 4},
        W::AbstractArray{<:AbstractFloat, 4},
        TEs::NTuple{N, Real}
    ) -> p

Weighted least squares for multi-echo data (phase offset = 0).

### Arguments
- `p::AbstractArray{<:AbstractFloat, 3}`: weighted least-squares estimate for phase
- `phas::AbstractArray{<:AbstractFloat, 4}`: unwrapped multi-echo phase
- `W::AbstractArray{<:AbstractFloat, 4}`: reciprocal of error variance of voxel
- `TEs::NTuple{N, Real}`: echo times

### Returns
- `p`: weighted least-squares estimate for phase
"""
function fit_echo_linear!(
    p::AbstractArray{Tp, 3},
    phas::AbstractArray{Tphas, 4},
    W::AbstractArray{TW, 4},
    TEs::NTuple{NT, Real}
) where {Tp<:AbstractFloat, Tphas<:AbstractFloat, TW<:Real, NT}
    nx, ny, nz, nt = size(phas)

    nt == NT || throw(DimensionMismatch())
    size(p) == (nx, ny, nz) || throw(DimensionMismatch())
    size(W) == size(phas)   || throw(DimensionMismatch())

    T = foldl(promote_type, (Tphas, TW), init=Tp)
    tes = convert.(T, TEs)

    _zeroT = zero(T)
    _zeroTp = zero(Tp)

    @threads for k in 1:nz
        @inbounds for j in 1:ny
            for i in 1:nx
                num = _zeroT
                den = _zeroT
                for t in Base.OneTo(NT)
                    w = W[i,j,k,t] * W[i,j,k,t] * tes[t]
                    num = muladd(w, phas[i,j,k,t], num)
                    den = muladd(w, tes[t], den)
                end
                p[i,j,k] = iszero(den) ? _zeroTp : num * inv(den)
            end
        end
    end

    return p
end

"""
    fit_echo_linear!(
        p::AbstractArray{<:AbstractFloat, 3},
        p0::AbstractArray{<:AbstractFloat, 3},
        phas::AbstractArray{<:AbstractFloat, 4},
        W::AbstractArray{<:AbstractFloat, 4},
        TEs::NTuple{N, Real}
    ) -> (p, p0)

Weighted least squares for multi-echo data (estimate phase offset).

### Arguments
- `p::AbstractArray{<:AbstractFloat, 3}`: weighted least-squares estimate for phase
- `p0::AbstractArray{<:AbstractFloat, 3}`: weighted least-squares estimate for phase offset
- `phas::AbstractArray{<:AbstractFloat, 4}`: unwrapped multi-echo phase
- `W::AbstractArray{<:AbstractFloat, 4}`: reciprocal of error variance of voxel
- `TEs::NTuple{N, Real}`: echo times

### Returns
- `p`: weighted least-squares estimate for phase
- `p0`: weighted least-squares estimate for phase offset
"""
function fit_echo_linear!(
    p::AbstractArray{Tp, 3},
    p0::AbstractArray{Tp0, 3},
    phas::AbstractArray{Tphas, 4},
    W::AbstractArray{TW, 4},
    TEs::NTuple{NT, Real}
) where {Tp<:AbstractFloat, Tp0<:AbstractFloat, Tphas<:AbstractFloat, TW<:Real, NT}
    nx, ny, nz, nt = size(phas)

    nt == NT || throw(DimensionMismatch())
    size(p)  == (nx, ny, nz) || throw(DimensionMismatch())
    size(p0) == (nx, ny, nz) || throw(DimensionMismatch())
    size(W)  == size(phas)   || throw(DimensionMismatch())

    T = foldl(promote_type, (Tp0, Tphas, TW), init=Tp)
    tes = convert.(T, TEs)

    _zeroT = zero(T)
    _zeroTp = zero(Tp)

    @threads for k in 1:nz
        @inbounds for j in 1:ny
            for i in 1:nx
                x = _zeroT
                y = _zeroT
                w = _zeroT
                for t in Base.OneTo(NT)
                    w += W[i,j,k,t]
                    x = muladd(W[i,j,k,t], tes[t], x)
                    y = muladd(W[i,j,k,t], phas[i,j,k,t], y)
                end

                w = inv(w)
                x *= w
                y *= w

                num = _zeroT
                den = _zeroT
                for t in Base.OneTo(NT)
                    xx = tes[t] - x
                    yy = phas[i,j,k,t] - y
                    ww = W[i,j,k,t] * xx
                    num = muladd(ww, yy, num)
                    den = muladd(ww, xx, den)
                end

                p[i,j,k] = iszero(den) ? _zeroTp : num * inv(den)
                p0[i,j,k] = y - p[i,j,k] * x
            end
        end
    end

    return p, p0
end


module Bisect_Tests

using Compat
using Compat.Test
using IntervalArithmetic
using EAGOBranchBound

S = BnBSolver()
B = BnBModel([Interval(1.0,2.0),Interval(3.0,4.0),Interval(3.0,5.0),Interval(3.0,6.0)])
N = [Interval(1.0,2.0),Interval(3.0,4.0),Interval(4.0,5.0),Interval(4.0,6.0)]
nx = 2

X1,X2 = EAGOBranchBound.Bisect_Rel(S,B,N)
Y1,Y2 = EAGOBranchBound.Bisect_Abs_Imp(S,B,N,nx)
Z1,Z2 = EAGOBranchBound.Bisect_Rel_Imp(S,B,N,nx)
@test X1[1] == Interval(1.0,1.5)
@test X2[1] == Interval(1.5,2.0)
@test X2[2:4] == [Interval(3.0,4.0),Interval(4.0,5.0),Interval(4.0,6.0)]
@test Y1[1:3] == [Interval(1.0,2.0),Interval(3.0,4.0),Interval(4.0,5.0)]
@test Y1[4] == Interval(4.0,5.0)
@test Y2[4] == Interval(5.0,6.0)
@test Z1[1:3] == [Interval(1.0,2.0),Interval(3.0,4.0),Interval(4.0,5.0)]
@test Z1[4] == Interval(4.0,5.0)
@test Z2[4] == Interval(5.0,6.0)

end


module TestDegenerate
  using Test
  using EcologicalNetworks
  using LinearAlgebra

  # Generate some data

  A = UnipartiteProbabilisticNetwork([1.0 0.0 0.0; 0.0 0.1 0.0; 0.0 0.0 0.1])
  @test isdegenerate(A)

  B = BipartiteProbabilisticNetwork([1.0 1.0 1.0 0.1; 0.0 1.0 0.1 0.0; 0.5 1.0 0.0 0.0])
  @test !isdegenerate(B)

  C = UnipartiteNetwork(zeros(Bool, (4,4)))
  @test isdegenerate(C)

  D = UnipartiteNetwork(Matrix(I, (4,4)))
  @test isdegenerate(D)

  K = simplify(UnipartiteNetwork([false false false; false true true; false true true]))
  @test !isdegenerate(K)

  N = BipartiteNetwork([false false false; true false true; true true false])
  M = simplify(N)
  @test isdegenerate(N)
  @test !isdegenerate(M)

end


using eFEMpart, JLD

using LinearAlgebra # for condition number calculation

function main()
  function computeNorms(N,order)
    start = time()
    # load mesh
    mesh = squareMesh([-2,2,0,1],N,order)

    OperatorType = :AdvDiffAS

    Pe = 2.3
    windX(x,y) = -x^4*y^2 + 5.9; windY(x,y) = x^3*y^3
    wx = [windX.(mesh.xy[i].x,mesh.xy[i].y) for i=1:length(mesh.xy)]
    wy = [windY.(mesh.xy[i].x,mesh.xy[i].y) for i=1:length(mesh.xy)]
    param  = AdvDiffParam(wx,wy,Pe*ones(length(mesh.xy)))

    dNodes = Dirichlet(:left,:right,:top,:bottom)
    dBCf = Dirichlet((x,y) -> x^4*y^4)
    ff   = Forcing((x,y) -> 4*5.9*x^3*y^4 - 4/Pe*x^2*y^2*(4*x^2 + 3*y^2))
    Nodes = [dNodes]
    bcfun = [dBCf,ff]

    prob = Problem(mesh,Nodes,bcfun,OperatorType)
    sol = solve(prob,mesh,param)

    elapsed = time() - start

    # compute condition number of operator matrix
    #LinOp = GenerateSystem(mesh,prob,param)
    #ApplyBC!(LinOp,mesh,prob,param,OperatorType)
    #κ = cond(Array(LinOp.Op),2)
    κ = 0.5

    # compute mesh h
    h = hCalc(mesh)

    # generate exact solution
    Uexact(x,y) = dBCf.f(x,y)
    UexactArr = [Uexact.(mesh.xy[i].x,mesh.xy[i].y) for i=1:length(mesh.xy)]

    # compute difference
    err = sol.u - UexactArr

    # compute Li norm
    L1   = DomainNorm(mesh.xy,mesh.cm,err;normID="1")
    L2   = DomainNorm(mesh.xy,mesh.cm,err;normID="2")
    Linf = DomainNorm(mesh.xy,mesh.cm,err;normID="Inf")

    return κ,h,L1,L2,Linf,elapsed
  end

  Narr    = [8,8,16,32]#,64]#,128]#,256]
  order   = [1,2]
  N       = length(Narr)
  κarr    = zeros(N,2)
  harr    = zeros(N,2)
  L1arr   = zeros(N,2)
  L2arr   = zeros(N,2)
  LInfarr = zeros(N,2)
  timearr = zeros(N,2)

  for i=1:N
    for j=1:2
      n = Narr[i]
      o = order[j]
      κarr[i,j],harr[i,j],L1arr[i,j],L2arr[i,j],LInfarr[i,j],timearr[i,j] = computeNorms(n,o)
      println("completed N=$(n), order=$(o)")
    end
  end

  κarr = κarr[2:end,:]
  harr = harr[2:end,:]
  L1arr = L1arr[2:end,:]
  L2arr = L2arr[2:end,:]
  LInfarr = LInfarr[2:end,:]
  timearr = timearr[2:end,:]

  # export output to *.jld file
  run(`mkdir -p data`)

  save("data/TEMP_advdiffAS_validation.jld", "harr",harr[:,1],
                                          "o1L1arr", L1arr[:,1],
                                          "o1L2arr", L2arr[:,1],
                                          "o1LInfarr",LInfarr[:,1],
                                          "o1timearr",timearr[:,1],
                                          "o2L1arr", L1arr[:,2],
                                          "o2L2arr", L2arr[:,2],
                                          "o2LInfarr",LInfarr[:,2],
                                          "o2timearr",timearr[:,2],
                                          "κarr",κarr[:,2])

  nothing
end

main()


module baffled_piston_examples
using FinEtools
using FinEtoolsAcoustics
using LinearAlgebra
using UnicodePlots

function baffled_piston_H8_transient()
	rho = 1.21*phun("kg/m^3");# mass density
	c  = 343.0*phun("m/s");# sound speed
	bulk =  c^2*rho;
	frequency=2000; # Hz
	omega=2*pi*frequency;
	P_piston = 1.0e-3*phun("MPa"); # amplitude of the piston pressure
	R = 50.0*phun("mm");# radius of the piston
	Ro = 400.0*phun("mm"); # radius of the enclosure
	nref = 2;#number of refinements of the sphere around the piston
	nlayers = 48;                     # number of layers of elements surrounding the piston
	tolerance = R/(2^nref)/10
	dt=1.0/frequency/20;
	tfinal=50*dt;
	nsteps=round(tfinal/dt)+1;

	println("""

		Baffled piston in a half-sphere domain with hard boundary condition.

		Hexahedral mesh. Algorithm version.
		""")

	t0  =  time()

	# Hexahedral mesh
	fens,fes  =  H8sphere(R,nref);
	bfes  =  meshboundary(fes)
	# File  =   "baffledabc_boundary.vtk"
	# vtkexportmesh(File, bfes.conn, fens.xyz, FinEtools.MeshExportModule.Q4)
	#  @async run(`"paraview.exe" $File`)

	l = selectelem(fens,bfes,facing = true,direction = [1.0 1.0  1.0])
	ex(xyz, layer) = (R+layer/nlayers*(Ro-R))*xyz/norm(xyz)
	fens1,fes1  =  H8extrudeQ4(fens, subset(bfes,l), nlayers, ex);
	fens,newfes1,fes2 =  mergemeshes(fens1, fes1, fens, fes, tolerance)
	fes = cat(newfes1,fes2)

	# Piston surface mesh
	bfes  =  meshboundary(fes)
	l1 = selectelem(fens, bfes, facing = true, direction = [-1.0 0.0 0.0])
	l2 = selectelem(fens, bfes, distance = R, from = [0.0 0.0 0.0], inflate = tolerance)
	piston_fes = subset(bfes,intersect(l1,l2));
	# File  =   "piston_fes.vtk"
	# vtkexportmesh(File, piston_fes.conn, fens.xyz, FinEtools.MeshExportModule.Q4)
	#  @async run(`"paraview.exe" $File`)

	# Outer spherical boundary
	louter = selectelem(fens, bfes, facing = true, direction = [1.0 1.0  1.0], dotmin= 0.001)
	outer_fes = subset(bfes,louter);
	# File  =   "outer_fes.vtk"
	# vtkexportmesh(File, outer_fes.conn, fens.xyz, FinEtools.MeshExportModule.Q4)
	#  @async run(`"paraview.exe" $File`)

	println("Pre-processing time elapsed  =  ",time() - t0,"s")

	t1  =  time()

	material = MatAcoustFluid(bulk, rho)
	# Region of the fluid
	femm = FEMMAcoust(IntegDomain(fes, GaussRule(3, 2)), material)

	geom = NodalField(fens.xyz);
	P = NodalField(zeros(nnodes(geom),1));
	piston_fenids = connectednodes(piston_fes);
	setebc!(P, piston_fenids, true, 1, 0.0);
	applyebc!(P);
	numberdofs!(P);

	S  =  acousticstiffness(femm, geom, P);
	C  =  acousticmass(femm, geom, P);
	D= (2.0/dt)*S +(dt/2.0)*C;

	P0 = deepcopy(P)
	Pdd0 = deepcopy(P)
	P1 = deepcopy(P)
	Pdd1 = deepcopy(P)
	TMPF = deepcopy(P)
	P = nothing # we don't need this field anymore
	vP0 = zeros(P0.nfreedofs)
	vP1 = deepcopy(vP0)
	vPd0 = deepcopy(vP0)
	vPd1 = deepcopy(vP0)

	t = 0.0
	P0.fixed_values[piston_fenids,1] .= P_piston*sin(omega*t)
	Pdd0.fixed_values[piston_fenids,1] .= P_piston*(-omega^2)*sin(omega*t)
	vP0 = gathersysvec!(P0, vP0)

	nh = selectnode(fens, nearestto = [R+Ro/2, 0.0, 0.0] )
	Pnh = [P1.values[nh, 1][1]]
	ts = [t]
	step =0;
	while t <=tfinal
		step = step  +1;
		t=t+dt;
		P1.fixed_values[piston_fenids,1] .= P_piston*sin(omega*t)
		Pdd1.fixed_values[piston_fenids,1] .= P_piston*(-omega^2)*sin(omega*t)
		TMPF.fixed_values = P0.fixed_values + P1.fixed_values
		F = nzebcloadsacousticmass(femm, geom, TMPF);
		TMPF.fixed_values = Pdd0.fixed_values + Pdd1.fixed_values
		F = F + nzebcloadsacousticstiffness(femm, geom, TMPF);
		# println("$(norm(F))")
		vPd1 = D\((2/dt)*(S*vPd0) - C*(2*vP0+(dt/2)*vPd0) + F);
		vP1 = vP0 + (dt/2)*(vPd0+vPd1);
		scattersysvec!(P1, vP1); # store current pressure
		# Swap variables for the next step
		copyto!(vP0, vP1)
		copyto!(vPd0, vPd1)
		copyto!(P0, P1)
		copyto!(Pdd0, Pdd1)
		# Graphics output
		File  =   "baffled_piston-$(step).vtk"
		vtkexportmesh(File, fes.conn, fens.xyz, FinEtools.MeshExportModule.H8;
			scalars = [(  "P", P1.values+P1.fixed_values)])
		#  @async run(`"paraview.exe" $File`)
		push!(Pnh, P1.values[nh, 1][1])
		push!(ts, t)
		println("step = $( step )")
	end

	# @pgf a = Axis({
	#     xlabel = "Time",
	#     ylabel = "Pnh",
	#     title = "Transient pressure"
	# },
	# Plot(Table([:x => vec(ts), :y => vec(Pnh)])))
	# display(a)
	plt = lineplot(vec(ts), vec(Pnh), canvas = DotCanvas, title = "Transient pressure", name = "P", xlabel = "Time", ylabel = "Pnh")
	display(plt)

	# using Plots
	# plotly()
	# plot(vec(ts), vec(Pnh))
	# gui()
	println("$(Pnh[end])")

	# # Surface for the ABC
	# abc1  =  FDataDict("femm"=>FEMMAcoustSurf(IntegDomain(outer_fes, GaussRule(2, 2)),
	#           material))
	#
	# # Surface of the piston
	# flux1  =  FDataDict("femm"=>FEMMAcoustSurf(IntegDomain(piston_fes, GaussRule(2, 2)),
	#           material),  "normal_flux"=> -rho*a_piston+0.0im);
	#
	# # Make model data
	# modeldata =  FDataDict("fens"=>  fens,
	#                  "omega"=>omega,
	#                  "regions"=>[region1],
	#                  "flux_bcs"=>[flux1], "ABCs"=>[abc1])
	#
	# # Call the solver
	# modeldata = FinEtools.AlgoAcoustModule.steadystate(modeldata)
	#
	# println("Computing time elapsed  =  ",time() - t1,"s")
	# println("Total time elapsed  =  ",time() - t0,"s")
	#
	# geom = modeldata["geom"]
	# P = modeldata["P"]
	#
	# File  =   "baffledabc.vtk"
	# vtkexportmesh(File, fes.conn, geom.values, FinEtools.MeshExportModule.H8;
	# scalars = [("absP", abs.(P.values))])
	# @async run(`"paraview.exe" $File`)
	#
	# # using Winston
	# # pl  =  FramedPlot(title = "Matrix",xlabel = "x",ylabel = "Re P, Im P")
	# # setattr(pl.frame, draw_grid = true)
	# # add(pl, Curve([1:length(C[:])],vec(C[:]), color = "blue"))
	#
	# # # pl = plot(geom.values[nLx,1][ix],scalars[nLx][ix])
	# # # xlabel("x")
	# # # ylabel("Pressure")
	# # display(pl)
	#
	# true
end # baffled_piston_H8_transient

function allrun()
    println("#####################################################")
    println("# baffled_piston_H8_transient ")
    baffled_piston_H8_transient()
    return true
end # function allrun

end # module baffled_piston_examples


##############################################################################
#
# TD-IDF
#
##############################################################################

tf_idf{T <: Real}(dtm::Matrix{T}) = tf_idf!(dtm, Array(Float64, size(dtm)...))

tf_idf{T <: Real}(dtm::SparseMatrixCSC{T}) =  tf_idf!(dtm, sparse(Int[], Int[], 0.0, size(dtm)...))

tf_idf!{T <: Real}(dtm::AbstractMatrix{T}) = tf_idf!(dtm, dtm)

tf_idf!{T <: Real}(dtm::SparseMatrixCSC{T}) = tf_idf!(dtm, dtm)

tf_idf(dtm::DocumentTermMatrix) = tf_idf(dtm.dtm)

# This does not make sense, since DocumentTermMatrix is based on an array of integers
#tf_idf!(dtm::DocumentTermMatrix) = tf_idf!(dtm.dtm)

# The second Matrix will be overwritten with the result
# Will work correctly if dtm and tfidf are the same matrix
function tf_idf!{T1 <: Real, T2 <: AbstractFloat}(dtm::AbstractMatrix{T1}, tfidf::AbstractMatrix{T2})
    n, p = size(dtm)

    # TF tells us what proportion of a document is defined by a term
    for i in 1:n
        words_in_document = 0
        for j in 1:p
            words_in_document += dtm[i, j]
        end
        tfidf[i, :] = dtm[i, :] ./ words_in_document
    end

    # IDF tells us how rare a term is in the corpus
    documents_containing_term = vec(sum(dtm .> 0, 1))
    idf = log(n ./ documents_containing_term)

    # TF-IDF is the product of TF and IDF
    for i in 1:n
        for j in 1:p
           tfidf[i, j] = tfidf[i, j] * idf[j]
        end
    end

    return tfidf
end


using Test, AutoARIMA

@testset "Forecast" begin

    @test forecast(MAModel{1}(0.0, [-0.7], 1.0), [3.0,8,2,5,6]) ≈ 1.01941

    @test forecast(ARModel{2}(1.0, [0.5, 0.5], 1.0), [3.0,2]) ≈ 3.5

    @test forecast(ARModel{2}(10.0, [0.5, 0.5], 1.0), [3.0,2]) ≈ 12.5

end


# linear softmax classifier (with stochastic gradient descent)
type LinearClassifier
    k::Int64 # number of outputs
    n::Int64 # number of inputs
    weights::Array{Float64, 2} # n * k weight matrix

    outputs :: Vector{Float64}
end

function LinearClassifier(k, n)
    weights = rand(n, k) * 2 - 1; # range is [-1, 1]
    LinearClassifier(k, n, weights, zeros(k))
end

function predict(c::LinearClassifier, x::Array{Float64})
    # the softmax() function from the NumericExtension package is more numeric stable
    return vec(softmax(x * c.weights))
end

function predict!(c::LinearClassifier, x::Array{Float64})
    # c.outputs = vec(softmax(x * c.weights))
    s = 0.0
    for i in 1:c.k
        s = 0.0
        for j in 1:c.n
            s += x[j] * c.weights[j, i]
        end
        c.outputs[i] = s
    end

    softmax!(c.outputs, c.outputs);
end

function train_one(c::LinearClassifier, x::Array{Float64}, y::Int64, α::Float64=0.025)
    # if !in(y, 1 : c.k)
    #     msg = @sprintf "A sample is discarded because the label y = %d is not in range of 1 to %d" y c.k
    #     warn(msg)
    #     return
    # end

    predict!(c, x)
    c.outputs[y] -= 1

    # c.weights -= α * x' * outputs;
    # BLAS.ger!(-α, vec(x), c.outputs, c.weights)
    m = 0.0
    j = 0
    limit = c.n - 4
    for i in 1:c.k
        m = α * c.outputs[i]
        j = 1
        while j <= limit
            @nexprs 4 (idx->c.weights[j + idx - 1, i] -= m * x[j + idx - 1])
            j+=4
        end
        while j <= c.n
            c.weights[j, i] -= m * x[j]
            j+=1
        end
    end
end

function train_one(c::LinearClassifier, x::Array{Float64}, y::Int64, input_gradient::Array{Float64}, α::Float64=0.025)
    predict!(c, x)
    c.outputs[y] -= 1

    # input_gradient = ( c.weights * outputs' )'
    # BLAS.gemv!('N', α, c.weights, c.outputs, 1.0, input_gradient)
    m = 0.0
    j = 0
    limit = c.n - 4
    for i in 1:c.k
        m = α * c.outputs[i]
        j = 1
        while j <= limit
            @nexprs 4 (idx->input_gradient[j+idx-1] += m * c.weights[j+idx-1, i])
            j+=4
        end
        while j <= c.n
            input_gradient[j] += m * c.weights[j, i]
            j+=1
        end
    end

    # c.weights -= α * x' * outputs;
    # BLAS.ger!(-α, vec(x), c.outputs, c.weights)
    for i in 1:c.k
        m = α * c.outputs[i]
        j = 1
        while j <= limit
            @nexprs 4 (idx->c.weights[j + idx - 1, i] -= m * x[j + idx - 1])
            j+=4
        end
        while j <= c.n
            c.weights[j, i] -= m * x[j]
            j+=1
        end
    end
end

# calculate the overall log likelihood. Mainly used for debugging
function log_likelihood(c, X, y)
    n = size(X, 1)
    l = 0
    for i in 1:n
        l += log(predict(c, X[i, :])[y[i]])
    end
    return l
end

# calculate the accuracy on the testing dataset
function accuracy(c::LinearClassifier, X::Array{Float64}, y::Array{Int64})
    n = size(X, 1)
    succ = 0
    for i in 1 : n
        output = predict(c, X[i, :])
        if maximum(output) == output[y[i]]
            succ += 1
        end
    end
    return succ / n
end


type GaussQuadratureRule
    weights::Vector{Float64}
    points::Vector{Vector{Float64}}
end

weights(qr::GaussQuadratureRule) = qr.weights
points(qr::GaussQuadratureRule) = qr.points

function integrate(qr::GaussQuadratureRule, f::Function)
    I = 0.0
    for (w, x) in zip(weights(qr), points(qr))
        I += w * f(x)
    end
    return I
end

# Maybe we could stage the coming three function and create ndim loops.
function make_quadrule(ref::ReferenceLine, order::Int)
    p, w = gausslegendre(order)
    weights = Array(Float64, order)
    points = Array(Vector{Float64}, order)
    count = 1
    for i = 1:order
        points[count] = [p[i]]
        weights[count] = w[i]
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceSquare, order::Int)
    p, w = gausslegendre(order)
    weights = Array(Float64, order^2)
    points = Array(Vector{Float64}, order^2)
    count = 1
    for i = 1:order, j = 1:order
        points[count] = [p[i], p[j]]
        weights[count] = w[i] * w[j]
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceCube, order::Int)
    p, w = gausslegendre(order)
    weights = Array(Float64, order^3)
    points = Array(Vector{Float64}, order^3)
    count = 1
    for i = 1:order, j = 1:order, k = 1:order
        points[count] = [p[i], p[j], p[k]]
        weights[count] = w[i] * w[j] * w[k]
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceTriangle, order::Int)
    p1, w1 = gaussjacobi(order, 0.0, 0.0)
    p2, w2 = gaussjacobi(order, 1.0, 0.0)

    affine(a, b) = [(1+a) * (1-b) / 4, (1+b) / 2]

    SCALE = 1/8
    weights = Array(Float64, order^2)
    points = Array(Vector{Float64}, order^2)
    count = 1
    for i = 1:order, j = 1:order
        weights[count] = w1[i] * w2[j] * SCALE
        points[count] = affine(p1[i], p2[j])
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceTetrahedron, order::Int)
    p1, w1 = gaussjacobi(order, 0.0, 0.0)
    p2, w2 = gaussjacobi(order, 1.0, 0.0)
    p3, w3 = gaussjacobi(order, 2.0, 0.0)
    error("Bugged!")
    function affine(a, b, c)
       a = (a+1) / 2
       b = (b+1) / 2
       c = (c+1) / 2
       [a*b*c, a*b*(1-c), a*(1-b)]
     end

    SCALE = 1/48
    weights = Array(Float64, order^3)
    points = Array(Vector{Float64}, order^3)
    count = 1
    for i = 1:order, j = 1:order, k = 1:order
        weights[count] = w1[i] * w2[j] * w2[k] * SCALE
        points[count] = affine(p1[i], p2[j], p3[k])
        count += 1
    end
    GaussQuadratureRule(weights, points)
end


# Testing gammatonegram response
s = read(open(joinpath(dirname(@__FILE__), "data", "sa2.raw")), Float64, 42701)
fs = 16000
(W, _) = AuditoryFilters.fft2gammatonemx(512, fs, 23, 1, 150, 4000)
G = gammatonegram(s, fs, 0.025, 0.010, 23, 150, 4000, 1)

G_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "G_matlab.csv")))
F_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "F_matlab.csv")))
W_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "W_matlab.csv")))

@test_approx_eq W W_matlab
@test_approx_eq G.frequencies F_matlab'
@test_approx_eq G.amplitude G_matlab


function igyp01()
#
#  generalization of Igarash and Ypma
#
    p = poly([2.35*[1;1;1];2.56])
    z = [2.35 3;2.56 1];
    p, PolyZeros(z)
end


function forsub(A::AbstractArray{T,2}, b::AbstractVector{T}) where {T<:Number}
#
# assume A is an lower triangular matrix
# forsub performs forward substitution to solve Ax=b
#
    n = size(A,1)
    LowerTriangular(A) \ b[1:n]
end


import Base: position
using FFTW

"""
    PositionBasis(xmin, xmax, Npoints)
    PositionBasis(b::MomentumBasis)

Basis for a particle in real space.

For simplicity periodic boundaries are assumed which means that
the rightmost point defined by `xmax` is not included in the basis
but is defined to be the same as `xmin`.

When a [`MomentumBasis`](@ref) is given as argument the exact values
of ``x_{min}`` and ``x_{max}`` are due to the periodic boundary conditions
more or less arbitrary and are chosen to be
``-\\pi/dp`` and ``\\pi/dp`` with ``dp=(p_{max}-p_{min})/N``.
"""
struct PositionBasis{T,X1,X2} <: Basis
    shape::Vector{T}
    xmin::Float64
    xmax::Float64
    N::T
    function PositionBasis{X1,X2}(xmin::Real, xmax::Real, N::T) where {X1,X2,T<:Int}
        @assert isa(X1, Real) && isa(X2, Real)
        new{T,X1,X2}([N], xmin, xmax, N)
    end
end
PositionBasis(xmin::Real, xmax::Real, N::Int) = PositionBasis{xmin,xmax}(xmin,xmax,N)

"""
    MomentumBasis(pmin, pmax, Npoints)
    MomentumBasis(b::PositionBasis)

Basis for a particle in momentum space.

For simplicity periodic boundaries are assumed which means that
`pmax` is not included in the basis but is defined to be the same as `pmin`.

When a [`PositionBasis`](@ref) is given as argument the exact values
of ``p_{min}`` and ``p_{max}`` are due to the periodic boundary conditions
more or less arbitrary and are chosen to be
``-\\pi/dx`` and ``\\pi/dx`` with ``dx=(x_{max}-x_{min})/N``.
"""
struct MomentumBasis{P1,P2} <: Basis
    shape::Vector{Int}
    pmin::Float64
    pmax::Float64
    N::Int
    function MomentumBasis{P1,P2}(pmin::Real, pmax::Real, N::Int) where {P1,P2}
        @assert isa(P1, Real) && isa(P2, Real)
        new([N], pmin, pmax, N)
    end
end
MomentumBasis(pmin::Real, pmax::Real, N::Int) = MomentumBasis{pmin,pmax}(pmin, pmax, N)

PositionBasis(b::MomentumBasis) = (dp = (b.pmax - b.pmin)/b.N; PositionBasis(-pi/dp, pi/dp, b.N))
MomentumBasis(b::PositionBasis) = (dx = (b.xmax - b.xmin)/b.N; MomentumBasis(-pi/dx, pi/dx, b.N))

==(b1::PositionBasis, b2::PositionBasis) = b1.xmin==b2.xmin && b1.xmax==b2.xmax && b1.N==b2.N
==(b1::MomentumBasis, b2::MomentumBasis) = b1.pmin==b2.pmin && b1.pmax==b2.pmax && b1.N==b2.N

"""
    gaussianstate(b::PositionBasis, x0, p0, sigma)
    gaussianstate(b::MomentumBasis, x0, p0, sigma)

Create a Gaussian state around `x0` and` p0` with width `sigma`.

In real space the gaussian state is defined as

```math
\\Psi(x) = \\frac{1}{\\pi^{1/4}\\sqrt{\\sigma}}
            e^{i p_0 (x-\\frac{x_0}{2}) - \\frac{(x-x_0)^2}{2 \\sigma^2}}
```

and is connected to the momentum space definition

```math
\\Psi(p) = \\frac{\\sqrt{\\sigma}}{\\pi^{1/4}}
            e^{-i x_0 (p-\\frac{p_0}{2}) - \\frac{1}{2}(p-p_0)^2 \\sigma^2}
```

via a Fourier-transformation

```math
\\Psi(p) = \\frac{1}{\\sqrt{2\\pi}}
            \\int_{-\\infty}^{\\infty} e^{-ipx}\\Psi(x) \\mathrm{d}x
```

The state has the properties

* ``⟨p⟩ = p_0``
* ``⟨x⟩ = x_0``
* ``\\mathrm{Var}(x) = \\frac{σ^2}{2}``
* ``\\mathrm{Var}(p) = \\frac{1}{2 σ^2}``

Due to the numerically necessary discretization additional scaling
factors ``\\sqrt{Δx}`` and ``\\sqrt{Δp}`` are used so that
``\\langle x_i|Ψ\\rangle = \\sqrt{Δ x} Ψ(x_i)`` and ``\\langle p_i|Ψ\\rangle = \\sqrt{Δ p} Ψ(p_i)`` so
that the resulting Ket state is normalized.
"""
function gaussianstate(b::PositionBasis, x0::Real, p0::Real, sigma::Real)
    psi = Ket(b)
    dx = spacing(b)
    alpha = 1.0/(pi^(1/4)*sqrt(sigma))*sqrt(dx)
    x = b.xmin
    for i=1:b.N
        psi.data[i] = alpha*exp(1im*p0*(x-x0/2) - (x-x0)^2/(2*sigma^2))
        x += dx
    end
    return psi
end

function gaussianstate(b::MomentumBasis, x0::Real, p0::Real, sigma::Real)
    psi = Ket(b)
    dp = spacing(b)
    alpha = sqrt(sigma)/pi^(1/4)*sqrt(dp)
    p = b.pmin
    for i=1:b.N
        psi.data[i] = alpha*exp(-1im*x0*(p-p0/2) - (p-p0)^2/2*sigma^2)
        p += dp
    end
    return psi
end

"""
    spacing(b::PositionBasis)

Difference between two adjacent points of the real space basis.
"""
spacing(b::PositionBasis) = (b.xmax - b.xmin)/b.N
"""
    spacing(b::MomentumBasis)

Momentum difference between two adjacent points of the momentum basis.
"""
spacing(b::MomentumBasis) = (b.pmax - b.pmin)/b.N

"""
    samplepoints(b::PositionBasis)

x values of the real space basis.
"""
samplepoints(b::PositionBasis) = (dx = spacing(b); Float64[b.xmin + i*dx for i=0:b.N-1])
"""
    samplepoints(b::MomentumBasis)

p values of the momentum basis.
"""
samplepoints(b::MomentumBasis) = (dp = spacing(b); Float64[b.pmin + i*dp for i=0:b.N-1])

"""
    position(b::PositionBasis)

Position operator in real space.
"""
position(b::PositionBasis) = SparseOperator(b, sparse(Diagonal(complex(samplepoints(b)))))

"""
    position(b:MomentumBasis)

Position operator in momentum space.
"""
function position(b::MomentumBasis)
    b_pos = PositionBasis(b)
    transform(b, b_pos)*dense(position(b_pos))*transform(b_pos, b)
end

"""
    momentum(b:MomentumBasis)

Momentum operator in momentum space.
"""
momentum(b::MomentumBasis) = SparseOperator(b, sparse(Diagonal(complex(samplepoints(b)))))

"""
    momentum(b::PositionBasis)

Momentum operator in real space.
"""
function momentum(b::PositionBasis)
    b_mom = MomentumBasis(b)
    transform(b, b_mom)*dense(momentum(b_mom))*transform(b_mom, b)
end

"""
    potentialoperator(b::PositionBasis, V(x))

Operator representing a potential ``V(x)`` in real space.
"""
function potentialoperator(b::PositionBasis, V::Function)
    x = samplepoints(b)
    diagonaloperator(b, V.(x))
end

"""
    potentialoperator(b::MomentumBasis, V(x))

Operator representing a potential ``V(x)`` in momentum space.
"""
function potentialoperator(b::MomentumBasis, V::Function)
    b_pos = PositionBasis(b)
    transform(b, b_pos)*dense(potentialoperator(b_pos, V))*transform(b_pos, b)
end

"""
    potentialoperator(b::CompositeBasis, V(x, y, z, ...))

Operator representing a potential ``V`` in more than one dimension.

# Arguments
* `b`: Composite basis consisting purely either of `PositionBasis` or
    `MomentumBasis`. Note, that calling this with a composite basis in
    momentum space might consume a large amount of memory.
* `V`: Function describing the potential. ATTENTION: The number of arguments
    accepted by `V` must match the spatial dimension. Furthermore, the order
    of the arguments has to match that of the order of the tensor product of
    bases (e.g. if `b=bx⊗by⊗bz`, then `V(x,y,z)`).
"""
function potentialoperator(b::CompositeBasis, V::Function)
    if isa(b.bases[1], PositionBasis)
        potentialoperator_position(b, V)
    elseif isa(b.bases[1], MomentumBasis)
        potentialoperator_momentum(b, V)
    else
        throw(IncompatibleBases())
    end
end
function potentialoperator_position(b::CompositeBasis, V::Function)
    for base=b.bases
        @assert isa(base, PositionBasis)
    end

    points = [samplepoints(b1) for b1=b.bases]
    dims = length.(points)
    n = length(b.bases)
    data = Array{ComplexF64}(undef, dims...)
    @inbounds for i=1:length(data)
        index = Tuple(CartesianIndices(data)[i])
        args = (points[j][index[j]] for j=1:n)
        data[i] = V(args...)
    end

    diagonaloperator(b, data[:])
end
function potentialoperator_momentum(b::CompositeBasis, V::Function)
    bases_pos = []
    for base=b.bases
        @assert isa(base, MomentumBasis)
        push!(bases_pos, PositionBasis(base))
    end
    b_pos = tensor(bases_pos...)
    transform(b, b_pos)*dense(potentialoperator_position(b_pos, V))*transform(b_pos, b)
end

"""
    FFTOperator

Abstract type for all implementations of FFT operators.
"""
abstract type FFTOperator{BL<:Basis, BR<:Basis, T} <: AbstractOperator{BL,BR} end

Base.eltype(x::FFTOperator) = promote_type(eltype(x.mul_before), eltype(x.mul_after))

"""
    FFTOperators

Operator performing a fast fourier transformation when multiplied with a state
that is a Ket or an Operator.
"""
mutable struct FFTOperators{BL<:Basis,BR<:Basis,T<:Array,P1,P2,P3,P4} <: FFTOperator{BL, BR, T}
    basis_l::BL
    basis_r::BR
    fft_l!::P1
    fft_r!::P2
    fft_l2!::P3
    fft_r2!::P4
    mul_before::T
    mul_after::T
    function FFTOperators(b1::BL, b2::BR,
        fft_l!::P1,
        fft_r!::P2,
        fft_l2!::P3,
        fft_r2!::P4,
        mul_before::T,
        mul_after::T) where {BL<:Basis,BR<:Basis,T,P1,P2,P3,P4}
        new{BL,BR,T,P1,P2,P3,P4}(b1, b2, fft_l!, fft_r!, fft_l2!, fft_r2!, mul_before, mul_after)
    end
end

"""
    FFTKets

Operator that can only perform fast fourier transformations on Kets.
This is much more memory efficient when only working with Kets.
"""
mutable struct FFTKets{BL<:Basis,BR<:Basis,T<:Array,P1,P2} <: FFTOperator{BL, BR, T}
    basis_l::BL
    basis_r::BR
    fft_l!::P1
    fft_r!::P2
    mul_before::T
    mul_after::T
    function FFTKets(b1::BL, b2::BR,
        fft_l!::P1,
        fft_r!::P2,
        mul_before::T,
        mul_after::T) where {BL<:Basis,BR<:Basis, T, P1, P2}
        new{BL, BR, T, P1, P2}(b1, b2, fft_l!, fft_r!, mul_before, mul_after)
    end
end

"""
    transform(b1::MomentumBasis, b2::PositionBasis)
    transform(b1::PositionBasis, b2::MomentumBasis)

Transformation operator between position basis and momentum basis.
"""
function transform(basis_l::MomentumBasis, basis_r::PositionBasis; ket_only::Bool=false)
    Lx = (basis_r.xmax - basis_r.xmin)
    dp = spacing(basis_l)
    dx = spacing(basis_r)
    if basis_l.N != basis_r.N || abs(2*pi/dp - Lx)/Lx > 1e-12
        throw(IncompatibleBases())
    end
    mul_before = exp.(-1im*basis_l.pmin*(samplepoints(basis_r) .- basis_r.xmin))
    mul_after = exp.(-1im*basis_r.xmin*samplepoints(basis_l))/sqrt(basis_r.N)
    x = Vector{ComplexF64}(undef, length(basis_r))
    if ket_only
        FFTKets(basis_l, basis_r, plan_bfft!(x), plan_fft!(x), mul_before, mul_after)
    else
        A = Matrix{ComplexF64}(undef, length(basis_r), length(basis_r))
        FFTOperators(basis_l, basis_r, plan_bfft!(x), plan_fft!(x), plan_bfft!(A, 2), plan_fft!(A, 1), mul_before, mul_after)
    end
end

"""
    transform(b1::CompositeBasis, b2::CompositeBasis)

Transformation operator between two composite bases. Each of the bases
has to contain bases of type PositionBasis and the other one a corresponding
MomentumBasis.
"""
function transform(basis_l::PositionBasis, basis_r::MomentumBasis; ket_only::Bool=false)
    Lx = (basis_l.xmax - basis_l.xmin)
    dp = spacing(basis_r)
    dx = spacing(basis_l)
    if basis_l.N != basis_r.N || abs(2*pi/dp - Lx)/Lx > 1e-12
        throw(IncompatibleBases())
    end
    mul_before = exp.(1im*basis_l.xmin*(samplepoints(basis_r) .- basis_r.pmin))
    mul_after = exp.(1im*basis_r.pmin*samplepoints(basis_l))/sqrt(basis_r.N)
    x = Vector{ComplexF64}(undef, length(basis_r))
    if ket_only
        FFTKets(basis_l, basis_r, plan_fft!(x), plan_bfft!(x), mul_before, mul_after)
    else
        A = Matrix{ComplexF64}(undef, length(basis_r), length(basis_r))
        FFTOperators(basis_l, basis_r, plan_fft!(x), plan_bfft!(x), plan_fft!(A, 2), plan_bfft!(A, 1), mul_before, mul_after)
    end
end

function transform(basis_l::CompositeBasis, basis_r::CompositeBasis; ket_only::Bool=false, index::Vector{Int}=Int[])
    @assert length(basis_l.bases) == length(basis_r.bases)
    if length(index) == 0
        check_pos = [isa.(basis_l.bases, PositionBasis)...]
        check_mom = [isa.(basis_l.bases, MomentumBasis)...]
        if any(check_pos) && !any(check_mom)
            index = [1:length(basis_l.bases);][check_pos]
        elseif any(check_mom) && !any(check_pos)
            index = [1:length(basis_l.bases);][check_mom]
        else
            throw(IncompatibleBases())
        end
    end
    if all(isa.(basis_l.bases[index], PositionBasis))
        @assert all(isa.(basis_r.bases[index], MomentumBasis))
        transform_xp(basis_l, basis_r, index; ket_only=ket_only)
    elseif all(isa.(basis_l.bases[index], MomentumBasis))
        @assert all(isa.(basis_r.bases[index], PositionBasis))
        transform_px(basis_l, basis_r, index; ket_only=ket_only)
    else
        throw(IncompatibleBases())
    end
end

function transform_xp(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Vector{Int}; ket_only::Bool=false)
    n = length(basis_l.bases)
    Lx = [(b.xmax - b.xmin) for b=basis_l.bases[index]]
    dp = [spacing(b) for b=basis_r.bases[index]]
    dx = [spacing(b) for b=basis_l.bases[index]]
    N = [length(b) for b=basis_l.bases]
    for i=1:n
        if N[i] != length(basis_r.bases[i])
            throw(IncompatibleBases())
        end
    end
    for i=1:length(index)
        if abs(2*pi/dp[i] - Lx[i])/Lx[i] > 1e-12
            throw(IncompatibleBases())
        end
    end

    if index[1] == 1
        mul_before = exp.(1im*basis_l.bases[1].xmin*(samplepoints(basis_r.bases[1]) .- basis_r.bases[1].pmin))
        mul_after = exp.(1im*basis_r.bases[1].pmin*samplepoints(basis_l.bases[1]))/sqrt(basis_r.bases[1].N)
    else
        mul_before = ones(N[1])
        mul_after = ones(N[1])
    end
    for i=2:n
        if any(i .== index)
            mul_before = kron(exp.(1im*basis_l.bases[i].xmin*(samplepoints(basis_r.bases[i]) .- basis_r.bases[i].pmin)), mul_before)
            mul_after = kron(exp.(1im*basis_r.bases[i].pmin*samplepoints(basis_l.bases[i]))/sqrt(basis_r.bases[i].N), mul_after)
        else
            mul_before = kron(ones(N[i]), mul_before)
            mul_after = kron(ones(N[i]), mul_after)
        end
    end
    mul_before = reshape(mul_before, (N...,))
    mul_after = reshape(mul_after, (N...,))

    x = Array{ComplexF64}(undef, N...)
    if ket_only
        FFTKets(basis_l, basis_r, plan_fft!(x, index), plan_bfft!(x, index), mul_before, mul_after)
    else
        A = Array{ComplexF64}(undef, [N; N]...)
        FFTOperators(basis_l, basis_r, plan_fft!(x, index), plan_bfft!(x, index), plan_fft!(A, [n + 1:2n;][index]), plan_bfft!(A, [1:n;][index]), mul_before, mul_after)
    end
end

function transform_px(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Vector{Int}; ket_only::Bool=false)
    n = length(basis_l.bases)
    Lx = [(b.xmax - b.xmin) for b=basis_r.bases[index]]
    dp = [spacing(b) for b=basis_l.bases[index]]
    dx = [spacing(b) for b=basis_r.bases[index]]
    N = [length(b) for b=basis_l.bases]
    for i=1:n
        if N[i] != length(basis_r.bases[i])
            throw(IncompatibleBases())
        end
    end
    for i=1:length(index)
        if abs(2*pi/dp[i] - Lx[i])/Lx[i] > 1e-12
            throw(IncompatibleBases())
        end
    end

    if index[1] == 1
        mul_before = exp.(-1im*basis_l.bases[1].pmin*(samplepoints(basis_r.bases[1]) .- basis_r.bases[1].xmin))
        mul_after = exp.(-1im*basis_r.bases[1].xmin*samplepoints(basis_l.bases[1]))/sqrt(N[1])
    else
        mul_before = ones(N[1])
        mul_after = ones(N[1])
    end
    for i=2:n
        if i in index
            mul_before = kron(exp.(-1im*basis_l.bases[i].pmin*(samplepoints(basis_r.bases[i]) .- basis_r.bases[i].xmin)), mul_before)
            mul_after = kron(exp.(-1im*basis_r.bases[i].xmin*samplepoints(basis_l.bases[i]))/sqrt(N[i]), mul_after)
        else
            mul_before = kron(ones(N[i]), mul_before)
            mul_after = kron(ones(N[i]), mul_after)
        end
    end
    mul_before = reshape(mul_before, (N...,))
    mul_after = reshape(mul_after, (N...,))

    x = Array{ComplexF64}(undef, N...)
    if ket_only
        FFTKets(basis_l, basis_r, plan_bfft!(x, index), plan_fft!(x, index), mul_before, mul_after)
    else
        A = Array{ComplexF64}(undef, [N; N]...)
        FFTOperators(basis_l, basis_r, plan_bfft!(x, index), plan_fft!(x, index), plan_bfft!(A, [n + 1:2n;][index]), plan_fft!(A, [1:n;][index]), mul_before, mul_after)
    end
end

DenseOperator(op::FFTOperator) = op*identityoperator(DenseOpType, op.basis_r)

dagger(op::FFTOperators) = transform(op.basis_r, op.basis_l)
dagger(op::FFTKets) = transform(op.basis_r, op.basis_l; ket_only=true)

tensor(A::FFTOperators, B::FFTOperators) = transform(tensor(A.basis_l, B.basis_l), tensor(A.basis_r, B.basis_r))
tensor(A::FFTKets, B::FFTKets) = transform(tensor(A.basis_l, B.basis_l), tensor(A.basis_r, B.basis_r); ket_only=true)

function mul!(result::Ket{B1},M::FFTOperator{B1,B2},b::Ket{B2},alpha_,beta_) where {B1<:Basis,B2<:Basis}
    alpha = convert(ComplexF64, alpha_)
    beta = convert(ComplexF64, beta_)
    N::Int = length(M.basis_r)
    if beta==Complex(0.)
        @inbounds for i=1:N
            result.data[i] = M.mul_before[i] * b.data[i]
        end
        M.fft_r! * reshape(result.data, size(M.mul_before))
        @inbounds for i=1:N
            result.data[i] *= M.mul_after[i] * alpha
        end
    else
        psi_ = Ket(M.basis_l, copy(b.data))
        @inbounds for i=1:N
            psi_.data[i] *= M.mul_before[i]
        end
        M.fft_r! * reshape(psi_.data, size(M.mul_before))
        @inbounds for i=1:N
            result.data[i] = beta*result.data[i] + alpha * psi_.data[i] * M.mul_after[i]
        end
    end
    result
end

function mul!(result::Bra{B2},b::Bra{B1},M::FFTOperator{B1,B2},alpha_,beta_) where {B1<:Basis,B2<:Basis}
    alpha = convert(ComplexF64, alpha_)
    beta = convert(ComplexF64, beta_)
    N::Int = length(M.basis_l)
    if beta==Complex(0.)
        @inbounds for i=1:N
            result.data[i] = conj(M.mul_after[i]) * conj(b.data[i])
        end
        M.fft_l! * reshape(result.data, size(M.mul_after))
        @inbounds for i=1:N
            result.data[i] = conj(result.data[i]) * M.mul_before[i] * alpha
        end
    else
        psi_ = Bra(M.basis_r, conj(b.data))
        @inbounds for i=1:N
            psi_.data[i] *= conj(M.mul_after[i])
        end
        M.fft_l! * reshape(psi_.data, size(M.mul_after))
        @inbounds for i=1:N
            result.data[i] = beta*result.data[i] + alpha * conj(psi_.data[i]) * M.mul_before[i]
        end
    end
    result
end

function mul!(result::Operator{B1,B3,T},A::Operator{B1,B2},B::FFTOperators{B2,B3},alpha_,beta_) where {B1<:Basis,B2<:Basis,B3<:Basis,T}
    alpha = convert(ComplexF64, alpha_)
    beta = convert(ComplexF64, beta_)
    if beta != Complex(0.)
        data = similar(result.data, size(result.data, 1), size(result.data, 2))
    else
        data = result.data
    end
    copyto!(data, A.data)
    @inbounds for j=1:length(B.mul_after), i=1:length(B.mul_after)
        data[i, j] *= B.mul_after[j]
    end
    conj!(data)
    n = size(B.mul_after)
    B.fft_l2! * reshape(data, n..., n...)
    conj!(data)
    N = prod(n)
    @inbounds for j=1:N, i=1:N
        data[i, j] *= B.mul_before[j]
    end
    if alpha != Complex(1.)
        lmul!(alpha, data)
    end
    if beta != Complex(0.)
        rmul!(result.data, beta)
        result.data += data
    end
    result
end

function mul!(result::Operator{B1,B3,T},A::FFTOperators{B1,B2},B::Operator{B2,B3},alpha_,beta_) where {B1<:Basis,B2<:Basis,B3<:Basis,T}
    alpha = convert(ComplexF64, alpha_)
    beta = convert(ComplexF64, beta_)
    if beta != Complex(0.)
        data = similar(result.data, size(result.data, 1), size(result.data, 2))
    else
        data = result.data
    end
    copyto!(data, B.data)
    @inbounds for j=1:length(A.mul_before), i=1:length(A.mul_before)
        data[i, j] *= A.mul_before[i]
    end
    n = size(A.mul_before)
    A.fft_r2! * reshape(data, n...,n...)
    N = prod(n)
    @inbounds for j=1:N, i=1:N
        data[i, j] *= A.mul_after[i]
    end
    if alpha != Complex(1.)
        lmul!(alpha, data)
    end
    if beta != Complex(0.)
        rmul!(result.data, beta)
        result.data += data
    end
    result
end


function deltasum(p0, ps, k, m2, T,δk=0.1)
    if p0 + δk > 2 * Epi(k, m2) > p0 - δk
        return (k * (-2 * k + ps) * coth(Epi(k, m2) / (2 * T))) /
               (8 * Epi(k, m2)^2 * (2 * Epi(k, m2) + p0) * pi^2)
    else
        return (k * (-2 * k + ps) * coth(sqrt(k^2 + m2) / (2 * T))) /
               (2 * sqrt(k^2 + m2) * (4 * (k^2 + m2) - p0^2) * pi^2)
    end
end

function F1tilde_delta(p0, p, k, m, T)
    (
        k *
        (2 * k - p) *
        csch(sqrt(k^2 + m) / (2 * T))^2 *
        (
            sqrt(k^2 + m) *
            (2 * k - p) *
            (4 * k + p) *
            (-4 * (k^2 + m) + p0^2) +
            (
                (2 * k - p) * (4 * k + p) * p0^2 +
                12 * (k^2 + m) * (8 * m + 2 * k * p + p^2 - 2 * p0^2)
            ) *
            T *
            sinh(sqrt(k^2 + m) / T)
        )
    ) / (192 * (k^2 + m)^(3 / 2) * (-4 * (k^2 + m) + p0^2)^2 * pi^2 * T)
end

function dkF1Tilde(p0, ps, k, m, T;kwarg...)
    if k > ps / 2
        hquadrature(
            x -> 2 * x * (-p0^2 + x^2)^-1 * ppfun(x, ps, k, Epi(k, m), T),
            2 * Epi(k, m),
            Epi(k + ps, m) + Epi(k, m);kwarg...
        )[1] +
        deltasum(p0, ps, k, m, T) +
        F1tilde_delta(p0, ps, k, m, T)
    elseif k <= ps / 2
        hquadrature(
            x -> 2 * x * (-p0^2 + x^2)^-1 * ppfun(x, ps, k, Epi(k, m), T),
            2 * Epi(k, m),
            Epi(k + ps, m) + Epi(k, m);kwarg...
        )[1]
    end
end

function dkF2Tilde(p0, ps, k, m, T;kwarg...)
    if k > ps / 2
        if Epi(k + ps, m) - Epi(k, m) < p0
            return hquadrature(
                x -> 2 * x * (-p0^2 + x^2)^-1 * pmfun(x, ps, k, Epi(k, m), T),
                0.0,
                Epi(k + ps, m) - Epi(k, m);kwarg...
            )[1]+pmfun_zerofix(p0, ps, k, m, T,9)
        else
            return quadgk_PV(
                x -> pmfun(x, ps, k, Epi(k, m), T),
                0.0,
                p0,
                Epi(ps + k, m) - Epi(k, m);kwarg...
            )[1]+pmfun_zerofix(p0, ps, k, m, T,9)
        end
    elseif k <= ps / 2
        if Epi(k + ps, m) - Epi(k, m) < p0 || p0 < Epi(k - ps, m) - Epi(k, m)
            return hquadrature(
                x -> 2 * x * (-p0^2 + x^2)^-1 * pmfun(x, ps, k, Epi(k, m), T),
                Epi(k - ps, m) - Epi(k, m),
                Epi(k + ps, m) - Epi(k, m);kwarg...
            )[1]
        else
            return quadgk_PV(
                x -> pmfun(x, ps, k, Epi(k, m), T),
                Epi(k - ps, m) - Epi(k, m),
                p0,
                Epi(ps + k, m) - Epi(k, m);kwarg...
            )[1]
        end
    end
end

dkF1TildeAll(q0, qs, k, m, T;kwarg...)=dkF1Tilde(abs(q0), qs, k, m, T;kwarg...)
dkF2TildeAll(q0, qs, k, m, T;kwarg...)=dkF2Tilde(abs(q0), qs, k, m, T;kwarg...)

function dklamdabar_Zero(q0, qs, k, m, T, lam4pik, Npi;kwarg...)
        lam4pik^2 *
        (
            (2 + Npi) * (dkF1TildeAll(k, m, T)+dkF2TildeAll(k, m, T)) +
             6 * dkF1TildeAll(q0, qs, k, m, T;kwarg...) +
             6 * dkF2TildeAll(q0, qs, k, m, T;kwarg...)
        )
end

function lamdabar_Zero(q0,qs,k, T, Npi, mfun,lamfun,UVScale;kwarg...)
    -hquadrature(
        kprim -> dklamdabar_Zero(q0, qs, kprim, mfun(kprim), T, lamfun(kprim), Npi;kwarg...),
        k,
        UVScale;kwarg...
        # initdiv=200,
    )[1] + lamfun(UVScale)
end


include("gridap_warmup.jl")

using PartitionedArrays
using PartitionedPoisson

#poisson(mode=:seq,nc=(10,10),np=(2,2))
#poisson(mode=:seq,nc=(10,10,10),np=(2,2,2))
poisson(mode=:mpi,nc=(10,10),np=(1,1),verbose=true)
poisson(mode=:mpi,nc=(10,10,10),np=(1,1,1),verbose=true)


mutable struct Core

    wksp      :: Workspace

#    mask_sT   :: AbstractArray{Float64, 3}
#    mask_T    :: AbstractArray{Float64, 3}

    amo_slab  :: Union{AdvancedMatrixOperators, Nothing}
    amo       :: Union{AdvancedMatrixOperators, Nothing}

    mtx       :: Dict

    vd        :: VerticalDiffusion

    cdatam     :: Union{CyclicDataManager, Nothing}

    function Core(
        ev :: Env,
        tmpfi :: TempField,
    )

        cfg = ev.config

        wksp = Workspace(Nx=ev.Nx, Ny=ev.Ny, Nz=ev.Nz)

        gf = PolelikeCoordinate.CurvilinearSphericalGridFile(
            cfg["domain_file"];
            R  = Re,
            Ω  = Ω,
        )

        gd      = ev.gd
        gd_slab = ev.gd_slab

        amo_slab = AdvancedMatrixOperators(;
            gd = gd_slab,
            mask_T     = ev.topo.sfcmask_sT,
            deepmask_T = ev.topo.sfcmask_sT,
        )

        amo = AdvancedMatrixOperators(;
            gd = gd,
            mask_T = ev.topo.mask_T,
            deepmask_T = ev.topo.deepmask_T,
        )

        # Build Advection Matrix
        function build!(id_mtx, idx)
            local result
            # using transpose speeds up by 100 times
            tp = transpose(id_mtx) |> sparse
            result = transpose(tp[:, view(idx, :)]) |> sparse
            dropzeros!(result)

            idx .= 0 # clean so that debug is easir when some girds are not assigned
            return result
        end

        # Build the matrix to broadcast sW to W grid
        # Notice that we cannot use W_pts because gd_slab has two W layers
        num_sT = reshape( collect(1:amo_slab.bmo.T_pts), amo_slab.bmo.T_dim...)
        mapping_T = repeat(num_sT, outer=(ev.Nz+1, 1, 1))
        W_broadcast_sT = build!(amo_slab.bmo.T_I_T, mapping_T)

        # Build Radiation Matrix
        swflx_factor_W =  cfg["rad_R"]  * exp.(gd.z_W / cfg["rad_ζ1"]) + (1.0 - cfg["rad_R"]) * exp.(gd.z_W / cfg["rad_ζ2"])

        Nz_bot = view(ev.topo.Nz_bot_sT, 1, :, :)
        #println("Nz_bot: ", Nz_bot[20, 114])
        #println("Before: ", swflx_factor_W[:, 20, 114])
        for j=1:ev.Ny, i=1:ev.Nx
            swflx_factor_W[(Nz_bot[i, j]+1):end, i, j] .= 0.0 # Bottom absorbs everything
        end
        #println("After: ", swflx_factor_W[:, 20, 114])

        #swflx_factor_W .= 0
        #swflx_factor_W[1, :, :] .= 1.0

        # Surface flux is immediately absorbed at the top layer
        sfcflx_factor_W = 0.0 * gd.z_W
        sfcflx_factor_W[1, :, :] .= 1.0

        #Ks_H matrix
#        Ks_H_U = zeros(Float64, amo.bmo.V_dim...)

        # This design is to prevent the equator develop wind-Ekman-SST feedback
        # that creates spurious rainband in Central Pacific
        Ks_H_func = (ϕ, z) -> 500.0 + (20000.0 - 500.0) * exp( - 0.5 * (ϕ/deg2rad(10.0))^2 ) * exp(z/100.0)

#        println(size(gd.ϕ_U))
#        println(size(gd.ϕ_V))
#        println(size(gd.ϕ_UV))
#        println(size(gd.z_V))

        Ks_H_U = reshape(Ks_H_func.(gd.ϕ_U, gd.z_U), amo.bmo.U_dim...)
        Ks_H_V = reshape(Ks_H_func.(gd.ϕ_V, gd.z_V), amo.bmo.V_dim...)
        Ks_H_T = reshape(Ks_H_func.(gd.ϕ_T, gd.z_T), amo.bmo.T_dim...)

        # f and ϵ matrices
        f_sT = 2 * gd.Ω * sin.(gd_slab.ϕ_T)
        β_sT = (2 * gd.Ω / gd.R) * cos.(gd_slab.ϕ_T)
        ϵ_sT = f_sT * 0 .+ cfg["ϵ"]
        D_sT = f_sT.^2 + ϵ_sT.^2
        invD_sT = D_sT.^(-1.0)

        mtx = Dict(
            :ones_T          => ones(Float64, amo.bmo.T_pts),
            :T_swflxConv_sT  => - amo.T_mask_T * amo.T_DIVz_W * spdiagm(0 => view(swflx_factor_W, :)) * W_broadcast_sT,
            :T_sfcflxConv_sT => - amo.T_mask_T * amo.T_DIVz_W * spdiagm(0 => view(sfcflx_factor_W, :)) * W_broadcast_sT,
            :invD_sT         => invD_sT,
            :f_sT            => f_sT,
            :β_sT            => β_sT,
            :ϵ_sT            => ϵ_sT,
            :D_sT            => D_sT,
            :Ks_H_T          => Ks_H_T,
            :Ks_H_U          => Ks_H_U,
            :Ks_H_V          => Ks_H_V,
        )

        tmpfi.check_usage[:Ks_H_U] .= Ks_H_U
        tmpfi.check_usage[:Ks_H_V] .= Ks_H_V

        vd = VerticalDiffusion(
            amo;
            K_iso=cfg["Ks_V"],
            K_cva=(cfg["convective_adjustment"] == "on") ? cfg["Ks_V_cva"] : cfg["Ks_V"],
        )

        # surface mask but is in 3D T grid. This one is different from the topo.sfcmask_sT.
        sfcmask_T = zeros(Float64, amo.bmo.T_dim...)
        sfcmask_T[1, :, :] .= 1
        mtx[:T_sfcmask_T] = spdiagm(0 => amo.T_mask_T * reshape(sfcmask_T, :))

        if length(ev.cdata_varnames) == 0
            cdatam = nothing
        else

            if cfg["cdata_var_file_map"] == nothing
                throw(ErrorException("Some config require cyclic data forcing file: $(ev.cdata_varnames)"))
            else
                function parseDateTime(timetype, str)
                    m = match(r"(?<year>[0-9]+)-(?<month>[0-9]{2})-(?<day>[0-9]{2})\s+(?<hour>[0-9]{2}):(?<min>[0-9]{2}):(?<sec>[0-9]{2})", str)
                    if m == nothing
                        throw(ErrorException("Unknown time format: " * (str)))
                    end

                    return timetype(
                        parse(Int64, m[:year]),
                        parse(Int64, m[:month]),
                        parse(Int64, m[:day]),
                        parse(Int64, m[:hour]),
                        parse(Int64, m[:min]),
                        parse(Int64, m[:sec]),
                    )
                end

                timetype   = getproperty(CFTime, Symbol(cfg["timetype"]))

                var_file_map = Dict()
                for varname in ev.cdata_varnames
                    var_file_map[varname] = cfg["cdata_var_file_map"][varname]
                end

                cdatam = CyclicDataManager(;
                    timetype     = timetype,
                    var_file_map = var_file_map,
                    beg_time     = parseDateTime(timetype, cfg["cdata_beg_time"]),
                    end_time     = parseDateTime(timetype, cfg["cdata_end_time"]),
                    align_time   = parseDateTime(timetype, cfg["cdata_align_time"]),
                    sub_yrng     = ev.sub_yrng,
                )

                tmpfi.datastream = makeDataContainer(cdatam)
            end
        end

        if cfg["weak_restoring"] == "on"
            mtx[:T_invτwk_TEMP_T] = - amo.T_mask_T * spdiagm(0 => ones(Float64, amo.bmo.T_pts)) / cfg["τwk_TEMP"]
            mtx[:T_invτwk_SALT_T] = - amo.T_mask_T * spdiagm(0 => ones(Float64, amo.bmo.T_pts)) / cfg["τwk_SALT"]
        end

        return new(
            wksp,

            #mask_sT,
            #ev.topo.mask_T,

            amo_slab,
            amo,

            mtx,

            vd,

            cdatam,

        )
    end

end


## Example script for linear-strength vortex panel method
using LinearAlgebra
using Base.Iterators
using AeroMDAO

## Airfoil
alpha_u = [0.2, 0.3, 0.2, 0.15, 0.2]
alpha_l = [-0.2, -0.1, -0.1, -0.001]
dzs     = (0., 0.)
# airfoil = kulfan_CST(alpha_u, alpha_l, dzs, 0.0, 60);
airfoil = naca4((0,0,1,2))
V, α    = 1., 5.
ρ       = 1.225
uniform = Uniform2D(V, α)
num_pans = 80
panels  = make_panels(airfoil, num_pans);

## Vortex panel
A = @time vortex_influence_matrix(panels)
b = @time neumann_boundary_condition(panels, velocity(uniform))
γs = A \ b

##
cl = 2 * sum(@. (γs[1:end-1] + γs[2:end]) / 2 * panel_length(panels)) / uniform.magnitude

## Velocities and pressure coefficient
# panel_vels = [ velocity(uniform) .+ sum(vortex_velocity.(γs[1:end-1], γs[2:end], panels, Ref(panel_i))) for panel_i in panels ]

## Plotting
using CairoMakie

CairoMakie.activate!()

using LaTeXStrings

const LS = LaTeXString

## Pressure coefficient
qts = γs
cps = @. 1 - qts^2 / uniform.magnitude^2

upper, lower = get_surface_values(panels, cps)
upper, lower = [ lower[end]; upper ], [ upper[end]; lower ]

# Plot
f1 = Figure(resolution = (900, 400))
ax = Axis(f1[1,1], aspect = 1, yreversed = true, xlabel = L"(x/c)", ylabel = L"C_p")
l1 = lines!(upper, label = L"\mathrm{Upper}")
l2 = lines!(lower, label = L"\mathrm{Lower}")

axislegend(L"\mathrm{Surface}")

# Velocity field
x_domain, y_domain = (-1, 2), (-1, 1)
grid_size = 50
x_dom, y_dom = range(x_domain..., length = grid_size), range(y_domain..., length = grid_size)
grid = product(x_dom, y_dom)
pts = panel_points(panels);

total_velocity(xs, ys) = map((x, y) -> Point2f(velocity(uniform) .+ sum(vortex_velocity.(γs[1:end-1], γs[2:end], panels, x, y))...), xs, ys)

vortex_vels = total_velocity(first.(grid), last.(grid))
speeds = @. sqrt(first(vortex_vels)^2 + last(vortex_vels)^2)
cps = @. 1 - speeds^2 / uniform.magnitude^2

# Plot
ax2 = Axis(f1[1,2], aspect = 1)
hm = contourf!(x_dom, y_dom, speeds, levels = 50)

streamplot!(total_velocity, first.(grid), last.(grid), colormap = Reverse(:viridis))
poly!(pts, color = "black")

Colorbar(f1[1,3], hm, label = L"Normalized Speed $(U/U_\infty)$")
colsize!(f1.layout, 2, Aspect(1, 1.0))

# f1[0,:] = Label(f1, LS("Linear Vortex Panel Method"))

# arrows!(x_dom, y_dom, first.(vortex_vels), last.(vortex_vels))
# CB2 = colorbar(orientation = "horizontal", label = "Relative Airspeed (U/U∞)")
# tight_layout()
# show()
f1

##
save("plots/LinearVortex.svg", f1, px_per_unit = 1.5)


function setup_TriangleMesh(::Type{T}, ::Type{U}) where {T, U}
    name = "tri"
    vertices = [Point{2,T}(   0, 0),
                Point{2,T}(   1, 0),
                Point{2,T}(1//2, 1),
                Point{2,T}(3//2, 1)]
    polytopes = [Triangle{U}(1, 2, 3), Triangle{U}(2, 4, 3)]
    groups = Dict{String, BitSet}()
    groups["A"] = BitSet([1])
    groups["B"] = BitSet([2])
    return PolytopeVertexMesh(name, vertices, polytopes, groups)
end

function setup_QuadrilateralMesh(::Type{T}, ::Type{U}) where {T, U}
    name = "quad"
    vertices = [Point{2,T}(0, 0),
                Point{2,T}(1, 0),
                Point{2,T}(1, 1),
                Point{2,T}(0, 1),
                Point{2,T}(2, 0),
                Point{2,T}(2, 1)]
    polytopes = [Quadrilateral{U}(1, 2, 3, 4), Quadrilateral{U}(2, 5, 6, 3)]
    groups = Dict{String, BitSet}()
    groups["A"] = BitSet([1])
    groups["B"] = BitSet([2])
    return PolytopeVertexMesh(name, vertices, polytopes, groups)
end

function setup_MixedPolygonMesh(::Type{T}, ::Type{U}) where {T, U}
    name = "tri_quad"
    vertices = [Point{2,T}(0, 0),
                Point{2,T}(1, 0),
                Point{2,T}(1, 1),
                Point{2,T}(0, 1),
                Point{2,T}(2, 0)]
    polytopes = [Quadrilateral{U}(1, 2, 3, 4), Triangle{U}(2, 5, 3)]
    groups = Dict{String, BitSet}()
    groups["A"] = BitSet([1])
    groups["B"] = BitSet([2])
    return PolytopeVertexMesh(name, vertices, polytopes, groups)
end

function setup_QuadraticTriangleMesh(::Type{T}, ::Type{U}) where {T, U}
    name = "tri6"
    vertices = [Point{2,T}(   0,    0),
                Point{2,T}(   1,    0),
                Point{2,T}(   0,    1),
                Point{2,T}(1//2,    0),
                Point{2,T}(2//5, 2//5),
                Point{2,T}(   0, 1//2),
                Point{2,T}(   1,    1),
                Point{2,T}(   1, 1//2),
                Point{2,T}(1//2,    1)]
    polytopes = [QuadraticTriangle{U}(1, 2, 3, 4, 5, 6),
                 QuadraticTriangle{U}(2, 7, 3, 8, 9, 5)]
    groups = Dict{String, BitSet}()
    groups["A"] = BitSet([1])
    groups["B"] = BitSet([2])
    return PolytopeVertexMesh(name, vertices, polytopes, groups)
end

function setup_QuadraticQuadrilateralMesh(::Type{T}, ::Type{U}) where {T, U}
    name = "quad8"
    vertices = [Point{2,T}(    0,    0),
                Point{2,T}(    1,    0),
                Point{2,T}(    1,    1),
                Point{2,T}(    0,    1),
                Point{2,T}(    2,    0),
                Point{2,T}(    2,    1),
                Point{2,T}( 1//2,    0),
                Point{2,T}(7//10, 1//2),
                Point{2,T}( 1//2,    1),
                Point{2,T}(    0, 1//2),
                Point{2,T}( 3//2,    0),
                Point{2,T}(    2, 1//2),
                Point{2,T}( 3//2,    1),
               ]
    polytopes = [QuadraticQuadrilateral{U}(1, 2, 3, 4, 7, 8, 9, 10),
                 QuadraticQuadrilateral{U}(2, 5, 6, 3, 11, 12, 13, 8)]
    groups = Dict{String, BitSet}()
    groups["A"] = BitSet([1])
    groups["B"] = BitSet([2])
    return PolytopeVertexMesh(name, vertices, polytopes, groups)
end

function setup_MixedQuadraticPolygonMesh(::Type{T}, ::Type{U}) where {T, U}
    name = "tri6_quad8"
    vertices = [Point{2,T}(    0,    0),
                Point{2,T}(    1,    0),
                Point{2,T}(    1,    1),
                Point{2,T}(    0,    1),
                Point{2,T}(    2,    0),
                Point{2,T}( 1//2,    0),
                Point{2,T}(7//10, 1//2),
                Point{2,T}( 1//2,    1),
                Point{2,T}(    0, 1//2),
                Point{2,T}( 3//2,    0),
                Point{2,T}( 3//2, 1//2),
               ]
    polytopes = [QuadraticQuadrilateral{U}(1, 2, 3, 4, 6, 7, 8, 9),
                 QuadraticTriangle{U}(2, 5, 3, 10, 11, 7)]
    groups = Dict{String, BitSet}()
    groups["A"] = BitSet([1])
    groups["B"] = BitSet([2])
    return PolytopeVertexMesh(name, vertices, polytopes, groups)
end


using Test
using QuantumPropagators
using LinearAlgebra
using UnicodePlots

@testset "TLS Rabi Cycling" begin

    # We use Rabi cycling in a TLS as a test case, which allows to to compare
    # the propagation with the known analytic solution.

    SHOWPLOT = false

    Ψ0 = ComplexF64[1, 0]
    Ĥ = ComplexF64[0  0.5; 0.5 0]
    tlist = collect(range(0, 1.5π, length=101)) # 3π/2 pulse

    genfunc(tlist, i; kwargs...) = Ĥ

    storage = init_storage(Ψ0, tlist)
    @test isa(storage, Matrix)
    @test eltype(storage) == ComplexF64

    Ψ_out = propagate(Ψ0, genfunc, tlist; storage=storage)
    Ψ_expected = ComplexF64[-1/√2, -1im/√2]  # note the phases

    pop0 = abs.(storage[1,:]).^2
    SHOWPLOT && println(
        lineplot(tlist ./ π, pop0, ylim=[0, 1], title="fw prop")
    )

    @test norm(Ψ_out - Ψ_expected) < 1e-12
    @test pop0[end] ≈ 0.5

    # Propagating backwards in time should exactly reverse the dyanmics (since
    # they are unitary). Thus, we should end up back at the initial state, and
    # the stored states (being filled back-to-front) should exactly match the
    # stored states from the forward propagation.

    storage_bw = init_storage(Ψ0, tlist)
    Ψ_out_bw = propagate(Ψ_out, genfunc, tlist; backwards=true,
                         storage=storage_bw)

    pop0_bw = abs.(storage_bw[1,:]).^2
    SHOWPLOT && println(
        lineplot(tlist ./ π, pop0_bw, ylim=[0, 1], title="bw prop")
    )

    @test norm(Ψ_out_bw - Ψ0) < 1e-12
    @test pop0_bw[1] ≈ 1.0

    @test norm(storage - storage_bw) < 1e-12

end

@testset "Optomech propagation" begin
    include("optomech.jl")
    Ψ0 = ket(0, N_cav) ⊗ ket(2, N_mech)
    H = H_cav + H_mech + H_int
    tlist = [0:0.2:50;]
    Ψ1 = propagate(Ψ0, (tlist, i; kwargs...)->H, tlist, method=:newton)
    @test (norm(Ψ1) - 1.0) < 1e-12
    Ψ2 = propagate(Ψ0, (tlist, i; kwargs...)->H, tlist, method=:cheby)
    @test norm(Ψ1 - Ψ2) < 1e-10
    Ψ3 = propagate(Ψ0, (tlist, i; kwargs...)->H, tlist, method=:auto)
    @test norm(Ψ1 - Ψ3) < 1e-10
end


export Kraskov1, Kraskov2
using Neighborhood

# naive application of estimators in Entropies.jl
function mutualinfo(x::Vector_or_Dataset, y::Vector_or_Dataset, est::NearestNeighborEntropyEstimator; base = 2)
    Hx = genentropy(Dataset(x), est; base = base)
    Hy = genentropy(Dataset(y), est; base = base)
    Hxy = genentropy(Dataset(x, y), est; base = base)

    return Hx + Hy - Hxy
end

abstract type KNNMutualInformationEstimator <: MutualInformationEstimator end

# bias-corrected estimators from Kraskov et al. (2004)

"""
    Kraskov1(k::Int = 1; metric_x = Chebyshev(), metric_y = Chebyshev()) <: MutualInformationEstimator

The ``I^{(1)}`` nearest neighbor based mutual information estimator from
Kraskov et al. (2004), using `k` nearest neighbors. The distance metric for
the marginals ``x`` and ``y`` can be chosen separately, while the `Chebyshev` metric
is always used for the `z = (x, y)` joint space.
"""
struct Kraskov1 <: KNNMutualInformationEstimator
    k::Int
    metric_x::Metric
    metric_y::Metric
    metric_z::Metric # always Chebyshev, otherwise estimator is not valid!

    function Kraskov1(k::Int = 1; metric_x = Chebyshev(), metric_y = Chebyshev())
        new(k, metric_x, metric_y, Chebyshev())
    end
end

"""
    Kraskov2(k::Int = 1; metric_x = Chebyshev(), metric_y = Chebyshev()) <: MutualInformationEstimator

The ``I^{(2)}(x, y)`` nearest neighbor based mutual information estimator from
Kraskov et al. (2004), using `k` nearest neighbors. The distance metric for
the marginals ``x`` and ``y`` can be chosen separately, while the `Chebyshev` metric
is always used for the `z = (x, y)` joint space.
"""
struct Kraskov2 <: KNNMutualInformationEstimator
    k::Int
    metric_x::Metric
    metric_y::Metric
    metric_z::Metric # always Chebyshev, otherwise estimator is not valid!

    function Kraskov2(k::Int = 1; metric_x = Chebyshev(), metric_y = Chebyshev())
        new(k, metric_x, metric_y, Chebyshev())
    end
end

function eval_dists_to_knns!(ds, pts, knn_idxs, metric)
    @inbounds for i in 1:length(pts)
        ds[i] = evaluate(metric, pts[i], pts[knn_idxs[i]])
    end

    return ds
end

# In the Kraskov1 estimator, ϵs are the distances in the Z = (X, Y) joint space
# In the Kraskov2 estimator, ϵs are the distances in the X and Y marginal spaces
function count_within_radius!(p, x, metric, ϵs, N)
    @inbounds for i in 1:N
        ϵ = ϵs[i] / 2
        xᵢ = x[i]
        p[i] = count(evaluate(metric, xᵢ, x[j]) < ϵ for j in 1:N)
    end

    return p
end

function mutualinfo(x::Vector_or_Dataset{D1, T}, y::Vector_or_Dataset{D2, T}, est::Kraskov1;
        base = MathConstants.e) where {D1, D2, T}
    @assert length(x) == length(y)
    z = Dataset(x, y)
    N = length(z)

    # Common for both kraskov estimators
    tree_z = KDTree(z, est.metric_z)
    k = est.k
    idxs_z, dists_z = bulksearch(tree_z, z.data, NeighborNumber(k + 1))

    # TODO: These quantities are never used in the subsequent code. I am therefore
    # commenting them out...
    # X = Dataset(x)
    # Y = Dataset(y)
    # tree_x = KDTree(X, est.metric_x)
    # tree_y = KDTree(Y, est.metric_y)
    # idxs_x, dists_x = bulksearch(tree_x, X.data, NeighborNumber(k + 1))
    # idxs_y, dists_y = bulksearch(tree_y, Y.data, NeighborNumber(k + 1))

    kth_nns_z = [idx_z[k + 1] for idx_z in idxs_z]
    ϵs_z = [dz[k + 1] for dz in dists_z]
    ϵs_x = zeros(Float64, N)
    ϵs_y = zeros(Float64, N)
    eval_dists_to_knns!(ϵs_x, x, kth_nns_z, est.metric_x)
    eval_dists_to_knns!(ϵs_y, y, kth_nns_z, est.metric_y)

    # if the following equality holds for all points, then things are correct until this point
    #ϵ_maxes = [max(a, b) for (a, b) in zip(ϵs_x, ϵs_y)]
    #@assert all(ϵ_maxes .== ϵs_z)

    # Kraskov1 estimator
    nx = zeros(Int, N)
    ny = zeros(Int, N)
    count_within_radius!(nx, x, est.metric_x, ϵs_z, N)
    count_within_radius!(ny, y, est.metric_y, ϵs_z, N)

    MI = digamma(est.k) - sum(digamma.(nx .+ 1) + digamma.(ny .+ 1))/N + digamma(N)

    # Kraskov uses the natural logarithm in their derivations, so need to convert in the last step
    if base != MathConstants.e
        return MI / log(base)
    else
        return MI
    end
end

function mutualinfo(x::Vector_or_Dataset{D1, T}, y::Vector_or_Dataset{D2, T}, est::Kraskov2;
        base = MathConstants.e) where {D1, D2, T}
    @assert length(x) == length(y)
    z = Dataset(x, y)
    N = length(z)

    # Common for both kraskov estimators
    tree_z = KDTree(z, est.metric_z)

    k = est.k
    idxs_z = bulkisearch(tree_z, z.data, NeighborNumber(k + 1))

    kth_nns_z = [idx_z[k + 1] for idx_z in idxs_z]
    ϵs_x = zeros(Float64, N)
    ϵs_y = zeros(Float64, N)
    eval_dists_to_knns!(ϵs_x, x, kth_nns_z, est.metric_x)
    eval_dists_to_knns!(ϵs_y, y, kth_nns_z, est.metric_y)

    # TODO: The following quantities are never used in this code,
    # so I am commenting them out...
    # idxs_z, dists_z = bulksearch(tree_z, z.data, NeighborNumber(k + 1))
    # X = Dataset(x)
    # Y = Dataset(y)
    # tree_x = KDTree(X, est.metric_x)
    # tree_y = KDTree(Y, est.metric_y)
    # idxs_x, dists_x = knn(tree_x, X.data, k + 1, true)
    # idxs_y, dists_y = knn(tree_y, Y.data, k + 1, true)
    # ϵs_z = [dz[k + 1] for dz in dists_z]

    # if the following equality holds for all points, then things are correct until this point
    #ϵ_maxes = [max(a, b) for (a, b) in zip(ϵs_x, ϵs_y)]
    #@assert all(ϵ_maxes .== ϵs_z)

    # Kraskov2 estimator
    nx = zeros(Int, N)
    ny = zeros(Int, N)
    count_within_radius!(nx, x, est.metric_x, ϵs_x, N)
    count_within_radius!(ny, y, est.metric_y, ϵs_y, N)

    MI = digamma(est.k) - 1/k - sum(digamma.(nx) .+ digamma.(ny))/N + digamma(N)

    # Kraskov uses the natural logarithm in their derivations, so need to convert in the last step
    if base != MathConstants.e
        return MI / log(base)
    else
        return MI
    end
end


module TestPI

using Test, PIMD

@testset "PI" begin
    β, m, V = 1., 1., harmonicPotential
    system = OneDimensionalQuantumSystem(β, m, V)
    @test system.potential(10.) ≈ 50.
    @test system.force(10.) ≈ -10.
    @test system.hessian(10.) ≈ 1.
end

end


using Abaco
using Test

Abaco.loginit()

# My time zone: Saturday, October 30, 2021 3:43:35 PM GMT+02:00 DST
window = 900
metric_ts = 1635600600
metric_sn = "Civetta"

expected_triggers = 10
actual_triggers = 0

#1
abaco = abaco_init(interval=window, ages=10) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[identity] age [$ts]: scope: [$ne] $name = $value"
    @test ts == metric_ts
    @test ne == metric_sn
    @test value == 100.0
    actual_triggers += 1
end

formula(abaco, "r = x")
ingest(abaco, metric_ts, metric_sn, "x", 100.0)

#2
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[negation] age [$ts]: scope: [$ne] $name = $value"
    @test value == -100
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = -x")
ingest(abaco, metric_ts, metric_sn, "x", 100)

#3
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[add] age [$ts]: scope: [$ne] $name = $value"
    @test value == 6.2
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = x + y")
ingest(abaco, metric_ts, metric_sn, "x", 1.0)
ingest(abaco, metric_ts, metric_sn, "y", 5.2)

#4
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[mult] age [$ts]: scope: [$ne] $name = $value"
    @test value == 3.1*2.1
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = x * y")
ingest(abaco, metric_ts, metric_sn, "x", 3.1)
ingest(abaco, metric_ts, metric_sn, "y", 2.1)

#5
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[log10] age [$ts]: scope: [$ne] $name = $value"
    @test value == 2
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = log10(x)")
ingest(abaco, metric_ts, metric_sn, "x", 100.0)

#6
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[exp] age [$ts]: scope: [$ne] $name = $value"
    @test value == exp(1)
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = exp(x)")
ingest(abaco, metric_ts, metric_sn, "x", 1)

#7
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[div] age [$ts]: scope: [$ne] $name = $value"
    @test value == 1
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = div(x, y)")
ingest(abaco, metric_ts, metric_sn, "x", 5)
ingest(abaco, metric_ts, metric_sn, "y", 3.1)

#8
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[x^(y-1)] age [$ts]: scope: [$ne] $name = $value"
    @test value == 4
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = x^(y-1)")
ingest(abaco, metric_ts, metric_sn, "x", 2)
ingest(abaco, metric_ts, metric_sn, "y", 3)

#9
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[x^2 + y^3 +3*x^4] age [$ts]: scope: [$ne] $name = $value"
    @test value == 60
    actual_triggers += 1
end
metric_ts += window
formula(abaco, "r = x^2 + x^3 + 3*x^4")
ingest(abaco, metric_ts, metric_sn, "x", 2)
ingest(abaco, metric_ts, metric_sn, "y", 3)

#10
oncomplete(abaco) do ts, ne, name, value, inputs
    global actual_triggers
    @debug "[factorial] age [$ts]: scope: [$ne] $name = $value"
    @test value == factorial(6)
    actual_triggers += 1
end
metric_ts += window
# if a function requires an integer argument try to convert the arg value to Int
formula(abaco, "r = factorial(Int(x+x))")
ingest(abaco, metric_ts, metric_sn, "x", 3)

@test actual_triggers === expected_triggers


import LinearAlgebra: lmul!, rmul!

function rmul!(A::AbstractVector, D::Diagonal)
    @assert !LinearAlgebra.has_offset_axes(A)
    A .= A .* transpose(D.diag)
    return A
end

function lmul!(D::Diagonal, B::AbstractVector)
    @assert !LinearAlgebra.has_offset_axes(B)
    B .= D.diag .* B
    return B
end


"""
    solve(
        equation::DiffusionEquation{ClosedIntervalDomain{T},F},
        u,
        tlist,
        M,
        N;
        method = "filterfirst",
        boundary_conditions = "exact",
        solver = QNDF(),
        abstol = 1e-4,
        reltol = 1e-3,
        degmax = 10,
        λ = 1e-6
    ) where {T,F}

Solve `equation` from `tlist[1]` to `tlist[2]` with initial conditions `u` and a
discretization of `M` and `N` points. If `method` is `"filterfirst"`, the equation is filtered then
discretized. If `method` is `"discretizefirst"`, the equation is discretized then filtered.
"""
function solve(
    equation::DiffusionEquation{ClosedIntervalDomain{T},F},
    u,
    tlist,
    M,
    N;
    method = "filterfirst",
    boundary_conditions = "exact",
    solver = QNDF(),
    abstol = 1e-4,
    reltol = 1e-3,
    degmax = 10,
    λ = 1e-6,
) where {T,F}

    @unpack domain, filter, f, g_a, g_b = equation
    x = discretize(domain, M)
    ξ = discretize(domain, N)

    # Get matrices
    D = diffusion_matrix(domain, N)
    W = filter_matrix(filter, domain, M, N; degmax, λ)
    R = reconstruction_matrix(filter, domain, M, N; degmax, λ)

    W₀ = W[:, 1]
    Wₙ = W[:, end]

    # Initial conditions
    ū = apply_filter(u, filter, domain)
    uₕ = u.(ξ)
    # ūₕ = ū.(x)
    ūₕ = filter_matrix(filter, domain, M, N) * uₕ

    function f!(fₕ, t)
        fₕ[1] = 0
        fₕ[2:end-1] .= f.(ξ[2:end-1], t)
        fₕ[end] = 0
    end

    if boundary_conditions == "exact"
        Mass = [
            spzeros(1, N + 1)
            spzeros(N - 1, 1) sparse(I, N - 1, N - 1) spzeros(N - 1, 1)
            spzeros(1, N + 1)
        ]
        J = [
            -1 spzeros(1, N)
            D[2:end-1, :]
            spzeros(1, N) -1
        ]
        γ_a = g_a
        γ_b = g_b
    elseif boundary_conditions == "derivative"
        Mass = sparse(I, N + 1, N + 1)
        # Zero out boundary u
        J = D
        J[[1, end], :] .= 0
        dropzeros!(J)
        γ_a = t -> derivative(g_a, t)
        γ_b = t -> derivative(g_b, t)
    else
        error("Unknown boundary conditions")
    end

    if method == "filterfirst"
        error("Not implemented")
    elseif method == "discretizefirst"
        # Solve discretized-then-filtered problem
        p = (; Ju = zero(x), J = W * J * R, fₕ = zero(ξ), Wf = zero(x))
        function Mdu!(du, u, p, t)
            # @show t
            @unpack Ju, J, fₕ, Wf = p
            f!(fₕ, t)
            mul!(Wf, W, fₕ)
            mul!(Ju, J, u)
            @. du = Ju + Wf + γ_a(t) * W₀ + γ_b(t) * Wₙ
        end
        odefunction = ODEFunction(
            Mdu!,
            jac = (J, u, p, t) -> (J .= p.J),
            jac_prototype = p.J,
            mass_matrix = W * Mass * R,
        )
        problem = ODEProblem(odefunction, ūₕ, tlist, p)
        solution = OrdinaryDiffEq.solve(problem, solver; abstol, reltol)
    else
        error("Unknown method")
    end

    solution
end

"""
    solve(
        equation::DiffusionEquation{PeriodicIntervalDomain{T},F},
        u,
        tlist,
        M,
        N;
        method = "filterfirst",
        solver = QNDF(),
        abstol = 1e-4,
        reltol = 1e-3,
        degmax = 10,
        λ = 1e-6
    ) where {T,F}

Solve `equation` from `tlist[1]` to `tlist[2]` with initial conditions `u` and a
discretization of `M` and `N` points. If `method` is `"filterfirst"`, the equation is filtered then
discretized. If `method` is `"discretizefirst"`, the equation is discretized then filtered.
"""
function solve(
    equation::DiffusionEquation{PeriodicIntervalDomain{T},F},
    u,
    tlist,
    M,
    N;
    method = "filterfirst",
    solver = QNDF(),
    abstol = 1e-4,
    reltol = 1e-3,
    degmax = 10,
    λ = 1e-6,
) where {T,F}
    @unpack domain, filter = equation

    # Domain
    x = discretize(domain, M)
    ξ = discretize(domain, N)

    # Get matrices
    D = diffusion_matrix(domain, N)
    W = filter_matrix(filter, domain, M, N; degmax, λ)
    R = reconstruction_matrix(filter, domain, M, N; degmax, λ)
    A = spdiagm(α.(ξ))

    # Initial conditions
    uₕ = u.(ξ)
    ūₕ = W * uₕ

    if method == "filterfirst"
        error("Not implemented")
        # Filter
        h = filter.width
        dh(ξ) = 0.0
        α(ξ) = 1 / 3 * dh(ξ) * h(ξ)
    elseif method == "discretizefirst"
        p = (; J = W * D * R)
        du!(du, u, p, t) = mul!(du, p.J, u)
        odefunction = ODEFunction(
            du!,
            jac = (J, u, p, t) -> (J .= p.J),
            jac_prototype = p.J,
            mass_matrix = W * R,
        )
        problem = ODEProblem(odefunction, ūₕ, tlist, p)
        solution = OrdinaryDiffEq.solve(problem, solver; abstol, reltol)
    else
        error("Unknown method")
    end

    solution
end


@inline fractional_x_index(x, ::Center, grid::RegularRectilinearGrid) = @inbounds (x - grid.xC[1]) / grid.Δx
@inline fractional_x_index(x, ::Face, grid::RegularRectilinearGrid) = @inbounds (x - grid.xF[1]) / grid.Δx

@inline fractional_y_index(y, ::Center, grid::RegularRectilinearGrid) = @inbounds (y - grid.yC[1]) / grid.Δy
@inline fractional_y_index(y, ::Face, grid::RegularRectilinearGrid) = @inbounds (y - grid.yF[1]) / grid.Δy

@inline fractional_z_index(z, ::Center, grid::RegularRectilinearGrid) = @inbounds (z - grid.zC[1]) / grid.Δz
@inline fractional_z_index(z, ::Face, grid::RegularRectilinearGrid) = @inbounds (z - grid.zF[1]) / grid.Δz

"""
    fractional_indices(x, y, z, loc, grid::RegularRectilinearGrid)

Convert the coordinates `(x, y, z)` to _fractional_ indices on a regular rectilinear grid located at `loc`
where `loc` is a 3-tuple of `Center` and `Face`. Fractional indices are floats indicating a location between
grid points.
"""
@inline function fractional_indices(x, y, z, loc, grid::RegularRectilinearGrid)
    i = fractional_x_index(x, loc[1], grid)
    j = fractional_y_index(y, loc[2], grid)
    k = fractional_z_index(z, loc[3], grid)
    return i, j, k
end

# Trilinear Lagrange polynomials
@inline ϕ₁(ξ, η, ζ) = (1 - ξ) * (1 - η) * (1 - ζ)
@inline ϕ₂(ξ, η, ζ) = (1 - ξ) * (1 - η) *      ζ
@inline ϕ₃(ξ, η, ζ) = (1 - ξ) *      η  * (1 - ζ)
@inline ϕ₄(ξ, η, ζ) = (1 - ξ) *      η  *      ζ
@inline ϕ₅(ξ, η, ζ) =      ξ  * (1 - η) * (1 - ζ)
@inline ϕ₆(ξ, η, ζ) =      ξ  * (1 - η) *      ζ
@inline ϕ₇(ξ, η, ζ) =      ξ  *      η  * (1 - ζ)
@inline ϕ₈(ξ, η, ζ) =      ξ  *      η  *      ζ

@inline _interpolate(field, ξ, η, ζ, i, j, k) =
    @inbounds (  ϕ₁(ξ, η, ζ) * field[i,   j,   k  ]
               + ϕ₂(ξ, η, ζ) * field[i,   j,   k+1]
               + ϕ₃(ξ, η, ζ) * field[i,   j+1, k  ]
               + ϕ₄(ξ, η, ζ) * field[i,   j+1, k+1]
               + ϕ₅(ξ, η, ζ) * field[i+1, j,   k  ]
               + ϕ₆(ξ, η, ζ) * field[i+1, j,   k+1]
               + ϕ₇(ξ, η, ζ) * field[i+1, j+1, k  ]
               + ϕ₈(ξ, η, ζ) * field[i+1, j+1, k+1])

"""
    interpolate(field, x, y, z)

Interpolate `field` to the physical point `(x, y, z)` using trilinear interpolation.
"""
@inline function interpolate(field, x, y, z)
    LX, LY, LZ = location(field)
    i, j, k = fractional_indices(x, y, z, (LX(), LY(), LZ()), field.grid)

    # Convert fractional indices to unit cell coordinates 0 <= (ξ, η, ζ) <=1
    # and integer indices (with 0-based indexing).
    ξ, i = modf(i)
    η, j = modf(j)
    ζ, k = modf(k)

    # Convert indices to proper integers and shift to 1-based indexing.
    return _interpolate(field.data, ξ, η, ζ, Int(i+1), Int(j+1), Int(k+1))
end

"""
    interpolate(field, LX, LY, LZ, grid, x, y, z)

Interpolate `field` to the physical point `(x, y, z)` using trilinear interpolation. The location of
the field is specified with `(LX, LY, LZ)` and the field is defined on `grid`.

Note that this is a lower-level `interpolate` method defined for use in CPU/GPU kernels.
"""
@inline function interpolate(field, LX, LY, LZ, grid, x, y, z)
    i, j, k = fractional_indices(x, y, z, (LX, LY, LZ), grid)

    # We use mod and trunc as CUDA.modf is not defined.
    # For why we use Base.unsafe_trunc instead of trunc see:
    # https://github.com/CliMA/Oceananigans.jl/issues/828
    # https://github.com/CliMA/Oceananigans.jl/pull/997
    ξ, i = mod(i, 1), Base.unsafe_trunc(Int, i)
    η, j = mod(j, 1), Base.unsafe_trunc(Int, j)
    ζ, k = mod(k, 1), Base.unsafe_trunc(Int, k)

    return _interpolate(field, ξ, η, ζ, i+1, j+1, k+1)
end


using DataFrames

srand(1234)

n,k = 1000,2
X = randn(n, k)
ϵ = randn(n)

β = [0.2, -1.0]
γ = [-0.4, 0.5]

ystar = X*β + ϵ
y = map((yi) -> searchsortedfirst(γ,yi), ystar)

df = convert(DataFrame, X)
df[:y] = y

writetable(joinpath(Pkg.dir("OrderedResponse")) * "/data/testdat.csv", df)


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM, LinearAlgebra, Test

# 2d heat problem (one element)

X = Dict(
    1 => [0.0,0.0],
    2 => [1.0,0.0],
    3 => [1.0,1.0],
    4 => [0.0,1.0])

# define volume element
element1 = Element(Quad4, (1, 2, 3, 4))

update!(element1, "geometry", X)
update!(element1, "thermal conductivity", 6.0)
update!(element1, "heat source", 12.0)

# define boundary element for flux
element2 = Element(Seg2, (1, 2))
update!(element2, "geometry", X)
# linear ramp from 0 -> 6 in time 0 -> 1
update!(element2, "heat flux", 0.0 => 0.0)
update!(element2, "heat flux", 1.0 => 6.0)

# define heat problem and add elements to problem
problem = Problem(PlaneHeat, "one element heat problem", 1)
add_elements!(problem, element1, element2)

# Set constant source f=12 with k=6. Accurate solution is
# T=1 on free boundary, u(x,y) = -1/6*(1/2*f*x^2 - f*x)
# when boundary flux not active (at t=0)
time = 0.0
assemble!(problem, time)
A = Matrix(problem.assembly.K)
b = Vector(problem.assembly.f)
A_expected = [
     4.0 -1.0 -2.0 -1.0
    -1.0  4.0 -1.0 -2.0
    -2.0 -1.0  4.0 -1.0
    -1.0 -2.0 -1.0  4.0]
free_dofs = [1, 2]
@test isapprox(A, A_expected)
@test isapprox(A[free_dofs, free_dofs] \ b[free_dofs], [1.0, 1.0])

# Set constant flux g=6 on boundary. Accurate solution is
# u(x,y) = x which equals T=1 on boundary.
# at time t=1.0 all loads should be on.
empty!(problem.assembly)
time = 1.0
assemble!(problem, time)
A = Matrix(problem.assembly.K)
b = Vector(problem.assembly.f)
@test isapprox(A[free_dofs, free_dofs] \ b[free_dofs], [2.0, 2.0])


using Revise
using Test: @testset, @test

include("Interface.jl")
include("ExampleFunctions.jl")

using .Interface: @ODE
using .StrangeAttractors: yuwang, TSUCS2, lorenz, TSUCS1

u = [[1., 2., 3.] [4., 5., 6.] [7., 8., 9.]]

lorenz(u)

@testset "Interface" begin
    @testset "Three-Dimensional Functions" begin

        @testset "Typing" begin
            u = [[1., 2., 3.] [4., 5., 6.] [7., 8., 9.]]

            @test lorenz(u) ≈ [[10., 10., 10.] [23., 83., 125.] [-6., 4., 32.]]
        end

        @testset "Time Functionality" begin
            u = [[1., 2., 3.] [4., 5., 6.] [7., 8., 9.]]

            time_ODE = Interface.@ODE quote
                dx = t
                dy = x + y
                dz = y + z
            end

            @test length(methods(time_ODE)) == 1

            @test time_ODE(0., u)[:,1] ≈ Vector{Float64}([0, 0, 0])
            @test time_ODE(2., u)[:,1] ≈ Vector{Float64}([2, 2, 2])
        end
    end

    @testset "Two-Dimensional Functions" begin

    end
end

lorenz([[1.,2.,3.] [3.,4.,5.] [6.,7., 8.]], σ = 10, ρ = 28, β = 8/3)
TSUCS2(0., [[1,2,3] [3,4,5] [6,7, 8]], α = 1, δ = 2, ς = 4, ζ = 8, β = 5, ϵ = 6)
yuwang(0., [[1,2,3] [3,4,5] [6,7, 8]], α = 1, β = 2, σ = 4, δ = 8)


using StructuralCausalModels, Test

d1 =OrderedDict(
  [:u, :w] => :v,
  [:u] => :x,
  :s1 => [:u],
  :w => :y,
  [:s2] => [:w]
)

e1 = StructuralCausalModels.edge_matrix(d1)

d2 = OrderedDict(
  :u => [:x, :v],
  :s1 => [:u],
  :w => [:v, :y],
  :s2 => [:w]
);

e2 = StructuralCausalModels.edge_matrix(d2)

@testset "dag_formulations" begin

  syms = names(e1, 1)
  for s1 in syms
    for s2 in syms
      @test e1[s1, s2] == e2[s1, s2]
    end
  end

end


using OrdinaryDiffEq
a=100.0
function df(du,u,p,t)
	du[1:end-1] .= a*(u[2:end] .- u[1:end-1])
	du[end] = -a*u[end]
end

u0 = [0.0,0.0,0.0,0.0,1.0]
tspan = (0,0.1)
prob = ODEProblem(df, u0, tspan)
sol = solve(prob, Tsit5(), saveat=0.001)

leg=["\$1^\\text{st}\$ tap", "\$2^\\text{nd}\$ tap", "\$3^\\text{rd}\$ tap", "\$4^\\text{th}\$ tap", "\$5^\\text{th}\$ tap"]

plts = []
for i in 1:5
	push!(plts, @pgf PlotInc(
			{ no_marks, very_thick },
			Table(; x = sol.t, y = sol[end+1-i,:])
		)
	)
	push!(plts, @pgf LegendEntry(leg[i]))
end

p = @pgf TikzPicture(
	Axis({
		height="4cm",
		width=raw"\textwidth",
		# "scale only axis",
		xmin=0,
		xmax=0.1,
		ymin=0,
		ymax=1,
		xlabel="time in seconds",
		ylabel="impulse response",
		"axis background/.style"={fill="none"},
		},
		plts...
	)
)


## AlmostBandedMatrix

immutable AlmostBandedMatrix{T} <: AbstractMatrix{T}
    bands::BandedMatrix{T}
    fill::LowRankMatrix{T}
    function AlmostBandedMatrix(bands::BandedMatrix{T},fill::LowRankMatrix{T})
        if size(bands) ≠ size(fill)
            error("Data and fill must be compatible size")
        end
        new(bands,fill)
    end
end

AlmostBandedMatrix(bands::BandedMatrix,fill::LowRankMatrix) =
    AlmostBandedMatrix{promote_type(eltype(bands),eltype(fill))}(bands,fill)

AlmostBandedMatrix{T}(::Type{T},n::Integer,m::Integer,l::Integer,u::Integer,r::Integer) =
    AlmostBandedMatrix(BandedMatrix(T,n,m,l,u),LowRankMatrix(T,n,m,r))

abzeros{T}(::Type{T},n::Integer,m::Integer,l::Integer,u::Integer,r::Integer) =
    AlmostBandedMatrix(bzeros(T,n,m,l,u),lrzeros(T,n,m,r))

Base.size(A::AlmostBandedMatrix) = size(A.bands)
Base.linearindexing{ABM<:AlmostBandedMatrix}(::Type{ABM}) =
    Base.LinearSlow()

function getindex(B::AlmostBandedMatrix,k::Integer,j::Integer)
    if j > k + bandwidth(B.bands,2)
        B.fill[k,j]
    else
        B.bands[k,j]
    end
end

# can only change the bands, not the fill
function Base.setindex!(B::AlmostBandedMatrix,v,k::Integer,j::Integer)
        B.bands[k,j] = v
end

function pad!(B::AlmostBandedMatrix,n::Integer,m::Integer)
    pad!(B.bands,n,m)
    pad!(B.fill,n,m)
    B
end


using HttpServer

function fibonacci(n)
  if n == 1 return 1 end
  if n == 2 return 1 end
  prev = BigInt(1)
  pprev = BigInt(1)
  for i=3:n
    curr = prev + pprev
    pprev = prev
    prev = curr
  end
  return prev
end

http = HttpHandler() do req::Request, res::Response
    m = match(r"^/fibo/(\d+)/?$",req.resource)
    if m == nothing return Response(404) end
    number = BigInt(m.captures[1])
    if number < 1 || number > 100_000 return Response(500) end
    return Response(string(fibonacci(number)))
end

http.events["error"]  = (client, err) -> println(err)
http.events["listen"] = (port)        -> println("Listening on $port...")

server = Server(http)
run(server, 8000)


#=
The following tests check the distribution of `rand(RotType)`.
To visualize the distribution, try the following script.

    using Rotations, Plots
    RotType = RotMatrix3{Float64}

    # Sampling
    n = 10000
    rs = rand(RotType,n)
    q = rand(RotType)
    angles = rotation_angle.([r/q for r in rs])/π

    # Cumulative distribution function
    cdf_sampled(t) = count(<(t), angles)/n
    cdf_true(t) = t-sinpi(t)/π

    # Probability density function and histogram of sampling distribution
    histogram(angles, normed=true, label=false)
    plot!(t->1-cospi(t), 0, 1, label=false)

    # Cumulative distribution function
    plot(cdf_sampled, 0, 1, label=false)
    plot!(cdf_true, 0, 1, label=false)
=#

@testset "Distribution" begin

    # TODO: consider to remove Euler rotations (#155)
    Type_SO3 =  (RotMatrix3, RotMatrix{3}, AngleAxis, RotationVec,
                QuatRotation, MRP, RodriguesParam,
                RotXYZ, RotYZX, RotZXY, RotXZY, RotYXZ, RotZYX,
                RotXYX, RotYZY, RotZXZ, RotXZX, RotYXY, RotZYZ)

    Type_SO2 = (RotMatrix2, RotMatrix{2}, Angle2d, RotX, RotY, RotZ)

    # Number of sampling
    N = 100000

    @testset "Distribution $(RotType)" for RotType in Type_SO3
        # Sampling
        rs = rand(RotType, N)
        q = rand(RotType)
        angles = rotation_angle.([r/q for r in rs])/π

        # Check sampled rotations are rotations
        @test all(isrotation.(rs))
        @test isrotation(q)

        # Cumulative distribution function
        cdf_sampled(t) = count(<(t), angles)/N
        cdf_true(t) = t-sinpi(t)/π

        # Check the CDFs are approximately equal
        ts = 0:0.01:1
        @test norm(cdf_true.(ts) - cdf_sampled.(ts), Inf) < 0.01
    end

    @testset "distribution $(RotType)" for RotType in Type_SO2
        # Sampling
        rs = rand(RotType, N)
        q = rand(RotType)
        angles = rotation_angle.([r/q for r in rs])/2π
        angles = mod.(angles, 1)

        # Check sampled rotations are rotations
        @test all(isrotation.(rs))
        @test isrotation(q)

        # Cumulative distribution function
        cdf_sampled(t) = count(<(t), angles)/N
        cdf_true(t) = t

        # Check the CDFs are approximately equal
        ts = 0:0.01:1
        @test norm(cdf_true.(ts) - cdf_sampled.(ts), Inf) < 0.01
    end

    @testset "distribution $(N)-dimensinal RotMatrix" for N in 2:5
        # Sampling
        RotType = RotMatrix{N}
        n = 1000000
        rs = rand(RotType, n)
        q = rand(RotType)
        norms1 = norm.([(r/q - one(SMatrix{N,N})) for r in rs])
        norms2 = norm.([(r - one(SMatrix{N,N})) for r in rs])

        # Check sampled rotations are rotations
        @test all(isrotation.(rs))
        @test isrotation(q)

        # Cumulative distribution function
        cdf_sampled1(t) = count(<(t), norms1)/n
        cdf_sampled2(t) = count(<(t), norms2)/n

        # Check the CDFs are approximately equal
        ts = 0:0.01:1
        @test norm(cdf_sampled1.(ts) - cdf_sampled2.(ts), Inf) < 0.01
    end
end