using BioFeatures
using DataFrames
using IntervalTrees

type Tile
  start::Int64
  stop::Int64
  value::Float64
end

function make_tiles( start, stop; num_tiles = 20, default_value = 0 )
     #bin_size
     #println("Making tiles $gene_start")
     length = stop - start
     bin_size = int64( length / num_tiles)
     tiles = Tile[]
     tile_start = start
     tile_end = tile_start + (bin_size - 1)
     for t = 1:num_tiles
         tile = Tile(tile_start,tile_end, default_value)
         #println(join([tile_start,tile_end, bin_size], "\t") )
         push!(tiles,tile)
         tile_start = tile_end + 1
         tile_end = tile_start + bin_size -1
     end
     return tiles
end

function make_tiles( feature::DataFrame; num_tiles = 20,default_value = 0  )
    start = int( feature[:start][1] )
    stop  = int( feature[:stop][1] )
    tiles = make_tiles(start,stop,num_tiles=num_tiles, default_value=default_value)
end

#  calculate_tile_value: Checks whether a tile overlaps any features and sets the value field of the tile if so
#  Input:
#     tile :: Tile
#     region_features:: Dictionary of IntervalTrees containing Features - indexed by the region name e.g. chr name
#  region_features should hold only one region in the features object
function calculate_tile_value( tile, region_features, value=1 )
    features_interval_trees = collect( values( region_features.features ) )
    #println(features_interval_trees[1])
    features_interval_tree = features_interval_trees[1]
    #println(tile)
    intersected_start = IntervalTrees.hasintersection( features_interval_tree, int32(tile.start) )
    if intersected_start # change this to complete interval, not just point
      tile.value = value
      #println("Overlap Tile")
    end
    intersected_stop = IntervalTrees.hasintersection( features_interval_tree, int32(tile.stop) )
    if intersected_stop # change this to complete interval, not just point
      tile.value = value
      #println("Overlap Tile")
    end
end

function set_tile_value( tile, region_features, type_field, value_map::Dict{ASCIIString,Float64} )
    features_interval_trees = collect(  values( region_features.features ) )
    #println(features_interval_trees[1])
    features_interval_tree = features_interval_trees[1]
    #println(tile)
    intersected = IntervalTrees.intersect( features_interval_tree, (int32(tile.start),int32(tile.stop)) )
    intersected_array = collect(intersected)
    if length(intersected_array) > 0
        current_type = ""
        for ia in intersected_array
            df = ia[2]
            #println(df)
            feature_type = df[type_field][1]
            value_to_assign = value_map[feature_type]
            #println(value_to_assign)
            if current_type == ""
                current_type=df[type_field][1]
            elseif current_type != df[type_field][1]
                value_to_assign = value_map["MIXED"]
            end
            tile.value = value_to_assign
        end
    end
    #return tile
    #intersected_start = IntervalTrees.hasintersection( features_interval_tree, int32(tile.start) )
    #if intersected_start # change this to complete interval, not just point
    #  tile.value = value
    #  println("Overlap Tile")
    #end
    #intersected_stop = IntervalTrees.hasintersection( features_interval_tree, int32(tile.stop) )
    #if intersected_stop # change this to complete interval, not just point
    #  tile.value = value
    #  println("Overlap Tile")
    #end
end

function set_tile_value1(tile, region_features)
    features_interval_trees = collect(  values( region_features.features ) )
    #println(features_interval_trees[1])
    features_interval_tree = features_interval_trees[1]
    #println(tile)
    intersected = IntervalTrees.intersect( features_interval_tree, (int32(tile.start),int32(tile.stop)) )
    intersected_array = collect(intersected)
    value_to_assign = 0
    tile_length = tile.stop - tile.start
    if length(intersected_array) > 0
        current_type = ""
        for ia in intersected_array
            df = ia[2]
            #println(df)
            #println(tile)
            if df[:start][1] <= tile.start && df[:stop][1] >= tile.stop
                value_to_assign = 1
                #println("Assigned $value_to_assign as feature is larger than tile")
                continue
            end

            if  df[:start][1] >= tile.start && df[:stop][1] <= tile.stop
                rat =  (df[:stop][1] - df[:start][1]) / tile_length
                #println("adding $rat")
                value_to_assign += rat # (df[:stop][1] - df[:start][1]) / tile_length )
                #println("assigned $value_to_assign as feature ovelaps completely with tile")
                continue
            end

            if df[:start][1] > tile.start && df[:start][1] < tile.stop
                rat =  ( tile.stop - df[:start][1] ) / tile_length
                #println("adding $rat")
                value_to_assign += rat
                #println("assigned $value_to_assign (feature start overlaps with tile")
                continue
            elseif df[:stop][1] > tile.start && df[:stop][1] < tile.stop
                rat = ( df[:stop][1] - tile.start  ) / tile_length
                #println("adding $rat")
                value_to_assign += rat #( (tile.stop - df[:stop][1]  ) / tile_length )
                #println("assigned $value_to_assign (feature start overlaps with tile")
                continue
            end

            #value_to_assign = value_map[feature_type]
            #println(value_to_assign)
            #if current_type == ""
            #    current_type=df[type_field][1]
            #elseif current_type != df[type_field][1]
            #    value_to_assign = value_map["MIXED"]
            #end
            #tile.value = value_to_assign
        end
        tile.value = value_to_assign
    end
end

# Adds the values of a contiguous set of tiles
# to a row in a matrix
function add_tiles_to_tiles_matrix(tiles_matrix,tiles)
    tiles_matrix_row = map( x -> x.value, tiles)
    #println("CONCAT TILES TO MATRIX")
    if(tiles_matrix == nothing)
       tiles_matrix = tiles_matrix_row
       println("TILES: $tiles_matrix")
    else
       tiles_matrix = hcat(tiles_matrix,tiles_matrix_row)
    end
    return tiles_matrix
end

function add_tiles_to_tiles_matrix(tiles_matrix,tiles,idx)
   i=1
   for tile in tiles
     tiles_matrix[idx,i] = tile.value
     i +=1
   end
   #tiles_matrix_row = map( x -> x.value, tiles)
   #println("CONCAT TILES TO MATRIX")
   # if(tiles_matrix == nothing)
   #    tiles_matrix = tiles_matrix_row
   #    println("TILES: $tiles_matrix")
   # else
   #    tiles_matrix = hcat(tiles_matrix,tiles_matrix_row)
   # end
   # return tiles_matrix
end

function add_tiles_to_predefined_tiles_matrix(tiles_matrix,tiles,idx)
   i=1
   for tile in tiles
     tiles_matrix[idx,i] = tile.value
     i +=1
   end
   #tiles_matrix_row = map( x -> x.value, tiles)
   #println("CONCAT TILES TO MATRIX")
   # if(tiles_matrix == nothing)
   #    tiles_matrix = tiles_matrix_row
   #    println("TILES: $tiles_matrix")
   # else
   #    tiles_matrix = hcat(tiles_matrix,tiles_matrix_row)
   # end
   # return tiles_matrix
end


@testset "handle" begin
    @jyro function m()
        x ~ Normal()
    end

    handle!(m, TraceHandler())
    @test length(m.handlers_stack) == 1
    @test isa(m.handlers_stack[1], TraceHandler)

    new_m = handle(m, LogJointHandler())
    @test length(m.handlers_stack) == 1
    @test length(new_m.handlers_stack) == 2
    @test isa(new_m.handlers_stack[1], TraceHandler)
    @test isa(new_m.handlers_stack[2], LogJointHandler)
end

@testset "trace" begin
    @jyro function m()
        x ~ Normal()
    end

    handle!(m, TraceHandler())
    t = m()
    @test haskey(t, :msgs)
    @test haskey(t[:msgs], :x)
    @test isa(t[:msgs][:x][:value], Float64)
end

@testset "logjoint" begin
    @jyro function m()
        x ~ Normal()
    end

    trace!(m)
    log_joint!(m)
    t = m()
    val = t[:msgs][:x][:value]

    @test logpdf(Normal(), val) == t[:logjoint]
end

@testset "condition" begin
    @jyro function m()
        x ~ Normal()
    end

    condition!(m, Dict(:x => 0.0))
    log_joint!(m)
    trace!(m)

    t = m()
    val = t[:msgs][:x][:value]

    @test val == 0.0
    @test t[:logjoint] == logpdf(Normal(), 0.0)
    @test t[:msgs][:x][:is_observed]
end

@testset "replay handler" begin
    # First sampling.
    @jyro function model1()
        x ~ Normal()
    end

    handle!(model1, TraceHandler())
    trace1 = model1()

    # Second sampling.
    @jyro function model2()
        x ~ Normal()
        y ~ Normal()
    end

    handle!(model2, ReplayHandler(trace1))
    handle!(model2, TraceHandler())
    trace2 = model2()

    @test trace1[:msgs][:x][:value] == trace2[:msgs][:x][:value]
end

@testset "escape handler" begin
    @jyro function model()
        x ~ Normal()
        {:y => 1} ~ Normal()
        z ~ Normal()
    end

    # Escape when we reach variable y.
    handle!(model, EscapeHandler(x -> x[:name] == (:y => 1)))
    handle!(model, TraceHandler())
    try
        t = model()
    catch e
        if isa(e, EscapeException)
            t = e.trace

            @test e.name == (:y => 1)
            @test haskey(t[:msgs], :x)
            @test isa(t[:msgs][:x][:value], Float64)
            @test t[:msgs][e.name][:value] == nothing
            @test !haskey(t[:msgs], :z)
        else
            throw(e)
        end
    end
end

@testset "queue" begin
    @jyro function model()
        a ~ Bernoulli(0.5)
        b ~ Bernoulli(0.5)
    end

    true_enum = Set([(true,true), (false,true), (true,false), (false,false)])
    empty_trace = Dict()
    empty_trace[:msgs] = Dict()
    q = [empty_trace]

    enum_model = queue(model, q)
    generated_enum = Set()
    while length(q) > 0
        t = enum_model()
        push!(generated_enum, (t[:msgs][:a][:value], t[:msgs][:b][:value]))
    end

    @test true_enum == generated_enum
end

@testset "generated handlers" begin
    @jyro function model()
        a ~ Normal()
    end

    conditioned_model = condition(model, Dict(:a => 1.0))
    trace!(model)
    t = model()
    cond_t = trace(conditioned_model)()

    # Traces from conditioned model should be different.
    @test cond_t[:msgs][:a][:value] == 1.0
    @test cond_t[:msgs][:a][:value] != t[:msgs][:a][:value] # This should technically never happen.
end


export PDText
export readpdf, tokenize

using CodecZlib
import Base.Unicode

mutable struct PDText
    str::String
    children::Vector{PDText}
    page::Int
    fcoord::Rectangle
    gcoord::Rectangle
end

function PDText(children::Vector{PDText})
    @assert all(c -> c.page == children[1].page, children)
    str = join(map(c -> c.str, children))
    fcoord = merge(map(c -> c.fcoord, children))
    gcoord = merge(map(c -> c.gcoord, children))
    PDText(str, children, children[1].page, fcoord, gcoord)
end

function leaves(text::PDText)
    l = PDText[]
    function f(t::PDText)
        if isempty(t.children)
            push!(l,t)
        else
            foreach(f, t.children)
        end
    end
    f(text)
    l
end

function readpdf(filepath::String)
    if endswith(filepath, ".txt")
        lines = open(readlines, filepath)
    elseif endswith(filepath, ".pdf")
        gzfilepath = "$filepath.0-3-1.txt.gz"
        if !isfile(gzfilepath)
            jar = download_jar("0.3.1")
            run(`java -classpath $jar paperai.pdfextract.PDFExtractor $filepath`)
        end
        lines = open(s -> readlines(GzipDecompressorStream(s)), gzfilepath)
    else
        throw("File error.")
    end

    texts = PDText[]
    for line in lines
        isempty(line) && continue
        items = split(line, '\t')
        page = parse(Int, items[1])
        str = Unicode.normalize(items[2], :NFKC)
        str[1] == '[' && str[end] == ']' && continue

        coord = map(x -> parse(Float64,x), split(items[3]," "))
        fcoord = Rectangle(coord...)
        coord = map(x -> parse(Float64,x), split(items[4]," "))
        gcoord = Rectangle(coord...)
        t = PDText(str, PDText[], page, fcoord, gcoord)
        push!(texts, t)
    end
    texts
end

function tokenize(texts::Vector{PDText})
    tokens = PDText[]
    buffer = [texts[1]]
    average = texts[1].fcoord.w
    for i = 2:length(texts)
        prev, curr = texts[i-1], texts[i]
        expected = prev.fcoord.x + prev.fcoord.w + 0.1*average
        if prev.page != curr.page || prev.fcoord.x > curr.fcoord.x || expected < curr.fcoord.x
            push!(tokens, PDText(buffer))
            buffer = [curr]
            average = curr.fcoord.w
        else
            push!(buffer, curr)
            average = (average + curr.fcoord.w) / 2
        end
    end
    isempty(buffer) || push!(tokens,PDText(buffer))
    tokens
end

function tokenize_line(texts::Vector{PDText})
end

export average_size
function average_size(texts::Vector{PDText})
    w = 0.0
    h = 0.0
    for t in texts
        w += t.fcoord.w
        h += t.fcoord.h
    end
    w /= length(texts)
    h /= length(texts)
end


using Caesar

robotId = "Munster"
sessionId = "Test"

fg = initfg(sessionname=sessionId, robotname=robotId)

# also add a PriorPose2 to pin the first pose at a fixed location
addVariable!(fg, :x0, Pose2) # , tags=[:POSE]
addFactor!(fg, [:x0], PriorPose2(MvNormal(zeros(3), 0.01*Matrix(LinearAlgebra.I, 3,3))) )

# Drive around in a hexagon
for i in 0:5
  psym = Symbol("x$i")
  nsym = Symbol("x$(i+1)")
  addVariable!(fg, nsym, Pose2) # , tags=[:VARIABLE;:POSE]
  addFactor!(fg, [psym;nsym], Pose2Pose2( MvNormal([10.0;0;pi/3], 0.01*Matrix(LinearAlgebra.I, 3,3)) ), autoinit=true )
  # Pose2Pose2_NEW(MvNormal([10.0;0;pi/3], diagm([0.1;0.1;0.1].^2)))
end

response = addVariable!(fg, :l1, Point2("LANDMARK"))


#####
##### Diagnostics
#####

#=
    io_dictionary_ref_state()
    io_dictionary_aux()

These functions return a dictionary whose
 - `keys` are the nc variable names
 - `values` are NamedTuples corresponding to
    - `dims` (`("z")`  or `("z", "t")`) and
    - `group` (`"reference"` or `"profiles"`)
=#

function io_dictionary_ref_state()
    DT = NamedTuple{(:dims, :group, :field), Tuple{Tuple{String}, String, Any}}
    cent_ref_state = center_ref_state # so that things nicely align :)
    io_dict = Dict{String, DT}(
        "ρ0_f" => (; dims = ("zf",), group = "reference", field = state -> face_ref_state(state).ρ0),
        "ρ0_c" => (; dims = ("zc",), group = "reference", field = state -> cent_ref_state(state).ρ0),
        "p0_f" => (; dims = ("zf",), group = "reference", field = state -> face_ref_state(state).p0),
        "p0_c" => (; dims = ("zc",), group = "reference", field = state -> cent_ref_state(state).p0),
        "α0_f" => (; dims = ("zf",), group = "reference", field = state -> face_ref_state(state).α0),
        "α0_c" => (; dims = ("zc",), group = "reference", field = state -> cent_ref_state(state).α0),
    )
    return io_dict
end

#! format: off
# TODO: use a better name, this exports fields from the prognostic state, and the aux state.
function io_dictionary_aux()
    DT = NamedTuple{(:dims, :group, :field), Tuple{Tuple{String, String}, String, Any}}
    io_dict = Dict{String, DT}(
        "updraft_area" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).area),
        "updraft_ql" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).q_liq),
        "updraft_qi" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).q_ice),
        "updraft_RH" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).RH),
        "updraft_qt" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).q_tot),
        "updraft_w" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_bulk(state).w),
        "updraft_temperature" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).T),
        "updraft_thetal" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).θ_liq_ice),
        "updraft_buoyancy" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).buoy),
        "H_third_m" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).H_third_m),
        "W_third_m" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).W_third_m),
        "QT_third_m" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).H_third_m),
        "cloud_fraction" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).cloud_fraction), # was this "cloud_fraction_mean"?
        "buoyancy_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).buoy),
        "temperature_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).T),
        "RH_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).RH),
        "s_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).s),
        "ql_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).q_liq),
        "qi_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).q_ice),
        "tke_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).tke),
        "Hvar_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).Hvar),
        "QTvar_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).QTvar),
        "HQTcov_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).HQTcov),
        "u_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_prog_grid_mean(state).u),
        "v_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_prog_grid_mean(state).v),
        "w_mean" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_prog_grid_mean(state).w),
        "qt_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_prog_grid_mean(state).q_tot),
        "thetal_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_prog_grid_mean(state).θ_liq_ice),
        "eddy_viscosity" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).KM),
        "eddy_diffusivity" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).KH),
        "env_tke" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).tke),
        "env_Hvar" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).Hvar),
        "env_QTvar" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).QTvar),
        "env_HQTcov" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).HQTcov),

        "tke_buoy" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.buoy),
        "tke_pressure" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.press),
        "tke_dissipation" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.dissipation),
        "tke_entr_gain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.entr_gain),
        "tke_detr_loss" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.detr_loss),
        "tke_shear" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.shear),
        "tke_interdomain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).tke.interdomain),

        "Hvar_dissipation" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).Hvar.dissipation),
        "Hvar_entr_gain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).Hvar.entr_gain),
        "Hvar_detr_loss" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).Hvar.detr_loss),
        "Hvar_interdomain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).Hvar.interdomain),
        "Hvar_rain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).Hvar.rain_src),
        "Hvar_shear" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).Hvar.shear),

        "QTvar_dissipation" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).QTvar.dissipation),
        "QTvar_entr_gain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).QTvar.entr_gain),
        "QTvar_detr_loss" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).QTvar.detr_loss),
        "QTvar_shear" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).QTvar.shear),
        "QTvar_rain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).QTvar.rain_src),
        "QTvar_interdomain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).QTvar.interdomain),

        "HQTcov_rain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).HQTcov.rain_src),
        "HQTcov_dissipation" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).HQTcov.dissipation),
        "HQTcov_entr_gain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).HQTcov.entr_gain),
        "HQTcov_detr_loss" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).HQTcov.detr_loss),
        "HQTcov_shear" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).HQTcov.shear),
        "HQTcov_interdomain" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment_2m(state).HQTcov.interdomain),

        "env_w" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_environment(state).w),
        "env_qt" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).q_tot),
        "env_ql" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).q_liq),
        "env_qi" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).q_ice),
        "env_area" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).area),
        "env_temperature" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).T),
        "env_RH" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).RH),
        "env_thetal" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).θ_liq_ice),
        "env_cloud_fraction" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_environment(state).cloud_fraction),
        "massflux_s" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_grid_mean(state).massflux_s),
        "diffusive_flux_s" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_grid_mean(state).diffusive_flux_s),
        "total_flux_s" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_grid_mean(state).massflux_s .+ face_aux_grid_mean(state).diffusive_flux_s),

        "qr_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_prog_precipitation(state).q_rai),
        "qs_mean" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_prog_precipitation(state).q_sno),

        "mixing_length" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).mixing_length),

        "updraft_cloud_fraction" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).cloud_fraction),

        "updraft_qt_precip" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).qt_tendency_precip_formation),
        "updraft_thetal_precip" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_bulk(state).θ_liq_ice_tendency_precip_formation),

        "massflux_tendency_h" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).massflux_tendency_h),
        "massflux_tendency_qt" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).massflux_tendency_qt),
        "diffusive_tendency_h" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).diffusive_tendency_h),
        "diffusive_tendency_qt" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).diffusive_tendency_qt),

        "total_flux_h" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_turbconv(state).diffusive_flux_h .+ face_aux_turbconv(state).massflux_h),
        "total_flux_qt" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_turbconv(state).diffusive_flux_qt .+ face_aux_turbconv(state).massflux_qt),

        "ed_length_scheme" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).mls),
        "mixing_length_ratio" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).ml_ratio),
        "entdet_balance_length" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_turbconv(state).l_entdet),

        "rad_dTdt" => (; dims = ("zc", "t"), group = "profiles", field = state -> center_aux_grid_mean(state).dTdt_rad),
        "rad_flux" => (; dims = ("zf", "t"), group = "profiles", field = state -> face_aux_grid_mean(state).f_rad),

    )
    return io_dict
end
#! format: on

function initialize_io(io_dict::Dict, Stats::NetCDFIO_Stats)
    for var_name in keys(io_dict)
        add_field(Stats, var_name; dims = io_dict[var_name].dims, group = io_dict[var_name].group)
    end
end

function io(io_dict::Dict, Stats::NetCDFIO_Stats, state)
    for var in keys(io_dict)
        write_field(Stats, var, io_dict[var].field(state); group = io_dict[var].group)
    end
end


module TestNewick

using Phylo
using Test

@testset "newick and nexus parsing" begin
    @testset "For $TreeType" for TreeType in [NamedBinaryTree]
        @test length(nodeiter(parsenewick("((,),(,));"))) == 7
        @test ["where", "when it's good", "Not mine", "MyLeaf"] ⊆
            nodenameiter(parsenewick("""((MyLeaf,"when it's good"),
                                             ('Not mine',where));""",
                                     TreeType))
        tree =
            parsenewick("((MyLeaf[&Real=23,'Not real'={5,4}]:4.0,)Parent,(,));",
                        TreeType)
        branches = branchfilter(tree) do tree, branch
            return getnodename(tree, src(tree, branch)) == "Parent" &&
                getnodename(tree, dst(tree, branch)) == "MyLeaf"
        end
        @test length(branches) == 1
        @test getlength(tree, first(branches)) ≈ 4.0
        @test getnodedata(tree, "MyLeaf")["Real"] == 23
        @test 5 ∈ getnodedata(tree, "MyLeaf")["Not real"]
        @test 4 ∈ getnodedata(tree, "MyLeaf")["Not real"]
        @test_throws Exception parsenewick("((,),(,)));", TreeType)
        @test_throws Exception parsenewick("((,),(,))", TreeType)
        @test_throws Exception parsenewick("((,),(,);", TreeType)
        @test_throws Exception parsenewick("((MyLeaf:-4.0,)Parent,(,));",
                                           TreeType)
        tree = open(f -> parsenewick(f, TreeType), Phylo.path("H1N1.newick"))
        @test nleaves(tree) == 507
        @test ntrees(tree) == 1
        treex = open(f -> parsenexus(f, TreeType), Phylo.path("H1N1.trees"))
        @test nleaves(tree) == nleaves(treex) ==
            nleaves(treex["TREE1"]) == nleaves(treex["TREE2"])
        @test ntrees(treex) == 2
    end

    @testset "For $TreeType" for TreeType in
        [NamedTree, RootedTree, ManyRootTree]
        @test nnodes(parsenewick("((,),(,,));", TreeType)) == 8
        @test ["where", "when it's good", "Not mine", "MyLeaf", "next"] ⊆
            nodenameiter(parsenewick("""((MyLeaf,"when it's good",next),
                                             ('Not mine',where));""",
                                     TreeType))
        tree = parsenewick("((MyLeaf:4.0,)Parent,());", TreeType)
        branches = branchfilter(tree) do tree, branch
            return getnodename(tree, src(tree, branch)) == "Parent" &&
                getnodename(tree, dst(tree, branch)) == "MyLeaf"
        end
        @test length(branches) == 1
        @test getlength(tree, first(branches)) ≈ 4.0
        @test_throws Exception parsenewick("((,),(,)));", TreeType)
        @test_throws Exception parsenewick("((,),(,))", TreeType)
        @test_throws Exception parsenewick("((,),(,);", TreeType)
        tree = open(f -> parsenewick(f, TreeType), Phylo.path("H1N1.newick"))
        @test nleaves(tree) == 507
        @test ntrees(tree) == 1
        treex = open(f -> parsenexus(f, TreeType), Phylo.path("H1N1.trees"))
        @test nleaves(tree) == nleaves(treex) ==
            nleaves(treex["TREE1"]) == nleaves(treex["TREE2"])
        @test ntrees(treex) == 2
    end
end

end


include("IOM/src/share/LogSystem.jl")
include("IOM/src/share/PolelikeCoordinate.jl")

include("IOM/src/models/EMOM/ENGINE_EMOM.jl")
include("IOM/src/driver/driver_working.jl")

using MPI
using CFTime
using Dates
using ArgParse
using .PolelikeCoordinate
using .LogSystem

function parse_commandline()

    s = ArgParseSettings()

    @add_arg_table s begin
        "--read-restart"
            help = "Configuration file."
            arg_type = Bool

        "--stop-n"
            help = "Core of the model."
            arg_type = Int64

        "--time-unit"
            help = "."
            arg_type = String
    end

    return parse_args(s)
end

parsed = parse_commandline()

t_start = DateTimeNoLeap(1, 1, 1)

time_unit = Dict(
    "year" => Dates.Year,
    "month" => Dates.Month,
    "day"   => Dates.Day,
)[parsed["time-unit"]]

t_simulation = time_unit(parsed["stop-n"])
t_end = t_start + t_simulation
MPI.Init()
comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)
Δt = Second(86400)

include("config.jl")

if rank == 0
    gf = PolelikeCoordinate.CurvilinearSphericalGridFile(
        config["MODEL_CORE"]["domain_file"],
        R  = 6371229.0,
        Ω  = 2π / (86400 / (1 + 365/365)),
    )

end

coupler_funcs = (
    master_before_model_init = function()
        return parsed["read-restart"], t_start
    end,

    master_after_model_init! = function(OMMODULE, OMDATA)
        writeLog("[Coupler] After model init")

    end,
    master_before_model_run! = function(OMMODULE, OMDATA)
        writeLog("[Coupler] Before model run")
        writeLog("[Coupler] This is where flux exchange happens.")

        global comm, rank, t_end

        is_master = rank == 0

        # Test if t_end is reached
        return_values = nothing
        if is_master
            t_end_reached = OMDATA.clock.time >= t_end

            if t_end_reached
                return_values = ( :END, Δt, t_end_reached )
            else
                return_values = ( :RUN, Δt, t_end_reached )
            end
        end

        # Deal with coupling
        if is_master

            writeLog("[Coupler] Need to broadcast forcing fields.")
            # compute flux
            lat = gf.yc
            lon = gf.xc

#            OMDATA.x2o["SWFLX"][1, :, :] .= - (cos.(deg2rad.(lat).+0.1) .+ 1) / 2 .* (sin.(deg2rad.(lon)) .+ 1)/2 * 100.0
#            OMDATA.x2o["TAUX_east"][1, :, :]   .= 0.2 * (cos.(deg2rad.(lat).+0.1) .+ 1) / 2
#            OMDATA.x2o["TAUY_north"][1, :, :]  .= 0.1 * (sin.(deg2rad.(lon)) .+ 1) / 2
#cos.(deg2rad.(lat))

        end

        return return_values

    end,

    master_after_model_run! = function(OMMODULE, OMDATA)
        writeLog("[Coupler] After model run")
    end,

    master_finalize! = function(OMMODULE, OMDATA)
        writeLog("[Coupler] Finalize")

    end
)

runModel(
    ENGINE_EMOM,
    coupler_funcs,
    config,
)


module JuMPComplementarity

using MathProgBase, AmplNLWriter
importall MathProgBase.SolverInterface

export AmplComplementaritySolver

struct AmplComplementaritySolver <: MathProgBase.AbstractMathProgSolver
    solver_command::String
    options::Vector{String}
    filename::String
end

import AmplNLWriter.NonlinearModel
import AmplNLWriter.loadproblem!

macro complementarity(m::Model, var::Expr, g::Expr)

#=
function printModel(m::)

end
=#

end # module


# Low-level generated functions for working with grids

# _update_buffers => NTuple{N,SMatrix}
# Generate an SArray from the main array and the last SArray
@generated function _update_buffers(
    buffers::Tuple, src::AbstractArray{T}, ::Val{R}, i, j
) where {T,R}
    B = 2R; S = 2R + 1; L = S^2
    newvals = Expr[]
    for n in 0:2B-1
        push!(newvals, :(@inbounds src[i + $n, j + 2R]))
    end
    newbuffers = Expr(:tuple)
    for b in 1:B
        bufvals = Expr(:tuple)
        for n in S+1:L
            push!(bufvals.args, :(@inbounds buffers[$b][$n]))
        end
        for n in b:b+B
            push!(bufvals.args, newvals[n])
        end
        push!(newbuffers.args, :(SArray{Tuple{$S,$S},$T,2,$L}($bufvals)))
    end
    return quote
        return $newbuffers
    end
end

# _initialise_buffers => NTuple{N,SMatrix}
# Generate an SArray from the main array
@generated function _initialise_buffers(src::AbstractArray{T}, ::Val{R}, i, j) where {T,R}
    B = 2R; S = 2R + 1; L = S^2
    columns = []
    zerocol = Expr[]
    for r in 1:2B
        push!(zerocol, :(zero(T)))
    end
    push!(columns, zerocol)
    for c in 0:S-2
        newcol = Expr[]
        for r in 0:2B-1
            push!(newcol, :(@inbounds src[i + $r, j + $c]))
        end
        push!(columns, newcol)
    end
    newbuffers = Expr(:tuple)
    for b in 1:B
        bufvals = Expr(:tuple)
        for c in 1:S, r in b:b+B
            exp = columns[c][r]
            push!(bufvals.args, exp)
        end
        push!(newbuffers.args, :(SArray{Tuple{$S,$S},$T,2,$L}($bufvals)))
    end
    return quote
        return $newbuffers
    end
end

# _getwindow => SMatrix
# Get a single window square from an array, as an SMatrix.
# We use this on GPUs to get a neighborhood window from the main
# grid, which is 10x or more faster than using a view.
# We could possible just use this instead of _update_buffers
# for the sake of simplicity, with some performance loss.
@generated function _getwindow(tile::AbstractArray{T}, ::Neighborhood{R}, i, j) where {T,R}
    S = 2R+1
    L = S^2
    vals = Expr(:tuple)
    for jj in 1:S, ii in 1:S
        push!(vals.args, :(@inbounds tile[$ii + i - 1, $jj + j - 1]))
    end
    return :(SMatrix{$S,$S,$T,$L}($vals))
end

# _readcell
# Returns a single value or NamedTuple of values.
# This occurs for every cell for every rule, so has to be very fast.
@generated function _readcell(data, ::K, I...) where K<:Tuple
    expr = Expr(:tuple)
    keys = map(_unwrap, Tuple(K.parameters))
    for (i, k) in enumerate(keys)
        push!(expr.args, :(@inbounds source(data[$(QuoteNode(k))])[I...]))
    end
    return quote
        keys = $keys
        vals = $expr
        NamedTuple{keys,typeof(vals)}(vals)
    end
end
@inline function _readcell(data::AbstractSimData, ::Val{K}, I...) where K
    @inbounds source(data[K])[I...]
end

# _writecell! => nothing
# Write values to grid/s at index `I`.
# This occurs for every cell for every rule, so has to be very fast.
@generated function _writecell!(data, ::Val{<:CellRule}, wkeys::K, vals::Union{Tuple,NamedTuple}, I...) where K<:Tuple
    expr = Expr(:block)
    keys = map(_unwrap, Tuple(K.parameters))
    for (i, k) in enumerate(keys)
        # MUST write to source(grid) - CellRule doesn't switch grids
        push!(expr.args, :(@inbounds source(data[$(QuoteNode(k))])[I...] = vals[$i]))
    end
    push!(expr.args, :(nothing))
    return expr
end
@inline function _writecell!(data, ::Val{<:CellRule}, wkeys::Val{K}, val, I...) where K
    # MUST write to source(grid) - CellRule doesn't switch grids
    @inbounds source(data[K])[I...] = val
    return nothing
end
@generated function _writecell!(data, ::Val, wkeys::K, vals::Union{Tuple,NamedTuple}, I...) where K<:Tuple
    expr = Expr(:block)
    keys = map(_unwrap, Tuple(K.parameters))
    for (i, k) in enumerate(keys)
        # MUST write to dest(grid) here, not grid K
        # setindex! has overrides for the grid
        push!(expr.args, :(@inbounds dest(data[$(QuoteNode(k))])[I...] = vals[$i]))
    end
    push!(expr.args, :(nothing))
    return expr
end
@inline function _writecell!(data, ::Val, wkeys::Val{K}, val, I...) where K
    # MUST write to dest(grid) here, not grid K
    # setindex! has overrides for the grid
    @inbounds dest(data[K])[I...] = val
    return nothing
end

# _getreadgrids => Union{ReadableGridData,Tuple{ReadableGridData,Vararg}}
# Retrieves `GridData` from a `SimData` object to match the requirements of a `Rule`.
# Returns a `Tuple` holding the key or `Tuple` of keys, and grid or `Tuple` of grids.
@generated function _getreadgrids(::Rule{R,W}, data::AbstractSimData) where {R<:Tuple,W}
    Expr(:tuple,
        Expr(:tuple, (:(Val{$(QuoteNode(key))}()) for key in R.parameters)...),
        Expr(:tuple, (:(data[$(QuoteNode(key))]) for key in R.parameters)...),
    )
end
@generated function _getreadgrids(::Rule{R,W}, data::AbstractSimData) where {R,W}
    :((Val{$(QuoteNode(R))}(), data[$(QuoteNode(R))]))
end

# _getwritegrids => Union{WritableGridData,Tuple{WritableGridData,Vararg}}
# Retrieves `GridData` from a `SimData` object to match the requirements of a `Rule`.
# Returns a `Tuple` holding the key or `Tuple` of keys, and grid or `Tuple` of grids.
@generated function _getwritegrids(::Rule{R,W}, data::AbstractSimData) where {R,W<:Tuple}
    Expr(:tuple,
        Expr(:tuple, (:(Val{$(QuoteNode(key))}()) for key in W.parameters)...),
        Expr(:tuple, (:(WritableGridData(data[$(QuoteNode(key))])) for key in W.parameters)...),
    )
end
@generated function _getwritegrids(::Rule{R,W}, data::AbstractSimData) where {R,W}
    :((Val{$(QuoteNode(W))}(), WritableGridData(data[$(QuoteNode(W))])))
end

# _combinegrids => AbstractSimData
# Combine grids into a new NamedTuple of grids depending
# on the read and write keys required by a rule.
@inline function _combinegrids(data::AbstractSimData, wkeys, wgrids)
    allkeys = map(Val, keys(data))
    allgrids = values(data)
    _combinegrids(data, allkeys, allgrids, wkeys, wgrids)
end
@inline function _combinegrids(data::AbstractSimData, rkeys, rgrids, wkeys, wgrids)
    @set data.grids = _combinegrids(rkeys, rgrids, wkeys, wgrids)
end
@inline function _combinegrids(rkey, rgrids, wkey, wgrids)
    _combinegrids((rkey,), (rgrids,), (wkey,), (wgrids,))
end
@inline function _combinegrids(rkey, rgrids, wkeys::Tuple, wgrids::Tuple)
    _combinegrids((rkey,), (rgrids,), wkeys, wgrids)
end
@inline function _combinegrids(rkeys::Tuple, rgrids::Tuple, wkey, wgrids)
    _combinegrids(rkeys, rgrids, (wkey,), (wgrids,))
end
@generated function _combinegrids(rkeys::Tuple, rgrids::Tuple, wkeys::Tuple, wgrids::Tuple)
    rkeys = _vals2syms(rkeys)
    wkeys = _vals2syms(wkeys)
    keysexp = Expr(:tuple, QuoteNode.(wkeys)...)
    dataexp = Expr(:tuple, :(wgrids...))

    for (i, key) in enumerate(rkeys)
        if !(key in wkeys)
            push!(dataexp.args, :(rgrids[$i]))
            push!(keysexp.args, QuoteNode(key))
        end
    end

    return quote
        keys = $keysexp
        vals = $dataexp
        NamedTuple{keys,typeof(vals)}(vals)
    end
end

# _replacegrids => AbstractSimData
# Replace grids in a NamedTuple with new grids where required.
function _replacegrids(data::AbstractSimData, newkeys, newgrids)
    @set data.grids = _replacegrids(grids(data), newkeys, newgrids)
end
@generated function _replacegrids(allgrids::NamedTuple, newkeys::Tuple, newgrids::Tuple)
    newkeys = map(_unwrap, newkeys.parameters)
    allkeys = allgrids.parameters[1]
    expr = Expr(:tuple)
    for key in allkeys
        if key in newkeys
            i = findfirst(k -> k == key, newkeys)
            push!(expr.args, :(newgrids[$i]))
        else
            push!(expr.args, :(allgrids.$key))
        end
    end

    return quote
        vals = $expr
        NamedTuple{$allkeys,typeof(vals)}(vals)
    end
end
@generated function _replacegrids(allgrids::NamedTuple, newkey::Val, newgrid::GridData)
    newkey = _unwrap(newkey)
    allkeys = allgrids.parameters[1]
    expr = Expr(:tuple)
    for key in allkeys
        if key == newkey
            push!(expr.args, :(newgrid))
        else
            push!(expr.args, :(allgrids.$key))
        end
    end

    return quote
        vals = $expr
        NamedTuple{$allkeys,typeof(vals)}(vals)
    end
end

# _vals2syms => Union{Symbol,Tuple}
# Get symbols from a Val or Tuple type
@inline _vals2syms(x::Type{<:Tuple}) = map(v -> _vals2syms(v), x.parameters)
@inline _vals2syms(::Type{<:Val{X}}) where X = X


# This file is a part of NeighborhoodApproximationIndex.jl

import SimilaritySearch: search, getpools, index!
using InvertedFiles, Intersections, KCenters, StatsBase, Parameters, LinearAlgebra
export KnrIndex, search, KnrOrderingStrategies, DistanceOrdering, InternalDistanceOrdering, DistanceOnTopKOrdering

mutable struct KnrOpt
    ksearch::Int32
end

abstract type KnrOrderingStrategy end

"""
    struct DistanceOrdering <: KnrOrderingStrategy end

Used as `ordering` parameter of `KnrIndex` specifies that each object `u` found by the inverted index will be evaluated against a given query.
This is the default ordering strategy.
"""
struct DistanceOrdering <: KnrOrderingStrategy end

"""
    struct InternalDistanceOrdering <: KnrOrderingStrategy end

Used as `ordering` parameter of `KnrIndex` specifies that the internal distance of the underlying inverted index will be used to ordering the `k` nearest neighbors.
Useful when the internal distance is quite representative of the real distance metric or whenever speed is the major objective.
"""
struct InternalDistanceOrdering <: KnrOrderingStrategy end

"""
    struct DistanceOnTopKOrdering <: KnrOrderingStrategy
        top::Int
    end

Used as `ordering` parameter of `KnrIndex` specifies that only the `top` elements are evaluated against the distance metric. Useful for very costly distance functions. If you are in doubt please use `DistanceOrdering` instead.
"""
mutable struct DistanceOnTopKOrdering <: KnrOrderingStrategy
    top::Int
end

"""
    struct KnrIndex <: AbstractSearchContext

The K nearest references inverted index

# Parameters

- `dist`: the distance function of the index
- `db`: the database of indexed objects
- `centers`: a search index for a set of references
- `invfile`: an inverted file data structure
- `kbuild`: the number of references to be computed and stored by each indexed object
- `ordering`: specifies how the index performs final `k` nn selection
- `opt`: the parameters to be optimized by `optimize!`
"""
struct KnrIndex{
            DistType<:SemiMetric,
            DataType<:AbstractDatabase,
            CentersIndex<:AbstractSearchContext,
            InvIndexType<:AbstractInvertedFile,
            OrderingType<:KnrOrderingStrategy
        } <: AbstractSearchContext
    dist::DistType
    db::DataType
    centers::CentersIndex
    invfile::InvIndexType
    kbuild::Int32
    ordering::OrderingType
    opt::KnrOpt
end

@inline Base.length(idx::KnrIndex) = length(idx.invfile)
Base.show(io::IO, idx::KnrIndex) = print(io, "{$(typeof(idx)) centers=$(typeof(idx.centers)), n=$(length(idx)), ordering=$(idx.ordering)}")

"""
    struct KnrCaches
        enc
    end

Caches used for `KnrIndex` (one per thread)

# Properties
- `enc`: `KnnResult` for encoding purposes
"""
struct KnrCaches
    enc::KnnResult
end

const GlobalKnrCachesPool = Vector{KnrCaches}(undef, 0)

function __init__()
    n = Threads.nthreads()

    while length(GlobalKnrCachesPool) < n
        push!(GlobalKnrCachesPool, KnrCaches(KnnResult(10)))
    end
end

getpools(idx::KnrIndex) = GlobalKnrCachesPool

function getencodeknnresult(k::Integer, pools::Vector{KnrCaches})
    reuse!(pools[Threads.threadid()].enc, k)
end

"""
    push!(idx::KnrIndex, obj; pools=getpools(idx), encpools=getpools(idx.centers))

Inserts `obj` into the indexed
"""
function Base.push!(idx::KnrIndex, obj; pools=getpools(idx), encpools=getpools(idx.centers))
    res = getencodeknnresult(idx.kbuild, pools)
    search(idx.centers, obj, res; pools=encpools)
    push!(idx.invfile, res)
    idx
end

"""
    get_parallel_block(n)

An heuristic to compute the `parallel_block` with respect with the number of elements to insert
"""
get_parallel_block(n) = min(n, 8 * Threads.nthreads())

"""
    append!(idx::KnrIndex, db; <kwargs>)

Appends all items in the database `db` into the index

# Arguments
- `idx`: the index structure
- `db`: the objects to be appended

# Keyword arguments
- `parallel_block`: the number of elements to be inserted in parallel
- `pools`: unused argument
- `verbose`: controls the verbosity of the procedure
"""
function Base.append!(idx::KnrIndex, db;
        parallel_block=get_parallel_block(length(db)),
        pools=getpools(idx),
        verbose=true
    )

    append!(idx.db, db)
    index!(idx; parallel_block, pools, verbose)
end

"""
    index!(idx::KnrIndex; parallel_block=get_parallel_block(length(idx.db)), pools=nothing, verbose=true)

Indexes all non indexed items in the database

# Arguments

- `idx`: the index structure

# Keyword arguments
- `parallel_block`: the number of elements to be inserted in parallel
- `pools`: unused parameter
- `verbose`: controls verbosity of the procedure

"""
function index!(idx::KnrIndex; parallel_block=get_parallel_block(length(idx.db)), pools=nothing, verbose=true)
    sp = length(idx) + 1
    n = length(idx.db)
    E = [KnnResult(idx.kbuild) for _ in 1:parallel_block]
    while sp < n
        ep = min(n, sp + parallel_block - 1)
        verbose && println(stderr, "$(typeof(idx)) appending chunk ", (sp=sp, ep=ep, n=n), " ", Dates.now())

        Threads.@threads for i in sp:ep
            begin
                res = reuse!(E[i - sp + 1], idx.kbuild)
                search(idx.centers, idx[i], res)
            end
        end

        append!(idx.invfile, VectorDatabase(E), ep-sp+1; parallel_block)
        sp = ep + 1
    end
end

"""
    KnrIndex(
        dist::SemiMetric,
        db::AbstractDatabase;
        invfiletype=BinaryInvertedFile,
        invfiledist=JaccardDistance(),
        initial=:dnet,
        maxiters=0,
        refs=references(dist, db; initial),
        centers=nothing,
        kbuild=3,
        ksearch=1,
        centersrecall::AbstractFloat=length(db) > 10^3 ? 0.95 : 1.0,
        ordering=DistanceOrdering(),
        pools=nothing,
        parallel_block=get_parallel_block(length(db)),
        verbose=false
    )

A convenient function to create a `KnrIndex`, it uses several default arguments. After the construction, use [`optimize!`](@ref) to adjust the index to some performance.

# Arguments
- `dist`: Distance object (a `SemiMetric` object, see `Distances.jl`)
- `db`: The database of objects to be indexed.

# Keyword arguments
- `invfiletype`: the type of the underlying inverted file (`BinaryInvertedFile` or `WeightedInvertedFile`)
- `invfiledist`: the distance of the underlying inverted file (see [`InvertedFiles.jl`](https://github.com/sadit/InvertedFiles.jl) package)
- `centers`: The index used for centers/references, if `centers === nothing` then a sample of `db` will be used.
- `initial`: indicates how references are selected, only used if `refs` will be computed; see [`references`](@ref) for more detail.
- `maxiters`: how many iterations of the Lloyd algorithm are applied to initial references, only used if `refs` will be computed; see [`references`](@ref) for more detail.
- `refs`: the set of reference, only used if `centers === nothing`
- `centersrecall`: used when `centers === nothing`; if `centersrecall == 1` then it will create an exact index on `refs` or an approximate if `0 < centersrecall < 1`
- `kbuild`: the number of references to compute and store on construction
- `ksearch`: the number of references to compute on searching
- `ordering`: the ordering strategy
- `pools`: an object with preallocated caches specific for `KnrIndex`, if `pools=nothing` it will use default caches.
- `parallel_block` Parallel construction works on batches, this is the size of these blocks
- `verbose` true if you want to see messages
"""
function KnrIndex(
        dist::SemiMetric,
        db::AbstractDatabase;
        invfiletype=BinaryInvertedFile,
        invfiledist=JaccardDistance(),
        initial=:dnet,
        maxiters=0,
        refs=references(dist, db; initial, maxiters),
        centers=nothing,
        kbuild::Integer=3,
        ksearch::Integer=1,
        centersrecall::AbstractFloat=length(db) > 10^3 ? 0.95 : 1.0,
        ordering::KnrOrderingStrategy=DistanceOrdering(),
        pools=nothing,
        parallel_block::Integer=get_parallel_block(length(db)),
        verbose=false
    )
    kbuild = convert(Int32, kbuild)
    ksearch = convert(Int32, ksearch)

    if centers === nothing
        if centersrecall == 1.0
            centers = ExhaustiveSearch(; db=refs, dist)
        else
            0 < centersrecall < 1 || throw(ArgumentError("the expected recall for centers index should be 0 < centersrecall < 0"))
            centers = SearchGraph(; db=refs, dist, verbose)
            index!(centers; parallel_block)
            optimize!(centers, OptimizeParameters(kind=MinRecall(centersrecall)))
        end
    end

    invfile = invfiletype(length(centers), invfiledist)
    idx = KnrIndex(dist, db, centers, invfile, kbuild, ordering, KnrOpt(ksearch))
    pools = pools === nothing ? getpools(idx) : pools
    index!(idx; parallel_block, pools, verbose)
    idx
end


module Model

using StaticArrays

# parameter types
export Image, SkyPatch, PsfComponent,
       GalaxyComponent, GalaxyPrototype,
       PriorParams, UnconstrainedParams,
       CanonicalParams, BrightnessParams, StarPosParams,
       GalaxyPosParams, GalaxyShapeParams,
       PsfParams, RawPSF, CatalogEntry,
       init_source, ParamSet

# functions
export align

# constants
export band_letters, D, Ia, B,
       galaxy_prototypes, prior,
       shape_standard_alignment,
       brightness_standard_alignment,
       gal_shape_alignment,
       ids_names, ids_free_names,
       ids, ids_free, star_ids, gal_ids, gal_shape_ids, psf_ids, bids

import Base.convert
import Base.+
import Distributions
import FITSIO, WCS
import WCS.WCSTransform
import ForwardDiff
import ..Log

import Base.length

const cfgdir = joinpath(Pkg.dir("Celeste"), "cfg")

const band_letters = ['u', 'g', 'r', 'i', 'z']

# The number of bands (colors).
const B = length(band_letters)

include("model/light_source_model.jl")
include("model/psf_model.jl")
include("model/image_model.jl")
include("model/param_set.jl")
include("model/imaged_sources.jl")
include("model/wcs_utils.jl")

import ..SensitiveFloats: SensitiveFloat, clear!
include("bivariate_normals.jl")
include("model/fsm_util.jl")
include("model/log_prob.jl")

end  # module


function cartesian_iter(state, res, A, Asize)
    for i in CartesianIndices(Asize)
        res[i] = A[i]
    end
    return
end

function clmap!(ctx, f, out, b)
    i = linear_index(ctx) # get the kernel index it gets scheduled on
    out[i] = f(b[i])
    return
end

function ntuple_test(ctx, result, ::Val{N}) where N
    result[1] = ntuple(Val(N)) do i
        Float32(i) * 77f0
    end
    return
end

function ntuple_closure(ctx, result, ::Val{N}, testval) where N
    result[1] = ntuple(Val(N)) do i
        Float32(i) * testval
    end
    return
end

@testsuite "base" (AT, eltypes)->begin
    @testset "copyto!" begin
        x = fill(0f0, (10, 10))
        y = rand(Float32, (20, 10))
        a = AT(x)
        b = AT(y)
        r1 = CartesianIndices((1:7, 3:8))
        r2 = CartesianIndices((4:10, 3:8))
        copyto!(x, r1, y, r2)
        copyto!(a, r1, b, r2)
        @test x == Array(a)
        r2 = CartesianIndices((4:11, 3:8))
        @test_throws ArgumentError copyto!(a, r1, b, r2)

        r2 = CartesianIndices((4:10, 3:8))
        x2 = fill(0f0, (10, 10))
        copyto!(x2, r1, b, r2)
        @test x2 == x

        fill!(a, 0f0)
        copyto!(a, r1, y, r2)
        @test Array(a) == x

        x = fill(0f0, (10,))
        y = rand(Float32, (20,))
        a = AT(x)
        b = AT(y)
        r1 = CartesianIndices((1:7,))
        r2 = CartesianIndices((4:10,))
        copyto!(x, r1, y, r2)
        copyto!(a, r1, b, r2)
        @test x == Array(a)
        r2 = CartesianIndices((4:11,))
        @test_throws ArgumentError copyto!(a, r1, b, r2)

        x = fill(0f0, (10,))
        y = rand(Float32, (20,))
        a = AT(x)
        b = AT(y)
        r1 = 1:7
        r2 = 4:10
        copyto!(x, r1, y, r2)
        copyto!(a, r1, b, r2)
        @test x == Array(a)

        x = fill(0f0, (10,))
        y = fill(1f0, (10,))
        a = AT(x)
        b = AT(y)
        copyto!(a, b)
        @test Float32.(y) == Array(a)

        # wrapped gpu array to wrapped gpu array
        x = rand(Float32, 4, 4)
        a = AT(x)
        b = view(a, 2:3, 2:3)
        c = AT{eltype(b)}(undef, size(b))
        copyto!(c, b)
        Array(c) == Array(b)

        # wrapped gpu array to cpu array
        z = Array{eltype(b)}(undef, size(b))
        copyto!(z, b)
        @test z == Array(b)

        # cpu array to wrapped gpu array
        copyto!(b, z)

        # bug in copyto!
        ## needless N type parameter
        @test compare((x,y)->copyto!(y, selectdim(x, 2, 1)), AT, ones(Float32, 2, 2, 2), zeros(Float32, 2, 2))
        ## inability to copyto! smaller destination
        ## (this was broken on Julia <1.5)
        @test compare((x,y)->copyto!(y, selectdim(x, 2, 1)), AT, ones(Float32, 2, 2, 2), zeros(Float32, 3, 3))

        if (Float32 in eltypes && Float64 in eltypes)
            # mismatched types
            let src = rand(Float32, 4)
                dst = AT{Float64}(undef, size(src))
                copyto!(dst, src)
                @test Array(dst) == src
            end
            let dst = Array{Float64}(undef, 4)
                src = AT(rand(Float32, size(dst)))
                copyto!(dst, src)
                @test Array(src) == dst
            end
        end
    end

    @testset "cat" begin
        @test compare(hcat, AT, rand(Float32, 3), rand(Float32, 3))
        @test compare(hcat, AT, rand(Float32, ), rand(Float32, 1, 3))
        @test compare(hcat, AT, rand(Float32, 1, 3), rand(Float32))
        @test compare(hcat, AT, rand(Float32, 3), rand(Float32, 3, 3))
        @test compare(hcat, AT, rand(Float32, 3, 3), rand(Float32, 3))
        @test compare(hcat, AT, rand(Float32, 3, 3), rand(Float32, 3, 3))
        #@test compare(hcat, AT, rand(Float32, ), rand(Float32, 3, 3))
        #@test compare(hcat, AT, rand(Float32, 3, 3), rand(Float32))

        @test compare(vcat, AT, rand(Float32, 3), rand(Float32, 3))
        @test compare(vcat, AT, rand(Float32, 3, 3), rand(Float32, 3, 3))
        @test compare(vcat, AT, rand(Float32, ), rand(Float32, 3))
        @test compare(vcat, AT, rand(Float32, 3), rand(Float32))
        @test compare(vcat, AT, rand(Float32, ), rand(Float32, 3, 3))
        #@test compare(vcat, AT, rand(Float32, 3, 3), rand(Float32))

        @test compare((a,b) -> cat(a, b; dims=4), AT, rand(Float32, 3, 4), rand(Float32, 3, 4))
    end

    @testset "reshape" begin
        @test compare(reshape, AT, rand(Float32, 10), Ref((10,)))
        @test compare(reshape, AT, rand(Float32, 10), Ref((10,1)))
        @test compare(reshape, AT, rand(Float32, 10), Ref((1,10)))

        @test_throws Exception reshape(AT(rand(Float32, 10)), (10,2))
    end

    @testset "reinterpret" begin
        a = rand(ComplexF32, 22)
        A = AT(a)
        af0 = reinterpret(Float32, a)
        Af0 = reinterpret(Float32, A)
        @test Array(Af0) == af0
        a = rand(ComplexF32, 4, 4)
        A = AT(a)

        a = rand(ComplexF32, 10 * 10)
        A = AT(a)
        af0 = reshape(reinterpret(Float32, vec(a)), (20, 10))
        Af0 = reshape(reinterpret(Float32, vec(A)), (20, 10))
        @test Array(Af0) == af0
    end

    AT <: AbstractGPUArray && @testset "ntuple test" begin
        result = AT(Vector{NTuple{3, Float32}}(undef, 1))
        gpu_call(ntuple_test, result, Val(3))
        @test Array(result)[1] == (77, 2*77, 3*77)
        x = 88f0
        gpu_call(ntuple_closure, result, Val(3), x)
        @test Array(result)[1] == (x, 2*x, 3*x)
    end

    AT <: AbstractGPUArray && @testset "cartesian iteration" begin
        Ac = rand(Float32, 32, 32)
        A = AT(Ac)
        result = fill!(copy(A), 0.0f0)
        gpu_call(cartesian_iter, result, A, size(A))
        Array(result) == Ac
    end

    AT <: AbstractGPUArray && @testset "Custom kernel from Julia function" begin
        x = AT(rand(Float32, 100))
        y = AT(rand(Float32, 100))
        gpu_call(clmap!, -, x, y; target=x)
        jy = Array(y)
        @test map!(-, jy, jy) ≈ Array(x)
    end

    @testset "map" begin
        @test compare((a, b)-> map(+, a, b),    AT, rand(Float32, 10), rand(Float32, 10))
        @test compare((a, b)-> map!(-, a, b),   AT, rand(Float32, 10), rand(Float32, 10))
        @test compare((a, b, c, d)-> map!(*, a, b, c, d), AT, rand(Float32, 10), rand(Float32, 10), rand(Float32, 10), rand(Float32, 10))
    end

    @testset "repeat" begin
        @test compare(a-> repeat(a, 5, 6),  AT, rand(Float32, 10))
        @test compare(a-> repeat(a, 5),     AT, rand(Float32, 10))
        @test compare(a-> repeat(a, 5),     AT, rand(Float32, 5, 4))
        @test compare(a-> repeat(a, 4, 3),  AT, rand(Float32, 10, 15))
        @test compare(a-> repeat(a, 0),     AT, rand(Float32, 10))
        @test compare(a-> repeat(a, 0),     AT, rand(Float32, 5, 4))
        @test compare(a-> repeat(a, 4, 0),  AT, rand(Float32, 10, 15))
    end

    @testset "permutedims" begin
        @test compare(x->permutedims(x, [1, 2]), AT, rand(Float32, 4, 4))

        inds = rand(1:100, 150, 150)
        @test compare(x->permutedims(view(x, inds, :), (3, 2, 1)), AT, rand(Float32, 100, 100))
    end

    @testset "circshift" begin
        @test compare(x->circshift(x, (0,1)), AT, reshape(Vector(1:16), (4,4)))
    end

    @testset "copy" begin
        a = AT([1])
        b = copy(a)
        fill!(b, 0)
        @test Array(b) == [0]
        @test Array(a) == [1]
    end

    @testset "input output" begin
        # compact=false to avoid type aliases
        replstr(x, kv::Pair...) = sprint((io,x) -> show(IOContext(io, :compact => false, :limit => true, :displaysize => (24, 80), kv...), MIME("text/plain"), x), x)
        showstr(x, kv::Pair...) = sprint((io,x) -> show(IOContext(io, :limit => true, :displaysize => (24, 80), kv...), x), x)

        @testset "showing" begin
            # vectors and non-vector arrays showing
            # are handled differently in base/arrayshow.jl
            A = AT(Int64[1])
            B = AT(Int64[1 2;3 4])

            msg = replstr(A)
            @test occursin(Regex("^1-element $AT{Int64,\\s?1.*}:\n 1\$"), msg)

            # # result of e.g. `print` differs on 32bit and 64bit machines
            # due to different definition of `Int` type
            # print([1]) shows as [1] on 64bit but Int64[1] on 32bit
            msg = showstr(A)
            @test msg == "[1]" || msg == "Int64[1]"

            msg = replstr(B)
            @test occursin(Regex("^2×2 $AT{Int64,\\s?2.*}:\n 1  2\n 3  4\$"), msg)

            msg = showstr(B)
            @test msg == "[1 2; 3 4]" || msg == "Int64[1 2; 3 4]"

            # the printing of Adjoint depends on global state
            msg = replstr(A')
            @test occursin(Regex("^1×1 Adjoint{Int64,\\s?$AT{Int64,\\s?1.*}}:\n 1\$"), msg) ||
                occursin(Regex("^1×1 LinearAlgebra.Adjoint{Int64,\\s?$AT{Int64,\\s?1.*}}:\n 1\$"), msg) ||
                occursin(Regex("^1×1 adjoint\\(::$AT{Int64,\\s?1.*}\\) with eltype Int64:\n 1\$"), msg)
        end
    end
end


# Utility methods for aggregating Select Registers together based on the limited number
# of counters available.
#
# This will allow a suite of counter values to be setup and will automatically keep
# iterating until all counters have been configured.
#
# If it becomes relevant, this will have to be refactored to be aware if we're targeting
# IMC/CHA etc because the rules vary slightly.

numconcurrent(x::T) where {T} = numconcurrent(T)
function numconcurrent(::Type{T}) where {T}
    error("Number of Concurrent Counters not defined for type $T")
end

numconcurrent(::Type{CoreSelectRegister}) = numcounters()

# Uncore PMUs fixed at 4 concurrent events.
# This logic gets more complicated when we throw on CHA Filters and the restriction that
# only a subset of the PMU counters can take some events - but for now we won't worry
# about that.
numconcurrent(::Type{UncoreSelectRegister}) = 4

"""
    group(x) -> itr

Return an iterator `itr` that partitions a collection of Selection Registers in a way
that matches the underlying capability of the hardware.
"""
group(x) = Iterators.partition(x, numconcurrent(eltype(x)))


function asap4(style = Rectangular(2,1); blacklist = ())
    arch = TopLevel{2}("asap4")
    arch.metadata["style"] = style

    # Extra parameters
    num_fifos = 2

    ####################
    # Normal Processor #
    ####################
    processor = build_processor_tile(style)
    for r in 0:23, c in 0:26
        in((r,c), blacklist) && continue
        add_child(arch, processor, CartesianIndex(r,c))
    end
    for r in 0:19, c in 27
        in((r,c), blacklist) && continue
        add_child(arch, processor, CartesianIndex(r,c))
    end

    ####################
    # Memory Processor #
    ####################
    memory_processor = build_processor_tile(style, include_memory = true)
    for r = 24, c = 0:26
        add_child(arch, memory_processor, CartesianIndex(r,c))
    end

    #################
    # 2 Port Memory #
    #################
    memory_2port = build_memory(2)
    for r = 25, c in (0, 2, 4, 6, 9, 11, 13, 15, 18, 20, 22, 24)
        add_child(arch, memory_2port, CartesianIndex(r,c))
    end

    #################
    # 1 Port Memory #
    #################
    memory_1port = build_memory(1)
    for r = 25, c in (8, 17)
        add_child(arch, memory_1port, CartesianIndex(r,c))
    end

    #################
    # Input Handler #
    #################
    input_handler = build_input_handler(1)
    for (r,c) ∈ zip((0,11,13,17), (-1, 28, -1, 28))
        add_child(arch, input_handler, CartesianIndex(r,c))
    end
    # for (r,c) ∈ zip((5,6,5,6), (-1,-1,28,28))
    #     add_child(arch, input_handler, CartesianIndex(r,c))
    # end
    ##################
    # Output Handler #
    ##################
    output_handler = build_output_handler(1)
    for (r,c) ∈ zip((0,11,13,17), (28, -1, 28, -1))
        add_child(arch, output_handler, CartesianIndex(r,c))
    end

    #######################
    # Global Interconnect #
    #######################
    connect_processors(arch, style)
    connect_io(arch, style)
    connect_memories(arch, style)
    return arch
end


mutable struct Tensor
    inner::Ptr{primitivTensor_t}

    function Tensor(inner::Ptr{primitivTensor_t})
        @assert inner != C_NULL
        tensor = new(inner)
        finalizer(tensor, obj -> @call(DeleteTensor, (Ptr{primitivTensor_t},), obj.inner))
        tensor
    end

    function Tensor()
        tensor_ptr_ptr = Ptr{primitivTensor_t}[0]
        @call(CreateTensor, (Ptr{Ptr{primitivTensor_t}},), tensor_ptr_ptr)
        Tensor(tensor_ptr_ptr[1])
    end
end

"""
    valid(tensor)

Check whether the object is valid or not.

# Arguments
- `tensor::Tensor` tensor object
"""
function valid(tensor::Tensor)
    retval_ptr = Ref{PRIMITIV_C_BOOL}(0)
    @call(IsValidTensor, (Ptr{primitivTensor_t}, Ptr{Cuint}), tensor.inner, retval_ptr)
    retval_ptr[] == PRIMITIV_C_TRUE
end

"""
    shape(tensor)

Returns the shape of the Tensor.

# Arguments
- `tensor::Tensor` tensor object
"""
function shape(tensor::Tensor)
    shape_ptr_ptr = Ptr{primitivShape_t}[0]
    @call(GetTensorShape, (Ptr{primitivTensor_t}, Ptr{Ptr{primitivShape_t}}), tensor.inner, shape_ptr_ptr)
    Shape(shape_ptr_ptr[])
end

"""
    to_float(tensor)

Retrieves one internal value in the tensor.

# Arguments
- `tensor::Tensor` tensor object
"""
function to_float(tensor::Tensor)
    retval_ptr = Ref{Float32}(0)
    @call(EvaluateTensorAsFloat, (Ptr{primitivTensor_t}, Ptr{Cfloat}), tensor.inner, retval_ptr)
    retval_ptr[]
end

"""
    reset!(tensor)

Reset internal values using a constant.

# Arguments
- `tensor::Tensor` tensor object
"""
function reset!(tensor::Tensor, k::AbstractFloat)
    @call(ResetTensor, (Ptr{primitivTensor_t}, Cfloat), tensor.inner, k)
end


# Takes full dataset and thins by a factor of 100:

using JLD
@load "tyc3559_mcmc_norv_fixedkernel_v03.jld"

nthin = 100
par_mcmc_thin = zeros(nwalkers,nthin,nparam)
log_like_mcmc_thin =  zeros(nwalkers,nthin)

for i=1:nwalkers
  for j=1:nthin
    par_mcmc_thin[i,j,:]=par_mcmc[i,(j-1)*100+1,:]
    log_like_mcmc_thin[i,j,:]=log_like_mcmc[i,(j-1)*100+1]
  end
end

save("tyc3559_mcmc_norv_fixedkernel_v03_thin.jld","par_mcmc",par_mcmc_thin,"log_like_mcmc",log_like_mcmc_thin,"nwalkers",nwalkers,"nsteps",nthin,"nparam",nparam,"nvary",nvary,"ivary",ivary,"time_phot",time_phot,"flux",flux,"nlast",nlast)


#jl #! format: off
# # Parsing Tabular Data

# **Originally Contributed by**: Clayton Barrows

# ## Introduction

# An example of how to parse tabular files (CSV) files similar to the format established in
# the [RTS-GMLC](github.com/gridmod/rts-gmlc/RTS_Data/SourceData) and create a `System` using
# [PowerSystems.jl](github.com/NREL-SIIP/PowerSystems.jl)

# ### Dependencies
using PowerSystems
using TimeSeries
using Dates

# ### Fetch Data
# PowerSystems.jl links to some test data that is suitable for this example.
# Let's download the test data
PowerSystems.download(PowerSystems.TestData; branch = "master") # *note* add `force=true` to get a fresh copy
base_dir = pkgdir(PowerSystems);

# ### The tabular data format relies on a folder containing `*.csv` files and a `user_descriptors.yaml` file
# First, we'll read the tabular data
RTS_GMLC_DIR = joinpath(base_dir, "data", "RTS_GMLC");
rawsys = PowerSystems.PowerSystemTableData(
    RTS_GMLC_DIR,
    100.0,
    joinpath(RTS_GMLC_DIR, "user_descriptors.yaml"),
    timeseries_metadata_file = joinpath(RTS_GMLC_DIR, "timeseries_pointers.json"),
    generator_mapping_file = joinpath(RTS_GMLC_DIR, "generator_mapping_multi_start.yaml"),
)

# ### Create a `System`
# Next, we'll create a `System` from the `rawsys` data. Since a `System` is predicated on a
# time series resolution and the `rawsys` data includes both 5-minute and 1-hour resolution
# time series, we also need to specify which time series we want to include in the `System`.
# The `time_series_resolution` kwarg filters to only include time series with a matching resolution.

sys = System(rawsys; time_series_resolution = Dates.Hour(1));
horizon = 24;
interval = Dates.Hour(24);
transform_single_time_series!(sys, horizon, interval);
sys


@contract begin
    require(a > 0, b > a, c)
    ensure(b < a, a < 0, !c)
    function foo(a::Int64, b::Int64, c::Bool)
        a *= -1
        b *= -1
        c = !c
        return
    end
end

@contract begin
    require(a > 0, b > a, c)
    ensure(b < a, a > b, !c)
    function flawedEnsureFoo(a::Int64, b::Int64, c::Bool)
        a -= 2
        b += 2
        c = true
        return
    end
end

@testset "Assertion Errors" begin
    @test_throws ContractBreachException foo(-1, 10, true)
    @test_throws ContractBreachException foo(5, 2, true)
    @test_throws ContractBreachException foo(5, 10, false)

    @test_throws ContractBreachException flawedEnsureFoo(10, 100, true)
    @test_throws ContractBreachException flawedEnsureFoo(5, 100, true)
    @test_throws ContractBreachException flawedEnsureFoo(2, 100, true)
end

@testset "Contract Breach Description" begin
    try
        foo(-1, 10, true)
    catch e
        b = IOBuffer()
        showerror(b, e)
        @test String(take!(b)) ==
              "ContractBreachException: Breach on Requirement Expression 'a > 0' in function 'foo'"
    end

    try
        flawedEnsureFoo(10, 100, true)
    catch e
        b = IOBuffer()
        showerror(b, e)
        @test String(take!(b)) ==
              "ContractBreachException: Breach on Ensure Expression 'b < a' in function 'flawedEnsureFoo'"
    end
end


using Test
using MathOptInterface
import Cbc

const MOI = MathOptInterface

const OPTIMIZER = Cbc.Optimizer()
MOI.set(OPTIMIZER, MOI.Silent(), true)

@testset "SolverName" begin
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
end

@testset "supports_default_copy_to" begin
    @test !MOI.Utilities.supports_allocate_load(OPTIMIZER, false)
    @test !MOI.Utilities.supports_allocate_load(OPTIMIZER, true)
    @test !MOI.Utilities.supports_default_copy_to(OPTIMIZER, false)
    @test !MOI.Utilities.supports_default_copy_to(OPTIMIZER, true)
end

const CACHED = MOI.Utilities.CachingOptimizer(
    MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    OPTIMIZER
)
const BRIDGED = MOI.Bridges.full_bridge_optimizer(CACHED, Float64)

const CONFIG = MOI.Test.TestConfig(duals = false, infeas_certificates = false)

@testset "basic_constraint_tests" begin
    MOI.Test.basic_constraint_tests(CACHED, CONFIG)
end

@testset "Unit" begin
    MOI.Test.unittest(BRIDGED, CONFIG, [
        # TODO(odow): implement attributes:
        "number_threads",

        # INFEASIBLE_OR_UNBOUNDED instead of DUAL_INFEASIBLE
        "solve_unbounded_model",

        # No quadratics
        "delete_soc_variables",
        "solve_qcp_edge_cases",
        "solve_qp_edge_cases",
    ])
end

@testset "ModelLike" begin
    @test MOI.get(CACHED, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
    @testset "default_objective_test" begin
         MOI.Test.default_objective_test(CACHED)
     end
     @testset "default_status_test" begin
         MOI.Test.default_status_test(CACHED)
     end
    @testset "nametest" begin
        MOI.Test.nametest(CACHED)
    end
    @testset "validtest" begin
        MOI.Test.validtest(CACHED)
    end
    @testset "emptytest" begin
        MOI.Test.emptytest(BRIDGED)
    end
    @testset "orderedindicestest" begin
        MOI.Test.orderedindicestest(CACHED)
    end
end

@testset "Continuous Linear" begin
    MOI.Test.contlineartest(BRIDGED, CONFIG)
end

@testset "Integer Linear" begin
    MOI.Test.intlineartest(BRIDGED, CONFIG, [
        # Cbc does not support indicator constraints.
        "indicator1", "indicator2", "indicator3", "indicator4",
        # TODO(odow): needs MOI at least 0.9.14.
        "semiconttest", "semiinttest"
    ])
end

@testset "Test params" begin
    # Note: we generate a non-trivial problem to ensure that Cbc struggles to
    # find a solution at the root node.
    knapsack_model = MOI.Utilities.Model{Float64}()
    N = 100
    x = MOI.add_variables(knapsack_model, N)
    MOI.add_constraint.(knapsack_model, MOI.SingleVariable.(x), MOI.ZeroOne())
    MOI.add_constraint(
        knapsack_model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1 + sin(i) for i = 1:N], x),
            0.0
        ),
        MOI.LessThan(10.0),
    )
    MOI.set(
        knapsack_model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([cos(i) for i = 1:N], x),
            0.0
        )
    )
    model = Cbc.Optimizer(
        maxSol = 1, presolve = "off", cuts = "off", heur = "off", logLevel = 0
    )
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.SOLUTION_LIMIT
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    MOI.empty!(model)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.SOLUTION_LIMIT
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
end

@testset "Test PrimalStatus" begin
    model = MOI.Utilities.Model{Float64}()
    x = MOI.add_variable(model)
    MOI.add_constraint(model, x, MOI.GreaterThan(1.0))
    MOI.add_constraint(model, x, MOI.LessThan(0.0))
    cbc = Cbc.Optimizer()
    MOI.copy_to(cbc, model)
    MOI.optimize!(cbc)
    MOI.get(cbc, MOI.PrimalStatus()) == MOI.NO_SOLUTION
end


module CImGui

using CSyntax
using CSyntax.CEnum
using CSyntax.CEnum: Cenum, name_value_pairs

include("LibCImGui.jl")
using .LibCImGui

# missing macros
const IM_COL32_R_SHIFT = 0
const IM_COL32_G_SHIFT = 8
const IM_COL32_B_SHIFT = 16
const IM_COL32_A_SHIFT = 24
const IM_COL32_A_MASK  = 0xff000000
IM_COL32(R,G,B,A) = (ImU32(A)<<IM_COL32_A_SHIFT) | (ImU32(B)<<IM_COL32_B_SHIFT) | (ImU32(G)<<IM_COL32_G_SHIFT) | (ImU32(R)<<IM_COL32_R_SHIFT)
const IM_COL32_WHITE       = IM_COL32(255,255,255,255)  # opaque white = 0xffffffff
const IM_COL32_BLACK       = IM_COL32(0,0,0,255)        # opaque black
const IM_COL32_BLACK_TRANS = IM_COL32(0,0,0,0)          # transparent black = 0x00000000
const IMGUI_PAYLOAD_TYPE_COLOR_3F = "_COL3F"
const IMGUI_PAYLOAD_TYPE_COLOR_4F = "_COL4F"
const FLT_MAX = igGET_FLT_MAX()

include("helper.jl")
include("wrapper.jl")

const IMGUI_VERSION = unsafe_string(GetVersion())

include("backend/GLFW/GLFWBackend.jl")
using .GLFWBackend

include("backend/OpenGL/OpenGLBackend.jl")
using .OpenGLBackend

end # module


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

export QuEST32, QuEST64

module _QuEST_jl_Build_Setup

    include(joinpath("..","deps","build_setup.jl"))

end #^ module _QuEST_jl_Build_Setup

using Libdl: dlopen, dlclose, RTLD_LAZY, RTLD_DEEPBIND, RTLD_GLOBAL

const _QUEST_LIB = Ref{Ptr{Cvoid}}(0)

"""
Function `_dlopen_QuEST()` — loads QuEST dll, returns precision code.
"""
function _dlopen_QuEST(num::Int) ::Cint
    @assert num ∈ [32,64]
    let numstr = string(num)

        if _QuEST_jl_Build_Setup.EXPERT_BUILD
            global _QUEST_LIB[] =
                dlopen("libQuEST_"*numstr,
                       RTLD_LAZY|RTLD_DEEPBIND|RTLD_GLOBAL)
        else
            global _QUEST_LIB[] =
                dlopen(joinpath(@__DIR__,"..","deps","QuEST","build"*numstr,"QuEST","libQuEST"),
                       RTLD_LAZY|RTLD_DEEPBIND|RTLD_GLOBAL)
        end

    end #^ let

    return ccall(:getQuEST_PREC , Cint, (),)
end #^ fn _dlopen_QuEST

function _dlclose_QuEST() ::Nothing
    dlclose(_QUEST_LIB[])
    _QUEST_LIB[] = 0
    nothing;
end

"""
Functions `_quest_mtx_𝑛()` for 𝑛 ∈ {2,4} — convert Julia matrix into QuEST matrix
"""
function _quest_mtx_2(U ::Matrix{Complex{FLT}}) ::QuEST_Types.ComplexMatrix2 where FLT <: Real
    @assert size(U) == (2,2)
    u = QuEST_Types.ComplexMatrix2(
        ( (real(U[1,1]), real(U[1,2])), (real(U[2,1]), real(U[2,2])) ),
        ( (imag(U[1,1]), imag(U[1,2])), (imag(U[2,1]), imag(U[2,2])) )  )
    return u
end

"""
Functions `_quest_mtx_𝑛()` for 𝑛 ∈ {2,4} — convert Julia matrix into QuEST matrix
"""
function _quest_mtx_4(U ::Matrix{Complex{FLT}}) ::QuEST_Types.ComplexMatrix4 where FLT <: Real
    @assert size(U) = (4,4)
    u = QuEST_Types.ComplexMatrix4(
        ( ( real(U[1,1]), real(U[1,2]), real(U[1,3]), real(U[1,4]) ),
          ( real(U[2,1]), real(U[2,2]), real(U[2,3]), real(U[2,4]) ),
          ( real(U[3,1]), real(U[3,2]), real(U[3,3]), real(U[3,4]) ),
          ( real(U[4,1]), real(U[4,2]), real(U[4,3]), real(U[4,4]) ) ),
        ( ( imag(U[1,1]), imag(U[1,2]), imag(U[1,3]), imag(U[1,4]) ),
          ( imag(U[2,1]), imag(U[2,2]), imag(U[2,3]), imag(U[2,4]) ),
          ( imag(U[3,1]), imag(U[3,2]), imag(U[3,3]), imag(U[3,4]) ),
          ( imag(U[4,1]), imag(U[4,2]), imag(U[4,3]), imag(U[4,4]) ) )  )
end

################################################################################
# 32 bit QuEST
################################################################################

"""
Module `QuEST`𝑥𝑦 — Julia wrapper for QuEST with 𝑥𝑦-bit floating point precision.
"""
module QuEST32
export QuEST_Types
export QubitIdx

import .._dlopen_QuEST

function QuEST_init() ::Nothing
    prec = _dlopen_QuEST(32)
    @assert prec == 1 "Wrong precision. Please rebuild the package"
end

module QuEST_Types  # 32 bit qreal
    using CEnum

    const MPI_FLOAT  = Cfloat   # This is needed, because ...
    const MPI_DOUBLE = Cdouble  # ...

    include(joinpath("..","deps","questclang_common_32.jl"))

end #^ module QuEST_Types

const QubitIdx = Cint
const Qreal    = Float32
@assert sizeof(Qreal)==sizeof(QuEST_Types.qreal)

#
# Include C-wrappers, which will use the correct data types
#

using .QuEST_Types

include("data_structure_functions.jl")
include("QASM_logging.jl")
include("debugging.jl")
include("operators.jl")
include("decoherence.jl")
include("state_init.jl")
include("unitaries.jl")
include("calculations.jl")
include("gates.jl")

end #^ module QuEST32

################################################################################
# 64 bit QuEST
################################################################################

"""
Module `QuEST`𝑥𝑦 — Julia wrapper for QuEST with 𝑥𝑦-bit floating point precision.
"""
module QuEST64
export QuEST_Types
export QubitIdx

import .._dlopen_QuEST

function QuEST_init() ::Nothing
    prec = _dlopen_QuEST(64)
    @assert prec == 2 "Wrong precision. Please rebuild the package"
end

module QuEST_Types  # 64 bit qreal
    using CEnum

    const MPI_FLOAT  = Cfloat   # This is needed, because ...
    const MPI_DOUBLE = Cdouble  # ...

    include(joinpath("..","deps","questclang_common_64.jl"))

end #^ module QuEST_Types

const QubitIdx = Cint
const Qreal    = Float64
@assert sizeof(Qreal)==sizeof(QuEST_Types.qreal)

#
# Include C-wrappers, which will use the correct data types
#

using .QuEST_Types

include("data_structure_functions.jl")
include("QASM_logging.jl")
include("debugging.jl")
include("operators.jl")
include("decoherence.jl")
include("state_init.jl")
include("unitaries.jl")
include("calculations.jl")
include("gates.jl")

end #^ module QuEST64

end # module QuEST_jl
# EOF


using GeoData, Test, Dates, Statistics
using GeoData: Start, Center, End,
      formatdims, dims, aggregate, upsample, downsample, span

@testset "upsample" begin
    @test upsample(1, 2) == 1
    @test upsample(2, 2) == 3
    @test upsample(3, 2) == 5
    @test upsample(1, 3) == 1
    @test upsample(2, 3) == 4
    @test upsample(3, 3) == 7
end

@testset "downsample" begin
    @test downsample(1, 2) == 1
    @test downsample(3, 2) == 2
    @test downsample(5, 2) == 3
    @test downsample(1, 3) == 1
    @test downsample(4, 3) == 2
    @test downsample(7, 3) == 3
end

data1 = [ 1  2  3  4  5  6 -1
          7  8  9 10 11 12 -1
         13 14 15 16 17 18 -1]
data2 = 2 * data1
data3 = 3 * data1
data4 = 4 * data1
dimz = Lon([30., 40., 50.]; mode=Sampled(Ordered(), Regular(10.0), Points())),
       Lat(LinRange(-10., 20., 7); mode=Sampled(Ordered(), Regular(5.0), Points()))
array1 = GeoArray(data1, dimz)
array2 = GeoArray(data2, dimz)
array1a = GeoArray(data3, dimz)
array2a = GeoArray(data4, dimz)
stack1 = GeoStack(array1, array2; keys=(:array1, :array2))
stack2 = GeoStack(array1a, array2a; keys=(:array1, :array2))
dates = [DateTime(2017), DateTime(2018)]
series = GeoSeries([stack1, stack2], (Ti(dates),));

@testset "Aggregate a dimension" begin
    lat = Lat(LinRange(3, 13, 6);
              mode=Sampled(Ordered(), Regular(2.0), Intervals(Start())))
    aglat = aggregate(Start(), lat, 3)
    @test span(mode(aglat)) == Regular(6.0)
    @test disaggregate(Start(), aglat, 3) == lat

    aglon = aggregate(Start(), dimz[1], 3)
    @test step(mode(aglon)) === 30.0
    @test val(aglon) == [30.0]
    disaglon = disaggregate(Start(), aglon, 3)
    @test val(disaglon) == val(dimz[1])
    @test step(disaglon) == step(dimz[1])
    @test mode(disaglon) == mode(dimz[1])

    aglat = aggregate(Start(), dimz[2], 3)
    @test step(mode(aglat)) === 15.0
    @test val(aglat) == LinRange(-10.0, 5.0, 2)
    disaglat = disaggregate(Start(), aglat, 3)
    # The last item is lost due to rounding in `aggregate`
    @test val(disaglat) != val(dimz[2])
    @test val(disaglat) === LinRange(-10.0, 15.0, 6)
    @test step(disaglat) == step(dimz[2])
    @test mode(disaglat) == mode(dimz[2])
end

@testset "aggregate and disaggregate at a locus" begin
    @testset "single scale single locus" begin
        scale = 3
        array1_aggstart = aggregate(Start(), array1, scale)
        @test array1_aggstart == [1 4]
        @test length.(dims(array1_aggstart)) == size(array1_aggstart)
        array1_disagstart = disaggregate(Start(), array1_aggstart, scale)
        disaggstart =
            [1 1 1 4 4 4
             1 1 1 4 4 4
             1 1 1 4 4 4]
        @test array1_disagstart == disaggstart
        array1_aggcenter = aggregate(Center(), array1, scale)
        @test array1_aggcenter == [8 11]
        array1_disagcenter = disaggregate(Center(), array1_aggcenter, scale)
        disaggcenter =
            [8 8 8 11 11 11
             8 8 8 11 11 11
             8 8 8 11 11 11]
        @test array1_disagcenter == disaggcenter
        array1_aggend = aggregate(End(), array1, scale)
        @test array1_aggend == [15 18]
        array1_disaggend = disaggregate(End(), array1_aggend, scale)
        @test array1_disaggend ==
            [15 15 15 18 18 18
             15 15 15 18 18 18
             15 15 15 18 18 18]

        stack1_aggstart = aggregate(Start(), stack1, scale)
        @test stack1_aggstart[:array2] == [2 8]
        stack1_disaggstart = disaggregate(Start(), stack1_aggstart, scale)
        @test stack1_disaggstart[:array1] == disaggstart
        @test aggregate(Center(), stack1, scale)[:array2] == [16 22]

        series_aggcenter = aggregate(Center(), series, scale)
        @test series_aggcenter[2][:array2] == [32 44]
        series_disaggcenter = disaggregate(Center(), series_aggcenter, scale)
        @test series_disaggcenter[2][:array2] ==
            [32 32 32 44 44 44
             32 32 32 44 44 44
             32 32 32 44 44 44]
        @test typeof(aggregate(Start(), series, scale)) <: GeoSeries
    end

    @testset "mixed scales" begin
        scale = (3, 2)
        @test aggregate(Start(), array1, scale) == [1 3 5]
        scale = (Lat(2), Lon(3))
        @test aggregate(Center(), array1, scale) == [8 10 12]
        @test aggregate(End(), array1, scale) == [14 16 18]
        A = aggregate(Start(), array1, scale)
        @test length.(dims(A)) == size(A)
    end

    @testset "mixed locus" begin
        scale = 3
        @test aggregate((End(), Start()), array1, 3) == [13 16]
        @test aggregate((End(), Start()), array1, (3, 2)) == [13 15 17]
        A = aggregate((End(), Start()), array1, scale)
        @test length.(dims(A)) == size(A)
    end

    @testset "dim scale" begin
        agg = aggregate(Start(), array1, (Lat(3), Lon(1)))
        @test agg == aggregate(Start(), array1, (1, 3))
        disagg = disaggregate(Start(), agg, (Lat(3), Lon(1)))
        @test disagg ==
            [1  1  1  4  4  4
             7  7  7 10 10 10
            13 13 13 16 16 16]
        agg = aggregate(Start(), array1, (Lon(1), Lat(Near(-4))))
        @test agg == aggregate(Start(), array1, (1, 2))
        @testset "scale 1 dims are unchanged" begin
            @test dims(agg, Lon) === dims(array1, Lon)
        end
    end

end

@testset "Aggregate with a function" begin
    @test aggregate(sum, array1, 3) == [72 99]
    @test aggregate(median, array1, 3) == [8 11]
    @test aggregate(sum, array1, (3, 2)) == [45 57 69]
    A = aggregate(sum, array1, (3, 2))
    @test length.(dims(A)) == size(A)
end

@testset "Aggregate different index modes" begin
    dimz = Band(1:3), Dim{:category}([:a, :b, :c]), X([10, 20, 30, 40])
    a1 = [1 2 3; 4 5 6; 7 8 9]
    A = cat(a1, a1 .+ 10, a1 .+ 20, a1 .+ 30, dims=3)
    da = DimensionalArray(A, dimz)
    @test vec(aggregate(sum, da, (3, 2, 2))) == [114, 354]
end


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

module TestIO

using MXNet
using Base.Test

using ..Main: rand_dims

function test_mnist()
  info("IO::MNIST")
  filenames = mx.get_mnist_ubyte()

  batch_size = 10
  mnist_provider = mx.MNISTProvider(image=filenames[:train_data],
                                    label=filenames[:train_label],
                                    batch_size=batch_size, silent=true, shuffle=false)
  data_spec = mx.provide_data(mnist_provider)
  label_spec = mx.provide_label(mnist_provider)
  @test data_spec == [(:data, (28,28,1,batch_size))]
  @test label_spec == [(:softmax_label, (batch_size,))]

  n_batch = 0
  for batch in mnist_provider
    if n_batch == 0
      data_array  = mx.empty(28,28,1,batch_size)
      label_array = mx.empty(batch_size)
      # have to use "for i=1:1" to get over the legacy "feature" of using
      # [ ] to do concatenation in Julia
      data_targets = [[(1:batch_size, data_array)] for i = 1:1]
      label_targets = [[(1:batch_size, label_array)] for i = 1:1]

      mx.load_data!(mnist_provider, batch, data_targets)
      mx.load_label!(mnist_provider, batch, label_targets)

      true_labels = [5,0,4,1,9,2,1,3,1,4] # the first 10 labels in MNIST train
      got_labels  = Int[copy(label_array)...]
      @test true_labels == got_labels
    end

    n_batch += 1
  end

  @test n_batch == 60000 / batch_size
end

function test_arrays_impl(data::Vector, label::Vector, provider::mx.ArrayDataProvider)
  data = convert(Vector{Array{Float64}}, data)
  label = convert(Vector{Array{Float64}}, label)

  sample_count = size(data[1])[end]
  batch_size   = mx.get_batch_size(provider)
  idx_all = 1:batch_size:sample_count

  for (d1, (_, d2)) in zip(data, mx.provide_data(provider))
    @test size(d1)[1:end-1] == d2[1:end-1]
    @test batch_size == d2[end]
  end
  for (d1, (_, d2)) in zip(label, mx.provide_label(provider))
    @test size(d1)[1:end-1] == d2[1:end-1]
    @test batch_size == d2[end]
  end

  info("IO::Array::#data=$(length(data)),#label=$(length(label)),batch_size=$batch_size")
  for (idx, batch) in zip(idx_all, provider)
    data_batch = [x[[Colon() for i=1:ndims(x)-1]..., idx:min(idx+batch_size-1,sample_count)] for x in data]
    data_get   = mx.get_data(provider, batch)

    for (d_real, d_get) in zip(data_batch, data_get)
      @test d_real ≈ copy(d_get)[[1:n for n in size(d_real)]...]
      @test mx.count_samples(provider, batch) == size(d_real)[end]
    end
  end
end

function test_arrays()
  sample_count = 15
  batch_size   = 4
  dims_data    = [rand_dims()..., sample_count]
  data         = rand(dims_data...)
  provider     = mx.ArrayDataProvider(data, batch_size=batch_size)
  test_arrays_impl(Array[data], [], provider)

  dims_label   = [rand_dims()..., sample_count]
  label        = rand(dims_label...)
  provider     = mx.ArrayDataProvider(data, label, batch_size=batch_size)
  test_arrays_impl(Array[data], Array[label], provider)

  provider     = mx.ArrayDataProvider(:data=>data, :my_label=>label, batch_size=batch_size)
  test_arrays_impl(Array[data], Array[label], provider)

  dims_data2   = [rand_dims()..., sample_count]
  data2        = rand(dims_data2...)
  provider     = mx.ArrayDataProvider((:data=>data, :data2=>data2), label, batch_size=batch_size)
  test_arrays_impl(Array[data,data2], Array[label], provider)
end

function test_arrays_shuffle()
  info("IO::Array::shuffle")

  sample_count = 15
  batch_size   = 4
  data         = rand(mx.MX_float, 1, sample_count)
  label        = collect(1:sample_count)
  provider     = mx.ArrayDataProvider(data, :index => label, batch_size=batch_size, shuffle=true)

  idx_all      = 1:batch_size:sample_count
  data_got     = similar(data)
  label_got    = similar(label)
  for (idx, batch) in zip(idx_all, provider)
    data_batch  = mx.get(provider, batch, :data)
    label_batch = mx.get(provider, batch, :index)
    ns_batch    = mx.count_samples(provider, batch)
    data_got[idx:idx+ns_batch-1]  = copy(data_batch)[1:ns_batch]
    label_got[idx:idx+ns_batch-1] = copy(label_batch)[1:ns_batch]
  end

  @test label_got != label
  @test sort(label_got) == label
  @test size(data_got) == size(data[:, Int[label_got...]])
  @test data_got ≈ data[:, Int[label_got...]]
end

@testset "IO Test" begin
  test_arrays_shuffle()
  test_arrays()
  test_mnist()
end

end


module PowerModelsInterface

import PowerSystems
import PowerModels
import Dates

# PMI exports
export get_pm_data
export apply_time_series
export apply_time_period!

# PowerModels exports
import PowerModels: run_opf
export run_opf
import PowerModels: run_mn_opf
export run_mn_opf
import PowerModels: run_pf
export run_pf

import PowerModels: run_ac_pf
export run_ac_pf
import PowerModels: run_dc_pf
export run_dc_pf
import PowerModels: compute_ac_pf
export compute_ac_pf
import PowerModels: compute_dc_pf
export compute_dc_pf

import PowerModels: run_model
export run_model
import PowerModels: instantiate_model
export instantiate_model

# exact non-convex models
import PowerModels: ACPPowerModel
export ACPPowerModel
import PowerModels: ACRPowerModel
export ACRPowerModel
import PowerModels: ACTPowerModel
export ACTPowerModel
import PowerModels: IVRPowerModel
export IVRPowerModel

# linear approximations
import PowerModels: DCPPowerModel
export DCPPowerModel
import PowerModels: DCMPPowerModel
export DCMPPowerModel
import PowerModels: BFAPowerModel
export BFAPowerModel
import PowerModels: NFAPowerModel
export NFAPowerModel

# quadratic approximations
import PowerModels: DCPLLPowerModel
export DCPLLPowerModel
import PowerModels: LPACCPowerModel
export LPACCPowerModel

# quadratic relaxations
import PowerModels: SOCWRPowerModel
export SOCWRPowerModel
import PowerModels: SOCWRConicPowerModel
export SOCWRConicPowerModel
import PowerModels: SOCBFPowerModel
export SOCBFPowerModel
import PowerModels: SOCBFConicPowerModel
export SOCBFConicPowerModel
import PowerModels: QCRMPowerModel
export QCRMPowerModel
import PowerModels: QCLSPowerModel
export QCLSPowerModel

# sdp relaxations
import PowerModels: SDPWRMPowerModel
export SDPWRMPowerModel
import PowerModels: SparseSDPWRMPowerModel
export SparseSDPWRMPowerModel

const PSY = PowerSystems
const IS = PSY.IS
const PM = PowerModels

include("definitions.jl")
include("pm_data_translator.jl")
include("devices/bus.jl")
include("devices/branch.jl")
include("devices/load.jl")
include("devices/gen.jl")
include("devices/storage.jl")
include("devices/areas.jl")
include("pm_wrapper.jl")

end # module


module FHColoredLightbeam

using ..Ahorn, Maple
using Random

@mapdef Entity "FrostHelper/ColoredLightbeam" ColoredLightbeam(x::Integer, y::Integer, width::Integer=32, height::Integer=24, flag::String="", rotation::Integer=0, color::String="ccffff")

const colors = sort(collect(keys(Ahorn.XNAColors.colors)))

function lightbeamFinalizer(entity::ColoredLightbeam)
    x, y = Ahorn.position(entity)
    nx, ny = Int.(entity.nodes[1])

    theta = atan(y - ny, x - nx) + pi / 2
    magnitude = sqrt((y - ny)^2 + (x - nx)^2)

    entity.rotation = round(Int, rad2deg(theta))
    entity.height = round(Int, max(24, magnitude))
end

Ahorn.editingOptions(entity::ColoredLightbeam) = Dict{String, Any}(
    "color" => colors
)

const placements = Ahorn.PlacementDict(
    "Colored Lightbeam (Frost Helper)" => Ahorn.EntityPlacement(
        ColoredLightbeam,
        "line",
        Dict{String, Any}(),
        lightbeamFinalizer
    )
)

Ahorn.minimumSize(entity::ColoredLightbeam) = 32, 24
Ahorn.resizable(entity::ColoredLightbeam) = true, true

function rotate(point::Tuple{Number, Number}, theta::Number)
    res = [
        cos(theta)  -sin(theta);
        sin(theta) cos(theta)
    ] * [point[1]; point[2]]

    return (res[1], res[2])
end

function getSelectionRect(entity::ColoredLightbeam)
    x, y = Ahorn.position(entity)

    theta = deg2rad(get(entity.data, "rotation", 0))
    width = round(Int, get(entity.data, "width", 32))
    height = round(Int, get(entity.data, "height", 24))

    points = Tuple{Number, Number}[
        (-width / 2, 0),
        (width / 2, 0),
        (-width / 2, height - 4),
        (width / 2, height - 4)
    ]

    rotated = rotate.(points, theta)

    tlx, tly = minimum(p -> p[1], rotated), minimum(p -> p[2], rotated)
    brx, bry = maximum(p -> p[1], rotated), maximum(p -> p[2], rotated)

    return Ahorn.Rectangle(x + tlx, y + tly, brx - tlx, bry - tly)
end

function Ahorn.selection(entity::ColoredLightbeam)
    return getSelectionRect(entity)
end

texture = "util/lightbeam"

function renderLightbeam(ctx::Ahorn.Cairo.CairoContext, entity::ColoredLightbeam)
    x, y = Ahorn.position(entity)

    lightbeam = Ahorn.getSprite(texture, "Gameplay")

    theta = deg2rad(get(entity.data, "rotation", 0))
    width = round(Int, get(entity.data, "width", 32))
    height = round(Int, get(entity.data, "height", 24))

    rng = Ahorn.getSimpleEntityRng(entity)

    Ahorn.Cairo.save(ctx)

    Ahorn.translate(ctx, x, y)
    Ahorn.rotate(ctx, theta)
    Ahorn.translate(ctx, width / 2, 0)
    Ahorn.rotate(ctx, pi / 2)
    Ahorn.scale(ctx, (height - 4) / lightbeam.width, 1)
    Ahorn.set_antialias(ctx, 1)

    Ahorn.Cairo.save(ctx)

    for i in 0:width - 1
        Ahorn.drawImage(ctx, texture, 0, i, alpha=0.3)
    end

    for i in 0:4:width - 1
        num = i * 0.6
        lineWidth = 4.0 + sin(num * 0.5 + 1.2) * 4.0
        length = height + sin(num * 0.25) * 8
        alpha = 0.6 + sin(num + 0.8) * 0.3
        offset = sin((num + i * 32) * 0.1 + sin(num * 0.05 + i * 0.1) * 0.25) * (width / 2.0 - lineWidth / 2.0)

        # Not part of official calculations, but makes the render a bit less boring
        offsetMultiplier = (rand(rng) - 0.5) * 2

        for j in 0:3
            Ahorn.drawImage(ctx, texture, 0, offset * offsetMultiplier + width / 2 + j, alpha=alpha / 1.5)
        end
    end

    Ahorn.Cairo.restore(ctx)
end

Ahorn.renderAbs(ctx::Ahorn.Cairo.CairoContext, entity::ColoredLightbeam, room::Maple.Room) = renderLightbeam(ctx, entity)

end


export QBasedPolicy, get_prob

using Flux: softmax

"""
    QBasedPolicy(learner::Q, selector::S) -> QBasedPolicy{Q, S}

Use a Q-`learner` to generate the estimations of actions and use `selector` to get the action.
"""
Base.@kwdef struct QBasedPolicy{
    Q<:AbstractLearner,
    S<:AbstractActionSelector,
} <: AbstractPolicy
    learner::Q
    selector::S
end

(π::QBasedPolicy)(obs::Observation) = obs |> π.learner |> π.selector

get_prob(π::QBasedPolicy, s) = get_prob(π.selector, π.learner(s))

function update!(
    π::QBasedPolicy{<:Union{PrioritizedDQNLearner,RainbowLearner}},
    buffer::AbstractTurnBuffer,
)
    indexed_batch = extract_transitions(buffer, π)
    if !isnothing(indexed_batch)
        inds, batch = indexed_batch
        priorities = update!(π.learner, batch)
        isnothing(priorities) || (priority(buffer)[inds] .= priorities)
    end
end

update!(π::QBasedPolicy, model::AbstractEnvironmentModel; kw...) =
    update!(π.learner, model; kw...)

priority(transition, π::QBasedPolicy) = priority(transition, π.learner)

#####
# dispatches
#####

extract_transitions(buffer, π::QBasedPolicy) = extract_transitions(buffer, π.learner)

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    π::QBasedPolicy{<:TDLearner{<:AbstractQApproximator,:ExpectedSARSA}},
)
    if length(buffer) > 0
        n = π.learner.n
        transitions = @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            actions = action(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n):end],
            terminals = terminal(buffer)[max(1, end - n):end],
            next_states = state(buffer)[max(1, end - n):end],
        )
        s′ = transitions.next_states[end]
        merge(transitions, (probs_of_a′ = get_prob(π, s′),))
    else
        nothing
    end
end