module XInts

export XInt

using Base.GMP: Limb, BITS_PER_LIMB, SLimbMax, ULimbMax, ClongMax, CulongMax, CdoubleMax
import Base.GMP: isneg, ispos
import Base.GMP.MPZ
import Base.MPFR
using Base.GC: @preserve
using Base: BitInteger
import Base: +, -, *, ^, &, |, ==, /, ~, <<, >>, >>>, <, <=,
             string, widen, hastypemax, tryparse_internal,
             unsafe_trunc, trunc, mod, rem, iseven, isodd, gcd, lcm, xor, div, fld, cld,
             invmod, count_ones, trailing_zeros, trailing_ones, cmp, isqrt,
             flipsign, powermod, gcdx, promote_rule, factorial, binomial,
             digits!, ndigits0zpb, signbit, sign, iszero, isone,
             AbstractFloat, BigFloat, float, _prevpow2, _nextpow2,
             hash, hash_integer, sum, divrem

mutable struct Wrap
    b::BigInt
    d::Ptr{Limb}

    function Wrap(isoutput::Bool=false)
        b = BigInt()
        @assert b.alloc == 1
        @assert b.size == 0
        wrap = new(b, b.d)
        if !isoutput
            # init! was called
            finalizer(wrap) do wrap
                wrap.b.d = wrap.d
                wrap.b.alloc = 1
                wrap.b.size = 0
            end
        end
        wrap
    end
end

if BITS_PER_LIMB == 32
    const SLimb = Int32
    const SLimbW = Int64
    const LimbW = UInt64

    @inline _divLimb(n) = n >>> 5
    @inline _modLimb(n) = n & 31
    @inline _divByteLimb(n) = n >>> 2
    @inline _modByteLimb(n) = n & 3
elseif BITS_PER_LIMB == 64
    const SLimb = Int64
    const SLimbW = Int128
    const LimbW = UInt128

    @inline _divLimb(n) = n >>> 6
    @inline _modLimb(n) = n & 63
    @inline _divByteLimb(n) = n >>> 3
    @inline _modByteLimb(n) = n & 7
else
    error()
end

const slimbmin = typemin(SLimb)
const slimbmax = typemax(SLimb)
# when there is only one limb in a vector, this limb must always be >= limb1min
# reciprocally, a direct integer must have its absolute value < limb1min
const limb1min = slimbmin % Limb # typically 0x8000000000000000

const Short = Union{Limb,SLimb}
const ShortMax = Union{ULimbMax,SLimbMax}

abstract type XSigned <: Signed end

# T must be a "bit integer", assumed for efficiency, and also sizeof
struct Static{T<:Integer} <: XSigned
    x::T
end

# we could re-use Static for this, but Static <: Number already has getindex defined,
# so we clearly separate the two concepts
struct SVec{U<:Unsigned}
    u::U
end

is_short(s::Static) = sizeof(s.x) < sizeof(Limb) ? true : fits(s.x)
short(s::Static) = s.x % SLimb

vec(s::Static{<:Signed}) = SVec(unsigned(abs(s.x)))
vec(s::Static{<:Unsigned}) = SVec(s.x)

function lenvec(s::Static)
    # @assert !is_short(s)
    v = vec(s)
    lz = leading_zeros(v.u) >>> 3 # count leading full zero bytes
    usedbytes = sizeof(v.u)-lz
    _divByteLimb(usedbytes - 1) + 1, v
end

len(s::Static) = lenvec(s)[1]
vec(s::Static) = lenvec(s)[2]

iszero(s::Static) = iszero(s.x)
ispos(s::Static) = s.x > 0
signbit(s::Static) = s.x < 0
_sign(s::Static) = ispos(s) ? SLimb(1) : SLimb(-1)

function Base.getindex(s::SVec, i::Integer)
    # getindex is called only when sizeof(x) >= sizeof(Limb)
    if sizeof(s.u) == sizeof(Limb)
        # we always assume @inbounds
        s.u % Limb
    else
        (s.u >> ((i-1)*BITS_PER_LIMB)) % Limb
    end
end

Base.show(io::IO, s::Static) =
    if is_short(s)
        print(io, "Static($(s.x))")
    else
        sl, sv = lenvec(s)
        print(io, "Static(", [sv[i] for i=1:sl], ")")
    end

struct XInt <: XSigned
    x::SLimb # immediate integer or sign+length
    v::Union{Nothing,Vector{Limb}}

    @inline XInt(x::SLimb) =
        x === typemin(SLimb) ?
            new(-one(SLimb), [limb1min]) :
            new(x, nothing)

    global _XInt

    # unsafe version which doesn't check for typemin(SLimb)
    _XInt(x::SLimb, ::Nothing=nothing) = new(x, nothing)

    function _XInt(x::SLimb, v::Vector{Limb}, reduce::Bool=false)
        if reduce
            # we still assume length(v) >= abs(x); reduce here means to not store
            # a too small integer in a vector representation
            xl = abs(x)
            if xl <= 1
                if xl == 0
                    return new(SLimb(0), nothing)
                else # xl == 1
                    limb = @inbounds v[1]
                    limb < limb1min &&
                        return new(flipsign(limb % SLimb, x), nothing)
                end
            end
        end
        new(x, v)
    end
end

XInt(x::XInt) = x
XInt(x::Static) = XInt(x.x)

# reducing version of XInt(::XInt)
_XInt(x::XInt) =
    if x.v === nothing
        x
    else
        _XInt(x.x, vec(x), true)
    end

_XInt(u::Limb) = _XInt(u % SLimb) # no check done that u fits

is_short(x::XInt) = x.v === nothing
is_short(x::XInt, y::XInt) = x.v === nothing === y.v

short(x::XInt) = x.x
len(x::XInt) = abs(x.x)

include("MPN.jl")
include("mutants.jl")

MPN.ptr(x::XInt) = pointer(x.v::Vector{Limb})
MPN.len(x::XInt) = abs(x.x) % MPN.Size

function _XInt(z::BigInt, extra::SLimb=0)
    # @assert extra >= 0
    len = abs(z.size)
    if len == 0
        XInt(zero(SLimb))
    elseif len == 1 && (x = unsafe_load(z.d)) < limb1min
        XInt(flipsign(x % SLimb, z.size))
    else
        x = _XInt(z.size % SLimb, Vector{Limb}(undef, len+extra))
        @preserve x z unsafe_copyto!(pointer(x.v), z.d, len)
        x
    end
end

XInt(z::BigInt) = _XInt(z)

fits(z::Limb) = !Core.is_top_bit_set(z) # z < limb1min
fits(z::SLimb) = z !== slimbmin
fits(z::Signed) = slimbmin < z <= slimbmax
fits(z::Unsigned) = z <= slimbmax

XInt(z::Limb) = XInt!(nothing, z)

# only SLimb and Limb in this Union need to be validated, but are taken care of by
# other more specific constructors
XInt(z::Union{ULimbMax,SLimbMax}) = _XInt(z % SLimb)

XInt(z::Union{SLimbW,LimbW}) = XInt!(nothing, z)
_XInt(z::Union{SLimbW,LimbW}, extra::SLimb=0) = XInt!(nothing, z, extra)

# TODO: this can be made quite faster for "bit integers": for some reason, allocating
# limbs according to sizeof(z) (and breaking the loop when iszero(zz))
# is more efficient than via ndigits (although ndigits seems to be fast)
# NOTE: also, it allocates more than expected for e.g. Int256, and this disappears
# if the function is re-evaled (then we get 1 allocation as expected)
function _XInt(z::Integer, extra::SLimb=0)::XInt
    slimbmin < z <= slimbmax &&
        return _XInt((z % SLimb)::SLimb)
    nd = ndigits(z, base=2)
    xl = _divLimb(nd)
    xl += !iszero(_modLimb(nd))
    limbs = _vec(xl+extra)
    zz = abs(z)
    for i=1:xl
        @inbounds limbs[i] = zz % Limb
        zz >>>= BITS_PER_LIMB
    end
    _XInt(flipsign(xl, z), limbs)
end

@inline XInt(z::Integer) = _XInt(z)

function XInt(x::Float64)
    isinteger(x) || throw(InexactError(:XInt, XInt, x))
    unsafe_trunc(XInt, x)
end

# no Union, to disambiguate
XInt(x::Float16) = XInt(Float64(x))
XInt(x::Float32) = XInt(Float64(x))

function init!(z::Wrap, x::XInt)
    b = z.b
    if is_short(x)
        b.alloc = 1
        b.size = sign(x.x)
        b.d = z.d
        unsafe_store!(z.d, abs(x.x) % Limb)
    else
        len = x.x
        b.alloc = abs(len)
        b.size = len
        b.d = pointer(x.v)
    end
    z.b::BigInt
end

struct Wraps
    # read-write temporary storage (output arguments)
    U::Wrap
    V::Wrap
    W::Wrap
    # read-only shallow wrappers of BigInt objects ("const" arguments in libgmp)
    X::Wrap
    Y::Wrap
    Z::Wrap

    Wraps() = new(Wrap(true), Wrap(true), Wrap(true), Wrap(), Wrap(), Wrap())
end

const tWraps = Wraps[]

mutable struct Buffer
    vec::Vector{Limb}
    # TODO: determine whether `used` is really needed; it's to try to work
    # around the problem mentioned in Julia#17015, that affects also usage
    # of the @bigint macro, related to GC
    used::Bool

    Buffer() = new(sizehint!(Limb[], 20), false)
end

const tBuffer = Buffer[]

function tmpvec(n::Integer)
    buf = tBuffer[Threads.threadid()]
    if buf.used
        # if a previous call failed to be followed by freetmpvec,
        # or maybe finalizers triggered by GC call tmpvec while the
        # resource is already in use
        _vec(n)
    else
        v = vec!(buf.vec, n)
        buf.used = true # after possible allocation above, which might fail
        v
    end
end

function freetmpvec(v::Vector)
    buf = tBuffer[Threads.threadid()]
    if v === buf.vec
        @assert buf.used
        buf.used = false
    end
    nothing
end

function __init__()
    copy!(tWraps, [Wraps() for _=1:Threads.nthreads()])
    copy!(tBuffer, [Buffer() for _=1:Threads.nthreads()])
end

macro bigint(args...)
    @assert length(args) >= 2
    out = args[1]
    xs = esc.(args[2:end-1])
    body = args[end]

    lets = :(let
                 $(esc(body))
             end)
    vars = lets.args[1].args

    @gensym mpzs
    push!(vars, :($(esc(mpzs)) = tWraps[Threads.threadid()]))
    if !(out isa Expr)
        out = Expr(:tuple, out)
    end
    for (i, outi) in enumerate(out.args)
        push!(vars, :($(esc(outi)) = getfield($(esc(mpzs)), $i).b))
    end
    for (i, xi) in enumerate(xs)
        push!(vars, :($xi = init!(getfield($(esc(mpzs)), $i+3), $xi)))
    end
    Expr(:gc_preserve, lets, xs...)
end

function BigInt(x::XInt)
    if is_short(x)
        BigInt(x.x)
    else
        xl, xv = lenvec(x)
        nbits = xl*8*sizeof(Limb)
        b = VERSION >= v"1.3" ? BigInt(; nbits=nbits) :
                                MPZ.realloc2(nbits)
        @preserve b xv unsafe_copyto!(b.d, pointer(xv), xl)
        b.size = x.x
        b
    end
end

Base.show(io::IO, x::XInt) = print(io, string(x))

function string(n::XInt; base::Integer = 10, pad::Integer = 1)
    is_short(n) && return string(n.x; base=base, pad=pad)
    base < 0 && return Base._base(Int(base), n, pad, (base>0) & isneg(n))
    2 <= base <= 62 || throw(ArgumentError("base must be 2 ≤ base ≤ 62, got $base"))
    nd1 = ndigits(n, base=base)
    nd  = max(nd1, pad)
    sv  = Base.StringVector(nd + isneg(n))
    xl, xv = lenvec(n)
    if !ispow2(base)
        # MPN.get_str might clobber the input
        xv = copy(xv)
    end
    str_size = @preserve sv Int(MPN.get_str(pointer(sv) + nd - nd1 + isneg(n), base, xv, xl))
    @assert str_size == nd1
    digits = base <= 36 ? Base.base36digits : Base.base62digits
    @inbounds for i = (1:nd-nd1) .+ isneg(n)
        sv[i] = '0' % UInt8
    end
    @inbounds for j = nd-nd1+isneg(n)+1:lastindex(sv)
        sv[j] = digits[1+sv[j]]
    end
    isneg(n) && (sv[1] = '-' % UInt8)
    String(sv)
end

widen(::Type{XInt}) = XInt

hastypemax(::Type{XInt}) = false

promote_rule(::Type{XInt}, ::Type{<:Integer}) = XInt
promote_rule(::Type{XInt}, ::Type{BigInt}) = XInt
promote_rule(::Type{BigInt}, ::Type{XInt}) = XInt
promote_rule(::Type{XInt}, ::Type{<:AbstractFloat}) = BigFloat

tryparse_internal(::Type{XInt}, s::AbstractString, startpos::Int, endpos::Int,
                  base_::Integer, raise::Bool) =
                      XInt(tryparse_internal(BigInt, s, startpos, endpos, base_, raise))

unsafe_trunc(::Type{XInt}, x::Union{Float32,Float64}) = XInt(unsafe_trunc(BigInt, x))

function trunc(::Type{XInt}, x::Union{Float32,Float64})
    isfinite(x) || throw(InexactError(:trunc, XInt, x))
    unsafe_trunc(XInt, x)
end

rem(x::XInt, ::Type{Bool}) =
    is_short(x) ? x.x % Bool :
                  (@inbounds x.v[1]) % Bool

rem(x::XInt, ::Type{T}) where T<:Union{SLimbMax,ULimbMax} =
    is_short(x) ? x.x % T : flipsign((@inbounds x.v[1] % T), x.x)

# TODO: extend to "AbstractBitInteger", when available in Base
function rem(x::XInt, ::Type{T}) where T<:Union{Base.BitUnsigned,Base.BitSigned}
    is_short(x) && return x.x % T
    u = zero(T)
    # TODO: this is missing an optimization that BigInt has; it would be faster here
    # to have a chain of if/elseif based on the value of the min below
    for l = 1:min(abs(x.x), cld(sizeof(T), sizeof(Limb)))
        u += (@inbounds x.v[l] % T) << ((sizeof(Limb)<<3)*(l-1))
    end
    flipsign(u, x.x)
end

rem(x::Integer, ::Type{XInt}) = XInt(x)

isodd(x::XInt) = is_short(x) ? isodd(x.x) : isodd(@inbounds x.v[1])
iseven(x::XInt) = !isodd(x)

function (::Type{T})(x::XInt) where T<:Base.BitUnsigned
    if is_short(x)
        convert(T, x.x)
    elseif sizeof(T) < sizeof(Limb)
        convert(T, convert(Limb, x))
    else
        0 <= x.x <= cld(sizeof(T), sizeof(Limb)) || throw(InexactError(nameof(T), T, x))
        x % T
    end
end

ispos(x::XInt) = x.x > 0
ispos(x::XInt, y::XInt) = x.x > 0 && y.x > 0
isneg(x::XInt) = x.x < 0
signbit(x::XInt) = x.x < 0
sign(x::XInt) = _XInt(sign(x.x))
_sign(x::XInt) = sign(x.x)
iszero(x::XInt) = iszero(x.x)
isone(x::XInt) = is_short(x) ? isone(x.x) : isone(@inbounds x.v[1])
# TODO: remove is_short check when we guarantee that short integers are always stored in x.x

function (::Type{T})(x::XInt) where T<:Base.BitSigned
    is_short(x) && return T(x.x)
    n = abs(x.x)
    if sizeof(T) < sizeof(Limb)
        # @assert SLimb == typeof(Signed(one(Limb)))
        convert(T, convert(SLimb, x))
    else
        0 <= n <= cld(sizeof(T), sizeof(Limb)) || throw(InexactError(nameof(T), T, x))
        y = x % T
        ispos(x) ⊻ (y > 0) && throw(InexactError(nameof(T), T, x)) # catch overflow
        y
    end
end

function Float64(x::XInt, ::RoundingMode{:ToZero})
    if is_short(x)
        Float64(x.x, RoundToZero)
    else
        MPN.get_d(x, len(x), x.x)
    end
end

function (::Type{T})(x::XInt, ::RoundingMode{:ToZero}) where T<:Union{Float16,Float32}
    T(Float64(x, RoundToZero), RoundToZero)
end

function (::Type{T})(n::XInt, ::RoundingMode{:Down}) where T<:CdoubleMax
    x = T(n, RoundToZero)
    x > n ? prevfloat(x) : x
end

function (::Type{T})(n::XInt, ::RoundingMode{:Up}) where T<:CdoubleMax
    x = T(n, RoundToZero)
    x < n ? nextfloat(x) : x
end

for T = [Float16, Float32, Float64]
    @eval (::Type{$T})(x::XInt, r::RoundingMode{:Nearest}=RoundNearest) =
        is_short(x) ? $T(x.x, r) : @bigint () x $T(x, r)
end

AbstractFloat(x::XInt) = BigFloat(x)

BigFloat(x::XInt, r::MPFR.MPFRRoundingMode=MPFR.ROUNDING_MODE[];
         precision::Integer=MPFR.DEFAULT_PRECISION[]) =
             is_short(x) ? BigFloat(x.x, r; precision=precision) :
                           @bigint () x BigFloat(x, r; precision=precision)

float(::Type{XInt}) = BigFloat

# Binary ops
for (fJ, fC) in ((:mod, :fdiv_r), (:rem, :tdiv_r),
                 (:gcd, :gcd), (:lcm, :lcm))
    fC! = Symbol(fC, :!)
    @eval begin
        ($fJ)(x::XInt, y::XInt) =
            is_short(x, y) ? XInt(($fJ)(widen(x.x), widen(y.x))) :
                             @bigint z x y XInt(MPZ.$fC!(z, x, y))
    end
end

# TODO: 3+ args specializations for some ops, like in gmp.jl
# (already done for +)

# we put @inline, as it seems these methods don't by themselves,
# maybe because add! is already inline and too big
@inline +(x::XInt, y::XInt)       = add!(nothing, x, y)
@inline +(x::XInt, y::BitInteger) = add!(nothing, x, y)
@inline +(x::XInt, y::Integer)    = add!(_XInt(y, 1), x)
@inline +(x::Integer, y::XInt)    = y + x

@inline -(x::XInt, y::XInt)       = sub!(nothing, x, y)
@inline -(x::BitInteger, y::XInt) = sub!(nothing, x, y)
@inline -(x::Integer, y::XInt)    = sub!(_XInt(x, 1), y)
@inline -(x::XInt, y::Integer)    = neg!(y-x)

@inline *(x::XInt, y::XInt)     = mul!(nothing, x, y)
@inline *(x::XInt, c::ShortMax) = mul!(nothing, x, c)
@inline *(x::XInt, y::Integer)  = mul!(_XInt(y, safe_len(x)), x)
@inline *(x::Integer, y::XInt)  = y * x

@inline (&)(x::XInt, y::XInt)       = and!(nothing, x, y)
@inline (&)(x::XInt, y::BitInteger) = and!(nothing, x, y)
@inline (&)(x::XInt, y::Integer)    = and!(XInt(y), x)
@inline (&)(x::Integer, y::XInt)    = y & x

@inline (|)(x::XInt, y::XInt)       = ior!(nothing, x, y)
@inline (|)(x::XInt, y::BitInteger) = ior!(nothing, x, y)
@inline (|)(x::XInt, y::Integer)    = ior!(XInt(y), x)
@inline (|)(x::Integer, y::XInt)    = y | x

@inline xor(x::XInt, y::XInt)       = xor!(nothing, x, y)
@inline xor(x::XInt, y::BitInteger) = xor!(nothing, x, y)
@inline xor(x::XInt, y::Integer)    = xor!(XInt(y), x)
@inline xor(x::Integer, y::XInt)    = y ⊻ x

# it doesn't seem worthwhile to define 4+ args specific + methods,
# the vararg one is better
@inline +(x::XInt, y::XInt, z::XInt) = add!((x+y), z)
+(x::XInt, y::XInt, z::XInt, r::XInt...) = sum(tuple(x, y, z, r...))

function sum(arr::Union{AbstractArray{XInt},Tuple{Vararg{XInt}}})
    sp = _XInt(0) # > 0
    sn = _XInt(0) # < 0
    s0 = zero(Int128)

    for x = arr
        if is_short(x)
            s0 += short(x)
        elseif ispos(x)
            sp = addbig_positive!(sp, x)
        else
            sn = addbig_positive!(sn, x)
        end
    end
    if !iszero(sp)
        if !iszero(sn)
            s = addbig!(sp, sp, neg!(sn))
            add!(s, s0)
        else
            add!(sp, s0)
        end
    elseif !iszero(sn)
        add!(neg!(sn), s0)
    else
        XInt(s0)
    end
end

for (r, f!) in ((RoundToZero, :tdiv_q!),
               (RoundDown, :fdiv_q!),
               (RoundUp, :cdiv_q!))
    @eval div(x::XInt, y::XInt, ::typeof($r)) =
        is_short(x, y) ? XInt(div(widen(x.x), widen(y.x), $r)) :
                         @bigint z x y XInt(MPZ.$f!(z, x, y))
end

# For compat only. Remove in 2.0.
div(x::XInt, y::XInt) = div(x, y, RoundToZero)
fld(x::XInt, y::XInt) = div(x, y, RoundDown)
cld(x::XInt, y::XInt) = div(x, y, RoundUp)

divrem(x::XInt, y::XInt) =
    is_short(x, y) ? XInt.(divrem(x.x, y.x)) :
                     @bigint (u, v) x y XInt.(MPZ.tdiv_qr!(u, v, x, y))

/(x::XInt, y::XInt) = float(x)/float(y)

invmod(x::XInt, y::XInt) =
    is_short(x, y) ? XInt(invmod(widen(x.x), widen(y.x))) :
                     @bigint () x y XInt(invmod(x, y))

/(x::XInt, y::Union{ClongMax,CulongMax}) = float(x)/y
/(x::Union{ClongMax,CulongMax}, y::XInt) = x/float(y)

# unary ops
(-)(x::XInt) = neg!(x)
(~)(x::XInt) = com!(nothing, x)

<<(x::XInt, c::UInt) = lshift!(nothing, x, c)
>>(x::XInt, c::UInt) = rshift!(nothing, x, c)
>>>(x::XInt, c::UInt) = x >> c

trailing_zeros(x::XInt) =
    if is_short(x)
        iszero(x.x) && throw(DomainError(x)) # InexactError for BigInt
        trailing_zeros(x.x)
    else
        @bigint () x MPZ.scan1(x, 0)
    end

trailing_ones(x::XInt) =
    if is_short(x)
        x.x == -1 && throw(DomainError(x))
        trailing_ones(x.x)
    else
        @bigint () x MPZ.scan0(x, 0)
    end

count_ones(x::XInt) =
    if is_short(x)
        x.x < 0 && throw(DomainError(x))
        count_ones(x.x)
    else
        @bigint () x MPZ.popcount(x)
    end

cmp(x::XInt, y::XInt) = is_short(x, y) ? cmp(x.x, y.x) : @bigint () x y cmp(x, y)
cmp(x::XInt, y::Integer) = is_short(x) ? cmp(x.x, y) : @bigint () x cmp(x, y)
cmp(x::XInt, y::BigInt)  = is_short(x) ? cmp(x.x, y) : @bigint () x cmp(x, y) # disambiguation
cmp(y::Integer, x::XInt) = -cmp(x, y)
cmp(y::BigInt, x::XInt)  = -cmp(x, y) # disambiguation
cmp(x::XInt, y::CdoubleMax) = is_short(x) ? cmp(x.x, y) : @bigint () x cmp(x, y)
cmp(y::CdoubleMax, x::XInt) = -cmp(x, y)

function ==(x::XInt, y::XInt)
    if is_short(x, y)
        x.x == y.x
    else
        @bigint () x y x == y
    end
end

==(x::XInt, y::Integer) =
    if is_short(x)
        x.x == y
    else
        @bigint () x x == y
    end

==(x::Integer, y::XInt) = y == x

# disambiguation
==(x::XInt, y::BigInt) = invoke(==, Tuple{XInt, Integer}, x, y)
==(x::BigInt, y::XInt) = y == x

==(x::XInt, f::CdoubleMax) = isnan(f) ? false : cmp(x, f) == 0
==(f::CdoubleMax, x::XInt) = isnan(f) ? false : cmp(x, f) == 0

<=(x::XInt, y::XInt) = cmp(x,y) <= 0
<=(x::XInt, i::Integer) = cmp(x,i) <= 0
<=(x::XInt, i::BigInt)  = cmp(x,i) <= 0 # disambiguation
<=(i::Integer, x::XInt) = cmp(x,i) >= 0
<=(i::BigInt, x::XInt)  = cmp(x,i) >= 0 # disambiguation
<=(x::XInt, f::CdoubleMax) = isnan(f) ? false : cmp(x,f) <= 0
<=(f::CdoubleMax, x::XInt) = isnan(f) ? false : cmp(x,f) >= 0

<(x::XInt, y::XInt) = cmp(x,y) < 0
<(x::XInt, i::Integer) = cmp(x,i) < 0
<(x::XInt, i::BigInt)  = cmp(x,i) < 0 # disambiguation
<(i::Integer, x::XInt) = cmp(x,i) > 0
<(i::BigInt, x::XInt)  = cmp(x,i) > 0 # disambiguation

<(x::XInt, f::CdoubleMax) = isnan(f) ? false : cmp(x,f) < 0
<(f::CdoubleMax, x::XInt) = isnan(f) ? false : cmp(x,f) > 0

isqrt(x::XInt) = is_short(x) ? XInt(isqrt(x.x)) : @bigint z x XInt(MPZ.sqrt!(z, x))

^(x::XInt, y::Culong) = @bigint z x XInt(MPZ.pow_ui!(z, x, y))

function xint_pow(x::XInt, y::Integer)
    if y<0; throw(DomainError(y, "`y` cannot be negative.")); end
    @noinline throw1(y) =
        throw(OverflowError("exponent $y is too large and computation will overflow"))
    if x== 1; return x; end
    if x==-1; return isodd(y) ? x : -x; end
    if y>typemax(Culong)
       x==0 && return x

       #At this point, x is not 1, 0 or -1 and it is not possible to use
       #gmpz_pow_ui to compute the answer. Note that the magnitude of the
       #answer is:
       #- at least 2^(2^32-1) ≈ 10^(1.3e9) (if Culong === UInt32).
       #- at least 2^(2^64-1) ≈ 10^(5.5e18) (if Culong === UInt64).
       #
       #Assume that the answer will definitely overflow.

       throw1(y)
    end
    return x^convert(Culong, y)
end

^(x::XInt   , y::XInt   ) = xint_pow(x, y)
^(x::XInt   , y::BigInt ) = xint_pow(x, y)
^(x::BigInt , y::XInt   ) = Base.GMP.bigint_pow(x, y)
^(x::XInt   , y::Bool   ) = y ? x : one(x)
^(x::XInt   , y::Integer) = xint_pow(x, y)
^(x::Integer, y::XInt   ) = xint_pow(XInt(x), y)
^(x::Bool   , y::XInt   ) = Base.power_by_squaring(x, y)

powermod(x::XInt, p::XInt, m::XInt) =
    is_short(x, p) && is_short(m) ? XInt(powermod(x.x, p.x, m.x)) :
                                    @bigint () x p m XInt(powermod(x, p, m))

powermod(x::Integer, p::Integer, m::XInt) = powermod(XInt(x), XInt(p), m)

function gcdx(a::XInt, b::XInt)
    is_short(a, b) &&
        return XInt.(gcdx(widen(a.x), widen(b.x)))
        # widen above necessary at least for (a, b) == (typemin(Int), 0)
    if iszero(b) # shortcut this to ensure consistent results with gcdx(a,b)
        return a < 0 ? (-a,-one(XInt),b) : (a,one(XInt),b)
    end
    g, s, t = XInt.(@bigint (x, y, z) a b MPZ.gcdext!(x, y, z, a, b))
    if t == 0
        # work around a difference in some versions of GMP
        if a == b
            return g, t, s
        elseif abs(a)==abs(b)
            return g, t, -s
        end
    end
    g, s, t
end

flipsign(x::XInt, y::Integer) = signbit(y) ? -x : x
flipsign(x::XInt, y::XInt)  = signbit(y) ? -x : x

factorial(x::XInt) = signbit(x) ? throw(DomainError(x)) :
                                  @bigint z x XInt(MPZ.fac_ui!(z, x))

binomial(x::XInt, k::UInt) = @bigint z x XInt(MPZ.bin_ui!(z, x, k))
binomial(x::XInt, k::Integer) = k < 0 ? XInt(0) : binomial(x, UInt(k))

function digits!(a::AbstractVector{T}, n::XInt; base::Integer = 10) where {T<:Integer}
    if 2 ≤ base ≤ 62
        s = codeunits(string(n; base=base))
        i, j = firstindex(a)-1, length(s)+1
        lasti = min(lastindex(a), firstindex(a) + length(s)-1 - isneg(n))
        while i < lasti
            # base ≤ 36: 0-9, plus a-z for 10-35
            # base > 36: 0-9, plus A-Z for 10-35 and a-z for 36..61
            x = s[j -= 1]
            a[i += 1] = base ≤ 36 ? (x>0x39 ? x-0x57 : x-0x30) : (x>0x39 ? (x>0x60 ? x-0x3d : x-0x37) : x-0x30)
        end
        lasti = lastindex(a)
        while i < lasti
            a[i+=1] = zero(T)
        end
        isneg(n) ? map!(-,a,a) : a
    else
        invoke(digits!, Tuple{typeof(a), Integer}, a, n; base=base) # slow generic fallback
    end
end

ndigits0zpb(x::XInt, b::Integer) = is_short(x) ? ndigits0zpb(x.x, b) :
                                                 @bigint () x ndigits0zpb(x, b)

function _prevpow2(x::XInt)
    if is_short(x)
        XInt(_prevpow2(x.x))
    else
        len = abs(x.x)
        @assert length(x.v) >= len
        high = @inbounds x.v[len]
        @assert !iszero(high) # like for BigInt/mpz_t
        shift = BITS_PER_LIMB - leading_zeros(high) - 1
        v = fill(zero(Limb), len)
        @inbounds v[len] = one(Limb) << shift
        _XInt(x.x, v)
    end
end

function _nextpow2(x::XInt)
    if is_short(x)
        y = _nextpow2(x.x)
        if sign(x) != sign(y)
            @assert y == typemin(Int)
            _XInt(1, Limb[typemin(Int) % Limb])
        else
            XInt(y)
        end
    else
        len = abs(x.x)
        @assert length(x.v) >= len
        popcount = GC.@preserve x MPZ.mpn_popcount(pointer(x.v), len)
        if popcount <= 1
            @assert popcount == 1
            return x
        end
        high = @inbounds x.v[len]
        @assert !iszero(high) # like for BigInt/mpz_t
        shift = (BITS_PER_LIMB - leading_zeros(high)) & (BITS_PER_LIMB-1)
        newlen = len + iszero(shift)
        v = fill(zero(Limb), newlen)
        @inbounds v[newlen] = one(Limb) << shift
        _XInt(flipsign(newlen, x.x), v)
    end
end

Base.checked_abs(x::XInt) = abs(x)
Base.checked_neg(x::XInt) = -x
Base.checked_add(a::XInt, b::XInt) = a + b
Base.checked_sub(a::XInt, b::XInt) = a - b
Base.checked_mul(a::XInt, b::XInt) = a * b
Base.checked_div(a::XInt, b::XInt) = div(a, b)
Base.checked_rem(a::XInt, b::XInt) = rem(a, b)
Base.checked_fld(a::XInt, b::XInt) = fld(a, b)
Base.checked_mod(a::XInt, b::XInt) = mod(a, b)
Base.checked_cld(a::XInt, b::XInt) = cld(a, b)
Base.add_with_overflow(a::XInt, b::XInt) = a + b, false
Base.sub_with_overflow(a::XInt, b::XInt) = a - b, false
Base.mul_with_overflow(a::XInt, b::XInt) = a * b, false

_copy(x::XInt) = _XInt(x.x, _copy(x.v))
_copy(v::Vector) = copy(v)
_copy(::Nothing) = nothing

Base.deepcopy_internal(x::XInt, d::IdDict) = get!(() -> _copy(x), d, x)

if Limb === UInt
    using Base: hash_uint

    function hash_integer(x::XInt, h::UInt)
        is_short(x) && return hash_integer(x.x, h)
        sz = x.x
        xv = x.v
        b = @inbounds xv[1]
        h ⊻= hash_uint(ifelse(sz < 0, -b, b) ⊻ h)
        for k = 2:abs(sz)
            h ⊻= hash_uint(@inbounds xv[k] ⊻ h)
        end
        h
    end

    function hash(x::XInt, h::UInt)
        GC.@preserve x begin
            is_short(x) && return hash(x.x, h)
            sz = x.x
            xv = x.v
            if sz == 1
                return hash(@inbounds(xv[1]), h)
            elseif sz == -1
                limb = @inbounds xv[1]
                limb <= typemin(Int) % UInt && return hash(-(limb % Int), h)
            end
            pow = trailing_zeros(x)
            nd = Base.ndigits0z(x, 2)
            idx = _divLimb(pow) + 1
            shift = _modLimb(pow) % UInt
            upshift = BITS_PER_LIMB - shift
            asz = abs(sz)
            if shift == 0
                limb = @inbounds xv[idx]
            else
                limb1 = @inbounds xv[idx]
                limb2 = idx < asz ? @inbounds(xv[idx+1]) : UInt(0)
                limb = limb2 << upshift | limb1 >> shift
            end
            if nd <= 1024 && nd - pow <= 53
                # TODO: add test for this branch
                return hash(ldexp(flipsign(Float64(limb), sz), pow), h)
            end
            h = hash_integer(1, h)
            h = hash_integer(pow, h)
            h ⊻= hash_uint(flipsign(limb, sz) ⊻ h)
            for idx = idx+1:asz
                if shift == 0
                    limb = @inbounds xv[idx]
                else
                    limb1 = limb2
                    if idx == asz
                        limb = limb1 >> shift
                        limb == 0 && break # don't hash leading zeros
                    else
                        limb2 = @inbounds xv[idx+1]
                        limb = limb2 << upshift | limb1 >> shift
                    end
                end
                h ⊻= hash_uint(limb ⊻ h)
            end
            return h
        end
    end
end

## Random

import Random: Sampler, rand, rand!
using Random: Random, Repetition, AbstractRNG

abstract type SamplerXInt <: Sampler{XInt} end

struct SamplerXIntBig{SP<:Sampler{Limb}} <: SamplerXInt
    a::XInt           # first
    m::XInt           # range length - 1 (generate z ∈ [0, m])
    nlimbsmax::SLimb  # max number of limbs for z+a
    highsp::SP        # sampler for the highest limb
end

safe_len(x::XInt) = is_short(x) ? SLimb(!iszero(x)) : abs(x.x)

@noinline function SamplerXIntBig(::Type{RNG}, r1::XInt, r2::XInt) where {RNG<:AbstractRNG}
    m = r2-r1
    isneg(m) && throw(ArgumentError("range must be non-empty"))
    if is_short(m)
        mvl = m.x % Limb
        nlimbsmax = 0
    else
        nlimbs, mv = lenvec(m)
        if nlimbs == 1
            mvl = @inbounds mv[1]
            nlimbsmax = 0
        else
            nlimbsmax = max(nlimbs, safe_len(r2), safe_len(r1))
            nlimbsmax += ispos(r1) # add! will request one more limb just in case
            mvl = @inbounds mv[nlimbs]
        end
    end
    highsp = Sampler(RNG, zero(Limb):mvl)
    SamplerXIntBig(r1, m, nlimbsmax, highsp)
end

struct SamplerXIntShort{SP<:Sampler{SLimb}} <: SamplerXInt
    sp::SP
end

@inline function Sampler(::Type{RNG}, r::AbstractUnitRange{XInt}, N::Repetition
                 ) where {RNG<:AbstractRNG}
    r1, r2 = first(r), last(r)
    if is_short(r1, r2)
        SamplerXIntShort(Sampler(RNG, r1.x:r2.x, N))
    else
        SamplerXIntBig(RNG, r1, r2)
    end
end

rand(rng::AbstractRNG, sp::SamplerXInt) = rand!(rng, _XInt(0), sp)

function rand!(rng::AbstractRNG, x::XInt, sp::SamplerXIntBig)
    iszero(sp.nlimbsmax) &&
        return add!(x, sp.a, _XInt(rand(rng, sp.highsp)))
    nlimbs, mv = lenvec(sp.m)
    limbs = vec!(x, sp.nlimbsmax)
    mvl = @inbounds mv[nlimbs]
    while true
        @inbounds rand!(rng, @view(limbs[1:nlimbs-1]))
        @inbounds h = limbs[nlimbs] = rand(rng, sp.highsp)
        h < mvl && break
        MPN.cmp(limbs, sp.m, nlimbs) <= 0 && break
    end
    len = normalize(limbs, nlimbs)
    add!(_XInt(len, limbs, false), _XInt(len, limbs, true), sp.a)
end

rand!(rng::AbstractRNG, ::XInt, sp::SamplerXIntShort) = _XInt(rand(rng, sp.sp))

end # module


struct POMCPOWTree{B,A,O,RB}
    # action nodes
    n::Vector{Int}
    v::Vector{Float64}
    generated::Vector{Vector{Pair{O,Int}}}
    a_child_lookup::Dict{Tuple{Int,O}, Int} # may not be maintained based on solver params
    a_labels::Vector{A}
    n_a_children::Vector{Int}

    # observation nodes
    sr_beliefs::Vector{B} # first element is #undef
    total_n::Vector{Int}
    tried::Vector{Vector{Int}}
    o_child_lookup::Dict{Tuple{Int,A}, Int} # may not be maintained based on solver params
    o_labels::Vector{O}

    # root
    root_belief::RB

    function POMCPOWTree{B,A,O,RB}(root_belief, sz::Int=1000) where{B,A,O,RB}
        sz = min(sz, 100_000)
        return new(
            sizehint!(Int[], sz),
            sizehint!(Int[], sz),
            sizehint!(Vector{Pair{O,Int}}[], sz),
            Dict{Tuple{Int,O}, Int}(),
            sizehint!(A[], sz),
            sizehint!(Int[], sz),

            sizehint!(Array{B}(undef, 1), sz),
            sizehint!(Int[0], sz),
            sizehint!(Vector{Int}[Int[]], sz),
            Dict{Tuple{Int,A}, Int}(),
            sizehint!(Array{O}(undef, 1), sz),

            root_belief
        )
    end
end

@inline function push_anode!(tree::POMCPOWTree{B,A,O}, h::Int, a::A, n::Int=0, v::Float64=0.0, update_lookup=true) where {B,A,O}
    anode = length(tree.n) + 1
    push!(tree.n, n)
    push!(tree.v, v)
    push!(tree.generated, Pair{O,Int}[])
    push!(tree.a_labels, a)
    push!(tree.n_a_children, 0)
    if update_lookup
        tree.o_child_lookup[(h, a)] = anode
    end
    push!(tree.tried[h], anode)
    tree.total_n[h] += n
    return anode
end

struct POWTreeObsNode{B,A,O,RB} <: BeliefNode
    tree::POMCPOWTree{B,A,O,RB}
    node::Int
end

isroot(h::POWTreeObsNode) = h.node==1
@inline function belief(h::POWTreeObsNode)
    if isroot(h)
        return h.tree.root_belief
    else
        return StateBelief(h.tree.sr_beliefs[h.node])
    end
end
function sr_belief(h::POWTreeObsNode)
    if isroot(h)
        error("Tried to access the sr_belief for the root node in a POMCPOW tree")
    else
        return h.tree.sr_beliefs[h.node]
    end
end
n_children(h::POWTreeObsNode) = length(h.tree.tried[h.node])


module DownloadArgo

using Printf, Dates, YAML, NetCDF, NCDatasets, CSV, DataFrames, Dierckx

"""
    GDAC_FTP(b::String,y::Int,m::Int)

Download Argo data files for one regional domain (b), year (y), and
month (m) from the `GDAC` FTP server (`ftp://ftp.ifremer.fr/ifremer/argo`
or, equivalently, `ftp://usgodae.org/pub/outgoing/argo`).

```
b="atlantic"; yy=2009:2009; mm=8:12;
for y=yy, m=mm;
    println("\$b/\$y/\$m"); DownloadArgo.GDAC_FTP(b,y,m)
end
```
"""
function GDAC_FTP(b::String,y::Int,m::Int)
    yy = @sprintf "%04d" y
    mm = @sprintf "%02d" m
    c=`wget --quiet -r ftp://ftp.ifremer.fr/ifremer/argo/geo/"$b"_ocean/$yy/$mm`
    run(c)
end

"""
    mitprof_interp_setup(fil::String)

Get parameters etc from yaml file (`fil`).
"""
function mitprof_interp_setup(fil::String)

    meta=YAML.load(open(fil))

    #1. file list

    d=meta["dirIn"]
    b=meta["subset"]["basin"]
    y=meta["subset"]["year"]

    list0=Array{Array,1}(undef,12)
    for m=1:12
        sd="$b"*"_ocean/$y/"*Printf.@sprintf("%02d/",m)
        tmp=readdir(d*sd)
        list0[m]=[sd*tmp[i] for i=1:length(tmp)]
    end

    nf=sum(length.(list0))
    list1=Array{String,1}(undef,nf)
    f=0
    for m=1:12
        for ff=1:length(list0[m])
            f+=1
            list1[f]=list0[m][ff]
        end
    end

    meta["fileInList"]=list1;

    #2. coordinate

    z_std=meta["depthLevels"]
    if length(z_std)>1
        tmp1=(z_std[2:end]+z_std[1:end-1])/2
        z_top=[z_std[1]-(z_std[2]-z_std[1])/2;tmp1]
        z_bot=[tmp1;z_std[end]+(z_std[end]-z_std[end-1])/2]
    else
        z_top=0.9*z_std
        z_bot=1.1*z_std
    end

    meta["z_std"]=z_std
    meta["z_top"]=z_top
    meta["z_bot"]=z_bot

    #3. various other specs

    meta["inclZ"] = false
    meta["inclT"] = true
    meta["inclS"] = true
    meta["inclU"] = false
    meta["inclV"] = false
    meta["inclPTR"] = false
    meta["inclSSH"] = false
    meta["TPOTfromTINSITU"] = true

    meta["doInterp"] = true
    meta["addGrid"] = true
    meta["outputMore"] = false
    meta["method"] = "interp"
    meta["fillval"] = -9999.0
    meta["buffer_size"] = 10000

    meta["fileOut"]=meta["name"]*"_"*meta["subset"]["basin"]*"_"*string(meta["subset"]["year"])*".nc"
    meta["var_out"]=meta["variables"]

    return meta
end

"""
    GetOneProfile(m)

Get one profile from netcdf file.
"""
function GetOneProfile(ds,m)

    #
    t=ds["JULD"][m]
    ymd=Dates.year(t)*1e4+Dates.month(t)*1e2+Dates.day(t)
    hms=Dates.hour(t)*1e4+Dates.minute(t)*1e2+Dates.second(t)

    lat=ds["LATITUDE"][m]
    lon=ds["LONGITUDE"][m]
    lon < 0.0 ? lon=lon+360.0 : nothing

    direction=ds["DIRECTION"][m]
    direc=0
    direction=='A' ? direc=1 : nothing
    direction=='D' ? direc=2 : nothing

    #
    pnum_txt=ds["PLATFORM_NUMBER"][:,m]
    #pnum_txt=pnum_txt[findall((!ismissing).(pnum_txt))]
    ii=findall(skipmissing(in.(pnum_txt,"0123456789")))
    ~isempty(ii) ? pnum_txt=String(vec(Char.(pnum_txt[ii]))) : pnum_txt="9999"
    pnum=parse(Int,pnum_txt)

    #
    p=ds["PRES_ADJUSTED"][:,m]
    p_QC=ds["PRES_ADJUSTED_QC"][:,m]
    if sum((!ismissing).(p))==0
        p=ds["PRES"][:,m]
        p_QC=ds["PRES_QC"][:,m]
    end

    #set qc to 5 if missing
    p_QC[findall(ismissing.(p))].='5'
    #avoid potential duplicates
    for n=1:length(p)-1
        if ~ismissing(p[n])
            tmp1=( (!ismissing).(p[n+1:end]) ).&( p[n+1:end].==p[n] )
            tmp1=findall(tmp1)
            p[n.+tmp1].=missing
            p_QC[n.+tmp1].='5'
        end
    end

    #position and date
    isBAD=0
    ~in(ds["POSITION_QC"][m],"1258") ? isBAD=1 : nothing
    ~in(ds["JULD_QC"][m],"1258") ? isBAD=1 : nothing

    #pressure
    tmp1=findall( (!in).(p_QC,"1258") )
    if (length(tmp1)<=5)&&(length(tmp1)>0)
        #omit these few bad points but keep the profile
        p[tmp1].=missing
    elseif length(tmp1)>5
        #flag the profile (will be omitted later)
        #but keep the bad points (for potential inspection)
        isBAD=1
    end

    #temperature
    t=ds["TEMP_ADJUSTED"][:,m]
    t_QC=ds["TEMP_ADJUSTED_QC"][:,m]
    t_ERR=ds["TEMP_ADJUSTED_ERROR"][:,m]
    t_ERR[findall( (ismissing).(t_ERR) )].=0.0

    if sum((!ismissing).(t))==0
        t=ds["TEMP"][:,m]
        t_QC=ds["TEMP_QC"][:,m]
    end

    #salinity
    if haskey(ds,"PSAL")
        s=ds["PSAL_ADJUSTED"][:,m]
        s_QC=ds["PSAL_ADJUSTED_QC"][:,m]
        s_ERR=ds["PSAL_ADJUSTED_ERROR"][:,m]
        s_ERR[findall( (ismissing).(t_ERR) )].=0.0
        if sum((!ismissing).(t))==0
            s=ds["PSAL"][:,m]
            s_QC=ds["PSAL_QC"][:,m]
        end
    else
        s=fill(missing,size(t))
        s_QC=Char.(32*ones(size(t_QC)))
    end

    if ismissing(t[1]) #this file does not contain temperature data...
        t=fill(missing,size(t))
        t_ERR=fill(0.0,size(t))
    else #apply QC
        tmp1=findall(skipmissing((!in).(t_QC,"1258")))
        t[tmp1].=missing
    end

    if ismissing(s[1]) #this file does not contain salinity data...
        s=fill(missing,size(s))
        s_ERR=fill(0.0,size(s))
    else #apply QC
        tmp1=findall(skipmissing((!in).(s_QC,"1258")))
        s[tmp1].=missing
    end

    prof=Dict()
    prof["pnum_txt"]=pnum_txt
    prof["ymd"]=convert(Union{Int,Missing},ymd)
    prof["hms"]=convert(Union{Int,Missing},hms)
    prof["lat"]=convert(Union{Float64,Missing},lat)
    prof["lon"]=convert(Union{Float64,Missing},lon)
    prof["direc"]=convert(Union{Int,Missing},direc)
    prof["T"]=convert(Array{Union{Float64,Missing}},t)
    prof["S"]=convert(Array{Union{Float64,Missing}},s)
    prof["p"]=convert(Array{Union{Float64,Missing}},p)
    prof["T_ERR"]=convert(Array{Union{Float64,Missing}},t_ERR)
    prof["S_ERR"]=convert(Array{Union{Float64,Missing}},s_ERR)
    prof["isBAD"]=isBAD
    prof["DATA_MODE"]=ds["DATA_MODE"][m]

    return prof
end

"""
    prof_PtoZ!(prof,meta)

Convert prof["p"] to prof["depth"]
"""
function prof_PtoZ!(prof,meta)
    l=prof["lat"]
    v=meta["var_out"][1]
    prof[v]=similar(prof["p"],Union{Missing,Float64})
    prof[v].=missing
    k=findall((!ismissing).(prof["p"]))
    prof[v][k]=[DownloadArgo.sw_dpth(Float64(prof["p"][kk]),Float64(l)) for kk in k]
end

"""
    prof_TtoΘ!(prof,meta)

Convert prof["T"] to potential temperature
"""
function prof_TtoΘ!(prof,meta)
    T=prof[meta["var_out"][2]]
    P=0.981*1.027*prof[meta["var_out"][1]]
    S=35.0*ones(size(T))
    k=findall( (!ismissing).(T) )
    T[k]=[DownloadArgo.sw_ptmp(Float64(S[kk]),Float64(T[kk]),Float64(P[kk])) for kk in k]
end

"""
    prof_interp!(prof,meta)

Interpolate from prof["depth"] to meta["z_std"]
"""
function prof_interp!(prof,meta)
    for ii=2:length(meta["var_out"])
        v=meta["var_out"][ii]
        v_e=v*"_ERR"

        z_std=meta["z_std"]
        t_std=similar(z_std,Union{Missing,Float64})
        e_std=similar(z_std,Union{Missing,Float64})

        z=prof["depth"]
        t=prof[v]
        do_e=haskey(prof,v_e)
        do_e ? e=prof[v_e] : e=[]

        kk=findall((!ismissing).(z.*t))
        if (meta["doInterp"])&&(length(kk)>1)
            z_in=z[kk]; t_in=t[kk]
            do_e ? e_in=e[kk] : nothing
            k=sort(1:length(kk),by= i -> z_in[i])
            z_in=z_in[k]; t_in=t_in[k]
            do_e ? e_in=e_in[k] : nothing
            #omit values outside observed range:
            D0=minimum(skipmissing(z_in))
            D1=maximum(skipmissing(z_in))
            msk1=findall( (z_std.<D0).|(z_std.>D1) )
            #avoid duplicates:
            msk2=findall( ([false;(z_in[1:end-1]-z_in[2:end]).==0.0]).==true )
            if length(kk)>5
                spl = Spline1D(z_in, t_in)
                t_std[:] = spl(z_std)
                t_std[msk1].=missing
                t_std[msk2].=missing
                if do_e
                    spl = Spline1D(z_in, e_in)
                    e_std[:] = spl(z_std)
                    e_std[msk1].=missing
                    e_std[msk2].=missing
                end
            else
                t_std = []
                e_std = []
            end
            prof[v]=t_std
        end
    end
end

"""
    sw_dpth(P,LAT)

Calculate depth in meters from pressure (P; in decibars) and
latitude (LAT; in °N)

```
d = sw_dpth(100.0,20.0)
```
"""
function sw_dpth(P,LAT)
    # Original author:  Phil Morgan 92-04-06  (morgan@ml.csiro.au)
    # Reference: Unesco 1983. Algorithms for computation of fundamental properties of
    # seawater, 1983. _Unesco Tech. Pap. in Mar. Sci._, No. 44, 53 pp. Eqn 25, p26

    c1 = 9.72659
    c2 = -2.2512E-5
    c3 = 2.279E-10
    c4 = -1.82E-15
    gam_dash = 2.184e-6

    LAT = abs.(LAT)
    X   = sin.(deg2rad.(LAT))
    X   = X.*X
    bot_line = 9.780318*(1.0+(5.2788E-3+2.36E-5*X).*X) + gam_dash*0.5*P
    top_line = (((c4*P+c3).*P+c2).*P+c1).*P
    DEPTHM   = top_line./bot_line

    return DEPTHM
end

"""
    sw_ptmp(S,T,P,PR)

Calculate potential temperature as per UNESCO 1983 report from salinity (S;
in psu), in situ temperature (T; in °C), and pressure (P; in decibar)
relative to PR (in decibar; 0 by default).

```
ptmp = sw_ptmp(S,T,P,PR=missing)
```
"""
function sw_ptmp(S,T,P,PR=0.0)
# Original author:  Phil Morgan
# Reference: Fofonoff, P. and Millard, R.C. Jr
#    Unesco 1983. Algorithms for computation of fundamental properties of
#    seawater. _Unesco Tech. Pap. in Mar. Sci._, No. 44. Eqn.(31) p.39
#    Bryden, H. 1973.
#    "New Polynomials for thermal expansion, adiabatic temperature gradient
#    and potential temperature of sea water." DEEP-SEA RES., 1973, Vol 20

# theta1
del_P  = PR - P
del_th = del_P.*sw_adtg(S,T,P);
th     = T + 0.5*del_th;
q      = del_th;

# theta2
del_th = del_P.*sw_adtg(S,th,P+0.5*del_P);
th     = th + (1 - 1/sqrt(2))*(del_th - q);
q      = (2-sqrt(2))*del_th + (-2+3/sqrt(2))*q;

# theta3
del_th = del_P.*sw_adtg(S,th,P+0.5*del_P);
th     = th + (1 + 1/sqrt(2))*(del_th - q);
q      = (2 + sqrt(2))*del_th + (-2-3/sqrt(2))*q;

# theta4
del_th = del_P.*sw_adtg(S,th,P+del_P);
PT     = th + (del_th - 2*q)/6;

return PT
end

"""
    sw_adtg(S,T,P)

Calculate adiabatic temperature gradient as per UNESCO 1983 routines from salinity
(S; in psu), in situ temperature (T; in °C), and pressure (P; in decibar)
```
adtg = sw_adtg(S,T,P)
```
"""
function sw_adtg(S,T,P)
    # Original author:  Phil Morgan
    # Reference: Fofonoff, P. and Millard, R.C. Jr
    #    Unesco 1983. Algorithms for computation of fundamental properties of
    #    seawater. _Unesco Tech. Pap. in Mar. Sci._, No. 44. Eqn.(31) p.39
    #    Bryden, H. 1973.
    #    "New Polynomials for thermal expansion, adiabatic temperature gradient
    #    and potential temperature of sea water." DEEP-SEA RES., 1973, Vol 20

    a0 =  3.5803E-5;
    a1 =  8.5258E-6;
    a2 = -6.836E-8;
    a3 =  6.6228E-10;

    b0 =  1.8932E-6;
    b1 = -4.2393E-8;

    c0 =  1.8741E-8;
    c1 = -6.7795E-10;
    c2 =  8.733E-12;
    c3 = -5.4481E-14;

    d0 = -1.1351E-10;
    d1 =  2.7759E-12;

    e0 = -4.6206E-13;
    e1 =  1.8676E-14;
    e2 = -2.1687E-16;

    ADTG =      a0 + (a1 + (a2 + a3.*T).*T).*T + (b0 + b1.*T).*(S-35.0) +
    ( (c0 + (c1 + (c2 + c3.*T).*T).*T) + (d0 + d1.*T).*(S-35.0) ).*P +
    (  e0 + (e1 + e2.*T).*T ).*P.*P

    return ADTG
end

end


# DiffEqSolution to Array
function diffeqsol_to_array(x)
    xarr = cpu(x)
    return reshape(xarr, size(xarr)[1:2])
end

function diffeqsol_to_cuarray(x)
    xarr = gpu(x)
    return reshape(xarr, size(xarr)[1:2])
end

function diffeqsol_to_trackedarray(x)
    xarr = x.u[end]
    return reshape(xarr, size(xarr)[1:2])
end

function diffeqsol_to_3dtrackedarray(x)
    return hcat(map(_x -> reshape(_x, size(_x, 1), 1, size(_x, 2)), x.u)...)
end

_convert_tspan(tspan, p) = eltype(p).(tspan)

_convert_tspan(tspan, p::TrackedArray) = Tracker.collect(eltype(p).(tspan))

norm_batched(x::AbstractArray) = sqrt.(sum(x .^ 2, dims = 1))

# Running Average Meter
mutable struct AverageMeter{T}
    last_value::T
    sum::T
    count::Int

    AverageMeter(T = Float32) = new{T}(T(0), T(0), 0)
end

function reset!(am::AverageMeter{T}) where {T}
    am.last_value = T(0)
    am.sum = T(0)
    am.count = 0
    return am
end

function update!(am::AverageMeter{T}, val::T) where {T}
    am.last_value = val
    am.sum += val
    am.count += 1
    return am.sum / am.count
end

(am::AverageMeter)() = am.sum / am.count

# Pretty Logger
function table_logger(header::Vector{String}, record::Vector{String} = [])
    n = length(header) + length(record)
    ind_lens = vcat(length.(header), length.(record))
    span = sum(ind_lens .+ 3) + 1
    println("="^span)
    for h in vcat(header, record)
        print("| $h ")
    end
    println("|")
    println("="^span)

    avg_meters = Dict{String,AverageMeter}(rec => AverageMeter() for rec in record)

    patterns = ["%$l.4f" for l in ind_lens]
    fmtrfuncs = generate_formatter.(patterns)
    function internal_logger(last::Bool, records::Dict, args::Vararg)
        if length(records) > 0
            for (rec, val) in records
                update!(avg_meters[rec], val)
            end
            return
        end
        if last
            println("="^span)
            return
        end
        for h in [
            fmtrfunc(arg) for (fmtrfunc, arg) in
            zip(fmtrfuncs, vcat([args...], [avg_meters[rec]() for rec in record]))
        ]
            print("| $h ")
        end
        println("|")
    end
end

# MvNormal for FFJORD
struct BatchedMultiVariateNormal{M,N,O,P,Q}
    μ::M
    cov::N
    inv_cov::N
    det_cov::O
    lt_decom::P
    k::Q
end

Flux.@functor BatchedMultiVariateNormal

function BatchedMultiVariateNormal(μ::AbstractVector, cov::AbstractMatrix)
    inv_cov = inv(cov)
    det_cov = det(cov)
    lt_decom = Array(cholesky(cov))
    return BatchedMultiVariateNormal(
        reshape(μ, :, 1),
        cov,
        inv_cov,
        det_cov,
        lt_decom,
        length(μ),
    )
end

function (mvnorm::BatchedMultiVariateNormal)(x::AbstractMatrix{T}) where {T}
    denom = sqrt((T(2π)^mvnorm.k) * T(mvnorm.det_cov))
    diff = x .- mvnorm.μ
    return exp.(diag(-permutedims(diff, (2, 1)) * (mvnorm.inv_cov * diff) / 2)) / denom
end

function sample(mvnorm::BatchedMultiVariateNormal{CuArray{T,2}}, nsamples::Int) where {T}
    samples = CUDA.randn(T, mvnorm.k, nsamples)::CuArray{T,2}
    return mvnorm.μ .+ mvnorm.lt_decom * samples
end

function sample(
    mvnorm::BatchedMultiVariateNormal{TrackedArray{T,2,CuArray{T,2}}},
    nsamples::Int,
) where {T}
    samples = CUDA.randn(T, mvnorm.k, nsamples)::CuArray{T,2}
    return mvnorm.μ .+ mvnorm.lt_decom * samples
end

function sample(mvnorm::BatchedMultiVariateNormal{Array{T,2}}, nsamples::Int) where {T}
    samples = randn(T, mvnorm.k, nsamples)
    return mvnorm.μ .+ mvnorm.lt_decom * samples
end

function sample(
    mvnorm::BatchedMultiVariateNormal{TrackedArray{T,2,Array{T,2}}},
    nsamples::Int,
) where {T}
    samples = randn(T, mvnorm.k, nsamples)
    return mvnorm.μ .+ mvnorm.lt_decom * samples
end

# Gradient Updates
function update_parameters!(ps, gs, opt)
    for (p, g) in zip(ps, gs)
        length(p) == 0 && continue
        # any(isnan.(g)) && error("NaN Gradient Detected. Terminating...")
        Flux.Optimise.update!(opt, data(p), data(g))
        p.tracker.f = Tracker.Call()
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    SourcePath(sources, batchsize=10^3)

Traverse an object from `sources` and outwards.
Optionally pass a `batchsize` for KD-tree evaluations.
"""
struct SourcePath <: Path
  sources::Vector{Int}
  batchsize::Int
end

SourcePath(sources) = SourcePath(sources, 10^3)

function traverse(object, path::SourcePath)
  sources   = path.sources
  batchsize = path.batchsize
  @assert allunique(sources) "non-unique sources"
  @assert all(1 .≤ sources .≤ nelements(object)) "sources must be valid locations"
  @assert length(sources) ≤ nelements(object) "more sources than points in object"

  # fit search tree
  kdtree = KDTree(coordinates.([centroid(object, s) for s in sources]))

  # other locations that are not sources
  others = setdiff(1:nelements(object), sources)

  # process points in batches
  batches = Iterators.partition(others, batchsize)

  # compute distances to sources
  dists = []
  for batch in batches
    coords = coordinates.([centroid(object, b) for b in batch])
    _, ds = knn(kdtree, coords, length(sources), true)
    append!(dists, ds)
  end

  [sources; view(others, sortperm(dists))]
end


using BinaryBuilder

name = "SLATE"
version = v"2022.05.00"

# Collection of sources required to build PETSc. Avoid using the git repository, it will
# require building SOWING which fails in all non-linux platforms.
sources = [
    GitSource("https://bitbucket.org/icl/slate.git", "5fb57877effa06d2ef090402a39341ebeb44f180")
]

# Bash recipe for building across all platforms

# Needs to add -Dcapi eventually once it's added to the cmake build system. Note yet available under CMAKAE toolchain.
script = raw"""
cd slate
git submodule update --init
mkdir build && cd build
cmake \
  -DCMAKE_INSTALL_PREFIX=${prefix} \
  -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
  -DCMAKE_BUILD_TYPE="Release" \
  -Dblas=openblas \
  -Dbuild_tests=no \
  -DMPI_RUN_RESULT_CXX_libver_mpi_normal="0" \
  -DMPI_RUN_RESULT_CXX_libver_mpi_normal__TRYRUN_OUTPUT="" \
  -Drun_result="0" \
  -Drun_result__TRYRUN_OUTPUT="ok" \
  -Dc_api=yes \
  ..
make -j${nproc}
make install
"""

# We attempt to build for all defined platforms
platforms = expand_gfortran_versions(expand_cxxstring_abis(supported_platforms(; exclude=!Sys.islinux)))
platforms = filter(p -> libgfortran_version(p) ≠ v"3", platforms)
products = [
    LibraryProduct("libslate", :libslate),
    LibraryProduct("libslate_lapack_api", :libslate_lapack_api)
    # LibraryProduct("libslate_scalapack_api, :libslate_scalapack_api) ** Not yet available under CMAKE toolchain.
]

dependencies = [
    Dependency("CompilerSupportLibraries_jll"),
    Dependency("OpenBLAS32_jll"),
    Dependency("MPICH_jll")
]

# Build the tarballs.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6", preferred_gcc_version=v"7")


sabb(x::String, val::Integer, happ::Integer)::String =
    +(count( c -> c ∈ "sabbatical",x),val,happ) > 22 ? "Sabbatical! Boom!" : "Back to your desk, boy."


__precompile__()
##===----------------------------------------------------------------------===##
##                                   ACCELERATE.JL                            ##
##                                                                            ##
##===----------------------------------------------------------------------===##
##                                                                            ##
##  This file checks for the available frameworks and imports only the        ##
##  frameworks available; it also coordinates access to three frameworks      ##
##  through the unified Vectorize.f() syntax                                  ##
##                                                                            ##
##===----------------------------------------------------------------------===##
module Vectorize

export @vectorize

# Cross-version compatibility
if VERSION < v"0.5.0-dev" # Julia v0.4
    OS = OS_NAME
else
    OS = Sys.KERNEL
end

## FUNCTION DICTIONARY
functions = Dict()

## Function for adding to function dictionary
function addfunction(d::Dict, k, v)
    if k in keys(d)
        append!(d[k], [v])
    else
        d[k] = []
        append!(d[k], [v])
    end
end

## START OSX
## On OS X, check if Accelerate can be found - future proofing
OS == :Darwin ? (ACCELERATE = true) : (ACCELERATE = false)
if ACCELERATE == true
    if isfile("/System/Library/Frameworks/Accelerate.framework/Accelerate")
        eval(:(include("Accelerate.jl")))
    else
        warn("Unable to locate the Apple Accelerate framework")
    end
end
## END OSX

## Find Yeppp library
libyeppp_ = Libdl.find_library(["libyeppp"])
if libyeppp_ != ""
    @eval const global libyeppp = libyeppp_
else
    currdir = @__FILE__
    @static if is_windows()
        bindir = currdir[1:end-16]*"deps\\src\\yeppp\\binaries\\"
    else
        bindir = currdir[1:end-16]*"deps/src/yeppp/binaries/"
    end

    if OS == :Darwin
        @eval const global libyeppp = bindir*"macosx/x86_64/libyeppp.dylib"
    elseif OS == :Linux
        @eval const global libyeppp = bindir*"linux/x86_64/libyeppp.so"
    elseif OS == :NT
        @eval const global libyeppp = bindir*"windows\\amd64\\yeppp.dll"
    end
end

if isfile(libyeppp)
    include("Yeppp.jl") # include Yeppp if present
end

## Find VML
const global libvml = Libdl.find_library(["libmkl_rt"], ["/opt/intel/mkl/lib"])
if libvml != ""
    include("VML.jl")
end

# vectorize macro
macro vectorize(ex)
    len = length(ex.args)
    if len == 2
        f = :(Vectorize.$(ex.args[1]))
        arg = ex.args[2]
        return esc(:($f($arg)))
    elseif len == 3
        f = :(Vectorize.$(ex.args[1]))
        arg1 = ex.args[2]
        arg2 = ex.args[3]
        return esc(:($f($arg1, $arg2)))
    else
        f = :(Vectorize.$(ex.args[1]))
        arg1 = ex.args[2]
        arg2 = ex.args[3]
        arg3 = ex.args[4]
        return esc(:($f($arg1, $arg2, $arg3)))
    end
end

# Include optimized functions
include("Functions.jl")

end # module


"""
This function does nothing if the product is not of Capacity type.
"""
function update_capacity_revenues!(product::T,
                                scenario_name::String,
                                price_years::NamedTuple{(:start_year, :end_year),
                                    Tuple{Int64, Int64}},
                                capacity_revenues::Vector{Float64},
                                capacity_prices::Dict{String, AxisArrays.AxisArray{Float64, 1}},
                                max_cap::Float64) where T <: Product
    return capacity_revenues
end

"""
This function calculates expected capacity revenues for each project.
Returns capacity revenues.
"""
function update_capacity_revenues!(product::Capacity,
                                scenario_name::String,
                                price_years::NamedTuple{(:start_year, :end_year),
                                    Tuple{Int64, Int64}},
                                capacity_revenues::Vector{Float64},
                                capacity_prices::Dict{String, AxisArrays.AxisArray{Float64, 1}},
                                max_cap::Float64)
    capacity_revenues += capacity_prices[scenario_name][price_years[:start_year]:price_years[:end_year]] .*
                         get_accepted_perc(product)[scenario_name][price_years[:start_year]:price_years[:end_year]] *
                         max_cap *
                         get_derating(product)
    return capacity_revenues
end

"""
This function updates the expected profit array of the project
if the product is a capacity product.
Returns nothing.
"""
function update_profit!(project::P,
                        scenario_name::String,
                        product::Capacity,
                        operating_profit::AxisArrays.AxisArray{Float64, 2},
                        capacity_revenues::Vector{Float64},
                        capacity_forward_years::Int64,
                        rec_revenues::Vector{Float64},
                        price_years::NamedTuple{(:start_year, :end_year),
                                     Tuple{Int64, Int64}},
                        update_years::NamedTuple{(:start_year, :end_year),
                                      Tuple{Int64, Int64}},
                        iteration_year::Int64) where P <: Project{<: BuildPhase}

    capacity_revenue_start_year = max(1, capacity_forward_years - price_years[:start_year] + 1)

        for i in capacity_revenue_start_year:length(capacity_revenues)
                set_scenario_profit!(get_finance_data(project),
                            scenario_name,
                            get_name(product),
                            iteration_year,
                            update_years[:start_year] + i - 1,
                            capacity_revenues[i])
        end

    return
end


## MODEL TRAITS

# model trait names:
const MODEL_TRAITS = [:input_scitype, :output_scitype, :target_scitype,
                      :is_pure_julia, :package_name, :package_license,
                      :load_path, :package_uuid,
                      :package_url, :is_wrapper, :supports_weights,
                      :supports_online, :docstring,
                      :name, :is_supervised, :prediction_type,
                      :implemented_methods, :hyperparameters,
                      :hyperparameter_types]
const SUPERVISED_TRAITS = filter(MODEL_TRAITS) do trait
    !(trait in [:output_scitype,])
end
const UNSUPERVISED_TRAITS = filter(MODEL_TRAITS) do trait
    !(trait in [:target_scitype, :prediction_type, :supports_weights])
end

# fallback trait declarations:
input_scitype(::Type) = Unknown
output_scitype(::Type) = Unknown
target_scitype(::Type) = Unknown  # used for measures too
is_pure_julia(::Type) = missing
package_name(::Type) = "unknown"
package_license(::Type) = "unknown"
load_path(::Type) = "unknown"
package_uuid(::Type) = "unknown"
package_url(::Type) = "unknown"
is_wrapper(::Type) = false
supports_online(::Type) = false
supports_weights(::Type) = false  # used for measures too
docstring(M::Type) = string(M)
docstring(M::Type{<:MLJType}) = name(M)
docstring(M::Type{<:Model}) =
    "$(name(M)) from $(package_name(M)).jl.\n"*
    "[Documentation]($(package_url(M)))."

# "derived" traits:
name(M::Type) = string(M)
name(M::Type{<:MLJType}) = split(string(coretype(M)), '.')[end] |> String
is_supervised(::Type{<:Type}) = false
is_supervised(::Type{<:Supervised}) = true
prediction_type(::Type) = :unknown # used for measures too
prediction_type(::Type{<:Deterministic}) = :deterministic
prediction_type(::Type{<:Probabilistic}) = :probabilistic
prediction_type(::Type{<:Interval}) = :interval
implemented_methods(M::Type{<:MLJType}) = map(f->f.name, methodswith(M))
hyperparameters(M::Type) = collect(fieldnames(M))
_fieldtypes(M) = [fieldtype(M, fld) for fld in fieldnames(M)]
hyperparameter_types(M::Type) = string.(_fieldtypes(M))
# function hyperparmeter_defaults(M::Type)
#     try
#         model = M()
#         return [Meta.parse(string(getproperty(model, fld)))
#                 for fld in fieldnames(M)]
#     catch
#         return []
#     end
# end

# following 5 lines commented out because they dissallow precompilation:
# for trait in MODEL_TRAITS
#     eval(quote
#         $trait(object) = $trait(typeof(object))
#     end)
# end

input_scitype(object) = input_scitype(typeof(object))
output_scitype(object) = output_scitype(typeof(object))
target_scitype(object) = target_scitype(typeof(object))
is_pure_julia(object) = is_pure_julia(typeof(object))
package_name(object) = package_name(typeof(object))
package_license(object) = package_license(typeof(object))
load_path(object) = load_path(typeof(object))
package_uuid(object) = package_uuid(typeof(object))
package_url(object) = package_url(typeof(object))
is_wrapper(object) = is_wrapper(typeof(object))
supports_weights(object) = supports_weights(typeof(object))
supports_online(object) = supports_online(typeof(object))
docstring(object) = docstring(typeof(object))
name(object) = name(typeof(object))
is_supervised(object) = is_supervised(typeof(object))
prediction_type(object) = prediction_type(typeof(object))
implemented_methods(object) = implemented_methods(typeof(object))
hyperparameters(object) = hyperparameters(typeof(object))
hyperparameter_types(object) = hyperparameter_types(typeof(object))


function +(a,b)
    println(a,"+",b)
    return a-b
end

i=8+9

function emmm(x::Int=9)
    println(x+8)
end

emmm()
emmm(1)


let
    biomass_f = 0.01

    # Feed jobs
    Ch = Channel(1) do ch
        cGLCs = Hd.val("cGLC")
        for (exp, cGLC)  in enumerate(cGLCs)
            put!(ch, (exp, cGLC))
        end
    end

    @threads for thid in 1:nthreads()
        for (exp, cGLC) in Ch
            # handle cache
            datfile = dat_file(string(DAT_FILE_PREFFIX, ME_Z_FIXXED_G_BOUNDED); exp)
            if isfile(datfile)
                lock(WLOCK) do
                    INDEX[ME_Z_FIXXED_G_BOUNDED, :DFILE, exp] = datfile
                    @info("Cached loaded (skipping)",
                        exp, cGLC, datfile, thid
                    ); println()
                end
                continue
            end

            # setup
            model = load_model(exp)
            objidx = ChU.rxnindex(model, iJR.BIOMASS_IDER)
            M, N = size(model)
            exp_growth = Hd.val("D", exp)
            ChU.ub!(model, iJR.BIOMASS_IDER, exp_growth * (1.0 + biomass_f))
            ChU.lb!(model, iJR.BIOMASS_IDER, exp_growth * (1.0 - biomass_f))
            model = ChLP.fva_preprocess(model,
                batchlen = 50,
                check_obj = iJR.BIOMASS_IDER,
                verbose = true
            )

            lock(WLOCK) do
                @info("Doing $(ME_Z_FIXXED_G_BOUNDED)",
                    exp, cGLC, thid
                ); println()
            end

            # maxent
            epout = ChEP.maxent_ep(model; alpha = Inf, damp = 0.985, epsconv = 1e-4,
                        verbose = false, maxiter = 5000
                    )

            # storing
            lock(WLOCK) do
                # Storing
                dat = Dict()
                dat[:exp_beta] = 0.0
                dat[:epouts] = Dict(0.0 => epout)
                dat[:model] = model |> ChU.compressed_model

                # caching
                serialize(datfile, dat)
                INDEX[ME_Z_FIXXED_G_BOUNDED, :DFILE, exp] = datfile

                @info("Finished ", exp, thid)
                println()
            end
        end # for (exp, cGLC) in Ch
    end # for thid in 1:nthreads()
end


###This file intends to demonstrate the rendering and basically follows basic.jl
import MuJoCo
import GLFW

using MuJoCo
using MuJoCo.time

mj_activate()
xmlfile = joinpath(@__DIR__, "cartpoleBalance.xml")
model = MJModel(xmlfile)
data = MJData(model)

window = GLFW.CreateWindow(640, 480, "demo")
GLFW.MakeContextCurrent(window)
GLFW.SwapInterval(1)

renderer = MJRenderer(data)

width = Int32(0)
height = Int32(0)
mj_step(model, data)
while !GLFW.WindowShouldClose(window)

    #simulate
    simstart = time(data)
    while (time(data) - simstart < 1.0/60.0 )
        mj_step(model, data)
    end

    #Render
    width, height = GLFW.GetFramebufferSize(window);
    mjv_updateScene(renderer)
    mjr_render(renderer, width, height)

    # Swap front and back buffers
	GLFW.SwapBuffers(window)

	# Poll for and process events
	GLFW.PollEvents()
end

GLFW.Terminate()


module Tracker

using MacroTools
using MacroTools: @q, @forward

import Base: ==

export TrackedArray, TrackedVector, TrackedMatrix, Params, param, back!

tracker(x) = nothing

istracked(x) = tracker(x) ≠ nothing
isleaf(x) = !istracked(x) || isleaf(tracker(x))
grad(x) = grad(tracker(x))
grad(::Nothing) = nothing
data(x) = x

struct Call{F,As<:Tuple}
  func::F
  args::As
end

Call(f::F, args::T) where {F,T} = Call{F,T}(f, args)
Call() = Call(nothing, ())

# When deserialising, the object_id changes
a::Call == b::Call = a.func == b.func && a.args == b.args

@inline (c::Call)() = c.func(data.(c.args)...)

mutable struct Tracked{T}
  ref::UInt32
  f::Call
  isleaf::Bool
  grad::T
  Tracked{T}(f::Call) where T = new(0, f, false)
  Tracked{T}(f::Call, grad::T) where T = new(0, f, false, grad)
  Tracked{T}(f::Call{Nothing}, grad::T) where T = new(0, f, true, grad)
end

istracked(x::Tracked) = true
isleaf(x::Tracked) = x.f == Call()
grad(x::Tracked) = x.grad

track(f::Call, x) = Tracked{typeof(x)}(f)

function _forward end

function track(f::F, xs...; kw...) where F
  y, back = _forward(f, xs...; kw...)
  track(Call(back, tracker.(xs)), y)
end

macro grad(ex)
  @capture(shortdef(ex), (name_(args__) = body_) |
                         (name_(args__) where {T__} = body_)) || error("Need a function definition")
  T == nothing && (T = [])
  isexpr(name, :(::)) || (name = :(::typeof($name)))
  insert!(args, 1+isexpr(args[1], :parameters) , name)
  @q(Tracker._forward($(args...)) where $(T...) = $body) |> esc
end

function update!(x, Δ)
  x.data .+= data(Δ)
  tracker(x).grad .= 0
  return x
end

include("idset.jl")
include("back.jl")
include("scalar.jl")
include("array.jl")
include("numeric.jl")

"""
    hook(f, x) -> x′

Hook into gradient backpropagation. `x` is unmodified, but when backpropagating
`f` will be applied to the incoming gradient. For example, `hook(-, x)` will reverse
the sign of the gradient applied to `x`."""
hook(f, x) = istracked(x) ? track(hook, f, x) : x
@grad hook(f, x) = data(x), Δ -> (nothing, f(Δ))

"""
    checkpoint(f, args...)

Behaves like `f(args...)`, but avoids storing the intermediate values needed for
calculating gradients. Instead, `f(args...)` will be called again during the
backward pass. This can be used to save memory in larger models.
"""
checkpoint(f, args...) = track(checkpoint, f, args...)

@grad function checkpoint(f, args...)
  data(f(args...)), function (Δ)
    y, back = forward(f, args...)
    (nothing, back(Δ)...)
  end
end

nobacksies(f, x) = track(nobacksies, f, x)
nobacksies(f, xs::Tuple) = map(x -> nobacksies(f, x), xs)
@grad nobacksies(f, x) = data(x), Δ -> error("Nested AD not defined for $f")

param(x::Number) = TrackedReal(float(x))
param(xs::AbstractArray) = TrackedArray(float.(xs))

@grad identity(x) = data(x), Δ -> (Δ,)
param(x::TrackedReal) = track(identity, x)
param(x::TrackedArray) = track(identity, x)

import Adapt.adapt

adapt(T, xs::TrackedArray) = param(adapt(T, data(xs)))

end


"""
The module `PDB` defines types and methods to work with protein structures inside Julia.
It is useful to link structural and sequential information, and
needed for measure the predictive performance at protein contact prediction of mutual information scores.

**Features**

- Read and parse PDF and PDBML files
- Calculate distance and contacts between atoms or residues
- Determine interaction between residues

```julia

using MIToS.PDB
```
"""
module PDB

using LightXML
using AutoHashEquals
using DataStructures
using Formatting
using MIToS.Utils
using PairwiseListMatrices
using FixedSizeArrays
using RecipesBase           # Plots for PDBResidue

import Base: ==, hash, length, size, -, +, ./, norm, dot, angle, cross, vec, any, print, show, parse
import MIToS.Utils: findobjects, isobject

export covalentradius, vanderwaalsradius, check_atoms_for_interactions,

PDBResidueIdentifier, Coordinates, PDBAtom, PDBResidue,
distance, contact, findheavy, findatoms, findCB, selectbestoccupancy, bestoccupancy,
angle, proximitymean,

ishydrophobic, isaromatic, iscationic, isanionic,
ishbonddonor, ishbondacceptor, hydrogenbond,
vanderwaals, vanderwaalsclash, covalent, disulphide,
aromaticsulphur, pication, aromatic, ionic, hydrophobic,
stridehydrogenbond, chimerahydrogenbond,

PDBFile, PDBML, downloadpdb, getpdbdescription,

kabsch, center!, rmsd,
getCA, CAmatrix, coordinatesmatrix, change_coordinates,
centeredcoordinates, centeredresidues,
superimpose,
mean_coordinates, rmsf,

# Mitos.Utils
isobject, findobjects, Is, Not, In, collectobjects, collectcaptures,

@residues, residues, @atoms, atoms, @residuesdict, residuesdict

include("PDBResidues.jl")
include("AtomsData.jl")
include("Interaction.jl")
include("PDBMLParser.jl")
include("PDBParser.jl")
include("Kabsch.jl")
include("Plots.jl")

@deprecate bestoccupancy! bestoccupancy

end


function device!(
    x,
    output_ode,
    voltage_r,
    voltage_i,
    current_r,
    current_i,
    ix_range::UnitRange{Int},
    ode_range::UnitRange{Int},
    dynamic_device::DynG,
    inputs::SimulationInputs,
) where {DynG <: PSY.DynamicGenerator}
    #Obtain local device states
    device_states = @view x[ix_range]

    #Obtain references
    sys_Sbase = get_base_power(inputs)
    sys_f0 = get_base_frequency(inputs)
    sys_ω = get_ω_sys(inputs)

    #Update Voltage data
    get_inner_vars(dynamic_device)[VR_gen_var] = voltage_r[1]
    get_inner_vars(dynamic_device)[VI_gen_var] = voltage_i[1]

    #Obtain ODEs and Mechanical Power for Turbine Governor
    mdl_tg_ode!(device_states, view(output_ode, ode_range), sys_ω, dynamic_device)

    #Obtain ODEs for AVR
    mdl_pss_ode!(device_states, view(output_ode, ode_range), sys_ω, dynamic_device)

    #Obtain ODEs for AVR
    mdl_avr_ode!(device_states, view(output_ode, ode_range), dynamic_device)

    #Obtain ODEs for Machine
    mdl_machine_ode!(
        device_states,
        view(output_ode, ode_range),
        current_r,
        current_i,
        sys_Sbase,
        sys_f0,
        dynamic_device,
    )

    #Obtain ODEs for PSY.Shaft
    mdl_shaft_ode!(
        device_states,
        view(output_ode, ode_range),
        sys_f0,
        sys_ω,
        dynamic_device,
    )

    return
end

function device!(
    voltage_r,
    voltage_i,
    current_r,
    current_i,
    device::PSY.Source,
    ::SimulationInputs,
)
    mdl_source!(voltage_r, voltage_i, current_r, current_i, device)
    return
end

function device!(
    voltage_r,
    voltage_i,
    current_r,
    current_i,
    device::PSY.PowerLoad,
    ::SimulationInputs,
)
    mdl_Zload!(voltage_r, voltage_i, current_r, current_i, device)
    return
end

function device!(
    x,
    output_ode::Vector{T},
    voltage_r,
    voltage_i,
    current_r,
    current_i,
    ix_range::UnitRange{Int},
    ode_range::UnitRange{Int},
    dynamic_device::DynI,
    inputs::SimulationInputs,
) where {DynI <: PSY.DynamicInverter, T <: Real}
    #Obtain local device states
    device_states = @view x[ix_range]

    #Obtain references
    Sbase = get_base_power(inputs)
    sys_f0 = get_base_frequency(inputs)
    sys_ω = get_ω_sys(inputs)

    #Update Voltage data
    get_inner_vars(dynamic_device)[VR_inv_var] = voltage_r[1]
    get_inner_vars(dynamic_device)[VI_inv_var] = voltage_i[1]

    #Update V_ref
    V_ref = PSY.get_ext(dynamic_device)[CONTROL_REFS][V_ref_index]
    get_inner_vars(dynamic_device)[V_oc_var] = V_ref

    #Obtain ODES for DC side
    mdl_DCside_ode!(
        device_states,
        view(output_ode, ode_range),
        sys_f0,
        sys_ω,
        dynamic_device,
    )

    #Obtain ODEs for PLL
    mdl_freq_estimator_ode!(
        device_states,
        view(output_ode, ode_range),
        sys_f0,
        sys_ω,
        dynamic_device,
    )

    #Obtain ODEs for OuterLoop
    mdl_outer_ode!(
        device_states,
        view(output_ode, ode_range),
        sys_f0,
        sys_ω,
        dynamic_device,
    )

    #Obtain inner controller ODEs and modulation commands
    mdl_inner_ode!(device_states, view(output_ode, ode_range), dynamic_device)

    #Obtain converter relations
    mdl_converter_ode!(dynamic_device)

    #Obtain ODEs for output filter
    mdl_filter_ode!(
        device_states,
        view(output_ode, ode_range),
        current_r,
        current_i,
        Sbase,
        sys_f0,
        sys_ω,
        dynamic_device,
    )

    return
end


# wraps variables found in EIGEN.INC for use in Julia

const NEIGX = 300
const NFREQX = 120
const NZX = 18*IIX + NGVX + NGPX + NGFX
const NDFFX = NFLPX + NENGX + 2*NYGUSX-1

struct EIGEN_s
    LROM::Array{Bool,1}
    NVTOT::Array{Int32,1}
    NEIGEN::Array{Int32,1}
    NEIDAT::Array{Int32,1}
    IPEIG::OffsetArrays.OffsetArray{Int32,1,Array{Int32,1}}
    IPDAT::OffsetArrays.OffsetArray{Int32,1,Array{Int32,1}}
    IPFRQ::Array{Int32,1}
    NFREQ::Array{Int32,1}
    NFDAT::Array{Int32,1}
    NROM::Array{Int32,1}
    NROMC::Array{Int32,1}
    KROM::Array{Int32,1}
    ZESEED::Array{ComplexF64,1}
    ZLAMDA::OffsetArrays.OffsetArray{ComplexF64,1,Array{ComplexF64,1}}
    ZLADAT::OffsetArrays.OffsetArray{ComplexF64,1,Array{ComplexF64,1}}
    ZQEIG::OffsetArrays.OffsetArray{ComplexF64,2,Array{ComplexF64,2}}
    ZLEIG::OffsetArrays.OffsetArray{ComplexF64,2,Array{ComplexF64,2}}
    ZWB::Array{ComplexF64,2}
    ZWBC::Array{ComplexF64,2}
    ZROML::Array{ComplexF64,2}
    ZROMCL::Array{ComplexF64,2}
    ZROM::Array{ComplexF64,1}
    ZROMC::Array{ComplexF64,1}
    ZQ::Array{ComplexF64,2}
    ZPARAM::Array{ComplexF64,1}
    ZQSENS::Array{ComplexF64,2}
    ZQSENSF::Array{ComplexF64,3}
    ZAN::Array{ComplexF64,1}
    ZQJOIN::Array{ComplexF64,2}
    ZAFORCE::Array{ComplexF64,1}
    ZRFORCE::Array{ComplexF64,1}
    ZGFORCE::Array{ComplexF64,1}
    ZTFORCE::Array{ComplexF64,1}
    ZAMOMNT::Array{ComplexF64,1}
    ZRMOMNT::Array{ComplexF64,1}
    ZGMOMNT::Array{ComplexF64,1}
    ZTMOMNT::Array{ComplexF64,1}
    ZFLIFT::Array{ComplexF64,1}
    ZFDRAG::Array{ComplexF64,1}
    ZFUNS::OffsetArrays.OffsetArray{ComplexF64,3,Array{ComplexF64,3}}
    ZDATS::OffsetArrays.OffsetArray{ComplexF64,3,Array{ComplexF64,3}}
    ZQSFUNS::OffsetArrays.OffsetArray{ComplexF64,4,Array{ComplexF64,4}}
    ZQSDATS::OffsetArrays.OffsetArray{ComplexF64,4,Array{ComplexF64,4}}
    AQEIG::OffsetArrays.OffsetArray{Float64,1,Array{Float64,1}}
    EPHASE::Array{Float64,1}
    FREQ::OffsetArrays.OffsetArray{Float64,1,Array{Float64,1}}
    FDAT::OffsetArrays.OffsetArray{Float64,2,Array{Float64,2}}
    function EIGEN_s()
        as_eigl = cglobal((:as_eigl_, libaswing), Int32); as_eigl_idx = as_eigl
        LROM = unsafe_wrap(Array, as_eigl_idx, NEIGX)

        as_eigi = cglobal((:as_eigi_, libaswing), Int32); as_eigi_idx = as_eigi
        NVTOT = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        NEIGEN = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        NEIDAT = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        ipeig = unsafe_wrap(Array, as_eigi_idx, NEIGX+1); as_eigi_idx += sizeof(Int32)*(NEIGX+1)
        IPEIG = OffsetArrays.OffsetArray(ipeig, 0:NEIGX)
        ipdat = unsafe_wrap(Array, as_eigi_idx, NEIGX+1); as_eigi_idx += sizeof(Int32)*(NEIGX+1)
        IPDAT = OffsetArrays.OffsetArray(ipdat, 0:NEIGX)
        IPFRQ = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        NFREQ = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        NFDAT = unsafe_wrap(Array, as_eigi_idx, NDFFX); as_eigi_idx += sizeof(Int32)*NDFFX
        NROM = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        NROMC = unsafe_wrap(Array, as_eigi_idx, 1); as_eigi_idx += sizeof(Int32)
        KROM = unsafe_wrap(Array, as_eigi_idx, NEIGX)

        as_eigz = cglobal((:as_eigz_, libaswing), ComplexF64); as_eigz_idx = as_eigz
        ZESEED = unsafe_wrap(Array, as_eigz_idx, 1); as_eigz_idx += sizeof(ComplexF64)
        zlamda = unsafe_wrap(Array, as_eigz_idx, NEIGX+1); as_eigz_idx += sizeof(ComplexF64)*(NEIGX+1)
        ZLAMDA = OffsetArrays.OffsetArray(zlamda, 0:NEIGX)
        zladat = unsafe_wrap(Array, as_eigz_idx, NEIGX+1); as_eigz_idx += sizeof(ComplexF64)*(NEIGX+1)
        ZLADAT = OffsetArrays.OffsetArray(zladat, 0:NEIGX)
        zqeig = unsafe_wrap(Array, as_eigz_idx, (NZX, NEIGX+1)); as_eigz_idx += sizeof(ComplexF64)*NZX*(NEIGX+1)
        ZQEIG = OffsetArrays.OffsetArray(zqeig, 1:NZX, 0:NEIGX)
        zleig = unsafe_wrap(Array, as_eigz_idx, (NZX, NEIGX+1)); as_eigz_idx += sizeof(ComplexF64)*NZX*(NEIGX+1)
        ZLEIG = OffsetArrays.OffsetArray(zleig, 1:NZX, 0:NEIGX)
        ZWB = unsafe_wrap(Array, as_eigz_idx, (NEIGX, NFRPX)); as_eigz_idx += sizeof(ComplexF64)*NEIGX*NFRPX
        ZWBC = unsafe_wrap(Array, as_eigz_idx, (NEIGX, NFRPX)); as_eigz_idx += sizeof(ComplexF64)*NEIGX*NFRPX
        ZROML = unsafe_wrap(Array, as_eigz_idx, (NEIGX, NFRPX)); as_eigz_idx += sizeof(ComplexF64)*NEIGX*NFRPX
        ZROMCL = unsafe_wrap(Array, as_eigz_idx, (NEIGX, NFRPX)); as_eigz_idx += sizeof(ComplexF64)*NEIGX*NFRPX
        ZROM = unsafe_wrap(Array, as_eigz_idx, NEIGX); as_eigz_idx += sizeof(ComplexF64)*NEIGX
        ZROMC = unsafe_wrap(Array, as_eigz_idx, NEIGX); as_eigz_idx += sizeof(ComplexF64)*NEIGX
        ZQ = unsafe_wrap(Array, as_eigz_idx, (18, IIX)); as_eigz_idx += sizeof(ComplexF64)*18*IIX
        ZPARAM = unsafe_wrap(Array, as_eigz_idx, KPTOT); as_eigz_idx += sizeof(ComplexF64)*KPTOT
        ZQSENS = unsafe_wrap(Array, as_eigz_idx, (KSTOT, NSENX)); as_eigz_idx += sizeof(ComplexF64)*KSTOT*NSENX
        ZQSENSF = unsafe_wrap(Array, as_eigz_idx, (KSTOT, NSENX, NFRPX)); as_eigz_idx += sizeof(ComplexF64)*KSTOT*NSENX*NFRPX
        ZAN = unsafe_wrap(Array, as_eigz_idx, NNX); as_eigz_idx += sizeof(ComplexF64)*NNX
        ZQJOIN = unsafe_wrap(Array, as_eigz_idx, (12, NJX)); as_eigz_idx += sizeof(ComplexF64)*12*NJX
        ZAFORCE = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZRFORCE = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZGFORCE = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZTFORCE = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZAMOMNT = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZRMOMNT = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZGMOMNT = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZTMOMNT = unsafe_wrap(Array, as_eigz_idx, 3); as_eigz_idx += sizeof(ComplexF64)*3
        ZFLIFT = unsafe_wrap(Array, as_eigz_idx, 1); as_eigz_idx += sizeof(ComplexF64)
        ZFDRAG = unsafe_wrap(Array, as_eigz_idx, 1); as_eigz_idx += sizeof(ComplexF64)
        zfuns = unsafe_wrap(Array, as_eigz_idx, (NFREQX+1, KPTOT, NDFFX)); as_eigz_idx += sizeof(ComplexF64)*(NFREQX+1)*KPTOT*NDFFX
        ZFUNS = OffsetArrays.OffsetArray(zfuns, 0:NFREQX, 1:KPTOT, 1:NDFFX)
        zdats = unsafe_wrap(Array, as_eigz_idx, (NFREQX+1, KPTOT, NDFFX)); as_eigz_idx += sizeof(ComplexF64)*(NFREQX+1)*KPTOT*NDFFX
        ZDATS = OffsetArrays.OffsetArray(zdats, 0:NFREQX, 1:KPTOT, 1:NDFFX)
        zqsfuns = unsafe_wrap(Array, as_eigz_idx, (NFREQX+1, KSTOT, NSENX, NDFFX)); as_eigz_idx += sizeof(ComplexF64)*(NFREQX+1)*KSTOT*NSENX*NDFFX
        ZQSFUNS = OffsetArrays.OffsetArray(zqsfuns, 0:NFREQX, 1:KSTOT, 1:NSENX, 1:NDFFX)
        zqsdats = unsafe_wrap(Array, as_eigz_idx, (NFREQX+1, KSTOT, NSENX, NDFFX)); as_eigz_idx += sizeof(ComplexF64)*(NFREQX+1)*KSTOT*NSENX*NDFFX
        ZQSDATS = OffsetArrays.OffsetArray(zqsdats, 0:NFREQX, 1:KSTOT, 1:NSENX, 1:NDFFX)

        as_eigr = cglobal((:as_eigr_, libaswing), Float64); as_eigr_idx = as_eigr
        aqeig = unsafe_wrap(Array, as_eigr_idx, (NEIGX+1)); as_eigr_idx += sizeof(Float64)*(NEIGX+1)
        AQEIG = OffsetArrays.OffsetArray(aqeig, 0:NEIGX)
        EPHASE = unsafe_wrap(Array, as_eigr_idx, 1); as_eigr_idx += sizeof(Float64)
        freq = unsafe_wrap(Array, as_eigr_idx, (NFREQX+1)); as_eigr_idx += sizeof(Float64)*(NFREQX+1)
        FREQ = OffsetArrays.OffsetArray(freq, 0:NFREQX)
        fdat = unsafe_wrap(Array, as_eigr_idx, (NFREQX+1, NDFFX))
        FDAT = OffsetArrays.OffsetArray(fdat, 0:NFREQX, 1:NDFFX)

        return new(LROM, NVTOT, NEIGEN, NEIDAT, IPEIG,
            IPDAT, IPFRQ, NFREQ, NFDAT, NROM, NROMC, KROM, ZESEED, ZLAMDA,
            ZLADAT, ZQEIG, ZLEIG, ZWB, ZWBC, ZROML, ZROMCL, ZROM, ZROMC, ZQ,
            ZPARAM, ZQSENS, ZQSENSF, ZAN, ZQJOIN, ZAFORCE, ZRFORCE, ZGFORCE,
            ZTFORCE, ZAMOMNT, ZRMOMNT, ZGMOMNT, ZTMOMNT, ZFLIFT, ZFDRAG, ZFUNS,
            ZDATS, ZQSFUNS, ZQSDATS, AQEIG, EPHASE, FREQ, FDAT)
    end
end


"""
Ase fileformat
ASCII Segment File
Described as route network. Each line represents a route segment

See EUROCONTROL NEST Manual for Ase fileformat description

Ase DataFrame Column Names
#   Field           Type        Size    Comment
1   :FLIGHTCOUNT    Float64             normally number of flights using this route segment, could be a load
2   :SEGMENTPARITY  Int64               0=NO, 1=ODD, 2=EVEN, 3=ODD_LOW, 4=EVEN_LOW, 5=ODD_HIGH, 6=EVEN_HIGH
3   :SEGMENTTYPE    Int64               0=NO, 1=NORMAL, 2=ARRIVAL, 3=DEPARTURE (permanent rte segment)
                                        20=NO, 21=NORMAL, 22=ARRIVAL, 23=DEPARTURE (CDR Generic)
                                        40=NO, 41=NORMAL, 42=ARRIVAL, 43=DEPARTURE (CDR 1 )
                                        60=NO, 61=NORMAL, 62=ARRIVAL, 63=DEPARTURE (CDR 2)
                                        80=NO, 81=NORMAL, 82=ARRIVAL, 83=DEPARTURE (CDR 3)
                                        100=NO, 101=NORMAL, 102=ARRIVAL, 103=DEPARTURE (CDR 1+2)
                                        120=NO, 121=NORMAL, 122=ARRIVAL, 123=DEPARTURE (CDR 1+3)
4   :LATBEGINSEGMENT_DEG Float64
5   :LONBEGINSEGMENT_DEG Float64
6   :LATENDSEGMENT_DEG   Float64
7   :LONENDSEGMENT_DEG   Float64
8   :SEGMENTNAME    String      11      routePointNameBegin_routePointNameEnd (separator '_')

%% = TBD
%  = TBD
*  = TBD
"""
module Ase

export read

using Format
using CSV
using Dates

const fileformat = Dict(1=>Float64, 2=>Int64, 3=>Int64, 4=>Float64,
5=>Float64, 6=>Float64, 7=>Float64, 8=>String)

const header_format = ["FLIGHTCOUNT", "SEGMENTPARITY", "SEGMENTTYPE",
"LATBEGINSEGMENT_DEG", "LONBEGINSEGMENT_DEG", "LATENDSEGMENT_DEG", "LONENDSEGMENT_DEG",
"SEGMENTNAME"]

function read(file)
    df = CSV.read(file, delim=" ", types=fileformat, header=header_format,
    copycols=true)
    reformat!(df)
    return df
end

function reformat!(df)
    df[:,:LATBEGINSEGMENT_DEG] = df[:,:LATBEGINSEGMENT_DEG] / 60.0
    df[:,:LONBEGINSEGMENT_DEG] = df[:,:LONBEGINSEGMENT_DEG] / 60.0
    df[:,:LATENDSEGMENT_DEG] = df[:,:LATENDSEGMENT_DEG] / 60.0
    df[:,:LONENDSEGMENT_DEG] = df[:,:LONENDSEGMENT_DEG] / 60.0
end

end # module


matrix = zeros(Float64, (1000,1000,1000))
# use matrix, then when done set variable to 0 to garbage collect the matrix:
matrix = 0 # large memory pool will now be collected when needed


using BioSequences
using FASTX
using Test
using DataFrames
using CSV
#tes Alt parsing and repr update saving writing to fasta
using st16SseqJuliaTools

primers = "/home/gediminas/working/primers/primers.fasta"
db_file = "/home/gediminas/Desktop/tm/GRCh37_latest_rna.fna"

#collect primer  st16SseqJuliaTools
primer_inf = Array{Tuple{String,String},1}()
reader = FASTA.Reader(open(primers, "r"))
record = FASTA.Record()

ctp = 0
while !eof(reader)
    global ctp += 1
    read!(reader, record)
    id = identifier(record)
    seq = sequence(record)
    push!(primer_inf,(id,seq))
    if ctp == 100
        break
    end
end
close(reader)

data = DataFrame(
 ID=Array{String,1}(),
 SEQ=Array{String,1}(),
 K=Array{Int64,1}(),
 K2=Array{Int64,1}(),
 N=Array{Int64,1}(),
 NMBBIND=Array{Int64,1}(),
 TAIL=Array{Float64,1}(),
 E=Array{Float64,1}(),
 KMER=Array{String,1}(),
 KMER_COUNT=Array{Int64,1}(),
 KMER_MSIMAT=Array{Int64,1}(),
 TM_DNA_DNA = Array{Float64,1}(),
 TM_DNA_RNA_HS = Array{Float64,1}(),
 TM_DNA_RNA_LS = Array{Float64,1}(),
 TM_DNA_DNA_KMER = Array{Float64,1}(),
)

ctn = 0
for (k,k2,n) in [(15,3,1),(15,2,1),(15,0,1),(15,3,2),(15,2,2),(15,0,2),(15,3,3),(15,2,3),(15,0,3),(6,3,1),(6,2,1),(6,0,1),(6,3,2),(6,2,2),(6,0,2),(6,3,3),(6,2,3),(6,0,3)]
    kmers_db = KmerData(db_file,k)
    sm = get_significance_model(k,k2,n,kmers_db)
    global ctn += 1
    println("Working with this set", (k,k2,n) )
    for (id,seq) in primer_inf
        parts = []
        push!(parts,id)
        push!(parts,seq)
        push!(parts,k)
        push!(parts,k2)
        push!(parts,n)
        rescal = collect_matches_data(seq,k,k2,n,kmers_db,sm)
        for r in rescal
            push!(parts,r)
        end
        push!(data,parts)
    end
    CSV.write("REZ_$ctn.csv",data)
end
println(data)


export Configuration,ReadConfig

function ReadConfig(filename::String)
    input = open(filename, "r")
    inputs=JSON.parse(input)
    close(input)
    return inputs
end

mutable struct Molecule
    type::Array{Int64,1}
    mass::Array{Float64,1}
    coords::Array{Float64,2}
    bonds::Array{Int64,2}
    angles::Array{Int64,2}
    vels::Array{Float64,2}
end

#function Molecule(type::Array{Int64,1},
#    mass::Array{Float64,1},
#    coords::Array{Float64,2},
#    bonds::Array{Array{Int64,2},1},
#    angles::Array{Array{Int64,2},1},
#    vels::Array{Float64,2})
#new(
#type::Array{Int64,1},
#mass::Array{Float64,1},
#coords::Array{Float64,2},
#bonds::Array{Array{Int64,2},1},
#angles::Array{Array{Int64,2},1},
#vels::Array{Float64,2}
#)
#end

struct Configuration
    start::Int64
    box::Array{Float64,1}
    moles::Array{Molecule,1}
    moles_length::Array{Int64,1}
    function Configuration(data::Dict)
        start=1
        box=zeros(3)
        if "box" in keys(data)
            box=convert(Array{Float64,1}, data["box"])
            @info "Box Lx Ly Lz: $(box[1]) $(box[2]) $(box[3])"
        end

        moles=Vector{Molecule}(undef,0)
        moles_length=zeros(Int64,0)

        if "molecules" in keys(data)
            molecules=data["molecules"]
            num_moles=length(keys(molecules))

            moles=Vector{Molecule}(undef,num_moles)

            moles_length=zeros(Int64,num_moles)
            for molei in 1:num_moles
                moles_length[molei]=length(keys(molecules[string(molei)]["atoms"]))
                #molecule=Molecule(moles_length[molei])
                type=zeros(Int64,moles_length[molei])
                mass=zeros(Int64,moles_length[molei])
                bonds=zeros(Int64,0,3)
                angles=zeros(Int64,moles_length[molei],4)
                coords=zeros(Float64,moles_length[molei],3)
                vels=zeros(Float64,moles_length[molei],3)
                for atomi in 1:moles_length[molei]
                    type[atomi]=molecules[string(molei)]["atoms"][string(atomi)]["type"]
                    mass[atomi]=molecules[string(molei)]["atoms"][string(atomi)]["mass"]
                    coords[atomi,1:3]=convert(Array{Float64,1},molecules[string(molei)]["atoms"][string(atomi)]["coords"])
                    if "vels" in keys(molecules[string(molei)]["atoms"][string(atomi)])
                        vels[atomi,1:3]=convert(Array{Float64,1},molecules[string(molei)]["atoms"][string(atomi)]["vels"])
                    end
                end
                if "bonds" in keys(molecules[string(molei)])
                    num_bonds=length(keys(molecules[string(molei)]["bonds"]))
                    bonds=zeros(Int64,num_bonds,3)
                    for bondi in 1:num_bonds
                        bonds[bondi,1:3]=convert(Array{Int64,1},molecules[string(molei)]["bonds"][string(bondi)])
                    end
                end
                moles[molei]=Molecule(type,mass,coords,bonds,angles,vels)
            end
        end

        new(start,box,
            moles,
            moles_length)
    end
end


using DataFrames
using Dates

@testset "simulate" begin
    @testset "basic" begin
        @system SSimulate(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
        end
        n = 10
        r = simulate(SSimulate, stop=n)
        @test r isa DataFrame
        @test size(r, 1) == (n+1)
        @test propertynames(r) == [:time, :a, :b]
        @test r[end, :time] == n*u"hr"
        @test r[end, :a] == 1
        @test r[end, :b] == n
        r = simulate(SSimulate, stop=n, config=(:SSimulate => :a => 2))
        @test r[end, :a] == 2
        @test r[end, :b] == 2n
        r = simulate(SSimulate, stop=n, target=[:b])
        @test size(r, 2) == 2
        @test propertynames(r) == [:time, :b]
        r = simulate(SSimulate, stop=n, index=:b, target=[:b])
        @test size(r, 2) == 1
        @test propertynames(r) == [:b]
    end

    @testset "stop number" begin
        @system SSimulateStopNumber(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
        end
        r = simulate(SSimulateStopNumber, stop=5)
        @test r[end-1, :b] == 4
        @test r[end, :b] == 5
    end

    @testset "stop number unit" begin
        @system SSimulateStopNumberUnit(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
        end
        r1 = simulate(SSimulateStopNumberUnit, stop=2u"hr")
        @test r1[end, :b] == 2 * 1
        r2 = simulate(SSimulateStopNumberUnit, stop=2u"d")
        @test r2[end, :b] == 2 * 24
        r3 = simulate(SSimulateStopNumberUnit, stop=2u"yr")
        @test r3[end, :b] == 2 * 24 * 365.25
    end

    @testset "stop count" begin
        @system SSimulateStopCount(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
            c => 5 ~ preserve::int
        end
        r = simulate(SSimulateStopCount, stop=:c)
        @test r[end-1, :b] == 4
        @test r[end, :b] == 5
    end

    @testset "stop count unit" begin
        @system SSimulateStopCountUnit(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
            c ~ preserve::Cropbox.Quantity(parameter)
        end
        r1 = simulate(SSimulateStopCountUnit, stop=:c, config=:0 => :c => 2u"hr")
        @test r1[end, :b] == 2 * 1
        r2 = simulate(SSimulateStopCountUnit, stop=:c, config=:0 => :c => 2u"d")
        @test r2[end, :b] == 2 * 24
        r3 = simulate(SSimulateStopCountUnit, stop=:c, config=:0 => :c => 2u"yr")
        @test r3[end, :b] == 2 * 24 * 365.25
    end

    @testset "stop boolean" begin
        @system SSimulateStopBoolean(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
            z(b) => b >= 5 ~ flag
        end
        r1 = simulate(SSimulateStopBoolean, stop=:z)
        @test r1[end-1, :b] == 4
        @test r1[end, :b] == 5
        r2 = simulate(SSimulateStopBoolean, stop="z")
        @test r1 == r2
        r3 = simulate(SSimulateStopBoolean, stop=s -> s.b' >= 5)
        @test r1 == r3
    end

    @testset "snap" begin
        @system SSimulateSnap(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
            f(b) => (b % 2 == 0) ~ flag
        end
        n = 10
        f(s) = s.b' % 2 == 0
        r0 = simulate(SSimulateSnap, stop=n, snap=nothing)
        @test !all(r0.b .% 2 .== 0)
        r1 = simulate(SSimulateSnap, stop=n, snap=f)
        @test all(r1.b .% 2 .== 0)
        r2 = simulate(SSimulateSnap, stop=n, snap=:f)
        @test all(r2.b .% 2 .== 0)
        r3 = simulate(SSimulateSnap, stop=n, snap="f")
        @test all(r3.b .% 2 .== 0)
    end

    @testset "snatch" begin
        @system SSimulateSnatch(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
        end
        n = 3
        i = 0
        f(D, s) = begin
            d = D[1]
            @test s.a' == d[:a] && s.b' == d[:b]
            i += 1
            d[:c] = i
        end
        r = simulate(SSimulateSnatch, stop=n, snatch=f)
        @test i == 1 + n
        @test r[!, :c] == [1, 2, 3, 4]
    end

    @testset "callback" begin
        @system SSimulateCallback(Controller) begin
            a => 1 ~ preserve(parameter)
            b(a) ~ accumulate
        end
        n = 3
        i = 0
        f(s, m) = begin
            r = m.result[end, :]
            @test s.a' == r.a && s.b' == r.b
            i += 1
        end
        simulate(SSimulateCallback, stop=n, callback=f)
        @test i == n
    end

    @testset "layout" begin
        @system SSimulateLayout(Controller) begin
            i => 1 ~ accumulate
            a(i) => i-1 ~ track
            b(i) => 2i ~ track
        end
        L = [
            (target=:a,),
            (index=[:t => "context.clock.time", "i"], target=["a", :B => :b]),
            (base="context.clock", index="time", target="step"),
            (target=:b, meta=(:c => 0, :d => :D)),
        ]
        n = 1
        r = simulate(SSimulateLayout, L, stop=n)
        @test propertynames(r[1]) == [:time, :a]
        @test propertynames(r[2]) == [:t, :i, :a, :B]
        @test propertynames(r[3]) == [:time, :step]
        @test propertynames(r[4]) == [:time, :b, :c, :d]
        @test r[1][end, :time] == r[2][end, :t] == r[3][end, :time] == r[4][end, :time] == n*u"hr"
        @test r[1][end, :a] == 0
        @test r[2][end, :B] == 2
        @test r[3][end, :step] == 1u"hr"
        @test all(r[4][!, :c] .== 0) && all(r[4][!, :d] .== :D)
    end

    @testset "layout and configs" begin
        @system SSimulateLayoutConfigs(Controller) begin
            p ~ preserve(parameter)
            i => 1 ~ accumulate
            a(i, p) => p*(i-1) ~ track
            b(i, p) => 2p*i ~ track
        end
        L = [
            (index=:i, target=:a),
            (index=:t => "context.clock.time", target=:b),
            (target=[:i, :a, :b],),
        ]
        p1, p2 = 1, 2
        C = [
            :SSimulateLayoutConfigs => :p => p1,
            :SSimulateLayoutConfigs => :p => p2,
        ]
        n = 10
        r = simulate(SSimulateLayoutConfigs, L, C, stop=n)
        @test length(r) == length(L)
        o = r[3]
        @test o[o.time .== n*u"hr", :i] == [n, n]
        @test o[o.time .== n*u"hr", :a] == [p1*(n-1), p2*(n-1)]
        @test o[o.time .== n*u"hr", :b] == [2p1*n, 2p2*n]
    end

    @testset "configs" begin
        @system SSimulateConfigs(Controller) begin
            a ~ preserve(parameter)
            b(a) ~ accumulate
        end
        p = [1, 2]
        C = @config !(:SSimulateConfigs => :a => p)
        n = 10
        r = simulate(SSimulateConfigs, configs=C, stop=n)
        @test r[r.time .== n*u"hr", :a] == p
        @test r[r.time .== n*u"hr", :b] == p .* n
    end

    @testset "meta" begin
        @system SSimulateMeta(Controller) begin
            a ~ preserve(parameter)
            b(a) ~ accumulate
        end
        C = (
            :SSimulateMeta => (a=1,),
            :Extra => (b=2, c=:C),
        )
        n = 10
        r1 = simulate(SSimulateMeta, config=C, index=(), meta=:Extra, stop=n)
        @test propertynames(r1) == [:a, :b, :c]
        @test all(r1.b .== 2)
        @test all(r1.c .== :C)
        r2 = simulate(SSimulateMeta, config=C, index=(), meta=(:Extra, :d => 0), stop=n)
        @test propertynames(r2) == [:a, :b, :c, :d]
        @test r1.b == r2.b
        @test r1.c == r2.c
        @test all(r2.d .== 0)
    end

    @testset "wildcard" begin
        @system SSimulateWildcardA begin
            x => 1 ~ preserve
            y => 2 ~ preserve
        end
        @system SSimulateWildcard(Controller) begin
            A(context) ~ ::SSimulateWildcardA
            a => 3 ~ preserve
            b => 4 ~ preserve
        end
        r1 = simulate(SSimulateWildcard; target="*")
        @test names(r1) == ["time", "a", "b"]
        @test collect(r1[1,:]) == [0u"hr", 3, 4]
        r2 = simulate(SSimulateWildcard; target="A.*")
        @test names(r2) == ["time", "A.x", "A.y"]
        @test collect(r2[1,:]) == [0u"hr", 1, 2]
        r3 = simulate(SSimulateWildcard; target=["*", "A.*"])
        @test names(r3) == ["time", "a", "b", "A.x", "A.y"]
        @test collect(r3[1,:]) == [0u"hr", 3, 4, 1, 2]
        r4 = simulate(SSimulateWildcard; target=["A.*", "*"])
        @test names(r4) == ["time", "A.x", "A.y", "a", "b"]
        @test collect(r4[1,:]) == [0u"hr", 1, 2, 3, 4]
    end

    @testset "options" begin
        @system SSimulateOptions(Controller) begin
            a ~ preserve(extern)
            b ~ ::int(override)
        end
        n = 10
        a, b = 1, 2
        o = (; a, b)
        r = simulate(SSimulateOptions, options=o, stop=n)
        @test r[end, :a] == a
        @test r[end, :b] == b
    end

    @testset "seed" begin
        @system SSimulateSeed(Controller) begin
            a => rand() ~ track
            b(a) ~ accumulate
        end
        n = 10
        r1 = simulate(SSimulateSeed, seed=0, stop=n)
        @test r1[end, :a] == 0.5392892841426182
        @test r1[end, :b] == 3.766035118243237
        r2 = simulate(SSimulateSeed, seed=0, stop=n)
        @test r1 == r2
    end

    @testset "no seed" begin
        @system SSimulateNoSeed(Controller) begin
            a => rand() ~ track
            b(a) ~ accumulate
        end
        n = 10
        r1 = simulate(SSimulateNoSeed, seed=nothing, stop=n)
        r2 = simulate(SSimulateNoSeed, seed=nothing, stop=n)
        @test r1 != r2
    end

    @testset "extractable" begin
        @system SSimulateExtractable(Controller) begin
            a => 1 ~ track
            b => :hello ~ track::sym
            c => "world" ~ track::str
            d => DateTime(2020, 3, 1) ~ track::DateTime
            e => Dict(:k => 0) ~ track::Dict
            f => [1, 2, 3] ~ track::Vector
            g => (1, 2) ~ track::Tuple
        end
        r = simulate(SSimulateExtractable)
        N = propertynames(r)
        # default = Union{Number,Symbol,AbstractString,AbstractDateTime}
        @test :a ∈ N
        @test :b ∈ N
        @test :c ∈ N
        @test :d ∈ N
        @test :e ∉ N
        @test :f ∉ N
        @test :g ∉ N
    end
end