"""
    easter_rata(y::Dates.Year) → Int

Returns Easter date as a *[Rata Die](https://en.wikipedia.org/wiki/Rata_Die)* number.

Based on *[Algo R](http://www.linuxtopia.org/online_books/programming_books/python_programming/python_ch38.html)*.
"""
function easter_rata(y::Dates.Year)

    # Algo R only works after 1582
    if y.value < 1582
        # Are you using this? Send me a postcard!
        error("Year cannot be less than 1582. Provided: $(y.value).")
    end

    # Century
    c = div(y.value , 100) + 1

    # Shifted Epact
    local se::Int = mod(14 + 11*(mod(y.value, 19)) - div(3*c, 4) + div(5+8*c, 25), 30)

    # Adjust Epact
    if (se == 0) || ((se == 1) && ( 10 < mod(y.value, 19) ))
       se += 1
    end

    # Paschal Moon
    p = Dates.Date(y.value, 4, 19).instant.periods.value - se

    # Easter: locate the Sunday after the Paschal Moon
    return p + 7 - mod(p, 7)
end

"""
    easter_date(y::Dates.Year) → Dates.Date

Returns result of `easter_rata` as a `Dates.Date` instance.
"""
@inline easter_date(y::Dates.Year) = Dates.Date(Dates.rata2datetime(easter_rata(y)))

"""
    findweekday(weekday_target::Integer, yy::Integer, mm::Integer, occurrence::Integer, ascending::Bool) → Date

Given a year `yy` and month `mm`, finds a date where a choosen weekday occurs.

`weekday_target` values are declared in module `Base.Dates`:
`Monday,Tuesday,Wednesday,Thursday,Friday,Saturday,Sunday = 1,2,3,4,5,6,7`.

If `ascending` is true, searches from the beginning of the month. If false, searches from the end of the month.

If `occurrence` is `2` and `weekday_target` is `Monday`, searches the 2nd Monday of the given month, and so on.
"""
function findweekday(weekday_target::Integer, yy::Integer, mm::Integer, occurrence::Integer, ascending::Bool) :: Dates.Date

    dt = Dates.Date(yy, mm, 1)
    local dt_dayofweek::Integer
    local offset::Integer

    @assert occurrence > 0 "occurrence must be > 0. Provided $(occurrence)."

    if ascending
        dt_dayofweek = Dates.dayofweek(dt)
        offset = mod(weekday_target + 7 - dt_dayofweek, 7)
    else
        dt = Dates.lastdayofmonth(dt)
        dt_dayofweek = Dates.dayofweek(dt)
        offset = mod(dt_dayofweek + 7 - weekday_target, 7)
    end

    if occurrence > 1
        offset += 7 * (occurrence - 1)
    end

    if ascending
        return dt + Dates.Day(offset)
    else
        return dt - Dates.Day(offset)
    end
end

"""

   adjustweekendholidayPost(dt, [adjust_saturdays]) → Date

In the UK and Canada, if a holiday falls on Saturday or Sunday, it's observed on the next business day.
This function will adjust to the next Monday.

`adjust_saturdays` kwarg defaults to `true`.
"""
function adjustweekendholidayPost(dt::Dates.Date; adjust_saturdays::Bool = true) :: Dates.Date

    if adjust_saturdays && (Dates.dayofweek(dt) == Dates.Saturday)
        return dt + Dates.Day(2)
    end

    if Dates.dayofweek(dt) == Dates.Sunday
        return dt + Dates.Day(1)
    end

    return dt
end

"""
    adjustweekendholidayUS(dt) → Date

In the United States, if a holiday falls on Saturday, it's observed on the preceding Friday.
If it falls on Sunday, it's observed on the next Monday.
"""
function adjustweekendholidayUS(dt::Dates.Date) :: Dates.Date

    if Dates.dayofweek(dt) == Dates.Saturday
        return dt - Dates.Day(1)
    end

    if Dates.dayofweek(dt) == Dates.Sunday
        return dt + Dates.Day(1)
    end

    return dt
end


# types and util funcitons used by the gaussian BP
# included by MessageGNN.jl

export GMM,GMMBP, randomSPMatrix, randomGaussian, randPrecision, randomGaussianBPConvergence, messagePassing, messagePassingOneStep, generateRandomBPConvgenceGraph,generateRandomBPConvgenceTree,getMarginalPrecision

import Distributions
import RandomMatrices

mutable struct GMM
    μ::Array{Float32,1}
    V::Array{Float32,2}
    P::Array{Float32,2}
    idx::Array{Int32,2}
    adj::Array{Int32,2}
    coeffs::Array{Float32,2}
    nvar::Int32
    nfactor::Int32
    GMM(μ, V;isPrecision = false) = constructGMM(new(), μ, V;isPrecision = isPrecision)
end

mutable struct GMMBP
    g::GMM
    marginalMeanList::Array{Float32,2} # list of intermediate marginal means
    marginalPList::Array{Float32,2} # list of intermediate marginal variances
    GMMBP(g, meanList, PList) = constructGMMBP(new(), g, meanList, PList)
end

function constructGMM(g::GMM, μ::AbstractVector, M::AbstractArray;isPrecision = false)
    if isPrecision
        g.P = M
        g.V = inv(M)
    else
        g.V = M
        g.P = inv(P)
    end
    g.μ = μ
    P = g.P
    g.nvar = size(μ, 1)
    g.nfactor = g.nvar * (g.nvar + 1) / 2 # will be updated later
    g.idx = zeros(Int32, g.nfactor, 2)
    g.adj = zeros(Int32, g.nfactor, g.nvar)
    g.coeffs = zeros(Float32, g.nfactor, 3)
    # singleton factors
    for i in 1:g.nvar
        g.idx[i,:] = [i,-1]
        g.adj[i,i] = 1
        g.coeffs[i,:] = [P[i,i], -2 * g.μ[i] * P[i,i], 0]
    end
    # pairwise factors
    fid = g.nvar + 1
    for i in 1:g.nvar - 1
        for j in i + 1:g.nvar
            if (P[i,j] != 0)
                g.idx[fid,:] = [i,j]
                g.adj[fid,i] = 1
                g.adj[fid,j] = 1
                g.coeffs[fid,:] =  [2 * P[i,j], -2 * g.μ[i] * P[i,j], -2 * g.μ[j] * P[i,j] ]
                fid += 1
            end
        end
    end
    g.nfactor = fid - 1
    g.idx = g.idx[1:g.nfactor,:]
    g.adj = g.adj[1:g.nfactor,:]
    g.coeffs = g.coeffs[1:g.nfactor,:]
    return g
end

function constructGMMBP(gbp::GMMBP, g::GMM, meanList::AbstractArray, PList::AbstractArray)
    gbp.g = g
    gbp.marginalMeanList = meanList
    gbp.marginalPList = PList
    return gbp
end

function randomSPMatrix(N::Integer)
    # A = Q'ΛQ
    # 2. random orthonormal used to transform the diagonal matrix
    # Harr matrix samples from beta-circular ensemble. beta: 1-orthogonal 2-unitary 4-symplectic
    Λ = diagm(rand(Distributions.Uniform(0, 10), N))
    Q = rand(RandomMatrices.Haar(1), N)
    return Q' * Λ * Q
end

function randomGaussian(N::Integer)
    # generate a dense random gaussian graphical model
    # 1. random mean vector
    μ =  rand(Uniform(-10, 10), N)
    # 2. random positive definite covariance matrix
    V = randomSPMatrix(N)
    # 3. return tuple of (mean,variance) as parametrization
    return (μ, V)
end

function randPrecision(ndim::Integer, nsample::Integer)
    ## generate random precision matrix that BP can converge to the true variance
    count = 0
    res = Array{Array{Float64,2},1}(nsample)
    total = 0
    while count < nsample
        total += 1
        _, Σ = randomGaussian(ndim)
        P = inv(Σ)
        P2 = P.^2
        crossSum = sum(P2, 2)
        crossMean = mean(crossSum)
        P[diagind(P)] += crossSum + rand(Uniform(0, crossMean), ndim)
        @assert all(eigvals(P) .> 0) "precision matrix not positive definite"
        # println("spectral radius: $ρ")
        count += 1
        res[count] = P
    end
    return res
end

function randPrecision(ndim::Integer)
    _, Σ = randomGaussian(ndim)
    P = inv(Σ)
    P2 = P.^2
    crossSum = sum(P2, 2)
    crossMean = mean(crossSum)
    P[diagind(P)] += crossSum + rand(Uniform(0, crossMean), ndim)
    @assert all(eigvals(P) .> 0) "precision matrix not positive definite"
    return P
end

function randomGaussianBPConvergence(ndim::Integer)
    # generate a random gaussian graphical model that can be exactly solved by BP
    # return: (mean, precision)
    # print("start generating bp graph")
    μ, Σ = randomGaussian(ndim)
    P = inv(Σ)
    P2 = P.^2
    crossSum = sum(P2, 2)
    crossMean = mean(crossSum)
    P[diagind(P)] += crossSum + rand(Uniform(0, crossMean), ndim)
    @assert all(eigvals(P) .> 0) "precision matrix not positive definite"
    # print(μ)
    # print("generated bp graph")
    return μ, P
end

function messagePassing(C::AbstractArray, niter::Integer;verbose = false)
    # message passing on a GGM with precision matrix C
    # C: original precision matrix
    # propagation rule:
    #   Pv_ij = C_ij^2/4Pf_ij
    #   Pf_jk = ∑_{i!=k} Pv_ij   # including Pv_jj
    niter -= 1
    n = size(C, 1)
    # initialize
    Pv = zeros(n, n)
    Pf = zeros(n, n)
    PvList = zeros(niter + 1, n, n)
    PfList = zeros(niter + 1, n, n)
    PrecisionList = zeros(niter + 1, n)

    for i in 1:n
        Pv[i,i] = C[i,i]
    end
    PvList[1,:,:] = Pv
    PfList[1,:,:] = Pf
    PrecisionList[1,:] = getMarginalPrecision(Pv)

    for q in 1:niter
        # 0.make a copy for convergence test
        oldPv = copy(Pv)
        oldPf = copy(Pf)
        # 1.update Pf
        for j in 1:n
            for k in 1:n
                if j != k
                    Pf[j,k] = sum(Pv[:,j]) - Pv[k,j]
                end
            end
        end
        # 2.update Pv
        for i in 1:n
            for j in 1:n
                if ( i != j ) && (C[i,j] != 0)
                    Pv[i,j] = -C[i,j]^2 / (4 * oldPf[i,j])
                end
            end
        end
        PvList[q + 1,:,:] = Pv
        PfList[q + 1,:,:] = Pf
        PrecisionList[q + 1,:] = getMarginalPrecision(Pv)
    end
    if verbose
        return Pv, Pf, PrecisionList
    else
        return Pv, Pf
    end
end

function messagePassingOneStep(Pv::AbstractArray, Pf::AbstractArray, C::AbstractArray)
    @assert size(Pv) == size(Pf) "size not match"
    @assert size(Pv, 1) == size(Pv, 2) "input should be square matrix"
    n = size(Pv, 1)
    oldPf = copy(Pf)
    for j in 1:n
        for k in 1:n
            if j != k
                Pf[j,k] = sum(Pv[:,j]) - Pv[k,j]
            end
        end
    end
    # 2.update Pv
    for i in 1:n
        for j in 1:n
            if (i != j) && (C[i,j] != 0)
                Pv[i,j] = -C[i,j]^2 / (oldPf[i,j])
            end
        end
    end
    return Pv, Pf
end

function getMarginalPrecision(Pv::AbstractMatrix)
    @assert size(Pv, 1) == size(Pv, 2) "input should be square matrix"
    return sum(Pv, 1)
end

function pruferToTree(seq)
    n = size(seq, 1)
    nnode = n + 2
    adj = zeros(nnode, nnode)
    degree = ones(nnode)
    for i in seq
        degree[i] += 1
    end
    for i in seq
        for jnode in 1:nnode
            if degree[jnode] == 1
                adj[i,jnode] = 1
                adj[jnode,i] = 1
                degree[i] -= 1
                degree[jnode] -= 1
                break
            end
        end
    end
    u = 0
    v = 0
    for inode in 1:nnode
        if degree[inode] == 1
            if u == 0
                u = inode
            else
                v = inode
                break
            end
        end
    end
    adj[u,v] = 1
    adj[v,u] = 1
    return adj
end

function randomTreeMatrix(ndim::Integer)
    n = ndim - 2
    seq = rand(1:ndim, n)
    adj = pruferToTree(seq)
    return adj
end

function generateRandomBPConvgenceGraph(ndim::Integer, bpIteration::Integer)
    # 1. generate a random precison matrix that will converge when applying BP on it
    μ, C = randomGaussianBPConvergence(ndim)
    # 2. run message passing algorithm and get extracted intermediate marginal precision vector
    Pv, Pf, PList = messagePassing(C, bpIteration;verbose = true)
    # 3. construct the data structure
    meanList = zeros(size(PList)...)
    g = GMM(μ, C;isPrecision = true)
    gbp = GMMBP(g, meanList, PList)
    return gbp
end

function generateRandomBPConvgenceTree(ndim::Integer, bpIteration::Integer)
    # 1. generate a random precison matrix that will converge when applying BP on it
    μ, C = randomGaussian(ndim)
    C .*= randomTreeMatrix(ndim) + eye(ndim)
    # 2. run message passing algorithm and get extracted intermediate marginal precision vector
    Pv = zeros(ndim, ndim)
    for k in 1:ndim
        Pv[k,k] = C[k,k]
    end
    Pf = zeros(ndim, ndim)
    for j in 1:bpIteration
        Pv, Pf = messagePassingOneStep(Pv, Pf, C)
    end
    # 3. construct the data structure
    meanList = zeros(size(PList)...)
    g = GMM(μ, C;isPrecision = true)
    gbp = GMMBP(g, meanList, PList)
    return gbp
end


module Splines2

export ns, ns_, bs, bs_, is, is_, ms, ms_, basis

using OffsetArrays
using LinearAlgebra
using Statistics

# utility functions
zeroArray(a :: Array{T}) where T<:Real = OffsetArray(a,(size(a).*0).-1)

# type definitions
abstract type AbstractSplineBasis{T<:Real} end
# Note: we have used OffsetArray for converting from C code
mutable struct SplineBasis{T<:Real} <: AbstractSplineBasis{T}
    order::Int # order of the spline
    nknots::Int # number of knots
    ncoef::Int # number of coefficients
    knots::OffsetArray{T,1} # knot vector
end
mutable struct BSplineBasis{T<:Real} <: AbstractSplineBasis{T}
    spline_basis::SplineBasis{T}
    boundary_knots::Tuple{T,T}
    interior_knots::Union{Array{T,1}, Nothing}
    intercept::Bool
    df::Int
end
mutable struct NSplineBasis{T<:Real} <: AbstractSplineBasis{T}
    b_spline_basis::BSplineBasis{T}
    qmat::Array{T,2}
    tl0::Array{T,1}
    tl1::Array{T,1}
    tr0::Array{T,1}
    tr1::Array{T,1}
end

# constructors
function SplineBasis(knots :: Array{T,1}, order :: Int = 4) where T<:Real
    SplineBasis(order,
                length(knots),
                length(knots)-order,
                zeroArray(knots))
end
function BSplineBasis(boundary_knots :: Tuple{T,T},
                      interior_knots :: Union{Array{T,1}, Nothing} = nothing,
                      order :: Int = 4,
                      intercept :: Bool = false) where T<:Real
    l_interior_knots = interior_knots == nothing ? 0 : length(interior_knots)
    df = Int(intercept) + order - 1 + l_interior_knots
    nknots = l_interior_knots + 2*order
    ncoef = nknots - order
    knots = zeros(T,nknots)
    for i=1:order
        knots[i] = boundary_knots[1]
        knots[nknots-i+1] = boundary_knots[2]
    end
    if (l_interior_knots>0)
        for i=1:l_interior_knots
            knots[i+order]=interior_knots[i]
        end
    end
    BSplineBasis(SplineBasis(knots,order), boundary_knots, interior_knots, intercept, df)
end
function NSplineBasis(boundary_knots :: Tuple{T,T},
                      interior_knots :: Union{Array{T,1}, Nothing} = nothing,
                      order :: Int = 4,
                      intercept :: Bool = false) :: NSplineBasis{T} where T<:Real
    bs = BSplineBasis(boundary_knots,interior_knots,order,intercept)
    co = basis(bs,[bs.boundary_knots[1],bs.boundary_knots[2]],2)
    qmat_ = LinearAlgebra.qr(transpose(co)).Q
    qmat = collect(transpose(qmat_[:, 3:size(qmat_,2)]))
    tl0 = qmat * basis(bs,boundary_knots[1])
    tl1 = qmat * basis(bs,boundary_knots[1],1)
    tr0 = qmat * basis(bs,boundary_knots[2])
    tr1 = qmat * basis(bs,boundary_knots[2],1)
    NSplineBasis(bs,
                 qmat,
                 tl0,tl1,tr0,tr1)
end

function basis(bs :: SplineBasis{T}, x :: T, ders :: Int = 0) where T<:Real
    function find_interval(x :: T)
        k = bs.order - 1
        n = bs.nknots - k - 1
        l = k
        while (x < bs.knots[l] && l != k)
            l = l - 1
        end
        l = l + 1
        while (x >= bs.knots[l] && l != n)
            l = l + 1
        end
        l - 1
    end
    t = bs.knots
    k = bs.order - 1
    m = ders
    hh = bs.order
    ell = find_interval(x)
    result = zeroArray(zeros(T,2*k+2))
    one = T(1)
    zero = T(0)
    result[0] = one
    for j = 1:(k-m)
        for n = 0:(j-1)
            result[hh+n] = result[n]
        end
        result[0] = zero
        for n = 1:j
            ind = ell + n
            xb = t[ind]
            xa = t[ind - j]
            if (xb == xa)
                result[n] = zero
                continue
            end
            w = result[hh + n - 1]/(xb - xa)
            result[n - 1] = result[n-1] + w*(xb - x)
            result[n] = w*(x - xa)
        end
    end
    for j = (k - m + 1):k
        for n = 0:(j-1)
            result[hh+n] = result[n]
        end
        result[0] = zero
        for n = 1:j
            ind = ell + n
            xb = t[ind]
            xa = t[ind - j]
            if (xb == xa)
                result[m] = zero
                continue
            end
            w = j*result[hh + n - 1]/(xb - xa)
            result[n - 1] = result[n - 1] - w
            result[n] = w
        end
    end
    offset = ell - k
    v = zeros(T,bs.ncoef)
    v[(1+offset):(k+1+offset)] = result[0:k]
    v
end

function basis(bs :: BSplineBasis{T}, x :: T, ders :: Int = 0) where T<:Real
    if (x < bs.boundary_knots[1] || x > bs.boundary_knots[2])
        if (x < bs.boundary_knots[1])
            k_pivot = T(0.75)*bs.boundary_knots[1]+
            T(0.25)*bs.spline_basis.knots[5-1] # 0-based
        else
            k_pivot = T(0.75)*bs.boundary_knots[2]+
            T(0.25)*bs.spline_basis.knots[length(bs.spline_basis.knots)-4-1] # 0-based
        end
        delta = x - k_pivot
        if (ders==0)
            vec = basis(bs.spline_basis,k_pivot,0) +
                basis(bs.spline_basis,k_pivot,1)*delta +
                basis(bs.spline_basis,k_pivot,2)*delta*delta/T(2.0) +
                basis(bs.spline_basis,k_pivot,3)*delta*delta*delta/T(6.0)
        elseif (ders==1)
            vec = splines.basis(bs.spline_basis,k_pivot,1) +
                splines.basis(bs.spline_basis,k_pivot,2)*delta +
                splines.basis(bs.spline_basis,k_pivot,3)*delta*delta/T(2.0)
        elseif (ders==2)
            vec = splines.basis(bs.spline_basis,k_pivot,2) +
                splines.basis(bs.spline_basis,k_pivot,3)*delta
        elseif (ders==3)
            vec = splines.basis(bs.spline_basis,k_pivot,3)
        else
            vec = k_pivot .* T(0)
        end
    else
        vec = basis(bs.spline_basis, x, ders)
    end
    if (!bs.intercept)
        vec = vec[2:length(vec)]
    end
    vec
end

function basis(ns :: NSplineBasis{T}, x::T, ders :: Int = 0) where T<:Real
    if (x < ns.b_spline_basis.boundary_knots[1])
        if (ders==0)
            vec = ns.tl0 + (x - ns.b_spline_basis.boundary_knots[1])*ns.tl1
        elseif (ders==1)
            vec = ns.tl1
        else
            vec = ns.tl1.* T(0)
        end
    elseif (x > ns.b_spline_basis.boundary_knots[2])
        if (ders==0)
            vec = ns.tr0 + (x - ns.b_spline_basis.boundary_knots[2])*ns.tr1
        elseif (ders==1)
            vsc = ns.tr1
        else
            vec = ns.tr1.* T(0)
        end
    else
        vec = ns.qmat*basis(ns.b_spline_basis, x, ders)
    end
    vec
end

function basis(bs :: AbstractSplineBasis{T}, x :: Array{T,1}, ders :: Int = 0) where T<:Real
    f(xi) = basis(bs, xi, ders)
    copy(transpose(reduce(hcat, f.(x))))
end

# utility function for processing the spline arguments
function spline_args(x :: Array{T,1};
                     boundary_knots :: Union{Tuple{T,T},Nothing} = nothing,
                     interior_knots :: Union{Array{T,1},Nothing} = nothing,
                     order :: Int = 4,
                     intercept :: Bool = false,
                     df :: Int = 3 + Int(intercept),
                     knots :: Union{Array{T,1}, Nothing} = nothing,
                     knots_offset :: Int = 0) where T<:Real
    if (interior_knots != nothing && boundary_knots != nothing)
        # pass
    elseif (knots != nothing)
        boundary_knots = extrema(knots)
        interior_knots = length(knots)==2 ? nothing : knots[2:(length(knots)-1)]
    else
        if (boundary_knots == nothing)
            boundary_knots = extrema(x)
        end
        iKnots = df - order + knots_offset + 1 - Int(intercept)
        if (iKnots>0)
            p = range(T(0), length=iKnots+2, stop=T(1))[2:(iKnots+1)]
            index = (x .>= boundary_knots[1]) .* (x .<= boundary_knots[2])
            interior_knots = Statistics.quantile(x[index], p)
        end
    end
    (boundary_knots, interior_knots)
end

"""
    bs_(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for B-splines and return a function with signature
`(x:: Array{T,1}; ders :: Int = 0)` for evaluation of `ders`
derivative for the splines at `x`.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `centre :: Union{T,Nothing} = nothing)`: value to centre the splines

# Examples
```jldoctest
julia> Splines2.bs_(collect(0.0:0.2:1.0), df=3)(collect(0.0:0.2:1.0))
6×3 Array{Float64,2}:
 0.0    0.0    0.0
 0.384  0.096  0.008
 0.432  0.288  0.064
 0.288  0.432  0.216
 0.096  0.384  0.512
 0.0    0.0    1.0
```
"""
function bs_(x :: Array{T,1};
            boundary_knots :: Union{Tuple{T,T},Nothing} = nothing,
            interior_knots :: Union{Array{T,1},Nothing} = nothing,
            order :: Int = 4,
            intercept :: Bool = false,
            df :: Int = order - 1 + Int(intercept),
            knots :: Union{Array{T,1}, Nothing} = nothing,
            centre :: Union{T,Nothing} = nothing) where T<:Real
    (boundary_knots, interior_knots) =
        spline_args(x, boundary_knots=boundary_knots, interior_knots=interior_knots,
                    order=order, intercept=intercept,
                    df=df, knots=knots)
    spline = BSplineBasis(boundary_knots, interior_knots, order, intercept)
    spline
end

"""
    bs(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for B-splines.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `centre :: Union{T,Nothing} = nothing)`: value to centre the splines
- `ders :: Int = 0`: derivatives of the splines

# Examples
```jldoctest
julia> Splines2.bs(collect(0.0:0.2:1.0), df=3)
6×3 Array{Float64,2}:
 0.0    0.0    0.0
 0.384  0.096  0.008
 0.432  0.288  0.064
 0.288  0.432  0.216
 0.096  0.384  0.512
 0.0    0.0    1.0
```
"""
function bs(x :: Array{T,1}; ders :: Int = 0, kwargs...) where T<:Real
    bs_(x; kwargs...)(x, ders=ders)
end

"""
    ns_(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for natural B-splines and return a function with signature
`(x:: Array{T,1}; ders :: Int = 0)` for evaluation of `ders`
derivative for the splines at `x`.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `centre :: Union{T,Nothing} = nothing)`: value to centre the splines

# Examples
```jldoctest
julia> Splines2.ns_(collect(0.0:0.2:1.0), df=3)(collect(0.0:0.2:1.0))
6×3 Array{Float64,2}:
  0.0       0.0        0.0
 -0.100444  0.409332  -0.272888
  0.102383  0.540852  -0.359235
  0.501759  0.386722  -0.172481
  0.418872  0.327383   0.217745
 -0.142857  0.428571   0.714286
```
"""
function ns_(x :: Array{T,1};
            boundary_knots :: Union{Tuple{T,T},Nothing} = nothing,
            interior_knots :: Union{Array{T,1},Nothing} = nothing,
            order :: Int = 4,
            intercept :: Bool = false,
            df :: Int = order - 3 + Int(intercept),
            knots :: Union{Array{T,1}, Nothing} = nothing,
            centre :: Union{T,Nothing} = nothing) where T<:Real
    (boundary_knots, interior_knots) =
        spline_args(x, boundary_knots=boundary_knots, interior_knots=interior_knots,
                    order=order, intercept=intercept, df=df, knots=knots, knots_offset=2)
    spline = NSplineBasis(boundary_knots, interior_knots, order, intercept)
    spline
end

function (spline::Union{BSplineBasis, NSplineBasis})(x; ders :: Int = 0)
    b = basis(spline, x, ders)
    #if (centre != nothing && ders==0)
    #    bc = basis(spline, centre, ders)
    #    for i=1:size(b,1)
    #        b[i,:] -= bc
    #    end
    #end
    b
end

"""
    ns(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for natural B-splines.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `centre :: Union{T,Nothing} = nothing)`: value to centre the splines
- `ders :: Int = 0`: derivatives of the splines

# Examples
```jldoctest
julia> Splines2.ns(collect(0.0:0.2:1.0), df=3)
6×3 Array{Float64,2}:
  0.0       0.0        0.0
 -0.100444  0.409332  -0.272888
  0.102383  0.540852  -0.359235
  0.501759  0.386722  -0.172481
  0.418872  0.327383   0.217745
 -0.142857  0.428571   0.714286
```
"""
function ns(x :: Array{T,1}; ders :: Int = 0, kwargs...) where T<:Real
    ns_(x; kwargs...)(x, ders=ders)
end

"""
    is_(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for I-splines and return a function with signature
`(x:: Array{T,1}; ders :: Int = 0)` for evaluation of `ders`
derivative for the splines at `x`.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)

# Examples
```jldoctest
julia> Splines2.is_(collect(0.0:0.2:1.0), df=3)(collect(0.0:0.2:1.0))
6×3 Array{Float64,2}:
 0.0     0.0     0.0
 0.1808  0.0272  0.0016
 0.5248  0.1792  0.0256
 0.8208  0.4752  0.1296
 0.9728  0.8192  0.4096
 1.0     1.0     1.0
```
"""
function is_(x :: Array{T,1};
            boundary_knots :: Union{Tuple{T,T},Nothing} = nothing,
            interior_knots :: Union{Array{T,1},Nothing} = nothing,
            order :: Int = 4,
            intercept :: Bool = false,
            df :: Int = order - 1 + Int(intercept),
            knots :: Union{Array{T,1}, Nothing} = nothing) where T<:Real
    (boundary_knots, interior_knots) =
        spline_args(x, boundary_knots=boundary_knots, order=order+1,
                    intercept=intercept, df=df+1, knots=knots)
    spline = BSplineBasis(boundary_knots, interior_knots, order+1, false)
    knots = parent(spline.spline_basis.knots)
    findj(x) = interior_knots==nothing ? (order+1) : searchsortedlast(knots,x)
    function eval(x :: Array{T,1}; ders :: Int = 0)
        b = basis(spline,x,ders)
        (nrow,ncol) = size(b)
        for i=1:nrow
            js = findj(x[i])
            for j in ncol:-1:1
                if j>js
                    b[i,j]=T(0)
                elseif j==ncol
                    b[i,j]=b[i,j]
                else
                    b[i,j]=b[i,j]+b[i,j+1]
                end
            end
            for j in ncol:-1:1
                if j < js-(order+1)
                    b[i,j]=T(1)
                end
            end
        end
        intercept ? b : b[:, 2:size(b,2)]
    end
    eval
end

"""
    is(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for I-splines.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `ders :: Int = 0`: derivatives of the splines

# Examples
```jldoctest
julia> Splines2.is(collect(0.0:0.2:1.0), df=3)
6×3 Array{Float64,2}:
 0.0     0.0     0.0
 0.1808  0.0272  0.0016
 0.5248  0.1792  0.0256
 0.8208  0.4752  0.1296
 0.9728  0.8192  0.4096
 1.0     1.0     1.0
```
"""
function is(x :: Array{T,1}; ders :: Int = 0, kwargs...) where T<:Real
    is_(x; kwargs...)(x, ders=ders)
end

"""
    ms_(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for M-splines and return a function with signature
`(x:: Array{T,1}; ders :: Int = 0)` for evaluation of `ders`
derivative for the splines at `x`.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `centre :: Union{T,Nothing} = nothing)`: value to centre the splines

# Examples
```jldoctest
julia> Splines2.ms_(collect(0.0:0.2:1.0), df=3)(collect(0.0:0.2:1.0))
6×3 Array{Float64,2}:
 0.0    0.0    0.0
 1.536  0.384  0.032
 1.728  1.152  0.256
 1.152  1.728  0.864
 0.384  1.536  2.048
 0.0    0.0    4.0
```
"""
function ms_(x :: Array{T,1};
            boundary_knots :: Union{Tuple{T,T},Nothing} = nothing,
            interior_knots :: Union{Array{T,1},Nothing} = nothing,
            order :: Int = 4,
            intercept :: Bool = false,
            df :: Int = order - 1 + Int(intercept),
            knots :: Union{Array{T,1}, Nothing} = nothing,
            centre :: Union{T,Nothing} = nothing) where T<:Real
    (boundary_knots, interior_knots) =
        spline_args(x, boundary_knots=boundary_knots, interior_knots=interior_knots,
                    order=order, intercept=intercept, df=df, knots=knots)
    spline = BSplineBasis(boundary_knots, interior_knots, order, intercept)
    function eval(x :: Array{T,1}; ders :: Int = 0)
        b = basis(spline, x, ders)
        knots = parent(spline.spline_basis.knots)
        function loop(j)
            denom = knots[j+order]-knots[j]
            denom > T(0) ? order/denom : T(0)
        end
        transCoef = map(loop,1:(length(knots)-order))
        if (!intercept)
            transCoef = transCoef[2:length(transCoef)]
        end
        if (centre != nothing && ders==0)
            bc = basis(spline, centre, ders)
            for i=1:size(b,1)
                b[i,:] -= bc
            end
        end
        for i=1:size(b,1)
            b[i,:] .*= transCoef
        end
        b
    end
    eval
end

"""
    ms(x :: Array{T,1}; <keyword arguments>) where T<:Real

Calculate a basis for M-splines.

The keyword arguments include one of:
1. `df`, possibly in combination with `intercept`
2. `boundary_knots` and `interior_knots`
3. `knots`

# Arguments
- `boundary_knots :: Union{Tuple{T,T},Nothing} = nothing`: boundary knots
- `interior_knots :: Union{Array{T,1},Nothing} = nothing`: interior knots
- `order :: Int = 4`: order of the spline
- `intercept :: Bool = false`: bool for whether to include an intercept
- `df :: Int = order - 1 + Int(intercept)`: degrees of freedom
- `knots :: Union{Array{T,1}, Nothing} = nothing`: full set of knots (excluding repeats)
- `centre :: Union{T,Nothing} = nothing)`: value to centre the splines
- `ders :: Int = 0`: derivatives of the splines

# Examples
```jldoctest
julia> Splines2.ms(collect(0.0:0.2:1.0), df=3)
6×3 Array{Float64,2}:
 0.0    0.0    0.0
 1.536  0.384  0.032
 1.728  1.152  0.256
 1.152  1.728  0.864
 0.384  1.536  2.048
 0.0    0.0    4.0
```
"""
function ms(x :: Array{T,1}; ders :: Int = 0, kwargs...) where T<:Real
    ms_(x; kwargs...)(x, ders=ders)
end

end # module


"""
n = DMatrix!(n)
n = DMatrix!(n, (:mode=>:symbolic))
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 1/15/2017, Last Modified: 12/06/2019 \n
--------------------------------------------------------------------------------------\n
"""
function DMatrix!(n::NLOpt, kwargs...)         #TODO make IMatrix and option
  kw = Dict(kwargs);

  if !haskey(kw,:mode); mode=:defaut;
  else; mode = get(kw,:mode,0);
  end

  if mode==:defaut
    check_ts = zeros(Float64,n.ocp.Ni);     # check input
    for int in 1:n.ocp.Ni; check_ts[int]=maximum(n.ocp.ts[int]); end
    if maximum(check_ts) < 10*eps()
      error("\n ts is full of zeros: make sure that you call create_intervals() first to calculate ts! \n
                NOTE: This may have occured because  (:finalTimeDV => true) and the final time dv is not working properly!! \n")
    end
    D = [zeros((n.ocp.Nck[int]+1),(n.ocp.Nck[int]+1)) for int in 1:n.ocp.Ni];
    for int in 1:n.ocp.Ni
        D[int] = polyDiff(n.ocp.ts[int]) # +1 is already appended onto ts
    end
    n.ocp.DMatrix = [zeros((n.ocp.Nck[int]),(n.ocp.Nck[int]+1)) for int in 1:n.ocp.Ni];
    DM = [zeros((n.ocp.Nck[int]),(n.ocp.Nck[int])) for int in 1:n.ocp.Ni];
    n.ocp.IMatrix = [zeros((n.ocp.Nck[int]),(n.ocp.Nck[int]+1)) for int in 1:n.ocp.Ni];
    for int in 1:n.ocp.Ni
        n.ocp.DMatrix[int] = D[int][1:end-1,:];    # [Nck]X[Nck+1]
        if n.s.ocp.integrationScheme==:lgrImplicit
          DM[int] = n.ocp.DMatrix[int][:,2:end];   # [Nck]X[Nck]
          n.ocp.IMatrix[int] = inv(DM[int]);       # I = inv(D[:,2:N_k+1])
        end
    end
  elseif mode==:symbolic # for validation only, too slow otherwise
    error(" \n cannot precompile with SymPy \n
             so this fucntion was turned off for typical use!! \n
               -> do a (using SymPy) in NLOptControl.jl then remove this error message and rerun \n")
    Dsym = [Array{Any}(undef, n.ocp.Nck[int],n.ocp.Nck[int]+1) for int in 1:n.ocp.Ni];
    n.ocp.DMatrix = [Array{Any}(undef, n.ocp.Nck[int],n.ocp.Nck[int]+1) for int in 1:n.ocp.Ni]
    test = [Array{Any}(undef, n.ocp.Nck[int]+1) for int in 1:n.ocp.Ni];
    val = 1; # since this is always = 1 this funtion is useful for testing, without scaling the problem from [-1,1] this was useful becuase tf was a design variable
    tf = Sym("tf")
    createIntervals!(n,tf); # gives symbolic expression
    for int in 1:n.ocp.Ni
      for i in 1:n.ocp.Nck[int]+1
        test[int][i] =  n.ocp.ts[int][i](tf=>val)
      end
    end
    for int in 1:n.ocp.Ni
        for idx in 1:n.ocp.Nck[int]+1
            for j in 1:n.ocp.Nck[int]
                f = lagrange_basis_poly(tf, test[int], idx)
                Dsym[int][j,idx] = diff(f,tf) # symbolic differentiation --> slow but useful TODO include this in test functions
                n.ocp.DMatrix[int][j,idx] = Dsym[int][j,idx](tf=>test[int][j])
            end
        end
    end
  end
  nothing
end

"""
scale_w(w,ta,tb)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/23/2017, Last Modified: 9/18/2017 \n
--------------------------------------------------------------------------------------\n
"""
function scale_w(w,ta,tb)
  (tb - ta)/2*w
end

"""
n = create_intervals(n)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/23/2017, Last Modified: 12/06/2019 \n
--------------------------------------------------------------------------------------\n
"""
function createIntervals!(n::NLOpt)
    tm = range(-1,1;length=n.ocp.Ni+1)     # create mesh points
    di = 2/n.ocp.Ni                           # interval size
    # go through each mesh interval creating time intervals; map [tm[i-1],tm[i]] --> [-1,1]
    n.ocp.ts = [[scale_tau(n.ocp.tau[int],tm[int],tm[int+1]);di*int-1] for int in 1:n.ocp.Ni]
    n.ocp.ws = [scale_w(n.ocp.w[int],tm[int],tm[int+1]) for int in 1:n.ocp.Ni]
    nothing
end

# NOTE this function was used for testing, but is currently depreciated. When it is used again figure out why and explain why
# di = (tf + 1).n.ocp.Ni
function createIntervals!(n::NLOpt, tf)
    tm = range(-1,1; length=n.ocp.Ni+1)       # create mesh points
    di = (tf + 1)/n.ocp.Ni                      # interval size
    # go through each mesh interval creating time intervals; map [tm[i-1],tm[i]] --> [-1,1]
    n.ocp.ts = [[scale_tau(n.ocp.tau[int],tm[int],tm[int+1]);di*int-1] for int in 1:n.ocp.Ni]
    n.ocp.ws = [scale_w(n.ocp.w[int],tm[int],tm[int+1]) for int in 1:n.ocp.Ni]
    nothing
end

"""
L = lagrange_basis_poly!(x,x_data,L)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/26/2016, Last Modified: 12/02/2019
Last Modified By: Alexander Buck
Citations: This function was influenced by the lagrange() function [located here](https://github.com/pjabardo/Jacobi.jl/blob/master/src/gauss_quad.jl)
--------------------------------------------------------------------------------------\n
# Input Arguments
* `x`: point to approximate function value at
* `x_data`: x data to used calculate basis polynomials
* `L`: array of interest

# Output Arguments
* `L`: Lagrange basis polynomials

A basic description of Lagrange interpolating polynomials is provided [here](http://127.0.0.1:8000/lagrange_poly.html#lagrange-poly)

# TODO get ride of Nc and replace it with Nck[int]
"""
function lagrange_basis_poly!(x::AbstractArray{T},x_data,L::AbstractArray{T}) where {T<:Number}
    Nc = length(x_data) - 1
    ns = length(x);
    L = zeros(Float64,Nc+1,ns);
    for idx in 1:Nc+1
      for j in 1:ns
        L[idx,j] = lagrange_basis_poly(x[j],x_data,Nc,idx);
      end
    end
    return L
end
lagrange_basis_poly(x::AbstractArray{T},x_data,Nc) where {T<:Number} = lagrange_basis_poly!(x::AbstractArray{T},x_data,Nc,zeros(x))

"""
D = polyDiff(x);
--------------------------------------------------------------------------\n
Last modifed for julia on 12/06/2019 by Mattias Fält\n
Original Author: JJ.A.C. Weideman, S.C. Reddy 1998\n
Original Function Name: poldif.m  |  Source: [matlabcentral](https://www.mathworks.com/matlabcentral/fileexchange/29-dmsuite)\n
https://pdfs.semanticscholar.org/bae2/1eb9458f194887bc8d7808383f56d7f4dca0.pdf
https://math.stackexchange.com/questions/1105160/evaluate-derivative-of-lagrange-polynomials-at-construction-points
--------------------------------------------------------------------------\n
# Input Arguments
* `x::Vector`: Vector of N distinct nodes.

# Output Arguments
* `D::Array{Float64,2}`: Differention Matrix

"""
function polyDiff(x)  #TODO get ride of B stuff
  N = length(x);
  x = x[:];                     # Make sure x is a column vector
  M = 1;                        # assume
  alpha = ones(N,1);
  B = zeros(M,N);
  L= Matrix{Bool}(LinearAlgebra.I, N, N);
  XX = repeat(x,1,N);
  DX = XX-XX';                 # DX contains entries x(k)-x(j).
  DX[L] = ones(N,1);           # Put 1's one the main diagonal.
  c = alpha.*prod(DX,dims=2);       # Quantities c(j).
  C = repeat(c,1,N);
  C = C./C';                   # Matrix with entries c(k)/c(j).
  Z = 1 ./ DX;                   # Z contains entries 1/(x(k)-x(j))
  Z[L] = zeros(N,1);           # with zeros on the diagonal.
  X = Z';                      # X is same as Z', but with
  # diagonal entries removed.
  flag = trues(size(X));
  flag[L] .= false;
  X = X[flag];
  X = reshape(X,N-1,N);
  Y = ones(N-1,N);             # Initialize Y and D matrices.
  D = Matrix{Float64}(LinearAlgebra.I,N,N);                  # Y is matrix of cumulative sums,
  temp = reshape(B[1,:],1,N)
  Y   = cumsum([temp; Y[1:N-1,:].*X], dims=1);     # Diagonals
  D   = Z.*(C.*repeat(LinearAlgebra.diag(D),1,N) - D);   # Off-diagonals
  D[L]   = Y[N,:];                         # Correct the diagonal
  return D
end


using JuMP
using Ipopt
using Hiop
using DelimitedFiles
using LinearAlgebra
using Test

include("opfdata.jl")

# export solve, model, initialPt_IPOPT, outputAll, computeAdmitances
function solve(opfmodel, opf_data)
    optimize!(opfmodel)
    status = termination_status(opfmodel)
    if status != MOI.LOCALLY_SOLVED
        println("Could not solve the model to optimality.")
    end
    return opfmodel,status
end

function model(opf_data; max_iter=100, solver="Ipopt", hiopalgebra = :Dense, scaling = 1.0)
    Pg0, Qg0, Vm0, Va0 = initialPt_IPOPT(opf_data)
    lines = opf_data.lines; buses = opf_data.buses; generators = opf_data.generators; baseMVA = opf_data.baseMVA
    busIdx = opf_data.BusIdx; FromLines = opf_data.FromLines; ToLines = opf_data.ToLines; BusGeners = opf_data.BusGenerators;

    nbus  = length(buses); nline = length(lines); ngen  = length(generators)

    #branch admitances
    YffR,YffI,YttR,YttI,YftR,YftI,YtfR,YtfI,YshR,YshI = computeAdmitances(lines, buses, baseMVA)

    #
    # JuMP model now
    #
    if solver == "Hiop"
        opfmodel = Model(optimizer_with_attributes(Hiop.Optimizer, "algebra" => hiopalgebra))
    else
        opfmodel = Model(optimizer_with_attributes(Ipopt.Optimizer, "max_iter" => max_iter))
    end

    @variable(opfmodel, generators[i].Pmin <= Pg[i=1:ngen] <= generators[i].Pmax, start = Pg0[i])
    @variable(opfmodel, generators[i].Qmin <= Qg[i=1:ngen] <= generators[i].Qmax, start = Qg0[i])

    @variable(opfmodel, Va[i=1:nbus], start = Va0[i])
    @variable(opfmodel, buses[i].Vmin <= Vm[i=1:nbus] <= buses[i].Vmax, start = Vm0[i])
    # fix the voltage angle at the reference bus
    set_lower_bound(Va[opf_data.bus_ref], buses[opf_data.bus_ref].Va)
    set_upper_bound(Va[opf_data.bus_ref], buses[opf_data.bus_ref].Va)

    # minimize active power
    coeff0 = Vector{Float64}(undef, ngen)
    for (i,v) in enumerate(generators)
        coeff0[i] = v.coeff[v.n]
    end
    coeff1 = Vector{Float64}(undef, ngen)
    for (i,v) in enumerate(generators)
        coeff1[i] = v.coeff[v.n - 1]
    end
    coeff2 = Vector{Float64}(undef, ngen)
    for (i,v) in enumerate(generators)
        coeff2[i] = v.coeff[v.n - 2]
    end

    @NLobjective(opfmodel, Min, scaling*sum( coeff2[i]*(baseMVA*Pg[i])^2
                                    + coeff1[i]*(baseMVA*Pg[i])
                                    + coeff0[i] for i=1:ngen))

    #
    # power flow balance
    #

    for b in 1:nbus
        # real part
        @NLconstraint(
        opfmodel,
        ( sum( YffR[l] for l in FromLines[b]) + sum( YttR[l] for l in ToLines[b]) + YshR[b] ) * Vm[b]^2
        + sum( Vm[b]*Vm[busIdx[lines[l].to]]  *( YftR[l]*cos(Va[b]-Va[busIdx[lines[l].to]]  ) + YftI[l]*sin(Va[b]-Va[busIdx[lines[l].to]]  )) for l in FromLines[b] )
        + sum( Vm[b]*Vm[busIdx[lines[l].from]]*( YtfR[l]*cos(Va[b]-Va[busIdx[lines[l].from]]) + YtfI[l]*sin(Va[b]-Va[busIdx[lines[l].from]])) for l in ToLines[b]   )
        - ( sum(baseMVA*Pg[g] for g in BusGeners[b]) - buses[b].Pd ) / baseMVA      # Sbus part
        ==0)
        # imaginary part
    end
    for b in 1:nbus
        @NLconstraint(
        opfmodel,
        ( sum(-YffI[l] for l in FromLines[b]) + sum(-YttI[l] for l in ToLines[b]) - YshI[b] ) * Vm[b]^2
        + sum( Vm[b]*Vm[busIdx[lines[l].to]]  *(-YftI[l]*cos(Va[b]-Va[busIdx[lines[l].to]]  ) + YftR[l]*sin(Va[b]-Va[busIdx[lines[l].to]]  )) for l in FromLines[b] )
        + sum( Vm[b]*Vm[busIdx[lines[l].from]]*(-YtfI[l]*cos(Va[b]-Va[busIdx[lines[l].from]]) + YtfR[l]*sin(Va[b]-Va[busIdx[lines[l].from]])) for l in ToLines[b]   )
        - ( sum(baseMVA*Qg[g] for g in BusGeners[b]) - buses[b].Qd ) / baseMVA      #Sbus part
        ==0)
    end
    #
    # branch/lines flow limits
    #
    nlinelim=0
    for l in 1:nline
        if lines[l].rateA!=0 && lines[l].rateA<1.0e10
        nlinelim += 1
        flowmax=(lines[l].rateA/baseMVA)^2

        #branch apparent power limits (from bus)
        Yff_abs2=YffR[l]^2+YffI[l]^2;        Yft_abs2=YftR[l]^2+YftI[l]^2
        Yre=YffR[l]*YftR[l]+YffI[l]*YftI[l]; Yim=-YffR[l]*YftI[l]+YffI[l]*YftR[l]
        @NLconstraint(
            opfmodel,
            Vm[busIdx[lines[l].from]]^2 *
            ( Yff_abs2*Vm[busIdx[lines[l].from]]^2 + Yft_abs2*Vm[busIdx[lines[l].to]]^2
            + 2*Vm[busIdx[lines[l].from]]*Vm[busIdx[lines[l].to]]
                *(Yre*cos(Va[busIdx[lines[l].from]]-Va[busIdx[lines[l].to]])
                -Yim*sin(Va[busIdx[lines[l].from]]-Va[busIdx[lines[l].to]])
                )
            )
            - flowmax <=0)

        #branch apparent power limits (to bus)
        Ytf_abs2=YtfR[l]^2+YtfI[l]^2; Ytt_abs2=YttR[l]^2+YttI[l]^2
        Yre=YtfR[l]*YttR[l]+YtfI[l]*YttI[l]; Yim=-YtfR[l]*YttI[l]+YtfI[l]*YttR[l]
        @NLconstraint(
            opfmodel,
            Vm[busIdx[lines[l].to]]^2 *
            ( Ytf_abs2*Vm[busIdx[lines[l].from]]^2 + Ytt_abs2*Vm[busIdx[lines[l].to]]^2
            + 2*Vm[busIdx[lines[l].from]]*Vm[busIdx[lines[l].to]]
                *(Yre*cos(Va[busIdx[lines[l].from]]-Va[busIdx[lines[l].to]])
                -Yim*sin(Va[busIdx[lines[l].from]]-Va[busIdx[lines[l].to]])
                )
            )
            - flowmax <=0)
        end
    end
    println("Buses: $nbus  Lines: $nline  Generators: $ngen")
    println("Lines with limits: $nlinelim")

    return opfmodel, Pg, Qg, Va, Vm
end
# Compute initial point for IPOPT based on the values provided in the case data
function initialPt_IPOPT(opfdata)
    Pg=zeros(length(opfdata.generators)); Qg=zeros(length(opfdata.generators)); i=1
    for g in opfdata.generators
        # set the power levels in in between the bounds as suggested by matpower
        # (case data also contains initial values in .Pg and .Qg - not used with IPOPT)
        Pg[i]=0.5*(g.Pmax+g.Pmin)
        Qg[i]=0.5*(g.Qmax+g.Qmin)
        i=i+1
    end
    @assert i-1==length(opfdata.generators)

    Vm=zeros(length(opfdata.buses)); i=1;
    for b in opfdata.buses
        # set the ini val for voltage magnitude in between the bounds
        # (case data contains initials values in Vm - not used with IPOPT)
        Vm[i]=0.5*(b.Vmax+b.Vmin);
        i=i+1
    end
    @assert i-1==length(opfdata.buses)

    # set all angles to the angle of the reference bus
    Va = opfdata.buses[opfdata.bus_ref].Va * ones(length(opfdata.buses))

    return Pg,Qg,Vm,Va
end

# cases with objective scaling
# Dense Hiop algebra
for case in [[0.0283286,"case9"],[0.00357143, "case118"],[0.00645161, "case300"]]
    casepath = joinpath(dirname(@__FILE__), "data", case[2])
    max_iter=100
    opfdata = opf_loaddata(casepath)
    Pg0, Qg0, Vm0, Va0 = initialPt_IPOPT(opfdata)
    opfmodel_ref, Pg_ref, Qg_ref, Va_ref, Vm_ref = model(opfdata; solver="Ipopt", scaling = case[1])
    opfmodel, Pg, Qg, Va, Vm = model(opfdata; solver="Hiop", hiopalgebra = :Dense, scaling = case[1])
    opfmodel_ref ,status = solve(opfmodel_ref,opfdata)
    opfmodel,status = solve(opfmodel,opfdata)
    @test value.(Pg) ≈ value.(Pg)
    @test value.(Qg) ≈ value.(Qg)
    @test value.(Vm) ≈ value.(Vm)
    @test value.(Va) ≈ value.(Va)
    @test objective_value(opfmodel) ≈ objective_value(opfmodel_ref)
end

# Sparse Hiop algebra
# This works with sparse algebra
for case in [[0.0283286,"case9"],[0.00357143, "case118"]]
    casepath = joinpath(dirname(@__FILE__), "data", case[2])
    max_iter=100
    opfdata = opf_loaddata(casepath)
    Pg0, Qg0, Vm0, Va0 = initialPt_IPOPT(opfdata)
    opfmodel_ref, Pg_ref, Qg_ref, Va_ref, Vm_ref = model(opfdata; solver="Ipopt", scaling = case[1])
    opfmodel, Pg, Qg, Va, Vm = model(opfdata; solver="Hiop", hiopalgebra = :Sparse, scaling = case[1])
    opfmodel_ref ,status = solve(opfmodel_ref,opfdata)
    opfmodel,status = solve(opfmodel,opfdata)
    @test value.(Pg) ≈ value.(Pg)
    @test value.(Qg) ≈ value.(Qg)
    @test value.(Vm) ≈ value.(Vm)
    @test value.(Va) ≈ value.(Va)
    @test objective_value(opfmodel) ≈ objective_value(opfmodel_ref)
end

# Crashes
# for case in [[0.00645161, "case300"]]
#     casepath = joinpath(dirname(@__FILE__), "data", case[2])
#     max_iter=100
#     opfdata = opf_loaddata(casepath)
#     Pg0, Qg0, Vm0, Va0 = initialPt_IPOPT(opfdata)
#     opfmodel_ref, Pg_ref, Qg_ref, Va_ref, Vm_ref = model(opfdata; solver="Ipopt", scaling = case[1])
#     opfmodel, Pg, Qg, Va, Vm = model(opfdata; solver="Hiop", hiopalgebra = :Sparse, scaling = case[1])
#     opfmodel_ref ,status = solve(opfmodel_ref,opfdata)
#     opfmodel,status = solve(opfmodel,opfdata)
#     @test_broken value.(Pg) ≈ value.(Pg)
#     @test_broken value.(Qg) ≈ value.(Qg)
#     @test_broken value.(Vm) ≈ value.(Vm)
#     @test_broken value.(Va) ≈ value.(Va)
#     @test_broken objective_value(opfmodel) ≈ objective_value(opfmodel_ref)
# end


using DoubletFinder
using Test

@testset "DoubletFinder" begin
test_array=[2.1,3.0,0.8,0.7,0.5,0.4,0.2,1.01,1.33,3.88,0.99,0.96,0.98,1.23,4.55]
@testset "Testing linear algebra results" begin
   @test 1 == 1
end;

@testset "findV" begin
    @test findV(3.12,3.12) == 0.0
    @test findV(3.0,3.12) ≈ 8858.92604 atol=1e-4
end;

@testset "findZAbs" begin
    @test findZAbs(0.0, 3.12) == 3.12
    @test findZAbs(1.0, 3.12) ≈ 3.119986 atol=1e-4
end;

@testset "getblocks1" begin
    @test getblocks1(test_array)[1] == (3, 7)
    @test getblocks1(test_array)[2] == (11, 13)
end;
@testset "iscontained" begin
    @test iscontained(test_array,66.0) == false
    @test iscontained(test_array,3.25) == true
end;

    end; #DoubletFinder


##################################################
# Setup
##################################################

bnd = MCBN.BayesNetDAI(5)
bnd2 = MCBN.BayesNetDAI(5)
bnd3 = MCBN.BayesNetDAI(5)

##################################################
# Entropy and KLD Tests
##################################################

@test_approx_eq MCBN.entropy(bnd) 5.0
@test_approx_eq MCBN.naive_entropy(bnd) 5.0
@test_approx_eq MCBN.kld(bnd,bnd2) 0.0

function test_net(bnd)
    MCBN.check_bnd(bnd)
    @test_approx_eq MCBN.naive_entropy(bnd) MCBN.entropy(bnd)
    @test_approx_eq MCBN.kld(bnd,bnd) 0.0
end

function test_kld(bnd1, bnd2)
    kldval = MCBN.kld(bnd1, bnd2)
    kldvalrev = MCBN.kld(bnd2, bnd1)
    @test kldval >= 0
    @test kldvalrev >= 0
end

test_net(bnd)
test_kld(bnd, bnd2)

MCBN.add_edge!(bnd3, 1, 2)
#MCBN.check_bnd(bnd3)

MCBN.set_factor!(bnd3, 2, [1.0, 0.5, 0.0, 0.5], false)
MCBN.set_factor!(bnd3, 1, [0.8, 0.2])
MCBN.check_bnd(bnd3)

@test_approx_eq MCBN.entropy(bnd3) 3.9219280948873623
@test_approx_eq MCBN.entropy(bnd3) MCBN.naive_entropy(bnd3)

test_kld(bnd, bnd3)

MCBN.restoreFactors!(bnd3.fg)
bnd3.dirty = true

@test_approx_eq MCBN.entropy(bnd3) MCBN.naive_entropy(bnd3)
@test_approx_eq MCBN.entropy(bnd3) 5.0

@test_approx_eq MCBN.kld(bnd,bnd3) 0.0
@test_approx_eq MCBN.kld(bnd3,bnd) 0.0

test_net(bnd3)
test_kld(bnd, bnd3)

##################################################
# Entropy and KLD Tests
##################################################

bns = MCBN.BayesNetSampler(3, rand(1:2, 5,3))
origE = MCBN.energy(bns)
origfvalue = copy(bns.fvalue)

@time for i=1:1000
    scheme = MCBN.propose!(bns)
    MCBN.check_bns(bns)
    MCBN.reject!(bns)
    MCBN.check_bns(bns)
    E = MCBN.energy(bns)
    @test_approx_eq bns.fvalue origfvalue
    @test_approx_eq E origE
end

@time for i=1:1000
    s = MCBN.propose!(bns)
    MCBN.check_bns(bns)
    E1 = MCBN.energy(bns)
    f1 = copy(bns.fvalue)
    bns.changelist = [1:3;]
    E2 = MCBN.energy(bns)
    f2 = copy(bns.fvalue)
    MCBN.save!(bns)
    @test_approx_eq f1 f2
    @test_approx_eq E1 E2
end


function eye(n::Int64)
    res = zeros(n,n)
    for i = 1:n res[i,i] = 1.0 end
    return res
end
#Pkg FillArrays (Eye) ?

function eye(dtype::DataType, n::Int64)
    res = zeros(dtype, n, n)
    I = 1.0 |> dtype
    for i = 1:n res[i,i] = I end
    return res
end

function delta(x::Real; c::Real = 0., η::Real = 0.05)
    num = η
    den = ((x-c)^2 + η^2) * π
    return num/den
end

function gaussian(x::Real, fac::AbstractVector)
    # f(x) = 1/(σ√(2π)) exp(-((x-μ)/σ)^2 /2 )
    μ = fac[1]; σ2 = fac[2]
    σ2 <= 0. ? (@error "σ2 should be positive.") :
    norm = √(2π * σ2)
    m = (x-μ)^2 / σ2
    return 1/norm*exp(-m/2)
end

function multi_gaussian(N::Int64; ms::Vector{Float64}=[-3.,3.], smax::Real=3.)
    fac = rand(N,2)
    fac[:,1] = rand(N) .* (ms[2] - ms[1]) .+ ms[1]
    fac[:,2] = rand(N) .* smax
    function f(x::Real)
        res = 0.
        for i = 1: N
            res += gaussian(x, fac[i,:])
        end
        res/N
    end
    return f, fac
end

function Masubara_freq(n::Int64, β::Real; type::Symbol= :f)
    if type == :b  N = 2n
    elseif type == :f  N = 2n + 1
    else @error "type should be :b for bosons and :f for fermions"
    end
    return N*π/β
end

function Masubara_GF(n::Int64, A::Function, β::Real;
    type::Symbol = :f, Λ::Float64=100., L::Int64=1000000)
    # G(iωn) = 1/2π ∫dΩ A(Ω)/(iωn - Ω)
    ωn = Masubara_freq(n, β, type=type)
    dL = 2*Λ/L; lowlim = -Λ
    Gn = 0.
    for n = 0:L
        ω = lowlim + dL*n
        Gn += A(ω)/(1.0im*ωn-ω) * dL
    end
    return Gn/(2π)
end

function h(z::Number)
    #Conformal mapping(Mobius transform)
    #\bar{C+} -> \bar{D}
    return (z - 1.0im)/(z + 1.0im)
end

function invh(z::Number)
    return 1.0im*(1 + z)/(1 - z)
end

function h1(z::Number, Y::Number)
    #Conformal mapping(Mobius transform)
    #C+ -> D
    return (z - Y)/(z - Y')
end

function invh1(z::Complex, Y::Number)
    return (z*Y'-Y)/(z-1)
end

function ispossemidef(A::Matrix)
    evals = eigvals(A)
    return all(evals .>= 0)
end


# Canonical form of multivariate normal

## generic types

immutable GenericMvNormalCanon{Prec<:AbstractPDMat} <: AbstractMvNormal
	dim::Int             # dimension of sample space
	zeromean::Bool       # whether the mean vector is zero
	μ::Vector{Float64}   # the mean vector
    h::Vector{Float64}   # potential vector, i.e. inv(Σ) * μ
    J::Prec              # precision matrix, i.e. inv(Σ)

    function GenericMvNormalCanon(μ::Vector{Float64}, h::Vector{Float64}, J::Prec, zmean::Bool)
    	d = dim(J)
    	length(μ) == length(h) == d || throw(ArgumentError("Inconsistent argument dimensions."))
    	new(d, zmean, μ, h, J)
    end

	function GenericMvNormalCanon(J::Prec)
		d = dim(J)
    	new(d, true, zeros(d), zeros(d), J)
    end
end

function GenericMvNormalCanon{P<:AbstractPDMat}(μ::Vector{Float64}, h::Vector{Float64}, J::P)
	GenericMvNormalCanon{P}(μ, h, J, allzeros(μ))
end

function GenericMvNormalCanon{P<:AbstractPDMat}(h::Vector{Float64}, J::P, zmean::Bool)
	μ = zmean ? zeros(length(h)) : (J \ h)
	GenericMvNormalCanon{P}(μ, h, J, zmean)
end

function GenericMvNormalCanon{P<:AbstractPDMat}(h::Vector{Float64}, J::P)
	GenericMvNormalCanon(h, J, allzeros(h))
end

function GenericMvNormalCanon{P<:AbstractPDMat}(J::P)
    d = dim(J)
    GenericMvNormalCanon{P}(zeros(d), zeros(d), J, true)
end

## type aliases and convenient constructors

typealias MvNormalCanon   GenericMvNormalCanon{PDMat}
typealias DiagNormalCanon GenericMvNormalCanon{PDiagMat}
typealias IsoNormalCanon  GenericMvNormalCanon{ScalMat}

MvNormalCanon(J::PDMat) = GenericMvNormalCanon(J)
MvNormalCanon(J::Matrix{Float64}) = GenericMvNormalCanon(PDMat(J))
MvNormalCanon(h::Vector{Float64}, J::PDMat) = GenericMvNormalCanon(h, J)
MvNormalCanon(h::Vector{Float64}, J::Matrix{Float64}) = GenericMvNormalCanon(h, PDMat(J))

DiagNormalCanon(J::PDiagMat) = GenericMvNormalCanon(J)
DiagNormalCanon(J::Vector{Float64}) = GenericMvNormalCanon(PDiagMat(J))
DiagNormalCanon(h::Vector{Float64}, J::PDiagMat) = GenericMvNormalCanon(h, J)
DiagNormalCanon(h::Vector{Float64}, J::Vector{Float64}) = GenericMvNormalCanon(h, PDiagMat(J))

IsoNormalCanon(J::ScalMat) = GenericMvNormalCanon(J)
IsoNormalCanon(d::Integer, prec::Real) = GenericMvNormalCanon(ScalMat(int(d), float64(prec)))
IsoNormalCanon(h::Vector{Float64}, J::ScalMat) = GenericMvNormalCanon(h, J)
IsoNormalCanon(h::Vector{Float64}, prec::Real) = GenericMvNormalCanon(h, ScalMat(length(h), float64(prec)))

# conversion between conventional form and canonical form

function Base.convert{C<:AbstractPDMat}(D::Type{GenericMvNormal{C}}, cf::GenericMvNormalCanon{C})
	GenericMvNormal{C}(cf.dim, cf.zeromean, cf.μ, inv(cf.J))
end

function Base.convert{C<:AbstractPDMat}(D::Type{GenericMvNormalCanon{C}}, d::GenericMvNormal{C})
	J::C = inv(d.Σ)
	h::Vector{Float64} = J * d.μ
	GenericMvNormalCanon{C}(d.μ, h, J, d.zeromean)
end

canonform{C<:AbstractPDMat}(d::GenericMvNormal{C}) = convert(GenericMvNormalCanon{C}, d)

# Basic statistics

dim(d::GenericMvNormalCanon) = d.dim

mean(d::GenericMvNormalCanon) = d.μ
mode(d::GenericMvNormalCanon) = d.μ
modes(d::GenericMvNormalCanon) = [mode(d)]

var(d::GenericMvNormalCanon) = diag(inv(d.J))
cov(d::GenericMvNormalCanon) = full(inv(d.J))
invcov(d::GenericMvNormalCanon) = full(d.J)
logdet_cov(d::GenericMvNormalCanon) = -logdet(d.J)

entropy(d::GenericMvNormalCanon) = 0.5 * (dim(d) * (float64(log2π) + 1.0) - logdet(d.J))

# PDF evaluation

function sqmahal(d::GenericMvNormalCanon, x::Vector{Float64})
    z::Vector{Float64} = d.zeromean ? x : x - d.μ
    quad(d.J, z)
end

function sqmahal!(r::Array{Float64}, d::GenericMvNormalCanon, x::Matrix{Float64})
    if !(size(x, 1) == dim(d) && size(x, 2) == length(r))
        throw(ArgumentError("Inconsistent argument dimensions."))
    end
    z::Matrix{Float64} = d.zeromean ? x : bsubtract(x, d.μ, 1)
    quad!(r, d.J, z)
end

# Sampling (for GenericMvNormal)

function rand!(d::GenericMvNormalCanon, x::Vector{Float64})
    unwhiten_winv!(d.J, randn!(x))
    if !d.zeromean
        add!(x, d.μ)
    end
    x
end

function rand!(d::GenericMvNormalCanon, x::Matrix{Float64})
    unwhiten_winv!(d.J, randn!(x))
    if !d.zeromean
        badd!(x, d.μ, 1)
    end
    x
end


using Zygote

begin
W = rand(2, 3); x = rand(3);
@show gradient(W -> sum(W*x), W)
# Since we only have W as parameter, the first element is the grad
# a 2x3 array
@show gradient(W -> sum(W*x), W)[1]
end

begin
d = Dict()
# Note the do is a x-> ....  function
# Gradient of that function is calculated at x=5
# function acts like x**2  so  its gradient is 2x=2*5=10
@show  gradient(5) do x
         d[:x] = x
         d[:x] * d[:x]
       end

@show  d[:x]
end


"""
    HubbardReal1D(address; u=1.0, t=1.0)

Implements a one-dimensional Bose Hubbard chain in real space.

```math
\\hat{H} = -t \\sum_{\\langle i,j\\rangle} a_i^† a_j + \\frac{u}{2}\\sum_i n_i (n_i-1)
```

# Arguments

* `address`: the starting address, defines number of particles and sites.
* `u`: the interaction parameter.
* `t`: the hopping strength.

# See also

* [`HubbardMom1D`](@ref)
* [`ExtendedHubbardReal1D`](@ref)

"""
struct HubbardReal1D{TT,A<:AbstractFockAddress,U,T} <: AbstractHamiltonian{TT}
    add::A
end

function HubbardReal1D(addr; u=1.0, t=1.0)
    U, T = promote(float(u), float(t))
    return HubbardReal1D{typeof(U),typeof(addr),U,T}(addr)
end

function Base.show(io::IO, h::HubbardReal1D)
    print(io, "HubbardReal1D($(h.add); u=$(h.u), t=$(h.t))")
end

function starting_address(h::HubbardReal1D)
    return getfield(h, :add)
end

LOStructure(::Type{<:HubbardReal1D{<:Real}}) = IsHermitian()

Base.getproperty(h::HubbardReal1D, s::Symbol) = getproperty(h, Val(s))
Base.getproperty(h::HubbardReal1D{<:Any,<:Any,U}, ::Val{:u}) where U = U
Base.getproperty(h::HubbardReal1D{<:Any,<:Any,<:Any,T}, ::Val{:t}) where T = T
Base.getproperty(h::HubbardReal1D, ::Val{:add}) = getfield(h, :add)

function num_offdiagonals(::HubbardReal1D, address::BoseFS)
    return 2 * num_occupied_modes(address)
end

"""
    bose_hubbard_interaction(address)

Return Σ_i *n_i* (*n_i*-1) for computing the Bose-Hubbard on-site interaction (without the
*U* prefactor.)

# Example

```jldoctest
julia> Hamiltonians.bose_hubbard_interaction(BoseFS{4,4}((2,1,1,0)))
2
julia> Hamiltonians.bose_hubbard_interaction(BoseFS{4,4}((3,0,1,0)))
6
```
"""
function bose_hubbard_interaction(b::BoseFS{<:Any,<:Any,A}) where A
    return bose_hubbard_interaction(Val(num_chunks(A)), b)
end

@inline function bose_hubbard_interaction(_, b::BoseFS)
    result = 0
    for (n, _, _) in occupied_modes(b)
        result += n * (n - 1)
    end
    return result
end

@inline function bose_hubbard_interaction(::Val{1}, b::BoseFS)
    # currently this ammounts to counting occupation numbers of modes
    chunk = chunks(b.bs)[1]
    matrixelementint = 0
    while !iszero(chunk)
        chunk >>>= (trailing_zeros(chunk) % UInt) # proceed to next occupied mode
        bosonnumber = trailing_ones(chunk) # count how many bosons inside
        # surpsingly it is faster to not check whether this is nonzero and do the
        # following operations anyway
        chunk >>>= (bosonnumber % UInt) # remove the counted mode
        matrixelementint += bosonnumber * (bosonnumber - 1)
    end
    return matrixelementint
end

function diagonal_element(h::HubbardReal1D, address::BoseFS)
    h.u * bose_hubbard_interaction(address) / 2
end

function get_offdiagonal(h::HubbardReal1D, add::BoseFS, chosen)
    naddress, onproduct = hopnextneighbour(add, chosen)
    return naddress, - h.t * sqrt(onproduct)
end


module Solvers

using LinearAlgebra
using NLSolvers,Roots
using PositiveFactorizations
using DiffResults, ForwardDiff
using StaticArrays
import NaNMath
    function cholesky_linsolve(d,B,∇f)
        cholesky!(Positive, B)
        Bchol = Cholesky(B,'L',0)
        d .=  Bchol\∇f
    end

    function cholesky_linsolve(B,∇f)
        Bchol =cholesky(Positive, B)
        Bchol\∇f
    end

    #
    Base.summary(::NLSolvers.Newton{<:Direct, typeof(cholesky_linsolve)}) = "Newton's method with Cholesky linsolve"
    CholeskyNewton() = NLSolvers.Newton(linsolve=cholesky_linsolve)
    export CholeskyNewton

    """
        det_22(a,b,c,d)

    Calculates `a*b - c*d` without less rounding error than doing it naively
    """
    function det_22(a,b,c,d)
        t = c*d
        e = muladd(c,d,-t) #cd - cd
        f = muladd(a,b,-t) #ab - cd
        return f-e  #ab - cd + cd - cd
    end

    function solve_cubic_eq(poly::AbstractVector{T}) where {T<:Real}
        # copied from PolynomialRoots.jl, adapted to be AD friendly
        # Cubic equation solver for complex polynomial (degree=3)
        # http://en.wikipedia.org/wiki/Cubic_function   Lagrange's method
        # poly = (a,b,c,d) that represents ax3 + bx2 + cx2 + d

        _1 = one(T)
        third = _1/3
        a1  =  complex(one(T) / poly[4])
        E1  = -poly[3]*a1
        E2  =  poly[2]*a1
        E3  = -poly[1]*a1
        s0  =  E1
        E12 =  E1*E1
        A   =  det_22(2*E1,E12,9*E1,E2) + 27*E3
        #A   =  2*E1*E12 - 9*E1*E2 + 27*E3 # = s1^3 + s2^3
        B = det_22(E1,E1,3,E2)
        #B   =  E12 - 3*E2                 # = s1 s2
        # quadratic equation: z^2 - Az + B^3=0  where roots are equal to s1^3 and s2^3
        Δ2 = det_22(A,A,4*B*B,B)
        Δ = sqrt(Δ2)
        #Δ = (A*A - 4*B*B*B)^0.5
        if real(conj(A)*Δ)>=0 # scalar product to decide the sign yielding bigger magnitude
            s1 = exp(log(0.5 * (A + Δ)) * third)
        else
            s1 = exp(log(0.5 * (A - Δ)) * third)
        end
        if s1 == 0
            s2 = s1
        else
            s2 = B / s1
        end
        zeta1 = complex(-0.5, sqrt(T(3.0))*0.5)
        zeta2 = conj(zeta1)
        return (third*(s0 + s1 + s2), third*(s0 + s1*zeta2 + s2*zeta1), third*(s0 + s1*zeta1 + s2*zeta2))
    end

    """
        roots3(pol)
    solves a cubic equation of the form pol[1] + pol[2]*x + pol[3]*x^2 + pol[4]*x^3
    """
    function roots3(pol)
        return SVector(solve_cubic_eq(pol))
    end

    function roots3(a,b,c,d)
        x = SVector(a,b,c,d)
        return roots3(x)
    end

    """
    function x_sol(res::NLSolvers.ConvergenceInfo)

        Returns the scalar or vector x that solves the system of equations or is the minimizer of an optimization procedure.
    """

    function x_sol(res::NLSolvers.ConvergenceInfo{<:NLSolvers.LineSearch, <:Any, <:NLSolvers.NEqOptions})
        res.info.solution
    end

    function x_sol(res::NLSolvers.ConvergenceInfo{<:NLSolvers.TrustRegion, <:Any, <:NLSolvers.NEqOptions})
        return res.info.zero
    end

    function x_sol(res::NLSolvers.ConvergenceInfo{<:Any, <:Any, <:NLSolvers.OptimizationOptions})
        return res.info.minimizer
    end

    @inline log(x::Number) = Base.log(x)
    @inline log(x::Complex) = Base.log(x)
    @inline log(x::Real) = ifelse(x>=zero(x),Base.log(x),zero(x)/zero(x))
    @inline log(x::Float64) = NaNMath.log(x)
    @inline log(x::Float32) = NaNMath.log(x)
    @inline log(x::Int8) = NaNMath.log(float(x))
    @inline log(x::Int16) = NaNMath.log(float(x))
    @inline log(x::Int32) = NaNMath.log(float(x))
    @inline log(x::Int64) = NaNMath.log(float(x))

    include("ADNewton.jl")
    include("nlsolve.jl")
    include("fixpoint/fixpoint.jl")
    include("optimize.jl")

end # module


using Colorsys
using Test

@testset "Colorsys.jl" begin
    @testset "_v" begin
        @test Colorsys._v(1,2,0.1) ≈ 1.6
        @test Colorsys._v(1,2,0.4) == 2
        @test Colorsys._v(1,2,0.6) ≈ 1.4
        @test Colorsys._v(1,2,0.8) == 1
    end
    y,i,q = rgb2yiq(120, 130, 140)
    @test y≈0.5023529411 && i≈-0.0360784313 && q≈0.0039215686
    r,g,b = yiq2rgb(0.8, 0.6, 0.4)
    @test (r,g,b) == (255,96,211)
    @testset "rgb2hls" begin
        h,l,s = rgb2hls(120, 130, 140)
        @test h≈210.0000000000 && l≈0.5098039215 && s≈0.0800000000
        h,l,s = rgb2hls(120, 120, 120)
        @test h≈0.0000000000 && l≈0.4705882352 && s≈0.0000000000
    end
    @testset "hls2rgb" begin
        r,g,b = hls2rgb(120, 0.6, 0.4)
        @test (r,g,b) == (112,193,112)
        r,g,b = hls2rgb(120, 0.6, 0.0)
        @test (r,g,b) == (0,0,0)
    end
    @testset "rgb2hsv" begin
        h,s,v = rgb2hsv(120, 130, 140)
        @test h≈210.0000000000 && s≈0.1428571428 && v≈0.5490196078
        h,s,v = rgb2hsv(120, 120, 120)
        @test h≈0.0 && s≈0.0 && v≈0.4705882352
        h,s,v = rgb2hsv(220, 130, 240)
        @test h≈289.0909090909 && s≈0.4583333333 && v≈0.9411764705
    end
    @testset "hsv2rgb" begin
        r,g,b = hsv2rgb(40, 0.6, 0.4)
        @test (r,g,b) == (102,81,40)
        r,g,b = hsv2rgb(80, 0.6, 0.4)
        @test (r,g,b) == (81,102,40)
        r,g,b = hsv2rgb(120, 0.6, 0.4)
        @test (r,g,b) == (40,102,40)
        r,g,b = hsv2rgb(200, 0.6, 0.4)
        @test (r,g,b) == (40,81,102)
        r,g,b = hsv2rgb(260, 0.6, 0.4)
        @test (r,g,b) == (61,40,102)
        r,g,b = hsv2rgb(320, 0.6, 0.4)
        @test (r,g,b) == (102,40,81)
        r,g,b = hsv2rgb(120, 0, 0.2)
        @test (r,g,b) == (51,51,51)
    end
    h,l,s = yiq2hls(0.8, 0.6, 0.4)
    @test h≈316.6037735849 && l≈0.6882352941 && s≈1.0000000000
    h,s,v = yiq2hsv(0.8, 0.6, 0.4)
    @test h≈316.6037735849 && s≈0.6235294117 && v≈1.0000000000
    y,i,q = hls2yiq(120, 0.6, 0.4)
    @test y≈0.6266274509 && i≈-0.0889411764 && q≈-0.1651764705
    h,s,v = hls2hsv(120, 0.6, 0.4)
    @test h≈120.0000000000 && s≈0.4196891191 && v≈0.7568627450
    y,i,q = hsv2yiq(120, 0.6, 0.4)
    @test y≈0.3003137254 && i≈-0.0680784313 && q≈-0.1264313725
    h,l,s = hsv2hls(120, 0.6, 0.4)
    @test h≈120.0000000000 && l≈0.2784313725 && s≈0.4366197183
end


using Mamba

## Data
equiv = Dict{Symbol, Any}(
  :group => [1, 1, 2, 2, 2, 1, 1, 1, 2, 2],
  :y =>
    [1.40 1.65
     1.64 1.57
     1.44 1.58
     1.36 1.68
     1.65 1.69
     1.08 1.31
     1.09 1.43
     1.25 1.44
     1.25 1.39
     1.30 1.52]
)
equiv[:N] = size(equiv[:y], 1)
equiv[:P] = size(equiv[:y], 2)

equiv[:T] = [equiv[:group] 3 - equiv[:group]]

## Model Specification
model = Model(

  y = Stochastic(2,
    (delta, mu, phi, pi, s2_1, T) ->
      begin
        sigma = sqrt(s2_1)
        UnivariateDistribution[
          begin
            m = mu + (-1)^(T[i, j] - 1) * phi / 2 + (-1)^(j - 1) * pi / 2 +
                delta[i, j]
            Normal(m, sigma)
          end
          for i in 1:10, j in 1:2
        ]
      end,
    false
  ),

  delta = Stochastic(2,
    s2_2 -> Normal(0, sqrt(s2_2)),
    false
  ),

  mu = Stochastic(
    () -> Normal(0, 1000)
  ),

  phi = Stochastic(
    () -> Normal(0, 1000)
  ),

  theta = Logical(
    phi -> exp(phi)
  ),

  pi = Stochastic(
    () -> Normal(0, 1000)
  ),

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

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

  equiv = Logical(
    theta -> Int(0.8 < theta < 1.2)
  )

)

## Initial Values
inits = [
  Dict(:y => equiv[:y], :delta => zeros(10, 2), :mu => 0, :phi => 0,
       :pi => 0, :s2_1 => 1, :s2_2 => 1),
  Dict(:y => equiv[:y], :delta => zeros(10, 2), :mu => 10, :phi => 10,
       :pi => 10, :s2_1 => 10, :s2_2 => 10)
]

## Sampling Scheme
scheme = [NUTS(:delta),
          Slice([:mu, :phi, :pi], 1.0),
          Slice([:s2_1, :s2_2], 1.0, Univariate)]
setsamplers!(model, scheme)

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


@testset "Cubic equation" begin
    Base.info("Testing cubic equation...")
    solver = Solver()
    @test isapprox(solver(1.0,-3.0,3.0,-1.0), ones(3))
    @test isapprox(solver(2.0,-3.0,-5.0,6.0), [-1.5,1.0,2.0])
    @test isapprox.(solver(1.0,1.0,1.0,-3.0), Union{Real,Complex}[1.0,-1.0-im*√(2),-1.0+im*√(2)]) == ones(Bool, 3)
    @test_throws ErrorException solver(0.0,1.0,1.0,-3.0)
    Base.info("Done testing cubic equation.")
end


# Define `training`, `data_dir`, `eki`

using Statistics
using CairoMakie
using ProgressBars

using ParameterEstimocean.PseudoSteppingSchemes: trained_gp_predict_function, ensemble_array
using ParameterEstimocean.Transformations: ZScore, normalize!, inverse_normalize!
using ParameterEstimocean.Parameters: transform_to_constrained

# Specify a directory to which to save the files generated in this script
dir = joinpath(directory, "emulate_sample_forward_map")
isdir(dir) || mkdir(dir)

"""
    truncate_forward_map_to_length_k_uncorrelated_points(k)

# Arguments
- G: d × M array whose M columns are individual forward map outputs ∈ Rᵈ.
- k: The number of dimensions to reduce the forward map outputs to.
# Returns
- Ĝ: k × M array whose M columns are individual forward map outputs in an
uncorrelated output space Rᵏ computed by PCA analysis using SVD as described
in Appendix A.2. of Cleary et al. "Calibrate Emulate Sample" (2021).
"""
function truncate_forward_map_to_length_k_uncorrelated_points(G, y, k)

    d, M = size(G)
    m = mean(G, dims=2) # d × 1

    # Center the columns of G at zero
    Gᵀ_centered = (G .- m)'

    # SVD
    # Gᵀ = Ĝᵀ Σ Vᵀ
    # (M × d) = (M × d)(d × d)(d × d)
    F = svd(Gᵀ_centered; full=false)
    Ĝᵀ = F.U
    Σ = Diagonal(F.S)
    Vᵀ = F.Vt

    @assert Gᵀ_centered ≈ Ĝᵀ * Σ * Vᵀ
    @show size(Gᵀ_centered), M, d
    @show size(Ĝᵀ), M, d
    @show size(Σ), d, d
    @show size(Vᵀ), d, d

    # Eigenvalue sum
    total_eigenval = sum(Σ .^ 2)

    # Keep only the `k` < d most important dimensions
    # corresponding to the `k` highest singular values.
    # Dimensions become (M × d) = (M × k)(k × k)(k × d)
    Ĝᵀ = Ĝᵀ[:, 1:k]
    Σ = Σ[1:k, 1:k]
    Vᵀ = Vᵀ[1:k, :]

    @info "Preserved $(sum(Σ .^ 2)*100/total_eigenval)% of the original variance in the output data by
           reducing the dimensionality of the output space from $d to $k."

    Ĝ = Ĝᵀ'
    #      Ĝᵀ = (Gᵀ - mᵀ) V Σ⁻¹
    # (M × k) = (M × d)(d × k)(k × k)
    # Therefore
    #       Ĝ = Σ⁻¹ Vᵀ (G - m)
    # (k × M) = (k × k)(k × d)(d × M)
    #
    # Therefore to transform the observations,
    # ŷ = Σ⁻¹ Vᵀ (y - m)
    ŷ = (1 ./ Σ) * Vᵀ * (y - m)
    return Ĝ, ŷ
end

# First, conglomerate all samples generated thus far by EKI.
# This will be the training data for the GP emulator.
# We will filter out all failed particles.
n = 5
# n = eki.iteration

# X = hcat([constrained_ensemble_array(eki, i) for i in 0:n]...) # constrained
X = hcat([ensemble_array(eki, iter) for iter in 0:n]...) # unconstrained
# G = hcat(getproperty.(eki.iteration_summaries, :forward_map_output)...)
@info "Performing emulation based on samples from the first $n iterations of EKI."
G = hcat([summary.forward_map_output for summary in eki.iteration_summaries[0:5]]...)

nan_values = vec(mapslices(any, isnan.(G); dims=1)) # bitvector
not_nan_indices = findall(.!nan_values) # indices of columns (particles) with no `NaN`s
X = X[:, not_nan_indices]
G = G[:, not_nan_indices]

# Transform forward map output samples to uncorrelated space.
# This will allow us to use the uncertainty estimates from each
# GP in the emulator.
k = 20
y = eki.mapped_observations
Ĝ, ŷ = truncate_forward_map_to_length_k_uncorrelated_points(G, y, k)

# We will approximately non-dimensionalize the inputs according to mean and variance
# computed across all generated training samples.
zscore_X = ZScore(mean(X, dims=2), std(X, dims=2))
normalize!(X, zscore_X)

# Ensemble covariance across all generated samples (diagonal should be all ones)
cov_θθ_all_iters = cov(X, X, dims = 2, corrected = false)

###
### Emulate and Sample
###

# The likelihood we wish to sample with MCMC is π(θ|y)=exp(-Φ(θ)), the posterior density on θ given y.
# The MCMC sampler takes in a function `nll` which maps θ to the negative log value Φ(θ).
# In the following example, we use a GP to emulate the forward map output G.

# We will take advantage of the parallelizability of our forward map
# by running parallel chains of MCMC.
n_chains = 128

# Length and burn-in length per chain
chain_length = 1000
burn_in = 100

chain_length_emulate = 10000
burn_in_emulate = 1000

###
### Sample from emulated loss landscape using parallel chains of MCMC
###

# vector of predict functions. Ĝ is k x Nsamples
# predicts = [trained_gp_predict_function(Ĝ[i,:]) for i in size(Ĝ,1)]
predicts = []

@info "Training $k gaussian processes for the emulator."
validation_results=[]
for i in ProgressBar(1:k) # forward map index

    # Values of the forward maps of each sample at index `i`
    yᵢ = Ĝ[i, :]

    # Reserve `validation_fraction` representative samples for the emulator
    # We will sort `yᵢ` and take evenly spaced samples so the samples are representative.
    M = length(yᵢ)
    lq = Int(round(M/4))
    uq = lq*3
    evenly_spaced_samples = Int.(round.(range(lq, uq, length = Int(round(0.03M)))))
    emulator_validation_indices = sort(eachindex(yᵢ), by = i -> yᵢ[i])[evenly_spaced_samples]
    not_emulator_validation_indices = [i for i in 1:M if !(i in emulator_validation_indices)]
    X_validation = X[:, emulator_validation_indices]
    yᵢ_validation = yᵢ[emulator_validation_indices]

    predict = trained_gp_predict_function(X[:, not_emulator_validation_indices], yᵢ[not_emulator_validation_indices]; standardize_X = false, zscore_limit = nothing)
    push!(predicts, predict)

    ŷᵢ_validation, Γgp_validation = predict(X_validation)
    push!(validation_results, (yᵢ_validation, ŷᵢ_validation, diag(Γgp_validation)))
end

n_columns = 5
N_axes = k
n_rows = Int(ceil(N_axes / n_columns))
fig = Figure(resolution = (300n_columns, 350n_rows), fontsize = 8)
ax_coords = [(i, j) for j = 1:n_columns, i = 1:n_rows]
for (forward_map_index, result) in enumerate(validation_results)

    yᵢ_validation, ŷᵢ_validation, Γgp_validation = result
    r = round(Statistics.cor(yᵢ_validation, ŷᵢ_validation); sigdigits=2)
    @info "Pearson R for predictions on reserved subset of training points for $(k)th entry in the transformed forward map output : $r"
    ax = Axis(fig[ax_coords[i]...], xlabel = "True",
                                    xticks = LinearTicks(2),
                                    ylabel = "Predicted",
                                    title = "Index $i. Pearson R: $r")

    scatter!(ax, yᵢ_validation, ŷᵢ_validation)
    lines!(ax, yᵢ_validation, yᵢ_validation; color=(:black, 0.5), linewidth=3)
    errorbars!(yᵢ_validation, ŷᵢ_validation, sqrt.(Γgp_validation), color = :red, linewidth=2)
    save(joinpath(dir, "emulator_validation_performance_linear_linear.png"), fig)
end

# θ is a vector of parameter vectors
function nll_emulator(θ_vector)

    θ_array = hcat(θ_vector...)
    results = [predict(θ_array) for predict in predicts]
    μ_gps = hcat(getindex.(results, 1)...) # length(θ) x k
    Γ_gps = cat(getindex.(results, 2)...; dims=3) # length(θ) x length(θ) x k

    # @show size(μ_gps)
    # @show size(Γ_gps)

    inv_sqrt_Γθ = eki.precomputed_arrays[:inv_sqrt_Γθ]
    μθ = eki.precomputed_arrays[:μθ]
    Γŷ = 0 # assume that Γŷ is negligible compared to Γgp

    Φs = []
    for (i, θ) in enumerate(θ_vector)

        Ggp = μ_gps[i, :] # length-k vector
        Γgp = diagm(Γ_gps[i, i, :])

        θ_unconstrained = copy(θ)
        inverse_normalize!(θ_unconstrained, zscore_X)
        # Φ₁ = (1/2)*|| (Γgp + Γŷ)^(-½) * (ŷ - Ggp) ||²
        Φ₁ = (1/2) * norm(inv(sqrt(Γgp .+ Γŷ)) * (ŷ .- Ggp))^2
        # Φ₂ = (1/2)*|| Γθ^(-½) * (θ - μθ) ||²
        Φ₂ = eki.tikhonov ? (1/2) * norm(inv_sqrt_Γθ * (θ_unconstrained .- μθ))^2 : 0
        Φ₃ = (1/2) * log(det(Γgp .+ Γŷ))

        push!(Φs, Φ₁ + Φ₂ + Φ₃)
    end

    return Φs
end

C = Matrix(Hermitian(cov_θθ_all_iters))
@assert C ≈ cov_θθ_all_iters
dist_θθ_all_iters = MvNormal(zeros(size(X, 1)), C)
perturb() = rand(dist_θθ_all_iters)
proposal(θ) = θ + perturb()
seed_X = [perturb() for _ in 1:n_chains] # Where to initialize θ
# initial_ensemble = eki.iteration_summaries[0].parameters_unconstrained
# seed_X = [initial_ensemble[:,j] for j in 1:size(initial_ensemble,2)]

chain_X_emulated, chain_nll_emulated = markov_chain(nll_emulator, proposal, seed_X, chain_length_emulate; burn_in = burn_in_emulate, n_chains)
# inverse_normalize!.(chain_X_emulated, zscore_X)

samples = hcat(chain_X_emulated...)
inverse_normalize!(samples, zscore_X)
unscaled_chain_X_emulated = collect.(transform_to_constrained(eki.inverse_problem.free_parameters.priors, samples))
# unscaled_chain_X_emulated = [samples[:,j] for j in 1:size(samples, 2)]

###
### Sample from true eki objective using parallel chains of MCMC
###

# θ is a vector of parameter vectors
function nll_true(θ)

    # vector of vectors to 2d array
    θ_mx = hcat(θ...)
    inverse_normalize!(θ_mx, zscore_X)

    G = forward_map(training, θ)

    # # Vector of (Φ₁, Φ₂) pairs, one for each ensemble member at the current iteration
    # objective_values = []
    # error = 0
    # for j in 1:size(θ_mx, 2)
    #     try
    #         error = sum(eki_objective(eki, θ_mx[:, j], G[:, j]; constrained=true))
    #     catch DomainError
    #         error = Inf
    #     end
    #     push!(objective_values, error)
    # end
    objective_values = [sum(eki_objective(eki, θ_mx[:, j], G[:, j]; constrained=false)) for j in 1:size(θ_mx, 2)]

    # @show length(objective_values)
    # @show findall(x -> isfinite(x), objective_values)

    return objective_values
end

chain_X, chain_nll = markov_chain(nll_true, proposal, seed_X, chain_length; burn_in, n_chains)

samples = hcat(chain_X...)
inverse_normalize!(samples, zscore_X)
unscaled_chain_X = collect.(transform_to_constrained(eki.inverse_problem.free_parameters.priors, samples))
# unscaled_chain_X = [samples[:,j] for j in 1:size(samples, 2)]

using FileIO
file = joinpath(dir, "markov_chains.jld2")
save(file, Dict("unscaled_chain_X" => unscaled_chain_X,
                "unscaled_chain_X_emulated" => unscaled_chain_X_emulated))

# G_sobol = forward_map(big_training, params)
# save(file, G)
# G_sobol = load(file)["G"]

begin
    n_columns = 3
    hist_fig, hist_axes = plot_mcmc_densities2(unscaled_chain_X_emulated, parameter_set.names;
                                    n_columns,
                                    directory = dir,
                                    filename = "mcmc_densities_hist.png",
                                    label = "Emulated",
                                    color = (:blue, 0.8),
                                    type = "hist")

    # color = Makie.LinePattern(; direction = [Vec2f(1), Vec2f(1, -1)], width = 2, tilesize = (20, 20),
    #         linecolor = :blue, background_color = (:blue, 0.2))
    # color = Makie.LinePattern(; background_color = (:blue, 0.2))

    density_fig, density_axes = plot_mcmc_densities2(unscaled_chain_X_emulated, parameter_set.names;
                                    n_columns,
                                    directory = dir,
                                    filename = "mcmc_densities_density_textured.png",
                                    label = "Emulated",
                                    show_means = true,
                                    color = (:blue, 0.8),
                                    type = "density")

    plot_mcmc_densities2!(hist_fig, hist_axes, unscaled_chain_X, parameter_set.names;
                                    n_columns,
                                    directory = dir,
                                    filename = "mcmc_densities_hist.png",
                                    label = "True",
                                    color = (:orange, 0.5),
                                    type = "hist")

    plot_mcmc_densities2!(density_fig, density_axes, unscaled_chain_X, parameter_set.names;
                                    n_columns,
                                    directory = dir,
                                    filename = "mcmc_densities_density_textured.png",
                                    label = "True",
                                    show_means = true,
                                    color = (:orange, 0.5),
                                    strokecolor = :orange, strokewidth = 3, strokearound = true,
                                    type = "density")
end

begin
    xticks=(collect(1:Nparam), string.(collect(parameter_set.names)))
    xticklabelrotation = pi/2

    fig = Figure(resolution = (2000, 600), fontsize=28)
    Nparam = length(unscaled_chain_X[1])

    unscaled_chain_X_mx = hcat(unscaled_chain_X...)
    cor_true = Statistics.cor(unscaled_chain_X_mx, dims=2)

    unscaled_chain_X_emulated_mx = hcat(unscaled_chain_X_emulated...)
    cor_emulated = Statistics.cor(unscaled_chain_X_emulated_mx, dims=2)

    lb = minimum([minimum(cor_true), minimum(cor_emulated)])
    ub = maximum([maximum(cor_true), maximum(cor_emulated)])
    lb = -1.0
    ub = 1.0

    ax1 = Axis(fig[1, 1]; xticks, yticks=xticks, title="Emulated", xticklabelrotation)
    hmap1 = heatmap!(ax1, cor_emulated; colormap = :balance, colorrange=(lb, ub))
    # Colorbar(fig[1, 2], hmap1; label="Correlation")

    ax2 = Axis(fig[1, 2]; xticks, yticks=xticks, title="True", xticklabelrotation)
    hmap2 = heatmap!(ax2, cor_true; colormap = :balance, colorrange=(lb, ub))
    Colorbar(fig[1, 3], hmap2; label="Pearson Correlation")

    ax4 = Axis(fig[1, 4]; xticks, yticks=xticks, title="Difference (Emulated - True)", xticklabelrotation)
    hmap4 = heatmap!(ax4, cor_emulated .- cor_true; colormap = :balance, colorrange=(lb, ub))

    colsize!(fig.layout, 3, Relative(1/25))

    save(joinpath(dir, "correlation_heatmaps.png"), fig)
end

# colsize!(fig.layout, 1, Aspect(1, 1.0, 1, 1.0))
# colgap!(fig.layout, 7)
# display(fig)


function trival((twos, threes, fives))
    BigInt(2)^twos * BigInt(3)^threes * BigInt(5)^fives
end

mutable struct Hammings
    ham532 :: Vector{Tuple{Float64,Tuple{Int,Int,Int}}}
    ham53 :: Vector{Tuple{Float64,Tuple{Int,Int,Int}}}
    ndx532 :: Int
    ndx53 :: Int
    next2 :: Tuple{Float64,Tuple{Int,Int,Int}}
    next3 :: Tuple{Float64,Tuple{Int,Int,Int}}
    next5 :: Tuple{Float64,Tuple{Int,Int,Int}}
    next53 :: Tuple{Float64,Tuple{Int,Int,Int}}
    Hammings() = new(
        Vector{Tuple{Float64,Tuple{Int,Int,Int}}}(),
        Vector{Tuple{Float64,Tuple{Int,Int,Int}}}(),
        1, 1,
        (1.0, (1, 0, 0)), (log(2, 3), (0, 1, 0)),
        (log(2, 5), (0, 0, 1)), (0.0, (0, 0, 0))
    )
end
Base.eltype(::Type{Hammings}) = Tuple{Int,Int,Int}
function Base.iterate(HM::Hammings, st = HM) # :: Union{Nothing,Tuple{Tuple{Float64,Tuple{Int,Int,Int}},Hammings}}
    log2of2, log2of3, log2of5 = 1.0, log(2,3), log(2,5)
    if st.next2[1] < st.next53[1]
        push!(st.ham532, st.next2); st.ndx532 += 1
        last, (twos, threes, fives) = st.ham532[st.ndx532]
        st.next2 = (log2of2 + last, (twos + 1, threes, fives))
    else
        push!(st.ham532, st.next53)
        if st.next3[1] < st.next5[1]
            st.next53 = st.next3; push!(st.ham53, st.next3)
            last, (_, threes, fives) = st.ham53[st.ndx53]; st.ndx53 += 1
            st.next3 = (log2of3 + last, (0, threes + 1, fives))
        else
            st.next53 = st.next5; push!(st.ham53, st.next5)
            last, (_, _, fives) = st.next5
            st.next5 = (log2of5 + last, (0, 0, fives + 1))
        end
    end
    len53 = length(st.ham53)
    if st.ndx53 > (len53 >>> 1)
        nlen53 = len53 - st.ndx53 + 1
        copyto!(st.ham53, 1, st.ham53, st.ndx53, nlen53)
        resize!(st.ham53, nlen53); st.ndx53 = 1
    end
    len532 = length(st.ham532)
    if st.ndx532 > (len532 >>> 1)
        nlen532 = len532 - st.ndx532 + 1
        copyto!(st.ham532, 1, st.ham532, st.ndx532, nlen532)
        resize!(st.ham532, nlen532); st.ndx532 = 1
    end
    _, tri = st.ham532[end]
    tri, st
#    convert(Union{Nothing,Tuple{Tuple{Float64,Tuple{Int,Int,Int}},Hammings}},(st.ham532[end], st))
#   (length(st.ham532), length(st.ham53)), st
end

foreach(x -> print(trival(x)," "), (Iterators.take(Hammings(), 20))); println()
let count = 1691; for t in Hammings() count <= 1 && (println(trival(t)); break); count -= 1 end end
let count = 1000000; for t in Hammings() count <= 1 && (println(trival(t)); break); count -= 1 end end


using JuliaGSGP
using CSV, DataFrames
using ScikitLearn: CrossValidation.train_test_split

# df = CSV.read("path/to/dataset", DataFrame)
df = df[2:end]
data = convert(Matrix, df)
X = data[:, 1:end-1]
y = data[:, end]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)
X_train = X_train' # columns corresponds to one instances
X_test = X_test'

aq(a, b) = a / sqrt(1+b^2)

pset = PrimitiveSet()
add_function!(pset, +, 2) #(pset, function, arity)
add_function!(pset, -, 2)
add_function!(pset, *, 2)
add_function!(pset, aq, 2)

for i in 1:size(X)[2]
    add_variable!(pset, i)
end

cfg = Config(population_size=100, num_generations=100, p_crossover=0.5, p_mutation=0.5, mutation_step=0.01, min_depth=1, max_depth=6, tournament_size=4, pset=pset, create_rand_tree=generate_grow, create_population=ramped_half_and_half, fitness=rmse)

best_idx, rec = gsgp(cfg, X_train, y_train, X_test, y_test);

# `reconstruct_tree()` will take long time to execute.
# tree = reconstruct_tree(best_idx, cfg.mutation_step)


# Node that implements a controller for a DC motor.
# Calling syntax: nodectrl.jl [MQTT server address]
using openbuildnet

# Process input arguments
mqttserver = length(ARGS)>0?String(ARGS[1]):"tcp://localhost:1883"

# Controller's states
x = zeros(3)
command = 0.0

# The parameters
A = [-0.82 1.0 0.82; 1.0 0.0 0.0; 0.0 1.0 0.0]
C = [12.62 -19.75 7.625]

# Create the node
workspace = "test2"
node = OBNNode("ctrl", workspace, mqttserver)

# Create inputs
u_v = create_input(node, "v", Float64)  # Create velocity input
u_sp = create_input(node, "sp", Float64) # Create setpoint input

# Create output
y_u = create_output(node, "u", Float64) # Create command output

# Create the CSV writer
dumpfile = open("controller.txt", "w")

# Create Callbacks
on_init(node) do
    set(y_u, command) # Initialize the output (not necessary)
    println("At ", sim_time(node), " simulation started.")
end

on_term(node) do
    close(dumpfile)  # Close the CSV file
    println("At ", sim_time(node), " simulation terminated.")
end

const MAINBLOCK = 0

on_block_output(node, MAINBLOCK) do
    global command
    command = C * x
    set(y_u, command[1])
end

on_block_state(node, MAINBLOCK) do
    # Get the inputs
    sp = get(u_sp)
    v = get(u_v)
    global x
    x = A * x
    x[1] += 32.0 * (sp - v)

    # Write to the CSV file
    write(dumpfile, join((sim_time(node), sp, v), "\t"), "\t",
          join(command, "\t"), "\t", join(x, "\t"), "\n")
end

println("Ready to run the controller node; please start all other nodes ...")
status = runsafe(node, 60)
println("Simulation stopped with status = ", status)


using Luna
import Test: @test, @testset

radius = 125e-6 # HCF core radius
flength = 3 # HCF length

gas = :Ar
pressure = 80e-3 # gas pressure in bar

λ0 = 800e-9 # central wavelength of the pump pulse
τfwhm = 10e-15 # FWHM duration of the pump pulse
energy = 60e-6 # energy in the pump pulse

τ0 = Tools.τfw_to_τ0(τfwhm, :gauss)

out = prop_capillary(radius, flength, gas, pressure; λ0, τfwhm, energy,
                     modes=:HE11, trange=1000e-15, λlims=(150e-9, 4e-6), kerr=false, plasma=false)

t, Et = Processing.getEt(out)
Ptend = abs2.(Et[:, end]) # pulse profile at end

fwend = Maths.fwhm(t, Ptend; method=:spline) # FWHM duration at end
τ0end = Tools.τfw_to_τ0(fwend, :gauss) # Gaussian τ0 duration at end

mode = Capillary.MarcatiliMode(radius, gas, pressure) # propagation mode
β2 = Modes.dispersion(mode, 2, PhysData.wlfreq(λ0)) # dispersion of the mode
GDD = β2*flength # total GDD accumulated

τ0calc = sqrt((GDD/τ0)^2 + τ0^2) # analytical formula for Gaussian duration after GDD

@test isapprox(τ0calc, τ0end; rtol=1e-2)

energy_calc = Capillary.transmission(radius, λ0, flength)*energy # calculated energy at end
@test isapprox(energy_calc, Processing.energy(out)[end]; rtol=1e-2)


using Plots, LaTeXStrings
pyplot() # gr()
x = LinRange(0,2π,100) # with 100 the number of p
plot(x, sin.(x), c = :black)
savefig("figSin1py.pdf")


export AffineSubspaceHomotopy, set_subspaces!

"""
    AffineSubspaceHomotopy(F::System, V::AffineSubspace, W::AffineSubspace)
    AffineSubspaceHomotopy(F::AbstractSystem, V::AffineSubspace, W::AffineSubspace)

Creates a homotopy ``H(x,t) = (F ∩ γ(t))(x)`` where ``γ(t)`` is a family of affine subspaces
such that ``H(x,1) = (F ∩ V)(x)`` and ``H(x,0) = (F ∩ W)(x)``.
Here ``γ(t)`` is the geodesic between `V` and `W` in the affine Grassmanian, i.e.,
it is the curve of minimal length connecting `V` and `W`.
See also [`AffineSubspace`](@ref) and [`geodesic`](@ref) and the references therein.
"""
Base.@kwdef struct AffineSubspaceHomotopy{S<:AbstractSystem} <: AbstractHomotopy
    system::S

    start::AffineSubspace{ComplexF64}
    target::AffineSubspace{ComplexF64}

    J::Matrix{ComplexF64}
    Q::Matrix{ComplexF64}
    Q_cos::Matrix{ComplexF64}
    Θ::Vector{Float64}
    U::Matrix{ComplexF64}
    γ1::Matrix{ComplexF64}

    x::Vector{ComplexF64}
    ẋ::Vector{ComplexF64}
    x_high::Vector{ComplexDF64}
    t_cache::Base.RefValue{ComplexF64}
    # For AD
    taylor_t_cache::Base.RefValue{ComplexF64}
    taylor_γ::NTuple{4,Matrix{ComplexF64}}
    v::Vector{ComplexF64}
    tx³::TaylorVector{4,ComplexF64}
    tx²::TaylorVector{3,ComplexF64}
    tx¹::TaylorVector{2,ComplexF64}
end

## Implementation details

# The computation is performed using (implicit) Stiefel coordinates. If `size(F) = m, n`
# and `V` and `W` are affine subspaces of dimension `k, then the computation is performed in
# the affine Grassmanian Graff(k,n) embedded in the Grassmanian(k+1,n+1) using the system
# ``[F(γ(t)v); (γ(t)v)[end] - 1]``. Here the `(γ(t)v)[k+1] - 1` condition ensures that
# we are in the affine Grassmanian and that `(γ(t)v)[1:k]` is the correct value in \C^n.

AffineSubspaceHomotopy(F::ModelKit.System, start::AffineSubspace, target::AffineSubspace) =
    AffineSubspaceHomotopy(ModelKitSystem(F), start, target)

function AffineSubspaceHomotopy(
    system::AbstractSystem,
    start::AffineSubspace,
    target::AffineSubspace,
)

    Q, Θ, U = geodesic_svd(target, start)
    Q_cos = target.intrinsic.Y * U
    tx³ = TaylorVector{4}(ComplexF64, size(Q, 1))

    AffineSubspaceHomotopy(
        system = system,
        start = copy(start),
        target = copy(target),
        J = zeros(ComplexF64, size(system) .+ (1, 1)),
        Q = Q,
        Q_cos = Q_cos,
        Θ = Θ,
        U = U,
        γ1 = Q_cos * LA.diagm(0 => cos.(Θ)) + Q * LA.diagm(0 => sin.(Θ)),
        x = zeros(ComplexF64, size(Q, 1)),
        ẋ = zeros(ComplexF64, size(Q, 1)),
        x_high = zeros(ComplexDF64, size(Q, 1)),
        t_cache = Ref(complex(NaN, NaN)),
        taylor_t_cache = Ref(complex(NaN, NaN)),
        taylor_γ = tuple((similar(Q) for i = 0:3)...),
        v = zeros(ComplexF64, size(Q, 2)),
        tx³ = tx³,
        tx² = TaylorVector{3}(tx³),
        tx¹ = TaylorVector{2}(tx³),
    )
end
Base.size(H::AffineSubspaceHomotopy) = (size(H.system)[1] + 1, dim(H.start) + 1)

"""
    set_subspaces!(H::AffineSubspaceHomotopy, start::AffineSubspace, target::AffineSubspace)

Update the homotopy `H` to track from the affine subspace `start` to `target`.
"""
function set_subspaces!(
    H::AffineSubspaceHomotopy,
    start::AffineSubspace,
    target::AffineSubspace,
)
    Q, Θ, U = geodesic_svd(target, start)
    copy!(H.start, start)
    copy!(H.target, target)
    H.Q .= Q
    H.Θ .= Θ
    H.U .= U
    LA.mul!(H.Q_cos, target.intrinsic.Y, U)
    LA.mul!(H.γ1, Q, LA.diagm(0 => sin.(Θ)))
    LA.mul!(H.γ1, H.Q_cos, LA.diagm(0 => cos.(Θ)), true, true)
    H
end

function γ!(H::AffineSubspaceHomotopy, t::Number)
    H.t_cache[] != t || return first(H.taylor_γ)
    if isreal(t)
        _γ!(H, real(t))
    else
        _γ!(H, t)
    end
    H.t_cache[] = t

    first(H.taylor_γ)
end
@inline function _γ!(H::AffineSubspaceHomotopy, t::Number)
    @unpack Q, Q_cos, Θ = H
    γ = first(H.taylor_γ)
    n, k = size(γ)
    @inbounds for j = 1:k
        Θⱼ = Θ[j]
        s, c = sincos(t * Θⱼ)
        for i = 1:n
            γ[i, j] = Q_cos[i, j] * c + Q[i, j] * s
        end
    end
    γ
end

γ̇!(H::AffineSubspaceHomotopy, t::Number) = isreal(t) ? _γ̇!(H, real(t)) : _γ̇!(H, t)
@inline function _γ̇!(H::AffineSubspaceHomotopy, t::Number)
    @unpack Q, Q_cos, Θ = H
    _, γ̇ = H.taylor_γ
    n, k = size(γ̇)
    @inbounds for j = 1:k
        Θⱼ = Θ[j]
        s, c = sincos(t * Θⱼ)
        ċ = -s * Θⱼ
        ṡ = c * Θⱼ
        for i = 1:n
            γ̇[i, j] = Q_cos[i, j] * ċ + Q[i, j] * ṡ
        end
    end
    γ̇
end

function set_solution!(u::Vector, H::AffineSubspaceHomotopy, x::AbstractVector, t)
    (length(x) == length(H.x) - 1) ||
        throw(ArgumentError("Cannot set solution. Expected extrinsic coordinates."))
    for i = 1:length(x)
        H.x[i] = x[i]
    end
    H.x[end] = 1

    if isone(t)
        LA.mul!(u, H.γ1', H.x)
    elseif iszero(t)
        LA.mul!(u, H.Q_cos', H.x)
    else
        LA.mul!(u, γ!(H, t), H.x)
    end
end

function get_solution(H::AffineSubspaceHomotopy, u::AbstractVector, t)
    if isone(t)
        (@view H.γ1[1:end-1, :]) * u
    elseif iszero(t)
        (@view H.Q_cos[1:end-1, :]) * u
    else
        γ = γ!(H, t)
        (@view γ[1:end-1, :]) * u
    end
end

function evaluate!(u, H::AffineSubspaceHomotopy, v::AbstractVector, t)
    γ = γ!(H, t)
    n = first(size(H.system))
    if eltype(v) isa ComplexDF64
        LA.mul!(H.x_high, γ, v)
        evaluate!(u, H.system, H.x_high)
        u[n+1] = H.x_high[end] - 1.0
    else
        LA.mul!(H.x, γ, v)
        evaluate!(u, H.system, H.x)
        u[n+1] = H.x[end] - 1.0
    end
    u
end

function evaluate_and_jacobian!(u, U, H::AffineSubspaceHomotopy, v::AbstractVector, t)
    γ = γ!(H, t)

    LA.mul!(H.x, γ, v)
    evaluate_and_jacobian!(u, H.J, H.system, H.x)
    LA.mul!(U, H.J, γ)

    n = first(size(H.system))
    u[n+1] = H.x[end] - 1

    m = length(v)
    for j = 1:m
        U[n+1, j] = γ[end, j]
    end

    nothing
end

function taylor!(u, ::Val{1}, H::AffineSubspaceHomotopy, v, t)
    γ = γ!(H, t)
    γ̇ = γ̇!(H, t)

    # apply chain rule
    #    d/dt [F(γ(t)v); (γ(t)v)[end] - 1] = [J_F(γ(t)v)* γ̇(t)*v;  (γ̇(t)v)[end]]
    LA.mul!(H.x, γ, v)
    LA.mul!(H.ẋ, γ̇, v)
    evaluate_and_jacobian!(u, H.J, H.system, H.x)
    LA.mul!(u, H.J, H.ẋ)
    M = size(H, 1)
    u[M] = H.ẋ[end]

    u
end

function _taylor_γ!(H::AffineSubspaceHomotopy, t::Number)
    @unpack Q, Q_cos, Θ, U, taylor_γ = H

    γ, γ¹, γ², γ³ = taylor_γ
    n, k = size(γ)
    @inbounds for j = 1:k
        Θⱼ = Θ[j]
        s, c = sincos(t * Θⱼ)
        c¹ = -s * Θⱼ
        s¹ = c * Θⱼ
        Θⱼ_2 = 0.5 * Θⱼ^2
        s² = -s * Θⱼ_2
        c² = -c * Θⱼ_2
        Θⱼ_3 = Θⱼ_2 * Θⱼ / 3
        s³ = -c * Θⱼ_3
        c³ = s * Θⱼ_3
        for i = 1:n
            γ[i, j] = Q_cos[i, j] * c + Q[i, j] * s
            γ¹[i, j] = Q_cos[i, j] * c¹ + Q[i, j] * s¹
            γ²[i, j] = Q_cos[i, j] * c² + Q[i, j] * s²
            γ³[i, j] = Q_cos[i, j] * c³ + Q[i, j] * s³
        end
    end

end

function taylor_γ!(H::AffineSubspaceHomotopy, t::Number)
    H.taylor_t_cache[] != t || return H.taylor_γ

    if isreal(t)
        _taylor_γ!(H, real(t))
    else
        _taylor_γ!(H, t)
    end
    H.taylor_t_cache[] = t

    H.taylor_γ
end

function taylor!(u, v::Val{2}, H::AffineSubspaceHomotopy, tv::TaylorVector, t, incr::Bool)
    γ, γ¹, γ², γ³ = taylor_γ!(H, t)
    x, x¹, x² = vectors(H.tx²)
    v, v¹ = vectors(tv)

    if !incr
        LA.mul!(x, γ, v)
        LA.mul!(x¹, γ¹, v)
    end

    H.v .= v¹
    LA.mul!(H.x, γ, v¹)
    x¹ .+= H.x

    LA.mul!(H.x, γ¹, H.v)
    H.v .= v
    LA.mul!(H.x, γ², H.v, true, true)
    x² .= H.x

    taylor!(u, Val(2), H.system, H.tx²)

    n = first(size(H.system))
    u[n+1] = x²[end]

    u
end

function taylor!(u, v::Val{3}, H::AffineSubspaceHomotopy, tv::TaylorVector, t, incr::Bool)
    γ, γ¹, γ², γ³ = taylor_γ!(H, t)
    x, x¹, x², x³ = vectors(H.tx³)
    v, v¹, v² = vectors(tv)

    if !incr
        LA.mul!(x, γ, v)
        LA.mul!(x¹, γ¹, v)
        LA.mul!(x¹, γ, v¹, true, true)
        LA.mul!(x², γ², v)
        LA.mul!(x², γ¹, v¹, true, true)
    end

    H.v .= v²
    LA.mul!(H.x, γ, H.v)
    x² .+= H.x

    LA.mul!(H.x, γ¹, H.v)
    H.v .= v
    LA.mul!(H.x, γ³, H.v, true, true)
    H.v .= v¹
    LA.mul!(H.x, γ², H.v, true, true)
    x³ .= H.x

    taylor!(u, Val(3), H.system, H.tx³)

    n = first(size(H.system))
    u[n+1] = x³[end]

    u
end