module Tools
using DataFrames

export make_Y_bus, Get

function admitance(value)
    return value^-1
end

function make_Y_bus(data, VN)

    n_bus = length(collect(skipmissing(data.Bus)))
    Y = zeros(Complex, n_bus, n_bus)
    Z = zeros(Complex, n_bus, n_bus) .+ Inf

    branch = dropmissing(data[:, [:FB, :TB, :Type, :R_Ohm, :X_ohm]])
    branch = branch[branch.Type.=="Fixed", :]

    Bshunt = 1im * (data.Bshunt_MVAr * 1e6) ./ (VN^2)

    for value in eachrow(branch)
        Z[value.FB, value.TB] = value.R_Ohm + 1im * value.X_ohm
        Z[value.TB, value.FB] = value.R_Ohm + 1im * value.X_ohm
    end

    for i in 1:n_bus
        for j in 1:n_bus
            if i == j
                Y[i, j] = sum(admitance.(Z[i, :])) + Bshunt[i]
            else
                Y[i, j] = -admitance(Z[i, j])
            end
        end
    end

    return Y

end

module Get
using JuMP: value
include("Estimator.jl")
using .Estimator: build_sets

function sets(sys)
    (Ω, bΩ, L, K, D, S) = build_sets(sys)
    return (Ω, bΩ, L, K, D, S)
end

function voltage_module(model, b, l, k, s)
    return value(model[:V_module][b, l, k, s])
end

function voltage(model, b, l, k, s)
    Vre = value(model[:V][:Re, b, l, k, s])
    Vim = value(model[:V][:Im, b, l, k, s])
    return Vre + 1im * Vim
end

function current(model, b, l, k, s)
    Ire = value(model[:I][:Re, b, l, k, s])
    Iim = value(model[:I][:Im, b, l, k, s])
    return Ire + 1im * Iim
end

function power_active(model, b, l, k, s)
    return value(model[:P][b, l, k, s])
end

function power_reactive(model, b, l, k, s)
    return value(model[:Q][b, l, k, s])
end

function power_active_DER(model, d, l, k, s)
    return value(model[:pᴰᴱᴿ][d, l, k, s])
end

function power_reactive_DER(model, d, l, k, s)
    return value(model[:qᴰᴱᴿ][d, l, k, s])
end

function active_losses(model, l, k, s)
    return value(model[:Ploss][l, k, s])
end

function reactive_losses(model, l, k, s)
    return value(model[:Qloss][l, k, s])
end

end

end


function disparate_impact(ft::Fairness.FairTensor; priv_grps::Array{String, 1})
	num, den = 0, 0
	priv_count, unpriv_count = 0, 0
	for grp in ft.labels
		if grp in priv_grps
			den += sum(ft[Fairness._ftIdx(ft, grp), 2, :])
			priv_count += sum(ft[Fairness._ftIdx(ft, grp), :, :])
		else
			num += sum(ft[Fairness._ftIdx(ft, grp), 1, :])
			unpriv_count += sum(ft[Fairness._ftIdx(ft, grp), :, :])
		end
	end
	return (num/unpriv_count)/(den/priv_count)
end

struct DisparateImpact <: Fairness.Measure
	grp::Symbol
	priv_grps::Array{String, 1}
end

function (di::DisparateImpact)(ŷ, X, y)
	if typeof(ŷ[1]) <: MLJBase.UnivariateFinite
		ŷ = StatsBase.mode.(ŷ)
	end
	grps = X[:, di.grp]
	ft = fair_tensor(categorical(ŷ), categorical(y), categorical(grps))
	return disparate_impact(ft, priv_grps=di.priv_grps)
end

MLJBase.name(::Type{<:DisparateImpact}) = "DisparateImpact"
# MLJBase.target_scitype(::Type{<:DisparateImpact}) = AbstractArray{Multiclass{2},1}
MLJBase.supports_weights(::Type{<:DisparateImpact}) = false # for now
# MLJBase.prediction_type(::Type{<:DisparateImpact}) = :deterministic # Not specifying it to have check_measures false
MLJBase.orientation(::Type{<:DisparateImpact}) = :other # other options are :score, :loss
MLJBase.reports_each_observation(::Type{<:DisparateImpact}) = false
MLJBase.aggregation(::Type{<:DisparateImpact}) = Fairness.Mean()
MLJBase.is_feature_dependent(::Type{<:DisparateImpact}) = true


# Defining the fitness function
# It calls fopa.jl which returns the gain profile and some performance parameters

# The input of fitness is vector with first element corrsponding to L, second element corresponding to λ_p1, and third element corresponding to λ_p2
# Output of the fitness function is defined as weighted sum of bandwidth (BW), peak gain (Gpeak) and gain ripply (G_var)
function fitness(A,fnc)
    if fnc == "ackley"
        x = A[1];
        y = A[2];
        f = -20*exp(-0.2* ( 0.5*(x^2 + y^2) )^0.5 ) - exp( 0.5*( cos(2*pi*x) + cos(2*pi*y) )) + exp(1) + 20;
        f = -f;
    end

    if fnc == "sphere"
        f = sum(A.^2);
        f = -f;
    end

    if fnc == "rosenbrock"
        D = length(A);
        sum = 0;
        for m = 1:D-1
            sum = sum + 100*(A[m+1]-A[m]^2)^2 + (1-A[m])^2;
        end
        f = -sum;
    end

    if fnc == "rastrigin"
        D = length(A);
        sum = 0;
        for m = 1:D
            sum = sum + A[m]^2 - 10*cos(2*pi*A[m])
        end
        f = -1*(sum + 10*D);
    end

    if fnc == "double_gaussian"
        x0 = 0.1;
        y0 = -0.2;
        z0 = 0.3;

        f1 = exp(-(A[1]-x0).^2 -(A[2]-y0).^2 -(A[3]-z0).^2);

        x0 = -0.2;
        y0 = 0.2;
        z0 = -0.1;

        f2 = exp(-(A[1]-x0).^2 -(A[2]-y0).^2 -(A[3]-z0).^2);
        f =f1 + 0.6*f2;
    end
    return f
end

# A = [.1 -.2 .3];     # test input

# println("Testing the fitness function:")
# println("Test input = ", A)
# println("Fitness value = ", fitness(A))


using Catlab.CategoricalAlgebra, Catlab.Graphs, Catlab.Present, Catlab.Graphics, Catlab.Theories
using Catlab.CategoricalAlgebra.FinCats: FinCatGraphEq

using Distributions: Exponential, Geometric, cdf
using Random: randsubseq
using StatsBase: sample

using Plots
using ProgressBars

include("utils.jl")

# schema
@present ThSIR(FreeSchema) begin
    (S,I,R)::Ob
end

@acset_type SIR(ThSIR)

# infection rules
I = @acset SIR begin I=1 end
I2 = @acset SIR begin I=2 end
SI = @acset SIR begin S=1; I=1 end
L_infect =  ACSetTransformation(I, SI; I=[1]) # fn from dom I to codom SI, I is the injective function mapping stuff in dom to codom
R_infect = ACSetTransformation(I, I2; I=[1])

# recovery rules
R = @acset SIR begin R=1 end
L_recovery = ACSetTransformation(SIR(), I)
R_recovery = ACSetTransformation(SIR(), R)

# parameters
N = 1000
I0 = 5
S0 = N - I0

Δt = 0.1
tmax = 100
steps = Int(tmax/Δt)
γ = 1/10 # recovery rate
R0 = 2.5
β = R0 * γ # R0 for corresponding ODEs

state = @acset SIR begin S=S0; I=I0; R=N-S0-I0 end

# simulation loop
out = fill(-1, (steps, 3))

for t = ProgressBar(1:steps)
    # matches
    infections_m = homomorphisms(SI, state)
    recoveries_m = homomorphisms(I, state)

    # queued updates objects
    infections = queued_updates(infections_m, L_infect, R_infect)
    recoveries = queued_updates(recoveries_m, L_recovery, R_recovery)

    # sample occurances
    sample_matches(infections, β/N, Δt)
    sample_matches(recoveries, γ, Δt)

    global state = fire_events(state, [infections, recoveries])

    # write output
    out[t, :] = [nparts(state, x) for x in [:S, :I, :R]]
end

plot(
    (1:steps) * Δt,
    out,
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number"
)


export AffineChartSystem, on_affine_chart

"""
    AffineChartSystem(F::AbstractSystem, v::PVector{T,N})

Given a system ``F(x): (ℙ^{m_1} × ⋯ × ℙ^{m_N}) → ℂⁿ`` this creates a new affine
system ``F′`` which operates on the affine chart defined by the vector
``v ∈ ℙ^{m_1} × ⋯ × ℙ^{m_N}`` and the augmented conditions ``vᵀx = 1``.
"""
struct AffineChartSystem{S<:AbstractSystem,N} <: AbstractSystem
    system::S
    chart::PVector{ComplexF64,N}
end

ModelKit.variables(F::AffineChartSystem) = variables(F.system)
ModelKit.parameters(F::AffineChartSystem) = parameters(F.system)
ModelKit.variable_groups(F::AffineChartSystem) = variable_groups(F.system)

"""
    on_affine_chart(F::Union{System,AbstractSystem}, dimensions)

Construct an `AffineChartSystem` on a randomly generated chart `v`. Each entry is drawn
idepdently from a univariate normal distribution.
"""
on_affine_chart(F::System, dims = nothing) = on_affine_chart(ModelKitSystem(F), dims)
function on_affine_chart(F::AbstractSystem, dims = nothing)
    vargroups = variable_groups(F)
    if vargroups === nothing
        dims = [size(F, 2) - 1]
    else
        dims = length.(vargroups,) .- 1
    end
    chart = PVector(randn(ComplexF64, sum(dims) + length(dims)), tuple(dims...))
    AffineChartSystem(F, chart)
end

function Base.size(F::AffineChartSystem{<:AbstractSystem,N}) where {N}
    m, n = size(F.system)
    (m + N, n)
end

function evaluate_chart!(u, v::PVector{<:Any,N}, x::AbstractVector) where {N}
    ranges = dimension_indices(v)
    n = length(u) - N
    for (k, range) in enumerate(ranges)
        out = zero(promote_type(eltype(u), eltype(x)))
        @inbounds for i in range
            out += v[i] * x[i]
        end
        u[n+k] = out - 1.0
    end
    nothing
end

function jacobian_chart!(U, v::PVector{<:Any,N}, x::AbstractVector) where {N}
    ranges = dimension_indices(v)
    n = size(U, 1) - N
    for j = 1:size(U, 2), i = (n+1):size(U, 1)
        U[i, j] = zero(eltype(U))
    end
    for (k, range) in enumerate(ranges)
        for j in range
            U[n+k, j] = v[j]
        end
    end
    nothing
end

function (F::AffineChartSystem{<:Any,N})(x::AbstractVector, p = nothing) where {N}
    v = PVector(x, dims(F.chart))
    u = F.system(x, p)
    append!(u, zeros(N))
    evaluate_chart!(u, F.chart, v)
    u
end
function (F::AffineChartSystem{<:Any,N})(x, p = nothing) where {N}
    u = F.system(x, p)
    append!(u, zeros(N))
    evaluate_chart!(u, F.chart, x)
    u
end

function evaluate!(u, F::AffineChartSystem{<:Any,N}, x, p = nothing) where {N}
    evaluate!(u, F.system, x)
    evaluate_chart!(u, F.chart, x)
    u
end

function evaluate_and_jacobian!(
    u,
    U,
    F::AffineChartSystem{<:Any,N},
    x,
    p = nothing,
) where {N}
    evaluate_and_jacobian!(u, U, F.system, x)
    evaluate_chart!(u, F.chart, x)
    jacobian_chart!(U, F.chart, x)
    nothing
end

function taylor!(u, v::Val, F::AffineChartSystem, tx, p = nothing)
    u .= zero(eltype(u))
    taylor!(u, v, F.system, tx, p)
    # affine chart part is always zero since it is an affine linear form
    u
end


using LinearAlgebra
using Test

using SetProg, SetProg.Sets
using Polyhedra
using MultivariatePolynomials

using JuMP
const MOIT = MOI.Test

@testset "Controlled invariant" begin
    config = MOIT.Config()
    @testset "Ellipsoid" begin
        # Q = [1 0
        #      0 1]
        # t = √det(Q) = 1
        Q = [1.0, -1/4, 1.0]
        t = √15/4
        @testset "Homogeneous" begin
            Tests.ci_ell_homogeneous_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, [Q; t])),
                                          config)
        end
        @testset "Non-homogeneous" begin
            @testset "Ellipsoid" begin
                β = -1.0
                b = [0.0, 0.0]
                Tests.ci_ell_nonhomogeneous_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, [Q; β; b; t])),
                                                 config)
            end
            @testset "PolySet" begin
                β = -1.0
                b = [0.0, 0.0]
                Q = [1.0, -0.5, 1.0]

                Tests.ci_quad_nonhomogeneous_test(bridged_mock(mock -> begin
                         # β+1 b[1] b[2]
                         #  .  Q[1] Q[2]
                         #  .   .   Q[3]
                         MOI.Utilities.mock_optimize!(mock, [β+1; b[1]; Q[1]; b[2]; Q[2]; Q[3]; 2Q])
                     end),
                config)
            end
        end
        @testset "Piecewise" begin
            @testset "Diamond/2D cross-polytope" begin
                Q1 = [1.0, -0.25, 1.0]
                Q2 = [1.0, -1.0, 1.0]
                Tests.ci_piecewise_semiell_homogeneous_test(
                    bridged_mock(mock -> MOI.Utilities.mock_optimize!(
                        mock,
                        [Q1; Q2; Q2; Q1; collect(1:MOI.get(mock, MOI.NumberOfVariables()) - 12)])),
                    config)
            end
            @testset "mci" begin
                Q1 = [1.0, -1.0, 1.0]
                Q2 = [1.0, 0.0, 0.0]
                Q3 = [0.25, 0.5, 1.0]
                # FIXME See https://github.com/blegat/SetProg.jl/runs/4384813153?check_suite_focus=true
                if Sys.WORD_SIZE != 32
                    Tests.ci_piecewise_semiell_mci_homogeneous_test(
                        bridged_mock(mock -> MOI.Utilities.mock_optimize!(
                            mock,
                            [Q1; Q1; Q2; Q3; Q2; Q3; collect(1:MOI.get(mock, MOI.NumberOfVariables()) - 18)])),
                        config)
                end
            end
        end
    end
    @testset "Quartic" begin
        # The PSD matrix for the variable  is 3 x 3 so 3 * (3+1) / 2 = 6
        # The PSD matrix for the convexity is 6 x 6 so 6 * (6+1) / 2 = 21
        # entries
        # 1 variable for t
        # hence 28 variables
        ci_quartic_α = -1/8
        ci_quartic_β =  1/4
        ci_quartic_hess = 6 * [2.0, ci_quartic_α, 2.0, ci_quartic_α, 2.0,
                               2.0, 2.0, ci_quartic_α, ci_quartic_α, 2.0]
        sol = [1.0; ci_quartic_α; 6 - 2ci_quartic_β; ci_quartic_β; ci_quartic_α; 1.0;
               ci_quartic_hess]
        Tests.ci_quartic_homogeneous_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, sol)),
                                          config)
    end
end


#= generate_testing dataset.jl script is used to generate data that will be use for testing.
1. Simulates dataset
2. Apply MatrixLMnet v0.1.0 L1-penalized
3. Save datasets and estimates to be comapred with newer version =#

###########
# Library #
###########
# using Distributions, Random, Statistics, LinearAlgebra, StatsBase
# using DataFrames, MLBase, Distributed
using MatrixLMnet #v0.1.0
# using Test
using Helium

####################
# External sources #
####################
include("sim_helpers.jl")
# include("l1_helpers.jl")

###################
#  Simulated Data #
###################

#=
Description:
-----------

Model: 𝐘 = 𝐗 𝛃 𝐙' + 𝜎𝜖, with 𝜖∼𝑁(0,1)

Simulate data set consisting  of 20/20/200 observations and 8 predictors.
We let 𝛽₁ = (3, 1.5, 0, 0, 2, 0, 0, 0), 𝛽₂ = (0, 1.5, 0, 3.5, 2, 0, 0 , 2) where
𝛃 = [𝛽₁, 𝛽₂] and 𝜎 = 3.
The pairwise correlation between 𝑋ᵢ and 𝑋ⱼ was set to be 𝑐𝑜𝑟(𝑖,𝑗)=(0.5)^|𝑖−𝑗|.
Here, the Z matrix is an identity matrix.
=#

rng = MersenneTwister(2021)

# Simulation parameters
p = 8; # Number of predictors
β1 = [3.5, 1.5, 0,   0, 2, 0, 0 , 0];
β2 = [  0, 1.5, 0, 3.5, 2, 0, 0 , 2];
σ = 3;
n = 240;

# Generate correlation matrix
matCor = zeros(p,p)
for j = 1:p, i = 1:p
    matCor[i,j] = 0.5^abs(i-j)
end

# Generate predictors
X = simulateCorrelatedData(matCor, n);

# Generate response
Y1 = X*β1 + σ*rand(Normal(0, 1), n);
Y2 = X*β2 + σ*rand(Normal(0, 1), n);
Y = hcat(Y1, Y2);

# Generate Z matrix
Z = 1.0*Matrix(I, 2, 2);

# Build raw data object from MatrixLM.jl
dat = MatrixLMnet.MatrixLM.RawData(Response(Y), Predictors(X, Z));

# Hyper parameters
λ = [10.0, 5.0, 3.0]

###############
# TEST Lasso  #
###############

# Lasso penalized regression - ista
est = mlmnet(ista!, dat, λ, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
est_B_Lasso_ista = est.B[:, :, 3];

# Lasso penalized regression - fista
est = mlmnet(fista!, dat, λ, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
est_B_Lasso_fista = est.B[:, :, 3];

# Lasso penalized regression - fista backtracking
est = mlmnet(fista_bt!, dat, λ, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
est_B_Lasso_fista_bt = est.B[:, :, 3];

# Lasso penalized regression - admm
est = mlmnet(admm!, dat, λ, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
est_B_Lasso_admm = est.B[:, :, 3];

# Lasso penalized regression - cd
Random.seed!(rng)
est = mlmnet(cd!, dat, λ, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
est_B_Lasso_cd = est.B[:, :, 3];

#################################
# TEST Lasso - Crossvalidation  #
#################################

# Lasso penalized regression - ista cv
Random.seed!(rng)
est = mlmnet_cv(ista!, dat, λ, 10, 1, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
smmr_Lasso = lambda_min_deprecated(est);
smmr_ista = hcat(smmr_Lasso.AvgMSE, smmr_Lasso.AvgPercentZero)

# Lasso penalized regression - fista cv
Random.seed!(rng)
est = mlmnet_cv(fista!, dat, λ, 10, 1, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
smmr_Lasso = lambda_min_deprecated(est);
smmr_fista = hcat(smmr_Lasso.AvgMSE, smmr_Lasso.AvgPercentZero)

# Lasso penalized regression - fista-bt cv
Random.seed!(rng)
est = mlmnet_cv(fista_bt!, dat, λ, 10, 1, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
smmr_Lasso = lambda_min_deprecated(est);
smmr_fistabt = hcat(smmr_Lasso.AvgMSE, smmr_Lasso.AvgPercentZero)

# Lasso penalized regression - admm cv
Random.seed!(rng)
est = mlmnet_cv(admm!, dat, λ, 10, 1, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
smmr_Lasso = lambda_min_deprecated(est);
smmr_admm = hcat(smmr_Lasso.AvgMSE, smmr_Lasso.AvgPercentZero)

# Lasso penalized regression - cd cv
Random.seed!(rng)
est = mlmnet_cv(cd!, dat, λ, 10, 1, hasZIntercept = false, hasXIntercept = false, isVerbose = false);
smmr_Lasso = lambda_min_deprecated(est);
smmr_cd = hcat(smmr_Lasso.AvgMSE, smmr_Lasso.AvgPercentZero)

###############################
#  Save Dataset and Estimates #
###############################

# Output directory name
dataDir = realpath(joinpath(@__DIR__,"data"))

# Save X, Y, and Z
Helium.writehe(X, joinpath(dataDir, "Xmat.he"))
Helium.writehe(Y, joinpath(dataDir, "Ymat.he"))
Helium.writehe(Z, joinpath(dataDir, "Zmat.he"))

# Save estimates results
Helium.writehe(est_B_Lasso_ista, joinpath(dataDir, "B_ista.he"))
Helium.writehe(est_B_Lasso_fista, joinpath(dataDir, "B_fista.he"))
Helium.writehe(est_B_Lasso_fista_bt, joinpath(dataDir, "B_fistabt.he"))
Helium.writehe(est_B_Lasso_admm, joinpath(dataDir, "B_admm.he"))
Helium.writehe(est_B_Lasso_cd, joinpath(dataDir, "B_cd.he"))

Helium.writehe(smmr_ista, joinpath(dataDir, "smmr_ista.he"))
Helium.writehe(smmr_fista, joinpath(dataDir, "smmr_fista.he"))
Helium.writehe(smmr_fistabt, joinpath(dataDir, "smmr_fistabt.he"))
Helium.writehe(smmr_admm, joinpath(dataDir, "smmr_admm.he"))
Helium.writehe(smmr_cd, joinpath(dataDir, "smmr_cd.he"))


"""
```
smooth_ait_gdp()
```

implements flexible AIT policy, which combines AIT with GDP targeting and discounting of the history.
The "smooth" part in the name describes how the degree to which history is discounted smooths
the response of monetary policy to deviations from the average inflation and GDP targets.
"""
function smooth_ait_gdp()
    AltPolicy(:smooth_ait_gdp, smooth_ait_gdp_eqcond, smooth_ait_gdp_solve,
              forecast_init = smooth_ait_gdp_forecast_init,
              color = RGB(0., 0., 0.5430)) # dark blue
end

function smooth_ait_gdp_replace_eq_entries(m::AbstractDSGEModel,
                                            Γ0::Matrix{Float64}, Γ1::Matrix{Float64},
                                            C::Vector{Float64}, Ψ::Matrix{Float64}, Π::Matrix{Float64})
    # add law of motion for pgap
    eq        = m.equilibrium_conditions
    endo      = m.endogenous_states

    ait_Thalf = haskey(get_settings(m), :ait_Thalf) ? get_setting(m, :ait_Thalf) : 10.
    gdp_Thalf = haskey(get_settings(m), :gdp_Thalf) ? get_setting(m, :gdp_Thalf) : 10.
    ρ_pgap    = exp(log(0.5) / ait_Thalf)
    ρ_ygap    = exp(log(0.5) / gdp_Thalf)
    ρ_smooth  = haskey(get_settings(m), :smooth_ait_gdp_ρ_smooth) ? get_setting(m, :smooth_ait_gdp_ρ_smooth) : 0.656 # m[:ρ]
    φ_π       = haskey(get_settings(m), :smooth_ait_gdp_φ_π) ? get_setting(m, :smooth_ait_gdp_φ_π) : 0.25
    φ_y       = haskey(get_settings(m), :smooth_ait_gdp_φ_y) ? get_setting(m, :smooth_ait_gdp_φ_y) : 0.25

    # This assumes that the inflation target is the model's steady state
    Γ0[eq[:eq_pgap], endo[:pgap_t]] = 1.
    Γ0[eq[:eq_pgap], endo[:π_t]]    = -1.
    Γ1[eq[:eq_pgap], endo[:pgap_t]] = ρ_pgap

    # Add in the GDP targeting rule
    Γ0[eq[:eq_ygap], endo[:ygap_t]] = 1.
    # Γ0[eq[:eq_ygap], endo[:π_t]]    = -1. # already accounted for by AIT
    Γ1[eq[:eq_ygap], endo[:ygap_t]] = ρ_ygap # 1.

    Γ0[eq[:eq_ygap], endo[:y_t]]    = -1.
    Γ0[eq[:eq_ygap], endo[:z_t]]    = -1.
    Γ1[eq[:eq_ygap], endo[:y_t]]    = -1.

    # Zero out old policy rule
    Γ0[eq[:eq_mp], :]              .= 0.
    Γ1[eq[:eq_mp], :]              .= 0.
    Ψ[eq[:eq_mp], :]               .= 0.

    # Replace monetary policy rule
    Γ0[eq[:eq_mp], endo[:R_t]]      = 1.
    Γ1[eq[:eq_mp], endo[:R_t]]      = ρ_smooth
    C[ eq[:eq_mp]]                  = 0.

    Γ0[eq[:eq_mp], endo[:pgap_t]]   = -φ_π * (1. / (1. - ρ_pgap)) * (1. - ρ_smooth) # (1 - m[:ρ]) # This is the AIT part
    Γ0[eq[:eq_mp], endo[:ygap_t]]   = -φ_y * (1. / (1. - ρ_ygap)) * (1. - ρ_smooth) # (1. - m[:ρ]) # This is the GDP part

    return Γ0, Γ1, C, Ψ, Π
end

"""
```
smooth_ait_gdp_eqcond(m::AbstractDSGEModel)
```

Solves for the transition equation of `m` under a price level
targeting rule (implemented by adding a price-gap state)
"""
function smooth_ait_gdp_eqcond(m::AbstractDSGEModel, reg::Int = 1)
    # get the old indices
    old_states = sort!(collect(values(m.endogenous_states)))
    old_eqs    = sort!(collect(values(m.equilibrium_conditions)))

    # Get equilibrium condition matrices
    Γ0_noaltpol, Γ1_noaltpol, C_noaltpol, Ψ_noaltpol, Π_noaltpol = eqcond(m, reg)

    # check to make sure this state hasn't already been added, then:
    # 1. add the pgap state to the model
    # 2. increment the indices of every augmented state

    if !in(:pgap_t, keys(m.endogenous_states))
        # add pgap state to endogenous_staets
        m.endogenous_states[:pgap_t]       = n_states(m) + 1
        m.equilibrium_conditions[:eq_pgap] = n_equilibrium_conditions(m) + 1
        # now n_states(m) and n_equilibrium_conditions(m) will calculate appropriately

        # increment every state in endogenous_states_augmented
        aug = m.endogenous_states_augmented
        map(x -> aug[x] = aug[x]+1, collect(keys(aug)))
    end
    if !in(:ygap_t, keys(m.endogenous_states))
        # add pgap state to endogenous_staets
        m.endogenous_states[:ygap_t]       = n_states(m) + 1
        m.equilibrium_conditions[:eq_ygap] = n_equilibrium_conditions(m) + 1
        # now n_states(m) and n_equilibrium_conditions(m) will calculate appropriately

        # increment every state in endogenous_states_augmented
        aug = m.endogenous_states_augmented
        map(x -> aug[x] = aug[x]+1, collect(keys(aug)))
    end

    nstates = n_states(m)

    # fill in new Γ0, Γ1, C, Ψ, Π
    Γ0 = zeros(Float64, nstates, nstates)
    Γ0[old_eqs, old_states] = Γ0_noaltpol

    Γ1 = zeros(Float64, nstates, nstates)
    Γ1[old_eqs, old_states] = Γ1_noaltpol

    C  = zeros(Float64, nstates)
    C[old_eqs, :] = C_noaltpol

    Ψ  = zeros(Float64, nstates, n_shocks_exogenous(m))
    Ψ[old_eqs, :] = Ψ_noaltpol

    Π  = zeros(Float64, nstates, n_shocks_expectational(m))
    Π[old_eqs, :] = Π_noaltpol

    for para in m.parameters
        if !isempty(para.regimes)
            ModelConstructors.toggle_regime!(para, reg)
        end
    end

    Γ0, Γ1, C, Ψ, Π = smooth_ait_gdp_replace_eq_entries(m, Γ0, Γ1, C, Ψ, Π)

    for para in m.parameters
        if !isempty(para.regimes)
            ModelConstructors.toggle_regime!(para, 1)
        end
    end

    return Γ0, Γ1, C, Ψ, Π
end

"""
```
smooth_ait_gdp_solve(m::AbstractDSGEModel)
```

Solves for the transition equation of `m` under a price level
targeting rule (implemented by adding a price-gap state)
"""
function smooth_ait_gdp_solve(m::AbstractDSGEModel; regime_switching::Bool = false, regimes::Vector{Int} = Int[1])
    # Get equilibrium condition matrices

    if length(regimes) == 1
        Γ0, Γ1, C, Ψ, Π  = smooth_ait_gdp_eqcond(m, regimes[1])
        TTT_gensys, CCC_gensys, RRR_gensys, eu = gensys(Γ0, Γ1, C, Ψ, Π, 1+1e-6, verbose = :low)

        # Check for LAPACK exception, existence and uniqueness
        if eu[1] != 1 || eu[2] != 1
            throw(GensysError())
        end

        TTT_gensys = real(TTT_gensys)
        RRR_gensys = real(RRR_gensys)
        CCC_gensys = real(CCC_gensys)

        # Augment states
        TTT, RRR, CCC = DSGE.augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys; regime_switching = regime_switching,
                                            reg = regimes[1])
        return TTT, RRR, CCC
    else
        Γ0s = Vector{Matrix{Float64}}(undef, length(regimes))
        Γ1s = Vector{Matrix{Float64}}(undef, length(regimes))
        Cs = Vector{Vector{Float64}}(undef, length(regimes))
        Ψs = Vector{Matrix{Float64}}(undef, length(regimes))
        Πs = Vector{Matrix{Float64}}(undef, length(regimes))
        for reg in regimes
            Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg]  = smooth_ait_gdp_eqcond(m, reg)
        end

        n_regimes = length(regimes)
        TTTs= Vector{Matrix{Float64}}(undef, n_regimes)
        RRRs = Vector{Matrix{Float64}}(undef, n_regimes)
        CCCs = Vector{Vector{Float64}}(undef, n_regimes)

        # Solve model
        for reg in regimes
            TTT_gensys, CCC_gensys, RRR_gensys, eu = gensys(Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg], 1+1e-6)

            if !((eu[1] == 1) & (eu[2] == 1))
                throw(GensysError("Gensys does not give existence"))
            end
            TTT_gensys = real(TTT_gensys)
            RRR_gensys = real(RRR_gensys)
            CCC_gensys = reshape(CCC_gensys, size(CCC_gensys, 1))

            # Augment states
            TTTs[reg], RRRs[reg], CCCs[reg] = DSGE.augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys;
                                                             regime_switching = regime_switching,
                                                             reg = reg)
        end

        return TTTs, RRRs, CCCs
    end
end

"""
```
init_smooth_ait_gdp_forecast(m::AbstractDSGEModel, shocks::Matrix{T}, final_state::Vector{T})
```

Adjust shocks matrix and final state vector for forecasting under the SMOOTH_AIT_GDP rule
"""
function smooth_ait_gdp_forecast_init(m::AbstractDSGEModel, shocks::Matrix{T}, final_state::Vector{T};
                                       cond_type::Symbol = :none) where {T<:AbstractFloat}

    pgap_t  = m.endogenous_states[:pgap_t]
    ygap_t = m.endogenous_states[:ygap_t]

    final_state = if pgap_t < ygap_t
        vcat(final_state[1:pgap_t - 1], -get_setting(m, :pgap_value), final_state[pgap_t + 1:ygap_t - 1],
             -get_setting(m, :ygap_value), final_state[ygap_t + 1:end])
    else
        vcat(final_state[1:ygap_t - 1], -get_setting(m, :ygap_value), final_state[ygap_t + 1:pgap_t - 1],
             -get_setting(m, :pgap_value), final_state[pgap_t + 1:end])
    end

    return shocks, final_state
end


using Echidna
using Plots

gr()

f1(θ) = [cos(θ), sin(θ)]

pts = [f1(π/2 * x) for x in rand(10)]
x = [p[1] for p in pts]
y = [p[2] for p in pts]
scatter(x, y)

zmin = [1.0, 1.0]
extremepts = find_extreme_points(pts, zmin)

xx = [x[i] for i in extremepts]
yy = [y[i] for i in extremepts]
scatter!(xx, yy)

f2(x) = [x, exp(-x)]
pts = [f2(5*x) for x in rand(10)]
x = [p[1] for p in pts]
y = [p[2] for p in pts]
scatter(x, y)

extremepts = find_extreme_points(pts)

xx = [x[i] for i in extremepts]
yy = [y[i] for i in extremepts]
scatter!(xx, yy)

zpts = [pts[i] for i in extremepts]
find_axis_intercepts(zpts)

#### TESTING POINT ASSOCIATION
canpoints = [[0.5, 1.0], [2.0, 0.0]]
refpoints = [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]

index, dists = associate_points(canpoints, refpoints)


# Julia1_3.6.jl Modules
# defining a Module

module Gravity      # introduces new global scope
#using somelib       # presuming somelib contains something you need

#import Base.show # if you just want to use show you could use using

# import allows you to extend show with new methods
g=9.80665           # m/s^2
v(t) = g*t            # velocity as a function of time
s(t) = 0.5*g*t^2   # Distance as a function of time
export v, s, g      # make these functions and variables global
end

using Main.Gravity
for t=0.0:0.5:5.0
    print("t=$t v=$(v(t))\n      s=$(s(t)) \n")
end
#
###
#
module Astronomy
# constants
const c = 299792.458     # speed of light km/s
const dsec=24.0*3600.0   # seconds in a day
const ly = 365.25*dsec*c # light year in km
const au = 149597870.7   # astronomical unit in km (Radius of earths orbit)
const parsec = au*648000.0/π   # parsec in km (Distance to create parallax
#                              of 1 arcsecond of parallax from earths orbit)
export c
# functions
parsectoly(ps)= ps*parsec/ly  #  cnv parsecs to light years
stardis(parallax)=parsectoly(1.0/parallax)  # calc star dist
# An angle of 1 arcsecond of a degree  1/3600 of a degree
# from an object
export stardis
end

using Printf
using Main.Astronomy
nearstars = Dict("ProximaCentauri"=>.76850, "AlphaCentauri"=>.74723,
 "Bernard's Star"=>.54745, "Luhman16"=>.501398,
 "Wise0855-714"=>.449, "Wolf"=>.41516,
 "Lalande21185"=>.30264, "Sirius A"=>.37668)
@printf("       Star name   parallax     d.lightyears\n")
for star in nearstars
    name,plax=star
    @printf("%16s %10.6f %10.4f\n",name,plax,stardis(plax))
end

println(ly)  # error since not exported

println(Astronomy.ly) # But can still get ly using Module name
println(Astronomy.parsectoly(1.0))  # same with functions


const USE_GPU = true
using ImplicitGlobalGrid
using ParallelStencil
using ParallelStencil.FiniteDifferences3D
@static if USE_GPU
    @init_parallel_stencil(package=CUDA, ndims=3);
else
    @init_parallel_stencil(package=Threads, ndims=3);
end

@parallel function diffusion3D_step!(T2::Data.Array{_T}, T::Data.Array{_T}, Ci::Data.Array{_T}, lam::_T, dt::_T, _dx::_T, _dy::_T, _dz::_T) where _T <: PSNumber
    @inn(T2) = @inn(T) + dt*(lam*@inn(Ci)*(@d2_xi(T)*_dx^2 + @d2_yi(T)*_dy^2 + @d2_zi(T)*_dz^2));
    return
end

function diffusion3D(lam::_T, c0::_T, lx::_T, ly::_T, lz::_T) where _T <: PSNumber
    # Numerics
    nx, ny, nz = 512, 512, 512;                              # Number of gridpoints in dimensions x, y and z
    nt         = 100;                                        # Number of time steps
    me, dims   = init_global_grid(nx, ny, nz);
    dx         = lx/(nx_g()-1);                              # Space step in x-dimension
    dy         = ly/(ny_g()-1);                              # Space step in y-dimension
    dz         = lz/(nz_g()-1);                              # Space step in z-dimension
    _dx, _dy, _dz = 1.0/dx, 1.0/dy, 1.0/dz;

    # Array initializations
    T   = @zeros(_T, nx, ny, nz);
    T2  = @zeros(_T, nx, ny, nz);
    Ci  = @zeros(_T, nx, ny, nz);

    # Initial conditions
    Ci .= 1/c0;                                              # 1/Heat capacity
    T  .= 1.7;
    T2 .= T;                                                 # Assign also T2 to get correct boundary conditions.

    # Time loop
    dt   = min(dx^2,dy^2,dz^2)/lam/maximum(Ci)/8.1;          # Time step for 3D Heat diffusion
    lam, dt, _dx, _dy, _dz = _T.((lam, dt, _dx, _dy, _dz))   # Convert all scalars to the required numbertype.
    for it = 1:nt
        if (it == 11) tic(); end                             # Start measuring time.
        @parallel diffusion3D_step!(T2, T, Ci, lam, dt, _dx, _dy, _dz);
        update_halo!(T2);
        T, T2 = T2, T;
    end
    time_s = toc()

    # Performance
    A_eff = (2*1+1)*1/1e9*nx*ny*nz*sizeof(_T);               # Effective main memory access per iteration [GB] (Lower bound of required memory access: T has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
    t_it  = time_s/(nt-10);                                  # Execution time per iteration [s]
    T_eff = A_eff/t_it;                                      # Effective memory throughput [GB/s]
    if (me==0) println("time_s=$time_s T_eff=$T_eff"); end

    finalize_global_grid();
end

################################################################################
# CALL OF THE FUNCTION

# Physics
lam        = 1.0;                                        # Thermal conductivity
c0         = 2.0;                                        # Heat capacity
lx, ly, lz = 1.0, 1.0, 1.0;                              # Length of computational domain in dimension x, y and z

lam, c0, lx, ly, lz = Float64.((lam, c0, lx, ly, lz))    # Convert all scalars to the desired numbertype.

diffusion3D(lam, c0, lx, ly, lz)


# This function converts GPS week and day to T_TAI (MJD referenced to International Atomic Time)
function gps2MJD(week, day)
    # GPS Epoch in Mod. Julian Date
    gpsMJDEpoch = 2444244.5 - 2400000.5

    # Add GPS week and days
    MJD = gpsMJDEpoch + week*7.0 + day + 19.0 / 86400.0

    return MJD
end


#!/usr/bin/env julia

module PE18

const str=
"75
95 64
17 47 82
18 35 87 10
20 04 82 47 65
19 01 23 75 03 34
88 02 77 73 07 63 67
99 65 04 28 06 16 70 92
41 41 26 56 83 40 80 70 33
41 48 72 33 47 32 37 16 94 29
53 71 44 65 25 43 91 52 97 51 14
70 11 33 28 77 73 17 78 39 68 17 57
91 71 52 38 17 14 91 43 58 50 27 29 48
63 66 04 68 89 53 67 30 73 16 69 87 40 31
04 62 98 27 23 09 70 98 73 93 38 53 60 04 23"

const data = map(split(str,"\n")) do line
    parse.(Int,split(line," "))
end

end

export pe18

function pe18(data=PE18.data)
    foldr(data) do next, prev
        map(max,prev[1:end-1].+next,prev[2:end].+next)
    end |> first
end

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe18()|>print
end


using DelimitedFiles
using IterTools

function find_tuple_2020(arr, size, def)
    for t in subsets(arr, size)
        if sum(t) == 2020
            return t
        end
    end
    def
end

day1part1(arr) = prod(find_tuple_2020(arr, 2, (0,0)))
day1part2(arr) = prod(find_tuple_2020(arr, 3, (0,0,0)))

function day1(path)
    arr = readdlm("day01/input.txt", Int)
    p1 = day1part1(arr)
    p2 = day1part2(arr)
    (p1, p2)
end

#=
a = [1000, 1721, 979, 366, 299, 675, 1456]
subsets(a, 2)

find_tuple_2020([1000, 1721, 979, 366, 299, 675, 1456], 2, (0,0))

day1("day01/day01.jl")
=#


@testset "Grid generation" begin

    @testset "Regular pattern" begin

        @test collect(MicroscopyTools.regular_pattern((4, 4), offset = -1, step = 2)) == [(-1, -1) (-1, 1); (1, -1) (1, 1)]

        @test collect(MicroscopyTools.regular_pattern((4, 4), step = 2)) == [(1, 1) (1, 3); (3, 1) (3, 3)]

        @test collect(MicroscopyTools.regular_pattern((3, 3))) == [(1, 1) (1, 2) (1, 3); (2, 1) (2, 2) (2, 3); (3, 1) (3, 2) (3, 3)]

    end

end


abstract BasisFun{T<:Number}
abstract BasisFun1d{T<:Number} <: BasisFun{T}

" Total number of modes of a basis function "
nmodes{T<:BasisFun}(b::T) = b.nmodes

" Total number of modes of a basis function "
length{T<:BasisFun1d}(b::T) = nmodes(b)

" Total number of modes of a basis function "
size{T<:BasisFun1d}(b::T) = (nmodes(b),)

" Can the modes be decoupled into boundary/interior modes? "
isbndryint{T<:BasisFun}(b::T) = false
isbndryint{T<:BasisFun}(::Type{T}) = false

bndidx{T<:BasisFun}(b::T) = bndidx(b.lnum)
nbndry{T<:BasisFun}(b::T) = nbndry(b.lnum)

intidx{T<:BasisFun}(b::T) = intidx(b.lnum)
ninterior{T<:BasisFun}(b::T) = ninterior(b.lnum)

seq2bi!{T}(b::BasisFun, x::AbstractVector{T}, y::AbstractVector{T}) = seq2bi!(b.lnum, x, y)
seq2bi{T}(b::BasisFun, x::AbstractVector{T}) = seq2bi(b.lnum, x)
bi2seq!{T}(b::BasisFun, x::AbstractVector{T}, y::AbstractVector{T}) = bi2seq!(b.lnum, x, y)
bi2seq{T}(b::BasisFun, x::AbstractVector{T}) = bi2seq(b.lnum, x)
bi2seq!{T}(b::BasisFun, xb::AbstractVector{T}, xi::AbstractVector{T}, y::AbstractVector{T}) =
    bi2seq!(b.lnum, xb, xi, y)
bi2seq{T}(b::BasisFun, xb::AbstractVector{T}, xi::AbstractVector{T}) = bi2seq(b.lnum, xb, xi)
seq2b!{T}(b::BasisFun, x::AbstractVector{T}, y::AbstractVector{T}) = seq2b!(b.lnum, x, y)
seq2i!{T}(b::BasisFun, x::AbstractVector{T}, y::AbstractVector{T}) = seq2i!(b.lnum, x, y)
seq2b{T}(b::BasisFun, x::AbstractVector{T}) = seq2b(b.lnum, x)
seq2i{T}(b::BasisFun, x::AbstractVector{T}) = seq2i(b.lnum, x)

"""
Compute the value of mode `i` at all points ξ

The values are returned on array `y`.
"""
function basis1d!{T<:BasisFun1d}(b::T, ξ::AbstractArray, y::AbstractArray, p::Integer)
    for i = 1:length(ξ)
        y[i] = basis1d(b, ξ[i], p)
    end
    y
end

basis1d{T<:BasisFun1d}(b::T, ξ::AbstractArray, p::Integer) = basis1d!(b, ξ, similar(ξ), p)

"""
Karniadakis/Sherwin 1D modal polynomial basis.
"""
immutable ModalC01d{T<:Number} <: BasisFun1d{T}
    nmodes::Int
    lnum::LocalNumSys1d
end
ModalC01d{T<:Number}(n, ::Type{T}=Float64) = ModalC01d{T}(n, LocalNumSys1d([1,2], [3:n;]))

isbndryint(b::ModalC01d) = true
isbndryint(::Type{ModalC01d}) = true

function basis1d{T<:Number}(b::ModalC01d, ξ::T, p::Integer)

    if p == 1
      ϕ = (one(T) - ξ) / 2
    elseif p == 2
      ϕ = (one(T) + ξ) / 2
    else
      ϕ = (one(T) - ξ)*(one(T) + ξ) / 4 * jacobi(ξ, p-3, 1, 1)
    end
    return ϕ
end

"""
Legendre orthogonal basis
"""
immutable Legendre1d{T<:Number} <: BasisFun1d{T}
    nmodes::Int
    lnum::LocalNumSys1d
end
Legendre1d(n) = Legendre1d(n+1, LocalNumSys1d( [1:(n+1);], Int[]) )

basis1d{T<:Number}(b::Legendre1d, ξ::T, p::Integer) = legendre(ξ, p-1)

"""
Lagrange polynomial basis.
"""
immutable Lagrange1d{T<:Number} <: BasisFun1d
    nmodes::Int
    lnum::LocalNumSys1d
    z::Array{T,1}
end

"""
Build a Lagrange polynomial basis at Gauss-Lobatto nodes.
"""
function Lagrange1d{T<:Number}(n, ::Type{T}=Float64)
    z = Jacobi.zglj(n, 0, 0, T)
    Lagrange1d{T}(n, LocalNumSys1d([1,n], [2:n-1;]), z)
end

"""
Build a Lagrange polynomial basis at any set of distinct nodes.
"""
function Lagrange1d{T<:Number}(x::AbstractVector{T})
    n = length(x)
    z = zeros(T, n)

    for i = 1:n
        z[i] = x[i]
    end
    bidx = Int[]
    iidx = Int[]

    for i = 1:n
        if z[i]==-1 || z[i]==1
            append!(bidx, [i])
        else
            append!(iidx, [i])
        end
    end
    Lagrange1d(n, LocalNumSys1d(bidx, iidx), z)
end

isbndryint(b::Lagrange1d) = isbndryint(b.lnum)

basis1d{T<:Number}(b::Lagrange1d{T}, ξ::T, p) = Jacobi.lagrange(p, ξ, b.z)


using Distributions
using LinearAlgebra
using StatsBase

num_bins = 20

function get_hist_edges(input_dict, nbins, population)
    all_outputs = reduce(vcat, [input_dict[i] for i in keys(input_dict)])

    max_output = maximum(all_outputs)
    min_output = minimum(all_outputs)

    if (population) # percent off
        # println("Population CC bin edges!")
        bin_edges = collect(-1:102.0/nbins:101.0)
    else # FITC
        # println("Single cell CC bin edges!")
        bin_edges = collect(min_output-1:(max_output+2-min_output)/nbins:max_output+1)
    end

    h = fit(Histogram, all_outputs, bin_edges)
    return collect(h.edges[1])
end

function find_bin_index(value, bin_edges) # wheeee no bounds checking here
    if debug
        print("Bin edges: ")
        println(bin_edges)
        print("Value: ")
        println(value)
    end
    for i in 1:length(bin_edges)-1
        if value >= bin_edges[i] && value < bin_edges[i+1]
            if debug
                print("Found index at: ")
                println(i)
            end
            return i
        else
            if debug
                print("Did not find index at ")
                println(i)
            end
        end
    end
end

# each row of matrix is input, each column is output
function build_conditional_matrix(input_dict, population)
    inputs = collect(keys(input_dict))
    output_bin_edges = get_hist_edges(input_dict, num_bins, population) # lower edges, bin into 20 things
    freqs = zeros(length(inputs), length(output_bin_edges)-1)

    if debug
        println("Frequencies: ")
        println(freqs)
    end

    for i in 1:length(inputs)
        input = inputs[i]
        for j in 1:length(input_dict[input])
            output = input_dict[input][j]
            bin_idx = find_bin_index(output, output_bin_edges)

            if isnothing(bin_idx)
                println(output_bin_edges)
                println(output)
            end
            if debug
                print("Adding count to ")
                print(i)
                print(", ")
                println(bin_idx)
            end
            freqs[i, bin_idx] += 1
        end
    end

    if debug
        println("Frequencies: ")
        println(freqs)
    end

    # now we have to normalize each row
    for i in 1:size(freqs)[1]
        tot = sum(freqs[i,:])
        for j in 1:length(output_bin_edges)-1
            freqs[i, j] = freqs[i, j]/tot
        end
    end

    if debug
        println("Frequencies: ")
        println(freqs)
    end

    return freqs
end

function blahut_arimoto(x)
    numrows, numcols = size(x)

    # check for negative entries
    for i in 1:numrows
        for j in 1:numcols
            if (x[i,j]<0)
                println("ERROR: negative entry in input matrix")
                return
            end
        end
    end

    # check for zero columns
    for i in 1:numcols
        arr = x[:,i]
        if arr == zeros(1, numrows)
            println("ERROR: there is a zero column in the matrix")
            return;
        end
    end

    # check for zero rows and make sure they sum to 1
    for i in 1:numrows
        arr = x[i,:]
        if arr == zeros(1, numcols)
            println("ERROR: there is a zero row in the matrix")
            return;
        end
        x[i,:] = x[i,:]/sum(x[i,:]) # normalize rows
    end

    r = ones(numrows)/numrows
    q = zeros(numrows, numcols)
    error_tolerance = 1e-5/numrows
    r1 = zeros(numrows)

    for i in 1:numrows
        x[i,:] = x[i,:]/sum(x[i,:])
    end
    for iter in 1:10000
        for j in 1:numcols
            q[:,j] = r .* x[:,j]
            q[:,j] = q[:,j]/sum(q[:,j])
        end
        for i in 1:numrows
            r1[i] = prod(q[i,:] .^ x[i,:])
        end
        r1 = r1/sum(r1)
        r = r1

        #  if (norm(r1-r) < error_tolerance)
        #      break
        #  else
        #      r = r1
        #  end
    end

    C = 0
    for i in 1:numrows
        for j in 1:numcols
            if r[i] > 0 && q[i,j] > 0
                C = C + r[i]*x[i,j]*log(q[i,j]/r[i])
            end
        end
    end

    C = C / log(2)
    return (C, r)
end

#= Computes the channel capacity given a dict of input => outputs
=#
function compute_channel_capacity(input_dict, population=true)
    m = build_conditional_matrix(input_dict, population)
    C, r = blahut_arimoto(m)
    return C
end


using PtFEM

ProjDir = dirname(@__FILE__)

l = 1.0       # Total length [m]
N = 5         # Number of nodes
els = N - 1   # Number of finite elements (in x direction)
nod = 2       # Number of nodes per finite elements
nodof = 1     # Degrees of freedom for each node
np_types = 1  # Number of property types
EA = 1.0e5    # Strain stiffness
nip = 1       # Number of integration points

data = Dict(
  # StructuralElement(nxe, np_types, nip, FiniteElement(nod, nodof))
  :struc_el => Rod(els, np_types, nip, Line(nod, nodof)),
  :properties => [EA;],
  # Compute x_coords using length l and number of elements, els
  :x_coords => 0.0:l/els:l,
  # Define a support for node N
  # Fix the single dof (x direction displacement) of node N
  :support => [(N, [0])],
  # External forces are applied to nodes 1 to 5.
  :loaded_nodes => [
      (1, [-0.625]),
      (2, [-1.25]),
      (3, [-1.25]),
      (4, [-1.25]),
      (5, [-0.625])
    ]
);

# Display the data dictionary
data |> display
println()

# Solve the FEM model
@time fem, dis_df, fm_df = p41(data)
println()

display(fem.displacements)
println()
display(fem.actions)
println()

display(dis_df)
println()
display(fm_df)
println()

using Plots
gr(size=(400,500))

x = 0.0:l/els:l
u = convert(Array, dis_df[:, :x_translation])

p = Vector{Plots.Plot{Plots.GRBackend}}(undef, 2)
titles = ["PtFEM Ex41.1 u(x)", "PtFEM Ex41.1 N(x)"]

p[1]=plot(ylim=(0.0, 1.0), xlim=(0.0, 5.0),
  yflip=true, xflip=false, xlab="Normal force [N]",
  ylab="x [m]", title=titles[2]
)
vals = convert(Array, fm_df[:, :normal_force_2])
for i in 1:els
    plot!(p[1],
      [vals[i], vals[i]],
      [(i-1)*l/els, i*l/els], color=:blue,
      fill=true, fillalpha=0.1, leg=false
    )
    delta = abs(((i-1)*l/els) - (i*l/els)) / 20.0
    y1 = collect(((i-1)*l/els):delta:(i*l/els))
    for j in 1:length(y1)
      plot!(p[1],
        [vals[i], 0.0],
        [y1[j], y1[j]], color=:blue, alpha=0.5
      )
    end
end

p[2] = plot(u, x, xlim=(-0.00003, 0.0), yflip=true,
  xlab="Displacement [m]", ylab="x [m]",
  fillto=0.0, fillalpha=0.1, leg=false, title=titles[1])

plot(p..., layout=(1, 2))
savefig(ProjDir*"/Ex41.1.png")


const R = MLJLinearModels
const CI = get(ENV, "CI", "false") == "true"

DO_COMPARISONS = DO_COMPARISONS && !CI
DO_COMPARISONS && (using PyCall; using RCall)
SKLEARN_LM = nothing
PY_RND     = nothing
if DO_COMPARISONS
    SKLEARN_LM = pyimport("sklearn.linear_model")
    PY_RND     = pyimport("random")
    QUANTREG   = rimport("quantreg")
end

m(s, p=true) = println("\n== $s ==" * ifelse(p, "\n", ""))
mm(s) = println("\n > $s < \n")

nnz(θ) = sum(abs.(θ) .> 0)

"""Make portion s of vector 0."""
sparsify!(θ, s) = (θ .*= (rand(length(θ)) .< s))

"""Add outliers to portion s of vector."""
outlify(y, s) = (n = length(y); y .+ 20 * randn(n) .* (rand(n) .< s))

"""Generate continuous (X, y) with and without intercept."""
function generate_continuous(n, p; seed=0, sparse=1)
    Random.seed!(seed)
    X  = randn(n, p)
    X1 = R.augment_X(X, true)
    θ  = randn(p)
    θ1 = randn(p+1)
    sparse < 1 && begin
        sparsify!(θ, sparse)
        sparsify!(θ1, sparse)
    end
    y  = X*θ + 0.1 * randn(n)
    y1 = X1*θ1 + 0.1 * randn(n)
    return ((X, y, θ), (X1, y1, θ1))
end

"""Generate continuous X and binary y with and without intercept."""
function generate_binary(n, p; seed=0, sparse=1)
    Random.seed!(seed)
    X  = randn(n, p)
    X1 = R.augment_X(X, true)
    θ  = randn(p)
    θ1 = randn(p+1)
    sparse < 1 && begin
        sparsify!(θ, sparse)
        sparsify!(θ1, sparse)
    end
    y  = rand(n) .< R.σ.(X*θ)
    y  = y .* ones(Int, n) .- .!y .* ones(Int, n)
    y1 = rand(n) .< R.σ.(X1*θ1)
    y1 = y1 .* ones(Int, n) .- .!y1 .* ones(Int, n)
    return ((X, y, θ), (X1, y1, θ1))
end

"""Simple function to sample from a multinomial."""
function multi_rand(Mp)
    # Mp[i, :] sums to 1
    n, c = size(Mp)
    be   = reshape(rand(length(Mp)), n, c)
    y    = zeros(Int, n)
    @inbounds for i in eachindex(y)
        rp = 1.0
        for k in 1:c-1
            if (be[i, k] < Mp[i, k] / rp)
                y[i] = k
                break
            end
            rp -= Mp[i, k]
        end
    end
    y[y .== 0] .= c
    return y
end

"""Generate continuous X and multiclass y with and without intercept."""
function generate_multiclass(n, p, c; seed=0, sparse=1)
    Random.seed!(seed)
    X   = randn(n, p)
    X1  = R.augment_X(X, true)
    θ   = randn(p * c)
    θ1  = randn((p+1) * c)
    sparse < 1 && begin
        sparsify!(θ, sparse)
        sparsify!(θ1, sparse)
    end
    y   = zeros(Int, n)
    y1  = zeros(Int, n)

    P   = R.apply_X(X, θ, c)
    M   = exp.(P)
    Mn  = M ./ sum(M, dims=2)
    P1  = R.apply_X(X, θ1, c)
    M1  = exp.(P1)
    Mn1 = M1 ./ sum(M1, dims=2)

    y  = multi_rand(Mn)
    y1 = multi_rand(Mn1)
    return ((X, y, θ), (X1, y1, θ1))
end


save_results = false

TO.diffmethod(::RobotZoo.Cartpole) = RD.ForwardAD()
@testset "Cartpole" begin
# load expected results
res = load(joinpath(@__DIR__,"cartpole.jld2"))

# Set up problem and solver
prob, opts = Problems.Cartpole()
t = TO.get_times(prob)
U = prob.constraints[1].z_max[end] * 2 * sin.(t .* 2)
U = [[u] for u in U]
initial_controls!(prob, U)
rollout!(prob)

solver = ALTROSolver(prob, opts, save_S=true)
n,m,N = size(solver)
ilqr = Altro.get_ilqr(solver)
Altro.initialize!(ilqr)
X = states(ilqr)
U = controls(ilqr)
@test X ≈ res["X"] atol=1e-6
@test U ≈ res["U"] atol=1e-6

# State diff Jacobian
TO.state_diff_jacobian!(ilqr.G, ilqr.model, ilqr.Z)
G = ilqr.G
@test G[1] == zeros(n,n)
@test G[N] == zeros(n,n)
@test G ≈ res["G"] atol=1e-6

# Dynamics Jacobians
TO.dynamics_expansion!(TO.integration(ilqr), ilqr.D, ilqr.model, ilqr.Z)
TO.error_expansion!(ilqr.D, ilqr.model, ilqr.G)
A = [Matrix(D.A_) for D in ilqr.D]
B = [Matrix(D.B_) for D in ilqr.D]
@test A ≈ res["A"] atol=1e-6
@test B ≈ res["B"] atol=1e-6

# Unconstrained Cost Expansion
TO.cost_expansion!(ilqr.quad_obj, ilqr.obj.obj, ilqr.Z, init=true, rezero=true)
TO.error_expansion!(ilqr.E, ilqr.quad_obj, ilqr.model, ilqr.Z, ilqr.G)
hess0 = [[E.Q E.H'; E.H E.R] for E in ilqr.E]
grad0 = [Vector([E.q; E.r]) for E in ilqr.E]
@test hess0 ≈ res["hess0"] atol=1e-6
@test grad0 ≈ res["grad0"] atol=1e-6

# AL Cost Expansion
TO.cost_expansion!(ilqr.quad_obj, ilqr.obj, ilqr.Z, init=true, rezero=true)
TO.error_expansion!(ilqr.E, ilqr.quad_obj, ilqr.model, ilqr.Z, ilqr.G)
hess = [[E.Q E.H'; E.H E.R] for E in ilqr.E]
grad = [Vector([E.q; E.r]) for E in ilqr.E]
@test hess ≈ res["hess"] atol=1e-6
@test grad ≈ res["grad"] atol=1e-6

# Backward Pass
ΔV = Altro.static_backwardpass!(ilqr)
K = Matrix.(ilqr.K)
d = Vector.(ilqr.d)
S = [Matrix(S.Q) for S in ilqr.S]
s = [Vector(S.q) for S in ilqr.S]
Qzz = [[E.Q E.H'; E.H E.R] for E in ilqr.Q]
Qz = [Vector([E.q; E.r]) for E in ilqr.Q]
@test K ≈ res["K"] atol=1e-6
@test d ≈ res["d"] atol=1e-6
@test S ≈ res["S"] atol=1e-6
@test s ≈ res["s"] atol=1e-6
@test Qzz ≈ res["Qzz"] atol=1e-6
@test Qz  ≈ res["Qz"] atol=1e-6

if save_results
    @save joinpath(@__DIR__,"cartpole.jld2") X U G A B hess0 grad0 hess grad K d S s Qzz Qz
end

end


using Test
using TightBindingLattice

@testset "OrthoCube" begin

    @testset "failures" begin
        @test_throws DimensionMismatch OrthoCube([1 2 3; 4 5 6])
        @test_throws ArgumentError OrthoCube([1 1; 1 1])
    end

    @testset "orthogonal" begin
        orthocube = OrthoCube([3 0; 0 3])
        @test dimension(orthocube) == 2
        @test isequiv(orthocube, OrthoCube([3 0; 3 3]))
        @test !isequiv(orthocube, OrthoCube([3 0; 1 3]))
        @test !isequiv(orthocube, OrthoCube([6 0; 0 6]))
        @test !isequiv(OrthoCube([6 0; 0 6]), orthocube)

        for i in 0:6, j in 0:6
            @test orthocube.wrap([i, j]) == ([i÷3, j÷3], [i%3, j%3])
        end
        let r = hcat([[i,j] for i in 0:6 for j in 0:6]...),
            R = hcat([[i÷3,j÷3] for i in 0:6 for j in 0:6]...),
            ρ = hcat([[i%3,j%3] for i in 0:6 for j in 0:6]...)
            @test orthocube.wrap(r) == (R, ρ)
        end
    end

    @testset "nonorthogonal" begin
        orthocube = OrthoCube([2 1; 1 2])
        @test orthocube.inverse_shape_matrix == [2//3 -1//3; -1//3 2//3]
        @test dimension(orthocube) == 2
        @test orthocube.wrap([0,0]) == ([0,  0], [0, 0])
        @test orthocube.wrap([1,0]) == ([0, -1], [2, 2])
        @test orthocube.wrap([2,0]) == ([1, -1], [1, 1])
        @test orthocube.wrap([3,0]) == ([2, -1], [0, 0])
    end

    @testset "generator" begin
        ortho = OrthoCube([4 0; 0 3])
        generate_coordinates(ortho, [ 1 4; -1 -3])
        generate_coordinates(ortho, [ 1 0; 0 1])
        @test_throws DimensionMismatch generate_coordinates(ortho, [ 1 0 0; 0 1 0]) # not square
        @test_throws ArgumentError generate_coordinates(ortho, [ 2 0; 0 1]) # unimodular
        @test_throws ArgumentError generate_coordinates(ortho, [ 1 0; 1 1])
        # c2 = generate_coordinates(ortho, [1 0; 0 1])
    end

end


struct SklarDist{CT,TplMargins} <: Distributions.ContinuousMultivariateDistribution
    C::CT
    m::TplMargins
    function SklarDist(C,m)
        d = length(C)
        @assert length(m) == d
        @assert all(mᵢ isa Distributions.UnivariateDistribution for mᵢ in m)
        return new{typeof(C),typeof(m)}(C,m)
    end
end
Base.length(S::SklarDist{CT,TplMargins}) where {CT,TplMargins} = length(S.C)
Base.eltype(S::SklarDist{CT,TplMargins}) where {CT,TplMargins} = Base.eltype(S.C)

function Distributions.cdf(S::SklarDist{CT,TplMargins},x) where {CT,TplMargins}
    return cdf(S.C,cdf.(S.m,x))
end
function Distributions._rand!(rng::Distributions.AbstractRNG, S::SklarDist{CT,TplMargins}, x::AbstractVector{T}) where {CT,TplMargins,T}
    Random.rand!(rng,S.C,x)
     x .= Distributions.quantile.(S.m,x)
end
function Base.rand(rng::Distributions.AbstractRNG,S::SklarDist{CT,TplMargins}) where {CT,TplMargins}
     x = Random.rand(rng,S.C)
     x .= Distributions.quantile.(S.m,x)
     return x
end
function Distributions._logpdf(S::SklarDist{CT,TplMargins},u) where {CT,TplMargins}
    sum(Distributions.logpdf(S.m[i],u[i]) for i in 1:length(u)) + Distributions.logpdf(S.C,Distributions.cdf.(S.m,u))
end

function Distributions.fit(::Type{SklarDist{CT,TplMargins}},x) where {CT,TplMargins}
    # The first thing to do is to fit the marginals :
    @assert length(TplMargins.parameters) == size(x,1)
    m = Tuple(Distributions.fit(TplMargins.parameters[i],x[i,:]) for i in 1:size(x,1))
    u = pseudos(x)
    C = Distributions.fit(CT,u)
    return SklarDist(C,m)
end


using FastTransforms, LowRankApprox, Compat
using Compat.Test, Compat.LinearAlgebra

if VERSION < v"0.7-"
    const mul! = Base.A_mul_B!
    const lmul! = Base.scale!
    const rmul! = Base.scale!
else
    Ac_mul_B!(C, A, B) = mul!(C, A', B)
    At_mul_B!(C, A, B) = mul!(C, transpose(A), B)
end

import FastTransforms: Butterfly

@testset "Butterfly algorithm" begin
    kernel = (x,y) -> exp(im*x*y)

    function randfft(m,n,σ)
        x = (range(0,stop=m-1,length=m)+σ*rand(m))
        y = (range(0,stop=n-1,length=n)+σ*rand(n))*2π/n
        Complex{Float64}[kernel(x,y) for x in x, y in y]
    end

    function randnfft(m,n,σ)
        x = (range(0,stop=m-1,length=m)+σ*randn(m))
        y = (range(0,stop=n-1,length=n)+σ*randn(n))*2π/n
        Complex{Float64}[kernel(x,y) for x in x, y in y]
    end

    println("Testing the butterflied FFT")

    N = 10
    A = Vector{Matrix{Complex{Float64}}}(undef,N)
    for n in 1:N
        A[n] = randnfft(2^n,2^n,0.0)
    end

    for n in 7:N
        @time B = Butterfly(A[n], n-5)
        nb = size(A[n], 2)
        b = rand(Complex{Float64}, nb)./(1:nb)
        u = zero(b)
        @time uf = A[n]*b
        @time mul!(u, B, b)
        w = zero(b)
        @time Ac_mul_B!(w, B, u)
        lmul!(inv(2^n), w)
        println(norm(u-uf)/2^n)
        println(norm(w-b))
        println(norm(w-A[n]\u))
    end

    println("Testing the butterflied Hilbert matrix")

    N = 10
    A = Vector{Matrix{Float64}}(undef,N)
    for n in 1:N
        A[n] = Float64[1/(i+j-1) for i = 1:2^n+50,j=1:2^n+50]
    end

    for n in 7:N
        @time B = Butterfly(A[n], n-5)
        nb = size(A[n], 2)
        b = rand(Float64, nb)./(1:nb)
        u = zero(b)
        @time uf = A[n]*b
        @time mul!(u, B, b)
        w = zero(b)
        @time At_mul_B!(w, B, b)
        println(norm(u-uf)/nb)
        println(norm(w-A[n]'b))
        println(norm(u-w))
    end

    println("Testing the butterflied NUFFT")

    N = 10
    A = Vector{Matrix{Complex{Float64}}}(undef,N)
    B = Vector{Butterfly{Complex{Float64}}}(undef,N)
    for n in 7:N
        A[n] = randnfft(2^n+50,2^n+50,0.1)
        @time B[n] = Butterfly(A[n], n-6)
    end

    for n in 7:N
        nb = size(A[n], 2)
        b = rand(Complex{Float64}, nb)./(1:nb)
        uf = zero(b)
        u = zero(b)
        @time mul!(uf, A[n], b)
        @time mul!(u, B[n], b)
        println(norm(u-uf))
    end
end


struct StrainData
  outside_multiplier::Float32
  household_multiplier::Float32
end

const StrainTable = SVector{3, StrainData}

function make_strains(;base_multiplier::Real=1.0, british_multiplier::Real=1.70, delta_multiplier::Real=1.7*1.5)::StrainTable
  SA[
    StrainData(base_multiplier, base_multiplier),
    StrainData(base_multiplier * british_multiplier, base_multiplier * british_multiplier),
    StrainData(base_multiplier* delta_multiplier, base_multiplier * delta_multiplier)
  ]
end

getdata(table::StrainTable, strain::StrainKind) = table[UInt8(strain)]


# gemm on various size matrices

function matmultest(n, iter)
    a = rand(n,n)
    b = similar(a)
    for i=1:iter
        A_mul_B!(b, a, a)
    end
    b
end

problemsizes = [(2, 10^6, "tiny"), (2^4, 10^5, "small"), (2^6, 10^4, "medium"), (2^8, 10^2, "large")]
testdata = [(gemvtest, "matmul", "matrix-matrix multiplication", problemsizes)]
include("../perfgeneric.jl")


#=
Copyright (c) 2015, Roberto Garrappa
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in
      the documentation and/or other materials provided with the distribution
    * Neither the name of the Department of Mathematics - University of Bari - Italy nor the names
      of its contributors may be used to endorse or promote products derived
      from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
=#

"""
@article{Garrappa2015NumericalEO,
  title={Numerical Evaluation of Two and Three Parameter Mittag-Leffler Functions},
  author={Roberto Garrappa},
  journal={SIAM J. Numer. Anal.},
  year={2015},
  volume={53},
  pages={1350-1369}
}
"""

myeps(x) = x |> one |> float |> eps
myeps(x::Complex) =  x |> real |> myeps

"""
    mittleffderiv(α, β, z)

Compute the derivative of the Mittag-Leffler function at `z` for parameters `α, β`.
"""
function mittleffderiv(α, β, z)
    q = 1//2
    #case 1, small z
    if abs(z) <= q
        ω = α + β - 3//2
        D = α^2 - 4*α*β + 6*α + 1
        #k1
        if α>1
            k₁ = ((2 - α - β)/(α - 1)) + 1
        elseif α > 0 && α <= 1 && D <= 0
            k₁ = ((3 - α - β)/α) + 1
        else
            k₁ = maximum([((3-α-β)/α) + 1, ((1-2*ω*α + sqrt(D))/(2*(α^2)))+1])
        end
        k₀ = maximum([k₁, log(myeps(z)*(1-abs(z)))/log(abs(z))])
        k₀ = ceil(Int, k₀)
        out = zero(z)
        for k in 0:k₀
            out = out + ((k + 1)*z^k)/(gamma(α + β + α*k))
        end
    #case 2, larger z
    else
        out = (mittleff(α, β-1, z) - (β - 1)*mittleff(α, β, z))/(α*z)
    end
    return out
end

"""
    mittleffderiv(α, z)

Compute mittleffderiv(α, 1, z)
"""
mittleffderiv(α, z) = mittleffderiv(α, 1, z)

"""
    mittleff(α, β, γ, z)

Compute three-parametric mittleff(α, β, γ, z).

```tex
@article{2015,
   title={Numerical Evaluation of Two and Three Parameter Mittag-Leffler Functions},
   volume={53},
   ISSN={1095-7170},
   url={http://dx.doi.org/10.1137/140971191},
   DOI={10.1137/140971191},
   number={3},
   journal={SIAM Journal on Numerical Analysis},
   publisher={Society for Industrial & Applied Mathematics (SIAM)},
   author={Garrappa, Roberto},
   year={2015},
   month={Jan},
   pages={1350–1369}
}
```
"""
function mittleff(α, β, γ, z)
    log_epsilon = log(10^(-15))
    E = zeros(Float64)
    abs(z) < 1e-15 ? E=1/gamma(β) : E=LTInversion(1, z, α, β, γ, log_epsilon)
    return E
end

"""
    mittleff(α, β, z)

Compute two-parametric Mittag Leffler function
"""
mittleff(α, β, z) = mittleff(α, β, 1, z)

"""
    mittleff(α, z)

Compute single-parametric Mittag Leffler function
"""
mittleff(α, z) = mittleff(α, 1, 1, z)

mittleff(α, vec::AbstractArray) = map(x -> mittleff(α, 1, x), vec)
mittleff(α, β, vec::AbstractArray) = map(x -> mittleff(α, β, x), vec)
mittleff(α, β, γ, vec::AbstractArray) = map(x -> mittleff(α, β, γ, x), vec)

function LTInversion(t, lambda, alpha, beta, gama, log_epsilon)
    theta = angle(lambda)
    kmin = ceil(Int, -alpha/2 - theta/2/pi)
    kmax = floor(Int, alpha/2 - theta/2/pi)
    k_vett = collect(kmin:kmax)
    s_star = abs(lambda)^(1/alpha) * exp.(im*(theta.+2*k_vett*pi)/alpha)

    phi_s_star = (real(s_star).+ abs.(s_star))/2
    index_s_star = sortperm(phi_s_star)
    phi_s_star = sort(phi_s_star)
    s_star = s_star[index_s_star]
    index_save = phi_s_star .> 1e-15

    s_star = s_star[index_save]
    phi_s_star = phi_s_star[index_save]
    s_star = [0; s_star]
    phi_s_star = [0; phi_s_star]
    J1 = length(s_star)
    J = J1 - 1
    p = [max(0,-2*(alpha*gama-beta+1)) ; ones(J)*gama]
    q = [ones(J)*gama ; +Inf]
    phi_s_star = [phi_s_star; +Inf]
    admissible_regions = findall((phi_s_star[1:end-1] .< (log_epsilon - log(eps()))/t) .& (phi_s_star[1:end-1] .< phi_s_star[2:end]))
    JJ1 = admissible_regions[end]
    mu_vett = ones(ComplexF64, JJ1)*Inf
    N_vett = ones(ComplexF64, JJ1)*Inf
    h_vett = ones(ComplexF64, JJ1)*Inf
    find_region = 0

    while find_region==0
        for j1 = admissible_regions
            muj=0; hj=0; Nj=0
            if j1 < J1
                j1=Int64(j1)
                (muj, hj, Nj) = OptimalParam_RB(t, phi_s_star[j1], phi_s_star[j1+1], p[j1], q[j1], log_epsilon) ;
            else
                (muj, hj, Nj) = OptimalParam_RU(t,phi_s_star[j1],p[j1],log_epsilon) ;
            end
            mu_vett[j1] = muj; h_vett[j1] = hj; N_vett[j1] = Nj
        end
        if minimum(real.(N_vett[:])) > 200
            log_epsilon = log_epsilon + log(10)
        else
            find_region = 1
        end
    end

    (N, iN) = findmin(real.(N_vett)); mu = mu_vett[iN]; h = h_vett[iN];

    k = collect(-N:N)
    u = h.*k
    z = mu*(1*im*u.+1).^2
    zd = -2*mu*u .+ 2*mu*1*im
    zexp = exp.(z*t)
    F = z.^(alpha*gama-beta)./(z.^alpha .- lambda).^gama.*zd
    S = zexp.*F
    Integral = h*sum(S)/2/pi/im

    ss_star = s_star[(iN[1]+1):end]
    Residues = sum(1/alpha*(ss_star).^(1-beta).*exp.(t*ss_star))

    E = Integral + Residues
    if isreal(lambda)
        E = real(E)
    end
    return E
end

function OptimalParam_RU(t, phi_s_star_j, pj, log_epsilon)
    sq_phi_s_star_j = sqrt(phi_s_star_j)
    phi_s_star_j > 0 ? phibar_star_j = phi_s_star_j*1.01 : phibar_star_j = 0.01
    sq_phibar_star_j = sqrt(phibar_star_j)

    f_min=1
    f_max=10
    f_tar=5

    stop=0
    sq_muj=0
    A=0
    Nj=0

    while stop==0
        phi_t = phibar_star_j*t
        log_eps_phi_t = log_epsilon/phi_t
        Nj = Complex(ceil(real(phi_t/pi*(1-3*log_eps_phi_t/2+sqrt(Complex(1-2*log_eps_phi_t))))), ceil(imag(phi_t/pi*(1-3*log_eps_phi_t/2+sqrt(Complex(1-2*log_eps_phi_t))))))
        A = pi*Nj/phi_t
        sq_muj = sq_phibar_star_j*abs(4-A)/abs(7-sqrt(1+12*A))
        fbar=((sq_phibar_star_j-sq_phi_s_star_j)/sq_muj)^(-pj)
        stop=(pj<1e-14) || (f_min<fbar && fbar<f_max)
        if stop == 0
            sq_phibar_star_j = f_tar^(-1/pj)*sq_muj+sq_phi_s_star_j
            phibar_star_j = sq_phibar_star_j^2
        end
    end
    muj = sq_muj^2
    hj=(-3*A-2+2*sqrt(1+12*A))/(4-A)/Nj

    log_eps=log(eps())
    threshold = (log_epsilon-log_eps)/t
    if muj > threshold
        abs(pj)<1e-14 ? Q=0 : Q=f_tar^(-1/pj)*sqrt(muj)
        phibar_star_j = (Q+sqrt(phi_s_star_j))^2
        if phibar_star_j<threshold
            w=sqrt(log_eps/(log_eps-log_epsilon))
            u=sqrt(-phibar_star_j*t/log_eps)
            muj = threshold
            Nj=ceil(w*log_epsilon/2/pi/(u*w-1))
            hj=sqrt(log_eps/(log_eps-log_epsilon))/Nj
        else
            Nj=+Inf
            hj=0
        end
    end
    return muj, hj, Nj
end

function OptimalParam_RB(t, phi_s_star_j, phi_s_star_j1, pj, qj, log_epsilon)
    log_eps = -36.043653389117154
    fac = 1.01
    conservative_error_analysis = 0

    f_max = exp(log_epsilon - log_eps)

    sq_phi_star_j = sqrt(phi_s_star_j)
    threshold = 2*sqrt((log_epsilon - log_eps)/t)
    sq_phi_star_j1 = min(sqrt(phi_s_star_j1), threshold-sq_phi_star_j)

    if pj < 1.0e-14 && qj < 1.0e-14
        sq_phibar_star_j = sq_phi_star_j
        sq_phibar_star_j1 = sq_phi_star_j1
        adm_region = 1
    end

    if pj < 1.0e-14 && qj >= 1.0e-14
        sq_phibar_star_j = sq_phi_star_j
        sq_phi_star_j > 0 ? f_min = fac*(sq_phi_star_j/(sq_phi_star_j1-sq_phi_star_j))^qj : f_min = fac
        if f_min < f_max
            f_bar = f_min + f_min/f_max*(f_max-f_min)
            fq = f_bar^(-1/qj)
            sq_phibar_star_j1 = (2*sq_phi_star_j1-fq*sq_phi_star_j)/(2+fq)
            adm_region = 1
        else
            adm_region = 0
        end
    end

    if pj >= 1.0e-14 && qj < 1.0e-14
        sq_phibar_star_j1 = sq_phi_star_j1
        f_min = fac*(sq_phi_star_j1/(sq_phi_star_j1-sq_phi_star_j))^pj
        if f_min < f_max
            f_bar = f_min + f_min/f_max*(f_max-f_min)
            fp = f_bar^(-1/pj)
            sq_phibar_star_j = (2*sq_phi_star_j+fp*sq_phi_star_j1)/(2-fp)
            adm_region = 1
        else
            adm_region = 0
        end
    end

    if pj >= 1.0e-14 && qj >= 1.0e-14
        f_min = fac*(sq_phi_star_j+sq_phi_star_j1)/(sq_phi_star_j1-sq_phi_star_j)^max(pj,qj) ;
        if f_min < f_max
            f_min = max(f_min,1.5)
            f_bar = f_min + f_min/f_max*(f_max-f_min)
            fp = f_bar^(-1/pj)
            fq = f_bar^(-1/qj)
            conservative_error_analysis==0 ? w = -phi_s_star_j1*t/log_epsilon : w = -2*phi_s_star_j1*t/(log_epsilon-phi_s_star_j1*t)
            den = 2+w - (1+w)*fp + fq
            sq_phibar_star_j = ((2+w+fq)*sq_phi_star_j + fp*sq_phi_star_j1)/den
            sq_phibar_star_j1 = (-(1+w)*fq*sq_phi_star_j + (2+w-(1+w)*fp)*sq_phi_star_j1)/den
            adm_region = 1
        else
            adm_region = 0
        end
    end
    if adm_region==1
        log_epsilon = log_epsilon - log(f_bar)
        conservative_error_analysis==0 ? w = -sq_phibar_star_j1^2*t/log_epsilon : w = -2*sq_phibar_star_j1^2*t/(log_epsilon-sq_phibar_star_j1^2*t)
        muj = (((1+w)*sq_phibar_star_j + sq_phibar_star_j1)/(2+w))^2
        hj = -2*pi/log_epsilon*(sq_phibar_star_j1-sq_phibar_star_j)/((1+w)*sq_phibar_star_j + sq_phibar_star_j1)
        Nj = Complex(ceil(real(sqrt(Complex(1-log_epsilon/t/muj))/hj)), ceil(imag(sqrt(Complex(1-log_epsilon/t/muj))/hj)))
    else
        muj = 0; hj = 0; Nj = +Inf
    end
    return muj, hj, Nj
end

function mlds(z,al,be,k)
    max_gamma_arg = 171.624
    Jmax = floor(Int, (max_gamma_arg - be)/al)
    G = gamma.(al.*(collect(0:Jmax)).+be)
    jj = collect(k:Jmax)
    f = ones(size(jj))
    for l = 1 : k
        f = f.*(jj.-l.+1)
    end
    c = f./G[k+1:Jmax+1]
    E = zeros(size(z))
    Err_Round = zeros(size(z))
    Err_Round1 = copy(Err_Round)
    Err_Round2 = copy(Err_Round)
    for n = 1 : length(z)
        if abs(z[n]) < eps()
            E[n] = factorial(k)/G[k+1]
        else
            sum_arg = c.*z[n].^(jj.-k)
            abs_sum_arg = abs.(sum_arg)
            i_abs_sum_arg = abs_sum_arg .> eps()/2
            if any(i_abs_sum_arg) == 0
                i_abs_sum_arg[1] = 1
            end
            abs_sum_arg = abs_sum_arg[i_abs_sum_arg]
            sum_arg = sum_arg[i_abs_sum_arg]
            i_abs_sum_arg = sortperm(abs_sum_arg)
            abs_sum_arg = sort(abs_sum_arg)
            sum_arg = sum_arg[i_abs_sum_arg]
            if length(sum_arg) == 1
                E[n] = sum_arg
            else
                S = cumsum(sum_arg)
                S = S[2:end]
                E[n] = S[end]
            end
            J = length(sum_arg) - 1
            JJ = [J; collect(J:-1:1)]
            Err_Round1[n] = sum(JJ.*abs_sum_arg)*eps()
            length(sum_arg)==1 ? Err_Round2[n] = Err_Round1[n] : Err_Round2[n] = sum(abs.(S))*eps()
            Err_Round[n] = exp((log(Err_Round1[n])+log(Err_Round2[n]))/2)
        end
    end
    return E, Err_Round
end

function mldr(t, s, α, β, k)
    omega = zeros(k+2)
    omega[1] = 1
    jj = collect(0:k)
    p = (α.-jj)
    pr = 1
    for j = 1:k+1
        pr = pr*p[j]
        omega[j+1] = pr/factorial(j)
    end
    Hk = zeros(k+1)
    Hk[1] = 1
    for j = 1:k
        ll = collect(1:j)
        Hk[j+1] = -1 ./α.*sum(omega[ll.+2].*(k.*ll./j.+1).*Hk[j.-ll.+1])
    end
    ck = zeros(k+1)
    for j = 0:k
        temp = 0
        for l = 0 : k-j
            if l == 0
                p = 1
            else
                ll = 0:l-1
                p = prod(α.-β.-ll)
            end
            temp = temp .+ p*Hk[k-j-l+1]./factorial(l)
        end
        ck[j+1] = temp/factorial(j)
    end
    tep = Polynomial(ck)
    result = tep.(s)
    R = 1 ./α^(k+1)*exp.(t.*s).*(s//1)^(1-α.*k.-β).*result
    return R
end
#=
"""
```tex
@article{Garrappa2018ComputingTM,
  title={Computing the Matrix Mittag-Leffler Function with Applications to Fractional Calculus},
  author={Roberto Garrappa and Marina Popolizio},
  journal={Journal of Scientific Computing},
  year={2018},
  volume={77},
  pages={129-153}
}
```
"""
function mldlt(z,alpha,beta,k)
    log_epsilon = log(10^(-15)) ;
    E = zeros(size(z)) ;
    for j = 1 : length(z)
        if abs(z[j] < 1.0e-14)
            E[j] = factorial(k)/gamma(alpha*k+beta)
        else
            E[j] = LTIversion(1,z[j],alpha,beta,k,log_epsilon) ;
        end
        if isreal(z[j])
            E[j] = real(E[j])
        end
    end
    return E
end

function mld(z, alpha, beta, k)
    Ek = zeros(size(z))
    tau = 1.0e-14 ;
    max_gamma_arg = 171.624 ;
    Jmax = floor((max_gamma_arg - beta)/alpha) ;
    z_abs_max = (tau*gamma(alpha*Jmax+beta)/prod(Jmax.-collect(0:k-1))).^(1/(Jmax-k)) ;
    i_z_se = abs.(z).<z_abs_max ;
    z_se = z[i_z_se] ;
    E_se = zeros(size(z)) ; Err_Round = zeros(size(z)) ; e = ones(size(z)) ;
    (E_se[i_z_se], Err_Round[i_z_se]) = mlds(z_se,alpha,beta,k) ;
    i_z_se_accept = (Err_Round./(e.+abs.(E_se)) .< tau) .& ~(E_se==0)  ;
    Ek[i_z_se_accept] = E_se[i_z_se_accept]
    i_ze_lt = .!i_z_se_accept ;
    if k <= 3
        p = 0
    elseif k <= 7
        p = 1
    else
        p = 2
    end
    c = zeros(k-p+1, k-p+1) ;
    c[1, 1] = 1 ;
    for kk = 1 : k-p
        c[kk+1,1] = (1-alpha*(kk-1)-beta)*c[kk,1]
        for j = 1 : kk-1
            c[kk+1,j+1] = c[kk,j] + (j+1-alpha*(kk-1)-beta)*c[kk,j+1]
        end
        c[kk+1,kk+1] = 1 ;
    end
    for j = 0 : k-p
        if abs(c[k-p+1,j+1]) > 1.0e-14
            Ek[i_ze_lt] = Ek[i_ze_lt] .+ c[k-p+1,j+1]*mldlt(z[i_ze_lt], alpha, (k-p)*alpha+beta-j, p) ;
        end
    end
    Ek[i_ze_lt] = Ek[i_ze_lt]./alpha^(k-p) ;
    return Ek
end

function mittleff(A, alpha, beta)
    fun(z, k) = mld(z, alpha, beta, k)
    E=funm(A)
end
#FIXME: Didn't see any matrix function support in Julia😟.
# Yingbo has an implementation: https://github.com/YingboMa/Funm.jl
=#


"""
This function defines how a dimension is split.
It is used everywhere dimensions are split, therefore it (and the cutoff) define the recursion pattern.

The current behavior is to try and get dimensions into primarily blocks of 8.
"""
function split_dim(N)
    N <= cutoff && return N, 0
    number_cutoff_blocks, remainder = divrem(N,cutoff)

    half_of_cutoff_blocks, odd = divrem(number_cutoff_blocks, 2)

    8(half_of_cutoff_blocks+odd), 8half_of_cutoff_blocks + remainder
end

function reasonably_square(mindim, maxdim)
    # mindim, maxdim = minmax(M, N)
    round(Int, √2 * mindim, RoundUp) >= maxdim
end

# function blocks_per_dim(N::Integer)
#     blocks = 1
#     block_size = N
#     cuts = 0
#     # bit twiddling is 50% faster
#     while cutoff < block_size
#         cuts += 1
#         blocks <<= 1 # *=2
#         block_size >>= 1 #  ÷=2
#     end
#     blocks + min(max(0, N - cutoff << cuts), blocks) #N - cutoff * 2^cuts
# end
# @generated blocks_per_dim(::Val{N}) where N = blocks_per_dim(N) # Is this needed?

# ### Implemented so that the leading blocks are larger. This is because we tend to do just slightly less with them.
# function block_position(N::Integer, i::Integer)
#     num_blocks = blocks_per_dim(N)
#     1 + num_blocks - cld(num_blocks*(N+1-i),N)
# end
# function block_position(N::Integer, i::Integer, j::Integer) #Not faster on 0.6?
#     num_blocks = blocks_per_dim(N)
#     1+num_blocks - cld(num_blocks*(N+1-i),N), 1+num_blocks - cld(num_blocks*(N+1-j),N)
# end

# function block_position(N, i)
#     position = 1
#     while N > cutoff
#         N, r = divrem(N, 2)
#         Npr = N + r
# #        @show i, N, Npr
#         if i > Npr
#             position += blocks_per_dim(Npr)#round(Int, log2(Npr/cutoff), RoundUp)
#             i -= Npr
#         else
#             N = Npr
#         end
#     end
#     position
# end

function dense_sub2ind(dims::NTuple{2}, i, j, offset = 0)
    M, N = dims
    minind, maxind = minmax(M, N)
    maxind <= cutoff && return (j-1)*M + i + offset
    if reasonably_square(minind, maxind) && minind > cutoff #divide both
        M1, M2 = split_dim(M)
        N1, N2 = split_dim(N)
        if i <= M1 && j <= N1 # upper left block
            return dense_sub2ind( (M1, N1), i, j, offset )
        elseif j <= N1 # lower left block
            return dense_sub2ind( (M2, N1), i - M1, j, offset + M1 * N1 )
        elseif i <= M1 # upper right block
            return dense_sub2ind( (M1, N2), i, j - N1, offset + M * N1 )
        else # lower right block
            return dense_sub2ind( (M2, N2), i - M1, j - N1, offset + M * N1 + M1 * N2 )
        end
    elseif M > N #divide M
        M1, M2 = split_dim(M)
        if i <= M1
            return dense_sub2ind( (M1, N), i, j, offset)
        else
            return dense_sub2ind( (M2, N), i - M1, j, offset + M1 * N)
        end
    else #divide N
        N1, N2 = split_dim(N)
        if j <= N1
            return dense_sub2ind( (M, N1), i, j, offset)
        else
            return dense_sub2ind( (M, N2), i, j - N1, offset + M * N1)
        end
    end
end

"""
Generates a recursive quote.
The function keeps track of
    M, number of rows
    N, number of columns
    linear_offset, linear offset into the data of original matrix
    row_offset, Cartesian offset of rows
    column_offset, Cartesian offset of columns
Through the recursion. When the recurion terminates, it walks the expression and replaces all instances of those symbols with their current value.
"""
function dense_recursion_quote(q, M, N, linear_offset = 0, row_offset = 0, column_offset = 0)
    minind, maxind = minmax(M, N)

    # The recursion terminates, and we substitute expressions for their value.
    if maxind <= cutoff
        return MacroTools.postwalk(x -> begin
            if isa(x, Symbol)
                if x == :M
                    return M
                elseif x == :N
                    return N
                elseif x == :linear_offset
                    return linear_offset
                elseif x == :row_offset
                    return row_offset
                elseif x == :column_offset
                    return column_offset
                else
                    return x
                end
            else
                return x
            end
        end, q)
    end

    if reasonably_square(minind, maxind) && minind > cutoff # we divide the matrix into four blocks.
        M1, M2 = split_dim(M)
        N1, N2 = split_dim(N)
        return :(if i <= $(M1 + row_offset) && j <= $(N1 + column_offset) # enter block 1
                $(dense_recursion_quote(q, M1, N1, linear_offset, row_offset, column_offset))
            elseif j <= $(N1 + column_offset) # enter block 2
                $(dense_recursion_quote(q, M2, N1, linear_offset + M1 * N1, row_offset + M1, column_offset))
            elseif i <= $(M1 + row_offset) # enter block 3
                $(dense_recursion_quote(q, M1, N2, linear_offset + M * N1, row_offset, column_offset + N1))
            else # enter block 4
                $(dense_recursion_quote(q, M2, N2, linear_offset + M * N1 + M1 * N2, row_offset + M1, column_offset + N1))
            end)
    elseif M > N # we are splitting the matrix into two blocks stacked on top of one another.
        M1, M2 = split_dim(M)
        return :(if i <= $(M1 + row_offset)
                $(dense_recursion_quote(q, M1, N, linear_offset, row_offset, column_offset))
            else
                $(dense_recursion_quote(q, M2, N, linear_offset + M1 * N, row_offset + M1, column_offset))
            end)
    else # we are splitting the matrix into two blocks, side by side.
        N1, N2 = split_dim(N)
        return :(if j <= $(N1 + column_offset)
                $(dense_recursion_quote(q, M, N1, linear_offset, row_offset, column_offset))
            else
                $(dense_recursion_quote(q, M, N2, linear_offset + M * N1, row_offset, column_offset + N1))
            end)
    end
end

function dense_sub2ind_quote(M, N)
    dense_recursion_quote(
        :( (j-(column_offset+1) ) * M + i + (linear_offset-row_offset) ),
        M, N)
end

"""
Cartesian indexing isn't recomended, but it is convenient for printing.
The approach for sub2ind here incurs branches, which would disable SIMD.
Therefore, Cartesian indexing is strongly discouraged for hot loops.
Still, it is reasonably fast -- only a handful of ns.
"""
@generated dense_sub2ind(::Val{M}, ::Val{N}, i, j) where {N,M} = dense_sub2ind_quote(M, N)
# @generated function dense_sub2ind(::Val{M}, ::Val{N}, i, j) where {N,M}
#     dense_sub2ind_quote(M, N)
# end

function triangle_sub2ind(N, i, j, offset = 0)
    N <= cutoff && return small_triangle(j) + i + offset
    N1, N2 = split_dim(N)
    if max(i,j) <= N1 # enter block 1 # triangle / symmetric block
        return triangle_sub2ind(N1, i, j, offset)
    elseif i <= N1 # enter block 3 # square block
        return dense_sub2ind((N1, N2), i, j - N1, offset + big_triangle(N1))
    elseif j <= N1 # enter block 2 #symmetric or 0 block # Not supported; check i, j before calling.
        throw("Indexing into this half not supported.")
    else # enter block 4 # triangle / symmetric block
        return triangle_sub2ind(N2, i - N1, j - N1, offset + big_triangle(N1) + N1*N2)
    end
end

"""
The triangular (or Symmetric) matrix is assumed to be square, with only upper triangular elements represented.
Lower triangular matrices are simply a view of an upper triangular matrix that pretends the storage is row-major.

This function takes in two quotes, to allow for different behavior in a triangular / diagonal block, and the offdiagonal blocks.

When the recursion terminates, the following symbols are replaced with their current value:
    M -- only used for the off diagonal blocks. See dense_recursion_quote for details.
    N -- the current size of the matrix block.
    linear_offset -- linear offset within the original data matrix.
    row_offset -- how many rows the block is offset
    column_offset -- how many columns the block has been offset
"""
function triangle_recursion_quote(triangle_q, rectangle_q, N, linear_offset = 0, row_offset = 0, column_offset = 0)
    if N <= cutoff
        return MacroTools.postwalk(x -> begin
            if isa(x, Symbol)
                if x == :N
                    return N
                elseif x == :linear_offset
                    return linear_offset
                elseif x == :row_offset
                    return row_offset
                elseif x == :column_offset
                    return column_offset
                else
                    return x
                end
            else
                return x
            end
        end, triangle_q)
    end

    N1, N2 = split_dim(N)
    :(if max( i - $row_offset, j - $column_offset) <= $N1 # enter block 1 # triangle / symmetric block
            $(triangle_recursion_quote(triangle_q, rectangle_q, N1, linear_offset, row_offset, column_offset))
        elseif i <= $(N1 + row_offset) # enter block 3 # square block
            $(dense_recursion_quote(rectangle_q, N1, N2, linear_offset + big_triangle(N1), row_offset, column_offset + N1))
        else # enter block 4 # triangle / symmetric block
            $(triangle_recursion_quote(triangle_q, rectangle_q, N2,
                linear_offset + big_triangle(N1) + N1*N2, row_offset + N1, column_offset + N1))
        end)

end

function triangle_sub2ind_quote(N)
    triangle_recursion_quote(
        :( small_triangle( j-column_offset ) + i + (linear_offset-row_offset) ),
        :( (j-(column_offset+1) ) * M + i + (linear_offset-row_offset) ),
        N
    )
end

@generated function triangle_sub2ind(::Val{N}, i, j) where N
    quote
        Base.@_inline_meta
        @boundscheck i > j && throw("Indexing into this half not supported.")
        $(triangle_sub2ind_quote(N))
    end
end

function block_dim(M,N)
    mindim, maxdim = minmax(M,N)
    # The "reasonably square" may need improvement
    # For example, a 91 x 64 breaks down into
    # 46 x 64 and 45 x 64; the 45 x 64 also fails
    # the criterion, and breaks into two 45 x 32 blocks
    # How bad is this wonky pattern?
    # Should I round in the comparison?
    if reasonably_square(mindim, maxdim) && mindim > cutoff
        M1, M2 = split_dim(M)
        N1, N2 = split_dim(N)
        out = [ (M1,N1) (M1, N2); (M2,N1) (M2,N2) ]
    elseif maxdim <= cutoff
        out = Matrix{Tuple{Int,Int}}(undef,1,1)
        out[1] = (M,N)
    elseif M > N
        out = Matrix{Tuple{Int,Int}}(undef,2,1)
        M1, M2 = split_dim(M)
        out[1] = (M1, N)
        out[2] = (M2, N)
    else
        out = Matrix{Tuple{Int,Int}}(undef,1,2)
        N1, N2 = split_dim(N)
        out[1] = (M, N1)
        out[2] = (M, N2)
    end
    out
end

# function split_dim!(n_dim::Vector{Int}, i)
#     N1, N2 = split_dim(n_dim[i])
#     n_dim[i] = N2
#     insert!(n_dim, i, N1)
# end
# function split_row(dims::Matrix{Tuple{Int}}, i)
#     N1, N2 = split_dim(n_dim[i])
#     n_dim[i] = N2
#     insert!(n_dim, i, N1)
# end
function split_col(dims::Matrix{Tuple{Int,Int}})
    num_rowblocks, num_colblocks = size(dims)
    @assert num_colblocks == 1
    vcat([block_dim(dims[i]...) for i = 1:length(dims)]...)
end
function split_row(dims::Matrix{Tuple{Int,Int}})
    num_rowblocks, num_colblocks = size(dims)
    @assert num_rowblocks == 1
    hcat([block_dim(dims[i]...) for i = 1:length(dims)]...)
end

# function block_structure(N, offset = 0)
#     bounds = Int[]
#     block_structure!(bounds, N)
#     # sort!(bounds)
# end
# function block_structure!(bounds, N, offset = 0)
#     if N > cutoff
#         N, r = divrem(N, 2)
#         Npr = N + r
#         push!(bounds, Npr + offset)
#         block_structure!(bounds, Npr, offset)
#         block_structure!(bounds, N, offset + Npr)
#         return bounds
#     else
#         return bounds
#     end
# end

# function identify_block_placement_quote(N)
#     bounds = block_structure(N)
#     num_bounds = length(bounds)
#     num_blocks = num_bounds + 1
#     q = quote
#         b_0 = 0
#         $(Symbol(:b_, num_blocks)) = $N
#     end
#     for k ∈ 1:num_bounds
#         push!(q.args, :( @inbounds $(Symbol(:b_, k)) = $(bounds[k])) )
#     end
#     push!(q.args,
#         :( Base.Cartesian.@nif $num_blocks d -> (i <= b_d)  d -> ( row_block_id = d; row_lbound = b_{d-1}; row_ubound = b_d ) d -> ( row_block_id = d; row_lbound = b_{d-1}; row_ubound = b_d ) )
#     )
#     push!(q.args,
#         :( Base.Cartesian.@nif $num_blocks d -> (j <= b_d)  d -> ( col_block_id = d; col_lbound = b_{d-1}; col_ubound = b_d ) d -> ( col_block_id = d; col_lbound = b_{d-1}; col_ubound = b_d ) )
#     )
#     push!(q.args, :(block_i = i - row_lbound))
#     push!(q.args, :(block_j = j - col_lbound))
#     q
# end
# function identify_block_placement_quote(M, N)
#     rowbounds = block_structure(M)
#     colbounds = block_structure(N)
#     num_rowbounds = length(rowbounds)
#     num_colbounds = length(colbounds)
#     num_rowblocks = num_rowbounds + 1
#     num_colblocks = num_colbounds + 1
#     q = quote
#         r_0 = 0
#         c_0 = 0
#         $(Symbol(:r_, num_rowblocks)) = $M
#         $(Symbol(:c_, num_colblocks)) = $N
#     end
#     for k ∈ 1:num_rowbounds
#         push!(q.args, :( @inbounds $(Symbol(:r_, k)) = $(rowbounds[k])) )
#     end
#     for k ∈ 1:num_colbounds
#         push!(q.args, :( @inbounds $(Symbol(:c_, k)) = $(colbounds[k])) )
#     end
#     push!(q.args,
#         :( Base.Cartesian.@nif $num_rowblocks d -> (i <= r_d)  d -> ( row_lbound = r_{d-1}; row_ubound = r_d ) d -> (row_lbound = r_{d-1}; row_ubound = r_d) )
#     )
#     push!(q.args,
#         :( Base.Cartesian.@nif $num_colblocks d -> (j <= c_d)  d -> ( col_lbound = c_{d-1}; col_ubound = c_d ) d -> (col_lbound = c_{d-1}; col_ubound = c_d) )
#     )
#     push!(q.args, :(block_i = i - row_lbound))
#     push!(q.args, :(block_j = j - col_lbound))
#     q
# end

# @generated function triangle_sub2ind(::Val{N}, i, j) where N
#     N <= cutoff && return :(small_triangle(j)+i)
#     q = identify_block_placement_quote(N)
#     push!(q.args, :(ind = big_triangle(col_lbound) + (col_ubound - col_lbound) * row_lbound))
#     push!(q.args, :(if row_block_id == col_block_id # then we're in a symmetric portion
#         ind += small_triangle(block_j) + block_i
#     else #then we are in square block
#         ind += (row_ubound-row_lbound) * (block_j-1) + block_i
#     end))
#     push!(q.args, :ind)
#     q
# end

# @generated dense_sub2ind(::Val{M}, ::Val{N}, i, j) where {M,N}
#     max(M,N) <= cutoff && return :(j*$(M-1)+i)
#     q = identify_block_placement_quote(M,N)
#     push!(q.args :( col_lbound*$N + row_lbound*(col_ubound-col_lbound) + (row_ubound - row_lbound)*(block_j-1) + block_i ))
#     q
# end

# @generated function triangle_sub2ind(::Val{N}, i, j) where N
#     N <= cutoff && return :(small_triangle(j)+i)
#     num_blocks = blocks_per_dim(N)
#     quote
#         row_block_id = $(1+num_blocks) - cld($num_blocks*($(N+1)-i),$N)
#         col_block_id = $(1+num_blocks) - cld($num_blocks*($(N+1)-j),$N)

#         row_block_lbound = cld((row_block_id-1)*$N,$num_blocks)
#         # row_block_ubound = cld( row_block_id   *$N,$num_blocks)
#         col_block_lbound = cld((col_block_id-1)*$N,$num_blocks)
#         col_block_ubound = cld( col_block_id   *$N,$num_blocks)

#         ind = big_triangle(col_block_lbound) + (col_block_ubound - col_block_lbound) * row_block_lbound

#         block_i = i - row_block_lbound
#         block_j = j - col_block_lbound

#         if row_block_id == col_block_id # then we're in a symmetric portion
#             ind += small_triangle(block_j) + block_i
#         else #then we are in square block
#             nrows_in_block = cld( row_block_id   *$N,$num_blocks)  - row_block_lbound
#             ind += nrows_in_block * (block_j-1) + block_i
#         end
#         ind
#     end
# end


# using Pkg
# Pkg.activate(@__DIR__)
# Pkg.instantiate()

using CompScienceMeshes, BEAST
using LinearAlgebra

# Γ = meshcuboid(1.0, 1.0, 1.0, 0.25)
Γ = meshsphere(1.0, 0.3)
X = raviartthomas(Γ)
Y = buffachristiansen(Γ, ibscaled=true)

sol = 1.0
T = TDMaxwell3D.singlelayer(speedoflight=sol, numdiffs=1)
K = TDMaxwell3D.doublelayer(speedoflight=sol, numdiffs=1)

Δt, Nt = 0.6, 200
δ = timebasisdelta(Δt, Nt)
T0 = timebasiscxd0(Δt, Nt)
T1 = timebasisshiftedlagrange(Δt,Nt,1)
T2 = timebasisshiftedlagrange(Δt, Nt, 2)
T3 = timebasisshiftedlagrange(Δt, Nt, 3)

duration = 20 * Δt
delay = 1.5 * duration
amplitude = 1.0
gaussian = creategaussian(duration, delay, amplitude)

direction, polarisation = ẑ, x̂
E = BEAST.planewave(polarisation, direction, derive(gaussian), 1.0)
H = direction × E

@hilbertspace j m
@hilbertspace k l

pmchwt = @discretise(
    2.0T[k,j] + 2.0K[k,m] -
    2.0K[l,j] + 2.0T[l,m] == H[k] + E[l],
    j∈X⊗T2, m∈Y⊗T2, k∈X⊗δ, l∈Y⊗δ)

u = solve(pmchwt)

# nX = numfunctions(X)
# uj = u[1:nX]
# um = u[nX+1:end]

# Θ, Φ = range(0.0,stop=2π,length=100), 0.0
# ffpoints = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for θ in Θ for ϕ in Φ]
#
# # Don't forgt the far field comprises two contributions
# κ, η = 1.0, 1.0
# ffm = potential(MWFarField3D(κ*im), ffpoints, um, X)
# ffj = potential(MWFarField3D(κ*im), ffpoints, uj, X)
# ff = η*im*κ*ffj + im*κ*cross.(ffpoints, ffm)
#
# using Plots
# plot(xlabel="theta")
# plot!(Θ,norm.(ffm),label="far field")
#
# import PlotlyJS
# using LinearAlgebra
# fcrj, _ = facecurrents(uj,X)
# PlotlyJS.plot(patch(Γ, norm.(fcrj)))


using OptimizationAlgorithms
const Optimization = OptimizationAlgorithms

# optimization of kernel hyperparameters of Gaussian Process
# could add optimization w.r.t. leave-one-out loss
# TODO: optimization of noise variance
# IDEA: pass inference method (Cholesky, Iterative, Stochastic, ...)
optimize(k, θ, x, y, σ²::Real) = optimize(k, θ, x, 𝒩(y, σ²))
function optimize(k, θ::AbstractVector, x::AbstractVector,
                  y::Union{AbstractVector, Multivariate})
    optimize!(k, copy(θ), x, y)
end
optimize!(k, θ, x, y, σ²::Real) = optimize!(k, θ, x, 𝒩(y, σ²))
function optimize!(k, θ::AbstractVector, x::AbstractVector,
                   y::Union{AbstractVector, Multivariate})
    val(θ) = nlml(k(θ), x, y) # make this "safe"?
    function valdir(θ)
        v, g = value_gradient(nlml, k, θ, x, y) # need to add tol if y is Multivariate
        g .*= -1
        v, g
    end
    # optimization based on L-BFGS with m stages
    m = 5
    D = Optimization.CustomDirection(val, valdir, θ)
    D = Optimization.LBFGS(D, θ, m, check = false)
    # scaling proposed Dection to unit length before doing line search,
    # avoids jumping to sub-optimal and very large parameter vectors
    D = Optimization.UnitDirection(D)
    # to accept a set, do line search for step size to ensure decrease
    D = Optimization.DecreasingStep(D, θ)
    # finally, find fixed point of the iterations
    _, t = Optimization.fixedpoint!(D, θ)
    return θ
end

# here, k(θ) is assumed to be a kernel function
# computes both value and gradient
function value_gradient(::typeof(nlml), k, θ::AbstractVector, x::AbstractVector,
                        y::Union{AbstractVector, Multivariate})
    val, push = pushforward(nlml, k(θ), x, y)

    ∇ = (x, y) -> reshape(ForwardDiff.gradient(z->k(z)(x, y), θ), :, 1)
    G = gramian(∇, x) # lazily evaluated matrix of kernel gradients
    G_matofvec = G.A
    G = matvec2vecmat(G_matofvec) # casts G to a vector of dense matrices. IDEA: make lazy!

    grad = similar(θ)
    @threads for i in eachindex(grad) # parallelizing here is efficient
        grad[i] = push(G[i])
    end
    return val, grad
end

function gradient(::typeof(nlml), k, θ::AbstractVector,
                    x::AbstractVector, y::Union{AbstractVector, Multivariate})
    value_gradient(nlml, k, θ, x, y)[2]
end
using CovarianceFunctions: Gramian
matvec2vecmat(A::Gramian{<:AbstractVecOrMat}) = matvec2vecmat(Matrix(A))
# convert matrix of vectors to vector of matrices
function matvec2vecmat(A::AbstractMatrix{<:AbstractVecOrMat})
    k = length(A[1])
    all(==(k), (length(Ai) for Ai in A)) || throw(DimensionMismatch("component vectors do not have the same size"))
    n, m = size(A)
    B = Vector{Matrix{eltype(A[1])}}(undef, k)
    for i in eachindex(B) # this is still inefficient, because it calculates the gradient multiple times
        B[i] = [a[i] for a in A]
    end
    return B
end

########################### optimizing noise variance ##########################
function value_gradient(::typeof(nlml), log_σ²::Real, K::Factorization, y::AbstractVector)
    n = length(y)
    Σ = exp(log_σ²)*I(n)
    W = Woodbury(Σ, K)
    val, push = pushforward(nlml, W, y)
    grad = push(Σ)
    return val, grad
end

function optimize(σ²::Real, k, x::AbstractVector, y::AbstractVector; tol::Real = tol)
    # pre-factorize kernel matrix
    n = length(x)
    K = gramian(k, x)
    K = factorize(K, tol = tol)
    function val(log_σ²::Real)
        Σ = exp(log_σ²)*I(n)
        nlml(Woodbury(Σ, K), y)
    end
    val(θ::AbstractVector) = val(θ[1])
    function valdir(log_σ²::Real)
        v, g = value_gradient(nlml, log_σ², K, y)
        g = -g
        v, g
    end
    function valdir(θ::AbstractVector)
        v, g = valdir(θ[1])
        v, [g]
    end
    θ = [log(σ²)]
    D = Optimization.CustomDirection(val, valdir, θ)
    D = Optimization.LBFGS(D, θ, 2, check = false) # does this reduce to secant?
    D = Optimization.UnitDirection(D)
    D = Optimization.DecreasingStep(D, θ)
    Optimization.fixedpoint!(D, θ)
    return σ² = exp(θ[1])
end


using SpreadsheetFunctions

sheet_1_a2 = 1
sheet_1_a3 = 2
sheet_1_a4 = 3
sheet_1_a2a4 = [sheet_1_a2; sheet_1_a3; sheet_1_a4]
sheet_1_a5 = SUM(sheet_1_a2a4)
sheet_1_b2 = 4
sheet_1_b3 = 5
sheet_1_b4 = 6
sheet_1_b2b4 = [sheet_1_b2; sheet_1_b3; sheet_1_b4]
sheet_1_b5 = SUM(sheet_1_b2b4)
sheet_1_c2 = 7
sheet_1_c3 = 8
sheet_1_c4 = 9
sheet_1_c2c4 = [sheet_1_c2; sheet_1_c3; sheet_1_c4]
sheet_1_c5 = SUM(sheet_1_c2c4)
sheet_1_a3a5 = [sheet_1_a3; sheet_1_a4; sheet_1_a5]
sheet_1_a6 = SUM(sheet_1_a3a5)
sheet_1_b11 = 1
sheet_1_c11 = 2
sheet_1_d11 = 3
sheet_1_b11d11 = [sheet_1_b11 sheet_1_c11 sheet_1_d11]
sheet_1_e11 = SUM(sheet_1_b11d11)
sheet_1_b12 = 4
sheet_1_c12 = 5
sheet_1_d12 = 6
sheet_1_b12d12 = [sheet_1_b12 sheet_1_c12 sheet_1_d12]
sheet_1_e12 = SUM(sheet_1_b12d12)
sheet_1_b13 = 7
sheet_1_c13 = 8
sheet_1_d13 = 9
sheet_1_b13d13 = [sheet_1_b13 sheet_1_c13 sheet_1_d13]
sheet_1_e13 = SUM(sheet_1_b13d13)
sheet_1_b11b13 = [sheet_1_b11; sheet_1_b12; sheet_1_b13]
sheet_1_b14 = SUM(sheet_1_b11b13)
sheet_1_c11c13 = [sheet_1_c11; sheet_1_c12; sheet_1_c13]
sheet_1_c14 = SUM(sheet_1_c11c13)
sheet_1_d11d13 = [sheet_1_d11; sheet_1_d12; sheet_1_d13]
sheet_1_d14 = SUM(sheet_1_d11d13)
sheet_1_e11e13 = [sheet_1_e11; sheet_1_e12; sheet_1_e13]
sheet_1_e14 = SUM(sheet_1_e11e13)
sheet_1_b14d14 = [sheet_1_b14 sheet_1_c14 sheet_1_d14]
sheet_1_e15 = SUM(sheet_1_b14d14)
sheet_1_b1 = "B"
sheet_1_b6 = nothing
sheet_1_b7 = nothing
sheet_1_b8 = nothing
sheet_1_b9 = nothing
sheet_1_b10 = "Widget"
sheet_1_b15 = nothing
sheet_1_b16 = nothing
sheet_1_bb = [sheet_1_b1; sheet_1_b2; sheet_1_b3; sheet_1_b4; sheet_1_b5; sheet_1_b6; sheet_1_b7; sheet_1_b8; sheet_1_b9; sheet_1_b10; sheet_1_b11; sheet_1_b12; sheet_1_b13; sheet_1_b14; sheet_1_b15; sheet_1_b16]
sheet_1_a16 = SUM(sheet_1_bb)+5
sheet_1_a1 = "A"
sheet_2_d3 = sheet_1_a1
sheet_2_e3 = sheet_1_b1
sheet_2_a2 = 1
sheet_2_a3 = 2
sheet_2_a4 = 6
sheet_2_a2a4 = [sheet_2_a2; sheet_2_a3; sheet_2_a4]
sheet_2_a5 = SUM(sheet_2_a2a4)
sheet_2_b2 = 4
sheet_2_b3 = 5
sheet_2_b4 = 6
sheet_2_b2b4 = [sheet_2_b2; sheet_2_b3; sheet_2_b4]
sheet_2_b5 = SUM(sheet_2_b2b4)
sheet_2_c2 = 7
sheet_2_c3 = 8
sheet_2_c4 = 9
sheet_2_c2c4 = [sheet_2_c2; sheet_2_c3; sheet_2_c4]
sheet_2_c5 = PRODUCT(sheet_2_c2c4)
sheet_2_a1 = "A"
sheet_2_b1 = "B"
sheet_2_c1 = "C"
sheet_2_d1 = 1
sheet_2_e1 = nothing
sheet_2_11 = [sheet_2_a1 sheet_2_b1 sheet_2_c1 sheet_2_d1 sheet_2_e1]
sheet_2_c6 = CONCATENATE(sheet_2_11)
sheet_2_d5 = nothing
sheet_2_e5 = 3
sheet_2_d6 = nothing
sheet_2_e6 = nothing
sheet_2_d7 = nothing
sheet_2_e7 = nothing
sheet_2_d5f7 = [sheet_2_d5 sheet_2_e5; sheet_2_d6 sheet_2_e6; sheet_2_d7 sheet_2_e7]
sheet_2_d8 = PRODUCT(sheet_2_d5f7)