using DiffEqSensitivity, OrdinaryDiffEq, RecursiveArrayTools, DiffEqBase,
      ForwardDiff, Calculus, QuadGK, LinearAlgebra, Zygote
using Test

function fb(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]*t
  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end
function foop(u,p,t)
  dx = p[1]*u[1] - p[2]*u[1]*u[2]*t
  dy = -p[3]*u[2] + p[4]*u[1]*u[2]
  [dx,dy]
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 * t
  J[2,1] = y
  J[1,2] = b * x * -1 * t
  J[2,2] = c * -1 + x
end

f = ODEFunction(fb,jac=jac)
p = [1.5,1.0,3.0,1.0]; u0 = [1.0;1.0]
prob = ODEProblem(f,u0,(0.0,10.0),p)
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
probb = ODEProblem(fb,u0,(0.0,10.0),p)
proboop = ODEProblem(foop,u0,(0.0,10.0),p)

solb = solve(probb,Tsit5(),abstol=1e-14,reltol=1e-14)
sol_end = solve(probb,Tsit5(),abstol=1e-14,reltol=1e-14,
          save_everystep=false,save_start=false)

sol_nodense = solve(probb,Tsit5(),abstol=1e-14,reltol=1e-14,dense=false)
soloop = solve(proboop,Tsit5(),abstol=1e-14,reltol=1e-14)
soloop_nodense = solve(proboop,Tsit5(),abstol=1e-14,reltol=1e-14,dense=false)

# Do a discrete adjoint problem
println("Calculate discrete adjoint sensitivities")
t = 0.0:0.5:10.0
# g(t,u,i) = (1-u)^2/2, L2 away from 1
function dg(out,u,p,t,i)
  (out.=2.0.-u)
end

_,easy_res = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14)
_,easy_res2 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14,
                                 sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14))
_,easy_res22 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=QuadratureAdjoint(autojacvec=false,abstol=1e-14,reltol=1e-14))
_,easy_res23 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14,compile=true))
_,easy_res3 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint())
_,easy_res32 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint(autojacvec=false))
_,easy_res4 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=BacksolveAdjoint())
_,easy_res42 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint(autojacvec=false))
_,easy_res43 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint(autojacvec=false,checkpointing=false))
_,easy_res5 = adjoint_sensitivities(sol,Kvaerno5(nlsolve=NLAnderson(), smooth_est=false),
                                 dg,t,abstol=1e-12,
                                 reltol=1e-10,
                                 sensealg=BacksolveAdjoint())
_,easy_res6 = adjoint_sensitivities(sol_nodense,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint(checkpointing=true),
                                  checkpoints=sol.t[1:500:end])
_,easy_res62 = adjoint_sensitivities(sol_nodense,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(checkpointing=true,autojacvec=false),
                                   checkpoints=sol.t[1:500:end])

# It should automatically be checkpointing since the solution isn't dense
_,easy_res7 = adjoint_sensitivities(sol_nodense,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint(),
                                  checkpoints=sol.t[1:500:end])

_,easy_res8 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                    reltol=1e-14,
                                    sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.TrackerVJP()))
_,easy_res9 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                    reltol=1e-14,
                                    sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.ZygoteVJP()))
_,easy_res10 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                    reltol=1e-14,
                                    sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.ReverseDiffVJP())
                                    )
_,easy_res11 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                     reltol=1e-14,
                                     sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.ReverseDiffVJP(true))
                                     )

adj_prob = ODEAdjointProblem(sol,QuadratureAdjoint(abstol=1e-14,reltol=1e-14),dg,t)
adj_sol = solve(adj_prob,Tsit5(),abstol=1e-14,reltol=1e-14)
integrand = AdjointSensitivityIntegrand(sol,adj_sol,QuadratureAdjoint(abstol=1e-14,reltol=1e-14))
res,err = quadgk(integrand,0.0,10.0,atol=1e-14,rtol=1e-12)

@test isapprox(res, easy_res, rtol = 1e-10)
@test isapprox(res, easy_res2, rtol = 1e-10)
@test isapprox(res, easy_res22, rtol = 1e-10)
@test isapprox(res, easy_res23, rtol = 1e-10)
@test isapprox(res, easy_res3, rtol = 1e-10)
@test isapprox(res, easy_res32, rtol = 1e-10)
@test isapprox(res, easy_res4, rtol = 1e-10)
@test isapprox(res, easy_res42, rtol = 1e-10)
@test isapprox(res, easy_res43, rtol = 1e-10)
@test isapprox(res, easy_res5, rtol = 1e-7)
@test isapprox(res, easy_res6, rtol = 1e-9)
@test isapprox(res, easy_res62, rtol = 1e-9)
@test all(easy_res6 .== easy_res7)  # should be the same!
@test isapprox(res, easy_res8, rtol = 1e-9)
@test_broken isapprox(res, easy_res9, rtol = 1e-9)
@test isapprox(res, easy_res10, rtol = 1e-9)
@test isapprox(res, easy_res11, rtol = 1e-9)

println("OOP adjoint sensitivities ")

_,easy_res = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14)
_,easy_res2 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14,
                                 sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14))
@test_broken easy_res22 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=QuadratureAdjoint(autojacvec=false,abstol=1e-14,reltol=1e-14))[1] isa AbstractArray
_,easy_res2 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                    reltol=1e-14,
                                    sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14,compile=true))
_,easy_res3 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint())
@test easy_res32 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint(autojacvec=false))[1] isa AbstractArray
_,easy_res4 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=BacksolveAdjoint())
@test easy_res42 = adjoint_sensitivities(soloop,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint(autojacvec=false))[1] isa AbstractArray
_,easy_res5 = adjoint_sensitivities(soloop,Kvaerno5(nlsolve=NLAnderson(), smooth_est=false),
                                 dg,t,abstol=1e-12,
                                 reltol=1e-10,
                                 sensealg=BacksolveAdjoint())
_,easy_res6 = adjoint_sensitivities(soloop_nodense,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint(checkpointing=true),
                                  checkpoints=soloop_nodense.t[1:5:end])
@test_broken easy_res62 = adjoint_sensitivities(soloop_nodense,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(checkpointing=true,autojacvec=false),
                                   checkpoints=soloop_nodense.t[1:5:end])

_,easy_res8 = adjoint_sensitivities(soloop_nodense,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.TrackerVJP()))
_,easy_res9 = adjoint_sensitivities(soloop_nodense,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.ZygoteVJP()))
_,easy_res10 = adjoint_sensitivities(soloop_nodense,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.ReverseDiffVJP())
                                   )
_,easy_res11 = adjoint_sensitivities(soloop_nodense,Tsit5(),dg,t,abstol=1e-14,
                                     reltol=1e-14,
                                     sensealg=InterpolatingAdjoint(autojacvec=DiffEqSensitivity.ReverseDiffVJP(true))
                                     )

@test isapprox(res, easy_res, rtol = 1e-10)
@test isapprox(res, easy_res2, rtol = 1e-10)
@test isapprox(res, easy_res22, rtol = 1e-10)
@test isapprox(res, easy_res23, rtol = 1e-10)
@test isapprox(res, easy_res3, rtol = 1e-10)
@test isapprox(res, easy_res32, rtol = 1e-10)
@test isapprox(res, easy_res4, rtol = 1e-10)
@test isapprox(res, easy_res42, rtol = 1e-10)
@test isapprox(res, easy_res5, rtol = 1e-9)
@test isapprox(res, easy_res6, rtol = 1e-10)
@test isapprox(res, easy_res62, rtol = 1e-9)
@test isapprox(res, easy_res8, rtol = 1e-9)
@test isapprox(res, easy_res9, rtol = 1e-9)
@test isapprox(res, easy_res10, rtol = 1e-9)
@test isapprox(res, easy_res11, rtol = 1e-9)

println("Calculate adjoint sensitivities ")

_,easy_res8 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  save_everystep=false,save_start=false,
                                  sensealg=BacksolveAdjoint())
_,easy_res82 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                   reltol=1e-14,
                                   save_everystep=false,save_start=false,
                                   sensealg=BacksolveAdjoint(checkpointing=false))

@test isapprox(res, easy_res8, rtol = 1e-9)
@test isapprox(res, easy_res82, rtol = 1e-9)

_,end_only_res = adjoint_sensitivities(sol_end,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,
                                  save_everystep=false,save_start=false,
                                  sensealg=BacksolveAdjoint())

@test isapprox(res, end_only_res, rtol = 1e-9)

println("Calculate adjoint sensitivities from autodiff & numerical diff")
function G(p)
  tmp_prob = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
  sol = solve(tmp_prob,Tsit5(),abstol=1e-14,reltol=1e-14,sensealg=DiffEqBase.SensitivityADPassThrough(),saveat=t)
  A = Array(sol)
  sum(((2 .- A).^2)./2)
end
G([1.5,1.0,3.0,1.0])
res2 = ForwardDiff.gradient(G,[1.5,1.0,3.0,1.0])
res3 = Calculus.gradient(G,[1.5,1.0,3.0,1.0])

import Tracker
res4 = Tracker.gradient(G,[1.5,1.0,3.0,1.0])[1]

import ReverseDiff
res5 = ReverseDiff.gradient(G,[1.5,1.0,3.0,1.0])

@test norm(res' .- res2) < 1e-7
@test norm(res' .- res3) < 1e-5
@test norm(res' .- res4) < 1e-6
@test norm(res' .- res5) < 1e-6

# check other t handling

t2 = [0.5, 1.0]
t3 = [0.0, 0.5, 1.0]
t4 = [0.5, 1.0, 10.0]

_,easy_res2 = adjoint_sensitivities(sol,Tsit5(),dg,t2,abstol=1e-14,
                                  reltol=1e-14)
_,easy_res3 = adjoint_sensitivities(sol,Tsit5(),dg,t3,abstol=1e-14,
                                  reltol=1e-14)
_,easy_res4 = adjoint_sensitivities(sol,Tsit5(),dg,t4,abstol=1e-14,
                                  reltol=1e-14)

function G(p,ts)
  tmp_prob = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
  sol = solve(tmp_prob,Tsit5(),abstol=1e-10,reltol=1e-10,sensealg=DiffEqBase.SensitivityADPassThrough(),saveat=ts)
  A = convert(Array,sol)
  sum(((2 .- A).^2)./2)
end
res2 = ForwardDiff.gradient(p->G(p,t2),[1.5,1.0,3.0,1.0])
res3 = ForwardDiff.gradient(p->G(p,t3),[1.5,1.0,3.0,1.0])
res4 = ForwardDiff.gradient(p->G(p,t4),[1.5,1.0,3.0,1.0])

@test easy_res2' ≈ res2
@test easy_res3' ≈ res3
@test easy_res4' ≈ res4

println("Adjoints of u0")

function dg(out,u,p,t,i)
  out .= 1 .- u
end

ū0,adj = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                         reltol=1e-14)

_,adjnou0 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                        reltol=1e-14)

ū02,adj2 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=BacksolveAdjoint(),
                                    reltol=1e-14)

ū022,adj22 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=BacksolveAdjoint(autojacvec=false),
                                    reltol=1e-14)

ū023,adj23 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=BacksolveAdjoint(autojacvec=false,checkpointing=false),
                                    reltol=1e-14)

ū03,adj3 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=InterpolatingAdjoint(),
                                    reltol=1e-14)

ū032,adj32 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=InterpolatingAdjoint(autojacvec=false),
                                    reltol=1e-14)

ū04,adj4 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=InterpolatingAdjoint(checkpointing=true),
                                    checkpoints=sol.t[1:500:end],
                                    reltol=1e-14)

@test_nowarn adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                      sensealg=InterpolatingAdjoint(checkpointing=true),
                                      checkpoints=sol.t[1:5:end],
                                      reltol=1e-14)

ū042,adj42 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=InterpolatingAdjoint(checkpointing=true,autojacvec=false),
                                    checkpoints=sol.t[1:500:end],
                                    reltol=1e-14)

ū05,adj5 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14),
                                    reltol=1e-14)

ū052,adj52 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=QuadratureAdjoint(autojacvec=false,abstol=1e-14,reltol=1e-14),
                                    reltol=1e-14)

ū05,adj53 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                    sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14,compile=true),
                                    reltol=1e-14)

ū0args,adjargs = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                        save_everystep=false, save_start=false,
                        sensealg=BacksolveAdjoint(),
                        reltol=1e-14)

ū0args2,adjargs2 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                        save_everystep=false, save_start=false,
                        sensealg=InterpolatingAdjoint(),
                        reltol=1e-14)

res = ForwardDiff.gradient(prob.u0) do u0
  tmp_prob = remake(prob,u0=u0)
  sol = solve(tmp_prob,Tsit5(),abstol=1e-14,reltol=1e-14,saveat=t)
  A = convert(Array,sol)
  sum(((1 .- A).^2)./2)
end

@test ū0 ≈ res rtol = 1e-10
@test ū02 ≈ res rtol = 1e-10
@test ū022 ≈ res rtol = 1e-10
@test ū023 ≈ res rtol = 1e-10
@test ū03 ≈ res rtol = 1e-10
@test ū032 ≈ res rtol = 1e-10
@test ū04 ≈ res rtol = 1e-10
@test ū042 ≈ res rtol = 1e-10
@test ū05 ≈ res rtol = 1e-10
@test ū052 ≈ res rtol = 1e-10
@test adj ≈ adjnou0 rtol = 1e-10
@test adj ≈ adj2 rtol = 1e-10
@test adj ≈ adj22 rtol = 1e-10
@test adj ≈ adj23 rtol = 1e-10
@test adj ≈ adj3 rtol = 1e-10
@test adj ≈ adj32 rtol = 1e-10
@test adj ≈ adj4 rtol = 1e-10
@test adj ≈ adj42 rtol = 1e-10
@test adj ≈ adj5 rtol = 1e-10
@test adj ≈ adj52 rtol = 1e-10
@test adj ≈ adj53 rtol = 1e-10

@test ū0args ≈ res rtol = 1e-10
@test adjargs ≈ adj rtol = 1e-10
@test ū0args2 ≈ res rtol = 1e-10
@test adjargs2 ≈ adj rtol = 1e-10

println("Do a continuous adjoint problem")

# Energy calculation
g(u,p,t) = (sum(u).^2) ./ 2
# Gradient of (u1 + u2)^2 / 2
function dg(out,u,p,t)
  out[1]= u[1] + u[2]
  out[2]= u[1] + u[2]
end

adj_prob = ODEAdjointProblem(sol,QuadratureAdjoint(abstol=1e-14,reltol=1e-14),g,nothing,dg)
adj_sol = solve(adj_prob,Tsit5(),abstol=1e-14,reltol=1e-10)
integrand = AdjointSensitivityIntegrand(sol,adj_sol,QuadratureAdjoint(abstol=1e-14,reltol=1e-14))
res,err = quadgk(integrand,0.0,10.0,atol=1e-14,rtol=1e-10)

println("Test the `adjoint_sensitivities` utility function")
_,easy_res = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                 reltol=1e-14)
println("2")
_,easy_res2 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint())
_,easy_res22 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(autojacvec=false))
println("23")
_,easy_res23 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14))
_,easy_res232 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14,compile=false))
_,easy_res24 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=QuadratureAdjoint(autojacvec=false,abstol=1e-14,reltol=1e-14))
println("25")
_,easy_res25 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint())
_,easy_res26 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint(autojacvec=false))
_,easy_res262 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint(autojacvec=false,checkpointing=false))
println("27")
_,easy_res27 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                  reltol=1e-14,
                                  checkpoints=sol.t[1:500:end],
                                  sensealg=InterpolatingAdjoint(checkpointing=true))
_,easy_res28 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                  reltol=1e-14,
                                  checkpoints=sol.t[1:500:end],
                                  sensealg=InterpolatingAdjoint(checkpointing=true,autojacvec=false))
println("3")
_,easy_res3 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                  reltol=1e-14,
                                  sensealg=InterpolatingAdjoint())
_,easy_res32 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=InterpolatingAdjoint(autojacvec=false))
println("33")
_,easy_res33 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=QuadratureAdjoint(abstol=1e-14,reltol=1e-14))
_,easy_res34 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=QuadratureAdjoint(autojacvec=false,abstol=1e-14,reltol=1e-14))
println("35")
_,easy_res35 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint())
_,easy_res36 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                   reltol=1e-14,
                                   sensealg=BacksolveAdjoint(autojacvec=false))
println("37")
_,easy_res37 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                  reltol=1e-14,
                                  checkpoints=sol.t[1:500:end],
                                  sensealg=InterpolatingAdjoint(checkpointing=true))
_,easy_res38 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                  reltol=1e-14,
                                  checkpoints=sol.t[1:500:end],
                                  sensealg=InterpolatingAdjoint(checkpointing=true,autojacvec=false))

@test norm(easy_res .- res) < 1e-8
@test norm(easy_res2 .- res) < 1e-8
@test norm(easy_res22 .- res) < 1e-8
@test norm(easy_res23 .- res) < 1e-8
@test norm(easy_res232 .- res) < 1e-8
@test norm(easy_res24 .- res) < 1e-8
@test norm(easy_res25 .- res) < 1e-8
@test norm(easy_res26 .- res) < 1e-8
@test norm(easy_res262 .- res) < 1e-8
@test norm(easy_res27 .- res) < 1e-8
@test norm(easy_res28 .- res) < 1e-8
@test norm(easy_res3 .- res) < 1e-8
@test norm(easy_res32 .- res) < 1e-8
@test norm(easy_res33 .- res) < 1e-8
@test norm(easy_res34 .- res) < 1e-8
@test norm(easy_res35 .- res) < 1e-8
@test norm(easy_res36 .- res) < 1e-8
@test norm(easy_res37 .- res) < 1e-8
@test norm(easy_res38 .- res) < 1e-8

println("Calculate adjoint sensitivities from autodiff & numerical diff")
function G(p)
  tmp_prob = remake(prob,u0=eltype(p).(prob.u0),p=p,
                    tspan=eltype(p).(prob.tspan))
  sol = solve(tmp_prob,Tsit5(),abstol=1e-14,reltol=1e-14)
  res,err = quadgk((t)-> (sum(sol(t)).^2)./2,0.0,10.0,atol=1e-14,rtol=1e-10)
  res
end
res2 = ForwardDiff.gradient(G,[1.5,1.0,3.0,1.0])
res3 = Calculus.gradient(G,[1.5,1.0,3.0,1.0])

@test norm(res' .- res2) < 1e-8
@test norm(res' .- res3) < 1e-6

# Buffer length test
f = (du, u, p, t) -> du .= 0
p = zeros(3); u = zeros(50)
prob = ODEProblem(f,u,(0.0,10.0),p)
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
@test_nowarn _,res = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14)

@testset "Checkpointed backsolve" begin
  using DiffEqSensitivity, OrdinaryDiffEq
  tf = 10.0
  function lorenz(du,u,p,t)
    σ, ρ, β = p
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
    return nothing
  end
  prob_lorenz = ODEProblem(lorenz, [1.0, 0.0, 0.0], (0, tf), [10, 28, 8/3])
  sol_lorenz = solve(prob_lorenz,Tsit5(),reltol=1e-6,abstol=1e-9)
  function dg(out,u,p,t,i)
    (out.=2.0.-u)
  end
  t = 0:0.1:tf
  _,easy_res1 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                    reltol=1e-9,
                                    sensealg=BacksolveAdjoint())
  _,easy_res2 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                    reltol=1e-9,
                                    sensealg=InterpolatingAdjoint())
  _,easy_res3 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                    reltol=1e-9,
                                    sensealg=BacksolveAdjoint(),
                                    checkpoints=sol_lorenz.t[1:10:end])
  _,easy_res4 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                    reltol=1e-9,
                                    sensealg=BacksolveAdjoint(),
                                    checkpoints=sol_lorenz.t[1:20:end])
  # cannot finish in a reasonable amount of time
  @test_skip adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                   reltol=1e-9,
                                   sensealg=BacksolveAdjoint(checkpointing=false))
  @test easy_res2 ≈ easy_res1 rtol=1e-5
  @test easy_res2 ≈ easy_res3 rtol=1e-5
  @test easy_res2 ≈ easy_res4 rtol=1e-4

  ū1,adj1 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                     reltol=1e-9,
                                     sensealg=BacksolveAdjoint())
  ū2,adj2 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                     reltol=1e-9,
                                     sensealg=InterpolatingAdjoint())
  ū3,adj3 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                     reltol=1e-9,
                                     sensealg=BacksolveAdjoint(),
                                     checkpoints=sol_lorenz.t[1:10:end])
  ū4,adj4 = adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                     reltol=1e-9,
                                     sensealg=BacksolveAdjoint(),
                                     checkpoints=sol_lorenz.t[1:20:end])
  # cannot finish in a reasonable amount of time
  @test_skip adjoint_sensitivities(sol_lorenz,Tsit5(),dg,t,abstol=1e-6,
                                      reltol=1e-9,
                                      sensealg=BacksolveAdjoint(checkpointing=false))
  @test ū2 ≈ ū1 rtol=1e-5
  @test adj2 ≈ adj1 rtol=1e-5
  @test ū2 ≈ ū3 rtol=1e-5
  @test adj2 ≈ adj3 rtol=1e-5
  @test ū2 ≈ ū4 rtol=1e-4
  @test adj2 ≈ adj4 rtol=1e-4

  # LQR Tests from issue https://github.com/SciML/DiffEqSensitivity.jl/issues/300
  x_dim = 2
  T = 40.0

  cost = (x, u) -> x'*x
  params =  [-0.4142135623730951, 0.0, -0.0, -0.4142135623730951, 0.0, 0.0]

  function dynamics!(du,u,p,t)
      du[1] = -u[1] + tanh(p[1]*u[1]+p[2]*u[2])
      du[2] = -u[2] + tanh(p[3]*u[1]+p[4]*u[2])
  end

  function backsolve_grad(sol, lqr_params, checkpointing)
      bwd_sol = solve(
          ODEAdjointProblem(
              sol,
              BacksolveAdjoint(checkpointing = checkpointing),
              (x, lqr_params, t) -> cost(x,lqr_params),
          ),
          Tsit5(),
          dense = false,
          save_everystep = false,
      )

      bwd_sol.u[end][1:end-x_dim]
      #fwd_sol, bwd_sol
  end

  x0 = ones(x_dim)
  fwd_sol = solve(
      ODEProblem(dynamics!, x0, (0, T), params),
      Tsit5(),abstol=1e-9, reltol=1e-9,
      u0 = x0,
      p = params,
      dense = false,
      save_everystep = true
  )

  backsolve_results = backsolve_grad(fwd_sol, params, false)
  backsolve_checkpointing_results = backsolve_grad(fwd_sol, params, true)

  @test backsolve_results != backsolve_checkpointing_results

  int_u0, int_p = adjoint_sensitivities(fwd_sol,Tsit5(),(x, params, t)->cost(x,params), nothing, sensealg=InterpolatingAdjoint())

  @test isapprox(-backsolve_checkpointing_results[1:length(x0)], int_u0, rtol=1e-10)
  @test isapprox(backsolve_checkpointing_results[(1:length(params)) .+ length(x0)], int_p', rtol=1e-10)
end

using Test
using LinearAlgebra, DiffEqSensitivity, OrdinaryDiffEq, ForwardDiff, QuadGK
@testset "Adjoint of differential algebric equations with mass matrix" begin
  function G(p, prob, ts, cost)
    tmp_prob_mm = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
    sol = solve(tmp_prob_mm,Rodas5(autodiff=false),abstol=1e-14,reltol=1e-14,saveat=ts)
    cost(sol)
  end
  alg = Rodas5(autodiff=false)
  @testset "Fully ranked mass matrix" begin
    @info "discrete cost"
    A = [1 2 3; 4 5 6; 7 8 9]
    function foo(du, u, p, t)
      mul!(du, A, u)
      du .= du .+ p
      du[2] += sum(p)
      return nothing
    end
    mm = -[1 2 4; 2 3 7; 1 3 41]
    u0 = [1, 2.0, 3]
    p = [1.0, 2.0, 3]
    prob_mm = ODEProblem(ODEFunction(foo, mass_matrix=mm), u0, (0, 1.0), p)
    sol_mm = solve(prob_mm, Rodas5(), reltol=1e-14, abstol=1e-14)

    ts = 0:0.01:1
    dg(out,u,p,t,i) = out .= -1
    _, res = adjoint_sensitivities(sol_mm,alg,dg,ts,abstol=1e-5,reltol=1e-5,sensealg=QuadratureAdjoint())
    reference_sol = ForwardDiff.gradient(p->G(p, prob_mm, ts, sum),vec(p))
    @test res' ≈ reference_sol rtol=1e-3

    _, res_interp = adjoint_sensitivities(sol_mm,alg,dg,ts,abstol=1e-5,reltol=1e-5,sensealg=InterpolatingAdjoint())
    @test res_interp ≈ res rtol = 1e-2
    _, res_interp2 = adjoint_sensitivities(sol_mm,alg,dg,ts,abstol=1e-5,reltol=1e-5,sensealg=InterpolatingAdjoint(checkpointing=true),checkpoints=sol_mm.t[1:10:end])
    @test res_interp2 ≈ res rtol = 1e-2

    # backsolve doesn't work
    _, res_bs = adjoint_sensitivities(sol_mm,alg,dg,ts,abstol=1e-5,reltol=1e-5,sensealg=BacksolveAdjoint(checkpointing=false))
    @test res_bs ≈ res rtol = 1e-2
    _, res_bs2 = adjoint_sensitivities(sol_mm,alg,dg,ts,abstol=1e-5,reltol=1e-5,sensealg=BacksolveAdjoint(checkpointing=true),checkpoints=sol_mm.t)
    @test res_bs2 ≈ res rtol = 1e-2

    @info "continuous cost"
    g_cont(u,p,t) = (sum(u).^2) ./ 2
    dg_cont(out,u,p,t) = out .= sum(u)
    _,easy_res_cont = adjoint_sensitivities(sol_mm,alg,g_cont,nothing,
                                     dg_cont,abstol=1e-10,reltol=1e-10,
                                     sensealg=QuadratureAdjoint())
    function G_cont(p)
      tmp_prob_mm = remake(prob_mm,u0=eltype(p).(prob_mm.u0),p=p,
                        tspan=eltype(p).(prob_mm.tspan))
      sol = solve(tmp_prob_mm,Rodas5(autodiff=false),abstol=1e-14,reltol=1e-14)
      res,err = quadgk((t)-> (sum(sol(t)).^2)./2,prob_mm.tspan...,atol=1e-14,rtol=1e-10)
      res
    end
    reference_sol_cont = ForwardDiff.gradient(G_cont, p)
    @test easy_res_cont' ≈ reference_sol_cont rtol=1e-3
  end

  @testset "Singular mass matrix" begin
    function rober(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(u,p,t)
      y₁,y₂,y₃ = u
      k₁,k₂,k₃ = p
      return [-k₁*y₁+k₃*y₂*y₃,
       k₁*y₁-k₂*y₂^2-k₃*y₂*y₃,
       y₁ + y₂ + y₃ - 1]
    end
    M = [1. 0  0
         0  1. 0
         0  0  0]
    for iip in [true, false]
      f = ODEFunction{iip}(rober,mass_matrix=M)
      p = [0.04,3e7,1e4]

      prob_singular_mm = ODEProblem(f,[1.0,0.0,0.0],(0.0,100),p)
      sol_singular_mm = solve(prob_singular_mm,Rodas5(),reltol=1e-8,abstol=1e-8)
      ts = [50, sol_singular_mm.t[end]]
      dg_singular(out,u,p,t,i) = (fill!(out, 0); out[end] = -1)
      _, res = adjoint_sensitivities(sol_singular_mm,alg,dg_singular,ts,abstol=1e-5,reltol=1e-5,sensealg=QuadratureAdjoint())
      reference_sol = ForwardDiff.gradient(p->G(p, prob_singular_mm, ts, sol->sum(last, sol.u)), vec(p))
      @test res' ≈ reference_sol rtol = 1e-2

      _, res_interp = adjoint_sensitivities(sol_singular_mm,alg,dg_singular,ts,abstol=1e-5,reltol=1e-5,sensealg=InterpolatingAdjoint())
      @test res_interp ≈ res rtol = 1e-2
      _, res_interp2 = adjoint_sensitivities(sol_singular_mm,alg,dg_singular,ts,abstol=1e-5,reltol=1e-5,sensealg=InterpolatingAdjoint(checkpointing=true),checkpoints=sol_singular_mm.t[1:10:end])
      @test res_interp2 ≈ res rtol = 1e-2

      # backsolve doesn't work
      _, res_bs = adjoint_sensitivities(sol_singular_mm,alg,dg_singular,ts,abstol=1e-5,reltol=1e-5,sensealg=BacksolveAdjoint(checkpointing=false))
      @test_broken res_bs ≈ res rtol = 1e-2
      _, res_bs2 = adjoint_sensitivities(sol_singular_mm,alg,dg_singular,ts,abstol=1e-5,reltol=1e-5,sensealg=BacksolveAdjoint(checkpointing=true),checkpoints=sol_singular_mm.t)
      @test_broken res_bs2 ≈ res rtol = 1e-2
    end
  end
end


# __precompile__(true)

module Expectations

# Load external dependencies.
using Compat, FastGaussQuadrature
using Compat.LinearAlgebra
using SpecialFunctions
using Distributions

# Load internal files.
include("types.jl"),
include("iterable.jl")

# Export
export expectation, Expectation, IterableExpectation, Gaussian, Trapezoidal, FiniteDiscrete, QuadratureAlgorithm, ExplicitQuadratureAlgorithm, nodes, weights, QuantileRange

end # module


"""
Function to simulate a galaxy depending on type of parameters specified given the sample size.
"""
function simulate_galaxy(p::NFWParameters{G}, samplesize::T; x0::G = 10.0^-8, αt::G = 1.0, rate::Bool = false) where {G <: AbstractFloat, T <: Integer}
    function nfwopt(x::Array{G, 1})
        return -(x[1] * x[1] * x[3] * profile_density(x[1], x[2], x[3], p, αt = αt))
    end
    #function nfwopt(x::Array{G, 1})
    #    return -(2.0 * log(x[1]) +  log(x[3]) + log_profile_density(x[1], x[2], x[3], p, αt = αt))
    #end

    minProb = Optim.optimize(nfwopt, [0.1, 1.0, 1.0], method = Optim.NelderMead())
    fmax = max(0.0,  -1.1 * minProb.minimum)

    vmax0 = escape_velocity(x0, p)
    sampledr = Array{Float64}(samplesize)
    sampledvr = Array{Float64}(samplesize)
    sampledvt = Array{Float64}(samplesize)
    if rate
        tested = 0
    end
    accepted = 0
    while accepted < samplesize
        x = (p.rlim  / p.rs) * rand()
        vr = vmax0 * rand()
        vt = vmax0 * rand()
        u = rand()
        #u = log(rand())
        sampleDensity = x * x * vt * profile_density(x, vr, vt, p, αt = αt)
        #sampleDensity = 2.0 * log(x) + log(vt) + log_profile_density(x, vr, vt, p, αt = αt)
        if rate
            tested += 1
        end
        if sampleDensity >= (u * fmax)
            #if sampleDensity >= (u + fmax)
            accepted += 1
            sampledr[accepted] = x * p.rs
            sampledvr[accepted] =  vr
            sampledvt[accepted] = vt
        end
    end
    if rate
        return SphericalGalaxy(sampledr, sampledvr, sampledvt, samplesize), accepted / tested
    end
    return SphericalGalaxy(sampledr, sampledvr, sampledvt, samplesize)
end

function simulate_galaxy(p::SFWParameters{G}, samplesize::T; x0::G = 10.0^-8, αt::G = 1.0, rate::Bool = false) where {G <: AbstractFloat, T <: Integer}
    function sfwopt(x::Array{G, 1})
        return -(x[1] * x[1] * x[3] * profile_density(x[1], x[2], x[3], p, αt = αt))
    end

    minProb = Optim.optimize(sfwopt, [0.1, 1.0, 1.0], method = Optim.NelderMead())
    fmax = max(0.0,  -1.1 * minProb.minimum)

    vmax0 = escape_velocity(x0, p)
    sampledr = Array{Float64}(samplesize)
    sampledvr = Array{Float64}(samplesize)
    sampledvt = Array{Float64}(samplesize)
    if rate
        tested = 0
    end
    accepted = 0
    while accepted < samplesize
        x = (p.rlim  / p.rs) * rand()
        vr = vmax0 * rand()
        vt = vmax0 * rand()
        u = rand()
        sampleDensity = x * x * vt * profile_density(x, vr, vt, p, αt = αt)
        if rate
            tested += 1
        end
        if sampleDensity >= (u * fmax)
            accepted += 1
            sampledr[accepted] = x * p.rs
            sampledvr[accepted] =  vr
            sampledvt[accepted] = vt
        end
    end
    if rate
        return SphericalGalaxy(sampledr, sampledvr, sampledvt, samplesize), accepted / tested
    end
    return SphericalGalaxy(sampledr, sampledvr, sampledvt, samplesize)
end

function simulate_galaxy(p::NFWParameters{G}, mμ::G, mσ::G, samplesize::T; x0::G = 10.0^-8, αt::G = 1.0, rate::Bool = false) where {G <: AbstractFloat, T <: Integer}
    metal = randn(samplesize) .* mσ .+ mμ
    if rate
        sg, acptrt = simulate_galaxy(p, samplesize, x0 = x0, αt = αt, rate = rate)
        return MetallicSphericalGalaxy(sg, metal), acptrt
    else
        return MetallicSphericalGalaxy(simulate_galaxy(p, samplesize, x0 = x0, αt = αt, rate = rate), metal)
    end
end

function simulate_galaxy(p::SFWParameters{G}, mμ::G, mσ::G, samplesize::T; x0::G = 10.0^-8, αt::G = 1.0, rate::Bool = false) where {G <: AbstractFloat, T <: Integer}
    metal = randn(samplesize) .* mσ .+ mμ
    if rate
        sg, acptrt = simulate_galaxy(p, samplesize, x0 = x0, αt = αt, rate = rate)
        return MetallicSphericalGalaxy(sg, metal), acptrt
    else
        return MetallicSphericalGalaxy(simulate_galaxy(p, samplesize, x0 = x0, αt = αt, rate = rate), metal)
    end
end

"""
Use a grid rejection sampler instead of a standard rejection sampler
"""
function simulate_galaxy_grid(p::SFWParameters{G}, samplesize::T; x0::G = 10.0^-8, αt::G = 1.0, rate::Bool = false, ngrid::Integer = 3) where {G <: AbstractFloat, T <: Integer}

    vmax0 = escape_velocity(x0, p)

    function sfwopt(x::Array{G, 1})
        return -(x[1] * x[1] * x[3] * profile_density(x[1], x[2], x[3], p, αt = αt))
    end

    ##grid bounds
    xgrid = linspace(0.0, p.rlim  / p.rs, ngrid + 1)
    vrgrid = linspace(0.0, vmax0, ngrid + 1)
    vtgrid = linspace(0.0, vmax0, ngrid + 1)

    ##midpoints
    xmids = xgrid[1:ngrid] + 0.5 * step(xgrid)
    vrmids = vrgrid[1:ngrid] + 0.5 * step(vrgrid)
    vtmids = vtgrid[1:ngrid] + 0.5 * step(vtgrid)

    ##optimization variables
    startpoints = Array{Float64}(ngrid^3, 3)
    lowers = Array{Float64}(ngrid^3, 3)
    uppers = Array{Float64}(ngrid^3, 3)
    fmaxes = Array{Float64}(ngrid^3)
    ii = 1
    for xx in 1:ngrid, rr in 1:ngrid, tt in 1:ngrid
        startpoints[ii, :] = [xmids[xx], vrmids[rr], vtmids[tt]]
        lowers[ii, :] = [xgrid[xx], vrgrid[rr], vtgrid[tt]]
        uppers[ii, :] = [xgrid[xx + 1], vrgrid[rr + 1], vtgrid[tt + 1]]
        minProb = Optim.optimize(sfwopt, startpoints[ii, :], lowers[ii, :], uppers[ii, :], Fminbox{Optim.NelderMead}())
        fmaxes[ii] = max(0.0,  -1.1 * minProb.minimum)
        ii += 1
    end
    w = StatsBase.Weights(fmaxes)
    #http://julianlsolvers.github.io/Optim.jl/latest/user/minimization/#box-minimization

    sampledr = Array{Float64}(samplesize)
    sampledvr = Array{Float64}(samplesize)
    sampledvt = Array{Float64}(samplesize)
    if rate
        tested = 0
    end
    accepted = 0
    while accepted < samplesize
        idx = StatsBase.sample(w)
        x = lowers[ii, 1] + step(xgrid) * rand()
        vr = lowers[ii, 2] + step(vrgrid) * rand()
        vt = lowers[ii, 3] + step(vtgrid) * rand()
        u = rand()
        sampleDensity = x * x * vt * profile_density(x, vr, vt, p, αt = αt)
        if rate
            tested += 1
        end
        if sampleDensity >= (u * fmax)
            accepted += 1
            sampledr[accepted] = x * p.rs
            sampledvr[accepted] =  vr
            sampledvt[accepted] = vt
        end
    end
    if rate
        return SphericalGalaxy(sampledr, sampledvr, sampledvt, samplesize), accepted / tested
    end
    return SphericalGalaxy(sampledr, sampledvr, sampledvt, samplesize)
end


using Test
using SymArrays
using SymArrays: symarrlength, _sub2grp, which_symgrp
using TensorOperations
using Random
using CUDA
CUDA.allowscalar(false)

@testset "SymArrays.jl" begin
    @testset "SymArray" begin
        @test symarrlength((3,6,4,3),(3,2,1,3)) == 8400

        S = SymArray{(2,),Float64}(5,5)
        @test nsymgrps(S) == 1
        @test symgrps(S) == (2,)
        @test _sub2grp(S,2,5) == _sub2grp(S,5,2)
        @test _sub2grp(S,3,4) != _sub2grp(S,5,3)
        # _sub2grp has to have same number of arguments as size of N
        @test_throws MethodError _sub2grp(S,3,5,1)
        # indexing the array allows having additional 1s at the end
        @test S[3,5,1] == S[3,5]
        @test_throws BoundsError S[3,5,3]
        # this calculation gives an index for data that is within bounds, but should not be legal
        # make sure this is caught
        @test_throws BoundsError S[0,6]
        @test (S[1,5] = 2.; S[1,5] == 2.)
        @test_throws BoundsError S[0,6] = 2.

        S = SymArray{(3,1,2,2),Float64}(3,3,3,2,4,4,4,4)
        @test nsymgrps(S) == 4
        @test symgrps(S) == (3,1,2,2)
        @test size(S) == (3,3,3,2,4,4,4,4)
        @test length(S) == 2000
        # iterating over all indices should give only the distinct indices,
        # i.e., give the same number of terms as the array length
        @test sum(1 for s in S) == length(S)
        @test sum(1 for I in eachindex(S)) == length(S)

        # calculating the linear index when iterating over Cartesian indices should give sequential access to the array
        @test 1:length(S) == [(LinearIndices(S.data)[_sub2grp(S,Tuple(I)...)...] for I in eachindex(S))...]

        @testset "_sub2grp" begin
            # test that permuting exchangeable indices accesses the same array element
            i1  = _sub2grp(S,1,2,3,1,4,3,2,1)
            @test _sub2grp(S,2,3,1,1,4,3,2,1) == i1
            @test _sub2grp(S,2,3,1,1,3,4,2,1) == i1
            @test _sub2grp(S,2,3,1,1,3,4,1,2) == i1
            @test _sub2grp(S,2,1,3,1,4,3,1,2) == i1
            # make sure that swapping independent indices gives a different array element
            @test _sub2grp(S,2,1,3,1,4,1,3,2) != i1
            @test _sub2grp(S,1,2,3,2,4,3,2,1) != i1
            @test _sub2grp(S,1,2,3,1,1,3,2,4) != i1

            SI = eachindex(S)
            @test first(SI) == CartesianIndex(1,1,1,1,1,1,1,1)

            NN = 4
            maxNdim = 40
            for Ndim = 1:maxNdim
                S = SymArray{(Ndim,),Float64}(ntuple(i->NN,Ndim)...)
                Is = ntuple(i->NN,Ndim)
                @test _sub2grp(S,Is...) == size(S.data)
                Is = ntuple(i->1,Ndim)
                @test _sub2grp(S,Is...) == (1,)
            end
        end

        A = rand(5,5)
        @test A' != A
        @test SymArray{(1,1)}(A) == A
        @test SymArray{(2,)}(A) != A

        A = A + A'
        # the standard equality test uses iteration over the arrays, (a,b) in zip(A,B),
        # which only accesses the actually different indices in SymArrays
        @test SymArray{(2,)}(A) != A
        # broadcasting goes over all indices
        @test all(SymArray{(2,)}(A) .== A)

        # views on existing vector-like types
        x = rand(3)
        S = SymArray{(2,)}(x,2,2)
        @test S.data === x
        fill!(S,8.)
        @test all(S .== 8.)

        S2 = SymArray{(2,)}(0*x,2,2)
        copyto!(S2,S)
        @test S == S2

        x = 5:10
        S = SymArray{(2,)}(x,3,3)
        @test S.data === x

        A = rand(10)
        S = SymArray{(1,)}(A,10)
        @test S.data === A

        # but if called without sizes, the copy constructor should be used
        S = SymArray{(1,)}(A)
        @test S.data !== A
    end

    @testset "which_symgrp" begin
        S = SymArray{(3,2),Float64}(5,5,5,3,3)
        @test which_symgrp(S,1) == (1,3)
        @test which_symgrp(S,3) == (1,3)
        @test which_symgrp(S,4) == (2,2)
        @test which_symgrp(S,5) == (2,2)
    end

    @testset "Contractions" begin
        @testset "Manual" begin
            N,M,O = 3, 5, 8
            for T in (Float64,ComplexF64)
                A = rand(T,N)

                for U in (Float64,ComplexF64)
                    for (n,dims) in enumerate([(N,M,O),(O,N,M),(M,O,N)])
                        B = rand(U,dims...)
                        n==1 && (@tensor D1[j,k] := A[i]*B[i,j,k])
                        n==2 && (@tensor D1[j,k] := A[i]*B[j,i,k])
                        n==3 && (@tensor D1[j,k] := A[i]*B[j,k,i])
                        D2 = contract(A,B,Val(n))
                        @test D1 ≈ D2
                        if T==U
                            contract!(D2,A,B,Val(n))
                            @test D1 ≈ D2
                        else
                            @test_throws MethodError contract!(D2,A,B,Val(n))
                        end
                    end

                    S = SymArray{(3,),U}(N,N,N)
                    rand!(S.data)
                    B = collect(S)
                    @test contract(A,S,Val(1)) == contract(A,S,Val(2))
                    @test contract(A,S,Val(1)) == contract(A,S,Val(3))
                    for n = 1:3
                        @test collect(contract(A,S,Val(n))) ≈ contract(A,B,Val(n))
                    end

                    S = SymArray{(2,1),U}(N,N,N)
                    rand!(S.data)
                    B = collect(S)
                    @test contract(A,S,Val(1)) == contract(A,S,Val(2))
                    for n = 1:3
                        @test collect(contract(A,S,Val(n))) ≈ contract(A,B,Val(n))
                    end
                    @test !(collect(contract(A,S,Val(1))) ≈ collect(contract(A,S,Val(3))))

                    S = SymArray{(2,),U}(N,N)
                    rand!(S.data)
                    B = collect(S)
                    @test contract(A,S,Val(1)) ≈ contract(A,S,Val(2))
                    for n = 1:2
                        @test collect(contract(A,S,Val(n))) ≈ contract(A,B,Val(n))
                    end
                end
            end
        end
        @testset "Generated" begin
            # use small dimension sizes here so the tests do not take too long
            N,M,O = 4, 5, 6
            arrTypes = has_cuda_gpu() ? (Array,CuArray) : (Array,)

            for arrType in arrTypes
                for T in (Float64,ComplexF64)
                    A = rand(T,N,M,O) |> arrType
                    for U in (Float64,ComplexF64)
                        S = SymArray{(2,3,1),U}(arrType,M,M,N,N,N,O)
                        @test storage_type(S) <: arrType

                        rand!(S.data)
                        # first collect GPU->CPU, then SymArray -> Array
                        B = collect(collect(S)) |> arrType
                        @test storage_type(B) <: arrType

                        TU = promote_type(U,T)
                        res21 = SymArray{(1,1,1,3,1),TU}(arrType,N,O,M,N,N,N,O)
                        contract!(res21,A,S,Val(2),Val(1))
                        @tensor res21_tst[i,k,l,m,n,o,p] := A[i,j,k] * B[j,l,m,n,o,p]
                        @test collect(collect(res21)) ≈ collect(res21_tst)
                        if T==U
                            contract!(res21_tst,A,B,Val(2),Val(1))
                            @test collect(collect(res21)) ≈ collect(res21_tst)
                        end

                        res13 = SymArray{(1,1,2,2,1),TU}(arrType,M,O,M,M,N,N,O)
                        contract!(res13,A,S,Val(1),Val(3))
                        @tensor res13_tst[j,k,l,m,n,o,p] := A[i,j,k] * B[l,m,i,n,o,p]
                        @test collect(collect(res13)) ≈ collect(res13_tst)
                        if T==U
                            contract!(res13_tst,A,B,Val(1),Val(3))
                            @test collect(collect(res13)) ≈ collect(res13_tst)
                        end

                        # dimension 3, 4, and 5 should be equivalent
                        contract!(res13,A,S,Val(1),Val(4))
                        @test collect(collect(res13)) ≈ collect(res13_tst)
                        contract!(res13,A,S,Val(1),Val(5))
                        @test collect(collect(res13)) ≈ collect(res13_tst)
                    end

                    A = rand(T,N) |> arrType
                    S = SymArray{(2,3,1),T}(arrType,N,N,N,N,N,N)
                    rand!(S.data)
                    # first collect GPU->CPU, then SymArray -> Array
                    B = collect(collect(S)) |> arrType

                    res11 = SymArray{(1,3,1),T}(arrType,N,N,N,N,N)
                    contract!(res11,A,S,Val(1),Val(1))
                    @tensor res11_tst[j,k,l,m,n] := A[i] * B[i,j,k,l,m,n]
                    @test collect(collect(res11)) ≈ collect(res11_tst)
                    contract!(res11_tst,A,B,Val(1),Val(1))
                    @test collect(collect(res11)) ≈ collect(res11_tst)

                    res13 = SymArray{(2,2,1),T}(arrType,N,N,N,N,N)
                    contract!(res13,A,S,Val(1),Val(3))
                    @tensor res13_tst[j,k,l,m,n] := A[i] * B[j,k,i,l,m,n]
                    @test collect(collect(res13)) ≈ collect(res13_tst)
                    contract!(res13_tst,A,B,Val(1),Val(3))
                    @test collect(collect(res13)) ≈ collect(res13_tst)

                    # dimension 3, 4, and 5 should be equivalent
                    contract!(res13_tst,A,B,Val(1),Val(4))
                    @test collect(collect(res13)) ≈ collect(res13_tst)
                    contract!(res13_tst,A,B,Val(1),Val(5))
                    @test collect(collect(res13)) ≈ collect(res13_tst)

                    res16 = SymArray{(2,3),T}(arrType,N,N,N,N,N)
                    contract!(res16,A,S,Val(1),Val(6))
                    @tensor res16_tst[j,k,l,m,n] := A[i] * B[j,k,l,m,n,i]
                    @test collect(collect(res16)) ≈ collect(res16_tst)
                    contract!(res16_tst,A,B,Val(1),Val(6))
                    @test collect(collect(res16)) ≈ collect(res16_tst)

                    # check that contraction from SymArray to Array works
                    A = rand(T,M,O) |> arrType
                    S = SymArray{(1,2,1),T}(arrType,N,M,M,O)
                    rand!(S.data)
                    B = collect(collect(S)) |> arrType

                    res12 = zeros(T,O,N,M,O) |> arrType
                    res12_tst = zeros(T,O,N,M,O) |> arrType
                    contract!(res12,A,S,Val(1),Val(2))
                    contract!(res12_tst,A,B,Val(1),Val(2))
                    @test collect(res12) ≈ collect(res12_tst)
                end
            end
        end
    end
end


using Match
using DataStructures
include("CherokeeQuantumGates.jl")

# author: Younis Shah

stage_deque = StageOutputDeque(Deque{StageOutput}())

function map_inputs(input::String)
  mapreduce(state -> rstrip(lstrip(state)) == "zero" ? b₀ : b₁ , kron, split(input, " "))
end

function tensor(ψ::Array{Symbol})
  ψ = detect_gate(ψ)
  accumulator = ψ[1]
  for i = 2:length(ψ)
    accumulator = round_v1(kron(accumulator, ψ[i]))
  end
  accumulator
end

# Convert the string of quantum gates to a Symbol Table/Matrix
function create_symbolic_ckt_matrix(ckt_string::String, nqubits::Int64, stages::Int64)
  symbolic_ckt_mat = Array(Symbol, nqubits, stages)
  ckt = split(ckt_string,";")
  for column = 1:length(ckt)
    rows = split(ckt[column])
    for row = 1:length(rows)
      symbolic_ckt_mat[row, column] = symbol(rows[row])
    end
  end
  symbolic_ckt_mat
end

function detect_gate(gate_symbol::Symbol)
  @match gate_symbol begin
    :H => H
    :I => Ι
    :X => σₓ
    :Y => Y
    :Z => Z
    :S => S
     _ => "NA"
  end
end

@vectorize_1arg Symbol detect_gate

function generate_basis_vectors_v1(n::Int64)
  basis_vecs = [0 : 2^n-1]
  digits_ = Array(String, 2^n)
  for i in basis_vecs
    digits_[i+1] = bin(i,n)
  end
  digits_
end

function gate_CNOT_positions(column::Array{Symbol})
  control_list, target_pos = Int64[], 0
  for i = 1:length(column)
    if column[i] == :C
      push!(control_list, i)
    elseif column[i] == :N
      target_pos = i
    end
  end
  control_list,target_pos
end

function simulate_CNOT(basis_vectors::Array, control, target::Int64, nqubits::Int64)
  vectors = map(vec -> parseint(vec, 2), map(vec -> xor(vec, control, target), basis_vectors))
  CNOT = zeros(Int64, 2^nqubits, 2^nqubits) # as Sparse matrix
  j=1
  for i in vectors
    CNOT[i+1,j] = 1
    j +=1
  end
  CNOT
end

function xor(s::String, control_qubits, target_qubit::Int64)
  digits_ = split(s,"")
  and_of_c_qubits = reduce(&, create_control_qubit_array(digits_, control_qubits)) # "And" the qubits -> Ummr Parry
  digits_[target_qubit] = string(and_of_c_qubits $ int(digits_[target_qubit]))
  join(digits_)
end

function create_control_qubit_array(s::Array, control_qubits)
  c_qubits = Int64[]
  for c in control_qubits
    push!(c_qubits, int(s[c]))
  end
  c_qubits
end

function round_v1(accumulator::Array)
  if length(filter(x -> x != 0, imag(accumulator))) == 0 # means no imaginary part
      accumulator = round(accumulator, 3)
    else
      accumulator = complex(round(real(accumulator),3), round(imag(accumulator),3))
    end
  accumulator
end

# expects a column vector
function router(mat::Array{Symbol}, basis_vecs::Array, nqubits_stages::(Int64,Int64), stage_deque::StageOutputDeque)
  if :C in mat && :N in mat && !(:H in mat || :X in mat || :Y in mat || :Z in mat || :S in mat) # only CNOT
    control, target = gate_CNOT_positions(mat)
    stage = StageOutput(nqubits_stages[2], simulate_CNOT(basis_vecs, control, target, nqubits_stages[1]))
    push!(stage_deque.stages, stage)
  elseif !(:C in mat && :N in mat)  # No CNOT
    stage = StageOutput(nqubits_stages[2], tensor(mat))
    push!(stage_deque.stages, stage)
  else
    return "NA"
  end
end

function multiply_stages(stages::StageOutputDeque)
  accumulator = pop!(stages.stages).output;
  for i = 1:length(stages.stages)
    stage = pop!(stage_deque.stages).output
    accumulator = accumulator * stage
  end
 round_v1(accumulator)
end

function simulate(ckt::String, nqubits::Int64, stages::Int64)
  #print_with_color(:blue, "SIMULATING...")
  mat = create_symbolic_ckt_matrix(ckt, nqubits, stages)
  basis_vecs = generate_basis_vectors_v1(nqubits)
  for i = 1:stages
    if router(mat[:,i], basis_vecs, (nqubits,i), stage_deque) == "NA"
      return "NA"
    end
  end
 # print_with_color(:blue, "DONE\n")
  #saved_stages = deepcopy(stage_deque) ############################### for WEB
  multiply_stages(stage_deque) # OUTPUT! Yaaaaayyy!!!!
end

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

#ckt = "H H H Y;Y X I I" # from web
##println(simulate(ckt, nqubits, stages))

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


#= Neural Network Variations
Mainly tuning the Multi-Layer Perceptron
A big limitation is that this doesn't use GPU at all
since ScikitLearn doesn't support it
So deeper models will take forever even on a good CPU
And I don't even want to try the 13-feature data set
-Matei
=#

using DataFrames, CSV, ScikitLearn, PyPlot, Statistics

# Data Loading
################################################################################

data = CSV.File("audio_recoded_combined_binary.csv") |> DataFrame

X = convert(Array, data[!,Not(:Class)])
y = convert(Array, data[!,:Class]) # :Class is our target variable

@sk_import model_selection: train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

@sk_import neural_network: MLPClassifier

#=
mlp_1layer10  = MLPClassifier(hidden_layer_sizes=(10))
mlp_1layer25  = MLPClassifier(hidden_layer_sizes=(25))
mlp_1layer40  = MLPClassifier(hidden_layer_sizes=(40))
mlp_1layer50  = MLPClassifier(hidden_layer_sizes=(50))
mlp_1layer100 = MLPClassifier(hidden_layer_sizes=(100))

mlp_6layer    = MLPClassifier(hidden_layer_sizes=(30, 50, 60, 10, 10, 10))
mlp_5layer    = MLPClassifier(hidden_layer_sizes=(50, 10, 10, 10, 10, 10))
=#

@sk_import metrics: classification_report
@sk_import metrics: plot_confusion_matrix
@sk_import model_selection: cross_val_score

function trainModel(model)
    fit!(model,X_train,y_train)
    # Training performance
    y_pred = predict(model,X_train)
    print(classification_report(y_train,y_pred))
    # Confusion matrices
    plot_confusion_matrix(model,X_test,y_test)
    PyPlot.gcf()
end

function crossValidateModel(model,folds=2)
    #fit!(model,X_train,y_train)
    @time cva = cross_val_score(model, X_train, y_train, cv=folds)
    println("mean accuracy: \n" *  string(mean(cva)))
    @time cvp = cross_val_score(model, X_train, y_train, cv=folds, scoring="precision")
    println("mean precision: \n" * string(mean(cvp)))
    @time cvr = cross_val_score(model, X_train, y_train, cv=folds, scoring="recall")
    println("mean recall: \n" *    string(mean(cvr)))
end

function testModel(model)
    # Testing performance
    y_pred = predict(model,X_test)
    print(classification_report(y_test,y_pred))
    # Confusion matrices
    plot_confusion_matrix(model,X_test,y_test)
    PyPlot.gcf()
end

mlp_2layer    = MLPClassifier(hidden_layer_sizes=(50, 10))
trainModel(         mlp_2layer)
crossValidateModel( mlp_2layer)
testModel(          mlp_2layer)

mlp_2layerlarge    = MLPClassifier(hidden_layer_sizes=(100, 50))
trainModel(         mlp_2layerlarge)
crossValidateModel( mlp_2layerlarge)
testModel(          mlp_2layerlarge)

mlp_3layer    = MLPClassifier(hidden_layer_sizes=(50, 10, 10))
trainModel(         mlp_3layer)
crossValidateModel( mlp_3layer)
testModel(          mlp_3layer)

mlp_4layer    = MLPClassifier(hidden_layer_sizes=(50, 10, 10, 10))
trainModel(         mlp_4layer)
crossValidateModel( mlp_4layer)
testModel(          mlp_4layer)

mlp_5layer    = MLPClassifier(hidden_layer_sizes=(50, 10, 10, 10, 10, 10))
trainModel(         mlp_5layer)
crossValidateModel( mlp_5layer)
testModel(          mlp_5layer)

testModel(mlp_1layer50)
@sk_import model_selection: cross_val_score
@time cva = cross_val_score(mlp_1layer50, X_train, y_train, cv=10)
println("mean accuracy:" * mean(cva))
mlp_1layer50_p  = MLPClassifier(hidden_layer_sizes=(50))
fit!(model,X_train,y_train)
@time cvp = cross_val_score(mlp_1layer50_p, X_train, y_train, cv=10, scoring="precision")
println("mean precision:" * mean(cvp))
@time cvr = cross_val_score(model, X_train, y_train, cv=10, scoring="recall")
println("mean recall:" * mean(cvr))

# Cross validation
################################################################################
#=
@sk_import model_selection: cross_val_score

function crossValidateModel(model)

    #fit!(model,X_train,y_train)
    @time cva = cross_val_score(model, X_train, y_train, cv=10)
    println("mean accuracy:" * mean(cva))
    @time cvp = cross_val_score(model, X_train, y_train, cv=10, scoring="precision")
    println("mean precision:" * mean(cvp))
    @time cvr = cross_val_score(model, X_train, y_train, cv=10, scoring="recall")
    println("mean recall:" * mean(cvr))

end

@time crossValidateModel(mlp_1layer50)
=#
# Validation and learning curves
# Will get this to work eventually -Matei
################################################################################

#=
@sk_import model_selection: validation_curve
@sk_import model_selection: learning_curve

using Plots, StatsPlots

validation_curve(mlp_1layer50, X, y,
    param_name="gamma", param_range=np.logspace(-6, -1, 5), scoring="accuracy")

fit!(mlp_1layer50,X_train,y_train)
train_sizes, train_scores, valid_scores = learning_curve(
    mlp_1layer50, X, y, train_sizes=[100, 500, 1000], cv=2)

plot(train_scores, valid_scores)
=#

# Random Forest Comparison
################################################################################

@sk_import ensemble: RandomForestClassifier
rfc    = RandomForestClassifier()
trainModel(         rfc)
crossValidateModel( rfc)
testModel(          rfc)


"""
Check if the matrix is square...
"""
function checkmatrix(m::Matrix)
    (size(m,1) != size(m,2) || !(eltype(m) <: Number)) && error("Only square matrices of Numbers accepted, sorry!!!");
    for r in 1:size(m,1)
        for c in 1:size(m,2)
            if r == c
                m[r,c] != 0 && error("The diagonal must be filled by 0s");
            else
                m[r,c] != m[c,r] && error("The matrix must be symmetrical");
            end
        end
    end
    return m;
end

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

"""
Return all the possibles combinations of these items
"""
function all_possible_rutes(items::Vector{T})::Vector{Vector{T}} where T

    #combs
    combs = Vector{Vector{T}}();

    #base case
    if length(items) == 2
        return push!(combs, [items[2],items[1]],[items[1],items[2]]);
    end

    #recursion block
    for i in 1:length(items)
        push!(combs, pushfirst!.(all_possible_rutes(copy_and_remove(items,i)),items[i])...);
    end

    return combs;
end

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

"""
get the total cost of a given rute. The cost include the travel from the last city back to home.
"""
function rute_total_cost(cmatrix::Matrix{T}, rute::Vector{Int}) where T

    cost = 0;
    #add travel cost
    for t in 1:length(rute)-1
        cost += cmatrix[rute[t],rute[t+1]]; # central part of the trip
    end
    return cost + cmatrix[rute[end], 1]; # the way home
end

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

"""
get the total costs of the given rutes. The cost of each rute include the travel from the last city back to home.
"""
function rutes_total_costs(cmatrix::Matrix{T}, rutes::Vector{Vector{Int}})::Vector where T
    costs = Vector();
    for rute in rutes
       push!(costs,rute_total_cost(cmatrix, rute)); # pushing all the rutes costs
    end
    return costs;
end

#################################################################################
"""
return the costs of each stretch of the rute...
"""
function rute_stretches_costs(cmatrix::Matrix{T}, rute::Vector{Int}) where T
   costs = zeros(T,length(rute));
    for t in 1:length(rute)-1
        costs[t] = cmatrix[rute[t],rute[t+1]]; # central part of the trip
    end
    costs[end] = cmatrix[rute[end],1];
    return costs;
end

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

"""
return the rute of minimal cost.
"""
function pick_cheaper(cmatrix::Matrix{T}, rutes::Vector{Vector{Int}})::Vector{Int} where T
    return rutes[findmin(rutes_total_costs(cmatrix, rutes))[2]];
end

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

function copy_and_remove(a::AbstractVector,i::Int)
    return lazy_remove!(copy(a),i);
end

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

"""
Remove an element from the array. Warning, it is very probable that this
methos unsort the array!!!
"""
function lazy_remove!(n::AbstractArray, i::Int)
    if i == length(n)
        pop!(n);
    else
        n[i] = pop!(n);
    end
    return n;
end

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

"""
Pick the rute from the bit matrix
"""
function pick_rute!(bm)
        r = Vector{Int}()
        bml = size(bm,1)
        d1 = 1;
        for i in 1:bml
            for d2 in 1:bml
                if bm[d1,d2] != 0
                    # println("Action in ", d1, " , ", d2);

                    #Save d1
                    push!(r,d1);
                    # println(d1, " pushed ", r);

                    #Clear
                    # println("Clearing")
                    # display(bm);
                    bm[d1,d2] = 0;
                    bm[d2,d1] = 0;
                    # println("Cleared")
                    # display(bm);

                    #Jump to d2
                    # println("Jumping to ", d2)
                    d1 = d2;

                    #break
                    # println("Break")
                    # println();
                    break;
                end
            end
        end
        return r;
end

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

"""
Get the number of possible rutes given a number of cities to visit...
"""
function rutescount(ncities::T)::T where {T <: Integer}
    if ncities <= 1 error("ncities must be greater than 1")
    elseif ncities == 2 return 1
    else
        return (ncities - 1) * rutescount(ncities - 1)
    end
end

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

function reverserute(r)
    rr = reverse(r);
    po = pop!(rr);
    return pushfirst!(rr,po);
end


using DiffEqSensitivity, QuasiMonteCarlo, Test

function ishi_batch(X)
    A= 7
    B= 0.1
    @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
end
function ishi(X)
    A= 7
    B= 0.1
    sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1])
end

function linear_batch(X)
    A= 7
    B= 0.1
    @. A*X[1,:]+B*X[2,:]
end
function linear(X)
    A= 7
    B= 0.1
    A*X[1]+B*X[2]
end

n = 600000
lb = -ones(4)*π
ub = ones(4)*π
sampler = SobolSample()
A,B = QuasiMonteCarlo.generate_design_matrices(n,lb,ub,sampler)

res1 = gsa(ishi,Sobol(),A,B)
res2 = gsa(ishi_batch,Sobol(),A,B,batch=true)

@test res1.S1 ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4
@test res2.S1 ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4

@test res1.ST ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4
@test res2.ST ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4

res1 = gsa(linear,Sobol(),A,B)
res2 = gsa(linear_batch,Sobol(),A,B,batch=true)

@test res1.S1 ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.S1 ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

@test res1.ST ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.ST ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

#=
library(sensitivity)
ishigami.fun <- function(X) {
  A <- 7
  B <- 0.1
  sin(X[, 1]) + A * sin(X[, 2])^2 + B * X[, 3]^4 * sin(X[, 1])
}
n <- 600000
X1 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
X2 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
sobol2007(ishigami.fun, X1, X2)
sobolSalt(ishigami.fun, X1, X2, scheme="A")
sobolSalt(ishigami.fun, X1, X2, scheme="B")

library(sensitivity)
ishigami.fun <- function(X) {
  A <- 7
  B <- 0.1
  A * X[, 1] + B * X[, 2]
}
n <- 6000000
X1 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
X2 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
sobol2007(ishigami.fun, X1, X2)
sobolSalt(ishigami.fun, X1, X2, scheme="A")
sobolSalt(ishigami.fun, X1, X2, scheme="B")
=#

function ishi_linear(X)
    A= 7
    B= 0.1
    [sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1]),A*X[1]+B*X[2]]
end

function ishi_linear_batch(X)
    A= 7
    B= 0.1
    X1 = @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
    X2 = @. A*X[1,:]+B*X[2,:]
    vcat(X1',X2')
end

res1 = gsa(ishi_linear,Sobol(),A,B)
res2 = gsa(ishi_linear_batch,Sobol(),A,B,batch=true)

# Now both tests together

@test res1.S1[1,:] ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4
@test res2.S1[1,:] ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4

@test res1.ST[1,:] ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4
@test res2.ST[1,:] ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4

@test res1.S1[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.S1[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

@test res1.ST[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.ST[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4


using Distributed
using LinearAlgebra #, Statistics
using Printf, WriteVTK, AbaqusReader
using PyCall, PyPlot, JLD, ProgressMeter

using AD4SM

@show Elements.setp(4)
;

mean(x) = sum(x)/length(x)

# this function is the constraint equation for the internal holes
function cnst_func(u,ucntr,r0,cntr0,Rsq)
    r    = r0+u
    cntr = cntr0+ucntr
    sum((r-cntr).^2)-Rsq
end
# this function find the centre of the constrained holes in the reference configuration
function find_centre(nodes; dTol = 1e-12, maxiter=21)

  c0    = mean(nodes)
  bdone = false
  iter  = 0

  while !bdone
    D2c0   = adiff.D2(c0)
    global deltas = [(D2c0[1]-p[1]).^2+(D2c0[2]-p[2]).^2 for p in nodes]
    r      = norm(deltas)
    grad   = adiff.grad(r)
    if (norm(grad)<dTol) || (iter>maxiter)
      bdone = true
    else
      c0  -= adiff.hess(r)\grad
    end
    iter +=1
  end
  Rsquares = adiff.val.(deltas)
  return (c0, Rsquares)
end

# some plotting functions
patch = pyimport("matplotlib.patches")
coll  = pyimport("matplotlib.collections")

function plot_model(elems, nodes;
                    u = zeros(length(nodes[1]), length(nodes)),
                    Φ = [],
                    linewidth = 0.25,
                    facecolor = :c,
                    edgecolor = :b,
                    alpha     = 1,
                    cmap      = :hsv,
                    clim      = [],
                    dTol      = 1e-6,
                    cfig      = figure(),
                    ax        = cfig.add_subplot(1,1,1))

  nodes     = [node + u[:,ii] for (ii,node) in enumerate(nodes)]
  patchcoll = coll.PatchCollection([patch.Polygon(nodes[elem.nodes])
                                    for elem ∈ elems], cmap=cmap)
  if !isempty(Φ)
    patchcoll.set_array(Φ)
    cfig.colorbar.(patchcoll, ax=ax)

    if isempty(clim)
      clim = patchcoll.get_clim()
      if abs(clim[2]-clim[1]) < dTol
        clim  = sum(clim)/2*[0.9, 1.1]
      end
    end
    patchcoll.set_clim(clim)
  else
    patchcoll.set_color(facecolor)
  end

  patchcoll.set_edgecolor(edgecolor)
  patchcoll.set_alpha(alpha)
  patchcoll.set_linewidth(linewidth)

  ax.set_aspect("equal")
  ax.add_collection(patchcoll)
  ax.autoscale()
  return (cfig, ax, patchcoll)
end
function get_I1(elems, u0)

    nElems = length(elems)
    F      = Elements.getinfo(elems,u0,info=:F)
    J      = [det(F)                       for F  in F]
    C      = [transpose(F)*F               for F  in F]
    L3     = [(C[1]C[4]-C[2]C[3])^-1       for C  in C]
    Ic     = [Materials.getInvariants(C[ii], L3[ii]) for ii in 1:nElems]

    I1     = [item[1] for item in Ic]
end
;

sMeshFile = "Pattern2D03FinerMesh02j.inp"
ϵn0       = 0.24
bfside    = true
mat       = Materials.NeoHooke(10)
bprogress = false
bwarmup   = true
blheqs    = false
bsheqs    = true
sFileName = splitext(sMeshFile)[1]
;

mymodel   = AbaqusReader.abaqus_read_mesh(sMeshFile)
@printf("\n\t loaded %s \n", sMeshFile)

(xmin, xmax) = extrema([item[2][1] for item ∈ mymodel["nodes"]])
(ymin, ymax) = extrema([item[2][2] for item ∈ mymodel["nodes"]])

nodes    = [mymodel["nodes"][ii]    for ii ∈ 1:length(mymodel["nodes"])]
el_nodes = [mymodel["elements"][ii] for ii ∈ 1:length(mymodel["elements"])]

elems = []
for item ∈ el_nodes
if length(item)==3
  push!(elems, Elements.Tria(item, nodes[item], mat=mat))
elseif length(item)==4
  push!(elems, Elements.Quad(item, nodes[item], mat=mat))
end
end

points   = hcat(nodes...)
cells    = [if length(nodes)==3
            MeshCell(VTKCellTypes.VTK_TRIANGLE, nodes)
          else length(nodes)==4
            MeshCell(VTKCellTypes.VTK_QUAD, nodes)
          end  for nodes in el_nodes ]

@show nNodes, nElems  = length(nodes), length(el_nodes)
@show Δx, Δy          = (xmax-xmin), (ymax-ymin)

nid_bndl = mymodel["node_sets"]["ID_L"]
nid_bndr = mymodel["node_sets"]["ID_R"]
nid_bndt = mymodel["node_sets"]["ID_T"]
nid_bndb = mymodel["node_sets"]["ID_B"]

nid_sh01 = mymodel["node_sets"]["ID_SH01"]
nid_sh02 = mymodel["node_sets"]["ID_SH02"]
nid_sh03 = mymodel["node_sets"]["ID_SH03"]
nid_sh04 = mymodel["node_sets"]["ID_SH04"]
nid_sh05 = mymodel["node_sets"]["ID_SH05"]
nid_sh06 = mymodel["node_sets"]["ID_SH06"]
nid_sh07 = mymodel["node_sets"]["ID_SH07"]
nid_sh08 = mymodel["node_sets"]["ID_SH08"]
nid_sh09 = mymodel["node_sets"]["ID_SH09"]

nid_shx  = [nid_sh01, nid_sh02, nid_sh03,
          nid_sh04, nid_sh05, nid_sh06,
          nid_sh07, nid_sh08, nid_sh09]
nSHs     = length(nid_shx)

nid_lh01 = mymodel["node_sets"]["ID_LH01"]
nid_lh02 = mymodel["node_sets"]["ID_LH02"]
nid_lh03 = mymodel["node_sets"]["ID_LH03"]
nid_lh04 = mymodel["node_sets"]["ID_LH04"]
nid_lh05 = mymodel["node_sets"]["ID_LH05"]
nid_lh06 = mymodel["node_sets"]["ID_LH06"]
nid_lh07 = mymodel["node_sets"]["ID_LH07"]
nid_lh08 = mymodel["node_sets"]["ID_LH08"]
nid_lh09 = mymodel["node_sets"]["ID_LH09"]
nid_lh10 = mymodel["node_sets"]["ID_LH10"]
nid_lh11 = mymodel["node_sets"]["ID_LH11"]
nid_lh12 = mymodel["node_sets"]["ID_LH12"]
nid_lh13 = mymodel["node_sets"]["ID_LH13"]
nid_lh14 = mymodel["node_sets"]["ID_LH14"]
nid_lh15 = mymodel["node_sets"]["ID_LH15"]
nid_lh16 = mymodel["node_sets"]["ID_LH16"]

nid_lhx  = [nid_lh01, nid_lh02, nid_lh03,
          nid_lh04, nid_lh05, nid_lh06,
          nid_lh07, nid_lh08, nid_lh09, nid_lh10,
          nid_lh11, nid_lh12, nid_lh13,
          nid_lh14, nid_lh15, nid_lh16]
nLHs     = length(nid_lhx)

cntrs_s = [find_centre(nodes[idx]) for idx in nid_shx]
cntrs_l = [find_centre(nodes[idx]) for idx in nid_lhx]
nodes = vcat(nodes,getindex.(cntrs_s,1),getindex.(cntrs_l,1))
;

plot_model(elems, nodes, facecolor=:c, edgecolor=:c)
PyPlot.title("undeformed model")
;

idxs  = LinearIndices(((2, nNodes+nSHs+nLHs)))
eqns  = Array{Elements.ConstEq}(undef, 0)
if bsheqs
    eqns  = vcat(eqns,
                 [ [ begin
                      r0      = nodes[nid]
                      cntr0   = nodes[nNodes+ii]
                      Rsq     = cntrs_s[ii][2][jj]
                      id_dofs = vcat(idxs[:,nid][:],idxs[:,nNodes+ii][:])
                      Elements.ConstEq(x->cnst_func(x[1:2],x[3:4],r0,cntr0,Rsq),
                                       id_dofs,adiff.D2)
                    end  for (jj,nid) in enumerate(nid_ii)]
                  for (ii,nid_ii) in enumerate(nid_shx)]...)
end
if blheqs
    eqns = vcat(eqns,
                [ [ begin
                     r0      = nodes[nid]
                     cntr0   = nodes[nNodes+nSHs+ii]
                     Rsq     = cntrs_l[ii][2][jj]
                     id_dofs = vcat(idxs[:,nid][:],idxs[:,nNodes+nSHs+ii][:])
                     Elements.ConstEq(x->cnst_func(x[1:2],x[3:4],r0,cntr0,Rsq),
                                      id_dofs,adiff.D2)
                   end  for (jj,nid) in enumerate(nid_ii)]
                 for (ii,nid_ii) in enumerate(nid_lhx)]... )
end
;

u                  = fill(NaN, 2, nNodes+nSHs+nLHs)
u[:,nid_bndb]     .= 0
u[1,nid_bndt]     .= 0
u[2,nid_bndt]     .= -ϵn0*Δy

if !bsheqs u[:,nNodes+1:nNodes+nSHs] .= NaN; end
if !blheqs u[:,nNodes+nSHs+1:end]    .= NaN; end

if !bfside
    u[1,nid_bndl] .= 0
    u[1,nid_bndr] .= 0
end

ifree     = isnan.(u)
icnst     = .!ifree
;

u0         = 1e-4Δx*randn(2, nNodes+nSHs+nLHs)
u0[icnst] .= 0
unew       = copy(u0)
@time Elements.solvestep!(elems, u0, unew, ifree, eqns=eqns, λ = zeros(length(eqns)),
                          bprogress=false, becho=true)
u[ifree] .= unew[ifree]
;

N       = 24
LF_c    = vcat(range(0.0, 0.9, length=2N÷8),
               range(0.9, 1.0, length=6N÷8))

allus_c = Elements.solve(elems, u, LF=LF_c, ifree=ifree, eqns=eqns,
                          bprogress=false, becho=true, bechoi=true)
;

rf_tot_c = [sum(item[2][2,nid_bndt])   for item in allus_c]
Δu_tot_c = [mean(item[1][2,nid_bndt])  for item in allus_c]
;

I1 = get_I1(elems, allus_c[end][1])

cfig = figure()
ax1   = cfig.add_subplot(2,1,1)

plot_model(elems, nodes, alpha=0.05,
  facecolor=:c, edgecolor=:c, cfig=cfig, ax=ax1)
plot_model(elems, nodes, u = allus_c[end][1],
  edgecolor=:c, Φ = get_I1(elems, allus_c[end][1]), cfig=cfig, ax=ax1)
title("deformed model - compression")

ax2   = cfig.add_subplot(2,1,2)
ax2.plot(Δu_tot_c, rf_tot_c)
xlabel("Δu_tot_")
ylabel("rf_tot_c")
title("force-displacement, compressive branch")

cfig.tight_layout(pad=2.0, w_pad=2.0, h_pad=2.0)
;

u                  = fill(NaN, 2, nNodes+nSHs+nLHs)
u[:,nid_bndb]     .= 0
u[1,nid_bndt]     .= 0
u[2,nid_bndt]     .= ϵn0*Δy

if !bsheqs u[:,nNodes+1:nNodes+nSHs] .= NaN; end
if !blheqs u[:,nNodes+nSHs+1:end]    .= NaN; end

if !bfside
    u[1,nid_bndl]     .= 0
    u[1,nid_bndr]     .= 0
end

ifree     = isnan.(u)
icnst     = .!ifree
;

u0         = 1e-4Δx*randn(2, nNodes+nSHs+nLHs)
u0[icnst] .= 0
unew       = copy(u0)
@time Elements.solvestep!(elems, u0, unew, ifree, eqns=eqns, λ = zeros(length(eqns)),
                          bprogress=false, becho=true)
u[ifree] .= unew[ifree]
;

N       = 10
LF_t    = vcat(range(0.0, 0.9, length=2N÷3),
               range(0.9, 1.0, length=N÷3+1))
println("\n\t starting the tensile branch \n"); flush(stdout)
allus_t = Elements.solve(elems, u, LF=LF_t, ifree=ifree, eqns=eqns,
                          bprogress=false, becho=true, bechoi=true)
;

rf_tot_t = [sum(item[2][2,nid_bndt])   for item in allus_t]
Δu_tot_t = [mean(item[1][2,nid_bndt])  for item in allus_t]
;

JLD.save(sFileName*".jld",
            "nodes", nodes, "elems", elems,
            "points", points, "cells", cells, "mat", mat,
            "nid_bndl", nid_bndl, "nid_bndr", nid_bndr,
            "nid_bndt", nid_bndt, "nid_bndb", nid_bndb,
            "allus_t", allus_t, "allus_c", allus_c,
            "LF_t", LF_t, "LF_c", LF_c,
            "rf_tot_t", rf_tot_t, "Δu_tot_t", Δu_tot_t,
            "rf_tot_c", rf_tot_c, "Δu_tot_c", Δu_tot_c,
            "cntrs_s", cntrs_s, "cntrs_l", cntrs_l,
            "nid_shx", nid_shx, "nid_lhx", nid_lhx,
            "Δx", Δx, "Δy", Δy, "nNodes", nNodes, "nElems", nElems,
            "sMeshFile", sMeshFile, "sFileName", sFileName)
@printf("results written to %s\n", sFileName); flush(stdout)
;

I1    = get_I1(elems, allus_t[end][1])

cfig  = figure()
ax1   = cfig.add_subplot(2,1,1)

plot_model(elems, nodes, alpha=0.05,
  facecolor=:c, edgecolor=:c, cfig=cfig, ax=ax1)
plot_model(elems, nodes, u = allus_t[end][1],
  edgecolor=:c, Φ = get_I1(elems, allus_t[end][1]), cfig=cfig, ax=ax1)
title("deformed model - tensile")

ax2   = cfig.add_subplot(2,1,2)
ax2.plot(Δu_tot_t, rf_tot_t)
xlabel("Δu_tot")
ylabel("rf_tot")
title("force-displacement, tensile branch")

cfig.tight_layout(pad=2.0, w_pad=2.0, h_pad=2.0)
;

PyPlot.figure()
PyPlot.plot(Δu_tot_c, rf_tot_c, Δu_tot_t, rf_tot_t)
PyPlot.xlabel("Δu_tot_")
PyPlot.ylabel("rf_tot_c")
PyPlot.title("force-displacement, total")
;


#=
Contains the pieces to build the stochastic Neoclassical growth model
described in the Arellano, Maliar, Maliar, Tsyrennikov 2016 paper.
=#

#
# The Model
#
"""
The stochastic Neoclassical growth model type contains parameters
which define the model

* α: Capital share in output
* β: Discount factor
* δ: Depreciation rate
* γ: Risk aversion
* ρ: Persistence of the log of the productivity level
* σ: Standard deviation of shocks to log productivity level
* A: Coefficient on C-D production function

* kgrid: Grid over capital
* zgrid: Grid over productivity
* grid: Grid of (k, z) pairs
* eps_nodes: Nodes used to integrate
* weights: Weights used to integrate
* z1: A grid of the possible z1s tomorrow given eps_nodes and zgrid
"""
@with_kw struct NeoclassicalGrowth
    # Parameters
    α::Float64 = 0.36
    β::Float64 = 0.99
    δ::Float64 = 0.02
    γ::Float64 = 2.0
    ρ::Float64 = 0.95
    σ::Float64 = 0.01
    A::Float64 = (1.0/β - (1 - δ)) / α

    # Grids
    kgrid::Vector{Float64} = collect(range(0.9, stop=1.1, length=10))
    zgrid::Vector{Float64} = collect(range(0.9, stop=1.1, length=10))
    grid::Matrix{Float64} = gridmake(kgrid, zgrid)
    eps_nodes::Vector{Float64} = qnwnorm(5, 0.0, σ^2)[1]
    weights::Vector{Float64} = qnwnorm(5, 0.0, σ^2)[2]
    z1::Matrix{Float64} = (zgrid.^(ρ))' .* exp.(eps_nodes)
end

# Helper functions
f(ncgm::NeoclassicalGrowth, k, z) = @. z * (ncgm.A * k^ncgm.α)
df(ncgm::NeoclassicalGrowth, k, z) = @. ncgm.α * z * (ncgm.A * k^(ncgm.α - 1.0))

u(ncgm::NeoclassicalGrowth, c) = c > 1e-10 ? @.(c^(1-ncgm.γ)-1)/(1-ncgm.γ) : -1e10
du(ncgm::NeoclassicalGrowth, c) = c > 1e-10 ? c.^(-ncgm.γ) : 1e10
duinv(ncgm::NeoclassicalGrowth, u) = u .^ (-1 / ncgm.γ)

expendables_t(ncgm::NeoclassicalGrowth, k, z) = (1-ncgm.δ)*k + f(ncgm, k, z)

# Types for solution methods
abstract type SolutionMethod end

struct IterateOnPolicy <: SolutionMethod end
struct VFI_ECM <: SolutionMethod end
struct VFI_EGM <: SolutionMethod end
struct VFI <: SolutionMethod end
struct PFI_ECM <: SolutionMethod end
struct PFI <: SolutionMethod end
struct dVFI_ECM <: SolutionMethod end
struct EulEq <: SolutionMethod end

#
# Type for Approximating Value and Policy
#
mutable struct ValueCoeffs{T <: SolutionMethod,D <: Degree}
    d::D
    v_coeffs::Vector{Float64}
    k_coeffs::Vector{Float64}
end

function ValueCoeffs(::Type{Val{d}}, method::T) where T <: SolutionMethod where d
    # Initialize two vectors of zeros
    deg = Degree{d}()
    n = n_complete(2, deg)
    v_coeffs = zeros(n)
    k_coeffs = zeros(n)

    return ValueCoeffs{T,Degree{d}}(deg, v_coeffs, k_coeffs)
end

function ValueCoeffs(
        ncgm::NeoclassicalGrowth, ::Type{Val{d}}, method::T
    ) where T <: SolutionMethod where d
    # Initialize with vector of zeros
    deg = Degree{d}()
    n = n_complete(2, deg)
    v_coeffs = zeros(n)

    # Policy guesses based on k and z
    k, z = ncgm.grid[:, 1], ncgm.grid[:, 2]
    css = ncgm.A - ncgm.δ
    yss = ncgm.A
    c_pol = f(ncgm, k, z) * (css/yss)

    # Figure out what kp is
    k_pol = expendables_t(ncgm, k, z) - c_pol
    k_coeffs = complete_polynomial(ncgm.grid, d) \ k_pol

    return ValueCoeffs{T,Degree{d}}(deg, v_coeffs, k_coeffs)
end

solutionmethod(::ValueCoeffs{T}) where T <:SolutionMethod = T

# A few copy methods to make life easier
Base.copy(vp::ValueCoeffs{T,D}) where T where D =
    ValueCoeffs{T,D}(vp.d, vp.v_coeffs, vp.k_coeffs)

function Base.copy(vp::ValueCoeffs{T1,D}, ::T2) where T1 where D where T2 <: SolutionMethod
    ValueCoeffs{T2,D}(vp.d, vp.v_coeffs, vp.k_coeffs)
end

function Base.copy(
        ncgm::NeoclassicalGrowth, vp::ValueCoeffs{T}, ::Type{Val{new_degree}}
    ) where T where new_degree
    # Build Value and policy matrix
    deg = Degree{new_degree}()
    V = build_V(ncgm, vp)
    k = build_k(ncgm, vp)

    # Build new Phi
    Phi = complete_polynomial(ncgm.grid, deg)
    v_coeffs = Phi \ V
    k_coeffs = Phi \ k

    return ValueCoeffs{T,Degree{new_degree}}(deg, v_coeffs, k_coeffs)
end

"""
Updates the coefficients for the value function inplace in `vp`
"""
function update_v!(vp::ValueCoeffs, new_coeffs::Vector{Float64}, dampen::Float64)
    vp.v_coeffs = (1-dampen)*vp.v_coeffs + dampen*new_coeffs
end

"""
Updates the coefficients for the policy function inplace in `vp`
"""
function update_k!(vp::ValueCoeffs, new_coeffs::Vector{Float64}, dampen::Float64)
    vp.k_coeffs = (1-dampen)*vp.k_coeffs + dampen*new_coeffs
end

"""
Builds either V or dV depending on the solution method that is given. If it
is a solution method that iterates on the derivative of the value function
then it will return derivative of the value function, otherwise the
value function itself
"""
build_V_or_dV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs) =
    build_V_or_dV(ncgm, vp, solutionmethod(vp)())

build_V_or_dV(ncgm, vp::ValueCoeffs, ::SolutionMethod) = build_V(ncgm, vp)
build_V_or_dV(ncgm, vp::ValueCoeffs, T::dVFI_ECM) = build_dV(ncgm, vp)

function build_dV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
    Φ = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())
    Φ*vp.v_coeffs
end

function build_V(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
    Φ = complete_polynomial(ncgm.grid, vp.d)
    Φ*vp.v_coeffs
end

function build_k(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
    Φ = complete_polynomial(ncgm.grid, vp.d)
    Φ*vp.k_coeffs
end

#
# Helper functions
#
function compute_EV!(cp_kpzp::Vector{Float64}, ncgm::NeoclassicalGrowth,
                     vp::ValueCoeffs, kp, iz)
    # Pull out information from types
    z1, weightsz = ncgm.z1, ncgm.weights

    # Get number nodes
    nzp = length(weightsz)

    EV = 0.0
    for izp in 1:nzp
        zp = z1[izp, iz]
        complete_polynomial!(cp_kpzp, [kp, zp], vp.d)
        EV += weightsz[izp] * dot(vp.v_coeffs, cp_kpzp)
    end

    return EV
end

function compute_EV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs, kp, iz)
    cp_kpzp = Array{Float64}(undef, n_complete(2, vp.d))

    return compute_EV!(cp_kpzp, ncgm, vp, kp, iz)
end

function compute_EV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
    # Get length of k and z grids
    kgrid, zgrid = ncgm.kgrid, ncgm.zgrid
    nk, nz = length(kgrid), length(zgrid)
    temp = Array{Float64}(undef, n_complete(2, vp.d))

    # Allocate space to store EV
    EV = Array{Float64}(undef, nk*nz)

    _inds = LinearIndices((nk, nz))

    for ik in 1:nk, iz in 1:nz
        # Pull out states
        k = kgrid[ik]
        z = zgrid[iz]
        ikiz_index = _inds[ik, iz]

        # Pass to scalar EV
        complete_polynomial!(temp, [k, z], vp.d)
        kp = dot(vp.k_coeffs, temp)
        EV[ikiz_index] = compute_EV!(temp, ncgm, vp, kp, iz)
    end

    return EV
end

function compute_dEV!(cp_dkpzp::Vector,  ncgm::NeoclassicalGrowth,
                      vp::ValueCoeffs, kp, iz)
    # Pull out information from types
    z1, weightsz = ncgm.z1, ncgm.weights

    # Get number nodes
    nzp = length(weightsz)

    dEV = 0.0
    for izp in 1:nzp
        zp = z1[izp, iz]
        complete_polynomial!(cp_dkpzp, [kp, zp], vp.d, Derivative{1}())
        dEV += weightsz[izp] * dot(vp.v_coeffs, cp_dkpzp)
    end

    return dEV
end

function compute_dEV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs, kp, iz)
    compute_dEV!(Array{Float64}(undef, n_complete(2, vp.d)), ncgm, vp, kp, iz)
end

function env_condition_kp!(cp_out::Vector{Float64}, ncgm::NeoclassicalGrowth,
                           vp::ValueCoeffs, k::Float64, z::Float64)
    # Compute derivative of VF
    dV = dot(vp.v_coeffs, complete_polynomial!(cp_out, [k, z], vp.d, Derivative{1}()))

    # Consumption is then computed as
    c = duinv(ncgm, dV / (1 - ncgm.δ .+ df(ncgm, k, z)))

    return expendables_t(ncgm, k, z) - c
end

function env_condition_kp(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,
                          k::Float64, z::Float64)
    cp_out = Array{Float64}(undef, n_complete(2, vp.d))
    env_condition_kp!(cp_out, ncgm, vp, k, z)
end

function env_condition_kp(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
    # Pull out k and z from grid
    k = ncgm.grid[:, 1]
    z = ncgm.grid[:, 2]

    # Create basis matrix for entire grid
    dPhi = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())

    # Compute consumption
    c = duinv(ncgm, (dPhi*vp.v_coeffs) ./ (1-ncgm.δ.+df(ncgm, k, z)))

    return expendables_t(ncgm, k, z) .- c
end


#=
Solves the "Offer Distribution Unknown" Model by value function
iteration and a second faster method discussed in the corresponding
quantecon lecture.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-08-14

References
----------

http://quant-econ.net/jl/odu.html

=#

using Distributions

"""
Unemployment/search problem where offer distribution is unknown

##### Fields

- `bet::Real` : Discount factor on (0, 1)
- `c::Real` : Unemployment compensation
- `F::Distribution` : Offer distribution `F`
- `G::Distribution` : Offer distribution `G`
- `f::Function` : The pdf of `F`
- `g::Function` : The pdf of `G`
- `n_w::Int` : Number of points on the grid for w
- `w_max::Real` : Maximum wage offer
- `w_grid::AbstractVector` : Grid of wage offers w
- `n_pi::Int` : Number of points on grid for pi
- `pi_min::Real` : Minimum of pi grid
- `pi_max::Real` : Maximum of pi grid
- `pi_grid::AbstractVector` : Grid of probabilities pi
- `quad_nodes::Vector` : Notes for quadrature ofer offers
- `quad_weights::Vector` : Weights for quadrature ofer offers

"""
type SearchProblem
    bet::Real
    c::Real
    F::Distribution
    G::Distribution
    f::Function
    g::Function
    n_w::Int
    w_max::Real
    w_grid::AbstractVector
    n_pi::Int
    pi_min::Real
    pi_max::Real
    pi_grid::AbstractVector
    quad_nodes::Vector
    quad_weights::Vector
end

"""
Constructor for `SearchProblem` with default values

##### Arguments

- `bet::Real(0.95)` : Discount factor in (0, 1)
- `c::Real(0.6)` : Unemployment compensation
- `F_a::Real(1), F_b::Real(1)` : Parameters of `Beta` distribution for `F`
- `G_a::Real(3), G_b::Real(1.2)` : Parameters of `Beta` distribution for `G`
- `w_max::Real(2)` : Maximum of wage offer grid
- `w_grid_size::Int(40)` : Number of points in wage offer grid
- `pi_grid_size::Int(40)` : Number of points in probability grid

##### Notes

There is also a version of this function that accepts keyword arguments for
each parameter

"""
function SearchProblem(bet=0.95, c=0.6, F_a=1, F_b=1, G_a=3, G_b=1.2,
                       w_max=2, w_grid_size=40, pi_grid_size=40)

    F = Beta(F_a, F_b)
    G = Beta(G_a, G_b)

    # NOTE: the x./w_max)./w_max in these functions makes our dist match
    #       the scipy one with scale=w_max given
    f(x) = pdf(F, x./w_max)./w_max
    g(x) = pdf(G, x./w_max)./w_max

    pi_min = 1e-3  # avoids instability
    pi_max = 1 - pi_min

    w_grid = linspace_range(0, w_max, w_grid_size)
    pi_grid = linspace_range(pi_min, pi_max, pi_grid_size)

    nodes, weights = qnwlege(21, 0.0, w_max)

    SearchProblem(bet, c, F, G, f, g,
                  w_grid_size, w_max, w_grid,
                  pi_grid_size, pi_min, pi_max, pi_grid, nodes, weights)
end

# make kwarg version
function SearchProblem(;bet=0.95, c=0.6, F_a=1, F_b=1, G_a=3, G_b=1.2,
                       w_max=2, w_grid_size=40, pi_grid_size=40)
    SearchProblem(bet, c, F_a, F_b, G_a, G_b, w_max, w_grid_size,
                  pi_grid_size)
end

function q(sp::SearchProblem, w, pi_val)
    new_pi = 1.0 ./ (1 + ((1 - pi_val) .* sp.g(w)) ./ (pi_val .* sp.f(w)))

    # Return new_pi when in [pi_min, pi_max] and else end points
    return clamp(new_pi, sp.pi_min, sp.pi_max)
end

"""
Apply the Bellman operator for a given model and initial value.

##### Arguments

- `sp::SearchProblem` : Instance of `SearchProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output.
- `;ret_policy::Bool(false)`: Toggles return of value or policy functions

##### Returns

None, `out` is updated in place. If `ret_policy == true` out is filled with the
policy function, otherwise the value function is stored in `out`.

"""
function bellman_operator!(sp::SearchProblem, v::Matrix, out::Matrix;
                          ret_policy::Bool=false)
    # Simplify names
    f, g, bet, c = sp.f, sp.g, sp.bet, sp.c
    nodes, weights = sp.quad_nodes, sp.quad_weights

    vf = CoordInterpGrid((sp.w_grid, sp.pi_grid), v, BCnan, InterpLinear)

    # set up quadrature nodes/weights
    # q_nodes, q_weights = qnwlege(21, 0.0, sp.w_max)

    for w_i=1:sp.n_w
        w = sp.w_grid[w_i]

        # calculate v1
        v1 = w / (1 - bet)

        for pi_j=1:sp.n_pi
            _pi = sp.pi_grid[pi_j]

            # calculate v2
            function integrand(m)
                quad_out = similar(m)
                for i=1:length(m)
                    mm = m[i]
                    quad_out[i] = vf[mm, q(sp, mm, _pi)] * (_pi*f(mm) +
                                                            (1-_pi)*g(mm))
                end
                return quad_out
            end
            integral = do_quad(integrand, nodes, weights)
            # integral = do_quad(integrand, q_nodes, q_weights)
            v2 = c + bet * integral

            # return policy if asked for, otherwise return max of values
            out[w_i, pi_j] = ret_policy ? v1 > v2 : max(v1, v2)
        end
    end
    return out
end

function bellman_operator(sp::SearchProblem, v::Matrix;
                          ret_policy::Bool=false)
    out_type = ret_policy ? Bool : Float64
    out = Array(out_type, sp.n_w, sp.n_pi)
    bellman_operator!(sp, v, out, ret_policy=ret_policy)
end

"""
Extract the greedy policy (policy function) of the model.

##### Arguments

- `sp::SearchProblem` : Instance of `SearchProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function get_greedy!(sp::SearchProblem, v::Matrix, out::Matrix)
    bellman_operator!(sp, v, out, ret_policy=true)
end

get_greedy(sp::SearchProblem, v::Matrix) = bellman_operator(sp, v,
                                                            ret_policy=true)

"""
Updates the reservation wage function guess phi via the operator Q.

##### Arguments

- `sp::SearchProblem` : Instance of `SearchProblem`
- `phi::Vector`: Current guess for phi
- `out::Vector` : Storage for output

##### Returns

None, `out` is updated in place to hold the updated levels of phi
"""
function res_wage_operator!(sp::SearchProblem, phi::Vector, out::Vector)
    # Simplify name
    f, g, bet, c = sp.f, sp.g, sp.bet, sp.c

    # Construct interpolator over pi_grid, given phi
    phi_f = CoordInterpGrid(sp.pi_grid, phi, BCnearest, InterpLinear)

    # set up quadrature nodes/weights
    q_nodes, q_weights = qnwlege(7, 0.0, sp.w_max)

    for (i, _pi) in enumerate(sp.pi_grid)
        integrand(x) = max(x, phi_f[q(sp, x, _pi)]).*(_pi*f(x) + (1-_pi)*g(x))
        integral = do_quad(integrand, q_nodes, q_weights)
        out[i] = (1 - bet)*c + bet*integral
    end
end

"""
Updates the reservation wage function guess phi via the operator Q.

See the documentation for the mutating method of this function for more details
on arguments
"""
function res_wage_operator(sp::SearchProblem, phi::Vector)
    out = similar(phi)
    res_wage_operator!(sp, phi, out)
    return out
end


using IterativeSolvers
using FactCheck

facts("Basic operations") do

context("Adivtype") do
    A = rand(1:10, 5, 5)
    b = rand(Float32, 5)
    @fact IterativeSolvers.Adivtype(A, b) --> Float32
end

context("Linear operator defined as a function") do
    # A = cycle-back operator
    function shiftback!(output, b)
        n = length(b)
        length(output) == n || throw(DimensionMismatch())
        for i = 1:n-1
            output[i+1] = b[i]
        end
        output[1] = b[n]
        output
    end

    # A' = cycle-forward operator
    function shiftfwd!(output, b)
        n = length(b)
        length(output) == n || throw(DimensionMismatch())
        for i = 2:n
            output[i-1] = b[i]
        end
        output[n] = b[1]
        output
    end

    b = rand(Float32, 5)
    Atimesb = [b[end]; b[1:end-1]]
    Aptimesb = [b[2:end]; b[1]]

    A = MatrixFcn{Int}(5, 5, shiftback!)
    @fact eltype(A) --> Int
    @fact size(A) --> (5,5)
    @fact size(A,1) --> 5
    @fact size(A,2) --> 5
    @fact length(A) --> 25
    @fact A*b --> Atimesb

    output = similar(b)
    @fact A_mul_B!(output, A, b) --> Atimesb
    @fact A'*b --> Atimesb

    A = MatrixCFcn{Int}(5, 5, shiftback!, shiftfwd!)
    @fact eltype(A) --> Int
    @fact size(A) --> (5,5)
    @fact size(A,1) --> 5
    @fact size(A,2) --> 5
    @fact length(A) --> 25
    @fact A*b --> Atimesb

    output = similar(b)
    @fact A_mul_B!(output, A, b) --> Atimesb
    @fact A'*b --> Aptimesb
    @fact Ac_mul_B!(output, A, b) --> Aptimesb
end

end


#################### Mathematical Operators ####################
function cummean(x::AbstractArray)
    return mapslices(cummean, x, dims = 1)
end

function cummean(x::AbstractVector)
    y = similar(x, Float64)
    xs = 0.0
    fill!(y, xs)

    c = 0
    for i in 1:length(x)
        if !ismissing(x[i])
            c += 1
            xs += x[i]
        end
        y[i] = xs / c
    end
    return y
end

## Csorgo S and Faraway JJ. The exact and asymptotic distributions of the
## Cramer-von Mises statistic. Journal of the Royal Statistical Society,
## Series B, 58: 221-234, 1996.
function pcramer(q::Real)
    p = 0.0
    for k in 0:3
        c1 = 4.0 * k + 1.0
        c2 = c1^2 / (16.0 * q)
        p += gamma(k + 0.5) / factorial(k) * sqrt(c1) * exp(-c2) * besselk(0.25, c2)
    end
    return p / (pi^1.5 * sqrt(q))
end

# This sorting stack was sourced from https://rosettacode.org/wiki/Natural_sorting#Julia
splitbynum(x) = split(x, r"(?<=\D)(?=\d)|(?<=\d)(?=\D)")
numstringtonum(arr) = [(n = tryparse(Float32, e)) != nothing ? n : e for e in arr]
function natural(x, y)
    xarr = numstringtonum(splitbynum(x))
    yarr = numstringtonum(splitbynum(y))
    for i in 1:min(length(xarr), length(yarr))
        if typeof(xarr[i]) != typeof(yarr[i])
            a = string(xarr[i]); b = string(yarr[i])
        else
             a = xarr[i]; b = yarr[i]
        end
        if a == b
            continue
        else
            return a < b
        end
    end
    return length(xarr) < length(yarr)
end

function _dict2namedtuple(d::Dict)
    t_keys = ntuple(x -> Symbol(collect(keys(d))[x]), length(keys(d)))
    t_vals = ntuple(x -> collect(values(d))[x], length(values(d)))
    return NamedTuple{t_keys}(t_vals)
end

function _namedtuple2dict(n::NamedTuple)
    return Dict(pairs(n))
end

function _ntdictmerge(n::NamedTuple, args::NamedTuple...)
    ndict = _namedtuple2dict(n)
    ndict2 = map(x -> _namedtuple2dict(x), args)
    return _dict2namedtuple(merge(ndict, ndict2...))
end

"""
    cskip(x)

Wrapper for `collect(skipmissing(x))`.
"""
cskip(x) = collect(skipmissing(x))


const naivebayesstanmodel = "
// supervised naive Bayes

data {
  // training data
  int<lower=1> K;               // num topics
  int<lower=1> V;               // num words
  int<lower=0> M;               // num docs
  int<lower=0> N;               // total word instances
  int<lower=1,upper=K> z[M];    // topic for doc m
  int<lower=1,upper=V> w[N];    // word n
  int<lower=1,upper=M> doc[N];  // doc ID for word n
  // hyperparameters
  vector<lower=0>[K] alpha;     // topic prior
  vector<lower=0>[V] beta;      // word prior
}
parameters {
  simplex[K] theta;   // topic prevalence
  simplex[V] phi[K];  // word dist for topic k
}
model {
  // priors
  theta ~ dirichlet(alpha);
  for (k in 1:K)
    phi[k] ~ dirichlet(beta);
  // likelihood, including latent category
  for (m in 1:M)
    z[m] ~ categorical(theta);
  for (n in 1:N)
    w[n] ~ categorical(phi[z[doc[n]]]);
}
"