using Revise

using DataFrames, Statistics, Distributions, StatsBase,
  LinearAlgebra, BenchmarkTools, Random, Optim, Zygote, UnPack

MFloat64 = Union{Missing, Float64}

Random.seed!(11)

# a simple AR(1) model
struct AR1{T<:Real, Tσ²<:Union{Nothing,Real}}
  ϕ::T
  b::T
  μ::T
  σ²::Tσ² #note- this is σ²e, not the overall σ² of the series
end

AR1(;ϕ, μ=missing, b=missing, σ²=nothing) = AR1(
  ϕ, coalesce(b, μ*(1-ϕ)), coalesce(μ, b/(1-ϕ)),σ²)

centerar1(ar1::AR1) = AR1(ar1.ϕ, 0.0, 0.0, ar1.σ²)

function estimatear1ols(x::Vector{T}) where T
  x̃ = @view x[2:end]
  Lx̃ = @view x[1:(end-1)]

  ϕ = cov(x̃, Lx̃)/(var(Lx̃))
  b = mean(x̃) - ϕ * mean(Lx̃)
  μ = b/(1-ϕ)

  ε = similar(x)
  ε[1] = x[1] - μ
  ε[2:end] = x̃ .- ϕ .* Lx̃
  σ² = var(ε)

  @assert b ≈ mean(x̃) - ϕ * mean(Lx̃)
  @assert ϕ ≈ cov(x̃ .- mean(x̃), Lx̃ .- mean(Lx̃))/var(
    Lx̃ .- mean(Lx̃))

  return AR1(;ϕ,b,σ²)
end

function estimatear1mle(x; tol=10^-10)

  #start by centering
  N=length(x)
  μ = mean(x)
  x̃ = @view(x[2:end]) .- μ
  Lx̃ = @view(x[1:(end-1)]) .- μ

  x₁ = x[1]
  #computes the next iteration of ϕ given a previous ϕ
  cσ²(ϕ) = (x₁^2*(1-ϕ^2) + sum((x̃ .- ϕ*Lx̃).^2))/N
  function objϕ(ϕ)
    σ² = cσ²(ϕ)
    return (sum((x̃ .- ϕ * Lx̃) .* Lx̃) - (σ²/(1-ϕ^2) - x₁^2)*ϕ)^2
  end

  ϕ = optimize(objϕ, -1.01, 1.01, Brent()) |> Optim.minimizer

  return AR1(;ϕ, μ, σ² = cσ²(ϕ))
end

function estimatear1mlecheck(x; tol=10^-10)

  #start by centering
  N=length(x)
  μ = mean(x)
  x̃ = @view(x[2:end]) .- μ
  Lx̃ = @view(x[1:(end-1)]) .- μ

  x₁ = x[1]

  l(ϕ, σ²) = (-N/2*log(2*π*σ²)
    + 1/2*log(1-ϕ^2)
    - x₁^2*(1-ϕ^2)/(2*σ²)
    - sum((x̃ .- ϕ .* Lx̃).^2)/(2*σ²)) * -1

  l(Φ) = l(Φ[1], Φ[2])

  ϕ₀ = cov(x̃, Lx̃)/var(Lx̃) #start by guessing the least squares solution
  ε₀= x̃ - ϕ₀*Lx̃
  σ²₀ = var(ε₀)

  ϕ, σ² = optimize(l,  [-1.0,0.0], [1.0,Inf], [ϕ₀, σ²₀], Fminbox(LBFGS()),
    autodiff = :forward) |> Optim.minimizer

  return AR1(;ϕ, μ, σ²)
end

AR1(x::AbstractVector; estimationmethod=estimatear1mle) = estimationmethod(x)

function checkmath2(x)
  if length(x) > 250
    x = x[1:250] |> deepcopy
    @warn "length(x) > 250 so truncating length to 250. This is only for the purposes
    of computing the full (not log) likelihood"
  end

  N=length(x)
  μ = mean(x)
  x̃ = @view(x[2:end]) .- μ
  Lx̃ = @view(x[1:(end-1)]) .- μ

  ϕ₀ = 0.75
  σ²₀ = 2.2
  x₁ = x̃[1]

  L(ϕ, σ²) =  (1-ϕ^2)^(1/2)/σ²^(1/2)*pdf(Normal(), x₁/(σ²/(1-ϕ^2))^(1/2))*prod(
    ((x̃ᵢ,Lx̃ᵢ)->1/σ²^(1/2)*pdf(Normal(), (x̃ᵢ-ϕ*Lx̃ᵢ)/σ²^(1/2))).(x̃,Lx̃))

  l(ϕ, σ²) = (-N/2*log(2*π*σ²)
    + 1/2*log(1-ϕ^2)
    - x₁^2*(1-ϕ^2)/(2*σ²)
    - sum((x̃ .- ϕ .* Lx̃).^2)/(2*σ²))

  @assert log(L(ϕ₀, σ²₀)) ≈ l(ϕ₀, σ²₀,) "log(L(ϕ, σ²)): $(log(L(ϕ₀, σ²₀)))
    l(ϕ, σ²): $(l(ϕ₀, σ²₀))"

  #now verify the derivitive and conditional wrt sigma^2
  lσ²(ϕ=ϕ₀, σ²=σ²₀,) = (-N/(2σ²)
    + x₁^2*(1-ϕ^2)/(2σ²^2)
    + sum((x̃ .- ϕ .* Lx̃).^2)/(2σ²^2))
  lσ²z(ϕ=ϕ₀, σ²=σ²₀,) = gradient(σ²->l(ϕ, σ²,), σ²)[1]
  lσ²n() =  (l(ϕ₀, σ²₀+1e-6,)-l(ϕ₀, σ²₀,))/1e-6
  @assert lσ²() ≈ lσ²z() "lσ²(...): $(lσ²())
    lσ²z(...): $(lσ²z())"

  cσ²(ϕ) = (x₁^2*(1-ϕ^2) + sum((x̃ .- ϕ*Lx̃).^2))/N
  @assert 1 + lσ²z(ϕ₀, cσ²(ϕ₀)) ≈ 1.0 "cσ²(ϕ₀): $(cσ²(ϕ₀))
    lσ²z(ϕ₀, cσ²(ϕ₀)): $(lσ²z(ϕ₀, cσ²(ϕ₀)))"

  #finally, verify the derivitive wrt ϕ
  lϕ(ϕ=ϕ₀, σ²=σ²₀) = (-ϕ/(1-ϕ^2) +
    x₁^2*ϕ/σ² +
    sum((x̃ .- ϕ * Lx̃) .* Lx̃)/σ²)
  lϕz(ϕ=ϕ₀, σ²=σ²₀,) = gradient(ϕ->l(ϕ, σ²), ϕ)[1]
  @assert lϕ() ≈ lϕz() "lϕ(...): $(lϕ())
    lϕz(...): $(lϕz())"

  ϕcerr(ϕ, σ²=σ²₀)= sum((x̃ .- ϕ * Lx̃) .* Lx̃) - (σ²/(1-ϕ^2) - x₁^2)*ϕ
  cϕ(σ²=σ²₀,) = optimize(ϕ->ϕcerr(ϕ)^2, -1.01, 1.01, Brent(), abs_tol=10^-14) |> Optim.minimizer
  @assert isapprox(lϕz(cϕ()), 0.0,atol=1e-5) "cϕ(ϕ₀): 2$(cϕ())
    lϕz(ϕcerr()): $(lϕz(ϕcerr()))"
  #lμ()
end

function simulatear1!(Xt::AbstractMatrix;
    ϕ::TV, σ::TV, μ::TV) where TV<:AbstractVector
  #@unpack ϕ, b, σ², μ = ar1

  #first simulate the centered distribution
  #we will do this all in place

  rand!(Normal(), Xt)
  Xt .*= σ

  Xt[:,1] .*= 1 ./ (1 .- ϕ.^2) .^ 0.5 # the first entry is simulated from the unconditional distribution

  #=Threads.@threads for c ∈ 1:size(X,2)
    for t ∈ 2:size(X,1)
      X[t,c] += ϕ[c] * X[t-1,c]
    end
  end=#

  for t ∈ 2:size(Xt,2)
    for c ∈ 1:size(Xt,1)
      Xt[c,t] += ϕ[c] * Xt[c,t-1]
    end
  end

  Xt .+= μ

  return nothing
end

#simulates an AR1 series
function simulatear1!(x::AbstractVector; ϕ, σ, μ)

  rand!(Normal(0,σ), x)
  x[1] *= 1/(1-ϕ^2)^0.5 # the first entry is simulated from the unconditional distribution
  for t ∈ 2:length(x)
    x[t] += ϕ * x[t-1]
  end
  x .+= μ

  return nothing
end

#non-inplace version for testing
function simulatear1(ar1::AR1{T,Tσ²}, N::Int) where {T,Tσ²}
  res = Vector{Float64}(undef, N)
  simulatear1!(res; σ=ar1.σ²^0.5, μ=ar1.μ, ϕ=ar1.ϕ)
  return res
end

function compareoncear1xy!(Xt::AbstractMatrix,
  yar1::AR1,  y::AbstractVector,
  comparefunc::Tcomparefunc;
  ϕx::Vector{Float64},σx::Vector{Float64}, μx::Vector{Float64}) where Tcomparefunc

  simulatear1!(Xt; ϕ=ϕx, σ=σx, μ=μx)
  simulatear1!(y, ϕ=yar1.ϕ, σ=yar1.σ²^0.5, μ=yar1.μ)
  comparefunc(Xt',y)
end

function compareoncear1xy!(Xt::AbstractVector,
  yar1::AR1,  y::AbstractVector,
  comparefunc::Tcomparefunc;
  ϕx::Vector{Float64},σx::Vector{Float64}, μx::Vector{Float64}) where Tcomparefunc

  simulatear1!(Xt; ϕ=ϕx[1], σ=σx[1], μ=μx[1])
  simulatear1!(y, ϕ=yar1.ϕ, σ=yar1.σ²^0.5, μ=yar1.μ)
  comparefunc(Xt,y)
end

#this runs the simulations- basically a parametric bootstrap under the null
#statfunc is computed for each pass, while aggfunc is computed on the statistics
#at the end of the pass
function comparear1undernull(X::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4,
    Xts = (size(X,2) > 1 ?
      [similar(X' |> Matrix{Float64}) for t ∈ 1:Threads.nthreads()] :
      [similar(X |> vec |> Vector{Float64}) for t ∈ 1:Threads.nthreads()])
    ) where {Tx<:AbstractMatrix, Tstatfunc, Taggfunc}

  x = collect(eachcol(X))
  K = length(x)
  xar1 = (AR1.(x)) .|> centerar1
  yar1 = AR1(y) |> centerar1

  results = Dict(s => missings(Float64, Nsimulations, K) for s ∈ keys(statfunc(X,y)))
  @assert (k->ismissing.(results[k]) |> all).(keys(results)) |> all

  #this assumes stationary results are meaningless
  if checkstationarity && max((xar1 .|> xar1i -> abs(xar1i.ϕ))..., abs(yar1.ϕ)) ≥ 1.0-stationaritytol
    xar1 .= ((xar1i)->AR1(NaN, NaN, xar1i.μ, xar1i.σ²)).(xar1)
    yar1 = AR1(NaN, NaN, yar1.μ, yar1.σ²)
  end

  ys = [similar(y) for t ∈ 1:Threads.nthreads()]

  ϕx::Vector{Float64} = xar1 .|> xar1i->xar1i.ϕ
  σx::Vector{Float64} = xar1 .|> xar1i->xar1i.σ²^0.5
  μx::Vector{Float64} = xar1 .|> xar1i->xar1i.μ

  #draw a sample and compute the aggregate function, sotring the reuslts in an array
  Threads.@threads for n ∈ 1:Nsimulations
  #for n ∈ 1:Nsimulations

    t::Int = Threads.threadid()
    #res = compareoncear1xy!(xar1, xs[1], yar1, ys[1], statfunc)
    res = compareoncear1xy!(Xts[t], yar1, ys[t], statfunc; ϕx, σx, μx)

    #this structure allows us to compute multiple statistics
    #println(res)
    for s ∈ keys(res)

      @assert ismissing.(results[s][n, :]) |> all
      results[s][n,:] .= res[s]
    end
  end

  return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
end

#this is just a wrapper for the 2-vector case
function compareundernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc, comparefunc::Tcomparefunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx<:AbstractVector, Tstatfunc, Taggfunc, Tcomparefunc}

  #=function vstatfunc(X,y)
    #println("X $(size(X))")
    rawstats = statfunc(X |> vec, y)
    return Dict(k=>rawstats[k] for (k,v) ∈ zip(keys(rawstats), values(rawstats)))
  end
  vaggfunc(vr,s) = aggfunc(vr |> vec, s)=#

  return comparefunc(reshape(x,length(x),1) |> Matrix{Float64}, y, statfunc, aggfunc;
    Nsimulations, checkstationarity, stationaritytol)
end

#wrapper for the wrapper
comparear1undernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx<:AbstractVector, Tstatfunc, Taggfunc} = compareundernull(
      x,y,statfunc,aggfunc, comparear1undernull;
      Nsimulations, checkstationarity, stationaritytol)

#=function comparearima11undernull_old(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4) where {Tx, Tstatfunc, Taggfunc}

    #difference everything
    dx = x[2:end] .- x[1:(end-1)]
    dy = x[2:end] .- y[1:(end-1)]

    dxar1 = AR1(dx) |> centerar1
    dyar1 = AR1(dy) |> centerar1

    results = Dict(s => Vector{MFloat64}(undef, Nsimulations) for s ∈ keys(statfunc(dx,dy)))

    #this assumes stationary results are meaningless
    if checkstationarity && max(abs(dxar1.ϕ), abs(dyar1.ϕ)) ≥ 1.0-stationaritytol
      dxar1 = AR1(NaN, NaN, dxar1.μ, dxar1.σ²)
      dyar1 = AR1(NaN, NaN, dxar1.μ, dxar1.σ²)
    end

    dxs = [similar(dx) for t ∈ 1:Threads.nthreads()]
    dys = [similar(dy) for t ∈ 1:Threads.nthreads()]

    #draw a sample and compute the aggregate function, sotring the reuslts in an array
    Threads.@threads for n ∈ 1:Nsimulations
      t::Int = Threads.threadid()
      res = compareoncear1xy!(dxar1, dxs[t], dyar1, dys[t], statfunc)

      #this structure allows us to compute multiple statistics
      for s ∈ keys(res)
        @assert results[s][n] ≡ missing
        results[s][n] = res[s]
      end
    end

    return (; Dict(s=>aggfunc(skipmissing(results[s]), s) for s ∈ keys(results))...)
end=#

function comparearima11undernull(X::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx<:AbstractMatrix, Tstatfunc, Taggfunc}

  #difference everything
  dX = reduce(hcat, (xi->xi[2:end] .- xi[1:(end-1)]).(eachcol(X)))
  dx = dX |> eachcol |> collect
  dy = y[2:end] .- y[1:(end-1)]

  K = length(dx)

  dxar1 = (AR1.(dx)) .|> centerar1
  dyar1 = AR1(dy) |> centerar1

  results = Dict(s => missings(Float64, Nsimulations, K) for s ∈ keys(statfunc(dX,dy)))
  @assert (k->ismissing.(results[k]) |> all).(keys(results)) |> all

  #this assumes stationary results are meaningless
  if checkstationarity && max((dxar1i -> abs(dxar1i.ϕ)).(dxar1)..., abs(dyar1.ϕ)) ≥ 1.0-stationaritytol
    dxar1 .= ((dxar1i)->AR1(NaN, NaN, dxar1i.μ, dxar1i.σ²)).(dxar1)
    dyar1 = AR1(NaN, NaN, dyar1.μ, dyar1.σ²)
  end

  dXs = [similar(dX) for t ∈ 1:Threads.nthreads()]
  dxs = dXs .|> eachcol .|> collect
  dys = [similar(dy) for t ∈ 1:Threads.nthreads()]

  #draw a sample and compute the aggregate function, sotring the reuslts in an array
  Threads.@threads for n ∈ 1:Nsimulations
  #for n ∈ 1:Nsimulations

    t::Int = Threads.threadid()
    #res = compareoncear1xy!(xar1, xs[1], yar1, ys[1], statfunc)
    res = compareoncear1xy!(dxar1, dXs[t], dxs[t], dyar1, dys[t], statfunc)

    #this structure allows us to compute multiple statistics
    for s ∈ keys(res)

      @assert ismissing.(results[s][n, :]) |> all
      results[s][n,:] .= res[s]
    end
  end

  return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
end

comparearima11undernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx<:AbstractVector, Tstatfunc, Taggfunc} = compareundernull(
      x,y,statfunc,aggfunc, comparearima11undernull;
      Nsimulations, checkstationarity, stationaritytol)

nonparametricnullp(v::AbstractVector{T}, x) where T<:Real = (
  x ≥ 0 ? sum(v .> x)/length(v) : sum(v .< x)/length(v))
nonparametricnullp(v::AbstractVector, x) = nonparametricnullp(v |> skipmissing |> collect, x)
nonparametricnullp(M::AbstractMatrix, x) = (
  v->nonparametricnullp(v |> skipmissing |> collect, x)).(eachcol(M))

function testar1(;N)
  ϕ = 0.8
  ε = rand(Normal(), N) + rand(Normal(1,1), N)
  x = missings(Float64, N)
  x[1] = ε[1]

  global x = x
  for t ∈ 2:N
    x[t] = ϕ*x[t-1] + ε[t]
  end
  @info "ϕ=$ϕ, b=$(mean(ε)), μ=$(mean(x))"

  checkmath2(x)

  #=estolsar1 = @btime AR1($x, estimationmethod=estimatear1ols)
  @info "ols conditional: $estolsar1"

  estmlear1 = @btime AR1($x, estimationmethod=estimatear1mle)
  @info "mle: $estmlear1"

  estmle2ar1 = @btime AR1($x, estimationmethod=estimatear1mlecheck)
  @info "mle2: $estmle2ar1"=#

  actar1=AR1(;ϕ, b=0.5, σ²=2)
  xsim = simulatear1(actar1, 10^6)
  @assert isapprox(mean(xsim), 2.5, atol=0.01)  "mean(xsim): $(mean(xsim))"
  @assert isapprox(var(xsim), 2/(1-ϕ^2), atol=0.01) "var(xsim): $(var(xsim))"
  ar1sim = AR1(xsim)
  @info "ar1sim: $(ar1sim)"
  @assert isapprox(ar1sim.ϕ, ϕ, atol=0.01)
  @assert isapprox(ar1sim.b, 0.5, atol=0.01)
  @assert isapprox(ar1sim.μ, 2.5, atol=0.01)
  @assert isapprox(ar1sim.σ², 2, atol=0.01)

  #=xsim! = simulatear1(actar1, Vector{Float64}(undef,10^6))
  @assert isapprox(mean(xsim), 2.5, atol=0.01)  "mean(xsim): $(mean(xsim))"
  @assert isapprox(var(xsim), 2/(1-ϕ^2), atol=0.01) "var(xsim): $(var(xsim))"
  ar1sim = AR1(xsim!)
  @info "ar1sim: $(ar1sim)"
  @assert isapprox(ar1sim.ϕ, ϕ, atol=0.01)
  @assert isapprox(ar1sim.b, 0.5, atol=0.01)
  @assert isapprox(ar1sim.μ, 2.5, atol=0.01)
  @assert isapprox(ar1sim.σ², 2, atol=0.01)=#

  Nsimulations=10^4
  dmeantest = 0.2
  #recall the compare func will pass the name of the relevant stat func when it is called
  aggfunc(v, s) = nonparametricnullp(v, Dict(:cor=>0.1, :dmean=>dmeantest)[s])
  statfunc(x,y) = Dict(:cor=>cor(x,y), :dmean=>mean(x)-mean(y))

  @info "std(x): $(std(x))"
  @info "expected std(x): $((2/(1-ϕ^2))^0.5)"
  #derive the distribution for the difference in means of two AR(1) dists
  R = [0.8^abs(i-j) for i ∈ 1:N, j ∈ 1:N]
  Sdmean = (2/N^2*(ones(N)'* ( R .* 2/(1-ϕ^2)) * ones(N)))^0.5
  Epdmean = cdf(Normal(), - dmeantest/Sdmean)
  res = comparear1undernull(x, deepcopy(x), statfunc, aggfunc; Nsimulations)
  @info "compareundernull result: $res (Epdmean: $Epdmean)"
  y = deepcopy(x)

  @btime comparear1undernull($x, $y, $statfunc, $aggfunc; Nsimulations=$Nsimulations)
  throw("simulation complete")

  X = hcat(x,x .* 2.0)
  mstatfunc(X,y) = Dict(:cor=>cor(X,y) |> vec, :cov=>cov(X,y) |> vec)
  maggfunc(M,s) = (v->nonparametricnullp(v, Dict(:cor=>0.1, :cov=>0.1)[s])).(eachcol(M))

  res = comparear1undernull(X, y, mstatfunc, maggfunc; Nsimulations)
  @info "compareundernull matrix result: $res"

  ix = cumsum(x .+ 0.2)
  iy = cumsum(x .+ 0.2)
  res = comparear1undernull(ix, deepcopy(ix), statfunc, aggfunc; Nsimulations)
  @info "integrated compareundernull (mispecified) dmean: $res"
  res = comparearima11undernull(ix, deepcopy(ix), statfunc, aggfunc; Nsimulations)
  @info "integrated compareundernull dmean: $res"
  res = comparear1undernull(X, deepcopy(x), mstatfunc, maggfunc; Nsimulations)
  @info "integrated compareundernull matrix result: $res"

end

@time testar1(N=1000)


import GeometricMultigrid: δ,inside

# set up examples
function static(;n=32,N=2,T::Type=Float32,m=rand(T,N))
    dims = ntuple(i->n+2,N)
    L,x = zeros(T,dims...,N), zeros(T,dims)
    m ./= √sum(abs2,m)
    for I ∈ inside(x)
        x[I] = sum( @. m*I.I)
        for i ∈ 1:N
            I.I[i] !=2 && (L[I,i] = 1)
        end
    end
    A,x = Poisson(L),FieldVector(x)
    A,A*x
end

function dipole(;n=32,N=2,T::Type=Float32,x0=rand(T,N),m=rand(T,N),s=rand(),Δt=1)
    dims = ntuple(i->n+2,N)
    L,b = zeros(T,dims...,N), zeros(T,dims)
    x0 .*= n/3; x0 .+= (n/3); s *= n/4; m ./= √sum(abs2,m)
    f(x) = sum(@. m*(x-x0))*exp(-s^2*sum(abs2,x.-x0))
    for I ∈ inside(b)
        b[I] = f(I.I)
        for i ∈ 1:N
            I.I[i] !=2 && (L[I,i] = Δt)
        end
    end
    b .-= sum(b)/n^N
    Poisson(L),FieldVector(b)
end

@fastmath kern₀(d) = 0.5+0.5d+0.5sin(π*d)/π
μ₀(d) = kern₀(clamp(d,-1,1))
function sphere(;n=32,N=2,T::Type=Float32,x0=rand(T,N),m=rand(T,N),s=rand(),Δt=1)
    dims = ntuple(i->n+2,N)
    L,b = zeros(T,dims...,N), zeros(T,dims)
    x0 .*= n/3; x0 .+= (n/3); s *= n/4; m ./= √sum(abs2,m)
    d(x) = √sum(abs2,x.-x0)-s
    for I ∈ inside(b), i ∈ 1:N
        μp = μ₀(d(I.I .+ 0.5 .* δ(i,N).I))
        μm = μ₀(d(I.I .- 0.5 .* δ(i,N).I))
        I.I[i] !=2 && (L[I,i] = Δt*μm)
        b[I] += Δt*m[i]*(μp-μm)
    end
    Poisson(L),FieldVector(b)
end

using DataStructures
function create_synthetic(;len=100,kw...)
    OrderedDict(
        "2D-static"=>[static(N=2;kw...) for i ∈ 1:len],
        "3D-static"=>[static(N=3;kw...) for i ∈ 1:len],
        "2D-dipole"=>[dipole(N=2;kw...) for i ∈ 1:len],
        "3D-dipole"=>[dipole(N=3;kw...) for i ∈ 1:len],
        "2D-sphere"=>[sphere(N=2;kw...) for i ∈ 1:len],
        "3D-sphere"=>[sphere(N=3;kw...) for i ∈ 1:len]
    )
end


"""
    braycurtis(se::SummarizedExperiment, assay_name::String)

Evaluates the braycurtis dissimilarity index for each sample and rturns it as a matrix.

# Arguments
- `se::SummarizedExperiment`: the experiment object of interest.
- `assay_name::String`: the name of the assay to base the evaluation on.
"""
function braycurtis(se::SummarizedExperiment, assay_name::String)

    pairwise(BrayCurtis(), assay(se, assay_name), dims = 2)

end

"""
    jaccard(se::SummarizedExperiment, assay_name::String)

Evaluates the jaccard dissimilarity index for each sample and rturns it as a matrix.

# Arguments
- `se::SummarizedExperiment`: the experiment object of interest.
- `assay_name::String`: the name of the assay to base the evaluation on.
"""
function jaccard(se::SummarizedExperiment, assay_name::String)

    pairwise(Jaccard(), assay(se, assay_name), dims = 2)

end

"""
    hellinger(se::SummarizedExperiment, assay_name::String)

Evaluates the hellinger dissimilarity index for each sample and rturns it as a matrix.

# Arguments
- `se::SummarizedExperiment`: the experiment object of interest.
- `assay_name::String`: the name of the assay to base the evaluation on.
"""
function hellinger(se::SummarizedExperiment, assay_name::String)

    pairwise(HellingerDist(), assay(se, assay_name), dims = 2)

end

"""
    pcoa(se::SummarizedExperiment, assay_name::String)

Runs a Multidimensional Scaling (MDS) for the selected assay of se and returns a
model that can be executed with `predict(model)`.

# Arguments
- `se::SummarizedExperiment`: the experiment object of interest.
- `assay_name::String`: the name of the assay to base the evaluation on.
"""
function pcoa(se::SummarizedExperiment, assay_name::String; dist = braycurtis, dim_number = 3)

    return fit(MDS, dist(se, assay_name), maxoutdim = dim_number, distances = true)

end


# Semidiscrete KT Scheme: Second-Order
# Based on:
# Kurganov, Tadmor. New High Resolution Central Schemes for Non Linear Conservation
# Laws and Convection-Difussion Equations. Journal of Comp Physics 160, pp 241-282. 2000

struct FVSKTScheme{ltype <: AbstractSlopeLimiter} <: AbstractFVAlgorithm
  slopeLimiter :: ltype
end

function FVSKTScheme(;slopeLimiter=GeneralizedMinmodLimiter())
  FVSKTScheme(slopeLimiter)
end

function update_flux_value(u,∇u,j,mesh,Flux, alg::FVSKTScheme)
    # Local speeds of propagation
    uminus=cellval_at_left(j,u,mesh)+0.5*cellval_at_left(j,∇u,mesh)
    uplus=cellval_at_right(j,u,mesh)-0.5*cellval_at_right(j,∇u,mesh)
    ul = cellval_at_left(j,u,mesh)
    ur = cellval_at_right(j,u,mesh)
    aa = max(fluxρ(uminus,Flux),fluxρ(uplus,Flux))
    # Numerical Fluxes
    0.5*(Flux(uplus)+Flux(uminus)) - aa/2*(uplus - uminus)
end

"""
compute_fluxes!(hh, Flux, u, mesh, dt, alg::FVSKTScheme, ::Type{Val{true}})
Numerical flux of Kurkanov Tadmor scheme in 1D
"""
function compute_fluxes!(fluxes, Flux, u, mesh, dt, alg::FVSKTScheme, noscalar::Bool, ::Type{Val{true}})
    slopeLimiter = alg.slopeLimiter
    #update vector
    # 1. slopes
    ∇u = compute_slopes(u, mesh, slopeLimiter, noscalar, Val{true})

    #update vector
    for j in node_indices(mesh)
        if noscalar
            fluxes[:,j] .= update_flux_value(u,∇u,j,mesh,Flux,alg)
        else
            fluxes[j] = update_flux_value(u,∇u,j,mesh,Flux,alg)
        end
    end
end

function compute_fluxes!(fluxes, Flux, u, mesh, dt, alg::FVSKTScheme, noscalar::Bool, ::Type{Val{false}})
    slopeLimiter = alg.slopeLimiter
    #update vector
    # 1. slopes
    ∇u = compute_slopes(u, mesh, slopeLimiter, noscalar, Val{false})

    #update vector
    for j in node_indices(mesh)
        if noscalar
            fluxes[:,j] .= update_flux_value(u,∇u,j,mesh,Flux,alg)
        else
            fluxes[j] = update_flux_value(u,∇u,j,mesh,Flux,alg)
        end
    end
end


function gpdps_enep!(u,v,data,params)
    σ,τ,ω,maxit = params
    S,Sᵀ,z₁,z₂,α,a,b,ue₁,ue₂ = data

    # preallocate
    uold = similar(u)
    ξ = similar(u)
    u₁,u₂ = view(u,:,1), view(u,:,2)
    v₁,v₂ = view(v,:,1), view(v,:,2)
    ξ₁,ξ₂ = view(ξ,:,1), view(ξ,:,2)
    ρ₁,ρ₂ = similar(u₁), similar(v₁)
    Kᵤ,Kᵥ = similar(u), similar(v)
    for it = 1:maxit
        # Kx
        yᵥᵤ = S(v₁,u₂)
        yᵤᵥ = S(u₁,v₂)
        yᵤᵤ = S(u₁,u₂)
        @. ρ₁ = 2yᵤᵤ - yᵤᵥ - z₁
        @. ρ₂ = 2yᵤᵤ - yᵥᵤ - z₂
        Kᵤ .= Sᵀ(ρ₁,ρ₂) .+ α.*u

        # print output
        Psi = (norm(yᵤᵤ-z₁)^2/2 + α/2*norm(u₁) - norm(yᵥᵤ-z₁)^2/2 - α/2*norm(v₁) +
               norm(yᵤᵤ-z₂)^2/2 + α/2*norm(u₂) - norm(yᵤᵥ-z₂)^2/2 - α/2*norm(v₂)) /
              length(z₁)
        err = (norm(u₁-ue₁)^2 + norm(u₂-ue₂)^2)/length(z₁)
        @printf("it %3d: Psi = %1.3e\t Err = %1.3e\n", it, Psi, err)

        # x = prox_G(x-τKx)
        @. uold = u
        @. u = clamp(u-τ*Kᵤ,a,b)
        @. ξ = u + ω*(u-uold)

        # Ky
        yᵥᵤ = S(v₁,ξ₂)
        yᵤᵥ = S(ξ₁,v₂)
        @. ρ₁ = z₁ - yᵤᵥ
        @. ρ₂ = z₂ - yᵥᵤ
        Kᵥ .= Sᵀ(ρ₁,ρ₂) .- α.*v

        # y = prox_F*(y+σKy)
        @. v = clamp(v+σ*Kᵥ,a,b)
    end
end

function setup_enep(N,α,a,b)
    x₁,x₂,A = setup_fdm(N)

    B₁ = x₂.<0.5
    B₂ = x₂.>0.5

    y  = @. sin(π*x₁)*sin(π*x₂)
    p₁ = @. sin(2π*x₁)*sin(2π*x₂)
    p₂ = @. sin(3π*x₁)*sin(3π*x₂)

    z₁ = y .+ A*p₁
    z₂ = y .+ A*p₂

    u₁ = clamp.(B₁.*p₁./α,a,b)
    u₂ = clamp.(B₂.*p₂./α,a,b)

    f = A*y .- B₁.*u₁ .- B₂.*u₂

    C  = cholesky(A)
    S  = (u₁,u₂) -> C\(B₁.*u₁ .+ B₂.*u₂ .+ f)
    Sᵀ = (q₁,q₂) -> hcat(B₁.*(C\q₁),B₂.*(C\q₂))

    return S,Sᵀ,z₁,z₂,u₁,u₂
end

function setup_fdm(N)
    dx = 1/(N+1)
    xm = range(dx,stop=1-dx,length=N)
    xx,yy = reshape(xm, 1, N), reshape(xm, N, 1)
    x₁,x₂ = repeat(xx,outer=(N,1))[:], repeat(yy,outer=(1,N))[:]

    e  = fill(1/dx^2,N)
    D₂ = spdiagm(-1=>-e[2:end],0=>2e,1=>-e[2:end])
    Id = sparse(I,N,N)
    A  = kron(D₂,Id)+kron(Id,D₂)
    return x₁,x₂,A
end


using Plots, Distributions, LinearAlgebra, ForwardDiff

n = 100
m = 50

Q = Diagonal(ones(n))

A = rand(m, n)
b = A * randn(n)
z_sol = 1.0 * [2 * Q A'; A zeros(m, m)] \ [zeros(n); b]
# @show norm(lagrangian_gradient(z_sol))
# @show lagrangian(z_sol)

function obj(x)
	x' * Q * x
end

function gradient(x)
	2.0 * Q * x
end

function lagrangian(z)
	x = view(z, 1:n)
	y = view(z, n .+ (1:m))

	obj(x) + y' * (A * x - b)
end

function lagrangian_gradient(z)
	x = view(z, 1:n)
	y = view(z, n .+ (1:m))

	[gradient(x) + A' * y; A * x - b]
end

dir = Diagonal([ones(n); -1.0 * ones(m)])

# adam
function adam(obj, gradient, x;
		max_iter = 1000)

	# history
	J_hist = [obj(x)]

	# parameters
	α = 1.0e-4
	β1 = 0.9
	β2 = 0.999
	ϵ = 10.0^(-8.0)
	m = zero(x)
	v = zero(x)
	m̂ = zero(x)
	v̂ = zero(x)

	# iterate
	for i = 1:max_iter
		g = gradient(x)
		# m .= β1 .* m + (1.0 - β1) .* g
		# v .= β2 .* v + (1.0 - β2) .* (g.^2.0)
		# m̂ .= m ./ (1.0 - β1^i)
		# v̂ .= v ./ (1.0 - β2^i)
		# x .-= α * m̂ ./ (sqrt.(v̂) .+ ϵ)
		x .-= α * dir * g
		push!(J_hist, obj(x))
		println("norm $(norm(gradient(x)))")
		i % 1000 == 0 && println("iter: $i")
	end

	return x, J_hist
end

x_sol, J_hist = adam(lagrangian, lagrangian_gradient, ones(n + m),
	max_iter = 100000)

plt = plot(J_hist,
	xlabel = "iteration", ylabel = "obj.", title = "sgd",
	# label = "α = $α (adam)",
	color = :red, width = 2.0)

display(plt)


@doc raw"""
    airydisk([T=Float64], point; x, y, fwhm, ratio=0, amp=1, theta=0, bkg=0)
    airydisk([T=Float64], px, py; x, y, fwhm, ratio=0, amp=1, theta=0, bkg=0)

An unnormalized Airy disk. The position can be specified in `(x, y)` coordinates as a `Tuple`, `AbstractVector`, or as separate arguments. If `theta` is given, the PSF will be rotated by `theta` degrees counter-clockwise from the x-axis. If `bkg` is given it will be added as a scalar to the PSF.

The `fwhm` can be a scalar (isotropic) or a vector/tuple (diagonal). Keep in mind that `theta` has no effect for isotropic distributions and is degenerate with the `fwhm` parameters (i.e., theta=90 is the same as reversing the `fwhm` tuple)

If `ratio` is supplied, this will be the Airy pattern for a centrally-obscured aperture (e.g., a Newtonian telescope). This has a slightly expanded functional form, and in general the central Airy disk will be smaller and the first Airy ring will be brighter.

# Functional form

The Airy disk is a distribution over the radius `r` (the square-Euclidean distance)

```
f(x | x̂, FWHM) = [ 2J₁(q) / q ]^2
```
where `J₁` is the first-order Bessel function of the first kind and
```
q ≈ π * r * D/ λ ≈ π * r / (0.973 * FWHM)
```

If user a non-zero central obscuration via `ratio`, the functional form becomes

```
f(x | x̂, FWHM, ϵ) = [ 2J₁(q) / q - 2ϵJ₁(ϵq) / q ]^2 / (1 - ϵ^2)^2
```
where `ϵ` is the ratio (0 ≤ `ϵ` < 1).
"""
airydisk(T, px, py; x, y, fwhm, ratio=0, amp=one(T), theta=0, bkg=0) = convert(T, _airydisk(px, py, x, y, fwhm, ratio, amp, theta, bkg))

# factor for scaling radius in terms of the fwhm
const AIRY_PRE = π / 0.973

# isotropic
function _airydisk(px, py, x, y, fwhm, ratio, amp, theta, background)
    # find offset from center
    dx = px - x
    dy = py - y
    # rotate
    !iszero(theta) && @warn "isotropic airydisk is not affected by non-zero rotation angle $theta"
    # unnormalized airydisk
    r = sqrt(dx^2 + dy^2)
    # short-circuit
    iszero(r) && return amp / (1 - ratio^2)^2 + background
    q = AIRY_PRE * r / fwhm
    I1 = 2 * besselj1(q) / q
    if iszero(ratio)
        return amp * I1^2 + background
    end
    I2 = 2 * ratio * besselj1(q * ratio) / q
    return amp * ((I1 - I2) / (1 - ratio^2))^2 + background
end

# bivariate
function _airydisk(px, py, x, y, fwhm::BivariateLike, ratio, amp, theta, background)
    # find offset from center
    dx = px - x
    dy = py - y
    # rotate
    if !iszero(theta)
        dx, dy = rotate_point(dx, dy, theta)
    end
    # unnormalized airy disk
    fwhmx, fwhmy = fwhm
    q = AIRY_PRE * sqrt((dx / fwhmx)^2 + (dy / fwhmy)^2)
    # short-circuit
    iszero(q) && return amp / (1 - ratio^2)^2 + background
    I1 = 2 * besselj1(q) / q
    if iszero(ratio)
        return amp * I1^2 + background
    end
    I2 = 2 * ratio * besselj1(q * ratio) / q
    return amp * ((I1 - I2) / (1 - ratio^2))^2 + background
end


module SinglePendulum

using NeuralNetworkAnalysis, LaTeXStrings
import Plots, DifferentialEquations
using Plots: plot, plot!, savefig, font, Measures.mm

# problem

m = 0.5
L = 0.5
c = 0.0
g = 1.0
gL = g/L
mL = 1/(m*L^2)

@taylorize function single_pendulum!(dx, x, p, t)
    dx[1] = x[2]
    dx[2] = gL * sin(x[1]) + mL*(x[3] - c*x[2])
    dx[3] = zero(x[3])
    return dx
end

controller = read_nnet_sherlock(@modelpath "" "controller_Single-Pendulum_sherlock")

X0 = Hyperrectangle([1.1, 0.1], [0.1, 0.1])
U0 = ZeroSet(1)
ivp = @ivp(x' = single_pendulum!(x), dim: 3, x(0) ∈ X0 × U0)
vars_idx = Dict(:state_vars=>1:2, :control_vars=>3)

period = 0.05
k = 20
T = k * period
T_warmup = 2 * period

prob = ControlledPlant(ivp, controller, vars_idx, period)

# simulation

println("simulation")
res = @timed simulate(prob, T=T, trajectories=10, include_vertices=true)
sim = res.value
print_timed(res)

# reachability analysis

alg = TMJets(abstol=1e-7, orderT=4, orderQ=1)
alg_nn = Ai2()

function benchmark(; T=T, silent::Bool=false)
    silent || println("flowpipe construction")
    res = @timed solve(prob, T=T, alg_nn=alg_nn, alg=alg)
    sol = res.value
    silent || print_timed(res)

    return sol
end

benchmark(T=T_warmup, silent=true)
res = @timed benchmark()
sol = res.value

# plotting

vars = (1, 2)
fig = plot(xlab=L"\theta", ylab=L"\theta'",
           xlims=(1, 2.6), ylims=(-0.5, 2.5),
           leg=:topleft,
           legendfontsize=5,
           tickfont=font(25, "Times"),
           guidefontsize=25,
           xguidefont=font(25, "Times"),
           yguidefont=font(25, "Times"),
           bottom_margin=0mm, left_margin=0mm, right_margin=0mm, top_margin=2mm, size=(500, 450))
plot!(fig, sol, vars=vars, color=:orange, lab="")
plot!(project(X0, vars), c=:blue, lab=L"X_0", alpha=0.6)
plot_simulation!(fig, sim; vars=vars, color=:black, lab="")
savefig("singlependulum_juliareach.png")

end
nothing


using SymbolicUtils: Sym, FnType, Term, Add, Mul, Pow, symtype, operation, arguments
using SymbolicUtils
using IfElse: ifelse
using Test

@syms w::Real z::Real a::Real b::Real x

addr = addr = Add(Real, 10, Dict(w=>1, z=>1))

substitute(addr, Dict(w=>1, z=>3.0)) |> display

fieldnames(typeof(addr)) |> display

addr.coeff |> display


using Test
using WavePropBase
using WavePropBase.Geometry
using WavePropBase.Integration
using WavePropBase.Interpolation
using WavePropBase.Mesh
using LinearAlgebra
using StaticArrays

@testset "Lagrange elements" begin
    @testset "LagrangeLine" begin
        ## line in 3d
        vtx = ((0.,0.,0.),(1.,1.,1.))
        l   = LagrangeLine(vtx)
        @test domain(l) == ReferenceLine()
        @test length(return_type(l)) == 3
        @test dimension(domain(l)) == 1
        # makes no sense to ask for normal here. Make sure error is thrown.
        @test_throws AssertionError normal(l,SVector(0.5)) == SVector(-1.,1.)/√2
        ## line in 2d
        a = (0.,0.)
        b = (1.,1.)
        l   = LagrangeLine((a,b))
        @test domain(l) == ReferenceLine()
        @test length(return_type(l)) == 2
        @test dimension(domain(l)) == 1
        @test normal(l,SVector(0.5)) == SVector(1.,-1.)/√2
    end
    @testset "LagrangeTriangle" begin
        # triangle in 2d
        vtx = SVector(SVector(0.,0.),SVector(0.,1.),SVector(-1.,0))
        t   = LagrangeTriangle(vtx)
        @test length(return_type(t)) == 2
        @test dimension(domain(t)) == 2
        # triangle in 3d
        vtx = SVector(SVector(0.,0.,0.),SVector(0.,1.,0.),SVector(-1.,0,0.))
        t   = LagrangeTriangle(vtx)
        @test length(return_type(t)) == 3
        @test dimension(domain(t)) == 2
        @test normal(t,SVector(0.1,0.1)) == SVector(0,0,1.)
    end
    @testset "Tetrahedron" begin
        # TODO: add tests
    end
end


module NNBasicFunction

  export hardlim , hardlims, purelin, satlin, satlins, logsig, tansig, poslin;

   function hardlim(x)
		if (x<0)
	     return 0
	    else
	    return  +1
	   end
	end

    function hardlims(x)
		if (x<0)
	     return -1
	    else
	    return +1
	   end
    end

    function purelin(x)
	   x
   	end

   function satlin(x)
	   if(x<0)
		   return 0
	   elseif(0<=x<=1)
		   return x
	   else
		   return 1
	   end
   end

	function satlins(x)
	   if(x<-1)
		   return -1
	   elseif(-1<=x<=1)
		   return x
	   else
		   return 1
	   end
   end

	function logsig(x)
		let top=1,bottom=1+ℯ^-x
			return top/bottom
		end
	end

	function tansig(x)
		let first= ℯ^x, last=ℯ^-x
			top= first-last
			bottom=first+last
		return top/bottom
		end
	end
	function poslin(x)
		if(x<0)
		  return 0
		else
		  return x
		end
	end

end


using ApproxFun, Base.Test

single_sin = Fun(sin,Interval(0.f0,1.f0))
double_sin = Fun(sin,Interval(0.,1.))
big_sin = Fun(sin,Interval(parse(BigFloat,"0.0"), parse(BigFloat,"1.0")))

@test ncoefficients(single_sin) <= ncoefficients(double_sin)
@test ncoefficients(double_sin) <= ncoefficients(big_sin)

@test eltype(coefficients(single_sin)) == Float32
@test eltype(coefficients(double_sin)) == Float64
@test eltype(coefficients(big_sin)) == BigFloat

single_double_err = coefficients(single_sin-double_sin)[1:ncoefficients(single_sin)]
@test norm(single_double_err) < 10eps(Float32)

single_double_err = coefficients(double_sin-big_sin)[1:ncoefficients(double_sin)]
@test norm(single_double_err) < 10eps(Float64)


function move(instructions, dir = Complex(1, 0), way = Complex(10, 1), pos = Complex(0, 0))
    pos1 = pos2 = pos
    directions = Dict([("N", Complex(0, 1)), ("S", Complex(0, -1)), ("E", Complex(1, 0)), ("W", Complex(-1, 0))])
    for instruction in instructions
        (cmd, num) = match(r"([NSEWLRF])([0-9]+)", instruction).captures
        val = parse(Int, num)
        if cmd == "F"
            pos1, pos2 = map(x -> x[1] += x[2] * val, [[pos1, dir], [pos2, way]])
        elseif cmd == "L"
            dir, way = map(x -> x *= Complex(0, 1)^div(val, 90), [dir, way])
        elseif cmd == "R"
            dir, way = map(x -> x *= Complex(0, -1)^div(val, 90), [dir, way])
        else
            pos1, way = map(x -> x += val * get(directions, cmd, nothing), [pos1, way])
        end
    end
    return map(x -> abs(real(x)) + abs(imag(x)), [pos1, pos2])
end

instructions = readlines("d12.in")
(pos1, pos2) = move(instructions)
println("P1: ", pos1)
println("P2: ", pos2)


using Distributions
#Using the NonCentral T Distribution
#Use an F-test to calculate power for a two-sided test
function CRT_clustersize(n1, rho, d, alpha, power_level)
    n2 = 3
    while n2 < 100
        ncp = d/(sqrt(4*(1+(n1-1)*rho)/(n1*n2)))
        #println(ncp)
        crit = quantile(FDist(1,n2-2),1-alpha)
        power_calc = 1 - cdf(NoncentralF(1,n2-2,ncp^2),crit)
        #println("p",power_calc)
        if power_calc > power_level
            cluster_num = n2
            return cluster_num
         end
        n2 +=1
    end
end


immutable ExponentialIntercept{T} <: Model
    β₀::T
    a::T
    δ::T
end

function Base.show(io::IO, ::Type{ExponentialIntercept})
    return print(io, "Intercept Exponential")
end

function Base.show(io::IO, m::ExponentialIntercept)
    @printf(
        io,
        "Intercept Exponential (β₀ = %.3f, a = %.3f, δ = %.3f)",
        m.β₀,
        m.a,
        m.δ,
    )
end

function initial_parameters(::Type{ExponentialIntercept})
    return Float64[0.0, log(1.0), logit(0.5)]
end

canonical_parameters(m::ExponentialIntercept) = (m.β₀, m.a, m.δ)

function ExponentialIntercept(parameters::Vector)
    β₀ = parameters[1]
    a = exp(parameters[2])
    δ = invlogit(parameters[3])
    return ExponentialIntercept(β₀, a, δ)
end

@inline function predict(
    m::ExponentialIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, δ = m.β₀, m.a, m.δ
    u1 = x1 * δ^t1
    u2 = x2 * δ^t2
    z = β₀ + a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::ExponentialIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, δ = m.β₀, m.a, m.δ
    u2 = x2 * δ^t2
    u1 = x1 * δ^t1
    gr = (
        1.0,
        (u2 - u1) * a,
        a * ((x2 * t2 * δ^t2 - x1 * t1 * δ^t1) / δ) * δ * (1 - δ),
    )
    return gr
end


using Test
using SumOfSquares
using DynamicPolynomials

function univariate_sum_test(optimizer,
                   config::MOIT.Config,
                   cone::SumOfSquares.PolyJuMP.PolynomialSet)
    atol = config.atol
    rtol = config.rtol

    model = _model(optimizer)

    @polyvar x y
    # (x - 1)^2 + (y + 1)^2
    cref = @constraint(model, x^2 + y^2 + 2(y - x) + 2 in cone, sparsity=Sparsity.Variable())

    optimize!(model)

    @test termination_status(model) == MOI.OPTIMAL

    @test primal_status(model) == MOI.FEASIBLE_POINT

    p = gram_matrix(cref)
    @test p isa SumOfSquares.SparseGramMatrix
    @test length(p.sub_gram_matrices) == 2
    @test getmat(p.sub_gram_matrices[1]) ≈ ones(2, 2) atol=atol rtol=rtol
    @test p.sub_gram_matrices[1].basis.monomials == [y, 1]
    @test getmat(p.sub_gram_matrices[2]) ≈ [1 -1; -1 1] atol=atol rtol=rtol
    @test p.sub_gram_matrices[2].basis.monomials == [x, 1]

    S = SumOfSquares.SOSPolynomialSet{
        SumOfSquares.FullSpace, Monomial{true}, MonomialVector{true}, SumOfSquares.Certificate.Sparsity.Ideal{Sparsity.Variable, SumOfSquares.Certificate.MaxDegree{typeof(cone), MonomialBasis}}
    }
    @test list_of_constraint_types(model) == [(Vector{AffExpr}, S)]
    test_delete_bridge(
        model, cref, 0,
        ((MOI.VectorOfVariables, MOI.Nonnegatives, 0),
         (MOI.VectorOfVariables, MOI.PositiveSemidefiniteConeTriangle, 0)
        ))
end
sos_univariate_sum_test(optimizer, config)   = univariate_sum_test(optimizer, config, SOSCone())
sd_tests["sos_univariate_sum"] = sos_univariate_sum_test
sdsos_univariate_sum_test(optimizer, config) = univariate_sum_test(optimizer, config, SDSOSCone())
soc_tests["sdsos_univariate_sum"] = sdsos_univariate_sum_test
dsos_univariate_sum_test(optimizer, config)  = univariate_sum_test(optimizer, config, DSOSCone())
linear_tests["dsos_univariate_sum"] = dsos_univariate_sum_test


using Test

n = 10;
A = rand(n, n); A = A + A';

include("ldlt_symm.jl")

# L = Array(Float64, 2)
# D = Array(Float64, 2)
# P = Array(Float64, 2)
# pp = Array(Int, 1)
global L     = Matrix{Float64}(undef, n, n)
global D     = Matrix{Float64}(undef, n, n)
global P     = Matrix{Float64}(undef, n, n)
global pp    = Vector{Float64}(undef, n)
global ρ     = Float64
global ncomp = Int64

try (L, D, pp, ρ, ncomp) = ldlt_symm(A)
catch except
    @test rethrow(except)
end

#@test_approx_eq P*A*P'  L*D*L'
# @test_approx_eq A[pp,pp]  L*D*L'
@test A[pp, pp] ≈ L * D * L'

v = rand(n,1);
#@test_approx_eq P*v  v[pp]
#@test_approx_eq P'*v  v[invperm(pp)]


using SymbolicTensors
using SymPy
using Test
using LinearAlgebra

@testset "SymbolicTensors.jl" begin
    @testset "types" begin
        field = TensorIndexType("field","f")
        @test typeof(field) == TensorIndexType
        @indices field i j
        @test typeof(i) == typeof(j) == TensorIndex
        A = TensorHead("A",[field])
        B = TensorHead("B",[field])
        @test typeof(A) == TensorHead
        @test typeof(A(i)) == IndexedTensor
        @test typeof(A(i)*A(-i)) == SymPy.Sym
        @test typeof(A(i) + B(i)) == TensAdd
        @test typeof(A(i) * B(j)) == TensMul
    end
    @testset "arithmetic" begin
        field = TensorIndexType("field","f")
        @indices field i j k l m n
        A = TensorHead("A",[field])
        B = TensorHead("B",[field])
        C = TensorHead("C",[field,field],get_tsymmetry.fully_symmetric(-2))
        @test A(i) + A(i) == 2*A(i)
        @test A(i) - A(i) == 0
        @test canon_bp(A(i)*A(j) + A(j)*A(i)) == 2*A(i)*A(j)
        @test canon_bp(C(-i,-j)*A(i)*A(j)) == 0
        # -- odd numbers of antisymmetric tensors contracted give zero
        @test canon_bp(C(i,-j)*C(j,-k)*C(k,-i)) == 0
        @test scalarIsEqual(A(i)*A(j)*field.metric(-i,-j),A(i)*A(-i),field.metric)
        @test scalarIsEqual(A(i)*B(j)*field.metric(-i,-j),A(i)*B(-i),field.metric)
        # should check properties of deltas, epsilons here too
        # -- should add some from xTensor
    end
    # -- should calculate some known Riemann scalars
    @testset "derivatives" begin
        field = TensorIndexType("field","f")
        @indices field i j k
        A = TensorHead("A",[field])
        @test diff(A(i),A(j)) == field.metric(i,-j)
        @test diff(A(-i),A(j)) == field.metric(-i,-j)
        @test diff(A(i),A(-j)) == field.metric(i,j)
        g = 1/(1+A(i)*A(-i))^2 * field.delta(-i,-j)
        gg = diff(g(-i,-j),A(k))
        @test diff(field.delta(-i,-j),A(k)) == 0
        # product rule w/ scalar chain rule
        @test scalarIsEqual( diff( A(i)*A(-i)*A(j), A(k)),
            A(i)*A(-i)*field.metric(j,-k) + 2*A(j)*A(-k) ,field.metric)
        # distribute over TensAdd
        @test canon_bp(diff( (A(i)*A(j) + A(j)*A(i))/2, A(k))) ==
            canon_bp(A(i)*field.metric(j,-k) + A(j)*field.metric(i,-k))
        #@test gg(-i,-j,-k) ==
        #h = canon_bp(1//2*(gg(-i,-j,-k) + gg(-i,-k,-j) - gg(-j,-k,-i) ))
        #GG = TensorHead("GG",[field,field,field])
        #h = canon_bp(1//2*(GG(-i,-j,-k) + GG(-i,-k,-j) - GG(-j,-k,-i) ))
        #@test h(i,-j,-k) == h(i,-k,-j)
    end

    @testset "errors" begin
    end

    @testset "replacements" begin
        field = TensorIndexType("field","f")
        @indices field i j k
        A = TensorHead("A",[field])
        field.set_metric(field.delta)
        @test replace_with_arrays(A(i)*A(j)*field.metric(-i,-j), Dict(field.delta(-i,-j)=>[1 0 ; 0 1], A(i)=>[0,4] )) == 16
    end

    @testset "quoting" begin
        vars = symbols("a b c d e f")
        a,b,c,d,e,f = vars
        function equalQuotes(q1,q2)
            if q1.head == :function
                q1 = q1.args[2]
                q2 = q2.args[2]
            end
            global a,b,c,d,e,f
            a,b,c,d,e,f = randn(6)
            return eval(q1) == eval(q2)
        end
        #Quote a scalar
        @test Quote(vars[1]) == convert(Expr,vars[1])
        cse_ex = sympy.cse(a^2 + 1/(1+a^2))
        @test equalQuotes(Quote(cse_ex),quote
              x0 = a ^ 2
              x0 + (1 + x0) ^ -1
        end)
        @test equalQuotes(Quote("hub",cse_ex,[a]),:(function hub(a)
              x0 = a ^ 2
              x0 + (1 + x0) ^ -1
        end))
        arr = [a^2,b^2+a^2,3c^2]
        cse_arr = sympy.cse(arr)
        @test equalQuotes(Quote(arr),quote
            reshape([a ^ 2, a ^ 2 + b ^ 2, __prod__(3, c ^ 2)], (3,))
        end)
        @test equalQuotes(Quote(cse_arr),quote
            x0 = a ^ 2
            reshape([x0, x0 + b ^ 2, __prod__(3, c ^ 2)], (3, 1))
        end)
        @test equalQuotes(Quote("_A",cse_arr,[a,b,c]),:(function _A(a, b, c)
            x0 = a ^ 2
            reshape([x0, x0 + b ^ 2, __prod__(3, c ^ 2)], (3, 1))
        end))
    end
    @testset "Riemann scalar computations" begin
        field = TensorIndexType("field","f")
        @indices field i j k l m n o
        F = TensorHead("F",[field])
        δ = field.delta

        simp(x) = contract_metric(canon_bp(x),field.metric)

        α = symbols("a",positive=true)
        g = 6 * α * δ(-i,-j) / (1-F(i)*F(j)*δ(-i,-j))^2
        gg = diff(g(-i,-j),F(k))
        gg = simp(gg)

        h = (gg(-i,-j,-k) - gg(-j,-k,-i) + gg(-k,-i,-j))/2
        h =  h(-i,-j,-k)
        h = simp(h)

        dh = diff(h(-i,-j,-k),F(l))
        dh = simp(dh)

        hh = (field.metric(m,n)*h(-m,-i,-j))*h(-n,-k,-l)
        hh = simp(hh)

        Riemann = dh(-i,-l,-j,-k) - dh(-i,-k,-l,-j) + hh(-j,-k,-i,-l) - hh(-j,-l,-i,-k)
        Riemann = simp(Riemann)

        Ric = Riemann(-i,-j,-k,-l)*field.metric(i,k)
        Ric = simp(Ric)

        R = Ric(-i,-j)*field.metric(j,i)
        R = simp(R)

        ## turn into arrays
        vars = symbols("x y z",real=true)
        x,y,z = vars
        d = length(vars)
        repl = Dict{Any,Any}(F(m) => vars, field.delta(-m,-n) => Matrix{Int64}(I,d,d))
        garr = simplify(replace_with_arrays(g(-i,-j),repl))

        repl[field] = garr
        #repl[field.metric(m,-n)] = Matrix{Int64}(I,d,d)
        harr = factor.(replace_with_arrays(h,repl))
        dharr = factor.(replace_with_arrays(dh,repl))
        hharr = factor.(replace_with_arrays(hh,repl))
        Rarr = factor(replace_with_arrays(R,repl))

        testpt = Dict(x=>0.1,y=>-0.8,z=>0.05,α=>1)
        @test float(subs(harr[1,1,1],testpt)) ≈ float(subs(-12*x / (-1+x^2+y^2+z^2)^3,testpt))
        @test float(subs(dharr[1,1,1,1],testpt)) ≈ float(subs(12*(1+5*x^2 - y^2 - z^2) / (-1+x^2+y^2+z^2)^4,testpt))
        @test float(subs(Rarr,testpt)) ≈ -4

        h = get_christoffel(F,field,g)
        Riemann = get_riemann(F,h,field,g)
        Ric = Riemann(-i,-j,-k,-l)*field.metric(i,k)
        Ric = simp(Ric)

        R = Ric(-i,-j)*field.metric(j,i)
        R = simp(R)
        vars = symbols("x y z",real=true)
        x,y,z = vars
        d = length(vars)
        repl = Dict{Any,Any}(F(m) => vars, field.delta(-m,-n) => Matrix{Int64}(I,d,d))
        garr = simplify(replace_with_arrays(g(-i,-j),repl))

        repl[field] = garr
        #repl[field.metric(m,-n)] = Matrix{Int64}(I,d,d)
        Rarr = factor(replace_with_arrays(R,repl))
        @test float(subs(Rarr,testpt)) ≈ -4
    end
end


using Hypatia
using Empirikos
using Test

# We compare the chi^2 F-localization intervals against the intervals reported
# by Lord and Cressie.

lord_cressie = LordCressie.load_table()
Zs = LordCressie.ebayes_samples()

idx_test = [1, 13, 20]
z_cressie = lord_cressie.x[idx_test]
lower_test = lord_cressie.Lower1[idx_test]
upper_test = lord_cressie.Upper1[idx_test]

gcal = DiscretePriorClass(range(0.0,stop=1.0,length=300));

postmean_targets = Empirikos.PosteriorMean.(BinomialSample.(z_cressie,20));
chisq_floc = Empirikos.ChiSquaredFLocalization(0.05)

floc_method_chisq = FLocalizationInterval(flocalization = chisq_floc,
                                       convexclass= gcal, solver=Hypatia.Optimizer)

chisq_cis = confint.(floc_method_chisq, postmean_targets, Zs)

lower_chisq_ci = getproperty.(chisq_cis, :lower)
upper_chisq_ci = getproperty.(chisq_cis, :upper)

@test lower_chisq_ci ≈ lower_test atol=1e-3
@test upper_chisq_ci ≈ upper_test atol=0.005

Zs_collapse = begin
	Zs_collapse = deepcopy(Zs)
	n0 = pop!(Zs_collapse.store, BinomialSample(0, 20))
	n1 = pop!(Zs_collapse.store, BinomialSample(1, 20))
	updated_keys =  [BinomialSample(Interval(0,1), 20); collect(keys(Zs_collapse))]
	updated_values = [n0+n1; collect(values(Zs_collapse))]
	Empirikos.MultinomialSummary(updated_keys, updated_values)
end

discr = integer_discretizer(1:20)
Zs_collapse2 = discr(Zs)
discr(Zs_collapse) == Zs_collapse2

@test nobs(Zs_collapse) == nobs(Zs_collapse2)

floc_method_dkw = FLocalizationInterval(
							flocalization = DvoretzkyKieferWolfowitz(;α=0.05),
                            convexclass= gcal, solver=Hypatia.Optimizer)

fitted_dkw = fit(DvoretzkyKieferWolfowitz(;α=0.05), Zs_collapse)
fitted_dkw2 = fit(DvoretzkyKieferWolfowitz(;α=0.05), Zs_collapse2)

postmean_target = postmean_targets[2]
dkw_ci = confint(floc_method_dkw, postmean_target, Zs_collapse);
dkw_ci2 = confint(floc_method_dkw, postmean_target, Zs_collapse2);

@test dkw_ci == dkw_ci2

lam_chisq_withF = Empirikos.AMARI(
							convexclass=gcal,
                            flocalization = Empirikos.ChiSquaredFLocalization(0.01),
                            solver=Hypatia.Optimizer, discretizer=discr,
                            modulus_model = Empirikos.ModulusModelWithF
                           )

postmean_ci_lam_withF = confint(lam_chisq_withF, postmean_target, Zs_collapse2)

lam_chisq_withoutF = Empirikos.AMARI(
							convexclass=gcal,
                            flocalization = Empirikos.ChiSquaredFLocalization(0.01),
                            solver=Hypatia.Optimizer, discretizer=discr,
                            modulus_model = Empirikos.ModulusModelWithoutF
                           )

postmean_ci_lam_withoutF = confint(lam_chisq_withoutF, postmean_target, Zs_collapse2)

@test postmean_ci_lam_withoutF.lower ≈ postmean_ci_lam_withF.lower
@test postmean_ci_lam_withoutF.upper ≈ postmean_ci_lam_withF.upper


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.7.0-beta3
#     language: julia
#     name: julia-1.7
# ---

# %%
using Plots
x = range(-5, 5; length=300)
y = range(-3, 3; length=200)
f(x, y) = exp(-x^2/2-y^2)
z = f.(x', y)
surface(x, y, z; camera=(30, 60), color=:gist_earth)

# %%
surface(x, y, f; camera=(30, 60), color=:gist_earth)

# %%


module _ApEn
export ApEn
using Statistics: mean, std
    """
        Ap, Phi = ApEn(Sig)

    Returns the approximate entropy estimates `Ap` and the log-average number of
    matched vectors `Phi` for `m` = [0,1,2], estimated from the data sequence `Sig`
    using the default parameters: embedding dimension = 2, time delay = 1,
    radius distance threshold = 0.2*SD(`Sig`), logarithm = natural

        Ap, Phi = ApEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    Returns the approximate entropy estimates `Ap` of the data sequence `Sig`
    for dimensions = [0,1,...,`m`] using the specified keyword arguments:

    # Arguments:
    `m`      - Embedding Dimension, a positive integer\n
    `tau`    - Time Delay, a positive integer\n
    `r`      - Radius Distance Threshold, a positive scalar  \n
    `Logx`   - Logarithm base, a positive scalar\n

    # See also `XApEn`, `SampEn`, `MSEn`, `FuzzEn`, `PermEn`, `CondEn`, `DispEn`

    # References:
        [1] Steven M. Pincus,
            "Approximate entropy as a measure of system complexity."
            Proceedings of the National Academy of Sciences
            88.6 (1991): 2297-2301.

    """
    function ApEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    N = length(Sig)
    (N>10) ? nothing : error("Sig:   must be a numeric vector with >10 samples")
    (m > 0) ? nothing : error("m:     must be an integer > 0")
    (tau > 0) ? nothing : error("tau:   must be an integer > 0")
    (r>=0) ? nothing :  error("r:     must be a positive value")
    (Logx>0) ? nothing : error("Logx:     must be a positive number > 0")

    Counter = 1*(abs.(Sig .- transpose(Sig)) .<= r)
    M = Int.([m*ones(N-m*tau); repeat(collect(m-1:-1:1),inner=tau)])
    Ap = zeros(m+1)
    Phi = zeros(m+2)

    for n = 1:N-tau
        ix = findall(Counter[n, :] .== 1)

        for k = 1:M[n]
            ix = ix[(ix .+ (k*tau)) .<= N]
            p1 = repeat(transpose(Sig[n:tau: n+(tau*k)]), length(ix))
            p2 = Sig[ix .+ transpose(collect(0:tau:(k*tau)))]
            ix = ix[findall(maximum(abs.(p1 - p2),dims=2) .<= r)]
            Counter[n, ix] .+= 1
        end
    end

    Phi[1] = log(Logx, N)/N
    Phi[2] = mean(log.(Logx, sum(Counter.>0,dims=2)/N))
    Ap[1]  = Phi[1] - Phi[2]

    for k = 0:m-1
        ai = sum(Counter.>k+1,dims=2)/(N-(k+1)*tau)
        bi = sum(Counter.>k,dims=2)/(N-(k*tau))
        ai = ai[ai.>0]
        bi = bi[bi.>0]
        Phi[k+3] = sum(log.(Logx,ai))/(N-(k+1)*tau)
        Ap[k+2]  = sum(log.(Logx,bi))/(N-(k*tau)) - Phi[k+3]
    end

    return Ap, Phi
    end

end

"""Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub"""


using CPUTime
using DifferentialEquations
using Plots

f = @ode_def begin
    dx = v
    dv = 1/m*(-c*v - k*x + u)
    du = -Kp*v + Ki*(xg-x) - Kd/m*(-c*v - k*x + u)
end m c k xg Kp Ki Kd

function do_1()
    """1つめのやつ"""
    m = 1.0
    c = 1.0
    k = 1.0
    xg = 1.0
    Kp = 3
    Ki = 1
    Kd = 0.05

    p = [m, c, k, xg, Kp, Ki, Kd]

    x0 = [0.0, 0.0, Kp*(xg - 1.0)]  # 初期値を設定
    tspan = (0.0, 20.0)  # 解く範囲を設定

    prob = ODEProblem(f ,x0 ,tspan, p)
    sol = solve(prob)  # ソルバで解く

    # plot
    plot(
        sol,
        linewidth=2,
        xaxis="time [sec]",
        yaxis="",
        label=["position" "velocity" "input"]
    )

end

function do_2()
    """2つめのやつ"""
    m = 1.0
    c = 1.0
    k = 1.0
    xg = 1.0

    n = 1000  # アニメの枚数
    Kp_min, Kp_max = 9.99, 10.0
    Ki_min, Ki_max = 4.99, 5.0
    Kd_min, Kd_max = 0.0, 5.0
    Kp = [i for i in Kp_min:(Kp_max-Kp_min)/n:Kp_max]
    Ki = [i for i in Ki_min:(Ki_max-Ki_min)/n:Ki_max]
    Kd = [i for i in Kd_min:(Kd_max-Kd_min)/n:Kd_max]

    sols = []
    tspan = (0.0, 5.0)  # 解く範囲を設定

    for j in 1:n
        x0 = [0.0, 0.0, Kp[j]*(xg - 1.0)]  # 初期値を設定
        p = [m, c, k, xg, Kp[j], Ki[j], Kd[j]]
        prob = ODEProblem(f ,x0 ,tspan, p)
        sol = solve(prob)  # ソルバで解く
        push!(sols, sol)
    end

    ano_xposi = (tspan[1] + tspan[2])/2
    anim = @animate for i in 1:1:n #length(sols)
        Kp_now = round(Kp[i], digits=1)
        Ki_now = round(Ki[i], digits=1)
        Kd_now = round(Kd[i], digits=1)
        s = "Kp = " * string(Kp_now) * ", Ki = " * string(Ki_now) * ", Kd = " * string(Kd_now)
        plot(
            sols[i],
            ylims = (-1, xg*1.5), annotate = (ano_xposi, -1, s),
        )
    end
    mp4(anim, "example_4_jl.mp4", fps = 30)

end

#@time do_1()

@time do_2()


@testset "interpolations" begin

for interpolation in (Lagrange{1, RefCube, 1}(),
                      Lagrange{1, RefCube, 2}(),
                      Lagrange{2, RefCube, 1}(),
                      Lagrange{2, RefCube, 2}(),
                      Lagrange{2, RefTetrahedron, 1}(),
                      Lagrange{2, RefTetrahedron, 2}(),
                      Lagrange{3, RefCube, 1}(),
                      Serendipity{2, RefCube, 2}(),
                      Serendipity{3, RefCube, 2}(),
                      Lagrange{3, RefTetrahedron, 1}(),
                      Lagrange{3, RefTetrahedron, 2}(),)

    # Test of utility functions
    ndim = Ferrite.getdim(interpolation)
    r_shape = Ferrite.getrefshape(interpolation)
    func_order = Ferrite.getorder(interpolation)
    @test typeof(interpolation) <: Interpolation{ndim,r_shape,func_order}
    @test typeof(Ferrite.getlowerdim(interpolation)) <: Interpolation{ndim-1}
    @test typeof(Ferrite.getlowerorder(interpolation)) <: Interpolation{ndim,r_shape,func_order-1}

    n_basefuncs = getnbasefunctions(interpolation)
    x = rand(Tensor{1, ndim})
    f = (x) -> Ferrite.value(interpolation, Tensor{1, ndim}(x))
    @test vec(ForwardDiff.jacobian(f, Array(x))') ≈
           reinterpret(Float64, Ferrite.derivative(interpolation, x))
    @test sum(Ferrite.value(interpolation, x)) ≈ 1.0

    coords = Ferrite.reference_coordinates(interpolation)
    for node in 1:n_basefuncs
        N_node = Ferrite.value(interpolation, coords[node])
        for k in 1:node
            if k == node
                @test N_node[k] ≈ 1.0
            else
                @test N_node[k] ≈ 0.0
            end
        end
    end
end

end


module HMC
using Distributions

"""
Get random momentum from a parameter.
"""
rand_momentum(x) = randn(size(x))

"""
Compute kinetic energy given momentum
"""
compute_kinetic_energy(momentum) = sum(momentum .^ 2) / 2

"""
See p. 14 of
https://arxiv.org/pdf/1206.1901.pdf
"""
# TODO: Check this
function hmc_update(curr_state,
                    logprob::Function, grad_logprob::Function,
                    num_leapfrog_steps::Int, eps::Float64;
                    momentum_sd::Real=1)
  # Position variables
  qs = copy(curr_state)

  # Momentum variables
  ps = rand_momentum(qs)

  # Current kinetic energy
  curr_K = compute_kinetic_energy(ps)

  # Potential enerdy
  U(s) = -logprob(s)

  # Current potential energy
  curr_U = U(qs)

  # Compute gradient of U
  grad_U(qs) = -grad_logprob(qs)

  # Make a half step for momentum at the beginning
  gs = grad_U(qs)
  ps .-= eps * qs / 2

  # Alternate full steps for position and momentum
  for i in 1:num_leapfrog_steps
    # Make a full step for the position
    qs .+= eps * ps

    # Make a full step for the momentum, except at end of trajectory
    if i < num_leapfrog_steps
      gs = grad_U(qs)
      ps .-= eps * qs
    end
  end

  # Make a half step for momentum at the end
  gs = grad_U(qs)
  ps .-= eps * qs / 2

  # Proposed potential energy
  cand_U = U(qs)

  # Proposed kinetic  energy
  cand_K = compute_kinetic_energy(ps)

  ### Metropolis step ###
  # log acceptance ratio.  NOTE: potential = -log_prob
  log_acceptance_ratio = curr_U + curr_K - cand_U - cand_K

  # Accept or reject
  if log_acceptance_ratio > log(rand())
    return qs, -cand_U
  else
    return curr_state, -curr_U
  end
  ### End of Metropolis step ###
end

end


"""
    multroot(p; tol, scale=false)

  Finds all roots of a real or complex polynomial p,

        p(x) = a_1 x^n + a_2 x^n-1 + ... + a_n x + a_n+1,

  given by an (n+1)-row vector `p = (a_1, ..., a_n+1)`

  METHOD : for details, contact the author

             Zhonggang Zeng
             Department of Mathematics
             Northeastern Illinois University
             Chicago, IL 60625

             email: zzeng@neiu.edu

 This code is freely released for research exchange only. The author
 is not responsible for any demage caused by using this code.

  INPUT :  `p = polynomial coefficients.`

  OUTPUT : `z = distinct roots of polynomial p`

            l = corresponding multiplicities
  CALL :
               z = multroot(p).
"""
function multroot(p::AbstractVector{T}; tol::S = 1e-10, scaling::Bool=false) where {T<:Number,S<:Real}
    #
    # clear leading/trailing zeros
    #
    jj = findall(!iszero, p)
    isempty(jj) && throw(ArgumentError("zero polynomial has no roots"))
	Z = zero(T)
	E = one(T)
    n = length(p)
    j1, j2 = extrema(jj)
    q = p[j1:j2]
    if j1 >= j2
        return PolyZeros(T[], Int[]), Inf, Inf, 0
    end

    scale = scale!(q, scaling)
	z0, bke = try
        gcdroot(q, tol)
    catch ex
        @warn ex
        S[], Inf
    end

    if bke < S(0.01)

        z1, bkerr, pjcnd, job = pejroot(q, z0.z, z0.mult)

        # z1, bkerr, pcjnd, job = z0, bke, Z, 1
        z1.z .*= scale
        z = z1.z
        l = z1.mult
        if j2 < n
            push!(z1.z,  Z)
            push!(z1.mult, n-j2)
        end
        if job == 1
            #
            # show off results
            #
			nz = length(z)
            @printf("\n");
            @printf("    !!!THE COMPUTATION IS SUCCESSFUL!!!\n")
            @printf("\n")
            @printf("THE PEJORATIVE CONDITION NUMBER:       %g \n", pjcnd)
            @printf("THE BACKWARD ERROR:                    %6.2e %6.2e \n", bkerr, bke)
            @printf("THE ESTIMATED FORWARD ROOT ERROR:      %6.2e \n", 2 * bkerr * pjcnd)
            @printf("\n");
            if norm(imag(z)) == Z
                @printf("        computed roots         multiplicities\n")
                @printf("\n")
				for j = 1:nz @printf("%25.15f \t \t \t %3d \n", z[j], l[j]) end
            else
                @printf("                  computed roots ")
                @printf("                 multiplicities\n");
                @printf("\n");
                for j = 1:nz
				    @printf("%22.15f + %22.15f i \t %3d \n", real(z[j]), imag(z[j]), l[j])
				end
            end
        else
            z = roots(p)	# fallback to standard rootfinder
			l = ones(Int, n - 1)
            z1 = PolyZeros(z, l)
        end
    else
        z = roots(p)		# fallback to standard rootfinder
		l = ones(Int, n - 1)
        z1 = PolyZeros(z, l)
		job = 0
		bkerr = S(Inf)
		pjcnd = S(Inf)
    end
    z1, bkerr, pjcnd, job
end

function scale!(p::Vector, scaling::Bool)
    n = length(p) - 1
    a = p[1]
    c = p[1]
    p ./= a
    if scaling
        b = abs(p[n+1])
        p[n÷2+2:n+1] ./= b
        c = b ^ (1 / n)
        cc = c
        for k = 2:(n+1)÷2
            p[k] /= cc
            p[n+2-k] *= cc
            cc *= c
        end
        if iseven(n)
            p[n÷2+1] /= cc
        end
    end
    c
end


using Test
using CompScienceMeshes

# G0 is the interior of the metal (no simulation required here)
# G1 is the unbounded exterior
# G2 is the bounded interior

fn = joinpath(@__DIR__, "assets", "thick_cladding.msh")
G01 = CompScienceMeshes.read_gmsh_mesh(fn, physical="Gamma01")
G02 = CompScienceMeshes.read_gmsh_mesh(fn, physical="Gamma02")
G12 = CompScienceMeshes.read_gmsh_mesh(fn, physical="Gamma12")

G = CompScienceMeshes.read_gmsh_mesh(fn)

@test numcells(G01) + numcells(G02) + numcells(G12) == numcells(G)

# Build region boundaries with normals pointing inward:
G1 = weld(-G01, -G12)
@test CompScienceMeshes.isoriented(G1)

G2 = weld(-G02, G12)
@test CompScienceMeshes.isoriented(G2)

isclosed(m) = (numcells(boundary(m)) == 0)

@test isclosed(G1)
@test isclosed(G2)


struct PosteriorGP{Tprior, Tdata} <: AbstractGP
    prior::Tprior
    data::Tdata
end

"""
    posterior(fx::FiniteGP, y::AbstractVector{<:Real})

Constructs the posterior distribution over `fx.f` given observations `y` at `x` made under
noise `fx.Σy`. This is another `AbstractGP` object. See chapter 2 of [1] for a recap on
exact inference in GPs. This posterior process has mean function
```julia
m_posterior(x) = m(x) + k(x, fx.x) inv(cov(fx)) (y - mean(fx))
```
and kernel
```julia
k_posterior(x, z) = k(x, z) - k(x, fx.x) inv(cov(fx)) k(fx.x, z)
```
where `m` and `k` are the mean function and kernel of `fx.f` respectively.
"""
function posterior(fx::FiniteGP, y::AbstractVector{<:Real})
    m, C_mat = mean_and_cov(fx)
    C = cholesky(Symmetric(C_mat))
    δ = y - m
    α = C \ δ
    return PosteriorGP(fx.f, (α=α, C=C, x=fx.x, δ=δ))
end

"""
    posterior(fx::FiniteGP{<:PosteriorGP}, y::AbstractVector{<:Real})

Constructs the posterior distribution over `fx.f` when `f` is itself a `PosteriorGP` by
updating the cholesky factorisation of the covariance matrix and avoiding recomputing it
from original covariance matrix. It does this by using `update_chol` functionality.

Other aspects are similar to a regular posterior.
"""
function posterior(fx::FiniteGP{<:PosteriorGP}, y::AbstractVector{<:Real})
    m2 = mean(fx.f.prior, fx.x)
    δ2 = y - m2
    C12 = cov(fx.f.prior, fx.f.data.x, fx.x)
    C22 = cov(fx.f.prior, fx.x) + fx.Σy
    chol = update_chol(fx.f.data.C, C12, C22)
    δ = vcat(fx.f.data.δ, δ2)
    α = chol \ δ
    x = vcat(fx.f.data.x, fx.x)
    return PosteriorGP(fx.f.prior , (α=α, C=chol, x=x, δ=δ))
end

# AbstractGP interface implementation.

function Statistics.mean(f::PosteriorGP, x::AbstractVector)
    return mean(f.prior, x) + cov(f.prior, x, f.data.x) * f.data.α
end

function Statistics.cov(f::PosteriorGP, x::AbstractVector)
    return cov(f.prior, x) - Xt_invA_X(f.data.C, cov(f.prior, f.data.x, x))
end

function cov_diag(f::PosteriorGP, x::AbstractVector)
    return cov_diag(f.prior, x) - diag_Xt_invA_X(f.data.C, cov(f.prior, f.data.x, x))
end

function Statistics.cov(f::PosteriorGP, x::AbstractVector, z::AbstractVector)
    C_xcond_x = cov(f.prior, f.data.x, x)
    C_xcond_y = cov(f.prior, f.data.x, z)
    return cov(f.prior, x, z) - Xt_invA_Y(C_xcond_x, f.data.C, C_xcond_y)
end

function mean_and_cov(f::PosteriorGP, x::AbstractVector)
    C_xcond_x = cov(f.prior, f.data.x, x)
    m_post = mean(f.prior, x) + C_xcond_x' * f.data.α
    C_post = cov(f.prior, x) - Xt_invA_X(f.data.C, C_xcond_x)
    return (m_post, C_post)
end

function mean_and_cov_diag(f::PosteriorGP, x::AbstractVector)
    C_xcond_x = cov(f.prior, f.data.x, x)
    m_post = mean(f.prior, x) + C_xcond_x' * f.data.α
    C_post_diag = cov_diag(f.prior, x) - diag_Xt_invA_X(f.data.C, C_xcond_x)
    return (m_post, C_post_diag)
end


function parse_line_segments(path::String)::Array{Int}
    line_segments = Array{Int}(undef, 1, 4)
    open(path) do file
        lines = readlines(file)
        for line in lines
            coordinates = split(line, " -> ")
            segment = [0 0 0 0]
            i = 1
            for coordinate in coordinates
                x = parse(Int, split(coordinate, ",")[1])
                y = parse(Int, split(coordinate, ",")[2])
                segment[i] = x
                segment[i+1] = y
                i += 2
            end
            line_segments = vcat(line_segments, segment)
        end
    end
    return line_segments[2:end, :]
end

function draw_map(line_segments)
    min, max = minimum(line_segments), maximum(line_segments)
    field = zeros(Int, max-min+1, max-min+1)
    line_segments = line_segments .+ (-min + 1)
    for i in 1:size(line_segments)[1]
        segment = line_segments[i, :]
        draw_line(field, segment)
    end
    return field
end

function draw_line(field, line_segment)
    x1, y1 = line_segment[1], line_segment[2]
    x2, y2 = line_segment[3], line_segment[4]
    x_step = x1 < x2 ? 1 : -1
    y_step = y1 < y2 ? 1 : -1
    if x1 == x2
        field[y1:y_step:y2, x1] = field[y1:y_step:y2, x1] .+ 1
    elseif y1 == y2
        field[y1, x1:x_step:x2] = field[y1, x1:x_step:x2] .+ 1
    elseif abs(x1-x2) == abs(y1-y2)
        for (x, y) in zip(x1:x_step:x2, y1:y_step:y2)
            field[y, x] += 1
        end
    end
end

function find_dangerous_areas(map)
    overlaps = map .> 1
    return sum(overlaps)
end

function main()
    path = "./inputs/puzzle.txt"
    segments = parse_line_segments(path)
    map = draw_map(segments)
    display(map)
    overlaps = find_dangerous_areas(map)
    println(overlaps)
end

main()


## Angle type definition and elementary functions
type Angle
  #type for angles. of form o + p * pi where p is rational
  o::Rational{Int64}
  p::Rational{Int64}
end

Base.show(io::IO, a::Angle) = print(io, a.o, "+", a.p, "π")

function Base.:+(a::Angle, b::Angle)
  return Angle(a.o + b.o, a.p + b.p)
end

function Base.:-(a::Angle)
  return Angle(-a.o, -a.p)
end

function Base.:-(a::Angle, b::Angle)
  return a + -b
end

function Base.:*(a::Angle, b::Angle)
  if a.p == 0//1 || b.p == 0//1 #only one nonzero pi coefficient => closed under multiplication
    return Angle(a.o * b.o, (a.o * b.p + b.o * a.p)) #pi^2 squared coefficient (a.p*b.p) is always zero
  end
  error("Cannot multiply. Must have at least one zero p term for type stability")
end

function Base.:*(a::Angle, b::Rational)
  return Angle(a.o * b, a.p * b)
end

function Base.:*(a::Rational, b::Angle)
  return b*a
end

function Base.:*(a::Array{Angle,1}, b::Angle)
  return b.*a
end

function Base.:*(a::Angle, b::Array{Angle,1})
  return b*a
end

function Base.:/(a::Angle, b::Angle)
  if b.p == 0//1 #invertable
    return Angle(a.o / b.o, a.p / b.o)
  end
  if a.o == 0//1 && b.o == 0//1 #pi terms will cancel
    return Angle(a.p / b.p, 0//1)
  end
  warn("Answer not expressible as angle data typel.")
  return (a.o + a.p * pi)/(b.o + b.p * pi)
end

function Base.:/(a::Array{Angle,1}, b::Angle)
  return a ./ b
end

function Base.:/(a::Int, b::Angle)
  return Angle(a//1, 0//1)/b
end

function Base.inv(a::Angle)
  return 1/a
end

function Base.zero(::Type{Angle})
  return Angle(zero(Rational), zero(Rational))
end

function Base.zero(a::Angle)
  return zero(Angle)
end

function Base.one(::Type{Angle})
  return Angle(one(Rational), zero(Rational))
end

import Base.== #the other syntax was throwing errors
==(a::Angle, b::Angle) = a.o == b.o && a.p == b.p
==(a::Angle, n:: Int) = a.o == n && a.p == 0

function Base.transpose(a::Angle)
  return a
end

function Base.norm(a::Angle)
  return norm(a.o + a.p * pi)
end

function Base.convert(::Type{Angle}, a::Int)
  return Angle(a//1, 0//1)
end

## Side type definition and elementary functions
type Side
  #type for sides. of form q + r *sqrt(m)
  q::Rational{Int64}
  r::Rational{Int64}
  m::Int

  function Side(q, r, m)
    if !isnull(m) && issquare(m) != 0 #m is square
      new(q + r * issquare(m), 0//1, m) #simplify to rational
    else
      new(q, r, m)
    end
  end
end

Base.show(io::IO, a::Side) = print(io, a.q, "+", a.r, "√", a.m)

function issquare(m)
  #returns if m is a perfect square. m will always be small and > 2
  #returns 0 if not square and the square root if square
  i = 2 #m is always greater than 2
  while i^2 <= m
    if i^2 == m
      return i
    end
    i += 1
  end
  return 0
end

function mfix(a::Side, b::Side)
  #checks if the value for m is null
  if isnull(a.m)
    if isnull(b.m)
      error("No non-null values for m.")
    else
      a.m = b.m
    end
  else
    if isnull(b.m)
      b.m = a.m
    end
  end
    return a, b
end

function Base.:+(a::Side, b::Side)
  #could add error check if m is the same
  a, b = mfix(a, b)
  return Side(a.q + b.q, a.r + b.r, a.m)
end

function Base.:-(a::Side)
  return Side(-a.q, -a.r, a.m)
end

function Base.:-(a::Side, b::Side)
  return a + -b
end

function Base.:*(a::Side, b::Rational)
  return Side(a.q * b, a.r * b, a.m)
end

function Base.:*(a::Rational, b::Side)
  return b*a
end

function Base.:*(a::Side, b::Side)
  a, b = mfix(a, b)
  return Side(a.q * b.q + a.r * b.r * a.m, a.r * b.q + a.q * b.r, a.m)
end

function Base.:*(a::Array{Side,2}, b::Side)
  return a .* b
end

function Base.:/(a::Side, b::Side)
  a, b = mfix(a, b)
  return Side((a.q*b.q-a.r*b.r*a.m)/(b.q^2-a.m*b.r^2),
   -(a.q*b.r-a.r*b.q)/(b.q^2-a.m*b.r^2), a.m) #never divides by zero because m cannot be square
end

function Base.:/(a::Int, b::Side)
  return Side(a//1, 0, b.m) / b
end

function Base.:/(a::Array{Side,1}, b::Side)
  return a ./ b
end

function Base.isless(a::Side, b::Side)
  a, b = mfix(a, b)
  return (a.q + a.r * sqrt(a.m)) < (b.q + b.r * sqrt(b.m))
end

function Base.inv(a::Side)
  return 1/a
end

import Base.== #the other syntax was throwing errors
==(a::Side, b::Side) = a.q == b.q &&  a.r == b.r #does not check for m equality
==(a::Side, b::Int) = a.q == b && a.r == 0//1

function Base.convert(::Type{Side}, a::Int, m)
  return Side(a//1, 0//1, m)
end

function Base.convert(::Float64, a::Side)
  return (a.p + a.r * a.m)::Float64
end

function Base.convert(::Type{Array{Side,2}}, a::Array{Int64,2}, m::Int64)
  output = zeros(Side, size(a), mvalue=m)
  for i in 1:length(a)
    output[i] = convert(Side, a[i], m)
  end
  return output
end

function Base.zero(::Type{Side})
  return Side(zero(Rational), zero(Rational), Nullable{Int}())
end

function Base.zero(::Type{Side}, m::Int)
  return Side(zero(Rational), zero(Rational), m)
end

function Base.zero(a::Side)
  return zero(Side)
end

function Base.zeros(::Type{Side}, size; mvalue::Int=Nullable{Int64})
  return repmat([zero(Side, mvalue)], size...)
end

function Base.zeros(a::Array{Side})
  return zeros(Side, size(a), mvalue=findm(a))
  error("All m-values in matrix null.")
end

function Base.one(::Type{Side})
  return Side(one(Rational), zero(Rational), Nullable{Int}())
end

function Base.one(::Type{Side}, m::Int)
  return Side(one(Rational), zero(Rational), m)
end

function Base.transpose(a::Side)
  return a
end

function Base.norm(a::Side)
  return a.q + a.r * sqrt(a.m)
end

function findm(a::Array{Side})
 #finds the first non-null value for m in the matrix
  for i in a
    if !isnull(i.m)
      return i.m
    end
  end
end

function sideeye(::Type{Side}, n, m)
  #returns nxn identity matrix of size with m
  output = zeros(Side, (n,n), mvalue=m)
  for i in 1:n
    output[i,i] = one(Side, m)
  end
  return output
end

## rref
function rref(mat)
  #returns modified reduced row eschelon form of mat
  #probably could be improved further
  j, k = size(mat)
  for i in 1:min(j,k) #loop through rows
    if mat[i,:] != zeros(mat[i,:]) #non zero row.
      pivotind = findfirst(mat[i,:])
      mat[i,:] /= mat[i,pivotind] #normalize leading element
      minus = repmat(mat[i,:]', j, 1) #base to zero the other rows
      minus .*= mat[:,pivotind] #scale to zero elements below pivot
      minus[i,:] = zeros(minus[i,:]) #zero pivot row to avoid cancelling it
      mat -= minus
    end
  end
  return rrefsort(mat)

end

function rrefsort(mat)
  #sorts rows to their aligned pivot location
  output = zeros(mat) #storage for output
  j, k = size(mat)
  for i in 1:min(j, k) #loop through each row
   if mat[i,:] != zeros(mat[i,:]) #non zero row.
    output[findfirst(mat[i,:]),:] = mat[i,:]
   end
  end
  return output
end


"""
Oja (non-Distributed) streaming PCA method

Author: Wayne Wang
Email: wayneyw@umich.edu
Last modified: 08/26/2021
"""

struct ADAGRAD end
struct ROBBINS_MONRO end

function oja(Xt::AbstractArray{<:Real}, V::AbstractArray{<:Real},
    α::AbstractArray{<:Real}; scheduling::AbstractString = "adagrad",
    stepsize_dim::Int = 1)
    """
    Oja's method
    """
    # initial setting
    if scheduling == "adagrad"
        scheduling = ADAGRAD()
    elseif scheduling == "robbins-monro"
        scheduling = ROBBINS_MONRO()
    end

    # update eigenvectors & learning rates
    V, α = ojaupdate(scheduling, Xt, V, α, stepsize_dim)
    # retraction
    V .= Array(qr!(V).Q)

    return V, α
end

function ojaupdate(scheduling::ROBBINS_MONRO, Xt::AbstractArray{<:Real}, V::AbstractArray{<:Real},
    α::AbstractArray{<:Real}, stepsize_dim::Int)
    """
    Oja × Robbins-Monro
    """
    V .+= α*Xt*(copy(Xt')*V)
    α = nothing
    return V, α
end

function ojaupdate(scheduling::ADAGRAD, Xt::AbstractArray{<:Real}, V::AbstractArray{<:Real},
    α::AbstractArray{<:Real}, stepsize_dim::Int)
    """
    Oja × Adagrad
    """
    grad = Xt*(copy(Xt')*V)
    if stepsize_dim == 1 # α ∈ k × 1
        α .= α .+ sum(abs2,grad,dims=1)[:]
        V .+= grad ./ reshape(sqrt.(α),(1,size(α,1)))
    elseif stepsize_dim == 2 # α ∈ d × k
        α .= α .+ grad .* grad
        V .+= grad ./ sqrt.(α)
    end
    return V, α
end


#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*--------------------------------------------------------------------
#* Author: Tiago Os��io
#* $Id$
#*-------------------------------------------------------------------
type FirstOrder
	FirstOrder()=begin
		new(
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"Time Constant",
				:Unit=>"s",
				:Default=>4
			)),
			DanaReal(Dict{Symbol,Any}(
				:Unit=>"1/s"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Unit=>"1/s"
			)),
			DanaReal(),
			DanaReal(Dict{Symbol,Any}(
				:Default=>0
			)),
			control_signal(Dict{Symbol,Any}(
				:Brief=>"State"
			)),
			control_signal(Dict{Symbol,Any}(
				:Brief=>"Input signal",
				:PosX=>0,
				:PosY=>0.5
			)),
			control_signal(Dict{Symbol,Any}(
				:Brief=>"Output signal",
				:PosX=>1,
				:PosY=>0.5
			)),
			[
				:(diff(x) = A*x + B*u),
				:(y = C*x + D*u),
			],
			[
				"","",
			],
			[:tau,:A,:B,:C,:D,],
			[:x,:u,:y,]
		)
	end
	tau::DanaReal
	A::DanaReal
	B::DanaReal
	C::DanaReal
	D::DanaReal
	x::control_signal
	u::control_signal
	y::control_signal
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export FirstOrder
function setEquationFlow(in::FirstOrder)
	addEquation(1)
	addEquation(2)
end
function atributes(in::FirstOrder,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Icon]="icon/PIDIncr"
	drive!(fields,_)
	return fields
end
FirstOrder(_::Dict{Symbol,Any})=begin
	newModel=FirstOrder()
	newModel.attributes=atributes(newModel,_)
	newModel
end


# Factorization Type definitions

# Reference: I.S. Dhillon, B.N. Parlett. Orthogonal Eigenvectors and Relative Gaps.
# Siam J. Matrix Anal. Appl. Vol. 25, No. 3, pp. 858-899 (2003)

module mr3

export LDL, NDN, UDU, getvec

# LDL factorization type for tridiagonal matrices
type LDL{T}
    d::Vector{T} # diagonal entries of D
    l::Vector{T} # lower diagonal entries of L
    function LDL{T}(d::Vector{T}, l::Vector{T})
        length(d)==length(l)+1 ? new(d,l) : throw(DimensionMismatch())
    end
end

# UDU factorization type for tridiagonal matrices
type UDU{T}
    d::Vector{T} # diagonal entries of D
    u::Vector{T} # upper diagonal entries of U
    function UDU{T}(d::Vector{T}, u::Vector{T})
        length(d)==length(u)+1 ? new(d,u) : throw(DimensionMismatch())
    end
end

# Storage of all twisted factorizations
type NDN{T}
    dl::Vector{T} # diagonal entries of D+
    l::Vector{T} # sub-diagonal entries of L+
    du::Vector{T} # diagonal entries of D-
    u::Vector{T} # super-diagonal entries of U-
    γ::Vector{T} # one choice for each twisted factorization
    function NDN{T}(dl::Vector{T}, l::Vector{T}, du::Vector{T}, u::Vector{T}, γ::Vector{T})
        b = length(dl)==length(du) && length(l)==length(u) && length(dl)==(length(l)+1) && length(dl)==length(γ)
        b ? new(dl, l, du, u, γ) : throw(DimensionMismatch())
    end
end

# LDL constructors
LDL{T}(d::AbstractVector{T}, l::AbstractVector{T})=LDL{T}(copy(d), copy(l))
LDL(n::Integer)=LDL(zeros(n), zeros(n-1)) # blank constructor

# UDU constructors
UDU{T}(d::AbstractVector{T}, u::AbstractVector{T})=UDU{T}(copy(d), copy(u))
UDU(n::Integer)=UDU(zeros(n), zeros(n-1)) # blank constructor

#NDN constructors
function NDN{T}(dl::AbstractVector{T}, l::AbstractVector{T}, du::AbstractVector{T}, u::AbstractVector{T}, γ::AbstractVector{T})
    return NDN{T}(copy(dl), copy(l), copy(du), copy(u), copy(γ))
end
NDN(n::Integer)=NDN(zeros(n), zeros(n-1), zeros(n), zeros(n-1), zeros(n)) # blank constructor

# get LDL factorization of a SymTridiagonal factorization
function LDL{T}(M::SymTridiagonal{T})
    n = length(M.dv)
    Md = copy(M.dv)
    Me = copy(M.ev)
    # vectors for LDL factorization
    d = Vector{T}(n)
    l = Vector{T}(n-1)
    for i=1:n-1
        d[i] = Md[i];
        l[i] = Me[i]/d[i];
        Md[i+1] = Md[i+1] - l[i]*d[i]*l[i]
    end
    d[n] = Md[n]
    return LDL(d,l)
end

function fullLDL(F::LDL)
    n = length(F.d)
    L = diagm(F.l,-1)
    L = L + I
    D = diagm(F.d)
    return L, D
end

function dstqds(F::LDL, μ::Real)
    # Transform F = LDL to G = LDL-μI
    # where G is also a LDL factorization
    # ref Dhillon Parlett 2003 pp.869
    n = length(F.d)
    G = LDL(n)
    s = -μ
    for i = 1:n-1
        G.d[i] = s + F.d[i]
        G.l[i] = F.l[i]*F.d[i]/G.d[i]
        s = G.l[i]*F.l[i]*s - μ
    end
    G.d[n] = s + F.d[n]
    return G
end

function dqds(F::LDL, μ::Real)
    # transform F = LDL to G = LDL - μI
    # where G is a UDU factorization
    # ref Dhillon Parlett 2003 pp.870
    n = length(F.d)
    G = UDU(n)
    p = zeros(n)
    p[n] = F.d[n] - μ
    for i = n-1:-1:1
        G.d[i+1] = F.d[i]*F.l[i]^2 + p[i+1]
        t = F.d[i]/G.d[i+1]
        G.u[i] = F.l[i]*t
        p[i] = p[i+1]*t - μ
    end
    G.d[1] = p[1]
    return G
end

function fullUDU(F::UDU)
    n = length(F.d)
    U = diagm(F.u,1)
    U = U + I
    D = diagm(F.d)
    return D, U
end

function dtwqds(F::LDL, μ::Real)
    # transform F = LDL to G = LDL - μI
    # where G is a NDN factorization
    # ref Dhillon Parlett 2003 pp.871
    n = length(F.d)
    G = NDN(zeros(n), zeros(n-1), zeros(n), zeros(n-1), zeros(n))
    s = Vector{Real}(n)
    p = Vector{Real}(n)

    # L+D+L+ factorization
    s[1] = -μ
    for i = 1:n-1
        G.dl[i] = s[i] + F.d[i]
        G.l[i] = F.l[i]*F.d[i]/G.dl[i]
        s[i+1] = G.l[i]*F.l[i]*s[i] - μ
    end
    G.dl[n] = s[n] + F.d[n]

    # U-D-U- factorization
    p[n] = F.d[n] - μ
    for i = n-1:-1:1
        G.du[i+1] = F.d[i]*F.l[i]^2 + p[i+1]
        t = F.d[i]/G.du[i+1]
        G.u[i] = F.l[i]*t
        p[i] = p[i+1]*t - μ
    end
    G.du[1] = p[1]

    # compute γ
    for r  = 1:n-1
        G.γ[r] = s[r] + F.d[r]*p[r+1]/G.du[r+1]
    end
    G.γ[n] = s[n] + F.d[n]

    return G
end

function twisted_eigenvector(F::NDN, r::Integer, G::LDL)
    # transform find smallest eigenvector of F = NDN
    # r tells which NDN factorization to use
    # INPUTS:
    #   F - NDN factorizations of LDL-μI.  Generate using dtwqds
    #   r - which twisted factorization should be used
    #   G - unshifted LDL factorization.  Used if shifted factorization has 0 on diagonal
    # ref Dhillon Parlett 2003 pp.866

    n = length(F.dl)
    z = zeros(n)
    z[r] = 1
    for i = r-1:-1:1
        if F.dl[i]!=0
            z[i] = -F.l[i]*z[i+1]
        else
            z[i] = -G.d[i+1] * G.l[i+1] / (G.d[i] * G.l[i]) * z[i+2]
        end
    end
    for j = r:n-1
        if F.du[j+1]!=0
            z[j+1] = -F.u[j]*z[j]
        else
            z[j] = -G.d[j-1] * G.l[j-1] / (G.d[j] * G.l[j]) * z[j-1]
        end
    end

    # normalize
    z = z / norm(z)
    return z
end

function argmin(v::AbstractVector)
    # return index with smallest value

    curmin = Inf # current minimum
    amin = 0 # current argmin
    for i = 1:length(v)
        if v[i] <= curmin
            amin = i
            curmin = v[i]
        end
    end
    return amin
end

function getvec(F::LDL, λ::Real)
    # get eigenvector with approximate eigenvalue λ
    # use twisted factorizations
    # ref Dhillon Parlett 2003 pp.865-866

    G = dtwqds(F, λ)
    r = argmin(abs(G.γ))
    return twisted_eigenvector(G, r, F)
end

end # module


export measure, dirac
export variables, monomials, moments

# If a monomial is not in x, it does not mean that the moment is zero, it means that it is unknown/undefined
struct Measure{T, MT <: AbstractMonomial, MVT <: AbstractVector{MT}} <: AbstractMeasure{T}
    a::Vector{T}
    x::MVT

    function Measure{T, MT, MVT}(a::Vector{T}, x::MVT) where {T, MT, MVT}
        @assert length(a) == length(x)
        new(a, x)
    end
end
Measure(a::AbstractVector{T}, x::AbstractVector{TT}) where {T, TT <: AbstractTermLike} = Measure{T, monomialtype(TT), monovectype(x)}(monovec(a, x)...)

"""
    measure(a, X::AbstractVector{<:AbstractMonomial})

Creates a measure with moments `moment(a[i], X[i])` for each `i`.
"""
measure(a, X) = Measure(a, X)

"""
    variables(μ::AbstractMeasureLike)

Returns the variables of `μ` in decreasing order. Just like in MultivariatePolynomials, it could contain variables of zero degree in every monomial.
"""
MP.variables(μ::Measure) = variables(μ.x)

"""
    monomials(μ::AbstractMeasureLike)

Returns an iterator over the monomials of `μ` sorted in the decreasing order.
"""
MP.monomials(μ::Measure) = μ.x

"""
    moments(μ::AbstractMeasureLike)

Returns an iterator over the moments of `μ` sorted in decreasing order of monomial.
"""
moments(μ::Measure) = map((α, x) -> moment(α, x), μ.a, μ.x)

Base.:(*)(α, μ::Measure) = measure(α * μ.a, μ.x)
Base.:(*)(μ::Measure, α) = measure(μ.a * α, μ.x)
Base.:(-)(μ::Measure) = measure(-μ.a, μ.x)
function Base.:(+)(μ::Measure, ν::Measure)
    @assert μ.x == ν.x
    measure(μ.a + ν.a, μ.x)
end

"""
    dirac(X::AbstractVector{<:AbstractMoment}, s::AbstractSubstitution...)

Creates the dirac measure by evaluating the moments of `X` using `s`.

## Examples

Calling `dirac([x*y, x*y^2], x=>3, y=>2)` should the measure with moment `x*y` of value `6` and moment `x*y^2` of value `12`.
"""
function dirac(x::AbstractVector{MT}, s::MP.AbstractSubstitution...) where {MT <: AbstractMonomial}
    Measure([m(s...) for m in x], x)
end