using Base.Threads
using LinearAlgebra

#local ttvec rank increase function with noise ϵ_wn
function tt_up_rks_noise(tt_vec,tt_ot_i,rkm,rk,ϵ_wn)
	vec_out = zeros(eltype(tt_vec),size(tt_vec,1),rkm,rk)
	vec_out[:,1:size(tt_vec,2),1:size(tt_vec,3)] = tt_vec
	if !iszero(ϵ_wn)
		if rkm == size(tt_vec,2) && rk>size(tt_vec,3)
			Q = rand_orthogonal(size(tt_vec,1)*rkm,rk-size(tt_vec,3))
			vec_out[:,:,size(tt_vec,3)+1:rk] = ϵ_wn*reshape(Q,size(tt_vec,1),rkm,rk-size(tt_vec,3))
			tt_ot_i =0
		elseif rk == size(tt_vec,3) && rkm>size(tt_vec,2)
			Q = rand_orthogonal(rkm-size(tt_vec,2),size(tt_vec,1)*rk)
			vec_out[:,size(tt_vec,2)+1:rkm,:] = ϵ_wn*reshape(Q,size(tt_vec,1),rkm-size(tt_vec,2),rk)
			tt_ot_i =0
		elseif rk>size(tt_vec,3) && rkm>size(tt_vec,2)
			Q = rand_orthogonal((rkm-size(tt_vec,2))*size(tt_vec,1),(rk-size(tt_vec,3)))
			vec_out[:,size(tt_vec,2)+1:rkm,size(tt_vec,3)+1:rk] = ϵ_wn*reshape(Q,size(tt_vec,1),rkm-size(tt_vec,2),rk-size(tt_vec,3))
		end
	end
	return vec_out
end

"""
returns the ttvector with ranks rks and noise ϵ_wn for the updated ranks
"""
function tt_up_rks(x_tt::ttvector{T},rk_max::Int;rks=vcat(1,rk_max*ones(Int,length(x_tt.ttv_dims)-1),1),ϵ_wn=0.0) where T<:Number
	d = length(x_tt.ttv_dims)
	vec_out = Array{Array{T}}(undef,d)
	out_ot = zeros(Int64,d)
	@assert(rk_max > maximum(x_tt.ttv_rks),"New bond dimension too low")
	n_in = 1
	n_out = prod(x_tt.ttv_dims)
	for i in 1:d
		n_in *= x_tt.ttv_dims[i]
		n_out = Int(n_out/x_tt.ttv_dims[i])
		rks[i+1] = min(rks[i+1],n_in,n_out)
		vec_out[i] = tt_up_rks_noise(x_tt.ttv_vec[i],x_tt.ttv_ot[i],rks[i],rks[i+1],ϵ_wn)
	end
	return ttvector{T}(vec_out,x_tt.ttv_dims,rks,x_tt.ttv_ot)
end

"""
returns the orthogonalized ttvector with root i
"""
function orthogonalize(x_tt::ttvector{T};i=1::Int) where T<:Number
	d = length(x_tt.ttv_dims)
	@assert(1≤i≤d, DimensionMismatch("Impossible orthogonalization"))
	y_tt = copy(x_tt)
	y_tt.ttv_ot[i]=0
	for j in 1:i-1
		y_tt.ttv_ot[j]=1
		y_vectemp = reshape(y_tt.ttv_vec[j],y_tt.ttv_dims[j]*y_tt.ttv_rks[j],y_tt.ttv_rks[j+1])
		F = qr(y_vectemp)
		y_tt.ttv_vec[j] = reshape(Matrix(F.Q)[:,1:y_tt.ttv_rks[j+1]],x_tt.ttv_dims[j],y_tt.ttv_rks[j],y_tt.ttv_rks[j+1])
		@threads for k in 1:x_tt.ttv_dims[j]
			y_tt.ttv_vec[j+1][k,:,:] = F.R[1:y_tt.ttv_rks[j+1],1:y_tt.ttv_rks[j+1]]*y_tt.ttv_vec[j+1][k,:,:]
		end
	end
	for j in d:-1:i+1
		y_tt.ttv_ot[j]=-1
		y_vectemp = reshape(permutedims(y_tt.ttv_vec[j],[2,1,3]),y_tt.ttv_rks[j],y_tt.ttv_dims[j]*y_tt.ttv_rks[j+1])
		F = lq(y_vectemp)
		y_tt.ttv_vec[j] = permutedims(reshape(Matrix(F.Q)[1:y_tt.ttv_rks[j],:],y_tt.ttv_rks[j],x_tt.ttv_dims[j],y_tt.ttv_rks[j+1]),[2 1 3])
		@threads for k in 1:x_tt.ttv_dims[j]
			y_tt.ttv_vec[j-1][k,:,:] = y_tt.ttv_vec[j-1][k,:,:]*F.L[1:y_tt.ttv_rks[j],1:y_tt.ttv_rks[j]]
		end
	end
	return y_tt
end

"""
returns a TT representation where the singular values lower than tol are discarded
"""
function tt_rounding(x_tt::ttvector{T};tol=1e-12) where T<:Number
	d = length(x_tt.ttv_dims)
	y_rks = copy(x_tt.ttv_rks)
	y_vec = copy(x_tt.ttv_vec)
	for j in 1:d-1
		A = zeros(T,x_tt.ttv_dims[j],y_rks[j],x_tt.ttv_dims[j+1],y_rks[j+2])
		@tensor A[a,b,c,d] = y_vec[j][a,b,z]*y_vec[j+1][c,z,d]
		u,s,v = svd(reshape(A,size(A,1)*size(A,2),:);alg=LinearAlgebra.QRIteration())
		rtol = norm(s)*tol
		Σ = s[s.>rtol]
		y_rks[j+1] = length(Σ)
		y_vec[j] = reshape(u[:,s.>rtol],x_tt.ttv_dims[j],y_rks[j],:)
		y_vec[j+1] = permutedims(reshape(Diagonal(Σ)*v'[s.>rtol,:],:,x_tt.ttv_dims[j+1],y_rks[j+2]),[2 1 3])
	end
	for j in d:-1:2
		A = zeros(T,x_tt.ttv_dims[j-1],y_rks[j-1],x_tt.ttv_dims[j],y_rks[j+1])
		@tensor A[a,b,c,d] = y_vec[j-1][a,b,z]*y_vec[j][c,z,d]
		F = svd(reshape(A,size(A,1)*size(A,2),:),alg=LinearAlgebra.QRIteration())
		rtol = norm(F.S)*tol
		Σ = F.S[F.S .>rtol]
		y_rks[j] = length(Σ)
		y_vec[j] = permutedims(reshape(F.Vt[F.S.>rtol,:],:,x_tt.ttv_dims[j],y_rks[j+1]),[2 1 3])
		y_vec[j-1] = reshape(F.U[:,F.S.>rtol]*Diagonal(Σ),x_tt.ttv_dims[j-1],y_rks[j-1],:)
	end
	return ttvector{T}(y_vec,x_tt.ttv_dims,y_rks,vcat(0,ones(Int64,d-1)))
end

"""
returns the rounding of the TT operator
"""
function tt_rounding(A_tto::ttoperator;tol=1e-12)
	return ttv_to_tto(tt_rounding(tto_to_ttv(A_tto);tol=tol))
end

"""
returns the singular values of the reshaped tensor x[μ_1⋯μ_k;μ_{k+1}⋯μ_d] for all 1≤ k ≤ d
"""
function tt_svdvals(x_tt::ttvector;tol=1e-14)
	d = length(x_tt.ttv_dims)
	Σ = Array{Array{Float64,1},1}(undef,d-1)
	y_tt = orthogonalize(x_tt)
	y_rks = y_tt.ttv_rks
	for j in 1:d-1
		A = zeros(y_tt.ttv_dims[j],y_rks[j],y_tt.ttv_dims[j+1],y_rks[j+2])
		@tensor A[a,b,c,d] = y_tt.ttv_vec[j][a,b,z]*y_tt.ttv_vec[j+1][c,z,d]
		u,s,v = svd(reshape(A,size(A,1)*size(A,2),:),alg=LinearAlgebra.QRIteration())
		Σ[j] = s[s.>tol]
		y_rks[j+1] = length(Σ[j])
		y_tt.ttv_vec[j+1] = permutedims(reshape(Diagonal(Σ[j])*v'[s.>tol,:],:,y_tt.ttv_dims[j+1],y_rks[j+2]),[2 1 3])
	end
	return Σ
end

function sv_trunc(s::Array{Float64},tol)
	if tol==0.0
		return s
	else
		d = length(s)
		i=0
		weight = 0.0
		norm2 = dot(s,s)
		while (i<d) && weight<tol*norm2
			weight+=s[d-i]^2
			i+=1
		end
		return s[1:(d-i+1)]
	end
end

function left_compression(A,B;tol=1e-12)
    dim_A = [i for i in size(A)]
    dim_B = [i for i in size(B)]

    B = permutedims(B, [2,1,3]) #B r_1 x n_2 x r_2
    U = reshape(A,:,dim_A[3])
    u,s,v = svd(U*reshape(B, dim_B[2],:),full=false) #u is the new A, dim(u) = n_1r_0 x tilde(r)_1
    s_trunc = sv_trunc(s,tol)
    dim_B[2] = length(s_trunc)
    U = reshape(u[:,1:dim_B[2]],dim_A[1],dim_A[2],dim_B[2])
    B = reshape(Diagonal(s_trunc)*v[:,1:dim_B[2]]',dim_B[2],dim_B[1],dim_B[3])
    return U, permutedims(B,[2,1,3])
end

"""
parallel compression of the ttvector
TODO refactoring
"""
function tt_compression_par(X::ttvector;tol=1e-14,Imax=2)
    Y = deepcopy(X.ttv_vec) :: Array{Array{Float64,3},1}
    rks = deepcopy(X.ttv_rks) :: Array{Int64}
    d = length(X.ttv_dims)
    rks_prev = zeros(Integer,d)
    i=0
    while norm(rks-rks_prev)>0.1 && i<Imax
        i+=1
        rks_prev = deepcopy(rks) :: Array{Int64}
        if mod(i,2) == 1
            @threads for k in 1:floor(Integer,d/2)
                Y[2k-1], Y[2k] = left_compression(Y[2k-1], Y[2k], tol=tol)
                rks[2k-1] = size(Y[2k-1],3)
            end
        else
            @threads for k in 1:floor(Integer,(d-1)/2)
                Y[2k], Y[2k+1] = left_compression(Y[2k], Y[2k+1], tol=tol)
                rks[2k] = size(Y[2k],3)
            end
        end
    end
    return ttvector(Y,X.ttv_dims,rks,zeros(Integer,d))
end

function tt_compression_par(A::ttoperator;tol=1e-14,Imax=2)
	return ttv_to_tto(tt_compression_par(tto_to_ttv(A);tol=tol,Imax=Imax))
end


using DiffEqSensitivity, OrdinaryDiffEq, ForwardDiff, Calculus
using Test
function fb(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -t*p[3]*u[2] + t*u[1]*u[2]
end
function jac(J,u,p,t)
  (x, y, a, b, c) = (u[1], u[2], p[1], p[2], p[3])
  J[1,1] = a + y * b * -1
  J[2,1] = t * y
  J[1,2] = b * x * -1
  J[2,2] = t * c * -1 + t * x
end
function paramjac(pJ,u,p,t)
  (x, y, a, b, c) = (u[1], u[2], p[1], p[2], p[3])
  pJ[1,1] = x
  pJ[2,1] = 0.0
  pJ[1,2] = - x * y
  pJ[2,2] = 0.0
  pJ[1,3] = 0.0
  pJ[2,3] = - t * y
end

f = ODEFunction(fb,jac=jac,paramjac=paramjac)
p = [1.5,1.0,3.0]
prob = ODEForwardSensitivityProblem(f,[1.0;1.0],(0.0,10.0),p)
probInpl = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p)
probnoad = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p,
                                        ForwardSensitivity(autodiff=false))
probnoadjacvec = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p,
                                        ForwardSensitivity(autodiff=false,autojacvec=true))
probnoad2 = ODEForwardSensitivityProblem(f,[1.0;1.0],(0.0,10.0),p,
                                        ForwardSensitivity(autodiff=false))
probvecmat = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p,
                                          ForwardSensitivity(autojacvec=false,autojacmat=true))
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
@test_broken solve(probInpl,KenCarp4(),abstol=1e-14,reltol=1e-14).retcode == :Success
solInpl = solve(probInpl,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solInpl2 = solve(probInpl,Rodas4(autodiff=false),abstol=1e-10,reltol=1e-10)
solnoad = solve(probnoad,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solnoadjacvec = solve(probnoadjacvec,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solnoad2 = solve(probnoad,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solvecmat = solve(probvecmat,Tsit5(),abstol=1e-14,reltol=1e-14)

x = sol[1:sol.prob.f.numindvar,:]

@test sol(5.0) ≈ solnoad(5.0)
@test sol(5.0) ≈ solnoad2(5.0)
@test sol(5.0) ≈ solnoadjacvec(5.0) atol=1e-6 rtol=1e-6
@test sol(5.0) ≈ solInpl(5.0)
@test isapprox(solInpl(5.0), solInpl2(5.0),rtol=1e-5)
@test sol(5.0) ≈ solvecmat(5.0)

# Get the sensitivities

da = sol[sol.prob.f.numindvar+1:sol.prob.f.numindvar*2,:]
db = sol[sol.prob.f.numindvar*2+1:sol.prob.f.numindvar*3,:]
dc = sol[sol.prob.f.numindvar*3+1:sol.prob.f.numindvar*4,:]

sense_res1 = [da[:,end] db[:,end] dc[:,end]]

prob = ODEForwardSensitivityProblem(f.f,[1.0;1.0],(0.0,10.0),p,
                                    ForwardSensitivity(autojacvec=true))
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,saveat=0.01)
x = sol[1:sol.prob.f.numindvar,:]

# Get the sensitivities

res = sol[1:sol.prob.f.numindvar,:]
da = sol[sol.prob.f.numindvar+1:sol.prob.f.numindvar*2,:]
db = sol[sol.prob.f.numindvar*2+1:sol.prob.f.numindvar*3,:]
dc = sol[sol.prob.f.numindvar*3+1:sol.prob.f.numindvar*4,:]

sense_res2 = [da[:,end] db[:,end] dc[:,end]]

function test_f(p)
  prob = ODEProblem(f,eltype(p).([1.0,1.0]),(0.0,10.0),p)
  solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,save_everystep=false)[end]
end

p = [1.5,1.0,3.0]
fd_res = ForwardDiff.jacobian(test_f,p)
calc_res = Calculus.finite_difference_jacobian(test_f,p)

@test sense_res1 ≈ sense_res2 ≈ fd_res
@test sense_res1 ≈ sense_res2 ≈ calc_res

## Check extraction

xall, dpall = extract_local_sensitivities(sol)
@test xall == res
@test dpall[1] == da

_,dpall_matrix = extract_local_sensitivities(sol,Val(true))
@test mapreduce(x->x[:, 2], hcat, dpall) == dpall_matrix[2]

x, dp = extract_local_sensitivities(sol,length(sol.t))
sense_res2 = reduce(hcat,dp)
@test sense_res1 == sense_res2

@test extract_local_sensitivities(sol,sol.t[3]) == extract_local_sensitivities(sol,3)

tmp = similar(sol[1])
@test extract_local_sensitivities(tmp,sol,sol.t[3]) == extract_local_sensitivities(sol,3)

# asmatrix=true
@test extract_local_sensitivities(sol, length(sol), true) == (x, sense_res2)
@test extract_local_sensitivities(sol, sol.t[end], true) == (x, sense_res2)
@test extract_local_sensitivities(tmp, sol, sol.t[end], true) == (x, sense_res2)

# Return type inferred
@inferred extract_local_sensitivities(sol, 1)
@inferred extract_local_sensitivities(sol, 1, Val(true))
@inferred extract_local_sensitivities(sol, sol.t[3])
@inferred extract_local_sensitivities(sol, sol.t[3], Val(true))
@inferred extract_local_sensitivities(tmp, sol, sol.t[3])
@inferred extract_local_sensitivities(tmp, sol, sol.t[3], Val(true))

### ForwardDiff version

prob = ODEForwardSensitivityProblem(f.f,[1.0;1.0],(0.0,10.0),p,
                                    ForwardDiffSensitivity())
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,saveat=0.01)

xall, dpall = extract_local_sensitivities(sol)
@test xall ≈ res
@test dpall[1] ≈ da atol=1e-9

_,dpall_matrix = extract_local_sensitivities(sol,Val(true))
@test mapreduce(x->x[:, 2], hcat, dpall) == dpall_matrix[2]

x, dp = extract_local_sensitivities(sol,length(sol.t))
sense_res2 = reduce(hcat,dp)
@test fd_res == sense_res2

@test extract_local_sensitivities(sol,sol.t[3]) == extract_local_sensitivities(sol,3)

tmp = similar(sol[1])
@test extract_local_sensitivities(tmp,sol,sol.t[3]) == extract_local_sensitivities(sol,3)

# asmatrix=true
@test extract_local_sensitivities(sol, length(sol), true) == (x, sense_res2)
@test extract_local_sensitivities(sol, sol.t[end], true) == (x, sense_res2)
@test extract_local_sensitivities(tmp, sol, sol.t[end], true) == (x, sense_res2)

# Return type inferred
@inferred extract_local_sensitivities(sol, 1)
@inferred extract_local_sensitivities(sol, 1, Val(true))
@inferred extract_local_sensitivities(sol, sol.t[3])
@inferred extract_local_sensitivities(sol, sol.t[3], Val(true))
@inferred extract_local_sensitivities(tmp, sol, sol.t[3])
@inferred extract_local_sensitivities(tmp, sol, sol.t[3], Val(true))

# Test mass matrix
function rober_MM(du, u, p, t)
    y₁, y₂, y₃ = u
    k₁, k₂, k₃ = p
    du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
    du[2] = k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃
    du[3] = y₁ + y₂ + y₃ - 1
    nothing
end
function rober_no_MM(du, u, p, t)
    y₁, y₂, y₃ = u
    k₁, k₂, k₃ = p
    du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
    du[2] = k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃
    du[3] = k₂*y₂^2
    nothing
end

M = [1.0 0 0; 0 1.0 0; 0 0 0]
p = [0.04, 3e7, 1e4]
u0 = [1.0, 0.0, 0.0]
tspan = (0.0, 12.0)

f_MM= ODEFunction(rober_MM, mass_matrix = M)
f_no_MM= ODEFunction(rober_no_MM)

prob_MM_ForwardSensitivity = ODEForwardSensitivityProblem(f_MM, u0, tspan, p, ForwardSensitivity())
sol_MM_ForwardSensitivity = solve(prob_MM_ForwardSensitivity , Rodas4(autodiff = false), reltol = 1e-14, abstol = 1e-14)

prob_MM_ForwardDiffSensitivity = ODEForwardSensitivityProblem(f_MM, u0, tspan, p, ForwardDiffSensitivity())
sol_MM_ForwardDiffSensitivity = solve(prob_MM_ForwardDiffSensitivity, Rodas4(autodiff = false), reltol = 1e-14, abstol = 1e-14)

prob_no_MM = ODEForwardSensitivityProblem(f_no_MM, u0, tspan, p, ForwardSensitivity())
sol_no_MM= solve(prob_no_MM, Rodas4(autodiff = false), reltol = 1e-14, abstol = 1e-14)

sen_MM_ForwardSensitivity = extract_local_sensitivities(sol_MM_ForwardSensitivity,10.0,true)
sen_MM_ForwardDiffSensitivity = extract_local_sensitivities(sol_MM_ForwardDiffSensitivity,10.0,true)
sen_no_MM = extract_local_sensitivities(sol_no_MM,10.0,true)

@test sen_MM_ForwardSensitivity[2] ≈ sen_MM_ForwardDiffSensitivity[2] atol=1e-10 rtol=1e-10
@test sen_MM_ForwardSensitivity[2] ≈ sen_no_MM[2] atol=1e-10 rtol=1e-10
# Test Float32

function f32(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + u[1]*u[2]
end
p = [1.5f0,1.0f0,3.0f0]
prob = ODEForwardSensitivityProblem(f32,[1.0f0;1.0f0],(0.0f0,10.0f0),p)
sol = solve(prob,Tsit5())


"""
    Line(z,a)

Line in the complex plane.

#### Arguments
- `z::Number`: Base point belonging to the line.
- `a::Real`: Angle (in radians) between the line and the X axis.
"""
struct Line{N <: Number, R <: Real} <: AbstractLine
  base::N
  a::R

  function Line{N,R}(z0::N, a0::R) where {N <: Number, R <: Real}
    new(z0,a0)
  end
end

Line(z0::N, a0::R) where {N <: Number, R <: Real} = Line{N,R}(z0,a0)

basepoint(l::Line) = l.base
basepointC(l::Line) = l.base
basepointR2(l::Line) = real(l.base), imag(l.base)
lineangle(l::Line) = l.a
normal(l::Line) = complex(cos(l.a), sin(l.a))
normalC(l::Line) = complex(cos(l.a), sin(l.a))
normalR2(l::Line) = cos(l.a), sin(l.a)

Base.show(io::IO, l::Line{N,R}) where {N,R} =
print(io, "Line{$N,$R}: base point ", l.base, ", angle with X ", l.a)

"""
    Ray(z,a)

A straight line ray.

#### Arguments
- `z::Number`: Base point belonging to the ray.
- `a::Real`: Angle (in radians) between the ray and the X axis.
"""
struct Ray{N <: Number, R <: Real} <: AbstractRay
  base::N
  a::R

  function Ray{N,R}(z0::N, a0::R) where {N <: Number, R <: Real}
    new(z0,a0)
  end
end

Ray(z0::N, a0::R) where {N <: Number, R <: Real} = Ray{N,R}(z0,a0)

basepoint(l::Ray) = l.base
basepointC(l::Ray) = l.base
basepointR2(l::Ray) = real(l.base), imag(l.base)
lineangle(l::Ray) = l.a
normal(l::Ray) = complex(cos(l.a), sin(l.a))
normalC(l::Ray) = complex(cos(l.a), sin(l.a))
normalR2(l::Ray) = cos(l.a), sin(l.a)

Base.show(io::IO, l::Ray{N,R}) where {N,R} =
print(io, "Ray{$N,$R}: base point", l.base, ", angle with X ", l.a)

"""
    BiRay(z1,z2,a)

A straight line bi-ray: the ray from \$z_1\$ to \$\\infty\$ union
the ray from \$z_2\$ to \$\\infty\$, with complementary angles.
"""
struct BiRay{N <: Number, R <: Real} <: AbstractLinearCurve
  z1::N
  z2::N
  a::R

  function BiRay{N,R}(z10::N, z20::N, a0::R) where {N <: Number, R <: Real}
    new(z10,z20,a0)
  end
end

BiRay(z10::N, z20::N, a0::R) where {N <: Number, R <: Real} = BiRay{N,R}(z10,z20,a0)
BiRay(z10::Number, z20::Number, a0::Real) = BiRay(promote(z10,z20)...,a0)

initialpoint(l::BiRay) = l.z1
finalpoint(l::BiRay) = l.z2
lineangle(l::BiRay) = l.a
normal(l::BiRay) = cos(l.a), sin(l.a)

torays(br::BiRay) = Ray(initialpoint(br),lineangle(br)), Ray(finalpoint(br),lineangle(br)+π)

Base.show(io::IO, l::BiRay{N,R}) where {N,R} =
  print(io, "BiRay{$N,$R}: initial point", l.z1, ", final point", l.z2, ", angle with X ", l.a)

"""
    LineSegment(z1,z2)

A straight line segment.
"""
struct LineSegment{N <: Number} <: AbstractLineSegment
  z1::N
  z2::N

  function LineSegment{N}(z10::N, z20::N) where N <: Number
    new(z10,z20)
  end
end

LineSegment(z10::N, z20::N) where {N <: Number} = LineSegment{N}(z10,z20)
LineSegment(z10::Number, z20::Number) = LineSegment(promote(z10,z20)...)

initialpoint(l::LineSegment) = l.z1
finalpoint(l::LineSegment) = l.z2
initialpointC(l::LineSegment) = l.z1
finalpointC(l::LineSegment) = l.z2
initialpointR2(l::LineSegment) = real(l.z1), imag(l.z1)
finalpointR2(l::LineSegment) = real(l.z2), imag(l.z2)
lineangle(l::LineSegment) = angle(z1-z2)
normal(l::LineSegment) = perp(z1-z2)
normalC(l::LineSegment) = perp(z1-z2)
#normalR2(l::LineSegment) = perp(z1-z2)

Base.show(io::IO, l::LineSegment{N}) where {N} =
print(io, "LineSegment{$N}: from ", l.z1, " to ", l.z2)

#=
ToDO!

LineR2
RayR2
ByRayR2
LineSegmentR2
=#

linepoints(x1::Real, y1::Real, x2::Real, y2::Real, numpts::Integer=100) =
[(x1+t*(x2-x1)/numpts, y1+t*(y2-y1)/numpts) for t ∈ 0:numpts]

linecomplexes(z1::Number, z2::Number, numpts::Integer=100) =
[z1+t*(z2-z1)/numpts for t ∈ 0:numpts]

function linepoints(l::AbstractLineSegment, numpts::Integer=100)
  x1, y1 = initialpointR2(l)
  x2, y2 = finalpointR2(l)
  [(x1+t*(x2-x1)/numpts, y1+t*(y2-y1)/numpts) for t ∈ 0:numpts]
end

function linecomplexes(l::AbstractLineSegment, numpts::Integer=100)
  z1 = initialpointC(l)
  z2 = finalpointC(l)
  [z1+t*(z2-z1)/numpts for t ∈ 0:numpts]
end

"""
Checks if `z` is inside the semi-plane determined by the line `l` with
  \$0 > normal \\cdot (z - base)\$.
"""
isinside(z::Number, l::Line) = 0 > dot(normal(l), z - l.base)

"""
Checks if `z` is outside the semi-plane determined by the line `l` with
  \$0 < normal \\cdot (z - base)\$.
"""
isoutside(z::Number, c::Line) = 0 < dot(normal(l), z - l.base)


#=
findpeaks.jl
Julia translation from Matlab code "findPeaks4.m" from
https://sethares.engr.wisc.edu/vocoders/matlabphasevocoder.html

The comments in that file say that the Matlab version was modified
from findPeaks.m by P. Moller-Nielson 28.3.03, pm-n.
See http://www.daimi.au.dk/~pmn/sound
=#

"""
    findpeaks(amp, max_peak::Int = 50, eps_peak::Real = 0.005)
Given a real-valued vector of amplitudes `amp`,
find up to `max_peak` peaks.
Ignores any amplitudes `≤ eps_peak * maxium(amp)`.

Returns a matrix with 3 columns where each row describes a peak.
"""
function findpeaks(
    amp::AbstractVector{<:Real},
    max_peak::Int = 50, # max number of peaks
    eps_peak::Real = 0.005, # minimum relative height of peaks
)

    N = length(amp)
    n = 3:(N-1)
    peakamp = amp[n] .* (amp[n] .> amp[n .- 1]) .* (amp[n] .> amp[n .+ 1])
    peakpos = zeros(Int, max_peak)

    maxamp = maximum(peakamp)
    npeak = 0
    for p in 1:max_peak
        (m, b) = findmax(peakamp)
        if m ≤ (eps_peak * maxamp)
            break
        end
        peakpos[p] = b + 2
        peakamp[b] = 0
        npeak = p
    end

    peakpos = sort(peakpos)
    peaks = zeros(Int, npeak, 3)

    last_b = 1
    for p in 1:npeak
        b = peakpos[max_peak - npeak + p]
        first_b = last_b + 1
        if p == npeak
            last_b = N
        else
            next_b = peakpos[max_peak - npeak + p + 1]
            rel_min = argmin(amp[b:next_b])
            last_b = b + rel_min - 1
        end
        peaks[p,:] = [first_b, b, last_b]
    end

    return peaks
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

# XXX: trailsign keyword not implemented, not sure it's worth the effort.
# Possible strategy: hold the sign of number in a variable (you can use
# copysign(1, number)), set result[1] equal to dd and multiply the first
# non-zero element (use findfirst for that) by the sign of number.
function _sixty{T<:AbstractFloat}(number::T)
    dd = trunc(abs(number))
    mm = abs(60.0*number)
    ss = abs(3600.0*number)
    result = zeros(typeof(number), 3)
    result[1] = trunc(number)
    result[2] = trunc(mm - 60.0*dd)
    result[3] = ss - 3600.0*dd - 60.0*result[2]
    return result
end

"""
    sixty(number) -> [deg, min, sec]

### Purpose ###

Converts a decimal number to sexagesimal.

### Explanation ###

The reverse of `ten` function.

### Argument ###

* `number`: decimal number to be converted to sexagesimal.

### Output ###

An array of three `AbstractFloat`, that are the sexagesimal counterpart
(degrees, minutes, seconds) of `number`.

### Example ###

``` julia
sixty(-0.615)
# => 3-element Array{Float64,1}:
#     -0.0
#     36.0
#     54.0
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
sixty(n::Real) = _sixty(float(n))


import PotentialFlow.Utils: @get, MappedVector

@testset "Utils" begin
    @testset "get" begin
        z = rand(ComplexF64)

        @get z (re, im) (r, i)
        @test i == imag(z)
        @test r == real(z)

    end

    @testset "MappedVector" begin
        x = [π, 0.0, π]
        y = MappedVector(cos, x, 1)
        @test y[0] == -1.0
        @test y[1] == 1.0
        @test y[2] == -1.0

        buff = IOBuffer()
        show(buff, y)
        @test String(take!(buff)) == "Vector{Float64} → typeof(cos) (0:2)"
        #@test String(take!(buff)) == "Array{Float64,1} → typeof(cos) (0:2)"

    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    TopologicalStructure

A data structure for constructing topological relations
in a [`Mesh`](@ref).

## References

* Floriani, L. & Hui, A. 2007. [Shape representations
  based on simplicial and cell complexes]
  (https://diglib.eg.org/handle/10.2312/egst.20071055.063-087)
"""
abstract type TopologicalStructure end

"""
    boundary(connectivity, rank, structure)

Boundary relation that maps the `connectivity` with a given
rank to a set of connectivities of smaller `rank` based on
the topological `structure`.

For example, the boundary of a `Connectivity{Triangle}` is a
set with three `Connectivity{Segment}`.
"""
boundary(c::Connectivity, rank::Integer, structure) =
  boundary(c, Val(rank), structure)

"""
    coboundary(connectivity, rank, structure)

Co-boundary relation that maps the `connectivity` with a given
rank to a set of connectivities of higher `rank` based on the
topological `structure`.

For example, the coboundary of a `Connectivity{Segment}` can
be a set with two `Connectivity{Triangle}`.
"""
coboundary(c::Connectivity, rank::Integer, structure) =
  coboundary(c, Val(rank), structure)

"""
    adjacency(connectivity, structure)

Adjacency relation that maps the `connectivity` with rank `p`
to a set of connectivities sharing a `p-1` connectivity.

For example, the adjacency of a `Connectivity{Triangle}` can
be a set of `Connectivity{Triangle}` sharing a `Connectivity{Segment}`.
"""
function adjacency end

# ----------------
# IMPLEMENTATIONS
# ----------------

include("topostructures/explicit.jl")
include("topostructures/halfedge.jl")

# conversions between structures
include("topostructures/conversions.jl")


struct Card
    suit::Int64
    rank::Int64
    function Card(suit::Int64, rank::Int64)
        @assert(1 <= suit <= 4, "suite is not between 1 and 4")
        @assert(1 <= rank <= 13, "rank is not between 1 and 13")
        new(suit, rank)
    end
end

# Global Variables
const suit_names = ["♣", "♦", "♥", "♠"]
const rank_names = ["A", "2", "3", "4", "5", "6", "7", "8", "9", "10", "J", "Q", "K"]

function Base.show(io::IO, card::Card)
    print(io, rank_names[card.rank], suit_names[card.suit])
end

import Base.< # Must use import
# this won't work:
# function base.<(...)
# end
function <(c1::Card, c2::Card)
    # suit is compared first
    # than rank
    (c1.suit, c1.rank) < (c2.suit, c2.rank)
end

# Abstract Types and Subtyping
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_abstract_types_and_subtyping
abstract type CardSet # <: inheritance
end

struct Deck <: CardSet
    cards :: Array{Card, 1} # 1-D array of Card
end

function Deck()
    deck = Deck(Card[])
    for suit = 1:4
        for rank = 1:13
            push!(deck.cards, Card(suit, rank))
        end
    end

    deck
end

function Base.show(io::IO, cs::CardSet)
    for card in cs.cards
        print(io, card, " ")
    end
    println()
end

function Base.pop!(cs::CardSet)
    pop!(cs.cards)
end

function Base.push!(cs::CardSet, card::Card)
    push!(cs.cards, card)
    nothing
end

import Random: shuffle!
function shuffle!(deck::Deck)
    shuffle!(deck.cards)
    deck
end

# Exercise 18-2
function Base.isless(x::Card, y::Card)
    x < y
end

function Base.sort!(deck::Deck)
    sort!(deck.cards)
    deck
end

# Abstract Types and Subtyping
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_abstract_types_and_subtyping
struct Hand <: CardSet
    cards :: Array{Card, 1}
    label :: String
end

function Hand(label::String="")
    Hand(Card[], label)
end

# Abstract Types and Functions
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_abstract_types_and_functions
"""
    move!(source::CardSet, target::CardSet, n::Int)

Move `n` number of cards from `source` to `target`.
"""
function move!(source::CardSet, target::CardSet, n::Int)
    @assert 1<= n <= length(source.cards)

    for i = 1:n
        card = pop!(source)
        push!(target, card)
    end

    nothing
end


typeof(2 + 3) #> Int64
:(2 + 3) #> :(2 + 3)
typeof(:(2 + 3)) #> Expr
exp = :(2 + 3)
names(exp)
#3-element Array{Symbol,1}:
# :head
# :args
# :typ
exp.head #> :call
exp.args #> 3-element Array{Any,1}:  :+ 2 3
exp.typ #> Any
2 + 3 == + (2, 3) #> true

exp2 = :(2 + a * b - c)
exp2.args #> 3-element Array{Any,1}: :- :(2 + a * b) :c
:(2 + a * b).args #> 3-element Array{Any,1}:  :+ 2  :(a * b)
:(a * b).args #> 3-element Array{Any,1}: :* :a :b

x = 5 #> 5
:x #> :x

quote
   a = 42
   b = a^2
   a - b
end
# quote  # none, line 2:
#     a = 42 # line 3:
#     b = a^2 # line 4:
#     a - b
# end
:(a = 42; b = a^2; a - b)
# returns:
# quote
#     a = 42
#     begin
#         b = a^2
#         a - b
#     end
# end

type Expr
  head::Symbol
  args::Array{Any,1}
  typ
end

dump(:(2 + 2))
# Expr
#   head: Symbol call
#   args: Array(Any,(3,))
#     1: Symbol +
#     2: Int64 2
#     3: Int64 2
#   typ: Any

dump(:(2 + a * b - c))
# Expr
#   head: Symbol call
#   args: Array(Any,(3,))
#     1: Symbol -
#     2: Expr
#       head: Symbol call
#       args: Array(Any,(3,))
#         1: Symbol +
#         2: Int64 2
#         3: Expr
#           head: Symbol call
#           args: Array(Any,(3,))
#           typ: Any
#       typ: Any
#     3: Symbol c
#   typ: Any

symbol('\'') #> :'
symbol("'") #> :'


# middle

@test middle(3) === 3.0
@test middle(2, 3) === 2.5
@test middle(1:8) === 4.5
@test middle([1:8]) === 4.5

# ensure type-correctness
for T in [Bool,Int8,Int16,Int32,Int64,Int128,Uint8,Uint16,Uint32,Uint64,Uint128,Float16,Float32,Float64]
    @test middle(one(T)) === middle(one(T), one(T))
end

# median
@test median([1.]) === 1.
@test median([1.,3]) === 2.
@test median([1.,3,2]) === 2.

@test median([1,3,2]) === 2.0
@test median([1,3,2,4]) === 2.5

@test median([0.0,Inf]) == Inf
@test median([0.0,-Inf]) == -Inf
@test median([0.,Inf,-Inf]) == 0.0
@test median([1.,-1.,Inf,-Inf]) == 0.0
@test isnan(median([-Inf,Inf]))

@test all(median([2 3 1 -1; 7 4 5 -4], 2) .== [1.5, 4.5])
@test all(median([2 3 1 -1; 7 4 5 -4], 1) .== [4.5 3.5 3.0 -2.5])

@test_throws ErrorException median([])
@test isnan(median([NaN]))
@test isnan(median([0.0,NaN]))
@test isnan(median([NaN,0.0]))
@test isequal(median([NaN 0.0; 1.2 4.5], 2), reshape([NaN; 2.85], 2, 1))

@test mean([1,2,3]) == 2.
@test mean([0 1 2; 4 5 6], 1) == [2.  3.  4.]
@test mean([1 2 3; 4 5 6], 1) == [2.5 3.5 4.5]

# test var & std

# edge case: empty vector
# iterable; this has to throw for type stability
@test_throws ErrorException var(())
@test_throws ErrorException var((); corrected=false)
@test_throws ErrorException var((); mean=2)
@test_throws ErrorException var((); mean=2, corrected=false)
# reduction
@test isnan(var(Int[]))
@test isnan(var(Int[]; corrected=false))
@test isnan(var(Int[]; mean=2))
@test isnan(var(Int[]; mean=2, corrected=false))
# reduction across dimensions
@test_approx_eq var(Int[], 1) [NaN]
@test_approx_eq var(Int[], 1; corrected=false) [NaN]
@test_approx_eq var(Int[], 1; mean=[2]) [NaN]
@test_approx_eq var(Int[], 1; mean=[2], corrected=false) [NaN]

# edge case: one-element vector
# iterable
@test isnan(@inferred(var((1,))))
@test var((1,); corrected=false) === 0.0
@test var((1,); mean=2) === Inf
@test var((1,); mean=2, corrected=false) === 1.0
# reduction
@test isnan(@inferred(var([1])))
@test var([1]; corrected=false) === 0.0
@test var([1]; mean=2) === Inf
@test var([1]; mean=2, corrected=false) === 1.0
# reduction across dimensions
@test_approx_eq @inferred(var([1], 1)) [NaN]
@test_approx_eq var([1], 1; corrected=false) [0.0]
@test_approx_eq var([1], 1; mean=[2]) [Inf]
@test_approx_eq var([1], 1; mean=[2], corrected=false) [1.0]

@test var(1:8) == 6.

@test_approx_eq varm([1,2,3], 2) 1.
@test_approx_eq var([1,2,3]) 1.
@test_approx_eq var([1,2,3]; corrected=false) 2.0/3
@test_approx_eq var([1,2,3]; mean=0) 7.
@test_approx_eq var([1,2,3]; mean=0, corrected=false) 14.0/3

@test_approx_eq varm((1,2,3), 2) 1.
@test_approx_eq var((1,2,3)) 1.
@test_approx_eq var((1,2,3); corrected=false) 2.0/3
@test_approx_eq var((1,2,3); mean=0) 7.
@test_approx_eq var((1,2,3); mean=0, corrected=false) 14.0/3

@test_approx_eq var([1 2 3 4 5; 6 7 8 9 10], 2) [2.5 2.5]'
@test_approx_eq var([1 2 3 4 5; 6 7 8 9 10], 2; corrected=false) [2.0 2.0]'

@test_approx_eq stdm([1,2,3], 2) 1.
@test_approx_eq std([1,2,3]) 1.
@test_approx_eq std([1,2,3]; corrected=false) sqrt(2.0/3)
@test_approx_eq std([1,2,3]; mean=0) sqrt(7.0)
@test_approx_eq std([1,2,3]; mean=0, corrected=false) sqrt(14.0/3)

@test_approx_eq stdm((1,2,3), 2) 1.
@test_approx_eq std((1,2,3)) 1.
@test_approx_eq std((1,2,3); corrected=false) sqrt(2.0/3)
@test_approx_eq std((1,2,3); mean=0) sqrt(7.0)
@test_approx_eq std((1,2,3); mean=0, corrected=false) sqrt(14.0/3)

@test_approx_eq std([1 2 3 4 5; 6 7 8 9 10], 2) sqrt([2.5 2.5]')
@test_approx_eq std([1 2 3 4 5; 6 7 8 9 10], 2; corrected=false) sqrt([2.0 2.0]')

A = Complex128[exp(i*im) for i in 1:10^4]
@test_approx_eq varm(A,0.) sum(map(abs2,A))/(length(A)-1)
@test_approx_eq varm(A,mean(A)) var(A)

# test covariance

function safe_cov(x, y, zm::Bool, cr::Bool)
    n = length(x)
    if !zm
        x = x .- mean(x)
        y = y .- mean(y)
    end
    dot(vec(x), vec(y)) / (n - int(cr))
end

X = [1. 2. 3. 4. 5.; 5. 4. 6. 2. 1.]'
Y = [6. 1. 5. 3. 2.; 2. 7. 8. 4. 3.]'

for vd in [1, 2], zm in [true, false], cr in [true, false]
    # println("vd = $vd: zm = $zm, cr = $cr")
    if vd == 1
        k = size(X, 2)
        Cxx = zeros(k, k)
        Cxy = zeros(k, k)
        for i = 1:k, j = 1:k
            Cxx[i,j] = safe_cov(X[:,i], X[:,j], zm, cr)
            Cxy[i,j] = safe_cov(X[:,i], Y[:,j], zm, cr)
        end
        x1 = vec(X[:,1])
        y1 = vec(Y[:,1])
    else
        k = size(X, 1)
        Cxx = zeros(k, k)
        Cxy = zeros(k, k)
        for i = 1:k, j = 1:k
            Cxx[i,j] = safe_cov(X[i,:], X[j,:], zm, cr)
            Cxy[i,j] = safe_cov(X[i,:], Y[j,:], zm, cr)
        end
        x1 = vec(X[1,:])
        y1 = vec(Y[1,:])
    end

    c = zm ? cov(x1; mean=0, corrected=cr) :
             cov(x1; corrected=cr)
    @test isa(c, Float64)
    @test_approx_eq c Cxx[1,1]

    C = zm ? cov(X; vardim=vd, mean=0, corrected=cr) :
             cov(X; vardim=vd, corrected=cr)
    @test size(C) == (k, k)
    @test_approx_eq C Cxx

    c = zm ? cov(x1, y1; mean=0, corrected=cr) :
             cov(x1, y1; corrected=cr)
    @test isa(c, Float64)
    @test_approx_eq c Cxy[1,1]

    C = zm ? cov(x1, Y; vardim=vd, mean=0, corrected=cr) :
             cov(x1, Y; vardim=vd, corrected=cr)
    @test size(C) == (1, k)
    @test_approx_eq C Cxy[1,:]

    C = zm ? cov(X, y1; vardim=vd, mean=0, corrected=cr) :
             cov(X, y1; vardim=vd, corrected=cr)
    @test size(C) == (k, 1)
    @test_approx_eq C Cxy[:,1]

    C = zm ? cov(X, Y; vardim=vd, mean=0, corrected=cr) :
             cov(X, Y; vardim=vd, corrected=cr)
    @test size(C) == (k, k)
    @test_approx_eq C Cxy
end

# test correlation

function safe_cor(x, y, zm::Bool)
    if !zm
        x = x .- mean(x)
        y = y .- mean(y)
    end
    x = vec(x)
    y = vec(y)
    dot(x, y) / (sqrt(dot(x, x)) * sqrt(dot(y, y)))
end

for vd in [1, 2], zm in [true, false]
    # println("vd = $vd: zm = $zm")
    if vd == 1
        k = size(X, 2)
        Cxx = zeros(k, k)
        Cxy = zeros(k, k)
        for i = 1:k, j = 1:k
            Cxx[i,j] = safe_cor(X[:,i], X[:,j], zm)
            Cxy[i,j] = safe_cor(X[:,i], Y[:,j], zm)
        end
        x1 = vec(X[:,1])
        y1 = vec(Y[:,1])
    else
        k = size(X, 1)
        Cxx = zeros(k, k)
        Cxy = zeros(k, k)
        for i = 1:k, j = 1:k
            Cxx[i,j] = safe_cor(X[i,:], X[j,:], zm)
            Cxy[i,j] = safe_cor(X[i,:], Y[j,:], zm)
        end
        x1 = vec(X[1,:])
        y1 = vec(Y[1,:])
    end

    c = zm ? cor(x1; mean=0) : cor(x1)
    @test isa(c, Float64)
    @test_approx_eq c Cxx[1,1]

    C = zm ? cor(X; vardim=vd, mean=0) : cor(X; vardim=vd)
    @test size(C) == (k, k)
    @test_approx_eq C Cxx

    c = zm ? cor(x1, y1; mean=0) : cor(x1, y1)
    @test isa(c, Float64)
    @test_approx_eq c Cxy[1,1]

    C = zm ? cor(x1, Y; vardim=vd, mean=0) : cor(x1, Y; vardim=vd)
    @test size(C) == (1, k)
    @test_approx_eq C Cxy[1,:]

    C = zm ? cor(X, y1; vardim=vd, mean=0) : cor(X, y1; vardim=vd)
    @test size(C) == (k, 1)
    @test_approx_eq C Cxy[:,1]

    C = zm ? cor(X, Y; vardim=vd, mean=0) : cor(X, Y; vardim=vd)
    @test size(C) == (k, k)
    @test_approx_eq C Cxy
end

# test hist

@test sum(hist([1,2,3])[2]) == 3
@test hist([])[2] == []
@test hist([1])[2] == [1]
@test hist([1,2,3],[0,2,4]) == ([0,2,4],[2,1])
@test hist([1,2,3],0:2:4) == (0:2:4,[2,1])
@test all(hist([1:100]/100,0.0:0.01:1.0)[2] .==1)
@test hist([1,1,1,1,1])[2][1] == 5
@test sum(hist2d(rand(100, 2))[3]) == 100
@test hist([1 2 3 4;1 2 3 4]) == (0.0:2.0:4.0, [2 2 0 0; 0 0 2 2])

@test midpoints(1.0:1.0:10.0) == 1.5:1.0:9.5
@test midpoints(1:10) == 1.5:9.5
@test midpoints(Float64[1.0:1.0:10.0]) == Float64[1.5:1.0:9.5]

@test quantile([1,2,3,4],0.5) == 2.5
@test quantile([1., 3],[.25,.5,.75])[2] == median([1., 3])
@test quantile([0.:100.],[.1,.2,.3,.4,.5,.6,.7,.8,.9])[1] == 10.0


module TestDecisionTree

using Test
import CategoricalArrays
using MLJBase

# load code to be tested:
import MLJModels
import DecisionTree
using MLJModels.DecisionTree_

# get some test data:
using RDatasets
iris = dataset("datasets", "iris")
const X = iris[:, 1:4]
const y = iris[:, 5]

baretree = DecisionTreeClassifier()

baretree.max_depth = 1
fitresult, cache, report = MLJBase.fit(baretree, 2, X, y);
baretree.max_depth = -1 # no max depth
fitresult, cache, report =
    MLJBase.update(baretree, 1, fitresult, cache, X, y);

# in this case decision tree is a perfect predictor:
yhat = MLJBase.predict_mode(baretree, fitresult, X);
@test yhat == y

# but pruning upsets this:
baretree.post_prune = true
baretree.merge_purity_threshold=0.1
fitresult, cache, report =
    MLJBase.update(baretree, 2, fitresult, cache, X, y)
yhat = MLJBase.predict_mode(baretree, fitresult, X);
@test yhat != y
yhat = MLJBase.predict(baretree, fitresult, X);

# check preservation of levels:
yyhat = predict_mode(baretree, fitresult, MLJBase.selectrows(X, 1:3))
@test MLJBase.classes(yyhat[1]) == MLJBase.classes(y[1])

info_dict(baretree)

# # testing machine interface:
# tree = machine(baretree, X, y)
# fit!(tree)
# yyhat = predict_mode(tree, MLJBase.selectrows(X, 1:3))
using Random: seed!
seed!(0)

n,m = 10^3, 5 ;
raw_features = rand(n,m);
weights = rand(-1:1,m);
labels = raw_features * weights;
features = MLJBase.table(raw_features);

R1Tree = DecisionTreeRegressor(min_samples_leaf=5, pruning_purity_threshold=0.1)
R2Tree = DecisionTreeRegressor(min_samples_split=5)
model1, = MLJBase.fit(R1Tree,1, features, labels)

vals1 = MLJBase.predict(R1Tree,model1,features)
R1Tree.post_prune = true
model1_prune, = MLJBase.fit(R1Tree,1, features, labels)
vals1_prune = MLJBase.predict(R1Tree,model1_prune,features)
@test vals1 !=vals1_prune

@test DecisionTree.R2(labels, vals1) > 0.8

model2, = MLJBase.fit(R2Tree, 1, features, labels)
vals2 = MLJBase.predict(R2Tree, model2, features)
@test DecisionTree.R2(labels, vals2) > 0.8

info_dict(R1Tree)

end
true


# The SteinMaternUnboundedPowerKernel has of the form:
#
# k(x, y) = ((1 + ||x||^alpha)*(1 + ||y||^alpha)) * km(x,y)
#
# where ||.|| is the L2 norm and km is the matern kernel.

type SteinMaternUnboundedPowerKernel <: SteinMaternWeightedKernel
    # the alpha parameter
    alpha::Float64
    # the gamma parameter
    gamma::Float64
end

# default is use gamma = sqrt{3} and alpha = 2.5
SteinMaternUnboundedPowerKernel() = SteinMaternUnboundedPowerKernel(2.5, sqrt(3.0))

# default is use gamma sqrt{3}
SteinMaternUnboundedPowerKernel(alpha::Float64) = SteinMaternUnboundedPowerKernel(alpha, sqrt(3.0))

# The term (1 + gamma*||z||^alpha)
function _Qz(ker::SteinMaternUnboundedPowerKernel, z::Array{Float64, 1})
    1.0 + norm(z,2)^ker.alpha
end

# The term ((1 + gamma*||x||^alpha)*(1 + gamma*||y||^alpha))^-1
function Q(ker::SteinMaternUnboundedPowerKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    _Qz(ker, x) * _Qz(ker, y)
end

function gradxQ(ker::SteinMaternUnboundedPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    fy = _Qz(ker, y)
    mx = norm(x,2)^(ker.alpha - 2.0)
    ker.alpha * mx * fy * x
end

function gradxyQ(ker::SteinMaternUnboundedPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    mx = norm(x,2)^(ker.alpha - 2)
    my = norm(y,2)^(ker.alpha - 2)
    ker.alpha^2 * mx * my * dot(x, y)
end


# VGG16 model for convolutional feature extraction
function vgg16(w, x; o=Dict())
    # get parameters
    fc6drop = get(o, :fc6drop, 0.0)
    fc7drop = get(o, :fc7drop, 0.0)
    mode = get(o, :cnnmode, 1)
    featuremaps = get(o, :featuremaps, false)

    conv1_1 = conv4(w[1], x; padding=1, mode=mode) .+ w[2]
    conv1_1 = relu.(conv1_1)
    conv1_2 = conv4(w[3], conv1_1; padding=1, mode=mode) .+ w[4]
    conv1_2 = relu.(conv1_2)
    pool1   = pool(conv1_2)

    conv2_1 = conv4(w[5], pool1; padding=1, mode=mode) .+ w[6]
    conv2_1 = relu.(conv2_1)
    conv2_2 = conv4(w[7], conv2_1; padding=1, mode=mode) .+ w[8]
    conv2_2 = relu.(conv2_2)
    pool2   = pool(conv2_2)

    conv3_1 = conv4(w[9], pool2; padding=1, mode=mode) .+ w[10]
    conv3_1 = relu.(conv3_1)
    conv3_2 = conv4(w[11], conv3_1; padding=1, mode=mode) .+ w[12]
    conv3_2 = relu.(conv3_2)
    conv3_3 = conv4(w[13], conv3_2; padding=1, mode=mode) .+ w[14]
    conv3_3 = relu.(conv3_3)
    pool3   = pool(conv3_3)

    conv4_1 = conv4(w[15], pool3; padding=1, mode=mode) .+ w[16]
    conv4_1 = relu.(conv4_1)
    conv4_2 = conv4(w[17], conv4_1; padding=1, mode=mode) .+ w[18]
    conv4_2 = relu.(conv4_2)
    conv4_3 = conv4(w[19], conv4_2; padding=1, mode=mode) .+ w[20]
    conv4_3 = relu.(conv4_3)
    pool4   = pool(conv4_3)

    conv5_1 = conv4(w[21], pool4; padding=1, mode=mode) .+ w[22]
    conv5_1 = relu.(conv5_1)
    conv5_2 = conv4(w[23], conv5_1; padding=1, mode=mode) .+ w[24]
    conv5_2 = relu.(conv5_2)
    conv5_3 = conv4(w[25], conv5_2; padding=1, mode=mode) .+ w[26]
    conv5_3 = relu.(conv5_3)

    if !featuremaps
        pool5 = pool(conv5_3)
        fc6 = w[27] * mat(pool5) .+ w[28]
        fc6 = relu.(fc6)
        fc6 = dropout(fc6, fc6drop)
        fc7 = w[29] * mat(fc6) .+ w[30]
        fc7 = relu.(fc7)
    end
end

# VGG16 model for convolutional feature extraction
function vgg19(w, x; o=Dict())
    # get parameters
    fc6drop = get(o, :fc6drop, 0.0)
    fc7drop = get(o, :fc7drop, 0.0)
    mode = get(o, :cnnmode, 1)
    featuremaps = get(o, :featuremaps, false)

    conv1_1 = conv4(w[1], x; padding=1, mode=mode) .+ w[2]
    conv1_1 = relu.(conv1_1)
    conv1_2 = conv4(w[3], conv1_1; padding=1, mode=mode) .+ w[4]
    conv1_2 = relu.(conv1_2)
    pool1   = pool(conv1_2)

    conv2_1 = conv4(w[5], pool1; padding=1, mode=mode) .+ w[6]
    conv2_1 = relu.(conv2_1)
    conv2_2 = conv4(w[7], conv2_1; padding=1, mode=mode) .+ w[8]
    conv2_2 = relu.(conv2_2)
    pool2   = pool(conv2_2)

    conv3_1 = conv4(w[9], pool2; padding=1, mode=mode) .+ w[10]
    conv3_1 = relu.(conv3_1)
    conv3_2 = conv4(w[11], conv3_1; padding=1, mode=mode) .+ w[12]
    conv3_2 = relu.(conv3_2)
    conv3_3 = conv4(w[13], conv3_2; padding=1, mode=mode) .+ w[14]
    conv3_3 = relu.(conv3_3)
    conv3_4 = conv4(w[15], conv3_3; padding=1, mode=mode) .+ w[16]
    conv3_4 = relu.(conv3_4)
    pool3   = pool(conv3_4)

    conv4_1 = conv4(w[17], pool3; padding=1, mode=mode) .+ w[18]
    conv4_1 = relu.(conv4_1)
    conv4_2 = conv4(w[19], conv4_1; padding=1, mode=mode) .+ w[20]
    conv4_2 = relu.(conv4_2)
    conv4_3 = conv4(w[21], conv4_2; padding=1, mode=mode) .+ w[22]
    conv4_3 = relu.(conv4_3)
    conv4_4 = conv4(w[23], conv4_3; padding=1, mode=mode) .+ w[24]
    conv4_4 = relu.(conv4_4)
    pool4   = pool(conv4_4)

    conv5_1 = conv4(w[25], pool4; padding=1, mode=mode) .+ w[26]
    conv5_1 = relu.(conv5_1)
    conv5_2 = conv4(w[27], conv5_1; padding=1, mode=mode) .+ w[28]
    conv5_2 = relu.(conv5_2)
    conv5_3 = conv4(w[29], conv5_2; padding=1, mode=mode) .+ w[30]
    conv5_3 = relu.(conv5_3)
    conv5_4 = conv4(w[31], conv5_3; padding=1, mode=mode) .+ w[32]
    conv5_4 = relu.(conv5_4)

    if !featuremaps
        pool5 = pool(conv5_4)
        fc6 = w[33] * mat(pool5) .+ w[34]
        fc6 = relu.(fc6)
        fc6 = dropout(fc6, fc6drop)
        fc7 = w[35] * mat(fc6) .+ w[36]
        fc7 = relu.(fc7)
        return fc7
    end

    return conv5_4
end

function get_vgg_weights(vggmat; last_layer="relu.7")
    ws = []
    for l in vggmat["layers"]
        if haskey(l, "weights") && length(l["weights"]) != 0
            w = l["weights"]
            if startswith(l["name"], "conv")
                w[2] = reshape(w[2], (1,1,length(w[2]),1))
            elseif startswith(l["name"], "fc")
                w[1] = mat(w[1])'
            end
            push!(ws, w...)
        end

        last_layer == l["name"] && break
    end
    map(KnetArray, ws)
end

# mode, 0=>train, 1=>test
function resnet50(w,x,ms; mode=1)
    # layer 1
    conv1  = conv4(w[1],x; padding=3, stride=2) .+ w[2]
    bn1    = batchnorm(w[3:4],conv1,ms; mode=mode)
    pool1  = pool(bn1; window=3, stride=2)

    # layer 2,3,4,5
    r2 = reslayerx5(w[5:34], pool1, ms; strides=[1,1,1,1], mode=mode)
    r3 = reslayerx5(w[35:73], r2, ms; mode=mode)
    r4 = reslayerx5(w[74:130], r3, ms; mode=mode) # 5
    r5 = reslayerx5(w[131:160], r4, ms; mode=mode)

    # fully connected layer
    pool5  = pool(r5; stride=1, window=7, mode=2)
    fc1000 = w[161] * mat(pool5) .+ w[162]
end

# mode, 0=>train, 1=>test
function resnet101(w,x,ms; mode=1)
    # layer 1
    conv1 = reslayerx1(w[1:3],x,ms; padding=3, stride=2, mode=mode)
    pool1 = pool(conv1; window=3, stride=2)

    # layer 2,3,4,5
    r2 = reslayerx5(w[4:33], pool1, ms; strides=[1,1,1,1], mode=mode)
    r3 = reslayerx5(w[34:72], r2, ms; mode=mode)
    r4 = reslayerx5(w[73:282], r3, ms; mode=mode)
    r5 = reslayerx5(w[283:312], r4, ms; mode=mode)

    # fully connected layer
    pool5  = pool(r5; stride=1, window=7, mode=2)
    fc1000 = w[313] * mat(pool5) .+ w[314]
end

# mode, 0=>train, 1=>test
function resnet152(w,x,ms; mode=1)
    # layer 1
    conv1 = reslayerx1(w[1:3],x,ms; padding=3, stride=2, mode=mode)
    pool1 = pool(conv1; window=3, stride=2)

    # layer 2,3,4,5
    r2 = reslayerx5(w[4:33], pool1, ms; strides=[1,1,1,1], mode=mode)
    r3 = reslayerx5(w[34:108], r2, ms; mode=mode)
    r4 = reslayerx5(w[109:435], r3, ms; mode=mode)
    r5 = reslayerx5(w[436:465], r4, ms; mode=mode)

    # fully connected layer
    pool5  = pool(r5; stride=1, window=7, mode=2)
    fc1000 = w[466] * mat(pool5) .+ w[467]
end

# Batch Normalization Layer
# works both for convolutional and fully connected layers
# mode, 0=>train, 1=>test
function batchnorm(w, x, ms; mode=1, epsilon=1e-5)
    mu, sigma = nothing, nothing
    if mode == 0
        d = ndims(x) == 4 ? (1,2,4) : (2,)
        s = reduce((a,b)->a*size(x,b), d)
        mu = sum(x,d) / s
        sigma = sqrt(epsilon + (sum(x.-mu,d).^2) / s)
    elseif mode == 1
        mu = shift!(ms)
        sigma = shift!(ms)
    end

    # we need getval in backpropagation
    push!(ms, AutoGrad.getval(mu), AutoGrad.getval(sigma))
    xhat = (x.-mu) ./ sigma
    return w[1] .* xhat .+ w[2]
end

function reslayerx0(w,x,ms; padding=0, stride=1, mode=1)
    b  = conv4(w[1],x; padding=padding, stride=stride)
    bx = batchnorm(w[2:3],b,ms; mode=mode)
end

function reslayerx1(w,x,ms; padding=0, stride=1, mode=1)
    relu.(reslayerx0(w,x,ms; padding=padding, stride=stride, mode=mode))
end

function reslayerx2(w,x,ms; pads=[0,1,0], strides=[1,1,1], mode=1)
    ba = reslayerx1(w[1:3],x,ms; padding=pads[1], stride=strides[1], mode=mode)
    bb = reslayerx1(w[4:6],ba,ms; padding=pads[2], stride=strides[2], mode=mode)
    bc = reslayerx0(w[7:9],bb,ms; padding=pads[3], stride=strides[3], mode=mode)
end

function reslayerx3(w,x,ms; pads=[0,0,1,0], strides=[2,2,1,1], mode=1) # 12
    a = reslayerx0(w[1:3],x,ms; stride=strides[1], padding=pads[1], mode=mode)
    b = reslayerx2(w[4:12],x,ms; strides=strides[2:4], pads=pads[2:4], mode=mode)
    relu.(a .+ b)
end

function reslayerx4(w,x,ms; pads=[0,1,0], strides=[1,1,1], mode=1)
    relu.(x .+ reslayerx2(w,x,ms; pads=pads, strides=strides, mode=mode))
end

function reslayerx5(w,x,ms; strides=[2,2,1,1], mode=1)
    x = reslayerx3(w[1:12],x,ms; strides=strides, mode=mode)
    for k = 13:9:length(w)
        x = reslayerx4(w[k:k+8],x,ms; mode=mode)
    end
    return x
end

function get_params(params)
    len = length(params["value"])
    ws, ms = [], []
    for k = 1:len
        name = params["name"][k]
        value = convert(Array{Float32}, params["value"][k])

        if endswith(name, "moments")
            push!(ms, reshape(value[:,1], (1,1,size(value,1),1)))
            push!(ms, reshape(value[:,2], (1,1,size(value,1),1)))
        elseif startswith(name, "bn")
            push!(ws, reshape(value, (1,1,length(value),1)))
        elseif startswith(name, "fc") && endswith(name, "filter")
            push!(ws, transpose(reshape(value,size(value,3,4))))
        elseif startswith(name, "conv") && endswith(name, "bias")
            push!(ws, reshape(value, (1,1,length(value),1)))
        else
            push!(ws, value)
        end
    end
    map(KnetArray, ws), map(KnetArray, ms)
end


struct PLSE <: ParametrisedConvexApproximator
    n::Int
    m::Int
    i_max::Int
    T::Real
    NN::Flux.Chain
    function PLSE(n::Int, m::Int, i_max::Int, T::Real, h_array::Vector{Int}, act)
        @assert T > 0
        node_array = [n, h_array..., i_max*(m+1)]
        new(n, m, i_max, T, construct_layer_array(node_array, act))
    end
end

"""
    (nn::PLSE)(x, u)

Infer using two input arguments `x` and `u`.
# Notes
size(x) = (n, d)
size(u) = (m, d)
# Types
x and u should be arrays.
For example, u = [1] for the one-element case.
"""
function (nn::PLSE)(x::AbstractArray, u::AbstractArray)
    @unpack T = nn
    is_vector = length(size(x)) == 1
    @assert is_vector == (length(size(u)) == 1)
    x = is_vector ? reshape(x, :, 1) : x
    u = is_vector ? reshape(u, :, 1) : u
    @assert size(x)[2] == size(u)[2]
    tmp = affine_map(nn, x, u)
    _res = T * Flux.logsumexp((1/T)*tmp, dims=1)
    res = is_vector ? reshape(_res, 1) : _res
end

function (nn::PLSE)(x::AbstractArray, u::Convex.AbstractExpr)
    @unpack T = nn
    tmp = affine_map(nn, x, u)
    res = [T * Convex.logsumexp((1/T)*tmp)]
end

Flux.params(approximator::PLSE) = Flux.params(approximator.NN)


immutable Hypergeometric <: DiscreteUnivariateDistribution
    ns::Float64 # number of successes in population
    nf::Float64 # number of failures in population
    n::Float64  # sample size
    function Hypergeometric(s::Real, f::Real, n::Real)
        isinteger(s) && zero(s) <= s || error("ns must be a non-negative integer")
        isinteger(f) && zero(f) <= f || error("nf must be a non-negative integer")
        isinteger(n) && zero(n) < n < s + f ||
            error("n must be a positive integer <= (ns + nf)")
        new(float64(s), float64(f), float64(n))
    end
end

@_jl_dist_3p Hypergeometric hyper

function insupport(d::Hypergeometric, x::Real)
    isinteger(x) && zero(x) <= x <= d.n && (d.n - d.nf) <= x <= d.ns
end

mean(d::Hypergeometric) = d.n * d.ns / (d.ns + d.nf)

function var(d::Hypergeometric)
    N = d.ns + d.nf
    p = d.ns / N
    d.n * p * (1.0 - p) * (N - d.n) / (N - 1.0)
end


function eval_exp(N)
    a = range(0, stop=2*pi, length=N)
    A = Matrix{ComplexF64}(undef, N, N)

    @inbounds Threads.@threads for j in 1:N
    for i in 1:N
        A[i,j] = exp((100+im)*im*sqrt(a[i]^2 + a[j]^2))
    end
    end

    return A

end

eval_exp(5)
print(string("running loop on ", Threads.nthreads(), " threads \n"))
for N in 1000:1000:10000
    @time begin
    A = eval_exp(N)
    end
end


"""
    ChebyshevImplementation

implementation of computation of second Poincare invariant by approximating surface with
Chebyshev polynomials
"""
module ChebyshevImplementation

using ...PoincareInvariants: @argcheck
import ...PoincareInvariants: compute!, getpoints, getpointnum

using Base: Callable
using LinearAlgebra

include("PaduaTransforms.jl")
using .PaduaTransforms

const AbstractArray3 = AbstractArray{<:Any, 3}

## Differentiation ##

function getdiffmat(::Type{T}, degree::Integer) where T
    D = zeros(T, degree+1, degree+1)

    D[1, 2:2:end] .= 1:2:degree

    for i in 3:2:degree+1
        D[2:2:i-1, i] .= 2 * (i-1)
    end

    for i in 4:2:degree+1
        D[3:2:i-1, i] .= 2 * (i-1)
    end

    D
end

struct DiffPlan{T}
    D::Matrix{T}
end

DiffPlan{T}(degree::Integer) where T = DiffPlan{T}(getdiffmat(T, degree))

# differentiate in the Chebyshev basis
function differentiate!(∂x::AbstractMatrix, ∂y::AbstractMatrix, P::DiffPlan, coeffs::AbstractMatrix)
    ∂x[:, :] = coeffs  # differentiate each row
    rmul!(∂x, LowerTriangular(P.D'))

    ∂y[:, :] = coeffs  # differentiate each column
    lmul!(UpperTriangular(P.D), ∂y)

    ∂x, ∂y
end

function differentiate!(∂x::AbstractArray3, ∂y::AbstractArray3, P::DiffPlan, coeffs::AbstractArray3)
    axes(∂x, 3) == axes(∂y, 3) == axes(coeffs, 3) || error()

    @views for i in axes(coeffs, 3)
        differentiate!(∂x[:, :, i], ∂y[:, :, i], P, coeffs[:, :, i])
    end

    ∂x, ∂y
end

## Integration ##

getintegrator(::Type{T}, n) where T = T[isodd(i) ? 0 : T(2) / T(1 - i^2) for i in 0:n]
getintegrator(n) = getintegrator(Float64, n)

integrate(coeffs, integrator) = dot(integrator, coeffs, integrator)

## getintegrand for Ω Callable and out-of-place ##

struct OOPIntPlan{T, IP, P}
    invpaduaplan::IP
    ∂xvals::Matrix{T}
    ∂yvals::Matrix{T}
    intvals::Vector{T}
    paduaplan::P
end

function OOPIntPlan{T}(D, degree) where T
    invpaduaplan = InvPaduaTransformPlan{T}(degree)
    ∂xvals = Matrix{T}(undef, getpaduanum(degree), D)
    ∂yvals = Matrix{T}(undef, getpaduanum(degree), D)
    intvals = Vector{T}(undef, getpaduanum(degree))
    paduaplan = PaduaTransformPlan{T}(degree)

    OOPIntPlan{T, typeof(invpaduaplan), typeof(paduaplan)}(
        invpaduaplan, ∂xvals, ∂yvals, intvals, paduaplan
    )
end

function getintegrand!(
    intcoeffs::AbstractMatrix, plan::OOPIntPlan, Ω::Callable, phasepoints::AbstractMatrix, t, p,
    ∂xcoeffs::AbstractArray3, ∂ycoeffs::AbstractArray3
)
    ∂xvals = invpaduatransform!(plan.∂xvals, plan.invpaduaplan, ∂xcoeffs)
    ∂yvals = invpaduatransform!(plan.∂yvals, plan.invpaduaplan, ∂ycoeffs)

    @views for i in axes(plan.intvals, 1)
        plan.intvals[i] = dot(∂yvals[i, :], Ω(phasepoints[i, :], t, p), ∂xvals[i, :])
    end

    paduatransform!(intcoeffs, plan.paduaplan, plan.intvals)

    intcoeffs
end

## ChebyshevPlan and _compute! ##

getintplan(::Type{T}, ::Callable, D, degree, ::Val{false}) where T = OOPIntPlan{T}(D, degree)
# getintplan(::Type{T}, ::Callable, D, degree, ::Val{true}) where T = IPIntPlan{T}(D, degree)
# getintplan(::Type{T}, Ω::AbstractMatrix, D, degree, ::Val{nothing}) where T = IPIntPlan{T}(Ω, D, degree)

struct ChebyshevPlan{T, IP, PP<:PaduaTransformPlan}
    degree::Int
    paduaplan::PP
    phasecoeffs::Array{T, 3}
    diffplan::DiffPlan{T}
    ∂x::Array{T, 3}
    ∂y::Array{T, 3}
    intplan::IP  # getting coefficients of integrand to integrate
    intcoeffs::Matrix{T}
    integrator::Vector{T}
end

function ChebyshevPlan{T}(Ω::Callable, D::Integer, N::Integer, ::Val{inplace}) where {T, inplace}
    degree = getdegree(nextpaduanum(N))

    paduaplan = PaduaTransformPlan{T}(degree)
    phasecoeffs = zeros(T, degree+1, degree+1, D)
    diffplan = DiffPlan{T}(degree)
    ∂x = Array{T, 3}(undef, degree+1, degree+1, D)
    ∂y = Array{T, 3}(undef, degree+1, degree+1, D)
    intplan = getintplan(T, Ω, D, degree, Val(inplace))
    intcoeffs = zeros(T, degree+1, degree+1)
    integrator = getintegrator(T, degree)

    ChebyshevPlan{T, typeof(intplan), typeof(paduaplan)}(
        degree, paduaplan, phasecoeffs, diffplan, ∂x, ∂y, intplan, intcoeffs, integrator
    )
end

function compute!(plan::ChebyshevPlan, Ω::Callable, phasepoints::AbstractMatrix, t, p)
    paduatransform!(plan.phasecoeffs, plan.paduaplan, phasepoints)
    differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
    getintegrand!(plan.intcoeffs, plan.intplan, Ω, phasepoints, t, p, plan.∂x, plan.∂y)
    integrate(plan.intcoeffs, plan.integrator)
end

# function _compute!(plan::ChebyshevPlan, Ω::AbstractMatrix, D, phasepoints)
#     paduatransform!(plan.phasecoeffs, plan.paduaplan, phasepoints)
#     differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
#     getintegrand!(plan.intcoeffs, plan.intplan, Ω, plan.∂x, plan.∂y)
#     integrate(plan.intcoeffs, plan.intplan)
# end

## getpoints and getpointnum ##

getpointnum(plan::ChebyshevPlan) = getpaduanum(plan.degree)
getpoints(plan::ChebyshevPlan) = getpaduapoints(plan.degree) do x, y
    (x + 1) / 2, (y + 1) / 2
end

getpoints(f::Function, plan::ChebyshevPlan) = getpaduapoints(plan.degree) do x, y
    f((x + 1) / 2, (y + 1) / 2)
end

end  # module ChebyshevImplementation


import Base.reshape, Base.vec
export reshape, vec

reshape(x::Constant, m::Int64, n::Int64) = Constant(reshape(x.value, m, n))

function reshape(x::AffineExpr, m::Int64, n::Int64)
  if x.size[1] * x.size[2] != m * n
    error("Cannot reshape expression of size $(x.size[1]) by $(x.size[2]) to $m by $n")
  end
  this = AffineExpr(:reshape, (x,), x.vars_to_coeffs_map, x.constant, (m, n))
  this.evaluate = ()->reshape(x.evaluate(), m, n)
  return this
end

vec(x::AffineOrConstant) = reshape(x, x.size[1] * x.size[2], 1)


using LinearAlgebra
using GPFlux
using Zygote
using Random
using GaussianProcesses
using Test

@testset "negative logliklihood" begin
	Random.seed!(4)
	d = 4
	n = 10
	X = rand(d, n)
	y = rand(n)

	# GaussianProcesses GP
	gp_zmean = MeanZero()
	gp_se_kernel = SE(4.0, 3.0)
	gp_lognoise = -1.0
	gp_gp = GP(X, y, gp_zmean, gp_se_kernel, gp_lognoise)
	PT = GaussianProcesses.FullCovariancePrecompute(n)
	CST = GaussianProcesses.FullCovariance()
	GaussianProcesses.update_mll_and_dmll!(gp_gp, PT)
	gp_nll = -gp_gp.mll
	gp_grads = -gp_gp.dmll

	# GPFlux GP
	zmean = ConstantMean()
	se_kernel = IsoGaussKernel([4.0], [3.0])
	noise2 = [-1.0]
	gp = GaussProcess(zmean, se_kernel, noise2)
	ps = params(gp)
	my_nll = negloglik(gp, X, y)
	gs = gradient(()->negloglik(gp, X, y), ps)
	my_grads = [gs[p][1] for p in ps if p!=[0.0]]

	@show gp_nll
	@show my_nll
	@show sort(gp_grads)
	@show sort(my_grads)

	@test isapprox(my_nll, gp_nll, rtol=1e-3, atol=1e-8)
	@test isapprox(sort(gp_grads), sort(my_grads), rtol=1e-4, atol=1e-8)
end


using DataFrames, Query, Compat, GLM, CSV
using PlotlyJS

function percentError(estimate,exact)
  return abs((estimate-exact)/exact)*100
end

function percentDifference(val1,val2)
  return abs((val1-val2)/((val1+val2)/2))*100
end

function readTestData(f)
  re = r"^\[(?<data>[-?\d+\.\d+\s*]+)\],\s+\[(?<label>[-?\d+\.\d+?\s*]+)\]$"
  data = Float64[]
  label = Float64[]
  # Read headerline
  header = readline(f)
  while !eof(f)
    s = readline(f)
    m = match(re,s)
    data = [data;readdlm(IOBuffer(m[:data]))]
    label = [label;readdlm(IOBuffer(m[:label]))]
  end
  data = convert(DataFrame,
                 Dict(zip(["data$i" for i in 1:size(data,2)],
                          [data[:,i] for i in 1:size(data,2)])))
  label = convert(DataFrame,
                  Dict(zip(["label$i" for i in 1:size(label,2)],
                           [label[:,i] for i in 1:size(label,2)])))
  return hcat(data,label)
end

function readLqrData(filename)
  data = Float64[]
  label = Float64[]
  df = CSV.read(filename; delim=',', header=1)
  return df
end

function getQuadAngleControl(df)
  controls = Float64[]
  controlSystem = [1 1 -1 -1;
                   1 -1 -1 1;
                   1 1 1 -1;
                   1 -1 1 1]

  for i=1:size(df,1)
    controls = [controls ;
      transpose(controlSystem \
                [df[:label1][i],df[:label2][i],df[:label3][i],df[:label4][i]])]
  end
  controls = convert(DataFrame,
                     Dict(zip(["thrustCtl" "pitchCtl" "rollCtl" "yawCtl"],
                              [controls[:,i] for i in 1:size(controls,2)])))

  return hcat(df,controls)
end

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

token = "LQR"
if token == "test"
  f = open("./data/test.txt","r")

  df = readTestData(f);
  df = getQuadAngleControl(df)

  linearModel = Dict("pitch" => lm(@formula(pitchCtl ~ data10 + data11 + data12 +
                                            data4 + data5 + data6), df),
                     "roll" => lm(@formula(rollCtl ~ data13 + data14 + data15 +
                                           data1 + data2 + data3), df),
                     "yaw" => lm(@formula(yawCtl ~ data7 + data8 + data9),df),
                     "thrust" => lm(@formula(thrustCtl ~ data16 + data17 + data18),df))

  print(linearModel["roll"])
  print(linearModel["yaw"])
  print(linearModel["thrust"])

elseif token == "PID"

elseif token == "LQR"

  # LQR gains
  gains = [0.0020 1.0000 -0.0000 -0.0418 1.3856 -0.0000 1.1039 0.0001 -0.0000 4.6346 0.0395 -0.0000;
           -1.0000 0.0020 -0.0043 -1.3583 -0.0458 -0.0015 0.0004 1.1348 -0.0028 -0.0247 4.6926 -0.0028;
           -0.0035 0.0000 0.1773 -0.0070 -0.0001 -0.0188 -0.0000 -0.0010 1.0131 -0.0002 -0.0142 0.9842;
           0.0037 -0.0000 -0.9842 0.0147 0.0001 -0.9920 0.0000 0.0044 0.2049 0.0001 0.2483 0.1773]

  f = Dict{Any,Any}()
  df = Dict{Any,Any}()
  gainDiff = Dict{Any,Any}()
  linearModels = Dict{Any,Dict{String,Any}}()
  f[1] = [0,"data/lqrNn0Loops.csv"]
  f[2] = [5,"data/lqrNn5Loops.csv"]
  f[3] = [10,"data/lqrNn10Loops.csv"]
  f[4] = [15,"data/lqrNn15Loops.csv"]
  f[5] = [20,"data/lqrNn20Loops.csv"]

  for i=1:5
    numLoops = f[i][1]
    df[numLoops] = readLqrData(f[i][2])
    linearModels[i] = Dict("del_lat" => lm(@formula(del_lat ~ x + y + z + u + v
                                                            + w + p + q + r +
                                                            phi + theta + psi),
                                           df[numLoops]),
                           "del_lon" => lm(@formula(del_lon ~ x + y + z + u + v
                                                            + w + p + q + r +
                                                            phi + theta + psi),
                                           df[numLoops]),
                           "del_yaw" => lm(@formula(del_yaw ~ x + y + z + u + v
                                                            + w + p + q + r +
                                                            phi + theta + psi),
                                           df[numLoops]),
                           "del_thrust" => lm(@formula(del_thrust ~ x + y + z + u + v
                                                                  + w + p + q + r +
                                                                  phi + theta +
                                                                  psi),df[numLoops]))

    gainDiff[numLoops] = zeros(size(gains))
    r=[]

    push!(r,[coef(linearModels[i]["del_lat"])[2:end] gains[1,:]])
    push!(r,[coef(linearModels[i]["del_lon"])[2:end] gains[2,:]])
    push!(r,[coef(linearModels[i]["del_yaw"])[2:end] gains[3,:]])
    push!(r,[coef(linearModels[i]["del_thrust"])[2:end] gains[4,:]])

    for i=1:4
      for j=1:size(r[i],1)
        gainDiff[numLoops][i,j] = percentDifference(r[i][j,1],r[i][j,2])
      end
    end
  end

  numLoopsSet = [0 5 10 15 20]'
  gainDiffChange = [mean(gainDiff[i]) for i in numLoopsSet]'

  # Plotting
  #=
  trace1 = scatter(;x=numLoopsSet,y=gainDiffChange,
            mode="markers",
            marker=attr(color="#02010f",size=10,symbol="square",
                        line=attr(color="#02010f",width=0)),
            name="Gain Percent Difference")

  layout = Layout(autosize=false, width=500, height=500,
                  margin=attr(l=0,r=0,b=0,t=65))
  p = plot([trace1],layout)
  =#

  plot(scatter(x=numLoopsSet,y=gainDiffChange,
               marker=attr(color="#b3010f",size=10,symbol="square",
                           line=attr(color="#02010f",width=0)),
               name="Gain Percent Difference"))
  display(p)

end


using Distributions, DataFrames, DataFramesMeta, Arrow, Statistics, Dates, CSV, ODBC, WRDSMerger
using Revise
using AbnormalReturns

##
conn = ODBC.Connection("wrds-pgdata-64")
##

df_ff_data = @chain ff_data(conn) begin
    @rtransform(:mkt = :mktrf + :rf)
end

ff_data_max = maximum(df_ff_data.date)

@time df_crsp_raw = @chain crsp_data(conn, Date(1990), ff_data_max; cols=["ret"]) begin
    sort([:permno, :date])
end

data = MarketData(
    df_ff_data,
    df_crsp_raw
)

##

mkt_mat = @chain df_ff_data begin
    select([:mktrf, :smb, :hml, :umd])
    dropmissing
    Matrix
    cov
end
# 4×4 Matrix{Float64}:
#   0.000116248  -1.03659e-5   1.1843e-5   -1.04678e-5
#  -1.03659e-5    3.49562e-5  -2.37018e-6   1.9103e-6
#   1.1843e-5    -2.37018e-6   3.77871e-5  -1.05367e-5
#  -1.04678e-5    1.9103e-6   -1.05367e-5   5.98879e-5

##

@time df_all_regs = @chain df_crsp_raw begin
    select([:permno, :date])
    dropmissing
    @by(
        :permno,
        :date_start = minimum(:date),
        :date_end = maximum(:date)
    )
    @transform(:reg = AbnormalReturns.vector_reg(data, :permno, :date_start, :date_end, @formula(ret ~ mktrf + smb + hml + umd); minobs=100, save_residuals=true))
end

##

function winsorize(vect::AbstractArray, low_per::AbstractFloat=0.01, high_per::AbstractFloat=0.99)
    lower = quantile(skipmissing(vect), low_per)
    upper = quantile(skipmissing(vect), high_per)
    returnVect = vect[:]
    for (i, v) in enumerate(vect)
        if typeof(v) == Missing
            continue
        end
        if v > upper
            returnVect[i] = upper
        elseif v < lower
            returnVect[i] = lower
        end
    end
    return returnVect
end

df_coefs = @chain df_all_regs begin
    @rtransform(:r2 = r2(:reg))
    dropmissing
    @rtransform(:coefs = NamedTuple{(Symbol.(coefnames(:reg))..., ^(:var))}((coef(:reg)..., var(:reg))))
     _.coefs
    DataFrame
    rename("(Intercept)" => "int")
    @transform(
        :int = winsorize(:int),
        :mktrf = winsorize(:mktrf),
        :smb = winsorize(:smb),
        :hml = winsorize(:hml),
        :umd = winsorize(:umd),
        :var = winsorize(:var)
    )
end
##
coef_means = mean(df_coefs[:, 1:5] |> Matrix, dims=1)[1, :]
# [
#     0.0003416652538979692,
#     0.7557301603393168,
#     0.6035365140401592,
#     0.13841355976862427,
#     -0.0917780093857371,
# ]
coef_cov = cov(df_coefs[:, 1:5] |> Matrix)
# 2.52326e-6  -4.71181e-5   3.31182e-5  -7.05013e-5    1.32415e-5
# -4.71181e-5   0.27552      0.173033     0.0742338    -0.0448335
#  3.31182e-5   0.173033     0.389159     0.0460528    -0.0280572
# -7.05013e-5   0.0742338    0.0460528    0.40226      -0.000262376
#  1.32415e-5  -0.0448335   -0.0280572   -0.000262376   0.183439

fit_mle(Gamma, @rsubset(df_coefs, :var > 0).var)
# Gamma{Float64}(α=0.5792772418365709, θ=0.00448646291687775)


## Functions for centered moments of moments feature maps
#const m1s2 = [1,2,3,4,5];
const m2lim2 = 10;
function apply_mofm_full(X::Array{Float64}, rowcounts::Int64, mom::StatisticalMoments,options)
   lenresponse = length(mom.moments[1]);
   momnums = options.aux;
   s=size(momnums);
   count =1; i =1;
    if(length(s)==1)
        for i in momnums
            mm  = mom.moments[i];
            for j = 1:m2lim2
                X[rowcounts,count] = get_moment_vec(mm,j);
                count = count +1;
            end
        end
   else
        for mm in mom.moments
            for j = 1:m2lim2
                if(options.aux[i,j]==1)
                   X[rowcounts,count] = get_moment_vec(mm,j);
                   count = count +1;
                end
            end
            i = i +1;
        end
   end
end

function get_num_mofm__full_features(mom::StatisticalMoments, options)
    s=size(options.aux);
    if(length(s)==1)
        lenresponse =  length(options.aux)*m2lim2;
     else
        lenresponse = sum(options.aux);
    end
    return lenresponse;
end

# Defines the feature map
fmofm_full = FeatureMap(apply_mofm_full,get_num_mofm__full_features, "mofm_full_map");

#"returns the  mom_num-th moment of resps, where resps is a vector whose elements can be either numbers or vectors."
function get_moment_vec(resps, mom_num)
    mu = mean(resps);
    if(mom_num==1) return mu end;
    num_sim = length(resps);
    if(num_sim<2)
        println("Not enough simulations to calculate momments, return zeros")
        return 0.0*mu
    end
    moment = sum((mu - r).^mom_num for r in resps)./(num_sim-1.0)
    return sign(moment).*abs(moment).^(1.0/mom_num) # I think the sign does make a difference, as it changes with different labels
end


include("activate.jl")

using LinearAlgebra

n = Int(1e3);
k = 1e5

xpi = rand(n);
total = sum(xpi);
const xp = xpi ./ total;

f(x) = 2 * norm(x - xp)^3 - norm(x)^2

# this is just for the example -> better explicitly define your gradient
grad!(storage, x) = ReverseDiff.gradient!(storage, f, x)

# pick feasible region
lmo = FrankWolfe.ProbabilitySimplexOracle(1.0); #radius needs to be float

# compute some initial vertex
x0 = collect(FrankWolfe.compute_extreme_point(lmo, zeros(n)))

# benchmarking Oracles
FrankWolfe.benchmark_oracles(f, grad!, () -> randn(n), lmo; k=100)

@time x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Nonconvex(),
    print_iter=k / 10,
    verbose=true,
);


module test_RotationMatrix

import ModiaMath

#  Desired:
#    using Test
#    using LinearAlgebra
#    using StaticArrays
#    using Unitful
#
#  In order that these packages need not to be defined in the user environment, they are included via ModiaMath:
using ModiaMath.Test
using ModiaMath.LinearAlgebra
using ModiaMath.StaticArrays
using ModiaMath.Unitful

angle1 = pi / 2
angle2 = 90u"°"

R1a = ModiaMath.rot1(angle1)
R1b = ModiaMath.rot1(angle2)
R1c = ModiaMath.rot_e([1,0,0], angle1)
R1d = ModiaMath.rot_e([1,0,0], angle2)
R1e = ModiaMath.rot_nxy([1,0,0], [0,0,1])

R2a = ModiaMath.rot2(angle1)
R2b = ModiaMath.rot2(angle2)
R2c = ModiaMath.rot_e([0,1,0], angle1)
R2d = ModiaMath.rot_e([0,1,0], angle2)

R3a = ModiaMath.rot3(angle1)
R3b = ModiaMath.rot3(angle2)
R3c = ModiaMath.rot_e([0,0,1], angle1)
R3d = ModiaMath.rot_e([0,0,1], angle2)

R4a = ModiaMath.rot123(angle1, angle1, angle1)
R4b = ModiaMath.rot123(angle2, angle2, angle2)
R4c = ModiaMath.rot_e([1,0,0], -angle1) *
      ModiaMath.rot_e([0,1,0], -angle1) *
      ModiaMath.rot_e([0,0,1], -angle1)

R5  = ModiaMath.rot_nxy([1,0,0], [0,0,1])

v1  = SVector{3,Float64}(2.0, -1.0, 4.0)

v2a = ModiaMath.resolve2(R1a, v1)
v1a = ModiaMath.resolve1(R1a, v2a)

R6a = ModiaMath.absoluteRotation(R1a, R2a)
R6b = ModiaMath.rot123(angle1, angle1, 0.0)
R7  = ModiaMath.relativeRotation(R1a, R6b)

@testset "ModiaMath.Frames: test RotationMatrix" begin
    @test isapprox(R1b, R1a)
    @test isapprox(R1c, R1a)
    @test isapprox(R1d, R1a)
    @test isapprox(R1e, R1a)

    @test isapprox(R2b, R2a)
    @test isapprox(R2c, R2a)
    @test isapprox(R2d, R2a)

    @test isapprox(R3b, R3a)
    @test isapprox(R3c, R3a)
    @test isapprox(R3d, R3a)

    @test isapprox(R4b, R4a)
    @test isapprox(R4c, R4a)

    @test isapprox(R5, R1a)

    @test isapprox(v1, v1a)

    @test isapprox(R6a, R6b)
    @test isapprox(R7, R2a)

    R4d = ModiaMath.inverseRotation(R4a)
    R4e = ModiaMath.inverseRotation(R4d)
    @test isapprox(R4a, R4e)

    angle4 = 45u"°"
    e      = normalize([1.0, 1.0, 1.0])
    R      = ModiaMath.rot_e(e, angle4)
    v3     = [1.0, 2.0, 3.0]
    v4     = ModiaMath.resolve2(R, v3)
    angle5 = ModiaMath.planarRotationAngle(e, v3, v4)
    isapprox(angle4, angle5)

    isapprox(-ModiaMath.eAxis(3), ModiaMath.eAxis(-3))
end

end


function prime_field(K::NumField)
  return FlintRationalField()
end

function prime_field(F::FqNmodFiniteField; cached::Bool = true)
  return GF(Int(characteristic(F)), cached = cached)
end

function prime_field(F::FqFiniteField; cached::Bool = true)
  return GF(characteristic(F), cached = cached)
end

function prime_field(F::T; cached::Bool = true) where T <: Union{GaloisField, GaloisFmpzField}
  return F
end

function Nemo.hom(F::FinField, K::FinField, a::FinFieldElem; check::Bool = true)
  @assert parent(a) == K

  if check
    @assert iszero(defining_polynomial(F)(a))
  end

  #We need a preimage function
  p = characteristic(K)
  Kp = prime_field(K, cached = false)
  Kpx = PolynomialRing(Kp, "x", cached = false)[1]
  Fp = prime_field(F, cached = false)
  Fpx = PolynomialRing(Fp, "x", cached = false)[1]
  M = zero_matrix(Kp, degree(F), degree(K))
  el = one(K)
  M[1, 1] = one(Kp)
  for i = 2:nrows(M)
    el = mul!(el, el, a)
    for j = 1:ncols(M)
      M[i, j] = coeff(el, j-1)
    end
  end

  aux = zero_matrix(Kp, 1, degree(F))
  aux1 = zero_matrix(Kp, 1, degree(K))
  function img(x::FinFieldElem)
    @assert parent(x) == F
    for i = 1:degree(F)
      aux[1, i] = Kp(coeff(x, i-1))
    end
    mul!(aux1, aux, M)
    pol = Kpx(elem_type(Kp)[aux1[1, j] for j = 1:degree(K)])
    return K(pol)
  end

  function preimg(x::FinFieldElem)
    @assert parent(x) == K
    for i = 1:degree(K)
      aux1[1, i] = Kp(coeff(x, i-1))
    end
    fl, y = can_solve_with_solution(M, aux1, side = :left)
    if !fl
      error("The element is not in the image!")
    end
    polF = Fpx(elem_type(Fp)[Fp(lift(y[1, j])) for j = 1:degree(F)])
    return F(polF)
  end

  return Nemo.FinFieldMorphism(F, K, img, preimg)
end

#Function that finds an embedding of k into K
function Nemo.embed_any(k::FinField, K::FinField)
  f = defining_polynomial(k)
  rt = roots(f, K)[1]
  return hom(k, K, rt)
end

function tr(a::FinFieldElem, mp::Nemo.FinFieldMorphism)
  @assert parent(a) == codomain(mp)
  q = order(domain(mp))
  d = divexact(degree(codomain(mp)), degree(domain(mp)))
  el = a
  aq = a
  for i = 1:d-1
    aq = aq^q
    el += aq
  end
  return mp\(el)
end

function norm(a::FinFieldElem, mp::Nemo.FinFieldMorphism)
  @assert parent(a) == codomain(mp)
  q = order(domain(mp))
  d = divexact(degree(codomain(mp)), degree(domain(mp)))
  el = a
  aq = a
  for i = 1:d-1
    aq = aq^q
    el *= aq
  end
  return mp\(el)
end

#TODO: compute via linear algebra.
function minpoly(a::FinFieldElem, mp::Nemo.FinFieldMorphism)
  conjs = Set{typeof(a)}()
  q = order(domain(mp))
  push!(conjs, a)
  el = a^q
  while !(el in conjs)
    push!(conjs, el)
    el = el^q
  end
  t = PolynomialRing(codomain(mp), "x", cached = false)[2]
  pol = prod(t-x for x in conjs)
  return map_coeffs(preimage_map(mp), pol)
end

function Nemo.generator_minimal_polynomial(mp::Nemo.FinFieldMorphism{FqFiniteField, FqFiniteField})
  return minpoly(gen(codomain(mp)), mp)
end

function Nemo.any_root(f::PolyElem{T}) where T <: FinFieldElem
  return roots(f)[1]
end


using FractalTools
using Plots

# Define function
f(x) = 4 - (x - 2)^2 + 1

# Interpolation points
xi, dx, xf = 0, 1, 4
x = collect(xi : dx : xf)
y = f.(x)
d = 0.5 * ones(length(x) - 1)

# Construct interpolation
interp = interpolate(x, y, d)
ws = interp.ifs.ws

# End points
xc = [0, 4, 4, 0, 0]
yc = [0, 0, 4, 4, 0]
pc = [xc yc]

for niter in 1 : niters
    xt = collect(xi : 0.1dx : xf)
    plt = plot(xt, f.(xt))
    pc = vcat([mapslices(wt, pc, dims=2) for wt in ws]...)
end

function transform_plot()
    # plt = plot(getindex.(pc, 1), getindex.(pc,2), marker=(:circle, 3))
    # xt = collect(xi : 0.1dx : xf)
    # plot!(xt, f.(xt))
    for wt in ws
        wp = []
        for (xi, yi) in zip(xc, yc)
            push!(wp, wt([xi, yi]))
        end
        plot!(getindex.(wp, 1), getindex.(wp,2), marker=(:circle, 3), subplot=1)
    end
    plt
end

transform_plot()


using GLPK
using JuMP
function scan_maker(A)
    m= Model(GLPK.Optimizer)

    level = size(A, 2)
    v = zeros(Int, level)
    ub = zeros(Int, level)
    lb = zeros(Int, level)

    @variable(m, x[1:level])
    @constraint(m, con, A*x.>=0)

    function setc(a)
        set_normalized_rhs.(con, float.(a))
    end

    function scan(c::Channel)
        i = 1
        init = 1
        while i > 0
            if i >= init
                @objective(m, Max, x[i])
                res = JuMP.optimize!(m)
                if termination_status(m) == MOI.OPTIMAL || termination_status(m) == MOI.DUAL_INFEASIBLE
                    ub[i] = round(Int, getvalue(x[i]))
                    set_objective_sense(m, MOI.MIN_SENSE)
                    res = JuMP.optimize!(m)
                    @assert termination_status(m) == MOI.OPTIMAL || termination_status(m) == MOI.DUAL_INFEASIBLE
                    lb[i] = round(Int, getvalue(x[i]))

                    v[i] = lb[i]
                    init += 1
                else
                    @assert termination_status(m) == MOI.INFEASIBLE
                    i -= 1
                    continue
                end
            elseif v[i] < ub[i]
                v[i] += 1
            else
                set_upper_bound(x[i], Inf)
                set_lower_bound(x[i], -Inf)
                init -= 1
                i -= 1
                continue
            end

            if i >= level
                put!(c, v)
                continue
            else
                set_upper_bound(x[i], v[i])
                set_lower_bound(x[i], v[i])
                i += 1
            end
        end
        close(c)
    end

    return setc, scan
end


using Revise
using AdFem
using PyPlot
using LinearAlgebra
using Statistics
using MAT
function f(x, y)
        1/(1+x^2) + x * y + y^2
    end

mmesh = Mesh(50, 50, 1/50, degree=2)

left = bcnode((x,y)->x<1e-5, mmesh)
right = bcedge((x1,y1,x2,y2)->(x1>0.049-1e-5) && (x2>0.049-1e-5), mmesh)

t1 = eval_f_on_boundary_edge((x,y)->1.0e-4, right, mmesh)
t2 = eval_f_on_boundary_edge((x,y)->0.0, right, mmesh)
rhs = compute_fem_traction_term(t1, t2, right, mmesh)

nu = 0.3
x = gauss_nodes(mmesh)
E = Variable(ones(get_ngauss(mmesh)))
# E = constant(eval_f_on_gauss_pts(f, mmesh))

D = compute_plane_stress_matrix(E, nu*ones(get_ngauss(mmesh)))
K = compute_fem_stiffness_matrix(D, mmesh)

bdval = [eval_f_on_boundary_node((x,y)->0.0, left, mmesh);
        eval_f_on_boundary_node((x,y)->0.0, left, mmesh)]
DOF = [left;left .+ mmesh.ndof]
K, rhs = impose_Dirichlet_boundary_conditions(K, rhs, DOF, bdval)
u = K\rhs
U = matread("fenics/data2.mat")["u"]

using Random; Random.seed!(233)
# idx = rand(1:mmesh.ndof, 1000)
idx = 1:mmesh.ndof
idx = [idx; idx .+ mmesh.ndof]

loss = sum((U[idx] - u[idx])^2) * 1e10
sess = Session(); init(sess)

_loss = Float64[]
cb = (vs, iter, loss)->begin
    global _loss
    push!(_loss, loss)
    printstyled("[#iter $iter] loss=$loss\n", color=:green)
    nu = vs[1]
    close("all")
    visualize_scalar_on_gauss_points(nu, mmesh)
    if mod(iter, 10)==1
        matwrite("fenics/data3-inverse$iter.mat", Dict("iter"=>iter,"loss"=>_loss, "E"=>nu))
        savefig("fenics/inverse3_pixel$iter.png")
    end
end

# run(sess, loss)
loss_ = BFGS!(sess, loss, vars = [E], callback = cb)


@testset "correlation feature selection" begin
    # 1. Test without k and threshold
    boston = dataset("MASS", "Boston")
    selector = UnivariateFeatureSelector(method = pearson_correlation)
    X = boston[:, Not(:MedV)]
    y = boston.MedV
    selected_features_all = select_features(selector, X, y)
    feature_scores_all = calculate_feature_importance(selector.method, X, y)
    expected = [
        (:LStat, -0.7376627261740151),
        (:Rm, 0.6953599470715394),
        (:PTRatio, -0.507786685537562),
        (:Indus, -0.4837251600283729),
        (:Tax, -0.46853593356776707),
        (:NOx, -0.4273207723732826),
        (:Crim, -0.38830460858681154),
        (:Rad, -0.38162623063977785),
        (:Age, -0.37695456500459634),
        (:Zn, 0.3604453424505432),
        (:Black, 0.3334608196570664),
        (:Dis, 0.24992873408590394),
        (:Chas, 0.1752601771902984),
    ]
    @test Symbol.(selected_features_all) == first.(expected)
    # Using isapprox because the result with Matrix vs Array is not the
    # exact same due to floating point computation. More detail:
    # https://github.com/JuliaLang/Statistics.jl/issues/24
    @test all(sort(collect(values(feature_scores_all)), by=abs, rev=true) .≈ last.(expected))

    # 2. Test with only k
    selector.k = 5
    selected_features = select_features(selector, X, y)
    @test selected_features == selected_features_all[1:5]

    # 3. Test with only threshold
    selector.k = nothing
    selector.threshold = 0.5
    selected_features_threshold = select_features(selector, X, y)
    feature_scores = calculate_feature_importance(selector.method, X, y)
    selected_features_score = filter(e->e[1] in selected_features_threshold, feature_scores)
    @test all(i -> abs(i) >= 0.5, values(selected_features_score))

    # 4. Test with both
    selector.k = 2
    selector.threshold = 0.5
    selected_features = select_features(selector, X, y)
    @test selected_features == selected_features_threshold[1:2]

    # 5. Test warn when setting k = 0
    selector.k = 0
    @test_logs(
        (:warn, r"k cannot be less than 1. Resetting value to 1"),
        select_features(selector, X, y)
    )

end

@testset "f-test feature selection" begin
    # 1. Test without k and threshold
    iris = dataset("datasets", "iris")
    selector = UnivariateFeatureSelector(method = f_test)
    X = iris[:, Not(:Species)]
    y = Vector{Int64}(replace(
        iris.Species,
        "setosa" => 1,
        "versicolor" => 2,
        "virginica" => 3,
    ))
    selected_features_all = select_features(selector, X, y)
    feature_scores_all = calculate_feature_importance(selector.method, X, y)
    expected = [
        (:PetalLength, 0.0),
        (:SepalWidth, 8.948397578478762e-14),
        (:PetalWidth, 0.37975655685608567),
        (:SepalLength, 0.8960092318703157),
    ]
    @test Symbol.(selected_features_all) == first.(expected)
    @test all(sort(collect(values(feature_scores_all)), by=abs) .≈ last.(expected))

    # Test also the p-values are not affected when target encoding swapped
    y = Vector{Int64}(replace(
        iris.Species,
        "setosa" => 3,
        "versicolor" => 1,
        "virginica" => 2,
    ))
    selected_features_all = select_features(selector, X, y)
    feature_scores_all = calculate_feature_importance(selector.method, X, y)
    @test Symbol.(selected_features_all) == first.(expected)
    @test all(sort(collect(values(feature_scores_all)), by=abs) .≈ last.(expected))

    # 2. Test with only k
    selector.k = 2
    selected_features = select_features(selector, X, y)
    @test selected_features == selected_features_all[1:2]

    # 3. Test with only threshold
    selector.k = nothing
    selector.threshold = 0.5
    selected_features_threshold = select_features(selector, X, y)
    feature_scores = calculate_feature_importance(selector.method, X, y)
    selected_features_score = filter(e->e[1] in selected_features, feature_scores)
    @test all(i -> abs(i) <= 0.5, values(selected_features_score))

    # 4. Test with both
    selector.k = 2
    selector.threshold = 0.5
    selected_features = select_features(selector, X, y)
    @test selected_features == selected_features_threshold[1:2]
end

@testset "chi-square feature selection" begin
    # 1. Test without k and threshold
    biopsy = dataset("MASS", "biopsy")[1:150, :] # Only use 150 data for test purpose
    selector = UnivariateFeatureSelector(method = chisq_test)
    # One-hot encode and only use 3 features
    X = one_hot_encode(biopsy[:, [:V1, :V2, :V3]]; drop_original = true)
    y = Vector{Int64}(replace(biopsy.Class, "benign" => 1, "malignant" => 2))
    selected_features_all = select_features(selector, X, y)
    feature_scores_all = calculate_feature_importance(selector.method, X, y)
    expected = [
        (:V2_1, 1.2657890724845934e-23),
        (:V3_1, 6.420486028446245e-22),
        (:V1_10, 5.522053763061195e-7),
        (:V1_1, 2.501514629267585e-6),
        (:V3_5, 3.034100127497124e-6),
        (:V2_10, 1.6011325170677212e-5),
        (:V2_5, 3.6277160399500314e-5),
        (:V3_10, 0.00018174534852021112),
        (:V1_8, 0.000940161869807205),
        (:V2_3, 0.0021417003254380123),
        (:V1_4, 0.004106630359262986),
        (:V1_7, 0.004251672172049395),
        (:V1_9, 0.004251672172049395),
        (:V3_7, 0.004251672172049395),
        (:V2_4, 0.004446485982989494),
        (:V3_6, 0.004446485982989494),
        (:V1_2, 0.005279083487939165),
        (:V3_4, 0.007455684495551331),
        (:V2_6, 0.009569951674053072),
        (:V2_7, 0.020438140781459755),
        (:V1_3, 0.05154325403026049),
        (:V3_8, 0.0924058047659373),
        (:V1_5, 0.10639608980657496),
        (:V3_3, 0.1697251182676055),
        (:V3_9, 0.25122678139322957),
        (:V3_2, 0.26413900085737735),
        (:V2_8, 0.4100233417486837),
        (:V2_2, 0.6279678908602435),
        (:V1_6, 0.695728316971443),
    ]
    @test Symbol.(selected_features_all) == first.(expected)
    @test all(sort(collect(values(feature_scores_all)), by=abs) .≈ last.(expected))

    # Test also the p-values are not affected when target encoding changed
    y = Vector{Int64}(replace(biopsy.Class, "benign" => 0, "malignant" => 1))
    selected_features_all = select_features(selector, X, y)
    feature_scores_all = calculate_feature_importance(selector.method, X, y)
    @test Symbol.(selected_features_all) == first.(expected)
    @test all(sort(collect(values(feature_scores_all)), by=abs) .≈ last.(expected))

    # 2. Test with only k
    selector.k = 2
    selected_features = select_features(selector, X, y)
    @test selected_features == selected_features_all[1:2]

    # 3. Test with only threshold
    selector.k = nothing
    selector.threshold = 0.5
    selected_features_threshold = select_features(selector, X, y)
    feature_scores = calculate_feature_importance(selector.method, X, y)
    selected_features_score = filter(e->e[1] in selected_features, feature_scores)
    @test all(i -> abs(i) <= 0.5, values(selected_features_score))

    # 4. Test with both
    selector.k = 2
    selector.threshold = 0.5
    selected_features = select_features(selector, X, y)
    @test selected_features == selected_features_threshold[1:2]
end


function nnls(A::AbstractMatrix, b::AbstractVector, ops::Options)
    # Set up the preconditioner as the inverse diagonal of the Hessian
    a = sum(A.^2, 1)
    @defaults ops precondprep=(P, x, l, u, mu)->precondprepnnls(P, x, mu, a)
    ops = copy(ops)
    @set_options ops precondprep=precondprep
    # Create the initial guess (an interior point)
    T = promote_type(eltype(A), eltype(b))
    x = fill(one(T), size(A, 2))
    # Set up constraints
    l = zeros(eltype(x), length(x))
    u = fill(inf(eltype(x)), length(x))
    # Perform the optimization
    func = (g,x) -> nnlsobjective(g, x, A, b)
    fminbox(func, x, l, u, ops)
end
nnls(A::AbstractMatrix, b::AbstractVector) = nnls(A, b, Options())

function nnlsobjective(g, x::AbstractVector, A::AbstractMatrix, b::AbstractVector)
    d = A*x - b
    val = sum(d.^2)/2
    if !(g === nothing)
        At_mul_B(g, A, d)
    end
    val
end

function precondprepnnls(P, x, mu, a)
    for i = 1:length(x)
        P[i] = 1/(mu/x[i]^2 + a[i])
    end
end

export nnls


grid = [
	 [1/24, 1/24, 1/24, 1/24, 1/24]
	,[1/24, 0.0, 0.0, 1/24, 1/24]
	,[1/24, 0.0, 1/24, 1/24, 1/24]
	,[1/24, 0.0, 0.0, 1/24, 1/24]
	,[1/24, 0.0, 1/24, 1/24, 1/24]
	,[1/24, 1/24, 1/24, 1/24, 1/24]

	]

const HEIGHT=6
const WIDTH=5

@enum SquareType OPEN CLOSED
@enum Direction WEST NORTH EAST SOUTH
@enum DriftType STRAIGHT LEFT RIGHT

const Drift=Dict(STRAIGHT=>.7, LEFT=>.15, RIGHT=>.15)
const AllDirects =(WEST, NORTH, EAST, SOUTH)

# Actual => Evidence
const Sense=Dict(OPEN=>Dict(OPEN=>.8, CLOSED=>.2), CLOSED=>Dict(OPEN=>.25, CLOSED=>.75))

function print_grid(grid::Array{Array{Float64,1},1})
	for i = 1:HEIGHT
		for j = 1:WIDTH
			if(!notblocked(grid, (i,j)))
				print("####","\t")
			else
				print(round(grid[i][j]*100, digits=2), "\t")
		   	end
		end
		println()
	end
end

# Check if the posistion is blocked (either out of boudns/obstacles) or a free space
function notblocked(grid::Array{Array{Float64,1},1}, pos::Tuple{Int64, Int64})
	if( (pos[1]>0 && pos[1]<=HEIGHT) && (pos[2]>0 && pos[2]<=WIDTH) )
		return (grid[pos[1]][pos[2]] != 0.0)
	end
	return false
end

# Given direction adn posistions return psostion that is pos+direction
function move(pos::Tuple{Int64, Int64}, dir::Direction)

	if(dir==WEST)
		return (pos[1], pos[2]-1)
	end

	if(dir==NORTH)
		return (pos[1]-1, pos[2])
	end

	if(dir==EAST)
		return (pos[1], pos[2]+1)
	end

	if(dir==SOUTH)
		return (pos[1]+1, pos[2])
	end
end

# Given posistion and evidecne return probabilty of being in that posistion
# PRECONDITION: Grid, position and evidecne
# POSTCONDITION: Conditonal Evidence Probabilty at that point
function evidence_Probabilty(grid::Array{Array{Float64,1},1},
							pos::Tuple{Int64, Int64},
							evidence::Tuple{SquareType, SquareType, SquareType, SquareType})
	prod = 1
	for i in 1:4
		tmp_pos = move(pos,AllDirects[i])
		block = notblocked(grid, tmp_pos)
		if (block)
			prod*= Sense[OPEN][evidence[i]]
		else
			prod*= Sense[CLOSED][evidence[i]]
		end
	end
	prod
end

# Get the evidecne contional probabilty of each posistoin*Pos(s_i)
# Then divide each posistion with the evidnece conditonal probabilty
# PRECONDITION: Given grid and evidence
# POSTCONDITION: Return filter of the grid, usign conditonal evdiecne probabilty
function filter(grid::Array{Array{Float64,1},1}, evidence::Tuple{SquareType, SquareType, SquareType, SquareType})
	tmp_grid = deepcopy(grid)
	for row in 1:6
		for col in 1:5
			if(notblocked(grid, (row,col)))
				tmp_grid[row][col]*=evidence_Probabilty(grid, (row, col), evidence)
			end
		end
	end
	total_sum = sum(sum(tmp_grid))
	tmp_grid / total_sum
end

# Gets the transitional probabilty of posistions going into a point,
# or transitional probablites from the point
# the first is used for prediction, the other for smoothing
# PRECONDITION: Given grid, position, a direction and type of transtional probabitly
# POSTCONDITION: Return probabilty of points going into pos, or probabitly of points going away from pos.
function transprob( grid::Array{Array{Float64,1},1}, pos::Tuple{Int64, Int64}, dir::Direction, getforward=false) # => array of  (pos::Tuple{Int64, Int64}, Drift[Direction] * P(s), Drift[Direction])
	arr = []
	function _gen_parts(straightDir::Direction ,
						leftDir::Direction,
						rightDir::Direction,
						behindDir::Direction,
						grid::Array{Array{Float64,1},1},
						pos::Tuple{Int64,Int64} )
		parent_pos = []
		behind = move(pos, behindDir)
		straight = move(pos, straightDir)
		left = move(pos, leftDir)
		right = move(pos, rightDir)

		if( !notblocked(grid, straight)) # Bounce
			push!(parent_pos, (pos,Drift[STRAIGHT]*grid[pos[1]][pos[2]], Drift[STRAIGHT] ))
		elseif (getforward) # If this is smoothing, want the probabilty in front
			push!(parent_pos, (straight,Drift[STRAIGHT]*grid[straight[1]][straight[2]], Drift[STRAIGHT] ))
		end
		if( notblocked(grid, behind) && !getforward) # Prob of square behind current ot move to current
			push!(parent_pos, (behind,Drift[STRAIGHT]*grid[behind[1]][behind[2]], Drift[STRAIGHT]))
		end
		if(notblocked(grid, left)) # Get probabitly of the left pos coming to curretn
			push!(parent_pos, (left,Drift[LEFT]*grid[left[1]][left[2]], Drift[LEFT] ))
		else #Bounce from left
			push!(parent_pos, (pos,Drift[LEFT]*grid[pos[1]][pos[2]], Drift[LEFT] ))
		end
		if(notblocked(grid, right)) #Get probabilty of right pos coming to current
			push!(parent_pos, (right,Drift[RIGHT]*grid[right[1]][right[2]],Drift[RIGHT]))
		else #Bounce
			push!(parent_pos, (pos,Drift[RIGHT]*grid[pos[1]][pos[2]], Drift[RIGHT]))
		end
	end

	if(dir==WEST)
		arr=_gen_parts(WEST, SOUTH, NORTH, EAST, grid, pos)
	elseif(dir==NORTH)
		arr=_gen_parts(NORTH, EAST, WEST, SOUTH, grid, pos)
	elseif(dir==SOUTH)
		arr=_gen_parts(SOUTH, EAST, WEST, NORTH, grid, pos)
	elseif(dir==EAST)
		arr=_gen_parts(EAST, SOUTH, NORTH, WEST, grid, pos)
	end

	arr
end

# Prediction is the sum of possiple transitional probabilties
# that can reach pos_i, * P(pos).j
# PRECONDITON: Given a grid and direction
# POSTCONDITION: Prediction of where the agent will go after this movement.
function predict(grid::Array{Array{Float64,1}}, dir::Direction)
	tmp_grid = deepcopy(grid)
	for row in 1:6
		for col in 1:5
			if(notblocked(grid, (row,col)))
				x= transprob(grid, (row, col), dir)
				val=sum([x[2] for x in transprob(grid, (row, col), dir)])
				tmp_grid[row][col]=val
			end
		end
	end

	tmp_grid
end

# Get the transitional probabilty of a point going OUT, not in
# an it's conditional probabilty, with it's inital probabilty
# returns 2 things: B at pos, and B*p(s)
# PRECONDITON: Grid
function smoothpart( grid::Array{Array{Float64,1}},
					Bgrid::Array{Array{Float64,1}},
					evidence::Tuple{SquareType, SquareType, SquareType, SquareType},
					dir::Direction, pos::Tuple{Int64, Int64})

	parent_pos=transprob(grid, pos, dir, true)
         x=(sum([evidence_Probabilty(grid, i[1], evidence)* Bgrid[i[1][1]][i[1][2]]* i[3] for i in parent_pos]))
	(x, x * grid[pos[1]][pos[2]])
end

# Get the smoothing part for each posistion in grid
# Then divide the whoel grid by the sum of it's parts
function smooth(                grid::Array{Array{Float64,1}},
				Bgrid::Array{Array{Float64,1}},
				evidence::Tuple{SquareType, SquareType, SquareType, SquareType},
				dir::Direction)

	NewGrid = deepcopy(grid)
	B = deepcopy(Bgrid)

	for row in 1:6
		for col in 1:5
			if(notblocked(grid, (row,col)))
				val=smoothpart(grid, Bgrid, evidence, dir, (row,col))
				B[row][col] = val[1]
				NewGrid[row][col] = val[2]
			else
				B[row][col]  = 0
				NewGrid[row][col] = 0
			end
		end
	end
	NewGrid/=sum(sum(NewGrid))
	(NewGrid, B)
end

tmp  = filter(grid, (OPEN, OPEN, OPEN, OPEN))
tmp2 = predict(tmp, WEST)
tmp3 = filter(tmp2, (CLOSED, CLOSED, OPEN, CLOSED))
tmp4 = predict(tmp3, NORTH)
tmp5 = filter(tmp4, (CLOSED, CLOSED, OPEN, CLOSED))

Bgrid = [
	 [1.0, 1.0, 1.0, 1.0, 1.0]
	 ,[1.0, 1.0, 1.0, 1.0, 1.0]
	 ,[1.0, 1.0, 1.0, 1.0, 1.0]
	 ,[1.0, 1.0, 1.0, 1.0, 1.0]
	 ,[1.0, 1.0, 1.0, 1.0, 1.0]
	 ,[1.0, 1.0, 1.0, 1.0, 1.0]

	 ]

(tmp6,b)=smooth(tmp3,Bgrid,(CLOSED,CLOSED,OPEN,CLOSED),NORTH)
(tmp7,c)=smooth(tmp,b,(CLOSED,CLOSED,OPEN,CLOSED),WEST)

println("Initial Location Probabilities")
print_grid(grid)
println("\nFiltering after Evidence [0, 0, 0, 0]")
print_grid(tmp)
println("\nPrediction after Action W")
print_grid(tmp2)
println("\nFiltering after Evidence [1, 1, 0, 1]")
print_grid(tmp3)
println("\nPrediction after Action N")
print_grid(tmp4)
println("\nFiltering after Evidence [1, 1, 0, 1]")
print_grid(tmp5)
println("\nLast position Smoothing with Evidence [1, 1, 0, 1] and north")
print_grid(tmp6)
println("\nSecond Last posistion smoothing with Evidence [1, 1, 0, 1] And west")
print_grid(tmp7)