using LinearAlgebra
"""
        (V,H,hist)=iar(prob,startvec,N)

### Infinite Arnoldi method
Infinite Arnoldi method, as described in Algorithm 2 in  "A linear eigenvalue algorithm for the nonlinear eigenvalue problem",
    by Jarlebring, Elias and Michiels, Wim and Meerbergen, Karl.

Runs `N` steps of the infinite Arnoldi method started with `startvec` with target `σ=0`.
The argument `prob` represents a problem `M(λ)` and has associated functions
`compute_Mder`, `compute_Mlincomb`:

* `M=compute_Mder(prob,s)` computes the matrix function value at `s=λ`.
* `z=compute_Mlincomb(prob,s,X)` computes a linear combination of derivatives. The linear combination coefficients are given in the matrix `X`:  `z=M(s)*X[:,1]+M'(s)*X[:,2]+...+M^{(k)}(s)*X[:,k+1]`.

"""
function iar(prob,startvec,N,precomputed_factorization=nothing)

    T=ComplexF64; # Problem type

    n=length(startvec);

    sigma=0;  # Only target implemented

    x0=zeros(T,n);
    x0[:]=normalize(startvec);
    H=zeros(T,N+1,N);
    print("IAR Iteration:");
    V=zeros(T,n*(N+1),N+1);
    V[1:n,1]=x0;

    start_time=time();
    hist=Dict(:time_count => NaN*zeros(N))

    if (isnothing(precomputed_factorization))
        Msolve=factorize(compute_Mder(prob,sigma));
    else
        Msolve=precomputed_factorization;
    end

    for k=1:N
        print("$k ")
        # Create L matrix: Equation (17) in nummath paper
        L=Diagonal(1 ./(1:k));

        vecx=V[1:(k*n),k];
        X=reshape(vecx,n,k);

        # Compute y1,..yn: Equation (13) in nummath paper
        Y=zeros(T,size(X,1),size(X,2)+1);
        Y[:,2:end]=X*L;

        # Compute y0: Equation (18) in nummath paper
        # y0=nep.M_lin_comb(X,Y,sigma);
        if (true)
            # Note: First col of Y is zero
            z0=compute_Mlincomb(prob,sigma,Y);
            #z0=Mlinprob.M_lin_comb(Y(:,2:end),1);
        else
            z0=compute_Mlincomb(prob,sigma,X[:,1:k]);
        end

        y0=-(Msolve\z0);
        Y[:,1]=y0;
        vecY=reshape(Y,n*(k+1),1);

        if (true)
            kk=k+1;
        else
            kk=1; # short vector orthogonalization
        end
        V0=V[1:(kk*n),1:k];
        h=V0'*vecY[1:(kk*n)];
        vecY=vecY-V[1:((k+1)*n),1:k]*h;

        g=V0'*vecY[1:(kk*n)];   # twice to be sure
        vecY=vecY-V[1:((k+1)*n),1:k]*g;
        h=h+g;

        beta=norm(vecY);  # normalize
        vecY=vecY/beta;
        H[1:(k+1),k]=[h;beta];
        V[1:((k+1)*n),k+1]=vecY;

        hist[:time_count][k]=time()-start_time;

    end

    println("done");

    return (V,H,hist);

end


"""
    BSplineKit.BoundaryConditions

Contains some boundary condition definitions.
"""
module BoundaryConditions

export Natural

abstract type BoundaryCondition end

"""
    Natural

Generalised natural boundary condition.

This boundary condition is convenient for spline interpolations, as it
provides extra constraints enabling to equate the number of unique B-spline
knots to the number of data points.

For cubic splines (order ``k = 4``), this corresponds to [natural cubic
splines](https://en.wikipedia.org/wiki/Spline_(mathematics)#Examples), imposing
the second derivatives to be zero at the boundaries (``S''(a) = S''(b) = 0``).

For higher-order splines, this boundary condition generalises the standard
natural cubic splines, by setting derivatives of order ``2, 3, …, k/2`` to
be zero at the boundaries.
For instance, for ``k = 6`` (quintic splines), this imposes ``S'' = S''' = 0``.
In practice, BSplineKit.jl achieves this by using [basis recombination](@ref
basis-recombination-api).

Note that, for symmetry reasons, only even-order splines are supported by this BC.
"""
struct Natural <: BoundaryCondition end

end


using Logging
using JuMP, Gurobi
using DecisionProgramming
using CSV, DataFrames, PrettyTables

# Setting subproblem specific parameters
const chosen_risk_level = "12%"

# Reading tests' technical performance data (dummy data in this case)
data = CSV.read("CHD_preventative_care_data.csv", DataFrame)

# Bayes posterior risk probabilities calculation function
# prior = prior risk level for which the posterior risk distribution is calculated for,
# t = test done
# returns a 100x1 vector with the probabilities of getting CHD given the prior risk level and test result
# for no test done (i.e. t = 3) returns a zero vector
function update_risk_distribution(prior::Int64, t::Int64)
    if t == 1 # the test is TRS
        # P(TRS = result | sick) = P(TRS_if_sick = result) * P(sick) = P(TRS_if_sick = result) * P(prior_risk)
        numerators = data.TRS_if_sick .* data.risk_levels[prior]

        # P(TRS = result) = P(TRS_if_sick = result) * P(sick) + P(TRS_if_healthy = result) * P(healthy)
        denominators = data.TRS_if_sick .* data.risk_levels[prior]  + data.TRS_if_healthy .* (1 - data.risk_levels[prior])

        posterior_risks = numerators./denominators

        # if the denominator is zero, post_risk is NaN, changing those to 0
        for i = 1:101
            if isnan(posterior_risks[i])
                posterior_risks[i] = 0
            end
        end

        return posterior_risks

    elseif t == 2 #the test is GRS
        numerators = (data.GRS_if_sick .* data.risk_levels[prior])
        denominators = data.GRS_if_sick .* data.risk_levels[prior]  + data.GRS_if_healthy .* (1 .- data.risk_levels[prior])

        posterior_risks =  numerators./denominators

        # if the denominator is zero, post_risk is NaN, changing those to 0
        for i = 1:101
            if isnan(posterior_risks[i])
                posterior_risks[i] = 0
            end
        end

        return posterior_risks

    else # no test performed
        risks_unchanged = zeros(100,1)

        return risks_unchanged

    end
end

# State probabilites calculation function
# risk_p = the resulting array from update_risk_distribution
# t = test done
# h = CHD or no CHD
# returns the probability distribution in 101x1 vector for the states of the R node given the prior risk level (must be same as to function update_risk_distribution), test t and health h
function state_probabilities(risk_p::Array{Float64}, t::Int64, h::Int64, prior::Int64)

    #if no test is performed, then the probabilities of moving to states (other than the prior risk level) are 0 and to the prior risk element is 1
    if t == 3
        state_probabilites = zeros(101)
        state_probabilites[prior] = 1.0
        return state_probabilites
    end

    # return vector
    state_probabilites = zeros(101)

    # copying the probabilities of the scores for ease of readability
    if h == 1 && t == 1    # CHD and TRS
        p_scores = data.TRS_if_sick
    elseif t ==1    # no CHD and TRS
        p_scores = data.TRS_if_healthy
    elseif h == 1 && t == 2 # CHD and GRS
        p_scores = data.GRS_if_sick
    else # no CHD and GRS
        p_scores = data.GRS_if_healthy
    end

    for i = 1:101 #iterating through all risk levels 0%, 1%, ..., 99%, 100% in data.risk_levels
        for j = 1:101 #iterates through all risk estimates in risk_p
            #finding all risk estimates risk_p[j] within risk level i
            # risk_level[i] <= risk_p < risk_level[i]
            if i < 101 && data.risk_levels[i] <= risk_p[j] && risk_p[j] < data.risk_levels[i+1]
                state_probabilites[i] += p_scores[j]
            elseif i == 101 && data.risk_levels[i] <= risk_p[j] #special case: the highest risk level[101] = 100%
                state_probabilites[i] += p_scores[j]
            end
        end
    end

    return state_probabilites
end

@info("Creating the influence diagram.")
diagram = InfluenceDiagram()

const H_states = ["CHD", "no CHD"]
const T_states = ["TRS", "GRS", "no test"]
const TD_states = ["treatment", "no treatment"]
const R_states = [string(x) * "%" for x in [0:1:100;]]

add_node!(diagram, ChanceNode("R0", [], R_states))
add_node!(diagram, ChanceNode("R1", ["R0", "H", "T1"], R_states))
add_node!(diagram, ChanceNode("R2", ["R1", "H", "T2"], R_states))
add_node!(diagram, ChanceNode("H", ["R0"], H_states))

add_node!(diagram, DecisionNode("T1", ["R0"], T_states))
add_node!(diagram, DecisionNode("T2", ["R1"], T_states))
add_node!(diagram, DecisionNode("TD", ["R2"], TD_states))

add_node!(diagram, ValueNode("TC", ["T1", "T2"]))
add_node!(diagram, ValueNode("HB", ["H", "TD"]))

generate_arcs!(diagram)

X_R0 = ProbabilityMatrix(diagram, "R0")
X_R0[chosen_risk_level] = 1
add_probabilities!(diagram, "R0", X_R0)

X_H = ProbabilityMatrix(diagram, "H")
X_H[:, "CHD"] = data.risk_levels
X_H[:, "no CHD"] = 1 .- data.risk_levels
add_probabilities!(diagram, "H", X_H)

X_R = ProbabilityMatrix(diagram, "R1")
for s_R0 = 1:101, s_H = 1:2, s_T1 = 1:3
    X_R[s_R0, s_H,  s_T1, :] =  state_probabilities(update_risk_distribution(s_R0, s_T1), s_T1, s_H, s_R0)
end
add_probabilities!(diagram, "R1", X_R)
add_probabilities!(diagram, "R2", X_R)

cost_TRS = -0.0034645
cost_GRS = -0.004
forbidden = 0     #the cost of forbidden test combinations is negligible
Y_TC = UtilityMatrix(diagram, "TC")
Y_TC["TRS", "TRS"] = forbidden
Y_TC["TRS", "GRS"] = cost_TRS + cost_GRS
Y_TC["TRS", "no test"] = cost_TRS
Y_TC["GRS", "TRS"] = cost_TRS + cost_GRS
Y_TC["GRS", "GRS"] = forbidden
Y_TC["GRS", "no test"] = cost_GRS
Y_TC["no test", "TRS"] = cost_TRS
Y_TC["no test", "GRS"] = cost_GRS
Y_TC["no test", "no test"] = 0
add_utilities!(diagram, "TC", Y_TC)

Y_HB = UtilityMatrix(diagram, "HB")
Y_HB["CHD", "treatment"] = 6.89713671259061
Y_HB["CHD", "no treatment"] = 6.65436854256236
Y_HB["no CHD", "treatment"] = 7.64528451705134
Y_HB["no CHD", "no treatment"] = 7.70088349200034
add_utilities!(diagram, "HB", Y_HB)

generate_diagram!(diagram)

@info("Creating the decision model.")
model = Model()
z = DecisionVariables(model, diagram)

# Defining forbidden paths to include all those where a test is repeated twice
forbidden_tests = ForbiddenPath(diagram, ["T1","T2"], [("TRS", "TRS"),("GRS", "GRS"),("no test", "TRS"), ("no test", "GRS")])
fixed_R0 = FixedPath(diagram, Dict("R0" => chosen_risk_level))
scale_factor = 10000.0
x_s = PathCompatibilityVariables(model, diagram, z; fixed = fixed_R0, forbidden_paths = [forbidden_tests], probability_cut=false)

EV = expected_value(model, diagram, x_s, probability_scale_factor = scale_factor)
@objective(model, Max, EV)

@info("Starting the optimization process.")
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "MIPFocus" => 3,
    "MIPGap" => 1e-6,
)
set_optimizer(model, optimizer)

optimize!(model)

@info("Extracting results.")
Z = DecisionStrategy(z)
S_probabilities = StateProbabilities(diagram, Z)
U_distribution = UtilityDistribution(diagram, Z)

@info("Printing decision strategy using tailor made function:")
print_decision_strategy(diagram, Z, S_probabilities)

@info("Printing state probabilities:")
# Here we can see that the probability of having a CHD event is exactly that of the chosen risk level
print_state_probabilities(diagram, S_probabilities, ["R0", "R1", "R2"])

@info("Printing utility distribution.")
print_utility_distribution(U_distribution)

print_statistics(U_distribution)


using Test, LinearAlgebra

@testset "(22)<-1->(22) bacon & toner facet inequalities" begin

using BellScenario: LocalPolytope, ConvexPolytope, Degeneracy

function bacon_toner_inequalities()
    inequalities = []

    # 16 positivity inequalities
    for i in 1:16
        ineq = zeros(Int64, (1,17))
        ineq[i+1] = -1
        push!(inequalities, ineq)
    end

    # 16 inequalities of second form
    a_set = [[0 1 0 1], [1 0 1 0], [0 1 1 0], [1 0 0 1]]
    b_set = [[0 0 1 1], [1 1 0 0], [0 1 1 0], [1 0 0 1]]
    for a in a_set
        for b in b_set
            ineq = zeros(Int64, (1,17))

            ineq[1] = -2

            p1_index = a[1]*8 + b[1]*4 + 1 + 1
            p2_index = a[2]*8 + b[2]*4 + 2 + 1
            p3_index = a[3]*8 + b[3]*4 + 3 + 1
            p4_index = a[4]*8 + b[4]*4 + 4 + 1

            ineq[p1_index] = 1
            ineq[p2_index] = 1
            ineq[p3_index] = 1
            ineq[p4_index] = 1

            push!(inequalities, ineq)
        end
    end

    # 16 inequalities of 3rd form
    for a in [0,1]
        for b in [0,1]
            for i in [0,1]
                for j in [0,1]
                    ineq = zeros(Int64, (1,17))

                    p1_index = a*8 + 0*4 + i*2 + j*1 + 2
                    p2_index = a*8 + 1*4 + i*2 + j*1 + 2
                    p3_index = 0*8 + b*4 + (1-i)*2 + (1-j)*1 + 2
                    p4_index = 1*8 + b*4 + (1-i)*2 + (1-j)*1 + 2
                    p5_index = a*8 + b*4 + i*2 + (1-j)*1 + 2

                    ineq[p1_index] = -1
                    ineq[p2_index] = -1
                    ineq[p3_index] = -1
                    ineq[p4_index] = -1
                    ineq[p5_index] = 1

                    push!(inequalities, ineq)
                end
            end
        end
    end

    # Transform Bacon and Toner results to normalization subspace
    transform = [
        1 -1 0 0 0 -1 0 0 0 -1 0 0 0;
        1 0 -1 0 0 0 -1 0 0 0 -1 0 0;
        1 0 0 -1 0 0 0 -1 0 0 0 -1 0;
        1 0 0 0 -1 0 0 0 -1 0 0 0 -1;
    ]

    reduced_inequalities = []
    for ineq in inequalities
        reduced_ineq = ineq[:,14:17]*transform + ineq[:,1:13]
        push!(reduced_inequalities, reduced_ineq)
    end

    reduced_inequalities
end

@testset "verifying polytope" begin

    dir = "./test/regression/files/"

    vertices = LocalPolytope.vertices((2,2), (2,2), bits=1, fixed_direction=false)

    # Bacon and Toner claim 112 unique vertices
    @test size(vertices) == (112,)
    @test size(vertices) == size(unique(vertices))

    constraints = ConvexPolytope.facet_constraints(vertices, dir=dir)
    @test constraints[2] == []
    @test constraints[3] == []

    inequalities = constraints[1]
    # Bacon and Toner claim 48 unique facet inequalities
    @test size(inequalities) == (48,)
    @test size(inequalities) == size(unique(inequalities))

    @test issetequal( inequalities, bacon_toner_inequalities())

    canonical_groups = Degeneracy.canonical_facets((2,2),(2,2),inequalities,"normalized")
    @test length(canonical_groups) == 3
    @test issetequal(inequalities, collect(Iterators.flatten(canonical_groups)))
end

end


# TODO:
# x midpoint rule
# x scattered/interpolation
# x GCV
# x L-curve
# - bounds on λ
# - initial guess for λ? will require mods to RegularizationTools
# - scaled λ? will need to work out equivalency with λ² formulation, and resolve with
#   derivative rather than difference matrix
# - optimize λ via standard deviation?
# - relative weights?
# - arbitrary weighting -- implemented but not yet tested
# - midpoint rule with scattered?
# x add argument types for `RegularizationSmooth` constructor methods (why isn't this done
#   for the other interpolaters?)
# - make use of `munge_data` features (allow for matrix rather than vector u & t arguments?)
# - validate data and t̂
# x unit tests

const LA = LinearAlgebra

### Regularization data smoothing and interpolation
struct RegularizationSmooth{uType,tType,FT,T} <: AbstractInterpolation{FT,T}
    u::uType
    û::uType
    t::tType
    t̂::tType
    wls::uType
    wr::uType
    d::Int              # derivative degree used to calculate the roughness
    λ::AbstractFloat    # regularization parameter
    alg::Symbol  # how to determine λ: `:fixed`, `:gcv_svd`, `:gcv_tr`, `L_curve`
    Aitp::AbstractInterpolation{FT,T}
    RegularizationSmooth{FT}(u,û,t,t̂,wls,wr,d,λ,alg,Aitp) where FT =
        new{typeof(u),typeof(t),FT,eltype(u)}(u,û,t,t̂,wls,wr,d,λ,alg,Aitp)
end

"""
# Arguments
- `u::Vector`:  dependent data
- `t::Vector`:  independent data

# Optional Arguments
- `t̂::Vector`: t-values to use for the smooth curve (useful when data has missing values or
               is "scattered"); if not provided, then `t̂ = t`; must be monotonically
               increasing
- `wls::{Vector,Symbol}`: weights to use with the least-squares fitting term; if set to
                          `:midpoint`, then midpoint-rule integration weights are used for
                          _both_ `wls` and `wr`
- `wr::Vector`: weights to use with the roughness term
- `d::Int = 2`: derivative used to calculate roughness; e.g., when `d = 2`, the 2nd
                derivative (i.e. the curvature) of the data is used to calculate roughness.

# Keyword Arguments
- `λ::{Number,Tuple} = 1.0`: regulariation parameter; larger values result in a smoother
                             curve; the provided value is used directly when `alg = :fixed`;
                             otherwise it is used as an initial guess for the optimization
                             method, or as bounds if a 2-tuple is provided (TBD)
- `alg::Symbol = :gcv_svd`: algorithm for determining an optimal value for λ; the provided λ
                            value is used directly if `alg = :fixed`; otherwise `alg =
                            [:gcv_svd, :gcv_tr, :L_curve]` is passed to the
                            RegularizationTools solver

## Example Constructors
Smoothing using all arguments
```julia
A = RegularizationSmooth(u, t, t̂, wls, wr, d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::AbstractVector,
                              wls::AbstractVector, wr::AbstractVector, d::Int=2;
                              λ::Real=1.0, alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    M = _mapping_matrix(t̂, t)
    Wls½ = LA.diagm(sqrt.(wls))
    Wr½ = LA.diagm(sqrt.(wr))
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂,wls,wr,d,λ,alg,Aitp)
end
"""
Direct smoothing, no `t̂` or weights
```julia
A = RegularizationSmooth(u, t, d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, d::Int=2; λ::Real=1.0,
                              alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    t̂ = t
    N = length(t)
    M = Array{Float64}(LA.I, N, N)
    Wls½ = Array{Float64}(LA.I, N, N)
    Wr½ = Array{Float64}(LA.I, N-d, N-d)
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂, LA.diag(Wls½), LA.diag(Wr½), d,λ,alg,Aitp)
end
"""
`t̂` provided, no weights
```julia
A = RegularizationSmooth(u, t, t̂, d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::AbstractVector,
                              d::Int=2; λ::Real=1.0, alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    N, N̂ = length(t), length(t̂)
    M = _mapping_matrix(t̂, t)
    Wls½ = Array{Float64}(LA.I, N, N)
    Wr½ = Array{Float64}(LA.I, N̂-d, N̂-d)
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂, LA.diag(Wls½), LA.diag(Wr½), d,λ,alg,Aitp)
end
"""
`t̂` and `wls` provided
```julia
A = RegularizationSmooth(u, t, t̂, wls, d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::AbstractVector,
                              wls::AbstractVector, d::Int=2; λ::Real=1.0,
                              alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    N, N̂ = length(t), length(t̂)
    M = _mapping_matrix(t̂, t)
    Wls½ = LA.diagm(sqrt.(wls))
    Wr½ = Array{Float64}(LA.I, N̂-d, N̂-d)
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂,wls, LA.diag(Wr½), d,λ,alg,Aitp)
end
"""
`wls` provided, no `t̂`
```julia
A = RegularizationSmooth(u, t, nothing, wls,d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::Nothing,
                              wls::AbstractVector, d::Int=2; λ::Real=1.0,
                              alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    t̂ = t
    N = length(t)
    M = Array{Float64}(LA.I, N, N)
    Wls½ = LA.diagm(sqrt.(wls))
    Wr½ = Array{Float64}(LA.I, N-d, N-d)
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂,wls, LA.diag(Wr½), d,λ,alg,Aitp)
end
"""
`wls` and `wr` provided, no `t̂`
```julia
A = RegularizationSmooth(u, t, nothing, wls, wr, d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::Nothing,
                              wls::AbstractVector, wr::AbstractVector, d::Int=2;
                              λ::Real=1.0, alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    t̂ = t
    N = length(t)
    M = Array{Float64}(LA.I, N, N)
    Wls½ = LA.diagm(sqrt.(wls))
    Wr½ = LA.diagm(sqrt.(wr))
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂,wls, LA.diag(Wr½), d,λ,alg,Aitp)
end
"""
Keyword provided for `wls`, no `t̂`
```julia
A = RegularizationSmooth(u, t, nothing, :midpoint, d; λ=[1.0], alg=[:gcv_svd])
```
"""
function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::Nothing,
                              wls::Symbol, d::Int=2; λ::Real=1.0, alg::Symbol=:gcv_svd)
    u, t = munge_data(u, t)
    t̂ = t
    N = length(t)
    M = Array{Float64}(LA.I, N, N)
    wls, wr = _weighting_by_kw(t, d, wls)
    Wls½ = LA.diagm(sqrt.(wls))
    Wr½ = LA.diagm(sqrt.(wr))
    û, λ, Aitp = _reg_smooth_solve(u, t̂, d, M, Wls½, Wr½, λ, alg)
    RegularizationSmooth{true}(u,û,t,t̂, LA.diag(Wls½), LA.diag(Wr½), d,λ,alg,Aitp)
end
# """ t̂ provided and keyword for wls  _TBD_ """
# function RegularizationSmooth(u::AbstractVector, t::AbstractVector, t̂::AbstractVector,
#                               wls::Symbol, d::Int=2; λ::Real=1.0, alg::Symbol=:gcv_svd)

""" Solve for the smoothed depedent variables and create spline interpolator """
function _reg_smooth_solve(u::AbstractVector, t̂::AbstractVector, d::Int, M::AbstractMatrix,
                           Wls½::AbstractMatrix, Wr½::AbstractMatrix, λ::Real, alg::Symbol)
    λ = float(λ) # `float` expected by RT
    D = _derivative_matrix(t̂,d)
    Ψ = RT.setupRegularizationProblem(Wls½*M, Wr½*D)
    Wls½u = Wls½*u

    if alg == :fixed
        b̄ = RT.to_standard_form(Ψ, Wls½u) # via b̄
        ū = RT.solve(Ψ, b̄, λ)
        û = RT.to_general_form(Ψ, Wls½u, ū)
    else
        # note that the provided λ is currently ignored; using as initial guess or bounds
        # TBD, JJS 12/21/21
        result = RT.solve(Ψ, Wls½u, alg=alg)
        û = result.x
        λ = result.λ
    end
    Aitp = CubicSpline(û,t̂)
    # It seems logical to use B-Spline of order d+1, but I am unsure if theory supports the
    # extra computational cost, JJS 12/25/21
    #Aitp = BSplineInterpolation(û,t̂,d+1,:ArcLen,:Average)
    return û, λ, Aitp
end

"""
Order d derivative matrix for the provided t vector
"""
function _derivative_matrix(t::AbstractVector, d::Int)
    N = length(t)
    if d == 0
        return Array{Float64}(LA.I, (N, N))
    end
    dt = t[d+1:end] - t[1:end-d]
    V = LA.diagm(1 ./ dt)
    Ddm1 = diff(_derivative_matrix(t, d-1), dims=1)
    D = d * V*Ddm1
    return D
end

"""Linear interpolation mapping matrix, which maps `û` to `u`."""
function _mapping_matrix(t̂::AbstractVector,t::AbstractVector)
    N = length(t)
    N̂ = length(t̂)
    # map the scattered points to the appropriate index of the smoothed points
    idx = searchsortedlast.(Ref(t̂),t)
    # allow for "extrapolation"; i.e. for t̂ extremum that are interior to t
    idx[idx.==0] .+= 1
    idx[idx.==N̂] .+= -1
    # create the linear interpolation matrix
    m2 = @. (t - t̂[idx])/(t̂[idx+1] - t̂[idx])
    M = zeros(eltype(t), (N,N̂))
    for i in 1:N
        M[i,idx[i]] = 1 - m2[i]
        M[i,idx[i]+1] = m2[i]
    end
    return M
end

""" Common-use weighting, currently only `:midpoint` for midpoint-rule integration """
function _weighting_by_kw(t::AbstractVector, d::Int, wls::Symbol)
    # `:midpoint` only for now, but plan to add functionality for `:relative` weighting
    N = length(t)
    if wls == :midpoint
        bmp = zeros(N)
        bmp[1] = -t[1] + t[2]
        for i in 2:N-1
            bmp[i] = -t[i-1] + t[i+1]
        end
        bmp[N] = -t[N-1] + t[N]
        # divide by 2 doesn't matter in the minimize step, but keeping for correctness if
        # used as a template elsewhere
        bmp = bmp/2
        start = floor(Int, d/2) + 1
        final = iseven(d) ? N - (start - 1)  : N - start
        b̃mp = bmp[start:final]
        return bmp, b̃mp
    else
        throw("Unknown `$(wls)` keyword used for weighting, use `:midpoint`.")
    end
end

_interpolate(A::RegularizationSmooth{<:AbstractVector{<:Number}}, t::Number) =
    _interpolate(A.Aitp, t)


using Plots
using Printf
using SpectralMethodFFT

function xy_kl_init(c::ConfigFFT)

    xy_X = xy_Xgen(c)
    xy_Y = xy_Ygen(c)

    xy_r = sqrt(xy_X^2 + xy_Y^2)
    xy_u = sinc(xy_r / 2)
    kl_u = kl_xy(xy_u)

    return xy_u, kl_u

end

function kl_develop_kl(kl_u::KFunc, dt::Real)

    c = kl_u.config
    xlen, ylen = xlens_xranges(c.xranges)

    kl_K = kl_Kgen(c)
    kl_L = kl_Lgen(c)

    return exp( - 4 * π^2 * dt * (
        (kl_K/xlen)^2 + (kl_L/ylen)^2
    )) * kl_u

end

function main()

    # -------------
    # Parameters
    # -------------
    # *** space ***
    xranges = ((-10., 10.), (-10., 10.))
    ngrids = (128, 128)
    c = ConfigFFT(ngrids, xranges)
    # *** time ***
    t_st = 0.
    t_ed = 2.0
    nt = 200
    dt = (t_ed - t_st) / nt

    # -------------
    # Time Integration
    # -------------
    xy_u, kl_u = xy_kl_init(c)
    xy_X = xy_Xgen(c)
    xy_Y = xy_Ygen(c)

    anim = Animation()
    for it = 0:nt
        t = @sprintf("%3.2f", t_st + it*dt)
        xy_u = xy_kl(kl_u)
        plt = plot(
                xy_X[:,1], xy_Y[1,:], xy_u,
                title="t=$t",
                zlims=(-0.25, 1.1),
                xlabel="x", ylabel="y", zlabel="u",
                st=:wireframe)
        frame(anim, plt)
        if it < nt
            kl_u = kl_develop_kl(kl_u, dt)
        end
    end
    gif(anim, "heat_2d.gif", fps=20)

end

main()


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../useful_routines.jl")

p = [0.0, 0.0]

mergemapper = CutCellDG.MergeMapper()

south = CutCellDG.cell_map_to_south()
@test allapprox(south(p), [0.0, -2.0])
@test allapprox(mergemapper[1](p), [0.0, -2.0])

southeast = CutCellDG.cell_map_to_south_east()
@test allapprox(southeast(p), [2.0, -2.0])
@test allapprox(mergemapper[5](p), [2.0, -2.0])

east = CutCellDG.cell_map_to_east()
@test allapprox(east(p), [2.0, 0.0])
@test allapprox(mergemapper[2](p), [2.0, 0.0])

northeast = CutCellDG.cell_map_to_north_east()
@test allapprox(northeast(p), [2.0, 2.0])
@test allapprox(mergemapper[6](p), [2.0, 2.0])

north = CutCellDG.cell_map_to_north()
@test allapprox(north(p), [0.0, 2.0])
@test allapprox(mergemapper[3](p), [0.0, 2.0])

northwest = CutCellDG.cell_map_to_north_west()
@test allapprox(northwest(p), [-2.0, 2.0])
@test allapprox(mergemapper[7](p), [-2.0, 2.0])

west = CutCellDG.cell_map_to_west()
@test allapprox(west(p), [-2.0, 0.0])
@test allapprox(mergemapper[4](p), [-2.0, 0.0])

southwest = CutCellDG.cell_map_to_south_west()
@test allapprox(southwest(p), [-2.0, -2.0])
@test allapprox(mergemapper[8](p), [-2.0, -2.0])

polyorder = 1
numqp = 2

elasticitybasis = LagrangeTensorProductBasis(2, polyorder)
levelsetbasis = HermiteTensorProductBasis(2)

x0 = [0.0, 0.0]
meshwidths = [2.0, 1.0]
nelements = [2, 1]
refpoints = interpolation_points(elasticitybasis)
dim,nf = size(interpolation_points(levelsetbasis))
dgmesh = CutCellDG.DGMesh(x0, meshwidths, nelements, refpoints)
cgmesh = CutCellDG.CGMesh(x0, meshwidths, nelements, nf)

normal = [1.0, 0.0]
xI = [1.1, 0.0]
levelset = CutCellDG.LevelSet(
    x -> plane_distance_function(x, normal, xI),
    cgmesh,
    levelsetbasis,
)

cutmesh = CutCellDG.CutMesh(dgmesh, levelset)
cellquads = CutCellDG.CellQuadratures(cutmesh, levelset, numqp)
interfacequads = CutCellDG.InterfaceQuadratures(cutmesh, levelset, numqp)
facequads = CutCellDG.FaceQuadratures(cutmesh, levelset, numqp)

mergedmesh = CutCellDG.MergedMesh!(cutmesh, cellquads, facequads, interfacequads)
@test CutCellDG.number_of_nodes(mergedmesh) == 8
@test CutCellDG.has_merged_cells(mergedmesh)
@test CutCellDG.merge_direction(mergedmesh,+1,1) == 0
@test CutCellDG.merge_direction(mergedmesh,+1,2) == 0
@test CutCellDG.merge_direction(mergedmesh,-1,1) == 0
@test CutCellDG.merge_direction(mergedmesh,-1,2) == 2


"""
    with_period(kernel::Kernel, period::Real)

Construct a transformed kernel with `period`.

# Examples

```jldoctest
julia> kernel = with_period(SqExponentialKernel(), π/2);

julia> x = rand();

julia> y = rand();

julia> kernel(x, y) ≈ (SqExponentialKernel() ∘ PeriodicTransform(2/π))(x, y)
true
```
"""
function with_period(kernel::Kernel, period::Real)
    return kernel ∘ PeriodicTransform(inv(period))
end


@testset "NumField/NfAbsNS/ConjugatesNS.jl" begin
  Qx, x = QQ["x"]
  K, a = NumberField([x^2 - 2, x^2 - 3, x^2 - 5], "a")
  r, s = @inferred signature(K)
  @test (r, s) == (8, 0)
  Kabs, KabstoK = simple_extension(K)
  b = rand(K, -1:1)

  c = conjugates(b, 64)
  cc = conjugates(KabstoK\b, 64)

  for _c in c
    @test count(_cc -> overlaps(_cc, _c), cc) == count(_cc -> overlaps(_cc, _c), c)
  end

  Kt, t = K["t"]
  L, b = NumberField(t^2 - a[1] * a[2])
  @test signature(L) == (8, 4)
end

@testset "NumField/NfAbsNS/Conjugates.jl" begin
  Qx, x = QQ["x"]
  K, a = wildanger_field(3, 13)
  r, s = @inferred signature(K)
  @test (r, s) == (1, 1)
  b = rand(K, -1:1)

  c = conjugates(b, ComplexField(64))
  cc = conjugates(b, 64)

  for _c in c
    @test count(_cc -> overlaps(_cc, _c), cc) == count(_cc -> overlaps(_cc, _c), c)
  end
end


@testset "Map" begin
  @testset "Existence of (pre)image" begin
    G = DiagonalGroup([3, 5, 0])
    H, mH = sub(G, [G[1]])
    a = G[1]
    b = G[3]
    bb, c = @inferred haspreimage(mH, a)
    @test bb
    @test mH(c) == a
    bb, c = @inferred haspreimage(mH, b)
    @test !bb
    # TODO: Test for hasimage missing
  end

  @testset "Homomorphisms" begin
    G = DiagonalGroup([4, 4, 4])
    H = DiagonalGroup([4, 4, 4])
    h = @inferred hom(gens(G), gens(H))
    @test h(G[1]) == H[1]
    @test h(G[2]) == H[2]
    @test h(G[3]) == H[3]

    h = @inferred hom(G, [2*h for h in gens(H)])
    @test h(G[1]) == 2*H[1]
    @test h(G[2]) == 2*H[2]
    @test h(G[3]) == 2*H[3]
  end

  @testset "Kernel" begin
    G = DiagonalGroup([4, 4, 4])
    H = DiagonalGroup([4, 4, 4])
    h = @inferred hom(G, [2*h for h in gens(H)])
    @test h(G[1]) == 2*H[1]
    @test h(G[2]) == 2*H[2]
    @test h(G[3]) == 2*H[3]

    K, mK = @inferred kernel(h)
    @test all(iszero(h(mK(k))) for k in K)
    @test order(K) == 8
  end

  @testset "Image" begin
    G = DiagonalGroup([4, 4, 4])
    H = DiagonalGroup([4, 4, 4])
    h = @inferred hom(G, [2*h for h in gens(H)])
    @test h(G[1]) == 2*H[1]
    @test h(G[2]) == 2*H[2]
    @test h(G[3]) == 2*H[3]

    I, mI = @inferred image(h)
    @test all(haspreimage(h, mI(i))[1] for i in I)
    @test order(I) == 8
  end

  @testset "Injectivity" begin
    G = DiagonalGroup([4, 4, 4])
    H = DiagonalGroup([4, 4, 4])

    h = @inferred hom(G, [2*h for h in gens(H)])
    b = @inferred isinjective(h)
    @test !b

    h = @inferred hom(G, [3*h for h in gens(H)])
    b = @inferred isinjective(h)
    @test b
  end

  @testset "Surjectivity" begin
    G = DiagonalGroup([4, 4, 4])
    H = DiagonalGroup([4, 4, 4])

    h = @inferred hom(G, [2*h for h in gens(H)])
    b = @inferred issurjective(h)
    @test !b

    h = @inferred hom(G, [3*h for h in gens(H)])
    b = @inferred issurjective(h)
    @test b
  end

  @testset "Bijectivity" begin
    G = DiagonalGroup([4, 4, 4])
    H = DiagonalGroup([4, 4, 4])

    h = @inferred hom(G, [2*h for h in gens(H)])
    b = @inferred isbijective(h)
    @test !b

    h = @inferred hom(G, [3*h for h in gens(H)])
    b = @inferred isbijective(h)
    @test b
  end
end


module problem109

description = """
In the game of darts a player throws three darts at a target board which is split into twenty equal sized sections numbered one to twenty.
The score of a dart is determined by the number of the region that the dart lands in. A dart landing outside the red/green outer ring scores zero. The black and cream regions inside this ring represent single scores. However, the red/green outer ring and middle ring score double and treble scores respectively.
At the centre of the board are two concentric circles called the bull region, or bulls-eye. The outer bull is worth 25 points and the inner bull is a double, worth 50 points.
There are many variations of rules but in the most popular game the players will begin with a score 301 or 501 and the first player to reduce their running total to zero is a winner. However, it is normal to play a "doubles out" system, which means that the player must land a double (including the double bulls-eye at the centre of the board) on their final dart to win; any other dart that would reduce their running total to one or lower means the score for that set of three darts is "bust".
When a player is able to finish on their current score it is called a "checkout" and the highest checkout is 170: T20 T20 D25 (two treble 20s and double bull).

There are exactly eleven distinct ways to checkout on a score of 6:

D3
D1 	D2
S2 	D2
D2 	D1
S4 	D1
S1 	S1 	D2
S1 	T1 	D1
S1 	S3 	D1
D1 	D1 	D1
D1 	S2 	D1
S2 	S2 	D1

Note that D1 D2 is considered different to D2 D1 as they finish on different doubles. However, the combination S1 T1 D1 is considered the same as T1 S1 D1.
In addition we shall not include misses in considering combinations; for example, D3 is the same as 0 D3 and 0 0 D3.
Incredibly there are 42336 distinct ways of checking out in total.

How many distinct ways can a player checkout with a score less than 100?
"""

# Use a triple nested loop to count all the solutions
# The outer loop is all the doubles that represent the final dart
# The two inner loops represent the preceding two throws.
# The count is incremented at each nested loop, as the game can be won in 1,2 or 3 throws.
function solve()
  count = 0
  singles = [1:20;25] # 25 for outer bull
  doubles = 2 * singles
  triples = 3 * singles[1:end-1] # no triple bull
  throws = [singles;doubles;triples]

  for i = 1:length(doubles)
    count += 1
    for j = 1:length(throws)
      if doubles[i] + throws[j] < 100
        count += 1
        for k = j:length(throws)
          if doubles[i] + throws[j] + throws[k] < 100
            count += 1
          end
        end
      end
    end
  end
  println("$count ways to go out on less than 100")
end

end


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

"""
    DevadossRourke()

Devadoss-Rourke triangulation.

The algorithm is described in the first chapter of Devadoss & Rourke 2011,
and is based on the Polygonal Jordan Curve theorem.

Because the algorithm relies on recursion, it is only appropriate for small
polygonal areas. Currently, the implementation does not support holes.

## References

* Devadoss, S & Rourke, J. 2011. [Discrete and computational geometry]
  (https://press.princeton.edu/books/hardcover/9780691145532/discrete-and-computational-geometry)
"""
struct DevadossRourke end

function discretize(polyarea::PolyArea, ::DevadossRourke)
  # build bridges in case the polygonal area has
  # holes, i.e. reduce to a single outer boundary
  𝒫 = polyarea |> unique |> bridge

  # points on resulting mesh
  points = collect(vertices(𝒫))

  # Devadoss-Rourke recursion
  connec = devadossrourke(points, 1:length(points))

  SimpleMesh(points, connec)
end

function devadossrourke(v::AbstractVector{Point{Dim,T}}, inds) where {Dim,T}
  I = CircularVector(inds)
  n = length(I)

  if n > 3 # split chain
    # find lowest vertex
    i  = 1
    yᵢ = last(coordinates(v[I[1]]))
    for j in 2:n
      yⱼ = last(coordinates(v[I[j]]))
      if yⱼ < yᵢ
        i, yᵢ = j, yⱼ
      end
    end

    # left/right chains
    linds = i-1:i+1
    rinds = i+1:i+n-1

    # check if candidate diagonal is valid
    Δ = Triangle(v[I[linds]])
    intriangle = findall(j -> v[I[j]] ∈ Δ, rinds[2:end-1])
    isdiag = signarea(Δ) > zero(T) && isempty(intriangle)

    # adjust diagonal if necessary
    if !isdiag
      l  = Line(v[I[i-1]], v[I[i+1]])
      js = rinds[intriangle .+ 1]
      k  = argmax([evaluate(Euclidean(), l, v[I[j]]) for j in js])
      j  = js[k]
      linds = i:j
      rinds = j:i+n
    end

    # we adjust the circular indices and
    # use `inds` instead of `I` in the
    # recursion to avoid memory copies
    linds = [mod1(ind,n) for ind in linds]
    rinds = [mod1(ind,n) for ind in rinds]

    # perform recursion
    left  = devadossrourke(v, inds[linds])
    right = devadossrourke(v, inds[rinds])
    [left; right]
  else
    # return the triangle
    [connect(Tuple(inds), Triangle)]
  end
end


# based on Gustavo Mellior's: https://github.com/GMellior/NGM_neuralnetwork
using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using ForwardDiff
import ModelingToolkit: Interval, infimum, supremum

@parameters x # x=K
@variables u(..), s(..) # u(x)=V(K) & s(x) = I(K) policy
Dx = Differential(x)  # Dx(u(x))  == V_K(K)
Vx(x) = Dx(u(x))
#
δ=0.10; ρ=0.05; z= ρ + δ +.10; Iss = (z-ρ-δ)/(ρ+δ); kss = Iss/δ;
eq = [s(x) ~ Vx(x)-1, ρ*u(x) ~ z*x - s(x) - 0.5*s(x)^2 + Vx(x)*(s(x) - δ*x)]
bcs  = [Vx(kss) ~ 1 + δ*kss]
domains = [x ∈ Interval(kss*0.1,kss*1.5)]
#
elems = size(eq)[1]
neurons = 20
chain             = [FastChain(FastDense(1,neurons,tanh)
                    ,FastDense(neurons,neurons,softplus)
                    ,FastDense(neurons,1)) for _ in 1:elems]
dx0               = kss*(1.5-0.1)/400
initθ             = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
discretization    = PhysicsInformedNN(chain, GridTraining(dx0),init_params =initθ)
@named pde_system = PDESystem(eq,bcs,domains,[x],[u(x),s(x)])
prob = discretize(pde_system,discretization)
const losses = []
cb = function (p,l)
    push!(losses, l)
    if length(losses)%500==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    return false
end
maxit1 = 1200; maxit2 = 4000;
res  = GalacticOptim.solve(prob, ADAM(0.06); atol=1e-7, cb = cb, maxiters=maxit1)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=maxit2)
phi  = discretization.phi

using Plots
kgrid      = LinRange(kss*0.1,kss*1.5,1_000)
I          = Int.(collect(size(res.minimizer)./elems))
Vfunc(x)   = first(phi[1](x,res.minimizer[1:I[1]]))
Cfunc(x)   = first(phi[2](x,res.minimizer[I[1]+1:end]))

dVk = ForwardDiff.derivative.(Vfunc, Array(kgrid))
II  = dVk .- 1

# Plot V
plot(kgrid ,Vfunc.(kgrid),
    #line =:dashdot,
    linewidth=3,label="Neural net",
    legend=:topleft,xlabel = "k",ylabel="V(k)",
    fg_legend = :false, xtick=0:2:10, ytick=-20:1:-8)
vline!([kss],lab="Kss")
hline!([(1/ρ)*(z*kss -Iss -.5*(Iss)^2 )],lab="Vss")

#display(plot!(kgrid,vfunc,label="Finite differences",legend=:topleft))

plot(kgrid ,Cfunc.(kgrid),
    line =:dashdot,linewidth=3,label="Neural net",
    legend=:topleft,xlabel = "k",ylabel="I(k)",
    fg_legend = :false, xtick=0:2:10, ytick=0.5:0.5:2)
plot!(kgrid ,x -> Iss, lab="I_SS", ylims=(0,1))
#

# Compute HJB error
HJBerror = z*kgrid .- II .- .5*II.^2 + dVk .*(II-δ*kgrid) - ρ*Vfunc.(kgrid)
plot(kgrid,HJBerror,linewidth=3,ylims=(-1.2e-4,1e-4),
    xlabel = "k",ylabel="HJB error",label="Neural net", fg_legend = :false)
########
display(plot(losses,legend = false, yaxis=:log))


## recurrent networks
## LSTM implementation is taken from:
## https://github.com/FluxML/Flux.jl/blob/7104fd933200833d3d73e6dec2f9572e3436f4dc/src/layers/recurrent.jl

################################################################################
#                                 Vanilla RNN                                  #
################################################################################

mutable struct RNNCell
  W_ih::AbstractMatrix
  W_hh::AbstractMatrix
  b_ih::AbstractVector
  b_hh::AbstractVector
end

function RNNCell(inp::Integer, hid::Integer)
    k_sqrt = sqrt(1 / hid)
    d = Uniform(-k_sqrt, k_sqrt)
    return RNNCell(rand(d, hid, inp), rand(d, hid, hid), rand(d, hid), rand(hid))
end

Base.show(io::IO, m::RNNCell) = print(io, "RNNCell($(size(m.W_ih, 2)) => $(size(m.W_ih, 1)))")

# using proxy function to be able to compute numeric gradient in tests
function rnn_forward(W_ih, W_hh, b_ih, b_hh, x, h)
    inp_v = W_ih * x .+ b_ih
    hid_v = W_hh * h .+ b_hh
    h_ = tanh.(inp_v .+ hid_v)
    return h_
end

# input should be of size (inp_size, batch)
function forward(m::RNNCell, x::AbstractMatrix, h::AbstractMatrix)
    rnn_forward(m.W_ih, m.W_hh, m.b_ih, m.b_hh, x, h)
end

(m::RNNCell)(x::AbstractMatrix, h::AbstractMatrix) = forward(m, x, h)

mutable struct RNN
    cell::RNNCell
end
RNN(inp_size, hid_size) = RNN(RNNCell(inp_size, hid_size))
RNN(inp_hid::Pair{Int}) = RNN(RNNCell(inp_hid[1], inp_hid[2]))

function init_hidden(m::RNN, batch_size::Integer)
    hid_size = length(m.cell.b_hh)
    return zeros(hid_size, batch_size)
end

# x_seq should be of size (inp_size, batch, seq_len)
# h should be of size (hid_size, batch)
function forward(m::RNN, x_seq::AbstractArray{T, 3}, h::AbstractArray{T, 2}) where T
    # device = Yota.guess_device([h])
    inp_size, batch, seq_len = size(x_seq)
    hid_size = length(m.cell.b_hh)
    h_all_sz = (size(h)..., 1)
    cell = m.cell
    # 1st element
    x = x_seq[:, :, 1]
    h = forward(cell, x, h)
    h_all = reshape(h, h_all_sz)
    # all other elements
    for i=2:size(x_seq, 3)
        x = x_seq[:, :, i]
        h = forward(cell, x, h)
        h_all = cat(h_all, reshape(h, h_all_sz); dims=3)
    end
    return h_all, h
end

(m::RNN)(x::AbstractArray{T,3}, h::AbstractArray{T,2}) where T = forward(m, x, h)

################################################################################
#                                  LSTM                                        #
################################################################################

mutable struct LSTMCell
  W_ih::AbstractMatrix
  W_hh::AbstractMatrix
  b_ih::AbstractVector
  b_hh::AbstractVector
end

function LSTMCell(inp::Integer, hid::Integer)
    k_sqrt = sqrt(1 / hid)
    d = Uniform(-k_sqrt, k_sqrt)
    return LSTMCell(rand(d, 4*hid, inp), rand(d, 4*hid, hid), rand(d, 4*hid), rand(4*hid))
end
LSTMCell(inp_hid::Pair{Int}) = LSTMCell(inp_hid[1], inp_hid[2])

Base.show(io::IO, m::LSTMCell) = print(io, "LSTMCell($(size(m.W_ih, 2)) => $(size(m.W_ih, 1)))")

slice(len, idx) = (1:len) .+ len*(idx-1)

function lstm_forward(W_ih, W_hh, b_ih, b_hh, x, h, c)
    hid_len = size(h, 1)
    σ = tanh
    # weight breakdown in PyTorch is described here:
    # https://github.com/pytorch/pytorch/blob/0c48092b2270d56cdab327bd1ff0ca89f5b7d569/torch/nn/modules/rnn.py#L479-L487
    # however, we don't split weights into W_ii, W_if, etc., but instead first calculate
    # all linear transformations as `y` and only after that take slices of this variable
    # to calculate i, f, g and o
    y = W_ih*x .+ b_ih .+ W_hh*h .+ b_hh
    i = σ.(y[slice(hid_len, 1), :])
    f = σ.(y[slice(hid_len, 2), :])
    g = tanh.(y[slice(hid_len, 3), :])
    o = σ.(y[slice(hid_len, 4), :])
    c_ = f .* c .+ i .* g
    h_ = o .* tanh.(c_)
    return h_, c_
end

# x should be of size (inp_size, batch_size)
# h and c should be of size (hid_size, batch_size)
function forward(m::LSTMCell, x::AbstractMatrix, h::AbstractMatrix, c::AbstractMatrix)
    return lstm_forward(m.W_ih, m.W_hh, m.b_ih, m.b_hh, x, h, c)
end

(m::LSTMCell)(x::AbstractMatrix, h::AbstractMatrix, c::AbstractMatrix) = forward(m, x, h, c)

mutable struct LSTM
    cell::LSTMCell
end
LSTM(inp_size, hid_size) = LSTM(LSTMCell(inp_size, hid_size))
LSTM(inp_hid::Pair{Int}) = LSTM(LSTMCell(inp_hid[1], inp_hid[2]))

function init_hidden(m::LSTM, batch_size::Integer)
    hid_size = length(m.cell.b_hh) ÷ 4
    return zeros(hid_size, batch_size), zeros(hid_size, batch_size)
end

# x_seq should be of size (inp_size, batch, seq_len)
# h and c should be of size (hid_size, batch)
function forward(m::LSTM, x_seq::AbstractArray{T, 3}, h::AbstractArray{T, 2}, c::AbstractArray{T, 2}) where T
    # device = Yota.guess_device([h])
    inp_size, batch, seq_len = size(x_seq)
    hid_size = size(h, 1)
    h_all_sz = (size(h)..., 1)
    cell = m.cell
    # 1st element
    x = x_seq[:, :, 1]
    h, c = forward(cell, x, h, c)
    h_all = reshape(h, h_all_sz)
    # all other elements
    for i=2:size(x_seq, 3)
        x = x_seq[:, :, i]
        h, c = forward(cell, x, h, c)
        h_all = cat(h_all, reshape(h, h_all_sz); dims=3)
    end
    return h_all, h, c
end

(m::LSTM)(x::AbstractArray{T,3}, h::AbstractArray{T,2}, c::AbstractArray{T,2}) where T = forward(m, x, h, c)

################################################################################
#                                  GRU                                         #
################################################################################

mutable struct GRUCell
  W_ih::AbstractMatrix
  W_hh::AbstractMatrix
  b_ih::AbstractVector
  b_hh::AbstractVector
end

function GRUCell(inp::Integer, hid::Integer)
    k_sqrt = sqrt(1 / hid)
    d = Uniform(-k_sqrt, k_sqrt)
    return GRUCell(rand(d, 3*hid, inp), rand(d, 3*hid, hid), rand(d, 3*hid), rand(3*hid))
end
GRUCell(inp_hid::Pair{Int}) = GRUCell(inp_hid[1], inp_hid[2])

Base.show(io::IO, m::GRUCell) = print(io, "GRUCell($(size(m.W_ih, 2)) => $(size(m.W_ih, 1)))")

function gru_forward(W_ih, W_hh, b_ih, b_hh, x, h)
    hid_len = size(h, 1)
    σ = tanh
    y_i = W_ih*x .+ b_ih
    y_h = W_hh*h .+ b_hh
    # parameters are slices of precalculated y_i and y_h
    s1 = slice(hid_len, 1); s2 = slice(hid_len, 2); s3 = slice(hid_len, 3)
    r = σ.(y_i[s1, :] .+ y_h[s1, :])
    z = σ.(y_i[s2, :] .+ y_h[s2, :])
    n = tanh.(y_i[s3, :] .+ r .* y_h[s3, :])
    h_ = (1 .- z) .* n .+ z .* h
    return h_
end

# x should be of size (inp_size, batch_size)  -- check
# h and c should be of size (hid_size, batch_size)  -- check
function forward(m::GRUCell, x::AbstractMatrix, h::AbstractMatrix)
    return gru_forward(m.W_ih, m.W_hh, m.b_ih, m.b_hh, x, h)
end

(m::GRUCell)(x::AbstractMatrix, h::AbstractMatrix) = forward(m, x, h)

mutable struct GRU
    cell::GRUCell
end
GRU(inp_size, hid_size) = GRU(GRUCell(inp_size, hid_size))
GRU(inp_hid::Pair{Int}) = GRU(GRUCell(inp_hid[1], inp_hid[2]))

function init_hidden(m::GRU, batch_size::Integer)
    hid_size = length(m.cell.b_hh) ÷ 3
    return zeros(hid_size, batch_size)
end

# x_seq should be of size (inp_size, batch, seq_len)
# h should be of size (hid_size, batch)
function forward(m::GRU, x_seq::AbstractArray{T, 3}, h::AbstractArray{T, 2}) where T
    # device = Yota.guess_device([h])
    inp_size, batch, seq_len = size(x_seq)
    hid_size = size(h, 1)
    h_all_sz = (size(h)..., 1)
    cell = m.cell
    # 1st element
    x = x_seq[:, :, 1]
    h = forward(cell, x, h)
    h_all = reshape(h, h_all_sz)
    for i=2:size(x_seq, 3)
        x = x_seq[:, :, i]
        h = forward(cell, x, h)
        h_all = cat(h_all, reshape(h, h_all_sz); dims=3)
    end
    return h_all, h
end

(m::GRU)(x::AbstractArray{T,3}, h::AbstractArray{T,2}) where T = forward(m, x, h)

## SAVING THIS CODE FOR FUTURE
## PyTorch provides optionally multilayer and biderectional RNN implementation.
## There are many ways to do the same thing, but to align the code with PyTorch
## I need to clearly understand their implementation, which is out of scope for me now.
## Thus I'm just saving this in-progress piece of code for future and stick
## with simple implementation for now.
## It's worth to note that at least multiple layers can be achieved on user level
## by stacking several RNNs

## mutable struct RNN
##     num_layers::Int
##     num_directions::Int
##     cells::RNNCell
## end

## RNN(inp::Integer, hid::Integer; num_layers::Int=1, biderectional::Bool=false) =
##     RNN(num_layers, biderectional ? 2 : 1, RNNCell(inp, hid))
## RNN(inp_hid::Pair{Int}; num_layers::Int=1, biderectional::Bool=false) =
##     RNN(inp_hid[1], inp_hid[2]; num_layers=num_layers, biderectional=biderectional)

## Base.show(io::IO, m::RNN) = print(io, "RNN($(m.cell)")

## function init_hidden(m::RNN, batch_size::Integer)
##     hid_size = length(m.cell.b_hh)
##     return zeros(hid_size, batch_size, m.num_layers * m.num_directions)
## end

## # x_seq should be of size (inp_size, batch_size, seq_len)
## # H should be of size (hid_size, batch_size, num_layers*num_directions)
## function forward(m::RNN, x_seq::AbstractArray{T, 3}, H::AbstractArray{T, 3}) where T
##     inp_size, batch_size, seq_len = size(x_seq)
##     hid_size = length(m.cell.b_hh)
##     # h_all = zeros(hid_size, batch_size, m.num_layers * m.num_directions)
##     # h_all_sz = (size(h)..., 1)
##     cell = m.cell
##     for l=1:m.num_layers
##         h = H[:, :, 2l - 1]           # l-th layer, forward direction
##         for i=1:size(x_seq, 3)
##             x = x_seq[:, :, i]
##             h = forward(cell, x, h)
##             # h_all = cat(h_all, reshape(h, h_all_sz); dims=3)
##         end
##         h = H[:, :, 2l]  # l-th layer, reverse direction
##         for i=size(x_seq, 3):1
##             x = x_seq[:, :, i]
##             h = forward(cell, x, h)
##             # h_all = cat(h_all, reshape(h, h_all_sz); dims=3)
##         end
##         # TODO: add another direction
##     end
##     # return h_all, h
##     return h
## end

## (m::RNN)(x::AbstractArray{T,3}, h::AbstractArray{T,2}) where T = forward(m, x, h)


using DSP, Test
using Statistics: mean

@testset "hilbert" begin
    # Testing hilbert transform
    t = (0:1/256:2-1/256)
    a0 = sinpi.(t)
    a1 = cospi.(t)
    a2 = sinpi.(2*t)
    a3 = cospi.(2*t)
    a = hcat(a0, a1, a2, a3)

    h = hcat(hilbert(a0), hilbert(a1), hilbert(a2), hilbert(a3))
    h_abs = abs.(h)
    h_angle = angle.(h)
    h_real = real.(h)
    #The real part should be equal to the original signals:
    @test h_real ≈ a

    #The absolute value should be one everywhere, for this input:
    @test h_abs ≈ ones(size(a))

    #For the 'slow' sine - the phase should go from -pi/2 to pi/2 in
    #the first 256 bins:
    @test h_angle[1:256,1] ≈ -pi/2:pi/256:pi/2-pi/256

    #For the 'slow' cosine - the phase should go from 0 to pi in the
    #same interval:
    @test h_angle[1:256,2] ≈ 0:pi/256:pi-pi/256

    #The 'fast' sine should make this phase transition in half the time:
    @test h_angle[1:128,3] ≈ -pi/2:pi/128:pi/2-pi/128

    #Ditto for the 'fast' cosine:
    @test h_angle[1:128,4] ≈ 0:pi/128:pi-pi/128

    #The imaginary part of hilbert(cos(t)) = sin(t) Wikipedia
    @test imag(h[:,2]) ≈ a0

    #Sanity check with odd number of samples
    h2 = hilbert([ones(10); zeros(9)])
    @test real(h2) ≈ [ones(10); zeros(9)]

    #Sanity check with integer arguments
    r = rand(1:20, 128)
    @test hilbert(r) == hilbert(map(Float64, r))

    # Test hilbert with 2D input
    @test h ≈ hilbert(a)
end

@testset "fft helpers" begin
    @test meanfreq(sin.(2*π*10*(0:1e-3:10*π)),1e3) ≈ 10.0 rtol=1e-3

    # nextfastfft
    @test nextfastfft(64) == 64
    @test nextfastfft(65) == 70
    @test nextfastfft(127) == 128
    @test nextfastfft((64,65,127)) == (64,70,128)
    @test nextfastfft(64,65,127) == nextfastfft((64,65,127))
end

## COMMON DSP TOOLS

@testset "dB conversion" begin
    # dB conversion
    @test 3dB ≈ db2pow(3)
    @test -3dB ≈ db2pow(-3)
    @test 3dBa ≈ db2amp(3)
    @test -3dBa ≈ db2amp(-3)
    @test isa(3e0dB, Float64)
    @test isa(3f0dB, Float32)
    test_num = convert(Float64, pi)
    @test pow2db(test_num) ≈ 10*log10(test_num)
    @test amp2db(test_num) ≈ 20*log10(test_num)
    @test test_num*dB ≈ db2pow(test_num)
    @test test_num*dBa ≈ db2amp(test_num)
    @test test_num ≈ db2pow(pow2db(test_num))
    @test test_num ≈ db2amp(amp2db(test_num))
end

@testset "rms" begin
    n = (5,6)
    for x in ( randn(n), randn(n)+randn(n)im )
        @test rms(x) ≈ sqrt(mean(abs.(x).^2))
        @test rmsfft(fft(x)) ≈ rms(x)
    end # for
end

@test shiftin!([1,2,3,4],[5,6]) ≈ [3,4,5,6]

## arraysplit
@test collect(DSP.arraysplit(collect(1:10), 3, 1)) == [ collect(1.0:3), collect(3.0:5), collect(5.0:7), collect(7.0:9)]

# DELAY FINDING UTILITIES

@testset "delay finders" begin
    # Test signals length
    L = 10
    # Test delay
    d = 3

    # Forcing valid test parameters values
    L = round(Int, abs(L))
    d = round(Int, min(abs(d), L))

    # Generating the test signal as random 0 mean white uniform noise
    x = 2.0 .* rand(L) .- 1.0

    # Tests

    @test finddelay([zeros(d); x], x) == d
    @test finddelay([zeros(d); x], -x) == d
    @test finddelay(x, [zeros(d); x]) == -d
    @test finddelay(-x, [zeros(d); x]) == -d

    @test shiftsignal(x, d) == [zeros(d); x[1:(end - d)]]
    @test shiftsignal(x, -d) == [x[(d + 1):end]; zeros(d)]

    y = copy(x)
    z = shiftsignal!(y, d)
    @test z == y
    @test y == [zeros(d); x[1:(end - d)]]

    y = copy(x)
    z = shiftsignal!(y, -d)
    @test z == y
    @test y == [x[(d + 1):end]; zeros(d)]

    y, s = alignsignals([zeros(d); x], x)
    @test y == [x; zeros(d)]
    @test s == d

    y, s = alignsignals(x, [zeros(d); x])
    @test y == [zeros(d); x[1:(end - d)]]
    @test s == -d

    y = copy(x)
    z, s = alignsignals!(y, [zeros(d); x])
    @test z == y
    @test s == -d
    @test y == [zeros(d); x[1:(end - d)]]

    y = [zeros(d); x]
    z, s = alignsignals!(y, x)
    @test z == y
    @test s == d
    @test y == [x; zeros(d)]

end


"""
Simulated Binary Crossover (SBX).

See Deb&Agrawal "Simulated binary crossover for continuous search space", 1994, Complex Systems
"""
struct SimulatedBinaryCrossover <: CrossoverOperator{2,2}
    p::Float64      # probability to modify a dimension
    η::Float64      # distribution index
    η_exp::Float64  # pre-processed index

    function SimulatedBinaryCrossover(p::Number, η::Number)
        (0.0 <= p <= 1.0) || throw(ArgumentError("p must be within [0,1]"))
        η > 0.0 || throw(ArgumentError("η must be positive"))
        new(p, η, (1/(η+1)))
    end
    SimulatedBinaryCrossover(params::Parameters) =
        SimulatedBinaryCrossover(params[:SBX_p], params[:SBX_η])
end

const SBX_DefaultOptions = ParamsDict(
    :SBX_p => 0.2,
    :SBX_η => 3.0
)

# sample SPX spread factor β
function randbeta(xover::SimulatedBinaryCrossover)
    u = rand()
    return (u < 0.5 ? 2*u : 1/(2*(1-u)))^xover.η_exp
end

function apply!(xover::SimulatedBinaryCrossover,
                targets::AbstractVector{<:AbstractIndividual}, targetIndices::Vector{Int},
                pop, parentIndices)
    @assert length(targets) == 2
    @assert length(targetIndices) == 2
    @assert length(parentIndices) == 2
    p1ix, p2ix = parentIndices
    @inbounds for i in 1:numdims(pop)
        if rand() <= xover.p
            # crossover
            β = randbeta(xover)
            parent_avg = 0.5 * (pop[i,p1ix] + pop[i,p2ix])
            delta = 0.5 * (pop[i,p1ix] - pop[i,p2ix])
            if rand(Bool) delta = -delta end # swap the parents
            targets[1][i] = parent_avg + β * delta
            targets[2][i] = parent_avg - β * delta
        else
            # just copy the parents
            targets[1][i] = pop[i,p1ix]
            targets[2][i] = pop[i,p2ix]
        end
    end
    return targets
end


using MPIClusterManagers # to import MPIManager
using Distributed #provides addprocs()

manager=MPIManager(np=2)
addprocs(manager)

#Setup worker environments
@everywhere using Pkg
@everywhere Pkg.activate((@__DIR__)*"/..")
@everywhere using Plasmo
@everywhere using PipsNLP

graph = OptiGraph()

n1 = @optinode(graph)
n2 = @optinode(graph)

@variable(n1, 0 <= x <= 2)
@variable(n1, 0 <= y <= 3)
@variable(n1, z >= 0)
@constraint(n1, x+y+z >= 4)
@objective(n1, Min, y)

@variable(n2,x)
@NLconstraint(n2,ref,exp(x) >= 2)
@variable(n2,z >= 0)
@constraint(n2,z + x >= 4)
@objective(n2,Min,x)

@linkconstraint(graph,n1[:x] == n2[:x])

#distribute the graph to workers
julia_workers = collect(values(manager.mpi2j))
remote_references = PipsNLP.distribute_optigraph(graph,julia_workers,remote_name = :pipsgraph)

# The remote optigraphs can be queried if they are fetched from the other workers
# NOTE: this will copy the optigraph to the main process
# r1 = fetch(remote_references[1])
# r2 = fetch(remote_references[2])

#Solve with pips-nlp
@mpi_do manager begin
    using MPI
    PipsNLP.pipsnlp_solve(pipsgraph)
end

#fill the local solution with the result on worker 2 (worker 2 is the root MPI rank)
PipsNLP.fill_solution!(graph,:pipsgraph,2)


using Plots
rosenbrock(x) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

linspace(a,b,N)= (a+ (b-a)*(i)/(N-1) for i=0:(N-1))

x = y = sort(collect(Iterators.flatten([linspace(-1,-0.1,25), linspace(0,2,1000), linspace(2.1,3,25)])))

dpi = 600

l = Vector{Float64}()
for i in x
   for j in y
       v = [i j]
       append!(l, rosenbrock(v))
   end
end

surface(x,y,l, dpi=dpi)
png("./rosenbrock_surface.png")
contour(x,y,l, dpi=dpi)
png("./rosenbrock_contour.png")


struct Dual{T}
    x::T # primal ?
    der::T # derivative
end

# """ a and b are functions
# $ c(x) = a(x) + b(x) $
# $ dc/dx = da/dx + db/dx $
# """

Base.:+(a::Dual, b::Dual) = Dual(a.x + b.x, a.der + b.der)

Base.:+(a::Dual, b::Number) = Dual(a.x + b, a.der)
Base.:+(a::Number, b::Dual) = b + a

Base.:-(a::Dual, b::Dual) = Dual(a.x - b.x, a.der - b.der)

Base.:-(a::Dual, b::Number) = Dual(a.x - b, a.der)
Base.:-(a::Number, b::Dual) = Dual(a - b.x, -b.der)

Base.:-(a::Dual) = Dual(-a.x, -a.der)
# Base.:inv(a::Dual) = Dual(1/a.x, -a.der)

Base.:*(a::Dual, b::Dual) = Dual(a.x * b.x, a.der * b.x + a.x * b.der) # product rule

Base.:*(a::Dual, b::Number) = Dual(a.x * b, a.der * b)
Base.:*(a::Number, b::Dual) = b * a

Base.:/(a::Dual, b::Dual) = Dual(a.x / b.x, (a.der * b.x - a.x * b.der) / (b.x ^ 2))
Base.:/(a::Dual, b::Number) = a * inv(b)
Base.:/(a::Number, b::Dual) = Dual(a / b.x, ( -a * b.der ) / ( b.x ^ 2 ) ) #

Base.:^(a::Dual, b::Integer) = Base.power_by_squaring(a, b)
Base.:^(a::Number, b::Dual) = Dual(Float64(a^b.x), a^b.x * log(a))

Base.:≈(a::Dual, b::Dual) = a.x ≈ b.x && a.der ≈ b.der

derivative(f, x) = f(Dual(x, one(x))).der

# primatives
Base.sin(a::Dual) = Dual(sin(a.x), a.der * cos(a.x))
Base.cos(a::Dual) = Dual(cos(a.x), -a.der * sin(a.x))
Base.tan(a::Dual) = sin(a)/cos(a)

Base.csc(a::Dual) = Dual(-csc(a.x)*cot(a.x), -a.der * sin(a.x))
Base.sec(a::Dual) = Dual(tan(a.x), -a.der * sin(a.x))
Base.cot(a::Dual) = 1/tan(a)

Base.exp(a::Dual) = Dual(exp(a.x), exp(a.x) * a.der)
Base.log(a::Dual) = Dual(log(a.x), a.der/a.x)

# used with ∇ and N-dimensional scalar valued functions
# struct MultiDual{N, T}
#     x::T
#     ders::SVector{N, T}
# end

# Base.:+(a::MultiDual, b::MultiDual) = MultiDual(a.x + b.x, a.ders .+ b.ders)
# Base.:*(a::MultiDual, b::MultiDual) = MultiDual(a.x * b.x, a.ders .* b.x .+ a.x .* b.ders)

# ∇(f, x) = f(MultiDual(x, ones(length(x)))).ders

# struct MultiDual2{M, N, T}
#     x::SVector{M, T}
#     J::SMatrix{M, N, T}
# end


"""
l1_adaptive_trend_filter{T<:Float64, S<:Int}(y::Vector{T}, components::Vector{S}; f::Vector{T} = Vector{T}(0), numλ::S=S(40), numγ::S=S(10), MAXITER::S=S(500), verbose::Bool=false,lower_bounds::Vector{T}=-10e+7*ones(T, TOTALCOMPONENTS), upper_bounds::Vector{T}=10e+7*ones(T, TOTALCOMPONENTS) )

The function `l1_adaptive_trend_filter` performs the filtering via coordinate descent and takes the following required inputs:

* y : Signal or time-series to be filtered; (Must be a vector of reals)
* components: List of integers corresponding to the types of components to be considered. (must be a vector of integers containing some of the following numbers: 1 = Step, 2 = Spike, 3 = Slope, 4 = Sine, 5 = Cossine )

Optional inputs:

* f: Vector of overcomplete frequencies (Default=ø);
* numλ: Size of the regularizer path for the parameter λ (Default=40);
* numγ: Size of the regularizer path for the parameter γ (Default=10);
* MAXITER: Maximum number of iterations (Default=500);
* verbose: Boolean flag for displaying progress of algorithm (Default=false);
* lower_bounds: List of lower bounds for each component type (Default=[-∞,-∞,-∞,-∞,-∞]);
* upper_bounds: List of upper bpunds for each component type (Default=[+∞,+∞,+∞,+∞,+∞]).

This function returns:

*  β_path: Path of components coefficients;
*  y_path: Path of filtered signals;
*  β_best: Best components coefficients according to the EBIC criteria;
*  y_best: Best filtered signal according to the EBIC criteria;
*  λ_best: Best value for the λ regularizer according to the EBIC criteria;
*  γ_best: Best value for the γ regularizer according to the EBIC criteria.

## Example

```
y = rand(18) # some inputs
components = [1,3] #(meaning we only consider step and slope components)

# run the algorithm
beta_path, y_path, beta_best, y_best, lambda_best, gamma_best = l1_adaptive_trend_filter(y,components)
```

"""
function l1_adaptive_trend_filter{T<:Float64, S<:Int}(
    y::Vector{T}, components::Vector{S}; f::Vector{T} = Vector{T}(0), numλ::S=S(40), numγ::S=S(10), MAXITER::S=S(500), verbose::Bool=false,
    lower_bounds::Vector{T}=-10e+7*ones(T, TOTALCOMPONENTS), upper_bounds::Vector{T}=10e+7*ones(T, TOTALCOMPONENTS) )

    # subtracting the mean
    y_mean = mean(y)
    y = y - y_mean

    # prepare check for dimension sizes
    const N = length(y)

    # iterator
    IT = initIT_range(N, components, f, MAXITER=MAXITER)::iterator

    # compute moments
    const d = initData(IT, f, f)::dataCD

    # inner products between response y and components
    const xdy = initXDY(IT, y, d)

    # build regularization path for λ (same path shape, different λs)
    path, λ_max = compute_λ_path(IT, xdy, numλ, d)

    # build regularization path for γ (unique path)
    const Γ = compute_γ_path(IT, xdy, numγ, d)

    # initialize weights
    w = Vector{Float64}[]
    for i in 1:TOTALCOMPONENTS
        if in(i,IT.components)
            push!(w, ones(IT.nelements[i]))
        else
            push!(w, ones(0))
        end
    end

    # lasso pass
    @fastmath β_path, y_path, β_best, y_best, λ_best, γ_best = coordinate_descent(
      IT, d, xdy, path, λ_max, [1.0,] , y, lower_bounds, upper_bounds, w, verbose
      )

    # exclude the components the lasso has set to zero
    update_components!(IT, w, β_best, d, xdy)

    @fastmath β_path, y_path, β_best, y_best, λ_best, γ_best = coordinate_descent(
      IT, d, xdy, path, λ_max, Γ, y, lower_bounds, upper_bounds, w, verbose
    )

    # adding back the mean
    y_best = y_best + y_mean

    if verbose
      print(string(
              "best regularization according to BIC was (λ=", round(λ_best, 3),
              ", γ=", round(γ_best, 3), ") \n"
              ))
    end

    return β_path, y_path, β_best, y_best, λ_best, γ_best
end

# coordinate descent algorithm for the regularization path (Λ x Γ)
function coordinate_descent(
    IT::iterator, d::dataCD, xdy::Vector{Vector{Float64}}, path::Vector{Float64}, λ_max::Vector{Float64}, Γ::Vector{Float64},
    y::Vector{Float64}, lower_bounds::Vector{Float64}, upper_bounds::Vector{Float64}, w::Vector{Vector{Float64}}, verbose::Bool; sparse::Int = 0
    )

  const N = IT.obs

  # initializations
  #if sparse == 1
  #  β_path, β_tilde, β, activeSet = initSparse(IT)
  #else
  β_path, β_tilde, β, activeSet = initDense(IT)
  #end

  # memory allocation
  BIC = Inf::Float64
  β_ols = 0.0::Float64
  partial_fit = 0.0
  β_best = 0.0
  y_best = 0
  λ_best = 0.0
  γ_best = 0.0
  path_iteration = 0

  y_path = Vector{Float64}[]
  #for i in 1:TOTALCOMPONENTS
  #  push!(w, zeros(IT.nelements[i]))
  #end

  # regularization path
  @inbounds for γ in Γ

    # clear warm-start
    clear_warmstart!(β_tilde, IT)

    @inbounds for path_iter in path

      λ = path_iter * λ_max

      change_flag = true

      if verbose
        path_iteration += 1
        print(string(
                "regularizers = (", round(λ, 3), ", ", round(γ, 3),
                ") path iteration = ", path_iteration, ".\n"
                ))
      end

      # loop until active set converges
      @inbounds for iter in 1:IT.maxIter

        if verbose
          println(iter)
        end

        # convergence check
        if !change_flag
          break
        end
        change_flag = false

        # cycle through every component
        @inbounds for c1 in IT.components, j in IT.elements[c1]

            # consider hoisting means and variaces, for better memoru use

            # compute the partial fit with the components in the active set
            partial_fit = 0.0
            @inbounds for c2 in IT.components, l in IT.elements[c2]
                @inbounds if activeSet[c2][l] && (c1, j) != (c2, l)
                    @inbounds partial_fit += GM2(c1, c2, j, l, d, N) * β_tilde[c2][l]
                end
            end

            # univariate ordinary leasts squares coefficient
            inner_prod_partial_residual = (xdy[c1][j] - partial_fit) / IT.obs

            # weighted penalty
            #w[c1][j] = 1.0 / (abs(β_ols)^γ)
            # w[c1][j] = 1.0 / abs(inner_prod_partial_residual)

            # # soft thresholding operator
            # if abs(inner_prod_partial_residual) <= λ[c1] * w[c1][j]^γ
            #     if activeSet[c1][j]
            #         β_tilde[c1][j] = 0.0
            #         activeSet[c1][j] = false
            #         change_flag = true
            #         #println("$(c1)  , $(j)")
            #     end
            # else
            #     β_tilde[c1][j] = sign(inner_prod_partial_residual) * (abs(inner_prod_partial_residual) - λ[c1] * w[c1][j]^γ) / (d.σ[c1][j]^2)

            #     # projection onto the box constraints [lower_bound, upper_bound]
            #     β_tilde[c1][j] = max(β_tilde[c1][j], lower_bounds[c1])
            #     β_tilde[c1][j] = min(β_tilde[c1][j], upper_bounds[c1])

            #     if !activeSet[c1][j] #&& β_tilde[c1][j] != 0.0
            #         activeSet[c1][j] = true
            #         change_flag = true
            #         #println("$(c1)  , $(j)")
            #     end
            # end
            change_flag = soft_treshold(j, inner_prod_partial_residual, λ[c1], γ, w[c1][j], activeSet[c1], β_tilde[c1], d.σ[c1][j], lower_bounds[c1], upper_bounds[c1])

          end
      end

      push!(β_path, deepcopy(β_tilde))

      # bayesian information criterion
      BIC_new, y_hat = compute_BIC(y, β_tilde, activeSet, IT, d, xdy)

      push!(y_path, copy(y_hat))

      if verbose
          print(string(" BIC = ", BIC_new))
      end

      # save the best fit so far
      if BIC_new < BIC
          BIC = BIC_new
          y_best = copy(y_hat)
          β_best = deepcopy(β_tilde)
          λ_best = λ
          γ_best = γ
      end

    end
  end

  return β_path, y_path, β_best, y_best, λ_best, γ_best
end

function clear_warmstart!(β_tilde::Vector{Vector{Float64}}, IT::iterator)

  # clear warm-start
    for i in IT.components
        for j in IT.elements[i]
            β_tilde[i][j] = 0.0
        end
    end

    nothing
end

function soft_treshold(j::Int, residual::Float64, λ::Float64, γ::Float64, w::Float64, activeSet::Vector{Bool}, β_tilde::Vector{Float64},
    var::Float64, lb::Float64, ub::Float64 )::Bool

    if abs(residual) <= λ * w^γ
        if activeSet[j]
            β_tilde[j] = 0.0
            activeSet[j] = false
            return true
            #println("$(c1)  , $(j)")
        else
            return false
        end
    else
        β_tilde[j] = sign(residual) * (abs(residual) - λ * w^γ) / (var^2)

        # projection onto the box constraints [lower_bound, upper_bound]
        β_tilde[j] = max(β_tilde[j], lb)
        β_tilde[j] = min(β_tilde[j], ub)

        if !activeSet[j] #&& β_tilde[c1][j] != 0.0
            activeSet[j] = true
            return true
            #println("$(c1)  , $(j)")
        else
            return false
        end
    end
    return false
end

function update_components!(IT, w, β, d, xdy)

  # clear given iterator
  for c in IT.components
    IT.elements[c] = Int[]
  end

  for c in IT.components
    for j in 1:IT.nelements[c]

      if β[c][j] != 0.0
        # pre-compute weight
        w[c][j] = abs(1.0 / β[c][j]) #/ (d.σ[c][j]^2)
        β[c][j] = 0.0
        # update iterator only with nonzero elements
        push!(IT.elements[c], j)
      end

    end
  end

  return nothing
end


"""
    Spectral{F,M,V1,V2,V3}

Create a spectral layer, whose forward pass is given by

    σ.(B*(λ_in.*x) .- λ_out.*(B*x) .+ b)

where `B` is a static matrix, and `λ_in`, `λ_out` and `b` are trainable vectors of parameters.

These layers can be used as drop-in replacement to traditional dense layers, and are usually as good as dense layers for similar amounts of parameters (but spectral layers are wider for the same number of parameters).

# Examples
```jldoctest
julia> sp = Spectral(5, 2)
Spectral(5 => 2)
```
"""
struct Spectral{F, M, V1, V2, V3}
    B::M
    λ_in::V1
    λ_out::V2
    b::V3
    σ::F
end

function Spectral(in::Integer, out::Integer, σ=identity;
    init=Flux.glorot_uniform,
    bias=true
    )
    B = init(out, in)
    λ_in = init(in)
    λ_out = init(out)
    b = bias ? init(out) : bias
    Spectral(B, λ_in, λ_out, b, σ)
end

Flux.trainable(s::Spectral) = (s.λ_in, s.λ_out, s.b,)

Flux.@functor Spectral

function (s::Spectral)(x::AbstractVecOrMat)
    B, λ_in, λ_out, b, σ = s.B, s.λ_in, s.λ_out, s.b, s.σ
    σ.(B*(λ_in.*x) .- λ_out.*(B*x) .+ b)
end

function Base.show(io::IO, sp::Spectral)
    print(io, "Spectral(", size(sp.B, 2), " => ", size(sp.B, 1))
    sp.σ == identity || print(io, ", ", sp.σ)
    sp.b == false && print(io, "; bias=false")
    print(io, ")")
end

"""
    spectralize(d::Dense)

Create a spectral layer with the same dimensions as `d`.
"""
function spectralize(d::Dense)
    in, out = size(d.weight)
    Spectral(in, out, d.σ)
end


using PositiveFactorizations

function positive_linsolve(d, B, ∇f)
    cholesky!(Positive, B)
    Bchol = Cholesky(B, 'L', 0)
    d .= Bchol \ ∇f
end
function positive_linsolve(B, ∇f)
    Bchol = cholesky(Positive, B)
    Bchol \ ∇f
end
function positive_factorize(B)
    cholesky!(Positive, B)
    Cholesky(B, 'L', 0)
end

Base.summary(::NLSolvers.Newton{<:Direct,typeof(positive_linsolve)}) =
    "Newton's method with PositiveFactorizations.jl linsolve"


function sumsqrtn(n)
r = 0
for i = 1:n
r = r + sqrt(i)
end
return r
end


using LinearAlgebra
using Statistics
"""
    InterpretMe.partial_dependence(data::Array{<:Any,2},
                                   features::Union{<:Integer,Vector{<:Integer}},
                                   estimate_fn, grid_size::Int,
                                   return_ice::Bool)
Calculate the partial dependance of a machine learning model.

`estimate_fn` is a function from the model that generates a vector of predictions from a matrix of inputs.
`feature_names` should be the feature names corresponding to the features in `data`
`return_ice` whether to return the ICE lines.
Returns
`pd` the partial dependancy results
`grid_values` the grid values for all the selected features
"""
function partial_dependence(data::Array{<:Any,2},
                            features::Union{<:Integer,Vector{<:Integer}},
                            estimate_fn, grid_size::Int, return_ice::Bool)
    # Mainly inspired by Sk-learn's function.
    # https://scikit-learn.org/stable/modules/generated/sklearn.inspection.partial_dependence.html
    #Step 1 Necessary checking
    if grid_size <=1
         raise(error("grid_size should be greater than 1, $grid_size was passed"))
    end
    #Step 2 Get_grid values
    grid, grid_values = get_grid(data[:,features], grid_size)
    #Step 3 evaluate partial_dependance
    probe = estimate_fn(data[1:2,:])
    pd = typeof(probe)(undef, size(grid,1), size(probe,2))
    # pd = []
    for grid_value_index in 1:size(grid,1)
        temp_data = copy(data)
        # (n,length(features)) .= (1,length(features))
        temp_data[:,features] .= grid[grid_value_index,:]
        predictions = estimate_fn(temp_data)
        pd[grid_value_index,:] = mean(predictions, dims=1)
    end
    pd, grid_values
end

function get_grid(data_features, grid_size)
    values = Nothing
    for feature in 1:size(data_features,2)
        unique_vals = unique(data_features[:,feature])
        value =
            if length(unique_vals) <= grid_size
                sort(unique_vals)
            else
                LinRange(minimum(data_features[:,feature]),
                 maximum(data_features[:,feature]), grid_size)
            end
        if values == Nothing
            values = [value]
        else
            push!(values, value)
        end
    end
    cartesian_prod_looped(values), values
end

function cartesian_prod(values)
    if length(values) == 1
        values[1]
    else
        cartesian_prod([
        hcat(repeat(values[1],inner=(size(values[2],1),
                                    [1 for i in 1:ndims(values[1])-1]...))
            ,repeat(values[2],outer=(size(values[1],1),
                                    [1 for i in 1:ndims(values[2])-1]...))
            )
         , values[3:end]...])
     end
 end

 function cartesian_prod_looped(values)
     while length(values) > 1
         values = [
         hcat(repeat(values[1],inner=(size(values[2],1),
                                     [1 for i in 1:ndims(values[1])-1]...))
             ,repeat(values[2],outer=(size(values[1],1),
                                     [1 for i in 1:ndims(values[2])-1]...))
             )
          , values[3:end]...]
      end
      values[1]
      # To return transposed.
      # res_values = typeof(values[1])(undef,size(transpose(values[1]))...)
      # using LinearAlgebra
      # transpose!(res_values, values[1])
      # or
      # typeof(values[1])(transpose(values[1]))
  end


using Base.Test, Distributions, DataFrames
using UnobservableQueue
const UQ = UnobservableQueue

@testset "Unobservable Queue" begin
    @testset "Queue Simulation" begin
        # M/M/2 queue
        c = 2 # number of servers
        μ = 1 / c # service rate (mean service time is 1/μ) (single server)
        λ = 0.9 # arrival rate (mean interarrival time is 1/λ)
        adist = Exponential(1/λ) # arrival distribution
        sdist = Exponential(1/μ) # service distribution (single server)
        ncust = 10_000 # total number of customers generated
        timelimit = 1_000 # time limit for simulation
        seed = 8710 # rng seed
        s1 = UQ.Sim(ncust, timelimit, seed)
        q1 = UQ.Queue(c, adist, sdist)
        df1 = UQ.runsim(s1, q1)
        @test mean(df1[:dtime] - df1[:stime]) ≈ (1 / μ) atol = 0.1
        @test df1[:atime] |> diff |> mean ≈ (1 / λ) atol = 0.1
        @test all(diff(df1[:atime]) .> 0)
        @test all(diff(sort(df1, :dorder)[:dtime]) .> 0)

        # G/G/15 queue
        c = 15 # number of servers
        μa = (1 / 2) * 0.99
        μs = 1 / 2
        σ = 0.9
        adist = LogNormal(exp(1/μa), σ) # arrival distribution
        sdist = LogNormal(exp(1/μs), σ) # service distribution (single server)
        ncust = 10_000 # total number of customers generated
        timelimit = 10_000 # time limit for simulation
        seed = 3001 # rng seed
        s2 = UQ.Sim(ncust, timelimit, seed)
        q2 = UQ.Queue(c, adist, sdist)
        df2 = UQ.runsim(s2, q2)
        @test mean(df2[:dtime] - df2[:stime]) ≈ mean(rand(sdist, 10_000)) atol = 100
        @test df2[:atime] |> diff |> mean ≈ mean(rand(adist, 10_000)) atol = 100
        @test all(diff(df2[:atime]) .> 0)
        @test all(diff(sort(df2, :dorder)[:dtime]) .> 0)
    end

    @testset "Server Estimation" begin
        # Introduce observation error into true data
        c = 2 # number of servers
        μ = 1 / c # service rate (mean service time is 1/μ) (single server)
        λ = 0.9 # arrival rate (mean interarrival time is 1/λ)
        adist = Exponential(1/λ) # arrival distribution
        sdist = Exponential(1/μ) # service distribution (single server)
        ncust = 10_000 # total number of customers generated
        timelimit = 1_000 # time limit for simulation
        seed = 8710 # rng seed
        s1 = UQ.Sim(ncust, timelimit, seed)
        q1 = UQ.Queue(c, adist, sdist)
        df1 = UQ.runsim(s1, q1)
        noiseout = UQ.disorder(df1, 2)
        @test count(noiseout[:DepartOrderErr] - noiseout[:DepartOrder] .== 0.0) / size(noiseout,1) ≈ 0.9 atol = 0.1
        @test all(sort(noiseout, :DepartMeas)[:DepartOrder] .== collect(1:size(noiseout,1)))
        @test all(sort(noiseout, :DepartMeas)[:DepartOrder] .== sort(noiseout, :Depart)[:DepartOrder])

        # Deltamax algorithm
        o1 = [1,2,3,5,4]
        o2 = [1,3,6,5,4]
        @test UQ.c_deltamax(o1) == 2 && UQ.c_deltamax(o2) == 3

        # Order-based algorithm
        c = 15 # number of servers
        μa = (1 / c) * 0.99
        μs = 1 / c
        σ = 0.9
        adist = LogNormal(exp(1/μ), σ) # arrival distribution
        sdist = LogNormal(exp(1/μ), σ) # service distribution (single server)
        ncust = 10_000 # total number of customers generated
        timelimit = 10_000 # time limit for simulation
        seed = 3001 # rng seed
        s2 = UQ.Sim(ncust, timelimit, seed)
        q2 = UQ.Queue(c, adist, sdist)
        df2 = UQ.runsim(s2, q2)
        outorder1 = sort(df1, :dorder)[:aorder]
        outorder2 = sort(df2, :dorder)[:aorder]
        outorder3 = randperm(1_000)
        @test UQ.c_order(outorder3) == UQ.c_order_slow(outorder3)
        @test UQ.c_order([1,3,5]) == 3
        @test UQ.c_order(outorder1) == 2

        # Variance and Uninformed algorithms
        cmax = 19
        A1, A2 = df1[:atime][1:1000], df2[:atime][1:1000]
        D1, D2 = df1[:dtime][1:1000], df2[:dtime][1:1000]
        (cvar1, cunf1, VS1) = UQ.c_var_unf(A1, D1, cmax)
        (cvar2, cunf2, VS2) = UQ.c_var_unf(A2, D2, cmax)
        @test cvar1 == 2 && cunf1 == 2

        # Build distributions
        (adist, sdist) = UQ.builddist(2, "Uniform", "LogNormal", 0.99)
        (adist2, sdist2) = UQ.builddist(15, "Exponential", "Beta", 0.9)
        @test (1/mean(adist)) / (2 * (1/mean(sdist))) ≈ 0.99 atol = 0.0001
        @test (1/mean(adist2)) / (15 * (1/mean(sdist2))) ≈ 0.9 atol = 0.0001
    end

    @testset "Metrics" begin
        # Convergence estimation
        settings = ["Order" "Exponential" "Exponential" 400 9 0.9;
                    "Order" "Uniform" "LogNormal" 400 15 0.9]
        n_runs = 10 # limit runs for debugging and timing
        n_methods = 3 # number of methods to compare
        max_servers = 19 # max number of servers
        obs_max = 1_000 # max observations available
        time_limit = 10_000 # max simulation time
        ϵ = 0.05 # convergence quality
        window = 20 # observation window for convergence estimate
        window_detail = 50 # observation window for error estimate
        step = 20 # how many observations to skip while calculating convergence
        seed = 8710
        param_inf = Paraminf(settings, n_runs, n_methods, max_servers, obs_max, time_limit, ϵ, window, window_detail, step, seed)
        (conv, conv_meas, ests, ests_meas) = UQ.convergence(param_inf,1)
        @test (conv[1] > 0) && (conv[1] < 400) && ((conv[2] == 0) || conv[2] > conv[1])
        @test (conv[1] == conv_meas[1]) && (conv_meas[2] == 0)

        # Error estimation
        settings = ["Order" "Exponential" "Exponential" 400 9 0.9;
                    "Order" "Uniform" "LogNormal" 400 15 0.9;
                    "Order" "LogNormal" "LogNormal" 60 9 0.9]
        n_runs = 10 # limit runs for debugging and timing
        n_methods = 3 # number of methods to compare
        max_servers = 19 # max number of servers
        obs_max = 1_000 # max observations available
        time_limit = 10_000 # max simulation time
        ϵ = 0.05 # convergence quality
        window = 20 # observation window for convergence estimate
        window_detail = 50 # observation window for error estimate
        step = 20 # how many observations to skip while calculating convergence
        seed = 8710
        param_inf = Paraminf(settings, n_runs, n_methods, max_servers, obs_max, time_limit, ϵ, window, window_detail, step, seed)
        (err, err_meas, detail, detail_meas) = UQ.esterror(param_inf, 3)
        @test err[1] == 0 && err[2] >= 1
        @test err_meas[1] == 0 && err_meas[2] >= err[2]

        # Comparison
        settings = ["Order" "Exponential" "Exponential" 400 9 0.9;
                    "Order" "Uniform" "LogNormal" 400 15 0.9;
                    "Order" "LogNormal" "LogNormal" 60 9 0.9]
        n_runs = size(settings,1) # experimental runs
        n_methods = 3 # number of methods to compare
        max_servers = 19 # max number of servers
        obs_max = 1_000 # max observations available
        time_limit = 10_000 # max simulation time
        ϵ = 0.05 # convergence quality
        window = 20 # observation window for convergence estimate
        window_detail = 50 # observation window for error estimate
        step = 20 # how many observations to skip while calculating convergence
        seed = 8710
        param_inf = Paraminf(settings, n_runs, n_methods, max_servers, obs_max, time_limit, ϵ, window, window_detail, step, seed)
        (rerr, rconv, rraw) = infer(param_inf)
        @test mean(rerr[1][:,1]) < mean(rerr[1][:,2])
        @test mean(rerr[2][:,1]) < mean(rerr[2][:,2])
        @test (rconv[1][1,1] < rconv[1][1,2]) || (rconv[1][1,2] == 0)
        @test (rconv[2][1,1] < rconv[2][1,2]) || (rconv[2][1,2] == 0)
    end
end

@testset "LCFS" begin
    # LCFS Order-based algorithm
    c = 7 # number of servers
    μ = 1 / c # service rate (mean service time is 1/μ) (single server)
    λ = 0.9 # arrival rate (mean interarrival time is 1/λ)
    adist = Exponential(1/λ) # arrival distribution
    sdist = Exponential(1/μ) # service distribution (single server)
    ncust = 10_000 # total number of customers generated
    timelimit = 10_000 # time limit for simulation
    seed = 3001 # rng seed
    s2 = UQ.Sim(ncust, timelimit, seed)
    q2 = UQ.Queue(c, adist, sdist)
    df = UQ.runsimLCFS(s2, q2)
    lcfsorder, dfl = UQ.lcfs(df, :dtime)
    chat = UQ.c_order_LCFS(lcfsorder)
    @test chat == c

    # LCFS Variance, Uninformed algorithm
    c = 7 # number of servers
    λ = 0.99 # arrival rate (mean interarrival time is 1/λ)
    μ = 1 / c # service rate (mean service time is 1/μ) (single server)
    adist = Exponential(1/λ) # arrival distribution
    sdist = Exponential(1/μ) # service distribution (single server)
    ncust = 500 # total number of customers generated
    timelimit = 10_000 # time limit for simulation
    seed = 1302 # rng seed
    cmax = 19
    s2 = UQ.Sim(ncust, timelimit, seed)
    q2 = UQ.Queue(c, adist, sdist)
    df1 = UQ.runsimLCFS(s2, q2)
    A1, D1 = df1[:atime], df1[:dtime]
    (cvar1, cunf1, VS1) = UQ.c_var_unf_LCFS(df1, A1, D1, cmax)
    @test cvar1 == c
    @test cunf1 == c

    # LCFS first depature function
    df = DataFrame()
    df[:dtime] = [1.1,1.2,1.3,1.4,1.5,1.6]
    df[:nd] = [1,2,3,4,5,6]
    @test UQ.firstdepart(df, 5, 1.2) == 1.5
    @test UQ.firstdepart(df, 5, 1.6) == 1.6

    # LCFS Introduce observation error into true data
    c = 2 # number of servers
    λ = 0.9 # arrival rate (mean interarrival time is 1/λ)
    μ = 1 / c # service rate (mean service time is 1/μ) (single server)
    adist = Exponential(1/λ) # arrival distribution
    sdist = Exponential(1/μ) # service distribution (single server)
    ncust = 10_000 # total number of customers generated
    timelimit = 1_000 # time limit for simulation
    seed = 8710 # rng seed
    s1 = UQ.Sim(ncust, timelimit, seed)
    q1 = UQ.Queue(c, adist, sdist)
    df1 = UQ.runsimLCFS(s1, q1)
    noiseout = UQ.disorderLCFS(df1)
    @test all(sort(noiseout, :dmeas)[:dorder] .== collect(1:size(noiseout,1)))
    @test all(sort(noiseout, :dmeas)[:dorder] .== sort(noiseout, :dtime)[:dorder])

    # LCFS Convergence estimation
    settings = ["Order" "Exponential" "Exponential" 400 9 0.9;
                "Order" "Uniform" "LogNormal" 400 15 0.9]
    n_runs = 10 # limit runs for debugging and timing
    n_methods = 3 # number of methods to compare
    max_servers = 19 # max number of servers
    obs_max = 1_000 # max observations available
    time_limit = 10_000 # max simulation time
    ϵ = 0.05 # convergence quality
    window = 20 # observation window for convergence estimate
    window_detail = 50 # observation window for error estimate
    step = 20 # how many observations to skip while calculating convergence
    seed = 8710
    param_inf = Paraminf(settings, n_runs, n_methods, max_servers, obs_max, time_limit, ϵ, window, window_detail, step, seed)
    @time (conv, conv_meas, ests, ests_meas) = UQ.convergenceLCFS(param_inf,1)
    @test (conv[1] > 0) && (conv[1] < 400) && ((conv[2] == 0) || conv[2] >= conv[1])
    @test (conv[1] == conv_meas[1]) && (conv_meas[2] >= conv_meas[1])

    # LCFS Error estimation
    settings = ["Order" "Exponential" "Exponential" 400 9 0.9;
                "Order" "Uniform" "LogNormal" 400 15 0.9;
                "Order" "LogNormal" "LogNormal" 60 9 0.9]
    n_runs = 10 # limit runs for debugging and timing
    n_methods = 3 # number of methods to compare
    max_servers = 19 # max number of servers
    obs_max = 1_000 # max observations available
    time_limit = 10_000 # max simulation time
    ϵ = 0.05 # convergence quality
    window = 20 # observation window for convergence estimate
    window_detail = 50 # observation window for error estimate
    step = 20 # how many observations to skip while calculating convergence
    seed = 8710
    param_inf = Paraminf(settings, n_runs, n_methods, max_servers, obs_max, time_limit, ϵ, window, window_detail, step, seed)
    @time (err, err_meas, detail, detail_meas) = UQ.esterrorLCFS(param_inf, 3)
    @test err[1] < err[2]
    @test err_meas[1] < err_meas[2]

    # LCFS Comparison
    settings = ["Order" "Exponential" "Exponential" 400 9 0.9;
                "Order" "Uniform" "LogNormal" 400 15 0.9;
                "Order" "LogNormal" "LogNormal" 60 9 0.9]
    n_runs = size(settings,1) # experimental runs
    n_methods = 3 # number of methods to compare
    max_servers = 19 # max number of servers
    obs_max = 1_000 # max observations available
    time_limit = 10_000 # max simulation time
    ϵ = 0.05 # convergence quality
    window = 20 # observation window for convergence estimate
    window_detail = 50 # observation window for error estimate
    step = 20 # how many observations to skip while calculating convergence
    seed = 8710
    param_inf = Paraminf(settings, n_runs, n_methods, max_servers, obs_max, time_limit, ϵ, window, window_detail, step, seed)
    @time (rerr, rconv, rraw) = inferLCFS(param_inf)
    @test mean(rerr[1][:,1]) < mean(rerr[1][:,2])
    @test mean(rerr[2][:,1]) < mean(rerr[2][:,2])
    @test (rconv[1][1,1] <= rconv[1][1,2]) || (rconv[1][1,2] == 0)
    @test (rconv[2][1,1] <= rconv[2][1,2]) || (rconv[2][1,2] == 0)
end


@inline function add_fpdd_dd(x::T, y::Tuple{T,T}) where {T<:AbstractFloat}
    yhi, ylo = y
    yhi, ylo = add_2(x, yhi, ylo)
    return yhi, ylo
end

@inline function sub_fpdd_dd(x::T, y::Tuple{T,T}) where {T<:AbstractFloat}
    yhi, ylo = y
    yhi, ylo = sub_2(x, yhi, ylo)
    return yhi, ylo
end

@inline function mul_fpdd_dd(x::T, y::Tuple{T,T}) where {T<:AbstractFloat}
    hi, lo = y
    hihi, hilo = mul_2(x, hi)
    lohi, lolo = mul_2(x, lo)
    hi, lo = add_2(hihi, hilo, lohi, lolo)
    return hi, lo
end

@inline function dvi_fpdd_dd(x::T, y::Tuple{T,T}) where {T<:AbstractFloat}
    yinv = inv_dd_dd(y)
    zhi, zlo = mul_fpdd_dd(x, yinv)
    return zhi, zlo
end

@inline function dvi_fpdd_dd_fast(x::T, y::Tuple{T,T}) where {T<:AbstractFloat}
    yinv = inv_dd_dd_fast(y)
    zhi, zlo = mul_fpdd_dd(x, yinv)
    return zhi, zlo
end


using Printf
# using Plots
using Statistics
using LinearAlgebra
using BenchmarkTools

using PyPlot
pygui(true);

dim=2
mult=3
n=30

i=5

tau=2*pi
omega=1

#Delayed Mathieu parameters
kappa=0.1
delta=3.2
epsilon=1
b=0.5

#Cauchy matrices
function A(t)
 [-kappa -(delta+epsilon*cos(omega*t)); 1 0]
end

B=[0 b; 0 0]

deltat=tau/n
T=2*pi/omega

function checkT(tau,T)
    if tau <= T
        println("OK")
    else
        println("Error")
    end
end

checkT(tau,T)

#initial function
function v_ini(t)
    transpose([1*sin(omega*(t+tau/2)), 2*sin(omega*(t+tau/3)), 3*sin(omega*(t+tau/4)), 4*sin(omega*(t+tau/5)), 5*sin(omega*(t+tau/6)), 6*sin(omega*(t+tau/7))])
end

vect=zeros(n+1,mult*dim)
for j = 0:n
        vect[j+1,:]=v_ini(-tau+j*deltat)
end
vect
function simulation(init)
    k=Int16(floor(i*T/deltat))
    solmatrix=zeros(k,mult*dim)
    sol=vcat(init,solmatrix)

    for m = 1:dim:(dim*mult-(dim-1))
        for j = 0:(k-1)
            Y1=sol[n+1+j,m:m+(dim-1)]
            Y2=Y1+deltat/2*(A(j*deltat)*Y1+B*sol[j+1,m:m+(dim-1)])
            Y3=Y1+deltat/2*(A((j+1/2)*deltat)*Y2+B*(sol[j+1,m:m+(dim-1)]+sol[j+2,m:m+(dim-1)])/2)
            Y4=Y1+deltat*(A((j+1/2)*deltat)*Y3+B*(sol[j+1,m:m+(dim-1)]+sol[j+2,m:m+(dim-1)])/2)
            Y=Y1+deltat/6*((A(j*deltat)*Y1+B*sol[j+1,m:m+(dim-1)])+2*(A((j+1/2)*deltat)*Y2+B*(sol[j+1,m:m+(dim-1)]+sol[j+2,m:m+(dim-1)])/2)+2*(A((j+1/2)*deltat)*Y3+B*(sol[j+1,m:m+(dim-1)]+sol[j+2,m:m+(dim-1)])/2)+(A((j+1)*deltat)*Y4+B*sol[j+2,m:m+(dim-1)]))
            sol[n+2+j,m:m+(dim-1)]=transpose(Y)
        end
    end
    tvect=transpose(zeros(1,n+1+k))
    for j = 1:(n+1+k)
        tvect[j]=-n*deltat+j*deltat
    end
    solt=hcat(tvect,sol)
    return(solt)
end

solt1=simulation(vect)
@btime simulation(vect)
@time simulation(vect)

tplot=[solt1[j,1] for j in 1:size(solt1,1)]
v1plot=[solt1[j,2] for j in 1:size(solt1,1)]
v2plot=[solt1[j,3] for j in 1:size(solt1,1)]
v3plot=[solt1[j,4] for j in 1:size(solt1,1)]
v4plot=[solt1[j,5] for j in 1:size(solt1,1)]
v5plot=[solt1[j,6] for j in 1:size(solt1,1)]
v6plot=[solt1[j,7] for j in 1:size(solt1,1)]
plot1=plot(tplot,hcat(v1plot,v2plot,v3plot,v4plot,v5plot,v6plot))

function Ha(solt)
    rat=Int16(floor(T/tau))
    ha=zeros(n,(i+1)*mult*dim)
    for j = 1:i+1
        ha[:,(j-1)*mult*dim+1:j*mult*dim]=solt[(j-1)*n+1:j*n,2:mult*dim+1]
    end
    return(ha)
end

ha=Ha(solt1)

function eigenv(haa)
    eigena=zeros(mult*dim,1)
    Hglob=zeros(mult*dim,i*mult*dim)
    for j = 1:i
        Hglob[:,(j-1)*mult*dim+1:j*mult*dim]=inv(transpose(haa[:,(j-1)*mult*dim+1:j*mult*dim])*haa[:,(j-1)*mult*dim+1:j*mult*dim])*(transpose(haa[:,(j-1)*mult*dim+1:j*mult*dim])*haa[:,j*mult*dim+1:(j+1)*mult*dim])
    end
    for j = 1:i
        eigena=hcat(eigena,eigvals(Hglob[:,(j-1)*mult*dim+1:j*mult*dim]))
    end
    eigena=eigena[:,2:i+1]
    return(eigena)
end

U=eigenv(ha)

eigRe1=[real(U[n,1]) for n in 1:size(U,1)]
eigIm1=[imag(U[n,1]) for n in 1:size(U,1)]
plot1=plot(eigRe1,eigIm1,seriestype=:scatter)
Re_circle=[cos(n) for n in 0:0.05:2*pi]
Im_circle=[sin(n) for n in 0:0.05:2*pi]
plot!(plot1,Re_circle,Im_circle)

eigRe2=[real(U[n,2]) for n in 1:size(U,1)]
eigIm2=[imag(U[n,2]) for n in 1:size(U,1)]
plot2=plot(eigRe2,eigIm2,seriestype=:scatter)
plot!(plot2,Re_circle,Im_circle)

eigRe3=[real(U[n,3]) for n in 1:size(U,1)]
eigIm3=[imag(U[n,3]) for n in 1:size(U,1)]
plot3=plot(eigRe3,eigIm3,seriestype=:scatter)
plot!(plot3,Re_circle,Im_circle)

eigRe4=[real(U[n,4]) for n in 1:size(U,1)]
eigIm4=[imag(U[n,4]) for n in 1:size(U,1)]
plot4=plot(eigRe4,eigIm4,seriestype=:scatter)
plot!(plot4,Re_circle,Im_circle)

eigRe5=[real(U[n,5]) for n in 1:size(U,1)]
eigIm5=[imag(U[n,5]) for n in 1:size(U,1)]
plot5=plot(eigRe5,eigIm5,seriestype=:scatter)
plot!(plot5,Re_circle,Im_circle)


function hessenberg(A::AbstractMatrix{T},m::Int)where{T}

    n = size(A,1)
    inner_tol = 10*eps(T)

    H = zeros(T,(m+1),m)
    V = zeros(T,n,(m+1))

    rand!(view(V, :, 1))

    nor = norm(view(V, :, 1))
    view(V, :, 1) .= view(V, :, 1)./nor

    for j in 1:m
        A_mul_B!(view(V, :, (j+1)),A,view(V, :, j))
        for i in 1:j
            H[i,j] = vecdot(view(V, :, i),view(V, :, (j+1)))
            view(V, :, (j+1)) .= view(V, :, (j+1)) .- H[i,j].*view(V, :, i)
        end
        H[j+1,j] = norm(view(V, :, (j+1)))
        if abs(H[j+1,j]) < inner_tol
            println("breakdown occurs")
            break
        end
        view(V, :, j+1) .= view(V, :, (j+1))./H[j+1,j]
    end
    return H, V
end


using Plots
R = 0.0821
T = 298
isoterm(x) = R * T / x
v = 5 : -0.1 : 2.5
p = isoterm.(v)
plot(xlabel = "V[L]", ylabel = "P[atm]", xlims = xlims = (0, 7),
		ylims = ylims = (0, 14),
		xticks = 0 : 0.5 : xlims[2], yticks = 0 : 1 : ylims[2])
plot!(v, p;
		color = :red, label = "A: Trinn 1", arrow = :arrow, linewidth = 2)
plot!([v[end], v[end]], [p[end], p[1]];
		color = :blue, label = "A: Trinn 2", arrow = :arrow, linewidth = 2)
plot!([v[1], v[end]], [p[1], p[1]];
		color = :green, label = "B: Trinn 1", arrow = :arrow, linewidth = 2)
#


"""
    CorrBijector <: Bijector{2}

A bijector implementation of Stan's parametrization method for Correlation matrix:
https://mc-stan.org/docs/2_23/reference-manual/correlation-matrix-transform-section.html

Basically, a unconstrained strictly upper triangular matrix `y` is transformed to
a correlation matrix by following readable but not that efficient form:

```
K = size(y, 1)
z = tanh.(y)

for j=1:K, i=1:K
    if i>j
        w[i,j] = 0
    elseif 1==i==j
        w[i,j] = 1
    elseif 1<i==j
        w[i,j] = prod(sqrt(1 .- z[1:i-1, j].^2))
    elseif 1==i<j
        w[i,j] = z[i,j]
    elseif 1<i<j
        w[i,j] = z[i,j] * prod(sqrt(1 .- z[1:i-1, j].^2))
    end
end
```

It is easy to see that every column is a unit vector, for example:

```
w3' w3 ==
w[1,3]^2 + w[2,3]^2 + w[3,3]^2 ==
z[1,3]^2 + (z[2,3] * sqrt(1 - z[1,3]^2))^2 + (sqrt(1-z[1,3]^2) * sqrt(1-z[2,3]^2))^2 ==
z[1,3]^2 + z[2,3]^2 * (1-z[1,3]^2) + (1-z[1,3]^2) * (1-z[2,3]^2) ==
z[1,3]^2 + z[2,3]^2 - z[2,3]^2 * z[1,3]^2 + 1 -z[1,3]^2 - z[2,3]^2 + z[1,3]^2 * z[2,3]^2 ==
1
```

And diagonal elements are positive, so `w` is a cholesky factor for a positive matrix.

```
x = w' * w
```

Consider block matrix representation for `x`

```
x = [w1'; w2'; ... wn'] * [w1 w2 ... wn] ==
[w1'w1 w1'w2 ... w1'wn;
 w2'w1 w2'w2 ... w2'wn;
 ...
]
```

The diagonal elements are given by `wk'wk = 1`, thus `x` is a correlation matrix.

Every step is invertible, so this is a bijection(bijector).

Note: The implementation doesn't follow their "manageable expression" directly,
because their equation seems wrong (7/30/2020). Insteadly it follows definition
above the "manageable expression" directly, which is also described in above doc.
"""
struct CorrBijector <: Bijector{2} end

function (b::CorrBijector)(x::AbstractMatrix{<:Real})
    w = cholesky(x).U  # keep LowerTriangular until here can avoid some computation
    r = _link_chol_lkj(w)
    return r + zero(x)
    # This dense format itself is required by a test, though I can't get the point.
    # https://github.com/TuringLang/Bijectors.jl/blob/b0aaa98f90958a167a0b86c8e8eca9b95502c42d/test/transform.jl#L67
end

(b::CorrBijector)(X::AbstractArray{<:AbstractMatrix{<:Real}}) = map(b, X)

function (ib::Inverse{<:CorrBijector})(y::AbstractMatrix{<:Real})
    w = _inv_link_chol_lkj(y)
    return w' * w
end
(ib::Inverse{<:CorrBijector})(Y::AbstractArray{<:AbstractMatrix{<:Real}}) = map(ib, Y)

function logabsdetjac(::Inverse{CorrBijector}, y::AbstractMatrix{<:Real})
    K = LinearAlgebra.checksquare(y)

    result = float(zero(eltype(y)))
    for j in 2:K, i in 1:(j - 1)
        @inbounds abs_y_i_j = abs(y[i, j])
        result += (K - i + 1) * (
            IrrationalConstants.logtwo - (abs_y_i_j + LogExpFunctions.log1pexp(-2 * abs_y_i_j))
        )
    end

    return result
end
function logabsdetjac(b::CorrBijector, X::AbstractMatrix{<:Real})
    #=
    It may be more efficient if we can use un-contraint value to prevent call of b
    It's recommended to directly call
    `logabsdetjac(::Inverse{CorrBijector}, y::AbstractMatrix{<:Real})`
    if possible.
    =#
    return -logabsdetjac(inv(b), (b(X)))
end
function logabsdetjac(b::CorrBijector, X::AbstractArray{<:AbstractMatrix{<:Real}})
    return mapvcat(X) do x
        logabsdetjac(b, x)
    end
end

function _inv_link_chol_lkj(y)
    K = LinearAlgebra.checksquare(y)

    w = similar(y)

    @inbounds for j in 1:K
        w[1, j] = 1
        for i in 2:j
            z = tanh(y[i-1, j])
            tmp = w[i-1, j]
            w[i-1, j] = z * tmp
            w[i, j] = tmp * sqrt(1 - z^2)
        end
        for i in (j+1):K
            w[i, j] = 0
        end
    end

    return w
end

"""
    function _link_chol_lkj(w)

Link function for cholesky factor.

An alternative and maybe more efficient implementation was considered:

```
for i=2:K, j=(i+1):K
    z[i, j] = (w[i, j] / w[i-1, j]) * (z[i-1, j] / sqrt(1 - z[i-1, j]^2))
end
```

But this implementation will not work when w[i-1, j] = 0.
Though it is a zero measure set, unit matrix initialization will not work.

For equivelence, following explanations is given by @torfjelde:

For `(i, j)` in the loop below, we define

    z₍ᵢ₋₁, ⱼ₎ = w₍ᵢ₋₁,ⱼ₎ * ∏ₖ₌₁ⁱ⁻² (1 / √(1 - z₍ₖ,ⱼ₎²))

and so

    z₍ᵢ,ⱼ₎ = w₍ᵢ,ⱼ₎ * ∏ₖ₌₁ⁱ⁻¹ (1 / √(1 - z₍ₖ,ⱼ₎²))
            = (w₍ᵢ,ⱼ₎ * / √(1 - z₍ᵢ₋₁,ⱼ₎²)) * (∏ₖ₌₁ⁱ⁻² 1 / √(1 - z₍ₖ,ⱼ₎²))
            = (w₍ᵢ,ⱼ₎ * / √(1 - z₍ᵢ₋₁,ⱼ₎²)) * (w₍ᵢ₋₁,ⱼ₎ * ∏ₖ₌₁ⁱ⁻² 1 / √(1 - z₍ₖ,ⱼ₎²)) / w₍ᵢ₋₁,ⱼ₎
            = (w₍ᵢ,ⱼ₎ * / √(1 - z₍ᵢ₋₁,ⱼ₎²)) * (z₍ᵢ₋₁,ⱼ₎ / w₍ᵢ₋₁,ⱼ₎)
            = (w₍ᵢ,ⱼ₎ / w₍ᵢ₋₁,ⱼ₎) * (z₍ᵢ₋₁,ⱼ₎ / √(1 - z₍ᵢ₋₁,ⱼ₎²))

which is the above implementation.
"""
function _link_chol_lkj(w)
    K = LinearAlgebra.checksquare(w)

    z = similar(w) # z is also UpperTriangular.
    # Some zero filling can be avoided. Though diagnoal is still needed to be filled with zero.

    # This block can't be integrated with loop below, because w[1,1] != 0.
    @inbounds z[1, 1] = 0

    @inbounds for j=2:K
        z[1, j] = atanh(w[1, j])
        tmp = sqrt(1 - w[1, j]^2)
        for i in 2:(j - 1)
            p = w[i, j] / tmp
            tmp *= sqrt(1 - p^2)
            z[i, j] = atanh(p)
        end
        z[j, j] = 0
    end

    return z
end