text:
"""
   Traits for the HARPS spectrograph.
    https://www.eso.org/sci/facilities/lasilla/instruments/harps.html
Author: Alex Wise and collaborators
Created: April 2021
"""

""" Delegates loading of code specifying types essential to the package.  """

import RvSpectMLBase: min_order, max_order, min_pixel_in_order, max_pixel_in_order, min_pixel, max_pixel

min_order(::HARPS2D) = 1
max_order(::HARPS2D) = 72

min_pixel_in_order(inst::HARPS2D) = 1
max_pixel_in_order(inst::HARPS2D) = 4096

min_pixel(::HARPS1D) = 1
max_pixel(::HARPS1D) = 300000 # TODO: Update once know size of HARPS's 1d extracted spectra

import RvSpectMLBase: bad_col_ranges

function bad_col_ranges(inst::HARPS2D, ord::Int)
    [0:0]#placeholder function to be updated if/when we decide certain pixels should be excluded from analysis
end

import RvSpectMLBase: orders_to_use_default, min_col_default, max_col_default

orders_to_use_default(::HARPS2D) = 1:72   # avoid last order due to oxygen tellurics, first 7 due to poor SNR

min_col_default(::HARPS2D, ord::Integer) = 300  # Avoid where continuum normalization effected by edges/scattered light
max_col_default(::HARPS2D, ord::Integer) = 3800  # Avoid where continuum normalization effected by edges/scattered light

import RvSpectMLBase: get_pixel_range
function get_pixel_range(inst::HARPS2D, ord::Integer)
    minc = max(min_col_default(inst, ord))
    maxc = min(max_col_default(inst, ord))
    return minc:maxc
end

import RvSpectMLBase: metadata_symbols_default, metadata_strings_default
# HARPS headers
metadata_symbols_default(::AnyHARPS) = Symbol[:bjd, :target, :exptime, :airmass, :ssb_rv_kmps, :wfile]
metadata_strings_default(::AnyHARPS) = String["MJD-OBS", "OBJECT", "EXPTIME", "AIRMASS", "ESO DRS BERV", "ESO DRS CAL TH FILE"]

import RvSpectMLBase: default_ccf_mask_v_width
default_ccf_mask_v_width(::AnyHARPS) = 820.0 #approx one HARPS pixel

import RvSpectMLBase: get_inst_module
get_inst_module(::AnyHARPS) = HARPS

#not sure what this does right now
import RvSpectMLBase: get_λ_range
function get_λ_range(data::CLT) where { T1<:Real, T2<:Real, T3<:Real, A1<:AbstractArray{T1,2}, A2<:AbstractArray{T2,2}, A3<:AbstractArray{T3,2},
                                       IT<:AnyHARPS, CLT<:Spectra2DBasic{T1,T2,T3,A1,A2,A3,IT} }
   (λmin, λmax) = NaNMath.extrema(data.λ)
   return (min=λmin, max=λmax)
end

default_λmin = 3800.0
default_λmax = 6900.0

max_stars_repo_name: RvSpectML/EchelleInstruments

max_issues_repo_name: RvSpectML/EchelleInstruments

max_forks_repo_name: RvSpectML/EchelleInstruments


text:
using BinaryAnalysis
using Base.Test

root = joinpath(JULIA_HOME,"julia")
dylibs = get_dltree(root)
@show dylibs

ext: jl


text:
export AffineAdd

"""
`AffineAdd(A::AbstractOperator, d, [sign = true])`

Affine addition to `AbstractOperator` with an array or scalar `d`.

Use `sign = false` to perform subtraction.

```julia
julia> A = AffineAdd(Sin(3),[1.;2.;3.])
sin+d  ℝ^3 -> ℝ^3

julia> A*[3.;4.;5.] == sin.([3.;4.;5.]).+[1.;2.;3.]
true

julia> A = AffineAdd(Exp(3),[1.;2.;3.],false)
e-d  ℝ^3 -> ℝ^3

julia> A*[3.;4.;5.] == exp.([3.;4.;5.]).-[1.;2.;3.]
true

```

"""
struct AffineAdd{L <: AbstractOperator, D <: Union{AbstractArray, Number}, S} <: AbstractOperator
  A::L
  d::D
  function AffineAdd(A::L, d::D, sign::Bool = true) where {L, D <: AbstractArray}
      if size(d) != size(A,1)
          throw(DimensionMismatch("codomain size of $A not compatible with array `d` of size $(size(d))"))
      end
      if eltype(d) != codomainType(A)
          error("cannot tilt opertor having codomain type $(codomainType(A)) with array of type $(eltype(d))")
      end
      new{L,D,sign}(A,d)
  end
  # scalar
  function AffineAdd(A::L, d::D, sign::Bool = true) where {L, D <: Number}
      if typeof(d) <: Complex && codomainType(A) <: Real
          error("cannot tilt opertor having codomain type $(codomainType(A)) with array of type $(eltype(d))")
      end
      new{L,D,sign}(A,d)
  end
end

# Mappings
# array
function mul!(y::DD, T::AffineAdd{L, D, true}, x) where {L <: AbstractOperator, DD, D}
    mul!(y,T.A,x)
    y .+= T.d
end

function mul!(y::DD, T::AffineAdd{L, D, false}, x) where {L <: AbstractOperator, DD, D}
    mul!(y,T.A,x)
    y .-= T.d
end

mul!(y, T::AdjointOperator{AffineAdd{L, D, S}}, x) where {L <: AbstractOperator, D, S} = mul!(y,T.A.A',x)

# Properties

size(L::AffineAdd) = size(L.A)

domainType(L::AffineAdd) = domainType(L.A)
codomainType(L::AffineAdd) = codomainType(L.A)

is_linear(L::AffineAdd) = is_linear(L.A)
is_null(L::AffineAdd) = is_null(L.A)
is_eye(L::AffineAdd) = is_diagonal(L.A)
is_diagonal(L::AffineAdd) = is_diagonal(L.A)
is_invertible(L::AffineAdd) = is_invertible(L.A)
is_AcA_diagonal(L::AffineAdd) = is_AcA_diagonal(L.A)
is_AAc_diagonal(L::AffineAdd) = is_AAc_diagonal(L.A)
is_full_row_rank(L::AffineAdd) = is_full_row_rank(L.A)
is_full_column_rank(L::AffineAdd) = is_full_column_rank(L.A)
is_sliced(L::AffineAdd) = is_sliced(L.A)

fun_name(T::AffineAdd{L,D,S}) where {L,D,S} = "$(fun_name(T.A))"*(S ? "+" : "-")*"d"
fun_type(L::AffineAdd) = fun_type(L.A)

diag(L::AffineAdd) = diag(L.A)
diag_AcA(L::AffineAdd) = diag_AcA(L.A)
diag_AAc(L::AffineAdd) = diag_AAc(L.A)

# utils
import Base: sign
sign(T::AffineAdd{L,D,false}) where {L,D} = -1
sign(T::AffineAdd{L,D, true}) where {L,D} =  1

function permute(T::AffineAdd{L,D,S}, p::AbstractVector{Int}) where {L,D,S}
    A = permute(T.A,p)
    return AffineAdd(A,T.d,S)
end

displacement(A::AffineAdd{L,D,true})  where {L,D} =  A.d .+ displacement(A.A)
displacement(A::AffineAdd{L,D,false}) where {L,D} = -A.d .+ displacement(A.A)

remove_displacement(A::AffineAdd) = remove_displacement(A.A)


text:
using SampleJuliaProj
using Test

@testset "SampleJuliaProj.jl" begin
    # Write your tests here.
    #1+1 == 2
end


text:
########################################################################
# Atoms of the Positive Simplex
########################################################################

"""
    PosSimplexAtom(n, z)

Atom in the positive simplex of dimension n.
"""
struct PosSimplexAtom <: AbstractAtom
    n::Int64
    z::Vector{Float64}
    function PosSimplexAtom(z)
        n = length(z)
        return new(n, z)
    end
end

"""
Multiply a Positive simplex atom by a linear map.
"""
Base.:(*)(M::AbstractLinearOp, a::PosSimplexAtom) = M*a.z
Base.vec(a::PosSimplexAtom) = a.z
Base.length(a::PosSimplexAtom) = a.n

########################################################################
# Atomic set for the positive simplex
########################################################################

"""
    PosSimplex(c, n, maxrank=n)

Atomic set defined by positive simplex in `n` variables.
The atomic set takes an optional parameters:

`maxrank` is the maximum dimension of the face exposed by a vector.
"""
struct PosSimplex <: AbstractAtomicSet
    c::Vector{Float64}
    n::Int64
    maxrank::Int64
    function PosSimplex(c, maxrank)
        n = length(c)
        minimum(c) > 0 || throw(DomainError(c,"c must be in non-negative orthant"))
        n ≥ maxrank ≥ 1 || throw(DomainError(maxrank,"maxrank must be ≥ 1"))
        return new(c, n, maxrank)
    end
end

function PosSimplex(c::Vector{Float64}; maxrank = length(findall(!iszero,c)) )
    index_c = findall(x -> x != 0, c)
    maxrank <= length(index_c) || throw(DomainError(maxrank,"maxrank should be less than sparstiy level of c"))
    return PosSimplex(c, maxrank)
end

PosSimplex(n::Int64; maxrank=n) = PosSimplex(ones(n), maxrank)

"""
    gauge(A::PosSimplex, x::Vector)

Gives the sum `dot(c, x)`.
"""
function gauge(A::PosSimplex, x::Vector)
    for i in eachindex(x)
        if (A.c[i] == 0 && x[i] != 0) || x[i] < 0
            return Inf
        end
    end
    return dot(A.c,x)
end
"""
    support(A::PosSimplex, z::Vector)

Gives the inf-norm of `c ⊙ z`.
"""
function support(A::PosSimplex, z::Vector)
    sup = 0.
    for i in eachindex(z)
        if A.c[i] != 0
            sup = max(sup, z[i]/A.c[i])
        end
    end

    return sup
end
"""
    expose(A::PosSimplex, z::Vector)

A non-overwriting version of [`expose!`](@ref).
"""
expose(A::PosSimplex, z; kwargs...) = expose!(A::PosSimplex, copy(z); kwargs...)

"""
    expose!(A::PosSimplex, z::Vector; tol=1e-12)

Obtain an atom in the face exposed by the vector `z`.
The vector `z` is overwritten. If `norm(z,Inf)<tol`, then
`z` is returned untouched.
"""

function expose!(A::PosSimplex, z::Vector; tol=1e-1)
    # z = Float64.(z)
    zsup = support(A, z)
    if maximum(z) < -1e-12
        return PosSimplexAtom(zero(z))
    end

    index_c = findall(x -> x != 0, A.c)
    index_c_comp = findall(x -> x == 0, A.c)

    if zsup < 1e-12
        nnz = 1
    else
        nnz = 0
    end

    # for i in eachindex(index_c_comp)
    #     z[index_c_comp[i]] = 0
    # end
    for i in eachindex(index_c)
        val = (A.c[index_c[i]]).^(-1) * z[index_c[i]]
        if abs(zsup - val) ≤ tol*zsup
            z[index_c[i]] = 1/A.c[index_c[i]]
            nnz += 1
        else
            z[index_c[i]] = zero(eltype(z))
        end
    end
    return PosSimplexAtom(z ./= nnz)
end

Base.length(A::PosSimplex) = A.n
rank(A::PosSimplex) = A.maxrank
atom_name(A::PosSimplex) = "Positive Simplex"
atom_description(A::PosSimplex) = "{ x ∈ ℝⁿ | ⟨c, x⟩ ≤ 1, x ≥ 0}"
atom_parameters(A::PosSimplex) = "c = cost vector, n = $(length(A)); maxrank = $(A.maxrank)"

########################################################################
# Face of the positive simplex.
########################################################################

struct PosSimplexFace{LM <: LinearMap} <: AbstractFace
    n::Int64
    k::Int64
    S::LM
end

PosSimplexFace(S::LinearMap) = PosSimplexFace(size(S)...,S)

"""
    face(A, z)

Return a face of the atomic set `A` exposed by the vector `z`. The dimension of
the exposed face is limited by `k=maxrank(A)`. If the dimension is being limited
by `k`, then the routine returns a set of rank `k` atoms that define the subset
of the exposed face.
"""
function face(A::PosSimplex, z::Vector; rTol=1e-1)

    t = support(A, z)
    a = expose(A, z, tol=rTol*t)
    idx = findall(!iszero, vec(a))

    k = length(idx)
    val = [(A.c[i]).^(-1) for i in idx]
    n = length(A)
    if support(A,z) < 1e-12
        S = sparse(idx, collect(1:k), val, n, k + 1)
    else
        S = sparse(idx, collect(1:k), val, n, k)
    end
    S = LinearMap(S)
    return PosSimplexFace(S)
end

"""
Given a face and a set of weights, reveal the corresponding
point on the face.
"""
Base.:(*)(F::PosSimplexFace, c::Vector) = F.S*c
Base.:(*)(M::AbstractLinearOp, F::PosSimplexFace) = M*F.S
Base.:(*)(λ::Real, F::PosSimplexFace) = λ*F.S
Base.length(F::PosSimplexFace) = F.n
rank(F::PosSimplexFace) = F.k
vec(F::PosSimplexFace) = F.S
face_name(F::PosSimplexFace) = "Face of positive simplex"
face_parameters(F::PosSimplexFace) = "rank = $(rank(F)); n = $(length(F))"


text:
# Statistics for interpreting actomyosin simulation results
# Alex Tam, 23/11/20

"Actin filament curvature"
function curvature(af, s, Lxx, Lxy, Lyx, Lyy)
    filament_curvatures = Vector{Float64}(); # Pre-allocate mean curvature data
    # Loop over filaments
    for j = 1:length(af)
        f::Actin_Filament = af[j]; # Extract current filament
        seg_lengths = get_segment_lengths(f, s, Lxx, Lxy, Lyx, Lyy); # Dimensional segment lengths
        total_filament_curvature = 0.0; # Pre-allocate integrated curvature of current filament
        # Compute curvature at interior nodes
        for i = 1:length(s.an[f.index])-2
            # Extract physical positions of plus node
            xp = s.an[f.index][i+2][1]*Lxx + s.an[f.index][i+2][2]*Lyx;
            yp = s.an[f.index][i+2][1]*Lxy + s.an[f.index][i+2][2]*Lyy;
            # Extract physical positions of centre node
            xc = s.an[f.index][i+1][1]*Lxx + s.an[f.index][i+1][2]*Lyx;
            yc = s.an[f.index][i+1][1]*Lxy + s.an[f.index][i+1][2]*Lyy;
            # Extract physical positions of minus node
            xm = s.an[f.index][i][1]*Lxx + s.an[f.index][i][2]*Lyx;
            ym = s.an[f.index][i][1]*Lxy + s.an[f.index][i][2]*Lyy;
            # Compute numerical second derivatives
            L01 = seg_lengths[i]; L12 = seg_lengths[i+1]; # Extract segment lengths
            avl = (L01 + L12)/2; # Average length
            ddfx = ((xp-xc)/L12 - (xc-xm)/L01)/avl; # Numerical second derivative (x)
            ddfy = ((yp-yc)/L12 - (yc-ym)/L01)/avl; # Numerical second derivative (y)
            # Add contribution to integrated curvature
            total_filament_curvature += sqrt(ddfx^2 + ddfy^2)*avl;
        end
        push!(filament_curvatures, total_filament_curvature);
    end
    # histogram(filament_curvatures)
    return mean(filament_curvatures)
end

"Two-filament index"
function two_filament_index(mm, s, Lxx, Lxy, Lyx, Lyy)
    index = Vector{Float64}(); # Pre-allocate index data
    theta::Float64 = 0.0; # Pre-allocate angle
    # Loop over motors
    for i = 1:length(mm)
        m::Myosin_Motor = mm[i]; # Extract current motor
        f1::Actin_Filament = m.f1; f2::Actin_Filament = m.f2; # Extract filaments attached to current motor
        # Calculate filament lengths
        L1 = sum(get_segment_lengths(f1, s, Lxx, Lxy, Lyx, Lyy));
        L2 = sum(get_segment_lengths(m.f2, s, Lxx, Lxy, Lyx, Lyy));
        # Calculate motor positions
        x1, y1, x2, y2 = get_motor_pos(m, s, Lxx, Lxy, Lyx, Lyy);
        m1x = (x1 - m.t1[1])*Lxx + (y1 - m.t1[2])*Lyx;
        m1y = (x1 - m.t1[1])*Lxy + (y1 - m.t1[2])*Lyy;
        m2x = (x2 - m.t2[1])*Lxx + (y2 - m.t2[2])*Lyx;
        m2y = (x2 - m.t2[1])*Lxy + (y2 - m.t2[2])*Lyy;
        # Find segments on which motor attaches
        seg1, seg2, r1, r2 = get_motor_relative_pos_segment(m, s, Lxx, Lxy, Lyx, Lyy);
        # Calculate segment positions
        mx1, my1, px1, py1 = get_segment_nodes(f1, f1.segments[seg1], s, m.t1, Lxx, Lxy, Lyx, Lyy);
        mx2, my2, px2, py2 = get_segment_nodes(f2, f2.segments[seg2], s, m.t2, Lxx, Lxy, Lyx, Lyy);
        # Calculate angle between vectors
        vec1 = [px1 - mx1, py1 - my1]; # Vector between motor and plus end of filament 1
        vec2 = [px2 - mx2, py2 - my2]; # Vector between motor and plus end of filament 2
        theta = acos( (vec1[1]*vec2[1] + vec1[2]*vec2[2])/( sqrt(vec1[1]^2 + vec1[2]^2)*sqrt(vec2[1]^2 + vec2[2]^2) ) ); # Angle between vectors
        # Compute index
        if all(s.mp[i] .<=1)
            Id = ( 2*(L1*s.mp[i][1] + L2*s.mp[i][2])/(L1+L2) - 1)*sin(theta/2)^2; # Heuristic index based on four filament branches
        else
            Id = 0; # Return zero if motor is detached
        end
        push!(index, Id);
    end
    # histogram(index)
    return mean(index), theta
end

"Actin filament speed"
function filament_speed(parN, af, s, s_old, Lxx, Lxy, Lyx, Lyy)
    node_speeds = Vector{Float64}(); # Pre-allocate vector of node speeds
    # Loop over filaments
    for f in af
        # Loop over nodes
        for i = 1:length(s.an[f.index])
            # Extract physical node position (new)
            xn = s.an[f.index][i][1]*Lxx + s.an[f.index][i][2]*Lyx;
            yn = s.an[f.index][i][1]*Lxy + s.an[f.index][i][2]*Lyy;
            # Extract physical node position (old). s_old can be used because it's computed after turnover
            xo = s_old.an[f.index][i][1]*Lxx + s_old.an[f.index][i][2]*Lyx;
            yo = s_old.an[f.index][i][1]*Lxy + s_old.an[f.index][i][2]*Lyy;
            # Calculate speed
            speed = sqrt((xn-xo)^2 + (yn-yo)^2)/parN.dt;
            push!(node_speeds, speed)
        end
    end
    # histogram(node_speeds)
    return mean(node_speeds)
end

"Myosin motor speed"
function motor_speed(parN, mm, s, s_old, Lxx, Lxy, Lyx, Lyy)
    motor_speeds = Vector{Float64}(); # Pre-allocate vector of motor head speeds
    # Loop over motors
    for m in mm
        # Calculate (current) filament lengths
        L1 = sum(get_segment_lengths(m.f1, s, Lxx, Lxy, Lyx, Lyy));
        L2 = sum(get_segment_lengths(m.f2, s, Lxx, Lxy, Lyx, Lyy));
        # Calculate motor speeds
        s1 = L1*(s.mp[m.index][1]-s_old.mp[m.index][1])/parN.dt;
        s2 = L2*(s.mp[m.index][2]-s_old.mp[m.index][2])/parN.dt;
        push!(motor_speeds, s1); push!(motor_speeds, s2);
    end
    # histogram(motor_speeds)
    return mean(motor_speeds)
end

"Paired distances between nodes"
function pcf(af, s, Lxx, Lyy)
    distances = Vector{Float64}(); # Pre-allocate vector of all paired distances
    # Loop over all filament pairs
    for i = 1:length(af)
        for j = 1:length(af)
            if j > i # Prevent double counting
                # Loop over all nodes
                for k = 1:length(s.an[i])
                    n1x = Lxx*(s.an[i][k][1] - floor(s.an[i][k][1])); # In-domain node position on 1st filament (x)
                    n1y = Lyy*(s.an[i][k][2] - floor(s.an[i][k][2])); # In-domain node position on 1st filament (y)
                    for l = 1:length(s.an[j])
                        n2x = Lxx*(s.an[j][l][1] - floor(s.an[j][l][1])); # In-domain node position on 2nd filament (x)
                        n2y = Lyy*(s.an[j][l][2] - floor(s.an[j][l][2])); # In-domain node position on 2nd filament (y)
                        dist_candidates = Vector{Float64}(); # Pre-allocate possible shortest distances given periodic BCs
                        test = 0; #[-1, 0, 1]; # Translations to test for shortest distance
                        # Loop over translations
                        for t1x in test
                            for t1y in test
                                for t2x in test
                                    for t2y in test
                                        push!(dist_candidates, sqrt( (n1x-n2x-t1x*Lxx+t2x*Lxx)^2 + (n1y-n2y-t1y*Lyy+t2y*Lyy)^2 ));
                                    end
                                end
                            end
                        end
                        push!(distances, minimum(dist_candidates)); # Store shortest distance between nodes
                    end
                end
            end
        end
    end
    histogram(distances)
end

"Compute integrated network statistics"
function integrated_statistics(parN, Curvature, Index)
    Curvature_Int = Vector{Float64}(); Index_Int = Vector{Float64}(); # Pre-allocate time-integrated variables
    times = (0:parN.nT-1).*parN.dt; # Vector of times at which we obtain measurements
    writedlm("times.csv", times);
    for i = 1:parN.nT
        push!(Curvature_Int, trapz(Curvature, times)); # Integrate force over time
        push!(Index_Int, trapz(Index, times)); # Integrate force over time
    end
    return Curvature_Int, Index_Int
end


text:
using LAMMPS
using LinearAlgebra:norm

"""
    SNAP_LAMMPS is a PotentialLearningProblem, it is based on the SNAP
    implementation of LAMMPS, which is accessed through LAMMPS.jl.
    Mathematical formulation: A. P. Thompson et al. (10.1016/j.jcp.2014.12.018)
"""
mutable struct SNAP_LAMMPS <: PotentialLearningProblem

    path::String                      # Path of the example folder
    β::Vector{Float64}                # Fitted parameters
    A::Matrix{Float64}                # Matrix of potentials, forces, and stresses.
    b::Vector{Float64}                # DFT training data - Reference data
    dft_data::Vector{Float64}         # DFT training data
    ref_data::Vector{Float64}         # Reference data
    no_atoms_per_type::Vector         # No. of atoms per type
    no_atoms_per_conf::Int64          # No. of atoms per atomic configurations
    no_atomic_conf::Int64             # No. of atomic configurations
    no_train_atomic_conf::Int64       # No. of trainning atomic configurations
    rcutfac::Float64                  # Scale factor applied to all cutoff radii (positive real)
    rfac0::Float64                    # Parameter in distance to angle conversion (0 < rcutfac < 1)
    twojmax::Int64                    # SNAP-LAMMPS parameter
    cutoff_radii::Vector{Float64}     # List of cutoff radii, one for each type
    neighbor_weights::Vector{Float64} # List of neighbor weights, one for each type
    ncoeff::Int64                     # ncoeff-1 = no. of param. associated to an atom type
    cols::Int64                       # No. of columns of matrix `A`
    numTypes::Int64                   # No. of atom types
    opt_pars::Dict                    # Optional parameters: rmin0, switchflag,
                                      # bzeroflag, quadraticflag, chem, chemflag,
                                      # bnormflag, wselfallflag, switchflag
    fit_forces::Bool                  # Enable/disable force fitting

end

"""
    SNAP_LAMMPS(dft_data::Vector{Float64}, ref_data::Vector{Float64}, params::Dict)

Creates a SNAP_LAMMPS.
"""
function SNAP_LAMMPS(dft_data::Vector{Float64}, ref_data::Vector{Float64}, params::Dict)

    β = []
    A = Matrix{Float64}(undef, 0, 0)
    b = dft_data - ref_data
    pars = params["global"]
    path = pars["path"]
    aux = pars["no_atoms_per_type"]
    no_atoms_per_type = length(aux) > 1 ? aux : [aux]
    no_atoms_per_conf = sum(no_atoms_per_type)
    no_atomic_conf = pars["no_atomic_conf"]
    no_train_atomic_conf = pars["no_train_atomic_conf"]
    rcutfac = pars["rcutfac"]
    rfac0 = pars["rfac0"]
    twojmax = pars["twojmax"]
    J = twojmax / 2.0
    ncoeff = round(Int, (J + 1) * (J + 2) * (( J + (1.5)) / 3. ) + 1)
    cols = 2 * ncoeff
    aux = pars["cutoff_radii"]
    cutoff_radii = length(aux) > 1 ? aux : [aux]
    aux = pars["neighbor_weights"]
    neighbor_weights = length(aux) > 1 ? aux : [aux]
    numTypes = pars["numTypes"]

    opt_pars = Dict()
    if haskey(pars, "rmin0") opt_pars["rmin0"] = pars["rmin0"] end
    if haskey(pars, "switchflag") opt_pars["switchflag"] = pars["switchflag"] end
    if haskey(pars, "bzeroflag") opt_pars["bzeroflag"] = pars["bzeroflag"] end
    if haskey(pars, "quadraticflag") opt_pars["quadraticflag"] = pars["quadraticflag"] end
    if haskey(pars, "chem") opt_pars["chem"] = pars["chem"] end
    if haskey(pars, "chemflag") opt_pars["chemflag"] = pars["chemflag"] end
    if haskey(pars, "bnormflag") opt_pars["bnormflag"] = pars["bnormflag"] end
    if haskey(pars, "wselfallflag") opt_pars["wselfallflag"] = pars["wselfallflag"] end
    if haskey(pars, "switchflag") opt_pars["switchflag"] = pars["switchflag"] end

    fit_forces = pars["fit_forces"]

    p = SNAP_LAMMPS( path, β, A, b, dft_data, ref_data, no_atoms_per_type, no_atoms_per_conf,
                     no_atomic_conf, no_train_atomic_conf, rcutfac, rfac0, twojmax,
                     cutoff_radii, neighbor_weights, ncoeff, cols, numTypes, opt_pars,
                     fit_forces)

    p.A = calc_A(p)

    return p
end

"""
    error(β::Vector{Float64}, s::SNAP_LAMMPS)

Error function needed to perform the learning process.
"""
function error(β::Vector{Float64}, s::SNAP_LAMMPS)
    return norm(s.A * β - s.b)
end

"""
    get_bispectrums(lmp, data_path::String, p::SNAP_LAMMPS)

Calculates the bispectrums components and its derivatives.
See https://docs.lammps.org/compute_sna_atom.html
"""
function get_bispectrums(lmp, data_path::String, p::SNAP_LAMMPS)

    args = ""
    # scale factor applied to all cutoff radii (positive real)
    args *= "$(p.rcutfac) "
    # parameter in distance to angle conversion (0 < rcutfac < 1)
    args *= "$(p.rfac0) "
    # band limit for bispectrum components (non-negative integer)
    args *= "$(p.twojmax) "
    # R_1, R_2,… = list of cutoff radii, one for each type (distance units)
    args *= join(["$r " for r in p.cutoff_radii])
    # w_1, w_2,… = list of neighbor weights, one for each type
    args *= join(["$r " for r in p.neighbor_weights])
    # zero or more keyword/value pairs may be appended
    # keyword = rmin0 or switchflag or bzeroflag or quadraticflag or chem or bnormflag or wselfallflag
    args *= haskey(p.opt_pars, "rmin0") ? "rmin0 $(p.opt_pars["rmin0"]) " : ""
    args *= haskey(p.opt_pars, "switchflag") ? "switchflag $(p.opt_pars["switchflag"]) " : ""
    args *= haskey(p.opt_pars, "bzeroflag") ? "bzeroflag $(p.opt_pars["bzeroflag"]) " : ""
    args *= haskey(p.opt_pars, "quadraticflag") ? "quadraticflag $(p.opt_pars["quadraticflag"]) " : ""
    args *= haskey(p.opt_pars, "chem") ? "chem " * join(["$r " for r in p.opt_pars["chem"]]) * " " : ""
    args *= haskey(p.opt_pars, "chemflag") ? "chemflag $(p.opt_pars["chemflag"]) " : ""
    args *= haskey(p.opt_pars, "bnormflag") ? "bnormflag $(p.opt_pars["bnormflag"]) " : ""
    args *= haskey(p.opt_pars, "wselfallflag") ? "wselfallflag $(p.opt_pars["wselfallflag"]) " : ""
    args *= haskey(p.opt_pars, "switchflag") ? "switchflag $(p.opt_pars["switchflag"]) " : ""

    read_data_str = "read_data " * data_path
    command(lmp, "log none")
    command(lmp, "units metal")
    command(lmp, "boundary p p p")
    command(lmp, "atom_style atomic")
    command(lmp, "atom_modify map array")
    command(lmp, read_data_str)
#    command(lmp, "pair_style hybrid/overlay zero 10.0 zbl 4.0 4.8")
#    command(lmp, "pair_coeff * * zero")
#    command(lmp, "pair_coeff * * zbl 73 73")
    command(lmp, "pair_style zero $(p.rcutfac)")
    command(lmp, "pair_coeff * *")
    command(lmp, "compute PE all pe")
    command(lmp, "compute S all pressure thermo_temp")
    command(lmp, "compute SNA all sna/atom " * args)
    command(lmp, "compute SNAD all snad/atom " * args)
    command(lmp, "compute SNAV all snav/atom " * args)
    command(lmp, "thermo_style custom pe")
    command(lmp, "run 0")

    # Extract bispectrum
    bs = extract_compute(lmp, "SNA",  LAMMPS.API.LMP_STYLE_ATOM,
                                      LAMMPS.API.LMP_TYPE_ARRAY)
    deriv_bs = extract_compute(lmp, "SNAD", LAMMPS.API.LMP_STYLE_ATOM,
                                            LAMMPS.API.LMP_TYPE_ARRAY)
    return bs, deriv_bs
end

"""
    calc_A(params::Dict, p::SNAP_LAMMPS)

Calculates the matrix `A` (See Eq. 13, 10.1016/j.jcp.2014.12.018).
This calculation requires accessing the SNAP implementation of LAMMPS.
"""
function calc_A(p::SNAP_LAMMPS)

    A = LMP(["-screen","none"]) do lmp

        A_potential = Array{Float64}(undef, p.no_train_atomic_conf, p.cols)

        A_forces = p.fit_forces ? Array{Float64}(undef, 3 * p.no_atoms_per_conf *
                                  p.no_train_atomic_conf, p.cols) : Array{Float64}(undef, 0, p.cols)

        for j in 1:p.no_train_atomic_conf
            data_path = joinpath(p.path, "DATA", string(j), "DATA")
            bs, deriv_bs = get_bispectrums(lmp, data_path, p)

            # Create a row of the potential energy block
            row = Vector{Float64}()
            for no_atoms in p.no_atoms_per_type
                push!(row, no_atoms)
                for k in 1:(p.ncoeff-1)
                    acc = 0.0
                    for n in 1:no_atoms
                        acc += bs[k, n]
                    end
                    push!(row, acc)
                end
            end
            A_potential[j, :] = row

            if p.fit_forces
                # Create a set of rows of the force block
                k = (j - 1) * 3 * p.no_atoms_per_conf + 1
                for n = 1:p.no_atoms_per_conf
                    for c = [1, 2, 3] # component x, y, z
                        row = Vector{Float64}()
                        for t = 1:p.numTypes # e.g. 2 : Ga and N
                            a = 1 + (p.ncoeff-1) * (c-1) + 3 * (p.ncoeff-1) * (t-1)
                            b = a + p.ncoeff - 2
                            row = [row; [0.0; deriv_bs[a:b, n]]]
                        end
                        A_forces[k, :] = row
                        k += 1
                    end
                end
            end

            command(lmp, "clear")
        end
        return [A_potential; A_forces]
    end

    return A
end

"""
    potential_energy(p::Potential, j::Int64)

Calculates the potential energy of a particular atomic configuration (j) using
the fitted parameters β.
This calculation requires accessing the SNAP implementation of LAMMPS.
"""
function potential_energy(p::PotentialLearningProblem, j::Int64)

    data_path = joinpath(p.path, "DATA", string(j), "DATA")
    lmp = LMP(["-screen","none"])
    bs, deriv_bs = get_bispectrums(lmp, data_path, p)

    E_tot_acc = 0.0
    for (i, no_atoms) in enumerate(p.no_atoms_per_type)
        for n in 1:no_atoms
            E_atom_acc = p.β[p.ncoeff*(i-1)+1]
            for k in p.ncoeff*(i-1)+2:i*p.ncoeff
                k2 = k - p.ncoeff * (i - 1) - 1
                E_atom_acc += p.β[k] * bs[k2, n]
            end
            E_tot_acc += E_atom_acc
        end
    end

    command(lmp, "clear")
    return E_tot_acc
end

"""
    forces(p::Potential, j::Int64)

Calculates the forces of a particular atomic configuration (j) using the
fitted parameters β.
This calculation requires accessing the SNAP implementation of LAMMPS.
See https://docs.lammps.org/compute_sna_atom.html#compute-snad-atom-command
"""
function forces(p::PotentialLearningProblem, j::Int64)

    data_path = joinpath(p.path, "DATA", string(j), "DATA")
    lmp = LMP(["-screen","none"])
    bs, deriv_bs = get_bispectrums(lmp, data_path, p)

    forces = Vector{Force}()

    for n = 1:p.no_atoms_per_conf
        f = [0.0, 0.0, 0.0]
        for c = [1, 2, 3] # component x, y, z
            deriv_bs_row = []
            for t = 1:p.numTypes
                a = 1 + (p.ncoeff-1) * (c-1) + 3 * (p.ncoeff-1) * (t-1)
                b = a + p.ncoeff - 2
                deriv_bs_row = [deriv_bs_row; [0.0; deriv_bs[a:b, n]]]
            end
            f[c] = sum(deriv_bs_row .* p.β)
        end
        push!(forces, Force(f))
    end

    command(lmp, "clear")
    return forces
end


text:
# ---
# title: 617. Merge Two Binary Trees
# id: problem617
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Tree
# link: <https://leetcode.com/problems/merge-two-binary-trees/description/>
# hidden: true
# ---
#
# Given two binary trees and imagine that when you put one of them to cover the
# other, some nodes of the two trees are overlapped while the others are not.
#
# You need to merge them into a new binary tree. The merge rule is that if two
# nodes overlap, then sum node values up as the new value of the merged node.
# Otherwise, the NOT null node will be used as the node of new tree.
#
# **Example 1:**
#
#
#
#     Input:
#     	Tree 1                     Tree 2
#               1                         2
#              / \                       / \
#             3   2                     1   3
#            /                           \   \
#           5                             4   7
#     Output:
#     Merged tree:
#     	     3
#     	    / \
#     	   4   5
#     	  / \   \
#     	 5   4   7
#
#
#
#
# **Note:** The merging process must start from the root nodes of both trees.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
Base.@kwdef struct RecordStateAction <: AbstractHook
    records::Any = VectorSATrajectory(; state = Vector{Float32})
end

function (h::RecordStateAction)(::PreActStage, policy, env, action)
    push!(h.records; state = copy(state(env)), action = action)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BC},
    ::Val{:CartPole},
    ::Nothing;
    seed = 123,
    save_dir = nothing,
)
    rng = StableRNG(seed)

    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(state(env)), length(action_space(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = BasicDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                batch_size = 32,
                min_replay_history = 100,
                loss_func = huber_loss,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularArraySARTTrajectory(
            capacity = 1000,
            state = Vector{Float32} => (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)
    hook = RecordStateAction()
    run(agent, env, stop_condition, hook)

    bc = BehaviorCloningPolicy(
        approximator = NeuralNetworkApproximator(
            model = Chain(
                Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                Dense(128, 128, relu; initW = glorot_uniform(rng)),
                Dense(128, na; initW = glorot_uniform(rng)),
            ) |> cpu,
            optimizer = ADAM(),
        ),
    )

    s = BatchSampler{(:state, :action)}(32;)

    for i in 1:300
        _, batch = s(hook.records)
        RLBase.update!(bc, batch)
    end

    description = """
    # Behavior Cloning with CartPole

    This experiment uses transitions during the experiment
    `JuliaRL_BasicDQN_CartPole` to train a behavior policy.
    """

    hook = ComposedHook(TotalRewardPerEpisode(), TimePerStep())

    Experiment(bc, env, StopAfterEpisode(100), hook, description)
end


text:
include("utils.jl")

@testset "Circle" begin
    M = Circle()
    @testset "Real Circle Basics" begin
        @test repr(M) == "Circle(ℝ)"
        @test representation_size(M) == ()
        @test manifold_dimension(M) == 1
        @test !is_manifold_point(M, 9.0)
        @test_throws DomainError is_manifold_point(M, 9.0, true)
        @test !is_tangent_vector(M, 9.0, 0.0)
        @test_throws DomainError is_tangent_vector(M, 9.0, 0.0, true)
        @test is_tangent_vector(M, 0.0, 0.0; check_base_point = false)
        @test get_coordinates(M, Ref(0.0), Ref(2.0), DefaultOrthonormalBasis())[] ≈ 2.0
        @test get_coordinates(
            M,
            Ref(0.0),
            Ref(2.0),
            DiagonalizingOrthonormalBasis(Ref(1.0)),
        )[] ≈ 2.0
        @test get_coordinates(
            M,
            Ref(0.0),
            Ref(-2.0),
            DiagonalizingOrthonormalBasis(Ref(1.0)),
        )[] ≈ -2.0
        @test get_coordinates(
            M,
            Ref(0.0),
            Ref(2.0),
            DiagonalizingOrthonormalBasis(Ref(-1.0)),
        )[] ≈ -2.0
        @test get_coordinates(
            M,
            Ref(0.0),
            Ref(-2.0),
            DiagonalizingOrthonormalBasis(Ref(-1.0)),
        )[] ≈ 2.0
        y = [0.0]
        get_coordinates!(M, y, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))
        @test y ≈ [2.0]
        @test get_vector(M, Ref(0.0), Ref(2.0), DefaultOrthonormalBasis())[] ≈ 2.0
        @test get_vector(M, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))[] ≈
              2.0
        @test get_vector(M, Ref(0.0), Ref(-2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))[] ≈
              -2.0
        @test get_vector(M, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(-1.0)))[] ≈
              -2.0
        @test get_vector(
            M,
            Ref(0.0),
            Ref(-2.0),
            DiagonalizingOrthonormalBasis(Ref(-1.0)),
        )[] ≈ 2.0
        @test number_of_coordinates(M, DiagonalizingOrthonormalBasis(Ref(-1.0))) == 1
        @test number_of_coordinates(M, DefaultOrthonormalBasis()) == 1
        @test flat(M, 0.0, FVector(TangentSpace, 1.0)) == FVector(CotangentSpace, 1.0)
        @test sharp(M, 0.0, FVector(CotangentSpace, 1.0)) == FVector(TangentSpace, 1.0)
        @test vector_transport_to(M, 0.0, 1.0, 1.0, ParallelTransport()) == 1.0
        @test retract(M, 0.0, 1.0) == exp(M, 0.0, 1.0)
        @test injectivity_radius(M) ≈ π
        @test injectivity_radius(M, Ref(-2.0)) ≈ π
        @test injectivity_radius(M, Ref(-2.0), ExponentialRetraction()) ≈ π
        @test injectivity_radius(M, ExponentialRetraction()) ≈ π
        @test mean(M, [-π / 2, 0.0, π]) ≈ -π / 2
        @test mean(M, [-π / 2, 0.0, π], [1.0, 1.0, 1.0]) == -π / 2
        z = project(M, 1.5 * π)
        z2 = [0.0]
        project!(M, z2, 1.5 * π)
        @test z2[1] == z
        @test project(M, z) == z
        @test project(M, 1.0, 2.0) == 2.0
    end
    TEST_STATIC_SIZED && @testset "Real Circle and static sized arrays" begin
        v = MVector(0.0)
        x = SVector(0.0)
        log!(M, v, x, SVector(π / 4))
        @test norm(M, x, v) ≈ π / 4
        @test is_tangent_vector(M, x, v)
        @test is_tangent_vector(M, [], v; check_base_point = false)
        @test project(M, 1.0) == 1.0
        x = MVector(0.0)
        project!(M, x, x)
        @test x == MVector(0.0)
        x .+= 2 * π
        project!(M, x, x)
        @test x == MVector(0.0)
        @test project(M, 0.0, 1.0) == 1.0
    end
    types = [Float64]
    TEST_FLOAT32 && push!(types, Float32)

    basis_types = (DefaultOrthonormalBasis(),)
    basis_types_real = (
        DefaultOrthonormalBasis(),
        DiagonalizingOrthonormalBasis(Ref(-1.0)),
        DiagonalizingOrthonormalBasis(Ref(1.0)),
    )
    for T in types
        @testset "Type $T" begin
            pts = convert.(Ref(T), [-π / 4, 0.0, π / 4])
            test_manifold(
                M,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_default_vector_transport = true,
                test_vee_hat = false,
                is_mutating = false,
            )
            ptsS = SVector.(pts)
            test_manifold(
                M,
                ptsS,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_default_vector_transport = true,
                vector_transport_methods = [
                    ParallelTransport(),
                    SchildsLadderTransport(),
                    PoleLadderTransport(),
                ],
                test_vee_hat = true,
                basis_types_vecs = basis_types_real,
                basis_types_to_from = basis_types_real,
            )
        end
    end
    Mc = Circle(ℂ)
    @testset "Complex Circle Basics" begin
        @test repr(Mc) == "Circle(ℂ)"
        @test representation_size(Mc) == ()
        @test manifold_dimension(Mc) == 1
        @test is_tangent_vector(Mc, 1im, 0.0)
        @test is_manifold_point(Mc, 1im)
        @test !is_manifold_point(Mc, 1 + 1im)
        @test_throws DomainError is_manifold_point(Mc, 1 + 1im, true)
        @test !is_tangent_vector(Mc, 1 + 1im, 0.0)
        @test_throws DomainError is_tangent_vector(Mc, 1 + 1im, 0.0, true)
        @test !is_tangent_vector(Mc, 1im, 2im)
        @test_throws DomainError is_tangent_vector(Mc, 1im, 2im, true)
        @test flat(Mc, 0.0 + 0.0im, FVector(TangentSpace, 1.0im)) ==
              FVector(CotangentSpace, 1.0im)
        @test sharp(Mc, 0.0 + 0.0im, FVector(CotangentSpace, 1.0im)) ==
              FVector(TangentSpace, 1.0im)
        @test norm(Mc, 1.0, log(Mc, 1.0, -1.0)) ≈ π
        @test is_tangent_vector(Mc, 1.0, log(Mc, 1.0, -1.0))
        v = MVector(0.0 + 0.0im)
        x = SVector(1.0 + 0.0im)
        log!(Mc, v, x, SVector(-1.0 + 0.0im))
        @test norm(Mc, SVector(1.0), v) ≈ π
        @test is_tangent_vector(Mc, x, v)
        @test project(Mc, 1.0) == 1.0
        project(Mc, 1 / sqrt(2.0) + 1 / sqrt(2.0) * im) ==
        1 / sqrt(2.0) + 1 / sqrt(2.0) * im
        x = MVector(1.0 + 0.0im)
        project!(Mc, x, x)
        @test x == MVector(1.0 + 0.0im)
        x .*= 2
        project!(Mc, x, x)
        @test x == MVector(1.0 + 0.0im)

        angles = map(x -> exp(x * im), [-π / 2, 0.0, π])
        @test mean(Mc, angles) ≈ exp(-π * im / 2)
        @test mean(Mc, angles, [1.0, 1.0, 1.0]) ≈ exp(-π * im / 2)
        @test_throws ErrorException mean(Mc, [-1.0 + 0im, 1.0 + 0im])
        @test_throws ErrorException mean(Mc, [-1.0 + 0im, 1.0 + 0im], [1.0, 1.0])
    end
    types = [Complex{Float64}]
    TEST_FLOAT32 && push!(types, Complex{Float32})

    for T in types
        @testset "Type $T" begin
            a = 1 / sqrt(2.0)
            pts = convert.(Ref(T), [a - a * im, 1 + 0im, a + a * im])
            test_manifold(
                Mc,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_default_vector_transport = true,
                is_mutating = false,
                test_vee_hat = false,
                exp_log_atol_multiplier = 2.0,
                is_tangent_atol_multiplier = 2.0,
            )
            ptsS = SVector.(pts)
            test_manifold(
                Mc,
                ptsS,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_default_vector_transport = true,
                test_vee_hat = true,
                exp_log_atol_multiplier = 2.0,
                is_tangent_atol_multiplier = 2.0,
                basis_types_vecs = basis_types,
                basis_types_to_from = basis_types,
            )
        end
    end
end


text:
module StatsMonitor

using Sockets
using Statistics
using Dates
using JSON3
using OrderedCollections
using AbstractTrees
using OnlineStats
using StatsBase
using RecipesBase

export Stats, fit!

include("stats.jl")

# #-----------------------------------------------------------------------------# Backend
# abstract type Backend end

# struct TerminalBackend <: Backend end

# #-----------------------------------------------------------------------------# functions
# const functions = OrderedDict(
#     1 => sum,
#     2 => mean,
#     3 => x -> quantile(x, [0, .25, .5, .75, 1])
# )

# #-----------------------------------------------------------------------------# Bucket
# struct Bucket{T}
#     f::Int
#     data::Vector{T}
# end
# calc(o::Bucket) = (value=functions[o.f](o.data), n=length(o.data), f=o.f)

# #-----------------------------------------------------------------------------# StatServer
# struct StatServer{B<:Backend}
#     buckets_int::Dict{Symbol, Bucket{Int}}
#     buckets_float::Dict{Symbol, Bucket{Float64}}
#     buckets_string::Dict{Symbol, Bucket{String}}
#     port::Int
#     backend::B
#     interval::Int
# end
# function StatServer(; port=8125, backend=TerminalBackend(), interval=5)
#     _b(T::Type) = Dict{Symbol, Bucket{T}}()
#     StatServer(_b(Int), _b(Float64), _b(String), port, backend, interval)
# end
# function Base.empty!(o::StatServer)
#     empty!(o.buckets_int)
#     empty!(o.buckets_float)
#     empty!(o.buckets_string)
# end

# function start(o::StatServer)
#     @info "Listening on port $(o.port)..."
#     server = listen(o.port)
#     @info "Creating Timer..."
#     timer = Timer(1, interval=o.interval) do x
#         flush(o)
#         empty!(o)
#     end
#     while true
#         sock = accept(server)
#         @async while isopen(sock)
#             s = readline(sock)
#             # try
#                 msg = Meta.parse(s)
#                 data = eval_msg(msg)
#                 push!(o, data)
#             # catch ex
#             #     @warn ex
#             # end
#         end
#     end
#     close(timer)
#     close(server)
# end

# function eval_msg(ex)
#     ok = ex.head === :tuple &&
#         ex.args[1] isa QuoteNode &&
#         ex.args[1].value isa Symbol &&
#         ex.args[2] isa Int &&
#         ex.args[3] isa Union{Int, Float64, String}
#     ok || error("Message has issues: $ex")
#     eval(ex)
# end

# function Base.push!(o::StatServer, data::Tuple)
#     id, function_id, value = data
#     b = getbucket(o, value)
#     haskey(b, id) ? push!(b[id].data, value) : (b[id] = Bucket(function_id, [value]))
# end

# getbucket(o::StatServer, ::Int) = o.buckets_int
# getbucket(o::StatServer, ::Float64) = o.buckets_float
# getbucket(o::StatServer, ::String) = o.buckets_string

# #-----------------------------------------------------------------------------# flush
# function flush(o::StatServer{TerminalBackend})
#     println("Output at time :$(now())")
#     for (k,v) in o.buckets_int
#         printstyled("  > $k: $(calc(v))\n", color=:green)
#     end
#     for (k, v) in o.buckets_float
#         printstyled("  > $k: $(calc(v))\n", color=:cyan)
#     end
#     for (k, v) in o.buckets_string
#         printstyled("  > $k: $(calc(v))\n", color=:yellow)
#     end
# end

end # module


text:
"""
    ScalarColumn(fmt, header, n)

Trace column for any kind of scalar, returned by the callback function
`n`. Alternatively, `n` can be a scalar set manually by the user.
"""
mutable struct ScalarColumn{R,N} <: TraceColumn
    fmt::FormatExpr
    header::String
    n::N
end

scalar_format(::Type{R}, signed) where R =
    FormatExpr("{1:$(signed ? '+' : "")10.3e}")

scalar_format(::Type{R}, signed) where {R<:Integer} =
    FormatExpr("{1:$(signed ? '+' : "")4d} ")

"""
    ScalarColumn(n, header, ::Type{R}[, signed=false]) where {R<:Real}

Construct a [`ScalarColumn`](@ref) for the callback function `n`, with
a format automatically generated depending on `R` is an integer or
not, and whether it can be signed.
"""
function ScalarColumn(n::Function, header::String, ::Type{R}, signed=false) where {R<:Real}
    fmt = scalar_format(R, signed)
    ScalarColumn{R,typeof(n)}(fmt, lpad(header, length(format(fmt, n(1)))), n)
end

"""
    ScalarColumn(n::R, header[, signed=false]) where {R<:Real}

Construct a [`ScalarColumn`](@ref) for the scalar value `n`, with
a format automatically generated depending on `R` is an integer or
not, and whether it can be signed.
"""
function ScalarColumn(n::R, header::String, signed=false) where {R<:Real}
    fmt = scalar_format(R, signed)
    ScalarColumn{R,R}(fmt, lpad(header, length(format(fmt, n))), n)
end

(c::ScalarColumn{R,<:Function})(i::Integer) where R = c.n(i)
(c::ScalarColumn{R})(i::Integer) where R = c.n

export ScalarColumn


text:
using CartesianDDM

dofs = (8, 12)
procs = (2, 3)
overs = (3, 2)

dims = CartesianDDMDimension.(dofs, procs, overs)
context = CartesianDDMContext(dims)

v = rand(prod(dofs))
dv = decompose(context, v)

dv[3] = 2
w = collect(dv)


text:
# test basic case, with mode plots

using Pseudospectra, Test

@testset "Basic" begin

    A = Pseudospectra.grcar(40)

    opts = Dict{Symbol,Any}()
    ps_data = new_matrix(A,opts)
    driver!(ps_data,opts,gs)
    @test iscomputed(ps_data)
    modeplot(ps_data,0,0.5+2.0im)
    modeplot(ps_data,1,0.5+2.0im)

end


text:
#!/usr/bin/julia

function start(n)
    for i = 1:n
        if i % 15 == 0
            print("Fizz-Buzz ")
        elseif i % 3 == 0
            print("Fizz ")
        elseif i % 5 == 0
            print("Buzz ")
        else
            print(i, " ")
        end
        if i % 10 == 0
            println()
        end
    end
end

@time start(100)


text:
"""
    quap(model::Turing.Model, args...; kwargs...)

Takes a Turing model and returns the auadratic (Laplace) approximation of the
posterior distribution.
To this end, it calculates the Hessian at the MAP solution and
returns a Gaussian distribution that approximates the true
posterior.

# Example
```julia
using Turing
@model height(heights) = begin
    μ ~ Normal(178, 20)
    σ ~ Uniform(0, 50)
    heights .~ Normal(μ, σ)
end
m = height(d2.height)
res = quap(m)
```
"""
function quap(model::Turing.Model, args...; kwargs...)
    opt = optimize(model, MAP(), args...; kwargs...)

    coef = opt.values.array
    var_cov_matrix = informationmatrix(opt)
    sym_var_cov_matrix = Symmetric(var_cov_matrix)  # lest MvNormal complains, loudly
    converged = Optim.converged(opt.optim_result)

    distr = if length(coef) == 1
        Normal(coef[1], √sym_var_cov_matrix[1])  # Normal expects stddev
    else
        MvNormal(coef, sym_var_cov_matrix)       # MvNormal expects variance matrix
    end

    params = StatsBase.params(model)
    params_tuple = tuple(Symbol.(params)...)

    (coef = NamedTuple{params_tuple}(coef), vcov = sym_var_cov_matrix, converged = converged,
        distr = distr, params = [params...])
end

function StatsBase.params(model::Turing.Model)
    nt = model |> Turing.VarInfo |> Turing.tonamedtuple
    p = String[]
    for (a, v) in nt
        arr = a[1]
        var = v[1]
        if length(arr) != 1
            append!(p, ["$var[$i]" for i in 1:length(arr)])
        else
            push!(p, var)
        end
    end
    return p
end


text:
struct Pending end
struct Rescinded end
struct Waiting end

const OfferFlags = Union{Pending,Rescinded,Waiting}
const OfferState{T} = Union{Pending,Rescinded,Waiting,T}

struct Offer{T,State<:Reagents.Ref{OfferState{T}}}
    state::State
    task::Base.Task
end

Offer{T}() where {T} =
    Offer(Reagents.Ref{OfferState{T}}(Pending()), current_task())::Offer{T}

offerid(offer::Offer) = objectid(offer.state)
offerid(::Nothing) = UInt(0)

struct Reaction
    caslist::ImmutableList{CAS}
    offers::ImmutableList{Offer}
    postcommithooks::ImmutableList{Any}
end

Reaction() = Reaction(nothing, nothing, nothing)

hascas(rx::Reaction) = rx.caslist !== nothing

uintptr(x) = UInt(pointer_from_objref(x.ref))

withcas(rx::Reaction, cas::CAS) =
    Reaction(pushsortedby(uintptr, rx.caslist, cas), rx.offers, rx.postcommithooks)

withoffer(rx::Reaction, offer::Offer) =
    Reaction(rx.caslist, pushfirst(rx.offers, offer), rx.postcommithooks)

withpostcommit(rx::Reaction, @nospecialize(f)) =
    Reaction(rx.caslist, rx.offers, pushfirst(rx.postcommithooks, f))

combine(rx1::Reaction, rx2::Reaction) = Reaction(
    combinesortedby(uintptr, rx1.caslist, rx2.caslist),
    combine(rx1.offers, rx2.offers),
    combine(rx1.postcommithooks, rx2.postcommithooks),
)

setcasing!(::Nothing) = true
function setcasing!(list::ImmutableListNode{CAS})
    (; ref, old) = list.value
    (_, success) = @atomicreplace ref.value old => CASing()
    success || return false
    if setcasing!(list.tail)
        return true
    else
        @atomic ref.value = old
        return false
    end
end

setall!(::Nothing) = true
function setall!(list::ImmutableListNode{CAS})
    (; ref, new) = list.value
    @atomic ref.value = new
    return setall!(list.tail)
end

commit!(rx::Reaction) = setcasing!(rx.caslist) && setall!(rx.caslist)

struct Commit end

hascas(::Commit) = false

const Reactable = Union{Reactor,Commit}

maysync(actr::Reactor) = maysync(actr.reagent) || maysync(actr.continuation)
maysync(::Commit) = false

maysync(::Reagent) = false


text:
seed!(12345)

function testCategorical1(n::Int64, p::Vector{Float64})
  vs::Vector{Int64} = sampleCategorical(n, p)
  counts = Vector{Int64}(undef, length(p))
  fill!(counts, 0)
  for i in 1:n
    v::Int64 = vs[i]
    @inbounds counts[vs[i]] += 1
  end
  return testGOFMultinomial(p, counts)
end

function testCategorical2(n::Int64, p::Vector{Float64})
  counts = Vector{Int64}(undef, length(p))
  fill!(counts, 0)
  for i in 1:n
    v::Int64 = sampleCategorical(p)
    @inbounds counts[v] += 1
  end
  return testGOFMultinomial(p, counts)
end

@test testCategorical1(100, [0.5; 0.1; 0.25; 0.15]) > 0.01
@test testCategorical1(10000, [0.5; 0.1; 0.25; 0.15]) > 0.01

@test testCategorical2(100, [0.5; 0.1; 0.25; 0.15]) > 0.01
@test testCategorical2(10000, [0.5; 0.1; 0.25; 0.15]) > 0.01


text:
function getTableauHSPARK(s, σ, o, tsym, g, h, lq, lp, ω, d=Nothing)
    # α_q_1 and α_p_2 need to be conjugate symplectic
    # α_q_1 and α_p_3 need to be conjugate symplectic
    # α_q_2 and α_p_2 need to be conjugate symplectic
    # α_q_2 and α_p_3 need to be conjugate symplectic
    # -> α_q_1 and α_q_2 need to be identical
    # -> α_p_2 and α_p_3 need to be identical

    α_q_2 = lq.a
    α_q_3 = lq.a
    α_p_2 = lp.a
    α_p_3 = lp.a

    b_q_2 = lq.b
    b_q_3 = lq.b
    b_p_2 = lp.b
    b_p_3 = lp.b

    # α_p_1 is free, but determines a_q_1 and a_q_2
    α_q_1 = h.a
    α_p_1 = h.a

    # a_p_1, a_p_2 and a_p_3 are free (they are not used in HSPARKpSecondary integrators anymore)
    # REALLY?
    a_q_1 = g.a
    b_q_1 = g.b

    a_p_1 = g.a
    b_p_1 = g.b

    a_p_2 = g.a
    a_p_3 = g.a

    β_q_1 = lq.b
    β_q_2 = lq.b
    β_q_3 = lq.b

    β_p_1 = lp.b
    β_p_2 = lp.b
    β_p_3 = lp.b

    # a_q_1_ij = b_q_1_j / b_p_1_i * (b_p_1_i - α_p_1_ji)
    # a_q_1 = zeros(s,σ)
    # for i in 1:s
    #     for j in 1:σ
    #         a_q_1[i,j] = b_q_1[j] / b_p_1[i] * (b_p_1[i] - α_p_1[j,i])
    #     end
    # end

    # a_q_1_ij = b_q_1_j / b_p_1_i * (b_p_1_i - α_p_1_ji)
    # TODO check this
    a_q_2 = zeros(s,σ)
    for i in 1:s
        for j in 1:σ
            a_q_2[i,j] = b_q_2[j] / b_p_1[i] * (b_p_1[i] - α_p_2[j,i])
        end
    end

    # a_q_2_ij = b_q_2_j / b_p_1_i * (b_p_1_i - α_p_1_ji)
    # TODO check this
    a_q_3 = zeros(s,σ)
    for i in 1:s
        for j in 1:σ
            a_q_3[i,j] = b_q_3[j] / b_p_1[i] * (b_p_1[i] - α_p_3[j,i])
        end
    end

    a_q = (a_q_1, a_q_2, a_q_3)
    b_q = (b_q_1, b_q_2, b_q_3)
    c_q = g.c

    a_p = (a_p_1, a_p_2, a_p_3)
    b_p = (b_p_1, b_p_2, b_p_3)
    c_p = g.c

    α_q = (α_q_1, α_q_2, α_q_3)
    β_q = (β_q_1, β_q_2, β_q_3)
    γ_q = lq.c

    α_p = (α_p_1, α_p_2, α_p_3)
    β_p = (β_p_1, β_p_2, β_p_3)
    γ_p = lp.c

    coeff_q = CoefficientsSPARK(tsym, o, s, σ, a_q, b_q, c_q)
    coeff_p = CoefficientsSPARK(tsym, o, s, σ, a_p, b_p, c_p)
    coeff_q̃ = CoefficientsSPARK(tsym, o, σ, s, α_q, β_q, γ_q)
    coeff_p̃ = CoefficientsSPARK(tsym, o, σ, s, α_p, β_p, γ_p)

    TableauHSPARKsecondary(tsym, o, s, σ, coeff_q, coeff_p, coeff_q̃, coeff_p̃, ω, d)
end

function TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIII"))
    o = 2s-2
    getTableauHSPARK(s, s, o, name, lq, lp, lq, lp, get_lobatto_ω_matrix(s), get_lobatto_nullvector(s))
end

function TableauHSPARKLobattoIIIAB(s)
    lq = TableauLobattoIIIA(s)
    lp = TableauLobattoIIIB(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIIAIIIB"))
end

function TableauHSPARKLobattoIIIBA(s)
    lq = TableauLobattoIIIB(s)
    lp = TableauLobattoIIIA(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIIBIIIA"))
end

function TableauHSPARKLobattoIIICC̄(s)
    lq = TableauLobattoIIIC(s)
    lp = TableauLobattoIIIC̄(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIICIIIC̄"))
end

function TableauHSPARKLobattoIIIC̄C(s)
    lq = TableauLobattoIIIC̄(s)
    lp = TableauLobattoIIIC(s)
    TableauHSPARKLobattoIII(s, lq, lp; name = Symbol("HSPARKLobattoIIIC̄IIIC"))
end

function TableauHSPARKLobattoIIID(s)
    l = TableauLobattoIIID(s)
    TableauHSPARKLobattoIII(s, l, l; name = Symbol("HSPARKLobattoIIID"))
end

function TableauHSPARKLobattoIIIE(s)
    l = TableauLobattoIIIE(s)
    TableauHSPARKLobattoIII(s, l, l; name = Symbol("HSPARKLobattoIIIE"))
end

function TableauHSPARKGLRKLobattoIII(s, σ, lq, lp; name = Symbol("HSPARKGLRKLobattoIII"))
    o = 2s
    g = TableauGauss(s)
    getTableauHSPARK(s, σ, o, name, g, get_lobatto_glrk_coefficients(s, σ), lq, lp, get_GLRK_ω_matrix(σ), get_lobatto_nullvector(σ))
end

function TableauHSPARKGLRKLobattoIIIAB(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIA(σ), TableauLobattoIIIB(σ); name = Symbol("HSPARKGLRKLobattoIIIAIIIB"))
end

function TableauHSPARKGLRKLobattoIIIBA(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIB(σ), TableauLobattoIIIA(σ); name = Symbol("HSPARKGLRKLobattoIIIBIIIA"))
end

function TableauHSPARKGLRKLobattoIIICC̄(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIC(σ), TableauLobattoIIIC̄(σ); name = Symbol("HSPARKGLRKLobattoIIICIIIC̄"))
end

function TableauHSPARKGLRKLobattoIIIC̄C(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIC̄(σ), TableauLobattoIIIC(σ); name = Symbol("HSPARKGLRKLobattoIIIC̄IIIC"))
end

function TableauHSPARKGLRKLobattoIIID(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIID(σ), TableauLobattoIIID(σ); name = Symbol("HSPARKGLRKLobattoIIID"))
end

function TableauHSPARKGLRKLobattoIIIE(s, σ=s+1)
    TableauHSPARKGLRKLobattoIII(s, σ, TableauLobattoIIIE(σ), TableauLobattoIIIE(σ); name = Symbol("HSPARKGLRKLobattoIIIE"))
end


text:
using Example2018
using Test

# Polynomials form a ring

# 1. addition is associative
# 2. additive identity
# 3. additive inverse
# 4. addition commutes
# 5. multipilcation is ... (not code this part on the code)

#@testset "Polynomials form a ring" begin
#
#    z = Polynomial([0])
#
#    for n in 1:100
#
#        deg = rand(0:100)
#        p = Polynomial(rand(-10000000:10000000,deg+1))
#
#        @test p + z == p
#        @test z + p == p
#        @test p + (-p) == z
#    end
#
#end