text:
# This file contain lowerest level-1 function
# this is a patch for CUDA.allowscalar(false)
USE_CUDA = false
export allowcuda
function allowcuda(b::Bool)
    global USE_CUDA = b
end

import Base.convert
convert(::Type{CuArray{T,2}},U::CUDA.CUSOLVER.CuQRPackedQ{T,CuArray{T,2}}) where T<: Number = CuArray(U)
convert(::Type{Array{T,2}}, U::LinearAlgebra.QRCompactWYQ{T,Array{T,2}}) where T<: Number = Array(U)

"""
    see multiply!
"""
function single_tensor_multiply(mpo_tensor::AbstractArray{T,4} where T<:Number,mps_tensor::AbstractArray{T,3} where T<:Number)
    (l,u,r,d)=size(mpo_tensor)
    (a,d,b)=size(mps_tensor)
    if USE_CUDA
        mpo_tensor = CuArray(mpo_tensor)
        mps_tensor = CuArray(mps_tensor)
    end
    temp_site = reshape(mpo_tensor,:,d)*reshape(permutedims(mps_tensor,(2,1,3)),d,:) #lurab
    temp_site = permutedims(reshape(temp_site,(l,u,r,a,b)),(1,4,2,3,5))
    return Array(temp_site)
end

"""
    sweep in a single site
"""
function single_tensor_sweep!(mps::AbstractMPS,site::Int,::LeftNormalization)
    l=mps.bdim[site-1] # left bond dimension
    r=mps.bdim[site]   # current bond dimension
    A=reshape(mps[site],(l*mps.S,r)) # A is a matrix unfolded from the current tensor
    if USE_CUDA
        A = CuArray(A)
    end
    U,R = qr!(A) # here we intent to do QR = A. However there is no BP, so we do SVD instead
    A = nothing # clear cache
    s = norm(R)
    R = R./s # devided by norm
    res = log(s)
    U = convert(typeof(R),U)
    Dnew = size(R)[1]

    U = reshape(U ,l,mps.S,Dnew)
    mps[site] = Array(U)   # U is LinearAlgebra.QRCompactWYQ type,it is not normal array
    temp = reshape(mps[site+1],r,:)
    if USE_CUDA
        temp = CuArray(temp)
    end
    mps[site+1] = Array(reshape(R*temp,:,mps.S,mps.bdim[site+1]))
    mps.bdim[site] = Dnew
    return res
end

function single_tensor_sweep!(mps::AbstractMPS,site::Int,::RightNormalization)
    l = mps.bdim[site-1]
    r = mps.bdim[site]
    A = copy(reshape(mps[site],(l, r*mps.S)))
    if USE_CUDA
        A = CuArray(A)
    end
    A = copy(transpose(A))
    U,R = qr!(A) # UR = A^{T} R^{T}U^{T} = A
    A = nothing
    Dnew = size(R)[1]
    U = convert(typeof(R),U)
    U = convert(typeof(R),transpose(U))
    R = convert(typeof(U),transpose(R))
    mps[site] = Array(reshape(U,(Dnew,mps.S,r)))
    temp = reshape(mps[site-1],mps.bdim[site-2]*mps.S,mps.bdim[site-1])
    if USE_CUDA
        temp = CuArray(temp)
    end
    mps[site-1] = Array(reshape(temp* R ,(mps.bdim[site-2], mps.S, Dnew)))
    mps.bdim[site-1] = Dnew
end

"""
    cut-off a single tensor
"""
function single_tensor_cutoff!(mps::AbstractMPS,site::Int,Dcut::Int,::LeftNormalization;ϵ=0)
    l = mps.bdim[site-1]
    r = mps.bdim[site]
    A = reshape(mps[site],(l, r*mps.S))
    if USE_CUDA
        A = CuArray(A)
    end
    U, S, V = svd!(A)
    A = nothing
    Dnew = min(Dcut, sum(S.>ϵ))
    res = sum(S[Dnew+1:min(l,r*mps.S)])
    V = copy(adjoint(V))[1:Dnew,:]
    V = reshape(V,(Dnew,mps.S,:))
    mps[site] = Array(V)
    temp = reshape(mps[site-1],mps.bdim[site-2]*mps.S,mps.bdim[site-1])
    if USE_CUDA
        temp = CuArray(temp)
    end
    mps[site-1] = Array(reshape(temp* U[:,1:Dnew] *Diagonal(S[1:Dnew]),(mps.bdim[site-2], mps.S, Dnew)))
    mps.bdim[site-1] = Dnew
    return res
end

"""
    find the sigular value spectrum for a single tensor
"""
function single_tensor_spectrum!(mps::AbstractMPS,site::Int;ϵ=1E-13)
    l = mps.bdim[site-1]
    r = mps.bdim[site]
    A = reshape(mps[site],(l, r*mps.S))
    if USE_CUDA
        A = CuArray(A)
    end
    U, S, V = svd!(A)
    A = nothing
    Dnew = min(sum(S.>ϵ))
    V = copy(adjoint(V))[1:Dnew,:]
    V = copy(reshape(V,(Dnew,mps.S,:)))
    mps[site] = Array(V)
    temp = reshape(mps[site-1],mps.bdim[site-2]*mps.S,mps.bdim[site-1])
    if USE_CUDA
        temp = CuArray(temp)
    end
    mps[site-1] = Array(reshape(temp* U[:,1:Dnew] *Diagonal(S[1:Dnew]),(mps.bdim[site-2], mps.S, Dnew)))
    mps.bdim[site-1] = Dnew
    return S
end

function _overlap(mps1::AbstractMPS,mps2::AbstractMPS,::Type{CuArray})
    ϵ = 1E-30
    FloatType = eltype(mps1[1])
    E = CuArray{FloatType}(transpose(reshape(mps1[1],(mps1.S, mps1.bdim[1])))) * CuArray{FloatType}(reshape(mps2[1],(mps2.S, mps2.bdim[1])))
    res = 0.0
    for i=2:mps1.L
        #E = torch.einsum('ab,ade,bdf->ef', (E, mps1[i], mps2[i])) # einsum version of following three lines
        a=size(mps1[i])[1]
        e=size(mps1[i])[3]
        E = transpose(reshape(transpose(E) * reshape(CuArray{FloatType}(mps1[i]),a,:),:,e))
        E = E * reshape(CuArray{FloatType}(mps2[i]),:,mps2.bdim[i])
        s = norm(E)
        res += log(s)
        E = E./s # devided by norm
    end
    temp = sum(E)
    @assert abs(imag(temp))<1E-10 # TODO: we do not calculate non-real overlap. May need complex support
    return real(exp(res)*temp)
end

function _overlap(mps1::AbstractMPS,mps2::AbstractMPS,::Type{Array})
    ϵ = 1E-13
    E = transpose(reshape(mps1[1],(mps1.S, mps1.bdim[1]))) * reshape(mps2[1],(mps2.S, mps2.bdim[1]))
    res = 0.0
    for i=2:mps1.L
        #E = torch.einsum('ab,ade,bdf->ef', (E, mps1[i], mps2[i])) # einsum version of following three lines
        a=size(mps1[i])[1]
        e=size(mps1[i])[3]
        E = transpose(reshape(transpose(E) * reshape(mps1[i],a,:),:,e))
        E = E * reshape(mps2[i],:,mps2.bdim[i])
        s = norm(E)
        res += log(s)
        E = E./s # devided by norm
    end
    temp = sum(E)
    @assert abs(imag(temp))<1E-10 # TODO: we do not calculate non-real overlap. May need complex support
    return real(exp(res)*temp)
end

max_stars_repo_name: qiyang-ustc/JMPS

max_issues_repo_name: qiyang-ustc/JMPS

max_forks_repo_name: qiyang-ustc/JMPS


text:
# Autogenerated wrapper script for Jansson_jll for i686-w64-mingw32
export libjansson

JLLWrappers.@generate_wrapper_header("Jansson")
JLLWrappers.@declare_library_product(libjansson, "libjansson-4.dll")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libjansson,
        "bin\\libjansson-4.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()

ext: jl


text:
# # The Passport Problem

# Let's now apply what we have learnt to solve a real modelling problem.

# The [Passport Index Dataset](https://github.com/ilyankou/passport-index-dataset)
# lists travel visa requirements for 199 countries, in `.csv` format. Our task
# is to find out the minimum number of passports required to visit all
# countries without  requiring a visa.

# ## Data handling

# Read the following file into a dataframe:

const DATA_FILENAME = joinpath(@__DIR__, "data", "passport-index-matrix.csv")

passport_data = nothing  # Replace this with Code

# In this dataset, the first column represents a passport (=from) and each
# remaining column represents a foreign country (=to).

# The values in each cell are as follows:
# * 3 = visa-free travel
# * 2 = eTA is required
# * 1 = visa can be obtained on arrival
# * 0 = visa is required
# * -1 is for all instances where passport and destination are the same

# Extract the list of country names

countries = nothing  # Replace this with code

# Let's check that you found all 199 countries

@assert length(countries) == 199

# The values we are interested in are -1 and 3. Modify the dataframe so that
# the -1 and 3 are `1` (true), and all others are `0` (false).

nothing  # Replace this with code

# The values in the dataframe should now represent:
#  * 1 = no visa required for travel
#  * 0 = visa required for travel

# ## Visualization

# Produce a visualization of the passport data.  For example, you might want to
# plot a histogram of the number of countries that have visa-free travel.

nothing

# ## Modelling

# We can formulate the passport problem in math as follows:
# ```math
# \begin{aligned}
# \min && \sum_{i \in I} x_i \\
# \text{s.t.} && \sum_{i \in I} a_{i,j} \cdot x_i \geq 1 && \forall j \in I \\
# && x_i \in \{0,1\} && \forall i \in I
# \end{aligned}
# ```
# where $x_i$ is a binary variable indicating if we want to select passport $i$,
# and $a_{i,j}$ is 1 if there is visa-free travel between countries $i$ and $j$,
# and 0 otherwise.

# Formulate and solve this problem with JuMP.

nothing

max_stars_repo_name: blegat/jump-training-materials

max_issues_repo_name: blegat/jump-training-materials

max_forks_repo_name: blegat/jump-training-materials


text:
using Dates, LibPQ, DataKnots, DataKnots4Postgres
using DataKnots: Query, Environment, Pipeline, ValueOf, BlockOf,
                 target, lookup, cover, uncover, lift, compose,
                 syntaxof, relabel, assemble, designate, fits,
                 relabel
import DataKnots: translate, lookup, Lift, Label
import Base: show

# So that queries can be used inside a macro without the dollar sign,
# let's create a macro that injects a translation. This permits
# query macros to be created and then reused later.

macro define(expr)
    @assert expr.head === Symbol("=")
    name = Expr(:quote, expr.args[1])
    body = :(translate($__module__, $(Expr(:quote, expr.args[2]))))
    return :(DataKnots.translate(mod::Module, ::Val{$(name)}) = $(body))
end

#
# Sometimes we wish to do dispatch by the type of input elements. For
# example, if the input is already a date we may wish to treat it a
# particular way, otherwise we may want to make it a date first.
#

DispatchByType(tests::Pair{DataType}...) =
    Query(DispatchByType, collect(Pair{DataType}, tests))

function DispatchByType(env::Environment, p::Pipeline,
                        tests::Vector{Pair{DataType}})
    for (typ, query) in tests
        if fits(target(uncover(p)), BlockOf(ValueOf(typ)))
            return assemble(env, p, query)
        end
    end
    error("doesn't match any type: $(syntaxof(target(p)))")
end

# This is a work-around till we have better naming support in the
# PostgreSQL adapter. Basically, we want to be able to use `concept` in
# a query and either access the main concept table in the root context,
# or the primary link to that table from other contexts.

CascadeGet(names...) =
    Query(CascadeGet, names...)

function CascadeGet(env::Environment, p::Pipeline, names...)
    tgt = target(p)
    for name in names
        q = lookup(tgt, name)
        if q != nothing
            return compose(p, cover(q))
        end
    end
    error("cannot find any of $(join(string.(names), ", "))" *
          " at\n$(syntaxof(tgt))")
end

const ObservationPeriod =
    CascadeGet(:observation_period,
               :observation_period_via_fpk_observation_period_person) >>
    Label(:observation_period)

translate(::Module, ::Val{:observation_period}) = ObservationPeriod

const Visit =
    CascadeGet(:visit, :fpk_drug_visit, :visit_occurrence,
               :visit_occurrence_via_fpk_visit_person,
               :visit_occurrence_via_fpk_drug_visit) >>
    Label(:visit)

translate(::Module, ::Val{:visit}) = Visit

const DrugExposure =
    CascadeGet(:drug_exposure,
              :drug_exposure_via_fpk_drug_person) >>
    Label(:drug_exposure)

translate(::Module, ::Val{:drug_exposure}) = DrugExposure

const DrugEra =
    CascadeGet(:drug_era,
              :drug_era_via_fpk_drug_era_person) >>
    Label(:drug_era)

translate(::Module, ::Val{:drug_era}) = DrugEra

const Concept =
    CascadeGet(:concept, :fpk_observation_concept,
               :fpk_visit_concept, :fpk_condition_concept,
               :fpk_device_concept, :fpk_procedure_concept,
               :fpk_drug_era_concept, :fpk_drug_concept) >>
    Label(:concept)

translate(::Module, ::Val{:concept}) = Concept

const SourceConcept =
    CascadeGet(:fpk_observation_concept_s,
               :fpk_visit_concept_s, :fpk_condition_concept_s,
               :fpk_device_concept_s, :fpk_procedure_concept_s,
               :fpk_drug_era_concept_s, :fpk_drug_concept_s) >>
    Label(:source_concept)

translate(::Module, ::Val{:source_concept}) = SourceConcept

const Person =
    CascadeGet(:person, :fpk_observation_person,
               :fpk_visit_person, :fpk_condition_person,
               :fpk_device_person, :fpk_procedure_person,
               :fpk_drug_era_person, :fpk_drug_person) >>
    Label(:person)

translate(::Module, ::Val{:person}) = Person

# Let's also make `condition` work globally and locally to retrieve
# condition_occurrence records by patient. Unfortunately, `Condition`
# is a built-in name of a type, so we won't make a native version.

translate(::Module, ::Val{:condition}) =
              CascadeGet(:condition, :condition_occurrence,
                  :condition_occurrence_via_fpk_condition_person) >>
              Label(:condition)

translate(::Module, ::Val{:days_supply}) =
              Get(:days_supply) .* Dates.Day(1)
              Label(:days_supply)

# Depending upon the type of record, the actual start date field
# has a different name.  To make this usable, let's normalize it.
# For procedures, let's take it as the start and ending date.

const StartDate =
        CascadeGet(:start_date, :observation_period_start_date,
                       :condition_start_date, :drug_era_start_date,
                       :drug_exposure_start_date, :visit_start_date,
                       :device_exposure_start_date, :procedure_date) >>
        Date.(It) >>
        Label(:start_date)

translate(::Module, ::Val{:start_date}) = StartDate

const EndDate =
        CascadeGet(:end_date, :observation_period_end_date,
                       :condition_end_date, :drug_era_end_date,
                       :drug_exposure_end_date, :visit_end_date,
                       :device_exposure_end_date, :procedure_date) >>
        Date.(It) >>
        Label(:end_date)

translate(::Module, ::Val{:end_date}) = EndDate

# Many query operations involve date intervals. We can't use native
# Julia range object since it's a vector, and vectors are lifted to a
# plural value rather than treated as a tuple. That said, we could
# create a custom type, `DateInterval`, lifted to combinators. This
# sort of interval is inclusive of endpoints.

struct DateInterval
    start_date::Date
    end_date::Date
end
Base.show(io::IO, i::DateInterval) =
    print(io, "$(i.start_date) to $(i.end_date)")

DateInterval(start_date::String, end_date::String) =
    DateInterval(Date(start_date), Date(end_date))
translate(mod::Module, ::Val{:date_interval}, args::Tuple{Any, Any}) =
    DateInterval.(translate.(Ref(mod), args)...)

Lift(::Type{DateInterval}) =
    DispatchByType(DateInterval => It,
                   Date => DateInterval.(It, It),
                   String => DateInterval.(Date.(It), Date.(It)),
                   Any => DateInterval.(StartDate, EndDate)) >>
    Label(:date_interval)
translate(::Module, ::Val{:date_interval}) = Lift(DateInterval)

lookup(ity::Type{DateInterval}, name::Symbol) =
    if name in (:start_date, :end_date)
        lift(getfield, name) |> designate(ity, Date)
    end

# We define `includes` to mean that a date falls within a date interval
# inclusively, or that one interval is completely subsumed by another.

includes(period::DateInterval, val::String) =
    includes(period, Date(val))

includes(period::DateInterval, val::Date) =
    period.end_date >= val >= period.start_date

includes(period::DateInterval, val::DateInterval) =
   (val.start_date >= period.start_date) &&
   (period.end_date >= val.end_date)

# we define

and_previous(init::Date, len::Day) =
    DateInterval(init - len, init)
and_subsequent(init::Date, len::Day) =
    DateInterval(init, init + len)
and_previous(di::DateInterval, len::Day) =
    DateInterval(di.start_date - len, di.end_date)
and_subsequent(di::DateInterval, len::Day) =
    DateInterval(di.start_date, di.end_date + len)

"""
    collapse_intervals(intervals, allowance)

This function collapses a vector of intervals based upon an
allowance, such as 180days between the end of a previous interval,
and the start of the next.
"""
function collapse_intervals(intervals::Vector{DateInterval},
                            allowance::Day)
    intervals′ = Vector{DateInterval}()
    c = nothing
    for i in sort(intervals, by=(i -> i.start_date))
        if c === nothing
            c = i
        elseif c.end_date + allowance < i.start_date
            push!(intervals′, c)
            c = i
        elseif c.end_date < i.end_date
            c = DateInterval(c.start_date, i.end_date)
        end
    end
    if c !== nothing
        push!(intervals′, c)
    end
    intervals′
end

CollapseIntervals(A, I) = collapse_intervals.(I >> DateInterval, A)
ReverseCollapse(I, A) = CollapseIntervals(A, I)
CollapseIntervals(A) = DataKnots.Then(ReverseCollapse, (A,))

translate(mod::Module, ::Val{:collapse_intervals},
          args::Tuple{Any, Any}) =
    CollapseIntervals(translate.(Ref(mod), args)...)

translate(mod::Module, ::Val{:collapse_intervals},
          args::Tuple{Any}) =
    CollapseIntervals(translate.(Ref(mod), args)...)

"""
X >> Includes(Y)

This combinator is true if the interval of `Y` is completely included
in the interval for `X`.  That is, if the starting point of `Y` is
greater than or equal to the starting point of `X`, and also if the
ending point of `Y` is less than or equal to the ending point of `X`.
This combinator accepts a `DateInterval` for its arguments, but also
any object that has `StartDate` and `EndDate` defined.
"""
Includes(Y) = includes.(It >> DateInterval, Y >> DateInterval)
translate(mod::Module, ::Val{:includes}, args::Tuple{Any}) =
    Includes(translate.(Ref(mod), args)...)

During(Y) = includes.(Y >> DateInterval, It >> DateInterval)
translate(mod::Module, ::Val{:during}, args::Tuple{Any}) =
    During(translate.(Ref(mod), args)...)

AndPrevious(Y) = and_previous.(It >> DateInterval, Lift(Day, (Y,)))
translate(mod::Module, ::Val{:and_previous}, args::Tuple{Any}) =
    AndPrevious(translate.(Ref(mod), args)...)

AndSubsequent(Y) = and_subsequent.(It >> DateInterval, Lift(Day, (Y,)))
translate(mod::Module, ::Val{:and_subsequent}, args::Tuple{Any}) =
    AndSubsequent(translate.(Ref(mod), args)...)

# Sometimes it's useful to list the concepts ancestors.

AncestorConcept =
    It.concept_ancestor_via_fpk_concept_ancestor_concept_2 >>
    It.fpk_concept_ancestor_concept_1 >>
    Label(:ancestor_concept)

translate(::Module, ::Val{:ancestor_concept}) = AncestorConcept

# This is used to define `HasCode` which checks if the concept
# associated with this entity matches the given codes, or if
# one of its ancestors matches.

AnyOf(Xs...) = Lift(|, (Xs...,))
OneOf(X, Ys...) = AnyOf((X .== Y for Y in Ys)...)
HasCode(vocabulary_id, codes...) =
    (It.vocabulary_id .== vocabulary_id) .&
    OneOf(It.concept_code, (string.(code) for code in codes)...) .&
    .! Exists(It.invalid_reason)
IsCoded(vocabulary_id, codes...) =
    Exists(
        Filter(HasCode(vocabulary_id, codes...) .|
               Exists(AncestorConcept >>
                      Filter(HasCode(vocabulary_id, codes...)))))

translate(mod::Module, ::Val{:hascode}, args::Tuple{Any,Vararg{Any}}) =
    HasCode(translate.(Ref(mod), args)...)
translate(mod::Module, ::Val{:iscoded}, args::Tuple{Any,Vararg{Any}}) =
    IsCoded(translate.(Ref(mod), args)...)

# In macros, which wish to write things like `90days`. For Julia
# this interpreted as "90 * days", hence we just need to make "days"
# be a constant of 1 day.

translate(::Module, ::Val{:days}) = Dates.Day(1)

# Temporary sort since it's not implemented yet, it happens that
# group provides the functionality needed though.

DataKnots.Label(::Nothing) = Query(Label, nothing);
DataKnots.Label(::Environment, p::Pipeline, ::Nothing) =
    relabel(p, nothing)

Sort(X) = Given(:source => It, Group(X) >> It.source >> Label(nothing))

translate(mod::Module, ::Val{:sort}, args::Tuple{Any}) =
    Sort(translate.(Ref(mod), args)...)

# This creates our test database for us.

sp10 = DataKnot(LibPQ.Connection(""))


text:
using Genie.Renderer

jsonview = raw"
Dict(@vars(:root) => Dict(lang => greet for (lang,greet) in @vars(:greetings)))
"

viewfile = mktemp()
write(viewfile[2], jsonview)
close(viewfile[2])

words = Dict(:en => "Hello", :es => "Hola", :pt => "Ola", :it => "Ciao")

@testset "JSON rendering" begin
  @testset "JSON view rendering with vars" begin
    r = json(Renderer.FilePath(viewfile[1]), root = "greetings", greetings = words)

    @test String(r.body) == """{"greetings":{"en":"Hello","it":"Ciao","pt":"Ola","es":"Hola"}}"""
  end;
end;


text:
module Julia_Functions

#import Pkg; Pkg.activate(@__DIR__); Pkg.instantiate();
using TrajectoryOptimization
using RobotDynamics
using StaticArrays, LinearAlgebra
using Rotations
using Altro

function Rotations.getproperty(q::QuatRotation, f::Symbol)
    if f == :w
        getfield(q,:q).s
    elseif f == :x
        getfield(q,:q).v1
    elseif f == :y
        getfield(q,:q).v2
    elseif f == :z
        getfield(q,:q).v3
    else
        getfield(q,f)
    end
end

struct Quadrotor{R} <: RigidBody{R}
    n::Int
    m::Int
    mass::Float64
    J::Diagonal{Float64,SVector{3,Float64}}
    Jinv::Diagonal{Float64,SVector{3,Float64}}
    gravity::SVector{3,Float64}
    motor_dist::Float64
    kf::Float64
    km::Float64
    bodyframe::Bool  # velocity in body frame?
    info::Dict{Symbol,Any}
end

function Quadrotor{R}(;
        mass=0.5,
        J=Diagonal(@SVector [0.0023, 0.0023, 0.004]),
        gravity=SVector(0,0,-9.81),
        motor_dist=0.1750,
        kf=1.0,
        km=0.0245,
        bodyframe=false,
        info=Dict{Symbol,Any}()) where R
    Quadrotor{R}(13,4,mass,J,inv(J),gravity,motor_dist,kf,km,bodyframe,info)
end

(::Type{Quadrotor})(;kwargs...) = Quadrotor{UnitQuaternion{Float64}}(;kwargs...)

RobotDynamics.control_dim(::Quadrotor) = 4

function RobotDynamics.forces(model::Quadrotor, x, u)
    q = orientation(model, x)
    kf = model.kf
    g = model.gravity
    m = model.mass

    # Extract motor speeds
    w1 = u[1]
    w2 = u[2]
    w3 = u[3]
    w4 = u[4]

    # Calculate motor forces
    F1 = max(0,kf*w1);
    F2 = max(0,kf*w2);
    F3 = max(0,kf*w3);
    F4 = max(0,kf*w4);
    F = @SVector [0., 0., F1+F2+F3+F4] #total rotor force in body frame

    m*g + q*F # forces in world frame
end

function RobotDynamics.moments(model::Quadrotor, x, u)

    kf, km = model.kf, model.km
    L = model.motor_dist

    # Extract motor speeds
    w1 = u[1]
    w2 = u[2]
    w3 = u[3]
    w4 = u[4]

    # Calculate motor forces
    F1 = max(0,kf*w1);
    F2 = max(0,kf*w2);
    F3 = max(0,kf*w3);
    F4 = max(0,kf*w4);

    # Calculate motor torques
    M1 = km*w1;
    M2 = km*w2;
    M3 = km*w3;
    M4 = km*w4;
    tau = @SVector [L*(F2-F4), L*(F3-F1), (M1-M2+M3-M4)] #total rotor torque in body frame
end

RobotDynamics.inertia(model::Quadrotor) = model.J
RobotDynamics.inertia_inv(model::Quadrotor) = model.Jinv
RobotDynamics.mass(model::Quadrotor) = model.mass

model = Quadrotor();
n,m = size(model)
N = 51                # number of knot points
tf = 10.0               # total time (sec)
dt = tf/(N-1)          # time step (sec)

x0 = RBState([0,1,1], UnitQuaternion(I), zeros(3), zeros(3))
xf = RBState([-1,-1.5,2], UnitQuaternion(I), zeros(3), zeros(3))

# objective
Q = Diagonal(@SVector fill(0.1, n))
R = Diagonal(@SVector fill(0.01, m))
Qf = Diagonal(@SVector fill(100.0, n))
obj = LQRObjective(Q,R,Qf,xf,N)

# constraints
cons = ConstraintList(n,m,N)
add_constraint!(cons, BoundConstraint(n,m, u_min=zeros(4), u_max=fill(10.0,4)), 1:N-1)
add_constraint!(cons, GoalConstraint(xf, SA[1,2,3]), N)

# problem
prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons);

u0 = @SVector fill(0.5*model.mass/m, m)
U_hover = [copy(u0) for k = 1:N-1]; # initial hovering control trajectory

#U_hover = load("/home/iconlab/iconlab/mcbhatt2/U_data.jld","U");

#print(U_hover)

U_hover = SVector{4, Float64}[[1.367968171814254, 1.3288459328855668, 1.4070908617102398, 1.4462141437387321], [1.9715547374597115, 2.016302740771755, 1.9268066202442293, 1.8820582747773127], [1.101539665993785, 1.1122955948108078, 1.0907837323752416, 1.080027653295336], [0.9785812329649376, 0.9776680989738087, 0.9794941673273266, 0.9804068966558835], [1.1292101850989422, 1.1245215555576424, 1.1338986241835711, 1.1385869472115653], [1.2141981728925115, 1.2087430632875573, 1.2196532176193535, 1.2251082719950694], [1.226494794448998, 1.222232240650811, 1.2307573491970238, 1.2350199424271557], [1.2182360120139846, 1.2157195168623933, 1.2207525124764211, 1.2232690237066424], [1.214498719096447, 1.2134448115516385, 1.2155526298777555, 1.2166065333813074], [1.2167951676792275, 1.2166873064363313, 1.2169030361237105, 1.217010897813051], [1.2208287344905977, 1.2212008525331624, 1.220456627373153, 1.2200845201657018], [1.2241163982776095, 1.2246418140785544, 1.2235909936153182, 1.2230655937401798], [1.2262920336792238, 1.226785803999927, 1.2257982724072862, 1.2253045177675832], [1.2277553631232958, 1.2281406014076521, 1.2273701315238497, 1.2269849066208238], [1.2288768258968328, 1.229143059101066, 1.2286105977251365, 1.2283443756470047], [1.229830961928901, 1.2299994402661116, 1.2296624877273563, 1.2294940188566406], [1.230668586066505, 1.2307693938866808, 1.2305677819495864, 1.2304669824190064], [1.2313980366082589, 1.2314580357481975, 1.2313380409375672, 1.2312780491961117], [1.2320228051606268, 1.232061236531608, 1.2319843770710248, 1.2319459523694354], [1.2325503753817022, 1.2325788336396764, 1.2325219202007534, 1.232493467982806], [1.2329912733626125, 1.2330155537916891, 1.2329669957765534, 1.2329427208266164], [1.2333569958225272, 1.23337927763559, 1.2333347165969604, 1.2333124397545459], [1.2336587517018904, 1.2336793809505775, 1.233638124777673, 1.2336175000356435], [1.2339068483065294, 1.2339255092319097, 1.2338881894482645, 1.2338695326060694], [1.2341104038227906, 1.2341267433878709, 1.2340940660842001, 1.2340777302164134], [1.2342772484796398, 1.2342911188912795, 1.234263379680955, 1.234249512615075], [1.2344139585717606, 1.2344254430182913, 1.2344024755632028, 1.2343909941416025], [1.2345259761025853, 1.234535321268843, 1.2345166322502137, 1.2345072898150597], [1.2346177675532075, 1.2346252976918919, 1.2346102386644313, 1.234602710986793], [1.234692986592227, 1.2346990345520883, 1.2346869398771154, 1.2346808941296725], [1.23475462108564, 1.234759486278137, 1.234749757171373, 1.2347448939628092], [1.2348051172052021, 1.2348090474027353, 1.2348011883124925, 1.234797259890199], [1.2348464806500759, 1.2348496699452716, 1.2348432926093942, 1.2348401049008293], [1.2348803581003898, 1.2348829535649506, 1.2348777636848676, 1.2348751696388238], [1.234908102769214, 1.2349102145457638, 1.2349059916299103, 1.2349038811240203], [1.2349308275938178, 1.2349325396383493, 1.234929115591899, 1.2349274046887955], [1.2349494489003654, 1.23495082951042, 1.2349480676945412, 1.234946688111845], [1.2349647225855094, 1.2349658348385575, 1.2349636093131053, 1.2349624979832279], [1.234977274070631, 1.2349781862101008, 1.2349763610092737, 1.2349754496920227], [1.2349876224980303, 1.2349884170618366, 1.234986827817505, 1.2349860339711924], [1.234996198855409, 1.2349969769418154, 1.2349954220105774, 1.2349946445278943], [1.2350033570264471, 1.2350042307922409, 1.235002485778577, 1.2350016124933072], [1.2350093764610997, 1.2350104396105166, 1.2350083160666505, 1.2350072532572183], [1.2350144558638096, 1.2350157203914138, 1.2350131925719725, 1.2350119280312268], [1.2350186983898652, 1.2350199916848168, 1.2350174053466831, 1.2350161147094239], [1.2350220960852338, 1.2350229116711917, 1.2350212749232432, 1.235020452205432], [1.2350244308944385, 1.2350238154808144, 1.2350250510130711, 1.2350256761268736], [1.2350246193165197, 1.2350210614475667, 1.235028173663761, 1.2350317192830789], [1.2350147241902587, 1.2350064721389382, 1.2350229734259566, 1.2350312370145065], [1.2349449621936592, 1.234955264417305, 1.2349346677445727, 1.2349243694146959]]

initial_controls!(prob, U_hover)
rollout!(prob);

opts = SolverOptions(
    penalty_scaling=5.,
    penalty_initial=2.,
    constraint_tolerance = 1e-2,
    cost_tolerance = 1e-2
    )

solver = ALTROSolver(prob,opts);
solve!(solver)
#println("Cost: ", cost(solver))
#println("Constraint violation: ", max_violation(solver))
#println("Iterations: ", iterations(solver))
#print((solver.stats).status)

function gen_tracking_problem(prob::Problem, N; Qk = 0.1, Rk = 0.01, Qfk = Qk,)
    n,m = size(prob)
    dt = prob.Z[1].dt
    tf = (N-1)*dt

    # Get sub-trajectory
    Z = Traj(prob.Z[1:N])
    x0 = state(Z[1])
    xf = state(Z[N])  # this actually doesn't effect anything

    # Generate a cost that tracks the trajectory
    Q = Diagonal(@SVector fill(Qk, n))
    R = Diagonal(@SVector fill(Rk, m))
    Qf = Diagonal(@SVector fill(Qfk, n))
    obj = TrajectoryOptimization.TrackingObjective(Q, R, Z, Qf=Qf)

    # Use the same constraints, except the Goal constraint
    cons = ConstraintList(n,m,N)
    for (inds, con) in zip(prob.constraints)
        if !(con isa GoalConstraint)
            if inds.stop > N
                inds = inds.start:N-(prob.N - inds.stop)
            end
            length(inds) > 0 && add_constraint!(cons, con, inds)
        end
    end

    prob = Problem(prob.model, obj, xf, tf, x0=x0, constraints=cons,
        integration=TrajectoryOptimization.integration(prob)
    )
    initial_trajectory!(prob, Z)
    return prob
end

function run_MPC_iterate(x_update,t0,k_mpc,prob_mpc,Z_track,X_traj,iters,times,i)
    # Update initial time
    #t0 += dt
    #k_mpc += 1
    global prob_mpc = prob_mpc
    #TrajectoryOptimization.set_initial_time!(prob_mpc, t0)

    global X_traj
    #print(X_traj)
    # Update initial state by using 1st control, and adding some noise
    r = x_update[1:3]
    q = UnitQuaternion(x_update[4],x_update[5],x_update[6],x_update[7], false)
    q_prev = UnitQuaternion(X_traj[i][4],X_traj[i][5],X_traj[i][6],X_traj[i][7],false)

    q_dot = (q-q_prev)/dt

    ang = inv(q_prev)*q_dot

    v = (r - X_traj[i][1:3])/dt
    w = [ang[3,2], ang[1,3], ang[2,1]]
    x0 = RBState(r, q, v, w)

    # Update the initial state after the dynamics are propogated.
    TrajectoryOptimization.set_initial_state!(prob_mpc, x0)
    X_traj[i+1] = x0

    global Z_track = Z_track

    # Update tracking cost
    TrajectoryOptimization.update_trajectory!(prob_mpc.obj, Z_track, k_mpc)

    # Shift the initial trajectory
    RobotDynamics.shift_fill!(prob_mpc.Z)

    #Shift the multipliers and penalties
    Altro.shift_fill!(get_constraints(altro))

    # Solve the updated problem
    Altro.solve!(altro)

    # Log the results and performance
    global iters[i,1] = iterations(altro)

    # ALTRO in ms and ECOS in s, by default
    global times[i,1] = altro.stats.tsolve

end

function run_MPC(prob_mpc, opts_mpc, Z_track,
                            num_iters = length(Z_track) - prob_mpc.N)
    # First, Let's generate the ALTRO problem
    altro = ALTROSolver(prob_mpc, opts_mpc)

    Qk = prob_mpc.obj[1].Q[1]
    Rk = prob_mpc.obj[1].R[1]
    Qfk = prob_mpc.obj[end].Q[1]

    # Solve initial iteration
    Altro.solve!(altro)

    err_traj = zeros(num_iters,2)
    err_x0 = zeros(num_iters,2)
    iters = zeros(Int, num_iters,2)
    times = zeros(num_iters,2)

    # Get the problem state size and control size
    n,m = size(prob_mpc)

    t0 = 0
    k_mpc = 1
    x0 = SVector(prob_mpc.x0)
    X_traj = [copy(x0) for k = 1:num_iters+1]

    # Begin the MPC LOOP
    for i = 1:num_iters
        # Update initial time
        t0 += dt
        k_mpc += 1
        TrajectoryOptimization.set_initial_time!(prob_mpc, t0)

        # Update initial state by using 1st control, and adding some noise
        x0 = discrete_dynamics(TrajectoryOptimization.integration(prob_mpc),
                                    prob_mpc.model, prob_mpc.Z[1])

        # Update the initial state after the dynamics are propogated.
        TrajectoryOptimization.set_initial_state!(prob_mpc, x0)
        X_traj[i+1] = x0

        # Update tracking cost
        TrajectoryOptimization.update_trajectory!(prob_mpc.obj, Z_track, k_mpc)

        # Shift the initial trajectory
        RobotDynamics.shift_fill!(prob_mpc.Z)

        # Shift the multipliers and penalties
        Altro.shift_fill!(get_constraints(altro))

        # Solve the updated problem
        Altro.solve!(altro)

        # Log the results and performance
        iters[i,1] = iterations(altro)

        # ALTRO in ms and ECOS in s, by default
        times[i,1] = altro.stats.tsolve

    end

    # Return the trajectory that was traced out by the MPC
    # X_ecos and U_ecos are needed to grab the ECOS variable information
    # All other logistics are in the dictionary.
    return X_traj, Dict(:time=>times, :iter=>iters)
end

N_mpc = 6

global prob_mpc = gen_tracking_problem(prob, N_mpc)

print(prob_mpc)

opts_mpc = SolverOptions(
    penalty_scaling=5.,
    penalty_initial=2.,
    constraint_tolerance = 1e-2,
    cost_tolerance = 1e-2
    );

#Z_track = load("/home/iconlab/iconlab/mcbhatt2/data.jld2")["Z_track"];
Z_track = TrajectoryOptimization.get_trajectory(solver)
#print(Z_track)

altro = ALTROSolver(prob_mpc, opts_mpc)

num_iters = length(Z_track) - prob_mpc.N

Qk = prob_mpc.obj[1].Q[1]
Rk = prob_mpc.obj[1].R[1]
Qfk = prob_mpc.obj[end].Q[1]

# Solve initial iteration
Altro.solve!(altro)

err_traj = zeros(num_iters,2)
err_x0 = zeros(num_iters,2)
iters = zeros(Int, num_iters,2)
times = zeros(num_iters,2)

# Get the problem state size and control size
n,m = size(prob_mpc)

t0 = 0
k_mpc = 1
x0 = SVector(prob_mpc.x0)
X_traj = [copy(x0) for k = 1:num_iters+1]

integrate = TrajectoryOptimization.integration

set_initial_time = TrajectoryOptimization.set_initial_time!
end

#N_mpc = 6

#prob_mpc = gen_tracking_problem(prob, N_mpc)

#opts_mpc = SolverOptions(
#    penalty_scaling=5.,
#    penalty_initial=2.,
#    constraint_tolerance = 1e-2,
#    cost_tolerance = 1e-2
#    );

#Z_track = TrajectoryOptimization.get_trajectory(solver);

#X_traj, res = run_MPC(prob_mpc, opts_mpc, Z_track);

#using Statistics

#times = mean(res[:time], dims=1)
#print("Average time:", times[1])

max_stars_repo_name: labicon/crazyswarm-labicon

max_issues_repo_name: labicon/crazyswarm-labicon

max_forks_repo_name: labicon/crazyswarm-labicon


text:
using XmlToDict
using Base.Test, DataStructures

# write your own tests here
@test 1 == 1

xmltest = """
<?xml version="1.0" encoding="UTF-8"?>
<bookstore>
  <book category="COOKING" tag="first">
    <title lang="en">Everyday Italian</title>
    <author>Giada De Laurentiis</author>
    <year>2005</year>
    <price>30.00</price>
  </book>
  <book category="CHILDREN">
    <title lang="en">Harry Potter</title>
    <author>J K. Rowling</author>
    <year>2005</year>
    <price>29.99</price>
  </book>
  <newspaper>
    <name>The Times</name>
  </newspaper>
</bookstore>
"""

xt1 = xmlparse_string(xmltest)

@test xt1["bookstore"]["book"][1]["@category"] == "COOKING"
@test xt1["bookstore"]["book"][1]["@tag"] == "first"
@test xt1["bookstore"]["book"][1]["title"]["@lang"] == "en"
@test xt1["bookstore"]["book"][1]["title"]["#text"] == "Everyday Italian"
@test xt1["bookstore"]["book"][1]["author"]["#text"] == "Giada De Laurentiis"
@test xt1["bookstore"]["book"][1]["year"]["#text"] == "2005"
@test xt1["bookstore"]["book"][1]["price"]["#text"] == "30.00"

@test xt1["bookstore"]["book"][2]["@category"] == "CHILDREN"
@test xt1["bookstore"]["book"][2]["title"]["@lang"] == "en"
@test xt1["bookstore"]["book"][2]["title"]["#text"] == "Harry Potter"
@test xt1["bookstore"]["book"][2]["author"]["#text"] == "J K. Rowling"
@test xt1["bookstore"]["book"][2]["year"]["#text"] == "2005"
@test xt1["bookstore"]["book"][2]["price"]["#text"] == "29.99"

println("newspaper -> ", xt1["bookstore"][x2d"newspaper+"])

@test_throws KeyError xt1["bookstore"][x2d"magazines"]
@test xt1["bookstore"][x2d"magazines*"] == []

println("x2d'newspaper+' -> ", xt1["bookstore"][x2d"newspaper+"])
println("x2d'newspaper*' -> ", xt1["bookstore"][x2d"newspaper*"])
println("x2d'newspaper' -> ", [ xt1["bookstore"][x2d"newspaper"] ] )

@test isa(xt1["bookstore"]["newspaper"], Array) == false
@test xt1["bookstore"][x2d"newspaper+"] == [ xt1["bookstore"]["newspaper"] ]
@test xt1["bookstore"][x2d"newspaper*"] == [ xt1["bookstore"]["newspaper"] ]


text:
export
    Path,
    samplepath,
    quadpathpoly,
    pathmean,
    esspath

"""
    AllowedTimeType

Syntactic sugar for union type of Vector{Real} and LinSpace{Real} (types
accepted for the `samplelocalpath` function).
"""
const AllowedTimeType = Union{Vector{<:Real}, AbstractRange{<:Real}}

"""
    Path

Type to store a path simulated via PDSampling. It stores the corners and the
times.
"""
mutable struct Path
    xs::Matrix{Real}  # corners (dim: nfeatures * nsteps)
    ts::Vector{Real}  # times (dim: nsteps)
    # -- implicit
    p::Int      # nfeatures
    nseg::Int   # number of segments
    Path(xs, ts) = new(xs, ts, size(xs, 1), length(ts) - 1)
end

"""
    Segment

Type to store the information relative to a single segment within a path in the
global BPS.
"""
struct Segment
    ta::Real
    tb::Real
    xa::Vector{Real}
    xb::Vector{Real}
    # -- implicit
    tau::Real
    v::Vector{Real}
    function Segment(ta, tb, xa, xb)
        @assert tb > ta "times of subsequent events should be distinct"
        new(ta, tb, xa, xb, tb - ta, (xb - xa) / (tb - ta))
    end
end

"""
    getsegment(p, j)

Retrieves segment j where j=1 corresponds to the first segment from the initial
time to the first corner.
"""
getsegment(p::Path, j::Int) =
    Segment(p.ts[j], p.ts[j+1], p.xs[:,j], p.xs[:,j+1])

"""
    samplepath(p, t)

Sample the piecewise linear trajectory defined by the corners `xs` and the
times `ts` at given time `t`. The matrix returned has dimensions `p*length(t)`.
"""
function samplepath(p::Path, t::AllowedTimeType)
    @assert t[1] >= p.ts[1] && t[end] <= p.ts[end]
    samples = zeros(p.p, length(t))
    #
    ti = t[1]
    j = searchsortedlast(p.ts, ti)
    seg = getsegment(p, j)
    i = 1
    while i <= length(t)
        ti = t[i]
        if ti <= seg.tb
            samples[:,i] = seg.xa + (ti - seg.ta) * seg.v
            # increment
            i += 1
        else
            j += searchsortedlast(p.ts[j+1:end], ti)
            # safeguard for out of bound: point to last segment
            j = (j+1 >= length(p.ts)) ? length(p.ts) - 1 : j
            seg = getsegment(p,j)
        end
    end
    samples
end
samplepath(p::Path, t::Real) = samplepath(p, [t])[:, 1]

"""
    quadpathpoly(path, poly)

Integrate a polynomial on the elements along a path. So if we are sampling in Rd
and have phi(X)=Poly(xk), then we can integrate the polynomial along the path.
Note that we can't do the same thing if phi(X) mixes 2 or more components.
"""
function quadpathpoly(path::Path, pol::Poly)
    dim = length(path.xs[:,1])
    res = zeros(dim)
    nseg = length(path.ts)-1
    T = path.ts[end]
    for segidx = 1:nseg-1
        # for each segment in the path
        segment = getsegment(path, segidx)
        # get the characteristics
        tau = segment.tau
        xa = segment.xa
        v = segment.v
        # integrate along that segment
        for d = 1:dim
            res[d] += polyint(polyval(pol, Poly([xa[d], v[d]])))(tau)
        end
    end
    # last segment
    lastsegment = getsegment(path, nseg)
    # characteristics
    tau = T - lastsegment.ta
    xa = lastsegment.xa
    v = lastsegment.v
    # integrate along that segment
    for d = 1:dim
        res[d] += polyint(polyval(pol, Poly([xa[d], v[d]])))(tau)
    end
    res/T
end
pathmean(path::Path) = quadpathpoly(path, Poly([0.0, 1.0]))

"""
    esspath(path, ns, rtol, limitfrac)

Increase the number of sample points on the path. The hope is that the ESS for
each dimension divided by the number of samples becomes larger than `limfrac`.
"""
function esspath(path::Path; ns::Int=1000, rtol::Real=0.1,
                 limfrac::Real=0.2, maxns::Int=100000)

    Tp = 0.999 * path.ts[end]
    gg = range(0, stop=Tp, length=ns)

    spath  = samplepath(path, gg)
    old_ess = [ess(spath[i, :]) for i in 1:path.p]
    flag = any(old_ess ./ ns .> limfrac)
    cur_ess = flag ? similar(old_ess) : old_ess

    while flag && (ns < maxns / 2)
        ns *= 2
        gg = range(0, stop=Tp, length=ns)
        spath = samplepath(path, gg)
        cur_ess = [ess(spath[i, :]) for i in 1:path.p]

        # cond1 : check whether the ESS/NS is big enough
        # cond2 : check whether the ESS changed significantly
        flag   = any(cur_ess ./ ns .> limfrac) ||
                 any((abs.(cur_ess - old_ess) ./ old_ess) .> rtol)
        old_ess = cur_ess
    end
    (cur_ess, ns)
end

"""
    ess(v)

Compute the ESS of a vector of samples. This code is adapted from Klara.jl.
"""
function ess(v::AbstractVector{T}) where T <: Real
    length_v = length(v)
    mcvar_iid = var(v) / length_v

    # = IMSE ============================
    maxlag = length_v - 1
    k = Int(floor((maxlag - 1) / 2))
    m = k + 1
    g = zeros(T, m)
    # empirical autocovariance
    acv = autocov(v, 0:maxlag)
    # calculate Ĝ_{n, m} from the acv
    for j ∈ 0:k
        g[j+1] = acv[2j + 1] + acv[2j + 2]
        if g[j+1] <= 0
            m = j
            break
        end
    end
    # create monotone sequence of g
    if m > 1
        for j ∈ 2:m
            if g[j] > g[j - 1]
                g[j] = g[j - 1]
            end
        end
    end
    mcvar_imse = (-acv[1] + 2sum(g[1:m])) / length_v
    # ===================================

    return length_v * mcvar_iid / mcvar_imse
end


text:
module CovarianceMatrices

# using Requires
#
# function __init__()
#     @info("INIT CALLED")
#     @require GLM="38e38edf-8417-5370-95a0-9cbb8c7f171a" include("glm.jl")
# end

using LinearAlgebra
using Statistics

#include("varhac.jl")
include("types.jl")
include("HAC.jl")
include("HC.jl")
include("CRHC.jl")
include("covariance.jl")
include("methods.jl")
include("glm.jl")

export QuadraticSpectralKernel, TruncatedKernel, ParzenKernel, BartlettKernel,
       TukeyHanningKernel, VARHAC, HC0, HC1, HC2, HC3, HC4, HC4m, HC5, CRHC0, CRHC1,
       CRHC2, CRHC3, NeweyWest, Andrews,
       HACCache, HCCache, CRHCCache, CovarianceMatrix, vcov, stderr

end # module


text:
function call_limit_stop_check(status, information, options)
    cond =  status.f_calls >= options.f_calls_limit
    options.debug && cond && @info("Stopped since call_limit was met.")
    cond
end

function iteration_stop_check(status, information, options)
    cond =  status.iteration >= options.iterations
    options.debug && cond && @info("Stopped since iteration limit was met.")
    cond
end

function time_stop_check(status, information, options)
    cond =  status.overall_time >= options.time_limit
    options.debug && cond && @info("Stopped since time limit was met.")
    cond
end

function accuracy_stop_check(status::State{xf_indiv}, information, options)
    cond =  !isnan(information.f_optimum) && abs(status.best_sol.f - information.f_optimum) < options.f_tol
    options.debug && cond && @info("Stopped since accuracy was met.")
    cond
end

function accuracy_stop_check(status::State{xfgh_indiv}, information, options)
    if isnan(information.f_optimum)
        return false
    end

    vio = violationsSum(status.best_sol.g, status.best_sol.h, ε=options.h_tol)
    cond =  vio ≈ 0.0 && abs(status.best_sol.f - information.f_optimum) < options.f_tol

    options.debug && cond && @info("Stopped since accuracy and feasibility was met.")
    cond
end

function accuracy_stop_check(status::State{xFgh_indiv}, information, options)
    return false
end

function var_stop_check(status, information, options)
    if typeof(status.best_sol.f) <: Array
        return false
    end
    cond =  var(map( s->s.f, status.population )) ≈ 0.0
    options.debug && cond && @info("Stopped since varF was met.")
    cond
end

function diversity_stop_check(status, information, options)
    if typeof(status.best_sol.f) <: Array
        return false
    end

    cond =  var(map( s->s.f, status.population )) ≈ 0.0
    options.debug && cond && @info("Stopped since varF was met.")
    cond
end

function stop_check(status, information, options)
    accuracy_stop_check(status, information, options) ||
    var_stop_check(status, information, options)
end


text:
using DataFrames: DataFrame

struct ROOTDirectory
    name::AbstractString
    header::ROOTDirectoryHeader
    keys::Vector{TKey}
end

struct ROOTFile
    filename::AbstractString
    format_version::Int32
    header::FileHeader
    fobj::IOStream
    tkey::TKey
    streamers::Streamers
    directory::ROOTDirectory
    branch_cache::Dict{String, TBranch}
end

function ROOTFile(filename::AbstractString)
    fobj = Base.open(filename)
    preamble = unpack(fobj, FilePreamble)
    String(preamble.identifier) == "root" || error("Not a ROOT file!")
    format_version = preamble.fVersion

    if format_version < 1000000
        @debug "32bit ROOT file"
        header = unpack(fobj, FileHeader32)
    else
        @debug "64bit ROOT file"
        header = unpack(fobj, FileHeader64)
    end

    # Streamers
    if header.fSeekInfo != 0
        @debug "Reading streamer info."
        seek(fobj, header.fSeekInfo)
        streamers = Streamers(fobj)
    else
        @debug "No streamer info present, skipping."
    end

    seek(fobj, header.fBEGIN)
    tkey = unpack(fobj, TKey)

    # Reading the header key for the top ROOT directory
    seek(fobj, header.fBEGIN + header.fNbytesName)
    dir_header = unpack(fobj, ROOTDirectoryHeader)
    if dir_header.fSeekKeys == 0
        ROOTFile(format_version, header, fobj, tkey, [])
    end

    seek(fobj, dir_header.fSeekKeys)
    header_key = unpack(fobj, TKey)

    n_keys = readtype(fobj, Int32)
    keys = [unpack(fobj, TKey) for _ in 1:n_keys]

    directory = ROOTDirectory(tkey.fName, dir_header, keys)

    ROOTFile(filename, format_version, header, fobj, tkey, streamers, directory, Dict())
end

function Base.show(io::IO, f::ROOTFile)
    n_entries = length(f.directory.keys)
    entries_suffix = n_entries == 1 ? "entry" : "entries"
    n_streamers = length(f.streamers)
    streamers_suffix = n_streamers == 1 ? "streamer" : "streamers"
    print(io, typeof(f))
    print(io, "(\"$(f.filename)\") with $n_entries $entries_suffix ")
    print(io, "and $n_streamers $streamers_suffix.")
end

UUID(f::ROOTFile) = f.header.fUUID

function streamerfor(f::ROOTFile, name::AbstractString)
    for e in f.streamers.elements
        if e.streamer.fName == name
            return e
        end
    end
    error("No streamer found for $name.")
end

function Base.getindex(f::ROOTFile, s::AbstractString)
    if '/' ∈ s
        @debug "Splitting path '$s' and getting items recursively"
        paths = split(s, '/')
        return f[first(paths)][join(paths[2:end], "/")]
    end
    tkey = f.directory.keys[findfirst(isequal(s), keys(f))]
    @debug "Retrieving $s ('$(tkey.fClassName)')"
    streamer = getfield(@__MODULE__, Symbol(tkey.fClassName))
    streamer(f.fobj, tkey, f.streamers.refs)
end

function Base.keys(f::ROOTFile)
    keys(f.directory)
end

function Base.keys(d::ROOTDirectory)
    [key.fName for key in d.keys]
end

function Base.keys(b::TBranchElement)
    [branch.fName for branch in b.fBranches.elements]
end

function Base.get(f::ROOTFile, k::TKey)
end

Base.keys(t::TTree) = [b.fName for b in t.fBranches.elements]

function Base.getindex(t::T, s::AbstractString) where {T<:Union{TTree, TBranchElement}}
    if '/' ∈ s
        @debug "Splitting path '$s' and getting branches recursively"
        paths = split(s, '/')
        return t[first(paths)][join(paths[2:end], "/")]
    end
    @debug "Searching for branch '$s' in $(length(t.fBranches.elements)) branches."
    for branch in t.fBranches.elements
        @debug branch.fName
        if branch.fName == s
            return branch
        end
    end
    missing
end

"""
    function array(f::ROOTFile, path)

Reads an array from a branch. Currently hardcoded to Int32
"""
function array(f::ROOTFile, path; raw=false)
    if path ∈ keys(f.branch_cache)
        branch = f.branch_cache[path]
    else
        branch = f[path]
        if ismissing(branch)
            error("No branch found at $path")
        end
    end

    if !raw && length(branch.fLeaves.elements) > 1
        error(
            "Branches with multiple leaves are not supported yet. Try reading with `array(...; raw=true)`.",
        )
    end

    leaf = first(branch.fLeaves.elements)
    rawdata, rawoffsets = readbasketsraw(f.fobj, branch)
    # https://github.com/scikit-hep/uproot3/blob/54f5151fb7c686c3a161fbe44b9f299e482f346b/uproot3/interp/auto.py#L144
    isjagged = (match(r"\[.*\]", leaf.fTitle) !== nothing)

    # there are two possibility, one is the leaf is just normal leaf but the title has "[...]" in it
    # magic offsets, seems to be common for a lot of types, see auto.py in uproot3
    # only needs when the jaggedness comes from TLeafElements, not needed when
    # the jaggedness comes from having "[]" in TLeaf's title
    jagg_offset = leaf isa TLeafElement ? 10 : 0
    if raw
        return rawdata, rawoffsets
    end
    # the other is where we need to auto detector T bsaed on class name
    if isjagged || !iszero(jagg_offset) # non-primitive jagged leaf
        T = autointerp_T(branch, leaf)

        # for each "event", the index range is `offsets[i] + jagg_offset + 1` to `offsets[i+1]`
        # this is why we need to append `rawoffsets` in the `readbasketsraw()` call
        # when you use this range to index `rawdata`, you will get raw bytes belong to each event
        # Say your real data is Int32 and you see 8 bytes after indexing, then this event has [num1, num2] as real data
        @views [
                ntoh.(reinterpret(
                                  T, rawdata[ (rawoffsets[i]+jagg_offset+1):rawoffsets[i+1] ]
                                 )) for i in 1:(length(rawoffsets) - 1)
               ]
    else # the branch is not jagged
        return ntoh.(reinterpret(primitivetype(leaf), rawdata))
    end
end

function autointerp_T(branch, leaf)
    if hasproperty(branch, :fClassName)
        classname = branch.fClassName # the C++ class name, such as "vector<int>"
        m = match(r"vector<(.*)>", classname)
        m===nothing && error("Cannot understand fClassName: $classname.")
        elname = m[1]
        elname = endswith(elname, "_t") ? lowercase(chop(elname; tail=2)) : elname  # Double_t -> double
        try
            elname == "bool" && return Bool #Cbool doesn't exist
            getfield(Base, Symbol(:C, elname))
        catch
            error("Cannot convert element of $elname to a native Julia type")
        end
    else
        primitivetype(leaf)
    end

end

"""
    function DataFrame(f::ROOTFile, path)

Reads a tree into a dataframe
"""
function DataFrame(f::ROOTFile, path)
    names = keys(f[path])
    cols = [array(f, path * "/" * n) for n in names]
    DataFrame(cols, names, copycols=false) #avoid double allocation
end

"""
    splitup(data::Vector{UInt8}, offsets, T::Type; skipbytes=0, primitive=false)

Given the `offsets` and `data` return by `array(...; raw = true)`, reconstructed the actual
array (can be jagged, or with custome struct).
"""
function splitup(data::Vector{UInt8}, offsets, T::Type; skipbytes=0, primitive=false)
    elsize = sizeof(T)
    out = sizehint!(Vector{Vector{T}}(), length(offsets))
    lengths = diff(offsets)
    push!(lengths, length(data) - offsets[end] + offsets[1])  # yay ;)
    io = IOBuffer(data)
    for (idx, l) in enumerate(lengths)
        # println("$idx / $(length(lengths))")
        if primitive
            error("primitive interpretation is buggy")
            push!(out, reinterpret(T, data[skipbytes+1:skipbytes+Int32((l - skipbytes))]))
        else
            skip(io, skipbytes)
            n = (l - skipbytes) / elsize
            push!(out, [readtype(io, T) for _ in 1:n])
        end
    end
    out
end

function readbasketsraw(io, branch)
    seeks = branch.fBasketSeek
    bytes = branch.fBasketBytes

    total_entries = branch.fEntries
    # Just to check if we have a jagged structure
    # streamer = streamerfor()

    max_len = sum(bytes)
    data = sizehint!(Vector{UInt8}(), max_len)
    offsets = sizehint!(Vector{Int32}(), total_entries+1) # this is always Int32
    idx = 1
    for (basket_seek, n_bytes) in zip(seeks, bytes)
        @debug "Reading raw basket data" basket_seek n_bytes
        basket_seek == 0 && break
        seek(io, basket_seek)
        idx += readbasketbytes!(data, offsets, io, idx)
    end
    data, offsets
end

# Thanks Jim and Philippe
# https://groups.google.com/forum/#!topic/polyglot-root-io/yeC0mAizQcA
# The offsets start at fKeylen - fLast + 4. A singe basket of data looks like this:
#                                           4 bytes          4 bytes
# ┌─────────┬────────────────────────────────┬───┬────────────┬───┐
# │ TKey    │ content                        │ X │ offsets    │ x │
# └─────────┴────────────────────────────────┴───┴────────────┴───┘
#           │←        fLast - fKeylen       →│                    │
#           │                                                     │
#           │←                       fObjlen                     →│
#
function readbasketbytes!(data, offsets, io, idx)
    basketkey = unpack(io, TBasketKey)

    s = datastream(io, basketkey)  # position(s) == 0, but offsets start at -basketkey.fKeylen
    start = position(s)
    contentsize = basketkey.fLast - basketkey.fKeylen
    offsetbytesize = basketkey.fObjlen - contentsize - 8
    offset_len = offsetbytesize ÷ 4 # these are always Int32

    if offsetbytesize > 0
        @debug "Offset data present" offsetbytesize
        skip(s, contentsize)
        skip(s, 4) # a flag that indicates the type of data that follows
        readoffsets!(offsets, s, offset_len, length(data), -basketkey.fKeylen)
        skip(s, 4)  # "Pointer-to/location-of last used byte in basket"
        seek(s, start)
    end

    @debug "Reading $(contentsize) bytes"
    readbytes!(s, data, idx, contentsize)
    # offsets starts at -fKeylen, same as the `local_offset` we pass in in the loop
    push!(offsets, basketkey.fLast - basketkey.fKeylen)

    contentsize
end

function readoffsets!(out, s, contentsize, global_offset, local_offset)
    for _ in 1:contentsize
        offset = readtype(s, Int32) + global_offset + local_offset
        push!(out, offset)
    end
end

"""
    function readbytes!(io, b, offset, nr)

Efficient read of bytes into an existing array at a given offset
"""
function readbytes!(io, b, offset, nr)
    resize!(b, offset + nr - 1)
    nb = UInt(nr)
    # GC.@preserve b unsafe_read(io, pointer(b, offset), nb)
    unsafe_read(io, pointer(b, offset), nb)
    nothing
end


text:
# power analysis functions for F2

"""
    powerF2Test(u::Real = 0, v::Real = 0, f2::Float64 = 0.0, alpha::Float64 = 0.05)

Compute power for the general linear model for a sample with `u` degress of
freedom for numerator and `v` degrees of freedom for denominator to test the
effect size `f2` at type I error = `alpha` (default: 0.05).
"""
function powerF2Test(;
    u::Real = 0,
    v::Real = 0,
    f2::Real = 0.0,
    alpha::Real = 0.05)

    check_args(u=u, v=v, f2=f2, alpha=alpha)

    λ = f2*(u + v +1)
    return ccdf(NoncentralF(u,v,λ),cquantile(FDist(u,v),alpha))
end

"""
    samplesizeF2Test(u::Real = 0, v::Real = 0, f2::Float64 = 0.0, alpha::Real = 0.05)

Compute power for the general linear model for a sample with `u` degress of
freedom for numerator and `v` degrees of freedom for denominator to test the
effect size `f2` at type I error = `alpha` (default: 0.05).
"""
function samplesizeF2Test(;
    u::Real = 0,
    v::Real = 0,
    f2::Real = 0.0,
    alpha::Real = 0.05,
    power::Real = 0.0)

    check_args(u=u, v=v, f2=f2, alpha=alpha, power=power)

    if u == 0
        return ceil(Int64,fzero(x->powerF2Test(u = x, v = v, f2 = f2, alpha = alpha) - power, 2.0, 10.0^7))
    end
    return ceil(Int64,fzero(x->powerF2Test(u = u, v = x, f2 = f2, alpha = alpha) - power, 2.0, 10.0^7))
end

"""
    effectsizeF2Test(u::Real = 0, v::Real = 0, alpha::Float64 = 0.05, power::Float64 = 0.8)

Compute the effect size for the general linear model that a sample with `u` degress of
freedom for numerator and `v` degrees of freedom for denominator can test with
power > `power` (default: 0.8) at type I error = `alpha` (default: 0.05).
"""
function effectsizeF2Test(;
    u::Real = 0,
    v::Real = 0,
    alpha::Real = 0.05,
    power::Real = 0.0)

    check_args(u=u,v=v,alpha=alpha,power=power)

    return fzero(x->powerF2Test(u = u, v = v, f2 = x, alpha = alpha) - power, 1e-07, 1e+07)
end

"""
    alphaF2Test(u::Real = 0, v::Real = 0, f2::Float64 = 0.0, power::Float64 = 0.8)

Compute the probability of type I error that the general linear model with
a sample with `u` degress of freedom for numerator and `v` degrees of freedom
for denominator has in testin the effect size `f2` with power > `power` (default: 0.8).
"""
function alphaF2Test(;
    u::Real = 0,
    v::Real = 0,
    f2::Real = 0.0,
    power::Real = 0.0)

    check_args(u=u, v=v, f2=f2, power=power)

    return fzero(x->powerTTest(u = u, v = v, f2 = f2, alpha = x) - power, 1e-10, 1 - 1e-10)
end

"""
    pwr.F2Test(u::Real = 0, v::Real = 0, f2::Float64 = 0.0, alpha::Float64 = 0.0, power::Float64 = 0.0)

Compute one of the test parameters such as sample size with `u` degrees of freedom
for the numerator and `v` degrees of freedom for the denominator,
effect size (`f2`), type I error (`alpha`), and power (`power`).
The parameter to be estimated must be set to zero (default).
"""
function F2Test(;
    u::Real = 0,
    v::Real = 0,
    f2::Real = 0.0,
    alpha::Real = 0.0,
    power::Real = 0.0)

    if sum([x == 0 for x in (u,v,f2,alpha,power)]) != 1
        error("exactly one of u, v, f2, power, and alpha must be zero")
    end

    if power == 0.0
        power = powerF2Test(u = u, v = v, f2 = f2, alpha = alpha)
    elseif alpha == 0.0
        alpha = alphaF2Test(u = u, v = v, f2 = f2, power = power)
    elseif f2 == 0.0
        f2 = effectsizeF2Test(u = u, v = v, alpha = alpha, power = power)
    elseif u == 0
        u = samplesizeF2Test(u = u, v = v, f2 = f2, alpha = alpha, power = power)
    elseif v == 0
        v = samplesizeF2Test(u = u, v = v, f2 = f2, alpha = alpha, power = power)
    end

    return htest(
        string("Multiple regression power calculation"),
        OrderedDict(
            "u" => u,
            "v" => v,
            "f2" => f2,
            "alpha" => alpha,
            "power" => power)
        )
end

max_stars_repo_name: mwsohn/pwr

max_issues_repo_name: mwsohn/pwr

max_forks_repo_name: mwsohn/pwr


text:
@testset "incomplete gamma ratios" begin
#Computed using Wolframalpha gamma(a,x)/gamma(a) ~ gamma_q(a,x,0) function.
    @test gamma_inc(10,10)[2] ≈ 0.45792971447185221
    @test gamma_inc(1,1)[2] ≈ 0.3678794411714423216
    @test gamma_inc(0.5,0.5)[2] ≈ 0.31731050786291410
    @test gamma_inc(BigFloat(30.5),BigFloat(30.5))[2] ≈ parse(BigFloat,"0.47591691193354987004") rtol=eps()
    @test gamma_inc(5.5,0.5)[2] ≈ 0.9999496100513121669
    @test gamma_inc(0.5,7.4)[2] ≈ 0.0001195355018130302
    @test gamma_inc(0.5,0.22)[2] ≈ 0.507122455359825146
    @test gamma_inc(0.5,0.8)[2] ≈ 0.20590321073206830887
    @test gamma_inc(11.5,0.5)[2] ≈ 0.999999999998406112
    @test gamma_inc(0.19,0.99)[2] ≈ 0.050147247342905857
    @test gamma_inc(0.9999,0.9999)[2] ≈ 0.3678730556923103
    @test gamma_inc(24,23.9999999999)[2] ≈ 0.472849720555859138
    @test gamma_inc(0.5,0.55)[2] ≈ 0.29426610430496289
    @test gamma_inc(Float32(0.5),Float32(0.55))[2] ≈ Float32(gamma_inc(0.5,0.55)[2])
    @test gamma_inc(Float16(0.5),Float16(0.55))[2] ≈ Float16(gamma_inc(0.5,0.55)[2])
    @test gamma_inc(30,29.99999)[2] ≈ 0.475717712451705704
    @test gamma_inc(30,29.9)[2] ≈ 0.482992166284958565
    @test gamma_inc(10,0.0001)[2] ≈ 1.0000
    @test gamma_inc(0.0001,0.0001)[2] ≈ 0.000862958131006599
    @test gamma_inc(0.0001,10.5)[1] ≈ 0.999999999758896146
    @test gamma_inc(1,1)[1] ≈ 0.63212055882855768
    @test gamma_inc(13,15.1)[2] ≈ 0.25940814264863701
    @test gamma_inc(0.6,1.3)[2] ≈ 0.136458554006505355
    @test gamma_inc((100),(80))[2] ≈ 0.9828916869648668
    @test gamma_inc((100),(80),1)[2] ≈ 0.9828916869
    @test Float16(gamma_inc((100),(80),2)[2]) ≈ Float16(.983)
    @test gamma_inc(13.5,15.1)[2] ≈ 0.305242642543419087
    @test gamma_inc(11,9)[1] ≈ 0.2940116796594881834
    @test gamma_inc(8,32)[1] ≈ 0.99999989060651042057
    @test gamma_inc(15,16)[2] ≈ 0.3675273597655649298
    @test gamma_inc(15.5,16)[2] ≈ 0.4167440299455427811
    @test gamma_inc(0.9,0.8)[1] ≈ 0.59832030278768172
    @test gamma_inc(1.7,2.5)[1] ≈ 0.78446115627678957
    @test gamma_inc(11.1,0.001)[2] ≈ 1.0000
    @test gamma_inc(1e7, (1e7)+1)[1] ≈ 0.5001682088254367
    @test gamma_inc(1e7, (1e7)+1)[2] ≈ 0.4998317911745633
    @test_throws DomainError gamma_inc(-1,2,2)
    @test_throws DomainError gamma_inc(0,0,1)
end
@testset "inverse of incomplete gamma ratios" begin
#Compared with Scipy.special.gammaincinv
    @test gamma_inc_inv(1.0,0.5,0.5) ≈ 0.69314718055994529
    ctr=1
    ans1 = [7.4153939596077105e-06, 4.1948837553001128e-05, 0.00011560348487144592, 0.00023733157806706144, 0.00041465371557152853, 0.00065420366900931063, 0.00096200074513191707, 0.0013436111982025845, 0.0018042543050560357, 0.0023488772409990069, 0.0029822108340951982, 0.0037088128669865747, 0.0045331028890989175, 0.0054593910326394268, 0.0064919024756292113, 0.0076347986776777021, 0.0088921961857776899, 0.010268183591667559, 0.011766837076460529, 0.013392234877244308, 0.015148470939019535, 0.01703966796416153, 0.019069990034434534, 0.021243654953428382, 0.02356494643740838, 0.026038226268114743, 0.028667946510727123, 0.031458661893117974, 0.034415042438025324, 0.037541886437423046, 0.040844133857854445, 0.044326880266606375, 0.047995391371211706, 0.051855118268798489, 0.055911713507252728, 0.060171048067033703, 0.064639229380863375, 0.069322620518498704, 0.074227860675536064, 0.079361887118881563, 0.084731958757381876, 0.090345681524433508, 0.09621103578051432, 0.10233640596793599, 0.10873061277817557, 0.11540294812451861, 0.12236321325012786, 0.12962176034488948, 0.13718953809449788, 0.14507814164343533, 0.15329986752124386, 0.16186777416055501, 0.17079574872787454, 0.18009858109673355, 0.18979204592069479, 0.19989299391476917, 0.21041945363284603, 0.22139074524173796, 0.23282760804678732, 0.24475234382888267, 0.25718897841979127, 0.27016344438650586, 0.28370378823436093, 0.29784040619656998, 0.3126063134848735, 0.32803745287127684, 0.34417304970517859, 0.36105602201014886, 0.37873345623790661, 0.39725716170037845, 0.41668431981048032, 0.43707824825000413, 0.45850930533751827, 0.4810559665889298, 0.50480611430358513, 0.52985859275268155, 0.55632509731250579, 0.58433248729003373, 0.61402564160686812, 0.64557101747749401, 0.67916113010215806, 0.71502025447176809, 0.75341177168885431, 0.79464776273185678, 0.83910172694543061, 0.8872257293787823, 0.93957396283390326, 0.99683583211997739, 1.0598835767754553, 1.1298428254726491, 1.2082007260001051, 1.296978498210863, 1.3990206562238474, 1.5185103158285946, 1.6619620177540984, 1.8403447442261882, 2.0743414883146705, 2.4107139457641296, 3.0000967446589719]
    ans2 = [0.0028980766274747166, 0.006908184939374363, 0.01149696382036067, 0.016518276652159908, 0.02189753286142773, 0.027589145508248085, 0.03356249352199024, 0.03979579246097344, 0.046272957709150415, 0.052981821800176826, 0.05991304249771414, 0.06705939521991744, 0.07441529398164964, 0.08197645553388541, 0.08973965716025696, 0.09770255795612678, 0.10586356446493835, 0.11422172813580914, 0.12277666614622669, 0.1315284997431042, 0.14047780597228823, 0.14962557982476085, 0.15897320462672093, 0.16852242906309764, 0.17827534962728914, 0.18823439758407795, 0.19840232975077693, 0.2087822225659563, 0.21937746904082012, 0.23019177828605702, 0.24122917738430874, 0.2524940154406757, 0.2639909696948641, 0.27572505362160177, 0.2877016269830214, 0.29992640782953106, 0.31240548647561406, 0.325145341505105, 0.33815285788768823, 0.35143534731545956, 0.36500057089608295, 0.3788567643680646, 0.3930126660346167, 0.40747754764618266, 0.4222612484987209, 0.4373742130560526, 0.4528275324509786, 0.4686329902724166, 0.48480311310582785, 0.5013512263630729, 0.5182915160173085, 0.5356390969506658, 0.5534100887296879, 0.5716216997488716, 0.5902923208297984, 0.6094416295366399, 0.6290907066738872, 0.649262166675578, 0.6699803038855925, 0.6912712570761026, 0.7131631949691678, 0.7356865260311793, 0.7588741364223205, 0.7827616607300203, 0.807387791030685, 0.832794630951811, 0.8590281028041281, 0.8861384175955134, 0.9141806199237333, 0.943215222504511, 0.9733089486005939, 1.0045356051136263, 1.0369771149071734, 1.0707247444923012, 1.1058805731481947, 1.1425592627436494, 1.1808902052215764, 1.2210201487102676, 1.2631164361809093, 1.307371036420099, 1.354005611796903, 1.4032779600790313, 1.4554903028704096, 1.5110000943860755, 1.5702343296128558, 1.633708805354231, 1.7020545444724036, 1.7760548371325608, 1.8566984660391717, 1.9452584157780999, 2.043412265612975, 2.153433912595464, 2.278514210551947, 2.4233308771016073, 2.595143613479916, 2.8061289679463384, 3.079145167578199, 3.4655956756034487, 4.12987911692835]
    ans3 = [0.1071755222650277, 0.16050998728955507, 0.20416656508566741, 0.24280678250939972, 0.27826797183157947, 0.31150317301391278, 0.34308358281638091, 0.37338428082062985, 0.40266774948514333, 0.43112656018759449, 0.45890727249428631, 0.48612476640606467, 0.51287131234095706, 0.53922256543527058, 0.5652416715299875, 0.5909821656505071, 0.61649007140344114, 0.64180545586647186, 0.66696360394811371, 0.69199592089152651, 0.71693063676844848, 0.7417933642658795, 0.76660754611766135, 0.7913948184012789, 0.81617530891938483, 0.84096788496285169, 0.86579036123554576, 0.89065967617260244, 0.91559204301304553, 0.94060308059774034, 0.96570792782021209, 0.99092134486698613, 1.0162578037796903, 1.04173157040596, 1.0673567794461716, 1.0931475040232679, 1.1191178209852426, 1.1452818729806347, 1.1716539282165654, 1.198248438708915, 1.2250800977595322, 1.2521638973416984, 1.2795151860393545, 1.3071497281656619, 1.3350837646807319, 1.363334076536002, 1.3919180510932918, 1.4208537523000897, 1.4501599953496895, 1.4798564266162111, 1.5099636097318834, 1.5405031187690257, 1.571497639604559, 1.6029710806837836, 1.6349486945666225, 1.6674572118386319, 1.7005249892070939, 1.7341821738873691, 1.7684608867260749, 1.8033954269180192, 1.839022501668433, 1.8753814847505423, 1.9125147086353582, 1.9504677957575041, 1.9892900355679812, 2.0290348153641435, 2.0697601145466438, 2.1115290740202663, 2.1544106550478053, 2.1984804051368831, 2.2438213526962998, 2.2905250575198943, 2.3386928510229508, 2.3884373090925655, 2.4398840121490508, 2.4931736625766741, 2.5484646505330266, 2.6059361873998439, 2.6657921648912399, 2.7282659516930696, 2.7936264154376702, 2.8621855665549849, 2.9343083789530127, 3.0104255776654698, 3.0910505401567816, 3.176802011280996, 3.2684352132712395, 3.3668853781968782, 3.4733301840058628, 3.5892819029224894, 3.7167280552758761, 3.8583548958347427, 4.0179202697917491, 4.2009145830871581, 4.4158272046973419, 4.6768376437172963, 5.0104347814570787, 5.4758322541221567, 6.2608605070659076]
    ans4 = [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, 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, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 9.793e-320, 8.8258281996309067e-311, 5.2442064082758777e-302, 2.0885820384164982e-293, 5.6626416070015106e-285, 1.0605839281955743e-276, 1.3913463845423177e-268, 1.2952838036620446e-260, 8.6636808810333524e-253, 4.2124146389233922e-245, 1.5054592477106387e-237, 3.9965668615698637e-230, 7.9602338168240791e-223, 1.20090997348131e-215, 1.3847294830684997e-208, 1.2309060164407832e-201, 8.5045125183290383e-195, 4.6029057651609662e-188, 1.9661243539196415e-181, 6.675413023480524e-175, 1.8137947775332813e-168, 3.9697776116742062e-162, 7.0423151317779386e-156, 1.0186480541842382e-149, 1.2082919222730254e-143, 1.1817770759861981e-137, 9.5806543461731422e-132, 6.4705147650865631e-126, 3.6582023269162791e-120, 1.7394045962948254e-114, 6.9867971004491818e-109, 2.3810402190694626e-103, 6.9129343541856114e-98, 1.71669965360591e-92, 3.6603939138610754e-87, 6.7261696481112057e-82, 1.0689601189656641e-76, 1.4743519045954981e-71, 1.7706253773856427e-66, 1.85749673681364e-61, 1.7074484121414459e-56, 1.3793741591778402e-51, 9.8216596440656967e-47, 6.1811283531832712e-42, 3.447488184610534e-37, 1.7085354519424486e-32, 7.5427383121540407e-28, 2.9735875496464905e-23, 1.0493212802006967e-18, 3.3220770969837717e-14, 9.4569508903874499e-10, 2.4259428385570885e-05]
    ans5 = [68.182564057578659, 70.357062784263817, 71.760877899786948, 72.829353353366415, 73.706386972437542, 74.458490388905602, 75.12219461238017, 75.719845624365121, 76.266166586236068, 76.771402301686194, 77.242986392330991, 77.686494958669414, 78.106225421683931, 78.50556502309793, 78.887234808310836, 79.253456568431432, 79.606070310177984, 79.946618932262979, 80.276410559274709, 80.596565284301676, 80.908050799453733, 81.211709956856978, 81.508282370701906, 81.798421552199031, 82.082708649987609, 82.361663579041377, 82.635754117812908, 82.905403408390271, 83.170996189584713, 83.432884016058068, 83.691389659635405, 83.946810846256611, 84.199423449687018, 84.449484238382098, 84.697233252859277, 84.942895876083071, 85.18668464780427, 85.428800864609457, 85.669436000206872, 85.908772974675117, 86.146987296778406, 86.384248099705346, 86.620719087619619, 86.856559407959935, 87.091924462499136, 87.326966668607852, 87.561836180921006, 87.796681582652582, 88.031650555063862, 88.266890533043721, 88.502549354439552, 88.738775910585431, 88.975720805449285, 89.213537031000897, 89.45238066667676, 89.692411611339708, 89.933794356792788, 90.176698812815317, 90.421301194840566, 90.667784986823875, 90.916341993645545, 91.167173499597013, 91.420491552216376, 91.676520394081336, 91.935498069297722, 92.197678236508722, 92.463332226600599, 92.732751391148227, 93.006249797557189, 93.284167339317776, 93.566873345624515, 93.854770794832618, 94.148301262289351, 94.447950766886493, 94.754256724998356, 95.067816279076155, 95.38929634645433, 95.719445839736153, 96.059110654771956, 96.409252222677836, 96.770970703934069, 97.145534304515252, 97.534416777401361, 97.939346035676579, 98.362368106482464, 98.805932668773451, 99.273009611654814, 99.767251266034179, 100.2932237698064, 100.85674648050276, 101.46540670265372, 102.12937183070848, 102.86273396770143, 103.68587349786873, 104.62994489700357, 105.74630631344544, 107.12941611870914, 108.98632200164798, 111.9556294726119]
    for x = 0.01:0.01:0.99
        @test gamma_inc_inv(0.4, x, 1.0-x) ≈ ans1[ctr]#scipy.special.gammaincinv(0.4,x)
        @test gamma_inc_inv(0.8, x, 1.0-x) ≈ ans2[ctr]#scipy.special.gammaincinv(0.8,x)
        @test gamma_inc_inv(1.8, x, 1.0-x) ≈ ans3[ctr]#scipy.special.gammaincinv(1.8,x)
        @test gamma_inc_inv(0.001, x, 1.0-x) ≈ ans4[ctr]#scipy.special.gammaincinv(.001,x)
        @test gamma_inc_inv(88.6, x, 1.0-x) ≈ ans5[ctr]#scipy.special.gammaincinv(88.6,x)
        ctr+=1
    end
    for x=-.5:.5:.9
        @test SpecialFunctions.loggamma1p(x) ≈ loggamma(1.0+x)
    end
    for x = .5:5.0:100.0
        @test SpecialFunctions.stirling(x) ≈ log(gamma(x)) - (x-.5)*log(x)+x- log(2*pi)/2.0
    end
end

double(x::Real) = Float64(x)
double(x::Integer) = Int(x)
double(x::Complex) = ComplexF64(x)

@testset "upper incomplete gamma function" begin
    setprecision(BigFloat, 256) do
        for a in Any[0:0.4:3; 1:3], x in 0:0.2:2
            @test gamma(a,x) ≈ gamma(big(a),big(x))
        end
        @test_throws DomainError gamma(-2.2, -1.3)
        for (a,x, exact) in (
            (2, big"+1.3", big"0.6268231239782289871813662853957039889809398944861850589869804057956189274569818"),
            (2, big"-1.3", big"-1.100789000285773266137246974803445875429455665367265440176867948378982424804222"),
            (big"2.2", big"1.3", big"0.7550934853735524106456916078787596171599416239996064979513848803118671763945577"),
            (big"-2.2", big"1.3", big"0.03938564959393195337328006806473296774233286427565465045140458665248322893076784"),
            (big"-2.2", big"-1.3"+0im, Complex(big"0.1688350167695890519002747177035271528716667324397453142169971691104641885370081", big"-1.724677939954414412829215929952389349929001266936564849801773346878175589599461")),
            (2, big"1.3"+2im, Complex(big"0.2347744561498965806069446787000456827042253434143074251879541754828136789074293", big"-0.7967951407674886396960561446265346226800340382780520672369163777061700801594228")),
            (big"2.2", big"1.3"+2im, Complex(big"0.4226969694490623767886715572749155659150206342283251276656790859161308520476838", big"-0.9507541450333763666696561254491461255244463810462729822382144498284499948680099")),
            (big"2.2"-2im, big"1.3"+2im, Complex(big"-3.858698824275578861673404086439928163791926530897603434146764316225759116935376", big"0.2105970967650200831829566202893410569725743566037534990454186505024476246014339")),
            (30, big"-1e2", big"-2.080142496354742431762923569133534773079637504483999408076147454085194055640495e+101"),
            )
            if !(exact isa Complex) # we don't yet have a Complex{BigFloat} gamma function
                @test gamma(a, x) ≈ exact atol=eps(abs(exact))*1000
            end
            @test gamma(double(a), double(x)) ≈ double(exact) rtol=1e-13
        end
    end
    @test gamma(30, big(-1000)) ≈ big"-1.914496653187781420578078670260160511219745144309147121829045802464973219500799e+521" rtol=1e-70
    @test gamma(30, -1000) == -Inf
    @test gamma(Inf, 51.2) == Inf
    @test gamma(-Inf, -51.2) == 0.0
    @test_throws DomainError gamma(Inf, -51.2)
    @test gamma(2.3, Inf) == 0.0
    @test gamma(2, -Inf) == -Inf
    @test_throws DomainError gamma(2.2, -Inf)
end

@testset "upper incomplete gamma function logarithm" begin
    for (a,z) in ((3,5), (3,-5), (3,5+4im), (3,-5+4im), (3,5-4im), (-3,5+4im), (-3,-5+4im))
        @test exp(loggamma(a,z)) ≈ gamma(a,z) rtol=1e-13
    end
    @test loggamma(50, 1e7) ≅ -9.9992102133082030308153473164868727954977460876571275797855e6
    @test real(loggamma(50, 1e7 + 1e8im)) ≅ -9.999097142860392e6
    @test cis(imag(loggamma(50, 1e7 + 1e8im))) ≈ cis(1.0275220422549918) rtol=1e-8
    @test real(loggamma(10+20im, -1e5 + 1e8im)) ≈ 100134.3502048662475864409896160245625409130538724704329203542339
    @test cis(imag(loggamma(10+20im, -1e5 + 1e8im))) ≈ cis(-2.6572071454623347) rtol=1e-8
    @test loggamma(-1e8, 1e9) ≅ -3.0723266045132171331933746054573197040165846554476396719312e9
    @test loggamma(3, Inf) == -Inf
    @test_throws DomainError loggamma(3, -Inf)
    @test loggamma(Inf, 3.2) == Inf
    @test loggamma(-Inf, 3.2) == -Inf
    @test_throws DomainError loggamma(Inf, -3.2)
    @test loggamma(117.3, 0) == loggamma(117.3)
    @test loggamma(7, -300.2) ≅ log(gamma(7, -300.2))
    @test_throws DomainError loggamma(6, -3.2)
end


text:
using Test

function test_grid()
    shape = (4, 4, 4)
    grid = Grid(shape, (1.0, 1.0, 1.0))
    all(size(grid) .== shape)
end

function test_fields()
    grid = Grid((4, 5, 6), (1.0, 1.0, 1.0))
    fields = Fields(grid)
    update_curl_E!(fields)
    true
end

function test_poisson()
    grid = Grid((30, 30, 3), (10.0, 10.0, 0.1))
    charge_density = zeros(size(grid))
    charge_density[6, 6, 2] = 1.0
    gs_potential = zeros(size(grid))
    solve_poisson!(
        gs_potential,
        charge_density,
        grid,
        Free,
        GaussSeidel(max_deviation = 1e-6, relaxation_factor = 1.6),
    )
    correct_potential = zeros(size(grid))
    solve_poisson!(correct_potential, charge_density, grid, Free, Analytical)
    maximum(abs.(gs_potential - correct_potential)[isfinite.(
        correct_potential,
    )]) / compute_solution_scale(sum(abs.(charge_density)), grid, DIM3) < 0.06
end

function test_particles()
    positions = zeros(2, 1)
    positions[:, 1] = [0.0, 0.0]
    grid = Grid((30, 30), (10.0, 10.0))
    charge_density = zeros(size(grid))
    potential = zeros(size(grid))

    n_steps = 20
    total_charge = zeros(n_steps)

    for i = 1:n_steps
        particles = Particles{2}(positions, 1.0, 1.0)
        charge_density .= 0
        sample_charge_density!(charge_density, grid, particles)
        total_charge[i] = sum(charge_density) * cell_volume(grid)
        positions .+= [0.1, 0.07]
    end

    maximum(abs.(total_charge .- 1.0)) < 1e-12
end

function test_interpolation()
    grid = Grid((4, 5, 6), (1.0, 1.0, 1.0))
    fields = Fields(grid)
    E_interp = FieldInterpolator(fields, ElectricField, DIM1)
    interpolate(E_interp, [0.3, 0.81, 0.02])
    true
end

function test_timestepping()
    grid = Grid((50, 50), (1.0, 1.0))
    positions = zeros(2, 1)
    positions[:, 1] = [0.2, 0.5]
    #positions[:, 2] = [0.5, 0.5]
    particles = Particles{2}(positions, 1.0, 1.0)
    fields = initialize_fields(particles, grid)
    savefig(
        scalarheatmap(sqrt.(
            fields.E[1] .^ 2 .+ fields.E[2] .^ 2 .+ fields.E[3] .^ 2,
        )),
        "Enorm.png",
    )
    n_steps = 100
    n_skips = 1
    anim = @animate for _ = 1:n_steps
        vectorfieldquiver(centers(grid)[1], centers(grid)[2], fields.E)
        advance!(particles, fields, 1e-2)
        @show(maximum(fields.E[1]))
    end every n_skips
    gif(anim, "E.gif", fps = 15)

    true
end

export test_timestepping

function runtests()

    @testset "grid.jl" begin
        @test test_grid()
    end

    @testset "fields.jl" begin
        @test test_fields()
    end

    @testset "poisson.jl" begin
        @test test_poisson()
    end

    @testset "particles.jl" begin
        @test test_particles()
    end

    @testset "interp.jl" begin
        @test test_interpolation()
    end

    @testset "timestepping.jl" begin
        @test test_timestepping()
    end
end
export runtests

max_stars_repo_name: lars-frogner/PIC

max_issues_repo_name: lars-frogner/PIC

max_forks_repo_name: lars-frogner/PIC


text:
"""
    `get_rs_mats!(REG_DATA)`

    Computes necessary matrices for transfer matrix methods.

    # Examples
    ```julia-repl
    julia> CpelNano.get_rs_mats!(x)
    ```
"""
function get_rs_mats!(rs::RegStruct)::Nothing

    # Get αs & βs
    αs, βs = get_αβ_from_ϕ(rs.ϕhat, rs)

    # Get u's
    rs.tm.u1 = get_u(αs[1])
    rs.tm.uN = get_u(αs[end])

    # Get W's
    Ws = Vector{Array{Float64,2}}()
    @inbounds for n = 1:length(βs)
        push!(Ws, get_W(αs[n], αs[n + 1], βs[n]))
    end

    # Record W's
    rs.tm.Ws = Ws

    # Return nothing
    return nothing

end
"""
    `get_rs_lgtrck_mats!(REG_DATA)`

    Computes necessary matrices for transfer matrix methods that relies on the log-sum trick.

    # Examples
    ```julia-repl
    julia> CpelNano.get_rs_lgtrck_mats!(x)
    ```
"""
function get_rs_lgtrck_mats!(rs::RegStruct)::Nothing

    # Get αs & βs
    αs, βs = get_αβ_from_ϕ(rs.ϕhat, rs)

    # Get u's
    rs.tm.log_u1 = get_log_u(αs[1])
    rs.tm.log_uN = get_log_u(αs[end])

    # Get W's
    logWs = Vector{Array{Float64,2}}()
    @inbounds for n = 1:length(βs)
        push!(logWs, get_log_W(αs[n], αs[n + 1], βs[n]))
    end

    # Record W's
    rs.tm.log_Ws = logWs

    # Return nothing
    return nothing

end
"""
    `get_rs_logZ!(REG_DATA)`

    Computes and stores the (log) partition function using log-sum trick transfer matrix.

    # Examples
    ```julia-repl
    julia> CpelNano.get_rs_logZ!(x)
    ```
"""
function get_rs_logZ!(rs::RegStruct)::Nothing

    # Store partition function
    rs.logZ = get_log_Z(rs.tm.log_u1, rs.tm.log_uN, rs.tm.log_Ws)

    # Return nothing
    return nothing

end
"""
    `comp_log_g_p(REG_DATA,P,Xp)`

    Computes log(g_1(xp)) values.

    # Examples
    ```julia-repl
    julia> CpelNano.comp_log_g_p(rs,p,xp)
    ```
"""
function comp_log_g_p(rs::RegStruct, p::Int64, xp::Float64)::Float64

    # If p=1 then return nothing
    p == 1 && return 0.00

    # Get αs & βs
    αs, βs = get_αβ_from_ϕ(rs.ϕhat, rs)

    # Get log(uN)
    aux = (αs[p - 1] + xp * βs[p - 1]) / 2.0
    log_uNp = [-aux,+aux]

    # If p=2 then return inner product of v's
    p == 2 && return log_vec_vec_mult(log_uNp, log_uNp)

    # Get log(W)
    log_Wp = [
        -αs[p - 2] / 2.0 + βs[p - 2] - aux -αs[p - 2] / 2.0 - βs[p - 2] + aux;
        +αs[p - 2] / 2.0 - βs[p - 2] - aux +αs[p - 2] / 2.0 + βs[p - 2] + aux
    ]

    # If p=3 then return inner product of v's
    p == 3 && return log_vec_vec_mult(log_vec_mat_mult(rs.tm.log_u1, log_Wp), log_uNp)

    # Return log(g1(xp))
    log_Ws = mult_log_mats(vcat(rs.tm.log_Ws[1:(p - 3)], [log_Wp]))
    return log_vec_vec_mult(log_vec_mat_mult(rs.tm.log_u1, log_Ws), log_uNp)

end
"""
    `comp_log_g_q(REG_DATA,Q,Xq)`

    Computes log(g_2(xq)) values.

    # Examples
    ```julia-repl
    julia> CpelNano.comp_log_g_q(rs,q,xq)
    ```
"""
function comp_log_g_q(rs::RegStruct, q::Int64, xq::Float64)::Float64

    # If q=N(=L after re-scaling) then return log(1)
    q == rs.L && return 0.00

    # Get αs & βs
    αs, βs = get_αβ_from_ϕ(rs.ϕhat, rs)

    # Get log(u1)
    aux = (αs[q + 1] + xq * βs[q]) / 2.0
    log_u1p = [-aux; +aux]

    # If q=N-1 then return inner product of v's
    q == rs.L - 1 && return log_vec_vec_mult(log_u1p, log_u1p)

    # Get Wp
    log_Wp = [
        -aux + βs[q + 1] - αs[q + 2] / 2.0 -aux - βs[q + 1] + αs[q + 2] / 2.0;
        +aux - βs[q + 1] - αs[q + 2] / 2.0 +aux + βs[q + 1] + αs[q + 2] / 2.0
    ]

    # If q=N-2 then return inner product of v's
    q == rs.L - 2 && return log_vec_vec_mult(log_vec_mat_mult(log_u1p, log_Wp), rs.tm.log_uN)

    # Return log(g2(xq))
    log_Ws = mult_log_mats(vcat([log_Wp], rs.tm.log_Ws[(q + 2):end]))
    return log_vec_vec_mult(log_vec_mat_mult(log_u1p, log_Ws), rs.tm.log_uN)

end
"""
    `get_rs_log_gs!(REG_DATA)`

    Computes and stores log(g_i(±⋅)) values.

    # Examples
    ```julia-repl
    julia> CpelNano.get_rs_log_gs!(rs)
    ```
"""
function get_rs_log_gs!(rs::RegStruct)::Nothing

    # Loop over analysis regions
    @inbounds for k = 1:rs.nls_rgs.num

        # Get p and q
        p = minimum(rs.nls_rgs.cpg_ind[k])
        q = maximum(rs.nls_rgs.cpg_ind[k])

        if p != 0 && q != 0
            # Compute log g1(±x_p) & log g2(±x_q)
            push!(rs.tm.log_gs.pm, comp_log_g_p(rs, p, -1.0))
            push!(rs.tm.log_gs.pp, comp_log_g_p(rs, p, +1.0))
            push!(rs.tm.log_gs.qm, comp_log_g_q(rs, q, -1.0))
            push!(rs.tm.log_gs.qp, comp_log_g_q(rs, q, +1.0))
        else
            # Push NaN
            push!(rs.tm.log_gs.pm, NaN)
            push!(rs.tm.log_gs.pp, NaN)
            push!(rs.tm.log_gs.qm, NaN)
            push!(rs.tm.log_gs.qp, NaN)
        end

    end

    # Return nothing
    return nothing

end
"""

    `get_rs_exp_log_g1!(RS_DATA)`

    Computes E[log(g1(Xp))] for all analysis regions using the transfer matrix method and log-trick.

    # Examples
    ```julia-repl
    julia> CpelNano.get_rs_exp_log_g1!(rs)
    ```
"""
function get_rs_exp_log_g1!(rs::RegStruct)::Nothing

    # Loop over analysis regions
    rs.exps.log_g1 = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num
        # Get index & next if no CpG sites
        p = minimum(rs.nls_rgs.cpg_ind[k])
        p == 0 && continue
        # Compute expectation
        if p == 1
            # If first CpG site (can happen k>1)
            rs.exps.log_g1[k] = 0.0
        else
            # Compute expectation using TM (using marginal seems to be unstable)
            logD = log.([rs.tm.log_gs.pm[k] 0.0;0.0 rs.tm.log_gs.pp[k]])
            log_WsDWs = mult_log_mats(vcat(rs.tm.log_Ws[1:(p - 1)], [logD], rs.tm.log_Ws[p:end]))
            log_u1WsDWs = log_vec_mat_mult(rs.tm.log_u1, log_WsDWs)
            rs.exps.log_g1[k] = exp(log_vec_vec_mult(log_u1WsDWs, rs.tm.log_uN) - rs.logZ)
        end
    end

    # Return nothing
    return nothing

end
"""

    `comp_crs_exps_log_g1(RS,RS_MIX)`

    Computes E[log(g1(Xp;θ̃));θi] for all analysis regions using the transfer matrix method and log-trick.

    # Examples
    ```julia-repl
    julia> CpelNano.comp_crs_exps_log_g1(rs,rs_mix)
    ```
"""
function comp_crs_exps_log_g1(rs::RegStruct, rs_mix::RegStruct)::Vector{Float64}

    # Loop over analysis regions
    exps_log_g1 = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num
        # Get index & next if no CpG sites
        p = minimum(rs.nls_rgs.cpg_ind[k])
        p == 0 && continue
        # Compute expectation
        if p == 1
            # If first CpG site (can happen k>1)
            exps_log_g1[k] = 0.0
        else
            # Compute expectation using TM (using marginal seems to be unstable)
            logD = log.([rs_mix.tm.log_gs.pm[k] 0.0;0.0 rs_mix.tm.log_gs.pp[k]])
            log_WsDWs = mult_log_mats(vcat(rs.tm.log_Ws[1:(p - 1)], [logD], rs.tm.log_Ws[p:end]))
            log_u1WsDWs = log_vec_mat_mult(rs.tm.log_u1, log_WsDWs)
            exps_log_g1[k] = exp(log_vec_vec_mult(log_u1WsDWs, rs.tm.log_uN) - rs.logZ)
        end
    end

    # Return vector
    return exps_log_g1

end
"""

    `get_rs_exp_log_g2!(RS_DATA)`

    Computes E[log(g2(Xq))] for all analysis regions using the transfer matrix method and log-trick.

    # Examples
    ```julia-repl
    julia> CpelNano.get_rs_exp_log_g2!(rs)
    ```
"""
function get_rs_exp_log_g2!(rs::RegStruct)::Nothing

    # Loop over analysis regions
    rs.exps.log_g2 = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num
        # Get index & next if no CpG sites
        q = maximum(rs.nls_rgs.cpg_ind[k])
        q == 0 && continue
        # Compute expectation
        if q == rs.N
            # If last CpG site (can happen k<K)
            rs.exps.log_g2[k] = 0.0
        else
            # Compute expectation using TM (using marginal seems to be unstable)
            logD = log.([rs.tm.log_gs.qm[k] 0.0;0.0 rs.tm.log_gs.qp[k]])
            log_WsDWs = mult_log_mats(vcat(rs.tm.log_Ws[1:(q - 1)], [logD], rs.tm.log_Ws[q:end]))
            log_u1WsDWs = log_vec_mat_mult(rs.tm.log_u1, log_WsDWs)
            rs.exps.log_g2[k] = exp(log_vec_vec_mult(log_u1WsDWs, rs.tm.log_uN) - rs.logZ)
        end
    end

    # Return nothing
    return nothing

end
"""

    `comp_crs_exps_log_g2(RS,RS_MIX)`

    Computes E[log(g2(Xq;θ̃));θi] for all analysis regions using the transfer matrix method and log-trick.

    # Examples
    ```julia-repl
    julia> CpelNano.comp_crs_exps_log_g2(rs,rs_mix)
    ```
"""
function comp_crs_exps_log_g2(rs::RegStruct, rs_mix::RegStruct)::Vector{Float64}

    # Loop over analysis regions
    exps_log_g2 = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num
        # Get index & next if no CpG sites
        q = maximum(rs.nls_rgs.cpg_ind[k])
        q == 0 && continue
        # Compute expectation
        if q == rs.N
            # If last CpG site (can happen k<K)
            exps_log_g2[k] = 0.0
        else
            # Compute expectation using TM (using marginal seems to be unstable)
            logD = log.([rs_mix.tm.log_gs.qm[k] 0.0;0.0 rs_mix.tm.log_gs.qp[k]])
            log_WsDWs = mult_log_mats(vcat(rs.tm.log_Ws[1:(q - 1)], [logD], rs.tm.log_Ws[q:end]))
            log_u1WsDWs = log_vec_mat_mult(rs.tm.log_u1, log_WsDWs)
            exps_log_g2[k] = exp(log_vec_vec_mult(log_u1WsDWs, rs.tm.log_uN) - rs.logZ)
        end
    end

    # Return vector
    return exps_log_g2

end
"""
    `comp_mml!(REG_DATA)`

    Function that computes mean methylation level (MML) vector.

    # Examples
    ```julia-repl
    julia> CpelNano.comp_mml!(rs)
    ```
"""
function comp_mml!(rs::RegStruct)::Nothing

    # Loop over analysis regions
    rs.mml = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num

        # Check analysis region has CpG sites
        p = minimum(rs.nls_rgs.cpg_ind[k])
        p == 0 && continue

        # Check number of CpG sites
        Nk = length(rs.nls_rgs.cpg_ind[k])

        # Set MML vector
        rs.mml[k] = 0.5 / Nk * (Nk + sum(rs.exps.ex[rs.nls_rgs.cpg_ind[k]]))

    end

    # Return nothing
    return nothing

end
"""
    `comp_nme!(REG_DATA)`

    Function that computes normalized methylation entropy (NME) vector for each analysis
    region in the estimation region.

    # Examples
    ```julia-repl
    julia> CpelNano.comp_nme!(rs)
    ```
"""
function comp_nme!(rs::RegStruct)::Nothing

    # Get αs & βs
    αs, βs = get_αβ_from_ϕ(rs.ϕhat, rs)

    # Loop over analysis regions
    rs.nme = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num

        # Get indices
        p = minimum(rs.nls_rgs.cpg_ind[k])
        q = maximum(rs.nls_rgs.cpg_ind[k])
        p == q == 0 && continue

        # Add terms and normalize
        nme = rs.logZ - rs.exps.log_g1[k] - rs.exps.log_g2[k] - αs[p:q]' * rs.exps.ex[p:q]
        nme += p == q ? 0.0 : - βs[p:(q - 1)]' * rs.exps.exx[p:(q - 1)]
        nme /= (q - p + 1) * LOG2

        # Store
        rs.nme[k] = nme

    end

    # Return nothing
    return nothing

end
"""
    `comp_crs_ent(RS,RS_MIX)`

    Function that computes cross-entropy H(θi,θ̃). This function is used in function
    `comp_cmd` to compute the Coefficient of Methylation Divergence (CMD).

    # Examples
    ```julia-repl
    julia> CpelNano.comp_crs_ent(rs,rs_mix)
    ```
"""
function comp_crs_ent(rs::RegStruct, rs_mix::RegStruct)::Vector{Float64}

    # Get αs & βs of geometric mixture of Isings
    α̃s, β̃s = get_αβ_from_ϕ(rs_mix.ϕhat, rs_mix)

    # Compute E[log g1(Xp;θ̃);θi] and E[log g2(Xq;θ̃);θi] for all analysis regions
    exps_log_g1 = comp_crs_exps_log_g1(rs, rs_mix)
    exps_log_g2 = comp_crs_exps_log_g2(rs, rs_mix)
        # print_log("exps_log_g1=$(exps_log_g1)")
        # print_log("exps_log_g2=$(exps_log_g2)")

    # Loop over analysis regions
    crs_ent_vec = fill(NaN, rs.nls_rgs.num)
    @inbounds for k = 1:rs.nls_rgs.num

        # Get indices
        p = minimum(rs.nls_rgs.cpg_ind[k])
        q = maximum(rs.nls_rgs.cpg_ind[k])
        p == q == 0 && continue

        # Compute cross entropy
        crs_ent = rs_mix.logZ - α̃s[p:q]' * rs.exps.ex[p:q] - exps_log_g1[k] - exps_log_g2[k]
        crs_ent += p == q ? 0.0 : - β̃s[p:(q - 1)]' * rs.exps.exx[p:(q - 1)]

        # Store
        crs_ent_vec[k] = crs_ent

    end

    # Return vector
    return crs_ent_vec

end
"""
    `get_stat_sums!(RS)`

    Function that computes and stores all statistical summaries for an estimation region.

    # Examples
    ```julia-repl
    julia> CpelNano.get_stat_sums!(RS)
    ```
"""
function get_stat_sums!(rs::RegStruct, config::CpelNanoConfig)::Nothing

    # Get matrices
    get_rs_lgtrck_mats!(rs)

    # Partition function
    get_rs_logZ!(rs)

    # Analysis region info
    get_nls_reg_info!(rs, config)

    # Get them all
    get_rs_log_gs!(rs)

    # Compute E[log g(X)]
    get_rs_exp_log_g1!(rs)
    get_rs_exp_log_g2!(rs)

    # Compute E[X] and E[XX]
    get_rs_exps!(rs)

    # Compute μ(X)
    comp_mml!(rs)

    # Compute h(X)
    comp_nme!(rs)

    # Return nothing
    return nothing

end
"""
    `comp_cmd(REG_DATA_1,REG_DATA_2)`

    Function that computes the normalized generalised Jensen-Shannon divergence (G-JSD), or
    Coefficient of Methylation Divergence (CMD), over each analysis region.

    # Examples
    ```julia-repl
    julia> comp_cmd(REG_DATA_1,REG_DATA_2)
    ```
"""
function comp_cmd(rs1::RegStruct, rs2::RegStruct)::Vector{Float64}

    ## Struct for geometric mixture

    # Create structure
    mix = CpelNano.RegStruct()
    mix.cpg_pos = rs1.cpg_pos;
    mix.N = rs1.N;
    mix.ρn = rs1.ρn;
    mix.dn = rs1.dn;
    mix.nls_rgs = rs1.nls_rgs;
    mix.Nl = rs1.Nl;
    mix.L = rs1.L;
    mix.ρl = rs1.ρl;
    mix.dl = rs1.dl;
    mix.ϕhat = 0.5 .* (rs1.ϕhat + rs2.ϕhat)

    # Required quantities
    get_rs_lgtrck_mats!(mix)
    get_rs_logZ!(mix)
    get_rs_exps!(mix)
    get_rs_log_gs!(mix)

    ## Compute cross-entropies

    # Cross-entropy w/ θ1 & θ2
    cross_1 = comp_crs_ent(rs1, mix)
    cross_2 = comp_crs_ent(rs2, mix)

    ## Compute CMD for each analysis region

    cmd_vec = fill(NaN, rs1.nls_rgs.num)
    @inbounds for k = 1:rs1.nls_rgs.num

        # Get indices
        p = minimum(rs1.nls_rgs.cpg_ind[k])
        q = maximum(rs1.nls_rgs.cpg_ind[k])

        # Skip if empty analysis region
        p == q == 0 && continue

        # Shannon's Entropy
        ent_1 = rs1.nme[k] * (q - p + 1) * LOG2
        ent_2 = rs2.nme[k] * (q - p + 1) * LOG2
            # print_log("H(θ1)=$(ent_1); H(θ2)=$(ent_2)")
            # print_log("H(θ1,θ̃)=$(cross_1[k]); H(θ2,θ̃)=$(cross_2[k])")

        # If CMD is undefined continue
        cross_1[k] == cross_2[k] == 0.0 && continue

        # Store CMD k-th component
        cmd_vec[k] = 1.0 - (ent_1 + ent_2) / (cross_1[k] + cross_2[k])

    end

    # Return CMD
    return cmd_vec

end
###################################################################################################
# OTHER QUANTITIES
###################################################################################################
"""
    `get_ρ(u1, uN, Ws, Z)`

    Computes Pearson correlation between x_n and x_(n+1)

    # Examples
    ```julia-repl
    julia> Ngrp = 10; αs = fill(0.0,Ngrp); βs = fill(0.0,Ngrp-1);
    julia> u1 = CpelNano.get_u(αs[1]); uN = CpelNano.get_u(αs[end]);
    julia> Ws = [CpelNano.get_W(αs[n],αs[n+1],βs[n]) for n=1:length(βs)];
    julia> Z = CpelNano.get_Z(u1,uN,Ws);
    julia> CpelNano.get_ρ(u1, uN, Ws, Z)
    9-element Array{Float64,1}:
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
 ```
 """
function get_ρ(u1::Vector{Float64}, uN::Vector{Float64}, Ws::Vector{Array{Float64,2}}, Z::Float64)::Vector{Float64}

    # Get E[X] & E[X{i}X{i+1}]
    exs = CpelNano.get_E_X(u1, uN, Ws, Z)
    exxs = CpelNano.get_E_XX(u1, uN, Ws, Z)

    # Compute √Var(Xi)
    std_vec = sqrt.(ones((length(Ws) + 1)) - exs.^2)

    # Compute covariance
    ex = exs[1:length(exs) - 1]
    ex_prime = exs[2:length(exs)]
    cov_vec = exxs - (ex .* ex_prime)

    # Return correlation
    return cov_vec ./ (std_vec[1:length(std_vec) - 1] .* std_vec[2:length(std_vec)])

end

"""
    `get_E_XiXj(u1,uN,Ws,Z,i,j)`

    Computes E[X_iX_j], where j is larger than i

    # Examples
    ```julia-repl
    julia> Ngrp = 10; αs = fill(0.0,Ngrp); βs = fill(0.0,Ngrp-1);
    julia> u1 = CpelNano.get_u(αs[1]); uN = CpelNano.get_u(αs[end]);
    julia> Ws = [CpelNano.get_W(αs[n],αs[n+1],βs[n]) for n=1:length(βs)];
    julia> Z = CpelNano.get_Z(u1,uN,Ws);
    julia> CpelNano.get_E_XiXj(u1,uN,Ws,Z,1,10)
    0.0
 ```
 """
function get_E_XiXj(u1::Vector{Float64}, uN::Vector{Float64}, Ws::Vector{Array{Float64,2}}, Z::Float64, i::Int64, j::Int64)::Float64

    # Scale Ws for numerical stability
    # Ws = US*Ws

    # Confirm that j is larger than i
    if i >= j
        return NaN
    end

    ## Compute E[X_iX_j]

    # Compute transfer matrix product up to i
    if i == 1
        exixj_temp = u1' * [-1 0; 0 1]
    else
        exixj_temp = u1' * prod(Ws[1:i - 1]) * [-1 0; 0 1]
    end

    # Compute transfer matrix product from i to N
    if j == (length(Ws) + 1)
        exixj = exixj_temp * prod(Ws[i:end]) * [-1 0; 0 1] * uN
    else
        exixj = exixj_temp * prod(Ws[i:j - 1]) * [-1 0; 0 1] * prod(Ws[j:end]) * uN
    end

    # Return E[X_iX_j]
    return exixj / (u1' * prod(Ws) * uN)

end

"""
    `Σmat(u1,uN,Ws,Z)`

    Computes the covariance matrix in upper triangular form, where the diagonal is the variance and the [i,j] element
    is the covariance between Xi and Xj

    # Examples
    ```julia-repl
    julia> Ngrp = 10; αs = fill(0.0,Ngrp); βs = fill(0.0,Ngrp-1);
    julia> u1 = CpelNano.get_u(αs[1]); uN = CpelNano.get_u(αs[end]);
    julia> Ws = [CpelNano.get_W(αs[n],αs[n+1],βs[n]) for n=1:length(βs)];
    julia> Z = CpelNano.get_Z(u1,uN,Ws);
    julia> CpelNano.Σmat(u1,uN,Ws,Z)
    10×10 UpperTriangular{Float64,Array{Float64,2}}:
     1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     ⋅   1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     ⋅    ⋅   1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     ⋅    ⋅    ⋅   1.0  0.0  0.0  0.0  0.0  0.0  0.0
     ⋅    ⋅    ⋅    ⋅   1.0  0.0  0.0  0.0  0.0  0.0
     ⋅    ⋅    ⋅    ⋅    ⋅   1.0  0.0  0.0  0.0  0.0
     ⋅    ⋅    ⋅    ⋅    ⋅    ⋅   1.0  0.0  0.0  0.0
     ⋅    ⋅    ⋅    ⋅    ⋅    ⋅    ⋅   1.0  0.0  0.0
     ⋅    ⋅    ⋅    ⋅    ⋅    ⋅    ⋅    ⋅   1.0  0.0
     ⋅    ⋅    ⋅    ⋅    ⋅    ⋅    ⋅    ⋅    ⋅   1.0
 ```
 """
function Σmat(u1::Vector{Float64}, uN::Vector{Float64}, Ws::Vector{Array{Float64,2}}, Z::Float64)::UpperTriangular

    # Num CpG sites
    N = length(Ws) + 1

    # Get E[X]
    exs = CpelNano.get_E_X(u1, uN, Ws, Z)

    # Compute Var(Xi)
    var_vec = ones(N) - exs.^2 # add square

    # Construct covariance matrix
    cov_matrix = Array{Float64}(undef, N, N)

    @inbounds for i in 1:N
        @inbounds for j in i:N
            # Diagonal elements are the variance
            if i == j
                cov_matrix[i,j] = var_vec[i]
            else
                cov_matrix[i,j] = CpelNano.get_E_XiXj(u1, uN, Ws, Z, i, j)
            end
        end
    end

    # Return Σ
    return UpperTriangular(cov_matrix)

end


text:
@testset "EuclideanVector"  begin

#---

using ACE
using Random, Printf, Test, LinearAlgebra, ACE.Testing
using ACE: evaluate, evaluate_d, SymmetricBasis, NaiveTotalDegree, PIBasis
using ACE.Random: rand_rot, rand_refl

# construct the 1p-basis
D = NaiveTotalDegree()
maxdeg = 6
ord = 3

B1p = ACE.Utils.RnYlm_1pbasis(; maxdeg=maxdeg, D = D)

# generate a configuration
nX = 10
Xs = rand(EuclideanVectorState, B1p.bases[1], nX)
cfg = ACEConfig(Xs)

#---

@info("SymmetricBasis construction and evaluation: EuclideanVector")

φ = ACE.EuclideanVector(Complex{Float64})
pibasis = PIBasis(B1p, ord, maxdeg; property = φ, isreal=false)
basis = SymmetricBasis(pibasis, φ)
@time SymmetricBasis(pibasis, φ);

BB = evaluate(basis, cfg)

# a stupid but necessary test
BB1 = basis.A2Bmap * evaluate(basis.pibasis, cfg)
println(@test isapprox(BB, BB1, rtol=1e-10))

Iz = findall(iszero, sum(norm, basis.A2Bmap, dims=1)[:])
if !isempty(Iz)
   @warn("The A2B map for EuclideanVector has $(length(Iz))/$(length(basis.pibasis)) zero-columns!!!!")
end

@info("Test equivariance properties")

tol = 1e-10

@info("check for rotation, permutation and inversion equivariance")
for ntest = 1:30
   Xs = rand(EuclideanVectorState, B1p.bases[1], nX)
   BB = evaluate(basis, ACEConfig(Xs))
   Q = rand([-1,1]) * ACE.Random.rand_rot()
   Xs_rot = Ref(Q) .* shuffle(Xs)
   BB_rot = evaluate(basis, ACEConfig(Xs_rot))
   print_tf(@test all([ norm(Q' * b1 - b2) < tol
                        for (b1, b2) in zip(BB_rot, BB)  ]))
end
println()

# ## keep for further profiling
#
# φ = ACE.EuclideanVector(Complex{Float64})
# pibasis = PIBasis(B1p, ord, maxdeg; property = φ, isreal = false)
# basis = SymmetricBasis(pibasis, φ)
# @time SymmetricBasis(pibasis, φ);
#
# Profile.clear(); # Profile.init(; delay = 0.0001)
# @profile SymmetricBasis(pibasis, φ);
# ProfileView.view()

##

@info(" ... derivatives")
for ntest = 1:30
   Us = randn(SVector{3, Float64}, length(Xs))
   C = randn(typeof(φ.val), length(basis))
   F = t -> sum( sum(c .* b.val)
                 for (c, b) in zip(C, ACE.evaluate(basis, ACEConfig(Xs + t[1] * Us))) )
   dF = t -> [ sum( sum(c .* db)
                    for (c, db) in zip(C, ACE.evaluate_d(basis, ACEConfig(Xs + t[1] * Us)) * Us) ) ]
   print_tf(@test fdtest(F, dF, [0.0], verbose=false))
end
println()

##

end