using Test, LinearAlgebra
using OrdinaryDiffEq
using DiffEqSensitivity, StochasticDiffEq, DiffEqBase
using ForwardDiff, Calculus, ReverseDiff
using Random
import Tracker, Zygote

@info "SDE Adjoints"

seed = 100
Random.seed!(seed)
abstol = 1e-4
reltol = 1e-4

u₀ = [0.5]
tstart = 0.0
tend = 0.1
dt = 0.005
trange = (tstart, tend)
t = tstart:dt:tend
tarray = collect(t)

function g(u,p,t)
  sum(u.^2.0/2.0)
end

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

p2 = [1.01,0.87]

@testset "SDE oop Tests" begin
  f_oop_linear(u,p,t) = p[1]*u
  σ_oop_linear(u,p,t) = p[2]*u

  function f_oop_linear(u::Tracker.TrackedArray,p,t)
    dx = p[1]*u[1]
    Tracker.collect([dx])
  end

  function σ_oop_linear(u::Tracker.TrackedArray,p,t)
    dx = p[2]*u[1]
    Tracker.collect([dx])
  end

  p = [1.01,0.0]

  # generate ODE adjoint results

  prob_oop_ode = ODEProblem(f_oop_linear,u₀,(tstart,tend),p)
  sol_oop_ode = solve(prob_oop_ode,Tsit5(),saveat=t,abstol=abstol,reltol=reltol)
  res_ode_u0, res_ode_p = adjoint_sensitivities(sol_oop_ode,Tsit5(),dg!,t
    ,abstol=abstol,reltol=reltol,sensealg=BacksolveAdjoint())

  function G(p)
    tmp_prob = remake(prob_oop_ode,u0=eltype(p).(prob_oop_ode.u0),p=p,
                    tspan=eltype(p).(prob_oop_ode.tspan),abstol=abstol, reltol=reltol)
    sol = solve(tmp_prob,Tsit5(),saveat=Array(t),abstol=abstol, reltol=reltol)
    res = g(sol,p,nothing)
  end
  res_ode_forward = ForwardDiff.gradient(G,p)
  #res_ode_reverse = ReverseDiff.gradient(G,p)

  res_ode_trackeru0, res_ode_trackerp = Zygote.gradient((u0,p)->sum(solve(prob_oop_ode,Tsit5();u0=u0,p=p,abstol=abstol,reltol=reltol,saveat=Array(t),sensealg=TrackerAdjoint()).^2.0/2.0),u₀,p)

  @test isapprox(res_ode_forward[1], sum(@. u₀^2*exp(2*p[1]*t)*t), rtol = 1e-4)
  #@test isapprox(res_ode_reverse[1], sum(@. u₀^2*exp(2*p[1]*t)*t), rtol = 1e-4)
  @test isapprox(res_ode_p'[1], sum(@. u₀^2*exp(2*p[1]*t)*t), rtol = 1e-4)
  @test isapprox(res_ode_p', res_ode_trackerp, rtol = 1e-4)

  # SDE adjoint results (with noise == 0, so should agree with above)

  Random.seed!(seed)
  prob_oop_sde = SDEProblem(f_oop_linear,σ_oop_linear,u₀,trange,p)
  sol_oop_sde = solve(prob_oop_sde,RKMil(interpretation=:Stratonovich),dt=1e-4,adaptive=false,save_noise=true)
  res_sde_u0, res_sde_p = adjoint_sensitivities(sol_oop_sde,
    EulerHeun(),dg!,t,dt=1e-2,sensealg=BacksolveAdjoint())

  @info res_sde_p

  function GSDE1(p)
    Random.seed!(seed)
    tmp_prob = remake(prob_oop_sde,u0=eltype(p).(prob_oop_sde.u0),p=p,
                    tspan=eltype(p).(prob_oop_sde.tspan))
    sol = solve(tmp_prob,RKMil(interpretation=:Stratonovich),dt=tend/10000,adaptive=false,sensealg=DiffEqBase.SensitivityADPassThrough(),saveat=t)
    A = convert(Array,sol)
    res = g(A,p,nothing)
  end
  res_sde_forward = ForwardDiff.gradient(GSDE1,p)
  res_sde_reverse = ReverseDiff.gradient(GSDE1,p)

  res_sde_trackeru0, res_sde_trackerp = Zygote.gradient((u0,p)->sum(solve(prob_oop_sde,RKMil(interpretation=:Stratonovich),dt=tend/1400,adaptive=false,u0=u0,p=p,saveat=Array(t),sensealg=TrackerAdjoint()).^2.0/2.0),u₀,p)

  noise = vec((@. sol_oop_sde.W(tarray)))
  Wfix = [W[1][1] for W in noise]
  @test isapprox(res_sde_forward[1], sum(@. u₀^2*exp(2*p[1]*t)*t), rtol = 1e-4)
  @test isapprox(res_sde_reverse[1], sum(@. u₀^2*exp(2*p[1]*t)*t), rtol = 1e-4)
  @test isapprox(res_sde_p'[1], sum(@. u₀^2*exp(2*p[1]*t)*t), rtol = 1e-4)
  @test isapprox(res_sde_p'[2], sum(@. (Wfix)*u₀^2*exp(2*(p[1])*tarray+2*p[2]*Wfix)), rtol = 1e-4)
  @test isapprox(res_sde_p'[1], res_sde_trackerp[1], rtol = 1e-4)

  # SDE adjoint results (with noise != 0)

  Random.seed!(seed)
  prob_oop_sde2 = SDEProblem(f_oop_linear,σ_oop_linear,u₀,trange,p2)
  sol_oop_sde2 = solve(prob_oop_sde2,RKMil(interpretation=:Stratonovich),dt=tend/1e6,adaptive=false,save_noise=true)

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint())

  @info res_sde_p2

  # test consitency for different switches for the noise Jacobian
  res_sde_u02a, res_sde_p2a = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=false))

  @test isapprox(res_sde_u02, res_sde_u02a, rtol = 1e-6)
  @test isapprox(res_sde_p2, res_sde_p2a, rtol = 1e-6)

  @info res_sde_p2a

  res_sde_u02a, res_sde_p2a = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

  @test isapprox(res_sde_u02, res_sde_u02a, rtol = 1e-6)
  @test isapprox(res_sde_p2, res_sde_p2a, rtol = 1e-6)

  @info res_sde_p2a

  res_sde_u02a, res_sde_p2a = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise()))

  @test isapprox(res_sde_u02, res_sde_u02a, rtol = 1e-6)
  @test isapprox(res_sde_p2, res_sde_p2a, rtol = 1e-6)

  @info res_sde_p2a

  function GSDE2(p)
    Random.seed!(seed)
    tmp_prob = remake(prob_oop_sde2,u0=eltype(p).(prob_oop_sde2.u0),p=p,
                      tspan=eltype(p).(prob_oop_sde2.tspan)
                      #,abstol=abstol, reltol=reltol
                      )
    sol = solve(tmp_prob,RKMil(interpretation=:Stratonovich),dt=tend/1e6,adaptive=false,sensealg=DiffEqBase.SensitivityADPassThrough(),saveat=Array(t))
    A = convert(Array,sol)
    res = g(A,p,nothing)
  end
  res_sde_forward2 = ForwardDiff.gradient(GSDE2,p2)
  res_sde_reverse2 = ReverseDiff.gradient(GSDE2,p2)

  Random.seed!(seed)
  res_sde_trackeru02, res_sde_trackerp2 = Zygote.gradient((u0,p)->sum(solve(prob_oop_sde2,RKMil(interpretation=:Stratonovich),dt=tend/1e3,adaptive=false,u0=u0,p=p,saveat=Array(t),sensealg=TrackerAdjoint()).^2.0/2.0),u₀,p2)

  Wfix = [sol_oop_sde2.W(t)[1][1] for t in tarray]
  resp1 = sum(@. tarray*u₀^2*exp(2*(p2[1])*tarray+2*p2[2]*Wfix))
  resp2 = sum(@. (Wfix)*u₀^2*exp(2*(p2[1])*tarray+2*p2[2]*Wfix))
  resp = [resp1, resp2]

  @test isapprox(res_sde_forward2, resp, rtol = 2e-6)
  @test isapprox(res_sde_reverse2, resp, rtol = 2e-6)
  @test isapprox(res_sde_trackerp2, resp, rtol = 4e-1)

  @test isapprox(res_sde_p2', res_sde_forward2, rtol = 1e-6)
  @test isapprox(res_sde_p2', resp, rtol = 2e-6)

  @info "ForwardDiff" res_sde_forward2
  @info "ReverseDiff" res_sde_reverse2
  @info "Exact" resp
  @info "Adjoint SDE" res_sde_p2

  # consistency check with respect to tracker
  Random.seed!(seed)
  prob_oop_sde2 = SDEProblem(f_oop_linear,σ_oop_linear,u₀,trange,p2)
  sol_oop_sde2 = solve(prob_oop_sde2,RKMil(interpretation=:Stratonovich),dt=tend/1e3,adaptive=false,save_noise=true)
  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e3,adaptive=false,sensealg=BacksolveAdjoint())

  @test isapprox(res_sde_p2', res_sde_trackerp2, rtol = 3e-4)

  # Free memory to help Travis

  Wfix = nothing
  res_sde_forward2 = nothing
  res_sde_reverse2 = nothing
  resp = nothing
  res_sde_trackerp2 = nothing
  res_sde_u02 = nothing
  sol_oop_sde2 = nothing
  res_sde_u02a = nothing
  res_sde_p2a = nothing
  res_sde_p2 = nothing
  sol_oop_sde = nothing
  GC.gc()

  # SDE adjoint results with diagonal noise

  Random.seed!(seed)
  prob_oop_sde2 = SDEProblem(f_oop_linear,σ_oop_linear,[u₀;u₀;u₀],trange,p2)
  sol_oop_sde2 = solve(prob_oop_sde2,EulerHeun(),dt=tend/1e6,adaptive=false,save_noise=true)

  @info "Diagonal Adjoint"

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint())

  res_sde_u03, res_sde_p3 = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=false))

  res_sde_u04, res_sde_p4 = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

  res_sde_u05, res_sde_p5 = adjoint_sensitivities(sol_oop_sde2,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise()))

  @info  res_sde_p2

  sol_oop_sde2 = nothing
  GC.gc()

  @test isapprox(res_sde_p2, res_sde_p3, rtol = 1e-7)
  @test isapprox(res_sde_p2, res_sde_p4, rtol = 1e-7)
  @test isapprox(res_sde_p2, res_sde_p5, rtol = 1e-7)

  @test isapprox(res_sde_u02, res_sde_u03, rtol = 1e-7)
  @test isapprox(res_sde_u02, res_sde_u04, rtol = 1e-7)
  @test isapprox(res_sde_u02, res_sde_u05, rtol = 1e-7)

  @info "Diagonal ForwardDiff"
  res_sde_forward2 = ForwardDiff.gradient(GSDE2,p2)
  #@info "Diagonal ReverseDiff"
  #res_sde_reverse2 = ReverseDiff.gradient(GSDE2,p2)

  #@test isapprox(res_sde_forward2, res_sde_reverse2, rtol = 1e-6)
  @test isapprox(res_sde_p2', res_sde_forward2, rtol = 1e-3)
  #@test isapprox(res_sde_p2', res_sde_reverse2, rtol = 1e-3)

  # u0
  function GSDE3(u)
    Random.seed!(seed)
    tmp_prob = remake(prob_oop_sde2,u0=u,p=eltype(p).(prob_oop_sde2.p),
                      tspan=eltype(p).(prob_oop_sde2.tspan)
                      #,abstol=abstol, reltol=reltol
                      )
    sol = solve(tmp_prob,RKMil(interpretation=:Stratonovich),dt=tend/1e6,adaptive=false,saveat=Array(t))
    A = convert(Array,sol)
    res = g(A,p,nothing)
  end

  @info "ForwardDiff u0"
  res_sde_forward2 = ForwardDiff.gradient(GSDE3,[u₀;u₀;u₀])

  @test isapprox(res_sde_u02, res_sde_forward2, rtol = 1e-4)

end

##
## Inplace
##
@testset "SDE inplace Tests" begin
  f!(du,u,p,t) = du.=p[1]*u
  σ!(du,u,p,t) = du.=p[2]*u

  Random.seed!(seed)
  prob_sde = SDEProblem(f!,σ!,u₀,trange,p2)
  sol_sde = solve(prob_sde,RKMil(interpretation=:Stratonovich),dt=tend/1e6,adaptive=false, save_noise=true)

  # Wfix = [sol_sde.W(t)[1][1] for t in tarray]
  # resp1 = sum(@. tarray*u₀^2*exp(2*(p2[1])*tarray+2*p2[2]*Wfix))
  # resp2 = sum(@. (Wfix)*u₀^2*exp(2*(p2[1])*tarray+2*p2[2]*Wfix))
  # resp = [resp1, resp2]
  # resu0 = sum(@. u₀*exp(2*(p2[1])*tarray+2*p2[2]*Wfix))

  function GSDE(p)
    Random.seed!(seed)
    tmp_prob = remake(prob_sde,u0=eltype(p).(prob_sde.u0),p=p,
                      tspan=eltype(p).(prob_sde.tspan))
    sol = solve(tmp_prob,RKMil(interpretation=:Stratonovich),dt=tend/1e6,adaptive=false,saveat=Array(t))
    A = convert(Array,sol)
    res = g(A,p,nothing)
  end

  res_sde_forward = ForwardDiff.gradient(GSDE,p2)

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint())

  @test isapprox(res_sde_p', res_sde_forward, rtol = 1e-5)

  @info res_sde_p

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=false))

  @info res_sde_p2

  @test isapprox(res_sde_p, res_sde_p2, rtol = 1e-5)
  @test isapprox(res_sde_u0, res_sde_u02, rtol = 1e-5)

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

  @info res_sde_p2

  @test isapprox(res_sde_p, res_sde_p2, rtol = 1e-5) # not broken here because it just uses the vjps
  @test isapprox(res_sde_u0 ,res_sde_u02, rtol = 1e-5)

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise()))

  @info res_sde_p2

  @test isapprox(res_sde_p, res_sde_p2, rtol = 1e-10)
  @test isapprox(res_sde_u0 ,res_sde_u02, rtol = 1e-10)

  # diagonal noise

  #compare with oop version
  f_oop_linear(u,p,t) = p[1]*u
  σ_oop_linear(u,p,t) = p[2]*u
  Random.seed!(seed)
  prob_oop_sde = SDEProblem(f_oop_linear,σ_oop_linear,[u₀;u₀;u₀],trange,p2)
  sol_oop_sde = solve(prob_oop_sde,EulerHeun(),dt=tend/1e6,adaptive=false,save_noise=true)
  res_oop_u0, res_oop_p = adjoint_sensitivities(sol_oop_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint())

  @info res_oop_p

  Random.seed!(seed)
  prob_sde = SDEProblem(f!,σ!,[u₀;u₀;u₀],trange,p2)
  sol_sde = solve(prob_sde,EulerHeun(),dt=tend/1e6,adaptive=false,save_noise=true)

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint())

  @test isapprox(res_sde_p, res_oop_p, rtol = 1e-6)
  @test isapprox(res_sde_u0 ,res_oop_u0, rtol = 1e-6)

  @info res_sde_p

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=false))

  @test isapprox(res_sde_p, res_oop_p, rtol = 1e-6)
  @test isapprox(res_sde_u0 ,res_oop_u0, rtol = 1e-6)

  @info res_sde_p

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

  @test_broken isapprox(res_sde_p, res_oop_p, rtol = 1e-6)
  @test isapprox(res_sde_u0 ,res_oop_u0, rtol = 1e-6)

  @info res_sde_p

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol_sde,EulerHeun(),dg!,Array(t)
      ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise()))

  @test isapprox(res_sde_p, res_oop_p, rtol = 1e-6)
  @test isapprox(res_sde_u0 ,res_oop_u0, rtol = 1e-6)

  @info res_sde_p
end

# scalar noise
@testset "SDE scalar noise tests" begin
  using DiffEqNoiseProcess

  f!(du,u,p,t) = (du .= p[1]*u)
  σ!(du,u,p,t) = (du .= p[2]*u)

  @info "scalar SDE"

  Random.seed!(seed)
  W = WienerProcess(0.0,0.0,0.0)
  u0 = rand(2)

  linear_analytic_strat(u0,p,t,W) = @.(u0*exp(p[1]*t+p[2]*W))

  prob = SDEProblem(SDEFunction(f!,σ!,analytic=linear_analytic_strat),σ!,u0,trange,p2,
    noise=W
    )
  sol = solve(prob,EulerHeun(), dt=tend/1e6, save_noise=true)

  @test isapprox(sol.u_analytic,sol.u, atol=2e-5)

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
    ,dt=tend/1e6,adaptive=false,sensealg=BacksolveAdjoint())

  function compute_grads(sol, scale=1.0)
    xdis = sol(tarray)
    helpu1 = [u[1] for u in xdis.u]
    tmp1 = sum((@. xdis.t*helpu1*helpu1))

    Wtmp = [sol.W(t)[1][1] for t in tarray]
    tmp2 = sum((@. Wtmp*helpu1*helpu1))

    tmp3 = sum((@. helpu1*helpu1))/helpu1[1]

    return [tmp3, scale*tmp3], [tmp1*(1.0+scale^2), tmp2*(1.0+scale^2)]
  end

  @test isapprox(compute_grads(sol, u0[2]/u0[1])[2], res_sde_p', atol=1e-6)
  @test isapprox(compute_grads(sol, u0[2]/u0[1])[1], res_sde_u0, atol=1e-6)
end


module Test_selinv

include("../src/ldlt.jl")
include("../src/selinv.jl")

using Test
using Random
using SparseArrays
using LinearAlgebra

@testset "selinv_from_ldlt" begin
    function test(A,conj,adj)
        Ap,Ai,Av = A.colptr,A.rowval,A.nzval
        Fp,Fi = symbolic_ldlt(Ap,Ai)
        Fv = numeric_ldlt(Ap,Ai,Av,Fp,Fi; conj=conj)
        P = SparseMatrixCSC(size(A)...,Fp,Fi,ones(Bool,length(Fi)))
        Bv = selinv_from_ldlt(Fp,Fi,Fv; conj=conj)
        B = SparseMatrixCSC(size(A)...,Fp,Fi,Bv)
        Bref = inv(Matrix(A)) .* P
        @test B ≈ Bref
    end

    for T in (Float32,Float64,ComplexF32,ComplexF64)
        for (conj,adj) in ((conj,adjoint),(identity,transpose))
            @testset "n = 1" begin
                Random.seed!(1)
                Ap = [1,2]
                Ai = [1]
                Av = rand(T,1)
                A = SparseMatrixCSC(1,1,Ap,Ai,Av)
                A = 2I + A + adj(A)
                test(A,conj,adj)
            end

            @testset "n = 2" begin
                @testset "A = [1 0; 0 1]" begin
                    Random.seed!(1)
                    Ap = [1,2,3]
                    Ai = [1,2]
                    Av = rand(T,2)
                    A = SparseMatrixCSC(2,2,Ap,Ai,Av)
                    A = 2I + A + adj(A)
                    test(A,conj,adj)
                end
                @testset "A = [1 1; 1 0]" begin
                    Ap = [1,3,3]
                    Ai = [1,2]
                    Av = rand(T,2)
                    A = SparseMatrixCSC(2,2,Ap,Ai,Av)
                    A = 2I + A + adj(A)
                    test(A,conj,adj)
                end
                @testset "A = [1 1; 1 1]" begin
                    Ap = [1,3,4]
                    Ai = [1,2,2]
                    Av = rand(T,3)
                    A = SparseMatrixCSC(2,2,Ap,Ai,Av)
                    A = 2I + A + adj(A)
                    test(A,conj,adj)
                end
            end

            @testset "n in 3:100" begin
                @testset for seed = 1:100
                    Random.seed!(seed)
                    n = rand(3:100)
                    A = sprand(n,n,rand())
                    A = 2I + A + adj(A)
                    test(A,conj,adj)
                end
            end
        end
    end
end

end # module


horner2(coef,x) = foldr((u,v) -> u + x*v, 0, coef)


using Test
# using Revise
using PolynomialBasis

PB = PolynomialBasis

function allequal(v1, v2)
    return all(v1 .≈ v2)
end

function allequal(v1, v2, tol)
    np = length(v1)
    f = length(v2) == np
    return f && all([isapprox(v1[i], v2[i], atol = tol) for i = 1:np])
end

coeffs = reshape([1.0], 1, 1)
b0 = PB.LagrangeBasis(0)
b1 = PB.LagrangeBasis(1)
tp1 = PB.LagrangeTensorProductBasis(2, 1)

@test_throws AssertionError PB.InterpolatingPolynomial(coeffs, b1)
ip = PB.InterpolatingPolynomial(coeffs, b0)

@test allequal(ip.coeffs, coeffs)

float_type = typeof(0.0)
ip = PB.InterpolatingPolynomial(float_type, 2, b1)
@test size(ip.coeffs) == (2, 2)

ip = PB.InterpolatingPolynomial(3, tp1)
@test size(ip.coeffs) == (3, 4)

ip = PB.InterpolatingPolynomial(1, 1, 2; start = 0.0)
@test allequal(ip.basis(0.0), [1.0, 0.0, 0.0])
@test allequal(ip.basis(0.5), [0.0, 1.0, 0.0])
@test allequal(ip.basis(1.0), [0.0, 0.0, 1.0])

v = [1.0 2.0 3.0]
PB.update!(ip, v)
@test allequal(ip.coeffs, v)
@test_throws DimensionMismatch PB.update!(ip, 1:4)

ip = PB.InterpolatingPolynomial(1, 2, 2)
PB.update!(ip, repeat([1.0, 2.0, 3.0], 3))
@test ip(0.0, -1.0) ≈ 1.0
@test ip(1.0, 0.0) ≈ 2.0
@test ip(-1.0, 1.0) ≈ 3.0
@test ip([-1.0, 1.0]...) ≈ 3.0

ip = PB.InterpolatingPolynomial(2, 2, 2)
coeffs = [
    1.0 1.0 1.0 2.0 2.0 2.0 3.0 3.0 3.0
    5.0 5.0 5.0 7.0 7.0 7.0 9.0 9.0 9.0
]
PB.update!(ip, coeffs)
@test allequal(ip(-1.0, 0.0), [1.0, 5.0])
@test allequal(ip(0.0, -1.0), [2.0, 7.0])
@test allequal(ip(1.0, 1.0), [3.0, 9.0])

P = PB.InterpolatingPolynomial(1, 1, 2)
f(x) = 3x^2 + 2x - 5
df(x) = 6x + 2
coeffs = f.(P.basis.points)
PB.update!(P, coeffs)
@test all(PB.gradient(P, [-1.0]) .≈ [df(-1.0)])
@test all(PB.gradient(P, [0.19]) .≈ [df(0.19)])

P = PB.InterpolatingPolynomial(1, 2, 2)
f(x, y) = 3x^2 * y + 2y^2 + 2x * y
dfx(x, y) = 6 * x * y + 2 * y
dfy(x, y) = 3x^2 + 4y + 2x
coeffs = [f((P.basis.points[:, i])...) for i = 1:9]
PB.update!(P, coeffs)
@test allequal(PB.gradient(P, [0.1, 0.2]), [dfx(0.1, 0.2) dfy(0.1, 0.2)])
@test allequal(PB.gradient(P, 1, [0.7, -0.3]), dfx(0.7, -0.3))
@test allequal(PB.gradient(P, 2, [-0.5, 0.75]), dfy(-0.5, 0.75))
@test allequal(PB.gradient(P, [0.1, 0.45]), [dfx(0.1, 0.45) dfy(0.1, 0.45)])

poly = PB.InterpolatingPolynomial(1, 1, 3)
f(x) = 7x^3 + 2x^2
coeffs = f.(poly.basis.points)
update!(poly, coeffs)
testh(x) = 42x + 4
@test all([
    allequal(PB.hessian(poly, p)[1], testh(p)) for p in poly.basis.points
])

poly = PB.InterpolatingPolynomial(1, 2, 2)
f(x) = 12x[1]^2 + 7x[1] * x[2] + 11x[2]^2
h = [24, 7, 22]
coeffs = mapslices(f, Array(poly.basis.points), dims = 1)
update!(poly, coeffs)
p = poly.basis.points
@test all([allequal(PB.hessian(poly, p[:, i]), h, 1e2eps()) for i = 1:9])

poly = PB.InterpolatingPolynomial(1, 2, 3)
f(x) = x[1]^3 + 2x[1]^2 * x[2] + 18.0
tp3h(x) = [6x[1] + 4x[2], 4x[1], 0.0]
p = poly.basis.points
coeffs = mapslices(f, Array(p), dims = 1)
update!(poly, coeffs)
@test all([
    allequal(PB.hessian(poly, p[:, i]), tp3h(p[:, i]), 1e3eps()) for i = 1:16
])


using IntervalUnionArithmetic
using Test

@testset "constructor" begin

    x = 1..2
    a = intervalUnion(x)
    @test length(a.v) == 1
    @test a.v[1] == x

    a2 = intervalUnion(1, 2)
    @test length(a.v) == 1
    @test a.v[1] == x

    x2 = 3 .. 4
    b = intervalUnion([x, x2])
    @test length(b.v) == 2
    @test b.v == [x, x2]

    b2 = x ∪ x2
    @test b2.v == b.v

    x3 = interval(5, 6)
    b3 = b2 ∪ x3
    @test b3.v == [x, x2, x3]

    # Remove empties
    c = x ∪ ∅
    @test length(c.v) == 1
    @test c.v[1] == x

    # Remove empties
    c2 = x ∪ ∅ ∪ ∅ ∪ ∅ ∪ ∅ ∪ ∅
    @test length(c2.v) == 1
    @test c2.v[1] == x

    # Don't remove if IntervalUnion is called
    c3 = IntervalUnion([x, ∅])
    @test length(c3.v) == 2
    @test c3.v == [x, ∅]

    # Condense
    d = interval(1, 2) ∪ interval(1.5, 3)
    @test length(d.v) == 1
    @test d.v[1] == interval(1, 3)

    d2 = interval(1, 2) ∪ interval(2, 3)
    @test length(d2.v) == 1
    @test d2.v[1] == interval(1, 3)

end

@testset "close gaps" begin

end


export PiecewiseConstants

"A dictionary of piecewise constant functions on `[0,1]`."
struct PiecewiseConstants{T} <: Dictionary1d{T,T}
    n       ::  Int
end

PiecewiseConstants(n = 10) = PiecewiseConstants{Float64}(n)

## Basic interface

# size should return a tuple (length is defined automatically in terms of size)
size(h::PiecewiseConstants) = (h.n,)

# By implementing `similar`, we support promotion of the domain type, as well
# as resizing of the dictionary.
# If resizing is not possible, one can add a statement such as
# "@assert n == length(dict)", which will cause an error upon any user attempt to resize.
# Same for promotion.
similar(::PiecewiseConstants, ::Type{S}, n::Int) where {S} = PiecewiseConstants{S}(n)

# Evaluation of the basis function:
unsafe_eval_element(h::PiecewiseConstants{T}, idx::Int, x) where {T} = one(T)

# The interpolation_grid returns a set of points such that the interpolation
# problem is uniquely solvable. This is useful for function approximation.
interpolation_grid(h::PiecewiseConstants) = MidpointEquispacedGrid(length(h), support(h))

# The support of our dictionary is the half-open interval [0,1)
support(h::PiecewiseConstants{T}) where {T} = Interval{:closed,:open,T}(0,1)

# We also implement the support of each basis function. By default, this would
# be the default of the dictionary itself.
support(h::PiecewiseConstants{T}, idx::Int) where {T} =
    Interval{:closed,:open,T}(0,one(T)/h.n) .+ (one(T)/h.n*(idx-1))

## Orthogonality properties

# The measure of a dictionary is used when computing projections (inner products
# of a function with a basis function)
measure(h::PiecewiseConstants{T}) where {T} = FourierWeight{T}()

isorthogonal(h::PiecewiseConstants, measure::FourierWeight) = true

# Is our dictionary orthogonal with respect to a discrete measure? We check all points.
function isorthogonal(h::PiecewiseConstants, μ::DiscreteWeight)
    for (i,xi) in zip(ordering(h),points(μ))
        if xi ∉ support(h, i)
            return false
        end
    end
    return true
end


using Gadfly, Distributions

set_default_plot_size(7inch, 3inch)

x = sort!(rand(Poisson(20), 1000))
p1 = plot(x=x, Scale.x_discrete, Geom.histogram)
p2 = plot(x=x, Scale.x_discrete, Geom.histogram(density=true))
hstack(p1, p2)


# Variation of Information

"""
    varinfo(k1::Int, a1::AbstractVector{Int}, k2::Int, a2::AbstractVector{Int})
    varinfo(R::ClusteringResult, k0::Int, a0::AbstractVector{Int})
    varinfo(R1::ClusteringResult, R2::ClusteringResult)

Compute the variation of information between the two clusterings.

Each clustering is provided either as an instance of [`ClusteringResult`](@ref)
subtype or as a pair of arguments:
 - a number of clusters (`k1`, `k2`, `k0`)
 - a vector of point to cluster assignments (`a1`, `a2`, `a0`).
"""
function varinfo(k1::Int, a1::AbstractVector{Int},
                 k2::Int, a2::AbstractVector{Int})

    # check input arguments
    n = length(a1)
    length(a2) == n || throw(DimensionMismatch("Inconsistent array length."))

    # count & compute probabilities
    p1 = zeros(k1)
    p2 = zeros(k2)
    P = zeros(k1, k2)

    for i = 1:n
        @inbounds l1 = a1[i]
        @inbounds l2 = a2[i]
        p1[l1] += 1.0
        p2[l2] += 1.0
        P[l1, l2] += 1.0
    end

    for i = 1:k1
        @inbounds p1[i] /= n
    end
    for i = 1:k2
        @inbounds p2[i] /= n
    end
    for i = 1:(k1*k2)
        @inbounds P[i] /= n
    end

    # compute variation of information

    H1 = entropy(p1)
    H2 = entropy(p2)

    I = 0.0
    for j = 1:k2, i = 1:k1
        pi = p1[i]
        pj = p2[j]
        pij = P[i,j]
        if pij > 0.0
            I += pij * log(pij / (pi * pj))
        end
    end

    return H1 + H2 - I * 2.0
end

varinfo(R::ClusteringResult, k0::Int, a0::AbstractVector{Int}) =
    varinfo(nclusters(R), assignments(R), k0, a0)

varinfo(R1::ClusteringResult, R2::ClusteringResult) =
    varinfo(nclusters(R1), assignments(R1),
            nclusters(R2), assignments(R2))


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.3.1
#     language: julia
#     name: julia-1.3
# ---

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # Ocean Color Data & Model Recipes
#
# Here we set out to compare [CBIOMES model output](https://cbiomes.readthedocs.io/en/latest/) with [ocean color data](https://www.oceancolour.org).
#
# <img src="../figs/cbiomes-01.png" alt="Drawing" style="height: 100px;"/>
#
# _Notebooks are written in [Julia](https://julialang.org) and reproducible via [binder](https://mybinder.org/v2/gh/gaelforget/Cbiomes2019Notebooks/master)._

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### Activate packages for later use
#
# It is assumed that listed packages have aleary been installed using `julia`'s package manager (documentation available [here](https://docs.julialang.org/en/)).
# -

using OceanColorData, Plots, Distributions

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### Model and data wavebands
#
# Currently, the `OC-CCI` [satellite data set](https://esa-oceancolour-cci.org) provides remotely sensed reflectance at 6 wavelengths (`wv_cci` in `nm`) while the `CBIOMES-global` [ocean model](https://cbiomes.readthedocs.io/) outputs irradiance reflectance at 13 wavelengths (`wv_drwn3` in `nm`).
# -

wv_cci=[412, 443, 490, 510, 555, 670]
wv_drwn3=[400,425,450,475,500,525,550,575,600,625,650,675,700];

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### Define interpolation factors
#
# Later on, `jj` and `ww` are used to interpolate, in wavelength space, model output from `wv_drwn3` to `wv_cci`.

# + {"slideshow": {"slide_type": "-"}}
jj=Array{Int64,1}(undef,6)
ww=Array{Float64,1}(undef,6)
for ii=1:6
    tmp=wv_cci[ii].-wv_drwn3
    kk=maximum(findall(tmp.>=0))
    jj[ii]=kk
    ww[ii]=tmp[kk]/(wv_drwn3[kk+1]-wv_drwn3[kk])
end

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### Define a test case
#
# A vector of 13 irradiance reflectances (`Rirr`) is used to represent one space-time location in the model. Later on, we derive `Rrs0` and `Rrs` from `Rirr`. The expected results of the recipe are `ref_Rrs0` and `ref_Rrs`.

# + {"slideshow": {"slide_type": "-"}}
siz=[1,1]

Rirr=Array{Float32,3}(undef,(siz[1],siz[2],13))
Rirr[1,1,:]= 1e-3*[23.7641,26.5037,27.9743,30.4914,28.1356,
    21.9385,18.6545,13.5100,5.6338,3.9272,2.9621,2.1865,1.8015]

ref_Rrs0=Array{Float32,3}(undef,(siz[1],siz[2],13))
ref_Rrs0=1e-3*[4.1753, 4.6640, 4.9270, 5.3781, 4.9558, 3.8505
    , 3.2680, 2.3598, 0.9796, 0.6822, 0.5143, 0.3795, 0.3126]

ref_Rrs=Array{Float32,3}(undef,(siz[1],siz[2],6))
ref_Rrs=1e-3*[4.4099, 4.8533, 5.1247, 4.5137, 3.0864, 0.4064];

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### Convert to remotely sensed reflectance
#
# The following recipe is from [Dutkiewicz et al 2018](https://doi.org/10.5194/bg-15-613-2018).
# -

tmp=Rirr/3
Rrs0=(0.52*tmp)./(1.0 .-1.7*tmp);

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### Interpolate in wavelength space
#
# Interpolating model output from `wv_drwn3` to `wv_cci` allows for direct comparison with satellite data.

# + {"slideshow": {"slide_type": "-"}}
Rrs=Array{Float32,3}(undef,(siz[1],siz[2],6))
for vv=1:6
    tmp0=Rrs0[:,:,jj[vv]]
    tmp1=Rrs0[:,:,jj[vv]+1]
    Rrs[:,:,vv]=tmp0.*(1-ww[vv])+tmp1.*ww[vv]
end

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### Verify result
#
# Let's visualize using the `Plots.jl` package that the resulting `Rrs` matches `ref_Rrs`.

# +
plot(vec(Rrs),linewidth=4,lab="recomputed Rrs")
plot!(ref_Rrs,linewidth=4,ls=:dash,lab="reference result")

#Or, equivalentaly, via OceanColorData.RemotelySensedReflectance :
RrsBis=RemotelySensedReflectance(Rirr,wv_drwn3,wv_cci)
scatter!(vec(RrsBis),lab="OceanColorData.jl")

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### Estimate chlorophyll  from reflectances
#
# Satellite `Chl_a` estimates provided by `OC-CCI` derive from `Rrs` using the blue/green reflectance ratio method as done with model output in the next code bloc (See [Dutkiewicz et al 2018](https://doi.org/10.5194/bg-15-613-2018) for details).

# +
RRSB=max.(Rrs[:,:,2],Rrs[:,:,3]) #blue
RRSG=Rrs[:,:,5] #green
X = log10.(RRSB./RRSG) #ratio of blue to green

C=[0.3272, -2.9940, +2.7218, -1.2259, -0.5683] #OC4 algorithms (SeaWifs, CCI)
a0=C[1]; a1=C[2]; a2=C[3]; a3=C[4]; a4=C[5];
chld=exp10.(a0.+a1*X+a2*X.^2+a3*X.^3+a4*X.^4); #apply polynomial recipe

#Or, equivalently, via OceanColorData.RrsToChla
[chld[1] RrsToChla(vec(Rrs))]

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### Optical classification using reflectances
#
# `Fuzzy logic` classifiers defined in [Moore et al 2009](https://doi.org/10.1016/j.rse.2009.07.016) and [Jackson et al 2017](http://dx.doi.org/10.1016/j.rse.2017.03.036) can be used to assign optical class memberships from an `Rrs` vector. While Moore et al define `n=8` classes using an in-situ database, Jackson et al instead define `n=14` classes using a satellite database. The latter benefits from better data coverage across all of the ecological provinces of the global ocean and is used in `OC-CCI`.
#
# In both cases the classifier is encoded in a mean reflectance spectra (`M[i][1:6]`) and a covariance matrix (`S[i][1:6,1:6]`) provided for each optical class (`i` in `1:n`). Class memberships are then derived by computing the squared Mahalanobis distance to each `M[i]` and passing the result to cumulative chi-squared distribution function (Equations 11 and 12 in [Moore et al 2011](https://doi.org/10.1109/36.942555)).

# + {"slideshow": {"slide_type": "subslide"}}
(M,Sinv)=Moore2009()
M09=Dict("M" => M, "Sinv" => Sinv)
plot(wv_cci,M,w=3); xlabel!("nm"); ylabel!("Rrs")

# + {"slideshow": {"slide_type": "subslide"}}
(M,Sinv)=Jackson2017()
J17=Dict("M" => M, "Sinv" => Sinv)
plot(wv_cci,M,w=3); xlabel!("nm"); ylabel!("Rrs")

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### Compute class memberships
#
# The `fcm` function below takes in a classifier matrix (`M` + `Sinv`) along with a vector of remotely sensed reflectcances (`Rrs`) as input and returns a membership vector (values between 0 and 1).

# + {"slideshow": {"slide_type": "subslide"}}
M09["Membership"]=FuzzyClassification(M09["M"],M09["Sinv"],Rrs)
bar(M09["Membership"]); xlabel!("class"); ylabel!("M09 membership")

# + {"slideshow": {"slide_type": "subslide"}}
J17["Membership"]=FuzzyClassification(J17["M"],J17["Sinv"],Rrs)
bar(J17["Membership"]); xlabel!("class"); ylabel!("J17 membership")


#---------------------------------------------------------------------#
#This function computes the RK Coefficients.
#Written by F.X. Giraldo on April 19, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
include("create_rhs.jl")

function ti_SSP!(q0,u,Dhat,Fhat,intma,periodicity,time,ntime,dt,stages,DFloat)

    #Initialize RK coefficients
    a0=zeros(DFloat,stages)
    a1=zeros(DFloat,stages)
    beta=zeros(DFloat,stages)
    if stages == 1 #RK1-SSP
        a0[1]=1; a1[1]=0; beta[1]=1
    elseif stages == 2 #RK2-SSP
        a0[1]=1; a1[1]=0; beta[1]=1; #SSP
        a0[2]=1.0/2.0; a1[2]=1.0/2.0; beta[2]=1.0/2.0
    elseif stages == 3 #RK3-SSP
        a0[1]=1; a1[1]=0; beta[1]=1;
        a0[2]=3.0/4.0; a1[2]=1.0/4.0; beta[2]=1.0/4.0
        a0[3]=1.0/3.0; a1[3]=2.0/3.0; beta[3]=2.0/3.0
    end

    #Initialize arrays
    q1=copy(q0); qp=copy(q0)
    Npoin=length(q0)

    #Time Integration
    for itime=1:ntime
        time=time + dt
        #RK Stages
        for s=1:stages
            #Create RHS Matrix
#            R=-Dhat*qp
            R=create_rhs(qp,u,Dhat,Fhat,intma,DFloat)

            #Solve System
            @inbounds for I=1:Npoin
                qp[I]=a0[s]*q0[I] + a1[s]*q1[I] + dt*beta[s]*R[I]
            end #s
            if periodicity[Npoin] == periodicity[1]
                qp[Npoin]=qp[1] #periodicity
            end
            #Update
            q1 .= qp
        end #ik
#        println(" itime = ",itime," time = ",time," qmax = ",maximum(qp)," qmin = ",minimum(qp))

        #Update Q
        q0 .= qp
    end #itime

    return (q0)
end


using XGrad
# using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile
using Test

function test_compare(f; inputs...)
    vals = ([val for (name, val) in inputs]...,)
    ctx = Dict()
    df = xdiff(f; ctx=ctx, inputs...)
    # compare_test runs in an older world age than the generated function
    # need to invoke the latest one in the test
    dvals = Base.invokelatest(df, vals...)
    dvals_a = [dvals...]

    # TODO: ReverseDiff doesn't support Julia 0.7 yet, use finite difference instead
    # f_tape = GradientTape(f, vals)
    # compiled_f_tape = compile(f_tape)
    # cfg = GradientConfig(vals)
    # results = map(similar, vals)
    # gradient!(results, compiled_f_tape, vals)
    # results_a = [results...]
    # @test isapprox(results_a, dvals_a[2:end]; atol=0.1)
end

logistic(x) = 1 ./ (1 + exp.(-x))
@diffrule logistic(x::Number) x (logistic(x) .* (1 .- logistic(x))) .* ds

include("tricky.jl")
include("aggr.jl")
include("linreg.jl")
include("ann.jl")
include("autoencoder.jl")
# include("vae.jl")  -- disabled until Distributions are updated to Julia 0.7
include("others.jl")
include("struct.jl")
include("xgrad.jl")


module GridDensities

using Distributions

export GridDensity, pdf

struct GridDensity
    lo::Vector{Float64}
    hi::Vector{Float64}
    bins::Vector{Int64}
    dist::Distributions.Categorical
    volume::Float64
    cart
    linear
end

function GridDensity(data, lo, hi, bins)
    n = prod(bins)
    @assert length(data) == n
    @assert length(lo) == length(hi) == length(bins)
    @assert all(hi .≥ lo)
    volume = prod(hi .- lo)/n # volume for each cell
    p = copy(data) ./ sum(data)
    dist = Categorical(p)
    cart = CartesianIndices(Tuple(bins))
    linear = LinearIndices(Tuple(bins))
    return GridDensity(lo, hi, bins, dist, volume, cart, linear)
end

function _bin(x, lo, hi, bins)
    if x < lo || x > hi
        return 0
    elseif x < hi
        return Int64(floor(bins*(x - lo)/(hi - lo)) + 1)
    else
        return bins
    end
end

function pdf(d::GridDensity, x)
    dims = length(d.bins)
    @assert length(x) == dims
    @assert all(d.lo .≤ x .≤ d.hi)
    inds = [_bin(x[i], d.lo[i], d.hi[i], d.bins[i]) for i = 1:dims]
    ind = d.linear[inds...]
    return Distributions.pdf(d.dist, ind) / d.volume
end

function _rand(bin::Int64, lo, hi, bins)
    return lo + (bin-1+Base.rand())*(hi-lo)/bins
end

function Base.rand(d::GridDensity)
    ind = rand(d.dist)
    dims = length(d.bins)
    coords = d.cart[ind]
    return [_rand(coords[i], d.lo[i], d.hi[i], d.bins[i]) for i = 1:dims]
end

function Base.rand(d::GridDensity, n::Int64)
    return [Base.rand(d) for i = 1:n]
end

end # module


export
    nthfibonacci, nfibonacci,
    allfibonacci, somefibonacci, exactfibonacci

struct FibonacciIterator{T<:Integer}
    x1::T
    x2::T
end

nthfibonacci(n::Int, S::Type = Int) = @pipe allfibonacci(S) |> drop(_, n - one(n)) |> first
nfibonacci(n::Int, S::Type = Int) = collect(exactfibonacci(n, S))

allfibonacci(T::Type = Int) = allfibonacci(one(T), one(T))
allfibonacci(x1::Integer, x2::Integer) = FibonacciIterator(promote(x1, x2)...)

somefibonacci(xmax::Integer) = somefibonacci(xmax, one(xmax), one(xmax))
somefibonacci(xmax::Integer, x1::Integer, x2::Integer) = somefibonacci(promote(xmax, x1, x2)...)
somefibonacci(xmax::T, x1::T, x2::T) where {T<:Integer} = @pipe allfibonacci(x1, x2) |> takewhile(x -> x <= xmax, _)

exactfibonacci(n::Int, S::Type = Int) = exactfibonacci(n, one(S), one(S))
exactfibonacci(n::Int, x1::Integer, x2::Integer) = exactfibonacci(n, promote(x1, x2)...)
exactfibonacci(n::Int, x1::S, x2::S) where {S<:Integer} = @pipe allfibonacci(x1, x2) |> take(_, n)

Base.start(it::FibonacciIterator) = (checked_sub(it.x2, it.x1), it.x1)
Base.next(::FibonacciIterator, state) = (state[2], (state[2], checked_add(state[1], state[2])))
Base.done(::FibonacciIterator, _) = false

Base.eltype(it::FibonacciIterator) = Base.eltype(typeof(it))
Base.eltype(::Type{FibonacciIterator{T}}) where {T} = T

Base.iteratorsize(::FibonacciIterator) = Base.IsInfinite()


import Base:isless
using LightGraphs: Graph, neighbors, degree, nv, ne
using Combinatorics: permutations

export are_isomorphic, enumerate_steiner, label_fst, enumerate_fst

"find unlabeled nodes with at most one unlabeled neighbor"
function find_next_nodes_to_label(graph::Graph, labels::Vector{Int})
    n = nv(graph)
    @assert n == length(labels)
    nolabel = (1:n)[labels .== 0]
    filter(v -> sum(labels[neighbors(graph, v)] .== 0) ≤ 1, nolabel)
end

"combine neighbor labels into tentative labels"
function tentative_labels(graph::Graph, labels::Vector{Int}, nodes::Vector{Int})
    @assert all(labels[nodes] .== 0)
    tent = Tuple{Vector{Int},Int}[]
    for u in nodes
        nblabels = filter(l -> l ≠ 0, labels[neighbors(graph, u)])
        sort!(nblabels)
        push!(tent, (nblabels, u))
    end
    sort!(tent)
    tent
end

"""
Are two given (undirected) trees isomorphic?

Reference: algorithm 2 from Fampa et al.: A specialized branch-and-bound
algorithm for the Euclidean Steiner tree problem in n-space.
"""
function are_isomorphic(tree1::Topology, tree2::Topology)
    @assert is_tree(tree1) && is_tree(tree2)
    g1 = Graph(digraph_from_topology(tree1))
    g2 = Graph(digraph_from_topology(tree2))

    if nv(g1) != nv(g2) || ne(g1) != ne(g2)
        return false
    end

    label1, label2 = fill(0, nv(g1)), fill(0, nv(g2))
    label = 1

    # first label all the leaf nodes
    leaves1 = [n for n=1:nv(g1) if degree(g1, n) == 1]
    leaves2 = [n for n=1:nv(g2) if degree(g2, n) == 1]
    if length(leaves1) ≠ length(leaves2)
        return false
    end
    label1[leaves1] .= label
    label2[leaves2] .= label
    label += 1

    # add more labels until difference or isomorphism is found
    while true
        # find next nodes to give labels to
        next1 = find_next_nodes_to_label(g1, label1)
        next2 = find_next_nodes_to_label(g2, label2)
        if length(next1) ≠ length(next2)
            return false
        elseif length(next1) == 0
            return true
        end

        # try labels from multiset of neighbors' labels
        tent1 = tentative_labels(g1, label1, next1)
        tent2 = tentative_labels(g2, label2, next2)
        if any(t1[1] ≠ t2[1] for (t1, t2) in zip(tent1, tent2))
            return false
        end

        # actually assign the label to smallest of tentatives
        smallest = tent1[1][1]
        label1[[t[2] for t in tent1 if t[1] == smallest]] .= label
        label2[[t[2] for t in tent2 if t[1] == smallest]] .= label
        label += 1
    end
end

"""
Enumerate all representative trees.

This considers only the trees interconnecting the Steiner nodes for full Steiner
trees of n terminals.

Following algorithm 1 in Fampa et al.: A specialized branch-and-bound algorithm
for the Euclidean Steiner tree problem in n-space.
"""
function enumerate_steiner(n::Int)
    nodes = [Node(i,i) for i=1:n-2]
    nclass = fill(0, n - 2) # by number of steiner nodes
    repr = Dict()

    # start with single node
    nclass[1] = 1
    repr[(1, 1)] = Topology(nodes[1:1], Arc[])

    # add more steiner nodes
    for j in 2:(n - 2)
        # look at trees of one size smaller
        for k in 1:nclass[j-1]
            t_base = repr[(j-1, k)]
            g_base = Graph(digraph_from_topology(t_base))

            # and all potential neighbors
            for i in 1:(j-1)
                # maximum degree
                if degree(g_base, i) == 3
                    continue
                end

                # connect new node here then for new tree
                t_new = Topology(nodes[1:j], vcat(t_base.arcs, [Arc(i,j)]))

                # avoid isomorphies
                if any(are_isomorphic(t_new, repr[(j,l)]) for l in 1:nclass[j])
                    continue
                end

                # add to collection
                nclass[j] += 1
                repr[(j, nclass[j])] = t_new
            end
        end
    end
    nclass, repr
end

# needed below
isless(x::T, y::Tuple{T,T}) where {T<:Int} = true
isless(x::Tuple{T,T}, y::T) where {T<:Int} = false

"""
Computes a labeling of a Full Steiner Tree.

This can be used to check for isomorphic trees (having same label).

Following algorithm 4 in Fampa et al.: A specialized branch-and-bound algorithm
for the Euclidean Steiner tree problem in n-space.
"""
function label_fst(fst::Topology)
    n = length(fst.nodes)
    g = Graph(digraph_from_topology(fst))
    terms = [v for v in 1:n if degree(g, v) == 1]
    stein = [v for v in 1:n if degree(g, v) == 3]
    @assert length(terms) + length(stein) == n

    label = Any[-1 for i = 1:n]
    label[terms] = terms

    # rooted binary tree with lexicographical ordering
    root = minimum(terms)
    current_level, next_level = [root], []
    bfs_edges = []
    visited = fill(false, n)
    # BFS
    while length(current_level) > 0
        visited[current_level] .= true
        empty!(next_level)
        for v in current_level
            for w in neighbors(g, v)
                if !visited[w]
                    push!(next_level, w)
                    push!(bfs_edges, (v,w))
                end
            end
        end
        current_level = next_level[:] # copy!
    end

    # work up from the leaves
    reverse!(bfs_edges)
    for depth in 1:2:length(bfs_edges)-1
        t1, h1 = bfs_edges[depth]
        t2, h2 = bfs_edges[depth + 1]
        @assert t1 == t2
        @assert label[t1] == -1
        l1, l2 = label[h1], label[h2]
        label[t1] = l1 < l2 ? (l1, l2) : (l2, l1)
    end

    # all labeled
    @assert all(label .!= -1)

    # label of root's adjacent Steiner node is label for tree
    t, h = bfs_edges[end]
    @assert t == root
    label[h]
end

"""
Enumerate all representative full steiner trees.

Following algorithm 3 in Fampa et al.: A specialized branch-and-bound algorithm
for the Euclidean Steiner tree problem in n-space.
"""
function enumerate_fst(terminals)
    nterms = length(terminals)
    @assert nterms ≥ 3
    nstein = nterms - 2

    stein = [Node(-1,i) for i=1:nstein]
    nodes = vcat(stein, terminals)

    # get all classes for trees of Steiner nodes
    ncls, repr = enumerate_steiner(nterms)
    nclasses = ncls[nstein] # only "full" trees

    # topologies indexed by label
    trees = Dict()

    # early exit for special case of 3 terminals
    if nterms == 3
        tree = Topology(nodes, [Arc(1,2), Arc(1,3), Arc(1,4)])
        trees[label_fst(tree)] = tree
        return trees
    end

    # iterate over all "inner" trees on Steiner nodes
    for c in nclasses
        inner = repr[nstein, c]
        graph = Graph(digraph_from_topology(inner))
        deg1 = [s for s in 1:nstein if degree(graph, s) == 1]
        deg2 = [s for s in 1:nstein if degree(graph, s) == 2]
        deg3 = [s for s in 1:nstein if degree(graph, s) == 3]
        @assert sum(map(length, [deg1, deg2, deg3])) == nstein
        @assert 2 * length(deg1) + length(deg2) == nterms
        K = length(deg1)

        # iterate over permutations of terminals
        termidx = collect((nstein+1):(nstein+nterms))
        for perm in permutations(termidx)
            # TODO: skip those permutations that yield the same tree,
            # i.e. where two terminals are swapped that are connected to
            # the same degree-1 Steiner node

            # connect terminals to Steiner nodes in order
            arcs = inner.arcs[:]
            for k in 1:K
                push!(arcs, Arc(perm[2k - 1], deg1[k]))
                push!(arcs, Arc(perm[2k],     deg1[k]))
            end
            for l in 1:length(deg2)
                push!(arcs, Arc(perm[2K + l], deg2[l]))
            end
            tree = Topology(nodes, arcs)

            # check if new tree is isomorphic to old one
            label = label_fst(tree)
            if haskey(trees, label)
                # TODO: select lexicographically smaller tree
                continue
            end

            trees[label] = tree
        end
    end
    trees
end


# # 1D and 2D Hawkes processes with exponential kernel

#=
In this example, I simulate first a 1D self-exciting Hawkes process and then a 2D one.
The interaction kernel is an exponentially decaying function, defined as:

```math
g(t) = H(t) \,  \frac{1}{\tau} \, \exp\left(- \frac{t}{\tau} \right)
```

where $H(t)$ is the Heaviside function: zero for $t<0$, one for $t\geq 0$.

Note that Kernels are always normalized so that their integral
between $-\infty$ and $+\infty$ is 1.
=#

# ## Initialization
using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:default) ; gr()
using FFTW

using ProgressMeter
using Random
Random.seed!(0)

using HawkesSimulator; const global H = HawkesSimulator

"""
    onedmat(x::R) where R
Generates a 1-by-1 Matrix{R} that contains `x` as only element
"""
function onedmat(x::R) where R
  return cat(x;dims=2)
end;

## #src
# First I define the kernel, and the self-interaction weight.
#
# `myw` is a scaling factor (the weight of the
# autaptic connection). The baseline rate is given by `myinput`

mytau = 0.5  # kernel time constant
myw = onedmat(0.85) # weight: this needs to be a matrix
myinput = [0.7,] # this needs to be a vector
mykernel = H.KernelExp(mytau);

# This is the plot of the (self) interaction kernel
# (before  the scaling by `myw`)

theplot = let ts = range(-1.0,5;length=150),
  y = map(t->H.interaction_kernel(t,mykernel) , ts )
  plot(ts , y ; linewidth=3,leg=false,xlabel="time (s)",
     ylabel="interaction kernel")
end;
plot(theplot)

# Now I build the network, using the simplified constructor
# for one-population networks.
# The object `PopulationState` initializes a population of
# Hawkes neurons of the specified size (1, in this case)
# and the kernel `mykernel` neurons in the same population
# all have the same kernel.
# (this might be an inconvenience if one wants qualitatively different
# self-connections... but it might be fixed in future versions)

popstate = H.PopulationState(mykernel,1)
ntw = H.RecurrentNetwork(popstate,myw,myinput);

## #src
# ## Simulation
# The length of the simulation is measured by the total number of spikes
# here called `num_spikes`.
# The function `flush_trains!(...)` is used to store older spikes as
# history and let them be ignored by the kernels.
# The time parameters should be regulated based on the kernel shape.

function run_simulation!(network,num_spikes,
    t_flush_trigger=300.0,t_flush=100.0)
  t_now = 0.0
  H.reset!(network) # clear spike trains etc
  @showprogress "Running Hawkes process..." for _ in 1:num_spikes
    t_now = H.dynamics_step_singlepopulation!(t_now,network)
    H.flush_trains!(network,t_flush_trigger;Tflush=t_flush)
  end
  H.flush_trains!(network) # flush everything into history
  return t_now
end

n_spikes = 100_000
Tmax = run_simulation!(ntw,n_spikes)
ratenum = H.numerical_rates(popstate;Tend=Tmax)[1]
@info "Simulation completed, mean rate $(round(ratenum;digits=2)) Hz"

##
# ## Visualize raster plot of the events

# the raster plot shows some correlation between the neural activities
# neuron one (lower row) excites neuron two (upper row)

function rasterplot(tlims = (1100.,1120.) )
  _train = popstate.trains_history[1]
  plt=plot()
  train = filter(t-> tlims[1]< t < tlims[2],_train)
  nspk = length(train)
  scatter!(plt,train,fill(0.1,nspk),markersize=30,
      markercolor=:black,markershape=:vline,leg=false)
  return plot!(plt,ylims=(0.0,0.2),xlabel="time (s)")
end

rasterplot()

# event times are always stored in `pops.trains_history`

##  #src
# ## Plot the instantaneous rate

#=
This is the probability of a spike given the past activity up until that
moment. It is usually denoted by $\lambda^*(t)$.
=#

function get_insta_rate(t,popstate)
  _train = popstate.trains_history[1]
  myinput[1] + H.interaction(t,_train,myw[1],popstate.unittype)
end
function plot_instarate(popstate,tlims=(1100,1120))
  tplot = range(tlims...;length=100)
  _train = popstate.trains_history[1]
  tspk = filter(t-> tlims[1]<=t<=tlims[2],_train) # add the exact spiketimes for cleaner plot
  tplot = sort(vcat(tplot,tspk,tspk .- 1E-4))
  plt=plot(xlabel="time (s)",ylabel="instantaneous rate")
  plot!(plt,t->get_insta_rate(t,popstate),tplot;linewidth=2,color=:black)
  scatter!(t->get_insta_rate(t,popstate),tspk;leg=false)
  avg_rate = H.numerical_rates(popstate;Tend=Tmax)[1]
  ylim2 = ylims()[2]
  plot!(plt,tplot, fill(avg_rate,length(tplot)),
     color=:red,linestyle=:dash,ylims=(0,ylim2))
end
plot_instarate(popstate)

# ## Plot the total event counts

# Total number of events as a function of time. It grows linearly, and the
# steepness is pretty much the rate.

function plot_counts(tlims=(0.,1000.))
  avg_rate = H.numerical_rates(popstate;Tend=Tmax)[1]
  tplot = range(tlims...;length=100)
  _train = popstate.trains_history[1]
  nevents(tnow::Real) = count(t-> t <= tnow,_train)
  plt=plot(xlabel="time (s)",ylabel="number of events",leg=false)
  plot!(plt,tplot,nevents.(tplot),color=:black,linewidth=2)
  plot!(plt,tplot , tplot .* avg_rate,color=:red,linestyle=:dash)
end
plot_counts()

# ## Rate
#=
Now I compare the numerical rate with the analytic solution.

The analytic rate corresponds to the stationary solution
of a linear dynamical system (assumung all stationary rates are above zero).
```math
\frac{\mathrm d \mathbf r}{\mathrm d t} = - \mathbf r + W\;r + \mathbf h \quad;
\qquad  r_\infty = (I-W)^{-1} \; \mathbf h
```
=#

rate_analytic = inv(I-myw)*myinput
rate_analytic = rate_analytic[1] # 1-D , just a scalar

@info "Mean rate -  numerical $(round(ratenum;digits=2)), analytic  $(round(rate_analytic;digits=2))"

##
# ## Covariance density

# TODO : write definition
#
# I compute the covariance density it numerically.
# The time inteval `mydt` should not be too small.

mytrain = popstate.trains_history[1]
mydt = 0.1
myτmax = 25.0
mytaus = H.get_times(mydt,myτmax)
ntaus = length(mytaus)
cov_num = H.covariance_self_numerical(mytrain,mydt,myτmax);

##
# Now I compute the covariance density analytically, at higher resolution,
# and I compare the two.

function four_high_res(dt::Real,Tmax::Real) # higher time resolution, longer time
  k1,k2 = 2 , 0.01
  myτmax = Tmax * k1
  dt *= k2
  mytaus = H.get_times(dt,myτmax)
  nkeep = div(length(mytaus),k1)
  myfreq = H.get_frequencies_centerzero(dt,myτmax)
  gfou = myw[1,1] .* H.interaction_kernel_fourier.(myfreq,Ref(popstate)) |> ifftshift
  ffou = let r=rate_analytic
    covf(g) = r/((1-g)*(1-g'))
    map(covf,gfou)
  end
  retf = real.(ifft(ffou)) ./ dt
  retf[1] *= dt  # first element is rate
  return mytaus[1:nkeep],retf[1:nkeep]
end

(taush,covfou)=four_high_res(mydt,myτmax)

function doplot()
  plt = plot(xlabel="time delay (s)",ylabel="Covariance density")
  plot!(plt,mytaus[2:end], cov_num[2:end] ; linewidth=3, label="simulation" )
  plot!(plt,taush[2:end],covfou[2:end]; label="analytic",linewidth=3,linestyle=:dash)
  return plt
end
doplot()

# 1D system completed !
##
# ## Same results, but in a 2D system

myτ = 1/2.33
mywmat = [ 0.31   0.3
           0.9  0.15 ]
myin = [1.0,0.1]
p1 = H.KernelExp(myτ)
ps1 = H.PopulationState(p1,2)
ntw = H.RecurrentNetwork(ps1,mywmat,myin)

# ## Start the simulation
# the function `run_simulation!(...)` has been defined above
# Note that `n_spikes` is the total number of spikes
# among **all** units in the system.

n_spikes = 500_000

Tmax = run_simulation!(ntw,n_spikes,100.0,10.0);

# ## Check the rates
# The analytic rate is from  Eq between 6 and  7 in Hawkes 1971

num_rates = H.numerical_rates(ps1)
myspikes_both = ps1.trains_history

ratefou = let G0 =  mywmat .* H.interaction_kernel_fourier(0,p1)
  inv(I-G0)*myin |> real
end
rate_analytic = inv(I-mywmat)*myin

@info "Total duration $(round(Tmax;digits=1)) s"
@info "Rates are $(round.(num_rates;digits=2))"
@info "Analytic rates are $(round.(rate_analytic;digits=2)) Hz"

##    #src
# ## Covariance density
# there are 4 combinations, therefore I will compare 4 lines.

mydt = 0.1
myτmax = 15.0
mytaus = H.get_times(mydt,myτmax)
ntaus = length(mytaus)
cov_num = H.covariance_density_numerical(myspikes_both,mydt,myτmax)

function doplot()
  plt=plot(xlabel="time delay (s)",ylabel="Covariance density")
  for i in 1:2, j in 1:2
    plot!(plt,mytaus[2:end-1],cov_num[i,j,2:end-1], linewidth = 3, label="cov $i-$j")
  end
  return plt
end

doplot()

# The analytic solution is eq 12 from Hawkes 1971

function four_high_res(dt::Real,Tmax::Real)
  k1 = 2
  k2 = 0.005
  myτmax,mydt = Tmax * k1, dt*k2
  mytaus = H.get_times(mydt,myτmax)
  nkeep = div(length(mytaus),k1)
  myfreq = H.get_frequencies_centerzero(mydt,myτmax)
  G_omega = map(mywmat) do w
    ifftshift( w .* H.interaction_kernel_fourier.(myfreq,Ref(p1)))
  end
  D = Diagonal(ratefou)
  M = Array{ComplexF64}(undef,2,2,length(myfreq))
  Mt = similar(M,Float64)
  for i in eachindex(myfreq)
    G = getindex.(G_omega,i)
    M[:,:,i] = (I-G)\D/(I-G')
  end
  for i in 1:2,j in 1:2
    Mt[i,j,:] = real.(ifft(M[i,j,:]))
    Mt[i,j,2:end] ./= mydt # diagonal of t=0 contains the rate
  end
  return mytaus[1:nkeep],Mt[:,:,1:nkeep]
end

taush,Cfou=four_high_res(mydt,myτmax)

function oneplot(i,j)
  plt=plot(xlabel="time delay (s)",ylabel="Covariance density",title="cov $i - $j")
  plot!(plt,mytaus[2:end],cov_num[i,j,2:end] ; linewidth = 3, label="simulation")
  plot!(plt,taush[2:end],Cfou[i,j,2:end]; linestyle=:dash, linewidth=3, label="analytic")
end

# 1
oneplot(1,1)

# 2
oneplot(1,2)

# 3
oneplot(2,1)

# 4
oneplot(2,2)

# **THE END**

# Literate.markdown("examples/exp_1and2D.jl","docs/src";documenter=true,repo_root_url="https://github.com/dylanfesta/HawkesSimulator.jl/blob/master") #src


import Base.eltypeof, Base.promote_eltypeof, LinearAlgebra.lmul!, LinearAlgebra.rmul!

export opEye, opOnes, opZeros, opDiagonal, opRestriction, opExtension, BlockDiagonalOperator

"""`opEye()`

Identity operator.
```
opI = opEye()
v = rand(5)
@assert opI * v === v
```
"""
struct opEye <: AbstractLinearOperator{Any} end

lmul!(op :: opEye, x::AbstractArray{T, 1}) where {T} = nothing
rmul!(x::AbstractArray{T, 1}, op :: opEye) where {T} = nothing
lmul!(op :: opEye, A::AbstractArray{T, 2}) where {T} = nothing
rmul!(A::AbstractArray{T, 2}, op :: opEye) where {T} = nothing
lmul!(op :: opEye, T::AbstractLinearOperator) = nothing
rmul!(T::AbstractLinearOperator, op :: opEye) = nothing

*(::opEye, x::AbstractArray{T, 1} where {T}) = x
*(x::AbstractArray{T, 1} where {T}, ::opEye) = x
*(::opEye, A::AbstractArray{T, 2} where {T}) = A
*(A::AbstractArray{T, 2} where {T}, ::opEye) = A
*(::opEye, T::AbstractLinearOperator) = T
*(T::AbstractLinearOperator, ::opEye) = T
*(::opEye, T::opEye) = T

function show(io::IO, op::opEye)
  println(io, "Identity operator")
end

function mulOpEye!(res, v, α, β::T, n_min) where T
  if β != zero(T)
    res[1:n_min] .= @views α .* v[1:n_min] .+ β .* res[1:n_min]
    res[n_min+1:end] .= β
  else
    res[1:n_min] .= @views α .* v[1:n_min]
    res[n_min+1:end] .= 0
  end
end

"""
    opEye(Mv)
    opEye(T, n)
    opEye(n)

Identity operator of order `n` and of data type `T` (defaults to `Float64`) with preallocated vector `Mv`.
"""
function opEye(Mv::AbstractVector{T}) where T
  n = length(Mv)
  prod! = @closure (res, v, α, β) -> mulOpEye!(res, v, α, β, n)
  LinearOperator{T}(n, n, true, true, prod!, prod!, prod!, Mv, Mv, Mv)
end

opEye(T::DataType, n::Int) = opEye(zeros(T, n))
opEye(n::Int) = opEye(Float64, n)

# TODO: not type stable
"""
    opEye(T, nrow, ncol)
    opEye(nrow, ncol)

Rectangular identity operator of size `nrow`x`ncol` and of data type `T`
(defaults to `Float64`).
"""
function opEye(Mv::AbstractVector{T}, nrow::I, ncol::I) where {T,I<:Integer}
  length(Mv) == nrow || throw(LinearOperatorException("shape mismatch"))
  if nrow == ncol
    return opEye(T, nrow)
  end
  prod! = @closure (res, v, α, β) -> mulOpEye!(res, v, α, β, min(nrow, ncol))
  return LinearOperator{T}(nrow, ncol, false, false, prod!, prod!, prod!, Mv, Mv, Mv)
end

opEye(T::DataType, nrow::I, ncol::I) where {I<:Integer} =  opEye(zeros(T, nrow), nrow, ncol)
opEye(nrow::I, ncol::I) where {I<:Integer} = opEye(Float64, nrow, ncol)

function mulOpOnes!(res, v, α, β::T) where T
  if β != zero(T)
    res .= (α * sum(v)) .+ β .* res
  else
    res .= (α * sum(v))
  end
end

"""
    opOnes(Mv, nrow, ncol)
    opOnes(T, nrow, ncol)
    opOnes(nrow, ncol)

Operator of all ones of size `nrow`-by-`ncol` of data type `T` (defaults to
`Float64`) and storage vector Mv.
"""
function opOnes(Mv::AbstractVector{T}, nrow::I, ncol::I) where {T,I<:Integer}
  length(Mv) == nrow || throw(LinearOperatorException("shape mismatch"))
  prod! = @closure (res, v, α, β) -> mulOpOnes!(res, v, α, β)
  LinearOperator{T}(nrow, ncol, nrow == ncol, nrow == ncol, prod!, prod!, prod!, Mv, Mv, Mv)
end

opOnes(T::DataType, nrow::I, ncol::I) where {I<:Integer} = opOnes(zeros(T, nrow), nrow, ncol)
opOnes(nrow::I, ncol::I) where {I<:Integer} = opOnes(Float64, nrow, ncol)

function mulOpZeros!(res, v, α, β::T) where T
  if β != zero(T)
    res .*= β
  end
end

"""
    opZeros(Mv, nrow, ncol)
    opZeros(T, nrow, ncol)
    opZeros(nrow, ncol)

Zero operator of size `nrow`-by-`ncol`, of data type `T` (defaults to
`Float64`) and storage vector Mv.
"""
function opZeros(Mv::AbstractVector{T}, nrow::I, ncol::I) where {T,I<:Integer}
  length(Mv) == nrow || throw(LinearOperatorException("shape mismatch"))
  prod! = @closure (res, v, α, β) -> mulOpZeros!(res, v, α, β)
  LinearOperator{T}(nrow, ncol, nrow == ncol, nrow == ncol, prod!, prod!, prod!, Mv, Mv, Mv)
end

opZeros(T::DataType, nrow::I, ncol::I) where {I<:Integer} = opZeros(zeros(T, nrow), nrow, ncol)
opZeros(nrow::I, ncol::I) where {I<:Integer} = opZeros(Float64, nrow, ncol)

function mulSquareOpDiagonal!(res, d, v, α, β::T) where T
  if β != zero(T)
    res .= α .* d .* v .+ β .* res
  else
    res .= α .* d .* v
  end
end

"""
    opDiagonal(Mv, Maw, d)
    opDiagonal(d)

Diagonal operator with the vector `d` on its main diagonal and storage vectors `Mv` and `Maw`.
"""
function opDiagonal(Mv::AbstractVector{T}, Maw::AbstractVector{T}, d::AbstractVector{T}) where {T}
  prod! = @closure (res, v, α, β) -> mulSquareOpDiagonal!(res, d, v, α, β)
  ctprod! = @closure (res, w, α, β) -> mulSquareOpDiagonal!(res, conj.(d), w, α, β)
  LinearOperator{T}(length(d), length(d), true, isreal(d), prod!, prod!, ctprod!, Mv, Mv, Maw)
end

function opDiagonal(d::AbstractVector{T}) where {T}
  nrow = length(d)
  Mv = zeros(T, nrow)
  Maw = isreal(d) ? Mv : zeros(T, nrow)
  return opDiagonal(Mv, Maw, d)
end

#TODO: not type stable
function mulOpDiagonal!(res, d, v, α, β, n_min)
  res[1:n_min] .= @views α .* d[1:n_min] .* v[1:n_min] .+ β .* res[1:n_min]
  res[n_min+1:end] .= 0
end
"""
    opDiagonal(Mv, Mtu, Maw, nrow, ncol, d)
    opDiagonal(nrow, ncol, d)

Rectangular diagonal operator of size `nrow`-by-`ncol` with the vector `d` on
its main diagonal and storage vectors `Mv`, `Mtu`, `Maw`.
"""
function opDiagonal(Mv::AbstractVector{T}, Mtu::AbstractVector{T}, Maw::AbstractVector{T},
                    nrow::I, ncol::I, d::AbstractVector{T}) where {T,I<:Integer}
  nrow == ncol <= length(d) && (return opDiagonal(d[1:nrow]))
  n_min = min(nrow, ncol)
  prod! = @closure (res, v, α, β) -> mulOpDiagonal!(res, d, v, α, β, n_min)
  tprod! = @closure (res, u, α, β) -> mulOpDiagonal!(res, d, u, α, β, n_min)
  ctprod! = @closure (res, w, α, β) -> mulOpDiagonal!(res, conj.(d), w, α, β, n_min)
  LinearOperator{T}(nrow, ncol, false, false, prod!, tprod!, ctprod!, Mv, Mtu, Maw)
end

function opDiagonal(nrow::I, ncol::I, d::AbstractVector{T}) where {T,I<:Integer}
  Mv = zeros(T, nrow)
  Mtu = zeros(T, ncol)
  Maw = isreal(d) ? Mtu : zeros(T, ncol)
  return opDiagonal(Mv, Mtu, Maw, nrow, ncol, d)
end

function mulRestrict!(res, I, v, α, β)
  res .= v[I]
end

function multRestrict!(res, I, u, α, β)
  res .= 0
  res[I] = u
end

"""
    Z = opRestriction(Mv, Mtu, I, ncol)
    Z = opRestriction(I, ncol)
    Z = opRestriction(:, ncol)

Creates a LinearOperator restricting a `ncol`-sized vector to indices `I` with storage Vectors `Mv`, `Mtu`.
The operation `Z * v` is equivalent to `v[I]`. `I` can be `:`.

    Z = opRestriction(k, ncol)

Alias for `opRestriction([k], ncol)`.
"""
function opRestriction(Mv::AbstractVector{T}, Mtu::AbstractVector{T}, Idx::LinearOperatorIndexType{I}, ncol::I) where {T,I<:Integer}
  all(1 .≤ Idx .≤ ncol) || throw(LinearOperatorException("indices should be between 1 and $ncol"))
  nrow = length(Idx)
  prod! = @closure (res, v, α, β) -> mulRestrict!(res, Idx, v, α, β)
  tprod! = @closure (res, u, α, β) -> multRestrict!(res, Idx, u, α, β)
  return LinearOperator{I}(nrow, ncol, false, false, prod!, tprod!, tprod!, Mv, Mtu, Mtu)
end

opRestriction(Idx::LinearOperatorIndexType{I}, ncol::I) where {I<:Integer} = opRestriction(zeros(length(Idx)), zeros(ncol), Idx, ncol)

opRestriction(::Colon, ncol::I) where {I<:Integer} = opEye(I, ncol)

opRestriction(k::I, ncol::I) where {I<:Integer} = opRestriction([k], ncol)

"""
    Z = opExtension(I, ncol)
    Z = opExtension(:, ncol)

Creates a LinearOperator extending a vector of size `length(I)` to size `ncol`,
where the position of the elements on the new vector are given by the indices
`I`.
The operation `w = Z * v` is equivalent to `w = zeros(ncol); w[I] = v`.

    Z = opExtension(k, ncol)

Alias for `opExtension([k], ncol)`.
"""
opExtension(Idx::LinearOperatorIndexType{I}, ncol::I) where {I<:Integer} = opRestriction(Idx, ncol)'

opExtension(::Colon, ncol::I) where {I<:Integer} = opEye(I, ncol)

opExtension(k::I, ncol::I) where {I<:Integer} = opExtension([k], ncol)

# indexing for linear operators
import Base.getindex
function getindex(
  op::AbstractLinearOperator,
  rows::Union{LinearOperatorIndexType{I}, I, Colon},
  cols::Union{LinearOperatorIndexType{I}, I, Colon},
) where {I<:Integer}
  R = opRestriction(rows, size(op, 1))
  E = opExtension(cols, size(op, 2))
  return R * op * E
end

eltypeof(op::AbstractLinearOperator) = eltype(op)  # need this for promote_eltypeof

"""
    BlockDiagonalOperator(Mv, Mtu, Maw, M1, M2, ..., Mn)
    BlockDiagonalOperator(M1, M2, ..., Mn)

Creates a block-diagonal linear operator:

    [ M1           ]
    [    M2        ]
    [       ...    ]
    [           Mn ]
"""
function BlockDiagonalOperator(Mv::AbstractVector{S}, Mtu::AbstractVector{S}, Maw::AbstractVector{S}, ops...) where {S}
  nrow = ncol = 0
  for op ∈ ops
    m, n = size(op)
    nrow += m
    ncol += n
  end
  T = promote_eltypeof(ops...)

  function prod!(y, x, α, β)
    # y = zeros(T, nrow)
    k = 0
    j = 0
    for op ∈ ops
      m, n = size(op)
      # y[(k + 1):(k + m)] .= op * x[(j + 1):(j + n)]
      @views mul!(y[(k + 1):(k + m)] , op, x[(j + 1):(j + n)], α, β)
      k += m
      j += n
    end
    # y
  end

  function tprod!(y, x, α, β)
    # y = zeros(T, ncol)
    k = 0
    j = 0
    for op ∈ ops
      m, n = size(op)
      # y[(k + 1):(k + n)] .= transpose(op) * x[(j + 1):(j + m)]
      @views mul!(y[(k + 1):(k + n)], transpose(op), x[(j + 1):(j + m)], α, β)
      k += n
      j += m
    end
    # y
  end

  function ctprod!(y, x, α, β)
    # y = zeros(T, ncol)
    k = 0
    j = 0
    for op ∈ ops
      m, n = size(op)
      # y[(k + 1):(k + n)] .= op' * x[(j + 1):(j + m)]
      @views mul!(y[(k + 1):(k + n)], adjoint(op), x[(j + 1):(j + m)], α, β)
      k += n
      j += m
    end
    # y
  end

  symm = all((issymmetric(op) for op ∈ ops))
  herm = all((ishermitian(op) for op ∈ ops))
  LinearOperator{T}(nrow, ncol, symm, herm, prod!, tprod!, ctprod!, Mv, Mtu, Maw)
end

function BlockDiagonalOperator(ops...)
  nrow = ncol = 0
  for op ∈ ops
    m, n = size(op)
    nrow += m
    ncol += n
  end
  T = promote_eltypeof(ops...)
  symm = all((issymmetric(op) for op ∈ ops))
  herm = all((ishermitian(op) for op ∈ ops))
  Mv = zeros(T, nrow)
  Mtu = symm ? Mv : zeros(T, ncol)
  Maw = herm ? Mv : zeros(T, ncol)
  return BlockDiagonalOperator(Mv, Mtu, Maw, ops...)
end


# Computes a polynomial Lyapunov for the switched system from Example 2.8 of
# [PJ08].
#   x_{k+1} = A_σ x_k, σ ∈ {1, 2}
# Where
#   A1 = [1 0   A2 = [0  1
#         1 0]        0 -1]
# There is a quartic Lyapunov function for the system but only a quadratic for
# the system with A1' = A1/√2, A2' = A2/√2.
# See https://github.com/blegat/SwitchOnSafety.jl for computing Lyapunov
# functions for switched systems.
#
# [PJ08] P. Parrilo and A. Jadbabaie
# Approximation of the joint spectral radius using sum of squares
# Linear Algebra and its Applications, Elsevier, 2008, 428, 2385-2402
# Inspired from the construction of:
# Ando, T. and Shih, M.-h.
# Simultaneous Contractibility.
# SIAM Journal on Matrix Analysis & Applications, 1998, 19, 487

using LinearAlgebra, Test
using SumOfSquares
using DynamicPolynomials

# Search for polynomial Lyapunov functions of degree `2d` for `A1/γ` and `A2/γ`
function lyapunov_switched_system_test(
    optimizer, config::MOIT.Config, degree::Int, γ, feasible::Bool, basis)

    atol = config.atol
    rtol = config.rtol

    A1 = [1 0
          1 0]
    A2 = [0  1
          0 -1]

    model = _model(optimizer)

    @polyvar x[1:2]

    # p is should be strictly positive so we create a nonnegative `p`
    # and sum it with a strictly positive `q`.
    # Since the problem is homogeneous (i.e. given any `λ > 0`, `p` is
    # feasible iff `λp` is feasible), this is wlog.
    p0 = @variable(model, variable_type=SOSPoly(basis(monomials(x, degree))))
    q = GramMatrix(SOSDecomposition(x.^degree))

    # Keep `p` in a `GramMatrix` form while `q + p0` would transform it to
    # a polynomial. It is not mandatory to keep it in its `GramMatrix` form
    # but it will allow us to test `JuMP.value(::GramMatrix{JuMP.AffExpr})`.
    p = gram_operate(+, q, p0)

    c1 = @constraint(model, p - p(x => (A1 / γ) * vec(x)) in SOSCone(), basis = basis)
    c2 = @constraint(model, p - p(x => (A2 / γ) * vec(x)) in SOSCone(), basis = basis)

    JuMP.optimize!(model)

    if feasible
        @test JuMP.termination_status(model) == MOI.OPTIMAL
        @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
        @test all(eigvals(Matrix(getmat(JuMP.value(p0)))) .≥ -atol)
        @test JuMP.value(p0).basis isa basis
        @test getmat(JuMP.value(p)) ≈ getmat(gram_operate(+, q, JuMP.value(p0))) atol=atol rtol=rtol
        @test gram_matrix(c1).basis isa basis
        @test gram_matrix(c2).basis isa basis
    else
        @test JuMP.termination_status(model) == MOI.INFEASIBLE
        @test JuMP.dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE
        for (μ1, μ2) in [(JuMP.dual(c1), JuMP.dual(c2)),
                         (moments(c1), moments(c2))]
            # The dual constraint should work on any polynomial.
            # Let's test it with q
            lhs = dot(μ1, q(x => A1 * vec(x))) + dot(μ2, q(x => A2 * vec(x)))
            rhs = dot(μ1, q) + dot(μ2, q)
            @test atol + rtol * max(abs(lhs), abs(rhs)) + lhs >= rhs
        end
        @test moment_matrix(c1).basis isa basis
        @test moment_matrix(c2).basis isa basis
    end
end

quadratic_infeasible_lyapunov_switched_system_test(optimizer, config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 - ε, false, MonomialBasis)
sd_tests["quadratic_infeasible_lyapunov_switched_system"] = quadratic_infeasible_lyapunov_switched_system_test
quadratic_infeasible_scaled_lyapunov_switched_system_test(optimizer, config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 - ε, false, ScaledMonomialBasis)
sd_tests["quadratic_infeasible_scaled_lyapunov_switched_system"] = quadratic_infeasible_scaled_lyapunov_switched_system_test
quadratic_feasible_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 + ε, true, MonomialBasis)
sd_tests["quadratic_feasible_lyapunov_switched_system"]   = quadratic_feasible_lyapunov_switched_system_test
quadratic_feasible_scaled_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 + ε, true, ScaledMonomialBasis)
sd_tests["quadratic_feasible_scaled_lyapunov_switched_system"]   = quadratic_feasible_scaled_lyapunov_switched_system_test
quartic_infeasible_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 - ε, false, MonomialBasis)
sd_tests["quartic_infeasible_lyapunov_switched_system"]   = quartic_infeasible_lyapunov_switched_system_test
quartic_infeasible_scaled_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 - ε, false, ScaledMonomialBasis)
sd_tests["quartic_infeasible_scaled_lyapunov_switched_system"]   = quartic_infeasible_scaled_lyapunov_switched_system_test
quartic_feasible_lyapunov_switched_system_test(optimizer,     config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 + ε, true, MonomialBasis)
sd_tests["quartic_feasible_lyapunov_switched_system"]     = quartic_feasible_lyapunov_switched_system_test
quartic_feasible_scaled_lyapunov_switched_system_test(optimizer,     config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 + ε, true, ScaledMonomialBasis)
sd_tests["quartic_feasible_scaled_lyapunov_switched_system"]     = quartic_feasible_scaled_lyapunov_switched_system_test


# run a Verlet tarjectory from a double well in 1D
using Plots
using Printf
using LinearAlgebra
using ForwardDiff
using QuadGK
using BasicMD

include("potentials.jl")

M = 1.0;
q₀ = [-0.5];
p₀ = [1.0];
Δt = 1e-1;
n_iters = 10^3;

V = x->DoubleWell(x);
cfg = ForwardDiff.GradientConfig(V, q₀);
gradV! = (gradV, x)-> ForwardDiff.gradient!(gradV, V, x, cfg);

sampler = Verlet(gradV!, M, Δt);

Q₀ = copy(q₀);
P₀ = copy(p₀);
sample_trajectory!([Q₀, P₀], sampler, options=MDOptions(n_iters=n_iters));
@printf("In Place (Q,P) after %d iterations: (%g,%g)\n",n_iters, Q₀[1], P₀[1]);

Xvals = sample_trajectory([q₀, p₀], sampler,
    options=MDOptions(n_iters=n_iters,n_save_iters=n_iters));
Q = Xvals[end][1][1]
P = Xvals[end][2][1]
@printf("(Q,P) after %d iterations: (%g,%g)\n",n_iters, Q[1], P[1]);

Xvals = sample_trajectory([q₀, p₀], sampler, options=MDOptions(n_iters=n_iters));
Q = Xvals[end][1][1]
P = Xvals[end][2][1]
@printf("(Q,P) after %d iterations: (%g,%g)\n",n_iters, Q[1], P[1]);
plot(Δt *(1:n_iters), [X[1][1] for X in Xvals],label="q")
plot!(Δt *(1:n_iters), [X[2][1] for X in Xvals],label="p")
xlabel!("t")


# This file is a part of SimilaritySearch.jl
# License is Apache 2.0: https://www.apache.org/licenses/LICENSE-2.0.txt

export cosine_similarity, cosine_distance, angle_distance

using LinearAlgebra
import LinearAlgebra: normalize, normalize!

function normalize(X::AbstractVector{S}) where S<:AbstractVector
    normalize.(X)
end

function normalize!(X::AbstractVector{S}) where S<:AbstractVector
    for x in X
        normalize!(x)
    end

    X
end

function normalize!(X::AbstractMatrix)
    for x in eachcol(X)
        normalize!(x)
    end

    X
end

"""
    cosine_distance(a, b)::Float64

Computes the cosine distance between two vectors, it expects normalized vectors (see [normalize!](@ref) method).
Please use angle_distance if you are expecting a metric function (cosine_distance is a faster
alternative whenever the triangle inequality is not needed)
"""
function cosine_distance(a, b)::Float64
    1 - dot(a, b)
end

const π_2 = π / 2
"""
    angle_distance(a, b)::Float64

Computes the angle  between two SparseVector objects (sparse vectors).
It supposes that all vectors are normalized (see `normalize!` function)

"""
function angle_distance(a, b)::Float64
    d = dot(a, b)

    if d <= -1.0
        return π
    elseif d >= 1.0
        return 0.0
    elseif d == 0  # turn around for zero vectors, in particular for denominator=0
        return π_2
    else
        return acos(d)
    end
end


export GaussianNoise

"""
    GaussianNoise(σ = 1.0)

Returns a Gaussian noise with standard deviation `σ`.
"""
struct GaussianNoise <: Sequence
  σ
end

GaussianNoise() = GaussianNoise(1.0)

sample(seq::GaussianNoise, t::AbstractVector) = seq.σ * randn(length(t))


struct SimpleNormal{Tm, Tσ}
    m::Tm
    σ::Tσ
end

function AbstractGPs.logpdf(d::SimpleNormal, y::Real)
    m = Zygote.literal_getfield(d, Val(:m))
    σ = Zygote.literal_getfield(d, Val(:σ))
    return -(log(2π) + 2log(σ) + ((y - m) / σ)^2) / 2
end

function ChainRulesCore.rrule(::typeof(AbstractGPs.mean), d::T) where {T<:Normal}
    function mean_pullback(Δ::Real)
        return NoTangent, Tangent{T}(μ=Δ)
    end
    return d.μ, mean_pullback
end

function ChainRulesCore.rrule(::typeof(AbstractGPs.var), d::T) where {T<:Normal}
    function var_pullback(Δ::Real)
        NoTangent, Tangent{T}(σ=2Δ * d.σ)
    end
    return d.σ^2, var_pullback
end

"""
    gaussian_reconstruction_term(
        y::AbstractVector{<:Real},
        σ²::AbstractVector{<:Real},
        m̃::AbstractVector{<:Real},
        σ̃²::AbstractVector{<:Real},
    )

Computes
```julia
∫ logpdf(Normal(f, σ²[n]), y[n]) pdf(Normal(m̃[n], σ̃²[n]), f) df
```
for each `n ∈ eachindex(y)`.
"""
function gaussian_reconstruction_term(
    y::AbstractVector{<:Real},
    Σy::Diagonal{<:Real},
    mq::AbstractVector{<:Real},
    Σq::Diagonal{<:Real},
)
    return sum(
        map(
            (y, σ², m̃, σ̃²) -> logpdf(SimpleNormal(m̃, sqrt(σ²)), y) - σ̃² / (2σ²),
            y, diag(Σy), mq, diag(Σq),
        ),
    )
end

"""
    approx_posterior(
        f::AbstractGP,
        x::AbstractVector,
        η1::AbstractVector{<:Real},
        η2::AbstractVector{<:Real},
    )

Compute the approximate posterior. This is just a posterior GP some kind, which is another
AbstractGP.
"""
function approx_posterior(
    f::AbstractGP,
    x::AbstractVector,
    η1::AbstractVector{<:Real},
    η2::Diagonal{<:Real},
)
    y, Σ = canonical_from_natural(η1, η2)
    return posterior(f(x, Σ + 1e-6 * I), y)
end

"""
    update_approx_posterior(
        f::AbstractGP,
        x::AbstractVector,
        η1::AbstractVector{<:Real},
        η2::AbstractVector{<:Real},
        r,
        ρ::Real,
    )
"""
function update_approx_posterior(
    f::AbstractGP,
    x::AbstractVector,
    η1::AbstractVector{<:Real},
    η2::Diagonal{<:Real},
    r,
    ρ::Real,
)
    # Check that the step size makes sense.
    (ρ > 1 || ρ <= 0) && throw(error("Bad step size"))

    # Compute the approximate posterior marginals.
    mq, σ²q = mean_and_var(approx_posterior(f, x, η1, η2)(x))

    # Compute both of the expectation parameters.
    m1, m2 = expectation_from_canonical(mq, Diagonal(σ²q))

    # Compute the gradient w.r.t. both of the expectation parameters. This is equivalent to
    # the natural gradient w.r.t. the natural parameters.
    g1, g2 = Zygote.gradient(
        (m1, m2) -> begin
            y, σ² = canonical_from_expectation(m1, m2)
            return r(y, σ²)
        end,
        m1, m2,
    )

    # Perform a step of NGA in the first natural pseudo-observation vector.
    η1_new = (1 - ρ) .* η1 .+ ρ .* g1
    η2_new = (1 - ρ) .* η2 .+ ρ .* g2

    return η1_new, η2_new, sqrt(sum(abs2, g1 - η1) + sum(abs2, g2 - η2))
end

"""
    optimise_approx_posterior(
        f::AbstractGP,
        x::AbstractVector,
        η1::AbstractVector{<:Real},
        η2::AbstractVector{<:Real},
        ∇r,
        ρ::Real;
        max_iterations=1_000,
        gtol=1e-8,
    )
"""
function optimise_approx_posterior(
    f::AbstractGP,
    x::AbstractVector,
    η1::AbstractVector{<:Real},
    η2::Diagonal{<:Real},
    r,
    ρ::Real;
    max_iterations=1_000,
    tol=1e-8,
)
    # Perform initial iteration.
    η1, η2, delta_norm = update_approx_posterior(f, x, η1, η2, r, ρ)
    iteration = 1

    # Iterate further until convergence met or max iterations exceeded.
    while delta_norm > tol && iteration < max_iterations
        η1, η2, delta_norm = update_approx_posterior(f, x, η1, η2, r, ρ)
        iteration += 1
    end
    return η1, η2, iteration, delta_norm
end

"""
    elbo(
        f::AbstractGP,
        x::AbstractVector,
        η1::AbstractVector{<:Real},
        η2::AbstractVector{<:Real},
        r,
    )

Compute the evidence lower bound associated with the GP `f`, with reconstruction term `r`,
variational parameters `η1` and `η2`, and collection of inputs `x`.
"""
function AbstractGPs.elbo(
    f::AbstractGP,
    x::AbstractVector,
    η1::AbstractVector{<:Real},
    η2::Diagonal{<:Real},
    r,
)
    ỹ, σ̃² = canonical_from_natural(η1, η2)
    mq, σ²q = mean_and_var(approx_posterior(f, x, η1, η2)(x))
    return logpdf(f(x, σ̃²), ỹ) + r(mq, Diagonal(σ²q)) - gaussian_reconstruction_term(ỹ, σ̃², mq, Diagonal(σ²q))
end

abstract type AbstractIntegrater end

"""
    GaussHermiteQuadrature(num_points::Int)

Specify that the expectation required for the reconstruction term is to be computed using
Gauss-Hermite quadrature with `num_points` quadrature points.
"""
struct GaussHermiteQuadrature <: AbstractIntegrater
    num_points::Int
end

"""
    UnivariateFactorisedLikelihood(build_lik)

A likelihood function which factorises across inputs.
Applying this likelihood to any vector `f::AbstractVector{<:Real}` yields a `Product`
distribution.
Exposing this structure is useful when working with quadrature methods.
"""
struct UnivariateFactorisedLikelihood{Tf}
    build_lik::Tf
end

function (l::UnivariateFactorisedLikelihood)(x::AbstractVector{<:Real})
    return AbstractGPs.Product(map(l.build_lik, x))
end

function build_reconstruction_term(
    integrater::GaussHermiteQuadrature,
    latent_gp::LatentGP,
    y::AbstractVector,
)
    integrands = build_integrands(latent_gp, y)
    function reconstruction_term(mq::AbstractVector{<:Real}, Σq::Diagonal{<:Real})
        return sum(batch_quadrature(integrands, mq, map(sqrt, diag(Σq)), integrater.num_points))
    end
    return reconstruction_term
end

function build_integrands(
    latent_gp::LatentGP{<:AbstractGP, <:UnivariateFactorisedLikelihood},
    y::AbstractVector,
)
    lik = latent_gp.lik.build_lik
    return map(y_ -> (f -> logpdf(lik(f), y_)), y)
end

"""
    optimize_elbo(
        build_latent_gp,
        integrater::AbstractIntegrater,
        x::AbstractVector,
        y::AbstractVector,
        θ0::AbstractVector{<:Real},
        optimiser,
        options,
    )

Optimise the elbo w.r.t. model parameters.

# Arguments
- `build_latent_gp`: a unary function accepting an `AbstractVector{<:Real}` returning a
    LatentGP
- `integrater`::AbstractIntegrater: an object specifying how to perform inference
- `x::AbstractVector`: a collection of inputs
- `y::AbstractVector`: a collection of outputs
- `θ0::AbstractVector{<:Real}`: initial model parameter values
- `optimiser`: an optimiser that can be passed to `Optim.optimize`
- `options`: options to be passed to `Optim.optimize`

# Returns
- `AbstractGP`: the approximate posterior at the optimum
- `NamedTuple`: results summary containing various useful bits of information
"""
function optimize_elbo(
    build_latent_gp,
    integrater::AbstractIntegrater,
    x::AbstractVector,
    y::AbstractVector,
    θ0::AbstractVector{<:Real},
    optimiser,
    options,
)
    # Initialise variational parameters. We'll be mutating these to enable warm-starts
    # during each outer iteration of the algorithm.
    __η1 = zeros(length(x))
    __η2 = Diagonal(-ones(length(x)))

    function objective(θ::AbstractVector{<:Real})

        # Unflatten θ and build the model at the current hyperparameters.
        l = build_latent_gp(θ)
        r = build_reconstruction_term(integrater, l, y)

        # Optimise the approximate posterior. Drop the gradient because we're
        # differentiating through the optimum.
        η1_opt, η2_opt = Zygote.ignore() do
            η1, η2, iters, delta = optimise_approx_posterior(
                l.f, x, __η1, __η2, r, 1; tol=1e-4,
            )
            __η1 .= η1
            __η2 .= η2
            return η1, η2
        end

        # Compute the negation of the elbo.
        return -elbo(l.f, x, η1_opt, η2_opt, r)
    end

    training_results = Optim.optimize(
        objective, θ -> only(Zygote.gradient(objective, θ)), θ0, optimiser, options;
        inplace=false,
    )

    l = build_latent_gp(training_results.minimizer)
    r = build_reconstruction_term(integrater, l, y)
    η1, η2, iters, delta = optimise_approx_posterior(l.f, x, __η1, __η2, r, 1; tol=1e-4)
    approx_post = approx_posterior(l.f, x, η1, η2)

    results_summary = (
        training_results=training_results,
        η1=__η1,
        η2=__η2,
        iters=iters,
        delta=delta,
    )

    return approx_post, results_summary
end


module SimpleRoots

export Bisection, Secant, Brent

export Lower, Upper, Both

export findzero, quad

abstract type AbstractRootFinder end

abstract type AbstractBisection <: AbstractRootFinder end

bracket(method::AbstractBisection) = method.bracket

struct Bisection{B} <: AbstractBisection
    bracket::B
end
Bisection(arg1, arg2) = Bisection((arg1, arg2))

"""
    Secant(bracket)
    Secant(bracket...)

Secant method of bracketed root-finding.

Bracket can be a length 2 vector, a tuple or 2 argument
"""
struct Secant{B} <: AbstractBisection
    bracket::B
end
Secant(arg1, arg2) = Secant((arg1, arg2))

struct Brent{B} <: AbstractBisection
    bracket::B
end
Brent(arg1, arg2) = Brent((arg1, arg2))

"""
    findzero(f, bracket; atol=1e-7, max_iter=100)
    findzero(f, method::AbstractBisection; atol=1e-7, max_iter=100)

Find roots using the passed in method. If a bracket is passed in without
specifying the method, use the [`Secant`](@ref) method.
"""
findzero(f, bracket; kwargs...) = findzero(f, Secant(bracket); kwargs...)
findzero(f, method::AbstractRootFinder; atol=1e-7, max_iter=100) =
    findzero(f, method, atol, max_iter)

"""
    findzero(f, method::Secant, atol, max_iter)
    findzero(f, method::Secant; atol=1e-7, max_iter=100)

Find root using the secant method.
"""
function findzero(f, method::Secant, atol, max_iter)
    local x
    x0, x1 = bracket(method)
    y0, y1 = f(x0), f(x1)

    for _ in 1:max_iter
        x = x1 - y1 * (x1 - x0) / (y1 - y0)
        if abs(x - x1) < atol
            return x, true
        end
        x0 = x1
        y0 = y1
        x1 = x
        y1 = f(x1)
    end
    return x, false
end

"""
    findzero(f, method::Bisection, atol, max_iter)
    findzero(f, method::Bisection; atol=1e-7, max_iter=100)

Find root using the bisection method.
"""
function findzero(f, method::Bisection, atol, max_iter)
    a, b = bracket(method)
    fa = f(a)
    local c
    for _ in 1:max_iter
        if (b - a) <= atol
            return b, true
        end

        c = (a + b)/2
        fc = f(c)
        if fa * fc > zero(fa * fc)
            a = c # Root is in the right half of [a, b].
            fa = fc
        elseif fc == zero(fc)
            return c, true
        else
            b = c # Root is in the left half of [a, b].
        end
    end

    return b, false
end

"""
    findzero(f, method::Brent, atol, maxiter)
    findzero(f, method::Brent; atol=1e-7, max_iter=100)

Find root using Brents method. Returns a tuple of the found value
and a Bool specifying wether the root was found within the tolerance, or not.
"""
function findzero(f, method::Brent, atol, maxiter)
    local e, d
    epsval = eps()
    a, b = bracket(method)
    fa, fb = f(a), f(b)
    c = b
    fc = fb
    for iter = 1:maxiter
        if (fb > zero(fb) && fc > zero(fc)) || (fb < zero(fb) && fc < zero(fc))
            c = a
            fc = fa
            d = b - a
            e = d
        end
        if abs(fc) < abs(fb)
            a = b
            b = c
            c = a
            fa = fb
            fb = fc
            fc = fa
        end
        tol1 = 2epsval * abs(b) + atol/2
        xm = (c - b)/2
        if (abs(xm) <= tol1 || fb == zero(fb))
            return b, true
        end
        if abs(e) >= tol1 && abs(fa) > abs(fb)
            s = fb / fa
            if (a == c)
                p = 2xm * s
                q = 1 - s
            else
                q = fa / fc
                r = fb / fc
                p = s * (2xm * q * (q - r) - (b - a) * (r - 1))
                q = (q - 1) * (r - 1) * (s - 1)
            end
            if p > zero(p)
                q = -q
            end
            p = abs(p)
            if (2p < min(3xm * q - abs(tol1 * q), abs(e * q)))
                e = d
                d = p / q
            else
                d = xm
                e = d
            end
        else
            d = xm
            e = d
        end
        a = b
        fa = fb
        if (abs(d) > tol1)
            b = b + d
        else
            b = b + tol1 * sign(xm)
        end
        fb = f(b)
    end

    return b, false
end

"""
Used to specifiy the result of quad, either
[`Lower`](@ref), [`Upper`](@ref) or [`Both`](@ref)
"""
abstract type QuadraticResult end

"""
Return the lower result value
"""
struct Lower <: QuadraticResult end

"""
Return the upper result value
"""
struct Upper <: QuadraticResult end

"""
Return both lower and upper result values as a tuple
"""
struct Both <: QuadraticResult end

struct QuadraticError <: Exception
   var::String
end

Base.showerror(io::IO, e::QuadraticError) = print(io, e.var, "!")

"""
    function quad(res::QuadraticResult, a, b, c)

Simple quadratic equation solver.

Returns the result specified with `res`
[`Lower`](@ref), [`Upper`](@ref) or [`Both`](@ref)
"""
function quad(res::QuadraticResult, a, b, c)
    x = (b^2 - 4 * a * c)
    x < zero(x) && throw(QuadraticError("Imaginary roots in quadratic"))
    if a == zero(a)
        if b == zero(b)
            c != zero(c) && throw(QuadraticError("Can't solve quadratic"))
            # return zero with correct units
            return out(res, zero(-c / b))
        else
            return out(res, -c / b)
        end
    else
        return side(res, x, a, b)
    end
end
quad(a, b, c) = quad(Both(), a, b, c)

@inline side(::Both, x, a, b) = side(Lower(), x, a, b), side(Upper(), x, a, b)
@inline side(::Upper, x, a, b) = (-b + sqrt(x)) / 2a
@inline side(::Lower, x, a, b) = (-b - sqrt(x)) / 2a

@inline out(::Both, x) = (x, x)
@inline out(::Upper, x) = x
@inline out(::Lower, x) = x

end # module


using LinearRegression
using Test

function check_type_stability(regressor, X)
    @testset "type-inference" begin
        @inferred coef(regressor)
        @inferred regressor(X)
        @inferred regressor(X[1, :])
    end
end

function test_linreg_multivariate(X, y, beta, X_element, y_element)
    intercept = size(X, 2) < size(beta, 1)

    @testset "weights $(repr(weights))" for weights in (nothing, ones(size(X, 1)))
        @testset "Solver $method" for method in (SolveQR(), SolveCholesky())
            regressor = linregress(X, y, weights; intercept=intercept, method=method)
            @test coef(regressor) ≈ beta
            @test regressor(X) ≈ y
            @test regressor(X_element) ≈ y_element
            check_type_stability(regressor, X)
        end
    end
end

@testset "LinearRegression.jl" begin
    @testset "univariate" begin
        X = [4 6]'
        y = [-2, 4]

        @testset "weights $(repr(weights))" for weights in (nothing, ones(2))
            @testset "Solver $method" for method in (SolveQR(), SolveCholesky())
                regressor = linregress(X, y, weights; intercept=true, method=method)
                @test regressor([5]) ≈ 1
                @test regressor([4 5 6]') ≈ [-2, 1, 4]
                @test coef(regressor) ≈ [3, -14]

                regressor = linregress(X, y, weights; intercept=false, method=method)
                @test regressor([0]) == 0
                check_type_stability(regressor, [4 5 6]')
            end
        end
    end

    @testset "multivariate" begin
        @testset "intercept $intercept" for intercept in (true, false)
            X = rand(100, 13)
            X0 = intercept ? [X ones(size(X, 1))] : X

            @testset "single-output" begin
                beta = rand(size(X0, 2))
                y = X0 * beta

                test_linreg_multivariate(X, y, beta, X[17, :], y[17])
            end

            @testset "multi-output" begin
                beta = rand(size(X0, 2), 5)
                y = X0 * beta

                test_linreg_multivariate(X, y, beta, X[17, :], y[17, :])
            end
        end
    end
end


export newton, lagrange

newton(f, x₀::Real; kwargs...) = newton_impl(f, x₀, derivative, abs; kwargs...)
newton(f, x₀::SReal; kwargs...) = newton_impl(f, x₀, jacobian, norm; kwargs...)

function newton_impl(f, x₀,
                     jfun,  # jacobian function
                     nfun;  # norm function
                     rtol::Real=τᵣ,
                     atol::Real=τₐ,
                     maxit=100,  # maximum iterations
                     maxit_ls=20)  # maximum iterations inside the line search
    # Solve f(x) = 0 using the Newton-Armijo method.
    # x₀: initial guess
    # f: function of x

    isconverged = true  # true if solution converged; false otherwise
    α = 1e-4
    ∆max = 1 / rtol  # maximum step size
    perturb_ls = atol^0.75  # ≈ 1e-12, between sqrt(eps) and eps; some perturbation allowed in line search

    # Initialize.
    n = 0
    xₙ = x₀
    fₙ = f(xₙ)
    τf = rtol*nfun(fₙ) + atol
    rx₀ = nfun(x₀)

    # Perform the Newton method.
    while nfun(fₙ) > τf
        λ = 1.
        n_ls = 0  # line search iteration counter

        f′ₙ = jfun(f, xₙ)

        local ∆x  # see https://discourse.julialang.org/t/cumbersome-scoping-rules-for-try-catch-finally/4582
        try
            ∆x = -f′ₙ \ fₙ
        catch err
            isa(err,SingularException) && @error "For xₙ = $xₙ at n = $n, f′ₙ = $f′ₙ is singular."
            break
        end

        # Avoid too large Newton steps.
        nfun(∆x) ≤ ∆max || (isconverged = false; break)
        xₙ₊₁ = xₙ + λ*∆x
        fₙ₊₁ = f(xₙ₊₁)

        # Perform the line search to determine λ.  The stopping criterion does not
        # have perturb_ls ≈ 1e-12 on the RHS, but I guess this kind of perturbation
        # allows update in xₙ even in the situation where line search is supposed
        # to fail.
        while nfun(fₙ₊₁) ≥ (1 - α*λ) * nfun(fₙ) + perturb_ls
            λ /= 2
            xₙ₊₁ = xₙ + λ*∆x
            fₙ₊₁ = f(xₙ₊₁)
            n_ls += 1

            # Too many iteration steps in line search
            n_ls ≤ maxit_ls || (isconverged = false; break)
        end

        # Step accepted; continue the Newton method.
        xₙ = xₙ₊₁
        n += 1

        # Too many iteration steps in Newton's method.
        n ≤ maxit || (isconverged = false; break)
        # n ≤ maxit || throw(ErrorException("Newton method fails to converge in $n iteration steps."))

        fₙ = f(xₙ)
    end
    isnan(nfun(fₙ)) && @error "For xₙ = $xₙ at n = $n, f(xₙ) = $fₙ has NaN."

    return (sol=xₙ, converged=isconverged)  # named tuple
end

# function newton_impl(f, x₀,
#                      jfun,  # jacobian function
#                      nfun;  # norm function
#                      rtol::Real=τᵣ,
#                      atol::Real=τₐ,
#                      maxit=100,  # maximum iterations
#                      maxit_ls=20)  # maximum iterations inside the line search
#     # Solve f(x) = 0 using the Newton-Armijo method.
#     # x₀: initial guess
#     # f: function of x
#     K = Size(x₀)[1]
#
#     isconverged = true  # true if solution converged; false otherwise
#     α = 1e-4
#     ∆max = 1 / rtol  # maximum step size
#     perturb_ls = atol^0.75  # ≈ 1e-12, between sqrt(eps) and eps; some perturbation allowed in line search
#
#     # Initialize.
#     n = 0
#     xₙ = x₀
#     fₙ = f(xₙ)
#     τf = rtol*nfun(fₙ) + atol
#     rx₀ = nfun(x₀)
#
#     # Perform the Newton method.
#     while nfun(fₙ) > τf
#         λ = 1.
#         n_ls = 0  # line search iteration counter
#
#         f′ₙ = jfun(f, xₙ)
#         bool_nz = SVec(ntuple(k->!iszero(f′ₙ[:,k]), Val(K)))
#         M = sum(bool_nz)
#         ind_nz = SVec{M}((1:K)[bool_nz])
#
#         local ∆x  # see https://discourse.julialang.org/t/cumbersome-scoping-rules-for-try-catch-finally/4582
#         try
#             ∆x = -f′ₙ[:,ind_nz] \ fₙ
#         catch err
#             isa(err,SingularException) && @error "For xₙ = $xₙ at n = $n, f′ₙ = $f′ₙ is singular."
#             break
#         end
#
#         # Avoid too large Newton steps.
#         nfun(∆x) ≤ ∆max || (isconverged = false; break)
#         # xₙ₊₁ = xₙ + λ*∆x
#         xₙ₊₁′ = MVec(xₙ)
#         xₙ₊₁′[ind_nz] .+= λ .* ∆x
#         xₙ₊₁ = SVec(xₙ₊₁′)
#         fₙ₊₁ = f(xₙ₊₁)
#
#         # # Perform the line search to determine λ.  The stopping criterion does not
#         # # have perturb_ls ≈ 1e-12 on the RHS, but I guess this kind of perturbation
#         # # allows update in xₙ even in the situation where line search is supposed
#         # # to fail.
#         # while nfun(fₙ₊₁) ≥ (1 - α*λ) * nfun(fₙ) + perturb_ls
#         #     λ /= 2
#         #     xₙ₊₁ = xₙ + λ*∆x
#         #     fₙ₊₁ = f(xₙ₊₁)
#         #     n_ls += 1
#         #
#         #     # Too many iteration steps in line search
#         #     n_ls ≤ maxit_ls || (isconverged = false; break)
#         # end
#
#         # Step accepted; continue the Newton method.
#         xₙ = xₙ₊₁
#         n += 1
#
#         # Too many iteration steps in Newton's method.
#         n ≤ maxit || (isconverged = false; break)
#         # n ≤ maxit || throw(ErrorException("Newton method fails to converge in $n iteration steps."))
#
#         fₙ = f(xₙ)
#     end
#     isnan(nfun(fₙ)) && @error "For xₙ = $xₙ at n = $n, f(xₙ) = $fₙ has NaN."
#
#     return (sol=xₙ, converged=isconverged)  # named tuple
# end

# Need to use different initial guesses for finding the left and right boundaries.
function lagrange(f_obj, f_con, x₀::SReal{K}; kwargs...) where {K}
    ind_x = SVec(ntuple(identity, Val(K)))

    y₀ = SVec(x₀.data..., 1)
    L(y) = f_obj(y[ind_x]) - y[K+1] * f_con(y[ind_x])
    ∇L(y) = gradient(L, y)

    res = newton(∇L, y₀; kwargs...)
    yₛ = res.sol
    converged = res.converged

    xₛ = yₛ[ind_x]
    λₛ = yₛ[K+1]
    Lₛ = L(yₛ)

    return (sol=xₛ, λ=λₛ, L=Lₛ, converged=converged)
end


function loadmatrix(filename::AbstractString)
    transpose(reshape(parse.(Bool, Iterators.flatten(eachline(filename))), length(readline(filename)), :))
end

frombits(bs) = foldl((a, b) -> 2a + b, bs)

function mostcommonbits(A::AbstractMatrix)
    mid = size(A, 1) / 2
    [mid <= s for s in sum(A, dims=1)]
end

function part1(filename::AbstractString)
    bits = mostcommonbits(loadmatrix(filename))
    γ = frombits(bits)
    ϵ = frombits(.!bits)
    γ * ϵ
end

A = loadmatrix("example.txt")
# @time part1("input.txt")

function part2(filename::AbstractString)
    A = loadmatrix(filename)
    O₂ = airfilter(A)
    CO₂ = airfilter(A, !)
    O₂ * CO₂
end

function airfilter(A, fn = identity)
    for i in 1:size(A, 2)
        bits = fn.(mostcommonbits(A))
        A = transpose(hcat((row for row in eachrow(A) if row[i] == bits[i])...))
        if size(A, 1) == 1
            break
        end
    end
    frombits(vec(A))
end


##
# Here, we copy a few definitions from TimeSpans.jl in order to not depend on internals
# https://github.com/beacon-biosignals/TimeSpans.jl/blob/e3c999021336e51a08d118e6defb792e38ac1cc7/src/TimeSpans.jl

const NS_IN_SEC = Dates.value(Nanosecond(Second(1)))  # Number of nanoseconds in one second

# Tweaked from TimeSpans version: https://github.com/beacon-biosignals/AlignedSpans.jl/pull/2#discussion_r829582819
nanoseconds_per_sample(sample_rate) = NS_IN_SEC / sample_rate

function time_from_index(sample_rate, sample_index)
    sample_index > 0 || throw(ArgumentError("`sample_index` must be > 0"))
    return Nanosecond(ceil(Int, (sample_index - 1) * nanoseconds_per_sample(sample_rate)))
end
#
##

# Helper to get the index and the rounding error in units of time
function index_and_error_from_time(sample_rate, sample_time::Period, mode::RoundingMode)
    time_in_nanoseconds = Dates.value(convert(Nanosecond, sample_time))
    time_in_nanoseconds >= 0 ||
        throw(ArgumentError("`sample_time` must be >= 0 nanoseconds"))
    time_in_seconds = time_in_nanoseconds / NS_IN_SEC
    floating_index = time_in_seconds * sample_rate + 1
    index = round(Int, floating_index, mode)
    return index, time_from_index(sample_rate, index) - sample_time
end

"""
    n_samples(sample_rate, duration::Period)

Returns the minimal number of samples that can occur in a span of duration `duration`.
"""
function n_samples(sample_rate, duration::Period)
    duration_in_nanoseconds = Dates.value(convert(Nanosecond, duration))
    duration_in_nanoseconds >= 0 ||
        throw(ArgumentError("`duration` must be >= 0 nanoseconds"))
    duration_in_seconds = duration_in_nanoseconds / NS_IN_SEC
    n_indices = duration_in_seconds * sample_rate
    return floor(Int, n_indices)
end