module GaussianDistributions
# there are fundamental operations with Gaussian distributions:
# 1. marginals
# 2. conditionals
# 3. linear functionals
# 4. sampling
# 5. optimizing hyper-parameters of Gaussian Processes
# this package provides convenient and efficient implementations for these operations
function marginal end
function conditional end

using LinearAlgebra
using LinearAlgebra: checksquare

using Base.Threads
using StaticArrays
using FillArrays
using Statistics
import StatsBase: sample

using CovarianceFunctions
using CovarianceFunctions: gramian, MercerKernel, MultiKernel, AbstractKernel
Statistics.cov(k::MercerKernel, x::AbstractVector) = gramian(k, x)

using ForwardDiff # necessary for value_gradient in util.jl
using DiffResults

export Gaussian, Univariate, Multivariate, VectorGaussian, 𝒩
export mean, cov, var, std, pdf, cdf, nld, nlml
export marginal, conditional, conditional!, sample

using LazyInverse: inverse, Inverse, pseudoinverse
using LazyLinearAlgebra
using LinearAlgebraExtensions: difference, AbstractMatOrFacOrUni, AbstractMatOrFac
using WoodburyIdentity

abstract type AbstractGaussian{M, S} end
Statistics.mean(G::AbstractGaussian) = G.μ
Statistics.cov(G::AbstractGaussian) = G.Σ # ensure Hermitian-ness if Σ is a matrix?
const tol = 1e-6 # default error tolerance for matrix factorizations and samples
const VecOfVec{T} = AbstractVector{<:AbstractVector{T}}

############################# Gaussian Distribution ############################
struct Gaussian{M, S} <: AbstractGaussian{M, S}
    μ::M
    Σ::S
	function Gaussian(μ::Real, σ²::Real)
		σ² ≥ 0 ? new{typeof(μ), typeof(σ²)}(float(μ), σ²) : error("σ² < 0")
	end
	function Gaussian(μ::AbstractVector, Σ::AbstractMatOrFacOrUni)
		length(μ) == checksquare(Σ) || error("dimensions of μ, Σ not compatible")
		new{typeof(μ), typeof(Σ)}(μ, Σ)
	end
	Gaussian(μ, Σ) = new{typeof(μ), typeof(Σ)}(μ, Σ)
end
const 𝒩 = Gaussian
# 1-argument input constructors, assuming mean zero
Gaussian() = Gaussian(0., 1.)
# passing variances only
Gaussian(σ²::Real) = Gaussian(zero(σ²), σ²)
Gaussian(σ²::AbstractVector) = Gaussian(Diagonal(σ²))
Gaussian(Σ::AbstractMatOrFacOrUni) = Gaussian(Zeros(size(Σ, 1)), Σ)
# passing mean and variances
Gaussian(μ::AbstractVector, σ²::Real) = Gaussian(μ, σ²*I(length(μ)))
Gaussian(μ::AbstractVector, σ²::AbstractVector) = Gaussian(μ, Diagonal(σ²))

zero_mean(x) = zero(eltype(x))
# TODO: make this a mutable static vector
zero_mean(x, d::Int) = zeros(eltype(x), d)
zero_mean(d::Int) = x->zero_mean(x, d)

Gaussian(k) = Gaussian(zero_mean, k)

function Base.length(G::Gaussian)
	if G isa Univariate || G isa Multivariate
		return length(G.μ)
	else
		return Inf # assumed to be process
	end
end

const UnivariateGaussian{M, S} = Gaussian{<:Real, <:Real}
const Univariate = UnivariateGaussian
const Uni = Univariate

# (G::Univariate)(i::Int) = marginal(G, i)
Base.getindex(G::Univariate, i) = marginal(G, i) # bracket-indexing

const MultivariateGaussian{M, S} = Gaussian{<:AbstractVector, <:AbstractMatOrFacOrUni}
const Multivariate = MultivariateGaussian
const Multi = Multivariate

# (G::Multivariate)(i::Int) = marginal(G, i)
Base.getindex(G::Multivariate, i) = marginal(G, i) # bracket-indexing

# Gaussian is assumed to be process-valued, if neither Uni, nor Multi
(G::Gaussian)(x) = Gaussian(mean(G)(x), cov(G)(x, x))

############################# convenience ######################################
Statistics.cov(G::Multivariate, i, j) = cov(G)[i, j] # view?
Statistics.var(G::Union{Uni, Multi}) = diag(cov(G))
Statistics.std(G::Union{Uni, Multi}) = sqrt.(max.(var(G), 0))

# variance (computes diagonal of covariance), should be specialized where
# more efficient computation is possible
Statistics.var(G::Gaussian) = x->max(G.Σ(x, x), 0)
Statistics.std(G::Gaussian) = x->sqrt(var(G)(x))
Statistics.mean(G::Gaussian, x::AbstractVector) = marginal(mean(G), x)
# IDEA: allow for randomized approximation of this function akin to "Scaling Gaussian Process Regression with Derivatives"
# Statistics.var(G::Gaussian, x::AbstractVector) = var(G).(x)
function Statistics.var(G::Gaussian, x::AbstractVector)
	σ² = fill(var(G)(x[1]), length(x))
	@threads for i in 2:length(x) # virtually perfect parallel scaling
		σ²[i] = var(G)(x[i])
	end
	return σ²
end
function Statistics.std(G::Gaussian, x::AbstractVector)
	σ = var(G, x)
	@. σ = sqrt(max(σ, 0))
end

Base.copy(G::Gaussian) = Gaussian(copy(G.μ), copy(G.Σ))
function Base.:(==)(G::Gaussian, H::Gaussian)
	return mean(G) == mean(H) && cov(G) == cov(H)
end
function Base.isapprox(G::Gaussian, H::Gaussian; atol = eps(), rtol = eps())
    isapprox(G.μ, H.μ, atol = atol, rtol = rtol) && isapprox(G.Σ, H.Σ, atol = atol, rtol = rtol)
end

function LinearAlgebra.cholesky(G::Multivariate; tol::Real = tol, check::Bool = false) # low rank via cholesky
	cov(G) isa AbstractMatrix || return G # to prevent calling cholesky on a factorization
	S = Symmetric(cov(G))
    C = cholesky(S, Val(true), tol = tol, check = check)
    if check
		issuccess(C) || throw("pivoted cholesky not successful")
	end
    return Gaussian(mean(G), C)
end
function LinearAlgebra.factorize(G::Multivariate; tol::Real = tol)
	cov(G) isa AbstractMatrix || return G # to prevent calling cholesky on a non-factorizable object
	return cholesky(G, tol = tol)
end

include("util.jl")
include("sample.jl")
include("marginal.jl")
include("conditional.jl")
include("multi.jl")
include("linear.jl")
include("density.jl")
include("optimization.jl")
include("plot.jl")

# non-essential extensions
include("inputnoise.jl")

end # GaussianDistributions


using FluxRC, Test

# Kou's benchmark results
"""
r
[-0.81684757,  0.63369515, -0.81684757, -0.10810302, -0.78379396, -0.10810302]

s
[0.63369515, -0.81684757, -0.81684757, -0.78379396, -0.10810302, -0.10810302]

rf
[-0.77459667, 0., 0.77459667, 0.77459667, 0., -0.77459667, -1., -1., -1.]

sf
[-1., -1., -1., -0.77459667, 0., 0.77459667, 0.77459667, 0., -0.77459667]

V
[[ 0.70710678,  1.45054272,  1.39329302,  0.        ,  0.        ,
    -0.04593312],
[ 0.70710678, -0.72527136,  0.43019448,  1.25620684, -0.83406781,
    1.03084378],
[ 0.70710678, -0.72527136,  0.43019448, -1.25620684,  0.83406781,
    1.03084378],
[ 0.70710678, -0.67569095,  0.30868284,  0.        , -0.        ,
    -1.08925616],
[ 0.70710678,  0.33784547, -0.70898933, -0.58516552, -0.88132995,
    0.04853591],
[ 0.70710678,  0.33784547, -0.70898933,  0.58516552,  0.88132995,
    0.04853591]]

correction_coeffs
[[ 0.3928371   0.62853936  0.3928371   0.55555556  0.88888889  0.55555556
  0.3928371   0.62853936  0.3928371 ]
[-0.55555556 -0.88888889 -0.55555556 -0.52003383  0.62853936  1.30570803
  0.923275    0.44444444 -0.36771945]
[ 0.68041382  1.08866211  0.68041382  0.21689446 -0.76980036  1.70760644
  1.20746009 -0.54433105  0.15336754]
[-0.74535599  0.          0.74535599  1.20746009  1.08866211  0.15336754
 -0.10844723 -0.76980036 -0.85380322]
[ 0.91287093 -0.         -0.91287093 -0.64549722  2.          0.64549722
 -0.45643546 -1.41421356  0.45643546]
[ 0.60858062 -1.21716124  0.60858062  1.6939963   0.86066297  0.02732963
  0.01932497  0.60858062  1.19783627]]

phifj_solution
[[ 0.39198253  0.72780648  0.39198253 -0.13710719  0.42817214  4.66476318
  3.29848568  0.30276343 -0.09694943]
[-0.09694943  0.30276343  3.29848568  4.66476318  0.42817214 -0.13710719
  0.39198253  0.72780648  0.39198253]
[ 3.29848568  0.30276343 -0.09694943  0.55434701  1.02927379  0.55434701
 -0.09694943  0.30276343  3.29848568]
[ 0.20029351  2.7069103   0.20029351 -1.03402506 -0.97126559  0.00792184
  0.00560158 -0.68678848 -0.73116613]
[-0.73116613 -0.68678848  0.00560158  0.00792184 -0.97126559 -1.03402506
  0.20029351  2.7069103   0.20029351]
[ 0.00560158 -0.68678848 -0.73116613  0.2832578   3.82814926  0.2832578
 -0.73116613 -0.68678848  0.00560158]]
"""

py_V = [
    0.70710678 1.45054272 1.39329302 0.0 0.0 -0.04593312
    0.70710678 -0.72527136 0.43019448 1.25620684 -0.83406781 1.03084378
    0.70710678 -0.72527136 0.43019448 -1.25620684 0.83406781 1.03084378
    0.70710678 -0.67569095 0.30868284 0.0 -0.0 -1.08925616
    0.70710678 0.33784547 -0.70898933 -0.58516552 -0.88132995 0.04853591
    0.70710678 0.33784547 -0.70898933 0.58516552 0.88132995 0.04853591
]

py_Vf = [
    0.70710678 -1.0 1.22474487 -1.34164079 1.64316767 1.09544512
    0.70710678 -1.0 1.22474487 0.0 -0.0 -1.36930639
    0.70710678 -1.0 1.22474487 1.34164079 -1.64316767 1.09544512
    0.70710678 -0.661895 0.27606157 1.5368458 -0.82158384 2.15610529
    0.70710678 0.5 -0.61237244 0.8660254 1.59099026 0.6846532
    0.70710678 1.661895 2.17342817 0.19520501 0.82158384 0.03478494
    0.70710678 1.661895 2.17342817 -0.19520501 -0.82158384 0.03478494
    0.70710678 0.5 -0.61237244 -0.8660254 -1.59099026 0.6846532
    0.70710678 -0.661895 0.27606157 -1.5368458 0.82158384 2.15610529
]

phifj_ref = [
    0.39198253 0.72780648 0.39198253 -0.13710719 0.42817214 4.66476318 3.29848568 0.30276343 -0.09694943
    -0.09694943 0.30276343 3.29848568 4.66476318 0.42817214 -0.13710719 0.39198253 0.72780648 0.39198253
    3.29848568 0.30276343 -0.09694943 0.55434701 1.02927379 0.55434701 -0.09694943 0.30276343 3.29848568
    0.20029351 2.7069103 0.20029351 -1.03402506 -0.97126559 0.00792184 0.00560158 -0.68678848 -0.73116613
    -0.73116613 -0.68678848 0.00560158 0.00792184 -0.97126559 -1.03402506 0.20029351 2.7069103 0.20029351
    0.00560158 -0.68678848 -0.73116613 0.2832578 3.82814926 0.2832578 -0.73116613 -0.68678848 0.00560158
]

# my workflow
N = deg = 2
Np = (N + 1) * (N + 2) ÷ 2

pl, wl = tri_quadrature(N)

V = vandermonde_matrix(N, pl[:, 1], pl[:, 2])
Vr, Vs = ∂vandermonde_matrix(N, pl[:, 1], pl[:, 2])
∂l = ∂lagrange(V, Vr, Vs)

ϕ = correction_field(N, V)

pf, wf = triface_quadrature(N)
ψf = zeros(3, N + 1, Np)
for i = 1:3
    ψf[i, :, :] .= vandermonde_matrix(N, pf[i, :, 1], pf[i, :, 2])
end

# Vandermonde -> solution points
V1 = hcat(V[1, :], V[3, :], V[5, :], V[2, :], V[4, :], V[6, :]) |> permutedims
@test V1 ≈ py_V

# Vandermonde -> flux points
Vf = zeros(9, 6)
for i = 1:3, j = 1:3
    Vf[(i-1)*3+j, :] .= ψf[i, j, :]
end
@test Vf ≈ py_Vf

# coefficients
ϕ1 = zeros(6, 9)
for i = 1:3, j = 1:3
    ϕ1[:, (i-1)*3+j] .= ϕ[i, j, :]
end
ϕ2 = hcat(ϕ1[1, :], ϕ1[3, :], ϕ1[5, :], ϕ1[2, :], ϕ1[4, :], ϕ1[6, :]) |> permutedims
@test ϕ2 ≈ phifj_ref

# reproduce from python script
cd(@__DIR__)
include("python_phi.jl")

py_sigma, py_fi, py_figrad =
    py"get_phifj_solution_grad_tri"(N, N + 1, 6, 3 * (N + 1), N + 1, py_V, py_Vf)
@test ϕ2 ≈ py_fi


function spdiags(B, d, m, n)
    d = d[:]
    p = length(d)

    len = zeros(p + 1, 1)
    for k in 1:p
        len[k+1] =
            int(len[k] + length(Base.max(1, 1 - d[k]):Base.min(m, n - d[k])))
    end
    a = zeros(int(len[p+1]), 3)
    for k in 1:p
        # Append new d[k]-th diagonal to compact form
        i = Base.max(1, 1 - d[k]):Base.min(m, n - d[k])
        a[(int(len[k])+1):int(len[k+1]), :] = [i i + d[k] B[i+(m>=n)*d[k], k]]
    end

    A = sparse(int(a[:, 1]), int(a[:, 2]), a[:, 3], m, n)

    return A
end

y = readcsv("snp_500.txt")
n = length(y)

e = ones(n, 1)
D = spdiags([e -2 * e e], 0:2, n - 2, n)

lambda = 50
start = time()
D * x
println(time() - start)

x = Variable(n)
p = minimize(lambda * sum(D * x))
println(time() - start)
solve!(p)
p.optval
println(time() - start)


using LinearAlgebra

function recursion_partition_function(num_elec, eff_spec_SpinUp, eff_spec_SpinDn)
    Z_up, Z_dn = zeros(num_elec), zeros(num_elec)
    zₖ_up, zₖ_dn = zeros(num_elec), zeros(num_elec)
    for k = 1 : num_elec
        zₖ_up[k] = sum(eff_spec_SpinUp .^ k)
        zₖ_dn[k] = sum(eff_spec_SpinDn .^ k)
    end
    for N = 1 : num_elec
        for k = 1 : N
            sgn_k = 2 * isodd(k) - 1
            if k == N
                Z_up[N] += sgn_k * zₖ_up[k]
                Z_dn[N] += sgn_k * zₖ_dn[k]
            else
                Z_up[N] += sgn_k * Z_up[N-k] * zₖ_up[k]
                Z_dn[N] += sgn_k * Z_dn[N-k] * zₖ_dn[k]
            end
        end
        Z_up[N] = Z_up[N] / N
        Z_dn[N] = Z_dn[N] / N
    end
    return Z_up, Z_dn
end

function recursion_OccNum(Z_up, eff_spec_SpinUp, Z_dn, eff_spec_SpinDn, num_elec, num_sites) # Calculating <nᵢ> * Z_N
    n_up_mean = zeros(num_elec, num_sites)
    n_dn_mean = zeros(num_elec, num_sites)
    n_up_mean[1, :], n_dn_mean[1, :] = eff_spec_SpinUp, eff_spec_SpinDn
    for N = 2 : num_elec
        n_up_mean[N, :] = eff_spec_SpinUp .* (Z_up[N-1] .- n_up_mean[N-1, :])
        n_dn_mean[N, :] = eff_spec_SpinDn .* (Z_dn[N-1] .- n_dn_mean[N-1, :])
    end
    return n_up_mean ./ Z_up, n_dn_mean ./ Z_dn
end

function measure_tot_energy(num_sites, t, U, U_array_SpinUp, U_array_SpinDn, n_up_mean, n_dn_mean, num_SpinUp, num_SpinDn, Z_up, Z_dn)
    E_onebody, E_twobody = 0, 0
    for i = 1 : num_sites
        for j = 1 : num_sites
            if abs(i - j) == 1 || abs(i - j) == num_sites - 1
                D_ij_up = sum(U_array_SpinUp[i, j, :] .* n_up_mean[num_SpinUp, :])
                D_ij_dn = sum(U_array_SpinDn[i, j, :] .* n_dn_mean[num_SpinDn, :])
                E_onebody += -t * ( D_ij_up + D_ij_dn )
            end
        end
    end
    for i = 1 : num_sites
        U_matrix = U_array_SpinUp[i, i, :] * U_array_SpinDn[i, i, :]'
        n_matrix = n_up_mean[num_SpinUp, :] * n_dn_mean[num_SpinDn, :]'
        D_iiii = sum(U_matrix .* n_matrix)
        E_twobody += U * D_iiii
    end
    return E_onebody, E_twobody
end

function recursion(self::Lattice, expϵup::Array{Float64,1}, expϵdn::Array{Float64,1})
    nup_mean = zeros(Float64, self.nsites, self.nsites)
    ndn_mean = zeros(Float64, self.nsites, self.nsites)
    Znup_ratio = zeros(Float64, self.ntot)
    Zndn_ratio = zeros(Float64, self.ntot)
    expϵup = expϵup ./ expϵup[self.nsites]
    expϵdn = expϵdn ./ expϵdn[self.nsites]
    Znup_ratio[1], Zndn_ratio[1] = sum(expϵup), sum(expϵdn)
    nup_mean[:, 1], ndn_mean[:, 1] = expϵup / Znup_ratio[1], expϵdn / Zndn_ratio[1]
    for N = 2 : max(self.nup, self.ndn)
        Znup_ratio[N] = sum(expϵup .* (1 .- nup_mean[:, N-1])) / N
        Zndn_ratio[N] = sum(expϵdn .* (1 .- ndn_mean[:, N-1])) / N
        nup_mean[:, N] = expϵup .* (1 .- nup_mean[:, N - 1]) / Znup_ratio[N]
        ndn_mean[:, N] = expϵdn .* (1 .- ndn_mean[:, N - 1]) / Zndn_ratio[N]
    end
    return nup_mean, ndn_mean, Znup_ratio, Zndn_ratio
end

function measure_energy(
    self::Lattice, Uup_array::Array{Float64,3}, Udn_array::Array{Float64,3},
    nup_mean::Array{Float64,2}, ndn_mean::Array{Float64,2}
)
    Eonebody, Etwobody = 0, 0
    for i = 1 : self.nsites
        for j = 1 : self.nsites
            if abs(i - j) == 1 || abs(i - j) == self.nsites - 1
                Dij_up = sum(Uup_array[:, i, j] .* nup_mean[:, self.nup])
                Dij_dn = sum(Udn_array[:, i, j] .* ndn_mean[:, self.ndn])
                Eonebody += -self.t * ( Dij_up + Dij_dn )
            end
        end
    end
    for i = 1 : self.nsites
        Umatrix = Uup_array[:, i, i] * Udn_array[:, i, i]'
        nmatrix = nup_mean[:, self.nup] * ndn_mean[:, self.ndn]'
        Diiii = sum(Umatrix .* nmatrix)
        Etwobody += self.U * Diiii
    end
    return Eonebody, Etwobody
end

# rank-one perturbation calculator
function rankone_perturb!(
    γ::Float64, σil::Int64, Ns::Int64,
    MatProdUp_new::Array{Float64,2}, MatProdDn_new::Array{Float64,2}
)
    Δup, Δdn = Matrix{Float64}(I, size(MatProdUp_new)), Matrix{Float64}(I, size(MatProdDn_new))
    Δup[Ns, Ns] = exp(-4 * γ * σil)
    Δdn[Ns, Ns] = exp(4 * γ * σil)
    MatProdUp_new = Δup * MatProdUp_new
    MatProdDn_new = Δdn * MatProdDn_new
    return MatProdUp_new, MatProdDn_new
end

# determine the sign of the accept_ratio
@inline function Rsign(R::Float64)
    Rsgn = 2 * (R > 0) - 1
end

function MCIntegration(lattice::Lattice, qmc::QMC, opt::Opt)
    # initialize vectors that store MC data
    Ek_array = Vector{Float64}() # kinetic energy
    Ep_array = Vector{Float64}() # potential energy
    E_array = Vector{Float64}() # total energy
    sgn_array = Vector{Int64}() # sign
    NiMax_array = Vector{Float64}()
    # initialize a random configuration
    σlist = 2 * (rand(Float64, (lattice.nsites, qmc.L)) .< 0.5) .- 1
    MatProdUp, MatProdDn, EigensUp, EigensDn = matrix_product_QR(σlist, lattice, qmc, opt.stab_interval)
    nup_mean, ndn_mean, Znup_ratio, Zndn_ratio = recursion(lattice, real(EigensUp.values), real(EigensDn.values))
    sgn = Rsign(prod(Znup_ratio[1 : lattice.nup]) * prod(Zndn_ratio[1 : lattice.nup]))
    for n = 1 : qmc.samples
        σlist_new = copy(σlist)
        # random flip a spin
        #flipL, flipNs = rand(1 : qmc.L), rand(1 : lattice.nsites)
        flip = rand(1 : qmc.L * lattice.nsites)
        σlist_new[flip] *= -1
        #rankone_perturb!(qmc.γ, σlist[flipNs, flipL], flipNs, MatProdUp_new, MatProdDn_new)
        MatProdUp_new, MatProdDn_new, EigensUp_new, EigensDn_new = matrix_product_QR(σlist_new, lattice, qmc, opt.stab_interval)
        nup_mean_new, ndn_mean_new, Znup_ratio_new, Zndn_ratio_new = recursion(lattice, real(EigensUp_new.values), real(EigensDn_new.values))
        Rup = prod(Znup_ratio_new[1 : lattice.nup] ./ Znup_ratio[1 : lattice.nup])
        Rdn = prod(Zndn_ratio_new[1 : lattice.ndn] ./ Zndn_ratio[1 : lattice.ndn])
        sgn_new = Rsign(Rup * Rdn) * sgn
        R = abs(Rup * Rdn)
        accept_ratio = R / (1 + R)
        if rand() <= accept_ratio   # accept
            σlist = σlist_new
            Znup_ratio, Zndn_ratio = Znup_ratio_new, Zndn_ratio_new
            EigensUp, EigensDn = EigensUp_new, EigensDn_new
            nup_mean, ndn_mean = nup_mean_new, ndn_mean_new
            sgn = sgn_new
            Uup_array, Udn_array = overlap_coefficient(lattice.nsites, EigensUp.vectors, EigensDn.vectors)
            Eonebody, Etwobody = measure_energy(lattice, Uup_array, Udn_array, nup_mean, ndn_mean)
            push!(sgn_array, sgn)
            push!(Ek_array, Eonebody * sgn)
            push!(Ep_array, Etwobody * sgn)
            push!(E_array, Eonebody * sgn + Etwobody * sgn)
            push!(NiMax_array, maximum(nup_mean[:, lattice.nup]))
        else    # reject
            Uup_array, Udn_array = overlap_coefficient(lattice.nsites, EigensUp.vectors, EigensDn.vectors)
            Eonebody, Etwobody = measure_energy(lattice, Uup_array, Udn_array, nup_mean, ndn_mean)
            push!(sgn_array, sgn)
            push!(Ek_array, Eonebody * sgn)
            push!(Ep_array, Etwobody * sgn)
            push!(E_array, Eonebody * sgn + Etwobody * sgn)
            push!(NiMax_array, maximum(nup_mean[:, lattice.nup]))
        end
    end
    E_mean = mean(E_array[opt.nburn_in : qmc.samples])
    E_error = std(E_array[opt.nburn_in : qmc.samples]) / sqrt(qmc.samples - opt.nburn_in)
    Ek_mean = mean(Ek_array[opt.nburn_in : qmc.samples])
    Ek_error = std(Ek_array[opt.nburn_in : qmc.samples]) / sqrt(qmc.samples - opt.nburn_in)
    Ep_mean = mean(Ep_array[opt.nburn_in : qmc.samples])
    Ep_error = std(Ep_array[opt.nburn_in : qmc.samples]) / sqrt(qmc.samples - opt.nburn_in)
    sgn_avg = mean(sgn_array[opt.nburn_in : qmc.samples])
    NiMax_avg = mean(NiMax_array[opt.nburn_in : qmc.samples])
    return E_mean, E_error, Ek_mean, Ek_error, Ep_mean, Ep_error, sgn_avg, NiMax_avg
end


"""
    RectangularCluster

Fills the landscape with rectangles containing a random value. The size of each
rectangle/patch is between `minimum` and `maximum` (the two can be equal for a
fixed size rectangle).
"""
struct RectangularCluster <: NeutralLandscapeMaker
    minimum::Integer
    maximum::Integer
    function RectangularCluster(x::T, y::T) where {T <: Integer}
        @assert 0 < x <= y
        new(x, y)
    end
end

RectangularCluster() = RectangularCluster(2,4)

function _landscape!(mat, alg::RectangularCluster)
    mat .= -1.0
    while minimum(mat) == -1.0
        width, height = rand(alg.minimum:alg.maximum, 2)
        row = rand(1:(size(mat,1)-(width-1)))
        col = rand(1:(size(mat,2)-(height-1)))
        mat[row:(row+(width-1)) , col:(col+(height-1))] .= rand()
    end
    return mat
end


using UnicodePlots

# a simple example for DNN_CTCLoss_With_Softmax usage
# by the way, it's for a single input case
S = 5
T = 100
E = 1000
x = 0.5*rand(S,T);
P = Variable(x,keepsgrad=true);
L = zeros(1, E)
for i = 1:E
    L[i] = DNN_CTC_With_Softmax(P, [2 3 4 5 2 3 4 5])
    backward()
    update!(P, 1E-3)
end

lineplot(vec(L))

# a simple example for DNN_Batch_CTCLoss_With_Softmax usage
# by the way, it's for a batch of inputs case
E = 100
S = 5
T1 = 10
T2 = 20
x1 = 0.5*rand(S,T1);  # first input
x2 = 0.3*rand(S,T2);  # second input
seq1 = [2 3 4];       # first input's label
seq2 = [3 2];         # second input's label

x   = [x1 x2];        # a batch of inputs
seq = [seq1 seq2];    # a batch of labels

P = Variable(x, keepsgrad=true);
L = zeros(1, E)
for i = 1:E
    L[i] = DNN_Batch_CTC_With_Softmax(P, seq, [T1 T2], [length(seq1) length(seq2)])
    backward()
    update!(P, 1E-3)
end

lineplot(vec(L))


# Truncated normal distribution
"""
    TruncatedNormal(mu, sigma, l, u)

The *truncated normal distribution* is a particularly important one in the family of truncated distributions.
We provide additional support for this type with `TruncatedNormal` which calls `Truncated(Normal(mu, sigma), l, u)`.
Unlike the general case, truncated normal distributions support `mean`, `mode`, `modes`, `var`, `std`, and `entropy`.
"""
TruncatedNormal(mu::Real, sigma::Real, a::Real, b::Real) =
    Truncated(Normal(mu, sigma), a, b)

### statistics

minimum(d::Truncated{Normal{T},Continuous}) where {T <: Real} = d.lower
maximum(d::Truncated{Normal{T},Continuous}) where {T <: Real} = d.upper

function mode(d::Truncated{Normal{T},Continuous}) where T <: Real
    μ = mean(d.untruncated)
    d.upper < μ ? d.upper :
    d.lower > μ ? d.lower : μ
end

modes(d::Truncated{Normal{T},Continuous}) where {T <: Real} = [mode(d)]

# do not export. Used in mean, var
function _F1(x::Real, y::Real; thresh=1e-7)
    @assert 0 < thresh < Inf
    -Inf < x < Inf && -Inf < y < Inf || throw(DomainError())
    ϵ = exp(x^2 - y^2)
    if abs(x) > abs(y)
        _F1(y,x)
    elseif abs(x - y) ≤ thresh
        δ = y - x
        √π*x + (√π/2 + (-√π*x/6 + (-√π/12 + x*(√π/90 + (√π*x^2)/90)δ)δ)δ)δ
    elseif max(x,y) < 0
        (1 - ϵ) / (ϵ * erfcx(-y) - erfcx(-x))
    elseif min(x,y) > 0
        (1 - ϵ) / (erfcx(x) - ϵ * erfcx(y))
    else
        exp(-x^2) * (1 - ϵ) / (erf(y) - erf(x))
    end
end

# do not export. Used in mean, var
function _F2(x::Real, y::Real; thresh=1e-7)
    @assert 0 < thresh < Inf
    -Inf < x < Inf && -Inf < y < Inf || throw(DomainError())
    ϵ = exp(x^2 - y^2)
    if abs(x) > abs(y)
        _F2(y,x)
    elseif abs(x - y) ≤ thresh
        δ = y - x
        √π*x^2 - √π/2 + (√π*x + (√π/3 - √π*x^2/3 + (((√π/30 + √π*x^2/45)x^2 - 4*√π/45)δ - √π*x/3)δ)δ)δ
    elseif max(x,y) < 0
        (x - ϵ * y) / (ϵ * erfcx(-y) - erfcx(-x))
    elseif min(x,y) > 0
        (x - ϵ * y) / (erfcx(x) - ϵ * erfcx(y))
    else
        exp(-x^2) * (x - ϵ * y) / (erf(y) - erf(x))
    end
end

# do not export. Used in mean
function _tnmean(a, b)
    -Inf < a ≤ b < Inf || throw(DomainError())
    √(2/π) * _F1(a/√2, b/√2)
end

# do not export. Used in mean
function _tnmean(a, b, μ, σ)
    -Inf < a ≤ b < Inf || throw(DomainError())
    -Inf < μ < Inf && 0 < σ < Inf || throw(DomainError())
    α = (a - μ)/σ; β = (b - μ)/σ
    μ + _tnmean(α, β) * σ
end

# do not export. Used in var
function _tnvar(a, b)
    -Inf < a ≤ b < Inf || throw(DomainError())
    1 + 2/√π * _F2(a/√2, b/√2) - 2/π * _F1(a/√2, b/√2)^2
end

# do not export. Used in var
function _tnvar(a, b, μ, σ)
    -Inf < a ≤ b < Inf || throw(DomainError())
    -Inf < μ < Inf && 0 < σ < Inf || throw(DomainError())
    α = (a - μ)/σ; β = (b - μ)/σ
    _tnvar(α, β) * σ^2
end

function mean(d::Truncated{Normal{T},Continuous}) where T <: Real
    d0 = d.untruncated
    μ = mean(d0)
    σ = std(d0)
    if isfinite(d.lower) && isfinite(d.upper) && isfinite(μ) && isfinite(σ)
        # avoids loss of significance when truncation is far from μ.
        # See https://github.com/cossio/TruncatedNormal.jl/blob/master/notes/normal.pdf.
        _tnmean(d.lower, d.upper, μ, σ)
    else
        a = (d.lower - μ) / σ
        b = (d.upper - μ) / σ
        μ + ((normpdf(a) - normpdf(b)) / d.tp) * σ
    end
end

function var(d::Truncated{Normal{T},Continuous}) where T <: Real
    d0 = d.untruncated
    μ = mean(d0)
    σ = std(d0)
    if isfinite(d.lower) && isfinite(d.upper) && isfinite(μ) && isfinite(σ)
        # avoids loss of significance when truncation is far from μ.
        # See https://github.com/cossio/TruncatedNormal.jl/blob/master/notes/normal.pdf.
        _tnvar(d.lower, d.upper, μ, σ)
    else
        a = (d.lower - μ) / σ
        b = (d.upper - μ) / σ
        z = d.tp
        φa = normpdf(a)
        φb = normpdf(b)
        aφa = isinf(a) ? 0.0 : a * φa
        bφb = isinf(b) ? 0.0 : b * φb
        t1 = (aφa - bφb) / z
        t2 = abs2((φa - φb) / z)
        abs2(σ) * (1 + t1 - t2)
    end
end

function entropy(d::Truncated{Normal{T},Continuous}) where T <: Real
    d0 = d.untruncated
    z = d.tp
    μ = mean(d0)
    σ = std(d0)
    a = (d.lower - μ) / σ
    b = (d.upper - μ) / σ
    aφa = isinf(a) ? 0.0 : a * normpdf(a)
    bφb = isinf(b) ? 0.0 : b * normpdf(b)
    0.5 * (log2π + 1.) + log(σ * z) + (aφa - bφb) / (2.0 * z)
end

### sampling

## Use specialized sampler, as quantile-based method is inaccurate in
## tail regions of the Normal, issue #343

function rand(rng::AbstractRNG, d::Truncated{Normal{T},Continuous}) where T <: Real
    d0 = d.untruncated
    μ = mean(d0)
    σ = std(d0)
    a = (d.lower - μ) / σ
    b = (d.upper - μ) / σ
    z = randnt(rng, a, b, d.tp)
    return μ + σ * z
end

# Rejection sampler based on algorithm from Robert (1995)
#
#  - Available at http://arxiv.org/abs/0907.4010

function randnt(rng::AbstractRNG, lb::Float64, ub::Float64, tp::Float64)
    local r::Float64
    if tp > 0.3   # has considerable chance of falling in [lb, ub]
        r = randn(rng)
        while r < lb || r > ub
            r = randn(rng)
        end
        return r

    else
        span = ub - lb
        if lb > 0 && span > 2.0 / (lb + sqrt(lb^2 + 4.0)) * exp((lb^2 - lb * sqrt(lb^2 + 4.0)) / 4.0)
            a = (lb + sqrt(lb^2 + 4.0))/2.0
            while true
                r = rand(rng, Exponential(1.0 / a)) + lb
                u = rand(rng)
                if u < exp(-0.5 * (r - a)^2) && r < ub
                    return r
                end
            end
        elseif ub < 0 && ub - lb > 2.0 / (-ub + sqrt(ub^2 + 4.0)) * exp((ub^2 + ub * sqrt(ub^2 + 4.0)) / 4.0)
            a = (-ub + sqrt(ub^2 + 4.0)) / 2.0
            while true
                r = rand(rng, Exponential(1.0 / a)) - ub
                u = rand(rng)
                if u < exp(-0.5 * (r - a)^2) && r < -lb
                    return -r
                end
            end
        else
            while true
                r = lb + rand(rng) * (ub - lb)
                u = rand(rng)
                if lb > 0
                    rho = exp((lb^2 - r^2) * 0.5)
                elseif ub < 0
                    rho = exp((ub^2 - r^2) * 0.5)
                else
                    rho = exp(-r^2 * 0.5)
                end
                if u < rho
                    return r
                end
            end
        end
    end
end


function label(colors::AbstractArray)

    # Set the color-number converter list
    color_key = ["black", "brown", "red", "orange", "yellow", "green", "blue", "violet", "grey", "white"]

    # Initialize the array that hold the converted numbers
    converted_colors = collect(1:3)

    # For every color, put the converted number in the converted_colors array
    for color_index in 1:3
        converted_colors[color_index] = (findfirst(x -> x == colors[color_index], color_key) - 1)
    end

    # If the second digit ends with zero, end the first two digits at the first digit, and increment the power of ten
    # else, keep the two digits and convert the third number into a power of 10
    if converted_colors[2] == 0
        first_two_digits = converted_colors[1]
        converted_colors[3] += 1
        power_of_ten = 10^converted_colors[3]
    else
        first_two_digits = 10*converted_colors[1] + converted_colors[2]
        power_of_ten = 10^converted_colors[3]
    end

    # If the resistance is an integral multiple of 1000 ohms, write the answer in terms of kiloohms
    # Else, write the answer in terms of ohms
    if power_of_ten >= 1000
        power_of_ten = 10^(converted_colors[3]-3)
        suffix = "kiloohms"
        first_two_digits *= power_of_ten
    else
        first_two_digits *= power_of_ten
        suffix = "ohms"
    end

    # print the final anwer with the appropriate suffix
    return ("$first_two_digits" * " " * "$suffix")

end


# Define abstract interface
"""
$(TYPEDEF)

Abstract supertype for magnetic fields.
"""
abstract type MagneticFields{T} end

"""
$(SIGNATURES)

Return the magnetic field strength H [T/μ₀] (μ₀ being the vacuum permeability)
as `SVector{3,T}` for a given location r [m] and time t [s].
"""
function magneticFieldStrength(magneticField::MagneticFields{T},r::SVector{3,T},t::T) where {T}
    error("magneticFieldStrength is not implemented for magneticField::$(typeof(magneticField)).")
end

"""
$(TYPEDEF)

Abstract supertype for static magnetic fields.
"""
abstract type StaticMagneticFields{T} <: MagneticFields{T} end

"""
$(SIGNATURES)

Return the magnetic field strength H [T/μ₀] (μ₀ being the vacuum permeability)
as `SVector{3,T}` for a given location r [m].
"""
function magneticFieldStrength(magneticField::StaticMagneticFields{T},r::SVector{3,T}) where {T}
    error("magneticFieldStrength is not implemented for magneticField::$(typeof(magneticField)).")
end

function magneticFieldStrength(magneticField::StaticMagneticFields{T},r::SVector{3,T},t::T) where {T}
    return magneticFieldStrength(magneticField,r)
end

# Define specific subtypes
## ideal gradient field
struct StaticGradientField{T} <: StaticMagneticFields{T}
    center::SVector{3,T}
    gradient::SMatrix{3,3,T}

    StaticGradientField(center::AbstractVector{T},gradient::AbstractMatrix{T}) where {T} = new{T}(SVector{3,T}(center),SMatrix{3,3,T}(gradient))
end

function magneticFieldStrength(magneticField::StaticGradientField{T},r::SVector{3,T}) where {T}
    return magneticField.gradient*(r-magneticField.center)
end

## ideal homogenous field
struct StaticHomogeneousField{T} <: StaticMagneticFields{T}
    magneticFieldStrength::SVector{3,T}

    StaticHomogeneousField(magneticFieldStrength::Vector{T}) where {T} = new{T}(SVector{3,T}(magneticFieldStrength))
end

function magneticFieldStrength(magneticField::StaticHomogeneousField{T},r::SVector{3,T}) where {T}
    return magneticField.magneticFieldStrength
end

## ideal homogenous field with sinusoidal 1D excitation
struct SinusoidalField{T} <: MagneticFields{T}
    spatialFieldProfile::StaticMagneticFields{T}
    baseFrequency::T
    divider::Int
    doubleFrequency::T
    relativePhase::T # = phase/π
end

function SinusoidalField(spatialFieldProfile::StaticMagneticFields{T},baseFrequency::Real,divider::Real,relativePhase::Real=0.0) where {T}
    doubleFrequency = 2//divider * baseFrequency
    return SinusoidalField{T}(spatialFieldProfile,T(baseFrequency),T(divider),T(doubleFrequency),T(relativePhase))
end

function magneticFieldStrength(magneticField::SinusoidalField{T},r::SVector{3,T},t::T) where {T}
    return magneticFieldStrength(magneticField.spatialFieldProfile,r)*sinpi(magneticField.doubleFrequency*t+magneticField.relativePhase)
end


# # Types
#
# Topics:
# 1. How types work in Julia
# 2. Defining types
# 3. Type parameters
# 4. Example: `OneHotVector`
# 5. Tuple types
# 6. Summary of kinds of types
# 7. The types of functions

#-

# ## How types work in Julia
#
# Every Julia value has a unique concrete type associated with it:

typeof(3)
typeof('a')

# Types are also first-class values:

typeof(Int)
typeof(DataType)

# The type of a value describes its memory layout as well as describing,
# semantically, "what kind of thing it is".

#-

# There are also *abstract types* that serve to group other types together.

supertype(Int)
supertype(Signed)
supertype(Integer)
supertype(Real)
supertype(Number)

# `isa` queries whether a value belongs to a type.
# `<:` queries whethe one type is a subtype (subset) of another

2 isa Int
2 isa Real
2 isa String
Int <: Integer

# `::` asserts that a value has a certain type, or else throws an error

2::Int
2::String

#-

# ## Defining types

struct Point
    x::Float64
    y::Float64
end

Point(1,2)

#-

# Types can be "called" and have methods added like functions.
# When used that way they are sometimes called *constructors*.

Point() = Point(0,0)

Point()

Int(2.0)

Int(2.1)

# ### Abstract types

abstract type PointLike end

struct Point2 <: PointLike
    x::Float64
    y::Float64
end

# ### Mutability
#
# Instances of `Point` cannot be modified after construction:

p = Point(1,2)
p.x = 0

# Use `mutable struct` to declare a mutable type:

mutable struct MutablePoint
    x  # note: default field type is `Any`
    y
end

# ### Primitive types
#
# Primitive types are "scalar"-like types whose representation is just a string
# of bits. They have no fields.
# These are not defined often in user code, but are needed in some cases.

primitive type MyInt32 <: Signed 32 end

# ## Type parameters
#
# Some types actually define a family of related types.
# The following is a concrete type, but the *parameters* in curly
# braces indicate that it is one of many possible related types:

Array{Float64,2}

# What is `Array` by itself?

typeof(Array)

# `Array` is an *iterated union* of dense, in-memory arrays over all
# element types and numbers of dimensions.

Array == Array{T,N} where N where T

# The curly braces substitute values for these variables:

Array{Int} == Array{Int,N} where N

Array{Int} <: Array

(Array{T,3} where T) <: Array

[1,2] isa Array{Int}

[1,2] isa Array{T,1} where T

[1,2] isa Array{T,2} where T

# `Vector` is an alias for 1-d Arrays:

Vector == Array{T,1} where T

#-

# Types with different parameters are just different, and have no subtype
# relationship. This is called *invariance*.

Array{Int} <: Array{Real}

# This is surprising at first, but makes sense if you think about memory
# representations.
# `Array{Real}` can hold any kind of real number so it needs to be an array of
# pointers (or some other structure containing type information for each
# element).
# `Array{Int}` only has `Int`s, so it can use an efficient "flat"
# representation.

# ### Defining types with parameters

struct GenericPoint{T<:Real}
    x::T
    y::T
end

GenericPoint(1,2)

GenericPoint(1,2.0)

GenericPoint(1.0,2.0)

GenericPoint{Float64}(1,2)

# Constructors can be defined for these too, using syntax that looks like
# how they are called:

function GenericPoint{T}() where T
    return GenericPoint{T}(0,0)
end

GenericPoint{Int}()

GenericPoint()

GenericPoint() = GenericPoint{Float64}()

GenericPoint()

# ## Example: `OneHotVector`
#
# Often used in machine learning, a "one hot" vector is a vector of all zeros,
# except for a single `1` entry.
# Representing it as a standard vector is memory-inefficient, so it cries out
# for a special implementation.

struct OneHotVector <: AbstractVector{Int}
    idx::Int
    len::Int
end

Base.size(v::OneHotVector) = (v.len,)

Base.getindex(v::OneHotVector, i::Integer) = Int(i == v.idx)

OneHotVector(3, 10)

A = rand(5,5)

A * OneHotVector(3, 5)

# ### Exercise: make it generic over element type

# ## Tuple types
#
# Tuples are very much like structs, but their types a bit special.
# 1. They have fields, but no names, only indices.
# 2. They are always immutable.
# 3. They can have an arbitrary number of parameters/fields.
# 4. They are *covariant*.

t = ("a", "b")
typeof(t)

# Field access is syntax for calls to `getfield` and `setfield!`

getfield(t, 2)

# For tuples `getindex` == `getfield`:

t[1]

getfield(p, :x)
getfield(p, 2)

Tuple{String} <: Tuple{Any}

# Remember: not true for `Vector{String} <: Vector{Any}`!

# ## Summary of kinds of types
#
# 1. Concrete types: struct, mutable struct, tuple, primitive
# 2. abstract types
# 3. UnionAll types, `Array{T} where T`
# 4. Union types, `Union{Int, Missing}`
# 5. The empty, or "bottom" type `Union{}`

# ## The types of functions
#
# How do functions fit into this scheme?
# Somewhat non-obviously, they are structs!
# In Julia, every "function" has its own struct type.

#-

f(x) = 2x

# Is equivalent to

struct typeof_f
end

(::typeof_f)(x) = 2x

const f = typeof_f()

f(3)

#-

# A simple, top-level function is a struct with 0 fields.

typeof(+)
typeof(cos)

# *Closures*, or inner functions, that capture surrounding
# variables, keep those variables as fields.

adder(x) = y->x+y

a = adder(100)

typeof(a)
dump(a)

a.x


"""
`Basics.integrate()`
    ... integrates a one- or two-dimensional function by different numerical methods, either on a given grid or for a general function.

  + `("function: on radial grid, Newton-Cotes", F::Array{Float64,1}, grid::Radial.Grid)`
    ... to integrate the (radial) function F over the given (radial) grid, int_0^infinity dr F(r) by using a 5-point Newton-Cotes
        integration formula; a value::Float64 is returned.

  + `("function: on radial grid, Simpson rule", F::Array{Float64,1}, grid::Radial.Grid)`
    ... to integrate the (radial) function F over the given (radial) grid, int_0^infinity dr F(r) by using Simpson's rule;
        a value::Float64 is returned.

  + `("function: on radial grid, trapez rule", F::Array{Float64,1}, grid::Radial.Grid)`
    ... to integrate the (radial) function F over the given (radial) grid, int_0^infinity dr F(r), by using a simple trapez rule;
        a value::Float64 is returned.
"""
function Basics.integrate(sa::String, F::Array{Float64,1}, grid::Radial.Grid)
    if       sa == "function: on radial grid, Newton-Cotes"     wa = integrateOnGridNewtonCotes(F, grid)
    elseif   sa == "function: on radial grid, Simpson rule"     wa = integrateOnGridSimpsonRule(F, grid)
    elseif   sa == "function: on radial grid, trapez rule"      wa = integrateOnGridTrapezRule( F, grid)
    else     error("Unsupported keystring = $sa.")
    end

end

"""
`Basics.integrateOnGridNewtonCotes()`  ... integrates by using a 5-point Newton-Cotes integration formula; see Basic.integrate().
"""
function Basics.integrateOnGridNewtonCotes(F::Array{Float64,1}, grid::Radial.Grid)
    coefficients = Array{Float64}([2 * 7, 32, 12, 32, 7]) * 2 / 45 * grid.h
    n = size(coefficients, 1)

    result = 0.;    i0     = 1

    while  abs(F[i0]) == 0
        i0 += 1
    end
    gamma  = log(F[i0]/F[i0+1] * grid.rp[i0+1]/grid.rp[i0]) / log(grid.r[i0]/grid.r[i0+1])
    result = 1/(gamma + 1) * grid.r[i0] * F[i0] / grid.rp[i0]

    result += F[i0] * coefficients[n]
    for i = i0+1:size(F, 1)
        result += coefficients[(i - i0) % (n - 1) + 1] * F[i]
    end

    if (size(F, 1) - i0) % (n - 1) + 1 == 1
        result -= coefficients[n] * F[size(F, 1)]
    end

    return( result )
end

"""
`Basics.integrateOnGridSimpsonRule()`  ... integrates by using Simpson's rule; see JAC.Basics.integrate().
"""
function Basics.integrateOnGridSimpsonRule(F::Array{Float64,1}, grid::Radial.Grid)
    coefficients = Array{Float64}([2 * 1, 4, 1]) / 3 * grid.h
    n = size(coefficients, 1)

    result = 0.
    result += F[1] * coefficients[n]
    for i = 2:size(F, 1)
        result += coefficients[(i - 1) % (n - 1) + 1] * F[i]
    end

    if (size(F, 1) - 1) % (n - 1) + 1 == 1
        result -= coefficients[n] * F[size(F, 1)]
    end

    return( result )
end

"""
`Basics.integrateOnGridTrapezRule()`  ... integrates by using a simple trapez rule; see JAC.Basics.integrate().
"""
function Basics.integrateOnGridTrapezRule(F::Array{Float64,1}, grid::Radial.Grid)
    coefficients = Array{Float64}([2 * 1, 1]) / 2 * grid.h
    n = size(coefficients, 1)

    result = 0.

    result += F[1] * coefficients[n]
    for i = 2:size(F, 1)
        result += coefficients[(i - 1) % (n - 1) + 1] * F[i]
    end

    if (size(F, 1) - 1) % (n - 1) + 1 == 1
        result -= coefficients[n] * F[size(F, 1)]
    end

    return( result )
end


include("../src/fdfd_local_field.jl")
using Plots
pyplot()

### reference simulation
refractive_indexes = ones(3) * 1.45
ps = [0.75, 0.5, 0.66, 0.33, 0.66, 0.5, 0.75, 0.5, 0.66, 0.33]
x, y, Ez, dpml, dsource, resolution = simulation_hole_layers_unit_cell(ps, refractive_indexes=refractive_indexes, frequency=0.25)

heatmap(real.(ϵ_hole_layers(x, y, ps, refractive_indexes=refractive_indexes)))
gui()

heatmap(real.(Ez), c=ColorGradient([:red, :white, :blue]))
gui()

### structure simulation
refractive_indexes = Float64[1.0, 1.0, 1.45]
x, y, Ez, dpml, dsource, resolution = simulation_hole_layers_unit_cell(ps, refractive_indexes=refractive_indexes)

heatmap(real.(ϵ_hole_layers(x, y, ps, refractive_indexes=refractive_indexes)))
gui()

heatmap(real.(Ez), c=ColorGradient([:red, :white, :blue]))
gui()


"""
psdplot(signal; fs = 1, maxf::Real = -1, method = "welch", order = 20)

Conviniently plots the power spectral density using one of two methods currently implemented in DSP.jl or a custom method using an AR model.
Futher plot parameters can be passed as keyword arguments.

# Args:

* 'signal::Vector': Data Vector containing te signal
* 'fs': Sample rate in Hz
* 'maxf': maximum frequency that shal be displayed
* 'method': Method for esitmation, either 'welch', 'ar' or 'fft'

"""
psdplot

@userplot PSDPlot
# psdplot(signal, fs; maxf = nothing, estimator = welch_pgram, kwargs...)
@recipe function f(p::PSDPlot; fs = 1, maxf::Real = -1, method = "welch", order = 20)
    y = (p.args[1]) .- mean(p.args[1]) # detrend/ remove dc component
    # Methods for estimation
    if method == "ar"
        psd = pyulear(vec(y), order, fs = fs)
    elseif method == "fft"
        psd = periodogram(vec(y), fs = fs)
        # psd.power[1] = psd.power[2]
    elseif method == "welch"
        psd = welch_pgram(vec(y), fs = fs)
    else
        throw(ArgumentError("Invalid method choosen. Supported are 'welch', 'fft' and 'yule'."))
    end

    # plot porps
    yaxis := :log
    title := "PSD"
    # legend := false
    xaxis := "f [Hz]"
    if maxf > 0
        xlims := (0, maxf)
    else
        xlims := (minimum(psd.freq), maximum(psd.freq))
    end
    @series begin
        psd.freq, psd.power
    end
end

# Get PSD from AR process with variance σ² (Burg 1967)
function S(f::Number, σ::Number, φ::Vector{<:Number})
    indices = 1:length(φ)
    sumXes = sum(φ .* exp.(-2im * pi * f * indices))
    return 2σ^2 / abs2(1 + (sumXes))
end

"""
pyulear(x::Vector{<:Number}, order::Int)

Matlab Äquivalent von pyulear. Bestimme das Spektrum parametrisch mittels Yule-Walker-Gleichungen

# Args
* 'x::Vector{<:Number}': Zeitreihe aus der das Spektrum bestimmt werden soll
* 'order::Int': Ordnung des zugrundeliegenden AR Modells

# Keyword Args
* 'fs::Real = 1' Sampling frequency to get the power at the respective frequencies
"""
function pyulear(x::Vector{<:Real}, order::Int; fs::Real = 1, Δf = 0.001)
    f = 0:Δf:0.5
    fReal = range(0, stop = fs/2, length = length(f))
    # fit ar model
    (a,g) = lpc(x, order)
    # estimate spectrum
    Sar(x) = S(x, sqrt(g), a)
    power = Sar.(f)
    return DSP.Periodograms.Periodogram(power, fReal)
end

"""
getSpectralComponent(signal::Vector{<:Real}, f_range::Tuple{<:Real, <:Real}; method::String = "ar", Δf::Real = 0.001, order::Int = 10)

Get the spectral power of a certain frequency range, in natural frequencies, using different methods ['ar', 'fft', 'welch'].

# Args:

* 'signal::Vector{<:Real}': The signal that is to be decomposed.
* 'f_range::Tuple{<:Real, <:Real}': The frequency range of interest in natural frequencies [0, 0.5].

# Kwargs

* 'method::String = "ar"': Method used for psd  estimation. Choose from 'ar', 'fft' and 'welch'.
* 'Δf::Real = 0.001': Only relevant for method ar, definfes distance of the grid of the underlying integration
* 'order::Int = 10': Only relevant for method ar, defines the order of the ar model

# Examples

```julia
julia> getSpectralComponent(signal, (0.1, 0.4))
```
"""
function getSpectralComponent(signal::Vector{<:Real}, f_range::Tuple{<:Real, <:Real}; method::String = "ar", Δf::Real = 0.001, order::Int = 10)
    0.0 <= f_range[1] < f_range[2] <= 0.5 || throw(DomainError("Start of range ($(f_range[1])) has to be smaller than end ($(f_range[2])) And in range of 0:0.5."))
    if method == "ar"
        psd = pyulear(signal, order, Δf = Δf)
    elseif method == "fft"
        psd = periodogram(signal)
        Δf = psd.freq[2] - psd.freq[1]
    elseif method == "welch"
        psd = welch_pgram(signal)
        Δf = psd.freq[2] - psd.freq[1]
    else
        throw(DomainError("Valid methods are: 'ar', 'fft' & 'welch'."))
    end
    pow = sum(Δf .* psd.power[findall(x -> f_range[1] <= x <= f_range[2], psd.freq)])
    return pow
end

"""
getSpectralComponentAR(ar::Union{ControlSystems.TransferFunction, AbstractVector}, f_range::Tuple{<:Real, <:Real}; σ::Real = 1,  Δf::Real = 0.001)

Get the spectral power of a certain frequency range, in natural frequencies, using different methods ['ar', 'fft', 'welch'].

# Args:

* 'ar::Union{ControlSystems.TransferFunction, AbstractVector}': Fitted AR model that schal be decomposed .
* 'f_range::Tuple{<:Real, <:Real}': The frequency range of interest in natural frequencies [0, 0.5].

# Kwargs

* 'Δf::Real = 0.001': Only relevant for method ar, definfes distance of the grid of the underlying integration
* 'σ::Real = 1': sqrt(innovation variance / prediction error)

# Examples

```julia
julia> getSpectralComponent(ar(iddata(signal), 10), (0.1, 0.4))
```
"""
function getSpectralComponentAR(ar::Union{ControlSystems.TransferFunction, AbstractVector}, f_range::Tuple{<:Real, <:Real}; σ::Real = 1,  Δf::Real = 0.001)
    0.0 <= f_range[1] < f_range[2] <= 0.5 || throw(DomainError("Start of range ($(f_range[1])) has to be smaller than end ($(f_range[2])) And in range of 0:0.5."))
    if typeof(ar) <: ControlSystems.TransferFunction
        arcoef = denvec(ar)[1]
        arcoef = arcoef[2:end]
    elseif typeof(ar) <: AbstractVector
        arcoef = ar
    end
    f = 0:Δf:0.5
    Sar(x) = S(x, σ, arcoef)
    power = Sar.(f)
    pow = sum(Δf .* power[findall(x -> f_range[1] <= x <= f_range[2], f)])
    return pow
end


struct BNet
    number_of_nodes::Int
    weights::Array{Float64,2}
    thresholds::Vector{Float64}
    BNet(ws, ts) = new(size(ws,1), ws, ts)
end

load(::Type{BNet}, wfile, tfile) = BNet(readdlm(wfile), vec(readdlm(tfile)))

Base.size(net) = net.number_of_nodes

volume(net) = 1 << size(net)

function fire!(net, state, next_state)
    temp = net.weights * state
    for i in 1:length(state)
        next_state[i] = if temp[i] > net.thresholds[i]
            1
        elseif temp[i] < net.thresholds[i]
            0
        else
            state[i]
        end
    end
    next_state
end

fire(net, state) = fire!(net, state, similar(state))

function encode(state)
    const N = length(state)
    code = 0
    for i in 1:N
        code = 2code + state[N - i + 1]
    end
    code + 1
end

function transitions(net)
    const N, V = size(net), volume(net)
    trans = Array{Int}(V)
    state = zeros(Int, N)
    next_state = fire(net, state)

    code, next_code = 1, encode(next_state)
    trans[code] = next_code

    while code < V
        for i in 1:N
            if state[i] == 0
                state[i] += 1
                state[1:i-1] = 0
                break
            end
        end
        code += 1

        fire!(net, state, next_state)
        next_code = encode(next_state)
        trans[code] = next_code
    end

    trans
end

function attractors(net)
    const V = volume(net)
    const trans = transitions(net)
    seen = zeros(Bool, V)
    basins = fill(-1, V)
    basin_number = 0

    attractors = Vector{Int}[]
    state_stack = Int[]
    cycle = Int[]
    for initial_state in trans
        if seen[initial_state]
            continue
        end

        state = initial_state
        in_cycle = false
        terminus = next_state = trans[state]
        seen[state] = true

        while !seen[next_state]
            push!(state_stack, state)
            state = next_state
            terminus = next_state = trans[state]
            seen[state] = true
        end

        basin = basins[next_state]
        if basins[next_state] == -1
            basin = basin_number
            basin_number += 1
            push!(cycle, state)
            in_cycle = (terminus != state)
        end

        basins[state] = basin
        while !isempty(state_stack)
            state = pop!(state_stack)
            basins[state] = basin
            if in_cycle
                push!(cycle, state)
                in_cycle = (terminus != state)
            end
        end

        if !isempty(cycle)
            push!(attractors, cycle)
            cycle = Int[]
        end
    end
    attractors
end

function main(wfile, tfile)
    const net = load(BNet, wfile, tfile)

    duration = @elapsed begin
        attrs = attractors(net)
    end
    println("First Invocation:  $(duration)s\n")

    duration = @elapsed begin
        attrs = attractors(net)
    end
    w = ceil(log2(length(attrs)))
    for (i, attr) in enumerate(attrs)
        print("Attractor $i: ")
        for state in attr
            print("$(state-1) ")
        end
        println()
    end
    println("Second Invocation: $(duration)s")
end

main(ARGS[1], ARGS[2])


using ReverseDiff
using CatViews

x = [-1. -2; -1 -2]
y = [1.; 2]
w = [2.; 3]
b = [-3.]

type Vars
    lookup::Dict
    grad_lookup::Dict
    flat
    grad_flat

    Vars() = Vars(Dict(), Dict(), nothing, nothing)
end
getindex(vars::Vars, varname::String) = vars.lookup[varname]
setindex!(vars::Vars, varname::String, var) = vars.lookup[varname] = var
function compile!(vars::Vars)
    for (k, v) in vars.lookup
        vars.grad_lookup[k] = similar(v)
    end
    vars.flat = CatView(map(x -> view(x, :), values(vars.lookup)))
    vars.grad_flat = CatView(map(x -> view(x, :), values(vars.grad_lookup)))
end

vars = Vars()
vars[:w] = w
vars[:b] = b
compile!(vars)

ReverseDiff.@forward sigmoid(x) = 1. ./ (1. + exp(-x))
mse(a, y) = mean((a - y) .* (a - y))
#= mse(a, y) = mean((a - y) .^ 2) =#
f1(x, y, vars::Vars) = mse(sigmoid(x*vars[:w] .+ vars[:b]), y)

out = f1(x, y, vars)
println(out)

inputs = (x, y, vars)
result = map(similar, inputs)
result_view = CatView(map(x -> view(x, :), result))
f1! = ReverseDiff.compile_gradient(f1, result)

@time f1!(result, inputs)
@time f1!(result, inputs)

for r in result
    println(r)
end

println(result_view)

#= Expected output =#
#=  =#
#= [array([[ -3.34017280e-05,  -5.01025919e-05], =#
#=        [ -6.68040138e-05,  -1.00206021e-04]]), array([[ 0.9999833], =#
#=        [ 1.9999833]]), array([[  5.01028709e-05], =#
#=        [  1.00205742e-04]]), array([ -5.01028709e-05])] =#


module Normalizers

using StatsBase
using Dates
using DataFrames
using Statistics
using MultivariateStats
using Random

using AutoMLPipeline.AbsTypes
using AutoMLPipeline.Utils
import AutoMLPipeline.AbsTypes: fit!, transform!

export fit!,transform!
export Normalizer

"""
    Normalizer(Dict(
       :method => :zscore
    ))

Transforms continuous features into normalized form such as zscore, unitrange, square-root, log, pca, ppca
with parameter:

- `:method` => `:zscore` or `:unitrange` or `:sqrt` or `:log` or `pca` or `ppca` or `fa`
- :zscore => standard z-score with centering and scaling
- :unitrange => unit range normalization with centering and scaling
- :sqrt => square-root transform
- :pca => principal component analysis transform
- :ppca => probabilistic pca
- :fa => factor analysis
- :log => log transform

Example:

    function generatedf()
        Random.seed!(123)
        gdate = DateTime(2014,1,1):Dates.Minute(15):DateTime(2016,1,1)
        gval1 = rand(length(gdate))
        gval2 = rand(length(gdate))
        gval3 = rand(length(gdate))
        X = DataFrame(Date=gdate,Value1=gval1,Value2=gval2,Value3=gval3)
        X
    end

    X = generatedf()
    norm = Normalizer(Dict(:method => :zscore))
    fit!(norm,X)
    res=transform!(norm,X)

Implements: `fit!`, `transform!`
"""
mutable struct Normalizer <: Transformer
  model
  args

  function Normalizer(args=Dict())
    default_args = Dict(
        :method => :zscore
    )
    new(nothing, mergedict(default_args, args))
  end
end

function Normalizer(st::Symbol)
  Normalizer(Dict(:method=>st))
end

"""
    fit!(st::Statifier, features::T, labels::Vector=[]) where {T<:Union{Vector,Matrix,DataFrame}}

Validate argument features other than dates are continuous.
"""
function fit!(norm::Normalizer, features::DataFrame, labels::Vector=[])
  # check features are in correct format and no categorical values
  (infer_eltype(features[:,1]) <: DateTime && infer_eltype(Matrix(features[:,2:end])) <: Real) ||
    (infer_eltype(Matrix(features)) <: Real) ||
    error("Normalizer.fit!: make sure features are purely float values or float values with Date on first column")
  norm.model = norm.args
end

"""
    transform!(norm::Normalizer, features::T) where {T<:Union{Vector,Matrix,DataFrame}}

Compute statistics.
"""
function transform!(norm::Normalizer, pfeatures::DataFrame)
  pfeatures != DataFrame() || return DataFrame()
  res = Array{Float64,2}(undef,0,0)
  if (infer_eltype(pfeatures[:,1]) <: DateTime && infer_eltype(Matrix(pfeatures[:,2:end])) <: Real)
    res = processnumeric(norm,Matrix{Float64}(pfeatures[:,2:end])) |> DataFrame
  elseif infer_eltype(Matrix(pfeatures)) <: Real
    features = pfeatures |> Array{Float64}
    res = processnumeric(norm,features) |> DataFrame
  else
    error("Normalizer.transform!: make sure features are purely float values or float values with Date on first column")
  end
  res
end

function processnumeric(norm::Normalizer,features::Matrix)
  if norm.args[:method] == :zscore
    ztransform(features)
  elseif norm.args[:method] == :unitrange
    unitrtransform(features)
  elseif norm.args[:method] == :pca
    pca(features)
  elseif norm.args[:method] == :ppca
    ppca(features)
  elseif norm.args[:method] == :ica
    ica(features)
  elseif norm.args[:method] == :fa
    fa(features)
  elseif norm.args[:method] == :sqrt
    sqrtf(features)
  elseif norm.args[:method] == :log
    logf(features)
  else
    error("arg's :method is mapped to unknown keyword")
  end
end

# apply sqrt transform
function sqrtf(X)
  sqrt.(X)
end

# apply log transform
function logf(X)
  log.(X)
end

# apply z-score transform
function ztransform(X)
  xp = X' |> collect |> Matrix{Float64}
  fit(ZScoreTransform, xp,dims=2; center=true, scale=true) |> dt -> StatsBase.transform(dt,xp)' |> collect
end

# unit-range
function unitrtransform(X)
  xp = X' |> collect |> Matrix{Float64}
  fit(UnitRangeTransform,xp,dims=2) |> dt -> StatsBase.transform(dt,xp)' |> collect
end

# pca
function pca(X)
  xp = X' |> collect |> Matrix{Float64}
  m = fit(PCA,xp)
  transform(m,xp)' |> collect
end

function ica(X,kk::Int=0)
  k = kk
  if k == 0
    k = size(X)[2]
  end
  xp = X' |> collect |> Matrix{Float64}
  m = fit(ICA,xp,k)
  transform(m,xp)' |> collect
end

# ppca
function ppca(X)
  xp = X' |> collect |> Matrix{Float64}
  m = fit(PPCA,xp)
  transform(m,xp)' |> collect
end

# fa
function fa(X)
  xp = X' |> collect |> Matrix{Float64}
  m = fit(FactorAnalysis,xp)
  transform(m,xp)' |> collect
end

function generatedf()
    Random.seed!(123)
    gdate = DateTime(2014,1,1):Dates.Minute(15):DateTime(2016,1,1)
    gval1 = rand(length(gdate))
    gval2 = rand(length(gdate))
    gval3 = rand(length(gdate))
    X = DataFrame(Date=gdate,Value1=gval1,Value2=gval2,Value3=gval3)
    X
end

function testnormalizer()
  X = generatedf()
  norm = Normalizer(Dict(:method => :zscore))
  fit!(norm,X)
  res=transform!(norm,X)
  @assert isapprox(mean(res[:,1]),0.0,atol=1e-8)
  @assert isapprox(mean(res[:,2]),0.0,atol=1e-8)
  @assert isapprox(std(res[:,1]),1.0,atol=1e-8)
  @assert isapprox(std(res[:,2]),1.0,atol=1e-8)
  norm = Normalizer(Dict(:method => :unitrange))
  fit!(norm,X)
  res=transform!(norm,X)
  @assert isapprox(minimum(res[:,1]),0.0,atol=1e-8)
  @assert isapprox(minimum(res[:,2]),0.0,atol=1e-8)
  @assert isapprox(maximum(res[:,1]),1.0,atol=1e-8)
  @assert isapprox(maximum(res[:,2]),1.0,atol=1e-8)
  norm = Normalizer(Dict(:method => :pca))
  fit!(norm,X)
  res=transform!(norm,X)
  @assert isapprox(std(res[:,1]),0.28996,atol=1e-2)
  norm = Normalizer(Dict(:method => :fa))
  fit!(norm,X)
  res = transform!(norm,X)
  @assert isapprox(std(res[:,1]),0.81670,atol=1e-2)
  norm = Normalizer(Dict(:method => :ppca))
  fit!(norm,X)
  res = transform!(norm,X)
  @assert isapprox(std(res[:,1]),0.00408,atol=1e-2)
  norm = Normalizer(Dict(:method => :log))
  fit!(norm,X)
  res = transform!(norm,X)
  @assert isapprox(std(res[:,1]),0.99941,atol=1e-2)
  norm = Normalizer(Dict(:method => :sqrt))
  fit!(norm,X)
  res = transform!(norm,X)
  @assert isapprox(std(res[:,1]),0.2355,atol=1e-2)
end

end


# %%
using Tmatrix

# inputs
rx, rz = 1e-6, 1.3e-6
n_max = 3
k1_r = 1e7;
k1_i = 0.0;
k2_r = 1.5e7;
k2_i = 1e3;

k1 = Complex(k1_r, k1_i)
k2 = Complex(k2_r, k2_i)

BigFloat_precision_list = 2 .^ [6, 7, 8, 9]

println("***************************************************************")
println("***************************************************************")
println("Calculating T-matrix using Complex numbers")
println("***************************************************************")
println("***************************************************************")

println("using Float64 =====================================")
println("No symmetry exploited:")
@time T_no_symmetry_float_64 = calculate_Tmatrix_for_spheroid(
    rx,
    rz,
    n_max,
    k1,
    k2,
    rotationally_symmetric = false,
    symmetric_about_plane_perpendicular_z = false,
);

# %%
println("exploiting symmetry:")
@time T_with_symmetry_float_64 = calculate_Tmatrix_for_spheroid(
    rx,
    rz,
    n_max,
    k1,
    k2,
    rotationally_symmetric = true,
    symmetric_about_plane_perpendicular_z = true,
);

println("using BigFloat =====================================")
println("exploiting symmetry:")
T_with_symmetry_big_float = []
for precision_here in BigFloat_precision_list
    println("using precision $precision_here")
    @time T = calculate_Tmatrix_for_spheroid(
        rx,
        rz,
        n_max,
        k1,
        k2,
        rotationally_symmetric = true,
        symmetric_about_plane_perpendicular_z = true,
        BigFloat_precision = precision_here,
    )
    append!(T_with_symmetry_big_float, T)
    println()
end

"""
n_array = collect(1:10)
elapsed_time = zeros(size(n_array))
for n = 1:length(n_array)
elapsed_time[n] = @elapsed calculate_Tmatrix_for_spheroid(rx, rz, n_max, k1, k2, rotationally_symmetric=true, symmetric_about_plane_perpendicular_z=true, BigFloat_precision=8)
end
"""
println()
println()
println()
println("***************************************************************")
println("***************************************************************")
println("Calculating T-matrix using real numbers")
println("***************************************************************")
println("***************************************************************")

println("using Float64 =====================================")
println("No symmetry exploited:")
@time T_no_symmetry_float_64_SeparateRealImag =
    calculate_Tmatrix_for_spheroid_SeparateRealImag(
        rx,
        rz,
        n_max,
        k1_r,
        k1_i,
        k2_r,
        k2_i;
        rotationally_symmetric = false,
        symmetric_about_plane_perpendicular_z = false,
    );

println("exploiting symmetry:")
@time T_with_symmetry_float_64_SeparateRealImag =
    calculate_Tmatrix_for_spheroid_SeparateRealImag(
        rx,
        rz,
        n_max,
        k1_r,
        k1_i,
        k2_r,
        k2_i;
        rotationally_symmetric = true,
        symmetric_about_plane_perpendicular_z = true,
    );

println("using BigFloat =====================================")
println("exploiting symmetry:")
T_with_symmetry_big_float_SeparateRealImag = []
for precision_here in BigFloat_precision_list
    println("using precision $precision_here")
    @time T = calculate_Tmatrix_for_spheroid_SeparateRealImag(
        rx,
        rz,
        n_max,
        k1_r,
        k1_i,
        k2_r,
        k2_i;
        rotationally_symmetric = true,
        symmetric_about_plane_perpendicular_z = true,
        BigFloat_precision = precision_here,
    )
    append!(T_with_symmetry_big_float_SeparateRealImag, T)
end


"""
    Module EntropyStableEuler

Includes general math tools
"""

module EntropyStableEuler

const γ=1.4
export logmean
export u_vfun, v_ufun, betafun, pfun
export euler_fluxes, wavespeed
export vortex, primitive_to_conservative

include("./logmean.jl")
include("./euler_fluxes.jl")
include("./euler_variables.jl")

# 2D isentropic vortex solution for testing. assumes domain around [0,20]x[-5,5]
function vortex(x,y,t,γ=1.4)

    x0 = 5
    y0 = 0
    beta = 5
    r2 = @. (x-x0-t)^2 + (y-y0)^2

    u = @. 1 - beta*exp(1.0-r2)*(y-y0)/(2.0*pi)
    v = @. beta*exp(1.0-r2)*(x-x0-t)/(2.0*pi)
    rho = @. 1.0 - (1.0/(8.0*γ*pi^2))*(γ-1.0)/2.0*(beta*exp(1.0-r2))^2
    rho = @. rho^(1/(γ-1))
    p = @. rho^γ

    return (rho, u, v, p)
end

end


#Structure for MD Simulations
module MD_Structure

export MD, CalculateForces!

using Parameters
using Random

@with_kw struct Params
    n::Int64;
    ρ::Float64;
    T::Float64;
    L::Float64;
end

mutable struct MD
    p::Params;
    n::Int64;
    r::Matrix{Float64};
    v::Matrix{Float64};
    f::Matrix{Float64};
    kinetic_energy::Float64;
    potential_energy::Float64;
    total_energy::Float64;
    function MD()
        n_ = 1000;
        ρ_ = 0.8;
        T_ = 1.0;
        L_ = (n_/ρ_)^(1.0/3.0);
        p_ = Params( n = n_, ρ = ρ_, T = T_, L = L_);#Default constructor: n = 1000, ρ = 0.8, T = 1.0.
        r_ = Matrix{Float64}(undef,3,n_);#Column-major order
        v_ = Matrix{Float64}(undef,3,n_);
        f_ = Matrix{Float64}(undef,3,n_);
        this = new(p_,p_.n,r_,v_,f_,0.0,0.0,0.0);
        Initalize!(this);
        return this;
    end
    function MD(params)
        p_ = params;#Call it with parameters defined elsewhere
        n_ = p_.n;
        r_ = Matrix{Float64}(undef,3,n_);
        v_ = Matrix{Float64}(undef,3,n_);
        f_ = Matrix{Float64}(undef,3,n_);
        this = new(p_,p_.n,r_,v_,f_,0.0,0.0,0.0);
        Initalize!(this);
        return this;
    end
end

function Initalize!(md::MD)
    @unpack n,L,T = md.p;
    nside = round(n^(1.0/3.0));
    a = L/nside;
    xhat = [1, 0, 0];#Column vectors for cardinal directions
    yhat = [0, 1, 0];
    zhat = [0, 0, 1];
    i = 1;
    for nx in 0:(nside-1)
        for ny in 0:(nside-1)
            for nz in 0:(nside-1)
                md.r[:,i] =  a*( (nx+0.5)*xhat + (ny+0.5)*yhat + (nz+ 0.5)*zhat );#Build lattice
                i += 1;
            end
        end
    end
    md.v = randn(size(md.v));#Choose velocities from a normal distribution.
    for k in 1:3
        sumv = 0.0;
        for i in 1:n
            sumv += md.v[k,i];
        end
        md.v[k,:] = md.v[k,:] .- sumv/n;#Subtract off center-of-mass
    end
    Thermalize!(md);
end

function KineticEnergy(md::MD)
    return 24.0*sum(md.v.*md.v);
end

function Temperature(md::MD)
    K = KineticEnergy(md);
    T = 2.0*K/(3.0*(md.n-1));
    return T;
end

function Thermalize!(md::MD)
    Tcurrent = Temperature(md);
    σ2 = md.p.T/Tcurrent;
    md.v .*= sqrt(σ2);
end

function CalculateForces!(md::MD)
    md.f .= 0.0; fij = 0.0; vij = 0.0; d2 = 0.0;
    @unpack n,L = md.p;
    Linv = 1.0/L;
    d2cut = 2.5*2.5;
    v = r -> 4.0*( r^(-12) - r^(-6) );
    vcut = v(2.5);
    ri = zeros(3); rj = zeros(3); rij = zeros(3);
    function potential_and_force(d2)
        d2_inv = 1.0/d2;
        d6_inv = d2_inv*d2_inv*d2_inv;
        Fij = 48.0*(d6_inv*d6_inv - 0.5*d6_inv)*d2_inv;
        Vij = 4.0*(d6_inv*d6_inv - d6_inv) - vcut;
        return Fij, Vij;
    end
    md.potential_energy = 0.0;
    for i in 1:(n-1)
        ri = md.r[:,i];
        for j in (i+1):n
            rj = md.r[:,j];
            rij .= ri .- rj;
            rij .-= L.*round.(rij*Linv);
            d2 = sum(rij.*rij);
            if d2 <= d2cut;
                fij,vij = potential_and_force(d2);
                md.f[:,i] += fij.*rij;
                md.f[:,j] -= fij.*rij;
                md.potential_energy += vij;
            end
        end
    end
end

end


using ApproxFun, Base.Test

c = rand(1000)
x=rand(10000)
f=Fun(c,Chebyshev)
y=f(x)
y=f(x)

@time y=f(x)
println("Clenshaw large coeffs, many points: Time should be ~0.024")
# 0.012482274  with unsafe_view
# 0.024306262 with inbounds

y=f(.1)
y=f(.1)
y=f(.1)

@time y=f(.1);
println("Clenshaw large coeffs, 1 point: Time should be ~6e-6")

# @time is 8.853e-6 seconds

f=Fun(exp)
x=sample(f,100000)
x=sample(f,100000)
@time x=sample(f,100000)
println("Sample: Time should be ~0.25")
# 0.213793292 with unsafe_view
# 0.268162181 with inbounds

f=Fun(x->cos(x),20)
roots(f)
roots(f)
@time for k=1:100
    roots(f)
end
println("Small roots: Time should be ~0.015")

f=Fun(x->cos(1000x),1000)
roots(f)
roots(f)
@time roots(f)
println("Roots: Time should be ~0.13")


#=
Provides a type called LQ for solving linear quadratic control
problems.

@author : Spencer Lyon <spencer.lyon@nyu.edu>
@author : Zac Cranko <zaccranko@gmail.com>

@date : 2014-07-05

References
----------

https://lectures.quantecon.org/jl/lqcontrol.html

=#

@doc doc"""
Linear quadratic optimal control of either infinite or finite horizon

The infinite horizon problem can be written

```math
\min \mathbb{E} \sum_{t=0}^{\infty} \beta^t r(x_t, u_t)
```

with

```math
r(x_t, u_t) := x_t' R x_t + u_t' Q u_t + 2 u_t' N x_t
```

The finite horizon form is

```math
\min \mathbb{E} \sum_{t=0}^{T-1} \beta^t r(x_t, u_t) + \beta^T x_T' R_f x_T
```

Both are minimized subject to the law of motion

```math
x_{t+1} = A x_t + B u_t + C w_{t+1}
```

Here ``x`` is n x 1, ``u`` is k x 1, ``w`` is j x 1 and the matrices are conformable for
these dimensions.  The sequence ``{w_t}`` is assumed to be white noise, with zero
mean and ``\mathbb{E} w_t w_t' = I``, the j x j identity.

For this model, the time ``t`` value (i.e., cost-to-go) function ``V_t`` takes the form

```math
x' P_T x + d_T
```

and the optimal policy is of the form ``u_T = -F_T x_T``.  In the infinite horizon
case, ``V, P, d`` and ``F`` are all stationary.

##### Fields

- `Q::ScalarOrArray` : `k x k` payoff coefficient for control variable u. Must be
  symmetric and nonnegative definite
- `R::ScalarOrArray` : `n x n` payoff coefficient matrix for state variable x.
  Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : `n x n` coefficient on state in state transition
- `B::ScalarOrArray` : `n x k` coefficient on control in state transition
- `C::ScalarOrArray` : `n x j` coefficient on random shock in state transition
- `N::ScalarOrArray` : `k x n` cross product in payoff equation
- `bet::Real` : Discount factor in `[0, 1]`
- `capT::Union{Int, Void}` : Terminal period in finite horizon problem
- `rf::ScalarOrArray` : `n x n` terminal payoff in finite horizon problem. Must be symmetric and nonnegative definite
- `P::ScalarOrArray` : `n x n` matrix in value function representation ``V(x) = x'Px + d``
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period

"""
mutable struct LQ
    Q::ScalarOrArray
    R::ScalarOrArray
    A::ScalarOrArray
    B::ScalarOrArray
    C::ScalarOrArray
    N::ScalarOrArray
    bet::Real
    capT::Union{Int, Nothing} # terminal period
    rf::ScalarOrArray
    P::ScalarOrArray
    d::Real
    F::ScalarOrArray # policy rule
end

"""
Main constructor for LQ type

Specifies default argumets for all fields not part of the payoff function or
transition equation.

##### Arguments

- `Q::ScalarOrArray` : `k x k` payoff coefficient for control variable u. Must be
  symmetric and nonnegative definite
- `R::ScalarOrArray` : `n x n` payoff coefficient matrix for state variable x.
  Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : `n x n` coefficient on state in state transition
- `B::ScalarOrArray` : `n x k` coefficient on control in state transition
- `;C::ScalarOrArray{zero(size(R}(1)))` : `n x j` coefficient on random shock in
  state transition
- `;N::ScalarOrArray{zero(size(B,1)}(size(A, 2)))` : `k x n` cross product in
  payoff equation
- `;bet::Real(1.0)` : Discount factor in `[0, 1]`
- `capT::Union{Int, Void}(Void)` : Terminal period in finite horizon
  problem
- `rf::ScalarOrArray{fill(NaN}(size(R)...))` : `n x n` terminal payoff in finite
  horizon problem. Must be symmetric and nonnegative definite.

"""
function LQ(Q::ScalarOrArray,
            R::ScalarOrArray,
            A::ScalarOrArray,
            B::ScalarOrArray,
            C::ScalarOrArray=fill(zero(eltype(R)), size(R, 1)),
            N::ScalarOrArray=zero(B'A);
            bet::ScalarOrArray=1.0,
            capT::Union{Int,Nothing}=nothing,
            rf::ScalarOrArray=fill(NaN, size(R)...))

    k = size(Q, 1)
    n = size(R, 1)
    F = k==n==1 ? zero(Float64) : fill(zero(Float64), k, n)
    P = copy(rf)
    d = 0.0

    LQ(Q, R, A, B, C, N, bet, capT, rf, P, d, F)
end

@doc doc"""
Update `P` and `d` from the value function representation in finite horizon case

##### Arguments

- `lq::LQ` : instance of `LQ` type

##### Returns

- `P::ScalarOrArray` : `n x n` matrix in value function representation
  ``V(x) = x'Px + d``
- `d::Real` : Constant in value function representation

##### Notes

This function updates the `P` and `d` fields on the `lq` instance in addition to
returning them

"""
function update_values!(lq::LQ)
    # Simplify notation
    Q, R, A, B, N, C, P, d = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C, lq.P, lq.d

    # Some useful matrices
    s1 = Q + lq.bet * (B'P*B)
    s2 = lq.bet * (B'P*A) + N
    s3 = lq.bet * (A'P*A)

    # Compute F as (Q + B'PB)^{-1} (beta B'PA)
    lq.F = s1 \ s2

    # Shift P back in time one step
    new_P = R - s2'lq.F + s3

    # Recalling that tr(AB) = tr(BA)
    new_d = lq.bet * (d + tr(P * C * C'))

    # Set new state
    lq.P, lq.d = new_P, new_d
end

@doc doc"""
Computes value and policy functions in infinite horizon model.

##### Arguments

- `lq::LQ` : instance of `LQ` type

##### Returns

- `P::ScalarOrArray` : n x n matrix in value function representation ``V(x) = x'Px + d``
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period

##### Notes

This function updates the `P`, `d`, and `F` fields on the `lq` instance in
addition to returning them

"""
function stationary_values!(lq::LQ)
    # simplify notation
    Q, R, A, B, N, C = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C

    # solve Riccati equation, obtain P
    A0, B0 = sqrt(lq.bet) * A, sqrt(lq.bet) * B
    P = solve_discrete_riccati(A0, B0, R, Q, N)

    # Compute F
    s1 = Q .+ lq.bet * (B' * P * B)
    s2 = lq.bet * (B' * P * A) .+ N
    F = s1 \ s2

    # Compute d
    d = lq.bet * tr(P * C * C') / (1 - lq.bet)

    # Bind states
    lq.P, lq.F, lq.d = P, F, d
end

"""
Non-mutating routine for solving for `P`, `d`, and `F` in infinite horizon model

See docstring for `stationary_values!` for more explanation
"""
function stationary_values(lq::LQ)
    _lq = LQ(copy(lq.Q),
             copy(lq.R),
             copy(lq.A),
             copy(lq.B),
             copy(lq.C),
             copy(lq.N),
             bet=copy(lq.bet),
             capT=lq.capT,
             rf=copy(lq.rf))

    stationary_values!(_lq)
    return _lq.P, _lq.F, _lq.d
end

"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence(lq::LQ, x0::T, policies) where T
    capT = length(policies)

    x_path = Vector{T}(undef, capT+1)
    u_path = Vector{T}(undef, capT)

    x_path[1] = x0
    u_path[1] = -(first(policies)*x0)
    w_path = lq.C * randn(capT+1)

    for t = 2:capT
        f = policies[t]
        x_path[t] = lq.A*x_path[t-1] + lq.B*u_path[t-1] + w_path[t]
        u_path[t] = -(f*x_path[t])
    end
    x_path[end] = lq.A*x_path[capT] + lq.B*u_path[capT] + w_path[end]

    x_path, u_path, w_path
end

"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence(lq::LQ, x0::Vector{T}, policies) where T
    # Ensure correct dimensionality
    n, j, k = size(lq.C, 1), size(lq.C, 2), size(lq.B, 2)
    capT = length(policies)

    A, B, C = lq.A, reshape(lq.B, n, k), reshape(lq.C, n, j)

    x_path = Matrix{T}(undef, n, capT+1)
    u_path = Matrix{T}(undef, k, capT)
    w_path = C*randn(j, capT+1)

    x_path[:, 1] = x0
    u_path[:, 1] = -(first(policies)*x0)

    for t = 2:capT
        f = policies[t]
        x_path[:, t] = A*x_path[: ,t-1] + B*u_path[:, t-1] + w_path[:, t]
        u_path[:, t] = -(f*x_path[:, t])
    end
    x_path[:, end] = A*x_path[:, capT] + B*u_path[:, capT] + w_path[:, end]

    x_path, u_path, w_path
end

@doc doc"""
Compute and return the optimal state and control sequence, assuming innovation ``N(0,1)``

##### Arguments

- `lq::LQ` : instance of `LQ` type
- `x0::ScalarOrArray`: initial state
- `ts_length::Integer(100)` : maximum number of periods for which to return process. If `lq` instance is finite horizon type, the sequenes are returned only for `min(ts_length, lq.capT)`

##### Returns

- `x_path::Matrix{Float64}` : An `n x T+1` matrix, where the t-th column represents ``x_t``
- `u_path::Matrix{Float64}` : A `k x T` matrix, where the t-th column represents ``u_t``
- `w_path::Matrix{Float64}` : A `n x T+1` matrix, where the t-th column represents `lq.C*N(0,1)`

"""
function compute_sequence(lq::LQ, x0::ScalarOrArray, ts_length::Integer=100)

    # Compute and record the sequence of policies
    if isa(lq.capT, Nothing)
        stationary_values!(lq)
        policies = fill(lq.F, ts_length)
    else
        capT = min(ts_length, lq.capT)
        policies = Vector{typeof(lq.F)}(undef, capT)
        for t = capT:-1:1
            update_values!(lq)
            policies[t] = lq.F
        end
    end

    _compute_sequence(lq, x0, policies)
end


@rule typeof(*)(:out, Marginalisation) (m_A::PointMass, m_in::PointMass) = PointMass(mean(m_A) * mean(m_in))

@rule typeof(*)(:out, Marginalisation) (m_A::PointMass{ <: Real }, m_in::GammaDistributionsFamily) = begin
    return GammaShapeRate(shape(m_in), rate(m_in) / mean(m_A))
end

@rule typeof(*)(:out, Marginalisation) (m_A::PointMass{ <: AbstractMatrix }, m_in::F) where { F <: NormalDistributionsFamily } = begin
    A = mean(m_A)
    μ_in, Σ_in = mean_cov(m_in)
    return convert(promote_variate_type(F, NormalMeanVariance), A * μ_in, A * Σ_in * A')
end

#------------------------
# AbstractVector * UnivariateNormalDistributions
#------------------------
# We consider the following updates as a special case of the MatrixVariate * Multivariate updates.
# Namely, Ax = y, where A ∈ R^{nx1}, x ∈ R^1, and y ∈ R^n. In this case, the matrix A
# can be represented by a n-dimensional vector, and x by a scalar. Before computation,
# quantities are converted to their proper dimensions (see situational sketch below).
#
#     | a ~ AbstractVector -> R^{nx1}
#     v  out ~ Multivariate -> R^n
# -->[x]-->
# in1 ~ Univariate -> R^1
@rule typeof(*)(:out, Marginalisation) (m_A::PointMass{ <: AbstractVector }, m_in::UnivariateNormalDistributionsFamily) = begin
    a = mean(m_A)

    μ_in, v_in = mean_var(m_in)

    # TODO: check, do we need correction! here? (ForneyLab does not have any correction in this case)
    # TODO: Σ in this rule is guaranteed to be ill-defined, has rank equal to one and has determinant equal to zero
    μ = μ_in * a
    Σ = mul_inplace!(v_in, a * a')

    return MvNormalMeanCovariance(μ, Σ)
end

#------------------------
# Real * UnivariateNormalDistributions
#------------------------
@rule typeof(*)(:out, Marginalisation) (m_A::PointMass{ <: Real }, m_in::UnivariateNormalDistributionsFamily) = begin
    a = mean(m_A)
    μ_in, v_in = mean_var(m_in)

    return NormalMeanVariance(a*μ_in, a^2*v_in)
end


#   Author: V. Vitale
#   Feb 2022
#   Conctraction utilities

## initial E and F matrices for the left and right vacuum states
function initial_L(W::Array)
    sW=size(W)
    L = ones(1,sW[1],1)#ones(sW[1],1,1)
    return L
end

function initial_R(W::Array)
    sW=size(W)
    R = ones(1,sW[2],1)#ones(sW[2],1,1)
    return R
end

function construct_R(A::MPS, W::MPO)
    R = Dict()
    R[A.N] = initial_R(W.data[A.N])
    for i in A.N:-1:1
        R[i-1] = contract_from_right(R[i], A.data[i], W.data[i])
    end
    return R
end

function construct_L(A::MPS, W::MPO)
    L = Dict()
    L[1] = initial_L(W.data[1])
    return L
end

## tensor contraction from the left hand side
## +-    +--X-
## |     |  |
## L' =  L--W-
## |     |  |
## +-    +--X-

function contract_from_left(E_L::Array,X::Array,W::Array)
    @tensor temp_1[:] := E_L[1,-3,-4]  * X[1,-2,-1]
    @tensor temp_2[:] := temp_1[-1, -2, -5, 1 ]* conj( X[1, -4,-3] )
    @tensor temp[:] := temp_2[ -1, 2, -3, 3, 1 ] * W[1,-2, 2, 3]
    return temp
end

## tensor contraction from the right hand side
##  -+     -Y--+
##   |      |  |
##  -R' =  -W--R
##   |      |  |
##  -+     -Y--+
function contract_from_right(E_R::Array,Y::Array,W::Array)
    @tensor temp_1[:] :=  Y[-1,-2,1]  * E_R[ 1, -3,-4]
    @tensor temp_2[:] := temp_1[ -1, -2, -5,  1] * conj( Y[-3,-4,1] )
    @tensor temp[:] := temp_2[ -1, 1, -3, 2, 6] * W[ -2 6 1 2 ]
    return temp
end


using PlanerTransformer
using Dierckx

in_3f4 =
 (25e3,(82,19),(210,360),
  100e3,(48,19),(170,1070),
  200e3,(35,19),(150,1300),
  400e3,(25,16),(140,1700),
  1e6,(15,20),(78,1900),
  2e6,(5.5,20),(40,1850),
  3e6,(4,40),(26,2000))

# x=log10(B), y=log10(frequency), z=log10(P)

x = log10.([82,   210,  48,    170,   35,    150,   25,    140,   15,  78,   5.5, 20,   4,   26])
y = log10.([25e3, 25e3, 100e3, 100e3, 200e3, 200e3, 400e3, 400e3, 1e6, 1e6,  2e6, 2e6,  3e6, 3e6])
z = log10.([19,   360,  19,    1070,  19,    1300,  16,    1700,  20,  1900, 20,  1850, 40,  2000])
Spline2D(x, y, z; kx=3, ky=3, s=0.0)


function zad35(d, el1, el2)
    A = zeros(d,d)
    A[:, 1:2:end] .= el1
    A[:, 2:2:end] .= el2
    return A

end


using FactCheck
facts("testing ols.jl") do
    # OLS
    x = [ones(10,1) (1:10)]
    srand(1)
    y = rand(10,1)
    b = ols(y,x,silent=true)[1]
    @fact b[1,1] --> roughly(0.07655377367377458)
    # samin
    include("../examples/samin_example.jl");
    @fact xopt[2] -->roughly(3.25)
    # npreg
    include("../examples/npreg_example.jl");
    @fact yhat[1,1] -->roughly(-0.8108308248103517)
end


using ClusteringMetrics
using Test

Z1 = fill(0.0, 2, 10)
Z1[:,1:5] = [0.1 0.2 0.15 0.2 0.22;
             0.15 0.21 0.13 0.21 0.19]
Z1[:,6:end] = Z1[:,1:5] .- [1.0;1.0]
ϵ = [-0.010114902986680585 0.004463792650270601 -0.00153613750975559 -0.0055050913393961485 0.0035611946194614437 -0.0009473083769016829 -0.0011402700190551592 0.00043139456531041 0.007363905500805968 -0.0036852874409562726;
     -0.004970181279034265 -0.008439692559205151 -0.002063148403333957 0.0027397663153164308 -0.009655682877822354 0.002952749805508587 0.010722686475465308 -0.008101446035381798 0.0024392528878517524 -0.007034803358359168]
Z1 .+= ϵ
label1 = [fill(1,5);fill(2,5)]

@testset "Isolation distance" begin
    # large isolation distance
    d = ClusteringMetrics.get_isolation_distance(Z1, label1)
    @test d ≈  [7346.57515900546, 11695.14306811699]
end

@testset "Cluster matching" begin
    Z2 = fill!(similar(Z1), 0.0)
    Z2 .= Z1
    label2 = label1
    dd = ClusteringMetrics.get_isolation_distance(Z1, label1, Z2, label2)
    # large cross distance, indicating high generalisability
    @test dd ≈ [296.9888641641518, 301.87857155988695]
    dd = ClusteringMetrics.get_cluster_distance(Z1, label1, Z2, label2)

    # zero confusion when the sets are the same
    @test dd ≈ [0.0, 0.0]

    n = ClusteringMetrics.get_cluster_isolation(Z1, label1, Z2, label2)
    @test n ≈ [1.0, 1.0]

    Z2[:,1:5] = Z1[:,6:end]
    Z2[:, 6:end] = Z1[:, 1:5]
    dd = ClusteringMetrics.get_isolation_distance(Z1, label1, Z2, label2)
    # low cross distance indicating low generlisability
    @test dd ≈ [0.0033671296154972326, 0.003312590207488839]
    dd = ClusteringMetrics.get_cluster_distance(Z1, label1, Z2, label2)
    @test dd ≈ [64.50554889037353, 64.5055488903735]
    n = ClusteringMetrics.get_cluster_isolation(Z1, label1, Z2, label2)
    @test n ≈ [0.0, 0.0]
end


using CGAL
using Test

using CGAL.Cxx

cxx"""
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>

typedef CGAL::Exact_predicates_inexact_constructions_kernel K;
typedef K::Point_3          Point_3;
typedef K::Vector_3         Vector_3;
typedef K::Triangle_3       Triangle_3;
typedef K::Tetrahedron_3    Tetrahedron_3;
"""

indexable_to_array(v::Cxx.CxxCore.CppValue, n) = [icxx"auto r =$(v)[$i]; r;" for i = 0:(n-1)]
indexable_to_array(v::Cxx.CxxCore.CppPtr, n) = [icxx"auto r =(*($(v)))[$i]; r;" for i = 0:(n-1)]

@testset "Vector_3 - construct/add" begin
    pv1 = @cxxnew Vector_3(10, 20, 30)
    pv2 = @cxxnew Vector_3(1, 2, 3)

    v3 = icxx"*$pv1 + *$pv2;"

    @test 11 == icxx"auto r = $(v3).x(); r;"
    @test 22 == icxx"auto r = $(v3).y(); r;"
    @test 33 == icxx"auto r = $(v3).z(); r;"

    @test indexable_to_array(pv1, 3) == [10, 20, 30]
    @test indexable_to_array(pv2, 3) == [1, 2, 3]
    @test indexable_to_array(v3, 3) == [11, 22, 33]
end

@testset "Triangle_3" begin
    tri1 = icxx"Triangle_3(Point_3(1,-1,0), Point_3(-1,-1,0), Point_3(0,1,0));"
    tri2 = icxx"Triangle_3(Point_3(1,-1,1), Point_3(-1,-1,1), Point_3(0,1,1));"
    tri3 = icxx"Triangle_3(Point_3(1,0,-1), Point_3(-1,0,-1), Point_3(0,0,1));"
    @test icxx"do_intersect($tri1, $tri1);" == true
    @test icxx"do_intersect($tri1, $tri2);" == false
    @test icxx"do_intersect($tri1, $tri3);" == true

    @test (@cxx CGAL::do_intersect(tri1, tri3)) == true

    @test icxx"$tri1.squared_area();" == 4
end

@testset "Tetrahedron_3" begin
    tet1 = icxx"Tetrahedron_3(Point_3(1,-1,0), Point_3(-1,-1,0), Point_3(0,1,0), Point_3(0,0,1));"
    tri4 = icxx"Triangle_3(Point_3(1,-1,0.5), Point_3(-1,-1,0.5), Point_3(0,1,0.5));"
    tri5 = icxx"Triangle_3(Point_3(1,-1,10.0), Point_3(-1,-1,10.0), Point_3(0,1,10.0));"
    @test (@cxx CGAL::do_intersect(tet1, tri4)) == true
    @test (@cxx CGAL::do_intersect(tet1, tri5)) == false
end


using  DataFrames
using LinearAlgebra

mutable struct HiddenLayer

    number_of_neurons::Integer
    weight_matrix::Matrix{Float64}
    weight_initializer::Function
    bias_vector::Vector{Float64}
    activation_function::Function

end

function return_name(x)
    return String(Symbol(x))
end

function activations(layer::HiddenLayer,x::Matrix{Float64})
    # Calculates the activations of the hidden layer neurons
    #
    # Parameters
    # ----------
    # `layer` is the HiddenLayer with neurons
    # `x` is the input matrix
    #
    # Returns
    # -------
    # Activation matrix after passing through hidden layer

    n_observations = size(x)[1]
    activation_matrix = zeros(layer.n_of_neurons, n_observations)

        for i = 1:n_observations
            activation_matrix[:, i] = layer.activation_func(layer.weight_matrix * x[i, :]' + layer.bias_vector)
        end

    activation_matrix
 end

 mutable struct ExtremeLearningMachine
    # Extreme Learning Machine
    #
    # Properties
    # ----------
    # `hidden_layer` is the hidden layer inside this ELM
    # `output_weights` is the weight matrix calculated during training

    hidden_layer::HiddenLayer
    output_weights::Matrix{Float64}

    function ExtremeLearningMachine(n_hidden_neurons::Integer)
	this = new()

	weight_matrix =Array{Float64}(undef,0,0)
	bias_vector = rand(n_hidden_neurons)
	this.hidden_layer = HiddenLayer(n_hidden_neurons, weight_matrix,glorot_uniform, bias_vector, sigmoid)
	println("######################### ELM STRUCTURE ##############################")
    println("# Number of Hidden Neurons: ",this.hidden_layer.number_of_neurons)
    println("# Activation Function: ",return_name(this.hidden_layer.activation_function))
    println("######################################################################")
    this

    end

end

function compile(elm::ExtremeLearningMachine,input_dim::Integer)
    elm.hidden_layer.weight_matrix=elm.hidden_layer.weight_initializer(input_dim,elm.hidden_layer.number_of_neurons)
    println("compiled")
end

function find_activations(layer::HiddenLayer,x::Matrix{Float64})
# Calculates the activations of the hidden layer neurons
#
# Parameters
# ----------
# `layer` is the HiddenLayer with neurons
# `x` is the input matrix
#
# Returns
# -------
# Activation matrix after passing through hidden layer

n_observations = size(x)[1]
act_matrix = zeros(layer.number_of_neurons, n_observations)
transposedx=x'
for i = 1:n_observations
    act_matrix[:, i] =layer.weight_matrix * transposedx[ :,i] + layer.bias_vector
end

map((x)->layer.activation_function(x),act_matrix)

end

function fit!(elm::ExtremeLearningMachine,x,y)
    # Trains the elm using the given training data
    #
    # Parameters
    # ----------
    # `elm` the ELM to train
    # `x` input data
    # `y` output data

    if typeof(x) == DataFrame
    x_mat = float(array(x))
    else
    x_mat = x
    end

    if typeof(y) != Vector{Float64}
    y_vec = float(array(y))
    else
    y_vec = y
    end

    n_observations, n_inputs = size(x_mat)
    weight_matrix = rand(elm.hidden_layer.number_of_neurons, n_inputs) * 2 .- 1
    elm.hidden_layer.weight_matrix = weight_matrix

    act_matrix = find_activations(elm.hidden_layer, x_mat)

    output_weights = y_vec' * pinv(act_matrix)
    elm.output_weights = output_weights
    end

    function fitSVD!(elm::ExtremeLearningMachine,x,y)
        # Trains the elm using the given training data
        #
        # Parameters
        # ----------
        # `elm` the ELM to train
        # `x` input data
        # `y` output data

        if typeof(x) == DataFrame
        x_mat = float(array(x))
        else
        x_mat = x
        end

        if typeof(y) != Vector{Float64}
        y_vec = float(array(y))
        else
        y_vec = y
        end

        n_observations, n_inputs = size(x_mat)
        weight_matrix = rand(elm.hidden_layer.number_of_neurons, n_inputs) * 2 .- 1
        elm.hidden_layer.weight_matrix = weight_matrix

        act_matrix = find_activations(elm.hidden_layer, x_mat)

        mySvd =  SVDCM(act_matrix)
        println(size(mySvd.V))
        println(size(((mySvd.S_inv))))
        c=mySvd.V*mySvd.S_inv
        d=c*mySvd.U'
        output_weights = y_vec' *d
        elm.output_weights = output_weights
        end

function predict(elm::ExtremeLearningMachine,x)
# Predicts the output
#
# Parameters
# ----------
# `elm` the trained ELM
# `x` input data to predict (Matrix)

if typeof(x) == DataFrame
x_mat = float(array(x))
else
x_mat = x
end

act_matrix = find_activations(elm.hidden_layer, x_mat)
vec(elm.output_weights * act_matrix)
end


# PROX_QUAD    The proximal operator of a quadratic.
#
#   prox_quad(v,lambda,A,b)
function prox_quad(
	v,
	lambda::Real,
	A,
	b
)
    rho = 1 / lambda
    n = size(A)

    if issparse(A)
        x = (A + rho * speye(m)) \ (rho * v .- b);
    else
        x = (A + rho * eye(m)) \ (rho * v .- b);
    end

    return x
end


#!/usr/bin/env julia

using Eigenbroetler

function magic_phase(x::Number, y::Number)
    r = sqrt(x * x + y * y)
    return exp(-(r / 100) ^ 2) * exp(2 * im * x * (x ^ 2 - y ^ 2) / 200 ^ 2)
end

eb = Eigenbrot(magic_phase, 512, 512)
save("magic_phase.png", eb, ImageSetting(Phase, LinearScale))
save("magic_phase_dft.png", fft(eb), ImageSetting(Phase, LinearScale))