module Map_removeEQ

export map_removeEQ

using Dates, JLD2, DSP, PlotlyJS, Printf, ORCA, FileIO

using SeisIO

include("get_kurtosis.jl")
include("remove_eq.jl")

using .Get_kurtosis, .Remove_eq

"""
    ParallelEQremoval(dlid, InputDict::Dict)
    remove earthquake and save it into jld2 file.

"""
function map_removeEQ(dlid, InputDict::Dict)

    #store data

    finame                 =InputDict["finame"]
    IsKurtosisRemoval      =InputDict["IsKurtosisRemoval"]
    max_edgetaper_duration =InputDict["max_edgetaper_duration"]
    kurtosis_timewindow    =InputDict["kurtosis_timewindow"]
    kurtosis_threshold     =InputDict["kurtosis_threshold"]
    IsSTALTARemoval        =InputDict["IsSTALTARemoval"]
    stalta_longtimewindow  =InputDict["stalta_longtimewindow"]
    stalta_threshold       =InputDict["stalta_threshold"]
    invert_tukey_α         =InputDict["invert_tukey_α"]
    max_wintaper_duration  =InputDict["max_wintaper_duration"]
    removal_shorttimewindow=InputDict["removal_shorttimewindow"]
    overlap                =InputDict["overlap"]
    plot_kurtosis_α        =InputDict["plot_kurtosis_α"]
    plot_boxheight         =InputDict["plot_boxheight"]
    plot_span              =InputDict["plot_span"]
    fodir                  =InputDict["fodir"]
    foname                 =InputDict["foname"]
    fopath                 =InputDict["fopath"]
    IsSaveFig              =InputDict["IsSaveFig"]

    DLtimestamplist        =InputDict["DLtimestamplist"]
    stationlist            =InputDict["stationlist"]
    NumofTimestamp         =InputDict["NumofTimestamp"]

    tstamp = DLtimestamplist[dlid]

    if mod(dlid, round(0.1*NumofTimestamp)+1) == 0
        println(@sprintf("start process %s", tstamp))
    end

    SRall = SeisData(length(stationlist))

    icount = 0

    bt_getkurtosis = 0.0
    bt_removeeq = 0.0

    for st = stationlist
        #S = t[joinpath(tstamp, st)]
        st1 = replace(st, "-"=>"")
        S = FileIO.load(finame, joinpath(tstamp, st1))

        if S.misc["dlerror"] == 0
            dt = 1/S.fs
            #tvec = collect(0:S.t[2,1]-1) * dt ./ 60 ./ 60
            tvec = collect(0:length(S.x)-1) * dt ./ 60 ./ 60

            #tapering to avoid instrumental edge artifacts
            SeisIO.taper!(S,  t_max = max_edgetaper_duration, α=0.05)
            S1 = deepcopy(S)

            #set long window length to user input since last window of previous channel will have been adjusted
            S1.misc["eqtimewindow"] = fill(true, length(S1.x))

            if IsKurtosisRemoval
                # compute kurtosis and detect earthqukes
                bt_1 = @elapsed S1 = Get_kurtosis.get_kurtosis(S1, float(kurtosis_timewindow))
                bt_2 = @elapsed S1 = Remove_eq.detect_eq_kurtosis(S1, tw=float(removal_shorttimewindow), kurtosis_threshold=float(kurtosis_threshold), overlap=float(overlap))

                btsta_1 = 0

                if IsSTALTARemoval
                    # detect earthquake and tremors by STA/LTA
                    btsta_1 = @elapsed S1 = Remove_eq.detect_eq_stalta(S1, float(stalta_longtimewindow), float(removal_shorttimewindow),
                                        float(stalta_threshold), float(overlap))
                end

                bt_3 = @elapsed S1 = Remove_eq.remove_eq(S1, S, float(invert_tukey_α), plot_kurtosis_α, max_wintaper_duration,
                                plot_boxheight, trunc(Int, plot_span), fodir, tstamp, tvec, IsSaveFig)

                bt_getkurtosis += bt_1
                bt_removeeq += bt_2 + bt_3 + btsta_1

                #if mod(dlid, round(0.1*NumofTimestamp)+1) == 0
                #    println([bt_1, bt_2, btsta_1, bt_3])
                #end

                #remove kurtosis for reduce size
                S1.misc["kurtosis"] = []
                S1.misc["eqtimewindow"] = []

            else
                #only STA/LTA
                if IsSTALTARemoval

                    bt_2 = @elapsed S1 = Remove_eq.detect_eq_stalta(S1, float(stalta_longtimewindow), float(removal_shorttimewindow),
                                        float(stalta_threshold), float(overlap))
                    bt_3 = @elapsed S1 = Remove_eq.remove_eq(S1, S, float(invert_tukey_α), plot_kurtosis_α, max_wintaper_duration,
                                    plot_boxheight, trunc(Int, plot_span), fodir, tstamp, tvec, IsSaveFig)

                    #remove kurtosis for reduce size
                    S1.misc["kurtosis"] = []
                    S1.misc["eqtimewindow"] = []

                    bt_getkurtosis += 0.0
                    bt_removeeq += bt_2 + bt_3

                else
                    #no removal process.
                    println("Both 'IsKurtosisRemoval' and 'IsKurtosisRemoval' are false. No removal process is executed.")
                    exit(0)
                end

            end

        else
            #download error found: save as it is.
            S1 = S
        end

        icount += 1
        SRall[icount] = S1

    end

    return (SRall, bt_getkurtosis, bt_removeeq)
end

end


function get_pv_coordinates(obj::ExtendedPVCoordinatesProvider, arg0::FieldAbsoluteDate, arg1::Frame)
    return jcall(obj, "getPVCoordinates", TimeStampedFieldPVCoordinates, (FieldAbsoluteDate, Frame), arg0, arg1)
end

function get_pv_coordinates(obj::PVCoordinatesProvider, arg0::AbsoluteDate, arg1::Frame)
    return jcall(obj, "getPVCoordinates", TimeStampedPVCoordinates, (AbsoluteDate, Frame), arg0, arg1)
end

function to_field_pv_coordinates_provider(obj::ExtendedPVCoordinatesProvider, arg0::Field)
    return jcall(obj, "toFieldPVCoordinatesProvider", FieldPVCoordinatesProvider, (Field,), arg0)
end


module TypedTables

using NullableArrays, Compat

export DefaultKey, Cell, Column, Row, Table, DenseTable, KeyTable, DenseKeyTable

export rename, name, eltypes, permutecols, ncol, nrow, head, tail, storagetype, storagetypes

export unique!

export readtable, readtable!, writetable

export @Cell, @Row, @Column, @Table
export @col, @select, @filter, @filter!, @filter_mask

import Base.(==)

# TODO Do we need a unique table? It could keep a persistent hash. (Perhaps a multi-column key table with no value columns?)
# TODO union!, intersect! and setdiff!
# TODO dplyr funcitons (should we use same names or julia interfaces or both??)
#      - arrange (sorting)
#      - select/rename (combination of rename with indexing)
#      - distinct (currently unique and is already O(N log N))
#      - mutate (perhaps compute or computecol or something, given mutate has a meaning in julia)
#      - summarize
#      - sample_n and sample_frac (could use rand() interface I think)
#      - left_join, right_join, full_join (aka outer_join), semi_join, anti_join
#      - think about group_by (generates something similar to KeyTable of Tables?)
# TODO Load from an object with a dataframe interface
# TODO Figure out namespace issues for head and tail (maybe both are in Base for Julia 0.5?)
# TODO sub/slice for generating a sub-table (i.e. simply a table with different StorageTypes)
# TODO possibly implement subtable type (with no ability to insert/delete rows, though can setindex!) - is this only necessary for Row?
# TODO fix TableKey so it always references the length of the current table, not its parent
# TODO Somehow make sense of the key mess for joins
# TODO Unary relational operations like select, project, etc (see dplyr for inspiration)
# TODO Binary relational operations like left_join, etc.
# TODO other types of computed joins, map, do, etc
# TODO finish DenseTable
# TODO implement KeyTable ?
# TODO implement DenseKeyTable ?

# Julia 0.5 agenda:
# TODO Use generalized comprehensions in @select (user specified data structure)
#      and @filter (BitArray)
# TODO Use @pure functions where beneficial

include("Cell.jl")
include("Column.jl")
include("Row.jl")
include("Table.jl")

include("show.jl")
include("datamacros.jl")
include("setalgorithms.jl")
include("join.jl")
include("IO.jl")

end # module


# m = (0,1,1,2)
# shuffle_24_bytes_to_32 = Val( ntuple(i->m[mod1(i,4)] +  3( (i-1)÷4 ),32) )
const shuffle_24_bytes_to_32 =  (
     0,  1,  1,  2,
     3,  4,  4,  5,
     6,  7,  7,  8,
     9, 10, 10, 11,
    12, 13, 13, 14,
    15, 16, 16, 17,
    18, 19, 19, 20,
    21, 22, 22, 23
)

# AVX 256-bit (32 bytes)
const shuffle_24_bytes_to_32_val = Val{shuffle_24_bytes_to_32}()
const shuffle_24_bytes_to_32_m4_val = Val{shuffle_24_bytes_to_32 .+ 4}()
const shuffle_24_bytes_to_32_m8_val = Val{shuffle_24_bytes_to_32 .+ 8}()
const choose_even_odd_ints_16 = Val( ntuple(i->isodd(i-1)*16 + i -1, 16) )

"""
    unpack_uint12_to_uint16(A::Vector{UInt8}, out::Vector{UInt16}, [ i ])

    Unpack 12-bit integers into 16-bit integers. Two 12-bit integers are packed consecutively into three bytes.
"""
@inline function unpack_uint12_to_uint16(
    A::SIMD.FastContiguousArray{UInt8,1},
    out::SIMD.FastContiguousArray{UInt16,1},
    i,
    offset = 0,
    shuffle_24_bytes_to_32::Val{T} = shuffle_24_bytes_to_32_val
    ) where T

    # Load 32 bytes (we only use the first 24)
    a = @inbounds vload(Vec{32,UInt8}, A, i-offset)

    a16 = reinterpret( Vec{16,UInt16}, shufflevector(a, shuffle_24_bytes_to_32) )

    a16 = shufflevector(a16 & 0xfff, a16 >> 4, choose_even_odd_ints_16)

    @inbounds vstore(a16, out, 1 + 2( ( i-1 ) ÷ 3) )
end
@inline unpack_uint12_to_uint16_offset_m4(A, out, i) = unpack_uint12_to_uint16(A, out, i, 4, shuffle_24_bytes_to_32_m4_val)
@inline unpack_uint12_to_uint16_offset_m8(A, out, i) = unpack_uint12_to_uint16(A, out, i, 8, shuffle_24_bytes_to_32_m8_val)

"""
    unpack_uint12_to_uint16(A::SIMD.FastContiguousArray{UInt8,1}, out::Vector{UInt16})

    Unpack entire array with a preallocated out buffer for compressed data
"""
function unpack_uint12_to_uint16(A::SIMD.FastContiguousArray{UInt8,1}, out::SIMD.FastContiguousArray{UInt16,1})
    unpack_uint12_to_uint16(A, out, 1)
    idx = 25:24:length(A)
    @inbounds for i = idx[1:end-1]
        unpack_uint12_to_uint16_offset_m4(A, out, i)
    end
    extra_bytes = rem(length(A), 3)
    unpack_uint12_to_uint16_offset_m8(A, out, length(A) - 23 - extra_bytes)
    if extra_bytes == 1
        out[length(A) * 2 ÷ 3] = A[end]
    elseif extra_bytes == 2
        out[length(A) * 2 ÷ 3] = reinterpret(UInt16, A[end-1:end])[1] & 0xfff
    end
    out
end

"""
    unpack_uint12_to_uint16(A::SIMD.FastContiguousArray{UInt8,1})

    Unpack entire array, allocates and returns output buffer
"""
function unpack_uint12_to_uint16(A::SIMD.FastContiguousArray{UInt8,1})
    out = Vector{UInt16}(undef, length(A) * 2 ÷ 3)
    unpack_uint12_to_uint16(A, out)
end

"""
    unpack_uint12_to_uint16(filepath::AbstractString)

    Load file at filepath, memory map it and unpack
"""
function unpack_uint12_to_uint16(filepath::AbstractString)
    ios = open(filepath, "r")
    A = Mmap.mmap(ios)
    unpack_uint12_to_uint16(A);
    close(ios)
end

function unpack_uint12_to_uint16_partitioned(inn, out = Array{UInt16,1}(undef, div(length(inn), 3, RoundUp)*2))
    partition_length = div( length(inn), Threads.nthreads(), RoundUp)
    # Ensure that partition length is a multiple of three
    partition_length += 3 - mod1(partition_length, 3)
    @debug "Partition Length: " partition_length partition_length ÷ 3 * 2
    inp = collect(Iterators.partition(inn, partition_length))
    outp = collect(Iterators.partition(out, partition_length ÷ 3 * 2))
    @inbounds Threads.@threads for i in eachindex(inp)
        unpack_uint12_to_uint16(inp[i], outp[i])
    end
    out
end

#=

=#
const shiftvec = reinterpret(Vec{16,UInt16},Vec{8,UInt32}(0x00040000))
const andvec = reinterpret(Vec{16,UInt16}, Vec{8,UInt32}(0xffff0fff))
@inline function unpack_uint12_to_uint16_one_shuffle(
    A::SIMD.FastContiguousArray{UInt8,1},
    out::SIMD.FastContiguousArray{UInt16,1},
    i,
    offset = 0,
    shuffle_24_bytes_to_32::Val{T} = shuffle_24_bytes_to_32_val,
    ) where T
    # Load 32 bytes (we only use the first 24)
    a = @inbounds vload(Vec{32,UInt8}, A, i-offset)

    a16 = reinterpret( Vec{16,UInt16}, shufflevector(a, shuffle_24_bytes_to_32) )

    a16 = a16 >> shiftvec & andvec

    @inbounds vstore(a16, out, 1 + 2( ( i-1 ) ÷ 3) )
end


"""
    ConstantAgent

Takes constant actions.
"""
struct ConstantAgent <: AbstractPolicy
    scale::Int
end
create_agent(::Val{ConstantAgent}; scale=50, kwargs...) = ConstantAgent(scale)

function (a::ConstantAgent)(env::AbstractEnv)
    fill(a.scale, length(action_space(env)))
end

"""
    RandomAgent

Takes random actions.
"""
struct RandomAgent <: AbstractPolicy
    rng::AbstractRNG
end
create_agent(::Val{:RandomAgent}; seed = 37, kwargs...) = RandomAgent(StableRNG(seed))

function (a::RandomAgent)(env::AbstractEnv)
    round.(Int, rand.(a.rng, action_space(env)))
end

"""
    SinusAgent

`period` is the period in steps, the whole curve will be between `0` and `2*amplitude`
"""
mutable struct SinusAgent <: AbstractPolicy
    period::Float64
    amplitude::Float64
    step::Int
end
create_agent(::Val{:SinusAgent}; period=1000, amplitude=40.0, kwargs...) = SinusAgent(period, amplitude, 0)

function (a::SinusAgent)(env::AbstractEnv)
    a.step += 1
    round.(Int, a.amplitude .* (1 .+ sin(a.step * 2π / a.period)) .+ ones(length(action_space(env))))
end

"""
    SimpleAgent

Something similar to the k8s scaler
"""
struct SimpleAgent <: AbstractPolicy
    down::Float64
    up::Float64
end
create_agent(::Val{:SimpleAgent}; down = 0.5, up = 0.8, kwargs...) = SimpleAgent(down, up)

function (agent::SimpleAgent)(env::AbstractEnv)
    env = env[!]
    a = Int[ms.running_nodes for ms in env.microservices]
    for (i, ms) in enumerate(env.microservices)
        util = length(ms.queue) / ms.running_nodes
        if util < agent.down
            a[i] -= 1
        elseif util > agent.up
            a[i] += 1
        end
    end
    lowas = [as.left for as in env.action_space]
    highas = [as.right for as in env.action_space]
    return clamp.(a, lowas, highas)
end

# Guided
function create_agent(::Val{:Guided};
    learner, teacher,
    kwargs...
)
    learner = create_agent(Val(learner); kwargs...)
    teacher = create_agent(Val(teacher); kwargs...)
    GuidedPolicy(learner, teacher; kwargs...)
end

"""
    K8sAgent

Something similar to the k8s scaler
"""
mutable struct K8sAgent <: AbstractPolicy
    desired_util::Float64
    step_tolerance::Float64
    scale_memory::Matrix{Int}
    mem_ptr::Int
end
create_agent(::Val{:K8sAgent}; desired_util=0.8, stabilization_steps=300, step_tolerance=0.1, env, kwargs...) = K8sAgent(desired_util, step_tolerance, zeros(Int, length(env[!].microservices), stabilization_steps), 1)

function (agent::K8sAgent)(env::AbstractEnv)
    env = env[!]

    newptr = mod(agent.mem_ptr, size(agent.scale_memory, 2)) + 1
    for (i, ms) in enumerate(env.microservices)
        rn = max(1e-6, ms.running_nodes) # To avoid problems when running_nodes == 0
        util = min(1, length(ms.queue) / rn)
        ratio = util / agent.desired_util
        if abs(ratio - 1) < agent.step_tolerance # No new action if within limit
            agent.scale_memory[i, newptr] = agent.scale_memory[i, agent.mem_ptr]
        else # Record new desired action
            agent.scale_memory[i, newptr] = ceil(Int, rn * ratio)
        end
    end

    agent.mem_ptr = newptr

    a = vec(maximum(agent.scale_memory, dims=2)) # Take max of desired action in window

    lowas = [as.left for as in env.action_space]
    highas = [as.right for as in env.action_space]
    return clamp.(a, lowas, highas)
end

"""
    Oracle

Optimal agent with assumption of knowledge of graph and delays, hardcoded for simpleflip env
"""
struct OracleAgent <: AbstractPolicy
end
create_agent(::Val{:OracleAgent}; kwargs...) = OracleAgent()

function (agent::OracleAgent)(env::AbstractEnv)
    env = env[!]
    a = zeros(Int, length(env.microservices))
    q = map(ms -> length(ms.queue), env.microservices)
    for tidx in 1:length(env.jobtypes)
        q = map(ms -> count(job -> job.typeindex == tidx, ms.queue), env.microservices[env.jobtypes[tidx].path])
        a[env.jobtypes[tidx].path] .+= [q[1] + 1; [max(q[i-1], q[i]) for i in 2:length(q)]]
    end
    lowas = [as.left for as in env.action_space]
    highas = [as.right for as in env.action_space]
    return clamp.(a, lowas, highas)
end

"""
    Loaded

Loads bson saved policy and acts on it using mean values of actions
"""
struct LoadedAgent{P<:GaussianNetwork} <: AbstractPolicy
    policy::P
end
function create_agent(::Val{:LoadedAgent}; loadpath::String, kwargs...)
    # Expects SAC, probably works with others with gaussian net
    # Though make sure the correct env is used then
    dict = BSON.load(loadpath)
    LoadedAgent(dict[:policy].policy.model)
end
function (agent::LoadedAgent)(env::AbstractEnv)
    s = state(env)
    action_mean, action_std = agent.policy(s)
    return action_mean
end

"""
    Wrapped

Wraps an env that returns from the inner action space, to a normalized action space.
"""
struct WrappedAgent{P<:AbstractPolicy,F<:Function} <: AbstractPolicy
    policy::P
    action_transform::F
end
function create_agent(::Val{:WrappedAgent}; env, wrapped_policy, kwargs...)
    A = action_space(env)
    alow = [x isa UnitRange ? x.start : x.left for x in A]
    ahigh = [x isa UnitRange ? x.stop : x.right for x in A]

    action_mapping(a) = 2 .* (a .- alow) ./ (ahigh .- alow) .- 1
    inner_agent = create_agent(Val(wrapped_policy); env=env[!], kwargs...)

    WrappedAgent(inner_agent, action_mapping)
end
function (agent::WrappedAgent)(env::AbstractEnv)
    a = agent.policy(env[!]) # Get action based on base env
    return agent.action_transform(a)
end


const PLANETOID_URL = "https://github.com/kimiyoung/planetoid/raw/master/data"
const PLANETOID_DATASETS = [:citeseer, :cora, :pubmed]
const EXTS = ["allx", "ally", "graph", "test.index", "tx", "ty", "x", "y"]
const DATAURLS = [joinpath(PLANETOID_URL, "ind.$(d).$(ext)") for d in PLANETOID_DATASETS, ext in EXTS]

planetoid_init() = register(DataDep(
    "Planetoid",
    """
    The citation network datasets "Cora", "CiteSeer", "PubMed" from
    "Revisiting Semi-Supervised Learning with Graph Embeddings"
    <https://arxiv.org/abs/1603.08861> paper.
    Nodes represent documents and edges represent citation links.
    """,
    reshape(DATAURLS, :),
    "f52b3d47f5993912d7509b51e8090b6807228c4ba8c7d906f946868005c61c18";
    post_fetch_method=preprocess_planetoid,
))

function preprocess_planetoid(local_path)
    for dataset in PLANETOID_DATASETS
        graph_file = @datadep_str "Planetoid/ind.$(dataset).graph"
        trainX_file = @datadep_str "Planetoid/ind.$(dataset).x"
        trainy_file = @datadep_str "Planetoid/ind.$(dataset).y"
        testX_file = @datadep_str "Planetoid/ind.$(dataset).tx"
        testy_file = @datadep_str "Planetoid/ind.$(dataset).ty"
        allX_file = @datadep_str "Planetoid/ind.$(dataset).allx"
        ally_file = @datadep_str "Planetoid/ind.$(dataset).ally"

        train_X = read_data(trainX_file)
        train_y = read_data(trainy_file)
        test_X = read_data(testX_file)
        test_y = read_data(testy_file)
        all_X = read_data(allX_file)
        all_y = read_data(ally_file)
        graph = read_graph(graph_file)

        num_V = length(graph)
        sg = to_simplegraph(graph, num_V)
        num_E = ne(sg)
        feat_dim = size(all_X, 2)
        label_dim = size(all_y, 2)
        train_X, train_y = sparse(train_X'), sparse(train_y')
        test_X, test_y = sparse(test_X'), sparse(test_y')
        all_X, all_y = sparse(all_X'), sparse(all_y')
        raw = Dict(:graph=>graph, :train_X=>train_X, :train_y=>train_y,
                   :test_X=>test_X, :test_y=>test_y, :all_X=>all_X, :all_y=>all_y)
        meta = (graph=(num_V=num_V, num_E=num_E),
                train=(features_dim=(feat_dim, size(train_X, 2)), labels_dim=(label_dim, size(train_y, 2))),
                test=(features_dim=(feat_dim, size(test_X, 2)), labels_dim=(label_dim, size(test_y, 2))),
                all=(features_dim=(feat_dim, size(all_X, 2)), labels_dim=(label_dim, size(all_y, 2)))
                )

        graphfile = replace(graph_file, "ind.$(dataset).graph"=>"$(dataset).graph.jld2")
        trainfile = replace(graph_file, "ind.$(dataset).graph"=>"$(dataset).train.jld2")
        testfile = replace(graph_file, "ind.$(dataset).graph"=>"$(dataset).test.jld2")
        allfile = replace(graph_file, "ind.$(dataset).graph"=>"$(dataset).all.jld2")
        rawfile = replace(graph_file, "ind.$(dataset).graph"=>"$(dataset).raw.jld2")
        metadatafile = replace(graph_file, "ind.$(dataset).graph"=>"$(dataset).metadata.jld2")

        @save graphfile sg
        @save trainfile train_X train_y
        @save testfile test_X test_y
        @save allfile all_X all_y
        @save rawfile raw
        @save metadatafile meta
    end
end

function read_data(filename)
    py"""
    import pickle
    from scipy.sparse import csr_matrix

    with open($filename,"rb") as f:
        u = pickle._Unpickler(f)
        u.encoding = "latin1"
        data = u.load()

    if type(data) is csr_matrix:
        data = data.toarray()
    """
    return SparseMatrixCSC(Array(py"data"))
end

read_index(filename) = map(x -> parse(Int64, x), readlines(filename))

function read_graph(filename)
    py"""
    import pickle

    with open($filename,"rb") as f:
        u = pickle._Unpickler(f)
        u.encoding = "latin1"
        data = u.load()
    """
    return Dict(py"data")
end

struct Planetoid <: Dataset
    dataset::Symbol

    function Planetoid(ds::Symbol)
        @assert ds in PLANETOID_DATASETS "`dataset` should be one of citeseer, cora, pubmed."
        new(ds)
    end
end

function graphdata(pla::Planetoid)
    file = @datadep_str "Planetoid/$(pla.dataset).graph.jld2"
    @load file sg
    sg
end

function traindata(pla::Planetoid)
    file = @datadep_str "Planetoid/$(pla.dataset).train.jld2"
    @load file train_X train_y
    train_X, train_y
end

function testdata(pla::Planetoid)
    file = @datadep_str "Planetoid/$(pla.dataset).test.jld2"
    @load file test_X test_y
    test_X, test_y
end

function alldata(pla::Planetoid)
    file = @datadep_str "Planetoid/$(pla.dataset).all.jld2"
    @load file all_X all_y
    all_X, all_y
end

function rawdata(pla::Planetoid)
    file = @datadep_str "Planetoid/$(pla.dataset).raw.jld2"
    @load file raw
    raw
end

function metadata(pla::Planetoid)
    file = @datadep_str "Planetoid/$(pla.dataset).metadata.jld2"
    @load file meta
    meta
end


#####
##### Interface for players
#####

"""
    AbstractPlayer{Game}

Abstract type for a player of `Game`.
"""
abstract type AbstractPlayer{Game} end

GameType(::AbstractPlayer{Game}) where Game = Game

"""
    think(::AbstractPlayer, game)

Return a probability distribution over actions as a `(actions, π)` pair. Also optionally return a value_estimate of the current state
"""
function think end

"""
    reset_player!(::AbstractPlayer)

Reset the internal memory of a player (e.g. the MCTS tree).
The default implementation does nothing.
"""
function reset_player!(::AbstractPlayer)
  return
end

"""
    player_temperature(::AbstractPlayer, game, turn_number)

Return the player temperature, given the number of actions that have
been played before by both players in the current game.

A default implementation is provided that always returns 1.
"""
function player_temperature(::AbstractPlayer, game, turn)
  return 1.0
end

"""
    select_move(player::AbstractPlayer, game, turn_number)

Return a single action. A default implementation is provided that samples
an action according to the distribution computed by [`think`](@ref), with a
temperature given by [`player_temperature`](@ref).
"""
function select_move(player::AbstractPlayer, game, turn_number)
  actions, π, _ = think(player, game)
  τ = player_temperature(player, game, turn_number)
  π = apply_temperature(π, τ)
  return actions[Util.rand_categorical(π)]
end

#####
##### Random Player
#####

"""
    RandomPlayer{Game} <: AbstractPlayer{Game}

A player that picks actions uniformly at random.
"""
struct RandomPlayer{Game} <: AbstractPlayer{Game} end

function think(player::RandomPlayer, game)
  actions = GI.available_actions(game)
  n = length(actions)
  π = ones(n) ./ length(actions)
  return actions, π, nothing
end

#####
##### Epsilon-greedy player
#####

"""
    EpsilonGreedyPlayer{Game, Player} <: AbstractPlayer{Game}

A wrapper on a player that makes it choose a random move
with a fixed ``ϵ`` probability.
"""
struct EpsilonGreedyPlayer{G, P} <: AbstractPlayer{G}
  player :: P
  ϵ :: Float64
  function EpsilonGreedyPlayer(p::AbstractPlayer{G}, ϵ) where G
    return new{G, typeof(p)}(p, ϵ)
  end
end

function think(p::EpsilonGreedyPlayer, game)
  actions, π, value_estimate = think(p.player, game)
  n = length(actions)
  η = ones(n) ./ n
  return actions, (1 - p.ϵ) * π + p.ϵ * η, value_estimate
end

function reset!(p::EpsilonGreedyPlayer)
  reset!(p.player)
end

function player_temperature(p::EpsilonGreedyPlayer, game, turn)
  return player_temperature(p.player, game, turn)
end

#####
##### Player with a custom temperature
#####

"""
    PlayerWithTemperature{Game, Player} <: AbstractPlayer{Game}

A wrapper on a player that enables overwriting the temperature schedule.
"""
struct PlayerWithTemperature{G, P} <: AbstractPlayer{G}
  player :: P
  temperature :: AbstractSchedule{Float64}
  function PlayerWithTemperature(p::AbstractPlayer{G}, τ) where G
    return new{G, typeof(p)}(p, τ)
  end
end

function think(p::PlayerWithTemperature, game)
  return think(p.player, game)
end

function reset!(p::PlayerWithTemperature)
  reset!(p.player)
end

function player_temperature(p::PlayerWithTemperature, game, turn)
  return p.temperature[turn]
end

#####
##### An MCTS-based player
#####

"""
    MctsPlayer{Game, MctsEnv} <: AbstractPlayer{Game}

A player that selects actions using MCTS.

# Constructors

    MctsPlayer(mcts::MCTS.Env; τ, niters, timeout=nothing)

Construct a player from an MCTS environment. When computing each move:

- if `timeout` is provided,
  MCTS simulations are executed for `timeout` seconds by groups of `niters`
- otherwise, `niters` MCTS simulations are run

The temperature parameter `τ` can be either a real number or a
[`AbstractSchedule`](@ref).

    MctsPlayer(oracle::MCTS.Oracle, params::MctsParams; timeout=nothing)

Construct an MCTS player from an oracle and an [`MctsParams`](@ref) structure.
If the oracle is a network, this constructor handles copying it, putting it
in test mode and copying it on the GPU (if necessary).
"""
struct MctsPlayer{G, M} <: AbstractPlayer{G}
  mcts :: M
  niters :: Int
  timeout :: Union{Float64, Nothing}
  τ :: AbstractSchedule{Float64} # Temperature
  function MctsPlayer(mcts::MCTS.Env{G}; τ, niters, timeout=nothing) where G
    @assert niters > 0
    @assert isnothing(timeout) || timeout > 0
    new{G, typeof(mcts)}(mcts, niters, timeout, τ)
  end
end

# Alternative constructor
function MctsPlayer(
    oracle::MCTS.Oracle{G}, params::MctsParams; timeout=nothing) where G
  fill_batches = false
  if isa(oracle, AbstractNetwork)
    oracle = Network.copy(oracle, on_gpu=params.use_gpu, test_mode=true)
    params.use_gpu && (fill_batches = true)
  end
  mcts = MCTS.Env{G}(oracle,
    nworkers=params.num_workers,
    gamma=params.gamma,
    fill_batches=fill_batches,
    cpuct=params.cpuct,
    noise_ϵ=params.dirichlet_noise_ϵ,
    noise_α=params.dirichlet_noise_α,
    prior_temperature=params.prior_temperature)
  return MctsPlayer(mcts,
    niters=params.num_iters_per_turn,
    τ=params.temperature,
    timeout=timeout)
end

# MCTS with random oracle
function RandomMctsPlayer(::Type{G}, params::MctsParams) where G
  oracle = MCTS.RandomOracle{G}()
  mcts = MCTS.Env{G}(oracle,
    nworkers=1,
    cpuct=params.cpuct,
    gamma=params.gamma,
    noise_ϵ=params.dirichlet_noise_ϵ,
    noise_α=params.dirichlet_noise_α)
  return MctsPlayer(mcts,
    niters=params.num_iters_per_turn,
    τ=params.temperature)
end

function think(p::MctsPlayer, game)
  v = nothing
  if isnothing(p.timeout) # Fixed number of MCTS simulations
    v = MCTS.explore!(p.mcts, game, p.niters)
  else # Run simulations until timeout
    start = time()
    while time() - start < p.timeout
      v = MCTS.explore!(p.mcts, game, p.niters)
    end
  end
  actions, π = MCTS.policy(p.mcts, game)
  return actions, π, v
end

function player_temperature(p::MctsPlayer, game, turn)
  return p.τ[turn]
end

function reset_player!(player::MctsPlayer)
  MCTS.reset!(player.mcts)
end

#####
##### Network player
#####

"""
    NetworkPlayer{Game, Net} <: AbstractPlayer{Game}

A player that uses the policy output by a neural network directly,
instead of relying on MCTS.
"""
struct NetworkPlayer{G, N} <: AbstractPlayer{G}
  network :: N
  function NetworkPlayer(nn::AbstractNetwork{G}; use_gpu=false) where G
    nn = Network.copy(nn, on_gpu=use_gpu, test_mode=true)
    return new{G, typeof(nn)}(nn)
  end
end

function think(p::NetworkPlayer, game)
  actions = GI.available_actions(game)
  state = GI.current_state(game)
  π, v = MCTS.evaluate(p.network, state)
  return actions, π, v
end

#####
##### Merging two players into one
#####

"""
    TwoPlayers{Game} <: AbstractPlayer{Game}

If `white` and `black` are two [`AbstractPlayer`](@ref), then
`TwoPlayers(white, black)` is a player that behaves as `white` when `white`
is to play and as `black` when `black` is to play.
"""
struct TwoPlayers{G, W, B} <: AbstractPlayer{G}
  white :: W
  black :: B
  function TwoPlayers(white, black)
    G = GameType(white)
    @assert G == GameType(black)
    return new{G, typeof(white), typeof(black)}(white, black)
  end
end

function think(p::TwoPlayers, game)
  if GI.white_playing(game)
    return think(p.white, game)
  else
    return think(p.black, game)
  end
end

function select_move(p::TwoPlayers, game, turn)
  if GI.white_playing(game)
    return select_move(p.white, game, turn)
  else
    return select_move(p.black, game, turn)
  end
end

function reset!(p::TwoPlayers)
  reset!(p.white)
  reset!(p.black)
end

function player_temperature(p::TwoPlayers, game, turn)
  if GI.white_playing(game)
    return player_temperature(p.white, game, turn)
  else
    return player_temperature(p.black, game, turn)
  end
end

#####
##### Players can play against each other
#####

"""
    play_game(player; flip_probability=0.) :: Trace

Simulate a game by an [`AbstractPlayer`](@ref) and return a trace.

- For two-player games, please use [`TwoPlayers`](@ref).
- If the `flip_probability` argument is set to ``p``, the board
  is _flipped_ randomly at every turn with probability ``p``,
  using [`GI.apply_random_symmetry`](@ref).
"""
function play_game(player; flip_probability=0., initial_state=nothing)
  Game = GameType(player)
  if initial_state === nothing
    game = Game()
  else
    game = Game(initial_state)
  end
  initial_state = GI.current_state(game)
  trace = Trace{Game}(GI.current_state(game))
  while true
    if GI.game_terminated(game)
      return initial_state, trace
    end
    if !iszero(flip_probability) && rand() < flip_probability
      game = GI.apply_random_symmetry(game)
    end
    actions, π_target, value_estimate = think(player, game)
    τ = player_temperature(player, game, length(trace))
    π_sample = apply_temperature(π_target, τ)
    a = actions[Util.rand_categorical(π_sample)]
    GI.play!(game, a)
    push!(trace, π_target, GI.white_reward(game), value_estimate, GI.current_state(game))
  end
end

#####
##### Evaluate two players against each other
#####

"""
    @enum ColorPolicy ALTERNATE_COLORS BASELINE_WHITE CONTENDER_WHITE

Policy for attributing colors in a duel between a baseline and a contender.
"""
@enum ColorPolicy ALTERNATE_COLORS BASELINE_WHITE CONTENDER_WHITE

"""
    pit(handler, contender, baseline, ngames)

Evaluate two `AbstractPlayer` against each other in a series of games.

Note that this function can only be used with two-player games.

# Arguments

  - `handler`: this function is called after each simulated
     game with three arguments: the game number `i`, the reward `r` for the
     contender player and the trace `t`
  - `ngames`: number of games to play

# Optional keyword arguments
  - `reset_every`: if set, players are reset every `reset_every` games
  - `color_policy`: determines the [`ColorPolicy`](@ref),
     which is `ALTERNATE_COLORS` by default
  - `flip_probability=0.`: see [`play_game`](@ref)
  - `play_swiss_style=false`: if true, the tournament is played Swiss Style
"""
function pit(
    handler, contender, baseline, num_games; gamma,
    reset_every=nothing, color_policy=ALTERNATE_COLORS, flip_probability=0., play_swiss_style=false)
  Game = GameType(contender)
  @assert GI.two_players(Game)
  baseline_white = (color_policy != CONTENDER_WHITE)
  zsum = 0.
  initial_state = nothing
  for i in 1:num_games
    white = baseline_white ? baseline : contender
    black = baseline_white ? contender : baseline
    if i % 2 == 0
      initial_state, trace =
        play_game(TwoPlayers(white, black), flip_probability=flip_probability, initial_state=initial_state)
    else
      initial_state, trace =
        play_game(TwoPlayers(white, black), flip_probability=flip_probability)
    end
    z = total_reward(trace, gamma)
    baseline_white && (z = -z)
    zsum += z
    handler(i, z, trace)
    if !isnothing(reset_every) && (i % reset_every == 0 || i == num_games)
      reset_player!(baseline)
      reset_player!(contender)
    end
    if color_policy == ALTERNATE_COLORS
      baseline_white = !baseline_white
    end
  end
  return zsum / num_games
end

#####
##### Redudancy analysis
#####

function compute_redundancy(Game, states)
  # TODO: Excluding the inial state from redundancy statistics
  # only makes sense for deterministic games.
  init_state = GI.current_state(Game())
  noninit = filter(!=(init_state), states)
  unique = Set(noninit)
  return 1. - length(unique) / length(noninit)
end

#####
##### Simple utilities to play an interactive game
#####

"""
    Human{Game} <: AbstractPlayer{Game}

Human player that queries the standard input for actions.

Does not implement [`think`](@ref) but instead implements
[`select_move`](@ref) directly.
"""
struct Human{Game} <: AbstractPlayer{Game} end

struct Quit <: Exception end

function select_move(::Human, game, turn)
  a = nothing
  while isnothing(a) || a ∉ GI.available_actions(game)
    print("> ")
    str = readline()
    print("\n")
    isempty(str) && throw(Quit())
    a = GI.parse_action(game, str)
  end
  return a
end

"""
    interactive!(game, white, black)

Launch an interactive session for `game::AbstractGame` between players
`white` and `black`. Both players have type `AbstractPlayer` and one of them
is typically [`Human`](@ref).
"""
function interactive!(game, player)
  try
  GI.render(game)
  turn = 0
  while !GI.game_terminated(game)
    action = select_move(player, game, turn)
    GI.play!(game, action)
    GI.render(game)
    turn += 1
  end
  catch e
    isa(e, Quit) || rethrow(e)
    return
  end
end

interactive!(game, white, black) = interactive!(game, TwoPlayers(white, black))

interactive!(game::G) where G = interactive!(game, Human{G}(), Human{G}())


using Revise

using Flux, Zygote
using Flux.Data.MNIST
using Flux: @epochs
using Flux: throttle
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated
using Statistics

using Images: channelview
using Base.Iterators: partition

using CuArrays

include("data.jl")
include("model.jl")
include("train.jl")

function test_BatchNorm()
    using Flux
    using CuArrays
    bn = BatchNorm(32);
    testimg = randn(Float32, 28, 28, 32, 1);
    @time bn(testimg);
    @time gpu(bn)(gpu(testimg));

    Flux.testmode!(bn)
    Flux.testmode!(bn, false)
    sum(bn(testimg))

    p = Flux.param(testimg);

    gs = Tracker.gradient(()->sum(bn(p)), Flux.params(bn))
    gs

    function f()
        sum(bn(p))
    end
    Flux.gradient(f, p)
end

function test_allowscalar()
    CuArrays.allowscalar(false)
    x = rand(Float32, 10, 3) |> gpu;
    y = Flux.onehotbatch(1:3, 1:10) |> gpu;
    accuracy(x, y) = Flux.onecold(x) .== Flux.onecold(y);
    accuracy(x, y)
    CuArrays.allowscalar(true)
end

function test_bn_grad()
    x = randn(Float32, 32, 32, 3, 16) |> gpu
    model = Chain(
        myConv((3,3), 3=>16, stride=(1,1), pad=1),
        BatchNorm(16),
        MeanPool((8,8)),
        x -> reshape(x, :, size(x)[end]),
        Dense(256,10)) |> gpu
    model(x)

    Flux.gradient() do
        sum(model(x))
    end

    Flux.testmode!(model)

    # In test mode, it would fail if I don't patch Flux.jl/src/cuda/cudnn.jl
    # cudnnBNBackward!. The root cause is sum(iter; dims) have dims as explicit
    # kwarg, and thus calling without dims=... will call the wrong sum. Also
    # squeeze is deprecated for dropdims. NOTE: this is in Flux v0.9, the master
    # branch currently removed manual testmode, see discussions:
    #
    # - Train/test mode discussion by MikeInnes https://github.com/FluxML/Flux.jl/issues/643
    # - https://github.com/FluxML/Flux.jl/issues/909
    # - https://github.com/FluxML/Flux.jl/issues/232
    Flux.gradient() do
        sum(model(x))
    end
end

# TODO learning rate?
function test_CIFAR10_ds()
    ds, test_ds = load_CIFAR10_ds(batch_size=128);
    x, y = next_batch!(ds) |> gpu;

    # model = get_CIFAR_CNN_model()
    model = resnet(20) |> gpu
    # model = resnet(32)
    # model = resnet(56)
    # model = resnet(68)

    # model = WRN(28, 10) |> gpu;
    model = WRN(16, 4) |> gpu;

    model(x)

    adv = attack_CIFAR10_PGD_k(7)(model, x, y);
    Flux.testmode!(model)
    adv = attack_CIFAR10_PGD_k(7)(model, x, y);

    size(adv)

    augment = Augment()
    augx = augment(cpu(x)) |> gpu;
    # augment(x)
    model(augx)
    accuracy_with_logits(model(augx), y)

    opt = ADAM(1e-3)

    train!(model, opt, ds, print_steps=50)

    accuracy_with_logits(model(x), y)

end

function test_free_train()
    ds, test_ds = load_MNIST_ds(batch_size=50);
    x, y = next_batch!(ds) |> gpu;

    model = get_Madry_model()
    model(x)
    opt = ADAM(1e-3)

    size(gpu(zeros(size(x)...)) + x)

    free_train!(model, opt, ds, 0.3)
end

function evaluate_AE(ae, ds; cnn=nothing)
    xx,yy = next_batch!(ds) |> gpu

    @info "Sampling BEFORE images .."
    sample_and_view(xx, yy, cnn)
    @info "Sampling AFTER images .."
    # NOTE: I need .data, currently handled in sample_and_view
    rec = ae(xx)
    sample_and_view(rec, yy, cnn)
end

function test_CIFAR_ae()
    ds, test_ds = load_CIFAR10_ds(batch_size=128);
    x, y = next_batch!(ds) |> gpu;
    x, y = next_batch!(test_ds) |> gpu;

    # get a pretrained WRN
    cnn = CIFAR10_pretrained_fn()
    model = CIFAR10_pretrained_fn()

    accuracy_with_logits(cnn(x), y)
    accuracy_with_logits(model(x), y)
    Flux.testmode!(model)
    Flux.testmode!(model, false)
    Flux.testmode!(cnn)
    Flux.testmode!(cnn, false)

    params(cnn[2].μ)

    cnn(x)

    onecold(cpu(cnn(x)))

    model[2].μ
    cnn[2].μ
    cnn[2].σ²

    cnn

    train!(cnn, opt, ds)

    @save "test.bson" MMM=cpu(cnn) opt
    @load "test.bson" MMM

    MMM=gpu(MMM);

    fieldnames(typeof(cnn[5]))
    cnn[5].norm_layers[1].σ²

    # get the AE

    # ae = cifar10_AE()
    ae = dunet()
    encoder, decoder = cifar10_deep_AE()

    ae = Chain(encoder, decoder)

    rx = randn(Float32, 32, 32, 3, 16) |> gpu;
    size(encoder(rx))
    size(decoder(encoder(rx)))
    size(ae(rx))

    opt = ADAM(1e-3)
    # train the AE
    aetrain!(ae, opt, ds)

    evaluate_AE(ae, test_ds, cnn=cnn)

    param_count(ae)
    param_count(cnn)
end

function test_AE()
    ds, test_ds = load_MNIST_ds(batch_size=50);
    x, y = next_batch!(ds) |> gpu;

    cnn = get_Madry_model()
    @show size(cnn(x))

    opt = ADAM(1e-3)

    Flux.testmode!(cnn, false)
    cnn = MNIST_pretrained_fn()

    accuracy_with_logits(cnn(x), y)

    cnn(x)

    size(cnn[1](x))

    # TODO schedule lr simply by setting field of opt, and maintain the
    # state. FIXME But it seems that the state is maintained by a dict with key
    # Param(x). I'm not sure if this aproach can resume the state.
    #
    # opt.eta = 1e-4

    train!(cnn, opt, ds)

    cnn = maybe_train("trained/pretrain-MNIST.bson", get_Madry_model(), train_fn)

    # ae = dense_AE()
    ae = CNN_AE()
    # ae = CNN2_AE()
    size(ae(x))

    # FIXME opt states?
    opt = ADAM(1e-3)
    opt = ADAM(1e-4)
    # FIXME Flux.mse(logits, x)
    # FIXME sigmoid mse?
    aetrain!(ae, opt, ds)

    evaluate_AE(ae, test_ds, cnn=cnn)

    advae_train!(ae, cnn, opt, attack_PGD_k(40), ds, print_steps=20)

    evaluate(cnn, test_ds)
    evaluate(cnn, test_ds, attack_fn=attack_PGD_k(40))

    evaluate(Chain(ae, cnn), test_ds)
    evaluate(Chain(ae, cnn), test_ds, attack_fn=attack_FGSM)
    evaluate(Chain(ae, cnn), test_ds, attack_fn=attack_PGD_k(40))
end

function test_log_image()
    logger = TBLogger("tensorboard_logs/test/image", tb_append, min_level=Logging.Info)

    ds, test_ds = load_CIFAR10_ds(batch_size=128);
    x, y = next_batch!(ds) |> gpu;
    img = x[:,:,:,1];

    # CAUTION if pass in raw img, the error message contains the whole array,
    # which is very slow
    showable("image/png", MyImage(img))
    log_image(logger, "testimage", MyImage(img));

    display(MyImage(img))
end

function sha_arr(x)
    # assuming 3 channel images in a batch
    res = map(collect(1:size(x)[end])) do i
        img = x[:,:,:,i]
        bytes2hex(sha256(repr(img)))
    end
    res
end

function get_wrong_sha(model, ds)
    # evaluate model, and get the wrong sample shas
    res = []
    @showprogress for i in 1:ds.nbatch
        x, y = next_batch!(ds) |> gpu
        logits = model(x)
        wrong_idx = onecold(cpu(logits)) .!= onecold(cpu(y))
        if length(wrong_idx) > 0
            # get sha
            shas = sha_arr(cpu(x)[:,:,:,wrong_idx])
            push!(res, shas...)
        end
    end
    res
end

function load_trained_model(file)
    @load file model
    return gpu(model)
end

# https://github.com/FluxML/Flux.jl/issues/160
function weight_params(m::Chain, ps=Flux.Params())
    map((l)->weight_params(l, ps), m.layers)
    ps
end
weight_params(m::Dense, ps=Flux.Params()) = push!(ps, m.W)
weight_params(m::Conv, ps=Flux.Params()) = push!(ps, m.weight)
weight_params(m::ConvTranspose, ps=Flux.Params()) = push!(ps, m.weight)
weight_params(m, ps=Flux.Params()) = ps

function test_weight_params()
    # get model
    weight_params(model)
end

function test_boundary()
    # test whether all the nail households
    model = CIFAR10_pretrained_fn()
    ds, test_ds = load_CIFAR10_ds(batch_size=128);

    itM = load_trained_model(
        "trained/CIFAR10/itadv-schedule/(0.0004=>40000,4.0e-5=>60000,4.0e-6=>80000).bson")
    dyM = load_trained_model(
        "trained/CIFAR10-dyattack/(0=>2000,1=>3000,2=>4000,3=>5000,4=>6000,7=>8000).bson")
    freeM = load_trained_model("trained/CIFAR10-free/test-2019-11-18T16:57:55.947.bson")

    sha_it = get_wrong_sha(itM, test_ds);
    sha_dy = get_wrong_sha(dyM, test_ds);
    sha_free = get_wrong_sha(freeM, test_ds);

    size(sha_it)
    size(unique(sha_it))
    @show length(sha_it) length(sha_dy) length(sha_free)
    @show length(unique(sha_it)) length(unique(sha_dy)) length(unique(sha_free))
    @show size(intersect(sha_it, sha_dy, sha_free))
    @show size(intersect(sha_it, sha_dy, sha_free))

    ps = params(model);

    model[1].weight in ps

    model

    Flux.mapparams(model) do p
        @show typeof(p)
        p
    end

    model(x)
    typeof(ps)
    typeof(keys(ps.params.dict))
    keytype(ps.params.dict)

    shas[1]
    size(shas)
    size(unique(shas))
    size(unique(shas2))
    # overlap
    size(intersect(shas, shas2))
end

function test()
    ds, test_ds = load_CIFAR10_ds(batch_size=128);
    x, y = next_batch!(ds) |> gpu;
    x, y = next_batch!(test_ds) |> gpu;
    model = load()
    model[2].μ
    # Flux.testmode!(model, false)
    model(x)

    # hash the x
    img = x[:,:,:,1];
    size(img)

    # using SHA
    @time bytes2hex(sha256("test"))

    @time bytes2hex(sha256(repr(img)))

    cat(res..., dims=4)
    size(res)
    res[1]
    size(unique(res))
    res[1]

    ss = map(x) do img
        bytes2hex(sha256(repr(img)))
    end

    for i in 1:10
        i
    end

    accuracy_with_logits(model(x), y)

    sum(onecold(cpu(model(x))) .== onecold(cpu(y)))

    adv = attack_CIFAR10_PGD_k(7)(model, x, y);

    view_preds(model, x, y)
    view_preds(model, adv, y)

    sample_and_view(wrong, cpu(y)[])
    sample_and_view(right)
    displayable("image/png")
end

function view_preds(model, x, y)
    @show accuracy_with_logits(model(x), y)
    right_idx = onecold(cpu(model(x))) .== onecold(cpu(y));
    wrong_idx = onecold(cpu(model(x))) .!= onecold(cpu(y));
    @info "right predictions:"
    sample_and_view(cpu(x)[:,:,:,right_idx], cpu(y)[:,right_idx], model)
    @info "wrong predictions:"
    sample_and_view(cpu(x)[:,:,:,wrong_idx], cpu(y)[:,wrong_idx], model)
    nothing
end


# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

# This is the builder for the library Dierckx in Netlib
# (http://www.netlib.org/dierckx/), which doesn't really have a versin number.
name = "Dierckx"
version = v"0.0.1"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/scipy/scipy/releases/download/v1.5.1/scipy-1.5.1.tar.xz",
                  "0728bd66a5251cfeff17a72280ae5a40ec14add217f94868d1415b3c469b610a"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/scipy-*/scipy/interpolate/fitpack
mkdir -p "${libdir}"
gfortran -o "${libdir}/libddierckx.${dlext}" -O3 -shared -fPIC *.f
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = expand_gfortran_versions(supported_platforms())

# The products that we will ensure are always built
products = [
    LibraryProduct("libddierckx", :libddierckx),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="CompilerSupportLibraries_jll", uuid="e66e0078-7015-5450-92f7-15fbd957f2ae")),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


module GeoMakie

using Reexport
using GeometryBasics, Colors, AbstractPlotting.MakieLayout, AbstractPlotting, ImageMagick
using GeoJSON

using AbstractPlotting.FileIO
using AbstractPlotting.MakieLayout: Formatting
using AbstractPlotting.MakieLayout.GridLayoutBase
using AbstractPlotting.MakieLayout.GridLayoutBase: Side
using GeoInterface: GeoInterface, coordinates, AbstractPolygon, AbstractMultiPolygon, features, geometry

using Proj4
import Proj4: Projection
import AbstractPlotting: convert_arguments, convert_attribute
using AbstractPlotting.DocStringExtensions

# fix conflicts
const AbstractGeometry = GeometryBasics.AbstractGeometry
const Point = AbstractPlotting.Point
const attributes = AbstractPlotting.attributes
const transform! = Proj4.transform!
const volume = AbstractPlotting.volume
const Mesh = GeometryBasics.Mesh
const Text = AbstractPlotting.Text

include("conversions.jl")
include("projections.jl")
include("data.jl")

export LonLat, PlateCarree, Mercator, TransverseMercator, UTM, AlbersEqualArea, Robinson, WinkelTripel

@reexport using Colors, AbstractPlotting
@reexport using Proj4

end # module


#=
MathOptInterface wrapper of Hypatia solver
=#

const SV = MOI.SingleVariable
const SAF = MOI.ScalarAffineFunction
const VV = MOI.VectorOfVariables
const VAF = MOI.VectorAffineFunction

export Optimizer

"""
$(TYPEDEF)

A MathOptInterface optimizer type for Hypatia.
"""
mutable struct Optimizer{T <: Real} <: MOI.AbstractOptimizer
    use_dense_model::Bool # make the model use dense A and G data instead of sparse

    solver::Solvers.Solver{T} # Hypatia solver object
    model::Models.Model{T} # Hypatia model object

    # result data
    x::Vector{T}
    s::Vector{T}
    y::Vector{T}
    z::Vector{T}

    # data for transforming certificates
    obj_sense::MOI.OptimizationSense
    num_eq_constrs::Int
    constr_offset_eq::Vector{Int}
    constr_prim_eq::Vector{T}
    constr_offset_cone::Vector{Int}
    constr_prim_cone::Vector{T}
    nonpos_idxs::UnitRange{Int}
    interval_idxs::UnitRange{Int}
    interval_scales::Vector{T}
    moi_other_cones_start::Int
    moi_other_cones::Vector{MOI.AbstractVectorSet}

    function Optimizer{T}(;
        use_dense_model::Bool = true, # TODO should depend on the size and sparsity of A, G in the model
        solver_options... # TODO allow passing in a solver?
        ) where {T <: Real}
        opt = new{T}()
        opt.use_dense_model = use_dense_model
        if !haskey(solver_options, :syssolver)
            # choose default system solver based on use_dense_model
            sstype = (use_dense_model ? Solvers.QRCholDenseSystemSolver :
                Solvers.SymIndefSparseSystemSolver)
            solver_options = (solver_options..., syssolver = sstype{T}())
        end
        if !haskey(solver_options, :preprocess)
            # only preprocess if using dense model # TODO maybe should preprocess if sparse
            solver_options = (solver_options..., preprocess = use_dense_model)
        end
        if !haskey(solver_options, :reduce)
            # only reduce if using dense model
            solver_options = (solver_options..., reduce = use_dense_model)
        end
        if !haskey(solver_options, :init_use_indirect)
            # only use indirect if not using dense model
            solver_options = (solver_options...,
                init_use_indirect = !use_dense_model)
        end
        opt.solver = Solvers.Solver{T}(; solver_options...)
        return opt
    end
end

Optimizer(; options...) = Optimizer{Float64}(; options...) # default to Float64

MOI.get(::Optimizer, ::MOI.SolverName) = "Hypatia"
MOI.get(opt::Optimizer, ::MOI.RawSolver) = opt.solver

MOI.is_empty(opt::Optimizer) = (opt.solver.status == Solvers.NotLoaded)

MOI.empty!(opt::Optimizer) = (opt.solver.status = Solvers.NotLoaded)

MOI.supports(
    ::Optimizer{T},
    ::Union{
        MOI.ObjectiveSense,
        MOI.ObjectiveFunction{SV},
        MOI.ObjectiveFunction{SAF{T}},
        },
    ) where {T <: Real} = true

MOI.supports_constraint(
    ::Optimizer{T},
    ::Type{<:Union{SV, SAF{T}}},
    ::Type{<:Union{
        MOI.EqualTo{T},
        MOI.GreaterThan{T},
        MOI.LessThan{T},
        MOI.Interval{T},
        }},
    ) where {T <: Real} = true

MOI.supports_constraint(
    ::Optimizer{T},
    ::Type{<:Union{VV, VAF{T}}},
    ::Type{<:SupportedCones{T}},
    ) where {T <: Real} = true

# build representation as min c'x s.t. A*x = b, h - G*x in K
function MOI.copy_to(
    opt::Optimizer{T},
    src::MOI.ModelLike;
    copy_names::Bool = false,
    ) where {T <: Real}
    @assert !copy_names
    idx_map = MOI.Utilities.IndexMap()

    # variables
    n = MOI.get(src, MOI.NumberOfVariables()) # columns of A
    j = 0
    for vj in MOI.get(src, MOI.ListOfVariableIndices()) # MOI.VariableIndex
        j += 1
        idx_map[vj] = MOI.VariableIndex(j)
    end
    @assert j == n

    # objective function
    F = MOI.get(src, MOI.ObjectiveFunctionType())
    if F == SV
        obj = SAF{T}(MOI.get(src, MOI.ObjectiveFunction{F}()))
    elseif F == SAF{T}
        obj = MOI.get(src, MOI.ObjectiveFunction{F}())
    else
        error("function type $F not supported")
    end
    (Jc, Vc) = (Int[], T[])
    for t in obj.terms
        push!(Jc, idx_map[t.variable_index].value)
        push!(Vc, t.coefficient)
    end
    obj_offset = obj.constant
    if MOI.get(src, MOI.ObjectiveSense()) == MOI.MAX_SENSE
        Vc .*= -1
        obj_offset *= -1
    end
    opt.obj_sense = MOI.get(src, MOI.ObjectiveSense())
    model_c = Vector(sparsevec(Jc, Vc, n))

    # constraints
    get_src_cons(F, S) = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
    get_con_fun(con_idx) = MOI.get(src, MOI.ConstraintFunction(), con_idx)
    get_con_set(con_idx) = MOI.get(src, MOI.ConstraintSet(), con_idx)
    i = 0 # MOI constraint index

    # equality constraints
    p = 0 # rows of A (equality constraint matrix)
    (IA, JA, VA) = (Int[], Int[], T[])
    (Ib, Vb) = (Int[], T[])
    (Icpe, Vcpe) = (Int[], T[]) # constraint set constants for opt.constr_prim_eq
    constr_offset_eq = Vector{Int}()

    for F in (SV, SAF{T}), ci in get_src_cons(F, MOI.EqualTo{T})
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.EqualTo{T}}(i)
        push!(constr_offset_eq, p)
        p += 1
        fi = get_con_fun(ci)
        si = get_con_set(ci)
        if F == SV
            push!(IA, p)
            push!(JA, idx_map[fi.variable].value)
            push!(VA, -1)
            push!(Vb, -si.value)
        else
            for vt in fi.terms
                push!(IA, p)
                push!(JA, idx_map[vt.variable_index].value)
                push!(VA, -vt.coefficient)
            end
            push!(Vb, fi.constant - si.value)
        end
        push!(Ib, p)
        push!(Icpe, p)
        push!(Vcpe, si.value)
    end

    for F in (VV, VAF{T}), ci in get_src_cons(F, MOI.Zeros)
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.Zeros}(i)
        push!(constr_offset_eq, p)
        fi = get_con_fun(ci)
        dim = MOI.output_dimension(fi)
        if F == VV
            append!(IA, (p + 1):(p + dim))
            append!(JA, idx_map[vi].value for vi in fi.variables)
            append!(VA, -ones(T, dim))
        else
            for vt in fi.terms
                push!(IA, p + vt.output_index)
                push!(JA, idx_map[vt.scalar_term.variable_index].value)
                push!(VA, -vt.scalar_term.coefficient)
            end
            append!(Ib, (p + 1):(p + dim))
            append!(Vb, fi.constants)
        end
        p += dim
    end

    push!(constr_offset_eq, p)
    model_A = dropzeros!(sparse(IA, JA, VA, p, n))
    model_b = Vector(sparsevec(Ib, Vb, p))
    opt.num_eq_constrs = i
    opt.constr_prim_eq = Vector(sparsevec(Icpe, Vcpe, p))
    opt.constr_offset_eq = constr_offset_eq

    # conic constraints
    q = 0 # rows of G (cone constraint matrix)
    (IG, JG, VG) = (Int[], Int[], T[])
    (Ih, Vh) = (Int[], T[])
    (Icpc, Vcpc) = (Int[], T[]) # constraint set constants for opt.constr_prim_eq
    constr_offset_cone = Vector{Int}()
    cones = Cones.Cone{T}[]

    # build up one nonnegative cone from LP constraints
    nonneg_start = q

    # nonnegative-like constraints
    for F in (SV, SAF{T}), ci in get_src_cons(F, MOI.GreaterThan{T})
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.GreaterThan{T}}(i)
        push!(constr_offset_cone, q)
        q += 1
        fi = get_con_fun(ci)
        si = get_con_set(ci)
        if F == SV
            push!(IG, q)
            push!(JG, idx_map[fi.variable].value)
            push!(VG, -1)
            push!(Vh, -si.lower)
        else
            for vt in fi.terms
                push!(IG, q)
                push!(JG, idx_map[vt.variable_index].value)
                push!(VG, -vt.coefficient)
            end
            push!(Vh, fi.constant - si.lower)
        end
        push!(Ih, q)
        push!(Vcpc, si.lower)
        push!(Icpc, q)
    end

    for F in (VV, VAF{T}), ci in get_src_cons(F, MOI.Nonnegatives)
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.Nonnegatives}(i)
        push!(constr_offset_cone, q)
        fi = get_con_fun(ci)
        if F == VV
            for vj in fi.variables
                q += 1
                push!(IG, q)
                push!(JG, idx_map[vj].value)
                push!(VG, -1)
            end
        else
            dim = MOI.output_dimension(fi)
            for vt in fi.terms
                push!(IG, q + vt.output_index)
                push!(JG, idx_map[vt.scalar_term.variable_index].value)
                push!(VG, -vt.scalar_term.coefficient)
            end
            append!(Ih, (q + 1):(q + dim))
            append!(Vh, fi.constants)
            q += dim
        end
    end

    # nonpositive-like constraints
    nonpos_start = q

    for F in (SV, SAF{T}), ci in get_src_cons(F, MOI.LessThan{T})
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.LessThan{T}}(i)
        push!(constr_offset_cone, q)
        q += 1
        fi = get_con_fun(ci)
        si = get_con_set(ci)
        if F == SV
            push!(IG, q)
            push!(JG, idx_map[fi.variable].value)
            push!(VG, 1)
            push!(Vh, si.upper)
        else
            for vt in fi.terms
                push!(IG, q)
                push!(JG, idx_map[vt.variable_index].value)
                push!(VG, vt.coefficient)
            end
            push!(Vh, -fi.constant + si.upper)
        end
        push!(Ih, q)
        push!(Vcpc, si.upper)
        push!(Icpc, q)
    end

    for F in (VV, VAF{T}), ci in get_src_cons(F, MOI.Nonpositives)
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.Nonpositives}(i)
        push!(constr_offset_cone, q)
        fi = get_con_fun(ci)
        if F == VV
            for vj in fi.variables
                q += 1
                push!(IG, q)
                push!(JG, idx_map[vj].value)
                push!(VG, 1)
            end
        else
            dim = MOI.output_dimension(fi)
            for vt in fi.terms
                push!(IG, q + vt.output_index)
                push!(JG, idx_map[vt.scalar_term.variable_index].value)
                push!(VG, vt.scalar_term.coefficient)
            end
            append!(Ih, (q + 1):(q + dim))
            append!(Vh, -fi.constants)
            q += dim
        end
    end

    # single nonnegative cone
    opt.nonpos_idxs = (nonpos_start + 1):q
    if q > nonneg_start
        push!(cones, Cones.Nonnegative{T}(q - nonneg_start))
    end

    # build up one L_infinity norm cone from two-sided interval constraints
    interval_start = q
    SV_intvl = get_src_cons(SV, MOI.Interval{T})
    SAF_intvl = get_src_cons(SAF{T}, MOI.Interval{T})
    num_intervals = length(SV_intvl) + length(SAF_intvl)
    interval_scales = zeros(T, num_intervals)

    if num_intervals > 0
        i += 1
        push!(constr_offset_cone, q)
        q += 1
        push!(Ih, q)
        push!(Vh, 1)
    end

    interval_count = 0
    for (F, Cs) in ((SV, SV_intvl), (SAF{T}, SAF_intvl)), ci in Cs
        i += 1
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.Interval{T}}(i)
        push!(constr_offset_cone, q)
        q += 1

        upper = get_con_set(ci).upper
        lower = get_con_set(ci).lower
        @assert isfinite(upper) && isfinite(lower)
        @assert upper > lower
        mid = (upper + lower) / T(2)
        scal = 2 / (upper - lower)

        fi = get_con_fun(ci)
        if F == SV
            push!(IG, q)
            push!(JG, idx_map[fi.variable].value)
            push!(VG, -scal)
            push!(Vh, -mid * scal)
        else
            for vt in fi.terms
                push!(IG, q)
                push!(JG, idx_map[vt.variable_index].value)
                push!(VG, -vt.coefficient * scal)
            end
            push!(Vh, (fi.constant - mid) * scal)
        end
        push!(Ih, q)
        push!(Vcpc, mid)
        push!(Icpc, q)
        interval_count += 1
        interval_scales[interval_count] = scal
    end

    opt.interval_idxs = (interval_start + 2):q
    opt.interval_scales = interval_scales
    if q > interval_start
        # exists at least one interval-type constraint
        push!(cones, Cones.EpiNormInf{T, T}(q - interval_start))
    end

    # non-LP conic constraints
    moi_other_cones_start = i + 1
    moi_other_cones = MOI.AbstractVectorSet[]

    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        if S <: LinearCones{T}
            continue # already copied these constraints
        end
        @assert S <: SupportedCones{T}

        for ci in get_src_cons(F, S)
            i += 1
            idx_map[ci] = MOI.ConstraintIndex{F, S}(i)
            push!(constr_offset_cone, q)
            fi = get_con_fun(ci)
            si = get_con_set(ci)
            push!(moi_other_cones, si)
            dim = MOI.output_dimension(fi)

            if F == VV
                JGi = (idx_map[vj].value for vj in fi.variables)
                IGi = permute_affine(si, 1:dim)
                VGi = rescale_affine(si, fill(-one(T), dim))
            else
                JGi = (idx_map[vt.scalar_term.variable_index].value
                    for vt in fi.terms)
                IGi = permute_affine(si, [vt.output_index for vt in fi.terms])
                VGi = rescale_affine(si, [-vt.scalar_term.coefficient
                    for vt in fi.terms], IGi)
                Ihi = permute_affine(si, 1:dim)
                Vhi = rescale_affine(si, fi.constants)
                Ihi = q .+ Ihi
                append!(Ih, Ihi)
                append!(Vh, Vhi)
            end

            IGi = q .+ IGi
            append!(IG, IGi)
            append!(JG, JGi)
            append!(VG, VGi)
            push!(cones, cone_from_moi(T, si))
            q += dim
        end
    end

    push!(constr_offset_cone, q)

    # finalize model
    model_G = dropzeros!(sparse(IG, JG, VG, q, n))
    model_h = Vector(sparsevec(Ih, Vh, q))

    opt.model = Models.Model{T}(model_c, model_A, model_b, model_G, model_h,
        cones; obj_offset = obj_offset)
    if opt.use_dense_model # convert A and G to dense
        Models.densify!(opt.model)
    end

    opt.constr_offset_cone = constr_offset_cone
    opt.constr_prim_cone = Vector(sparsevec(Icpc, Vcpc, q))
    opt.moi_other_cones_start = moi_other_cones_start
    opt.moi_other_cones = moi_other_cones

    return idx_map
end

function MOI.optimize!(opt::Optimizer{T}) where {T <: Real}
    # build and solve the model
    model = opt.model
    solver = opt.solver

    Solvers.load(solver, model)
    Solvers.solve(solver)

    status = Solvers.get_status(solver)
    primal_obj = Solvers.get_primal_obj(solver)
    dual_obj = Solvers.get_dual_obj(solver)
    opt.x = Solvers.get_x(solver)
    opt.y = Solvers.get_y(solver)
    opt.s = Solvers.get_s(solver)
    opt.z = Solvers.get_z(solver)

    # transform solution for MOI conventions
    opt.constr_prim_eq += model.b - model.A * opt.x
    opt.s[opt.nonpos_idxs] .*= -1
    opt.z[opt.nonpos_idxs] .*= -1
    opt.s[opt.interval_idxs] ./= opt.interval_scales
    i = opt.moi_other_cones_start - opt.num_eq_constrs
    for cone in opt.moi_other_cones
        if needs_untransform(cone)
            os = opt.constr_offset_cone
            idxs = (os[i] + 1):os[i + 1]
            @views untransform_affine(cone, opt.s[idxs])
            @views untransform_affine(cone, opt.z[idxs])
        end
        i += 1
    end
    opt.constr_prim_cone .+= opt.s
    opt.z[opt.interval_idxs] .*= opt.interval_scales

    return
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
MOI.set(opt::Optimizer, ::MOI.Silent, value::Bool) = (opt.solver.verbose = !value)
MOI.get(opt::Optimizer, ::MOI.Silent) = !opt.solver.verbose

MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
MOI.set(opt::Optimizer, ::MOI.TimeLimitSec, value::Real) =
    (opt.solver.time_limit = value)
MOI.set(opt::Optimizer, ::MOI.TimeLimitSec, ::Nothing) =
    (opt.solver.time_limit = Inf)

function MOI.get(opt::Optimizer, ::MOI.TimeLimitSec)
    isfinite(opt.solver.time_limit) && return opt.solver.time_limit
    return
end

function MOI.get(opt::Optimizer, ::MOI.SolveTime)
    if opt.solver.status in (:NotLoaded, :Loaded)
        error("solve has not been called")
    end
    return Solvers.get_solve_time(opt.solver)
end

MOI.get(opt::Optimizer, ::MOI.RawStatusString) = string(opt.solver.status)

MOI.get(opt::Optimizer, ::MOI.BarrierIterations) = opt.solver.num_iters

function MOI.get(opt::Optimizer, ::MOI.TerminationStatus)
    status = opt.solver.status
    if status in (Solvers.NotLoaded, Solvers.Loaded)
        return MOI.OPTIMIZE_NOT_CALLED
    elseif status == Solvers.Optimal
        return MOI.OPTIMAL
    elseif status == Solvers.PrimalInfeasible
        return MOI.INFEASIBLE
    elseif status == Solvers.DualInfeasible
        return MOI.DUAL_INFEASIBLE
    elseif status == Solvers.SlowProgress
        return MOI.SLOW_PROGRESS
    elseif status == Solvers.IterationLimit
        return MOI.ITERATION_LIMIT
    elseif status == Solvers.TimeLimit
        return MOI.TIME_LIMIT
    else
        @warn("Hypatia status $(opt.solver.status) not handled")
        return MOI.OTHER_ERROR
    end
end

function MOI.get(opt::Optimizer, ::MOI.PrimalStatus)
    status = opt.solver.status
    if status == Solvers.Optimal
        return MOI.FEASIBLE_POINT
    elseif status == Solvers.PrimalInfeasible
        return MOI.INFEASIBLE_POINT
    elseif status == Solvers.DualInfeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == Solvers.IllPosed
        return MOI.OTHER_RESULT_STATUS
    else
        return MOI.UNKNOWN_RESULT_STATUS
    end
end

function MOI.get(opt::Optimizer, ::MOI.DualStatus)
    status = opt.solver.status
    if status == Solvers.Optimal
        return MOI.FEASIBLE_POINT
    elseif status == Solvers.PrimalInfeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == Solvers.DualInfeasible
        return MOI.INFEASIBLE_POINT
    elseif status == Solvers.IllPosed
        return MOI.OTHER_RESULT_STATUS
    else
        return MOI.UNKNOWN_RESULT_STATUS
    end
end

function MOI.get(opt::Optimizer, ::MOI.ObjectiveValue)
    raw_obj_val = Solvers.get_primal_obj(opt.solver)
    return ((opt.obj_sense == MOI.MAX_SENSE) ? -1 : 1) * raw_obj_val
end

function MOI.get(
    opt::Optimizer,
    ::Union{MOI.DualObjectiveValue, MOI.ObjectiveBound},
    )
    raw_dual_obj_val = Solvers.get_dual_obj(opt.solver)
    return ((opt.obj_sense == MOI.MAX_SENSE) ? -1 : 1) * raw_dual_obj_val
end

MOI.get(opt::Optimizer, ::MOI.ResultCount) = 1

MOI.get(
    opt::Optimizer,
    ::MOI.VariablePrimal,
    vi::MOI.VariableIndex,
    ) = opt.x[vi.value]

MOI.get(
    opt::Optimizer,
    a::MOI.VariablePrimal,
    vi::Vector{MOI.VariableIndex},
    ) = MOI.get.(opt, a, vi)

function MOI.get(
    opt::Optimizer,
    ::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{F, S},
    ) where {F <: MOI.AbstractFunction, S <: MOI.AbstractScalarSet}
    # scalar set
    i = ci.value
    if i <= opt.num_eq_constrs
        # constraint is an equality
        return opt.y[opt.constr_offset_eq[i] + 1]
    else
        # constraint is conic
        i -= opt.num_eq_constrs
        return opt.z[opt.constr_offset_cone[i] + 1]
    end
end

function MOI.get(
    opt::Optimizer,
    ::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{F, S},
    ) where {F <: MOI.AbstractFunction, S <: MOI.AbstractVectorSet}
    # vector set
    i = ci.value
    if i <= opt.num_eq_constrs
        # constraint is an equality
        os = opt.constr_offset_eq
        return opt.y[(os[i] + 1):os[i + 1]]
    else
        # constraint is conic
        i -= opt.num_eq_constrs
        os = opt.constr_offset_cone
        return opt.z[(os[i] + 1):os[i + 1]]
    end
end

MOI.get(
    opt::Optimizer,
    a::MOI.ConstraintDual,
    ci::Vector{MOI.ConstraintIndex},
    ) = MOI.get.(opt, a, ci)

function MOI.get(
    opt::Optimizer,
    ::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{F, S},
    ) where {F <: MOI.AbstractFunction, S <: MOI.AbstractScalarSet}
    # scalar set
    i = ci.value
    if i <= opt.num_eq_constrs
        # constraint is an equality
        return opt.constr_prim_eq[opt.constr_offset_eq[i] + 1]
    else
        # constraint is conic
        i -= opt.num_eq_constrs
        return opt.constr_prim_cone[opt.constr_offset_cone[i] + 1]
    end
end

function MOI.get(
    opt::Optimizer,
    ::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{F, S},
    ) where {F <: MOI.AbstractFunction, S <: MOI.AbstractVectorSet}
    # vector set
    i = ci.value
    if i <= opt.num_eq_constrs
        # constraint is an equality
        os = opt.constr_offset_eq
        return opt.constr_prim_eq[(os[i] + 1):os[i + 1]]
    else
        # constraint is conic
        i -= opt.num_eq_constrs
        os = opt.constr_offset_cone
        return opt.constr_prim_cone[(os[i] + 1):os[i + 1]]
    end
end

MOI.get(
    opt::Optimizer,
    a::MOI.ConstraintPrimal,
    ci::Vector{MOI.ConstraintIndex},
    ) = MOI.get.(opt, a, ci)


using Base.Test
import COBRAUtils

@testset "E. Coli" begin
    model = COBRAUtils.loadEColiTestModel()

    @test COBRAUtils.metabolite_index(model, "6pgc[c]") == 4
    @test COBRAUtils.reaction_index(model, "ACONTa") == 4
    @test sort(model.mets[COBRAUtils.reaction_metabolites(model, "ACALD")]) == sort(["acald[c]", "accoa[c]", "coa[c]", "h[c]", "nad[c]", "nadh[c]"])
    @test sort(model.mets[COBRAUtils.reaction_substrates(model, "ACALD")]) == sort(["acald[c]", "coa[c]", "nad[c]"])
    @test sort(model.mets[COBRAUtils.reaction_products(model, "ACALD")]) == sort(["accoa[c]", "h[c]", "nadh[c]"])
    @test sort(model.rxns[COBRAUtils.metabolite_reactions(model, "ac[c]")]) == sort(["ACt2r", "ACKr"])

    @test COBRAUtils.reaction_substrates(model, "EX_ac(e)") == [7]
    @test isempty(COBRAUtils.reaction_products(model, "EX_ac(e)"))

    @test COBRAUtils.exchange_reactions(model, "ac[e]") == [20]

    status, μ, v = COBRAUtils.optimize(model)
    @test μ ≈ 0.8739215069684305
end

@testset "reduce_model E. Coli" begin
    ecoli = COBRAUtils.loadEColiTestModel();
    ecolired = COBRAUtils.reduce_model(ecoli);
    @test size(ecolired.S) == (68, 87)
    status0, μ0, v0 = COBRAUtils.optimize(ecoli)
    status1, μ1, v1 = COBRAUtils.optimize(ecolired)
    @test μ0 ≈ μ1
end


# This file is the successor to "images/GenerateImages.jl" (v0.5.1).
module Colormaps

using Colors

struct ColormapSVG <: Main.SVG
    buf::IOBuffer
end

function ColormapSVG(cm::AbstractVector{T}, w="96mm", h="4mm") where T<:Color
    io = IOBuffer()
    n = length(cm)
    write(io,
        """
        <svg xmlns="http://www.w3.org/2000/svg" version="1.1"
             width="$w" height="$h"
             viewBox="0 0 $n 1" preserveAspectRatio="none"
             shape-rendering="crispEdges" stroke="none">
        """)
    for i in 1:n
        c = hex(cm[i])
        write(io,
            """
            <rect width="$(n+1-i)" height="1" x="$(i-1)" y="0" fill="#$c" />
            """)
    end
    write(io, "</svg>")
    ColormapSVG(io)
end

end


function OffPolicyGAIL(;π,
                        S,
                        𝒟_demo,
                        𝒟_ndas::Array{ExperienceBuffer} = ExperienceBuffer[],
                        normalize_demo::Bool=true,
                        D::ContinuousNetwork,
                        solver=SAC,
                        d_opt::NamedTuple=(epochs=5,),
                        log::NamedTuple=(;),
                        kwargs...)

    # Define the training parameters for the desciminator
    d_opt = TrainingParams(;name="discriminator_", loss=()->nothing, d_opt...)

    # Setup NDA parameters
    N_nda = length(𝒟_ndas)
    λ_nda = Float32(-1 / N_nda)
    N_datasets = 2 + N_nda

    # Normalize and/or change device of expert and NDA data
    dev = device(π)
    A = action_space(π)
    normalize_demo && (𝒟_demo = normalize!(deepcopy(𝒟_demo), S, A))
    𝒟_demo = 𝒟_demo |> dev
    for i in 1:N_nda
        𝒟_ndas[i] = normalize_demo && normalize!(deepcopy(𝒟_ndas[i]), S, A)
        𝒟_ndas[i] = 𝒟_ndas[i] |> dev
    end

    # Build the solver
    𝒮 = solver(;π=π,
                S=S,
                log=(dir="log/offpolicygail", period=500, log...),
                kwargs...)

    # Setup the training of the discriminator
    B = d_opt.batch_size

    # These are minibatch buffers
    𝒟_batch = buffer_like(𝒮.buffer, capacity=B, device=dev)
    𝒟_demo_batch = deepcopy(𝒟_batch)
    𝒟_demo_π_batch = deepcopy(𝒟_batch)
    𝒟_ndas_batch = [deepcopy(𝒟_batch) for 𝒟_nda in 𝒟_ndas]
    𝒟_ndas_π_batch = [deepcopy(𝒟_batch) for 𝒟_nda in 𝒟_ndas]

    function GAIL_callback(𝒟; info=Dict(), kwargs...)
        for i=1:d_opt.epochs

            # Sample minibatchs
            rand!(𝒟_demo_batch, 𝒟_demo)
            # rand!(𝒟_demo_π_batch, 𝒟_demo)
            # 𝒟_demo_π_batch.data[:a] = action(π, 𝒟_demo_π_batch[:s])
            rand!(𝒟_batch, 𝒮.buffer)
            for i in 1:N_nda
                rand!(𝒟_ndas_batch[i], 𝒟_ndas[i])
                # rand!(𝒟_ndas_π_batch[i], 𝒟_ndas[i])
                # 𝒟_ndas_π_batch[i].data[:a] = action(π, 𝒟_ndas_π_batch[i][:s])
            end
            # Concatenate minibatches into on buffer
            𝒟_full = hcat(𝒟_demo_batch, 𝒟_batch, 𝒟_ndas_batch...)

            # concat inputs
            x = cat(flatten(𝒟_full[:s]), 𝒟_full[:a], dims=1)

            # Add some noise (Optional)
            # x .+= Float32.(rand(Normal(0, 0.2f0), size(x))) |> dev

            # Create labels
            y_demo = Flux.onehotbatch(ones(Int, B), 1:N_datasets)
            # y_demo_π = Flux.onehotbatch(2*ones(Int, B), 1:N_datasets)
            y_π = Flux.onehotbatch(2*ones(Int, B), 1:N_datasets)
            y_ndas = [Flux.onehotbatch((i+2)*ones(Int, B), 1:N_datasets) for i=1:N_nda]
            # y_ndas_π = [Flux.onehotbatch(2*ones(Int, B), 1:N_datasets) for i=1:N_nda]

            y = cat(y_demo, y_π, y_ndas..., dims=2) |> dev

            # println(gradient_penalty(D, x_demo, x_π))
            # + 10f0 * gradient_penalty(D, x_demo, x_π)
            train!(D, (;kwargs...) -> Flux.Losses.logitcrossentropy(D(x), y), d_opt, info=info)
        end

        # ## replace the bufffer
        # rand!(𝒟_demo_batch, 𝒟_demo)
        # # rand!(𝒟_demo_π_batch, 𝒟_demo)
        # # 𝒟_demo_π_batch.data[:a] = action(π, 𝒟_demo_π_batch[:s])
        # rand!(𝒟_batch, 𝒮.buffer)
        # for i in 1:N_nda
        #     rand!(𝒟_ndas_batch[i], 𝒟_ndas[i])
        #     # rand!(𝒟_ndas_π_batch[i], 𝒟_ndas[i])
        #     # 𝒟_ndas_π_batch[i].data[:a] = action(π, 𝒟_ndas_π_batch[i][:s])
        # end
        #
        # 𝒟_full = hcat(𝒟_demo_batch, 𝒟_batch, 𝒟_ndas_batch...)
        #
        # for k in keys(𝒟)
        #     𝒟.data[k] = 𝒟_full.data[k]
        # end
        # 𝒟.elements = 𝒟_full.elements
        # 𝒟.next_ind = 𝒟_full.next_ind

        ## Compute the rewards
        D_out = Flux.softmax(value(D, flatten(𝒟[:s]), 𝒟[:a]))
        w = [1f0, 0f0, λ_nda*ones(Float32, N_nda)...] |> dev

        𝒟[:r] .= sum((Base.log.(D_out .+ 1f-5) .- Base.log.(1f0 .- D_out .+ 1f-5)) .* w, dims=1)
    end

    𝒮.post_batch_callback = GAIL_callback
    𝒮
end


"""
    QuantumGlue

Interface between quantum science and technology packages.

Yao: quantum information/computing
QuantumOptics: ...

Spin up in `QuantumOptics.SpinBasis(1//2)` corresponds to `0` in the default basis of `Yao`.

`QuantumOptics` is aliased to `QO`.
"""
module QuantumGlue
import Yao
import QuantumOptics
const QO = QuantumOptics

export Ket, register

"""
    qubitbasis(n::Integer)

Create an `n`-qubit basis of type QO.CompositeBasis. The basis is the tensor
product of built-in `Spin(1//2)` bases.
"""
function qubitbasis(n::Integer)
    return QO.CompositeBasis((QO.SpinBasis(1//2) for i in 1:n)...)
end

"""
    Ket(reg::Yao.AbstractRegister)

Convert `reg` to a instance of type `QuantumOptics.Ket`. The input is not checked.
This method is an inverse of `register` acting on some objects of type `QuantumOptics.Ket`.
"""
function QuantumOptics.Ket(reg::Yao.AbstractRegister)
    return QuantumOptics.Ket(reg |> Yao.nqubits |> qubitbasis, reg |> Yao.state |> Yao.Intrinsics.matvec)
end

"""
    register(s::QO.Ket)

Convert `s` to an instance of `Yao.DefaultRegister`. The data is not copied, but
is shared between the input and output. The input is not checked. `s` must be a
`QuantumOptics.CompositeBasis` of `SpinBasis(1//2)`.  This method is an inverse
of `spins`.
"""
function Yao.register(s::QO.Ket)
    return Yao.register(s.data)
end

end # module QuantumGlue


immutable FESection{P <: AbstractFElement, T <: AbstractMaterial}
    elements::Vector{P}
    material::T
end

function FESection{P <: AbstractFElement, T <: AbstractMaterial,
                   Q <: AbstractMaterialStatus}(mat::T, ele_type::Type{P}, matstat_type::Type{Q})
    elements = ele_type{matstat_type}[]
    FESection(elements, mat)
end

immutable MaterialSection{P <: AbstractMaterial}
    material::P
    elements::Set{Int}
end
MaterialSection{P <: AbstractMaterial}(mat::P) = MaterialSection(mat, Set{Int}())

immutable ElementSection{P <: AbstractFElement}
    element_type::Type{P}
    elements::Set{Int}
end

ElementSection{P <: AbstractFElement}(ele_type::Type{P}) = ElementSection{P}(ele_type, Set{Int}())

# TODO: Clean this up
function push!(section::FESection, elem::AbstractFElement)
    push!(section.elements, elem)
end

function push!{P <: AbstractMaterial}(section::MaterialSection{P}, elemset::ElementSet)
    for i in elemset.element_ids
        push!(section.elements, i)
    end
end

function push!{P <: AbstractFElement}(section::ElementSection{P}, elemset::ElementSet)
    for i in elemset.element_ids
        push!(section.elements, i)
    end
end


### Implements a list of RLEVectors like Bioconductor's RLEList
### Not doing compressed versions where the latter concatenates the underlying
###   data into one vector for cheap vectorization

type RLEVectorList{T1,T2 <: Integer} <: AbstractArray{T1, 1}
    els::Vector{ RLEVector{T1,T2} }
    function RLEVectorList(rles...)
        rvl = ""
        return(rvl)
    end
end


InputFolder = './Images/VesselsSpots/';
OutputFolder = './Results/Images/VesselsSpots/Spots/';
Fill = -1;

@iA = '*.tif';

@fxm_mFilterObjSize [iA] > [M];
params.MinArea = 1;
params.MaxArea = 1;
/endf

/show iA > M;
/keep M > tif;


using ARDESPOT
using Test

using POMDPs
using POMDPModels
using POMDPSimulators
using Random
using POMDPModelTools
using ParticleFilters

include("memorizing_rng.jl")
include("independent_bounds.jl")

pomdp = BabyPOMDP()
pomdp.discount = 1.0

K = 10
rng = MersenneTwister(14)
rs = MemorizingSource(K, 50, rng)
Random.seed!(rs, 10)
b_0 = initialstate(pomdp)
scenarios = [i=>rand(rng, b_0) for i in 1:K]
o = false
b = ScenarioBelief(scenarios, rs, 0, o)
pol = FeedWhenCrying()
r1 = ARDESPOT.branching_sim(pomdp, pol, b, 10, (m,x)->0.0)
r2 = ARDESPOT.branching_sim(pomdp, pol, b, 10, (m,x)->0.0)
@test r1 == r2
tval = 7.0
r3 = ARDESPOT.branching_sim(pomdp, pol, b, 10, (m,x)->tval)
@test r3 == r2 + tval*length(b.scenarios)

scenarios = [1=>rand(rng, b_0)]
b = ScenarioBelief(scenarios, rs, 0, false)
pol = FeedWhenCrying()
r1 = ARDESPOT.rollout(pomdp, pol, b, 10, (m,x)->0.0)
r2 = ARDESPOT.rollout(pomdp, pol, b, 10, (m,x)->0.0)
@test r1 == r2
tval = 7.0
r3 = ARDESPOT.rollout(pomdp, pol, b, 10, (m,x)->tval)
@test r3 == r2 + tval

# AbstractParticleBelief interface
@test n_particles(b) == 1
s = particle(b,1)
@test rand(rng, b) == s
@test pdf(b, rand(rng, b_0)) == 1
sup = support(b)
@test length(sup) == 1
@test first(sup) == s
@test mode(b) == s
@test mean(b) == s
@test first(particles(b)) == s
@test first(weights(b)) == 1.0
@test first(weighted_particles(b)) == (s => 1.0)
@test weight_sum(b) == 1.0
@test weight(b, 1) == 1.0
@test currentobs(b) == o
@test_deprecated previous_obs(b)
@test history(b)[end].o == o

pomdp = BabyPOMDP()

# constant bounds
bds = (reward(pomdp, true, false)/(1-discount(pomdp)), 0.0)
solver = DESPOTSolver(bounds=bds)
planner = solve(solver, pomdp)
hr = HistoryRecorder(max_steps=2)
@time hist = simulate(hr, pomdp, planner)

# policy lower bound
bds = IndependentBounds(DefaultPolicyLB(FeedWhenCrying()), 0.0)
solver = DESPOTSolver(bounds=bds)
planner = solve(solver, pomdp)
hr = HistoryRecorder(max_steps=2)
@time hist = simulate(hr, pomdp, planner)

# policy lower bound with final value
fv(m::BabyPOMDP, x) = reward(m, true, false)/(1-discount(m))
bds = IndependentBounds(DefaultPolicyLB(FeedWhenCrying(), final_value=fv), 0.0)
solver = DESPOTSolver(bounds=bds)
planner = solve(solver, pomdp)
hr = HistoryRecorder(max_steps=2)
@time hist = simulate(hr, pomdp, planner)

# Type stability
pomdp = BabyPOMDP()
bds = IndependentBounds(reward(pomdp, true, false)/(1-discount(pomdp)), 0.0)
solver = DESPOTSolver(epsilon_0=0.1,
                      bounds=bds,
                      rng=MersenneTwister(4)
                     )
p = solve(solver, pomdp)

b0 = initialstate(pomdp)
D = @inferred ARDESPOT.build_despot(p, b0)
@inferred ARDESPOT.explore!(D, 1, p)
@inferred ARDESPOT.expand!(D, length(D.children), p)
@inferred ARDESPOT.prune!(D, 1, p)
@inferred ARDESPOT.find_blocker(D, length(D.children), p)
@inferred ARDESPOT.make_default!(D, length(D.children))
@inferred ARDESPOT.backup!(D, 1, p)
@inferred ARDESPOT.next_best(D, 1, p)
@inferred ARDESPOT.excess_uncertainty(D, 1, p)
@inferred action(p, b0)

bds = IndependentBounds(reward(pomdp, true, false)/(1-discount(pomdp)), 0.0)
rng = MersenneTwister(4)
solver = DESPOTSolver(epsilon_0=0.1,
                      bounds=bds,
                      rng=rng,
                      random_source=MemorizingSource(500, 90, rng),
                      tree_in_info=true
                     )
p = solve(solver, pomdp)
a = action(p, initialstate(pomdp))

include("random_2.jl")

# visualization
show(stdout, MIME("text/plain"), D)
a, info = action_info(p, initialstate(pomdp))
show(stdout, MIME("text/plain"), info[:tree])

# from README:
using POMDPs, POMDPModels, POMDPSimulators, ARDESPOT

pomdp = TigerPOMDP()

solver = DESPOTSolver(bounds=(-20.0, 0.0))
planner = solve(solver, pomdp)

for (s, a, o) in stepthrough(pomdp, planner, "s,a,o", max_steps=10)
    println("State was $s,")
    println("action $a was taken,")
    println("and observation $o was received.\n")
end


# This file contains a function to generate time series

"""
    @array(array::AbstractString)

Wrap an array for zero-based indexing.

# Usage

```jldoctest; output = false
using OffsetArrays
using Scats.Internal.Input: InputStruct
using Scats.Internal.Gen: @array

file, _ = mktemp()
Input = InputStruct()

Input.example(file)
Input(file)

N₋₁ = Input.N - 1
@array("Input.t") === OffsetArray(Input.t, 0:N₋₁)

# output

true
```
"""
macro array(array::AbstractString)
    return esc(Meta.parse("OffsetArray($array, 0:N₋₁)"))
end

"""
    gen!(Gen::GenStruct, Input::InputStruct)

Generate time series for an instance of [`InputStruct`](@ref) using generator
parameters from an instance of [`GenStruct`](@ref).

# Usage

```jldoctest; output = false
using Scats
s = Scats.API()
file, _ = mktemp()
s.Gen.example(file)
s.gen!()

# output

```
"""
function gen!(Gen::GenStruct, Input::InputStruct)

    @unpack "Gen" "N" "Δt" "q" "α" "β" "r" "A" "ν" "ϕ" "γ"

    # Save some values in input
    @pack "Input" "N" "Δt" "q"

    # Calculate standard deviation
    σ = √(sum(A .* A) / (2 * γ))

    # Calculate auxiliary variable
    N₋₁ = N - 1

    # Check sizes of arrays
    if size(Input.t, 1) != N
        Input.t = Vector{RT}(undef, N)
    end
    if size(Input.x, 1) != N
        Input.x = Vector{RT}(undef, N)
    end

    # Wrap arrays from input
    ta = @array "Input.t"
    xa = @array "Input.x"

    # Generate random numbers array
    rng = MersenneTwister()
    rand = @array "randn(rng, RT, N)"

    # Generate time series
    for k in 0:N₋₁

        # Calculate current time moment
        t = Δt * k

        # Write it in the array
        ta[k] = t

        # Calculate the linear trend part
        c = α + β * t

        # Add the harmonics part
        for l in 1:r
            c = c + A[l] * cos(2 * π * ν[l] * t - ϕ[l])
        end

        # Add the noise part
        c = c + rand[k]

        # Write a result in the array
        xa[k] = c

    end

end


### Multi-Lane intersection with fixed obstacles

"""
Container type for the environment parameters
"""
@with_kw mutable struct TInterParams
    # geometry
    x_min::Float64 = -20.0
    x_max::Float64 = 20.0
    y_min::Float64 = -20.0
    inter_x::Float64 = 0.
    inter_y::Float64 = 0.
    inter_width::Float64 = 6.0
    inter_height::Float64 = 1.0
    lane_width::Float64 = 3.0
    nlanes_main::Int64 = 2
    nlanes::Int64 = 1

    stop_line::Float64 = 10. # in m along the ego car lane
    # cars
    speed_limit::Float64 = 10.0 # in m/s
    car_rate::Float64 = 0.3 # probability of a car appearing every time steps
end

"""
Generate a T intersection given a geometry specify by TInterParams
"""
function gen_T_roadway(params::TInterParams)
    #=
    ________________________________________________________
        <---- seg 2     main street <----- seg 1
        -----> seg 3                 -----> seg 4
    _________________            ___________________________
                    | seg . seg  |
                    |  5  .  6   |
                    |     .      |
                    |     .      |
                    |     .      |

    =#
    roadway = Roadway()
    # aliases
    x_min = params.x_min
    x_max = params.x_max
    y_min = params.y_min
    inter_x = params.inter_x
    inter_y = params.inter_y
    inter_width = params.inter_width
    inter_height = params.inter_height
    lane_width = params.lane_width
    nlanes_main = params.nlanes_main
    nlanes = params.nlanes

    seg1 = VecSE2(x_max, inter_y+nlanes_main*lane_width - lane_width/2, pi)
    length1 = x_max - inter_x - inter_width
    add_line!(seg1, nlanes_main, length1, roadway)

    seg2 = VecSE2(inter_x - inter_width, inter_y+nlanes_main*lane_width - lane_width/2, pi)
    length2 =  inter_x - inter_width - x_min
    add_line!(seg2, nlanes_main, length2, roadway)

    seg3 = VecSE2(x_min, inter_y-nlanes_main*lane_width + lane_width/2, 0.0)
    length3 = inter_x - x_min - inter_width
    add_line!(seg3, nlanes_main, length3, roadway)

    seg4 = VecSE2(inter_x + inter_width, inter_y-nlanes_main*lane_width + lane_width/2, 0.0)
    length4 = x_max - inter_x - inter_width
    add_line!(seg4, nlanes_main, length4, roadway)

    seg5 = VecSE2(inter_x-nlanes*lane_width + lane_width/2, inter_y - nlanes_main*lane_width - inter_height, -pi/2)
    length5 = inter_y - nlanes_main*lane_width - y_min - inter_height
    add_line!(seg5, nlanes, length5, roadway)

    seg6 = VecSE2(inter_x+nlanes*lane_width - lane_width/2, y_min, pi/2)
    length6 = inter_y - nlanes_main*lane_width - y_min - inter_height
    add_line!(seg6, nlanes, length6, roadway)

    connections = [Connection(1,2), Connection(3,4)]
    push!(connections, Connection(1, 5, 0, [(nlanes_main, 1)]))
    push!(connections, Connection(3, 5, 0, [(1, 1)]))
    # what follows is the ego car left turn
    push!(connections, Connection(6, 2, 0, [(1, nlanes_main)]))
    push!(connections, Connection(6, 4, 0, [(1, 1)]))

    add_junction!(Junction(connections), roadway)
    return roadway
end

function build_obstacles(params::TInterParams)#TODO
    # HEIGHT = 6
    # WIDTH = 6
    #
    # #Left Building
    # left_building = ConvexPolygon(4)
    # top_right = VecSE2(3*params.lane_width/4, -params.lane_width/2 - 1.5, 0.)
    # left_building.pts = [top_right, top_right + VecSE2(-WIDTH, 0., 0.),
    #                 top_right + VecSE2(-WIDTH, -HEIGHT, 0.), top_right + VecSE2(0., -HEIGHT, 0.)]
    # left_building.npts = 4; # a rectangle
    #
    # #Right Building
    # right_building = ConvexPolygon(4)
    # top_right = top_right + VecSE2(top_right.x + 2*params.lane_width + WIDTH, 0., 0.)
    # right_building.pts = [top_right, top_right + VecSE2(-WIDTH, 0., 0.),
    #                 top_right + VecSE2(-WIDTH, -HEIGHT, 0.), top_right + VecSE2(0., -HEIGHT, 0.)]
    # right_building.npts = 4; # a rectangle
    # return [left_building, right_building]
    return Vector{ConvexPolygon}[]
end

"""
T intersection environment
"""
mutable struct IntersectionEnv <: OccludedEnv
    roadway::Roadway
    obstacles::Vector{ConvexPolygon}
    params::TInterParams
    lane_map::Dict{String, Lane}
    priorities::Dict{Tuple{LaneTag, LaneTag}, Bool}
    directions::Dict{Symbol, Tuple{LaneTag, LaneTag}}
end

function IntersectionEnv(params::TInterParams = TInterParams())
    roadway = gen_T_roadway(params)
    lane_map = Dict{String, Lane}()
    obstacles = build_obstacles(params)
    priorities, directions = priority_map(roadway)
    return IntersectionEnv(roadway, obstacles, params, lane_map, priorities, directions)
end

# return a mapping between a label and a route for a 2 lanes T intersection
function priority_map(roadway::Roadway)
    priorities = Dict{Tuple{LaneTag, LaneTag}, Bool}()
    directions = Dict{Symbol, Tuple{LaneTag, LaneTag}}()

    # straight from right
    route = (LaneTag(1,1),LaneTag(2,1))
    priorities[route] = true
    directions[:straight_from_right] = route
    route = (LaneTag(1,2),LaneTag(2,2))
    priorities[route] = true
    directions[:straight_from_right2] = route

    # straight from left
    route = (LaneTag(3,1),LaneTag(4,1))
    priorities[route] = true
    directions[:straight_from_left] = (LaneTag(1,1),LaneTag(2,1))
    route = (LaneTag(3,2),LaneTag(4,2))
    priorities[route] = true
    directions[:straight_from_left2] = (LaneTag(1,2),LaneTag(2,2))

    # right turns
    route = (LaneTag(3,1), LaneTag(5,1))
    priorities[route] = true
    directions[:right_from_left] = route
    # route = (LaneTag(3,2), LaneTag(5,1))
    # priorities[route] = true
    # directions[:right_from_left2] = route
    route = (LaneTag(6,1), LaneTag(4,1))
    priorities[route] = false
    directions[:right_from_bot] = route
    # route = (LaneTag(6,1), LaneTag(4,2))
    # priorities[route] = false
    # directions[:right_from_bot2] = route

    # left turns
    route = (LaneTag(6,1), LaneTag(2,2))
    priorities[route] = false
    route = (LaneTag(6,1), LaneTag(2,1))
    priorities[route] = false
    route = (LaneTag(1,1), LaneTag(5,1))
    priorities[route] = false
    route = (LaneTag(1,2), LaneTag(5,1))
    priorities[route] = false
    return priorities, directions
end