# Rust (1987) bus engine replacement problem
# Notation: I use underscore to denote empirical counterparts

# Import packages
using Optim
using Distributions
using Statistics
using DataFrames
using CSV
using LinearAlgebra

function compute_U(θ::Vector, s::Vector)::Matrix
    """Compute static utility"""
    u1 = - θ[1]*s - θ[2]*s.^2       # Utility of not investing
    u2 = - θ[3]*ones(size(s))       # Utility of investing
    U = [u1 u2]                     # Combine in a matrix
    return U
end;

function compute_Vbar(θ::Vector, λ::Number, β::Number, s::Vector)::Matrix
    """Compute value function by Bellman iteration"""
    k = length(s)                                 # Dimension of the state space
    U = compute_U(θ, s)                           # Static utility
    index_λ = Int[1:k [2:k; k]];                  # Mileage index
    index_A = Int[1:k ones(k,1)];                 # Investment index
    γ = Base.MathConstants.eulergamma             # Euler's gamma

    # Iterate the Bellman equation until convergence
    Vbar = zeros(k, 2);
    Vbar1 = Vbar;
    dist = 1;
    iter = 0;
    while dist>1e-8
        V = γ .+ log.(sum(exp.(Vbar), dims=2))     # Compute value
        expV = V[index_λ] * [1-λ; λ]               # Compute expected value
        Vbar1 =  U + β * expV[index_A]             # Compute v-specific
        dist = max(abs.(Vbar1 - Vbar)...);         # Check distance
        iter += 1;
        Vbar = Vbar1                               # Update value function
    end
    return Vbar
end;

function generate_data(θ::Vector, λ::Number, β::Number, s::Vector, N::Int)::Tuple
    """Generate data from primitAes"""
    Vbar = compute_Vbar(θ, λ, β, s)             # Solve model
    ε = rand(Gumbel(0,1), N, 2)                 # Draw shocks
    St = rand(s, N)                             # Draw states
    A = (((Vbar[St,:] + ε) * [-1;1]) .> 0)     # Compute investment decisions
    δ = (rand(Uniform(0,1), N) .< λ)            # Compute mileage shock
    St1 = min.(St .* (A.==0) + δ, max(s...))   # Compute nest state
    df = DataFrame(St=St, A=A, St1=St1)
    CSV.write("../data/rust.csv", df)
    return St, A, St1
end;

function logL_Rust(θ0::Vector, λ::Number, β::Number, s::Vector, St::Vector, A::BitVector)::Number
    """Compute log-likelihood functionfor Rust problem"""
    # Compute value
    Vbar = compute_Vbar(θ0, λ_, β, s)

    # Implied choice probabilities
    EP = exp.(Vbar[:,2]) ./ (exp.(Vbar[:,1]) + exp.(Vbar[:,2]))

    # Likelihood
    logL = sum(log.(EP[St[A.==1]])) + sum(log.(1 .- EP[St[A.==0]]))
    return -logL
end;

function compute_T(k::Int, λ_::Number)::Array
    """Compute transition matrix"""
    T = zeros(k, k, 2);

    # Conditional on not investing
    T[k,k,1] = 1;
    for i=1:k-1
        T[i,i,1] = 1-λ_
        T[i,i+1,1] = λ_
    end

    # Conditional on investing
    T[:,1,2] .= 1-λ_;
    T[:,2,2] .= λ_;

    return(T)
end;

function HM_inversion(CCP::Matrix, T::Array, U::Matrix, β::Number)::Vector
    """Perform HM inversion"""

    # Compute LHS (to be inverted)
    γ = Base.MathConstants.eulergamma
    LEFT = I - β .* (CCP[:,1] .* T[:,:,1] + CCP[:,2] .* T[:,:,2])

    # Compute LHS (not to be inverted)
    RIGHT = γ .+ sum(CCP .* (U .- log.(CCP)) , dims=2)

    # Compute V
    EV_ = inv(LEFT) * RIGHT
    return vec(EV_)
end;

function from_EV_to_EP(EV_::Vector, T::Array, U::Matrix, β::Number)::Vector
    """Compute expected policy from expected value"""
    E = exp.( U + β .* [(T[:,:,1] * EV_) (T[:,:,2] * EV_)] )
    EP_ = E[:,2] ./ sum(E, dims=2)
    return vec(EP_)
end;

function logL_HM(θ0::Vector, λ::Number, β::Number, s::Vector, St::Vector, A::BitVector, T::Array, CCP::Matrix)::Number
    """Compute log-likelihood function for HM problem"""
    # Compute static utility
    U = compute_U(θ0, s)

    # Espected value by inversion
    EV_ = HM_inversion(CCP, T, U, β)

    # Implies choice probabilities
    EP_ = from_EV_to_EP(EV_, T, U, β)

    # Likelihood
    logL = sum(log.(EP_[St[A.==1]])) + sum(log.(1 .- EP_[St[A.==0]]))
    return -logL
end;

function logL_AM(θ0::Vector, λ::Number, β::Number, s::Vector, St::Vector, A::BitVector, T::Array, CCP::Matrix, K::Int)::Number
    """Compute log-likelihood function for AM problem"""
    # Compute static utility
    U = compute_U(θ0, s)
    EP_ = CCP[:,2]

    # Iterate HM mapping
    for _=1:K
        EV_ = HM_inversion(CCP, T, U, β)    # Espected value by inversion
        EP_ = from_EV_to_EP(EV_, T, U, β)   # Implies choice probabilities
        CCP = [(1 .- EP_) EP_]
    end

    # Likelihood
    logL = sum(log.(EP_[St[A.==1]])) + sum(log.(1 .- EP_[St[A.==0]]))
    return -logL
end;

## Main

# Set parameters
θ = [0.13; -0.004; 3.1];
λ = 0.82;
β = 0.95;

# State space
k = 10;
s = Vector(1:k);

# Compute value function
Vbar = compute_Vbar(θ, λ, β, s);

# Generate data
N = Int(1e5);
St, A, St1 = generate_data(θ, λ, β, s, N);
print("\n\nwe observe ", sum(A), " investment decisions in ", N, " observations")

# Estimate lambda
Δ = St1 - St;
λ_ = mean(Δ[(A.==0) .& (St.<10)]);

print("\n\nEstimated lambda: $λ_ (true = $λ)")

# True likelihood value
logL_trueθ = logL_Rust(θ, λ, β, s, St, A);
print("\n\nThe likelihood at the true parameter is $logL_trueθ")

# Select starting values
θ0 = Float64[0,0,0];

# Optimize
θ_R = optimize(x -> logL_Rust(x, λ, β, s, St, A), θ0).minimizer;
print("\n\nEstimated thetas: $θ_R (true = $θ)")

# Not all initial values are equally good
θ0 = Float64[1,1,1];

# Optimize
θ_R2 = optimize(x -> logL_Rust(x, λ, β, s, St, A), θ0).minimizer;
print("\n\nEstimated thetas: $θ_R2 (true = $θ)")

# Estimate CCP
P = [mean(A[St.==i]) for i=s]
CCP = [(1 .- P) P]

# Compute T
T = compute_T(k, λ_)

# Optimize
θ0 = Float64[0,0,0];
θ_HM = optimize(x -> logL_HM(x, λ, β, s, St, A, T, CCP), θ0).minimizer;
print("\n\nEstimated thetas: $θ_HM (true = $θ)")

# Aguirregabiria & Mira (2002)
K = 2

# Optimize
θ0 = Float64[0,0,0];
θ_AM = optimize(x -> logL_AM(x, λ, β, s, St, A, T, CCP, K), θ0).minimizer;
print("\n\nEstimated thetas: $θ_AM (true = $θ)")

# Compare times
θ0 = Float64[0,0,0];
time_Rust = optimize(x -> logL_Rust(x, λ, β, s, St, A), θ0).time_run;
time_HM = optimize(x -> logL_HM(x, λ, β, s, St, A, T, CCP), θ0).time_run;
time_AM = optimize(x -> logL_AM(x, λ, β, s, St, A, T, CCP, K), θ0).time_run;
print("Time Rust: $time_Rust\nTime HM: $time_HM\nTime AM: $time_AM")


### A Pluto.jl notebook ###
# v0.12.3

using Markdown
using InteractiveUtils

# ╔═╡ 8e3d151e-0b8b-11eb-1f36-319328a7772c
using Flux

# ╔═╡ 92dfefce-0b8b-11eb-32b4-391aecbbb876
using Distributions

# ╔═╡ 54417ed0-0b8c-11eb-2064-230c0020f268
using IterTools: ncycle

# ╔═╡ 622d190a-0b8c-11eb-1f00-63d9bfdca449
using Random

# ╔═╡ 65e0990a-0b8c-11eb-01ce-233699db160b
using LinearAlgebra

# ╔═╡ a68b270e-0b8c-11eb-050e-b704f7137d1c
using Parameters: @with_kw

# ╔═╡ 20ef48ee-0b8e-11eb-3f94-313efbfc9478
using PlutoUI

# ╔═╡ 055b2642-0b8f-11eb-2777-c1159e9e0e21
using Plots

# ╔═╡ 0f73843e-0b90-11eb-1425-5dce556a16a9
using BenchmarkTools

# ╔═╡ c63890b6-0b8a-11eb-3771-5bde29bc994b
md"""
# Regularization

Weight attenuation is equivalent to regularization of $L_2$ norm, Regularization is a common method to deal with overfitting by adding penalty term to model loss function to make the learned model parameter values smaller.

$l(w, b) + \frac{\lambda}{2n}||w||^2$

among them，$||w||^2 = w_{1}^2 + w_{2}^{2}$

## Higer-dimensional linear regression experiments

Taking high dimensional linear regression as an example, an overfitting problem is introduced.The feature dimension of the data sample is $P$, and the set feature of the data of the training set and the test set is $x_1, x_2, ..., x_p$.

$y = 0.05 + \sum_{i=1}^{p} 0.01x_i + \epsilon$

so we set dimension $p = 200$, smaple number $n=20$:

"""

# ╔═╡ 81ed2efc-0b8b-11eb-1b99-33cb760768fb
md"""
- first import some packages:
"""

# ╔═╡ 6d6b9062-0b8c-11eb-1c2c-d1a65f7ff889
md"""
- do simulation:
"""

# ╔═╡ 949227d2-0b8c-11eb-3ce8-211d2b5010b1
begin
	@with_kw mutable struct Args
		n::Int = 20
		p::Int = 200
		lr::Float64 = 0.003
	end

	args = Args()

	w, b = ones(args.p)*0.01, 0.05
	features = rand(Normal(0,1), (120, args.p))
	labels = features * w .+ b + rand(120)/100
	train_data = Flux.Data.DataLoader((Array(features[1:20, :]'),Array(labels[1:20]')),batchsize=1)
	test_x,test_y = Array(features[21:120, :]'), Array(labels[21:120]')
end

# ╔═╡ 37796f14-0b8d-11eb-350a-cb297a03f349
begin
	train_loss = []
	test_loss = []
	for λ in 0:0.1:4
    	model = Chain(Dense(200,1))
    	ps = Flux.params(model)
    	opt = Descent(args.lr)
    	sqnorm(x) = sum(abs2, x)
    	loss(x,y) = Flux.mse(model(x), y) + λ*sum(sqnorm, Flux.params(model))
    	Flux.train!(loss, ps, train_data, opt);
    	push!(train_loss, Flux.mse(model(Array(features[1:20, :]')), Array(labels[1:20]')))
    	push!(test_loss, Flux.mse(model(test_x), test_y))
	end
	plot(0:0.1:4, [train_loss test_loss], label=["train_loss" "test_loss"], xlabel="λ", ylabel="Loss", title="regularization λ influence the loss")
end

# ╔═╡ 80566366-0b8f-11eb-3ee6-37f6064d46fe
md"""
As we can see, the coefficient of regularization is $\lambda$, if $0<\lambda<1$, we have high loss in test but at train set is great,absolutely, we overfitting train set. when the $\lambda$ become more larger, the test loss start to fall down.

## Dropout
In multilayer perceptron, we describe a multilayer perceptron with a single hidden layer, assuming that number of inputs is 4, and the number of hidden layer unit is 5. So

$h_i = \phi(x_1w_{1i} + x_2w_{2i} + x_3w_{3i}+x_4w_{4i} + b_i)$

among them, $\phi$ is activative function, $x_1,...x_4$ is input value, $w_{1i}, ..., w_{4i}$ is hidden layer parameters.

When use dropout way to this hidden layer, hidden units in this layer will have the chance to be dropped.set the probability is $p$. drop out don't change the expect.

Because in training process, hidden units will be dropped out by $p$, so output calculate can't rely too much on anyone in $h_1,..h_5$.

**For example:**
"""

# ╔═╡ ebc4d248-0b8f-11eb-1123-810b0d1586ab
md"**_Import data:_**"

# ╔═╡ 38c61114-0b90-11eb-3c6f-2d622821ce69
begin
	images = Flux.Data.FashionMNIST.images()
	labels_i = Flux.Data.FashionMNIST.labels()
	k = sort(unique(labels_i))
	onehot_labels = Flux.onehotbatch(labels_i, k)
	train_x = rand(784, 6000)
	for i in 1:6000
		train_x[:,i] = Float64.(reshape(images[i],(784,1)))
	end
	train_data_i = Flux.Data.DataLoader(Array(train_x[:, 1:4000]), Array(onehot_labels[:, 1:4000]), batchsize=10, shuffle=true)
end

# ╔═╡ 8fdaa60c-0b90-11eb-0fca-510bd890f89e
md"**_dropout some layers_**:
every layers has the same probability $p$ of being dropped."

# ╔═╡ d7dc8f6a-0b90-11eb-371b-d9be51947ec4
begin
	m_d = Chain(
    Dense(784, 256, relu),
    Dropout(0.2),
    Dense(256, 256, relu),
    Dropout(0.5),
    Dense(256, 10, relu))
	loss1(x, y) = Flux.logitcrossentropy(m_d(x), y)
	opt_d = Descent(0.05)
	ps_d = Flux.params(m_d)
	@btime Flux.train!(loss1, ps_d, ncycle(train_data_i, 10), opt_d)
end

# ╔═╡ 850218e0-0b91-11eb-0fd9-f13ca40fcb84
md"define accuracy function:"

# ╔═╡ 9f6b4ccc-0b91-11eb-09f4-eb564e81dfc1
accuracy(x, y, model) = Flux.mean(Flux.onecold(model(x)) .== Flux.onecold(y))

# ╔═╡ b9dc1886-0b91-11eb-1bc3-9bb26f25af96
md"show the train accuracy:"

# ╔═╡ aa557fa6-0b91-11eb-0743-233eff7ac496
accuracy(Array(train_x[:, 1:4000]), Array(onehot_labels[:, 1:4000]), m_d)

# ╔═╡ d3021716-0b91-11eb-012c-2f69cb71d076
md"show the test accuracy:"

# ╔═╡ e6b006d0-0b91-11eb-2d8d-29558ea1d719
accuracy(Array(train_x[:, 4000:6000]), Array(onehot_labels[:, 4000:6000]), m_d)

# ╔═╡ Cell order:
# ╟─c63890b6-0b8a-11eb-3771-5bde29bc994b
# ╟─81ed2efc-0b8b-11eb-1b99-33cb760768fb
# ╠═8e3d151e-0b8b-11eb-1f36-319328a7772c
# ╠═92dfefce-0b8b-11eb-32b4-391aecbbb876
# ╠═54417ed0-0b8c-11eb-2064-230c0020f268
# ╠═622d190a-0b8c-11eb-1f00-63d9bfdca449
# ╠═65e0990a-0b8c-11eb-01ce-233699db160b
# ╠═a68b270e-0b8c-11eb-050e-b704f7137d1c
# ╠═20ef48ee-0b8e-11eb-3f94-313efbfc9478
# ╠═055b2642-0b8f-11eb-2777-c1159e9e0e21
# ╠═0f73843e-0b90-11eb-1425-5dce556a16a9
# ╟─6d6b9062-0b8c-11eb-1c2c-d1a65f7ff889
# ╠═949227d2-0b8c-11eb-3ce8-211d2b5010b1
# ╠═37796f14-0b8d-11eb-350a-cb297a03f349
# ╟─80566366-0b8f-11eb-3ee6-37f6064d46fe
# ╟─ebc4d248-0b8f-11eb-1123-810b0d1586ab
# ╠═38c61114-0b90-11eb-3c6f-2d622821ce69
# ╟─8fdaa60c-0b90-11eb-0fca-510bd890f89e
# ╠═d7dc8f6a-0b90-11eb-371b-d9be51947ec4
# ╠═850218e0-0b91-11eb-0fd9-f13ca40fcb84
# ╟─9f6b4ccc-0b91-11eb-09f4-eb564e81dfc1
# ╟─b9dc1886-0b91-11eb-1bc3-9bb26f25af96
# ╠═aa557fa6-0b91-11eb-0743-233eff7ac496
# ╠═d3021716-0b91-11eb-012c-2f69cb71d076
# ╠═e6b006d0-0b91-11eb-2d8d-29558ea1d719


using FileIO
using Parameters

include("random_matrices.jl")
include("network_sims_gpu.jl")
include("LS_calculations.jl")

#------------------------------------------------------------------------------
# This is the part where we find which eta and network realization to

J_seeds = collect(5000:5019)
etas = collect(-1.0f0:0.1f0:1.0f0)

ind = parse(Int64, ARGS[1])
ps = [(eta, J_seed) for eta in etas, J_seed in J_seeds]
eta, J_seed = ps[ind]
#------------------------------------------------------------------------------

# Where the results will be saved
dir = "results1"

N = 3000
g = 1.5f0
s = 1f0

h0_seed = 6000
dh_seed = 7000
Q0_seed = 8000

J = generate_J(N, eta, seed=J_seed)
h = generate_nv(N, Float32, seed=h0_seed)
dh = generate_nv(N, Float32, seed=dh_seed)
Q = generate_ONB(N, Float32, seed=Q0_seed)

dt = 1f-2
t_bound = 2f3
t_ONS = 1f0
wait = 2f2

@time Lspec, max_LE = sim_lyapunov_gpu(J, g, h, dh, Q; s=s, dt=dt,
		t_bound=t_bound, t_ONS=t_ONS, wait=wait, return_max_LE=true,
		printstuff=false, p_factor=1e-3, scale_limit=100)

dimensionality = KY_attractor_dim(Lspec)
entropy_rate = KS_entropy_rate(Lspec)

fn = "Lspec_E$(ind)_N=$N,g=$g,eta=$eta,J_seed=$J_seed.jld2"
out = Dict{String,Any}()
@pack! out = N, g, eta, s, J_seed, h0_seed, dh_seed, Q0_seed, dt, t_bound,
		t_ONS, wait, Lspec, max_LE, dimensionality, entropy_rate

save(joinpath(dir, fn), out)


cfr(k) = print("Kelvin: $k, ",
  "Celsius: $(round(k-273.15,2)), ",
  "Fahrenheit: $(round(k*1.8-459.67,2)), ",
  "Rankine: $(round(k*1.8,2))")


using CITIfile: rdciti
using Test

s,f = rdciti("seg_list.cit")
@testset "SEG_LIST-style file" begin
    @test size(s) == (1,1,10)
    @test size(f) == (10,)
    @test f == range(1, 4, length=10)
    @test s[1,1,1] ≈ complex(0.86303E-1,-8.98651E-1)
    @test s[1,1,end] ≈ complex(-7.78350E-1,5.72082E-1)
end


using Test
using Distributions
include("transforms.jl")
include("legpts.jl")
include("statistics.jl")

@testset "transforms.jl" begin
    @test invBoxCox(boxCox(2, 3), 3) ≈ 2.0 atol = 1e-15
    @test invBoxCox(boxCox(2, 1), 1) ≈ 2.0 atol = 1e-15
    @test boxCox(invBoxCox(2, 3), 3) ≈ 2.0 atol = 1e-15
    @test boxCox(invBoxCox(2, 1), 1) ≈ 2.0 atol = 1e-15
    @test boxCox(invBoxCox([1 2 3], 1), 1) ≈ [1 2 3] atol = 1e-15
    @test boxCox(invBoxCox([1; 2], 1), 1) ≈ [1;2] atol = 1e-15
    @test boxCox(invBoxCox(2., -0.45), -0.45) ≈ 2.0 atol = 1e-15
    @test invBoxCox(boxCox(2., -2), -2) ≈ 2.0 atol = 1e-15
end

@testset "legpts.jl" begin
    f = x -> exp(x)+x^2
    @test_broken int1D(f, -1, 1) ≈ 3.01706 atol = 1e-5
end

@testset "statistics" begin
    log = LogNormal(2, 1)
    y = x-> Distributions.cdf(log, x)-0.5
    med = median(log)
    @test_broken bisection(y, 0, 15) ≈  med atol = 1e-3
end

@testset ""


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using StaticArrays, ForwardDiff

@testset "forwarddiff_utils" begin
    @test @inferred(ForwardDiff.value(BAT.forwarddiff_eval(sin, 0.5))) == sin(0.5)
    @test @inferred(ForwardDiff.partials(BAT.forwarddiff_eval(sin, 0.5))[1]) == cos(0.5)
    @test length(ForwardDiff.partials(BAT.forwarddiff_eval(sin, 0.5))) == 1

    @test @inferred(BAT.forwarddiff_vjp(0.7, BAT.forwarddiff_eval(sin, 0.5))) == 0.7 * ForwardDiff.derivative(sin, 0.5)

    f(x) = SVector(1.1 * x[1] + 2.1 * x[2] + 3.1 * x[3], -1.2 * x[1] - 2.2 * x[2] + -3.2 * x[3])

    x = SVector(0.1, 0.2, 0.3)
    ΔΩ = SVector((10, 20))

    @test @inferred(BAT.forwarddiff_vjp(ΔΩ, BAT.forwarddiff_eval(f, x))) == ForwardDiff.jacobian(f, x)' * ΔΩ

    @test @inferred((x -> BAT.forwarddiff_pullback(f, x)[1])(x)) == f(x)
    @test @inferred(BAT.forwarddiff_pullback(f, x)[2](ΔΩ)) == ForwardDiff.jacobian(f, x)' * ΔΩ

    X = [SVector(0.1i, 0.2i, 0.3i) for i in 1:7]
    ΔΩs = [SVector(10i, 20i) for i in 1:7]
    @test @inferred((X -> BAT.forwarddiff_broadcast_pullback(Ref(f), X)[1])(X)) == broadcast(f, X)
    @test @inferred(BAT.forwarddiff_broadcast_pullback(f, X)[2](ΔΩs)) == broadcast((f, x, ΔΩ) -> ForwardDiff.jacobian(f, x)' * ΔΩ, Ref(f), X, ΔΩs)
end


using Econometrics, Statistics, Random

# version which generates shock internally
function SVmodel(θ, n, burnin)
    ϕ = θ[1]
    ρ = θ[2]
    σ = θ[3]
    hlag = 0.0
    ys = zeros(n)
    for t = 1:burnin+n
        # bound the log variance at a high value,
        # too bullet-proof things which use this
        h = min(ρ*hlag + σ*randn(),20.0)
        y = ϕ*exp(h/2.0)*randn()
        if t > burnin
            ys[t-burnin] = y
        end
        hlag = h
    end
    ys, ϕ*exp(hlag/2.0) # return the sample of returns, plus the final period volatility
end

# auxiliary model: HAR-RV
# Corsi, Fulvio. "A simple approximate long-memory model
# of realized volatility." Journal of Financial Econometrics 7,
# no. 2 (2009): 174-196.
function HAR(y)
    ylags = lags(y,10)
    X = [ones(size(y,1)) ylags[:,1]  mean(ylags[:,1:4],dims=2) mean(ylags[:,1:10],dims=2)]
    # drop missings
    y = y[11:end]
    X = X[11:end,:]
    βhat = X\y
    σhat = std(y-X*βhat)
    vcat(βhat,σhat)
end

function aux_stat(y)
    s = std(y)
    y = abs.(y)
    m = mean(y)
    s2 = std(y)
    k = std(y.^2.0)
    c = cor(y[1:end-1],y[2:end])
    # ratios of quantiles of moving averages to detect clustering
    q = try
        q = quantile(ma(y/s2,4), [0.1, 0.25, 0.75, 0.9])
    catch
        q = [1.0, 1.0, 0.0, 0.0]
    end
        c1 = q[4]/q[1]
        c2 = q[3]/q[2]
    vcat(m, s, s2, k, c, c1, c2, HAR(y))
    # four noise stats
    #vcat(m, s, s2, k, c, c1, c2, HAR(y), randn(4)./sqrt(size(y,1)))
end

# asymptotic Gaussian likelihood function of statistic
function logL(θ, m, n, η, ϵ, withdet=true)
    S = size(η,2)
    k = size(m,1)
    ms = zeros(S, k)
    Threads.@threads for s = 1:S
        y, junk = SVmodel(θ, n, η[:,s], ϵ[:,s])
        y = min.(y, 100.0)
        y = max.(y,-100.0)
        ms[s,:] = sqrt(n)*aux_stat(y)
    end
    mbar = mean(ms,dims=1)[:]
    if ~any(isnan.(mbar))
        Σ = cov(ms)
        x = (m .- mbar)
        lnL = try
            if withdet
                lnL = -0.5*log(det(Σ)) - 0.5*x'*inv(Σ)*x # for Bayesian
            else
                lnL = 0.5*x'*inv(Σ)*x # for classic indirect inference (note sign change)
            end
        catch
            lnL = -Inf
        end
     else
         lnL = -Inf
     end
     return lnL
end

# prior just checks that we're in the bounds
function prior(theta, lb, ub)
    a = 0.0
    if(all((theta .>= lb) .& (theta .<= ub)))
        a = 1.0
    end
    return a
end

# uniform random walk in one dimension
function proposal1(current, tuning, lb, ub)
    trial = copy(current)
    if rand() > 0.1
        i = rand(1:size(trial,1))
        trial[i] += tuning[i].*randn()
    else
        trial = lb + (ub - lb).*rand(size(lb,1))
    end
    return trial
end

# MVN random walk, or occasional draw from prior
function proposal2(current, cholV, lb, ub)
    trial = copy(current)
    if rand() > 0.1
        trial += cholV'*randn(size(trial))
    else
        trial = lb + (ub - lb).*rand(size(lb,1))
    end
    return trial
end


abstract type KDEMethod end
kde(m::KDEMethod, args...; kwargs...) = error("kde not implemented for $m, args=$args, kwargs=$kwargs")

struct FixedBandwidth <: KDEMethod
    GMMType
    h::Float64
end
FixedBandwidth(h::Float64) = FixedBandwidth(UnivariateGMM, h)
bandwidth(m::FixedBandwidth, b::AbstractParticleBelief) = m.h

struct SilvermansRule <: KDEMethod
    GMMType
end
SilvermansRule() = SilvermansRule(UnivariateGMM)
bandwidth(m::SilvermansRule, b::AbstractParticleBelief) = max((4/(3*n_particles(b)))^(0.2) * std(b), sqrt(eps()))

kde(m::Union{FixedBandwidth, SilvermansRule}, b::AbstractParticleBelief) = GMMType(b, bandwidth(m, b))


# This file is part of the IntervalArithmetic.jl package; MIT licensed

+(a::IntervalBox, b::IntervalBox) = IntervalBox( a.v .+ b.v )
+(a::IntervalBox, b::Real) = IntervalBox( a.v .+ b )
+(a::Real, b::IntervalBox) = IntervalBox( a .+ b.v )

-(a::IntervalBox, b::IntervalBox) = IntervalBox( a.v .- b.v )
-(a::IntervalBox, b::Real) = IntervalBox( a.v .- b )
-(a::Real, b::IntervalBox) = IntervalBox( a .- b.v )
-(a::IntervalBox) = IntervalBox( .- a.v )

*(a::IntervalBox, b::Real) = IntervalBox( a.v .* b )
*(a::Real, b::IntervalBox) = IntervalBox( a .* b.v )

/(a::IntervalBox, b::Real) = IntervalBox( a.v ./ b )

# broadcasting:

# wrap decides whether to wrap the result in an IntervalBox or not, based on the return type
wrap(v::SVector{N,T} where {N,T<:Interval}) = IntervalBox(v)
wrap(v) = v

@inline Base.broadcast(f, X::IntervalBox) = wrap(f.(X.v))
@inline Base.broadcast(f, X::IntervalBox, Y::IntervalBox) = wrap(f.(X.v, Y.v))
@inline Base.broadcast(f, X::IntervalBox, y) = wrap(f.(X.v, y))

for op in (:+, :-, :∩, :∪, :⊆, :isinterior, :dot, :setdiff, :×)
    @eval $(op)(a::SVector, b::IntervalBox) = $(op)(IntervalBox(a), b)
end

for op in (:+, :-, :∩, :∪, :⊆, :isinterior, :dot, :setdiff, :×)
    @eval $(op)(a::IntervalBox, b::SVector) = $(op)(a, IntervalBox(b))
end


# LogLinear Transformation

"""
    LogLinearFormula(df::Int64)

Returns GLM formula can be used in `glm` with `FreqTab` data frame.
`df`, is abbreviation for degrees of formula, represents degree of polynomial log-linear method.
This function works very slowly. If a fixed degree formula will be used repeatedly, define the formula as a variable.
"""
function LogLinearFormula(df)
    fml = "@formula(freq ~ 1 +"
    for d in 1:df
        fml *= "scale^$d"
        if d != df
            fml *= " + "
        else
            fml *= ")"
        end
    end
    return eval(Meta.parse(fml))
end

const fml₁ = @formula(freq ~ 1 + scale)
const fml₂ = @formula(freq ~ 1 + scale + scale^2)
const fml₃ = @formula(freq ~ 1 + scale + scale^2 + scale^3)
const fml₄ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4)
const fml₅ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4 + scale^5)
const fml₆ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4 + scale^5 + scale^6)
const fml₇ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4 + scale^5 + scale^6 + scale^7)
const fml₈ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4 + scale^5 + scale^6 + scale^7 + scale^8)
const fml₉ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4 + scale^5 + scale^6 + scale^7 + scale^8 + scale^9)
const fml₁₀ = @formula(freq ~ 1 + scale + scale^2 + scale^3 + scale^4 + scale^5 + scale^6 + scale^7 + scale^8 + scale^9 + scale^10)

mutable struct SmoothedFreqTab <: EG
    table
    raw
    interval
    fit
end
"""
    presmoothing(F::FreqTab, @LogLinearFormula(df::Int64))

Returns presmoothed frequency table as `SmoothedFreqTab` and `glm` fitted object.

Preserving first C moments of original frequency data, passed `LogLinearFormula(C)` to `fml`.
C is a degree of freedom

# Examples
```julia
julia> tab = freqtab(expandtable(ACTmath.scale, ACTmath.xcount));
julia> presmoothing(tab, Equate.fml₄)
Equate.SmoothedFreqTab(40×5 DataFrame
│ Row │ scale   │ freq     │ cumfreq  │ prob        │ cumprob     │
│     │ Float64 │ Float64? │ Float64? │ Float64     │ Float64     │
├─────┼─────────┼──────────┼──────────┼─────────────┼─────────────┤
│ 1   │ 1.0     │ 0.864929 │ 0.864929 │ 0.000199799 │ 0.000199799 │
│ 2   │ 2.0     │ 2.47529  │ 3.34022  │ 0.000571794 │ 0.000771592 │
⋮
│ 38  │ 38.0    │ 31.9351  │ 4291.07  │ 0.00737703  │ 0.991238    │
│ 39  │ 39.0    │ 22.7842  │ 4313.85  │ 0.00526315  │ 0.996501    │
│ 40  │ 40.0    │ 15.1484  │ 4329.0   │ 0.00349928  │ 1.0         │, [1, 2, 3, 3, 3, 4, 4, 4, 4, 4  …  40, 40, 40, 40, 40, 40, 40, 40, 40, 40], 1.0, StatsModels.TableRegressionModel{GLM.GeneralizedLinearModel{GLM.GlmResp{Array{Float64,1},Distributions.Poisson{Float64},GLM.LogLink},GLM.DensePredChol{Float64,LinearAlgebra.Cholesky{Float64,Array{Float64,2}}}},Array{Float64,2}}

freq ~ 1 + scale + :(scale ^ 2) + :(scale ^ 3) + :(scale ^ 4)

Coefficients:
─────────────────────────────────────────────────────────────────────────────────
                   Coef.   Std. Error       z  Pr(>|z|)    Lower 95%    Upper 95%
─────────────────────────────────────────────────────────────────────────────────
(Intercept)  -1.35982     0.310243      -4.38    <1e-4   -1.96789     -0.751757
scale         1.30127     0.0728808     17.85    <1e-70   1.15843      1.44412
scale ^ 2    -0.089081    0.00589528   -15.11    <1e-50  -0.100636    -0.0775265
scale ^ 3     0.00254824  0.000195003   13.07    <1e-38   0.00216604   0.00293044
scale ^ 4    -2.67698e-5  2.25114e-6   -11.89    <1e-31  -3.1182e-5   -2.23576e-5
─────────────────────────────────────────────────────────────────────────────────)

```
"""
function presmoothing(F::EG, fml)
    fit1 = glm(fml, F.table, Poisson(), LogLink())
    freq = predict(fit1, DataFrame(scale = F.table.scale))
    tab = DataFrame(scale = F.table.scale, freq = freq, cumfreq = cumsum(freq),
                    prob = freq ./ sum(freq), cumprob = cumsum(freq) ./ sum(freq))
    return SmoothedFreqTab(tab, F.raw, F.interval, fit1)
end

"""
    presmoothing(F::EG, degree::Int64)
Examine various degree of freedom for fitting the log liner models to the raw score histogram.

# Examples
```julia
julia> tab = freqtab(expandtable(ACTmath.scale, ACTmath.xcount))
Equate.FreqTab(40×5 DataFrame
│ Row │ scale   │ freq  │ cumfreq │ prob       │ cumprob  │
│     │ Float64 │ Int64 │ Int64   │ Float64    │ Float64  │
├─────┼─────────┼───────┼─────────┼────────────┼──────────┤
│ 1   │ 1.0     │ 1     │ 1       │ 0.000231   │ 0.000231 │
│ 2   │ 2.0     │ 1     │ 2       │ 0.000231   │ 0.000462 │
⋮
│ 38  │ 38.0    │ 38    │ 4291    │ 0.00877801 │ 0.991222 │
│ 39  │ 39.0    │ 23    │ 4314    │ 0.00531301 │ 0.996535 │
│ 40  │ 40.0    │ 15    │ 4329    │ 0.003465   │ 1.0      │, [1, 2, 3, 3, 3, 4, 4, 4, 4, 4  …  40, 40, 40, 40, 40, 40, 40, 40, 40, 40], 1.0)

julia> fit = presmoothing(tab, 10)
Fitting dof = 1 model.
Fitting dof = 2 model.
Fitting dof = 3 model.
Fitting dof = 4 model.
Fitting dof = 5 model.
Fitting dof = 6 model.
Fitting dof = 7 model.
Fitting dof = 8 model.
Fitting dof = 9 model.
Fitting dof = 10 model.
10×7 DataFrame. Omitted printing of 1 columns
│ Row │ degree │ aic     │ aicc    │ bic     │ loglik   │ deviance │
│     │ Int64  │ Float64 │ Float64 │ Float64 │ Float64  │ Float64  │
├─────┼────────┼─────────┼─────────┼─────────┼──────────┼──────────┤
│ 1   │ 1      │ 2235.18 │ 2235.5  │ 2238.56 │ -1115.59 │ 1988.28  │
│ 2   │ 2      │ 661.558 │ 662.225 │ 666.625 │ -327.779 │ 412.654  │
⋮
│ 8   │ 8      │ 290.675 │ 296.675 │ 305.875 │ -136.338 │ 29.7704  │
│ 9   │ 9      │ 292.573 │ 300.159 │ 309.462 │ -136.286 │ 29.6683  │
│ 10  │ 10     │ 294.051 │ 303.48  │ 312.629 │ -136.026 │ 29.1466  │

```
"""
function presmoothing(F::EG, degree::Int64)
    if degree > 11
        throw(ArgumentError("Maximum degree for examining must be less than ore equal to 10."))
    end
    fmls = (fml₁, fml₂, fml₃, fml₄, fml₅, fml₆, fml₇, fml₈, fml₉, fml₁₀)
    AIC, BIC, AICC, DEVIANCE, LL = zeros.(Float64, fill(degree, 5))
    FIT = Any[]
    for d in 1:degree
        println("Fitting dof = $d model.")
        fit = glm(fmls[d], F.table, Poisson(), LogLink())
        push!(FIT, fit)
        AIC[d] = StatsBase.aic(fit)
        BIC[d] = StatsBase.bic(fit)
        AICC[d] = StatsBase.aicc(fit)
        LL[d] = StatsBase.loglikelihood(fit)
        DEVIANCE[d] = StatsBase.deviance(fit)
    end
    DataFrame(degree = [1:1:degree;], aic = AIC, aicc = AICC, bic = BIC, loglik = LL, deviance = DEVIANCE, fit = FIT)
end

mutable struct SmoothedNEATFreqTab <: NEAT
    tableX
    tableV
    rawX
    rawV
    intervalX
    intervalV
    fitX
    fitV
    marginal
end
"""
    presmoothing(F::NEAT, LogLinearFormula(df::Int64), LogLinearFormula(df::Int64))

Returns presmoothed frequency table as `SmoothedNEATFreqTab` and `glm` fitted object,
but `interval` and `marginal` elements, which are returned, are not based on smoothed score.

Preserving first C moments of original frequency data, passed `LogLinearFormula(C)` to `fml`.
C is a degree of freedom
"""
function presmoothing(F::NEAT, fmlX, fmlV)
    # Smoothing X (independent part)
    fitX = glm(fmlX, F.tableX, Poisson(), LogLink())
    freqX = predict(fitX, DataFrame(scale = F.tableX.scale))
    tabX = DataFrame(scale = F.tableX.scale, freq = freqX, cumfreq = cumsum(freqX),
                    prob = freqX ./ sum(freqX), cumprob = cumsum(freqX) ./ sum(freqX))
    # Smoothing V (common part)
    fitV = glm(fmlV, F.tableV, Poisson(), LogLink())
    freqV = predict(fitV, DataFrame(scale = F.tableV.scale))
    tabV = DataFrame(scale = F.tableV.scale, freq = freqV, cumfreq = cumsum(freqV),
                    prob = freqV ./ sum(freqV), cumprob = cumsum(freqV) ./ sum(freqV))
    return SmoothedNEATFreqTab(tabX, tabV, F.rawX, F.rawV, F.intervalX, F.intervalV, fitX, fitV, F.marginal)
end

# Kernel method
function RjX(x, xⱼ, a, μ, hX)
    return (x - a * xⱼ - (1-a)*μ) / (a*hX)
end
struct KernelFreqTab <: EG
    table
    raw
    interval
    Bandwidth
end
"""
    KernelSmoothing(X::SmoothedFreqTab; kernel = :Gaussian, hX = 0.622, scale = X.table.scale)

Returns Kernel smoothed frequency table as `KernelFreqTab`.

# Arguments

- `X` the object which class is `freqtab`

## Optional Arguments

- `kernel`
- `hX` The band width.
- `newint` A new interval to depict the continuized probability line.

# Value

`KernelFreqTab`

In the kernel smoothing, the choice of bandwidth `hX` is an important consideration. When bandwidth is large, the smoothed distribution becomes linear equating fucnction. Contrary, bandwidth is small, it becomes spikye shaped distribution. We recommend to use default value `hX = 6.22` or select it by using `EstBandwidth`.
"""
function KernelSmoothing(X::EG; kernel = :Gaussian, hX = 0.622, newint = X.interval/100)
    scale = (minimum(X.table.scale) - X.interval/2):newint:(maximum(X.table.scale) + X.interval/2)
    # hX = bandwidth of cumulative distribution function
    N = sum(X.table.freq)
    μ = mean(X.raw); σ² = var(X.raw)
    a² = σ² / (σ² + hX^2)
    a = sqrt(a²)
    𝒇hX = zeros(Float64, length(scale))
    FhX = zeros(Float64, length(scale))
    for (i, x) in enumerate(scale), (j, xⱼ) in enumerate(X.table.scale)
        𝒇hX[i] += X.table.prob[j]*pdf.(Normal(0, 1), RjX(x, xⱼ, a, μ, hX)) / (a*hX)
        FhX[i] += X.table.prob[j]*cdf.(Normal(0, 1), RjX(x, xⱼ, a, μ, hX))
    end
    𝒇hX = 𝒇hX ./ sum(𝒇hX) # normalize
    tbl = DataFrame(scale = scale, freq = N .* 𝒇hX, cunfreq = cumsum(N .* 𝒇hX), prob = 𝒇hX, cumprob = cumsum(𝒇hX))
    return KernelFreqTab(tbl, X.raw, X.interval, hX)
end

function BandwidthPenalty(hX, X::SmoothedFreqTab, lprob, rprob, K; kernel = :Gaussian)
    hX = exp(hX[1])
    r = X.table.prob
    μ = mean(X.raw); σ² = var(X.raw)
    a² = σ² / (σ² + hX^2)
    a =sqrt(a²)
    𝒇hX = zeros(Float64, length(X.table.scale))
    r𝒇′hX = zeros(Float64, length(X.table.scale))
    l𝒇′hX = zeros(Float64, length(X.table.scale))
    for (i, x) in enumerate(X.table.scale), (j, xⱼ) in enumerate(X.table.scale)
        R = RjX(x, xⱼ, a, μ, hX)
        𝒇hX[i] += X.table.prob[j]*pdf.(Normal(0, 1), R) / (a*hX)
        r𝒇′hX[i] -= rprob[j]*pdf.(Normal(0, 1), R) / (a*hX)^2 * R
        l𝒇′hX[i] -= lprob[j]*pdf.(Normal(0, 1), R) / (a*hX)^2 * R
    end
    pen1 = sum((r .- 𝒇hX) .^2)
    pen2 = sum((l𝒇′hX .< 0) .* (r𝒇′hX .> 0))
    return pen1 + K * pen2
end

"""
    EstBandwidth(X::EG; kernel = :Gaussian)

Estimate optimal bandwidth `hX` in Gaussian kernel.

# Arguments

- `X` is the object which the class is `freqtab`.

# Value

Optimize information.

"""
function EstBandwidth(X::SmoothedFreqTab; kernel = :Gaussian, K = 1)
    x = X.table.scale
    lscore = x .- X.interval/4
    rscore = x .+ X.interval/4
    lpred = predict(X.fit, DataFrame(scale = lscore))
    rpred = predict(X.fit, DataFrame(scale = rscore))
    opt = optimize(hX -> BandwidthPenalty(hX, X, lpred, rpred, K; kernel = kernel), [0.5], method = BFGS())
    println("Minimizer sould be transformed `exp()` before interpretation. \nMinimizer (after taking `exp()`) = $(exp(opt.minimizer[1]))\n\n")
    return opt
end


# Test to rewrite the reconstructions functions with inputs the data instead of some noise level.

push!(LOAD_PATH, "../models")
push!(LOAD_PATH, "../")
using SparseInverseProblems
using SuperResModels
using TestCases
using Utils

using PyCall
@pyimport numpy as np
@pyimport matplotlib.pyplot as plt

# Containers to handle parameters
type StaticParameters
    # allowed frequencies
    f_c::Int64
    # spatial max val.
    x_max::Float64
    # ???
    num_x::Int64
    # Weights on the allowed frequencies.
    filter_x::Array{Float64}
end
type DynamicParameters
    # 2K+1 is the number of time samples.
    K::Int64
    # sampling rate
    tau::Float64
    # maxium allowed speed
    v_max::Float64
    # ????
    num_v::Int64
end

### Static parameters and container for use.
f_c = 20
x_max = 1.0
num_x = 10*f_c
filter_x =  ones(2*f_c+1);
static_parameters = StaticParameters(f_c, x_max, num_x, filter_x)

model_static = SuperResModels.Fourier1d(static_parameters.f_c, static_parameters.x_max, static_parameters.filter_x, static_parameters.num_x)
### Dynamic parameters and container for use.
K = 2
tau = 1.0/(K*1.0)
# the maximum velocity on the domain is given by x_max/2 + K tau v_max = 1
v_max = 0.5
num_v = 10
dynamic_parameters = DynamicParameters(K, tau, v_max, num_v)
model_dynamic = SuperResModels.DynamicFourier1d(model_static, dynamic_parameters.v_max, dynamic_parameters.tau, dynamic_parameters.K, dynamic_parameters.num_v*dynamic_parameters.K)

# We check the dimensions of each case.
(thetas, weights) = TestCases.cloud_1d_full(x_max, v_max, 0.1,10,K, tau, 1)
println(" thetas = ", thetas)
println(" weights = ", weights)

# Size of the dynamic reconstruction
target_dynamic = Utils.generate_target(model_dynamic, thetas, weights)
println(" Size of the dynamic reconstruction is: ", size(target_dynamic) )

# Size of a single static reconstruction
target_static = Utils.target_to_static(target_dynamic, model_dynamic.times)
#for k = 1: length(model_dynamic.times)
#thetas_t = Utils.to_static(thetas, model_dynamic.times[k], model_static.x_max)
#push!(target_static, Utils.generate_target(model_static,thetas_t, weights))
#println(" Size of a single static reconstruction is: ", size(target_static[k]) )
#end

res = zeros(size(target_dynamic))
for k = 1:length(model_dynamic.times)
    for j = 1:length(target_static[1])
	index = (k-1)*length(target_static[1]) + j;
        res[index] = target_dynamic[index] - target_static[k][j]
	println("res = ", res[(k-1)*length(target_static[1]) + j], "target static = ", target_static[k][j], "target dynamic = ", target_dynamic[(k-1)*length(target_static[1]) + j] )
    end
end
println(" The difference is: ", norm(res)/norm(target_dynamic))
#println(" Target dynamic = ", target_dynamic)
#println(" Target static = ", target_static)


using DSP
include("extend_to_dc.jl")
include("make_conjugate_symmetric.jl")
include("gen_impulse_response.jl")
include("clean_array.jl")
include("gen_step_response.jl")

function s_parameter_to_time_domain(freqs,s11; kb_beta=6)
   # function [timePts, step_response] = SParams2TDR(Freq,S11,varargin)
   # SParams2TDR takes in measured frequency data and outputs a step response

   # prepare s parameter for transformation
   res = extend_to_dc(freqs,s11);
   FreqToDC = res[1]; S11toDC = res[2]
   doubleS11 = make_conjugate_symmetric(S11toDC); # runs from DC -> - deltaF
   doubleFreq = vcat(FreqToDC, -flipdim(FreqToDC[2:end],1));

   filter_vals = kaiser(length(doubleS11),kb_beta);

   doubleS11 = flipdim(fftshift(filter_vals),1) .* doubleS11;
   res = gen_impulse_response(doubleFreq, doubleS11);
   timePts = res[1]; imp_response = res[2];
   # Generate impulse response
   imp_response = clean_array(imp_response,threshold=1e-6);
   # Check if it's real
   if ~isreal(imp_response)
      warn("Impulse response is not real. Results may not make sense.");
   end
   # Generate step response
   res = gen_step_response(timePts, imp_response);
   step_response = res[2]

   # Convert the step response to a vector of floats (if possible)
   step_response = convert(Vector{Float64}, step_response)
   println("Step response type: ", typeof(step_response))
   return (timePts, step_response)
end


module AdDynamics
    using DifferentialEquations
    using ForwardDiff
    using DiffResults
    using StaticArrays
    using ..ProbInfo, ..LinPoint
    struct LinRes
        endpoint::SArray{Tuple{14},Float64,1,14}
        derivative::SArray{Tuple{14,21},Float64,2,294}
    end

    const mass_idx = 1
    const r_idx = SVector(2,3,4)
    const v_idx = SVector(5,6,7)
    const qbi_idx = SVector(8,9,10,11)
    const omb_idx = SVector(12,13,14)

    struct IntegratorParameters{T}
        dt :: Float64
        uk :: SArray{Tuple{3}, T, 1, 3}
        up :: SArray{Tuple{3}, T, 1, 3}
        sigma :: T
        pinfo :: ProbInfo
    end

    function DCM(quat::SArray{Tuple{4}, T, 1, 4} where T)
        q0 = quat[1]
        q1 = quat[2]
        q2 = quat[3]
        q3 = quat[4]

        return transpose(@SMatrix [(1-2*(q2^2 + q3^2))         (2*(q1 * q2 + q0 * q3))  (2*(q1*q3 - q0 * q2));
                        (2*(q1 * q2 - q0 * q3))     (1-2*(q1^2 + q3^2))      (2*(q2*q3 + q0 * q1));
                        (2*(q1 * q3 + q0 * q2))     (2*(q2*q3 - q0 * q1))    (1-2*(q1^2 + q2^2)) ])
    end

    function Omega{T}(omegab::SArray{Tuple{3}, T, 1, 3})
        z = zero(T)
        return @SMatrix [z        -omegab[1] -omegab[2] -omegab[3];
                        omegab[1] z           omegab[3] -omegab[2];
                        omegab[2] -omegab[3]  z          omegab[1];
                        omegab[3]  omegab[2] -omegab[1]  z         ]
    end

    function dx(state::SArray{Tuple{14}, T, 1, 14} where T, u::SArray{Tuple{3}, T, 1, 3} where T, info::ProbInfo)
        qbi = state[qbi_idx]
        omb = state[omb_idx]
        thr_acc = DCM(qbi) * u/state[mass_idx]
        grv_acc = @SVector [-info.g0,0,0]
        acc = thr_acc + grv_acc
        rot_vel = 0.5*Omega(omb)*qbi
        rot_acc = info.jBi*(cross(info.rTB,u) - cross(omb,info.jB*omb))
        return @SVector [-info.a*norm(u),
                state[5],state[6],state[7],
                acc[1], acc[2], acc[3],
                rot_vel[1], rot_vel[2], rot_vel[3], rot_vel[4],
                rot_acc[1], rot_acc[2], rot_acc[3]]
    end

    function dynamics_sim(state::SArray{Tuple{14}, T, 1, 14} where T, p::IntegratorParameters, t)
        lkm = (p.dt-t)/p.dt
        lkp = t/p.dt
        u = p.uk*lkm + p.up*lkp
        p.sigma*dx(state,u,p.pinfo)
    end

    macro genIdx(start,last)
        return :(@SVector [i for i=$start:$last])
    end
    const stateC = @genIdx(1,14)
    const ukC = @genIdx(15,17)
    const upC = @genIdx(18,20)
    const sigmaC = 21
    function fdiff(pinfo,dt)
        function f(p)
            state = p[stateC]
            uk = p[ukC]
            up = p[upC]
            sigma = p[sigmaC]
            prob = ODEProblem(dynamics_sim, state, eltype(p).((0.0,dt)), IntegratorParameters(dt,uk,up,sigma,pinfo))
            DifferentialEquations.solve(prob, Tsit5(),reltol=1e-8,abstol=1e-8,force_dtmin=true,dtmin=0.0001,save_everystep=false,saveat=[dt])[end]
        end
        return f
    end

    function linearize_segment(fun, sigma_lin::Float64, istate::SVector{14}, c1::SVector{3}, c2::SVector{3})
        ip = SVector{21}(vcat(istate, c1, c2, sigma_lin))
        res = DiffResults.JacobianResult((@MVector zeros(14)),MVector(ip))
        ForwardDiff.jacobian!(res, fun, ip)
        jcb = DiffResults.jacobian(res)
        return LinRes(DiffResults.value(res),SMatrix{14,21}(jcb))
    end

    function linearize_dynamics(states::Array{LinPoint,1}, sigma_lin::Float64, dt::Float64, info::ProbInfo)
        results = Array{LinRes,1}(length(states)-1)
        fun = fdiff(info,dt)
        for i=1:length(states)-1
            cstate = states[i]
            nstate = states[i+1]
            results[i] = linearize_segment(fun, sigma_lin, cstate.state, cstate.control, nstate.control)
        end
        return results
    end
    function next_step(dynam, ab, abn, state, control_k, control_kp, sigma, sigHat, relax)
        ctrl = vcat(state-ab.state, control_k-ab.control, control_kp-abn.control,sigma-sigHat)
        return dynam.derivative *ctrl + dynam.endpoint + relax
    end
    export linearize_dynamics, next_step, LinRes
end


using ApproxFunOrthogonalPolynomials, ApproxFunBase, BlockBandedMatrices,  LinearAlgebra, Test
import ApproxFunBase: Multiplication,InterlaceOperator, Block, ∞,
        testfunctional, testbandedoperator, testraggedbelowoperator, testinfoperator, testblockbandedoperator
import ApproxFunOrthogonalPolynomials: JacobiZ

@testset "Operator" begin
    @testset "Evaluation" begin
        testfunctional(Evaluation(Ultraspherical(1),0.1))
        d = -4 .. 4
        f = Fun(exp, Ultraspherical(1,d))
        @test f(-4) ≈ Number(ldirichlet()*f) ≈ Number(Evaluation(Ultraspherical(1,d),-4)*f)

        testfunctional(Evaluation(Chebyshev(),0.1,1))
        testfunctional(Evaluation(Chebyshev(),0.1,1)-Evaluation(Chebyshev(),0.1,1))

        let f = Fun(cos)
            @test (Evaluation(Chebyshev(),0.1,1)*f)(0.1)  ≈ f'(0.1)
        end
    end

    @testset "Conversion" begin
        C=Conversion(Ultraspherical(1),Ultraspherical(2))
        testbandedoperator(C)

        @test Matrix(C[1:5,1:5])  ≈   [1.0 0.0 -0.3333333333333333 0.0  0.0
                                              0.0 0.5  0.0               -0.25 0.0
                                              0.0 0.0  0.3333333333333333 0.0 -0.2
                                              0.0 0.0  0.0                0.25 0.0
                                              0.0 0.0  0.0                0.0  0.2]
        d=ChebyshevInterval()
        A=Conversion(Chebyshev(d),Ultraspherical(2,d))
        x = Fun()
        f=Fun(exp)
        @test AbstractMatrix(view(A.op, Block.(1:3), Block.(1:3))) isa BlockSkylineMatrix
        testbandedoperator(A)

        @test norm(A\Fun(x.*f,rangespace(A))-(x.*f)) < 100eps()

        @test (I : Chebyshev() → Chebyshev()) * Fun() ≈ Fun()
    end

    @testset "Derivative" begin
        testbandedoperator(Derivative(Ultraspherical(1)))

        d=Interval(-10.,5.);
        S=Chebyshev(d)

        @test norm(Fun(Fun(Fun(exp,S),Ultraspherical(2,d)),S)-Fun(exp,S)) < 100eps()

        @test copy(view(Derivative(Ultraspherical(1)),1:2,1:2))[1,2] ≈ Derivative(Ultraspherical(1))[1,2]
        @test exp(0.1) ≈ (Derivative()*Fun(exp,Ultraspherical(1)))(0.1)
        D=Derivative(d)
        f = Fun(exp,d)
        @test norm((Conversion(Chebyshev(d),Ultraspherical(2,d))\(D^2*f))-f'') < 100eps()

        f=Fun(x->x^2)
        D=Derivative(domain(f))
        @test norm(D*f-f')<100eps()
    end

    @testset "Toeplitz" begin
        C=ToeplitzOperator([1.,2.,3.],[4.,5.,6.])

        @time testbandedoperator(C)

        @test Matrix(C[1:5,1:5])  ≈  [4.0 5.0 6.0 0.0  0.0
                                             1.0 4.0 5.0 6.0 0.0
                                             2.0 1.0 4.0 5.0 6.0
                                             3.0 2.0 1.0 4.0 5.0
                                             0.0 3.0 2.0 1.0 4.0]

         testbandedoperator(HankelOperator([1.,2.,3.,4.,5.,6.,7.]))
    end

    @testset "Multiplication" begin
        testbandedoperator(Multiplication(Fun(Chebyshev(),[1.,2.,3.]),Chebyshev()))

        x=Fun(identity)
        X=Multiplication(x,space(x))
        f = Fun(exp)
        testbandedoperator(X)

        @test norm(X*f-(x.*f)) < 100eps()

        A=Conversion(Chebyshev(),Ultraspherical(2))*X
        @time testbandedoperator(A)

        @test norm((mul_coefficients(A,f.coefficients))-coefficients(x.*f,rangespace(A))) < 100eps()
    end

    @testset "Integral" begin
        f=Fun(exp)
        d=domain(f)
        Q=Integral(d)
        D=Derivative(d)

        @time testbandedoperator(Q)

        @test norm((Q+I)*f-(integrate(f)+f)) < 100eps()
        @test norm((Q)*f-(integrate(f))) < 100eps()
    end

    @testset "Special functions" begin
        x=Fun(identity)
        @test norm(cos(x)-Fun(cos))<10eps()
        @test norm(sin(x)-Fun(sin))<10eps()
        @test norm(exp(x)-Fun(exp))<10eps()
        @test norm(sin(x)./x-Fun(x->sinc(x/π)))<100eps()
    end

    @testset "Permutation" begin
        P=ApproxFunBase.PermutationOperator([2,1])
        testbandedoperator(P)
        @test P[1:4,1:4] ≈ [0 1 0 0; 1 0 0 0; 0 0 0 1; 0 0 1 0]
    end

    @testset "Mixed" begin
        d = ChebyshevInterval()
        D = Derivative(d)
        x = Fun(identity,d)
        A = D*(x*D)
        B = D+x*D^2
        C = x*D^2+D

        testbandedoperator(A)
        testbandedoperator(B)
        testbandedoperator(C)
        @time testbandedoperator(x*D)

        f=Fun(exp)
        @test (A.ops[end]*f)(0.1) ≈ f'(0.1)
        @test ((x*D)*f)(0.1) ≈ 0.1*f'(0.1)
        @test (A*f)(0.1) ≈ f'(0.1)+0.1*f''(0.1)
        @test (B*f)(0.1) ≈ f'(0.1)+0.1*f''(0.1)
        @test (C*f)(0.1) ≈ f'(0.1)+0.1*f''(0.1)

        testbandedoperator(A-B)
        testbandedoperator(B-A)
        testbandedoperator(A-C)

        @test norm((A-B)[1:10,1:10]|>Matrix) < eps()
        @test norm((B-A)[1:10,1:10]|>Matrix) < eps()
        @test norm((A-C)[1:10,1:10]|>Matrix) < eps()
        @test norm((C-A)[1:10,1:10]|>Matrix) < eps()
        @test norm((C-B)[1:10,1:10]|>Matrix) < eps()
        @test norm((B-C)[1:10,1:10]|>Matrix) < eps()
    end

    @testset "Cached" begin
        @test cache(Derivative(Chebyshev(),2))[1,1] == 0

        S=Chebyshev()
        D=Derivative(S)
        @time for padding = [true,false]
          co=ApproxFunBase.CachedOperator(D,ApproxFunBase.RaggedMatrix(Float64[],Int[1],0),(0,0),domainspace(D),rangespace(D),bandwidths(D),padding) #initialise with empty RaggedMatrix
          @test co[1:20,1:10] == D[1:20,1:10]
          @test size(co.data) == (20,10)
          ApproxFunBase.resizedata!(co,10,30)
          @test size(co.data)[2] == 30 && size(co.data)[1] ≥ 20
        end
    end

    @testset "Reverse" begin
        testbandedoperator(ApproxFunBase.Reverse(Chebyshev()))
        testbandedoperator(ApproxFunBase.ReverseOrientation(Chebyshev()))

        @test ApproxFunBase.Reverse(Chebyshev())*Fun(exp) ≈ Fun(x->exp(-x))
        @test ApproxFunBase.ReverseOrientation(Chebyshev())*Fun(exp) ≈ Fun(exp,Segment(1,-1))
    end

    @testset "Sub-operator re-view bug" begin
        D = Derivative(Chebyshev())
        S = view(D[:, 2:end], Block.(3:4), Block.(2:4))
        @test parent(S) == D
        @test parentindices(S) == (3:4,2:4)
        @test bandwidths(S)  == (-2,2)
    end

    @testset "Sub-operators" begin
        f = Fun(exp)
        D = Derivative(Chebyshev())
        testbandedoperator(D[:, 2:end])

        u = D[:,2:end] \ f
        @test u(0.1) ≈ exp(0.1)-coefficient(f,1)

        D̃ = Derivative(space(u))
        @test bandwidths(D̃) == (0,0)
        testbandedoperator(D̃)
        @test norm(u'-f) < 10eps()

        testbandedoperator(D[1:end,2:end])
        u = D[1:end,2:end] \ f
        @test u(0.1) ≈ exp(0.1)-f.coefficients[1]

        testbandedoperator(D[1:∞,2:∞])
        u = D[1:∞,2:∞] \ f
        @test u(0.1) ≈ exp(0.1)-f.coefficients[1]
    end

    @testset "InterlaceOperator" begin
        A = InterlaceOperator(Diagonal([Matrix(I,2,2),Derivative(Chebyshev())]))

        @test A[Block(1):Block(2), Block(1):Block(2)] isa BlockBandedMatrix
        @test Matrix(view(A, Block(1), Block(1))) == A[1:3,1:3]
        @test Matrix(view(A, Block(1):Block(2), Block(1):Block(2))) == A[1:4,1:4]
        testblockbandedoperator(A)
    end

    @testset "Projection and subspaces" begin
        S=Chebyshev()
        SS = S|(2:5)
        @test ApproxFunBase.block(SS,3) == Block(4)

        for C in (Operator(I,S)[3:end,:], Operator(I,S)[3:end,1:end])
            @test ApproxFunBase.domaindimension(domainspace(C)) == 1
            @test union(S,domainspace(C)) == S

            B=Dirichlet(S)

            Ai=[B;C]

            @test ApproxFunBase.colstop(Ai,1) == 2

            x=Fun()
            f=exp(x)
            u=[B;C]\[[0.,0.],f]

            @test abs(u(-1)) ≤ 10eps()
            @test abs(u(1)) ≤ 10eps()

            f=(1-x^2)*exp(x)
            u=[B;C]\[[0.,0.],f]

            @test u ≈ f
        end
    end

    @testset "Zero operator has correct bandwidths" begin
        Z=ApproxFunBase.ZeroOperator(Chebyshev())
        @test ApproxFunBase.bandwidths(Z) == ApproxFunBase.bandwidths(Z+Z)
    end

    @testset "hcat of functionals (#407)" begin
        x = Fun()
        B = [1 ldirichlet()]
        @test (B*[1;x])[1] == Fun(ConstantSpace(ApproxFunBase.Point(-1.0)),[0.0])
    end

    @testset "views of views" begin
        A = Derivative(Chebyshev()) + I
        B = A[1:2:∞,1:2:∞]
        C = B[2:∞,3:∞]
        @test A[3:2:∞,5:2:∞] == C
    end

    @testset "Multiplication functions" begin
        x = Fun()
        M = Multiplication(x, Chebyshev())
        @test exp(M).f == Multiplication(exp(x), Chebyshev()).f
        @test (-M).f == Multiplication(-x, Chebyshev()).f
        @test exp(-M).f == Multiplication(exp(-x), Chebyshev()).f
        @test (M/3).f == (3\M).f == Multiplication(x/3, Chebyshev()).f
        @test (M*3).f == (3*M).f == Multiplication(x*3, Chebyshev()).f
    end

    @testset "lastindex" begin
        Z = Operator(I,Chebyshev())
        S = view(Z,1:10,1:10)
        @test lastindex(S,1) == lastindex(S,2) == 10
        @test lastindex(S) == 100
        @test S[end,end] ≈ 1
        @test S[end-1,end] ≈ 0
    end

    @testset "Chebyshev Neumann" begin
        E = Evaluation(Chebyshev(),rightendpoint,1)
        @test E[1:4] == [0,1,4,9]
    end

    @testset "Chebyshev Dirichlet" begin
        S=Chebyshev()
        @time for io in (
                [Dirichlet(S);Derivative(Chebyshev());lneumann(S)],
                [Dirichlet(S);Derivative(Chebyshev())+Fun(cos);lneumann(S)],
                [Dirichlet(S);Derivative(Chebyshev())],
                [Dirichlet(S);Derivative(Chebyshev())+Fun(cos)],
                [Derivative(Chebyshev());Dirichlet(S)],
                [Derivative(Chebyshev())+Fun(cos);Dirichlet(S)])
            testraggedbelowoperator(io)
        end
    end

    @testset "JacobiZ" begin
        testbandedoperator(JacobiZ(Chebyshev(), 0.1))
    end
end


function addto!(jm::ModelView, m::EnergyModel, sn::SubNetwork)
    T = axis(sn.model, :snapshots)
    branches = sn[PassiveBranch]
    B = axis(branches)
    length(B) == 0 && return

    p = get(branches, :p, B, T)
    effimp = mapcat(effectiveimpedance, branches)

    C = cycle_matrix(sn)
    Cl = rowvals(C)
    Cv = nonzeros(C)

    @constraint(jm, cycles[c=1:size(C,2),t=T], sum(Cv[j] * effimp[Cl[j]] * p[Cl[j],t] for j=nzrange(C,c)) == 0)
end

addcomponent(SubNetwork, :subnetworks)


py"""

import numpy as np
import scipy.stats as sst

def bestcent(inpcent):
    '''
    Calculate the best peak and centroid and their uncertainties using the median of the
    distributions
    '''
    perclim = 84.1344746
    centau = sst.scoreatpercentile(inpcent, 50)
    centau_uperr = (sst.scoreatpercentile(inpcent, perclim))-centau
    centau_loerr = centau-(sst.scoreatpercentile(inpcent, (100.-perclim)))
    #print('Centroid, error: %10.3f  (+%10.3f -%10.3f)'%(centau, centau_loerr, centau_uperr))
    return centau,centau_loerr,centau_uperr

"""

function centerror(centdistr)
    centau,centau_loerr,centau_uperr = py"bestcent($centdistr)"
    return centau,centau_loerr,centau_uperr
end


using LinearAlgebra
using KrylovKit
using OMEinsum
using IterTools
"""
tensor order graph: from left to right, top to bottom.
```
a ────┬──── c    a──────┬──────b
│     b     │    │      │      │
├─ d ─┼─ e ─┤    │      c      │
│     g     │    │      │      │
f ────┴──── h    d──────┴──────e
```
"""

"""
    λ, FL = leftenv(Au, Ad, M, FL = _arraytype(Au)(rand(eltype(Au), size(Au,1), size(M,1), size(Au,1))); kwargs...)

Compute the left environment tensor for MPS `AL` and MPO `M`, by finding the left fixed point
of `Au - M - Ad` contracted Aung the physical dimension.
```
┌──  Au─       ┌──       a ────┬──── c
│    │         │         │     b     │
FL ─ M ─  = λL FL─       ├─ d ─┼─ e ─┤
│    │         │         │     g     │
┕──  Ad─       ┕──       f ────┴──── h
```
"""
function leftenv(Au, Ad, M, FL = _arraytype(Au)(rand(eltype(Au), size(Au,1), size(M,1), size(Ad,1))); kwargs...)
    λs, FLs, info = eigsolve(FL -> ein"((adf,abc),dgeb),fgh -> ceh"(FL,Au,M,Ad), FL, 1, :LM; ishermitian = false, kwargs...)
    if length(λs) > 1 && norm(abs(λs[1]) - abs(λs[2])) < 1e-12
        @show λs
        if real(λs[1]) > 0
            return λs[1], FLs[1]
        else
            return λs[2], FLs[2]
        end
    end
    return λs[1], FLs[1]
end

"""
    λ, FR = rightenv(Au, Ad, M, FR = _arraytype(Au)(randn(eltype(Au), size(Au,1), size(M,3), size(Au,1))); kwargs...)

Compute the right environment tensor for MPS `AR` and MPO `M`, by finding the right fixed point
of `Au - M - Ad` contracted Aung the physical dimension.
```
 ─ Au ──┐         ──┐      a ────┬──── c
    │   │           │      │     b     │
 ─  M ──FR   = λR ──FR     ├─ d ─┼─ e ─┤
    │   │           │      │     g     │
 ─ Ad ──┘         ──┘      f ────┴──── h
```
"""
function rightenv(Au, Ad, M, FR = _arraytype(Au)(randn(eltype(Au), size(Au,1), size(M,3), size(Ad,1))); kwargs...)
    Au = permutedims(Au,(3,2,1))
    Ad = permutedims(Ad,(3,2,1))
    ML = permutedims(M,(3,2,1,4))
    return leftenv(Au, Ad, ML, FR; kwargs...)
end

"""
    λ, FL = norm_FL(Au, Ad, FL; kwargs...)

Compute the left environment tensor for normalization, by finding the left fixed point
of `Au - Ad` contracted Aung the physical dimension.
```
┌── Au─      ┌──         a──────┬──────b
FL  │  =  λL FL                 │
┕── Ad─      ┕──                c
                                │
                         d──────┴──────e
```
"""
function norm_FL(Au, Ad, FL = _arraytype(Au)(rand(eltype(Au), size(Au,1), size(Ad,1))); kwargs...)
    λs, FLs, info = eigsolve(FL -> ein"(ad,acb),dce -> be"(FL,Au,Ad), FL, 1, :LM; ishermitian = false, kwargs...)
    if length(λs) > 1 && norm(abs(λs[1]) - abs(λs[2])) < 1e-12
        @show λs
        if real(λs[1]) > 0
            return λs[1], FLs[1]
        else
            return λs[2], FLs[2]
        end
    end
    return λs[1], FLs[1]
end

"""
    λ, FR = norm_FR(Au, Ad, FR; kwargs...)

Compute the right environment tensor for normalization, by finding the right fixed point
of `Au - Ad` contracted Aung the physical dimension.
```
 ─ AR──┐       ──┐       a──────┬──────b
   │   FR  = λR  FR             │
 ─ AR──┘       ──┘              c
                                │
                         d──────┴──────e
```
"""
function norm_FR(Au, Ad, FR = _arraytype(Au)(randn(eltype(Au), size(Au,3), size(Ad,3))); kwargs...)
    Au = permutedims(Au,(3,2,1))
    Ad = permutedims(Ad,(3,2,1))
    return norm_FL(Au, Ad, FR; kwargs...)
end


import Base.+, Base.-

export +, -

# Unary Negation

function -(x::Constant)
  return Constant(-x.value)
end

function -(x::AffineExpr)
  vars_to_coeffs_map = Dict{UInt64, Constant}()
  for (v, c) in x.vars_to_coeffs_map
    vars_to_coeffs_map[v] = -c
  end
  this = AffineExpr(:-, (x,), vars_to_coeffs_map, -x.constant, x.size)
  this.evaluate = ()->(-x.evaluate())
  return this
end

# Binary Addition

function +(x::Constant, y::Constant)
  if x.size == (1, 1) || y.size == (1, 1) || x.size == y.size
    return convert(Constant, x.value + y.value)
  else
    error("Cannot add two expressions with sizes $(x.size) and $(y.size)")
  end
end

function +(x::AffineExpr, y::Constant)
  if x.size != (1, 1) && y.size != (1, 1) && x.size != y.size
    error("Cannot add two expressions with sizes $(x.size) and $(y.size)")
  end

  if x.size == (1, 1) && y.size != (1, 1)
    x = repmat(x, y.size[1], y.size[2])
  end

  vars_to_coeffs_map = copy(x.vars_to_coeffs_map)
  constant = x.constant + vec(y)
  this = AffineExpr(:+, (x, y), vars_to_coeffs_map, constant, x.size)
  this.evaluate = ()->x.evaluate() + y.evaluate()
  return this
end

function +(x::Constant, y::AffineExpr)
  return y + x
end

function +(x::AffineExpr, y::AffineExpr)
  if x.size != (1, 1) && y.size != (1, 1) && x.size != y.size
    error("Cannot add two expressions with sizes $(x.size) and $(y.size)")
  end

  if x.size == (1, 1) && y.size != (1, 1)
    x = repmat(x, y.size[1], y.size[2])
  elseif y.size == (1, 1) && x.size != (1, 1)
    y = repmat(y, x.size[1], x.size[2])
  end

  vars_to_coeffs_map = copy(x.vars_to_coeffs_map)
  for (v, c) in y.vars_to_coeffs_map
    if haskey(vars_to_coeffs_map, v)
      vars_to_coeffs_map[v] = vars_to_coeffs_map[v] + c
    else
      vars_to_coeffs_map[v] = c
    end
  end
  constant = x.constant + y.constant
  this = AffineExpr(:+, (x, y), vars_to_coeffs_map, constant, x.size)
  this.evaluate = ()->x.evaluate() + y.evaluate()
  return this
end

+(x::AffineExpr, y::Numeric) = +(x, convert(Constant, y))
+(x::Numeric, y::AffineExpr) = +(y, convert(Constant, x))

function +(x::SumSquaresExpr, y::SumSquaresExpr)
  affines = copy(x.affines)
  multipliers = copy(x.multipliers)
  append!(affines, y.affines)
  append!(multipliers, y.multipliers)
  this = SumSquaresExpr(:+, affines, multipliers, x.scalar + y.scalar)
  return this
end

function +(x::Number, y::SumSquaresExpr)
  try
    x = convert(Float64, x)
  catch
    error("Only real scalars can be added to Sum Squares expressions")
  end
  if x < 0
    error("Only nonnegative scalars can be added to Sum Squares expressions")
  end
  this = SumSquaresExpr(:+, y.affines, y.multipliers, y.scalar + x)
  return this
end

+(x::SumSquaresExpr, y::Number) = +(y, x)

# Binary Subtraction
-(x::AffineExpr, y::AffineExpr) = +(x, -y)
-(x::AffineExpr, y::Constant) = +(x, -y)
-(x::Constant, y::AffineExpr) = +(-y, x)
-(x::AffineExpr, y::Numeric) = -(x, convert(Constant, y))
-(x::Numeric, y::AffineExpr) = -(convert(Constant, x), y)


using Test
using StaticArrays
using Random
using Nystrom
using LinearAlgebra
Random.seed!(1)

@testset "Density test" begin
    Geometry.clear_entities!()
    pde   = Helmholtz(;dim=3,k=1)
    Ω     = ParametricSurfaces.Sphere()
    Γ     = boundary(Ω) |> Geometry.Domain
    M     = ParametricSurfaces.meshgen(Γ,(10,10))
    mesh  = NystromMesh(view(M,Γ),order=1)
    σ     = Density(target->norm(coords(target)),mesh)
    @test eltype(σ) == Float64
    σ     = Density(mesh) do target
        x = coords(target)
        exp(im*2*norm(x))
    end
    @test eltype(σ) == ComplexF64
end

@testset "TangentialDensity test" begin
    Geometry.clear_entities!()
    Ω   = ParametricSurfaces.Sphere(;radius=1) |> Geometry.Domain
    Γ   = boundary(Ω)
    M   = ParametricSurfaces.meshgen(Γ,(2,2))
    mesh = NystromMesh(view(M,Γ),order=2)
    pde = Elastostatic(;dim=3,μ=2,λ=3)
    T = Nystrom.default_density_eltype(pde)
    xout = SVector(3,3,3)
    c = rand(T)

    σ = Density(mesh) do target  # density defined with a tangential field
        x = coords(target)
        cross(normal(target), SingleLayerKernel(pde)(xout,target)*c)
    end
    tan_σ = TangentialDensity(σ)
    @test Density(tan_σ) ≈ σ

    ncross_σ = ncross(σ)
    ncross_tan_σ = ncross(tan_σ)
    @test TangentialDensity(ncross_σ) ≈ ncross_tan_σ
    @test Density(ncross_tan_σ) ≈ ncross_σ
end


## Define a particle consisting of all cluster variables for a Gaussian mixture
## N is no. of potential clusters
## d is dimension of data
## n_obs is the number of observations
## n_levels is the number of levels, used in multinomial
## Priors:
# β_0 = 0.5
# α_0 = 1
# κ_0 = 0.001
# μ_0 = 0
mutable struct GaussianCluster
  n::Int64
  μ::Vector{Float64}   ## mean of observations in clusters
  Σ::Vector{Float64}   ## sum of observations in clusters
  λ::Vector{Float64}
  β::Vector{Float64}
  GaussianCluster(dataFile) =         new(0,
                                      zeros(Float64, size(dataFile, 2)),
                                      zeros(Float64, size(dataFile, 2)),
                                      ones(Float64, size(dataFile, 2)),
                                      ones(Float64, size(dataFile, 2)) .* 0.5)
end

function copy_particle(particle::GaussianCluster, dataFile)
    new_particle = GaussianCluster(dataFile)
    new_particle.n = particle.n
    for i in eachindex(particle.μ)
      new_particle.μ[i] = particle.μ[i]
      new_particle.Σ[i] = particle.Σ[i]
      new_particle.λ[i] = particle.λ[i]
      new_particle.β[i] = particle.β[i]
    end
    return new_particle
end

function calc_logprob(obs, cl::GaussianCluster, featureFlag)
    @fastmath out = sum(featureFlag) * (log(1 / sqrt(pi)) +
                          SpecialFunctions.loggamma(0.5 * cl.n + 1.0) -
                          SpecialFunctions.loggamma(0.5 * cl.n + 0.5))
    # Iterate over features
      for q in 1:size(obs, 1)
        if featureFlag[q]
        out += 0.5 * (log(cl.λ[q] / (cl.n + 1.0)))
        @inbounds out -= (0.5 * cl.n + 1.0) *
              log(1.0 + (1.0 / (cl.n + 1.0)) *
                  ((obs[q] - cl.μ[q]) ^ 2.0) * cl.λ[q])
        end
      end
      return out
  # """
end

function cluster_add!(cl::GaussianCluster, obs, featureFlag)
  cl.n     += 1
  @inbounds  for q = 1:length(obs)
    if featureFlag[q]
      cl.Σ[q]  += obs[q]
      cl.β[q] += (cl.n - 1 + 0.001) * (obs[q] - cl.μ[q]) ^ 2 / (2 * (cl.n + 0.001))
      cl.μ[q]  = cl.Σ[q] / (cl.n + 0.001)
      cl.λ[q]    = ((0.5 * cl.n + 0.5) * (cl.n + 0.001)) /
                      (cl.β[q] * (cl.n + 1.001))
    end
  end
  return
end

function calc_logmarginal(cl::GaussianCluster)
  # This returns the log of the marginal likelihood of the cluster
  # Used for feature selection
  α_n = (cl.n / 2 + 0.5)
  α_0 = 0.5
  β_n = cl.β
  β_0 = 0.5
  κ_0 = 0.001
  κ_n = cl.n + κ_0

  lm =  - α_n * log.(β_n) .+ ((α_0 * log(β_0)) .+
        SpecialFunctions.loggamma(α_n) - SpecialFunctions.loggamma(α_0) +
        0.5 * (log(κ_0) - log(κ_n)) -
        (cl.n * 0.5) * log(2 * π))
  return lm
end

function gaussian_normalise!(dataFile::Array)
  for d = 1:size(dataFile, 2)
    # μ = mean(dataFile[:, d])
    # σ = std(dataFile[:, d]) + eps(Float64)
    μ = median(dataFile[:, d])
    σ = 0.5 * (μ - quantile(dataFile[:, d], 0.05)) + eps(Float64)
    dataFile[:, d] = (dataFile[:, d] .- μ) ./ σ
  end
  return
end


using TomoForward
using Images
using Plots
using XfromProjections
using StaticArrays
using PyCall
using Logging

include("../tv_primaldual_flow.jl")
include("../optical_flow.jl")
include("../../phantoms/sperm_phantom.jl")

H,W = 609, 609
detmin, detmax = -38.0, 38.0
function radon_operator(height, width, detcount, angles)
    proj_geom = ProjGeom(0.5, detcount, angles)
    A = fp_op_parallel2d_line(proj_geom, height, width, detmin,detmax, detmin,detmax)
    return A
end
grid = collect(detmin:0.1:detmax)
r(s) = 1.0

images, tracks = get_sperm_phantom(301,r,grid)
frames = images[:,:,collect(271:274)]
@info size(frames)

w_tv = 0.01
w_flow  = 0.001
c=10.0

detcount = Int(floor(H*1.4))

cwd = @__DIR__
path = normpath(joinpath(@__DIR__, "result"))
!isdir(path) && mkdir(path)
cd(path)

#4 angles evenly spaced between 0 and pi
angles = [rand(0.0:0.001:π, 3), rand(0.0:0.001:π, 3), rand(0.0:0.001:π, 3), rand(0.0:0.001:π, 3)]#collect(range(π/2,3*π/2,length=ang_nr+1))[1:end-1]
@info angles
#match size of input image (generating data)
As = map(t -> radon_operator(size(frames[:,:,1],1),size(frames[:,:,1],2),detcount, angles[t]),1:size(frames)[3])
bs = zeros(size(As[1])[1],size(frames)[3])
map(t -> bs[:,t] = As[t]*vec(frames[:,:,t]), 1:size(frames)[3])

u0s = zeros(H,W,size(frames)[3])

As  = map(t -> radon_operator(H,W, detcount, angles[t]),1:size(u0s)[3])

@info "Reconstruction using tv regularization frame by frame"
niter=450
us_tv = zeros(H,W,size(frames)[3])
for t = 1:size(frames)[3]
    A = As[t]
    p = bs[:,t]
    u0 = u0s[:,:,t]
    us_tv[:,:,t] .= recon2d_tv_primaldual!(us_tv[:,:,t], A, p, niter, w_tv, c)
end

img2 = us_tv[:,:,2]#zeros(128,128)
img1 = us_tv[:,:,1]
H, W = size(img2)
#img2[30:60,40:50] .= 1.0

#flow = zeros(H, W, 2)
#flow[:,:, 1] .= -13.0
#flow[:,:, 2] .= 10.0
#flow[:,:, 1] +=rand(H,W)*4
#flow[:,:, 2] +=rand(H,W)*4

#Wop = compute_warping_operator(flow)
#img1_ = Wop * vec(img2)
#img1 = reshape(img1_, H, W)

us = zeros(H,W,2)
us[:,:,2] = img2
us[:,:,1] = img1

v = get_flows(us)

W_list = mapslices(f -> compute_warping_operator(f), v,dims=[1,2,3])
@info size(W_list)
#t = 1
#Wuv .= mul!(Wuv, W_list[t], vec(us[:,:,t+1])) .- vec(us[:,:,t])

#flow_x, flow_y, im2Warped = py"py_flow"(img1,img2)
#v = zeros(H,W,2)
#v[:,:,1] = flow_x
#v[:,:,2] = flow_y#W_list = mapslices(f -> compute_warping_operator(f), v, dims=[1,2,3])
#Woptical = W_list[1]#compute_warping_operator(v)
#img3_ = Woptical * vec(img2)
#img3 = reshape(img3_,H, W)
#W_list = mapslices(f -> compute_warping_operator(f), v,dims=[1,2,3])

t = 1
_ubar_1 = view(us,:,:,t+1)
_ubar = view(us,:,:,t)
Wu = W_list[t]*(collect(Iterators.flatten(_ubar_1))) - (collect(Iterators.flatten(_ubar)))
p3_ascent = reshape(Wu, H, W)

p3_adjoint = reshape(W_list[t]'*vec(p3_ascent),H,W)-p3_ascent

img3 = max.(us[:,:,1].-p3_adjoint, 0.0)

p1 = plot(Gray.(img1[:,:,1]), aspect_ratio=:equal, framestyle=:none, title="img1")
p2 = plot(Gray.(img2[:,:,1]), aspect_ratio=:equal, framestyle=:none, title="img2")
p3 = plot(Gray.(img3), aspect_ratio=:equal, framestyle=:none, title="rec")
p4 = plot(Gray.(p3_ascent), aspect_ratio=:equal, framestyle=:none, title="ascent")
p5 = plot(Gray.(p3_adjoint), aspect_ratio=:equal, framestyle=:none, title="adjoint")
l = @layout [a b c d e]
plot(p1, p2, p3, p4, p5, layout = l, size=(2000,600), linewidth=5)
#XX, YY = repeat(collect(1:W)', H, 1), repeat(collect(1:H), 1, W)
#p1 = quiver([(XX...)...][1:20:end], [(YY...)...][1:20:end], quiver=([(v[:,:,1]...)...][1:20:end]*0.001, [(v[:,:,2]...)...][1:20:end]*0.001), title="v")
#p2 = quiver([(XX...)...][1:20:end], [(YY...)...][1:20:end], quiver=([(flow[:,:,1]...)...][1:20:end]*0.1, [(flow[:,:,2]...)...][1:20:end]*0.1), title="flow")
#l = @layout [a b]
#plot(p1, p2, layout = l, size=(2000,600), linewidth=5)


module WeightedNetwork

using SparseArrays
using Graphs
using SimpleWeightedGraphs
using DataStructures
using LinearAlgebra
using Plots
using GraphPlot
using DataStructures
using ShowGraphviz
# using Cairo
# using Fontconfig
using Random
using Tokenize

export
    # generators
    rand_network, rand_directed_network, rand_dag, mk_connected_layered_dag,

    # graph properties
    children, parent, neighbors,

    # dataformat
    adjacency_matrix, AdjacencyList,

    # visualization mit GraphPlot:
    plot_network, plot_graph,

    # visualization mit Graphviz (for small graphs or use file_export:)
    plot_graphviz, write_dot_file, read_dot_file, AttributeDict, default_attributes,

    # algorithm:
    dfs_path, bfs_path, toplogicalsort, node_color, path_to_adj_matrix, shortest_path

include("./data_types/graphdata_types.jl")
include("./generators/random_networks.jl")
include("./visualization/gplots.jl")
include("./visualization/dots.jl")
include("./algorithm/traversel.jl")
include("./algorithm/graph_utils.jl")

end


struct Beta{T<:Real} <: Prior
    pivots::Vector{T}
    weights::Vector{T}
    pdf::Vector{T}
    a::T
    b::T
    low::T
    high::T
end

function Beta(a::Real, b::Real; low::Real = 0, high::Real = 1)
    p, ω = gauss_legendre_25(low, high)
    pdf  = dbeta.(p, convert(Float64, a), convert(Float64, b))
    z    = (high - low)/2 * sum(pdf .* ω)
    pdf  = pdf ./ z # normalize
    return Beta{Float64}(
        convert(Vector{Float64}, p),
        convert(Vector{Float64}, ω),
        convert(Vector{Float64}, pdf),
        convert(Float64, a),
        convert(Float64, b),
        convert(Float64, low),
        convert(Float64, high)
    )
end

# convenience function to fit a and b
function Beta(;mean::Real = .5, sd::Real = sqrt(1/12))
    sd >= .5 ? error("sd must be strictly smaller than .5") : nothing
    # solve for a, b
    a = (-mean^3 + mean^2 - mean*sd^2) / sd^2
    b = (mean - 1) * (mean^2 - mean + sd^2) / sd^2
    Beta(a, b; low = 0, high = 1)
end

function pdf(p::T, prior::Beta{T})::T where {T<:Real}
    if !(prior.low <= p <= prior.high); return 0.0 end
    F(p) = cdf(Distributions.Beta(prior.a, prior.b), p)
    return pdf(Distributions.Beta(prior.a, prior.b), p) / ((F(prior.high) - F(prior.low)))
end

function cdf(p::T, prior::Beta{T})::T where {T<:Real}
    p < prior.low ? (return 0.0) : nothing
    p > prior.high ? (return 1.0) : nothing
    F(p) = cdf(Distributions.Beta(prior.a, prior.b), p)
    return (F(p) - F(prior.low)) / ((F(prior.high) - F(prior.low)))
end

function expectation(f::Function, prior::Beta{T})::T where {T<:Real}

    return (prior.high - prior.low)/2 * sum( prior.pdf .* f.(prior.pivots) .* prior.weights )
end

function mean(prior::Beta{T})::T where {T<:Real}

    return expectation(p -> p, prior)
end

function condition(prior::Beta{T}; low::T = prior.low, high::T = prior.high)::Beta{T} where {T<:Real}

    return Beta(prior.a, prior.b, low = max(low, prior.low), high = min(high, prior.high))
end

function update(prior::Beta{T}, x::TI, n::TI)::Beta{T} where {T<:Real,TI<:Integer}

    !(0 <= x <= n) ? error(@sprintf "invalid x=%i / n=%i " x n) : nothing
    return Beta(prior.a + x, prior.b + n - x; low = prior.low, high = prior.high)
end

function string(prior::Beta)
    if (prior.low != 0) | (prior.high != 1)
        @sprintf "Beta|[%.2f,%.2f](a=%.2f,b=%.2f)" prior.low prior.high prior.a prior.b
    else
        @sprintf "Beta(a=%.2f,b=%.2f)" prior.a prior.b
    end
end

struct BetaMixture{T<:Real} <: Prior
    ω::Vector{T}
    priors::Vector{Beta{T}}
end

BetaMixture(ω::T, priors::Vector{Beta{T}}) where {T<:Real} = BetaMixture{T}(ω, priors)

bounds(prior::BetaMixture) = [minimum(minimum.(bounds.(prior.priors))), maximum(maximum.(bounds.(prior.priors)))]

*(ω::T, prior::Beta{T}) where{T<:Real} = BetaMixture{T}([ω], [prior])
*(ω::T, prior::BetaMixture{T}) where{T<:Real} = BetaMixture{T}(ω.*prior.ω, prior.priors)
+(φ::BetaMixture{T}, η::BetaMixture{T}) where {T<:Real} = BetaMixture{T}(vcat(φ.ω, η.ω), vcat(φ.priors, η.priors))
is_proper(mprior::BetaMixture) = sum(mprior.ω) == 1

function pdf(p::T, mprior::BetaMixture{T})::T where {T<:Real}

    return sum( mprior.ω .* pdf.(p, mprior.priors) )
end

function pdf(x1::TI, x2::TI, mprior::BetaMixture{T}) where {T<:Real,TI<:Integer}

    return sum( mprior.ω .* pdf.(x1, x2, mprior.priors) )
end

function cdf(p::T, mprior::BetaMixture{T})::T where {T<:Real}

    return sum( mprior.ω .* cdf.(p, mprior.priors) )
end

function expectation(f::Function, mprior::BetaMixture{T})::T where {T<:Real}

    return sum( mprior.ω .* expectation.(f::Function, mprior.priors) )
end

function mean(mprior::BetaMixture{T})::T where {T<:Real}

    sum( mprior.ω .* mean.(mprior.priors) )
end

condition(mprior::BetaMixture; low::T1 = 0., high::T1 = 1.) where {T1<:Real, T2<:Real} =
    BetaMixture(mprior.ω,  condition.(mprior.priors; low = low, high = high))

update(mprior::BetaMixture{T}, x::TI, n::TI) where {T<:Real, TI<:Integer} = BetaMixture{T}(mprior.ω,  update.(mprior.priors, x, n))

function string(mprior::BetaMixture)
    n = length(mprior.priors)
    res = ""
    for i in 1:(n - 1)
        res *= (@sprintf "%.2f*" mprior.ω[i]) * string(mprior.priors[i]) * " + "
    end
    res *= (@sprintf "%.2f*" mprior.ω[n]) * string(mprior.priors[n])
    res
end


@testset "Integral Data Generation (Univariate)" begin
    m = InfiniteModel();
    @infinite_parameter(m, t in [-Inf, Inf], supports = [0, .5, 1])
    @infinite_parameter(m, x[1:2] in [-Inf, Inf])
    # test _trapezoid_coeff
    @testset "_trapezoid_coeff" begin
        @test InfiniteOpt.MeasureToolbox._trapezoid_coeff([1,2,3]) == [0.5, 1., 0.5]
        @test_throws ErrorException InfiniteOpt.MeasureToolbox._trapezoid_coeff([1])
    end
    # test generate_integral_data (trapezoid)
    @testset "generate_integral_data (trapezoid)" begin
        @test generate_integral_data(t, 0, 1, UniTrapezoid()) ==
                    FunctionalDiscreteMeasureData(t, IOMT._trapezoid_coeff, 0, All,
                               NoGenerativeSupports(),default_weight, 0, 1, false)
    end
    # test _ensure_independent_param
    @testset "_ensure_independent_param" begin
        @test IOMT._ensure_independent_param(t, Automatic()) isa Nothing
        @test_throws ErrorException IOMT._ensure_independent_param(x[1], Automatic())
    end
    # test generate_integral_data (Gauss-Legendre)
    @testset "generate_integral_data (Gauss-Legendre)" begin
        @test generate_integral_data(t, 0, 1, GaussLegendre()) isa DiscreteMeasureData
        @test generate_integral_data(t, 0, 1, GaussLegendre()).coefficients == FastGaussQuadrature.gausslegendre(3)[2] * 1/2
    end
    # test generate_integral_data (Gauss-Lobatto)
    @testset "generate_integral_data (Gauss-Lobato)" begin
        @test generate_integral_data(t, 0, 1, GaussLobatto()) isa DiscreteMeasureData
        @test generate_integral_data(t, 0, 1, GaussLobatto()).coefficients == FastGaussQuadrature.gausslobatto(3)[2] * 1/2
    end
    # test _lobatto_coeff
    @testset "_lobatto_coeff" begin
        @test InfiniteOpt.MeasureToolbox._lobatto_coeff([1,2], 2) == [0.5, 0.5]
        @test_throws ErrorException InfiniteOpt.MeasureToolbox._lobatto_coeff([1], 3)
    end
    # test generate_integral_data (FEGaussLobatto)
    @testset "generate_integral_data (FEGaussLobatto)" begin
        @test_throws ErrorException generate_integral_data(t, 0, Inf, FEGaussLobatto())
        @test_throws ErrorException generate_integral_data(t, 0, 1, FEGaussLobatto(), num_nodes = 1)
        @test_throws ErrorException generate_integral_data(x[1], 0, 1, FEGaussLobatto(), num_nodes = 3)
        @test generate_integral_data(t, 0, 1, FEGaussLobatto()) isa FunctionalDiscreteMeasureData
        @test generate_integral_data(t, 0, 1, FEGaussLobatto()).label == All
        expected = [0.08333333333333333, 0.3333333333333333, 0.16666666666666666,
                    0.3333333333333333, 0.08333333333333333]
        @test all(abs.(generate_integral_data(t, 0, 1, FEGaussLobatto()).coeff_function([0, 0.5, 1]) .- expected) .< 0.00005)
        @variable(m, flob, Infinite(t))
        integral(flob, t, 0, 1, eval_method = FEGaussLobatto(), num_nodes = 3)
        add_generative_supports(t)
        lobb_supps = supports(t, label = All)
        @test all(abs.(lobb_supps .- [0, .25, .5, .75, 1] .< .000005))
    end
    # test generate_integral_data (Gauss-Radau)
    @testset "generate_integral_data (Gauss-Radau)" begin
        @test generate_integral_data(t, 0, 1, GaussRadau()) isa DiscreteMeasureData
        @test generate_integral_data(t, 0, 1, GaussRadau()).coefficients == FastGaussQuadrature.gaussradau(3)[2] * 1/2
    end
    # test generate_integral_data (Gauss-Jacobi)
    @testset "generate_integral_data (Gauss-Jacobi)" begin
        @test generate_integral_data(t, 0, 1, GaussJacobi(5.0, 4.0)) isa DiscreteMeasureData
        @test generate_integral_data(t, 0, 1, GaussJacobi(5.0, 4.0)).coefficients == FastGaussQuadrature.gaussjacobi(3, 5.0, 4.0)[2] * 1/2
        @test_throws ErrorException GaussJacobi(-5, -5)
    end
     # test generate_integral_data (Gauss-Chebyshev)
    @testset "generate_integral_data (Gauss-Chebyshev)" begin
        @test generate_integral_data(t, 0, 1, GaussChebyshev(3)) isa DiscreteMeasureData
        @test generate_integral_data(t, 0, 1, GaussChebyshev(3)).coefficients == FastGaussQuadrature.gausschebyshev(3, 3)[2] * 1/2
        @test_throws ErrorException GaussChebyshev(5)
    end
    # test generate_integral_data (Gauss-Laguerre)
    @testset "generate_integral_data (Gauss-Laguerre)" begin
        @test generate_integral_data(t, -Inf, 0, GaussLaguerre()) isa DiscreteMeasureData
        @test generate_integral_data(t, 0, Inf, GaussLaguerre()) isa DiscreteMeasureData
    end
    # test generate_integral_data (Gauss-Hermite)
    @testset "generate_integral_data (Gauss-Hermite)" begin
        @test generate_integral_data(t, -Inf, Inf, GaussHermite()) isa DiscreteMeasureData
    end
    # test generate_integral_data (Quadrature)
    @testset "generate_integral_data (Quadrature)" begin
        @test generate_integral_data(t, 0, 1, Quadrature()) isa FunctionalDiscreteMeasureData
        @test generate_integral_data(t, -Inf, 0, Quadrature()) isa DiscreteMeasureData
        @test generate_integral_data(t, -Inf, Inf, Quadrature()) isa DiscreteMeasureData
    end
    # test handling of inappropriate quadrature method use
    @testset "generate_integral_data (improper quadrature method usage)" begin
        warn = "The `GaussLegendre` method can only be applied on finite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, 0, GaussLegendre())) isa DiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, Inf, GaussLegendre())) isa DiscreteMeasureData
        warn = "The `GaussLobatto` method can only be applied on finite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, 0, GaussLobatto())) isa DiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, Inf, GaussLobatto())) isa DiscreteMeasureData
        warn = "The `GaussRadau` method can only be applied on finite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, 0, GaussRadau())) isa DiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, Inf, GaussRadau())) isa DiscreteMeasureData
        warn = "The `GaussJacobi` method can only be applied on finite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, 0, GaussJacobi(5.0, 4.0))) isa DiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, Inf, GaussJacobi(5.0, 4.0))) isa DiscreteMeasureData
        warn = "The `GaussChebyshev` method can only be applied on finite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, 0, GaussChebyshev(3))) isa DiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, Inf, GaussChebyshev(3))) isa DiscreteMeasureData
        warn = "`GaussLaguerre` quadrature can only be applied on semi-infinite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, 0, 1, GaussLaguerre())) isa DiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, Inf, GaussLaguerre())) isa DiscreteMeasureData
        warn = "`GaussHermite` quadrature can only be applied on infinite intervals, " *
               "switching to an appropriate method."
        @test (@test_logs (:warn, warn) generate_integral_data(t, 0, 1, GaussHermite())) isa FunctionalDiscreteMeasureData
        @test (@test_logs (:warn, warn) generate_integral_data(t, -Inf, 0, GaussHermite())) isa DiscreteMeasureData
    end
    # test generate_integral_data (uniform Monte Carlo)
    @testset "generate_integral_data (UniMCSampling())" begin
        @test generate_integral_data(t, 0, 1, UniMCSampling()) isa FunctionalDiscreteMeasureData
        @test_throws ErrorException generate_integral_data(t, -Inf, 0, UniMCSampling())
        warn = "Cannot specify a nonzero `num_supports` when making an " *
               "integral with repsect to an individual dependent parameter. " *
               "Setting `num_supports = 0`."
        @test_logs (:warn, warn) generate_integral_data(x[1], 0, 1, UniMCSampling(), num_supports = 1)
    end
    # test generate_integral_data (independent Monte Carlo)
    @testset "generate_integral_data (UniIndepMCSampling())" begin
        @test generate_integral_data(t, 0, 1, UniIndepMCSampling()) isa DiscreteMeasureData
        @test_throws ErrorException generate_integral_data(t, -Inf, 0, UniIndepMCSampling())
        @test_throws ErrorException generate_integral_data(t, -Inf, Inf, UniIndepMCSampling())
    end
    # test generate_integral_data (Automatic)
    @testset "generate_integral_data (Automatic)" begin
        @test generate_integral_data(t, 0, 1, Automatic()) isa FunctionalDiscreteMeasureData
        @test generate_integral_data(t, -Inf, 1, Automatic()) isa DiscreteMeasureData
        @test generate_integral_data(t, -Inf, Inf, Automatic()) isa DiscreteMeasureData
        @test generate_integral_data(x[1], 0, 1, Automatic()) isa FunctionalDiscreteMeasureData
        @test_throws ErrorException generate_integral_data(x[1], -Inf, 0, Automatic())
        @test_throws ErrorException generate_integral_data(x[1], -Inf, Inf, Automatic())
    end
    # test fallback
    @testset "generate_integral_data (fallback)" begin
        @test_throws ErrorException generate_integral_data(x, 1, 2, Quadrature())
    end
end

@testset "Integral Data Generation (Multivariate)" begin
    m = InfiniteModel();
    @infinite_parameter(m, x[1:2] in [-Inf, Inf])
    # test generate_integral_data (MultiMCSampling)
    @testset "generate_integral_data (MultiMCSampling)" begin
        @test generate_integral_data(x, [0, 0], [1, 1], MultiMCSampling()) isa FunctionalDiscreteMeasureData
        @test_throws ErrorException generate_integral_data(x, [0, 0], [Inf, 1], MultiMCSampling())
        @test_throws ErrorException("`MultiMCSampling` is not applicable to (semi-)infinite intervals.") generate_integral_data(x, [0, -Inf], [1, 1], MultiMCSampling())
    end
    # test _make_multi_mc_supports
    # @testset "_make_multi_mc_supports" begin
    #     dx = dispatch_variable_ref.(x)
    #     @test IOMT._make_multi_mc_supports(dx, [0, 0], [1, 1], 5) isa Matrix{Float64}
    #     @test size(IOMT._make_multi_mc_supports(dx, [0, 0], [1, 1], 5)) == (2, 5)
    # end
    # test generate_integral_data (MultiIndepMCSampling)
    @testset "generate_integral_data (MultiIndepMCSampling)" begin
        @test generate_integral_data(x, [0, 0], [1, 1], MultiIndepMCSampling()) isa DiscreteMeasureData
        @test_throws ErrorException generate_integral_data(x, [0, 0], [Inf, 1], MultiIndepMCSampling())
        @test_throws ErrorException generate_integral_data(x, [0, -Inf], [1, 1], MultiIndepMCSampling())
    end
    # test generate_integral_data (Automatic)
    @testset "generate_integral_data (Automatic)" begin
        @test generate_integral_data(x, [0, 0], [1, 1], Automatic()) isa FunctionalDiscreteMeasureData
    end
end

@testset "Univariate Integrals" begin
    m = InfiniteModel();
    @infinite_parameter(m, t in [0, 10])
    @variable(m, inf1, Infinite(t))
    @testset "uni_integral_defaults" begin
        @test uni_integral_defaults() == IOMT.UniIntegralDefaults
    end
    @testset "[set/clear]_uni_integral_defaults" begin
        @test uni_integral_defaults() == Dict{Symbol, Any}(:eval_method => Automatic())
        @test set_uni_integral_defaults(num_nodes = 5, new_kwarg = true) isa Nothing
        @test uni_integral_defaults()[:num_nodes] == 5
        @test uni_integral_defaults()[:new_kwarg]
        @test clear_uni_integral_defaults() isa Nothing
        @test uni_integral_defaults() == Dict{Symbol, Any}(:eval_method => Automatic())
    end
    @testset "integral" begin
        @test InfiniteOpt._index_type(integral(inf1, t)) == MeasureIndex
        @test_throws ErrorException integral(inf1, t, -1, 1)
        @test_throws ErrorException integral(inf1, t, 9, 11)
        @test_throws ErrorException integral(inf1, t, 3, 1)
        @test name(integral(inf1, t)) == "integral"
    end
    @testset "∫" begin
        @test InfiniteOpt._index_type(∫(inf1, t)) == MeasureIndex
        @test_throws ErrorException ∫(inf1, t, -1, 1)
        @test_throws ErrorException ∫(inf1, t, 9, 11)
        @test_throws ErrorException ∫(inf1, t, 3, 1)
        @test name(∫(inf1, t)) == "integral"
    end
end

@testset "Multivariate Integrals" begin
    m = InfiniteModel();
    @infinite_parameter(m, x[1:2] in [0, 10])
    @infinite_parameter(m, xi[1:2] in [0, 5])
    @infinite_parameter(m, y[2:3] in [0, 1])
    @variable(m, inf1, Infinite(x))
    @variable(m, inf2, Infinite(xi))
    @variable(m, inf3, Infinite(y))
    @testset "multi_integral_defaults" begin
        @test multi_integral_defaults() == IOMT.MultiIntegralDefaults
    end
    @testset "[set/clear]_multi_integral_defaults" begin
        @test multi_integral_defaults() == Dict{Symbol, Any}(:eval_method => Automatic())
        @test set_multi_integral_defaults(num_supports = 5, new_kwarg = true) isa Nothing
        @test multi_integral_defaults()[:num_supports] == 5
        @test multi_integral_defaults()[:new_kwarg]
        @test clear_multi_integral_defaults() isa Nothing
        @test multi_integral_defaults() == Dict{Symbol, Any}(:eval_method => Automatic())
    end
    @testset "integral" begin
        @test InfiniteOpt._index_type(integral(inf1, x)) == MeasureIndex
        @test InfiniteOpt._index_type(integral(inf2, xi)) == MeasureIndex
        @test InfiniteOpt._index_type(integral(inf1, x, 1, 2)) == MeasureIndex
        @test InfiniteOpt._index_type(integral(inf2, xi, 1, 2)) == MeasureIndex
        @test_throws ErrorException integral(inf3, y, [0, 0], [1, 1])
        @test_throws ErrorException integral(inf1, x, 3, 1)
        @test_throws ErrorException integral(inf1, x, -1, 1)
        @test_throws ErrorException integral(inf1, x, 9, 11)
        @test name(integral(inf1, x)) == "integral"
    end
    @testset "∫" begin
        @test InfiniteOpt._index_type(∫(inf1, x)) == MeasureIndex
        @test InfiniteOpt._index_type(∫(inf2, xi)) == MeasureIndex
        @test InfiniteOpt._index_type(∫(inf1, x, 1, 2)) == MeasureIndex
        @test InfiniteOpt._index_type(∫(inf2, xi, 1, 2)) == MeasureIndex
        @test_throws ErrorException ∫(inf3, y, [0, 0], [1, 1])
        @test_throws ErrorException ∫(inf1, x, 3, 1)
        @test_throws ErrorException ∫(inf1, x, -1, 1)
        @test_throws ErrorException ∫(inf1, x, 9, 11)
        @test name(∫(inf1, x)) == "integral"
    end
end

@testset "Macro" begin
    m = InfiniteModel();
    @infinite_parameter(m, t in [0, 10])
    @infinite_parameter(m, x[1:2] in [0, 10])
    @variable(m, inf1, Infinite(t))
    @variable(m, inf2, Infinite(x))
    @testset "@integral" begin
        @test InfiniteOpt._index_type(@integral(inf1, t)) == MeasureIndex
        @test InfiniteOpt._index_type(@integral(inf2, x)) == MeasureIndex
        @test InfiniteOpt._index_type(@integral(inf2, t)) == MeasureIndex
        @test InfiniteOpt._index_type(@integral(inf1, t, 3, 7)) == MeasureIndex
        @test InfiniteOpt._index_type(@integral(inf2, x, [0, 2], [3, 5])) == MeasureIndex
        @test_macro_throws ErrorException @integral(inf1, t, 1)
    end
    @testset "@∫" begin
        @test InfiniteOpt._index_type(@∫(inf1, t)) == MeasureIndex
        @test InfiniteOpt._index_type(@∫(inf2, x)) == MeasureIndex
        @test InfiniteOpt._index_type(@∫(inf2, t)) == MeasureIndex
        @test InfiniteOpt._index_type(@∫(inf1, t, 3, 7)) == MeasureIndex
        @test InfiniteOpt._index_type(@∫(inf2, x, [0, 2], [3, 5])) == MeasureIndex
        @test_macro_throws ErrorException @∫(inf1, t, 1)
    end
end


@info "Starting utilities tests"
@testset "Utilities" begin
    t = collect(-2:0.01:2)
    U = [cos.(t).*exp.(-t.^2) sin.(2*t)]
    S = Diagonal([2.; 3.])
    V = [sin.(t).*exp.(-t) cos.(t)]
    A = U*S*V'
    σ = 0.5
    Â = A + σ*randn(401, 401)
    n_1 = norm(A-Â)
    B = optimal_shrinkage(Â)
    optimal_shrinkage!(Â)
    @test norm(A-Â) < n_1
    @test norm(A-B) == norm(A-Â)

    X = randn(3, 100)
    Y = randn(3, 100)
    k = 3

    @test AIC(k, X, Y) == 2*k-2*log(sum(abs2, X- Y))
    @test AICC(k, X, Y) == AIC(k, X, Y)+ 2*(k+1)*(k+2)/(size(X)[2]-k-2)
    @test BIC(k, X, Y) == -2*log(sum(abs2, X -Y)) + k*log(size(X)[2])
    @test AICC(k, X, Y, likelihood = (X,Y)->sum(abs, X-Y)) == AIC(k, X, Y, likelihood = (X,Y)->sum(abs, X-Y))+ 2*(k+1)*(k+2)/(size(X)[2]-k-2)

   end


using ITensors,
      ITensorsGPU,
      LinearAlgebra, # For tr()
      Combinatorics, # For permutations()
      Random,
      CUDA,
      Test

@testset "cuITensor $T Contractions" for T ∈ (Float64,ComplexF64)
  mi,mj,mk,ml,ma = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  a = Index(ma,"a")
  @testset "Test contract cuITensors" begin
      A = cuITensor(randomITensor(T))
      B = cuITensor(randomITensor(T))
      Ai = cuITensor(randomITensor(T,i))
      Bi = cuITensor(randomITensor(T,i))
      Aj = cuITensor(randomITensor(T,j))
      Aij = cuITensor(randomITensor(T,i,j))
      Aji = cuITensor(randomITensor(T,j,i))
      Bij = cuITensor(randomITensor(T,i,j))
      Aik = cuITensor(randomITensor(T,i,k))
      Ajk = cuITensor(randomITensor(T,j,k))
      Ajl = cuITensor(randomITensor(T,j,l))
      Akl = cuITensor(randomITensor(T,k,l))
      Aijk = cuITensor(randomITensor(T,i,j,k))
      Ajkl = cuITensor(randomITensor(T,j,k,l))
      Aikl = cuITensor(randomITensor(T,i,k,l))
      Akla = cuITensor(randomITensor(T,k,l,a))
      Aijkl = cuITensor(randomITensor(T,i,j,k,l))
    @testset "Test contract cuITensor (Scalar*Scalar -> Scalar)" begin
      C = A*B
      @test scalar(C)≈scalar(A)*scalar(B)
      C = cuITensor(T(2.0))*cuITensor(T(2.0))
      @test scalar(C)≈T(4.0)
    end
    @testset "Test contract cuITensor (Scalar*Vector -> Vector)" begin
      C = A*Ai
      @test cpu(C)≈scalar(A)*cpu(Ai)
      C = cuITensor(T(2.0))*Ai
      @test cpu(C)≈T(2.0)*cpu(Ai)
      C = Ai*cuITensor(T(2.0))
      @test cpu(C)≈T(2.0)*cpu(Ai)
    end
    @testset "Test contract cuITensor (Vector*Scalar -> Vector)" begin
      C = Aj*A
      @test cpu(C)≈scalar(A)*cpu(Aj)
    end
    @testset "Test contract cuITensors (Vectorᵀ*Vector -> Scalar)" begin
      C = Ai*Bi
      Ccpu = cpu(Ai) * cpu(Bi)
      @test scalar(Ccpu)≈scalar(C)
    end
    @testset "Test contract cuITensors (Vector*Vectorᵀ -> Matrix)" begin
      C = Ai*Aj
      Ccpu = cpu(Ai)*cpu(Aj)
      @test Ccpu≈cpu(permute(C,i,j))
    end
    @testset "Test contract cuITensors (Matrix*Scalar -> Matrix)" begin
      Aij = permute(Aij,i,j)
      C = Aij*A
      @test cpu(permute(C,i,j))≈scalar(A)*cpu(Aij)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Vector -> Vector)" begin
      cAij = permute(copy(Aij),j,i)
      Ccpu = cpu(Aij)*cpu(Aj)
      C = cAij*Aj
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Vector -> Vector)" begin
      cpAij = permute(copy(Aij),i,j)
      Ccpu = cpu(cpAij)*cpu(Aj)
      C   = copy(cpAij)*copy(Aj)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Vector*Matrix -> Vector)" begin
      Aij = permute(Aij,i,j)
      C = Ai*Aij
      Ccpu = cpu(Ai) * cpu(Aij)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Vector*Matrixᵀ -> Vector)" begin
      C = Ai*Aji
      Ccpu = cpu(Ai) * cpu(Aji)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrix -> Scalar)" begin
      Aij = permute(Aij,i,j)
      Bij = permute(Bij,i,j)
      C = Aij*Bij
      Ccpu = cpu(Aij) * cpu(Bij)
      @test scalar(Ccpu)≈scalar(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrix -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      Ccpu = cpu(Aij)*cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Matrix -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Scalar -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aijk*A
      @test cpu(permute(C,i,j,k))≈scalar(A)*cpu(Aijk)
    end
    @testset "Test contract cuITensors (3-Tensor*Vector -> Matrix)" begin
      cAijk = permute(copy(Aijk),i,j,k)
      C = cAijk*Ai
      Ccpu = cpu(cAijk) * cpu(Ai)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Vector*3-Tensor -> Matrix)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aj*Aijk
      Ccpu = cpu(Aj)*cpu(Aijk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Matrix -> Vector)" begin
      Aijk = permute(Aijk,i,j,k)
      Aik = permute(Aik,i,k)
      C = Aijk*Aik
      Ccpu = cpu(Aijk)*cpu(Aik)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Matrix -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajl = permute(Ajl,j,l)
      C = Aijk*Ajl
      Ccpu = cpu(Aijk)*cpu(Ajl)
      @test Ccpu≈cpu(permute(C,i,k,l))
    end
    @testset "Test contract cuITensors (Matrix*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Akl  = permute(Akl,k,l)
      C = Akl*Aijk
      Ccpu = cpu(Aijk)*cpu(Akl)
      @test Ccpu≈cpu(permute(C,l,i,j))
    end
    @testset "Test contract cuITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajkl = permute(Ajkl,j,k,l)
      C = Aijk*Ajkl
      Ccpu = cpu(Aijk)*cpu(Ajkl)
      @test Ccpu≈cpu(permute(C,i,l))
    end
    @testset "Test contract cuITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijk ∈ permutations([i,j,k]), inds_jkl ∈ permutations([j,k,l])
        Aijk = permute(Aijk,inds_ijk...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijk
        Ccpu = cpu(Ajkl) * cpu(Aijk)
        @test Ccpu≈cpu(C)
      end
    end
    @testset "Test contract cuITensors (4-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijkl ∈ permutations([i,j,k,l]), inds_kla ∈ permutations([k,l,a])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Akla = permute(Akla,inds_kla...)
        C = Akla*Aijkl
        Ccpu = cpu(Akla) * cpu(Aijkl)
        @test Ccpu≈cpu(C)
      end
    end
    @testset "Test contract cuITensors (4-Tensor*3-Tensor -> 1-Tensor)" begin
      for inds_ijkl ∈ permutations([i,j,k,l]), inds_jkl ∈ permutations([j,k,l])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijkl
        Ccpu = cpu(Ajkl)*cpu(Aijkl)
        @test Ccpu≈cpu(C)
      end
    end
    @testset "Test supersized contract cuITensors (14-Tensor*14-Tensor -> 14-Tensor)" begin
        a_only_inds = [Index(2) for ii in 1:7]
        b_only_inds = [Index(2) for ii in 1:7]
        shared_inds = [Index(2) for ii in 1:7]
        A = randomITensor(IndexSet(vcat(a_only_inds, shared_inds)))
        B = randomITensor(IndexSet(vcat(b_only_inds, shared_inds)))
        cA = cuITensor(A)
        cB = cuITensor(B)
        inds_a = vcat(a_only_inds, shared_inds)
        inds_b = vcat(b_only_inds, shared_inds)
        cA_ = permute(cA,inds_a...)
        cB_ = permute(cB,inds_b...)
        C = cA_*cB_
        Ccpu = cpu(cA_)*cpu(cB_)
        @test Ccpu≈cpu(C)
        for shuffles in 1:1 # too many permutations to test all
            inds_a = shuffle(vcat(a_only_inds, shared_inds))
            inds_b = shuffle(vcat(b_only_inds, shared_inds))
            cA_ = permute(cA,inds_a...)
            cB_ = permute(cB,inds_b...)
            C = cA_*cB_
            Ccpu = cpu(cA_)*cpu(cB_)
            @test Ccpu≈cpu(C)
        end
    end
  end # End contraction testset
end