# This file is a part of Julia. License is MIT: https://julialang.org/license

module XoshiroSimd
# Getting the xoroshiro RNG to reliably vectorize is somewhat of a hassle without Simd.jl.
import ..Random: TaskLocalRNG, rand, rand!, Xoshiro, CloseOpen01, UnsafeView,
                 SamplerType, SamplerTrivial
using Base: BitInteger_types
using Core.Intrinsics: llvmcall

# Vector-width. Influences random stream.
xoshiroWidth() = Val(8)
# Simd threshold. Influences random stream.
simdThreshold(::Type{T}) where T = 64
simdThreshold(::Type{Bool}) = 640

@inline _rotl45(x::UInt64) = (x<<45)|(x>>19)
@inline _shl17(x::UInt64) = x<<17
@inline _rotl23(x::UInt64) = (x<<23)|(x>>41)
@inline _plus(x::UInt64,y::UInt64) = x+y
@inline _xor(x::UInt64,y::UInt64) = xor(x,y)
@inline _and(x::UInt64, y::UInt64) = x & y
@inline _or(x::UInt64, y::UInt64) = x | y
@inline _lshr(x, y::Int32) = _lshr(x, y % Int64)
@inline _lshr(x::UInt64, y::Int64) = llvmcall("""
    %res = lshr i64 %0, %1
    ret i64 %res
    """,
    UInt64,
    Tuple{UInt64, Int64},
    x, y)

@inline _bits2float(x::UInt64, ::Type{Float64}) = reinterpret(UInt64, Float64(x >>> 11) * 0x1.0p-53)
@inline function _bits2float(x::UInt64, ::Type{Float32})
    #=
    # this implementation uses more high bits, but is harder to vectorize
    x = x >>> 16  # discard low 16 bits
    u = Float32(x >>> 24) * Float32(0x1.0p-24)
    l = Float32(x & 0x00ffffff) * Float32(0x1.0p-24)
    =#
    ui = (x>>>32) % UInt32
    li = x % UInt32
    u = Float32(ui >>> 8) * Float32(0x1.0p-24)
    l = Float32(li >>> 8) * Float32(0x1.0p-24)
    (UInt64(reinterpret(UInt32, u)) << 32) | UInt64(reinterpret(UInt32, l))
end

# required operations. These could be written more concisely with `ntuple`, but the compiler
# sometimes refuses to properly vectorize.
for N in [4,8,16]
    let code, s, fshl = "llvm.fshl.v$(N)i64",
        VT = :(NTuple{$N, VecElement{UInt64}})

        s = ntuple(_->VecElement(UInt64(45)), N)
        @eval @inline _rotl45(x::$VT) = ccall($fshl, llvmcall, $VT, ($VT, $VT, $VT), x, x, $s)

        s = ntuple(_->VecElement(UInt64(23)), N)
        @eval @inline _rotl23(x::$VT) = ccall($fshl, llvmcall, $VT, ($VT, $VT, $VT), x, x, $s)

        code = """
        %lshiftOp = shufflevector <1 x i64> <i64 17>, <1 x i64> undef, <$N x i32> zeroinitializer
        %res = shl <$N x i64> %0, %lshiftOp
        ret <$N x i64> %res
        """
        @eval @inline _shl17(x::$VT) = llvmcall($code, $VT, Tuple{$VT}, x)

        code = """
        %res = add <$N x i64> %1, %0
        ret <$N x i64> %res
        """
        @eval @inline _plus(x::$VT, y::$VT) = llvmcall($code, $VT, Tuple{$VT, $VT}, x, y)

        code = """
        %res = xor <$N x i64> %1, %0
        ret <$N x i64> %res
        """
        @eval @inline _xor(x::$VT, y::$VT) = llvmcall($code, $VT, Tuple{$VT, $VT}, x, y)

        code = """
        %res = and <$N x i64> %1, %0
        ret <$N x i64> %res
        """
        @eval @inline _and(x::$VT, y::$VT) = llvmcall($code, $VT, Tuple{$VT, $VT}, x, y)

        code = """
        %res = or <$N x i64> %1, %0
        ret <$N x i64> %res
        """
        @eval @inline _or(x::$VT, y::$VT) = llvmcall($code, $VT, Tuple{$VT, $VT}, x, y)

        code = """
        %tmp = insertelement <1 x i64> undef, i64 %1, i32 0
        %shift = shufflevector <1 x i64> %tmp, <1 x i64> %tmp, <$N x i32> zeroinitializer
        %res = lshr <$N x i64> %0, %shift
        ret <$N x i64> %res
        """
        @eval @inline _lshr(x::$VT, y::Int64) = llvmcall($code, $VT, Tuple{$VT, Int64}, x, y)

        code = """
        %shiftamt = shufflevector <1 x i64> <i64 11>, <1 x i64> undef, <$N x i32> zeroinitializer
        %sh = lshr <$N x i64> %0, %shiftamt
        %f = uitofp <$N x i64> %sh to <$N x double>
        %scale = shufflevector <1 x double> <double 0x3ca0000000000000>, <1 x double> undef, <$N x i32> zeroinitializer
        %m = fmul <$N x double> %f, %scale
        %i = bitcast <$N x double> %m to <$N x i64>
        ret <$N x i64> %i
        """
        @eval @inline _bits2float(x::$VT, ::Type{Float64}) = llvmcall($code, $VT, Tuple{$VT}, x)

        code = """
        %as32 = bitcast <$N x i64> %0 to <$(2N) x i32>
        %shiftamt = shufflevector <1 x i32> <i32 8>, <1 x i32> undef, <$(2N) x i32> zeroinitializer
        %sh = lshr <$(2N) x i32> %as32, %shiftamt
        %f = uitofp <$(2N) x i32> %sh to <$(2N) x float>
        %scale = shufflevector <1 x float> <float 0x3e70000000000000>, <1 x float> undef, <$(2N) x i32> zeroinitializer
        %m = fmul <$(2N) x float> %f, %scale
        %i = bitcast <$(2N) x float> %m to <$N x i64>
        ret <$N x i64> %i
        """
        @eval @inline _bits2float(x::$VT, ::Type{Float32}) = llvmcall($code, $VT, Tuple{$VT}, x)
    end
end

function forkRand(rng::Union{TaskLocalRNG, Xoshiro}, ::Val{N}) where N
    # constants have nothing up their sleeve. For more discussion, cf rng_split in task.c
    # 0x02011ce34bce797f == hash(UInt(1))|0x01
    # 0x5a94851fb48a6e05 == hash(UInt(2))|0x01
    # 0x3688cf5d48899fa7 == hash(UInt(3))|0x01
    # 0x867b4bb4c42e5661 == hash(UInt(4))|0x01
    s0 = ntuple(i->VecElement(0x02011ce34bce797f * rand(rng, UInt64)), Val(N))
    s1 = ntuple(i->VecElement(0x5a94851fb48a6e05 * rand(rng, UInt64)), Val(N))
    s2 = ntuple(i->VecElement(0x3688cf5d48899fa7 * rand(rng, UInt64)), Val(N))
    s3 = ntuple(i->VecElement(0x867b4bb4c42e5661 * rand(rng, UInt64)), Val(N))
    (s0, s1, s2, s3)
end

_id(x, T) = x

@inline function xoshiro_bulk(rng::Union{TaskLocalRNG, Xoshiro}, dst::Ptr{UInt8}, len::Int, T::Union{Type{UInt8}, Type{Bool}, Type{Float32}, Type{Float64}}, ::Val{N}, f::F = _id) where {N, F}
    if len >= simdThreshold(T)
        written = xoshiro_bulk_simd(rng, dst, len, T, Val(N), f)
        len -= written
        dst += written
    end
    if len != 0
        xoshiro_bulk_nosimd(rng, dst, len, T, f)
    end
    nothing
end

@noinline function xoshiro_bulk_nosimd(rng::Union{TaskLocalRNG, Xoshiro}, dst::Ptr{UInt8}, len::Int, ::Type{T}, f::F) where {T, F}
    if rng isa TaskLocalRNG
        task = current_task()
        s0, s1, s2, s3 = task.rngState0, task.rngState1, task.rngState2, task.rngState3
    else
        (; s0, s1, s2, s3) = rng::Xoshiro
    end

    i = 0
    while i+8 <= len
        res = _plus(_rotl23(_plus(s0,s3)),s0)
        unsafe_store!(reinterpret(Ptr{UInt64}, dst + i), f(res, T))
        t = _shl17(s1)
        s2 = _xor(s2, s0)
        s3 = _xor(s3, s1)
        s1 = _xor(s1, s2)
        s0 = _xor(s0, s3)
        s2 = _xor(s2, t)
        s3 = _rotl45(s3)
        i += 8
    end
    if i < len
        res = _plus(_rotl23(_plus(s0,s3)),s0)
        t = _shl17(s1)
        s2 = _xor(s2, s0)
        s3 = _xor(s3, s1)
        s1 = _xor(s1, s2)
        s0 = _xor(s0, s3)
        s2 = _xor(s2, t)
        s3 = _rotl45(s3)
        ref = Ref(f(res, T))
        # TODO: This may make the random-stream dependent on system endianness
        ccall(:memcpy, Ptr{Cvoid}, (Ptr{UInt8}, Ptr{UInt64}, Csize_t), dst+i, ref, len-i)
    end
    if rng isa TaskLocalRNG
        task.rngState0, task.rngState1, task.rngState2, task.rngState3 = s0, s1, s2, s3
    else
       rng.s0, rng.s1, rng.s2, rng.s3 =  s0, s1, s2, s3
    end
    nothing
end

@noinline function xoshiro_bulk_nosimd(rng::Union{TaskLocalRNG, Xoshiro}, dst::Ptr{UInt8}, len::Int, ::Type{Bool}, f)
    if rng isa TaskLocalRNG
        task = current_task()
        s0, s1, s2, s3 = task.rngState0, task.rngState1, task.rngState2, task.rngState3
    else
        (; s0, s1, s2, s3) = rng::Xoshiro
    end

    i = 0
    while i+8 <= len
        res = _plus(_rotl23(_plus(s0,s3)),s0)
        shift = 0
        while i+8 <= len && shift < 8
            resLoc = _and(_lshr(res, shift), 0x0101010101010101)
            unsafe_store!(reinterpret(Ptr{UInt64}, dst + i), resLoc)
            i += 8
            shift += 1
        end

        t = _shl17(s1)
        s2 = _xor(s2, s0)
        s3 = _xor(s3, s1)
        s1 = _xor(s1, s2)
        s0 = _xor(s0, s3)
        s2 = _xor(s2, t)
        s3 = _rotl45(s3)
    end
    if i < len
        # we may overgenerate some bytes here, if len mod 64 <= 56 and len mod 8 != 0
        res = _plus(_rotl23(_plus(s0,s3)),s0)
        resLoc = _and(res, 0x0101010101010101)
        ref = Ref(resLoc)
        ccall(:memcpy, Ptr{Cvoid}, (Ptr{UInt8}, Ptr{UInt64}, Csize_t), dst+i, ref, len-i)
        t = _shl17(s1)
        s2 = _xor(s2, s0)
        s3 = _xor(s3, s1)
        s1 = _xor(s1, s2)
        s0 = _xor(s0, s3)
        s2 = _xor(s2, t)
        s3 = _rotl45(s3)
    end
    if rng isa TaskLocalRNG
        task.rngState0, task.rngState1, task.rngState2, task.rngState3 = s0, s1, s2, s3
    else
        rng.s0, rng.s1, rng.s2, rng.s3 = s0, s1, s2, s3
    end
    nothing
end

@noinline function xoshiro_bulk_simd(rng::Union{TaskLocalRNG, Xoshiro}, dst::Ptr{UInt8}, len::Int, ::Type{T}, ::Val{N}, f::F) where {T,N,F}
    s0, s1, s2, s3 = forkRand(rng, Val(N))

    i = 0
    while i + 8*N <= len
        res = _plus(_rotl23(_plus(s0,s3)),s0)
        t = _shl17(s1)
        s2 = _xor(s2, s0)
        s3 = _xor(s3, s1)
        s1 = _xor(s1, s2)
        s0 = _xor(s0, s3)
        s2 = _xor(s2, t)
        s3 = _rotl45(s3)
        unsafe_store!(reinterpret(Ptr{NTuple{N,VecElement{UInt64}}}, dst + i), f(res, T))
        i += 8*N
    end
    return i
end

@noinline function xoshiro_bulk_simd(rng::Union{TaskLocalRNG, Xoshiro}, dst::Ptr{UInt8}, len::Int, ::Type{Bool}, ::Val{N}, f) where {N}
    s0, s1, s2, s3 = forkRand(rng, Val(N))
    msk = ntuple(i->VecElement(0x0101010101010101), Val(N))
    i = 0
    while i + 64*N <= len
        res = _plus(_rotl23(_plus(s0,s3)),s0)
        t = _shl17(s1)
        s2 = _xor(s2, s0)
        s3 = _xor(s3, s1)
        s1 = _xor(s1, s2)
        s0 = _xor(s0, s3)
        s2 = _xor(s2, t)
        s3 = _rotl45(s3)
        for k=0:7
            tmp = _lshr(res, k)
            toWrite = _and(tmp, msk)
            unsafe_store!(reinterpret(Ptr{NTuple{N,VecElement{UInt64}}}, dst + i + k*N*8), toWrite)
        end
        i += 64*N
    end
    return i
end

function rand!(rng::Union{TaskLocalRNG, Xoshiro}, dst::Array{Float32}, ::SamplerTrivial{CloseOpen01{Float32}})
    GC.@preserve dst xoshiro_bulk(rng, convert(Ptr{UInt8}, pointer(dst)), length(dst)*4, Float32, xoshiroWidth(), _bits2float)
    dst
end

function rand!(rng::Union{TaskLocalRNG, Xoshiro}, dst::Array{Float64}, ::SamplerTrivial{CloseOpen01{Float64}})
    GC.@preserve dst xoshiro_bulk(rng, convert(Ptr{UInt8}, pointer(dst)), length(dst)*8, Float64, xoshiroWidth(), _bits2float)
    dst
end

for T in BitInteger_types
    @eval function rand!(rng::Union{TaskLocalRNG, Xoshiro}, dst::Union{Array{$T}, UnsafeView{$T}}, ::SamplerType{$T})
        GC.@preserve dst xoshiro_bulk(rng, convert(Ptr{UInt8}, pointer(dst)), length(dst)*sizeof($T), UInt8, xoshiroWidth())
        dst
    end
end

function rand!(rng::Union{TaskLocalRNG, Xoshiro}, dst::Array{Bool}, ::SamplerType{Bool})
    GC.@preserve dst xoshiro_bulk(rng, convert(Ptr{UInt8}, pointer(dst)), length(dst), Bool, xoshiroWidth())
    dst
end

end # module


abstract PCModification

abstract Rotation <: PCModification
abstract Reflection <: PCModification

abstract CWRotation <: Rotation
abstract counterCWRotation <: Rotation

abstract NoModification <: PCModification

#######################################################
## resolving pdf and cdf arguments for modifications ##
#######################################################

function resolve_args(mod::PCModification,
                      u1::FloatVec, u2::FloatVec)
    error("pdf modification still needs to be implemented for this
    type of copula modification")
end

function resolve_args(mod::PCModification, u1::Float64, u2::Float64)
    return resolve_args(mod, [u1], [u2])
end

function resolve_pval(mod::PCModification,
                      u1::FloatVec, u2::FloatVec)
    error("cdf modification still needs to be implemented for this
    type of copula modification")
end

function resolve_pval(mod::PCModification, u1::Float64, u2::Float64)
    return resolve_pval(mod, [u1], [u2])
end


"""
spin29.jl - unscented transform robustness for the δfq problem, this is the old working
spin23.jl
"""

WDIR = joinpath(@__DIR__, "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using Distributions
using HDF5
using Hyperopt
using ForwardDiff
using LinearAlgebra
using Random
using RobotDynamics
using StaticArrays
using Zygote
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin29"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
# state indices
const STATE1_IDX = SVector{HDIM_ISO}(1:HDIM_ISO)
const STATE2_IDX = SVector{HDIM_ISO}(STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO)
const INTCONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
const CONTROLS_IDX = INTCONTROLS_IDX[end] + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT
const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const S1_IDX = SVector{HDIM_ISO}(DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO)
const S2_IDX = SVector{HDIM_ISO}(S1_IDX[end] + 1:S1_IDX[end] + HDIM_ISO)
const S3_IDX = SVector{HDIM_ISO}(S2_IDX[end] + 1:S2_IDX[end] + HDIM_ISO)
const S4_IDX = SVector{HDIM_ISO}(S3_IDX[end] + 1:S3_IDX[end] + HDIM_ISO)
const S5_IDX = SVector{HDIM_ISO}(S4_IDX[end] + 1:S4_IDX[end] + HDIM_ISO)
const S6_IDX = SVector{HDIM_ISO}(S5_IDX[end] + 1:S5_IDX[end] + HDIM_ISO)
const S7_IDX = SVector{HDIM_ISO}(S6_IDX[end] + 1:S6_IDX[end] + HDIM_ISO)
const S8_IDX = SVector{HDIM_ISO}(S7_IDX[end] + 1:S7_IDX[end] + HDIM_ISO)
const S9_IDX = SVector{HDIM_ISO}(S8_IDX[end] + 1:S8_IDX[end] + HDIM_ISO)
const S10_IDX = SVector{HDIM_ISO}(S9_IDX[end] + 1:S9_IDX[end] + HDIM_ISO)
# control indices
const D2CONTROLS_IDX = 1:CONTROL_COUNT
# misc
const SAMPLE_STATES = [IS3_ISO]
const SAMPLE_STATE_COUNT = 1
const SAMPLES_PER_STATE = 10
const ASTATE_SIZE = ASTATE_SIZE_BASE + SAMPLE_STATE_COUNT * SAMPLES_PER_STATE * HDIM_ISO
const ACONTROL_SIZE = CONTROL_COUNT

# model
module Data
using RobotDynamics
using StaticArrays
const RD = RobotDynamics
const HDIM_ISO = 4
mutable struct Model <: RD.AbstractModel
    fq_cov::Float64
    alpha::Float64
end
end
Model = Data.Model
@inline RD.state_dim(model::Model) = ASTATE_SIZE
@inline RD.control_dim(model::Model) = ACONTROL_SIZE
@inline astate_sample_inds(sample_state_index::Int, sample_index::Int) = (
    SVector{HDIM_ISO}((
        ASTATE_SIZE_BASE + (sample_state_index - 1) * SAMPLES_PER_STATE * HDIM_ISO
        + (sample_index - 1) * HDIM_ISO + 1
    ):(
        ASTATE_SIZE_BASE + (sample_state_index - 1) * SAMPLES_PER_STATE * HDIM_ISO
        + sample_index * HDIM_ISO
    ))
)
const SAMPLE_IDXS = [astate_sample_inds(i, j)
                     for i = 1:SAMPLE_STATE_COUNT for j = 1:SAMPLES_PER_STATE]

# dynamics
function RD.discrete_dynamics(::Type{RK3}, model::Model, astate::SVector{ASTATE_SIZE},
                              acontrol::SVector{ACONTROL_SIZE}, time::Real, dt::Real)
    # base dynamics
    negi_hc = astate[CONTROLS_IDX[1]] * NEGI_H1_ISO
    h_prop = exp(dt * (FQ_NEGI_H0_ISO + negi_hc))
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    intcontrols = astate[INTCONTROLS_IDX[1]] + dt * astate[CONTROLS_IDX[1]]
    controls = astate[CONTROLS_IDX[1]] + dt * astate[DCONTROLS_IDX[1]]
    dcontrols = astate[DCONTROLS_IDX[1]] + dt * acontrol[D2CONTROLS_IDX[1]]

    # unscented transform
    s1 = astate[S1_IDX]
    s2 = astate[S2_IDX]
    s3 = astate[S3_IDX]
    s4 = astate[S4_IDX]
    s5 = astate[S5_IDX]
    s6 = astate[S6_IDX]
    s7 = astate[S7_IDX]
    s8 = astate[S8_IDX]
    s9 = astate[S9_IDX]
    s10 = astate[S10_IDX]
    # compute state mean
    sm = 1//10 .* (
        s1 + s2 + s3 + s4 + s5
        + s6 + s7 + s8 + s9 + s10
    )
    # compute state covariance
    d1 = s1 - sm
    d2 = s2 - sm
    d3 = s3 - sm
    d4 = s4 - sm
    d5 = s5 - sm
    d6 = s6 - sm
    d7 = s7 - sm
    d8 = s8 - sm
    d9 = s9 - sm
    d10 = s10 - sm
    s_cov = 1 / (2 * model.alpha^2) .* (
        d1 * d1' + d2 * d2' + d3 * d3' + d4 * d4'
        + d5 * d5' + d6 * d6' + d7 * d7' + d8 * d8'
        + d9 * d9' + d10 * d10'
    )
    # perform cholesky decomposition on joint covariance
    cov = zeros(eltype(s_cov), HDIM_ISO + 1, HDIM_ISO + 1)
    cov[1:HDIM_ISO, 1:HDIM_ISO] .= s_cov
    cov[HDIM_ISO + 1, HDIM_ISO + 1] = model.fq_cov
    # TOOD: cholesky! requires writing zeros in upper triangle
    cov_chol = model.alpha * cholesky(Symmetric(cov)).L
    s_chol1 = cov_chol[1:HDIM_ISO, 1]
    s_chol2 = cov_chol[1:HDIM_ISO, 2]
    s_chol3 = cov_chol[1:HDIM_ISO, 3]
    s_chol4 = cov_chol[1:HDIM_ISO, 4]
    fq_chol1 = cov_chol[HDIM_ISO + 1, 1]
    fq_chol2 = cov_chol[HDIM_ISO + 1, 2]
    fq_chol3 = cov_chol[HDIM_ISO + 1, 3]
    fq_chol4 = cov_chol[HDIM_ISO + 1, 4]
    fq_chol5 = cov_chol[HDIM_ISO + 1, 5]
    # propagate transformed states
    s1 = exp(dt * ((FQ + fq_chol1) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol1)
    s2 = exp(dt * ((FQ + fq_chol2) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol2)
    s3 = exp(dt * ((FQ + fq_chol3) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol3)
    s4 = exp(dt * ((FQ + fq_chol4) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol4)
    s5 = exp(dt * ((FQ + fq_chol5) * NEGI_H0_ISO + negi_hc)) * sm
    s6 = exp(dt * ((FQ - fq_chol1) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol1)
    s7 = exp(dt * ((FQ - fq_chol2) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol2)
    s8 = exp(dt * ((FQ - fq_chol3) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol3)
    s9 = exp(dt * ((FQ - fq_chol4) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol4)
    s10 = exp(dt * ((FQ - fq_chol5) * NEGI_H0_ISO + negi_hc)) * sm
    # normalize
    s1 = s1 ./sqrt(s1's1)
    s2 = s2 ./sqrt(s2's2)
    s3 = s3 ./sqrt(s3's3)
    s4 = s4 ./sqrt(s4's4)
    s5 = s5 ./sqrt(s5's5)
    s6 = s6 ./sqrt(s6's6)
    s7 = s7 ./sqrt(s7's7)
    s8 = s8 ./sqrt(s8's8)
    s9 = s9 ./sqrt(s9's9)
    s10 = s10 ./sqrt(s10's10)

    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
        s1; s2; s3; s4; s5; s6; s7; s8; s9; s10
    ]

    return astate_
end

# This cost puts a gate error cost on
# the sample states and a LQR cost on the other terms.
# The hessian w.r.t the state and controls is constant.
struct Cost{N,M,T} <: TO.CostFunction
    Q::Diagonal{T, SVector{N,T}}
    R::Diagonal{T, SVector{M,T}}
    q::SVector{N, T}
    c::T
    target_state1::SVector{HDIM_ISO, T}
    q_ss::T
    active_samples::Array{Int,1}
end

function Cost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
              xf::SVector{N,T}, target_state1::SVector{HDIM_ISO,T}, q_ss::T,
              active_samples::Array{Int,1}) where {N,M,T}
    q = -Q * xf
    c = 0.5 * xf' * Q * xf
    return Cost{N,M,T}(Q, R, q, c, target_state1, q_ss, active_samples)
end

@inline TO.state_dim(cost::Cost{N,M,T}) where {N,M,T} = N
@inline TO.control_dim(cost::Cost{N,M,T}) where {N,M,T} = M
@inline Base.copy(cost::Cost{N,M,T}) where {N,M,T} = Cost{N,M,T}(
    cost.Q, cost.R, cost.q, cost.c,
    cost.target_state1, cost.q_ss, cost.active_samples
)

function TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}) where {N,M,T}
    cost_ = 0.5 * astate' * cost.Q * astate + cost.q'astate + cost.c
    for i = 1:SAMPLE_STATE_COUNT
        for j in cost.active_samples
            sample_idx = astate_sample_inds(i, j)
            cost_ = cost_ + gate_error_iso2(astate, cost.target_state1; s1o=sample_idx[1] - 1)
        end
    end
    return cost_
end

@inline TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}, acontrol::SVector{M}) where {N,M,T} = (
    TO.stage_cost(cost, astate) + 0.5 * acontrol' * cost.R * acontrol
)

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T},
                      astate::SVector{N,T}) where {N,M,T}
    E.q = cost.Q * astate + cost.q
    for i = 1:SAMPLE_STATE_COUNT
        for j in cost.active_samples
            sample_idx = astate_sample_inds(i, j)
            E.q[sample_idx] = (cost.q_ss * jacobian_gate_error_iso2(astate, cost.target_state1;
                                                                   s1o=sample_idx[1] - 1))
        end
    end
    return false
end

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                      acontrol::SVector{M,T}) where {N,M,T}
    TO.gradient!(E, cost, astate)
    E.r = cost.R * acontrol
    E.c = 0
    return false
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T},
                     astate::SVector{N,T}) where {N,M,T}
    hess_sample = -1 * cost.q_ss * hessian_gate_error_iso2(cost.target_state1)
    for i = 1:SAMPLE_STATE_COUNT
        for j in cost.active_samples
            sample_idx = astate_sample_inds(i, j)
            E.Q[sample_idx, sample_idx] = hess_sample
        end
    end
    E.Q = E.Q + cost.Q
    return true
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                     acontrol::SVector{M,T}) where {N,M,T}
    TO.hessian!(E, cost, astate)
    E.R = cost.R
    E.H .= 0
    return true
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=60., solver_type=altro,
                  sqrtbp=false, integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e0, 1e-1],
                  dt_inv=Int64(1e1), smoke_test=false, constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true,
                  fq_cov=FQ * 1e-2, max_iterations=Int64(2e5), gradient_tol_int=1,
                  dJ_counter_limit=Int(1e2), state_cov=1e-2, seed=0, alpha=1.,)
    Random.seed!(seed)
    model = Model(fq_cov, alpha)
    n = RD.state_dim(model)
    m = RD.control_dim(model)
    t0 = 0.

    # initial state, target state
    x0 = zeros(n)
    xf = zeros(n)
    target_states = Array{SVector{HDIM_ISO},1}(undef, SAMPLE_STATE_COUNT)
    gate = GT_GATE_ISO[gate_type]
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    xf[STATE1_IDX] = gate * IS1_ISO_
    xf[STATE2_IDX] = gate * IS2_ISO_
    state_dist = Distributions.Normal(0., state_cov)
    for i = 1:SAMPLE_STATE_COUNT
        sample_state = SAMPLE_STATES[i]
        target_state = gate * sample_state
        target_states[i] = target_state
        for j = 1:SAMPLES_PER_STATE
            sample_idx = astate_sample_inds(i, j)
            sample = sample_state .+ rand(state_dist, HDIM_ISO)
            x0[sample_idx] = sample ./ sqrt(sample'sample)
            xf[sample_idx] = target_state
        end
    end
    x0 = SVector{n}(x0)
    xf = SVector{n}(xf)

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    x_max = SVector{n}(x_max)
    x_min = SVector{n}(x_min)

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    x_max_boundary = SVector{n}(x_max_boundary)
    x_min_boundary = SVector{n}(x_min_boundary)

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}([
        fill(1e-4, CONTROL_COUNT);
    ]) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], CONTROL_COUNT); # ∫a
        fill(qs[3], CONTROL_COUNT); # a
        fill(qs[4], CONTROL_COUNT); # ∂a
        fill(0, SAMPLES_PER_STATE * HDIM_ISO);
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[6], CONTROL_COUNT); # ∂2a
    ]))
    # objective = LQRObjective(Q, R, Qf, xf, N)
    active_samples = Array(1:SAMPLES_PER_STATE)
    cost_k = Cost(Q, R, xf, target_states[1], qs[5], active_samples)
    cost_f = Cost(Qf, R, xf, target_states[1], N * qs[5], active_samples)
    objective = TO.Objective(cost_k, cost_f, N)

    # constraints
    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_idxs = copy(SAMPLE_IDXS)
    push!(norm_idxs, STATE1_IDX)
    push!(norm_idxs, STATE2_IDX)
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx) for idx in norm_idxs]
    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N)
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(
        solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
        projected_newton_tolerance=al_tol, n_steps=n_steps,
        penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
        projected_newton=projected_newton, iterations_inner=iterations_inner,
        iterations_outer=iterations_outer, iterations=max_iterations,
        gradient_tolerance_intermediate=gradient_tol_int,
        dJ_counter_limit=dJ_counter_limit,
    )
    Altro.solve!(solver)

    # Post-process.
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "qs" => qs,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gradient_tolerance_intermediate" => gradient_tol_int,
        "dJ_counter_limit" => dJ_counter_limit,
        "integrator_type" => Integer(integrator_type),
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "iterations" => iterations_,
        "seed" => seed,
        "fq_cov" => fq_cov,
        "max_iterations" => max_iterations,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    return result
end


function build_kalman_model(obj::IntegratedPropagatorBuilder, arg0::List, arg1::List, arg2::ParameterDriversList)
    return jcall(obj, "buildKalmanModel", KalmanODModel, (List, List, ParameterDriversList), arg0, arg1, arg2)
end

function build_ls_model(obj::IntegratedPropagatorBuilder, arg0::Vector{IntegratedPropagatorBuilder}, arg1::List, arg2::ParameterDriversList, arg3::ModelObserver)
    return jcall(obj, "buildLSModel", BatchLSODModel, (Vector{IntegratedPropagatorBuilder}, List, ParameterDriversList, ModelObserver), arg0, arg1, arg2, arg3)
end

function get_frame(obj::PropagatorBuilder)
    return jcall(obj, "getFrame", Frame, ())
end

function get_initial_orbit_date(obj::PropagatorBuilder)
    return jcall(obj, "getInitialOrbitDate", AbsoluteDate, ())
end

function get_orbit_type(obj::PropagatorBuilder)
    return jcall(obj, "getOrbitType", OrbitType, ())
end

function get_orbital_parameters_drivers(obj::PropagatorBuilder)
    return jcall(obj, "getOrbitalParametersDrivers", ParameterDriversList, ())
end

function get_position_angle(obj::PropagatorBuilder)
    return jcall(obj, "getPositionAngle", PositionAngle, ())
end

function get_propagation_parameters_drivers(obj::PropagatorBuilder)
    return jcall(obj, "getPropagationParametersDrivers", ParameterDriversList, ())
end

function get_selected_normalized_parameters(obj::PropagatorBuilder)
    return jcall(obj, "getSelectedNormalizedParameters", Vector{jdouble}, ())
end

function reset_orbit(obj::IntegratedPropagatorBuilder, arg0::Orbit)
    return jcall(obj, "resetOrbit", void, (Orbit,), arg0)
end


using Distributed
# using ClusterManagers
using Logging

# logger = SimpleLogger(stdout, Logging.Debug)
# old_logger = global_logger(logger)

try
    num_tasks = parse(Int, ENV["SLURM_NTASKS"])
    @info "Number of slurm tasks: $num_tasks"
#    addprocs(SlurmManager(num_tasks), unbuffered="")#, topology=:master_worker)
catch KeyError
    @warn "Not inside a slurm job"
end

@info "Number of processes $(nprocs())"

@everywhere begin
    using Pkg
    Pkg.activate(".")
end

@everywhere using QContinuousMeasurement
@everywhere using TimerOutputs

using HDF5
using DelimitedFiles
using Logging

using ArgParse
using ProgressMeter

# Parse the arguments passed to the script
s = ArgParseSettings()

@add_arg_table! s begin
    "--Nj"
        help = "Number of spins"
        default = 1
        arg_type = Int64
    "--Ntraj"
        help = "Number of trajectories"
        default = 1
        arg_type = Int64
    "--Tfinal", "-T"
        help = "Final time"
        arg_type = Float64
        default = 1.0
    "--dt"
        help = "dt"
        arg_type = Float64
        default = 0.0001
    "--kind"
        help = "Independent noise rate k_ind"
        arg_type = Float64
        default = 1.0
    "--kcoll"
        help = "Collective noise rate k_coll"
        arg_type = Float64
        default = 0.0
    "--Gamma"
        help = "Monitoring noise rate Gamma"
        arg_type = Float64
        default = 1.0
    "--omega"
        help = "Hamiltonian frequency omega"
        arg_type = Float64
        default = 1.0
    "--eta"
        help = "Measurement efficiency"
        arg_type = Float64
        default = 1.0
    "--mu"
        help = "Initial squeezing parameter"
        arg_type = Float64
        default = 0.0
    "--outpoints"
        help = "Number of output points"
        arg_type = Int64
        default = 200
    "--liouvillianfile"
        help = "npz file with the liovuillian data"
        arg_type = String
        default = nothing
    "--outdir"
        help = "Output directory"
        arg_type = String
        default = "."
end

args = parse_args(s)

jobid = "SLURM_JOB_ID" in keys(ENV) ? jobid = ENV["SLURM_JOB_ID"] : ""
filename = string("$(args["outdir"])/sim_Nj_$(args["Nj"])_Ntraj_$(args["Ntraj"])_Tf_$(args["Tfinal"])_" *
            "dt_$(args["dt"])_kind_$(args["kind"])_kcoll_$(args["kcoll"])_Gamma_$(args["Gamma"])_" *
            "omega_$(args["omega"])_eta_$(args["eta"])_mu_$(args["mu"])_$(jobid)")

let pkgstr = ""
    for (uuid, pkg) in Pkg.dependencies()
        if pkg.is_direct_dep
            pkgstr *= "\t$(pkg.name) $(pkg.version)"
            if pkg.is_tracking_repo
                pkgstr *= " #$(pkg.git_revision) ($(pkg.git_source))"
            elseif pkg.is_tracking_path
                pkgstr *= " $(pkg.source)"
            end
            pkgstr *= "\n"
        end
    end
    @info "Installed packages\n$pkgstr"
end

@info "Output filename: $filename"

Ntraj = args["Ntraj"]

@everywhere to = TimerOutput()

modelparams = ModelParameters(Nj=args["Nj"],
                              kind=args["kind"],
                              Gamma=args["Gamma"],
                              omega=args["omega"],
                              eta=args["eta"],
                              dt=args["dt"],
                              Tfinal=args["Tfinal"],
                              outpoints=args["outpoints"],
                              mu=args["mu"])

@info "Output every $(modelparams._outsteps) steps"

@info "Initializing model..."

init_time = @elapsed begin
    if args["kind"] == 0.0
        @everywhere model = CollectiveDephasingModel($modelparams)
        @everywhere initial_state = fixedj_sss($modelparams.Nj, $modelparams.mu)
    else
        @assert args["kcoll"] == 0.0
        @everywhere model = LocalDephasingModel($modelparams, $args["liouvillianfile"])
        @everywhere initial_state = blockdiag_sss($modelparams.Nj, $modelparams.mu)
    end
end

xi = squeezing_param(initial_state)
@info "Initial spin squeezing $(xi)"
# println(model)
# println(initial_state)

@info "Done in $init_time seconds..."

# Opens the FileWriter
writer = FileWriter("$filename.h5", modelparams, args["Ntraj"], ["FI", "QFI", "FIstrong", "xi2y", "Jx"])

progress_channel = RemoteChannel(() -> Channel{Bool}(1000))

p = Progress(Ntraj * length(get_time(model)), barglyphs=BarGlyphs("[=> ]"), dt=60, barlen=10)

@async while take!(progress_channel)
    next!(p)
end

if "JULIA_NUM_THREADS" in keys(ENV)
    @everywhere numthreads = $(parse(Int, ENV["JULIA_NUM_THREADS"]))
    @info("Number of threads: $numthreads")
else
    @info("JULIA_NUM_THREADS not set, using 1 thread")
    @everywhere numthreads = 1
end

# @everywhere begin
#     numthreads = 1
#     try
#         if "JULIA_NUM_THREADS" in keys(ENV)
#             numthreads = parse(Int, ENV["JULIA_NUM_THREADS"])
#             @info("Number of threads: $numthreads")
#         else
#             @info("JULIA_NUM_THREADS not set, using 1 thread")
#             numthreads = 1
#         end
#     catch
#         @info "Uffa!"
#     end
# end

@everywhere function thread_simulate_trajectory(model, initial_state, ntraj, file_channel, progress_channel)
    pid = Distributed.myid()
    nth = Threads.nthreads()
    Threads.@threads for i in 1:ntraj
        tid = Threads.threadid()

        simulate_trajectory(model, initial_state, file_channel, progress_channel)
        @debug "Trajectory done from thread $tid of $nth on worker $pid."
    end

    @debug "Batch of $ntraj done!"
end

file_channel = writer.channel

function prepare_batches(n, batch_size)
    if n % batch_size == 0
        return [batch_size for i = 1:(n ÷ batch_size)]
    else
        m = Int(floor(n / batch_size))
        vcat([[batch_size for i = 1:m], n - (m * batch_size)]...)
    end
end

trajectory_time = @elapsed begin
   pmap(x -> thread_simulate_trajectory(model, initial_state, x, file_channel, progress_channel), prepare_batches(Ntraj, numthreads), on_error=identity)
end

@info "pmap done"
put!(progress_channel, false)

@info "Trajectory simulation time: $trajectory_time"

for i in eachindex(workers())
    fetch(@spawnat i begin
        try
            result = read(`grep VmHWM /proc/$(getpid())/status`, String)
            peakmem = tryparse(Int, String(match(r"(\d+)", result)[1]))

            @info "Peak memory in GB: $(peakmem / 1024^2)"
        catch er
            @info er
        end
    end)
end

close(writer)


################################################################################
# Observations

@doc raw"""
    mutable struct Observation{T <: Real}

An *observation* is a sequence of time-ordered data (TOD) that has
been acquired by some detector. It implements the following fields:

- `time`: an array of `N` elements, representing the time of the
  sample. Being defined as an `AbstractArray`, a range (e.g.,
  `0.0:0.1:3600.0`) can be used to save memory
- `pixidx`: an array of `N` elements, each being a pixel index in a
  map. In no way the Healpix pixelization scheme is enforced; in fact,
  the map must not even be spherical.
- `psi_angle`: an array of `N` elements, each containing the
  orientation of the polarization angle with respect to some fixed
  reference system (e.g., the Ecliptic plane). The angles must be in
  **radians**.
- `tod`: an array of `N` elements, containing the actual data measured
  by the instrument and already calibrated to some physical units
  (e.g., K_CMB)
- `sigma_squared`: an array of `N` elements, containing the squared
  RMS of each sample. To save memory, you should usually use a
  `RunLengthArray` here.
- `flagged`: a Boolean array of `N` elements, telling whether the
  sample should be discarded (`true`) or not (`false`) during the
  map-making process.
- `name`: a string representing the detector. This is used only for
  debugging purposes.

To create an observation, the constructor takes *keyword arguments*
instead of parameters; in this way, the code should be more
readable. Also, sensible defaults will be provided for the missing
fields:

```julia
# We do not specify "sigma_squared" nor "flagged", so they will be
# initialized to a vector of ones and a vector of `false`
obs = Observation{Float64}(
    time = 0.0:0.1:3600.0,
    pixidx = 1:3601,
    psi_angle = zeros(3601),
    tod = randn(3601),
)
```

"""
mutable struct Observation{T <: Real}
    time::AbstractVector{T}
    pixidx::Vector{Int}
    psi_angle::Vector{T}

    tod::Vector{T}
    sigma_squared::AbstractVector{T}
    flagged::AbstractVector{Bool}

    name::String

    function Observation{T}(
        ;
        time = T[],
        pixidx = Int[],
        psi_angle = T[],
        tod = T[],
        sigma_squared = T[],
        flagged = Bool[],
        name = "unnamed",
    ) where {T <: Real}
        @assert !isempty(pixidx)

        nsamples = length(pixidx)

        int_time = isempty(time) ? (0:(nsamples - 1)) : time
        int_tod = isempty(tod) ? zeros(T, nsamples) : tod
        int_psi_angle = isempty(psi_angle) ? zeros(T, nsamples) : psi_angle
        int_sigma_squared = isempty(sigma_squared) ? RunLengthArray{Int,T}(nsamples, 1.0) : sigma_squared
        int_flagged = isempty(flagged) ? zeros(Bool, nsamples) : flagged

        @assert length(int_time) == nsamples
        @assert length(int_psi_angle) == nsamples
        @assert length(int_tod) == nsamples
        @assert length(int_sigma_squared) == nsamples
        @assert length(int_flagged) == nsamples

        new(
            int_time,
            pixidx,
            int_psi_angle,
            int_tod,
            int_sigma_squared,
            int_flagged,
            name,
        )
    end
end

function Base.show(io::IO, obs::Observation)
    print(io, "Observation(name => \"$(obs.name)\", samples = $(length(obs.time)))")
end

function Base.show(io::IO, ::MIME"text/plain", obs::Observation)

    isempty(obs.time) && println(io, "Observation for detector %(obs.name): empty")

    println(
        io,
        @sprintf(
            """Observation for detector %s:
Number of samples....................... %d
Start time.............................. %f
Duration................................ %f
Number of flagged samples............... %d (%.1f%%)
Average level of the TOD................ %e
Average noise^2 level................... %e
""",
            obs.name,
            length(obs.time),
            obs.time[1],
            obs.time[end] - obs.time[1],
            length(obs.flagged[obs.flagged]),
            100 * length(obs.flagged[obs.flagged]) / length(obs.flagged),
            mean(obs.tod),
            mean(obs.sigma_squared),
        )
    )
end


workspace()
using JukaiNLP
using JukaiNLP.DepParsing
using HDF5

path = joinpath(Pkg.dir("JukaiNLP"), ".corpus/nyt100.h5")
forms = h5read(path, "str")
formdict = IdDict(forms)

conv(form, cat, head) = push!(formdict,form), push!(catdict,cat), parse(Int,head)

trainpath = joinpath(Pkg.dir("JukaiNLP"), ".corpus/wsj_02-21.conll")
testpath = joinpath(Pkg.dir("JukaiNLP"), ".corpus/wsj_23.conll")
traindata = CoNLL.read(trainpath, 2, 5, 7)
testdata = CoNLL.read(testpath, 2, 5, 7)

#=
# parser for unlabeled dependency tree
path = Pkg.dir("JukaiNLP")
parser = DepParser(Unlabeled, "$(path)/dict/en-word_nyt.dict")

# parser for labeled dependency tree
parser = DepParser(Labeled, "$(path)/dict/en-word_nyt.dict")
sents = readconll(parser, "$(path)/corpus/mini-training-set.conll")

# supports structured perceptron and feedforward
# neural network (Chen and Manning, 2014) models

n = div(length(sents), 10) * 8
trainsents, testsents = sents[1:n], sents[n+1:end]

#train2!(FeedForward, parser, trainsents, iter=20)
#modelpath = "C:/Users/hshindo/Desktop/depparser.jld"
#JLD.save(modelpath, "depparser", parser.model)

train!(FeedForward, parser, trainsents, iter=20)

#train!(parser, trainsents, nonlinear=tanh, hiddensizes=[200])

res = parser(testsents)
evaluate(parser, res)
=#


# TODO: think of a better name for this? and consider using a special container
# type (akin to ModelFrame), since we'd probably like to hold onto the schema,
# too.

################################################################################
# With Formula as a vector of Terms:
#
# entry point is to set_schema on Vector{Term}.  initialize set of terms seen
# already as a Set{Term}.  then map set_schema! on each term.
#
# for individual Term, default method is to do nothing (return the term), and
# add it to the set of terms already seen.
#
# for dealing with categorical terms, need to keep track of what's been seen so
# far, and the context in which the Term we're currently processing occurred (by
# itself, or as part of an interaction, or otherwise).  All that matters for
# determining whether an aliased term is present is the _set_ of variables
# contained in previously encountered terms, so we also need a way of extracting
# that.

function set_schema!(f::Formula, sch::Data.Schema)
    @argcheck !f.schema_set "Schema already set for this formula!"
    already = Set()
    f.term = set_schema(f.term, Set(), sch)
    f.schema_set = true
    f
end

set_schema(terms::AbstractVector{<:Terms.Term}, already::Set, sch::Data.Schema) =
    map(t->set_schema(t, already, sch), terms)

set_schema(::Nothing, ::Set, ::Data.Schema) = nothing

set_schema(t::Terms.FormulaTerm, already::Set, sch::Data.Schema) =
    Terms.FormulaTerm(set_schema(t.lhs, Set(), sch),
                      set_schema(t.rhs, already, sch))

set_schema(t::Terms.Intercept, already::Set, ::Data.Schema) = (push!(already, termsyms(t)); t)

set_schema(t::Terms.FunctionTerm, already::Set, ::Data.Schema) =
    (push!(already, termsyms(t)); t)

function set_schema(t::Terms.Interaction, already::Set, sch::Data.Schema)
    new_t = Terms.Interaction(map(s -> set_schema(s, t, already, sch), t.terms))
    push!(already, termsyms(new_t))
    new_t
end

set_schema(t::Terms.Eval, already::Set, sch::Data.Schema) = set_schema(t, t, already, sch)

function set_schema(t::Terms.Eval, context::T, already::Set, sch::Data.Schema) where T<:Terms.Term
    @debug "$t in context of $context, already seen $already"
    push!(already, termsyms(t))
    if is_categorical(t, sch)
        aliased = aliassyms(t, context)
        @debug "  aliases: $aliased"
        if aliased in already
            # TODO: allow custom contrasts here
            contr = DEFAULT_CONTRASTS()
        else
            # lower-order term that is aliased by full-rank contrasts for this
            # term is NOT present, so use full-rank contrasts and add the
            # aliased term to the set of terms we've seen
            contr = FullDummyCoding()
            push!(already, aliased)
        end
        uniq = sch.metadata[:summaries][name(t)][:unique]
        mat = ContrastsMatrix(contr, uniq)
        invindex = Dict(lev => i for (i,lev) in enumerate(mat.levels))
        N = size(mat, 2)
        Terms.Categorical{N}(name(t), mat, invindex)
    else
        Terms.Continuous(name(t))
    end
end

"""
    termsyms(t::Terms.Term)

Extract the Set of symbols referenced in this term.

This is needed in order to determine when a categorical term should have
standard (reduced rank) or full rank contrasts, based on the context it occurs
in and the other terms that have already been encountered.
"""

termsyms(t::Terms.Term) = Set()
termsyms(t::Terms.Intercept) = Set([1])
termsyms(t::Union{Terms.Eval, Terms.Categorical, Terms.Continuous}) = Set([t.name])
termsyms(t::Terms.Interaction) = mapreduce(termsyms, union, t.terms)
termsyms(t::Terms.FunctionTerm) = Set([t.name])

"""
    aliassyms(t::T, context::S) where {T<:Terms.Term, S<:Terms.Term}

Get the Set of symbols that this term potentially aliases in this context.
"""
aliassyms(::Terms.Eval, context::Terms.Eval) = Set([1])
function aliassyms(t::Terms.Eval, context::Terms.Interaction)
    mapreduce(termsyms, union, c for c in context.terms if t ≠ c)
end


using JLD2, LinearAlgebra, Oceananigans, Printf
record_interaction = false
include(pwd() * "/scripts/vizinanigans.jl")
include(pwd() * "/scripts/states.jl")
include(pwd() * "/scripts/compare.jl")
include(pwd() * "/scripts/zonalstates.jl")

include(pwd() * "/analysis_scripts/" * "post_analysis.jl")

files = [
    pwd() * "/Channel_1_zonal_averages.jld2",
    pwd() * "/Channel_16_zonal_averages.jld2",
    pwd() * "/Channel_24_zonal_averages.jld2",
]

new_files = [
    pwd() * "/Channel_8_zonal_averages.jld2",
    pwd() * "/Ridge_8_zonal_averages.jld2",
    pwd() * "/Relaxation_Channel_16_zonal_averages.jld2",
    pwd() * "/Abernathy_16_zonal_averages.jld2",
    pwd() * "/Abernathy_16_τ0.05_Q10_zonal_averages.jld2",
    pwd() * "/Abernathy_1_τ0.20_Q10_zonal_averages.jld2",
    pwd() * "/Abernathy_16_τ0.20_Q10_zonal_averages.jld2",
    pwd() * "/Abernathy_16_τ0.10_Q5_zonal_averages.jld2",
    pwd() * "/Abernathy_1_τ0.10_Q5_zonal_averages.jld2",

]
##
file = new_files[end] # new_files[1]
file = pwd() * "/Abernathy_16_τ0.20_Q10_zonal_averages.jld2"
zonalstatistics = jldopen(file)
tkeys = keys(zonalstatistics["timeseries"]["t"])
states, statenames, units, domain = grabzonalstates(file, ghost = 3, startind = 20)
li = 1 # bottom
mval = 0 # positive
ui = length(domain[2])-mval # top
newstates = [state[:, li:end-mval] for state in states]
xlims = (domain[1][1], domain[1][end])
ylims = (domain[2][li], domain[2][ui])
scene = visualize(newstates, statenames = statenames, xlims = xlims, ylims = ylims, units = units)

##
seconds = 20
fps = 10
frames = round(Int, fps * seconds )
if record_interaction
record(scene, pwd() * "/zonalandtemporal.mp4"; framerate = fps) do io
    for i = 1:frames
        sleep(1/fps)
        recordframe!(io)
    end
end
end


export mutate_weight, mutate_connect, mutate_neuron, mutate_enabled, mutate

"Mutate the weight of genes"
function mutate_weight(ind::NEATIndiv, cfg::Dict)
    #TODO Add mutation power
    # https://github.com/FernandoTorres/NEAT/blob/master/src/genome.cpp
    ind_mut = NEATIndiv(ind)
    for g in values(ind_mut.genes)
        g.weight = g.weight + randn()*cfg["weight_factor"]
    end
    ind_mut.fitness.= -Inf
    ind_mut
end

function indexof(a::Array{Float64}, f::Float64)
    findall(x->x==f, a)[1]
end

"Add a connection between 2 random neurons"
function mutate_connect(ind::NEATIndiv, cfg::Dict)
    sort!(ind.neuron_pos)

    ind_mut = NEATIndiv(ind)
    nb_neur = length(ind_mut.neuron_pos)
    # println(nb_neur, " ", ind.neuron_pos)

    if cfg["hyperNEAT"] # HyperNEAT
        n_in = 4
        n_out = 1
    else # NEAT
        n_in = cfg["n_in"]
        n_out = cfg["n_out"]
    end

    # Valid neuron pairs
    valid = trues(nb_neur, nb_neur)

    # Remove existing ones
    for g in values(ind.genes)
        i_origin = indexof(ind.neuron_pos, g.origin)
        i_dest = indexof(ind.neuron_pos, g.destination)
        valid[i_origin, i_dest]=false
    end

    for orig in 1:nb_neur
        # Remove links to self
        valid[orig, orig]=false

        orig_pos = ind.neuron_pos[orig]

        for dest in 1:nb_neur
            dest_pos = ind.neuron_pos[dest]

            # Remove links towards input neurons and bias neuron
            if dest_pos <= 0
                valid[orig, dest]=false
            end

            # Remove links between output neurons (would not support adding a neuron)
            if orig_pos >= 1 && dest_pos >= 1
                valid[orig, dest]=false
            end

            # Remove recurrence if needed
            if !cfg["allow_recurrence"] && orig_pos > dest_pos
                valid[orig, dest]=false
            end
        end
    end

    ind_mut.fitness.= -Inf
    ind_mut.network = Network(n_in, n_out, Dict()) # Reset network

    # Filter invalid ones
    conns = findall(valid)
    if length(conns) > 0
        shuffle!(conns) # Pick random
        cfg["innovation_max"] += 1

        i = ind.neuron_pos[conns[1][1]]
        j = ind.neuron_pos[conns[1][2]]
        if j <= 0 # Catching error where destination is an input
            return ind_mut # Refuse mutation
        end

        g = Gene(i, j, cfg["innovation_max"])

        ind_mut.genes[cfg["innovation_max"]]=g
    end
    ind_mut
end

"Remove a random connection"
function mutate_disconnect(ind::NEATIndiv, cfg::Dict)
    ind_mut = NEATIndiv(ind)
    if length(ind_mut.genes)<=1 # Always keep 1 gene
        return ind_mut
    end
    k = rand(collect(keys(ind_mut.genes))) # pick a random gene
    pop!(ind_mut.genes, k) # remove it
    if cfg["hyperNEAT"] # HyperNEAT
        n_in = 4
        n_out = 1
    else # NEAT
        n_in = cfg["n_in"]
        n_out = cfg["n_out"]
    end
    ind_mut.network = Network(n_in, n_out, Dict()) # Reset network
    ind_mut
end

"Split a connection into 2 connections with a new neuron"
function mutate_neuron(ind::NEATIndiv, cfg::Dict)
    ind_mut = NEATIndiv(ind)

    if length(ind.genes)==0
        return ind_mut
    end

    # Connection to split
    g = rand(collect(values(ind_mut.genes)))
    g.enabled = false

    # Create neuron between origin and destination, in [0; 1]
    n = random_position(g.origin, g.destination)
    push!(ind_mut.neuron_pos, n)
    sort!(ind_mut.neuron_pos)
    ind_mut.activ_functions[n] = rand(cfg["activation_functions"])

    # Create connections
    i = cfg["innovation_max"]
    ind_mut.genes[i+1] = Gene(g.origin, n, i+1)
    ind_mut.genes[i+2] = Gene(n, g.destination, i+2)
    cfg["innovation_max"] +=2
    ind_mut.fitness.= -Inf
    ind_mut
end

"Switch enabled"
function mutate_enabled(ind::NEATIndiv, cfg::Dict)
    ind_mut = NEATIndiv(ind)
    g = rand(collect(values(ind_mut.genes)))
    g.enabled = !g.enabled
    ind_mut
end

function mutate(ind::NEATIndiv, cfg::Dict)
    # println("mutation")
    if rand() < cfg["p_mutate_add_neuron"]
        c =  mutate_neuron(ind, cfg)
    elseif rand() < cfg["p_mutate_add_connection"]
        c = mutate_connect(ind, cfg)
    elseif rand() < cfg["p_mutate_remove_connection"]
        c = mutate_disconnect(ind, cfg)
    elseif rand() < cfg["p_mutate_weights"]
        c = mutate_weight(ind, cfg)
    elseif rand() < cfg["p_mutate_enabled"]
        c = mutate_enabled(ind, cfg)
    else
        # return clone if no mutation occurs
        c = NEATIndiv(ind)
    end
    build!(c)
    c
end


using Test
using HolidayCalendars
using Dates: Date, dayofweek
using RDates

@testset "Holiday Calendars" verbose=true begin
    @testset "United Kingdom → Settlements" begin
        holidays_2004 = Set([
            Date(2004,1,1),
            Date(2004,4,9),
            Date(2004,4,12),
            Date(2004,5,3),
            Date(2004,5,31),
            Date(2004,8,30),
            Date(2004,12,27),
            Date(2004,12,28),
        ])

        cal = RDates.calendar(CALENDARS, ["UK/SETTLEMENT"])
        for d in range(Date(2004,1,1), Date(2004,12,31), rd"1b@WEEKEND"; cal_mgr=CALENDARS)
            @test is_holiday(cal, d) == (d in holidays_2004 || dayofweek(d) > 5)
        end

        holidays_2005 = Set([
            Date(2005,1,3),
            Date(2005,3,25),
            Date(2005,3,28),
            Date(2005,5,2),
            Date(2005,5,30),
            Date(2005,8,29),
            Date(2005,12,26),
            Date(2005,12,27),
        ])

        for d in range(Date(2005,1,1), Date(2005,12,31), rd"1b@WEEKEND"; cal_mgr=CALENDARS)
            @test is_holiday(cal, d) == (d in holidays_2005 || dayofweek(d) > 5)
        end

        holidays_2021 = Set([
            Date(2021,1,1),
            Date(2021,4,2),
            Date(2021,4,5),
            Date(2021,5,3),
            Date(2021,5,31),
            Date(2021,8,30),
            Date(2021,12,27),
            Date(2021,12,28),
        ])

        for d in range(Date(2021,1,1), Date(2021,12,31), rd"1b@WEEKEND"; cal_mgr=CALENDARS)
            @test is_holiday(cal, d) == (d in holidays_2021 || dayofweek(d) > 5)
        end

        holidays_2022 = Set([
            Date(2022,1,3),
            Date(2022,4,15),
            Date(2022,4,18),
            Date(2022,5,2),
            Date(2022,6,2),
            Date(2022,6,3),
            Date(2022,8,29),
            Date(2022,12,26),
            Date(2022,12,27),
        ])

        for d in range(Date(2022,1,1), Date(2022,12,31), rd"1b@WEEKEND"; cal_mgr=CALENDARS)
            @test is_holiday(cal, d) == (d in holidays_2022 || dayofweek(d) > 5)
        end
    end
end


# struct TemperedModel <: AbstractPPL.AbstractProbabilisticProgram
#     model :: DynamicPPL.Model
#     β     :: AbstractFloat
# end

function make_tempered_model end


function backwards_tfunc(@nospecialize(f), primal::IRCode, inst::Expr, @nospecialize(Δ))
    if f === Core.tuple
        return tuple_tfunc(Any[NoTangent, tuple_type_fields(Δ)...])
    elseif f == Core.getfield
        tt = argextype(inst.args[2], primal, primal.sptypes)
        ot = widenconst(tt)
        if ot <: Tuple
            if length(ot.parameters) == 1
                # Only one choice
                rt = tuple_tfunc(Any[Δ])
            else
                idxt = argextype(inst.args[3], primal, primal.sptypes)
                isa(idxt, Const) || error()
                args = Any[ZeroTangent for i = 1:length(ot.parameters)]
                args[idxt.val] = Δ
                rt = tuple_tfunc(args)
            end
        else
            idxt = argextype(inst.args[3], primal, primal.sptypes)
            if !isa(idxt, Const)
                @show idxt
                error()
            end
            if isa(idxt.val, Symbol)
                symn = idxt.val
            elseif isa(idxt.val, Int)
                @show tt
                symn = fieldname(tt, idxt.val)
            else
                error()
            end
            rt = NamedTuple{(symn,), Tuple{widenconst(Δ)}}
        end
        return tuple_tfunc(Any[NoTangent, rt, NoTangent, NoTangent])
    elseif f == Core.apply_type
        return tuple_tfunc(Any[NoTangent for i = 1:length(inst.args)])
    elseif f == Core.typeof
        return tuple_tfunc(Any[ZeroTangent for i = 1:length(inst.args)])
    elseif f == (===)
        return tuple_tfunc(Any[ZeroTangent for i = 1:length(inst.args)])
    elseif f == nfields
        return tuple_tfunc(Any[ZeroTangent for i = 1:length(inst.args)])
    elseif f == fieldtype
        return tuple_tfunc(Any[ZeroTangent for i = 1:length(inst.args)])
    elseif f === Core._apply_iterate
        ft = argextype(inst.args[3], primal, primal.sptypes)
        f = singleton_type(ft)
        @assert f === Core.tuple
        return backwards_tfunc(f, primal, inst, Δ)
    end
    @show f
    error()
end

function forward_tfunc(@nospecialize(f), primal::IRCode, inst::Expr, @nospecialize(Δ))
    if f === Core.tuple
        return tuple_tfunc(Any[tuple_type_fields(Δ)[2:end]...])
    elseif f == Core.getfield
        idxt = argextype(inst.args[3], primal, primal.sptypes)
        isa(idxt, Const) || error()
        val = getfield_tfunc(Δ, Const(2))
        fieldval = (isa(val, Const) && (isa(val.val, ZeroTangent) || isa(val.val, NoTangent))) ?
            Const(ZeroTangent()) : getfield_tfunc(val, idxt)
        return fieldval
    elseif f === Core.apply_type || f === Core.typeof
        return NoTangent
    end
    @show f
    error()
end

function getfield_prop_zero_tfunc(@nospecialize(s), @nospecialize(x))
    if isa(s, Const)
        if isa(s.val, NoTangent) || isa(s.val, ZeroTangent)
            return s
        end
    end
    if s === ZeroTangent || s === NoTangent
        return s
    end
    return getfield_tfunc(s, x)
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@with_kw mutable struct AndersonAccelerationData{T}
    current::Int = 1
    Q::T = 1e10
end

@with_kw mutable struct AndersonAccelerationParameters
    m::Int = 10
end

"""
    AndersonAcceleratedProximal

Functor object for using anderson accleration in the prox step of a quasigradient algorithm. Create by supplying an [`AndersonAcceleration`](@ref) object through `prox ` to `QuasiGradient.Optimizer` or by setting the [`Prox`](@ref) attribute.

...
# Parameters
- `prox::AbstractProx = Polyhedron`: Inner prox step
- `m::Integer = 10`: Anderson memory
...
"""
struct AndersonAcceleratedProximal{T <: AbstractFloat, P <: AbstractProximal} <: AbstractProximal
    data::AndersonAccelerationData{T}
    parameters::AndersonAccelerationParameters

    prox::P
    x::Vector{T}
    save_g::Vector{T}
    grad_mapping::Vector{T}
    residuals::Vector{T}
    y::Vector{Vector{T}}
    g::Vector{Vector{T}}

    function AndersonAcceleratedProximal(proximal::AbstractProximal, ::Type{T}; kw...) where T <: AbstractFloat
        P = typeof(proximal)
        return new{T,P}(AndersonAccelerationData{T}(),
                        AndersonAccelerationParameters(; kw...),
                        proximal,
                        Vector{T}(),
                        Vector{T}(),
                        Vector{T}(),
                        Vector{T}(),
                        Vector{Vector{T}}(),
                        Vector{Vector{T}}())
    end
end

function initialize_prox!(quasigradient::AbstractQuasiGradient, anderson::AndersonAcceleratedProximal)
    # Initialize inner
    initialize_prox!(quasigradient, anderson.prox)
    return nothing
end

function restore_proximal_master!(quasigradient::AbstractQuasiGradient, anderson::AndersonAcceleratedProximal)
    # Restore inner
    restore_proximal_master!(quasigradient, anderson.prox)
    return nothing
end

function prox!(quasigradient::AbstractQuasiGradient, anderson::AndersonAcceleratedProximal, x::AbstractVector, ∇f::AbstractVector, γ::AbstractFloat)
    if length(anderson.y) == 0
        # First iteration standard prox
        push!(anderson.y, copy(x))
        push!(anderson.y, x - γ*∇f)
        push!(anderson.g, copy(anderson.y[2]))
        push!(anderson.g, zero(x))
        append!(anderson.save_g, copy(anderson.y[2]))
        prox!(quasigradient, anderson.prox, x, ∇f, γ)
        append!(anderson.x, copy(x))
        append!(anderson.grad_mapping, fill(Inf, length(x)))
        anderson.data.current = 2
        anderson.data.Q = quasigradient.data.Q
        return nothing
    else
        # Guard step
        @unpack Q = quasigradient.data
        if Q - anderson.data.Q > -0.5 * γ * norm(anderson.grad_mapping / γ)^2
            # Revert to regular proximal step
            x .= anderson.x
            anderson.y[anderson.data.current] .= anderson.save_g
            # Recompute objective and gradient
            Q = resolve_subproblems!(quasigradient)
            anderson.data.Q = Q
            if Q <= quasigradient.data.Q
                quasigradient.data.Q = Q
            end
            # Skip anderson step if this is the final iteration
            if terminate(quasigradient,
                         quasigradient.data.iterations,
                         Q,
                         x,
                         quasigradient.gradient)
                return nothing
            end
        end
    end
    # Inner prox step
    anderson.g[anderson.data.current] .=  x - γ*∇f
    anderson.save_g .= anderson.g[anderson.data.current]
    _x = copy(x)
    prox!(quasigradient, anderson.prox, _x, ∇f, γ)
    anderson.grad_mapping .= x - _x
    anderson.x .= _x
    # Anderson update
    anderson_update!(anderson)
    # Proximal update on y
    x .= anderson.y[anderson.data.current]
    prox!(quasigradient, anderson.prox, x, ∇f, 0.0)
    anderson.data.Q = quasigradient.data.Q
    return nothing
end

function anderson_update!(anderson::AndersonAcceleratedProximal)
    mk = min(anderson.parameters.m, length(anderson.y))
    index_set = collect(1:mk)
    R = reduce(hcat, [anderson.g[i] - anderson.y[i] for i in reverse(index_set)])
    RR = transpose(R)*R
    RR ./= norm(RR)
    x = (RR + 1e-10I) \ ones(size(RR, 1))
    α = x / sum(x)
    # Update active index and memory
    anderson.data.current = (anderson.data.current % anderson.parameters.m) + 1
    # Reserve memory if still smaller than m
    if length(anderson.y) < anderson.parameters.m
        push!(anderson.y, zero(anderson.y[1]))
        push!(anderson.g, zero(anderson.y[1]))
    end
    anderson.y[anderson.data.current] .= sum([α[i]*anderson.g[mk-i+1] for i in index_set])
    push!(anderson.residuals, norm(R*α))
    return nothing
end

# API
# ------------------------------------------------------------
"""
    AndersonAcceleration

Factory object for [`AndersonAcceleratedProximal`](@ref). Pass to `prox` in `Quasigradient.Optimizer` or set the [`Prox`](@ref) attribute. See ?AndersonAcceleratedProximal for parameter descriptions.

"""
mutable struct AndersonAcceleration <: AbstractProx
    prox::AbstractProx
    parameters::AndersonAccelerationParameters
end
AndersonAcceleration(; prox::AbstractProx = Polyhedron(), kw...) = AndersonAcceleration(prox, AndersonAccelerationParameters(; kw...))

function (anderson::AndersonAcceleration)(structure::StageDecompositionStructure, x₀::AbstractVector, ::Type{T}) where T <: AbstractFloat
    proximal = anderson.prox(structure, x₀, T)
    return AndersonAcceleratedProximal(proximal, T; type2dict(anderson.parameters)...)
end

function str(::AndersonAcceleration)
    return ""
end


time_modeling(model::ModelAll, srcGeometry::Geometry, srcData, recGeometry::Geometry, recData, perturbation, srcnum::UnitRange{Int64}, op::Char, mode::Int64) =
	time_modeling(model, srcGeometry, srcData, recGeometry, recData, perturbation, srcnum, op, mode, Options())

function time_modeling(model::ModelAll, srcGeometry::Geometry, srcData, recGeometry::Geometry, recData, perturbation, srcnum::UnitRange{Int64}, op::Char, mode::Int64, options)
# time_modeling function for multiple sources. Depending on the operator and mode, this function distributes the sources
# and if applicable the input data amongst the available workers.

	p = default_worker_pool()
	time_modeling = remote(TimeModeling.time_modeling)
	# time_modeling = wrap_retry(time_modeling, options.retry_n)

	numSources = length(srcnum)
	results = Array{Any}(numSources)

	# Process shots from source channel asynchronously
	@sync begin
		for j=1:numSources

			# local geometry for current position
			srcGeometryLocal = subsample(srcGeometry,j)
			recGeometryLocal = subsample(recGeometry,j)

			# Parallelization
			if op=='F' && mode==1
				srcDataLocal = Array{Any}(1)
				srcDataLocal[1] = srcData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, srcDataLocal, recGeometryLocal, [], [], j, op, mode, options)
			elseif op=='F' && mode==-1
				recDataLocal = Array{Any}(1)
				recDataLocal[1] = recData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, [], recGeometryLocal, recDataLocal, [], j, op, mode, options)
			elseif op=='J' && mode==1
				srcDataLocal = Array{Any}(1)
				srcDataLocal[1] = srcData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, srcDataLocal, recGeometryLocal, [], perturbation, j, op, mode, options)
			elseif op=='J' && mode==-1
				srcDataLocal = Array{Any}(1)
				srcDataLocal[1] = srcData[j]
				recDataLocal = Array{Any}(1)
				recDataLocal[1] = recData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, srcDataLocal, recGeometryLocal, recDataLocal, [], j, op, mode, options)
			end
		end
	end

	if op=='F' || (op=='J' && mode==1)
		argout1 = results[1]
		for j=2:numSources
			argout1 = [argout1; results[j]]
		end
	elseif op=='J' && mode==-1
		argout1 = results[1]
		for j=2:numSources
			argout1 += results[j]
		end
	else
		error("operation no defined")
	end
	return argout1
end


macro cutensor(ex::Expr)
    cuwrapdict = Dict{Any,Any}()
    parser = TensorParser()
    parser.preprocessors[end] = ex->extracttensorobjects(ex, cuwrapdict)
    push!(parser.postprocessors, ex->addcutensorwraps(ex, cuwrapdict))
    return esc(parser(ex))
end

macro cutensor(ex::Expr, orderex::Expr)
    cuwrapdict = Dict{Any,Any}()
    parser = TensorParser()
    if !(orderex.head == :(=) && orderex.args[1] == :order &&
            orderex.args[2] isa Expr && orderex.args[2].head == :tuple)
        throw(ArgumentError("unkown first argument in @tensor, should be `order = (...,)`"))
    end
    indexorder = map(normalizeindex, orderex.args[2].args)
    parser.contractiontreebuilder = network->indexordertree(network, indexorder)
    parser.preprocessors[end] = ex->extracttensorobjects(ex, cuwrapdict)
    push!(parser.postprocessors, ex->addcutensorwraps(ex, cuwrapdict))
    return esc(parser(ex))
end

function extracttensorobjects(ex, cuwrapdict)
    inputtensors = getinputtensorobjects(ex)
    outputtensors = getoutputtensorobjects(ex)
    newtensors = getnewtensorobjects(ex)
    existingtensors = unique!(vcat(inputtensors, outputtensors))
    alltensors = unique!(vcat(existingtensors, newtensors))
    tensordict = Dict{Any,Any}(a => gensym() for a in alltensors)
    pre = Expr(:block, [Expr(:(=), tensordict[a], a) for a in existingtensors]...)
    cutensordict = Dict{Any,Any}(a => gensym() for a in alltensors)
    ex = replacetensorobjects((obj,leftind,rightind)->get(cutensordict, obj, obj), ex)
    post = Expr(:block,
                [Expr(:call, :copyto!, a, cutensordict[a]) for a in outputtensors]...,
                [Expr(:(=), a, cutensordict[a]) for a in newtensors]...)
    for k in inputtensors
        a = tensordict[k]
        b = cutensordict[k]
        push!(cuwrapdict, b=>(a,:($b = CuArray($a))))
    end
    for k in setdiff(outputtensors, inputtensors)
        a = tensordict[k]
        b = cutensordict[k]
        push!(cuwrapdict, b=>(a,:($b = CuArray{eltype($a)}(undef, size($a)))))
    end
    return Expr(:block, pre, ex, post)
end

function addcutensorwraps(ex, cuwrapdict)
    for (b,(a,assignb)) in cuwrapdict
        ex = _replace_in_eltype(ex, b, a)
        ex = _splice(ex, b, assignb)
    end
    return ex
end

function _replace_in_eltype(ex::Expr, b, a)
    if ex.head == :call && ex.args[1] == :eltype && ex.args[2] == b
        return Expr(:call, :eltype, a)
    else
        return Expr(ex.head, (_replace_in_eltype(arg, b, a) for arg in ex.args)...)
    end
end
_replace_in_eltype(ex, b, a) = ex

# check if a subexpression contains/uses the variable s, but ignore `eltype` calls
_contains(ex::Expr, s) = any(e->_contains(e, s), ex.args)
_contains(ex::Symbol, s) = ex == s
_contains(ex, s) = false

function _splice(ex::Expr, s1, assign)
    if ex.head == :block
        args = copy(ex.args)
        i = 1
        while i <= length(args)
            if _contains(args[i], s1)
                insert!(args, i, assign)
                break
            end
            i += 1
        end
        return Expr(ex.head, args...)
    elseif !(ex.head == :call && ex.args[1] == :throw)
        error("unexpected expression: $ex")
    end
end


export FormConstraintCheckEach, FormConstraintCheckLast, FormConstraintCheckPickDefault
export constrain_form, default_form_check_strategy, is_point_mass_form_constraint

using TupleTools

import Base: +

# Form constraints are preserved during execution of the `prod` function
# There are two major strategies to check current functional form
# We may check and preserve functional form of the result of the `prod` function
# after each subsequent `prod`
# or we may want to wait after all `prod` functions in the equality chain have been executed

abstract type AbstractFormConstraint end

"""
    FormConstraintCheckEach

This form constraint check strategy checks functional form of the messages product after each product in an equality chain.
Usually if a variable has been connected to multiple nodes we want to perform multiple `prod` to obtain a posterior marginal.
With this form check strategy `constrain_form` function will be executed after each subsequent `prod` function.

See also: [`FormConstraintCheckLast`](@ref), [`default_form_check_strategy`](@ref), [`constrain_form`](@ref), [`multiply_messages`](@ref)
"""
struct FormConstraintCheckEach end

"""
    FormConstraintCheckEach

This form constraint check strategy checks functional form of the last messages product in the equality chain.
Usually if a variable has been connected to multiple nodes we want to perform multiple `prod` to obtain a posterior marginal.
With this form check strategy `constrain_form` function will be executed only once after all subsequenct `prod` functions have been executed.

See also: [`FormConstraintCheckLast`](@ref), [`default_form_check_strategy`](@ref), [`constrain_form`](@ref), [`multiply_messages`](@ref)
"""
struct FormConstraintCheckLast end

"""
    FormConstraintCheckPickDefault

This form constraint check strategy simply fallbacks to a default check strategy for a given form constraint.

See also: [`FormConstraintCheckEach`](@ref), [`FormConstraintCheckLast`](@ref), [`default_form_check_strategy`](@ref)
"""
struct FormConstraintCheckPickDefault end

"""
    default_form_check_strategy(form_constraint)

Returns a default check strategy (e.g. `FormConstraintCheckEach` or `FormConstraintCheckEach`) for a given form constraint object.

See also: [`FormConstraintCheckEach`](@ref), [`FormConstraintCheckLast`](@ref), [`constrain_form`](@ref)
"""
function default_form_check_strategy end

"""
    is_point_mass_form_constraint(form_constraint)

Specifies whether form constraint always returns PointMass estimates or not. For a given `form_constraint` returns either `true` or `false`.

See also: [`FormConstraintCheckEach`](@ref), [`FormConstraintCheckLast`](@ref), [`constrain_form`](@ref)
"""
function is_point_mass_form_constraint end

"""
    constrain_form(form_constraint, something)

Checks that the functional form of `something` object fits `form_constraint` constraint.
If functional form of `something` object does not fit the given constraint this function
should try to approximate `something` object to be in line with the given `form_constraint`.

See also: [`FormConstraintCheckEach`](@ref), [`FormConstraintCheckLast`](@ref), [`default_form_check_strategy`](@ref), [`is_point_mass_form_constraint`](@ref)
"""
function constrain_form end

struct CompositeFormConstraint{C} <: AbstractFormConstraint
    constraints :: C
end

constrain_form(composite::CompositeFormConstraint, something) = reduce((form, constraint) -> constrain_form(constraint, form), composite.constraints, init = something)

function default_form_check_strategy(composite::CompositeFormConstraint)
    strategies = map(default_form_check_strategy, composite.constraints)
    if !(all(e -> e === first(strategies), TupleTools.tail(strategies)))
        error("Different default form check strategy for composite form constraints found. Use `form_check_strategy` options to specify check strategy.")
    end
    return first(strategies)
end

function is_point_mass_form_constraint(composite::CompositeFormConstraint)
    is_point_mass = map(is_point_mass_form_constraint, composite.constraints)
    pmindex       = findnext(is_point_mass, 1)
    if pmindex !== nothing && pmindex !== length(is_point_mass)
        error("Composite form constraint supports point mass constraint only at the end of the form constrains specification.")
    end
    return last(is_point_mass)
end

Base.:+(constraint::AbstractFormConstraint) = constraint

Base.:+(left::AbstractFormConstraint,  right::AbstractFormConstraint)  = CompositeFormConstraint((left, right))
Base.:+(left::AbstractFormConstraint,  right::CompositeFormConstraint) = CompositeFormConstraint((left, right.constraints...))
Base.:+(left::CompositeFormConstraint, right::AbstractFormConstraint)  = CompositeFormConstraint((left.constraints..., right))
Base.:+(left::CompositeFormConstraint, right::CompositeFormConstraint) = CompositeFormConstraint((left.constraints..., right.constraints...))


"""
    DynamicAggregation

Functor object for using dynamic aggregation in an L-shaped algorithm. Create by supplying a [`DynamicAggregate`](@ref) object through `aggregate ` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

The following selection rules are available
- [`SelectUniform`](@ref)
- [`SelectDecaying`](@ref)
- [`SelectRandom`](@ref
- [`SelectClosest`](@ref)
- [`SelectClosestToReference`](@ref)

...
# Parameters
- `naggregates::Int`: Number of aggregates
- `rule::SelectionRule`: Rule that determines which aggregate an incoming cut should be placed in
- `lock_after::Function = (τ,n)->false`: Function that determines if the current aggregation scheme should be fixed, based on the current optimality gap `τ` and the number of iterations `n`
...
"""
struct DynamicAggregation{T <: AbstractFloat, S <: SelectionRule} <: AbstractAggregation
    aggregates::Vector{AggregatedOptimalityCut{T}}
    rule::S
    partitioning::Dict{Int,Int}
    lock::Function

    function DynamicAggregation(naggregates::Integer, rule::SelectionRule, lock::Function, ::Type{T}) where T <: AbstractFloat
        S = typeof(rule)
        aggregates = [zero(AggregatedOptimalityCut{T}) for _ = 1:naggregates]
        return new{T,S}(aggregates, rule, Dict{Int,Int}(), lock)
    end
end

function aggregate_cut!(lshaped::AbstractLShapedSolver, ::DynamicAggregation, cut::HyperPlane)
    return add_cut!(lshaped, cut)
end

function aggregate_cut!(lshaped::AbstractLShapedSolver, aggregation::DynamicAggregation{T}, cut::HyperPlane{OptimalityCut}) where T <: AbstractFloat
    if aggregation.lock(gap(lshaped),niterations(lshaped)) && haskey(aggregation.partitioning, cut.id)
        aggregation.aggregates[aggregation.partitioning[cut.id]] += cut
        return false
    end
    (idx, full) = select(aggregation.rule, aggregation.aggregates, cut)
    aggregation.aggregates[idx] += cut
    aggregation.partitioning[cut.id] = idx
    if full
        aggregate = aggregation.aggregates[idx]
        aggregation.aggregates[idx] = zero(AggregatedOptimalityCut{T})
        return add_cut!(lshaped, aggregate)
    end
    return false
end

function aggregate_cut!(cutqueue::CutQueue, aggregation::DynamicAggregation, ::MetaData, t::Integer, cut::HyperPlane, x::AbstractArray)
    put!(cutqueue, (t, cut))
    return nothing
end

function aggregate_cut!(cutqueue::CutQueue, aggregation::DynamicAggregation{T}, metadata::MetaData, t::Integer, cut::HyperPlane{OptimalityCut}, x::AbstractArray) where T <: AbstractFloat
    gap = fetch(metadata, t, :gap)
    if aggregation.lock(gap, t) && haskey(aggregation.partitioning, cut.id)
        aggregation.aggregates[aggregation.partitioning[cut.id]] += cut
        return nothing
    end
    (idx, full) = select(aggregation.rule, aggregation.aggregates, cut)
    aggregation.partitioning[cut.id] = idx
    aggregation.aggregates[idx] += cut
    if full
        aggregate = aggregation.aggregates[idx]
        put!(cutqueue, (t, aggregate))
        aggregation.aggregates[idx] = zero(AggregatedOptimalityCut{T})
    end
    return nothing
end

function nthetas(nscenarios::Integer, ::DynamicAggregation)
    return nscenarios
end

function nthetas(nscenarios::Integer, ::DynamicAggregation, ::AbstractScenarioProblems)
    return nscenarios
end

function flush!(lshaped::AbstractLShapedSolver, aggregation::DynamicAggregation{T}) where T <: AbstractFloat
    added = false
    for (i,aggregate) in enumerate(aggregation.aggregates)
        if !iszero(aggregate)
            added |= add_cut!(lshaped, aggregate)
            aggregation.aggregates[i] = zero(AggregatedOptimalityCut{T})
        end
    end
    reset!(aggregation.rule)
    return added
end

function flush!(cutqueue::CutQueue, aggregation::DynamicAggregation{T}, ::MetaData, t::Integer, x::AbstractArray) where T <: AbstractFloat
    for (i,aggregate) in enumerate(aggregation.aggregates)
        if !iszero(aggregate)
            put!(cutqueue, (t, aggregate))
            aggregation.aggregates[i] = zero(AggregatedOptimalityCut{T})
        end
    end
    reset!(aggregation.rule)
    return nothing
end

# API
# ------------------------------------------------------------
"""
    DynamicAggregate(naggregates::Integer, rule::SelectionRule; lock_after::Function = (τ,n)->false)

Factory object for [`DynamicAggregation`](@ref). Pass to `aggregate ` in the `LShapedSolver` factory function. See ?DynamicAggregation for parameter descriptions.

"""
struct DynamicAggregate{S <: SelectionRule} <: AbstractAggregator
    naggregates::Int
    rule::S
    lock::Function

    function DynamicAggregate(naggregates::Integer, rule::SelectionRule; lock_after = (τ,n)->false)
        S = typeof(rule)
        return new{S}(naggregates, rule, lock_after)
    end
end

function (aggregator::DynamicAggregate)(nscenarios::Integer, T::Type{<:AbstractFloat})
    if aggregator.rule isa SelectRandom && aggregator.rule.max == 1
        return NoAggregation()
    end
    return DynamicAggregation(aggregator.naggregates, aggregator.rule, aggregator.lock, T)
end

function remote_aggregator(aggregation::DynamicAggregation, ::AbstractScenarioProblems, ::Integer)
    return DynamicAggregate(length(aggregation.aggregates), aggregation.rule; lock_after = aggregation.lock)
end

function str(aggregator::DynamicAggregate)
    return "dynamic aggregation ruled by $(str(aggregator.rule))"
end


 #-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*
#*---------------------------------------------------------------------
#* Author: Estefane Horn, N�bia do Carmo Ferreira
#*$Id$
#*-------------------------------------------------------------------
#-------------------------------------------------------------------
#* Model of a valve (simplified)
#*--------------------------------------------------------------------
#*
#* Author: Paula B. Staudt
#*--------------------------------------------------------------------
type valve_flow
	valve_flow()=begin
		PP=outers.PP
		NComp=outers.NComp
		new(
			DanaPlugin(Dict{Symbol,Any}(
				:Brief=>"External Physical Properties",
				:Type=>"PP"
			)),
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of Components"
			)),
			flow_mol(Dict{Symbol,Any}(
				:Brief=>"Minimum Flow Allowable in the valve",
				:Default=>0
			)),
			flow_mol(Dict{Symbol,Any}(
				:Brief=>"Maximum Flow Allowable in the valve",
				:Default=>1000
			)),
			stream(Dict{Symbol,Any}(
				:Brief=>"Inlet stream",
				:PosX=>0,
				:PosY=>0.7365,
				:Symbol=>"_{in}",
				:Protected=>true
			)),
			stream(Dict{Symbol,Any}(
				:Brief=>"Outlet stream",
				:PosX=>1,
				:PosY=>0.7365,
				:Symbol=>"_{out}",
				:Protected=>true
			)),
			control_signal(Dict{Symbol,Any}(
				:Brief=>"Flow Signal",
				:PosX=>0.5,
				:PosY=>0,
				:Protected=>true
			)),
			[
				:(Outlet.F = Inlet.F),
				:(Outlet.T = Inlet.T),
				:(Outlet.P = Inlet.P),
				:(Outlet.h = Inlet.h),
				:(Outlet.v = Inlet.v),
				:(Outlet.z = Inlet.z),
				:(Outlet.F = MinFlow + FlowFraction*(MaxFlow-MinFlow)),
			],
			[
				"Overall Molar Balance","Temperature","Pressure","Energy Balance","Vapour fraction","Componente Molar Balance","Flow computation",
			],
			[:PP,:NComp,:MinFlow,:MaxFlow,],
			[:Inlet,:Outlet,:FlowFraction,]
		)
	end
	PP::DanaPlugin
	NComp::DanaInteger
	MinFlow::flow_mol
	MaxFlow::flow_mol
	Inlet::stream
	Outlet::stream
	FlowFraction::control_signal
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export valve_flow
function setEquationFlow(in::valve_flow)
	addEquation(1)
	addEquation(2)
	addEquation(3)
	addEquation(4)
	addEquation(5)
	addEquation(6)
	addEquation(7)
end
function atributes(in::valve_flow,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=true
	fields[:Icon]="icon/Valve"
	fields[:Brief]="Model of a very simple valve for setting the flow with a controller."
	fields[:Info]="== ASSUMPTIONS ==
* nothing happens in this valve

== SET ==
* MinFlow: the Minimum Flow Allowable in the valve;
* MaxFlow: the Maximum Flow Allowable in the valve;

== SPECIFY ==
* the Inlet stream
* the FlowFraction (the model requires an inlet signal, also you can use a controller for setting the FlowFraction)
"
	drive!(fields,_)
	return fields
end
valve_flow(_::Dict{Symbol,Any})=begin
	newModel=valve_flow()
	newModel.attributes=atributes(newModel,_)
	newModel
end


#= ForwardDiff is removed as a dependency in this version.
import MechanicalSketch
import MechanicalSketch: PALETTE, O, HE, WI, EM, FS, PT, ∙, settext
import MechanicalSketch: @import_expand, set_scale_sketch
import MechanicalSketch: generate_complex_potential_source, generate_complex_potential_vortex
import MechanicalSketch: clamped_velocity_matrix
import MechanicalSketch: place_image
import MechanicalSketch: LicSceneOpts
import MechanicalSketch: ColorSchemes, color_with_lumin
import MechanicalSketch: BinLegend, setfont, FS, fontsize, background, sethue, dimension_aligned, line
import MechanicalSketch: Movie, Scene, animate, color_matrix_current, draw_legend, @layer
import MechanicalSketch: fontsize, fontface, setfont, color_with_lumin
import MechanicalSketch: convolution_matrix, ColorSchemes
import ColorSchemes:     isoluminant_cgo_80_c38_n256

let

if !@isdefined m²
    @import_expand(~m, s, °)
end
include("test_functions_33.jl")

# Scaling and placement
physwidth = 10.0m
physheight = 4.0m
framedheight = physheight * 1.2
totheight = framedheight * 3
set_scale_sketch(totheight, HE)
totwidth = totheight * WI / HE
Δx = totwidth / 4
Δy = framedheight
Oadj = O + (0, EM /3)
cent1 = Oadj + (-Δx, Δy)
cent2 = Oadj + ( Δx, Δy)
cent3 = Oadj + (-Δx, 0.0m)
cent4 = Oadj + ( Δx, 0.0m)
cent5 = Oadj + (-Δx, -Δy)

# Reused velocity field from earlier tests
max_velocity = 0.5m/s
velocity_matrix = clamped_velocity_matrix(ϕ_33; physwidth, physheight, cutoff = max_velocity);
speedmatrix = hypot.(velocity_matrix)
c1 = convolution_matrix(velocity_matrix)
lum = 50
cols = isoluminant_cgo_80_c38_n256.colors .|> co -> color_with_lumin(co, lum)
legen = BinLegend(maxlegend = max_velocity, binwidth = 0.1m/s, colorscheme = cols, name = :Speed)

function scenepic(scene, framenumber)
    colmat = color_matrix_current(scene, framenumber)
    place_image(scene.opts.O, colmat)
end

function sc1(scene, framenumber)
    ulp, lrp = scenepic(scene, framenumber)
    str = "    1: Luminosity variation = 1 / 6"
    settext(str, ulp, markup = true)

end
function sc2(scene, framenumber)
    ulp, _ = scenepic(scene, framenumber)
    str = "    2: Luminosity variation = 2 / 6"
    settext(str, ulp, markup = true)
end

function sc3(scene, framenumber)
    ulp, _ = scenepic(scene, framenumber)
    str = "    3: Luminosity variation = 3 / 6"
    settext(str, ulp, markup = true)
end
function sc4(scene, framenumber)
    ulp, _ = scenepic(scene, framenumber)
    str = "    4: Luminosity variation = 4 / 6"
    settext(str, ulp, markup = true)
end

function sc5(scene, framenumber)
    ulp, lrp = scenepic(scene, framenumber)
    str = "    5: Luminosity variation = 5 / 6"
    settext(str, ulp, markup = true)
    # Centre of this scene's rectangle
    O = scene.opts.O
    # Top left of legend
    legendpos = O + ((lrp[1] - ulp[1]) / 2 + EM, (ulp[2] - lrp[2]) / 2)
    @layer begin
        sethue("white")
        # Font for the 'toy' text interface
        fontface("Calibri")
        # 1 pt font size = 12/72 inch - by the book.
        # Letter spacing works differently here than in Word, so we adjust a little.
        fontsize(FS)
        draw_legend(legendpos, scene.opts.legend)
    end
end

function backdrop(scene, framenumber)
    setfont("Calibri", FS)
    fontsize(FS)
    O = scene.opts.O
    background(color_with_lumin(PALETTE[6], 30))
    sethue(PALETTE[3])
    Δt = scene.opts.cycle_duration
    frames_per_cycle = Δt ∙ scene.opts.framerate
    # Where we are in the repeating cycle, [0, 1 - frame_duration]
    normalized_time = framenumber / frames_per_cycle
    t = Δt ∙ normalized_time

    dimension_aligned(O + (-2.0m, physheight / 2), O + (-1.0m, physheight / 2),
        fromextension = (0, EM), toextension = (0, EM), offset = -EM)
    dimension_aligned(O + (-3.0m, physheight / 2), O + (-2.0m, physheight / 2),
        fromextension = (0, EM), toextension = (0, EM), offset = -EM)
    dimension_aligned(O + (-4.0m, physheight / 2), O + (-3.0m, physheight / 2),
        fromextension = (0, EM), toextension = (0, EM), offset = -EM)
    for i = 0:4
        x = t ∙ 0.5m/s + i ∙ 1.0m - 4.0m
        line(O + (x, 1.1 * physheight / 2), O + (x, physheight / 2), :stroke)
    end
end
movie = Movie(WI, HE, "Five flow fields with varying luminosity variation", 15:15)
Δt = 2.0s
framerate = 30/s
endframe = Δt ∙ framerate - 1
o1 = LicSceneOpts(Δt,  cent1; phase_magnitude_matrix = c1, data = speedmatrix, legend = legen, luminosity_variation = 1 / 6)
o2 = LicSceneOpts(Δt,  cent2; phase_magnitude_matrix = c1, data = speedmatrix, legend = legen, luminosity_variation = 2 / 6)
o3 = LicSceneOpts(Δt,  cent3; phase_magnitude_matrix = c1, data = speedmatrix, legend = legen, luminosity_variation = 3 / 6)
o4 = LicSceneOpts(Δt,  cent4; phase_magnitude_matrix = c1, data = speedmatrix, legend = legen, luminosity_variation = 4 / 6)
o5 = LicSceneOpts(Δt,  cent5; phase_magnitude_matrix = c1, data = speedmatrix, legend = legen, luminosity_variation = 5 / 6)

scenes = [  Scene(movie, backdrop, 0:endframe, optarg = LicSceneOpts(Δt, O)),
            Scene(movie, sc1, 0:endframe, optarg = o1),
            Scene(movie, sc2, 0:endframe, optarg = o2),
            Scene(movie, sc3, 0:endframe, optarg = o3),
            Scene(movie, sc4, 0:endframe, optarg = o4),
            Scene(movie, sc5, 0:endframe, optarg = o5)];

animate(movie, scenes;
    creategif = true,
    pathname = joinpath(@__DIR__, "test_34.png"),
    framerate = framerate * s)

end # Let
=#


# Floating point data format invariants

const SIGN_BIT = 0x80000000
const SIGN_BIT_64 = UInt64(SIGN_BIT)

const bmask16 = 0xFFFF
const bmask32 = 0xFFFFFFFF

const UNO = one(UInt32)
const UNO64 = one(UInt64)

# VAX floating point data formats (see VAX Architecture Reference Manual)

const VAX_F_SIGN_BIT      = SIGN_BIT
const VAX_F_EXPONENT_MASK = Int32(0x7F800000)
const VAX_F_EXPONENT_SIZE = UInt32(8)
const VAX_F_EXPONENT_BIAS = Int32(128)
const VAX_F_MANTISSA_MASK = UInt32(0x007FFFFF)
const VAX_F_MANTISSA_SIZE = UInt32(23)
const VAX_F_HIDDEN_BIT    = UInt32( UNO << VAX_F_MANTISSA_SIZE )

const VAX_D_EXPONENT_MASK = Int64(VAX_F_EXPONENT_MASK)
const VAX_D_EXPONENT_SIZE = UInt64(VAX_F_EXPONENT_SIZE)
const VAX_D_EXPONENT_BIAS = Int64(VAX_F_EXPONENT_BIAS)
const VAX_D_MANTISSA_MASK = UInt64(VAX_F_MANTISSA_MASK)
const VAX_D_MANTISSA_SIZE = UInt64(VAX_F_MANTISSA_SIZE)
const VAX_D_HIDDEN_BIT    = UInt64(VAX_F_HIDDEN_BIT)

const VAX_G_EXPONENT_MASK = Int64(0x7FF00000)
const VAX_G_EXPONENT_SIZE = UInt64(11)
const VAX_G_EXPONENT_BIAS = Int64(1024)
const VAX_G_MANTISSA_MASK = UInt64(0x000FFFFF)
const VAX_G_MANTISSA_SIZE = UInt64(20)
const VAX_G_HIDDEN_BIT    = UInt64( UNO << VAX_G_MANTISSA_SIZE )

# IEEE floating point data formats (see Alpha Architecture Reference Manual)

const IEEE_S_SIGN_BIT      = SIGN_BIT
const IEEE_S_EXPONENT_MASK = Int32(0x7F800000)
const IEEE_S_EXPONENT_SIZE = UInt32(8)
const IEEE_S_EXPONENT_BIAS = Int32(127)
const IEEE_S_MANTISSA_MASK = UInt32(0x007FFFFF)
const IEEE_S_MANTISSA_SIZE = UInt32(23)
const IEEE_S_HIDDEN_BIT    = UInt32( UNO << IEEE_S_MANTISSA_SIZE )

const IEEE_T_EXPONENT_MASK = Int64(0x7FF00000)
const IEEE_T_EXPONENT_SIZE = UInt64(11)
const IEEE_T_EXPONENT_BIAS = Int64(1023)
const IEEE_T_MANTISSA_MASK = UInt64(0x000FFFFF)
const IEEE_T_MANTISSA_SIZE = UInt64(20)
const IEEE_T_HIDDEN_BIT    = UInt64( 1 << IEEE_T_MANTISSA_SIZE )


"""
$(TYPEDEF)

Partial prior belief (absolute data) on any variable, given `<:SamplableBelief` and which dimensions of the intended variable.

Notes
- If using [`AMP.ManifoldKernelDensity`](@ref), don't double partial.  Only define the partial in this `PartialPrior` container.
  - Future TBD, consider using `AMP.getManifoldPartial` for more general abstraction.
"""
struct PartialPrior{T <: SamplableBelief,P <: Tuple} <: AbstractPrior
  Z::T
  partial::P
end

# TODO, standardize, but shows error on testPartialNH.jl
getSample(cf::CalcFactor{<:PartialPrior}) = samplePoint(cf.factor.Z)   # remove in favor of ManifoldSampling.jl
# getManifold(pp::PartialPrior) = TranslationGroup(length(pp.partial)) # uncomment

getManifold(pp::PartialPrior{<:PackedManifoldKernelDensity}) = pp.Z.manifold

"""
$(TYPEDEF)

Serialization type for `PartialPrior`.
"""
Base.@kwdef struct PackedPartialPrior <: AbstractPackedFactor
  Z::PackedSamplableBelief
  partials::Vector{Int}
end

function convert(::Type{PackedPartialPrior}, d::PartialPrior)
  PackedPartialPrior(convert(PackedSamplableBelief, d.Z), [d.partial...;])
end
function convert(::Type{PartialPrior}, d::PackedPartialPrior)
  PartialPrior(convert(SamplableBelief, d.Z),(d.partials...,))
end


#!/usr/bin/env julia

import Celeste.AccuracyBenchmark
import Celeste.GalsimBenchmark

test_case_names = String[]
for arg in ARGS
    if ismatch(r"^--test_case_names", arg)
        # test case name: function name in test_case_definitions.py, or CL_DESCR
        # field from FITS header
        test_case_names =
            split(replace(arg, "--test_case_names=", ""), ",")
    else
        println("Usage: run_galsim_benchmark.jl " *
                "[--test_case_names=test1,test2,...]")
        error("Invalid argument: ", arg)
    end
end

srand(12345)

truth_catalog, single_predictions = GalsimBenchmark.run_benchmarks(
    test_case_names=test_case_names,
    joint_inference=false,
)
unused, joint_predictions = GalsimBenchmark.run_benchmarks(
    test_case_names=test_case_names,
    joint_inference=true,
)

score_df = AccuracyBenchmark.score_predictions(
    truth_catalog,
    [single_predictions, joint_predictions],
)
println(repr(score_df))


#cd("/home/user/github/TOF-Tracer2/")

include("$(pwd())/startupAPiTOF.jl")

fp = "./ExampleFiles/APiTOFDATA/" # All files in this path will be processed
filefilterRegexp = r"\.h5$"
rf = "./ExampleFiles/APiTOFDATA/Data_09_46_20_04mm.h5"  # The mass scale from this file defines the mass scale of all

masslist = MasslistFunctions.loadMasslist("./ExampleFiles/MASSLISTS/exampleMasslist.csv")
cr = [37.028406 55.038971 282.118343] # pos
#cr = [62 125] # neg

# alternatively: use an auto generated masslist
# masslistMasses, masslistElements, masslistElementsMasses, masslistCompositions = createMassList(C=0:20, O=0:20, N=0:1, allowRadicals=false) #
(masslistMasses, masslistElements, masslistElementsMasses, masslistCompositions) = masslist
s = (masslistMasses.>17) .& ( masslistMasses.<900)
masslistMasses = masslistMasses[s]
masslistCompositions = masslistCompositions[s,:]

####################### END OF SETTINGS ###############################################################

####################### Processing sequence ###########################################################

correctMassScaleAndExtractSumSpec(
    fp,
    masslistMasses,
    masslistElements,
    masslistElementsMasses,
    masslistCompositions,
    rf,
    cr,
    filefilterRegexp=filefilterRegexp,
    onlyUseAverages = true,
    plotControlMass = true,
    testRangeStart = 54.5, # the mass shift of this region will be shown if plot control mass is set true. Should not be part of calibRegions
    testRangeEnd = 55.5,
    recalibInterval = 900,
    resolution = 5500,
    firstNFiles= 1,
    lastNFiles = 0,
    UMRpeaks = true, # create Unit Mass Resolution peaks
    UMRmasses = length(s),
    UMRmassLow = 0.2,
    UMRmassHigh = 0.4,
    umrBinWidth =100
    #binWidth = 10
    )
#
baselineAndPeakshape(
    fp,
    peakshapeRegions=4,
    peakshapeRegionStretch=1.0,
    peakshapeQuantileValue = 0.7,
    peakfindingNoiseThresholdValue = 5,
    peakWindowWidth = 100,
    peakfindingSignalLimit=0.5
    )
    #
#
mtrx = deconvolute(
    fp,
    calcTransposed = false
    )
#