using Oceananigans.Advection: boundary_buffer
using Oceananigans.Operators: ℑxᶠᵃᵃ, ℑxᶜᵃᵃ, ℑyᵃᶠᵃ, ℑyᵃᶜᵃ, ℑzᵃᵃᶠ, ℑzᵃᵃᶜ
using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, AbstractTimeDiscretization

const ATC = AbstractTurbulenceClosure
const ATD = AbstractTimeDiscretization

struct RasterDepthMask end

struct GridFittedBoundary{M, S} <: AbstractImmersedBoundary
    mask :: S
    mask_type :: M
end

GridFittedBoundary(mask; mask_type=nothing) = GridFittedBoundary(mask, mask_type)

@inline is_immersed(i, j, k, underlying_grid, ib::GridFittedBoundary) = ib.mask(node(c, c, c, i, j, k, underlying_grid)...)
@inline is_immersed(i, j, k, underlying_grid, ib::GridFittedBoundary{<:RasterDepthMask}) = ib.mask(i, j)

const IBG = ImmersedBoundaryGrid

const c = Center()
const f = Face()

#####
##### ImmersedBoundaryGrid
#####

@inline solid_cell(i, j, k, ibg) = is_immersed(i, j, k, ibg.grid, ibg.immersed_boundary)

@inline solid_node(LX, LY, LZ, i, j, k, ibg) = solid_cell(i, j, k, ibg) # fallback (for Center or Nothing LX, LY, LZ)

@inline solid_node(::Face, LX, LY, i, j, k, ibg) = solid_cell(i, j, k, ibg) || solid_cell(i-1, j, k, ibg)
@inline solid_node(LX, ::Face, LZ, i, j, k, ibg) = solid_cell(i, j, k, ibg) || solid_cell(i, j-1, k, ibg)
@inline solid_node(LX, LY, ::Face, i, j, k, ibg) = solid_cell(i, j, k, ibg) || solid_cell(i, j, k-1, ibg)

@inline solid_node(::Face, ::Face, LZ, i, j, k, ibg) = solid_node(c, f, c, i, j, k, ibg) || solid_node(c, f, c, i-1, j, k, ibg)
@inline solid_node(::Face, LY, ::Face, i, j, k, ibg) = solid_node(c, c, f, i, j, k, ibg) || solid_node(c, c, f, i-1, j, k, ibg)
@inline solid_node(LX, ::Face, ::Face, i, j, k, ibg) = solid_node(c, f, c, i, j, k, ibg) || solid_node(c, f, c, i, j, k-1, ibg)

@inline solid_node(::Face, ::Face, ::Face, i, j, k, ibg) = solid_node(c, f, f, i, j, k, ibg) || solid_node(c, f, f, i-1, j, k, ibg)

@inline conditional_flux_ccc(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(c, c, c, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))
@inline conditional_flux_ffc(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(f, f, c, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))
@inline conditional_flux_fcf(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(f, c, f, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))
@inline conditional_flux_cff(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(c, f, f, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))

@inline conditional_flux_fcc(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(f, c, c, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))
@inline conditional_flux_cfc(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(c, f, c, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))
@inline conditional_flux_ccf(i, j, k, ibg::IBG{FT}, grid, flux, args...) where FT = ifelse(solid_node(c, c, f, i, j, k, ibg), zero(FT), flux(i, j, k, grid, args...))

#####
##### Advective fluxes
#####

# ccc, ffc, fcf
 @inline viscous_flux_ux(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_ccc(i, j, k, ibg, ibg.grid, viscous_flux_ux, disc, clo, args...)
 @inline viscous_flux_uy(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_ffc(i, j, k, ibg, ibg.grid, viscous_flux_uy, disc, clo, args...)
 @inline viscous_flux_uz(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_fcf(i, j, k, ibg, ibg.grid, viscous_flux_uz, disc, clo, args...)

 # ffc, ccc, cff
 @inline viscous_flux_vx(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_ffc(i, j, k, ibg, ibg.grid, viscous_flux_vx, disc, clo, args...)
 @inline viscous_flux_vy(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_ccc(i, j, k, ibg, ibg.grid, viscous_flux_vy, disc, clo, args...)
 @inline viscous_flux_vz(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_cff(i, j, k, ibg, ibg.grid, viscous_flux_vz, disc, clo, args...)

 # fcf, cff, ccc
 @inline viscous_flux_wx(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_fcf(i, j, k, ibg, ibg.grid, viscous_flux_wx, disc, clo, args...)
 @inline viscous_flux_wy(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_cff(i, j, k, ibg, ibg.grid, viscous_flux_wy, disc, clo, args...)
 @inline viscous_flux_wz(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_ccc(i, j, k, ibg, ibg.grid, viscous_flux_wz, disc, clo, args...)

# fcc, cfc, ccf
@inline diffusive_flux_x(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_fcc(i, j, k, ibg, ibg.grid, diffusive_flux_x, disc, clo, args...)
@inline diffusive_flux_y(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_cfc(i, j, k, ibg, ibg.grid, diffusive_flux_y, disc, clo, args...)
@inline diffusive_flux_z(i, j, k, ibg::IBG, disc::ATD, clo::ATC, args...) = conditional_flux_ccf(i, j, k, ibg, ibg.grid, diffusive_flux_z, disc, clo, args...)

#####
##### Advective fluxes
#####

@inline _advective_momentum_flux_Uu(i, j, k, ibg::IBG, args...) = conditional_flux_ccc(i, j, k, ibg, ibg, advective_momentum_flux_Uu, args...)
@inline _advective_momentum_flux_Vu(i, j, k, ibg::IBG, args...) = conditional_flux_ffc(i, j, k, ibg, ibg, advective_momentum_flux_Vu, args...)
@inline _advective_momentum_flux_Wu(i, j, k, ibg::IBG, args...) = conditional_flux_fcf(i, j, k, ibg, ibg, advective_momentum_flux_Wu, args...)

@inline _advective_momentum_flux_Uv(i, j, k, ibg::IBG, args...) = conditional_flux_ffc(i, j, k, ibg, ibg, advective_momentum_flux_Uv, args...)
@inline _advective_momentum_flux_Vv(i, j, k, ibg::IBG, args...) = conditional_flux_ccc(i, j, k, ibg, ibg, advective_momentum_flux_Vv, args...)
@inline _advective_momentum_flux_Wv(i, j, k, ibg::IBG, args...) = conditional_flux_cff(i, j, k, ibg, ibg, advective_momentum_flux_Wv, args...)

@inline _advective_momentum_flux_Uw(i, j, k, ibg::IBG, args...) = conditional_flux_fcf(i, j, k, ibg, ibg, advective_momentum_flux_Uw, args...)
@inline _advective_momentum_flux_Vw(i, j, k, ibg::IBG, args...) = conditional_flux_cff(i, j, k, ibg, ibg, advective_momentum_flux_Vw, args...)
@inline _advective_momentum_flux_Ww(i, j, k, ibg::IBG, args...) = conditional_flux_ccc(i, j, k, ibg, ibg, advective_momentum_flux_Ww, args...)

@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, args...) = conditional_flux_fcc(i, j, k, ibg, ibg, advective_tracer_flux_x, args...)
@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, args...) = conditional_flux_cfc(i, j, k, ibg, ibg, advective_tracer_flux_y, args...)
@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, args...) = conditional_flux_ccf(i, j, k, ibg, ibg, advective_tracer_flux_z, args...)

#####
##### "Boundary-aware" interpolation
#####
##### Don't interpolate dead cells.
#####

@inline near_x_boundary(i, j, k, ibg, buffer) = any(ntuple(δ -> solid_cell(i - buffer - 1 + δ, j, k, ibg), Val(2buffer + 1)))
@inline near_y_boundary(i, j, k, ibg, buffer) = any(ntuple(δ -> solid_cell(i, j - buffer - 1 + δ, k, ibg), Val(2buffer + 1)))
@inline near_z_boundary(i, j, k, ibg, buffer) = any(ntuple(δ -> solid_cell(i, j, k - buffer - 1 + δ, ibg), Val(2buffer + 1)))

for bias in (:symmetric, :left_biased, :right_biased)

    for (d, ξ) in enumerate((:x, :y, :z))

        code = [:ᵃ, :ᵃ, :ᵃ]

        for loc in (:ᶜ, :ᶠ)
            code[d] = loc
            second_order_interp = Symbol(:ℑ, ξ, code...)
            interp = Symbol(bias, :_interpolate_, ξ, code...)
            alt_interp = Symbol(:_, interp)

            near_boundary = Symbol(:near_, ξ, :_boundary)

            # Conditional high-order interpolation in Bounded directions
            @eval begin
                import Oceananigans.Advection: $alt_interp
                using Oceananigans.Advection: $interp

                @inline $alt_interp(i, j, k, ibg::ImmersedBoundaryGrid, scheme, ψ) =
                    ifelse($near_boundary(i, j, k, ibg, boundary_buffer(scheme)),
                           $second_order_interp(i, j, k, ibg.grid, ψ),
                           $interp(i, j, k, ibg.grid, scheme, ψ))
            end
        end
    end
end


# xvecs_read and xvecs_write
@testset "xvecs" begin
  fn = tempname()
  d, n, m, h = 32, Int(1e3), 4, 256
  X, C, B = generate_random_dataset(Float32, Int16, d, n, m, h)

  # fvecs
  fvecs_write(X, fn)
  X2 = fvecs_read(n, fn)
  @test all(X .== X2)
  rm(fn)

  Xint = convert(Matrix{Int32}, floor.(X.-0.5f0)*1000)
  ivecs_write(Xint, fn)
  Xint2 = ivecs_read(n, fn)
  @test all(Xint .== Xint2)
  rm(fn)
end


abstract type AbstractPrimalHeuristic end
abstract type RoundingHeuristic <: AbstractPrimalHeuristic end
abstract type AllBlockHeuristic <: AbstractPrimalHeuristic end

function add!(htype::Type{T}, LD::AbstractLagrangeDual) where T <: AbstractPrimalHeuristic
    push!(LD.heuristics, htype)
end

"""
    run!

This run a rounding heuristic.
"""
function run!(::Type{RoundingHeuristic}, LD::AbstractLagrangeDual, val, ub, lb)
    #set weights of each model if none exists
    num_all_blocks = parallel.sum(num_blocks(LD))
    all_block_ids = Set([ks[1] for (ks, i) in LD.var_to_index])
    if num_all_blocks != LD.block_model.combined_weights.count
        LD.block_model.combined_weights = Dict(block_id => 1/num_all_blocks for block_id in all_block_ids)
    end

    # get the mean value of the coupling variables at the root process and broadcast to other processes
    new_primal_solution = Dict() #maps coupling_id to value
    if parallel.is_root()
        for (ks, i) in LD.var_to_index
            coupling_id = ks[2]
            block_id = ks[1]
            new_primal_solution[coupling_id] = haskey(new_primal_solution, coupling_id) ? new_primal_solution[coupling_id] + LD.block_model.combined_weights[block_id] * val[block_id][i] : LD.block_model.combined_weights[block_id] * val[block_id][i]
        end
        parallel.bcast(new_primal_solution)
    else
        new_primal_solution = parallel.bcast(nothing)
    end

    #fix variables with new primal solution and enforce integrality through rounding
    for variables in LD.block_model.coupling_variables
        if JuMP.is_integer(variables.ref) || JuMP.is_binary(variables.ref)
            JuMP.fix(variables.ref, round(new_primal_solution[variables.key.coupling_id]), force=true)
            new_primal_solution[variables.key.coupling_id] = round(new_primal_solution[variables.key.coupling_id])
        else
            JuMP.fix(variables.ref, new_primal_solution[variables.key.coupling_id], force=true)
            new_primal_solution[variables.key.coupling_id] = new_primal_solution[variables.key.coupling_id]
        end
    end

    #obtain primal bound by solving the subproblems in parallel
    cur_primal_bound = 0.0
    for (id,m) in block_model(LD)
        JuMP.optimize!(m)
        if ! (JuMP.termination_status(m)  in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED])
            cur_primal_bound = + Inf
            break
        else
            cur_primal_bound += JuMP.objective_value(m)
        end
    end
    cur_primal_bound_sum = parallel.sum(cur_primal_bound)

    #update primal bound and solution
    if cur_primal_bound_sum < LD.block_model.primal_bound
        LD.block_model.primal_bound = cur_primal_bound_sum
        LD.block_model.primal_solution = copy(new_primal_solution)
        if parallel.is_root()
            println("  Found new best primal bound: $(LD.block_model.primal_bound)")
        end
    end

    #unfix variables and recover their original bounds
    for var in coupling_variables(LD)
        JuMP.unfix(var.ref)
        JuMP.set_lower_bound(var.ref, lb[var.key.block_id][index_of_λ(LD, var)])
        JuMP.set_upper_bound(var.ref, ub[var.key.block_id][index_of_λ(LD, var)])
    end
end

"""
    run!

This run a heuristic that simply fixes the coupling variable values.
"""
function run!(::Type{AllBlockHeuristic}, LD::AbstractLagrangeDual, val, ub, lb)
    all_block_ids = Set([ks[1] for (ks, i) in LD.var_to_index])
    all_coupling_ids = Set([ks[2] for (ks, i) in LD.var_to_index])

    # broadcast the val
    if parallel.is_root()
        parallel.bcast(val)
    else
        val = parallel.bcast(nothing)
    end

    #iterate over all blocks. Fix the coupling variables to the optimal solution of each block.
    for block_id in all_block_ids
        new_primal_solution = Dict() #maps coupling_id to value
        for coupling_id in all_coupling_ids
            new_primal_solution[coupling_id] = val[block_id][LD.var_to_index[block_id, coupling_id]]
        end

        #fix variables with new primal solution and enforce integrality
        for variables in LD.block_model.coupling_variables
            JuMP.fix(variables.ref, new_primal_solution[variables.key.coupling_id], force=true)
        end

        #obtain primal bound by solving the subproblems in parallel
        cur_primal_bound = 0.0
        for (id,m) in block_model(LD)
            JuMP.optimize!(m)
            if ! (JuMP.termination_status(m)  in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED])
                cur_primal_bound = + Inf
                break
            else
                cur_primal_bound += JuMP.objective_value(m)
            end
        end

        #update bounds and solution
        cur_primal_bound_sum = parallel.sum(cur_primal_bound)
        if cur_primal_bound_sum < LD.block_model.primal_bound
            LD.block_model.primal_bound = cur_primal_bound_sum
            LD.block_model.primal_solution = copy(new_primal_solution)

            if parallel.is_root()
                println("  Found new best primal bound: $(LD.block_model.primal_bound)")
            end
        end
    end

    #unfix variables
    for var in coupling_variables(LD)
        JuMP.unfix(var.ref)
        JuMP.set_lower_bound(var.ref, lb[var.key.block_id][index_of_λ(LD, var)])
        JuMP.set_upper_bound(var.ref, ub[var.key.block_id][index_of_λ(LD, var)])
    end
end


@testset "reshape" begin
    A = rand(4, 5)
    B, pullback = rrule(reshape, A, (5, 4))
    @test B == reshape(A, (5, 4))
    Ȳ = randn(4, 5)

    (s̄, Ā, d̄) = pullback(Ȳ)
    @test s̄ == NO_FIELDS
    @test d̄ isa DoesNotExist
    @test extern(Ā) == reshape(Ȳ, (5, 4))

    B, pullback = rrule(reshape, A, 5, 4)
    @test B == reshape(A, 5, 4)

    Ȳ = randn(4, 5)
    (s̄, Ā, d̄1, d̄2) = pullback(Ȳ)
    @test s̄ == NO_FIELDS
    @test d̄1 isa DoesNotExist
    @test d̄2 isa DoesNotExist
    @test extern(Ā) == reshape(Ȳ, 5, 4)
end

@testset "hcat" begin
    A = randn(3, 2)
    B = randn(3)
    C = randn(3, 3)
    H, pullback = rrule(hcat, A, B, C)
    @test H == hcat(A, B, C)
    H̄ = randn(3, 6)
    (ds, dA, dB, dC) = pullback(H̄)
    @test ds == NO_FIELDS
    @test dA ≈ view(H̄, :, 1:2)
    @test dB ≈ view(H̄, :, 3)
    @test dC ≈ view(H̄, :, 4:6)
end

@testset "reduce hcat" begin
    A = randn(3, 2)
    B = randn(3, 1)
    C = randn(3, 3)
    x = [A, B, C]
    H, pullback = rrule(reduce, hcat, x)
    @test H == reduce(hcat, x)
    H̄ = randn(3, 6)
    x̄ = randn.(size.(x))
    rrule_test(reduce, H̄, (hcat, nothing), (x, x̄))
end

@testset "vcat" begin
    A = randn(2, 4)
    B = randn(1, 4)
    C = randn(3, 4)
    V, pullback = rrule(vcat, A, B, C)
    @test V == vcat(A, B, C)
    V̄ = randn(6, 4)
    (ds, dA, dB, dC) = pullback(V̄)
    @test ds == NO_FIELDS
    @test dA ≈ view(V̄, 1:2, :)
    @test dB ≈ view(V̄, 3:3, :)
    @test dC ≈ view(V̄, 4:6, :)
end

@testset "reduce vcat" begin
    A = randn(2, 4)
    B = randn(1, 4)
    C = randn(3, 4)
    x = [A, B, C]
    V, pullback = rrule(reduce, vcat, x)
    @test V == reduce(vcat, x)
    V̄ = randn(6, 4)
    x̄ = randn.(size.(x))
    rrule_test(reduce, V̄, (vcat, nothing), (x, x̄))
end

@testset "fill" begin
    y, pullback = rrule(fill, 44, 4)
    @test y == [44, 44, 44, 44]
    (ds, dv, dd) = pullback(ones(4))
    @test ds === NO_FIELDS
    @test dd isa DoesNotExist
    @test extern(dv) == 4

    y, pullback = rrule(fill, 2.0, (3, 3, 3))
    @test y == fill(2.0, (3, 3, 3))
    (ds, dv, dd) = pullback(ones(3, 3, 3))
    @test ds === NO_FIELDS
    @test dd isa DoesNotExist
    @test dv ≈ 27.0
end


module ADFPEPS

export HamiltonianModel,hamiltonian
export Hubbard,hop_pair,Hubbard_hand,hop_pair_hand,THubbard
export observable

include("hamiltonianmodels.jl")
include("fermion.jl")
include("autodiff.jl")
include("variationalipeps.jl")
include("observable.jl")

end


using DataFrames

"""
    padItem(item, nchar)

Pads an item with an appropriate amount of "." characters so that n-grams can be
created. Provides two methods:

    padItem(item::Array{String}, nchar::Integer)
    padItem(item::String, nchar::Integer)

# Arguments
* **item** The item that needs to be padded; may be an `Array` of `Strings` or
    simply a `String`
* **nchar** The number of characters contained in the n-gram that is being
    examined (e.g., 2 for diphones)

# Returns
The item with appropriate padding
"""
function padItem(item::Array{String}, nchar::Integer)
    padding = ["." for x in 1:(nchar-1)]
    item = vcat(padding, item, padding)
    return item
end

function padItem(item::String, nchar::Integer)
    padding = repeat(".", nchar-1)
    item = padding * item * padding
    return item
end

"""
    createGrams(item, nchar::Integer)

Creates the n-grams that will be analyzed. If looking at diphones, will produce
diphones from the input, for example. Here, n-grams is used in a general sense
as an extension of the n-gram in natural language processing.

# Arguments
* **item** The item for which to create the list of n-grams
* **nchar** The number of characters for each gram (e.g., 2 for diphones)

# Returns
A `Vector` of the created n-grams
"""
function createGrams(item, nchar::Integer)
    grams = Vector()
    for i=1:(length(item)-(nchar-1))
        push!(grams, item[i:i+(nchar-1)])
    end
    return grams
end

"""
    countSymbols(corpus::Array, frequencies::Array, positional::Bool,
        nchar::Integer, pad::Bool)

Counts the number of occurrences of each symbol in the passed in corpus;
produces, for example, the number of times "K" appears in the corpus.

# Arguments
* **corpus** The corpus from which to derive the counts
* **frequencies** The number of times each item in the corpus has been observed;
    this variable must correspond with `corpus` such that the n-th element of
    `corpus` is related to the n-th element of `frequencies`.
* **positional** Whether to consider where in a word a given phone appears
(e.g., should "K" as the first sound be considered a different category than "K"
    as the second sound?)
* **nchar** The number of characters for each n-gram that will be examined
    (e.g., 2 for diphones)
* **pad** Whether to add padding to each query or not

# Returns
A `Dict` where each key is a differnt observed "gram," and each paired value is
the number of times that n-gram was observed
"""
function countSymbols(corpus::Array, frequencies::Array, positional::Bool,
    nchar::Integer, pad::Bool)

    counts = Dict()
    total = 0

    for item in corpus

        if pad
            item = padItem(item, nchar)
        end
        item = createGrams(item, nchar)

        for (i, ch) in enumerate(item)

            key = positional ? (i, ch) : ch

            if haskey(counts, key)
                counts[key] += frequencies[i]
            else
                counts[key] = frequencies[i]
            end

            total += 1
        end
    end

    return counts
end

"""
    calcProb(counts, query, positional, total)

Calculates the phonotactic probability for a query by determining the count of
each n-gram contained in the query, dividing each count by the total number of
n-grams observed in the corpus, and taking the product of those proportions

# Arguments
* **counts** A `Dict` that can be indexed with n-grams and produce the number
    of occurrences for that item
* **query** The item for which the phonotactic probability should be calculated
* **positional** Whether to consider where in the query a given phone appears
    (e.g., should "K" as the first sound be considered a different category
    than "K" as the second sound?)

# Returns
The calculated phonotactic probability for the given query item
"""
function calcProb(counts::Dict, query, positional::Bool, total)
    unitCounts = Int64[]

    for (i, ch) in enumerate(query)

        key = positional ? (i, ch) : ch

        if haskey(counts, key)
            push!(unitCounts, counts[key])
        else
            push!(unitCounts, 0)
        end
    end

    return prod(unitCounts / total)
end

"""
    phnprb(corpus::Array, frequencies::Array, queries::Array; positional=false,
        nchar=1, pad=true)

Calculates the phonotactic probability for each item in a list of queries based
on a corpus

# Arguments
* **corpus** The corpus on which to base the probability calculations
* **frequencies** The frequencies associated with each element in `corpus`
* **queries** The items for which the probability should be calculated

## Keyword arguments
* **positional**  Whether to consider where in the query a given phone appears
(e.g., should "K" as the first sound be considered a different category than "K"
    as the second sound?)
* **nchar** The number of characters for each n-gram that will be examined
    (e.g., 2 for diphones)
* **pad** Whether to add padding to each query or not

# Returns
A `DataFrame` with the queries in the first column and the probability values
    in the second
"""
function phnprb(corpus::Array, frequencies::Array{Int64}, queries::Array;
    positional=false, nchar=1, pad=true)

    counts = countSymbols(corpus, frequencies, positional, nchar, pad)
	total = sum(values(counts))

    results = Vector()

    for query in queries
        local paddedQuery
        if pad
            paddedQuery = padItem(query, nchar)
        else
            paddedQuery = query
        end
        grams = createGrams(paddedQuery, nchar)
        prob = calcProb(counts, grams, positional, total)
        push!(results, prob)
    end

    return DataFrame(Query=queries, Probability=results)
end


"""
    Cell{T}

preserve 'fmtId' for Cell
currently there is no functionality for fmtId after DataType of cell is determined
"""
struct Cell{T}
    numFmtId::UInt16
    v::T
end
Cell(id, v::T) where T = Cell{T}(id, v)

function Cell(wb::WorkBook, x::EzXML.Node)
    # @show "cell data / " x["r"]
    fmt_id, T = find_cell_datatype(wb, x)

    cell_values = filter(iselement, nodes(x))
    # sometimes there can be no value, even through it has cell address
    value = if isempty(cell_values)
                missing
            else # Excel always store value `v` node at the end
                v = nodecontent.(cell_values)[end]
                v = get_cellvalue(T, wb, v)
            end
    Cell(fmt_id, value)
end

# fallback methods
Base.Symbol(x::Cell) = convert(Symbol, x)
Base.convert(::Type{Symbol}, x::Cell) = Symbol(x.v)

function find_cell_datatype(wb::WorkBook, x)
    fmt_id = haskey(x, "s") ? wb.xf_index_to_number_format[parse(Int, x["s"])+1] : 0
    T = if haskey(x, "t")
            ECMA_CELLTYPE[x["t"]]
        else
            fmt = haskey(ECMA_NUMBER_FORMAT, fmt_id) ? ECMA_NUMBER_FORMAT[fmt_id] :
                  wb.custom_number_format[fmt_id]
            fmt[1]
        end
    (fmt_id, T)
end

function get_cellvalue(::Type{Bool}, wb, x)::Bool
    if x == "1" || x == "0"
        x = parse(Int8, x)
    end
    Bool(x)
end
function get_cellvalue(::Type{DateTime}, wb, x)
    x
end
get_cellvalue(::Type{ExcelDateTime}, wb, x) = parse(ExcelDateTime, x)
get_cellvalue(::Type{Missing}, wb, x) = missing
get_cellvalue(::Type{String}, wb, x) = String(x)
get_cellvalue(::Type{Float64}, wb, x) = parse(Float64, x)
get_cellvalue(::Type{T}, wb, x) where T<:Integer= parse(T, x)
function get_cellvalue(::Type{SharedString}, wb, x)
    wb.shared_strings[parse(Int, x)+1]
end
get_cellvalue(::Type{GeneralFormat}, wb, x) = parse(GeneralFormat, x)
get_cellvalue(::Type{ScientificFormat}, wb, x) = parse(ScientificFormat, x)

# peel off 'Cell{T}' and returns T(value)
value(x::Cell) = x.v
value(x::Missing) = missing
function peeloff(ws::WorkSheet)
    x = ws.data
    isa(x, Array{Union{EzExcel.Cell, Missing}}) ? _peeloff(x) : x
end
function _peeloff(datas)
    v = value.(datas)
    T = unique(typeof.(v))
    if length(T) < 3
        convert(Array{Union{T...}}, v)
    else
        convert(Array{Any}, v)
    end
end


using Base: Symbol
import LinearAlgebra

# We wanna implement some special stuff which might be bad for Base.Expr
struct MyExpr
    e::Expr

    function MyExpr(e::Expr, should_cleanup::Bool = true)
        if should_cleanup
            new(cleanup(e))
        else
            new(flatten_once(e))
        end
    end
end
MyExpr(args...; cleanup = true) = MyExpr(Expr(args...), cleanup)
Base.getproperty(e::MyExpr, field::Symbol) = getproperty(getfield(e, :e), field)
Base.:(==)(a::MyExpr, b::MyExpr) = getfield(a, :e) == getfield(b, :e)
expr2myexpr(x::Expr) = MyExpr(x)
expr2myexpr(x) = x
function Base.show(io::IO, me::MyExpr)
    print(io, "MyExpr(")
    e = getfield(me, :e)
    show(io, e.head)
    for arg in e.args
        print(io, ", ")
        show(io, arg)
    end
    print(io, ")")
end

function cleanup(e::Expr)
    if e.head == :call
        head = e.args[1]
        args = e.args[2:end]
    else
        head = e.head
        args = e.args
    end

    new_args = Any[]

    for child in args
        if child isa Expr
            child = cleanup(child)
        end
        if (child isa Expr || child isa MyExpr) && child.head == e.head
            append!(new_args, child.args)
        elseif child isa LineNumberNode; nothing;
        else
            push!(new_args, child)
        end
    end
    Expr(head, new_args...)
end

# I.e. a c or c^†
struct Operator{T1, T2}
    name::Symbol
    daggered::Bool
    index::T1
    time::T2
    function Operator(n, d, i, t)
        _i = expr2myexpr(i); _t = expr2myexpr(t)
        new{typeof(_i), typeof(_t)}(n, d, _i, _t)
    end
end

# constructors
create(idx, time=0; name = :c) = Operator(name, true, idx, time)
annihilate(idx, time=0; name = :c) = Operator(name, false, idx, time)
c(idx, time=0; name = :c) = annihilate(idx, time, name=name)
# cd(idx, time=0; name = :c) = create(idx, time, name=name)
function Base.:(==)(a::Operator, b::Operator)
    a.name == b.name && a.daggered == b.daggered && a.index == b.index && a.time == b.time
end

ExpectationValue(x::MyExpr) = MyExpr(:ExpectationValue, x)

struct ArrayElement
    name::Symbol
    indices::Vector{Any}
    ArrayElement(name::Symbol, idxs::Vector) = new(name, expr2myexpr.(idxs))
end
ArrayElement(name::Symbol, idxs...) = ArrayElement(name, collect(idxs))
function Base.:(==)(a::ArrayElement, b::ArrayElement)
    a.name == b.name && a.indices == b.indices
end

struct Sym
    name::Symbol
end
Base.:(==)(a::Sym, b::Sym) = a.name == b.name

# base extensions
LinearAlgebra.adjoint(o::Operator) = Operator(o.name, !o.daggered, o.index, o.time)
const OpOrExpr = Union{ArrayElement, Operator, MyExpr, Sym}
Base.:(-)(a::OpOrExpr) = MyExpr(:*, -1, a)
for T1 in (OpOrExpr, Number), T2 in (OpOrExpr, Number)
    T1 == Number && T2 == Number && continue
    Base.:(*)(a::T1, b::T2) = MyExpr(:*, a, b)
    Base.:(+)(a::T1, b::T2) = MyExpr(:+, a, b)
    Base.:(-)(a::T1, b::T2) = MyExpr(:+, a, -b)
    Base.:(*)(a::T1, b::T2, c, ds...) = *(MyExpr(:*, a, b), c, ds...)
    Base.:(+)(a::T1, b::T2, c, ds...) = +(MyExpr(:+, a, b), c, ds...)
end

########################################
### Printing
########################################

function Base.show(io::IO, o::Operator)
    print(io, o.name)
    o.daggered && print(io, "^†")
    print(io, "_{")
    _print_expr(io, o.index)
    print(io, "}(", o.time, ")")
end

print_expr(e::MyExpr) = print_expr(stdout, e)
function print_expr(io::IO, e::MyExpr)
    _print_expr(io, e)
    println(io)
end
function _print_expr(io, e::MyExpr)
    e.head == :+ && print(io, "(")
    e.head == :ExpectationValue && print(io, "⟨")
    for (i, child) in enumerate(e.args)
        i > 1 && print(io, " ", e.head, " ")
        _print_expr(io, child)
    end
    e.head == :ExpectationValue && print(io, "⟩")
    e.head == :+ && print(io, ")")
    nothing
end
function _print_expr(io, e::ArrayElement)
    print(io, e.name, "_{")
    for (i, idx) in enumerate(e.indices)
        i > 1 && print(io, ", ")
        _print_expr(io, idx)
    end
    print(io, "}")
end
_print_expr(io, e::Sym) = print(io, e.name)
_print_expr(io, e) = print(io, e)

########################################
### Expand brackets & flatten tree
########################################

# expand a * (b + c) -> a*b + a*c
# flatten nested operations
function expand(e::MyExpr)
    for i in eachindex(e.args)
        e.args[i] = expand(e.args[i])
    end
    new_expr = _expand(e)
    while new_expr != e
        e = new_expr
        new_expr = _expand(e)
    end
    return e
end
expand(e) = e

function _expand(e::MyExpr)
    # expand *(+(a, b), x) -> +(*(a, x), *(b, x))
    if e.head == :*
        i = 1
        while i <= length(e.args)
            child = e.args[i]
            if child isa MyExpr && child.head == :+
                products = map(child.args) do x
                    MyExpr(:*, e.args[1:i-1]..., x, e.args[i+1:end]..., cleanup=false)
                end
                new_expr = MyExpr(:+, products..., cleanup=false)
                # deal with other possible +(c, d) terms
                for j in eachindex(products)
                    new_expr.args[j] = _expand(new_expr.args[j])
                end
                # possibly flatten additions
                return _expand(new_expr)
            elseif child isa MyExpr && child.head == :*
                deleteat!(e.args, i)
                for x in child.args
                    insert!(e.args, i, x)
                    i += 1
                end
            else
                i += 1
            end
        end
    elseif e.head == :+
        i = 1
        while i < length(e.args)
            child = e.args[i]
            if child isa MyExpr && child.head == :+
                deleteat!(e.args, i)
                append!(e.args, child.args)
            else
                i += 1
            end
        end
    elseif e.head == :ExpectationValue
        if e.args[1].head == :+
            terms = [MyExpr(:ExpectationValue, x, cleanup=false) for x in e.args[1].args]
            return MyExpr(:+, terms..., cleanup=false)
        end
    end
    return e
end
_expand(e) = e

# flatten one level of nested operators
function flatten_once(e::Expr)
    new_args = Any[]
    for child in e.args
        if (child isa Expr || child isa MyExpr) && child.head == e.head
            append!(new_args, child.args)
        else
            push!(new_args, child)
        end
    end
    Expr(e.head, new_args...)
end

########################################
### Apply Wicks theorem
########################################

# Wicks theorem
# ⟨abcd⟩ ->  ⟨ab⟩⟨cd⟩ - ⟨ac⟩⟨bd⟩ + ⟨ad⟩⟨bc⟩ etc
function wicks(e)
    if e.head == :+
        products = wicks.(e.args)
        MyExpr(:+, products...)
    elseif e.head == :ExpectationValue
        wicks(e.args[1])
    else
        _wicks(e)
    end
end

function _wicks(e)
    @assert e.head == :*
    skip = filter(i -> !(e.args[i] isa Operator), eachindex(e.args))
    groups = generate_pairings(e.args, copy(skip))
    products = map(enumerate(groups)) do (i, g)
        evs = map(p -> MyExpr(:ExpectationValue, MyExpr(:*, p...)), g)
        if isodd(i)
            # index 1 3 5 ... have even numbers pair-permutations
            MyExpr(:*, e.args[skip]..., evs...)
        else
            MyExpr(:*, -1, e.args[skip]..., evs...)
        end
    end
    MyExpr(:+, products...)
end

function generate_pairings(v::Vector{T}, skip=Int[]) where T
    length(v) - length(skip) == 0 && return Vector{Any}[]
    @assert (length(v) - length(skip)) % 2 == 0 "$v"
    combos = Vector{Any}[]
    first_idx = findfirst(i -> !(i in skip), eachindex(v))
    push!(skip, first_idx)

    for i in first_idx+1:length(v)
        if !(i in skip)
            if length(skip)+1 < length(v)
                rest = generate_pairings(v, [skip; i])
                append!(combos, [vcat(Pair(v[first_idx], v[i]), r) for r in rest])
            else
                push!(combos, [Pair(v[first_idx], v[i])])
            end
        end
    end

    return combos
end

########################################
### convert expectation values to Greens functions
########################################

function to_greens(e::MyExpr)
    if e.head == :+
        new_args = Any[]
        for i in eachindex(e.args)
            converted = _to_greens(e.args[i])
            converted == 0.0 || push!(new_args, converted)
        end
        return MyExpr(:+, new_args...)
    else
        return _to_greens(e)
    end
end

function _to_greens(e::MyExpr)
    @assert e.head == :*

    outputs = []
    for (i, x) in enumerate(e.args)
        if x isa MyExpr && x.head == :ExpectationValue
            y = x.args[1]
            y.head == :* || error(y)
            length(y.args) == 2 || error(y)
            a, b = y.args
            if a.daggered && !b.daggered
                # c^† c - like
                delta = if a.time == b.time
                    ifelse(a.index == b.index, 1, ArrayElement(:I, a.index, b.index))
                elseif a.time isa Number && b.time isa Number
                    0
                else
                    ArrayElement(:I, a.index, b.index) * ArrayElement(:I, a.time, b.time)
                end
                x = delta - ArrayElement(Symbol(:G, b.time, a.time), b.index, a.index)
                push!(outputs, x)
            elseif !a.daggered && b.daggered
                # c c^†
                # GreensElement(a.indices[1], b.indices[2])
                push!(outputs, ArrayElement(Symbol(:G, a.time, b.time), a.index, b.index))
            else
                return 0.0
            end
        else
            push!(outputs, x)
        end
    end

    return *(outputs...)
end

########################################
### Replacements
########################################

struct PLACEHOLDER{T} end
PLACEHOLDER() = PLACEHOLDER{Any}()
struct REPLACEMENT
    x::Int64
end

replace(e, replacements::Pair...) = _replace1(e, replacements)
function _replace1(e, replacements::Tuple)
    e = _replace2(e, replacements)
    for (matchable, replacement) in replacements
        is_match, captures = match(e, matchable)
        if is_match
            return apply_replace(replacement, captures)
        end
    end
    return e
end
function _replace2(e::MyExpr, replacements)
    MyExpr(
        _replace1(e.head, replacements),
        [_replace1(child, replacements) for child in e.args]...
    )
end
function _replace2(ae::ArrayElement, replacements)
    ArrayElement(
        _replace1(ae.name, replacements),
        [_replace1(idx, replacements) for idx in ae.indices]...
    )
end
_replace2(e, replacements) = e

match(e, ::PLACEHOLDER{Any}) = true, Any[e]
match(e::T, ::PLACEHOLDER{T}) where {T} = true, Any[e]
match(a, b) = a == b, Any[]
function match(e::ArrayElement, m::ArrayElement)
    passthrough = Any[]
    if e.name != m.name || length(e.indices) != length(m.indices)
        return false, passthrough
    end
    for (a, b) in zip(e.indices, m.indices)
        is_match, x = match(a, b)
        is_match || return false, passthrough
        passthrough = vcat(passthrough, x)
    end
    return true, passthrough
end
function match(e::MyExpr, m::MyExpr)
    passthrough = Any[]
    if e.head != m.head || length(e.args) != length(m.args)
        return false, passthrough
    end
    for (a, b) in zip(e.args, m.args)
        is_match, x = match(a, b)
        is_match || return false, passthrough
        passthrough = vcat(passthrough, x)
    end
    return true, passthrough
end

apply_replace(r::REPLACEMENT, captures::Vector) = captures[r.x]
apply_replace(x, captures::Vector) = x
function apply_replace(r::MyExpr, captures::Vector)
    MyExpr(
        apply_replace(r.head, captures),
        (apply_replace(x, captures) for x in r.args)...
    )
end
function apply_replace(r::ArrayElement, captures::Vector)
    ArrayElement(
        apply_replace(r.name, captures),
        (apply_replace(x, captures) for x in r.indices)...
    )
end

function cleanup(e::MyExpr)
    replace(
        e,
        MyExpr(:+, PLACEHOLDER(), 0) => REPLACEMENT(1),
        MyExpr(:+, 0, PLACEHOLDER()) => REPLACEMENT(1),
        MyExpr(:*, PLACEHOLDER(), 0) => 0,
        MyExpr(:*, 0, PLACEHOLDER()) => 0,
        MyExpr(:*, PLACEHOLDER(), 1) => REPLACEMENT(1),
        MyExpr(:*, 1, PLACEHOLDER()) => REPLACEMENT(1),
    )
end

########################################
### To code
########################################

function to_expr(e::MyExpr)
    if e.head in (:+, :*)
        return Expr(:call, e.head, to_expr.(e.args)...)
    elseif e.head == :ExpectationValue
        error("Expectation values must be transformed.")
    else
        error("What is a $(e.head)?")
    end
end
to_expr(ae::ArrayElement) = Expr(:ref, ae.name, to_expr.(ae.indices)...)
to_expr(s::Symbol) = s
to_expr(x::Number) = x

function generate_kernel_function(
        e::Expr; mc_type = Any, model_type = Any,
        index_unpacking = default_index_unpacking(e),
        greens_unpacking = default_greens_unpacking(e),
    )

    func_head = Expr(
        :tuple,
        Expr(Symbol("::"), :mc, Symbol(mc_type)),
        Expr(Symbol("::"), :model, Symbol(model_type)),
        :packed_indices, :packed_greens,
    )
    func_body = Expr(
        :block,
        :(N = length(lattice(model))),
        index_unpacking, greens_unpacking,
        e
    )
    Expr(Symbol("->"), func_head, func_body)
end

function default_index_unpacking(e)
    indices = filter(idx -> idx != :N, collect_indices(e))
    if length(indices) == 2
        @assert indices == Set([:i, :j])
        return :(i, j = packed_indices)
    elseif length(indices) == 4
        @assert indices == Set([:i, :j, :k, :l])
        return :(i, j, k, l = packed_indices)
    else
        error("We can only deal with 2 or 4 indices atm. $(length(indices))")
    end
end

function default_greens_unpacking(e)
    Gs = filter(idx -> idx != :N, collect_greens(e))
    if length(Gs) == 1
        @assert Gs == Set([:G00])
        return :(G00 = packed_greens)
    elseif length(Gs) == 4
        @assert Gs == Set([:G00, :G0l, :Gl0, :Gll])
        return :(G00, G0l, Gl0, Gll = packed_greens)
    else
        error("We can only deal with 2 or 4 indices atm. $(length(Gs))")
    end
end

function collect_indices(e::Expr, indices=Set{Symbol}(), is_index=false)
    if e.head in (:ref, :call)
        for arg in e.args[2:end]
            collect_indices(arg, indices, true)
        end
    elseif e.head == :block
        for arg in e.args
            collect_indices(arg, indices, is_index)
        end
    else
        error("Failed to parse $(e.head)")
    end
    return indices
end
function collect_indices(s::Symbol, indices, is_index)
    is_index && push!(indices, s)
end
collect_indices(_, _, _) = nothing

function collect_greens(e::Expr, names=Set{Symbol}())
    if e.head == :ref
        @assert length(e.args) > 1
        push!(names, e.args[1])
    elseif e.head == :call
        for arg in e.args[2:end]
            collect_greens(arg, names)
        end
    elseif e.head == :block
        for arg in e.args
            collect_greens(arg, names)
        end
    else
        error("Failed to parse $(e.head)")
    end
    return filter(name -> startswith(string(name), 'G'), names)
end
collect_greens(_, _) = nothing

macro expand2formula(code)
    code |> expand |> wicks |> to_greens |> to_expr
end

macro expand2kernel(code)
    code |> expand |> wicks |> to_greens |> to_expr |> generate_kernel_function
end


function get_frame(f,frame)

    stream = f.avin.video_info[1].stream
    stream_info = f.avin.video_info[1]
    seek_stream_index = stream_info.stream_index0

    #timeBase = stream.time_base.num * VideoIO.AV_TIME_BASE / f.avin.video_info[1].codec_ctx.time_base.den
    timeBase = stream.time_base.num * VideoIO.AV_TIME_BASE / 500

    VideoIO.av_seek_frame(f.avin.apFormatContext[1],0,round(Int64,frame),VideoIO.AVSEEK_FLAG_ANY | VideoIO.AVSEEK_FLAG_FRAME)
    #VideoIO.avcodec_flush_buffers(f.pVideoCodecContext)
    #while !VideoIO.have_frame(f)
            #idx = VideoIO.pump(f.avin)
            #idx == f.stream_index0 && break
            #idx == -1 && throw(EOFError())
    #end
    #VideoIO.reset_frame_flag!(f)
    nothing
end

function plot_spikes(gui)

    #The number of lines to plot is equal to the number of spikes
    # + threshold line
    # + event triggered average
    # + event_triggered average axis
    gui.spikes=[Point2f0(0.0,0.0) for i=1:50,j=1:length(gui.spikes_ts)+3]

    for i=1:length(gui.spikes_ts)
        for j=1:50
            gui.spikes[j,i] = Point2f0(j*5+200,gui.y_data[gui.spikes_ts[i]-10+j,1].*gui.y_scales[1]+300)
        end
    end

    for j=1:50
        gui.spikes[j,end-2] = Point2f0(j*5+200,gui.s.thres*gui.y_scales[1]+300)
    end

    for j=1:length(gui.myhist)
        norm_hist = gui.myhist[j]./(sum(gui.myhist)+1)
        gui.spikes[j,end-1] = Point2f0(j*10+200,norm_hist.*100+50)
    end

    for j=1:length(gui.myhist)
        gui.spikes[j,end] = Point2f0(j*10+200,50)
    end

    nothing
end

function plot_lines(gui,t; increment=2,p_per_frame=100, xoff=0, yoff=100)

    #Since this is automatically advancing, we can slow things down here if we wish
    #By only plotting if a certain number of frames have passed
    #If we want to speed things up, we could alternatively change more than 1 point on the axis by adding
    #Points from between the previous slider value and current slider value.

    #For example, if we are at real time, 1000 data points pass every frame (30 fps)
    #If we are 1/10 real time, then 100
    #1/100 real time then 10 data point every frame
    #1/1000 real time then 1 data point every frame

    #We have 500 data points on the x axis.
    #However long our t_duration is, should be flushed completely after that amount
    #of data has been processed. At 1/10 real time, we are passing 3000 every second
    #So 30000, or one second means that 30000 / 30 = 1000 points per frame updated

    #First, our total number of points should be closest whole number
    #Points / (# frames for data)
    # 600 / (300)  = 16 or 480 points total

    total_points = 500
    remainder_points = total_points - increment

    for mycov=1:3

        for i=1:remainder_points
            gui.cov1[i,mycov] = Point2f0(i+xoff,gui.cov1[i+increment,mycov][2])
        end

        x_inds = linspace(1,p_per_frame,increment+1)
        for i=1:increment

            x_point = remainder_points+i+xoff

            l_bound = t-p_per_frame + round(Int64,x_inds[i])
            r_bound = t-p_per_frame + round(Int64,x_inds[i+1])-1

            if (l_bound >= r_bound)
                y_val = gui.y_data[l_bound,mycov]
            else
                min_y = minimum(gui.y_data[l_bound:r_bound,mycov])
                max_y = maximum(gui.y_data[l_bound:r_bound,mycov])

                if (abs(min_y)>max_y)
                    y_val = min_y
                else
                    y_val = max_y
                end
            end

            y_point = y_val*gui.y_scales[mycov]+yoff*mycov

            gui.cov1[i+remainder_points,mycov] = Point2f0(x_point,y_point)
        end

    end
end


module CAS_module

push!(LOAD_PATH, "./")
push!(LOAD_PATH, "../")
push!(LOAD_PATH, "../ACASXU_Dependencies")
push!(LOAD_PATH, "../../VICAS/src")
push!(LOAD_PATH, "../../DeepCorrection")
push!(LOAD_PATH, "../../DeepCorrection/src")

using Simulation_Env_module
using JLD, CAS_QMDP, DeepQLearning

export CAS
export ACAS_Xu_Network, Simple_CAS, No_CAS, VICAS, Corrected_VICAS, VICASWeighted # NNCAS
export CAS_DICT
export acas_xu_network_data, evaluate_network
export alphas, get_qvals, norm_angle
export policy, STATE_DIM, AUG_STATE_DIM, TERM_VAR
export correction_version, VICAS_policyname, correction_type, correction_weight

abstract type CAS end

include("../ACASXU_Dependencies/load_network.jl")
const nnfile_name = "../ACASXU_Dependencies/ACASXU_TF12_run3_DS_Minus15_Online_08172017_200Epochs_largerRange.nnet"
const acas_xu_network_data = NNet(nnfile_name)

mutable struct ACAS_Xu_Network <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function ACAS_Xu_Network()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :COC, 2 => :WL, 3 => :WR, 4 => :SL, 5 => :SR)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => 0., :WL => 5., :WR => -5., :SL => 10., :SR => -10.)
		return this
	end
end

# const VICAS_policyname = "test_mid3"
const VICAS_policyname = "mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0"
const alphas_filename =
	"../../VICAS/policies/" * VICAS_policyname * ".jld" # requires disc = mid3
const alphas = load_alphas(alphas_filename)

mutable struct VICAS <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function VICAS()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
		return this
	end
end

mutable struct VICASWeighted <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function VICASWeighted()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
		return this
	end
end

# const correction_type = :sector
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-13T18\:55\:43.491" # best sector
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-15T22\:38\:58.856"
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-28T10\:55\:10.697"
const correction_type = :closest
const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-28T10\:42\:31.084" # best closest
# const correction_version = "LOWFI_mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0_2018-11-29T22\:18\:30.068"
# requires disc = mid3
const problem_file = "../../DeepCorrection/logs/" * correction_version * "/final_problem.jld"
const weights_file = "../../DeepCorrection/logs/" * correction_version * "/final_weights.jld"
policy = restore(problem_file=problem_file, weights_file=weights_file)
correction_weight = policy.env.correction_weight

mutable struct Corrected_VICAS <: CAS
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function Corrected_VICAS()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
		return this
	end
end

# const nn_version = ""
# const nn_problem_file = "../../DeepCorrection/logs/" * nn_version * "/final_problem.jld"
# const nn_weights_file = "../../DeepCorrection/logs/" * nn_version * "/final_weights.jld"
# nn_policy = restore(problem_file=nn_problem_file, weights_file=nn_weights_file)
# nn_policy.env.correction = true

# mutable struct NNCAS <: CAS
# 	Ind_to_Advisory_dict::Dict{Int64, Symbol}
# 	Advisory_to_Ind_dict::Dict{Symbol, Int64}
# 	Advisory_to_Action_dict::Dict{Symbol, Float64}

# 	function NNCAS()
# 		this = new()
# 		this.Ind_to_Advisory_dict =
# 				Dict(1 => :SR, 2 => :WR, 3 => :KEEP, 4 => :WL, 5 => :SL, 6 => :COC)
# 		this.Advisory_to_Ind_dict =
# 				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
# 		this.Advisory_to_Action_dict =
# 				Dict(:COC => -1.0, :WL => 5., :WR => -5., :SL => 10., :SR => -10., :KEEP => 0.0)
# 		return this
# 	end
# end

"""
	outdated
"""
mutable struct Simple_CAS <: CAS
	alert_threshold::Int64
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function Simple_CAS()
		this = new()
		set_alert_threshold = 1320. * 0.6 # [ft] ≈ 400 [m]
		this.alert_threshold = set_alert_threshold
		this.Ind_to_Advisory_dict =
				Dict(1 => :COC, 2 => :WL, 3 => :WR, 4 => :SL, 5 => :SR)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => 0., :WL => 5., :WR => -5., :SL => 10., :SR => -10.)
		return this
	end
end

mutable struct No_CAS <: CAS
	dt::Float64
	Ind_to_Advisory_dict::Dict{Int64, Symbol}
	Advisory_to_Ind_dict::Dict{Symbol, Int64}
	Advisory_to_Action_dict::Dict{Symbol, Float64}

	function No_CAS()
		this = new()
		this.Ind_to_Advisory_dict =
				Dict(1 => :COC, 2 => :WL, 3 => :WR, 4 => :SL, 5 => :SR)
		this.Advisory_to_Ind_dict =
				Dict(value => key for (key, value) in this.Ind_to_Advisory_dict)
		this.Advisory_to_Action_dict =
				Dict(:COC => 0., :WL => 5., :WR => -5., :SL => 10., :SR => -10.)
		return this
	end
end

# construct all CAS' in advance to avoid repetitive contruction
CAS_DICT = Dict{DataType, CAS}()
for i = 1:length(subtypes(CAS))
	CAS_DICT[subtypes(CAS)[i]] = subtypes(CAS)[i]()
end

println("All CAS loaded!")

end # module


## ----------------------------------------------------------------------------
function plot_inner_plots(METHODS, METHODS_LABELS = string.(METHODS))

    # This is just for Nanchen
    Data = Nd
    iJR = NiJR
    src = nameof(Data)

    DAT = ChE.load_DAT(src)

    # experiments with approx D
    Dintervals = Dict(
        0.05 => 1:3,
        0.1 => 4:9,
        0.15 => 4:11,
        0.2 => 10:11,
        0.25 => 10:15,
        0.3 => 12:15,
        0.35 => 12:17,
        0.4 => 16:17,
        :tot => 1:17
    )

    int_ps = Dict()
    for (avD, EXPS) in Dintervals

        Nd_rxns_map = iJR.load_rxns_map()
        iJR_ider_subs = iJR.load_inner_rxns_subs()

        internals = iJR.load_inners_idermap() |> keys
        color_pool = Plots.distinguishable_colors(length(internals))
        ider_colors = Dict(id => c for (id, c) in zip(internals, color_pool))

        ps = Plots.Plot[]
        for (method_label, method) in zip(METHODS_LABELS, METHODS)

            p = plot(;
                xlabel = _textbf("exp abs. flux"),
                ylabel = _textbf("model abs. flux")
            )

            vals = []

            for sub in NiJR.subs()
                iJR_iders = iJR_ider_subs[sub]

                (sub in ["others"]) && continue
                marker = (10, subs_markers[sub])

                for iJR_ider in iJR_iders

                    Nd_ider = Nd_rxns_map[iJR_ider]

                    @info "Doing" Nd_ider iJR_ider

                    exp_exglcs = abs.(DAT[:exp, :flx, "GLC", EXPS])
                    exp_vals = abs.(DAT[:exp, :flx, Nd_ider, EXPS]) ./ exp_exglcs
                    model_vals = abs.(DAT[method, :flx, iJR_ider, EXPS]) ./ exp_exglcs
                    model_errs = abs.(get.([DAT], 0.0, method, :err, iJR_ider, EXPS)) ./ exp_exglcs
                    model_errs = map(model_errs) do err
                        isnan(err) && return 0.0
                        err > 1.0 && return 0.0
                        err
                    end

                    push!(vals, exp_vals..., model_vals...)

                    scatter!(p, exp_vals , model_vals;
                        label = "",
                        yerr = model_errs,
                        color = ider_colors[iJR_ider],
                        alpha = 0.8,
                        marker,
                    )

                end # for iJR_ider

            end # for (sub, iJR_iders)

            sort!(vals)
            plot!(p, [vals], [vals]; label = "",
                ls = :dash, alpha = 0.8, lw = 3,
                color = :black,
            )

            # title
            m, M = minimum(vals), maximum(vals)
            x = m + (M - m) * 0.01
            y = M - (M - m) * 0.05
            Dstr = (avD == :tot) ? "all" : round(avD; sigdigits=2)
            text = _textbf(method_label)
            annotate!(p, [(x, y, (text, 26, :left, :top, :black))])
            text = _textbf("D=", Dstr)
            annotate!(p, [(x, y - y * 0.13, (text, 26, :left, :top, :black))])

            push!(ps, p)

        end # for (method_label

        # sfig(ChE, ps,
        #     @fileid, "inner_flx_corr", (;avD), ".png",
        #     layout = (1, length(ps))
        # )

        int_ps[avD] = ps
    end

    return int_ps
end


# using Stars
# using Test
# cd("test")

@testset "SpatialPixelsDataFrame" begin
    b = bbox(70, 15, 140, 55)
    proj = 4326
    coords = [1 2; 3 4]
    data = [1.0 2; 3 4]

    s = Spatial(b, proj)
    @test typeof(s) == Spatial{bbox, Int64} #

    sp = SpatialPoints(coords, b, proj)
    @test typeof(sp) == SpatialPoints{Matrix{Int64}, bbox, Int64}

    sp_df = SpatialPixelsDataFrame(data, coords, b, proj)
    @test typeof(sp_df) == SpatialPixelsDataFrame{Matrix{Float64}, Matrix{Int64}, bbox, Int64}

    @test st_bbox(s) == b
    @test st_bbox(sp) == b
    @test st_bbox(sp_df) == b

    @test st_proj(s) == proj
    @test st_proj(sp) == proj
    @test st_proj(sp_df) == proj

    @test st_coords(sp) == coords
    @test st_coords(sp_df) == coords
end


module ODDS

    using OutlierDetection: CLASS_NORMAL, CLASS_OUTLIER, to_categorical
    using DataFrames: DataFrame, convert
    using DataDeps
    using MAT

    export list, load

    const ODDS_DATASETS = [
        (
            name = "annthyroid",
            md5 = "17a1012d005f2179eca3b87c0028eb5ca807794b359c33fa1f8ca5ac846c514f",
            url = "https://www.dropbox.com/s/aifk51owxbogwav/annthyroid.mat?dl=1"
        ),
        (
            name = "arrhythmia",
            md5 = "e4d47e499591ae31e5d7f0cd983367092353eac240f48e14c9b36457edb3460d",
            url = "https://www.dropbox.com/s/lmlwuspn1sey48r/arrhythmia.mat?dl=1"
        ),
        (
            name = "breastw",
            md5 = "a2e3f7c7f61ddde87c4373f0c0539de996227b9659c70b3c227948ebf0c6062b",
            url = "https://www.dropbox.com/s/g3hlnucj71kfvq4/breastw.mat?dl=1"
        ),
        (
            name = "cardio",
            md5 = "5e568b6a73711012c481fa6cc984ab9c68e9a4fa44d6aeb3b43efa381c3f695b",
            url = "https://www.dropbox.com/s/galg3ihvxklf0qi/cardio.mat?dl=1"
        ),
        (
            name = "cover",
            md5 = "39acf6da4d31f08c2c28dc047249d8cd8ccd9d6fc299d9b5a08874ac92290c98",
            url = "https://www.dropbox.com/s/awx8iuzbu8dkxf1/cover.mat?dl=1"
        ),
        # the MAT import for ecoli is currently not working
        # (
        #     name = "ecoli",
        #     md5 =  "76886bf14f549680450a0720c58b461db965f8b762a5ab4d165c03bbdcb355fc",
        #     url =  "https://www.dropbox.com/s/f38mquv0g08xxgp/ecoli.mat?dl=1"
        # ),
        (
            name = "glass",
            md5 = "2c506e6e1e2ca1e82a905d10a90773151ff812c4bd969c71a12db5b2bb6e7921",
            url = "https://www.dropbox.com/s/iq3hjxw77gpbl7u/glass.mat?dl=1"
        ),
        (
            name = "http",
            md5 = "81d689cfec5407305ab98cd26794e2cefb3aef5a521bf9c59744107aabcaed7e",
            url = "https://www.dropbox.com/s/iy9ucsifal754tp/http.mat?dl=1"
        ),
        (
            name = "ionosphere",
            md5 = "8ac5dfa6f47db4acdfe10615489b7f7dd015d79d5b381ce641e72baca3cfefb0",
            url = "https://www.dropbox.com/s/lpn4z73fico4uup/ionosphere.mat?dl=1"
        ),
        (
            name = "letter",
            md5 = "f4a229b0e2a6b0c9157ce27fa760bcbac0f948cd690b2f5d51c120e5450f2195",
            url = "https://www.dropbox.com/s/rt9i95h9jywrtiy/letter.mat?dl=1"
        ),
        (
            name = "lympho",
            md5 = "031b907ecb50347720b08b38f8139d05635adeca9335a064d89d5569c27a5834",
            url = "https://www.dropbox.com/s/ag469ssk0lmctco/lympho.mat?dl=1"
        ),
        (
            name = "mammography",
            md5 = "271ebb568314a856666d3504b4882e21b0ea6e1ba9e648ad256d572a36df597e",
            url = "https://www.dropbox.com/s/tq2v4hhwyv17hlk/mammography.mat?dl=1"
        ),
        (
            name = "mnist",
            md5 = "9479e221f236684c09f3e11c7ad124725d34ecd301b9450425e8b113368413a6",
            url = "https://www.dropbox.com/s/n3wurjt8v9qi6nc/mnist.mat?dl=1"
        ),
        (
            name = "mulcross",
            md5 = "f64a10efb641229356e7ae06e8805f9025e1ce29a2b0abb1207447423746276b",
            url = "https://www.dropbox.com/s/tmfan6f7qf46bkx/mulcross.mat?dl=1"
        ),
        (
            name = "musk",
            md5 = "50bbf414caceb66cfe6f6b1eb7c1ff7c6eb1e6d87173f7bb597ba55cd3dc5493",
            url = "https://www.dropbox.com/s/we6aqhb0m38i60t/musk.mat?dl=1"
        ),
        (
            name = "optdigits",
            md5 = "30c664cd6f3740a173277a27cd1028d471b06f0835a9902c9712749cec58c1c4",
            url = "https://www.dropbox.com/s/w52ndgz5k75s514/optdigits.mat?dl=1"
        ),
        (
            name = "pendigits",
            md5 = "7bc32ac83b3d583443acea8b1342a54aaa59eaedb9996af405322f473f3c22bc",
            url = "https://www.dropbox.com/s/1x8rzb4a0lia6t1/pendigits.mat?dl=1"
        ),
        (
            name = "pima",
            md5 = "b6a5be8d548b9293d6958beceb1c77055591714d25d40d957eba87808456468f",
            url = "https://www.dropbox.com/s/mvlwu7p0nyk2a2r/pima.mat?dl=1"
        ),
        (
            name = "satellite",
            md5 = "6feac3112b9c14e1c3e60afc437f2f3d29dc1000119c9f950c628078778d6aa0",
            url = "https://www.dropbox.com/s/dpzxp8jyr9h93k5/satellite.mat?dl=1"
        ),
        (
            name = "satimage-2",
            md5 = "95c37a07cea2cb201f6900d14f521510e10b6b2f0074c06ac46d3157440ec67b",
            url = "https://www.dropbox.com/s/hckgvu9m6fs441p/satimage-2.mat?dl=1"
        ),
        (
            name = "shuttle",
            md5 = "b6e396c92354d5dcf5036901e36dc58e2913f2d1a5e138c35a2bedf9c69db899",
            url = "https://www.dropbox.com/s/mk8ozgisimfn3dw/shuttle.mat?dl=1"
        ),
        (
            name = "smtp",
            md5 = "5b6d13dd23e8c5e4235081ea58b606df216a2b9b86e7f5403cde26f62dafc0d8",
            url = "https://www.dropbox.com/s/dbv2u4830xri7og/smtp.mat?dl=1"
        ),
        (
            name = "speech",
            md5 = "193d85404c5aff8b33e7e0540a8be211a7c94948d1a33162ce500b7e526ea937",
            url = "https://www.dropbox.com/s/w6xv51ctea6uauc/speech.mat?dl=1"
        ),
        (
            name = "thyroid",
            md5 = "49155d7861b977e90808faf68343ff67a67f63090c2744d514dbb67ac68e7e33",
            url = "https://www.dropbox.com/s/bih0e15a0fukftb/thyroid.mat?dl=1"
        ),
        (
            name = "vertebral",
            md5 = "84b50b055f118f54b5c976650ccd284ab7c1adee822f8b3eb517c7e2c890d15e",
            url = "https://www.dropbox.com/s/5kuqb387sgvwmrb/vertebral.mat?dl=1"
        ),
        (
            name = "vowels",
            md5 = "51d5397bad1ceeda22b7cd44fd58125c5d603c9c2edb53aa59933330e259ac7b",
            url = "https://www.dropbox.com/s/pa26odoq6atq9vx/vowels.mat?dl=1"
        ),
        (
            name = "wbc",
            md5 = "afe6b26953cc196b9a480318c903e0edef22529c2f61bd02082438c6e70630a7",
            url = "https://www.dropbox.com/s/ebz9v9kdnvykzcb/wbc.mat?dl=1"
        ),
        (
            name = "wine",
            md5 = "d1b7cc643379e14397ee25bc1125ba08e4c929f2ec44d8782b1340d7045d4cf7",
            url = "https://www.dropbox.com/s/uvjaudt2uto7zal/wine.mat?dl=1"
        )
    ]

    list() = [name for (name, _, _) in ODDS_DATASETS]
    to_name(dataset::AbstractString) = "odds-$dataset"

    function load(dataset::AbstractString)
        @assert dataset in list()

        dep = @datadep_str to_name(dataset)
        data = MAT.matread("$dep/$dataset.mat")
        X = DataFrame(data["X"], :auto)
        y = ifelse.(data["y"][:, 1] .== 0, CLASS_NORMAL, CLASS_OUTLIER)
        X, to_categorical(y)
    end

    function __init__()
        for (name, md5, url) in ODDS_DATASETS
            register(DataDep(
                to_name(name),
                """Dataset: $name
                Collection: Outlier Detection DataSets (ODDS)
                Authors: Shebuti Rayana
                Website: http://odds.cs.stonybrook.edu""",
                url,
                md5
            ))
        end
    end

end


function ShapiroInterSatelliteRangeModifier(arg0::jdouble)
    return ShapiroInterSatelliteRangeModifier((jdouble,), arg0)
end

function get_parameters_drivers(obj::ShapiroInterSatelliteRangeModifier)
    return jcall(obj, "getParametersDrivers", List, ())
end

function modify(obj::ShapiroInterSatelliteRangeModifier, arg0::EstimatedMeasurement)
    return jcall(obj, "modify", void, (EstimatedMeasurement,), arg0)
end


using ProjAssistant
@quickactivate

## -------------------------------------------------------------------
@time begin

    import SparseArrays
    using ArgParse

    import Chemostat_Nanchen2006
    const ChN = Chemostat_Nanchen2006

    const iJR = ChN.iJR904
    const Nd = ChN.NanchenData # experimental data
    const Bd = ChN.BegData    # cost data

    import Chemostat
    const Ch = Chemostat
    const ChU = Ch.Utils
    const ChSS = Ch.SteadyState
    const ChLP = Ch.LP
end

## ----------------------------------------------------------------------------
# arg settings
ARGSET = ArgParseSettings()
@ArgParse.add_arg_table! ARGSET begin
    "--ignore-cached"
        help = "Ingnore on disk version of data"
        action = :store_true
end

ARGS_DICT = ArgParse.parse_args(ARGSET)
if isinteractive()
    # dev values
    ignore_cached = true
else
    ignore_cached = ARGS_DICT["ignore-cached"]
end
@info("ARGS", ignore_cached)

## -------------------------------------------------------------------
# Helper function
include("1.0.1_functions.jl")

## -------------------------------------------------------------------
MODELS_FILE = procdir(iJR, "base_models.bson")
MODELS = ldat(() -> Dict{Any, Any}(), iJR, MODELS_FILE);

## -------------------------------------------------------------------
# file globals
use_cached = false

## -------------------------------------------------------------------
let
    compress(model) = ChU.compressed_model(model)

    ## --------------------------------------
    @info "COST_CONSTRAINED MODELS";

    base_model = load_raw_model()
    exchs = ChU.exchanges(base_model)

    rescale_bounds!(base_model)
    close_exchanges!(base_model, exchs)
    base_model = add_cost(base_model)
    base_model = reset_exchanges(base_model, exchs)

    MODELS["base_model"] = compress(base_model)

    ## --------------------------------------
    scale_factor = 1000.0
    # fva_preprocessed models
    for (exp, D) in Nd.val(:D) |> enumerate

        DAT = get!(MODELS, "fva_models", Dict())
        !ignore_cached && haskey(DAT, exp) && continue # cached

        fva_model = make_fva_model(base_model, exp, D;
            scale_factor
        )

        # storing
        DAT[exp] = compress(fva_model)

        # caching
        sdat(iJR, MODELS, MODELS_FILE);
        GC.gc()
    end

    ## --------------------------------------
    # MAX MODEL
    max_model = make_max_model(base_model; scale_factor)

    # saving
    MODELS["max_model"] = compress(max_model)
    sdat(iJR, MODELS, MODELS_FILE)

    println()
end


import AllenNeuropixels as AN
using WGLMakie
using DataFrames
using Statistics
using FileIO

session_table = AN.getsessiontable()

function numprobes(sessionid)
	session_table[session_table[!, :id].==sessionid, :probe_count][1]
end

function targetintersections(sessionid)
	s = session_table[session_table.id.==sessionid, :ecephys_structure_acronyms][1]
	targets = ["VISp", "VISl", "VISrl", "VISal", "VISpm", "VISam", "CA1", "CA3", "DG", "SUB", "ProS", "LGd", "LP", "APN"];
	structures = eachmatch(r"'(.*?)'", s)
	structures = [replace(x.match, r"'"=>s"") for x ∈ structures]
	targetsintersected = [x ∈ structures for x ∈ targets]
	return mean(targetsintersected)
end;

metrics = AN.getunitanalysismetricsbysessiontype("brain_observatory_1.1")

session_metrics = combine(groupby(metrics, :ecephys_session_id),
		:ecephys_session_id => numprobes∘unique => :num_probes,
		:ecephys_session_id => targetintersections∘unique => :target_intersections,
		:has_lfp_data=>all,
		:genotype => (x->all(isequal.(x, ("wt/wt",)))) => :is_wt,
		:max_drift=>median∘skipmissing,
		:d_prime=>median∘skipmissing,
		:isolation_distance=>median∘skipmissing,
		:silhouette_score=>median∘skipmissing,
		:snr=>median∘skipmissing)

oursession = subset(session_metrics,
						:num_probes => ByRow(==(6)),
						:target_intersections => ByRow(>(0.9)),
						:is_wt => ByRow(==(true)),
						:has_lfp_data_all => ByRow(==(true)),
						:max_drift_median_skipmissing => ByRow(<(25.0)))

session = AN.Session(oursession.ecephys_session_id[1])

session = AN.Session(757216464)
probeid = 769322749

# --------------------- Do the plots and save to dropbox --------------------- #
AN.Plots.formattedreferencevolume(session, "C:\\Users\\Brendan\\Dropbox (Sydney Uni Student)\\AllenNeuropixels\\probelocations.html")

data = AN.getdownsampledlfp(session, probeid);
fig = AN.Plots.neuroslidingcarpet(data; resolution=(800, 800));
AN.Plots.htmlexport(fig)


# Optimization and solution query.

"""
Call Knitro to solve the problem. The return value indicates
the solution status.

"""
function KN_solve(m::Model)
    # Check sanity. If model has Julia callbacks, we need to ensure
    # that Knitro is not multithreaded. Otherwise, the code will segfault
    # as we have trouble calling Julia code from multithreaded C
    # code. See issue #93 on https://github.com/jump-dev/KNITRO.jl.
    if has_callbacks(m)
        KN_set_param(m, "par_numthreads", 1)
        KN_set_param(m, "par_msnumthreads", 1)
    end
    # For KN_solve, we do not return an error if ret is different of 0.
    m.status = @kn_ccall(solve, Cint, (Ptr{Cvoid},), m.env)
    return m.status
end

##################################################
# Solution status and info
##################################################
"Return the solution status, objective, primal and dual variables."
function KN_get_solution(m::Model)
    # we first check that the model is well defined to avoid segfault
    @assert m.env != C_NULL
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    x = zeros(Cdouble, nx)
    lambda = zeros(Cdouble, nx + nc)
    status = Cint[0]
    obj = Cdouble[0.]

    ret = @kn_ccall(get_solution, Cint,
                    (Ptr{Cvoid}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
                    m.env, status, obj, x, lambda)

    # Keep solution in cache.
    m.status = status[1]
    m.x = x
    m.mult = lambda
    m.obj_val = obj[1]
    return status[1], obj[1], x, lambda
end

# some wrapper functions for MOI
function get_status(m::Model)
    @assert m.env != C_NULL
    if m.status != 1
        return m.status
    end
    status = Cint[0]
    obj = Cdouble[0.]
    ret = @kn_ccall(get_solution, Cint,
                    (Ptr{Cvoid}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
                    m.env, status, obj, C_NULL, C_NULL)
    # Keep status in cache.
    m.status = status[1]
    return status[1]
end

function get_objective(m::Model)
    @assert m.env != C_NULL
    if isfinite(m.obj_val)
        return m.obj_val
    end
    status = Cint[0]
    obj = Cdouble[0.]
    ret = @kn_ccall(get_solution, Cint,
                    (Ptr{Cvoid}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
                    m.env, status, obj, C_NULL, C_NULL)
    # Keep objective value in cache.
    m.obj_val = obj[1]
    return obj[1]
end

function get_solution(m::Model)
    # We first check that the model is well defined to avoid segfault.
    @assert m.env != C_NULL
    if !isempty(m.x)
        return m.x
    end
    nx = KN_get_number_vars(m)
    x = zeros(Cdouble, nx)
    status = Cint[0]
    obj = Cdouble[0.]
    ret = @kn_ccall(get_solution, Cint,
                    (Ptr{Cvoid}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
                    m.env, status, obj, x, C_NULL)
    # Keep solution in cache.
    m.x = x
    return x
end
get_solution(m::Model, ix::Int) = isempty(m.x) ? get_solution(m)[ix] : m.x[ix]

function get_dual(m::Model)
    # we first check that the model is well defined to avoid segfault
    @assert m.env != C_NULL
    if !isempty(m.mult)
        return m.mult
    end
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)
    lambda = zeros(Cdouble, nx + nc)
    status = Cint[0]
    obj = Cdouble[0.]
    ret = @kn_ccall(get_solution, Cint,
                    (Ptr{Cvoid}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
                    m.env, status, obj, C_NULL, lambda)
    # Keep multipliers in cache.
    m.mult = lambda
    return lambda
end
get_dual(m::Model, ix::Int) = isempty(m.mult) ? get_dual(m)[ix] : m.mult[ix]

# New getters for Knitro >= 12.0
if KNITRO_VERSION >= v"12.0"
    @define_getters get_var_primal_values
    @define_getters get_var_dual_values
    @define_getters get_con_dual_values
end


mutable struct Model <: Renderable
	mesh::Mesh
	material::Material
	transform::Matrix{Float32}
	boundDirty::Bool
	bound::Shapes.Shape

	Model(mesh::Mesh, material::Material) = new(mesh, material, Matrix{Float32}(I, 4, 4), true, similar(mesh.bound))
end

isvalid(model::Model) = isvalid(model.mesh) && isvalid(model.material)

gettransform(model::Model) = model.transform
function settransform(model::Model, m::Matrix{Float32})
	@assert isvalid(model)
	model.transform[:] = m
	model.boundDirty = true
end

localbound(model::Model) = model.mesh.bound

function getbound(model::Model)
	@assert isvalid(model)
	if model.boundDirty
		Shapes.transform(model.bound, model.transform, model.mesh.bound)
		model.boundDirty = false
	end
	model.bound
end

function setbound(model::Model, bound::Shapes.Shape)
	@assert isvalid(model)
	model.bound = bound
	model.boundDirty = false
end

function render(model::Model, renderer::Renderer)
	@assert isvalid(model)
	model.material.shader.setupMaterial(model)
	apply(model.material, renderer)
	render(model.mesh)
end