#=
Aodwt.jl
2019-02-23 Jeff Fessler, University of Michigan
=#

export Aodwt

using LinearMapsAA: LinearMapAA, LinearMapAM, LinearMapAO
using Wavelets: dwt!, idwt!, wavelet, WT

"""
    A, levels, mfun = Aodwt(dims ; level::Int=3, wt=wavelet(WT.haar))

Create orthogonal discrete wavelet transform (ODWT) `LinearMapAA`

in
- `dims::Dims` tuple of dimensions

option
- `level::Int` # of levels; default 3
- `wt` wavelet transform type (see `Wavelets` package); default Haar
- `operator::Bool=true` default to `LinearMapAO`

out
- `A` a `LinearMapAX` object
- `scales` array of size `dims` showing the scale of each coefficient
which is useful when imposing scale-dependent regularization
- `mfun` convenience function for A*X when X is a Matrix or Array (not vector)

2019-02-23 Jeff Fessler, University of Michigan
"""
function Aodwt(
	dims::Dims ;
	T::DataType = Float32,
	level::Int = 3,
	wt = wavelet(WT.haar),
	operator::Bool = true, # !
)

	function mfunA(lev)
		# todo: avoiding convert involves `dwt!` limitation; see:
		# https://github.com/JuliaDSP/Wavelets.jl/issues/53
		# https://github.com/JuliaDSP/Wavelets.jl/pull/54
	#	forw!(y,x) = dwt!(y, convert(AbstractArray{T},x), wt, level)
	#	back!(x,y) = idwt!(x, convert(AbstractArray{T},y), wt, level)
		forw!(y,x) = dwt!(y, x, wt, lev)
		back!(x,y) = idwt!(x, y, wt, lev)

		if operator
			mfun = (A, x) -> A * x
			return mfun, LinearMapAA(forw!, back!,
				(prod(dims), prod(dims)) ;
				prop = (wt=wt, level=lev), T=T,
				operator = true, idim=dims, odim=dims,
			)
		else
			mfun = (A, x) -> reshape(A * vec(x), dims)
			return mfun, LinearMapAA(
				(y,x) -> vec(forw!(reshape(y, dims), reshape(x, dims))),
				(x,y) -> vec(back!(reshape(x, dims), reshape(y, dims))),
			#	x -> vec(dwt(reshape(x, dims), wt, level)),
			#	y -> vec(idwt(reshape(y, dims), wt, level)),
				(prod(dims), prod(dims)) ;
				prop = (wt=wt, level=lev), T=T,
			)
		end
	end

	mfun, A = mfunA(level)

	scales = zeros(dims)
	for il=1:level
		_,Al = mfunA(il)
		tmp = mfun(Al, ones(dims)) .== 0
		scales += il * (tmp .& (scales .== 0))
	end

	return A, scales, mfun
end


# # clip-02-05.jl

# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking, Optim
#gr(size=(600,600));

# ### snippet 3.2

p_grid = range(0, step=0.001, stop=1)
prior = ones(length(p_grid))
likelihood = [pdf(Binomial(9, p), 6) for p in p_grid]
posterior = likelihood .* prior
posterior = posterior / sum(posterior)
samples = sample(p_grid, Weights(posterior), length(p_grid));
samples[1:5]

# ### snippet 3.3
# Draw 10000 samples from this posterior distribution

N = 10000
samples = sample(p_grid, Weights(posterior), N);

# In StatisticalRethinkingJulia samples will always be stored
# in an MCMCChains.Chains object.

chn = MCMCChains.Chains(reshape(samples, N, 1, 1), ["toss"]);

# Describe the chain

MCMCChains.describe(chn)

# Plot the chain

plot(chn)

# ### snippet 3.4

# Create a vector to hold the plots so we can later combine them

p = Vector{Plots.Plot{Plots.GRBackend}}(undef, 2)
p[1] = scatter(1:N, samples, markersize = 2, ylim=(0.0, 1.3), lab="Draws")

# ### snippet 3.5

# Analytical calculation

w = 6
n = 9
x = 0:0.01:1
p[2] = density(samples, ylim=(0.0, 5.0), lab="Sample density")
p[2] = plot!( x, pdf.(Beta( w+1 , n-w+1 ) , x ), lab="Conjugate solution")

# Add quadratic approximation

plot(p..., layout=(1, 2))

# End of `03/clip-02-05.jl`


"""
2D Predator-Prey Model

Taken from:
[14.3 pp2: A 2D Predator-Prey Model][auto07p/14.3]
of AUTO-07p manual.

[auto07p/14.3]: http://www.dam.brown.edu/people/sandsted/auto/auto07p.pdf#page=138
"""
module PredatorPrey

using DiffEqBase: ODEProblem
using StaticArrays: SVector
using Setfield: @lens
import Setfield

using ...Bifurcations: BifurcationProblem
using ...Codim2: cast_container, DiffEqCodim2Problem

f1(u, p) =
    p[2] * u[1] * (1 - u[1]) - u[1] * u[2] - p[1] * (1 - exp(- p[3] * u[1]))
f2(u, p) =
    - u[2] + p[4] * u[1] * u[2]

f(u::SVector, p, t) = SVector(f1(u, p), f2(u, p))

function f(du, u, p, t)
    du[1] = f1(u, p)
    du[2] = f2(u, p)
    nothing
end

const _f = f

function make_prob(
        p = (0.0, 3, 5, 3);
        f = _f,
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(f, u0, tspan, p),
        param_axis = (@lens _[1]),
        t_domain = (0.0, 1.0),
        kwargs...)
    x_min = SVector(-0.2, -Inf)
    x_max = SVector(1.0, Inf)
    x_min = cast_container(typeof(ode.u0), x_min)
    x_max = cast_container(typeof(ode.u0), x_max)
    return BifurcationProblem(
        ode, param_axis, t_domain;
        phase_space = (x_min, x_max),
        kwargs...)
end

prob = make_prob()
ode = prob.p.de_prob
p = ode.p
u0 = ode.u0
param_axis = prob.p.param_axis

function make_codim2_prob(
        p = ode.p;
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(f, u0, tspan, p),
        param_axis1 = (@lens _[1]),
        param_axis2 = (@lens _[3]),
        t_domain = ([-Inf, -Inf], [Inf, Inf]),
        kwargs...)
    return DiffEqCodim2Problem(
        ode,
        param_axis1,
        param_axis2,
        t_domain;
        kwargs...)
end

end  # module


using Plots
using Random
using Pkg; Pkg.add("LaTeXStrings")
using LaTeXStrings

# Set working directory
cd("/Users/elvis/Documents/DSI/courses/content/stochastic-approximations/julia-code/")

# Set seed
Random.seed!(123)

# Generate MC points
n = 4000
points_x = 3 .* rand(n)
points_y = -3 .+ 7 .* rand(n)

# Define functions
f1(x) = 3*abs(cos(x)) .+ 2*sin(x)
f2(x) = -3*abs(cos(x)) .+ 2*sin(x)

# Label points in circle
labels = [f2(x) <= y <= f1(x) for (x,y) in zip(points_x, points_y)]

# Create curves
x_vals = 0:.01:3
f1_x = 3*abs.(cos.(x_vals)) .+ 2*sin.(x_vals)
f2_x = -3*abs.(cos.(x_vals)) .+ 2*sin.(x_vals)

# Create plot of MC points
plot(points_x, points_y, seriestype = :scatter,
     group = .!labels,
     legend = false,
     ms = 3,
     bg = RGB(247/255, 236/255, 226/255))

# Add curves to plot
plot!(x_vals, f1_x, seriestype= :line, label = L"f_1(x)", lc = :blue,
    legend = false)
plot!(x_vals, f2_x, seriestype= :line, label = L"f_2(x)", lc = :violet)
plot!(x_vals, f1_x, fillrange = [f2_x, f1_x], fillcolor = :blue, alpha = .12)

# Area estimate
area = 21*sum(labels)/n

savefig("../images/area-ex1-monte-carlo.svg")


using Turing, ReverseDiff, Memoization
using DifferentialEquations
using Plots
using Statistics
using StatsPlots
Turing.setadbackend(:reversediff)
Turing.setrdcache(true)
using Distributions
using LinearAlgebra

# here we are simulating the signal of two coupled oscillators
# with amplitudes A1 and A2 of kBT/k and kBT/(k+c) where c is the coupling constant
# D = kBT/gamma
# in terms of the standard parameters of a OU process:
# sigma = sqrt(2*D)
# Theta = D/A

function corr_osc(c)
    # function that returns two time series that are correlated through a coupling coefficient c
    μ = 0.0 # mean is zero
    σ = sqrt(2) # D=1
    Θ1 = 1.0
    Θ2 = 1.0+abs(c)

    W1 = OrnsteinUhlenbeckProcess(Θ1,μ,σ,0.0,1.0)
    W2 = OrnsteinUhlenbeckProcess(Θ2,μ,σ,0.0,1.0)
    prob1 = NoiseProblem(W1,(0.0,100.0))
    prob2 = NoiseProblem(W2,(0.0,100.0))
    sol1 = solve(prob1;dt=0.1)
    sol2 = solve(prob2;dt=0.1)

    # creating the two correlated
    x1 = (sol1.u .+ sol2.u)/2
    if c>0
        x2 = (sol1.u .- sol2.u)/2
    else
        x2 = (sol2.u .- sol1.u)/2
    end
    return x1,x2
end

# Ornstein-Uhlenbeck process
@model ou(rn,T,delta_t) = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ beta(5.0,1.0)

    rn[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        rn[i] ~ Normal(rn[i-1]*b,sqrt(ampl*(1-b^2)))
    end
end

# Ornstein-Uhlenbeck process of two coupled oscillators
@model ou_corr(rn1,rn2,T,delta_t) = begin
    ampl1 ~ Uniform(0.0,5.0)
    ampl2 ~ Uniform(0.0,5.0)
    d ~ Uniform(0.0,5.0)
    b1 = exp(-delta_t*d/ampl1)
    b2 = exp(-delta_t*d/ampl2)

    rn1[1] ~ Normal(0,sqrt(ampl1))
    rn2[1] ~ Normal(0,sqrt(ampl2))

    for i=2:T
        rn1[i] ~ Normal(rn1[i-1]*b1,sqrt(ampl1*(1-b1^2)))
        rn2[i] ~ Normal(rn2[i-1]*b2,sqrt(ampl2*(1-b2^2)))
    end
end

# Ornstein-Uhlenbeck process with added Gaussian noise
@model oupn(rn,T,delta_t,::Type{R}=Vector{Float64}) where {R} = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ Beta(5.0,1.0)
    noise_ampl ~ Uniform(0.0,1)

    b = exp(-delta_t/tau)
    r = R(undef, T)

    r[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        r[i] ~ Normal(r[i-1]*b,sqrt(ampl*(1-b^2)))
    end
    rn ~ MvNormal(r,sqrt(noise_ampl))
end

x1, x2 = corr_osc(-0.8)

p1 = Plots.plot(x1)
p1 = Plots.plot!(x2)

pearson = Statistics.cor(x1,x2)
println(pearson)

# lets see whether we can estimate c from the data
y1 = x1 .+ x2
y2 = x1 .- x2
chn = sample(ou_corr(y1,y2,length(y1),0.1), NUTS(0.65), 10000)

print(describe(chn))
p2 = plot(chn)

ampl1 = Array(chn[:ampl1])
ampl2 = Array(chn[:ampl2])

if ampl1>ampl2
    c = (ampl1 .- ampl2)./ampl2
else
    c = (ampl1 .- ampl2)./ampl1
end

println("A1,A2: ",mean(ampl1),",",mean(ampl2))
println("c estimate: ",mean(c),"std: ",std(c))


T = 3
n = 2
m = 1

A = [1.0 1.0; 0.0 1.0]
B = [0.0; 1.0]

Q = [1.0 0.0; 0.0 1.0]
R = [1.0]

x_init = [1.0; 0.0]

nz = T * n + (T - 1) * m + T * n

function unpack(z)
	x3 = z[1:n]
	u2 = z[n .+ (1:m)]
	x2 = z[n + m .+ (1:n)]
	u1 = z[n + m + n .+ (1:m)]
	x1 = z[n + m + n + m .+ (1:n)]
	y3 = z[n + m + n + m + n .+ (1:n)]
	y2 = z[n + m + n + m + n + n .+ (1:n)]
	y1 = z[n + m + n + m + n + n + n .+ (1:n)]

	return x3, u2, x2, u1, x1, y3, y2, y1
end

z0 = rand(nz)

function lagrangian(z, x_init)
	x3, u2, x2, u1, x1, y3, y2, y1 = unpack(z)

	L = 0.0

	# objective
	L += transpose(x3) * Q * x3
	L += (transpose(u2) * R * u2)[1]
	L += transpose(x2) * Q * x2
	L += (transpose(u1) * R * u1)[1]
	L += transpose(x1) * Q * x1

	# constraints
	L += transpose(y3) * (A * x2 + B * u2[1] - x3)
	L += transpose(y2) * (A * x1 + B * u1[1] - x2)
	L += transpose(y1) * (x_init - x1)

	return L
end

@variables z_sym[1:nz]
@variables x_sym[1:n]

L = lagrangian(z_sym, x_sym)
L = simplify.(L)

dL = Symbolics.gradient(L, z_sym)
ddL = Symbolics.hessian(L, z_sym)
ddLx = Symbolics.jacobian(dL, x_sym)

L_grad = eval(Symbolics.build_function(dL, z_sym, x_sym)[1])
L_hess = eval(Symbolics.build_function(ddL, z_sym, x_sym)[1])
Lx_hess = eval(Symbolics.build_function(ddLx, z_sym, x_sym)[1])

function solve(z0, x_init)
	z = copy(z0)

	∇L = L_grad(z, x_init)

	println()
	for i = 1:10
		r_norm = norm(∇L)
		println("iter ($i)")
		println("	grad norm: $(r_norm)")
		r_norm < 1.0e-8 && break

		∇²L = L_hess(z, x_init)

		Δ = ∇²L \ ∇L

		α = 1.0
		ẑ = z - α * Δ

		∇L̂ = L_grad(ẑ, x_init)

		iter = 1
		while norm(∇L̂) > norm(∇L)
			α *= 0.5
			ẑ = z - α * Δ
			∇L̂ = L_grad(ẑ, x_init)
			iter += 1

			if iter > 100
				break
			end
		end

		z = ẑ
		∇L = ∇L̂
	end

	return z, ∇L, L_hess(z, x_init)
end

z_sol, ∇L_sol, ∇²L_sol = solve(z0, x_init)
K_direct = (∇²L_sol \ Lx_hess(z_sol, x_init))[n + m + n .+ (1:m), :]

# LQR solution
K, P = tvlqr([A for t = 1:T-1], [B for t = 1:T-1], [Q for t = 1:T], [R[1,1] for t = 1:T-1])

@assert norm(K[1] - K_direct) < 1.0e-8


using Pkg; Pkg.activate(pwd())
using DEMC
using Random
using LinearAlgebra
using Distributions
Random.seed!(319531501)

# generate data
nobs = 1000
npar = 25
Σ = ones(npar, npar)
for i = 1:npar
    for iprime = 1:i
        if i != iprime
            Σ[i, iprime] = 0.25
            Σ[iprime, i] = 0.25
        end
    end
end

DataDistribution = MvNormal(zeros(npar), Σ)
X = zeros(nobs, npar+1)
X[:,1] .= 1.
for iobs = 1:nobs
    X[iobs, 2:end] .= rand(DataDistribution)
end

β = 1 .+ rand(npar+1)*3
y = X*β + randn(nobs)

# log objective function - for correct std errors need to use optimal weights
log_obj(b) = -0.5 * sum((y.-X*b).^2)

# set up of DEMCz chain
ndim = length(β)
opts = DEMC.demcopt(ndim)
opts.blockindex = [1:ndim] # parameter blocks: here choose all parameters to be updated simultaenously
opts.Nblocks = length(opts.blockindex)
opts.eps_scale = 1e-5*ones(ndim) # scale of random error around DE update
opts.γ = 2. # scale of DE update, 2.38 is the "optimal" number for a normal distribution
opts.N = 5 # number of chains
opts.K = 10 # every K steps add current N draws to Z
# Number of iterations in Chain
opts.Ngeneration = 100000
opts.autostop = :Rhat
opts.autostop_every = 2000
opts.autostop_Rhat = 1.1
opts.print_step = 1000
opts.verbose = false

Z = randn((10*ndim, ndim))

# run chain
mc, Z = DEMC.demcz_sample(log_obj, Z, opts)

# drop first half of chain
N, Npar, Ntot = size(mc.chain)
keep = Int(Ntot-opts.autostop_every)+1:Ntot
Ngen_burned = length(keep)
chain_burned = mc.chain[:,:,keep]
chainflat = DEMC.flatten_chain(chain_burned, N, Ngen_burned, Npar)'

# covariance of estimates
b, Σb = DEMC.mean_cov_chain(chain_burned)
# correct std errors for incorrect weighting matrix
# OLS
bols = (X'*X)\(X'*y)
ui = y .- X*bols
s2hat = (1/(nobs-npar)) * sum(ui.^2)
Σols = inv(X'X)*s2hat

# difference simulated to OLS
println("\n MCMC estimates: ", b, "\n dist to OLS: ", b .- bols)

accept_ratio, Rhat = DEMC.convergence_check(mc.chain, mc.log_obj, "none"; verbose = false, parnames = [])


# Test Benders Algorithm for Mixed Integer Linear Programming
# Edward J. Xu, edxu96@outlook.com
# April 5th, 2019
# ---------------------------------------------------------------------------
push!(LOAD_PATH, "$(homedir())/Desktop/Benders Decomposition, MILP")
cd("$(homedir())/Desktop/Benders Decomposition, MILP")
using Benders
# ---------------------------------------------------------------------------
# Test Example 1: mat_a and mat_b being column vector
n_x = 1
n_y = 1
vec_min_y = hcat([0])
vec_max_y = hcat([10])
vec_c = hcat([5])
vec_f = hcat([-3])
vec_b = hcat([4; 0; -13])
mat_a = hcat([1; 2; 1])
mat_b = hcat([2; -1; -3])
Benders.milp(n_x = n_x,
             n_y = n_y,
             vec_min_y = vec_min_y,
             vec_max_y = vec_max_y,
             vec_c = vec_c,
             vec_f = vec_f,
             vec_b = vec_b,
             mat_a = mat_a,
             mat_b = mat_b,
             epsilon = 0,
             timesIterationMax = 5)
# ---------------------------------------------------------------------------
# Test Example 2: mat_a and mat_b being matrix
n_x_2 = 2
n_y_2 = 2
vec_min_y_2 = hcat([0; 0])
vec_max_y_2 = hcat([10; 10])
vec_c_2 = hcat([5; 3])
vec_f_2 = hcat([-3; 1])
vec_b_2 = hcat([4; 0; -13])
mat_a_2 = [1 3; 2 1; 1 -5]
mat_b_2 = [2 -4; -1 2; -3 1]
Benders.milp(n_x = n_x_2,
             n_y = n_y_2,
             vec_min_y = vec_min_y_2,
             vec_max_y = vec_max_y_2,
             vec_c = vec_c_2,
             vec_f = vec_f_2,
             vec_b = vec_b_2,
             mat_a = mat_a_2,
             mat_b = mat_b_2,
             epsilon = 0.000001,
             timesIterationMax = 10)
# ---------------------------------------------------------------------------
# Test Example 3: sub and ray problem
n_x_3 = 2
n_y_3 = 2
vec_min_y_3 = hcat([0; 0])
vec_max_y_3 = hcat([2; 2])
vec_c_3 = hcat([2; 6])
vec_f_3 = hcat([2; 3])
vec_b_3 = hcat([5; 4])
mat_a_3 = [-1 2; 1 -3]
mat_b_3 = [3 -1; 2 2]
Benders.milp(n_x = n_x_3,
             n_y = n_y_3,
             vec_min_y = vec_min_y_3,
             vec_max_y = vec_max_y_3,
             vec_c = vec_c_3,
             vec_f = vec_f_3,
             vec_b = vec_b_3,
             mat_a = mat_a_3,
             mat_b = mat_b_3,
             epsilon = 0.000001,
             timesIterationMax = 10)


#
# Denoise an SPD Example with Douglas Rachford
#
# Denoise an SPD Example with Parallel Douglas-Rachford to minimize the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data and cost as for SPD_Image_CP_vs_CPPA
#
# This example and its data is used for comparison in Section 6.2 of
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez:
# > _Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds_,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using Manopt, Manifolds
using Images, CSV, DataFrames, LinearAlgebra, JLD2, Dates
#
# Settings
experiment_name = "SPD_Image_DR"
export_orig = true
export_result = true
export_table = true
use_debug = true
asy_render_detail = 2
results_folder = joinpath(@__DIR__, "Image_TV")
comparison_data = joinpath(
    @__DIR__, "..", "CyclicProximalPoint", "Image_TV", "SPD_Image_CPPA-cost.jld2"
)
!isdir(results_folder) && mkdir(results_folder)
#
# Parameters
η = 0.58
λ = 0.93
α = 6.0
#
# Manifold & Data
f = artificial_SPD_image2(32)
if export_orig
    fn = joinpath(results_folder, experiment_name * "-orig.asy")
    asymptote_export_SPD(fn; data=f, scale_axes=(7.5, 7.5, 7.5))
    render_asymptote(fn; render=asy_render_detail)
end
sC = StopAfterIteration(400)
try
    cppa_data = load(comparison_data)
    cost_threshold = cppa_data["cost_function_value"]
    cppa_iter = cppa_data["iterations"]
    global sC = StopWhenAny(StopAfterIteration(400), StopWhenCostLess(cost_threshold))
    @info "Comparison to CPPA (`SPDImage_CPPA.jl`) and its cost after $(cppa_iter) iterations of $cost_threshold."
catch e
    msg = sprint(showerror, e)
    @info "Comparison to CPPA (`CyclicProximalPoint/SPDImage_CPPA.jl`) not possible, data file $(comparison_data) either missing or corrupted.\n Error: $(msg).\n\n Starting with a default stopping criterion."
end
#
# Build Problem for L2-TV
pixelM = SymmetricPositiveDefinite(3);
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
d = length(size(f))
rep(d) = (d > 1) ? [ones(Int, d)..., d] : d
fidelity(x) = 1 / 2 * distance(M, x, f)^2
Λ(x) = forward_logs(M, x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM), repeat(x, rep(d)...), Λ(x)), 1)
#
# Setup & Optimize
print("--- Douglas–Rachford with η: $(η) and λ: $(λ) ---\n")
cost(x) = 1 / α * fidelity(x[1]) + prior(x[1])
N = PowerManifold(M, NestedPowerRepresentation(), 5)
prox1 = (η, x) -> [prox_distance(M, η, f, x[1]), prox_parallel_TV(M, α * η, x[2:5])...]
prox2 = (η, x) -> fill(mean(M, x, GradientDescentEstimation(); stop_iter=4), 5)
x0 = fill(f, 5)
@time o = DouglasRachford(
    N,
    cost,
    [prox1, prox2],
    x0;
    λ=i -> η,
    α=i -> λ, # map from Paper notation of BPS16 to toolbox notation
    debug=use_debug ? [:Iteration, " | ", :Cost, "\n", 10, :Stop] : missing,
    record=export_table ? [:Iteration, :Cost] : missing,
    stopping_criterion=sC,
    return_options=true,
)
y = get_solver_result(o)[1]
export_table && (r = get_record(o))
#
# Result
numIter = length(r)
if export_result
    fn = joinpath(
        results_folder,
        experiment_name * "img-result-$(numIter)-α$(replace(string(α), "." => "-")).asy",
    )
    asymptote_export_SPD(fn; data=y, scale_axes=(7.5, 7.5, 7.5))
    render_asymptote(fn; render=asy_render_detail)
end
if export_table
    # scale cost back for saving such that its comparable with the other two results
    A = cat([ri[1] for ri in r], [ri[2] / α for ri in r]; dims=2)
    CSV.write(
        joinpath(results_folder, experiment_name * "-Cost.csv"),
        DataFrame(A);
        writeheader=false,
    )
end


# # [Linear Time Continuous System with Controls](@id linear_continuous_controls)
#
# Now we will extend the [`previous example`](@ref linear_continuous) by adding some exegeneous control signals.
# As always, we will generate some data via `OrdinaryDiffEq.jl`

using DataDrivenDiffEq
using ModelingToolkit
using LinearAlgebra
using OrdinaryDiffEq
#md using Plots

A = [-0.9 0.2; 0.0 -0.2]
B = [0.0; 1.0]
u0 = [10.0; -10.0]
tspan = (0.0, 10.0)

f(u,p,t) = A*u .+ B .* sin(0.5*t)

sys = ODEProblem(f, u0, tspan)
sol = solve(sys, Tsit5(), saveat = 0.05);

# We will use the data provided by our problem, but add the control signal `U = sin(0.5*t)` to it.
X = Array(sol)
t = sol.t
control(u,p,t) = [sin(0.5*t)]
prob = ContinuousDataDrivenProblem(X, t, U = control)

# And plot the problems data.

#md plot(prob)

# Again, we will use `gDMD` to estimate the systems dynamics. Since we have a control signal
# defined in the problem, the algorithm will detect it automatically and use `gDMDc`:

res = solve(prob, DMDSVD(), digits = 1)
#md println(res)

# We see that the system has been recovered correctly, indicated by the small error and high AIC score of the result. We can confirm this by looking at the resulting [`Basis`](@ref)

#md
system = result(res)
#md println(system)

# And also plot the prediction of the recovered dynamics

#md plot(res)

# Again, we can have a look at the generator of the system, which is independent from the inputs.

generator(system)

# Sticking to the same procedure as earlier, we now use a linear sparse regression to solve the problem

@parameters t
@variables x[1:2](t) u[1:1](t)

basis = Basis([x; u], x, controls = u, independent_variable = t, name = :LinearBasis)
#md print(basis) #hide

# Note that we added a new variable `u[1](t)` as a control to both the equations and the basis constructor.
# Afterwards, we simply `solve` the already defined problem with our `Basis` and a `SparseOptimizer`

sparse_res = solve(prob, basis, STLSQ(1e-1))
#md println(sparse_res)

# Which holds the same equations
sparse_system = result(sparse_res)
#md println(sparse_system)

# Again, we can have a look at the result

#md plot(
#md     plot(prob), plot(sparse_res), layout = (1,2)
#md )

# Both results can be converted into an `ODESystem`. To include the control signal, we simply
# substitute the control variables in the corresponding equations.

subs_control = (u[1] => sin(0.5*t))

eqs = map(equations(sparse_system)) do eq
    eq.lhs ~ substitute(eq.rhs, subs_control)
end

@named sys = ODESystem(
    eqs,
    get_iv(sparse_system),
    states(sparse_system),
    parameters(sparse_system)
    );

# And simulated using `OrdinaryDiffEq.jl` using the (known) initial conditions and the parameter mapping of the estimation.

x0 = [x[1] => u0[1], x[2] => u0[2]]
ps = parameter_map(sparse_res)

ode_prob = ODEProblem(sys, x0, tspan, ps)
estimate = solve(ode_prob, Tsit5(), saveat = prob.t);

# And look at the result
#md plot(sol, color = :black)
#md plot!(estimate, color = :red, linestyle = :dash)

#md # ## [Copy-Pasteable Code](@id linear_continuous_copy_paste)
#md #
#md # ```julia
#md # @__CODE__
#md # ```

@test all(aic(sparse_res) .>= 1e3) #src
@test all(l2error(sparse_res) .<= 5e-1) #src
@test all(determination(sparse_res) .>= 0.97) #src
@test Array(sol) ≈ Array(estimate) rtol = 5e-2 #src


#*************************************************************************#
#
# AbstractType
#
#*************************************************************************#
abstract type ElementwiseOperator <: AparametricModel end

#*************************************************************************#
#
# Generate Buffers
#
#*************************************************************************#
function dx_buffer(f::ElementwiseOperator)
    return Diagonal{Float64}(zeros(f.noutputs))
end

#*************************************************************************#
#
# LogisticFunction
#
#*************************************************************************#
"""
    LogisticFunction(n)

Initialize `ParametricModel` representing logistic function. It applies
`f(x) = 1 / (1+exp(-x))` elementwise to a given input vector.
`n` is the input vector dimension.
"""
struct LogisticFunction <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
end

function LogisticFunction(n::Int)
    LogisticFunction(n, n)
end

function _evaluate!(f::LogisticFunction,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})

    z .= 1 ./ (1+exp.(-x))

    if !isempty(dx)
        dx.diag .=  z .* (1-z)
    end
end

#*************************************************************************#
#
# Hyperbolic Tangent
#
#*************************************************************************#
"""
    HyperbolicTangent(n)

Initialize `ParametricModel` representing hyperbolic tangent. It
applies `f(x) = tanh(x)` elementwise to a given input vector.
 `n` is the input vector dimension.
"""
struct HyperbolicTangent <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
end

function HyperbolicTangent(n::Int)
    HyperbolicTangent(n, n)
end

function _evaluate!(f::HyperbolicTangent,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})
    z .= tanh.(x)

    if !isempty(dx)
        dx.diag .=  1 - z.^2
    end
end

#*************************************************************************#
#
# AffineMap
#
#*************************************************************************#
"""
    AffineMap(a[, b])

Initialize `ParametricModel` representing affine function. It applies
`f(x) = a*x + b` elementwise to a given input vector. `a` and
`b` are vectors.
"""
struct AffineMap <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
    a::Vector{Float64}
    b::Vector{Float64}
end

function AffineMap{R<:Real, S<:Real}(a::Vector{R}, b::Vector{S}=zeros(length(a)))
    if length(a) != length(b)
        throw(ArgumentError("Vector `a` and `b` should have the same size."))
    end

    n = length(a)
    AffineMap(n, n, a, b)
end

function _evaluate!(f::AffineMap,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})
    z .= f.a.*x + f.b

    if !isempty(dx)
        dx.diag .=  f.a
    end
end

#*************************************************************************#
#
# Identity
#
#*************************************************************************#
"""
    Identity(n)

Initialize `ParametricModel` representing identity function.  It applies
`f(x) = x` elementwise to a given input vector. `n` is the input vector
dimension.
"""
struct Identity <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
end

function Identity(n::Int)
    Identity(n, n)
end

function _evaluate!(f::Identity,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})
    z .= x

    if !isempty(dx)
        dx.diag .=  1
    end
end


using LinearAlgebra;
using Printf;

function linspace(start::Number, stop::Number, n::Int64)
	linx = zeros(Float64, n);
	delta = (stop - start) / (n - 1);
	for i = 1:n
		linx[i] = start + delta * (i - 1)
	end
	return linx;
end

lx = 1.0;
t = 0.0
nx = 11;
dx = lx / (nx - 1);
dt = 1e-2;
tlim = 0.1;
kappa = 1.0
beta = dt * kappa / 2.0dx^2;
t_save = linspace(0.0, tlim, nx);
save_count = 1;

function solve_linear(n::Int64, a::Float64, b::Float64, c::Float64, r::Array)
	A = zeros(n, n);
	for i = 2:n+1
		for j = 2:n+1
			if(i-1 == j-1-1)
				A[i-1, j-1] = c;
			elseif (i-1==j-1)
				A[i-1,j-1] = b;
			elseif (i-1-1 == j-1)
				A[i-1,j-1] = a;
			end
		end
	end
	return A \ r;
end

function crank_nicolson(u::Array)
	a = Float64(0.0);
	b = Float64(0.0);
	c = Float64(0.0);
	r = zeros(Float64, nx - 2);
	for i = 2:nx-1
		a = -beta;
		b = 1.0 + 2.0beta;
		c = -beta;
		r[i-1] = u[i] + beta*(u[i+1] - 2.0u[i] + u[i-1]);
	end
	r[begin] = r[begin] - a*u[begin]
	r[end] = r[end] - c * u[end];
	u[2:end-1] = solve_linear(nx - 2, a, b, c, r)
end

function output(x::Array, u::Array, t::Float64)
	filename = @sprintf("cn_t=%4.3f.dat", t);
	open(filename, "w") do io
		for i = 1:nx
			write(io, @sprintf("%.5f\t%.5f\n", x[i], u[i]));
		end
	end;
end

function main()
	@printf("beta = %.5f, dt = %.5f, dx = %.5f\n", beta, dt, dx);
	u = ones(Float64, nx);
	u[begin] = u[end] = 0.0;
	x = linspace(0, lx, nx);
	while t < tlim
		output(x, u, t);
		crank_nicolson(u);
		global t = t + dt;
	end
end

main()


using LinearAlgebra
include("poisson_linear_system.jl")
include("simpson_integrate.jl")

function poisson_error(numelements)
    xL = -1.0
    xR = +1.0
    wavenumber = pi / 2
    bcleft = bcright = 0.0

    numpts = numelements + 1
    stepsize = (xR - xL) / numelements
    xrange = range(xL, stop = xR, length = numpts)

    rhs = cos.(wavenumber * xrange)
    op = -poisson_linear_system(numpts)
    rhs[1] = bcleft / stepsize^2
    rhs[end] = bcright / stepsize^2
    sol = stepsize^2 * (op \ rhs)
    exactsolution =
        1.0 / wavenumber^2 * (cos.(wavenumber * xrange) .- cos(wavenumber))
    err = simpson_integrate(sol - exactsolution, stepsize)
    return err
end

function convergence_rate(err,dx)
    return diff(log.(err)) ./ diff(log.(dx))
end

powers = 1:10
numelements = 2 .^ powers
stepsize = 2.0 ./ numelements
err = poisson_error.(numelements)
rate = convergence_rate(err,stepsize)

# using Plots
# plot(stepsize,err,xscale = :log, yscale = :log)
# scatter!(stepsize,err)


#=============================================

Definitions of various similarity functions

=============================================#

using QuadGK
using LinearAlgebra: dot, norm

#====================
Definitions of built-in similarity functions
====================#

#====================
Cosine similarity
====================#

@doc raw"""
    cossim(x,y)

Computes the cosine similarity between two inputs ``x`` and ``y``. Cosine similarity is defined as

``\text{cossim}(x,y) = \frac{\left\langle x,y\right\rangle}{\|x\|\cdot\|y\|}``

where ``\left\langle\cdot,\cdot\right\rangle`` is an inner product (e.g. dot product) and ``\|\cdot\|`` is its derived norm. This is roughly interpreted as being related to the angle between the inputs ``x`` and ``y``: when ``x`` and ``y`` have low angle between them, `cossim(x,y)` is high (close to ``1``). When ``x`` and ``y`` have large angle between them, `cossim(x,y)` is low (close to ``-1``).

# Arguments
- `x` and `y`: two inputs for which `dot(x,y)`, `norm(x)`, and `norm(y)` are defined.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> using LinearAlgebra: dot, norm;

julia> x, y = rand(4), rand(4);

julia> cossim(x,y) == dot(x,y) / (norm(x) * norm(y))
true

julia> z = rand(5);

julia> cossim(x,z)
ERROR: DimensionMismatch("dot product arguments have lengths 4 and 5")
```

See also: [`SimHash`](@ref)
"""
function cossim(x::AbstractVector, y::AbstractVector)
    norm_x = norm(x)
    norm_y = norm(y)

    if norm_x == 0 || norm_y == 0
        "x and y must be nonzero" |> ErrorException |> throw
    end

    dot(x,y) / (norm_x * norm_y)
end

function cossim(f, g, interval::LSHFunctions.RealInterval)
    norm_f = L2_norm(f, interval)
    norm_g = L2_norm(g, interval)

    if norm_f == 0 || norm_g == 0
        "f and g must be nonzero" |> ErrorException |> throw
    end

    inner_prod(f, g, interval) / (norm_f * norm_g)
end

#====================
L^p distance
====================#

@doc raw"""
    ℓp(x::AbstractVector, y::AbstractVector, p::Real=2)
    ℓ1(x::AbstractVector, y::AbstractVector)
    ℓ2(x::AbstractVector, y::AbstractVector)

Computes the ``\ell^p`` distance between a pair of vectors, given by

``\ell^p(x,y) \coloneqq \|x - y\|_p = \left(\sum_i \left|x_i - y_i\right|^p\right)^{1/p}``

`ℓ1(x,y)` is the same as `ℓp(x,y,1)`, and `ℓ2(x,y)` is the same as `ℓp(x,y,2)`.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = [1, 2, 3];

julia> y = [4, 5, 6];

julia> ℓp(x,y,2) == (abs(1-4)^2 + abs(2-5)^2 + abs(3-6)^2)^(1/2)
true

julia> ℓp(x,y,3) == (abs(1-4)^3 + abs(2-5)^3 + abs(3-6)^3)^(1/3)
true
```

See also: [`ℓp_norm`](@ref), [`L1Hash`](@ref), [`L2Hash`](@ref)
"""
ℓp(x::AbstractVector, y::AbstractVector, p::Real=2) = Lp(x, y, p)

@doc (@doc ℓp)
ℓ1(x::AbstractVector, y::AbstractVector) = L1(x, y)

@doc (@doc ℓp)
ℓ2(x::AbstractVector, y::AbstractVector) = L2(x, y)

@doc raw"""
    Lp(x::AbstractVector, y::AbstractVector, p::Real=2)
    L1(x::AbstractVector, y::AbstractVector)
    L2(x::AbstractVector, y::AbstractVector)

Computes the ``ℓ^p`` distance between a pair of vectors ``x`` and ``y``. Identical to `ℓp(x,y,p)`, `ℓ1(x,y)`, and `ℓ2(x,y)`, respectively.

See also: [`ℓp`](@ref)
"""
function Lp(x::AbstractVector{T}, y::AbstractVector, p::Real=2) where T
    if p ≤ 0
        "p must be positive" |> ErrorException |> throw
    elseif length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs(x[ii] - y[ii])^p
    end

    return result^(1/p)
end

@doc (@doc Lp)
function L1(x::AbstractVector{T}, y::AbstractVector) where T
    if length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs(x[ii] - y[ii])
    end

    return result
end

@doc (@doc Lp)
function L2(x::AbstractVector{T}, y::AbstractVector) where T
    if length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs2(x[ii] - y[ii])
    end

    return √result
end

# Function space L^p distances

@doc raw"""
    Lp(f, g, interval::LSHFunctions.RealInterval, p)
    L1(f, g, interval::LSHFunctions.RealInterval)
    L2(f, g, interval::LSHFunctions.RealInterval)

Computes the ``L^p`` distance between two functions, given by

``L^p(f,g) \coloneqq \|f - g\|_p = \left(\int_a^b \left|f(x) - g(x)\right|^p \hspace{0.15cm} dx\right)^{1/p}``

# Examples
Below we compute the ``L^1``, ``L^2``, and ``L^3`` distances between ``f(x) = x^2 + 1`` and ``g(x) = 2x`` over the interval ``[0,1]``. The distances are computed by evaluating the integral

``\left(\int_0^1 \left|f(x) - g(x)\right|^p \hspace{0.15cm}dx\right)^{1/p} = \left(\int_0^1 \left|x^2 - 2x + 1\right|^p \hspace{0.15cm}dx\right)^{1/p} = \left(\int_0^1 (x - 1)^{2p} \hspace{0.15cm}dx\right)^{1/p}``

for ``p = 1``, ``p = 2``, and ``p = 3``.

```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = x^2 + 1; g(x) = 2x;

julia> interval = @interval(0 ≤ x ≤ 1);

julia> Lp(f, g, interval, 1) ≈ L1(f, g, interval) ≈ 3^(-1)
true

julia> Lp(f, g, interval, 2) ≈ L2(f, g, interval) ≈ 5^(-1/2)
true

julia> Lp(f, g, interval, 3) ≈ 7^(-1/3)
true
```

See also: [`Lp_norm`](@ref), [`ℓp`](@ref)
"""
Lp(f, g, interval::LSHFunctions.RealInterval, p::Real=2) =
    Lp_norm(x -> f(x) - g(x), interval, p)

@doc (@doc Lp)
L1(f, g, interval::LSHFunctions.RealInterval) = L1_norm(x -> f(x) - g(x), interval)

@doc (@doc Lp)
L2(f, g, interval::LSHFunctions.RealInterval) = L2_norm(x -> f(x) - g(x), interval)

#====================
Jaccard similarity
====================#

@doc raw"""
    jaccard(A::Set, B::Set) :: Float64

Computes the Jaccard similarity between sets ``A`` and ``B``, which is defined as

``\text{Jaccard}(A,B) = \frac{\left|A \cap B\right|}{\left|A \cup B\right|}``

# Arguments
- `A::Set`, `B::Set`: two sets whose Jaccard similarity we would like to compute.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> A, B = Set([1, 2, 3]), Set([2, 3, 4]);

julia> jaccard(A,B)
0.5

julia> jaccard(A,B) == length(A ∩ B) / length(A ∪ B)
true
```

See also: [`MinHash`](@ref)
"""
function jaccard(A::Set, B::Set) :: Float64
    if isempty(A)
        # Use the convention that if A = B = ∅, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        length(A ∩ B) / length(A ∪ B)
    end
end

@doc raw"""
    function jaccard(x::BitArray{1}, y::BitArray{1})

Computes the Jaccard similarity between a pair of binary vectors:

``J(x, y) = \frac{\sum_{i} \min{(x_i,y_i)}}{\sum_{i} \max{(x_i,y_i)}}``

# Arguments
- `x::BitArray{1}`, `y::BitArray{1}`: two binary vectors, in the form of `BitArray`s.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = BitArray([true, false, true, true, false]);

julia> y = BitArray([false, false, true, true, true]);

julia> jaccard(x,y)
0.5
```
"""
function jaccard(x::BitArray{1}, y::BitArray{1}) :: Float64
    union = sum(x .| y)
    if union == 0
        # Use the convention that if x and y are full of zeros, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        intersection = sum(x .& y)
        intersection / union
    end
end

@doc raw"""
    function jaccard(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Computes the Jaccard similarity between a pair of vectors of real numbers:

``J(x, y) = \frac{\sum_{i} \min{(x_i,y_i)}}{\sum_{i} \max{(x_i,y_i)}}``

# Arguments
- `x::AbstractVector{<:Real}`, `y::AbstractVector{<:Real}`: a pair of vectors containing real numbers (subtypes of `Real`).

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = [0.8, 0.1, 0.3, 0.4, 0.1];

julia> y = [1.0, 0.6, 0.0, 0.4, 0.5];

julia> jaccard(x,y)
0.5
```
"""
function jaccard(x::AbstractVector{T},
                 y::AbstractVector) :: Float64 where {T <: Real}
    if length(x) != length(y)
        DimensionMismatch("dimensions must match") |> throw
    end

    intersection = T(0)
    union = T(0)

    @inbounds @simd for ii = 1:length(x)
        if 0 ≤ x[ii] ≤ y[ii]
            intersection += x[ii]
            union += y[ii]
        elseif 0 ≤ y[ii] < x[ii]
            intersection += y[ii]
            union += x[ii]
        else
            ErrorException("vectors must have non-negative elements") |> throw
        end
    end

    if union == T(0)
        # Use the convention that if x and y are full of zeros, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        Float64(intersection / union)
    end
end

jaccard(x::AbstractVector{<:Integer}, y::AbstractVector{<:AbstractFloat}) =
    jaccard(y, x)

@doc raw"""
    function jaccard(A::Set{<:K},
                     B::Set{<:K},
                     weights::Dict{K,V}) where {K,V<:Number}

Computes the weighted Jaccard similarity between two sets:

``J(x, y) = \frac{\sum_{x\in A\cap B} w_x}{\sum_{y\in A\cup B} w_y}``

# Arguments
- `A::Set`, `B::Set`: two sets whose Jaccard similarity we would like to compute.
- `weights::Dict`: a dictionary mapping symbols in the sets `A` and `B` to numerical weights. These weights must be positive.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> A = Set(["a", "b", "c"]);

julia> B = Set(["b", "c", "d"]);

julia> W = Dict("a" => 0.2, "b" => 2.4, "c" => 0.6, "d" => 1.8);

julia> jaccard(A,B,W)
0.6
```
"""
function jaccard(A::Set{<:K},
                 B::Set{<:K},
                 weights::Dict{K,V}) :: Float64 where {K,V<:Real}

    union_weight = V(0)

    for el in A ∪ B
        w = weights[el]
        if w < 0
            ErrorException("weights must be non-negative") |> throw
        end
        union_weight += w
    end

    intersection_weight = sum(weights[el] for el in A ∩ B)

    # By convention, if A = B = ∅, their Jaccard similarity is zero
    if union_weight == V(0)
        Float64(0)
    else
        Float64(intersection_weight / union_weight)
    end
end

#====================
Inner product and norms
====================#

### Inner products

@doc raw"""
    inner_prod(x::AbstractVector, y::AbstractVector)

Computes the ``\ell^2`` inner product (dot product)

``\left\langle x, y\right\rangle = \sum_i x_iy_i``

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> using LinearAlgebra: dot;

julia> x, y = randn(4), randn(4);

julia> inner_prod(x,y) ≈ dot(x,y)
true
```
"""
inner_prod(x::AbstractVector, y::AbstractVector) = dot(x,y)

# 1-dimensional inner product between L^2 functions
@doc raw"""
    inner_prod(f, g, interval::LSHFunctions.RealInterval)

Computes the ``L^2`` inner product

``\left\langle f, g\right\rangle = \int_a^b f(x)g(x) \hspace{0.15cm} dx``

where the interval we're integrating over is specified by the `interval` argument.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = cos(x); g(x) = sin(x);

julia> inner_prod(f, g, @interval(0 ≤ x ≤ π/2)) ≈ 1/2
true
```
"""
inner_prod(f, g, interval::LSHFunctions.RealInterval) =
    quadgk(x -> f(x)g(x), interval.lower, interval.upper)[1]

### L^p norms
@doc raw"""
    Lp_norm(x::AbstractVector, p::Real = 2)
    L1_norm(x::AbstractVector)
    L2_norm(x::AbstractVector)

Compute the ``\ell^p`` norm of a vector ``x``. Identical to `ℓp_norm(x,p)`, `ℓ1_norm(x)`, and `ℓ2_norm(x)`, respectively.

See also: [`ℓp_norm`](@ref)
"""
Lp_norm(x::AbstractVector, p::Real = 2) = norm(x,p)

@doc (@doc Lp_norm)
L1_norm(x::AbstractVector) = norm(x,1)

@doc (@doc Lp_norm)
L2_norm(x::AbstractVector) = norm(x)

@doc raw"""
    ℓp_norm(x::AbstractVector, p::Real = 2)
    ℓ1_norm(x::AbstractVector)
    ℓ2_norm(x::AbstractVector)

Compute the ``\ell^p`` norm of a point ``x``, defined as

``\|x\|_p = \left(\sum_i \left|x_i\right|^p\right)^{1/p}``

# Examples

```jldoctest; setup = :(using LSHFunctions)
julia> x = randn(4);

julia> ℓp_norm(x, 1) ≈ ℓ1_norm(x) ≈ (map(u -> abs(u)^1, x) |> sum)^(1/1)
true

julia> ℓp_norm(x, 2) ≈ ℓ2_norm(x) ≈ (map(u -> abs(u)^2, x) |> sum)^(1/2)
true

julia> ℓp_norm(x, 3) ≈ (map(u -> abs(u)^3, x) |> sum)^(1/3)
true
```

See also: [`ℓp`](@ref), [`Lp_norm`](@ref)
"""
ℓp_norm(x::AbstractVector, p::Real = 2) = Lp_norm(x, p)

@doc (@doc ℓp_norm)
ℓ1_norm(x::AbstractVector) = L1_norm(x)

@doc (@doc ℓp_norm)
ℓ2_norm(x::AbstractVector) = L2_norm(x)

# 1-dimensional L^p norms

@doc raw"""
    Lp_norm(f, interval::LSHFunctions.RealInterval, p::Real=2)
    L1_norm(f, interval::LSHFunctions.RealInterval)
    L2_norm(f, interval::LSHFunctions.RealInterval)

Computes the ``L^p`` function-space norm of a function ``f``, which is given by the equation

``\|f\|_p = \left(\int_a^b \left|f(x)\right|^p \hspace{0.15cm} dx\right)^{1/p}``

`L1_norm(f, interval)` is the same as `Lp_norm(f, interval, 1)`, and `L2_norm(f, interval)` is the same as `Lp_norm(f, interval, 2)`.

# Examples

```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = x;

julia> interval = @interval(0 ≤ x ≤ 1);

julia> Lp_norm(f, interval, 1) ≈ L1_norm(f, interval) ≈ 2^(-1/1)
true

julia> Lp_norm(f, interval, 2) ≈ L2_norm(f, interval) ≈ 3^(-1/2)
true

julia> Lp_norm(f, interval, 3) ≈ 4^(-1/3)
true
```
"""
Lp_norm(f, interval::LSHFunctions.RealInterval, p::Real=2) = (quadgk(x -> abs(f(x)).^p, interval.lower, interval.upper)[1])^(1/p)

@doc (@doc Lp_norm)
L1_norm(f, interval::LSHFunctions.RealInterval) = quadgk(x -> abs(f(x)), interval.lower, interval.upper)[1]

@doc (@doc Lp_norm)
L2_norm(f, interval::LSHFunctions.RealInterval) = √quadgk(x -> abs2(f(x)), interval.lower, interval.upper)[1]


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

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

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

# Functional form
```
f(x | x̂, FWHM) = exp[-4ln(2) * ||x - x̂|| / FWHM^2]
```
where `x̂` and `x` are position vectors (indices) `||⋅||` represents the square-distance, and `FWHM` is the full width at half-maximum. If `FWHM` is a scalar, the Gaussian distribution will be isotropic. If `FWHM` is a vector or tuple, the weighting is applied along each axis (diagonal).
"""
gaussian(T, px, py; x, y, fwhm, amp=one(T), theta=0, bkg=0) = convert(T, _gaussian(px, py, x, y, fwhm, amp, theta, bkg))

"""
    normal

An alias for [`gaussian`](@ref)
"""
const normal = gaussian

# this is the factor to convert 1/(2σ²) to 1/(2fwhm²)
const GAUSS_PRE = -4 * log(2)

# isotropic
function _gaussian(px, py, x, y, fwhm, amp, theta, background)
    # find offset from center
    dx = px - x
    dy = py - y
    # rotate
    !iszero(theta) && @warn "isotropic gaussian is not affected by non-zero rotation angle $theta"
    # unnormalized gaussian likelihood
    sqmahab = dx^2 + dy^2
    sqmahab /= fwhm^2
    return amp * exp(GAUSS_PRE * sqmahab) + background
end

# bivariate
function _gaussian(px, py, x, y, fwhm::BivariateLike, amp, theta, background)
    # find offset from center
    dx = px - x
    dy = py - y
    # rotate
    if !iszero(theta)
        dx, dy = rotate_point(dx, dy, theta)
    end
    # unnormalized gaussian likelihood
    fwhmx, fwhmy = fwhm
    sqmahab = (dx / fwhmx)^2 + (dy / fwhmy)^2
    return amp * exp(GAUSS_PRE * sqmahab) + background
end


module YaWorkingData

using LinearAlgebra

export de_μ, directional_σ, ∇directional_σ,
  first_principal_component, pca, transform

const DT{T} = Vector{T} where {T <: Real}
const VDT{T} = Vector{DT{T}} where {T <: Real}

"""
Re-center to have mean 0. across all dim

```julia-repl
julia> de_μ([[1., 3., 2., 4.], [1., 3., 2., 4.]])
2-element Vector{Vector{Float64}}:
 [-1.5, 0.5, -0.5, 1.5]
 [-1.5, 0.5, -0.5, 1.5]
```
"""
function de_μ(data::VDT)::VDT
  """
  Re-center to have mean 0. across all dim
  """
  μ₀ = μ.(data)
  map(((d, μ)=t) -> d .- μ, zip(data, μ₀))
end

"""
Returns the variance of data in the direction of v
"""
function directional_σ(data::VDT, v::DT)::Float64
  dir_v = direction(v)
  sum(dot(u, dir_v) .^ 2 for u ∈ data)
end

function ∇directional_σ(data::VDT, v::DT)::DT
  """
  The gradient of directional variance with respect to v
  """
  v_dir = direction(v)
  sum(map(v -> 2 * dot(v, v_dir), data) .* data)
end

function first_principal_component(data::VDT, n::Integer=100, η=0.1)::DT
  # Start with a random guess
  guess = ones(eltype(data[1]), length(data[1]))
  for _ ∈ 1:n
    dσ = directional_σ(data, guess)
    ∇ = ∇directional_σ(data, guess)
    guess .+= ∇ * η
  end
  direction(guess)
end

function pca(data::VDT, ncomp::Integer)::VDT
  comps = VDT{eltype(data[1])}()
  for _ ∈ 1:ncomp
    comp = first_principal_component(data)
    push!(comps, comp)
    data = rm_project(data, comp)
  end
  comps
end

function transform(data::VDT, comps::VDT)::VDT
  transform_vect.(data, Ref(comps))
end

## Internals
"""
```julia-repl
julia> μ([1., 3., 2., 4.])
2.5
```
"""
function μ(data::DT)::Float64
  sum(data) / length(data)
end

function direction(v::DT)::DT
  v ./ norm(v)
end

function project(u::DT, v::DT)::DT
  """return the projection of u onto the direction v"""
  proj_len = dot(u, v)
  proj_len * v
end

function rm_project_from_vect(u::DT, v::DT)::DT
  """projects u onto v and subtracts the result from u"""
  u - project(u, v)
end

function rm_project(data::VDT, v::DT)::VDT
  rm_project_from_vect.(data, Ref(v))
end

function transform_vect(v::DT, comps::VDT)::DT
  dot.(Ref(v), comps)
end

end  # module

Data Science From Scratch/src/YaWorkingData.jl


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Statistics, DataFrames, CSV
	using Random, Distributions, Plots
	using HypothesisTests
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.06"

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	mu1, sig1, n1 = 0, 2, 8
	mu2, sig2, n2 = 0, 30, 15
	dist1 = Normal(mu1, sig1)
	dist2 = Normal(mu2, sig2)

	N = 10^4
	tdArray = Array{Tuple{Float64,Float64}}(undef,N)

	df(s1,s2,n1,n2) =
		(s1^2/n1 + s2^2/n2)^2 / ( (s1^2/n1)^2/(n1-1) + (s2^2/n2)^2/(n2-1) )
end;

# ╔═╡ fd06acdc-3b1d-11eb-1d5f-03f5cf0f595a
begin
	for i in 1:N
		x1Data = rand(dist1, n1)
		x2Data = rand(dist2, n2)
		x1Bar,x2Bar = mean(x1Data),mean(x2Data)
		s1,s2 = std(x1Data),std(x2Data)
		tStat = (x1Bar - x2Bar) / sqrt(s1^2/n1 + s2^2/n2)
		tdArray[i] = (tStat , df(s1,s2,n1,n2))
	end
	sort!(tdArray, by = first)

	invVal(v,i) = quantile(TDist(v),i/(N+1))

	xCoords  = Array{Float64}(undef,N)
	yCoords1 = Array{Float64}(undef,N)
	yCoords2 = Array{Float64}(undef,N)

	for i in 1:N
		xCoords[i] = first(tdArray[i])
		yCoords1[i] = invVal(last(tdArray[i]), i)
		yCoords2[i] = invVal(n1+n2-2, i)
	end
end;

# ╔═╡ fd06e9ae-3b1d-11eb-1646-d9c440ed4b98
begin
	scatter(xCoords, yCoords1, c=:blue, label="Calculated v", msw=0)
	scatter!(xCoords, yCoords2, c=:red, label="Fixed v", msw=0)
	plot!([-10,10], [-10,10],
		 c=:black, lw=0.3, xlims=(-8,8), ylims=(-8,8), ratio=:equal, label="",
		 xlabel="Theoretical t-distribution quantiles",
		 ylabel="Simulated t-distribution quantiles", legend=:topleft)
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.06"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═fd06acdc-3b1d-11eb-1d5f-03f5cf0f595a
# ╠═fd06e9ae-3b1d-11eb-1646-d9c440ed4b98
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"""
    Univariate optimization and root finding examples with a few different packages
"""

using Plots
#=
 You should load all packages at beginning of file!
 I just load the optimizers later to make it clear which package is being demonstrated.
=#

# Function definition
# f(x) = x[1]^2 - 0.5
f(x) = x[1]^3 - x[1]
xlim_low = -1.5
xlim_high = 1.5
plot1 = plot(f, xlim_low, xlim_high, label = "f(x)")
plot1 = plot!(x -> 0, xlim_low, xlim_high, label = "")
plot(plot1)

# Optimization
initial_value = -0.01
scatter!([initial_value],
         [0.],
         label = "Initial Value",
         markershape = :x)

## Optim
using Optim

optim_result = Optim.optimize(f, [initial_value], BFGS())
print("Optim minimum: ")
if optim_result.ls_success
    println(optim_result.minimizer)
    scatter!([optim_result.minimizer], [optim_result.minimum], label = "Optim minimum")
else
    println("Failed to converge")
end

# Root finding

## Roots
using Roots
roots_zero = find_zero(f, initial_value)
println("Roots 'find_zero' zero: $(roots_zero)")
scatter!([roots_zero],[f(roots_zero)], label = "Roots 'find_zero' zero")

## Alternative syntax
roots_fzero_zero = fzero(f, initial_value)
println("Roots 'fzero' zero: $(roots_fzero_zero)")
scatter!([roots_fzero_zero],[f(roots_zero)], label = "Roots 'fzero' zero")

## Alternative syntax
fx = ZeroProblem(f, initial_value)
roots_solve_zero = solve(fx)
println("Roots 'solve' zero: $(roots_solve_zero)")
scatter!([roots_solve_zero],[f(roots_zero)], label = "Roots 'solve' zero")

## NLsolve
using NLsolve
f!(F, x) = F[1] = f(x)
nls_zero = nlsolve(f!, [initial_value])
println("NLsolve zero: $(nls_zero.zero)")
scatter!([nls_zero.zero], [f(nls_zero.zero)], label = "NLsolve zero")

## LeastSquaresOptim
### Nonlinear least squares
### minₓ F(x)ᵀF(x)
using LeastSquaresOptim

ls_optim_result = LeastSquaresOptim.optimize(f, [initial_value], Dogleg())
print("LeastSquaresOptim minimum: ")
if ls_optim_result.converged
    println(ls_optim_result.minimizer)
    scatter!([ls_optim_result.minimizer], [f(ls_optim_result.minimizer)], label = "LeastSquaresOptim minimum")
else
    println("Failed to converge")
end


# # Custom halftoning with SDFs
# Ordered dithering & halftoning algorithms typically use predefined threshold matrices
# with unique values ranging from `1//n` to `(n-1)/n`, where `n` is the number of elements
# in the matrix.
#
# In this example, we will be generating these threshold matrices
# from signed distance functions (SDFs). A good resource for those is Inigo Quilez'
# blogpost ["2D distance functions"](https://www.iquilezles.org/www/articles/distfunctions2d/distfunctions2d.htm).
#
# Here we reimplement his function `sdStar5`:
using Images
using DitherPunk
using ColorSchemes
using TestImages

function star_sdf(x, y; r=0.5, rf=2.0)
    k1 = [0.809016994375, -0.587785252292]
    k2 = [-k1[1], k1[2]]
    p = [abs(x), y]
    p -= 2.0 * maximum([k1 ⋅ p, 0.0]) * k1
    p -= 2.0 * maximum([k2 ⋅ p, 0.0]) * k2
    p = [abs(p[1]), p[2] - r]
    ba = rf * [-k1[2], k1[1]] - [0, 1]
    h = clamp((p ⋅ ba) / (ba ⋅ ba), 0.0, r)
    return norm(p - ba * h) * sign(p[2] * ba[1] - p[1] * ba[2])
end;

# This SDF can be visualized by its contour:
#
# ![](https://user-images.githubusercontent.com/20258504/144755939-19b13e29-e50b-4744-bdeb-234f33790ca8.png)
#
# To convert an SDF to an `OrderedDither` algorithm, use this function:
function sdf2halftone(sdf, n)
    rg = range(-1, 1; length=n)
    A = [sdf(x, y) for y in rg, x in rg]
    p = sortperm(reshape(-A, :))
    B = Vector{Int}(undef, n^2)
    B[p] .= 1:(n^2)
    return OrderedDither(reshape(B, size(A)...)//(n^2 + 1))
end;

# Let's try it on a test image:
img = testimage("fabio_gray_512")
alg = sdf2halftone(star_sdf, 7)
dither(img, alg)

# We can also change the size of the matrix.
# Large halftoning matrices look great on high resolution images.
alg = sdf2halftone(star_sdf, 15)
dither(img, alg)


@defcomp co2forcing begin
    c_CO2concentration = Parameter(index=[time], unit="ppbv")
    f0_CO2baseforcing = Parameter(unit="W/m2", default=1.68)
    fslope_CO2forcingslope = Parameter(unit="W/m2", default=5.5)
    c0_baseCO2conc = Parameter(unit="ppbv", default=400859.5833)
    f_CO2forcing = Variable(index=[time], unit="W/m2")

    function run_timestep(p, v, d, t)

        # eq.13 in Hope 2006
        # the max() condition was added to prevent numerical issues due to  non-positive CO2 concentration levels which occurs in 2300 for low-emission scenarios
        # if the growth effects are paired with GDP-emission feedback
        if p.c_CO2concentration[t] > 0
            v.f_CO2forcing[t] = max(0, p.f0_CO2baseforcing + p.fslope_CO2forcingslope * log(p.c_CO2concentration[t] / p.c0_baseCO2conc) )
        else
            v.f_CO2forcing[t] = 0
        end
    end
end


# K=funLtoK(L)
# The function funLtoK(L) converts a kernel L matrix into a (normalized)
# kernel K matrix. The K matrix has to be semi-positive definite.
#
# Code available here:
# Keeler, 2018, https://github.com/hpaulkeeler/DetPoisson_Julia
#
# #TEST
# B=[3, 2, 1; 4, 5,6; 9, 8,7];
# L=B'*B
# L =
#
#   106    98    90
#    98    93    88
#    90    88    86
# K=funLtoK(L)
# K =
#
#     0.7602    0.3348   -0.0906
#     0.3348    0.3320    0.3293
#    -0.0906    0.3293    0.7492

function funLtoK(L)
    eigenVectLK=eigvecs(L); #eigen decomposition -- vectors
    eigenValL=(eigvals(L)); #eigen decomposition -- values
    eigenValK = eigenValL./(1 .+eigenValL); #eigenvalues of K
    eigenValK=Diagonal(eigenValK); #eigenvalues of L as diagonal matrix
    K=eigenVectLK*eigenValK*(eigenVectLK'); #recombine from eigen components
    K=real(K); #make sure all values are real
    return K;
end


function solve_lyapunov(A,Q,guess = I;iter_lim=100,epsilon=1e-14)
    iter = 1
    err = 1
    while iter <= iter_lim && err > epsilon
        new_guess = A*guess*A' + Q
        err = maximum(abs.(new_guess - guess))
        guess = new_guess
        iter += 1
    end
    return(guess)
end

using LinearAlgebra

function solve_lyapunov_vec(A,Q)
    aux = (I - kron(A,A))
    vec_ans = \(aux,vec(Q))
    res = reshape(vec_ans,size(Q))
    return res
end

using ToeplitzMatrices

function build_variance(g,h,Q,T) #only works for a single series
    sig_x = solve_lyapunov_vec(h,Q)
    y_autocov(t) = g*h*sig_x*h'^(t)*g'
    vec_st = map(x->first(y_autocov(x)),1:(T))
    res = Toeplitz(vec_st,vec_st)
    return res
end


function discreteIntegration(values::Array{Float64,1}, coordinates::AbstractArray{Float64,2}, l::Int, m::Int)

  @polyvar x y z

  sum = 0
  p = ylm(l,m,x,y,z)

  for k in 1:length(values)
    sum += values[k] * p((x,y,z)=>(coordinates[k,1],coordinates[k,2],coordinates[k,3]))
  end
  return sum * ((4*pi)/length(values))
end

function sphericalQuadrature(values::Array{Float64,1}, coordinates::AbstractArray{Float64,2}, L)

  C = [discreteIntegration(values,coordinates,l,m) for l=0:L for m=-l:l]
  return SphericalHarmonicCoefficients(C)
end

# Error propagation
function errorDiscreteIntegration(values::Array{Float64,1}, coordinates::AbstractArray{Float64,2}, l::Int, m::Int)

  @polyvar x y z

  sum = 0
  p = ylm(l,m,x,y,z)

  for k in 1:length(values)
    sum += (values[k] * p((x,y,z)=>(coordinates[k,1],coordinates[k,2],coordinates[k,3])))^2
  end
  return sqrt(sum) * ((4*pi)/length(values))
end

function errorSphericalQuadrature(values::Array{Float64,1}, coordinates::AbstractArray{Float64,2}, L)

  C = [errorDiscreteIntegration(values,coordinates,l,m) for l=0:L for m=-l:l]
  return SphericalHarmonicCoefficients(C)
end


# Statistical model for analyzing Radial Velocity observations
# Assumes existance of data in arrays times, obs, sigma_obs
# Assumes existance of function calc_model_rv(theta, time)

export ploglikelihood, plogprior, plogtarget
export set_times, set_obs, set_sigma_obs

# Observational data to compare model to
global times # ::Array{Float64,1}
global obs # ::Array{Float64,1}
global sigma_obs #::Array{Float64,1}

# Functions to set global data within module
function set_times(t::Array{Float64,1}) global times = t   end
function set_obs(o::Array{Float64,1})  global obs = o   end
function set_sigma_obs(so::Array{Float64,1}) global sigma_obs = so  end

function ploglikelihood(p::Vector)
  num_pl = num_planets(p)
  @assert num_pl >= 1
  if !is_valid(p) return -Inf end  # prempt model evaluation
  # Set t, o, and so to point to global arrays with observational data, while enforcing types
  t::Array{Float64,1} = times
  o::Array{Float64,1} = obs
  so::Array{Float64,1} = sigma_obs
  @assert length(times) == length(obs) == length(sigma_obs)
  jitter_sq = RvModelKeplerian.num_jitters >=1 ? extract_jitter(p)^2 : 0.
  chisq = zero(eltype(p))
  log_normalization = -0.5*length(t)*log(2pi)
  for i in 1:length(t)
    model::eltype(p) = calc_model_rv(p,t[i])
    sigma_eff_sq = so[i]^2+jitter_sq
    chisq += abs2(model-o[i])/sigma_eff_sq
    log_normalization -= 0.5*log(sigma_eff_sq)
  end
  return -0.5*(chisq)+log_normalization
end

function plogprior(p::Vector)
  num_pl = num_planets(p)
  @assert num_pl >= 1
  if !is_valid(p) return -Inf end  # prempt model evaluation
  logp = zero(eltype(p))
  logp -= 2*log(2pi)
  const max_period = 10000.0
  const max_amplitude = 10000.0
  for plid in 1:num_pl
    P::eltype(p) = extract_period(p,plid=plid)
    K::eltype(p) = extract_amplitude(p,plid=plid)
    logp += -log((1+P/P0::Float64)*log1p(max_period/P0::Float64)*
                 (1+K/K0::Float64)*log1p(max_amplitude/K0::Float64) )

  end
  if RvModelKeplerian.num_jitters >=1
     const max_jitter = 10000.0
     jitter::eltype(p) = extract_jitter(p)
     logp += -log((1+jitter/Jitter0::Float64)*log1p(max_jitter/Jitter0::Float64))
  end

  return logp::eltype(p)
end

plogtarget(p::Vector) = ploglikelihood(p) + plogprior(p)


using Distributions

# ρ(x) in integrand
ρ(x) = 1/sqrt(π)*exp(-x^2)

#=
Implementation of random walk sampler Metropolis-Hastings.
Inputs:
- n: The number of states to return
- b: The standard deviation of the normal distribution used in the proposal
     distribution
Outputs:
- Returns an array of length n denoting states S_i in Metropolis-Hastings alg.
=#
function metropolis_hastings(n, b)

    # Samples
    samples = zeros(n)
    samples[1] = rand()

    # Conditional Q distribution
    d = Normal(0,b)

    acceptances = 0

    for i=2:n
        # Propose Y
        ξ = rand(d)
        Y = samples[i-1] + ξ

        # Acceptance probability
        α = min(1, ρ(Y)/ρ(samples[i-1]))

        # Next state
        if rand() < α
            samples[i] = Y
            acceptances += 1
        else
            samples[i] = samples[i-1]
        end
    end

    println(acceptances)

    return samples
end

# Obtain 1000 samples with b = 2
S = metropolis_hastings(10000,1)

# f(x) in integrand
f(x) = -1/2 < x < 2 ? abs(cos(x)) : 0

# Estimate integral
int_estimate = sqrt(π)*mean([f(x) for x in S])


"Provide routines for making fake Gaussian models which have analytic Fourier transforms, for the purposes of testing."
module gauss

export imageGauss, FTGauss

using LinearAlgebra
using ..constants

"
    imageGauss(ll::AbstractVector{Float64}, mm::AbstractVector{Float64}, p::Vector{Float64}, k::Real; theta=0)

Given two arrays of ``l`` and ``m`` coordinates, corresponding to ``x`` and ``y``, fill an array of the Gaussian image
following the MATLAB convention for images (each row corresponds to a different y value).

`p0` is a vector of `[mu_RA, mu_DEC, sigma_x, sigma_y, rho]` in units of arcseconds.
``\\rho`` is the correlation of the Gaussian, ranging from 0 to 1.

```math
\\rho = \\frac{\\sigma_{xy}}{\\sigma_x \\sigma_y}
```

``\\mu_\\alpha`` and ``\\mu_\\delta`` are the locations of the centroid emission relative to the
image origin (RA=0, DEC=0).

``k`` is a scaling pre-factor to adjust the amplitude of the Gaussian.

The image intensity as a function of ``l`` and ``m`` (sky plane coordinates) is

```math
    I(l,m) = \\frac{k}{2 \\pi \\sqrt{|\\boldsymbol{\\Sigma}|}} \\exp \\left \\{ -\\frac{1}{2} \\boldsymbol{R}^\\mathrm{T} \\boldsymbol{\\Sigma}^{-1} \\boldsymbol{R} \\right \\}
```

Where ``\\boldsymbol{\\Sigma}`` is
```math
\\boldsymbol{\\Sigma} = \\left [ \\begin{array}{cc}
\\sigma_x & \\sigma_{xy} \\\\
\\sigma_{xy} & \\sigma_y \\\\
\\end{array} \\right ]
```

``\\theta`` is a desired angle to rotate the Gaussian about the origin. To make sense of the order of operations, the rotation must occur about the origin. A positive value means a counter-clockwise rotation of the Gaussian from North towards East.
"
function imageGauss(ll::AbstractVector{Float64}, mm::AbstractVector{Float64}, p::Vector{Float64}, k::Real; theta=0)

    # Both ll and mm increase with array index
    nx = length(ll)
    ny = length(mm)

    img = Array{Float64}(undef, ny, nx)
    mu = p[1:2] * arcsec #ll and mm shifts
    sigma_x = p[3] * arcsec
    sigma_y = p[4] * arcsec
    rho = p[5]
    sigma_xy = rho * sigma_x * sigma_y
    Sigma = Float64[[sigma_x^2, sigma_xy] [sigma_xy, sigma_y^2]]
    rot = [[cosd(theta), +sind(theta)] [-sind(theta), cosd(theta)]]
    # Sigma = Diagonal((p[3:4] * arcsec).^2) #Convert from arcsec to radians
    pre = 1.0 / (2pi * sqrt(det(Sigma))) * k
    for j=1:ny
        for i=1:nx
            # Rotation of Gaussian occurs about origin, so subtract off mean first.
            R = rot * (Float64[ll[i] , mm[j]] - mu)
            img[j, i] = pre * exp(-0.5 * (R' * (Sigma\R))[1]) # backslash solves for Sigma^{-1}R
        end
    end
    return img
end

"
    FTGauss(uu::Float64, vv::Float64, p::Vector{Float64}, k::Real)

Given u and v coordinates in [kλ], evaluate the analytic FT of the
aforementioned Gaussian.

`p` is a length 5 vector of `[mu_RA, mu_DEC, sigma_x, sigma_y, rho]` in units of arcseconds, corresponding to the **image** plane.

N.B. Here Sigma refers to the (same) covariance matrix in the *image* domain.

This function always returns a complex value."
function FTGauss(uu::Float64, vv::Float64, p::Vector{Float64}, k::Real; theta=0)
    uu = uu .* 1e3 #[λ]
    vv = vv .* 1e3 #[λ]

    rot = Float64[[cosd(theta), -sind(theta)] [sind(theta), cosd(theta)]]

    mu = p[1:2] * arcsec #ll and mm shifts
    sigma_x = p[3] * arcsec
    sigma_y = p[4] * arcsec
    rho = p[5]
    sigma_xy = rho * sigma_x * sigma_y
    Sigma = Float64[[sigma_x^2, sigma_xy] [sigma_xy, sigma_y^2]]

    # mu_RA, mu_DEC = p[1:2]
    # mu = Float64[mu_RA, mu_DEC] * arcsec #ll and mm shifts

    R0 = Float64[uu, vv]

    # Position to evaluate the rotated Gaussian centered at the origin. Evaluating the Gaussian as if we were
    # in uu', vv'.
    R = rot * R0
    model = k * exp(-2 * (pi^2) * (R' * Sigma * R)[1])
    # in this case, Sigma serves as an inverse matrix, since the covariance matrix corresponds to the *image* plane.
    # Sigma = Diagonal((p[3:4] * arcsec).^2) #Convert from arcsec to radians

    # Phase shift should be evaluated using the original coordinates and mu, since this is the direction that we are actually moving things. And, the interferometer samples uu,vv, not uu', vv'.
    phase_shift = exp(-2pi * 1.0im * (R0' * mu)[1]) # Not actually in polar phase form

    return model * phase_shift
end

"
    FTGauss(uu::AbstractVector{Float64}, vv::AbstractVector{Float64}, p::Vector{Float64}, k::Int)

`p` is a length 5 vector of `mu_RA`, `mu_DEC`, `sigma_RA`, `sigma_DEC`, `rho`.

Given two arrays of u and v coordinates in [kλ], fill an array with the
analytic FT of `imageGauss` evaluated at every pairwise (u,v) pair."
function FTGauss(uu::AbstractVector{Float64}, vv::AbstractVector{Float64}, p::Vector{Float64}, k::Real; theta=0)
    nu = length(uu)
    nv = length(vv)
    # Both uu and vv increase with array index
    img = Array{ComplexF64}(undef, nv, nu)
    for j=1:nv
        for i=1:nu
            img[j, i] = FTGauss(uu[i], vv[j], p, k, theta=theta)
        end
    end
    return img
end

end # module


type SimpleOG
    B :: Int64
    M :: Int64
    alpha :: Float64
    beta :: Float64
    R :: Array{Float64}
    Q :: Array{Float64}
end

function SimpleOG(;B=10, M=5, alpha=0.5, beta=0.9)

    u(c) = c^alpha
    n = B + M + 1
    m = M + 1

    R = Array(Float64,n,m)
    Q = zeros(Float64,n,m,n)

    for a in 0:M
        Q[:, a + 1, (a:(a + B)) + 1] = 1 / (B + 1)
        for s in 0:(B + M)
            R[s + 1, a + 1] = a<=s ? u(s - a) : -Inf
        end
    end

    return SimpleOG(B, M, alpha, beta, R, Q)
end


import Distributions

function laplace_mechanism(result::Number; sensitivity::Number, epsilon::Number)
    @assert sensitivity > 0
    @assert epsilon > 0
    lap = Distributions.Laplace(0.0, sensitivity/epsilon) # mean 0, scale=sensitivity/epsilon
    result + rand(lap)
end

function laplace_mechanism(result::Array; sensitivity::Number, epsilon::Number)
    [laplace_mechanism(x, sensitivity=sensitivity, epsilon=epsilon) for x in result]
end

function geometric_mechanism(result::Int; sensitivity::Int, epsilon::Number)
    @assert sensitivity > 0
    @assert epsilon > 0
    p = 1-exp(-epsilon/sensitivity)
    geo = Distributions.Geometric(p)
    result + rand(geo) - rand(geo)
end

function noisy_max(answers::Vector{S}; sensitivity::Number, epsilon::Number) where S <: Number
    noisy_answers = [laplace_mechanism(x, sensitivity=2*sensitivity, epsilon=epsilon) for x in answers]
    return argmax(noisy_answers)
end

function sparse_vector(;T::Number, N::Int, answers::Vector{S}, sensitivity::Number, epsilon::Number) where S <: Number
    noisy_T = laplace_mechanism(T; sensitivity=sensitivity, epsilon=epsilon/2)
    answered = 0
    i=1
    output = Bool[] # empty array of booleans
    while(answered < N && i <= length(answers))
        noisy_query = laplace_mechanism(answers[i]; sensitivity=2*N, epsilon=epsilon/2)
	    i = i + 1
	    if noisy_query >= noisy_T
	       append!(output, true)
	       answered = answered + 1
	    else
           append!(output, false)
	    end
    end
    output
end


## PART 7

include("autocorrelation.jl");

println("\nPart 7");
println("\n\tPick subjets v1 and v2 of each dataset,");
println("\tsizes N1 = 1,000 & N2 = 10,000.");

v1 = v[[1:N1],:];
v2 = v[[1:N2],:];

println("\n\tEstimate autocorrelations");
println("\n\t...for N1...");
xcorr = [0:10:400];
ycorr1 = zeros( length( xcorr ), 5 );
for a in [1:5]
  # Calculate autocorrelations
  ycorr1[:,a] = autocorrelation(N1, v1, a, xcorr);
  # Normalize by dividing through by Ĉ[1,a]
  ycorr1[:,a] /= ycorr1[1,a];
end
println("\n\t...for N2 (See plots below)");
ycorr2 = zeros( length( xcorr ), 5 );
for a in [1:5]
  # Calculate autocorrelations
  ycorr2[:,a] = autocorrelation(N2, v2, a, xcorr);
  # Normalize by dividing through by Ĉ[1,a]
  ycorr2[:,a] /= ycorr2[1,a];
end

# Make plots for each alpha
println("\n\tPlacing estimated autocorrelation plots 'plots1' and 'plots2':");
using Gadfly
plots1 = Plot[];
plots2 = Plot[];
for i in [1:5]
  # Plot
  push!(plots1, Gadfly.plot ( x = xcorr, y = ycorr1[:,i], Geom.line, Geom.point,
      Guide.xlabel("Separation"), Guide.ylabel("Normalized Autocorrelation"),
      Guide.title(join(["Estimated Autocorrelation for Dataset",string(i),
                        "with Sample Size N1 = $N1"]," "," "))))
  push!(plots2, Gadfly.plot ( x = xcorr, y = ycorr2[:,i], Geom.line, Geom.point,
      Guide.xlabel("Separation"), Guide.ylabel("Normalized Autocorrelation"),
      Guide.title(join(["Estimated Autocorrelation for Dataset",string(i),
                        "with Sample Size N2 = $N2"]," "," "))))
end

# Estimate autocorrelation times
println("\n\tEstimating autocorrelation times τ1 and τ2");
τ1 = Float64[];
τ2 = Float64[];
for a in [1:5]
  corra = autocorrelation(N1, v1, a, ns);
  corra /= corra[1];                          # Normalize
  push!(τ1,sum(corra) - 0.5);                 # Get autocorrelation time
end
for a in [1:5]
  corra = autocorrelation(N2, v2, a, ns);
  corra /= corra[1];                          # Normalize
  push!(τ2,sum(corra) - 0.5);                 # Get autocorrelation time
end

println("\n\tEstimated Autocorrelation Times for N1 = $N1:\n");
for i in [1:5]
  println("\t\tτ̂1v$i:\t$(τ1[i])");
end
println("\n\tEstimated Autocorrelation Times for N2 = $N2:\n");
for i in [1:5]
  println("\t\tτ̂2v$i:\t$(τ2[i])");
end
println("\n\tThe estimated autocorrelation times are clearly far better for the larger sample.");

# Estimate standard deviation
σ1 = Float64[];
σ2 = Float64[];
println("\n\tCreating estimators σ1 for the standard deviation for N1 = $N1");
for a in [1:5]
  diff1 = v1[:,a] .- mean(v1[:,a]);
  σ1a = sqrt( dot( diff1, diff1 ) / (N1-1) );
  println("\t\tσ̂1[$a] = $σ1a");
  push!(σ1, σ1a);
end
println("\n\tCreating estimators σ2 for the standard deviation for N2 = $N2");
for a in [1:5]
  diff2 = v2[:,a] .- mean(v2[:,a]);
  σ2a = sqrt( dot( diff2, diff2 ) / (N2-1) );
  println("\t\tσ̂2[$a] = $σ2a");
  push!(σ2, σ2a);
end

# Estimate standard deviation of the sample mean
σN1 = sqrt( 2 / N1 .* τ1 ) .* σ1;
σN2 = sqrt( 2 / N2 .* τ2 ) .* σ2;
println("\n\tEstimating standard deviations of the sample mean for N1 = $N1:");
for i in [1:5]
  println("\n\t\tσ̂N1[$i] = $(σN1[i])");
end
println("\n\tEstimating standard deviations of the sample mean for N2 = $N2:");
for i in [1:5]
  println("\n\t\tσ̂N2[$i] = $(σN2[i])");
end

# Compare to the actual measured values of the standard deviations of the sample mean
sigmean1 = σ̂1./σN1;
sigmean2 = σ̂2./σN2;
println("\n\tTake ratio of measured standard deviation of sample means over the estimated");
println("\tstandard deviation just found in order to measure the accuracy of the estimate.");
println("\n\tFind the ratio for N1 = $N1:")
for i in [1:5]
  println("\n\t\tσ̂1[$i] / σ̂N1[$i] = $(sigmean1[i])");
end
println("\n\tFind the ratio for N2 = $N2:")
for i in [1:5]
  println("\n\t\tσ̂2[$i] / σ̂N2[$i] = $(sigmean2[i])");
end

println("\n\tFor both sample sizes, these estimates are close to the true values of");
println("\tthe standard deviation of the sample means, though the agreement is far");
println("\tbetter for the larger sample size.");

# Find sample means
v̄1 = zeros(5);
v̄2 = zeros(5);
for i in [1:5]
  v̄1[i] = mean(v1[:,i]);
  v̄2[i] = mean(v2[:,i]);
end

# Finally, compare the normalized covariance matrices ρ1 and ρ2 to the true value ρ̂.
println("\n\tFinally, estimate the normalized covariance matrices ρ1 and ρ2.\n");
println("\tEstimating the covariance matrix c1 for sample size N1 = $N1");
D1 = v1 - ones(N1) * transpose(v̄1);
c1 = transpose(D1) * D1 / N1;
println("\tEstimating the covariance matrix c2 for sample size N2 = $N2");
D2 = v2 - ones(N2) * transpose(v̄2);
c2 = transpose(D2) * D2 / N2;

println("\tEstimating the normalized covariant matrix ρ1 for N1 = $N1");
ρ1 = c1 ./ ( σ1 * transpose(σ1) );
println("\tEstimating the normalized covariant matrix ρ2 for N2 = $N2");
ρ2 = c2 ./ ( σ2 * transpose(σ2) );


#
# Computes the density function of a list of values
#
struct Density
  x :: Vector{Float64}
  d :: Vector{Float64}
  step :: Float64
  norm :: Int64
end

"""
```
fitdensity(x; step, norm)
```

Obtains the density function given data sampled.

Use `step=(Float64)` to control the bin step. Use `norm=(0 or 1)` to set
if the number of data points or the probability of finding a data point
within ± step/2 will be output.

By default, `norm=1` (probability) and the step is `(xmax-xmin)/100`.

# Examples
```jldoctest
julia> x = randn(1000)

julia> d = fitdensity(x)

 ------------------- Density -------------

  `d` contains the number of data points within x ± 0.028325979340904375

 -----------------------------------------

```
"""
function fitdensity(v;nbins=nothing,
                      step=nothing, steptype="absolute",
                      vmin=nothing, vmax=nothing,
                      norm :: Int64 = 1)

  ndata = length(v)

  if vmin == nothing
    vmin = minimum(v)
  end
  if vmax == nothing
    vmax = maximum(v)
  end
  if nbins == nothing
    nbins = 100
  end
  if step == nothing
    step = (vmax - vmin)/nbins
  else
    # By default, the step size is absolute
    if steptype == "relative"
      step = step*(vmax-vmin)/nbins
    elseif steptype != "absolute"
      error(" steptype must be \"relative\" or \"absolute\"")
    end
  end

  x = Vector{Float64}(undef,nbins)
  df = Vector{Float64}(undef,nbins)

  binstep = (vmax - vmin)/nbins
  for i in 1:nbins
    x[i] = vmin + (i-1)*binstep + binstep/2
    nv = 0
    for j in 1:ndata
      if ( v[j] > x[i] - step/2 ) && ( v[j] <= x[i] + step/2 )
        nv = nv + 1
      end
    end
    binsize = min(vmax,x[i]+step/2) - max(vmin,x[i]-step/2)
    if norm == 0
      df[i] = nv
    elseif norm == 1
      df[i] = nv/(binsize*ndata)
    end
  end

  return Density(x,df,step,norm)
end
export fitdensity

function Base.show( io :: IO, d :: Density )
  if d.norm == 0
    s = "number of"
  end
  if d.norm == 1
    s = "probability of finding"
  end
  println("
 ------------------- Density -------------

  d contains the $s data points within x ± $(round(d.step/2,sigdigits=3))

 ----------------------------------------- ")
end


using Plots, Printf

"""
    shallow_water_1D(;nx=1024, do_visu=false)

1D shallow water equations solver for an instantaneous dam break.
The Lax-Friedrichs Method was applied to the continuity equation.
Geometry (length of 40 meters) and initial conditions (half of domain
have initial water level of 5 meters, other half is dry) match
BASEMENT version 2.8 Test Case H_1 "Dam break in a closed channel

# Arguments
    - `nx`: number of discretised cells.
    - `do_visu`: if true, each physical time step will be ploted.

# Return values
    - `H`: The solution arrray (Water surface height in m).
"""
@views function shallow_water_1D(;
    # Numerics
    nx = 1024,
    # Visualisation
    do_visu = false)
    # Physics
    Lx = 40.0
    g = 9.81
    u_max = 6 #from review of results with very small timesteps
    ttot = 20.0
    # Numerics
    nout = 100
    # Derived numerics
    dx = Lx / nx
    dt = dx / u_max / 10 #stability condition of Lax-Friedrichs Method, further divided by 10 for stability of initially discontinuous wave front
    #println(dt)
    nt = cld(ttot, dt)
    xc = LinRange(dx / 2, Lx - dx / 2, nx)
    # Array initialisation
    H = zeros(Float64, nx)
    H[1:round(Int64(nx / 2))] = 5.0 .+ zeros(Float64, round(Int64(nx / 2)))
    u = zeros(Float64, nx)
    dHdx = zeros(Float64, nx - 1)
    dudx = zeros(Float64, nx - 1)
    dudt = zeros(Float64, nx - 1)
    t_tic = 0.0
    niter = 0
    # Time loop
    for it = 0:nt
        if (it == 11)
            t_tic = Base.time()
            niter = 0
        end
        dHdx .= diff(H) ./ dx
        dudx .= diff(u) ./ dx
        dudt .= -1 / 2 .* (u[1:end-1] .+ u[2:end]) .* dudx .- (g .* dHdx) #momentum
        H[2:end-1] .= 1 / 2 .* (H[3:end] .+ H[1:end-2]) .- dt / 2 / dx .* ((H[3:end] .* u[3:end]) .- (H[1:end-2] .* u[1:end-2])) #continuity with Lax-Friedrichs Method
        H[1] = H[2]
        H[end] = H[end-1]
        u[2:end-1] .= u[2:end-1] .+ 1 / 2 .* (dudt[1:end-1] .+ dudt[2:end]) .* dt
        u[1] = -u[2]
        u[end] = -u[end-1]

        niter += 1
        if do_visu && (it % nout == 0)
            p1 = plot(xc, H, xlims = (xc[1], xc[end]), ylims = (0, 10),
                xlabel = "Lx (m)", ylabel = "water surface elevation (m)", label = "h",
                title = "time = $(round(it*dt, sigdigits=3)) s, stability: $(round(maximum(abs.(u))*dt/dx,sigdigits=3)) /1",
                linewidth = :1.0, framestyle = :box)
            plot!(xc, u, label = "u", linewidth = :1.0)
            display(p1)
        end
    end
    t_toc = Base.time() - t_tic
    A_eff = (2 * 2) / 1e9 * nx * sizeof(Float64)  # Effective main memory access per iteration [GB]
    t_it = t_toc / niter                      # Execution time per iteration [s]
    T_eff = A_eff / t_it                       # Effective memory throughput [GB/s]
    @printf("Time = %1.3f sec, T_eff = %1.2f GB/s (niter = %d)\n", t_toc, round(T_eff, sigdigits = 3), niter)
    return H
end

# shallow_water_1D(;do_visu=true)


export srDistance

"""
function srDistance

Unnormalized distance (reconstruction error) between two pairs (M,w):

min_P ||M1*diagm(w1) - M2*diagm(w2)*P ||_p

where P is a permutation matrix and || ||_p denotes the entrywise matrix pn-norm
(i.e., pn=2 corresponds to the Frobenius norm).

Before computing the error, the first dimensions of the binary matrices are
augmented from m*(p-1) to m*p so that the column sums of each block become 1.
This can be disabled by setting p=1.

If n has different values among the inputs, then the largest n is chosen and
the rest are augmented with zeros to make the dimensions compatible.

Input:
    M1 - binary array of size (m*(p-1),n) or (m*(p-1),n,N)
    w1 - frequency vector of length n or an array of size (n,N)
    M2 - as M1
    w2 - as w1

Keyword arguments:
    p - number of categories (default=2)
    pn - exponent in the matrix norm (default=2.0)

Output:
    distance(s) as a scalar or an N-vector
"""

function srDistance(M1::BitArray{2},w1::Vector,M2::BitArray{2},w2::Vector;p::Int64=2,pn::Float64=2.0)
    n = maximum((size(M1,2),length(w1),size(M2,2),length(w2)))
    M1 = [M1 falses(size(M1,1),n-size(M1,2))]
    w1 = [w1; zeros(n-length(w1))]
    M2 = [M2 falses(size(M2,1),n-size(M2,2))]
    w2 = [w2; zeros(n-length(w2))]

    perms = permutations(1:n)

    if p>1
        M1 = fullCategoricalMatrix(M1,p)
        M2 = fullCategoricalMatrix(M2,p)
    end

    distance = Inf
    for perm in perms
        M2perm = M2[:,perm]
        w2perm = w2[perm]

        mat = M1*diagm(w1) - M2perm*diagm(w2perm)
        dist = (sum(abs.(mat[:]).^pn))^(1/pn)
        if dist < distance
            distance = dist
        end
    end

    return distance
end

function srDistance(M1::BitArray{3},w1::Array{Float64,2},M2::BitArray{3},w2::Array{Float64,2};p::Int64=2,pn::Float64=2.0)
    N = size(M1,3)
    distances = SharedArray{Float64}(N)

    @sync @parallel for idx = 1:N
        distances[idx] = srDistance(M1[:,:,idx],w1[:,idx],M2[:,:,idx],w2[:,idx],p=p,pn=pn)
    end

    return distances
end

function fullCategoricalMatrix(M::BitArray{2},p::Int64)
    n = size(M,2)

    M1 = reshape(M,(p-1,:))
    blockcolsums = sum(M1,1)
    Mfull = [1-blockcolsums; M1]

    return reshape(Mfull,(:,n))
end


using Distributions, Random, Statistics, Plots; pyplot()
Random.seed!(2)

n, N, alpha = 10, 10^7, 0.05
mActual = 0.75
dist0, dist1 = Uniform(0,1), Uniform(0,mActual)

ts(sample) = maximum(sample) - minimum(sample)

empiricalDistUnderH0 = [ts(rand(dist0,n)) for _ in 1:N]
rejectionValue = quantile(empiricalDistUnderH0,alpha)

sample = rand(dist1,n)
testStat = ts(sample)
pValue = sum(empiricalDistUnderH0 .<= testStat)/N

if testStat > rejectionValue
    print("Didn't reject: ", round(testStat,digits=4))
    print(" > ", round(rejectionValue,digits=4))
else
    print("Reject: ", round(testStat,digits=4))
    print(" <= ", round(rejectionValue,digits=4))
end
println("\np-value = $(round(pValue,digits=4))")

stephist(empiricalDistUnderH0, bins=100, c=:blue, normed=true, label="")
plot!([testStat, testStat], [0,4], c=:red, label="Observed test statistic")
plot!([rejectionValue, rejectionValue], [0,4], c=:black, ls=:dash,
	label="Critical value boundary", legend=:topleft, ylims=(0,4),
    	xlabel = "x", ylabel = "Density")


#= convex.jl
Philip Barrett, pobarrett@gmail.com
27may2016, Chicago

Calculates the convex hull of a set of points
=#

"""
    acw( p1, p2, p3 )
Positive if p1 -> p2 -> p3 is anti-clockwise
"""
function acw( p1, p2, p3 )
  return (p2[1] - p1[1]) * (p3[2] - p1[2]) - (p2[2] - p1[2]) * (p3[1] - p1[1])
end

"""
    grahamScan( pts::Matrix )
Computes a convex hull using the Graham Scan algortithm NB: Returns a polygon object even when a matrix is submitted
"""

function grahamScan( pts::Matrix )

  N = size(pts)[1]
      # Number of points
  lex = sortrows( pts[ :, [2, 1 ] ] )[ :, [ 2, 1 ] ]
      # Sorts the points lexicographically by y then x
  P = lex[ 1, : ]
      # The initial point
  otherpts = lex[ 2:end, : ]
      # The other points
  orderpts = [ P; otherpts ; P ]
      # Initialize the anti-clockwise ordered points
  cah = zeros( N - 1 )
      # The cosine of the angle between each point and P.  Will be in
      # [-1,1] because of choice of P

  ## Order the points anti-clockwise ##
  for( i in 1:(N-1) )
    diff = otherpts[i,:] - P
    cah[i] = diff[1] / norm(diff)
        # Cosine is adjacent over hypotenuse
  end
  orderpts[2:N,:] = otherpts[ sortperm(cah, rev=true), : ]
      # Order the points by the angle measure

  ## Create the output ##
  out = similar( pts )
  out[ 1:2, : ] = orderpts[1:2,:]
      # Initialize
  M = 2
      # Counts number of rows in convex hull
  for( i in 3:N )
    while( acw( out[ M-1, : ], out[ M, : ], orderpts[ i, : ] ) < 0 )
      if( M > 2 )
        M -= 1
      elseif( i == N )
        break
      else
        i += 1
      end
    end
    M += 1
        # Increment counter
    out[ M, : ] = orderpts[ i, : ]
        # Add to the output if we have an acw angle
  end
  return Polygon( pts=out[ 1:M, : ] )
end

function grahamScan( poly::Polygon )
    return grahamScan( poly.pts )
end

"""
    convexhull( poly::Polygon )
Computes a convex hull using the Graham Scan algortithm from the CHull2D package
"""
function chull( poly::Polygon )
  return Polygon( pts = chull(poly.pts) )
end

function chull( pts::Matrix )
  N = size(pts)[1]
      # Number of points
  oldpts = Point{2, Float64}[Point(pts[i, 1], pts[i, 2])
                            for i=1:N]
      # Convert poly.pts to an array of points
  ch = CHull2d.convexhull(oldpts, _at=false)
      # The convex hull object
  N_newpts = length(ch.extremepoints)
  newpts = Array(Float64, N_newpts, 2)
      # Create the container for the new points
  for i=1:N_newpts
      curr_pt = ch.extremepoints[i]
      newpts[i, 1] = curr_pt[1]
      newpts[i, 2] = curr_pt[2]
  end
  return newpts
end


function ABCDQR_simula(nt::Int,seed::Int)
	#
	# simulated model to test the functions
	#
	# nt: number of time instants
	# seed: seed for the random numbers

	# random seed
	# ------------------------------
	Random.seed!(seed)

	# simulated system
	m = simula01()

	# vibration parameters
	# ----------------------------
	# time interval (s)
	dt = m["dt"]

	# state-space matrices
	# --------------------------------------
	A = m["A"]
	B = m["B"]
	C = m["C"]
	D = m["D"]

	# input definition
	# ---------------------------------------
	# system input
	u = 100*randn(2,nt) # if 100 is not used, Q and R are too little

	# state equation, no noise
	# ---------------------------------
	x = zeros(4,nt)
	for t=1:nt-1
		x[:,t+1] = A*x[:,t] + B*u[:,t]
	end
	P = cov(x')
	Q = 0.01*P # std approx 0.1
	ch = cholesky(Q)
	w = ch.L*randn(4,nt)

	# state equation with noise
	# ------------------------------------------
	x = zeros(4,nt)
	for t = 1:nt-1
		x[:,t+1] = A*x[:,t] + B*u[:,t] + w[:,t]
	end

	# observation equation
	# -------------------------------------
	y = C*x + D*u
	L = cov(y')
	R = 0.01*L # std approx 0.1
	ch = cholesky(R)
	v = ch.L*randn(2,nt)
	y = y + v

	#
	S = cov(w',v')

	# starting values
	m1 = C\y[:,1:1]
	P1 = 0*A

	m = Dict("dt"=>dt,
	"y"=>y,"u"=>u,
	"A"=>A,"B"=>B,"C"=>C,"D"=>D,
	"Q"=>Q,"R"=>R,"S"=>S,
	"m1"=>m1,"P1"=>P1)

	return m

end


"""Calculate the Fermi Function given some energy and some temperature.

Parameters
----------
en : float
    Energy relative to the fermi energy (E-Ef).

kbt : Temperature in same units as Energy (using kB).

Returns
-------
ffun : float
    The Fermi Function at en and temp."""
function fermi_fun(en, kbt)
    if en == 0
        ffun = 0.5
    elseif kbt == 0
        if en < 0
            ffun = 1.0
        elseif en > 0
            ffun = 0.0
        else
            throw(DomainError("Energy must be a real number."))
        end
    elseif kbt > 0
        #Using a tanh here instead of the more traditional formulation because
        #the function is better behaved for very large or small values of en/kbt
        ffun = 0.5*(1-tanh(0.5*en/kbt))

    elseif kbt < 0
        throw(DomainError("Temperature must be >= 0."))
    end

    return ffun
end


function add_mass_matrix!{T<:Number}(bas::GenBasis1d, el::Element1d, mass::AbstractMatrix{T}, λ=one(T))
    M = nmodes(bas)
    Q = nquad(bas)
    ϕ = qbasis(bas)
    wJ = jacweights(el)
    for k = 1:M
        for i = k:M
            m = 0.0
            for q = 1:Q
                m += ϕ[q,k] * ϕ[q, i] * wJ[q]
            end
            mass[k,i] += m*λ
            if k != i
                mass[i,k] += m*λ
            end
        end
    end
    return mass
end

function add_mass_matrix!{T<:Number}(bas::SEM1d{T}, el::Element1d, mass::AbstractMatrix{T}, λ=one(T))

    M = nmodes(bas)
    wJ = jacweights(el)

    for i = 1:M
        mass[i,i] += wJ[i]*λ
    end
end

mass_matrix{T<:Number}(bas::GenBasis1d, el::Element1d{T}, λ=one(T)) =
    add_mass_matrix!(bas, el, zeros(T, nmodes(bas), nmodes(bas)), λ)

function add_mass_matrix!{T<:Number}(bas::SEM1d{T}, el::Element1d, mass::AbstractMatrix{T}, λ::AbstractVector{T})

    M = nmodes(bas)
    wJ = jacweights(el)

    for i = 1:M
        mass[i,i] += wJ[i]*λ[i]
    end
end

function add_mass_matrix!{T<:Number}(bas::GenBasis1d, el::Element1d, mass::AbstractMatrix{T}, λ::AbstractVector{T})
    M = nmodes(bas)
    Q = nquad(bas)
    ϕ = qbasis(bas)
    wJ = jacweights(el)
    for k = 1:M
        for i = k:M
            m = zero(T)
            for q = 1:Q
                m += ϕ[q,k] * ϕ[q, i] * wJ[q] * λ[q]
            end
            mass[k,i] += m
            if k != i
                mass[i,k] += m
            end
        end
    end
    return mass
end

function add_stiff_matrix!{T<:Number}(bas::GenBasis1d, el::Element1d, mat::AbstractMatrix{T}, λ=one(T))
    M = nmodes(bas)
    Q = nquad(bas)
    dϕ = dqbasis(bas)
    D = diffmat(bas)
    wJ = jacweights(el)
    dξdx = deriv_ξ(el)
    for k = 1:M
        for i = k:M
            L = zero(T)
            for q = 1:Q
                L += dϕ[q,i] * dϕ[q,k] *wJ[q] * (dξdx[q]^2)*λ
            end
            mat[i,k] += L
            if k != i
                mat[k,i] += L
            end
        end
    end
    return mat
end

stiff_matrix{T<:Number}(bas::GenBasis1d, el::Element1d{T}, λ=one(T)) =
    add_stiff_matrix!(bas, el, zeros(T, nmodes(bas), nmodes(bas)), λ)

function add_stiff_matrix!{T<:Number}(bas::GenBasis1d, el::Element1d, mat::AbstractMatrix{T},
                                      λ::AbstractVector{T})
    M = nmodes(bas)
    Q = nquad(bas)
    dϕ = dqbasis(bas)
    D = diffmat(bas)
    wJ = jacweights(el)
    dξdx = deriv_ξ(el)

    for k = 1:M
        for i = k:M
            L = zero(T)
            for q = 1:Q
                L += dϕ[q,i] * dϕ[q,k] *wJ[q] * (dξdx[q]^2)*λ[q]
            end
            mat[i,k] += L
            if k != i
                mat[k,i] += L
            end
        end
    end

    return mat
end

# RHS

function add_rhs!{T<:Number}(bas::GenBasis1d{T}, el::Element1d{T}, f::AbstractVector{T},
                            Fe::AbstractVector{T})

    wJ = jacweights(el)
    ϕ = qbasis(bas)
    M = nmodes(bas)
    Q = nquad(bas)

    for k = 1:M
        F = zero(T)
        for q = 1:Q
            F += f[q] * ϕ[q,k] * wJ[q]
        end
        Fe[k] += F
    end
    return Fe
end

function add_rhs!{T<:Number}(bas::SEM1d{T}, el::Element1d{T}, f::AbstractVector{T},
                            Fe::AbstractVector{T})

    wJ = jacweights(el)
    M = nmodes(bas)

    for k = 1:M
        Fe[k] = f[k] * wJ[k]
    end
    return Fe
end


## Multivariate linear Gaussian model

module MVLinearGaussian

using SequentialMonteCarlo
using RNGPool
using StaticArrays
import SMCExamples.Particles.MVFloat64Particle
using LinearAlgebra
using Random

struct MVLGTheta{d}
  A::SMatrix{d, d, Float64}
  Q::SMatrix{d, d, Float64}
  C::SMatrix{d, d, Float64}
  R::SMatrix{d, d, Float64}
  x0::SVector{d, Float64}
  v0::SVector{d, Float64}
end

## scratch space for computations
struct MVLGPScratch{d}
  t1::MVector{d, Float64}
  t2::MVector{d, Float64}
end
MVLGPScratch{d}() where d = MVLGPScratch{d}(MVector{d, Float64}(undef),
  MVector{d, Float64}(undef))

function makeMVLGModel(theta::MVLGTheta, ys::Vector{SVector{d, Float64}}) where
  d
  n = length(ys)
  cholQ = SMatrix{d, d, Float64}(cholesky(theta.Q).L)
  invRover2 = 0.5 * inv(theta.R)
  sqrtv0 = sqrt.(theta.v0)
  logncG = - 0.5 * d * log(2 * π) - 0.5 * logdet(theta.R)
  @inline function lG(p::Int64, particle::MVFloat64Particle{d},
    scratch::MVLGPScratch{d})
    mul!(scratch.t1, theta.C, particle.x)
    @inbounds scratch.t2 .= scratch.t1 .- ys[p]
    mul!(scratch.t1, invRover2, scratch.t2)
    return logncG - dot(scratch.t1,scratch.t2)
  end
  @inline function M!(newParticle::MVFloat64Particle{d}, rng::RNG, p::Int64,
    particle::MVFloat64Particle{d}, scratch::MVLGPScratch{d})
    if p == 1
      randn!(rng, scratch.t1)
      scratch.t2 .= sqrtv0 .* scratch.t1
      newParticle.x .= theta.x0 .+ scratch.t2
    else
      randn!(rng, scratch.t1)
      mul!(scratch.t2, cholQ, scratch.t1)
      mul!(scratch.t1, theta.A, particle.x)
      newParticle.x .= scratch.t1 .+ scratch.t2
    end
  end
  return SMCModel(M!, lG, length(ys), MVFloat64Particle{d}, MVLGPScratch{d})
end

function simulateMVLGModel(theta::MVLGTheta{d}, n::Int64) where d
  model = makeMVLGModel(theta, Vector{SVector{d, Float64}}(undef, 0))
  ys = Vector{SVector{d, Float64}}(undef, n)
  xParticle = MVFloat64Particle{d}()
  xScratch = MVLGPScratch{d}()
  cholR = cholesky(theta.R).L
  rng = getRNG()
  for p in 1:n
    model.M!(xParticle, rng, p, xParticle, xScratch)
    ys[p] = theta.C*(xParticle.x) + cholR * randn(rng,d)
  end
  return ys
end

function defaultMVLGModel(d::Int64, n::Int64)
  function toeplitz(d::Int64, a::Float64, C::Float64)
    M = Matrix{Float64}(undef, d, d)
    for i = 1:d
      for j = 1:d
        M[i,j] = C * a^abs(i-j)
      end
    end
    return SMatrix{d, d, Float64}(M)
  end

  tA = SMatrix{d, d, Float64}(0.9 * Matrix{Float64}(I, d, d))
  tC = toeplitz(d, 0.5, 1.2)
  tQ = toeplitz(d, 0.2, 0.6)
  tR = toeplitz(d, 0.3, 1.5)
  tx0 = SVector{d,Float64}(range(1, step=1, length=d))
  tv0 = SVector{d,Float64}(range(2, step=1, length=d))
  theta = MVLGTheta(tA, tQ, tC, tR, tx0, tv0)
  ys = simulateMVLGModel(theta, n)

  ko = kalmanMV(theta, ys)

  mvlgModel = makeMVLGModel(theta, ys)

  return mvlgModel, theta, ys, ko
end

include("mvlgKalman.jl")

end


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

# %%
using GaloisFields, LinearAlgebra
const F5 = @GaloisField 5
A = F5[1 2; 0 3]

# %%
det(A)

# %%
B = F5[1 2; 3 4]

# %%
det(B)

# %%
F5(3) |> dump

# %%
F5 |> dump

# %%
Base.abs(x::GaloisFields.PrimeField) = x.n # type-piracy!
det(B)

# %%
F3 = @GaloisField 3
[det(F3.(rand(0:2, 3, 3))) for _ in 1:20]

# %%
Base.abs(x::GaloisFields.ExtensionField) = x.coeffs # type-piracy!
Base.isless(x::GaloisFields.PrimeField, y::GaloisFields.PrimeField) = isless(x.n, y.n) # type-piracy!
F9 = @GaloisField! 3 a^2 + 1
C = F9[1+a 2+a; 2+a 2+a]

# %%
det(C)

# %%


#########################################################################################################################################
############################## quadrature.jl contains basic univariate quadrature routines ##############################################
#########################################################################################################################################
#
# Types:
#       QuadRule               - object that stores n quadrature points and weights on the interval [-1,1]
#
# Functions:
#       dim()                  - determine the order of the quadrature rule
#
#
#########################################################################################################################################
#
# File part of the Julia IsoGeometric Analysis toolbox
# R.R.Hiemstra
# 19-04-2016
#
#########################################################################################################################################

"container for quadrature formula"
struct QuadRule
    points  :: Vector{Float64}
    weights :: Vector{Float64}
end

"order of the quadrature rule - number of quadrature points"
dim(qr::QuadRule) = length(qr.points)

"""
compute an 'n'-point Gausian integration rule of type 'method' on the interval [-1,1]

## input:
    n       :: Int             - number of quadrature points
    method  :: String          - choose among the following methods: "legendre", "lobatto", "chebyshev", "hermite", "jacobi", #                                 "radau". If left blank then the standard is "legendre"

## output:
    u       :: Vector{Float64} - quadrature points
    w       :: Vector{Float64} - quadrature weights
"""
function QuadRule(n::Int,method="legendre")

    if method=="legendre"
        u, w = gausslegendre(n)
    elseif method=="lobatto"
        u, w = gausslobatto(n)
    elseif method=="chebyshev"
        u, w = gausschebyshev(n)
    elseif method=="hermite"
        u, w = gausshermite(n)
    elseif method=="jacobi"
        u, w = gaussjacobi(n)
    elseif method=="radau"
        u, w = gaussradau(n)
    else
        error("Not a valid quadrature method")
    end

    return QuadRule(u,w)
end


# Information about mathematical functions needed for expression parsing

#################################################
#
#    Additional math functions
#
#################################################

sqr(x::Number) = x * x
rcp(x::Number) = one(x) / x
rsqrt(x::Number) = one(x) / sqrt(x)
rcbrt(x::Real) = one(x) / cbrt(x)

logit(x::Real) = log(x / (one(x) - x))
xlogx(x::Real) = x > zero(x) ? x * log(x) : zero(x)
xlogy(x::Real, y::Real) = x > zero(x) ? x * log(y) : zero(x)

logistic(x::Real) = one(x) / (one(x) + exp(-x))
invlogistic(y::Real) = -log(one(y) / y - one(y))

softplus(x::Real) = log(one(x) + exp(x))
invsoftplus(x::Real) = log(exp(x) - one(x))

logsumexp{T<:Real}(x::T, y::T) = x > y ? x + log(one(T) + exp(y - x)) : y + log(one(T) + exp(x - y))

@vectorize_1arg Number sqr
@vectorize_1arg Number rcp
@vectorize_1arg Real rsqrt
@vectorize_1arg Real rcbrt

@vectorize_1arg Real logit
@vectorize_1arg Real logistic


using Statistics

export bachelierFormula

"Price of a vanilla option under the Bachelier/Normal model"
function bachelierFormula(isCall::Bool, strike::T, forward::T, variance::T, discountDf::T) where {T}
    sign = 1.0
    if !isCall
        sign = -1.0
    end
    if variance == 0
        return Max(sign * (forward - strike), 0) * discountDf
    end
    sqrtvar = sqrt(variance)
    d = sign * (forward - strike) / sqrtvar
    if forward == strike
        return sqrtvar / sqrt(2 * pi)
    end
    normal = Normal()
    Nd = cdf(normal, d)
    nd = pdf(normal, d)
    return discountDf * (sign * (forward - strike) * Nd + sqrtvar * nd)
end


module Marcos

import GreyDecision.GreyNumbers: GreyNumber
import GreyDecision.Utility: unitize, weightize

function marcos(decisionMat::Array{GreyNumber{T}, 2}, weights::Array{Float64,1},
    fns::Array{Function,1}) where {T <: Real}

    row, col = size(decisionMat)

    w = unitize(weights)

    AAI = zeros(GreyNumber{T}, col)
    AI  = zeros(GreyNumber{T}, col)

    temp = [decisionMat; AI'; AAI']

    normalizedDecisionMat = similar(temp)

    @inbounds for i in 1:col
        if fns[i] == maximum
            AI[i] = maximum(decisionMat[:, i])
            temp[row + 1, i] = AI[i]
            AAI[i] = minimum(decisionMat[:, i])
            temp[row + 2, i] = AAI[i]
            normalizedDecisionMat[:, i] = map(x -> x / AI[i], temp[:, i])
        elseif fns[i] == minimum
            AI[i] = minimum(decisionMat[:, i])
            temp[row + 1, i] = AI[i]
            AAI[i] = maximum(decisionMat[:, i])
            temp[row + 2, i] = AAI[i]
            normalizedDecisionMat[:, i] = map(x -> AI[i] / x, temp[:, i])
        end
    end

    S  = zeros(GreyNumber{T}, col)

    for i in 1:(row + 2)
        S[i] = w .* normalizedDecisionMat[i, :] |> sum
    end

    KPlus  = map(x -> x / S[row + 1], S[1:row])
    KMinus = map(x -> x / S[row + 2], S[1:row])

    fKPlus = KPlus ./ (KPlus .+ KMinus)
    fKMinus = KMinus ./ (KPlus .+ KMinus)

    scores = zeros(GreyNumber{T}, row)

    for i in 1:row
        scores[i] = (KPlus[i] + KMinus[i]) / ((GreyNumber(1.0) + (GreyNumber(1.0) - fKPlus[i]) / fKPlus[i]) + ((GreyNumber(1.0) - fKMinus[i]) / fKMinus[i]))
    end

    rankings = sortperm(scores)

    bestIndex = rankings |> last

    result = Dict(
        :scores => scores,
        :rankings => rankings,
        :bestIndex => bestIndex
    )

    return result
end

end # end of module