#===============================================================================

    AssignGenosResult.jl

    Author: Tom Price
    Date:   July 2019

===============================================================================#

import Clustering.ClusteringResult

struct AssignGenosResult
    cluster_result      ::ClusteringResult
    best_i              ::Int_T
    best_genos_combins  ::Vector{Matrix{Int_T}}
    expected_genos_all  ::Matrix{Int_T}
    ucc_dict            ::OrderedDict{Set{Vector{Float_T}},UniqCombinCenters}
end

## constants >>

const EMPTY_BEST_GENO_COMBINS = Vector{Matrix{Int_T}}()

## default values for processing preset calibration data with 4 groups
const DEFAULT_AMP_ENCGR = Array{Int_T,2}(0, 0)
## const DEFAULT_AMP_ENCGR = [0 1 0 1; 0 0 1 1] ## NTC, homo ch1, homo ch2, hetero
const DEFAULT_AMP_INIT_FACTORS = [1, 1, 1, 1] ## sometimes "hetero" may not have very high end-point fluo
const DEFAULT_AMP_APG_LABELS = ["ntc", "homo_1", "homo_2", "hetero", "unclassified"] ## [0 1 0 1; 0 0 1 1]
## const DEFAULT_AMP_APG_LABELS = ["hetero", "homo_2", "homo_1", "ntc", "unclassified"] ## [1 0 1 0; 1 1 0 0]

## 3 groups without NTC (non-template control)
# const DEFAULT_EGR = [1 0 1; 0 1 1] # homo ch1, homo ch2, hetero
# const DEFAULT_AMP_INIT_FACTORS = [1, 1, 1] # sometimes "hetero" may not have very high end-point fluo
# const DEFAULT_AMP_EG_LABELS = ["homo_a", "homo_b", "hetero", "unclassified"]

const DEFAULT_AMP_CTRL_WELL_DICT = OrderedDict{Vector{Int_T},Vector{Int_T}}()
## key is genotype (Vector{Int}), value is well numbers (Vector{Int_T})
## example
# const DEFAULT_AMP_CTRL_WELL_DICT = OrderedDict(
#     [0, 0] => [1, 2], # NTC, well 1 and 2
#     [1, 0] => [3, 4], # homo ch1, well 3 and 4
#     [0, 1] => [5, 6], # homo ch2, well 5 and 6
#     [1, 1] => [7, 8]  # hetero, well 7 and 8
# )
## old approach
# const DEFAULT_AMP_CTRL_WELL_DICT = DefaultOrderedDict(Vector{Int_T}, Vector{Int_T}, Vector{Int_T}())

const DEFAULT_AMP_CATEG_WELL_VEC = Vector{Pair{Symbol,Any}}([
    :rbbs_ary3      => Colon(),     ## calibrated_data
    :blsub_fluos    => Colon(), ## baseline_subtracted_data
    :d0             => Colon(),
    :cq             => Colon()])

# const DEFAULT_AMP_CYCS                  = 0
const DEFAULT_AMP_CLUSTER_METHOD        = k_means_medoids
const DEFAULT_AMP_NORM_L                = 2


import Random: rand, seed!
import RandomNumbers: AbstractRNG
export Chacha8SIMD, ChachaSIMD
# Chacha Random Number Generator with SIMD-128 in mind
# DOUBLEROUNDS=10 means Chacha20
# Generates 64-bit random numbers

mutable struct ChachaSIMD{DOUBLEROUNDS,T} <: AbstractRNG{T}
    v::Vector{NTuple{4,UInt32}}
    x::Vector{NTuple{4,UInt32}}
    subCounter::Int8
end

function Chacha8SIMD(T = UInt64) #::Union{Type{UInt32},Type{UInt64}
    ChachaSIMD(
        [
            0x243F6A88,
            0x85A308D3,
            0x13198A2E,
            0x03707344,
            0xA4093822,
            0x299F31D0,
            0x082EFA98,
            0xEC4E6C89,
        ],
        UInt64(0),
        4,
        T,
    )
end

function ChachaSIMD(
    key::Vector{UInt32} = [
        0x243F6A88,
        0x85A308D3,
        0x13198A2E,
        0x03707344,
        0xA4093822,
        0x299F31D0,
        0x082EFA98,
        0xEC4E6C89,
    ],
    skipLen::UInt64 = UInt64(0),
    DOUBLEROUNDS::Integer = 10,
    T = UInt64,
)
    r = ChachaSIMD{DOUBLEROUNDS,T}(
        Vector{NTuple{4,UInt32}}(undef, 4),
        Vector{NTuple{4,UInt32}}(undef, 4),
        Int8(1),
    )
    v = r.v
    v[1] = (0x61707865, 0x3320646e, 0x79622d32, 0x6b206574)
    v[2] = (key[1], key[2], key[3], key[4])
    v[3] = (key[5], key[6], key[7], key[8])
    v[4] = (skipLen & 0xffffffff, skipLen >> 32, 0, 0)
    raw(r)
    r
end

@inline @inbounds function round(x::Vector{NTuple{4,UInt32}})
    x[1] = x[1] .+ x[2]
    a = (x[4] .⊻ x[1])
    x[4] = (a .<< 16) .| (a .>> (32 - 16))
    x[3] = x[3] .+ x[4]
    a = x[2] .⊻ x[3]
    x[2] = (a .<< 12) .| (a .>> (32 - 12))
    x[1] = x[1] .+ x[2]
    a = x[4] .⊻ x[1]
    x[4] = (a .<< 8) .| (a .>> (32 - 8))
    x[3] = x[3] .+ x[4]
    a = (x[2] .⊻ x[3])
    x[2] = (a .<< 7) .| (a .>> (32 - 7))

end

@inline @inbounds function shuffle(x::Vector{NTuple{4,UInt32}})
    x[2] = (x[2][2], x[2][3], x[2][4], x[2][1])
    x[3] = (x[3][3], x[3][4], x[3][1], x[3][2])
    x[4] = (x[4][4], x[4][1], x[4][2], x[4][3])
end

@inline @inbounds function unshuffle(x::Vector{NTuple{4,UInt32}})
    x[2] = (x[2][4], x[2][1], x[2][2], x[2][3])
    x[3] = (x[3][3], x[3][4], x[3][1], x[3][2])
    x[4] = (x[4][2], x[4][3], x[4][4], x[4][1])
end

@inline function roundPair(x::Vector{NTuple{4,UInt32}})
    round(x)
    shuffle(x)
    round(x)
    unshuffle(x)
end

@inline function raw(r::ChachaSIMD{DOUBLEROUNDS}) where {DOUBLEROUNDS}
    x = r.x
    x .= r.v
    @inbounds @simd for i = 1:DOUBLEROUNDS
        roundPair(x)
    end
    @inbounds @simd for i = 1:4
        x[i] = x[i] .+ r.v[i]
    end
    @inbounds r.v[4] = (r.v[4][1] + 1, r.v[4][2], r.v[4][3], r.v[4][4])
    @inbounds if (r.v[4][1] == 0)
        r.v[4] = (r.v[4][1], r.v[4][2] + 1, r.v[4][3], r.v[4][4])
    end
end

@inline function skip(r::ChachaSIMD{DOUBLEROUNDS}, skipLength::UInt64) where {DOUBLEROUNDS}
    actualSkip = (skipLength + r.subCounter - 1) / 8
    r.v[4] = (
        r.v[4][1] - 1 + (actualSkip & 0xffffffff) % UInt32,
        (r.v[4][2] >> 32) % UInt32,
        r.v[4][3],
        r.v[4][4],
    )
    raw(r)
    r.subCounter = (((r.subCounter + (skipLength)) % 8) + 1) % Int8
end

@inline @inbounds function rand(r::ChachaSIMD{R,UInt32}, ::Type{UInt32}) where {R}
    if (r.subCounter == 17)
        r.subCounter = 1
        raw(r)
    end
    xIndex = ((r.subCounter - 1) >> 2) + 1
    tIndex = (r.subCounter - 1 - ((xIndex - 1) << 2)) + 1
    xvalue = r.x[xIndex]
    value = xvalue[tIndex]
    r.subCounter += 1
    value
end

@inline @inbounds function rand(r::ChachaSIMD{R,UInt64}, ::Type{UInt32}) where {R}
    if (r.subCounter == 17)
        r.subCounter = 1
        raw(r)
    end
    xIndex = ((r.subCounter - 1) >> 2) + 1
    tIndex = (r.subCounter - 1 - ((xIndex - 1) << 2)) + 1
    xvalue = r.x[xIndex]
    value = xvalue[tIndex]
    r.subCounter += 2 #skip one int32 for compatibility with 64-bit
    value
end

@inline function rand(r::ChachaSIMD{R}, ::Type{UInt64}) where {R}
    if (r.subCounter == 17)
        r.subCounter = 1
        raw(r)
    end
    xIndex = ((r.subCounter - 1) >> 2) + 1
    tIndex = (r.subCounter - 1 - ((xIndex - 1) << 2)) + 1
    @inbounds xvalue = r.x[xIndex]
    @inbounds value = (UInt64(xvalue[tIndex]) << 32) | xvalue[tIndex+1]
    r.subCounter += 2
    value
end

@inline function rand(r::ChachaSIMD, ::Type{Float64})
    if (r.subCounter == 17)
        r.subCounter = 1
        raw(r)
    end
    xIndex = ((r.subCounter - 1) >> 2) + 1
    tIndex = (r.subCounter - 1 - ((xIndex - 1) << 2)) + 1
    @inbounds xvalue = r.x[xIndex]
    @inbounds value = (
        xvalue[tIndex] / 0x1p32 +
        (0.5 + 1.1102230246251565e-16) +
        (xvalue[tIndex+1] & UInt32(0x000FFFFF)) / 0x1p52
    )
    r.subCounter += 2
    value
      # (Float64(rand(r, UInt64) >> 12) + 0.5) /  0x1p52
end


using Pkg; Pkg.activate(@__DIR__)
# add Lyceum registry: add registry https://github.com/Lyceum/LyceumRegistry.git
using MuJoCo
mj_activate("/home/taylor/.mujoco/bin/mjkey.txt") # set location to MuJoCo key path

using LyceumMuJoCo, LyceumMuJoCoViz

# ## load MuJoCo
path = joinpath(@__DIR__, "../../../environments/atlas/deps/atlas_v5_10Hz.xml")
# path = joinpath(@__DIR__, "../../../environments/atlas/deps/atlas_v5_50Hz.xml")
# path = joinpath(@__DIR__, "../../../environments/atlas/deps/atlas_v5_100Hz.xml")
# path = joinpath(@__DIR__, "../../../environments/atlas/deps/atlas_v5_250Hz.xml")
# path = joinpath(@__DIR__, "../../../environments/atlas/deps/atlas_v5_1000Hz.xml")

include("mujoco_model.jl")
atlas = MuJoCoModel(path)
sim = LyceumMuJoCo.MJSim(atlas.m, atlas.d)

T = 10 * 1
states = Array(undef, statespace(sim), T)
# sim.d.qpos = copy(sim.m.qpos0)
sim.d.qpos[3] += 0.5 # = zero(sim.d.qpos)
@show sim.d.qpos
sim.d.qvel .= 0.0
max_viol = Float64[]
for t = 1:T
    LyceumMuJoCo.step!(sim)
    states[:, t] .= getstate(sim)
    push!(max_viol, !isempty(sim.d.contact) ? minimum([c.dist for c in sim.d.contact]) : 0.0)
end

visualize(sim, trajectories=[states])

@show minimum(max_viol)


struct UnrollSymbols
    u₁loopsym::Symbol
    u₂loopsym::Symbol
    vectorized::Symbol
end
struct UnrollArgs{T <: Union{Nothing,Int}}
    u₁::Int
    u₁loopsym::Symbol
    u₂loopsym::Symbol
    vectorized::Symbol
    u₂max::Int
    suffix::T
end
function UnrollArgs(u₁::Int, unrollsyms::UnrollSymbols, u₂max::Int, suffix)
    @unpack u₁loopsym, u₂loopsym, vectorized = unrollsyms
    UnrollArgs(u₁, u₁loopsym, u₂loopsym, vectorized, u₂max, suffix)
end
function UnrollArgs(ua::UnrollArgs, u::Int)
    @unpack u₁loopsym, u₂loopsym, vectorized, u₂max, suffix = ua
    UnrollArgs(u, u₁loopsym, u₂loopsym, vectorized, u₂max, suffix)
end
# UnrollSymbols(ua::UnrollArgs) = UnrollSymbols(ua.u₁loopsym, ua.u₂loopsym, ua.vectorized)

# isfirst(ua::UnrollArgs{Nothing}) = iszero(ua.u₁)
# isfirst(ua::UnrollArgs{Int}) = iszero(ua.u₁) & iszero(ua.suffix)

struct UnrollSpecification
    u₁loopnum::Int
    u₂loopnum::Int
    vectorizedloopnum::Int
    u₁::Int
    u₂::Int
end
# UnrollSpecification(ls::LoopSet, u₁loop::Loop, vectorized::Symbol, u₁, u₂) = UnrollSpecification(ls, u₁loop.itersymbol, vectorized, u₁, u₂)
function UnrollSpecification(us::UnrollSpecification, u₁, u₂)
    @unpack u₁loopnum, u₂loopnum, vectorizedloopnum = us
    UnrollSpecification(u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂)
end
# function UnrollSpecification(us::UnrollSpecification; u₁ = us.u₁, u₂ = us.u₂)
#     @unpack u₁loopnum, u₂loopnum, vectorizedloopnum = us
#     UnrollSpecification(u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂)
# end
isunrolled1(us::UnrollSpecification, n::Int) = us.u₁loopnum == n
isunrolled2(us::UnrollSpecification, n::Int) = !isunrolled1(us, n) && us.u₂loopnum == n
isvectorized(us::UnrollSpecification, n::Int) = us.vectorizedloopnum == n
function unrollfactor(us::UnrollSpecification, n::Int)
    @unpack u₁loopnum, u₂loopnum, u₁, u₂ = us
    (u₁loopnum == n) ? u₁ : ((u₂loopnum == n) ? u₂ : 1)
end

struct Loop
    itersymbol::Symbol
    starthint::Int
    stophint::Int
    startsym::Symbol
    stopsym::Symbol
    startexact::Bool
    stopexact::Bool
end
startstopint(s::Int, start) = s
startstopint(s::Symbol, start) = start ? 1 : 1024
startstopsym(s::Int) = Symbol("##UNDEFINED##")
startstopsym(s::Symbol) = s
function Loop(itersymbol::Symbol, start::Union{Int,Symbol}, stop::Union{Int,Symbol})
    Loop(
        itersymbol, startstopint(start,true), startstopint(stop,false),
        startstopsym(start), startstopsym(stop), start isa Int, stop isa Int
    )
end
Base.length(loop::Loop) = 1 + loop.stophint - loop.starthint
isstaticloop(loop::Loop) = loop.startexact & loop.stopexact

function startloop(loop::Loop, itersymbol)
    startexact = loop.startexact
    if startexact
        Expr(:(=), itersymbol, loop.starthint - 1)
    else
        Expr(:(=), itersymbol, Expr(:call, lv(:staticm1), Expr(:call, lv(:unwrap), loop.startsym)))
    end
end
addexpr(ex, incr) = Expr(:call, lv(:vadd), ex, incr)
function addexpr(ex, incr::Number)
    if iszero(incr)
        incr
    elseif incr > 0
        Expr(:call, lv(:vadd), ex, incr)
    else
        Expr(:call, lv(:vsub), ex, -incr)
    end
end
addexpr(ex::Number, incr::Number) = ex + incr
subexpr(ex, incr) = Expr(:call, lv(:vsub), ex, incr)
subexpr(ex::Number, incr::Number) = ex - incr
subexpr(ex, incr::Number) = addexpr(ex,  -incr)

staticmulincr(ptr, incr) = Expr(:call, lv(:staticmul), Expr(:call, :eltype, ptr), incr)
callpointerforcomparison(sym) = Expr(:call, lv(:pointerforcomparison), sym)
function vec_looprange(loopmax, UF::Int, mangledname::Symbol, ptrcomp::Bool)
    if ptrcomp
        vec_looprange(loopmax, UF, callpointerforcomparison(mangledname), staticmulincr(mangledname, VECTORWIDTHSYMBOL))
    else
        vec_looprange(loopmax, UF, mangledname, VECTORWIDTHSYMBOL)
    end
end
function vec_looprange(loopmax, UF::Int, mangledname, W)
    incr = if isone(UF)
        Expr(:call, lv(:valsub), W, 1)
    else
        Expr(:call, lv(:valmulsub), W, UF, 1)
    end
    compexpr = subexpr(loopmax, incr)
    Expr(:call, :<, mangledname, compexpr)
end

function looprange(stopcon, incr::Int, mangledname)
    if iszero(incr)
        Expr(:call, :≤, mangledname, stopcon)
    else
        Expr(:call, :≤, mangledname, subexpr(stopcon, incr))
    end
end
function looprange(loop::Loop, incr::Int, mangledname)
    loop.stopexact ? looprange(loop.stophint, incr, mangledname) : looprange(loop.stopsym, incr, mangledname)
end
function terminatecondition(
    loop::Loop, us::UnrollSpecification, n::Int, mangledname::Symbol, inclmask::Bool, UF::Int = unrollfactor(us, n)
)
    if !isvectorized(us, n)
        looprange(loop, UF, mangledname)
    elseif inclmask
        looprange(loop, 1, mangledname)
    elseif loop.stopexact
        vec_looprange(loop.stophint, UF, mangledname, false) # may not be u₂loop
    else
        vec_looprange(loop.stopsym, UF, mangledname, false) # may not be u₂loop
    end
end
function incrementloopcounter(us::UnrollSpecification, n::Int, mangledname::Symbol, UF::Int = unrollfactor(us, n))
    if isvectorized(us, n)
        if isone(UF)
            Expr(:(=), mangledname, Expr(:call, lv(:valadd), VECTORWIDTHSYMBOL, mangledname))
        else
            Expr(:(=), mangledname, Expr(:call, lv(:valmuladd), VECTORWIDTHSYMBOL, UF, mangledname))
        end
    else
        Expr(:(=), mangledname, Expr(:call, lv(:vadd), mangledname, UF))
    end
end
function incrementloopcounter!(q, us::UnrollSpecification, n::Int, UF::Int = unrollfactor(us, n))
    if isvectorized(us, n)
        if isone(UF)
            push!(q.args, Expr(:call, lv(:Static), VECTORWIDTHSYMBOL))
        else
            push!(q.args, Expr(:call, lv(:valmul), VECTORWIDTHSYMBOL, Expr(:call, Expr(:curly, lv(:Static), UF))))
        end
    else
        push!(q.args, Expr(:call, Expr(:curly, lv(:Static), UF)))
    end
end
function looplengthexpr(loop::Loop)
    if loop.stopexact
        if loop.startexact
            length(loop)
        else
            Expr(:call, lv(:vsub), loop.stophint + 1, loop.startsym)
        end
    elseif loop.startexact
        if isone(loop.starthint)
            loop.stopsym
        else
            Expr(:call, lv(:vsub), loop.stopsym, loop.starthint - 1)
        end
    else
        Expr(:call, lv(:vsub), loop.stopsym, Expr(:call, lv(:staticm1), loop.startsym))
    end
end
# use_expect() = false
use_expect() = true
function looplengthexpr(loop, n)
    le = looplengthexpr(loop)
    # if false && use_expect() && isone(n) && !isstaticloop(loop)
    #     le = expect(le)
    #     push!(le.args, Expr(:call, Expr(:curly, :Val, length(loop))))
    # end
    le
end

# load/compute/store × isunrolled × istiled × pre/post loop × Loop number
struct LoopOrder <: AbstractArray{Vector{Operation},5}
    oporder::Vector{Vector{Operation}}
    loopnames::Vector{Symbol}
    bestorder::Vector{Symbol}
end
# function LoopOrder(N::Int)
#     LoopOrder(
#         [ Operation[] for _ ∈ 1:8N ],
#         Vector{Symbol}(undef, N), Vector{Symbol}(undef, N)
#     )
# end
LoopOrder() = LoopOrder(Vector{Operation}[],Symbol[],Symbol[])
Base.empty!(lo::LoopOrder) = foreach(empty!, lo.oporder)
function Base.resize!(lo::LoopOrder, N::Int)
    Nold = length(lo.loopnames)
    resize!(lo.oporder, 8N)
    for n ∈ 8Nold+1:8N
        lo.oporder[n] = Operation[]
    end
    resize!(lo.loopnames, N)
    resize!(lo.bestorder, N)
    lo
end
Base.size(lo::LoopOrder) = (2,2,2,length(lo.loopnames))
Base.@propagate_inbounds Base.getindex(lo::LoopOrder, i::Int) = lo.oporder[i]
Base.@propagate_inbounds Base.getindex(lo::LoopOrder, i...) = lo.oporder[LinearIndices(size(lo))[i...]]

@enum NumberType::Int8 HardInt HardFloat IntOrFloat INVALID

struct LoopStartStopManager
    terminators::Vector{Int}
    incrementedptrs::Vector{Vector{ArrayReferenceMeta}}
    uniquearrayrefs::Vector{ArrayReferenceMeta}
end
# Must make it easy to iterate
# outer_reductions is a vector of indices (within operation vectors) of the reduction operation, eg the vmuladd op in a dot product
# O(N) search is faster at small sizes
struct LoopSet
    loopsymbols::Vector{Symbol}
    loopsymbol_offsets::Vector{Int}  # symbol loopsymbols[i] corresponds to loops[lso[i]+1:lso[i+1]] (CartesianIndex handling)
    loops::Vector{Loop}
    opdict::Dict{Symbol,Operation}
    operations::Vector{Operation} # Split them to make it easier to iterate over just a subset
    operation_offsets::Vector{Int}
    outer_reductions::Vector{Int} # IDs of reduction operations that need to be reduced at end.
    loop_order::LoopOrder
    preamble::Expr
    prepreamble::Expr # performs extractions that must be performed first, and don't need further registering
    preamble_symsym::Vector{Tuple{Int,Symbol}}
    preamble_symint::Vector{Tuple{Int,Int}}
    preamble_symfloat::Vector{Tuple{Int,Float64}}
    preamble_zeros::Vector{Tuple{Int,NumberType}}
    preamble_funcofeltypes::Vector{Tuple{Int,Symbol}}
    includedarrays::Vector{Symbol}
    includedactualarrays::Vector{Symbol}
    syms_aliasing_refs::Vector{Symbol}
    refs_aliasing_syms::Vector{ArrayReferenceMeta}
    cost_vec::Matrix{Float64}
    reg_pres::Matrix{Float64}
    included_vars::Vector{Bool}
    place_after_loop::Vector{Bool}
    unrollspecification::Base.RefValue{UnrollSpecification}
    loadelimination::Base.RefValue{Bool}
    lssm::Base.RefValue{LoopStartStopManager}
    vector_width::Base.RefValue{Int}
    isbroadcast::Base.RefValue{Bool}
    mod::Symbol
end

function cost_vec_buf(ls::LoopSet)
    cv = @view(ls.cost_vec[:,2])
    @inbounds for i ∈ 1:4
        cv[i] = 0.0
    end
    cv
end
function reg_pres_buf(ls::LoopSet)
    ps = @view(ls.reg_pres[:,2])
    @inbounds for i ∈ 1:5
        ps[i] = 0
    end
    ps
end
function save_tilecost!(ls::LoopSet)
    @inbounds for i ∈ 1:4
        ls.cost_vec[i,1] = ls.cost_vec[i,2]
        ls.reg_pres[i,1] = ls.reg_pres[i,2]
    end
    ls.reg_pres[5,1] = ls.reg_pres[5,2]
end

pushprepreamble!(ls::LoopSet, ex) = push!(ls.prepreamble.args, ex)
function pushpreamble!(ls::LoopSet, op::Operation, v::Symbol)
    if v !== mangledvar(op)
        push!(ls.preamble_symsym, (identifier(op),v))
    end
    nothing
end
function pushpreamble!(ls::LoopSet, op::Operation, v::Number)
    typ = v isa Integer ? HardInt : HardFloat
    id = identifier(op)
    if iszero(v)
        push!(ls.preamble_zeros, (id, typ))
    elseif isone(v)
        push!(ls.preamble_funcofeltypes, (id, :one))
    elseif v isa Integer
        push!(ls.preamble_symint, (id, convert(Int,v)))
    else
        push!(ls.preamble_symfloat, (id, convert(Float64,v)))
    end
end
pushpreamble!(ls::LoopSet, ex::Expr) = push!(ls.preamble.args, ex)
# function pushpreamble!(ls::LoopSet, op::Operation, RHS::Expr)
#     c = gensym(:licmconst)
#     if RHS.head === :call && first(RHS.args) === :zero
#         push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
#     elseif RHS.head === :call && first(RHS.args) === :one
#         push!(ls.preamble_funcofeltypes, (identifier(op), :one))
#     else
#         pushpreamble!(ls, Expr(:(=), c, RHS))
#         pushpreamble!(ls, op, c)
#     end
#     nothing
# end
function zerotype(ls::LoopSet, op::Operation)
    opid = identifier(op)
    for (id,typ) ∈ ls.preamble_zeros
        id == opid && return typ
    end
    INVALID
end

includesarray(ls::LoopSet, array::Symbol) = array ∈ ls.includedarrays

function LoopSet(mod::Symbol)
    LoopSet(
        Symbol[], [0], Loop[],
        Dict{Symbol,Operation}(),
        Operation[], [0],
        Int[],
        LoopOrder(),
        Expr(:block),Expr(:block),
        Tuple{Int,Symbol}[],
        Tuple{Int,Int}[],
        Tuple{Int,Float64}[],
        Int[],Int[],
        Symbol[], Symbol[], Symbol[],
        ArrayReferenceMeta[],
        Matrix{Float64}(undef, 4, 2),
        Matrix{Float64}(undef, 5, 2),
        Bool[], Bool[], Ref{UnrollSpecification}(),
        Ref(false), Ref{LoopStartStopManager}(),
        Ref(0), Ref(false), mod
    )
end

cacheunrolled!(ls::LoopSet, u₁loop, u₂loop, vectorized) = foreach(op -> setunrolled!(op, u₁loop, u₂loop, vectorized), operations(ls))

num_loops(ls::LoopSet) = length(ls.loops)
function oporder(ls::LoopSet)
    N = length(ls.loop_order.loopnames)
    reshape(ls.loop_order.oporder, (2,2,2,N))
end
names(ls::LoopSet) = ls.loop_order.loopnames
reversenames(ls::LoopSet) = ls.loop_order.bestorder
function getloopid_or_nothing(ls::LoopSet, s::Symbol)
    for (loopnum,sym) ∈ enumerate(ls.loopsymbols)
        s === sym && return loopnum
    end
end
getloopid(ls::LoopSet, s::Symbol) = getloopid_or_nothing(ls, s)::Int
getloop(ls::LoopSet, s::Symbol) = ls.loops[getloopid(ls, s)]
# getloop(ls::LoopSet, i::Integer) = ls.loops[i]
getloopsym(ls::LoopSet, i::Integer) = ls.loopsymbols[i]
Base.length(ls::LoopSet, s::Symbol) = length(getloop(ls, s))

# isstaticloop(ls::LoopSet, s::Symbol) = isstaticloop(getloop(ls,s))
# looprangehint(ls::LoopSet, s::Symbol) = length(getloop(ls, s))
# looprangesym(ls::LoopSet, s::Symbol) = getloop(ls, s).rangesym

"""
getop only works while construction a LoopSet object. You cannot use it while lowering.
"""
getop(ls::LoopSet, var::Number, elementbytes) = add_constant!(ls, var, elementbytes)
function getop(ls::LoopSet, var::Symbol, elementbytes::Int)
    get!(ls.opdict, var) do
        add_constant!(ls, var, elementbytes)
    end
end
function getop(ls::LoopSet, var::Symbol, deps, elementbytes::Int)
    get!(ls.opdict, var) do
        add_constant!(ls, var, deps, gensym(:constant), elementbytes)
    end
end
getop(ls::LoopSet, i::Int) = ls.operations[i]

# """
# Returns an operation with the same name as `s`.
# """
# function getoperation(ls::LoopSet, s::Symbol)
#     for op ∈ Iterators.Reverse(operations(ls))
#         name(op) === s && return op
#     end
#     throw("Symbol $s not found among operations(ls).")
# end

function Operation(
    ls::LoopSet, variable, elementbytes, instruction,
    node_type, dependencies, reduced_deps, parents, ref = NOTAREFERENCE
)
    Operation(
        length(operations(ls)), variable, elementbytes, instruction,
        node_type, dependencies, reduced_deps, parents, ref
    )
end
function Operation(ls::LoopSet, variable, elementbytes, instr, optype, mpref::ArrayReferenceMetaPosition)
    Operation(length(operations(ls)), variable, elementbytes, instr, optype, mpref)
end

operations(ls::LoopSet) = ls.operations
function pushop!(ls::LoopSet, op::Operation, var::Symbol = name(op))
    for opp ∈ operations(ls)
        if matches(op, opp)
            ls.opdict[var] = opp
            return opp
        end
    end
    push!(ls.operations, op)
    ls.opdict[var] = op
    op
end
function add_block!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
    for x ∈ ex.args
        x isa Expr || continue # be that general?
        x.head === :inbounds && continue
        push!(ls, x, elementbytes, position)
    end
end
function maybestatic!(expr::Expr)
    if expr.head === :call
        f = first(expr.args)
        if f === :length
            expr.args[1] = lv(:maybestaticlength)
        elseif f === :size && length(expr.args) == 3
            i = expr.args[3]
            if i isa Integer
                expr.args[1] = lv(:maybestaticsize)
                expr.args[3] = Expr(:call, Expr(:curly, :Val, convert(Int, i)))
            end
        end
    end
    expr
end
function add_loop_bound!(ls::LoopSet, itersym::Symbol, bound, upper::Bool = true)
    (bound isa Symbol && upper) && return bound
    bound isa Expr && maybestatic!(bound)
    N = gensym(string(itersym) * (upper ? "_loop_upper_bound" : "_loop_lower_bound"))
    pushprepreamble!(ls, Expr(:(=), N, bound))
    N
end

"""
This function creates a loop, while switching from 1 to 0 based indices
"""
function register_single_loop!(ls::LoopSet, looprange::Expr)
    itersym = (looprange.args[1])::Symbol
    r = looprange.args[2]
    if isexpr(r, :call)
        f = first(r.args)
        loop::Loop = if f === :(:)
            lower = r.args[2]
            upper = r.args[3]
            lii::Bool = lower isa Integer
            liiv::Int = lii ? convert(Int, lower) : 1
            uii::Bool = upper isa Integer
            if lii & uii # both are integers
                Loop(itersym, liiv, convert(Int, upper))
            elseif lii # only lower bound is an integer
                if upper isa Symbol
                    Loop(itersym, liiv, upper)
                elseif upper isa Expr
                    Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
                else
                    Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
                end
            elseif uii # only upper bound is an integer
                uiiv = convert(Int, upper)
                Loop(itersym, add_loop_bound!(ls, itersym, lower, false), uiiv)
            else # neither are integers
                L = add_loop_bound!(ls, itersym, lower, false)
                U = add_loop_bound!(ls, itersym, upper, true)
                Loop(itersym, L, U)
            end
        elseif f === :OneTo || isscopedname(f, :Base, :OneTo)
            otN = r.args[2]
            if otN isa Integer
                Loop(itersym, 1, otN)
            else
                otN isa Expr && maybestatic!(otN)
                N = gensym("loop" * string(itersym))
                pushprepreamble!(ls, Expr(:(=), N, otN))
                Loop(itersym, 1, N)
            end
        else
            N = gensym("loop" * string(itersym))
            pushprepreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
            L = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticfirst), N), false)
            U = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticlast), N), true)
            Loop(itersym, L, U)
        end
    elseif isa(r, Symbol)
        # Treat similar to `eachindex`
        N = gensym("loop" * string(itersym))
        pushprepreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
        L = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticfirst), N), false)
        U = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticlast), N), true)
        loop = Loop(itersym, L, U)
    else
        throw("Unrecognized loop range type: $r.")
    end
    add_loop!(ls, loop, itersym)
    nothing
end
function register_loop!(ls::LoopSet, looprange::Expr)
    if looprange.head === :block # multiple loops
        for lr ∈ looprange.args
            register_single_loop!(ls, lr)
        end
    else
        @assert looprange.head === :(=)
        register_single_loop!(ls, looprange)
    end
end
function add_loop!(ls::LoopSet, q::Expr, elementbytes::Int)
    register_loop!(ls, q.args[1])
    body = q.args[2]
    position = length(ls.loopsymbols)
    if body.head === :block
        add_block!(ls, body, elementbytes, position)
    else
        push!(ls, q, elementbytes, position)
    end
end
function add_loop!(ls::LoopSet, loop::Loop, itersym::Symbol = loop.itersymbol)
    push!(ls.loopsymbols, itersym)
    push!(ls.loops, loop)
    nothing
end

# function instruction(x)
#     x isa Symbol ? x : last(x.args).value
# end
# instruction(ls::LoopSet, f::Symbol) = instruction!(ls, f)
function instruction!(ls::LoopSet, x::Expr)
    x isa Symbol && return x
    instr = last(x.args).value
    if instr ∉ keys(COST)
        instr = gensym(:f)
        pushpreamble!(ls, Expr(:(=), instr, x))
    end
    Instruction(Symbol(""), instr)
end
instruction!(ls::LoopSet, x::Symbol) = instruction(x)

function maybe_const_compute!(ls::LoopSet, LHS::Symbol, op::Operation, elementbytes::Int, position::Int)
    # return op
    if iscompute(op) && iszero(length(loopdependencies(op)))
        ls.opdict[LHS] = add_constant!(ls, LHS, ls.loopsymbols[1:position], gensym(instruction(op).instr), elementbytes, :numericconstant)
    else
        op
    end
end
function strip_op_linenumber_nodes(q::Expr)
    filtered = filter(x -> !isa(x, LineNumberNode), q.args)
    if VERSION ≥ v"1.4"
        only(filtered)
    else
        @assert isone(length(filtered))
        first(filtered)
    end
end

function add_operation!(ls::LoopSet, LHS::Symbol, RHS::Symbol, elementbytes::Int, position::Int)
    add_constant!(ls, RHS, ls.loopsymbols[1:position], LHS, elementbytes)
end
function add_operation!(
    ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int, position::Int
)
    if RHS.head === :ref
        add_load_ref!(ls, LHS, RHS, elementbytes)
    elseif RHS.head === :call
        f = first(RHS.args)
        if f === :getindex
            add_load_getindex!(ls, LHS, RHS, elementbytes)
        elseif f === :zero || f === :one
            c = gensym(f)
            op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS, elementbytes, :numericconstant)
            if f === :zero
                push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
            else
                push!(ls.preamble_funcofeltypes, (identifier(op), :one))
            end
            op
        else
            # maybe_const_compute!(ls, add_compute!(ls, LHS, RHS, elementbytes, position), elementbytes, position)
            add_compute!(ls, LHS, RHS, elementbytes, position)
        end
    elseif RHS.head === :if
        add_if!(ls, LHS, RHS, elementbytes, position)
    elseif RHS.head === :block
        add_operation!(ls, LHS, strip_op_linenumber_nodes(RHS), elementbytes, position)
    else
        println(RHS)
        throw("Expression not recognized.")
    end
end
add_operation!(ls::LoopSet, RHS::Expr, elementbytes::Int, position::Int) = add_operation!(ls, gensym(:LHS), RHS, elementbytes, position)
function add_operation!(
    ls::LoopSet, LHS_sym::Symbol, RHS::Expr, LHS_ref::ArrayReferenceMetaPosition, elementbytes::Int, position::Int
)
    if RHS.head === :ref# || (RHS.head === :call && first(RHS.args) === :getindex)
        array, rawindices = ref_from_expr!(ls, RHS)
        RHS_ref = array_reference_meta!(ls, array, rawindices, elementbytes, gensym(LHS_sym))
        op = add_load!(ls, RHS_ref, elementbytes)
        iop = add_compute!(ls, LHS_sym, :identity, [op], elementbytes)
        # pushfirst!(LHS_ref.parents, iop)
    elseif RHS.head === :call
        f = first(RHS.args)
        if f === :getindex
            add_load!(ls, LHS_sym, LHS_ref, elementbytes)
        elseif f === :zero || f === :one
            c = gensym(f)
            op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS_sym, elementbytes, :numericconstant)
            # op = add_constant!(ls, c, Symbol[], LHS_sym, elementbytes, :numericconstant)
            if f === :zero
                push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
            else
                push!(ls.preamble_funcofeltypes, (identifier(op), :one))
            end
            op
        else
            add_compute!(ls, LHS_sym, RHS, elementbytes, position, LHS_ref)
        end
    elseif RHS.head === :if
        add_if!(ls, LHS_sym, RHS, elementbytes, position, LHS_ref)
    elseif RHS.head === :block
        add_operation!(ls, LHS, strip_op_linenumber_nodes(RHS), elementbytes, position)
    else
        println(RHS)
        throw("Expression not recognized.")
    end
end

function prepare_rhs_for_storage!(ls::LoopSet, RHS::Symbol, array, rawindices, elementbytes::Int, position::Int)
    add_store!(ls, RHS, array, rawindices, elementbytes)
end
function prepare_rhs_for_storage!(ls::LoopSet, RHS::Expr, array, rawindices, elementbytes::Int, position::Int)
    mpref = array_reference_meta!(ls, array, rawindices, elementbytes)
    cachedparents = copy(mpref.parents)
    ref = mpref.mref.ref
    # id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
    # lrhs = id === nothing ? gensym(:RHS) : ls.syms_aliasing_refs[id]
    lrhs = gensym(:RHS)
    mpref.varname = lrhs
    add_operation!(ls, lrhs, RHS, mpref, elementbytes, position)
    mpref.parents = cachedparents
    add_store!(ls, mpref, elementbytes)
end

function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
    if ex.head === :call
        finex = first(ex.args)::Symbol
        if finex === :setindex!
            array, rawindices = ref_from_setindex!(ls, ex)
            prepare_rhs_for_storage!(ls, ex.args[3], array, rawindices, elementbytes, position)
        else
            throw("Function $finex not recognized.")
        end
    elseif ex.head === :(=)
        LHS = ex.args[1]
        RHS = ex.args[2]
        if LHS isa Symbol
            if RHS isa Expr
                maybe_const_compute!(ls, LHS, add_operation!(ls, LHS, RHS, elementbytes, position), elementbytes, position)
            else
                add_constant!(ls, RHS, ls.loopsymbols[1:position], LHS, elementbytes)
            end
        elseif LHS isa Expr
            if LHS.head === :ref
                if RHS isa Symbol
                    add_store_ref!(ls, RHS, LHS, elementbytes)
                elseif RHS isa Expr
                    # need to check if LHS appears in RHS
                    # assign RHS to lrhs
                    array, rawindices = ref_from_expr!(ls, LHS)
                    prepare_rhs_for_storage!(ls, RHS, array, rawindices, elementbytes, position)
                else
                    add_store_ref!(ls, RHS, LHS, elementbytes)
                end
            elseif LHS.head === :tuple
                @assert length(LHS.args) ≤ 9 "Functions returning more than 9 values aren't currently supported."
                lhstemp = gensym(:lhstuple)
                vparents = Operation[maybe_const_compute!(ls, lhstemp, add_operation!(ls, lhstemp, RHS, elementbytes, position), elementbytes, position)]
                for i ∈ eachindex(LHS.args)
                    f = (:first,:second,:third,:fourth,:fifth,:sixth,:seventh,:eighth,:ninth)[i]
                    lhsi = LHS.args[i]
                    if lhsi isa Symbol
                        add_compute!(ls, lhsi, f, vparents, elementbytes)
                    elseif lhsi isa Expr && lhsi.head === :ref
                        tempunpacksym = gensym(:tempunpack)
                        add_compute!(ls, tempunpacksym, f, vparents, elementbytes)
                        add_store_ref!(ls, tempunpacksym, lhsi, elementbytes)
                    else
                        println(lhsi)
                        throw("Unpacking the above expression in the left hand side was not understood/supported.")
                    end
                end
                first(vparents)
            else
                println(LHS)
                throw("LHS not understood; only `:ref`s and `:tuple`s are currently supported.")
            end
        else
            println(LHS)
            throw("LHS not understood.")
        end
    elseif ex.head === :block
        add_block!(ls, ex, elementbytes, position)
    elseif ex.head === :for
        add_loop!(ls, ex, elementbytes)
    elseif ex.head === :&&
        add_andblock!(ls, ex, elementbytes, position)
    elseif ex.head === :||
        add_orblock!(ls, ex, elementbytes, position)
    elseif ex.head === :local # Handle locals introduced by `@inbounds`; using `local` with `@avx` is not recomended (nor is `@inbounds`; which applies automatically regardless)
        @assert length(ex.args) == 1 # TODO replace assert + first with "only" once support for Julia < 1.4 is dropped
        localbody = first(ex.args)
        @assert localbody.head === :(=)
        @assert length(localbody.args) == 2
        LHS = (localbody.args[1])::Symbol
        RHS = push!(ls, (localbody.args[2]), elementbytes, position)
        if isstore(RHS)
            RHS
        else
            add_compute!(ls, LHS, :identity, [RHS], elementbytes)
        end
    else
        println(ex)
        throw("Don't know how to handle expression.")
    end
end

function UnrollSpecification(ls::LoopSet, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol, u₁, u₂)
    order = names(ls)
    nu₁ = findfirst(isequal(u₁loop), order)::Int
    nu₂ = u₂ == -1 ? nu₁ : findfirst(isequal(u₂loop), order)::Int
    nv = findfirst(isequal(vectorized), order)::Int
    UnrollSpecification(nu₁, nu₂, nv, u₁, u₂)
end

"""
  looplengthprod(ls::LoopSet)

Convert to `Float64` for the sake of non-64 bit platforms.
"""
looplengthprod(ls::LoopSet) = prod(Float64 ∘ length, ls.loops)

function looplength(ls::LoopSet, s::Symbol)
    # search_tree(parents(operations(ls)[i]), name(op)) && return true
    id = getloopid_or_nothing(ls, s)
    if isnothing(id)
        l = 0.0
        # TODO: we could double count a loop.
        for op ∈ operations(ls)
            name(op) === s || continue
            for opp ∈ parents(op)
                if isloopvalue(opp)
                    oppname = first(loopdependencies(opp))
                    l += looplength(ls, oppname)
                elseif iscompute(opp)
                    oppname = name(opp)
                    l += looplength(ls, oppname)
                    # TODO elseif isconstant(opp)
                end
            end
            l += 1 - length(parents(op))
        end
        l
    else
        Float64(length(ls.loops[id]))
    end
end

# function getunrolled(ls::LoopSet)
#     order = names(ls)
#     us = ls.unrollspecification[]
#     @unpack u₁loopnum, u₂loopnum = us
#     order[u₁loopnum], order[u₂loopnum]
# end


module MaximumGeneratorProfitTCRDD

#  Packages
using Cbc
using D3TypeTrees
using DataFrames
using Dates
using DataStructures
using GLPK
using Ipopt
using JuMP
using LinearAlgebra
using Plots
using PowerSimulations
using PowerSystems
using TimeSeries

#  Package Constants
const PSY = PowerSystems
const PSI = PowerSimulations

#  Functions
# Power System Auxiliary Functions for TCRDD Algorithm
include("powersystemaux/assign_lines_index.jl")
include("powersystemaux/change_slack!.jl")
include("powersystemaux/get_binding_lines.jl")
#include("powersystemaux/get_Pg_thermal_bindingPg.jl")
#include("powersystemaux/get_Pg_thermal_constprice.jl")
#include("powersystemaux/get_Pg_thermal_nonzeroslope.jl")
#include("powersystemaux/get_Pg_thermal_slack.jl")
include("powersystemaux/get_PTDF_bindingLines.jl")
include("powersystemaux/get_PTDF_load.jl")
include("powersystemaux/get_PTDF_thermal_bindingPg.jl")
include("powersystemaux/get_PTDF_thermal_constprice.jl")
include("powersystemaux/get_PTDF_thermal_nonzeroslope.jl")
include("powersystemaux/get_PTDF_thermal_slack.jl")
#include("powersystemaux/get_thermal_bindingPg.jl")
#include("powersystemaux/get_thermal_constprice.jl")
#include("powersystemaux/get_thermal_nonzeroslope.jl")
include("powersystemaux/get_thermal_slack.jl")
include("powersystemaux/opf_PTDF.jl")

# Functions for TCRDD Algorithm

include("algorithm/bid_argmax_profit.jl")
include("algorithm/bisection_loop!.jl")
include("algorithm/compare_tan_inv_RDC.jl")
include("algorithm/d2f_PgPg_nonzeroslope.jl")
include("algorithm/disp_plots.jl")
include("algorithm/f_TCRDD.jl")
include("algorithm/residual_demand.jl")

# Main Functions
include("maxGenProfit_tcrdd.jl")
include("plots_TCRDD.jl")

# Type

# Accessor functions

# Useful Calculation Functions
#export change_slack!

# Main Functions
export maxGenProfit_tcrdd
export plots_TCRDD

end


__precompile__()

module Xpress

    using Libdl

    # Load in `deps.jl`, complaining if it does not exist
    const depsjl_path = joinpath(@__DIR__, "..", "deps", "deps.jl")

    if !isfile(depsjl_path)
        error("XPRESS cannot be loaded. Please run Pkg.build(\"Xpress\").")
    end

    include(depsjl_path)

    ### imports

    import Base.show, Base.copy

    # ### exports
    # export

    # ## xprs_env
    # free_env,

    # ## xprs_params
    # getparam, setparam!, setparams!,

    # ## xprs_model
    # set_sense!, update_model!, reset_model!,
    # read_model, write_model,
    # copy,

    # ## xprs_attrs
    # model_name, model_sense, model_type,
    # num_vars, num_constrs, num_sos, num_qconstrs,
    # num_cnzs, num_qnzs, num_qcnzs,
    # is_qp, is_qcp, is_mip,

    # lowerbounds, upperbounds, objcoeffs, set_objcoeffs!,

    # ## xprs_vars
    # XPRS_CONTINUOUS, XPRS_BINARY, XPRS_INTEGER,
    # add_var!, add_vars!, add_cvar!, add_cvars!,
    # add_bvar!, add_bvars!, add_ivar!, add_ivars!,
    # del_vars!,

    # ## xprs_constrs
    # add_constr!, add_constrs!, add_constrs_t!,
    # add_rangeconstr!, add_rangeconstrs!, add_rangeconstrs_t!,
    # get_constrmatrix, add_sos!, del_constrs!, chg_coeffs!,

    # ## xprs_quad
    # add_qpterms!, add_qconstr!,

    # ## higher level
    # xpress_model,

    # ## xprs_solve
    # optimize, computeIIS, get_solution,
    # get_status, OptimInfo, get_optiminfo, get_objval

    ### include source files
    module Lib

        import ..Xpress
        const libxprs = Xpress.libxprs

        include("ctypes.jl")
        include("common.jl")
        include("lib.jl")

        # include("custom.jl")

    end

    include("utils.jl")

    include("helper.jl")

    include("api.jl")

    # # license checker
    include("license.jl")

    function initialize()
        userlic()
        init() # Call XPRSinit for initialization
        atexit(free) # Call free when process terminates
    end
    # include("xprs_full_defines.jl")

    # include("xprs_common.jl")
    # include("xprs_env.jl")

    # include("xprs_model.jl")
    # include("xprs_params.jl")
    # include("xprs_attrs.jl")

    # include("xprs_vars.jl")

    # include("xprs_constrs.jl")
    # include("xprs_quad.jl")
    # include("xprs_highlevel.jl")

    # include("xprs_solve.jl")
    # include("xprs_callbacks.jl")
    # include("xprs_iis.jl")

    # include("XpressSolverInterface.jl")
    include("MOI/MOI_Wrapper.jl")

end


export StripTod, split_into_n, allocate_tod
export pwrq1, pwrq2, pwru1, pwru2
export demq1, demq2, demu1, demu2
export PWR_Q1_RANK, PWR_Q2_RANK, PWR_U1_RANK, PWR_U2_RANK
export DEM_Q1_RANK, DEM_Q2_RANK, DEM_U1_RANK, DEM_U2_RANK

import Printf: @printf
using AstroTime

"""
    split_into_n(length, num_of_segments)

Given the `length` of the array, the function returns a split of the
elements in `num_of_segments` subsets whose length is roughly the
same.

The result is an array containing the number of elements which must go
in each section; all the elements in the array have the property of
being similar and their sum being equal to `length`.

# Examples
```julia-repl
julia> split = split_into_n(20, 3)
3-element Array{Int64,1}:
6
7
7

julia> sum(split)
20
```
"""
function split_into_n(length, num_of_segments)
    @assert num_of_segments > 0
    @assert length >= num_of_segments
    start_pos = zeros(Int, num_of_segments+1)

    for i in 1:num_of_segments+1
        start_pos[i] =  floor(((i-1)*length/num_of_segments))
    end

    return start_pos[2:end]-start_pos[1:end-1]
end

@doc raw"""
A matrix of time-ordered data for a subset of polarimeters

The fields of this structure are the following:

- `time_range`: a range representing the time of each sample in the
  TOD.

- `samples`: a N×8×m matrix, where N is the number of samples in the
  timeline, and m is the number of polarimeters. The middle dimension
  spans the three Stokes parameters `I`, `Q`, and `U` (in this order).
  The value `N` is equal to `length(time_range)`

"""
struct StripTod{T <: Real, S}
    polarimeters::Any
    time_range::AbstractRange{S}
    samples::Array{T,3}
end

function Base.show(io::IO, tod::StripTod{T, S}) where {T, S}
    println(io, string(
        "TOD($(length(tod.polarimeters)) polarimeters, ",
        "$(length(tod.time_range)) samples/polarimeter, ",
        "time range from $(minimum(tod.time_range)) to ",
        "$(maximum(tod.time_range)))",
    ))
end

const PWR_Q1_RANK = 1
const PWR_Q2_RANK = 2
const PWR_U1_RANK = 3
const PWR_U2_RANK = 4
const DEM_Q1_RANK = 5
const DEM_Q2_RANK = 6
const DEM_U1_RANK = 7
const DEM_U2_RANK = 8

for (symb, ch_name, value) in [
    (:PWR_Q1_RANK, "PWR_Q1", PWR_Q1_RANK),
    (:PWR_Q2_RANK, "PWR_Q2", PWR_Q2_RANK),
    (:PWR_U1_RANK, "PWR_U1", PWR_U1_RANK),
    (:PWR_U2_RANK, "PWR_U2", PWR_U2_RANK),
    (:DEM_Q1_RANK, "DEM_Q1", DEM_Q1_RANK),
    (:DEM_Q2_RANK, "DEM_Q2", DEM_Q2_RANK),
    (:DEM_U1_RANK, "DEM_U1", DEM_U1_RANK),
    (:DEM_U2_RANK, "DEM_U2", DEM_U2_RANK),
]
    @eval begin
        @doc """
    $($ch_name)_RANK = $($value)

The rank of `$($ch_name)` in the `samples` field of a [`StripTod`](@ref) variable.

"""
        $symb
    end
end

for (index, ch_name, fn_name) in [
    (PWR_Q1_RANK, "PWR_Q1", :pwrq1),
    (PWR_Q2_RANK, "PWR_Q2", :pwrq2),
    (PWR_U1_RANK, "PWR_U1", :pwru1),
    (PWR_U2_RANK, "PWR_U2", :pwru2),
    (DEM_Q1_RANK, "DEM_Q1", :demq1),
    (DEM_Q2_RANK, "DEM_Q2", :demq2),
    (DEM_U1_RANK, "DEM_U1", :demu1),
    (DEM_U2_RANK, "DEM_U2", :demu2),
]
    @eval begin
        function ($fn_name)(t::StripTod, pol::T) where T <: Integer
            @view t.samples[:, $index, pol]
        end

        @doc """
    $($fn_name)(tod::Stripeline.StripTod, pol::T) where T <: Integer

Return a view of the timeline for channel `$($ch_name)` in the
time-ordered data for the polarimeter indexed by `pol` in `tod`. It is
a shortcut for

    tod.samples[:, $($ch_name)_RANK, pol]

"""
        $fn_name
    end
end

@doc raw"""
    allocate_tod(t, time_range, polarimeters; mpi_rank=0, mpi_size=1, zero_tod=true)

Allocate memory for a time-ordered data matrix appropriate for the set
of polarimeters in the variable `polarimeters` (an array or range,
only its length is used and it is copied as-is in the returned
[`StripTod`](@ref) value). The type of each sample is set to `t`. The
variable `time_range` must be a range of floating-point values or
epochs (as defined by the
[AstroTime](https://juliaastro.github.io/AstroTime.jl/stable/)
package) which represent the full set of samples acquired during an
observation.

The optional parameters `mpi_rank` and `mpi_size` can be used in MPI
scripts to allocate only the subset of samples that must be processed
by the current MPI process.

If `zero_tod` is true, each sample in the TOD is set to zero;
otherwise, its value is undefined. Using `zero_tod = false` is
slightly faster, but it should be used only if you are going to set by
hand every sample in the TOD once `allocate_tod` returns.

Return a [`StripTod`](@ref) value.

### Example ###

```jldoctest; setup = :(using Stripeline)
julia> tod = allocate_tod(Float32, 0.0:0.01:100.0, 1:5)
TOD(5 polarimeters, 10001 samples/polarimeter, time range from 0.0 to 100.0)
```

"""
function allocate_tod(
    t::Type{T},
    time_range::R,
    polarimeters;
    mpi_rank = 0,
    mpi_size = 1,
    zero_tod = true,
) where {S, V, R <: AbstractRange, T <: Real}

    rank_idx = mpi_rank + 1 # Same as rank, but it starts from 1

    nsamples_arr = split_into_n(length(time_range), mpi_size)
    nsamples = nsamples_arr[rank_idx]

    first_idx = sum(nsamples_arr[1:(rank_idx - 1)]) + 1
    last_idx = first_idx + nsamples - 1

    # Determine the time range for the current rank
    cur_range = time_range[first_idx]:step(time_range):time_range[last_idx]

    # Size of the field `samples`
    matrix_size = (length(cur_range), 8, length(polarimeters))

    StripTod(
        polarimeters,
        cur_range,
        zero_tod ? zeros(T, matrix_size) : Array{T}(undef, matrix_size),
    )
end


using Rheologies ; const R = Rheologies
using JuAFEM ; const J = JuAFEM
using Test
using SafeTestsets

@testset "variables interpolations" begin
    @test R.get_variables_interpolation((:u,), (2,), Cell{2,3,3}) ==
        Lagrange{2,RefTetrahedron,2}()
    @test R.get_variables_interpolation((:u,), (1,), Cell{2,3,3}) ==
        Lagrange{2,RefTetrahedron,1}()
    @test R.get_variables_interpolation((:u,:p), (2,1), Cell{2,3,3}) ==
        (Lagrange{2,RefTetrahedron,2}(), Lagrange{2,RefTetrahedron,1}())
    @test R.get_variables_interpolation((:u,:p), (2,1), Cell{2,4,4}) ==
        (Lagrange{2,RefCube,2}(), Lagrange{2,RefCube,1}())
end

@testset "get quad rules" begin
    @test typeof(R.get_quadrature_rules(2, :legendre, Cell{2,4,4})[1]) ==
        typeof(QuadratureRule{2,RefCube}(2))
    @test typeof(R.get_quadrature_rules(2, :legendre, Cell{2,4,4})[2]) ==
        typeof(QuadratureRule{1,RefCube}(2))
end

@testset "create values" begin
    el_geom = Cell{2,3,3}
    qr, qr_face = R.get_quadrature_rules(3, :legendre, el_geom)
    interp_u1 = R.get_variables_interpolation((:u,), (2,), el_geom)
    interp_u2 = R.get_variables_interpolation((:u,:p), (2,1), el_geom)
    interp_geom = J.default_interpolation(el_geom)
    @test length(R.create_values(qr, qr_face, interp_geom, interp_u1)) == 2
    @test length(R.create_values(qr, qr_face, interp_geom, interp_u2...)) == 3
end

@testset "create material properties" begin
    func_1D(x) = x[1]
    func_2D(x) = x[1] + x[2]
    elasticity1D = IncompressibleElasticity(E = func_1D, ν =0.4999)
    elasticity2D = IncompressibleElasticity(E = func_2D, ν =0.4999)

    r1D = Rheology(nothing,nothing,elasticity1D,nothing)
    r2D = Rheology(nothing,nothing,elasticity2D,nothing)

    grid1D = generate_grid(Line,(2,), Vec(0.0), Vec(10.0))
    grid2D = generate_grid(Quadrilateral,(2, 2), Vec(0.0,0.0), Vec(10.0,10.0))

    mp1D = create_material_properties(grid1D, r1D)
    mp2D = create_material_properties(grid1D, r1D)

    @test r1D(1.0).elasticity.E == 1.0
    @test r2D([1.0,2.0]).elasticity.E == 3.0

end


"""\nFunction that creates Utils struct from model definition\n# Arguments\n- `mdl::Model`: mutable struct Model.\n- `method_cube::String="CPU"`: the method for the utils cube. Try "CPU", "GPU","Mmap","Function".\n# Examples\n```julia-repl\njulia> execute_cube(my_model, "CPU")\n```\n"""
function execute_cube(mdl::Model,method_cube::String="CPU")

    (method_cube == "CPU") && ((cubefun, precompilation) = (cube_CPU, true))
    (method_cube == "GPU") && ((cubefun, precompilation) = (cube_GPU, true))
    (method_cube == "Mmap") && ((cubefun, precompilation) = (cube_Mmap, true))
    (method_cube == "Function") && ((cubefun, precompilation) = (cube_function, false))

    ( ! @isdefined cubefun ) && error("1.Method provided for Cube is incorrect. Check for spelling errors.")

    # precompilation = true

    # if (method_cube == "Function")

    #     precompiletion = false

    # end

    if precompilation == true

        (method_cube == "CPU") && (qb = cube_CPU_qb(mdl))
        (method_cube == "GPU") && (qb = cube_GPU_qb(mdl))
        (method_cube == "Mmap") && (qb = cube_Mmap_qb(mdl))
        return( Utils(cubefun,qb) )

    end

    return( Utils(cubefun) )

end


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

"""
    AbstractData

An abstract type to aid with custom spatial data types. Users can
subtype their domains from this type, and implement the methods in
`geotraits/data.jl`.
"""
abstract type AbstractData end

"""
    sdata₁ == sdata₂

Tells whether or not the spatial `sdata₁` and `sdata₂` are equal.
"""
==(sdata₁::AbstractData, sdata₂::AbstractData) =
  domain(sdata₁) == domain(sdata₂) &&
  values(sdata₁) == values(sdata₂)

# -----------
# TABLES API
# -----------

Tables.istable(::Type{<:AbstractData}) = true
Tables.materializer(sdata::AbstractData) = Tables.materializer(values(sdata))
Tables.columnaccess(sdata::AbstractData) = Tables.columnaccess(values(sdata))
Tables.rowaccess(sdata::AbstractData) = Tables.rowaccess(values(sdata))
Tables.schema(sdata::AbstractData) = Tables.schema(values(sdata))
Tables.columns(sdata::AbstractData) = Tables.columns(values(sdata))
Tables.columnnames(sdata::AbstractData) = Tables.columnnames(values(sdata))
Tables.getcolumn(sdata::AbstractData, c::Symbol) = Tables.getcolumn(values(sdata), c)
Tables.rows(sdata::AbstractData) = Tables.rows(values(sdata))

# -------------
# VARIABLE API
# -------------

function variables(sdata::AbstractData)
  s = Tables.schema(sdata)
  ns, ts = s.names, s.types
  @. Variable(ns, nonmissing(ts))
end

Base.getindex(sdata::AbstractData, var::Symbol) =
  Tables.getcolumn(sdata, var)

# ---------
# VIEW API
# ---------

Base.view(sdata::AbstractData, inds::AbstractVector{Int}) =
  DataView(sdata, inds, collect(name.(variables(sdata))))
Base.view(sdata::AbstractData, vars::AbstractVector{Symbol}) =
  DataView(sdata, 1:nelms(sdata), vars)
Base.view(sdata::AbstractData, inds, vars) =
  DataView(sdata, inds, vars)

#------------------
# IMPLEMENTATIONS
#------------------
"""
    SpatialData(domain, data)

Tabular `data` georeferenced in a spatial `domain`.
"""
struct SpatialData{𝒟,𝒯} <: AbstractData
  domain::𝒟
  table::𝒯

  function SpatialData{𝒟,𝒯}(domain, table) where {𝒟,𝒯}
    ne = nelms(domain)
    nr = length(Tables.rows(table))
    @assert ne == nr "number of table rows ≠ number of mesh elements"
    new(domain, table)
  end
end

SpatialData(domain::𝒟, table::𝒯) where {𝒟,𝒯} =
  SpatialData{𝒟,𝒯}(domain, table)

domain(sdata::SpatialData) = sdata.domain
values(sdata::SpatialData) = sdata.table

# ------------
# IO methods
# ------------
function Base.show(io::IO, sdata::SpatialData)
  N = ncoords(sdata)
  T = coordtype(sdata)
  n = nelms(sdata)
  print(io, "$n SpatialData{$T,$N}")
end

function Base.show(io::IO, ::MIME"text/plain", sdata::SpatialData)
  𝒟 = domain(sdata)
  𝒯 = values(sdata)
  s = Tables.schema(𝒯)
  vars = zip(s.names, s.types)
  println(io, 𝒟)
  println(io, "  variables")
  varlines = ["    └─$var ($V)" for (var,V) in vars]
  print(io, join(sort(varlines), "\n"))
end


const ADSUPPORT = (:ForwardDiff, :Zygote)
const ADAVAILABLE = Dict{Module, Function}()

Hamiltonian(metric::AbstractMetric, ℓπ, m::Module) = ADAVAILABLE[m](metric, ℓπ)

function Hamiltonian(metric::AbstractMetric, ℓπ)
    available = collect(keys(ADAVAILABLE))
    if length(available) == 0
        support_list_str = join(ADSUPPORT, " or ")
        error("MethodError: no method matching Hamiltonian(metric::AbstractMetric, ℓπ) because no backend is loaded. Please load an AD package ($support_list_str) first.")
    elseif length(available) > 1
        available_list_str = join(keys(ADAVAILABLE), " and ")
        constructor_list_str = join(map(m -> "Hamiltonian(metric, ℓπ, $m)", available), "\n  ")
        error("MethodError: Hamiltonian(metric::AbstractMetric, ℓπ) is ambiguous because multiple AD pakcages are available ($available_list_str). Please use AD explictly. Candidates:\n  $constructor_list_str")
    else
        return Hamiltonian(metric, ℓπ, first(available))
    end
end

### ForwardDiff

@require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin

import .ForwardDiff, .ForwardDiff.DiffResults

function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractVector)
    res = DiffResults.GradientResult(θ)
    ForwardDiff.gradient!(res, ℓπ, θ)
    return DiffResults.value(res), DiffResults.gradient(res)
end

# Implementation 1
function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractMatrix)
    jacob = similar(θ)
    res = DiffResults.JacobianResult(similar(θ, size(θ, 2)), jacob)
    ForwardDiff.jacobian!(res, ℓπ, θ)
    jacob_full = DiffResults.jacobian(res)

    d, n = size(jacob)
    for i in 1:n
        jacob[:,i] = jacob_full[i,1+(i-1)*d:i*d]
    end
    return DiffResults.value(res), jacob
end

# Implementation 2
# function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractMatrix)
#     local densities
#     f(x) = (densities = ℓπ(x); sum(densities))
#     res = DiffResults.GradientResult(θ)
#     ForwardDiff.gradient!(res, f, θ)
#     return ForwardDiff.value.(densities), DiffResults.gradient(res)
# end

# Implementation 3
# function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractMatrix)
#     v = similar(θ, size(θ, 2))
#     g = similar(θ)
#     for i in 1:size(θ, 2)
#         res = GradientResult(θ[:,i])
#         gradient!(res, ℓπ, θ[:,i])
#         v[i] = value(res)
#         g[:,i] = gradient(res)
#     end
#     return v, g
# end

function ForwardDiffHamiltonian(metric::AbstractMetric, ℓπ)
    ∂ℓπ∂θ(θ::AbstractVecOrMat) = ∂ℓπ∂θ_forwarddiff(ℓπ, θ)
    return Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
end

ADAVAILABLE[ForwardDiff] = ForwardDiffHamiltonian

end # @require

### Zygote

@require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin

import .Zygote

function ∂ℓπ∂θ_zygote(ℓπ, θ::AbstractVector)
    res, back = Zygote.pullback(ℓπ, θ)
    return res, first(back(Zygote.sensitivity(res)))
end

function ∂ℓπ∂θ_zygote(ℓπ, θ::AbstractMatrix)
    res, back = Zygote.pullback(ℓπ, θ)
    return res, first(back(ones(eltype(res), size(res))))
end

function ZygoteADHamiltonian(metric::AbstractMetric, ℓπ)
    ∂ℓπ∂θ(θ::AbstractVecOrMat) = ∂ℓπ∂θ_zygote(ℓπ, θ)
    return Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
end

ADAVAILABLE[Zygote] = ZygoteADHamiltonian

# Zygote.@adjoint

end # @require


@testset "Threading" begin
    @testset "Check number of threads" begin
        @test ccall(Arblib.@libflint(flint_get_num_threads), Cint, ()) == 1

        Arblib.flint_set_num_threads(2)
        @test ccall(Arblib.@libflint(flint_get_num_threads), Cint, ()) == 2

        Arblib.flint_set_num_threads(4)
        @test ccall(Arblib.@libflint(flint_get_num_threads), Cint, ()) == 4

        Arblib.flint_set_num_threads(1)
        @test ccall(Arblib.@libflint(flint_get_num_threads), Cint, ()) == 1
    end

    @testset "Disabled threading" begin
        # This manually unsets the threading flag and checks that we
        # get an error when trying to set the number of threads
        Arblib.__isthreaded[] = false

        @test_throws ErrorException Arblib.flint_set_num_threads(2)

        Arblib.__isthreaded[] = true
    end

    @testset "Matrix multiplication" begin
        A = ArbMatrix(reshape(1:10000, 100, 100))

        Arblib.flint_set_num_threads(1)
        B1 = Arblib.mul_threaded!(similar(A), A, A)

        Arblib.flint_set_num_threads(2)
        B2 = Arblib.mul_threaded!(similar(A), A, A)

        @test B1 == B2 == A * A

        Arblib.flint_set_num_threads(2)
    end
end


include("codec.jl")

using MsgPack
using CodecZlib

export decodedata,encodedata,parsemmtf,writemmtf,fetchmmtf

"""
Decodes the MsgPack unpacked dict of the MMTF file.
"""
function decodedata(input_data::Dict)
    # Initialize optional field's default values
    dict = Dict(
        "bFactorList" => Array{Float32,1}[],
        "occupancyList" => Array{Float32,1}[],
        "atomIdList" => Array{Int32,1}[],
        "altLocList" => Array{Char,1}[],
        "insCodeList" => Array{Char,1}[],
        "sequenceIndexList" => Array{Int32,1}[],
        "chainNameList" => Array{String,1}[],
        "spaceGroup" => "",
        "bondAtomList" => Array{Int32,1}[],
        "bondOrderList" => Array{Int8,1}[],
        "mmtfVersion" => "",
        "mmtfProducer" => "",
        "structureId" => "",
        "title" => "",
        "experimentalMethods" => Array{Any,1}[],
        "depositionDate" => "",
        "releaseDate" => "",
        "entityList" => Array{Any,1}[],
        "bioAssemblyList" => Array{Any,1}[],
        "rFree" => "",
        "rWork" => "",
        "resolution" => "",
        "unitCell" => Array{Any,1}[],
        "secStructList" => Array{Int8,1}[],
        "ncsOperatorList" => Array{Any,1}[]
    )
    # Check if all mandatory fields are present in the input data.
    mandateFields = ["groupTypeList","xCoordList","yCoordList","zCoordList","groupIdList","groupList","chainsPerModel","groupsPerChain","chainIdList","numBonds","numChains","numModels","numAtoms","numGroups"]
    for field in mandateFields
        if !haskey(input_data, field)
            throw(ArgumentError("Mandatory field \"$field\" not available in input data!"))
        end
    end
    # decode the necessary fields alone
    toDecodeFields = ["groupTypeList","xCoordList","yCoordList","zCoordList","bFactorList","occupancyList","atomIdList","altLocList","insCodeList","groupIdList","sequenceIndexList","chainNameList","chainIdList","bondAtomList","bondOrderList","secStructList"]
    for (key,value) in input_data
        if key in toDecodeFields
            dict[key] = decodearray(value)
        else
            dict[key] = value
        end
    end
    # return the final dict
    return dict
end

"""
Encodes the MMTF Dict.
"""
function encodedata(MMTFDict::Dict)
    encodedMMTFDict = Dict()
    encodedMMTFDict["groupTypeList"] = encodearray(MMTFDict["groupTypeList"], 4, 0)
    encodedMMTFDict["xCoordList"] = encodearray(MMTFDict["xCoordList"], 10, COORD_DIVIDER)
    encodedMMTFDict["yCoordList"] = encodearray(MMTFDict["yCoordList"], 10, COORD_DIVIDER)
    encodedMMTFDict["zCoordList"] = encodearray(MMTFDict["zCoordList"], 10, COORD_DIVIDER)
    encodedMMTFDict["bFactorList"] = encodearray(MMTFDict["bFactorList"], 10, OCC_B_FACTOR_DIVIDER)
    encodedMMTFDict["occupancyList"] = encodearray(MMTFDict["occupancyList"], 9, OCC_B_FACTOR_DIVIDER)
    encodedMMTFDict["atomIdList"] = encodearray(MMTFDict["atomIdList"], 8, 0)
    encodedMMTFDict["altLocList"] = encodearray(MMTFDict["altLocList"], 6, 0)
    encodedMMTFDict["insCodeList"] = encodearray(MMTFDict["insCodeList"], 6, 0)
    encodedMMTFDict["groupIdList"] = encodearray(MMTFDict["groupIdList"], 8, 0)
    encodedMMTFDict["sequenceIndexList"] = encodearray(MMTFDict["sequenceIndexList"], 8, 0)
    encodedMMTFDict["chainNameList"] = encodearray(MMTFDict["chainNameList"], 5, 4)
    encodedMMTFDict["chainIdList"] = encodearray(MMTFDict["chainIdList"], 5, 4)
    encodedMMTFDict["bondAtomList"] = encodearray(MMTFDict["bondAtomList"], 4, 0)
    encodedMMTFDict["bondOrderList"] = encodearray(MMTFDict["bondOrderList"], 2, 0)
    encodedMMTFDict["secStructList"] = encodearray(MMTFDict["secStructList"], 2, 0)
    encodedMMTFDict["groupList"] = MMTFDict["groupList"]
    encodedMMTFDict["chainsPerModel"] = MMTFDict["chainsPerModel"]
    encodedMMTFDict["groupsPerChain"] = MMTFDict["groupsPerChain"]
    encodedMMTFDict["spaceGroup"] = MMTFDict["spaceGroup"]
    encodedMMTFDict["mmtfVersion"] = MMTFDict["mmtfVersion"]
    encodedMMTFDict["mmtfProducer"] = MMTFDict["mmtfProducer"]
    encodedMMTFDict["structureId"] = MMTFDict["structureId"]
    encodedMMTFDict["entityList"] = MMTFDict["entityList"]
    encodedMMTFDict["bioAssemblyList"] = MMTFDict["bioAssemblyList"]
    encodedMMTFDict["rFree"] = MMTFDict["rFree"]
    encodedMMTFDict["rWork"] = MMTFDict["rWork"]
    encodedMMTFDict["resolution"] = MMTFDict["resolution"]
    encodedMMTFDict["title"] = MMTFDict["title"]
    encodedMMTFDict["experimentalMethods"] = MMTFDict["experimentalMethods"]
    encodedMMTFDict["depositionDate"] = MMTFDict["depositionDate"]
    encodedMMTFDict["releaseDate"] = MMTFDict["releaseDate"]
    encodedMMTFDict["unitCell"] = MMTFDict["unitCell"]
    encodedMMTFDict["numBonds"] = MMTFDict["numBonds"]
    encodedMMTFDict["numChains"] = MMTFDict["numChains"]
    encodedMMTFDict["numModels"] = MMTFDict["numModels"]
    encodedMMTFDict["numAtoms"] = MMTFDict["numAtoms"]
    encodedMMTFDict["numGroups"] = MMTFDict["numGroups"]
    encodedMMTFDict["ncsOperatorList"] = MMTFDict["ncsOperatorList"]
    return encodedMMTFDict
end

"""
Parses the MMTF file into a MMTFDict.
filepath: Path of the input file.
"""
function parsemmtf(filepath::AbstractString; gzip::Bool=false)
    return open(filepath) do input
        parsemmtf(input,gzip=gzip)
    end
end

function parsemmtf(input::IO; gzip::Bool=false)
    if gzip
        decompressed = IOBuffer()
        write(decompressed, GzipDecompressorStream(input))
        decodedata(unpack(take!(decompressed)))
    else
        decodedata(unpack(input))
    end
end

"""
Writes the MMTFDict back to a file.
MMTFDict: Dict containing the decoded MMTF data to be written to file.
filepath: Path to write the file
"""
function writemmtf(MMTFDict::Dict, filepath::AbstractString; gzip::Bool=false)
    open(filepath, "w") do output
        writemmtf(MMTFDict,output,gzip=gzip)
    end
end

function writemmtf(MMTFDict::Dict, output::IO; gzip::Bool=false)
    if gzip
        stream=IOBuffer(pack(encodedata(MMTFDict)))
        write(output,GzipCompressorStream(stream))
    else
        pack(output,encodedata(MMTFDict))
    end
end

"""
Fetches a file from RCSB server and decodes it into a MMTFDict.
pdbid: ID of the PDB file to be fetched and decoded.
"""
function fetchmmtf(pdbid)
    tempfile = tempname()
    download("$(BASE_URL)/$(pdbid).mmtf.gz",tempfile)
    parsemmtf(tempfile,gzip=true)
end


sparseaccess(A::SparseMatrixCSC) = A
sparseaccess(A::SubArray{<:Any,2,<:SparseMatrixCSC{<:Any,<:Any},<:Tuple{I,<:Any},false}) where
    I<:Union{Base.Slice{<:Base.OneTo},AbstractUnitRange} = A

sparseaccess(A::SubArray) = XSubArray(A)

# the actual nonzero entries belonging to the view
nzvalview(S::SparseMatrixCSCView) = view(S.parent.nzval, first(nzrange(S, first(axes(S, 2)))):last(nzrange(S, last(axes(S, 2)))))

#
nonzeros(S::SubArray{<:Any,2,<:SparseMatrixCSC}) = nonzeros(S.parent)
nonzeros(S::XSubArray{<:Any,<:Integer,2}) = nonzeros(S.sub)

nnz(S::SparseMatrixCSCView) = nzrange(S.parent, last(S.indices[2]))[end] - nzrange(S.parent, first(S.indices[2]))[1] + 1
nzrange(S::SparseMatrixCSCView, col::Integer) = nzrange(S.parent, S.indices[2][col])
function nzrange(S::SubArray{<:Any,2,<:SparseMatrixCSC,<:Tuple{I,J},false}, i::Integer) where {I<:AbstractUnitRange,J<:AbstractVector{<:Integer}}
    A = S.parent
    r = nzrange(A, S.indices[2][i])
    rvA = rowvals(A)
    m = axes(A, 1)
    r1 = first(r)
    r2 = last(r)
    ri = S.indices[1]
    i1 = first(ri)
    i2 = last(ri)
    if i1 > first(m)
        r1 = searchsortedfirst(rvA, i1, r1, r2, Forward)
    end
    if i2 < last(m)
        r2 = searchsortedlast(rvA, i2, r1, r2, Forward)
    end
    r1:r2
end

function nzrange(xS::XSubArray{<:Any,<:Integer,2,<:SparseMatrixCSC,<:Tuple{I,J},false}, i::Integer) where {I<:AbstractVector{<:Integer},J<:AbstractVector{<:Integer}}
    S = xS.sub
    A = S.parent
    r = nzrange(A, S.indices[2][i])
    rvA = rowvals(A)
    m = axes(A, 1)
    rangestrip(rvA, r, m, xS)
end

function rangestrip(rvA, r::AbstractVector{Int}, m::AbstractUnitRange, S::XSubArray{<:Any,<:Integer,2,<:Any,<:Tuple{I,<:Any}}) where I<:StepRange
    r1 = first(r)
    r2 = last(r)
    ri = S.sub.indices[1]
    i1 = first(ri)
    i2 = last(ri)
    minmax(i1, i2)
    if i1 > first(m)
        r1 = searchsortedfirst(rvA, i1, r1, r2, Forward)
    end
    if i2 < last(m)
        r2 = searchsortedlast(rvA, i2, r1, r2, Forward)
    end
    nzr = Vector{eltype(ri)}(undef, r2 - r1 + 1)
    rv = rowvals(S)
    k = 0
    for i in r1:r2
        j = rv[i]
        if j != 0
            nzr[k+=1] = i
        end
    end
    resize!(nzr, k)
    step(ri) < 0 ? reverse(nzr) : nzr
end

function rangestrip(rvA, r::AbstractVector{Int}, m::AbstractUnitRange, S::XSubArray{<:Any,<:Integer,2,<:Any,<:Tuple{I,<:Any}}) where I<:AbstractVector{<:Integer}
    r1 = first(r)
    r2 = last(r)
    ri = S.sub.indices[1]
    i1, i2 = extrema(ri)
    if i1 > first(m)
        r1 = searchsortedfirst(rvA, i1, r1, r2, Forward)
    end
    if i2 < last(m)
        r2 = searchsortedlast(rvA, i2, r1, r2, Forward)
    end
    vind = rowvals(S)[r1:r2]
    p = sortperm(vind)
    k = searchsortedlast(vind[p], 0)
    (r1:r2)[p][k+1:end]
end

rowvals(S::SparseMatrixCSCView) = S.parent.rowval
rowvals(xs::XSubArray) = xs.rowval

struct RowIndexVector{U,T,D} <: AbstractVector{Int}
    rvp::T
    ind::U
    dict::D
end

function rowvals(S::SubArray{<:Any,2,<:Any,<:Tuple{<:Base.Slice,<:AbstractVector},false})
    rowvals(S.parent)
end
function rowvals(S::SubArray{<:Any,2,<:Any,<:Tuple{<:UnitRange,<:AbstractVector},false})
    RowIndexVector(rowvals(S.parent), S.indices[1], nothing)
end
function rowvals(S::SubArray{<:Any,2,<:Any,<:Tuple{<:OrdinalRange,<:AbstractVector},false})
    RowIndexVector(rowvals(S.parent), S.indices[1], nothing)
end
function rowvals(S::SubArray{<:Any,2,<:Any,<:Tuple{<:AbstractVector,<:AbstractVector},false})
    ind = S.indices[1]
    RowIndexVector(rowvals(S.parent), ind, Dict(reverse.(enumerate(ind))))
end

import Base: getindex, size
size(rv::RowIndexVector) = size(rv.rvp)
size(s::XSubArray) = size(s.sub)
getindex(s::XSubArray, ind...) = getindex(s.sub, ind...)

function getindex(rv::RowIndexVector{<:UnitRange}, k::Integer)
    i = rv.rvp[k]
    ind = rv.ind
    a = first(ind)
    b = last(ind)
    ( a <= i <= b ) ? ( i - a + 1 ) : 0
end
function getindex(rv::RowIndexVector{<:StepRange}, k::Integer)
    i = rv.rvp[k]
    ind = rv.ind
    a = first(ind)
    b = last(ind)
    a <= i <= b || return 0
    s = step(ind)
    i -= a
    rem(i, s) == 0 ? 0 : div(i, s)
end
function getindex(rv::RowIndexVector{<:AbstractVector}, k::Integer)
    i = rv.rvp[k]
    ind = rv.ind
    a, b = extrema(ind)
    get(rv.dict, i, 0)
end

SparseArrays.nzrange(X::SparseVector, i) = 1:length(X.nzind)
SparseArrays.rowvals(X::SparseVector) = X.nzind


export TEBDOptions, TEBDResults, storeTEBDResult, runTEBD

struct TEBDOptions
    order::Int64
    tfinal::Float64
    dt::Float64
    substeps::Int64
    cutoff::Float64
end

struct TEBDResults
    time::Vector{Float64}
    observables::Dict{String,Tuple{Vector{ComplexF64},Vector{ComplexF64},Vector{ComplexF64}}}
    localOperators::Dict{String,Vector{Vector{ComplexF64}}}
    correlationFunctions::Dict{String,Vector{Matrix{ComplexF64}}}
    maxBondDimension::Vector{Int64}
end

function buildGatesTEBD1(model::Model, dt::Float64)::Vector{ITensor}
    return vcat(getGatesEven(model, dt), getGatesOdd(model, dt))
end

function buildGatesTEBD2(model::Model, dt::Float64)::Vector{ITensor}
    even = getGatesEven(model, dt / 2.0)
    return vcat(even, getGatesOdd(model, dt), reverse(even))
end

function buildGatesTEBD3(model::Model, dt::Float64)::Vector{ITensor}
    s = 1.0 / (2.0 - (2.0^(1.0 / 3.0)))
    dt1 = s * dt
    dt2 = (1 - 2.0 * s) * dt
    return vcat(
        buildGatesTEBD2(model, dt1),
        buildGatesTEBD2(model, dt2),
        buildGatesTEBD2(model, dt1)
    )
end

function storeTEBDResult(file::String, result::TEBDResults)
    mkpath(dirname(file))
    fptr = HDF5.h5open(file, "w")

    HDF5.write(fptr, "time", result.time)

    grp_observables = HDF5.create_group(fptr, "observables")
    for (name, values) in result.observables
        grp_observable = HDF5.create_group(grp_observables, string(name))
        HDF5.write(grp_observable, "value", values[1])
        HDF5.write(grp_observable, "squared", values[2])
        HDF5.write(grp_observable, "variance", values[3])
    end

    grpLocalOperators = HDF5.create_group(fptr, "local_operators")
    for (name, values) in result.localOperators
        HDF5.write(grpLocalOperators, name, reduce(hcat, values))
    end

    grpCorrelationFunctions = HDF5.create_group(fptr, "correlation_functions")
    for (name, values) in result.correlationFunctions
        # FIXME: there should be definitely a clever way to do this using reduce, hcat and reshape
        mat = Array{ComplexF64,3}(undef, size(values[1])[1], size(values[1])[1], length(result.time))
        for (i, _) in enumerate(result.time)
            mat[:, :, i] = values[i]
        end
        HDF5.write(grpCorrelationFunctions, name, mat)
    end

    HDF5.write(fptr, "max_bond_dimension", result.maxBondDimension)

    HDF5.close(fptr)
end

function runTEBD(psi0::MPS, model::Model, options::TEBDOptions)::TEBDResults
    sites = model.sites

    step = options.dt / options.substeps
    @info "Set substep size for TEBD to $step"

    @info "Construct TEBD gates …"
    gates = ITensor[]
    if options.order == 1
        gates = buildGatesTEBD1(model, step)
    elseif options.order == 2
        gates = buildGatesTEBD2(model, step)
    elseif options.order == 3
        gates = buildGatesTEBD3(model, step)
    else
        throw(DomainError(order, "TEBD not implemented for specified order"))
    end
    num_gates = length(gates)
    @info "Constructed $num_gates gates"

    gates = repeat(gates, options.substeps)

    time = 0.0
    psi = deepcopy(psi0)
    results = TEBDResults(
        Float64[],
        Dict{String,Tuple{Vector{ComplexF64},Vector{ComplexF64},Vector{ComplexF64}}}(),
        Dict{String,Vector{Vector{ComplexF64}}}(),
        Dict{String,Vector{Matrix{ComplexF64}}}(),
        Int64[]
    )

    push!(results.time, time)
    push!(results.maxBondDimension, maxlinkdim(psi))

    # measure initial observables
    for observable in getObservables(model)
        results.observables[observable.name] = (ComplexF64[], ComplexF64[], ComplexF64[])
        values = expect(psi, observable)
        push!(results.observables[observable.name][1], values[1])
        push!(results.observables[observable.name][2], values[2])
        push!(results.observables[observable.name][3], values[3])
    end

    # measure initial local operators
    for localOperator in getLocalOperators(model)
        results.localOperators[localOperator.name] = Vector{ComplexF64}[]
        push!(results.localOperators[localOperator.name], expect(psi, localOperator))
    end

    # measure initial correlation functions
    for correlationFunction in getCorrelationFunctions(model)
        results.correlationFunctions[correlationFunction.name] = Matrix{ComplexF64}[]
        push!(results.correlationFunctions[correlationFunction.name], expect(psi, correlationFunction))
    end

    maxBondDimension = 0

    while time < options.tfinal
        psi = apply(gates, psi; cutoff = options.cutoff)
        time += options.dt
        push!(results.time, time)
        push!(results.maxBondDimension, maxBondDimension)

        # measure current observables
        for observable in getObservables(model)
            values = expect(psi, observable)
            push!(results.observables[observable.name][1], values[1])
            push!(results.observables[observable.name][2], values[2])
            push!(results.observables[observable.name][3], values[3])
            maxBondDimension = maxlinkdim(psi)
        end

        # measure initial local operators
        for localOperator in getLocalOperators(model)
            push!(results.localOperators[localOperator.name], expect(psi, localOperator))
        end

        # measure initial correlation functions
        for correlationFunction in getCorrelationFunctions(model)
            push!(results.correlationFunctions[correlationFunction.name], expect(psi, correlationFunction))
        end

        percentage = time / options.tfinal * 100.0

        timestamp = Dates.format(now(), "YYYY-mm-dd HH:MM:SS")

        @info "$timestamp: Finished time=$time ($percentage%)"
        @info "                     max bond dimension=$maxBondDimension"
    end

    return results
end


using DiffEqBase: CallbackSet

using NQCDynamics: InitialConditions, DynamicsOutputs

abstract type AbstractSelection end

"""
Select the initial conditions from the distribution in order.
"""
struct OrderedSelection{D,I} <: AbstractSelection
    "Distribution that is sampled."
    distribution::D
    indices::I
end

"""
Obtain initial conditions by randomly sampling the distribution.
"""
struct RandomSelection{D} <: AbstractSelection
    "Distribution that is sampled."
    distribution::D
end

function Selection(distribution, selection::AbstractVector)
    @info "Sampling the provided distribution in the range $selection."
    OrderedSelection(distribution, selection)
end

function Selection(distribution, ::Any)
    @info "Sampling randomly from provided distribution."
    RandomSelection(distribution)
end

function (select::OrderedSelection)(prob, i, repeat)
    j = select.indices[i]
    u0 = sample_distribution(prob.p, select.distribution, j)
    DynamicsMethods.create_problem(u0, prob.tspan, prob.p)
end

function (select::RandomSelection)(prob, i, repeat)
    j = rand(1:lastindex(select.distribution))
    u0 = sample_distribution(prob.p, select.distribution, j)
    DynamicsMethods.create_problem(u0, prob.tspan, prob.p)
end

struct SelectWithCallbacks{S<:AbstractSelection,C1,C2,V} <: AbstractSelection
    selection::S
    standard_callbacks::C1
    changing_callbacks::C2
    values::V
    function SelectWithCallbacks(selection, standard_callbacks, output, trajectories; saveat=[])

        callbacks = []
        values = []
        for i=1:trajectories
            cb, vals = DynamicsOutputs.create_saving_callback(output; saveat=saveat)
            push!(callbacks, cb)
            push!(values, vals)
        end

        new{typeof(selection), typeof(standard_callbacks), typeof(callbacks), typeof(values)}(
            selection, standard_callbacks, callbacks, values)
    end
end

function (select::SelectWithCallbacks)(prob, i, repeat)
    prob = select.selection(prob, i, repeat)
    remake(prob, callback=CallbackSet(select.standard_callbacks, select.changing_callbacks[i]))
end


using TabularTDLearning
using POMDPs
using POMDPModels
using POMDPPolicies
using POMDPTesting
using Random
using Test
import POMDPLinter: @show_requirements

mdp = SimpleGridWorld()
Random.seed!(1)

@testset "qlearning" begin
    rng = MersenneTwister(1)
    solver = QLearningSolver(exploration_policy=EpsGreedyPolicy(mdp, 0.5, rng=rng), learning_rate=0.1, n_episodes=500, max_episode_length=50, eval_every=50, n_eval_traj=100, rng=rng)

    println("Test QLearning requirements: ")
    @show_requirements solve(solver, mdp)

    r = test_solver(solver, mdp, max_steps = 100)
    @test r > 0.0
end

@testset "sarsa" begin
    rng = MersenneTwister(1)
    solver = SARSASolver(exploration_policy=EpsGreedyPolicy(mdp, 0.5, rng=rng), learning_rate=0.1, n_episodes=500, max_episode_length=50, eval_every=50, n_eval_traj=100, rng=rng)

    println("Test SARSA requirements: ")
    @show_requirements solve(solver, mdp)

    r = test_solver(solver, mdp, max_steps = 100)
    @test r > 0.0
end

@testset "sarsa λ" begin
    rng = MersenneTwister(2)
    solver = SARSALambdaSolver(exploration_policy=EpsGreedyPolicy(mdp, 0.5, rng=rng), learning_rate=0.1, lambda=0.9, n_episodes=500, max_episode_length=50, eval_every=50, n_eval_traj=100, rng=rng)

    println("Test SARSALambda requirements: ")
    @show_requirements solve(solver, mdp)

    r = test_solver(solver, mdp, max_steps = 100)
    @test r > 0.0
end


using Revise
using Plots
using Pkg
# Pkg.activate(".")
using RadiativeTransfer
using RadiativeTransfer.Architectures
using RadiativeTransfer.Absorption
using RadiativeTransfer.Scattering
using RadiativeTransfer.vSmartMOM
using RadiativeTransfer.SolarModel
using InstrumentOperator
using Interpolations
using Polynomials
using ForwardDiff
using Distributions
using NCDatasets

## Atmospheric Radiative Transfer

# Load parameters from file
parameters = vSmartMOM.parameters_from_yaml("test/test_parameters/3BandParameters.yaml")
#parameters.architecture = CPU()
FT = Float64

# Load OCO Data:
# File names:
L1File   = "/net/fluo/data1/group/oco2/L1bSc/oco2_L1bScND_26780a_190715_B10003r_200429212407.h5"
metFile  = "/net/fluo/data1/group/oco2/L2Met/oco2_L2MetND_26780a_190715_B10003r_200429212406.h5"
dictFile = "/home/cfranken/code/gitHub/InstrumentOperator.jl/json/oco2.yaml"
# Load L1 file (could just use filenames here as well)
oco = InstrumentOperator.load_L1(dictFile,L1File, metFile);

# Pick some bands as tuple (or just one)
bands = (1,2,3);
#bands = (1,3);
# Indices within that band:
indices = (92:885,114:845,50:916);
#indices = (92:885,50:916);
# Geo Index (footprint,sounding):
GeoInd = [5,5000];

# Get data for that sounding:
oco_sounding = InstrumentOperator.getMeasurement(oco, bands, indices, GeoInd);

###################################################
# Produce black-body in wavenumber range

Tsolar = solar_transmission_from_file("/home/rjeyaram/RadiativeTransfer/src/SolarModel/solar.out")
Tsolar_interp = LinearInterpolation(Tsolar[:, 1], Tsolar[:, 2])

function getSolar(grid, Tsolar)
    T = 5777.0
    λ_grid = reverse(1e4 ./ grid) #collect(757:0.01:777)
    black_body = planck_spectrum_wl(T, λ_grid) * 2.1629e-05 * π
    black_body = SolarModel.watts_to_photons(λ_grid, black_body)

    # Get solar transmittance spectrum
    #solar_transmission = solar_transmission_from_file("/home/rjeyaram/RadiativeTransfer/src/solar_merged_20160127_600_26316_100.out", parameters.spec_bands[1])
    solar_transmission = Tsolar(grid)
    # Get outgoing solar radiation
    sun_out = reverse(solar_transmission) .* black_body;
    return sun_out
end
###################################################

# Runner is used to set AD fields as duals
function runner!(y, x, parameters=parameters, oco_sounding= oco_sounding, Tsolar = Tsolar_interp)

    # Set parameters fields as the dual numbers
    parameters.brdf = [vSmartMOM.LambertianSurfaceLegendre([x[1],x[3],x[4]]),vSmartMOM.LambertianSurfaceLegendre([x[7],x[8],x[5]]),vSmartMOM.LambertianSurfaceLegendre([x[9],x[10],x[6]])]

    parameters.scattering_params.rt_aerosols[1].τ_ref = exp(x[2]);
    parameters.scattering_params.rt_aerosols[1].p₀    = 800.0; #x[4]

    parameters.p   = oco_sounding.p_half
    parameters.q   = oco_sounding.q
    parameters.T   = oco_sounding.T# .+ 1.0 #.+ x[15]
    parameters.sza = oco_sounding.sza
    parameters.vza = [oco_sounding.vza]
    parameters.absorption_params.vmr["H2O"] = [parameters.q[1:65]*x[11] * 1.8; parameters.q[66:end]*x[15] * 1.8]
    a1 = zeros(7) .+ x[12]
    a2 = zeros(7) .+ x[13]
    a3 = zeros(6) .+ x[14]
    parameters.absorption_params.vmr["CO2"] = [a1; a2; a3]
    model = model_from_parameters(parameters);

    for i = 1:length(oco_sounding.BandID)
        println("Modelling band $(i)")
        # Run the model to obtain reflectance matrix
        R = vSmartMOM.rt_run(model, i_band=i);
        RR = oco_sounding.mueller[1]*R[1,1,:] + oco_sounding.mueller[2]*R[1,2,:] + oco_sounding.mueller[2]*R[1,3,:]

        # Get sun:
        @time sun_out = getSolar(parameters.spec_bands[i],Tsolar)
        # Apply Earth reflectance matrix
        earth_out = sun_out .* reverse(RR[:])

        # Re-interpolate I from ν_grid to new grid/resolution
        λ_grid = reverse(1e4 ./ parameters.spec_bands[i])
        @time interp_I = LinearInterpolation(λ_grid, earth_out);
        res = 0.001e-3;
        off = 0.5e-3
        wl = oco_sounding.ils[i].ν_out[1]-off:res:oco_sounding.ils[i].ν_out[end]+off;
        @show wl[1],wl[end], λ_grid[1],λ_grid[end]
        I_wl = interp_I(wl);

        # Convolve input spectrum with variable kernel
        @time I_conv = InstrumentOperator.conv_spectra(oco_sounding.ils[i], wl, I_wl)
        y[oco_sounding.BandID[i]] = I_conv
    end

end

# State vector
x = FT[0.2377, -3, -0.00244, 0, 0,0, 0.4, 0, 0.4, 0,1,400e-6,400e-6,400e-6, 1.0]#,

# Run FW model:
# @time runner(x);
y = oco_sounding.SpectralMeasurement;
Fx = zeros(length(y));
#ind = 92:885

for i=1:5
    K = ForwardDiff.jacobian(runner!, Fx, x);
    dx = K \ (y-Fx);
    x += dx;
    @show x;
end
runner!(Fx,x)

ν = oco_sounding.SpectralGrid*1e3
plot(y/1e20, label="Meas")
plot!(Fx/1e20, label="Mod")
plot!((y-Fx)/1e20, label="Meas-mod")

a = Dataset("/net/fluo/data1/group/oco2/oco2_L2DiaND_26780a_190715_B10004r_200618191413.h5")
g = a.group["RetrievalGeometry"]
s = a.group["SpectralParameters"]
i = findall(abs.(g["retrieval_latitude"][:].-41.736668).<0.0001)
fp_meas = s["measured_radiance"][:,i]
fp_mod  = s["modeled_radiance"][:,i]


module FalseColor
using Devectorize
include("colormaps.jl")

# ColorMap type for regularly spaced colormaps
immutable ColorMapRegular{T<:FloatingPoint}
colors::Array{T, 2}
end

# Lookup colors in regularly spaced colormaps - assumes x is within (0,size(cmap.data, 2)-1)
function lookupindex(cmap::ColorMapRegular, x)
xi = ceil(x)
w = xi - x
(xi, w)
end

# ColorMap type for irregularly spaced colormaps
immutable ColorMapIrregular{T<:FloatingPoint}
colors::Array{T, 2}
x::Vector{T}
w::Vector{T}
end

# Lookup colors in irregularly spaced colormaps - assumes x is within (0,size(cmap.data, 2)-1)
function lookupindex{T<:FloatingPoint}(cmap::ColorMapIrregular{T}, x::T)
xi = searchsortedlast(cmap.x, x)
@inbounds begin
w = (x - cmap.x[xi]) * cmap.w[xi]
end
(xi, w)
end

# Store a colormap plus its range
immutable ColorMapRange
cmap::Union(ColorMapRegular, ColorMapIrregular)
low
high
f

function ColorMapRange(cmap::Union(ColorMapRegular, ColorMapIrregular), low, high)
    T = eltype(cmap.colors)
    new(cmap, convert(T, low), convert(T, high), convert(T, (size(cmap.colors, 2) - 1)) / convert(T, high - low))
end
end

# A generic colormap constructor
function ColorMap{T<:FloatingPoint}(table::Array{T,2}, varargs...)
if size(table, 2) < 2
    error("Colormaps require at least 2 entries (2nd dimension");
end
n = size(table, 1)
if n != 3 && n != 4
    error("Colormaps require 2 or 3 channels (1st dimension)")
end
if n == 3
   return ColorMapRegular(table)
else
   w = vec(table[end,1:end-1])
   x = cumsum(w)
   x = [0.0, (x[1:end-1] * (length(x) / x[end]))]
   return ColorMapIrregular(table[1:end-1,:], x, 1.0./w)
end
end

# Determine the data range
function getrange{T<:Real,N}(data::AbstractArray{T,N}, varargs...)
# Check if the range was given as input
for i = 1:length(varargs)
    if (length(range[i]) == 2 && typeof(range[i][1]) <: Real && typeof(range[i][2]) <: Real)
        return (range[i][1], range[i][2])
    end
end
# Compute the range of the data
try
    return extremum(data)
catch
    # Support v0.2
    return (minimum(data), maximum(data))
end
end

function real2rgb(data::Real, cmap::ColorMapRange)
# Per-pixel conversion
@inbounds begin
data = convert(eltype(cmap.f), data);
if data <= cmap.low
    return cmap.cmap.colors[:,1]
elseif data >= cmap.high
    return cmap.cmap.colors[:,end]
else
    (xi, w) = lookupindex(cmap.cmap, (data - cmap.low) * cmap.f)
    xi_ = xi + 1
    w_ = 1.0 - w
    @devec out = cmap.cmap.colors[:,xi] .* w + cmap.cmap.colors[:,xi_] .* w_
    return out
end
end
end

function real2rgb!{T<:FloatingPoint,U<:Real}(out_::AbstractArray{T,3}, data::Array{U,2}, cmap::ColorMapRange)
# Check the output size
assert(size(out_, 1) == size(data, 1) && size(out_, 2) == size(data, 2) && size(out_, 3) == 3)
out = reshape(out_, length(data), 3)
# Go over each pixel
@inbounds begin
for i = 1:length(data)
    out[i,:] = real2rgb(data[i], cmap)
end
end
end

function real2rgb{T<:Real}(data::Array{T,2}, cmap::ColorMapRange)
out = Array(typeof(cmap.cmap.colors[1]), size(data, 1), size(data, 2), 3)
real2rgb!(out, data, cmap)
out
end

function real2rgb{T<:Real,U<:FloatingPoint}(data::AbstractArray{T,2}, cmap::Array{U,2}, varargs...)
# Check for a range
(low, high) = getrange(data, varargs...)
# Do the conversion
real2rgb(data, ColorMapRange(ColorMap(cmap), low, high))
end

function sc!{T<:FloatingPoint,U<:Number}(out::AbstractArray{T,3}, data::AbstractArray{U,2}, varargs...)
assert(size(out, 1) == size(data, 1) && size(out, 2) == size(data, 2) && size(out, 3) == 3 && ndims(data) <= 3)
(low, high) = getrange(data, varargs...)
real2rgb!(out, data, ColorMapRange(ColorMap(jet()), low, high))
end

function sc{T<:Number}(data::Array{T,2}, varargs...)
out = Array(Float64, size(data, 1), size(data, 2), 3, size(data)[4:end]...)
n = length(out) / (size(data, 1) * size(data, 2) * 3)
if n == 1
    sc!(out, data, varargs...)
else
    # Do multiple images at once
    for i = 1:n
        sc!(sub(out,:,:,:,i), sub(data,:,:,:,i), varargs...)
    end
end
out
end

# Functions to be visible outside
export real2rgb, real2rgb!, sc, sc!
end


using Test
using LinearAlgebra
using NCDatasets
using Random
const NC = NCDatasets
using CalibrateEDMF
using CalibrateEDMF.ModelTypes
using CalibrateEDMF.ReferenceModels
using CalibrateEDMF.ReferenceStats
using CalibrateEDMF.DistributionUtils
using CalibrateEDMF.TurbulenceConvectionUtils
using CalibrateEDMF.NetCDFIO
const CN = CalibrateEDMF.NetCDFIO
using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions

@testset "NetCDFIO_Diags" begin
    # Choose same SCM to speed computation
    data_dir = mktempdir()
    t_max = 2 * 3600.0
    namelist_args = [
        ("time_stepping", "t_max", t_max),
        ("time_stepping", "dt_max", 30.0),
        ("time_stepping", "dt_min", 20.0),
        ("grid", "dz", 150.0),
        ("grid", "nz", 20),
        ("stats_io", "frequency", 120.0),
    ]
    scm_dirs = repeat([joinpath(data_dir, "Output.Bomex.000000")], 2)
    kwargs_ref_model = Dict(
        :y_names => [["u_mean"], ["v_mean"]],
        :y_dir => scm_dirs,
        :scm_dir => scm_dirs,
        :case_name => repeat(["Bomex"], 2),
        :t_start => repeat([t_max - 2.0 * 3600], 2),
        :t_end => repeat([t_max], 2),
        :namelist_args => repeat([namelist_args], 2),
    )
    # Generate ref_stats
    ref_models = construct_reference_models(kwargs_ref_model)
    run_reference_SCM.(ref_models, overwrite = false, run_single_timestep = false)
    ref_stats = ReferenceStatistics(ref_models, y_type = SCM(), Σ_type = SCM())
    # Generate config
    config = Dict()
    config["process"] = Dict()
    config["process"]["N_ens"] = 10
    config["process"]["batch_size"] = nothing
    config["prior"] = Dict()
    config["prior"]["constraints"] = Dict("foo" => [bounded(0.0, 0.5)], "bar" => [bounded(0.0, 0.5)])
    config["reference"] = Dict()
    priors = construct_priors(config["prior"]["constraints"])
    ekp = EnsembleKalmanProcess(rand(2, 10), ref_stats.y, ref_stats.Γ, Inversion())
    N_ens = size(get_u_final(ekp), 2)
    diags = NetCDFIO_Diags(config, data_dir, ref_stats, N_ens, priors)

    # Test constructor
    @test isa(diags, NetCDFIO_Diags)
    @test isempty(diags.vars)

    # Test reference diagnostics
    io_reference(diags, ref_stats, ref_models)
    NC.Dataset(diags.filepath, "r") do root_grp
        ref_grp = root_grp.group["reference"]
        @test ref_grp["Gamma"] ≈ ref_stats.Γ
        @test ref_grp["Gamma_full"] ≈ ref_stats.Γ_full
        @test ref_grp["y"] ≈ ref_stats.y
        @test ref_grp["y_full"] ≈ ref_stats.y_full
    end

    # Test iteration-dependent diagnostics
    CN.init_iteration_io(diags)
    CN.open_files(diags)
    CN.write_iteration(diags)
    CN.close_files(diags)
    NC.Dataset(diags.filepath, "r") do root_grp
        ensemble_grp = root_grp.group["ensemble_diags"]
        @test length(ensemble_grp["iteration"]) == 2
    end
end