function SpecialFunctions.besselj(ν::Arblib.ArbOrRef, z::Arblib.ArbSeries)
    deg = Arblib.degree(z)
    res = zero(z)

    x = z[0]
    tmp = zero(x)
    ai = zero(x)

    # ai = besselj(ν, x)
    Arblib.hypgeom_bessel_j!(ai, ν, x)

    res[0] = ai

    if deg >= 1
        # ai = (besselj(ν - 1, x) - besselj(ν + 1, x)) / 2

        # ai = besselj(ν - 1, x)
        Arblib.sub!(tmp, ν, 1)
        Arblib.hypgeom_bessel_j!(ai, tmp, x)

        # tmp = besselj(ν + 1, x)
        # ai += tmp
        Arblib.add!(tmp, ν, 1)
        Arblib.hypgeom_bessel_j!(tmp, tmp, x)
        Arblib.sub!(ai, ai, tmp)

        # ai /= 2
        Arblib.mul_2exp!(ai, ai, -1)

        res[1] = ai

        if deg >= 2
            # ai = (besselj(ν - 2, x) + besselj(ν + 2, x) - 2a0) / 8

            # ai = besselj(ν - 2, x)
            Arblib.sub!(tmp, ν, 2)
            Arblib.hypgeom_bessel_j!(ai, tmp, x)

            # tmp = besselj(ν + 2, x)
            # ai += tmp
            Arblib.add!(tmp, ν, 2)
            Arblib.hypgeom_bessel_j!(tmp, tmp, x)
            Arblib.add!(ai, ai, tmp)

            # ai -= 2a0
            Arblib.mul_2exp!(tmp, Arblib.ref(res, 0), 1)
            Arblib.sub!(ai, ai, tmp)

            # ai /= 8
            Arblib.mul_2exp!(ai, ai, -3)

            res[2] = ai

            if deg >= 3
                # ai = (besselj(ν - 3, x) - besselj(ν + 3, x) - 6a1) / 48

                # ai = besselj(ν - 3, x)
                Arblib.sub!(tmp, ν, 3)
                Arblib.hypgeom_bessel_j!(ai, tmp, x)

                # tmp = besselj(ν + 3, x)
                # ai += tmp
                Arblib.add!(tmp, ν, 3)
                Arblib.hypgeom_bessel_j!(tmp, tmp, x)
                Arblib.sub!(ai, ai, tmp)

                # ai -= 6a1
                Arblib.mul!(tmp, Arblib.ref(res, 1), 6)
                Arblib.sub!(ai, ai, tmp)

                # ai /= 48
                Arblib.div!(ai, ai, 48)

                res[3] = ai

                if deg >= 4
                    x2 = x^2
                    x2mν2 = x2 - ν^2
                    for i = 4:deg
                        k = i - 4

                        # tmp = 2x * res[k + 1]
                        Arblib.mul_2exp!(tmp, x, 1)
                        Arblib.mul!(tmp, tmp, Arblib.ref(res, k + 1))

                        # ai = res[k] + tmp
                        Arblib.add!(ai, Arblib.ref(res, k), tmp)

                        # tmp = (k^2 + 4k + 4 - ν^2 + x^2)
                        # ai += tmp * res[k + 2]
                        Arblib.add!(tmp, x2mν2, k^2 + 4k + 4)
                        Arblib.addmul!(ai, tmp, Arblib.ref(res, k + 2))

                        # tmp = (2k^2 + 11k + 15) * x
                        # ai += tmp * res[k + 3]
                        Arblib.mul!(tmp, x, 2k^2 + 11k + 15)
                        Arblib.addmul!(ai, tmp, Arblib.ref(res, k + 3))

                        # tmp = x^2 * (k^2 + 7k + 12)
                        # ai /= tmp
                        Arblib.mul!(tmp, x2, k^2 + 7k + 12)
                        Arblib.div!(ai, ai, tmp)

                        # ai = -ai
                        Arblib.neg!(ai, ai)

                        res[i] = ai
                    end
                end
            end
        end
    end

    # Compose the Taylor series for the Bessel function with that of z

    # First set the constant in z to 0 temporarily
    Arblib.zero!(Arblib.ref(z, 0))
    Arblib.compose_series!(res, res, z, deg + 1)
    Arblib.set!(Arblib.ref(z, 0), x)

    return res
end

SpecialFunctions.besselj0(z::ArbSeries) = besselj(zero(Arb), z)
SpecialFunctions.besselj1(z::ArbSeries) = besselj(one(Arb), z)


#=
# 3D Poisson, Dirichlet bc
# Uses Dendro imported as a custom gen target
=#
if !@isdefined(Femshop)
    include("../Femshop.jl");
    using .Femshop
end
init_femshop("dendrotest");

useLog("dendrotestlog", level=3)

# default values (max_depth=6, wavelet_tol = 0.1, partition_tol = 0.3, solve_tol = 1e-6, max_iters = 100)
generateFor("target_dendro_cg.jl", params=(8, 0.001, 0.3, 0.000001, 500))

domain(3)
functionSpace(order=2)

u = variable("u")
testSymbol("v")

boundary(u, 1, DIRICHLET, "0")

# Write the weak form
coefficient("k1", "1.1")
coefficient("k2", "0.1")
coefficient("k3", "0.1")
coefficient("k4", "0.1")
coefficient("k5", "0.1")
f = coefficient("f", "-1.1*14*pi*pi*sin(3*pi*x)*sin(2*pi*y)*sin(pi*z) + 0.1*3*pi*cos(3*pi*x)*sin(2*pi*y)*sin(pi*z) + 0.1*2*pi*sin(3*pi*x)*cos(2*pi*y)*sin(pi*z) + 0.1*pi*sin(3*pi*x)*sin(2*pi*y)*cos(pi*z) - 0.1*sin(3*pi*x)*sin(2*pi*y)*sin(pi*z)")
weakForm(u, "-k1*dot(grad(u),grad(v)) + k2*deriv(u,1)*v + k3*deriv(u,2)*v + k4*deriv(u,3)*v - k5*u*v - f*v")

build_octree_with(f);
solve(u);

finalize_femshop()


function formulation(
    i::BipartiteMatchingInstance{Int},
    ::DefaultLinearFormulation;
    solver=nothing,
)
    indices = collect((src(e), dst(e)) for e in edges(i.graph))

    model = Model(solver)
    x = @variable(solver.model, [indices], binary = true)

    for i in 1:n_arms # Left nodes.
        @constraint(solver.model, sum(solver.x[(i, j)] for j in 1:n_arms) == 1)
    end
    for j in 1:n_arms # Right nodes.
        @constraint(solver.model, sum(solver.x[(i, j)] for i in 1:n_arms) == 1)
    end

    @objective(
        model,
        Max,
        sum(rewards[(i, j)] * solver.x[(i, j)] for (i, j) in keys(rewards))
    )
    @constraint(model, sum(x) <= i.m)

    set_silent(model)

    return m, x
end


# This file is part of Kpax3. License is MIT.

function logmarglik(y::Vector{Float64},
                    n::Real,
                    α::Vector{Float64},
                    β::Vector{Float64})
  SpecialFunctions.logbeta.(α .+ y, β .- y .+ n) .-
  SpecialFunctions.logbeta.(α, β)
end

function logmarglik(y::Real,
                    n::Real,
                    α::Real,
                    β::Real)
  SpecialFunctions.logbeta(α + y, β + n - y) - SpecialFunctions.logbeta(α, β)
end

function logcondmarglik(x::Vector{UInt8},
                        y::Vector{Float64},
                        n::Real,
                        α::Vector{Float64},
                        β::Vector{Float64})
  float(x) .* log.(α .+ y) .+
  (one(Float64) .- float(x)) .* log.(β .- y .+ n) .-
  log.(α .+ β .+ n)
end

function logcondmarglik(x::Real,
                        y::Real,
                        n::Real,
                        α::Real,
                        β::Real)
  (x * log(α + y) + (one(x) - x) * log(β + n - y) - log(α + β + n))::Float64
end

function marglik(y::Vector{Float64},
                 n::Real,
                 α::Vector{Float64},
                 β::Vector{Float64})
  exp.(logmarglik(y, n, α, β))
end

function marglik(y::Real,
                 n::Real,
                 α::Real,
                 β::Real)
  exp(logmarglik(y, n, α, β))
end

function condmarglik(x::Vector{UInt8},
                     y::Vector{Float64},
                     n::Real,
                     α::Vector{Float64},
                     β::Vector{Float64})
  exp.(logcondmarglik(x, y, n, α, β))
end

function condmarglik(x::Real,
                     y::Real,
                     n::Real,
                     α::Real,
                     β::Real)
  exp(logcondmarglik(x, y, n, α, β))
end


function autograd{T <: Number}(
    f::Function,
    x::T,
)
    return epsilon(f(Dual(x, one(T))))
end

function autograd!{T <: Real}(
    gr::Vector{T},
    x_d::Vector{Dual{T}},
    f::Function,
    x::Vector{T},
)
    n = length(x)
    for i in 1:n
        x_d[i] = Dual(x[i], zero(T))
    end
    for i in 1:n
        # Why isn't this x[i]
        x_d[i] = Dual(real(x_d[i]), one(T))
        result = f(x_d)
        gr[i] = epsilon(result)
        # Why isn't this x[i]
        x_d[i] = Dual(real(x_d[i]), zero(T))
    end
    return gr
end


module TrueSkill

global const BETA = 1.0::Float64
global const MU = 0.0::Float64
global const SIGMA = (BETA * 6)::Float64
global const GAMMA = (BETA * 0.03)::Float64
global const P_DRAW = 0.0::Float64
global const EPSILON = 1e-6::Float64
global const ITERATIONS = 30::Int64
global const sqrt2 = sqrt(2)
global const sqrt2pi = sqrt(2*pi)
global const PI = 1/(SIGMA^2)
global const TAU = MU*PI
global const minInt64 = (-9223372036854775808)::Int64
global const maxInt64 = ( 9223372036854775807)::Int64

function erfc(x::Float64)
    """Complementary error function (thanks to http://bit.ly/zOLqbc)"""
    z = abs(x)
    t = 1.0 / (1.0 + z / 2.0)
    r = begin
        a = -0.82215223 + t * 0.17087277
        b =  1.48851587 + t * a
        c = -1.13520398 + t * b
        d =  0.27886807 + t * c
        e = -0.18628806 + t * d
        f =  0.09678418 + t * e
        g =  0.37409196 + t * f
        h =  1.00002368 + t * g
        t * exp(-z * z - 1.26551223 + t * h)
        end
    if x < 0
        r = 2.0 - r
    end
    return r
end
function erfcinv(y::Float64)
    """The inverse function of erfc."""
    if y >= 2
        return -Inf
    elseif y < 0
        throw(DomainError(y, "argument must be nonnegative"))
    elseif y == 0
        return Inf
    end
    zero_point = y < 1
    if ! zero_point
        y = 2 - y
    end
    t = sqrt(-2 * log(y / 2.0))
    x = -0.70711 * ((2.30753 + t * 0.27061) / (1.0 + t * (0.99229 + t * 0.04481)) - t)
    for _ in 0:2
        err = erfc(x) - y
        x += err / (1.12837916709551257 * exp(-(x^2)) - x * err)
    end
    if zero_point
        r = x
    else
        r = -x
    end
    return r
end
function tau_pi(mu::Float64, sigma::Float64)
    if sigma > 0.
        _pi = sigma^-2
        _tau = _pi * mu
    elseif (sigma + 1e-5) < 0.0
        error("sigma should be greater than 0")
    else
        _pi = Inf
        _tau = Inf
    end
    return _tau, _pi
end
function mu_sigma(_tau::Float64, _pi::Float64)
    if _pi > 0.0
        sigma = sqrt(1/_pi)
        mu = _tau / _pi
    elseif (_pi + 1e-5) < 0.
        error("Precision should be greater than 0")
    else
        sigma = Inf
        mu = 0.0
    end
    return mu, sigma
end

struct Gaussian
    mu::Float64
    sigma::Float64
    function Gaussian(mu::Float64=MU, sigma::Float64=SIGMA)
        if sigma>=0.0
            return new(mu, sigma)
        else
            throw(error("Require: (sigma >= 0.0)"))
        end
    end
    function Gaussian(;mu::Float64=MU, sigma::Float64=SIGMA)
        return new(mu, sigma)
    end
end

global const N01 = Gaussian(0.0, 1.0)
global const Ninf = Gaussian(0.0, Inf)
global const N00 = Gaussian(0.0, 0.0)

Base.show(io::IO, g::Gaussian) = print("Gaussian(mu=", round(g.mu,digits=3), ", sigma=", round(g.sigma,digits=3), ")")
function _pi_(N::Gaussian)
    if N.sigma > 0.
        return N.sigma^-2
    else
        return Inf
    end
end
function _tau_(N::Gaussian)
    if N.sigma > 0.
        return N.mu * (N.sigma^-2)
    else
        return Inf
    end
end
function cdf(N::Gaussian, x::Float64)
    z = -(x - N.mu) / (N.sigma * sqrt2)
    return (0.5 * erfc(z))::Float64
end
function pdf(N::Gaussian, x::Float64)
    normalizer = (sqrt(2*pi) * N.sigma)^-1
    functional = exp( -((x - N.mu)^2) / (2*N.sigma ^2) )
    return (normalizer * functional)::Float64
end
function ppf(N::Gaussian, p::Float64)
    return N.mu - N.sigma * sqrt2  * erfcinv(2 * p)
end
function approx(;N::Gaussian, margin::Float64, tie::Bool)
    approx(N, margin, tie)
end
function approx(N::Gaussian, margin::Float64, tie::Bool)
    #The range is [alpha, beta]
    if !tie
        _alpha = (margin-N.mu)/N.sigma
        v = pdf(N01,-_alpha) / cdf(N01,-_alpha)
        w = v * (v + (-_alpha))
    else
        _alpha = (-margin-N.mu)/N.sigma
        _beta  = (margin-N.mu)/N.sigma
        v = (pdf(N01,_alpha)-pdf(N01,_beta))/(cdf(N01,_beta)-cdf(N01,_alpha))
        u = (_alpha*pdf(N01,_alpha)-_beta*pdf(N01,_beta))/(cdf(N01,_beta)-cdf(N01,_alpha))
        w =  - ( u - v^2 )
    end
    mu = N.mu + N.sigma * v
    sigma = N.sigma*sqrt(1-w)
    return Gaussian(mu, sigma)
end
function delta(N::Gaussian, M::Gaussian)
    return abs(N.mu - M.mu) , abs(N.sigma - M.sigma)
end
function exclude(N::Gaussian,M::Gaussian)
    return Gaussian(N.mu - M.mu, sqrt(N.sigma^2 - M.sigma^2) )
end
function Base.:+(N::Gaussian, M::Gaussian)
    mu = N.mu + M.mu
    sigma = sqrt(N.sigma^2 + M.sigma^2)
    return Gaussian(mu, sigma)
end
function Base.:-(N::Gaussian, M::Gaussian)
    mu = N.mu - M.mu
    sigma = sqrt(N.sigma^2 + M.sigma^2)
    return Gaussian(mu, sigma)
end
function Base.:*(N::Gaussian, M::Gaussian)
    _pi = _pi_(N) + _pi_(M)
    _tau = _tau_(N) + _tau_(M)
    mu, sigma = mu_sigma(_tau, _pi)
    return Gaussian(mu, sigma)
end
function Base.:/(N::Gaussian, M::Gaussian)
    _pi = _pi_(N) - _pi_(M)
    _tau = _tau_(N) - _tau_(M)
    mu, sigma = mu_sigma(_tau, _pi)
    return Gaussian(mu, sigma)
end
function Base.isapprox(N::Gaussian, M::Gaussian, atol::Real=0)
    return (abs(N.mu - M.mu) < atol) & (abs(N.sigma - M.sigma) < atol)
end
function forget(N::Gaussian, gamma::Float64, t::Int64)
    return Gaussian(N.mu, sqrt(N.sigma^2 + t*gamma^2))
end
function compute_margin(p_draw::Float64, sd::Float64)
    _N = Gaussian(0.0, sd )
    res = abs(ppf(_N, 0.5-p_draw/2))
    return res
end
struct Player
    prior::Gaussian
    beta::Float64
    gamma::Float64
    draw::Gaussian
    function Player(prior::Gaussian=Gaussian(MU,SIGMA), beta::Float64=BETA, gamma::Float64=GAMMA, draw::Gaussian=Ninf)
        new(prior, beta, gamma, draw)
    end
    function Player(;prior::Gaussian=Gaussian(MU,SIGMA), beta::Float64=BETA, gamma::Float64=GAMMA, draw::Gaussian=Ninf)
        Player(prior, beta, gamma, draw)
    end
end
Base.show(io::IO, r::Player) = print("Player(Gaussian(mu=", round(r.prior.mu,digits=3),", sigma=", round(r.prior.sigma,digits=3), "), beta=", round(r.beta, digits=3), ", gamma=" , round(r.gamma, digits=3),")")
function performance(R::Player)
    return Gaussian(R.prior.mu, sqrt(R.prior.sigma^2 + R.beta^2))
end

function likelihood(R::Player, likelihood_performance::Gaussian)
    return Gaussian(likelihood_performance.mu, sqrt(likelihood_performance.sigma^2 + R.beta^2))
end

mutable struct team_messages
    prior::Gaussian
    likelihood_lose::Gaussian
    likelihood_win::Gaussian
    likelihood_draw::Gaussian
end
function p(tm::team_messages)
    return tm.prior*tm.likelihood_lose*tm.likelihood_win*tm.likelihood_draw
end
function posterior_win(tm::team_messages)
    return tm.prior*tm.likelihood_lose*tm.likelihood_draw
end
function posterior_lose(tm::team_messages)
    return tm.prior*tm.likelihood_win*tm.likelihood_draw
end
function posterior_draw(tm::team_messages)
    return tm.prior*tm.likelihood_win*tm.likelihood_lose
end
function likelihood(tm::team_messages)
    return tm.likelihood_win*tm.likelihood_lose*tm.likelihood_draw
end
mutable struct draw_messages
    prior::Gaussian
    prior_team::Gaussian
    likelihood_lose::Gaussian
    likelihood_win::Gaussian
end
function p(um::draw_messages)
    return um.prior_team*um.likelihood_lose*um.likelihood_win
end
function posterior_win(um::draw_messages)
    return um.prior_team*um.likelihood_lose
end
function posterior_lose(um::draw_messages)
    return um.prior_team*um.likelihood_win
end
function likelihood(um::draw_messages)
    return um.likelihood_win*um.likelihood_lose
end
mutable struct diff_messages
    prior::Gaussian
    likelihood::Gaussian
end
function p(dm::diff_messages)
    return dm.prior*dm.likelihood
end

function Base.max(tuple1::Tuple{Float64,Float64}, tuple2::Tuple{Float64,Float64})
    return max(tuple1[1],tuple2[1]), max(tuple1[2],tuple2[2])
end
function Base.:>(tuple::Tuple{Float64,Float64}, threshold::Float64)
    return (tuple[1] > threshold) | (tuple[2] > threshold)
end

mutable struct Game
    teams::Vector{Vector{Player}}
    result::Vector{Float64}
    p_draw::Float64
    likelihoods::Vector{Vector{Gaussian}}
    evidence::Float64
    function Game(teams::Vector{Vector{Player}}, result::Vector{Float64}=Float64[],p_draw::Float64=0.0)
        ((0.0 > p_draw) | (1.0 <= p_draw)) &&  throw(error("0.0 <= Draw probability < 1.0"))
        (length(result)>0) && (length(teams)!= length(result)) && throw(error("(length(result)>0) & (length(teams)!= length(result))"))
        (p_draw == 0.0) && (length(result)>0) && (length(unique(result)) != length(result)) && throw(error("(p_draw == 0.0) && (length(result)>0) && (length(unique(result)) != length(result))"))

        _g = new(teams,result,p_draw,[],0.0)
        likelihoods(_g)
        return _g
    end
    function Game(teams::Vector{Vector{Player}}, result::Vector{Float64}=Float64[];p_draw::Float64=0.0)
        Game(teams, result, p_draw)
    end
    function Game(teams::Vector{Vector{Player}}; result::Vector{Float64}=Float64[],p_draw::Float64=0.0)
        Game(teams, result, p_draw)
    end
end
Base.length(G::Game) = length(G.teams)
function size(G::Game)
    return [length(team) for team in g.teams]
end
function performance(team::Vector{Player})
    p = N00
    for r in team
        p += performance(r)
    end
    return p
end
#TODO: reducir perf(game)
function performance(G::Game,i::Int64)
    res = N00
    for r in G.teams[i]
        res += performance(r)
    end
    return res
end
function draw_performance(G::Game,i::Int64)
    res = N00
    for r in G.teams[i]
        res += r.draw.sigma < Inf ? approx(r.draw,0.,false) : Ninf
    end
    return res
end
function likelihood_teams(g::Game)
    r = g.result == Float64[] ? [Float64(i) for i in length(g.teams):-1:1] : g.result
    o = sortperm(r, rev=true)
    t = [team_messages(performance(g,o[e]), Ninf, Ninf, Ninf) for e in 1:length(g)]
    d = [diff_messages(t[e].prior - t[e+1].prior, Ninf) for e in 1:length(g)-1]
    tie = [r[o[e]]==r[o[e+1]] for e in 1:length(d)]
    margin = [ g.p_draw==0.0 ?  0.0 :
               compute_margin(g.p_draw, sqrt( sum([a.beta^2 for a in g.teams[o[e]]]) + sum([a.beta^2 for a in g.teams[o[e+1]]]) ))
               for e in 1:length(d)]
    g.evidence = 1
    for e in 1:length(d)
        g.evidence *= !tie[e] ? 1-cdf(d[e].prior, margin[e]) : cdf(d[e].prior, margin[e])-cdf(d[e].prior, -margin[e])
    end
    step = (Inf, Inf)::Tuple{Float64,Float64}; iter = 0::Int64
    while (step > 1e-6) & (iter < 10)
        step = (0., 0.)
        for e in 1:length(d)-1#e=1
            d[e].prior = posterior_win(t[e]) - posterior_lose(t[e+1])
            d[e].likelihood = approx(d[e].prior,margin[e],tie[e])/d[e].prior
            likelihood_lose = posterior_win(t[e]) - d[e].likelihood
            step = max(step,delta(t[e+1].likelihood_lose,likelihood_lose))
            t[e+1].likelihood_lose = likelihood_lose
        end
        for e in length(d):-1:2
            d[e].prior = posterior_win(t[e]) - posterior_lose(t[e+1])
            d[e].likelihood = approx(d[e].prior,margin[e],tie[e])/d[e].prior
            likelihood_win = (posterior_lose(t[e+1]) + d[e].likelihood)
            step = max(step,delta(t[e].likelihood_win,likelihood_win))
            t[e].likelihood_win = likelihood_win
        end
        iter += 1
    end
    if length(d)==1
        d[1].prior = posterior_win(t[1]) - posterior_lose(t[2])
        d[1].likelihood = approx(d[1].prior,margin[1],tie[1])/d[1].prior
    end
    t[1].likelihood_win = (posterior_lose(t[2]) + d[1].likelihood)
    t[end].likelihood_lose = (posterior_win(t[end-1]) - d[end].likelihood)

    return [ likelihood(t[o[e]]) for e in 1:length(t)]
end
function likelihoods(g::Game)
    m_t_ft = likelihood_teams(g)
    g.likelihoods = [[
        likelihood(g.teams[e][i],
                   m_t_ft[e] - exclude(performance(g,e),performance(g.teams[e][i])))
        for i in 1:length(g.teams[e])] for e in 1:length(g)]
    return g.likelihoods
end
function posteriors(g::Game)
    return [[ g.likelihoods[e][i] * g.teams[e][i].prior for i in 1:length(g.teams[e])] for e in 1:length(g)]
end
mutable struct Skill
    forward::Gaussian
    backward::Gaussian
    likelihood::Gaussian
    elapsed::Int64
    online::Gaussian
    function Skill(forward::Gaussian=Ninf, backward::Gaussian=Ninf, likelihood::Gaussian=Ninf, elapsed::Int64=0)
        return new(forward, backward, likelihood, elapsed)
    end
    function Skill(;forward::Gaussian=Ninf, backward::Gaussian=Ninf, likelihood::Gaussian=Ninf, elapsed::Int64=0)
        return new(forward, backward, likelihood, elapsed)
    end
end
mutable struct Agent
    player::Player
    message::Gaussian
    last_time::Int64
end
function receive(agent::Agent, elapsed::Int64)
    if agent.message != Ninf
        res = forget(agent.message, agent.player.gamma, elapsed)
    else
       res = agent.player.prior
    end
    return res
end
function clean(agents::Dict{String,Agent}, last_time::Bool=false)
    for a in keys(agents)
        agents[a].message = Ninf
        if last_time agents[a].last_time = minInt64 end
    end
end
mutable struct Item
    agent::String
    likelihood::Gaussian
end
mutable struct Team
    items::Vector{Item}
    output::Float64
end
mutable struct Event
    teams::Vector{Team}
    evidence::Float64
end
function outputs(event::Event)
    return [ team.output for team in event.teams]
end
function get_composition(events::Vector{Event})
    return [[[i.agent for i in t.items] for t in e.teams] for e in events]
end
function get_results(events::Vector{Event})
    return [ [t.output for t in e.teams ] for e in events]
end
function compute_elapsed(last_time::Int64, actual_time::Int64)
    return last_time == minInt64 ? 0 : ( last_time == maxInt64 ? 1 : (actual_time - last_time))
end
mutable struct Batch
    time::Int64
    events::Vector{Event}
    skills::Dict{String,Skill}
    agents::Dict{String,Agent}
    p_draw::Float64
    function Batch(composition::Vector{Vector{Vector{String}}}, results::Vector{Vector{Float64}}=Vector{Vector{Float64}}() ,time::Int64=0, agents::Dict{String,Agent}=Dict{String,Agent}(), p_draw::Float64 = 0.0)
        (length(results)>0) & (length(composition)!= length(results)) && throw(error("(length(results)>0) & (length(composition)!= length(results))"))

        this_agents = Set(vcat((composition...)...))
        elapsed = Dict([ (a, compute_elapsed(agents[a].last_time, time) ) for a in this_agents  ])
        skills = Dict([ (a, Skill(receive(agents[a],elapsed[a]) ,Ninf ,Ninf , elapsed[a])) for a in this_agents  ])
        events = [Event([Team([Item(composition[e][t][a], Ninf) for a in 1:length(composition[e][t]) ]
                              , length(results) > 0 ? results[e][t] : Float64(length(composition[e])-t)  ) for t in 1:length(composition[e]) ]
                        ,0.0) for e in 1:length(composition) ]

        b = new(time, events , skills, agents, p_draw)

        iteration(b)
        return b
    end
    function Batch(;composition::Vector{Vector{Vector{String}}}, results::Vector{Vector{Float64}}
                 ,time::Int64=0 , agents::Dict{String,Agent}=Dict{String,Agent}(), p_draw::Float64 = 0.0)
        Batch(composition, results, time, agents, p_draw)
    end
end

Base.show(io::IO, b::Batch) = print("Batch(time=", b.time, ", events=", b.events, ")")
Base.length(b::Batch) = length(b.events)

function add_events(b::Batch, composition::Vector{Vector{Vector{String}}}, results::Vector{Vector{Float64}}=Vector{Vector{Float64}}())
    this_agents = Set(vcat((composition...)...))
    for a in this_agents#a="c"
        elapsed = compute_elapsed(b.agents[a].last_time , b.time )
        if !haskey(b.skills,a)
            b.skills[a] = Skill(receive(b.agents[a],elapsed) ,Ninf ,Ninf , elapsed)
        else
            b.skills[a].elapsed = elapsed
            b.skills[a].forward = receive(b.agents[a],elapsed)
        end
    end
    from = length(b)+1
    for e in 1:length(composition)
        event = Event([Team([Item(composition[e][t][a], Ninf) for a in 1:length(composition[e][t]) ]
                              , length(results) > 0 ? results[e][t] : Float64(length(composition[e])-t) ) for t in 1:length(composition[e]) ] , 0.0)
        push!(b.events, event)
    end
    iteration(b, from)
end
function posterior(b::Batch, agent::String)#agent="a_b"
    return b.skills[agent].likelihood*b.skills[agent].backward*b.skills[agent].forward
end
function posteriors(b::Batch)
    res = Dict{String,Gaussian}()
    for (a, s) in b.skills
        res[a] = posterior(b,a)
    end
    return res
end
function within_prior(b::Batch, item::Item, online = false)
    r = b.agents[item.agent].player
    if online
        return Player(b.skills[item.agent].forward,r.beta,r.gamma)
    else
        wp = posterior(b,item.agent)/item.likelihood
        return Player(wp,r.beta,r.gamma)
    end
end
function within_priors(b::Batch, event::Int64; online = false)#event=1
    return [ [within_prior(b,item,online) for item in team.items ] for team in b.events[event].teams ]
end
function iteration(b::Batch, from::Int64 = 1)
    for e in from:length(b)#e=1

        g = Game(within_priors(b, e), outputs(b.events[e]), b.p_draw)

        for (t, team) in enumerate(b.events[e].teams)#(t,team) = (2, b.events[e].teams[2])
            for (i, item) in enumerate(team.items)#(i, item) = (2, team.items[2])
                b.skills[item.agent].likelihood = (b.skills[item.agent].likelihood / item.likelihood) * g.likelihoods[t][i]
                item.likelihood = g.likelihoods[t][i]
            end
        end

        b.events[e].evidence = g.evidence
    end
end
function log_evidence2(b::Batch; online::Bool = false, agents::Vector{String} = Vector{String}())
    if isempty(agents)
        if online
            return sum([log(Game(within_priors(b, e, online=true), outputs(b.events[e]), b.p_draw).evidence) for e in 1:length(b)])
        else
            return sum([log(event.evidence) for event in b.events])
        end
    else
        filter = [!isdisjoint(vcat((comp...)...),agents) for comp in get_composition(b.events)]
        if online
            return sum([log(Game(within_priors(b, e, online=true), outputs(b.events[e]), b.p_draw).evidence) for e in 1:length(b) if filter[e] ])
        else
            return sum([log(b.events[e].evidence) for e in 1:length(b.events) if filter[e]])
        end
    end
end
function convergence(b::Batch, epsilon::Float64=1e-6, iterations::Int64 = 20)
    iter = 0::Int64
    step = (Inf, Inf)
    while (step > epsilon) & (iter < iterations)
        old = copy(posteriors(b))
        iteration(b)
        step = diff(old, posteriors(b))
        iter += 1
    end
    return iter
end
function forward_prior_out(b::Batch, agent::String)
    return b.skills[agent].forward * b.skills[agent].likelihood
end
function backward_prior_out(b::Batch, agent::String)
    N = b.skills[agent].likelihood*b.skills[agent].backward
    return forget(N, b.agents[agent].player.gamma, b.skills[agent].elapsed)
end
function new_backward_info(b::Batch)
    for a in keys(b.skills)
        b.skills[a].backward = b.agents[a].message
    end
    return iteration(b)
end
function new_forward_info(b::Batch)
    for a in keys(b.skills)
        b.skills[a].forward = receive(b.agents[a], b.skills[a].elapsed)
    end
    return iteration(b)
end
mutable struct History
    size::Int64
    batches::Vector{Batch}
    agents::Dict{String,Agent}
    time::Bool
    mu::Float64
    sigma::Float64
    beta::Float64
    gamma::Float64
    p_draw::Float64
    online::Bool
    function History(composition::Vector{Vector{Vector{String}}}, results::Vector{Vector{Float64}}=Vector{Vector{Float64}}(), times::Vector{Int64}=Int64[], priors::Dict{String,Player}=Dict{String,Player}(); mu::Float64=MU, sigma::Float64=SIGMA, beta::Float64=BETA, gamma::Float64=GAMMA, p_draw::Float64=P_DRAW, online::Bool=false)
        (length(results) > 0) & (length(composition) != length(results)) && throw(error("(length(times) > 0) & (length(composition) != length(results))"))
        (length(times) > 0) & (length(composition) != length(times)) && throw(error("length(times) > 0) & (length(composition) != length(times))"))

        agents = Dict([ (a, Agent(haskey(priors, a) ? priors[a] : Player(Gaussian(mu, sigma), beta, gamma), Ninf, minInt64)) for a in Set(vcat((composition...)...)) ])
        h = new(length(composition), Vector{Batch}(), agents, length(times)>0, mu, sigma, beta, gamma, p_draw, online)
        trueskill(h, composition, results, times, online)
        return h
    end
    function History(;composition::Vector{Vector{Vector{String}}},results::Vector{Vector{Float64}}=Vector{Vector{Float64}}(),times::Vector{Int64}=Int64[],priors::Dict{String,Player}=Dict{String,Player}(), mu::Float64=MU, sigma::Float64=SIGMA, beta::Float64=BETA, gamma::Float64=GAMMA, p_draw::Float64=P_DRAW, online::Bool=false)
        History(composition, results, times, priors, mu=mu, sigma=sigma, beta=beta, gamma=gamma, p_draw=p_draw, online=online)
    end
end

# History(composition::Vector{Vector{Vector{String}}}, results::Vector{Vector{Int64}}=Vector{Vector{Int64}}(), times::Vector{Int64}=Int64[], priors::Dict{String,Player}=Dict{String,Player}(); mu::Float64=MU, sigma::Float64=SIGMA, beta::Float64=BETA, gamma::Float64=GAMMA, p_draw::Float64=P_DRAW, online::Bool=false) = History(composition, convert(Vector{Vector{Float64}},results), times, priors, mu=mu, sigma=sigma, beta=beta, gamma=gamma, p_draw=p_draw, online=online)
#
# History(;composition::Vector{Vector{Vector{String}}}, results::Vector{Vector{Int64}}=Vector{Vector{Int64}}(), times::Vector{Int64}=Int64[], priors::Dict{String,Player}=Dict{String,Player}(), mu::Float64=MU, sigma::Float64=SIGMA, beta::Float64=BETA, gamma::Float64=GAMMA, p_draw::Float64=P_DRAW, online::Bool=false) = History(composition, convert(Vector{Vector{Float64}},results), times, priors, mu=mu, sigma=sigma, beta=beta, gamma=gamma, p_draw=p_draw, online=online)

Base.length(h::History) = h.size
Base.show(io::IO, h::History) = print("History(Events=", h.size
                                     ,", Batches=", length(h.batches)
                                    ,", Agents=", length(h.agents), ")")
function trueskill(h::History, composition::Vector{Vector{Vector{String}}},results::Vector{Vector{Float64}}, times::Vector{Int64}, online::Bool)
    o = length(times)>0 ? sortperm(times) : [i for i in 1:length(composition)]
    i = 1::Int64
    last = 0.0
    while i <= length(h)
        j, t = i, length(times) == 0 ? i : times[o[i]]
        while ((length(times)>0) & (j < length(h)) && (times[o[j+1]] == t)) j += 1 end
        if length(results)>0
            b = Batch(composition[o[i:j]],results[o[i:j]], t, h.agents, h.p_draw)
        else
            b = Batch(composition[o[i:j]], results , t, h.agents, h.p_draw)
        end
        push!(h.batches,b)
        if online
            new = round(100*(i/length(h)))
            if new != last
                print("\r",new,"%")
                last = new
            end
            for a in keys(b.skills)
                b.skills[a].online = b.skills[a].forward
            end
            convergence(h)
        end
        for a in keys(b.skills)
            h.agents[a].last_time = length(times) == 0 ? maxInt64 : t
            h.agents[a].message = forward_prior_out(b,a)
        end
        i = j + 1
    end
    if online println("\r100.0%") end
end
function diff(old::Dict{String,Gaussian}, new::Dict{String,Gaussian})
    step = (0., 0.)
    for a in keys(old)
        step = max(step, delta(old[a],new[a]))
    end
    return step
end
function iteration(h::History)
    step = (0., 0.)

    clean(h.agents)
    for j in length(h.batches)-1:-1:1# j = 2
        for a in keys(h.batches[j+1].skills)
            h.agents[a].message = backward_prior_out(h.batches[j+1],a)
        end
        old = copy(posteriors(h.batches[j]))
        new_backward_info(h.batches[j])
        step = max(step, diff(old, posteriors(h.batches[j])))
    end

    clean(h.agents)
    for j in 2:length(h.batches)# j = 3
        for a in keys(h.batches[j-1].skills)
            h.agents[a].message = forward_prior_out(h.batches[j-1],a)
        end
        old = copy(posteriors(h.batches[j]))
        new_forward_info(h.batches[j])
        step = max(step, diff(old, posteriors(h.batches[j])))
    end

    if (length(h.batches) == 1)
        old = copy(posteriors(h.batches[1]))
        iteration(h.batches[1])
        step = max(step, diff(old, posteriors(h.batches[1])))
    end

    return step
end
function convergence(h::History; epsilon::Float64=EPSILON, iterations::Int64=ITERATIONS, verbose = true)
    step = (Inf, Inf)::Tuple{Float64,Float64}
    iter = 1::Int64
    while (step > epsilon) & (iter <= iterations)
        verbose && print("Iteration = ", iter)
        step = iteration(h)
        iter += 1
        verbose && println(", step = ", step)
    end
    verbose && println("End")
    return step, iter
end
# function posteriors(h::History)
#     res = Vector{Vector{Vector{Gaussian}}}()
#     for b in h.batches
#         for e in b.events
#             ps = Vector{Vector{Gaussian}}()
#             for a in get_composition(e)
#
#                 t_p = (b.time, posterior(b,a))
#                 res[a] = haskey(res, a) ? push!(res[a],t_p) : [t_p]
#         end
#     end
# end
function learning_curves(h::History)
    res = Dict{String,Array{Tuple{Int64,Gaussian}}}()
    for b in h.batches
        for a in keys(b.skills)
            t_p = (b.time, posterior(b,a))
            res[a] = haskey(res, a) ? push!(res[a],t_p) : [t_p]
        end
    end
    return res
end
function log_evidence(h::History; online::Bool = false, agents::Vector{String} = Vector{String}() )
    return sum([log_evidence2(b, online=online, agents = agents) for b in h.batches])
end
function add_events(h::History,composition::Vector{Vector{Vector{String}}};results::Vector{Vector{Float64}}=Vector{Vector{Float64}}(),times::Vector{Int64}=Int64[],priors::Dict{String,Player}=Dict{String,Player}())
    add_events(h, composition, results, times, priors)
end
function add_events(h::History,composition::Vector{Vector{Vector{String}}},results::Vector{Vector{Float64}}=Vector{Vector{Float64}}(),times::Vector{Int64}=Int64[],priors::Dict{String,Player}=Dict{String,Player}())

    (length(times)>0) & !h.time && throw(error("length(times)>0 but !h.time"))
    (length(times)==0) & h.time && throw(error("length(times)==0 but h.time"))
    (length(results) > 0) & (length(composition) != length(results)) && throw(error("(length(results) > 0) & (length(composition) != length(results))"))
    (length(times) > 0) & (length(composition) != length(times)) && throw(error("length(times) > 0) & (length(composition) != length(times))"))

    this_agents = Set(vcat((composition...)...))
    for a in this_agents
        if !haskey(h.agents,a)
            h.agents[a] = Agent(haskey(priors, a) ? priors[a] : Player(Gaussian(h.mu, h.sigma), h.beta, h.gamma), Ninf, minInt64)
        end
    end

    clean(h.agents,true)
    n = length(composition)
    o = length(times)>0 ? sortperm(times) : [i for i in 1:length(composition)]
    i = 1::Int64; k = 1::Int64
    while i <= n
        j, t = i, length(times) == 0 ? i : times[o[i]]
        while ((length(times)>0) & (j < n) && (times[o[j+1]] == t)) j += 1 end
        while (!h.time & (h.size > k) ) || (h.time && (length(h.batches) >= k) && (h.batches[k].time < t))
            b = h.batches[k]
            if (k>1) new_forward_info(b) end
            for a in intersect(keys(b.skills), this_agents)#a ="a"
                b.skills[a].elapsed = compute_elapsed(b.agents[a].last_time , b.time )
                h.agents[a].last_time = length(times) == 0 ? maxInt64 : b.time
                h.agents[a].message = forward_prior_out(b,a)
            end
            k += 1
        end
        if (h.time && (length(h.batches) >= k) && (h.batches[k].time == t))
            b = h.batches[k]
            if length(results)>0
                add_events(b, composition[o[i:j]],results[o[i:j]])
            else
                add_events(b, composition[o[i:j]])
            end
        else
            if !h.time k = k + 1 end
            if length(results)>0
                b =  Batch(composition[o[i:j]],results[o[i:j]], t, h.agents, h.p_draw)
            else
                b =  Batch(composition[o[i:j]], results , t, h.agents, h.p_draw)
            end
            insert!(h.batches,  k , b)
            if h.time k = k + 1 end
        end
        for a in keys(b.skills)#a="a"
            h.agents[a].last_time = length(times) == 0 ? maxInt64 : t
            h.agents[a].message = forward_prior_out(b,a)
        end
        i = j + 1
    end
    while h.time && (length(h.batches) >= k)
        b = h.batches[k]
        new_forward_info(b)
        for a in intersect(keys(b.skills), this_agents)#a ="a"
            b.skills[a].elapsed = compute_elapsed(b.agents[a].last_time , b.time )
            h.agents[a].last_time = length(times) == 0 ? maxInt64 : b.time
            h.agents[a].message = forward_prior_out(b,a)
        end
        k += 1
    end
    h.size = h.size + n
    iteration(h)
end

end # module


using DifferentialEquations
using LinearAlgebra
using Plots
using DeltaRobot

# u0 =  [pi/4; pi/2; 0.0; pi/4; pi/2; 0.0; pi/4; pi/2; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]; # what should lambda be?
u0 =  [pi/4; 3*pi/4; 0.0; pi/4; 3*pi/4; 0.0; pi/4; 3*pi/4; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]; # what should lambda be?
du0 = [0.0; 0.0 ; 0.0 ; 0.0; 0.0 ; 0.0 ; 0.0; 0.0 ; 0.0 ; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]; # why must I give this?
differential_vars = [true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, false, false, false, false, false, false]
@assert length(u0) == length(du0)
@assert length(u0) == length(differential_vars)
@assert u0[10:18] == du0[1:9] # make sure initial conditions are consistent
p = DeltaRobotArmParams(0.3, 0.2, 0.1, 0.5, 0.5, 9.8, 0.09, 0.09, 0.09);
tspan = (0.0, 20.0)
prob = DAEProblem(delta_robot!, du0, u0, tspan, p, differential_vars=differential_vars)
dt = 0.05
sol = solve(prob, saveat=dt)

gr()
anim = animate_robot(sol, p)
gif(anim, "animations/robot_1.gif", fps=convert(Int64, 1 / dt))

plotly()
plot(sol, vars=(1, 2, 3))
plot(sol.t, sol[1,:])
plot(sol.t, sol[2,:])
plot(sol.t, sol[3,:])

gr()
prob_damped = DAEProblem(delta_robot_damped!, du0, u0, tspan, p, differential_vars=differential_vars)
sol_damped = solve(prob_damped, saveat=dt)
anim = animate_robot(sol_damped, p)
gif(anim, "animations/damp_robot_1.gif", fps=convert(Int64, 1 / dt))


@testset "lq_approx.jl" begin

    @testset "QuantEcon LQ Lecture Example" begin
        # Set parameters
        α, β, ρ1, ρ2, σ = 10.0, 0.95, 0.9, 0.0, 1.0

        R = 1 / β
        A = [1.0 0.0 0.0;
               α  ρ1  ρ2;
            0.0 1.0 0.0]
        A12 = zeros(3,1)
        ALQ_l = hcat(A, A12)
        ALQ_r = [0 -R 0 R]
        ALQ = vcat(ALQ_l, ALQ_r)

        RLQ = [0.0 0.0 0.0 0.0;
               0.0 0.0 0.0 0.0;
               0.0 0.0 0.0 0.0;
               0.0 0.0 0.0 1e-9]

        QLQ = 1.0
        BLQ = [0.0; 0.0; 0.0; R]
        β_LQ = β

        f(s, x) = -([1, s...]' * RLQ * [1, s...] + x' * QLQ * x);
        g(s, x) = ALQ * [1, s...] + BLQ * x;

        s_star, x_star = [1.0, 10.9, 1.0], 0.0;

        f_star = f(s_star, x_star)
        Df_star = [0.0, 0.0, -2.0e-9, 0.0]
        D²f_star =  [0.0 0.0 0.0 0.0;
                     0.0 0.0 0.0 0.0;
                     0.0 0.0 -2.0e-9 0.0;
                     0.0 0.0 0.0 -2.0]
        g_star = g(s_star, x_star)[2:end]
        Dg_star = [0.9 0.0 0.0 0.0;
                   1.0 0.0 0.0 0.0;
                   -1.052631578947368 0.0 1.052631578947368 1.052631578947368]

       lq = approx_lq(s_star, x_star, f_star, Df_star, D²f_star, g_star,
                      Dg_star, β)

       @test isapprox(lq.R, RLQ)
       @test isapprox(lq.Q, [QLQ])
       @test isapprox(lq.A, ALQ)
       @test isapprox(lq.B, BLQ)
    end

    @testset "One State and One Action Variable Example" begin
        A = [1.0 0.0;
             0.5 0.9];

        R = [5.0 0.0;
             0.0 0.3]

        N = [0.05 0.1]

        Q = [0.1];

        B = [0.0; 1.5];

        f(s, x) = -([1, s...]' * R * [1, s...] .+ x' * Q * x
                    .+ 2 * x' * N * [1, s...]);
        g(s, x) = A * [1, s...] + B * x;

        s_star = 5.0;
        x_star = 0.0;

        g_star = g(s_star, x_star)[2];
        f_star = f(s_star, x_star)[1];

        Df_star = [-3.0;
                   -1.1];
        D²f_star = [-0.6 -0.2;
                    -0.2 -0.2]
        Dg_star = [0.9 1.5];

        lq = approx_lq(s_star, x_star, f_star, Df_star, D²f_star, g_star,
                       Dg_star, 1.0)

        @test isapprox(lq.R, R)
        @test isapprox(lq.Q, Q)
        @test isapprox(lq.A, A)
        @test isapprox(lq.B, B)
        @test isapprox(lq.N, N)
    end

    @testset "One State and Two Action Variables Example" begin
        A = [1.0 0.0;
               0.5 0.9];

        R = [5.0 0.0;
               0.0 0.3]

        N = [0.05 0.1;
                0.01 0.02]

        Q = [0.1 -0.05;
                -0.05 0.3];

        B = [0.0 0.0; 1.5 -0.9];

        f(s, x) = -([1, s...]' * R * [1, s...] .+ x' * Q * x
                    .+ 2 * x' * N * [1, s...]);
        g(s, x) = A * [1, s...] + B * x;

        x_star = [0.0, 0.0];
        s_star = 5.0;

        g_star = g(s_star, x_star)[2];
        f_star = f(s_star, x_star)[1];

        Df_star = [-3.0;
                   -1.1;
                   -0.22];
        D²f_star = [-0.6 -0.2 -0.04;
                    -0.2 -0.2 0.1;
                    -0.04 0.1 -0.6]
        Dg_star = [0.9 1.5 -0.9];

        lq = approx_lq(s_star, x_star, f_star, Df_star, D²f_star, g_star,
                       Dg_star, 1.0)

        @test isapprox(lq.R, R)
        @test isapprox(lq.Q, Q)
        @test isapprox(lq.A, A)
        @test isapprox(lq.B, B)
        @test isapprox(lq.N, N)

    end

    @testset "Three States and Two Action Variables Example" begin
        A = [1.0 0.0 0.0 0.0;
             -0.2 0.5 0.01 0.0;
             0.05 -0.5 0.8 0.3;
            -0.1 -0.2 0.3 0.6];

        R = [5.0 0.00 0.00 0.00;
             0.0 0.15 0.21 -0.05;
             0.0 0.21 0.5 -0.01;
             0.0 -0.05 -0.01 0.8]

        N = [-0.1 0.005 -0.15 -0.2;
             0.005 -0.15 -0.2 0.2]

        Q = [1.0 -0.5;
             -0.5 0.9];

        B = [0.0 0.0;
             0.3 0.5;
             0.0 0.2;
             0.1 0.3];

       f(s, x) = ([1, s...]' * R * [1, s...] + x' * Q * x
                  + 2 * x' * N * [1, s...]);
       g(s, x) = A * [1, s...] + B * x;

       s_star = [-0.7916666666666763, -19.583333333333226, -14.541666666666579];
       x_star = [0.0, 0.0];

       g_star = g(s_star, x_star)[2:end];  # Drop constant
       f_star = f(s_star, x_star);

       Df_star = [-7.008333333332452;
                  -19.62499999999752;
                  -22.79583333333047;
                  11.483749999998526;
                  2.264166666666381];

       D²f_star = [0.3 0.42 -0.1 0.01 -0.3;
                   0.42 1.0 -0.02 -0.3 -0.4;
                   -0.1 -0.02 1.6 -0.4 0.4;
                   0.01 -0.3 -0.4 2.0 -1.0;
                   -0.3 -0.4 0.4 -1.0 1.8];

       Dg_star = [0.5 0.01 0.0 0.3 0.5;
                  -0.5 0.8 0.3 0.0 0.2;
                  -0.2 0.3 0.6 0.1 0.3];

       lq = approx_lq(s_star, x_star, f_star, Df_star, D²f_star, g_star,
                      Dg_star, 1.0);

        @test isapprox(-lq.R, R)
        @test isapprox(-lq.Q, Q)
        @test isapprox(lq.A, A)
        @test isapprox(lq.B, B)
        @test isapprox(-lq.N, N)

    end

    @testset "Monetary Policy Example" begin  # See 9.4.6
        vlq =  1.0e+02 * [-1.516771891935916  # Obtained from Fackler's code
                          -1.421873224803633
                          -1.253495277719767
                          -1.054476825045721
                          -0.919716533913779
                          -0.849214404323941
                          -0.842970436276206
                          -0.900984629770574
                          -0.975359737401322
                          -1.023256984807047
                          -1.350214126639643
                          -1.257350153002016
                          -1.093041592907463
                          -0.900127220717386
                          -0.771471010069413
                          -0.707072960963543
                          -0.706933073399777
                          -0.771051347378114
                          -0.849495841998174
                          -0.899427782898555
                          -1.056517641369012
                          -0.967723054720698
                          -0.811553268604770
                          -0.630847057382630
                          -0.514399007702594
                          -0.462209119564662
                          -0.474277392968833
                          -0.550603827915108
                          -0.637187096513793
                          -0.691188424403486
                          -0.714520526767850
                          -0.631830020603504
                          -0.487868395455513
                          -0.325474425685280
                          -0.227338617457150
                          -0.193460970771124
                          -0.223841485627201
                          -0.318480162025382
                          -0.417271591592004
                          -0.477376999965666
                          -0.490780548132429
                          -0.414194122452051
                          -0.282440658271998
                          -0.138358929953671
                          -0.058535363177447
                          -0.042969957943327
                          -0.091662714251310
                          -0.204613632101397
                          -0.315613222635957
                          -0.381822711493587
                          -0.385297705462748
                          -0.314815360266339
                          -0.195270057054223
                          -0.069500570187802
                          -0.007989244863485
                          -0.010736081081270
                          -0.077741078841160
                          -0.209004238143153
                          -0.332211989645650
                          -0.404525558987249
                          -0.398071998758808
                          -0.333693734046368
                          -0.226356591802190
                          -0.118899346387675
                          -0.075700262515263
                          -0.096759340184955
                          -0.182076579396751
                          -0.331651980150650
                          -0.467067892621084
                          -0.545485542446652
                          -0.529103428020609
                          -0.470829243792138
                          -0.375700262515897
                          -0.286555258553288
                          -0.261668416132782
                          -0.301039735254381
                          -0.404669215918082
                          -0.572556858123887
                          -0.720180931562259
                          -0.804702661871796
                          -0.682156123064999
                          -0.627951325825841
                          -0.540961118528224
                          -0.464024275533553
                          -0.451345594080985
                          -0.502925074170520
                          -0.618762715802159
                          -0.798858518975902
                          -0.954621366392899
                          -1.043212483691748
                          -0.778391993248151
                          -0.726221889503649
                          -0.643301069195345
                          -0.572468306684642
                          -0.565893705716043
                          -0.623577266289547
                          -0.745518988405155
                          -0.931718872062866
                          -1.091551106469175
                          -1.182176917262681]

        bet = [0.8 0.5; 0.2 0.6]
        gamm = [-0.8, 0.0]
        Omega = [0.3 0.0; 0.0 1.0]
        s_target = [0., 1.]
        alpha = [0.9, 0.4]
        Sigma = 0.04 * Matrix(I, 2, 2)
        discount = 0.9;

        k = [3, 3];
        m = [10, 10];
        breaks = m - (k.-1);
        s_min = [-15, -10];
        s_max = [15, 10];
        basis = Basis(map(SplineParams, breaks, s_min, s_max, k)...);

        s_star = [0., 1.];  # steady-state states
        x_star = (s_star[1] - alpha[1] - dot(bet[1,:],  s_star)) / gamm[1];
        e_star = [0.0, 0.0];

        f(s, x) = -(s - s_target)' * Omega * (s - s_target) / 2
        g(s, x, e) = alpha + bet * s + gamm * x + e

        f_star = f(s_star, x_star)
        Df_star = [0.0, 0.0, 0.0]
        D²f_star = [-0.3 0.0 0.0;
                    0.0 -1.0 0.0;
                    0.0 0.0 0.0]
        g_star = g(s_star, x_star, e_star)
        Dg_star = [0.8 0.5 -0.8;
                   0.2 0.6 0.0]

        lq = approx_lq(s_star, x_star, f_star, Df_star, D²f_star, g_star,
                       Dg_star, discount);

        P, F, d = stationary_values(lq)

        states = [ones(100) nodes(basis)[1]]
        vlq_approx = -[states[i, 1:3]' * P * states[i, 1:3] for i ∈ 1:100]

        @test isapprox(vlq_approx, vlq)
    end

end


function ybus!(
    ybus::SparseMatrixCSC{ComplexF64, Int64},
    b::PSY.ACBranch,
    num_bus::Dict{Int64, Int64},
    mult::Float64 = 1.0,
)
    arc = PSY.get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]

    Y_l = (1 / (PSY.get_r(b) + PSY.get_x(b) * 1im))
    Y11 = Y_l + (1im * PSY.get_b(b).from)

    ybus[bus_from_no, bus_from_no] += mult * Y11
    Y12 = -Y_l
    ybus[bus_from_no, bus_to_no] += mult * Y12
    #Y21 = Y12
    ybus[bus_to_no, bus_from_no] += mult * Y12
    Y22 = Y_l + (1im * PSY.get_b(b).to)
    ybus[bus_to_no, bus_to_no] += mult * Y22
    return
end

function ybus!(
    ybus::SparseMatrixCSC{ComplexF64, Int64},
    b::PSY.TapTransformer,
    num_bus::Dict{Int64, Int64},
    mult::Float64 = 1.0,
)
    arc = PSY.get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]

    Y_t = 1 / (PSY.get_r(b) + PSY.get_x(b) * 1im)
    c = 1 / PSY.get_tap(b)
    b = PSY.get_primary_shunt(b)

    Y11 = (Y_t * c^2)
    ybus[bus_from_no, bus_from_no] += mult * Y11
    Y12 = (-Y_t * c)
    ybus[bus_from_no, bus_to_no] += mult * Y12
    #Y21 = Y12
    ybus[bus_to_no, bus_from_no] += mult * Y12
    Y22 = Y_t
    ybus[bus_to_no, bus_to_no] += mult * (Y22 + (1im * b))

    return
end

function ybus!(
    ybus::SparseMatrixCSC{ComplexF64, Int64},
    b::PSY.PhaseShiftingTransformer,
    num_bus::Dict{Int64, Int64},
    mult::Float64 = 1.0,
)
    arc = PSY.get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]

    Y_t = 1 / (PSY.get_r(b) + PSY.get_x(b) * 1im)
    tap = (PSY.get_tap(b) * exp(PSY.get_α(b) * 1im))
    c_tap = (PSY.get_tap(b) * exp(-1 * PSY.get_α(b) * 1im))
    b = PSY.get_primary_shunt(b)

    Y11 = (Y_t / abs(tap)^2)
    ybus[bus_from_no, bus_from_no] += mult * Y11
    Y12 = (-Y_t / c_tap)
    ybus[bus_from_no, bus_to_no] += mult * Y12
    Y21 = (-Y_t / tap)
    ybus[bus_to_no, bus_from_no] += mult * Y21
    Y22 = Y_t
    ybus[bus_to_no, bus_to_no] += mult * (Y22 + (1im * b))
    return
end

function ybus!(
    ybus::SparseMatrixCSC{ComplexF64, Int64},
    fa::PSY.FixedAdmittance,
    num_bus::Dict{Int64, Int64},
    mult::Float64 = 1.0,
)
    bus = PSY.get_bus(fa)
    bus_ix = num_bus[PSY.get_number(bus)]
    ybus[bus_ix, bus_ix] += mult * fa.Y
    return
end


using Test, Parameters, BenchmarkTools

using FourierGPE

import FourierGPE:V
# ==== Units: ========================
# Units of ξ for length, 1/μ for time
# related by μ ≡ ħ²/mξ² = mc^2 where
# c = ħ/mξ is the speed of sound of
# the uniform system.

# ==== set simulation parameters ====
L = (16.,16.,16.)
N = (64,64,64)
γ = 0.05
tf = 4/γ
Nt = 200
t = LinRange(0.,tf,Nt)
# ========= Initialize simulation ======
X = xvecs(L,N)
x,y,z = X
y = y'
z = reshape(z,1,1,length(z))
V0 = @. zero(x*y*z)

sim = Sim(L,N)
# sim = Sim(sim,V0=StaticPotential(V0))

V(x,y,z,t) = 0.25*(x^2 + y^2 + z^2)*(1+0.01*sin(t))

@pack! sim = γ,tf,Nt,t
@unpack μ,g = sim

# ========= useful state functions

ψ0(x,y,z,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,z,0)/μ,0.0)+im*0.0)
healing(x,y,z,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,z,μ,g)))

#make initial state
ψi = ψ0.(x,y,reshape(z,1,1,length(z)),μ,g)

dψi = similar(ψi)

ϕi = kspace(ψi,sim)
dϕi = similar(ϕi)
@pack! sim = ϕi,γ,tf,t

@time nlin!(dψi,ψi,sim,0.1)


## Taken from version 0.4
cld{T<:Integer }(x::T, y::T) = div(x,y)+(!signbit(x$y)&(rem(x,y)!=0))

function makedata!(y::Vector{Float64}, x::Matrix{Float64})
    n = length(y)
    size(x)[1] == n || error("x and y have incompatible dimensions")
    randn!(y)
    for i in 1:n
        x[i] = 1
    end
    if length(x) > n
        for i in (n+1):length(x)
            x[i] = randn()
        end
    end
end

## Let's just do one for forecast unbiasedness first.
function oosstat!(f::Vector{Float64}, coef::Vector{Float64},
                  y::Vector{Float64}, x::Matrix{Float64})
    R = length(y) - length(f)
    for t = R:(length(y) - 1)
        coef[:] = x[1:t,:] \ y[1:t]
        f[t-R+1] = y[t+1] - (x[t+1,:] * coef)[1]
    end
    mean(f)
end

function bootmean!(ffull::Vector{Float64}, coef::Vector{Float64},
                   y::Vector{Float64}, x::Matrix{Float64})
    coef[:] = x \ y
    ffull[:] = y - x * coef
    mean(ffull)
end

## Let's just do one for absolute forecast error.
function oosstat2!(f::Vector{Float64}, coef::Vector{Float64},
                  y::Vector{Float64}, x::Matrix{Float64})
    R = length(y) - length(f)
    for t = R:(length(y) - 1)
        coef[:] = x[1:t,:] \ y[1:t]
        f[t-R+1] = abs(y[t+1] - (x[t+1,:] * coef)[1])
    end
    mean(f)
end

function bootmean2!(ffull::Vector{Float64}, coef::Vector{Float64},
                   y::Vector{Float64}, x::Matrix{Float64})
    coef[:] = x \ y
    ffull[:] = abs(y - x * coef)
    mean(ffull)
end

function runmc!(mcresults, nboot, P, R, k, α,
                oosstat!, bootmean!, truemean)
    y = Array(Float64, P+R)
    x = Array(Float64, P+R, k)
    yboot = similar(y)
    xboot = similar(x)
    f = Array(Float64,  P)
    ffull = Array(Float64, P+R)
    coef = Array(Float64, k)
    bootvalue = Array(Float64, nboot)

    for i in 1:length(mcresults)
        makedata!(y, x)
        oosvalue = oosstat!(f, coef, y, x) - truemean
        bootnull = bootmean!(ffull, coef, y, x)

        for j in 1:nboot
            blockboot!(yboot, xboot, y, x, 8)
            @inbounds bootvalue[j] = oosstat!(f, coef, yboot, xboot) - bootnull
        end
        crit = quantile(bootvalue, [α/2, 1 - α/2])
        @inbounds mcresults[i] = oosvalue < crit[1] || oosvalue > crit[2]
    end
end

mcresults1 = Array(Bool, 600)
@time runmc!(mcresults1, 499, 120, 240, 5, 0.05, oosstat!, bootmean!, 0)
mcresults2 = Array(Bool, 1500)
@time runmc!(mcresults2, 499, 120, 240, 5, 0.05, oosstat2!, bootmean2!, sqrt(2/π))


@testset "GS" begin

m = 20
n = 10
A = randn(m,n);

Aold  = copy(A)
Qc,Rc = cgs(A)
@test all(Aold.==A) # make sure input is unchanged
Qt,Rt = cgs!(A)
@test all(Qt.==A) # check that input IS overwritten
@test maximum(abs.(Qt-Qc)) < 1e-15
@test norm(Aold-Qc*Rc,Inf)/norm(A,Inf) < 1e-14
@test norm(Qc'*Qc-I,Inf) < 1e-14

Aold  = copy(A)
Qm,Rm = mgs(A)
@test all(Aold.==A) # make sure input is unchanged
Qt,Rt = mgs!(A)
@test all(Qt.==A) # check that input IS overwritten
@test maximum(abs.(Qt-Qm)) < 1e-15
@test norm(Aold-Qm*Rm,Inf)/norm(A) < 1e-14
@test norm(Qm'*Qm-I,Inf) < 1e-14

end


using Distributed
@everywhere using Mamba

## Data
pumps = Dict{Symbol, Any}(
  :y => [5, 1, 5, 14, 3, 19, 1, 1, 4, 22],
  :t => [94.3, 15.7, 62.9, 126, 5.24, 31.4, 1.05, 1.05, 2.1, 10.5]
)
pumps[:N] = length(pumps[:y])

## Model Specification
model = Model(

  y = Stochastic(1,
    (theta, t, N) ->
      UnivariateDistribution[(
        lambda = theta[i] * t[i];
        Poisson(lambda)) for i in 1:N
      ],
    false
  ),

  theta = Stochastic(1,
    (alpha, beta) -> Gamma(alpha, 1 / beta),
    true
  ),

  alpha = Stochastic(
    () -> Exponential(1.0)
  ),

  beta = Stochastic(
    () -> Gamma(0.1, 1.0)
  )

)

## Initial Values
inits = [
  Dict(:y => pumps[:y], :alpha => 1.0, :beta => 1.0,
       :theta => rand(Gamma(1.0, 1.0), pumps[:N])),
  Dict(:y => pumps[:y], :alpha => 10.0, :beta => 10.0,
       :theta => rand(Gamma(10.0, 10.0), pumps[:N]))
]

## Sampling Scheme
scheme = [Slice([:alpha, :beta], 1.0, Univariate),
          Slice(:theta, 1.0, Univariate)]
setsamplers!(model, scheme)

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

## Posterior Predictive Distribution
ppd = predict(sim, :y)
describe(ppd)


simple_loss(W::AbstractMatrix, b::AbstractVector, x::AbstractArray) = sum(W * x .+ b)

@testset "simplegrad" begin
    W, b, x = rand(128, 784), rand(128), rand(784, 100)

    val1, g = grad(simple_loss, W, b, x)

    simple_loss_grad = simplegrad(simple_loss, W, b, x)
    val2, dW, db, dx = simple_loss_grad(W, b, x)

    @test val1 == val2
    @test g[1] == dW
    @test g[2] == db
    @test g[3] == dx
end


using PyCMA
using Test

f(x) = sum(abs2, x)
res = PyCMA.cma.fmin(f, ones(3), .25)
@test f(res[1]) ≈ 0 atol = 1e-12


# This file is part of RLinearAlgebra.jl

module ProceduralTestLSVCDetermCycle

using Test, RLinearAlgebra, Random, LinearAlgebra

Random.seed!(1010)

@testset "LSVC Deterministic Cyclic -- Procedural" begin

    # Verify appropriate super type
    @test supertype(LinSysVecColDetermCyclic) == LinSysVecColSampler

    # Test for appropriate cycling of columns
    @test let
        A = rand(3, 10)
        b = rand(3)
        x = rand(10)

        cyc = LinSysVecColDetermCyclic()

        flag = true

        for i = 11:20
            v, M, res = RLinearAlgebra.sample(cyc, A, b, x, i)
            flag = flag & (v[i-10] == 1.0)
            flag = flag & (res ≈ dot(M * v, A * x - b))
        end

        flag
    end
end

end # End module


### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ f13a03a0-24e6-11ec-2825-8123926abe5d
begin
	import Pkg
	Pkg.activate(joinpath(@__DIR__))
	using SimLib, Plots, PlutoUI, Statistics, LinearAlgebra, LaTeXStrings, XXZNumerics, SpinSymmetry, LsqFit, SparseArrays
	gr()
	LinearAlgebra.BLAS.set_num_threads(8)
	#cd(joinpath(@__DIR__, ".."))
end

# ╔═╡ 3eef0c5c-9ef5-4e0e-ab9b-0721e68f8413
using BenchmarkTools

# ╔═╡ 25cbd399-a199-4469-be1d-f402df522511
md"""
# Entanglement Entropy implementation
## Basics
Have density matrix $\rho = \sum_i p_i |i\rangle\langle i|$ on full system. Can divide system in subsystems $A$ and $B$.

### Partial trace
```math
\rho_B = \mathrm{Tr}_A \rho = \sum_{j} \sum_i \langle a_i |\langle b_j| \rho |a_i \rangle b_j \rangle |b_j\rangle\langle b_j| = \sum_{i,j} \langle a_i b_j| \rho |a_i b_j \rangle |b_j\rangle\langle b_j|
```
where $\{a_i\}$ ($\{b_i\}$) are a basis of $A$ ($B$).

### Von-Neumann entropy:
```math
S(\rho) = \mathrm{Tr} \rho \log_2 \rho = \sum_i p_i \log_2 p_i
```

### Entanglement entropy
Entanglement entropy between subsystem $A$ and $B$ is
```math
S_{A} = S(\rho_A) = S(\mathrm{Tr}_B \rho) = S(\rho_B) = S_B
```
→ Need to compute partial trace and compute eigenvalues.

→ Better way for pure states: SVD
```math
|\psi\rangle = \sum_{i,j} c_{i,j} |a_i,b_j\rangle \overset{SVD}{=} \sum_i \sigma_i |\tilde{a}_i\rangle|\tilde{b}_i\rangle
```
Here we basically *reshaped* the state vector into a matrix and used singular value decomposition on it. This allows easy computation of the eigenvalues of the reduced density matrix. Consider:
```math
\begin{align}
\rho_B = \mathrm{Tr}_A \rho &= \mathrm{Tr}_A \sum_{i,j} \sigma_i \sigma_j^\star |\tilde{a}_i\rangle|\tilde{b}_i\rangle \langle\tilde{a}_j |\langle\tilde{b}_j|\\
&= \sum_j |\sigma_j|^2 |\tilde{b}_j\rangle \langle \tilde{b}_j|
\end{align}
```
**Note:** This also shows that $S(\rho_A) = S(\rho_B)$.

**Note 2:** The SVD way is computationally much faster than partial trace + eigen decomposition and saves a lot of memory as well.
"""

# ╔═╡ 3f654c7b-ff3b-423c-91ad-8c942ca5ece9
md"""
## Additional symmetries
For this computation having $|\psi\rangle$ in a subspace of a symmetry does not make the computation more straight-forward as the *reshape* operation does not work too well in all symmetrized bases.

What follows is explicitly only for spin-1/2 chains of $N$ spins but can probably be generalized more or less easily.

### z-Block basis
Let's denote the eigenspaces of the magnetization $\hat{S}_z$ by $L_k^N$ (where $k$ counts the number of $\uparrow$ in that space) and states by $|k, i\rangle$. The second index just enumerates the different basis vectors in this sector. The whole Hilbert space $\mathcal{H}$ decomposes:
```math
\mathcal{H} = \bigoplus_{k=0}^{N} L^N_k
```

Assume $|\psi\rangle \in L_k$ and thus $|\psi\rangle = \sum_i c_i |k,i\rangle$ and we want to trace out the first $N_1$ spins of the chain.

In this case we are lucky because $L^N_k$ again decompose nicely into smaller magnetization eigenspaces:
```math
L_k^N = \bigoplus_{k_1=k_{min}}^{k_{max}} L^{N_1}_{k_1} \otimes L^{N-N_1}_{k-k_1}
```
with $k_{min} = \max\{0, N_1+k-N\}$ and $k_{max} = \min\{k, N_1\}$.

This means, that the reduced density matrix is actually block diagonal! And the computation of the Schmidt values can be simplified as well as for every $k$ we again can reshape the corresponding vectors into a rectangular matrix.

```math
|\psi\rangle = \sum_i c_i |k,i\rangle = \sum_{k_1=k_{min}}^{k_{max}} \sum_{i,j} c_{i,j} |k_1,i\rangle|k-k_1,j\rangle
```

**Implementation note:**
The only thing left is to determine the correct index $i \in L^N_k$ from $(i,j) \in L^{N_1}_{k_1} \otimes L^{N-N_1}_{k-k_1}$. I currently do this via a lookup table and some index math. I already have functions to compute the indices of $L^N_k \subset \mathcal{H}$.

So what I do is:
1. Compute the indices $I$ and $J$ of $L^{N_1}_{k_1} \subset L^N_k$ and $L^{N-N_1}_{k-k_1} \subset L^N_k$
2. This gives the indices of $L^{N_1}_{k_1} \otimes L^{N-N_1}_{k-k_1} \subset L^N_k$ by $I\otimes J = \{2^{N-N_1}i+j | i\in I, j\in J\}$.
3. Thus we have the subspace's indices $I\otimes J \subset L^N_k \subset \mathcal{H}$
4. This I map via lookup table to the actual indices of the state vector.
"""

# ╔═╡ fc4f5bbb-cc20-45ed-aa1a-928ae12b1566
md"""
### Parity (spin flip)

This is not as easy although we can find a decomposition. Let $P^N_\pm$ denote the symmetry sectors. We can see:
```math
P^N_+ = P^{N_1}_+ \otimes P^{N-N_1}_+ \oplus P^{N_1}_- \otimes P^{N-N_1}_-
```

Again we have 2 blocks not talking to each other that may be treated seperately. Finding the indices is a bit more involved this time around.

Consider a state $|s\rangle \in \mathcal{H} = |s_1, s_2\rangle$ and its flipped partner $|\bar{s}\rangle = |\bar{s}_1, \bar{s}_2\rangle$. Also consider the half-flipped state $|s^\star\rangle = |s_1, \bar{s}_2\rangle$ and its partner $|\bar{s}^\star\rangle = |\bar{s}_1, s_2\rangle$. Thus we have $|S\rangle = |s\rangle + |\bar{s}\rangle \in P^N_+$ and $|S^\star\rangle = |s^\star\rangle + |\bar{s}^\star\rangle \in P^N_+$. One can directly construct:
```math
\begin{align}
|t_+\rangle &= |S\rangle + |S^\star\rangle = (|s_1\rangle + |\bar{s}_1\rangle)\otimes (|s_2\rangle + |\bar{s}_2\rangle) \in P^{N_1}_+ \otimes P^{N-N_1}_+\\
|t_-\rangle &= |S\rangle - |S^\star\rangle = (|s_1\rangle - |\bar{s}_1\rangle)\otimes (|s_2\rangle - |\bar{s}_2\rangle) \in P^{N_1}_- \otimes P^{N-N_1}_-
\end{align}
```
"""

# ╔═╡ 119fb04d-219b-4a08-a50b-584e470f517b
N = 11

# ╔═╡ 6ad0d260-5284-41f2-b7d7-273023d9114c
H = xxzmodel(Symmetric((1 .+ (2*rand(N,N))) .^ 6),-0.7)

# ╔═╡ 942f4e4c-8f21-4a60-ae78-e183e003c748
eig = eigen!(Hermitian(Matrix(H)));

# ╔═╡ a29ef3cf-b4c1-4320-b580-f3cda39e7fac
eig.vectors[:, 5]

# ╔═╡ dd2a8d83-7251-402c-81ea-7b9e52138080
size(eig.vectors,1)

# ╔═╡ 0ea97d53-929e-4683-9742-3183a527abf4
k = 5

# ╔═╡ 4539ba8a-ec44-444d-8ad0-ccc07e7443f5
abs2.(svdvals(reshape(eig.vectors[:, k], 2^div(N-1,2), 2^(N-div(N-1,2)))))

# ╔═╡ f7145441-e14b-47dc-b8c4-c7d243b38c8e
entropy(abs2.(svdvals(reshape(eig.vectors[:, k], 2^div(N-1,2), 2^(N-div(N-1,2))))))

# ╔═╡ 6472d25b-94f7-42ee-8a9d-935b31748eb9
eig.values[:,1]

# ╔═╡ 9bf89fcc-81f9-40c2-9fa5-ddd2eea47941
function halfchain_entropy(state, N=nspins(state), N1=div(N,2))
	entropy(abs2.(svdvals(reshape(state, 2^N1, 2^(N-N1)))))
end

# ╔═╡ d9b06c60-56be-4128-9422-fbfcda7f53e8
plot(halfchain_entropy.(eachcol(eig.vectors)))

# ╔═╡ 2bf00bc0-1d51-4134-afe3-9d5cf897f495
halfchain_entropy.(eachcol(eig.vectors))

# ╔═╡ 0f9e0ab5-7c6e-4e7a-9e58-2185c0b55e4f
halfchain_entropy(eig.vectors[:, k])

# ╔═╡ 1274618f-d59e-4af7-ad6a-81e9c4b1e8bd
begin
	function partial_trace(vec)
		N = round(Int, log2(length(vec)))
		Nhalf = div(N,2)
		ρ = vec * vec'
		tmp = reshape(ρ, 2^Nhalf,2^(N-Nhalf), 2^Nhalf,2^(N-Nhalf))
		res = zeros(eltype(ρ), 2^Nhalf, 2^Nhalf)
		for i in 1:2^(N-Nhalf)
			res += tmp[:,i,:,i]
		end
		Hermitian(res)
	end
end

# ╔═╡ 6910adba-dea2-4095-b888-12bf5a259fd5
reverse(eigvals(partial_trace(eig.vectors[:,k])))

# ╔═╡ efc5f35f-bb74-40ab-b335-b91a1ac7b6b6
NB = N

# ╔═╡ bf58d3eb-46fe-41b8-819a-1827bb29ab51
kB = div(N-1,2)

# ╔═╡ 18820e48-e53b-4d06-ac56-a2a958737156
N1 = 3

# ╔═╡ ae78b8ca-2a42-401d-ac1e-1a0507d015a1
symm_basis = ZBlockBasis(N,kB)

# ╔═╡ 1a6fdac3-7f6f-4c7b-b379-977268b54707
ψsymm = normalize!(ones(basissize(symm_basis)));

# ╔═╡ ae734d90-65cf-4e37-9c6c-aa725d31131c
ψfull = Vector(sparsevec(SpinSymmetry._indices(symm_basis), ψsymm, 2^NB));

# ╔═╡ 44a313c0-3c1b-4769-9176-81acbeffb017
entropy(abs2.(svdvals(reshape(ψfull, 2^N1, 2^(NB-N1)))))

# ╔═╡ 1957e200-71fe-45e3-aee5-39b8eecba31a
SZ = sum(op_list(σz, NB))

# ╔═╡ fff4c13f-003c-49f8-968a-a11a984ab8b0
diag(SZ[SpinSymmetry._indices(symm_basis), SpinSymmetry._indices(symm_basis)])

# ╔═╡ 70d4e7b4-4149-4094-ae75-9c6e7eb097ea
SpinSymmetry._indices(symm_basis)

# ╔═╡ 90000a7c-e927-4de3-b7ab-f238bd5f7af5
indexLookup = let inds = SpinSymmetry._indices(symm_basis)
	sparsevec(inds, 1:length(inds))
end

# ╔═╡ c1aac45a-ca0a-4188-af09-b082ae14cd7c
let res = []
	for k in reverse(0:min(N1,kB))
		let b1 = SpinSymmetry._indices(ZBlockBasis(N1,k)),
			b2 = SpinSymmetry._indices(ZBlockBasis(NB-N1, kB-k))
			push!(res, vec(b1 .+  2^N1 * (b2 .- 1)'))
		end
	end
	vcat(res...)
end

# ╔═╡ b121dd76-ecde-415b-8d29-e92c452b305b
let inds = SpinSymmetry._indices(symm_basis),
	v = Dict((i, ind) for (ind, i) in enumerate(inds))
	@btime $v[118]
end

# ╔═╡ db4667e3-0ed5-430b-b9e2-b245608d8d7b
v = Dict((i, ind) for (ind, i) in enumerate(SpinSymmetry._indices(symm_basis)))

# ╔═╡ 5db25eaa-8fe0-45ed-9029-a42ebc968c63
let v = Vector(indexLookup),
	k = [56,111], res = zeros(Int, length(k))
	@btime $res .= getindex.(Ref($v), $k)
end

# ╔═╡ b9928ff9-8c98-4498-b1d3-3597357a1ffd
let v = Dict((i, ind) for (ind, i) in enumerate(SpinSymmetry._indices(symm_basis))),
	k = [56,111], res = zeros(Int, length(k))
	@btime $res .= getindex.(Ref($v), $k)
end

# ╔═╡ 1e976790-2e7c-4e43-a222-a1e0d13c4f7f
let v = indexLookup,
	k = [56,111], res = zeros(Int, length(k))
	@btime $res .= getindex.(Ref($v), $k)
end

# ╔═╡ ea33e44e-65fd-46a7-b077-cfac2a0ef768
allinds =
	[SpinSymmetry._indices(ZBlockBasis(N1,k)) .+
		2^N1 .* (SpinSymmetry._indices(ZBlockBasis(NB-N1, kB-k)) .- 1)'
		 for k in max(0,N1+kB-NB):min(N1,kB)]

# ╔═╡ 3c1b94c5-2279-40bf-8b69-3b4dac8ea490
reducedρ2 = let res = [], indexLookup = Vector(indexLookup)
	for indmat in allinds
		l = size(indmat, 1)
		temp = zeros(eltype(ψsymm), l, l)
		vals = zeros(eltype(ψsymm), l)
		for inds in eachcol(indmat)
			vals .= getindex.(Ref(ψsymm), getindex.(Ref(indexLookup), inds))
			temp = mul!(temp, vals, vals', 1, 1)
		end
		push!(res, temp)
	end
	res
end

# ╔═╡ 5c262eea-0622-4046-a60d-432f24daf651
eigvals.(reducedρ2)

# ╔═╡ 7cc24ef6-a45b-4f87-a17f-91f7b196f5ad
svdvals(ψsymm[getindex.(Ref(indexLookup), allinds[2])]).^2

# ╔═╡ a84a4a13-ba18-4697-a643-817147e597bd
typeof(allinds)

# ╔═╡ 49481125-553b-48df-907b-ef98051cac47
sum(length, allinds)

# ╔═╡ 7f0db04e-5701-4978-bc8c-e98753bbb749
reducedρ = let res = []
	for k in reverse(0:min(N1,kB))
		let b1 = SpinSymmetry._indices(ZBlockBasis(N1,k)),
			b2 = 2^N1 * (SpinSymmetry._indices(ZBlockBasis(NB-N1, kB-k)) .- 1),
			temp = zeros(eltype(ψsymm), length(b2), length(b2))
			for indA in b1
				fullBasisInds = indA .+ b2
				symmBasisInds = indexLookup[fullBasisInds]
				vals = ψsymm[symmBasisInds]
				temp += vals * vals'
			end
			push!(res, temp)
		end
	end
	res
end

# ╔═╡ 2ba4c3c4-ec8a-4791-b940-4251fa9ae145
entropy(vcat(svdvals.(reducedρ)...))

# ╔═╡ 8cfdd892-56c1-49f8-a993-2f426f9ec690
function entanglement_entropy_zblock(state, N, kB, N1)
	(N == N1 || N1 == 0) && return 0
	(N == k  ||  k == 0) && return 0
	reducedρ = []
	for k in reverse(0:min(N1,kB))
		let b1 = SpinSymmetry._indices(ZBlockBasis(N1,k)),
			b2 = 2^N1 * (SpinSymmetry._indices(ZBlockBasis(NB-N1, kB-k)) .- 1),
			temp = zeros(eltype(ψsymm), length(b2), length(b2))
			for indA in b1
				fullBasisInds = indA .+ b2
				symmBasisInds = indexLookup[fullBasisInds]
				vals = state[symmBasisInds]
				temp += vals * vals'
			end
			push!(reducedρ, temp)
		end
	end
	entropy(vcat(eigvals!.(reducedρ)...))
end

# ╔═╡ 4b3461ca-73e6-4c1a-819a-53ab3774a6f4
entanglement_entropy_zblock(ψsymm, N, kB, N1)

# ╔═╡ c1343cbc-37b9-446e-9035-57983934c99b
Hsymm = symmetrize_operator(H, ZBlockBasis(N, kB))

# ╔═╡ 2cf277f4-38f9-43c6-90b4-20c1a0c5a5df
eigsymm = eigen!(Hermitian(Matrix(Hsymm)));

# ╔═╡ 98f5eff9-d71b-4a2d-ab88-a88b9fffa944
plot(entanglement_entropy_zblock.(eachcol(eigsymm.vectors), N, div(N-1,2), kB))

# ╔═╡ 79893553-b8b2-4d90-80f8-3b0d01d5096d

# ╔═╡ Cell order:
# ╠═25cbd399-a199-4469-be1d-f402df522511
# ╠═3f654c7b-ff3b-423c-91ad-8c942ca5ece9
# ╠═fc4f5bbb-cc20-45ed-aa1a-928ae12b1566
# ╠═f13a03a0-24e6-11ec-2825-8123926abe5d
# ╠═119fb04d-219b-4a08-a50b-584e470f517b
# ╠═6ad0d260-5284-41f2-b7d7-273023d9114c
# ╠═942f4e4c-8f21-4a60-ae78-e183e003c748
# ╠═a29ef3cf-b4c1-4320-b580-f3cda39e7fac
# ╠═dd2a8d83-7251-402c-81ea-7b9e52138080
# ╠═0ea97d53-929e-4683-9742-3183a527abf4
# ╠═4539ba8a-ec44-444d-8ad0-ccc07e7443f5
# ╠═f7145441-e14b-47dc-b8c4-c7d243b38c8e
# ╠═d9b06c60-56be-4128-9422-fbfcda7f53e8
# ╠═6472d25b-94f7-42ee-8a9d-935b31748eb9
# ╠═2bf00bc0-1d51-4134-afe3-9d5cf897f495
# ╠═0f9e0ab5-7c6e-4e7a-9e58-2185c0b55e4f
# ╠═9bf89fcc-81f9-40c2-9fa5-ddd2eea47941
# ╠═6910adba-dea2-4095-b888-12bf5a259fd5
# ╠═1274618f-d59e-4af7-ad6a-81e9c4b1e8bd
# ╠═efc5f35f-bb74-40ab-b335-b91a1ac7b6b6
# ╠═bf58d3eb-46fe-41b8-819a-1827bb29ab51
# ╠═18820e48-e53b-4d06-ac56-a2a958737156
# ╠═1a6fdac3-7f6f-4c7b-b379-977268b54707
# ╠═ae734d90-65cf-4e37-9c6c-aa725d31131c
# ╠═44a313c0-3c1b-4769-9176-81acbeffb017
# ╠═2ba4c3c4-ec8a-4791-b940-4251fa9ae145
# ╠═ae78b8ca-2a42-401d-ac1e-1a0507d015a1
# ╠═1957e200-71fe-45e3-aee5-39b8eecba31a
# ╠═fff4c13f-003c-49f8-968a-a11a984ab8b0
# ╠═70d4e7b4-4149-4094-ae75-9c6e7eb097ea
# ╠═90000a7c-e927-4de3-b7ab-f238bd5f7af5
# ╠═c1aac45a-ca0a-4188-af09-b082ae14cd7c
# ╠═3eef0c5c-9ef5-4e0e-ab9b-0721e68f8413
# ╠═b121dd76-ecde-415b-8d29-e92c452b305b
# ╠═db4667e3-0ed5-430b-b9e2-b245608d8d7b
# ╠═5db25eaa-8fe0-45ed-9029-a42ebc968c63
# ╠═b9928ff9-8c98-4498-b1d3-3597357a1ffd
# ╠═1e976790-2e7c-4e43-a222-a1e0d13c4f7f
# ╠═ea33e44e-65fd-46a7-b077-cfac2a0ef768
# ╠═3c1b94c5-2279-40bf-8b69-3b4dac8ea490
# ╠═5c262eea-0622-4046-a60d-432f24daf651
# ╠═7cc24ef6-a45b-4f87-a17f-91f7b196f5ad
# ╠═a84a4a13-ba18-4697-a643-817147e597bd
# ╠═49481125-553b-48df-907b-ef98051cac47
# ╠═7f0db04e-5701-4978-bc8c-e98753bbb749
# ╠═8cfdd892-56c1-49f8-a993-2f426f9ec690
# ╠═4b3461ca-73e6-4c1a-819a-53ab3774a6f4
# ╠═c1343cbc-37b9-446e-9035-57983934c99b
# ╠═2cf277f4-38f9-43c6-90b4-20c1a0c5a5df
# ╠═98f5eff9-d71b-4a2d-ab88-a88b9fffa944
# ╠═79893553-b8b2-4d90-80f8-3b0d01d5096d


function pointmin(_I :: Array{Float64,2})
    if isempty(_I)
        return false
    end
    I = copy(_I)
    Fx=zeros(size(I))
    Fy=zeros(size(I))
    Fz=zeros(size(I))

    J=zeros( tuple([size(I,i)+2 for i = 1:ndims(I)]...) )
    # J[:,:]=maximum(I[:])
    J .= maximum(I[:])

    J[2:end-1,2:end-1]=I
    Ne = [-1 -1; -1  0; -1  1; 0 -1; 0  1; 1 -1;  1  0; 1  1]
    for i=1:maximum(size(Ne))
        In=J[2+Ne[i,1]:end-1+Ne[i,1],2+Ne[i,2]:end-1+Ne[i,2]]
        check = In.<I #TODO Check this is what is intended (In<I)
        I[check]= copy(In[check])
        D=Ne[i,:]
        D=D./sqrt(sum(D.^2))
        Fx[check] .= D[1]
        Fy[check] .= D[2]
    end

    return (Fx,Fy)
end

function pointmin(_I :: Array{Float64,3})
    I = copy(_I)
    Fx=zeros(size(I))
    Fy=zeros(size(I))
    Fz=zeros(size(I))

    J=zeros( tuple([size(I,i)+2 for i = 1:ndims(I)]...) )
    J[:,:]=maximum(I[:])

    J[2:end-1,2:end-1,2:end-1]=I
    Ne=[-1 -1 -1; -1 -1  0; -1 -1  1; -1  0 -1; -1  0  0; -1  0  1; -1  1 -1; -1  1  0; -1  1  1;
    0 -1 -1;  0 -1  0;  0 -1  1;  0  0 -1;            0  0  1;  0  1 -1;  0  1  0;  0  1  1;
        1 -1 -1;  1 -1  0;  1 -1  1;  1  0 -1;  1  0  0;  1  0  1;  1  1 -1;  1  1  0;  1  1  1];

    for i=1:maximum(size(Ne))
        In = J[2+Ne[i,1]:end-1+Ne[i,1],2+Ne[i,2]:end-1+Ne[i,2],2+Ne[i,3]:end-1+Ne[i,3]]
        check = In.<I
        I[check]= In[check]
        D=Ne[i,:]
        D=D./sqrt(sum(D.^2))
        Fx[check]= D[1]
        Fy[check]= D[2]
        Fz[check]= D[3]
    end

    return (Fx,Fy,Fz)
end


@symbol_func function K_LP(cur_reactor::AbstractReactor)
  cur_K = cur_reactor.C_ejima

  cur_K += cur_reactor.l_i / 2

  cur_K += _external_inductance_term(cur_reactor)

  cur_K *= ( 4e-7 * cur_reactor.pi )

  cur_K
end

@symbol_func function _external_inductance_term(cur_reactor::AbstractReactor)
  cur_epsilon = cur_reactor.epsilon

  cur_kappa = kappa_x(cur_reactor)

  cur_numerator = _b_HN(cur_epsilon)

  cur_numerator -= _a_HN(cur_epsilon)

  cur_denominator = cur_kappa

  cur_denominator *= _d_HN(cur_epsilon)

  cur_denominator /= ( 1 - cur_epsilon )

  cur_denominator += 1

  cur_term = cur_numerator

  cur_term /= cur_denominator

  cur_term
end

function _a_HN(cur_epsilon::SymEngine.Basic, use_default::Bool=false)
  use_default && return _a_HN_default(cur_epsilon)
  symbols(:a_HN)
end

function _b_HN(cur_epsilon::SymEngine.Basic, use_default::Bool=false)
  use_default && return _b_HN_default(cur_epsilon)
  symbols(:b_HN)
end

function _d_HN(cur_epsilon::SymEngine.Basic, use_default::Bool=false)
  use_default && return _d_HN_default(cur_epsilon)
  symbols(:d_HN)
end

function _a_HN(cur_epsilon::AbstractFloat)
  _a_HN_default(cur_epsilon)
end

function _b_HN(cur_epsilon::AbstractFloat)
  _b_HN_default(cur_epsilon)
end

function _d_HN(cur_epsilon::AbstractFloat)
  _d_HN_default(cur_epsilon)
end

function _a_HN_default(cur_epsilon::AbstractSymbol)
  cur_a = 2.0

  cur_a += 9.25 * sqrt(cur_epsilon)

  cur_a -= 1.21 * cur_epsilon

  cur_a
end

function _b_HN_default(cur_epsilon::AbstractSymbol)
  cur_b = 1.0

  cur_b += 1.81 * sqrt(cur_epsilon)

  cur_b += 2.05 * cur_epsilon

  cur_b *= log( 8.0 / cur_epsilon )

  cur_b
end

function _d_HN_default(cur_epsilon::AbstractSymbol)
  cur_d = 1.0

  cur_d += 2 * cur_epsilon ^ 4

  cur_d -= 6 * cur_epsilon ^ 5

  cur_d += 3.7 * cur_epsilon ^ 6

  cur_d *= 0.73 * sqrt(cur_epsilon)

  cur_d
end


using Distributions, Plots, Zygote
using LinearAlgebra, Random
using Turing: MvNormal
using StatsBase

# run it for the funnel distribution
include("hmc.jl")
include("funnel.jl")

# define the state
z = HMCState(ones(2), 0.001, 300)
M = z.M
# The potential U and ∂U are given from l and ∇l
# define the kinetic energy K = p^⊤ Minv p / 2 i.e a multivariate Normal
K(p) = p' * (M \ p) / 2
∂K(p) = M \ p
hmcstep!(z, l, ∇l, K, ∂K)

nsamples, burnin = 60000, 10000
Random.seed!(1234)
samples=zeros(2, nsamples-burnin)
for s in 1:1:nsamples
    hmcstep!(z, l, ∇l, K, ∂K)
    if s > burnin
        samples[:,s-burnin] = z.q
    end
end

v_range = range(-50, stop=50, length=500)
x_range = range(-15, stop=15, length=500)
Z = [l([x,y], σ=3.0) for x in x_range, y in x_range]'

heatmap(x_range, x_range, exp.(-Z), color=:deep)
plot!(samples[1,:], samples[2,:], color=:red)

acf1 = autocor(samples[1,1:500:end], 1:20);
pltacf1 = plot(acf1, title="Autocorrelation function", line=:stem);
acf2 = autocor(samples[2,1:1000:end], 1:20);
pltacf2 = plot(acf2, title="Autocorrelation function", line=:stem);
plot(pltacf1, pltacf2, layout=(1,2))


using StatsBase
using DSP

"""
    surroplot(x, s; W = 100, gfs = 8, lfs = 6,
        cx = :black, cs = :red, size = (400, 600); kwargs...)

Plot a time series along with its surrogate realization `s`, and compare the
periodogram and histogram of the two time series.

## Arguments

- **`x`**: the time series for which to generate a surrogate.

- **`s`** is the surrogate realization of the time series.

## Keyword arguments

- **`W`**: Sets the number of windows for binning the periodogram.

- **`gfs`** and **`lfs`**: Fontsizes of the guides and legends, respectively.

- **`cx`** and **`cs`**: Colors of the original and the surrogate time series, respectively.

- **`size`**: A tuple giving the size of the assembled plot.

- **`nbins`**: The number of bins for the histograms.

- **`kwargs`**: Plots.jl keyword arguments that are supplied to the final `plot` call that assembles the subplots.
"""
function surroplot(x, s; W = 100, gfs = 8, lfs = 6,
        cx = :black, cs = :red, size = (500, 600),
        nbins = 50, kwargs...)

    # Time series
    p1 = Plots.plot()
    Plots.plot!(x, label = "", c = cx);
    Plots.plot!(s, label = "", c = cs)

    # Autocorrelation
    p2 = Plots.plot()
    Plots.plot!(autocor(x), label = "", c = cx)
    Plots.plot!(autocor(s), label = "", c = cs)

    # Binned multitaper periodograms
    p, psurr = DSP.mt_pgram(x), DSP.mt_pgram(s)
    p3 = Plots.plot()
    Plots.plot!(interp([x for x in p.freq], p.power, W), label = "", c = cx, yaxis = :log10)
    Plots.plot!(interp([x for x in psurr.freq], psurr.power, W), label = "", c = cs, yaxis = :log10)

    # Histograms
    p4 = Plots.plot()
    Plots.histogram!(x, label = "Original", alpha = 0.5, nbins = nbins, c = cx)
    Plots.histogram!(s, label = "Surrogate", alpha = 0.5, nbins = nbins, c = cs)

    Plots.xlabel!(p1, "Time step")
    Plots.ylabel!(p1, "Value")

    Plots.xlabel!(p2, "Lag")
    Plots.ylabel!(p2, "Autocorrelation")

    Plots.xlabel!(p3, "Binned frequency")
    Plots.ylabel!(p3, "Power")

    Plots.xlabel!(p4, "Binned value")
    Plots.ylabel!(p4, "Frequency")

    l = Plots.@layout [a{0.3h}; b{0.25h}; c{0.25h}; d{0.2h}]
    Plots.plot(p1, p2, p3, p4;
                layout = l,
                guidefont = (gfs, gfs, gfs, gfs),
                legendfont = (lfs, lfs, lfs, lfs),
                size = size,
                kwargs...
                )
end

export surroplot


# Roots of the characteristic equation of an SDE

"""
$SIGNATURES

Split roots into real and complex-valued roots. For the complex-valued roots,
only provides one of a complex conjugate pair. An error is raised if there is
no single complex conjugate
"""
function splitroots(r)
  rw = copy(r)
  rrc = empty(rw); ri = empty(rw)
  while !isempty(rw)
    rc = popfirst!(rw)
    if isreal(rc)
      push!(rrc,rc)
    else
      push!(ri,rc)
      is_conj = rw.==conj(rc)
      if count(is_conj)==0  error("No complex conjugate found for $rc")  end
      rw = rw[.!is_conj]
    end
  end
  rr = real(rrc)
  if 2length(ri)+length(rr)≠length(r)
    error("Order mismatch; complex roots may not occur in conjugate pairs")
  end
  return rr,ri
end

function splitroots(θ::NamedTuple)
  r_re = θ.r_re
  ri = if haskey(θ,:r_im)
    # zip not supported by Zygote; using buffer.
    # r_re_c = r_re[1:length(θ.r_im)]
    # map(x->complex(x...),zip(r_re_c,θ.r_im))
    nc = length(θ.r_im)
    B = zeros(complex(eltype(θ.r_im)),nc)
    ri_buf = Zygote.bufferfrom(B)
    for i in 1:nc
      ri_buf[i] = complex(r_re[i],θ.r_im[i])
    end
    copy(ri_buf)
  else
    (complex(eltype(r_re)))[]
  end
  nc = length(ri)
  rr = r_re[nc+1:end]
  return rr,ri
end

"
$SIGNATURES

Combine roots
"
function combineroots(rr,ri)
  if  !(eltype(rr)<:Real && eltype(ri)<:Complex) error("r_re must be real, and r_im complex")  end
  # ru = vcat(rr,ri,conj(ri)) # Zygote does not support
  nr = length(rr); ni = length(ri)
  ru_buf = Zygote.bufferfrom(zeros(eltype(ri),nr+2ni))
  ru_buf[1:nr] = rr
  ru_buf[nr.+(1:ni)] = ri
  ru_buf[nr.+ni.+(1:ni)] = conj(ri)
  ru = copy(ru_buf)
  r = sortroots(ru)
  return  r
end

sortroots(r) = sort(r;by=x->(real(x), imag(x)))

issorted_roots(r) = issorted(r;by=x->(real(x), imag(x)))

function issorted_rootstuple(θ)
  rr,ri = splitroots(θ)
  return issorted(rr) && issorted_roots(ri)
end

"""
$SIGNATURES

Random roots
"""
function randroots(rng,r,f,N,T)
  rreal,rcomplex = splitroots(r)
  Tr,Tc = eltype(rreal),eltype(rcomplex)
  rreald = map(rreal) do rc
    sdα = normparstd(rc*T)
    s = max(1,min(sdα,10)/T)
    σ = f*s/√N
    rd = uniform_rand(rng,rc,σ)
  end |> Vector{Tr}
  rcomplexd = map(rcomplex) do rc
    # Real component
    x = real(rc)
    sdα = normparstd(x*T)
    s = max(1,min(sdα,10)/T)
    σ = f*s/√N
    rdr = uniform_rand(rng,x,σ)
    # Imaginary component
    y = imag(rc)
    σ = f*1/T/√N
    rdi = uniform_rand(rng,y,σ)
    Complex(rdr,rdi)
  end |> Vector{Tc}
  rd = combineroots(rreald,rcomplexd)
  return rd
end

"""
$FIELDS

Roots representation of a Stochastic Differential Equation `Sde`.
"""
struct SdeRoots{T<:Complex,U} <: ScalarSde{T}
  "Roots"
  r::Vector{T}
  "Driving noise standard deviation"
  σϵ::U
  function SdeRoots(r::Vector{T},σϵ::U) where {T,U<:Real}
    if !(σϵ>0)  error("σϵ must be positive")  end
    new{T,U}(r,complex(σϵ))
  end
end

function Base.show(io::IO,eq::SdeRoots)
  p = order(eq)
  print(io,"SDE($p) roots "*join(rootstr.(eq.r),"; "))
  print(io,@sprintf " | σϵ = %0.3g" eq.σϵ)
end

function Base.show(io::IO,::MIME"text/plain",eq::SdeRoots)
  print(io,"Roots of Scalar Stochastic Differential Equation\n",eq)
end

function SdeRoots(θ::NamedTuple,σϵ)
  rr,ri = splitroots(θ)
  r = combineroots(rr,ri)
  return SdeRoots(r,σϵ)
end

function SdeRoots(r::AbstractVector{T};σy=one(real(T))) where {T}
  σϵ1 = one(real(T))
  g = std(SdeRoots(r,σϵ1))
  σϵ = σy/g
  return SdeRoots(r,σϵ)
end

function Statistics.std(eq::SdeRoots)
  if order(eq)==0 return eq.σϵ end
  mv = convert(VectorSde,eq)
  Σs = statcovz(mv)
  σy = sum(Σs) |> real |> √
  return σy
end

function SdeRoots(eq::Sde{T}) where {T}
  p = Polynomial([eq.a; one(T)])
  r = complex(roots(p))
  return SdeRoots(sortroots(r),eq.σϵ)
end

Polynomials.roots(eq::AbstractSde) = SdeRoots(eq).r

function Base.convert(::Type{VectorSde},eq::SdeRoots)
  r = eq.r
  if length(unique(r))<length(r) error("Roots must be unique") end
  p = order(eq)
  A = Diagonal(r)
  C = ones(eltype(r),1,p)
  U = [rc^(j-1) for j in 1:p, rc in r]
  v = U\[zeros(p-1);1]
  Q = eq.σϵ^2*v*v'
  R = zeros(eltype(r),1,1)
  veq = VectorSde(A,C,Q,R)
  return veq
end

function Base.convert(::Type{AbstractSde{U}},eq::SdeRoots{T}) where {T,U<:Complex}
  r = convert(Vector{U},eq.r)
  σϵ = convert(real(U),eq.σϵ)
  return SdeRoots(r,σϵ)
end

function Statistics.mean(eqs::Vector{<:SdeRoots})
  r = mean(getproperty.(eqs,:r))
  σϵ = mean(getproperty.(eqs,:σϵ))
  return SdeRoots(r,σϵ)
end

"""
Matrix of eigenvectors `U` of the `A` matrix for the SDE roots `eq`
"""
function eigenvecA(eq::SdeRoots)
  r = eq.r
  p = order(eq)
  # Next line not supported by Zygote, using Buffer instead
  # U = [rc^(i-1) for i in 1:p, rc in r]
  U_buf = Zygote.Buffer(r,p,p)
  for i in 1:p, (j,rc) in enumerate(r)
    U_buf[i,j] = rc^(i-1)
  end
  U = copy(U_buf)
  return U
end

function statcovz(eq::SdeRoots,U)
  p = order(eq)
  v = U\[zeros(p-1);1]
  Q = eq.σϵ^2*v*v'
  r = eq.r
  # Next line not supported by Zygote, using Buffer instead
  # Σs = [-Q[i,j]/(r[i]+conj(r[j])) for i in 1:p,j in 1:p]
  Σs_buf = Zygote.Buffer(v,p,p)
  for i in 1:p, j in 1:p
    Σs_buf[i,j] = -Q[i,j]/(r[i]+conj(r[j]))
  end
  Σs = copy(Σs_buf)
  return Σs
end

statcovz(eq::SdeRoots) = statcovz(eq,eigenvecA(eq))

"Complex to real conversion tolerance"
const def_tol = 1e-8

function roots2sdepar(r,tol=0)
  T = real(eltype(r))
  p = coeffs(fromroots(r))[1:end-1]
  a = map(p) do pc
    if !(abs(imag(pc))≤tol)
      error("Imaginary part exceeds tolerance, abs(",imag(pc),") > ",tol)
    end
    T(real(pc))
  end
  return a
end

function Sde(eq::SdeRoots,tol=def_tol)
  a = roots2sdepar(eq.r,tol)
  return Sde(a,eq.σϵ)
end

SdeRoots(eq::SdeRoots) = eq

SdeRoots(m::WhiteNoise{T}) where {T} = SdeRoots(complex(T)[],m.σ)

SdeRoots(proc::SdeProcess) = SdeRoots(proc.eq)

order(eq::SdeRoots) = length(eq.r)

"print roots"
function printroots(r)
  rstr = map(r) do rr
    s = @sprintf "%0.3g" real(rr)
    if !isreal(rr)
      s *= @sprintf " ± 2π%0.3gim" abs(imag(rr))/2π
    end
    s
  end
  w = maximum(length,rstr)
  for (i,rr) in enumerate(r)
    any(conj(rr).==r[1:i-1]) && continue
    rc = rstr[i]
    p = repeat(' ',w-length(rc))
    irc = @sprintf "-1/%0.3g" -1/real(rr)
    if !isreal(rr)
      irc *= @sprintf " ± (2π/%0.3g)im" 2π/abs(imag(rr))
    end
    s = "$rc$p = $irc"
    if real(rr)>0
      printstyled(s,"\n",color=:yellow)
    else
      println(s)
    end
  end
end

"compact root string"
function rootstr(x)
  return if isreal(x)
    @sprintf "%0.3g" real(x)
  else
    @sprintf "%0.3g%+0.3gim" real(x) imag(x)
  end
end

"
$SIGNATURES

Truncate real part of roots to `-R`.
"
function truncate_roots(r,R;warn=true)
  c = 0
  rt = map(r) do x
    a = real(x)
    return if a<-R
      c+=1
      isreal(r) ? -R : Complex(-R,imag(x))
    else
      x
    end
  end
  if c>0 && warn
    printroots(r)
    @warn "Truncated $c roots"
  end
  return rt
end