function makePlots()

println("Generating plots")

plotsDir = (pwd()[end-3:end] == "docs") ? "build/plots" : "docs/build/plots"
mkdir(plotsDir)
Plots.pyplot()

# PID design functions
P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
Plots.savefig(plotsDir*"/pidgofplot.svg")

ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=false)
gangoffourplot(P, [tf(1), C])
Plots.savefig(plotsDir*"/pidgofplot2.svg")
nyquistplot([P, P*C])
Plots.savefig(plotsDir*"/pidnyquistplot.svg")

ωp = 2
kp,ki,C60 = loopshapingPI(P,ωp,rl=1,phasemargin=60, doplot=true)
gangoffourplot(P, [tf(1), C60])
Plots.savefig(plotsDir*"/pidgofplot3.svg")
nyquistplot([P, P*C60])
Plots.savefig(plotsDir*"/pidnyquistplot2.svg")

# Advanced pole placement
ζ = 0.2
ω = 1
B = [1]
A   = [1, 2ζ*ω, ω^2]
P  = tf(B,A)
# Control design
ζ0 = 0.7
ω0 = 2
Am = [1, 2ζ0*ω0, ω0^2]
Ao = conv(2Am, [1/2, 1]) # Observer polynomial
AR = [1,0] # Force the controller to contain an integrator

B⁺  = [1] # The process numerator polynomial can be facored as B = B⁺B⁻ where B⁻ contains the zeros we do not want to cancel (non-minimum phase and poorly damped zeros)
B⁻  = [1]
Bm  = conv(B⁺, B⁻) # In this case, keep the entire numerator polynomial of the process

R,S,T = rstc(B⁺,B⁻,A,Bm,Am,Ao,AR) # Calculate the 2-DOF controller polynomials

Gcl = tf(conv(B,T),zpconv(A,R,B,S)) # Form the closed loop polynomial from reference to output

stepplot([P,Gcl]) # Visualize the open and closed loop responses.
Plots.savefig(plotsDir*"/ppstepplot.svg")
gangoffourplot(P, tf(-S,R)) # Plot the gang of four to check that all tranfer functions are OK
Plots.savefig(plotsDir*"/ppgofplot.svg")

P1 = "exp(-sqrt(s))"
f1 = stabregionPID(P1,logspace(-5,1,1000)); Plots.savefig(plotsDir*"/stab1.svg")
P2 = "100*(s+6).^2./(s.*(s+1).^2.*(s+50).^2)"
f2 = stabregionPID(P2,logspace(-5,2,1000)); Plots.savefig(plotsDir*"/stab2.svg")
P3 = tf(1,[1,1])^4
f3 = stabregionPID(P3,logspace(-5,0,1000)); Plots.savefig(plotsDir*"/stab3.svg")

# PID plots

P = tf([1.],[1., 1])
ζ = 0.5 # Desired damping
ws = logspace(-1,2,8) # A vector of closed-loop bandwidths
kp = 2*ζ*ws-1 # Simple pole placement with PI given the closed-loop bandwidth, the poles are placed in a butterworth pattern
ki = ws.^2
pidplots(P,:nyquist,;kps=kp,kis=ki, ω= logspace(-2,2,500))
Plots.savefig(plotsDir*"/pidplotsnyquist1.svg")
pidplots(P,:gof,;kps=kp,kis=ki, ω= logspace(-2,2,500))
Plots.savefig(plotsDir*"/pidplotgof1.svg")

kp = linspace(-1,1,8) # Now try a different strategy, where we have specified a gain crossover frequency of 0.1 rad/s
ki = sqrt(1-kp.^2)/10
pidplots(P,:nyquist,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsnyquist2.svg")
pidplots(P,:gof,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsgof2.svg")

#AR = [1,0]
#B⁺  = [1]
#B⁻  = [1]
#Bm  = conv(B⁺, B⁻)
#R,S,T = rstc(B⁺,B⁻,A,bm,Am,Ao,AR)
#Gcl = tf(conv(B,T),zpconv(A,R,B,S))
#f1 = stepplot([P,Gcl])
#f2 = gangoffourplot(P, tf(-S,R))
#Plots.savefig(f1, "$plotsDir/rstcstepplot.svg")
#Plots.savefig(f2, "$plotsDir/rstcgofplot.svg")
end


"""
    expectation_step(X, z, dists)

Computes the expectation step with:
- `X`: (m×n) multi-variate observations (rows are observations, columns are variables)
- `dists`: distributions, one for each category (unique label)
- `prior`: prior probability of each label
Returns `Z`: (m×n) matrix of probability distributions of each label for each observation
"""
function expectation_step(X, dists, prior)
    n = size(X, 1)
    Z = Matrix{Float64}(undef, n, length(prior))
    for k in eachindex(prior)
        for i in 1:n
            Z[i,k] = prior[k] * pdf(dists[k], X[i,:]) /
                     sum(
                        prior[j] * pdf(dists[j], X[i,:])
                        for j in eachindex(prior)
                     )
        end
    end
    return Z
end

"""
    maximization_step(::Type{<:D}, X, Z)

Maximization step, parameterized by the type of distribution used.
Returns a tuple (dists, π, loglike) with:
- `dists`: the created vector of distributions
- `prior` vector of updated prior probabilities of each label
"""
function maximization_step(::Type{<:MvNormal}, X, Z)
    n = size(X, 1)
    Nk = size(Z, 2)
    µ = map(1:Nk) do k
        num = sum(Z[i,k] .* X[i,:] for i in 1:n)
        den = sum(Z[i,k] for i in 1:n)
        num / den
    end

    Σ = map(1:Nk) do k
        num = zeros(size(X, 2), size(X, 2))
        for i in 1:n
            r = X[i,:] .- µ[k]
            num .= num .+ Z[i,k] .* (r * r')
        end
        den = sum(Z[i,k] for i in 1:n)
        num ./ den
    end
    prior = [inv(n) * sum(Z[i,k] for i in 1:n)
            for k in 1:Nk
        ]
    dists = map(1:Nk) do k
        MvNormal(µ[k], Σ[k] + 10e-7I)
    end
    return (dists, prior)
end

"""
    expectation_maximization(D::Type{<:Distribution}, X, Nk, maxiter::Integer = 500, loglike_diff = 10e-5; all_dists = nothing)

Complete EM algorithm, takes a type of distribution `D` used to
determine the maximization step, observations `X` and number of classes `Nk`.
"""
function expectation_maximization(D::Type{<:Distribution}, X, Nk, maxiter::Integer = 500, loglike_diff = 10e-5; all_dists = nothing)
    # initialize classes
    n = size(X,1)
    Z = zeros(n, Nk)
    for i in 1:n
        j0 = mod(i,Nk)+1
        j1 = j0 > 1 ? j0-1 : 2
        Z[i,j0] = 0.75
        Z[i,j1] = 0.25
    end
    (dists, prior) = maximization_step(D, X, Z)
    all_dists isa Vector && push!(all_dists, (dists, prior))
    l = loglike_mixture(X, dists, prior)
    lprev = 0.0
    iter = 0
    # EM iterations
    while iter < maxiter && abs(lprev-l) > loglike_diff
        Z = expectation_step(X, dists, prior)
        (dists, prior) = maximization_step(D, X, Z)
        all_dists isa Vector && push!(all_dists, (dists, prior))
        lprev = l
        l = loglike_mixture(X, dists, prior)
        iter += 1
    end
    return (dists, prior, Z, l, iter)
end

function loglike_mixture(X, dists, prior)
    l = zero(eltype(X))
    n = size(X,1)
    for i in 1:n
        l += log(
            sum(prior[k] * pdf(dists[k], X[i,:]) for k in eachindex(prior))
        )
    end
    return l
end

let
    Random.seed!(42)
    X = [randn(1000, 2); 3.0 .+ randn(1000, 2)]
    n = size(X, 1)
    for maxiter in (1, 5, 10, 100, 150, 250)
        (dists, prior, Z, l, iter) = expectation_maximization(MvNormal, X, 2, maxiter)
        contour_mixture(x1, x2) = maximum([pdf(dists[k], [x1,x2]) for k in 1:2])
        colors = map(1:n) do i
            (_, j) = findmax(Z[i,:])
            j == 1 ? :red : :blue
        end
        plt = Plots.contour(-2.8:0.05:6.0, -2.8:0.05:6.0, contour_mixture, contours = false, legend = nothing)
        Plots.scatter!(plt, X[:,1],X[:,2], label = "", color = colors)
        Plots.savefig(joinpath(base_img, "mixture_iter$(maxiter)_$iter.pdf"))
    end
end


using DiscreteFiltering
using OrdinaryDiffEq: ODEProblem, solve, Tsit5
using LinearAlgebra
using Plots
using SparseArrays

## Domain
a = 0.0
b = 2π
domain = PeriodicIntervalDomain(a, b)
# domain = ClosedIntervalDomain(a, b)

## Discretization
N = 100
M = N
x = discretize(domain, M)
ξ = discretize(domain, N)
Δx = (b - a) / N

## Filter
h₀ = 2.1Δx
h(x) = h₀ * (1 - 1 / 2 * cos(x))
dh(x) = h₀ / 2 * sin(x)
α(x) = 1 / 3 * dh(x) * h(x)
f = TopHatFilter(h)

# h₀ = 5.1Δx
# h(x) = h₀ # * (1 - 1 / 2 * cos(x))
# dh(x) = 0.0 # h₀ / 2 * sin(x)
# α(x) = 1 / 3 * dh(x) * h(x)
# f = GaussianFilter(h, Δx / 2)

## Time
T = 1.5
t = T

## Plot filter
plot(x, h)
ylims!((0, ylims()[2]))

## Get matrices
C = advection_matrix(domain, N)
D = diffusion_matrix(domain, N)
W = filter_matrix(f, domain, M, N)
R = reconstruction_matrix(f, domain, M, N)
A = spdiagm(α.(ξ))

## Inspect matrices
spy(W)
spy(R)

## Exact solutions
u(ξ, t) = sin(ξ - t) + 3 / 5 * cos(5(ξ - t)) + 1 / 25 * sin(20(ξ - 1 - t))
u_int(ξ, t) = -cos(ξ - t) + 3 / 25 * sin(5(ξ - t)) - 1 / 25 / 20 * cos(20(ξ - 1 - t))
ū(x, t) = 1 / 2h(x) * (u_int(x + h(x), t) - u_int(x - h(x), t))

## Discrete initial conditions
uₕ = u.(ξ, 0.0)
ūₕ = ū.(x, 0.0)
uₕ_allbar = W * uₕ

plot(x, uₕ, label = "Discretized")
plot!(x, ūₕ, label = "Filtered-then-discretized")
plot!(x, uₕ_allbar, label = "Discretized-then-filtered")

## Solve discretized problem
duₕ(uₕ, p, t) = -C * uₕ
prob = ODEProblem(duₕ, uₕ, (0, T))
sol = solve(prob, Tsit5(), abstol = 1e-6, reltol = 1e-4)

plot(x, uₕ, label = "Initial conditions")
plot!(x, sol(t), label = "Discretized")
plot!(x, u.(x, t), label = "Exact")

## Solve filtered-and-then-discretized problem
dūₕ(ūₕ, p, t) = (-C + A * D) * ūₕ
prob_bar = ODEProblem(dūₕ, ūₕ, (0, T))
sol_bar = solve(prob_bar, Tsit5(), abstol = 1e-6, reltol = 1e-4)

plot(x, uₕ, label = "Initial")
plot!(x, ūₕ, label = "Initial filtered")
plot!(x, sol_bar(t), label = "Filtered-then-discretized")
plot!(x, ū.(x, t), label = "Filtered exact")
# plot!(x, 500*α.(ξ))

## Solve discretized-and-then-filtered problem
duₕ_allbar(uₕ_allbar, p, t) = -W * (C * (R * uₕ_allbar))
# duₕ_allbar(uₕ_allbar, p, t) = -W * (C * (W \ uₕ_allbar))
prob_allbar = ODEProblem(duₕ_allbar, W * uₕ, (0, T))
sol_allbar = solve(prob_allbar, Tsit5(), abstol = 1e-6, reltol = 1e-4)

plot(x, uₕ, label = "Initial")
plot!(x, uₕ_allbar, label = "Initial discretized-then-filtered")
plot!(x, sol_allbar(t), label = "Discretized-then-filtered")
plot!(x, W * u.(ξ, t), label = "Exact")

## PGFPlotsX
pgfplotsx()

## Comparison
p = plot(
    size = (400, 300),
    xlabel = "\$x\$",
    legend = :topright,
)
plot!(p, x, uₕ, label = "\$u(x, t = 0.0)\$")
# plot!(p, x, ūₕ, label = "Initial filtered")
# plot!(p, x, sol(t), label = "Discretized")
# plot!(p, x, sol_bar(t), label = "Filtered-then-discretized")
# plot!(p, x, sol_allbar(t), label = "Discretized-then-filtered")
# plot!(p, x, [u.(x, t), ū.(x, t)], label = "Exact")
plot!(p, x, u.(x, t), label = "\$u(x, t = $t)\$")
plot!(p, x, ū.(x, t), label = "\$u(x, t = $t)\$")
# ylims!(p, minimum(uₕ), maximum(uₕ))
display(p)
savefig(p, "output/advection/solution.tikz")

## Relative error
u_exact = u.(ξ, t)
ū_exact = ū.(x, t)
err = abs.(sol(t) - u_exact) ./ maximum(abs.(u_exact))
err_bar = abs.(sol_bar(t) - u_exact) ./ maximum(abs.(u_exact))
err_allbar = abs.(sol_allbar(t) - u_exact) ./ maximum(abs.(u_exact))

plot()
plot!(ξ, err, label = "Unfiltered discretized")
plot!(x, err_bar, label = "Filtered-then-discretized")
plot!(x, err_allbar, label = "Discretized-then-filtered")


# # Variational problems
#
# In this example, we will numerically simulate an entropy-regularised Wasserstein gradient flow
# approximating the Fokker-Planck and porous medium equations.
#
# The connection between Wasserstein gradient flows and (non)-linear PDEs is due to Jordan, Kinderlehrer and Otto [^JKO98], and
# an easy-to-read overview of the topic is provided in Section 9.3 [^PC19]
#
# [^JKO98]: Jordan, Richard, David Kinderlehrer, and Felix Otto. "The variational formulation of the Fokker--Planck equation." SIAM journal on mathematical analysis 29.1 (1998): 1-17.
# [^PC19]: Peyré, Gabriel, and Marco Cuturi. "Computational optimal transport: With applications to data science." Foundations and Trends® in Machine Learning 11.5-6 (2019): 355-607.
#
# ## Fokker-Planck equation as a $W_2$ gradient flow
# For a potential function $\Psi$ and noise level $\sigma^2$, the Fokker-Planck equation (FPE) is
# ```math
# \partial_t \rho_t = \nabla \cdot (\rho_t \nabla \Psi) + \frac{\sigma^2}{2} \Delta \rho_t,
# ```
# and we take no-flux (Neumann) boundary conditions.
#
# This describes the evolution of a massless particle undergoing both diffusion (with diffusivity $\sigma^2$) and drift (along potential $\Psi$) according to the stochastic differential equation
# ```math
# dX_t = -\nabla \Psi(X_t) dt + \sigma dB_t.
# ```
# The result of Jordan, Kinderlehrer and Otto (commonly referred to as the JKO theorem) states that
# $\rho_t$ evolves following the 2-Wasserstein gradient flow of the Gibbs free energy functional
# ```math
#   F(\rho) = \int \Psi d\rho + \int \log(\rho) d\rho.
# ```
#
# ## Implicit schemes for gradient flows
# In an Euclidean space, the gradient flow of a functional $F$ is simply the solution of an ordinary differential equation
# ```math
#  \dfrac{dx(t)}{dt} = -\nabla F(x(t)).
# ```
# Of course, there is a requirement that $F$ is smooth. A more general formulation of a gradient flow that allows
# $F$ to be non-smooth is the implicit scheme
# ```math
#   x_{t+\tau} = \operatorname{argmin}_x \frac{1}{2} \| x - x_t \|_2^2 + \tau F(x).
# ```
# As the timestep $\tau$ shrinks, $x_t$ becomes a better and better approximation to the gradient flow of $F$.
#
# ## Wasserstein gradient flow
# In the context of the JKO theorem, we seek $\rho_t$ that is the gradient flow of $F$ with
# respect to the 2-Wasserstein distance. This can be achieved by choosing the $W_2$ metric in the implicit step:
# ```math
#   \rho_{t + \tau} = \operatorname{argmin}_{\rho} d_{W_2}^2(\rho_{t}, \rho) + \tau F(\rho).
# ```
# Finally, a numerical scheme for computing this gradient flow can be developed by using the entropic regularisation
# of optimal transport on a discretised domain
# ```math
#   \rho_{t + \tau} = \operatorname{argmin}_{\rho} \operatorname{OT}_\varepsilon(\rho_{t}, \rho) + \tau F(\rho),
# ```
# where
# ```math
#   \operatorname{OT}_\varepsilon(\alpha, \beta) = \min_{\gamma \in \Pi(\alpha, \beta)} \sum_{i,j} \frac{1}{2} \| x_i - x_j \|_2^2 \gamma_{ij} + \varepsilon \sum_{i, j} \gamma_{ij} \log(\gamma_{ij}).
# ```
# Each step of this problem is a minimisation problem with respect to $\rho$.
# Since we use entropic optimal transport which is differentiable, this can be solved using gradient-based methods.

# ## Problem setup
#
using OptimalTransport
using Distances
using LogExpFunctions
using Optim
using Plots
using StatsBase

using LinearAlgebra
using Logging

# Here, we set up the computational domain that we work on - we discretize the interval $[-1, 1]$.
# The natural boundary conditions to use will be Neumann (zero flux), see e.g. [^Santam2017]
#
# [^Santam2017]: Santambrogio, Filippo. "{Euclidean, metric, and Wasserstein} gradient flows: an overview." Bulletin of Mathematical Sciences 7.1 (2017): 87-154.

support = range(-1, 1; length=64)
C = pairwise(SqEuclidean(), support');

# Now we set up various functionals that we will use.
#
# We define the generalised entropy (Equation (4.4) of [^Peyre2015]) as follows. For $m = 1$ this is just the "regular" entropy, and $m = 2$ this is squared $L_2$.
#
# [^Peyre2015]: Peyré, Gabriel. "Entropic approximation of Wasserstein gradient flows." SIAM Journal on Imaging Sciences 8.4 (2015): 2323-2351.
function E(ρ; m=1)
    if m == 1
        return sum(xlogx.(ρ)) - sum(ρ)
    elseif m > 1
        return dot(ρ, @. (ρ^(m - 1) - m) / (m - 1))
    end
end;

# Now define $\psi(x)$ to be a potential energy function that has two potential wells at $x = ± 0.5$.
ψ(x) = 10 * (x - 0.5)^2 * (x + 0.5)^2;
plot(support, ψ.(support); color="black", label="Scalar potential")

# Having defined $\psi$, this induces a potential energy functional $\Psi$ on probability distributions $\rho$:
# ```math
#    \Psi(\rho) = \int \psi(x) \rho(x) dx = \langle \psi, \rho \rangle .
# ```
Ψ = ψ.(support);

# Define the time step $\tau$ and entropic regularisation level $\varepsilon$, and form the associated Gibbs kernel $K = e^{-C/\varepsilon}$.
τ = 0.05
ε = 0.01
K = @. exp(-C / ε)

# We define the (non-smooth) initial condition $\rho_0$ in terms of step functions.
H(x) = x > 0
ρ0 = @. H(support + 0.25) - H(support - 0.25)
ρ0 = ρ0 / sum(ρ0)
plot(support, ρ0; label="Initial condition ρ0", color="blue")

# `G_fpe` is the objective function for the implicit step scheme
# ```math
# G_\mathrm{fpe}(\rho) = \operatorname{OT}_\varepsilon(\rho_{t}, \rho) + \tau F(\rho),
# ```
# and we seek to minimise in $\rho$.
function G_fpe(ρ, ρ0, τ, ε, C)
    return sinkhorn2(ρ, ρ0, C, ε; regularization=true, maxiter=250) + τ * (dot(Ψ, ρ) + E(ρ))
end;

# `step` solves the implicit step problem to produce $\rho_{t + \tau}$ from $\rho_t$.
function step(ρ0, τ, ε, C, G)
    ## only print error messages
    opt = with_logger(SimpleLogger(stderr, Logging.Error)) do
        optimize(
            u -> G(softmax(u), ρ0, τ, ε, C),
            ones(size(ρ0)),
            LBFGS(),
            Optim.Options(; iterations=50, g_tol=1e-6);
            autodiff=:forward,
        )
    end
    return softmax(Optim.minimizer(opt))
end
# Now we simulate `N = 10` iterates of the gradient flow and plot the result.
N = 10
ρ = similar(ρ0, size(ρ0, 1), N)
ρ[:, 1] = ρ0
for i in 2:N
    ρ[:, i] = step(ρ[:, i - 1], τ, ε, C, G_fpe)
end
colors = range(colorant"red"; stop=colorant"blue", length=N)
plot(
    support,
    ρ;
    title=raw"$F(\rho) = \langle \psi, \rho \rangle + \langle \rho, \log(\rho) \rangle$",
    palette=colors,
    legend=nothing,
)

# ## Porous medium equation
#
# The porous medium equation (PME) is the nonlinear PDE
# ```math
# \partial_t \rho = \nabla \cdot (\rho \nabla \Psi) + \Delta \rho^m,
# ```
# again with Neumann boundary conditions. The value of $m$ in the PME corresponds to picking $m$ in the generalised entropy functional.
# Now, we will solve the PME with $m = 2$ as a Wasserstein gradient flow.
#
function G_pme(ρ, ρ0, τ, ε, C)
    return sinkhorn2(ρ, ρ0, C, ε; regularization=true, maxiter=250) +
           τ * (dot(Ψ, ρ) + E(ρ; m=2))
end;

# set up as previously
N = 10
ρ = similar(ρ0, size(ρ0, 1), N)
ρ[:, 1] = ρ0
for i in 2:N
    ρ[:, i] = step(ρ[:, i - 1], τ, ε, C, G_pme)
end
plot(
    support,
    ρ;
    title=raw"$F(\rho) = \langle \psi, \rho \rangle + \langle \rho, \rho - 1\rangle$",
    palette=colors,
    legend=nothing,
)


function parametrized_kuramoto(du,u,p,t)
    # NumberOfOscillators, CouplingConstant, and NaturalFreqs
    N,K,w = p
    # Centroid
    u[1] = 0
    for i in 1:N
        u[1] += cos(u[i+1]) + sin(u[i+1])*im
    end
    u[1] = u[1]/N
    r = abs(u[1])
    ψ = angle(u[1])
    # Oscillators
    for i in 1:N
    	du[i+1] = w[i] + K*r*sin(ψ-u[i+1])
    end
end

function getparameters_kuramoto(N,K,γ,distribution)
    distribution == "Uniform" ? nothing : error("$(distribution) is not yet supported")
    j = collect(1:N)
    w = broadcast(j->((2j-N-1)/(N-1))*((pi*γ)/2),j)
    p = (N=N, K=K, w=w)
    return p
end

function getstartpoints_kuramoto(N)
    d = Uniform(-pi,pi)
    u0::Array{ComplexF64} = rand(d,N)
    pushfirst!(u0,mean(u0))
    return u0
end

function kuramoto(N,K,T,fs,γ,distribution)
    tspan = (0.0,T)
    u0 = getstartpoints_kuramoto(N)
    p = getparameters_kuramoto(N,K,γ,distribution)
    prob = ODEProblem(parametrized_kuramoto,u0,tspan,p)
    sol = solve(prob,RK4(),saveat=collect(0:1/fs:T),progress=true)
end


# ---
# title: 96. Unique Binary Search Trees
# id: problem96
# author: zhwang
# date: 2022-02-23
# difficulty: Medium
# categories: Dynamic Programming, Tree
# link: <https://leetcode.com/problems/unique-binary-search-trees/description/>
# hidden: true
# ---
#
# Given _n_ , how many structurally unique **BST 's** (binary search trees) that
# store values 1 ...  _n_?
#
# **Example:**
#
#
#
#     Input: 3
#     Output: 5
#     Explanation: Given _n_ = 3, there are a total of 5 unique BST's:
#
#        1         3     3      2      1
#         \       /     /      / \      \
#          3     2     1      1   3      2
#         /     /       \                 \
#        2     1         2                 3
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 19`
#
#
## @lc code=start
using LeetCode

num_of_trees(n::Int) = binomial(2 * n, n) - binomial(2 * n, n + 1)
## @lc code=end


# 2D diffusion problem

# Packages and inclusions
using ModelingToolkit,DiffEqOperators,LinearAlgebra,Test,OrdinaryDiffEq
using ModelingToolkit: Differential

# Tests
@testset "Test 00: Dt(u(t,x)) ~ Dxx(u(t,x,y)) + Dyy(u(t,x,y))" begin
    @parameters t x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2
    Dt = Differential(t)
    t_min= 0.
    t_max = 2.0
    x_min = 0.
    x_max = 2.
    y_min = 0.
    y_max = 2.

    # 3D PDE
    eq  = Dt(u(t,x,y)) ~ Dxx(u(t,x,y)) + Dyy(u(t,x,y))

    analytic_sol_func(t,x,y) = exp(x+y)*cos(x+y+4t)
    # Initial and boundary conditions
    bcs = [u(t_min,x,y) ~ analytic_sol_func(t_min,x,y),
        u(t,x_min,y) ~ analytic_sol_func(t,x_min,y),
        u(t,x_max,y) ~ analytic_sol_func(t,x_max,y),
        u(t,x,y_min) ~ analytic_sol_func(t,x,y_min),
        u(t,x,y_max) ~ analytic_sol_func(t,x,y_max)]

    # Space and time domains
    domains = [t ∈ IntervalDomain(t_min,t_max),
            x ∈ IntervalDomain(x_min,x_max),
            y ∈ IntervalDomain(y_min,y_max)]
    pdesys = PDESystem([eq],bcs,domains,[t,x,y],[u(t,x,y)])

    # Method of lines discretization
    dx = 0.1; dy = 0.2
    for order in [2,4,6]
        discretization = MOLFiniteDifference([x=>dx,y=>dy],t;centered_order=order)
        prob = ModelingToolkit.discretize(pdesys,discretization)
        sol = solve(prob,Tsit5())

        # Test against exact solution
        # TODO: do this properly when sol[u] with reshape etc works
        @test sol.u[1][1] ≈ analytic_sol_func(sol.t[1],0.1,0.2)
        @test sol.u[1][2] ≈ analytic_sol_func(sol.t[1],0.2,0.2)

    end
end


module TF

import ..d; import ..dxasy

sigmoid{F <: AbstractFloat}(x::F) = one(F)/(one(F) + exp(-x))
dasysigmoid{F <: AbstractFloat}(y::F) = y * (one(F) - y)
#explicitly reinterpret the dasysigmoid.
dsigmoid{F <: AbstractFloat}(x::F) = dasysigmoid(sigmoid(x))

id{F <: AbstractFloat}(x::F) = x

softplus{F <: AbstractFloat}(x::F) = log(one(F) + exp(x))
dasysoftplus{F <: AbstractFloat}(y::F) = (exp(y) - one(F)) / exp(y)

#broadcast functions over arrays.
sigmoid{F <: AbstractFloat}(a::Array{F})      = sigmoid.(a)
dasysigmoid{F <: AbstractFloat}(a::Array{F})  = dasysigmoid.(a)
dsigmoid{F <: AbstractFloat}(a::Array{F})     = dsigmoid.(a)
id{F <: AbstractFloat}(a::Array{F})           = a
softplus{F <: AbstractFloat}(a::Array{F})     = softplus.(a)
dasysoftplus{F <: AbstractFloat}(a::Array{F}) = dasysoftplus.(a)

#non-simple transfer functions.
function softmax{F <: AbstractFloat}(a::Array{F})
  #go over the entire array and add up the values.
  totalexpsum = zero(F)
  for idx = 1:length(a)
    totalexpsum += exp(a[idx])
  end
  exp.(a) / totalexpsum
end
dsoftmax{F <: AbstractFloat}(a::Array{F}) = dasysigmoid.(softmax(a))

dxasy(::typeof(sigmoid)) = dasysigmoid
dxasy(::typeof(id)) = one
dxasy(::typeof(softplus)) = dasysoftplus
dxasy(::typeof(softmax)) = dasysigmoid

d(::typeof(sigmoid)) = dsigmoid
d(::typeof(id)) = one
d(::typeof(softplus)) = sigmoid
d(::typeof(softmax)) = dsoftmax

end


function nlp_cvx_103_010(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    # Test Goals:
    # - linear objective
    # - intersection convex quadratic constraints
    # Variants
    #   010 - one binding constraint (inflection point)
    #   011 - one binding constraint (inflection point)
    #   012 - one binding constraint (non-inflection point)
    #   013 - one binding constraint (non-inflection point)
    #   014 - intersection point

    model = Model(optimizer)

    @variable(model, x)
    @variable(model, y)

    @objective(model, Min, y)
    @NLconstraint(model, x^2 <= y)
    @NLconstraint(model, -x^2 + 1 >= y)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, 0, tol = objective_tol)
    return check_solution([x, y], [0, 0], tol = primal_tol)
end


module supply_allocation

using JuMP
using Cbc
using CSV

function solve_supply_allocation()

    #------
    # DATA
    #------

    n = 4;
    m = 3;
    Q = [30 15 20 35];
    d = [40 10 35];

    c = [
        40	39	40;
        35	36	31;
        33	43	42;
        37	35	38
    ];

    #------
    # MODEL
    #------

    model = Model(with_optimizer(Cbc.Optimizer));

    @variable(model, x[1:n,1:m] >= 0)

    @objective(model, Min, sum(c[i,j] * x[i,j] for i in 1:n, j in 1:m));

    @constraint(model,[j = 1:m],sum(x[i,j] for i in 1:n) == d[j]);

    @constraint(model,[i = 1:n],sum(x[i,j] for j in 1:m) <= Q[i]);

    f = open("supply_allocation.lp", "w")
    print(f, model)
    close(f)

    #-------
    # SOLVE
    #-------

    optimize!(model)

    #------------------------
    # WRITE SOLUTION TO FILE
    #------------------------

    f = open("supply_allocation.csv","w");

    print(f,",");
    print(f,"Munich,Salzburg,Prague");
    println(f,"");

    ports = ["Rotterdam" "Wilhelmshaven" "Bremerhaven" "Hamburg" ];

    for i = 1:n
        print(f,ports[i],",");
        for j = 1:m
            print(f,value(x[i,j]),",");
        end
        println(f,"");
    end

    close(f);

end

solve_supply_allocation();

end


#The key to the non-parametric method describedhere
#is that the sum of squared distances between pointsand
#their centroid is equal to (and can be calculateddirectly from)
#the sum of squared interpoint distancesdivided by the number of points
# https://onlinelibrary.wiley.com/doi/epdf/10.1111/j.1442-9993.2001.01070.pp.x
SSₜ(D,N) = sum(D .^2)/N # given a triangular dist matrix
function SS_W(D,N,n)
    return  SSₜ(D,N)/n
end

function F(D,W,N,n,a)
    SST = SSₜ(D,N)
    SSW = SS_W(W,N,n)
    A = SST - SSW
    F = (A/(a-1))/(W/(N-a))
end

permutest(D,group,N,n,a, n_perm = 1000)
    N =  length(group)
    a = length(unique(group))
    n = N/a
    W_inds = [CartesianIndex(i,j) for i in 1:(N-1) for j in i+1:N if group[i] == group[j]]
    W = D[W_inds]
    notW = D[.!W_inds]
    D = vcat(W,notW)

    f = F(D,W,N,n,a)
    Fs = Vector{Float64}(undef,n_perm)
    inds = 1:length(W)

    for i in 1:n_perm
        shuffle!.(D)
        Fs[i] = Fknown(view(D,inds),N,n,a)
    end
    P = sum(Fs .>= f)/n_perm
    return P
end


function nlp_cvx_106_011(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    model = Model(optimizer)

    @variable(model, -3 <= x <= 3)
    @variable(model, -1 <= y <= 1)

    @objective(model, Min, x + y)
    @NLconstraint(model, sin(-x - 1.0) + x / 2 + 0.5 <= y)
    @NLconstraint(model, cos(x - 0.5) + x / 4 - 0.5 >= y)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, -0.7868226265935826, tol = objective_tol)
    return check_solution(
        [x, y],
        [-0.5955231764562057, -0.1912994501373769],
        tol = primal_tol,
    )
end


using Test
include("Wasserstein.jl")

#	test Wasserstein
println("\n____________ Test Wasserstein.jl")
p1= [.3, 0.7]
s1= [-.7, 1.1]
p2= [0.2, 0.3, 0.2, 0.1, 0.2]
s2= [-.3, .9, 1.7, 3.2, 1.9]
d= distFunction(s1, s2)

@show A= Wasserstein(p1, p2, distFunction(s1, s2), 2.)

@test A.distance ≈ 0.9679876032264049
@test A.π ≈ [0.2 0.1 0.0 0.0 0.0; 0.0 0.2 0.2 0.1 0.2]

println()
@show A= Sinkhorn(p1, p2, distFunction(s1, s2), 2., 4.3)
@test A.distance ≈ 0.967988 atol= 1e-4
@test A.π ≈ [0.2 0.1 0.0 0.0 0.0; 0.0 0.2 0.2 0.1 0.2] atol = 1e-6

println()
@show A= Sinkhorn(p1, p2, distFunction(s1, s2), 2., 1.)
@test A.distance ≈ 0.9923876770989377 atol = .01
@test A.π ≈ [0.2 0.1 0.0 0.0 0.0; 0.0 0.2 0.2 0.1 0.2] atol = 0.02


#=
def sample_actual_demand(price):
    avg_demand = 65 + (-0.8) * price
    theta = 0.1/4
    k = avg_demand / theta
    return np.random.gamma(k*theta, k*theta**2, 1)[0]
=#

function sample_actual_demand(price)
    avg_demand = 65 + (-0.8) * price
    theta = 0.1/4
    k = avg_demand / theta
    return rand(Gamma(k*theta, k*theta^2), 1)
end

#=
def emperical_mean(sampler, n):
    mean = 0
    for i in range(1, n):
        mean = mean + sampler()
    return mean/n
=#

function emperical_mean(sampler, n)
    mean = 0
    for i in 1:n-1
        mean += sampler()[1]
    end
    return mean / n
end

#=
def emperical_demand_curve(min_price, max_price, n):
    prices = np.linspace(min_price, max_price, n)
    sampling = 5000
    demands = map(lambda p: emperical_mean(functools.partial(sample_actual_demand, p), sampling), prices)
    return np.dstack((prices, list(demands)))[0]
=#

function emperical_demand_curve(min_price, max_price, n)
    prices = range(min_price, max_price, length=n)
    sampling = 5000
    demands = []
    for p in prices
        sampler() = sample_actual_demand(p)
        demands = append!(demands, emperical_mean(sampler, sampling))
    end
    return DataFrame(
        :prices => prices,
        :demands => demands
        )
end


"""
Solve the heat equation ∂ₜu = Δu on [0,1] using OrdinaryDiffEq for time stepping.
Boundary conditions are such that an analytic solution is given by:
    sin(2π*x) * exp(-(2*π)^2*t)

This example demonstrates how to solve time dependent PDE's combining `fenics` and `DiffEq`.

"""
module HeatEquationExample

using Test
using OrdinaryDiffEq
using FEniCS

mesh = UnitIntervalMesh(50)
V = FunctionSpace(mesh,"P",1)
u_str = "sin(2*pi*x[0]) * exp(-t*pow(2*pi, 2))"

tspan = (0.0, 0.1)
u = interpolate(Expression(u_str, degree=2, pi=Float64(pi), t=tspan[1]), V)

bc = DirichletBC(V, 0.0, "on_boundary")
dudt = TrialFunction(V)
v = TestFunction(V)
Dudt = FeFunction(V)
F = dot(dudt, v)*dx + dot(grad(u), grad(v))*dx
a = lhs(F)
L = rhs(F)
p = (u=u, Dudt=Dudt, bc=bc, a=a, L=L)

function f!(dudt_vec, u_vec, p, t)
    assign(p.u, u_vec)
    lvsolve(p.a, p.L, p.Dudt, p.bc)
    copy!(dudt_vec, get_array(p.Dudt))
end

u0_vec = get_array(u)
t_final = 0.1
prob = ODEProblem(f!, u0_vec, tspan, p)
u_true = interpolate(Expression(u_str, degree=2, pi=Float64(pi), t=tspan[2]), V)

# some algorithms to try:
# alg = ImplicitEuler(autodiff=false)
# alg = Rosenbrock23(autodiff=false)
# alg = Rodas5(autodiff=false)

alg = KenCarp4(autodiff=false)
sol = OrdinaryDiffEq.solve(prob, alg, reltol=1e-6, abstol=1e-6)

vec_sol = get_array(p.u)
vec_true = get_array(u_true)
@test vec_sol ≈ vec_true rtol=1e-2

end#module


#############################
##  Numerical experiments associated with section 5.3
#############################
## This file contains the numerical experiments used for the highly oscillatory Bessel kernel example,
## which was originally featured by Nick Hale in https://doi.org/10.1093/imanum/dry042
##

using ApproxFun, MultivariateOrthogonalPolynomials, BandedMatrices, BlockBandedMatrices, SpecialFunctions, Plots
using SparseVolterraExamples

####
## First we define the Kernel and g(x) as in section 5.3.
## As mentioned in the paper, it is sensible to use pre-existing implementations to compute the Bessel functions
Kfun(x,y)=besselj(μ,ω*(x-y))
g(x) = besselj(μ+ν,ω*x)+1/(2*x^2)*((ν-1)*(ν-2)*besselj(ν-1,ω*x)+(ν+1)*(ν+2)*besselj(ν+1,ω*x))

#####
## Solver function with step-by-step explanation for given polynomial degree 'n'
## Note the factor 10^(-3) placed as discussed in the paper.
function solveSec54(n,g,Kfun,ν,μ,ω)
    gF = Fun(x->g(x),Jacobi(2,3, 0..1))
    V = triVolterraFullKernelOpP01(Kfun,n+2,true,339)
        V = reflectPabtoPba(n+2)*WLoweringP01P00(n+2)*V[1:n+2,1:n+2]
        V = Conversion(Jacobi(0,0,0..1),Jacobi(2,3,0..1))[1:n+2,1:n+2]*V[1:n+2,1:n+2]
        V = 10^(-3)*Derivative(Jacobi(0,1,0..1),2)[1:n+2,1:n+2]+ω^2*Conversion(Jacobi(0,1,0..1),Jacobi(2,3,0..1))[1:n+2,1:n+2]+ω*V
    coeff = [DirectEvalLHSP10at0(n+2);DirectEvalLHSP10atPrime0(n+2)*Derivative(Jacobi(0,1,0..1),1)[1:n+2,1:n+2];V[1:n,1:n+2]] \ [0.0;0.0;pad(gF.coefficients,n)]
    return Fun(Jacobi(0,1, 0..1), coeff)
end

####
## Compute the solution with same parameters as discussed by Hale and in our paper
## This correponds to computing the solution in Figure 7(a).
ν = 3
μ = 2
ω = 20
u = solveSec54(2000,g,Kfun,ν,μ,ω)
plot(u,color=:black, legend=:false, xlabel = "x", ylabel = "u(x)", legendfontsize=12, tickfontsize=10, thickness_scaling = 1 , grid=:none)

####
## A basic loop to compute errors between the above high order solution and increasingly accurate approximations
function errorvals(max,step,u)
    errorvec = []
    for n=10:step:max
        diff = maximum(abs.([maximum(solveSec54(n,g,Kfun,ν,μ,ω)-u),minimum(solveSec54(n,g,Kfun,ν,μ,ω)-u)]))
        errorvec = push!(errorvec,diff)
    end
    return errorvec
end
####
## Plot successive errors to produce something like Figure 7(b). The exact shape of the plot depends on how dense we make the steps in n.
## Note that for low polynomial orders, since the order is too low to resolve the oscillations, the error varies a lot.
## Only once the polynomial order can resolve the oscillations do we see sensible convergence.
errorvec = errorvals(450,40,u)
plot(Array(10:40:450),errorvec,color=:black, legend=:bottomleft , yscale=:log10 , xlabel = "n", label="sparse method" , ylabel = "error", legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)

examples/Section 5-3 - Bessel kernels with highly oscillatory solutions.jl


function eye(tpe::Type,n::S) where {S <: Integer}

    if tpe <: Number
        return Matrix{tpe}(I,n,n)
    end

end

eye(n::Integer) = eye(Float64,n::Integer)

function tracem(x::Array{T,2}) where {T <: Real}

  # Computes the matrix-trace as defined by Gomme and Klein (2011)

  # We require the number of rows to be greater than the number of columns, so
  # that m is greater than one.

  trans = false

  if size(x,1) < size(x,2)
    x = Matrix(x')
    trans = true
  end

  n = size(x,2)
  m = Int(size(x,1)/n)

  y = zeros(m,1)  # We want this to be a 2-d array for subsequent matrix multiplication

  @inbounds for i = 1:m
    @views y[i,1] = tr(x[(n*(i-1)+1):i*n,1:n])
  end

  if trans == true
    y = y'
  end

  return Matrix(y)

end

function kron_prod_times_vector(a::AbstractArray{T,2},b::AbstractArray{T,2},v::AbstractArray{T,1}) where {T <: Real}

    # This function efficiently computes kron(a,b)*v where a and b are matrices
    # and v is a vector.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    n5      = length(v)

    if n5 != n2*n4
        error("a, b, and v do not have the correct sizes")
    end

    p = vec(b*reshape(v,n4,n2)*a')

    return p

end

function vector_times_kron_prod(v::AbstractArray{T,1},a::AbstractArray{T,2},b::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_vector(a',b',v')')

    return product

end

function kron_prod_times_matrix(a::AbstractArray{T,2},b::AbstractArray{T,2},v::AbstractArray{T,2}) where {T <: Real}

    # This function efficiently computes kron(a,b)*v where a, b, and v are
    # conformable matrices.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    (n5,n6) = size(v)

    if n5 != n2*n4
        error("a, b, and v do not have the correct sizes")
    end

    p = zeros(n1*n3,n6)
    for i = 1:n6
        @views p[:,i] = kron_prod_times_vector(a,b,v[:,i])
    end

    return p

end

function matrix_times_kron_prod(v::AbstractArray{T,2},a::AbstractArray{T,2},b::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_matrix(a',b',v')')

    return product

end

function kron_prod_times_vector(a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2},v::AbstractArray{T,1}) where {T <: Real}

    # This function efficiently computes kron(a,kron(b,c))*v where a, b, and c
    # are matrices and v is a vector.  Included for completeness, but not
    # actually used.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    (n5,n6) = size(c)
    n7      = length(v)

    if n7 != n2*n4*n6
        error("a, b, c, and v do not have the correct sizes")
    end

    v_tilda = reshape(v,n4*n6,n2)*a'
    p = vec(kron_prod_times_matrix(b,c,v_tilda))

    return p

end

function vector_times_kron_prod(v::AbstractArray{T,1},a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_vector(a',b',c',v')')

    return product

end

function kron_prod_times_matrix(a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2},v::AbstractArray{T,2}) where {T <: Real}

    # This function efficiently computes kron(a,kron(b,c))*v where a, b, c and
    # v are comformable matrices.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    (n5,n6) = size(c)
    (n7,n8) = size(v)

    if n7 != n2*n4*n6
        error("a, b, c, and v do not have the correct sizes")
    end

    p = zeros(n1*n3*n5,n8)
    for i = 1:n8
        @views v_tilda = reshape(v[:,i],n4*n6,n2)*a'
        p[:,i] .= vec(kron_prod_times_matrix(b,c,v_tilda))
    end

    return p

end

function matrix_times_kron_prod(v::AbstractArray{T,2},a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_matrix(a',b',c',v')')

    return product

end

function dsylvester(a::AbstractArray{T,2}, b::AbstractArray{T,2}, c::Union{AbstractArray{T,1},AbstractArray{T,2}}) where {T <: Real}

    #= Uses the Hessenberg-Schur method to find the bounded solution of the
       discrete Sylvester equation:

             X + A*X*B = C

       Based on Golub, Nash, and Van Loan (1979).

    =#

    n = size(a, 1)
    m = size(b, 1)
    x = zeros(size(c))

    (s,u) = schur(Matrix(b'))
    (v,t) = hessenberg(a)

    c = v'*c*u

    j = m
    while j > 0
        j1 = j
        if j == 1
            block_size = 1
        elseif isequal(s[j,j-1],0.0) == false
            block_size = 2
            j -= 1
        else;
            block_size = 1
        end
        @views ajj = kron(s[j:j1,j:j1],t) + I
        @views rhs = vec(c[:,j:j1])
        if j1 < m
            @views rhs2 = t*(x[:,(j+1):m]*s[j:j1,(j+1):m]')
            rhs -= vec(rhs2)
        end
        w = ajj\rhs
        @views x[:,j] = w[1:n]
        if block_size == 2
            @views x[:,j1] = w[(n+1):2*n]
        end
        j -= 1
    end

    x = v*x*u'

    return x

end

function trm(x::AbstractArray{T,2}) where {T <: Real}

    # Computes the matrix trace as defined by Binning (2013).  Used for
    # computing the second-order terms, hss, gss.

    (n1,n2) = size(x)
    k = Int(round(sqrt(n2)))

    y = zeros(n1)
    for i = 1:k
        @views y += x[:,i+(i-1)*k]
    end

    return y

end

function trm2(x::AbstractArray{T,2}) where {T <: Real}

    # Computes the matrix trace as defined by Binning (2013).  Used for
    # computing the third-order terms, hssx, gssx.

    (n1,n2) = size(x)
    k = Int(round(n2^(1//3)))

    y = zeros(n1,k)
    for j = 1:k
        for i = 1:k
            @views y[:,j] += x[:,(j-1)+i+(i-1)*k^2]
        end
    end

    return y

end

function create_omega3(n::S) where {S <: Integer}

    # Creates the combination matrix for a third-order perturbation as defined in Levintal (2017).

    # This function is a simplified version of the create_OMEGA function originally written in
    # Matlab by Oren Levintal for his paper "Fifth Order Perturbation Solution to DSGE Models"
    # published in the Journal of Economic Dynamics and Control, 2017.  Permission to translate
    # this function into Julia and release it within the SolveDSGE module was granted by Oren
    # Levintal on February 5, 2020.

    ind = [1:n^3;]
    M = reshape(ind,1,n,n,n)
    Ix = eye(S,n^3)
    omega3 = (reshape(Ix[:,PermutedDimsArray(M,[1,4,2,3])],n^3,n^3)
           + reshape(Ix[:,PermutedDimsArray(M,[1,2,4,3])],n^3,n^3)
           + reshape(Ix[:,PermutedDimsArray(M,[1,2,3,4])],n^3,n^3))

    return omega3

end

function kron_prod_times_vector(A::Union{Array{Array{T,2},1},Array{Array{Complex{T},2},1}},x::Union{Array{T,1},Array{Complex{T},1}},n::Array{S,1},p::S) where {T <: Real, S <: Integer}

    # Computes y = (A[p] * A[p-1] * ... * A[1] )*x

    N = prod(n[1:p])
    z = copy(x)
    for i = 1:p
        z = (A[i]*reshape(z,n[i],Int(N/n[i])))'
    end
    y = reshape(Matrix(z),N)

    return y

end

function kron_prod_times_matrix(A::Union{Array{Array{T,2},1},Array{Array{Complex{T},2},1}},x::Union{Array{T,2},Array{Complex{T},2}},n::Array{S,1},p::S) where {T <: Real, S <: Integer}

    # Computes y = (A[p] * A[p-1] * ... * A[1] )*x

    N = prod(n[1:p])
    y = Array{Complex{T}}(undef,N,size(x,2))
    for j = 1:size(x,2)
        @views z = x[:,j]
        for i = 1:p
            z = (A[i]*reshape(z,n[i],Int(N/n[i])))'
        end
        y[:,j] .= reshape(Matrix(z),N)
    end

    return y

end

function KPShiftSolve(TT::Union{Array{Array{T,2},1},Array{Array{Complex{T},2},1}},n::Array{S,1},c::Union{Array{T,1},Array{Complex{T},1}},lambda::T,alpha::Union{T,Complex{T}}) where {T <: Real, S <: Integer}

    p = length(n)
    N = prod(n)

    c = copy(c)

    TT[p] = alpha*TT[p]
    if p == 1
        y = (TT[1] + lambda*Matrix{T}(I,n[1],n[1]))\c
    else
        y = Array{Complex{T}}(undef,N)
        mp = Int(N/n[p])
        for i = n[p]:-1:1
            idx = ((i-1)*mp+1):(i*mp)
            y[idx] = KPShiftSolve(TT[1:(p-1)],n[1:(p-1)],c[idx],lambda,TT[p][i,i])
            z = kron_prod_times_vector(TT[1:p],y[idx],n,p-1)
            for j = 1:(i-1)
                jdx = ((j-1)*mp+1):(j*mp)
                c[jdx] = c[jdx] - TT[p][j,i]*z
            end
        end
    end

    return y

end

function martin_van_loan(a::Array{T,2},b::Array{T,2},c::Array{T,2},d::Array{T,2},k::S) where {T <: Real, S <: Integer}

    #= Uses a recursive Schur method to find the bounded solution of the Sylvester
       equation:

       AX + B*X*(Kron^(k)C) = D

       Based on Martin and Van Loan (2006).  This is a simplified implementation
       of their algorithm, but it captures most of the gains over Golub, Nash,
       and van Loan (1979).

    =#

    a = copy(a)
    b = a\copy(b)
    c = copy(c)
    d = a\copy(d)

    (v,s) = hessenberg(b)       # v*s*v' = b
    (t,q) = schur(complex(c'))  # q*t*q' = c'

    v = Matrix(v)

    p = k + 2
    TT = Array{typeof(t)}(undef,p)
    TT[1] = s
    for i = 2:p
        TT[i] = conj(t)
    end

    Q     = fill(q,k+1)
    inv_Q = fill(Matrix(q'),k+1)

    n = fill(size(c,1),p)
    n[1] = size(b,1)
    N = prod(n)

    lambda = 1.0

    e = vec(v'*kron_prod_times_matrix(inv_Q,Matrix(d'),n[2:end],k+1)')

    y = reshape(KPShiftSolve(TT,n,e,lambda,1.0),size(d))
    x = real(Matrix(v*kron_prod_times_matrix(Q,Matrix(y'),n[2:end],k+1)'))

    return x

end

function dlyap(a::Array{T,2}, b::Array{T,2}) where {T <: Real}

    n = size(a,1)
    x = zeros(n,n)
    j = n

    (s,u) = schur(a)
    b = u'b*u

    while j > 0
        j1 = j
        if j == 1
            block = 1
        elseif !isequal(s[j,j-1],0.0)
            block = 2
            j -= 1
        else
            block = 1
        end
        @views lhs = kron(s[j:j1,j:j1],s) - I # I = eye(block*n)
        @views rhs = vec(b[:,j:j1])
        if j1 < n
            @views rhs2 = s*(x[:,(j1+1):n]*s[j:j1,(j1+1):n]')
            rhs += vec(rhs2)
        end
        w = -lhs\rhs
        @views x[:,j] = w[1:n]
        if block == 2
            @views x[:,j1] = w[(n+1):block*n]
        end
        j -= 1
    end

    x = u*x*u'

    return x

end

function ind2sub(i::S,dims::Tuple{S,Vararg{S}}) where {S <: Integer}

    if i < 1 || i > prod(dims)
        error("index is out of bounds.")
    end

    subs = Tuple(CartesianIndices(dims)[i])

    return subs

end

function compute_variances(soln::FirstOrderSolutionStoch)

    hx = soln.hx
    k = soln.k
    gx = soln.gx
    sigma = soln.sigma

    var_states = dlyap(hx,k*sigma*k')
    var_jumps  = gx*var_states*gx'

  return var_states,var_jumps

end

function compute_chebyshev_integrals(eps_nodes,eps_weights,nodes,order,rho,sigma)

  # A simplified implementation of the integration described in Judd et al (2017, section 5.1);
  # it uses the mean integral across the nodes.
  # Using integrals[i] = sum(exp.(sqrt(2)*sigma*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
  # leads to something much simplier and equally accurate, but is less recognizably appropriate
  # for the case where the shocks are AR(1) processes and ordinary polynomials are not being used.

  terms_num  = Array{Float64}(undef,length(eps_nodes))
  integrals  = Array{Float64}(undef,order+1,length(nodes))
  integrals2 = Array{Float64}(undef,order+1)
  for i = 1:(order+1)
    integrals2[i] = sum(exp.(sqrt(2)*sigma*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
    for j = 1:length(nodes)
      terms_num     .= rho*nodes[j] .+ sqrt(2)*sigma*eps_nodes
      terms_den      = rho*nodes[j]
      terms_num     .= chebyshev_polynomial(i,terms_num)[:,i]
      terms_den      = chebyshev_polynomial(i,terms_den)[i]
      integrals[i,j] = sum((terms_num/terms_den).*eps_weights)*pi^(-1/2)
     end
  end

  nodetoosmall = abs.(nodes) .< sqrt(eps())
  if sum(nodetoosmall) > 0
    if length(nodes) == 1
      integrals[:,1] .= integrals2
    else
      for i = 1:length(nodes)
        if nodetoosmall[i] == 1
          if i == 1
            integrals[:,i] .= integrals[:,i+1]
        elseif i == length(nodes)
            integrals[:,i] .= integrals[:,i-1]
          else
            integrals[:,i] .= (integrals[:,i-1]+integrals[:,i+1])/2
          end
        end
      end
    end
  end

  return reshape(sum(integrals,dims=2)/length(nodes),order+1)

#  integrals = Array{Float64}(undef,order+1)
#  for i = 1:(order+1)
#    integrals[i] = sum(exp.(sqrt(2)*sigma*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
#  end
#  return integrals

end

function compute_smolyak_integrals(eps_nodes,eps_weights,nx,order,grid,RHO,sigma)

  integrals = ones(nx,order+1)
  for j = 1:size(sigma,2)
      nodes = unique(grid[:,j])
      integrals[j,:] .= compute_chebyshev_integrals(eps_nodes,eps_weights,nodes,order,RHO[j,j],sigma[j,j])
  end

#  for j = 1:size(sigma,2)
#      for i = 1:(order+1)
#          integrals[j,i] = sum(exp.(sqrt(2)*sigma[j,j]*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
#      end
#  end

  return integrals

end

function weight_scale_factors(eps_nodes,eps_weights,multi_index,nx,grid,RHO,sigma)

  unique_multi_index = sort(unique(multi_index))
  unique_orders = SmolyakApprox.m_i(unique_multi_index).-1

  # Here we construct the base integrals

  base_integrals = Array{Array{Float64,2}}(undef,length(unique_orders))
  for i = 1:length(unique_orders)
    base_integrals[i] = compute_smolyak_integrals(eps_nodes,eps_weights,nx,unique_orders[i],grid,RHO,sigma)
  end

  # Compute the unique polynomial terms from the base polynomials

  unique_base_integrals = Array{Array{Float64,2}}(undef,length(unique_orders))
  for i = length(unique_orders):-1:2
    unique_base_integrals[i] = base_integrals[i][:,size(base_integrals[i-1],2)+1:end]
  end
  unique_base_integrals[1] = base_integrals[1]

  # Construct the first row of the interplation matrix

  new_integrals = unique_base_integrals[multi_index[1,1]][1,:]
  for i = 2:size(multi_index,2)
    new_integrals = kron(new_integrals,unique_base_integrals[multi_index[1,i]][i,:])
  end

  weight_scale_factor = copy(new_integrals)

  # Iterate over nodes, doing the above three steps at each iteration

  for j = 2:size(multi_index,1)

    new_integrals = unique_base_integrals[multi_index[j,1]][1,:]
    for i = 2:size(multi_index,2)
      new_integrals = kron(new_integrals,unique_base_integrals[multi_index[j,i]][i,:])
    end
    weight_scale_factor = [weight_scale_factor; new_integrals]

  end

  return weight_scale_factor

end

function scale_weights(weights,weight_scale_factor)

    scaled_weights = weights.*weight_scale_factor

    return scaled_weights

end

function compute_piecewise_linear_integrals(eps_nodes,eps_weights,sigma)

  integral = 1.0
  #integral = sum(exp.(sqrt(2)*sigma*eps_nodes).*eps_weights)*pi^(-1/2)

  return integral

end


"""
    eom_linear3d_nonlinearcoupling(x, p, n) -> Function

Equations of motion for a 3d linear system with nonlinear coupling [1].
The difference equations are

```math
\\begin{aligned}
x_1(t+1) &= a_1 x_1 (1-x_1(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{1}(t) \\
x_2(t+1) &= a_1 x_2 (1-x_2(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{2}(t) + b x_1 x_2 \\
x_3(t+1) &= a_3 x_3 (1-x_3(t))^2  e^{-x_3(t)^2} + 0.4 \\xi_{3}(t) + c x_{2}(t) \\
         &+ d x_{1}(t)^2.
\\end{aligned}
```

Here, ``\\xi_{1,2,3}(t)`` are independent normally distributed noise processes,
representing dynamical noise in the system, with zero mean and standard
deviations ``\\sigma_1``, ``\\sigma_2``, ``\\sigma_3``, respectively.

# References
1. Gourévitch, B., Le Bouquin-Jeannès, R., & Faucon, G. (2006). Linear and
nonlinear causality between signals: methods, examples and neurophysiological
applications. Biological Cybernetics, 95(4), 349–369.
"""
function eom_linear3d_nonlinearcoupling(x, p, n)
    x₁, x₂, x₃ = (x...,)
    a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃ = (p...,)
    ξ₁ = rand(Normal(0, σ₁))
    ξ₂ = rand(Normal(0, σ₂))
    ξ₃ = rand(Normal(0, σ₃))

    dx₁ = a₁*x₁*(1-x₁)^2 * exp(-x₁^2) + 0.4*ξ₁
    dx₂ = a₂*x₂*(1-x₂)^2 * exp(-x₂^2) + 0.4*ξ₂ + b*x₁*x₂
    dx₃ = a₃*x₃*(1-x₃)^2 * exp(-x₃^2) + 0.4*ξ₃ + c*x₂ + d*x₁^2

    return SVector{3}(dx₁, dx₂, dx₃)
end

function linear3d_nonlinearcoupling(uᵢ, a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃)
    p = [a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃]
    s = DiscreteDynamicalSystem(eom_linear3d_nonlinearcoupling, uᵢ, p)
    return s
end

"""
    linear3d_nonlinearcoupling(;uᵢ = rand(3), σ₁ = 1.0, σ₂ = 1.0, σ₃ = 1.0,
        a₁ = 3.4, a₂ = 3.4, a₃ = 3.4, b = 0.5, c = 0.3, d = 0.5) -> DiscreteDynamicalSystem

A 3d linear system with nonlinear coupling [1]. The difference equations are

```math
\begin{aligned}
x_1(t+1) &= a_1 x_1 (1-x_1(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{1}(t) \\
x_2(t+1) &= a_1 x_2 (1-x_2(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{2}(t) + b x_1 x_2 \\
x_3(t+1) &= a_3 x_3 (1-x_3(t))^2  e^{-x_3(t)^2} + 0.4 \\xi_{3}(t) + c x_{2}(t) \\
         &+ d x_{1}(t)^2.
\end{aligned}
```

Here, ``\\xi_{1,2,3}(t)`` are independent normally distributed noise processes,
representing dynamical noise in the system, with zero mean and standard
deviations ``\\sigma_1``, ``\\sigma_2``, ``\\sigma_3``, respectively.

# References
1. Gourévitch, B., Le Bouquin-Jeannès, R., & Faucon, G. (2006). Linear and
nonlinear causality between signals: methods, examples and neurophysiological
applications. Biological Cybernetics, 95(4), 349–369.
"""
linear3d_nonlinearcoupling(;uᵢ = rand(3), σ₁ = 1.0, σ₂ = 1.0, σ₃ = 1.0,
        a₁ = 3.4, a₂ = 3.4, a₃ = 3.4, b = 0.5, c = 0.3, d = 0.5) =
    linear3d_nonlinearcoupling(uᵢ, a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃)


"""
    trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)

Return the circumcenter of the triangle defined by `pt1`,
`pt2`, and `pt3`. The circumcenter is the center of a circle
that passes through the vertices of the triangle.

"""
function trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)
    d = 2(pt1.x * (pt2.y - pt3.y) + pt2.x * (pt3.y - pt1.y) + pt3.x * (pt1.y - pt2.y))
    x = ((pt1.x^2 + pt1.y^2) * (pt2.y - pt3.y) + (pt2.x^2 + pt2.y^2) * (pt3.y - pt1.y) + (pt3.x^2 + pt3.y * pt3.y) * (pt1.y - pt2.y)) / d
    y = ((pt1.x^2 + pt1.y^2) * (pt3.x - pt2.x) + (pt2.x^2 + pt2.y^2) * (pt1.x - pt3.x) + (pt3.x^2 + pt3.y * pt3.y) * (pt2.x - pt1.x)) / d
    return Point(x, y)
end

# """
#     trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)
# """
# function trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)
#     ap1, ap2 = perpendicular(pt1, pt2)
#     bp1, bp2 = perpendicular(pt2, pt3)
#     flag, ip = intersectionlines(ap1, ap2, bp1, bp2)
#     if flag
#         return ip
#     else
#         throw(error("trianglecircumcenter(): triangle is not valid"))
#     end
# end

"""
    triangleincenter(pt1::Point, pt2::Point, pt3::Point)

Return the incenter of the triangle defined by `pt1`, `pt2`,
and `pt3`. The incenter is the center of a circle inscribed
inside the triangle.
"""
function triangleincenter(pt1::Point, pt2::Point, pt3::Point)
    a = distance(pt2, pt3)
    b = distance(pt1, pt3)
    c = distance(pt1, pt2)
    x = ((a * pt1.x) + (b * pt2.x) + (c * pt3.x)) / (a + b + c)
    y = ((a * pt1.y) + (b * pt2.y) + (c * pt3.y)) / (a + b + c)
    return Point(x, y)
end

"""
    trianglecenter(pt1::Point, pt2::Point, pt3::Point)

Return the centroid of the triangle defined by `pt1`, `pt2`, and `pt3`.
"""
function trianglecenter(pt1::Point, pt2::Point, pt3::Point)
    return 1/3 * (pt1 + pt2  + pt3)
end

"""
    triangleorthocenter(pt1::Point, pt2::Point, pt3::Point)

Return the orthocenter of the triangle defined by `pt1`, `pt2`, and `pt3`.
"""
function triangleorthocenter(pt1::Point, pt2::Point, pt3::Point)
    return pt1 + pt2 + pt3 - 2trianglecircumcenter(pt1, pt2, pt3)
end


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

mu, sig = 50, 5
dist = Normal(mu, sig)
alpha = 0.01
nMax = 40

observations = rand(dist,1)
piLarray, piUarray = [], []

for _ in 2:nMax
    xNew = rand(dist)
    push!(observations,xNew)

    xbar, sd = mean(observations), std(observations)
    n = length(observations)
    tVal = quantile(TDist(n-1),1-alpha/2)
    delta = tVal * sd * sqrt(1+1/n)
    piL, piU = xbar - delta, xbar + delta

    push!(piLarray,piL); push!(piUarray,piU)
end

scatter(1:nMax, observations,
	c=:blue, msw=0, label="Observations")
plot!(2:nMax, piUarray,
	c=:red, shape=:xcross, msw=0, label="Prediction Interval")
plot!(2:nMax, piLarray,
	c=:red, shape=:xcross, msw=0, label="",
	ylims=(0,100), xlabel="Number of observations", ylabel="Value")


"""
    second-order cone
"""
function second_order_cone_projection(z)
    n = length(z)

    z0 = z[1]
    z1 = view(z, 2:n)

    if norm(z1) <= z0
        return z, true
    elseif norm(z1) <= -z0
        return zero(z), false
    else
        a = 0.5 * (1.0 + z0 / norm(z1))
        z_proj = zero(z)
        z_proj[1] = a * norm(z1)
        z_proj[2:end] = a * z1
        return z_proj, false
    end
end

function second_order_cone_product(z, s)
    n = length(z)
    SVector{n}([z' * s; z[1] * view(s, 2:n) + s[1] * view(z, 2:n)])
end

# check that second-order cone constraints are satisfied
function second_order_cone_check(x, idx_soc)
    for idx in idx_soc
        !second_order_cone_projection(view(x, idx))[2] && return false
    end
    return true
end

# check that inequality (non-negative orthant) constraints are satisfied
function inequality_check(x, idx_ineq)
    for i in idx_ineq
        if x[i] <= 0.0
            return false
        end
    end
    return true
end

function cone_check(x, idx_ineq, idx_soc)
    !inequality_check(x, idx_ineq) && return true
    !second_order_cone_check(x, idx_soc) && return true
end


import Remez
using SpecialFunctions: erfcx
using PyPlot

T = Float64

function g(y::T) where T <: AbstractFloat
    two = parse(T, "2")
    rootpi = sqrt(convert(T, π))
    if -one(T) ≤ y ≤ -one(T) + eps(T)
        return ( 1 / two - 1 / rootpi )
    elseif one(T) - eps(T) ≤ y ≤ one(T)
        return ( -1 + 1/rootpi ) / 2
    else
        x = (1+y) / (1-y)
        return ( 2*erfcx(x) / (1-y) - 1 ) / (1+y)
    end
end

function erfcx_approx(x)
    y = (x-1) / (x+1)
    return (1-y)*(1+(1+y)*g_approx(y)) / 2
end

n = 15
iterations = 6
pt, ddp, E = Remez.minimax(T, g, n, iterations)

g_approx(y) = Remez.Newton_poly(ddp[0:n], pt[0:n], y)

figure(1)
y = range(-one(T), one(T), length=201)

plot(y, g.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"$g(y)$")

figure(2)
plot(y, g.(y)-g_approx.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"error in $g$")

figure(3)
x = range(0, 10, length=201)
plot(x, erfcx_approx.(x) - erfcx.(x))
grid(true)
xlabel(L"$x$")
ylabel("error in erfcx")


using FdeSolver
using Plots

## inputs
tSpan = [0, 10]     # [intial time, final time]
x0 = [1, 1]             # intial value
β = 2            # order of the derivative
par = [16.0, 4.0] # [spring constant for a mass on a spring, inertial mass]
h = 0.01

## Equation
function F(t, n, β, x, par)

      K = par[1]
      m = par[2]

      - K ./ m .* x[n]

end

## Numerical solution
t, Yapp = FDEsolver(F, tSpan, x0, β, nothing, par, h = h)

#plot
plot(t, Yapp, linewidth = 5, title = "Simple harmonic motion (order=2)",
     xaxis = "Time (t)", yaxis = "x(t)", label = "Approximation")
a = x0[1] .* map(cos, sqrt(par[1] / par[2]) .* t) .+ x0[2] ./ sqrt(par[1] / par[2]) .* map(sin, sqrt(par[1] / par[2]) .* t)
plot!(t, a, lw = 3, ls = :dash, label = "Exact solution")


using SMM, DataFrames, CSV, HTTP, Random, LinearAlgebra, Statistics, Optim, DataFramesMeta, GLM

function PS7()
	#question 1
	url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS1-julia-intro/nlsw88.csv"
	df = CSV.read(HTTP.get(url).body, DataFrames)
	X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
	y = log.(df.wage);

	function ols_gmm(beta, X, y)
		return (y.-X*beta)'*(y.-X*beta)
	end

	function ols_gmm_with_sigma(beta, X, y)
		g = vcat(y .- X*beta[1:end-1],( (size(y,1)-1)/(size(y,1)-size(X,2)) )*var(y .- X*beta[1:end-1]) .- beta[end]^2)
		return g'*g
	end

	beta_gmm = optimize(b -> ols_gmm(b, X, y), rand(size(X,2)), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000))
	println(beta_gmm.minimizer)

	beta_gmm_sigma = optimize(b -> ols_gmm_with_sigma(b, X, y), rand(size(X,2)+1), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000))
	println(beta_gmm_sigma.minimizer)

	#Question 2
	df = dropmissing(df, :occupation)
	df[df.occupation.>7,:occupation] .= 7
	X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
	y = df.occupation;

	#from PS2
	function mlogit_mle(alpha, X, y)

			K = size(X,2)
			J = length(unique(y))
			N = length(y)
			bigY = zeros(N,J)
			for j=1:J
				bigY[:,j] = y.==j
			end
			A = [reshape(alpha,K,J-1) zeros(K)]

			P = exp.(X*A) ./ sum.(eachrow(exp.(X*A)))

			return -sum( bigY.*log.(P) )
	end

	alpha_hat= optimize(a -> mlogit_mle(a, X, y), rand(6*size(X,2)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))

	print(alpha_hat.minimizer)

	function mlogit_gmm(α, X, y)
			K = size(X,2)
			J = length(unique(y))
			N = length(y)
			bigY = zeros(N,J)
			for j=1:J
				bigY[:,j] = y.==j
			end
			bigα = [reshape(α,K,J-1) zeros(K)]

			P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

			g = bigY[:] .- P[:]

			J = g'*I*g
			return J
	end

	α_hat_gmm = optimize(a -> mlogit_gmm(a, X, y), alpha_hat.minimizer, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))
	print(α_hat_gmm.minimizer)

	td = TwiceDifferentiable(b -> mlogit_gmm(b, X, y), alpha_hat.minimizer; autodiff = :forward)
	α_hat_gmm_random=optimize(td, rand(6*size(X,2)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))

	print(α_hat_gmm_random.minimizer)

	#Question 3

	#Question4
	MA = SMM.parallelNormal()
	dc = SMM.history(MA.chains[1])
	dc = dc[dc[:accepted].==true, :]
	println(describe(dc))

	#Question 5
end

PS7()


using ChaosTools
using ChaosTools.DynamicalSystemsBase
using Test
using StatsBase
using Statistics
using Entropies

test_value = (val, vmin, vmax) -> @test vmin <= val <= vmax

println("\nTesting generalized entropy (genentropy) & linear scaling...")
@testset "Generalized Dimensions" begin
    @testset "Henon Map" begin
        ds = Systems.henon()
        ts = trajectory(ds, 200000)
        es = 10 .^ range(-0, stop = -3, length = 7)
        dd = zero(es)
        for q in [0,2,1, 2.56]
            for (i, ee) in enumerate(es)
                dd[i] = genentropy(ts, ee; q)
            end
            linr, dim = linear_region(-log.(es), dd)
            test_value(dim, 1.1, 1.3)
        end
    end
    @testset "Lorenz System" begin
        ds = Systems.lorenz()
        ts = trajectory(ds, 5000)
        es = 10 .^ range(1, stop = -3, length = 11)
        dd = zero(es)
        for q in [0,1,2]
            for (i, ee) in enumerate(es)
                dd[i] = genentropy(ts, ee; q)
            end
            linr, dim = linear_region(-log.(es), dd)
            if q == 0
                test_value(dim, 1.7, 1.9)
            else
                test_value(dim, 1.85, 2.0)
            end
        end
    end
end

println("\nTesting generalized entropy using Molteno's boxing method...")
@testset "Molteno's boxing method" begin
    @testset "Henon Map" begin
        ds = Systems.henon()
        ts = trajectory(ds, 200000)
        boxes, ϵs = molteno_boxing(ts)
        boxes = Probabilities.(boxes)
        for q in [0,2,1, 2.56]
            dd = genentropy.(boxes; q)
            linr, dim = linear_region(-log.(ϵs), dd)
            test_value(dim, 1.1, 1.3)
        end
    end
    @testset "Lorenz System" begin
        ds = Systems.lorenz()
        ts = trajectory(ds, 5000)
        boxes, ϵs = molteno_boxing(ts)
        boxes = Probabilities.(boxes)
        for q in [0,2,1, 2.56]
            dd = genentropy.(boxes; q)
            linr, dim = linear_region(-log.(ϵs), dd)
            test_value(dim, 1.85, 2.2)
        end
    end
end

println("\nTesting dimension calls (all names)...")
@testset "Dimension calls" begin
    ds = Systems.henon()
    ts = trajectory(ds, 50000)
    test_value(generalized_dim(ts; q = 1.32), 1.0, 1.3)
    test_value(generalized_dim(ts; base = 2, q = 1.32), 1.0, 1.3)
end


module problem114

description = """
Counting block combinations I
A row measuring seven units in length has red blocks with a minimum length of three units placed on it,
 such that any two red blocks (which are allowed to be different lengths) are separated by at least one black square.
There are exactly seventeen ways of doing this.
How many ways can a row measuring fifty units in length be filled?
"""

function block_combinations(n :: Integer, memo = zeros(Int, n))
  # trivial cases
  if n < 1 return 1 end

  # Count the total with each block size, and block position,
  # recursively counting the ways to fill the remaining space at the end of the row.
  if memo[n] == 0
    total = 1
    for blocksize = 3:n, blockpos = 1:(n-blocksize+1)
      total += block_combinations(n - blockpos - blocksize, memo)
    end
    memo[n] = total
  end
  return memo[n]
end

using Base.Test

@test block_combinations(3) == 2 #empty or 1 block
@test block_combinations(4) == 4 #empty or length 3 at pos 1, or length 3 at pos 2, or length 4 at pos 1
@test block_combinations(7) == 17

end


#  Note:  evalVolume is tested under test_linearFFD.jl since it is needed
#         everytime a control point is perturbed to compute the new location of
#         the geometric point

@testset "--- Checking B-spline Formulation ---" begin

  U = [0,0,0,0,2/5,3/5,3/5,1,1,1,1]
  order = 4
  p = order - 1
  u = 0.42
  nctl = length(U) - order

  @test ( nctl )== 7

  @testset "Checking knot span index evaluation" begin
    span = FFD.findSpan(u, U, order, nctl)
    @test ( span )== 5
  end

  @testset "Checking basis function evaluations" begin
    N = zeros(order)
    span = FFD.findSpan(u, U, order, nctl)
    FFD.basisFunctions(U, order, u, span, N)
    @test isapprox( N[1], 0.081) atol= 1e-15
    @test isapprox( N[2], 0.405) atol= 1e-15
    @test isapprox( N[3], 0.5136666666666666) atol= 1e-15
    @test isapprox( N[4], 0.00033333333333333153) atol= 1e-15
  end

  @testset "Checking curve point evaluation" begin
    P = 0:1/(nctl-1):1
    C = [0.0]
    FFD.evalCurve([u], U, order, P, C)
    @test isapprox( C[1], 0.40555555555555556) atol= 1e-16
  end

end  # End facts("--- Check Basis Function Calculations ---")

@testset "--- Checking B-spline Derivatives ---" begin

  @testset "Checking Basis Function Derivatives" begin

    U = [0,0,0,0,2/5,3/5,3/5,1,1,1,1]
    order = 4
    p = order - 1
    u = 0.42
    nctl = length(U) - order
    N = zeros(AbstractFloat, order)
    Nderiv = zeros(AbstractFloat, order)
    span = FFD.findSpan(u, U, order, nctl)
    FFD.derivBasisFunctions(u, U, order, span, N, Nderiv)

    @test isapprox( N[1], 0.081) atol= 1e-15
    @test isapprox( N[2], 0.405) atol= 1e-15
    @test isapprox( N[3], 0.5136666666666666) atol= 1e-15
    @test isapprox( N[4], 0.00033333333333333153) atol= 1e-15
    @test isapprox( Nderiv[1], -1.3500000000000003) atol= 1e-15
    @test isapprox( Nderiv[2], -2.25) atol= 1e-15
    @test isapprox( Nderiv[3], 3.55) atol= 1e-15
    @test isapprox( Nderiv[4], 0.04999999999999982) atol= 1e-15

  end  # End context("Checking Basis Function Derivatives")

end  # End facts("--- Checking B-spline Derivatives ---")


using LinearAlgebra, LowRankApprox, Formatting, ImageFiltering, Images, Plots, ImageMagick, Colors, TestImages, ImageView, ImageTransformations

I=load("rgb.jpg");

G=Gray.(I);
# size(G)

A=imresize(G, ratio=1/3);
# rank(A), size(A), eltype(A)

(p,q)=size(A);
B=Array{Float64}(A);

using DSP, Wavelets
B1=dwt(B, wavelet(WT.haar), 1);
B2=dwt(B, wavelet(WT.haar), 2);
B3=dwt(B, wavelet(WT.haar), 3);
B4=dwt(B, wavelet(WT.haar), 4);

B1a=zeros(Float64, p, q);
B2a=zeros(Float64, p, q);
B3a=zeros(Float64, p, q);
B4a=zeros(Float64, p, q);

B1a[1:Int(p/2),1:Int(q/2)]=B1[1:Int(p/2),1:Int(q/2)];
B2a[1:Int(p/4),1:Int(q/4)]=B2[1:Int(p/4),1:Int(q/4)];
B3a[1:Int(p/8),1:Int(q/8)]=B3[1:Int(p/8),1:Int(q/8)];
B4a[1:Int(p/16),1:Int(q/16)]=B4[1:Int(p/16),1:Int(q/16)];
B1d=B1-B1a;
B2d=B2-B2a;
B3d=B3-B3a;
B4d=B4-B4a;

C1a=idwt(B1a, wavelet(WT.haar), 1);
C2a=idwt(B2a, wavelet(WT.haar), 2);
C3a=idwt(B3a, wavelet(WT.haar), 3);
C4a=idwt(B4a, wavelet(WT.haar), 4);
C1d=idwt(B1d, wavelet(WT.haar), 1);
C2d=idwt(B2d, wavelet(WT.haar), 2);
C3d=idwt(B3d, wavelet(WT.haar), 3);
C4d=idwt(B4d, wavelet(WT.haar), 4);

A1a=Gray.(C1a);
A1d=Gray.(C1d);
A2a=Gray.(C2a);
A2d=Gray.(C2d);
A3a=Gray.(C3a);
A3d=Gray.(C3d);
A4a=Gray.(C4a);
A4d=Gray.(C4d);

plot(Gray.(A),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 0: Original Grayscale Image")
savefig("grayscale_haar_level0.png")

plot(Gray.(B1),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 1 (haar)")
savefig("grayscale_haar_level1.png")

plot(Gray.(C1a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 1 Approximation (haar)")
savefig("grayscale_haar_level1_a.png")

plot(Gray.(C2a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 2 Approximation (haar)")
savefig("grayscale_haar_level2_a.png")

plot(Gray.(C3a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 3 Approximation (haar)")
savefig("grayscale_haar_level3_a.png")

plot(Gray.(C4a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 4 Approximation (haar)")
savefig("grayscale_haar_level4_a.png")

plot(Gray.(C4d),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 4 Detail (haar)")
savefig("grayscale_haar_level4_d.png")


"""
    MVCalc(μstar,μ,Σ)

Calculate the std and weights of a portfolio (with mean return μstar) on MVF of risky assets.

# Remark
- Only (λ,δ) and thus (w,stdRp) depend on μstar. We could therefore speed up the computations a bit
by doing the loop over different μstar values inside the function (and thus not recalculate Σ_1,a,b,c).
"""
function MVCalc(μstar,μ,Σ)
    n    = length(μ)
    Σ_1  = inv(Σ)
    a    = μ'Σ_1*μ
    b    = μ'Σ_1*ones(n)
    c    = ones(n)'Σ_1*ones(n)
    λ    = (c*μstar - b)/(a*c-b^2)
    δ    = (a-b*μstar)/(a*c-b^2)
    w    = Σ_1 *(μ*λ.+δ)
    StdRp = sqrt(w'Σ*w)
    return StdRp,w
end

"""
    MVCalcRf(μstar,μ,Σ,Rf)

Calculate the std and portfolio weights of a portfolio (with a given mean, μstar) on MVF of (risky assets,riskfree)
"""
function MVCalcRf(μstar,μ,Σ,Rf)
    μe    = μ .- Rf
    Σ_1   = inv(Σ)
    w     = (μstar-Rf)/(μe'Σ_1*μe) * Σ_1*μe
    StdRp = sqrt(w'Σ*w)
    return StdRp,w                    #std and portfolio weights
end

"""
    MVTangencyP(μ,Σ,Rf)

Calculate the tangency portfolio
"""
function MVTangencyP(μ,Σ,Rf)           #calculates the tangency portfolio
    n    = length(μ)
    μe   = μ .- Rf                    #expected excess returns
    Σ_1  = inv(Σ)
    w    = Σ_1 *μe/(ones(n)'Σ_1*μe)
    muT  = w'μ + (1-sum(w))*Rf
    StdT = sqrt(w'Σ*w)
    return w,muT,StdT                  #portolio weights, mean and std
end


function asinh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && return x
    signbit(x) && return - asinh(-x)
    result = abs(x)
    result = result + sqrt(square(result) + 1.0)
    result = log(result)
    return copysign(result, x)
end

function acosh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    (x < 1.0 || isinf(x)) && throw(DomainError("$x"))
    result = x + sqrt(square(x) - 1.0)
    result = log(result)
    return result
end

function atanh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    (abs(x) > 1.0 || isinf(x)) && throw(DomainError("$x"))
    twox = DoubleFloat{T}(x.hi+x.hi, x.lo+x.lo)
    result = 1.0 + twox / (1.0 - x)
    result = log(result)
    result = DoubleFloat{T}(result.hi*0.5, result.lo*0.5)
    return result
end

function acsch(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && return zero(DoubleFloat{T})
    iszero(x) && return signbit(x) ? -DoubleFloat{T}(Inf) : DoubleFloat{T}(Inf)
    invx = inv(x)
    invx2 = inv(x*x)
    result = log(sqrt(1 + invx2) + invx)
    return result
end

function asech(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    (x < 0.0 || x > 1.0 || isinf(x)) && throw(DomainError("$x"))
    iszero(x) && return inf(DoubleFloat{T})
    isone(x) && return zero(DoubleFloat{T})
    invx = inv(x)
    result = fma(sqrt(invx + 1), sqrt(invx - 1), invx)
    return log(result)
end

function acoth(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && return copysign(zero(DoubleFloat{T}), x)
    isone(x) && return inf(DoubleFloat{T})
    isone(-x) && return -inf(DoubleFloat{T})
    (-1.0 < x < 1.0) && throw(DomainError("$x"))
    invx = inv(x)
    result = (log(1+invx) - log(1-invx))/2
    return result
end


## Julia program for Hierarchical Modeling of means and variances
## author: weiya <szcfweiya@gmail.com>
## date: 27 August, 2018

using Distributions
using SpecialFunctions
using StatsBase
using DelimitedFiles

function higibbs(Y, T, mu0 = 50.0, gamma20 = 25.0, nu0 = 1.0, sigma20 = 100.0, eta0 = 1.0, tau20 = 100.0, a = 1.0, b = 1/100.0, alpha = 1.0, NUMAX = 40)
    m = size(unique(Y[:,1]), 1)
    # starting value
    ybar = ones(m)
    sv = ones(m)
    n = ones(m)
    for j = 1:m
        yj = Y[ [Y[i,1] == j for i = 1:end], 2]
        ybar[j] = mean(yj)
        sv[j] = var(yj)
        n[j] = size(yj, 1)
    end
    theta = ybar
    sigma2 = copy(sv)
#    sigma20 = 1 / mean(sigma2)
#    nu0 = 2 * mean(sigma2)^2 / var(sigma2)
    mu = mean(theta)
    tau2 = var(theta)

    THETA = ones(T, m)
    SIGMA2 = ones(T, m)
    # mu tau2 sigma20 nu0
    MTSN = ones(T, 4)

    for t = 1:T
        # sample mu
        varmu = 1 / (m / tau2 + 1 / gamma20)
        meanmu = varmu * (m * mean(theta) / tau2 + mu0 / gamma20)
        rnorm = Normal(meanmu, sqrt(varmu))
        mu = rand(rnorm, 1)[1]

        # sample 1/tau2
        shapetau = (eta0 + m) / 2
        ratetau = ( eta0 * tau20 + sum((theta .- mu).^2) ) / 2
        rgamma = Gamma(shapetau, 1/ratetau)
        tau2 = 1 / rand(rgamma, 1)[1]

        # sample theta
        for j = 1:m
            vartheta = 1 / (n[j] / sigma2[j] + 1 / tau2)
            meantheta = vartheta * ( n[j] * mean(Y[ [Y[i,1] == j for i = 1:end], 2]) / sigma2[j] + mu / tau2)
            rnorm = Normal(meantheta, sqrt(vartheta))
            theta[j] = rand(rnorm, 1)[1]
        end
        THETA[t, :] .= theta

        # sample sigma2
        for j = 1:m
            shapesig = (nu0 + n[j])/2
            yj = Y[ [Y[i,1] == j for i = 1:end], 2]
            ratesig = ( nu0*sigma20 + sum( (yj .- theta[j]).^2 ) )/2
            rgamma = Gamma(shapesig, 1/ratesig)
            sigma2[j] = 1 / rand(rgamma, 1)[1]
        end
        SIGMA2[t, :] .= sigma2

        # sample sigma20
        shapesig = a + 0.5 * m * nu0
        ratesig = b + 0.5 * nu0 * sum(1 ./ sigma2)
        rgamma = Gamma(shapesig, 1/ratesig)
        sigma20 = rand(rgamma, 1)[1]

        # sample nu0
        x = 1:NUMAX
        lpnu0 = ones(NUMAX)
        lpnu0 .= m * ( .5 * x .* log.(sigma20 * x / 2) .- lgamma.(x/2) ) .+ (x / 2 .+ 1) * sum(log.(1 ./ sigma2)) .- x .* (alpha + .5 * sigma20 * sum(1 ./ sigma2))
        #println(lpnu0)
        nu0 = sample(x, pweights(exp.(lpnu0 .- maximum(lpnu0))))
#        println(pweights(exp.(lpnu0 .- maximum(lpnu0))))

        # store results
        MTSN[t, :] .= [mu, tau2, sigma20, nu0]
    end
    return THETA, SIGMA2, MTSN, sv, n
end

# run
Y = readdlm("math-score-Y.csv")
THETA, SIGMA2, MTSN, sv, n = higibbs(Y, 5000)

using PyPlot
# histogram
plt[:subplot](221)
plt[:hist](MTSN[:,2])
ylabel(L"$\mu$")
plt[:subplot](222)
plt[:hist](MTSN[:,2])
ylabel(L"$\tau^2$")
plt[:subplot](223)
plt[:hist](MTSN[:,3])
ylabel(L"$\nu_0$")
plt[:subplot](224)
plt[:hist](MTSN[:,4])
ylabel(L"$\sigma_0^2$")
plt[:tight_layout]()
show()

# shrinkage
f, (ax1, ax2) = plt[:subplots](1, 2)
ax1[:scatter](sv, SIGMA2[end,:])
ax1[:plot](sv, sv)
ax1[:set_xlabel](L"$s^2$")
ax1[:set_ylabel](L"$\hat \sigma^2$")
ax2[:scatter](n, sv-SIGMA2[end,:])
ax2[:plot](n, zeros(size(n, 1)))
ax2[:set_xlabel]("sample size")
ax2[:set_ylabel](L"$s^2-\hat \sigma^2$")
plt[:tight_layout]()
show()


using ForwardDiff
function cournot_newton(eta, c1, c2, guess, tolerance) # note that guess is a 1*2 array
    profit_1(guess) = (sum(guess))^(-1/eta)*guess[1] - 0.5*c1*guess[1]^2
    profit_2(guess) = (sum(guess))^(-1/eta)*guess[2] - 0.5*c2*guess[2]^2

    diff = 1e3
    while diff > tolerance #here diff is the distance matrix
        println("Now the guess for q1 is $(guess[1]), the guess for q2 is $(guess[2]).")
        f1 = guess -> ForwardDiff.gradient(profit_1,guess) # calculate FOC w.r.t q1
        f2 = guess -> ForwardDiff.gradient(profit_2,guess) # calculate FOC w.r.t q2
        f1_prime = guess -> ForwardDiff.hessian(profit_1,guess) # calculate SOC w.r.t q1
        f2_prime = guess -> ForwardDiff.hessian(profit_2,guess) # calculate SOC w.r.t q2
        new_guess_q1 = guess[1] - f1(guess)[1] / f1_prime(guess)[1,1]
        new_guess_q2 = guess[2] - f2(guess)[2] / f2_prime(guess)[2,2]
        diff = sqrt((guess[1]-new_guess_q1)^2 + (guess[2]-new_guess_q2)^2) # eclidean distance from old guess
        guess = [new_guess_q1 new_guess_q2] # form new guess
    end
    println("The optimal quantity for firm 1 is $(guess[1]), for firm 2 is $(guess[2])")
end


#=

A Pythagorean triplet is a set of three natural numbers, a < b < c, for which,
a2 + b2 = c2

For example, 3^2 + 4^2 = 9 + 16 = 25 = 5^2.

There exists exactly one Pythagorean triplet for which a + b + c = 1000.
Find the product abc.

=#

p = 12

function pythagorean_triplet(perimeter)
    for b in 2:Int(ceil((perimeter-1)/2))
        for a in 1:b-1
            if 2*perimeter*(a+b) == perimeter^2 + 2*a*b
                #println(a, ", ", b, ", ", perimeter-a-b)
                return a, b, perimeter-a-b
            end
            #println(a, ", ", b, ", ", perimeter-a-b)
        end
    end
end

res = pythagorean_triplet(1000)
prod(res)


function mandelbrot(a)
    z = 0
    for i=1:50
        z = z^2 + a
    end
    return z
end

for y=1.0:-0.05:-1.0
    for x=-2.0:0.0315:0.5
        abs(mandelbrot(complex(x, y))) < 2 ? print("*") : print(" ")
    end
    println()
end

# Taken from: https://rosettacode.org/wiki/Mandelbrot_set#Julia


using DrWatson
@quickactivate "StatReth"

# %% 3.1
Pr_Positive_Vampire = 0.95
Pr_Positive_Mortal = 0.01
Pr_Vampire = 0.001
Pr_Positive = Pr_Positive_Vampire * Pr_Vampire + Pr_Positive_Mortal * (1 - Pr_Vampire)
Pr_Vampire_Positive = Pr_Positive_Vampire * Pr_Vampire / Pr_Positive


export NumberTheory

module NumberTheory

export Z_star, Eulerφ, continued_fraction, mod_inverse, rand_primeto, factor_a_power_b
export is_order, order_from_float, find_order

"""
    Z_star(N::Int) -> Vector

returns the Z* group elements of `N`, i.e. {x | gcd(x, N) == 1}
"""
Z_star(N::Int) = filter(i->gcd(i, N)==1, 0:N-1)
Eulerφ(N) = length(Z_star(N))

"""
    continued_fraction(ϕ, niter::Int) -> Rational

obtain `s` and `r` from `ϕ` that satisfies `|s/r - ϕ| ≦ 1/2r²`
"""
continued_fraction(ϕ, niter::Int) = niter==0 ? floor(Int, ϕ) : floor(Int, ϕ) + 1//continued_fraction(1/mod(ϕ, 1), niter-1)
continued_fraction(ϕ::Rational, niter::Int) = niter==0 || ϕ.den==1 ? floor(Int, ϕ) : floor(Int, ϕ) + 1//continued_fraction(1/mod(ϕ, 1), niter-1)

"""
    mod_inverse(x::Int, N::Int) -> Int

Return `y` that `(x*y)%N == 1`, notice the `(x*y)%N` operations in Z* forms a group and this is the definition of inverse.
"""
function mod_inverse(x::Int, N::Int)
    for i=1:N
        (x*i)%N == 1 && return i
    end
    throw(ArgumentError("Can not find the inverse, $x is probably not in Z*($N)!"))
end

"""
    is_order(r, x, N) -> Bool

Returns true if `r` is the order of `x`, i.e. `r` satisfies `x^r % N == 1`.
"""
is_order(r, x, N) = powermod(x, r, N) == 1

"""
    find_order(x::Int, N::Int) -> Int

Find the order of `x` by brute force search.
"""
function find_order(x::Int, N::Int)
    findfirst(r->is_order(r, x, N), 1:N)
end

"""
    rand_primeto(N::Int) -> Int

Returns a random number `2 ≦ x < N` that is prime to `N`.
"""
function rand_primeto(N::Int)
    while true
        x = rand(2:N-1)
        d = gcd(x, N)
        if d == 1
            return x
        end
    end
end

"""
    order_from_float(ϕ, x, L) -> Int

Estimate the order of `x` to `L`, `r`, from a floating point number `ϕ ∼ s/r` using the continued fraction method.
"""
function order_from_float(ϕ, x, L)
    k = 1
    rnum = continued_fraction(ϕ, k)
    while rnum.den < L
        r = rnum.den
        if is_order(r, x, L)
            return r
        end
        k += 1
        rnum = continued_fraction(ϕ, k)
    end
    return nothing
end

"""
    factor_a_power_b(N::Int) -> (Int, Int) or nothing

Factorize `N` into the power form `a^b`.
"""
function factor_a_power_b(N::Int)
    y = log2(N)
    for b = 2:ceil(Int, y)
        x = 2^(y/b)
        u1 = floor(Int, x)
        u1^b == N && return (u1, b)
        (u1+1)^b == N && return (u1+1, b)
    end
end

end


using LinearAlgebra, Optim

export tracenorm, tracedist, frobdist, isunitary, ishermitian_tol,
        isspecial, offdiag, jacobi_simultaneous_diag

ϵ = sqrt(eps(Float64))

"""
Compute the trace norm of a matrix ``A``
"""
function tracenorm(A::Matrix{T}) where T<:Number
    #Compute the trace norm of a matrix A
    if ishermitian(A)
        return sum(abs.(eigvals(A)))
    else
        return sum(svd(A).S)
    end
end

"""
Compute the trace distance between matrices ``A``, ``B``.
"""
function tracedist(A::Matrix{T}, B::Matrix{T}) where T<:Number
    #Compute trace distance between matrices A, B
    return 0.5*tracenorm(A-B)
end

"""
Compute the Frobenius distance between matrices ``A``, ``B``.
"""
function frobdist(A::Matrix{T}, B::Matrix{T}) where T<:Number
    return norm(A-B, p=2)
end

"""
Check if a matrix is unitary up to a tolerance ``\epsilon``.
"""
function isunitary(A::Matrix{T}; tol=ϵ) where T <: Number
    return tracenorm(A'A - I) < tol
end

"""
Check if a matrix is Hermitian up to a tolerance ``\epsilon``.
"""
function ishermitian_tol(A::Matrix{T}; tol=ϵ) where T <: Number
    return tracedist(A, Matrix(A')) < tol
end

"""
Check if a matrix is special up to a tolerance ``\epsilon``.
"""
function isspecial(A::Matrix{T}; tol=ϵ) where T <: Number
    return isapprox(det(A), 1.0, rtol=tol)
end

"""
Sum of the off-diagonal elements of a matrix.
"""
function offdiag(A::Matrix{T}) where T <: Number
    return sum(tril(A,-1).^2) + sum(triu(A,1).^2)
end

"""
Function which implements the Jacobi algorithm for simultaneous (possibly approxiamate)
diagonalization of two real, symmetric matrices A, B.

Algorithm is from:

V. Kuleshov, A. Chaganty, P. Liang. "Simultaneous diagonalization: the asymmetric,
low-rank, and noisy settings" arXiv: 1051.06318 (2015)

J.-F. Cardoso, A. Soulomiac. "Jacobi angles for simultaneous diagonalization"
SIAM J. Matrix Anal. Appl. 17(1), 161-164. (1996)

"""
function jacobi_simultaneous_diag(A::Matrix{T}, B::Matrix{T};
        tol=sqrt(eps(T)), maxiter=1000) where T<:Real

    @assert ishermitian_tol(A) && ishermitian_tol(B) "A, B must be symmetric!"

    (m,mm) = size(A)
    done = false

    U = Float64.(Matrix(I, (m,mm)))
    Λa = A
    Λb = B

    #Cost function for minimization of "off-diagonal-ness"
    #TODO: I was not able to succesfully implement the exact solution for the minimizer from the
    #Cardoso paper...

    cost(X,Y) = offdiag(X)+offdiag(Y)

    start_cost = cost(A,B)

    function minims_lsq(X, Y, i, j)
        function f(x)
            g = LinearAlgebra.Givens(i,j, cos(x), sin(x))
            return cost(g'*X*g, g'*Y*g)
        end
        xopt = Optim.minimizer(optimize(f, -π/4, π/4)) #are these the correct limits?
        return (cos(xopt), sin(xopt))
    end

    target = tol*(norm(A,2)+norm(B,2))

    n = 1
    while ((cost(Λa,Λb) > tol) && (n < maxiter))
        for j=1:m-1
            for k=j+1:m
                (c, s) = minims_lsq(Λa, Λb, j, k)
                Γ = LinearAlgebra.Givens(j, k, c, s)
                U = U*Γ
                Λa = Γ'*Λa*Γ
                Λb = Γ'*Λb*Γ
            end
        end
    end

    final_cost = cost(Λa,Λb)

    success = true

    if ((n >= maxiter) || (final_cost > tol))
        @warn "Was unable to find a simultaneous diagonalization: start
                cost function: $(start_cost), final: $(final_cost)."
        success = false
    end

    return (U, Diagonal(Λa), Diagonal(Λb), success)
end


using SymPy

function ry_mat(ang)
    return [SymPy.cos(ang/2) -SymPy.sin(ang/2); SymPy.sin(ang/2) SymPy.cos(ang/2)]
end

function rz_mat(ang)
    return [SymPy.exp(-1im*ang/2) 0; 0 SymPy.exp(1im*ang/2)]
end

function gen_mat(a,b,c,d)
    return SymPy.exp(1im*d).*ry_mat(c)*rz_mat(b)*ry_mat(a)
end

function _params_zyz(mat::Matrix{<:Symbol})
    """Return the euler angles and phase for the ZYZ basis."""
    # We rescale the input matrix to be special unitary (det(U) = 1)
    # This ensures that the quaternion representation is real
    coeff = sympy.sqrt(sympy.det(mat))
    phase = -sympy.atan2(real(coeff),imag(coeff))
    su_mat = coeff * mat  # U in SU(2)
    # OpenQASM SU(2) parameterization:
    # U[0, 0] = exp(-i(phi+lambda)/2) * cos(theta/2)
    # U[0, 1] = -exp(-i(phi-lambda)/2) * sin(theta/2)
    # U[1, 0] = exp(i(phi-lambda)/2) * sin(theta/2)
    # U[1, 1] = exp(i(phi+lambda)/2) * cos(theta/2)
    theta = 2 * sympy.atan2(abs(su_mat[2, 1]), abs(su_mat[1, 1]))
    phiplambda = 2 * sympy.atan2(real(su_mat[2, 2]), imag(su_mat[2, 2]))
    phimlambda = 2 * sympy.atan2(real(su_mat[2, 1]), imag(su_mat[2, 1]))
    phi = (phiplambda + phimlambda) / 2
    lam = (phiplambda - phimlambda) / 2
    return theta, phi, lam, phase
end

function _params_zyz(mat::Matrix{<:Number})
    """Return the euler angles and phase for the ZYZ basis."""
    # We rescale the input matrix to be special unitary (det(U) = 1)
    # This ensures that the quaternion representation is real
    coeff = sqrt(det(mat) +0im)
    phase = -atan(real(coeff), imag(coeff))
    su_mat = coeff * mat  # U in SU(2)
    # OpenQASM SU(2) parameterization:
    # U[0, 0] = exp(-i(phi+lambda)/2) * cos(theta/2)
    # U[0, 1] = -exp(-i(phi-lambda)/2) * sin(theta/2)
    # U[1, 0] = exp(i(phi-lambda)/2) * sin(theta/2)
    # U[1, 1] = exp(i(phi+lambda)/2) * cos(theta/2)
    theta = 2 * atan(abs(su_mat[2, 1]), abs(su_mat[1, 1]))
    phiplambda = 2 * atan(real(su_mat[2, 2]), imag(su_mat[2, 2]))
    phimlambda = 2 * atan(real(su_mat[2, 1]), imag(su_mat[2, 1]))
    phi = (phiplambda + phimlambda) / 2
    lam = (phiplambda - phimlambda) / 2
    return theta, phi, lam, phase
end


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

"""
    variation(table)

Return the variation matrix `Τ` of the `table` such that:

- `Τ[i,j] = Var(log(x[i]/x[j]))` for `i, j = 1, ..., D`
"""
function variation(table)
  X = Tables.matrix(table)
  n = Tables.columnnames(table) |> collect
  D = size(X, 2)
  L = log.(X .+ eps())

  T = Matrix{Float64}(undef, D, D)
  for j in 1:D
    for i in j+1:D
      lr = L[:,i] .- L[:,j]
      T[i,j] = var(lr)
    end
    T[j,j] = 0.0
    for i in 1:j-1
      T[i,j] = T[j,i]
    end
  end

  AxisArray(T, row=n, col=n)
end

"""
    alrcov(table)

Return the log-ratio covariance matrix `Σ` of the `table` such that:

- `Σ[i,j] = cov(log(x[i]/x[D]), log(x[j]/x[D]))` for `i, j = 1, ..., d`
"""
function alrcov(table)
  alrtable = table |> ALR()

  Σ = cov(Tables.matrix(alrtable), dims=1)

  vars = Tables.columnnames(alrtable) |> collect
  AxisArray(Σ, row=vars, col=vars)
end

"""
    clrcov(table)

Return the centered log-ratio covariance matrix `Γ` of the `table` such that:

- `Γ[i,j] = cov(log(x[i]/g(x)), log(x[j]/g(x)))` for `i, j = 1, ..., D`,
where `g(x)` is the geometric mean.
"""
function clrcov(table)
  clrtable = table |> CLR()

  Γ = cov(Tables.matrix(clrtable), dims=1)

  vars = Tables.columnnames(clrtable) |> collect
  AxisArray(Γ, row=vars, col=vars)
end

"""
    lrarray(table)

Return the variation array `A` of the `table` such that:

- `A[i,j] = E[log(x[i]/x[j])]` for `i > j`
- `A[i,j] = Var(log(x[i]/x[j]))` for `i < j`
- `A[i,j] = 0` for `i = j`
"""
function lrarray(table)
  X = Tables.matrix(table)
  D = size(X, 2)
  L = log.(X .+ eps())

  A = Matrix{Float64}(undef, D, D)
  for i in 1:D
    for j in i+1:D
      lr = L[:,i] .- L[:,j]
      A[j,i] = mean(lr)
      A[i,j] = var(lr)
    end
    A[i,i] = 0.0
  end

  vars = Tables.columnnames(table) |> collect
  AxisArray(A, row=vars, col=vars)
end


using HypothesisTests
using Test
using DelimitedFiles
using StatsBase

@testset "Correlation" begin
    # Columns are line number, calcium, iron
    nutrient = readdlm(joinpath(@__DIR__, "data", "nutrient.txt"))[:, 1:3]
    w = CorrelationTest(nutrient[:,2], nutrient[:,3])
    let out = sprint(show, w)
        @test occursin("reject h_0", out) && !occursin("fail to", out)
    end
    let ci = confint(w)
        @test first(ci) ≈ 0.3327138 atol=1e-6
        @test last(ci) ≈ 0.4546639 atol=1e-6
    end
    @test nobs(w) == 737
    @test dof(w) == 735
    @test pvalue(w) < 1e-25

    x = CorrelationTest(nutrient[:,1], nutrient[:,2])
    @test occursin("fail to reject", sprint(show, x))
    let ci = confint(x)
        @test first(ci) ≈ -0.1105478 atol=1e-6
        @test last(ci) ≈ 0.0336730 atol=1e-6
    end
    @test pvalue(x) ≈ 0.2948405 atol=1e-6
end

@testset "Partial correlation" begin
    # Columns are information, similarities, arithmetic, picture completion
    wechsler = readdlm(joinpath(@__DIR__, "data", "wechsler.txt"))[:,2:end]
    w = CorrelationTest(wechsler[:,1], wechsler[:,2], wechsler[:,3:4])
    let out = sprint(show, w)
        @test occursin("reject h_0", out) && !occursin("fail to", out)
    end
    let ci = confint(w)
        @test first(ci) ≈ 0.4963917 atol=1e-6
        @test last(ci) ≈ 0.8447292 atol=1e-6
    end
    @test nobs(w) == 37
    @test dof(w) == 33
    @test pvalue(w) < 0.00001

    X = [ 2 1 0
          4 2 0
         15 3 1
         20 4 1]
    x = CorrelationTest(view(X,:,1), view(X,:,2), view(X,:,3))
    @test occursin("fail to reject", sprint(show, x))
    @test confint(x) == (-1.0, 1.0)
    @test nobs(x) == 4
    @test dof(x) == 1
    @test pvalue(x) ≈ 0.25776212 atol=1e-6
end


using DifferentialEquations

function parameterized_lorenz(du,u,p,t)
  x,y,z = u
  σ,ρ,β = p
  du[1] = dx = σ*(y-x)
  du[2] = dy = x*(ρ-z) - y
  du[3] = dz = x*y - β*z
end

using ParameterizedFunctions
g = @ode_def begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ ρ β

# using LinearAlgebra
# f(du,u,p,t) = mul!(du,A,u)

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
p = [10.0,28.0,8/3]
prob = ODEProblem(g,u0,tspan,p)

using StaticArrays, DifferentialEquations
A  = @SMatrix [ 1.0  0.0 0.0 -5.0
                4.0 -2.0 4.0 -3.0
               -4.0  0.0 0.0  1.0
                5.0 -2.0 2.0  3.0]
u0 = @SMatrix rand(4,2)
tspan = (0.0,1.0)
f(u,p,t) = A*u
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)
using Plots;
plot(sol)

# const κ = 1; const tau = 2 * pi
const out = zeros(2) # Define a cache variable
function delay_mathieu_model(du,u,h,p,t)
  δ, ϵ, b, κ = p
  h(out, p, t-tau)
  du[1] = - κ *u[1] - (δ+ϵ*cos(t)) *u[2] -b * out[2]
  du[2] = u[1]
end

p = [1.0,2.0,1,0.1]
tau=2*pi
T=2*pi
lags = [tau]
h(out, p, t) = (out.=1.0)
# h(out, p, t, deriv::Type{Val{i}})
tspan = (0.0,T)
u0 = [1.0,1.0]
prob = DDEProblem(delay_mathieu_model,h,tspan,p; constant_lags=lags)

alg = MethodOfSteps(Tsit5())
# alg = MethodOfSteps(BS3())
# alg = MethodOfSteps(Vern6())
sol = solve(prob,alg)#,reltol=1e-8,abstol=1e-8)

using Plots
plot(sol,vars=(0,2))

# plot(sol,linewidth=5,title="Solution to the linear ODE with a thick line",
#      xaxis="Time (t)",yaxis="u(t) (in μm)",label="My Thick Line!") # legend=false
# plot!(sol.t, t->0.5*exp(1.01t),lw=3,ls=:dash,label="True Solution!")
sol[2][3]
plot(sol,vars=(0,2))


n = 10 #rand(1:10)
m = n ÷ 2 # rand(1:n)
@test matrixdepot("vand", Int, n) == matrixdepot("vand", [1:n;])

A = matrixdepot("vand", Int, n)
@test A[:, m] == [1:n;].^(m-1)
println("'vand' passed test...")


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module BinaryInteger

export ModuleBinaryInteger
export V001855, V003314, V033156, V054248, V061168, V083652, V097383, V123753
export V295513, BinaryIntegerLength, Bil

"""
For positive n, BinaryIntegerLength is ``⌊ log[2](n) ⌋ + 1``, BinaryIntegerLength(0) = 0.

* BinaryIntegerLength, Bil, V001855, V003314, V033156, V054248, V061168, V083652, V097383, V123753, V295513
"""
const ModuleBinaryInteger = ""

"""
Return the length of the binary extension of an integer ``n``, which is defined as ``0`` if ``n = 0`` and for ``n > 0`` as ``⌊ log[2](n) ⌋ + 1``.
"""
BinaryIntegerLength(n) = n == 0 ? 0 : floor(Int, log2(n)) + 1

"""
Alias for the function BinaryIntegerLength.
"""
Bil(n) = BinaryIntegerLength(n)

"""
Return ``n`` Bil``(n) - 2^{\\text{Bil}(n)}`` where Bil``(n)`` is the binary integer length of ``n``.
"""
V295513(n) = n*Bil(n) - 2^Bil(n)

"""
Maximal number of comparisons for sorting ``n`` elements by binary insertion.
"""
V001855(n) = V295513(n) + 1

"""
Return the sum of lengths of binary expansions of ``0`` through ``n``.
"""
V083652(n) = V295513(n+1) + 2

"""
Recurrence ``a(n) = a(n-1) + ⌊ a(n-1)/(n-1) ⌋ + 2`` for ``m ≥ 2`` and ``a(1) = 1``.
"""
V033156(n) = V295513(n) + 2n

"""
Binary entropy function: ``a(n) = n + `` min ``( a(k) + a(n-k) : 1 ≤ k ≤ n-1 )`` for ``n > 1,`` and ``a(1) = 0``.
"""
V003314(n) = V295513(n) + n

"""
Binary entropy: ``a(n) = n +`` min ``{ a(k) + a(n-k) : 1 ≤ k ≤ n-1 }.``
"""
V054248(n) = V295513(n) + n + rem(n, 2)

"""
Minimum total number of comparisons to find each of the values ``1`` through ``n`` using a binary search with ``3``-way comparisons.
"""
V097383(n) = V295513(n+1) - div(n-1, 2)

"""
Partial sums of the sequence ``⌊ log[2](n) ⌋``.
"""
V061168(n) = V295513(n+1) - n + 1

"""
Partial sums of the sequence of length of the binary expansion of ``2n+1``.
"""
V123753(n) = V295513(n+1) + n + 2

#START-TEST-########################################################

using Test, SeqTests

function test()
    @testset "BinaryInteger" begin

        @test V295513(0) == -1
        @test V295513(1) == -1
        @test V295513(2) == 0
        @test V295513(3) == 2

        if is_oeis_installed()
            V = [V001855, V003314, V033156, V054248, V061168, V097383]
            for v in V SeqTest(v , 'V', 1) end

            V = [V083652, V123753]
            for v in V SeqTest(v , 'V', 0) end
        end
    end
end

function demo()
    println([V295513(n) for n in 0:12])
    println([V123753(n) for n in 0:12])
    println([V001855(n) for n in 1:12])
    println([V083652(n) for n in 1:12])
    println([V033156(n) for n in 1:12])
    println([V003314(n) for n in 1:12])
    println([V054248(n) for n in 1:12])
    println([V097383(n) for n in 1:12])
    println([V061168(n) for n in 1:12])
end

"""
[V295513(k) for k in 0:100000]
    0.014412 seconds (2 allocations: 781.391 KiB)
"""
function perf()
    @time [V295513(k) for k in 0:100000]
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


# this function shows the behavior of the
# likelihood, the log-likelihood and the average
# log-likelihood as the sample size increases
using Distributions
n = 10 # try increasing this to 50, 1000, etc.
p = 0.3
y = rand(Bernoulli(p),n)
f =  pdf.(Bernoulli(p), y)
@show n
@show ℒ = prod(f)
@show lnℒ = sum(log.(f))
@show avglnℒ = (1/n)*sum(log.(f))
nothing


using LPVSpectral, Statistics, Random
export Sinusoidal

"""
    Sinusoidal(fs)
Decompose a timeseries `s` into a **sum** `x + r`, where `x`
are sinusoidal components with the given frequencies `fs` that minimize
coefficients ``A, \\phi`` of the expression
```math
s \\approx A_0 + \\sum_i A_i \\cos(2\\pi f_i t + \\phi_i)
```
with ``A_0 = \\bar{s}`` the mean.

This method uses a new least-squares algorithm in frequency domain using the package
[LPVSpectral.jl](https://github.com/baggepinnen/LPVSpectral.jl), see[^Bagge2017].
It works for non-equispaced `t` axis (and also normal), is generally very accurate
(if choosen frequencies are not too close), but has performance scaling of
O(N^2.4) instead of O(n log(n)) of [`Fourier`](@ref).

Because it can work with arbitrary signal length the method always estimates the
zero-frequency Fourier component, and attributes it to `x`.
The fitted coefficients ``A, \\phi`` are available as fields `.A` and `.φ` of the struct
(first entry is zero-frequency component ``A_0, \\phi_0``).

[^Bagge2017]: F. Bagge Carlson et al., [Linear Parameter-Varying Spectral Decomposition](https://lup.lub.lu.se/search/publication/ac32368e-e199-44ff-b76a-36668ac7d595).
"""
struct Sinusoidal{F<:AbstractArray{<:AbstractFloat}} <: Decomposition
    fs::F
    A::Vector{Float64}
    φ::Vector{Float64}
end
Sinusoidal(fs) = Sinusoidal(fs, zeros(length(fs)+1), zeros(length(fs)+1))

function decompose(t, s, method::Sinusoidal)
    @assert 0 ∉ method.fs
    fs = [0, method.fs...]
    χ, = ls_spectral(s, t, fs)
    As = abs.(χ) ./ sqrt(2*length(fs))
    φs = angle.(χ)
    method.A .= As
    method.φ .= φs
    x = fill(As[1], length(s))
    for i in 2:length(fs)
        @. x += As[i]*cos(2π*fs[i]*t + φs[i])
    end
    return x, s .- x
end


#!/usr/bin/julia

# Trizen
# 15 November 2021
# https://github.com/trizen

# Compute the multiplicative order of `a` modulo `n`: znorder(a, n).
# This is the smallest positive integer k such that a^k == 1 (mod n).

using Primes

function divisors(n)

    d = Int64[1]

    for (p,e) in factor(n)
        t = Int64[]
        r = 1

        for i in 1:e
            r *= p
            for u in d
                push!(t, u*r)
            end
        end

        append!(d, t)
    end

    return sort(d)
end

function znorder(a, n)

    if isprime(n)
        for d in divisors(n-1)
            if (powermod(a, d, n) == 1)
                return d
            end
        end
    end

    f = factor(n)

    if (length(f) == 1)         # is prime power

        p = first(first(f))
        z = znorder(a, p)

        while (powermod(a, z, n) != 1)
            z *= p
        end

        return z
    end

    pp_orders = Int64[]

    for (p,e) in f
        push!(pp_orders, znorder(a, p^e))
    end

    return lcm(pp_orders)
end

isequal(znorder(97, factorial(14)), 25920)          || print("error")
isequal(znorder(53, factorial(15)), 2419200)        || print("error")
isequal(znorder(37, factorial(16)), 116121600)      || print("error")
isequal(znorder(31, factorial(17)), 6220800)        || print("error")
isequal(znorder(89, factorial(18)), 1045094400)     || print("error")
isequal(znorder(101,factorial(19)), 1254113280)     || print("error")
isequal(znorder(97, factorial(20)), 2239488000)     || print("error")

println(znorder(2, 341))        #> 10
println(znorder(97, 5040))      #> 12


# Polynomial kernel function

"""
    Poly <: Kernel

Polynomial kernel (covariance)
```math
k(x,x') = σ²(xᵀx' + c)ᵈ
```
with signal standard deviation ``σ``, additive constant ``c``, and degree ``d``.
"""
mutable struct Poly{T<:Real} <: Kernel
    "Constant"
    c::T
    "Signal variance"
    σ2::T
    "Degree of polynomial"
    deg::Int
    "Priors for kernel parameters"
    priors::Array
end

"""
Polynomial kernel function

    Poly(lc::Real, lσ::Real, deg::Int)

# Arguments
  - `lc::Real`: additive constant (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
  - `deg::Int`: degree of polynomial
"""
Poly(lc::T, lσ::T, deg::Int) where T = Poly{T}(exp(lc), exp(2 * lσ), deg, [])

function KernelData(k::Poly, X1::AbstractMatrix, X2::AbstractMatrix)
    XtX=X1'*X2
	if X1==X2
    	LinearAlgebra.copytri!(XtX, 'U')
	end
    LinIsoData(XtX)
end
kernel_data_key(k::Poly, X1::AbstractMatrix, X2::AbstractMatrix) = "LinIsoData"

_cov(poly::Poly, xTy) = poly.σ2*(poly.c.+xTy).^poly.deg
function cov(poly::Poly, x::AbstractVector, y::AbstractVector)
    K = _cov(poly, dot(x,y))
end
function cov!(cK::AbstractMatrix, poly::Poly, X::AbstractMatrix, data::LinIsoData)
    cK .= _cov(poly, data.XtX)
end
@inline @inbounds function cov_ij(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, i::Int, j::Int, dim::Int)
    return _cov(poly, dotij(X1, X2, i, j, dim))
end
@inline @inbounds function cov_ij(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, data::LinIsoData, i::Int, j::Int, dim::Int)
    return _cov(poly, data.XtX[i, j])
end
cov(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, data::LinIsoData) = _cov(poly, data.XtX)

get_params(poly::Poly) = [log(poly.c), log(poly.σ2) / 2]
get_param_names(poly::Poly) = [:lc, :lσ]
num_params(poly::Poly) = 2

function set_params!(poly::Poly, hyp::AbstractVector)
    length(hyp) == 2 || throw(ArgumentError("Polynomial function has two parameters"))
    poly.c = exp(hyp[1])
    poly.σ2 = exp(2 * hyp[2])
end

@inline dk_dlc(poly::Poly, xTy::Real) = poly.c*poly.deg*poly.σ2*(poly.c+xTy).^(poly.deg-1)
@inline dk_dlσ(poly::Poly, xTy::Real) = 2 * _cov(poly,xTy)
@inline function dKij_dθp(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, i::Int, j::Int, p::Int, dim::Int)
    if p==1
        return dk_dlc(poly, dotij(X1,X2,i,j,dim))
    else
        return dk_dlσ(poly, dotij(X1, X2,i,j,dim))
    end
end
@inline function dKij_dθp(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, data::LinIsoData, i::Int, j::Int, p::Int, dim::Int)
    if p==1
        return dk_dlc(poly, data.XtX[i,j])
    else
        return dk_dlσ(poly, data.XtX[i,j])
    end
end


# =========================================================================== #
# Error Functions
# =========================================================================== #

abstract type AbstractErrorFunction
end

struct LeastSquaresErrorFunction <: AbstractErrorFunction
end

function (error_fcn::LeastSquaresErrorFunction)(meas::T, x::T) where {T <: Number}
    dx = -meas + x
    return dx^2
end
function (error_fcn::LeastSquaresErrorFunction)(meas::Vector{T}, x::Vector{T}) where {T}
    dx = sum( (-meas .+ x).^2 )
    return dx
end
function (error_fcn::LeastSquaresErrorFunction)(meas, x::Vector{T}) where {T}
    dx = sum( (-meas .+ x).^2 )
    return dx
end

function differentiate(error_fcn::LeastSquaresErrorFunction, meas, x)
    return -meas + x
end

function eval(error_fcn::LeastSquaresErrorFunction, meas, x)
    return error_fcn(meas, x)
end

function gradient!(error_fcn::LeastSquaresErrorFunction, y_meas, y_out, delta_e)
    for kk = 1:length(delta_e)
        delta_e[kk] = differentiate(error_fcn, y_meas[kk], y_out[kk])
    end
    return delta_e
end

# function error_gradient!(nn::NeuralNetwork, y_meas, y_out, delta_e)
#     return gradient!(nn.error_function, y_meas, y_out, delta_e)
# end