using BinaryLambdaCalculus
using Test

@testset "Encoding/decoding" begin
    for t in terms(0, 10)
        @test decode(encode(t)) == t
        @test decompress(compress(t)) == t
    end
end

@testset "Enumeration" begin
    function S(m::T, n::T) where T
        # explicit recursive formula for testing
        if n <= 1
            return 0
        else
            return T(m >= n - 1) + S(m + 1, n - 2) + sum(S(m, k) * S(m, n - 2 - k) for k = 0:(n-2))
        end
    end

    S₀ₙ = [0, 0, 0, 0, 1, 0, 1, 1, 2, 1, 6, 5, 13, 14, 37, 44, 101, 134, 298, 431]
    @test all(S₀ₙ[i+1] == S(0, i) for i = 0:19) # ensure the reference is correct...
    @test [grygiel_lescanne(0, n) for n = 0:19] == S₀ₙ

    for m = 1:3, n = 0:20
        @test grygiel_lescanne(m, n) == S(m, n)
        @test length(terms(m, n)) == S(m, n)
    end

    let m = 3, n = 100,
        s = (a, b) -> grygiel_lescanne(a, b, BigInt)
        @test s(m, n) == s(m + 1, n - 2) + sum(s(m, k) * s(m, n - 2 - k) for k = 0:(n-2))
    end
end


"""
Normalised mean square error on predictions
"""
function NMSE(test::Data, gp::GaussianProcess)
    yp, _ = posterior(test.X, gp)
    y = reduce(vcat, test.Y)

    var_y = sum((y .- (sum(y) / 200)).^2)
    N = size(y)[1]

    (100 / (var_y * N)) * sum((yp .- y).^2)
end


using VoronoiCells
using Base.Test
using GeometricalPredicates

@test (cells = poissonvoronoicells(10^3) ; true)

cells = poissonvoronoicells(10^2)

@test typeof(sarea(cells,1)) == Float64
@test typeof(area(cells[1])) == Float64
@test typeof(random_cell(cells)) == VoronoiCells.VoronoiCell
@test isvalid(findindex(cells,1.5,1.5))

# Test that any point in the bulk is covered by a cell.
function testrand(cells,n)
    c = 0
    s = 1.2
    f = 0.6
    for i in 1:n
        (x,y) = (f*rand()+s,f*rand()+s)
        if VoronoiCells.isexternal(cells, x,y) c += 1 end
    end
    (c,c/n)
end

@test testrand(poissonvoronoicells(10^5),10^5) == (0,0.0)

# Test mean area of cells in the bulk
function meanarea(cells::Array{VoronoiCell,1}, cutoff)
    ma = 0.0
    cnt = 0
    lc = 1 + cutoff
    hc = 2 - cutoff
    for i in 1:length(cells)
        c = cells[i]
        g = getgenerator(c)
        getx(g) < lc && continue
        getx(g) > hc && continue
        gety(g) < lc && continue
        gety(g) > hc && continue
        cnt += 1
        ma += VoronoiCells.area(cells[i])
    end
    ma / cnt
end

@test abs(meanarea(poissonvoronoicells(10^5)._cells,.1) / 1e-5 -1) < 1e-2


@testset "Data" begin
    @test artificialIn_SAR_image(2) == 2 * π * ones(2, 2)

    @test artificial_S1_slope_signal(20, 0.0) == repeat([-π / 2], 20)

    @test ismissing(artificial_S1_signal(-1.0))
    @test ismissing(artificial_S1_signal(2.0))
    @test artificial_S1_signal(2) == [-3 * π / 4, -3 * π / 4]

    # for the remainder check data types only
    @test length(artificial_S1_signal(20)) == 20

    @test size(artificial_S2_whirl_image(64)) == (64, 64)
    @test length(artificial_S2_whirl_image(64)[1, 1]) == 3

    @test size(artificial_S2_rotation_image(64)) == (64, 64)
    @test length(artificial_S2_rotation_image(64)[1, 1]) == 3

    @test size(artificial_S2_whirl_patch(8)) == (8, 8)
    @test length(artificial_S2_whirl_patch(8)[1, 1]) == 3

    @test size(artificial_SPD_image(8)) == (8, 8)
    @test size(artificial_SPD_image(8)[1, 1]) == (3, 3)

    @test size(artificial_SPD_image2(8)) == (8, 8)
    @test size(artificial_SPD_image2(8)[1, 1]) == (3, 3)
    @test eltype(artificial_SPD_image2(8)) == Array{Float64,2}

    @test length(artificial_S2_lemniscate([0.0, 0.0, 1.0], 20)) == 20
    @test length(artificial_S2_lemniscate([0.0, 0.0, 1.0], 20)[1]) == 3
end


"""
    BadStatesProperty{N<:Real} <: Property

Type that represents a safety property characterized by a set of bad states.
The property is satisfied by a given set of states if the intersection with the
set of bad states is empty.

### Fields

- `bad`     -- convex set representing the bad states
- `witness` -- witness point (empty vector if not set)

### Notes

The following formula characterizes whether a set ``X`` satisfies a safety
property characterized by a set of bad states 𝑃:

```math
    X \\models 𝑃 \\iff X ∩ 𝑃.\\texttt{bad} = ∅
```
"""
mutable struct BadStatesProperty{N<:Real} <: Property
    bad::LazySet
    witness::Vector{N}

    BadStatesProperty{N}(bad::LazySet) where {N<:Real} = new(bad, N[])
end

# type-less convenience constructor
BadStatesProperty(bad::LazySet{N}) where {N<:Real} =
    BadStatesProperty{N}(bad)

"""
    dim(𝑃::BadStatesProperty)::Int

Return the dimension of a property with bad states.

### Input

- `𝑃` -- safety property with bad states

### Output

The dimension of the bad states.
"""
function dim(𝑃::BadStatesProperty)::Int
    return dim(𝑃.bad)
end

"""
    check(𝑃::BadStatesProperty, X::LazySet; witness::Bool=false)

Checks whether a convex set is disjoint from the set of bad states.

### Input

- `𝑃`       -- safety property with bad states
- `X`       -- convex set
- `witness` -- (optional, default: `false`) flag for returning a counterexample
               if the property is violated

### Output

Let ``Y`` be the bad states represented by 𝑃.
* If `witness` option is deactivated: `true` iff ``X ∩ Y = ∅``
* If `witness` option is activated:
  * `(true, [])` iff ``X ∩ Y = ∅``
  * `(false, v)` iff ``X ∩ Y ≠ ∅`` and ``v ∈ X ∩ Y``
"""
@inline function check(𝑃::BadStatesProperty, X::LazySet; witness::Bool=false)
    return isdisjoint(X, 𝑃.bad, witness)
end


# # [Quantum Circuit Born Machine](@id qcbm)

# Yao is designed with variational quantum circuits in mind, and this tutorial
# will introduce how to use Yao for this kind of task by implementing a quantum
# circuit born machine described in [Jin-Guo Liu, Lei Wang (2018)](https://arxiv.org/abs/1804.04168)

# let's use the packages first

using Yao, LinearAlgebra, Plots

# ## Training Target

# In this tutorial, we will ask the variational circuit to learn the most basic
# distribution: a guassian distribution. It is defined as follows:

# ```math
# f(x \left| \mu, \sigma^2\right) = \frac{1}{\sqrt{2\pi\sigma^2}} e^{-\frac{(x-\mu)^2}{2\sigma^2}}
# ```

# We implement it as `gaussian_pdf`:

function gaussian_pdf(x, μ::Real, σ::Real)
    pl = @. 1 / sqrt(2pi * σ^2) * exp(-(x - μ)^2 / (2 * σ^2))
    pl / sum(pl)
end
pg = gaussian_pdf(1:1<<6, 1<<5-0.5, 1<<4);

# We can plot the distribution, it looks like

plot(pg)

# ## Create the Circuit

# A quantum circuit born machine looks like the following:

# ![differentiable ciruit](assets/differentiable.png)

# It is composited by two different layers: **rotation layer** and **entangler layer**.

# ## Rotation Layer

# Arbitrary rotation is built with **Rotation Gate** on **Z, X, Z** axis
# with parameters.

# ```math
# Rz(\theta) \cdot Rx(\theta) \cdot Rz(\theta)
# ```

# Since our input will be a ``|0\dots 0\rangle`` state.
# The first layer of arbitrary rotation can just
# use ``Rx(\theta) \cdot Rz(\theta)`` and the last
# layer of arbitrary rotation could just
# use ``Rz(\theta)\cdot Rx(\theta)``

# In **幺**, every Hilbert operator is a **block** type, this
# ncludes all **quantum gates** and **quantum oracles**.
# In general, operators appears in a quantum circuit
# can be divided into **Composite Blocks** and **Primitive Blocks**.

# We follow the low abstraction principle and
# thus each block represents a certain approach
# of calculation. The simplest **Composite Block**
# is a **Chain Block**, which chains other blocks
# (oracles) with the same number of qubits together.
# It is just a simple mathematical composition of
# operators with same size. e.g.

# ```math
# \text{chain(X, Y, Z)} \iff X \cdot Y \cdot Z
# ```

# We can construct an arbitrary rotation block by chain ``Rz``, ``Rx``, ``Rz`` together.

chain(Rz(0.0), Rx(0.0), Rz(0.0))

# `Rx`, `Rz` will construct new rotation gate,
# which are just shorthands for `rot(X, 0.0)`, etc.

# Then let's chain them up

layer(nbit::Int, x::Symbol) = layer(nbit, Val(x))
layer(nbit::Int, ::Val{:first}) = chain(nbit, put(i=>chain(Rx(0), Rz(0))) for i = 1:nbit);

# We do not need to feed the first `n` parameter into `put` here.
# All factory methods can be **lazy** evaluate **the first arguements**, which is the number of qubits.
# It will return a lambda function that requires a single interger input.
# The instance of desired block will only be constructed until all the information is filled.
# When you filled all the information in somewhere of the declaration, 幺 will be able to infer the others.
# We will now define the rest of rotation layers

layer(nbit::Int, ::Val{:last}) = chain(nbit, put(i=>chain(Rz(0), Rx(0))) for i = 1:nbit)
layer(nbit::Int, ::Val{:mid}) = chain(nbit, put(i=>chain(Rz(0), Rx(0), Rz(0))) for i = 1:nbit);

# ## Entangler

# Another component of quantum circuit born machine are
# several **CNOT** operators applied on different qubits.

entangler(pairs) = chain(control(ctrl, target=>X) for (ctrl, target) in pairs);

# We can then define such a born machine

function build_circuit(n, nlayers, pairs)
    circuit = chain(n)
    push!(circuit, layer(n, :first))
    for i in 2:nlayers
        push!(circuit, cache(entangler(pairs)))
        push!(circuit, layer(n, :mid))
    end
    push!(circuit, cache(entangler(pairs)))
    push!(circuit, layer(n, :last))
    return circuit
end

# We use the method `cache` here to tag the entangler block that it
# should be cached after its first run, because it is actually a
# constant oracle. Let's see what will be constructed

build_circuit(4, 1, [1=>2, 2=>3, 3=>4])

# ## MMD Loss & Gradients

# The MMD loss is describe below:

# ```math
# \begin{aligned}
# \mathcal{L} &= \left| \sum_{x} p \theta(x) \phi(x) - \sum_{x} \pi(x) \phi(x) \right|^2\\
#             &= \langle K(x, y) \rangle_{x \sim p_{\theta}, y\sim p_{\theta}} - 2 \langle K(x, y) \rangle_{x\sim p_{\theta}, y\sim \pi} + \langle K(x, y) \rangle_{x\sim\pi, y\sim\pi}
# \end{aligned}
# ```

# We will use a squared exponential kernel here.

struct RBFKernel
    σ::Float64
    m::Matrix{Float64}
end

function RBFKernel(σ::Float64, space)
    dx2 = (space .- space').^2
    return RBFKernel(σ, exp.(-1/2σ * dx2))
end

kexpect(κ::RBFKernel, x, y) = x' * κ.m * y

# There are two different way to define the loss:

# In simulation we can use the probability distribution of the state directly

get_prob(qcbm) = probs(zero_state(nqubits(qcbm)) |> qcbm)

function loss(κ, c, target)
    p = get_prob(c) - target
    return kexpect(κ, p, p)
end

# Or if you want to simulate the whole process with measurement (which is entirely
# physical), you should define the loss with measurement results, for convenience
# we directly use the simulated results as our loss

# ### Gradients
# the gradient of MMD loss is

# ```math
# \begin{aligned}
# \frac{\partial \mathcal{L}}{\partial \theta^i_l} &= \langle K(x, y) \rangle_{x\sim p_{\theta^+}, y\sim p_{\theta}} - \langle K(x, y) \rangle_{x\sim p_{\theta}^-, y\sim p_{\theta}}\\
# &- \langle K(x, y) \rangle _{x\sim p_{\theta^+}, y\sim\pi} + \langle K(x, y) \rangle_{x\sim p_{\theta^-}, y\sim\pi}
# \end{aligned}
# ```

# which can be implemented as

function gradient(qcbm, κ, ptrain)
    n = nqubits(qcbm)
    prob = get_prob(qcbm)
    grad = zeros(Float64, nparameters(qcbm))

    count = 1
    for k in 1:2:length(qcbm), each_line in qcbm[k], gate in content(each_line)
        dispatch!(+, gate, π/2)
        prob_pos = probs(zero_state(n) |> qcbm)

        dispatch!(-, gate, π)
        prob_neg = probs(zero_state(n) |> qcbm)

        dispatch!(+, gate, π/2) # set back

        grad_pos = kexpect(κ, prob, prob_pos) - kexpect(κ, prob, prob_neg)
        grad_neg = kexpect(κ, ptrain, prob_pos) - kexpect(κ, ptrain, prob_neg)
        grad[count] = grad_pos - grad_neg
        count += 1
    end
    return grad
end

# Now let's setup the training

using Flux.Optimise
qcbm = build_circuit(6, 10, [1=>2, 3=>4, 5=>6, 2=>3, 4=>5, 6=>1])
dispatch!(qcbm, :random) # initialize the parameters

κ = RBFKernel(0.25, 0:2^6-1)
pg = gaussian_pdf(1:1<<6, 1<<5-0.5, 1<<4);
opt = ADAM()

function train(qcbm, κ, opt, target)
    history = Float64[]
    for _ in 1:100
        push!(history, loss(κ, qcbm, target))
        ps = parameters(qcbm)
        Optimise.update!(opt, ps, gradient(qcbm, κ, target))
        popdispatch!(qcbm, ps)
    end
    return history
end

history = train(qcbm, κ, opt, pg)
trained_pg = probs(zero_state(nqubits(qcbm)) |> qcbm)

# The history of training looks like below

title!("training history")
xlabel!("steps"); ylabel!("loss")
plot(history)

# And let's check what we got

fig2 = plot(1:1<<6, trained_pg; label="trained")
plot!(fig2, 1:1<<6, pg; label="target")
title!("distribution")
xlabel!("x"); ylabel!("p")

# So within 50 steps, we got a pretty close estimation of our
# target distribution!


k = Knots([1,2,2,3])

@testset "Knots" begin
    @test Knots([1,2,3]) == Knots(1:3) == Knots(1,2,3)
    @test Knots(2,4,5) == Knots([2,4,5])
    @test zero(Knots) == Knots() == Knots([])
    @test Knots(1:3) == Knots([3,2,1])

    @test ♯(k) == length(k) == 4

    @test Knots([-1,2,3]) + 2 * Knots([2,5]) == Knots([-1,2,2,2,3,5,5])
    @test Knots([1,2,3]) + Knots([2,4,5]) == Knots([1,2,2,3,4,5])
    @test 2 * Knots([2,3]) == Knots([2,2,3,3])

    @test unique(k) == Knots([1,2,3])

    @testset "inclusive relation" begin
        @test Knots([1,2,3])     ⊆ Knots([1,2,3])
        @test Knots([1,2,3])     ⊇ Knots([1,2,3])
        @test Knots([1,2,2,3])   ⊆ Knots([1,2,2,3,5])
        @test Knots([1,2,2,3])   ⊈ Knots([1,2,3,5])
        @test Knots([1,2,2,3,5]) ⊇ Knots([1,2,2,3])
        @test Knots([1,2,3,5])   ⊉ Knots([1,2,2,3])
    end

    @test 𝔫(k, 0.3) == 0
    @test 𝔫(k, 1.0) == 1
    @test 𝔫(k, 2.0) == 2
    @test 1 ∈ k
    @test 1.5 ∉ k

    @testset "knotindex for usual case" begin
        @test BasicBSpline._knotindex(k,1.5) == 1
        @test BasicBSpline._knotindex(k,2.5) == 3
        @test BasicBSpline._knotindex₋₀(k,1.5) == 1
        @test BasicBSpline._knotindex₋₀(k,2.5) == 3
        @test BasicBSpline._knotindex₊₀(k,1.5) == 1
        @test BasicBSpline._knotindex₊₀(k,2.5) == 3
    end

    @testset "knotindex for corner case" begin
        @test BasicBSpline._knotindex(k,0) == 0
        @test BasicBSpline._knotindex(k,1) == 1
        @test BasicBSpline._knotindex(k,2) == 3
        @test BasicBSpline._knotindex(k,3) == 3
        @test BasicBSpline._knotindex(k,4) == 4

        @test BasicBSpline._knotindex₋₀(k,0) == 0
        @test BasicBSpline._knotindex₋₀(k,1) == 0
        @test BasicBSpline._knotindex₋₀(k,2) == 1
        @test BasicBSpline._knotindex₋₀(k,3) == 3
        @test BasicBSpline._knotindex₋₀(k,4) == 4

        @test BasicBSpline._knotindex₊₀(k,0) == 0
        @test BasicBSpline._knotindex₊₀(k,1) == 1
        @test BasicBSpline._knotindex₊₀(k,2) == 3
        @test BasicBSpline._knotindex₊₀(k,3) == 4
        @test BasicBSpline._knotindex₊₀(k,4) == 4
    end

    @test string(k) == "Knots([1.0, 2.0, 2.0, 3.0])"
    @test string(Knots()) == "Knots([])"
end


## path to TRPL folder --> set it before running the code!
path_to_trpl = "D:\\Projects\\TRPL" # some path to the TRPL folder (just an example here)

## activate package environment
cd(@__DIR__)
using Pkg
Pkg.activate("Project.toml")

## loading packages
using OrdinaryDiffEq, LinearAlgebra, StaticArrays, Random, NLopt, DelimitedFiles,
	  BlackBoxOptim, Printf, PoissonRandom, DataFrames, Plots, CSV

## backend quick plots
pyplot()
default(size = (800, 600))

## save output? (true or false)
out = false

## noise? (true or false)
noise = true

## experimental parameters (example for synthetic data set 1)
background = [1, 2, 12, 26, 18]
ND         = [1, 1,  2,  3,  4]
C_scaling  = 1e6

## parameters (example for synthetic data set 1)
# 		        τ_r0,  τ_nr0,   τ_d,  τ_0,   τ_1,  τ_l,     α,     β,    r,      C
parameters = [ 175.0,  350.0,  60.0, 20.0, 800.0, 250.0, 70.0, 110.0, 20.0, 1000.0]
eta0s      = [4.0e-5, 0.0004, 0.004, 0.04, 0.4]

## time (example for synthetic data set 1)
tspan = [0.0, 1750.0]
dt    = 0.25

## number of data set (example for synthetic data set 1)
idx_syn = 1

## rng seeds (to get the same noise when recalculating a data set: rng_seed = abs(rand(Int)))
rng_seeds = [4057480306083410532, 7737535138175179370, 3066106408424540935,
			 3829281771415197884, 1025260804200922037]

## useful variables
num_pars = 10
num_traj = length(eta0s)
dims =  num_pars + num_traj

## output name
name_syn_data = "syn_data_$(idx_syn)"

## definition rate equations
function rate_equations(du, u, p, t)
	du[1] = -((1 - u[2]) / p[3]) * u[1] - ((1 - u[3]) / p[4]) * u[1] + (1 / p[5]) * u[3] * 1 / p[7] -
			((1 + u[1]) / p[1]) * u[1] - ((1 + u[1]) / (p[2] * (1 + p[9] * u[1]))) * u[1]
	du[2] = ((1 - u[2]) / (p[3])) * p[8] * u[1]
	du[3] = ((1 - u[3]) / p[4]) * p[7] * u[1] - u[3] / p[5] - u[3] / p[6]
end

## ODEProblem definition with temporary u0 and parameters
u0 = [0.1, 0, 0]
t = [tspan[1]:dt:tspan[2];]
p_true = vcat(parameters,eta0s)
prob = ODEProblem(rate_equations, u0, tspan, p_true)

## function to generate the model for plotting
function generateipl(par)
    function prob_func(prob, i, repeat)
        remake(prob, u0 = [par[num_pars+i], 0, 0])
    end
    ensemble_simu = EnsembleProblem(prob, prob_func = prob_func)
    sol = Array(solve(ensemble_simu, Rosenbrock23(), p = par[1:num_pars-1], EnsembleThreads(),
		     		  trajectories = num_traj, saveat = t, save_idxs = [1],
		     		  reltol = 1e-6, abstol = 1e-8))[1,:,:]
    ipl = zeros(size(sol))
    for i = 1:num_traj
    	ipl[:,i] = ((C_scaling * par[num_pars] / (10 ^ ND[i])) .* sol[:,i] .* (sol[:,i] .+ 1)) .+ background[i]
    end
    return ipl
end

## get synthetic data with poisson noise if wanted
IPL = generateipl(p_true)
if noise
	IPL = pois_rand.(Random.seed!(rng_seeds[idx_syn]),IPL)
end

## get non-zero indicies
idx_IPL = zeros(Bool,size(IPL))
for i = 1:num_traj
	idx_IPL[:,i] = IPL[:,i] .> 0
end

## setup loss function to calculate loss goal
function loss(par)
	l = 0.0
	function prob_func(prob, i, repeat)
		remake(prob, u0 = [par[num_pars+i], 0, 0])
	end
	ensemble_simu = EnsembleProblem(prob, prob_func = prob_func)
	sol = Array(solve(ensemble_simu, Rosenbrock23(), p = par[1:num_pars-1], EnsembleThreads(),
					  trajectories = num_traj, saveat = t,save_idxs = [1],
					  reltol=1e-6,abstol=1e-8))[1,:,:]
	for i = 1:num_traj
		model = ((C_scaling * par[num_pars] / (10 ^ ND[i])) .* sol[:,i] .* (sol[:,i] .+ 1)) .+ background[i]
		idx_used = idx_IPL[:,i] .& (model .> 0)
		l += sum((model[idx_used]) .- (IPL[idx_used,i] .* log.(model[idx_used])))
	end
	return l
end

## calculate and print loss goal
l = loss(p_true)
println("loss goal = $l")

## plotting font size
fs = 16

## plot data
ylim = (10^0,10^(ceil(log10(maximum(IPL)))))
plt = plot(palette=palette(:seaborn_colorblind, num_traj),xlims=(t[1],t[end]),yscale=:log10,
		   ylims=ylim,title="Synthetic data set $idx_syn",titlefontsize=fs,grid=true,
		   xguidefontsize=fs,yguidefontsize=fs,dpi=300,framestyle=:box,xtickfontsize=fs,
		   ytickfontsize=fs,xlabe="t/ns",ylabel="Photon Counts",legendfontsize=fs-4,
		   leg=:topright,margin=10Plots.mm)
for i = num_traj:-1:1
	plot!(t, IPL[:,i],linealpha=0.75,linewidth=1.5,label=@sprintf("η₀ = %.2E", p_true[num_pars+i]))
end
display(plt)

## save output
if out
	# save input parameters
	io = open(joinpath(path_to_trpl,"Data\\Syn\\Params\\$name_syn_data.txt"), "w")
		println(io, "parameters")
		println(io, "τ_r0, τ_nr0, τ_2t, τ_1t,  τ_1e, τ_1d, α, β, r, C, η0s")
		println(io, p_true)
		println(io, "C scaling")
		println(io, C_scaling)
		println(io, "background")
		println(io, background)
		println(io, "ND")
		println(io, ND)
		println(io, "time span")
		println(io, tspan)
		println(io, "dt")
		println(io, dt)
		println(io, "loss goal")
		println(io, l)
		if noise
			println(io, "rng seed")
			println(io, rng_seed)
		end
	close(io)

	# save data for fit
	output = Array{Any}(undef,(size(IPL,1)+3,num_traj+1))
	output[1,2] = C_scaling
	output[2,2:end] = ND
	output[3,2:end] = background
	output[4:end,1] = t
	output[4:end,2:end] = IPL
	output[1:3,1] .= 0
	output[1,3:end] .= 0
	df = DataFrame(output,:auto)
	CSV.write(joinpath(path_to_trpl,"Data\\Syn\\CSV\\$name_syn_data.csv"), df)
end


### A Pluto.jl notebook ###
# v0.19.0

using Markdown
using InteractiveUtils

# ╔═╡ d18893e3-89ec-46fc-9265-0daac744caea
using Manopt, Manifolds, Colors, PlutoUI

# ╔═╡ 3b55c91e-bda9-11ec-0399-b77442d8ff45
md"""
# Illustration of the Gradient of a Second Order Difference

This example explains how to compute the gradient of the second order
difference mid point model using `adjoint_Jacobi_field`

This example also illustrates the `PowerManifold` manifold as well as `ArmijoLinesearch`.
"""

# ╔═╡ b2a115cc-5d4e-4052-afc1-2972ed493958
md"We define some colors from [Paul Tol](https://personal.sron.nl/~pault/)"

# ╔═╡ c19c8ede-ea81-4b10-85c2-f73051d8090a
begin
    black = RGBA{Float64}(colorant"#000000")
    TolVibrantBlue = RGBA{Float64}(colorant"#0077BB")
    TolVibrantOrange = RGBA{Float64}(colorant"#EE7733")
    TolVibrantMagenta = RGBA{Float64}(colorant"#EE3377")
    TolVibrantCyan = RGBA{Float64}(colorant"#33BBEE")
    TolVibrantTeal = RGBA{Float64}(colorant"#009988")
end;

# ╔═╡ b5ba2630-e247-46b9-b306-faf8a531f22a
begin
    T = [0:0.1:1.0...]
    #render asy yes/no. If not, images included w/ markdown are assumed to be prerendered
    render_asy = false
    localpath = join(splitpath(@__FILE__)[1:(end - 1)], "/") # remove file to get files folder
    image_prefix = localpath * "/second_order_difference"
    @info image_prefix
end;

# ╔═╡ 2a51bd75-c47e-452c-bfea-f6d7cfcd3e45
md"""
Assume we have two points ``p,q`` on the equator of the
[Sphere](https://juliamanifolds.github.io/Manifolds.jl/stable/manifolds/sphere.html) ``\mathcal M = \mathbb S^2``
and a point ``r`` near the north pole
"""

# ╔═╡ 3c8a6104-d380-4a32-907c-aaf205ea0bdd
begin
    M = Sphere(2)
    p = [1.0, 0.0, 0.0]
    q = [0.0, 1.0, 0.0]
    c = mid_point(M, p, q)
    r = shortest_geodesic(M, [0.0, 0.0, 1.0], c, 0.1)
end;

# ╔═╡ 9cfecd72-0c1a-4432-9628-54b71db9c22d
md"""
Now the second order absolute difference can be stated as [^BacakBergmannSteidlWeinmann2016]

```math
d_2(p_1,p:2,p_3) := \min_{c ∈ \mathcal C_{p_1,p_3}} d_{\mathcal M}(c,p_2),\qquad p_1,p_2,p_3∈\mathcal M,
```

where ``\mathcal C_{p,q}`` is the set of all mid points ``g(\frac{1}{2};p,q)``, between `p`and `q`, i.e. where ``g``
is a (not necessarily minimizing) geodesic connecting both.

For illustration we further define the point opposite of the mid point `c` defined above
"""

# ╔═╡ a9596c67-1cfb-4d94-a7c8-71bf36fe06bb
c2 = -c;

# ╔═╡ 2dfaebe0-e706-4eb6-b661-aed9d8ae19c1
md"To illustrate the second order difference let‘s look at the geodesic connecting ``r`` and the mid point ``c``"

# ╔═╡ 1c484736-bb94-4f64-bcfd-ad7ed52ade39
geoPts_rc = shortest_geodesic(M, r, c, T);

# ╔═╡ 80aa3913-7639-4752-a762-58c17a7e1c47
render_asy && begin
    asymptote_export_S2_signals(
        image_prefix * "/SecondOrderData.asy";
        curves=[geoPts_rc],
        points=[[p, r, q], [c, c2]],
        colors=Dict(:curves => [TolVibrantTeal], :points => [black, TolVibrantBlue]),
        dot_size=3.5,
        line_width=0.75,
        camera_position=(1.2, 1.0, 0.5), #src
    )
    render_asymptote(image_prefix * "/SecondOrderData.asy"; render=2)
end;

# ╔═╡ 717ba868-7da5-4ce1-9163-2b20f8cdc253
PlutoUI.LocalResource(image_prefix * "/SecondOrderData.png")

# ╔═╡ dcf6bbeb-5b09-4689-8bb9-e35f089cff45
md" **Figure.** _The three poins ``p, q``, and ``r`` (neart north pole, all black), which is connected by a geodesic (teal) to the mid point of the first two, ``c``, (blue). The length of this geodesic is the cost of the second order total variation._"

# ╔═╡ 07a44b4c-876c-4fff-ad72-2baedc06d090
md"""
Since we moved ``r`` 10% along the geodesic from the north pole to ``c``, the distance
to ``c`` is ``\frac{9\pi}{20}\approx 1.4137``, and this is also what the second order total variation cost, see [`costTV2`](https://manoptjl.org/stable/functions/costs.html#Manopt.costTV2), yields.
"""

# ╔═╡ ebab946c-ba27-422d-a191-56294e57177b
costTV2(M, (p, r, q))

# ╔═╡ db26c554-31ac-4d62-93a1-9440361fa821
md"""
But also its gradient can be
easily computed since it is just a distance with respect to ``r`` and a
concatenation of a geodesic, where the start or end point is the argument,
respectively, with a distance.
Hence the adjoint differentials
[`adjoint_differential_geodesic_startpoint`](https://manoptjl.org/stable/functions/adjointdifferentials.html#Manopt.adjoint_differential_geodesic_startpoint-Tuple{AbstractManifold,%20Any,%20Any,%20Any,%20Any}) and [`adjoint_differential_geodesic_endpoint`](https://manoptjl.org/stable/functions/adjointdifferentials.html#Manopt.adjoint_differential_geodesic_endpoint-Tuple{AbstractManifold,%20Any,%20Any,%20Any,%20Any}) can be employed.
The gradient is also directly implemented, see [`grad_TV2`](https://manoptjl.org/stable/functions/gradients.html#Manopt.grad_TV2).
we obtain
"""

# ╔═╡ b685cd13-8b06-4d6d-a001-576f68951bd9
(Xp, Xr, Xq) = grad_TV2(M, (p, r, q))

# ╔═╡ dddd7edf-47b7-4b6d-8b20-e49b4db211a9
render_asy && begin
    asymptote_export_S2_signals(
        image_prefix * "/SecondOrderGradient.asy";
        points=[[p, r, q], [c, c2]],
        tangent_vectors=[Tuple.([[p, -Xp], [r, -Xr], [q, -Xq]])],
        colors=Dict(:tvectors => [TolVibrantCyan], :points => [black, TolVibrantBlue]),
        dot_size=3.5,
        line_width=0.75,
        camera_position=(1.2, 1.0, 0.5),
    )
    render_asymptote(image_prefix * "/SecondOrderGradient.asy"; render=2)
end;

# ╔═╡ 18bc72b8-fef9-4890-bdb8-dc477e0109db
PlutoUI.LocalResource(image_prefix * "/SecondOrderGradient.png")

# ╔═╡ cdf26e32-2383-4ac3-b814-ef7c98860b55
md"""
**Figure.** _The negative gradient of the second order difference cost indicates the movement of the three points in order to reduce their cost._
"""

# ╔═╡ 77149fb5-f3cb-438f-843c-9221ae5c08a9
md"If we now perform a gradient step with constant step size 1, we obtain the three points"

# ╔═╡ e53622ca-3e08-4e16-a1a3-6536a46c8a9f
pn, rn, qn = exp.(Ref(M), [p, r, q], [-Xp, -Xr, -Xq])

# ╔═╡ 6b50fb13-45db-4cbb-a1e0-d62187c0b2d3
md"as well we the new mid point"

# ╔═╡ 0bbd56ce-9041-4974-8525-a0b5bd03e9fe
cn = mid_point(M, pn, qn)

# ╔═╡ 2eee7e24-20aa-4ae4-a336-b2eec90ce24c
md"Let‘s also again consider the geodesic connecting the new point ``r_n`` and the new mid point ``c_n`` as well as the gradient"

# ╔═╡ ce758b03-0177-45c2-b895-e64d133ca4af
begin
    geoPts_rncn = shortest_geodesic(M, rn, cn, T)
    (Xpn, Xrn, Xqn) = grad_TV2(M, (pn, rn, qn))
end;

# ╔═╡ 77e2409a-0168-4715-8770-23272e469f87
md"The new configuration of the three points looks as follows"

# ╔═╡ d90730e1-7111-42d9-8688-b4ff9cca9c6e
render_asy && begin
    asymptote_export_S2_signals(
        image_prefix * "/SecondOrderMin1.asy";
        points=[[p, r, q], [c, c2, cn], [pn, rn, qn]],
        curves=[geoPts_rncn],
        tangent_vectors=[
            Tuple.([[p, -Xp], [r, -Xr], [q, -Xq]]),
            Tuple.([[pn, -Xpn], [rn, -Xrn], [qn, -Xqn]]),
        ],
        colors=Dict(
            :tvectors => [TolVibrantCyan, TolVibrantMagenta],
            :points => [black, TolVibrantBlue, TolVibrantOrange],
            :curves => [TolVibrantTeal],
        ),
        dot_size=3.5,
        line_width=0.75,
        camera_position=(1.2, 1.0, 0.5),
    )
    render_asymptote(image_prefix * "/SecondOrderMin1.asy"; render=2)
end;

# ╔═╡ 533f324b-40ca-4b74-968e-e9f1ce4262eb
PlutoUI.LocalResource(image_prefix * "/SecondOrderMin1.png")

# ╔═╡ a959be34-c0fb-4dc3-b46d-1702f616a8bb
md"""
**Figure.** *The new situation of ``p_n, q_n``, and ``r_n`` (orange) and the miod point of the first two, ``c_n`` (blue), which is again connected by a geodesic to ``r_n``. Note that this geodesic is shorter, but also that ``c`` and ``r`` switched places*. The new gradient (magenta) is only slightly reduced in magnitude.*
"""

# ╔═╡ 5b4aa1ae-7ab1-4efc-a796-1fa4e771e2b2
md"""
One can see, that this step slightly “overshoots”, i.e. ``r`` is now even below ``c``.
and the cost function is still at
"""

# ╔═╡ ea3ceaad-d032-42fd-84ca-a08b2fdd6db0
costTV2(M, (pn, rn, qn))

# ╔═╡ 96600890-af72-4842-973e-438cb930f57b
md"""
But we can also search for the best step size using [`linesearch_backtrack`](https://manoptjl.org/stable/plans/index.html#Manopt.linesearch_backtrack-Union{Tuple{T},%20Tuple{TF},%20Tuple{AbstractManifold,%20TF,%20Any,%20T,%20Any,%20Any,%20Any},%20Tuple{AbstractManifold,%20TF,%20Any,%20T,%20Any,%20Any,%20Any,%20AbstractRetractionMethod},%20Tuple{AbstractManifold,%20TF,%20Any,%20T,%20Any,%20Any,%20Any,%20AbstractRetractionMethod,%20T},%20Tuple{AbstractManifold,%20TF,%20Any,%20T,%20Any,%20Any,%20Any,%20AbstractRetractionMethod,%20T,%20Any}}%20where%20{TF,%20T})
on the `PowerManifold` manifold ``\mathcal N = \mathcal M^3 = (\mathbb S^2)^3``
"""

# ╔═╡ 4b662610-0812-4ae2-809c-52f7a8979398
begin
    x = [p, r, q]
    N = PowerManifold(M, NestedPowerRepresentation(), 3)
    s = linesearch_backtrack(N, x -> costTV2(N, x), x, grad_TV2(N, x), 1.0, 0.957, 0.999)
end

# ╔═╡ 17ccdc00-35ad-4ec5-a480-6f6527e38834
md"and we obtain the new points"

# ╔═╡ db8916e6-0ee7-4e37-a879-5f35aa11d38d
begin
    pm, rm, qm = exp.(Ref(M), [p, r, q], s * [-Xp, -Xr, -Xq])
    cm = mid_point(M, pm, qm)
    geoPts_pmqm = shortest_geodesic(M, pm, qm, T)
end;

# ╔═╡ 78d44ce1-9d7b-4cb1-b9cb-2d402b1dfabb
md"we obtain"

# ╔═╡ fdf9aa5e-a25a-42f3-a74c-cef5877f2a1d
render_asy && begin
    asymptote_export_S2_signals(
        image_prefix * "/SecondOrderMin2.asy";
        points=[[p, r, q], [c, c2, cm], [pm, rm, qm]],
        curves=[geoPts_pmqm],
        tangent_vectors=[Tuple.([[p, -Xp], [r, -Xr], [q, -Xq]])],
        colors=Dict(
            :tvectors => [TolVibrantCyan],
            :points => [black, TolVibrantBlue, TolVibrantOrange],
            :curves => [TolVibrantTeal],
        ),
        dot_size=3.5,
        line_width=0.75,
        camera_position=(1.2, 1.0, 0.5),
    )
    render_asymptote(image_prefix * "/SecondOrderMin2.asy"; render=2)
end;

# ╔═╡ b0a75ff9-9ff7-4db6-ad82-76f183663dc0
PlutoUI.LocalResource(image_prefix * "/SecondOrderMin2.png")

# ╔═╡ 9c37f43b-dc3e-45fb-a99f-756c83f2f658
md"""
**Figure.** *For the best step size found by line search, ``r_m`` and ``c_m`` nearly agree, i.e. ``r_m`` lies on the geodesic between ``p_m`` and ``q_m`` as the geodesic drawn here indicates.*
"""

# ╔═╡ 428d1422-7521-42fd-993d-e617115edf14
md"Here, the cost function yields"

# ╔═╡ 25804825-c9f6-4cfa-b209-918ad19d2ae1
costTV2(M, (pm, rm, qm))

# ╔═╡ 718f744e-e6af-4158-9418-7d2f4b757903
md"which is much closer to zero, as one can also see, since the new center ``c_m`` and ``r_m`` are quite close."

# ╔═╡ 962936da-c154-444f-8d48-ca1176dee1a1
md"""
## Literature

[^BacakBergmannSteidlWeinmann2016]:
    > Bačák, M., Bergmann, R., Steidl, G. and Weinmann, A.: _A second order nonsmooth
    > variational model for restoring manifold-valued images_,
    > SIAM Journal on Scientific Computations, Volume 38, Number 1, pp. A567–597,
    > doi: [10.1137/15M101988X](https://doi.org/10.1137/15M101988X),
    > arXiv: [1506.02409](https://arxiv.org/abs/1506.02409)
"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
Colors = "~0.12.8"
Manifolds = "~0.7.8"
Manopt = "~0.3.23"
PlutoUI = "~0.7.38"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.2"
manifest_format = "2.0"

[[deps.AbstractFFTs]]
deps = ["ChainRulesCore", "LinearAlgebra"]
git-tree-sha1 = "6f1d9bc1c08f9f4a8fa92e3ea3cb50153a1b40d4"
uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
version = "1.1.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.Adapt]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "af92965fb30777147966f58acb05da51c5616b5f"
uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
version = "3.3.3"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.ArnoldiMethod]]
deps = ["LinearAlgebra", "Random", "StaticArrays"]
git-tree-sha1 = "62e51b39331de8911e4a7ff6f5aaf38a5f4cc0ae"
uuid = "ec485272-7323-5ecc-a04f-4719b315124d"
version = "0.2.0"

[[deps.ArrayInterface]]
deps = ["Compat", "IfElse", "LinearAlgebra", "Requires", "SparseArrays", "Static"]
git-tree-sha1 = "c933ce606f6535a7c7b98e1d86d5d1014f730596"
uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
version = "5.0.7"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "9950387274246d08af38f6eef8cb5480862a435f"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.14.0"

[[deps.ChangesOfVariables]]
deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
git-tree-sha1 = "bf98fa45a0a4cee295de98d4c1462be26345b9a1"
uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
version = "0.1.2"

[[deps.ColorSchemes]]
deps = ["ColorTypes", "Colors", "FixedPointNumbers", "Random"]
git-tree-sha1 = "12fc73e5e0af68ad3137b886e3f7c1eacfca2640"
uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
version = "3.17.1"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.Colors]]
deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
git-tree-sha1 = "417b0ed7b8b838aa6ca0a87aadf1bb9eb111ce40"
uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
version = "0.12.8"

[[deps.Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "b153278a25dd42c65abbf4e62344f9d22e59191b"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.43.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.CovarianceEstimation]]
deps = ["LinearAlgebra", "Statistics", "StatsBase"]
git-tree-sha1 = "a3e070133acab996660d31dcf479ea42849e368f"
uuid = "587fd27a-f159-11e8-2dae-1979310e6154"
version = "0.2.7"

[[deps.DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[deps.DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "3daef5523dd2e769dad2365274f760ff5f282c7d"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.11"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.DensityInterface]]
deps = ["InverseFunctions", "Test"]
git-tree-sha1 = "80c3e8639e3353e5d2912fb3a1916b8455e2494b"
uuid = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
version = "0.4.0"

[[deps.Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[deps.Distributions]]
deps = ["ChainRulesCore", "DensityInterface", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SparseArrays", "SpecialFunctions", "Statistics", "StatsBase", "StatsFuns", "Test"]
git-tree-sha1 = "5a4168170ede913a2cd679e53c2123cb4b889795"
uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
version = "0.25.53"

[[deps.DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "b19534d1895d702889b219c382a6e18010797f0b"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.6"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.Einsum]]
deps = ["Compat"]
git-tree-sha1 = "4a6b3eee0161c89700b6c1949feae8b851da5494"
uuid = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
version = "0.4.1"

[[deps.FillArrays]]
deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"]
git-tree-sha1 = "246621d23d1f43e3b9c368bf3b72b2331a27c286"
uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
version = "0.13.2"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Graphs]]
deps = ["ArnoldiMethod", "Compat", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"]
git-tree-sha1 = "57c021de207e234108a6f1454003120a1bf350c4"
uuid = "86223c79-3864-5bf0-83f7-82e725a168b6"
version = "1.6.0"

[[deps.HybridArrays]]
deps = ["LinearAlgebra", "Requires", "StaticArrays"]
git-tree-sha1 = "eb6b23460f5544c5d09efae0818b86736cefcd3d"
uuid = "1baab800-613f-4b0a-84e4-9cd3431bfbb9"
version = "0.4.10"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.IfElse]]
git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1"
uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
version = "0.1.1"

[[deps.Inflate]]
git-tree-sha1 = "f5fc07d4e706b84f72d54eedcc1c13d92fb0871c"
uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9"
version = "0.1.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.InverseFunctions]]
deps = ["Test"]
git-tree-sha1 = "91b5dcf362c5add98049e6c29ee756910b03051d"
uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
version = "0.1.3"

[[deps.IrrationalConstants]]
git-tree-sha1 = "7fd44fd4ff43fc60815f8e764c0f352b83c49151"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.1"

[[deps.JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "abc9885a7ca2052a736a600f7fa66209f96506e1"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.4.1"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "3c837543ddb02250ef42f4738347454f95079d4e"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.3"

[[deps.Kronecker]]
deps = ["LinearAlgebra", "NamedDims", "SparseArrays", "StatsBase"]
git-tree-sha1 = "a51f46415c844dee694cb8b20a3fcbe6dba342c2"
uuid = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
version = "0.5.1"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.LogExpFunctions]]
deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "a970d55c2ad8084ca317a4658ba6ce99b7523571"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.12"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "3d3e902b31198a27340d0bf00d6ac452866021cf"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.9"

[[deps.Manifolds]]
deps = ["Colors", "Distributions", "Einsum", "Graphs", "HybridArrays", "Kronecker", "LinearAlgebra", "ManifoldsBase", "Markdown", "Random", "RecipesBase", "RecursiveArrayTools", "Requires", "SimpleWeightedGraphs", "SpecialFunctions", "StaticArrays", "Statistics", "StatsBase"]
git-tree-sha1 = "cc5e34ecd061bf9bedd62a8b8a28abe96c3b7c6e"
uuid = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
version = "0.7.8"

[[deps.ManifoldsBase]]
deps = ["LinearAlgebra", "Markdown"]
git-tree-sha1 = "b8518cbf0f06f1375c95479759b266bd75d89ad9"
uuid = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb"
version = "0.12.13"

[[deps.Manopt]]
deps = ["ColorSchemes", "ColorTypes", "Colors", "DataStructures", "Dates", "LinearAlgebra", "ManifoldsBase", "Markdown", "Printf", "Random", "Requires", "SparseArrays", "StaticArrays", "Statistics", "Test"]
git-tree-sha1 = "a8a292a422b0a6924f46d4e022fc8263721120c5"
uuid = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
version = "0.3.23"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "bf210ce90b6c9eed32d25dbcae1ebc565df2687f"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.2"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NamedDims]]
deps = ["AbstractFFTs", "ChainRulesCore", "CovarianceEstimation", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
git-tree-sha1 = "0856b62716585eb90cc1dada226ac9eab5f69aa5"
uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
version = "0.2.47"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.OpenLibm_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "05823500-19ac-5b8b-9628-191a04bc5112"

[[deps.OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.PDMats]]
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
git-tree-sha1 = "e8185b83b9fc56eb6456200e873ce598ebc7f262"
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
version = "0.11.7"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "621f4f3b4977325b9128d5fae7a8b4829a0c2222"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.2.4"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "ColorTypes", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "670e559e5c8e191ded66fa9ea89c97f10376bb4c"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.38"

[[deps.Preferences]]
deps = ["TOML"]
git-tree-sha1 = "d3538e7f8a790dc8903519090857ef8e1283eecd"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.5"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.QuadGK]]
deps = ["DataStructures", "LinearAlgebra"]
git-tree-sha1 = "78aadffb3efd2155af139781b8a8df1ef279ea39"
uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
version = "2.4.2"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.RecipesBase]]
git-tree-sha1 = "6bf3f380ff52ce0832ddd3a2a7b9538ed1bcca7d"
uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
version = "1.2.1"

[[deps.RecursiveArrayTools]]
deps = ["Adapt", "ArrayInterface", "ChainRulesCore", "DocStringExtensions", "FillArrays", "LinearAlgebra", "RecipesBase", "Requires", "StaticArrays", "Statistics", "ZygoteRules"]
git-tree-sha1 = "bfe14f127f3e7def02a6c2b1940b39d0dabaa3ef"
uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
version = "2.26.3"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.3.0"

[[deps.Rmath]]
deps = ["Random", "Rmath_jll"]
git-tree-sha1 = "bf3188feca147ce108c76ad82c2792c57abe7b1f"
uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa"
version = "0.7.0"

[[deps.Rmath_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "68db32dff12bb6127bac73c209881191bf0efbb7"
uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
version = "0.3.0+0"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[deps.SimpleTraits]]
deps = ["InteractiveUtils", "MacroTools"]
git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231"
uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
version = "0.9.4"

[[deps.SimpleWeightedGraphs]]
deps = ["Graphs", "LinearAlgebra", "Markdown", "SparseArrays", "Test"]
git-tree-sha1 = "a6f404cc44d3d3b28c793ec0eb59af709d827e4e"
uuid = "47aef6b3-ad0c-573a-a1e2-d07658019622"
version = "1.2.1"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "b3363d7460f7d098ca0912c69b082f75625d7508"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.0.1"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.SpecialFunctions]]
deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "5ba658aeecaaf96923dce0da9e703bd1fe7666f9"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "2.1.4"

[[deps.Static]]
deps = ["IfElse"]
git-tree-sha1 = "87e9954dfa33fd145694e42337bdd3d5b07021a6"
uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
version = "0.6.0"

[[deps.StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "4f6ec5d99a28e1a749559ef7dd518663c5eca3d5"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "1.4.3"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.StatsAPI]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "8d7530a38dbd2c397be7ddd01a424e4f411dcc41"
uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
version = "1.2.2"

[[deps.StatsBase]]
deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
git-tree-sha1 = "8977b17906b0a1cc74ab2e3a05faa16cf08a8291"
uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
version = "0.33.16"

[[deps.StatsFuns]]
deps = ["ChainRulesCore", "InverseFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"]
git-tree-sha1 = "5950925ff997ed6fb3e985dcce8eb1ba42a0bbe7"
uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
version = "0.9.18"

[[deps.SuiteSparse]]
deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.ZygoteRules]]
deps = ["MacroTools"]
git-tree-sha1 = "8c1a8e4dfacb1fd631745552c8db35d0deb09ea0"
uuid = "700de1a5-db45-46bc-99cf-38207098b444"
version = "0.2.2"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─3b55c91e-bda9-11ec-0399-b77442d8ff45
# ╠═d18893e3-89ec-46fc-9265-0daac744caea
# ╟─b2a115cc-5d4e-4052-afc1-2972ed493958
# ╠═c19c8ede-ea81-4b10-85c2-f73051d8090a
# ╠═b5ba2630-e247-46b9-b306-faf8a531f22a
# ╟─2a51bd75-c47e-452c-bfea-f6d7cfcd3e45
# ╠═3c8a6104-d380-4a32-907c-aaf205ea0bdd
# ╟─9cfecd72-0c1a-4432-9628-54b71db9c22d
# ╠═a9596c67-1cfb-4d94-a7c8-71bf36fe06bb
# ╟─2dfaebe0-e706-4eb6-b661-aed9d8ae19c1
# ╠═1c484736-bb94-4f64-bcfd-ad7ed52ade39
# ╠═80aa3913-7639-4752-a762-58c17a7e1c47
# ╟─717ba868-7da5-4ce1-9163-2b20f8cdc253
# ╟─dcf6bbeb-5b09-4689-8bb9-e35f089cff45
# ╟─07a44b4c-876c-4fff-ad72-2baedc06d090
# ╠═ebab946c-ba27-422d-a191-56294e57177b
# ╟─db26c554-31ac-4d62-93a1-9440361fa821
# ╠═b685cd13-8b06-4d6d-a001-576f68951bd9
# ╠═dddd7edf-47b7-4b6d-8b20-e49b4db211a9
# ╟─18bc72b8-fef9-4890-bdb8-dc477e0109db
# ╟─cdf26e32-2383-4ac3-b814-ef7c98860b55
# ╟─77149fb5-f3cb-438f-843c-9221ae5c08a9
# ╠═e53622ca-3e08-4e16-a1a3-6536a46c8a9f
# ╟─6b50fb13-45db-4cbb-a1e0-d62187c0b2d3
# ╟─0bbd56ce-9041-4974-8525-a0b5bd03e9fe
# ╠═2eee7e24-20aa-4ae4-a336-b2eec90ce24c
# ╠═ce758b03-0177-45c2-b895-e64d133ca4af
# ╟─77e2409a-0168-4715-8770-23272e469f87
# ╠═d90730e1-7111-42d9-8688-b4ff9cca9c6e
# ╠═533f324b-40ca-4b74-968e-e9f1ce4262eb
# ╟─a959be34-c0fb-4dc3-b46d-1702f616a8bb
# ╟─5b4aa1ae-7ab1-4efc-a796-1fa4e771e2b2
# ╠═ea3ceaad-d032-42fd-84ca-a08b2fdd6db0
# ╟─96600890-af72-4842-973e-438cb930f57b
# ╠═4b662610-0812-4ae2-809c-52f7a8979398
# ╟─17ccdc00-35ad-4ec5-a480-6f6527e38834
# ╠═db8916e6-0ee7-4e37-a879-5f35aa11d38d
# ╟─78d44ce1-9d7b-4cb1-b9cb-2d402b1dfabb
# ╠═fdf9aa5e-a25a-42f3-a74c-cef5877f2a1d
# ╠═b0a75ff9-9ff7-4db6-ad82-76f183663dc0
# ╟─9c37f43b-dc3e-45fb-a99f-756c83f2f658
# ╟─428d1422-7521-42fd-993d-e617115edf14
# ╠═25804825-c9f6-4cfa-b209-918ad19d2ae1
# ╟─718f744e-e6af-4158-9418-7d2f4b757903
# ╟─962936da-c154-444f-8d48-ca1176dee1a1
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


module TestReturns

using Base.Test
using TimeData
using EconDatasets
using DataFrames
using Dates
using DataArrays
using GLM

include(joinpath(Pkg.dir("Econometrics"), "src/Econometrics.jl"))

###################
## get test data ##
###################

fname = joinpath(Pkg.dir("Econometrics"), "test/data/no_git_intRates.csv")
intData = readTimedata(fname)

#########################
## test type interface ##
#########################

# constructor using tuple
origParams = (0.4, 3.2, 0.2)
cirMod = Econometrics.CIR(origParams...)
params = Econometrics.getParams(cirMod)
@test params == origParams

origParams = (0.4, 3.2, 0.2)
cirMod = Econometrics.CIR(origParams..., Dates.Day(1))
params = Econometrics.getParams(cirMod)
@test params == origParams
@test cirMod.scale == Dates.Day(1)

# constructor using array
origParamsArr = [origParams...]
cirMod = Econometrics.CIR(origParamsArr)
params = Econometrics.getParams(cirMod)
@test params == origParams

##############################
## test likelihood function ##
##############################

## load matlab results
res_llh_fname = joinpath(Pkg.dir("Econometrics"), "test/data/matlab_cir_llhs.csv")
llhs = readcsv(res_llh_fname)

## reproduce values in Julia
##--------------------------

data = intData.vals[1] |>
         x -> x[!isna(x)]

## testcase 1
params = (0.4, 3.2, 0.2)
cirMod = Econometrics.CIR([params...])

jl_llh = Econometrics.cirNllhx(data, cirMod, 1/250)
@test_approx_eq_eps jl_llh llhs[1] 0.08

jl_llh = Econometrics.cirNllhx(data, [params...], 1/250)
@test_approx_eq_eps jl_llh llhs[1] 0.08

## testcase 2
params = (0.8, 3.0, 0.8)
cirMod = Econometrics.CIR([params...])

jl_llh = Econometrics.cirNllhx(data, cirMod, 1/250)
@test_approx_eq_eps jl_llh llhs[2] 0.08

jl_llh = Econometrics.cirNllhx(data, [params...], 1/250)
@test_approx_eq_eps jl_llh llhs[2] 0.08

## testcase 3
params = (1.2, 2.0, 0.4)
cirMod = Econometrics.CIR([params...])

jl_llh = Econometrics.cirNllhx(data, cirMod, 1/250)
@test_approx_eq_eps jl_llh llhs[3] 0.08

## test with missing data - not reproducible in matlab
##----------------------------------------------------

# with missing data - no values for comparison available
Econometrics.cirNllhx(intData.vals[1], cirMod, 1/250)
Econometrics.cirNllhx(intData.vals[1], [params...], 1/250)

####################################
## test conditional distributions ##
####################################

params = (0.4, 3.2, 0.2)
cirMod = Econometrics.CIR([params...])

u = 0.5
r0 = 4.5

# cdf(icdf(u)) == u
kk = Econometrics.quantile(cirMod, u, r0, 1/250)
@test_approx_eq_eps u Econometrics.cdf(cirMod, kk, r0, 1/250) 1e-10

# get R results and compare them to actual Julia test values
resultsFname = joinpath(Pkg.dir("Econometrics"), "test/data/r_cir_results.csv")
rRes = readtable(resultsFname)

## compare cdf values to R
cdf1 = Econometrics.cdf(cirMod, 4.5, r0, 1/250)
@test_approx_eq_eps cdf1 rRes[1, 1] 1e-10

cdf2 = Econometrics.cdf(cirMod, 5.2, r0, 1.)
@test_approx_eq_eps cdf2 rRes[1, 2] 1e-10

## compare inverse cdf values to R
icdf1 = Econometrics.quantile(cirMod, 0.5, r0, 1/250)
@test_approx_eq_eps icdf1 rRes[1, 3] 1e-10

icdf2 = Econometrics.quantile(cirMod, 0.5, r0, 1.)
@test_approx_eq_eps icdf2 rRes[1, 4] 1e-10

## cdf for array data
expOut = [Econometrics.cdf(cirMod, 4.4, 4.5, 1/250);
          Econometrics.cdf(cirMod, 4.3, 4.4, 1/250)]
actOut = Econometrics.cdf(cirMod, [4.5; 4.4; 4.3], 1/250)
@test expOut == actOut

################
## simulation ##
################

include(joinpath(Pkg.dir("Econometrics"), "src/Econometrics.jl"))
params = (0.4, 3.2, 0.2)
cirMod = Econometrics.CIR([params...])
vals = Econometrics.simulate(cirMod, 4.5, 100, 1/250)

##############
## test OLS ##
##############

# get matlab results and compare them to actual Julia test values
resultsFname = joinpath(Pkg.dir("Econometrics"), "test/data/r_results.RData")
## rRes = read_rda(resultsFname)

## test ols
##---------

## expParams = Float64[rRes["olsParams"].data[ii].data[1] for ii=1:3]
## actParams = Econometrics.cirOls(intData.vals[:, 1], 1/250)

## @test_approx_eq expParams[1] actParams[1]
## @test_approx_eq expParams[2] actParams[2]
## @test_approx_eq expParams[3] actParams[3]

end


using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
using DynamicalSystems, PyPlot, Statistics
include(srcdir("style.jl"))

# Plot dependence on time and i.c.
function plot_return_dim(ds, εs, ic, e; Tmax = 20_000_000, ici=1)
    fig, axs = subplots(2, 1; sharex = true)
    mcts = []
    for (i, T) ∈ enumerate(ic)
        println("ic $i ...")
        u0 = trajectory(ds, T; dt = T, Ttr = 100)[end]
        n, c = mean_return_times(ds, u0, εs, Tmax)
        j = findfirst(isequal(0), c)
        j = isnothing(j) ? length(c) : j-1
        push!(mcts, n[1:j])
        o, d = linreg(log.(εs[1:j]), log.(n[1:j]))
        axs[1].plot(log.(εs), log.(n);
                    marker = "o", ls="-", lw = 1.0, ms = 5, color = "C$(i)",
                    label = "i.c. $(i), \$d = $(rdspl(abs(d)))\$"
        )
    end
    for (i, T) ∈ enumerate(e)
        println("Time $i ...")
        u0 = trajectory(ds, ic[ici]; Ttr = 100)[end]
        n, c = mean_return_times(ds, u0, εs, 10^T)
        j = findfirst(isequal(0), c)
        j = isnothing(j) ? length(c) : j-1
        o, d = linreg(log.(εs[1:j]), log.(n[1:j]))
        axs[2].plot(log.(εs), log.(n);
                    marker = "o", ls="-", lw = 1.0, ms = 5, color = lighten_color("C0", 1.2i),
                    label = "\$T=10^$(T), d = $(rdspl(abs(d)))\$"
        )
    end

    minlength = minimum(length.(mcts))
    average_mct = [mean(mcts[j][i] for j in 1:length(ic)) for i in 1:minlength]
    x, y = log.(εs[1:minlength]), log.(average_mct)
    _, d = linreg(x, -y)
    # _, d = linear_region(x, -y) # this doesn't work well. The line is too jagged.
    axs[1].plot(x, y; ls = "dashed", label = "avg., \$d=$(rdspl(d))\$", lw = 3.0)

    axs[1].legend(); axs[2].legend()
    axs[1].set_ylabel("log(⟨t⟩)")
    axs[2].set_ylabel("log(⟨t⟩)")
    axs[2].set_xlabel("log(ε)")
    fig.tight_layout()
    return fig, axs
end

# %% Henon map
# first plot random initial conditions
ds = Systems.henon()
ic = (500, 1000, 1500, 2000, 2500)
e = (3, 5, 7)
εs = sort!(ℯ .^ (-9:0.5:-3); rev = true)
fig, axs = plot_return_dim(ds, εs, ic, e)
axs[1].set_title("Henon map")

tr = trajectory(ds, 50000; Ttr = 100)
figure()
scatter(columns(tr)..., s = 2)
for (i, T) ∈ enumerate(ic)
    u0 = trajectory(ds, T; Ttr = 100)[end]
    scatter(u0[1], u0[2], color = "C$i")
end

# wsave(plotsdir("returntime", "henon"), fig)

# %% Towel map
# first plot random initial conditions
ds = Systems.towel()
εs = sort!(ℯ .^ (-7:0.5:-3); rev = true)
ic = (500, 1000, 1500)
e = (4, 6, 8)

tr = trajectory(ds, 50000; Ttr = 100)
using3D(); figure();
scatter3D(columns(tr)...; s = 2)
for (i, T) ∈ enumerate(ic)
    u0 = trajectory(ds, T; Ttr = 100)[end]
    scatter3D([u0[1]], [u0[2]], [u0[3]]; color = "C$i", s = 1000)
end

fig, axs = plot_return_dim(ds, εs, ic, e)
axs[1].set_title("Towel map")
# wsave(plotsdir("returntime", "towel"), fig)

# %% Roessler system
# first plot random initial conditions
ds = Systems.roessler([1,1,1.0])
εs = sort!(ℯ .^ (-4:0.5:-1); rev = true)
ic = (563.0, 1300.0, 1000.0, 1500.0, 222.0)
e = (4, 5, 6)

tr = trajectory(ds, 1000; Ttr = 100)
using3D(); figure();
plot3D(columns(tr)...; lw = 0.5)
for (i, T) ∈ enumerate(ic)
    u0 = trajectory(ds, T; Ttr = 100)[end]
    scatter3D([u0[1]], [u0[2]], [u0[3]]; color = "C$i", s = 1000)
end

# fig, axs = plot_return_dim(ds, εs, ic, e; Tmax = 10^6, ici=2)
# axs[1].set_title("Roessler system")
# wsave(plotsdir("returntime", "roessler"), fig)

# %% Roessler system
# first plot random initial conditions
ds = Systems.henonheiles([0.0, -0.25, 0.42081, 0.0])
εs = sort!(ℯ .^ (-4:0.5:-1); rev = true)
ic = (563.0, 1300.0, 1002.0, 1500.5, 222.0)
e = (4, 5, 6)

tr = trajectory(ds, 1000; Ttr = 100)
using3D(); figure();
x, y, z = columns(tr)
plot3D(x, y, z; lw = 0.5)
for (i, T) ∈ enumerate(ic)
    u0 = trajectory(ds, T; Ttr = 100)[end]
    scatter3D([u0[1]], [u0[2]], [u0[3]]; color = "C$i", s = 1000)
end

# fig, axs = plot_return_dim(ds, εs, ic, e; Tmax = 10^6, ici=2)
# axs[1].set_title("Roessler system")
# wsave(plotsdir("returntime", "roessler"), fig)

# %% Lorenz system
# first plot random initial conditions
ds = Systems.lorenz([1,10.0,1.0])
εs = sort!(ℯ .^ (-5:0.5:-0); rev = true)
ic = (563, 1000, 1500.5)
e = (4, 5, 6)

tr = trajectory(ds, 500; Ttr = 100, dt = 0.02)
using3D(); figure();
plot3D(columns(tr)...; lw = 1.0)
for (i, T) ∈ enumerate(ic)
    u0 = trajectory(ds, T; Ttr = 100)[end]
    scatter3D([u0[1]], [u0[2]], [u0[3]]; color = "C$i", s = 1000)
end

fig, axs = plot_return_dim(ds, εs, ic, e; Tmax = 10^5, ici=2)
axs[1].set_title("Lorenz system")
# wsave(plotsdir("returntime", "lorenz"), fig)

# %% Lorenz-96 system
D = 4
ds = Systems.lorenz96(D, range(0; length = D, step = 0.1); F = 24.0)
εs = sort!(ℯ .^ (-5:0.5:-0); rev = true)
ic = (26.0, 563.0, 876.8, 1000.0)
e = (4, 5, 6)

# TODO: Must increase max return time.

tr = trajectory(ds, 50; Ttr = 100.0, dt = 0.01)
using3D(); figure();
x, y, z = columns(tr)
plot3D(x, y, z; lw = 0.5)
for (i, T) ∈ enumerate(ic)
    u0 = trajectory(ds, T; Ttr = 100)[end]
    scatter3D([u0[1]], [u0[2]], [u0[3]]; color = "C$i", s = 1000)
end

fig, axs = plot_return_dim(ds, εs, ic, e; Tmax = 10^5, ici=2)
axs[1].set_title("Lorenz96, D = $D")
wsave(plotsdir("returntime", "lorenz96_D=$D"), fig)


module SimplexFacesTests

using Test
using STLCutters
using Gridap
using Gridap.ReferenceFEs

using STLCutters: center
using STLCutters: distance
using STLCutters: projection
using STLCutters: measure
using STLCutters: normal
using STLCutters: signed_distance
using STLCutters: contains_projection
using STLCutters: voxel_intersection
using STLCutters: simplex_face
using STLCutters: min_height
using STLCutters: Plane

# Face{1,2}

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
f = simplex_face(p1,p2)

@test num_dims(f) == 1
@test num_point_dims(f) == 2
@test center(f) == Point(0.5,0.5)
@test measure(f) ≈ √2
@test normal(f) ≈ (√2)/2 * Point(1,1)
@test normal(f,1) ≈ (√2)/2 * Point(-1,1)
@test normal(f,2) ≈ (√2)/2 * Point(1,-1)

pmin,pmax = get_bounding_box(f)

@test measure(pmin,pmax) == 1

# Face{1,3}

p1 = Point(0.0,1.0,0.0)
p2 = Point(1.0,0.0,1.0)
f = simplex_face(p1,p2)

@test num_dims(f) == 1
@test num_point_dims(f) == 3
@test center(f) == Point(0.5,0.5,0.5)
@test measure(f) ≈ √3
@test normal(f,1) ≈ 1/(√3) * Point(-1,1,-1)
@test normal(f,2) ≈ 1/(√3) * Point(1,-1,1)

pmin,pmax = get_bounding_box(f)

@test measure(pmin,pmax) == 1

# Face{2,2}

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
p3 = Point(0.0,0.0)
f = simplex_face(p1,p2,p3)

@test num_dims(f) == 2
@test num_point_dims(f) == 2
@test center(f) ≈ Point(1/3,1/3)
@test measure(f) ≈ 1/2
@test min_height(f) ≈ 1/√2

pmin,pmax = get_bounding_box(f)

@test measure(pmin,pmax) == 1

# Face{2,3}

p1 = Point(0.0,1.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3)

@test center(f) ≈ Point(1/3,1/3,1/3)
@test measure(f) ≈ (√3)/2

@test num_dims(f) == 2
@test num_point_dims(f) == 3
@test normal(f) ≈ 1/√3 * Point(-1,-1,-1)
@test min_height(f) ≈ √(3/2)

pmin,pmax = get_bounding_box(f)

@test measure(pmin,pmax) == 1

# Face{3,3}

p1 = Point(0.0,0.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,1.0,0.0)
p4 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3,p4)

@test center(f) ≈ Point(1/4,1/4,1/4)
@test measure(f) ≈ 1/6
@test min_height(f) ≈ 1/√3

pmin,pmax = get_bounding_box(f)

@test measure(pmin,pmax) == 1

# Plane{2}

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
f = simplex_face(p1,p2)

Π = Plane(f)
p = Point(0.0,0.0)

@test signed_distance(p,Π) ≈ -1/√2

# Plane{3}

p1 = Point(0.0,1.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3)

Π = Plane(f)
p = Point(0.0,0.0,0.0)

@test signed_distance(p,Π) ≈ 1/√3

# Distance Point - Point

p1 = Point(0.0,0.0)
p2 = Point(1.0,1.0)

distance(p1,p2) ≈ √2

p1 = Point(0.0,0.0,0.0)
p2 = Point(1.0,1.0,1.0)

distance(p1,p2) ≈ √3

# Distance Point - Face{1}

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
f = simplex_face(p1,p2)

a = Point(0.0,0.0)
b = Point(2.0,0.0)
c = Point(0.5,0.5)

@test distance(a,f) ≈ 1/√(2)
@test distance(b,f) ≈ 1.0
@test abs( distance(c,f) ) < 1e-9

# Distance Point - Face{2}

p1 = Point(0.0,1.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3)

a = Point(0.0,0.0,0.0)
b = Point(2.0,0.0,0.0)
c = Point(1/3,1/3,1/3)

@test distance(a,f) ≈ 1/√(3)
@test distance(b,f) ≈ 1.0
@test abs( distance(c,f) ) < 1e-9

# Distance Point - Face{3}

p1 = Point(0.0,0.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,1.0,0.0)
p4 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3,p4)

a = Point(0.0,0.0,0.0)
b = Point(2.0,0.0,0.0)
c = Point(1/4,1/4,1/4)
d = Point(1.0,1.0,1.0)

@test abs(distance(a,f)) < 1e-9
@test distance(b,f) ≈ 1.0
@test distance(c,f) == 0
@test distance(d,f) ≈ 2/√3

# Projection Point - Point

a = Point(0.0,1.0)
b = Point(1.0,0.0)

@test projection(a,b) == b

# Projection Point - Face{1}

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
f = simplex_face(p1,p2)

a = Point(0.0,0.0)
b = Point(0.5,0.0)

@test projection(a,f) ≈ Point(0.5,0.5)
@test projection(b,f) ≈ Point(0.75,0.25)

# Projection Point - Face{2}

p1 = Point(0.0,1.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3)

a = Point(0.0,0.0,0.0)
@test projection(a,f) ≈ Point(1/3,1/3,1/3)

# Contains projection Point - Face{1}

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
f = simplex_face(p1,p2)

a = Point(0.0,0.0)
b = Point(2.0,0.0)

@test contains_projection(f,a)
@test !contains_projection(f,b)

# Contains projection Point - Face{2}

p1 = Point(0.0,1.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3)

a = Point(0.0,0.0,0.0)
b = Point(2.0,0.0,0.0)

@test contains_projection(f,a)
@test !contains_projection(f,b)

# Contains projection Point - Face{3}

p1 = Point(0.0,0.0,0.0)
p2 = Point(1.0,0.0,0.0)
p3 = Point(0.0,1.0,0.0)
p4 = Point(0.0,0.0,1.0)
f = simplex_face(p1,p2,p3,p4)

a = Point(1/4,1/4,1/4)
b = Point(1.0,1.0,1.0)

@test contains_projection(f,a)
@test !contains_projection(f,b)

# Voxel intersection Face{2}

p = QUAD

pmin,pmax = Point(0.0,0.0),Point(1.0,1.0)
pmin -= 0.1
pmax += 0.1

p1 = Point(0.5,0.5)
p2 = Point(2.0,5.0)
a = simplex_face(p1,p2)

p1 = Point(0.0,1.0)
p2 = Point(1.0,0.0)
b = simplex_face(p1,p2)

p1 = Point(0.0,2.0)
p2 = Point(2.0,0.0)
c = simplex_face(p1,p2)

p1 = Point(1.0,2.0)
p2 = Point(2.0,1.0)
d = simplex_face(p1,p2)

p1 = Point(1.5,0.5)
p2 = Point(2.0,0.0)
e = simplex_face(p1,p2)

@test voxel_intersection(a,pmin,pmax,p)
@test voxel_intersection(b,pmin,pmax,p)
@test voxel_intersection(c,pmin,pmax,p)
@test !voxel_intersection(d,pmin,pmax,p)
@test !voxel_intersection(e,pmin,pmax,p)

# Voxel intersection Face{3}

p = HEX

pmin,pmax = Point(0.0,0.0,0.0),Point(1.0,1.0,1.0)
pmin -= 0.1
pmax += 0.1

p1 = Point(1.0,0.0,0.0)
p2 = Point(0.0,1.0,0.0)
p3 = Point(0.0,0.0,1.0)
a = simplex_face(p1,p2,p3)

p1 = Point(9.0,-11.0,-11.0)
p2 = Point(-11.0,9.0,9.0)
p3 = Point(10.0,10.0,10.0)
b = simplex_face(p1,p2,p3)

p1 = Point(0.5,0.5,0.5)
p2 = Point(0.0,2.0,0.0)
p3 = Point(0.0,0.0,2.0)
c = simplex_face(p1,p2,p3)

p1 = Point(4.0,0.0,0.0)
p2 = Point(0.0,4.0,0.0)
p3 = Point(0.0,0.0,4.0)
d = simplex_face(p1,p2,p3)

p1 = Point(2.0,-0.2,-0.2)
p2 = Point(1.5,0.5,0.0)
p3 = Point(1.5,0.0,0.5)
e = simplex_face(p1,p2,p3)

@test voxel_intersection(a,pmin,pmax,p)
@test voxel_intersection(b,pmin,pmax,p)
@test voxel_intersection(c,pmin,pmax,p)
@test !voxel_intersection(d,pmin,pmax,p)
@test !voxel_intersection(e,pmin,pmax,p)

end # module


#=

  Euler #4 in Julia.

  Problem 4
  """
  A palindromic number reads the same both ways. The largest palindrome made
  from the product of two 2-digit numbers is 9009 = 91 × 99.

  Find the largest palindrome made from the product of two 3-digit numbers.
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
include("Euler.jl")

# 0.00958293s
function euler4a()
    max = 0
    from = 100
    to = 999
    for i in from:to, j in i:to
        ij = i*j
        if ij > max && palindromic_number(ij)
            max = ij
        end
    end
    return max;
end

# 0.313482985s
function euler4b()
    from = 100
    to = 999
    a = []

    # Note: It seems that one cannot assign a variable (ij=i*j)
    # in a function pipe
    (from:to).|> i->
        (i:to).|> j->
            # ij = i * j; # this don't work! (Why?)
            if palindromic_number(i*j)
                push!(a,i*j)
            end
    return maximum(a)
end

# 0.085608391s
function euler4c()
    from = 100
    to = 999

    a = (from:to).|>i->(i:to).|>j->i*j
    # vcat(a...) flattens the array
    return maximum(filter(ij->palindromic_number(ij),vcat(a...)))
end

# "Full" chaining, much slower
# 0.12105172s
function euler4d()
    from = 100
    to = 999
    a = (from:to).|>i->(i:to).|>j->i*j
    return vcat(a...)|>x->filter(ij->palindromic_number(ij),x)|>maximum
end

run_euler(euler4a)
# run_euler(euler4b)
# run_euler(euler4c)
# run_euler(euler4d)


using Revise
using BNPRegressionGGA2021
using CSV
using DataStructures
using DataFrames
using Distributions
using Gadfly
using Random
using Statistics
const BNP = BNPRegressionGGA2021

# Define the DGP
function simulate_sample(N0, N1)
    X0 = [ones(N0) rand([0, 1], N0, 5)]
    # X1 = kron([ones(2, 5) [1, 0]], ones(N1, 1))

    X1 = rand(2, 5), rand(2, 4)
    α = 4
    δ = α / √(1 + α^2)
    ω = 1
    ξ = - ω * δ * √(2 / π)
    ϵ = √(1 - δ^2) .* randn(N0) + δ .* abs.(randn(N0)) .+ ξ
    y0 = 0.8 * X0[:, end] + ϵ
    y1 = LinRange(-3, 3, N1) |> collect |> x -> repeat(x, size(X1, 1) ÷ N1)
    return y0, X0, y1, X1
end

N0, N1, Niter = 100, 2, 1;
df = zeros(Bool, Niter, 6);
Random.seed!(1);
y0, X0, y1, X1 = simulate_sample(N0, N1);
plot(x = y0, Geom.density())
m = BNP.DGSBPNormal(; y0, X0, y1, X1);
_, _, chaing = BNP.sample!(m; mcmcsize = 10000, burnin = 5000);
mean(chaing)

# Get the MAP of γ for 100 simulated samples
begin
    N0, N1, Niter = 500, 2, 100;
    df = zeros(Bool, Niter, 6);
    Random.seed!(1);
    for iter in 1:Niter
        y0, X0, y1, X1 = simulate_sample(N0, N1);
        m = BNP.DGSBPNormal(; y0, X0, y1, X1);
        _, _, chaing = BNP.sample!(m; mcmcsize = 10000, burnin = 5000);
        df[iter, :] .= begin
            chaing |>
            DataStructures.counter |>
            collect |>
            x -> sort(x, by = x -> x[2], rev = true) |>
            x -> getindex(x, 1) |>
            x -> getindex(x, 1)
        end
        println(iter)
        println(df[iter, :])
    end
    CSV.write("data/simulation-normal-skewnormal-gamma-bnp.csv", DataFrame(df, :auto))
end

# Save the 100 simulated samples
Random.seed!(1);
df = map(1:100) do iter
    y0, X0, _, _ = simulate_sample(N0, N1);
    df = DataFrame(X0, :auto)
    df[!, :y0] = y0
    df
end;
df = reduce(vcat, df);
CSV.write("data/simulation-normal-skewnormal-data.csv", df);


using Revise
using Makie, CairoMakie, ZigZagBoomerang, SparseArrays, LinearAlgebra
using Random

function zz_sticky_events(xs0, i)
    first = true
    y = []
    for x in xs0
        if x[i] == 0.0  && first == true
            push!(y, x)
            first = false
        elseif x[i] == 0.0 && first == false
            continue
        else
            first = true
        end
    end
    return y
end

    Random.seed!(14)
    function ϕ(x, i, μ)
        x[i] - μ[i]
    end
    κ = 2.0
    n = 2
    μ = zeros(n)
    x0 = randn(n)
    θ0 = rand((-1.0,1.0), n)
    T = 10.0
    c = 0.1*ones(n)
    @time trace0, _ = ZigZagBoomerang.sspdmp(ϕ, 0.0, x0, θ0, T, c, κ, ZigZag(sparse(1.0I,n,n), zeros(n)), μ)
    ts0, xs0 = splitpairs(trace0)
    p1  = Scene()
        lines!(getindex.(xs0,1), getindex.(xs0,2), color=(:black, 0.4), markersize=0.1)
        y1 = zz_sticky_events(xs0, 1)
        y2 = zz_sticky_events(xs0, 2)
        scatter!(getindex.(y1,1), getindex.(y1,2),
            color = (:red, 0.3), strokecolor = (:red, 0.3), markersize = 15, marker = 'x')
        scatter!(getindex.(y2,1), getindex.(y2,2),
            color = (:blue, 0.3), strokecolor = (:blue, 0.3), markersize = 15, marker = 'x')
        scatter!([xs0[1][1]], [xs0[1][2]], color = (:black, 1.0), strokecolor = (:black, 1.0), markersize = 2.0)
        text!("x(0),y(0)",textsize = 0.15, position = (xs0[1][1]- 0.05, xs0[1][2] + 0.05))

p2 = lines(ts0, getindex.(xs0,1), color = (:red, 0.5), label = "x1", leg = true)
    lines!(p2, ts0, getindex.(xs0,2), color = (:blue, 0.5), label = "x2")
    xlabel!(p2, "t")
    ylabel!(p2, "")
p = vbox(p1,p2)

save("./figures/2d_sticky_samplers.png", title(p, "2d sticky Zig-Zag"))


"""
Utility functions for Tensor Graphical Models

Author: Wayne Wang
Last modified: 07/14/2021
"""

function kroneckersum_alt(A::AbstractMatrix, B::AbstractMatrix)
    """
    Alternate Kronecker sum defined as I ⊗ A + B ⊗ I
    """
    return kroneckersum(B,A)
end

function kroneckersum_list(mat_list::AbstractVector{<:Any})
    """
    Kronecker sum for a list of matrices using package Kronecker.jl
    In
        - mat_list::AbstractVector{<:Any}: a list of matrices
    Out
        - mat::KroneckerSum: mat_list[1] ⊕ ... ⊕ mat_list[end]
    """
    K = length(mat_list)
    mat = copy(mat_list[1])
    if K == 1
        return mat
    else
        for j=2:K
            mat = kroneckersum_alt(mat, mat_list[j])
        end
    end
    return mat
end

function mul_alt!(C::StridedMatrix, X::StridedMatrix, A::SparseMatrixCSC)
    mX, nX = size(X)
    nX == A.m || throw(DimensionMismatch())
    fill!(C, zero(eltype(C)))
    rowval = A.rowval
    nzval = A.nzval
    @inbounds for  col = 1:A.n, k=A.colptr[col]:(A.colptr[col+1]-1)
        ki=rowval[k]
        kv=nzval[k]
        for multivec_row=1:mX
            C[multivec_row, col] += X[multivec_row, ki] * kv
        end
    end
    C
end

import Base.*
function *(B::StridedMatrix, A::SparseMatrixCSC)
    mB, nB = size(B)
    nB == A.m || throw(DimensionMismatch())
    C = zeros(mB,A.n)
    rowval = A.rowval
    nzval = A.nzval
    @inbounds for  col = 1:A.n, k=A.colptr[col]:(A.colptr[col+1]-1)
        ki=rowval[k]
        kv=nzval[k]
        for multivec_row=1:mB
            C[multivec_row, col] += B[multivec_row, ki] * kv
        end
    end
    C
end

function offdiag(X::AbstractMatrix)
    """
    Extract the off-diagonal matrix of X
    """
    return X - Diagonal(X)
end

function offdiag!(X::AbstractMatrix)
    """
    In placce version of offdiag
    """
    X .= X .- Diagonal(X)
    return X
end


using LightGraphs
using Test

include("../../src/utils/filters.jl")

@testset "utils/filters" begin
    @testset "graph degree" begin
        graph = path_graph(10)
        add_edge!(graph, 1, 4)
        add_edge!(graph, 3, 8)
        add_edge!(graph, 4, 8)
        add_edge!(graph, 5, 7)
        add_edge!(graph, 6, 9)
        @test degreeisgreater(graph, 4, 4)
        @test !degreeisgreater(graph, 1, 3)
        @test !degreeislesser(graph, 4, 3)
        @test degreeislesser(graph, 1, 3)
        @test degreeis(graph, 1, [2])
        @test degreeis(graph, 1, [1, 2])
        @test !degreeis(graph, 1, [3])
        @test !degreeis(graph, 1, [3, 4])
    end
    @testset "max/min size subarray" begin
        arrays = [[1, 2, 3], [4], [5, 6, 7, 8]]
        @test getmaxsize(arrays) == 4
        @test getminsize(arrays) == 1
        arrays = []
        @test getmaxsize(arrays) == -1
        @test getminsize(arrays) == -1
    end
    @testset "getsubset" begin
         arrays = [[1, 2, 3], [1], [3], [1, 6, 3, 8]]
         @test getsubset(arrays, [1, 3]) == [[1, 2, 3], [1, 6, 3, 8]]
         @test getsubset(arrays, [2, 6]) == []
    end
end


function sort_lists!(main_list, other_lists...; cur_func::Function=identity)
  cur_indices = sortperm(main_list, by=cur_func)

  cur_list = collect(main_list)[cur_indices]
  _arrange_lists!(main_list, other_lists, cur_indices)
  cur_list
end

function shuffle_lists!(main_list, other_lists...)
  cur_indices = shuffle(1:length(main_list))

  cur_list = main_list[cur_indices]
  _arrange_lists!(main_list, other_lists, cur_indices)
  cur_list
end

function _arrange_lists!(main_list, other_lists, cur_indices)
  cur_lists = Any[ isa(main_list, Vector) ? main_list : collect(main_list) ]
  isempty(other_lists) || append!(cur_lists, other_lists)

  for cur_list in cur_lists
      cur_list .= cur_list[cur_indices]
  end
  main_list
end

function bisect_reorder!(cur_vector::Vector)
  cur_indices = breadth_indices(length(cur_vector))
  cur_vector .= cur_vector[cur_indices]
  cur_vector
end

function collect_bisection_indexes!(ix, level, a, b)
    a ≤ b || return
    mid = (a+b) ÷ 2
    push!(ix, level => mid)
    collect_bisection_indexes!(ix, level + 1, a, mid - 1)
    collect_bisection_indexes!(ix, level + 1, mid + 1, b)
end

function breadth_indices(n)
    ix = Vector{Pair{Int,Int}}()
    sizehint!(ix, n)
    collect_bisection_indexes!(ix, 0, 1, n)
    last.(sort!(ix, by = first))
end

function out_in_reorder!(cur_vector::Vector)
  cur_indices = centering_indices(length(cur_vector))
  cur_vector .= cur_vector[cur_indices]
  cur_vector
end

function centering_indices(num_points::Number)
  cur_indices = []

  cur_half_point = Int(floor(num_points/2))

  for cur_index in 1:cur_half_point
    push!(cur_indices, cur_index)
    push!(cur_indices, num_points-(cur_index-1))
  end
  isodd(num_points) && push!(cur_indices, cur_half_point+1)

  cur_indices
end

function in_out_reorder!(cur_vector::Vector)
  cur_indices = dividing_indices(length(cur_vector))
  cur_vector .= cur_vector[cur_indices]
  cur_vector
end

function dividing_indices(num_points::Number)
  cur_indices = []

  cur_half_point = Int(floor(num_points/2))

  isodd(num_points) && push!(cur_indices, cur_half_point+1)
  for cur_index in reverse(1:cur_half_point)
    push!(cur_indices, cur_index)
    push!(cur_indices, num_points-(cur_index-1))
  end

  cur_indices
end


using Combinatorics: combinations

function build_excursion(arity, n_way, disallow, seed = nothing)
    test_list = [zeros(eltype(arity), length(arity), 1)]
    for way_idx in 1:n_way
        push!(test_list, all_combinations(arity .- 1, way_idx))
    end
    test_cases = hcat(test_list...)
    test_cases .+= 1

    if seed !== nothing
        seeded = unique(hcat(seed, test_cases), dims = 2)
    else
        seeded = test_cases
    end

    allowed = zeros(Int, length(seeded))
    allow_cnt = 0
    for case_idx in 1:size(seeded, 2)
        if !disallow(seeded[:, case_idx])
            allow_cnt += 1
            allowed[allow_cnt] = case_idx
        end
    end
    seeded[:, allowed[1:allow_cnt]]
end

function build_excursion_multi(arity, n_way, levels, disallow, seed = nothing)
    param_cnt = length(arity)
    levels[n_way] = [collect(1:length(arity))]
    # excursion_set = Set{Tuple{Vararg{Int64, N} where N}}()
    excursion_set = Set{Array{Int, 1}}()
    for level in keys(levels)
        for indices in levels[level]
            # combos = Set(collect(tuple(x...) for x in combinations(indices, level)))
            for subset_idx in 1:level
                combos = Set(collect(combinations(indices, subset_idx)))
                union!(excursion_set, combos)
            end
        end
    end
    essity = arity .- 1  # 1 is the origin of the excursion
    excursions = sort(collect(excursion_set), by = x -> (length(x), x))
    include_origin = 1
    test_cnt = sum(prod(essity[inds]) for inds in excursions) + include_origin

    test_cases = ones(Int, param_cnt, test_cnt)
    test_cases[:, 1] .= 1
    test_idx = include_origin
    for exc_indices in excursions
        walks = ones(Int, length(exc_indices))
        for sub_idx in 1:prod(essity[exc_indices])
            test_idx += 1
            test_cases[exc_indices, test_idx] = walks .+ 1
            next_multiplicative!(walks, essity[exc_indices])
        end
        @assert walks == ones(Int, length(exc_indices))
    end
    @assert test_idx == test_cnt

    if seed !== nothing
        seeded = unique(hcat(seed, test_cases), dims = 2)
    else
        seeded = test_cases
    end

    allowed = zeros(Int, length(seeded))
    allow_cnt = 0
    for case_idx in 1:size(seeded, 2)
        if !disallow(seeded[:, case_idx])
            allow_cnt += 1
            allowed[allow_cnt] = case_idx
        end
    end
    seeded[:, allowed[1:allow_cnt]]
end


#!/usr/bin/env julia

# julia script to run 1 simulation
# argument: integer that will be mapped to a combination of parameters
# assumption: directory "simresults/" created prior to running the script

# This is to be used with slurm, which defines a shell variable SLURM_ARRAY_TASK_ID.
# This variable will serve the integer argument to the script,
# to determine the combination of parameters that each slurm task will run.

# example use, with default of 2 replicates per parameter combination, combination 14:
# julia onesimulation.jl 14
# or simply:
# ./onesimulation.jl 14
# then look at result file in the result directory

# example use with 200 replicates per parameter combination:
# julia onesimulation.jl 14 200

resultdirectory = "simresults"

"""
    onesimulation(N, n_rare_cat, mu, seed)

Simulate data from a multinomial distribution with 3 categories
and `N` units total.
- `n_rare_cat` = 1 or 2 rare categories, with very small probabilities
- `mu`: expected (or mean) number of observations in each rare category.
  The probability that each observation is in a given rare category is
  `p=mu/N`.

With 2 rare categories, the probabilities are [p,p,1-2p].
With 1 rare category, the probabilities are [p,(1-p)/2,(1-p)/2].

The data (counts in each categories) are then analyzed for a
goodness-of-fit test:
1. first using the Pearson test statistics: sum of (oi - ei)^2 / ei
   where ei = expected count in category i, oi = observed count
2. second using Qlog statistics: 2 * sum of L(ni,ei) * (oi - ei)/ei
   where L(x,y) is the logarithmic mean of x and y:
   L(x,y) = (x-y) / (log(x)-log(y)).
   L(x,y) is always between x and y, as expected for a mean.
   See for example Lorenzen (1995) for a justification for using Qlog.

Both test statistics are compared to a chi-square distribution
with 2 degrees of freedom to obtain p-values.

output: (pearson_stat, Qlog_stat, pearson_pvalue, Qlog_pvalue)

examples:

```julia
onesimulation( 50, 1, 0.1)
onesimulation(500, 2, 2.0)
```
"""
function onesimulation(N, nrarecat, mu, seed)
  Random.seed!(seed)
  p = mu/N
  @assert p>=0.0 && p<=0.5 "the rare category/ies cannot have a negative or >.5 probability"
  if nrarecat==1
    probs = [p,(1.0-p)/2,(1.0-p)/2]
    @assert p<=1.0
  else
    probs = [p, p, (1.0-2p)]
  end
  df = 2 # because 3 categories
  observed = rand(Multinomial(N,probs), 1) # 3x1 array
  phat = observed./N # estimated probabilities of the 3 categories
  pearson = N * sum((phat - probs).^2 ./probs)
  Qlog = 0.0
  for i in 1:3
    if observed[i] == 0 continue; end # no need to do xxx / log(0) = 0.0
    ratio = phat[i]/probs[i]
    ratio ≉ 1.0 || continue
    dev = ratio - 1.0 # deviation on ratio scale
    # if ratio ≈ 1 then dev/log(ratio) ≈ 1, but NaN if ratio = 1 exactly
    Qlog += probs[i] * dev^2 / log(ratio)
  end
  Qlog *= 2N
  # cdf = cumulative distribution function: P(random value <= x)
  # ccdf = complementary cdf:               P(random value >  x)
  p_pearson = ccdf(Chisq(df), pearson)
  p_Qlog    = ccdf(Chisq(df), Qlog)
  return pearson,Qlog, p_pearson,p_Qlog
end

"""
    arrayID_to_parameters(arrayID, total_n_reps)

Convert a single integer into 4 parameters to be used for the simulation:
`sample_size, n_rare_categories, mu, replicate_number`.
The integer is meant to be used as the "array ID" when submitting all these
simulations with slurm and with a slurm array.

Assumption: parameters vary in
- sample_size: 30, 1000
- n_rare_categories: 1, 2
- mu: 0.1, 1.0, 2.0
- replicate_number: from 1 to total_n_reps (input: to start small when developing the script)
"""
function arrayID_to_parameters(arrayID::Int, Nrep::Int)
  # start with Nrep = 2 to develop the julia script + slurm script,
  # then ramp up to Nrep = 200 or more
  samplesizes = [30, 1000]
  nrarecats = [1, 2]
  mus = [0.1, 1., 2.]

  Nssize    = length(samplesizes)
  Nnrarecat = length(nrarecats)
  Nmu       = length(mus)

  ntotal = Nrep * Nssize * Nnrarecat * Nmu
  @assert arrayID <= ntotal "ID must be between 1 and $ntotal"
  I = CartesianIndices( (Nrep,Nssize,Nnrarecat,Nmu) )
  I = I[arrayID].I # coordinates corresponding to the linear index arrayID
  rep = I[1]
  samplesize = samplesizes[I[2]]
  nrarecat = nrarecats[I[3]]
  mu = mus[I[4]]
  return rep, samplesize, nrarecat, mu
end

#---------------------------------------#
#  finally: read script argument        #
#           and run the simulation      #
#---------------------------------------#

# parse the integer argument
@assert length(ARGS)>0 "need 1 parameters: arrayID"
arrayID = parse(Int, ARGS[1])
if length(ARGS)==1
  Nreps = 2
  @info "number of replicates will be 2 by default"
else
  Nreps = parse(Int, ARGS[2])
  @info "number of replicates: $Nreps"
end
rep, samplesize, nrarecat, mu = arrayID_to_parameters(arrayID, Nreps)
@info ""  arrayID replicate=rep samplesize number_rare_categories=nrarecat mu

# load the packages that the simulation function needs
using Distributions # needs to be installed prior
using Random        # no install: standard library
using Printf        # standard library

# run the simulation. use arrayID as seed: will be different for each simulation
pearson, qlog, p_pearson, p_qlog = onesimulation(samplesize, nrarecat, mu, arrayID)

# save the result in tiny csv-formatted file
# (later, all these files will be concatenated with "cat")
outputfile = joinpath(resultdirectory, "simulation_" * @sprintf("%04d", arrayID) * ".csv")
open(outputfile, "w") do g
  write(g, "$arrayID,$rep,$samplesize,$nrarecat,$mu,") # input
  write(g, "$pearson,$qlog,$p_pearson,$p_qlog\n")      # output
end
@info "ID $arrayID is done"


export rand_rot, set_rot, get_rot, check_rot, faces, euler_char
export embed_rot, NF, dual

"""
`rand_rot(G::SimpleGraph)` creates a random rotation system for
the graph.
"""
function rand_rot(G::SimpleGraph{T}) where {T}
    d = Dict{T,RingList{T}}()
    for v in G.V
        a = shuffle(RingList(G[v]))
        d[v] = a
    end
    set_rot(G, d)
end

"""
`_default_rot(G::SimpleGraph)` creates a default
rotation system for the graph.
"""
function _default_rot(G::SimpleGraph{T}) where {T}
    d = Dict{T,RingList{T}}()
    for v in G.V
        d[v] = RingList(G[v])
    end
    return d
end

"""
`set_rot(G::SimpleGraph,d)` makes `d` the rotation system
for this graph (held in the graph's cache).

If `d` is omitted, the a default rotation is used.
"""
function set_rot(G::SimpleGraph, d)
    if !check_rot(G, d)
        error("Not a valid rotation system for this graph")
    end
    cache_save(G, :RotationSystem, d)
    nothing
end
set_rot(G::SimpleGraph) = set_rot(G, _default_rot(G))

"""
`get_rot(G::SimpleGraph,v)` returns a `RingList` of the neighbors of `v`.
This assumes that `G` has an associate rotation system. If not, one is
generated and an alert is generated.

`get_rot(G::SimpleGraph)` returns  a copy of the rotation system associated with `G`
(or an error if there is no rotation system).
"""
function get_rot(G::SimpleGraph, v)
    if !has(G, v)
        error("No such vertex $v in this graph")
    end
    d = get_rot(G)
    return d[v]
end
function get_rot(G::SimpleGraph)
    if cache_check(G, :RotationSystem)
        return cache_recall(G, :RotationSystem)
    end

    if cache_check(G, :GraphEmbedding)
        embed_rot(G)
        @info "Giving this graph, $G, a rotation system based on its xy-embedding."
    else
        set_rot(G)
        @info "Giving this graph, $G, a default rotation system."
    end
    return get_rot(G)
end

"""
`embed_rot(G::SimpleGraph)` assigns a rotation system to `G`
corresponding to its current `xy` embedding.
"""
function embed_rot(G::SimpleGraph{T}) where {T}
    xy = getxy(G)
    d = Dict{T,RingList{T}}()
    for v in G.V
        o = Complex(xy[v]...)    # location of v as a complex number
        Nv = G[v]  # neighbors of v as a list

        Z = [Complex(xy[w]...) for w in Nv]  # neighbor locations as complex nums
        th = [mod(angle(o - z), 2pi) for z in Z]  # angles of edges from v
        p = sortperm(th)
        d[v] = RingList(Nv[p])
    end
    set_rot(G, d)
end

"""
`check_rot(G::SimpleGraph,d::Dict)`  checks if `d` is a valid
rotation system for `G`.
"""
function check_rot(G::SimpleGraph{T}, d::Dict{T,RingList{T}}) where {T}
    for v in G.V
        if !haskey(d, v)
            return false
        end
        if Set(d[v]) != Set(G[v])
            return false
        end
    end
    return true
end

function _next_edge(G::SimpleGraph{T}, uv::Tuple{T,T}) where {T}
    u, v = uv
    r = get_rot(G, v)
    w = r(u)
    return (v, w)
end

"""
`_trace_face(G::SimpleGraph, v,w)` uses the graph's rotation system to find a
face starting with the edge `(u,v)` (in that order). Also may be called
by `_trace_face(G,(u,v))`.
"""
function _trace_face(G::SimpleGraph{T}, uv::Tuple{T,T}) where {T}
    data = T[]
    u, v = uv
    push!(data, u)
    push!(data, v)
    e = uv
    while true
        vw = _next_edge(G, e)
        if vw == uv
            break
        end
        push!(data, vw[2])
        e = vw
    end

    boundary = [(data[i], data[i+1]) for i = 1:length(data)-1]

    return RingList(boundary)
end
_trace_face(G::SimpleGraph{T}, u::T, v::T) where {T} = _trace_face(G, (u, v))

"""
`faces(G::SimpleGraph)` returns the set of faces of this graph
(given its rotation system). The rotation system is a
planar embedding iff this returns `2`.

Each face is a `RingList` of the (directed) edges bordering the face.

*Requires that the graph is connected and has at least one edge.*
"""
function faces(G::SimpleGraph{T}) where {T}
    FT = RingList{Tuple{T,T}}   # type of an oriented face
    result = Set{FT}()

    for e in G.E
        u, v = e
        F = _trace_face(G, u, v)
        push!(result, F)
        F = _trace_face(G, v, u)
        push!(result, F)
    end

    return result
end

"""
`NF(G)` returns the number of faces in the graph `G`
given its current rotation system.

*Requires that the graph is connected and has at least one edge.*
"""
NF(G::SimpleGraph) = length(faces(G))

"""
`euler_char(G::SimpleGraph)` computes the Euler characteristic
of the graph `G` with its associated rotation system.
Specifically, `euler_char(G)` returns `NV(G) - NE(G) + NF(G)`.

*Requires that the graph is connected and has at least one edge.*
"""
euler_char(G::SimpleGraph) = NV(G) - NE(G) + NF(G)

_shorten(F::RingList) = first.(F)

"""
`dual(G::SimpleGraph)` returns the dual graph of `G`.
The vertices of the dual are the faces of `G` and they are
adjacent if and only if they share a common edge.

*Requires that the graph is connected and has at least one edge.*
"""
function dual(G::SimpleGraph{T}) where {T}
    Flist = collect(faces(G))    # list of faces of the graph
    FT = RingList{Tuple{T,T}}    # data type of faces
    VT = Vector{T}               # data type of shortened faces
    GG = SimpleGraph{VT}()       # graph to return

    # the faces of G are the vertices of GG
    for f in Flist
        add!(GG, _shorten(f))
    end

    d = Dict{Tuple{T,T},FT}()    # mapping from (directed) edges to faces
    for f in Flist
        for e in f
            d[e] = f
        end
    end

    # add edges
    for e in G.E
        F1 = d[e]
        F2 = d[reverse(e)]
        add!(GG, _shorten(F1), _shorten(F2))
    end

    try
        # create a rotation system
        R = Dict{VT,RingList{VT}}()   # mapping from a face to a cycle of its neighbors
        for f in Flist
            bdy = [_shorten(d[reverse(e)]) for e in f]
            R[_shorten(f)] = RingList(bdy)
        end

        set_rot(GG, R)
    catch
        @info "Unable to transfer rotation system to the dual; assigning a default instead."
        set_rot(GG)
    end

    try
        embed(GG, :tutte)
    catch
    end

    name(GG, "Dual of $(name(G))")

    return GG
end


let
    arr = ones(2,2)
    pad_matrix(arr,(2,2),0)
end
let
    a = [1.0,2,3]
    b = [-4.0,5,6]
    @test array_isapprox(cross(a,b),cross_product_operator(a)*b)
end
let
    n = 4
    for i in 1:n
        v = one_hot(n,i)
        @test v[i] == 1
        v = one_hot(i,n)
        @test v[i] == 1
        for j in 1:n
            if i != j
                @test v[j] == 0
            end
        end
    end
    @test_throws AssertionError one_hot(-1,3)
    @test_throws AssertionError one_hot(1,-3)
    @test_throws AssertionError one_hot(-1,-3)
end
let
    vec = [1,2,3,4]
    for i in 1:4
        @test wrap_idx(length(vec),i) == i
        @test wrap_get(vec,i)[1] == i
    end
    @test wrap_idx(length(vec),5) == 1
    @test wrap_get(vec,5)[1] == 1
    @test wrap_idx(length(vec),9) == 1
    @test wrap_idx(length(vec),0) == 4
    @test wrap_idx(length(vec),-4) == 4
    @test wrap_idx(length(vec),-5) == 3
end
let
    a = [
        1 2;
        3 4
        ]
    @test wrap_get(a,[1,1]) == 1
    @test wrap_get(a,[3,3]) == 1
    @test wrap_get(a,[4,4]) == 4
    @test wrap_get(a,[0,-1]) == 3
end


using MagSim
using Plots;
gr();
using DataFrames
using CSV
using LaTeXStrings

# %%

βs = range(1e-3, stop = 1, length = 201)
Ls = [4, 8, 16, 32, 64];
prefix = "2020-5-14-simulating-magnetism-part-4";

q = 2
@time sims = simulate_Potts_parallel(q, metropolis!, Ls, βs);
df = sim_dict_to_df(sims);
CSV.write(prefix * "_q=$(q).csv", df);

# %%

function plot_data(df, Os, y_limits, y_ticks)
    plots = Dict()
    labels = latexstring.(Os)
    Ls = unique(df.L)

    for (i, O) in enumerate(Os)
        O ∈ [:e] ? legend_bool = true : legend_bool = false
        p = plot(
            dpi = 200,
            xaxis = (latexstring("\\beta")),
            yaxis = (labels[i], y_limits[i], y_ticks[i]),
            legend = legend_bool,
        )

        colors = [:firebrick1, :orange, :forestgreen, :cyan3, :dodgerblue]
        for (i, L) in enumerate(Ls)
            df_ = df[(df.L.==L), :]
            x = df_[:, :β]
            y = df_[:, O]
            plot!(
                x,
                y,
                marker = (:star4, colors[i], 4, stroke(0)),
                line = (colors[i], 2),
                label = latexstring("L = $(string(L))"),
            )
        end

        plots[O] = p
    end

    return plots
end

Os = [:e, :c]
y_limits = [(-2.1, -0.9), (-0.1, 2.6)]
y_ticks = [-2:0.2:-1, 0:0.5:2.5]

plots = plot_data(df, Os, y_limits, y_ticks);
plot_e_c = plot(plots[:e], plots[:c], layout = (2, 1), size = (600, 500));
plot_e_c
savefig(plot_e_c, "../images/" * prefix * "/potts_2_metropolis_e_c.png")


import Base: numerator, denominator, (==), zero, one, big

export SimpleRationalFunction

struct SimpleRationalFunction
    num::SimplePolynomial
    den::SimplePolynomial
    function SimpleRationalFunction(top::SimplePolynomial, bot::SimplePolynomial)
        if bot == 0
            error("Denominator cannot be zero")
        end
        d = gcd(top, bot)
        t, r = divrem(top, d)
        b, r = divrem(bot, d)

        # make the denominator monic
        a = b.data[end]
        t = t / a
        b = b / a

        new(t, b)
    end
end

SimpleRationalFunction(f::SimplePolynomial) = SimpleRationalFunction(f, SimplePolynomial(1))
SimpleRationalFunction(a::T) where {T<:CoefX} = SimpleRationalFunction(SimplePolynomial(a))
SimpleRationalFunction(a::T, f::SimplePolynomial) where {T<:CoefX} =
    SimpleRationalFunction(SimplePolynomial(a), f)
SimpleRationalFunction(f::SimplePolynomial, a::T) where {T<:CoefX} =
    SimpleRationalFunction(f, SimplePolynomial(a))
SimpleRationalFunction(a::T, b::S) where {T<:CoefX,S<:CoefX} =
    SimpleRationalFunction(a // b)
SimpleRationalFunction() = SimpleRationalFunction(0)

function make_function(f::SimpleRationalFunction)
    F(x) = (f.num.func(x)) / (f.den.func(x))
end

numerator(f::SimpleRationalFunction) = f.num
denominator(f::SimpleRationalFunction) = f.den

numerator(f::SimplePolynomial) = f
denominator(f::SimplePolynomial) = one(eltype(f))

(==)(f::SimpleRationalFunction, g::SimpleRationalFunction) =
    (f.num == g.num) && (f.den == g.den)
(==)(f::SimpleRationalFunction, p::SimplePolynomial) = f == SimpleRationalFunction(p)
(==)(p::SimplePolynomial, f::SimpleRationalFunction) = f == p
(==)(f::SimpleRationalFunction, a::T) where {T<:CoefX} = f == SimpleRationalFunction(a)
(==)(a::T, f::SimpleRationalFunction) where {T<:CoefX} = f == a

(f::SimpleRationalFunction)(x) = f.num(x) / f.den(x)

zero(::Type{SimpleRationalFunction}) = SimpleRationalFunction(0, 1)
zero(::SimpleRationalFunction) = SimpleRationalFunction(0, 1)
one(::Type{SimpleRationalFunction}) = SimpleRationalFunction(1, 1)
one(::SimpleRationalFunction) = SimpleRationalFunction(1, 1)

# big(f::SimpleRationalFunction) = SimpleRationalFunction(big(f.num),big(f.den))

# include("arithmetic.jl")
# include("show.jl")