### mcvar and mcse stand for Monte Carlo variance and Monte Carlo error respectively

## Monte Carlo variance assuming IID samples

mcvar(v::AbstractArray, ::Type{Val{:iid}}) = var(v)/length(v)

mcvar(v::AbstractArray, ::Type{Val{:iid}}, region) = mapslices(x -> mcvar(x, Val{:iid}), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:iid}}) = mcvar(s.value, Val{:iid})

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, i::Integer) = mcvar(s.value[i, :], Val{:iid})

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, r::AbstractVector=1:s.size) =
  eltype(s)[mcvar(s, Val{:iid}, i) for i in r]

## Monte Carlo standard error assuming IID samples

mcse(v::AbstractArray, ::Type{Val{:iid}}) = sqrt(mcvar(v, Val{:iid}))

mcse(v::AbstractArray, ::Type{Val{:iid}}, region) = mapslices(x -> mcse(x, Val{:iid}), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:iid}}) = mcse(s.value, Val{:iid})

mcse(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, i::Integer) = mcse(s.value[i, :], Val{:iid})

mcse(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, r::AbstractVector=1:s.size) =
  eltype(s)[mcse(s, Val{:iid}, i) for i in r]

## Monte Carlo variance using batch means
## Reference:
## Flegal J.M and Jones G.L
## Batch Means and Spectral Variance Estimators in Markov chain Monte Carlo
## Annals of Statistics, 2010, 38 (2), pp 1034-1070

function mcvar(v::AbstractVector{T}, ::Type{Val{:bm}}, batchlen::Integer=100) where T
  nbatches = div(length(v), batchlen)
  @assert nbatches > 1 "Choose batch size such that the number of batches is greather than one"
  nbsamples = nbatches*batchlen
  batchmeans = T[mean(v[((j-1)*batchlen+1):(j*batchlen)]) for j = 1:nbatches]
  return batchlen*var(batchmeans)/nbsamples
end

mcvar(v::AbstractArray, ::Type{Val{:bm}}, region, batchlen::Integer=100) =
  mapslices(x -> mcvar(vec(x), Val{:bm}, batchlen), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:bm}}, batchlen::Integer=100) = mcvar(s.value, Val{:bm}, batchlen)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, i::Integer, batchlen::Integer=100) =
  mcvar(s.value[i, :], Val{:bm}, batchlen)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, r::AbstractVector=1:s.size, batchlen::Integer=100) =
  eltype(s)[mcvar(s, Val{:bm}, i, batchlen) for i in r]

## Monte Carlo standard error using batch means

mcse(v::AbstractVector, ::Type{Val{:bm}}, batchlen::Integer=100) = sqrt(mcvar(v, Val{:bm}, batchlen))

mcse(v::AbstractArray, ::Type{Val{:bm}}, region, batchlen::Integer=100) =
  mapslices(x -> mcse(vec(x), Val{:bm}, batchlen), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:bm}}, batchlen::Integer=100) = mcse(s.value, Val{:bm}, batchlen)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, i::Integer, batchlen::Integer=100) =
  mcse(s.value[i, :], Val{:bm}, batchlen)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, r::AbstractVector=1:s.size, batchlen::Integer=100) =
  eltype(s)[mcse(s, Val{:bm}, i, batchlen) for i in r]

## Initial monotone sequence estimator (IMSE) of Monte Carlo variance
## Reference:
## Geyer C.J
## Practical Markov Chain Monte Carlo
## Statistical Science, 1992, 7 (4), pp 473-483

function mcvar(v::AbstractVector{T}, ::Type{Val{:imse}}, maxlag::Integer=length(v)-1) where T
  k = convert(Integer, floor((maxlag-1)/2))
  m = k+1

  # Preallocate memory
  g = Array(T, k+1)

  # Calculate empirical autocovariance
  acv = autocov(v, 0:maxlag)

  # Calculate \hat{G}_{n, m} from the autocovariances, see pp 477 in Geyer
  for j = 0:k
    g[j+1] = acv[2*j+1]+acv[2*j+2]
    if g[j+1] <= 0
      m = j
      break
    end
  end

  # Create the monotone sequence of g
  if m > 1
    for j = 2:m
      if g[j] > g[j-1]
        g[j] = g[j-1]
      end
    end
  end

  # Calculate the initial monotone sequence estimator
  return (-acv[1]+2*sum(g[1:m]))/length(v)
end

mcvar(v::AbstractArray, ::Type{Val{:imse}}, region) = mapslices(x -> mcvar(vec(x), Val{:imse}, length(x)-1), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:imse}}, maxlag::Integer=s.n-1) = mcvar(s.value, Val{:imse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcvar(s.value[i, :], Val{:imse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcvar(s, Val{:imse}, i, maxlag) for i in r]

## Initial monotone sequence estimator (IMSE) of Monte Carlo standard error

mcse(v::AbstractVector, ::Type{Val{:imse}}, maxlag::Integer=length(v)-1) = sqrt(mcvar(v, Val{:imse}, maxlag))

mcse(v::AbstractArray, ::Type{Val{:imse}}, region) = mapslices(x -> mcse(vec(x), Val{:imse}, length(x)-1), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:imse}}, maxlag::Integer=s.n-1) = mcse(s.value, Val{:imse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcse(s.value[i, :], Val{:imse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcse(s, Val{:imse}, i, maxlag) for i in r]

## Initial positive sequence estimator (IPSE) of Monte Carlo variance
## Reference:
## Geyer C.J
## Practical Markov Chain Monte Carlo
## Statistical Science, 1992, 7 (4), pp 473-483

function mcvar(v::AbstractVector{T}, ::Type{Val{:ipse}}, maxlag::Integer=length(v)-1) where T
  k = convert(Integer, floor((maxlag-1)/2))

  # Preallocate memory
  g = Array(T, k+1)
  m = k+1

  # Calculate empirical autocovariance
  acv = autocov(v, 0:maxlag)

  # Calculate \hat{G}_{n, m} from the autocovariances, see pp. 477 in Geyer
  for j = 0:k
    g[j+1] = acv[2*j+1]+acv[2*j+2]
    if g[j+1] <= 0
      m = j
      break
    end
  end

  # Calculate the initial positive sequence estimator
  return (-acv[1]+2*sum(g[1:m]))/length(v)
end

mcvar(v::AbstractArray, ::Type{Val{:ipse}}, region) = mapslices(x -> mcvar(vec(x), Val{:ipse}, length(x)-1), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:ipse}}, maxlag::Integer=s.n-1) = mcvar(s.value, Val{:ipse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcvar(s.value[i, :], Val{:ipse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcvar(s, Val{:ipse}, i, maxlag) for i in r]

## Initial positive sequence estimator (IPSE) of Monte Carlo standard error

mcse(v::AbstractVector, ::Type{Val{:ipse}}, maxlag::Integer=length(v)-1) = sqrt(mcvar(v, Val{:ipse}, maxlag))

mcse(v::AbstractArray, ::Type{Val{:ipse}}, region) = mapslices(x -> mcse(vec(x), Val{:ipse}, length(x)-1), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:ipse}}, maxlag::Integer=s.n-1) = mcse(s.value, Val{:ipse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcse(s.value[i, :], Val{:ipse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcse(s, Val{:ipse}, i, maxlag) for i in r]

## Wrapper function for estimating Monte Carlo variance using various approaches

mcvar(v::AbstractArray, vtype::Symbol, args...) = mcvar(v, Val{vtype}, args...)

mcvar(s::VariableNState{Univariate}, vtype::Symbol, args...) = mcvar(s, Val{vtype}, args...)

mcvar(s::VariableNState{Multivariate}, vtype::Symbol, i::Integer, args...) = mcvar(s, Val{vtype}, i, args...)

mcvar(s::VariableNState{Multivariate}, vtype::Symbol, r::AbstractVector=1:s.size, args...) = mcvar(s, Val{vtype}, r, args...)

mcvar(v::AbstractArray, args...) = mcvar(v, Val{:imse}, args...)

mcvar(s::VariableNState{Univariate}, args...) = mcvar(s, Val{:imse}, args...)

mcvar(s::VariableNState{Multivariate}, i::Integer, args...) = mcvar(s, Val{:imse}, i, args...)

mcvar(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size, args...) = mcvar(s, Val{:imse}, r, args...)

## Wrapper function for estimating Monte Carlo standard error using various approaches

mcse(v::AbstractArray, vtype::Symbol, args...) = mcse(v, Val{vtype}, args...)

mcse(s::VariableNState{Univariate}, vtype::Symbol, args...) = mcse(s, Val{vtype}, args...)

mcse(s::VariableNState{Multivariate}, vtype::Symbol, i::Integer, args...) = mcse(s, Val{vtype}, i, args...)

mcse(s::VariableNState{Multivariate}, vtype::Symbol, r::AbstractVector=1:s.size, args...) = mcse(s, Val{vtype}, r, args...)

mcse(v::AbstractArray, args...) = mcse(v, Val{:imse}, args...)

mcse(s::VariableNState{Univariate}, args...) = mcse(s, Val{:imse}, args...)

mcse(s::VariableNState{Multivariate}, i::Integer, args...) = mcse(s, Val{:imse}, i, args...)

mcse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size, args...) = mcse(s, Val{:imse}, r, args...)


"""
    WeightedMean(T = Float64)

Simple weighted mean, tracked as type `T`.

# Example:
    o = fit!(WeightedMean(), rand(100), rand(100))
    sum(o)
    mean(o)
"""
mutable struct WeightedMean{T} <: WeightedOnlineStat{T}
    μ::T
    W::T
    n::Int
    function WeightedMean{T}(μ = T(0), W = T(0), n = 0) where T
        new{T}(T(μ), T(W), Int(n))
    end
end

WeightedMean(μ::T, W::T, n::Int) where T = WeightedMean{T}(μ, W, n)
WeightedMean(::Type{T}) where T = WeightedMean(T(0), T(0), 0)
WeightedMean() = WeightedMean(Float64)

function OnlineStatsBase._fit!(o::WeightedMean{T}, x, w) where T
    xx = convert(T, x)
    ww = convert(T, w)

    o.n += 1
    o.W = smooth(o.W, ww, T(1) / o.n)
    o.μ = smooth(o.μ, xx, ww / (o.W * o.n))

    o
end

function OnlineStatsBase._merge!(o::WeightedMean{T}, o2::WeightedMean) where T
    o2_W = convert(T, o2.W)
    o2_μ = convert(T, o2.μ)

    o.n += o2.n
    o.W = smooth(o.W, o2_W, o2.n / o.n)
    o.μ = smooth(o.μ, o2_μ, (o2_W * o2.n) / (o.W * o.n))

    o
end

OnlineStatsBase.value(o::WeightedMean) = o.μ
Statistics.mean(o::WeightedMean) = value(o)
Base.sum(o::WeightedMean) = mean(o) * weightsum(o)
Base.copy(o::WeightedMean) = WeightedMean(o.μ, o.W, o.n)


function brinco1()
    r = rand()
    if r < 0.5
        1
    else
        -1
    end
end

function primera_llegada_a_0(L::Int=100)
    pos = 1
    t = 0

    while pos > 0

        pos = pos + brinco1()

        if pos > L
            pos = L
        end

        t += 1
    end

    t
end


#!/usr/bin/env julia

using Luxor, Random

using Test

using Random
Random.seed!(42)

function test_hypotrochoid(fname)
    Drawing(1400, 1400, fname)
    origin()
    background("grey20")
    sethue("antiquewhite")
    setopacity(0.8)
    tiles = Tiler(1200, 1200, 6, 6)
    for (pos, n) in tiles
        gsave()
        setline(1.5)
        translate(pos)

        R = 45; r = n; d = 20
        sethue("antiquewhite")
        poly(hypotrochoid(R, n, d, :stroke, stepby=pi/300, vertices=true), :stroke, close=true)
        text(string("R = $(R), r = $(r), d = $(d)"), halign=:center, O.x, O.y+tiles.tileheight/2)

        # test perimeterlength
        ptslist = hypotrochoid(R, n, d, :stroke, stepby=pi/300, vertices=true)
        polyperimeter(ptslist)

        R = 2n; r = 12; d = 5
        sethue("orange")
        (R == r) && (r += 1)
        p = hypotrochoid(R, r, d, :none, stepby=pi/305, vertices=true)
        polydistances(p) # just testing
        poly(p, :stroke, close=true)
        text(string("R = $(R), r = $(r), d = $(d)"), halign=:center, O.x, O.y+tiles.tileheight/2 + 10)

        # testing portions of polys
        portion = rand()
        setline(.25)
        sethue("magenta")
        section = polyportion(p, portion)
        poly(section, :stroke, close=false)
        sethue("green")
        circle(section[end], 2, :stroke)
        poly(polyremainder(p, portion), :stroke, close=false)

        text(string("portion ", round(portion, digits=2)), halign=:center, O.x, O.y+tiles.tileheight/2 + 30)

        sethue("cyan")
        setline(0.3)
        poly(offsetpoly(p, 15), :stroke, close=true)
        text("offset 15", halign=:center, O.x, O.y+tiles.tileheight/2 + 20)

        grestore()
    end
    @test finish() == true
end

using Random

Random.seed!(42)

fname = "hypotrochoid-test.pdf"
test_hypotrochoid(fname)
println("...finished test: output in $(fname)")


function scale_nodes(nodes::Array{T,1},domain::Array{T,1}) where {T<:AbstractFloat}

  @inbounds for i = 1:length(nodes)
    nodes[i] = domain[2] + (1.0+nodes[i])*(domain[1]-domain[2])/2
  end

  return nodes

end

function scale_nodes(nodes::Array{T,2},domain::Array{T,2}) where {T<:AbstractFloat}

  @inbounds for i = 1:size(nodes,1)
    @inbounds for j = 1:size(nodes,2)
      nodes[i,j] = domain[2,j] + (1.0+nodes[i,j])*(domain[1,j]-domain[2,j])/2
    end
  end

  return nodes

end

function normalize_node(node::T,domain::Array{T,1}) where {T<:AbstractFloat}

  if domain[1] == domain[2]
    norm_node = zero(T)
    return norm_node
  else
    norm_node = 2.0*(node-domain[2])/(domain[1]-domain[2])-1.0
    return norm_node
  end

end

function normalize_node(node::Array{T,1},domain::Array{T,1}) where {T<:AbstractFloat}

  norm_nodes = similar(node)
  @inbounds for i in eachindex(node)
    norm_nodes[i] = normalize_node(node[i],domain)
  end

  return norm_nodes

end


include("../src/Thirian/fractional_shift.jl")

# include("../src/FractionalDelayFilters.jl")
# import .FractionalDelayFilters

using LinearAlgebra
using FFTW
import PyPlot

import Utilities # https://gitlab.com/RoyCCWang/utilities

PyPlot.close("all")
fig_num = 1

PyPlot.matplotlib["rcParams"][:update](["font.size" => 22, "font.family" => "serif"])

# sine for debug.
f = sin
P = LinRange(-pi, pi, 3000)
U = P
f_U = f.(U)

## filter.
s = real.(f_U)

N_thirian = 15
b_thirian = Vector{Float64}(undef, N_thirian)

τ = 0.31
applydelaythirian!(b_thirian, s, τ)

## visualize.

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P, real.(f_U), label = "f")
PyPlot.plot(P, s, label = "s", "--")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("reversible fractional delay filtering")

x_ranges = [ P ]
s_itp, d_s_itp, d2_s_itp = Utilities.setupcubicitp(real.(f_U), x_ranges, 1.0)
s_func = xx->s_itp([xx])

S = FiniteSignalType(real.(f_U), P[1], P[end])

# no shifts.
j = 14 # test this time index.
println("no shift: j ", j)
s0 = real.(f_U)
println("zero: evalsignal( S, j ) - s0[j] = ", evalsignal( S, j ) - s0[j])
println("zero: s0[j] - s_func(P[j] = ", s0[j] - s_func(P[j]))
println()

# integer shifts.
d = 19
applyintegershift!(S, d)

j = d+4 # test this time index.
println("shift: j ", j)
evalsignal( S, j ) # evaluates s0[j-d].
println("zero: s0[j-d] - evalsignal( S, j ) = ", s0[j-d] - evalsignal( S, j ))
println()

# fractional shift, τ ∈ [0, 0.5].
ΔP = P[2]-P[1]

d = 2 # this is the integer delay (positive).
applyintegershift!(S, d)

τ = 0.319 # fractional delay (positive).

resetsignal!(S)
applydelaythirian!(b_thirian, S.s, τ)
evals_S = collect( evalsignal(S, j) for j = 1:length(S.s) )

evals_s_func = collect( s_func( P[j]-ΔP*τ ) for j = 1:length(P) )

println("norm difference between cubic itp and fractional delay: ", norm(evals_s_func - evals_S))
#[evals_S evals_s_func]

# visualize.
PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P, real.(f_U), label = "f")
PyPlot.plot(P, evals_S, label = "evals_S", "--")
PyPlot.plot(P, evals_s_func, label = "evals_s_func", "--")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("cubic itp vs. fractional delay filtering")

### test reversible property of thirian delay filter. see https://lcondat.github.io/publis/condat_icip08_rotations.pdf
τ = 0.21
resetsignal!(S)
applydelaythirian!(b_thirian, S.s, τ)

tmp_s = copy(S.s)

S.s[:] = reverse(S.s)
applydelaythirian!(b_thirian, S.s, τ)
S.s[:] = reverse(S.s)

#reverse!(tmp_s)
applyadvancethirian!(b_thirian, tmp_s, τ) # bake reverse!() into this function.
#reverse!(tmp_s)
println("test applyadvancethirian!: ", norm(tmp_s-S.s))

## timing.
# import BenchmarkTools
# BenchmarkTools.@btime applydelaythirian!(b_thirian, tmp_s, τ)
# BenchmarkTools.@btime applyadvancethirian!(b_thirian, tmp_s, τ)

evals_S = collect( evalsignal(S, j) for j = 1:length(S.s) )
println("reversible test: ", norm(evals_S - s0))
println()

### arb shift, test reversible, frac part ∈ [0, 0.5]
resetsignal!(S)

shift = 1.23
applyshift!(S, shift, b_thirian)

shift = -1.23
applyshift!(S, shift, b_thirian)
evals_S = collect( evalsignal(S, j) for j = 1:length(S.s) )
println("reversible test: ", norm(evals_S - s0))
println()

### arb shift, test reversible, frac part ∈ [0.5, 1]
resetsignal!(S)

shift = 1.9
applyshift!(S, shift, b_thirian)

shift = -1.9
applyshift!(S, shift, b_thirian)
evals_S = collect( evalsignal(S, j) for j = 1:length(S.s) )
println("reversible test: ", norm(evals_S - s0))
println()

######## get shift coordinates from Δp, which is in units of P.
Δp = 2.0
a = P[1]
b = P[end]

resetsignal!(S)

shift = Δp*length(S.s)/(b-a)
applyshift!(S, shift, b_thirian)

# visualize.
evals_S = collect( evalsignal(S, j) for j = 1:length(S.s) )
evals_s_func = collect( s_func( P[j]-ΔP*τ ) for j = 1:length(P) )

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P, real.(f_U), label = "f")
PyPlot.plot(P, evals_S, label = "evals_S", "--")
#PyPlot.plot(P, evals_s_func, label = "evals_s_func", "--")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("fractional shift with Δp = $(Δp)")


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# %%
using SymPy

# %%
function minpoly(x, y, f, g, h, N)
    m = sympy.Poly(f, x).degree()
    n = sympy.Poly(g, y).degree()
    K = sympy.Poly(h, x).degree()
    L = sympy.Poly(h, y).degree()
    λ = symbols("λ[0:$N]")
    z = symbols("z")
    P = sum(λ[begin+k]*z^k for k in 0:N-1) + z^N
    p = P(z => h).expand()
    @time for k in K*N:-1:m
        p = p(x^k => x^k - x^(k-m)*f).expand()
    end
    @time for l in L*N:-1:n
        p = p(y^l => y^l - y^(l-n)*g).expand()
    end
    Q = sympy.Poly(p, x, y)
    C = Q.coeffs()
    @time sol = solve(C, λ)
    P = sum(get(sol, λ[begin+k], λ[begin+k])*z^k for k in 0:N-1) + z^N
    [
        sympy.Poly(f, x)
        sympy.Poly(g, y)
        sympy.Poly(h, x, y)
        sympy.Poly(P, z)
    ]
end

# %%
@vars a p x y
f = x^2 - a
g = y^2 - p
h = x + y
minpoly(x, y, f, g, h, 4)

# %%
@vars a p x y
f = x^2 - a
g = y^2 - p
h = x * y
minpoly(x, y, f, g, h, 2)

# %%
@vars a p x y
f = x^2 - a
g = y^3 - p
h = x + y
minpoly(x, y, f, g, h, 6)

# %%
@vars a p x y
f = x^2 - a
g = y^3 - p
h = x * y
minpoly(x, y, f, g, h, 6)

# %%
@vars a p q x y
f = x^2 - a
g = y^3 + p*y + q
h = x + y
minpoly(x, y, f, g, h, 6)

# %%
@vars a p q x y
f = x^2 - a
g = y^3 + p*y + q
h = x * y
minpoly(x, y, f, g, h, 6)

# %%
@vars a p q r x y
f = x^2 - a
g = y^4 + p*y^2 + q*y + r
h = x + y
minpoly(x, y, f, g, h, 8)

# %%
@vars a p q r x y
f = x^2 - a
g = y^4 + p*y^2 + q*y + r
h = x * y
minpoly(x, y, f, g, h, 8)

# %%
@vars a p q r s x y
f = x^2 - a
g = y^5 + p*y^3 + q*y^2 + r*y + s
h = x + y
minpoly(x, y, f, g, h, 10)

# %%
@vars a p q r s x y
f = x^2 - a
g = y^5 + p*y^3 + q*y^2 + r*y + s
h = (x * y)^2
minpoly(x, y, f, g, h, 5)

# %%
@vars a p x y
f = x^3 - a
g = y^3 - p
h = (x + y)^3
minpoly(x, y, f, g, h, 3)

# %%
@vars a p x y
f = x^3 - a
g = y^3 - p
h = x * y
minpoly(x, y, f, g, h, 3)

# %%
@vars a b p q x y
f = x^3 + a*x + b
g = y^3 + p*y + q
h = x + y
minpoly(x, y, f, g, h, 9)

# %%
@vars a b p q x y
f = x^3 + a*x + b
g = y^3 + p*y + q
h = x * y
minpoly(x, y, f, g, h, 9)

# %%
@vars a p q x y
f = x^3 - a
g = y^4 + p*y + q
h = x + y
minpoly(x, y, f, g, h, 12)

# %%
@vars a p q x y
f = x^3 - a
g = y^4 + p*y + q
h = x * y
minpoly(x, y, f, g, h, 12)

# %%

0023/minimal polynomials of sums and products.jl


using Catlab
using Catlab.Doctrines
using Catlab.Graphics
using Catlab.WiringDiagrams
using Catlab.Programs
import Base.Multimedia: display
import Catlab.Graphics: to_graphviz, LeftToRight
import Base: (==), length, show
using Petri
using SemanticModels.CategoryTheory
import SemanticModels.CategoryTheory: undecorate, ⊔
using SemanticModels.PetriModels
using SemanticModels.PetriCospans
import SemanticModels.PetriCospans: otimes_ipm, compose_pushout
import SemanticModels: model

import Catlab.Doctrines:
    Ob,
    Hom,
    dom,
    codom,
    compose,
    ⋅,
    ∘,
    id,
    oplus,
    otimes,
    ⊗,
    ⊕,
    munit,
    mzero,
    braid,
    dagger,
    dunit,
    dcounit,
    mcopy,
    Δ,
    delete,
    ◊,
    mmerge,
    ∇,
    create,
    □,
    plus,
    zero,
    coplus,
    cozero,
    meet,
    top,
    join,
    bottom

export wd, draw, model, Disease

wd = to_wiring_diagram
draw(d::WiringDiagram) =
    to_graphviz(add_junctions(d), orientation = LeftToRight, labels = true)
draw(d::HomExpr) = draw(wd(d))

model(c::PetriCospan) = left(c.f).d[1].model

println("Making COVID Model")

@present Disease(FreeEpidemiology) begin
    S::Ob
    E::Ob
    I::Ob
    R::Ob
end


"""
    RandomWalks

A module that supports random walks on both integer lattices (of arbitrary finite dimension)
and on graphs.
"""
module RandomWalks
using LinearAlgebra: Diagonal
using LightGraphs: AbstractGraph, adjacency_matrix, degree, grid, nv, vertices
using Distributions: rand, Sampleable
using StochasticProcesses
using StochasticProcesses.MarkovChains
using StochasticProcesses.Samplers

export

# concrete types
MultiRandomWalk, RandomWalk,

# methods
nwalkers, rand, randomwalkmc

# include scripts
include("randomwalks.jl")
include("multirandomwalks.jl")
include("simulaterandomwalks.jl")
end


if !@isdefined TEST
    using Revise
end
using Mitosis

using Mitosis: meancov, flat
using Random, Test, LinearAlgebra, Statistics

Random.seed!(1)
K = 10000 # repetitions
atol = 2sqrt(K)
# State space system
#
# x[0] ∼ N(x0, P0)
# x[k] = Φx[k−1] + w[k],    w[k] ∼ N(0, Q)
# y[k] = Hx[k] + v[k],    v[k] ∼ N(0, R)

ξ0 = [1., 0.]
P0 = Matrix(1.0I, 2, 2)
prior = Gaussian{(:μ,:Σ)}(ξ0, P0)

Φ = [0.8 0.5; -0.1 0.8]
β = [0.1, 0.2]
Q = [0.2 0.0; 0.0 1.0]

Φ̃ = [1.0 0.5; -0.1 1.0]
β̃ = [0.2, 0.0]
Q̃ = [0.3 0.0; 0.0 2.0]

yshadow = [0.0]
H = [1.0 0.0]
R = Matrix(1.0I, 1, 1)
INoise = Gaussian(μ=zero(ξ0), Σ=Q)
Noise = Gaussian(μ=zero(yshadow), Σ=R)

@test AffineMap(Φ, β)(ξ0) == Φ*ξ0 + β
@test ConstantMap(β)(ξ0) == β

priortransition = kernel(Gaussian; μ=LinearMap(I(2)), Σ=ConstantMap(P0))
@test priortransition(ξ0) ≈ prior

transition = kernel(Gaussian; μ=AffineMap(Φ, β), Σ=ConstantMap(Q))
transition2 = kernel(Gaussian; μ=AffineMap(Φ, β), Σ=ConstantMap(Q))
transition1 = kernel(Gaussian; μ=LinearMap(Φ), Σ=ConstantMap(Q))
transition1 = transition2
transitiontilde = kernel(Gaussian; μ=AffineMap(Φ̃, β̃), Σ=ConstantMap(Q̃))

observation = kernel(Gaussian; μ=LinearMap(H), Σ=ConstantMap(R))

@test 10/sqrt(K) > norm(Φ*ξ0 + β - mean(rand(transition(ξ0)) for x in 1:K))

@test transition.ops isa NamedTuple{(:μ, :Σ)}

x0 = rand(prior)
y0 = rand(observation(x0))

x1 = rand(transition(x0))
y1 = rand(observation(x1))

x2 = rand(transition(x1))
y2 = rand(observation(x2))

p0 = Gaussian(μ=ξ0, Σ=P0)

@test mean(AffineMap(Φ, β)(p0)) == Φ*ξ0 + β
@test cov(AffineMap(Φ, β)(p0)) == Φ*P0*Φ'

q0 = observation(p0)
@test cov(q0) == H*P0*H' + R
p1 = transition2(p0)
@test cov(p1) == Φ*P0*Φ' + Q

q1 = observation(p1)
p2 = transition2(p1)
@test cov(p2) == Φ*(Φ*P0*Φ' + Q)*Φ' + Q
q2 = observation(p2)

# Check that this is all correct:

# Write down joint distribution of x's and y's by hand... (yes I am fine ;-))
# Define mean and covariance of the flattened vector of states and observations [x0 x1 x2 y0 y1 y2]
#μ_ = [1.0, 0.0, 0.8, -0.1, 0.59, -0.16, 1.0, 0.8, 0.59]
μ_ = [1.0, 0.0, 0.9, 0.1, 0.87, 0.19, 1.0, 0.9, 0.87] # with beta
μ = [ ξ0; Φ*ξ0 + β; Φ*(Φ*ξ0 + β)+β; H*ξ0; H*(Φ*ξ0 + β); H*(Φ*(Φ*ξ0 + β)+β); ]
@test μ_ ≈ μ

P0L = cholesky(P0).L
QL = cholesky(Q).L
RL = cholesky(R).L
Z(m,n) = zeros(m,n)
# Lower cholesky factor from "innovation form"
L = [   [     P0L  Z(2,2) Z(2,2)  Z(2,1)  Z(2,1)  Z(2,1)]
    [    Φ*P0L     QL Z(2,2)  Z(2,1)  Z(2,1)  Z(2,1)]
    [  Φ*Φ*P0L   Φ*QL     QL  Z(2,1)  Z(2,1)  Z(2,1)]
    [    H*P0L Z(1,2) Z(1,2)     RL  0I  0I]
    [  H*Φ*P0L   H*QL Z(1,2)     0I  RL  0I]
    [H*Φ*Φ*P0L H*Φ*QL   H*QL     0I  0I  RL]
]

Σ = L*L'

@testset "Uncertainty propagation" begin
    @test μ ≈ flat(mean.([p0, p1, p2, q0, q1, q2]))
    @test diag(Σ) ≈ flat(diag.(cov.([p0, p1, p2, q0, q1, q2])))
end

# Compute the conditional distribution of vector `x0` given data.
π0 = Mitosis.conditional(Gaussian(;μ=μ, Σ=Σ), 1:2, 7:9, vcat(y0, y1, y2))

# Compute the conditional distribution of vector `x0` given data.
π2 = Mitosis.conditional(Gaussian(;μ=μ, Σ=Σ), 5:6, 7:9, vcat(y0, y1, y2))

@testset "Kalman filter" begin
    p0f = correct(p0, observation, y0)[1]
    p1p = transition2(p0f)
    p1f = correct(p1p, observation, y1)[1]
    p2p = transition2(p1f)
    p2f = correct(p2p, observation, y2)[1]

    @test mean(p2f) ≈ mean(π2)
    @test cov(p2f) ≈ cov(π2)
end

@testset "right' linear gaussian case" begin
    p0 = backward(BF(), transition1, p1)[2]

    ν, P = meancov(p0)

    @test mean(transition1(ν)) ≈ mean(p1)
    @test Φ*cov(p0)*Φ' ≈ (cov(p1) + Q)

    H = inv(cov(p1))
    q1 = WGaussian(F=H*mean(p1), Γ=H, c=0.0)
    @test logdensity(q1, x0) ≈ logdensity(p1, x0)

    q0 = backward(BFFG(), transition1, q1)[2]

    ν0 = q0.Γ\q0.F
    P0 = inv(q0.Γ)
    @test mean(transition1(ν0)) ≈ mean(p1)
    @test Φ*P0*Φ' ≈ (cov(p1) + Q)

    pp = transition1(x0)⊕Gaussian(μ=0*p1.μ, Σ=p1.Σ)
    @test logdensity(q0, x0) ≈ logdensity(pp, p1.μ)

    @test 0.0 ≈ -logdet(pp.Σ)/2 - (q0.c + logdet(q0.Γ)/2) atol=1e-10
end

#=
# Next we consider the following model, in Soss notation
m = @model ξ0 begin
           x0 ~ MvNormal(ξ0, P0) # priortransition
           y0 ~ MvNormal(H*x0, R) # observation
           x1 ~ MvNormal(Φ*x0, Q) # transition2
           y1 ~ MvNormal(H*x1, R) # observation
           x2 ~ MvNormal(Φ*x1, Q) # transition2
           return y0, y1, x2
end
=#

@testset "Backward filter with fusion" begin
    global evi
    # forward model
    x0 = rand(priortransition(ξ0))
    y0 = rand(observation(x0))
    x1 = rand(transition(x0))
    y1 = rand(observation(x1))
    x2 = rand(transition(x1))

    # run backward filter
    m1a, p1a = backwardfilter(observation, y1; unfused=true)
    m1b, p1b = backwardfilter(transition2, x2; unfused=true)
    m1, p1 = fuse(p1a, p1b)
    m0a, p0a = backwardfilter(observation, y0; unfused=true)
    m0b, p0b = backwardfilter(transition2, p1; unfused=true)
    m0, p0 = fuse(p0a, p0b)
    m, evi = backwardfilter(priortransition, p0)

    # as byproduct this just computed the model evidence as function
    # of ξ0
    @test logdensity(evi, ξ0) ≈ logdensity(Gaussian(;μ=μ[5:8], Σ=Σ[5:8,5:8]), vcat(x2, y0, y1))

    # Alternative: p0_ is the conditional distribution of x0
    prior_ = WGaussian{(:F,:Γ,:c)}(P0\ξ0, inv(P0), 0.0)
    m0_, p0_ = fuse(p0a, p0b, prior_)
    π0 = Mitosis.conditional(Gaussian(;μ=μ, Σ=Σ), 1:2, 5:8, vcat(x2, y0, y1))
    @test mean(π0) ≈ mean(p0_)
    @test cov(π0) ≈ cov(p0_)

    # alternative with fusion
    _, evi_ = backwardfilter(priortransition, p0; unfused=true)
    _, evi2 = fuse(evi_)
    @test evi2 ≈ evi

    # run forward marginal smoother
    kᵒ = forward(BFFG(), priortransition, m)
    p0ᵒ = kᵒ(ξ0)
    # skip copy
    k0ᵒ = forward(BFFG(), transition2, m0b)
    p1ᵒ = k0ᵒ(p0ᵒ)

    # first step gives the conditional marginal of x0
    @test mean(π0) ≈ mean(p0ᵒ)
    @test cov(π0) ≈ cov(p0ᵒ)

    # second step gives the conditional marginal of x1
    π1 = Mitosis.conditional(Gaussian(;μ=μ, Σ=Σ), 3:4, 5:8, vcat(x2, y0, y1))
    @test mean(π1) ≈ mean(p1ᵒ)
    @test cov(π1) ≈ cov(p1ᵒ)

    # some more tests
    @test logdensity(backwardfilter(transition2, x2)[2], x1) ≈ logdensity(transition2(x1), x2)
    @test logdensity(fuse(backwardfilter(transition2, x2; unfused=true)[2])[2], x1) ≈ logdensity(transition2(x1), x2)

    _, q = fuse(backwardfilter(transition2, x2; unfused=true)[2], backwardfilter(transition2, x2; unfused=true)[2])
    @test logdensity(q, x1) ≈ 2logdensity(transition2(x1), x2)
    @test logdensity(p1, x1) ≈ logdensity(observation(x1), y1) + logdensity(transition2(x1), x2)

end

@testset "Change of measure" begin
    # forward model
    x0 = rand(priortransition(ξ0))
    y0 = rand(observation(x0))
    x1 = rand(transition(x0))
    y1 = rand(observation(x1))
    x2 = rand(transition(x1))

    # run backward filter
    m1a, p1a = backwardfilter(observation, y1; unfused=true)
    m1b, p1b = backwardfilter(transition2, x2; unfused=true)
    m1, p1 = fuse(p1a, p1b)
    m0a, p0a = backwardfilter(observation, y0; unfused=true)
    m0b, p0b = backwardfilter(transitiontilde, p1; unfused=true)
    m0, p0 = fuse(p0a, p0b)
    m, evi2 = backwardfilter(priortransition, p0)

    # as byproduct this just computed the model evidence as function
    # of ξ0
    #@test logdensity(evi, ξ0) ≈ logdensity(Gaussian(;μ=μ[5:8], Σ=Σ[5:8,5:8]), vcat(x2, y0, y1))

    # run forward marginal smoother
    kᵒ = forward(BFFG(), priortransition, m)
    p0ᵒ = kᵒ(ξ0)
    # skip copy

    Y = Any[]
    for i in 1:K
        x = Mitosis.rand(p0ᵒ)
        @assert x.ll == 0
        p1ᵒ = forward(BFFG(), transition2, transitiontilde, m0b, x[])
        y = Mitosis.rand(p1ᵒ)
        @assert y.ll != 0
        push!(Y, (y[], y.ll))
    end

    # second step gives the conditional marginal of x1
    π1 = Mitosis.conditional(Gaussian(;μ=μ, Σ=Σ), 3:4, 5:8, vcat(x2, y0, y1))
    @test mean(π1) ≈ mean(first.(Y) .* exp.(last.(Y)))*exp(logdensity(evi2, ξ0)-logdensity(evi, ξ0)) atol=atol

end


# # Supervised 3D Mesh Reconstruction
#
# ## Problem Description:
# We are given an initial source shape (sphere in this case) and we want to deform
# this source shape to fit target shape (dolphin in this case). For this
# demonstration, we will be using Triangle Mesh for the representation of source
# and target shape.
#
# Triangle Mesh has two main components, vertices and faces. Deformation of source
# shape to fit target shape can be achieved by offsetting source's vertices to fit
# target surface. Also, the number of vertices and faces is not equal in source
# and target shape.
#
# !!!note
#     For visualization purpose we will require to install Makie and
#     compatible backend (GLMakie or WGLMakie). To install it simply run
#     `] add AbstractPlotting GLMakie` in the julia prompt.

using Flux3D, Zygote, Flux, FileIO, Statistics
using Makie, GLMakie, WGLMakie, GeometricalPredicates, GeometryBasics

GLMakie.activate!()
set_theme!(show_axis = false)

# ## Downloading obj file of sphere and dolphin

download(
    "https://github.com/nirmalsuthar/public_files/raw/master/dolphin.obj",
    "dolphin.obj",
)
download("https://github.com/nirmalsuthar/public_files/raw/master/sphere.obj", "sphere.obj")

# ## Loading Triangle Mesh
# Triangle Mesh is handled by TriMesh in Flux3D. TriMesh also supports batched
# format, namely padded, packed and list, which allow us to use fast batched
# operations. We can load TriMesh with load_trimesh function
# (supports `obj`, `stl`, `ply`, `off` and `2DM`)

dolphin = load_trimesh("dolphin.obj")
src = load_trimesh("sphere.obj")

# ## Preprocessing data
# Preprocessing tgt (dolphin), such that its mean is zero and also scale it
# according to the bounding box of src (sphere), So that src can converge at
# greater speed.

tgt = deepcopy(dolphin)
verts = get_verts_packed(tgt)
center = mean(verts, dims = 2)
verts = verts .- center
scale = maximum(abs.(verts))
verts = verts ./ scale
tgt._verts_packed = verts

# ## Visualizing TriMesh
# We will use `visualize` function for visualizing TriMesh. This function uses
# **Makie** for plotting. In fact, we can also visualize PointCloud using this
# function, which makes this function handy dealing with different 3D format.
vbox(visualize(src), visualize(tgt))

# ```@raw html
# <p align="center">
#     <img width=480 height=270 src="../../assets/fitmesh_initial.png">
# </p>
# ```

# ## Defining loss objective
# Starting from the src mesh, we will deform src mesh by offsetting its vertices
# (by offset array), such that new deformed mesh is close to target mesh.
# Therefore, our loss function will optimize the offset array. We will be using
# the following metrics to define loss objective:
#
# * `chamfer_distance` - the distance between the deformed mesh and target mesh, which is calculated by taking randomly 5000 points from the surface of each mesh and calculating chamfer_distance between these two pointcloud.
# * `laplacian_loss` - also known as Laplacian smoothing will act as a regularizer.
# * `edge_loss` - this will minimize edges length in deformed mesh, also act as a regularizer.

function loss_dolphin(x::AbstractArray, src::TriMesh, tgt::TriMesh)
    src = Flux3D.offset(src, x)
    loss1 = chamfer_distance(src, tgt, 5000)
    loss2 = laplacian_loss(src)
    loss3 = edge_loss(src)
    return loss1 + 0.1 * loss2 + loss3
end
# ## Defining learning rate and optimizer

lr = 1.0
opt = Flux.Optimise.Momentum(lr, 0.9)
_offset = zeros(Float32, size(get_verts_packed(src))...)

# ## Using GPU for fast training [**Optional**]
# We can convert the TriMesh structure to GPU or CPU using`gpu` and `cpu`
# function which is exactly the same syntax as Flux.

tgt = tgt |> gpu
src = src |> gpu
_offset = _offset |> gpu

# ## Optimizing the offset array
# We first initialize offset array as zeros, hence deformed mesh is equivalent to
# src mesh (sphere). Next, we calculate loss using this offset array and we
# compute derivatives wrt. offset array and finally optimize the array.

@info("Training...")
θ = Zygote.Params([_offset])
for itr = 1:2001
    gs = gradient(θ) do
        loss_dolphin(_offset, src, tgt)
    end
    Flux.update!(opt, _offset, gs[_offset])
    if (itr % 10 == 1)
        loss = loss_dolphin(_offset, src, tgt)
        @show itr, loss
        save("src_$(itr).png", visualize(Flux3D.offset(src, _offset)))
    end
end

anim = @animate for i ∈ 1:8
    Plots.plot(load("src_$(1+250*(i-1)).png"), showaxis = false)
end
gif(anim, "src_deform.gif", fps = 2)

# ```@raw html
# <p align="center">

#     <img width=256 height=256 src="../../assets/fitmesh_anim.gif">
# </p>
# ```
# ## Postprocessing the predicted mesh
# We create a new TriMesh by offsetting src by final offset array and scale up
# the final_mesh by the same scaling factor we scale down tgt, such that
# final_mesh has similar bounding box as dolphin mesh.

final_mesh = Flux3D.offset(src, _offset)
final_mesh = Flux3D.scale!(final_mesh, scale)

# ## Saving the final_mesh
# Flux3D provide IO function `save_trimesh`  to save TriMesh
# (supports `obj`, `stl`, `ply`, `off` and `2DM`)

save_trimesh("results/final_mesh.off", final_mesh)
save("results/final_mesh.png", visualize(final_mesh))
Flux3D.AbstractPlotting.vbox(visualize(final_mesh), visualize(dolphin))

# ```@raw html
# <p align="center">
#     <img width=480 height=270 src="../../assets/fitmesh_final.png">
# </p>
# ```

# ## Finally..
# * Look into the other examples in `examples/`
# * Read more the TriMesh in TriMesh section in documentation and
# Metrics/Transforms section for manipulating TriMesh and computing standard
# metrics.


# ---
# title: 1512. Number of Good Pairs
# id: problem1512
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array, Hash Table, Math
# link: <https://leetcode.com/problems/number-of-good-pairs/description/>
# hidden: true
# ---
#
# Given an array of integers `nums`.
#
# A pair `(i,j)` is called _good_ if `nums[i]` == `nums[j]` and `i` < `j`.
#
# Return the number of _good_ pairs.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3,1,1,3]
#     Output: 4
#     Explanation: There are 4 good pairs (0,3), (0,4), (3,4), (2,5) 0-indexed.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,1,1,1]
#     Output: 6
#     Explanation: Each pair in the array are _good_.
#
#
# **Example 3:**
#
#
#
#     Input: nums = [1,2,3]
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 100`
#   * `1 <= nums[i] <= 100`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


todisplay = VERSION <= VersionNumber("1.2") ? true : false

using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using CAGD
if todisplay
    using ViewerGL
    GL = ViewerGL
end

function viewExplode(model; complete = false)
    V,CVs,FVs,EVs = Lar.pols2tria(model.G, model.T[1], model.T[2], model.T[3])

    GL.VIEW(GL.GLExplode(V,FVs,1.5,1.5,1.5,99,1));
    GL.VIEW(GL.GLExplode(V,EVs,1.5,1.5,1.5,99,1));
    GL.VIEW(GL.GLExplode(V,CVs[2 - complete:end],5,5,5,99,0.5));
end

function displayModel(model, exp = 1.5)
    V = model.G
    EV = Lar.cop2lar(model.T[1])
    FE = Lar.cop2lar(model.T[2])
    CF = Lar.cop2lar(model.T[3])
    FV = Lar.cop2lar(map(x -> Int8(x/2), abs.(model.T[2]) * abs.(model.T[1])))
    CV = []
    triangulated_faces = Lar.triangulate(convert(Lar.Points, V'), [model.T[1], model.T[2]])
    FVs = convert(Array{Lar.Cells}, triangulated_faces)
    EVs = Lar.FV2EVs(model.T[1], model.T[2])

    GL.VIEW([
        GL.GLPol(V,EV, GL.COLORS[1])
        GL.GLFrame
    ]);
    GL.VIEW(GL.GLExplode(V,FVs,exp,exp,exp,99));
    GL.VIEW(GL.GLExplode(V,EVs,exp,exp,exp,99,1));
end

function viewNum(model, numsize = 0.25)
	V = model.G
	EV = Lar.cop2lar(model.T[1])
	FV = Lar.cop2lar(map(x -> floor(Int8, x / 2), abs.(model.T[2]) * abs.(model.T[1])))
	VV = [[k] for k=1:size(V,2)]
	GL.VIEW([
        GL.numbering(numsize)((V,[VV, EV, FV]))
        GL.GLFrame
    ])
end

function viewNumEdge(model::CAGD.Model, τ::Int, numsize=0.25)
    V  = model.G
    EV = Lar.cop2lar(model.T[1])
    VV = [[k] for k=1:size(V,2)]
    faces = model.T[2][:, τ].nzind
    EVidx = unique(vcat([model.T[2][f, :].nzind for f in faces]...))
    EVs = Lar.cop2lar(model.T[1][EVidx, :])
    FVs = Lar.cop2lar(map(x -> floor(Int8, x / 2), abs.(model.T[2][faces, :]) * abs.(model.T[1])))
    VV = [[k] for k in unique(vcat(EVs...))]
    GL.VIEW([
        GL.GLPol(V, EV, GL.COLORS[1])
        GL.numbering(numsize)((V, [VV, EVs, FVs]))
        GL.GLFrame
    ])
end

function viewNumFace(model::CAGD.Model, τ::Int, numsize=0.25)
    V  = model.G
    EV = Lar.cop2lar(model.T[1])
    VV = [[k] for k=1:size(V,2)]
    EVidx = model.T[2][τ, :].nzind
    EVs = Lar.cop2lar(model.T[1][EVidx, :])
    FVs = Lar.cop2lar(map(x -> floor(Int8, x / 2), abs.(model.T[2][τ:τ, :]) * abs.(model.T[1])))
    VV = [[k] for k in unique(vcat(EVs...))]
    GL.VIEW([
        GL.GLPol(V, EV, GL.COLORS[1])
        GL.numbering(numsize)((V, [VV, EVs, FVs]))
        GL.GLFrame
    ])
end


export tarjan

mutable struct Node
    id::Int
    index::Int
    lowlink::Int
    onstack::Bool
    tails::Vector{Node}
end

doc"""
    tarjan(heads::Vector{Int})

Computes Tarjan's algorithm for finding strongly connected components (cycles) of a graph.

# Example
```julia
heads = [3, 4, 2, 1, 4, 5]
tarjan(heads)
```
"""
function tarjan(heads::Vector{Int})
    for h in heads
        1 <= h <= length(heads) || throw("Invalid input: $h")
    end

    nodes = [Node(i,0,0,false,Node[]) for i=1:length(heads)]
    for i = 1:length(heads)
        push!(nodes[heads[i]].tails, nodes[i])
    end
    stack = Node[]
    cycles = Vector{Int}[]
    index = 1
    function strongconnect!(node::Node)
        node.index = index
        node.lowlink = index
        index += 1
        push!(stack, node)
        node.onstack = true

        for t in node.tails
            if t.index == 0
                # Successor t has not yet been visited
                strongconnect!(t)
                node.lowlink = min(node.lowlink, t.lowlink)
            elseif t.onstack
                # Successor t is in stack and hence in the current SCC
                node.lowlink = min(node.lowlink, t.index)
            end
        end

        # If node is a root, pop the stack and generate an SCC
        if node.lowlink == node.index
            cycle = Int[]
            while true
                t = pop!(stack)
                t.onstack = false
                push!(cycle, t.id)
                t == node && break
            end
            length(cycle) > 1 && push!(cycles,sort!(cycle))
        end
    end

    for node in nodes
        node.index == 0 && strongconnect!(node)
    end
    cycles
end


using Test
import DifferentiableBackwardEuler
import DPFEHM
import Zygote

doplot = false
if doplot == true
	import PyPlot
end

#Theis solution
function W(u)
	if u <= 1
		return -log(u) + -0.57721566 + 0.99999193u^1 + -0.24991055u^2 + 0.05519968u^3 + -0.00976004u^4 + 0.00107857u^5
	else
		return (u^2 + 2.334733u + 0.250621) / (u^2 + 3.330657u + 1.681534) * exp(-u) / u
	end
end
function theis(t, r, T, S, Q)
	return Q * W(r^2 * S / (4 * T * t)) / (4 * pi * T)
end
#utitlity for checking gradients
function checkgradientquickly(f, x0, gradf, n; delta::Float64=1e-8, kwargs...)
	indicestocheck = sort(collect(1:length(x0)), by=i->abs(gradf[i]), rev=true)[1:n]
	#indicestocheck = [indicestocheck; rand(1:length(x0), n)]
	f0 = f(x0)
	for i in indicestocheck
		x = copy(x0)
		x[i] += delta
		fval = f(x)
		grad_f_i = (fval - f0) / delta
		@test isapprox(gradf[i], grad_f_i; kwargs...)
	end
end

#test Theis solution against groundwater model
steadyhead = 1e3
sidelength = 50.0
thickness = 10.0
n = 101
ns = [n, n]
coords, neighbors, areasoverlengths, volumes = DPFEHM.regulargrid2d([-sidelength, -sidelength], [sidelength, sidelength], ns, thickness)
k = 1e-5
Ks = fill(k, length(neighbors))
T = thickness * k
Q = 1e-3
Ss = 0.1
specificstorage = fill(Ss, size(coords, 2))
S = Ss * thickness
Qs = zeros(size(coords, 2))
Qs[ns[2] * (div(ns[1] + 1, 2) - 1) + div(ns[2] + 1, 2)] = -Q#put a fluid source in the middle
dirichletnodes = Int[]
dirichleths = zeros(size(coords, 2))
for i = 1:size(coords, 2)
	if abs(coords[1, i]) == sidelength || abs(coords[2, i]) == sidelength
		push!(dirichletnodes, i)
		dirichleths[i] = steadyhead
	end
end
isfreenode, nodei2freenodei, freenodei2nodei = DPFEHM.getfreenodes(prod(ns), dirichletnodes)
function unpack(p)
	@assert length(p) == length(neighbors)
	Ks = p[1:length(neighbors)]
	return Ks
end
function f_gw(u, p, t)
	Ks = unpack(p)
	return -DPFEHM.groundwater_residuals(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end
function f_gw_u(u, p, t)
	Ks = unpack(p)
	return -DPFEHM.groundwater_h(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end
function f_gw_p(u, p, t)
	Ks = unpack(p)
	return -DPFEHM.groundwater_Ks(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end
f_gw_t(u, p, t) = zeros(length(u))
t0 = 0.0
tfinal = 60 * 60 * 24 * 1e1
h0 = fill(steadyhead, size(coords, 2) - length(dirichletnodes))
p = Ks
print("groundwater forward")
@time h_gw = DifferentiableBackwardEuler.steps_diffeq(h0, f_gw, f_gw_u, f_gw_p, f_gw_t, p, t0, tfinal; abstol=1e-6, reltol=1e-6)
dpfehm_solution = map(h->DPFEHM.addboundaryconditions(h, dirichletnodes, dirichleths, isfreenode, nodei2freenodei), h_gw.u)
r0 = 0.1
goodnodes = collect(filter(i->coords[2, i] == 0 && coords[1, i] > r0 && coords[1, i] <= sidelength / 2, 1:size(coords, 2)))
rs = coords[1, goodnodes]
theis_drawdowns = theis.(h_gw.t[end], rs, T, S, Q)
gw_drawdowns = -dpfehm_solution[end][goodnodes] .+ steadyhead
if doplot
	fig, ax = PyPlot.subplots()
	ax.plot(rs, theis_drawdowns, "r.", ms=20, label="Theis")
	ax.plot(rs, gw_drawdowns, "k", linewidth=3, label="DPFEHM groundwater")
	ax.set_xlabel("x [m]")
	ax.set_ylabel("drawdown [m]")
	ax.legend()
	display(fig)
	println()
	PyPlot.close(fig)
end
@test isapprox(theis_drawdowns, gw_drawdowns; atol=1e-1)
g_gw(p) = DifferentiableBackwardEuler.steps(h0, f_gw, f_gw_u, f_gw_p, f_gw_t, p, t0, tfinal; abstol=1e-6, reltol=1e-6)[nodei2freenodei[goodnodes[round(Int, 0.25 * end)]], end]
print("groundwater gradient")
@time grad_gw_zygote = Zygote.gradient(g_gw, p)[1]
checkgradientquickly(g_gw, p, grad_gw_zygote, 3; delta=1e-8, rtol=1e-1)

#test Theis solution against richards equation model
coords_richards = vcat(coords, zeros(size(coords, 2))')
alphas = fill(0.5, length(neighbors))
Ns = fill(1.25, length(neighbors))
function f_richards(u, p, t)
	Ks = unpack(p)
	return DPFEHM.richards_residuals(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, coords_richards, alphas, Ns, Qs, specificstorage, volumes)
end
function f_richards_u(u, p, t)
	Ks = unpack(p)
	return DPFEHM.richards_psi(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, coords_richards, alphas, Ns, Qs, specificstorage, volumes)
end
function f_richards_p(u, p, t)
	Ks = unpack(p)
	return DPFEHM.richards_Ks(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, coords_richards, alphas, Ns, Qs, specificstorage, volumes)
end
f_richards_t(u, p, t) = zeros(length(u))
print("richards forward")
h0 = fill(steadyhead, size(coords, 2))
@time h_richards = DifferentiableBackwardEuler.steps_diffeq(h0, f_richards, f_richards_u, f_richards_p, f_richards_t, p, t0, tfinal; abstol=1e-6, reltol=1e-6)
richards_drawdowns = -h_richards[end][goodnodes] .+ steadyhead
if doplot
	fig, ax = PyPlot.subplots()
	ax.plot(rs, theis_drawdowns, "r.", ms=20, label="Theis")
	ax.plot(rs, richards_drawdowns, "k", linewidth=3, label="DPFEHM richards")
	ax.set_xlabel("x [m]")
	ax.set_ylabel("drawdown [m]")
	ax.legend()
	display(fig)
	println()
	PyPlot.close(fig)
end
@test isapprox(theis_drawdowns, richards_drawdowns; atol=1e-1)
#@test isapprox(h_richards[isfreenode, :], h_gw[:, :])#make sure richards and groundwater are giving the same thing -- commenting this out for now because the recent changes in gw mean they shouldn't be exactly the same -- uncomment when those changes get propagated into richards as well
g_richards(p) = DifferentiableBackwardEuler.steps(h0, f_richards, f_richards_u, f_richards_p, f_richards_t, p, t0, tfinal; abstol=1e-6, reltol=1e-6)[goodnodes[round(Int, 0.25 * end)], end]
print("richards gradient")
@time grad_richards_zygote = Zygote.gradient(g_richards, p)[1]
checkgradientquickly(g_richards, p, grad_richards_zygote, 3; delta=1e-8, rtol=1e-1)


n = rand(1:10)
Z = Companion(randn(n))

#Special properties
@test Matrix(inv(Z)) ≈ inv(Matrix(Z))

#Matvec product
b = randn(n)
@test Z*b ≈ Matrix(Z)*b

m = rand(1:10)
A = randn(m, n)
@test A*Z ≈ A*Matrix(Z)


import Distributed: @distributed

"""Returns an array of subsequent values from the Bernoulli map with parameter β, starting
from x over N time steps."""
function bernoulli_map(x::T,β::T,N::Int) where T
    oone = one(T)                   # 1 in format T
    xout = Array{T,1}(undef,N+1)    # preallocate
    xout[1] = x                     # store initial condition
    @inbounds for i in 2:N+1
        x = β*x                     # actual Bernoulli map
        x = x >= oone ? x-oone : x  # x mod 1
        xout[i] = x                 # store iteration
    end
    return xout
end

"""Find the Bernoulli orbit length that is reached starting from x with parameter β.
Recursive function that starts testing for orbits over period length N, but enlarges
that testing period if no orbit was found."""
function orbit_length(  x::T,            # initial condition
                        β::T,            # Bernoulli parameter
                        N::Int) where T  # max period length
    oone = one(T)                        # 1 in format T
    x0 = x                          # new initial condition
    n = 0                           # orbit length (0 = not found yet)
    j = 0                           # iteration counter

    while n == 0 && j < N           # stop when orbit is fousnd (n >0) or when max period is reached
        j += 1
        x = β*x                     # actual Bernoulli map
        x = x >= oone ? x-oone : x  # x mod 1
        n = x0 == x ? j : 0         # check for periodicity
    end

    # recursive call on orbit_length
    if n == 0         # in case no orbit was found in N steps, repeat with 5N steps
        n,x = orbit_length(x,β,5N)
    else
        return n,x    # return orbit length and last x that's on the orbit
    end
end

"""Returns the minimum x on the orbit of length N, starting from x (which is assumed to be on the orbit."""
function orbit_minimum( x::T,               # initial condition on the orbit
                        β::T,               # Bernoulli parameter
                        N::Int) where T     # the period length of the orbit
    oone = one(T)                   # 1 in format T
    xmin = x                        # store the orbit's minimum in xmin
    for i in 2:N+1
        x = β*x
        x = x >= oone ? x-oone : x  # x mod 1
        xmin = x < xmin ? x : xmin  # always store the smallest x on the orbit
    end
    return xmin
end

"""Returns both the length and minimum of the orbit reached from initial condition x."""
function orbit_length_minimum(  x0::T,              # initial condition on the orbit
                                β::Real;            # Bernoulli parameter
                                basin::Symbol=:x0_proportional,
                                N::Int=10) where T  # initial period length to test for
    o,x = orbit_length(x0,β,N)
    x = orbit_minimum(x,β,o)

    # for starting from all floats in [0,1) use x0_proportional basins
    # for starting from random uniform in [0,1) use uniform basins
    b = basin == :x0_proportional ? Float64(eps(x0)) : 1.0
    return Orbit(β,o,x,b)
end

"""For a given number format T, parameter β and n initial conditions, find orbits of the
Bernoulli map and test for their uniqueness."""
function find_orbits(   ::Type{T},                      # Number format
                        β::Real,                        # Bernoulli parameter
                        n::Int=10000;                   # n initial conditions
                        verbose::Bool=false) where T    # feedback on orbits found and time?

    # pre-allocate empty array of orbits
    orbits = Orbit[]
    uint_type = eval(Symbol("UInt"*repr(sizeof(T)*8)))
    Tβ = T(β)

    tic = time()
    orbits = @distributed (reduce_orbits) for i in 0:n-1            # for n ICs calculate orbit lengths & x
        orbit_length_minimum(reinterpret(T,uint_type(i)),β,basin=:x0_proportional)
    end

    orbits = reduce_orbits(orbits)      # convert Orbit to Vector{Orbit} of length 1 otherwise pass
    sort!(orbits)                       # from shortest to longest orbit
    normalise_basins!(orbits)

    if verbose
        toc = time()
        time_elapsed = readable_secs(toc-tic)
        println("Found $(length(orbits)) orbits in $time_elapsed.")
    end
    return orbits
end

"""For a given number format T, parameter β and n initial conditions, find orbits of the
Bernoulli map and test for their uniqueness."""
function find_orbits_rand(  ::Type{T},                      # Number format
                            β::Real,                        # Bernoulli parameter
                            n::Int=10000;                   # n initial conditions
                            verbose::Bool=false) where T    #

    # pre-allocate empty array of orbits
    orbits = Orbit[]
    Tβ = T(β)

    tic = time()
    orbits = @distributed (reduce_orbits) for i in 1:n            # for n ICs calculate orbit lengths & x
        orbit_length_minimum(T(rand()),β,basin=:uniform)
    end

    orbits = reduce_orbits(orbits)      # convert Orbit to Vector{Orbit} of length 1 otherwise pass
    sort!(orbits)                       # from shortest to longest orbit
    normalise_basins!(orbits)

    if verbose
        toc = time()
        time_elapsed = readable_secs(toc-tic)
        println("Found $(length(orbits)) orbits in $time_elapsed.")
    end
    return orbits
end

import StatsBase: fit, Histogram, mean, pweights

"""Convert a vector of orbits to the invariant measure,
a histogram with `nbins` bins in [0,1)."""
function invariant_measure(orbits::Vector{Orbit{T}},nbins::Int=40) where T

    norbits = length(orbits)
    histograms = fill(0.0,norbits,nbins)
    bin_width = 1/nbins
    bin_edges = collect(0:bin_width:1)

    for (i,orbit) in enumerate(orbits)
        X = bernoulli_map(orbit.min,orbit.β,orbit.length-1)
        X = Float64.(X)         # convert to Float64 to avoid type promoting in fit
        H = fit(Histogram,X,bin_edges).weights
        histograms[i,:] = H / sum(H) / bin_width
    end

    # weighted means
    basins = [orbit.b for orbit in orbits]
    return mean(histograms,pweights(basins),dims=1)[1,:]
end


using DifferentialEquations
using DiffEqSensitivity
using LinearAlgebra
using Zygote
using Statistics
using Distributions
using Plots
using Random

# set random seed for reproducible results
Random.seed!(1)

## ----- setup true solution ------------------------------------

# define function for RHS of lorenz system
function lorenz!(du, u, p, t)
    x,y,z = u
    σ,ρ,β = p

    du[1] = σ*(y - x)
    du[2] = x*(ρ - z) - y
    du[3] = x*y - β*z
end

θ = [10.0, 28.0, (8/3)]  # true parameter values
dt = 0.01  # time step for simulation (this may be unnecessary)
tend = 10.0  # max time for total simulation
ts_sim = (0.0, tend)
u0True = [1.0, 1.0, 1.0]  # true initial condition

# integrate to obtain the true solution
prob_true = ODEProblem(lorenz!, u0True, ts_sim, θ, saveat=dt)
sol_true = solve(prob_true)

# visualize the data
# px = plot(sol_true, vars=(0,1), color=:black, alpha=0.75, linewidth=2, label="true system")
# xlabel!("t")
# ylabel!("x(t)")
# py = plot(sol_true, vars=(0,2), color=:black, alpha=0.75, linewidth=2, label="true system")
# xlabel!("t")
# ylabel!("y(t)")

# pz = plot(sol_true, vars=(0,3), color=:black, alpha=0.75, linewidth=2, label="true system")
# xlabel!("t")
# ylabel!("z(t)")

# train_plot = plot(px, py, pz, layout=(3,1))

## ----- setup observation data ------------------------------------

# set up observation uncertainties
σ_m = 0.15  # i.e. observation std
# set up observation error covariance matrix
R = σ_m^2 .* I(3)

# set up details for measurement points
dt_m = 0.2
tm_end = 2.0  # final observation time
ts_m = dt_m:dt_m:tm_end

# define observation operator
function h(u)
    w = u
    return w
end

# define observation operator jacobian. We should probably generate this symbolically
function Dh(u)
    n = size(u, 1)
    D = I(n)
    return D
end

# populate observation vector
w = hcat([h(sol_true(tᵢ))+rand(MvNormal(zeros(3), R)) for tᵢ ∈ ts_m]...)

# update the graphs with observation stuff
# plot!(px, ts_m, w[1, :], seriestype=:scatter, color="light green", label="observation" )
# plot!(py, ts_m, w[2, :], seriestype=:scatter, color="light green", label="observation" )
# plot!(pz, ts_m, w[3, :], seriestype=:scatter, color="light green", label="observation" )

# train_plot = plot(px, py, pz, layout=(3,1))

## ----- perform the assimilation ------------------------------------

# set up data assimilation
u0b = [2.0; 3.0; 4.0]  # guess at initial condition
σ_b = 0.1  # model error std
B = σ_b^2 .* I(3)  # initial model error covariance
Q = 0.0 * I(3)  # process noise (assumed zero)

# integrate to obtain the true solution
prob_no_analysis = ODEProblem(lorenz!, u0b, ts_sim, θ, saveat=dt)
sol_no_analysis = solve(prob_no_analysis)

# set up arrays for solution outputs... we should probably pre-allocate these
ua = zeros(3, size(0:dt:tend, 1))
ua[:,1] = u0b

# function for getting indices from time value
function timeindex(t, Δt)
    return Int(round(t/Δt)) + 1
end

# helper function to update the analysis result
function update!(sol, t_now, t_next, dt)
    idx_now = timeindex(t_now, dt)
    idx_next = timeindex(t_next, dt)

    ua[:, idx_now:idx_next] = hcat(sol(t_now:dt:t_next)...)
end
Zygote.@nograd update!  # tell zygote to ignore update! when computing derivatives

function model_forward!(uₖ)
    prob = ODEProblem(lorenz!, uₖ, (t_now, t_next), θ)
    sol = solve(prob, Tsit5(), reltol=1e-6,abstol=1e-6)

    update!(sol, t_now, t_next, dt)

    return sol[:, end]
end

# ------------- perform the analysis ---------------------------------------------
# times = hcat(collect(0:dt_m:tm_end)..., tend)
times = 0:dt_m:tend

t_now = times[1]
t_next = times[2]
#model_forward!(u0b)

res = Zygote.withjacobian(model_forward!, u0b)
res.val
res.grad[1]

for k ∈ 1:(length(times)-1)
    t_now = times[k]
    t_next = times[k+1]

    idx_now = timeindex(t_now, dt)
    u_now = ua[:, idx_now]

    # u_next, DM = model_forward!(u_now)
    u_next, DM = Zygote.withjacobian(model_forward!, u_now)
    DM = DM[1]

    # update error covariance matrix
    B = DM * B * DM' + Q

    if t_next ∈ ts_m
        println("doing the update!")
        n = size(ua, 1)
        idx_next = timeindex(t_next, dt)
        idx_h = timeindex(t_next, dt_m) - 1 # -1 since this doesn't include 0

        # compute the Kalman gain matrix
        DH = Dh(ua[:,idx_next])
        println("\t Dₕ=$(DH)")
        denom = DH*B*DH' + R
        println("\t B:= $(B)")
        println("\t denom:= $(denom)")
        K = B*DH'*inv(denom)

        println("\tK= $(K)")
        # perform the analysis

        println("\tua = $(ua[:, idx_next])")
        ua[:,idx_next] .= ua[:, idx_next] + K*(w[:, idx_h] - h(ua[:, idx_next]))

        println("\tua = $(ua[:, idx_next])")
        B = (I(n) - K*DH)*B
    end
end

var_names = ["x(t)", "y(t)", "z(t)"]
plots = []
for i ∈ 1:3
    p = plot(sol_true, vars=(0,i), color=:black, alpha=0.75, linewidth=2, label="true system")
    plot!(p, ts_m, w[i, :], seriestype=:scatter, color="light green", label="observation" )
    plot!(p, sol_no_analysis, vars=(0, i), linestyle=:dashdot, color="purple", label="forecast")
    plot!(p, 0:dt:tend, ua[i, :], color="royal blue", alpha=1, linewidth = 1, label="analysis")
    xlabel!(p, "t")
    ylabel!(p, var_names[i])
    push!(plots, p)
end

final_plot = plot(plots..., layout=(3,1))

savefig("EKF_using_zygote.svg")
savefig("EKF_using_zygote.png")
display(final_plot)


# returns a lazy "vector" of the (real) value of `flat` over the entire unit cell, using
# `nsamples` per dimension; avoids double counting at unit cell edges
function _lazy_fourier_eval_over_uc_as_vector(flat::AbstractFourierLattice{D}, nsamples::Int64) where D
    step = 1.0/nsamples
    samples = range(-0.5, 0.5-step, length=nsamples) # `-step` to avoid double counting ±0.5 (cf. periodicity)
    if D == 2
        itr = (real(flat(x,y)) for x in samples for y in samples)
    elseif D == 3
        itr = (real(flat(x,y,z)) for x in samples for y in samples for z in samples)
    end
end

"""
    $(TYPEDSIGNATURES)

Return the isovalue of `flat` such that the interior of the thusly defined levelset
isosurface encloses a fraction `filling` of the unit cell.

The keyword argument `nsamples` specifies the grid-resolution used in evaluating this
answer (via `quantile`).
"""
function filling2isoval(flat::AbstractFourierLattice{D}, filling::Real=0.5, nsamples::Int64=51) where D
    itr = _lazy_fourier_eval_over_uc_as_vector(flat, nsamples)
    return quantile(itr, filling)
end

"""
    $(TYPEDSIGNATURES)

Return the filling fraction of the interior of the isosurface defined by `flat` for an
isovalue `isoval`.

The keyword argument `nsamples` specifies the grid-resolution used in evaluating this
answer (using staircase integration).
"""
function isoval2filling(flat::AbstractFourierLattice{D}, isoval::Real, nsamples::Int64=51) where D
    itr = _lazy_fourier_eval_over_uc_as_vector(flat, nsamples)
    return count(<(isoval), itr)/nsamples^D
end


module Test_bouncingBall_with_measurement

#=
Strange result of bouncing ball with Measurement.jl

https://github.com/JuliaPhysics/Measurements.jl/issues/64

Issue fixed in Measurements version 2.2.0
=#

using DifferentialEquations, Measurements, PyPlot

const g   = 9.81       # Gravity constant
const h0  = 1.0 ± 0.1  # Initial height
const cor = 0.7 ± 0.05  # Coefficient of restitution

function f(du,u,p,t)
    du[1] = u[2]
    du[2] = -g
end

function condition(u,t,integrator)
    z = u[1] + 1e-12
    return z
end

function affect_neg(integrator)
    println("Event at time = ", integrator.t, ", h = ", integrator.u[1])

    v_before = integrator.u[2]
    println("   v_before = ", v_before)
    v_after = -cor*v_before
    println("   v_after  = ", v_after)
    integrator.u[2] = v_after

    auto_dt_reset!(integrator)
    set_proposed_dt!(integrator, integrator.dt)
end

cb = ContinuousCallback(condition,nothing,affect_neg! = affect_neg)

u0 = [h0, 0.0]
tspan = (0.0±0.0, 2.0)
prob = ODEProblem(f,u0,tspan)
sol  = solve(prob,Tsit5(),tolerance=1e-6,saveat=0.01,callback=cb)

println("\n... Analytic solution upto second bounce")
t_ev1        = sqrt(2*h0/g)
h_ev1        = -g/2*t_ev1^2 + h0
v_ev1_before = -g*t_ev1
v_ev1_after  = -cor*v_ev1_before
v_ev1_0_46   = -g*(0.46-t_ev1) + v_ev1_after

t_ev2        = t_ev1 +2*v_ev1_after/g
h_ev2        = -g/2*(t_ev2-t_ev1)^2 + v_ev1_after*(t_ev2-t_ev1)
v_ev2_before = -v_ev1_after
v_ev2_after  = -cor*v_ev2_before

println("Event at time = ", t_ev1, ", h = ", h_ev1)
println("   v_before = ", v_ev1_before)
println("   v_after  = ", v_ev1_after)
println("   v(0.46)  = ", v_ev1_0_46)

println("Event at time = ", t_ev2, ", h = ", h_ev2)
println("   v_before = ", v_ev2_before)
println("   v_after  = ", v_ev2_after)

# Plot result
tm = Measurements.value.(sol.t)
tu = Measurements.uncertainty.(sol.t)

function plotUncertaintyOfEvents(tm, tu, vm, vu; color=[0.0, 0.0, 1.0, 0.5], linewidth=3.0, color_uncertainty=[0, 0.16, 0.78, 0.2])
    first = true
    for i in 1:length(tm)-1
        if tm[i] == tm[i+1] && tu[i] > 0.0
            tt = [tm[i]-tu[i], tm[i]+tu[i]]
            vlow = vm[i] - vu[i]
            vup  = vm[i+1] + vu[i+1]
            fill_between(tt,[vlow, vlow], [vup, vup], color=color_uncertainty)
            if first
                first = false
                plot(tt, [0.0, 0.0], color=color, linewidth=linewidth, label="uncertainty(time)")
            else
                plot(tt, [0.0, 0.0], color=color, linewidth=linewidth)
            end
        end
    end
end

figure(1)
clf()

h = getindex.(sol.u, 1)
hm = Measurements.value.(h)
hu = Measurements.uncertainty.(h)
hmax = hm + hu
hmin = hm - hu
#plot(tt, hmax, label="\$h_{max} in [m]\$")
#plot(tt, hmin, label="\$h_{min} in [m]\$")

# Remove uncertainty values with h < 0
for i in 1:length(hmin)
    if hmin[i] < 0
        hmin[i] = 0
    end
end
fill_between(tm,hmin,hmax, color=[0, 0.16, 0.78, 0.2], label="uncertainty(h)")
plot(tm, hm, color=[0.0, 0.0, 1.0, 1.0], label="\$h_{mean}\$ in [m]")
grid(true)
plotUncertaintyOfEvents(tm,tu,hm,hu)
xlabel("time in [s]")
legend()
text(0.25, 1.1, "\$h_{0} = \$" * string(h0))
text(0.25, 1.05, "\$cor = \$" * string(cor))
title("Bouncing ball with Measurement{Float64}")

figure(2)
clf()

v = getindex.(sol.u, 2)
vm = Measurements.value.(v)
vu = Measurements.uncertainty.(v)
vmax = vm + vu
vmin = vm - vu
#plot(tm, vmax, label="\$v_{max}\$")
#plot(tm, vmin, label="\$v_{min}\$")
fill_between(tm,vmin,vmax, color=[0, 0.16, 0.78, 0.2], label="uncertainty(v)")
plot(tm, vm, color=[0.0, 0.0, 1.0, 1.0], label="\$v_{mean}\$ in [m/s]")
grid(true)
plotUncertaintyOfEvents(tm,tu,vm,vu)
xlabel("time in [s]")
legend(loc="lower right")
text(0, 3.5, "\$h_{0} = \$" * string(h0))
text(0, 3.0, "\$cor = \$" * string(cor))
title("Bouncing ball with Measurement{Float64}")

figure(3)
clf()
plot(tm, tu, label="uncertainty(time) in [s]")
grid(true)
legend()
title("Bouncing ball with Measurement{Float64}")
end


import StaticArrays: SVector

#########################################################
# Helper functions to plot 3D rectangles.
#########################################################
"""
    rectangle3dvertices(x, y, z, ϵx, ϵy, ϵz)

Compute the coordinates of the vertices of a 3D rectangle, given the
coordinates of the origin (`x`, `y`, and `z`) and the corresponding edge lengths
(`ϵx`, `ϵy`, `ϵz`).
"""
function rectangle3dvertices(x, y, z, ϵx, ϵy, ϵz)
    v1b = [x,    y,    z]
    v2b = [x+ϵx, y,    z]
    v3b = [x+ϵx, y+ϵy, z]
    v4b = [x,    y+ϵy, z]
    v1t = [x,    y,    z+ϵz]
    v2t = [x+ϵx, y,    z+ϵz]
    v3t = [x+ϵx, y+ϵy, z+ϵz]
    v4t = [x,    y+ϵy, z+ϵz]
    (v1b, v2b, v3b, v4b, v1t, v2t, v3t, v4t)
end

"""
    rectangle3dvertices(o, ϵ)

Given the origin `o` (3-element vector) and edge lengths `ϵ` (also a
3-element vector) of a 3D rectangle, return the coordinates of its
vertices.
"""
function rectangle3dvertices(o, ϵ)
    x, y, z = (o...,)
    ϵx, ϵy, ϵz = (ϵ...,)
    rectangle3dvertices(x, y, z, ϵx, ϵy, ϵz)
end

"""
    connectvertices3d(o, ϵ)

Given the origin `o` (3-element vector) and edge lengths `ϵ` (also a
3-element vector) of a 3D rectangle, return a vector of line segments
(each a dim-by-n_vertices array) that when plotted together yields
the entire rectangle.
"""
function connectvertices3d(o, ϵ)
    # Get the vertices
    v1b, v2b, v3b, v4b, v1t, v2t, v3t, v4t = rectangle3dvertices(o, ϵ)

    # Connect vertices in top and bottom planes
    connect_top = zeros(3, 5)
    connect_bottom = zeros(3, 5)

    # Connect corners of the top and bottom planes
    corner1 = zeros(3, 2)
    corner2 = zeros(3, 2)
    corner3 = zeros(3, 2)
    corner4 = zeros(3, 2)
    connect_bottom[:, 1:5] = [v1b v2b v3b v4b v1b]
    connect_top[:, 1:5]    = [v1t v2t v3t v4t v1t]
    corner1[:, 1:2] = [v1b v1t]
    corner2[:, 1:2] = [v2b v2t]
    corner3[:, 1:2] = [v3b v3t]
    corner4[:, 1:2] = [v4b v4t]

    return [connect_top, connect_bottom,
            corner1, corner2, corner3, corner4]
end

export connectvertices3d

"""
    splitaxes(x)

Return a vector of the individual components of an array of points
provided as an array where each point is a column.
"""
splitaxes(x) = ([x[k, :] for k = 1:size(x, 1)]...,)

"""
    plot_3D_rect!(p, origin, edgelengths;
             lc = :black, lw = 1.0, ls = :solid)

Append a a 3D rectangle to an existing plot `p`, given the
origin `o` (3-element vector) and edge lengths
`ϵ` (3-element vector) of the rectangle.
"""
function plot_3D_rect!(p, o, ϵ;
             lc = :black, lw = 1.0, ls = :solid, lα = 0.7)

    linesegments = connectvertices3d(o, ϵ)
    for segment in linesegments
        plot!(p, splitaxes(segment),
            lc = lc, lw = lw, ls = ls, lα = lα)
    end
end

"""
Given the `origin` of a hyperrectangle and the `stepsizes` along each
coordinate axis, divide each axis into `N` stepsizes starting
from `origin[i] + (stepsizes[i]/N)/2`
"""
function fill_hyperrectangle_3D(origin, stepsizes, N = 10)
    xs, ys, zs = ([LinRange(origin[i] + stepsizes[i]/N/2, origin[i] + stepsizes[i] - stepsizes[i]/N/2, N) for i = 1:length(origin)]...,)
    pts = zeros(Float64, 3, length(xs)^3)
    i = 0
    for x in xs
        for y in ys
            for z in zs
                i += 1
                pts[:, i] = [x, y, z]
            end
        end
    end
    pts
end

struct Rectangle{D, T}
    origin::SVector{D, T}
    edgelengths::SVector{D, T}
end

rectangle3dvertices(r::Rectangle) = rectangle3dvertices((r.origin...,), (r.edgelengths...,))

function linesegments(r::Rectangle)
    # Get the vertices
    v1b, v2b, v3b, v4b, v1t, v2t, v3t, v4t = rectangle3dvertices(r.origin, r.edgelengths)

    # Connect vertices in top and bottom planes
    connect_top = zeros(3, 5)
    connect_bottom = zeros(3, 5)

    # Connect corners of the top and bottom planes
    corner1 = zeros(3, 2)
    corner2 = zeros(3, 2)
    corner3 = zeros(3, 2)
    corner4 = zeros(3, 2)
    connect_bottom[:, 1:5] = [v1b v2b v3b v4b v1b]
    connect_top[:, 1:5]    = [v1t v2t v3t v4t v1t]
    corner1[:, 1:2] = [v1b v1t]
    corner2[:, 1:2] = [v2b v2t]
    corner3[:, 1:2] = [v3b v3t]
    corner4[:, 1:2] = [v4b v4t]

    linesegments = [connect_top, connect_bottom, corner1, corner2, corner3, corner4]
    return linesegments
end

function connectrect(r::Rectangle)
    v1b, v2b, v3b, v4b, v1t, v2t, v3t, v4t = rectangle3dvertices(r.origin, r.edgelengths)

    x = hcat(v1b, v2b, v2t, v3t, v3b, v4b, v4t, v1t, v1b, v1t, v2t, v2b, v3b, v3t, v4t, v4b, v1b)
    splitaxes(x)
end

export
Rectangle,
connectrect,
rectangle3dvertices,
linesegments,
connectvertices3d


function smooth_max(k_smooth::Float64, level_sideline)

    # Compute lateral noise
    smooth_max = maximum(level_sideline) + 1/k_smooth.*log(sum(exp.(k_smooth*(level_sideline.-maximum(level_sideline)))))

    # Return
    return smooth_max
end


function smooth(y, win_len=11, win_method=2)
# Based on http://www.simulkade.com/posts/2015-05-07-how-to-smoothen-noisy-data.html
# This function requires the DSP package to be installed
# 1: flat
# 2: hanning
# 3: hamming ...
if win_len%2==0
  win_len+=1 # only use odd numbers
end
if win_method == 1
  w=ones(win_len)
elseif win_method==2
  w=DSP.hanning(win_len)
elseif win_method==3
  w=DSP.hamming(win_len)
end

if win_len<3
  return y
elseif length(y)<win_len
  return y
else
        y_new = [2*y[1].-reverse(y[1:win_len],1); y[:]; 2*y[end].-reverse(y[end-win_len:end],1)]
  y_smooth = conv(y_new, w/sum(w))
  ind = floor(Int, 1.5*win_len)
  return y_smooth[1+ind:end-ind-1]
end

end # end of function


using Cbc

include(joinpath(Pkg.dir("MathProgBase"),"test","mixintprog.jl"))

mixintprogtest(CbcSolver(LogLevel=1))


include("../src/pressurecorrelations.jl")

@testset "Friction factors" begin
    @test ChenFrictionFactor(35700, 2.259, 0.001)/4 ≈ 0.0063 atol = 0.001

    @test ChenFrictionFactor(5.42e5, 2.295, 0.0006)/4 ≈ 0.0046 atol=0.001 # per economides
    @test SerghideFrictionFactor(5.42e5, 2.295, 0.0006)/4 ≈ 0.0046 atol=0.001

    @test ChenFrictionFactor(71620, 2.441, 0.0002) ≈ 0.021 atol = 0.005 #per Takacs
    @test SerghideFrictionFactor(71620, 2.441, 0.0002) ≈ 0.021 atol = 0.005

    @test ChenFrictionFactor(101000, 2.295, 0.0002)/4 ≈ 0.005 atol = 0.005 #per economides
    @test SerghideFrictionFactor(101000, 2.295, 0.0002)/4 ≈ 0.005 atol = 0.005
end

@testset "Superficial velocities" begin

    @test liquidvelocity_superficial(375, 0, 2.441, 1.05, 0) ≈ 0.79 atol = 0.01
    @test gasvelocity_superficial(375, 0, 480, 109.7, 2.441, 0.0389) ≈ 1.93 atol = 0.01

end

#%% Beggs and Brill flowmap

#= B&B flowmap graphical test, to make sure the correct mapping is generated
# Re-run after any modifications to B&B flowmap function.
# Note that this does not establish stability at the region boundaries.

λ_l = 10 .^ collect(-4:0.1:0);
N_Fr = 10 .^ collect(-1:0.1:3);

function crossjoin(x, y)

    output = Array{Float64, 2}(undef, length(x) * length(y), 2)

    index = 1
    @inbounds for i in x
        @inbounds for j in y
            output[index, 1] = i
            output[index, 2] = j
            index += 1
        end
    end

    return output
end

testdata = crossjoin(λ_l, N_Fr);
flowpattern = map(BeggsAndBrillFlowMap, testdata[:,1], testdata[:,2]);

using Gadfly

plot(x = testdata[:,1], y = testdata[:,2], color = flowpattern,
Geom.rectbin,
Scale.x_log10, Scale.y_log10, Coord.Cartesian(xmin=-4, xmax=0, ymin=-1, ymax=3))
=#

@testset "Beggs and Brill vertical - single segment" begin

    inclination = 0
    v_sl = 0.79
    v_sg = 1.93
    v_m = 2.72
    ρ_l = 49.9
    ρ_g = 1.79
    id = 2.441
    σ_l = 8
    μ_l = 3
    μ_g = 0.02
    roughness = 0.0006
    pressure_est = 346.6

    @test BeggsAndBrill(1, 1, inclination, id, v_sl, v_sg, ρ_l, ρ_g, σ_l, μ_l, μ_g, roughness, pressure_est, SerghideFrictionFactor, true, false) ≈ 0.170 atol = 0.01

end


function no_whitening(
    dataset::DataSet{T, I})::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    datamean = zeros(T, dataset.P)

    transform_data!(dataset, Array{T, 2}(LinearAlgebra.I, dataset.P, dataset.P), datamean)
end

function cholesky_whitening!(
    dataset::DataSet{T, I})::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    datamean = zeros(T, dataset.P)

    for p in eachindex(datamean)
        datamean[p] = mean(view(dataset.data, p, :))
    end

    for n in 1:dataset.N
        buffer = view(dataset.data, :, n) - datamean
        setindex!(dataset.data, buffer, :, n)
    end

    #cov of a Transpose object yields the wrong (0.0) result in 1D, if all weights are equal (e.g. direct sampling)!
    covmatrix = cov(convert(typeof(dataset.data), transpose(dataset.data)), FrequencyWeights(dataset.weights), corrected=true)
    symcovmatrix = Symmetric(covmatrix)
    covmatrix_inv = inv(symcovmatrix)
    w = cholesky(covmatrix_inv).U
    wres = convert(Matrix{T}, w)

    transform_data!(dataset, wres, datamean)
end

function statistical_whitening!(
    dataset::DataSet{T, I})::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    datamean = zeros(T, dataset.P)

    for p in eachindex(datamean)
        datamean[p] = mean(view(dataset.data, p, :))
    end

    dataset.data .-= datamean

    covmatrix = cov(transpose(dataset.data), FrequencyWeights(dataset.weights), corrected=true)

    E = eigfact(covmatrix).vectors
    w_d = transpose(E)
    D = Diagonal(inv(E) * covmatrix * E)
    w = inv(full(sqrt.(D))) * w_d
    wres = convert(Matrix{T}, w)

    transform_data!(dataset, wres, datamean)
end

function transform_data!(
    dataset::DataSet{T, I},
    W::Matrix{T},
    datamean::Vector{T},
    substractmean::Bool = false)::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    local determinant::T

    if substractmean
        dataset.data .-= datamean
    end

    if diag(W) == ones(T, dataset.P)
        determinant = 1.0
    else
        dataset.data = W * dataset.data

        determinant = abs(det(W))
    end

    maxP::T = maximum(dataset.logprob)
    suggTargetProb::T = exp(maxP - partialsort(dataset.logprob, floor(Int64, dataset.N * 0.2)))

    dataset.iswhitened = true

    @debug "Determinant:\t" * string(determinant)
    @debug "Suggested Target Probability Factor:\t" * string(suggTargetProb)

    WhiteningResult(determinant, suggTargetProb, W, datamean)
end


cd(@__DIR__); include("setups/grid7x3.jl");

## DAG
D = natural_eigdist(𝚽, 𝛌, Q; distance = :DAG)
E = transform(fit(MDS, D, maxoutdim=2, distances=true))
plt = grid7x3_mds_heatmaps(E, 𝚽; backend = :pyplot)
savefig(plt, "../figs/Grid7x3_MDS_DAG.png")
display(plt)


@testset "Ideal class monoid of orders in number fields" begin

  Qx, x = FlintQQ["x"]
  f = x^3 + 31*x^2 + 43*x + 77
  K, a = number_field(f, "a")
  O = equation_order(K)
  icm = ideal_class_monoid(O)
  @test length(icm) == 59
  @test isisomorphic(evaluate(icm[1]), evaluate(icm[2]))[1] == false
  t, b = isisomorphic(O(2)*O, O(3)*O)
  @test t == true
  @test (O(3)*O)*(b*O) == K(2)*O
  t, b = isisomorphic(K(2)*O, K(3)*O)
  @test t == true
  @test (O(3)*O)*(b*O) == K(2)*O

  f = x^3 - 1000*x^2 - 1000*x - 1000
  K, a = number_field(f, "a")
  O = equation_order(K)
  icm = ideal_class_monoid(O)
  @test length(icm) == 69116

end

@testset "Ideal class monoid of orders in algebras" begin

  Qx, x = FlintQQ["x"]
  f = x^3 + 31*x^2 + 43*x + 77
  A = AlgAss(f)
  O = Order(A, basis(A))
  icm = ideal_class_monoid(O)
  @test length(icm) == 59
  @test isisomorphic(evaluate(icm[1]), evaluate(icm[2]))[1] == false
  t, b = isisomorphic(O(2)*O, O(3)*O)
  @test t == true
  @test (O(3)*O)*(b*O) == A(2)*O
  t, b = isisomorphic(A(2)*O, A(3)*O)
  @test t == true
  @test (O(3)*O)*(b*O) == A(2)*O

  f1 = x^2 + 4*x + 7
  f2 = x^3 - 9*x^2 - 3*x - 1
  A = AlgAss(f1*f2)
  O = Order(A, basis(A))
  icm = ideal_class_monoid(O)
  @test length(icm) == 852

end

@testset "Conjugacy of integral matrices" begin

  Zx, x = FlintZZ["x"]
  f = x^3 - 4x^2 + 8x + 5

  M = companion_matrix(f)
  U = identity_matrix(FlintZZ, nrows(M))
  for i = 1:nrows(M)
    for j = i + 1:ncols(M)
      U[i, j] = rand(-10:10)
    end
  end
  N = U*M*inv(U)

  b, V = isconjugate(M, N)
  @test b == true
  @test M == V*N*inv(V)

  N = identity_matrix(FlintZZ, nrows(M))
  b, V = isconjugate(M, N)
  @test b == false

  g = x^2 + 10x - 1
  h = x^2 - x + 5

  M = companion_matrix(f*g*h)
  U = identity_matrix(FlintZZ, nrows(M))
  for i = 1:nrows(M)
    for j = i + 1:ncols(M)
      U[i, j] = rand(-10:10)
    end
  end
  N = U*M*inv(U)

  b, V = isconjugate(M, N)
  @test b == true
  @test M == V*N*inv(V)

  N = identity_matrix(FlintZZ, nrows(M))
  b, V = isconjugate(M, N)
  @test b == false

end