abstract type SpectralEstimate{D,T} <: AbstractVector{T} end

@doc raw"""
    Periodogram(timeseries, Δ)

Compute the periodogram for the provided timeseries with sampling rate Δ.

# Arguments
- `timeseries`: A `Vector` if univariate and an `n` by `d` `Matrix` if multivariate, where `n` is the number of observations and `d` is the dimension of the timeseries.
- `Δ`: A positive real number.

The periodogram is defined as
```math
\boldsymbol I(ω)&=\boldsymbol J(ω) \boldsymbol J(ω)^H \quad \text{where} \quad \boldsymbol J(ω) = \sqrt{\frac{Δ}{2π n}}\sum_{t=0}^{n-1} \boldsymbol{P}_{tΔ}e^{-itΔ ω}
```

Note the periodogram is in terms of angular frequency here, and uses the normalisation ``Δ/2π``.
The choice of normalisation is essentially arbitrary; however, this matches our definition for the spectral density function.
"""
struct Periodogram{D,T,V} <: SpectralEstimate{D,T}
    Ω::V
    ordinates::Vector{T}
    function Periodogram(timeseries::Matrix{T}, Δ::Real) where {T}
        Δ > 0 || throw(ArgumentError("Δ should be a positive."))
        J = fftshift(fft(timeseries, 1),1)
        n = size(J,1)
        D = size(J,2)
        ordinates = [(J[ii, :]*J[ii, :]'.*(Δ / (2π * n))) for ii ∈ 1:size(J,1)]
        Ω = fftshift(fftfreq(n, 2π/Δ))
        new{D, eltype(ordinates), typeof(Ω)}(Ω, ordinates)
    end
    function Periodogram(timeseries::Vector{T}, Δ::Real) where {T}
        Δ > 0 || throw(ArgumentError("Δ should be a positive."))
        J = fftshift(fft(timeseries, 1),1)
        n = size(J,1)
        ordinates = abs.(J.^2).*(Δ / (2π * n))
        Ω = fftshift(fftfreq(n, 2π/Δ))
        new{1, eltype(ordinates), typeof(Ω)}(Ω, ordinates)
    end
    function Periodogram(Ω, ordinates::Vector{T}) where {T}
        new{size(ordinates[1],1), T, typeof(Ω)}(Ω, ordinates)
    end
end

@doc raw"""
    BartlettPeriodogram(timeseries, Δ, segmentlength)

Compute the Bartlett periodogram for the provided timeseries with sampling rate Δ.

# Arguments
- `timeseries`: A `Vector` if univariate and an `n` by `d` `Matrix` if multivariate, where `n` is the number of observations and `d` is the dimension of the timeseries.
- `Δ`: A positive real number.
- `segmentlength`: the length of series used in each segment.

Computes an estimate of the spectral density function using Bartlett's method. Using the same normalisation as `Periodogram`.

# External links

* [Bartlett's method on Wikipedia](https://en.wikipedia.org/wiki/Bartlett%27s_method)

"""
struct BartlettPeriodogram{D,T,V} <: SpectralEstimate{D,T}
    Ω::V
    ordinates::Vector{T}
    function BartlettPeriodogram(timeseries::Matrix{T}, Δ::Real, segmentlength::Int = Int(100÷Δ)) where {T}
        nsegments = size(timeseries, 1) ÷ segmentlength
        P = Periodogram(timeseries[1:segmentlength, :], Δ)

        for ii = 1:nsegments-1
            P += Periodogram(timeseries[segmentlength*ii.+(1:segmentlength), :], Δ)
        end

        BartlettPeriodogram(P/nsegments)
    end
    function BartlettPeriodogram(timeseries::Vector{T}, Δ::Real, segmentlength::Int = Int(100÷Δ)) where {T}
        nsegments = size(timeseries, 1) ÷ segmentlength
        P = Periodogram(timeseries[1:segmentlength], Δ)

        for ii = 1:nsegments-1
            P += Periodogram(timeseries[segmentlength*ii.+(1:segmentlength)], Δ)
        end

        BartlettPeriodogram(P/nsegments)
    end
    function BartlettPeriodogram(p::Periodogram{D,T}) where {D,T}
        new{D, T, typeof(p.Ω)}(p.Ω, p.ordinates)
    end
    function BartlettPeriodogram(Ω, ordinates::Vector{T}) where {T}
        new{size(ordinates[1],1), T, typeof(Ω)}(Ω, ordinates)
    end
end

struct CoherancyEstimate{D,T,V} <: SpectralEstimate{D,T}
    Ω::V
    ordinates::Vector{T}

    function CoherancyEstimate(timeseries::Matrix{T}, Δ::Real, segmentlength::Int = Int(100÷Δ)) where {T}
        S = BartlettPeriodogram(timeseries, Δ, segmentlength)
        CoherancyEstimate(S)
    end

    function CoherancyEstimate(S::SpectralEstimate{D,T}) where {D,T}
        newordinates = [[s[i,j] / sqrt(s[i,i]*s[j,j]) for i ∈ 1:D, j ∈ 1:D] for s ∈ S.ordinates]
        new{D,eltype(newordinates),typeof(S.Ω)}(S.Ω, newordinates)
    end

    function CoherancyEstimate(Ω, ordinates::Vector{T}) where {T}
        new{size(ordinates[1],1), T, typeof(Ω)}(Ω, ordinates)
    end
end

ndims(::SpectralEstimate{D,T}) where {D,T} = D
size(ŝ::SpectralEstimate) = (length(getfreq(ŝ)),)
getindex(ŝ::T, inds) where {T<:SpectralEstimate} = getconstructor(T)(getfreq(ŝ)[inds], getordinate(ŝ)[inds])
getindex(ŝ::SpectralEstimate, ind::Int) = (getfreq(ŝ)[ind],getordinate(ŝ)[ind])
log10(ŝ::T) where {T<:SpectralEstimate} = getconstructor(T)(getfreq(ŝ), log10.(getordinate(ŝ)))
getfreq(p::Periodogram) = p.Ω
getfreq(b::BartlettPeriodogram) = b.Ω
getfreq(c::CoherancyEstimate) = c.Ω
getordinate(p::Periodogram) = p.ordinates
getordinate(b::BartlettPeriodogram) = b.ordinates
getordinate(c::CoherancyEstimate) = c.ordinates
getconstructor(::Type{<:Periodogram}) = Periodogram
getconstructor(::Type{<:BartlettPeriodogram}) = BartlettPeriodogram
getconstructor(::Type{<:CoherancyEstimate}) = CoherancyEstimate

@recipe function f(ŝ::SpectralEstimate)
    HermitianPlot(getfreq(ŝ), getordinate(ŝ))
end

@recipe function f(ŝ::SpectralEstimate{1,T}) where {T} # univariate plotting
    getfreq(ŝ), getordinate(ŝ)
end

function Base.:+(ŝ₁::T, ŝ₂::T) where {T<:SpectralEstimate}
    getfreq(ŝ₁) == getfreq(ŝ₂) || error("Frequencies must be the same to add spectral estimates.")
    getconstructor(T)(getfreq(ŝ₁), getordinate(ŝ₁) .+ getordinate(ŝ₂))
end

function Base.:/(ŝ::T, a::Real) where {T<:SpectralEstimate}
    getconstructor(T)(getfreq(ŝ), getordinate(ŝ)./a)
end


using Compat

if VERSION.minor == 7
  @eval using LinearAlgebra, SparseMatrices
end

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

A = Float64[
  1/2 1/3 1/4 1/5 1/6;
  1/3 1/4 1/5 1/6 1/7;
  1/4 1/5 1/6 1/7 1/8;
  1/5 1/6 1/7 1/8 1/9;
  1/6 1/7 1/8 1/9 1/10;
]
b = ones(5)

(Al, Au, Ap) = lu(A)

@show y = Al \ b[Ap]
println()

@show x = Au \ y
println()

@show As = sparse(A)
println()

@show As\b
println()

@show F = lufact(As)
println()

if VERSION.minor == 7
  @show F.L * F.U == F.Rs .* A[F.p, :]
  println()

  @show F.L*F.U
  println()

  @show sparse(F.Rs.* A[F.p, F.q])
  println()
else
  @show ys = F[:L] \ (F[:Rs] .* b)
  println()

  @show x = F[:U] \ ys
  println()

  @show F[:L]*F[:U] == (F[:Rs] .* A)[F[:p], F[:q]]
  println()

  @show F[:L]*F[:U]
  println()

  @show sparse((F[:Rs] .* A)[F[:p], F[:q]])
  println()
end

@show lufact(A)
println()

end


export
    LinearRegressionModel,
    fit!,
    predict

mutable struct LinearRegressionModel
    df::DataFrame
    label::Symbol
    features::Vector{Symbol}
    argv::Vector{Float64}
end

function LinearRegressionModel(df::DataFrame, label::Symbol, features::Vector{Symbol})
    return LinearRegressionModel(df, label, features, rand(length(features)+1))
end

function LinearRegressionModel(df::DataFrame, label::Symbol, features::Symbol)
    return LinearRegressionModel(df, label, [features], rand(2))
end

function g(model::LinearRegressionModel, row_n::Int64)
    xs = collect(model.df[row_n, model.features])
    pushfirst!(xs, 1)
    y = sum(model.argv .* xs)

    return y
end

function loss(model::LinearRegressionModel)
    n = nrow(model.df)
    l = 0
    for i in 1:n
        l += (g(model, i) - model.df[i, model.label])^2
    end

    return l/n
end

function fit!(model::LinearRegressionModel; lr=1e-4, atol::Float64=1e-6, show=false)
    while (l = loss(model)) > atol
        show && println("Loss: $l")
        dl_da = 0
        dl_db = zeros(length(model.features))
        for i in 1:nrow(model.df)
            # intersection
            dl_da += g(model, i) - model.df[i, model.label]

            # sloaps
            for (j, f) in enumerate(model.features)
                dl_db[j] += (g(model, i) - model.df[i, model.label]) * model.df[i, f]
            end
        end

        # intersection
        model.argv[1] -= lr * 2 * dl_da
        # slopes
        for j in 2:length(model.argv)
            model.argv[j] -= lr * 2 * dl_db[j-1]
        end
    end
end

function predict(model::LinearRegressionModel, xs::Vector{<:Real})
    pushfirst!(xs, 1)
    y = sum(model.argv .* xs)

    return y
end

function predict(model::LinearRegressionModel, x::Real)
    return predict(model, [x])
end


#=
Agent dynamics in the lake model
=#

using QuantEcon
using Plots
pyplot()
include("lake_model.jl")

srand(42)
lm = LakeModel(d=0, b=0)
T = 5000     # Simulation length

alpha, lambda = lm.alpha, lm.lambda
P = [(1 - lambda) lambda; alpha (1 - alpha)]

mc = MarkovChain(P, [0; 1])     # 0=unemployed, 1=employed
xbar = rate_steady_state(lm)

s_path = simulate(mc, T; init=2)
s_bar_e = cumsum(s_path) ./ (1:T)
s_bar_u = 1 - s_bar_e
s_bars = [s_bar_u s_bar_e]

titles = ["Proportion of time unemployed" "Proportion of time employed"]
dates = collect(1:T)

plot(dates, s_bars, layout=(2, 1), title=titles, legend=:none)
hline!(reverse(xbar)', layout=(2, 1), color=:red, linestyle=:dash)


# Returns a partitioning of `range` into sub-ranges of length `partition_size`.
# If `range` is not divisible by `partition_size`, the last sub-range will be
# of length `range` mod `partition_size`.
function partition_range(range::UnitRange, partition_size::Int)
    n_partitions = ceil(Int, last(range) / partition_size)
    partitions = Vector{UnitRange{Int}}(undef, n_partitions)

    for p in 1:n_partitions
        l::Int = first(range) + (p-1) * partition_size
        r::Int = min(l - 1 + partition_size, last(range))
        partitions[p] = l:r
    end

    return partitions
end

function calculate_sample_weights(targets::T) where {T <: AbstractVector}
    t = targets
    N = length(t)
    ψ₀ = sum(t .!= 1) / N
    ψ₁ = sum(t .== 1) / N
    ψ = [q == 1 ? ψ₁ : ψ₀ for q in t]
end

function gaussian_projection_matrix(::Type{T},
                                    n_neurons::Int,
                                    n_features::Int) where {T <: Number}
    randn(T, n_neurons, n_features) ./ sqrt(n_features)
end

function project(samples::T₁,
                 weights::T₂,
                 activation_functions::Vector{T₃}) where {T₁ <: AbstractMatrix,
                                                          T₂ <: AbstractMatrix,
                                                          T₃ <: ActivationFunction}
    X = samples
    W = weights
    fs = activation_functions
    L = length(fs)

    H = W * X
    for l in 1:L
        H[l,:] .= fs[l].(H[l,:])
    end

    H
end


using PkgSRA, Test, Random
Random.seed!(13)

###
### First, Lorenz
###
include("../examples/example_lorenz.jl")
dat = Array(solve_lorenz_system())
true_grad = core_dyn_true(dat)

# Numerical derivative: the main function to be tested
numerical_grad = numerical_derivative(dat, ts)

# Basic
@test !isempty(numerical_grad)
@test size(numerical_grad) == size(true_grad)

# Accuracy
@test isapprox(numerical_grad, true_grad, rtol=1e-2)
@test !isapprox(numerical_grad, true_grad, rtol=1e-3)

###
### Second, SIR
###
include("../examples/example_sir.jl")
dat = Array(solve_sir_system())
true_grad = core_dyn_true(dat)

# Numerical derivative: the main function to be tested
numerical_grad = numerical_derivative(dat, ts)

# Basic
@test !isempty(numerical_grad)
@test size(numerical_grad) == size(true_grad)

# Accuracy
@test isapprox(numerical_grad, true_grad, rtol=1e-2)
@test !isapprox(numerical_grad, true_grad, rtol=1e-3)

# Check why the above doesn't work...
plot(numerical_grad[1,:], label="numerical");
    plot!(true_grad[1,:], label="true")

dt = ts[2] - ts[1]
bad_grad = (dat[:,2:end] .- dat[:,1:end-1]) / dt
plot(ts, numerical_grad[1,:], label="numerical");
    plot!(ts, bad_grad[1,:], label="one-step")
    plot!(ts, true_grad[1,:], label="true")


using JuMP, AmplNLWriter, MathOptInterfaceToMathProgBase, Test, Ipopt
using Random, Statistics, Test

const MOI = JuMP.MOI

model = JuMP.Model(
    with_optimizer(
        MathOptInterfaceToMathProgBase.Optimizer,
        solver = AmplNLWriter.AmplNLSolver(Ipopt.amplexe)
    )
)

N = 1000
h = 1/N
alpha = 350

@variables(model, begin
    -1 <= t[1:(N + 1)] <= 1
    -0.05 <= x[1:(N + 1)] <= 0.05
    u[1:(N + 1)]
end)
@NLobjective(model, Min, sum(0.5 * h * (u[i + 1]^2 + u[i]^2) +
    0.5 * alpha * h * (cos(t[i + 1]) + cos(t[i])) for i in 1:N)
)
@NLconstraint(model, [i = 1:N],
    x[i + 1] - x[i] - 0.5 * h * (sin(t[i + 1]) + sin(t[i])) == 0
)
@constraint(model, [i = 1:N],
    t[i + 1] - t[i] - 0.5 * h * u[i + 1] - 0.5 * h * u[i] == 0
)
JuMP.optimize!(model)

@test JuMP.termination_status(model) == MOI.OPTIMAL
@test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
@test JuMP.objective_value(model) ≈ 350.0

function example_mle(; verbose = true)
    n = 1_000
    Random.seed!(1234)
    data = randn(n)

    model = JuMP.Model(
        with_optimizer(
            MathOptInterfaceToMathProgBase.Optimizer,
            solver = AmplNLWriter.AmplNLSolver(Ipopt.amplexe)
        )
    )

    @variable(model, μ, start = 0.0)
    @variable(model, σ >= 0.0, start = 1.0)
    @NLobjective(model, Max, n / 2 * log(1 / (2 * π * σ^2)) -
        sum((data[i] - μ)^2 for i in 1:n) / (2 * σ^2)
    )
    JuMP.optimize!(model)
    if verbose
        println("μ = ", JuMP.value(μ))
        println("mean(data) = ", mean(data))
        println("σ^2 = ", JuMP.value(σ)^2)
        println("var(data) = ", var(data))
        println("MLE objective: ", JuMP.objective_value(m))
    end
    @test JuMP.value(μ) ≈ mean(data) atol = 1e-3
    @test JuMP.value(σ)^2 ≈ var(data) atol = 1e-2

    # constrained MLE?
    @NLconstraint(model, μ == σ^2)

    JuMP.optimize!(model)
    if verbose
        println("\nWith constraint μ == σ^2:")
        println("μ = ", JuMP.value(μ))
        println("σ^2 = ", JuMP.value(σ)^2)
        println("Constrained MLE objective: ", JuMP.objective_value(model))
    end
    @test JuMP.value(μ) ≈ JuMP.value(σ)^2
end

example_mle(verbose = false)


include("11_two_site_state_preparation.jl")

println("
#######################################################
# Tutorial 12                                         #
#                                                     #
# n-site circuit optimization with MPS                #
#######################################################
")

# Beyond full state, use approximations
n = 30
i = [Index(2, "S=1/2") for j in 1:n]

# Product state |Z+Z+Z+…Z+⟩
Zp = MPS(i, "Z+")

# Product state |Z-Z-Z-…Z-⟩
Zm = MPS(i, "Z-")

inner(Zp, Zp)
inner(Zm, Zp)

# Do some arithmatic with them, like tensors:
# TODO: fix this when normalize PR is merged!
#Cat = (Zp + Zm) / √2
#
# Performs a truncated MPS addition,
# returns an MPS approximate up to some tolerance.
Cat = (Zp + Zm) * (1 / √2)
@show maxlinkdim(Zp)
@show maxlinkdim(Zm)
@show maxlinkdim(Cat)
@show norm(Cat)
@show inner(Cat, Zp)
@show inner(Cat, Zm)

# Apply operators:
j = n ÷ 2
Xⱼ = op("X", i[j])
XⱼZp = apply(Xⱼ, Zp)

XⱼZp = MPS(i, [k == j ? "Z-" : "Z+" for k in 1:n])
@show inner(XⱼZp, apply(Xⱼ, Zp))

function ising(n; h)
  # Symbolic representation
  # of the Hamiltonian ("operator sum")
  H = OpSum()
  for j in 1:(n - 1)
    H -= "Z", j, "Z", j + 1
  end
  for j in 1:n
    H += h, "X", j
  end
  return H
end

h = 0.5
H = MPO(ising(n; h=h), i)

# Performs a truncated MPO*MPS contraction,
# returns an MPS approximate up to some tolerance.
HZp = H * Zp
@show inner(Zp', HZp)
@show inner(Zp', H, Zp)

# Gradient energy minimization
# Already written above! New Hamiltonian.
function E(ψ)
  ψHψ = inner(ψ', H, ψ)
  ψψ = inner(ψ, ψ)
  return ψHψ / ψψ
end

function minimize(f, ∂f, x; nsteps, γ, kwargs...)
  for n in 1:nsteps
    println("n = ", n, ", f_x = ", f(x))
    x = -(x, γ * ∂f(x); kwargs...)
  end
  return x
end

ψ⁰ = MPS(i, "Z+")
∂E(x) = gradient(E, x)[1]
ψ = minimize(E, ∂E, ψ⁰; nsteps=50, γ=0.1, maxdim=10, cutoff=1e-5)

# TODO: change to:
# ψ /= norm(ψ)
ψ *= inv(norm(ψ))

Eᵈᵐʳᵍ, ψᵈᵐʳᵍ = dmrg(H, ψ⁰; nsweeps=10, maxdim=10, cutoff=1e-5)

@show maxlinkdim(ψ⁰)
@show maxlinkdim(ψ)
@show maxlinkdim(ψᵈᵐʳᵍ)
@show E(ψ⁰), norm(∂E(ψ⁰))
@show E(ψ), norm(∂E(ψ))
@show E(ψᵈᵐʳᵍ), norm(∂E(ψᵈᵐʳᵍ))

#
# Circuit optimization
#

# Form the circuit from parameters
Ry_layer(θ, i) = [op("Ry", i[j]; θ=θ[j]) for j in 1:n]
CX_layer(i) = [op("CX", i[j], i[j+1]) for j in 1:2:(n-1)]

function U(θ, i; nlayers)
  n = length(i)
  Uᶿ = Ry_layer(θ[1:n], i)
  for l in 1:(nlayers - 1)
    Uᶿ = [Uᶿ; CX_layer(i)]
    Uᶿ = [Uᶿ; Ry_layer(θ[(1:n) .+ l * n], i)]
  end
  return Uᶿ
end

nlayers = 6
maxdim = 10
cutoff = 1e-5

# Find the circuit `U(θ)` that minimizes:
# E(θ) = ⟨0|U(θ)† H U(θ)|0⟩ = ⟨θ|H|θ⟩
function E(θ)
  # Apply the circuit:
  # |θ⟩ = U(θ)|0⟩
  ψᶿ = apply(U(θ, i; nlayers=nlayers), ψ⁰; maxdim=maxdim, cutoff=cutoff)

  # Compute the expecation value: ⟨θ|H|θ⟩
  # No need to normalize!
  return inner(ψᶿ', H, ψᶿ)
end

println()
println("Circuit optimization")
θ⁰ = zeros(nlayers * n)
∂E(ψ) = gradient(E, ψ)[1]
θ = minimize(E, ∂E, θ⁰; nsteps=20, γ=0.1)

ψᶿ = apply(U(θ, i; nlayers=nlayers), ψ⁰; maxdim=maxdim, cutoff=cutoff)

@show maxlinkdim(ψ⁰)
@show maxlinkdim(ψᶿ)
@show E(θ⁰), norm(∂E(θ⁰))
@show E(θ), norm(∂E(θ))

# TODO: Finish this implementation
# Example of preparing the ground state
# Find `U(θ)` that minimizes
#
# F(θ) = -|⟨ψ|U(θ)|0⟩|²
function F(θ)
  # Apply the circuit:
  # |θ⟩ = U(θ)|0⟩
  ψᶿ = apply(U(θ, i; nlayers=nlayers), ψ⁰; maxdim=maxdim, cutoff=cutoff)

  # -|⟨ψ|θ⟩|²
  return -abs(inner(ψ, ψᶿ))^2
end

println()
println("State preparation")
θ⁰ = zeros(nlayers * n)
∂F(ψ) = gradient(F, ψ)[1]
θ = minimize(F, ∂F, θ⁰; nsteps=20, γ=0.1)

ψᶿ = apply(U(θ, i; nlayers=nlayers), ψ⁰; maxdim=maxdim, cutoff=cutoff)

@show maxlinkdim(ψ⁰)
@show maxlinkdim(ψᶿ)
@show F(θ⁰), norm(∂F(θ⁰))
@show F(θ), norm(∂F(θ))


## R code 5.3
m5.1 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bA * A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.6
d$M <- scale( d$Marriage )
m5.2 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM * M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.10
m5.3 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M + bA*A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.13
m5.4 <- quap(
    alist(
        M ~ dnorm( mu , sigma ) ,
        mu <- a + bAM * A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bAM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.19
data(WaffleDivorce)
d <- list()
d$A <- standardize( WaffleDivorce$MedianAgeMarriage )
d$D <- standardize( WaffleDivorce$Divorce )
d$M <- standardize( WaffleDivorce$Marriage )

m5.3_A <- quap(
    alist(
      ## A -> D <- M
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M + bA*A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 ),
      ## A -> M
        M ~ dnorm( mu_M , sigma_M ),
        mu_M <- aM + bAM*A,
        aM ~ dnorm( 0 , 0.2 ),
        bAM ~ dnorm( 0 , 0.5 ),
        sigma_M ~ dexp( 1 )
    ) , data = d )

## R code 5.29
m5.5_draft <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N ,
        a ~ dnorm( 0 , 1 ) ,
        bN ~ dnorm( 0 , 1 ) ,
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 5.32
m5.5_draft <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N ,
        a ~ dnorm( 0 , 1 ) ,
        bN ~ dnorm( 0 , 1 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.34
m5.5 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N ,
        a ~ dnorm( 0 , 0.2 ) ,
        bN ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.37
m5.6 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.38
m5.7 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N + bM*M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bN ~ dnorm( 0 , 0.5 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.41
# M -> K <- N
# M -> N
n <- 100
M <- rnorm( n )
N <- rnorm( n , M )
K <- rnorm( n , N - M )
d_sim <- data.frame(K=K,N=N,M=M)

m5.7.A <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N + bM*M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bN ~ dnorm( 0 , 0.5 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=d_sim )

## R code 5.47
m5.8 <- quap(
    alist(
        height ~ dnorm( mu , sigma ) ,
        mu <- a[sex] ,
        a[sex] ~ dnorm( 178 , 20 ) ,
        sigma ~ dunif( 0 , 50 )
    ) , data=d )
precis( m5.8 , depth=2 )

## R code 5.51
d$K <- scale( d$kcal.per.g )
m5.9 <- quap(
    alist(
        K ~ dnorm( mu , sigma ),
        mu <- a[clade_id],
        a[clade_id] ~ dnorm( 0 , 0.5 ),
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 5.53
m5.10 <- quap(
    alist(
        K ~ dnorm( mu , sigma ),
        mu <- a[clade_id] + h[house],
        a[clade_id] ~ dnorm( 0 , 0.5 ),
        h[house] ~ dnorm( 0 , 0.5 ),
        sigma ~ dexp( 1 )
    ) , data=d )


# This file investigates integration error with respect to free variables in hidden 1D integration

using FractalTools
using Makie

# Construct interpolation data
f(x) = sin(2π * x)  + 1
g(x) = cos(2π * x)  + 1
dx   = 0.1
xi   = 0.
xf   = 1.
x    = collect(xi : dx : xf)
y    = f.(x)
z    = g.(x)
pts  = collect.(zip(x, y, z))

# Compute errors
fvals = (cos(2π * xi) - cos(2π * xf)) / (2π) + (xf - xi)
freevars = 0.001 : 0.01 : 0.999
mse = map(freevars) do freevar
    ival = integrate(pts, HInterp1D(fill(freevar, 2, 2)))[1]
    abs(fval - ival)
end

# Plot mse
fig = Figure()
ax = fig[1, 1] = Axis(fig, xlabel="Free Variable", ylabel="MSE", title="1D Hidden Interpolation MSE")
stem!(ax, freevars, mse, color=:black)
save(joinpath(@__DIR__, "hinteg1d_error.png"), fig)
display(fig)


# ------------------------------------------------
# basic math operators like dot mul,dot add,etc ..
# ------------------------------------------------
import Base.+
import Base.-
import Base.*
import Base.^

export dotAdd
export dotMul
export matAddVec
export matMulVec

function Base.:+(x::Variable{T}, constant) where T
    # a matrix add a constant element by element
    C = eltype(ᵛ(x))(constant)
    y = Variable{T}(ᵛ(x) .+ C, x.backprop)
    if x.backprop
        function matAddScalarBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y)
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, matAddScalarBackward)
    end
    return y
end

function Base.:+(constant, var::Variable{T}) where T
    return var + constant;
end

function Base.:-(x::Variable{T}, constant) where T
    # a matrix minus a constant element by element
    C = eltype(ᵛ(x))(constant)
    y = Variable{T}(ᵛ(x) .- C, x.backprop)
    if x.backprop
        function matMinusScalarBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y)
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, matMinusScalarBackward)
    end
    return y
end

function Base.:-(constant, x::Variable{T}) where T
    # a matrix minus a constant element by element
    C = eltype(ᵛ(x))(constant)
    y = Variable{T}(C .- ᵛ(x), x.backprop)
    if x.backprop
        function scalarMinusMatBackward()
            if need2computeδ!(x)
                δ(x) .-= δ(y)
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, scalarMinusMatBackward)
    end
    return y
end

function Base.:*(x::Variable{T}, constant) where T
    # a matrix multiplies a constant element by element
    C = eltype(ᵛ(x))(constant)
    y = Variable{T}(ᵛ(x) .* C, x.backprop)
    if x.backprop
        function matMulScalarBackward()
            if need2computeδ!(x)
                δ(x) .+= δ(y) .* constant
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, matMulScalarBackward)
    end
    return y
end

function Base.:*(constant, var::Variable{T}) where T
    return var * constant
end

function Base.:^(x::Variable{T}, n::Int) where T
    # 矩阵、列向量与常数按元素做幂指数运算
    n = eltype(ᵛ(x))(n)
    y = Variable{T}(ᵛ(x) .^ n, x.backprop)
    if x.backprop
        function powerBackward()
            if need2computeδ!(x)
                δ(x) .+= n .* ᵛ(y) ./ ᵛ(x) .* δ(y);
            end
            ifNotKeepδThenFreeδ!(y);
        end
        push!(graph.backward, powerBackward)
    end
    return y
end

function Base.:+(x::Variable{T1}, y::Variable{T2}) where {T1,T2}
    # a matrix add a matrix element by element: z = x + y
   @assert T1 <: T2 || T1 >: T2
   T = T1 <: T2 ? T1 : T2
   @assert (x.shape == y.shape) "2 inputs shall be the same size"
   backprop = (x.backprop || y.backprop)
   z = Variable{T}(ᵛ(x) + ᵛ(y), backprop)
   if backprop
       function add2varBackward()
           if need2computeδ!(x) δ(x) .+= δ(z) end
           if need2computeδ!(y) δ(y) .+= δ(z) end
           ifNotKeepδThenFreeδ!(z);
       end
       push!(graph.backward, add2varBackward)
   end
   return z
end

function Base.:-(x::Variable{T1}, y::Variable{T2}) where {T1,T2}
    # a matrix minus a matrix element by element : z = x - y
    @assert T1 <: T2 || T1 >: T2
    T = T1 <: T2 ? T1 : T2
    @assert (x.shape == y.shape) "2 inputs shall be the same size"
    backprop = (x.backprop || y.backprop)
    z = Variable{T}(ᵛ(x) - ᵛ(y), backprop)
    if backprop
        function minus2varBackward()
            if need2computeδ!(x) δ(x) .+= δ(z) end
            if need2computeδ!(y) δ(y) .-= δ(z) end
            ifNotKeepδThenFreeδ!(z);
        end
        push!(graph.backward, minus2varBackward)
    end
    return z
end

"""
    dotAdd(var1::Variable{T1}, var2::Variable{T2}) where {T1,T2}
a tensor add a tensor element by element
"""
function dotAdd(x::Variable{T1}, y::Variable{T2}) where {T1,T2}
    # a tensor add a tensor element by element: z = x .+ y
    @assert T1 <: T2 || T1 >: T2
    T = T1 <: T2 ? T1 : T2
    @assert (x.shape == y.shape) "2 inputs shall be the same size"
    backprop = (x.backprop || y.backprop)
    z = Variable{T}(ᵛ(x) .+ ᵛ(y), backprop)
    if backprop
        function dotAddBackward()
            if need2computeδ!(x) δ(x) .+= δ(z) end
            if need2computeδ!(y) δ(y) .+= δ(z) end
            ifNotKeepδThenFreeδ!(z);
        end
        push!(graph.backward, dotAddBackward)
    end
    return z
end

"""
    dotMul(var1::Variable{T1}, var2::Variable{T2}) where {T1,T2}
a tensor multiplies a tensor element by element
"""
function dotMul(x::Variable{T1}, y::Variable{T2}) where {T1,T2}
    # a tensor multiplies a tensor element by element: z = x .* y
    @assert T1 <: T2 || T1 >: T2
    T = T1 <: T2 ? T1 : T2
    @assert (x.shape == y.shape) "2 inputs shall be the same size"
    backprop = (x.backprop || y.backprop)
    z = Variable{T}(ᵛ(x) .* ᵛ(y), backprop)
    if backprop
        function dotMulBackward()
            if need2computeδ!(x) δ(x) .+= δ(z) .* ᵛ(y) end
            if need2computeδ!(y) δ(y) .+= δ(z) .* ᵛ(x) end
            ifNotKeepδThenFreeδ!(z);
        end
        push!(graph.backward, dotMulBackward)
    end
    return z
end

function Base.:*(W::Variable{T1}, X::Variable{T2}) where {T1,T2}
    # matrix W multiplies matrix X
    # 矩阵相乘 Y[i,j] = sum(W[i,k]*X[k,j],k=...)
    # W -- 权重矩阵
    # X -- n个输入列向量组成的矩阵
    # Y -- n个输出列向量组成的矩阵
    @assert T1 <: T2 || T1 >: T2
    T = T1 <: T2 ? T1 : T2
    backprop = (W.backprop || X.backprop)
    Y = Variable{T}(ᵛ(W) * ᵛ(X), backprop)
    if backprop
        function matMulBackward()
            if need2computeδ!(W) δ(W) .+= δ(Y)  * ᵛ(X)' end
            if need2computeδ!(X) δ(X) .+= ᵛ(W)' * δ(Y)  end
            ifNotKeepδThenFreeδ!(Y);
        end
        push!(graph.backward, matMulBackward)
    end
    return Y
end

"""
    matAddVec(var1::Variable{T1}, var2::Variable{T2}) where {T1,T2}
a matrix tensor `var1` adds a vector tensor `var2`
"""
function matAddVec(M::Variable{T1}, V::Variable{T2}) where {T1,T2}
    # M -- 充当和节点，非学习的参数
    # V -- 偏置列向量，要学习的参数
    # Z = M .+ V
    @assert T1 <: T2 || T1 >: T2
    T = T1 <: T2 ? T1 : T2
    @assert (M.shape[1]==V.shape[1] && V.shape[2]==1)
    backprop = (M.backprop || V.backprop)
    Z = Variable{T}(ᵛ(M) .+ ᵛ(V), backprop)
    if backprop
        function matAddVecBackward()
            if need2computeδ!(M)
                δ(M) .+= δ(Z)
            end
            if need2computeδ!(V)
                δ(V) .+= sum(δ(Z), dims=2)
            end
            ifNotKeepδThenFreeδ!(Z);
        end
        push!(graph.backward, matAddVecBackward)
    end
    return Z
end

"""
    matAddVec(var1::Variable{T1}, var2::Variable{T2}) where {T1,T2}
a matrix tensor `var1` multiplies a vector tensor `var2`
"""
function matMulVec(M::Variable{T1}, V::Variable{T2}) where {T1,T2}
    # M -- 一般充当激活节点，非网络需要学习的参数
    # V -- 列向量，循环权重，是网络需要学习的参数
    # Z = M .* V
    @assert T1 <: T2 || T1 >: T2
    T = T1 <: T2 ? T1 : T2
    @assert (M.shape[1]==V.shape[1] && V.shape[2]==1)
    backprop = (M.backprop || V.backprop)
    Z = Variable{T}(ᵛ(M) .* ᵛ(V), backprop)
    if backprop
        function matMulVecBackward()
            if need2computeδ!(M) δ(M) .+=     δ(Z) .* ᵛ(V)          end
            if need2computeδ!(V) δ(V) .+= sum(δ(Z) .* ᵛ(M), dims=2) end
            ifNotKeepδThenFreeδ!(Z);
        end
        push!(graph.backward, matMulVecBackward)
    end
    return Z
end


using ADCME
using PyPlot

function residual_and_jacobian(θ, u)
    X = ae(u, config, θ) + 1.0     # (1)
    Xp = tf.gradients(X, u)[1]
    Xpp = tf.gradients(Xp, u)[1]
    up = [u[2:end];constant(zeros(1))]
    un = [constant(zeros(1)); u[1:end-1]]
    R = Xp .* ((up-un)/2h)^2 + X .* (up+un-2u)/h^2 - φ
    dRdu = Xpp .* ((up-un)/2h)^2 + Xp.*(up+un-2u)/h^2 - 2/h^2*X
    dRdun = -Xp[2:end]/h .* (up-un)[2:end]/2h + X[2:end]/h^2
    dRdup = Xp[1:end-1]/h .* (up-un)[1:end-1]/2h + X[1:end-1]/h^2
    J = spdiag(n-1,
        -1=>dRdun,
        0=>dRdu,
        1=>dRdup)     # (2)
    return R, J
end

config = [20,20,20,1]
n = 100
h = 1/n
x = collect(LinRange(0, 1.0, n+1))

φ = @. (1 - 2*x)*(-100*x^2*(2*x - 2) - 200*x*(1 - x)^2)/(100*x^2*(1 - x)^2 + 1)^2 - 2 - 2/(100*x^2*(1 - x)^2 + 1)
φ = φ[2:end-1]
θ = Variable(ae_init([1,config...]))
u0 = constant(zeros(n-1))
function L(u)    # (3)
  u_obs = (@. x * (1-x))[2:end-1]
  loss = mean((u - u_obs)^2)
end
loss, solution, grad = NonlinearConstrainedProblem(residual_and_jacobian, L, θ, u0)
X_pred = ae(collect(LinRange(0.0,0.25,100)), config, θ) + 1.0

sess = Session(); init(sess)
BFGS!(sess, loss, grad, θ)
x_pred, sol = run(sess, [X_pred, solution])

figure(figsize=(10,4))
subplot(121)
s = LinRange(0.0,0.25,100)
x_exact = @. 1/(1+100*s^2) + 1
plot(s, x_exact, "-", linewidth=3, label="Exact")
plot(s, x_pred, "o", markersize=2, label="Estimated")
legend()
xlabel("u")
ylabel("X(u)")

subplot(122)
s = LinRange(0.0,1.0,101)[2:end-1]
plot(s, (@. s * (1-s)), "-", linewidth=3, label="Exact")
plot(s, sol, "o", markersize=2, label="Estimated")
legend()
xlabel("x")
ylabel("u")
savefig("nn.png")


## Packages
using NLsolve
using Parameters
using BenchmarkTools
using ForwardDiff

## Files
include("NewtonRoot.jl")
include("Riccati.jl")

######################## Preliminary ###################################
# Construct gradient and hessian at SS
function getM(md,kss,hss)
    @unpack θ,ψ,δ,γ_n,γ_z,ρ,β_hat = md

    f(x::Vector) = log((x[1]^θ)*((exp(x[2])*x[5])^(1-θ))-(1+γ_n)*(1+γ_z)*x[4]+(1-δ)*x[1])+ψ*log(1-x[5]);
    z_bar = [kss,0.0,1.0,kss,hss];
    grad = ForwardDiff.gradient(f,z_bar);
    hess = ForwardDiff.hessian(f,z_bar);

    # Apply Kydland and Prescott's method
    e = [0;0;1;0;0];

    M = e.*(f(z_bar)-grad'*z_bar.+(0.5.*z_bar'*hess*z_bar))*e' +
        0.5*(grad*e'-e*z_bar'*hess-hess*z_bar*e'+e*grad') +
        0.5*hess;

    # Translating M into the matrices we need:
    Q = M[1:3,1:3];
    W = M[1:3,4:5];
    R = M[4:5,4:5];

    A = [0 0 0; 0 ρ 0; 0 0 1];
    B = [1 0; 0 0; 0 0];
    C = [0; 1; 0];

    #Mapping to the problem without discounting (1 VARIABLES ARE ~ IN LECTURE NOTES)
    A_tld = sqrt(β_hat)*(A-B*(R\W'));
    B_tld = sqrt(β_hat)*B;
    Q_tld = Q-W*(R\W');

    return A_tld,B_tld,Q_tld,R,W
end

################### LQ Approximation ##################################
function LQ(A_tld,B_tld,R,Q_tld;tol=1.0e-10)

    Pn, Fn = Riccati(A_tld, B_tld, R, Q_tld, tol);
    F = Fn + R\W';
    P = Pn;

    return P,F
end

################### Vaughan ##################################
function Vaughan(A_tld,B_tld,Q_tld,R,W)

    L = size(A_tld)[1]
    H = [inv(A_tld)  (A_tld\B_tld)*(R\B_tld');
        Q_tld/A_tld Q_tld*(A_tld\B_tld)*(R\B_tld')+A_tld'];

    V = eigen(H).vectors
    #Note that Julia puts the eigenvalues out of the unit circle in the bottom of the matrix,
    #while in the lecture notes they are at the top
    P = V[L+1:end,L+1:end]/(V[1:L,L+1:end]);
    F1 = (R+B_tld'*P*B_tld)\B_tld'*P*A_tld;
    F = F1+R\W';

    return P,F
end


using CSV
using Plots
using StatPlots
using DataFrames
using GLM

gr()
cd(raw"D:\workspace\github\blog-py\blog\static\essay_resources\Notes_on_Computer_Age_Statistical_Inference")

kidney = CSV.read("kidney.csv", nullable=false)
@df kidney scatter(:age, :Tot)

X = hcat(ones(nrow(kidney)), kidney[:age])
y = kidney[:Tot]
OLS = fit(LinearModel, X, y)
# GLM.LinearModel{GLM.LmResp{Array{Float64,1}},GLM.DensePredChol{Float64,Base.LinAlg.Cholesky{Float64,Array{Float64,2}}}}:

# Coefficients:
#        Estimate Std.Error  t value Pr(>|t|)
# x1      2.86067  0.359561  7.95603   <1e-12
# x2   -0.0786009 0.0090557 -8.67972   <1e-14

age_samples = collect(20:10:90)
Xtest = hcat(ones(length(age_samples)), age_samples)
pred = predict(OLS, Xtest, :confint)

for i in 1:size(pred, 1)
    y_pred, y_lower, y_upper = pred[i, :]
    display(plot!([age_samples[i],age_samples[i]], [ y_lower, y_upper], linewidth = 3))
end

plot!(age_samples[[1, end]], pred[[1, end], 1], legend=:none, linewidth=3)
savefig("Figure_1_1.png")


module SpectralMethodsTrefethen

# Analogs for named functions in the book.
export cheb, chebfft, clencurt, gauss

# Required by the scripts, so must be in parent scope
eval(module_parent(current_module()),:(using PyPlot, Dierckx, Polynomials, LaTeXStrings, SpecialFunctions))

# Required to define view() below
using PyPlot: gca

# See if MATLAB is available/working.
is_matlab_running = try
    using MATLAB
    true
catch exc
    warn("MATLAB is not available. Continuing without MATLAB.");
    false
end

# Now see if it can find our files. We have to prevent the compiler from
# dealing with the MATLAB package functions, if that was not loaded successfully.
set_dir = quote
    is_matlab_running = try
        local d = joinpath(Pkg.dir("SpectralMethodsTrefethen"),"matlab");
        eval_string("mfiledir = cd('$d');");
        @mget mfiledir;
        # At some point MATLAB started warning for 2D 'cubic' interp.
        eval_string("warning off MATLAB:griddedInterpolant:CubicUniformOnlyWarnId");
        true
    catch
        warn("MATLAB cannot find package files. Continuing without MATLAB.");
        false
    end
end;

if is_matlab_running
    eval(set_dir)
end

## Analogs of functions from the book.

"""
    cheb(N)

Chebyshev differentiation matrix and grid.
"""
function cheb(N)
    N==0 && return 0,1;
    x = [ cos(pi*k/N) for k=0:N ];
    c = [2;ones(N-1);2] .* (-1).^(0:N);
    dX = x .- x';
    D  = (c*(1./c)') ./ (dX+eye(N+1));      # off-diagonal entries
    D  = D - diagm(squeeze(sum(D,2),2));    # diagonal entries
    return D,x
end

"""
    chebfft(v)

Differentiate values given at Chebyshev points via the FFT.
"""
function chebfft(v)
    # Simple, not optimal. If v is complex, delete "real" commands.
    N = length(v)-1;
    N==0 && return 0;
    x = [ cos(pi*k/N) for k=0:N ];
    ii = collect(0:N-1);
    V = [v; flipdim(v[2:N],1)];              # transform x -> theta
    U = real(fft(V));
    W = real(ifft(1im*[ii;0;1-N:-1].*U));
    w = zeros(N+1);
    w[2:N] = -W[2:N]./sqrt.(1-x[2:N].^2);    # transform theta -> x
    w[1] = sum(ii.^2.*U[ii+1])/N + .5*N*U[N+1];
    w[N+1] = sum((-1).^(ii+1).*ii.^2.*U[ii+1])/N + .5*(-1)^(N+1)*N*U[N+1];
    return w
end

"""
    clencurt(N)

Nodes and weights for Clenshaw-Curtis quadrature
"""
function clencurt(N)
    θ = [ pi*i/N for i=0:N ];
    x = cos.(θ);
    w = zeros(N+1);
    ii = 2:N;
    v = ones(N-1);
    if mod(N,2)==0
        w[1] = w[N+1] = 1/(N^2-1);
        for k = 1:N/2-1
            v = v - 2*cos.(2*k*θ[ii]) / (4*k^2-1);
        end
        v = v - cos.(N*θ[ii]) / (N^2-1);
    else
        w[1] = w[N+1] = 1/N^2;
        for k = 1:(N-1)/2
            v = v - 2*cos.(2*k*θ[ii]) / (4*k^2-1);
        end
    end
    w[ii] = 2*v/N;
    return x,w
end

"""
    gauss(N)

Nodes and weights for Gauss quadrature
"""
function gauss(N)
    β = [ .5/sqrt(1-1/(2*i)^2) for i = 1:N-1 ];
    T = diagm(β,1) + diagm(β,-1);
    (x,V) = eig(T);
    i = sortperm(x);
    w = 2*V[1,i].^2;
    return x[i],w
end

## Stand-ins for native functions in MATLAB.
import Base.view
export toeplitz, view
"""
    toeplitz(col[,row])

Construct Toeplitz matrix from first column and first row. If the row is not
given, the result is symmetric.
"""
function toeplitz(col,row=col)
    m,n = length(col),length(row);
    col[1]==row[1] || warn("Column wins conflict on the diagonal.");
    x = [ row[end:-1:2]; col ];
    return [ x[i-j+n] for i=1:m, j=1:n ]
end

"""
    view(az,el)

Sets the 3D viewing orientation azimuth and elevation (in degrees).
"""
function view(az::Real,el::Real)
    gca()[:view_init](el,az-90);
end

## Create callable functions for each of the Julia scripts.

export p1, p2, p3, p4, p5, p6, p7, p8, p9, p10
export p11, p12, p13, p14, p15, p16, p17, p18, p19, p20
export p21, p22, p23, p24, p25, p26, p27, p28, p29, p30
export p31, p32, p33, p34, p35, p36, p37, p38, p39, p40
export p6u, p23a, p24fine, p28b, p30b, p30c
for (root, dirs, files) in walkdir(joinpath(Pkg.dir("SpectralMethodsTrefethen"),"src","scripts"))
    for file in filter(x->endswith(x,".jl"),files)
        basename = file[1:end-3];
        fundef = quote
            function $(Symbol(basename))(;julia=true,matlab=$is_matlab_running,source=false)
                println(join(["Running script ",$basename,"..."]));
                if julia
                    # Packages are needed in Main scope for include()
                    println("Julia version:");
                    close("all")
                    tic(); include(joinpath($root,$file)); t=toc();
                end
                if matlab
                    println("MATLAB version:");
                    eval_string("close all");
                    tic(); eval_string($basename); t=[t;toc()];
                end
                if source
                    julia ? edit(joinpath($root,$file)) : nothing;
                    matlab ? edit(joinpath($root,"..","..","matlab",$basename * ".m")) : nothing;
                end
                return t

            end
        end
        eval(fundef);
    end
end

end


"""
    PGeneralizedGaussian(α, μ, p)

The *p-Generalized Gaussian distribution*, more commonly known as the exponential
power or the generalized normal distribution, with scale `α`, location `μ`, and
shape `p` has the probability density function

```math
f(x, \\mu, \\alpha, p) = \\frac{p}{2\\alpha\\Gamma(1/p)} e^{-(\\frac{|x-\\mu|}{\\alpha})^p} \\quad x \\in (-\\infty, +\\infty) , \\alpha > 0, p > 0
```

The p-Generalized Gaussian (GGD) is a parametric distribution that incorporates the
Normal and Laplacian distributions as special cases where `p = 1` and `p = 2`. As `p → ∞`,
the distribution approaches the Uniform distribution on `[μ-α, μ+α]`.

```julia
PGeneralizedGaussian()           # GGD with shape 2, scale 1, location 0, (the Normal distribution)
PGeneralizedGaussian(α, μ, p)    # GGD with location α, scale μ, and shape p

params(d)                       # Get the parameters, i.e. (α, μ, p)
shape(d)                        # Get the shape parameter, p
scale(d)                        # Get the scale parameter, α
location(d)                     # Get the location parameter, μ
```

External Links
 * [Generalized Gaussian on Wikipedia](http://en.wikipedia.org/wiki/Generalized_normal_distribution)
 * [Reference implementation paper](https://www.researchgate.net/publication/254282790_Simulation_of_the_p-generalized_Gaussian_distribution)
 """
struct PGeneralizedGaussian{T1<:Real, T2<:Real, T3<:Real} <: ContinuousUnivariateDistribution
    μ::T1
    α::T2
    p::T3
    PGeneralizedGaussian{T1,T2,T3}(μ::T1,α::T2,p::T3) where {T1<:Real, T2<:Real, T3<:Real} = new{T1,T2,T3}(µ, α, p)
end

function PGeneralizedGaussian(μ::T1,α::T2,p::T3; check_args::Bool=true) where {T1<:Real, T2<:Real, T3<:Real}
    check_args && @check_args(PGeneralizedGaussian, α > zero(α) && p > zero(p))
    return PGeneralizedGaussian{T1,T2,T3}(μ,α,p)
end

"""
    PGeneralizedGaussian(p)

Builds a p-generalized Gaussian with `μ=0.0, α=1.0`
"""
PGeneralizedGaussian(p::Real; check_args::Bool=true) = PGeneralizedGaussian(zero(p), one(p), p; check_args=check_args)

"""
    PGeneralizedGaussian()

Builds a default p-generalized Gaussian with `μ=0.0, α=√2, p=2.0`, corresponding
to the normal distribution with `μ=0.0, σ=1.0`.
"""
PGeneralizedGaussian() = PGeneralizedGaussian(0.0, √2, 2.0, check_args=false) # approximate scale with unity std deviation and shape 2

#### Conversions

convert(::Type{PGeneralizedGaussian{T1,T2,T3}}, μ::S1, α::S2, p::S3) where {T1 <: Real, T2 <: Real, T3 <:Real, S1 <: Real, S2 <: Real, S3 <: Real} = PGeneralizedGaussian(T1(μ),T2(α),T3(p))
function convert(::Type{PGeneralizedGaussian{T1,T2,T3}}, d::PGeneralizedGaussian{S1,S2,S3}) where {T1 <: Real, T2 <: Real, T3 <: Real, S1 <: Real, S2 <: Real, S3 <: Real}
    return PGeneralizedGaussian(T1(d.μ), T2(d.α), T3(d.p), check_args=false)
end

@distr_support PGeneralizedGaussian -Inf Inf

#### Parameters
partype(::PGeneralizedGaussian{T1,T2,T3}) where {T1,T2,T3} = promote_type(T1,T2,T3)

params(d::PGeneralizedGaussian) = (d.μ, d.α, d.p)
location(d::PGeneralizedGaussian) = d.μ
shape(d::PGeneralizedGaussian) = d.p
scale(d::PGeneralizedGaussian) = d.α

#### Statistics

mean(d::PGeneralizedGaussian) = d.μ
median(d::PGeneralizedGaussian) = d.μ
mode(d::PGeneralizedGaussian) = d.μ

var(d::PGeneralizedGaussian) = (d.α^2) * (gamma(3.0 * inv(d.p)) / gamma(inv(d.p)))
std(d::PGeneralizedGaussian) = (d.α) * sqrt(gamma(3.0 * inv(d.p)) / gamma(inv(d.p)))

skewness(d::PGeneralizedGaussian{T1, T2, T3}) where {T1,T2,T3} = zero(T1)
kurtosis(d::PGeneralizedGaussian) = gamma(5.0 * inv(d.p)) * gamma(inv(d.p)) / (gamma(3.0 * inv(d.p))^2) - 3.0
entropy(d::PGeneralizedGaussian) = inv(d.p) - log( d.p / (2.0 * d.α * gamma(inv(d.p))))

#### Evaluation

function pdf(d::PGeneralizedGaussian, x::Real)
    (μ, α, p) = params(d)
    return ( p / ( 2.0 * α * gamma(1 / p) ) ) * exp( -( abs(x - μ) / α )^p )
end
logpdf(d::PGeneralizedGaussian, x::Real) = log(pdf(d, x))

"""
    cdf(d, x)

Calculates the CDF of the distribution. To determine the CDF, the incomplete
gamma function is required. The CDF  of the Gamma distribution provides this,
with the necessary 1/Γ(a) normalization.
"""
function cdf(d::PGeneralizedGaussian, x::Real)
    (μ, α, p) = params(d)
    v = cdf(Gamma(inv(p), 1), (abs(x - μ) / α)^p) * inv(2)
    return typeof(v)(1/2) + sign(x - μ) * v
end

#### Sampling

"""
    rand(rng, d)

Extract a sample from the p-Generalized Gaussian distribution 'd'. The sampling
procedure is implemented from from [1].
[1]  Gonzalez-Farias, G., Molina, J. A. D., & Rodríguez-Dagnino, R. M. (2009).
Efficiency of the approximated shape parameter estimator in the generalized
Gaussian distribution. IEEE Transactions on Vehicular Technology, 58(8),
4214-4223.
"""
function rand(rng::AbstractRNG, d::PGeneralizedGaussian)

    # utilizing the sampler from the Gamma distribution.
    g = Gamma(inv(d.p), 1)

    # random variable with value -1 or 1 with probability (1/2).
    b = 2.0 * rand(Bernoulli()) -1

    return d.μ + inv(sqrt(d.α)) * rand(rng, g)^inv(d.p) * b
end


export distance

"""
    distance(H1::AbstractHyperrectangle, H2::AbstractHyperrectangle;
             [p]::Real=2)

Compute the standard distance between two hyperrectangular sets, defined as

```math
    \\inf_{x \\in H_1, y \\in H_2} \\{ d(x, y) \\}.
```

### Input

- `H1` -- hyperrectangular set
- `H2` -- hyperrectangular set
- `p`  -- (optional; default: `2`) value of the ``p``-norm

### Output

The distance, which is zero if the sets intersect and otherwise the ``p``-norm
of the shortest line segment between any pair of points.

### Notes

See also [`hausdorff_distance`](@ref) for an alternative distance notion.
"""
function distance(H1::AbstractHyperrectangle, H2::AbstractHyperrectangle;
                  p::Real=2)
    n = dim(H1)
    @assert n == dim(H2) "incompatible set dimensions $n and $(dim(H2))"

    N = promote_type(eltype(H1), eltype(H2))
    d = Vector{N}(undef, n)
    @inbounds for i in 1:n
        # find relative position in dimension i
        # (if c1 == c2, the results are equivalent independent of the branch)
        if center(H1, i) >= center(H2, i)
            lhs = low(H1, i)
            rhs = high(H2, i)
        else
            lhs = low(H2, i)
            rhs = high(H1, i)
        end
        if _leq(lhs, rhs)
            d[i] = zero(N)
        else
            d[i] = rhs - lhs
        end
    end
    return norm(d, p)
end

"""
    distance(S::AbstractSingleton, X::LazySet; [p]::Real=2.0)

Compute the distance between the singleton `S` and the set `X` with respect to
the given `p`-norm.

### Input

- `S` -- singleton, i.e., a set with one element
- `X` -- set
- `p` -- (optional, default: `2.0`) the `p`-norm used; `p = 2.0` corresponds to
         the usual Euclidean norm

### Output

A scalar representing the distance between the element wrapped by `S` and the
set `X`.
"""
function distance(S::AbstractSingleton, X::LazySet; p::Real=2.0)
    return distance(element(S), X; p=p)
end

distance(X::LazySet, S::AbstractSingleton; p::Real=2.0) = distance(S, X; p=p)

distance(S1::AbstractSingleton, S2::AbstractSingleton; p::Real=2.0) =
    distance(element(S1), element(S2); p=p)

# disambiguation
distance(S::AbstractSingleton, H::AbstractHyperrectangle; p::Real=2.0) =
    distance(element(S), H; p=p)
distance(H::AbstractHyperrectangle, S::AbstractSingleton; p::Real=2.0) =
    distance(element(S), H; p=p)


# polar plots

#JL Plots
using Plots

r(θ) = 1 + cos(θ) * sin(θ)^2
Plots.plot(r, 0, 2π, proj=:polar, lims=(0,1.5))

r(θ) = 1 + cos(θ) * sin(θ)^2
θ = range(0, stop=2π, length=50)
Plots.plot(θ, r.(θ), proj=:polar, lims=(0,1.5), c=:steelblue,  linewidth=2)

# Polar with markers

Θ = range(0, stop = 2π, length = 100)
pr = abs.(0.1 * randn(100) + sin.(3Θ))
Plots.plot(Θ, pr, proj = :polar, m = 2, c=:steelblue,  linewidth=2)


#####################################################################
# Base functions.
#####################################################################
function sphere(x)
  sum(x.^2)
end

# Schwefel's ellipsoid.
function ellipsoid(x)
  res = 0
  for(i in 1:length(x))
    res += sum(x[1:i])^2
  end
  res
end

function elliptic(x)
  D = length(x)
  condition = 1e+6
  coefficients = condition .^ linspace(0, 1, D)
  sum(coefficients .* x.^2)
end

function rastrigin(x)
  D = length(x)
  10 * D + sum( x.^2 ) - 10 * sum( cos( 2 * π * x ) )
end

function ackley(x)
  D = length(x)
  try
    20 - 20.*exp(-0.2.*sqrt(sum(x.^2)/D)) - exp(sum(cos(2 * π * x))/D) + e
  catch e
    # Sometimes we have gotten a DomainError from the cos function so we protect this call
    println(e)
    println("For input x = ", x)
    # Return a very large fitness value to indicate that this is NOT the solution we want.
    # TODO: Fix better way to handle this!
    9.99e100
  end
end

function schwefel1_2(x)
  D = length(x)
  partsums = zeros(D)
  partsum = 0
  for i in 1:D
    partsum += x[i]
    partsums[i] = partsum
  end
  sum(partsums.^2)
end

function rosenbrock(x)
  n = length(x)
  return( sum( 100*( x[2:n] - x[1:(n-1)].^2 ).^2 + ( x[1:(n-1)] - 1 ).^2 ) )
end

function step(x)
  sum(ceil(x + 0.5))
end

function griewank(x)
  n = length(x)
  1 + (1/4000)*sum(x.^2) - prod(cos(x ./ sqrt(1:n)))
end

function schwefel2_22(x)
  ax = abs(x)
  sum(ax) + prod(ax)
end

function schwefel2_21(x)
  maximum(abs(x))
end

# I'm unsure about this one since it does not return the expected minima at
# [1.0, 1.0].
function schwefel2_26(x)
  D = length(x)
  418.98288727243369 * D - sum(x .* sin(sqrt(abs(x))))
end

function cigar(x)
  x[1]^2 + 1e6 * sum(x[2:end].^2)
end

function cigtab(x)
  x[1]^2 + 1e8 * x[end]^2 + 1e4 * sum(x[2:(end-1)].^2)
end

# Shekel10 is a 4D, multi-minima, non-separable test problem. Our implementation
# is based on the C code in:
#   http://www.math.ntu.edu.tw/~wwang/cola_lab/test_problems/multiple_opt/multiopt_prob/Shekel10/Shekel10.c
Shekel10_A = [4 4 4 4; 1 1 1 1; 8 8 8 8; 6 6 6 6; 3 7 3 7; 2 9 2 9; 5 5 3 3; 8 1 8 1; 6 2 6 2; 7 3.6 7 3.6]
Shekel10_C = [0.1, 0.2, 0.2, 0.4, 0.4, 0.6, 0.3, 0.7, 0.5, 0.5]
function shekel(x, a, c)
  sum = 0.0
  for i in 1:length(c)
    den = 0.0
    for j in 1:size(a, 2)
      den += (x[j] - a[i,j])^2
    end
    sum = sum - 1 / (den + c[i])
  end
  return sum
end
shekel10(x) = shekel(x, Shekel10_A, Shekel10_C)

# Shekel7 is a 4D, multi-minima, non-separable test problem. Our implementation
# is based on the C code in:
#   http://www.math.ntu.edu.tw/~wwang/cola_lab/test_problems/multiple_opt/multiopt_prob/Shekel7/Shekel7.c
Shekel7_A = [4 4 4 4; 1 1 1 1; 8 8 8 8; 6 6 6 6; 3 7 3 7; 2 9 2 9; 5 5 3 3]
Shekel7_C = [0.1, 0.2, 0.2, 0.4, 0.4, 0.6, 0.3]
shekel7(x) = shekel(x, Shekel7_A, Shekel7_C)

# Shekel5 is a 4D, multi-minima, non-separable test problem. Our implementation
# is based on the C code in:
#   http://www.math.ntu.edu.tw/~wwang/cola_lab/test_problems/multiple_opt/multiopt_prob/Shekel5/Shekel5.c
Shekel5_A = [4 4 4 4; 1 1 1 1; 8 8 8 8; 6 6 6 6; 3 7 3 7]
Shekel5_C = [0.1, 0.2, 0.2, 0.4, 0.4]
shekel5(x) = shekel(x, Shekel5_A, Shekel5_C)

# Hartman 6D is a multi-minima, non-separable test problem. Our implementation is based on:
#  http://www.sfu.ca/~ssurjano/hart6.html
Hartman6_alpha = [1.0 1.2 3.0 3.2]
Hartman6_A = [10 3 17 3.50 1.7 8; 0.05 10 17 0.1 8 14; 3 3.5 1.7 10 17 8; 17 8 0.05 10 0.1 14]
Hartman6_P = 1e-4 * [1312 1696 5569 124 8283 5886; 2329 4135 8307 3736 1004 9991; 2348 1451 3522 2883 3047 6650; 4047 8828 8732 5743 1091 381]
function hartman(x, alpha, A, P)
  sum = 0.0
  for i in 1:length(alpha)
    isum = 0.0
    for j in 1:size(A, 2)
      isum += A[i,j] * (x[j] - P[i,j])^2
    end
    sum -= alpha[i] * exp(-isum)
  end
  sum
end
hartman6(x) = hartman(x, Hartman6_alpha, Hartman6_A, Hartman6_P)

# Hartman 3D is a multi-minima, non-separable test problem. Our implementation is based on:
#  http://www.sfu.ca/~ssurjano/hart3.html
# However, we get a different global minima than the one stated on that page.
Hartman3_alpha = [1.0 1.2 3.0 3.2]
Hartman3_A = [3.0 10 30; 0.1 10 35; 3.0 10 30; 0.1 10 36]
Hartman3_P = 1e-4 * [3689 1170 2673; 4699 4387 7470; 1091 8732 5547; 381 5743 8828]
hartman3(x) = hartman(x, Hartman3_alpha, Hartman3_A, Hartman3_P)
# The minima should be -3.86278 but we get a different one so use that in the problem spec:

#####################################################################
# S2 functions in addition to the base functions above. As stated
# in table II of the JADE paper: http://150.214.190.154/EAMHCO/pdf/JADE.pdf
#####################################################################

function quartic(x)
  D = length(x)
  sum( (1:D) .* x.^4 )
end

function noisy_quartic(x)
  quartic(x) + rand()
end

function s2_step(x)
  sum( ceil(x + 0.5).^2 )
end

#####################################################################
# Misc other interesting optimization functions and families.
#####################################################################

# Schwefel 2.13 is a hard one...
#function f=schwefel_213(x)
#global initial_flag
#persistent a b A alpha
#[ps,D]=size(x);
#if initial_flag==0
#    initial_flag=1;
#    load schwefel_213_data
#    if length(alpha)>=D
#        alpha=alpha(1:D);a=a(1:D,1:D);b=b(1:D,1:D);
#    else
#        alpha=-3+6*rand(1,D);
#        a=round(-100+200.*rand(D,D));
#        b=round(-100+200.*rand(D,D));
#    end
#    alpha=repmat(alpha,D,1);
#    A=sum(a.*sin(alpha)+b.*cos(alpha),2);
#end
#for i=1:ps
#    xx=repmat(x(i,:),D,1);
#    B=sum(a.*sin(xx)+b.*cos(xx),2);
#    f(i,1)=sum((A-B).^2,1);
#end

# This is a generator for the family of deceptive functions from the
# Cuccu2011 paper on novelty-based restarts. We have vectorized it to allow
# more than 1D versions. The Cuccu2011 paper uses the following values for
# (l, w) = [(5, 0),  (15, 0),  (30, 0),
#           (5, 2),  (15, 2),  (30, 2),
#           (5, 10), (15, 10), (30, 10)]
# and notes that (15, 2) and (30, 2) are the most difficult instances.
function deceptive_cuccu2011(l, w)
  (x) -> begin
    absx = abs(x)
    sumabsx = sum(absx)
    if sumabsx <= 1
      return sum(x.^2)
    elseif sumabsx >= l+1
      return sum((absx - l).^2)
    else
      return (1 - 0.5 * sum(sin( (π * w * (absx - 1)) / l ).^2))
    end
  end
end

# Deceptive/hardest instances:
deceptive_cuccu2011_15_2 = deceptive_cuccu2011(15, 2)
deceptive_cuccu2011_30_2 = deceptive_cuccu2011(30, 2)


# ---
# title: 1491. Average Salary Excluding the Minimum and Maximum Salary
# id: problem1491
# author: zhwang
# date: 2022-03-18
# difficulty: Easy
# categories: Array, Sort
# link: <https://leetcode.com/problems/average-salary-excluding-the-minimum-and-maximum-salary/description/>
# hidden: true
# ---
#
# Given an array of **unique** integers `salary` where `salary[i]` is the salary
# of the employee `i`.
#
# Return the average salary of employees excluding the minimum and maximum
# salary.
#
#
#
# **Example 1:**
#
#
#
#     Input: salary = [4000,3000,1000,2000]
#     Output: 2500.00000
#     Explanation: Minimum salary and maximum salary are 1000 and 4000 respectively.
#     Average salary excluding minimum and maximum salary is (2000+3000)/2= 2500
#
#
# **Example 2:**
#
#
#
#     Input: salary = [1000,2000,3000]
#     Output: 2000.00000
#     Explanation: Minimum salary and maximum salary are 1000 and 3000 respectively.
#     Average salary excluding minimum and maximum salary is (2000)/1= 2000
#
#
# **Example 3:**
#
#
#
#     Input: salary = [6000,5000,4000,3000,2000,1000]
#     Output: 3500.00000
#
#
# **Example 4:**
#
#
#
#     Input: salary = [8000,9000,2000,3000,6000,1000]
#     Output: 4750.00000
#
#
#
#
# **Constraints:**
#
#   * `3 <= salary.length <= 100`
#   * `10^3 <= salary[i] <= 10^6`
#   * `salary[i]` is unique.
#   * Answers within `10^-5` of the actual value will be accepted as correct.
#
#
## @lc code=start
using LeetCode

function average_salary_excluding_mami(salary::Vector{Int})
    return (sum(salary) - minimum(salary) - maximum(salary)) / (length(salary) - 2)
end
## @lc code=end


using UncertaintyQuantification, Random
Random.seed!(2547);
N = 200
mu = [0, 0]
sig = [1 -0.55; -0.55 1]
mvnormal = MvNormal(mu, sig)

sam = zeros(200,2)
for i in 1:200 sam[i,:] = rand(mvnormal) end

Random.seed!()

function likelihood(x)
    if x[1] >1 || x[1]<-1
        rv = 0
    else
        rv = 1
        for i in 1:N
            rv = rv/(2*pi*sqrt(1-x[1]^2))*exp(-1/(2*(1-x[1]^2))*(sam[i,1]^2-2*x[1]*
            sam[i,1]*sam[i,2]+sam[i,2]^2))
        end
    end
    return rv
end

function prior(x)
    if x[1] >1 || x[1]<-1
        rv = 0
    else
        rv = 1/(1-x[1]^2)^(3/2)
    end
    return rv
end
priorsample(n) = rand(Uniform(-1,1),(n))

tmcmcsample = tmcmc(likelihood, prior, priorsample, 2000, 0.2)


using LinearAlgebra
using Plots

#%
#% The 'illustrative example' from Section 1.4, Chapter 1.
#%
#% written by John Bardsley 2016.
#%
#% Define the matrix A using the SVD
v1 = [1/sqrt(2) 1/sqrt(2)]';
v2 = [-1/sqrt(2) 1/sqrt(2)]';
s1 = 1;
s2 = 1e-2;
A  = s1*v1*v1'+s2*v2*v2';

#% Define true x and noise free b
x   = [1 1]';
b_e = A*x;

#% Create realization from the data model
nsamp  = 10000;
σ  = 0.1;
b_samp = repeat(b_e,1,nsamp)+σ*randn(2,nsamp);

#figure(1),
  scatter(b_samp[1,:],b_samp[2,:]) #,'k*')
#% Create the corresponding least squares solutions
x_LS = A\b_samp;
#figure(2)
  scatter(x_LS[1,:],x_LS[2,:]) #,'k*')

#% Check ill-posedness inequality (1.28)
ill_posed_lhs = 0.5*(1/s1+1/s2)
ill_posed_rhs = norm(x)/norm(b_e)


import RowEchelon
import LinearAlgebra

#=
Homework 7
Section 3.2
Problems 10
=#

A = [1 3 -1 0 -2 ;
     0 2 -4 -1 -6 ;
     -2 -6 2 3 9 ;
     3 7 -3 8 -7 ;
     3 5 5 2 7]

detA = LinearAlgebra.det(A)

println("Original Matrix:")
println(A)
println("Determinate:")
println(detA)


function draw_wireframe_sphere(pos, r; npts=11, color=(255,255,255), name=nothing, opacity=1.)
    ϕs = range(0, stop=2π, length=npts)
    θs = range(-π/2, stop=π/2, length=npts)
    # vertical
    v_xs = [[r * cos(θ) * sin(ϕ) + pos[1] for θ in θs] for ϕ in ϕs[1:end-1]]
    v_ys = [[r * cos(θ) * cos(ϕ) + pos[2] for θ in θs] for ϕ in ϕs[1:end-1]]
    v_zs = [[r * sin(θ) + pos[3] for θ in θs] for ϕ in ϕs[1:end-1]]
    # horizontal
    h_xs = [[r * cos(θ) * sin(ϕ) + pos[1] for ϕ in ϕs] for θ in θs]
    h_ys = [[r * cos(θ) * cos(ϕ) + pos[2] for ϕ in ϕs] for θ in θs]
    h_zs = [[r * sin(θ) + pos[3] for ϕ in ϕs] for θ in θs]

    hover = isnothing(name) ? "skip" : nothing

    xs, ys, zs = Float64[], Float64[], Float64[]
    for i=1:length(v_xs)
        append!(xs, v_xs[i])
        push!(xs, NaN)
        append!(ys, v_ys[i])
        push!(ys, NaN)
        append!(zs, v_zs[i])
        push!(zs, NaN)
    end
    for j=1:length(h_xs)
        append!(xs, h_xs[j])
        push!(xs, NaN)
        append!(ys, h_ys[j])
        push!(ys, NaN)
        append!(zs, h_zs[j])
        push!(zs, NaN)
    end

    return scatter3d(;x=xs, y=ys, z=zs,
                      mode="lines",
                      line=attr(color="rgb$color"),
                      opacity=opacity,
                      showlegend=false, showscale=false, name=name,
                      hovertemplate=name, hoverinfo=hover,
            )
end

draw_central_body(bd::Union{CelestialBody,Star}; kwargs...
    ) = draw_wireframe_sphere(SVector{3,Float64}(0.,0.,0.), bd.eqradius; color=bd.color, name=bd.name, kwargs...)

draw_orbiting_body(bd::CelestialBody, t; kwargs...
    ) = draw_wireframe_sphere(time_orbital_position(t, bd.orbit), bd.eqradius; color=bd.color, name=bd.name, kwargs...)

draw_soi(bd::CelestialBody, t; kwargs...
    ) = draw_wireframe_sphere(time_orbital_position(t, bd.orbit), bd.SoI; color=bd.color, name=bd.name, opacity=0.25, kwargs...)


n = rand(1:10)

A = rand(n,n)

r = polarfact(A, alg = :newton);

# Test unitary matrix U

U = r.U
H = r.H

@test_approx_eq_eps U'*U eye(n) 1e-7

# Test Hermitian positive semifefinite matrix H

@test issym(H)

for i in eigvals(H)
    @test i >= 0.
end

@test_approx_eq A U*H

println("Newton method passed test...")


# Test script to test solver for a qp

using Test
using COSMO, SparseArrays, LinearAlgebra

# Quadratic program example from OSQP Doc
# min 0.5 * x'Px +  q'x
# s.t. l <= Ax <= u
q = [1; 1.]
P = sparse([4. 1; 1 2])
A = [1. 1; 1 0; 0 1]
l = [1.; 0; 0]
u = [1; 0.7; 0.7]

# Define the constraint l <= Ax <= u with the help of a Nonnegatives set
Aa = [-A; A]
ba = [u; -l]
constraint1 = COSMO.Constraint(Aa, ba, COSMO.Nonnegatives)

# define example problem
settings = COSMO.Settings(verbose=true, eps_abs = 1e-4, eps_rel = 1e-4)

model = COSMO.Model()
assemble!(model, P, q, constraint1, settings = settings)
res = COSMO.optimize!(model);

# solve again by defining the constraints with the help of a box
constraint1 = COSMO.Constraint(A, zeros(3), COSMO.Box(l, u))

model = COSMO.Model()
assemble!(model, P, q, constraint1, settings = settings)
res_box = COSMO.optimize!(model);

@testset "QP Problem" begin
  @test norm(res.x[1:2] - [0.3; 0.7], Inf) < 1e-3
  @test norm(res_box.x[1:2] - [0.3; 0.7], Inf) < 1e-3
  @test abs(res.obj_val - 1.88) < 1e-3
  @test abs(res_box.obj_val - 1.88) < 1e-3
end
nothing


@testset "Product with non-commutative monomials and variables" begin
    Mod.@ncpolyvar x y z
    m1 = x * (y * x)
    @test variables(m1)[1] == x
    @test variables(m1)[2] == y
    @test variables(m1)[3] == x
    @test all(exponents(m1) .== 1)
    m2 = y * (y * x)
    @test variables(m2)[1] == y
    @test variables(m2)[2] == x
    @test collect(exponents(m2)) == [2, 1]
    for m3 in [(x*x)*(y^2*x), (x*x*y)*(y*x), x*(x*y^2*x), (x^2*y^2)*x]
        @test variables(m3)[1] == x
        @test variables(m3)[2] == y
        @test variables(m3)[3] == x
        @test collect(exponents(m3)) == [2, 2, 1]
    end
    mv = monovec([x, y, x*y])
    @test x * mv == [x^2, x*y, x^2*y]
    @test mv * x == [x^2, y*x, x*y*x]
    @test y^2 * mv == [y^2*x, y^3, y^2*x*y]
    @test mv * y^2 == [x*y^2, y^3, x*y^3]
    m4 = y*x^2
    @test variables(m4)[1] == y
    @test variables(m4)[2] == x
    @test collect(exponents(m4)) == [1, 2]

    @test (x*y) * (3x^2 + 2x*y + 4y^2) == 3x*y*x^2 + 2x*y*x*y + 4x*y^3

    P = [2 3 4;
         3 4 5;
         4 5 6]
    Q = [4 3 5;
         3 2 4
         5 4 6]
    p = @inferred polynomial(P, [x*y, x^2, y^2])
    @test coefficients(p) == [4, 3, 5, 4, 3, 2, 6, 5, 4]
    @test monomials(p) == [x^4, x^3*y, x^2*y^2, x*y^3, x*y*x^2, x*y*x*y, y^4, y^2*x^2, y^2*x*y]
    p = @inferred polynomial(Q, monovec([x*y, x^2, y^2]))
    @test coefficients(p) == [4, 3, 5, 4, 3, 2, 6, 5, 4]
    @test monomials(p) == [x^4, x^3*y, x^2*y^2, x*y^3, x*y*x^2, x*y*x*y, y^4, y^2*x^2, y^2*x*y]
    q = @inferred polynomial(zeros(Float64, 0, 0), emptymonovec(x))
    @test q isa AbstractPolynomialLike{Float64}
    @test q == 0
end
@testset "Noncommutative quadratic" begin
    Mod.@ncpolyvar x[1:2]
    Q = Hermitian([1 2 + 3im; 2 - 3im 4])
    p = 1x[1]^2 + (2 + 3im) * x[1] * x[2] + (2 - 3im) * x[2] * x[1] + 4x[2]^2
    @test polynomial(Q, x) == p
    @test polynomial(Q, monovec(x)) == p
end


module FreudenthalTriangulations

using LinearAlgebra
using SparseArrays

export FreudenthalTriangulation, vertices, simplex, belief_vertices, belief_simplex

# Construct the Freudenthal triangulation of the belief simplex

struct FreudenthalTriangulation
        n::Int
        m::Int
end

"""
    vertices(T::FreudenthalTriangulation)
Construct the list of Freudenthal vertices in an `n` dimensional space with grid resolution `m`.
The vertices are represented by a list of `n` dimensional vectors.
"""
function vertices(T::FreudenthalTriangulation)
    V = Vector{Int}[]
    v = Vector{Int}(undef, T.n)
    v[1] = T.m
    _vertices!(V, v, 2)
    return V
end

function _vertices!(V::Vector{Vector{Int64}}, v::Vector{Int64}, i::Int64)
    n = length(v)
    if i > n
        push!(V, copy(v))
        return
    end
    for k in 0 : v[i-1]
        v[i] = k
        _vertices!(V, v, i+1)
    end
end

"""
    _freudenthal_simplex(x::Vector{Float64})
Returns the list of vertices of the simplex of point `x` in the Freudenthal grid.
"""
function _freudenthal_simplex(x::Vector{Float64})
    n = length(x)
    V = Vector{Vector{Int}}(undef, n+1)
    V[1] = floor.(Int, x)
    d = x - V[1]
    p = sortperm(d, rev=true)
    for i in 2 : n+1
        V[i] = copy(V[i-1])
        V[i][p[i-1]] += 1
    end
    return V
end

"""
    _barycentric_coordinates(x::Vector{Float64}, V::Vector{Vector{Int64}})
Given a point `x` and its simplex `V` in the Freudenthal grid, returns the barycentric coordinates
of `x` in the grid. `V` must be in the same order as provided by the output of `freudenthal_simplex`
"""
function _barycentric_coordinates(x::Vector{Float64}, V::Vector{Vector{Int64}})
    d = x - V[1]
    p = sortperm(d, rev=true)
    n = length(x)
    λ = Vector{Float64}(undef, n+1)
    λ[n+1] = d[p[n]]
    for i in n:-1:2
        λ[i] = d[p[i-1]] - d[p[i]]
    end
    λ[1] = 1.0 - sum(λ[2:end])
    return λ
end

"""
    simplex(T::FreudenthalTriangulation, x::Vector{Float64})
Given a point `x`, returns the simplex of the point `x` and the barycentric coordinates of `x` in the grid.
"""
function simplex(T::FreudenthalTriangulation, x::Vector{Float64})
    V = _freudenthal_simplex(x)
    return V, _barycentric_coordinates(x, V)
end

"""
    _to_belief(x)
Transform a point `x` in the Freudenthal space to a point in the belief space.
"""
_to_belief(x) = (push!(x[1:end-1] - x[2:end], x[end]))./x[1]

"""
    _to_freudenthal(b, m::Int64)
Transform a point `b` in the belief space to a point in the Freudenthal space.
`m` is the resolution of the Freudenthal grid.
"""
_to_freudenthal(b, m::Int64) = [sum(b[k] for k in i : length(b))*m for i in 1 : length(b)]

"""
    belief_vertices(T::FreudenthalTriangulation)
    Converts the list of Freudenthal vertices with dimension `T.n` and granularity `T.m` into
    belief space and returns the list
"""
belief_vertices(T::FreudenthalTriangulation) = _to_belief.(vertices(T))

"""
    belief_simplex(T::FreudenthalTriangulation, b)
    Constructs the belief simplex surrounding a belief `b` and the barycentric coordinates
    corresponding to the location of the belief within the simplex
"""
function belief_simplex(T::FreudenthalTriangulation, b)
    x = _to_freudenthal(b, T.m)
    V, λ = simplex(T, x)
    B = _to_belief.(V)
    valid =  λ .> sqrt(eps())
    return B[valid], λ[valid]
end

end # module


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

function _flux2mag(flux::T, zero_point::T, ABwave::T) where {T<:AbstractFloat}
    if isnan(ABwave)
        return -2.5*log10(flux) - zero_point
    else
        return -2.5*log10(flux) - 5 * log10(float(ABwave)) - 2.406
    end
end

"""
    flux2mag(flux[, zero_point, ABwave=number]) -> magnitude

### Purpose ###

Convert from flux expressed in erg/(s cm² Å) to magnitudes.

### Explanation ###

This is the reverse of `mag2flux`.

### Arguments ###

* `flux`: the flux to be converted in magnitude, expressed in
  erg/(s cm² Å).
* `zero_point`: the zero point level of the magnitude.  If not
 supplied then defaults to 21.1 (Code et al 1976).  Ignored if the `ABwave`
 keyword is supplied
* `ABwave` (optional numeric keyword): wavelength in Angstroms.
 If supplied, then returns Oke AB magnitudes (Oke & Gunn 1983, ApJ, 266, 713;
 http://adsabs.harvard.edu/abs/1983ApJ...266..713O).

### Output ###

The magnitude.

If the `ABwave` keyword is set then magnitude is given by the expression

\$\$\\text{ABmag} = -2.5\\log_{10}(f) - 5\\log_{10}(\\text{ABwave}) - 2.406\$\$

Otherwise, magnitude is given by the expression

\$\$\\text{mag} = -2.5\\log_{10}(\\text{flux}) - \\text{zero point}\$\$

### Example ###

```jldoctest
julia> using AstroLib

julia> flux2mag(5.2e-15)
14.609991640913002

julia> flux2mag(5.2e-15, 15)
20.709991640913003

julia> flux2mag(5.2e-15, ABwave=15)
27.423535345634598
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
flux2mag(flux::Real, zero_point::Real=21.1; ABwave::Real=NaN) =
    _flux2mag(promote(float(flux), float(zero_point), float(ABwave))...)


using DifferentialEquations
using Plots
gr()

function lif(u,p,t);
    gL, EL, C, Vth, I = p
    (-gL*(u-EL)+I)/C
end

function thr(u,t,integrator)
    integrator.u > integrator.p[4]
end

function reset!(integrator)
    integrator.u = integrator.p[2]
end

threshold = DiscreteCallback(thr,reset!)
current_step= PresetTimeCallback([2,15],integrator -> integrator.p[5] += 210.0)
cb = CallbackSet(current_step,threshold)

u0 = -75
tspan = (0.0, 40.0)
# p = (gL, EL, C, Vth, I)
p = [10.0, -75.0, 5.0, -55.0, 0]

prob = ODEProblem(lif, u0, tspan, p, callback=cb)

sol = solve(prob)

plot(sol)

#Izhikevichch Model
using DifferentialEquations
using Plots

function izh!(du,u,p,t);
    a, b, c, d, I = p

    du[1] = 0.04*u[1]^2+5*u[1]+140-u[2]+I
    du[2] = a*(b*u[1]-u[2])
end

function thr(u,t,integrator)
    integrator.u[1] >= 30
end

function reset!(integrator)
    integrator.u[1] = integrator.p[3]
    integrator.u[2] += integrator.p[4]
end

threshold = DiscreteCallback(thr,reset!)
current_step= PresetTimeCallback(50,integrator -> integrator.p[5] += 10)
cb = CallbackSet(current_step,threshold)

p = [0.02, 0.2, -50, 2, 0]
u0 = [-65, p[2]*-65]
tspan = (0.0, 300)

prob = ODEProblem(izh!, u0, tspan, p, callback=cb)

sol = solve(prob);
plot(sol, vars=1)

plot(sol, vars=2)

p = [0.02, 0.2, -65, 8, 0]
u0 = [-65, p[2]*-65]
tspan = (0.0, 300)

prob = ODEProblem(izh!, u0, tspan, p, callback=cb)
sol = solve(prob);
plot(sol, vars=1)

using DifferentialEquations
using Plots

# Potassium ion-channel rate functions
alpha_n(v) = (0.02 * (v - 25.0)) / (1.0 - exp((-1.0 * (v - 25.0)) / 9.0))
beta_n(v) = (-0.002 * (v - 25.0)) / (1.0 - exp((v - 25.0) / 9.0))

# Sodium ion-channel rate functions
alpha_m(v) = (0.182*(v + 35.0)) / (1.0 - exp((-1.0 * (v + 35.0)) / 9.0))
beta_m(v) = (-0.124 * (v + 35.0)) / (1.0 - exp((v + 35.0) / 9.0))

alpha_h(v) = 0.25 * exp((-1.0 * (v + 90.0)) / 12.0)
beta_h(v) = (0.25 * exp((v + 62.0) / 6.0)) / exp((v + 90.0) / 12.0)

function HH!(du,u,p,t);
    gK, gNa, gL, EK, ENa, EL, C, I = p
    v, n, m, h = u

    du[1] = (-(gK * (n^4.0) * (v - EK)) - (gNa * (m ^ 3.0) * h * (v - ENa)) - (gL * (v - EL)) + I) / C
    du[2] = (alpha_n(v) * (1.0 - n)) - (beta_n(v) * n)
    du[3] = (alpha_m(v) * (1.0 - m)) - (beta_m(v) * m)
    du[4] = (alpha_h(v) * (1.0 - h)) - (beta_h(v) * h)
end

current_step= PresetTimeCallback(100,integrator -> integrator.p[8] += 1)

# n, m & h steady-states
n_inf(v) = alpha_n(v) / (alpha_n(v) + beta_n(v))
m_inf(v) = alpha_m(v) / (alpha_m(v) + beta_m(v))
h_inf(v) = alpha_h(v) / (alpha_h(v) + beta_h(v))

p = [35.0, 40.0, 0.3, -77.0, 55.0, -65.0, 1, 0]
u0 = [-60, n_inf(-60), m_inf(-60), h_inf(-60)]
tspan = (0.0, 1000)

prob = ODEProblem(HH!, u0, tspan, p, callback=current_step)

sol = solve(prob);
plot(sol, vars=1)

plot(sol, vars=[2,3,4], tspan=(105.0,130.0))

function gSyn(max_gsyn, tau, tf, t);
    if t-tf >= 0
        return max_gsyn * exp(-(t-tf)/tau)
    else
        return 0.0
    end
end
function HH!(du,u,p,t);
    gK, gNa, gL, EK, ENa, EL, C, I, max_gSyn, ESyn, tau, tf = p
    v, n, m, h = u

    ISyn = gSyn(max_gSyn, tau, tf, t) * (v - ESyn)

    du[1] = (-(gK * (n^4.0) * (v - EK)) - (gNa * (m ^ 3.0) * h * (v - ENa)) - (gL * (v - EL)) + I - ISyn) / C
    du[2] = (alpha_n(v) * (1.0 - n)) - (beta_n(v) * n)
    du[3] = (alpha_m(v) * (1.0 - m)) - (beta_m(v) * m)
    du[4] = (alpha_h(v) * (1.0 - h)) - (beta_h(v) * h)
end

p = [35.0, 40.0, 0.3, -77.0, 55.0, -65.0, 1, 0, 0.008, 0, 20, 100]
tspan = (0.0, 200)
prob = ODEProblem(HH!, u0, tspan, p)
sol = solve(prob);
plot(sol, vars=1)

p = [35.0, 40.0, 0.3, -77.0, 55.0, -65.0, 1, 0, 0.01, 0, 20, 100]
tspan = (0.0, 200)
prob = ODEProblem(HH!, u0, tspan, p)
sol = solve(prob);
plot!(sol, vars=1)

function HH!(du,u,p,t);
    gK, gNa, gL, EK, ENa, EL, C, I, tau, tau_u, tau_R, u0, gmax, Esyn  = p
    v, n, m, h, u, R, gsyn = u

    du[1] = ((gK * (n^4.0) * (EK - v)) + (gNa * (m ^ 3.0) * h * (ENa - v)) + (gL * (EL - v)) + I + gsyn * (Esyn - v)) / C
    du[2] = (alpha_n(v) * (1.0 - n)) - (beta_n(v) * n)
    du[3] = (alpha_m(v) * (1.0 - m)) - (beta_m(v) * m)
    du[4] = (alpha_h(v) * (1.0 - h)) - (beta_h(v) * h)

    # Synaptic variables
    du[5] = -(u/tau_u)
    du[6] = (1-R)/tau_R
    du[7] = -(gsyn/tau)
end

function epsp!(integrator);
    integrator.u[5] += integrator.p[12] * (1 - integrator.u[5])
    integrator.u[7] += integrator.p[13] * integrator.u[5] * integrator.u[6]
    integrator.u[6] -= integrator.u[5] * integrator.u[6]

end

epsp_ts= PresetTimeCallback(100:100:500, epsp!)

p = [35.0, 40.0, 0.3, -77.0, 55.0, -65.0, 1, 0, 30, 1000, 50, 0.5, 0.005, 0]
u0 = [-60, n_inf(-60), m_inf(-60), h_inf(-60), 0.0, 1.0, 0.0]
tspan = (0.0, 700)
prob = ODEProblem(HH!, u0, tspan, p, callback=epsp_ts)
sol = solve(prob);
plot(sol, vars=1)

plot(sol, vars=7)

plot(sol, vars=[5,6])

epsp_ts= PresetTimeCallback(100:1000:5100, epsp!)

p = [35.0, 40.0, 0.3, -77.0, 55.0, -65.0, 1, 0, 30, 500, 50, 0.5, 0.005, 0]
u0 = [-60, n_inf(-60), m_inf(-60), h_inf(-60), 0.0, 1.0, 0.0]
tspan = (0.0, 5300)
prob = ODEProblem(HH!, u0, tspan, p, callback=epsp_ts)
sol = solve(prob);
plot(sol, vars=7)

plot(sol, vars=[5,6])

epsp_ts= PresetTimeCallback(100:100:500, epsp!)

p = [35.0, 40.0, 0.3, -77.0, 55.0, -65.0, 1, 0, 30, 100, 1000, 0.5, 0.005, 0]
u0 = [-60, n_inf(-60), m_inf(-60), h_inf(-60), 0.0, 1.0, 0.0]
tspan = (0.0, 700)
prob = ODEProblem(HH!, u0, tspan, p, callback=epsp_ts)
sol = solve(prob);
plot(sol, vars=7)

plot(sol, vars=[5,6])

using SciMLTutorials
SciMLTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    CubicVariogram(range=r, sill=s, nugget=n)
    CubicVariogram(ball; sill=s, nugget=n)

A cubic variogram with range `r`, sill `s` and nugget `n`.
Optionally, use a custom metric `ball`.
"""
struct CubicVariogram{V,B} <: Variogram
  sill::V
  nugget::V
  ball::B
end

CubicVariogram(ball; sill=1.0, nugget=zero(typeof(sill))) =
  CubicVariogram(sill, nugget, ball)

CubicVariogram(; range=1.0, sill=1.0, nugget=zero(typeof(sill))) =
  CubicVariogram(sill, nugget, MetricBall(range))

function (γ::CubicVariogram)(h::T) where {T}
  r = radius(γ.ball)
  s = γ.sill
  n = γ.nugget

  # constants
  c1 = T(35) / T(4)
  c2 = T(7)  / T(2)
  c3 = T(3)  / T(4)
  s1 = 7*(h/r)^2 - c1*(h/r)^3 + c2*(h/r)^5 - c3*(h/r)^7
  s2 = T(1)

  (h < r) * (s - n) * s1 +
  (h ≥ r) * (s - n) * s2 +
  (h > 0) * n
end

isstationary(::Type{<:CubicVariogram}) = true


function Test_AdaptiveSteppers(t0, tf)
    function func(t, y)
        [exp(t)]
    end

    y0, exact = [exp(t0)], exp(tf)
    methods = DopriStep, CarpStep, RKFStep

    for a=methods
        ts, ys = AdaptiveStepper(func, a, t0, copy(y0), tf)
        println("Actual error of ", a, ": ", abs(exact - ys[end][1]))
    end
end

#Test_AdaptiveSteppers(20, 21)

#### Systems ####

function SimpleSystem(t, y)
    [cos(t)]
end

function HarmonicOscillator(t, y)
    y1, ydot = y
    [ydot, 100*cos(20*t) - ydot/4.0 - y1]
end

function FourthOrderSystem(t, y)
    y1, ydot, yddot, ydddot = y
    [ydot, yddot, ydddot, exp(t)*cos(100*t)-7*ydddot-17*yddot-17*ydot-6*y1]
end

#t0, tf = 0, 15
#y0 = [0, 0, 0, 0]

#ts, ys = Stepper(FourthOrderSystem, BackwardEulerStep, t0, y0, tf, 0.01)
#ts, ys = MultiStepper(FourthOrderSystem, AdamsBash4, t0, y0, tf, 0.001)
#plot(ts, [a[1] for a=ys], "o-")


export leaky_relu

doc"""
    leaky_relu(x::Var, alpha=0.1)

Leaky Rectified Linear Unit.

```math
f(x) =
\begin{cases}
x & x > 0 \\
\alpha x & x \leq 0
\end{cases}
```

# References
* Maas et al., ["Rectifier Nonlinearities Improve Neural Network Acoustic Models"](http://web.stanford.edu/~awni/papers/relu_hybrid_icml2013_final.pdf), ICML 2013.
"""
leaky_relu(x::Var, alpha=0.1) = Var(leaky_relu.(x.data,eltype(x)(alpha)), (leaky_relu,x,alpha))
leaky_relu(x::T, alpha::T) where T = x >= zero(T) ? x : x*alpha
leaky_relu(x::Node, alpha=0.1) = Node(leaky_relu, x, alpha)

function addgrad!(y::Var, ::typeof(leaky_relu), x::Var, alpha::Float64)
    isvoid(x.grad) && return
    ∇leaky_relu!(y.grad, x.data, x.grad, eltype(x)(alpha))
end

function ∇leaky_relu!(gy::Array{T}, x::Array{T}, gx::Array{T}, alpha::T) where T
    @inbounds for i = 1:length(x)
        gx[i] += x[i] >= zero(T) ? gy[i] : gy[i]*alpha
    end
end


function nlsolve(f, j, x, iterations = 10^5, r_norm = norm, r_abstol = sqrt(eps(eltype(x))))
    xnext = x
    r = f(x)
    for i = 1:iterations
        if r_norm(r) <= r_abstol
            return r, xnext, :success
        end
        x = copy(xnext)
        s = - j(x)\r
        xnext = x + s
        r = f(xnext)
    end
    if r_norm(r) <= r_abstol
        return r, xnext, :success
    end
    r, xnext, :failure
end


#towel map:
println("Testing discrete system evolution...")
using StaticArrays, Base.Test, DynamicalSystems

@testset "Logistic Map" begin

  d1 = Systems.logistic(0.1)
  d2 = DiscreteDS1D(0.1, d1.eom)
  d3 = DiscreteDS1D(big(0.1), d1.eom, d1.deriv)

  @testset "Evolution & Timeseries" begin
    dd1 = evolve!(d1)
    dd2 = evolve!(d2)
    dd3 = evolve!(d3)
    @test dd1.state == dd2.state
    @test dd1.state ≈ dd3.state
    @test typeof(dd3.state) == BigFloat
    ts1 = timeseries(dd1, 100, mutate = false)
    ts3 = timeseries(dd3, 100, mutate = false)
    @test ts1[10] ≈ ts3[10]
    @test eltype(ts3) == BigFloat
  end
  @testset "Derivatives" begin
    f1 = d1.deriv(d1.state)
    f2 = d2.deriv(d2.state)
    f3 = d3.deriv(d3.state)
    @test isapprox(f1, f2;rtol = 1e-12)
    @test isapprox(f1, f3;rtol = 1e-12)
    @test typeof(f3) == BigFloat
  end
end

@testset "Folded-Towel Map" begin

  s1 = Systems.towel(0.1ones(3))
  s2 = DiscreteDS(0.1ones(3), s1.eom)
  s4 = DiscreteDS(round.(big.(0.1ones(3)),3), s1.eom, s1.jacob)

  @testset "Evolution & Timeseries" begin
    st1 = evolve(s1)
    st2 = evolve(s2)
    st4 = evolve(s4)

    @test isapprox.(st1, st2; rtol = 1e-12) == trues(s1.state)
    @test isapprox.(st1, st4; rtol = 1e-12) == trues(s1.state)

    evolve!(s1); evolve!(s2); evolve!(s4);
    @test s1.state ≈ s2.state
    @test s2.state ≈ s4.state

    ts = timeseries(s1, 100; mutate = false)
    @test size(ts) == (100, 3)
    ts4 = timeseries(s4, 100; mutate = false)
    @test size(ts4) == (100, 3)
    @test eltype(ts4) == BigFloat
    @test isapprox.(ts[10, :],ts4[10, :]) == trues(3)
  end
  @testset "Jacobians" begin

    J1 = jacobian(s1)
    @test typeof(J1) <: SMatrix
    J2 = jacobian(s2)
    J4 = jacobian(s4)
    @test typeof(J4) <: SMatrix

    @test isapprox.(J1, J2; rtol = 1e-6) == trues(J1)
    @test isapprox.(J1, J4; rtol = 1e-6) == trues(J1)
    @test eltype(J4) == BigFloat
  end
end


using MultivariateStats
using Base.Test

srand(34568)

## PCA with zero mean

X = randn(5, 10)
Y = randn(3, 10)

P = qr(randn(5, 5))[1][:, 1:3]
pvars = [5., 4., 3.]
M = PCA(Float64[], P, pvars, 15.0)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == zeros(5)
@test projection(M) == P
@test principalvars(M) == pvars
@test principalvar(M, 2) == pvars[2]
@test tvar(M) == 15.0
@test tprincipalvar(M) == 12.0
@test tresidualvar(M) == 3.0
@test principalratio(M) == 0.8

@test_approx_eq transform(M, X[:,1]) P'X[:,1]
@test_approx_eq transform(M, X) P'X

@test_approx_eq reconstruct(M, Y[:,1]) P * Y[:,1]
@test_approx_eq reconstruct(M, Y) P * Y

## PCA with non-zero mean

mv = rand(5)
M = PCA(mv, P, pvars, 15.0)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == mv
@test projection(M) == P
@test principalvars(M) == pvars
@test principalvar(M, 2) == pvars[2]
@test tvar(M) == 15.0
@test tprincipalvar(M) == 12.0
@test tresidualvar(M) == 3.0
@test principalratio(M) == 0.8

@test_approx_eq transform(M, X[:,1]) P' * (X[:,1] .- mv)
@test_approx_eq transform(M, X) P' * (X .- mv)

@test_approx_eq reconstruct(M, Y[:,1]) P * Y[:,1] .+ mv
@test_approx_eq reconstruct(M, Y) P * Y .+ mv

## prepare training data

d = 5
n = 1000

R = qr(randn(d, d))[1]
@test_approx_eq R'R eye(5)
scale!(R, sqrt([0.5, 0.3, 0.1, 0.05, 0.05]))

X = R'randn(5, n) .+ randn(5)
mv = vec(mean(X, 2))
Z = X .- mv

C = cov(X; vardim=2)
pvs0 = sort(eigvals(C); rev=true)
tv = sum(diag(C))

## pcacov (default using cov when d < n)

M = fit(PCA, X)
P = projection(M)
pvs = principalvars(M)

@test indim(M) == 5
@test outdim(M) == 5
@test mean(M) == mv
@test_approx_eq P'P eye(5)
@test_approx_eq C * P scale(P, pvs)
@test issorted(pvs; rev=true)
@test_approx_eq pvs pvs0
@test_approx_eq tvar(M) tv
@test_approx_eq sum(pvs) tvar(M)
@test_approx_eq reconstruct(M, transform(M, X)) X

M = fit(PCA, X; mean=mv)
@test_approx_eq projection(M) P

M = fit(PCA, Z; mean=0)
@test_approx_eq projection(M) P

M = fit(PCA, X; maxoutdim=3)
P = projection(M)

@test indim(M) == 5
@test outdim(M) == 3
@test_approx_eq P'P eye(3)
@test issorted(pvs; rev=true)

M = fit(PCA, X; pratio=0.85)

@test indim(M) == 5
@test outdim(M) == 3
@test_approx_eq P'P eye(3)
@test issorted(pvs; rev=true)

## pcastd

M = fit(PCA, X; method=:svd)
P = projection(M)
pvs = principalvars(M)

@test indim(M) == 5
@test outdim(M) == 5
@test mean(M) == mv
@test_approx_eq P'P eye(5)
@test_approx_eq_eps C * P scale(P, pvs) 1.0e-3
@test issorted(pvs; rev=true)
@test_approx_eq_eps pvs pvs0 1.0e-3
@test_approx_eq_eps tvar(M) tv 1.0e-3
@test_approx_eq sum(pvs) tvar(M)
@test_approx_eq reconstruct(M, transform(M, X)) X

M = fit(PCA, X; method=:svd, mean=mv)
@test_approx_eq projection(M) P

M = fit(PCA, Z; method=:svd, mean=0)
@test_approx_eq projection(M) P

M = fit(PCA, X; method=:svd, maxoutdim=3)
P = projection(M)

@test indim(M) == 5
@test outdim(M) == 3
@test_approx_eq P'P eye(3)
@test issorted(pvs; rev=true)

M = fit(PCA, X; method=:svd, pratio=0.85)

@test indim(M) == 5
@test outdim(M) == 3
@test_approx_eq P'P eye(3)
@test issorted(pvs; rev=true)


function alphaj(a::Complex128)
    return sqrt(1 - abs(a)^2)
end

function dictionaryElement(z::Complex128,a::Complex128)
    return alphaj(a)/(z-a)
end

function blaschkeElement(z::Complex128,a::Complex128)
    return (1-conj(a)*z)/(z-a)
end

function invBlaschkeElement(z::Complex128,a::Complex128)
    return (z-a)/(1-conj(a)*z)
end

function blaschkeProd(z::Complex128,aj::Vector{Complex128})
    r = one(z)
    for i=1:length(aj)
        r *= blaschkeElement(z,aj[i])
    end
    return r
end

function blaschkeProdExceptQ(z::Complex128,aj::Vector{Complex128},q::Int)
    r = one(z)
    for i=1:length(aj)
        if i!=q
            r *= blaschkeElement(z,aj[i])
        end
    end
    return r
end

function impulseTakenakaMalmquist(n,aj::Vector{Complex128})
    if n<=zero(n)
        return zero(Complex128)
    end
    temp1 = zero(Complex128)
    temp2 = zero(Complex128)
    for i=1:(length(aj)-1)
        temp1 += aj[i]^(n-1)*dictionaryElement(aj[i],aj[end])*blaschkeProdExceptQ(aj[i],aj[1:(length(aj)-1)],i)*(1-abs(aj[i])^2)
    end
    temp2 += aj[end]^(n-1)*alphaj(aj[end])*blaschkeProd(aj[end],aj[1:(length(aj)-1)])
    return temp1 + temp2
end


# coding: utf-8

# In[1]:

#global list is the key
#otherwise we will lose the list throughout iterations
global factors=[]

# In[2]:

function factorization(num)

    #negative and float should be excluded
    if (num<0) || !(isinteger(num))

        error("negative or float is not allowed")

    end

    #if n is smaller than 4
    #prime number it is
    if num>4

        #exclude 1 and itself
        #the largest factor of n can not exceed the half of n
        #because 2 is the smallest factor
        #the range of factors we are gonna try starts from 2 to fld(num,2)+1
        #int function to solve the odd number problem
        for i in 2:(fld(num,2)+1)

            #the factorization process
            if num%i==0

                push!(factors,i)

                #return is crucial
                #if the number is not a prime number
                #it will stop function from appending non-prime factors
                #the next few lines will be ignored
                return factorization(Int32(num/i))

            end

        end

    end

    #append the last factor
    #it could be n itself if n is a prime number
    #in that case there is only one element in the list
    #or it could be the last indivisible factor of n which is also a prime number
    push!(factors,num)

    if length(factors)==1

        println(num," is a prime number")
        empty!(factors)

    end

end

# In[3]:

factorization(71392643)

# In[4]:

print(factors)


module AffineFEOperatorsTests

using Test
using Gridap.Arrays
using Gridap.Algebra
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.Fields
using Gridap.FESpaces
using LinearAlgebra
using Gridap.CellData

domain =(0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)
order = 1
reffe = ReferenceFE(lagrangian,Float64,order)
V = FESpace(model,reffe;dirichlet_tags=[1,10])
U = V

trian = get_triangulation(model)
degree = 4
quad = CellQuadrature(trian,degree)

#f(x) = x[2]
#v = get_cell_shapefuns(V)
#u = get_cell_shapefuns_trial(U)
#
#cellmat = integrate(∇(v)⊙∇(u),quad)
#cellvec = integrate(v*f,quad)
#cellids = collect(1:num_cells(trian))
#
#assem = SparseMatrixAssembler(U,V)
#matdata = ([cellmat],[cellids],[cellids])
#vecdata = ([cellvec],[cellids])
#A =  assemble_matrix(assem,matdata)
#b =  assemble_vector(assem,vecdata)
#
#op = AffineFEOperator(U,V,A,b)
#@test A === get_matrix(op)
#@test b === get_vector(op)
#
#x = ones(length(b))
#r = A*x - b
#
#test_fe_operator(op,x,r,≈,jac=A)
#@test isa(get_algebraic_operator(op), AffineOperator)

#

tol = 1.0e-9

u_sol(x) = x[1]+x[2]
f_fun(x) = 0

V = FESpace(model,reffe;dirichlet_tags="boundary")
U = TrialFESpace(V,u_sol)

dΩ = Measure(quad)

a(u,v) = ∫(∇(v)⊙∇(u))*dΩ
l(v) = ∫(v*f_fun)*dΩ

assem = SparseMatrixAssembler(U,V)

op = AffineFEOperator(U,V,assem) do u,v
  ∫(∇(v)⊙∇(u))*dΩ, ∫(v*f_fun)*dΩ
end
uh = solve(op)
e = u_sol - uh

@test sum(∫(e*e)*dΩ) < tol
@test ∑(∫(e*e)*dΩ) < tol

op = AffineFEOperator(U,V) do u,v
  ∫(∇(v)⊙∇(u))*dΩ, ∫(v*f_fun)*dΩ
end
uh = solve(op)
e = u_sol - uh
@test sum(∫(e*e)*dΩ) < tol

op = AffineFEOperator(a,l,U,V)
uh = solve(op)
e = u_sol - uh
@test sum(∫(e*e)*dΩ) < tol

end # module


using Statistics
using Random

################################################################

obj(Y, X, b) = mean((Y .- X * b) .^ 2.0) # objective to minimize

# modified example from GradDescent.jl documentation
@testset "linreg" begin

    Random.seed!(1) # set seed
    n = 1000 # number of observations
    d = 10   # number of covariates
    X = randn(n, d) # simulated covariates
    b = randn(d)    # generated coefficients
    ϵ = randn(n) * 0.1 # noise
    Y = X * b + ϵ # observed outcome

    epochs = 100 # number of epochs

    θ = randn(d) # initialize model parameters

    loss_first, g = grad(obj, Y, X, θ)
    loss_last = loss_first
    for i in 1:epochs
        loss_last, g = grad(obj, Y, X, θ)
        # println("Epoch: $i; loss = $loss_last")
        δ = 0.01 * g[4]
        θ -= δ
    end
    @test loss_last < loss_first / 10   # loss reduced at least 10 times

end

mutable struct Linear2{T}
    W::AbstractMatrix{T}
    b::AbstractVector{T}
end

linear_loss(m::Linear2, X, y) = mean((m.W * X .+ m.b .- y) .^ 2.0)

@testset "linreg 2" begin

    n_vars = 10
    n_out = 2
    n_obs = 1000
    X = randn(n_vars, n_obs)
    W_true = rand(n_out, n_vars)
    b_true = rand(n_out)
    ϵ = randn(n_out, n_obs) * 0.1 # noise
    Y = W_true * X .+ b_true .+ ϵ

    m = Linear2(rand(n_out, n_vars), rand(n_out))

    epochs = 100

    for i in 1:epochs
        val, g = grad(linear_loss, m, X, Y)
        # println("Epoch: $i; loss = $val")
        update!(m, g[2], (x, gx) -> x - 0.1 * gx)
    end

    @test isapprox(m.W, W_true; atol=0.1)
    @test isapprox(m.b, b_true; atol=0.1)

end


export elu

doc"""
    elu(x::Var)

Exponential Linear Unit.

# References
* Clevert et al., ["Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)"](https://arxiv.org/abs/1511.07289), arXiv 2015.

```math
f(x) =
\begin{cases}
x & x > 0 \\
\alpha (e^{x}-1) & x\leq0
\end{cases}
```
where ``\alpha=1``.
"""
elu(x::Var) = Var(elu.(x.data), (elu,x))
elu(x::T) where T = x > zero(T) ? x : exp(x)-1
elu(x::Node) = Node(elu, x)

function addgrad!(y::Var, ::typeof(elu), x::Var)
    isvoid(x.grad) && return
    ∇elu!(y.data, y.grad, x.data, x.grad)
end

function ∇elu!(y::Array{T}, gy::Array{T}, x::Array{T}, gx::Array{T}) where T
    alpha = one(T)
    @inbounds for i = 1:length(x)
        gx[i] += x[i] > zero(T) ? gy[i] : gy[i]*(y[i]+alpha)
    end
end


using EngEconomics

# Given
# 12 months
# month1 - 15000 cost
# g = 5% per month
# grant is received in 6 monthy installments starting at month 2
# i = 12%

n = 1
m = 12
i = 0.12
nc = n * m
ic = i / m
Ac = 15000
g = 0.05
F = Ac * presentWorthConversionFactorGeometric(g, ic, nc) * compoundAmountFactor(ic, nc)
Ag = F * capitalRecoveryFactor(ic, 6)


###########################################################################################################
###########################################################################################################
###########################################################################################################
###########################################################################################################
###########################################################################################################
###########################################################################################################
#                                         RADIAL COEFFICIENTS                                             #
###########################################################################################################
###########################################################################################################
###########################################################################################################
###########################################################################################################
###########################################################################################################
###########################################################################################################
export SlaterON, SlaterSF, SlaterSA, SlaterSB, SlaterSN, SlaterSX
###########################################################################################################
########################################## SLATER TYPE ORBITALS ###########################################
SlaterON(n::arb, ζ::arb) = ((2*ζ)^(n + 1//2))//Sqrt(Gamma(2*n + One(n)))
SlaterON(n::Real, ζ::Real) = SlaterON(RF(n), RF(ζ))

function SlaterON(n1::arb, n2::arb, ρ::arb, τ::arb)
    res1 = ((ρ*(One(τ)+τ))^(n1 + 1//2)) * ((ρ*(One(τ)-τ))^(n2 + 1//2))
    res2 = Sqrt(Gamma(2*n1 + One(n1)) * Gamma(2*n2 + One(n2)))
    res1 // res2
end
SlaterON(n1::Real, n2::Real, ρ::Real, τ::Real) = SlaterON(RF(n1), RF(n2), RF(ρ), RF(τ))
###########################################################################################################
###########################################################################################################
###########################################################################################################
#                                     SLATER TYPE SPINOR ORBITALS                                         #
###########################################################################################################
###########################################################################################################
###########################################################################################################
########################################## NORMALIZATION FACTOR ###########################################
################# Relativistic normalization factor N_{n\kappa}
SlaterSF(n::arb, κ::Int) = ifelse(κ < 0, Abs(κ), Sqrt(κ^2 + (2*n + 1)))
SlaterSF(n::Real, κ::Int) = SlaterSF(RF(n), κ)
###########################################################################################################
############################################## COEFFICIENT A ##############################################
################# Normalization coeeficients so called A^{\beta}_{n\kappa}
SlaterSA(β::Int, n::arb, κ::Int) = ifelse(κ < 0, (β*κ)//(2*n), (β//2)*(SlaterSF(n, κ)-κ) - (RF(1)//RF(2)))
SlaterSA(β::Int, n::Real, κ::Int) = SlaterSA(β, RF(n), RF(κ))
###########################################################################################################
############################################## COEFFICIENT B ##############################################
################# Normalization coeeficients so called B^{\beta}_{n\kappa}
SlaterSB(β::Int, n::arb, κ::Int) = ifelse(κ < 0, Zero(RF), -β*((SlaterSF(n, κ) - κ)//(2*n + 1)))
SlaterSB(β::Int, n::Real, κ::Int) = SlaterSB(β, RF(n), RF(κ))
###########################################################################################################
############################################## COEFFICIENT N ##############################################
################# 1/Sqrt(N) part, the dominator is used for both wave and charge density
function SlaterSN(β::Int, n::arb, κ::Int)
    res1 = SlaterSA(β, n, κ) * SlaterSA(β, n, κ) * Gamma(2*n + RF(1))
    res2 = SlaterSA(β, n, κ) * SlaterSB(β, n, κ) * Gamma(2*n + RF(2))
    res3 = SlaterSB(β, n, κ) * SlaterSB(β, n, κ) * (Gamma(2*n + RF(3)) // 4)
    res1 + res2 + res3
end
SlaterSN(β::Int, n::Real, κ::Int) = SlaterSN(β, RF(n), κ)

SlaterSN(β::Int, n::arb, κ::Int, ζ::arb) = ((2*ζ)^(n + 1//2))//Sqrt(SlaterSN(β, n, κ))
SlaterSN(β::Int, n::Real, κ::Int, ζ::Real) = SlaterSN(β, RF(n), κ, RF(ζ))

function SlaterON(β1::Int, n1::arb, κ1::Int, β2::Int, n2::arb, κ2::Int, ρ::arb, τ::arb)
    res1 = ((ρ*(One(τ)+τ))^(n1 + 1//2)) * ((ρ*(One(τ)-τ))^(n2 + 1//2))
    res2 = Sqrt(SlaterSN(β1, n1, κ1) * SlaterSN(β2, n2, κ2))
    res1 // res2
end
SlaterSN(β1::Int, n1::Real, κ1::Int, β2::Int, n2::Real, κ2::Int, ρ::Real, τ::Real) =
SlaterSN(β1, RF(n1), κ1, β2, RF(n2), κ2, RF(ρ), RF(τ))

function SlaterSX(β1::Int, n1::arb, κ1::Int, ζ1::arb, β2::Int, n2::arb, κ2::Int, ζ2::arb)
    a1 = SlaterSA(β1, n1, κ1)
    a2 = SlaterSA(β2, n2, κ2)
    b1 = SlaterSB(β1, n1, κ1)
    b2 = SlaterSB(β2, n2, κ2)

    x1 = ζ1 * ζ1 * a1 * a2
    x2 = ζ1 * ζ2 * b1 * a2
    x3 = ζ1 * ζ2 * a1 * b2
    x4 = ζ2 * ζ2 * b1 * b2

    x = [x1, x2, x3, x4]
    x = transpose(x)
    return x
end
SlaterSX(β1::Int, n1::Real, κ1::Int, ζ1::Real, β2::Int, n2::Real, κ2::Int, ζ2::Real) =
SlaterSX(β1, RF(n1), κ1, RF(ζ1), β2, RF(n2), κ2, RF(ζ2))


function primitiven{T<:Integer}(m::T)
    1 < m || return T[]
    m != 2 || return T[1]
    !isprime(m) || return T[2:2:m-1]
    rp = trues(m-1)
    if isodd(m)
        rp[1:2:m-1] = false
    end
    for p in keys(factor(m))
        rp[p:p:m-1] = false
    end
    T[1:m-1][rp]
end

function pythagoreantripcount{T<:Integer}(plim::T)
    primcnt = 0
    fullcnt = 0
    11 < plim || return (primcnt, fullcnt)
    for m in 2:plim
        p = 2m^2
        p+2m <= plim || break
        for n in primitiven(m)
            q = p + 2m*n
            q <= plim || break
            primcnt += 1
            fullcnt += div(plim, q)
        end
    end
    return (primcnt, fullcnt)
end

println("Counting Pythagorian Triplets within perimeter limits:")
println("    Limit          All   Primitive")
for om in 1:10
    (pcnt, fcnt) = pythagoreantripcount(10^om)
    println(@sprintf "    10^%02d  %11d   %9d" om fcnt pcnt)
end


function lorentz_likelihood_hermitian_band_save(p,y::Vector{Float64},aex,al_small,indx,logdeta)
# Computes the likelihood of a sum of Exponential/Cosine kernels (which have a
# Lorentzian power spectrum) utilizing the approach of Ambikasaran (2015)
#   Numer. Linear Algebra Appl. 2015; 22:1102-1114 DOI: 10.1002/nla
# Uses (2.30) from Rasmussen & Williams (2006) to compute the log likelihood.
# The form of the kernel is:
#
#   K(t_i,t_j) = w_i \delta_{ij} + \sum_{k=1}^{p} \alpha_k \exp{-\beta_k |t_i-t_j|}
#
# This is represented by the matrix A_{ij} = K(t_i,t_j).
# The log likelihood is given by:
# ln(L) = -1/2 y^T A^{-1} y - 1/2 |det(A)|
# Note: if a \beta_k is complex, then its complex conjugate *must* be included as well.
#
# First, define band-diagonal matrix for kernel:
# [x_i,{r_{i,1},...,r_{i,p}},{l_{i,1},...,l_{i,p}}]

#tic()
n = length(y)
# There are p+1 sub-diagonals, p+1 super-diagonals + diagonal
# for a total of 2*p+3 non-zero diagonals:
nex = (2p+1)*n-2p
bex = zeros(Complex{Float64},nex)
for i=1:n
# Compute actual indices:
  irow =(i-1)*(1+2p)+1
  bex[irow] = y[i]
# Diagonal noise:
  for j=1:p
    if i < n
      bex[irow+  j] = 0.0
      bex[irow+p+j] = 0.0
    end
  end
end

# Specify the number of bands below & above the diagonal:
m1 = p+1
m2 = p+1
# Solve the equation A^{-1} y = b using band-diagonal LU back-substitution on
# the extended equations: A_{ex}^{-1} y_{ex} = b_{ex}:
banbks(aex,nex,m1,m2,al_small,indx,bex)
# Now select solution to compute the log likelihood:
# The equation A^{-1} y = b has been solved in bex (the extended vector).
# So, I need to pick out the b portion from bex, and take the dot product
# with y (which is the residuals of the data minus model, which is correlated
# noise that we are modeling with the multi-Lorentzian covariance function):
log_like = 0.0
for i=1:n
  log_like += real(bex[(i-1)*(2p+1)+1])*y[i]
end
# Convert this to log likelihood:
log_like = -0.5*log_like
return log_like
end


module PhysicalOceanography

export f, β

Ω = 7.2921 * 10^(-5) #rad/s
R = 6.3781 * 10^(6) #m

"""
    f(φ)

Coriolis Parameter
"""
f(φ) = 2Ω * sind(φ)

"""
    β(φ)

Rossby Parameter
"""
β(φ) = 2Ω / R * cosd(φ)

end