mutable struct Citizen <: Agents.AbstractAgent
    id::Int64
    pos::Int64
    home::Int64
    neuroticism::Float64
    trust_authorities::Float64
    fear::Float64
    socialnorm::Float64
    socialnorm_memory::Array{Any, 1}
    behavior::Bool
    behavior_buffer::Bool
    quarantined::Bool
    state::Symbol
    ticks_exposed::Int64
    ticks_infected::Int64
    ticks_quarantined::Int64
    incubation_period::Int64
    infection_duration::Int64
end

Citizen(;
    id::Int64, pos::Int64, home::Int64,
    neuroticism::Float64, trust_authorities::Float64, fear::Float64,
    socialnorm::Float64, socialnorm_memory::Array{Float64, 1},
    behavior::Bool, behavior_buffer::Bool,
    quarantined::Bool, state::Symbol,
    ticks_exposed::Int64, ticks_infected::Int64, ticks_quarantined::Int64,
    incubation_period::Int64, infection_duration::Int64
) = Citizen(
    id, pos, home,
    neuroticism, trust_authorities, fear,
    socialnorm, socialnorm_memory,
    behavior, behavior_buffer,
    quarantined, state,
    ticks_exposed, ticks_infected, ticks_quarantined,
    incubation_period, infection_duration
)

Citizen(id, pos, state) = Citizen(
    id=id, pos=pos, home=pos,
    neuroticism=randn_unit(), trust_authorities=randn_unit(), fear=0.0,
    socialnorm=0.0, socialnorm_memory=Float64[],
    behavior=false, behavior_buffer=false,
    quarantined=false, state=state,
    ticks_exposed=0, ticks_infected=0, ticks_quarantined=0,
    incubation_period=rand_incubation(), infection_duration=rand_duration()
)

function randn_unit()
    normal_distribution = Distributions.Normal(0.5, 0.2)
    random_normal_number = Distributions.rand(normal_distribution)
    if (random_normal_number < 0) | (random_normal_number > 1)
        random_normal_number = randn_unit()
    end
    return random_normal_number
end

function rand_incubation(α=2.14532, θ=1.5626)
    gamma_distribution = Distributions.Gamma(α, θ)
    random_gamma_number = Distributions.rand(gamma_distribution)
    random_incubation = trunc(Int64, round(random_gamma_number * 10))
    return random_incubation
end

function rand_duration(λ=17.5)
    poisson_distribution = Distributions.Poisson(λ)
    random_poisson_number = Distributions.rand(poisson_distribution)
    random_duration = trunc(Int64, round(random_poisson_number * 10))
    return random_duration
end


using Plots
using LinearAlgebra
using Polynomials

using SparseArrays
using Arpack

using ForwardDiff

using BenchmarkTools

function tria6(l1,l2,l3,xs,ys)
    L1 = l1; L2=l2; L3 = l3;
    N1 = (2*L1-1)*L1; N2 = (2*L2-1)*L2; N3 = (2*L3-1)*L3;
    N4 = 4*L1*L2; N5 = 4*L2*L3; N6 = 4*L3*L1;
    x = dot(xs',[N1 N2 N3 N4 N5 N6])
    y = dot(ys',[N1 N2 N3 N4 N5 N6])
    return x,y
end
xs = 10*[-1 0.3 0 -.5 0.2 -.5]
ys = 10*[0 0 1 0.2 .5 .8]
scatter()
for l1 = 0:0.02:1
    for l2 = 0:0.02:1-l1
        scatter!(tria6(l1,l2,1-l1-l2,xs,ys))
    end
end
scatter!()

function tria6jacobian(l1,l2,l3,xs,ys)
    #calculates jacobian matrix dN1/dl1 ....
    L1 = l1; L2=l2; L3=l3;
    N1 = (2*L1-1)*L1; N2 = (2*L2-1)*L2; N3 = (2*L3-1)*L3;
    N4 = 4*L1*L2
    N5 = 4*L2*L3
    N6 = 4*L3*L1
    xy =[xs;ys];
    Ndot = [4*l1-1 0 0;
            0 4*l2-1 0;
            0 0 4*l3-1;
            4*l2 4*l1 0;
            0 4*l2 4*l3;
            4*l3 0 4*l1]
    jacobian = xy*Ndot;
    return jacobian
end
YM = 230e9; #YOUNGS MODULUS
ν = 0.33 # POISON RATIO
d11 = YM*(1-ν)/(1+ν)/(1-2*ν); d12 = ν*d11/(1-ν); d33 = YM/(2*(1+ν));
E = [d11 d12 0;
     d12 d11 0;
     0 0 d33]
function calcBEB(l1,l2,l3,xs,ys)
    L1 = l1; L2= l2; L3 = l3;
    jacobian = tria6jacobian(l1,l2,l3,xs,ys)
    JP = [1 1 1; jacobian] #according to p368 of Fellipa C. Colorado isop
    P = JP\[0 0; 1 0; 0 1]
    Ndot = [4*l1-1 0 0;
            0 4*l2-1 0;
            0 0 4*l3-1;
            4*l2 4*l1 0;
            0 4*l2 4*l3;
            4*l3 0 4*l1]
    Nderivxy = P'*Ndot'; # is [dN1/dx ... dN6/dx;
                        #     dN1/dy ... dN6/dy]
    N1 = (2*L1-1)*L1; N2 = (2*L2-1)*L2; N3 = (2*L3-1)*L3;
    N4 = 4*L1*L2
    N5 = 4*L2*L3
    N6 = 4*L3*L1
    N = [N1 0 N2 0 N3 0 N4 0 N5 0 N6 0;
         0 N1 0 N2 0 N3 0 N4 0 N5 0 N6]
    N1x = Nderivxy[1,1]; N2x = Nderivxy[1,2]; N3x = Nderivxy[1,3]; N4x = Nderivxy[1,4]; N5x = Nderivxy[1,5]; N6x = Nderivxy[1,6];
    N1y = Nderivxy[2,1]; N2y = Nderivxy[2,2]; N3y = Nderivxy[2,3]; N4y = Nderivxy[2,4]; N5y = Nderivxy[2,5]; N6y = Nderivxy[2,6];

    B = [N1x 0 N2x 0 N3x 0 N4x 0 N5x 0 N6x 0;
         0  N1y 0 N2y 0 N3y 0 N4y 0 N5y 0 N6y;
         N1y N1x N2y N2x N3y N3x N4y N4x N5y N5x N6y N6x]
    BEB = B'*E*B
    return BEB
end

function tria6Ke(xs,ys)
    Ke  = zeros(12,12)
    g1 = [0.108103, 0.445948, 0.445948]; w1 =0.223382;
    g2 = [0.445948, 0.108103, 0.445948]; w2 =0.223382;
    g3 = [0.445948, 0.445948, 0.108103]; w3 =0.223382;
    g4 = [0.816848, 0.0915762, 0.0915762]; w4 = 0.109952;
    g5 = [0.0915762, 0.816848, 0.0915762]; w5 = 0.109952;
    g6 = [0.0915762, 0.0915762, 0.816848]; w6 = 0.109952;
    g = [[g1] [g2] [g3] [g4] [g5] [g6]]; w = [w1 w2 w3 w4 w5 w6];

    for i=1:6
        Ke = Ke + w[i]*calcBEB(g[i][1],g[i][2],g[i][3],xs,ys)
    end
    return Ke
end

#f = open("C:\\User\\billykon\\Desktop\\Mesh_1.unv","r")

#function edgeToLs(xedge,yedge,TOPOLOGY):


abstract type AbstractActivation end

struct Sigmoid <: AbstractActivation end
struct SoftPlus <: AbstractActivation end
struct Relu <: AbstractActivation end
struct Tanh <: AbstractActivation end
struct Identity <: AbstractActivation end

activate(::Sigmoid, x::Number) = 1 ./ (1 + exp(-x))
activate(::SoftPlus, x::Number) = log(1 + exp(x))
activate(::Relu, x::Number) = max(x, 0)
activate(::Tanh, x::Number) = tanh(x)
activate(::Identity, x::Number) = x

activate(a::AbstractActivation, x::AbstractArray) =
    map(_1->activate(a, _1), x)

activate!(out::AbstractArray, a::AbstractActivation, x::AbstractArray) =
    map!(_1->activate(a, _1), out, x)

for T in InteractiveUtils.subtypes(AbstractActivation)
    @eval (ta::$(T))(x) where T = activate(ta, x)
end

deriv(::Sigmoid, x) = exp(-x) ./((1 + exp(-x)).^2)
deriv(::SoftPlus, x) = exp(x) ./ (1 + exp(x))
deriv(::Relu, x::Number) = x >= 0.0 ? 1.0 : 0.0
deriv(r::Relu, x::AbstractArray) = map(_1->deriv(r, _1), x)
deriv(::Tanh, x) = 1 - tanh(x).^2
deriv(::Identity, x::Number) = one(x)
deriv(::Identity, x::AbstractArray) = ones(x)


using Printf
using LinearAlgebra
using Random

include("../common/PWGrid_v01.jl")
include("../common/wrappers_fft.jl")
include("../common/gen_lattice_pwscf.jl")

include("structure_factor.jl")
include("gen_dr.jl")
include("gen_rho.jl")
include("calc_ewald.jl")

function test_main()
    Ns = [64,64,64]
    LatVecs = gen_lattice_sc(16.0)
    pw_grid = PWGrid(Ns,LatVecs)
    # Atomic positions and nuclear charge
    Xpos = reshape( [0.0, 0.0, 0.0,
                     1.5, 0.0, 0.0], (3,2) )
    println("Xpos =")
    println(Xpos)

    Z = 1.0
    # Calculate structure factor
    Sf = structure_factor( Xpos, pw_grid.G )

    print("sum(Sf)="); println(sum(Sf))

    E_nn = calc_ewald( pw_grid, Xpos, Sf )
    @printf("E_nn = %18.10f\n", E_nn)
end

@time test_main()


"""
    gevfit()

Estimate the GEV parameters by maximum likelihood.

# Implementation

The function uses Nelder-Mead solver implemented in the [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl)
package to find the point where the log-likelihood is maximal.

The GEV parameters can be modeled as function of covariates as follows:
```math
μ = X₁ × β₁,
```
```math
ϕ = X₂ × β₂,
```
```math
ξ = X₃ × β₃.
```

The covariates are standardized before estimating the parameters to help fit the
 model. They are transformed back on their original scales before returning the
 fitted model.

See also [`gevfit`](@ref) for the other methods, [`gevfitpwm`](@ref), [`gevfitbayes`](@ref) and [`BlockMaxima`](@ref).

"""
function gevfit end

"""
    gevfit(y,
        locationcov = Vector{Variable}(),
        logscalecov = Vector{Variable}(),
        shapecov = Vector{Variable}()
        )

Estimate the GEV parameters.

# Arguments

- `y::Vector{<:Real}`: The vector of block maxima.
- `locationcov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the location parameter.
- `logscalecov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the log-scale parameter.
- `shapecov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the shape parameter.
"""
function gevfit(y::Vector{<:Real};
    locationcov::Vector{<:DataItem} = Vector{Variable}(),
    logscalecov::Vector{<:DataItem} = Vector{Variable}(),
    shapecov::Vector{<:DataItem} = Vector{Variable}())::MaximumLikelihoodEVA

    locationcovstd = standardize.(locationcov)
    logscalecovstd = standardize.(logscalecov)
    shapecovstd = standardize.(shapecov)

    model = BlockMaxima(Variable("y", y), locationcov = locationcovstd, logscalecov = logscalecovstd, shapecov = shapecovstd)

    fittedmodel = fit(model)

    return transform(fittedmodel)

end

"""
    gevfit(y,
        initialvalues,
        locationcov = Vector{Variable}(),
        logscalecov = Vector{Variable}(),
        shapecov = Vector{Variable}()
        )

Estimate the GEV parameters.

# Arguments

- `y::Vector{<:Real}`: the vector of block maxima.
- `initialvalues::Vector{<:Real}`: Vector of parameters initial values.
- `locationcov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the location parameter.
- `logscalecov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the log-scale parameter.
- `shapecov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the shape parameter.
"""
function gevfit(y::Vector{<:Real}, initialvalues::Vector{<:Real};
    locationcov::Vector{<:DataItem} = Vector{Variable}(),
    logscalecov::Vector{<:DataItem} = Vector{Variable}(),
    shapecov::Vector{<:DataItem} = Vector{Variable}(),)::MaximumLikelihoodEVA

    model = BlockMaxima(Variable("y", y), locationcov = locationcov, logscalecov = logscalecov, shapecov = shapecov)

    return fit(model, initialvalues)

end

"""
    gevfit(df::DataFrame,
        datacol::Symbol,
        locationcovid = Vector{Symbol}(),
        logscalecovid = Vector{Symbol}(),
        shapecovid = Vector{Symbol}()
        )

Estimate the GEV parameters.

# Arguments
- `df::DataFrame`: The dataframe containing the data.
- `datacol::Symbol`: The symbol of the column of `df` containing the block maxima data.
- `locationcovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the location parameter.
- `logscalecovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the log-scale parameter.
- `shapecovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the shape parameter.
"""
function gevfit(df::DataFrame, datacol::Symbol;
    locationcovid::Vector{Symbol}=Symbol[],
    logscalecovid::Vector{Symbol}=Symbol[],
    shapecovid::Vector{Symbol}=Symbol[])::MaximumLikelihoodEVA

    locationcovstd = standardize.(buildVariables(df, locationcovid))
    logscalecovstd = standardize.(buildVariables(df, logscalecovid))
    shapecovstd = standardize.(buildVariables(df, shapecovid))

    model = BlockMaxima(Variable(string(datacol), df[:, datacol]), locationcov = locationcovstd, logscalecov = logscalecovstd, shapecov = shapecovstd)

    fittedmodel = fit(model)

    return transform(fittedmodel)

end

"""
    gevfit(df::DataFrame,
        datacol::Symbol,
        locationcovid = Vector{Symbol}(),
        logscalecovid = Vector{Symbol}(),
        shapecovid = Vector{Symbol}()
        )

Estimate the GEV parameters.

# Arguments
- `df::DataFrame`: The dataframe containing the data.
- `datacol::Symbol`: The symbol of the column of `df` containing the block maxima data.
- `initialvalues::Vector{<:Real}`: Vector of parameters initial values.
- `locationcovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the location parameter.
- `logscalecovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the log-scale parameter.
- `shapecovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the shape parameter.
"""
function gevfit(df::DataFrame, datacol::Symbol, initialvalues::Vector{<:Real};
    locationcovid::Vector{Symbol}=Symbol[],
    logscalecovid::Vector{Symbol}=Symbol[],
    shapecovid::Vector{Symbol}=Symbol[])::MaximumLikelihoodEVA

    locationcov = buildVariables(df, locationcovid)
    logscalecov = buildVariables(df, logscalecovid)
    shapecov = buildVariables(df, shapecovid)

    model = BlockMaxima(Variable(string(datacol), df[:, datacol]), locationcov = locationcov, logscalecov = logscalecov, shapecov = shapecov)

    return fit(model, initialvalues)

end

"""
    gevfit(model::BlockMaxima, initialvalues::Vector{<:Real})

Estimate the parameters of the `BlockMaxima` model using the given initialvalues.
"""
function gevfit(model::BlockMaxima, initialvalues::Vector{<:Real})::MaximumLikelihoodEVA

    fit(model, initialvalues)

end


using MatterPower
using DoubleExponentialFormulas
using Plots, LaTeXStrings
using Dierckx
using DelimitedFiles, Printf

# %% Specify a redshift
redshift = [0, 0.3, 0.5, 0.7, 1, 1.3, 1.5]
ired = 1

# %% Read in linear and non-linear power spectra computed by CLASS and spline interpolate in k
filename = @sprintf("data/matterpower05_z%1d_pk.dat", ired)
d = readdlm(filename, comments = true)
pklin_class = Spline1D(d[:, 1], d[:, 2])
filename = @sprintf("data/matterpower05_z%1d_pk_nl.dat", ired)
d = readdlm(filename, comments = true)
pknl_class = Spline1D(d[:, 1], d[:, 2])

# %% Define a function to return a linear matter power spectrum (in units of Mpc^3/h^3)
# as a function of the comoving wavenumber, kovh, in units of h/Mpc.
# Here is an example using Einstein & Hu's analytical transfer function

# Cosmological parameters
As, ns, kpivot = 2.097e-9, 0.9652, 0.05
Ωm, ΩΛ, Ωb = 0.315, 0.685, 0.049
Ωk = 1 - Ωm - ΩΛ
h0 = 0.674
ωm, fb = Ωm * h0^2, Ωb / Ωm

# Tabulate linear growth factor as a function of scale factor, a
sol = setup_growth(Ωm, ΩΛ)
a = 1 / (1 + redshift[ired])
D1 = sol(a)[1]

pk(kovh) =
   D1^2 *
   As *
   (kovh * h0 / kpivot)^(ns - 1) *
   (2 * kovh^2 * 2998^2 / 5 / Ωm)^2 *
   t_nowiggle(kovh * h0, ωm, fb)^2 *
   2 *
   π^2 / kovh^3

# %% Alternatively you may read in pre-computed data and define a spline function
# using Dierckx
# pk = Spline1D(tabulated_kovh, tabulated_power_spectrum)

# %% Compute the non-linear power spectrum using halofit
# Get three parameters [kσ, neff, C] needed for the halofit transform
p = kσ, neff, C = setup_halofit(pk)

# Calculate Ωm(z) = Ωm(z=0)(1+z)^3 / E^2(z), where E^2(z) = H^2(z) / H0^2
z = redshift[ired]
E2 = Ωm * (1 + z)^3 + Ωk * (1 + z)^2 + ΩΛ
Ωmz = Ωm * (1 + z)^3 / E2

# Define a function to return a halofit non-linear power spectrum
pknl(kovh) = halofit(pk, p, Ωmz, kovh) # Mpc^3/h^3

# %% Compute correlation functions
kmin, kmax = 1e-4, Inf
R = Float64.(10:2:200)
ξ1 = zeros(length(R))
ξ2 = zeros(length(R))
ξ3 = zeros(length(R))
ξ4 = zeros(length(R))
for i = 1:length(R)
   func1(k) = k^2 * pk(k) * sin(k * R[i]) / (k * R[i]) / 2 / π^2
   ξ1[i], err = quaddeo(func1, R[i], 0, kmin, kmax)
   func2(k) = k^2 * pknl(k) * sin(k * R[i]) / (k * R[i]) / 2 / π^2
   ξ2[i], err = quaddeo(func2, R[i], 0, kmin, kmax)
   func3(k) = k^2 * pklin_class(k) * sin(k * R[i]) / (k * R[i]) / 2 / π^2
   ξ3[i], err = quaddeo(func3, R[i], 0, kmin, kmax)
   func4(k) = k^2 * pknl_class(k) * sin(k * R[i]) / (k * R[i]) / 2 / π^2
   ξ4[i], err = quaddeo(func4, R[i], 0, kmin, kmax)
end

# %% Plot results and save to "xi.pdf"
p = plot(
   R,
   R .^ 2 .* ξ4,
   m = 2,
   c = :red,
   lab = "Halofit ξ(R) with BAO",
   ylab = L"R^2\xi(R)",
   xlab = L"R~[Mpc/h]",
   title = "Redshift = $z",
)
p =
   plot!(R, R .^ 2 .* ξ2, ls = :dash, c = :red, lab = "Halofit ξ(R) w/o BAO")
p = plot!(R, R .^ 2 .* ξ3, m = 2, c = 1, lab = "Linear ξ(R) with BAO")
p = plot!(R, R .^ 2 .* ξ1, ls = :dash, c = 1, lab = "Linear ξ(R) w/o BAO")
savefig("xi.pdf")
display(p)


using JuMP
using Ipopt
using Random
# for plots
using Gadfly
using DataFrames

Random.seed!(1234);

x = rand(10)
a = rand(10)
b = rand()

projection = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
@variable(projection, x0[1:10])
@objective(projection, Min, sum((x - x0) .* (x - x0))) # We minimize the square of the distance here
@constraint(projection, x0' * a == b)                  # Point must lie on the hyperplane

optimize!(projection)
@show objective_value(projection);
@show value.(x0);

A_1 = rand(10, 10)
A_2 = rand(10, 10)
b_1 = rand(10)
b_2 = rand(10)

polyhedra_distance = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
@variable(polyhedra_distance, x[1:10])                       # Point closest on the first polyhedron
@variable(polyhedra_distance, y[1:10])                       # Point closest on the second polyhedron
@objective(polyhedra_distance, Min, sum((x - y) .* (x - y))) # We minimize the square of the distance here as above
@constraint(polyhedra_distance, A_1 * x .<= b_1)             # Point x must lie on the first polyhedron
@constraint(polyhedra_distance, A_2 * y .<= b_2)             # Point y must lie on the second polyhedron

optimize!(polyhedra_distance)
@show objective_value(polyhedra_distance);

fixed = [ 1   1  -1 -1    1   -1  -0.2  0.1;         # coordinates of fixed points
          1  -1  -1  1 -0.5 -0.2    -1    1]

M = size(fixed,2)                                    # number of fixed points
N = 6                                                # number of free points

A = [ 1  0  0 -1  0  0    0  0  0  0  0  0  0  0;    # Matrix on links
      1  0 -1  0  0  0    0  0  0  0  0  0  0  0;
      1  0  0  0 -1  0    0  0  0  0  0  0  0  0;
      1  0  0  0  0  0   -1  0  0  0  0  0  0  0;
      1  0  0  0  0  0    0 -1  0  0  0  0  0  0;
      1  0  0  0  0  0    0  0  0  0 -1  0  0  0;
      1  0  0  0  0  0    0  0  0  0  0  0  0 -1;
      0  1 -1  0  0  0    0  0  0  0  0  0  0  0;
      0  1  0 -1  0  0    0  0  0  0  0  0  0  0;
      0  1  0  0  0 -1    0  0  0  0  0  0  0  0;
      0  1  0  0  0  0    0 -1  0  0  0  0  0  0;
      0  1  0  0  0  0    0  0 -1  0  0  0  0  0;
      0  1  0  0  0  0    0  0  0  0  0  0 -1  0;
      0  0  1 -1  0  0    0  0  0  0  0  0  0  0;
      0  0  1  0  0  0    0 -1  0  0  0  0  0  0;
      0  0  1  0  0  0    0  0  0  0 -1  0  0  0;
      0  0  0  1 -1  0    0  0  0  0  0  0  0  0;
      0  0  0  1  0  0    0  0 -1  0  0  0  0  0;
      0  0  0  1  0  0    0  0  0 -1  0  0  0  0;
      0  0  0  1  0  0    0  0  0  0  0 -1  0  0;
      0  0  0  1  0  0    0  0  0  0  0 -1  0  0;
      0  0  0  0  1 -1    0  0  0  0  0  0  0  0;
      0  0  0  0  1  0   -1  0  0  0  0  0  0  0;
      0  0  0  0  1  0    0  0  0 -1  0  0  0  0;
      0  0  0  0  1  0    0  0  0  0  0  0  0 -1;
      0  0  0  0  0  1    0  0 -1  0  0  0  0  0;
      0  0  0  0  0  1    0  0  0  0 -1  0  0  0;]

placement = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
@variable(placement, p[1:M + N, 1:2])                     # A variable array for the coordinates of each point
@constraint(placement, p[N + 1:N + M, :] .== fixed')      # We had a constraint for the fixed points
dist = A * p                                         # Matrix of differences between coordinates of 2 points with a link
@objective(placement, Min, sum(dist .* dist))            # We minimize the sum of the square of the distances

optimize!(placement)
@show value.(p);
@show objective_value(placement);

# Plotting the points
df = DataFrame()
df.x = value.(p)[:,1]
df.y = value.(p)[:,2]
df.type = vcat(fill("Free points", N), fill("Fixed points", M))
plt = plot(df, x = "x", y = "y", color = "type", Geom.point)
draw(SVG(6inch, 6inch), plt)

n = 5;

Amin = [                                        # We'll try this problem with 4 times with different minimum area constraints
100 100 100 100 100;
 20  50  80 150 200;
180  80  80  80  80;
 20 150  20 200 110]

r = 1

figs=[]

for i = 1:4
    local A = Amin[i, :]

    floor_planning = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
    local x, y, w, h, W, H
    @variables(floor_planning, begin
        x[1:n] >= r
        y[1:n] >= r
        w[1:n] >= 0
        h[1:n] >= 0
        W
        H
    end)

    @constraints(floor_planning, begin
        x[5] + w[5] + r <= W    # No rectangles at the right of Rectangle 5
        x[1] + w[1] + r <= x[3] # Rectangle 1 is at the left of Rectangle 3
        x[2] + w[2] + r <= x[3] # Rectangle 2 is at the left of Rectangle 3
        x[3] + w[3] + r <= x[5] # Rectangle 3 is at the left of Rectangle 5
        x[4] + w[4] + r <= x[5] # Rectangle 4 is at the left of Rectangle 5
        y[4] + h[4] + r <= H    # No rectangles on top of Rectangle 4
        y[5] + h[5] + r <= H    # No rectangles on top of Rectangle 5
        y[2] + h[2] + r <= y[1] # Rectangle 2 is below Rectangle 1
        y[1] + h[1] + r <= y[4] # Rectangle 1 is below Rectangle 4
        y[3] + h[3] + r <= y[4] # Rectangle 3 is below Rectangle 4
        w .<= 5*h               # Aspect ratio constraint
        h .<= 5*w               # Aspect ratio constraint
        A .<= h .*  w           # Area constraint
    end)

    @objective(floor_planning, Min, W + H)

    optimize!(floor_planning)

    @show objective_value(floor_planning);

    D = DataFrame(x = value.(x), y = value.(y), x2 = value.(x) .+ value.(w), y2 = value.(y) .+ value.(h))
    local plt = plot(D, xmin = :x, ymin = :y, xmax = :x2, ymax = :y2, Geom.rect)
    push!(figs, plt)
end

draw(SVG(6inch, 6inch), vstack(hstack(figs[1], figs[2]), hstack(figs[3], figs[4])))


using Brownian
using Base.Test

@test convert(FGN, FBM(0:0.5:10, 0.4)) == FGN(0.757858283255199,0.4)

p = FBM(0:1/2^10:1, 0.4)

# fBm using FFT approach

rand(p)
rand(p, fbm=false)
rand([p, p])

# fBm using Riemann-Liouville approach

rand(p, method=:rl)
rand(p, method=:rl, wtype=:improved)
rand([p, p], method=:rl)
rand([p, p], method=:rl, wtype=:improved)

# fBm using Cholesky approach

rand(p, method=:chol)
rand(p, fbm=false, method=:chol)
rand([p, p], method=:chol)

# Checking function for efficient Cholesky update

p = FBM(0:.1:.1*100, 0.4)
q = FBM(0:.1:.1*101, 0.4)

P = autocov(p)
Q = autocov(q)
c = chol(P)'
chol_update(convert(Array{Float64, 2}, c), Q)


function weighted_mean(numbers, weights)
    tmp_weights = []
    tmp_number = []
    for i in 1:length(weights)
        if ismissing(numbers[i])
            continue
        else
            push!(tmp_weights, weights[i])
            push!(tmp_number, numbers[i])
        end
    end
    weights = tmp_weights
    numbers = tmp_number
    if length(weights) == 0
        return 0
    end
    if sum(tmp_weights) == 0
        return 0 # Such a pixel is background noise for us
    end

    wm = 0
    for i in 1:length(tmp_number)
        wm = wm + tmp_number[i]*tmp_weights[i]
    end
    return wm/sum(tmp_weights)
end


export Non_negative
export Map_0_1
export Piecewise_positive
export Pow4_positive
export Abs_positive

 # All these functions return a mapping function and
 # the inverse of it
 # they are used to map the real numbers to non-negative real numbers

"""
    Non_negative()

Returns a function and an inverse function inverse function
to map numbers to non-negative numbers.
We use a parabola.

# Examples
```julia-repl
julia> p, p_inv = Non_negative()
(DeconvOptim.var"#5#7"(), DeconvOptim.var"#6#8"())

julia> x = [-1, 2, -3]
3-element Array{Int64,1}:
 -1
  2
 -3

julia> p(x)
3-element Array{Int64,1}:
 1
 4
 9

julia> p_inv(p(x))
3-element Array{Float64,1}:
 1.0
 2.0
 3.0
```
"""
function Non_negative()
    #return x -> map(abs2, x), parab_inv
    return parab, parab_inv
end

parab(x) = abs2.(x)
parab_inv(x) = sqrt.(x)
 # define custom adjoint for parab because of
 # slow broadcasting
function ChainRulesCore.rrule(::typeof(parab), x)
    Y = parab(x)
    function aux_pullback(barx)
        return zero(eltype(Y)), (2 .* barx) .* x
    end
    return Y, aux_pullback
end

"""
    Map_0_1()

Returns a function and an inverse function
to map numbers to an interval between 0 and 1.
via an exponential function.

"""
function Map_0_1()
    return f01, f01_inv
end

f01(x) = 1 .- exp.(.- x.^2)
f01_inv(y) = sqrt.(.- log.(1 .- y))

"""
    Piecewise_positive()

Returns a function and an inverse function
to map numbers to larger than 0 via
two function stitched together.
"""
function Piecewise_positive()
    return f_pw_pos, f_pw_pos_inv
end

f_pw_pos(x) = ifelse.(x .> 0, one(eltype(x)) .+ x,one(eltype(x))./(one(eltype(x)).-x))
f_pw_pos_grad(x) = ifelse.(x .> 0, one(eltype(x)) , one(eltype(x))./abs2.(one(eltype(x)).-x))
f_pw_pos_inv(y) = ifelse.(y .> 1, y .- one(eltype(y)), one(eltype(y)) .- one(eltype(y))./y)

function ChainRulesCore.rrule(::typeof(f_pw_pos), x)
    Y = f_pw_pos(x)
    function aux_pullback(barx)
        return zero(eltype(Y)), barx .* f_pw_pos_grad(x)
    end
    return Y, aux_pullback
end

"""
    Pow4_positive()

Returns a function and an inverse function
to map numbers to larger than 0 with `abs2.(abs2.(x))`

"""
function Pow4_positive()
    return f_pow4, f_pow4_inv
end

f_pow4(x) = abs2.(abs2.(x))
f_pow4_inv(y) = sqrt.(sqrt.(y))

"""
    Abs_positive()

Returns a function and an inverse function
to map numbers to larger than 0.
"""
function Abs_positive()
    return f_abs, f_abs_inv
end

f_abs(x) = abs.(x)
f_abs_inv(y) = abs.(y)


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

function _altaz2hadec{T<:AbstractFloat}(alt::T, az::T, lat::T)
    # Convert to radians.
    alt_r = deg2rad(alt)
    az_r = deg2rad(az)
    lat_r = deg2rad(lat)
    # Find local hour angle (in degrees, from 0. to 360.).
    ha = rad2deg(atan2(-sin(az_r)*cos(alt_r),
                       -cos(az_r)*sin(lat_r)*cos(alt_r) +
                       sin(alt_r)*cos(lat_r)))
    ha < 0.0 && (ha += 360.0)
    ha = ha.%360.0
    # Find declination (positive if north of Celestial Equator, negative if
    # south)
    sindec = sin(lat_r)*sin(alt_r) + cos(lat_r)*cos(alt_r)*cos(az_r)
    dec = rad2deg(asin(sindec))  # convert dec to degrees
    return ha, dec
end

"""
    altaz2hadec(alt, az, lat) -> ha, dec

### Purpose ###

Convert Horizon (Alt-Az) coordinates to Hour Angle and Declination.

### Explanation ###

Can deal with the NCP singularity.  Intended mainly to be used by program
`hor2eq`.

### Arguments ###

Input coordinates may be either a scalar or an array, of the same dimension.

* `alt`: local apparent altitude, in degrees, scalar or array.
* `az`: the local apparent azimuth, in degrees, scalar or vector, measured
 *east* of *north*!!!  If you have measured azimuth west-of-south (like the book
 Meeus does), convert it to east of north via: `az = (az + 180) % 360`.
* `lat`: the local geodetic latitude, in degrees, scalar or array.

`alt` and `az` can be given as a 2-tuple `(alt, az)`.

### Output ###

2-tuple `(ha, dec)`

* `ha`: the local apparent hour angle, in degrees.  The hour angle is the time
  that right ascension of 0 hours crosses the local meridian.  It is
  unambiguously defined.
* `dec`: the local apparent declination, in degrees.

The output coordinates are always floating points and have the same type (scalar
or array) as the input coordinates.

### Example ###

Arcturus is observed at an apparent altitude of 59d,05m,10s and an azimuth
(measured east of north) of 133d,18m,29s while at the latitude of +43.07833
degrees.  What are the local hour angle and declination of this object?

``` julia
ha, dec = altaz2hadec(ten(59,05,10), ten(133,18,29), 43.07833)
# => (336.6828582472844,19.182450965120402)
```
The widely available XEPHEM code gets:

    Hour Angle = 336.683
    Declination = 19.1824

### Notes ###

`hadec2altaz` converts Hour Angle and Declination to Horizon (Alt-Az)
coordinates.

Code of this function is based on IDL Astronomy User's Library.
"""
altaz2hadec(alt::Real, az::Real, lat::Real) =
    _altaz2hadec(promote(float(alt), float(az), float(lat))...)

altaz2hadec(altaz::Tuple{Real, Real}, lat::Real) = altaz2hadec(altaz..., lat)

function altaz2hadec{R1<:Real, R2<:Real, R3<:Real}(alt::AbstractArray{R1},
                                                   az::AbstractArray{R2},
                                                   lat::AbstractArray{R3})
    @assert length(alt) == length(az) == length(lat)
    typealt = typeof(float(one(R1)))
    ha  = similar(alt, typealt)
    dec = similar(alt, typealt)
    for i in eachindex(alt)
        ha[i], dec[i] = altaz2hadec(alt[i], az[i], lat[i])
    end
    return ha, dec
end


################################################################################
#                             Euler Angle and Axis
################################################################################

export angleaxis_to_angle, angleaxis_to_dcm, angleaxis_to_quat

################################################################################
#                                  Operations
################################################################################

"""
    function *(ea₂::EulerAngleAxis{T1}, ea₁::EulerAngleAxis{T2}) where {T1,T2}

Compute the composed rotation of `ea₁ -> ea₂`. Notice that the rotation will be
represented by a Euler angle and axis (see `EulerAngleAxis`). By convention, the
output angle will always be in the range [0, π] [rad].

Notice that the vector representing the axis in `ea₁` and `ea₂` must be unitary.
This function neither verifies this nor normalizes the vector.

"""
function *(ea₂::EulerAngleAxis{T1}, ea₁::EulerAngleAxis{T2}) where {T1,T2}
    # Auxiliary variables.
    sθ₁o2, cθ₁o2 = sincos(ea₁.a/2)
    sθ₂o2, cθ₂o2 = sincos(ea₂.a/2)

    v₁ = ea₁.v
    v₂ = ea₂.v

    # Compute `cos(θ/2)` in which `θ` is the new Euler angle.
    cθo2 = cθ₁o2*cθ₂o2 - sθ₁o2*sθ₂o2 * dot(v₁, v₂)

    T = promote_type( T1, T2, typeof(sθ₁o2), typeof(sθ₂o2), typeof(cθo2) )

    if abs(cθo2) >= 1-eps()
        # In this case, the rotation is the identity.
        return EulerAngleAxis( T(0), SVector{3,T}(0,0,0) )
    else
        # Compute `sin(θ/2)` in which `θ` is the new Euler angle.
        sθo2 = sqrt(1 - cθo2*cθo2)

        # Compute the θ angle between [0, 2π].
        θ = 2acos(cθo2)

        # Keep the angle between [0, π].
        s = +1
        if θ > π
            θ = T(2)*π - θ
            s = -1
        end

        v = s*( sθ₁o2*cθ₂o2*v₁ + cθ₁o2*sθ₂o2*v₂ + sθ₁o2*sθ₂o2*(v₁ × v₂) )/sθo2

        return EulerAngleAxis(θ, v)
    end
end

"""
    @inline function inv(ea::EulerAngleAxis)

Compute the inverse rotation of `ea`. The Euler angle returned by this function
will always be in the interval [0, π].

"""
@inline function inv(ea::EulerAngleAxis{T}) where T<:Number
    # Make sure that the Euler angle is always in the inverval [0,π]
    s = -1
    θ = mod(ea.a, T(2)*π)

    if θ > π
        s = 1
        θ = T(2)π - θ
    end

    EulerAngleAxis(θ, s*ea.v)
end

################################################################################
#                                      IO
################################################################################

"""
    function display(ea::EulerAngleAxis{T}) where T
    function show(io::IO, mime::MIME"text/plain", ea::EulerAngleAxis{T}) where T

Display in `stdout` the Euler angle and axis `ea`.

"""
function show(io::IO, ea::EulerAngleAxis{T}) where T
    θ   = ea.a
    v   = ea.v
    str = @sprintf "%8.4f rad, [%8.4f, %8.4f, %8.4f]" ea.a v[1] v[2] v[3]

    print(io, "EulerAngleAxis{$T}: $str")

    nothing
end

function show(io::IO, mime::MIME"text/plain", ea::EulerAngleAxis{T}) where T
    # Check if the user wants colors.
    color = get(io, :color, false)

    b = (color) ? "\x1b[1m"         : ""
    d = (color) ? "\x1b[0m"         : ""
    g = (color) ? "\x1b[1m\x1b[32m" : ""
    y = (color) ? "\x1b[1m\x1b[33m" : ""

    str_a = @sprintf "%8.4f rad (%8.4f deg)" ea.a rad2deg(ea.a)
    str_v = @sprintf "[%8.4f, %8.4f, %8.4f]" ea.v[1] ea.v[2] ea.v[3]

    println(io, "EulerAngleAxis{$T}:")
    println(io, "$y  Euler angle: $d" * str_a)
      print(io, "$y   Euler axis: $d" * str_v)

    nothing
end

################################################################################
#                                 Conversions
################################################################################

# Euler Angles
# ==============================================================================

"""
    @inline function angleaxis_to_angle(θ::Number, v::AbstractVector, rot_seq::Symbol)
    @inline function angleaxis_to_angle(ea::EulerAngleAxis, rot_seq::Symbol)

Convert the Euler angle `θ` [rad]  and Euler axis `v`, which must be a unit
vector, to Euler angles with rotation sequence `rot_seq`. Those values can also
be passed inside the structure `ea` (see `EulerAngleAxis`).

The rotation sequence is defined by a `:Symbol`. The possible values are:
`:XYX`, `XYZ`, `:XZX`, `:XZY`, `:YXY`, `:YXZ`, `:YZX`, `:YZY`, `:ZXY`, `:ZXZ`,
`:ZYX`, and `:ZYZ`. If no value is specified, then it defaults to `:ZYX`.

# Example

```julia-repl
julia> ea = EulerAngleAxis(45*pi/180, [1;0;0]);

julia> angleaxis_to_angles(ea, :ZXY)
EulerAngles{Float64}:
  R(Z):   0.0000 rad (   0.0000 deg)
  R(X):   0.7854 rad (  45.0000 deg)
  R(Y):   0.0000 rad (   0.0000 deg)

```
"""
@inline function angleaxis_to_angle(θ::Number, v::AbstractVector, rot_seq::Symbol)
    # First we convert to quaternions then to Euler angles.
    quat_to_angle( angleaxis_to_quat(θ, v), rot_seq )
end

@inline angleaxis_to_angle(ea::EulerAngleAxis, rot_seq::Symbol) =
    angleaxis_to_angle(ea.a, ea.v, rot_seq)

# DCM
# ==============================================================================

"""
    @inline function angleaxis_to_dcm(a::Number, v::AbstractVector)
    @inline function angleaxis_to_dcm(ea::EulerAngleAxis)

Convert the Euler angle `a` [rad] and Euler axis `v`, which must be a unit
vector to a DCM. Those values can also be passed inside the structure `ea` (see
`EulerAngleAxis`).

# Remarks

It is expected that the vector `v` is unitary. However, no verification is
performed inside the function. The user must handle this situation.

# Example

```julia-repl
julia> v = [1;1;1];

julia> v /= norm(v);

julia> angleaxis_to_dcm(pi/2,v)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  0.333333   0.910684  -0.244017
 -0.244017   0.333333   0.910684
  0.910684  -0.244017   0.333333

julia> ea = EulerAngleAxis(pi/2,v);

julia> angleaxis_to_dcm(ea)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
  0.333333   0.910684  -0.244017
 -0.244017   0.333333   0.910684
  0.910684  -0.244017   0.333333
```

"""
@inline function angleaxis_to_dcm(a::Number, v::AbstractVector)
    sθ, cθ = sincos(a)
    aux    = 1 - cθ

    DCM(  cθ + v[1]*v[1]*aux   , v[1]*v[2]*aux + v[3]*sθ, v[1]*v[3]*aux - v[2]*sθ,
        v[1]*v[2]*aux - v[3]*sθ,    cθ + v[2]*v[2]*aux  , v[2]*v[3]*aux + v[1]*sθ,
        v[1]*v[3]*aux + v[2]*sθ, v[2]*v[3]*aux - v[1]*sθ,   cθ + v[3]*v[3]*aux)'
end

@inline angleaxis_to_dcm(ea::EulerAngleAxis) = angleaxis_to_dcm(ea.a, ea.v)

# Quaternions
# ==============================================================================

"""
    function angleaxis_to_quat(θ::Number, v::AbstractVector)

Convert the Euler angle `θ` [rad] and Euler axis `v`, which must be a unit
vector, to a quaternion.

# Remarks

It is expected that the vector `v` is unitary. However, no verification is
performed inside the function. The user must handle this situation.

# Example

```julia-repl
julia> v = [1;1;1];

julia> v /= norm(v);

julia> angleaxis_to_quat(pi/2,v)
Quaternion{Float64}:
  + 0.7071067811865476 + 0.408248290463863.i + 0.408248290463863.j + 0.408248290463863.k
```

"""
function angleaxis_to_quat(θ::Number, v::AbstractVector)
    # Check the arguments.
    (length(v) > 3) && throw(ArgumentError("The provided vector for the Euler axis must have 3 elements."))

    cθo2 = cos(θ/2)

    # Keep `q0` positive.
    s = (cθo2 < 0) ? -1 : +1

    # Create the quaternion.
    Quaternion( s*cos(θ/2), s*sin(θ/2)*v )
end

"""
    function angleaxis_to_quat(angleaxis::EulerAngleAxis)

Convert a Euler angle and Euler axis `angleaxis` (see `EulerAngleAxis`) to a
quaternion.

# Remarks

It is expected that the vector `angleaxis.v` is unitary. However, no
verification is performed inside the function. The user must handle this
situation.

# Example

```julia-repl
julia> v = [1;1;1];

julia> v /= norm(v);

julia> angleaxis_to_quat(EulerAngleAxis(pi/2,v))
Quaternion{Float64}:
  + 0.7071067811865476 + 0.408248290463863.i + 0.408248290463863.j + 0.408248290463863.k
```

"""
angleaxis_to_quat(angleaxis::EulerAngleAxis) = angleaxis_to_quat(angleaxis.a,
                                                                 angleaxis.v)


using Test
using LinearAlgebra
using PolynomialBasis
using IntervalArithmetic
using Revise
using ImplicitDomainQuadrature

IDQ = ImplicitDomainQuadrature

function corner_distance_function(x::V,xc) where {V<:AbstractVector}
    v = xc - x
    if all(x .<= xc)
        minimum(v)
    elseif all(x.> xc)
        return -sqrt(v'*v)
    elseif x[2] > xc[2]
        return v[2]
    else
        return v[1]
    end
end

function corner_distance_function(points::M,xc) where {M<:AbstractMatrix}
    return vec(mapslices(x->corner_distance_function(x,xc),points,dims=1))
end

corner = [0.,0.]
polyorder = 2
numqp = 5
poly = InterpolatingPolynomial(1,2,polyorder)
coeffs = corner_distance_function(poly.basis.points,corner)
update!(poly,coeffs)

quad = area_quadrature(poly,+1,[-1.,-1.],[1.,1.],numqp,numsplits=3)
squad = surface_quadrature(poly,[-1.,-1.],[1.,1.],numqp,numsplits=3)

using Plots
xrange = -1:1e-2:1
contour(xrange,xrange,(x,y)->poly(x,y),levels=[0.0],aspect_ratio=:equal)
scatter!(quad.points[1,:],quad.points[2,:])
scatter!(squad.points[1,:],squad.points[2,:])


import Pkg; Pkg.activate("../../")

using StatsFuns
using BenchmarkTools
import Random
using Flux
using Zygote
include(joinpath(@__DIR__, "BnpUtil_old.jl"))

function log_sum_exp(logx::T) where T <: AbstractArray
  mx = maximum(logx)
  return mx + log(sum(exp.(logx .- mx)))
end

Random.seed!(0);
x = randn(10)
@assert StatsFuns.logsumexp(x) == BnpUtil.logsumexpdd(x)[1]

# Logsumexp benchmarks
@benchmark StatsFuns.logsumexp(x)
@benchmark BnpUtil.logsumexp(x)[1]
@benchmark log_sum_exp(x)

### Gmm benchmarks ###
Random.seed!(0);
N, K = 101, 10
X = randn(N, K)
m = randn(K)
s = rand(K)
w = let
    _w = rand(K)
    _w / sum(_w)
end

lpdf_gmm_sf(x, m, s, w) = sum(logsumexp(normlogpdf.(m, s, x) .+ log.(w), dims=2))
@assert sum(BnpUtil.lpdf_gmm(X, m', s', w', dims=2)) == lpdf_gmm_sf(X, m', s', w')
@btime lpdf_gmm_sf(X, m', s', w')
@btime sum(BnpUtil.lpdf_gmm(X, m', s', w', dims=2))

# NOTE:
# logsumexp and lpdf_gmm seems to be comparable. In fact, for forward
# computaton, StatsFuns seems to be a little faster for logsumexp.

### BACKWARD: ###
# BnpUtil
@benchmark gs = Flux.gradient(Flux.params(m)) do
    sum(BnpUtil.lpdf_gmm(X, m', s', w', dims=2))
end

# StatsFuns (error?!)
@benchmark gs = Flux.gradient(Flux.params(m)) do
    lpdf_gmm_sf(X, m', s', w')
end

# Diagnose problem
x, loc, scale = 0.0, 2.0, 1.0
normlogpdf(loc, scale, x)
Flux.gradient(mu -> normlogpdf(mu, scale, x), loc)

function my_normlogpdf(loc, scale, x)
    z = (x - loc) / scale
    return -z * z * 0.5 - 0.5 * log(2 * pi * scale*scale)
end
Flux.gradient(mu -> my_normlogpdf(mu, scale, x), loc)


################################################################################
#  Copyright 2022, Tom Van Acker (BASF)                                        #
################################################################################
# TFT.jl                                                                       #
# A Julia package for Taylor-Fourier Transform.                                #
# See http://github.com/timmyfaraday/TFT.jl                                    #
################################################################################

"""
    TFT.sample_ospline(D::Int, K::Int, N::Int)

Function to obtain samples of the up-to-Dth-degree derivatives of Kth-degree
o-spline.

Input:
- `D::Int`  | maximum degree of the derivative [-]
- `K::Int`  | degree of the o-spline [-]
- `N::Int`  | number of samples of the fundamental cycle [-]

Output
- `Φ::Matrix{<:Real}`   | samples of the up-to-Dth-degree derivatives of the
                        | Kth-degree o-spline
"""
function sample_ospline(D::Int, K::Int, N::Int)
    # define the normalized time range `rU`
    ΔU  = 1 / N
    rU  = 0.0:ΔU:(1.0-ΔU)
    # define the knots range `rK`
    rK  = [-K:-1..., 1:K...]

    # initialize the Φ-matrix
    Φ = ones((K+1) * N, D+1)

    # fill the Φ-matrices
    for ctr in 0:K
        # one-based counter `cnt` wrt zero-based counter `ctr`
        ctn = ctr + 1
        # range of the first dimension of Φ wrt zero-based counter `ctr`
        rΦ  = (N * ctr + 1):(N * (ctr + 1))
        # base o-spline
        u   = -(K + 1) / 2 + ctr .+ rU
        ψ   = [1, -rK[ctn]] ./ -rK[ctn]
        for nk in 1:K
            Φ[rΦ, 1] .*= (u .- rK[ctr + nk]) ./ -rK[ctr + nk]
            if nk > 1
                ψ   = _DSP.conv(ψ, [1, -rK[ctr + nk]] ./ -rK[ctr + nk])
            end
        end
        # derivative by Horner scheme
        # NB: the Julia Polynomials pkg, takes its coefficients in reverse,
        # hence the need for the double reserve when determining ψ.
        for nd in 2:(D+1)
            ψ   = reverse(_POL.derivative(_POL.Polynomial(reverse(ψ))).coeffs)
            Φ[rΦ, nd] .*= ψ[1]
            for nk in 2:(K-nd+2)
                Φ[rΦ, nd] .*= u
                Φ[rΦ, nd] .+= ψ[nk]
            end
        end
    end

    # return the samples of the o-spline and its derivatives: Φ
    return Φ
end


# ---
# title: 1276. Number of Burgers with No Waste of Ingredients
# id: problem1276
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Greedy
# link: <https://leetcode.com/problems/number-of-burgers-with-no-waste-of-ingredients/description/>
# hidden: true
# ---
#
# Given two integers `tomatoSlices` and `cheeseSlices`. The ingredients of
# different burgers are as follows:
#
#   * **Jumbo Burger:** 4 tomato slices and 1 cheese slice.
#   * **Small Burger:** 2 Tomato slices and 1 cheese slice.
#
# Return `[total_jumbo, total_small]` so that the number of remaining
# `tomatoSlices` equal to 0 and the number of remaining `cheeseSlices` equal to
# 0. If it is not possible to make the remaining `tomatoSlices` and
# `cheeseSlices` equal to 0 return `[]`.
#
#
#
# **Example 1:**
#
#
#
#     Input: tomatoSlices = 16, cheeseSlices = 7
#     Output: [1,6]
#     **Explantion:** To make one jumbo burger and 6 small burgers we need 4*1 + 2*6 = 16 tomato and 1 + 6 = 7 cheese. There will be no remaining ingredients.
#
#
# **Example 2:**
#
#
#
#     Input: tomatoSlices = 17, cheeseSlices = 4
#     Output: []
#     **Explantion:** There will be no way to use all ingredients to make small and jumbo burgers.
#
#
# **Example 3:**
#
#
#
#     Input: tomatoSlices = 4, cheeseSlices = 17
#     Output: []
#     **Explantion:** Making 1 jumbo burger there will be 16 cheese remaining and making 2 small burgers there will be 15 cheese remaining.
#
#
# **Example 4:**
#
#
#
#     Input: tomatoSlices = 0, cheeseSlices = 0
#     Output: [0,0]
#
#
# **Example 5:**
#
#
#
#     Input: tomatoSlices = 2, cheeseSlices = 1
#     Output: [0,1]
#
#
#
#
# **Constraints:**
#
#   * `0 <= tomatoSlices <= 10^7`
#   * `0 <= cheeseSlices <= 10^7`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#=

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-06-27

References
----------

Simple port of the file quantecon.models.jv

http://quant-econ.net/jl/jv.html
=#

using Distributions
using QuantEcon

# NOTE: only brute-force approach is available in bellman operator. Waiting on a simple constrained optimizer to be written in pure Julia

"""
A Jovanovic-type model of employment with on-the-job search.

The value function is given by

\\[V(x) = \\max_{\\phi, s} w(x, \\phi, s)\\]

for

    w(x, phi, s) := x(1 - phi - s) + beta (1 - pi(s)) V(G(x, phi)) +
                    beta pi(s) E V[ max(G(x, phi), U)

where

* `x`: : human capital
* `s` : search effort
* `phi` : investment in human capital
* `pi(s)` : probability of new offer given search level s
* `x(1 - \phi - s)` : wage
* `G(x, \phi)` : new human capital when current job retained
* `U` : Random variable with distribution F -- new draw of human capital

##### Fields

- `A::Real` : Parameter in human capital transition function
- `alpha::Real` : Parameter in human capital transition function
- `bet::Real` : Discount factor in (0, 1)
- `x_grid::Vector` : Grid for potential levels of x
- `G::Function` : Transition `function` for human captial
- `pi_func::Function` : `function` mapping search effort to the probability of getting a new job offer
- `F::UnivariateDistribution` : A univariate distribution from which the value of new job offers is drawn
- `quad_nodes::Vector` : Quadrature nodes for integrating over phi
- `quad_weights::Vector` : Quadrature weights for integrating over phi
- `epsilon::Float64` : A small number, used in optimization routine

"""
type JvWorker
    A::Real
    alpha::Real
    bet::Real
    x_grid::Vector{Float64}
    G::Function
    pi_func::Function
    F::UnivariateDistribution
    quad_nodes::Vector
    quad_weights::Vector
    epsilon::Float64
end

"""
Constructor with default values for `JvWorker`

##### Arguments

 - `A::Real(1.4)` : Parameter in human capital transition function
 - `alpha::Real(0.6)` : Parameter in human capital transition function
 - `bet::Real(0.96)` : Discount factor in (0, 1)
 - `grid_size::Int(50)` : Number of points in discrete grid for `x`
 - `epsilon::Float(1e-4)` : A small number, used in optimization routine

##### Notes

There is also a version of this function that accepts keyword arguments for
each parameter

"""
function JvWorker(A=1.4, alpha=0.6, bet=0.96, grid_size=50, epsilon=1e-4)
    G(x, phi) = A .* (x .* phi).^alpha
    pi_func = sqrt
    F = Beta(2, 2)

    # integration bounds
    a, b, = quantile(F, 0.005), quantile(F, 0.995)

    # quadrature nodes/weights
    nodes, weights = qnwlege(21, a, b)

    # Set up grid over the state space for DP
    # Max of grid is the max of a large quantile value for F and the
    # fixed point y = G(y, 1).
    grid_max = max(A^(1.0 / (1.0 - alpha)), quantile(F, 1 - epsilon))

    # range for linspace(epsilon, grid_max, grid_size). Needed for
    # CoordInterpGrid below
    x_grid = collect(linspace(epsilon, grid_max, grid_size))

    JvWorker(A, alpha, bet, x_grid, G, pi_func, F, nodes, weights, epsilon)
end

# make kwarg version
JvWorker(;A=1.4, alpha=0.6, bet=0.96, grid_size=50, epsilon=1e-4) = JvWorker(A, alpha, bet, grid_size, epsilon)

"""
Apply the Bellman operator for a given model and initial value, returning only the value function

##### Arguments

- `jv::JvWorker` : Instance of `JvWorker`
- `V::Vector`: Current guess for the value function
- `new_V::Vector` : Storage for updated value function

##### Returns

None, `new_V` is updated in place with the value function.

##### Notes

Currently, only the brute-force approach is available. We are waiting on a simple constrained optimizer to be written in pure Julia

"""
function bellman_operator!(jv::JvWorker, V::Vector, new_V::Vector)

    # simplify notation
    G, pi_func, F, bet, epsilon = jv.G, jv.pi_func, jv.F, jv.bet, jv.epsilon
    nodes, weights = jv.quad_nodes, jv.quad_weights

    # prepare interpoland of value function
    Vf = LinInterp(jv.x_grid, V)

    # instantiate the linesearch variables
    max_val = -1.0
    cur_val = 0.0
    max_s = 1.0
    max_phi = 1.0
    search_grid = linspace(epsilon, 1.0, 15)

    for (i, x) in enumerate(jv.x_grid)

        function w(z)
            s, phi = z
            function h(u)
              out = similar(u)
              for j in 1:length(u)
                out[j] = Vf(max(G(x, phi), u[j])) * pdf(F, u[j])
              end
              out
            end
            integral = do_quad(h, nodes, weights)
            q = pi_func(s) * integral + (1.0 - pi_func(s)) * Vf(G(x, phi))

            return - x * (1.0 - phi - s) - bet * q
        end

        for s in search_grid
            for phi in search_grid
                if s + phi <= 1.0
                    cur_val = -w((s, phi))
                else
                    cur_val = -1.0
                end
                if cur_val > max_val
                    max_val, max_s, max_phi = cur_val, s, phi
                end
            end
        end

        new_V[i] = max_val
    end
end

"""
Apply the Bellman operator for a given model and initial value, returning policies

##### Arguments

- `jv::JvWorker` : Instance of `JvWorker`
- `V::Vector`: Current guess for the value function
- `out::Tuple{Vector, Vector}` : Storage for the two policy rules

##### Returns

None, `out` is updated in place with the two policy functions.

##### Notes

Currently, only the brute-force approach is available. We are waiting on a simple constrained optimizer to be written in pure Julia

"""
function bellman_operator!(jv::JvWorker, V::Vector, out::Tuple{Vector, Vector})

    # simplify notation
    G, pi_func, F, bet, epsilon = jv.G, jv.pi_func, jv.F, jv.bet, jv.epsilon
    nodes, weights = jv.quad_nodes, jv.quad_weights

    # prepare interpoland of value function
    Vf = LinInterp(jv.x_grid, V)

    # instantiate variables
    s_policy, phi_policy = out[1], out[2]

    # instantiate the linesearch variables
    max_val = -1.0
    cur_val = 0.0
    max_s = 1.0
    max_phi = 1.0
    search_grid = linspace(epsilon, 1.0, 15)

    for (i, x) in enumerate(jv.x_grid)

        function w(z)
            s, phi = z
            function h(u)
              out = similar(u)
              for j in 1:length(u)
                out[j] = Vf(max(G(x, phi), u[j])) * pdf(F, u[j])
              end
              out
            end
            integral = do_quad(h, nodes, weights)
            q = pi_func(s) * integral + (1.0 - pi_func(s)) * Vf(G(x, phi))

            return - x * (1.0 - phi - s) - bet * q
        end

        for s in search_grid
            for phi in search_grid
                if s + phi <= 1.0
                    cur_val = -w((s, phi))
                else
                    cur_val = -1.0
                end
                if cur_val > max_val
                    max_val, max_s, max_phi = cur_val, s, phi
                end
            end
        end

      s_policy[i], phi_policy[i] = max_s, max_phi
  end
end

function bellman_operator(jv::JvWorker, V::Vector; ret_policies=false)
    if ret_policies
        out = (similar(V), similar(V))
    else
        out = similar(V)
    end
    bellman_operator!(jv, V, out)
    return out
end


# Duffing_orbit_diagram_vs_period_zooom_v1.jl
# ===========================================

using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, OrdinaryDiffEq
using StaticArrays
using Random
using Statistics

function Duffing_ODE(du,u,p,t)
    d, a, omega = p
    du[1] =  u[2]
    du[2] =  - u[1] - u[1]*u[1]*u[1] - d*u[2] + a*sin(omega*t)
end

function Poincare_orbit(d,a,omega,y_init,ntrans,nattr,pcol)

    p = [d a omega]
    period = 2*pi / omega

 # transient to attractor
 # ----------------------
    tspan = (0., ntrans*period)
    prob = ODEProblem(Duffing_ODE, y_init, tspan, p)
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)

  # attractor  with root of Poincare map
  # ---------
    tspan = (0., nattr*period)
    y_init = sol.u[end]
    prob = ODEProblem(Duffing_ODE, y_init, tspan, p)
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8, saveat = period/2)

    u1vec = [u[1] for u in sol.u]
    npts = length(u1vec)
    for j = 2:2:npts
        u1vec[j] = -u1vec[j]
    end

  # plot
  # ----
    plt.scatter(period*ones(npts),u1vec,s=1, c = pcol)

    return y_init

end # of function

# parameters and initial values
# -----------------------------
d = 0.1
avec = [19.7 19.75]
n_T = 1500

ntrans = 800   # no. of transient periods
nattr = 64     # no. of points on the attractor
y_init = [1 0]

# open plot
# ---------
fig = figure()

# subplot loop
# ------------
@time for iplt = 1:2

    if iplt  == 1
        a = avec[1]
        T_min = 4.1 # 1
        T_max = 5.8 # 32
        T_vec = LinRange(T_min, T_max, n_T)
    else
        a = avec[2]
        T_min = 5.06 # 1
        T_max = 5.42 # 32
        T_vec = LinRange(T_max, T_min, n_T)
    end

    subplot(1,2,iplt)

    xlabel(L"T = 2\pi / \omega")
    ylabel(L"u_1")
    pcol = COLORS[1]

#  control parameter loop (period)
#  ----------------------
   for i_T = 1:length(T_vec)
       period = T_vec[i_T]
       omega = 2*pi / period
       y_init =  Poincare_orbit(d,a,omega,y_init,ntrans,nattr,pcol)
   end # of period loop

#   tight_layout()
#   ax = gca()
Tbmin = 5.1
Tbmax = 5.3
u1bmin = -0.94
u1bmax = -0.7

#  plot zoom box
if iplt == 1
#    xbox = [Tbmin, Tbmax, Tbmax , Tbmin, Tbmin]
#    ybox = [u1bmin, u1bmin, u1bmax, u1bmax, u1bmin]
#    plt.plot(xbox,ybox,linewidth=1.0,color="k")
   ax = gca()
   ax.set_xlim([T_min,T_max])
   ax.set_ylim([-1.7,1.4])
end
if iplt == 2
    ax = gca()
    ax.set_xlim([Tbmin,Tbmax])
    ax.set_ylim([u1bmin,u1bmax])
end

end # of subplot loop

# %%

add_identifiers!(fig)
fig.tight_layout(pad=0.3)

wsave(plotsdir("9", "duffing_orbitdiagram_zoom"), fig)


#=

  Asia medical problem in Turing.jl

  Port of *Church model https://github.com/stuhlmueller/jschurch/blob/master/tests/asia.church

  WebPPL model (exact probabilities)
  [ 'smoker', 0.6700973677390694 ],
  [ 'tb', 0.04072414942060647 ],
  [ 'cancer', 0.07774610343933955 ],
  [ 'tb_or_cancer', 0.11473103550405392 ],
  [ 'bronchitis', 1 ],
  [ 'xray', 0.15669986301877015 ],
  [ 'bronchitis', 1 ] ]

  This Turing model (using SMC())
  """
  Summary Statistics
    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

        smoker    0.6759    0.4681     0.0047    0.0073   4091.3819    1.0002     1207.2534
            tb    0.0400    0.1960     0.0020    0.0031   4124.4085    1.0004     1216.9987
        cancer    0.0719    0.2583     0.0026    0.0041   4258.4295    1.0001     1256.5446
  tb_or_cancer    0.1068    0.3089     0.0031    0.0052   4182.5913    1.0004     1234.1668
    bronchitis    1.0000    0.0000     0.0000    0.0000         NaN       NaN           NaN
          xray    0.1504    0.3575     0.0036    0.0057   4220.3910    1.0004     1245.3205
      dyspnoea    0.8080    0.3939     0.0039    0.0060   4220.4111    0.9999     1245.3264
  """

=#

using Turing
include("jl_utils.jl")

@model function asia()
    smoker ~ flip(0.5)
    tb ~ flip(smoker ? 0.05 : 0.01)
    cancer ~ flip(smoker ? 0.1 : 0.01)
    tb_or_cancer ~ Dirac(tb || cancer)
    bronchitis ~ flip(smoker ? 0.6 : 0.3)
    xray ~ flip(tb_or_cancer ? 0.98 : 0.05)
    dyspnoea ~ flip(tb_or_cancer ? bronchitis ? 0.9 : 0.7
                    : bronchitis ? 0.8 : 0.1)
    true ~ Dirac(bronchitis == true)
    true ~ Dirac(dyspnoea == true)

end

model = asia()

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, IS(), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, SMC(), MCMCThreads(), 10_000, 4)

display(chns)


module TestRandomMC

using QuantEcon
using Base.Test
using FactCheck

facts("Testing random_mc.jl") do
    context("Test random_markov_chain") do
        n, k = 5, 3
        mc_dicts = (Dict("P" => random_markov_chain(n).p, "k" => n),
                    Dict("P" => random_markov_chain(n, k).p, "k" => k))
        for d in mc_dicts
            P = d["P"]
            @fact size(P) --> (n, n)
            @fact all(x->(countnz(x)==d["k"]),
                      [P[i, :] for i in 1:size(P)[1]]) --> true
        end
    end

    context("Test random_stochastic_matrix") do
        n, k = 5, 3
        Ps = (random_stochastic_matrix(n), random_stochastic_matrix(n, k))
        for P in Ps
            @fact all(P .>= 0) --> true
            @fact all(x->isapprox(sum(x), 1),
                      [P[i, :] for i in 1:size(P)[1]]) --> true
        end
    end

    context("Test errors properly thrown") do
        # n <= 0
        @fact_throws random_markov_chain(0)

        # k <= 0
        @fact_throws random_markov_chain(2, 0)

        # k > n
        @fact_throws random_markov_chain(2, 3)
    end
end  # facts

end  # module


f(u,p,t) = p.*u

using DifferentialEquations, Plots
u0 = [10.0]
p = [-0.3]
tspan = (0.0,10.0)
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob)
plot(sol)
ylims!(0.0,10.0)

using Distributions
u0_dist = [Uniform(-10.0,10.0)]

prob_func(prob,i,repeat) = remake(prob, u0 = rand.(u0_dist))
ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)

ensemblesol = solve(ensemble_prob,Tsit5(),EnsembleThreads(),trajectories=100000)

plot(ensemblesol, vars = (0,1), lw=1,alpha=0.1, label=nothing, idxs = 1:250)

g(sol) = sol(4.0)
mean([g(sol) for sol in ensemblesol])

using DiffEqUncertainty
expectation(g, prob, u0_dist, p, MonteCarlo(), Tsit5(); trajectories=100000)

expectation(g, prob, u0_dist, p, Koopman(), Tsit5())

exp(p[1]*4.0)*mean(u0_dist[1])

@time expectation(g, prob, u0_dist, p, MonteCarlo(), Tsit5(); trajectories=100000)

@time expectation(g, prob, u0_dist, p, Koopman(), Tsit5())

u0_dist = [Uniform(0.0,10.0)]
@time expectation(g, prob, u0_dist, p, MonteCarlo(), Tsit5(); trajectories=100_000)

@time expectation(g, prob, u0_dist, p, Koopman(), Tsit5())[1]

exp(p[1]*4.0)*mean(u0_dist[1])

u0_dist = [Normal(3.0,2.0)]
expectation(g, prob, u0_dist, p, Koopman(), Tsit5())

exp(p[1]*4.0)*mean(u0_dist[1])

u0_dist = [truncated(Normal(3.0,2.0),-1000,1000)]
expectation(g, prob, u0_dist, p, Koopman(), Tsit5())

u0_dist = [truncated(Normal(3.0,2.0),-5,11)]
expectation(g, prob, u0_dist, p, Koopman(), Tsit5())

u0_dist = [truncated(Normal(3.0,2.0),3-1000,3+1000)]
expectation(g, prob, u0_dist, p, Koopman(), Tsit5())

g(sol) = [sol(4.0)[1], sol(6.0)[1]]
expectation(g, prob, u0_dist, p, Koopman(), Tsit5(); nout = 2)

exp.(p.*[4.0,6.0])*mean(u0_dist[1])

saveat = tspan[1]:.5:tspan[2]
g(sol) = Matrix(sol)
mean_koop = expectation(g, prob, u0_dist, p, Koopman(), Tsit5(); nout = length(saveat), saveat=saveat)

plot(t->exp(p[1]*t)*mean(u0_dist[1]),tspan..., xlabel="t", label="analytical")
scatter!(collect(saveat),mean_koop.u[:],marker=:o, label=nothing)

function g(sol, power, counter)
    counter[power] = counter[power] + 1
    sol(4.0)[1]^power
end

counters = [0,0,0]
x_koop = expectation(s->g(s,1,counters), prob, u0_dist, p, Koopman(), Tsit5())
x2_koop = expectation(s->g(s,2,counters), prob, u0_dist, p, Koopman(), Tsit5())
x3_koop = expectation(s->g(s,3,counters), prob, u0_dist, p, Koopman(), Tsit5())
counters

function g(sol, counter)
    counter[1] = counter[1] + 1
    v = sol(4.0)[1]
    [v, v^2, v^3]
end

counter = [0]
expectation(s->g(s,counter), prob, u0_dist, p, Koopman(), Tsit5(); nout = 3)
counter

function g(sol)
    x = sol(4.0)[1]
    [x, x^2]
end

koop = expectation(g, prob, u0_dist, p, Koopman(), Tsit5(); nout = 2)
mean_koop = koop[1]
var_koop = koop[2] - mean_koop^2

exp(2*p[1]*4.0)*var(u0_dist[1])

saveat = tspan[1]:.5:tspan[2]
g(sol) = [Matrix(sol)'; (Matrix(sol).^2)']

koop = expectation(g, prob, u0_dist, p, Koopman(), Tsit5(); nout = length(saveat)*2, saveat=saveat)
μ = koop.u[1:length(saveat)]
σ = sqrt.(koop.u[length(saveat)+1:end] - μ.^2)

plot(t->exp(p[1]*t)*mean(u0_dist[1]),tspan..., ribbon = t->-sqrt(exp(2*p[1]*t)*var(u0_dist[1])), label="Analytical Mean, 1 std bounds")
scatter!(collect(saveat),μ,marker=:x, yerror = σ, c=:black, label = "Koopman Mean, 1 std bounds")

function g(sol)
    v = sol(4.0)[1]
    [v, v^2, v^3]
end

koop = expectation(g, prob, u0_dist, p, Koopman(), Tsit5(); nout = 3)
mean_koop = koop[1]
var_koop = koop[2] - mean_koop^2
(koop[3] - 3.0*mean_koop*var_koop - mean_koop^3) / var_koop^(3/2)

g(sol) = sol(4.0)[1]
centralmoment(5, g, prob, u0_dist, p, Koopman(), Tsit5(),
                ireltol = 1e-9, iabstol = 1e-9)

using Quadrature, Cuba

u0_dist = [truncated(Normal(3.0,2.0),-5,11)]
p_dist = [truncated(Normal(-.7, .1), -1,0)]

g(sol) = sol(6.0)[1]

expectation(g, prob, u0_dist, p_dist, Koopman(), Tsit5(), EnsembleThreads();
                quadalg = CubaSUAVE(), batch=10)[1]

using BenchmarkTools

@btime expectation(g, prob, u0_dist, p_dist, Koopman(), Tsit5();
                quadalg = CubaSUAVE())[1]

@btime expectation(g, prob, u0_dist, p_dist, Koopman(), Tsit5(), EnsembleThreads();
                quadalg = CubaSUAVE(), batch=10)[1]

using DiffEqGPU

function f(du, u,p,t)
    @inbounds begin
        du[1] = p[1]*u[1];
    end
    nothing
end

u0 = Float32[10.0]
p = Float32[-0.3]
tspan = (0.0f0,10.0f0)
prob = ODEProblem(f,u0,tspan,p)

g(sol) = sol(6.0)[1]

u0_dist = [truncated(Normal(3.0f0,2.0f0),-5f0,11f0)]
p_dist = [truncated(Normal(-.7f0, .1f0), -1f0,0f0)]

@btime expectation(g, prob, u0_dist, p_dist, Koopman(), Tsit5(), EnsembleGPUArray();
                   quadalg = CubaSUAVE(), batch=1000)[1]

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


import .StatsFuns
using .StatsFuns: xlogx, logistic, logit, log1psq, log1pexp, logsumexp

@adjoint function xlogx(x::Real)
    y = xlogx(x)
    return y, function(Δ::Real)
        return (x > zero(x) ? Δ * (log(x) + one(y)) : zero(y),)
    end
end

@adjoint function logistic(x::Real)
    y = logistic(x)
    return y, Δ->(Δ * y * (1 - y),)
end

@adjoint logit(x::Real) = logit(x), Δ->(Δ / (x * (1 - x)),)

@adjoint log1psq(x::Real) = log1psq(x), Δ->(Δ * 2x / (1 + abs2(x)),)

@adjoint function log1pexp(x::Float64)
    return log1pexp(x), Δ->(Δ * (x < 18.0 ? logistic(x) : x < 33.3 ? 1 - exp(-x) : 1),)
end

@adjoint function log1pexp(x::Float32)
    return log1pexp(x), Δ->(Δ * (x < 9f0 ? logistic(x) : x < 16f0 ? 1 - exp(-x) : 1),)
end

@adjoint function logsumexp(X::AbstractArray{<:Real}; dims=:)
    lse = logsumexp(X; dims=dims)
    return lse, Δ -> (Δ .* exp.(X .- lse),)
end


using OmegaGrav
using MatterPower
using Dierckx
using Plots, LaTeXStrings
using OrdinaryDiffEq
# %% Compute the matter power spectrum using CLASS
include("compute_pk_class.jl")
# %% Compute the analytical linear theory result for K/W
# Reference: Equation (12) of Davis, Miller, White, ApJ, 490, 63 (1997).
zini = 9
a = 1/(1+zini):0.01:1.0 # scale factor
nred = length(a)
sol_δθ = setup_growth(Ωm, 1 - Ωm)
KovW = zeros(nred)
for i = 1:nred
    KovW[i] = -2 / 3 * (sol_δθ(a[i])[2] / sol_δθ(a[i])[1])^2 * a[i] / Ωm
end

# %% Compute `Ωgrav = Ωm * W / 2` at various redshifts
Ωgpk = zeros(nred)
Ωglin = zeros(nred)
for ired = 1:nred
    z = 1 / a[ired] - 1
    # Define functions to return linear and non-linear power spectra
    # Note: The CLASS code takes wavenumbers in units of 1/Mpc (no h) and
    # return power spectra in units of Mpc^3 (no 1/h^3).
    pknl_class(kovh) = cosmo.pk(kovh * h0, z) * h0^3
    pklin_class(kovh) = cosmo.pk_lin(kovh * h0, z) * h0^3
    # Spline interpolate in log(k)
    lnk = log(1e-4):0.05:log(100)
    pknl = Spline1D(lnk, pknl_class.(exp.(lnk)))
    pklin = Spline1D(lnk, pklin_class.(exp.(lnk)))
    # %% Compute Ωgrav from non-linear total matter P(k)
    Ωgpk[ired] = ograv_pk(x -> pknl(log(x)), z, Ωm)
    # %% Compute Ωgrav from linear total matter P(k)
    Ωglin[ired] = ograv_pk(x -> pklin(log(x)), z, Ωm)
end
Ωgrav_pklin = Spline1D(a, Ωglin)
Ωgrav_pknl = Spline1D(a, Ωgpk)

# %% Solve the Layzer-Irvine equation
tspan = (1 / (1 + zini), 1.0)
t0 = tspan[1]
## Non-linear P(k)
W(x) = 2 * Ωgrav_pknl(x)
dWda(x) = 2 * derivative(Ωgrav_pknl, x)
f(u, p, t) = -(2 * u + W(t)) / t - dWda(t)
u0 = KovW[1] * W(t0)
prob = ODEProblem(f, u0, tspan)
sol_pknl = solve(prob, Tsit5())
## Linear P(k): This result must agree with the analytical linear theory result
W(x) = 2 * Ωgrav_pklin(x)
dWda(x) = 2 * derivative(Ωgrav_pklin, x)
f(u, p, t) = -(2 * u + W(t)) / t - dWda(t)
u0 = KovW[1] * W(t0)
prob = ODEProblem(f, u0, tspan)
sol_pklin = solve(prob, Tsit5())

#%% Plot results and save to figure4.pdf
## Plot against the scale factor
#p = plot(
#    a,
#    sol_pknl.(a)[:, 1],
#    lw = 2,
#    xlab = "Scale factor, a",
#    ylab = "Comiving Density Parameters",
#    lab = L"\Omega_K: Nonlinear",
#    legend = :topleft,
#    legendfontsize = 12,
#    labelfontsize = 15,
#    ylims = [1e-7, 8e-7],
#    xlims = [0.2, 1.0],
#)
#p = plot!(a, -Ωgrav_pknl.(a), lw = 2, lab = L"-\Omega_W/2: Nonlinear")
#p = plot!(a, sol_pklin.(a)[:, 1], c = 1, lw = 2, ls = :dot, lab = L"\Omega_K: Linear")
#p = plot!(a, -Ωgrav_pklin.(a), c = 2, lw = 2, ls = :dot, lab = L"-\Omega_W/2: Linear")
#p = plot!(a, KovW .* W.(a), ls = :dot, lab = L"K: Linear,~Analytical")
## Plot against the redshift
redshift = 1 ./ a .- 1
p = plot(
    redshift,
    sol_pknl.(a)[:, 1],
    lw = 2,
    xlab = "Redshift, z",
    ylab = "Comiving Density Parameters",
    lab = L"\Omega_K: Nonlinear",
    legend = :topright,
    legendfontsize = 12,
    labelfontsize = 15,
    ylims = [1e-7, 8e-7],
    xlims = [0, 3.0],
)
p = plot!(redshift, -Ωgrav_pknl.(a), lw = 2, lab = L"-\Omega_W/2: Nonlinear")
p = plot!(
    redshift,
    sol_pklin.(a)[:, 1],
    c = 1,
    lw = 2,
    ls = :dot,
    lab = L"\Omega_K: Linear",
)
p = plot!(
    redshift,
    -Ωgrav_pklin.(a),
    c = 2,
    lw = 2,
    ls = :dot,
    lab = L"-\Omega_W/2: Linear",
)
savefig("figure4.pdf")
display(p)

# %% Clean CLASS (the equivalent of the struct_free() in the `main`
# of CLASS. This step is primordial when running in a loop over different
# cosmologies, as you will saturate your memory very fast if you ommit
# it.
cosmo.struct_cleanup()
# If you want to change completely the cosmology, you should also
# clean the arguments, otherwise, if you are simply running on a loop
# of different values for the same parameters, this step is not needed
cosmo.empty()


using FFTW

function split_and_zeropad(a, b)
    a_digits = digits(a)
    b_digits = digits(b)

    return vcat(a_digits, zeros(Int, length(b_digits))),
           vcat(b_digits, zeros(Int, length(a_digits)))
end

# TODO: there is probably a better way to do the carrying operations
#       check biginteger source code.
function vector_to_int(v)
    sum = 0
    for i = 1:length(v)
        sum += 10^(i-1)*v[i]
    end

    return sum
end

function carry!(v::Vector{Int})
    for i = 1:length(v)-1
        digs = digits(v[i])
        v[i+1] += vector_to_int(digs[2:end])
        v[i] = digs[1]
    end

    return v
end

function schonhage_strassen(a, b)
    va, vb = split_and_zeropad(a, b)

    vf = round.(Int, real(ifft(fft(va).*fft(vb))))

    carry!(vf)

    return vector_to_int(vf)
end


#1-dim CAs

#Elementary & totalistic CA with r-nearest neighbor and k states
struct CellularAutomaton

    #user given values
    N::Int #code
    ruleset::Array{Int8,1} #set of rules
    k::Int #number of states
    r::Int #r-nearest neighbor

    #internal variables
    cells::Array{Int8,2}

    function CellularAutomaton(N::Int, init::Array{Int,1}, gen::Int; k::Int=2, r::Int=1)
        ruleset = rule(N, k, r)

        w = length(init)
        cells = zeros(Int8, gen, w)
        cells[1,:] = Array{Int8}(init[:])'

        if k == 2
            #Elementary CA
            mp = reverse(Int8[k^i for i = 0:2r+1])
        else
            #Totalistic CA
            mp = Array{Int8}(ones(2r+1))
        end

        for i = 2:gen
            for j = 1:w
                ind = 0 ; s = 1
                for q = (j-r):(j+r)

                    #Cyclic boundary conditions
                    if q < 1
                        q = w-q
                    end
                    if q > w
                        q = q-w
                    end

                    ind += mp[s]*cells[i-1, q]
                    s += 1
                end
                cells[i,j] = ruleset[ind+1]
            end
        end

        new(N, ruleset, k, r, cells)
    end

    #Quick setup for one seed cell in the center
    function CellularAutomaton(N::Int, gen::Int; kvs...)
        init = Array{Int}(zeros(2*gen))
        init[gen] = 1

        CellularAutomaton(N, init, gen; kvs...)
    end
end

#Parse rule numbering according to the Wolfram code
function rule(n::Int, k=2, r=1)
    if k < 2
        error("number of states (k) must be larger than 1; you gave $k")
    elseif k == 2
        #Elementary CA
        return digits(n, base=k, pad=k^(2r+1))
    else
        #Totalistic CA
        return digits(n, base=k, pad=(2r+1)k-2)
    end
end


"""
    EXP3 Implementation
"""

type EXP3 <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    γ::Float64
    wVec::Vector
    pDist::Categorical

    function EXP3( noOfArms::Integer, γ::Real )
        new( noOfArms,
             0,
             0,
             γ,
             ones(noOfArms),
             Categorical(1/noOfArms*ones(noOfArms))
        )
    end
end

function get_arm_index( agent::EXP3 )
    agent.lastPlayedArm = rand( agent.pDist )
    return agent.lastPlayedArm
end

function update_reward!( agent::EXP3, r::Real )
    agent.noOfSteps     = agent.noOfSteps + 1;
    # Calculate estimated reward
    r_est = r/agent.pDist.p[agent.lastPlayedArm]

    # Update weight of arm
    agent.wVec[agent.lastPlayedArm] = agent.wVec[agent.lastPlayedArm] * exp(agent.γ*r_est/agent.noOfArms)

    # Calculate new probabilties
    p = (1-agent.γ) * agent.wVec/sum(agent.wVec) + agent.γ/agent.noOfArms

    # Make it a distrubution
    agent.pDist = Categorical( p )

    nothing
end

function reset!( agent::EXP3 )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0
    agent.wVec          = ones( agent.noOfArms )
    agent.pDist         = Categorical( 1/agent.noOfArms*ones(agent.noOfArms) )

    nothing
end

function info_str( agent::EXP3, latex::Bool )
    if latex
        return @sprintf( "EXP3 (\$\\gamma = %4.3f\$)", agent.γ )
    else
        return @sprintf( "EXP3 (γ = %4.3f)", agent.γ )
    end
end

"""
    EXP3.1 Implementation
    Based on : Fig.2 in Auer, P., Bianchi, N. C., Freund, Y., & E.Schapire, R. (2002). The Non-Stochastic Multi-Armed bandit problem. SIAM Journal of Computing, 22, 322–331. http://doi.org/10.1109/CDC.1983.269708
"""

type EXP31 <: BanditAlgorithmBase
    G_hat::Vector
    r::Int64
    g_r::Float64

    _EXP3::EXP3

    function EXP31( noOfArms )
        new( zeros( noOfArms ),
             0,
             (noOfArms*log(noOfArms))/(e-1),
             EXP3( noOfArms, 1.00 )  # Initial value for γ will simplify to 1.00
        )
    end
end

function get_arm_index( agent::EXP31 )
    get_arm_index( agent._EXP3 )
end

function update_reward!( agent::EXP31, r::Real )
    # Update G_hat for pulled arm
    agent.G_hat[agent._EXP3.lastPlayedArm] += r/agent._EXP3.pDist.p[agent._EXP3.lastPlayedArm]
    # println( "    G_hat = ", maximum(agent.G_hat), "    LHS = ", (agent.g_r - agent._EXP3.noOfArms/agent._EXP3.γ) )
    if maximum(agent.G_hat) > (agent.g_r - agent._EXP3.noOfArms/agent._EXP3.γ)
        # Calculate new g_r and reset EXP3 with new γ
        agent.r     += 1
        agent.g_r   = (agent._EXP3.noOfArms*log(agent._EXP3.noOfArms))/(e-1) * (4^agent.r)
        # println( "        g_r = ", agent.g_r )
        # println( "        New γ = ", √((agent._EXP3.noOfArms*log(agent._EXP3.noOfArms))/((e-1)*agent.g_r )) )
        agent._EXP3 = EXP3( agent._EXP3.noOfArms,
                            min(
                                1.00,
                                √((agent._EXP3.noOfArms*log(agent._EXP3.noOfArms))/((e-1)*agent.g_r ))
                            )
                        )
    else
        # Update reward to EXP3
        update_reward!( agent._EXP3, r )
    end

    nothing
end

function reset!( agent::EXP31 )
    reset!( agent._EXP3 )
    agent.G_hat = zeros( agent._EXP3.noOfArms)
    agent.r     = 0

    nothing
end

function info_str( agent::EXP31, latex::Bool )
    return @sprintf( "EXP3.1" )
end

"""
    REXP3 Implementation
    Based on Besbes, O., Gur, Y., & Zeevi, A. (2014). Stochastic Multi-Armed-Bandit Problem with Non-stationary Rewards. Advances in Neural Information Processing Systems, 2, 1–9.
"""

type REXP3 <: BanditAlgorithmBase
    noOfSteps::Int64    # Timesteps passed
    j::Int64            # Batch index
    Δ::Int64            # Batch size

    _EXP3::EXP3

    function REXP3( noOfArms::Integer, γ::Real, Δ::Integer )
        new( 0,
             1,
             Δ,
             EXP3(noOfArms,γ)
        )
    end
end

function get_arm_index( agent::REXP3 )
    get_arm_index( agent._EXP3 )
end

function update_reward!( agent::REXP3, r::Real )
    agent.noOfSteps = agent.noOfSteps + 1;
    update_reward!( agent._EXP3, r );
    # Reset if necessary
    if agent._EXP3.noOfSteps >= agent.Δ
        reset!( agent._EXP3 );
        agent.j = agent.j + 1;
    end

    nothing
end

function reset!( agent::REXP3 )
    reset!( agent._EXP3 )
    agent.noOfSteps     = 0;
    agent.j             = 1;

    nothing
end

function info_str( agent::REXP3, latex::Bool )
    if latex
        return @sprintf( "REXP3 (\$\\gamma = %4.3f, \\Delta = %d\$)", agent._EXP3.γ, agent.Δ )
    else
        return @sprintf( "REXP3 (γ = %4.3f, Δ = %d)", agent._EXP3.γ, agent.Δ )
    end
end

"""
    EXP3-IX Implementation
    Based on Neu, G. (2015). Explore no more: improved high-probability regret bounds for non-stochastic bandits. In Advances in Neural Information Processing Systems (pp. 1–9).
"""

type EXP3IX <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    η::Float64
    γ::Float64
    wVec::Vector
    pDist::Categorical

    function EXP3IX( noOfArms::Integer, η::Real, γ::Real )
        new( noOfArms,
             0,
             0,
             η,
             γ,
             ones(noOfArms),
             Categorical(1/noOfArms*ones(noOfArms))
        )
    end
end

function get_arm_index( agent::EXP3IX )
    agent.lastPlayedArm = rand( agent.pDist )
    return agent.lastPlayedArm
end

function update_reward!( agent::EXP3IX, r::Real )
    agent.noOfSteps     = agent.noOfSteps + 1;

    # Calculate loss
    l   = 1 - r

    # Calculate estimated reward
    l_est = l/(agent.pDist.p[agent.lastPlayedArm]+agent.γ)

    # Update weight of arm
    agent.wVec[agent.lastPlayedArm] = agent.wVec[agent.lastPlayedArm] * exp(-agent.η*l_est)

    # Calculate new probabilties
    p = agent.wVec/sum(agent.wVec)

    # Make it a distrubution
    agent.pDist = Categorical( p )

    nothing
end

function reset!( agent::EXP3IX )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0
    agent.wVec          = ones( agent.noOfArms )
    agent.pDist         = Categorical( 1/agent.noOfArms*ones(agent.noOfArms) )

    nothing
end

function info_str( agent::EXP3IX, latex::Bool )
    if latex
        return @sprintf( "EXP3-IX (\$\\eta = %4.3f, \\gamma = %4.3f\$)", agent.η, agent.γ )
    else
        return @sprintf( "EXP3-IX (η = %4.3f, γ = %4.3f)", agent.η, agent.γ )
    end
end


# Adapted from:
# SOSDEMO6 --- MAX CUT
# Section 3.6 of SOSTOOLS User's Manual

using Test
using SumOfSquares
using DynamicPolynomials

function maxcut_test(optimizer, config::MOIT.Config)
    atol = config.atol
    rtol = config.rtol

    @polyvar x[1:5]

    # Number of cuts
    f = 2.5 - 0.5*x[1]*x[2] - 0.5*x[2]*x[3] - 0.5*x[3]*x[4] - 0.5*x[4]*x[5] - 0.5*x[5]*x[1]

    # Boolean constraints
    bc = vec(x).^2 .- 1

    model = _model(optimizer)

    Z = monomials(x, 0:1)
    @variable(model, p1, SOSPoly(Z))

    Z = monomials(x, 0:2)
    @variable(model, p[1:5], Poly(Z))

    @testset "with γ=$γ it should be $(feasible ? "feasible" : "infeasible")" for (γ, feasible) in [(3.9, false), (4.1, true)]
        cref = @constraint(model, p1*(γ-f) + dot(p, bc) - (γ-f)^2 in SOSCone())

        JuMP.optimize!(model)

        if feasible
            @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
        else
            @test JuMP.dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE
        end
        delete(model, cref)
    end
end
sd_tests["maxcut"] = maxcut_test


using Random
using Statistics
using StatsBase
using Plots
import Base.+
import Base.-

include("Point_Struct.jl")
include("fun.jl")
include("constants.jl")

# This script reproduce the turbulance distribution function (a gaussian)
# Figure 4 in Errard's article

p_0 = [Point{Float64}(20000, 0.9, i, undef, undef, undef)() for i = (1:25000)*( (2*π) / 25000  )  ] # must be constant
p_1 = [Point{Float64}(20000, i, 0.9, undef, undef, undef)() for i = (1:25000)*( (2*π) / 25000 )  ] # must be constant
p_2 = [Point{Float64}(i, 0.9, 0.9, undef, undef, undef)() for i = (10:1000000)/1000000 * 100000  ] # must be constant

w   = [ω(o) for o in p_2]
BB = [B(i) for i in p_2]

z   = [i.z for i in p_2]

plot(  BB[2:length(BB)] / (2*π*(ω₀^2)) , z[2:length(z)], yaxis=:log)#, xlims=(0,1050))
#plot( w[2:length(z)] .* ( (ω₀ * π) / ( λ * 100000 )), z[2:length(z)])#, yaxis=:log)#, xlims=(0,1050))


import Base: sin, cos, tan, cot, sec, csc
import Base: real, imag, reim, ==
import Base: asin, acos, atan
import Base: asinh, acosh, atanh
import Base: sinh, cosh, tanh, coth, sech, csch
import Base: abs, exp, log, zero
import Base: +, -, *, /, ^

## By constructing bicomplex, we can represent multi-complex numbers using recursion
struct bicomplex
    i1
    i2
end

"""
Using matrix to represent bicomplex.
"""
function mat(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    return [i1 -i2; i2 i1]
end

function mat2bicomplex(matrix)
    out1, out2 = matrix[1, 1], matrix[2, 1]
    return bicomplex(out1, out2)
end

reim(bi::bicomplex) = (bi.i1, bi.i2)

==(a::bicomplex, b::bicomplex) = a.i1 == b.i1 && a.i2 == b.i2

## Basic operators: +, -, *, /, ^

+(bi1::bicomplex, bi2::bicomplex) = mat2bicomplex(mat(bi1)+mat(bi2))
+(bi::bicomplex, n) = mat2bicomplex(n.+mat(bi))
+(n, bi::bicomplex) = +(bi, n)

-(bi1::bicomplex, bi2::bicomplex) = mat2bicomplex(mat(bi1)-mat(bi2))
-(bi::bicomplex, n) = mat2bicomplex(mat(bi).-n)
-(n, bi::bicomplex) = -(bi, n)
-(bi::bicomplex) = mat2bicomplex(-mat(bi))

*(bi1::bicomplex, bi2::bicomplex) = mat2bicomplex(mat(bi1)*mat(bi2))
*(n, bi::bicomplex) = mat2bicomplex(n*mat(bi))
*(bi::bicomplex, n) = *(n, bi)

/(bi1::bicomplex, bi2::bicomplex) = mat2bicomplex(mat(bi1)/mat(bi2))
/(n, bi::bicomplex) = mat2bicomplex(n*inv(mat(bi)))
/(bi::bicomplex, n) = mat2bicomplex(mat(bi)/n)

^(bi::bicomplex, power) = mat2bicomplex(mat(bi)^power)

zero(bi::bicomplex) = bicomplex(zero(bi.i1), zero(bi.i2))

## Basic functions: abs, exp, log

abs(bi::bicomplex) = sqrt(bi.i1^2 + bi.i2^2)

function exp(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    out1 = exp(i1)*cos(i2)
    out2 = exp(i1)*sin(i2)
    return bicomplex(out1, out2)
end

function log(bi::bicomplex)
    out1 = log(abs(bi))
    out2 = angle(bi)
    return bicomplex(out1, out2)
end

angle(bi::bicomplex) = atan(bi.i2/bi.i1)

## Basic math functions: sin, cos, tan, cot, sec, csc

function sin(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    out1 = cosh(i2)*sin(i1)
    out2 = sinh(i2)*cos(i1)
    return bicomplex(out1, out2)
end

function cos(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    out1 = cosh(i2)*cos(i1)
    out2 = -sinh(i2)*sin(i1)
    return bicomplex(out1, out2)
end

tan(bi::bicomplex) = sin(bi)/cos(bi)
cot(bi::bicomplex) = cos(bi)/sin(bi)
sec(bi::bicomplex) = 1/cos(bi)
csc(bi::bicomplex) = 1/sin(bi)

## Basic inverse hyperbolic trig function: asinh, acosh, atanh, acoth, asech, acsch

function atan(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    temp = atanh(bicomplex(-i2, i1))
    return bicomplex(temp.i2, temp.i1)
end

## Basic hyperbolic functions: sinh, cosh, tanh, coth, sech, csch

function sinh(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    out1 = cosh(i1)*cos(i2)
    out2 = sinh(i1)*sin(i2)
    return bicomplex(out1, out2)
end

function cosh(bi::bicomplex)
    i1, i2 = bi.i1, bi.i2
    out1 = sinh(i1)*cos(i2)
    out2 = cosh(i1)*sin(i2)
    return bicomplex(out1, out2)
end

tanh(bi::bicomplex) = sinh(bi)/cosh(bi)
coth(bi::bicomplex) = cosh(bi)/sinh(bi)
sech(bi::bicomplex) = 1/cosh(bi)
csch(bi::bicomplex) = 1/sinh(bi)

image2(bi::bicomplex) = imag(bi.i2)

biderivative(f, point, h) = image2(f(bicomplex(point+im*h, h)))/h^2


path = @get_scratch!("rocket")

# rocket
nx = rocket.nq
nu = rocket.nu
nw = rocket.nw

nz = nx
nθ = nx + nu + 1

@variables z[1:nz] θ[1:nθ] κ[1:1]

# Residual
function residual(rocket, z, θ, κ)
    y = z[1:nx]
    x = θ[1:nx]
    u = θ[nx .+ (1:nu)]
    h = θ[nx + nu .+ (1:1)]
    w = zeros(rocket.nw)

    return y - (x + h[1] * f(rocket, 0.5 * (x + y), u, w))
end

r_rocket = residual(rocket, z, θ, κ)
r_rocket = Symbolics.simplify.(r_rocket)
rz_rocket = Symbolics.jacobian(r_rocket, z, simplify=true)
rθ_rocket = Symbolics.jacobian(r_rocket, θ, simplify=true)

# Build function
r_rocket_func = build_function(r_rocket, z, θ, κ)[2]
rz_rocket_func = build_function(rz_rocket, z, θ)[2]
rθ_rocket_func = build_function(rθ_rocket, z, θ)[2]

rz_rocket_array = similar(rz_rocket, Float64)
rθ_rocket_array = similar(rθ_rocket, Float64)

@save joinpath(path, "residual.jld2") r_rocket_func rz_rocket_func rθ_rocket_func rz_rocket_array rθ_rocket_array
@load joinpath(path, "residual.jld2") r_rocket_func rz_rocket_func rθ_rocket_func rz_rocket_array rθ_rocket_array

# projection
nz = 3 + 1 + 1 + 1 + 1 + 3
nθ = 3 + 1

function residual_projection(z, θ, κ)
    u = z[1:3]
    p = z[4:4]
    s = z[5:5]
    w = z[6:6]
    y = z[7:7]
    v = z[8:10]

    ū = θ[1:3]
    uu = θ[3 .+ (1:1)]
    idx = [3; 1; 2]
    [
     u - ū - v - [0.0; 0.0; y[1] + p[1]];
     uu[1] - u[3] - s[1];
     -y[1] - w[1];
     w[1] * s[1] - κ[1];
     p[1] * u[3] - κ[1];
     cone_product(u[idx], v[idx]) - [κ[1]; 0.0; 0.0]
    ]
end

@variables z[1:nz], θ[1:nθ], κ[1:1]

r_proj = residual_projection(z, θ, κ)
r_proj .= simplify(r_proj)
r_proj_func = Symbolics.build_function(r_proj, z, θ, κ)[2]

rz_proj = Symbolics.jacobian(r_proj, z)
rz_proj = simplify.(rz_proj)
rz_proj_func = Symbolics.build_function(rz_proj, z, θ)[2]
rθ_proj = Symbolics.jacobian(r_proj, θ)
rθ_proj = simplify.(rθ_proj)
rθ_proj_func = Symbolics.build_function(rθ_proj, z, θ)[2]

rz_proj_array = similar(rz_proj, Float64)
rθ_proj_array = similar(rθ_proj, Float64)

@save joinpath(path, "projection.jld2") r_proj_func rz_proj_func rθ_proj_func rz_proj_array rθ_proj_array
@load joinpath(path, "projection.jld2") r_proj_func rz_proj_func rθ_proj_func rz_proj_array rθ_proj_array


@testset "Test Saving and Loading PLS2 Models" begin

	Xtr        = [1 2;2 4;3 6;6 12;7 14.0]
	Ytr        = [2 2;4 4;6 6;12 12;14 14.0]
	Xt         = [4 8;5 10.0]
	model1    = PLSRegressor.fit(Xtr,Ytr,nfactors=2)
	pred1     = PLSRegressor.predict(model1,Xt)

	PLSRegressor.save(model1)
	model2    = PLSRegressor.load()

	pred2     = PLSRegressor.predict(model2,Xt)
    rm(PLSRegressor.MODEL_FILENAME)
	@test all(pred1 .== pred2)

end

@testset "PLS2 Prediction Tests (in sample)" begin

	@testset "Single Column Prediction Test" begin

		X        = [1; 2; 3.0]
		Y        = [1 1; 2 2; 3 3.0]
		model    = PLSRegressor.fit(X,Y,nfactors=1)
		pred     = PLSRegressor.predict(model,X)

		@test isequal(round.(pred),[1 1; 2 2; 3 3.0])

	end

	@testset "Constant Values Prediction Tests (Ax + b) | A=0, b=1 " begin

		X        = [1 3;2 1;3 2.0]
		Y        = [1 1; 1 1; 1 1.0]
		try
			PLSRegressor.fit(X,Y,nfactors=2)
		catch
			@test true
		end

	end

	@testset "Linear Prediction Tests " begin

		X        = [1 2; 2 4; 4 6.0]
		Y        = [4 2;6 4;8 6.0]
		model    = PLSRegressor.fit(X,Y,nfactors=2)
		pred     = PLSRegressor.predict(model,X)
		@test isequal(round.(pred),[4 2;6 4;8 6.0])

		X           = [1 -2; 2 -4; 4 -6.0]
		Y           = [-4 -2;-6 -4;-8 -6.0]
		model       = PLSRegressor.fit(X,Y,nfactors=2)
		pred        = PLSRegressor.predict(model,X)
		@test isequal(round.(pred),[-4 -2;-6 -4;-8 -6.0])

	end

end

@testset "PLS2 Pediction Tests (out of sample)" begin

	@testset "Linear Prediction Tests (Ax + b) | A>0" begin

		Xtr        = [1 2;2 4;3 6;6 12;7 14.0]
		Ytr        = [2 2;4 4;6 6;12 12;14 14.0]
		Xt         = [4 8;5 10.0]
		model    = PLSRegressor.fit(Xtr,Ytr,nfactors=2)
		pred     = PLSRegressor.predict(model,Xt)
		@test isequal(round.(pred),[8 8;10 10.0])

		Xtr        = [1 2;2 4;3 6;6 12;7 14.0]
		Ytr        = [2 4;4 6;6 8;12 14;14 16.0]
		Xt         = [4 8;5 10.0]
		model    = PLSRegressor.fit(Xtr,Ytr,nfactors=2)
		pred     = PLSRegressor.predict(model,Xt)
		@test isequal(round.(pred),[8 10;10 12.0])

	end

	@testset "Linear Prediction Tests (Ax + b) | A<0" begin

		Xtr        = [1 -2;2 -4;3 -6;6 -12;7 -14.0]
		Ytr        = [2 -2;4 -4;6 -6;12 -12;14 -14.0]
		Xt         = [4 -8;5 -10.0]
		model    = PLSRegressor.fit(Xtr,Ytr,nfactors=2)
		pred     = PLSRegressor.predict(model,Xt)
		@test isequal(round.(pred),[8 -8;10 -10.0])

		Xtr        = [1 -2;2 -4;3 -6;6 -12;7 -14.0]
		Ytr        = [2 -4;4 -6;6 -8;12 -14;14 -16.0]
		Xt         = [4 -8;5 -10.0]
		model    = PLSRegressor.fit(Xtr,Ytr,nfactors=2)
		pred     = PLSRegressor.predict(model,Xt)
		@test isequal(round.(pred),[8 -10;10 -12.0])

	end

end;


using IrrationalExpressions
using Base.Test

@test BigFloat(-π) == -BigFloat(π)
@test BigFloat(-(-π)) == BigFloat(π)
@test BigFloat(2π+1//3) == BigFloat(2)*BigFloat(π)+BigFloat(1//3)

@test isa(π, Irrational)
@test isa(π+1, IrrationalExpressions.IrrationalExpr)
@test isa(π+1.0, Float64)
@test isa(2π+1.0, Float64)
@test isa(π+BigFloat(1), BigFloat)
@test +(1+π) == 1+π
@test string(1+pi) == "1 + π = 4.1415926535897..."
@test string(2π+0.0) == "6.283185307179586"
@test string(2π+BigFloat(0)) == "6.283185307179586476925286766559005768394338798750211641949889184615632812572396"


# %%
using Tmatrix

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

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

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

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

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

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

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

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

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

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

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


"""
    glaser_liu_rokhlin_gauss_nodes(π::P, x0=pqr_start(P,degree(π));
                                         symmetry=pqr_symmetry(P),
                                         m=5)

Find Gauss nodes and weights of the  basis  polynomial  `π = basis(P,n)`. The nodes  are the  roots of the  polynomial and the weights are computed  from known formulas.

The method from  [Glaser,  Liu, and Rokhlin](DOI: 10.1137/06067016X) is  used. This is  an  O(n)  method (whereas the method  basedon the Jacobi  matrix is O(n^2)). This method fits easily  into the framework  provided  through  the `AbstractCCOP` types.  The [FastGaussQuadrature](https://github.com/JuliaApproximation/FastGaussQuadrature.jl) package provides even more efficient  algorithms and  pays attention to numerical  issues, examples of  which  can  be found in  [Hale and Townsend](https://core.ac.uk/download/pdf/9403469.pdf). The `FastGuassQuadrature` package is used when  available.

An example:

```
function  gauss_nodes_weights(P::Type{<:Jacobi{α,β}}, n)  where {α,β}
    # we don't  have a  good starting point  unless α=β
    if α == β
        xs, ws = glaser_liu_rokhlin_gauss_nodes(basis(MonicJacobi{α,β},n))
        λ = kn(P,n)^2
        xs, ws/λ
    else
        J = jacobi_matrix(P, n)
        eig = eigen(J, extrema(P)...)
        wts =  π̃βn(P,0) * (eig.vectors[1,:]).^2
        eig.values,  wts
    end

end

pqr_start(P::Type{MonicJacobi{α,α}}, n) where {α} =  0.0 # zero(eltype(P))
pqr_symmetry(P::Type{<:MonicJacobi{α,α}}) where {α} = true
function pqr_weight(P::Type{<:MonicJacobi{α,α}}, n, x, dπx) where {α}
    # https://www.chebfun.org/publications/HaleTownsend2013a.pdf
    # Eqn  (1.4)
    # Cnαβ should be  computed  using asymptotic formula  for larger n (§3.2.3)
    # XXX TThere i ss ome constant  that  makes this not work....
    β = α
    Cnαβ = 2^(α+β+1)
    Cnαβ *= gamma(n + α +  1) * gamma(n + β + 1) / gamma(n + α + β + 1)
    Cnαβ /= gamma(1 + n)
    val = Cnαβ / (1-x^2) / dπx^2
    val
end
```

"""
glaser_liu_rokhlin_gauss_nodes()

"""
    pqr

Compute  `p,q,r,p',q',r'` where `p(x)y'' + q(x)y' + r(x)  = 0`. Uses
the `a,b,c,d,e` values characterizing the  system.

"""
function pqr(P::Type{<:AbstractCCOP}, n, x)
    a,b,c,d,e = abcde(P)

    p =  a*x^2 + b*x  +  c
    dp = 2a*x + b
    q = d*x +  e
    dq = d
    r  = -(a*n*(n-1)  + d*n)
    dr = 0
    (p,q,r,dp,dq,dr)
end

"""
    pqr_start(p::P, n)

A starting  value for finding the gauss nodes
"""
pqr_start(::Type{P}, n) where {P <: AbstractCCOP} = 0

"""
    pqr_symmetry(p::P)

Boolean to specify if symmetry should apply to output
"""
pqr_symmetry(::Type{P}) where {P <: AbstractCCOP} = false

"""
    pqr_weight(p::P, x, dx)

Compute weight from x, dπx, x a node and dπx the derivative's value at the node
"""
pqr_weight(p::P, n, x, dx) where {P} = throw(MethodError())

# We just use clenshaw to  evaluate  `p(x), dp(x)` for Newton's method below,
# but implementing the following can be more efficient
# ## Evaluate basis(P,n)  and its  derivative using the three-point recursion representation
# function orthogonal_polyval_derivative(p::P, x::S) where {P <: AbstractOrthogonalPolynomial, S}

#     T = eltype(p)
#     oS = one(x)
#     R = eltype(one(T) * oS)

#     length(p) == 0 && return (zero(R), zero(R))
#     d = length(p) - 1

#     # P0,  P_{-1} case
#     pn,  pn_1 = P0(P,x), zero(x)
#     dpn, dpn_1  = dP0(P,x), zero(x)

#     ptot = p[0]*P0(P,x)
#     dptot  = p[0]*dP0(P,x)

#     ## we compute  forward here, finding pi, dp_i,  p_{i+1}, dp_{i+1}, ...
#     for i in 0:d-1
#         an, bn, cn=  An(p,i), Bn(p,i), Cn(p,i)
#         dpn_1, dpn = dpn,  an*pn + (an*x + bn)*dpn + cn*dpn_1
#         pn_1,  pn =  pn, (an*x + bn)*pn + cn*pn_1
#         ptot += p[i+1]*pn
#         dptot += p[i+1]*dpn
#     end
#     return (ptot, dptot)
# end

# run Newton's method to find zero of p
function newton(p::P, x0::S) where  {P <: AbstractCCOP, S}
    maxsteps = 25
    dp  = derivative(p)
    while maxsteps > 0
        px, dpx = p(x0),  dp(x0) #orthogonal_polyval_derivative(p, x0)
        Δ = px/dpx
        isnan(Δ) && return (x0, dpx)
        if isinf(Δ)
            x0 += sqrt(eps())  # give a nudge from minimum or maximum
            continue
        end
        x0 -= Δ
        min(abs(Δ),abs(px)) <= sqrt(eps(S))/100 && return  x0, dp(x0)# orthogonal_polyval_derivative(p, x0)[2]
        maxsteps -= 1
    end
    @warn "Newton's method did not converge from $x0"
    NaN
end

# Second-order Runge-Kutta with global truncation error of h^2
function RK(t0,  x0, F, h, n)
    k0 = h*F(t0,x0)
    local x1
    for i in 1:n
        t1 = t0 + h
        k1 = h * F(t0+h,x0+k0)
        x1 = x0 + 1/2 * (k0+k1)
        t0, x0, k0 = t1, x1, k1
    end
    x1
end

function prufer(Π::Type{P},n) where {P <: AbstractCCOP}
    (θ, x) -> begin
        dom = domain(Π)
        a, b = first(dom)+eps(), last(dom)-eps()
        x = clamp(x, a, b)
        p,q,r,dp,dq,dr = pqr(Π, n, x)
        -inv(sqrt(r/p) + (dr*p - dp*r + 2r*q)/(2r*p) * sin(2θ)/2)
    end
end

## Compute gauss node using  algorithm from
## A Fast Algorithm for the Calculation of the Roots of Special Functions
## by Glaser, Liu, Rokhlin
## DOI: 10.1137/06067016X
## specialized to the orthogonal polynomial case
function glaser_liu_rokhlin_gauss_nodes(π::P, x0=pqr_start(P,degree(π));
                                        symmetry=pqr_symmetry(P),
                                        m=5) where {P <: AbstractCCOP}

    n = degree(π)
    F = prufer(P, n)
    x = float(x0)

    # step 1 to get initial might be newton or might be RK
    if symmetry && iseven(n)
        # a max at π(x0)
        x = RK(0, x0, F, -(pi/2)/m, m)
    end
    x1, dπ1 = newton(π, x)
    rts, dπrts = [x1], [dπ1]
    N = symmetry ? div(n+1,2) : n
    for i in 2:N
        xx = RK(pi/2, rts[end], F, -pi/m, m)
        x, dπx = newton(π, xx)
        push!(rts, x)
        push!(dπrts, dπx)
    end
    if symmetry
        if iseven(n)
            rts = vcat(-reverse(rts), rts)
            dπrts = vcat(-reverse(dπrts), dπrts)
        else
            rts = vcat(-reverse(rts[2:end]), rts)
            dπrts = vcat(-reverse(dπrts[2:end]), dπrts)
        end
    end
    # get weights
    weights = pqr_weight.(P, n, rts, dπrts)
    rts,  weights
end


"""
    valueiterationsolver(w::FHWrapper, stage::Int64, util)

Perform one iteration of Value Iteration
"""
function valueiterationsolver(m::MDP, stage::Int64, util)
    next_stage_value = util               # maximum value in each row
    stage_q = zeros(length(stage_states(m, stage)), length(actions(m)))

    for s in stage_states(m, stage)
        isterminal(m, s) && continue

        for a in actions(m)
            si = stage_stateindex(m, s)
            ai = actionindex(m, a)
            for (sp, p) in weighted_iterator(transition(m, s, a))
                spi = stage_stateindex(m, sp)
                stage_q[si, ai] += p * (reward(m, s, a, sp) + discount(m) * next_stage_value[spi])
            end
        end
    end

    util = vec(maximum(stage_q; dims=2))
    pol = vec([i[2] for i in findmax(stage_q; dims=2)[2]])

    return stage_q, util, pol
end


# This is an example script demonstrating how PARDISO works on a
# medium-sized Hermitian positive definite matrix.
using Pardiso
using LinearAlgebra # for norm
using SparseArrays
using Random

# Script parameters.
# -----------------
verbose = false
n       = 100
lambda  = 3

# Create the Hermitian positive definite matrix A and the vector b in the
# linear system Ax = b.
e = ones(n)
e2 = ones(n-1)
A = spdiagm(-1 => im*e2, 0 => lambda*e, 1 => -im*e2)
b = rand(n) + im * zeros(n)

# Initialize the PARDISO internal data structures.
# ps = PardisoSolver()
ps = MKLPardisoSolver()

if verbose
    set_msglvl!(ps, Pardiso.MESSAGE_LEVEL_ON)
end

# If we want, we could just solve the system right now.
# Pardiso.jl will automatically detect the correct matrix type,
# solve the system and free the data
X1 = solve(ps, A, b)

# First set the matrix type to handle general complex
# hermitian positive definite matrices
set_matrixtype!(ps, Pardiso.COMPLEX_HERM_POSDEF)

# Initialize the default settings with the current matrix type
pardisoinit(ps)

# Get the correct matrix to be sent into the pardiso function.
# :N for normal matrix, :T for transpose, :C for conjugate
A_pardiso = get_matrix(ps, A, :N)

# Analyze the matrix and compute a symbolic factorization.
set_phase!(ps, Pardiso.ANALYSIS)
pardiso(ps, A_pardiso, b)
@printf("The factors have %d nonzero entries.\n", get_iparm(ps, 18))

# Compute the numeric factorization.
set_phase!(ps, Pardiso.NUM_FACT)
pardiso(ps, A_pardiso, b)

# Compute the solutions X using the symbolic factorization.
set_phase!(ps, Pardiso.SOLVE_ITERATIVE_REFINE)
x = similar(b) # Solution is stored in X
pardiso(ps, x, A_pardiso, b)
@printf("PARDISO performed %d iterative refinement steps.\n", get_iparm(ps, 7))

# Compute the residuals.
r = abs.(A*x - b)
@printf("The maximum residual for the solution is %0.3g.\n",maximum(r))
@test norm(R) < 1e-10

# Free the PARDISO data structures.
set_phase!(ps, Pardiso.RELEASE_ALL)
pardiso(ps)