# EXCLUDE FROM TESTING

using CuArrays, CUDAnative, CUDAdrv, CUDAapi

CUDAnative.initialize()
const dev = device()
const cap = capability(dev)

using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 10
BenchmarkTools.DEFAULT_PARAMETERS.gcsample = true

using SpecialFunctions

## scalar CPU version

@inline cndf2(in::Float32) = 0.5f0 + 0.5f0 * erf(0.707106781f0 * in)

function blackscholes_cpu(sptprice::Float32, strike::Float32, rate::Float32,
                      volatility::Float32, time::Float32)
    logterm = log10(sptprice / strike)
    powterm = .5f0 * volatility * volatility
    den = volatility * sqrt(time)
    d1 = (((rate + powterm) * time) + logterm) / den
    d2 = d1 - den
    NofXd1 = cndf2(d1)
    NofXd2 = cndf2(d2)
    futureValue = strike * exp(-rate * time)
    c1 = futureValue * NofXd2
    call = sptprice * NofXd1 - c1
    return call - futureValue + sptprice
end

## vectorized CPU version

@inline cndf2(in::AbstractArray{Float32}) = 0.5f0 .+ 0.5f0 .* erf.(0.707106781f0 .* in)

function blackscholes_cpu(sptprice::AbstractArray{Float32},
                      strike::AbstractArray{Float32},
                      rate::AbstractArray{Float32},
                      volatility::AbstractArray{Float32},
                      time::AbstractArray{Float32})
    logterm = log10.(sptprice ./ strike)
    powterm = .5f0 .* volatility .* volatility
    den = volatility .* sqrt.(time)
    d1 = (((rate .+ powterm) .* time) .+ logterm) ./ den
    d2 = d1 .- den
    NofXd1 = cndf2(d1)
    NofXd2 = cndf2(d2)
    futureValue = strike .* exp.(- rate .* time)
    c1 = futureValue .* NofXd2
    call = sptprice .* NofXd1 .- c1
    return call .- futureValue .+ sptprice
end

## native CUDA version

@inline cndf2_cuda(in::Float32) = 0.5f0 + 0.5f0 * CUDAnative.erf(0.707106781f0 * in)

function blackscholes_kernel(sptprice::AbstractArray{Float32},
                             strike::AbstractArray{Float32},
                             rate::AbstractArray{Float32},
                             volatility::AbstractArray{Float32},
                             time::AbstractArray{Float32},
                             out::AbstractArray{Float32})
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x

    if i <= size(sptprice, 1)
        logterm = CUDAnative.log10(sptprice[i] / strike[i])
        powterm = 0.5f0 * volatility[i] * volatility[i]
        den = volatility[i] * CUDAnative.sqrt(time[i])
        d1 = (((rate[i] + powterm) * time[i]) + logterm) / den
        d2 = d1 - den
        NofXd1 = cndf2_cuda(d1)
        NofXd2 = cndf2_cuda(d2)
        futureValue = strike[i] * CUDAnative.exp(-rate[i] * time[i])
        c1 = futureValue * NofXd2
        call = sptprice[i] * NofXd1 - c1
        out[i] = call - futureValue + sptprice[i]
    end

    return
end

## scalar CuArrays version

function blackscholes_gpu(sptprice::Float32, strike::Float32, rate::Float32,
                          volatility::Float32, time::Float32)
    logterm = CUDAnative.log10(sptprice / strike)
    powterm = .5f0 * volatility * volatility
    den = volatility * CUDAnative.sqrt(time)
    d1 = (((rate + powterm) * time) + logterm) / den
    d2 = d1 - den
    NofXd1 = cndf2_cuda(d1)
    NofXd2 = cndf2_cuda(d2)
    futureValue = strike * CUDAnative.exp(-rate * time)
    c1 = futureValue * NofXd2
    call = sptprice * NofXd1 - c1
    return call - futureValue + sptprice
end

## vectorized CuArrays version

@inline cndf2_cuarr(in::AbstractArray{Float32}) = 0.5f0 .+ 0.5f0 .* CUDAnative.erf.(0.707106781f0 .* in)

function blackscholes_gpu(sptprice::AbstractArray{Float32},
                          strike::AbstractArray{Float32},
                          rate::AbstractArray{Float32},
                          volatility::AbstractArray{Float32},
                          time::AbstractArray{Float32})
    logterm = CUDAnative.log10.(sptprice ./ strike)
    powterm = .5f0 .* volatility .* volatility
    den = volatility .* CUDAnative.sqrt.(time)
    d1 = (((rate .+ powterm) .* time) .+ logterm) ./ den
    d2 = d1 .- den
    NofXd1 = cndf2_cuarr(d1)
    NofXd2 = cndf2_cuarr(d2)
    futureValue = strike .* CUDAnative.exp.(- rate .* time)
    c1 = futureValue .* NofXd2
    call = sptprice .* NofXd1 .- c1
    return call .- futureValue .+ sptprice
end

## non-native CUDA C version

const cuda_source = "$(tempname()).cu"
const cuda_ptx = "$(tempname()).ptx"

open(cuda_source, "w") do io
    print(io, """
        extern "C" __global__ void blackscholes_kernel(const float *sptprice,
                                                       const float *strike,
                                                       const float *rate,
                                                       const float *volatility,
                                                       const float *time,
                                                       float *out,
                                                       size_t n)
        {
            int i = blockIdx.x * blockDim.x + threadIdx.x;
            if (i < n) {
                float logterm = log10(sptprice[i] / strike[i]);
                float powterm = 0.5 * volatility[i] * volatility[i];
                float den = volatility[i] * sqrt(time[i]);
                float d1 = (((rate[i] + powterm) * time[i]) + logterm) / den;
                float d2 = d1 - den;
                float NofXd1 = 0.5 + 0.5 * erf(0.707106781 * d1);
                float NofXd2 = 0.5 + 0.5 * erf(0.707106781 * d2);
                float futureValue = strike[i] * exp(-rate[i] * time[i]);
                float c1 = futureValue * NofXd2;
                float call = sptprice[i] * NofXd1 - c1;
                out[i] = call - futureValue + sptprice[i];
            }
        }
    """)
end

toolkit = CUDAapi.find_toolkit()
nvcc = CUDAapi.find_cuda_binary("nvcc", toolkit)
toolchain = CUDAapi.find_toolchain(toolkit)
flags = `-ccbin=$(toolchain.host_compiler) -arch=sm_$(cap.major)$(cap.minor)`
run(`$nvcc $flags -ptx -o $cuda_ptx $cuda_source`)

const cuda_module = CuModuleFile(cuda_ptx)
const cuda_function = CuFunction(cuda_module, "blackscholes_kernel")

## main

function checksum(reference, result)
    reference_sum = sum(reference)
    result_sum = sum(result)
    diff = abs(1-reference_sum/result_sum)
    if diff>0.01
        warn("checksum failed: $result_sum instead of $reference_sum (relative difference: $diff)")
        println(stacktrace())
    end
end

function main(iterations)
    sptprice   = Float32[ 42.0 for i = 1:iterations ]
    strike     = Float32[ 40.0 + (i / iterations) for i = 1:iterations ]
    rate       = Float32[ 0.5 for i = 1:iterations ]
    volatility = Float32[ 0.2 for i = 1:iterations ]
    time       = Float32[ 0.5 for i = 1:iterations ]

    timings = Dict()

    reference = blackscholes_cpu.(sptprice, strike, rate, volatility, time)

    let benchmark = @benchmarkable begin
                out = blackscholes_cpu.($sptprice, $strike, $rate,
                                        $volatility, $time)
            end setup=(
                out = zeros($sptprice)
            ) teardown=(
                checksum($reference, out)
            )
        timings["Single-threaded (scalar)"] = run(benchmark)
    end

    let benchmark = @benchmarkable begin
                out = blackscholes_cpu($sptprice, $strike, $rate,
                                       $volatility, $time)
            end setup=(
                out = zeros($sptprice)
            ) teardown=(
                checksum($reference, out)
            )
        timings["Single-threaded (vectorized)"] = run(benchmark)
    end

    let benchmark = @benchmarkable begin
                cudacall(cuda_function,
                         Tuple{Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},
                               Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Csize_t},
                         sptprice_dev, strike_dev, rate_dev, volatility_dev,
                         time_dev, out, n; blocks=grid, threads=block)
                synchronize()
            end setup=(
                sptprice_dev = CUDAdrv.CuArray($sptprice);
                strike_dev = CUDAdrv.CuArray($strike);
                rate_dev = CUDAdrv.CuArray($rate);
                volatility_dev = CUDAdrv.CuArray($volatility);
                time_dev = CUDAdrv.CuArray($time);
                out = CUDAdrv.CuArray{Float32}(size($strike));

                n = size($sptprice, 1);
                block = min(n, 1024);
                grid = ceil(Integer, n/block)
            ) teardown=(
                checksum($reference, Array(out))
            )
        timings["CUDA C (kernel)"] = run(benchmark)
    end

    let benchmark = @benchmarkable begin
                @cuda blocks=grid threads=block blackscholes_kernel(sptprice_dev, strike_dev, rate_dev,
                                                                    volatility_dev, time_dev, out)
                synchronize()
            end setup=(
                sptprice_dev = CUDAdrv.CuArray($sptprice);
                strike_dev = CUDAdrv.CuArray($strike);
                rate_dev = CUDAdrv.CuArray($rate);
                volatility_dev = CUDAdrv.CuArray($volatility);
                time_dev = CUDAdrv.CuArray($time);
                out = CUDAdrv.CuArray{Float32}(size($strike));

                n = size($sptprice, 1);
                block = min(n, 1024);
                grid = ceil(Integer, n/block)
            ) teardown=(
                checksum($reference, Array(out))
            )
        timings["CUDAnative.jl (kernel)"] = run(benchmark)
    end

    let benchmark = @benchmarkable begin
                out = blackscholes_gpu.(sptprice_dev, strike_dev, rate_dev,
                                        volatility_dev, time_dev)
                synchronize()
            end setup=(
                sptprice_dev = CuArrays.CuArray($sptprice);
                strike_dev = CuArrays.CuArray($strike);
                rate_dev = CuArrays.CuArray($rate);
                volatility_dev = CuArrays.CuArray($volatility);
                time_dev = CuArrays.CuArray($time);
                out = zeros($sptprice)
            ) teardown=(
                checksum($reference, Array(out))
            )
        timings["CuArrays.jl (scalar)"] = run(benchmark)
    end

    let benchmark = @benchmarkable begin
                out = blackscholes_gpu(sptprice_dev, strike_dev, rate_dev,
                                       volatility_dev, time_dev)
                synchronize()
            end setup=(
                sptprice_dev = CuArrays.CuArray($sptprice);
                strike_dev = CuArrays.CuArray($strike);
                rate_dev = CuArrays.CuArray($rate);
                volatility_dev = CuArrays.CuArray($volatility);
                time_dev = CuArrays.CuArray($time);
                out = zeros($sptprice)
            ) teardown=(
                checksum($reference, Array(out))
            )
        timings["CuArrays.jl (vectorized)"] = run(benchmark)
    end

    return timings
end

function driver()
    iterations = 10^7
    timings = main(iterations)

    println()
    println("Timings:")
    for (test, trials) in timings
        println("* $test: ", BenchmarkTools.prettytime(time(trials)))
    end

    println()
    println("Rates:")
    for (test, trials) in timings
        println("* $test: ", 1e9*iterations/time(trials), " ops/sec")
    end
end

driver()

rm(cuda_source)
rm(cuda_ptx)


struct utility_function{T} end

function utility_function{T}(n::Integer) where T
    u=zeros(T,n)
    @. u=max(0,log(n/2+1)-log(1:n))
    su=sum(u)
    @. u=u./su .- 1/n
    u
end

function randchisq(rn::AbstractRNG,::Type{T},ν::Integer) where T
    ν < 0 && throw(ArgumentError("ν (DOF) must be atleast 0."))
    s=zero(T)
    for i in 1:ν
        s+=abs2(randn(rn,T))
    end
    s
end

function randchisq(::Type{T},ν::Integer) where T
    randchisq(Random.default_rng(),T,ν)
end

function geo_mean(x::AbstractVector{T}) where T
    gm=one(T)
    ds=0
    for e in x
        e<=0 && continue
        ds+=1
        gm*=e
    end
    ds==0 && return zero(T)
    gm^(1/ds)
end

Base.@kwdef mutable struct Termination{T}
    atol::T=T(1e-8)
    iterations::Int=1e4
end

function termination_criteria(C::Termination,atol,iterations)
    atol <= C.atol && return true
    iterations >= C.iterations && return true
    return false
end


@testset "Embeddings" begin
  Qx, x = QQ["x"]
  K, a = NumberField(x^3 - 2, "a")
  K2, (a2, ) = NumberField([x^3 - 2], "a")
  QQx, x = Hecke.rationals_as_number_field()[1]["x"]
  K3, a3 = NumberField(x^3 - 2, "a")
  K4, (a4, ) = NumberField([x^3 - 2], "a")
  # Let's go a bit crazy with the type of the base field
  for (K, a) in [(K, a), (K2, a2), (K3, a3), (K4, a4)]
    Kt, t = K["t"]
    g = t^3 + a*t + 5
    L, b = NumberField(g, "b")
    r, s = @inferred signature(L)
    @test (r, s) == (1, 4)

    emb = @inferred complex_embeddings(L)
    @test emb === complex_embeddings(L)
    @test (@inferred number_field(emb[1])) === L

    # isreal
    @test length(real_embeddings(L)) == 1
    @test @inferred isreal(emb[1])
    @test count(isreal, emb) == 1
    @test count(!isreal, emb) == 8
    @test count(isimaginary, emb) == 8

    for e in emb
      if isreal(e)
        @test (@inferred conj(e)) == e
      else
        @test conj(e) in emb && conj(e) != e
      end
    end

    # print
    sprint(show, emb[1])
    sprint(show, MIME"text/plain"(), emb[1])

    # restrict
    @test number_field(@inferred restrict(emb[1], K)) === K
    Set([restrict(e, K) for e in emb]) == Set(complex_embeddings(K))
    Set([restrict(e, L) for e in emb]) == Set(emb)
    @test_throws ArgumentError restrict(emb[1], number_field(x^2 - 2, "a", cached = false)[1])

    # compare with complex_embeddings of absolute field
    Labs, LabstoL = @inferred absolute_simple_field(L)
    embabs = complex_embeddings(Labs)
    ainLabs = LabstoL\(L(a))
    binLabs = LabstoL\b

    for e in emb
      imageofb = @inferred e(b)
      @test overlaps(conj(e)(b), conj(e(b)))
      imageofa = e(L(a))
      @test count(ee -> overlaps(ee(binLabs), imageofb) && overlaps(ee(ainLabs), imageofa), embabs) == 1
      @test Hecke.radiuslttwopower(e(b, 128), -128)
    end

    # restrict along morphism
    LtoLabs = inv(LabstoL)
    for e in embabs
      ee = @inferred restrict(e, LtoLabs)
      @test overlaps(ee(LabstoL(gen(Labs))), e(gen(Labs)))
      @test ee == LtoLabs * e
      @test_throws ArgumentError restrict(e, LabstoL)
    end

    for e in embabs
      ee = @inferred extend(e, LabstoL)
      @test length(ee) == 1
      @test LabstoL * ee[1] == e
    end

    for e in emb
      ee = @inferred extend(e, LtoLabs)
      @test length(ee) == 1
      @test LtoLabs * ee[1] == e
    end

    z = (a + b)^100
    for p in [32, 128, 256, 1024]
      for e in emb
        fn = evaluation_function(e, p)
        @test Hecke.radiuslttwopower(fn(z), -p)
        c = e(z, p)
        @test Hecke.radiuslttwopower(c, -p)
      end
    end
  end

  # More complicated restriction

  Qx, x = QQ["x"]
  K, a = number_field(x^2 - 2, "a")
  L, b = number_field(polynomial(K, [-3, 0, 0, 1], cached = false))
  M, c = number_field(polynomial(L, [-5, 0, 1], cached = false))
  emb = complex_embeddings(M)
  @test Set([restrict(e, L) for e in emb]) == Set(complex_embeddings(L))
  @test Set([restrict(e, K) for e in emb]) == Set(complex_embeddings(K))

end


#=
  From
  "Probabilistic logic programming and its applications"
  Luc De Raedt, Leuven
  https://www.youtube.com/watch?v=3lnVBqxjC88
  @ 3:44

  """
  Mike has a bag of marbles with 4 white, 8 blue, and
  6 red marbles. He pulls out one marble from the bag
  and it is red. What is the probability that the
  second marble he pulls out of the bag is white?

  The answer is 0.234931.
  """

   draw1: The probability of drawing white is the probability for 1.0
   Distributions of variable draw[1]
   red        =>  100000  (1.000000)

   Distributions of variable draw[2]
   blue       =>   47339  (0.473390)
   red        =>   29353  (0.293530)
   white      =>   23308  (0.233080)

  Cf ~/cplint/bag_of_marbles.pl
     ~/blog/bag_of_marbles.blog
     ~/psi/bag_of_marbles.psi
     ~/webppl/bag_of_marbles.wppl

  Cf bag_of_marbles.jl. This current model use a for loop.

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

#=
draw1: The probability of drawing white is the probability for 1.0

Distributions of variable draw[2] (num:0)
2.00000 =>  188281  (0.470703)
3.00000 =>  117185  (0.292963)
1.00000 =>   94534  (0.236335)
=#

@model function bag_of_marbles()

    white = 1
    blue  = 2
    red   = 3

    start = [4,8,6]

    draw = tzeros(2)
    for t in 1:2
        # We have to use sum(Int8[...]) since sum([]) throws an error
        draw[t] ~ Categorical(
                     simplex([
                            start[white] - sum([draw[i] == white ? 1 : 0 for i in 1:t-1]),
                            start[blue]  - sum([draw[i] == blue  ? 1 : 0 for i in 1:t-1]),
                            start[red]   - sum([draw[i] == red   ? 1 : 0 for i in 1:t-1])
                            ])
                            )
    end

    true ~ Dirac(draw[1] == red)

end

model = bag_of_marbles()
num_chns = 4

# chns = sample(model, Prior(), 100_000)
chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(),10_000)

display(chns)

println("\ndraw1: The probability of drawing white is the probability for 1.0")
show_var_dist_pct(chns,Symbol("draw[1]"),["white","blue","red"])
show_var_dist_pct(chns,Symbol("draw[2]"),["white","blue","red"])


#df = readtime(Pkg.dir("PortfolioModels/test/data/monthlies.csv"))
df = readtime(Pkg.dir("PortfolioModels/test/data/monthlies_parsable.csv"))

g  = globalmin(df)

@test_approx_eq 0.456204551678408   g[1,1] #  0.4562 from Zivot's derivation
@test_approx_eq 0.2232518040053831  g[1,2] #  0.2233
@test_approx_eq 0.223251804005383   g[1,2] #  0.2233
@test_approx_eq 0.05122725260835459 g[1,3] #  0.0512
@test_approx_eq 0.2693163917078543  g[1,4] #  0.2693

# > muhat.vals
#    Boeing Nordstrom Starbucks Microsoft
#  0.011300  0.001545  0.028461  0.042712


include("transit_poly.jl")

# Randomizer seed
using Random
Random.seed!(42)

nd = 2+ceil(Int64,rand()*20)
r=rand(); b=rand()*(1+r); d = rand(nd)
#r=1.2*rand(); b=0.0; d = rand(nd)
#r = rand(); b = rand(); d = [1.0,1.0,10.0]
#r=0.7821623254303613; b= 0.5; d = [0.0,1.0]

flux = transit_poly_d(r,b,d)
@time flux = transit_poly_d(r,b,d)

# Now, integrate by hand:
function transit_poly_int_d(r,b,d)
s_1 = maximum([0.0,b-r]); s_2=minimum([1.0,b+r])
ns = 500; ds=(s_2-s_1)/ns
s = linspace(s_1+.5*ds,s_2-.5*ds,ns)
fobs = zero(r)
for j=1:ns
  if s[j] < r-b
    dphi = 2pi
  else
    dphi = 2*acos((s[j]^2+b^2-r^2)/(2*b*s[j]))
  end
  z = sqrt(1-s[j]^2)
  imu = d[1] + d[2]*z
  for n=2:length(d)-1
    imu += d[n+1]*(z^n*(n+2)-n*z^(n-2))
  end
  fobs += s[j]*ds*dphi*imu
end
fobs /= pi*(d[1]+2/3*d[2])
fobs = 1-fobs
#println("r: ",r," b: ",b,"f_an: ",flux," flux: ",flux0," f_num: ",fobs," f_exp: ",pi*(d[1]+2/3*d[2]))
return fobs
end
f_num = transit_poly_int_d(r,b,d)
@time f_num = transit_poly_int_d(r,b,d)
println("r: ",r," b: ",b," f_an: ",flux," f_num: ",f_num)

#nphi = 500; ns = 500; dphi = 2pi/nphi; ds=1/ns
#phi = linspace(.5*dphi,2pi-.5*dphi,nphi); s = linspace(.5*ds,1-.5*ds,ns)
#flux0 = zero(r); fobs = zero(r)
#for i=1:nphi, j=1:ns
#  x=s[j]*cos(phi[i]); y = s[j]*sin(phi[i])
#  z = sqrt(1-s[j]^2)
#  imu = d[1]+d[2]*z
#  for k=2:length(d)-1
#    imu += d[k+1]*(z^k*(k+2)-k*z^(k-2))
#  end
#  flux0 += s[j]*ds*dphi*imu
#  if (y-b)^2+x^2 > r^2
#    fobs += s[j]*ds*dphi*imu
#  end
#end
##println("r: ",r," b: ",b,"f_an: ",flux," flux: ",flux0," f_num: ",fobs/flux0," f_exp: ",pi*(d[1]+2/3*d[2]))
#println("r: ",r," b: ",b," f_an: ",flux," f_num: ",fobs/flux0)


raw"""
# Conjugate Gradient Method

## References

- [wiki]: https://en.wikipedia.org/wiki/Conjugate_gradient_method
- [painless]: Shewchuk, Jonathan Richard. "An introduction to the conjugate gradient method without
  the agonizing pain." (1994): 1.

The derivation in <https://en.wikipedia.org/wiki/Derivation_of_the_conjugate_gradient_method> is
rather complicated and unintuitive, so it's better to go through [painless].
"""
module ConjugateGradientMethod

using LinearAlgebra: norm, I, (⋅)

export conjugate_gradient, conjugate_directions

#=

At first "conjugate" sounds confusing, but it becomes clear after realizing it is actually "$p_i$
and $p_j$ (i ≠ j) are conjugate **w.r.t. $A$**" (while $r_i^T r_j = 0$ for i ≠ j). So the whole
iteration is indeed doing Gram-Schmidt process to find $A$-orthogonal search directions.

Recall some definitions:

    x_* = \sum _{i=1}^{n} α_i p_i  ...mutually conjugate vectors, w.r.t. $A$
    α_i = <p_i, r_i> / <p_i, p_i>_A

    r_i = b - A x_i  ...the residual, negative gradient of $f(x) = 1/2 x^T A x - b^T x$ at $x_i$
    e_i = x_i - x_*  ...the error

Some notes taken from [painless]:

- "In fact, if the search vectors are constructed by conjugation of the axial unit vectors,
  Conjugate Directions becomes equivalent to performing Gaussian elimination."
- In Conjugate Directions, $e_{i + 1}$ is $A$-orthogonal to $p_{i}$. This is equivalent to finding
  the minimum point along $p_i$.
- Conjugate Directions chooses value from $e_i = e_0 + P_i$ that minimizes $<e_0, e_0>_A$, where
  $P_i$ is `span {p_0, p_1, ..., p_{i - 1}}`. (See [painless] 7.3)
- Moreover, not only $<p_{i - 1}, e_i>_A = 0$, but for $i < j$, $<p_i, e_j>_A = 0$ also holds. As
  $p_i^T A e_j = 0$ and $r_j = -A e_j$, we know $p_i$ and $r_j$ are orthogonal.
- Finally, CG is simply a special case of Conjugate Directions where search directions are
  constructed by conjugation of the residuals. This choice makes $r_i^T r_j = 0$ and $p_i^T r_i =
  r_i^T r_i$ and $r_i^T A d_i = d_i^T A d_i$ and $r_{i + 1}$ being $A$-orthogonal to $P_i$, etc.,
  many nice properties, so there is no need to store all the previous search directions for doing
  Gram-Schmidt process.

Recapitulate some derivations of equations:

    α_i = <p_i, r_i> / <p_i, p_i>_A    (CD, given $p_i$, minimization)
        = <r_i, r_i> / <p_i, p_i>_A    (CG, u_i = r_i, c.f. [painless] Eq. 42)

    β_ij = -<u_i, p_j>_A / <p_j, p_j>_A    (CD, Gram-Schmidt, c.f. [painless] Eq. 37)

    β_i := β_{i}{i - 1}
        = -<r_{i + 1}, p_i>_A / <p_i, p_i>_A    (CG, Gram-Schmidt in Krylov subspace)
        = <r_{i + 1}, r_{i + 1}> / <r_i, r_i>    (CD, c.f. [painless] Eq. 43, definition of α_i)

=#

# Just a proof-of-concept. VERY slow.
#
# <https://en.wikipedia.org/wiki/Derivation_of_the_conjugate_gradient_method#The_conjugate_direction_method>
function conjugate_directions(
    A::AbstractMatrix{T},
    b::Vector{T},
)::Vector{T} where T<:AbstractFloat
    n = length(b)
    x = zeros(T, n)
    if norm(b) < eps(T)
        return x
    end

    U = Matrix(T.(I(n)))
    P = zeros(T, n, n)
    u(k::Int) = @view U[:, k]
    p(k::Int) = @view P[:, k]

    r = b
    for i in 1:n
        # Gram-Schmidt orthogonalization, directly translated from its definition.
        t = u(i)
        for j = 1:(i - 1)
            Ap = A * p(j)
            β = -(u(i)' * Ap) / (p(j)' * Ap)
            t += β * p(j)
        end
        p(i) .= t
        Ap = A * p(i)
        α = (p(i)' * r) / (p(i)' * Ap)
        x += α * p(i)
        r -= α * Ap
    end
    x
end

function conjugate_gradient(
    A::AbstractMatrix{T},
    b::Vector{T},
)::Vector{T} where T<:AbstractFloat
    n = length(b)
    x = zeros(T, n)
    if norm(b) < eps(T)
        return x
    end

    r = b - A * x
    p = r
    rr = r ⋅ r

    MAX_ITERS = n * 10
    for i in 1:MAX_ITERS
        Ap = A * p
        α = rr / (p ⋅ Ap)
        x += α * p
        r -= α * Ap
        rr_new = r ⋅ r
        if sqrt(rr_new) < eps(T)
            break
        end
        β = rr_new / rr
        p = r + β * p
        rr = rr_new
    end
    x
end

# TODO: Preconditioning.

end  # module


"""
    munkres(costMat) -> Zs

Find an optimal solution of the assignment problem represented by the square
matrix `costMat`. Return an sparse matrix illustrating the optimal matching.

# Examples

```julia
julia> costMat = ones(3, 3) - eye(3, 3)
3×3 Array{Float64,2}:
 0.0  1.0  1.0
 1.0  0.0  1.0
 1.0  1.0  0.0

julia> matching = Hungarian.munkres(costMat)
3×3 sparse matrix with 3 Int64 nonzero entries:
	[1, 1]  =  2
	[2, 2]  =  2
	[3, 3]  =  2

julia> full(matching)
3×3 Array{Int64,2}:
 2  0  0
 0  2  0
 0  0  2
```
"""
function munkres(costMat::AbstractMatrix{T}) where {T<:Real}
    rowNum, colNum = size(costMat)
    colNum ≥ rowNum ||
    throw(ArgumentError("Non-square matrix should have more columns than rows."))

    # preliminaries:
    # "no lines are covered;"
    rowCovered = falses(rowNum)
    colCovered = falses(colNum)

    # for tracking indices
    rowCoveredIdx = Int[]
    colCoveredIdx = Int[]
    rowUncoveredIdx = Int[]
    colUncoveredIdx = Int[]
    sizehint!(rowCoveredIdx, rowNum)
    sizehint!(colCoveredIdx, colNum)
    sizehint!(rowUncoveredIdx, rowNum)
    sizehint!(colUncoveredIdx, colNum)

    # "no zeros are starred or primed."
    # use a sparse matrix Zs to store these three kinds of zeros:
    # 0 => NON   => Non-zero
    # 1 => Z     => ordinary zero
    # 2 => STAR  => starred zero
    # 3 => PRIME => primed zero
    Zs = spzeros(Int8, rowNum, colNum)

    # for tracking changes per row/col of A
    Δrow = zeros(T, rowNum)
    Δcol = zeros(T, colNum)

    # "consider a row of the matrix A;
    #  subtract from each element in this row the smallest element of this row.
    #  do the same for each row."
    @inbounds for i = 1:rowNum
        m = costMat[i, 1]
        @simd for j = 2:colNum
            m = min(m, costMat[i, j])
        end
        Δrow[i] = -m
    end

    # "then consider each column of the resulting matrix and subtract from each
    #  column its smallest entry."
    # Note that, this step should be omitted if the input matrix is not square.
    if rowNum == colNum
        @inbounds for j = 1:colNum
            m = costMat[1, j] + Δrow[1]
            @simd for i = 2:rowNum
                m = min(m, costMat[i, j] + Δrow[i])
            end
            Δcol[j] = -m
        end
    end

    # for tracking those starred zero
    rowSTAR = falses(rowNum)
    colSTAR = falses(colNum)
    # since looping through a row in a SparseMatrixCSC is costy(not cache-friendly),
    # a row to column index mapping of starred zeros is also tracked here.
    row2colSTAR = Dict{Int,Int}()
    for ii in CartesianIndices(size(costMat))
        # "consider a zero Z of the matrix;"
        r, c = ii.I
        cost = costMat[r, c] + Δrow[r] + Δcol[c]
        if cost == 0
            Zs[r, c] = Z
            # "if there is no starred zero in its row and none in its column, star Z.
            #  repeat, considering each zero in the matrix in turn;"
            if !colSTAR[c] && !rowSTAR[r]
                Zs[r, c] = STAR
                rowSTAR[r] = true
                colSTAR[c] = true
                row2colSTAR[r] = c
                # "then cover every column containing a starred zero."
                colCovered[c] = true
            end
        end
    end

    # preliminaries done, go to step 1
    stepNum = 1

    # if the assignment is already found, exit
    stepNum = exit_criteria(colCovered, size(Zs))

    # pre-allocation
    sequence = Tuple{Int,Int}[]       # used in step 2
    minLocations = Tuple{Int,Int}[]   # used in step 3

    # these three steps are parallel with those in the paper:
    # J. Munkres, "Algorithms for the Assignment and Transportation Problems",
    # Journal of the Society for Industrial and Applied Mathematics, 5(1):32–38, 1957 March.
    while stepNum != 0
        if stepNum == 1
            stepNum = step1!(Zs, rowCovered, colCovered, rowSTAR, row2colSTAR)
        elseif stepNum == 2
            empty!(sequence)
            stepNum = step2!(Zs, sequence, rowCovered, colCovered, rowSTAR, row2colSTAR)
        elseif stepNum == 3
            empty!(rowCoveredIdx)
            empty!(colCoveredIdx)
            empty!(rowUncoveredIdx)
            empty!(colUncoveredIdx)
            empty!(minLocations)
            stepNum = step3!(costMat, Zs, minLocations, rowCovered, colCovered, rowSTAR, row2colSTAR, Δrow, Δcol, rowCoveredIdx, colCoveredIdx, rowUncoveredIdx, colUncoveredIdx)
        end
    end

    return Zs
end

function munkres(costMat::AbstractMatrix{S}) where {T<:Real,S<:Union{Missing,T}}
    # replace forbidden edges (i.e. those with a missing cost) by a very large cost, so that they
    # are never chosen for the matching (except if they are the only possible edges)
    costMatReal = [ismissing(x) ? typemax(T) : x for x in costMat]

    # compute the assignment with the standard procedure. the return value is guaranteed to be sparse
    assignment = munkres(costMatReal)

    # remove the forbidden edges, would they be in the assignment
    colLen = size(assignment, 2)
    rows = rowvals(assignment)
    for c = 1:colLen, i in nzrange(assignment, c)
        r = rows[i]
        if assignment[r, c] == STAR && ismissing(costMat[r, c])
            assignment[r, c] = Z # standard zero, no assignment made
        end
    end

    return assignment
end

"""
    exit_criteria(colCovered, ZsDims) -> stepNum
We adjust Munkres's algorithm in order to deal with rectangular matrices,
so only K columns are counted here, where K = min(size(Zs))
"""
function exit_criteria(colCovered, ZsDims)
    count = 0
    @inbounds for i in eachindex(colCovered)
        colCovered[i] && (count += 1)
    end
    if count == ZsDims[1]
        # algorithm exits
        return 0
    else
        # "otherwise, return to step 1."
        return 1
    end
end

"""
Step 1 of the original Munkres' Assignment Algorithm
"""
function step1!(Zs, rowCovered, colCovered, rowSTAR, row2colSTAR)
    colLen = size(Zs, 2)
    rows = rowvals(Zs)
    # step 1:
    zeroCoveredNum = 0
    # "repeat until all zeros are covered."
    while zeroCoveredNum < nnz(Zs)
        zeroCoveredNum = 0
        for c = 1:colLen, i in nzrange(Zs, c)
            r = rows[i]
            # "choose a non-covered zero and prime it"
            if colCovered[c] == false && rowCovered[r] == false
                Zs[r, c] = PRIME
                # "consider the row containing it."
                # "if there is a starred zero Z in this row"
                if rowSTAR[r]
                    # "cover this row and uncover the column of Z"
                    rowCovered[r] = true
                    colCovered[row2colSTAR[r]] = false
                else
                    # "if there is no starred zero in this row,
                    #  go at once to step 2."
                    return 2
                end
            else
                # otherwise, this zero is covered
                zeroCoveredNum += 1
            end
        end
    end
    # "go to step 3."
    return 3
end

"""
Step 2 of the original Munkres' Assignment Algorithm
"""
function step2!(Zs, sequence, rowCovered, colCovered, rowSTAR, row2colSTAR)
    ZsDims = size(Zs)
    rows = rowvals(Zs)
    # step 2:
    flag = false
    # "there is a sequence of alternating primed and starred zeros, constructed
    #  as follows:"
    # "let Z₀ denote the uncovered 0′.[there is only one.]"
    for c = 1:ZsDims[2], i in nzrange(Zs, c)
        r = rows[i]
        # note that Z₀ is an **uncovered** 0′
        if Zs[r, c] == PRIME && colCovered[c] == false && rowCovered[r] == false
            # push Z₀(uncovered 0′) into the sequence
            push!(sequence, (r, c))
            # "let Z₁ denote the 0* in the Z₀'s column.(if any)"
            # find 0* in Z₀'s column
            for j in nzrange(Zs, c)
                Z₁r = rows[j]
                if Zs[Z₁r, c] == STAR
                    # push Z₁(0*) into the sequence
                    push!(sequence, (Z₁r, c))
                    # set sequence continue flag => true
                    flag = true
                    break
                end
            end
            break
        end
    end
    # "let Z₂ denote the 0′ in Z₁'s row(there will always be one)."
    # "let Z₃ denote the 0* in the Z₂'s column."
    # "continue until the sequence stops at a 0′ Z₂ⱼ, which has no 0* in its column."
    while flag
        flag = false
        r = sequence[end][1]
        # find Z₂ in Z₃'s row (always exits)
        for c = 1:ZsDims[2]
            if Zs[r, c] == PRIME
                # push Z₂ into the sequence
                push!(sequence, (r, c))
                # find 0* in Z₂'s column
                for j in nzrange(Zs, c)
                    Z₃r = rows[j]
                    if Zs[Z₃r, c] == STAR
                        # push Z₃(0*) into the sequence
                        push!(sequence, (Z₃r, c))
                        # set sequence continue flag => true
                        flag = true
                        break
                    end
                end
                break
            end
        end
    end

    # "unstar each starred zero of the sequence;"
    for i = 2:2:length(sequence)-1
        Zs[sequence[i]...] = Z
    end

    # clean up
    fill!(rowSTAR, false)
    empty!(row2colSTAR)

    # "and star each primed zero of the sequence."
    for i = 1:2:length(sequence)
        Zs[sequence[i]...] = STAR
    end

    for c = 1:ZsDims[2], i in nzrange(Zs, c)
        r = rows[i]
        # "erase all primes;"
        if Zs[r, c] == PRIME
            Zs[r, c] = Z
        end
        # "and cover every column containing a 0*."
        if Zs[r, c] == STAR
            colCovered[c] = true
            rowSTAR[r] = true
            row2colSTAR[r] = c
        end
    end

    # "uncover every row"
    fill!(rowCovered, false)

    # "if all columns are covered, the starred zeros form the desired independent set."
    return exit_criteria(colCovered, ZsDims)
end

"""
Step 3 of the original Munkres' Assignment Algorithm
"""
function step3!(costMat::AbstractMatrix{T}, Zs, minLocations, rowCovered, colCovered, rowSTAR, row2colSTAR,
                Δrow, Δcol, rowCoveredIdx, colCoveredIdx, rowUncoveredIdx, colUncoveredIdx) where T <: Real
    # step 3(Step C):
    # "let h denote the smallest uncovered element of the matrix;
    #  add h to each covered row; then subtract h from each uncovered column."
    #             -h                    ||
    #              |     no change      ||  reduce old zeros
    #              | change:(+h)+(-h)=0 ||     change: +h
    #       +h ---------Covered Row-----||----Covered Row--------
    #              |  Uncovered Column  ||   Covered Column
    #==============|====================||================================#
    #              |   Uncovered Row    ||    Uncovered Row
    #              |  Uncovered Column  ||   Covered Column
    #              |    change: -h      ||     change: 0
    #              | produce new zeros  ||     no change
    #              |                    ||
    # since we always apply add/substract h to a whole row/column, it's unnecessary
    # to apply every operation to every entry of A's row/column, we only need a row
    # and a column vector to keep tracking those changes and use it when necessary.

    # find h and track the location of those new zeros
    @inbounds for i in eachindex(rowCovered)
        rowCovered[i] ? push!(rowCoveredIdx, i) : push!(rowUncoveredIdx, i)
    end

    @inbounds for i in eachindex(colCovered)
        colCovered[i] ? push!(colCoveredIdx, i) : push!(colUncoveredIdx, i)
    end

    h = typemax(T)
    @inbounds for j in colUncoveredIdx, i in rowUncoveredIdx
        cost = costMat[i, j] + Δcol[j] + Δrow[i]
        if cost <= h
            if cost != h
                h = cost
                empty!(minLocations)
            end
            push!(minLocations, (i, j))
        end
    end

    for i in rowCoveredIdx
        Δrow[i] += h
    end

    for i in colUncoveredIdx
        Δcol[i] -= h
    end

    # mark new zeros in Zs and remove those elements that will no longer be zero
    # produce new zeros
    for loc in minLocations
        Zs[loc...] = Z
    end
    # reduce old zeros
    rows = rowvals(Zs)
    for c in colCoveredIdx, i in nzrange(Zs, c)
        r = rows[i]
        if rowCovered[r]
            if Zs[r, c] == STAR
                rowSTAR[r] = false
                delete!(row2colSTAR, r)
            end
            Zs[r, c] = NON
        end
    end

    dropzeros!(Zs)

    # "return to step 1."
    return stepNum = 1
end


function jointplot(; x::AbstractVector = Int[], y::AbstractVector = Int[], smoothing::Int = 4)

    @assert sum(Missings.ismissing.(x)) == 0 "x has missing data"
    @assert sum(Missings.ismissing.(y)) == 0 "y has missing data"

    #Calculate kde
    a = kde((collect(skipmissing(x)),collect(skipmissing(y))))

    #Parse data from grid to 1D array
    x_mod = []
    y_mod = []
    d_mod = []

    for _x in enumerate(a.x)
        for _y in enumerate(a.y)
            push!(x_mod, _x[2])
            push!(y_mod, _y[2])
            push!(d_mod, a.density[_x[1], _y[1]])
        end
    end

    min_d = minimum(d_mod)
    max_d = maximum(d_mod)
    intervals = (max_d - min_d)/10
    center =  (max_d - min_d)/2

    #Plot
    a = heatmap(x = x_mod, y = y_mod, color = d_mod)
    a.width = a.height = 500
    a.scales[1]._type = a.scales[2]._type = "linear"
    a.marks[1].properties.enter.width = a.marks[1].properties.enter.height = VegaValueRef(value = smoothing)

    a.scales[2].zero = false

    a.scales[1].nice = a.scales[2].nice = true
    a.scales[3].domain =  collect(min_d:intervals:max_d)
    a.legends[1].values = [min_d, center, max_d]

    return a

end


using SpecialFunctions, ForwardDiff, FiniteDifferences
include("shared.jl")

const BIG_FD    = central_fdm(10,1)
const BIG_FD_O2 = central_fdm(10,2)

xvbesselkdv(v,x)  = ForwardDiff.derivative(_v->BesselK.adbesselkxv(_v,x), v)
xvbesselkdv2(v,x) = ForwardDiff.derivative(_v->xvbesselkdv(_v,x), v)

beskxv(v,x) = besselk(v,x)*(x^v)
dxvbesselkdv(v, x)  = BIG_FD(_v->beskxv(_v,x), v)
dxvbesselkdv2(v, x) = BIG_FD_O2(_v->beskxv(_v,x), v)

fastdxvbesselkdv2(v, x) = (beskxv(v+2e-6, x) - 2*beskxv(v+1e-6, x) + beskxv(v, x))/1e-12

const VGRID = range(0.25, 10.0, length=101)        # to avoid integer v.
const XGRID = range(0.0,  50.0, length=201)[2:end] # to avoid zero x.

const BASELINE = [dxvbesselkdv2(z[1], z[2])     for z in Iterators.product(VGRID, XGRID)]
const OURSOL   = [xvbesselkdv2(z[1], z[2])      for z in Iterators.product(VGRID, XGRID)]
const FASTFD   = [fastdxvbesselkdv2(z[1], z[2]) for z in Iterators.product(VGRID, XGRID)]
const TOLS     = atolfun.(zip(BASELINE, OURSOL))
const FDTOLS   = atolfun.(zip(BASELINE, FASTFD))
const DIFTOLS  = log10.(TOLS) .- log10.(FDTOLS)

gnuplot_save_matrix!("../plotdata/atols_deriv2_xv_fd.csv",    FDTOLS, VGRID, XGRID)
gnuplot_save_matrix!("../plotdata/atols_deriv2_xv_ad.csv",      TOLS, VGRID, XGRID)
gnuplot_save_matrix!("../plotdata/atols_deriv2_xv_fdad.csv", DIFTOLS, VGRID, XGRID)


@testset "Common" begin
    @test get_code_center_frequency_ratio(GPSL1) ≈ 1/1540
end


function sensorloc_data(seed, n)

    rng = Random.MersenneTwister(seed)

    # n = number of sensors points
    # m = number of anchor points
    m = floor(Int, 0.1 * n)

    # Sensor true position (2 dimensional)
    x_true = Random.rand(rng, Float64, (2, n))
    # Distances from sensors to sensors
    d = Dict((i, j) => norm(x_true[:, i] - x_true[:, j]) for i in 1:n for j in 1:i)

    # Anchor positions
    a = Dict(i => Random.rand(rng, Float64, (2, 1)) for i in 1:m)
    # Distances from anchor to sensors
    d_bar = Dict((k, j) => norm(x_true[:, j] - a[k]) for k in 1:m for j in 1:n)

    return m, x_true, a, d, d_bar
end

function sensorloc_eval(n, m, x_true, XX)
    @show norm(x_true - XX[1:2, 3:n + 2])
    @show rank = length([eig for eig in eigen(XX).values if eig > 1e-7])
    return nothing
end


function OMPTimeDomain(V::Array{Complex128,2},Precision::Int64,OrthoTol::Float64,gmethod,lmethod,ltol,gWallClockTime)
	set_bigfloat_precision(Precision)
	#Add a zero at the beginning of the trace to put make it
	#compatible with a strictly proper transfer function
	#which is what we are using to fit the data
	V = cat(1,zeros(Complex128,(1,size(V,2))),V)
	L = size(V,1)
	N = size(V,1)
	zarb = genZArb(L)
	zdbl = convert(Vector{Complex128},zarb)
	#pole vectors one at arbitrary precsion, one at dbl
	parb = zeros(ComplexPair{BigFloat},N-1)
	Larb = zeros(ComplexPair{BigFloat},N-1)
	Ldbl = zeros(ComplexPair{Float64},N-1)
	pdbl = zeros(Complex128,N-1)
	AICc = zeros(Float64,N-1)
    RSS = zeros(Float64,N-1)
	B = ones(ComplexPair{BigFloat},(L,N))/sqrt(L)
	Bdbl = ones(Complex128,(L,N))/sqrt(L)
	#Set up some fitting parameters for NLOPT
	opt = Opt(gmethod,2)
    lopt = Opt(lmethod,2)
    ftol_rel!(lopt, ltol) #local stopping criteria
    maxtime!(opt,gWallClockTime) #global stopping critera
    betalb = [0.01; 0]
    betaub = [0.99; (2*pi-eps(Float64))]
    lower_bounds!(opt,betalb)
    upper_bounds!(opt,betaub)
	terminate_fit = 0;
	k = 0;
	Resid = copy(V)
	while terminate_fit<1
		k += 1
		if k==1
			anm1 = zeros(Complex128,0)
		else
			anm1 = [pdbl[k-1]]
		end
		Bnm1 = convert(Vector{Complex128},B[:,k])
		max_objective!(opt,(x,g)->ompObj(x,g,zdbl,anm1,Bnm1,Resid))
		beta0 = genStart(betalb,betaub)
		(oVal,betaOpt,exitcode) = optimize(opt,beta0)
		println(oVal)
		parb[k] = convert(ComplexPair{BigFloat},betaOpt[1]*exp(im*betaOpt[2]))
		pdbl[k] = convert(Complex128,parb[k])
		(Bnext,Ldbl[k]) = phiNext(zdbl,[pdbl[1:k]],Bnm1)
		(B[:,k+1],Larb[k]) = phiNext(zarb,[parb[1:k]],B[:,k])
		Bdbl[:,k+1] = convert(Vector{Complex128},B[:,k+1])
		OrthogonalityMetric = maximum(abs(abs(Bnext'*Bdbl[:,1:k+1]) - cat(1,zeros(Float64,k),one(Float64))'))
		println(OrthogonalityMetric)
		if OrthogonalityMetric>OrthoTol
			#terminate the fit if we lose orthogonality when operating at double precision
			terminate_fit = 2;
		end
		Tdbl = ifft(Bdbl[:,1:k+1],1)*sqrt(L)
		Resid = V-Tdbl*Tdbl'*V
		RSS[k] = norm(Resid)
		#We have L-1 actual data points per column x size(V,2) columns, 2 numbers per element
		#And we have 2 parameters per pole, +1 for the variance estimate, plus 2 per linear coefficient
		AICc[k] = akaikeInformationCriterion(RSS[k],2*k+1+2*k*size(V,2),2*size(V,2)*(size(V,1)-1))
		if k==(N-1)
			terminate_fit = 1
		end
	end
    Tdbl = ifft(Bdbl[:,1:k+1],1)*sqrt(L)
    A = Tdbl[:,1:(k+1)]'*V
	return (zarb,parb,Larb,A,AICc,terminate_fit)
end

function OMPTimeDomainIndPoles(V::Array{Complex128,2},Precision::Int64,OrthoTol::Float64,gmethod,lmethod,ltol,gWallClockTime)
	for j=1:size(V,2)
		OMPTimeDomain(V,Precision,OrthoTol,gmethod,lmethod,ltol,gWallClockTime)
	end
end

function akaikeInformationCriterion(RSS,numParams,dataLen)
    #This gives the sample size corrected Akaike Information Criteria
    #commonly known as AICc in the literature
    K = numParams+1
    return dataLen*log(RSS/dataLen)+2*K+2*K*((K+1)/(dataLen-K-1))
end

function genStart(lb,ub)
	bstart = zero(lb)
	for k=1:length(lb)
		bstart[k] = (ub[k]-lb[k])*rand(1)[1]+lb[k]
	end
	return bstart
end

function ompObj(beta::Vector,grad::Vector,z::Vector{Complex128},anm1::Vector{Complex128},Bnm1::Vector{Complex128},D::Array{Complex128,2})
    R = beta[1];
    T = beta[2];
    p = [anm1; R*exp(im*T)]
    (Bn,L) = phiNext(z,p,Bnm1)
    Tn = ifft(Bn,1)*sqrt(length(z))
    Proj = Tn'*D;
    return sum(abs(Proj),2)[1]
end

function createFreqAxis(N::Int,dT)
	#For some time data with index values 0:(N-1) this function creates the corresponding frequency axis
	#in units of 1/time spacing
    fs = 1/dT
    df = fs/N
    f = [0:df:(fs-df)] - (fs-mod(N,2)*df)/2
    return f
end


function nelder_mead(fcn::Function,
                     x0::Array,
                     args...;
                     parameters            = Optim.AdaptiveParameters(),
                     initial_simplex       = Optim.AffineSimplexer(),
                     iterations::Int       = 1000,
                     store_trace::Bool     = false,
                     show_trace::Bool      = false,
                     extended_trace::Bool  = false,
                     kwargs...)

    Optim.optimize(fcn, x0,
                   method = Optim.NelderMead(parameters = parameters, initial_simplex = initial_simplex),
                   iterations = iterations, store_trace = store_trace, show_trace = show_trace,
                   extended_trace = extended_trace)
end

type MatlabSimplexer <: Optim.Simplexer
    a::Float64
    b::Float64
end
MatlabSimplexer(;a = 0.00025, b = 0.05) = MatlabSimplexer(a, b)

function Optim.simplexer{T, N}(A::MatlabSimplexer, initial_x::Array{T, N})
    n = length(initial_x)
    initial_simplex = Array{T, N}[initial_x for i = 1:n+1]
    for j = 1:n
        initial_simplex[j+1][j] += initial_simplex[j+1][j] == zero(T) ? A.b * initial_simplex[j+1][j] : A.a
    end
    initial_simplex
end


module BenchmarkingLinearInterpolators

using Printf
using LinearInterpolators
using LazyAlgebra
using BenchmarkTools

mimeshow(io::IO, x) = show(io, MIME"text/plain"(), x)
mimeshow(x) = mimeshow(stdout, x)

prt(op::AbstractString, t) = begin
    @printf("%s:\n", op)
    mimeshow(t)
end

prt(op::AbstractString, t, nops::Integer) = begin
    # times in nanoseconds, speed in Gigaflops
    @printf("%s [%.6f Gflops]:\n", op, nops/minimum(t.times))
    mimeshow(t)
    println()
end

n1, n2 = 300, 200
m1, m2 =  43,  29
y = randn(n1, n2)
ker = CatmullRomSpline()
A1 = SparseUnidimensionalInterpolator(ker, 1, 1:n1, range(1,n1,length=m1))
A2 = SparseUnidimensionalInterpolator(ker, 2, 1:n2, range(1,n2,length=m2))
H = A1*A2

function LazyAlgebra.vcreate(::Type{LazyAlgebra.Direct},
                             ::Gram,
                             x::AbstractArray{T,N},
                             scratch::Bool) where {T<:AbstractFloat,N}
    return similar(x)
end

r1 = 2*length(A1.C)
r2 = 2*length(A2.C)
x = conjgrad(H'*H, H'*y; verb=true, ftol=1e-8, maxiter=50)
z = vcopy(y)
prt("A1*x",      @benchmark($A1*$x),        r1*m2)
prt("A2*x",      @benchmark($A2*$x),        r2*m1)
prt("A2*A1*x",   @benchmark($(A2*A1)*$x),   r1*m2 + r2*n1)
prt("A1*A2*x",   @benchmark($(A1*A2)*$x),   r2*m1 + r1*n2)
prt("A1'*y",     @benchmark($(A1')*$y),     r1*m2)
prt("A2'*y",     @benchmark($(A2')*$y),     r2*m1)
prt("A2'*A1'*y", @benchmark($(A2'*A1')*$y), r1*m2 + r2*n1)
prt("A1'*A2'*y", @benchmark($(A1'*A2')*$y), r2*m1 + r1*n2)
println("")

n = 100
pos = [1.1*i + 0.2*j for i in 1:n1, j in 1:n2]
pmin, pmax = extrema(pos)
A = SparseInterpolator(ker, pos, range(pmin, pmax, length=n))
x = randn(n)
y = randn(n1,n2)
nops = 2*length(A.C)

prt("A*x",  @benchmark($A*$x), nops)
prt("A'*x", @benchmark($(A')*$y), nops)

end # module


struct CellMap
    yL::Any
    yR::Any
    jacobian::Any
    dim::Any
    function CellMap(yL, yR)
        dim = length(yL)
        @assert length(yR) == dim
        @assert all(yR .>= yL)
        jacobian = 0.5 * (yR - yL)
        new(yL, yR, jacobian, dim)
    end
end

function Base.show(io::IO, cellmap::CellMap)
    dim = cellmap.dim
    yL = cellmap.yL
    yR = cellmap.yR
    str = "CellMap\n\tLeft : $yL\n\tRight: $yR"
    print(io, str)
end

function dimension(C::CellMap)
    return C.dim
end

function jacobian(C::CellMap)
    return C.jacobian
end

function inverse_jacobian(C::CellMap)
    return 1.0 ./ jacobian(C)
end

function determinant_jacobian(C)
    return prod(jacobian(C))
end

function face_determinant_jacobian(C)
    jac = jacobian(C)
    return [jac[1], jac[2], jac[1], jac[2]]
end

function (C::CellMap)(x)
    dim = dimension(C)
    return C.yL .+ (jacobian(C) .* (x .+ ones(dim)))
end

function inverse(c::CellMap,x)
    jac = jacobian(c)
    return ((x .- c.yL) ./ jac) .- 1.0
end


using Plexifications
using Base.Test: @test, @test_throws

@test_throws ErrorException begin
    inv(DualPerplex(0, 0.0))
end

@test begin
    !iszerodivisor(DualPerplex(1,2,3,4))
end

@test begin
    iszerodivisor(DualPerplex(1,1,3,4))
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(DualPerplex(a))
end

@test begin
    a = 1
    !isreal(DualPerplex(a, a))
end

@test begin
    z = random(DualPerplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(DualPerplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, DualPerplex(1, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1, R: 2, A: 3, RA: 4]"
    l == r
end

@test begin
    x = random(DualPerplex{BigFloat})
    y = random(DualPerplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(DualPerplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(DualPerplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    barA(barA(x)) == x
end

@test begin
    x = random(DualPerplex{BigInt})
    barR(barR(x)) == x
end

@test begin
    x = random(DualPerplex{BigInt})
    barR(barA(x)) == barA(barR(x))
end

@test begin
    x = random(DualPerplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = barA(x * y)
    r = barA(x) * barA(y)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = barA(x + y)
    r = barA(x) + barA(y)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = barA(x - y)
    r = barA(x) - barA(y)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = barR(x * y)
    r = barR(x) * barR(y)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = barR(x + y)
    r = barR(x) + barR(y)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = barR(x - y)
    r = barR(x) - barR(y)
    l == r
end

@test begin
    x = random(DualPerplex{Rational{BigInt}})
    y = random(DualPerplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(DualPerplex{Rational{BigInt}})
    y = random(DualPerplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(DualPerplex{Rational{BigInt}})
    y = random(DualPerplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(DualPerplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    n = norm(x)
    n > big(0) || n == big(0)
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = norm(x * y)
    r = norm(x) * norm(y)
    l == r
end

@test_throws ErrorException begin
    inv(DualPerplex(0, 0.0))
end

@test_throws ErrorException begin
    random(DualPerplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(DualPerplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(DualPerplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(DualPerplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(DualPerplex{BigInt})
    z == z + zero(DualPerplex{BigInt})
end

@test begin
    z = random(DualPerplex{BigInt})
    z == z * one(DualPerplex{BigInt})
end

@test begin
    z = random(DualPerplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(DualPerplex{BigInt})
    z == z * one(z)
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = commutator(x, y)
    r = zero(DualPerplex{BigInt})
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = alternatorL(x, y)
    r = zero(DualPerplex{BigInt})
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = alternatorR(x, y)
    r = zero(DualPerplex{BigInt})
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    l = flexator(x, y)
    r = zero(DualPerplex{BigInt})
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = associator(x, y, z)
    r = zero(DualPerplex{BigInt})
    l == r
end

@test begin
    x = random(DualPerplex{BigInt})
    y = random(DualPerplex{BigInt})
    z = random(DualPerplex{BigInt})
    l = jacobiator(x, y, z)
    r = 3 * x * y * z
    l == r
end

@test begin
    l = DualPerplex(Prplex(1,2))
    r = DualPerplex(1,2,0)
    l == r
end

@test begin
    l = DualPerplex(1, 2, 3.0)
    r = DualPerplex(1.0, 2.0, 3.0)
    l == r
end

@test begin
    l = DualPerplex(1, 2, 3, 4.0)
    r = DualPerplex(1.0, 2.0, 3.0, 4.0)
    l == r
end


#############################################################
#
#  add extra quaternion operations for Quaternions
#  some should probably be merged back into Quaternions
#
#############################################################

#=
"""
    mean(qvec, [method = 0])

Function to find the mean of a set of rotations.

Args:

    qvec    - Vector of `Rotation`s

    method  - if 0 (default), the returned rotation minimizes sum_i=1:n (||R(Qi) - R(Qbar)||)  with || the Frobenius norm, or
              equivalently minimizes sum_i=1:n (sin(phi_i / 2))^2, where phi_i is the angle of Qe_i, with  Qbar = Q_i x Qe_i

            - if 1, the returned rotation minimizes sum_i=1:n (phi_i ^2) where phi_i is the angle of Qe_i, with  Qbar = Q_i x Qe_i
              this should be about the same as method == 0 if phi_i are all small (note: still not implemented)
"""
=#

"""
    mean(r::AbstractVector{Rotation})  ->  r̄

Find the mean `r̄` of a set of rotations `r`. The returned rotation minimizes `sum_i=1:n (||r[i] - r̄||)`  with `||` the Frobenius norm, or
equivalently minimizes `sum_i=1:n (sin(ϕ[i] / 2))^2`, where `ϕ[i] = rotation_angle(r[i] / r̄)`.
"""
function Statistics.mean(qvec::AbstractVector{QuatRotation{T}}, method::Integer = 0) where T
    #if (method == 0)
        M = zeros(4, 4)
        for i in 1:length(qvec)
            q = qvec[i]
            Qi = @SVector [q.w, q.x, q.y, q.z]  # convert types to ensure we don't get quaternion multiplication
            M .+= Qi * (Qi')
        end
        evec = eigfact(Symmetric(M), 4:4)
        Qbar = QuatRotation(evec.vectors[1], evec.vectors[2], evec.vectors[3], evec.vectors[4]) # This will renormalize the quaternion...
    #else
    #    error("I haven't coded this")
    #end

    return Qbar
end

function Statistics.mean(vec::AbstractVector{R}) where R<:Rotation
    R(mean(convert(Vector{QuatRotation{eltype(R)}}, vec)))
end


cd(@__DIR__)
using Pkg
Pkg.activate(".")

using Flux
using Plots
using Flux: @epochs, @show
using BSON: @save
using Random
using Statistics

#comment to test branch: test_branch
#function for parsing data
function get_data(filename)
    f1 = open(filename)
    tempdata = readlines(f1)
    tempdata[1] = tempdata[1][4:end] #depends on file TODO: may have to change
    close(f1)
    data = []
    regvec = zeros(length(tempdata),
                   length(split(tempdata[1], ",")))
    #parsing datavector to float array... not regularized
    for i in 1:length(tempdata)
        regvec[i,:] =  parse.(Float64,split(tempdata[i], ","))
    end
    return regvec
end

#normalizing data
function norm_data(data, scalingmatrix=nothing)
    #getting regularization values and normalizing data
    regvec = data
    normvec = zeros(length(regvec[:,1]),
                    length(regvec[1,:]))
    #checkign if optional scalingmatrix was passed, if not will calc mean and std
    if scalingmatrix == nothing
        scalingmatrix = zeros(length(regvec[1,:]),2)
        for i in 1:length(regvec[1,:])
            scalingmatrix[i,1] = Statistics.mean(regvec[:,i]) #first index is mean
            scalingmatrix[i,2] = Statistics.std(regvec[:,i])  #second index is std
        end
        for i in 1:length(regvec[1,:])
            #normalizing regvec
            normvec[:,i] = (regvec[:,i] .- scalingmatrix[i,1]) ./ scalingmatrix[i,2]
        end
    else
        #for case of normalizing test data by training std and mean
        for i in 1:length(regvec[1,:])
            #normalizing regvec
            normvec[:,i] = (regvec[:,i] .- scalingmatrix[i,1]) ./ scalingmatrix[i,2]
        end
    end
    return normvec, scalingmatrix
end

function vecvec_to_matrix(vecvec)
    dim1 = length(vecvec)
    dim2 = length(vecvec[1])
    my_array = zeros(dim1, dim2)
    for i in 1:dim1
        for j in 1:dim2
            my_array[i,j] = vecvec[i][j]
        end
    end
    return my_array
end

#getting training data
fname = "Dataset/FINAL_DATA.txt"
data_tmp = get_data(fname)
#seperating test and train data
train_data_tmp = []; test_data_tmp = [];
for i in 1:length(data_tmp[:,1])
    if data_tmp[i,2] in [0.3,0.38,0.46,0.53,0.61,0.65] #training case
        push!(train_data_tmp, data_tmp[i,:])
    else #testing case
        push!(test_data_tmp, data_tmp[i,:])
    end
end
train_data = vecvec_to_matrix(train_data_tmp)
test_data = vecvec_to_matrix(test_data_tmp)
#normalizing data
train_norm, scalingmatrix = norm_data(train_data)

#converting data to format which works with network
data = []
for i in 1:length(train_norm[:,1])
    #adjust x and y for fault detection here
    push!(data,(train_norm[i,[1,2,4]],train_norm[i,3]))
end
#network initialization function (allows for easier iterations)
function networkitr(data,Q,wd,iterations)
    #model... must adjust if you want a different structure
    itrmodel = Chain(Dense(3,Q),
                        Dense(Q,Q,gelu),
                        Dense(Q,1))
    opt = ADAM()
    para = Flux.params(itrmodel) # variable to represent all of our weights and biases
    l1(x) = sum(x .^ 2)
    loss(x, y) = Flux.mse(itrmodel(x), y) + wd * sum(l1, para) # cost function
    loss_wout_reg(x,y) = Flux.mse(itrmodel(x), y)
    #=
    for 500 iterations of our entire data set (@epochs 500) model is "trained"
    Flux.train!( cost function, weights and biases, training data, optimizer setting)
    =#
    lowestmse = 1.0
    for i in 1:iterations
        trainset = Random.shuffle(data)[1:64]
        Flux.train!(loss, para, trainset, opt)
        Err = 0.0
        for datapoint in data
            Err += loss_wout_reg(datapoint[1],datapoint[2])
        end
        Err = Err/length(data)
        #calculating lowestmse
        if lowestmse > Err
            global iterationbest = itrmodel #setting best model to lowestmse
            lowestmse = Err
        end
        #print iteration data
        print("\rError $lowestmse Itr $i/$iterations         ")
    end
    return iterationbest, lowestmse
end

#iterating training diff networks
network_dict = Dict{Int8,Float64}(8 => 1.0, 16 => 1.0, 32 => 1.0, 64 => 1.0)

lowestmse_overall = 1.0
for q in [8,16,32,64]
    for wd in [0,0.0001,0.00001,0.000001,0.0000001]
        #training netowrk iteration
        print("\nTesting q=$q,wd=$wd\n")
        global iteration_best, lowestmse_itr = networkitr(data,q,wd,1000)
        if network_dict[q] > lowestmse_itr
            network_dict[q] = lowestmse_itr
        end
        if lowestmse_itr < lowestmse_overall
            print("\nlowestmse = $lowestmse_itr  < best overall = $lowestmse_overall\n")
            print("Replacing Network with best network: ")
            global best_string = "$q Nodes, $wd regularization parameter\n"
            print(best_string)
            global overall_best = iteration_best
            global lowestmse_overall = lowestmse_itr
        end
    end
end
print("Best Network: ", best_string, "with $lowestmse_overall MSE")


using Distances, StructuredDistances, InteractionNetworkModels
using Plots, BenchmarkTools, Distributions, StatsPlots

d = Cityblock()

V = 9
τ = 0.2
mode = [rand() < 0.1 ? rand(1:4) : 0 for i in 1:V,j in 1:V]
γ = 1.9

model = SNF(mode, γ, d)
model = vectorise(model)

# Gibbs scan
gibbs_move_rand = GibbsRandMove(ν=1)
gibbs_move = GibbsScanMove(ν=1)

mcmc = McmcSampler(gibbs_move_rand)
mcmc_scan = McmcSampler(gibbs_move)

@time x = mcmc(model, desired_samples=10000, lag=1)
plot(x)

sum(mode)

acceptance_prob(mcmc_scan)
acceptance_prob(mcmc)
typeof(model)

# Testing posterior sampler
@time out  = mcmc_scan(model, desired_samples=100, lag=5, burn_in=500)
plot(out)
similar(model, rand(1:10, 20), 1.0)

model.mode
data = out.sample
γ_prior = Gamma(model.γ, 1)
plot(γ_prior)
G_prior = SNF(
    model.mode,
    0.1,
    model.d,
    directed=model.directed,
    self_loops=model.self_loops
)
posterior = SnfPosterior(data, G_prior, γ_prior)
mode_move = GibbsScanMove(ν=1)
aux_mcmc = McmcSampler(GibbsRandMove(ν=1), burn_in=3000, lag=50)

@time x = aux_mcmc(model, desired_samples=100)
acceptance_prob(aux_mcmc)
plot(x)

mcmc_posterior = SnfPosteriorSampler(
    mode_move, aux_mcmc, posterior,
    ε=0.1,
    aux_init_at_prev=true
)
mcmc_posterior.aux.data

@time x = mcmc_posterior(posterior, desired_samples=500, γ_init=4.0)

plot(x, model.mode)
x_sample_mat = get_sample_matrices(x)[end]

data

# Plot graph
using CairoMakie, GraphMakie, Graphs

g = SimpleDiGraph(x_sample_mat[1])

f,ax,p = graphplot(
    g,
    node_color=:blue,
    edge_plottype=:linesegments,
    showaxis=false
)
hidedecorations!(ax)  # hides ticks, grid and lables
hidespines!(ax)  # hide the frame
f

V = 15
g = SimpleDiGraph(
    [rand() < 0.5 ? true : false for i in 1:V, j in 1:V]
)

f, ax, p = graphplot(
    g,
    node_size=10,
    node_color=:red,
    edge_plottype=:linesegments
)
hidedecorations!(ax)  # hides ticks, grid and lables
hidespines!(ax)  # hide the frame
f


Base.@kwdef struct LinearAnalysis{I}
    npoints::Int=200
    dropcollinear::Bool=false
    interval::I=automatic
end

add_intercept_column(x::AbstractVector{T}) where {T} = [ones(T, length(x)) x]

# TODO: add multidimensional version
function (l::LinearAnalysis)(input::ProcessedLayer)
    output = map(input) do p, n
        x, y = p
        weights = get(n, :weights, similar(x, 0))
        default_interval = length(weights) > 0 ? nothing : :confidence
        interval = l.interval === automatic ? default_interval : l.interval
        # FIXME: handle collinear case gracefully
        lin_model = GLM.lm(add_intercept_column(x), collect(y); wts=weights, l.dropcollinear)
        x̂ = range(extrema(x)..., length=l.npoints)
        pred = GLM.predict(lin_model, add_intercept_column(x̂); interval)
        return if !isnothing(interval)
            ŷ, lower, upper = map(vec, pred) # GLM prediction returns matrices
            (x̂, ŷ), (; lower, upper)
        else
            ŷ = vec(pred) # GLM prediction returns matrix
            (x̂, ŷ), (;)
        end
    end
    default_plottype = isempty(output.named) ? Lines : LinesFill
    plottype = Makie.plottype(output.plottype, default_plottype)
    return ProcessedLayer(output; plottype)
end

"""
    linear(; interval=automatic, dropcollinear=false, npoints=200)

Compute a linear fit of `y ~ 1 + x`. An optional named mapping `weights` determines the weights.
Use `interval` to specify what type of interval the shaded band should represent.
Valid values of interval are `:confidence` delimiting the uncertainty of the predicted
relationship, and `:prediction` delimiting estimated bounds for new data points.
By default, this analysis errors on singular (collinear) data. To avoid that,
it is possible to set `dropcollinear=true`.
`npoints` is the number of points used by Makie to draw the line
"""
linear(; options...) = transformation(LinearAnalysis(; options...))


using .BifurcationKit
using .Setfield

# ------------------------------------------------------------------------------------------------ #
#                           Generate a bifurcation diagram from a Process                          #
# ------------------------------------------------------------------------------------------------ #
"""Ok, so this one plots a bifurcation diagram using the equations defining a Process.
You can give the single parameter you want to vary, param, and the range over which you want to vary it, prange. If you don't supply prange, I'll take the parameter range as the extrema of the parameter function in the Process. Note that this will only work for Processes defined using ODEproblems from DifferentialEquations. We'll discard the jacobian and compute it automatically.
"""
function bifurcation_diagram(P::Process; param=1, prange=extrema(parameterseries(P, p=param)), kwargs...)
    f = constantine(P; param)
    ode = ODEProblem(f, P.X0, (P.transient_t0, P.tmax); atol = 1e-10, rtol = 1e-9)
    opts = ContinuationPar(pMin = min(prange...)|>Float64,
                            pMax = max(prange...)|>Float64,
                            ds = 0.01,
                            dsmax = 0.05,
	                        # nInversion = 16,
                            # detectBifurcation = 3,
                            # maxBisectionSteps = 50,
                            nev = 3,
                            kwargs...) # parameterseries(P, p=param)[1]
    br, = continuation((x, p) -> Array(f(x, p, 0.0)), P.X0, [min(prange...)],
                            (@lens _[1]), opts;
                            plot = true,
                            verbosity = 0)
end
export bifurcation_diagram

# jet = BifurcationKit.get3Jet(f)
# opts = ContinuationPar(pMin = min(prange...)|>Float64,
#                         pMax = max(prange...)|>Float64,
#                         ds = 0.01,
#                         dsmax = 0.05,
#                         # nInversion = 16,
#                         # detectBifurcation = 3,
#                         # maxBisectionSteps = 50,
#                         nev = 3)
# bifurcationdiagram(jet...,
#     P.X0, 10.0, (@lens _[1]),
# 	# important argument: this is the maximal
# 	# recursion level
# 	5,
#     opts)

function constantine(P::Process; param)
    # Want to turn our special non-stationary EOM to a normal, constant parameter one
    eom = P.process
    p_profiles = forcevec(deepcopy(P.parameter_profile))
    p_parameters= forcevec(deepcopy(P.parameter_profile_parameters))
    p_profiles′ = (p_profiles[1:param-1]..., constantParameter, p_profiles[param+1:end]...)
    function f(X::AbstractArray, 𝑝, 𝑡)
        p_parameters′ = (p_parameters[1:param-1]..., 𝑝[1], p_parameters[param+1:end]...)
        if p_profiles′ isa Union{Tuple, Vector} && length(p_profiles′) == 1
            p_profiles′ = p_profiles′[1]
        end
        pf = tuplef2ftuple(p_profiles′, p_parameters′)
        eom(X, pf, 𝑡)
    end
    return f
end

# optnew = NewtonPar(verbose = true, tol = 1e-12)
# opts = ContinuationPar(dsmin = 0.0001, dsmax = 0.01, ds = 0.01, pMax = 1.,
# 	newtonOptions = setproperties(optnew; maxIter = 30, tol = 1e-8),
# 	maxSteps = 300, plotEveryStep = 40,
# 	detectBifurcation = 3, nInversion = 4, tolBisectionEigenvalue = 1e-17, dsminBisection = 1e-7)

#     function optrec(x, p, l; opt = opts)
#         return setproperties(opt; maxSteps = 300, detectBifurcation = 3, detectLoop = false)
#     end

# g = (x, p) -> collect(f(x, p, 0.0))
# J = (x, p) -> x0 -> BifurcationKit.ForwardDiff.jacobian(y->g(y, p), x0)
# jet = BifurcationKit.get3Jet(g, J)
# diagram = @time bifurcationdiagram(jet...,
# 	[0.1], 1.0, (@lens _[1]),
# 	# here we specify a maximum branching level of 4
# 	4, optrec)


function exposurize_model(model; exposure = 1)
    result = Array{Union{Nothing, LRMoE.AnyExpert}}(nothing, size(model)[1], size(model)[2], length(exposure))
    for i in 1:length(exposure)
        result[:,:,i] = LRMoE.exposurize_expert.(model, exposure = exposure[i])
    end
    return result
end

function loglik_aggre_dim(ls)
    return Matrix(sum(ls, dims = 1)[1,:,:]')
end

function loglik_aggre_gate_dim(gate, ll_mat)
    return gate + ll_mat
end

function loglik_aggre_gate_dim_comp(ll_mat)
    return rowlogsumexp(ll_mat)
end

# Exact case
function expert_ll_ind_mat_exact(Y, model)
    result = fill(NaN, size(model)[1], size(model)[2])
    for j in 1:size(model)[1]
        for k in 1:size(model)[2]
            result[j, k] = LRMoE.expert_ll_exact.(model[j, k], Y[j])
        end
    end
    return result # cat([[LRMoE.expert_ll_exact.(model[j, k], Y[j]) for j in 1:size(model)[1]] for k in 1:size(model)[2]]..., dims = 2)
end

function expert_ll_list_exact(Y, model)
    result = fill(NaN, size(model)[1], size(model)[2], size(Y)[1])
    for i in 1:size(Y)[1]
        result[:,:,i] = expert_ll_ind_mat_exact(Y[i,:], model[:,:,i])
    end
    return result # cat([expert_ll_ind_mat_exact(Y[i,:], model[:,:,i]) for i in 1:size(Y)[1]]..., dims = 3)
end


__precompile__()

"""
*MultipleTesting* package

* Adjusting p-values for multiple testing
* Estimating the fraction π₀ of tests under the null hypothesis
* Combination of p-values
"""
module MultipleTesting

using StatsBase
import StatsBase: fit

using Distributions
import Distributions: estimate

import SpecialFunctions: digamma

import Random: shuffle!

export
    PValues,
    adjust,
    PValueAdjustment,
    Bonferroni,
    BenjaminiHochberg,
    BenjaminiHochbergAdaptive,
    BenjaminiYekutieli,
    BenjaminiLiu,
    Hochberg,
    Holm,
    Hommel,
    Sidak,
    ForwardStop,
    BarberCandes,
    estimate_pi0,
    Pi0Estimator,
    Storey,
    StoreyBootstrap,
    LeastSlope,
    Oracle,
    TwoStep,
    RightBoundary,
    CensoredBUM,
    CensoredBUMFit,
    BUM,
    BUMFit,
    FlatGrenander,
    ConvexDecreasing,
    ConvexDecreasingFit,
    isin,
    fit,
    BetaUniformMixtureModel,
    PValueCombination,
    combine,
    FisherCombination,
    LogitCombination,
    StoufferCombination,
    TippettCombination,
    SimesCombination,
    WilkinsonCombination,
    MinimumCombination,
    estimate,
    HigherCriticismScores,
    HigherCriticismThreshold

include("types.jl")
include("utils.jl")
include("pval-adjustment.jl")
include("pi0-estimators.jl")
include("model.jl")
include("combinations.jl")
include("higher-criticism.jl")

end


"""
Load data from disk into a vector of tuples of the form (X2, T1, X1, T1, Y)
and a vector of frequency counts. This compressed representation of the data
allows approximately a 10x reduction in processing time because redundant
computations are avoided.
"""
function load()
    # Load raw data from CSV file.
    df = DataFrames.readtable("data/choices.csv")

    # Drop nulls.
    df = df[!DataFrames.isna(df[:LaterOptionChosen]), :]

    # Convert to non-nullable arrays.
    X = hcat(
        convert(Array{Float64}, df[:X2]),
        convert(Array{Float64}, df[:T2]),
        convert(Array{Float64}, df[:X1]),
        convert(Array{Float64}, df[:T1]),
    )
    y = convert(Array{Float64}, df[:LaterOptionChosen])

    # Normalize money measurements.
    X[:, 1] = X[:, 1] ./ maximum(X[:, 1])
    X[:, 3] = X[:, 3] ./ maximum(X[:, 3])

    # Translate data into a tuple format to maximize cache locality.
    n = size(X, 1)
    all_tuples = Array(NTuple{5, Float64}, n)
    for i in 1:n
        all_tuples[i] = (X[i, 1], X[i, 2], X[i, 3], X[i, 4], y[i])
    end

    # Work with the unique values in the input and their frequencies to
    # compress data for faster computations.
    unique_tuples_dict = StatsBase.countmap(all_tuples)
    unique_tuples = collect(keys(unique_tuples_dict))
    w = collect(values(unique_tuples_dict))

    return unique_tuples, w
end


NORM_TYPE = Inf

function lb_s_thres(iter::Class_iterate,dir::Class_point,pars::Class_parameters)
  ex = pars.ls.fraction_to_boundary_predict_exp
  #alpha_P_tilde = simple_max_step(iter.point.s,dir.s,zeros(length(dir.s)))
  #x_thres = max.(2.0 * LinearAlgebra.norm(dir.x,NORM_TYPE)^(2 + ex), LinearAlgebra.norm(dir.x,NORM_TYPE)^(1.0 + ex) + LinearAlgebra.norm(dir.y,NORM_TYPE)^0.5 * LinearAlgebra.norm(dir.x,NORM_TYPE))
  x_thres = LinearAlgebra.norm(dir.x,NORM_TYPE)^ex #+ LinearAlgebra.norm(dir.y,NORM_TYPE) + get_delta(iter)

  return min.(iter.point.s, LinearAlgebra.norm(dir.x,NORM_TYPE) * x_thres)
end

function lb_s_predict(iter::Class_iterate,dir::Class_point,pars::Class_parameters)
  #@show iter.frac_bd_predict
  return iter.frac_bd_predict .* lb_s_thres(iter,dir,pars)
end

function lb_y(iter::Class_iterate,dir::Class_point,pars::Class_parameters)
  #alpha_P_tilde = simple_max_step(iter.point.s,dir.s,zeros(length(dir.s)))
  return iter.frac_bd .*  iter.point.y * min(1.0, LinearAlgebra.norm(dir.x,NORM_TYPE) )
end

function lb_s(iter::Class_iterate,dir::Class_point,pars::Class_parameters)
    #alpha_P_tilde = simple_max_step(iter.point.s,dir.s,zeros(length(dir.s)))
    lb_s = iter.frac_bd .* lb_s_thres(iter,dir,pars)
    #min.(iter.point.s, LinearAlgebra.norm(dir.x,NORM_TYPE)^2)
    #@show lb_s
    return lb_s
end

function simple_max_step(val::Array{Float64,1}, dir::Array{Float64,1}, lb::Array{Float64,1})
    ratio = maximum( [1.0; -dir ./ (val - lb) ] ) #; -q * dir.y ./ iter.point.y] );
    max_step = 1.0 / ratio
    #@show max_step
    return max_step
end


using LinearAlgebra

ex=Ten([1,2,3],:i)*Ten([90,80,70],:i)
@test simplify(ex)==460

ex=Ten(:A,:i)*Ten(:A,:j)
@test simplify(ex&@equs(A=[1,2,3], j=i))==14

a=Alt([:i,:j])
@test a.x[1,2]==1
@test a.x[2,1]==-1
a=Alt([:i,:j,:k])
@test a.x[1,2,3]==1
@test a.x[3,2,1]==-1
a=Alt([:i,:j,:k,:l])
@test a.x[1,2,3,4]==1
@test a.x[4,2,3,1]==-1

ex=Alt([:i,:j,:k])*Ten(:a,:j)*Ten(:b,:k);a=[5,123,-12];b=[90,80,70];ex=ex&Equation(:a,a)&Equation(:b,b)
@test ex==Ten(Any[9570.0,-1430.0,-10670.0],Any[:i])
c=cross(a,b)
@test c[1]==ex&@equ(i=1)
@test c[2]==ex&@equ(i=2)
@test c[3]==ex&@equ(i=3)
ex=Alt([1,2,:k])*Ten([1,2,3],2)*Ten([90,80,70],:k);sex=simplify(ex)
@test sex==140
ex=Alt([1,3,:k])*Ten([1,2,3],3)*Ten([90,80,70],:k);sex=simplify(ex)
@test sex==-240
ex=Alt([1,:j,:k])*Ten([1,2,3],:j)*Ten([90,80,70],:k);sex=simplify(ex)
@test sex==-100
r=Alt([:i,:j,:k])*Ten([:a1,:a2,:a3],:j)*Ten([:b1,:b2,:b3],:k)&@equ i=1
@test r==simplify(:a2*:b3+-1.0*:a3*:b2)

ex=:c+Ten([:a1,:a2,:a3],:i)*Ten([:b1,:b2,:b3],:i)+:c;nex=simplify(ex)
@test nex==simplify(2*:c+:a1*:b1+:a2*:b2+:a3*:b3)

@test simplify(Ten(zeros(3,3).+Diagonal([1,1,1]),[:i,:i]))==3
@test simplify(Ten(zeros(4,4).+Diagonal([1,1,1,1]),[:i,:i]))==4
@test simplify(Ten(zeros(5,5).+Diagonal([1,1,1,1,1]),[:i,:i]))==5
r=Ten(:A,[:i,:i])&@equ A=[:a 0;0 :b]
@test r==:a+:b

io=IOBuffer()
print(io,simplify(Ten([1,2],:j)*Ten([3 2;1 -1],[:i,:j])))
str=String(take!(io))
#@test str=="Any[3,1](i) + 2 [2,-1](i)"
io=IOBuffer()
print(io,simplify(Ten([1,2],:j)*Ten([3 2;1 -1],[:j,:i])))
str=String(take!(io))
#@test str=="Any[3 2](i) + 2 [1 (-1)](i)"

@test duplicates([1,2,3,4,2])==[2,5]
@test duplicates([1,2,3,4,2,2])==[5,6]
@test duplicates([1,2,3],[0,3,0])==[3,2]
@test duplicates([1,3,3],[3,3,3])==[3,3]
@test Equations.arrduplicates([1,3,5],[12,241,13,3],[3,1,1])==([1,2],[2,4])
@test Equations.arrduplicates([1,3,55],[12,241,13],[33,1324,5,6,1])==([1,3],[1,5])

ex=Ten(:A,[:j,:i,:i])*Ten(:B,:j);r=ex&@equs(A=ones(3,3,3), B=[1,2,3])
@test r==18

eye3=zeros(3,3).+Diagonal([1,1,1])
ex=Ten(:A,[:i,:i]);eq=@equ A=$eye3;r=ex&eq
@test r==3

r=simplify(Ten([1,0],:i)+Ten([0,1],:i))
@test r==Ten([1,1],:i)
r=simplify(Ten([:c,:d],:i)+Ten([:a,:b],:i))
@test r==Ten(Any[:c+:a,:d+:b],Any[:i])
r=simplify(Alt([:i,:j,:k])*Ten([:a1,:a2,:a3],:j)*Ten([:b1,:b2,:b3],:k))

tt=Term[Factor[3,Ten([:a,:b],:i)],Factor[5,Ten([:c,:d],:i)]];stt=Equations.sumlify(tt)
@test stt[1][1].x[1]==simplify(Any[3*:a+5*:c,3*:b+5*:d][1])
@test stt[1][1].x[2]==simplify(Any[3*:a+5*:c,3*:b+5*:d][2])
#stt==simplify(Term[Factor[Equations.Ten(Any[3*:a+5*:c,3*:b+5*:d],Any[:i])]])

ex=Equations.sumlify(Equations.untensify!(sumconv(Alt([:i,:j])*Ten([10,100],:j)).terms))
@test sum(ex[1][1].x)==90
#ex==Term[Factor[Equations.Ten(Any[100.0,-10.0],Any[:i]),10]]

ex=Alt([:i,:j,:k])*Ten([1,0,0],:j)*Ten([0,0,1],:k);ex=sumconv(ex);ex=sumconv(ex);tt=Equations.untensify!(ex.terms);ttt=Equations.sumlify(tt)
@test ttt[1][1]&(@equ i=2)==-1

a=rand(Int,3).%9;b=rand(Int,3).%9
ab=[Equation(:a,a),Equation(:b,b)]
ex=Ten(:a,:i)*Ten(:b,:i)
@test ex&ab==dot(a,b)
ex=Alt([:i,:j,:k])*Ten(:a,:j)*Ten(:b,:k)
c=cross(a,b)
x=(ex&ab).x
@test c==x
ex=Alt([:i,:j,:k])*Ten(:a,:j)*Ten(:b,:k)*Ten(:a,:m)*Ten(:b,:m)
r=cross(a,b)*dot(a,b)
x=(ex&ab).x
@test r==x
ex=Alt([:i,:j,:k])*Ten(:a,:j)*Ten(:b,:k)*Ten(:a,:m)*Ten(:b,:m)+Ten(:b,:i)
r=cross(a,b)*dot(a,b)+b
x=(ex&ab).x
@test r==x
c=rand(Int,3).%9;d=rand(Int,3).%9
abcd=[Equation(:a,a),Equation(:b,b),Equation(:c,c),Equation(:d,d)]
ex=Alt([:i,:j,:k])*Ten(:a,:j)*Ten(:b,:k)*Alt([:m,:n,:o])*Ten(:c,:n)*Ten(:d,:o)
r=dot(cross(a,b),cross(c,d))
x=ex&abcd&Equation(:m,:i)
@test r==x

B=[:b1 :b2;:b3 :b4];ex=Alt([:i,:j])*Alt([:k,:l])*Ten(B,[:i,:k])*Ten(B,[:j,:l]);r=simplify(ex)
@test r==simplify(2*:b1*:b4+-2.0*:b2*:b3)


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

"""
    intersecttype(f, b1, b2)

Compute the intersection type of two boxes `b1` and `b2`
and apply function `f` to it.

The intersection type can be one of four types:

1. overlap with non-zero measure
2. intersect at one of the boundaries
3. intersect at corner point
4. do not overlap nor intersect
"""
function intersecttype(f::Function, b1::Box{Dim,T}, b2::Box{Dim,T}) where {Dim,T}
  m1, M1 = coordinates.(extrema(b1))
  m2, M2 = coordinates.(extrema(b2))

  # relevant vertices
  u = Point(max.(m1, m2))
  v = Point(min.(M1, M2))

  # branch on possible configurations
  if u ≺ v
    return OverlappingBoxes(Box(u, v)) |> f
  elseif u ≻ v
    return NoIntersection() |> f
  elseif isapprox(u, v, atol=atol(T))
    return CornerTouchingBoxes(u) |> f
  else
    return FaceTouchingBoxes(Box(u, v)) |> f
  end
end

"""
    intersecttype(function, ray, box)

Calculate the intersection type of a `ray`` and a `box`` and apply function `f` to it.

## References

* Williams A, Barrus S, Morley R K, et al., 2005. [An efficient and robust ray-box
  intersection algorithm]
  (https://dl.acm.org/doi/abs/10.1145/1198555.1198748)
"""
function intersecttype(f::Function, r::Ray{Dim,T}, b::Box{Dim,T}) where {Dim,T}
  invdir = one(T) ./ direction(r)
  lo, up = coordinates.(extrema(b))
  orig = coordinates(origin(r))

  tmin = zero(T)
  tmax = typemax(T)

  # check for intersection with slabs along with each axis
  for i in 1:Dim
    imin = (lo[i] - orig[i]) * invdir[i]
    imax = (up[i] - orig[i]) * invdir[i]

    # swap variables if necessary
    invdir[i] < zero(T) && ((imin, imax) = (imax, imin))

    # the ray is on a face of the box, avoid NaN
    (isnan(imin) || isnan(imax)) && continue

    (tmin > imax || imin > tmax) && return NoIntersection() |> f

    tmin = max(tmin, imin)
    tmax = min(tmax, imax)
  end

  tmin ≈ tmax && return TouchingRayBox(r(tmin)) |> f

  return CrossingRayBox(Segment(r(tmin), r(tmax))) |> f
end


# # CO2 and temperature

using AbstractGPs
using AbstractGPsMakie
using CairoMakie
using CSV
using DataDeps
using DataFrames
using Dates
using Distributions
using KernelFunctions
using LinearAlgebra
using Optim
using ParameterHandling
using Random
using Statistics
using Stheno
using Zygote

using CairoMakie: RGB

ENV["DATADEPS_ALWAYS_ACCEPT"] = "true"

# ## Obtain the data

# Use DataDeps.jl to obtain the data.
register(DataDep(
    "mauna_loa",
    "The classic Mauna Loa CO2 data.",
    "https://scrippsco2.ucsd.edu/assets/data/atmospheric/stations/in_situ_co2/monthly/" *
        "monthly_in_situ_co2_mlo.csv",
))
register(DataDep(
    "hadcrut5_global",
    "Globally-averaged surface temperatures",
    "https://www.metoffice.gov.uk/hadobs/hadcrut5/data/current/analysis/diagnostics/" *
        "HadCRUT.5.0.1.0.analysis.summary_series.global.monthly.csv",
))

# Load and pre-process the data.
ml_df, Ttr_df, Tte_df = let

    mauna_loa_co2 = let
        mauna_loa_data_raw = CSV.read(
            joinpath(datadep"mauna_loa", "monthly_in_situ_co2_mlo.csv"), DataFrame;
            skipto=58, header=false,
        )

        data = DataFrame(
            date=map(Date, mauna_loa_data_raw.Column1, mauna_loa_data_raw.Column2),
            CO2=map(x -> x == -99.99 ? missing : x, mauna_loa_data_raw.Column5),
        )
        data = filter(:CO2 => !ismissing, data)

        ## standardise CO2 data to have unit mean.
        m = mean(data.CO2)
        σ = std(data.CO2)
        data.CO2 = Float64.(data.CO2 .- m) ./ σ

        ## Convert dates into floats.
        data.float_date = map(x -> x.value, data.date .- Date(1)) ./ 365

        data
    end

    temperature = let
        hadcrut5_data_raw = CSV.read(
            joinpath(
                datadep"hadcrut5_global",
                "HadCRUT.5.0.1.0.analysis.summary_series.global.monthly.csv",
            ),
            DataFrame;
            header=1,
        )

        tmp = DataFrame(
            date=hadcrut5_data_raw[!, "Time"],
            temperature=hadcrut5_data_raw[!, "Anomaly (deg C)"],
        )

        ## Convert dates into floats.
        tmp.float_date = map(x -> x.value, tmp.date .- Date(1)) ./ 365

        ## Remove any data from before the first Mauna Loa data.
        time_min = minimum(mauna_loa_co2.date)
        filter(:date => date -> date >= time_min, tmp)
    end

    ## Split temperature data into train and test.
    Ntr = 425
    Ttr_raw = temperature[1:Ntr, :]
    Tte_raw = temperature[(Ntr + 1):end, :]

    ## Standard temperature data using training data.
    mtr = mean(Ttr_raw.temperature)
    σtr = std(Ttr_raw.temperature)
    Ttr = DataFrame(
        date=Ttr_raw.date,
        float_date=Ttr_raw.float_date,
        temperature=(Ttr_raw.temperature .- mtr) ./ σtr,
    )
    Tte = DataFrame(
        date=Tte_raw.date,
        float_date=Tte_raw.float_date,
        temperature=(Tte_raw.temperature .- mtr) ./ σtr,
    )

    mauna_loa_co2, Ttr, Tte
end;

# ## Build and train a model for the data

# Model for CO2 and temperature as a function of time.
function build_gp(θ)
    return @gppp let

        ## Shared trend process.
        f_trend = stretch(GP(SEKernel()), θ.λ_trend)

        ## Specify model for CO2.
        f_co2_latent = θ.CO2.σ_latent * f_trend
        f_co2_wiggle = θ.CO2.σ_wiggle * stretch(GP(SEKernel()), θ.CO2.λ_wiggle)
        f_co2_period = θ.CO2.σ_period * GP(SEKernel() ∘ PeriodicTransform(θ.CO2.freq))
        f_co2 = f_co2_latent + f_co2_wiggle + f_co2_period + θ.CO2.σm * GP(ConstantKernel())

        ## Specify model for temperature.
        f_T_trend = θ.T.σ_trend * f_trend
        f_T_wiggle = θ.T.σ_wiggle * stretch(GP(SEKernel()), θ.T.λ_wiggle)
        f_T = f_T_trend + f_T_wiggle + θ.T.σm * GP(ConstantKernel())
    end
end

x_co2 = ml_df.float_date;
x_T_tr = Ttr_df.float_date;
x_T_te = Tte_df.float_date;

y_co2 = ml_df.CO2;
y_T_tr = Ttr_df.temperature;
y_T_te = Tte_df.temperature;

x = BlockData(GPPPInput(:f_co2, x_co2), GPPPInput(:f_T, x_T_tr))
y = vcat(y_co2, y_T_tr)

function build_obs_cov(θ)
    σ²_CO2 = θ.CO2.σ_noise^2
    σ²_T = θ.T.σ_noise^2
    return Diagonal(vcat(fill(σ²_CO2, length(y_co2)), fill(σ²_T, length(y_T_tr))))
end

function nlml(θ)
    f = build_gp(θ)
    return -logpdf(f(x, build_obs_cov(θ)), y)
end

init_params = (
    λ_trend = positive(1e-1),
    CO2 = (
        σ_latent = positive(0.3),
        σ_wiggle = positive(0.3),
        λ_wiggle = positive(5.0),
        σ_period = positive(0.3),
        freq = fixed(1.0),
        σ_noise = positive(0.1),
        σm = positive(1.0),
    ),
    T = (
        σ_trend = positive(1.0),
        σ_wiggle = positive(1.0),
        λ_wiggle = positive(0.5),
        σ_noise = positive(0.1),
        σm = positive(1.0),
    ),
)

default_optimizer = LBFGS(;
    alphaguess=Optim.LineSearches.InitialStatic(; scaled=true),
    linesearch=Optim.LineSearches.BackTracking(),
)

function optimize_loss(loss, θ_init; optimizer=default_optimizer, maxiter=1_000)
    options = Optim.Options(; iterations=maxiter, show_trace=true)

    θ_flat_init, unflatten = ParameterHandling.value_flatten(θ_init)
    loss_packed = loss ∘ unflatten

    ## https://julianlsolvers.github.io/Optim.jl/stable/#user/tipsandtricks/#avoid-repeating-computations
    function fg!(F, G, x)
        if F !== nothing && G !== nothing
            val, grad = Zygote.withgradient(loss_packed, x)
            G .= only(grad)
            return val
        elseif G !== nothing
            grad = Zygote.gradient(loss_packed, x)
            G .= only(grad)
            return nothing
        elseif F !== nothing
            return loss_packed(x)
        end
    end

    result = optimize(Optim.only_fg!(fg!), θ_flat_init, optimizer, options; inplace=false)

    return unflatten(result.minimizer), result
end

θ_opt, result = optimize_loss(nlml, init_params)

# ## Plot the resulting model fit.

function colours()
    return Dict(
        :blue => RGB(0/255, 107/255, 164/255),
        :cyan => RGB(75/255, 166/255, 251/255),
        :red => RGB(200/255, 82 / 255, 0 / 255),
        :pink => RGB(169/255, 90/255, 161/255),
        :black => RGB(0.0, 0.0, 0.0),
        :orange => RGB(245/255, 121/255, 58/255),
    )
end

band_alpha() = 0.3
sample_alpha() = 0.2
point_alpha() = 1.0

x_plot = range(
    minimum(vcat(ml_df.float_date, Ttr_df.float_date)),
    maximum(vcat(ml_df.float_date, Tte_df.float_date));
    length=1_500,
);

f_post = posterior(build_gp(θ_opt)(x, build_obs_cov(θ_opt)), y);

function plot_band!(ax, x_plot, fx, colour, label)
    ms = marginals(fx)
    symband!(
        ax, x_plot, mean.(ms), std.(ms);
        bandscale=3,
        color=(colours()[colour], 0.5 * band_alpha()),
        label=label,
    )
end

function plot_sample!(ax, x_plot, fx, colour)
    gpsample!(
        ax, x_plot, fx;
        samples=4, color=(colours()[colour], sample_alpha()),
    )
end

function plot_gp!(ax, x_plot, fx, colour, label)
    plot_band!(ax, x_plot, fx, colour, label)
    plot_sample!(ax, x_plot, fx, colour)
end

let
    fig = Figure()
    ga = fig[1, 1] = GridLayout()

    jitter = 1e-6

    ## Plot CO2 posterior.
    let
        ax = Axis(ga[1, 1])
        f_post_x = f_post(GPPPInput(:f_co2, x_plot), jitter)

        plot_band!(ax, x_plot, f_post(GPPPInput(:f_co2_latent, x_plot), jitter), :orange, "CO2 (latent)")
        plot_band!(ax, x_plot, f_post(GPPPInput(:f_co2_wiggle, x_plot), jitter), :black, "CO2 (wiggle)")
        plot_band!(ax, x_plot, f_post(GPPPInput(:f_co2_period, x_plot), jitter), :pink, "CO2 (period)")

        plot_gp!(ax, x_plot, f_post_x, :blue, "CO2")
        scatter!(ax, x_co2, y_co2; color=(colours()[:blue], point_alpha()), markersize=3)

        axislegend(ax; position=:lt)

        hidexdecorations!(ax; grid=false)
    end

    ## Plot temperature posterior.
    let
        ax = Axis(ga[2, 1])
        f_post_x = f_post(GPPPInput(:f_T, x_plot), jitter)

        plot_band!(ax, x_plot, f_post(GPPPInput(:f_T_trend, x_plot), jitter), :orange, "T (trend)")
        plot_band!(ax, x_plot, f_post(GPPPInput(:f_T_wiggle, x_plot), jitter), :black, "T (wiggle)")

        plot_gp!(ax, x_plot, f_post_x, :blue, "T")
        scatter!(ax, x_T_tr, y_T_tr; color=(colours()[:blue], point_alpha()), markersize=3)
        scatter!(
            ax, x_T_te, y_T_te;
            color=(colours()[:pink], point_alpha()),
            markersize=5,
            marker=:xcross,
            label="T (test)"
        )

        axislegend(ax; position=:lt)
    end

    rowgap!(ga, 5)

    fig
end


using Manifolds, Manopt, Plots, Test

@testset "Test Gradient checks" begin
    M = Sphere(10)
    q = zeros(11)
    q[1] = 1.0
    p = zeros(11)
    p[1:4] .= 1 / sqrt(4)
    r = log(M, p, q)

    F(M, p) = 1 / 2 * distance(M, p, q)^2
    gradF(M, p) = -log(M, p, q)

    @test check_gradient(M, F, gradF, p, r)

    gradF2(M, p) = -0.5 * log(M, p, q)
    @test_throws ErrorException check_gradient(M, F, gradF2, p, r; error=true)
    @test !check_gradient(M, F, gradF2, p, r)

    check_gradient(M, F, gradF, p, r; plot=true)

    #test windowsize error
    @test_throws ErrorException Manopt.find_best_slope_window(zeros(2), zeros(2), 20)
    @test_throws ErrorException Manopt.find_best_slope_window(zeros(2), zeros(2), [2, 20])
end