include("./Measurements.jl")

abstract type AbstractCovarianceFunction{Tv} end

# automatically implements a mutating batched version of a given covariance function
function (cov::AbstractCovarianceFunction{Tv})(out::AbstractMatrix{Tv}, x_vec::AbstractVector{<:AbstractMeasurement}, y_vec::AbstractVector{<:AbstractMeasurement}) where Tv
    for cartesian in CartesianIndices(out)
        out[cartesian] = cov(x_vec[cartesian[1]], y_vec[cartesian[2]])
    end
end

# automatically implements a mutating batched version of a given covariance function, using symmetry
function (cov::AbstractCovarianceFunction{Tv})(out::AbstractMatrix{Tv}, x_vec::AbstractVector{<:AbstractMeasurement}) where Tv
    for cartesian in CartesianIndices(out)
        if cartesian[1] >= cartesian[2]
            out[cartesian] = cov(x_vec[cartesian[1]], x_vec[cartesian[2]])
        else
            out[cartesian] = out[cartesian[2], cartesian[1]]
        end
    end
end

struct MaternCovariance1_2{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end
function (cov::MaternCovariance1_2)(x::PointMeasurement, y::PointMeasurement)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return exp(-dist/sigma)
end

struct MaternCovariance3_2{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end
function (cov::MaternCovariance3_2)(x::PointMeasurement, y::PointMeasurement)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return (1+sqrt(3)*dist/sigma)*exp(-sqrt(3)*dist/sigma)
end

struct MaternCovariance5_2{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end

function (cov::MaternCovariance5_2)(x::PointMeasurement, y::PointMeasurement)
    # d = length(x.coordinate)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    F(t,a) = (1 + sqrt(5)*t/a + 5*t^2/(3*a^2))*exp(-sqrt(5)*t/a)
    return F(dist,sigma)
end

function (cov::MaternCovariance5_2)(x::ΔδPointMeasurement, y::ΔδPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;

    F(t,a) = (1 + sqrt(5)*t/a + 5*t^2/(3*a^2))*exp(-sqrt(5)*t/a)
    D2F(t,a) = -5*(d*a^2+sqrt(5)*d*a*t-5*t^2)/(3*a^4) * exp(-sqrt(5)*t/a);
    D4F(t,a) = 25*(d*(d+2)*a^2-(3+2*d)*sqrt(5)*a*t+5*t^2)/(3*a^6) * exp(-sqrt(5)*t/a);
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma)
end

function (cov::MaternCovariance5_2)(x::Δ∇δPointMeasurement, y::Δ∇δPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    wg_x = x.weight_∇
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    wg_y = y.weight_∇

    F(t,a) = (1 + sqrt(5)*t/a + 5*t^2/(3*a^2))*exp(-sqrt(5)*t/a)
    D2F(t,a) = -5*(d*a^2+sqrt(5)*d*a*t-5*t^2)/(3*a^4) * exp(-sqrt(5)*t/a)
    D4F(t,a) = 25*(d*(d+2)*a^2-(3+2*d)*sqrt(5)*a*t+5*t^2)/(3*a^6) * exp(-sqrt(5)*t/a)
    DF(t,a) = -5*(a+sqrt(5)*t)*exp(-sqrt(5)*t/a)/(3*a^3)
    D3F(t,a) = 25*exp(-sqrt(5)*t/a)*(a*(2+d)-sqrt(5)*t)/(3*a^5)
    DDF(t,a) = 25*exp(-sqrt(5)*t/a)/(3*a^4)
    vec = x.coordinate - y.coordinate
    dist = norm(vec);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma) - w1_x*D3F(dist,sigma)*sum(vec.*wg_y) + w1_y*D3F(dist,sigma)*sum(vec.*wg_x) - w2_x*DF(dist,sigma)*sum(vec.*wg_y) + w2_y*DF(dist,sigma)*sum(vec.*wg_x) + (sum(-wg_x.*wg_y)*DF(dist,sigma)+sum(wg_x.*vec)*sum(-wg_y.*vec)*DDF(dist,sigma))
end

struct MaternCovariance7_2{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end

function (cov::MaternCovariance7_2)(x::PointMeasurement, y::PointMeasurement)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    F(t,a) = (15*a^3+15*sqrt(7)*a^2*t+42*a*t^2+7*sqrt(7)*t^3)/(15*a^3)*exp(-sqrt(7)*t/a);
    return F(dist,sigma)
end

function (cov::MaternCovariance7_2)(x::ΔδPointMeasurement, y::ΔδPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    F(t,a) = (15*a^3+15*sqrt(7)*a^2*t+42*a*t^2+7*sqrt(7)*t^3)/(15*a^3)*exp(-sqrt(7)*t/a);
    D2F(t,a) = -7*(3*d*a^3+3*sqrt(7)*a^2*d*t+7*a*(d-1)*t^2-7*sqrt(7)*t^3)/(15*a^5)*exp(-sqrt(7)*t/a);
    D4F(t,a) = 49*(d*(d+2)*a^3+d*(d+2)*sqrt(7)*a^2*t-14*a*(2+d)*t^2+7*sqrt(7)*t^3)/(15*a^7)*exp(-sqrt(7)*t/a);
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma)
end

function (cov::MaternCovariance7_2)(x::Δ∇δPointMeasurement, y::Δ∇δPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    wg_x = x.weight_∇
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    wg_y = y.weight_∇

    F(t,a) = (15*a^3+15*sqrt(7)*a^2*t+42*a*t^2+7*sqrt(7)*t^3)/(15*a^3)*exp(-sqrt(7)*t/a);
    D2F(t,a) = -7*(3*d*a^3+3*sqrt(7)*a^2*d*t+7*a*(d-1)*t^2-7*sqrt(7)*t^3)/(15*a^5)*exp(-sqrt(7)*t/a);
    D4F(t,a) = 49*(d*(d+2)*a^3+d*(d+2)*sqrt(7)*a^2*t-14*a*(2+d)*t^2+7*sqrt(7)*t^3)/(15*a^7)*exp(-sqrt(7)*t/a);
    DF(t,a) = -7*(3*a^2+3*sqrt(7)*a*t+7*t^2)*exp(-sqrt(7)*t/a)/(15*a^4)
    D3F(t,a) = 49*exp(-sqrt(7)*t/a)*(a^2*(2+d)+sqrt(7)*a*(2+d)*t-7*t^2)/(15*a^6)
    DDF(t,a) = 49*exp(-sqrt(7)*t/a)*(a+sqrt(7)*t)/(15*a^5)
    vec = x.coordinate - y.coordinate
    dist = norm(vec);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma) - w1_x*D3F(dist,sigma)*sum(vec.*wg_y) + w1_y*D3F(dist,sigma)*sum(vec.*wg_x) - w2_x*DF(dist,sigma)*sum(vec.*wg_y) + w2_y*DF(dist,sigma)*sum(vec.*wg_x) + (sum(-wg_x.*wg_y)*DF(dist,sigma)+sum(wg_x.*vec)*sum(-wg_y.*vec)*DDF(dist,sigma))
end

struct MaternCovariance9_2{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end

function (cov::MaternCovariance9_2)(x::PointMeasurement, y::PointMeasurement)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    F(t,a) = (35*a^4+105*a^3*t+135*a^2*t^2+90*a*t^3+27*t^4)/(35*a^4)*exp(-3*t/a);
    return F(dist,sigma)
end

function (cov::MaternCovariance9_2)(x::ΔδPointMeasurement, y::ΔδPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    F(t,a) = (35*a^4+105*a^3*t+135*a^2*t^2+90*a*t^3+27*t^4)/(35*a^4)*exp(-3*t/a);
    D2F(t,a) = -9*(5*d*a^4+15*d*a^3*t+9*a^2*(2*d-1)*t^2+9*a*(d-3)*t^3-27*t^4)/(35*a^6)*exp(-3*t/a);
    D4F(t,a) = 81*(d*(d+2)*a^4+3*a^3*d*(d+2)*t+3*a^2*(d^2-4)-18*a*(d+2)*t^3+27*t^4)/(35*a^8)*exp(-3*t/a);
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma)
end

function (cov::MaternCovariance9_2)(x::Δ∇δPointMeasurement, y::Δ∇δPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    wg_x = x.weight_∇
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    wg_y = y.weight_∇

    F(t,a) = (35*a^4+105*a^3*t+135*a^2*t^2+90*a*t^3+27*t^4)/(35*a^4)*exp(-3*t/a);
    D2F(t,a) = -9*(5*d*a^4+15*d*a^3*t+9*a^2*(2*d-1)*t^2+9*a*(d-3)*t^3-27*t^4)/(35*a^6)*exp(-3*t/a);
    D4F(t,a) = 81*(d*(d+2)*a^4+3*a^3*d*(d+2)*t+3*a^2*(d^2-4)*(t^2)-18*a*(d+2)*t^3+27*t^4)/(35*a^8)*exp(-3*t/a);
    DF(t,a) = -9*(5*a^3+15*a^2*t+18*a*t^2+9*t^3)*exp(-3*t/a)/(35*a^5)
    D3F(t,a) = 81*exp(-3*t/a)*(a^3*(2+d)+3*a^2*(2+d)*t+3*a*(1+d)*t^2-9*t^3)/(35*a^7)
    DDF(t,a) = 81*exp(-3*t/a)*(a^2+3*a*t+3*t^2)/(35*a^6)
    vec = x.coordinate - y.coordinate
    dist = norm(vec);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma) - w1_x*D3F(dist,sigma)*sum(vec.*wg_y) + w1_y*D3F(dist,sigma)*sum(vec.*wg_x) - w2_x*DF(dist,sigma)*sum(vec.*wg_y) + w2_y*DF(dist,sigma)*sum(vec.*wg_x) + (sum(-wg_x.*wg_y)*DF(dist,sigma)+sum(wg_x.*vec)*sum(-wg_y.*vec)*DDF(dist,sigma))
end

struct MaternCovariance11_2{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end

# Matern covariance function
function (cov::MaternCovariance11_2)(x::PointMeasurement, y::PointMeasurement)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    F(t,a) = (945*a^5+945*sqrt(11)*a^4*t+4620*a^3*t^2+1155*sqrt(11)*a^2*t^3+1815*a*t^4+121*sqrt(11)*t^5)/(945*a^5)*exp(-sqrt(11)*t/a);
    return F(dist,sigma)
end

function (cov::MaternCovariance11_2)(x::ΔδPointMeasurement, y::ΔδPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    F(t,a) = (945*a^5+945*sqrt(11)*a^4*t+4620*a^3*t^2+1155*sqrt(11)*a^2*t^3+1815*a*t^4+121*sqrt(11)*t^5)/(945*a^5)*exp(-sqrt(11)*t/a);
    D2F(t,a) = -11*(105*d*a^5+105*d*sqrt(11)*a^4*t+165*(3*d-1)*a^3*t^2+55*sqrt(11)*(2*d-3)*a^2*t^3+121*(d-6)*a*t^4-121*sqrt(11)*t^5)/(945*a^7)*exp(-sqrt(11)*t/a);
    D4F(t,a) = 121*(15*d*(2+d)*a^5+15*d*(2+d)*sqrt(11)*a^4*t+66*(d^2+d-2)*a^3*t^2+11*sqrt(11)*a^2*(d^2-4*d-12)*t^3-121*(3+2*d)*a*t^4+121*sqrt(11)*t^5)/(945*a^9)*exp(-sqrt(11)*t/a);
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma)
end

function (cov::MaternCovariance11_2)(x::Δ∇δPointMeasurement, y::Δ∇δPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    wg_x = x.weight_∇
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    wg_y = y.weight_∇

    F(t,a) = (945*a^5+945*sqrt(11)*a^4*t+4620*a^3*t^2+1155*sqrt(11)*a^2*t^3+1815*a*t^4+121*sqrt(11)*t^5)/(945*a^5)*exp(-sqrt(11)*t/a);
    D2F(t,a) = -11*(105*d*a^5+105*d*sqrt(11)*a^4*t+165*(3*d-1)*a^3*t^2+55*sqrt(11)*(2*d-3)*a^2*t^3+121*(d-6)*a*t^4-121*sqrt(11)*t^5)/(945*a^7)*exp(-sqrt(11)*t/a);
    D4F(t,a) = 121*(15*d*(2+d)*a^5+15*d*(2+d)*sqrt(11)*a^4*t+66*(d^2+d-2)*a^3*t^2+11*sqrt(11)*a^2*(d^2-4*d-12)*t^3-121*(3+2*d)*a*t^4+121*sqrt(11)*t^5)/(945*a^9)*exp(-sqrt(11)*t/a);
    DF(t,a) = -11*(105*a^4+105*sqrt(11)*a^3*t+495*a^2*t^2+110*sqrt(11)*a*t^3+121*t^4)*exp(-sqrt(11)*t/a)/(945*a^6)
    D3F(t,a) = 121*exp(-sqrt(11)*t/a)*(15*a^4*(2+d)+15*sqrt(11)*a^3*(2+d)*t+33*a^2*(3+2*d)*t^2+11*sqrt*(11)*a*(d-1)*t^3-121*t^4)/(945*a^8)
    DDF(t,a) = 121*exp(-sqrt(11)*t/a)*(15*a^3+15*sqrt(11)*a^2*t+66*a*t^2+11*sqrt(11)*t^3)/(945*a^7)
    vec = x.coordinate - y.coordinate
    dist = norm(vec);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma) - w1_x*D3F(dist,sigma)*sum(vec.*wg_y) + w1_y*D3F(dist,sigma)*sum(vec.*wg_x) - w2_x*DF(dist,sigma)*sum(vec.*wg_y) + w2_y*DF(dist,sigma)*sum(vec.*wg_x) + (sum(-wg_x.*wg_y)*DF(dist,sigma)+sum(wg_x.*vec)*sum(-wg_y.*vec)*DDF(dist,sigma))
end

# The Gaussian covariance function
struct GaussianCovariance{Tv}<:AbstractCovarianceFunction{Tv}
    length_scale::Tv
end

function (cov::GaussianCovariance)(x::PointMeasurement, y::PointMeasurement)
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return exp(-dist^2/(2*sigma^2))
end

function (cov::GaussianCovariance)(x::ΔδPointMeasurement, y::ΔδPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    F(t,a) = exp(-t^2/(2*a^2))
    D2F(t,a) = (t^2 - a^2*d)/(a^4)*exp(-t^2/(2*a^2));
    D4F(t,a) = (a^4*d*(2+d)-2*a^2*(2+d)*t^2+t^4)*exp(-t^2/(2*a^2))/a^8
    dist = norm(x.coordinate - y.coordinate);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma)
end

function (cov::GaussianCovariance)(x::Δ∇δPointMeasurement, y::Δ∇δPointMeasurement)
    d = length(x.coordinate);
    w1_x = x.weight_Δ;
    w2_x = x.weight_δ;
    wg_x = x.weight_∇
    w1_y = y.weight_Δ;
    w2_y = y.weight_δ;
    wg_y = y.weight_∇

    F(t,a) = exp(-t^2/(2*a^2))
    D2F(t,a) = (t^2 - a^2*d)/(a^4)*exp(-t^2/(2*a^2));
    D4F(t,a) = (a^4*d*(2+d)-2*a^2*(2+d)*t^2+t^4)*exp(-t^2/(2*a^2))/a^8
    DF(t,a) = -exp(-t^2/(2*a^2))/a^2
    D3F(t,a) = exp(-t^2/(2*a^2))*(a^2*(2+d)*t-t^3)/a^6
    DDF(t,a) = exp(-t^2/(2*a^2))/a^4
    vec = x.coordinate - y.coordinate
    dist = norm(vec);
    sigma = cov.length_scale;
    return w1_x*w1_y*D4F(dist,sigma) + (w2_x*w1_y+w1_x*w2_y)*D2F(dist,sigma) + w2_x*w2_y*F(dist,sigma) - w1_x*D3F(dist,sigma)*sum(vec.*wg_y) + w1_y*D3F(dist,sigma)*sum(vec.*wg_x) - w2_x*DF(dist,sigma)*sum(vec.*wg_y) + w2_y*DF(dist,sigma)*sum(vec.*wg_x) + (sum(-wg_x.*wg_y)*DF(dist,sigma)+sum(wg_x.*vec)*sum(-wg_y.*vec)*DDF(dist,sigma))
end

function (cov::AbstractCovarianceFunction)(x::ΔδPointMeasurement, y::PointMeasurement)
    return cov(x, ΔδPointMeasurement(y))
end
function (cov::AbstractCovarianceFunction)(x::PointMeasurement, y::ΔδPointMeasurement)
    return (cov::AbstractCovarianceFunction)(y,x)
end

function (cov::AbstractCovarianceFunction)(x::Δ∇δPointMeasurement, y::PointMeasurement)
    return cov(x, Δ∇δPointMeasurement(y))
end
function (cov::AbstractCovarianceFunction)(x::PointMeasurement, y::Δ∇δPointMeasurement)
    return (cov::AbstractCovarianceFunction)(y,x)
end


"""
This module calculates (generalized) hypergeometric functions:

    mFn(a;b;z) = ∑_{k=0}^∞ (a_1,k) ⋯ (a_m,k) / (b_1,k) ⋯ (b_n,k) zᵏ/k!
"""
module HypergeometricFunctions

using DualNumbers
import ApproxFun: @clenshaw, real, eps
import FastTransforms: pochhammer

export _₂F₁, _₃F₂, mFn

include("Gauss.jl")
include("specialfunctions.jl")

end #module

using .HypergeometricFunctions

export _₂F₁, _₃F₂, mFn


Zero = [1, 0]
One = [0, 1]

I = [1 0
     0 1]

X = [0 1
     1 0]

Y = [0 -im
     im 0]

Z = [1  0
     0 -1]

H = [1  1
     1 -1]/√2

S = classical((a, b) -> (b, a), 2)

control(U) = Zero*Zero' ⊗ eye(U) + One*One' ⊗ U

CX = control(X)

T = control(CX)


function init_observable_mappings!(m::RealBond)

    observables = OrderedDict{Symbol,Observable}()
    population_mnemonic = get(get_setting(m, :population_mnemonic))

    ############################################################################
    ## 1. GDP
    ############################################################################
    gdp_fwd_transform = function (levels)
        # FROM: Level of nominal GDP (FRED :GDP series)
        # TO:   The cyclical component of log real GDP

        levels[:rGDP] = nominal_to_real(:GDP, levels)
        levels[:lGDP] = log.(levels[:rGDP])
        _, levels[:lGDP_cyc_comp] = hpfilter(levels[:lGDP], 1600)

        levels[:lGDP_cyc_comp]
    end

    gdp_rev_transform = logleveltopct_annualized

    observables[:obs_gdp] = Observable(:obs_gdp, [:GDP__FRED, population_mnemonic, :GDPDEF__FRED],
                                       gdp_fwd_transform, gdp_rev_transform,
                                       "Log Real GDP", "")

    m.observable_mappings = observables
end


using LinearAlgebra
using LineSearches, Optim

"""
Implements the Cyclic Coordinate Search optimization method. Uses a Hager-Zhang
linesearch to determine step size.
"""
function ccs(fn, x0; ϵ=.001, max_iter=100)
    algo_hz = Newton(linesearch = HagerZhang())
    i = 1
    k = 0
    n = length(x0)
    results = [(x0, fn(x0))]
    e_hat = Matrix(1I, n, n)

    while k < max_iter
        x, _ = results[end]

        if i == n + 1
            d_k = x - results[k-n+1][1]
            i = 1
        else
            d_k = e_hat[:, i]
        end

        α_fn(α) = fn(x + α .* d_k)
        res = Optim.optimize(α_fn, x, method=algo_hz)

        α = res.minimizer
        x_new = x + α .* d_k
        f_new = fn(x_new)
        push!(results, (x_new, f_new))

        k += 1
        i += 1

        if norm(x_new-x) < ϵ
            return Response(x_new, f_new, k, true)
        end
    end

    println("Maximum iterations reached.")
    x, fx = results[end]
    return Response(x, fx, k, false)
end


isunitary(op) = op' * op ≈ IMatrix(size(op, 1))
isunitary(op::Number) = op' * op ≈ one(op)

isreflexive(op) = op * op ≈ IMatrix(size(op, 1))
isreflexive(op::Number) = op * op ≈ one(op)

"""
    ishermitian(op) -> Bool

check if this operator is hermitian.
"""
LinearAlgebra.ishermitian(op) = op' ≈ op

function iscommute(ops...)
    n = length(ops)
    for i = 1:n
        for j = i+1:n
            iscommute(ops[i], ops[j]) || return false
        end
    end
    true
end

iscommute(op1, op2) = op1 * op2 ≈ op2 * op1


# Intersect
# ---------------------------------------------------------------------------------------------
#
# Intersection of a line segment and polygon in 2D
function intersect(l::LineSegment2D{T}, poly::Polygon{N, 2, T}) where {N,T}
    # Create the line segments that make up the polygon and intersect each one
    # until 2 unique points have been found
    p₁ = nan(Point2D{T})
    npoints = 0x0000
    for i ∈ 1:N-1
        hit, point = l ∩ LineSegment2D(poly[i], poly[i + 1])
        if hit
            if npoints === 0x0000
                npoints = 0x0001
                p₁ = point
            elseif !(p₁ ≈ point)
                return true, SVector(p₁, point)
            end
        end
    end
    hit, point = l ∩ LineSegment2D(poly[N], poly[1])
    if hit
        if npoints === 0x0000
            npoints = 0x0001
            p₁ = point
        elseif !(p₁ ≈ point)
            return true, SVector(p₁, point)
        end
    end
    return false, SVector(p₁, point)
end

# Möller, T., & Trumbore, B. (1997). Fast, minimum storage ray-triangle intersection.
function intersect(l::LineSegment3D{T}, tri::Triangle3D{T}) where {T}
    p = Point3D{T}(0, 0, 0)
    𝗲₁ = tri[2] - tri[1]
    𝗲₂ = tri[3] - tri[1]
    𝗱 = l.𝘂
    𝗽 = 𝗱 × 𝗲₂
    det = 𝗽 ⋅ 𝗲₁
    (det > -1e-8 && det < 1e-8) && return (false, p)
    inv_det = 1/det
    𝘁 = l.𝘅₁ - tri[1]
    u = (𝘁 ⋅ 𝗽)*inv_det
    (u < 0 || u > 1) && return (false, p)
    𝗾 = 𝘁 × 𝗲₁
    v = (𝗾 ⋅ 𝗱)*inv_det
    (v < 0 || u + v > 1) && return (false, p)
    t = (𝗾 ⋅ 𝗲₂)*inv_det
    (t < 0 || t > 1) && return (false, p)
    return (true, l(t))
end


module ntest
using  ..nonlineq

function test_broyden()

    fcn(x) = [ (x[1] + 3) * (x[2]^2 - 7) + 18,  sin( x[2] * exp(x[1]) - 1) ]

    x0 = [-0.5, 1.4]

    (x_sol, iter) = nonlineq.broyden(fcn, x0)

    println("Solution at: $(x_sol)")
    println("Iter: $(iter)" )

end

end # ntest


_diffusion_jacobian(model::AbstractSDE{D,M}, t0, x0) where {D,M} = ForwardDiff.jacobian(x->diffusion(model, t0, x), x0)
_diffusion_jacobian(model::AbstractSDE{1,M}, t0, x0) where {M} = ForwardDiff.derivative(x->diffusion(model, t0, x), x0)

function step(model::AbstractSDE{D,1}, scheme::Milstein, t0, x0, Δw) where D
  μ = drift(model, t0, x0)
  σ = diffusion(model, t0, x0)
  ∂σ = _diffusion_jacobian(model, t0, x0)
  x = x0 + μ * scheme.Δt + σ * Δw + ∂σ * σ * (Δw.^2 - scheme.Δt) / 2
  x
end


using Test
using MOSLab
SemiUstrip()
S = Silicon()

F = BoltzmanDist()
T = 300.0*Tunit

niB(T,S)
n(T,F,S)
@test isapprox(Eg(T,S.Eg),1.05; atol=0.1)
@test isapprox(log10(ni(T,F,S)),10.0; atol=0.5)
@test n(T,F,S) ≈ ni(T,F,S)
@test p(T,F,S) ≈ ni(T,F,S)
@test Nv(T,S) == 1.04e19
@test Nc(T,S) == 2.8e19
@test p(T,F,S) ≈ n(T,F,S)
@test Ef(T,F,S) ≈ Ei(T,F,S)
@test isapprox(Ei(T,F,S), 0.5*(Ec(T,S)+Ev(T,S)) + kb*T*0.5*(Nv(T,S)/Nc(T,S));atol=0.1)

PD = PDoppingC(1e15)
@test Na(PD) == 1e15
@test Nd(PD) == 0
ND = NDoppingC(1e15)
@test Nd(ND) == 1e15
@test Na(ND) == 0

PS = PSiliconNU(1e15)
@test Na(PS.Dopping) == 1e15
@test P(T,F,PS) ≈ N(T,F,PS)
@test p(T,F,PS) ≈ Na(PS.Dopping)
@test Ef(T,F,PS) < Ei(T,F,PS)
P(T,F,PS)
n(T,F,PS)
NS = NSiliconNU(1e15)
@test Nd(NS.Dopping) == 1e15
@test P(T,F,NS)≈N(T,F,NS)
@test n(T,F,NS) ≈ Nd(NS.Dopping)
@test Ef(T,F,NS) > Ei(T,F,NS)

SDI = SemiconductorData(T,F,Silicon();psi=0.0)
@test isapprox(Eg(SDI),1.05; atol=0.1)
@test isapprox(log10(ni(SDI)),10.0; atol=0.5)
@test n(SDI) ≈ ni(SDI)
@test p(SDI) ≈ ni(SDI)
@test p(SDI) ≈ n(SDI)
@test Ef(SDI) ≈ Ei(SDI)
@test isapprox(Ei(SDI), 0.5*(Ec(SDI)+Ev(SDI)) + kb*T*0.5*(Nv(T,S)/Nc(T,S));atol=0.1)


export ℎ, ħ, ℏ, 𝑒, Φ₀, ϕ₀, ln2,
    sigma0, σ0, sigmax, σx, sigmay, σy, sigmaz, σz, sigmaplus, σ₊, sigmaminus, σ₋,
    ⊗

# Physical and mathematical constants
const ℎ = 6.626069934E-34 # Planck constant, in Joules*seconds, NIST 2017
const ħ = 1.054571783E-34 # Reduced Planck constant
const ℏ = ħ # Actual unicode symbol
const 𝑒 = 1.6021766208E-19 # Elementary charge, in Coulomb, 2014 CODATA
const Φ₀ = 2.067833831E-15 # Magnetic flux quantum (ℎ/2𝑒), in Weber, 2014 CODATA
const ϕ₀ = 3.291059757E-16 # Reduced flux quantum (ℏ/2𝑒), in Weber, 2014 CODATA
const ln2 = 0.693147180559945309 # Natural logarithm of 2
const sqrtln16 = 1.66510922231539551 # Square root of ln(16)

# Pauli matrices (spin-1/2 operators)
const sigma0 = qeye(2)
const sigmax = Operator(sparse([1,2],[2,1],[1.0,1.0],2,2),(2,))
const sigmay = Operator(sparse([1,2],[2,1],[-1.0im,1.0im],2,2),(2,))
const sigmaz = Operator(sparse([1,2],[1,2],[1.0,-1.0],2,2),(2,))
const σ0 = sigma0
const σx = sigmax
const σy = sigmay
const σz = sigmaz
const sigmaplus  = Operator(sparse([2],[1],[1.0],2,2),(2,)) # σ₊*[1,0] = [0,1]
const sigmaminus = Operator(sparse([1],[2],[1.0],2,2),(2,))
const σ₊ = sigmaplus
const σ₋ = sigmaminus

# Function aliases
const ⊗ = kron
const tensor = kron
const fock = basis
const qzeros = qzero

# Table of square root of factorials √(n!) for n ∈ 0:127 (1KiB)
const sqrtfact_table = Float64[
    1.0000000000000000000, 1.0000000000000000000, 1.4142135623730950488,
    2.4494897427831780982, 4.8989794855663561964, 10.954451150103322269,
    26.832815729997476357, 70.992957397195392511, 200.79840636817813151,
    602.39521910453439454, 1904.9409439665052252, 6317.9743589223278835,
    21886.105181141755630, 78911.474450804681438, 295259.70128007648582,
    1.1435359058639129588E6, 4.5741436234556518350E6,
    1.8859677306253148084E7, 8.0014834285449844953E7,
    3.4877657663442939413E8, 1.5597762686284978865E9,
    7.1477928181858656891E9, 3.3526120082371710076E10,
    1.6078562354540587669E11, 7.8768547132293829008E11,
    3.9384273566146914504E12, 2.0082117944245961319E13,
    1.0434974580907397764E14, 5.5216695356722848749E14,
    2.9735100460129106405E15, 1.6286585271694955843E16,
    9.0679869067935485290E16, 5.1296280268036349125E17,
    2.9467469553410734788E18, 1.7182339742875652406E19,
    1.0165209277917570217E20, 6.0991255667505421303E20,
    3.7099532465014090857E21, 2.2869687743093501008E22,
    1.4282115417961529469E23, 9.0328029052332240864E23,
    5.7838159214452708158E24, 3.7483411234209726053E25,
    2.4579516484946125896E26, 1.6304206741784307882E27,
    1.0937194378152021970E28, 7.4179661362209580728E28,
    5.0855013667402369566E29, 3.5233386996620226535E30,
    2.4663370897634158575E31, 1.7439636808636061170E32,
    1.2454391808865586995E33, 8.9809896543167155890E33,
    6.5382591597917144388E34, 4.8046196242703894246E35,
    3.5632012788584194610E36, 2.6664556771205919519E37,
    2.0131298891248228833E38, 1.5331540468207617694E39,
    1.1776379687564843276E40, 9.1219444817107882595E40,
    7.1244663931920178495E41, 5.6098104478126475754E42,
    4.4526490041372451123E43, 3.5621192033097960898E44,
    2.8718723147247460219E45, 2.3331200978034608324E46,
    1.9097411059666879697E47, 1.5748128594969087944E48,
    1.3081378078327271991E49, 1.0944666130115569586E50,
    9.2221396029764281460E50, 7.8252449403763769254E51,
    6.6858922078602825325E52, 5.7514219472399922436E53,
    4.9808775141931966697E54, 4.3422283469044442401E55,
    3.8102899106011063425E56, 3.3651569321810681095E57,
    2.9910169058002623210E58, 2.6752468492881886262E59,
    2.4077221643593697636E60, 2.1802851503903891306E61,
    1.9863343046226278804E62, 1.8205054612841328324E63,
    1.6784231035053557904E64, 1.5565055535934567420E65,
    1.4518117296604018405E66, 1.3619201234191322394E67,
    1.2848328747704294744E68, 1.2188994890809338170E69,
    1.1627560052213890684E70, 1.1152763807523813626E71,
    1.0755335917960171153E72, 1.0427685057848376926E73,
    1.0163650175128549387E74, 9.9583027412855333880E74,
    9.8077907646157562109E75, 9.7092175013660332284E76,
    9.6605494379949297313E77, 9.6605494379949297313E78,
    9.7087320283538361861E79, 9.8053387065593642395E80,
    9.9513319291872585027E81, 1.0148407138632952789E83,
    1.0399022830248479171E84, 1.0706449288791818466E85,
    1.1074837259283487726E86, 1.1509308491181513962E87,
    1.2016070835354182374E88, 1.2602561412358770729E89,
    1.3277622343967479571E90, 1.4051714689748984594E91,
    1.4937177607097713179E92, 1.5948541417547209534E93,
    1.7102905289724946877E94, 1.8420392733196266468E95,
    1.9924701154117019517E96, 2.1643765498993678241E97,
    2.3610560690520684628E98, 2.5864073371085860453E99,
    2.8450480708194446498E100, 3.1424583053452915085E101,
    3.4851548555301426772E102, 3.8809042007129483894E103,
    4.3389828034793201975E104, 4.8704961173170509835E105,
    5.4887703709097355030E106
]


################################################################################
#
#  Free part of the unit group
#
################################################################################

@doc Markdown.doc"""
    regulator(x::Array{T, 1}, abs_tol::Int) -> arb

Compute the regulator $r$ of the elements in $x$, such that the radius of $r$
is less then `-2^abs_tol`.
"""
function regulator(x::Array{T, 1}, abs_tol::Int) where T
  K = _base_ring(x[1])
  deg = degree(K)
  r1, r2 = signature(K)
  rr = r1 + r2
  r = rr - 1 # unit rank

  p = 32

  conlog = Vector{Vector{arb}}(undef, r)

  while true
    q = 2
    for i in 1:r
      conlog[i] = conjugates_arb_log(x[i], p)
      for j in 1:r
        q = max(q, bits(conlog[i][j]))
      end
    end

    A = zero_matrix(ArbField(q, false), r, r)

    for i in 1:r
      for j in 1:r
        A[i, j] = conlog[i][j]
      end
    end

    z = abs(det(A))

    if isfinite(z) && radiuslttwopower(z, -abs_tol)
      return z
    end

    p = 2*p
  end
end

function lower_regulator_bound(K::AnticNumberField)
  return ArbField(64, false)("0.054")
end


using Test, MPI
include(joinpath("..", "..", "testhelpers.jl"))

include("iterativesolvers.jl")

@testset "Linear Solvers Poisson" begin
    runmpi(joinpath(@__DIR__, "columnwiselu.jl"))
    runmpi(joinpath(@__DIR__, "poisson.jl"))
    runmpi(joinpath(@__DIR__, "bandedsystem.jl"))
    runmpi(joinpath(@__DIR__, "cg.jl"))
    runmpi(joinpath(@__DIR__, "bgmres.jl"))
end


######
# This file is part of the MomentArithmetic.jl package, for performing arithmetic operations
# between moments of uncertain numbers, when the moments and the uncertain numbers are only
# partially known.
#
#   Definition of Moment Arithmetic
#
#           University of Liverpool,
#           Institute for Risk and Unertainty
#
#                                           Authors: Ander Gray, Scott Ferson
#                                           Email:  ander.gray@liverpool.ac.uk
######

##
#   Rowe mean
##

function roweFormula(EX, VX, LX , GX, t :: Function)

    P = 1 / (1 + (EX - LX)^2/ VX)
    Q = 1 / (1 + (GX - EX)^2/ VX)

    MIN = P * t(LX) + (1 - P) * t(EX + VX/(EX - LX))
    MAX = Q * t(GX) + (1 - Q) * t(EX + VX/(EX - GX))

    return env(MIN, MAX)

end

function rowe(x :: AbstractMoment, t :: Function; Nsub = Nsub[1])

    if !hasIntMoments(x); return roweNoSub(x,t); end

    EXs = x.mean; VXs = x.var;
    LX = left(x); GX = right(x);

    if typeof(EXs) <: Interval;  EXs = split(EXs, Nsub); end    # Sub-intervalise
    if typeof(VXs) <: Interval;  VXs = split(VXs, Nsub); end

    outs = [roweFormula(EX, VX, LX, GX, t) for EX in EXs, VX in VXs]

    return hull(outs[:])        # Same as env

end

# No sub-intervalisation
function roweNoSub(x :: AbstractMoment, t :: Function)

    EX = x.mean; VX = x.var;
    LX = left(x); GX = right(x);

    return roweFormula(EX, VX, LX, GX, t)

end

###
#   Rowe Variance
##

function rowevarFormula(EX, VX, LX , GX, t :: Function, invT :: Function)

    v = invT(roweFormula(EX, VX, LX , GX, t))

    Lv = left(v); Gv = right(v);

    MIN = ((t(Lv) - t(LX))/(Lv - LX))^2 * (VX + (Lv - EX)^2)
    MAX = ((t(Gv) - t(GX))/(Gv - GX))^2 * (VX + (Gv - EX)^2)

    if isnan(MIN); MIN =0; end
    if isnan(MAX); MAX =0; end

    return env(MIN, MAX)

end

function rowevar(x :: AbstractMoment, t :: Function, invT :: Function; Nsub = Nsub[1])

    if !hasIntMoments(x); return rowevarNoSub(x,t, invT); end

    EXs = x.mean; VXs = x.var;
    LX = left(x); GX = right(x);

    if typeof(EXs) <: Interval;  EXs = split(EXs, Nsub); end    # Sub-intervalise
    if typeof(VXs) <: Interval;  VXs = split(VXs, Nsub); end

    outs = [rowevarFormula(EX, VX, LX, GX, t, invT) for EX in EXs, VX in VXs]

    out = hull(outs[:])
    if isnan(out); out = interval(0, Inf); end
    return out

end

# No sub-intervalisation
function rowevarNoSub(x :: AbstractMoment, t :: Function, invT :: Function)

    EX = x.mean; VX = x.var;
    LX = left(x); GX = right(x);

    return rowevarFormula(EX, VX, LX, GX, t, invT)

end

###
#   Univariate transformations
###
-(x :: AbstractMoment) = Moments(-x.mean, x.var, -x.range)

function reciprocal(x :: AbstractMoment)

    if 0 ∈ x.range
        throw(ArgumentError("Tried to evaluate 1/X, with 0 ∈ X. Interval range is [-∞, ∞], moments not defined"))
    end

    f(x) = 1/x
    invF(x) = 1/x

    MEAN = rowe(x, f)
    VAR = rowevar(x, f, invF)
    RANGE = f(x.range)

    return Moments(MEAN, VAR, RANGE)

end

reciprocal(x :: Real) = 1/x

inv(x :: AbstractMoment) = reciprocal(x)

function exp(x :: AbstractMoment)

    MEAN = rowe(x, exp)
    VAR = rowevar(x, exp, log)
    RANGE = exp(x.range)

    MEAN2 = rowe(2 * x, exp)
    homespun_var = MEAN2 - MEAN^2

    VAR = VAR ∩ homespun_var

    return Moments(MEAN, VAR, RANGE)

end

#^(x :: AbstractMoment) = throw(ArgumentError("rowe exponentiate must be implemented"))

function ^(x :: AbstractMoment, a :: Integer)

    if a == 0;
        return Moments(1, 0, interval(1));
    end

    if a == 1; return x; end

    if a < 0;
        if 0 ∈ x.range
            throw(ArgumentError("Tried to evaluate 1/0, when 0 ∈ X when performing X^a where a < 0. Interval range is [-∞, ∞], Moments not defined"))
        end
        return reciprocal(x)^a;
    end

    if isodd(a) && 0 > x.range.lo && 0 < x.range.hi;
        return multFrechet(x^(a-1), x);
    end

    f(x) = x^a
    invF(x) = x^(1/a)

    MEAN = rowe(x, f)
    VAR = rowevar(x, f, invF)
    RANGE = f(x.range)

    if a == 2
        var_hs = homespun_var_square(x)
        VAR = VAR ∩ var_hs
    end

    return Moments(MEAN, VAR, RANGE)
end

function homespun_var_square( x :: Moments )
    that = x^4
    this = that.mean - (x.mean^2 + x.var)^2
    return max(0,this)
end

function ^(x :: AbstractMoment, a :: Float64)

    if isinteger(a); return x^Integer(a); end

    if x.range.lo < 0;
        throw(ArgumentError("Cannot take the power of a negative number to a real number. Yields imaginary numbers, for which moments are currently undefined"))
    end

    f(x) = x^a
    invF(x) = x^(1/a)

    MEAN = rowe(x, f)
    VAR = rowevar(x, f, invF)
    RANGE = f(x.range)

    return Moments(MEAN, VAR, RANGE)
end

function ^(x :: AbstractMoment, a :: Interval)

    if isinteger(a); return x^Integer(a); end

    if x.range.lo < 0;
        throw(ArgumentError("Cannot take the power of a negative number to a real number. Yields imaginary numbers, for which moments are currently undefined"))
    end

    x1 = x^(a.lo)
    x2 = x^(a.hi)

    out = env(x1, x2)

    unitMom = Moments(1,0,interval(1))

    if 0 ∈ a; out = env(out, unitMom); end
    if 1 ∈ x.range; out = env(out, unitMom); end

    return out
end

function sqrt(x :: AbstractMoment)

    f(x) = sqrt(x)
    invF(x) = x^2

    MEAN = rowe(x, f)
    VAR = rowevar(x, f, invF)
    RANGE = f(x.range)

    homespun_var = x.mean - MEAN^2
    VAR = VAR ∩ homespun_var;

    return Moments(MEAN, VAR, RANGE)
end

function ln(x :: AbstractMoment)

    f(x) = log(x)
    invF(x) = exp(x)

    MEAN = rowe(x, f)
    VAR = rowevar(x, f, invF)
    RANGE = f(x.range)

    return Moments(MEAN, VAR, RANGE)
end

function log(x = missing :: AbstractMoment, base = ℯ :: Real)
    if base == ℯ; return ln(x); end

    f(x) = log(base, x)
    invF(x) = base^x

    MEAN = rowe(x, f)
    VAR = rowevar(x, f, invF)
    RANGE = f(x.range)

    return Moments(MEAN, VAR, RANGE)
end

function abs(x :: AbstractMoment)
    if x.range >= 0;
        return x
    elseif x.range < 0;
        EY = -x.mean;
    else
        sqvar = sqrt(x.var)
        EY1 = abs(x.mean)
        EY2 = EY1 + sqvar * (pi - atan(EY1/sqvar));
        EY = env(EY1, EY2)
    end

    Yvar = max(0, x.mean^2 + x.var - EY);
    Yrange = abs(x.range)

    return Moments(EY, Yvar, Yrange)

end

###
#   Scalar arithmetic
###
+(x :: AbstractMoment, y :: Real) = Moments(x.mean + y, x.var, x.range + y)
+(x :: Real, y:: AbstractMoment) = y + x;
-(x :: AbstractMoment, y :: Real) = x + (-y);
-(x :: Real, y :: AbstractMoment) = x + (-y);

*(x :: AbstractMoment, y :: Real) = Moments(x.mean * y, x.var * y^2, x.range * y)
*(x :: Real, y:: AbstractMoment) = y * x;
/(x :: AbstractMoment, y :: Real) = x * (1/y);
/(x :: Real, y :: AbstractMoment) = x * reciprocal(y);

###
#   Moment Arithmetic
###

###
# Independent
###
function sumIndep(x :: AbstractMoment, y :: AbstractMoment)

    zMean  = x.mean  + y.mean;
    zVar   = x.var   + y.var;
    zRange = x.range + y.range;

    # Add dependent variables to dict
    return Moments(zMean, zVar, zRange);
end

function subIndep(x :: AbstractMoment, y :: AbstractMoment)

    zMean  = x.mean  - y.mean;
    zVar   = x.var   + y.var;
    zRange = x.range - y.range;

    # Add dependent variables to dict
    return Moments(zMean, zVar, zRange);
end

function multIndep(x :: AbstractMoment, y :: AbstractMoment)

    zMean  = x.mean  * y.mean;
    zVar   = (x.mean^2 * y.var) + (y.mean^2 * x.var) + (x.var * y.var);     # No subint required
    zRange = x.range * y.range;

    # Add dependent variables to dict
    return Moments(zMean, zVar, zRange);

end

function divIndep(x :: AbstractMoment, y :: AbstractMoment)
    multIndep(x,1/y);
end

function minIndep(x :: AbstractMoment, y :: AbstractMoment)

    if x.range.hi < y.range.lo
        zMean = x.mean;
        zVar = x.var;
        zRange = min(x.range, y.range)
        return Moments(zMean, zVar, zRange)
    end

    if y.range.hi < x.range.lo
        zMean = y.mean;
        zVar = y.var;
        zRange = min(x.range, y.range)
        return Moments(zMean, zVar, zRange)
    end

    return minFrechet(x, y)
end

function maxIndep(x :: AbstractMoment, y :: AbstractMoment)

    if x.range.hi < y.range.lo
        zMean = y.mean;
        zVar = y.var;
        zRange = max(x.range, y.range)
        return Moments(zMean, zVar, zRange)
    end

    if y.range.hi < x.range.lo
        zMean = x.mean;
        zVar = x.var;
        zRange = max(x.range, y.range)
        return Moments(zMean, zVar, zRange)
    end

    return maxFrechet(x, y)
end

+(x :: AbstractMoment, y :: AbstractMoment) = sumFrechet(x,y);
-(x :: AbstractMoment, y :: AbstractMoment) = subFrechet(x,y);
*(x :: AbstractMoment, y :: AbstractMoment) = multFrechet(x,y);
/(x :: AbstractMoment, y :: AbstractMoment) = divFrechet(x,y);
min(x :: AbstractMoment, y :: AbstractMoment) = minFrechet(x,y);
max(x :: AbstractMoment, y :: AbstractMoment) = maxFrechet(x,y);

##
# With Intervals (Moments with just ranges)
##
+(x :: AbstractMoment, y :: Interval) = sumFrechet(x, Moments(missing, missing, y));
-(x :: AbstractMoment, y :: Interval) = subFrechet(x, Moments(missing, missing, y));
*(x :: AbstractMoment, y :: Interval) = multFrechet(x, Moments(missing, missing, y));
/(x :: AbstractMoment, y :: Interval) = divFrechet(x,  Moments(missing, missing, y));
min(x :: AbstractMoment, y :: Interval) = minFrechet(x,Moments(missing, missing, y));
max(x :: AbstractMoment, y :: Interval) = maxFrechet(x,Moments(missing, missing, y));

+(x :: Interval, y :: AbstractMoment) = sumFrechet(Moments(missing, missing, x), y);
-(x :: Interval, y :: AbstractMoment) = subFrechet(Moments(missing, missing, x), y);
*(x :: Interval, y :: AbstractMoment) = multFrechet(Moments(missing, missing, x),y);
/(x :: Interval, y :: AbstractMoment) = divFrechet(Moments(missing, missing, x), y);
min(x :: Interval, y :: AbstractMoment) = minFrechet(Moments(missing, missing, x),y);
max(x :: Interval, y :: AbstractMoment) = maxFrechet(Moments(missing, missing, x),y);

###
#   Dependence unknown
###
function sumFrechet(x :: AbstractMoment, y :: AbstractMoment)

    zMean = x.mean + y.mean

    zVarLb = (sqrt(x.var) - sqrt(y.var))^2
    zVarUb = (sqrt(x.var) + sqrt(y.var))^2

    zVarLb = max(zVarLb, 0)

    zVar = env(zVarLb, zVarUb);

    zRange = x.range + y.range;

    return Moments(zMean, zVar, zRange);
end

function subFrechet(x :: AbstractMoment, y :: AbstractMoment)

    zMean = x.mean - y.mean

    zVarLb = (sqrt(x.var) - sqrt(y.var))^2
    zVarUb = (sqrt(x.var) + sqrt(y.var))^2

    zVarLb = max(zVarLb, 0)

    zVar = env(zVarLb, zVarUb);

    zRange = x.range - y.range;

    return Moments(zMean, zVar, zRange);
end

###
#   Frechet multiplication. Not sub-intervalising required.
###
function multFrechet(x :: AbstractMoment, y :: AbstractMoment)

    EX = x.mean; EY = y.mean;
    VX = x.var; VY = y.var;

    zMeanLb = EX * EY - sqrt(VX * VY);
    zMeanUb = EX * EY + sqrt(VX * VY);
    zMean = env(zMeanLb, zMeanUb);

    x2 = x^2; y2 = y^2

    E11 = cov(x2,y2) + (x2.mean * y2.mean);
    E22 = (cov(x,y) + (x.mean * y.mean)) ^2

    zVar = E11 - E22

    zRange = x.range * y.range;

    Moments(zMean, zVar, zRange)

end

# Requires alot of subintervaling for Gvar to have an effect
function multFrechetOld(x :: AbstractMoment, y :: AbstractMoment; Nsub = Nsub[1])

    EX = x.mean; EY = y.mean;
    VX = x.var; VY = y.var;

    zMeanLb = EX * EY - sqrt(VX * VY);
    zMeanUb = EX * EY + sqrt(VX * VY);
    zMean = env(zMeanLb, zMeanUb);

    if typeof(EX) <: Interval && !isvacuous(EX);  EX = split(EX, Nsub); end    # Sub-intervalise
    if typeof(EY) <: Interval && !isvacuous(EY);  EY = split(EY, Nsub); end

    zVar = [Gvar(Moments(ex ,VX ,x.range), Moments(ey, VY, y.range)) for ex in EX, ey in EY]

    zVar = hull(zVar[:])

    zRange = x.range * y.range;

    return Moments(zMean, zVar, zRange);
end

divFrechet(x :: AbstractMoment, y :: AbstractMoment) = multFrechet(x, 1/y)  #Not best possible

function minFrechet(x :: AbstractMoment, y :: AbstractMoment)

    zRange = min(x.range, y.range)
    zVar = env(0, max(x.var, y.var))

    EX = x.mean; EY = y.mean;

    LX = left(x.range); GX = right(x.range)
    LY = left(y.range); GY = right(y.range)

    zMean = Bmin(EX, LX, GX, EY, LY, GY);
    return Moments(zMean, zVar, zRange)

end

function maxFrechet(x :: AbstractMoment, y :: AbstractMoment)

    zRange = max(x.range, y.range)
    zVar = env(0, max(x.var, y.var))

    EX = x.mean; EY = y.mean;

    LX = left(x.range); GX = right(x.range)
    LY = left(y.range); GY = right(y.range)

    zMean = Bmax(EX, LX, GX, EY, LY, GY);
    return Moments(zMean, zVar, zRange)

end

###
#   Perfect Arithmetic
###

function sumPerfect(x :: Moments, y :: Moments)

    meanZ = x.mean + y.mean
    varZ  = (sqrt(x.var) + sqrt(y.var))^2
    rangeZ = x.range + y.range
    return Moments(meanZ, varZ, rangeZ)
end

function subPerfect(x :: Moments, y :: Moments)

    meanZ = x.mean - y.mean
    varZ  = (sqrt(x.var) - sqrt(y.var))^2
    rangeZ = x.range - y.range
    return Moments(meanZ, varZ, rangeZ)
end

multPerfect(x, y) = multCor(x, y, 1)

###
#   Opposite Arithmetic
###

function sumOpposite(x :: Moments, y :: Moments)

    meanZ = x.mean + y.mean
    varZ  = (sqrt(x.var) - sqrt(y.var))^2
    rangeZ = x.range + y.range
    return Moments(meanZ, varZ, rangeZ)
end

function subOpposite(x :: Moments, y :: Moments)

    meanZ = x.mean - y.mean
    varZ  = (sqrt(x.var) + sqrt(y.var))^2
    rangeZ = x.range - y.range
    return Moments(meanZ, varZ, rangeZ)
end

###
#   known Dependence, or interval Dependence
###

function sumCor(x :: Moments, y :: Moments, corr :: Interval, nSub = 10)

    xVar = x.var; yVar = y.var; R = corr;

    if typeof(x.var) <: Interval; xVar = mince(xVar, nSub); end
    if typeof(y.var) <: Interval; yVar = mince(yVar, nSub); end
    #if typeof(corr) <: Interval; R = mince(R, nSub); end

    varZ = [xs + ys + 2 * sqrt(xs) * sqrt(ys) * corrs for xs in xVar, ys in yVar, corrs in R]

    meanZ = x.mean + y.mean
    rangeZ = x.range + y.range

    varZ = hull(varZ[:]);

    return Moments(meanZ, varZ, rangeZ)
end

sumCor(x :: Moments, y :: Moments, corr :: Real, nSub = 10) = sumCor(x,y, interval(corr), nSub)

function subCor(x :: Moments, y :: Moments, corr :: Interval, nSub = 10)

    xVar = x.var; yVar = y.var; R = corr;

    if typeof(x.var) <: Interval; xVar = mince(xVar, nSub); end
    if typeof(y.var) <: Interval; yVar = mince(yVar, nSub); end
    #if typeof(corr) <: Interval; R = mince(R, nSub); end

    varZ = [xs + ys - 2 * sqrt(xs) * sqrt(ys) * corrs for xs in xVar, ys in yVar, corrs in R]

    meanZ = x.mean - y.mean
    rangeZ = x.range - y.range

    varZ = hull(varZ[:]);

    return Moments(meanZ, varZ, rangeZ)
end

subCor(x :: Moments, y :: Moments, corr :: Real, nSub = 10) = subCor(x,y, interval(corr), nSub)

function multCor(x :: Moments, y :: Moments, corr :: Interval)

    EX = x.mean; EY = y.mean;
    VX = x.var; VY = y.var;

    zMean = EX * EY + corr * sqrt(VX * VY);

    x2 = x^2; y2 = y^2

    E11 = cov(x2,y2) + (x2.mean * y2.mean);
    E22 = (corr * sqrt(VX * VY) + (x.mean * y.mean)) ^2

    zVar = E11 - E22

    zRange = x.range * y.range;

    return Moments(zMean, zVar, zRange)
end

multCor(x :: Moments, y :: Moments, corr :: Real) = multCor(x,y, interval(corr))

function divCor(x :: Moments, y :: Moments, corr :: Interval)
    return multCor(x, 1/y, -1 * corr)
end

divCor(x :: Moments, y :: Moments, corr :: Real) = divCor(x,y, interval(corr))

###
#   Using Covariances
###

function sumCov(x :: Moments, y :: Moments, cov :: Interval)

    xVar = x.var; yVar = y.var;

    varZ = xVar + yVar + 2 * cov;

    meanZ = x.mean + y.mean
    rangeZ = x.range + y.range

    return Moments(meanZ, varZ, rangeZ)
end

sumCov(x :: Moments, y :: Moments, cov :: Real) = sumCov(x,y, interval(cov))

function subCov(x :: Moments, y :: Moments, cov :: Interval)

    xVar = x.var; yVar = y.var;

    meanZ = x.mean - y.mean

    varZ = xVar + yVar - 2 * cov;

    rangeZ = x.range - y.range

    return Moments(meanZ, varZ, rangeZ)
end

subCov(x :: Moments, y :: Moments, cov :: Real) = subCov(x,y, interval(cov))

function multCov(x :: Moments, y :: Moments, Cov :: Interval)

    EX = x.mean; EY = y.mean;

    zMean = EX * EY + Cov;

    x2 = x^2; y2 = y^2

    E11 = cov(x2,y2) + (x2.mean * y2.mean);
    E22 = (Cov + (x.mean * y.mean))^2

    zVar = E11 - E22

    zRange = x.range * y.range;

    return Moments(zMean, zVar, zRange)
end

multCov(x :: Moments, y :: Moments, cov :: Real) = multCov(x,y, interval(cov))

###
#   Bertsimas mean for min and max. Do not require sub-intervalisation due to env and intersect
###

function Bmin(EX, LX, GX, EY, LY, GY)

    B1 = env(min(EX, EY), min(LX, LY))
    B2 = EX + EY - env(max(EX, EY), max(GX, GY))

    return intersect(B1,B2)
end

function Bmax(EX, LX, GX, EY, LY, GY)

    B1 = env(max(EX, EY), max(GX, GY))
    B2 = EX + EY - env(min(EX, EY), min(LX, LY))

    return intersect(B1,B2)
end

###
#   Goodman formula for variance under unknown dependence
###

#   Looks like a complidated repeat variables problem, but only EX and EY need sub-ints
function Gvar(x :: AbstractMoment, y :: AbstractMoment)

    EX = x.mean; VX = y.var
    EY = y.mean; VY = y.var

    EXY = multFrechetMean(x, y)
    EX2Y = multFrechetMean(x^2, y)
    EXY2 = multFrechetMean(x, y^2)
    EX2Y2 = multFrechetMean(x^2, y^2)

    EX2 = (x^2).mean
    EY2 = (y^2).mean

    E11 = EXY - EX*EY
    E12 = EX2Y - EX2*EY + 2*EX^2*EY - 2*EX*EXY
    E21 = EXY2 - EX*EY2 + 2*EX*EY^2 - 2*EY*EXY
    E22 = -3*EX^2*EY^2 + EX2*EY^2 + EX^2*EY2 + 4*EX*EY*EXY - 2*EY*EX2Y - 2*EX*EXY2 + EX2Y2

    VXY = EX^2*VX + EY^2*VY + 2*EX*EY*E11 + 2*EX*E12 + 2*EY*E21 + E22 - E11^2

    #VXY2 = EX^2*VX + EY^2*VY + 2*EX*EX2Y - 2*EY*EX2Y + 4*EX*EY^3 - 6*EX^2*EY^2 + 4*EX^3*EY - EXY^2 + 8*EX*EY*EXY - 4*EX^2*EXY - 4*EY^2*EXY -2*EX*EY*EX2 + EY^2*EX2 - 2*EX*EY*EY2 + EX^2*EY2 + 2*EY*EXY2 - 2*EX*EXY2 + EX2Y2

    return max(VXY, 0)

end

#=

    E11 = EXY - EX*EY
    E12 = EX2Y - EX2*EY + 2*EX^2*EY - 2*EX*EXY
    E21 = EXY2 - EX*EY2 + 2*EX*EY^2 - 2*EY*EXY
    E22 = -3*EX^2*EY^2 + EX2*EY^2 + EX^2*EY2 + 4*EX*EY*EXY - 2*EY*EX2Y - 2*EX*EXY2 + EX2Y2

    VXY = EX^2*VX + EY^2*VY + 2*EX*EY*E11 + 2*EX*E12 + 2*EY*E21 + E22 - E11^2

    A = EX; B = EY;
    C = EXY; D = EX2;
    E = EY2; F = EX2Y;
    G = EXY2; H = EX2Y2

    2 * A * B * (C - A*B) + 2 * A * (F -D*B +2 * A^2 * B - 2 * A * C) + 2 * B * (G - A * E + 2 * A * B^2 - 2 * B * C) + (-3 * A^2 * B^2 + D * B^2 + A^2 * E + 4*A*B*C -2*B*F - 2*A*G + H) - (C - A*B)^2

    simplify gives:

    2AF-2BF+4AB^3-6A^2B^2+4A^3B-C^2+8ABC-4A^2C-4B^2C-2ABD+B^2D-2ABE+A^2E+2BG-2AG+H

    2*EX*EX2Y - 2*EY*EX2Y + 4*EX*EY^3 - 6*EX^2*EY^2 + 4*EX^3*EY - EXY^2 + 8*EX*EY*EXY - 4*EX^2*EXY - 4*EY^2*EXY -2*EX*EY*EX2 + EY^2*EX2 - 2*EX*EY*EY2 + EX^2*EY2 + 2*EY*EXY2 - 2*EX*EXY2 + EX2Y2
=#

###
#   Woodpile
###

#=
###
#   Bertsimas with subintervalistion (does not require)
###
function minFrechetSub(x :: Moments, y ::Moments)

    zRange = min(x.range, y.range)
    zVar = env(0, max(x.var, y.var))

    EX = x.mean; EY = y.mean;

    LX = left(x.range); GX = right(x.range)
    LY = left(y.range); GY = right(y.range)

    if !hasIntMean(x) && !hasIntMean(y);
        zMean = Bmin(EX, LX, GX, EY, LY, GY);
        return Moments(zMean, zVar, zRange)
    end

    if typeof(EX) <: Interval;  EX = Split(EX, Nsub); end
    if typeof(EY) <: Interval;  EY = Split(EY, Nsub); end

    zMeans = [Bmin(EX, LX, GX, EY, LY, GY) for EX in EX, EY in EY];

    zMean = hull(zMeans[:])

    return Moments(zMean, zVar, zRange)

end

F(xmean, ymean, xvar, yvar) = (xmean^2 * yvar) + (ymean^2 * xvar) + (xvar * yvar)

xmean = interval(0,100)
xvar = interval(0,3)
ymean = interval(0,100)
yvar = interval(0,3)

xs = Split(xvar, 10)
ys = Split(yvar, 10)

outs  = [F(xmean, ymean, xvar, yvar) for xvar in xs, yvar in ys]
hull(outs[:])

function multFrechet(x :: AbstractMoment, y :: AbstractMoment; sub = false, Nsub = Nsub[1])

    EX = x.mean; EY = y.mean;
    VX = x.var; VY = y.var;

    zMeanLb = EX * EY - sqrt(VX * VY);
    zMeanUb = EX * EY + sqrt(VX * VY);
    zMean = env(zMeanLb, zMeanUb);

    if sub
        if typeof(VX) <: Interval && !isvacuous(VX);  VX = split(VX, Nsub); end    # Sub-intervalise
        if typeof(VY) <: Interval && !isvacuous(VY);  VY = split(VY, Nsub); end
    end

    COR = interval(-1, 1)           # cor = [-1, 1] for Frechet

    zVar = [(VX + EX^2) * (VY + EY^2) + COR*VX*VY - (COR * sqrt(VX) * sqrt(VY) + EY * EX)^2 for VX in VX, VY in VY]
    zVar = hull(zVar[:])

    zRange = x.range * y.range;

    Moments(zMean, zVar, zRange)

end

=#


Nemo.fit!(::fmpq_rel_series, Int) = nothing
Nemo.fit!(::fmpq_abs_series, Int) = nothing

@doc Markdown.doc"""
    integral(f::RelSeriesElem{T}) -> RelSeriesElem

Return the integral of the power series $f$.
"""
function Nemo.integral(f::RelSeriesElem{T}) where T
  g = parent(f)()
  fit!(g, precision(f)+1)
  set_precision!(g, precision(f)+1)
  v = valuation(f)
  Nemo.set_valuation!(g, v+1)
  for i=0:Nemo.pol_length(f)
    c = Nemo.polcoeff(f, i)
    if !iszero(c)
      setcoeff!(g, i, divexact(c, i+v+1))
    end
  end
  Nemo.renormalize!(g)
  return g
end


abstract type BOOL end
struct TRUE <: BOOL end
struct FALSE <: BOOL end

(t::TRUE)(x::BOOL, y::BOOL) = x
(f::FALSE)(x::BOOL, y::BOOL) = y

show(io::IO, ::TRUE) = print(io, "true")
show(io::IO, ::FALSE) = print(io, "false")


# HINT network from Kruse et al. (2020)
# Author: Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
# Date: August 2020

using DrWatson
@quickactivate "UncertaintyQuantificationAVP"
using UncertaintyQuantificationAVP
using InvertibleNetworks
using LinearAlgebra, Test, Random
Random.seed!(11)

# Define network
nx = 1
ny = 1
n_in = 2^3
n_hidden = 64
batchsize = 16
L = 20

# Multi-scale and single scale network
H = NetworkGlow1D(n_in, batchsize, n_hidden, L; logdet=true)

###################################################################################################
# Invertibility

# Test layers
test_size = 10
X = randn(Float32, nx, ny, n_in, test_size)

# Forward-backward
Z, logdet = H.forward(X)
_, X_ = H.backward(0f0.*Z, Z)
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-3)

# Forward-inverse
Z, logdet = H.forward(X)
X_ = H.inverse(Z)
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-3)

###################################################################################################
# Gradient test

# Loss
function loss(H, X)
    Z, logdet = H.forward(X)
    f = -log_likelihood(Z) - logdet
    ΔZ = -∇log_likelihood(Z)
    ΔX = H.backward(ΔZ, Z)[1]
    return f, ΔX
end

# Gradient test w.r.t. input
H = NetworkGlow1D(n_in, batchsize, n_hidden, L; logdet=true)
X = randn(Float32, nx, ny, n_in, batchsize)
X0 = randn(Float32, nx, ny, n_in, batchsize)
dX = X - X0

f0, ΔX = loss(H, X0)
h = 0.1f0
maxiter = 6
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test cond. HINT net: input\n")
for j=1:maxiter
    global h
    f = loss(H, X0 + h*dX)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test w.r.t. parameters
H = NetworkGlow1D(n_in, batchsize, n_hidden, L; logdet=true)
X = randn(Float32, nx, ny, n_in, batchsize); H.forward(X)
H0 = NetworkGlow1D(n_in, batchsize, n_hidden, L; logdet=true)
X0 = randn(Float32, nx, ny, n_in, batchsize); H0.forward(X0)
θ0 = get_pars(H0)
θ = get_pars(H)
dθ = θ .- θ0

f0, _ = loss(H0, X0); Δθ = deepcopy(get_grads(H0))
h = 0.1f0
maxiter = 6
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test cond. Glow net: pars\n")
for j=1:maxiter
    global h
    set_pars!(H0, θ0+h*dθ)
    f = loss(H0, X0)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dθ, Δθ))
    print(err1[j], "; ", err2[j], "\n")
    h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)


# Inspired by https://mtk.sciml.ai/dev/tutorials/acausal_components/

# a point in the circuit with a given voltage and current
@connector function Pin(; name)
    sts = @variables v(t)=1.0 i(t)=1.0
    ODESystem(Equation[], t, sts, []; name)
end

# ground pin, where voltage is 0
function Ground(; name)
    @named g = Pin()
    eqs = [g.v ~ 0.]
    compose(ODESystem(eqs, t, [], []; name), g)
end

# abstraction for all 2-pin components
function OnePort(; name)
    # two pins
    @named p = Pin()
    @named n = Pin()
    # component has a voltage across it, and current through
    sts = @variables v(t)=1.0 i(t)=1.0
    eqs = [
        v ~ p.v - n.v    # KVL
        0. ~ p.i + n.i   # KCL
        i ~ p.i          # Current through component is current through +ve pin
    ]
    compose(ODESystem(eqs, t, sts, []; name), p, n)
end

function ModelingToolkit.connect(::Type{Pin}, ps...)
    eqs = [
        0. ~ sum(p->p.i, ps)    # KCL
    ]
    # KVL
    for i in 1:length(ps)-1
        push!(eqs, ps[i].v ~ ps[i+1].v)
    end
    return eqs
end

connect_star(comps...) = connect([comp.n for comp in comps]...)

function connect_series(comps...)
    eqs = Equation[]
    for i in 1:length(comps)-1
        eqs = vcat(eqs, connect(comps[i].p, comps[i+1].n))
    end
    return eqs
end

function connect_parallel(comps...)
    return [
        connect((comp.p for comp in comps)...)
        connect((comp.n for comp in comps)...)
    ]
end


using Checkpointing
using ReverseDiff
using Zygote
using ForwardDiff
using Enzyme
using Test

function f(x)
    return [(x[1]*(x[2]-2.0))^2, (x[1]-1.0)^2*x[2]^2]
end
x = [2.0,6.0]

J_Re = Checkpointing.jacobian(f, x, ReverseDiffADTool())
J_Zy = Checkpointing.jacobian(f, x, ZygoteADTool())
J_Fo = Checkpointing.jacobian(f, x, ForwardDiffADTool())
J_En = Checkpointing.jacobian(f, x, EnzymeADTool())
# J_Di = Checkpointing.jacobian(f, x, DiffractorADTool())

@test J_Re ≈ J_Zy ≈ J_Fo ≈ J_En


 #Duffing Oscillator
 #Controller: RFPT (SISO)
 #using FPTControl toolbox

using FPTControl

Adaptive=1
LONG=Int(1e4)
δt=1e-3

#Kinematic Parameter
Λ=3

#Nominal Trajectory
ω=5
Ampl=0.4

#Controller Parameters
K=1e5
B=-1
A=1e-5

# Exact model parameters
αₑ=1
δₑ=0.2
βₑ=1

# Approx model parameters
αₐ=0.8
δₐ=0.1
βₐ=0.9

###############
# Exact Model #
###############

function exact(q,q_p,u)
  out=αₑ*q+δₑ*q_p+βₑ*q^3+u
  return out #q_pp
end

################
# Approx Model #
################

function approx(q,q_p,q_pp)
  out=q_pp#-αₐ*q-δₐ*q_p-βₐ*q^3
  return out #u
end

######################################
# Arrays for Plotting and simulation #
######################################

qN=zeros(LONG)
qN_p=zeros(LONG)
qN_pp=zeros(LONG)

qDes_pp=zeros(LONG)

qDef_pp=zeros(LONG)
u=zeros(LONG)

q=zeros(LONG)
q_p=zeros(LONG)
q_pp=zeros(LONG)

t=zeros(LONG)

error_int=0
past_input=0
past_response=0

for i=1:LONG-1
  t[i]=i*δt
  #Nominal trajectory
  qN[i]=Ampl*sin(ω*t[i])#sin(ω*time_mem[t])
  qN_p[i]=ω*Ampl*cos(ω*t[i])
  qN_pp[i]=-ω^2*Ampl*sin(ω*t[i])
  #Kinematic Block

  qDes_pp[i]=Λ^3*error_int+3*Λ^2*(qN[i]-q[i])+3*Λ*(qN_p[i]-q_p[i])+qN_pp[i]
  qDes_pp[i]
  # Deformation
  if Adaptive==1 && i>10
    qDef_pp[i]=G_SISO(past_input,past_response,qDes_pp[i],K,B,A)
  #  println(qDef_pp[i])
  else
    qDef_pp[i]=qDes_pp[i]
  end
  #Approx model
  u[i]=approx(q[i],q_p[i],qDef_pp[i])
  #u=qDef_pp_mem[t]
  past_input=qDef_pp[i]
  #Exact model
  q_pp[i]=exact(q[i],q_p[i],u[i])
  past_response=q_pp[i]
  #Integrals
  #q_p[i+1]=q_p[i]+δt*q_pp[i]
   q_p[i+1]=Integ(q_p[i],δt,q_pp[i])
  #q[i+1]=q[i]+δt*q_p[i]
  q[i+1]=Integ(q[i],δt,q_p[i])
  error_int=Integ(error_int,δt,qN[i]-q[i])
  #error_int=error_int+δt*(qN[i]-q[i])
end #for (simulation)

using PyPlot

figure(1)
title("Tracking error")
plot(t[3:LONG-1],qN[3:LONG-1]-q[3:LONG-1])
grid("on")
xlabel("Time [s]")
ylabel("error [m]")
#xlabel(L"\frac{x}{y}")   ##(L"") latex-es felirathoz

figure(2)
title("Nominal and Simulated")
plot(t[3:LONG-1],qN[3:LONG-1],color="red")
plot(t[3:LONG-1],q[3:LONG-1],color="green","r--")
grid("on")
xlabel("Time [s]")
ylabel("q [m]")

examples/RFPT_SISO/Duffing Oscillator with FPTControl Toolbox/duffing_FPTControl.jl


import ThorinDistributions
using Test

@testset "UnivariateGammaConvolution.jl tests" begin

    dist = ThorinDistributions.UnivariateGammaConvolution([1,0.5, 3.7],[4,2, 10])

    @testset "functions are working and consistant with past selves." begin
        spl = [
            44.21448353310021,
            40.67778979017133,
            38.67640776714249,
            23.525094134635108,
            23.38644058822338,
            29.498656708345194,
            35.340080215403574,
            56.380727733998995,
            25.52972573138749,
            51.20505953884155
            ]
            x = pdf.(dist,spl)
            x_fixed = [
            0.018757112919480665,
            0.020559416575509366,
            0.02140738452673978,
            0.019147239099563968,
            0.019031193841015707,
            0.02236341049555192,
            0.02240803708392265,
            0.011765753186399592,
            0.020617044066218452,
            0.014700161862923838
            ]

            y = cdf.(dist,spl)
            y_fixed = [
            0.6082512562725385,
            0.5386338347308577,
            0.49661228006226665,
            0.16557008552908928,
            0.16292327691403613,
            0.29129135121181804,
            0.4233565994115752,
            0.7937252726673726,
            0.20549205556263866,
            0.725339366143653
            ]

        @test x ≈ x_fixed
        @test y ≈ y_fixed
    end

    @testset "results of pdf and cdf are consistant with R coga" begin
        # We check against COGA from R on this simple example.
        # > library(coga)
        # > options(digits=20)
        # > coga::dcoga(1:10,c(1,0.5,3.7),1/c(4,2,10))
        #  [1] 9.1697877166357581e-07 1.4306872316747290e-05 6.6823198132184901e-05
        #  [4] 1.9070137722253317e-04 4.1580841085495259e-04 7.6517030290220020e-04
        #  [7] 1.2530854177965063e-03 1.8849283018035432e-03 2.6579560778065099e-03
        # [10] 3.5626340599333940e-03
        # > coga::pcoga(1:10,c(1,0.5,3.7),1/c(4,2,10))
        #  [1] 1.8115256558630823e-07 5.8062605413247208e-06 4.1780076813316473e-05
        #  [4] 1.6327100296424765e-04 4.5700519074037604e-04 1.0364052739079920e-03
        #  [7] 2.0336306415139858e-03 3.5906428800766086e-03 5.8506245787791631e-03
        # [10] 8.9504969486005927e-03
        # >

        values = Float64.(1:10)
        p = pdf.(dist,values)
        c = cdf.(dist,values)
        pᵣ = [
        9.1697877166357581e-07,
        1.4306872316747290e-05,
        6.6823198132184901e-05,
        1.9070137722253317e-04,
        4.1580841085495259e-04,
        7.6517030290220020e-04,
        1.2530854177965063e-03,
        1.8849283018035432e-03,
        2.6579560778065099e-03,
        3.5626340599333940e-03
        ]
        cᵣ = [
        1.8115256558630823e-07,
        5.8062605413247208e-06,
        4.1780076813316473e-05,
        1.6327100296424765e-04,
        4.5700519074037604e-04,
        1.0364052739079920e-03,
        2.0336306415139858e-03,
        3.5906428800766086e-03,
        5.8506245787791631e-03,
        8.9504969486005927e-03
        ]
        @test round.(p,digits=10) ≈ round.(pᵣ,digits=10)
        @test round.(c,digits=10) ≈ round.(cᵣ,digits=10)
    end

end


using LinearAlgebra, Tullio, Test, ForwardDiff, Random, Zygote, BenchmarkTools
using LoopVectorization, TensorOperations, TensorCast, CUDA
using Zygote: @adjoint
# 2-tensors
# from Michael Abbott
# donate to ChainRules
@adjoint function dot(x, A, y)
    Ay = A * y
    z = adjoint(x) * Ay
    return z, dz -> begin
               dx = lmul!(dz, conj!(Ay))
               dA = dz .* x .* adjoint(y)
               dy = lmul!(dz, adjoint(A) * x)
               (dx, dA, dy)
           end
end

quad_form_raw(x, A::AbstractMatrix) = x' * A * x
quad_form_ten(x, A::AbstractArray{<:Number,2}) = @tullio c := x[j] * A[j, k] * x[k]
quad_form_dot(x, A::AbstractMatrix) = dot(x, A, x)

# 3-tensors
quad_form_raw(x, A::AbstractArray{<:Number,3}) = [x' * A[i, :, :] * x for i in 1:size(A, 1)]
function quad_form_mul(x, A::AbstractArray{<:Number,3})
    @matmul c[l] := sum(j) (@matmul [l, j] := sum(k) A[l, j, k] * x[k]) * x[j]
end
function quad_form_avx(x, A::AbstractArray{<:Number,3})
    @tullio c[l] := x[j] * A[l, j, k] * x[k] tensor = false
end;
quad_form_ten(x, A::AbstractArray{<:Number,3}) = @tullio c[l] := x[j] * A[l, j, k] * x[k];
function quad_form_base(x, A::AbstractArray{<:Number,3})
    @tullio c[l] := x[j] * A[l, j, k] * x[k] tensor = false avx = false
end;

#tests
scalar_objective(f, x, A) = norm(f(x, A))

function print_benchmarks(N, use_cuda = true)
    A_2 = rand(N, N)
    A_3 = rand(N, N, N)
    x = rand(N)
    # 2-tensors
    printstyled("Raw Julia 2-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_raw, $x, $A_2)
    @btime gradient(scalar_objective, quad_form_raw, $x, $A_2)

    printstyled("Dot Julia 2-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_dot, $x, $A_2)
    @btime gradient(scalar_objective, quad_form_dot, $x, $A_2)

    printstyled("Tullio with Tensor Julia 2-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_ten, $x, $A_2)
    @btime gradient(scalar_objective, quad_form_ten, $x, $A_2)

    # three tensor
    printstyled("Raw Julia 3-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_raw, $x, $A_3)
    @btime gradient(scalar_objective, quad_form_raw, $x, $A_3)

    printstyled("TensorCast Julia 3-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_mul, $x, $A_3)
    @btime gradient(scalar_objective, quad_form_mul, $x, $A_3)

    printstyled("AVX Julia 3-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_avx, $x, $A_3)
    @btime gradient(scalar_objective, quad_form_avx, $x, $A_3)

    printstyled("Tullio with Tensor Julia 3-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_ten, $x, $A_3)
    @btime gradient(scalar_objective, quad_form_ten, $x, $A_3)

    printstyled("Tullio Base Julia 3-tensor for N = $N\n"; color = :red)
    @btime scalar_objective(quad_form_base, $x, $A_3)
    @btime gradient(scalar_objective, quad_form_base, $x, $A_3)

    if use_cuda
        A_3 = CuArray(A_3)
        x = CuArray(x)

        printstyled("CUDA TensorCast Julia 3-tensor for N = $N\n"; color = :red)
        @btime CUDA.@sync blocking = false scalar_objective(quad_form_mul, $x, $A_3)
        @btime CUDA.@sync blocking = false gradient(scalar_objective, quad_form_mul, $x,
                                                    $A_3)

        # Not sure why it is not working?
        # printstyled("CUDA Tullio  Julia 3-tensor for N = $N\n", color=:red)
        # @btime CUDA.@sync blocking=false scalar_objective(quad_form_ten, $x, $A_3)
        # @btime CUDA.@sync blocking=false gradient(scalar_objective, quad_form_ten, $x, $A_3)
    end
end

Zygote.refresh()
# try for different sizes
print_benchmarks(100, false)
#print_benchmarks(100)
#print_benchmarks(200)

# Quadratic forms for larger ones as well
A = rand(100, 100)
A_transpose = Array(A')
B = rand(100, 100)
C = zeros(100, 100)
C_vec = vec(copy(C))
temp = zeros(100, 100)

function quadform_1!(C, A, B)
    return C .= A' * B * A
end

function quadform_2!(C, temp, A, B)
    return mul!(C, mul!(temp, A', B), A)
end

@benchmark quadform_1!($C, $A, $B)
@benchmark quadform_2!($C, $temp, $A, $A_transpose, $B)


# =========================================================================== #
# ErrorEstimationOptions
# =========================================================================== #

struct ErrorEstimationOptions{c}
    norm_type::Val{c}

    function ErrorEstimationOptions{c}() where {c}
        norm_type = Val(c)
        return new(norm_type)
    end
end

function ErrorEstimationOptions(;
    norm_type = 2)

    return ErrorEstimationOptions{norm_type}()
end

function ErrorEstimationOptions(norm_type = 2)
    return ErrorEstimationOptions{norm_type}()
end

mutable struct ErrorEstimation
    current::Float64
    last::Float64
end

function ErrorEstimation(current)
    return ErrorEstimation(current, current)
end

function ErrorEstimation()
    return ErrorEstimation(0.0)
end

function update!(err_est::ErrorEstimation)
    err_est.last = err_est.current
    return err_est
end

function is_step_acceptable(err_est::ErrorEstimation)
    return is_step_acceptable(err_est.current)
end

function is_step_acceptable(err_est::T) where {T <: Number}
    return err_est <= 1.0;
end

function error_estimation(
    x_new,
    x_old,
    x_est,
    basic_options::BasicOdeOptions,
    error_estimation_options::ErrorEstimationOptions{2})

    n_dim = length(x_old)
    err_est = 0.0
    rfac    = sqrt( 1/n_dim * max( dot(x_new, x_new), dot(x_old, x_old) ) )
    for kk = 1:length(x_new)
        atol     = basic_options.abs_tol[kk]
        rtol     = basic_options.rel_tol[kk] * rfac
        scaling  = max(atol, rtol)
        tmp      = (x_new[kk] - x_est[kk]) / scaling
        err_est += tmp^2
    end
    err_est = sqrt(err_est)

    return err_est
end

function error_estimation(
    x_new,
    x_old,
    x_est,
    basic_options::BasicOdeOptions,
    error_estimation_options::ErrorEstimationOptions{Inf})

    err_est = 0.0
    for kk = 1:length(x_new)
        rfac     = max(abs(x_new[kk]), abs(x_old[kk]))
        atol     = basic_options.abs_tol[kk]
        rtol     = basic_options.rel_tol[kk] * rfac
        scaling  = max(atol, rtol)
        tmp      = abs(x_new[kk] - x_est[kk]) / scaling

        if err_est < tmp
            err_est = tmp
        end
    end

    return err_est
end


function xxz_hamiltonian(lx::Int64, delta::Float64=1.0;
                         boundary::Symbol=:OBC, mode::Symbol=:FULL)
    @assert lx > 1
    if mode == :FULL
        return xxz_explicit(lx, delta, boundary)
    elseif mode == :U1
        return xxz_u1sym(lx, delta, boundary=boundary)
    else
        error("unrecognized mode :", mode)
    end
end

function xxz_explicit(lx::Int64, delta::Float64, boundary::Symbol)

    sz, sp, sm = sparse.(spinoperators(1/2))
    heis_term = 0.5 * (kron(sp, sm) + kron(sm, sp)) +
        delta * kron(sz,sz)

    Hmat = heis_term
    if lx > 2
        # recursive generation of open chain
        for i=3:lx
            Hmat = kron(Hmat, I(2)) + kron(I(2^(i-2)), heis_term)
        end
    end

    if boundary == :OBC
        return Hmat
    elseif boundary == :PBC
        Hmat = Hmat +
            0.5 * kron(kron(sp, I(2^(lx-2))), sm) +
            0.5 * kron(kron(sm, I(2^(lx-2))), sp) +
            delta * kron(kron(sz, I(2^(lx-2))), sz)
        return Hmat
    else
        error("unrecognized boundary condition :", boundary)
    end
end

"""
    xxz_u1sym(Lx, delta; boundary, zsector)
"""

# # Generates the XXZ Hamiltonian of size `lx` in the given `zsector`
# # block. This is possible because ``\left[ S^z_{\text{tot}},
# # H\right]=0`` and its eigenvalues can be used to label separate
# # magnetization blocks in the Hamiltonian matrix and each block can be
# # diagonalized separately to significantly reduce time complexity. Each
# # `zsector` block consists of Ising vectors with exactly `M = (zsector
# # + lx)/2` spins up, so the size of the block is ``\binom{lx}{M}``. Using
# # Sterling formula the size of the largest block, ``M=lx/2``, is ``√{π
# # lx/2}`` times smaller than the full Hilbert space, which For typical
# # sizes of ED calculation ,``lx∼ 30``, is roughly about ``∼\!7`` times.

# """
function xxz_u1sym(lx::Int, delta::Float64=1.0;
                   boundary::Symbol=:OBC, zsector::Int=(lx % 2))

    @assert lx > 1
    @assert (lx-zsector) % 2 == 0

    if boundary == :OBC
        bond_range = 0:lx-2
    elseif boundary == :PBC
        bond_range = 0:lx-1
    else
        error("unrecognized boundary condition :", boundary)
    end

    M = div(lx + zsector, 2)
    block_size = binomial(lx, M)

    I = Int32[]
    J = Int32[]
    V = Float64[]

    szblock_states = Vector{Int}(undef, block_size)
    index = 0
    for state=0:2^lx-1
        if count_ones(state) == M
            index += 1
            szblock_states[index] = state
        end
    end

    # For all states (of magnetization M):
    for state_index=1:block_size
        state = szblock_states[state_index]
        # For all bond terms of the Hamiltonian
        for site=bond_range
            site_nn = (site + 1) % lx
            # if spin are parallel:
            if xor((state >> site) & 1, (state >> site_nn) & 1) == 0
                push!(I, state_index)
                push!(J, state_index)
                push!(V, delta * 0.25)
                # if spin are not parallel:
            else
                push!(I, state_index)
                push!(J, state_index)
                push!(V, - delta * 0.25)

                flipped_state = xor(state, 1 << site | 1 << site_nn)
                flipped_state_index = searchsortedfirst(szblock_states, flipped_state)
                push!(I, state_index)
                push!(J, flipped_state_index)
                push!(V, 0.5)
            end
        end
    end

    return sparse(I, J, V, block_size, block_size, +), (szblock_states .+ 1)
 end

"""
    j1j2_explicit(Lx, j1, j2)

explicit construction of j1j2 sparse full Hamiltonian.

"""
function j1j2_explicit(Lx::Int64,
                       j1::Float64,
                       j2::Float64;
                       boundary::Symbol=:open)
    @assert Lx > 2
    Sz, Sp, Sm = sparse.(spinoperators(1/2))
    I2 = I(2)

    heis_term1 = j1 * 0.5 * (kron(Sp, Sm) + kron(Sm, Sp)) + j1 * kron(Sz,Sz)
    heis_term2 = j2 * 0.5 * (kron(kron(Sp, I2), Sm) + kron(kron(Sm, I2), Sp)) +
        j2 * kron(kron(Sz, I2), Sz)

    Hmat = heis_term1
    for i=3:Lx
        Hmat = kron(Hmat, I2) + kron(I(2^(i-2)), heis_term1) +
            kron(I(2^(i-3)), heis_term2)
    end

    if boundary == :open
        return Hmat
    elseif boundary == :periodic
        Hmat = Hmat +
            j1 * 0.5 * kron(kron(Sp,I(2^(Lx-2))), Sm) +
            j1 * 0.5 * kron(kron(Sm,I(2^(Lx-2))), Sp) +
            j1 * kron(kron(Sz,I(2^(Lx-2))), Sz) +
            j2 * 0.5 * kron(I2, kron(kron(Sp,I(2^(Lx-3))), Sm)) +
            j2 * 0.5 * kron(I2, kron(kron(Sm,I(2^(Lx-3))), Sp)) +
            j2 * kron(I2, kron(kron(Sz,I(2^(Lx-3))), Sz)) +
            j2 * 0.5 * kron(kron(kron(Sp,I(2^(Lx-3))), Sm), I2) +
            j2 * 0.5 * kron(kron(kron(Sm,I(2^(Lx-3))), Sp), I2) +
            j2 * kron(kron(kron(Sz,I(2^(Lx-3))), Sz), I2)
        return Hmat
    else
        error("unrecognized boundary conditions :", boundary)
    end
end


q = Quaternion(1, 0, 0, 0)
r = λmap(q)
p = compressedλmap(q)

@test isapprox(r, ℝ³(1, 0, 0))
@test isapprox(p, ℝ³(0.7615941559557649, 0, 0))


# Sum-Product Networks as described by Poon&Domingos, Gens&Domingos etc.

abstract SPNode{T} # SPN node with names of type T

immutable Sum{T} <: SPNode{T}
  children::Vector{SPNode{T}}
  weights::Vector{Float64}
end

immutable Prod{T} <: SPNode{T}
  children::Vector{SPNode{T}}
end

# Variables have names of a given type.
immutable Var{T} <: SPNode{T}
  name::T
end

immutable NVar{T} <: SPNode{T}
  name::T
  negname::T
end
negname(name::Symbol) = symbol("__N__" * string(name))
negname{I<:Integer}(name::I) = -name
NVar(name) = NVar{typeof(name)}(name, negname(name))

# A SPN with variable names of type T.
type SumProductNetwork{T}
  root::SPNode{T}
end
SPN{T}(r::SPNode{T}) = SumProductNetwork{T}(r)

# Direct calculation of probabilities
calc{T, N <: Number}(spn::SumProductNetwork{T}, vals::Dict{T,N}) = calc(spn.root, vals)
function calc{T, N <: Number}(s::Sum{T}, vals::Dict{T,N})
  sum(calcchildren(s.children, vals) .* s.weights)
end
calcchildren{T, N <: Number}(children::Vector{SPNode{T}}, vals::Dict{T,N}) = map(c->calc(c, vals), children)
calc{T, N <: Number}(s::Prod{T}, vals::Dict{T,N}) = prod(calcchildren(s.children, vals))
calc{T, N <: Number}(v::Var{T}, vals::Dict{T,N}) = haskey(vals, v.name) ? vals[v.name] : 1
function calc{T, N <: Number}(v::NVar{T}, vals::Dict{T,N})
  haskey(vals, v.negname) ? vals[v.negname] : (haskey(vals, v.name) ? (1-vals[v.name]) : 1)
end

# But we also want to be able to calculate probabilities
# when values of some variables is not known. This is done by setting both
# the vars and negvars of the unknown variables to 1. Eas

# Simple examples
s1 = SPN(Var(1))
calc(s1, Dict{Int64,Int64}(1 => 1))

s2 = SPN(NVar(1))
calc(s2, Dict{Int64,Int64}(1 => 1))
calc(s2, Dict{Int64,Int64}(1 => 0))
calc(s2, Dict{Int64,Int64}(-1 => 1))
calc(s2, Dict{Int64,Int64}(-1 => 0))
calc(s2, Dict{Int64,Int64}())

s3 = SPN(Var(:a))
calc(s3, Dict{Symbol,Int64}(:a => 1))

# Fig 1 from cs.ubc.ca paper:
at = Sum(SPNode{Symbol}[Var(:a), NVar(:a)], Float64[0.5, 0.5])
bt = Sum(SPNode{Symbol}[Var(:b), NVar(:b)], Float64[0.5, 0.5])
ct = Sum(SPNode{Symbol}[Var(:c), NVar(:c)], Float64[0.5, 0.5])
dt = Sum(SPNode{Symbol}[Var(:d), NVar(:d)], Float64[0.5, 0.5])
left = Prod(SPNode{Symbol}[at, bt])
right = Prod(SPNode{Symbol}[ct, dt])
r = Sum(SPNode{Symbol}[left, right], Float64[0.5, 0.5])
@time calc(r, Dict{Symbol,Int64}(:a => 1, :b => 0, :c => 1, :d => 1))
calc(r, Dict{Symbol,Int64}(:a => 1, :b => 0, :c => 1))


"""
    statscholesky(xtx::Symmetric{T}, tol::Real=-1) where {T<:AbstractFloat}
Return a `CholeskyPivoted` object created from `xtx` where the pivoting scheme
retains the original order unless singularity is detected.  Columns that are
(computationally) linearly dependent on columns to their left are moved to the
right hand side in a left circular shift.
"""
function statscholesky(xtx::Symmetric{T}, tol::Real = -1) where {T<:AbstractFloat}
    n = size(xtx, 2)
    chpiv = cholesky(xtx, Val(true), tol = T(tol), check = false)
    chunp = cholesky(xtx, check = false);

    piv = LinearAlgebra.BlasInt.(collect(1:n))
    r = chpiv.rank

    if r < n

        k = chunp.info
        if k > r
            @warn """Fixed-effects matrix may be ill conditioned.
                     Left-circular shift may fail."""
        end

        nleft = n
        while nleft > r
            # the 0 lowerbound is for MKL compatibility
            # this arises when the unpivoted Cholesky succeeds but the pivoted
            # Chokesky estimates less than full rank (see #367)
            if 0 < k < nleft
                piv = piv[[1:k-1; k+1:n; k]]
                chunp = cholesky!(Symmetric(xtx[piv, piv]), check = false)
            end
            k = chunp.info
            nleft -= 1
        end

        for j = (r+1):n   # an MKL <-> OpenBLAS difference
            for i = (r+1):j
                chunp.factors[i, j] = zero(T)
            end
        end

    end

    CholeskyPivoted(chunp.factors, chunp.uplo, piv, r, tol, chpiv.info)
end

"""
    statsqr(x::Matrix{T}, ranktol::Real=1e-8) where {T<:AbstractFloat}
Return a `QRPivoted` object created from `x` where the pivoting scheme
retains the original order unless singularity is detected.
"""
function statsqr(x::Matrix{T}; ranktol=1e-8) where {T<:AbstractFloat}
    n = size(x, 2)

    qrpiv = qr(x, Val(true))
    rank =  searchsortedlast(abs.(diag(qrpiv.R)), ranktol, rev=true);
    if rank < n
        piv = qrpiv.p
    else
        piv = LinearAlgebra.BlasInt.(collect(1:n))
    end

    if rank < n && piv[1] ≠ 1
        # make sure the first column isn't moved
        # this is usually the intercept and it's desirable to avoid
        # pivoting the intercept out. note that this inflation may
        # not always be enough, but if it isn't, then you probably
        # have other problems

        x = deepcopy(x) # this is horrible
        x[:, 1] .*= 1e6
        qrpiv = qr(x, Val(true))
        piv = qrpiv.p
    end

    # preserve as much of the original column ordering as possible
    piv = [sort!(piv[1:rank]); sort!(piv[rank+1:n])];
    qrunp = qr(x[:, piv], Val(false));

    QRPivoted(qrunp.factors, qrpiv.τ, piv)
end


function barely_solve(input)
    dict = Dict{Int64, Array{Int64}}()
    for (i, n) in enumerate(input)
        dict[n] = [i]
    end
    latest = input[end]
    for i in length(input) + 1:2020
        latest = length(dict[latest]) < 2 ? 0 : dict[latest][end] - dict[latest][end - 1]
        !haskey(dict, latest) ? dict[latest] = [i] : push!(dict[latest], i)
    end
    latest
end

# TO DO: Find an efficient way to do this

function slow_solve(input)
    dict = Dict{Int64, Array{Int64}}()
    for (i, n) in enumerate(input)
        dict[n] = [i, 0]
    end
    latest = input[end]
    for i in length(input) + 1:30000000
        latest = dict[latest][2] == 0 ? 0 : dict[latest][1] - dict[latest][2]
        if !haskey(dict, latest)
            dict[latest] = [i, 0]
        else
            dict[latest][2] = dict[latest][1]
            dict[latest][1] = i
        end
    end
    latest
end

@assert barely_solve([0, 3, 6]) == 436 ["test case 1 failed"]
@assert barely_solve([1, 3, 2]) == 1   ["test case 1a failed"]
@assert barely_solve([2, 1, 3]) == 10  ["test case 1b failed"]
@assert slow_solve([0, 3, 6]) == 175594 ["test case 2 failed"]
@assert slow_solve([1, 3, 2]) == 2578   ["test case 2a failed"]
@assert slow_solve([2, 1, 3]) == 3544142  ["test case 2b failed"]

println("First star: " * string(barely_solve([15, 12, 0, 14, 3, 1])))
println("Second star: " * string(slow_solve([15, 12, 0, 14, 3, 1])))


#=
Definition of measures of cross-correlation and power-law correlation
=#
using MIDI, MusicManipulations
using Statistics, Random
using InformationMeasures
using DCCA

"""
    mutualinformation(x, y [, N]) → MI, MICI, center

Compute MI and it's confidence interval w.r.t the null hypothesis
(i.e no correlations), via bootstrap procedure (total of `N` times),
between `x` and `y`.

The confidence interval for the mutual information corresponding to
the null-hypothesis of no correlations. If the value of MI lies outside this
interval, then `x` and `y` are co-dependant.

Return MI, it's confidence interval and the center of the interval's distribution.
"""
function mutualinformation(MTD::Vector, vel::Vector, N = 10000)
    MTD = copy(MTD); vel = copy(vel)
    MI = round(get_mutual_information(MTD,vel), digits = 2)
    # direct estimation of error bars for MI is too complex.
    # test the null-hypothesis instead.
    # bootstraping procedure to get 95% confidence interval for MI under null-hypothesis.
    bootstrap = zeros(N)
    for i in 1:N
        bootstrap[i] = get_mutual_information(shuffle!(MTD), shuffle!(vel))
    end
    sort!(bootstrap)
    MiCI = (bootstrap[max(N÷40, 1)], bootstrap[N - N÷40])
    center = bootstrap[N÷2]
    return MI, MiCI, center
end

function normalized_mi(MTD::Vector, vel::Vector, N = 10000)
    mi, mici, cen = mutualinformation(MTD, vel, N)
    normalized_mi(mi, mici, cen)
end
function normalized_mi(mi::Real, mici::Tuple, cen::Real)
    mic = mi-cen
    if mic < 0
        return mic/(cen - mici[1])
    else
        return mic/(mici[2] - cen)
    end
end

"""
    pearson(x, y) → ρ, CI
Compute ρ, and it's 95% confidence interval between `x` and `y`.

Notice that `CI` is a 2-tuple (for low and upper bounds).
"""
function pearson(MTD::Vector,vel::Vector)
    r = cor(MTD,vel)
    #computing z transform and standard deviation for r'S error bars.
    #going to the z_transform space makes the distribution symmetric
    z_transform = 0.5*log((1+r)/(1-r))
    sigma = sqrt(1/(length(vel)-3))
    z_critical = 1.96*sigma
    #going back to the r space.
    rCI = (tanh(z_transform - z_critical), tanh(z_transform + z_critical))
    return r, rCI
end

"""
    bin_array(x, bin_size) → binned_array
Bin 'x' with a bin width of 'bin_size'.
"""
function bin_array(x::Array{Int,1}, bin_size::Int)
    new_bins = collect(unique(sort(x))[1]:bin_size:unique(sort(x))[end])
    binned_array = Int[]
    for value in x
        for index in 1:length(new_bins)-1
            if value >= new_bins[index] && value < new_bins[index+1]
                push!(binned_array,new_bins[index])
            end
            if value >= new_bins[end]
                push!(binned_array, new_bins[end])
                break
            end
        end
    end
    return binned_array
end

"""
    correlation_intervals(c)
Given container `c`, integrate with DrWatson's `produce_or_load` and
make a file that contains the Pearson and MI data.
"""
function correlation_intervals(c)
    mtds, vels = get_dataset(c)

    if get(c, :absolute, false)
        for i in 1:length(mtds)
            mtds[i] = abs.(mtds[i])
        end
    end

    # Pearson
    ρs = Float64[]
    ρus = Float64[]
    ρds = Float64[]
    # Mutual information stuff
    MIs = Float64[]
    centers = Float64[]
    MIus = Float64[]
    MIds = Float64[]

    for (m, v) in zip(mtds, vels)
        ρ, (d, u) = pearson(m, v)
        push!(ρs, ρ); push!(ρus, u - ρ); push!(ρds, ρ - d)

        MI, MICI, center = mutualinformation(m, v, 10000)
        push!(MIs, MI)
        push!(MIds, center - MICI[1])
        push!(MIus, MICI[2] - center)
        push!(centers, center)
    end
    return @dict( ρs, ρus, ρds, MIs, centers, MIus, MIds )
end

"""
    get_DCCA(x,y,start,stop,pts)
compute the rho_DCCA coefficients for the given window sizes, as well as a c
onfidence interval for the null hypothesis. Return two tuples :
- (points::Vector, values::Vector) which is the results of the DCCA analyses
  for the different time-scales contained in windows
- (l::Float64, u::Float64) the lower and upper bound for the 95% confidence
  interval
"""
function get_DCCA(x::Array{Float64,1},y::Array{Float64,1})
    windows, values = rhoDCCA(x,y)
    l, u = rhoDCCA_CI(x,y)
    return (windows, values), (l,u)
end


export GraphShortestPathEnv

using Random
using SparseArrays
using LinearAlgebra

mutable struct GraphShortestPathEnv{G, R} <: AbstractEnv
    graph::G
    pos::Int
    goal::Int
    max_steps::Int
    rng::R
    reward::Int
    step::Int
end

"""
    GraphShortestPathEnv([rng]; n=10, sparsity=0.1, max_steps=10)

Quoted **A.3** in the the paper [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://arxiv.org/abs/2106.01345).

> We give details of the illustrative example discussed in the introduction.
> The task is to find theshortest path on a fixed directed graph, which can be
> formulated as an MDP where reward is0whenthe agent is at the goal node
> and−1otherwise.  The observation is the integer index of the graphnode the
> agent is in. The action is the integer index of the graph node to move to
> next. The transitiondynamics transport the agent to the action’s node index if
> there is an edge in the graph, while theagent remains at the past node
> otherwise. The returns-to-go in this problem correspond to negativepath
> lengths and maximizing them corresponds to generating shortest paths.

"""
function GraphShortestPathEnv(rng=Random.GLOBAL_RNG; n=20, sparsity=0.1, max_steps=10)
    graph = sprand(rng, Bool, n, n, sparsity) .| I(n)

    goal = rand(rng, 1:n)
    pos = rand(rng, 1:n)
    while pos == goal
        pos = rand(rng, 1:n)
    end
    GraphShortestPathEnv(graph, pos, goal, max_steps, rng, 0, 0)
end

function (env::GraphShortestPathEnv)(action)
    env.step += 1
    if env.graph[action, env.pos]
        env.pos = action
    end
    env.reward = env.pos == env.goal ? 0 : -1
end

RLBase.state(env::GraphShortestPathEnv) = env.pos
RLBase.state_space(env::GraphShortestPathEnv) = axes(env.graph, 2)
RLBase.action_space(env::GraphShortestPathEnv) = axes(env.graph, 2)
RLBase.legal_action_space(env::GraphShortestPathEnv) = (env.graph[:, env.pos]).nzind
RLBase.reward(env::GraphShortestPathEnv) = env.reward
RLBase.is_terminated(env::GraphShortestPathEnv) = env.pos == env.goal || env.step >= env.max_steps

function RLBase.reset!(env::GraphShortestPathEnv)
    env.step = 0
    env.reward = 0
    env.goal = rand(env.rng, state_space(env))
    env.pos = rand(env.rng, state_space(env))
    while env.pos == env.goal
        env.pos = rand(env.rng, state_space(env))
    end
end

Random.seed!(env::GraphShortestPathEnv, seed) = Random.seed!(env.rng, seed)

function floyd_shortest_path(env::GraphShortestPathEnv)
    n = size(env.graph, 1)
    M = fill(Inf, n, n)
    for idx in CartesianIndices(M)
        if idx[1] == idx[2]
            M[idx] = 0
        elseif env.graph[idx]
            M[idx] = 1
        end
    end
    for k in 1:n
        for i in 1:n
            for j in 1:n
                if M[i, j] > M[i, k] + M[k, j]
                    M[i, j] = M[i, k] + M[k, j]
                end
            end
        end
    end
    M
end

#=

using ReinforcementLearning
using Random

rng = MersenneTwister(123)

env = GraphShortestPathEnv(rng)
policy = RandomPolicy(rng=rng)
M = RLEnvs.floyd_shortest_path(env)

Base.@kwdef struct ShortestPathCount <: AbstractHook
    shortest_paths::Vector{Float64} = []
end

(h::ShortestPathCount)(::PreEpisodeStage, policy, env) = push!(h.shortest_paths, M[env.goal, env.pos])

h = run(policy, env, StopAfterEpisode(1_000), ComposedHook(StepsPerEpisode(), ShortestPathCount()))

using UnicodePlots

barplot(x, [sum(h[2].shortest_paths .== i) for i in x])  # shortest path

#       ┌                                        ┐
#   1.0 ┤■■■■■■■■■■■ 113
#   2.0 ┤■■■■■■■■■■■■■ 134
#   3.0 ┤■■■■■■■■■■■■■■ 147
#   4.0 ┤■■■■■■■■■■■ 111
#   5.0 ┤■■■■■■■ 76
#   6.0 ┤■■■■ 40
#   7.0 ┤■■ 19
#   8.0 ┤ 3
#   9.0 ┤ 0
#   Inf ┤■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■ 357
#       └                                        ┘
#

barplot(1:10, [sum(h[1].steps .== i) for i in 1:10])  # random walk

#      ┌                                        ┐
#    1 ┤■■ 39
#    2 ┤■■ 38
#    3 ┤■ 17
#    4 ┤■ 23
#    5 ┤■ 28
#    6 ┤■ 24
#    7 ┤■ 25
#    8 ┤■ 21
#    9 ┤■ 16
#   10 ┤■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■ 769
#      └                                        ┘
#
=#


using StaticArrays, BenchmarkTools, ForwardDiff
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 10
show_ind_bench = true
println("Benchmarking Jacobian calculation of 3D Discrete systems for different types of e.o.m.")

u = rand(3); un = rand(3); su = SVector{3}(u); mu = MVector{3}(u)
mun = MVector{3}(un)
J = rand(3,3); MJ = MMatrix{3,3}(J); MI = MMatrix{3,3}(J);
mutable struct system1
  J::StaticArrays.MArray{Tuple{3,3},Float64,2,9}
end
sys = system1(MI)

println("\nVersion 1: StaticArray returned: eom_1(u) -> un::SVector")
@inline @inbounds function eom_1(x)
  x1 = x[1]; x2=x[2]; x3 = x[3]
  SVector{3}(
  3.8*x1*(1-x1)-0.05*(x2+0.35)*(1-2*x3),
  0.1*( (x2+0.35)*(1-2*x3)-1 )*(1-1.9*x1),
  3.78*x3*(1-x3)+0.2*x2
  )
end
println("--In-place jacobian")
@inline jac1!(J, x) = ForwardDiff.jacobian!(J, eom_1, x)
bj1 = @benchmark jac1!($sys.J, $su)
show_ind_bench && display(bj1)
println("--Return-value jacobian")
@inline jac1(x) = (j = ForwardDiff.jacobian(eom_1, x); sys.J .= j)
bjr1 = @benchmark jac1($su)
show_ind_bench && display(bjr1)

println("\nVersion 2: Base Array returned eom_2(x) -> xn::Vector ")
@inline @inbounds function eom_2(x)
  x1 = x[1]; x2=x[2]; x3 = x[3];
  xn = zeros(x)
  xn[1] = 3.8*x1*(1-x1)-0.05*(x2+0.35)*(1-2*x3)
  xn[2] = 0.1*( (x2+0.35)*(1-2*x3)-1 )*(1-1.9*x1)
  xn[3] = 3.78*x3*(1-x3)+0.2*x2
  return xn
end
cfg = ForwardDiff.JacobianConfig(eom_2, u)
println("--In-place jacobian")
@inline jac2!(J, x, cfg) = ForwardDiff.jacobian!(J, eom_2, x, cfg)
bj2 = @benchmark jac2!($sys.J, $su, $cfg)
show_ind_bench && display(bj2)
println("--Return-value jacobian")
@inline jac2(x, cfg) = (j = ForwardDiff.jacobian(eom_2, x, cfg); sys.J .= j)
bjr2 = @benchmark jac2($su, $cfg)
show_ind_bench && display(bjr2)

println("\nVersion 3: MVector returned eom_3(x) -> xn::MVector ")
@inline @inbounds function eom_3(x)
  x1 = x[1]; x2=x[2]; x3 = x[3];
  xn = MVector{3}(x)
  xn[1] = 3.8*x1*(1-x1)-0.05*(x2+0.35)*(1-2*x3)
  xn[2] = 0.1*( (x2+0.35)*(1-2*x3)-1 )*(1-1.9*x1)
  xn[3] = 3.78*x3*(1-x3)+0.2*x2
  return xn
end
cfg = ForwardDiff.JacobianConfig(eom_3, mu)
println("--In-place jacobian")
@inline jac3!(J, x, cfg) = ForwardDiff.jacobian!(J, eom_3, x, cfg)
bj3 = @benchmark jac3!($MJ, $mu, $cfg)
show_ind_bench && display(bj3)
println("--Return-value jacobian")
@inline jac3(x, cfg) = (j = ForwardDiff.jacobian(eom_3, x, cfg); sys.J .= j)
bjr3 = @benchmark jac3($mu, $cfg)
show_ind_bench && display(bjr3)

println("\nVersion 4: In place with 2 arguments: eom_towel!(xn, x) (changes xn)")
@inline @inbounds function eom_4!(xn, x)
  xn[1] = 3.8*x[1]*(1-x[1])-0.05*(x[2]+0.35)*(1-2*x[3])
  xn[2] = 0.1*( (x[2]+0.35)*(1-2*x[3])-1 )*(1-1.9*x[1])
  xn[3] = 3.78*x[3]*(1-x[3])+0.2*x[2]
end
cfg = ForwardDiff.JacobianConfig(eom_4!, un, u)
println("--In-place jacobian")
@inline jac4!(J, y, x, cfg) = ForwardDiff.jacobian!(J, eom_4!, y, x, cfg)
bj4 = @benchmark jac4!($J, $un, $u, $cfg)
show_ind_bench && display(bj4)
println("--Return-value jacobian")
@inline jac4(y, x, cfg) = (j = ForwardDiff.jacobian(eom_4!, y, x, cfg); sys.J .= j)
bjr4 = @benchmark jac4($un, $u, $cfg)
show_ind_bench && display(bjr4)

println("\nVersion 5: In place with 2 arguments & using MVector")
@inline @inbounds function eom_5!(xn::MVector, x::MVector)
  xn[1] = 3.8*x[1]*(1-x[1])-0.05*(x[2]+0.35)*(1-2*x[3])
  xn[2] = 0.1*( (x[2]+0.35)*(1-2*x[3])-1 )*(1-1.9*x[1])
  xn[3] = 3.78*x[3]*(1-x[3])+0.2*x[2]
end
cfg5 = ForwardDiff.JacobianConfig(eom_5!, mun, mu)
println("--In-place jacobian")
@inline jac5!(J, y, x, cfg) = ForwardDiff.jacobian!(J, eom_5!, y, x, cfg)
bj5 = @benchmark jac5!($MJ, $mun, $mu, $cfg5)
show_ind_bench && display(bj5)
println("--Return-value jacobian")
@inline jac5(y, x, cfg) = (j = ForwardDiff.jacobian(eom_5!, y, x, cfg); sys.J .= j)
bjr5 = @benchmark jac5($mun, $mu, $cfg5)
show_ind_bench && display(bjr5)

println("comparison of in-place jacobian call:")
for (i, b) in enumerate([bj2,bj3,bj4,bj5])
  println("v$(i+1) versus v1")
  sleep(0.1)
  display(judge(median(b), median(bj1)))
  sleep(0.1)
end

println("Conclusion about ForwardDiff jacobians:")
println("All methods for Jacobian are slower than the first (which uses SVector)")
println("However method 4 comes extremely close to it!")
println("It is actually even faster than the one with MVector")


function hypergraph_density(H, number_of_species)
    temp = [count(h -> length(h) == i, H)/binomial(number_of_species, i) for i=1:number_of_species]

    return temp
end

function hypergraph_dimension(H)
    return maximum(length.(H))
end

function get_descendents(G, h)
    return [union(h,v) for v in neighbors(G, last(h))]
end

function missing_edges(hypergraph)
    return setdiff(minimal_simplicial_complex(hypergraph), hypergraph)
end

"""
read_hypergraph(filename)
"""
function read_hypergraph(filename)
    Hp = readlines(filename)
    temp = [ vcat([parse(Int, ss) for ss in split.(h, ",")]...) for h in Hp]
    return temp
end

"""
hypergraph_subdivide(H; expansion = false)
"""
function hypergraph_subdivide(H; expansion = false)

    # Build graph of inclusion between hyperedges.
    # convention is [i,j] for (i -> j)
    G = []
    for size = reverse(2:hypergraph_dimension(H))
            nodes_s = findall(h -> length(h) == size , H) # v ∈ V such that length(H[v]) == size
            nodes_sm1 = findall(h -> length(h) < size, H) # v ∈ V such that length(H[v]) < size

            for p in nodes_s, c in nodes_sm1
                    if H[c] ⊆ H[p]
                            push!(G, [p,c])
                    end
            end
    end

    if expansion .== false # return just G
            return G
    else # expand G to the Vietoris-Rips comples

            # Expand G into K_H
            roots = setdiff(1:length(H), last.(G)) # roots of G, all nodes without incoming edges (characterized by last item of pairs [i,j])
            tops = setdiff(1:length(H), first.(G)) # tops of G, all nodes without outgoing edges

            list = [[i] for i=1:length(H)] #initialize list of hyperedges
            let temp = [[h] for h in roots]
                    @inbounds @showprogress "Subdividing: " for i =1:hypergraph_dimension(H)-1
                            #temp = [union(h, v) for h in temp for v in neighbors(G, last(h)) ]
                            temp = vcat(map(h -> get_descendents(G, h), temp)...)
                            union!(list, temp)
                            # remove all nodes that are already at the top of G
                            temp = filter(x -> .!issubset(last(x), tops), temp)
                    end
            end
            return list
    end
end

"""
betti_hypergraph_ripscomplex(H; max_dim = 3)

Computes Betti numbers of hypergraph.
Based on computing the hypergraph subdivision without expansion, which returns the inclusion graph
"""
function betti_hypergraph_ripscomplex(H; max_dim = 3)

    if length(H) == 0
        return Int.(zeros(max_dim + 1))
    elseif length(H) == 1
        return Int.(vcat(1, zeros(max_dim)))
    end

    _max_dim = min(max_dim, length(H) - 2);

    G = hypergraph_subdivide(H; expansion = false)
    # build distance matrix to give as input to Ripserer.
    M = 2.0*ones(length(H), length(H))
    [ M[g[1], g[2]] = 1 for g in G ]
    [ M[g[2], g[1]] = 1 for g in G ]
    M[diagind(M)] .= 0;

    #call Ripser
    flt = Rips(M; threshold=1.5, sparse=true)
    barcodes = ripserer(flt; dim_max = _max_dim)
    # interpret the outout. Barcodes is an array of tuples of dim = max_dim + 1.
    # All tuples where the second entry is Inf are the persistent ones.
    # Betti numbers just count the number of those persistent ones.
    betti = [ count(x -> x[2] .== Inf, barcodes[i]) for i in 1:length(barcodes)]

    padding = (max_dim+1) - length(betti);
    return vcat(betti, Int.(zeros(padding)))

    ## old function using Ripser library: https://github.com/mtsch/Ripser.jl#master
    # if length(H) == 0
    #     return zeros(max_dim + 1)
    # else
    #
    #     G = hypergraph_subdivide(H; expansion = false)
    #     # build distance matrix to give as input to Ripser.
    #     M = 2.0*ones(length(H), length(H))
    #     [ M[g[1], g[2]] = 1 for g in G ]
    #
    #     #call Ripser
    #     barcodes = ripser(M; dim_max = max_dim, threshold = 1.5)
    #     # interpret the outout. Barcodes is an array of tuples of dim = max_dim + 1.
    #     # All tuples where the second entry is Inf are the persistent ones.
    #     # Betti numbers just count the number of those persistent ones.
    #     betti = [ count(x -> x[2] .== Inf, barcodes[i]) for i in 1:length(barcodes)]
    #
    #     return betti
    # end
end

"""
minimal_simplicial_complex(hypergraph)

Computes the minimal simplicial complex containig the hypergraph H
"""
function minimal_simplicial_complex(hypergraph)
    K = Array{Int64, 1}[]
    #K = []
    for h in hypergraph
        union!(K, collect(combinations(h)))
    end
    return unique(Set, K) # delete duplicated simplices
end

"""
disassembly_hypergraph(hypergraph)

"""
function disassembly_hypergraph(hypergraph)
    M = missing_edges(hypergraph)
    temp = filter(h -> length(h) == 1, hypergraph) # hyperedges of dimension = 1 (i.e., vertices)

    return union(temp, M)
end