struct FuzzyVector
    numbers::Vector{FuzzyNumber}
end

Base.length(FV::FuzzyVector) = size(FV.numbers, 1)
Base.getindex(FV::FuzzyVector, i::Int64) = FV.numbers[i]
Base.show(io::IO, FV::FuzzyVector) = println(io, "fuzzy vector length $(length(FV))")
Base.iterate(FV::FuzzyVector, state=1) = state > length(FV.numbers) ? nothing : FV.numbers[state]

# FUZZY ARITHMETIC ============================================================
function Base.:+(A⃗::FuzzyVector, B⃗::FuzzyVector)
	p = length(A⃗)
	result_vector = Vector{FuzzyNumber}(undef, p)
	for k = 1:p
		result_vector[k] = A⃗[k] + B⃗[k]
	end
	FuzzyVector(result_vector)
end
function Base.:-(A⃗::FuzzyVector)
	p = length(A⃗)
	result_vector = Vector{FuzzyNumber}(undef, p)
	for k = 1:p
		result_vector[k] = -A⃗[k]
	end
	FuzzyVector(result_vector)
end
Base.:-(A⃗::FuzzyVector, B⃗::FuzzyVector) = A⃗ + (-B⃗)
# =============================================================================

function Base.:(==)(A⃗::FuzzyVector, B⃗::FuzzyVector)
    if length(A⃗) != length(B⃗)
        return false
    end
    for i = 1:length(A⃗)
        if A⃗[i] != B⃗[i]
            return false
        end
    end
    return true
end

function Base.:∪(FV::FuzzyVector, vec::Vector{FuzzyNumber})
   append!(FV.numbers, vec)
   return FV
end

function cut(X::FuzzyVector, α::Real)
	p = length(X)
	cuts = Vector{Interval}(undef, p)
	for j = 1:p
		cuts[j] = cut(X[j], α)
	end
	cuts
end

function cut(X::Vector{FuzzyVector}, p::Int, α::Real)
	N = length(X)
	cuts = Vector{Interval}(undef, N)
	for i = 1:N
		cuts[i] = cut(X[i][p], α)
	end
	cuts
end

function cut(X::Vector{FuzzyVector}, α::Real)
	N = length(X)
	p = length(X[1])
	cuts = Vector{Vector{Interval}}(undef, N)
	for i = 1:N
		cuts[i] = [cut(X[i][j], α) for j = 1:p]
	end
	cuts
end

function draw(ax::Axis3, FV::FuzzyVector; step=0.001, colormap=:jet1)
	A₁ = FV[1]
	A₂ = FV[2]
	X = collect(-4:step:4)
	Y = collect(-3:step:3)
	# X = collect(A₁.grades[1][1]:step:A₁.grades[1][2])
	# Y = collect(A₂.grades[1][1]:step:A₂.grades[1][2])
	Z = [min(A₁(x), A₂(y)) for x in X, y in Y]
	surface!(ax, X, Y, Z, colormap=colormap)#, axis=(type=Axis3,))

    # marking peak
	# max_indices = argmax(Z)
	# max_x = X[max_indices[1]]
	# max_y = Y[max_indices[2]]
	# μ = Z[max_indices]
    # annotate!(fig, [(max_x, max_y, μ, (μ, 8, :black, :center))])
end

function draw2d(FV::FuzzyVector; step::Float64=0.01, fig=nothing, c=:jet1, alpha::Real=nothing, marker=nothing, peak_text::Bool=false, font=Plots.font("Times", 8))
    gr(xtickfont=font, ytickfont=font, legendfont=font)
	A₁ = FV[1]
	A₂ = FV[2]
	# x1 = max(A₁.grades[1].left, -3.5)
	# x2 = min(A₁.grades[1].right, 3.5)
	# y1 = max(A₂.grades[1].left, -3.5)
	# y2 = min(A₂.grades[1].right, 3.5)
	x1 = A₁.grades[1].left
	x2 = A₁.grades[1].right
	y1 = A₂.grades[1].left
	y2 = A₂.grades[1].right
	X = collect(x1:step:x2)
	Y = collect(y1:step:y2)
	f(x, y) = min(A₁(x), A₂(y))
	if isnothing(fig)
		fig = Plots.contour(X, Y, f, c=c, aspect_ratio=1.0, seriesalpha=alpha, fill=true)
	else
		fig = Plots.contour(fig, X, Y, f, c=c, aspect_ratio=1.0, seriesalpha=alpha, fill=true)
		fig = Plots.contour(fig, [0, 0], [0, 0], [0, 1], c=c, aspect_ratio=1.0, seriesalpha=1.0, fill=true, colorbar=true)
	end

	# plot peak
	center_x = peak_center(A₁) # or can be centroid()
	center_y = peak_center(A₂) # or can be centroid()
	if !isnothing(marker)
		fig = Plots.scatter(fig, (center_x, center_y), legend=false, m=marker, color=:grey)
	end
	if peak_text
		Plots.annotate!([(center_x, center_y, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, :black, :center))])
	end

	fig = Plots.plot(fig, dpi=200)
	fig
end

function draw_contour_2d(FVs::Vector{FuzzyVector}; step::Float64=0.01, fig=nothing, contour_colors=[:jet1, :jet1], peak_colors=["black", "red"], alpha::Vector{Float64}=nothing, marker=nothing, peak_text::Bool=false, offset::Real=0, xlim=nothing, ylim=nothing, font=Plots.font("Times", 8))
    gr(xtickfont=font, ytickfont=font, legendfont=font)

	centers = []
	for (i, FV) in enumerate(FVs)
		A₁ = FV[1]
		A₂ = FV[2]
		x1 = A₁.grades[1].left
		x2 = A₁.grades[1].right
		y1 = A₂.grades[1].left
		y2 = A₂.grades[1].right
		X = collect(x1:step:x2)
		Y = collect(y1:step:y2)
		f(x, y) = min(A₁(x), A₂(y))
		center_x = peak_center(A₁) # or can be centroid()
		center_y = peak_center(A₂) # or can be centroid()
		push!(centers, (center_x, center_y))
		if isnothing(fig)
			fig = Plots.contour(X, Y, f, c=contour_colors[i], aspect_ratio=1.0, seriesalpha=alpha[i], xlim=xlim, ylim=ylim, fill=true)
		else
			fig = Plots.contour(fig, X, Y, f, c=contour_colors[i], aspect_ratio=1.0, seriesalpha=alpha[i], fill=true, xlim=xlim, ylim=ylim)
			fig = Plots.contour(fig, [0, 0], [0, 0], [0, 1], c=contour_colors[i], aspect_ratio=1.0, seriesalpha=1.0, fill=true, colorbar=true, xlim=xlim, ylim=ylim)
		end
	end

	# plot peak
	for (i, (center_x, center_y)) in enumerate(centers)
		if !isnothing(marker)
			fig = Plots.scatter(fig, (center_x, center_y), legend=false, m=marker, c=peak_colors[i])
		end
		if peak_text
			# fig = Plots.annotate(fig, [(center_x, center_y+offset, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, peak_colors[i], :center, "Times"))])
			Plots.annotate!([(center_x, center_y+offset, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, peak_colors[i], :center, "Times"))])
		end
	end

	fig = Plots.plot(fig, dpi=200)
	fig
end

function draw_3d(FVs::Vector{FuzzyVector}; step::Float64=0.01, fig=nothing, color=:grey, alpha::Float64=nothing, peak_text::Bool=false, offset::Real=0, xlim=nothing, ylim=nothing, font=Plots.font("Times", 8))
	centers = []
	for (i, FV) in enumerate(FVs)
		A₁ = FV[1]
		A₂ = FV[2]
		x1 = A₁.grades[1].left
		x2 = A₁.grades[1].right
		y1 = A₂.grades[1].left
		y2 = A₂.grades[1].right
		X = collect(x1:step:x2)
		Y = collect(y1:step:y2)
		f(x, y) = min(A₁(x), A₂(y))
		center_x = peak_center(A₁) # or can be centroid()
		center_y = peak_center(A₂) # or can be centroid()
		push!(centers, (center_x, center_y))
		fig = Plots.surface(fig, X, Y, f, aspect_ratio=1.0, color=color, alpha=alpha, fill=true, xlim=xlim, ylim=ylim)
		# fig = Plots.surface(fig, [0, 0], [0, 0], [0, 1], c=color, aspect_ratio=1.0, alpha=1.0, fill=true, colorbar=true, xlim=xlim, ylim=ylim)
	end

	# plot peak
	# for (i, (center_x, center_y)) in enumerate(centers)
	# 	if peak_text
	# 		# fig = Plots.annotate(fig, [(center_x, center_y+offset, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, peak_colors[i], :center, "Times"))])
	# 		Plots.annotate!([(center_x, center_y+offset, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, peak_colors[i], :center, "Times"))])
	# 	end
	# end

	fig = Plots.plot(fig, dpi=200)
	fig
end

function draw_center_3d(FVs::Vector{FuzzyVector}; step::Float64=0.01, fig=nothing, contour_colors=[:jet1, :jet1], peak_colors=["black", "red"], alpha::Vector{Float64}=nothing, marker=nothing, peak_text::Bool=false, offset::Real=0, xlim=nothing, ylim=nothing, font=Plots.font("Times", 8))
	centers = []
	for (i, FV) in enumerate(FVs)
		A₁ = FV[1]
		A₂ = FV[2]
		x1 = A₁.grades[1].left
		x2 = A₁.grades[1].right
		y1 = A₂.grades[1].left
		y2 = A₂.grades[1].right
		X = collect(x1:step:x2)
		Y = collect(y1:step:y2)
		f(x, y) = min(A₁(x), A₂(y))
		center_x = peak_center(A₁) # or can be centroid()
		center_y = peak_center(A₂) # or can be centroid()
		push!(centers, (center_x, center_y))
		fig = Plots.surface(fig, X, Y, f, c=contour_colors[i], aspect_ratio=1.0, seriesalpha=alpha[i], fill=true, xlim=xlim, ylim=ylim)
		fig = Plots.surface(fig, [0, 0], [0, 0], [0, 1], c=contour_colors[i], aspect_ratio=1.0, seriesalpha=1.0, fill=true, colorbar=true, xlim=xlim, ylim=ylim)
	end

	# plot peak
	# for (i, (center_x, center_y)) in enumerate(centers)
	# 	if !isnothing(marker)
	# 		fig = Plots.scatter(fig, (center_x, center_y), legend=false, m=marker, c=peak_colors[i])
	# 	end
	# 	if peak_text
	# 		# fig = Plots.annotate(fig, [(center_x, center_y+offset, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, peak_colors[i], :center, "Times"))])
	# 		Plots.annotate!([(center_x, center_y+offset, ("$(round(center_x, digits=2)), $(round(center_y, digits=2))", 8, peak_colors[i], :center, "Times"))])
	# 	end
	# end

	fig = Plots.plot(fig, dpi=200)
	fig
end

function draw2d_makie(ax::Axis, FV::FuzzyVector; step=0.01, colormap=:jet1, alpha=nothing)
	A₁ = FV[1]
	A₂ = FV[2]
	x1 = max(A₁.grades[1].left, -3.5)
	x2 = min(A₁.grades[1].right, 3.5)
	y1 = max(A₂.grades[1].left, -3.5)
	y2 = min(A₂.grades[1].right, 3.5)
	X = collect(x1:step:x2)
	Y = collect(y1:step:y2)
	Z = [min(A₁(x), A₂(y)) for x in X, y in Y]
	contour!(ax, X, Y, Z, colormap=colormap)#, alpha=alpha)
end

function u_uncertainty(FV::FuzzyVector; agg::String="average")
	p = length(FV)
	Us = Vector{Float64}(undef, p)
	for k = 1:p
		Us[k] = u_uncertainty(FV[k])
	end
	if agg == "average"
		return sum(Us) / p
	elseif agg == "max"
		return maximum(Us)
	elseif agg == "sum"
		return sum(Us)
	end
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function vactoair(wave_vac::AbstractFloat)
    if wave_vac >= 2000
        sigma2 = (1e4/wave_vac)^2  # Convert to wavenumber squared
        # Computer conversion factor.
        fact = 1 + 5.792105e-2/(238.0185 - sigma2) +
            1.67917e-3/(57.362 - sigma2)
        return wave_vac/fact # Convert wavelength
    else
        return wave_vac
    end
end

"""
    vactoair(wave_vacuum) -> wave_air

### Purpose ###

Converts vacuum wavelengths to air wavelengths.

### Explanation ###

Corrects for the index of refraction of air under standard conditions.
Wavelength values below \$2000 Å\$ will not be altered.  Uses relation of Ciddor
(1996).

### Arguments ###

* `wave_vacuum`: vacuum wavelength in angstroms.  Wavelengths are corrected for the index of
  refraction of air under standard conditions.  Wavelength values below \$2000 Å\$ will
  *not* be altered, take care within \$[1 Å, 2000 Å]\$.

### Output ###

Air wavelength in angstroms.

### Method ###

Uses relation of Ciddor (1996), Applied Optics 35, 1566
(http://adsabs.harvard.edu/abs/1996ApOpt..35.1566C).

### Example ###

If the vacuum wavelength is `w = 2000`, then `vactoair(w)` yields an air
wavelength of `1999.353`.

```jldoctest
julia> using AstroLib

julia> vactoair(2000)
1999.3526230448367
```

### Notes ###

`airtovac` converts air wavelengths to vacuum wavelengths.

Code of this function is based on IDL Astronomy User's Library.
"""
vactoair(wave_vac::Real) = vactoair(float(wave_vac))


using ModelingToolkit, SparseArrays, Test, GalacticOptim, GalacticOptimJL

@variables x y
@parameters a b
loss = (a - x)^2 + b * (y - x^2)^2
sys1 = OptimizationSystem(loss, [x, y], [a, b], name = :sys1)
sys2 = OptimizationSystem(loss, [x, y], [a, b], name = :sys2)

@variables z
@parameters β
loss2 = sys1.x - sys2.y + z * β
combinedsys = OptimizationSystem(loss2, [z], [β], systems = [sys1, sys2],
                                 name = :combinedsys)

equations(combinedsys)
states(combinedsys)
parameters(combinedsys)

calculate_gradient(combinedsys)
calculate_hessian(combinedsys)
generate_function(combinedsys)
generate_gradient(combinedsys)
generate_hessian(combinedsys)
ModelingToolkit.hessian_sparsity(combinedsys)

u0 = [sys1.x => 1.0
      sys1.y => 2.0
      sys2.x => 3.0
      sys2.y => 4.0
      z => 5.0]
p = [sys1.a => 6.0
     sys1.b => 7.0
     sys2.a => 8.0
     sys2.b => 9.0
     β => 10.0]

prob = OptimizationProblem(combinedsys, u0, p, grad = true)
sol = solve(prob, NelderMead())
@test sol.minimum < -1e5

prob2 = remake(prob, u0 = sol.minimizer)
sol = solve(prob, BFGS(initial_stepnorm = 0.0001), allow_f_increases = true)
@test sol.minimum < -1e8
sol = solve(prob2, BFGS(initial_stepnorm = 0.0001), allow_f_increases = true)
@test sol.minimum < -1e9

rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
x0 = zeros(2)
_p = [1.0, 100.0]

f = OptimizationFunction(rosenbrock, GalacticOptim.AutoModelingToolkit())
prob = OptimizationProblem(f, x0, _p)
sol = solve(prob, Newton())

# issue #819
@testset "Combined system name collisions" begin
    sys2 = OptimizationSystem(loss, [x, y], [a, b], name = :sys1)
    @test_throws ArgumentError OptimizationSystem(loss2, [z], [β], systems = [sys1, sys2])
end

# observed variable handling
@variables OBS
@named sys2 = OptimizationSystem(loss, [x, y], [a, b]; observed = [OBS ~ x + y])
OBS2 = OBS
@test isequal(OBS2, @nonamespace sys2.OBS)
@unpack OBS = sys2
@test isequal(OBS2, OBS)


using Revise
using AdFem
using PyPlot
using LinearAlgebra
using Statistics
using MAT
SMALL_SIZE = 20
MEDIUM_SIZE = 20
BIGGER_SIZE = 20

matplotlib.rc("font", size=SMALL_SIZE)          # controls default text sizes
matplotlib.rc("axes", titlesize=SMALL_SIZE)     # fontsize of the axes title
matplotlib.rc("axes", labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
matplotlib.rc("xtick", labelsize=SMALL_SIZE)    # fontsize of the tick labels
matplotlib.rc("ytick", labelsize=SMALL_SIZE)    # fontsize of the tick labels
matplotlib.rc("legend", fontsize=SMALL_SIZE)    # legend fontsize
matplotlib.rc("figure", titlesize=BIGGER_SIZE)  # fontsize of the figure title

function f(x, y)
    1/(1+x^2) + x * y + y^2
end

mmesh = Mesh(50, 50, 1/50, degree=2)

E = eval_f_on_gauss_pts(f, mmesh)

E2 = matread("fenics/data2-inverse10001.mat")["E"]
E3 = matread("pixel/data2-inverse981.mat")["E"]

close("all")
figure(figsize=(30, 8))
subplot(131)
title("Exact")
visualize_scalar_on_gauss_points(E, mmesh)
gca().set_rasterized(true)
subplot(132)
title("Difference (DNN)")
visualize_scalar_on_gauss_points(abs.(E-E2), mmesh)
xlabel("x")
subplot(133)
title("Difference (Discretization on Gauss Points)")
visualize_scalar_on_gauss_points(abs.(E-E3), mmesh)
xlabel("x")
tight_layout()
savefig("LinearElasticity.png")
savefig("LinearElasticity.pdf")


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

nearest correlation matrix, in the quantum relative entropy sense, adapted from
https://github.com/hfawzi/cvxquad/blob/master/examples/nearest_correlation_matrix.m
=#

struct NearestCorrelationJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    side::Int
end

function build(inst::NearestCorrelationJuMP{T}) where {T <: Float64}
    side = inst.side
    M = randn(T, side, side)
    M = M * M'
    vec_dim = Cones.svec_length(side)
    m_vec = zeros(T, vec_dim)
    Cones.smat_to_svec!(m_vec, M, sqrt(T(2)))

    model = JuMP.Model()
    JuMP.@variable(model, x_vec[1:vec_dim])
    X = zeros(JuMP.AffExpr, side, side)
    Cones.svec_to_smat!(X, one(T) * x_vec, sqrt(T(2)))
    JuMP.@constraint(model, diag(X) .== 1)

    JuMP.@variable(model, y)
    JuMP.@objective(model, Min, y)
    JuMP.@constraint(model, vcat(y, x_vec, m_vec) in
        Hypatia.EpiTrRelEntropyTriCone{T}(1 + 2 * vec_dim))

    return model
end


@doc raw"""
    update!(filter, y, u)

Update the internal states of the filter in accordance with the output y
and input u.

# Methods
    update!(filter::KalmanFilter, y, u)

Update the internal states of the KalmanFilter in according to
measured output $y$ and controlled input $u$.

The internal states are updated as follows

```math
\begin{aligned}
P_{t+1} & = Q + FP_{t}F^\top \quad - FP_{t}H^\top(R + HP_{t}H^\top)^{-1}H P_{t}F^\top \\
\breve{x}_{t+1}   & = F \breve{x}_{t} + K_{t}(y_t - H\breve{x}_{t})\\
K_{t}     & = FP_{t}H^\top(R + H P_{t} H^\top)^{-1} \\
c_{t+1} & = |H_i\breve x_{t} -y_t|^2_{(R + HP_{t}H^\top)^{-1}} + c_{t} \\
p(i|y_t) & = \frac{1}{(2\pi)^{n/2}\sqrt{|R + HP_tH^\top|}}e^{-1/2 (H\hat x_t - y_t)^\top (R + HP_tH^\top)^{-1}(Hx_t-y_t)}p(i|y_{t-1}).
\end{aligned}
```

The cummulative costs $c_t$ are states used by Minimax Adaptive Estimators, and $p(i|y_t)$ are states used by the standard adaptive estimator.

    update!(filter::MinimaxFilter, y, u)

Update each filter in the internal filterbank of the minimaxfilter object.

    update!(filter::BayesianFilter, y, u)

Update each filter in the internal filterbank of the standard multiple model adaptive filter
and normalize the probability of a given state being active conditioned on past measurements,
$p(i, y_t)$.
"""
function update!(filter::AbstractFilter, y, u)
end

function update!(
    filter::KalmanFilter{T},
    y::AbstractVector{T},
    u::AbstractVector{T},
    ) where {T<:Number}

    x = filter.x
    F = filter.F
    H = filter.H
    B = filter.B
    P = filter.P
    n = length(x)
    offset = filter.offset
    Re = (filter.R + H * P * H')
    K = F * P * H' / Re

    filter.c[] += ((H * x - y)' / Re) * (H * x - y)
    filter.pdy[] *= 1 / (2*π)^(n/2) / sqrt(det(Re)) * exp(-1/2 * ((H*x -y)' / Re)*(H*x - y))
    filter.x[:] = F*x + B*u + K*(y-H*x) + offset
    filter.P[:,:] = filter.Q + F*P*F' - F*P*(H' / Re) * H *P*F'
    return nothing
end

function update!(
    filter::MinimaxMMAE{T},
    y::AbstractVector{T},
    u::AbstractVector{T},
) where {T<:Number}
    for f in filter.filterbank
        update!(f, y, u)
    end

    return nothing
end

function update!(
    filter::BayesianMMAE{T},
    y::AbstractVector{T},
    u::AbstractVector{T},
) where {T<:Number}
    sum_pdys = 0
    for f in filter.filterbank
        update!(f, y, u)
        sum_pdys += f.pdy[]
    end

    for f in filter.filterbank
        f.pdy[] = f.pdy[]/sum_pdys
    end

    return nothing
end

@doc raw"""
    predict(filter)

Predict the states at the next time instance.

# Methods

    xhat = predict(filter::KalmanFilter)

Get the kalman filter prediction of the state at the next time-step, $\hat x_{t+1}$.

    yhat, val = predict(minimaxfilter)

Predict the output of the next time-step as the minimizing argument of
 the quadratically constrainde convex program
```math
J_N^\star(y)= \min_{\hat y} \max_i |\hat y - H_i\breve x_{N,i}|^2_{(I-\gamma^{-2}H_iP_{N,i}H_i^\top)^{-1}} - \gamma^2 c_{N,i}.
```

    yhat = predict(bayesianfilter)

Predict the output at the next time-step as the expected value
`` E[\hat y_{t+1}] = \sum_{i=1}^K \breve y_i p(i|y_t) ``, where $\breve y_i$ are the
corresponding KalmanFilter estimates.
"""
function predict(filter::AbstractFilter)
end

function predict(filter::KalmanFilter)
    return copy(filter.x)
end

function predict(filter::MinimaxMMAE, γ::Real; silent = false)
    γ <= 0 ? DomainError("γ must be positive") : nothing
    m = size(filter.filterbank[1].H)[1]
    K = length(filter.filterbank)

    # Setting up optimization model
    model = Model(filter.optimizer)
    silent ? set_silent(model) : unset_silent(model)
    @variable(model, yhat[1:m])
    @variable(model, t)
    for f in filter.filterbank
        W = (I - γ^(-2)*f.H*f.P*f.H')
        @constraint(
            model,
            (yhat - f.H*f.x)'*pinv(W)*(yhat - f.H*f.x) - γ^2*f.c[] <= t
        )
    end

    @objective(model, Min, t)
    optimize!(model)
    yhat = value.(yhat)
    val = value(t)
    return yhat, val
end

function predict(filter::MinimaxMMAE; γmax = 100, silent = true)
    γmin = sqrt(maximum([maximum(eigvals(kf.H*kf.P*kf.H')) for kf in filter.filterbank]))
    γ = (γmin + γmax)/2
    yhat = []
    val = 0
    for k = 1:20
        yhat, val = predict(filter, γ, silent = silent)
        val < 0 ? γmax = γ : γmin = γ
        γ = (γmin + γmax)/2
    end
    return yhat, val, γ
end

function predict(filter::BayesianMMAE)
    yhat = [f.pdy[]*f.H*f.x for f in filter.filterbank] |> sum
    return yhat
end

# Lazy wrappers

function update!(filter::AbstractFilter{T}, y::Number) where {T<:Number}
    return update!(filter, [T(y)], 0)
end

function update!(filter::AbstractFilter{T}, y::AbstractVector) where {T<:Number}
    return update!(filter, T.(y), 0)
end

function update!(filter::AbstractFilter{T}, y::Number, u::Number) where {T<:Number}
    return update!(filter, [T(y)], [T(u)])
end

function update!(filter::AbstractFilter{T}, y::Number, u::AbstractVector) where {T<:Number}
    return update!(filter, [T(y)], T.(u))
end

function update!(filter::AbstractFilter{T}, y::AbstractVector, u::Number) where {T<:Number}
    return update!(filter, T.(y), [T.(u)])
end

function update!(
    filter::AbstractFilter{T},
    y::AbstractVector,
    u::AbstractVector,
) where {T<:Number}

    return update!(filter, T.(y), T.(u))
end


module Integration

using Compat
using FastGaussQuadrature

export qdata

function qdata(beta::Vector{Float64},nqpts::Int64)

  n = length(beta)
  qnode = zeros(nqpts,n+1)
  qwght = zeros(nqpts,n+1)
  for j = findall(beta.>-1)
    qnode[:,j],qwght[:,j] = gaussjacobi(nqpts,0,beta[j])
  end
  qnode[:,n+1],qwght[:,n+1] = gaussjacobi(nqpts,0,0)
  return qnode, qwght

end

end


@doc raw"""
    Sphere{N} <: AbstractEmbeddedManifold{DefaultEmbeddingType}

The unit sphere manifold $𝕊^n$ represented by $n+1$-Tuples, i.e. in by
vectors in $ℝ^{n+1}$ of unit length. A sphere is an embedded manifold
and inherits the inner product of its embedding $\mathbb R^{n+1}$.
The tangent space at point p is given by

````math
T_p𝕊^n := \bigl\{ X ∈ ℝ^{n+1} : ⟨p,X⟩ = 0 \bigr \},
````

where $⟨\cdot,\cdot⟩$ denotes the inner product in the embedding $\mathbb R^{n+1}$.

This manifold is modeled as an embedded manifold to the [`Euclidean`](@ref), i.e.
several functions like the [`inner`](@ref inner(::Euclidean, ::Any...)) product and the
[`zero_tangent_vector`](@ref zero_tangent_vector(::Euclidean, ::Any...)) are inherited from the embedding.

# Constructor

    Sphere(n)

Generate $𝕊^{n} ⊂ ℝ^{n+1}$.
"""
struct Sphere{N} <: AbstractEmbeddedManifold{DefaultIsometricEmbeddingType} end

Sphere(n::Int) = Sphere{n}()

"""
    check_manifold_point(M, p; kwargs...)

Check whether `p` is a valid point on the [`Sphere`](@ref) `M`, i.e. is a vector
of length [`manifold_dimension`](@ref)`(M)+1` (approximately) of unit length.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_manifold_point(M::Sphere{N}, p; kwargs...) where {N}
    mpv = invoke(check_manifold_point, Tuple{supertype(typeof(M)), typeof(p)}, M, p; kwargs...)
    mpv === nothing || return mpv
    if !isapprox(norm(p), 1.0; kwargs...)
        return DomainError(
            norm(p),
            "The point $(p) does not lie on the sphere $(M) since its norm is not 1.",
        )
    end
    return nothing
end

"""
    check_tangent_vector(M, p, X; check_base_point = true, kwargs... )

Check whether `X` is a tangent vector to `p` on the [`Sphere`](@ref) `M`, i.e.
after [`check_manifold_point`](@ref)`(M,p)`, `X` has to be of same dimension as `p`
and orthogonal to `p`.
The optional parameter `check_base_point` indicates, whether to call [`check_manifold_point`](@ref)  for `p`.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_tangent_vector(
    M::Sphere{N},
    p,
    X;
    check_base_point = true,
    kwargs...,
) where {N}
    if check_base_point
        mpe = check_manifold_point(M, p; kwargs...)
        mpe === nothing || return mpe
    end
    mpv = invoke(
        check_tangent_vector,
        Tuple{supertype(typeof(M)), typeof(p), typeof(X)},
        M,
        p,
        X;
        check_base_point = false, # already checked above
        kwargs...
    )
    mpv === nothing || return mpv
    if !isapprox(abs(dot(p, X)), 0.0; kwargs...)
        return DomainError(
            abs(dot(p, X)),
            "The vector $(X) is not a tangent vector to $(p) on $(M), since it is not orthogonal in the embedding.",
        )
    end
    return nothing
end

decorated_manifold(M::Sphere) = Euclidean(representation_size(M)...; field=ℝ)

@doc raw"""
    distance(M::Sphere, p, q)

Compute the geodesic distance betweeen `p` and `q` on the [`Sphere`](@ref) `M`.
The formula is given by the (shorter) great arc length on the (or a) great circle
both `p` and `q` lie on.

````math
d_{𝕊^n}(p,q) = \arccos(⟨p,q⟩).
````
"""
distance(::Sphere, x, y) = acos(clamp(dot(x, y), -1, 1))

embed!(::Sphere, q, p) = (q .= p)

embed!(::Sphere, Y, p, X) = (Y .= X)

@doc raw"""
    exp(M::Sphere, p, X)

Compute the exponential map from `p` in the tangent direction `X` on the [`Sphere`](@ref)
`M` by following the great arc eminating from `p` in direction `X`.

````math
\exp_p X = \cos(\lVert X \rVert_p)p + \sin(\lVert X \rVert_p)\frac{X}{\lVert X \rVert_p}X,
````
where $\lVert X \rVert_p$ is the [`norm`](@ref norm(::Sphere,p,X)) on the
[`Sphere`](@ref) `M`.
"""
exp(::Sphere, ::Any...)

function exp!(M::Sphere, q, p, X)
    θ = norm(M, p, X)
    q .= cos(θ) .* p .+ usinc(θ) .* X
    return q
end

flat!(M::Sphere, ξ::CoTFVector, p, X::TFVector) = copyto!(ξ, X)

function get_basis(M::Sphere{N}, p, B::DiagonalizingOrthonormalBasis) where {N}
    A = zeros(N + 1, N + 1)
    A[1, :] = transpose(p)
    A[2, :] = transpose(B.frame_direction)
    V = nullspace(A)
    κ = ones(N)
    if !iszero(B.frame_direction)
        # if we have a nonzero direction for the geodesic, add it and it gets curvature zero from the tensor
        V = cat(B.frame_direction / norm(M, p, B.frame_direction), V; dims = 2)
        κ[1] = 0 # no curvature along the geodesic direction, if x!=y
    end
    Ξ = [V[:, i] for i = 1:N]
    return CachedBasis(B, κ, Ξ)
end

@doc raw"""
    get_coordinates(M::Sphere, p, X, B::DefaultOrthonormalBasis)

Represent the tangent vector `X` at point `p` from the [`Sphere`](@ref) `M` in
an orthonormal basis by rotating the vector `X` using the rotation matrix
$2\frac{q q^\mathrm{T}}{q^\mathrm{T} q} - I$ where $q = p + (1, 0, …, 0)$.
"""
function get_coordinates(M::Sphere{N}, p, X, B::DefaultOrthonormalBasis) where {N}
    if isapprox(p[1], 1)
        return X[2:end]
    else
        xp1 = p .+ ntuple(i -> ifelse(i == 1, 1, 0), N + 1)
        return (2*xp1*dot(xp1, X)/dot(xp1, xp1)-X)[2:end]
    end
end

function get_coordinates!(M::Sphere, Y, p, X, B::DefaultOrthonormalBasis)
    return copyto!(Y, get_coordinates(M, p, X, B))
end

function get_vector(M::Sphere{N}, p, X, B::DefaultOrthonormalBasis) where {N}
    p[1] ≈ 1 && return vcat(0, X)
    xp1 = p .+ ntuple(i -> ifelse(i == 1, 1, 0), N + 1)
    X0 = vcat(0, X)
    return 2 * xp1 * dot(xp1, X0) / dot(xp1, xp1) - X0
end

function get_vector!(M::Sphere, Y::AbstractVector, p, X, B::DefaultOrthonormalBasis)
    return copyto!(Y, get_vector(M, p, X, B))
end

@doc raw"""
    injectivity_radius(M::Sphere[, p])

Return the injectivity radius for the [`Sphere`](@ref) `M`, which is globally $π$.

    injectivity_radius(M::Sphere, x, ::ProjectionRetraction)

Return the injectivity radius for the [`ProjectionRetraction`](@ref) on the
[`Sphere`](@ref), which is globally $\frac{π}{2}$.
"""
injectivity_radius(::Sphere) = π
injectivity_radius(::Sphere, ::ExponentialRetraction) = π
injectivity_radius(::Sphere, ::ProjectionRetraction) = π / 2
injectivity_radius(::Sphere, ::Any) = π
injectivity_radius(::Sphere, ::Any, ::ExponentialRetraction) = π
injectivity_radius(::Sphere, ::Any, ::ProjectionRetraction) = π / 2
eval(quote
    @invoke_maker 1 Manifold injectivity_radius(M::Sphere, rm::AbstractRetractionMethod)
end)

@doc raw"""
    inverse_retract(M::Sphere, p, q, ::ProjectionInverseRetraction)

Compute the inverse of the projection based retraction on the [`Sphere`](@ref),
i.e. rearranging $p+X = q\lVert p+X\rVert_2$ yields
since $⟨p,X⟩ = 0$ and when $d_{𝕊^2}(p,q) ≤ \frac{π}{2}$ that

````math
\operatorname{retr}_p^{-1}(q) = \frac{q}{⟨p, q⟩} - p.
````
"""
inverse_retract(::Sphere, ::Any, ::Any, ::ProjectionInverseRetraction)

function inverse_retract!(::Sphere, X, p, q, ::ProjectionInverseRetraction)
    return (X .= q ./ dot(p, q) .- p)
end

@doc raw"""
    log(M::Sphere, p, q)

Compute the logarithmic map on the [`Sphere`](@ref) `M`, i.e. the tangent vector,
whose geodesic starting from `p` reaches `q` after time 1.
The formula reads for $x ≠ -y$

````math
\log_p q = d_{𝕊^n}(p,q) \frac{q-⟨p,q⟩ p}{\lVert q-⟨p,q⟩ p \rVert_2},
````

and a deterministic choice from the set of tangent vectors is returned if $x=-y$, i.e. for
opposite points.
"""
log(::Sphere, ::Any...)

function log!(S::Sphere, X, p, q)
    cosθ = dot(p, q)
    if cosθ ≈ -1 # appr. opposing points, return deterministic choice from set-valued log
        fill!(X, 0)
        if p[1] ≈ 1
            X[2] = 1
        else
            X[1] = 1
        end
        copyto!(X, X .- dot(p, X) .* p)
        X .*= π / norm(X)
    else
        cosθ = cosθ > 1 ? one(cosθ) : cosθ
        θ = acos(cosθ)
        X .= (q .- cosθ .* p) ./ usinc(θ)
    end
    return project!(S, X, p, X)
end

@doc raw"""
    manifold_dimension(M::Sphere)

Return the dimension of the [`Sphere`](@ref)`(n) `M`, i.e. $𝕊^n$, which is $\dim(𝕊^n) = n$.
"""
manifold_dimension(S::Sphere{N}) where {N} = N

"""
    mean(
        S::Sphere,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolationWithinRadius(π/2);
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::Manifold, args...)) of `x` using
[`GeodesicInterpolationWithinRadius`](@ref).
"""
mean(::Sphere, ::Any...)

function mean!(S::Sphere, p, x::AbstractVector, w::AbstractVector; kwargs...)
    return mean!(S, p, x, w, GeodesicInterpolationWithinRadius(π / 2); kwargs...)
end

"""
    normal_tvector_distribution(S::Sphere, p, σ)

Normal distribution in ambient space with standard deviation `σ`
projected to tangent space at `p`.
"""
function normal_tvector_distribution(S::Sphere, p, σ)
    d = Distributions.MvNormal(zero(p), σ)
    return ProjectedFVectorDistribution(TangentBundleFibers(S), p, d, project_vector!, p)
end

@doc raw"""
    project(M::Sphere, p)

Project the point `p` from the embedding onto the [`Sphere`](@ref) `M`.

````math
\operatorname{proj}_{𝕊^n}(p) = \frac{p}{\lVert p \rVert_2}.
````
"""
project(::Sphere, ::Any)

project!(S::Sphere, q, p) = copyto!(q, p./ norm(p))

@doc raw"""
    project(M::Sphere, p, X)

Project the point `X` onto the tangent space at `p` on the [`Sphere`](@ref) `M`.

````math
\operatorname{proj}_{p}(X) = X - ⟨p, X⟩p
````
"""
project(::Sphere, ::Any, ::Any)

project!(S::Sphere, Y, p, X) = (Y .= X .- dot(p, X) .* p)

@doc raw"""
    representation_size(M::Sphere)

Return the size points on the [`Sphere`](@ref) `M` are represented as, i.e.
for the `n`-dimensional [`Sphere`](@ref) it is vectors of size `(n+1,)`.
"""
@generated representation_size(::Sphere{N}) where {N} = (N + 1,)

@doc raw"""
    retract(M::Sphere, p, X, ::ProjectionRetraction)

Compute the retraction that is based on projection, i.e.

````math
\operatorname{retr}_p(X) = \frac{p+X}{\lVert p+X \rVert_2}
````
"""
retract(::Sphere, ::Any, ::Any, ::ProjectionRetraction)

function retract!(M::Sphere, q, p, X, ::ProjectionRetraction)
    q .= p .+ X
    return project!(M, q, q)
end

show(io::IO, ::Sphere{N}) where {N} = print(io, "Sphere($(N))")

"""
    uniform_distribution(M::Sphere, p)

Uniform distribution on given [`Sphere`](@ref) `M`. Generated points will be of
similar type as `p`.
"""
function uniform_distribution(M::Sphere, p)
    d = Distributions.MvNormal(zero(p), 1.0)
    return ProjectedPointDistribution(M, d, project!, p)
end

@doc doc"""
    vector_transport_to(M, p, X, q, ::ParallelTransport)

Compute the paralllel transport of the tangent vector `X` at `p` to `q`,
provided, the [`geodesic`](@ref) between `p` and `q` is unique

````math
P_{p←q}(X) = X - \frac{\langle \log_p q,X\rangle_p}{d^2_{𝕊^n}(p,q)}
\bigl(\log_xy + \log_yx \bigr).
````
"""
vector_transport_to(::Sphere, ::Any, ::Any, ::Any, ::Any, ::ParallelTransport)

function vector_transport_to!(M::Sphere, Y, p, X, q, ::ParallelTransport)
    X_pq = log(M, p, q)
    Xl = norm(M, p, X_pq)
    copyto!(Y, X)
    if Xl > 0
        factor = 2 * dot(X, q) / (norm(p + q)^2)
        Y .-= factor .* (p .+ q)
    end
    return Y
end


using Test, PolynomialBases

ufunc(x) = sinpi(x)

for basis_type in subtypes(PolynomialBases.NodalBasis{PolynomialBases.Line})
    for p in 1:5
        basis = basis_type(p)
        u = ufunc.(basis.nodes)
        D, M, R, B, MinvRtB = utility_matrices(basis)

        @test integrate(u, basis) ≈ sum(M * u)
        Ru = R*u
        @test Ru[1] ≈ interpolate(-1, u, basis)
        @test Ru[2] ≈ interpolate(+1, u, basis)

        @inferred includes_boundaries(basis)
        if includes_boundaries(basis) == Val{true}()
            @test Ru[1] ≈ u[1]
            @test Ru[end] ≈ u[end]
        else
            @test !(Ru[1] ≈ u[1])
            @test !(Ru[end] ≈ u[end])
        end

        # SBP property
        if satisfies_sbp(basis) == Val{true}()
            @test M*D + D'*M ≈ R'*B*R
        end
    end
end


import RRSTExperiments: InputEdge

function get_objective_value(w::Dict{Tuple{Int64, Int64}, Float64}, t_x::Vector{Edge{Int}})
    return sum(map(e -> haskey(w, (e.src, e.dst)) ? w[(e.src, e.dst)] : 0, t_x))
  end

  function get_solution(t_x::Vector{Edge{Int}})
    return map(e -> (e.src, e.dst), t_x)
  end

function solve_minmax_st(n::Int64, A::Vector{InputEdge})
    weight_map::Dict{Tuple{Int64, Int64}, Float64} = Dict()
    g = SimpleGraph(n)
    for e in A
      add_edge!(g, e.i, e.j)
      push!(weight_map, minmax(e.i, e.j) => (e.C + e.c + e.d))
    end

    w_mat = spzeros(Float64, n, n)
    for e in A
      w_mat[e.i, e.j] = e.C + e.c + e.d
    end

    t_x = kruskal_mst(g, w_mat)

    obj_value = get_objective_value(weight_map, t_x)
    x = get_solution(t_x)
    return obj_value, x
end


using Revise
using LinearAlgebra
using ControlSystems
using Plots
using Printf
#pgfplots()
plotly()
# constant definitions
ρ_∞ = 1.112
θ_0 = -2
V_∞ = 28
m = 318
S = 6.11
l_t = 4.63
I_yy = 432
S_t = 1.14
l_w = 0.3
C_Lαw = 5.55
C_Lαt = 4.3
C_Lδt = 0.45
b = 15
T_0 = 0
e = 0.8
g = 9.81

function getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy, S_t, l_w, C_Lαw, C_Lαt, C_Lδt, b, T_0, e, g)
    C_Lα = C_Lαw + S_t * C_Lαt / S
    q_∞ = 0.5 * ρ_∞ * V_∞^2

    # state matrix
    A = [
          [
          2 * ( m * g * sind(θ_0) -T_0) / (m * V_∞)
          g * cosd(θ_0) * (1 - 2 * S * C_Lα / (π * b^2 * e) ) / V_∞
          0
          -g * cosd(θ_0)
          ]';

          [
          -2 * g * cosd(θ_0) / V_∞
          (m * g * sind(θ_0) - T_0) / (m * V_∞) - ρ_∞ * V_∞ * S * C_Lα / (2 * m)
          V_∞ - ρ_∞ * V_∞ * S_t * l_t * C_Lαt / (2 * m)
          -g * sind(θ_0)
          ]';

          [
          0
          ρ_∞ * V_∞ * (S * l_w * C_Lαw - S_t * l_t * C_Lαt) / (2 * I_yy)
          -ρ_∞ * V_∞ * S_t * l_t^2 * C_Lαt / (2 * I_yy)
          0
          ]';

          [
          0
          0
          1
          0
          ]'
        ]

    # input matrix
     B = [
            0
            -q_∞ * S_t * C_Lδt / m
            -q_∞ * S_t * l_t * C_Lδt / I_yy
            0
         ]

    # sensory matrix
    C = Matrix(I,4,4)

    # direct term
    D = 0
    return A,B,C,D
end

# Question 1
A,B,C,D = getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy, S_t, l_w, C_Lαw, C_Lαt, C_Lδt, b, T_0, e, g)

# eigenvalue plots
λ = eigvals(A)
scatter(λ, xlabel="Re(λ)", ylabel="Im(λ)", legend = false)

# Question 2
# state space
sys = ss(A,B,C,D)
stepplot(sys, 200, legend = false)
plot!(ylabel = ["V_x (m/s)" "V_z (m/s)" "ω_y (deg/s)" "θ (deg)"], xlabel = ["" "" "" "Time (s)"], title = ["Step Response" "" "" ""], size = [600,600])

phugω = imag(λ[4]) # phugiod frequency
@printf("The phugoid frequency is %f and the time period is %f", phugω, 2 * π / phugω)

# Question 3
λ_C_Lαw = [
           eigvals(getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy, S_t, l_w, C_Lαw*0.8, C_Lαt, C_Lδt, b, T_0, e, g)[1])';
           eigvals(getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy, S_t, l_w, C_Lαw*1.2, C_Lαt, C_Lδt, b, T_0, e, g)[1])'
          ]'

λ_C_Lαt = [
           eigvals(getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy, S_t, l_w, C_Lαw, C_Lαt*0.8, C_Lδt, b, T_0, e, g)[1])';
           eigvals(getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy, S_t, l_w, C_Lαw, C_Lαt*1.2, C_Lδt, b, T_0, e, g)[1])'
          ]'

λ_I_yy = [
           eigvals(getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy*0.8, S_t, l_w, C_Lαw, C_Lαt, C_Lδt, b, T_0, e, g)[1])';
           eigvals(getABCD(ρ_∞, θ_0, V_∞, m, S, l_t, I_yy*1.2, S_t, l_w, C_Lαw, C_Lαt, C_Lδt, b, T_0, e, g)[1])'
          ]'

λ_V_∞ = [
           eigvals(getABCD(ρ_∞, θ_0, V_∞*0.8, m, S, l_t, I_yy, S_t, l_w, C_Lαw, C_Lαt, C_Lδt, b, T_0, e, g)[1])';
           eigvals(getABCD(ρ_∞, θ_0, V_∞*1.2, m, S, l_t, I_yy, S_t, l_w, C_Lαw, C_Lαt, C_Lδt, b, T_0, e, g)[1])'
          ]'

# pole plots
function q3polePlot(λ,λ_ref)
    scatter([λ[:,1][1:2] λ[:,1][3:4]], xlabel="Re(λ)", ylabel="Im(λ)",markercolor = :red, layout = 2, title=["SPPO" "PHUGOID"], labels = ["80%" ""])
    scatter!([λ[:,2][1:2] λ[:,2][3:4]], xlabel="Re(λ)", ylabel="Im(λ)",markercolor = :green, layout = 2, title=["SPPO" "PHUGOID"], labels = ["120%" ""])
    scatter!([λ_ref[1:2] λ_ref[3:4] ], xlabel="Re(λ)", ylabel="Im(λ)",markercolor = :blue, layout = 2, title=["SPPO" "PHUGOID"], labels = ["Unperturbed" ""])
    scatter!(size = [800,400])
end

q3polePlot(λ_C_Lαw, λ)
q3polePlot(λ_C_Lαt, λ)
q3polePlot(λ_I_yy, λ)
q3polePlot(λ_V_∞, λ)

# freq and damp plots

function q3fdPlot(λ, λ_ref, mcolor, mname)
    λ = [λ λ_ref]
    sppo_freq = [sqrt(real(λ[:,i][1])^2 + imag(λ[:,i][1])^2) for i = 1:3]
    phug_freq = [sqrt(real(λ[:,i][3])^2 + imag(λ[:,i][3])^2) for i = 1:3]
    sppo_damp = [abs(real(λ[:,i][1])) / sppo_freq[i] for i = 1:3]
    phug_damp = [abs(real(λ[:,i][3])) / phug_freq[i] for i = 1:3]

    scatter!([sppo_freq[1]/sppo_freq[3] phug_freq[1]/phug_freq[3]], [sppo_damp[1]/sppo_damp[3] phug_damp[1]/phug_damp[3]], xlabel="Natural frequency (rad/s)", ylabel="Damping ratio",markercolor = mcolor, layout = 2, title=["SPPO" "PHUGOID"], markershape = :cross, labels = ["80% "*mname ""])
    scatter!([sppo_freq[2]/sppo_freq[3] phug_freq[2]/phug_freq[3]], [sppo_damp[2]/sppo_damp[3] phug_damp[2]/phug_damp[3]], xlabel="Natural frequency (rad/s)", ylabel="Damping ratio",markercolor = mcolor, layout = 2, title=["SPPO" "PHUGOID"], markershape = :xcross, labels = ["120% "*mname ""])
end

scatter([1 1], [1 1], xlabel="Natural frequency (rad/s)", ylabel="Damping ratio",markercolor = :black, layout = 2, title=["SPPO" "PHUGOID"], labels = ["Unperturbed" ""], size = [800,400])
q3fdPlot(λ_C_Lαw, λ, :blue, "C_Lαw")
q3fdPlot(λ_C_Lαt, λ, :red, "C_Lαt")
q3fdPlot(λ_I_yy, λ, :green, "I_yy")
q3fdPlot(λ_V_∞, λ, :pink, "V_∞")

# Question 4
l_β = (l_w + l_t) / 2 # position of spring
σ = 5:0.1:50
k_β = 0.5 * ρ_∞ * V_∞^2 * S_t * l_t * σ

A_s = -0.5 * ρ_∞ * V_∞ ^2 * S_t * l_β * C_Lαt
A_rs = ρ_∞ * V_∞^2 * S_t * C_Lαt / (2 * m) * [0,1,l_t * m / I_yy,0]
A_sr = 0.5 * ρ_∞ * V_∞ * S_t * l_β * C_Lαt * [0 1 l_t 0]

A_k = [A + A_rs * A_sr / (k_β[i] - A_s) for i = 1:length(k_β)] # state matrix with k spring
λ_k = map((x) -> eigvals(x), A_k)
kLength = length(λ_k)
# get modes

sppo_k = [λ_k[i][2] for i = 1:kLength]
phug_k = [λ_k[i][4] for i = 1:kLength]

sppo_freq = sqrt(real(λ[1])^2 + imag(λ[1])^2)
phug_freq = sqrt(real(λ[3])^2 + imag(λ[3])^2)
sppo_damp = abs(real(λ[1])) / sppo_freq
phug_damp = abs(real(λ[3])) / phug_freq

sppo_freq_k = [sqrt(real(λ_k[i][1])^2 + imag(λ_k[i][1])^2) for i = 1:kLength]
phug_freq_k = [sqrt(real(λ_k[i][3])^2 + imag(λ_k[i][3])^2) for i = 1:kLength]
sppo_damp_k = [abs(real(λ_k[i][1])) / sppo_freq_k[i] for i = 1:kLength]
phug_damp_k = [abs(real(λ_k[i][3])) / phug_freq_k[i] for i = 1:kLength]

# plot poles
scatter([sppo_k phug_k], xlabel="Re(λ)", ylabel="Im(λ)", layout = 2, legend = false, title=["SPPO" "PHUGOID"], size = [800,400], markercolor = :red)
# frequency and damping plots
scatter([sppo_freq_k/sppo_freq phug_freq_k/phug_freq], [sppo_damp_k/sppo_damp phug_damp_k/phug_damp], xlabel="Natural frequency (rad/s)", ylabel="Damping ratio",markercolor = :red, layout = 2, title=["SPPO" "PHUGOID"], legend = false, size = [800,400])

# Question 5a
Q = C

# starting value

R=10^( -3.2 + 0.2 * 1 )
K = lqr(sys, Q, R)
P = feedback(sys,K)
λ_closed = pole(P)
scatter([λ_closed[1:2] λ_closed[3:4]], xlabel="Re(λ)", ylabel="Im(λ)", layout = 2, markercolor = :red, labels = ["closed loop poles" ""], size = [800,400], title = ["SPPO" "PHUGOID"])
for i = 2:25
    R=10^( -3.2 + 0.2 * i )
    K = lqr(sys, Q, R)
    P = feedback(sys,K)
    λ_closed = pole(P)
    scatter!([λ_closed[1:2] λ_closed[3:4]], xlabel="Re(λ)", ylabel="Im(λ)", layout = 2, markercolor = :red, labels = ["" ""])
end
scatter!([λ[1:2] λ[3:4]], xlabel="Re(λ)", ylabel="Im(λ)", layout = 2, markercolor = :blue, labels = ["open loop poles" ""], markershape = :cross)

# Question 5b
R = 1
Q = Diagonal([5,0,0,5])
K = lqr(sys, Q, R)
P = feedback(sys,K)
stepplot(P,200, legend = false)
plot!(ylabel = ["V_x (m/s)" "V_z (m/s)" "ω_y (deg/s)" "θ (deg)"], xlabel = ["" "" "" "Time (s)"], title = ["Step Response" "" "" ""], size = [600,600])
#png("Files/MechanicsOfFlightCoursework/q5bstep.png")


function CND(X)

  y = abs(X)

  if y > 37
    cum = 0
  else
    Exponential = exp(-y^2/2)

    if y < 7.07106781186547
      SumA = 0.0352624965998911 * y + 0.700383064443688
      SumA = SumA * y + 6.37396220353165
      SumA = SumA * y + 33.912866078383
      SumA = SumA * y + 112.079291497871
      SumA = SumA * y + 221.213596169931
      SumA = SumA * y + 220.206867912376
      SumB = 0.0883883476483184 * y + 1.75566716318264
      SumB = SumB * y + 16.064177579207
      SumB = SumB * y + 86.7807322029461
      SumB = SumB * y + 296.564248779674
      SumB = SumB * y + 637.333633378831
      SumB = SumB * y + 793.826512519948
      SumB = SumB * y + 440.413735824752
      cum = Exponential * SumA / SumB
    else
      SumA = y + 0.65
      SumA = y + 4 / SumA
      SumA = y + 3 / SumA
      SumA = y + 2 / SumA
      SumA = y + 1 / SumA
      cum = Exponential / (SumA * 2.506628274631)

      end
  end

  if X > 0
    cum = 1 - cum
  else
          cum
  end

end


## MIT license (c) 2019 by Andrew Lyasoff
##
## note the formula
## $\int_a^bf(t)dt={1\over2}(b-a)\int_{-1}^{+1}f\bigl({1\over2}(b-a)s+{1\over2}(a+b)\bigr)d s$
##(Used by FastGaussQuadrature)
##

begin
    @everywhere using SpecialFunctions
    @everywhere using Interpolations
    @everywhere using Roots
    @everywhere using FastGaussQuadrature

    @everywhere function EUcall(S::Float64,t::Float64,K::Float64,σ::Float64,r::Float64,δ::Float64)
        local v1,v2
        v1=(δ-r+σ^2/2.0)*t+log(K/S)
        v2=(-δ + r + σ^2/2.0)*t - log(K/S)
        v3=σ*sqrt(2.0*t)
        return -K*(exp(-t*r)/2)+K*(exp(-t*r)/2)*erf(v1/v3)+S*(exp(-t*δ)/2)*erf(v2/v3)+S*(exp(-t*δ)/2)
    end

    @everywhere function F(ϵ::Int64,t::Float64,u::Float64,v::Float64,r::Float64,δ::Float64,σ::Float64)
        v1=(r-δ+ϵ*σ^2/2)*(u-t)-log(v)
        v2=σ*sqrt(2*(u-t))
        return 1.0+erf(v1/v2)
    end

    @everywhere function ah(t::Float64,z::Float64,r::Float64,δ::Float64,σ::Float64,f)
        return (exp(-δ*(z-t))*(δ/2)*F(1,t,z,f(z)/f(t),r,δ,σ))
    end

    @everywhere function bh(t::Float64,z::Float64,r::Float64,δ::Float64,σ::Float64,K::Float64,f)
        return (exp(-r*(z-t))*(r*K/2)*F(-1,t,z,f(z)/f(t),r,δ,σ))
    end

    @everywhere function mke_grd(bgn::Float64,step::Float64,size::Int64)
        return bgn:step:(bgn+step+(size-1)*step)
    end

    @everywhere function workerF(loc_range::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},r::Float64,δ::Float64,σ::Float64,T::Float64,K::Float64,nds::Array{Float64,1},wghts::Array{Float64,1},f)
        local loc,LRT,aaa,bbb,an,bn,lnlr
        lnlr=length(loc_range)
        loc=zeros(lnlr)
        for i=lnlr:-1:1
            t=loc_range[i]
            an=[(1/2)*(T-t)*ah(t,(1/2)*(T-t)*s+(1/2)*(t+T),r,δ,σ,f) for s in nds]
            bn=[(1/2)*(T-t)*bh(t,(1/2)*(T-t)*s+(1/2)*(t+T),r,δ,σ,K,f) for s in nds]
            aaa=wghts'*an
            bbb=wghts'*bn;
            LRT=find_zero(x->x-K-EUcall(x,T-t,K,σ,r,δ)-aaa*x+bbb,(K-10,K+20));
            loc[i]=LRT
        end
        return loc
    end

    function cmbn(varg::Array{Array{Float64,1},1})
        if length(varg)>1
            loc=varg[1]
            for i=2:length(varg)
                loc=vcat(loc,varg[i])
            end
            return loc
        else
            return varg
        end
    end

end

function masterF(prnmb::Int64, start_iter::Int64,nmb_of_iter::Int64,conv_tol::Float64,K::Float64,σ::Float64,δ::Float64,r::Float64,T::Float64,Δ::Float64,no_grd_pnts::Int64,vls::Array{Float64,1},nds::Array{Float64,1},wghts::Array{Float64,1})
    local no_iter,conv_check,absc,absc0,val,valPrev,loc,f
    absc=range(0.0,length=pnmb*no_grd_pnts+1,stop=T)
    absc0=range(0.0,length=pnmb*no_grd_pnts,stop=(T-Δ))
    all_grid=[range((T/prnmb)*i,length=no_grd_pnts,stop=((T/prnmb)*(i+1)-Δ)) for i=0:(pnmb-1)]
    val=vls;
    f=CubicSplineInterpolation(absc,val,extrapolation_bc = Interpolations.Line())
    for i=2:(prnmb+1)
        remotecall_fetch(()->K, i)
        remotecall_fetch(()->σ, i)
        remotecall_fetch(()->δ, i)
        remotecall_fetch(()->r, i)
        remotecall_fetch(()->T, i)
        remotecall_fetch(()->all_grid, i)
        remotecall_fetch(()->f, i)
        remotecall_fetch(()->nds, i)
        remotecall_fetch(()->wghts, i)
    end
    no_iter=start_iter;
    conv_check=1000.0
    while no_iter<nmb_of_iter&&conv_check>conv_tol
        valPrev=val;
        no_iter+=1
        lst=[@spawnat i workerF(all_grid[myid()-1],r,δ,σ,T,K,nds,wghts,f) for i=2:(prnmb+1)]
        val0=[fetch(lst[i]) for i=1:prnmb];
        val=cmbn(val0)
        push!(val,max(K,K*(r/δ)))
        f=CubicSplineInterpolation(absc,val,extrapolation_bc = Interpolations.Line())
        for i=2:(prnmb+1)
            remotecall_fetch(()->f, i)
        end
        conv_check=maximum(abs.(valPrev-val))
    end
    return absc,val,conv_check,no_iter
end


import Distributions
import Base.rand

function rand(gp::FullModel, locations)
	Kpp = kernel(gp.k, locations, locations)
	Kdp = kernel(gp.k, gp.X, locations)
	K = Kpp - Kdp' * (gp.C \ Kdp)
	mean = predict(gp, locations)[1]
	rand(Distributions.MvNormal(mean, K))
end


using Distributions, Random, Test

## Helper: check a decision by simulation
function check_by_sim(d :: ExtremeValueDecision{T},  value_ixM :: Array{T,2},
    prob_ixM :: Array{T,2},  eVal_iV :: Array{T,1})    where T <: AbstractFloat

    rng = MersenneTwister(123);
    nTypes, n = size(value_ixM);
    success = true;

    for iType = 1 : nTypes
        success1 = check_one_by_sim(d, vec(value_ixM[iType,:]),
            vec(prob_ixM[iType,:]),  eVal_iV[iType],  rng);
        success = success && success1;
    end

    return success
end

## Check one type by simulation
# Returns a Bool (success or failure)
function check_one_by_sim(d :: ExtremeValueDecision{F},  valueV :: Vector{F},
    probV :: Vector{F},  eVal :: F,  rng :: AbstractRNG)  where  F <: AbstractFloat

    # Random vars for all alternatives
    n = length(valueV);
    nSim = Int64(5e5);
    simProbV, simVal = sim_extreme_value(d, valueV, nSim, rng);

    success = true;
    if !isapprox(simProbV, probV, atol = 5e-3)
       println("True probs:  ", probV);
       println("Sim probs:   ", simProbV);
       success = false;
    end

    # Check expected value
    if abs(simVal ./ eVal - 1.0) > 5e-3
       @warn "Discrepancy in values (demeaned = $(demeaned(d))): $simVal,  $eVal"
       success = false;
    end

    return success
end

function sim_extreme_value(d :: ExtremeValueDecision{F},  valueV :: Vector{F},
    nSim :: Integer, rng :: AbstractRNG) where F

    # Random vars for all alternatives
    n = length(valueV);
    maxV = zeros(F, nSim);
    countV = zeros(Int, n);
    for i1 = 1 : nSim
        randV = draw_gumbel_shocks(rng, pref_scale(d), n; demeaned = demeaned(d));
        # Shocks must be added, not subtracted.
        maxV[i1], maxIdx = findmax(valueV .+ randV);
        countV[maxIdx] += 1;
    end
    simProbV = countV ./ nSim;
    simVal = mean(maxV);
    return simProbV, simVal
end

function one_option_test()
    @testset "One option" begin
        prefScale = 1.4;
        nTypes = 3;
        value_iV = collect(range(32.0, 48.0, length = nTypes));

        # Demeaned
        d = ExtremeValueDecision(prefScale, true, true);
        prob_iV, eVal_iV = extreme_value_decision(d, value_iV);
        @test prob_iV == ones(Float64, nTypes)
        @test isapprox(eVal_iV, value_iV, rtol = 0.02)

        # Not demeaned
        d2 = ExtremeValueDecision(prefScale, false, true);
        prob_iV, eVal2_iV = extreme_value_decision(d2, value_iV);
        @test prob_iV == ones(Float64, nTypes)
        # If Gumbel shocks are subtracted (they have positive mean), demeaned expected values should be higher
        @test isapprox(eVal2_iV .- prefScale * EconLH.EulerConst, eVal_iV)
    end
end

function one_type_test()
    @testset "One type" begin
        prefScale = 0.8;
        nx = 5;
        valueV = (-0.2, 0.3, 0.9, -1.8, 2.0);
        @assert length(valueV) == nx;
        probV, eVal = extreme_value_decision_one(valueV, prefScale; demeaned = true);
        prob2V, eVal2 = extreme_value_decision_one([valueV...], prefScale;
            demeaned = true);
        @test isapprox(probV, prob2V);
        @test isapprox(eVal, eVal2);
    end
end

function many_types_test()
    @testset "Many types" begin
        prefScale = 0.8;
        nx = 3;
        nTypes = 5;
        value_ixV = [collect(LinRange(1.0, 2.0, nx) .+ 0.1 * j)  for j = 1 : nTypes];

        # Tuple input
        valueTuples = (value_ixV..., );
        prob1_ixV, eVal1_jV = extreme_value_decision(valueTuples, prefScale);

        # Matrix input. Columns are alternative
        value_ixM = zeros(nTypes, nx);
        for j = 1 : nTypes
            value_ixM[j,:] = value_ixV[j];
        end
        prob_ixM, eVal_jV = extreme_value_decision(value_ixM, prefScale);

        for j = 1 : nTypes
            @test isapprox(prob1_ixV[j], prob_ixM[j,:]);
        end
        @test isapprox(eVal_jV, eVal1_jV);
    end
end

function check_by_sim_test()
    @testset "Check by simulation" begin
        for demeaned in [true, false]
            n = 4;
            prefScale = 1.4;
            d = ExtremeValueDecision(prefScale, demeaned, true);

            # Row vector of values (as matrix)
            value_xV = collect(range(50.0, 55.0, length = n)');

            for nTypes in [1, 3]
                if nTypes == 1
                    value_ixM = value_xV;
                else
                    value_ixM = collect(range(1.0, 0.9, length = nTypes)) *
                        collect(range(50.0, 55.0, length = n)');
                end

                prob_ixM, eVal_iV = EconLH.extreme_value_decision(d, value_ixM);
                # This makes sure that there is something meaningful to simulate
                @test all(prob_ixM .< 0.95)
                @test check_by_sim(d, value_ixM, prob_ixM, eVal_iV);
            end
        end
    end
end

function draw_gumbel_test()
    @testset "Draw gumbel" begin
        rng = MersenneTwister(123);
        prefScale = 0.4;
        nSim = Int(1e3);

        shockV = draw_gumbel_shocks(rng, prefScale, nSim; demeaned = true);
        @test isapprox(mean(shockV), 0.0, atol = 0.05)
        @test size(shockV) == (nSim,)

        shockM = draw_gumbel_shocks(rng, prefScale, (nSim, 2); demeaned = false);
        @test isapprox(mean(shockM), prefScale * EconLH.EulerConst, atol = 0.05)
        @test size(shockM) == (nSim, 2)
    end
end

@testset "ExtremeValueDecision" begin
    one_option_test();
    one_type_test();
    many_types_test();
    check_by_sim_test();
    draw_gumbel_test();
end

# ---------------


export ¬, ∧, ∨, →, ⟷
export AbstractExpression, LogicalSymbol, LogicalExpression, istree, isnode, name, metadata, variables, operation, operations, arguments, parents, left, right, isassociative, iscommutative
export isunary, isbinary, argument_count
export @symbols, @unique_symbols

"""
    AbstractExpression

The abstract type which all expression trees are built from.
"""
abstract type AbstractExpression end

"""
    LogicalSymbol(name::Symbol, metadata::Any=nothing)

Represent a logical symbolically with a provided name. By default the attached metadata is set to `nothing`.
"""
struct LogicalSymbol <: AbstractExpression
    name::Symbol
    metadata::Any
end
LogicalSymbol(name::Symbol) = LogicalSymbol(name, nothing)
"""
    name(sym::LogicalSymbol)

The name of the symbol provided at instantiation. Equivalent to `symbol(sym)`.
"""
name(sym::LogicalSymbol) = sym.name
"""
    metadata(sym::LogicalSymbol)

The metadata of the symbol provided at instantiation, if any. Returns `nothing` if none was provided.
"""
metadata(sym::LogicalSymbol) = sym.metadata

Base.show(io::IO, sym::LogicalSymbol) = print(io, string(sym.name))
Base.hash(sym::LogicalSymbol, h::UInt) = hash(sym.name, hash(sym.metadata, h))
Base.:(==)(sym1::LogicalSymbol, sym2::LogicalSymbol) = sym1.name == sym2.name && isequal(metadata(sym1), metadata(sym2))
Base.isless(sym1::LogicalSymbol, sym2::LogicalSymbol) = Base.isless(sym1.name, sym2.name)

# fake copy methods since symbols are immutable
Base.copy(sym::LogicalSymbol) = sym
Base.deepcopy(sym::LogicalSymbol) = LogicalSymbol(name(sym), deepcopy(metadata(sym)))

# convenience macros
"""
Define any number of [`LogicalSymbols`](@ref Deductive.LogicalSymbol) with the names provided.

# Examples
```julia-repl
julia> @symbols a  # defines symbol `a` in the current scope

julia> @symbols b c d  # defines symbols `a`, `b`, and `c`

julia> @symbols α β γ  # defines symbols `α`, `β`, and `γ`
```
"""
macro symbols(syms...)
    definitions = [:(
        $(esc(sym)) = LogicalSymbol($(esc(Symbol))($(string(sym))))
    ) for sym ∈ syms]
    quote
        $(definitions...)
        nothing
    end
end

_unique_index = 1
macro unique_symbols(syms...)
    global _unique_index

    unique_names = [Symbol("u" * subscript_number(_unique_index + i)) for i ∈ 0:(length(syms) - 1)]
    definitions = [:(
        $(esc(sym)) = LogicalSymbol($(esc(Symbol))($(string(unique_name))))
    ) for (sym, unique_name) ∈ zip(syms, unique_names)]

    _unique_index += length(syms)

    quote
        $(definitions...)
        nothing
    end
end

"""
    LogicalOperation(bool_fn::Function, name::Symbol, argument_count::Int, associative::Bool, commutative::Bool)

Defines a logical operator which can be used to form expressions with. Usually you don't need to define your own
operators, just use the ones built-in to the package if possible.
"""
struct LogicalOperation
    bool_fn::Function
    name::Symbol
    argument_count::Int
    associative::Bool
    commutative::Bool
end
LogicalOperation(bool_fn::Function, name::Symbol, argument_count::Int) = LogicalOperation(bool_fn, name, argument_count, false, false)
name(op::LogicalOperation) = op.name
"""
    argument_count(op::LogicalOperation)

The number of arguments which the given operation expects to receive.
"""
argument_count(op::LogicalOperation) = op.argument_count
"""
    isunary(op::LogicalOperation)

True if the given operation only receives one argument, false otherwise.
"""
isunary(op::LogicalOperation) = argument_count(op) == 1
"""
    isbinary(op::LogicalOperation)

True if the given operation receives two arguments, false otherwise.
"""
isbinary(op::LogicalOperation) = argument_count(op) == 2
"""
    isassociative(op::LogicalOperation)

True if the given operation is associative, false otherwise.
"""
isassociative(op::LogicalOperation) = op.associative
"""
    iscommutative(op::LogicalOperation)

True if the given operation is commutative, false otherwise.
"""
iscommutative(op::LogicalOperation) = op.commutative
Base.show(io::IO, op::LogicalOperation) = print(io, string(op.name))
Base.hash(op::LogicalOperation, h::UInt) = hash(op.name, hash(argument_count(op), hash(op.associative, hash(op.commutative, h))))
function Base.:(==)(op1::LogicalOperation, op2::LogicalOperation)
    op1.name == op2.name && argument_count(op1) == argument_count(op2) && isassociative(op1) == isassociative(op2) && iscommutative(op1) == iscommutative(op2)
end

"""
    (op::LogicalOperation)(args::AbstractExpression...)

This function is the foundation of expression building, as it allows instances of `LogicalOperations` to be treated as real
Julia functions. When called this function returns a `LogicalExpression` representing the application of the given operator
on the provided arguments.
"""
function (op::LogicalOperation)(args::AbstractExpression...)
    if length(args) != argument_count(op)
        throw(ErrorException("Invalid argument count $(length(args)). Expected $(argument_count(op)) arguments."))
    end
    return LogicalExpression(AbstractExpression[args...], op)
end
(op::LogicalOperation)(args::Bool...) = op.bool_fn(args...)
(op::LogicalOperation)(args::BitVector) = op.bool_fn(args...)

recursivevariables(args::Vector{AbstractExpression}) = reduce(∪, variables.(args))
recursiveoperations(args::Vector{AbstractExpression}, rootop::LogicalOperation) = reduce(∪, operations.(args)) ∪ Set([rootop])
"""
    LogicalExpression(arguments::Vector{AbstractExpression}, operation::LogicalOperation)

Constructs an expression with given arguments and logical operation. Please refrain from using this syntax, instead using
the "operators as functions" syntax, where a `LogicalOperation` instance can be called to produce a `LogicalExpression`.
"""
mutable struct LogicalExpression <: AbstractExpression
    arguments::Vector{AbstractExpression}
    operation::LogicalOperation

    parents::Set{LogicalExpression}

    cached_variables::Set{LogicalSymbol}
    cached_operations::Set{LogicalOperation}
    cached_variables_valid::Bool
    cached_operations_valid::Bool

    function LogicalExpression(arguments::Vector{AbstractExpression}, operation::LogicalOperation)
        expr = new(arguments, operation, Set{LogicalExpression}(), recursivevariables(arguments), recursiveoperations(arguments, operation), true, true)
        add_to_parents!(expr)
        expr
    end
end
"""
    operation(expr::LogicalExpression)

The operation which is performed on the arguments of the given `LogicalExpression`.
"""
operation(expr::LogicalExpression) = getfield(expr, :operation)
"""
    parents(expr::LogicalExpression)

The parent expressions of a given subexpression. If an expression is never used within another, it will have no parents.
In most expressions there will only be one parent, but it is possible for an expression assigned to a variable to have
multiple parents by using it as a "named subexpression".
"""
parents(expr::LogicalExpression) = getfield(expr, :parents)

"""
    add_to_parents!(expr::LogicalExpression)

Adds an expression to the `parents` set of all its arguments. For internal use only.
"""
add_to_parents!(expr::LogicalExpression) = add_to_parents!(expr, arguments(expr))
function add_to_parents!(expr::LogicalExpression, relevant_args::Vector{AbstractExpression})
    for arg ∈ relevant_args
        if arg isa LogicalExpression
            push!(parents(arg), expr)
        end
    end
end

"""
    remove_from_parents!(expr::LogicalExpression)

Removes an expression from the `parents` set of all its arguments. For internal use only.
"""
remove_from_parents!(expr::LogicalExpression) = remove_from_parents!(expr, arguments(expr))
function remove_from_parents!(expr::LogicalExpression, relevant_args::Vector{AbstractExpression})
    for arg ∈ relevant_args
        if arg isa LogicalExpression
            delete!(parents(arg), expr)
        end
    end
end
metadata(::LogicalExpression) = nothing
"""
    left(expr::LogicalExpression)

If the expression is binary, this method returns the left-hand operand.
"""
left(expr::LogicalExpression) = isbinary(operation(expr)) ? arguments(expr)[1] : throw(ErrorException("Operation $(operation(expr)) is not binary"))
"""
    right(expr::LogicalExpression)

If the expression is binary, this method returns the right-hand operand.
"""
right(expr::LogicalExpression) = isbinary(operation(expr)) ? arguments(expr)[2] : throw(ErrorException("Operation $(operation(expr)) is not binary"))
"""
    isassociative(expr::LogicalExpression)

Checks if the entire expression is associative based on the associative property of its constituent operations.
"""
isassociative(expr::LogicalExpression) = length(operations(expr)) == 1 && isassociative(operation(expr))
"""
    iscommutative(expr::LogicalExpression)

Checks if the entire expression is commutative based on the commutative property of its constituent operations.
"""
iscommutative(expr::LogicalExpression) = length(operations(expr)) == 1 && iscommutative(operation(expr))
Base.hash(expr::LogicalExpression, h::UInt) = hash(arguments(expr), hash(operation(expr), h))
Base.:(==)(expr1::LogicalExpression, expr2::LogicalExpression) = operation(expr1) == operation(expr2) && all(arguments(expr1) .== arguments(expr2))

# copy methods
Base.copy(expr::LogicalExpression) = LogicalExpression(Vector{AbstractExpression}(arguments(expr)), operation(expr))
Base.deepcopy(expr::LogicalExpression) = LogicalExpression(Vector{AbstractExpression}(deepcopy.(arguments(expr))), operation(expr))

# mutability methods
function Base.setproperty!(expr::LogicalExpression, name::Symbol, x)
    if name == :arguments
        remove_from_parents!(expr)
        setfield!(expr, name, x)
        add_to_parents!(expr)
        expr.cached_variables_valid = false
    elseif name == :operation
        @assert argument_count(operation(expr)) == argument_count(x) "cannot assign operation with $(argument_count(x)) arguments to an expression with $(argument_count(operation(expr)))"
        setfield!(expr, name, x)
        expr.cached_operations_valid = false
    elseif name == :cached_variables_valid || name == :cached_operations_valid
        setfield!(expr, name, x)
        if !x  # cache invalidation, so propagate the invalidation up the expression tree
            for parent ∈ parents(expr)
                setproperty!(parent, name, x)
            end
        end
    elseif name == :cached_variables || name == :cached_operations
        throw(ErrorException("type LogicalExpression has protected field `$(name)`"))
    else
        throw(ErrorException("type LogicalExpression has no field `$(name)`"))
    end
end
function Base.getproperty(expr::LogicalExpression, name::Symbol)
    if name == :arguments
        return FakeVector(expr, name, getfield(expr, name))
    end

    throw(ErrorException("use method `$(name)(my_expr)` instead"))
end
# see utils.jl#FakeVector
function setvectorindex!(expr::LogicalExpression, name::Symbol, x, index::Int)
    if name == :arguments
        remove_from_parents!(expr, arguments(expr)[index:index])
        getfield(expr, :arguments)[index] = x
        add_to_parents!(expr, arguments(expr)[index:index])
        expr.cached_variables_valid = false
    else
        @warn "unexpected `setvectorindex!` call for field `$(name)`"
    end
end
# this method is probably not optimized but it doesn't matter since expressions usually have two or fewer arguments, and this method
# will hardly ever be used regardless
function setvectorindex!(expr::LogicalExpression, name::Symbol, x, index::Union{UnitRange{Int}, StepRange{Int, Int}})
    li = 1  # linear index, to account for StepRange
    for i ∈ index
        setvectorindex!(expr, name, x[li], i)
        li += 1
    end
end

function Base.show(io::IO, expr::LogicalExpression)
    showparens(expr) = (expr isa LogicalExpression) && !isunary(operation(expr))

    if isunary(operation(expr))
        arg = first(arguments(expr))
        show(io, operation(expr))
        if showparens(arg)
            print(io, "(")
        end
        print(io, arg)
        if showparens(arg)
            print(io, ")")
        end
    elseif isbinary(operation(expr))
        args = arguments(expr)

        if showparens(args[1])
            print(io, "(")
        end
        show(io, arguments(expr)[1])
        if showparens(args[1])
            print(io, ")")
        end

        print(io, " ")
        show(io, operation(expr))
        print(io, " ")

        if showparens(args[2])
            print(io, "(")
        end
        show(io, arguments(expr)[2])
        if showparens(args[2])
            print(io, ")")
        end
    end
end

# Standard AbstractExpression interface methods
"""
    istree(expr::T) where {T <: AbstractExpression}

Returns true when `expr` is a `LogicalExpression` and false otherwise.
"""
istree(::LogicalSymbol) = false
istree(::LogicalExpression) = true

"""
    isnode(expr::T) where {T <: AbstractExpression}

Returns true when `expr` is a `LogicalSymbol` and false otherwise.
"""
isnode(::LogicalSymbol) = true
isnode(::LogicalExpression) = false

"""
    variables(expr::T) where {T <: AbstractExpression}

Returns variables present in an entire expression tree. When the first argument is a `LogicalSymbol`, this is singleton
set containing only the provided `LogicalSymbol`. When the argument is a `LogicalExpression`, this is a set containing the union
of all the variables present in each argument.
"""
variables(sym::LogicalSymbol) = Set{LogicalSymbol}(LogicalSymbol[sym])
function variables(expr::LogicalExpression)
    if !getfield(expr, :cached_variables_valid)
        setfield!(expr, :cached_variables, recursivevariables(arguments(expr)))
        expr.cached_variables_valid = true
    end
    getfield(expr, :cached_variables)
end

"""
    operations(expr::T) where {T <: AbstractExpression}

Returns the operations present in an entire expression tree. When given a `LogicalSymbol`, this is an empty set, whereas
when given a `LogicalExpression`, this is a set containing the expression's own `LogicalOperation` and the operations set
of each of its arguments.
"""
operations(::LogicalSymbol) = Set{LogicalOperation}([])
function operations(expr::LogicalExpression)
    if !getfield(expr, :cached_operations_valid)
        setfield!(expr, :cached_operations, recursiveoperations(arguments(expr), operation(expr)))
        expr.cached_operations_valid = true
    end
    getfield(expr, :cached_operations)
end

"""
    arguments(expr::T) where {T <: AbstractExpression}

Returns the arguments which an `AbstractExpression` contains. `LogicalSymbols` contain no arguments and `LogicalExpressions`
can contain any number of arguments.
"""
arguments(::LogicalSymbol) = AbstractExpression[]
arguments(expr::LogicalExpression) = getfield(expr, :arguments)

# ASSOCIATION (perhaps move this elsewhere eventually?)
"""
    associative_ordering(expr::LogicalExpression)

Descends the expression tree with a left-side-first depth first search. Each symbol encountered is added to a list in
the order it appears in this search and is returned by this function.
"""
associative_ordering(expr::LogicalExpression) = reduce(vcat, associative_ordering.(arguments(expr)))
associative_ordering(sym::LogicalSymbol) = [sym]

"""
    isequal_associative(expr1::LogicalExpression, expr2::LogicalExpression)

Checks whether an expression is equal with another ignoring the associative ordering of each expression.
"""
function isequal_associative(expr1::LogicalExpression, expr2::LogicalExpression)
    length(operations(expr1)) == 1 && isassociative(first(operations(expr1))) && operations(expr1) == operations(expr2) && associative_ordering(expr1) == associative_ordering(expr2)
end
isequal_associative(sym1::LogicalSymbol, sym2::LogicalSymbol) = isequal(sym1, sym2)
isequal_associative(::AbstractExpression, ::AbstractExpression) = false

_associative_tree_count_cache = Dict()
"""
    associative_tree_count(nodes::Int)

Function which calculates how many ways to arrange parenthesis in an expression there are with a given node count.
"""
function associative_tree_count(nodes::Int)
    if nodes == 0
        return 1
    end
    if nodes == 1
        return 1
    end

    if haskey(_associative_tree_count_cache, nodes)
        return _associative_tree_count_cache[nodes]
    end

    agg = 0
    for i ∈ 0:(nodes - 1)
        agg += associative_tree_count(nodes - i - 1) * associative_tree_count(i)
    end

    _associative_tree_count_cache[nodes] = agg

    agg
end

# OPERATORS
"""
    ¬(x)

Logical negation operator, typed with `\\neg`. Boolean equivalent is the `!` operator.
"""
const ¬ = LogicalOperation(x -> !x, :¬, 1, false, false)

# binary operators
"""
    x ∧ y

Logical conjunction operator, typed with `\\wedge`. Boolean equivalent is the `&&` operator.
"""
const ∧ = LogicalOperation((x, y) -> x && y, :∧, 2, true, true)
"""
    x ∨ y

Logical disjunction operator, typed with `\\vee`. Boolean equivalent is the `||` operator.
"""
const ∨ = LogicalOperation((x, y) -> x || y, :∨, 2, true, true)
"""
    x → y

Logical implication operator, typed with `\\rightarrow`.
"""
const → = LogicalOperation((x, y) -> (¬x ∨ y), :→, 2, false, false)
"""
    x ⟷ y

Logical equivalence operator, typed with `\\longleftrightarrow`.
"""
const ⟷ = LogicalOperation((x, y) -> (x ∧ y) ∨ (¬x ∧ ¬y), :⟷, 2, true, true)

"""
    random_expression(subchance=0.5, symbol_alphabet::Vector{LogicalSymbol}=LogicalSymbol.([:a, :b, :c]), operator_alphabet::Vector{LogicalOperation}=[∧, ∨, →, ⟷, ¬])

Creates a random expression from a symbol and operator alphabet. The `subchance` parameter describes how likely it is that
an argument will be a subexpression as opposed to an expression.
"""
function random_expression(subchance=0.5, symbol_alphabet::Vector{LogicalSymbol}=LogicalSymbol.([:a, :b, :c]), operator_alphabet::Vector{LogicalOperation}=[∧, ∨, →, ⟷, ¬])
    if rand() > subchance
        return rand(symbol_alphabet)
    else
        rand_op = rand(operator_alphabet)
        return LogicalExpression(AbstractExpression[random_expression(subchance, symbol_alphabet, operator_alphabet) for _ ∈ 1:argument_count(rand_op)], rand_op)
    end
end


x = range(0, stop=1, length=100)
plot(x, x.^2, x.^3, x.^4)


@connector function HeatPort(; name, T_start=273.15 + 20.0, Q_flow_start=0.0)
    @variables T(t)=T_start
    @variables Q_flow(t)=Q_flow_start [connect = Flow]
    ODESystem(Equation[], t, [T, Q_flow], [], name=name)
end
Base.@doc """
    HeatPort(; name, T_start=273.15 + 20.0, Q_flow_start=0.0)

Port for a thermal system.

# Parameters:
- `T_start`: [K] Temperature of the port
- `Q_flow_start`: [W] Heat flow rate at the port

# States:
- `T`: [K] Temperature of the port
- `Q_flow`: [W] Heat flow rate at the port
""" HeatPort

"""
    Element1D(;name, dT0=0.0, Q_flow0=0.0)

This partial model contains the basic connectors and variables to allow heat transfer models to be created that do not
store energy. This model defines and includes equations for the temperature drop across the element, `dT`, and the heat
flow rate through the element from `port_a` to `port_b`, `Q_flow`.

# Parameters:
- `dT_start`:  [K] Initial temperature difference across the component a.T - b.T
- `Q_flow_start`: [W] Initial heat flow rate from port a -> port b

# States:
- `dT`:  [K] Temperature difference across the component a.T - b.T
- `Q_flow`: [W] Heat flow rate from port a -> port b
"""
function Element1D(;name, dT_start=0.0, Q_flow_start=0.0)
    @named port_a = HeatPort()
    @named port_b = HeatPort()
    sts = @variables begin
        dT(t)=dT_start
        Q_flow(t)=Q_flow_start
    end
    eqs = [
        dT ~ port_a.T - port_b.T
        port_a.Q_flow ~ Q_flow
        port_a.Q_flow + port_b.Q_flow ~ 0
    ]

    return compose(ODESystem(eqs, t, sts, []; name=name), port_a, port_b)
end


function shrink_regularizer(A, lambda,stepSz, type,a)
    #type = L1; SCAD; MCP
    if type == :L1
        a_shr = shrinkL1(A, lambda * stepSz)
    elseif type == :SCAD
        a_shr = shrinkL1(A - stepSz * SCAD_derivative(A, a, lambda), lambda * stepSz)
    elseif type == :MCP
        a_shr = shrinkL1(A - stepSz * MCP_derivative(A, a, lambda), lambda * stepSz)
    end
    return a_shr
end

function shrinkL1(A, lambda)
    dg = diag(A)
    soft = (x, t) -> sign(x) * max(0, abs(x) - t)
    #A = soft.(A, lambda) # A .= modifies Psi in place.
    B = soft.(A, lambda)
    B[diagind(A)] = dg
    return B
end

function SCAD_derivative(A, a, rho)
    #Function to evaluate derivative of [SCAD penalty - L1] with parameter a .> 2 &
    #regularization parameter rho. Matricized with no offdiagonal penalty.
    A = A - diagm(diag(A))
    trm1 = (abs.(A) .> rho) .& (abs.(A) .<= a*rho).*(-2*A./(2*(a-1)) + rho.*sign.(A)/(a-1))
    trm2 = (abs.(A) .> a*rho).*(-rho.*sign.(A))
    dy = trm1 + trm2
    return dy
end

function MCP_derivative(A, a, rho)
    #Function to evaluate derivative of [MCP penalty - L1] with parameter a .> 0 &
    #regularization parameter rho. Matricized with no offdiagonal penalty.
    A = A - diagm(diag(A))
    dy = (abs.(A) .< rho*a).*sign.(A).*rho.*(-abs.(A)/(rho*a))
    return dy
end

"""
PROJECT_KSUM Projects the inverse of a Kronecker sum matrix [stored as factor
matrices Xs, where Xs[k] is of size ps[k] x ps[k]] onto a
Kronecker sum of dimensions also given by the vector ps.
"""
function project_ksum(Xs, ps)
    K = length(Xs)
    SigInvs = Array{Any}(undef, K)
    Uinvs = Array{Any}(undef, K)
    Xinvs = Array{Any}(undef, K)

    if K > 1 # If K = 1 no projection required
        diagmat = zeros(ps[end:-1:1]...)
        for i = 1:K
            s, U = eigen(Xs[i])
            Uinvs[i] = U
            less = prod(ps[1:i-1])
            more = prod(ps[i+1:end])

            # more & less should swap in the following line.
            diagmat = diagmat + reshape(
                kron(ones(less,1), kron(s, ones(more,1))),
                ps[end:-1:1]...
            )
        end
        diagmat = permutedims(diagmat,K:-1:1)

        trA = undef
        for i = 1:K
            # matricize. Column is ps[i]-dim. Each row is a fiber.
            mtt = reshape(
                permutedims(1 ./ diagmat, [(1:i-1)... (i+1:K)... i]),
                prod(ps) ÷ ps[i],
                ps[i],
            )
            SigInvs[i] = dropdims(mean(mtt, dims=1); dims=1) # mean(mtt,1) # mtt is not callable.
            if i == 1
               trA = mean(SigInvs[i])
            end
            SigInvs[i] = SigInvs[i] .- trA * (K - 1) / K
        end

        for i = 1:K
            Xinvs[i] = Uinvs[i] * diagm(SigInvs[i]) * Uinvs[i]'
        end
    else
        #[U,s,V] = svd(Xs[1])
        s, U = eigen(Xs[1])
        Uinvs[1] = U
        SigInvs[1] = 1 ./ s
        Xinvs[1] = inv(Xs[1])
    end
    return Uinvs, SigInvs, Xinvs
end

function heatmaps(S)
    plot(
        [heatmap(s, aspect_ratio=:equal, yflip=true) for s in S]...,
        layout=(1, length(S)),
        size=(300 * length(S),300)
    )
end

## Below are some potentially useful functions

"""
Submatrix notation A(i,j|k)

A(i,j|k)_{r,s} = (1/mk) < A, I ⊗ E_{r,s} ⊗ I >
"""
function submatrix(A, ps, k, i, j)
    K = length(ps)
    pre, post = prod(ps[1:k-1]), prod(ps[k+1:K])

     # Ksum def in the paper requires reversed order of Julia dim
    i_post, i_pre = Base._ind2sub((post, pre), i)
    j_post, j_pre = Base._ind2sub((post, pre), j)
    # @vectorize_2nd_arg? Ref?
    # (j_pre - 1) * ps[k] * post + j_post .+ post * (0 : ps[k]-1)
    i_grid = [Base._sub2ind((post, ps[k], pre), i_post, r, i_pre) for r = 1:ps[k]]
    j_grid = [Base._sub2ind((post, ps[k], pre), j_post, s, j_pre) for s = 1:ps[k]]
    # println((i, j), ' ',
    #         collect((i_pre - 1) * ps[k] .+ i_post * (1 : ps[k])), ' ',
    #         collect((j_pre - 1) * ps[k] .+ j_post * (1 : ps[k])))
    return A[i_grid, j_grid]
end

function test_submatrix_repr()
    ps = [2, 2, 2]
    p = prod(ps)

    #A = reshape(1:p^2, p, p)
    A = rand(p, p)

    K = length(ps)
    for k = 1:K #1:K
        mk = p ÷ ps[k]
        pre, post = prod(ps[1:k-1]), prod(ps[k+1:end])
        B = zeros(size(A))
        for i = 1:mk
            i_post, i_pre = Base._ind2sub((post, pre), i)
            for j = 1:mk
                j_post, j_pre = Base._ind2sub((post, pre), j)
                sub = submatrix(A, ps, k, i, j)
                E_pre = zeros(pre, pre)
                E_pre[i_pre, j_pre] = 1
                E_post = zeros(post, post)
                E_post[i_post, j_post] = 1
                B += reduce(kron, [E_pre, sub, E_post])
                T = reduce(kron, [E_pre, sub, E_post])
                display(heatmap([A T],
                        yflip=true,
                        aspect_ratio=:equal,
                        title=string([i, j])))
            end
        end
        print(all(A == B))
    end
end

"""
The projection operation in Lemma 33 in the paper
"""
function project_ksum_v1(A, ps)
    p = prod(ps)
    K = length(ps)
    Ak = Array{Any}(undef, K)
    for k = 1:K
        mk = p ÷ ps[k]
        sub = [submatrix(A, ps, k, i, i) for i = 1:mk] # Diag mode-k submatrices
        Ak[k] = sum(sub) / mk  # Statistics.mean
    end
    tau = tr(A) / p
    A_proj = reduce(⊕, Ak) - (K - 1) * tau * I
    return A_proj
end

# function heatmaps(Ss)
#     plots = [heatmap(S) for S in Ss]
#
#     layout = @layout [a b]
#     p = plot(plots**, layout=layout, size=(1100,300))
#     return p
# end


# Xiao, Zhang, "A proximal stochastic gradient method  with progressive
# variance reduction.", SIAM Journal on Optimization 24.4 (2014): 2057-2075.
#
# Reddi, Sra, Poczos, and Smola, "Proximal stochastic methods for nonsmooth
# nonconvex finite-sum optimization." In Advances in Neural Information
# Processing Systems (2016), pp. 1145–1153.
#
# Allen-Zhu, Yuan, "Improved SVRG for non-strongly-convex or
# sum-of-non-convex objectives." In Proceedings of the 33rd
# International Conference on Machine Learning (2016): 1080–1089.
#

using LinearAlgebra
using ProximalOperators
using ProximalAlgorithms.IterationTools
using Printf
using Base.Iterators
using Random
using Flux

export solution

include("SVRG_basic.jl")
include("SVRG_basic_DNN.jl")

struct SVRG{R<:Real}
    γ::Maybe{R}
    maxit::Int
    verbose::Bool
    freq::Int
    m::Maybe{Int}
    plus::Bool
    DNN::Bool
    function SVRG{R}(;
        γ::Maybe{R} = nothing,
        maxit::Int = 10000,
        verbose::Bool = false,
        freq::Int = 1000,
        m::Maybe{Int} = nothing,
        plus::Bool = false,
        DNN::Bool = false,
    ) where {R}
        @assert γ === nothing || γ > 0
        @assert maxit > 0
        @assert freq > 0
        new(γ, maxit, verbose, freq, m, plus, DNN)
    end
end

function (solver::SVRG{R})(
    x0::AbstractArray{C};
    F = nothing,
    g = ProximalOperators.Zero(),
    L = nothing,
    μ = nothing,
    N = N,
) where {R,C<:RealOrComplex{R}}

    stop(state::SVRG_basic_state) = false
    disp(it, state) = @printf "%5d | %.3e  \n" it state.γ

    F === nothing && (F = fill(ProximalOperators.Zero(), (N,)))
    m = solver.m === nothing ? m = N : m = solver.m

    maxit = solver.maxit
    if solver.plus && solver.maxit > 25
        maxit = 25
        @warn "exponential number of inner updates...reverted to 25 maximum iterations"
    end

    # dispatching the structure
    iter = SVRG_basic_iterable(F, g, x0, N, L, μ, solver.γ, m, solver.plus)
    iter = take(halt(iter, stop), maxit)
    iter = enumerate(iter)
    num_iters, state_final = nothing, nothing
    for (it_, state_) in iter  # unrolling the iterator
        # see https://docs.julialang.org/en/v1/manual/interfaces/index.html
        if solver.verbose && mod(it_, solver.freq) == 0
            disp(it_, state_)
        end
        num_iters, state_final = it_, state_
    end
    if solver.verbose && mod(num_iters, solver.freq) !== 0
        disp(num_iters, state_final)
    end # for the final iteration
    return solution(state_final), num_iters
end

"""
    SVRG([γ, maxit, verbose, freq, m, plus])

Instantiate the SVRG algorithm  for solving (strongly) convex optimization problems of the form

    minimize 1/N sum_{i=1}^N f_i(x) + g(x)

If `solver = SVRG(args...)`, then the above problem is solved with

	solver(x0, [F, g, N, L, μ])

where F is an array containing f_i's, x0 is the initial point, and L, μ are arrays of
smoothness and strong convexity moduli of f_i's; they are optional when γ is provided.

Optional keyword arguments are:
* `γ`: stepsize
* `L`: an array of smoothness moduli of f_i's
* `μ`: (if strongly convex) an array of strong convexity moduli of f_i's
* `maxit::Integer` (default: `10000`), maximum number of iterations to perform.
* `verbose::Bool` (default: `true`), whether or not to print information during the iterations.
* `freq::Integer` (default: `100`), frequency of verbosity.
* `plus::Bool` !

"""

SVRG(::Type{R}; kwargs...) where {R} = SVRG{R}(; kwargs...)
SVRG(; kwargs...) = SVRG(Float64; kwargs...)

"""
If `solver = SVRG(args...)`, then

    itr = iterator(solver, x0, [F, g, N, L, μ])

is an iterable object. Note that [maxit, verbose, freq] fields of the solver are ignored here.

The solution at any given state can be obtained using solution(state), e.g.,
for state in Iterators.take(itr, maxit)
    # do something using solution(state)
end

See https://docs.julialang.org/en/v1/manual/interfaces/index.html
and https://docs.julialang.org/en/v1/base/iterators/ for a list of iteration utilities
"""

function iterator(
    solver::SVRG{R},
    x0::AbstractArray{C};
    F = nothing,
    g = ProximalOperators.Zero(),
    L = nothing,
    μ = nothing,
    N = N,
) where {R,C<:RealOrComplex{R}}
    F === nothing && (F = fill(ProximalOperators.Zero(), (N,)))
    m = solver.m === nothing ? m = N : m = solver.m
    # dispatching the iterator
    if solver.DNN
        iter = SVRG_basic_DNN_iterable(F, g, x0, N, L, μ, solver.γ, m, solver.plus)
    else
        iter = SVRG_basic_iterable(F, g, x0, N, L, μ, solver.γ, m, solver.plus)
    end

    return iter
end


using BitSAD
using DataStructures

circuit = @circuit SVFilter begin
    parameters : [
        f::Float64 => 0.125,
        q::Float64 => 1.875
    ]

    submodules : [
        delay₁::CircularBuffer{DBit},
        delay₂::CircularBuffer{DBit},
        ΣΔ₁::SDM,
        ΣΔ₂::SDM
    ]

    initialize : begin
        svfilter = new(f, q, delay₁, delay₂, ΣΔ₁, ΣΔ₂)
        fill!(svfilter.delay₁, zero(DBit))
        fill!(svfilter.delay₂, zero(DBit))

        return svfilter
    end

    circuit : (filter::SVFilter)(x::DBit) -> begin
        # get delay buffer values
        d₁old = popfirst!(filter.delay₁)
        d₂old = popfirst!(filter.delay₂)

        # calculate new buffer values and filter output
        d₂ = filter.ΣΔ₂(filter.f * d₁old + d₂old)
        d₁ = filter.ΣΔ₁(filter.f * (x - d₂ - filter.q * d₁old) + d₁old)

        # push new values into delay buffers
        push!(filter.delay₁, d₁)
        push!(filter.delay₂, d₂)

        return d₁
    end
end

SVFilter(f::Real, q::Real, delay::Integer) =
    SVFilter(f, q, CircularBuffer{DBit}(delay), CircularBuffer{DBit}(delay), SDM(), SDM())

function readdat(filename)
    s = DBitstream()
    samplerate = 0
    nchannels = 0

    for line in eachline(filename)
        m = match(r"; Sample Rate (\d+)", line)
        if !isnothing(m)
            samplerate = parse(Int, m.captures[1])
        end

        m = match(r"; Channels (\d+)", line)
        if !isnothing(m)
            nchannels = parse(Int, m.captures[1])
        end

        m = match(r"\s+[\de\-\.]+\s+([\d\.\-]+)\s+", line)
        if !isnothing(m)
            bit = round(parse(Float64, m.captures[1]))
            push!(s, DBit(bit))
        end
    end

    return s, samplerate, nchannels
end

function writedat(filename, s::DBitstream; samplerate, nchannels)
    open(filename, "w") do io
        write(io, "; Sample Rate $samplerate\n")
        write(io, "; Channels $nchannels\n")
        for i in 1:length(s)
            bit = pop!(s)
            write(io, " $((i - 1) * (1 / samplerate))  $(float(bit)) \n")
        end
    end
end

svfilter = SVFilter(0.125, 1.875, 1)
x, samplerate, nchannels = readdat("./examples/Expchirp.dat")
y = DBitstream()

for i in 1:length(x)
    push!(y, svfilter(pop!(x)))
end

writedat("./examples/Expchirp-output.dat", y; samplerate = samplerate, nchannels = nchannels)


import PyPlot.plot

"""
    plot_field(field::Array{<:Real,2}, a1::Array{<:Real,1}, a2::Array{<:Real,1}; cmap="coolwarm", vmin=nothing, vmax=nothing)

Plot the real-valued `field` on a unit cell with lattice vectors `a1` and `a2`.
"""
function plot_field(field::Array{<:Real,2}, a1::Array{<:Real,1}, a2::Array{<:Real,1}; cmap="coolwarm", vmin=nothing, vmax=nothing, label=nothing)
    ps = range(-0.5, stop=0.5, length=size(field,1)+1)
    qs = range(-0.5, stop=0.5, length=size(field,2)+1)
    xs = [p*a1[1]+q*a2[1] for p in ps, q in qs]
    ys = [p*a1[2]+q*a2[2] for p in ps, q in qs]
    if cmap == "coolwarm" && vmin == nothing && vmax == nothing
        vlim = maximum(abs.(field))
        vmin = -vlim
        vmax = +vlim
    end
    if cmap == "viridis" && vmin == nothing && vmax == nothing
        vmin = 0
        vmax = maximum(abs.(field))
    end
    pcolormesh(xs, ys, field, cmap=cmap, vmin=vmin, vmax=vmax)
    gca().set_aspect("equal")
    axis("off")
    colorbar()
    if label != nothing
        title(label)
    end
end

"""
    plot_field(field::Array{<:Real,2}, a1::Array{<:Real,1}, a2::Array{<:Real,1}; cmap="coolwarm", vmin=nothing, vmax=nothing)

Plot the complex-valued `field` on a unit cell with lattice vectors `a1` and `a2`.
"""
function plot_field(field::Array{<:Complex,2}, a1::Array{<:Real,1}, a2::Array{<:Real,1}; cmap="coolwarm", vmin=nothing, vmax=nothing, label=nothing)
    # Plot real and imaginary parts with the same color scale
    if cmap == "coolwarm" && vmin == nothing && vmax == nothing
        vlim = maximum(abs.(field))
        vmin = -vlim
        vmax = +vlim
    end
    if cmap == "viridis" && vmin == nothing && vmax == nothing
        vmin = 0
        vmax = maximum(abs.(field))
    end
    figure(figsize=(4,2))
    subplot(1,2,1)
    plot_field(real(field), a1, a2, cmap=cmap, vmin=vmin, vmax=vmax, label="Re[$label]")
    subplot(1,2,2)
    plot_field(imag(field), a1, a2, cmap=cmap, vmin=vmin, vmax=vmax, label="Im[$label]")
end

"""
    get_field(data::AbstractVector{<:Complex}, basis::PlaneWaveBasis;
                        k0=[0,0], t1s=-0.5:0.01:0.5, t2s=-0.5:0.01:0.5)

Convert the `data` from a `PlaneWaveBasis` to a real space grid.
"""
function get_field(data::AbstractVector{<:Complex}, basis::PlaneWaveBasis;
                    k0=[0,0], t1s=-0.5:0.01:0.5, t2s=-0.5:0.01:0.5)
    # Generate real space coordinates
    b1, b2 = basis.b1, basis.b2
    a1, a2 = bs_to_as(b1, b2)
    xs = [t1*a1[1]+t2*a2[1] for t1 in t1s, t2 in t2s]
    ys = [t1*a1[2]+t2*a2[2] for t1 in t1s, t2 in t2s]
    out = zeros(ComplexF64, length(t1s), length(t2s))
    # Inverse Fourier transform
    kxs = k0[1] .+ basis.kxs
    kys = k0[2] .+ basis.kys
    for (kx,ky,u) in zip(kxs, kys, data)
        for c in CartesianIndices(out)
            phase = kx*xs[c] + ky*ys[c]
            out[c] += u * exp(1im*phase)
        end
    end
    return out
end

function get_field_FDFD(data::AbstractVector{<:Complex}, basis::FDFDBasis; k0 = [0, 0])
    dx, dy = basis.resolution[1], basis.resolution[2]
    a1, a2 = basis.a1, basis.a2
    Nx, Ny = Int(1/dx), Int(1/dy)
    out = zeros(ComplexF64, Nx, Ny)
    for m = 1:Nx
        for n = 1:Ny
            Lx = (m - 0.5) * dx
            Ly = (n - 0.5) * dy
            index = Int((n - 1)*Nx + m)
            out[m, n] = data[index] * exp(-1im * k0[1] * Lx) * exp(-1im * k0[2] *Ly) #u function
        end
    end
    return out
end

function plot_field(data::AbstractVector{<:Complex}, basis::FDFDBasis; k0 = [0, 0], cmap="coolwarm", vmin=nothing, vmax=nothing, label=nothing)
    field = get_field_FDFD(data, basis, k0 = k0)
    plot_field(field, basis.a1, basis.a2, cmap=cmap, vmin=vmin, vmax=vmax, label=label)
end

"""
    plot_field(data::AbstractVector{<:Complex}, basis::PlaneWaveBasis;
                        k0=[0,0], cmap="coolwarm", vmin=nothing, vmax=nothing, label=nothing)

Converts the plane-wave amplitudes to real space using [`get_field`](@ref)
and then plots the field.
"""
function plot_field(data::AbstractVector{<:Complex}, basis::PlaneWaveBasis; k0=[0,0], cmap="coolwarm", vmin=nothing, vmax=nothing, label=nothing)
    field = get_field(data, basis, k0=k0)
    a1, a2 = bs_to_as(basis.b1, basis.b2)
    plot_field(field, a1, a2, cmap=cmap, vmin=vmin, vmax=vmax, label=label)
end

"""
    plot(geometry::Geometry)

Plot the permittivity and permeability of the `geometry` in real space.
"""
function PyPlot.plot(geometry::Geometry)
    # Plot permittivity
    plot_field(geometry.ep.epsxx, geometry.a1, geometry.a2, cmap="viridis", label="ϵ")
    # Plot permeability
    plot_field(geometry.mu.muxx, geometry.a1, geometry.a2, cmap="viridis", label="μ")
end

"""
    plot(solver::Solver)

Plot the representation of the `solver`'s geometry in real space.

The `solver` approximates the geometry using a truncated basis of plane waves,
so `plot(solver)` lets us judge how accurately the geometry is represented.
"""
function PyPlot.plot(solver::Solver)
    col = findfirst((solver.basis.ps .==0) .& (solver.basis.qs .== 0))
    # Plot permittivity
    plot_field(solver.epc[:,col], solver.basis, cmap="viridis", label="ϵ")
    # Plot permeability
    plot_field(solver.muc[:,col], solver.basis, cmap="viridis", label="μ")
end

"""
    plot(mode::Mode, [bloch_phase=true])

Plot the mode in real space.

To plot only the cell-periodic part of the Bloch wave, set `bloch_phase=false`.
"""
function PyPlot.plot(mode::Mode; bloch_phase=true)
    if bloch_phase
        label = mode.label
        plot_field(mode.data, mode.basis, label=label, k0=mode.k0)
    else
        label = L"e^{-i k \cdot r} \cdot" * mode.label
        plot_field(mode.data, mode.basis, label=label)
    end
end

function PyPlot.plot(mode::Mode_FDFD; bloch_phase=true)
    if bloch_phase
        label = mode.label
        plot_field(mode.data, mode.basis, label=label)
    else
        label = L"e^{-i k \cdot r} \cdot" * mode.label
        plot_field(mode.data, mode.basis, label=label , k0=mode.k0)
    end
end


"""
    AbstractMaxEntropyModel

An abstract type for a MaxEntropyModel. Each model has one or more structural constraints
that are fixed while the rest of the network is completely random.
"""
abstract type AbstractMaxEntropyModel end

"""
    UBCM{T,N} <: AbstractMaxEntropyModel where {T<:Real, N<:UInt}

# Undirected Binary Configuration Model (UBCM)
Maximum entropy model with a fixed degree sequence. Uses the model where ``x_i = e^{-\\theta_i}``

See also: [AbstractMaxEntropyModel](@ref)
"""
mutable struct UBCM{T,N} <: AbstractMaxEntropyModel where {T<:Real, N<:UInt}
    idx::Vector{N}
    κ::Vector{T}
    f::Vector{T}
    method::Symbol
    x0::Vector{T} # initial guess
    F::Vector{T}  # buffer for solution => used?
    f!::Function  # function that will need to be solved
    xs::Vector{T} # solution vector
    x::IndirectArray{T,1, N, Vector{N}, Vector{T}}  # solution vector expanded
end

Base.show(io::IO, model::UBCM{T,N}) where {T,N} = print(io, """UBCM{$(T)} model ($(length(model.κ) < length(model.idx)  ? "$(round((1- length(model.κ)/length(model.idx)) * 100,digits=2))% compression in $(N)" : "uncompressed"))""")

"""
	UBCM(k::Vector) where {T<:Int}

Generate a UBCM model from degree vector of a graph.
### Input
- `k` degree vector of the graph
- keyword arguments:
    - `compact` (optional, default: `true`) flag to indicate compression usage (as the same constraint leads to the same value)
	- `P` (optional, default: `Float64`) datatype used for storage
	- `method` (optional, default: `:newton`) method to uses to solve the system of equations; should be one of `:newton` or `fixedpoint`
	- `initial` (optional, default: `:degree`) vector of initial values or symbol indication what method tu use for initial guess

### Output
`UBCM` object

### Examples
```Julia
UBCM([1;2;2;3])
````

### See also
[`UBCM`](@ref)
"""
function UBCM(k::Vector{T}; compact::Bool=true,
							P::DataType=Float64,
							method::Symbol=:newton,
							initial::Union{Symbol, Vector}=:degree, kwargs...) where {T<:Int}
	# check input
	method ∈ [:newton, :fixedpoint] ? nothing : throw(ArgumentError("'$(method)' is not a valid solution method"))

	# convert datatypes
	k = convert.(P, k)

	# compress or not
	if compact
		idxtype = length(unique(k)) < typemax(UInt8) ? UInt8 : length(unique(k)) < typemax(UInt16) ? UInt16 : UInt32 # limited to 4.29e9 nodes
		idx, κ = IndirectArray{idxtype}(k).index, IndirectArray{idxtype}(k).values
		f = countmap(k, ones(P, size(k))) # using weight vector for proper type inference
		f = [f[v] for v in κ] # addition
	else
		κ = k
		idxtype = length(κ) < typemax(UInt8) ? UInt8 : length(κ) < typemax(UInt16) ? UInt16 : UInt32 # limited to 4.29e9 nodes
		idx = collect(one(idxtype):idxtype(length(κ)))
		f = ones(P,length(κ)) #     Dict(v => one(P) for v in κ)
	end

	# initial vector - need more work (current conditions are unstable)
	if isa(initial, Vector)
		@warn "THIS IS NOT WORKING!\n$(initial)\n$(initial[idx])\n$(κ)"
		length(initial[idx]) == length(κ) ? nothing : throw(DimensionMismatch("Length of initial vector $(length(initial)) does not match the length of κ $(length(κ))"))
		x0 = P.(initial[idx])
	elseif isa(initial, Symbol)
		initial ∈ [:random, :uniform, :degree, :nodes, :edges] ? nothing : throw(ArgumentError("'$(initial)' is not a valid initial argument method in combination with `$(method)` as a solution method"))
		# should be evaluated more
		if isequal(initial, :random)      		# random in  [0,1]
			x0 = rand(P, length(κ))
		elseif isequal(initial, :uniform) 		# uniform unique random value in [0,1]
			x0 = rand(P) .* ones(P, length(κ))
		elseif isequal(initial, :degree) 		# proportional to degree in [0,1]
			x0 = κ ./ maximum(κ)
		elseif isequal(initial, :nodes) 		# proportional to sqrt of number or nodes
			x0 = κ ./ P.(sqrt(sum(f)))
		elseif isequal(initial, :edges)			# proportional to sqrt of number of edges
			x0 = κ ./ P.(sqrt(sum(f .* κ)/2))
		end
	else
		throw(ArgumentError("`initial` should be either a Vector or a Symbol"))
	end

	# functions to compute
	if isequal(method, :newton)
		f! = (F::Vector, x::Vector) -> UBCM_∂ℒ_∂x!(F, x, κ, f)
	elseif isequal(method, :fixedpoint)
		f! = (F::Vector, x::Vector) -> UBCM_∂ℒ_∂x_it!(F, x, κ, f)
	end

	# buffers (required?)
	F = similar(x0) # modify in place method
	xs = similar(x0) # solution value

	# outresult with indirect indexing for further use
	x = IndirectArray(idx, xs)

	return UBCM(idx, κ, f, method, x0, F, f!, xs, x)
end

"""
	UBCM(G::T; kwargs...) where T<:AbstractGraph

Generate UBCM model directly from a graph.

### See also:
[`UBCM`](@ref)
"""
function UBCM(G::T; kwargs...) where T<:AbstractGraph
	if is_directed(G)
		@warn "Using a directed graph ($(G)) with an undirected network model"
	end

	return  UBCM(degree(G); kwargs...)
end

"""
	fffff(x)

helper function for UBCM gradient computation
"""
fffff(x) = x / (1 + x)

"""
	UBCM_∂ℒ_∂x!(F::Vector, x::Vector, κ::Vector, f::Vector)

Gradient of the likelihood function of the UBCM model using the ``x_{i}`` formulation.

F value of gradients
x value of parameters
κ value of (reduced) degree vector
f frequency associated with each value in κ

### See also:
[`UBCM`](@ref)

"""
function UBCM_∂ℒ_∂x!(F::Vector, x::Vector, k::Vector,f::Vector)
	@tturbo for i in eachindex(x)
		F[i] = -fffff(x[i]*x[i]) - k[i]
		for j in eachindex(x)
			F[i] += f[j] * fffff(x[i]*x[j])
		end
	end

	return F
end

"""
	UBCM_∂ℒ_∂x_it!(F::Vector, x::Vector, κ::Vector, f::Vector)

Iterative gradient of the likelihood function of the UBCM model using the `x_i` formulation.

F value of function
x value of parameters
κ value of (reduce) degree vector
f frequency associated with each value in κ

### See also:
[`UBCM`](@ref)

"""
function UBCM_∂ℒ_∂x_it!(F::Vector{T}, x::Vector{T}, κ::Vector{T}, f::Vector{T}) where {T}
    @tturbo for i in eachindex(x)
		fx = -x[i] / (1 + x[i] * x[i])
		for j in eachindex(x)
			fx += f[j] * x[j] / (1 + x[j] * x[i])
		end

		F[i] = κ[i] / fx
	end

	return F
end

"""
	solve!model::UBCM; , kwargs...)

Solve the equations of a UBCM model to obtain the maximum likelihood parameters.

### See also [`UBCM`](@ref)
"""
function solve!(model::UBCM; kwargs...)
	if isequal(model.method, :newton)
		res = nlsolve(model.f!, model.x0; autodiff=:forward, kwargs...)
	elseif isequal(model.method, :fixedpoint)
		res = fixedpoint(model.f!, model.x0; kwargs...)
	end

	if res.f_converged || res.x_converged
		model.x = res.zero[model.idx]
	else
		@warn "Method `$(model.method)` did not converge. Verify initial conditions or try another method."
	end

	return res
end

"""
    DBCM{T,N} <: AbstractMaxEntropyModel where {T<:Real, N<:UInt}

# Directed Binary Configuration Model (DBCM)
Maximum entropy model with a fixed in- and outdegree sequence.
"""
mutable struct DBCM{T,N} <: AbstractMaxEntropyModel where {T<:Real, N<:UInt}

end

"""
	UBCM_∂ℒ_∂x!(F::Vector, x::Vector, κ::Vector, f::Vector)

Gradient of the likelihood function of the UBCM model using the ``x{i}`` formulation. Used to generate a function
`(F, x) -> UBCM_∂ℒ_∂x!(F, x, κ, f)` by a `UBCM` instance.

F value of gradients
x value of parameters
κ value of (reduced degree vector)
f frequency associated with each value in κ

## see also:
[`UBCM`](@ref)

"""
function UBCM_∂ℒ_∂θ!(F::Vector, x::Vector, k::Vector,f::Vector)
	@tturbo for i in eachindex(x)
		F[i] = -fffff(x[i]*x[i]) - k[i]
		for j in eachindex(x)
			F[i] += f[j] * fffff(x[i]*x[j])
		end
	end
	return F
end


## runtests.jl : run all test files

# using PyPlot   # *****
using GaussianRandomFields
using Test
using Suppressor
using DelimitedFiles
using Pkg

# test indexsets
include("test_covariance_functions.jl")
include("test_cholesky.jl")
include("test_spectral.jl")
include("test_karhunen_loeve.jl")
include("test_circulant_embedding.jl")
include("test_gaussian_random_fields.jl")
include("test_separable.jl")
include("test_fem.jl")
#include("test_plot.jl")    # *****
include("test_show.jl")