function (+){T<:AbstractFloat, R<:RoundingMode}(a::T, b::T, rounding::R)
    hi,lo = eftSum2(a,b)
    eftRound(hi, lo, rounding)
end

function (-){T<:AbstractFloat, R<:RoundingMode}(a::T, b::T, rounding::R)
    hi,lo = eftDiff2(a,b)
    eftRound(hi, lo, rounding)
end

function (*){T<:AbstractFloat, R<:RoundingMode}(a::T, b::T, rounding::R)
    hi,lo = eftProd2(a,b)
    eftRound(hi, lo, rounding)
end

function (/){T<:AbstractFloat, R<:RoundingMode}(a::T, b::T, rounding::R)
    hi,lo = eftDiv2Approx(a,b)
    eftRound(hi, lo, rounding)
end

function (square){T<:AbstractFloat, R<:RoundingMode}(a::T, rounding::R)
    hi,lo = eftSquare(a)
    eftRound(hi, lo, rounding)
end

function (sqrt){T<:AbstractFloat, R<:RoundingMode}(a::T, rounding::R)
    hi,lo = eftSqrtApprox(a)
    eftRound(hi, lo, rounding)
end

for op in (:(+), :(-), :(*), :(/))
    @eval ($op){T<:Integer, R<:RoundingMode}(a::Rational{T}, b::Rational{T}, rounding::R) = ($op)(a,b)
end

if isdefined(Main,:FloatFloat)
FF=Main.FloatFloat.FF
for op in (:(+), :(-), :(*), :(/))
    @eval begin
        function ($op){T<:FF, R<:RoundingMode}(a::T, b::T, rounding::R)
            c = zero(T)
            with_rounding(Float64,rounding) do
                c = ($op)(a, b)
            end
            c
        end
    end
end
end
function (+){T<:Real, R<:RoundingMode}(a::T, b::T, rounding::R)
    c = zero(T)
    with_rounding(T,R) do
        c = a + b
    end
    c
end

function (-){T<:Real, R<:RoundingMode}(a::T, b::T, rounding::R)
    c = zero(T)
    with_rounding(T,R) do
        c = a - b
    end
    c
end

function (*){T<:Real, R<:RoundingMode}(a::T, b::T, rounding::R)
    c = zero(T)
    with_rounding(T,R) do
        c = a * b
    end
    c
end

function (/){T<:Real, R<:RoundingMode}(a::T, b::T, rounding::R)
    c = zero(T)
    with_rounding(T,R) do
        c = a / b
    end
    c
end

function (square){T<:Real, R<:RoundingMode}(a::T, rounding::R)
    c = zero(T)
    with_rounding(T,R) do
        c = a*a
    end
    c
end

function (sqrt){T<:Real, R<:RoundingMode}(a::T, rounding::R)
    c = zero(T)
    with_rounding(T,R) do
        c = sqrt(a)
    end
    c
end


using Distributed
@everywhere using Mamba

## Data
eyes = Dict{Symbol, Any}(
  :y =>
    [529.0, 530.0, 532.0, 533.1, 533.4, 533.6, 533.7, 534.1, 534.8, 535.3,
     535.4, 535.9, 536.1, 536.3, 536.4, 536.6, 537.0, 537.4, 537.5, 538.3,
     538.5, 538.6, 539.4, 539.6, 540.4, 540.8, 542.0, 542.8, 543.0, 543.5,
     543.8, 543.9, 545.3, 546.2, 548.8, 548.7, 548.9, 549.0, 549.4, 549.9,
     550.6, 551.2, 551.4, 551.5, 551.6, 552.8, 552.9, 553.2],
  :N => 48,
  :alpha => [1, 1]
)

## Model Specification
model = Model(

  y = Stochastic(1,
    (lambda, T, s2, N) ->
      begin
        sigma = sqrt(s2)
        UnivariateDistribution[(
          mu = lambda[Int(T[i])];
          Normal(mu, sigma)) for i in 1:N
        ]
      end,
    false
  ),

  T = Stochastic(1,
    (P, N) -> UnivariateDistribution[Categorical(P) for i in 1:N],
    false
  ),

  P = Stochastic(1,
    alpha -> Dirichlet(alpha)
  ),

  lambda = Logical(1,
    (lambda0, theta) -> Float64[lambda0; lambda0 + theta]
  ),

  lambda0 = Stochastic(
    () -> Normal(0.0, 1000.0),
    false
  ),

  theta = Stochastic(
    () -> Uniform(0.0, 1000.0),
    false
  ),

  s2 = Stochastic(
    () -> InverseGamma(0.001, 0.001)
  )

)

## Initial Values
inits = [
  Dict(:y => eyes[:y], :T => fill(1, eyes[:N]), :P => [0.5, 0.5],
       :lambda0 => 535, :theta => 5, :s2 => 10),
  Dict(:y => eyes[:y], :T => fill(1, eyes[:N]), :P => [0.5, 0.5],
       :lambda0 => 550, :theta => 1, :s2 => 1)
]

## Sampling Scheme
scheme = [DGS(:T),
          Slice([:lambda0, :theta], [5.0, 1.0]),
          Slice(:s2, 2.0, transform=true),
          SliceSimplex(:P, scale=0.75)]
setsamplers!(model, scheme)

## MCMC Simulations
sim = mcmc(model, eyes, inits, 10000, burnin=2500, thin=2, chains=2)
describe(sim)


"""
    AcrobotEnv(;kwargs...)

# Keyword arguments
- `T = Float64`
- `link_length_a = T(1.)`
- `link_length_b = T(1.)`
- `link_mass_a = T(1.)`
- `link_mass_b = T(1.)`
- `link_com_pos_a = T(0.5)`
- `link_com_pos_b = T(0.5)`
- `link_moi = T(1.)`
- `max_vel_a = T(4 * π)`
- `max_vel_b = T(9 * π)`
- `g = T(9.8)`
- `dt = T(0.2)`
- `max_steps = 200`
- `book_or_nips = 'book'`
- `avail_torque = [T(-1.), T(0.), T(1.)]`
"""
function AcrobotEnv(;
    T=Float64,
    link_length_a=T(1.0),
    link_length_b=T(1.0),
    link_mass_a=T(1.0),
    link_mass_b=T(1.0),
    link_com_pos_a=T(0.5),
    link_com_pos_b=T(0.5),
    link_moi=T(1.0),
    max_torque_noise=T(0.0),
    max_vel_a=T(4 * π),
    max_vel_b=T(9 * π),
    g=T(9.8),
    dt=T(0.2),
    max_steps=200,
    rng=Random.GLOBAL_RNG,
    book_or_nips="book",
    avail_torque=[T(-1.0), T(0.0), T(1.0)],
)

    params = AcrobotEnvParams{T}(
        link_length_a,
        link_length_b,
        link_mass_a,
        link_mass_b,
        link_com_pos_a,
        link_com_pos_b,
        link_moi,
        max_torque_noise,
        max_vel_a,
        max_vel_b,
        g,
        dt,
        max_steps,
    )

    env = AcrobotEnv(
        params,
        zeros(T, 4),
        0,
        false,
        0,
        rng,
        T(0.0),
        book_or_nips,
        [T(-1.0), T(0.0), T(1.0)],
    )
    reset!(env)
    env
end

acrobot_observation(s) = [cos(s[1]), sin(s[1]), cos(s[2]), sin(s[2]), s[3], s[4]]

RLBase.action_space(env::AcrobotEnv) = Base.OneTo(3)

function RLBase.state_space(env::AcrobotEnv{T}) where {T}
    high = [1.0, 1.0, 1.0, 1.0, env.params.max_vel_a, env.params.max_vel_b]
    Space(ClosedInterval{T}.(-high, high))
end

RLBase.is_terminated(env::AcrobotEnv) = env.done
RLBase.state(env::AcrobotEnv) = acrobot_observation(env.state)
RLBase.reward(env::AcrobotEnv) = env.reward

function RLBase.reset!(env::AcrobotEnv{T}) where {T <: Number}
    env.state[:] = T(0.1) * rand(env.rng, T, 4) .- T(0.05)
    env.t = 0
    env.action = 2
    env.done = false
    env.reward = -1
    nothing
end

# governing equations as per python gym
function (env::AcrobotEnv{T})(a) where {T <: Number}
    env.action = a
    env.t += 1
    torque = env.avail_torque[a]

    # noise to the force action
    noise_range = env.params.max_torque_noise
    if noise_range > 0
        torque += T(2.0 * noise_range) * rand(env.rng, T, 1) .- T(noise_range)
    end

    # augmented state for derivative function
    s_augmented = [env.state..., torque]

    ode = OrdinaryDiffEq.ODEProblem(dsdt, s_augmented, (0.0, env.params.dt), env)
    ns = OrdinaryDiffEq.solve(ode, OrdinaryDiffEq.RK4())
    # only care about final timestep of integration returned by integrator
    ns = ns.u[end]
    ns = ns[1:4]  # omit action

    # wrap the solution
    ns[1] = wrap(ns[1], -π, π)
    ns[2] = wrap(ns[2], -π, π)
    ns[3] = bound(ns[3], -env.params.max_vel_a, env.params.max_vel_a)
    ns[4] = bound(ns[4], -env.params.max_vel_b, env.params.max_vel_b)
    env.state = ns
    # termination criterion
    succeeded = -cos(ns[1]) - cos(ns[2] + ns[1]) > 1.0
    env.done = succeeded || env.t > env.params.max_steps
    env.reward = succeeded ? 0.0 : -1.0
    nothing
end

function dsdt(du, s_augmented, env::AcrobotEnv, t)
    # extract params
    m1 = env.params.link_mass_a
    m2 = env.params.link_mass_b
    l1 = env.params.link_length_a
    lc1 = env.params.link_com_pos_a
    lc2 = env.params.link_com_pos_b
    I1 = env.params.link_moi
    I2 = env.params.link_moi
    g = env.params.g

    # extract action and state
    a = s_augmented[end]
    s = s_augmented[1:end - 1]

    # writing in standard form
    theta1 = s[1]
    theta2 = s[2]
    dtheta1 = s[3]
    dtheta2 = s[4]
    ddtheta1 = 0.0
    ddtheta2 = 0.0

    # governing equations
    d1 = (m1 * lc1^2 + m2 * (l1^2 + lc2^2 + 2 * l1 * lc2 * cos(theta2)) + I1 + I2)
    d2 = m2 * (lc2^2 + l1 * lc2 * cos(theta2)) + I2
    phi2 = m2 * lc2 * g * cos(theta1 + theta2 - pi / 2.0)
    phi1 = (
        -m2 * l1 * lc2 * dtheta2^2 * sin(theta2) -
        2 * m2 * l1 * lc2 * dtheta2 * dtheta1 * sin(theta2) +
        (m1 * lc1 + m2 * l1) * g * cos(theta1 - pi / 2) +
        phi2
    )
    if env.book_or_nips == "nips"
        # the following line is consistent with the description in the
        # paper
        ddtheta2 = ((a + d2 / d1 * phi1 - phi2) / (m2 * lc2^2 + I2 - d2^2 / d1))
    elseif env.book_or_nips == "book"
        # the following line is consistent with the java implementation and the
        # book
        ddtheta2 = (
            (a + d2 / d1 * phi1 - m2 * l1 * lc2 * dtheta1^2 * sin(theta2) - phi2) /
            (m2 * lc2^2 + I2 - d2^2 / d1)
        )
        ddtheta1 = -(d2 * ddtheta2 + phi1) / d1
    end

    # return the values
    du[1] = dtheta1
    du[2] = dtheta2
    du[3] = ddtheta1
    du[4] = ddtheta2
    du[5] = 0.0
end

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

# wrap as per python gym
function wrap(x, m, M)
    """
    Wraps ``x`` so m <= x <= M; but unlike ``bound()`` which
    truncates, ``wrap()`` wraps x around the coordinate system defined by m,M.\n
    For example, m = -180, M = 180 (degrees), x = 360 --> returns 0.
    Args:
        x: a scalar
        m: minimum possible value in range
        M: maximum possible value in range
    Returns:
        x: a scalar, wrapped
    """
    diff = M - m
    while x > M
        x = x - diff
    end
    while x < m
        x = x + diff
    end
return x
end

function bound(x, m, M)
    """Either have m as scalar, so bound(x,m,M) which returns m <= x <= M *OR*
    have m as length 2 vector, bound(x,m, <IGNORED>) returns m[0] <= x <= m[1].
    Args:
        x: scalar
    Returns:
        x: scalar, bound between min (m) and Max (M)
    """

    # bound x between min (m) and Max (M)
    return min(max(x, m), M)
end


# TODO: make more efficient
"""
    GaussNewton solver
"""
struct GNSolver <: LinearSolver end

function gn_solver(A)
    GNSolver()
end

function linear_solve!(solver::GNSolver, x::Vector{T}, A::SparseMatrixCSC{T,Int}, b::Vector{T}) where T
    ldiv!(x, lu!(Array(A' * A)), Array(A' * b)) # solve
end

function linear_solve!(solver::GNSolver, x::Matrix{T}, A::AbstractMatrix{T}, B::AbstractMatrix{T}) where T
    ldiv!(x, lu!(A' * A), A' * B) # solve
end


module Geometry

using StaticArrays, LinearAlgebra, DocStringExtensions
using ..Grids: _ξ1x1, _ξ2x1, _ξ3x1, _ξ1x2, _ξ2x2, _ξ3x2, _ξ1x3, _ξ2x3, _ξ3x3, _M, _MI,
       _x1, _x2, _x3, _JcV
export LocalGeometry, lengthscale, resolutionmetric

"""
    LocalGeometry

The local geometry at a nodal point.

# Constructors

    LocalGeometry(polynomial::Val, vgeo::AbstractArray{T}, n::Integer, e::Integer)

Extracts a `LocalGeometry` object from the `vgeo` array at node `n` in element `e`.

# Fields

$(DocStringExtensions.FIELDS)
"""
struct LocalGeometry{T,P}
  "Polynomial interpolant: currently this is assumed to be `Val{polyorder}`, but this may change in future."
  polynomial::P
  "Cartesian coordinates"
  coord::SVector{3,T}
  "Jacobian from Cartesian to element coordinates: `invJ[i,j]` is ``∂ξ_i/∂x_j``"
  invJ::SMatrix{3,3,T,9}
end

function LocalGeometry(polynomial::Val, vgeo::AbstractArray{T}, n::Integer, e::Integer) where {T}
  coord = @SVector T[vgeo[n, _x1, e], vgeo[n, _x2, e], vgeo[n, _x3, e]]
  invJ = @SMatrix T[vgeo[n, _ξ1x1, e] vgeo[n, _ξ1x2, e] vgeo[n, _ξ1x3, e];
                    vgeo[n, _ξ2x1, e] vgeo[n, _ξ2x2, e] vgeo[n, _ξ2x3, e];
                    vgeo[n, _ξ3x1, e] vgeo[n, _ξ3x2, e] vgeo[n, _ξ3x3, e]]

  LocalGeometry(polynomial, coord, invJ)
end

"""
    resolutionmetric(g::LocalGeometry)

The metric tensor of the discretisation resolution. Given a unit vector `u` in Cartesian
coordinates and `M = resolutionmetric(g)`, `sqrt(u'*M*u)` is the degree-of-freedom density
in the direction of `u`.
"""
function resolutionmetric(g::LocalGeometry{T, Val{polyorder}}) where {T,polyorder}
  S = polyorder*g.invJ/2
  S'*S # TODO: return an eigendecomposition / symmetric object?
end

"""
    lengthscale(g::LocalGeometry)

The effective grid resolution at the point.
"""
function lengthscale(g::LocalGeometry{T, Val{polyorder}}) where {T,polyorder}
  2/(cbrt(det(g.invJ))*polyorder)
end

end # module


module IterativeSolvers

using LinearAlgebra
using Parameters
using Printf

import LinearAlgebra: ldiv!

export Continuous, IterativeSolver, ConjugateGradient, BiCGStab, GMRES, solve!

"""
Allows the identity operator `I` to be used as a preconditioner.
"""
function ldiv!(vout::AbstractVector,I::UniformScaling,vin::AbstractVector)
    copyto!(vout,vin)
    return nothing
end

"""
Abstract type to represent continuous real or complex numbers.
"""
Continuous = Union{AbstractFloat,Complex{<:AbstractFloat}}

"""
Abstract type to represent Iterative Solver methods.
"""
abstract type IterativeSolver{Ttol<:AbstractFloat,Tdata<:Continuous} end

#######################################
## Conjugate Gradient Implementation ##
#######################################

"""
To represent Conjugate Gradient algorithm.
"""
mutable struct ConjugateGradient{Ttol,Tdata} <: IterativeSolver{Ttol,Tdata}

    tol::Ttol
    maxiter::Int
    κmax::Ttol
    N::Int
    r::Vector{Tdata}
    p::Vector{Tdata}
    z::Vector{Tdata}

    function ConjugateGradient(z::AbstractVector{Tdata};tol::Ttol=1e-4,maxiter::Int=0,κmax::Ttol=1e12) where {Ttol<:AbstractFloat,Tdata<:Continuous}

        N = length(z)
        if maxiter<1
            maxiter = N
        end
        r = zeros(Tdata,N)
        p = zeros(Tdata,N)
        z = zeros(Tdata,N)
        return new{Ttol,Tdata}(tol,maxiter,κmax,N,r,p,z)
    end
end

"""
Solve `A⋅x=b` using the Conjugate Gradient method with a split preconditioner
such that `[L⁻¹⋅A⋅L⁻ᵀ]⋅u=L⁻¹⋅b` where `u=L⁻ᵀ⋅x`.
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},cg::ConjugateGradient{Ttol,Tdata},L,Lt;
                maxiter::Int=0,tol::Ttol=0.0,κmax::Ttol=0.0)::Int where {Ttol,Tdata}

    r = cg.r
    p = cg.p
    z = cg.z

    # set to default maxiter
    if iszero(maxiter)
        maxiter = cg.maxiter
    end

    # set to default tolerance
    if iszero(tol)
        tol = cg.tol
    end

    # set default max condition number
    if iszero(κmax)
        κmax = cg.κmax
    end

    # r₀ = b - A⋅x₀
    mul!(r,A,x)
    axpby!(1.0,b,-1.0,r)

    # r₀ = L⁻¹⋅r₀ = L⁻¹⋅(b - A⋅x₀)
    ldiv!(L,r)

    # p₀ = L⁻ᵀ⋅r₀ = L⁻ᵀ⋅L⁻¹⋅(b - A⋅x₀)
    ldiv!(p,Lt,r)

    # |L⁻ᵀ⋅L⁻¹⋅b|
    ldiv!(z,L,b)
    ldiv!(Lt,z)
    normL⁻ᵀL⁻¹b = norm(z)

    # calcualte initial tolerance
    ϵ₀ = norm(p)/normL⁻ᵀL⁻¹b
    ϵ  = ϵ₀

    # initial lower bound for condition number
    κmin  = 0.0

    # r₀⋅r₀
    rdotr = dot(r,r)

    @fastmath @inbounds for j in 1:maxiter

        # αⱼ = (rⱼ⋅rⱼ)/(pⱼ⋅A⋅pⱼ)
        mul!(z,A,p)
        α = rdotr/dot(p,z)

        # xⱼ₊₁ = xⱼ + αⱼ⋅pⱼ
        axpy!(α,p,x)

        # rⱼ₊₁ = rⱼ - αⱼ⋅L⁻¹⋅A⋅pⱼ
        ldiv!(L,z)
        axpy!(-α,z,r)

        # βⱼ = (rⱼ₊₁⋅rⱼ₊₁)/(rⱼ⋅rⱼ)
        new_rdotr = dot(r,r)
        β     = new_rdotr/rdotr
        rdotr = new_rdotr

        # pⱼ₊₁ = L⁻ᵀ⋅rⱼ₊₁ + βⱼ⋅pⱼ
        ldiv!(z,Lt,r)
        axpby!(1.0,z,β,p)

        # ϵ = |pⱼ₊₁|/|L⁻ᵀ⋅L⋅b| = |L⁻ᵀ⋅L⋅(A⋅xⱼ₊₁-b)|/|L⁻ᵀ⋅L⋅b|
        ϵ = norm(p)/normL⁻ᵀL⁻¹b

        # approximate lower bound for condition numbers
        κmin = max( κmin , (2*j/log(2*ϵ₀/ϵ))^2 )

        # check stop criteria
        if ϵ < tol  || κmin > κmax
            return j
        end
    end

    return maxiter

end

"""
Solve `A⋅x=b` using the Conjugate Gradient method with a preconditioner `P`.
Based on pseudocode from: https://www-users.cs.umn.edu/~saad/Calais/PREC.pdf
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},cg::ConjugateGradient{Ttol,Tdata},P;
                maxiter::Int=0,tol::Ttol=0.0,κmax::Ttol=0.0)::Int where {Ttol,Tdata}

    r = cg.r
    p = cg.p
    z = cg.z

    # set to default maxiter
    if iszero(maxiter)
        maxiter = cg.maxiter
    end

    # set to default tolerance
    if iszero(tol)
        tol = cg.tol
    end

    # set default max condition number
    if iszero(κmax)
        κmax = cg.κmax
    end

    # |b|
    normb = norm(b)

    # r₀ = b - A⋅x₀
    mul!(r,A,x)
    axpby!(1.0,b,-1.0,r)

    # z₀ = P⁻¹⋅r₀ = P⁻¹⋅(b - A⋅x₀)
    ldiv!(z,P,r)

    # p₀ = z₀
    copyto!(p,z)

    # r₀⋅z₀
    rdotz = dot(r,z)

    # calcualte initial tolerance
    ϵ₀ = norm(r)/normb
    ϵ  = ϵ₀

    # initial lower bound for condition number
    κmin  = 0.0

    @fastmath @inbounds for j in 1:maxiter

        # αⱼ = (rⱼ⋅zⱼ)/(pⱼ⋅A⋅pⱼ)
        mul!(z,A,p)
        α = rdotz/dot(p,z)

        # xⱼ₊₁ = xⱼ + αⱼ⋅pⱼ
        axpy!(α,p,x)

        # rⱼ₊₁ = rⱼ - αⱼ⋅A⋅pⱼ
        axpy!(-α,z,r)

        # ϵ = |rⱼ₊₁|/|b|
        ϵ = norm(r)/normb

        # approximate lower bound for condition numbers
        κmin = max( κmin , (2*j/log(2*ϵ₀/ϵ))^2 )

        # check stop criteria
        if ϵ < tol  || κmin > κmax
            return j
        end

        # zⱼ₊₁ = P⁻¹⋅rⱼ₊₁
        ldiv!(z,P,r)

        # βⱼ = (rⱼ₊₁⋅zⱼ₊₁)/(rⱼ⋅zⱼ)
        new_rdotz = dot(r,z)
        β     = new_rdotz/rdotz
        rdotz = new_rdotz

        # pⱼ₊₁ = zⱼ₊₁ + βⱼ⋅pⱼ
        axpby!(1.0,z,β,p)
    end

    return maxiter
end

"""
Solve `A⋅x=b` using the Conjugate Gradient method with no preconditioning.
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},cg::ConjugateGradient{Ttol,Tdata};
                maxiter::Int=0,tol::Ttol=0.0,κmax::Ttol=0.0)::Int where {Ttol,Tdata}

    r = cg.r
    p = cg.p
    z = cg.z

    # set to default maxiter
    if iszero(maxiter)
        maxiter = cg.maxiter
    end

    # set to default tol
    if iszero(tol) || κmin > κmax
        tol = cg.tol
    end

    # set default max condition number
    if iszero(κmax)
        κmax = cg.κmax
    end

    # |b|
    normb = norm(b)

    # r₀ = b - A⋅x₀
    mul!(r,A,x)
    axpby!(1.0,b,-1.0,r)

    # p₀ = r₀
    copyto!(p,r)

    # r₀⋅r₀
    rdotr = dot(r,r)

    # calcualte initial tolerance
    ϵ₀ = norm(r)/normb
    ϵ  = ϵ₀

    # initial lower bound for condition number
    κmin = 0.0

    @fastmath @inbounds for j in 1:maxiter

        # αⱼ = (rⱼ⋅rⱼ)/(pⱼ⋅A⋅pⱼ)
        mul!(z,A,p)
        α = rdotr/dot(p,z)

        # xⱼ₊₁ = xⱼ + αⱼ⋅pⱼ
        axpy!(α,p,x)

        # rⱼ₊₁ = rⱼ - αⱼ⋅A⋅pⱼ
        axpy!(-α,z,r)

        # ϵ = |rⱼ₊₁|/|b| = |b-A⋅xⱼ₊₁|/|b|
        ϵ = norm(r)/normb

        # approximate lower bound for condition numbers
        κmin = max( κmin , (2*j/log(2*ϵ₀/ϵ))^2 )

        # check stop criteria
        if ϵ < tol || κmin > κmax
            return j
        end

        # βⱼ = (rⱼ₊₁⋅rⱼ₊₁)/(rⱼ⋅rⱼ)
        new_rdotr = dot(r,r)
        β     = new_rdotr/rdotr
        rdotr = new_rdotr

        # pⱼ₊₁ = rⱼ₊₁ + βⱼ⋅pⱼ
        axpby!(1.0,r,β,p)
    end

    return maxiter
end

#############################
## BiCGStab Implementation ##
#############################

"""
To represent BiCGStab algorithm.
"""
mutable struct BiCGStab{Ttol,Tdata} <: IterativeSolver{Ttol,Tdata}

    tol::Ttol
    maxiter::Int
    r::Vector{Tdata}
    r̃::Vector{Tdata}
    p::Vector{Tdata}
    p̂::Vector{Tdata}
    s::Vector{Tdata}
    ŝ::Vector{Tdata}
    v::Vector{Tdata}
    t::Vector{Tdata}

    function BiCGStab(x::AbstractVector{Tdata};tol::Ttol=1e-4,maxiter::Int=0) where {Ttol<:AbstractFloat,Tdata<:Continuous}

        N = length(x)
        r = zeros(Tdata,N)
        r̃ = zeros(Tdata,N)
        p = zeros(Tdata,N)
        p̂ = zeros(Tdata,N)
        s = zeros(Tdata,N)
        ŝ = zeros(Tdata,N)
        v = zeros(Tdata,N)
        t = zeros(Tdata,N)
        return new{Ttol,Tdata}(tol,maxiter,r,r̃,p,p̂,s,ŝ,v,t)
    end
end

"""
Solve linear system using preconditioned BiCGStab.
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},bicgstab::BiCGStab{Ttol,Tdata},P=I;
                maxiter::Int=0,tol::Ttol=0.0)::Int where {Ttol,Tdata}

    @unpack r, r̃, p, p̂, s, ŝ, v, t = bicgstab

    # r = b - A⋅x
    mul!(r,A,x)
    @. r = b - r

    # r̃ = r
    copyto!(r̃,r)

    # |b|
    b̄ = norm(b)

    # intializing values
    β    = 0.0
    α    = 0.0
    ρᵢ₋₁ = 1.0
    ρᵢ₋₂ = 1.0
    ω    = 1.0

    if iszero(maxiter)
        maxiter = bicgstab.maxiter
    end

    if iszero(tol)
        tol = bicgstab.tol
    end

    # iterate to convergence
    @fastmath @inbounds for i in 1:maxiter
        ρᵢ₋₂ = ρᵢ₋₁
        ρᵢ₋₁ = dot(r̃,r)
        if ρᵢ₋₁==0.0
            break
        end
        β    = (ρᵢ₋₁/ρᵢ₋₂)*(α/ω)
        @. p = r + β*(p-ω*v)
        ldiv!(p̂,P,p)
        mul!(v,A,p̂)
        α    = ρᵢ₋₁/dot(r̃,v)
        @. s = r - α*v
        ϵ    = norm(s)/b̄
        if ϵ<tol
            @. x += α*p̂
            return i
        end
        ldiv!(ŝ,P,s)
        mul!(t,A,ŝ)
        ω     = dot(t,s)/dot(t,t)
        @. x += α*p̂ + ω*ŝ
        @. r  = s - ω*t
        ϵ     = norm(r)/b̄
        if ϵ<tol
            return i
        end
        if ω==0.0
            break
        end
    end

    return maxiter
end

##########################
## GMRES Implementation ##
##########################

"""
To represent GMRES algorithm.
"""

mutable struct GMRES{Ttol,Tdata} <: IterativeSolver{Ttol,Tdata}

    ndim::Int
    maxiter::Int
    restart::Int
    tol::Ttol
    H::Matrix{Tdata}
    V::Matrix{Tdata}
    r::Vector{Tdata}
    w::Vector{Tdata}
    s::Vector{Tdata}
    y::Vector{Tdata}
    sn::Vector{Tdata}
    cs::Vector{Tdata}

    function GMRES(x::AbstractVector{Tdata}; maxiter::Int=-1,restart::Int=-1,tol::Ttol=1e-4) where {Ttol<:AbstractFloat,Tdata<:Continuous}

        ndim = length(x)
        if maxiter<0
            maxiter = ndim
        end
        if restart<0
            restart = min(20,ndim)
        end
        H  = zeros(Tdata,restart+1,restart)
        V  = zeros(Tdata,ndim,restart+1)
        r  = zeros(Tdata,ndim)
        w  = zeros(Tdata,ndim)
        s  = zeros(Tdata,restart+1)
        y  = zeros(Tdata,restart+1)
        sn = zeros(Tdata,restart+1)
        cs = zeros(Tdata,restart+1)
        return new{Ttol,Tdata}(ndim,maxiter,restart,tol,H,V,r,w,s,y,sn,cs)
    end
end

function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},gmres::GMRES{Ttol,Tdata},M=I;
                maxiter::Int=0,tol::Ttol=0.0)::Int where {Ttol,Tdata}

    H  = gmres.H
    V  = gmres.V
    r  = gmres.r
    w  = gmres.w
    s  = gmres.s
    y  = gmres.y
    sn = gmres.sn
    cs = gmres.cs

    if iszero(maxiter)
        maxiter = gmres.maxiter
    end

    if iszero(tol)
        tol = gmres.tol
    end

    # calculating norm of b
    copyto!(r,b)
    ldiv!(M,r)
    normb = norm(r)
    if normb==0.0
        normb = 1.0
    end

    # calculating residual vector r and getting its norm
    mul!(r,A,x)  # r = A⋅x
    @. r = b - r # r = b - A⋅x
    ldiv!(M,r)   # r = M \ (b - A⋅x)
    β = norm(r)

    # iteration counter
    iter = 0

    # initialize error
    ϵ = β/normb
    if ϵ < tol
        return iter
    end

    @fastmath @inbounds while iter < maxiter
        @. V[:,1] = r/β
        fill!(s,0.0)
        s[1]  = β
        for i in 1:gmres.restart
            iter += 1
            vi    = @view V[:,i]
            mul!(w,A,vi)
            ldiv!(M,w)
            for k in 1:i
                vk     = @view V[:,k]
                H[k,i] = dot(vk,w)
                @. w  -= H[k,i] * vk
            end
            H[i+1,i]    = norm(w)
            @. V[:,i+1] = w / H[i+1,i]
            for k in 1:i-1
                H[k,i], H[k+1,i] = apply_plane_rotation(H[k,i], H[k+1,i], cs[k], sn[k])
            end
            cs[i], sn[i]     = generate_plane_rotation(H[i,i], H[i+1,i])
            H[i,i], H[i+1,i] = apply_plane_rotation(H[i,i], H[i+1,i], cs[i], sn[i])
            s[i], s[i+1]     = apply_plane_rotation(s[i], s[i+1], cs[i], sn[i])
            ϵ                = abs(s[i+1])/normb
            if ϵ < tol
                update!(x,i,H,s,y,V)
                return iter
            end
            if iter==gmres.maxiter
                break
            end
        end
        update!(x,gmres.restart,H,s,y,V)
        mul!(r,A,x)  # r = A⋅x
        @. r = b - r # r = b - A⋅x
        ldiv!(M,r)   # r = M \ (b - A⋅x) = M⁻¹⋅(b - A⋅x)
        β    = norm(r)
        ϵ    = β/normb
        if ϵ < tol
            return iter
        end
    end

    return iter
end

function update!(x::AbstractVector,k::Int,H::Matrix,s::AbstractVector,y::AbstractVector,V::AbstractMatrix)

    copyto!(y,s)

    @fastmath @inbounds for i in k:-1:1
        y[i] /= H[i,i]
        for j in i-1:-1:1
            y[j] -= H[j,i] * y[i]
        end
    end

    @fastmath @inbounds for j in 1:k
        for i in 1:length(x)
            x[i] += y[j] * V[i,j]
        end
    end

    return nothing
end

function apply_plane_rotation(dx::T,dy::T,cs::T,sn::T)::Tuple{T,T} where {T<:Continuous}

    dxp =  cs * dx + conj(sn) * dy
    dyp = -sn * dx +      cs  * dy
    return dxp, dyp
end

function generate_plane_rotation(dx::T,dy::T)::Tuple{T,T} where {T<:Real}

    dz = dx + im*dy
    θ  = angle(dz)
    return cos(θ), sin(θ)
end

function generate_plane_rotation(dx::T,dy::T)::Tuple{T,T} where {T<:Complex}

    c = 0.0
    s = 1.0
    if dx != 0.0
        c = abs(dx)/sqrt(abs2(dx)+abs2(dy))
        s = (dy/dx)*c
    end
    return c, s
end

end


__precompile__()

module EmpiricalOrthogonalFunctions

using Statistics
using LinearAlgebra
using Logging
import Base.truncate

include("eoftypes.jl")
include("eoffuncs.jl")

export EmpiricalOrthogonalFunction,
    pcs,
    eofs,
    correlationmap,
    covariancemap,
    eigenvalues,
    variancefraction,
    totalanomalyvar,
    northtest,
    projectfield,
    reconstruct,
    orthorotation,
    orthorotation!,
    truncate,
    truncate!

end # module


module exactQmodule

using Base, SparseArrays, Arpack, LinearAlgebra
using Combinatorics, KrylovKit, JLD

export AbstractQBasis

#from (1)
export bitSwap, bitCount, findFirstBit, getBitPos, FermionicBitSwap

#from (2)
export QState, shiftStateBasis, getCijMatrix, projectOntoStates, entanglementHalfCut

#from (3)
export QBosonicBasis, addParticleBosonic, constructBosonicState, bosonicHashfunc

#from 4
export fermionicCr, QFermionicBasis, fermionicCrMatrix, addParticleFermionic, constructFermionicState, fermionicHashfunc

#from 5
export loadFermionicCr, saveFermionicCr

abstract type AbstractQBasis end

include("./bitOperations.jl")
include("./Qstate.jl")
include("./bosonic.jl")
include("./fermionic.jl")
include("./preBuild.jl")

end #modlue


using TomoForward
using Images
using Plots
using XfromProjections
using StaticArrays
using Logging
using ImageTransformations
include("../../phantoms/simple_phantoms.jl")

replace_nan(v) = map(x -> isnan(x) ? zero(x) : x, v)

#Translation
function transformation(ϕ)
    v1(x) = x
    v2(x) = x - ϕ*10

    t(x) = SArray{Tuple{2},Float64,1,2}(v1(x[1]),v2(x[2]))
    return t
end

H, W = 50,50#128, 128
img = zeros(H,W)
img = disc_phantom(img,15,25,8)
#img = imresize(img, H, W)
#p_img = plot(Gray.(img), aspect_ratio=:equal, framestyle=:none, title="Original")

function radon_operator(img)
    nangles = 2
    detcount = Int(floor(size(img,1)*1.4))
    angles = rand(0.0:0.001:π, nangles)
    proj_geom = ProjGeom(1.0, detcount, angles)
    A = fp_op_parallel2d_line(proj_geom, size(img, 1), size(img, 2))
    return A
end

frames = zeros(H,W,20)
map(t -> frames[:,:,t+1]=replace_nan(warp(img, transformation(t*0.1), axes(img))), 0:size(frames)[3]-1)
As = map(t -> radon_operator(frames[:,:,t]),1:size(frames)[3])
bs = zeros(size(As[1])[1],size(frames)[3])
map(t -> bs[:,t] = As[t]*vec(frames[:,:,t]), 1:size(frames)[3])
niter=500
u0s = zeros(H,W,size(frames)[3])

mask = disc_phantom(u0s[:,:,1], 25.5, 25.5, 23)
mask = repeat(mask, outer=[1,1,20])

w_tv = 0.1
w_flow  = 0.5

@info "Reconstructing using joint motion estimation and reconstruction"
c=1.1
us_flow = recon2d_tv_primaldual_flow(As, bs, u0s, 20, niter, w_tv, w_flow, c, mask, 0.012)

@info "Reconstruction using tv regularization frame by frame"
us_tv = zeros(H,W,size(frames)[3])
for t = 1:size(frames)[3]
    A = As[t]
    p = bs[:,t]
    us_tv[:,:,t] .= recon2d_tv_primaldual!(us_tv[:,:,t], A, p, niter, w_tv, c)
end

@info "Preparing results in human readable format"
anim = @animate for t=1:size(frames)[3]
    l = @layout [a b c]
    p1 = plot(Gray.(frames[:,:,t]), aspect_ratio=:equal, framestyle=:none, title="Ground truth")
    p2 = plot(Gray.(us_flow[:,:,t]), aspect_ratio=:equal, framestyle=:none, title="Flow")
    p3 = plot(Gray.(us_tv[:,:,t]), aspect_ratio=:equal, framestyle=:none, title="TV")
    plot(p1, p2, p3, layout = l)
end

cwd = @__DIR__
path = normpath(joinpath(@__DIR__, "result"))
!isdir(path) && mkdir(path)
cd(path)
gif(anim, "reconstruction_flow.gif", fps = 1)
cd(cwd)


################################################################################
#                            Planar and Radial Flows                           #
#             Ref: Variational Inference with Normalizing Flows,               #
#               D. Rezende, S. Mohamed(2015) arXiv:1505.05770                  #
################################################################################

###############
# RadialLayer #
###############

mutable struct RadialLayer{T1<:Union{Real, AbstractVector{<:Real}}, T2<:AbstractVector{<:Real}} <: Bijector{1}
    α_::T1
    β::T1
    z_0::T2
end
function Base.:(==)(b1::RadialLayer, b2::RadialLayer)
    return b1.α_ == b2.α_ && b1.β == b2.β && b1.z_0 == b2.z_0
end

function RadialLayer(dims::Int, wrapper=identity)
    α_ = wrapper(randn(1))
    β = wrapper(randn(1))
    z_0 = wrapper(randn(dims))
    return RadialLayer(α_, β, z_0)
end

# all fields are numerical parameters
Functors.@functor RadialLayer

h(α, r) = 1 ./ (α .+ r)     # for radial flow from eq(14)
#dh(α, r) = .- (1 ./ (α .+ r)) .^ 2   # for radial flow; derivative of h()

# An internal version of transform that returns intermediate variables
function _transform(flow::RadialLayer, z::AbstractVecOrMat)
    return _radial_transform(first(flow.α_), first(flow.β), flow.z_0, z)
end
function _radial_transform(α_, β, z_0, z)
    α = LogExpFunctions.log1pexp(α_)            # from A.2
    β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
    if z isa AbstractVector
        r = norm(z .- z_0)
    else
        r = vec(sqrt.(sum(abs2, z .- z_0; dims = 1)))
    end
    transformed = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
    return (transformed = transformed, α = α, β_hat = β_hat, r = r)
end

(b::RadialLayer)(z::AbstractMatrix{<:Real}) = _transform(b, z).transformed
(b::RadialLayer)(z::AbstractVector{<:Real}) = vec(_transform(b, z).transformed)

function forward(flow::RadialLayer, z::AbstractVecOrMat)
    transformed, α, β_hat, r = _transform(flow, z)
    # Compute log_det_jacobian
    d = size(flow.z_0, 1)
    h_ = h(α, r)
    if transformed isa AbstractVector
        T = eltype(transformed)
    else
        T = typeof(vec(transformed))
    end
    log_det_jacobian::T = @. (
        (d - 1) * log(1 + β_hat * h_)
        + log(1 +  β_hat * h_ + β_hat * (- h_ ^ 2) * r)
    )   # from eq(14)
    return (rv = transformed, logabsdetjac = log_det_jacobian)
end

function (ib::Inverse{<:RadialLayer})(y::AbstractVector{<:Real})
    flow = ib.orig
    z0 = flow.z_0
    α = LogExpFunctions.log1pexp(first(flow.α_))            # from A.2
    α_plus_β_hat = LogExpFunctions.log1pexp(first(flow.β))  # from A.2

    # Compute the norm ``r`` from A.2.
    y_minus_z0 = y .- z0
    r = compute_r(y_minus_z0, α, α_plus_β_hat)
    γ = (α + r) / (α_plus_β_hat + r)

    return z0 .+ γ .* y_minus_z0
end

function (ib::Inverse{<:RadialLayer})(y::AbstractMatrix{<:Real})
    flow = ib.orig
    z0 = flow.z_0
    α = LogExpFunctions.log1pexp(first(flow.α_))            # from A.2
    α_plus_β_hat = LogExpFunctions.log1pexp(first(flow.β))  # from A.2

    # Compute the norm ``r`` from A.2 for each column.
    y_minus_z0 = y .- z0
    rs = mapvcat(eachcol(y_minus_z0)) do c
        return compute_r(c, α, α_plus_β_hat)
    end
    γ = reshape((α .+ rs) ./ (α_plus_β_hat .+ rs), 1, :)

    return z0 .+ γ .* y_minus_z0
end

"""
    compute_r(y_minus_z0::AbstractVector{<:Real}, α, α_plus_β_hat)

Compute the unique solution ``r`` to the equation
```math
\\|y_minus_z0\\|_2 = r \\left(1 + \\frac{α_plus_β_hat - α}{α + r}\\right)
```
subject to ``r ≥ 0`` and ``r ≠ α``.

Since ``α > 0`` and ``α_plus_β_hat > 0``, the solution is unique and given by
```math
r = (\\sqrt{(α_plus_β_hat - γ)^2 + 4 α γ} - (α_plus_β_hat - γ)) / 2,
```
where ``γ = \\|y_minus_z0\\|_2``. For details see appendix A.2 of the reference.

# References

D. Rezende, S. Mohamed (2015): Variational Inference with Normalizing Flows.
arXiv:1505.05770
"""
function compute_r(y_minus_z0::AbstractVector{<:Real}, α, α_plus_β_hat)
    γ = norm(y_minus_z0)
    a = α_plus_β_hat - γ
    r = (sqrt(a^2 + 4 * α * γ) - a) / 2
    return r
end

logabsdetjac(flow::RadialLayer, x::AbstractVecOrMat) = forward(flow, x).logabsdetjac


using ClassicalOrthogonalPolynomials, QuasiArrays, InfiniteArrays
using BlockBandedMatrices, BlockArrays, LazyArrays, Test, FillArrays, InfiniteLinearAlgebra, ArrayLayouts, LazyBandedMatrices
import BlockBandedMatrices: _BlockBandedMatrix
import ClassicalOrthogonalPolynomials: Inclusion, ShuffledRFFT
import QuasiArrays: MulQuasiArray

@testset "Fourier" begin
    @testset "ShuffledRFFT" begin
        ret = randn(3,5)
        p = ShuffledRFFT{Float64}(size(ret,1))
        P = ShuffledRFFT{Float64}(size(ret),1)
        @test size(p) == (3,)
        @test size(P) == (3,5)
        for k = 1:size(ret,2)
            @test (P * ret)[:,k] ≈ p * ret[:,k]
        end

        p = ShuffledRFFT{Float64}(size(ret,2))
        P = ShuffledRFFT{Float64}(size(ret),2)
        for k = 1:size(ret,1)
            @test (P * ret)[k,:] ≈ p * ret[k,:]
        end
    end

    @testset "Evaluation" begin
        F = Fourier()

        @test F == F

        @test axes(F,2) isa BlockedUnitRange{InfiniteArrays.InfStepRange{Int,Int}}
        @test axes(F) isa Tuple{<:Inclusion,BlockedUnitRange{InfiniteArrays.InfStepRange{Int,Int}}}
        @test axes(F,2)[Block(2)] == 2:3
        @test F[0.1,1] == 1.0
        @test F[0.1,2] == sin(0.1)
        @test F[0.1,1:4] == [1,sin(0.1),cos(0.1),sin(2*0.1)]
        @test F[0.1,Block(4)] == [sin(3*0.1),cos(3*0.1)]
        @test F[0.1,Block.(1:3)] == [1,sin(0.1),cos(0.1),sin(2*0.1),cos(2*0.1)]

        u = F * PseudoBlockVector([[1,2,3]; zeros(∞)], (axes(F,2),));
        @test u[0.1] == 1 + 2sin(0.1) + 3cos(0.1)
    end

    @testset "Transform" begin
        F = Fourier()
        θ = axes(F,1)
        @test F[:,Base.OneTo(5)] \ cos.(θ) ≈ [0,0,1,0,0]
        @test F[:,Block.(Base.OneTo(5))] \ cos.(θ) ≈ [0,0,1,0,0,0,0,0,0]

        @test (F \ cos.(θ))[Block(2)] ≈ [0,1]
        u = F * (F \ exp.(cos.(θ)))
        @test u[0.1] ≈ exp(cos(0.1))
        @test F[:,Base.OneTo(5)] \ [cos.(θ) sin.(θ)] ≈ [0 0; 0 1; 1 0; 0 0; 0 0]
        @test F[:,Block.(Base.OneTo(3))] \ [cos.(θ) sin.(θ)] ≈ [0 0; 0 1; 1 0; 0 0; 0 0]
        U = F / F \ [exp.(cos.(θ)) cos.(cos.(θ))]
        @test U[0.1,:] ≈ [exp(cos(0.1)),cos(cos(0.1))]
        @test U[[0.1,0.2],:] ≈ [exp(cos(0.1)) cos(cos(0.1)); exp(cos(0.2)) cos(cos(0.2))]
        @test U[0.1,1] ≈ exp(cos(0.1))
        @test U[[0.1,0.2],1] ≈ exp.(cos.([0.1,0.2]))
    end

    @testset "Derivative" begin
        F = Fourier()
        @test F\F === Eye((axes(F,2),))
        @test (F'F)[1:10,1:10] == Diagonal(Vcat(2π,Fill(1.0π,9)))
        @test blockisequal(axes(F'F), (axes(F,2),axes(F,2)))
        D = Derivative(axes(F,1))
        D̃ = (D*F).args[2]
        @test (F\F)*D̃ isa BlockArray
        @test (F \ (D*F))[Block.(1:3),Block.(1:3)] == [0 0 0 0 0; 0 0.0 -1 0 0; 0 1 0 0 0; 0 0 0 0 -2; 0 0 0 2 0]

        u = F * PseudoBlockVector([[1,2,3,4,5]; zeros(∞)], (axes(F,2),));
        @test blockisequal(axes(D̃,2),axes(u.args[2],1))
        @test (D*u)[0.1] ≈ 2cos(0.1) - 3sin(0.1) + 8cos(2*0.1) - 10sin(2*0.1)
    end

    @testset "cos(θ)" begin
        F = Fourier()
        θ = axes(F,1)
        c,s = cos.(θ),sin.(θ)
        @test c .* F isa MulQuasiArray
        @test s .* F isa MulQuasiArray
        @test cos.(θ) .* F isa MulQuasiArray
        @test sin.(θ) .* F isa MulQuasiArray
        X,Y = F \ (c .* F),F \ (s .* F)
        N = 10
        XY = X*Y
        @test XY isa MulMatrix
        @test MemoryLayout(XY) isa LazyBandedMatrices.ApplyBlockBandedLayout{typeof(*)}
        @test MemoryLayout(view(XY,Block.(1:N),Block.(1:N))) isa LazyBandedMatrices.ApplyBlockBandedLayout{typeof(*)}
        @test XY[Block.(1:N),Block.(1:N)] isa BlockSkylineMatrix
        @test blockbandwidths(XY[Block.(1:N),Block.(1:N)]) == (2,2)
        @test XY[Block.(1:N),Block.(1:N)] ≈ ApplyMatrix(*,Y,X)[Block.(1:N),Block.(1:N)]
        @test ApplyMatrix(*,X,X)[Block.(1:N),Block.(1:N)] + ApplyMatrix(*,Y,Y)[Block.(1:N),Block.(1:N)] ≈ Eye(19)
    end
end


#=

  The BUGS Book
  Example 3.5.1 Heart transplants: learning from data

  ~/jags/bugs_book_3_5_1.R
  """
  model {
    yT ~ dbin(pT, nT)
    pT ~ dunif(0, 1)
    for (i in 1:8) {
      sP[i] ~ dexp(theta)
    }
    theta ~ dgamma(0.001, 0.001)
    surv.t <- pT/theta # expected survival with transplant
    Is <- surv.t - 2
  }

  Data: list(yT=8, nT=10, sP=c(2,3,4,4,6,7,10,12))

  Output:
         Mean      SD  Naive SE Time-series SE
Is     3.1442 2.28250 0.0080699      0.0085034
pT     0.7500 0.12014 0.0004248      0.0005856
surv.t 5.1442 2.28250 0.0080699      0.0085034
theta  0.1667 0.05899 0.0002086      0.0002086
  """

  Note that most value are fixed/derived.

  Summary Statistics
  parameters      mean       std   naive_se      mcse       ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64   Float64   Float64       Float64

       theta    0.1473    0.0000     0.0000    0.0000   20.5530    0.9999        6.1813
          pT    0.8623    0.0000     0.0000    0.0000   20.5530    0.9999        6.1813
          yT    8.0000    0.0000     0.0000    0.0000       NaN       NaN           NaN
      surv_t    5.8559    0.0000     0.0000    0.0000       NaN       NaN           NaN
          Is    3.8559    0.0000     0.0000    0.0000   20.5530    0.9999        6.1813

  Cf ~/webppl/bugs_book_3_5_1.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function bugs_book_3_5_1(sPs)
    n = length(sPs)
    nT = 10

    theta ~ Gamma(0.001,1/0.001)
    pT ~ Uniform(0,1)
    yT ~ Binomial(nT,pT)

    surv_t ~ Dirac(pT/theta) # expected survival with transplant
    Is ~ Dirac(surv_t - 2.0)

    true ~ Dirac(yT == 8)
    sPs ~ filldist(Exponential(theta),n)

end

sPs = [2,3,4,4,6,7,10,12]
model = bugs_book_3_5_1(sPs)

num_chains = 4

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

# chns = sample(model, NUTS(), 10_000)
chns = sample(model, HMC(0.1,10), 10_000)

display(chns)
# display(plot(chns))


struct Taylens{N,T<:Real} <: ImplicitOp{Basis,Spin,Pix}
    # pixel remapping
    i::Matrix{Int}
    j::Matrix{Int}

    # residual displacement
    rx::Matrix{T}
    ry::Matrix{T}

    # precomputed quantities
    kα::Dict{Any,Matrix{Complex{T}}}
    xα::Dict{Any,Matrix{T}}

end

function Taylens(ϕ::FlatS0{P,T}, N) where {P,T}

    @unpack kx,ky,Nx,Ny,Δx = fieldinfo(ϕ)

    # total displacement
    d = ∇*ϕ
    dx, dy = d[1][:Ix], d[2][:Ix]

    # nearest pixel displacement
    indexwrap(ind::Int64, uplim)  = mod(ind - 1, uplim) + 1
    di, dj = (round.(Int,d/Δx) for d=(dx,dy))
    i = indexwrap.(di .+ (1:Nx)', Nx)
    j = indexwrap.(dj .+ (1:Ny) , Ny)

    # residual displacement
    rx, ry = ((d - i.*Δx) for (d,i)=[(dx,di),(dy,dj)])

    # precomputation
    kα = Dict{Any,Matrix{Complex{T}}}()
    xα = Dict{Any,Matrix{T}}()
    for n in 1:N, α₁ in 0:n
        kα[n,α₁] = im ^ n .* kx' .^ α₁ .* ky[1:Ny÷2+1] .^ (n - α₁)
        xα[n,α₁] = rx .^ α₁ .* ry .^ (n - α₁) ./ factorial(α₁) ./ factorial(n - α₁)
    end

    Taylens{N,T}(i,j,rx,ry,kα,xα)
end

# our implementation of Taylens
function *(L::Taylens{N}, f::FlatS0{P,T}) where {N,P,T}

    intlense(fx) = getindex.(Ref(fx), L.j, L.i)
    fl = f[:Il]

    # lens to the nearest whole pixel
    Lfx = intlense(f[:Ix])

    # add in Taylor series correction
    for n in 1:N, α₁ in 0:n
        Lfx .+= L.xα[n,α₁] .* intlense(fieldinfo(f).FFT \ (L.kα[n,α₁] .* fl))
    end

    FlatMap{P}(Lfx)
end

function *(L::Taylens, f::FieldTuple)
    Łf = Ł(f)
    F = typeof(Łf)
    F(map(f->L*f, Łf.fs))
end


include("euler/euler.jl")
using .Game: dice_combinations
using .Math: fast_max
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(a::Tuple{Int, Int}, b::Tuple{Int, Int})::Float64
    limit = fast_max(a[1] * a[2], b[1] * b[2])
    looses = zeros(Int, limit)
    result = 0
    for i ∈ b[1]:b[1] * b[2] looses[i] = dice_combinations(i, b[1], b[2]) + looses[i - 1] end
    for i ∈ a[1]:a[1] * a[2] result += dice_combinations(i, a[1], a[2]) * looses[i - 1] end
    return round(result / a[2] ^ a[1] / b[2] ^ b[1], digits=7)
end

compute((9, 4), (6, 6))

@benchmark compute((9, 4), (6, 6))


using ModelingToolkit, Catalyst #, LinearAlgebra
#using DiffEqBase, DiffEqJump
#using Plots, SpecialFunctions

## Parameter
N = 10                       # maximum cluster size
Vₒ = (4π/3)*(10e-06*100)^3   # volume of a monomers in cm³
Nₒ = 1e-06/Vₒ                # initial conc. = (No. of init. monomers) / bulk volume
uₒ = 10000                   # No. of monomers initially
V = uₒ/Nₒ                    # Bulk volume of system in cm³

integ(x) = Int(floor(x))
n        = integ(N/2)
nr       = N%2 == 0 ? (n*(n + 1) - n) : (n*(n + 1)) # No. of forward reactions

# possible pairs of reactant multimers
pair = []
for i = 2:N
    push!(pair,[1:integ(i/2)  i .- (1:integ(i/2))])
end
pair = vcat(pair...)
vᵢ = @view pair[:,1]   # Reactant 1 indices
vⱼ = @view pair[:,2]   # Reactant 2 indices
volᵢ = Vₒ*vᵢ           # cm⁻³
volⱼ = Vₒ*vⱼ           # cm⁻³
sum_vᵢvⱼ = @. vᵢ + vⱼ  # Product index

# set i to  1 for additive kernel, 2  for constant
i = 1
if i==1
    B = 1.53e03                # s⁻¹
    kv = @. B*(volᵢ + volⱼ)/V  # dividing by volume as its a bi-molecular reaction chain
elseif i==2
    C = 1.84e-04               # cm³ s⁻¹
    kv = fill(C/V, nr)
end

# state variables are X, pars stores rate parameters for each rx
@parameters t
@variables k[1:nr]  X[1:N](t)
pars = Pair.(collect(k), kv)

# time-span
if i == 1
    tspan = (0. ,2000.)
elseif i == 2
    tspan = (0. ,350.)
end

 # initial condition of monomers
u₀    = zeros(Int64, N)
u₀[1] = uₒ
u₀map = Pair.(collect(X), u₀)   # map variable to its initial value

# vector to store the Reactions in
rx = []
for n = 1:nr
    # for clusters of the same size, double the rate
    if (vᵢ[n] == vⱼ[n])
        push!(rx, Reaction(k[n], [X[vᵢ[n]]], [X[sum_vᵢvⱼ[n]]], [2], [1]))
    else
        push!(rx, Reaction(k[n], [X[vᵢ[n]], X[vⱼ[n]]], [X[sum_vᵢvⱼ[n]]],
                           [1, 1], [1]))
    end
end
@named rs = ReactionSystem(rx, t, collect(X), collect(k))
Catalyst.Graph(rs)


function Base.:^(a::UncertainValueDataset, b::UncertainValueDataset; n = 30000)
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    UncertainValueDataset([a[i] ^ b[i] for i = 1:N])
end

function Base.:^(a::T, b::UncertainValueDataset; n = 30000) where {T <: Number}
    N = length(a)
    UncertainValueDataset([a ^ b[i] for i = 1:N])
end

function Base.:^(a::Vector{T}, b::UncertainValueDataset; n = 30000) where {T <: Number}
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    UncertainValueDataset([a[i] ^ b[i] for i = 1:N])
end

function Base.:^(a::UncertainValueDataset, b::T; n = 30000) where {T <: Number}
    N = length(a)

    UncertainValueDataset([a[i] ^ b for i = 1:N])
end

function Base.:^(a::UncertainValueDataset, b::Vector{T}; n = 30000) where {T <: Number}
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    UncertainValueDataset([a[i] ^ b[i] for i = 1:N])
end


# States.jl
# Define electronic and nuclear states

"""
Function to define diabatic states as gaussian functions of width σ
--------------
R_n = bond length
σ = width of states
--------------
phi_l = gaussian centered on left site
phi_r = gaussian centered on right site
"""
function diabatic_state(R_n::Float64,σ::Float64=0.998)
    gaussian = function (R) return exp(-R^2/(2*σ^2)) end
    gaussian2(R) = gaussian(R)^2
    A = riemann(gaussian2, -10,10,1000)
    A = 1/sqrt(A)
    phi(R) = A*gaussian(R-R_n)
    return phi
end

"""
Overlap of diabatic states to determine coupling strength
--------------
phi_l = gaussian centered on left site
phi_r = gaussian centered on right site
--------------
J_lr = coupling strength to be used in diabatic Hamiltonian
"""
function overlap_phi(phi_l::Function, phi_r::Function)
    overlap_phi(R) = conj(phi_l(R))*phi_r(R)
    J_lr = riemann(overlap_phi, -10,10,1000)
    return J_lr
end

"""
Function for the derivative of a gaussian function
    (to be used in NACV when the derivative of phi w.r.t R is needed)
--------------
R_n = bond length
σ = width of states
--------------
d_phi = gaussian derivative
"""
function diabatic_derivative(σ::Float64=0.374)
    dphi = function (R) return (R)*exp(-(R)^2/(2*σ^2))/σ^2 end
end

"""
Function to define the adiabatic states from the diabatic states and
transformation Matrix
------------
U_nk = Unitary transformation matrix
phi_l = diabatic state on left site
phi_r = diabatic state on right site
------------
psi_g = adiabatic lower state
phi_e = adiabatic upper state
"""
function adiabatic_states(phi_l::Function, phi_r::Function, U_nk::Array)
    psi_g = function (R) return U_nk[1]*phi_l(R) + U_nk[2]*phi_r(R) end
    psi_e = function (R) return U_nk[3]*phi_l(R) + U_nk[4]*phi_r(R) end
    return psi_g, psi_e
end

"""
Function to define the non-adiabatic coupling vectors for all possible transisions
--------------
phi_l = diabatic state on left site
phi_r = diabatic state on right site
dphi = generic derivative of gaussian function
U_nk = Unitary transformation matrix
R_n = bond length
dR_l = change of nuclear coordinate on the left site
dR_r = change of nuclear coordinate on the right site
dt = time step this change has occured in
--------------
d_ge = Non-adiabatic coupling vector from lower to upper adiabatic states
d_eg = Non-adiabatic coupling vector from upper to lower adiabatic states
d_gg = Non-adiabatic coupling vector from lower to lower adiabatic states
d_ee = Non-adiabatic coupling vector from upper to upper adiabatic states
"""
function NACV(phi_l::Function, phi_r::Function, dphi::Function, U_nk::Array, R_n::Float64, dR_l, dR_r, dt)
    dphi_l = function (R) return dphi((R)) end
    dphi_r = function (R) return dphi((R-R_n)) end

    # perform integrations for each pair of diabatic states
    int_ll = function (R) return conj(phi_l(R))*dphi_l(R) end
    ll = riemann(int_ll,-10,10,1000)
    int_lr = function (R) return conj(phi_l(R))*dphi_r(R) end
    lr = riemann(int_lr,-10,10,1000)
    int_rl = function (R) return conj(phi_r(R))*dphi_l(R) end
    rl = riemann(int_rl,-10,10,1000)
    int_rr = function (R) return conj(phi_r(R))*dphi_r(R) end
    rr = riemann(int_rr,-10,10,1000)

    l_ge = dR_l*(conj(U_nk[1]*ll + U_nk[2]*rl)*U_nk[3])/dt
    r_ge = dR_r*(conj(U_nk[1]*lr + U_nk[2]*rr)*U_nk[4])/dt
    d_ge = l_ge+r_ge

    l_eg = dR_l*(conj(U_nk[3]*ll + U_nk[4]*rl)*U_nk[1])/dt
    r_eg = dR_r*(conj(U_nk[3]*lr + U_nk[4]*rr)*U_nk[2])/dt
    d_eg = l_eg+r_eg

    l_gg = dR_l*(conj(U_nk[1]*ll + U_nk[2]*rl)*U_nk[1])/dt
    r_gg = dR_r*(conj(U_nk[1]*lr + U_nk[2]*rr)*U_nk[2])/dt
    d_gg = l_gg+r_gg

    l_ee = dR_l*(conj(U_nk[3]*ll + U_nk[4]*rl)*U_nk[3])/dt
    r_ee = dR_r*(conj(U_nk[3]*lr + U_nk[4]*rr)*U_nk[4])/dt
    d_ee = l_ee+r_ee

    return d_ge,d_eg,d_gg,d_ee, [l_ge*dt/dR_l, r_ge*dt/dR_r, l_eg*dt/dR_l, r_eg*dt/dR_r]
end

"""
Function to determine adiabatic coefficients from diabatic coefficients
    and unitary transformation matrix
-------------
U_nk = Unitary Transformation Matrix
c1 = diabatic coefficient of left site
c2 = diabatic coefficient of right site
-------------
ag = adiabatic coefficient of ground state
ae = adiabatic coefficient of excited state
"""
function a_mn(U_nk::Array,cl::Complex{Float64},cr::Complex{Float64})
    ag = (cl*U_nk[4] - cr*U_nk[3])/(U_nk[1]*U_nk[4]-U_nk[2]*U_nk[3])
    ae = (cl*U_nk[2] - cr*U_nk[1])/(U_nk[2]*U_nk[3]-U_nk[1]*U_nk[4])
    return ag,ae
end

"""
Function to determine the probability of switching states
------------
a_m = adiabatic coefficient of state m
a_n = adiabatic coefficient of state n
d_mn = Non-adiabatic coupling vector for transition from state n to state m
dt = time step
------------
g_mn = probability of a transition from adiabatic surface m to adiabatic surface n
"""
function g_mn(a_m::Complex{Float64}, a_n::Complex{Float64}, d_mn::Float64, dt::Float64)
    g_mn =  -2*real(dt*(conj(a_n)*a_m*d_mn)/(conj(a_m)*a_m))
    return g_mn
end


using GeometricFlux
using Flux
using Flux: onehotbatch, onecold, logitcrossentropy, throttle
using Flux: @epochs
using JLD2  # use v0.1.2
using Statistics: mean
using SparseArrays
using LightGraphs.SimpleGraphs
using LightGraphs: adjacency_matrix
using CUDA

@load "data/cora_features.jld2" features
@load "data/cora_labels.jld2" labels
@load "data/cora_graph.jld2" g

num_nodes = 2708
num_features = 1433
hidden = 16
target_catg = 7
epochs = 20

## Preprocessing data
train_X = Float32.(features) |> gpu  # dim: num_features * num_nodes
train_y = Float32.(labels) |> gpu  # dim: target_catg * num_nodes

adj_mat = Matrix{Float32}(adjacency_matrix(g)) |> gpu

## Model
model = Chain(GCNConv(adj_mat, num_features=>hidden, relu),
              Dropout(0.5),
              GCNConv(adj_mat, hidden=>target_catg),
              softmax) |> gpu

## Loss
loss(x, y) = logitcrossentropy(model(x), y)
accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

## Training
ps = Flux.params(model)
train_data = [(train_X, train_y)]
opt = ADAM(0.05)
evalcb() = @show(accuracy(train_X, train_y))

@epochs epochs Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))


using NGWP, Plots, LightGraphs, JLD, MTSG, MAT

barbara = JLD.load("../datasets/barbara_gray_matrix.jld", "barbara")

## Build weighted graph
G, L, X = SunFlowerGraph(N = 400); N = nv(G)
W = 1.0 * adjacency_matrix(G)
if runapprox
    Q = incidence_matrix(G; oriented = true)
    edge_weight = [e.weight for e in edges(G)]

    ## eigenvectors of L(G)
    𝛌, 𝚽 = eigen(Matrix(L))
    standardize_eigenvectors!(𝚽)

    ## eigenvectors of Lsym(G)
    deg = sum(W, dims = 1)[:]  # weighted degree vector
    Lsym = diagm(deg.^(-1/2)) * (diagm(deg) - W) * diagm(deg.^(-1/2))
    𝛌sym, 𝚽sym = eigen(Lsym)
    standardize_eigenvectors!(𝚽sym)

    ## Build Dual Graph by DAG metric
    distDAG = eigDAG_Distance(𝚽, Q, N; edge_weight = edge_weight)
    Gstar_Sig = dualgraph(distDAG)
    G_Sig = GraphSig(W, xy = X)
    GP_dual = partition_tree_fiedler(Gstar_Sig; swapRegion = false)
    GP_primal = pairclustering(𝚽, GP_dual)
    jmax = size(GP_dual.rs, 2) - 1  # zero-indexed

    if allNGWPs
        # 54.986524 seconds (1.19 M allocations: 25.127 GiB, 2.36% gc time)
        @time VM_NGWP = vm_ngwp(𝚽, GP_dual)
        # 0.611176 seconds (225.41 k allocations: 844.488 MiB, 11.71% gc time)
        @time PC_NGWP = pc_ngwp(𝚽, GP_dual, GP_primal)
        # 50.939912 seconds (7.67 M allocations: 28.051 GiB, 3.11% gc time)
        @time LP_NGWP = lp_ngwp(𝚽, Gstar_Sig.W, GP_dual; ϵ = 0.3)
    end

    ## Build Dual Graph by DAG metric (Lsym)
    distDAG_Lsym = eigDAG_Distance(𝚽sym, Q, N; edge_weight = edge_weight)
    Gstar_Sig_Lsym = dualgraph(distDAG_Lsym)
    GP_dual_Lsym = partition_tree_fiedler(Gstar_Sig_Lsym; swapRegion = false)
    GP_primal_Lsym = pairclustering(𝚽sym, GP_dual_Lsym)
    jmax_Lsym = size(GP_dual_Lsym.rs, 2) - 1

    if allNGWPs
        # 61.328441 seconds (1.34 M allocations: 29.058 GiB, 2.30% gc time)
        @time VM_NGWP_Lsym = vm_ngwp(𝚽sym, GP_dual_Lsym)
        # 0.605507 seconds (237.87 k allocations: 876.689 MiB, 10.83% gc time)
        @time PC_NGWP_Lsym = pc_ngwp(𝚽sym, GP_dual_Lsym, GP_primal_Lsym)
        # 77.918189 seconds (2.12 M allocations: 42.872 GiB, 2.91% gc time)
        @time LP_NGWP_Lsym = lp_ngwp(𝚽sym, Gstar_Sig_Lsym.W, GP_dual_Lsym; ϵ = 0.3)
    end
else
    using VoronoiDelaunay, VoronoiCells, GeometricalPredicates
    ## Voronoi tessellation
    width_x = maximum(abs.(X[:, 1])) * 2; width_y = maximum(abs.(X[:, 2])) * 2;
    width = VoronoiDelaunay.max_coord - VoronoiDelaunay.min_coord
    center_coord = (VoronoiDelaunay.min_coord + VoronoiDelaunay.max_coord)/2
    X_transform = zeros(N,2)
    for i in 1:N
        X_transform[i,:] = X[i,:] ./ [width_x/width, width_y/width] + [center_coord, center_coord]
    end

    pts = [Point2D(X_transform[i,1], X_transform[i,2]) for i in 1:N]
    tess = DelaunayTessellation(N)
    push!(tess, pts)
end


using JuMP
using GLPKMathProgInterface
using Ipopt

n = 3;
lowRes = -20;
hiRes  = 20;
m = Model(solver=IpoptSolver())
@defVar(m, lowRes <= α[i=1:n]  <=hiRes)
@defVar(m, lowRes <= β1[i=1:n] <=hiRes)
@defVar(m, lowRes <= β2[i=1:n] <=hiRes)
@defVar(m, lowRes <= β3[i=1:n] <=hiRes)
@defVar(m, lowRes <= β4[i=1:n] <=hiRes)

@defVar(m, lowRes <= A011 <= hiRes)
@defVar(m, lowRes <= A012 <= hiRes)
@defVar(m, lowRes <= A013 <= hiRes)
@defVar(m, lowRes <= A021 <= hiRes)
@defVar(m, lowRes <= A022 <= hiRes)
@defVar(m, lowRes <= A023 <= hiRes)
@defVar(m, lowRes <= A031 <= hiRes)
@defVar(m, lowRes <= A032 <= hiRes)
@defVar(m, lowRes <= A033 <= hiRes)

@defVar(m, lowRes <= A111 <= hiRes)
@defVar(m, lowRes <= A112 <= hiRes)
@defVar(m, lowRes <= A113 <= hiRes)
@defVar(m, lowRes <= A121 <= hiRes)
@defVar(m, lowRes <= A122 <= hiRes)
@defVar(m, lowRes <= A123 <= hiRes)
@defVar(m, lowRes <= A131 <= hiRes)
@defVar(m, lowRes <= A132 <= hiRes)
@defVar(m, lowRes <= A133 <= hiRes)

@defVar(m, lowRes <= B011 <= hiRes)
@defVar(m, lowRes <= B012 <= hiRes)
@defVar(m, lowRes <= B013 <= hiRes)
@defVar(m, lowRes <= B021 <= hiRes)
@defVar(m, lowRes <= B022 <= hiRes)
@defVar(m, lowRes <= B023 <= hiRes)
@defVar(m, lowRes <= B031 <= hiRes)
@defVar(m, lowRes <= B032 <= hiRes)
@defVar(m, lowRes <= B033 <= hiRes)

@defVar(m, lowRes <= B111 <= hiRes)
@defVar(m, lowRes <= B112 <= hiRes)
@defVar(m, lowRes <= B113 <= hiRes)
@defVar(m, lowRes <= B121 <= hiRes)
@defVar(m, lowRes <= B122 <= hiRes)
@defVar(m, lowRes <= B123 <= hiRes)
@defVar(m, lowRes <= B131 <= hiRes)
@defVar(m, lowRes <= B132 <= hiRes)
@defVar(m, lowRes <= B133 <= hiRes)

A0 = [A011 A012
     A021 A022]
A1 = [A111 A112
     A121 A122]
B0 = [B011 B012
      B021 B022]
B1 = [B111 B112
      B121 B122]
e = [1 1]'

#@setObjective(m, Max, 0)
@addConstraints(m,begin
                  sum{α[i], i=1:n}==1
                  sum{β1[i], i=1:n}==1
                  sum{β1[i], i=1:n}==1
                  sum{β2[i], i=1:n}==0
                  sum{β3[i], i=1:n}==0
                  sum{β4[i], i=1:n}==0
                  sum{β4[i], i=1:n}==0
                end)
@addNLConstraints(m,begin
                  B111*β1[1]+B112*β1[1]+B113*β1[1] + B121*β1[2]+B122*β1[2]+B123*β1[2] + B131*β1[3]+B132*β1[3]+B133*β1[3] ==0
                  B111*β2[1]+B112*β2[1]+B113*β2[1] + B121*β2[2]+B122*β2[2]+B123*β2[2] + B131*β2[3]+B132*β2[3]+B133*β2[3] ==1
                  B111*β3[1]+B112*β3[1]+B113*β3[1] + B121*β3[2]+B122*β3[2]+B123*β3[2] + B131*β3[3]+B132*β3[3]+B133*β3[3] ==0
                  B111*β4[1]+B112*β4[1]+B113*β4[1] + B121*β4[2]+B122*β4[2]+B123*β4[2] + B131*β4[3]+B132*β4[3]+B133*β4[3] ==0
                  A011*α[1]+A012*α[1]+A013*α[1]    + A021*α[2]+A022*α[2]+A023*α[2]    + A031*α[3]+A032*α[3]+A033*α[3] ==1//2
                  B011*α[1]+B012*α[1]+B013*α[1]    + B021*α[2]+B022*α[2]+B023*α[2]    + B031*α[3]+B032*α[3]+B033*α[3] ==1
                  A111*β1[1]+A112*β1[1]+A113*β1[1] + A121*β1[2]+A122*β1[2]+A123*β1[2] + A131*β1[3]+A132*β1[3]+A133*β1[3]==1
                  A111*β2[1]+A112*β2[1]+A113*β2[1] + A121*β2[2]+A122*β2[2]+A123*β2[2] + A131*β2[3]+A132*β2[3]+A133*β2[3]==0
                  A111*β3[1]+A112*β3[1]+A113*β3[1] + A121*β3[2]+A122*β3[2]+A123*β3[2] + A131*β3[3]+A132*β3[3]+A133*β3[3]==-1
                  A111*β4[1]+A112*β4[1]+A113*β4[1] + A121*β4[2]+A122*β4[2]+A123*β4[2] + A131*β4[3]+A132*β4[3]+A133*β4[3]==0
                end)

@addNLConstraint(m,α[1]  * (B011+B012+B013)*(B011+B012+B013)  + α[2]  * (B021+B022+B123)*(B021+B022+B123) + α[3]  * (B031+B032+B133)*(B031+B032+B133)  == 3//2)
@addNLConstraint(m,β1[1] * (B111+B112+B013)*(B111+B112+B013)  + β1[2] * (B121+B122+B123)*(B121+B122+B123) + β1[3] * (B131+B132+B133)*(B131+B132+B133)  == 1)
@addNLConstraint(m,β2[1] * (B111+B112+B013)*(B111+B112+B013)  + β2[2] * (B121+B122+B123)*(B121+B122+B123) + β2[3] * (B131+B132+B133)*(B131+B132+B133)  == 0)
@addNLConstraint(m,β3[1] * (B111+B112+B013)*(B111+B112+B013)  + β3[2] * (B121+B122+B123)*(B121+B122+B123) + β3[3] * (B131+B132+B133)*(B131+B132+B133)  == -1)
@addNLConstraint(m,β4[1] * (B111+B112+B013)*(B111+B112+B013)  + β4[2] * (B121+B122+B123)*(B121+B122+B123) + β4[3] * (B131+B132+B133)*(B131+B132+B133)  == 2)

@addNLConstraint(m,β1[1] * (B111*(B111+B112+B113) + B112*(B121+B122+B123) + B113*(B131+B132+B133))  + β1[2] * (B121*(B111+B112+B113) + B122*(B121+B122+B123) + B123*(B131+B132+B133))  + β1[3] * (B131*(B111+B112+B113) + B132*(B121+B122+B123) + B133*(B131+B132+B133))  == 0)
@addNLConstraint(m,β2[1] * (B111*(B111+B112+B113) + B112*(B121+B122+B123) + B113*(B131+B132+B133))  + β2[2] * (B121*(B111+B112+B113) + B122*(B121+B122+B123) + B123*(B131+B132+B133))  + β2[3] * (B131*(B111+B112+B113) + B132*(B121+B122+B123) + B133*(B131+B132+B133))  == 0)
@addNLConstraint(m,β3[1] * (B111*(B111+B112+B113) + B112*(B121+B122+B123) + B113*(B131+B132+B133))  + β3[2] * (B121*(B111+B112+B113) + B122*(B121+B122+B123) + B123*(B131+B132+B133))  + β3[3] * (B131*(B111+B112+B113) + B132*(B121+B122+B123) + B133*(B131+B132+B133))  == 0)
@addNLConstraint(m,β4[1] * (B111*(B111+B112+B113) + B112*(B121+B122+B123) + B113*(B131+B132+B133))  + β4[2] * (B121*(B111+B112+B113) + B122*(B121+B122+B123) + B123*(B131+B132+B133))  + β4[3] * (B131*(B111+B112+B113) + B132*(B121+B122+B123) + B133*(B131+B132+B133))  == 1)

@addNLConstraint(m,(1/2)*((A111*(B011+B012+B013) + A112*(B021+B022+B023) + A113*(B031+B032+B033))*β1[1] +
      (A121*(B011+B012+B013) + A122*(B021+B022+B023) + A123*(B031+B032+B033))β1[2] +
      (A131*(B011+B012+B013) + A132*(B021+B022+B023) + A133*(B031+B032+B033))β1[3]) +
      (1/3)*((A111*(B011+B012+B013) + A112*(B021+B022+B023) + A113*(B031+B032+B033))*β3[1] +
            (A121*(B011+B012+B013) + A122*(B021+B022+B023) + A123*(B031+B032+B033))β3[2] +
            (A131*(B011+B012+B013) + A132*(B021+B022+B023) + A133*(B031+B032+B033))β3[3]) == 0)

print(m)

status = solve(m)

println("Objective value: ", getObjectiveValue(m))
println("β1 = ", getValue(β1))
println("β2 = ", getValue(β2))
println("β3 = ", getValue(β3))
println("β4 = ", getValue(β4))


@doc raw"""
    Coeffgamm2Simple(k, T, m)

The coefficient before the two point function flow.
"""
function Coeffgamm2Simple(k, T, m)
    (
        k * (
            -(coth(Epi(k, m) / (2 * T)) / Epi(k, m)^3) -
            csch(Epi(k, m) / (2 * T))^2 / (2 * T * Epi(k, m)^2)
        )
    ) / (16 * pi^2)
end

@doc raw"""
    VImSimple(p0, ps, q0, k, m, T, Npi, lam4pik)

compute $\mathrm{Im}V_k(q_0)$, the `k` dependence of $\lambda_{4\pi}$ and $m$ are neglected.

`VImSimple` only contains $V(q_0)$, no $V(-q_0)$

`VImSimple` contains type-1 delta function

`VImSimple` doesn't contains type-2 delta function
"""
function VImSimple(p0, ps, q0, k, m, T, Npi, lam4pik, UVScale)
    lam4pik^2 *
    (2 + Npi) *
    π *
    (
        3 * (
            F1All(p0 - q0, ps, k, m, T) - F1All(p0 - q0, ps, UVScale, m, T) +
            F1All(p0 + q0, ps, k, m, T) - F1All(p0 + q0, ps, UVScale, m, T)
            +
            F2All(p0 - q0, ps, k, m, T) - F2All(p0 - q0, ps, UVScale, m, T) +
            F2All(p0 + q0, ps, k, m, T) - F2All(p0 + q0, ps, UVScale, m, T)
        )
        # +
        # (2 + Npi) * (
        #     F1All(1e-8 - 1e-14, 1e-8, k, m, T) -
        #     F1All(1e-8 - 1e-14, 1e-8, UVScale, m, T) +
        #     F2All(1e-8 - 1e-14, 1e-8, k, m, T) -
        #     F2All(1e-8 - 1e-14, 1e-8, UVScale, m, T)
        # )
    )
end

@doc raw"""
    VImintqsSimple(p0, ps, k, T, Npi, m, lamda,UVScale)

compute $\int_0^kq_s^2dqs\int_{-1}^{1}d\cos\theta\mathrm{Im}V_k(q_0)$, the `k` dependence of $\lambda_{4\pi}$ and $m$ are neglected.

`VImintqsSimple` only contains $V(q_0)$, no $V(-q_0)$

`VImintqsSimple` contains type-1 delta function

`VImintqsSimple` doesn't contains type-2 delta function
"""
function VImintqsSimple(p0, ps, k, T, Npi, m, lamda,UVScale)
    hcubature(
        x ->
            x[1]^2 * VImSimple(
                p0,
                sqrt(x[1]^2 + ps^2 + 2 * x[1] * x[2] * ps),
                Epi(k, m),
                k,
                m,
                T,
                Npi,
                lamda,UVScale,
            ),
        [0.0, -1.0],
        [k, 1.0],atol=1e-6,rtol=1e-6,initdiv=10,maxevals=8000,
    )[1]
end

function VImintqsSimple(p0, ps,q0, k, T, Npi, m, lamda,UVScale)
    hcubature(
        x ->
            x[1]^2 * VImSimple(
                p0,
                sqrt(x[1]^2 + ps^2 + 2 * x[1] * x[2] * ps),
                q0,
                k,
                m,
                T,
                Npi,
                lamda,UVScale,
            ),
        [0.0, -1.0],
        [k, 1.0],atol=1e-2,rtol=1e-2,initdiv=40,maxevals=16000,
    )[1]
end

@doc raw"""
    propImSimple(p0, ps, T,IRScale,UVScale, Npi, m, lamda)
`x[1]` is `qs`, `x[2]` is `costh`, `x[3]` is `k`

# Arguments
- `m`: mass square, it's a constant number.
- `lamda`: $\lambda_{4\pi}$, it's a constant number.
"""
function propImSimple(p0, ps, T, IRScale, UVScale, Npi, m, lamda)
    -hquadrature(
        x ->
            2*Coeffgamm2Simple(x, T, m) *
            VImintqsSimple(p0, ps, x, T, Npi, m, lamda,UVScale),
        IRScale,
        UVScale,
        atol = 1e-4,
        rtol = 1e-4,maxevals=200,
    )[1]
end


using LaTeXStrings
using Juqbox
using Plots

# Function to loop over ϵ values and create a plot of the objective function for pertrubed Hamiltonians
function ep_plot(pcof::Vector{Float64}, params:: Juqbox.objparams, wa::Working_Arrays, ep_vals::AbstractArray)
    results = zeros(length(ep_vals),4)

    for i =1:length(ep_vals)
        epsh = ep_vals[i]

        # Additive noise
        for j = 2:size(params.Hconst,2)
            params.Hconst[j,j] += 0.01*epsh*(10.0^(j-2))
        end

        obj, _, _, secondaryobjf, traceinfid = Juqbox.traceobjgrad(pcof,params,wa,false, true)
        results[i,1] = obj
        results[i,2] = secondaryobjf
        results[i,3] = traceinfid
        results[i,4] = traceinfid+secondaryobjf

        # Reset
        for j = 2:size(params.Hconst,2)
            params.Hconst[j,j] -= 0.01*epsh*(10.0^(j-2))
        end

    end

    pl1 = Plots.plot(ep_vals,results[:,1],yaxis=:log,xlabel = L"\epsilon",ylabel = "Objective Function")
    return results,pl1
end

function evalctrl_no_carrier(params::objparams, pcof0:: Array{Float64, 1}, jFunc:: Int64, D1::Int64)

    # Evaluate the ctrl functions on this grid in time
    nplot = round(Int64, params.T*32)
    # is this resolution sufficient for the lab frame ctrl functions so we can get meaningful FFTs?
    td = collect(range(0, stop = params.T, length = nplot+1))

    nfreq = length(params.Cfreq)
    offset = (jFunc-1)*2*D1
    nCoeff = 2*D1
    pcof = copy(pcof0[offset+1:offset+nCoeff])

    if (params.use_bcarrier)
        # B-splines with carrier waves (zero out the frequency to plot just the splines)
        splinepar = Juqbox.bcparams(params.T, D1, params.Ncoupled, params.Nunc, zeros(1,1), pcof)
    else
        # regular B-splines
        splinepar = Juqbox.splineparams(params.T, D1, 2*(params.Ncoupled + params.Nunc), pcof)
    end

    # define inline function to enable vectorization over t
    controlplot(t, splinefunc) = Juqbox.controlfunc(t, splinepar, splinefunc)

    fact = 1.0/(2*pi) # conversion factor to GHz
    fact = 1.0 # conversion factor to rad/ns
    pj = fact.*controlplot.(td, 0)
    qj = fact.*controlplot.(td, 1)
    return pj, qj

end

# Maximum shift in Hamiltonian (in rad*GHz)
ep_max = 2*pi*2e-2

# For plotting purposes
max_ep_sweep = 2*pi*3e-2
len = 1001
ep_vals = range(-max_ep_sweep,stop=max_ep_sweep,length=len)
freshOptim = true

# Usual optimization
nquad = 1
include("swap-02-risk-neutral.jl")
if(freshOptim)
    pcof = Juqbox.run_optimizer(prob, pcof0)
    results,pl1 = ep_plot(pcof, params, wa, ep_vals)
    data = zeros(size(results,1),4)
    data[:,1] = ep_vals
    for j = 1:3
        data[:,j+1] = results[:,j]
    end
    writedlm("usual_control_GLQ2.dat", pcof)
    writedlm("usual_optim_OF_sweep_GLQ2.dat", data)
else
    pcof = vec(readdlm("usual_control_GLQ2.dat"))
    results = readdlm("usual_optim_OF_sweep_GLQ2.dat")
end

pcof_old = copy(pcof)

# Plot control functions
scalefactor = 1000/(2*pi)
unitStr = "MHz"

nplot = round(Int64, params.T*32)
td = collect(range(0, stop = params.T, length = nplot+1))
p_NF_1,q_NF_1 = evalctrl_no_carrier(params, pcof_old, 1, D1)
p_NF_2,q_NF_2 = evalctrl_no_carrier(params, pcof_old, 2, D1)

pfunc1 = scalefactor .* p_NF_1
qfunc1 = scalefactor .* q_NF_1
pmax1 = maximum(abs.(pfunc1))
qmax1 = maximum(abs.(qfunc1))
pfunc2 = scalefactor .* p_NF_2
qfunc2 = scalefactor .* q_NF_2
pmax2 = maximum(abs.(pfunc2))
qmax2 = maximum(abs.(qfunc2))
pmax = maximum([pmax1,pmax2])
qmax = maximum([qmax1,qmax2])

fnt = Plots.font("Helvetica", 12)
lfnt = Plots.font("Helvetica", 12)
Plots.default(titlefont=fnt, guidefont=fnt, tickfont=fnt, legendfont=lfnt, linewidth=1.5, size=(650, 350))

titlestr = "Rotating frame NF ctrl " * " Max-p=" *@sprintf("%.3e", pmax) * " Max-q=" *@sprintf("%.3e", qmax) * " " * unitStr

pl_ctrl_NF = Plots.plot(td, pfunc1, lab=L"p_{1,1}(t)", title = titlestr, xlabel="Time [ns]",
                                  ylabel=unitStr, legend= :outerright, linewidth=1.5, legendfontsize=14)
# add in the control function for the anti-symmetric Hamiltonian
Plots.plot!(pl_ctrl_NF,td, qfunc1, lab=L"q_{1,1}(t)", linewidth=1.5, legendfontsize=14)
Plots.plot!(pl_ctrl_NF,td, pfunc2, lab=L"p_{1,2}(t)", linewidth=1.5, legendfontsize=14)
Plots.plot!(pl_ctrl_NF,td, qfunc2, lab=L"q_{1,2}(t)", linewidth=1.5, legendfontsize=14)

# save plots of control functions in rotating frame without carrier waves
Plots.savefig(pl_ctrl_NF,  "robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"_D1_NF_ctrl.png")

# Risk-neutral optimization
nquad = 9
include("swap-02-risk-neutral.jl")
if(freshOptim)
    pcof = Juqbox.run_optimizer(prob, pcof0)
    results2,pl2 = ep_plot(pcof, params, wa, ep_vals)
    data2 = zeros(size(results2,1),4)
    data2[:,1] = ep_vals
    for j = 1:3
        data2[:,j+1] = results2[:,j]
    end
    writedlm("robust_control_GLQ2.dat", pcof)
    writedlm("robust_optim_OF_sweep_GLQ2.dat", data2)
else
    pcof = vec(readdlm("robust_control_GLQ2.dat"))
    results2 = readdlm("robust_optim_OF_sweep_GLQ2.dat")
end

# Plot objective functions for various Hamiltonian perturbations
plc = Plots.plot(scalefactor.*ep_vals,results[:,4],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="NF Infidelity")
Plots.plot!(plc,scalefactor.*ep_vals,results2[:,4],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="RN Infidelity")
Plots.plot!(plc,scalefactor.*ep_vals,results[:,3],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="NF Guard Level Pop.",linestyle=:dash)
Plots.plot!(plc,scalefactor.*ep_vals,results2[:,3],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="RN Guard Level Pop.",linestyle=:dash,legend= :outerright)

plc_short = Plots.plot(scalefactor.*ep_vals,results[:,4],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="NF Infidelity")
Plots.plot!(plc_short,scalefactor.*ep_vals,results2[:,4],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="RN Infidelity")
Plots.plot!(plc_short,scalefactor.*ep_vals,results[:,3],yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="NF Guard Level Pop.",linestyle=:dash)
Plots.plot!(plc_short,scalefactor.*ep_vals,results2[:,3],xlims=((-20,20)),yaxis=:log,xlabel = "Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="RN Guard Level Pop.",linestyle=:dash,legend= :outerright)

plc2 = Plots.plot(scalefactor.*ep_vals,results[:,2].+results[:,3],yaxis=:log,xlabel ="Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="Total NF Objective")
Plots.plot!(plc2,scalefactor.*ep_vals,results2[:,2].+results2[:,3],yaxis=:log,xlabel ="Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="Total RN Objective",legend= :outerright)

plc2_short = Plots.plot(scalefactor.*ep_vals,results[:,2].+results[:,3],yaxis=:log,xlabel ="Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="Total NF Objective")
Plots.plot!(plc2_short,scalefactor.*ep_vals,results2[:,2].+results2[:,3],xlims=((-20,20)),yaxis=:log,xlabel ="Hamiltonian Perturbation [MHz]",ylabel = "Objective Function", lab="Total RN Objective",legend= :outerright)

# Save plots to file
Plots.savefig(plc,  "robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"_D1_separate.png")
Plots.savefig(plc_short,  "robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"_D1_separate_short.png")
Plots.savefig(plc2,  "robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"_D1_total.png")
Plots.savefig(plc2_short,  "robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"_D1_total_short.png")

p_RN_1,q_RN_1 = evalctrl_no_carrier(params, pcof, 1, D1)
p_RN_2,q_RN_2 = evalctrl_no_carrier(params, pcof, 2, D1)
pfunc1 = scalefactor .* p_RN_1
qfunc1 = scalefactor .* q_RN_1
pmax1 = maximum(abs.(pfunc2))
qmax1 = maximum(abs.(qfunc2))
pfunc2 = scalefactor .* p_RN_2
qfunc2 = scalefactor .* q_RN_2
pmax1 = maximum(abs.(pfunc2))
qmax1 = maximum(abs.(qfunc2))
pmax = maximum([pmax1,pmax2])
qmax = maximum([qmax1,qmax2])

titlestr = "Rotating frame RN ctrl " * " Max-p=" *@sprintf("%.3e", pmax) * " Max-q=" *@sprintf("%.3e", qmax) * " " * unitStr
pl_ctrl_RN = Plots.plot(td, pfunc1, lab=L"p_{1,1}(t)", title = titlestr, xlabel="Time [ns]",
                                  ylabel=unitStr, legend= :topleft, linewidth=1.5, legendfontsize=12)
# add in the control function for the anti-symmetric Hamiltonian
Plots.plot!(pl_ctrl_RN,td, qfunc1, lab=L"q_{1,1}(t)", linewidth=1.5, legendfontsize=12)
Plots.plot!(pl_ctrl_RN,td, pfunc2, lab=L"p_{1,2}(t)", linewidth=1.5, legendfontsize=12)
Plots.plot!(pl_ctrl_RN,td, qfunc2, lab=L"q_{1,2}(t)", linewidth=1.5, legendfontsize=12)
Plots.savefig(pl_ctrl_RN,  "robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"D1_RN_ctrl.png")

# Save coefficients
save_pcof("robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"D1_NF_pcof.jld2",pcof_old)
save_pcof("robust_comparison_"* @sprintf("%1.1e", ep_max/(2*pi)) * "_"*@sprintf("%3.1f",T)*"_T_"*@sprintf("%3d",maxIter)*"_iters_"*@sprintf("%d",nquad)*"_N_"*@sprintf("%d",D1)*"D1_RN_pcof.jld2",pcof)


#!/usr/bin/env julia

using LogMath, Base.Test

include("test-math.jl")
include("test-LogNum.jl")


"""
$(SIGNATURES)

Fill empty sparse matrix A with random nonzero elements from interval [1,2]
using incremental assembly.

"""
function sprand!(A::AbstractSparseMatrix{Tv,Ti},xnnz::Int) where {Tv,Ti}
    m,n=size(A)
    for i=1:xnnz
        i=rand((1:m))
        j=rand((1:n))
        a=1.0+rand(Tv)
        A[i,j]+=a
    end
end


## cholesky.jl : Gaussian random field generator based on the Cholesky decomposition of the covariance matrix

"""
    Cholesky <: GaussianRandomFieldGenerator

A [`GaussiandRandomFieldGenerator`](@ref) based on a Cholesky factorization of the covariance matrix.

# Examples
```jldoctest
julia> m = Matern(0.1,1.0)
Matérn (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> c = CovarianceFunction(2,m)
2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> pts1 = 0:0.02:1; pts2 = 0:0.02:1
0.0:0.02:1.0

julia> grf = GaussianRandomField(c,Cholesky(),pts1,pts2)
Gaussian random field with 2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0) on a 51x51 structured grid, using a Cholesky decomposition

julia> plot(grf)
[...]

```
Note that the Cholesky factorization requires the covariance matrix to be `Symmetric` and positive definite. If the covariance matrix is not `Symmetric`, a warning will be thrown but the method will try to continue with an approximate symmetric matrix. If the covariance matrix is not positive definite, an error will be thrown. Try using the `Spectral` method in that case.

See also: [`Spectral`](@ref), [`KarhunenLoeve`](@ref), [`CirculantEmbedding`](@ref)
"""
struct Cholesky <: GaussianRandomFieldGenerator end

# compose a GaussianRandomField using Cholesky factorization
function _GaussianRandomField(mean, cov::CovarianceFunction, method::Cholesky, pts...)
    C = apply(cov,pts,pts)

    # error checking
    issymmetric(C) || @warn "to use a Cholesky factorization, the covariance matrix must be symmetric/hermitian"

    # compute Cholesky factorization
    isposdef!(Symmetric(C)) || throw(ArgumentError("to use a Cholesky factorization, the covariance matrix must be positive definite"))
    L = LowerTriangular(C')

    GaussianRandomField{Cholesky,typeof(cov),typeof(pts),typeof(mean),typeof(L)}(mean,cov,pts,L)
end

# returns the required dimension of the random points
randdim(grf::GaussianRandomField{Cholesky}) = size(grf.data, 1)

# sample from the GaussianRandomField using Cholesky factorization
function _sample(grf::GaussianRandomField{Cholesky}, xi)
	grf.mean + std(grf.cov) * reshape(grf.data * xi, size(grf.mean))
end

Base.show(io::IO, ::Cholesky) = print(io, "Cholesky decomposition")


using CSV
using Dates
using DataFrames
include("../src/TrueSkill.jl")
using .TrueSkill
global const ttt = TrueSkill
using DataFrames

base = Dates.value(Date("1970-01-01") - Date("1900-01-01"))

data = CSV.read("input/history.csv", DataFrame)
djokovic = "d643"
federer = "f324"
sampras = "s402"
lendl = "l018"
connors = "c044"
nadal = "n409"
john_mcenroe = "m047"
bjorn_borg = "b058"
aggasi = "a092"
hewitt = "h432"
edberg = "e004"
vilas = "v028"
nastase = "n008"

courier = "c243"
kuerten = "k293"
murray = "mc10"
wilander = "w023"
roddick = "r485"
#data.w1_id[occursin.("roddick",data.w1_name)]
#data.w1_name[occursin.("pi74",data.w1_id)]

days = Dates.value.(data[:,"time_start"] .- Date("1900-01-01"))
#dates = Date.(Dates.UTD.(days .+ Dates.value(Date("1900-01-01"))))

composition = [ r.double == "t" ? [[r.w1_id,r.w2_id],[r.l1_id,r.l2_id]] : [[r.w1_id],[r.l1_id]] for r in eachrow(data) ]
function fit()
    h = ttt.History(composition=composition, times = days, sigma = 1.6, gamma = 0.036)
    ttt.convergence(h,epsilon=0.01, iterations=10)
    return h
end
h = ttt.History(composition=composition, times = days, sigma = 1.6, gamma = 0.036)
lc = ttt.learning_curves(h)

dict = Dict{String,Vector}()
dict["federer_mu"] = [tp[2].mu for tp in lc["f324"]]
dict["federer_sigma"]= [tp[2].sigma for tp in lc["f324"]]
dict["federer_time"] = [tp[1] for tp in lc["f324"]]
dict["nadal_mu"] = [tp[2].mu for tp in lc["n409"]]
dict["nadal_sigma"] = [tp[2].sigma for tp in lc["n409"]]
dict["nadal_time"] = [tp[1] for tp in lc["n409"]]
dict["djokovic_mu"] = [tp[2].mu for tp in lc["d643"]]
dict["djokovic_sigma"] = [tp[2].sigma for tp in lc["d643"]]
dict["djokovic_time"] = [tp[1] for tp in lc["d643"]]
dict["sampras_mu"] = [tp[2].mu for tp in lc["s402"]]
dict["sampras_sigma"] = [tp[2].sigma for tp in lc["s402"]]
dict["sampras_time"] = [tp[1] for tp in lc["s402"]]
dict["aggasi_mu"] = [tp[2].mu for tp in lc["a092"]]
dict["aggasi_sigma"] = [tp[2].sigma for tp in lc["a092"]]
dict["aggasi_time"] = [tp[1] for tp in lc["a092"]]
dict["vilas_mu"] = [tp[2].mu for tp in lc["v028"]]
dict["vilas_sigma"] = [tp[2].sigma for tp in lc["v028"]]
dict["vilas_time"] = [tp[1] for tp in lc["v028"]]
dict["borg_mu"] = [tp[2].mu for tp in lc["b058"]]
dict["borg_sigma"] = [tp[2].sigma for tp in lc["b058"]]
dict["borg_time"] = [tp[1] for tp in lc["b058"]]
dict["mcenroe_mu"] = [tp[2].mu for tp in lc["m047"]]
dict["mcenroe_sigma"] = [tp[2].sigma for tp in lc["m047"]]
dict["mcenroe_time"] = [tp[1] for tp in lc["m047"]]

dict["lendl_mu"] = [tp[2].mu for tp in lc["l018"]]
dict["lendl_sigma"] = [tp[2].sigma for tp in lc["l018"]]
dict["lendl_time"] = [tp[1] for tp in lc["l018"]]

dict["connors_mu"] = [tp[2].mu for tp in lc["c044"]]
dict["connors_sigma"] = [tp[2].sigma for tp in lc["c044"]]
dict["connors_time"] = [tp[1] for tp in lc["c044"]]

dict["hewitt_mu"] = [tp[2].mu for tp in lc["h432"]]
dict["hewitt_sigma"] = [tp[2].sigma for tp in lc["h432"]]
dict["hewitt_time"] = [tp[1] for tp in lc["h432"]]

dict["edberg_mu"] = [tp[2].mu for tp in lc["e004"]]
dict["edberg_sigma"] = [tp[2].sigma for tp in lc["e004"]]
dict["edberg_time"] = [tp[1] for tp in lc["e004"]]

dict["nastase_mu"] = [tp[2].mu for tp in lc["n008"]]
dict["nastase_sigma"] = [tp[2].sigma for tp in lc["n008"]]
dict["nastase_time"] = [tp[1] for tp in lc["n008"]]

dict["courier_mu"] = [tp[2].mu for tp in lc["c243"]]
dict["courier_sigma"] = [tp[2].sigma for tp in lc["c243"]]
dict["courier_time"] = [tp[1] for tp in lc["c243"]]

dict["kuerten_mu"] = [tp[2].mu for tp in lc["k293"]]
dict["kuerten_sigma"] = [tp[2].sigma for tp in lc["k293"]]
dict["kuerten_time"] = [tp[1] for tp in lc["k293"]]

dict["murray_mu"] = [tp[2].mu for tp in lc["mc10"]]
dict["murray_sigma"] = [tp[2].sigma for tp in lc["mc10"]]
dict["murray_time"] = [tp[1] for tp in lc["mc10"]]

dict["wilander_mu"] = [tp[2].mu for tp in lc["w023"]]
dict["wilander_sigma"] = [tp[2].sigma for tp in lc["w023"]]
dict["wilander_time"] = [tp[1] for tp in lc["w023"]]

dict["roddick_mu"] = [tp[2].mu for tp in lc["r485"]]
dict["roddick_sigma"] = [tp[2].sigma for tp in lc["r485"]]
dict["roddick_time"] = [tp[1] for tp in lc["r485"]]

maxlen = maximum([length(value) for (key, value) in dict])

df = DataFrame(Dict(key => [value;repeat([missing],maxlen-length(value))] for (key, value) in dict))

CSV.write("output/atp_learning_curves_trueskill.csv", df; header=true)

ttt.convergence(h,epsilon=0.01, iterations=10)
#@time h = fit()
lc = ttt.learning_curves(h)

top = Set([key  for (key, value) in lc for s in value if ((s[2].mu > 6.0) & (s[1]>base))])

dict = Dict{String,Vector}()
dict["federer_mu"] = [tp[2].mu for tp in lc["f324"]]
dict["federer_sigma"]= [tp[2].sigma for tp in lc["f324"]]
dict["federer_time"] = [tp[1] for tp in lc["f324"]]
dict["nadal_mu"] = [tp[2].mu for tp in lc["n409"]]
dict["nadal_sigma"] = [tp[2].sigma for tp in lc["n409"]]
dict["nadal_time"] = [tp[1] for tp in lc["n409"]]
dict["djokovic_mu"] = [tp[2].mu for tp in lc["d643"]]
dict["djokovic_sigma"] = [tp[2].sigma for tp in lc["d643"]]
dict["djokovic_time"] = [tp[1] for tp in lc["d643"]]
dict["sampras_mu"] = [tp[2].mu for tp in lc["s402"]]
dict["sampras_sigma"] = [tp[2].sigma for tp in lc["s402"]]
dict["sampras_time"] = [tp[1] for tp in lc["s402"]]
dict["aggasi_mu"] = [tp[2].mu for tp in lc["a092"]]
dict["aggasi_sigma"] = [tp[2].sigma for tp in lc["a092"]]
dict["aggasi_time"] = [tp[1] for tp in lc["a092"]]
dict["vilas_mu"] = [tp[2].mu for tp in lc["v028"]]
dict["vilas_sigma"] = [tp[2].sigma for tp in lc["v028"]]
dict["vilas_time"] = [tp[1] for tp in lc["v028"]]
dict["borg_mu"] = [tp[2].mu for tp in lc["b058"]]
dict["borg_sigma"] = [tp[2].sigma for tp in lc["b058"]]
dict["borg_time"] = [tp[1] for tp in lc["b058"]]
dict["mcenroe_mu"] = [tp[2].mu for tp in lc["m047"]]
dict["mcenroe_sigma"] = [tp[2].sigma for tp in lc["m047"]]
dict["mcenroe_time"] = [tp[1] for tp in lc["m047"]]

dict["lendl_mu"] = [tp[2].mu for tp in lc["l018"]]
dict["lendl_sigma"] = [tp[2].sigma for tp in lc["l018"]]
dict["lendl_time"] = [tp[1] for tp in lc["l018"]]

dict["connors_mu"] = [tp[2].mu for tp in lc["c044"]]
dict["connors_sigma"] = [tp[2].sigma for tp in lc["c044"]]
dict["connors_time"] = [tp[1] for tp in lc["c044"]]

dict["hewitt_mu"] = [tp[2].mu for tp in lc["h432"]]
dict["hewitt_sigma"] = [tp[2].sigma for tp in lc["h432"]]
dict["hewitt_time"] = [tp[1] for tp in lc["h432"]]

dict["edberg_mu"] = [tp[2].mu for tp in lc["e004"]]
dict["edberg_sigma"] = [tp[2].sigma for tp in lc["e004"]]
dict["edberg_time"] = [tp[1] for tp in lc["e004"]]

dict["nastase_mu"] = [tp[2].mu for tp in lc["n008"]]
dict["nastase_sigma"] = [tp[2].sigma for tp in lc["n008"]]
dict["nastase_time"] = [tp[1] for tp in lc["n008"]]

dict["courier_mu"] = [tp[2].mu for tp in lc["c243"]]
dict["courier_sigma"] = [tp[2].sigma for tp in lc["c243"]]
dict["courier_time"] = [tp[1] for tp in lc["c243"]]

dict["kuerten_mu"] = [tp[2].mu for tp in lc["k293"]]
dict["kuerten_sigma"] = [tp[2].sigma for tp in lc["k293"]]
dict["kuerten_time"] = [tp[1] for tp in lc["k293"]]

dict["murray_mu"] = [tp[2].mu for tp in lc["mc10"]]
dict["murray_sigma"] = [tp[2].sigma for tp in lc["mc10"]]
dict["murray_time"] = [tp[1] for tp in lc["mc10"]]

dict["wilander_mu"] = [tp[2].mu for tp in lc["w023"]]
dict["wilander_sigma"] = [tp[2].sigma for tp in lc["w023"]]
dict["wilander_time"] = [tp[1] for tp in lc["w023"]]

dict["roddick_mu"] = [tp[2].mu for tp in lc["r485"]]
dict["roddick_sigma"] = [tp[2].sigma for tp in lc["r485"]]
dict["roddick_time"] = [tp[1] for tp in lc["r485"]]

maxlen = maximum([length(value) for (key, value) in dict])

using DataFrames
df = DataFrame(Dict(key => [value;repeat([missing],maxlen-length(value))] for (key, value) in dict))

CSV.write("output/atp_learning_curves.csv", df; header=true)


module IntegralArrays

using IntervalSets
using FixedPointNumbers: FixedPoint, floattype
using ColorTypes: Colorant

const .. = IntervalSets.:(..)
export IntegralArray, ..

"""
    iA = IntegralArray([buffer,] A)

Construct the integral array of the input array `A`. If buffer of the same shape as `A`
is provided, then this is a non-allocating version.

The integral array is calculated by assigning to each cell the sum of all cells above it and
to its left, i.e. the rectangle from origin point to the pixel position. For example, in 1-D
case, `iA[i] = sum(A[1:i])` if the vector `A`'s origin is 1.

```jldoctest; setup=:(using IntegralArrays)
julia> A = collect(reshape(1:25, 5, 5))
5×5 $(Matrix{Int}):
 1   6  11  16  21
 2   7  12  17  22
 3   8  13  18  23
 4   9  14  19  24
 5  10  15  20  25

julia> iA = IntegralArray(A)
5×5 IntegralArray{$Int, 2, $(Matrix{Int})}:
  1   7   18   34   55
  3  16   39   72  115
  6  27   63  114  180
 10  40   90  160  250
 15  55  120  210  325

julia> iA[3, 3] == sum(A[1:3, 1:3])
true

julia> iA[3..5, 3..5] == sum(A[3:5, 3:5])
true

julia> IntegralArray(A, A); # in-place modifying A itself

julia> A
5×5 $(Matrix{Int}):
  1   7   18   34   55
  3  16   39   72  115
  6  27   63  114  180
 10  40   90  160  250
 15  55  120  210  325
```

The closed interval `a..b` is used to support the integral array `iX` with a different
indexing semantic; `iX[a:b]` gives a subarray of `iX`, while `iX[a..b]` calculates the sum
of original array `X` over region `a:b`. `a ± σ` (can be typed by `\\pm<tab>`) is a
convenient constructor of the closed interval `a-σ..a+σ`.

When one needs to compute the sum/average of all blocks extracted from an image, pre-building
the integral array usually provides a more efficient computation.

```julia
using BenchmarkTools, IntegralArrays

# simplified 3x3 mean filter; only for demo purpose
function mean_filter_naive!(out, X)
    Δ = CartesianIndex(1, 1)
    for i in CartesianIndex(2, 2):CartesianIndex(size(X).-1)
        block = @view X[i-Δ: i+Δ]
        out[i] = mean(block)
    end
    return out
end
function mean_filter_integral!(out, X)
    iX = IntegralArray(X)
    for i in CartesianIndex(2, 2):CartesianIndex(size(X).-1)
        x, y = i.I
        out[i] = iX[x±1, y±1]/9
    end
    return out
end

X = Float32.(rand(1:5, 64, 64));
m1 = copy(X);
m2 = copy(X);

@btime mean_filter_naive!(\$m1, \$X); # 65.078 μs (0 allocations: 0 bytes)
@btime mean_filter_integral!(\$m2, \$X); # 12.161 μs (4 allocations: 16.17 KiB)
m1 == m2 # true
```
"""
struct IntegralArray{T, N, A} <: AbstractArray{T, N}
    data::A
end

IntegralArray{T}(A::AbstractArray) where {T} = IntegralArray{T}(similar(A, T), A)
function IntegralArray{T}(data::AbstractArray, A::AbstractArray) where {T}
    axes(data) == axes(A) || throw(DimensionMismatch("integral data axes $(axes(data)) should be equal original array axes $(axes(A))."))
    cumsum!(data, A; dims=1)
    for i = 2:ndims(A)
        cumsum!(data, data; dims=i)
    end
    IntegralArray{T, ndims(data), typeof(data)}(data)
end

IntegralArray(A::AbstractArray) = IntegralArray{_maybe_floattype(eltype(A))}(A)
IntegralArray(data::AbstractArray, A::AbstractArray) =
    IntegralArray{_maybe_floattype(eltype(data))}(data, A)

# small integer eltypes almost certainly hit overflow issues (#15)
const SmallInt = @static if Int === Int64
    Union{Int8, UInt8, Int16, UInt16, Int32, UInt32}
else
    Union{Int8, UInt8, Int16, UInt16}
end
_widen(::Type{T}) where T<:SmallInt = T <: Signed ? Int : UInt
IntegralArray(A::AbstractArray{T}) where T<:SmallInt =
    IntegralArray{_maybe_floattype(_widen(T))}(A)
IntegralArray(data::AbstractArray{T}, A::AbstractArray) where T<:SmallInt =
    throw(ArgumentError("Small integer eltype $(T) would cause potential overflow issue, please use more bits for buffer array."))

Base.IndexStyle(::Type{IntegralArray{T,N,A}}) where {T,N,A} = IndexStyle(A)
Base.size(A::IntegralArray) = size(A.data)
Base.axes(A::IntegralArray) = axes(A.data)
Base.@propagate_inbounds Base.getindex(A::IntegralArray, i::Int) = A.data[i]
Base.@propagate_inbounds Base.getindex(A::IntegralArray{T,N}, i::Vararg{Int,N}) where {T,N} = A.data[i...]

Base.@propagate_inbounds function Base.getindex(A::IntegralArray{T,N}, i::Vararg{ClosedInterval{Int},N}) where {T,N}
    ret = zero(T)
    @boundscheck checkbounds(A, map(maximum, i)...)
    return _getindex(ret, A, axes(A), 1, (), i)
end

@inline _getindex(ret, A, ::Tuple{}, s, iint, ::Tuple{}) = (@inbounds(Ai = A[iint...]); ret + s * Ai)
@inline function _getindex(ret, A, axs, s, iint::Dims{M}, iiv::Tuple{ClosedInterval{Int},Vararg{ClosedInterval{Int}}}) where M
    ax1, axtail = axs[1], Base.tail(axs)
    iv1, ivtail = iiv[1], Base.tail(iiv)
    isempty(iv1) && return ret
    val  = _getindex(ret, A, axtail,  s, (iint..., maximum(iv1)), ivtail)
    imin = minimum(iv1) - 1
    checkindex(Bool, ax1, imin) || return val
    return _getindex(val, A, axtail, -s, (iint..., imin), ivtail)
end

Base.@propagate_inbounds function Base.getindex(A::IntegralArray{T,N}, r::ClosedInterval{CartesianIndex{N}}) where {T,N}
    A[map(.., r.left.I, r.right.I)...]
end

# only promote FixedPointNumber types(e.g., `N0f8`) to floats
# keeps consistant result for `Int` to Base
_maybe_floattype(::Type{CT}) where CT<:Colorant = floattype(CT)
_maybe_floattype(::Type{T}) where T<:FixedPoint = floattype(T)
_maybe_floattype(::Type{T}) where T = T

end # module


"""
    proj(v)

Returns a projector onto subspace spanned by `v`.
"""
function proj(v::SparseVector{T}) where T<:Number
   SparseMatrixCSC{T}(v*v')
end

function proj(v::AbstractVector{<:Number})
   v*v'
end

"""
    szegedy_walk_operators(szegedy)

Return pair of Szegedy walk operators based on `sqrtstochastic` from `szegedy`.
Definition of those can be found in https://arxiv.org/abs/1611.02238 on page 4.
The operators should be applied in given order.
"""

function szegedy_walk_operators(szegedy::AbstractSzegedy)
   order = nv(szegedy.graph)
   projectors = map(x->2. * proj(szegedy.sqrtstochastic[:,x]), 1:order)

   r1 = cat(projectors..., dims=(1, 2))

   r2 = spzeros(eltype(szegedy.sqrtstochastic), order^2, order^2)
   for x=1:order
      r2[(1:order:order^2) .+ (x-1),(1:order:order^2) .+ (x-1)] = projectors[x]
   end

   r1 -= I
   r2 -= I

   (r1, r2)
end

"""
    szegedy_oracle_operators(szegedy, marked)

Return a Szegedy oracle operators based on `sqrtstochastic` from `szegedy` and
collection of `marked` vertices. Definition of those can be found in
https://arxiv.org/abs/1611.02238 on page 6. The operators should be applied in
given order.
"""
function szegedy_oracle_operators(szegedy::AbstractSzegedy,
                                marked::Vector{Int})
   order = nv(szegedy.graph)

   markedidentity = SparseMatrixCSC{eltype(szegedy.sqrtstochastic)}(I, order, order)
   for v=marked
      markedidentity[v,v] = -1.
   end

   q1 = kron(markedidentity, SparseMatrixCSC{eltype(szegedy.sqrtstochastic)}(I, order, order))
   q2 = kron(SparseMatrixCSC{eltype(szegedy.sqrtstochastic)}(I, order, order), markedidentity)

   (q1, q2)
end


# # Load some packages

using MinimaxCalibratedEBayes
using RCall
using Plots
using StatsPlots
using LaTeXStrings
using ExponentialFamilies
using MosekTools
using Random

const MCEB = MinimaxCalibratedEBayes
gr() # plotting backend
GR.settextfontprec(232, 3)

# # Prostate data analysis

# ## Preliminaries

# Let us make some analysis choices from the beginning

# ### Empirical Bayes Targets
# Our goal is to estimate the Posterior mean and the local false sign rates:

target_grid = -3:0.2:3
post_means = PosteriorMean.(StandardNormalSample.(target_grid))
lfsrs =  LFSR.(StandardNormalSample.(target_grid));

# ### Empirical Bayes Prior Class
# For the prior class we consider a mixture of Gaussians
# that have standard deviation $0.2$ and with mixing measure
# supported on $[-3,3]$. Note that we also specify the solver
# to be used for solving the convex programming problems;
# here we use Mosek.
gcal = GaussianMixturePriorClass(0.2, -3:0.02:3, Mosek.Optimizer, ["QUIET" =>true])

# ### Discretization of marginal observations
marginal_grid = -4:0.02:4

# ## Load the data

# The dataset is from the following reference:
#   >Dinesh Singh, Phillip G. Febbo, Kenneth Ross, Donald G. Jackson, Judith Manola, Christine Ladd, Pablo Tamayo, Andrew A. Renshaw, Anthony V. D’Amico, Jerome P. Richie, Eric S. Lander, Massimo Loda, Philip W. Kantoff, Todd R. Golub, and William R. Sellers. Gene expression correlates of clinical prostate cancer behavior. Cancer cell, 1(2): 203–209, 2002.

# See the following monograph for further illustration of empirical Bayes methods on this dataset
#   >Bradley Efron. Large-scale inference: Empirical Bayes methods for estimation, testing, and prediction. Cambridge University Press, 2012.
prostz_file = MinimaxCalibratedEBayes.prostate_data_file
R"load($prostz_file)"
@rget prostz;
Zs = StandardNormalSample.(prostz);

# We will implement the preprocessing steps manually here.
# Afterwards, we will use wrapper that automates all of these steps.

# First we split our data into two folds
Random.seed!(200)
Zs_train, Zs_test, _, _ = MCEB._split_data(Zs);

# Let us quickly check the range of values in the folds
extrema(response(Zs_train)), extrema(response(Zs_test))

# Now we use the first fold to estimate an $L_\infty$ neighborhood
# of the marginal density. We also fit a logspline model of the
# prior density to get our pilot estimators.

# ### Neighborhood construction:

infinity_band_options = KDEInfinityBandOptions(a_min=-4.0, a_max=4.0);
fkde_train = fit(infinity_band_options, response.(Zs_train));

# Let us visualize the result by showing a histogram of the datapoints
# the estimated density function and the $L_\infty$ bands.
prostate_marginal_plot = histogram(response.(Zs_train), bins=20, normalize=true,
          label="", alpha=0.4, fillcolor=:lightgray)
plot!(prostate_marginal_plot, fkde_train, fillalpha=0.45, #L"\bar{f}(x)\pm c_m",
      linewidth=1, label=nothing, #L"\bar{f}(x)\pm c_m",
	  legend=:topleft,  bg_legend=:transparent,
	  fg_legend=:transparent)

# ### G-model fitting

# We choose the following class
exp_family = ContinuousExponentialFamilyModel(Uniform(-3.6,3.6),
                         collect(-3.6:0.1:3.6); df=5, scale=true);
exp_family_solver = MCEB.ExponentialFamilyDeconvolutionMLE(cefm = exp_family, c0=0.001);

# Let us fit to the training data
exp_family_fit = fit(exp_family_solver, Zs_train);

# Let us look at the pilot estimates we got
lfsr_logspline_estimate = estimate.(lfsrs , exp_family_fit)
post_mean_logspline_estimate = estimate.(post_means , exp_family_fit);
#src expfamily_fits_plot = plot(plot(target_grid, lfsr_logspline_estimate),
#src                           plot(target_grid, post_mean_logspline_estimate),
#src						   label="")

# ### Implement discretization

# Discretize our test samples to implement minimax calibration methodology
Zs_test_discr = DiscretizedStandardNormalSamples(Zs_test, marginal_grid);
# Set the $L_{\infty}$ band we computed earlier.
Zs_test_discr = set_neighborhood(Zs_test_discr, fkde_train);

# ### Put everything together
# We put all of the above in one struct that will be used in the final step
# of confidence interval construction

# We set `cache_target=true` to speed up computations for the tutorial, cf.
# Appendix C.3.5 of the manuscript for an explanation.

cache_target = true
mceb_setup = MinimaxCalibratorSetup(Zs_train = Zs_train,
                                    Zs_test = Zs_test,
						            prior_class = gcal,
						            fkde_train = fkde_train,
						            Zs_test_discr = Zs_test_discr,
						            pilot_method = exp_family_fit,
									delta_tuner =  MCEB.HalfCIWidth,
									cache_target = cache_target);

# Fit he calibrators for the posterior means.
post_means_fits = fit.(mceb_setup, post_means)
prostate_postmeans_plot = plot(post_means_fits, ylims=(-3.2,3.2),
							fg_legend=:transparent, bg_legend=:transparent,
							legend = (0.35,0.84))
# Now repeat the same for the local false sign rates.
lfsrs_fits = fit.(mceb_setup, lfsrs)
prostate_lfsrs_plot = plot(lfsrs_fits, label=nothing)

# # The impact of Neighborhoods: Moving to opportunity

# We next seek to repeat the analysis in the second dataset. Here we will also
# see how most of the steps manually conducted for the prostate dataset
# can be automated.

# The references for this dataset is the following:
#  >Raj Chetty and Nathaniel Hendren. The impacts of neighborhoods on intergenerational mobility II: County-level estimates. The Quarterly Journal of Economics, 133(3):1163– 1228, 2018.

# ## Loading the dataset
# Let us start by loading the dataset and filtering out missing observations.

using CSV
nbhood_file = MinimaxCalibratedEBayes.chetty_hendren_file
nbhoods_df = CSV.read(nbhood_file, types=Dict(3=>String));
missing_dx = [!ismissing(x) for x in nbhoods_df[!,:p25_coef]]
β_hat = [nbhoods_df[i,:p25_coef] for i in findall(missing_dx)]
se_hat = [nbhoods_df[i,:p25_se] for i in findall(missing_dx)]
nbhood_Zs = StandardNormalSample.(β_hat./se_hat);
prostate_marginal_grid = -4:0.02:4
length(nbhood_Zs)

# ## Setup the minimax calibrator
# The above steps repeat all the setup we conducted for the prostate data
# to compute the `mceb_setup` object.
Random.seed!(34)
mceb_options = MinimaxCalibratorOptions(prior_class = gcal,
                               marginal_grid = prostate_marginal_grid,
							   pilot_options = exp_family_solver,
							   cache_target=cache_target)
nbhood_mceb_setup = fit(mceb_options, nbhood_Zs);

# Again us a sanity check, let us check the range of data in each fold.
extrema(response.(nbhood_mceb_setup.Zs_train)),extrema(response.(nbhood_mceb_setup.Zs_test))

# Also let us plot again the marginal distribution and the estimated $L_{\infty}$ band.
nbhood_marginal_plot = histogram(response.(nbhood_mceb_setup.Zs_train), bins=20, normalize=true,
	            label=nothing, alpha=0.4, fillcolor=:lightgray)
plot!(nbhood_marginal_plot, nbhood_mceb_setup.fkde_train, fillalpha=0.45,
                  label=nothing, linewidth=1)

# Compute the posterior means
nbhood_post_means_fits = fit.(nbhood_mceb_setup, post_means)
nbhood_postmeans_plot = plot(nbhood_post_means_fits,  ylims=(-3.5,3.5), label=nothing)

# Compute the local false sign rates
nbhood_lfsrs_fits = fit.(nbhood_mceb_setup, lfsrs)
nbhood_lfsrs_plot = plot(nbhood_lfsrs_fits,  label=nothing)

# # Construct final plot
using Plots.Measures
dataset_plot = plot(prostate_marginal_plot, prostate_postmeans_plot, prostate_lfsrs_plot,
	 nbhood_marginal_plot, nbhood_postmeans_plot, nbhood_lfsrs_plot,
	 title = ["a) Prostate" "b)" "c)" " d) Neighborhoods" "e)" "f)"],
	 layout= (2,3), size=(750,500))
plot(dataset_plot)

#src savefig(dataset_plot, "dataset_plots.svg")	#src
#src savefig(dataset_plot, "dataset_plots.pdf")	#src


module TestUtil

using Test
using AMLPipelineBase
using AMLPipelineBase.Utils
using DataFrames: nrow
using Random

function test_utils()
  data = getiris()

  catnum=find_catnum_columns(data)
  @test catnum == ([5],[1,2,3,4])

  (r,l)=holdout(nrow(data),0.9)
  @test length.((r,l)) == (15,135)

  kf=kfold(nrow(data),10)
  @test size.(kf) |> unique == [(135,)]

  Y=data.Species |> collect
  @test score(:accuracy,Y,Y) == 100.0
  @test score(:accuracy,Y,reverse(Y)) |> round == 33.0

  @test (infer_eltype.(eachcol(data)) .<: [Float64,Float64,Float64,Float64,AbstractString]) |> sum == 5

  dc = Dict(Dict(:a=>:c,:d=>Dict(:e=>:f)))
  @test nested_dict_to_tuples(dc) == Set(Any[([:d, :e], :f), ([:a], :c)])

  dx = deepcopy(dc)
  nested_dict_set!(dx,[:d,:e],:h)
  @test dx[:d][:e] == :h

  @test nested_dict_merge(dx,dc) == dc
  @test mergedict(dc,dx) == dx

  x = [2,200,332,201,missing,50,missing,6,missing,7,8,109]
  @test skipmean(x) |> round == 102.0
  @test skipmedian(x) == 50.0
  @test skipstd(x) |> round == 118.0
end
@testset "Utils" begin
  test_utils()
end

function test_metric()
   data = getiris()
   X=data[:,2:end]
   Y=data[:,1] |> collect
   rf = RandomForest()
   catf = CatFeatureSelector()
   numf = NumFeatureSelector()
   ohe = OneHotEncoder()

   pl = (catf |> ohe) + (numf) |> rf

   rmse(X,Y) = score(:rmse,X,Y)
   @test crossvalidate(pl,X,Y,rmse,10,false).mean < 0.5
   @test crossvalidate(pl,X,Y;metric=rmse,verbose=false).mean < 0.5

   mse(X,Y) = score(:mse,X,Y)
   @test crossvalidate(pl,X,Y,mse,10,false).mean < 0.2
   @test crossvalidate(pl,X,Y;metric=mse,verbose=false).mean < 0.2

   mae(X,Y) = score(:mae,X,Y)
   @test crossvalidate(pl,X,Y,mae,10,false).mean < 0.2
   @test crossvalidate(pl,X,Y;metric=mae,verbose=false).mean < 0.2

   @test crossvalidate(pl,X,Y,mae,10,false).mean < 0.2
   @test crossvalidate(pl,X,Y;metric=mae,verbose=false).mean < 0.2

   tx,ty,tsx,tsy = train_test_split(data[:,1:4],data.Species |> collect)
   fit!(pl,tx,ty)
   score(:accuracy,transform!(pl,tsx),tsy) > 0.50
   pl1 = fit(pl,tx,ty)
   score(:accuracy,transform(pl1,tsx),tsy) > 0.50
end
@testset "Utils Metrics" begin
  test_metric()
end

end