using Flux, Zygote
abstract type CurvatureInterface end

"""
    jacobians(curvature::CurvatureInterface, X::AbstractArray)

Computes the Jacobian `∇f(x;θ)` where `f: ℝᴰ ↦ ℝᴷ`.
"""
function jacobians(curvature::CurvatureInterface, X::AbstractArray)
    nn = curvature.model
    # Output:
    ŷ = nn(X)
    # Jacobian:
    𝐉 = jacobian(() -> nn(X),Flux.params(nn))
    𝐉 = reduce(hcat,[𝐉[θ] for θ ∈ curvature.params])
    return 𝐉, ŷ
end

"""
    gradients(curvature::CurvatureInterface, X::AbstractArray, y::Number)

Compute the gradients with respect to the loss function: `∇ℓ(f(x;θ),y)` where `f: ℝᴰ ↦ ℝᴷ`.
"""
function gradients(curvature::CurvatureInterface, X::AbstractArray, y::Number)
    model = curvature.model
    𝐠 = gradient(() -> curvature.loss(X,y),Flux.params(model))
    return 𝐠
end
struct EmpiricalFisher <: CurvatureInterface
    model::Any
    loss::Function
    params::AbstractArray
end

"""
    full(curvature::EmpiricalFisher, d::Union{Tuple,NamedTuple})

Compute the full empirical Fisher.
"""
function full(curvature::EmpiricalFisher, d::Tuple)
    x, y = d
    𝐠 = gradients(curvature, x, y)
    𝐠 = reduce(vcat,[vec(𝐠[θ]) for θ ∈ curvature.params])

    # Empirical Fisher:
    H = 𝐠 * 𝐠'
    return H

end


Random.seed!(17)
N = 1000
net = CompleteBipartiteGraph(N, N, Network)
eprop!(net, "w", e -> rand())
res = Matching.run_bp(net, γ=0.1, maxiters=1000, verbose=false)
@test abs(res.energy - π^2/6) < 0.1
@test res.num_violations == 0 # no violated constraints

N = 1000
net = CompleteGraph(N, Network)
eprop!(net, "w", e -> rand())
res = Matching.run_bp(net, γ=0.1, maxiters=1000, verbose=false)
@test abs(res.energy - π^2/12) < 0.1
@test res.num_violations <= 5 # if graph is non-bipartite may fail


# Turing.jl version of model at https://github.com/stan-dev/blob/master/basic_estimators/normal_mixture.stan

using StatsFuns
using ForwardDiff: Dual

@model nmmodel(y) = begin
  N = length(y)

  theta ~ Uniform(0, 1)

  mu = tzeros(2) # mu is sampled by HMC
  for i = 1:2
    mu[i] ~ Normal(0, 10)
  end

  k = tzeros(Int64, N)
  for i = 1:N
    k[i] ~ Categorical([theta, 1.0 - theta])
  end

  for i = 1:N
    y[i] ~ Normal(mu[k[i]], 1.0)
  end

  # logtheta_p = map(yᵢ -> [log(theta) + logpdf(Normal(mu[1], 1.0), yᵢ), log(1 - theta) + logpdf(Normal(mu[2], 1.0), yᵢ)], y)
  # map!(logtheta_pᵢ -> logtheta_pᵢ - logsumexp(logtheta_pᵢ), logtheta_p)   # normalization
  # for i = 1:N
  #   k[i] ~ Categorical(exp.(logtheta_p[i]))
  # end
end


using MAT
using Dates
using GLPK
using JuMP
using Statistics
using CartesianGeneticProgramming
using Cambrian
import Cambrian.mutate
using Plots
using Plots.PlotMeasures
using StatsPlots
include("../utils.jl")

function get_wave_t_in_range(param, min, max)
    indices_more = findall(>=(min), param)
    indices_less = findall(<(max), param)

    indices_to_keep = intersect(indices_more, indices_less)
    param = param[indices_to_keep]

    param, indices_to_keep
end

function fit_poly_inf_norm(x,y,n)
	m = Model(GLPK.Optimizer)

    @variable(m, alpha[1:n+1])
	@variable(m, t >= 0.)
	@constraint(m, [i=1:length(y)], sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i] <= t);
	@constraint(m, [i=1:length(y)], -t <= sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i]);
	@objective(m, Min, t)

	optimize!(m) # added new 2020-04-08

    return JuMP.value.(alpha)
end

f2(x, args) = (args[2] * x) .+ args[1]

const MATLAB_EPOCH = Dates.DateTime(-1,12,31)

date2num(d::Dates.DateTime) = Dates.value(d-MATLAB_EPOCH)/(1000*60*60*24)
num2date(n::Number) =  MATLAB_EPOCH + Dates.Millisecond(round(Int64, n*1000*60*60*24))

function plot_bathy(xs, profile; show_points=false, x_last_positive=NaN, x_first_negative=NaN, value_last_positive=NaN, value_first_negative=NaN, x_at_zero=NaN, title=NaN)
    p1 = plot(xs, profile, label="profile_t0")
    if show_points
        scatter!([x_last_positive, x_first_negative], [value_last_positive, value_first_negative], label="A, B")
        scatter!([x_at_zero], [0], label="Shoreline")
    end
    display(plot(p1, title=title))
end

function find_x_at_zero(profile, xs, i)
    t0_index_last_positive = findlast(profile.>0)
    if t0_index_last_positive == nothing t0_index_last_positive = 1 end

    t0_index_first_negative = findfirst(profile.<0)
    if t0_index_first_negative == nothing t0_index_first_negative = 1 end

    t0_x_last_positive = xs[t0_index_last_positive]
    t0_x_first_negative = xs[t0_index_first_negative]
    t0_value_last_positive = profile[t0_index_last_positive]
    t0_value_first_negative = profile[t0_index_first_negative]

    t0_coeffs = fit_poly_inf_norm([t0_value_last_positive, t0_value_first_negative], [t0_x_last_positive, t0_x_first_negative], 1)
    t0_x_at_zero = f2([0], t0_coeffs)[1]
    t0_x_at_zero
end

struct Day
    bathy0
    bathy1

    tp
    hs
    direction

    shoreline_change
    shoreline_change_class

    input_row
end

dir_path = get_dataset_path(;dataset_name="grand_popo")
kalman_bathy_vars = matread(string(dir_path, "Kalman_bathy.mat"))
bathy_xb = kalman_bathy_vars["xb"][1, :]# 1×213 Matrix{Float64}
# bathy_tn = kalman_bathy_vars["tn"]      # 1×1214 Matrix{Float64}
bathy_hbK = kalman_bathy_vars["hbK"]    # 213×1214 Matrix{Float64}

min_target_matdate = minimum(filter(!isnan, kalman_bathy_vars["tn"]))
max_target_matdate = maximum(filter(!isnan, kalman_bathy_vars["tn"]))

wave_vars = matread(string(dir_path, "Wave_GPP_ERA.mat"))
wave_t = wave_vars["t"]
wave_Tm = wave_vars["Tm"]
wave_hs = wave_vars["hs"]
wave_dir = wave_vars["dir"]

wave_t, indices = get_wave_t_in_range(wave_t, min_target_matdate, max_target_matdate)
wave_Tm = wave_Tm[indices]
wave_hs = wave_hs[indices]
wave_dir = wave_dir[indices]

bathy_dates = map(num2date, collect(min_target_matdate:max_target_matdate))  # collect() instead of using bathy_tn because bathy_tn includes NaN's, even though bathy_hbK has no NaN
wave_dates = map(num2date, wave_t)

days = []
n_offshore = 0
n_landward = 0
n_none = 0

all_tm = []
all_hs = []
all_dir = []
for i in 2:length(bathy_dates)
    _d0 = bathy_dates[i-1]
    _d1 = bathy_dates[i]

    _t0 = date2num(_d0)
    _t1 = date2num(_d1)

    _bathy_t0 = bathy_hbK[:, i-1]
    _bathy_t1 = bathy_hbK[:, i]

    _wave_t = deepcopy(wave_t)
    _wave_Tm = deepcopy(wave_Tm)
    _wave_dir = deepcopy(wave_dir)
    _wave_hs = deepcopy(wave_hs)

    _wave_t, _indices = get_wave_t_in_range(_wave_t, _t0, _t1)
    _wave_Tm = _wave_Tm[_indices]
    _wave_hs = _wave_hs[_indices]
    _wave_dir = _wave_dir[_indices]

    t0_x_at_zero = find_x_at_zero(_bathy_t0, bathy_xb, i)
    t1_x_at_zero = find_x_at_zero(_bathy_t1, bathy_xb, i)
    change = (t1_x_at_zero - t0_x_at_zero) / 100

    if change > 1 || change < -1
        println(i, "    ", change)
        exit()
    end

    if change > 0
        class = (1, 0, 0)  # (offshore, landward, none)
        global n_offshore
        n_offshore += 1
    elseif change < 0
        class = (0, 1, 0)
        global n_landward
        n_landward += 1
    else
        class = (0, 0, 1)
        global n_none
        n_none += 1
    end

    _wave_Tm = _wave_Tm ./ 100
    _wave_hs = _wave_hs ./ 10
    _wave_dir = _wave_dir ./ 100

    input_row = vcat(
        mean(_wave_Tm),
        minimum(_wave_Tm),
        maximum(_wave_Tm),
        std(_wave_Tm),
        mean(_wave_hs),
        minimum(_wave_hs),
        maximum(_wave_hs),
        std(_wave_hs),
        mean(_wave_dir),
        minimum(_wave_dir),
        maximum(_wave_dir),
        std(_wave_dir)
    )

    push!(all_tm, vcat(mean(_wave_Tm), minimum(_wave_Tm), maximum(_wave_Tm), std(_wave_Tm))...)
    push!(all_hs, vcat(mean(_wave_hs), minimum(_wave_hs), maximum(_wave_hs), std(_wave_hs))...)
    push!(all_dir, vcat(mean(_wave_dir), minimum(_wave_dir), maximum(_wave_dir), std(_wave_dir))...)

    day = Day(_bathy_t0, _bathy_t1, _wave_Tm, _wave_hs, _wave_dir, change, class, input_row)
    push!(days, day)
end

println("normalized(all_tm)     ", minimum(all_tm), "  ", maximum(all_tm))      # RAW: 0.007926686233555585     14.506896347331828
println("normalized(all_hs)     ", minimum(all_hs), "  ", maximum(all_hs))      # RAW: 0.0007081265658494296    2.6867510628524585
println("normalized(all_dir)     ", minimum(all_dir), "  ", maximum(all_dir))   # RAW: -13.212138751720701      17.384923209385946

println("n_offshore: ", n_offshore, ", n_landward: ", n_landward, ", n_none: ", n_none)  # n_offshore: 274, n_landward: 243, n_none: 696

function evaluate(ind::CGPInd, days::Vector{Any}, cfg::NamedTuple)
    errors = Vector{Float64}(undef, length(days))
    reset!(ind)
    for di in 1:length(days)
        day = days[di]

        new_change = process(ind, day.input_row)[1]
        error = (new_change*100 - day.shoreline_change*100)^2

        errors[di] = error
    end
    [mean(errors) * -1]
end

function make_configs(template)
    lst_n_cols = collect(20:10:50)
    lst_n_pop = collect(40:20:200)
    lst_f_recur = collect(0:0.1:1)
    println(length(lst_n_cols))
    println(length(lst_n_pop))
    println(length(lst_f_recur))
    configs = []
    for n_pop in lst_n_pop
        n_elite = 0.1 * n_pop
        for n_cols in lst_n_cols
            for f_recur in lst_f_recur
                config = deepcopy(template)
                two_arity = falses(length(config["functions"]))
                functions = Array{Function}(undef, length(config["functions"]))
                for i in eachindex(config["functions"])
                    fname = config["functions"][i]
                    if CGPFunctions.arity[fname] == 2
                        two_arity[i] = true
                    end
                    functions[i] = eval(Meta.parse(string("CGPFunctions.", fname)))
                end
                config["two_arity"] = two_arity
                config["functions"] = functions

                config["n_population"] = n_pop
                config["n_elite"] = n_elite
                config["columns"] = n_cols
                config["recur"] = f_recur
                config["id"] = string("auto_", "npop", n_pop, "ncols", n_cols, "frecur", f_recur)
                push!(configs, deepcopy(config))
            end
        end
    end
    configs
end

include("cfg_template.jl")
configs = make_configs(template)
println(length(configs))

cfg = configs[1]
mutate(i::CGPInd) = goldman_mutate(cfg, i)

for config in configs
    global cfg
    cfg = (; (Symbol(k)=>v for (k, v) in config)...)
    println(cfg)
    fit(i::CGPInd) = evaluate(i, days, cfg)
    e = CGPEvolution(cfg, fit)
    run!(e)
end


using PolynomialFactors
using AbstractAlgebra
P = PolynomialFactors

## Wilkinson (good out to 50 or so, then issues)

@testset "Test factor over Poly{Int}" begin

## test of factoring over Z[x]
R, x = ZZ["x"]
f = (x-1)*(x-2)*(x-3)
U = P.poly_factor(f)
@test prod(collect(keys(U))) - f == zero(x)

f = (x-1)*(x-2)^2*(x-3)^3
U = P.poly_factor(f)
@test U[x-3] == 3
@test f - prod([k^v for (k,v) in U]) == zero(f)

f = (x-2)*(3x-4)^2*(6x-7)^2
U = P.poly_factor(f)
@test U[6x-7] == 2

f = (x^5 - x - 1)
U = P.poly_factor(f)
@test U[f] == 1

f = x^4
U = P.poly_factor(f)
@test U[x] == 4

d = Dict(x-1 => 3, 2x-3=>4, x^2+x+1=>3)
f = prod([k^v for (k,v) in  d])
U = P.poly_factor(f)
for (k,val) in d
   @test U[k] == val
end

end

@testset "Test factor over Poly{BigInt}" begin
## BigInt
## issue #40 in Roots.jl
R,x = ZZ["x"]
f = x^2 - big(2)^256
U = P.poly_factor(f)
@test U[x - big(2)^128] == 1

p = x^15 - 1
@test length(P.poly_factor(p)) == 4

p = 1 + x^3 + x^6 + x^9 + x^12
@test length(P.poly_factor(p)) == 2
@test p - prod([k^v for (k,v) in P.poly_factor(p)]) == zero(p)

W(n) = prod([x-i for i in 1:20])
U = P.poly_factor(W(20))
@test U[x-5] == 1

## Swinnerton-Dyer Polys are slow to resolve, as over `p` they factor into linears, but over Z are irreducible.
## so the "fish out" step exhausts all possibilities.
S1 = x^2 - 2
U = P.poly_factor(S1)
@test U[S1] == 1

S2 = x^4 - 10*x^2 + 1
U = P.poly_factor(S2)
@test U[S2] == 1

S3 = x^8 - 40x^6 + 352x^4 - 960x^2 + 576
U = P.poly_factor(S3)
@test length(U) == 1

## Cyclotomic polys are irreducible over Z[x] too (https://en.wikipedia.org/wiki/Cyclotomic_polynomial)
C5 = x^4 + x^3 + x^2 + x + 1
U = P.poly_factor(C5)
@test U[C5] == 1

C10 = x^4 - x^3 + x^2 -x + 1
U = P.poly_factor(C10)
@test U[C10] == 1

C15 = x^8 - x^7 + x^5 - x^4 + x^3 - x + 1
U = P.poly_factor(C15)
@test U[C15] == 1

C25 = x^20 + x^15 + x^10 + x^5 + 1
U = P.poly_factor(C25)
@test U[C25] == 1

println("Test factor over Poly{Rational{BigInt}}")
## Rational
R,x = QQ["x"]
f = -17 * (x - 1//2)^3 * (x-3//4)^4
U = P.poly_factor(f)
@test U[-1 + 2x] == 3
@test f - prod([k^v for (k,v) in P.poly_factor(f)]) == zero(f)

end

# @testset "Test rational_roots" begin
# ### Rational roots
# R,x = ZZ["x"]
# W(n) = prod([x-i for i in 1:20])
# V = rational_roots(W(20))
# @test all(V .== 1:20)

# f = (2x-3)^4 * (5x-6)^7
# V = rational_roots(f)
# @test 3//2 in V
# @test 6//5 in V

# end

@testset "Test factormod" begin
    ## factormod
    R,x = ZZ["x"]

    # factormod has elements in GF(q)
    C10 = x^4 - x^3 + x^2 -x + 1
    U = P.factormod(C10, 5)
    # (1+x)^4
    @test length(U) == 1

    C25 = x^20 + x^15 + x^10 + x^5 + 1
    U = P.factormod(C25, 5)
    # (x+4)^20
    @test length(U) == 1

    U = P.factormod(x^4 + 1, 5)
    # (x^2+2)*(x^2+3)
    @test length(U) == 2

    U = P.factormod(x^4 + 1, 2)
    # (x+1)^4
    @test length(U) == 1

    p = 5x^5 - x^4 - 1
    U = P.factormod(p, 7)
    # 5 * (x^4+3*x^3+4*x^2+3*x+4) * (x + 1)
    V = 5 * (x^4+3*x^3+4*x^2+3*x+4) * (x + 1)
    vs = PolynomialFactors.poly_coeffs(V)
    ps = PolynomialFactors.poly_coeffs(p)
    ## v-p is not zero, v-pmod 7 should be:
    @test !all(iszero(vs - ps))
    @test all(iszero.(mod.(vs-ps,7)))

end


abstract type LinearSolver end

struct GaussElimination <: LinearSolver end

struct HansenBliekRohn <: LinearSolver end

function (hbr::HansenBliekRohn)(A, b)
    n = length(b)
    compA = comparison_matrix(A)
    u = compA\mag.(b)
    d = diag(inv(compA))
    α = diag(compA) .- 1 ./d
    α = Interval.(-α, α) #TODO: probably directed rounded is needed here, need to check
    β = @. u/d - mag(b)
    β = Interval.(-β, β)
    x = (b .+ β)./(diag(A) .+ α)

end
function (ge::GaussElimination)(A, b)
    n = length(b)
    A = MMatrix{n, n}(A)
    b = MVector{n}(b)

end
## JACOBI
struct Jacobi <: LinearSolver
    max_iterations::Int
    atol::Float64
end

Jacobi() = Jacobi(20, 0.0)

function (jac::Jacobi)(x, A, b)

    n = length(b)
    @inbounds for _ in 1:jac.max_iterations
        xold = copy(x)
        @inbounds for i in 1:n
            x[i] = b[i]
            @inbounds for j in 1:n
                (i == j) || (x[i] -= A[i, j] * xold[j])
            end
            x[i] = (x[i]/A[i, i]) ∩ xold[i]
        end
        all(x .== xold) && break
        #all(isapprox.(x, xold; atol=atol)) && break
    end
    nothing
end

## GAUSS SEIDEL
struct GaussSeidel <: LinearSolver
    max_iterations::Int
    atol::Float64
end

GaussSeidel() = GaussSeidel(20, 0.0)

function (gs::GaussSeidel)(x, A, b)
    n = length(b)

    @inbounds for _ in 1:gs.max_iterations
        xold = copy(x)
        @inbounds for i in 1:n
            x[i] = b[i]
            @inbounds for j in 1:n
                (i == j) || (x[i] -= A[i, j] * x[j])
            end
            x[i] = (x[i]/A[i, i]) .∩ xold[i]
        end
        all(x .== xold) && break
        #all(isapprox.(x, xold; atol=atol)) && break
    end
    nothing
end

## KRAWCZYK
struct Krawczyk <: LinearSolver
    max_iterations::Int
    atol::Float64
end

Krawczyk() = Krawczyk(20, 0.0)

function (kra::Krawczyk)(x, A, b)
    Ac = mid.(A)
    for i = 1:kra.max_iterations
        xnew  = (Ac\b  - Ac\(A*x) + x) .∩ x
        all( x .== xnew ) && return x
        x = xnew
    end
    return x
end

## wrapper

function solve(A, b, method)

    A, b = precondition(A, b)
    x = enclose(A, b)

    method(x, A, b)

    return x
end

function solve(A, b, method::HansenBliekRohn)
    A, b = precondition(A, b)
    return method(A, b)
end


# ===================================================
# The following schemes are described in the paper
#   Weak exponential schemes for stochastic
#   differential equations with additive noise
#   CM Mora - IMA journal of numerical analysis, 2005
# ===================================================

_drift_jacobian(model::AbstractSDE{D,M}, t0, x0) where {D,M} = ForwardDiff.jacobian(x->drift(model, t0, x), x0)
_drift_jacobian(model::AbstractSDE{1,M}, t0, x0) where {M} = ForwardDiff.derivative(x->drift(model, t0, x), x0)

function step(model::StateIndependentDiffusion, scheme::EulerExponential1, t0, x0, Δw)
  Jμ = _drift_jacobian(model, t0, x0)
  μ = drift(model, t0, x0)
  σ = diffusion(model, t0, x0)
  x = exp(Jμ * scheme.Δt) * (x0 + (μ - Jμ * x0) * scheme.Δt + σ * Δw)
  x
end

function step(model::StateIndependentDiffusion, scheme::EulerExponential2, t0, x0, Δw)
  Jμ = _drift_jacobian(model, t0, x0)
  μ = drift(model, t0, x0)
  σ = diffusion(model, t0, x0)
  x = (I - Jμ * scheme.Δt) \ (x0 + (μ - Jμ * x0) * scheme.Δt + σ * Δw)
  x
end

function step(model::StateIndependentDiffusion, scheme::EulerExponential3, t0, x0, Δw)
  Jμ = _drift_jacobian(model, t0, x0)
  expmJμ = exp(0.5 * Jμ * scheme.Δt)
  μ = drift(model, t0, x0)
  σ = diffusion(model, t0, x0)
  x = expmJμ * (expmJμ * x0 + (μ - Jμ * x0) * scheme.Δt + σ * Δw)
  x
end

function _normal_transition_params(model, scheme::EulerExponential3, t0, x0, x1)
  Jμ = _drift_jacobian(model, t0, x0)
  expmJμ = exp(0.5 * Jμ * scheme.Δt)
  μ = expmJμ * (expmJμ * x0 + (drift(model, t0, x0) - Jμ * x0) * scheme.Δt)
  σ = expmJμ * diffusion(model, t0, x0)
  Σ = scheme.Δt * σ * σ'
  μ, Σ
end


function Newton_G_Stop(nlp :: AbstractNLPModel;
                       x   :: Vector{T}=copy(nlp.meta.x0),
                       stp :: NLPStopping = NLPStopping(nlp,
                                                        NLPAtX(nlp.meta.x0)),
                       NwtDirection :: Function = NwtdirectionSpectral,
                       hessian_rep :: Function = hessian_dense,
                       kwargs...      # eventually options for the direction func
                       ) where T

    n = nlp.meta.nvar
    @info log_header([:iter, :f, :dual, :step, :slope], [Int, T, T, T, T],
                     hdr_override=Dict(:f=>"f(x)", :dual=>"‖∇f‖", :slope=>"∇fᵀd"))
    f = obj(nlp,x)
    ∇f = grad(nlp, x)

    OK = update_and_start!(stp, x = x, fx = f, gx = ∇f)

    τ₀ = 0.0005
    τ₁ = 0.9999

    @info log_row(Any[0, f, norm(∇f)])

    while !OK
        #Hx = hess(nlp, x)
        #H = Matrix(Symmetric(Hx,:L))

        H = hessian_rep(nlp, x)

        #Stopping.update!(stp.current_state; Hx = Hx)

        #Δ, O = eigen(H)
        # Boost negative values of Δ to 1e-8
        # devise an adaptative value for γ
        #γ = 1e-6
        #D = abs.(Δ) + max.((γ .- abs.(Δ)), 0.0) .*ones(n)
        #d = - O*diagm(1.0 ./ D)*O'*∇f

        d = NwtDirection(H, ∇f; kwargs...)

        hp0 = ∇f'*d

        #------------------------------------------
        # Hard coded line search
        hp0 = ∇f'*d
        t=1.0
        # Simple Wolfe forward tracking
        xt = x + t*d
        ∇ft = grad(nlp,xt)
        hp = ∇ft'*d
        ft = obj(nlp, xt)
        #  while  ~wolfe & armijo
        nbW = 0
        while (hp <= τ₁ * hp0) && (ft <= ( f + τ₀*t*hp0)) && (nbW < 10)
            t *= 5
            xt = x + t*d
            ∇ft = grad(nlp,xt)
            hp = ∇ft'*d
            ft = obj(nlp, xt)
            nbW += 1
            @debug "W", ft
        end
        tw = t

        # Simple Armijo backtracking
        nbk = 0
        while (ft > ( f + τ₀*t*hp0)) && (nbk < 20)
            t *= 0.5
            xt = x + t*d
            ft = obj(nlp, xt)
            nbk += 1
            @debug "A", ft
        end
        #------------------------------------------

        x += t*d
        f = ft
        if t!=tw   ∇ft = grad(nlp, xt) end
        ∇f = ∇ft

        OK = update_and_stop!(stp, x = x, gx = ∇f, fx = f)

        @info log_row(Any[stp.meta.nb_of_stop, f, norm(∇f), t, hp0])
    end

    if !stp.meta.optimal
        #@warn "Optimalité non atteinte"
        @warn status(stp,list=true)
    end

    return stp
    #return f, norm(∇f), stp
end


@testset "maximization wrapper" begin
    @testset "univariate" begin
        resmax = maximize(x->x^3, -1, 9)
        resmin = optimize(x->-x^3, -1, 9)
        @test Optim.maximum(resmax) == -Optim.minimum(resmin)
        @test resmax.res.minimum == resmin.minimum
        for meth in (Brent(), GoldenSection())
            resmax = maximize(x->x^3, -1, 9, meth)
            resmin = optimize(x->-x^3, -1, 9, meth)
            @test Optim.maximum(resmax) == -Optim.minimum(resmin)
            @test resmax.res.minimum == resmin.minimum
        end
    end
    @testset "multivariate" begin
        resmax = maximize(x->x[1]^3+x[2]^2, [3.0, 0.0])
        resmin = optimize(x->-x[1]^3-x[2]^2, [3.0, 0.0])
        @test Optim.maximum(resmax) == -Optim.minimum(resmin)
        @test resmax.res.minimum == resmin.minimum
        for meth in (NelderMead(), BFGS(), LBFGS(), GradientDescent(), Newton(), NewtonTrustRegion(), SimulatedAnnealing())
            resmax = maximize(x->x[1]^3+x[2]^2, [3.0, 0.0])
            resmin = optimize(x->-x[1]^3-x[2]^2, [3.0, 0.0])
            @test Optim.maximum(resmax) == -Optim.minimum(resmin)
            @test resmax.res.minimum == resmin.minimum
        end
    end
end


using GeoJSON, FactCheck, Turf

point1 = GeoJSON.Feature(GeoJSON.Point(Float64[-75.343, 39.984]))
point2 = GeoJSON.Feature(GeoJSON.Point(Float64[-75.534, 39.123]))

@fact Turf.distance(point1, point2, :miles) --> roughly(60.37218405837491)
@fact Turf.distance(point1, point2, :kilometers) --> roughly(97.15957803131901)
@fact Turf.distance(point1, point2, :radians) --> roughly(0.015245501024842149)
@fact Turf.distance(point1, point2, :degrees) --> roughly(0.8735028650863799)
@fact Turf.distance(point1, point2) --> distance(point1, point2, :kilometers) # default

polygon = GeoJSON.Polygon(Vector{Vector{Float64}}[Vector{Float64}[[0,0],[10,0],[10,10],[0,10],[0,0]]])
@fact Turf.area(polygon) --> roughly(1232921098571.292)
point = GeoJSON.Point(Float64[0, 0])
@fact Turf.area(point) --> roughly(0)
line = GeoJSON.LineString(Vector{Float64}[[0,0],[1,2]])
@fact Turf.area(line) --> roughly(0)
polygonFC = GeoJSON.FeatureCollection(GeoJSON.Feature[GeoJSON.Feature(polygon, Dict{AbstractString,Any}())])
@fact Turf.area(polygonFC) --> roughly(1232921098571.292)
polygonFC = GeoJSON.FeatureCollection(GeoJSON.Feature[GeoJSON.Feature(polygon, Dict{AbstractString,Any}()),
                                                      GeoJSON.Feature(polygon, Dict{AbstractString,Any}())])
@fact Turf.area(polygonFC) --> roughly(1232921098571.292 * 2)
feature = GeoJSON.Feature(polygon, Dict{AbstractString,Any}())
@fact Turf.area(feature) --> roughly(1232921098571.292)

pt = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/pt.geojson"))
@fact typeof(pt) --> GeoJSON.Feature
pts = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/pts.geojson"))
@fact typeof(pts) --> GeoJSON.FeatureCollection
closest = Turf.nearest(pt, pts)
@fact typeof(closest) --> GeoJSON.Feature
@fact typeof(closest.geometry) --> GeoJSON.Point
@fact closest.geometry.coordinates --> roughly([-75.33, 39.44])

polygon = Turf.bboxpolygon([0.0,0.0,10.0,10.0])
@fact typeof(polygon) --> GeoJSON.Polygon
@fact length(polygon.coordinates) --> 5
@fact polygon.coordinates[1] --> roughly(polygon.coordinates[end])

fc = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/fc.geojson"))
@fact typeof(fc) --> GeoJSON.FeatureCollection
bbox = Turf.envelope(fc)
@fact typeof(bbox) --> GeoJSON.Polygon

pt = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/Point.geojson"))
line = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/LineString.geojson"))
poly = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/Polygon.geojson"))
multiline = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/MultiLineString.geojson"))
multipoly = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/MultiPolygon.geojson"))
@fact Turf.extent(fc) --> roughly([20, -10, 130, 4])
@fact Turf.extent(pt) --> roughly([102, 0.5, 102, 0.5])
@fact Turf.extent(line) --> roughly([102, -10, 130, 4])
@fact Turf.extent(poly) --> roughly([100, 0, 101, 1])
@fact Turf.extent(multiline) --> roughly([100, 0, 103, 3])
@fact Turf.extent(multipoly) --> roughly([100, 0, 103, 3])

@fact Turf.square(Float64[0,0,5,10]) --> roughly([-2.5,0,7.5,10])
@fact Turf.square(Float64[0,0,10,5]) --> roughly([0,-2.5,10,7.5])

@fact Turf.size(Float64[0,0,10,10], 2) --> roughly([-5,-5,15,15])
@fact Turf.size(Float64[0,0,4,4], 1) --> roughly([0,0,4,4])
@fact Turf.size(Float64[0,0,4,4], 2) --> roughly([-2,-2,6,6])
@fact Turf.size(Float64[0,0,4,4], 0.5) --> roughly([1,1,3,3])
@fact Turf.size(Float64[-10,-10,0,0], 2) --> roughly([-15,-15,5,5])
@fact Turf.size(Float64[0,0,10,10], 1.5) --> roughly([-2.5,-2.5,12.5,12.5])
@fact Turf.size(Float64[0,0,10,10], 0.5) --> roughly([2.5,2.5,7.5,7.5])

feature = Turf.center(fc)
@fact typeof(feature) --> GeoJSON.Feature
@fact typeof(feature.geometry) --> GeoJSON.Point
@fact feature.geometry.coordinates --> roughly([75, -3])

pt1 = GeoJSON.Point(Float64[0, 0])
pt2 = GeoJSON.Point(Float64[10, 0])
midpt = Turf.midpoint(pt1, pt2)
@fact typeof(midpt) --> GeoJSON.Point
@fact midpt.coordinates --> roughly([0, 5])
pt1 = GeoJSON.Feature(GeoJSON.Point(Float64[1, 1]))
pt2 = GeoJSON.Feature(GeoJSON.Point(Float64[11, 11]))
@fact typeof(midpt) --> GeoJSON.Feature
@fact midpt.geometry.coordinates --> roughly([6, 6])

pt1 = GeoJSON.Feature(GeoJSON.Point(Float64[-75.4, 39.4]))
pt2 = GeoJSON.Feature(GeoJSON.Point(Float64[-75.534, 39.123]))
bearing = Turf.bearing(pt1, pt2) # what's the expected value?

pt1 = GeoJSON.Feature(GeoJSON.Point(Float64[-75.0, 39.0]))
dist = 100.0
bearing = 180.0
pt2 = Turf.destination(pt1, dist, bearing) # what's the expected value?

route1 = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/route1.geojson"))
route2 = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/route2.geojson"))
Turf.linedistance(route1, :miles)
# t.equal(Math.round(lineDistance(route1, 'miles')), 202);
#     t.true((lineDistance(route2, 'kilometers') - 742) < 1 &&
#             (lineDistance(route2, 'kilometers') - 742) > (-1) );

line = GeoJSON.parsefile(joinpath(dirname(@__FILE__),"geojson/dc-line.geojson"))
pt1 = GeoJSON.Feature(Turf.along(line, 1, :miles))
pt2 = GeoJSON.Feature(Turf.along(line, 1.2, :miles))
pt3 = GeoJSON.Feature(Turf.along(line, 1.4, :miles))
pt4 = GeoJSON.Feature(Turf.along(line, 1.6, :miles))
pt5 = GeoJSON.Feature(Turf.along(line, 1.8, :miles))
pt6 = GeoJSON.Feature(Turf.along(line, 2, :miles))
pt7 = GeoJSON.Feature(Turf.along(line, 100, :miles))
pt8 = GeoJSON.Feature(Turf.along(line, 0, :miles))
fc = GeoJSON.FeatureCollection(GeoJSON.Feature[pt1, pt2, pt3, pt4, pt5, pt6, pt7, pt8])
for feature in fc.features
    @fact typeof(feature.geometry) --> GeoJSON.Point
end
@fact length(fc.features) --> 8
@fact fc.features[end].geometry.coordinates --> roughly(pt8.geometry.coordinates)


function irfvar(y,p,H)

# Luca Brugnolini 2014
# Based on Kilian code
# Point estimates and stdev of structural IRFs identified by recursive structure
# Input : data y, lag order p, horizon for impulse response H
# Output: IRF, Std error of IRF, A0inv for structural identification,
#         slope parameter estimates B, cov matrix of estimated residual SIG
#         estimated residual U, reduced form IRF

(T,K) = size(y);

# OLS
Y = (y[p+1:T,:])';                           # LHS variable for OLS
X = ones(1,T-p);                             # constant term

for j=1:p
    X = [X; y[p+1-j:T-j,:]'];                # RHS variable for OLS

end

B   = (Y*X')/(X*X');                         # OLS estimates **
A   = B[:,2:end];                            # OLS estimates w/o intercept
U   = Y-B*X;                                 # estimated residual
SIG = (Y-B*X)*(Y-B*X)'/(T-p-p*K-1);          # estimated sigma(U)
SIGb = kron(inv(X*X'/(T-p)),SIG);            # varcov of vec(B):(K^2*p+K)^2
SIGa = SIGb[K+1:end,K+1:end];                # varcov of vec(A):(K^2*p)^2

# IRF and stdev(IRF)
#  Calculation of stdev follows Lutkepohl(2005) p.111,93
A = [A; eye(K*(p-1)) zeros(K*(p-1),K)];         # Companion form
A0inv = chol(SIG)';                              # recursive str assumption
rfIRF = zeros(K^2,H+1);                          # reduced form IRF
IRF   = zeros(K^2,H+1);                          # structural IRF
STD   = zeros(K^2,H+1);                          # stdev of structural IRF
COV2   = zeros(K^2,H+1);                         # COV of structural IRF due to SIG
J = [eye(K) zeros(K,K*(p-1))];                   # J matrix
L = elimat(K);                                   # elimination matrix
Kk = commutation(K,K);                           # commutation matrix
Hk = L'/(L*(eye(K^2)+Kk)*kron(A0inv,eye(K))*L');  # H matrix **
D = duplication(K);                              # duplication matrix
Dplus = (D'*D)\D';
SIGsig = 2*Dplus*kron(SIG,SIG)*Dplus';           # sigma(vech(sigma))
Cbar0 = kron(eye(K),J*eye(K*p)*J')*Hk;           # Cbar
rfIRF[:,1] = vec((J*eye(K*p)*J')');
IRF[:,1] = vec((J*eye(K*p)*J'*A0inv)');
STD[:,1] = vec((reshape(diag(real(sqrt(complex(Cbar0*SIGsig*Cbar0'/(T-p))))),K,K))')
COV2[:,1] = vec((reshape(diag((Cbar0*SIGsig*Cbar0'/(T-p))),K,K))');

for h=1:H
    Gi = zeros(K^2,K^2*p);

    for m=0:(h-1)
        Gi += kron(J*(A')^(h-1-m),J*(A^m)*J');
    end;

    C = kron(A0inv',eye(K))*Gi;
    Cbar = kron(eye(K),J*A^h*J')*Hk;

    rfIRF[:,h+1] = vec((J*A^h*J')');
    IRF[:,h+1] = vec((J*A^h*J'*A0inv)');
    STD[:,h+1] = vec((reshape(diag(real(sqrt(complex(C*SIGa*C'+Cbar*SIGsig*Cbar')/(T-p)))),K,K))')
    COV2[:,h+1] = vec((reshape(diag(((Cbar*SIGsig*Cbar')/(T-p))),K,K))');

end

    return IRF,STD,A0inv,B,SIG,U,rfIRF,COV2;

end


using JuMP
using Clp

DEMANDS = [10011.22115007576, 9993.540857825825, 10029.493305561082, 9995.513032624265, 10040.407110819071, 10054.429367942921, 10059.76519849408, 9987.000868933652, 10035.54111237392, 9979.076532648018, 9973.268512810797, 10001.801162849199, 9977.013028384847, 10043.66957251199, 10052.449661342594, 9967.208355572673, 10000.69477091062, 10030.05030028466, 10004.993099108486, 9990.907712391863, 10048.51255532599, 10006.528026036742, 9995.689091324659, 9990.868630720339, 9997.195144776251, 9971.848670139489, 9945.536501292985, 9948.336004602528, 9997.577323187325, 9970.52023558712, 10069.401333593078, 10043.51650796613, 10047.132683479025, 10072.330774070775, 10031.89101682498, 10020.955889705769, 9934.338643331117, 10014.21570856934, 10000.845689612031, 10002.553802750448, 10029.850884936972, 9948.279191163045, 9971.233600233412, 9921.28110096277, 9969.291956593264, 9941.660187262343, 9970.955683172295, 10121.63150429881, 9890.3148430648, 10055.434949187049]

EXCESS = [5262.38201199916, 5860.755189729461, 4984.212894426032, 5363.8531442275635, 4620.34431447503, 3825.9519209367336, 5343.680950738127, 5233.478472018467, 5123.639554039856, 4635.80130213247, 4505.55692176387, 4218.191006304447, 5501.7375350818, 5206.40743566552, 4848.244095467865, 4922.152296817367, 5405.514169790265, 4156.03729448171, 4824.382398257035, 4000.2715617359813, 4390.76908613548, 4842.974771246174, 5220.6774383278635, 4178.847030718337, 5881.9659847155, 5380.195019617506, 4621.415119581128, 6104.345451062469, 5160.047534233344, 5370.740725813598, 4229.219081632929, 5048.335218378297, 4820.895310902216, 5144.298907980797, 4765.881280652281, 4707.799430473926, 4008.493355905563, 4681.96931929508, 6194.13279529005, 5615.828913696583, 4279.739726326251, 4650.89876777662, 4966.17758135978, 5151.236396106044, 5177.558350548357, 5293.24954830797, 4730.71724368456, 5493.7205140204305, 4524.476628100214, 4795.57008939593]
NUC_COST = 300
COAL_COST = 250
T = length(DEMANDS)

m = Model(Clp.Optimizer)
@variable(m, x[1:T] >= 0)
@variable(m, y[1:T] >= 0)
@variable(m, w[1:T] >= 0)
@variable(m, z[1:T] >= 0)
@objective(m, Min, sum(NUC_COST * y[i] + COAL_COST * x[i] for i = 1:T))

for t = 1:T
    @constraint(m, w[t] == sum(x[t:-1:max(1, t-19)]))
    @constraint(m, z[t] == sum(y[t:-1:max(1, t-14)]))
end

@constraint(m, w + z + EXCESS .≥ DEMANDS)
@constraint(m, 0.8 * z - 0.2 * w .≤ 0.2 * EXCESS)

println(m)

optimize!(m)

#status = termination_status(m)

#println("Solution status: ", status)

#println("Objective value: ", objective_value(m))
println("x = ", value.(x))
println("y = ", value.(y))


nq = 3 + 4 + 3 # position + quaternion + leg joints

function quaternion_rotation_matrix(q)
	r, i, j, k  = q

	r11 = 1.0 - 2.0 * (j^2.0 + k^2.0)
	r12 = 2.0 * (i * j - k * r)
	r13 = 2.0 * (i * k + j * r)

	r21 = 2.0 * (i * j + k * r)
	r22 = 1.0 - 2.0 * (i^2.0 + k^2.0)
	r23 = 2.0 * (j * k - i * r)

	r31 = 2.0 * (i * k - j * r)
	r32 = 2.0 * (j * k + i * r)
	r33 = 1.0 - 2.0 * (i^2.0 + j^2.0)

	[r11 r12 r13;
     r21 r22 r23;
     r31 r32 r33]
end

function rotation_matrix_x(a)
    [1.0 0.0    0.0
     0.0 cos(a) -sin(a)
     0.0 sin(a)  cos(a)]
end

leg1_offset = [1.0; 1.0; 0.0]
l_shoulder = 1.0
l_thigh = 1.0
l_calf = 1.0
function kinematics_thigh_calf(a1, a2)
    [
     l_thigh * sin(a1) + l_calf * sin(a2);
     0.0;
    -l_thigh * cos(a1) - l_calf * cos(a2);
    ]
end

function foot_kinematics(q)
    pos = q[1:3]
    quat = q[4:7]
    joint = q[8:10]

    k_thigh_calf = kinematics_thigh_calf(q[9], q[10])
    k_shoulder_thigh_calf = rotation_matrix_x(q[8]) * (k_thigh_calf + [0.0; l_shoulder; 0.0])
    k_offset_shoulder_thigh_calf = k_shoulder_thigh_calf + leg1_offset
    k_body_offset_shoulder_thigh_calf = quaternion_rotation_matrix(quat) * k_offset_shoulder_thigh_calf
    k = k_body_offset_shoulder_thigh_calf + pos

    return k
end

q0 = [0.0; 0.0; 0.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0]
foot_kinematics(q0)

@variables q[1:nq] q̇[1:nq]

k = foot_kinematics(q)
j = Symbolics.jacobian(k, q)

k_func = eval(Symbolics.build_function(k, q)[1])
j_func = eval(Symbolics.build_function(j, q)[1])

function lagrangian(q, q̇)
    L = 0.0
    mass = 1.0
    inertia = 1.0
    gravity = 9.81
    k_foot = k_func(q)
    v_foot = j_func(q) * q̇
    L += mass * transpose(v_foot) * v_foot
    L -= mass * gravity * k_foot[3]

    return L
end

L = lagrangian(q, q̇)

ddL = Symbolics.hessian(L, [q; q̇])#, simplify=true)
dLq = Symbolics.gradient(L, q)#, simplify=true)
ddLq̇q = ddL[nq .+ (1:nq), 1:nq]

M = ddL[nq .+ (1:nq), nq .+ (1:nq)]
C = ddLq̇q * q̇ - dLq

mass_matrix_sym = eval(Symbolics.build_function(M, q)[1])
dynamics_bias_sym = eval(Symbolics.build_function(C, q, q̇)[1])

function lagrangian_derivatives(q, v)
	D1L = -1.0 * dynamics_bias_sym(q, v)
    D2L = mass_matrix_sym(q) * v
	return D1L, D2L
end

function dynamics(h, q0, q1, u1, w1, λ1, q2)
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 + λ1)
end

@variables h[1:1] q0[1:nq] q1[1:nq] u1[1:0] w1[1:0] λ1[1:nq] q2[1:nq]
d = dynamics(h, q0, q1, u1, w1, λ1, q2)


module ScalingDimspecTests

using Interpolations, DualNumbers, Base.Test

xs = -pi:(2pi/10):pi-2pi/10
ys = -2:.1:2
f(x,y) = sin(x) * y^2

itp = interpolate(Float64[f(x,y) for x in xs, y in ys], (BSpline(Quadratic(Periodic())), BSpline(Linear())), OnGrid())
sitp = scale(itp, xs, ys)

for (ix,x) in enumerate(xs), (iy,y) in enumerate(ys)
    @test_approx_eq_eps sitp[x,y] f(x,y) sqrt(eps(1.0))

    g = gradient(sitp, x, y)
    fx = epsilon(sitp[dual(x,1), dual(y,0)])
    fy = epsilon(sitp[dual(x,0), dual(y,1)])

    @test_approx_eq_eps g[1] fx sqrt(eps(1.0))
    @test_approx_eq_eps g[2] fy sqrt(eps(1.0))
end

end


include("cantorPairing.jl");

function cantorPairingRecursive(arrayTuple)
	len = length(arrayTuple);
	if (len == 2)
		return cantorPairing(arrayTuple[1], arrayTuple[2]);
	else
		return cantorPairing(arrayTuple[1], cantorPairingRecursive(arrayTuple[2:len]));
	end
end


include("graph.jl")
include("zadatak71BFS.jl")

M = [0 1 1 0 0 0;
     1 0 0 1 0 1;
     1 0 0 0 1 0;
     0 1 0 0 0 0;
     0 0 1 0 0 1;
     0 1 0 0 1 0]
G = nonEmptyGraph(M)
path = BFS!(G, 1, 6)

println("Path = $path");


export excess_solution_check, excess_solution_check!, OverdeterminedTracker
"""
    excess_solution_check!(path_result::PathResult,
                           F::RandomizedSystem,
                           newton_cache = NewtonCache(F.system))

Assigns to the [`PathResult`](@ref) `path_result` the `return_code` `:excess_solution` if
the `path_result` is a solution of the randomized system `F` but not of
the polynomial system underlying `F`.
This is performed by using Newton's method for non-singular solutions and comparing the
residuals of the solutions for singular solutions.
"""
function excess_solution_check!(
    path_result::PathResult,
    F::RandomizedSystem,
    newton_cache::NewtonCache = NewtonCache(F.system),
)
    is_success(path_result) || return path_result
    if is_nonsingular(path_result)
        acc = max(accuracy(path_result), eps())
        max_abs_norm_first_update = 100 * sqrt(acc)
        res = newton(
            F.system,
            solution(path_result),
            nothing,
            InfNorm(),
            newton_cache;
            extended_precision = true,
            atol = 1e3 * acc,
            rtol = 1e3 * acc,
            max_abs_norm_first_update = max_abs_norm_first_update,
        )
        if !is_success(res)
            path_result.return_code = :excess_solution
            path_result.accuracy = res.accuracy
            path_result.residual = res.residual
        end
    else
        # for singular solutions compare residuals due to lack of something better right now
        evaluate!(newton_cache.r, F.system, solution(path_result))
        system_residual = LA.norm(newton_cache.r, InfNorm())
        if system_residual > 100 * residual(path_result)
            path_result.return_code = :excess_solution
            path_result.residual = system_residual
        end
    end

    path_result
end

struct ExcessSolutionCheck{S<:RandomizedSystem,M} <: Function
    system::S
    newton_cache::NewtonCache{M}
end
(check::ExcessSolutionCheck)(R::PathResult) =
    excess_solution_check!(R, check.system, check.newton_cache)

"""
    excess_solution_check(F::RandomizedSystem)

Returns a function `λ(::PathResult)` which performs the excess solution check.
The call `excess_solution_check(F)(path_result)` is identical to
`excess_solution_check!(F, path_result)`. See also [`excess_solution_check!`](@ref).
"""
excess_solution_check(F::RandomizedSystem) = ExcessSolutionCheck(F, NewtonCache(F.system))

"""
    OverdeterminedTracker(tracker::AbstractPathTracker, F::RandomizedSystem)

Wraps the given [`AbstractPathTracker`](@ref) `tracker` to apply
[`excess_solution_check`](@ref) for the given randomized system `F` on each path result.
"""
struct OverdeterminedTracker{T<:AbstractPathTracker,E<:ExcessSolutionCheck} <:
       AbstractPathTracker
    tracker::T
    excess_solution_check::E
end
OverdeterminedTracker(T::AbstractPathTracker, F::RandomizedSystem) =
    OverdeterminedTracker(T, excess_solution_check(F))

track(T::OverdeterminedTracker, x; kwargs...) =
    T.excess_solution_check(track(T.tracker, x; kwargs...))


export poisson_trunc_point, poisson_cum_rtp

function poisson_trunc_point(λ, tol)
    k = 0.0
    expo = -λ
    term = exp(expo)
    term_sum = term
    log_λ = log(λ)
    probs = Vector{Float64}()

    while term_sum < tol / 2.0
        k += 1.0
        expo += log_λ - log(k)
        term = exp(expo)
        term_sum += term
    end

    ltp = Int(k)
    push!(probs, term)

    while 1.0 - term_sum > tol / 2.0
        k += 1.0
        expo += log_λ - log(k)
        term = exp(expo)
        term_sum += term
        push!(probs, term)
    end

    rtp = Int(k)
    @assert length(probs) == (rtp - ltp + 1)

    return ltp, rtp, probs
end

function poisson_cum_rtp(λ, t, tol)
    k = 0.0
    log_λ = log(λ)
    expo = -λ
    term = exp(expo)
    term_sum = term
    probs = Vector{Float64}()

    push!(probs, term)

    while tol < t * (1.0 - term_sum)
        k += 1.0
        expo += log_λ - log(k)
        term = exp(expo)
        term_sum += term
        push!(probs, term)
    end
    rtp = Int(k)
    @assert rtp + 1 == length(probs)
    return rtp, probs
end


using EngEconomics, Roots

# Given
aInit = -1000000
aAnnualSavings = 100000
aAnnualG = 50000
MARR = 0.10
aLife = 6

# Determine Payback Period


function get_num_trees(input_file, num_right_steps, num_down_steps, tree_symbol)
    num_trees = open(input_file) do file
        line_width = -1
        current_position = 1 # we start on the top left
        num_trees = 0
        lines_to_skip = 0

        for line in eachline(file)
            if lines_to_skip > 0
                lines_to_skip -= 1
                continue
            end

            cleaned_line = collect(strip(line))
            if line_width == -1
                line_width = length(cleaned_line)
            end
            equality = string(cleaned_line[current_position]) == tree_symbol

            if equality
                num_trees += 1
            end

            current_position = current_position + num_right_steps

            # we have periodic boundary positions
            if current_position > line_width
                current_position = current_position % line_width
            end

            # this is just a way of coding skipping lines
            lines_to_skip += (num_down_steps - 1)
        end
        num_trees
    end
    num_trees
end

input_steps = [
    (input_file = "day3_input.txt", num_right_steps = 1, num_down_steps = 1, tree_symbol = "#"),
    (input_file = "day3_input.txt", num_right_steps = 3, num_down_steps = 1, tree_symbol = "#"),
    (input_file = "day3_input.txt", num_right_steps = 5, num_down_steps = 1, tree_symbol = "#"),
    (input_file = "day3_input.txt", num_right_steps = 7, num_down_steps = 1, tree_symbol = "#"),
    (input_file = "day3_input.txt", num_right_steps = 1, num_down_steps = 2, tree_symbol = "#")]

all_trees = []
for input_settings in input_steps
    num_trees = get_num_trees(input_settings...)
    println("$input_settings num_trees: $num_trees")
    append!(all_trees, num_trees)
end

prod(all_trees)


using Flux: onehot, onehotbatch
using Random

const alphabet = [false, true]  # 0, 1

parity(x) = reduce(xor, x)

gendata(n::Int, k::Int) = gendata(n, k:k)
function gendata(n::Int, k::UnitRange{Int})
    X = bitrand.(rand(k, n))
    return [(onehotbatch(x, alphabet), onehot(y, alphabet)) for (x, y) in zip(X, parity.(X))]
end


# closest points using a particular metric, using a K-D tree implemented by NearestNeighbors
function closest_points(P,Q; metric=Euclidean())
    kdtree = KDTree(Q, metric)
    nearestidx, dists = nn(kdtree, P)
    return [(i,nearestidx[i]) for i=1:size(P,2)]
end

# generate matrices for Kabsch from a list of correspondences
matched_points(P::AbstractMatrix, Q::AbstractMatrix, matches) = (hcat([P[:,i] for (i,j) in matches]...), hcat([Q[:,j] for (i,j) in matches]...))
matched_points(P, Q, matches) = ([P[i] for (i,j) in matches], [Q[j] for (i,j) in matches])
matched_points(P,Q; correspondence=closest_points, kwargs...) = matched_points(P, Q, correspondence(P,Q; kwargs...))


if nworkers()<2
	addprocs(3-nworkers())
end

using jInv.Mesh
using jInv.LinearSolvers
using jInv.ForwardShare
using ForwardHelmholtz
using FWI
using Base.Test

# get a small mesh
domain = [0 2.0 0 2.2]
n      = [12;8;]
M      = getRegularMesh(domain,n)
omega  = 0.1;

# sources / receivers on top edge
idx    = reshape(collect(1:prod(n+1)),tuple(n+1...))
ib     = idx[:,1];
I      = speye(prod(n+1))
Q      = I[:,vec(ib)]
R      = copy(Q)

# get param without parallelizaion
ABLpad = 2;
gamma = getABL(M,true,ones(Int64,M.dim)*ABLpad,8*pi);
attenuation = 0.01*4*pi;
gamma += attenuation; # adding Attenuation.
gamma = gamma[:];
Ainv = getJuliaSolver();

pFor = getFWIparam(omega, one(Complex128),gamma,Q,R,M,Ainv,[workers()[1]])[1];

m0   = rand(tuple(n+1...))+1.0
dobs, = getData(vec(m0),fetch(pFor[1]));

# parallelize over sources
pForp,continuationDivision,SourcesSubInd = getFWIparam(omega, one(Complex128),gamma,Q,R,M,Ainv,workers())
dobsRF, = getData(vec(m0),pForp)
Dobs = zeros(eltype(dobs),size(dobs));
for k=1:length(dobsRF)
	dobsk = fetch(dobsRF[k]);
	Dobs[:,SourcesSubInd[k]] = dobsk;
end

@test vecnorm(Dobs-dobs)/vecnorm(dobs) < 1e-8


# Test the gaussian pre-image code
# Danny Perry (dperry@cs.utah.edu)
# May 2015

using Debug
using PyPlot
using Kernel

@debug function main()

	n = 100
	rad = 10
	thetas = 2*pi*rand(n)
	X = rad*[cos(thetas) sin(thetas)]
	Xtest = mean(X,1)

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot(Xtest[1],Xtest[2],"bo")

	iters = 1
	neibsize = 2
	sigma = 10.0
	#result,neibs = GaussianKernelDenoise(X, Xtest, iters, neibsize, sigma)
	K = GaussianKernel(X,sigma)
	Ktest = GaussianKernel(X,Xtest,sigma)
	result,neibs = GaussianKernelPreImage(X, K, Ktest, neibsize, sigma)
	println("==================")
	println("neibsize: ", neibsize)
	println("input: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))
	println("==================")

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot([Xtest[1],result[1]],[Xtest[2],result[2]],"bx--")
	plot(result[1],result[2],"gs")
	plot(X[vec(neibs[1,:]),1],X[vec(neibs[1,:]),2],"y^")
	title(@sprintf("%d",neibsize))

	neibsize = 4
	result,neibs = GaussianKernelDenoise(X, Xtest, iters, neibsize, sigma)
	println("==================")
	println("neibsize: ", neibsize)
	println("input: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))
	println("==================")

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot([Xtest[1],result[1]],[Xtest[2],result[2]],"bx--")
	plot(result[1],result[2],"gs")
	plot(X[vec(neibs[1,:]),1],X[vec(neibs[1,:]),2],"y^")
	title(@sprintf("%d",neibsize))

	neibsize = 10
	result,neibs = GaussianKernelDenoise(X, Xtest, iters, neibsize, sigma)
	println("==================")
	println("neibsize: ", neibsize)
	println("input: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))
	println("==================")

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot([Xtest[1],result[1]],[Xtest[2],result[2]],"bx--")
	plot(result[1],result[2],"gs")
	plot(X[vec(neibs[1,:]),1],X[vec(neibs[1,:]),2],"y^")
	title(@sprintf("%d",neibsize))

	theta = rand()*2*pi
	Xtest = rad*[cos(theta) sin(theta)]
	neibsize = 2
	result,neibs = GaussianKernelDenoise(X, Xtest, iters, neibsize, sigma)
	println("neibsize: ", neibsize)
	println("input: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot([Xtest[1],result[1]],[Xtest[2],result[2]],"bx--")
	plot(result[1],result[2],"gs")
	plot(X[vec(neibs[1,:]),1],X[vec(neibs[1,:]),2],"y^")
	title(@sprintf("%d",neibsize))

	neibsize = 4
	result,neibs = GaussianKernelDenoise(X, Xtest, iters, neibsize, sigma)
	println("neibsize: ", neibsize)
	println("input: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot([Xtest[1],result[1]],[Xtest[2],result[2]],"bx--")
	plot(result[1],result[2],"gs")
	plot(X[vec(neibs[1,:]),1],X[vec(neibs[1,:]),2],"y^")
	title(@sprintf("%d",neibsize))

	neibsize = 10
	result,neibs = GaussianKernelDenoise(X, Xtest, iters, neibsize, sigma)
	println("neibsize: ", neibsize)
	println("exact: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))

	figure()
	plot(X[:,1],X[:,2],"ro")
	hold(:on)
	plot([Xtest[1],result[1]],[Xtest[2],result[2]],"bx--")
	plot(result[1],result[2],"gs")
	plot(X[vec(neibs[1,:]),1],X[vec(neibs[1,:]),2],"y^")
	title(@sprintf("%d",neibsize))

	PyPlot.plt.show()
	@bp

end

function testdata()
	X = reshape([1:33],11,3)
	Xtest = X[5,:]
	X = X[[1:4,6:end],:]

	println("Xtest: ", Xtest)
	println("X:\n", X)

	return X,Xtest
end

@debug function simple(neibsize = 2)
	println("~~~~~~~~~~~~~~~~~~~")
	println(" neighbors approach")
	println("~~~~~~~~~~~~~~~~~~~")

	X,Xtest = testdata()

	sigma = 1.0
	K = GaussianKernel(X,sigma)
	Ktest = GaussianKernel(X,Xtest,sigma)
	result,neibs = GaussianKernelPreImage(X, K, Ktest, neibsize, sigma)

	println("neibs: ", neibs)

	println("===================")
	println("neibsize: ", neibsize)
	println("exact: ", Xtest)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))
	println("===================")

	return result
end

@debug function simple2()
	println("~~~~~~~~~~~~~~~~~~~")
	println(" fixed point approach")
	println("~~~~~~~~~~~~~~~~~~~")

	X,Xtest = testdata()

	sigma = 1.0
	K = GaussianKernel(X,sigma)
	U,S,V = svd(K)
	Ktest = GaussianKernel(X,Xtest,sigma)
	y = vec(Ktest' * V)
	z_init = mean(X,1)
	result,iters = GaussianKernelPreImage(y,V, X,z_init, sigma, 1e-3,100)

	println("===================")
	println("exact: ", Xtest)
	println("iters: ", iters)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))
	println("===================")
end

@debug function simple3()
	println("~~~~~~~~~~~~~~~~~~~")
	println(" both approach")
	println("~~~~~~~~~~~~~~~~~~~")

	X,Xtest = testdata()

	neibsize = 3
	sigma = 1.0
	K = GaussianKernel(X,sigma)
	U,S,V = svd(K)
	Ktest = GaussianKernel(X,Xtest,sigma)
	y = vec(Ktest' * V)
	z_init = simple(neibsize)
	println("zinit: ", z_init)
	println("diff: ", norm(Xtest-z_init))
	result,iters = GaussianKernelPreImage(y,V, X,z_init, sigma, 1e-20,10000)

	println("===================")
	println("neibsize: ", neibsize)
	println("exact: ", Xtest)
	println("iters: ", iters)
	println("result: ", result)
	println("diff: ", norm(Xtest-result))
	println("===================")
end

simple() # method 1
simple2() # method 2
simple3() # method 1 and 2
#main()


using Random
using GLPK
using MathOptInterface

const MOI  = MathOptInterface
const MOIU = MathOptInterface.Utilities;

using Test

D = 10  # variable dimension
N = 20  # no of inequality constraints

s = rand(N)
s = 2*s.-1
λ = max.(-s, 0)
s = max.(s, 0)
x̂ = rand(D)
A = rand(N, D)
b = A*x̂ + s
c = -A' * λ

# can feed dual problem to optimizer like this:
# model = MOI.instantiate(dual_optimizer(GLPK.Optimizer), with_bridge_type=Float64)

model = GLPK.Optimizer()
x = MOI.add_variables(model, D)

# define objective
objective_function = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(c, x), 0.0)
MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objective_function)
MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

# will be useful later
constraint_indices = []

# set constraints
for i in 1:N
    push!(constraint_indices, MOI.add_constraint(model,MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(A[i,:], x), 0.),MOI.LessThan(b[i])))
end

for i in 1:D
    push!(constraint_indices, MOI.add_constraint(model,MOI.SingleVariable(x[i]),MOI.GreaterThan(0.)))
end

MOI.optimize!(model)

@assert MOI.get(model, MOI.TerminationStatus()) in [MOI.LOCALLY_SOLVED, MOI.OPTIMAL]

x̄ = MOI.get(model, MOI.VariablePrimal(), x);  # solution

@assert abs(c'x̄ - c'x̂) <= 1e-8   # sanity check

##
## Checking KKT conditions
##

is_less_than(set::S) where {S<:MOI.AbstractSet} = false
is_less_than(set::MOI.LessThan{T}) where T = true

for con_index in constraint_indices
    con_value = MOI.get(model, MOI.ConstraintPrimal(), con_index)
    set = MOI.get(model, MOI.ConstraintSet(), con_index)
    μ = MOI.get(model, MOI.ConstraintDual(), con_index)

    if is_less_than(set)
        # μ[i]*(Ax - b)[i] = 0
        @assert μ*(con_value - set.upper) < 1e-10
    else
        # μ[j]*x[j] = 0
        @assert μ*(con_value - set.lower) < 1e-10
    end
end


#!/usr/bin/env julia

# A Julia script that reproduces expected results of the test module
# 'test/intfunctionTest.cpp'.
#
# From a shell, run the script as:
#   julia /path/to/intfunction.jl
#
# From Julia, the script can be run as:
#   include("path/to/intfunction.jl")

function intfunctionTest()
    nums = [ 12, 100, 37423 ]
    for n = nums
       println("floor(sqrt(", n, ")) = ", convert(Integer, floor(sqrt(n))) )
    end

    println()
    for n = nums
       println("ceil(sqrt(", n, ")) = ", convert(Integer, ceil(sqrt(n))) )
    end

    println()
    nums = [ 1, 8, 63, 64, 1024, 1025, 65000 ]
    for n = nums
       println("floor(log2(", n, ")) = ", convert(Integer, floor(log2(n))) )
    end

    println()
    for n = nums
       println("ceil(log2(", n, ")) = ", convert(Integer, ceil(log2(n))) )
    end

    println()
    nums = [ 1, 16, 65 ]
    for n = nums
       println("floor(log2(", n, ")) = ", convert(Integer, floor(log2(n))) )
    end

    println()
    for n = nums
       println("ceil(log2(", n, ")) = ", convert(Integer, ceil(log2(n))) )
    end
end

intfunctionTest()


using QuasiGeostrophy, Test
include(pwd() * "/test/test_utils.jl")
@wrapper u=1 σ=1

@macroexpand  equ = @to_equation σ=u
equ = @to_equation σ=u
@pde_system pde_system = [
    σ = ∂x(u),
    ∂t(u) = -∂x(u * u - ∂x(σ)),
]

@testset "Basic Equation Test" begin
    @pde_system pde_system = [
        σ = ∂x(u),
        ∂t(u)= -∂x(u * u - ∂x(σ)),
    ]
    @test pde_system[1].rhs == ∂x(u)
end


#####
# This includes short-hands that are convenient but perhaps confusing
#####

export chebyshevt, chebyshevu, legendre, ∫, ⨜, ⨍, ChebyshevWeight, 𝕀, 𝕌, 𝒟

## Constructors

Fun() = Fun(identity)
Fun(d::Domain) = Fun(identity,d)
Fun(d::Space) = Fun(identity,d)

## Chebyshev & Legendre polynomials

chebyshevt(n::Int,d::IntervalOrSegment{T}) where {T<:Number} = Fun(Chebyshev(d),[zeros(T,n);one(T)])
chebyshevu(n::Int,d::IntervalOrSegment{T}) where {T<:Number} =
    mod(n,2) == 1 ? Fun(Chebyshev(d),interlace(zeros(T,div(n+2,2)),2fill(one(T),div(n+2,2)))) :
                    Fun(Chebyshev(d),interlace(2fill(one(T),div(n+2,2)),zeros(T,div(n+2,2)))[1:n+1]-[one(T);zeros(T,n)])
legendre(n::Int,d::IntervalOrSegment{T}) where {T<:Number} = Fun(Legendre(d),[zeros(T,n);one(T)])

for poly in (:chebyshevt,:chebyshevu,:legendre)
    @eval begin
        $poly(::Type{T}, n::Int) where {T<:Number} = $poly(n,ChebyshevInterval{T}())
        $poly(n::Int) = $poly(Float64,n)
        $poly(n::AbstractRange, d::IntervalOrSegment{T}) where {T<:Number} = map(i->$poly(i,d),n)
    end
end

ChebyshevWeight(d,k)=k==0 ? JacobiWeight(-0.5,-0.5,d) : JacobiWeight(0.5,0.5,d)
ChebyshevWeight(d)=ChebyshevWeight(d,0)
ChebyshevWeight(k::Integer)=ChebyshevWeight(Segment(),k)
ChebyshevWeight()=ChebyshevWeight(0)

# shorthand for second order

ivp(d) = ivp(d,2)
bvp(d) = bvp(d,2)

## diff

# use conj(transpose(f)) for ArraySpace
adjoint(f::Fun)=differentiate(f)

∫(f::Fun)=integrate(f)
⨜(f::Fun)=cumsum(f)

for OP in (:Σ,:∮,:⨍,:⨎)
    @eval $OP(f::Fun)=sum(f)
end

∇(F::MultivariateFun) = grad(F)
dot(∇::Function,F::Vector{M}) where {M<:MultivariateFun} = div(F)
cross(∇::Function,F::Vector{M}) where {M<:MultivariateFun} = curl(F)

## Domains

const 𝕀 = ChebyshevInterval()
const ℝ = Line()
const 𝕌 = Circle()

𝒟 = Derivative()
Δ = Laplacian()


## The standard univariate linear Gaussian state-space model example

module LinearGaussian

using SequentialMonteCarlo
using RNGPool
import SMCExamples.Particles.Float64Particle

struct LGTheta
  A::Float64
  Q::Float64
  C::Float64
  R::Float64
  x0::Float64
  v0::Float64
end

function makeLGModel(theta::LGTheta, ys::Vector{Float64})
  n::Int64 = length(ys)
  sqrtQ::Float64 = sqrt(theta.Q)
  invRover2::Float64 = 0.5/theta.R
  sqrtv0::Float64 = sqrt(theta.v0)
  logncG::Float64 = -0.5 * log(2 * π * theta.R)
  @inline function lG(p::Int64, particle::Float64Particle, ::Nothing)
    @inbounds v::Float64 = theta.C*particle.x - ys[p]
    return logncG - v * invRover2 * v
  end
  @inline function M!(newParticle::Float64Particle, rng::RNG, p::Int64,
    particle::Float64Particle, ::Nothing)
    if p == 1
      newParticle.x = theta.x0 + sqrtv0*randn(rng)
    else
      newParticle.x = theta.A*particle.x + sqrtQ*randn(rng)
    end
  end
  return SMCModel(M!, lG, n, Float64Particle, Nothing)
end

function makelM(theta::LGTheta)
  tQ::Float64 = theta.Q
  negInvQover2::Float64 = -0.5/tQ
  tA::Float64 = theta.A
  @inline function lM(::Int64, particle::Float64Particle,
    newParticle::Float64Particle, ::Nothing)
    x::Float64 = particle.x
    y::Float64 = newParticle.x
    t::Float64 = tA*x - y
    return negInvQover2 * t * t
  end
  return lM
end

function simulateLGModel(theta::LGTheta, n::Int64)
  model = makeLGModel(theta, Vector{Float64}(undef, 0))
  ys = Vector{Float64}(undef, n)
  xParticle = Float64Particle()
  rng = getRNG()
  for p in 1:n
    model.M!(xParticle, rng, p, xParticle, nothing)
    ys[p] = theta.C*xParticle.x + sqrt(theta.R)*randn(rng)
  end
  return ys
end

function defaultLGModel(n::Int64)
  theta = LGTheta(0.9, 0.6, 1.2, 1.5, 1.0, 2.0)
  ys = simulateLGModel(theta, n)

  ko = kalman(theta, ys)

  lgModel = makeLGModel(theta, ys)

  return lgModel, theta, ys, ko
end

include("lgKalman.jl")
include("lgLOPModel.jl")
include("lgAPFModel.jl")

end


@testset "SGLD" begin
    #-----------------------
    # testing inplace addition
    #-----------------------
    function sample_reference!(s::SGLDState,grad::Function)
        # increment iteration
        s.t += 1

        # improve readability
        t = s.t
        x = s.x
        stepsize = s.stepsize(t)

        x[:] += stepsize*grad(x) + sqrt(2.0*stepsize)*randn(length(x))
        s
    end

    s = SGLDState(ones(2),0.01)
    ref_s = SGLDState(ones(2),0.01)
    sgldgrad(x) = x.*x

    srand(1)
    sample!(s,sgldgrad)

    srand(1)
    sample_reference!(ref_s,sgldgrad)

    @test s.x == ref_s.x

    #----------------------
    # Gaussian sanity check.
    #----------------------
    dm = GaussianMixtureModel([0.],[1.],[1.])

    gradient = DataModel.getgrad(dm)
    s = SGLDState(zeros(1),0.1)

    x_avg = [0.]
    x2_avg = [0.]
    for i in 1:50000
        sample!(s,gradient)
        x_avg =(i-1)/i*x_avg + s.x/i
        x2_avg = (i-1)/i*x2_avg + s.x.^2/i
    end

    @test abs(x_avg-0)[1] < 0.1
    @test abs(x2_avg-1)[1] < 0.2
end


using StatsBase

function stack(
    method,
    images::AbstractArray{<:AbstractArray},
    contrasts;
    pad=nothing
)

    # If lock is true (deafult) and the images have different sizes (and are 2D), and
    # padding has not been disabled, turn padding on.
    if length(images) > 1 && isnothing(pad)  && all(==(2), length.(size.(images))) && !all(Ref(size.(images)) .== size(first(images)))
        @warn "Padding images to a common size"
        pad = true
    else
        pad = false
    end

    # If pad is set, pad out the images with NaN to be the same size.
    if pad
        # Sort of an ungly line. paddedviews returns a tuple of views.
        # We need to convert to a vector of arrays.
        images = centered.([collect.(Images.paddedviews(NaN, images...))...])
    end

    contrast_maps = map(zip(images, contrasts)) do (img, cont)
        (1 ./ profile2image(cont, img)).^2
    end

    pixels = zeros(eltype(first(images)), length(images))
    weights= zeros(eltype(first(contrast_maps)), length(contrast_maps))
    mask= falses(length(contrast_maps))

    out = similar(first(images))
    for I in eachindex(first(images))
        for i in eachindex(images)
            pixels[i] = images[i][I]
            weights[i] = contrast_maps[i][I]
        end
        mask .= isfinite.(pixels) .& isfinite.(weights)
        if any(mask) && sum(view(weights,mask)) > 0
            out[I] = method(view(pixels,mask), aweights(view(weights,mask)))
        else
            out[I] = eltype(out)(NaN)
        end
    end
    return out
end

function stack(
    method,
    images::AbstractArray{<:AbstractArray};
    pad=nothing
)

    # If lock is true (deafult) and the images have different sizes (and are 2D), and
    # padding has not been disabled, turn padding on.
    if length(images) > 1 && isnothing(pad)  && all(==(2), length.(size.(images))) && !all(Ref(size.(images)) .== size(first(images)))
        @warn "Padding images to a common size."
        pad = true
    else
        pad = false
    end

    # If pad is set, pad out the images with NaN to be the same size.
    if pad
        # Sort of an ungly line. paddedviews returns a tuple of views.
        # We need to convert to a vector of arrays.
        images = centered.([collect.(Images.paddedviews(NaN, images...))...])
    end

    pixels = zeros(eltype(first(images)), length(images))
    mask= falses(length(images))

    out = similar(first(images))
    for I in eachindex(first(images))
        for i in eachindex(images)
            pixels[i] = images[i][I]
        end
        mask .= isfinite.(pixels)
        if any(mask)
            out[I] = method(view(pixels,mask))
        else
            out[I] = eltype(out)(NaN)
        end
    end
    return out
end
export stack


for N in [Float64, Rational{Int}, Float32]
    # compact set
    @test Ball1 <: CompactSet && !(Ball1 <: NonCompactSet)
    b = Ball1(N[0, 1], N(1))
    @test b isa CompactSet && !(b isa NonCompactSet)

    # non-compact set
    @test Line <: NonCompactSet && !(Line <: CompactSet)
    l = Line(N[0, 1], N(1))
    @test l isa NonCompactSet && !(l isa CompactSet)

    # lazy operations are neither compact nor non-compact by default
    @test !(MinkowskiSum <: CompactSet) && !(MinkowskiSum <: NonCompactSet)
    ms = b + l
    @test !(ms isa CompactSet) && !(ms isa NonCompactSet)
end


include("lsssolve.jl")
"""
Implements the NILSS algorithm due to Ni and Wang JCP 2017
Below, u -> F(u) is the time-1 primal map and u_trj refers
to a primal trajectory of length n.
Inputs:
	du_trj: dxdxn n-length timeseries of Jacobian matrices
			along u_trj
	X: dxn n-length timeseries of perturbation vectors.
			X[:,i] = dF/ds(u_trj[i,:]) which is in the
			tangent space of u_trj[i+1,:].
	f: dxn n-length timeseries of the neutral CLV. By
			assumption, the only neutral CLV is the RHS of
			the primal ode: du/dt = f(u,s).
			In case of maps, u_n = F(u_{n+1}), f is zeros.
	J:mxn n-length timeseries of m scalar functions J(u)
			evaluated along u_trj.
	dJ: mxdxn n-length timeseries of DJ(u) evaluated along
			u_trj
	s: array of parameters
	d_u: integer dimension of the unstable subspace. You can 		set a guess value to obtain correct sensitivities
		as long as the guess value is greater than or equal
		to the true d_u.
Outputs:
	vsh: dxn n-length timeseries of shadowing tangent
		vectors that are bounded solutions of the
		conventional tangent equation:
			v[:,i+1] = du_trj[:,i]*v[:,i] + X[:,i]

	dJds: scalar-valued sensitivity of interest:
		d/ds sum(J) ≈ d/ds <J, mu>, where mu is
		the SRB measure.
"""
function lss(du_trj, X, f, s, d_u)
	d, n = size(X)
	lyap_exps = zeros(d_u)
    R = zeros(d_u,d_u,n)
    Q = zeros(d,d_u,n)
	v = zeros(d,1,n) #assume one parameter
	v[:,1,1] = zeros(d)
	A = qr!(rand(d,d_u))
    Q[:,:,1] = Array(A.Q)
    R[:,:,1] = A.R

	b = zeros(d_u,1,n) #1 is the # parameters
	pf_Q = zeros(d_u,1,n) #1 is the # parameters
	pf_v = zeros(n)

	ff = zeros(n)
	if all(isapprox.(f,0.))
		ff = ones(n)
	else
		[ff[i] = sum(x -> x^2, f[:,i]) for i = 1:n]
	end
	pf_Q[:,1,1] = Q[:,:,1]'*f[:,1]
	[Q[:,j,1] = Q[:,j,1] - (pf_Q[j,1,1]*
	 		f[:,1]/ff[1]) for j=1:d_u]

	for i=2:n
	    v[:,:,i] = du_trj[:,:,i-1]*v[:,:,i-1] +
					X[:,i-1]
		Q[:,:,i] = du_trj[:,:,i-1]*Q[:,:,i-1]
		pf_Q[:,:,i] = Q[:,:,i]'*f[:,i]
		[Q[:,j,i] = Q[:,j,i] - pf_Q[j,1,i]*f[:,i]/
						ff[i] for j=1:d_u]
		A = qr!(Q[:,:,i])
        Q[:,:,i] = Array(A.Q)
        R[:,:,i] = A.R
		pf_v[i] = dot(v[:,1,i],f[:,i])
		v[:,:,i] = v[:,:,i] - (pf_v[i]/
							   ff[i]*f[:,i])
        b[:,:,i] = (Q[:,:,i]')*v[:,:,i]
		v[:,:,i] = v[:,:,i] - Q[:,:,i]*b[:,:,i]
		lyap_exps .+= log.(abs.(diag(R[:,:,i])))./n
    end
	#println(lyap_exps/0.01)

	b = reshape(collect(b),d_u, n)
	#println("Solving the least squares problem... ")
	a = lsssolve(R,b)
	#println("Computing shadowing direction...")
	v = reshape(collect(v),d,n)
	vsh = zeros(d, n)
	xi = zeros(n)
	vsh[:,1] = Q[:,:,1]*reshape(a[:,1], d_u, 1) +
				  v[:,1]
	for i = 2:n
		vsh[:,i] = v[:,i] + Q[:,:,i]*reshape(a[:,i],d_u,
					1)
		xi[i] = xi[i-1] + dot(-vsh[:,i] + du_trj[:,:,i-1]*
					vsh[:,i-1] + X[:,i-1], f[:,i])/ff[i]
	end
	return vsh, xi
end
function compute_sens(vsh, xi, dJ, f)
	#println("Computing sensitivity...")
	m, d, n = size(dJ)
	dJds = zeros(m)
	for i = 1:n
		dJds .= dJds .+ dJ[:,:,i]*(vsh[:,i] .+
					 xi[i]*f[:,i])/n
	end
	return dJds
end


using Flutes
using Test

@test round(Flutes.soundspeed(0.0); sigdigits=6) == 331.316
@test round(Flutes.soundspeed(25.0); sigdigits=6) == 346.146
@test Flutes.tone("A4") == 440.0
@test Flutes.tone("A5"; A=442.0) == 884.0
@test Flutes.tone("C♯0") == 17.323914
@test Flutes.tone("C♭8") == 3951.06641


using SpecialFunctions

function test()

##z = 1+2im
##w = erfcx(z)

ntest = 1000000
z = rand(ntest) + 1im*rand(ntest)
z = z * 10

println("n = ", ntest)

println("erf:")
@time w = erf.(z)
@time w = erf.(z)
@time w = erf.(z)
println("isfinite(w) = ", all(isfinite.(w)))

println("erfc:")
@time w = erfc.(z)
@time w = erfc.(z)
@time w = erfc.(z)
println("isfinite(w) = ", all(isfinite.(w)))

println("erfcx:")
@time w = erfcx.(z)
@time w = erfcx.(z)
@time w = erfcx.(z)
println("isfinite(w) =", all(isfinite.(w)))

println("exp:")
@time w = exp.(z)
@time w = exp.(z)
@time w = exp.(z)
println("isfinite(w) = ", all(isfinite.(w)))

println("done.")

end

test()
test()
test()


"""
   Kernels for GPs

Author: Eric Ford
Adapted from: https://github.com/eford/RvSpectraKitLearn.jl/blob/master/src/deriv_spectra_gp.jl

Could update to use KernelFunctions.jl, etc.
If upgrade to make use of other packages, please keep this so we ahve somet that gets the job done with minimal dependancies.
Author: Eric Ford
"""

function kernel_matern32(d::T; rho::T=one(T), sigmasq::T =one(T)) where { T<:Real }
  x = abs(sqrt(3)*d/rho)
  sigmasq * (1+x) * exp(-x)
end

function dkerneldx_matern32(d::T; rho::T=one(T), sigmasq::T =one(T)) where { T<:Real }
  x = abs(sqrt(3)*d/rho)
  -sigmasq * (3*d/rho^2) * exp(-x)
end

function d2kerneldx2_matern32(d::T; rho::T=one(T), sigmasq::T =one(T)) where { T<:Real }
  x = abs(sqrt(3)*d/rho)
  -sigmasq * (3/rho^2) * (1-x)*exp(-x)
end

function kernel_matern52(d::T; rho::T=one(T), sigmasq::T =one(T)) where { T<:Real }
  x = abs(sqrt(5)*d/rho)
  sigmasq * (1+x*(1+x/3)) * exp(-x)
end

function dkerneldx_matern52(d::T; rho::T=one(T), sigmasq::T =one(T)) where { T<:Real }
  x = abs(sqrt(5)*d/rho)
  -sigmasq/3 * (5*d/rho^2) *(1+x)*exp(-x)
end

function d2kerneldx2_matern52(d::T; rho::T=one(T), sigmasq::T =one(T)) where { T<:Real }
  x = abs(sqrt(5)*d/rho)
  -sigmasq/3 * (5/rho^2) * (1+x*(1-x))*exp(-x)
end

function nuttall_kernel(x::T; rho::T=one(T)) where { T<:Real }
      a0 = 0.355768
      a1 = 0.487396
      a2 = 0.144232
      a3 = 0.012604
   if abs(x)>rho return zero(T) end
   cpx = cospi(x/rho)
   spx = sinpi(x/rho)
   c2px = cpx*cpx-spx*spx
   return a0+a1*cpx+(a2+a3*cpx)*cpx*cpx-(a2+3*a3*cpx)*spx*spx
end

function nuttall_dkerneldx(x::T; rho::T=one(T)) where { T<:Real }
      a0 = 0.355768
      a1 = 0.487396
      a2 = 0.144232
      a3 = 0.012604
   if abs(x)>rho return zero(T) end
   cpx = cospi(x/rho)
   spx = sinpi(x/rho)
   c2px = cpx*cpx-spx*spx
   s2px = 2*cpx*spx
   s3px = spx*c2px+cpx*s2px
   return -pi/rho*(a1*spx+2*a2*s2px+3*a3*s3px)
end

function nuttall_d2kerneldx2(x::T; rho::T=one(T)) where { T<:Real }
      a0 = 0.355768
      a1 = 0.487396
      a2 = 0.144232
      a3 = 0.012604
   if abs(x)>rho return zero(T) end
   cpx = cospi(x/rho)
   spx = sinpi(x/rho)
   c2px = cpx*cpx-spx*spx
   c3px = cpx*(cpx*cpx-3*spx*spx)
   return -(pi/rho)^2*(a1*cpx+4*a2*c2px+9*a3*c3px)
end

function matern32_sparse_kernel(x::T; rho::T=one(T), bandwidth::T=4*one(T), sigmasq::T =one(T)) where { T<:Real }
  km = kernel_matern32(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*bandwidth)
  km*kn
end

function dkerneldx_matern32_sparse(x::T; rho::T=one(T), bandwidth::T=4*one(T), sigmasq::T =one(T)) where { T<:Real }
  km = kernel_matern32(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*bandwidth)
  dkmdx = dkerneldx_matern32(x,rho=rho,sigmasq=sigmasq)
  dkndx = nuttall_dkerneldx(x,rho=rho*bandwidth)
  km*dkndx+kn*dkmdx
end

function d2kerneldx2_matern32_sparse(x::T; rho::T=one(T), bandwidth::T=4*one(T), sigmasq::T =one(T)) where { T<:Real }
  km = kernel_matern32(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*bandwidth)
  dkmdx = dkerneldx_matern32(x,rho=rho,sigmasq=sigmasq)
  dkndx = nuttall_dkerneldx(x,rho=rho*bandwidth)
  d2kmdx2 = d2kerneldx2_matern32(x,rho=rho,sigmasq=sigmasq)
  d2kndx2 = nuttall_d2kerneldx2(x,rho=rho*bandwidth)
  km*d2kndx2+2*dkndx*dkmdx+kn*d2kmdx2
end

function matern52_sparse_kernel(x::T; rho::T=one(T), bandwidth::T=4*one(T), sigmasq::T =one(T)) where { T<:Real }
  km = kernel_matern52(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*bandwidth)
  km*kn
end

function dkerneldx_matern52_sparse(x::T; rho::T=one(T), bandwidth::T=4*one(T), sigmasq::T =one(T)) where { T<:Real }
  km = kernel_matern52(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*bandwidth)
  dkmdx = dkerneldx_matern52(x,rho=rho,sigmasq=sigmasq)
  dkndx = nuttall_dkerneldx(x,rho=rho*bandwidth)
  km*dkndx+kn*dkmdx
end

function d2kerneldx2_matern52_sparse(x::T; rho::T=one(T), bandwidth::T=4*one(T), sigmasq::T =one(T)) where { T<:Real }
  km = kernel_matern52(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*bandwidth)
  dkmdx = dkerneldx_matern52(x,rho=rho,sigmasq=sigmasq)
  dkndx = nuttall_dkerneldx(x,rho=rho*bandwidth)
  d2kmdx2 = d2kerneldx2_matern52(x,rho=rho,sigmasq=sigmasq)
  d2kndx2 = nuttall_d2kerneldx2(x,rho=rho*bandwidth)
  km*d2kndx2+2*dkndx*dkmdx+kn*d2kmdx2
end

#=
function make_kernel_matern32_data(x::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = zeros(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_data(x, kernel_matern32, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_matern32_obs_pred(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_obs_pred(xobs, xpred, kernel_matern32, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_dmatern32dx_obs_pred(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_obs_pred(xobs, xpred, dkerneldx_matern32, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_d2matern32dx2_obs_pred(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_obs_pred(xobs, xpred, d2kerneldx2_matern32, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_matern32_sparse_data(x::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = zeros(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_data(x, matern32_sparse_kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_matern32_sparse_obs_pred(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_obs_pred(xobs, xpred, matern32_sparse_kernel, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_dmatern32dx_sparse_obs_pred(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_obs_pred(xobs, xpred, dkerneldx_matern32_sparse, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_d2matern32dx2_sparse_obs_pred(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T=one(T))  where { T<:Real }
  make_kernel_obs_pred(xobs, xpred, d2kerneldx2_matern32_sparse, sigmasq_cor=sigmasq_cor, rho=rho)
end

=#


# Calculate structure factor
# special case: only for 1 species with Z = 1
function structure_factor( Xpos::Array{Float64,2}, G::Array{Float64,2} )
    Ng = size(G)[2]
    Na = size(Xpos)[2]
    Sf = zeros(ComplexF64,Ng)
    for ia = 1:Na
        for ig = 1:Ng
            GX = Xpos[1,ia]*G[1,ig] + Xpos[2,ia]*G[2,ig] + Xpos[3,ia]*G[3,ig]
            Sf[ig] = Sf[ig] + cos(GX) - im*sin(GX)
        end
    end
    return Sf
end


"""
makes a soft analogue of the heaviside step function. useful for inputs to differential equations, as it's easier on the numerics.
"""
soft_heaviside(t, nastiness, step_time) = 1 / (1 + exp(nastiness * (step_time - t)))
soft_heaviside(nastiness, step_time) = t -> soft_heaviside(t, nastiness, step_time)

get_ids_names(opArray) = repr.(opArray)

"""
    l2_hessian(nom_sol)
gets hessian according to L2 loss under the assumption that loss(θ₀) = 0. nom_sol is the solution of the nominal ODEProblem.
The Hessian then only requires first derivatives: it is sum_ij dyi/dθ * dyj/dtheta
"""
function l2_hessian(nom_sol)
    prob = nom_sol.prob
    function pToL2(p)
        pprob = remake(prob, p=p)
        psol = solve(pprob, nom_sol.alg, saveat=nom_sol.t) |> Array
        psol = reshape(psol, 1, :)
        return psol
    end

    gr = ForwardDiff.jacobian(pToL2, prob.p)
    u, d, v = svd(gr)
    return v * diagm(d.^2) * v'
    # = gr'*gr but a bit more accurate
end


using Flux, Flux.Data.MNIST
using Flux: @epochs, onehotbatch, mse, throttle
using Base.Iterators: partition
using Juno: @progress
# using CuArrays

# Encode MNIST images as compressed vectors that can later be decoded back into
# images.

imgs = MNIST.images()

# Partition into batches of size 1000
data = [float(hcat(vec.(imgs)...)) for imgs in partition(imgs, 1000)]
data = gpu.(data)

N = 32 # Size of the encoding

# You can try to make the encoder/decoder network larger
# Also, the output of encoder is a coding of the given input.
# In this case, the input dimension is 28^2 and the output dimension of
# encoder is 32. This implies that the coding is a compressed representation.
# We can make lossy compression via this `encoder`.
encoder = Dense(28^2, N, leakyrelu) |> gpu
decoder = Dense(N, 28^2, leakyrelu) |> gpu

m = Chain(encoder, decoder)

loss(x) = mse(m(x), x)

evalcb = throttle(() -> @show(loss(data[1])), 5)
opt = ADAM(params(m))

@epochs 10 Flux.train!(loss, zip(data), opt, cb = evalcb)

# Sample output

using Images

img(x::Vector) = Gray.(reshape(clamp.(x, 0, 1), 28, 28))

function sample()
  # 20 random digits
  before = [imgs[i] for i in rand(1:length(imgs), 20)]
  # Before and after images
  after = img.(map(x -> cpu(m)(float(vec(x))).data, before))
  # Stack them all together
  hcat(vcat.(before, after)...)
end

cd(@__DIR__)

save("sample.png", sample())


using ComputerVision
using Test
using StaticArrays
using GeometryTypes
using LinearAlgebra
#using Makie # You have to add this as a dependency to your environment.

world = PrimitiveWorld()
@inferred PrimitiveWorld()

points = get_points(world)

planes = get_planes(world)
plane = first(planes)
𝐧 = get_normal(plane)
d = get_distance(plane)

# Verify that the default points lie on the plane
𝛑 = push(𝐧, -d) # 𝛑 =[n -d]
for 𝐱 in points
    @test isapprox(dot(𝛑, hom(𝐱)), 0.0; atol = 1e-14)
end

camera₁ = DigitalCamera()
camera₂ = DigitalCamera()

#cameras = get_cameras(world)
#camera₁ = cameras[1]
#camera₂ = cameras[2]


function Predict(θ;μ=0,σ=1)
   P=function (X)
      m=size(X,1)
      X=(X.-μ)./σ
      X=[ones(m) X]
      p=X*θ
        return p
    end

        return P
end


Base.exponent(G::GroupsCore.Group) = exponent(conjugacy_classes(G))
Base.exponent(cclasses::AbstractArray) = lcm(order.(first.(cclasses)))
dixon_prime(G::GroupsCore.Group) = dixon_prime(order(G), exponent(G))

function dixon_prime(cclasses::AbstractArray)
    ordG = sum(length, cclasses)
    m = exponent(cclasses)
    return dixon_prime(ordG, m)
end

function dixon_prime(ordG::Integer, exponent::Integer)
    # `FiniteFields.GF{p}` needs `p` to be `Int` so no need to do the
    # computations of this function over `BigInt` if `ordG` is so we convert
    # to `Int`.
    p = 2 * convert(Int, isqrt(ordG))
    while true
        p = nextprime(p + 1)
        isone(p % exponent) && break # we need -1 to be in the field
    end
    return p
end

function common_esd(Ns, F::Type{FiniteFields.GF{q}}) where q
    @assert isprime(q)
    itr = iterate(Ns)
    @assert itr !== nothing
    esd = EigenSpaceDecomposition(F.(first(itr)))
    for N in Iterators.rest(Ns, last(itr))
        esd = refine(esd, F.(N, false))
        @debug N esd.eigspace_ptrs
        isdiag(esd) && return esd
    end
    return esd
end

function _multiplicities(
    chars::AbstractVector{<:Character{F}},
    cclasses = conjugacy_classes(first(chars)),
) where {F<:FiniteFields.GF}

    e = Int(exponent(cclasses))
    ie = inv(F(e))

    ω⁻¹ = inv(FiniteFields.rootofunity(F, e))
    ωs = Matrix{typeof(ω⁻¹)}(undef, e, e)
    Threads.@threads for k in 0:e-1
        ω⁻ᵏ = ω⁻¹^k
        for l in 0:e-1
            ωs[l+1, k+1] = ω⁻ᵏ^l
        end
    end

    multiplicities = zeros(Int, length(chars), length(cclasses), e)
    pmap = powermap(table(first(chars)))
    # pmap is a lazy object, computed on demand; to ensure its thread-safety
    # we precompute all its values here so in the loops below it is accessed read-only
    collect(pmap)

    for (i, χ) in enumerate(chars)
        Threads.@threads for j = 1:length(cclasses)
            for k = 0:e-1
                val = Int(ie * sum(χ[pmap[j, l]] * ωs[l+1, k+1] for l = 0:e-1))
                multiplicities[i, j, k+1] = val
            end
        end
    end

    return multiplicities
end


@testset "maximumlikelihoodeva.jl" begin

    n = 1000

    x₁ = Variable("x₁", randn(n)/3)
    x₂ = Variable("x₂", randn(n))
    x₃ = Variable("x₂", randn(n))

    θ = [1; 1; 1; -.5; 1; .1]

    μ = θ[1] .+ x₂.value*θ[2] .+ x₃.value*θ[3]
    ϕ = θ[4] .+ x₁.value*θ[5]
    ξ = θ[6]

    pd = GeneralizedExtremeValue.(μ, exp.(ϕ), ξ)

    y = rand.(pd)

    model = BlockMaxima(Variable("y", y), locationcov=[x₂; x₃], logscalecov = [x₁])

    fm = Extremes.MaximumLikelihoodEVA(model, θ)

    @testset "getdistribution(fittedmodel)" begin
        @test all(Extremes.getdistribution(fm) .== pd)
    end

    @testset "loglike(fd)" begin
        # Test with known values
        @test Extremes.loglike(fm) ≈ sum(logpdf.(pd,y))
    end

    @testset "quantile(fm, p)" begin
        # p not in [0, 1] throws
        @test_throws AssertionError Extremes.quantile(fm, -1)

        # Test with known values
        @test quantile(fm, .99) ≈ quantile.(pd,.99)
    end

    @testset "returnlevel(fm, returnPeriod)" begin
        # returnPeriod < 0 throws
        @test_throws AssertionError Extremes.returnlevel(fm, -1)

        # Test with known values
        @test quantile(fm, .99) ≈ quantile.(pd,.99)
    end

    @testset "hessian(fm)" begin
        # Tested in cint(fm)
    end

    @testset "parametervar(fm)" begin
        # Tested in cint(fm)
    end

    @testset "cint(fm)" begin

        @test_throws AssertionError cint(fm, 1)

        # Test with known values
        # Evaluating the confidence interval empirical level using a Monte Carlo approach.
        # Use of the hessian(fm) and parametervar function.

        niter = 1000

        cover = falses(niter,6)

        Threads.@threads  for i in 1:niter
            y = rand.(pd)
            m = BlockMaxima(Variable("y",y), locationcov = [x₂, x₃], logscalecov=[x₁])
            fm = gevfit(m, θ)
            ci = cint(fm)
            for j in eachindex(θ)
                cover[i,j] = ci[j][1] < θ[j] < ci[j][2]
            end
        end

        @test all(count.(eachcol(cover))/niter .> .9)

    end

    @testset "quantilevar(rl, level)" begin
        # Tested in cint(rl)
    end

    @testset "cint(rl, confidencelevel)" begin
        # confidencelevel not in [0, 1] throws
        @test_throws AssertionError cint(fm, 1)

        # Test with known values

        θ = [100; log(5); .1]

        pd = GeneralizedExtremeValue(θ[1],exp(θ[2]),θ[3])

        q = quantile(pd, 1-1/100)

        niter = 1000

        cover = falses(niter)

        Threads.@threads  for i in 1:niter
            y = rand(pd, 300)
            m = BlockMaxima(Variable("y",y))
            fm = gevfit(m, θ)
            rl = returnlevel(fm, 100)
            ci = cint(rl)
            cover[i] = ci[][1] < q < ci[][2]
        end

        @test count(cover)/niter .> .9

    end

    n = 1000
    nobservation = 10000
    nobsperblock = 1

    x = Variable("x",randn(n))

    θ = [-.5; 1; 0.1]

    ϕ = θ[1] .+ θ[2]*x.value
    σ = exp.(ϕ)
    ξ = θ[3]

    pd = GeneralizedPareto.(σ, ξ)

    y = rand.(pd)

    model = ThresholdExceedance(Variable("y", y), logscalecov = [x])

    fm = Extremes.MaximumLikelihoodEVA(model, θ)

    @testset "getdistribution(fittedmodel)" begin
        @test all(Extremes.getdistribution(fm) .== pd)
    end

    @testset "loglike(fd)" begin
        # Test with known values
        @test Extremes.loglike(fm) ≈ sum(logpdf.(pd,y))
    end

    @testset "quantile(fm, p)" begin
        # p not in [0, 1] throws
        @test_throws AssertionError Extremes.quantile(fm, -1)

        # Test with known values
        @test quantile(fm, .99) ≈ quantile.(pd,.99)
    end

    @testset "returnlevel(fm, returnPeriod)" begin
        # returnPeriod < 0 throws
        @test_throws AssertionError returnlevel(fm, 0, nobservation, nobsperblock, -100)

        # Test with known values
        rl = returnlevel(fm, 0, nobservation, nobsperblock, 100)
        p = 1-nobservation/(100 * nobsperblock * n)
        @test rl.value ≈ quantile.(pd, p)
    end

    @testset "hessian(fm)" begin
        # Tested in cint(fm)
    end

    @testset "parametervar(fm)" begin
        # Tested in cint(fm)
    end

    @testset "cint(fm)" begin

        @test_throws AssertionError cint(fm, 1)

        # Test with known values
        # Evaluating the confidence interval empirical level using a Monte Carlo approach.
        # Use of the hessian(fm) and parametervar function.

        niter = 1000

        cover = falses(niter,length(θ))

        Threads.@threads  for i in 1:niter
            y = rand.(pd)
            m = ThresholdExceedance(Variable("y",y), logscalecov=[x])
            fm = gpfit(m, θ)
            ci = cint(fm)
            for j in eachindex(θ)
                cover[i,j] = ci[j][1] < θ[j] < ci[j][2]
            end
        end

        @test all(count.(eachcol(cover))/niter .> .9)

    end

    @testset "quantilevar(rl, level)" begin
        # Tested in cint(rl)
    end

    @testset "cint(rl, confidencelevel)" begin
        # confidencelevel not in [0, 1] throws
        @test_throws AssertionError cint(fm, 1)

        # Test with known values

        θ = [log(5); .1]

        pd = GeneralizedPareto(exp(θ[1]),θ[2])

        n = 100
        nobservation = 1000
        nobsperblock = 1

        p = 1-nobservation/(100 * nobsperblock * n)

        q = quantile(pd, p)

        niter = 1000

        cover = falses(niter)

        Threads.@threads  for i in 1:niter
            y = rand(pd, n)
            m = ThresholdExceedance(Variable("y",y))
            fm = gpfit(m, θ)
            rl = returnlevel(fm, 0, nobservation, nobsperblock, 100)
            ci = cint(rl)
            cover[i] = ci[][1] < q < ci[][2]
        end

        @test count(cover)/niter .> .9

    end

    @testset "showfittedEVA(io, obj, prefix)" begin
        # print does not throw
        buffer = IOBuffer()
        @test_logs Extremes.showfittedEVA(buffer, fm, prefix = "\t")
    end

end


using TrialSampleSize

# Calculate the power for one-sample proportion test
#=
True proportion: 0.9
Proportion under the null: 0.5
Std: 1
Sample Size: 50
=#
# Calculate power
power(OneSampleProp(0.9, 0.6), n = 50)
# Calculate sample size
sample_size(OneSampleProp(0.9, 0.6), power = 0.8, alpha = 0.05)


using AbstractGPs
using ConjugateComputationVI
using Distributions
using Optim
using ParameterHandling
using Plots
using Random
using RDatasets
using StatsFuns
using Zygote

using ConjugateComputationVI:
    batch_quadrature,
    GaussHermiteQuadrature,
    UnivariateFactorisedLikelihood

# Prepare data.

# Get the (infamous) Iris dataset, and only keep two labels -- we'll just do binary
# classification for now.
iris = filter(
    row -> row.Species == "setosa" || row.Species == "virginica",
    dataset("datasets", "iris"),
)

setosa = filter(row -> row.Species == "setosa", iris)
virginica = filter(row -> row.Species == "virginica", iris)

scatter(setosa.SepalLength, setosa.SepalWidth; label="Setosa")
scatter!(virginica.SepalLength, virginica.SepalWidth; label="Virginica")

# Check that there are only two features.
@assert length(unique(iris.Species)) == 2

# Encode the specifies using Ints. Could also have used Booleans.
y = map(x -> x == "setosa" ? 1 : 0, iris.Species)

# Construct input matrix from feature columns.
X = hcat(iris.SepalLength, iris.SepalWidth);
x = ColVecs(collect(X'));

# Split into train and test sets.
train_indices = randperm(MersenneTwister(123456), length(y))[1:50];
test_indices = setdiff(eachindex(y), train_indices);
x_tr = x[train_indices];
x_te = x[test_indices];
y_tr = y[train_indices];
y_te = y[test_indices];

# Specify and perform inference in model.

θ_init = (scale=fixed(1.0), stretch=positive.(ones(2)));

θ_init_flat, unflatten = ParameterHandling.flatten(θ_init);

function build_latent_gp(θ::AbstractVector{<:Real})
    return build_latent_gp(ParameterHandling.value(unflatten(θ)))
end
function build_latent_gp(θ::NamedTuple)
    gp = GP(θ.scale * SEKernel() ∘ AbstractGPs.ScaleTransform(θ.stretch))
    lik = UnivariateFactorisedLikelihood(f -> Bernoulli(logistic(f)))
    return LatentGP(gp, lik, 1e-9)
end

f_approx_post, results_summary = ConjugateComputationVI.optimize_elbo(
    build_latent_gp,
    GaussHermiteQuadrature(10),
    x_tr,
    y_tr,
    θ_init_flat,
    BFGS(
        alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
        linesearch = Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(
        show_trace = true,
        iterations=25,
        f_calls_limit=50,
    ),
);

# Visualise results of model.

function post_mean_prediction(x::AbstractVector)
    m, v = mean_and_var(f_approx_post(x))
    conditionals = fill(f -> pdf(Bernoulli(logistic(f)), 1), eachindex(m))
    return batch_quadrature(conditionals, m, sqrt.(v), 25)
end

# Compute predictions using the optimal hyperparameters.
function MAP_predict(x::AbstractVector)
    posterior_mean_predictions = post_mean_prediction(x)
    MAP_predictions = map(x -> x > 0.5 ? 1 : 0, posterior_mean_predictions)
    return MAP_predictions
end

# Compute training set accuracy as a %age.
100 * mean(MAP_predict(x_tr) .== y_tr)

# Compute test set accuracy as a %age.
100 * mean(MAP_predict(x_te) .== y_te)

# Compute prediction surface.
x1_range = range(minimum(X[:, 1]) - 3, maximum(X[:, 1]) + 3; length=100);
x2_range = range(minimum(X[:, 2]) - 3, maximum(X[:, 2]) + 3; length=100);
x_pr_vecs = [[x1, x2] for x1 in x1_range, x2 in x2_range];
X_pr = reduce(hcat, x_pr_vecs);
x_pr = ColVecs(X_pr);

# Plot classification.
MAP_pr = MAP_predict(x_pr);
p1 = heatmap(
    x1_range, x2_range, reshape(MAP_pr, length(x2_range), length(x1_range))';
    label="classifier",
);
scatter!(p1, setosa.SepalLength, setosa.SepalWidth; label="");
scatter!(p1, virginica.SepalLength, virginica.SepalWidth; label="");

# Plot posterior mean prob.
post_mean_pr = post_mean_prediction(x_pr);
p2 = heatmap(
    x1_range, x2_range, reshape(post_mean_pr, length(x2_range), length(x1_range))';
    label="prob",
);
scatter!(p2, setosa.SepalLength, setosa.SepalWidth; label="");
scatter!(p2, virginica.SepalLength, virginica.SepalWidth; label="");

# Plot posterior marginals.
ms = marginals(f_approx_post(x_pr));
p3 = heatmap(
    x1_range, x2_range, reshape(mean.(ms), length(x2_range), length(x1_range))';
    label="latent mean",
);
scatter!(p3, setosa.SepalLength, setosa.SepalWidth; label="");
scatter!(p3, virginica.SepalLength, virginica.SepalWidth; label="");

p4 = heatmap(
    x1_range, x2_range, reshape(std.(ms), length(x2_range), length(x1_range))';
    label="latent std",
);
scatter!(p4, setosa.SepalLength, setosa.SepalWidth; label="");
scatter!(p4, virginica.SepalLength, virginica.SepalWidth; label="");

plot(p1, p2, p3, p4; layout=(2, 2))


function dq_ri(δ)
    ## Uses the referenceframe of the Kundur page 852 of dq to RI
    return [
        sin(δ) cos(δ)
        -cos(δ) sin(δ)
    ]
end

function ri_dq(δ)
    #Uses the reference frame of the Kundur page 852 of RI to dq
    return [
        sin(δ) -cos(δ)
        cos(δ) sin(δ)
    ]
end


using DiffEqBase, OrdinaryDiffEq, Sundials, ApproxFun
using Base.Test

S=Fourier()
u0=Fun(θ->cos(cos(θ-0.1)),S)
c=Fun(cos,S)

prob = SpectralTimeSteppingProblem((t,u)->u''+(c+1)*u',u0,(0.,1.))
@time u=solve(prob,Tsit5());
@time u=solve(prob,CVODE_BDF());
@time u=solve(prob,ode45());
@time u=solve(prob,radau());


using MeasureTheory, Static, BenchmarkTools, LinearAlgebra, StaticArrays
using MultivariateMeasures
x = randn(200,200);
Σ = x' * x;

using PositiveFactorizations
CΣ = cholesky(Positive, Σ)

function benchmark()
    σtimes = Float64[]
    ωtimes = Float64[]
    y = randn(200);
    for k in 1:100
        U = UpperTriangular(SizedMatrix{k,k}(CΣ.U.data[1:k,1:k]));
        C = Cholesky(U, 'U', 0);
        yk = y[1:k]
        push!(σtimes, @belapsed logdensity(MvNormal(σ=$C), $yk))
        push!(ωtimes, @belapsed logdensity(MvNormal(ω=$C), $yk))
    end

    (σtimes, ωtimes)
end

(σtimes, ωtimes) = benchmark()

# using PDMats
# C = cholesky(Σ[1:20,1:20])
# Dists.MvNormal(PDMat(C))

# @btime Dists.logpdf(Dists.MvNormalCanon(PDMat($C)), $y)
# @btime Dists.logpdf(Dists.MvNormal(PDMat($C)), $y)

using Plots

plot(σtimes[1:end] .* 1e6; dpi=300, legend=150, label="σ parameterization")
plot!(ωtimes[1:end] .* 1e6, label="ω parameterization")
xlabel!("Dimensions")
ylabel!("Time (μs)")
title!("MeasureTheory.jl Log-density for MvNormal")
savefig("mvnormal.svg")
savefig("mvnormal.png")

# C_unsized = Cholesky(C.UL.data.data, 'U', )

# using PDMats
# PD = PDMat(C_unsized)

# julia> @btime Dists.logpdf(Dists.MvNormalCanon($PD), $y)
#   4.944 μs (6 allocations: 1.97 KiB)
# -124.71152557023147

# julia> @btime Dists.logpdf(Dists.MvNormalCanon(PDMat($C_unsized)), $y)
#   16.552 μs (10 allocations: 27.25 KiB)
# -124.71152557023147

# julia> @btime logdensity(MvNormal(ω=$C), $y)
#   214.948 ns (0 allocations: 0 bytes)
# -124.53160601776467

# L = inv(C.L)
# Σ = L' * L

# @btime Dists.logpdf(Dists.MvNormalCanon($PD), $y)
# @btime Dists.logpdf(Dists.MvNormalCanon(PDMat($C_unsized)), $y)

Dists.logpdf(Dists.MvNormal(Σ), y)
logdensity(MvNormal(ω=C), y)

# n = 40;
# t = CorrCholesky(StaticInt(n));
# C = transform(t, randn(dimension(t)));
@btime MeasureTheory.logdet_pos($C)

# @btime logdensity(MvNormal(ω=$C), $(randn(n)))


function calculate_overlap(
    dataset::DataSet{T, I},
    volumes::Array{IntegrationVolume{T, I, V}, 1},
    integralestimates::IntermediateResults)::Array{T, 2} where{T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    M = length(integralestimates)

    overlap = zeros(T, M, M)

    sortedsets = SortedSet.([volumes[integralestimates.volumeID[i]].pointcloud.pointIDs for i = 1:M])

    idx = workpart(1:M, ParallelProcessingTools.workers(), ParallelProcessingTools.myid())
    @mt for i in workpart(idx, mt_nthreads(), threadid())
        for j = 1:M
            intersectpts = intersect(sortedsets[i], sortedsets[j])
            unionpts = union(sortedsets[i], sortedsets[j])

            overlap[i, j] = sum(dataset.weights[collect(intersectpts)]) / sum(dataset.weights[collect(unionpts)])
        end
    end
    overlap
end

function pdf_gauss(Z::Float64, μ::Float64,  σ_sq::Float64)::Float64
    return 1.0 / sqrt(2 * pi * σ_sq) * exp(-0.5 * (Z - μ)^2 / σ_sq)
end

function binomial_p_estimate_wald(n_total, n_success, nsigmas = 1)
    p_hat = n_success / n_total
    @assert p_hat >= 0 && p_hat <= 1
    @assert n_total > 0 && n_success > 0
    p_err = nsigmas * sqrt(p_hat * (1 - p_hat) / n_total)
end

function binomial_p_estimate_wilson(n_total, n_success, nsigmas = 1)
    n = n_total
    n_S = n_success
    n_F = n_total - n_S
    z = nsigmas
    z2 = z^2

    p_val = (n_S + z2/2) / (n + z2)

    p_err = z / (n + z2) * sqrt(n_S * n_F / n + z2 / 4)
end

function calculateuncertainty(dataset::DataSet{T, I}, volume::IntegrationVolume{T, I}, determinant::T, integral::T) where {T<:AbstractFloat, I<:Integer}

    #resorting the samples to undo the reordering of the space partitioning tree
    reorderscheme = sortperm(dataset.sortids[volume.pointcloud.pointIDs], rev = true)
    original_ordered_sample_ids = volume.pointcloud.pointIDs[reorderscheme]

    vol_samples = dataset.data[:, original_ordered_sample_ids]
    vol_weights = dataset.weights[original_ordered_sample_ids]

    x_min = 1.0 / dataset.N
    x_max = 1.0

    vol_weight = sum(vol_weights)
    total_weight = sum(dataset.weights)
    r = vol_weight / total_weight
    y_r = (x::Float64) -> Float64(x_min + (x_max - x_min) * x)

    ess = BAT.effective_sample_size(vol_samples, vol_weights)
    ess_total = BAT.effective_sample_size(dataset.data[:, dataset.sortids], dataset.weights[dataset.sortids])

    f = 1 ./ exp.(dataset.logprob[volume.pointcloud.pointIDs])

    μ_Z = mean(f)
    @assert ess>0

    σ_μZ_sq = var(f) / ess

    f_max = maximum(f)
    f_min = minimum(f)

    y = (x::Float64) -> Float64(f_min + (f_max - f_min) * x)

    integrand1(x) = pdf_gauss(y(x), μ_Z, σ_μZ_sq) / y(x)^2 * (f_max - f_min)
    integrand2(x) = pdf_gauss(y(x), μ_Z, σ_μZ_sq) / y(x) * (f_max - f_min)

    integral1 = QuadGK.quadgk(integrand1,0,1, rtol=0.001)[1]
    integral2 = QuadGK.quadgk(integrand2,0,1, rtol = 0.001)[1]

    uncertainty_Z = (integral1 - integral2^2) * volume.volume / r / determinant
    #set uncertainty to 0 if it is negative (numerical error)
    if uncertainty_Z < 0
        uncertainty_Z = 0
    else
        uncertainty_Z = sqrt(uncertainty_Z)
    end

    scaling_factor = ess / vol_weight
    scaling_factor_total = ess_total / total_weight

    uncertainty_r = binomial_p_estimate_wald(scaling_factor_total * total_weight, scaling_factor * vol_weight)
    uncertainty_r *= integral / r #error propagation

    uncertainty_tot = sqrt(uncertainty_Z^2 + uncertainty_r^2)

    return uncertainty_r, uncertainty_Z, uncertainty_tot, f_min, f_max, μ_Z, σ_μZ_sq, integral1, integral2, ess
end


export
    GaussianSuffStats,
    DiagGaussianSuffStats,
    suffstats

immutable GaussianSuffStats <: SuffStats
    ss::Tuple{Vector{Float}, Symmetric{Float, Matrix{Float}}}
    function GaussianSuffStats(s1::Vector{Float}, s2::Matrix{Float})
        @assert length(s1)==size(s2)[1]==size(s2)[2] "Inconsistent dimensions"
        new((s1, Symmetric(s2)))
    end
end

immutable DiagGaussianSuffStats <: SuffStats
    ss::Tuple{Vector{Float}, Vector{Float}}
    function DiagGaussianSuffStats(s1::Vector{Float}, s2::Vector{Float})
        @assert length(s1)==length(s2) "Inconsistent dimensions"
        new((s1, s2))
    end
end

const FGaussSS = GaussianSuffStats
const DGaussSS = DiagGaussianSuffStats
const GaussSS  = Union{FGaussSS, DGaussSS}

suffstats(::Type{GaussNP}, x::Vector{Float})=GaussianSuffStats(x, 0.5x*x')
suffstats(::Type{GaussMP}, x::Vector{Float})=GaussianSuffStats(x, 0.5x*x')
suffstats(::Type{DGaussNP},x::Vector{Float})=DiagGaussianSuffStats(x, 0.5x.^2)
suffstats(::Type{DGaussMP},x::Vector{Float})=DiagGaussianSuffStats(x, 0.5x.^2)

Base.getindex(s::GaussSS, ind::Int) = s.ss[ind]

Base.vec(s::GaussSS)  = vcat(s[1], vec(s[2]))
Base.norm(s::GaussSS) = norm(vec(s))

# The operation that matters the most is the sum and the product for a weighted
# sum (monte carlo estimator)

+(sa::FGaussSS, sb::FGaussSS) =
        GaussianSuffStats(sa[1]+sb[1], sa[2]+sb[2])
+(sa::DGaussSS, sb::DGaussSS) =
        DiagGaussianSuffStats(sa[1]+sb[1], sa[2]+sb[2])

*(a::Real, s::FGaussSS) = GaussianSuffStats(a*s[1], a*s[2].data)
*(a::Real, s::DGaussSS) = DiagGaussianSuffStats(a*s[1], a*s[2])

*(s::GaussSS, a::Real) = a*s
/(s::GaussSS, a::Real) = (1.0/a)*s

# Minus may be useful for

-(sa::FGaussSS, sb::FGaussSS) =
        GaussianSuffStats(sa[1]-sb[1], sa[2]-sb[2])
-(sa::DGaussSS, sb::DGaussSS) =
        DiagGaussianSuffStats(sa[1]-sb[1], sa[2]-sb[2])


using Test
using RenyiEntropyOfMagic
using OMEinsum

function naive_two_site(A1, A2, O1, O2)
    psi = vec(ein"ji,jk->ik"(A1, A2))
    return psi' * reshape(ein"ik,jl->ijkl"(O1, O2), 4, 4) * psi
end

function naive_three_site(A1, A2, A3, O1, O2, O3)
    psi = vec(ein"ai,abj,bk->ijk"(A1,A2,A3))
    O = reshape(ein"il,jm,kn->ijklmn"(O1,O2,O3), length(psi), length(psi))
    return psi' * O * psi
end

A1, A2 = rand(3, 2), rand(3, 2)
O1, O2 = rand(2, 2), rand(2, 2)

@test naive_two_site(A1, A2, O1, O2) ≈ boundary_vector(A1, O1)' * boundary_vector(A2, O2)

A1, A2, A3 = rand(3, 2), rand(3, 3, 2), rand(3, 2)
O1, O2, O3 = rand(2, 2), rand(2, 2), rand(2, 2)

result = boundary_vector(A1, O1)' * transfer_matrix(A2, O2) * boundary_vector(A3, O3)
target = naive_three_site(A1, A2, A3, O1, O2, O3)
@test result ≈ target

function enumerate_paulistring(L::Int)
    labels_a = Iterators.product([0:1 for _ in 1:L]...)
    labels_b = Iterators.product([0:1 for _ in 1:L]...)

    return map(Iterators.product(labels_a, labels_b)) do (as, bs)
        map(zip(as, bs)) do (a, b)
            if a == 0 && b == 0
                Consts.I2
            elseif a == 0 && b == 1
                Consts.Z
            elseif a == 1 && b == 0
                Consts.X
            else # 11
                Consts.X * Consts.Z
            end
        end
    end |> vec
end

As = rand(3, 2), rand(3, 3, 2), rand(3, 2)
sum(expectation(As, Os)^4 for Os in enumerate_paulistring(3))

using Statistics

mean(1:10^6) do _
    Os = sample_paulistring(3)
    expectation(As, Os)^4
end * 4^3
renyi_entropy(As)

exact_renyi_entropy(As, Os)
A = rand(2, 4, 2)
O = rand(2, 2)
transfer_matrix(A, O)