includet("simulation_model.jl")
Random.seed!(1)
runSCExperiment(1000)


# This file is a part of SimilaritySearch.jl

"""
    find_neighborhood(index::SearchGraph{T}, item, res, vstate)

Searches for `item` neighborhood in the index, i.e., if `item` were in the index whose items should be its neighbors (intenal function).
`res` is always reused since `reduce` creates a new KnnResult from it (a copy if `reduce` in its simpler terms)
"""
function find_neighborhood(index::SearchGraph, item, res, vstate)
    n = length(index)

    if n > 0
        empty!(res, index.neighborhood.ksearch)
        reduce(index.neighborhood.reduce, search(index, item, res; vstate), index)
    else
        KnnResult(index.neighborhood.ksearch)
    end
end

"""
    push_neighborhood!(index::SearchGraph, item, neighbors::KnnResult; apply_callbacks=true)

Inserts the object `item` into the index, i.e., creates an edge from items listed in L and the
vertex created for ìtem` (internal function)
"""
function push_neighborhood!(index::SearchGraph, item, neighbors::KnnResult; apply_callbacks=true)
    push!(index.db, item)
    push!(index.links, neighbors)
    push!(index.locks, Threads.SpinLock())
    n = length(index)
    k = index.neighborhood.k

    @inbounds for (id, dist) in neighbors
        vertex = index.links[id]
        vertex.k = max(maxlength(vertex), k) # adjusting maximum size to the current allowed neighborhood size
        push!(vertex, n, dist)
    end

    apply_callbacks && callbacks(index)

    if index.verbose && length(index) % 100_000 == 0
        println(stderr, "added n=$(length(index)), neighborhood=$(length(neighbors)), $(string(index.search_algo)), $(Dates.now())")
    end
end

const GlobalSatKnnResult = [KnnResult(1)]
function __init__neighborhood()
    for i in 2:Threads.nthreads()
        push!(GlobalSatKnnResult, KnnResult(1))
    end
end

"""
    SatNeighborhood()

New items are connected with a small set of items computed with a SAT like scheme (**cite**).
It starts with `k` near items that are reduced to a small neighborhood due to the SAT partitioning stage.
"""
struct SatNeighborhood <: NeighborhoodReduction end

Base.copy(::SatNeighborhood) = SatNeighborhood()

"""
    reduce(sat::SatNeighborhood, res::KnnResult, index::SearchGraph)

Reduces `res` using the SAT strategy.
"""
function Base.reduce(sat::SatNeighborhood, res::KnnResult, index::SearchGraph)
    near = GlobalSatKnnResult[Threads.threadid()]
    N = KnnResult(ceil(Int, maxlength(res)/2))

    #@inbounds for i in lastindex(res):-1:firstindex(res)  # DistSat => works a little better but also produces a bit larger neighborhoods
    #id, dist = res[i]
    @inbounds for (id, dist) in res
        pobj = index[id]
        empty!(near)
        push!(near, 0, dist)
        for nearID in keys(N)
            d = evaluate(index.dist, index[nearID], pobj)
            push!(near, nearID, d)
        end

        argmin(near) == 0 && push!(N, id, dist)
    end

    N
end

"""
    struct IdentityNeighborhood

It does not modifies the given neighborhood
"""
struct IdentityNeighborhood <: NeighborhoodReduction end

Base.copy(::IdentityNeighborhood) = IdentityNeighborhood()
Base.reduce(sat::IdentityNeighborhood, res::KnnResult, index::SearchGraph) = copy(res)


@testset "random.jl VariableSelection" begin
    @testset "RandomVariableSelection{TakeObjective=true}" begin
        trailer = SeaPearl.Trailer()
        cpmodel = SeaPearl.CPModel(trailer)
        x = SeaPearl.IntVar(1, 2, "x", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x)
        SeaPearl.addVariable!(cpmodel, y)
        SeaPearl.addObjective!(cpmodel,y)

        rng = MersenneTwister(10)

        variableselection = SeaPearl.RandomVariableSelection{true}()

        # This condition is there because of the way random are generated can change from one version to another
        if VERSION >= v"1.6.0"
            @test variableselection(cpmodel; rng=rng) == x
            @test variableselection(cpmodel; rng=rng) == x
            @test variableselection(cpmodel; rng=rng) == y
        end

        # With branchable variables
        trailer = SeaPearl.Trailer()
        cpmodel = SeaPearl.CPModel(trailer)
        x = SeaPearl.IntVar(1, 2, "x", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x)
        SeaPearl.addVariable!(cpmodel, y; branchable=false)
        SeaPearl.addObjective!(cpmodel,y)

        rng = MersenneTwister(10)

        variableselection = SeaPearl.RandomVariableSelection{true}()

        # This condition is there because of the way random are generated can change from one version to another
        if VERSION >= v"1.6.0"
            @test variableselection(cpmodel; rng=rng) == x
            @test variableselection(cpmodel; rng=rng) == x
            @test variableselection(cpmodel; rng=rng) == x
        end
    end
    @testset "RandomVariableSelection{TakeObjective=false}" begin
        trailer = SeaPearl.Trailer()
        cpmodel = SeaPearl.CPModel(trailer)
        x = SeaPearl.IntVar(1, 2, "x", trailer)
        z = SeaPearl.IntVar(1, 2, "z", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x; branchable=false)
        SeaPearl.addVariable!(cpmodel, y)
        SeaPearl.addVariable!(cpmodel, z)
        SeaPearl.addObjective!(cpmodel,y)

        rng = MersenneTwister(10)

        variableselection = SeaPearl.RandomVariableSelection{false}()

        # This condition is there because of the way random are generated can change from one version to another
        if VERSION >= v"1.6.0"
            @test variableselection(cpmodel; rng=rng) == z
            @test variableselection(cpmodel; rng=rng) == z
            @test variableselection(cpmodel; rng=rng) == z
        end
    end
end


using Morsel
using Meddle
using JSON
using WebSockets
using REPL
using Gadfly
using Compose
using Color
import YAML

f = open("settings.yml")
settings = YAML.load(f)
close(f)

const dark_theme =
    Gadfly.Theme(color("steel blue"),      # default_color
          0.6mm,                    # default_point_size
          0.2mm,                    # line_width
          color("#000000"),         # panel_fill
          color("#4d87c7"),         # panel_stroke
          color("#333333"),         # grid_color
          color("#f0f0f0"),         # grid_color_focused
          0.2mm,                    # grid line width
          Gadfly.default_font_desc,        # minor_label_font
          9pt,                      # minor_label_font_size
          color("#dddddd"),         # minor_label_color
          Gadfly.default_font_desc,        # major_label_font
          11pt,                     # major_label_font_size
          color("#dddddd"),         # major_label_color
          Gadfly.default_font_desc,        # point_label_font
          8pt,                      # point_label_font_size
          color("#dddddd"),         # point_label_color
          0.0mm,                    # bar_spacing
          1mm,                      # boxplot_spacing
          0.3mm,                    # highlight_width
          Gadfly.default_highlight_color,  # highlight_color
          Gadfly.default_middle_color,     # middle_color
          1000,                     # label_placement_iterations
          10.0,                     # label_out_of_bounds_penalty
          0.5,                      # label_hidden_penalty
          0.2)                      # label_visibility_flip_pr

include("msg_types.jl")

app = Morsel.app()

import Base: repl_show

function repl_show(io::WebSocket,result,user_id)
	message = JSON.to_json([MSG_OUTPUT_EVAL_RESULT,user_id,sprint(repl_show,result)])
	for sock in clients
		write(sock,message)
	end
end
function repl_show(io::WebSocket,p::Union(Plot,Canvas),user_id)
	p.theme = dark_theme
	div_id = "plot_$(dec(uint16(rand(Uint64))))"
	out = IOBuffer()
	draw(D3(out,6inch,6inch; emit_on_finish=false),p)
	message = JSON.to_json([MSG_OUTPUT_HTML,user_id,"""

	"""]) #"
	for sock in clients
		isopen(sock) && write(sock,message)
	end
end

repl_channel, response_channel = RemoteRef(),RemoteRef()
REPL.start_repl_backend(repl_channel,response_channel)

clients = Array(WebSocket,0)

websocket(app,"/") do req, res
	socket = res.res
	push!(clients,socket)
	while true
        msg = read(socket)
        julia_msg = JSON.parse(bytestring(msg))
        msg_type = julia_msg[1]
        if msg_type == MSG_INPUT_EVAL
        	username = julia_msg[2]
        	user_id = julia_msg[3]
        	ast = Base.parse_input_line(julia_msg[4])
        	put(repl_channel, (ast,1))
        	@async begin
        		message = JSON.to_json([MSG_OUTPUT_EVAL_INPUT,user_id,username,julia_msg[4]])
        		for sock in clients
        			isopen(sock) && write(sock, message)
        		end
        		(val, bt) = take(response_channel)
        		repl_show(socket,val,user_id)
        	end
        end
    end
    splice!(clients,findfirst(clients,socket))
end

start(app, parseint(setting["worker_port"]))


#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*
#*--------------------------------------------------------------------
#* Models to simulate a power plant.
#*--------------------------------------------------------------------
#* Author: Argimiro R. Secchi
#* $Id: power_plant.mso 195 2007-03-07 20:30:12Z arge $
#*-------------------------------------------------------------------
# Declaracao de tipos
type SumidouroQ
	SumidouroQ()=begin
		new(
			Potencia(Dict{Symbol,Any}(
				:Symbol=>"_{in}",
				:PosX=>0,
				:PosY=>0.5
			)),
			[:Qin,]
		)
	end
	Qin::Potencia
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export SumidouroQ
function atributes(in::SumidouroQ,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=true
	fields[:Icon]="icon/sumidouroQ"
	fields[:Brief]="Sumidouro de calor"
	drive!(fields,_)
	return fields
end
SumidouroQ(_::Dict{Symbol,Any})=begin
	newModel=SumidouroQ()
	newModel.attributes=atributes(newModel,_)
	newModel
end


# Define plotting recipe
@recipe function f(s::SynthControlModel; kind = "overall")
    legend --> :topleft

    if kind == "weights"

        wp = sort(DataFrame(comp = s.treatment_panel.comp_labels,
                  weight = s.w), :weight, rev = true)

        @series begin
            seriestype := :bar
            label --> ""
            xguide --> "Donor"
            yguide --> "Weight"
            linecolor := "white"
            wp[wp.weight .> 0.05, :comp], wp[wp.weight .> 0.05, :weight]
        end

    elseif kind == "overall"

        @series begin
            label --> s.treatment_panel.treatment[1]
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), [s.treatment_panel.y₁₀; s.treatment_panel.y₁₁]
        end

        @series begin
            label --> "Control"
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), s.ŷ₁
        end

        @series begin
            label --> "Impact"
            seriestype --> :bar
            seriesalpha --> 0.5
            linecolor --> "white"
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var]))[findfirst(x -> x >= s.treatment_panel.treatment[2],
                                    sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var]))):end], s.τ̂
        end

        @series begin
            label --> "Intervention"
            seriestype := :vline
            linestyle := :dash
            seriescolor := "black"
            xguide := "Time"
            yguide := "Outcome"
            [s.treatment_panel.treatment[2]]
        end

    elseif kind == "diffplot"
        @series begin
            label --> s.treatment_panel.treatment[1]
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), [s.treatment_panel.y₁₀; s.treatment_panel.y₁₁] .- s.ŷ₁
        end

        @series begin
            label --> ""
            seriestype := :scatter
            seriescolor := 1
            markerstrokecolor := "white"
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), [s.treatment_panel.y₁₀; s.treatment_panel.y₁₁] .- s.ŷ₁
        end

        @series begin
            label --> "Intervention"
            seriestype := :vline
            linestyle := :dash
            seriescolor := "black"
            xguide := "Time"
            yguide := "Outcome"
            [s.treatment_panel.treatment[2]]
        end
    end

     nothing
end


"""
    worldclim(layer::Integer; resolution::Float64=10.0, left=nothing, right=nothing, bottom=nothing, top=nothing)

Download and prepare a WorldClim 2.1 bioclimatic variable, and returns it as an
`SimpleSDMPredictor`. Layers are called by their number, from 1 to 19. The list
of available layers is given in a table below.

The keywords are `resolution`, which must be a floating point value, and either
`2.5`, `5.0`, or `10.0`, as well as optionally `left`, `right`, `bottom`, and
`top`, which will allow to only load parts of the data.

Internally, this function will download the main zip file for the required
resolution from the WordlClim website, extract it, and parse the required
layers.

It is recommended to *keep* the content of the `path` folder, as it will
eliminate the need to download and/or extract the tiff files. For example,
calling `wordlclim(1:19)` will download and extract everything, and future calls
will be much faster.

| Variable | Description                                                |
| ------   | ------                                                     |
| 1        | Annual Mean Temperature                                    |
| 2        | Mean Diurnal Range (Mean of monthly (max temp - min temp)) |
| 3        | Isothermality (BIO2/BIO7) (* 100)                          |
| 4        | Temperature Seasonality (standard deviation *100)          |
| 5        | Max Temperature of Warmest Month                           |
| 6        | Min Temperature of Coldest Month                           |
| 7        | Temperature Annual Range (BIO5-BIO6)                       |
| 8        | Mean Temperature of Wettest Quarter                        |
| 9        | Mean Temperature of Driest Quarter                         |
| 10       | Mean Temperature of Warmest Quarter                        |
| 11       | Mean Temperature of Coldest Quarter                        |
| 12       | Annual Precipitation                                       |
| 13       | Precipitation of Wettest Month                             |
| 14       | Precipitation of Driest Month                              |
| 15       | Precipitation Seasonality (Coefficient of Variation)       |
| 16       | Precipitation of Wettest Quarter                           |
| 17       | Precipitation of Driest Quarter                            |
| 18       | Precipitation of Warmest Quarter                           |
| 19       | Precipitation of Coldest Quarter                           |

Original data: https://www.worldclim.org/data/worldclim21.html
"""
function worldclim(layer::Integer; resolution::Float64=10.0, left=nothing, right=nothing, bottom=nothing, top=nothing)
    return raster(SimpleSDMPredictor, WorldClim(resolution), layer=layer, left=left, right=right, bottom=bottom, top=top)
end

worldclim(layers::Vector{T}; args...) where {T <: Integer} = [worldclim(l; args...) for l in layers]
worldclim(layers::UnitRange{T}; args...) where {T <: Integer} = [worldclim(l; args...) for l in layers]


using GeckoScreen
using Test
using DataFrames
using CSV
using Query

testcsv(path::AbstractString; kwargs...) = CSV.read(path, DataFrame; kwargs...)

@testset "GeckoScreen.jl" begin
    # Write your tests here.

    gmap = readGeCKOlibrary()
    first(gmap,6)
    #@test size(readGeCKOlibrary, 1) == 34

end


module System

using ComponentArrays
import Flight.Plotting: plots

export assemble_x0, init_x0, init_y0, init_u0, init_d0, f_cont!, f_disc!
export AbstractComponent, AbstractSystem, System

no_extend_error(f::Function, ::Type{S}) where {S} = error(
    "Function $f not implemented for type $S or incorrect call signature")
no_extend_warning(f::Function, ::Type{S}) where {S} = println(
    "Warning: Function $f not implemented for type $S or incorrect call signature")

#anything around which we can build a System
abstract type AbstractComponent end #anything that can go in a System

#init_x0 must return either nothing, an AbstractVector{<:Real} or a NT with a
#Union{AbstractVector{<:Real},Nothing}
init_x0(::C) where {C<:AbstractComponent} = nothing
init_y0(::C) where {C<:AbstractComponent} = nothing
init_u0(::C) where {C<:AbstractComponent} = nothing #sytems are not required to have control inputs
init_d0(::C) where {C<:AbstractComponent} = nothing #systems are not required to have discrete states

abstract type AbstractSystem{C<:AbstractComponent} end

#need the C type parameter for dispatch, the rest for type stability
mutable struct System{C, X <: Union{Nothing, AbstractVector{<:Real}},
                    Y, U, D, P, S} <: AbstractSystem{C}
    ẋ::X #continuous state vector derivative
    x::X #continuous state vector (to be used as a buffer for f_cont! evaluation)
    y::Y #output state
    u::U #control inputs
    d::D #discrete state
    t::Base.RefValue{Float64} #this allows propagation of t updates down the subsystem hierarchy
    params::P
    subsystems::S
end

f_cont!(::S, args...) where {S<:AbstractSystem} = no_extend_error(f_cont!, S)
(f_disc!(::S, args...)::Bool) where {S<:AbstractSystem} = no_extend_error(f_disc!, S)

#f_disc! is free to modify a Hybrid system's discrete state, control inputs and
#continuous state. if it modifies the latter, it must return true, false
#otherwise. it is dangerous to provide a default fallback for f_disc!, because
#if the intended f_disc! implementation for the System has the wrong interface,
#the dispatch will revert to the fallback, which may not be obvious at all. it
#is safer to force each concrete System that does not require an actual f_disc!
#to implement a trivial f_disc! that returns false

#when the System constructor for a certain Component is passed no
#parameters for dx and x, it calls the init_x0 method for that Component, which
#may return

assemble_x0(c::AbstractComponent) = assemble_x0(init_x0(c))

assemble_x0(::Nothing) = (s = nothing, ss = nothing)

assemble_x0(x::AbstractVector{<:Real}) = (s = x, ss = nothing)

function assemble_x0(x::ComponentVector)

    x_ss = Dict{Symbol,Any}()

    for id in keys(x)
        x_ss[id] = view(x, id)
    end

    return (s = x, ss = NamedTuple{Tuple(keys(x_ss))}(values(x_ss)))

end

# function assemble_x0(nt::NamedTuple{L,X}
#     ) where {L, X <: NTuple{N, Union{AbstractVector{<:Real}, Nothing}} where {N}}
function assemble_x0(nt::NamedTuple)
    #assembles the System's state vector from the non-empty subsystems' state
    #vectors. returns the System's state vector and the views to be assigned to
    #the subsystems
    x_ss = Dict{Symbol,Any}(pairs(nt))

    #para cada entry del nt, llamamos a assemble_x0, que nos devolvera un x_s y
    #un x_ss

    x_s_blocks = filter(x_ss) do item
        isa(item.second, AbstractVector{<:Real})
    end
    x_s = (!isempty(x_s_blocks) ? ComponentVector{Float64}(x_s_blocks) : nothing)

    for id in keys(x_ss)
        x_ss[id] = (id in keys(x_s_blocks) ? view(x_s, id) : x_ss[id])
    end

    return (s = x_s, ss = NamedTuple{Tuple(keys(x_ss))}(values(x_ss)))

end

function plots(::AbstractVector{<:Real}, ::AbstractVector{T}) where {T}
    no_extend_warning(plots, T) #nothing to plot by default, warn about it
end

end


include("instructions.jl")

function manager(instruction::T, model::GeneModel) where {T <: IndividualInstruction}
    if hasfield(T, :tag)
        tag = instruction.tag
        for id in collect(keys(model.individuals))
            if tag = nothing
                instruction.name(model, instruction, id)
            end
        end
    end
end

instruction = RandomGenome(2,"binary",1.0, nothing, nothing, false, generate_genotype)

individuals = Dict(1,BasicIndividual(zeros(2,2)))

model = GeneModel(individuals)

manager(instruction, model)


"""
    _make_jvp_call(fdm, f, y, xs, ẋs, ignores)

Call `FiniteDifferences.jvp`, with the option to ignore certain `xs`.

# Arguments
- `fdm::FiniteDifferenceMethod`: How to numerically differentiate `f`.
- `f`: The function to differentiate.
- `y`: The primal output `y=f(xs...)` or at least something of the right type
- `xs`: Inputs to `f`, such that `y = f(xs...)`.
- `ẋs`: The directional derivatives of `xs` w.r.t. some real number `t`.
- `ignores`: Collection of `Bool`s, the same length as `xs` and `ẋs`.
   If `ignores[i] === true`, then `ẋs[i]` is ignored for derivative estimation.

# Returns
- `Ω̇`: Derivative of output w.r.t. `t` estimated by finite differencing.
"""
function _make_jvp_call(fdm, f, y, xs, ẋs, ignores)
    f2 = _wrap_function(f, xs, ignores)

    ignores = collect(ignores)
    all(ignores) && return ntuple(_ -> NoTangent(), length(xs))
    sigargs = zip(xs[.!ignores], ẋs[.!ignores])
    return ProjectTo(y)(jvp(fdm, f2, sigargs...))
end

"""
    _make_j′vp_call(fdm, f, ȳ, xs, ignores) -> Tuple

Call `FiniteDifferences.j′vp`, with the option to ignore certain `xs`.

# Arguments
- `fdm::FiniteDifferenceMethod`: How to numerically differentiate `f`.
- `f`: The function to differentiate.
- `ȳ`: The adjoint w.r.t. output of `f`.
- `xs`: Inputs to `f`, such that `y = f(xs...)`.
- `ignores`: Collection of `Bool`s, the same length as `xs`.
  If `ignores[i] === true`, then `xs[i]` is ignored; `∂xs[i] === NoTangent()`.

# Returns
- `∂xs::Tuple`: Derivatives estimated by finite differencing.
"""
function _make_j′vp_call(fdm, f, ȳ, xs, ignores)
    f2 = _wrap_function(f, xs, ignores)

    ignores = collect(ignores)
    args = Any[NoTangent() for _ in 1:length(xs)]
    all(ignores) && return (args...,)
    sigargs = xs[.!ignores]
    arginds = (1:length(xs))[.!ignores]
    fd = j′vp(fdm, f2, ȳ, sigargs...)
    @assert length(fd) == length(arginds)

    for (dx, ind) in zip(fd, arginds)
        args[ind] = ProjectTo(xs[ind])(dx)
    end
    return (args...,)
end

"""
    _wrap_function(f, xs, ignores)

Return a new version of `f`, `fnew`, that ignores some of the arguments `xs`.

# Arguments
- `f`: The function to be wrapped.
- `xs`: Inputs to `f`, such that `y = f(xs...)`.
- `ignores`: Collection of `Bool`s, the same length as `xs`.
  If `ignores[i] === true`, then `xs[i]` is ignored; `∂xs[i] === NoTangent()`.
"""
function _wrap_function(f, xs, ignores)
    function fnew(sigargs...)
        callargs = Any[]
        j = 1

        for (i, (x, ignore)) in enumerate(zip(xs, ignores))
            if ignore
                push!(callargs, x)
            else
                push!(callargs, sigargs[j])
                j += 1
            end
        end
        @assert j == length(sigargs) + 1
        @assert length(callargs) == length(xs)
        return f(callargs...)
    end
    return fnew
end


using ArgParse
using PowerModelsSecurityConstrained

function scopf_main(parsed_args)
    ini_file = parsed_args["file"]
    scenario = parsed_args["scenario"]
    scoring_method = parsed_args["scoring-method"]
    time_limit_sol1 = parsed_args["time-limit"]
    time_limit_total = parsed_args["hard-time-limit"]
    distribute = parsed_args["distribute"]
    skip_solution2 = parsed_args["skip-solution2"]
    remove_solutions = parsed_args["remove-solutions"]
    gurobi = parsed_args["gurobi"]

    files, scenario_id = find_goc_files(ini_file, scenario_id=scenario)

    ini_dir = dirname(ini_file)
    output_dir = joinpath(ini_dir, scenario_id)

    info(LOGGER, "output dir: $(output_dir)")

    time_sol1_start = time()
    compute_solution1(files["con"], files["inl"], files["raw"], files["rop"], time_limit_sol1-60, scoring_method, "network name"; output_dir=output_dir, scenario_id=scenario_id, gurobi=gurobi)
    time_sol1 = time() - time_sol1_start

    if !skip_solution2
        compute_solution2(files["con"], files["inl"], files["raw"], files["rop"], trunc(Int, time_limit_total-time_sol1-60), scoring_method, "network name"; output_dir=output_dir, scenario_id=scenario_id)
    end

    if remove_solutions
        remove_solution_files(output_dir=output_dir)
        remove_detail_file(output_dir=output_dir)
    end
end

function parse_scopf_commandline()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--file", "-f"
            help = "the data initiation file (.ini)"
            required = true
        "--scenario", "-s"
            help = "the scenario to run (directory)"
            default = ""
        "--scoring-method", "-m"
            help = "the objective function (1,2,3,4)"
            arg_type = Int
            default = 2
        "--time-limit", "-t"
            help = "solution 1 runtime limit (seconds)"
            arg_type = Int
            default = 600000 #2700
        "--hard-time-limit"
            help = "solution 1 and 2 runtime limit (seconds)"
            arg_type = Int
            default = 600000 #34200
        "--distribute", "-d"
            help = "run on multiple processes"
            action = :store_true
        "--skip-solution2"
            help = "skip computation of solution2 file"
            action = :store_true
        "--remove-solutions"
            help = "delete the solution files after competition"
            action = :store_true
        "--gurobi", "-g"
            help = "use Gurobi for solving lp, qp and mip problems"
            action = :store_true
            #default = false
    end

    return parse_args(s)
end


export MeshPartitionTree
export tree, leaf_meshes
# A data structure to hold a hierarchical partition of a mesh.
# Since the mesh is partitioned, we only need to store the leaves of the partition
# tree to reconstruct the mesh.
#
# Each node has the name of the partition and an ID.
# If the node is an internal node, it has id = 0.
# If the node is a leaf node, it has a non-zero id, corresponding to the index by
# which the mesh may be found in the MeshPartitionTree type's leaf_meshes.
struct MeshPartitionTree{M <: AbstractMesh}
    partition_tree::Tree{Tuple{Int64,String}}
    leaf_meshes::Vector{M}
end

tree(mpt::MeshPartitionTree) = mpt.partition_tree
leaf_meshes(mpt::MeshPartitionTree) = mpt.leaf_meshes

# Partitions a mesh based upon the names of its groups
# For a hierarchical partition, all partitions in the hierarchy must contain face sets
# of the form "<name>_L<N>" where name is a string, and N is an integer.
# N is the level of the node in the tree: 1, 2, 3, etc...
# Example: "Grid_L1_triangle", or "Partition_L3"
function MeshPartitionTree(mesh::AbstractMesh; by::String="MPACT")
    @info "Partitioning mesh: "*name(mesh)
    # Extract the names of all face sets that contain 'by' (the variable)
    partition_names = _get_partition_names(mesh, by)
    # Create a tree to store the partition hierarchy.
    root = _create_mesh_partition_tree(mesh, partition_names)
    # Construct the leaf meshes
    leaf_meshes = _create_leaf_meshes(mesh, root)
    return MeshPartitionTree(root, leaf_meshes)
end

function _create_leaf_meshes(mesh::AbstractMesh, root::Tree)
    leaf_nodes = sort!(leaves(root), by=x->x.data)
    leaf_meshes = Vector{typeof(mesh)}(undef, length(leaf_nodes))
    leaf_ctr = 1
    node_children = children(root)
    if !isnothing(node_children)
        for child ∈ sort(node_children, by=x->x.data)
            child_mesh = submesh(mesh, data(child)[2])
            leaf_ctr = _create_leaf_meshes!(child_mesh, child, leaf_meshes, leaf_ctr)
        end
    else
        # Remove any Lattice, Module, or Coarse_Cell groups, since this info should now
        # be encoded in the tree
        mpact_groups = filter(x->isa_MPACT_partition_name(x), keys(groups(mesh)))
        for grp in mpact_groups
            pop!(groups(mesh), grp)
        end
        name = data(root)[2]
        root.data = (leaf_ctr, name)
        leaf_meshes[leaf_ctr] = mesh
    end
    return leaf_meshes
end

function _create_leaf_meshes!(mesh::AbstractMesh,
                              node::Tree,
                              leaf_meshes,
                              leaf_ctr::Int64=1
                             )
    node_children = children(node)
    if !isnothing(node_children)
        for child ∈ sort(node_children, by=x->x.data)
            child_mesh = submesh(mesh, data(child)[2])
            leaf_ctr = _create_leaf_meshes!(child_mesh, child, leaf_meshes, leaf_ctr)
        end
    else
        # Remove any Lattice, Module, or Coarse_Cell groups, since this info should now
        # be encoded in the tree
        mpact_groups = filter(x->isa_MPACT_partition_name(x), keys(groups(mesh)))
        for grp in mpact_groups
            pop!(mesh.groups, grp)
        end
        leaf_meshes[leaf_ctr] = mesh
        name = data(node)[2]
        node.data = (leaf_ctr, name)
        leaf_ctr += 1
    end
    return leaf_ctr
end

# Create a tree to store grid relationships.
function _create_mesh_partition_tree(mesh::AbstractMesh, partition_names::Vector{String})
    root = Tree((0, name(mesh)))
    parents = [root]
    tree_type = typeof(root)
    next_parents = tree_type[]
    remaining_names = copy(partition_names)
    # Extract the sets that are not a subset of any other set.
    this_level = Int64[]
    mesh_groups = groups(mesh)
    while length(remaining_names) > 0
        for i in eachindex(remaining_names)
            name_i = remaining_names[i]
            i_isa_lattice = startswith(name_i, "Lattice")
            i_isa_module = startswith(name_i, "Module")
            isa_subset = false
            for j in eachindex(remaining_names)
                if i === j
                    continue
                end
                name_j = remaining_names[j]
                j_isa_coarsecell = startswith(name_j, "Coarse")
                if j_isa_coarsecell && (i_isa_module || i_isa_lattice)
                    continue
                end
                j_isa_module = startswith(name_j, "Module")
                if j_isa_module && i_isa_lattice
                    continue
                end
                if mesh_groups[name_i] ⊆ mesh_groups[name_j]
                    isa_subset = true
                    break
                end
            end
            if !isa_subset
                push!(this_level, i)
            end
        end
        # Add the groups that are not a subset to the tree
        # and next_parents
        for id in this_level
            name = remaining_names[id]
            for parent in parents
                if isroot(parent)
                    node = Tree((0, name), parent)
                    push!(next_parents, node)
                else
                    parent_name = data(parent)[2]
                    if mesh_groups[name] ⊆ mesh_groups[parent_name]
                        node = Tree((0, name), parent)
                        push!(next_parents, node)
                        break
                    end
                end
            end
        end
        parents = next_parents
        next_parents = tree_type[]
        # Remove the groups added to the tree from the remaining names
        deleteat!(remaining_names, this_level)
        this_level = Int64[]
    end
    return root
end

function isa_MPACT_partition_name(x::String)
    return startswith(x, "Coarse_Cell_(") ||
           startswith(x, "Module_("     ) ||
           startswith(x, "Lattice_("    )
end
# Extract partition names
function _get_partition_names(mesh::AbstractMesh, by::String)
    # If by === MPACT, then we will partition by Lattice, Module, and Coarse Cell
    # Otherwise, partition by the string contained in by
    partition_names = collect(keys(groups(mesh)))
    if length(partition_names) === 0
        error("The mesh does not have any groups.")
    end
    if by === "MPACT"
        filter!(x->isa_MPACT_partition_name(x), partition_names)
        if length(partition_names) === 0
            error("The mesh does not have any MPACT grid hierarchy groups.")
        end
    else
        filter!(x->occursin(by, x), partition_names)
        if length(partition_names) === 0
            error("The mesh does not have any groups containing '", by, "'.")
        end
    end
    return sort!(partition_names)
end

function Base.show(io::IO, mpt::MeshPartitionTree{M}) where {M}
    println("MeshPartitionTree{",M,"}")
    println(tree(mpt))
end


using Iterators
using GeometryTypes, GPUArrays
using Visualize: STDCanvas
using Visualize, GeometryTypes, ModernGL
using Visualize.GLRasterization
using Visualize: Debugging
using Visualize.GLRasterization: swapbuffers!, destroy!, show!
using Visualize: View, Projection, ProjectionView

resolution = (1024, 900)
window = GLFWWindow(Area => (800, 600), Debugging => true)

for event in Visualize.NativeWindowEvents
    add!(window, event)
end
add!(window, Mouse.Drag)
window[Visualize.Open] = true
show!(window)

cam = PerspectiveCamera(
    TranslationSpeed => 1f0,
    LookAt => Vec3f0(0),
    EyePosition => Vec3f0(6, 6, 8),
    Rotation => Vec3f0(0),
    Area => window[Area],
    RotationSpeed => 0.1f0
)

add!(cam, Translation, window,
    Mouse.right, Set([Keyboard.left_control])
)
add!(cam, Rotation, window,
    Mouse.left, Set([Keyboard.left_control])
)
add!(cam, Projection, View)
add!(cam, EyePosition, LookAt)

canvas = STDCanvas(
    cam[View],
    cam[Projection],
    cam[ProjectionView],
    Vec2f0(resolution),
    cam[EyePosition],
)

uniforms = Uniforms(
    eye(Mat4f0),
    eye(Mat4f0),
    0f0,
    0f0,
    x-> x > 1.0f0,
    (x, u)-> Vec4f0(x, 1f0, 1f0, 1f0),
    mip
)
GLBackend.init(window[Window])
intensities = GLBackend.GLSampler(rand(Float32, 32, 32, 32));
mesh = GLUVWMesh(AABB(Vec3f0(0), Vec3f0(1)))
vertices = map(x-> UVWVertex(x[1], x[2]), zip(mesh.vertices, mesh.texturecoordinates))

vbo = VertexArray(view(vertices, mesh.faces));
program_args = (UniformBuffer(canvas), UniformBuffer(uniforms), GPUArrays.buffer(intensities));
draw_vol = GLRasterizer(
    vbo,
    program_args,
    vert_volume, frag_volume
)

glEnable(GL_DEPTH_TEST)
glDepthMask(GL_TRUE)
glDepthFunc(GL_LEQUAL)
glEnable(GL_CULL_FACE)
glCullFace(GL_FRONT)
GLAbstraction.enabletransparency()
glClearColor(1,1,1,1)

@async begin
    while isopen(window)
        GLFW.PollEvents()
        glViewport(0, 0, widths(window[Area])...)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        program_args[1][1] = STDCanvas(
            cam[View],
            cam[Projection],
            cam[ProjectionView],
            Vec2f0(resolution),
            cam[EyePosition],
        )
        draw_vol(vbo, program_args)
        swapbuffers!(window)
        sleep(0.001)
        yield()
    end
    destroy!(window)
end


################################################################################
### CUSTOM DATA STRUCTURES
#=
Custom data structures are defined here
(structures are immutable by default; if modification is required, use "mutable struct")
=#

################################################################################
### PRIMARY INDEX SETS
struct ArcDataStruct
    ID::Array
    node_s::Dict
    node_r::Dict
    cap::Dict
    len::Dict
end

struct NodeDataStruct
    ID::Array
    alias::Dict
    lon::Dict
    lat::Dict
end

struct ProductDataStruct
    ID::Array
    alias::Dict
    transport_cost::Dict
end

struct DemandDataStruct
    ID::Array
    node::Dict
    prod::Dict
    bid::Dict
    cap::Dict
end

struct SupplyDataStruct
    ID::Array
    node::Dict
    prod::Dict
    bid::Dict
    cap::Dict
end

struct TechDataStruct
    ID::Array
    Outputs::Dict
    Inputs::Dict
    OutputYields::Dict
    InputYields::Dict
    InputRef::Dict
    bid::Dict
    cap::Dict
    alias::Dict
end

struct TechsiteDataStruct
    ID::Array
    node::Dict
    tech::Dict
end

################################################################################
### SECONDARY SETS (SUBSETS, OR INTERSECTIONS OF TWO OR MORE PRIMARY INDICES)
struct SetStruct
    Ain::Dict
    Aout::Dict
    TPQ::Array
    NQT::Dict
    NPt::Dict
    NQt::Dict
end

################################################################################
### PARAMETERS
struct ParStruct
    dMAX::Dict
    sMAX::Dict
    α::Dict
    ξgenMAX::Dict
    ξconMAX::Dict
    fMAX::Dict
end

################################################################################
### MODEL OUTPUTS
struct OutputStruct
    z::Float64
    si::JuMP.Containers.DenseAxisArray
    dj::JuMP.Containers.DenseAxisArray
    snp::JuMP.Containers.DenseAxisArray
    dnp::JuMP.Containers.DenseAxisArray
    f::JuMP.Containers.DenseAxisArray
    ξcon::JuMP.Containers.DenseAxisArray
    ξgen::JuMP.Containers.DenseAxisArray
    πNP::JuMP.Containers.DenseAxisArray
    πA::Dict
    πT::Dict
    Φd::Dict
    Φs::Dict
    Φf::Dict
    Φξ::Dict
end

struct StatStruct
    NumVars::Float64
    TotalIneqCons::Float64
    TotalEqCons::Float64
    NumVarBounds::Float64
    ModelIneqCons::Float64
    ModelEqCons::Float64
end


include("../required.jl")

## Get layers
# Define coordinates range
coords = (left=-145.0, right=-50.0, bottom=20.0, top=75.0)

# Test loading variables
lc_vars = SimpleSDMPredictor(Copernicus, LandCover, 1:10; resolution=10.0, coords...)

# Load worldclim variables to compare
wc_vars = SimpleSDMPredictor(WorldClim, BioClim, 1:19; resolution=10.0, coords...);

## Temporary dimensions fix
# Check which line is extra
wc_test = SimpleSDMPredictor(WorldClim, BioClim, 1)[coords]
wc_vars[1].grid
wc_test.grid # last line seems extra
isequal(wc_test.grid[1:(end - 1), :], wc_vars[1].grid) # true

# Test clip fix
wc_clip = wc_test[top=coords.top - stride(wc_test, 2)]
isequal(wc_clip.grid, wc_vars[1].grid) # true
wc_vars[1].top
wc_clip.top # not exactly the same
isapprox(wc_vars[1].top, wc_clip.top) # approximately the same

# Test clip from world layer instead
wc_test2 = SimpleSDMPredictor(WorldClim, BioClim, 1)[coords][top=coords.top -
                                                                 stride(wc_test, 2)] # same limits
wc_test2.top # not helping

# Fix landcover layers
lc_vars = [l[top=coords.top - stride(l, 2)] for l in lc_vars]
isequal(latitudes(wc_vars[1]), latitudes(lc_vars[1])) # false
isapprox(latitudes(wc_vars[1]), latitudes(lc_vars[1])) # true so should be ok

## Export to CSV
# Combine environmental data
env_vars = [wc_vars; lc_vars]
# Create env matrix
env_mat = mapreduce(x -> vec(x.grid), hcat, env_vars)
replace!(x -> isnothing(x) ? NaN : x, env_mat)
# Create env dataframe
env_df = DataFrame(env_mat, :auto)
rename!(env_df, vcat(Symbol.("wc", 1:size(wc_vars, 1)), Symbol.("lc", 1:size(lc_vars, 1))))
insertcols!(env_df, 1, :site => 1:nrow(env_df))

# Get sites latitudes
lats = repeat(collect(latitudes(wc_vars[1])); outer=size(wc_vars[1].grid, 2))
# Get sites longitudes
lons = repeat(collect(longitudes(wc_vars[1])); inner=size(wc_vars[1].grid, 1))
# Create spa matrix
spa_mat = [lats lons]
# Create spa dataframe
spa_df = DataFrame(; site=eachindex(lats), lat=lats, lon=lons)

## Check difference with previously saved data

# Copy data
run(
    `cp ./data/proc/distributions_env_full.csv ./data/proc/distributions_env_full_backup.csv`,
)
run(
    `cp ./data/proc/distributions_spa_full.csv ./data/proc/distributions_spa_full_backup.csv`,
)
run(`cp ./data/proc/distributions_spa_qc.csv ./data/proc/distributions_spa_qc_backup.csv`)

# Export new files
CSV.write(joinpath("data", "proc", "distributions_env_full.csv"), env_df; delim="\t")
CSV.write(joinpath("data", "proc", "distributions_spa_full.csv"), spa_df; delim="\t")
# Git diff says spa files are different
# Sizes are different too
# Looking at the diff, it seems to be a float approximation difference, which should be fine

# Let's make sure it's just that
spa_df_backup = CSV.read(
    joinpath("data", "proc", "distributions_spa_full_backup.csv"), DataFrame
)
spa_df_backup
spa_df
isequal(spa_df_backup, spa_df) # false
isapprox(spa_df_backup, spa_df) # true
isequal(spa_df_backup.site, spa_df.site) # true
isapprox(spa_df_backup.lat, spa_df.lat) # true
isapprox(spa_df_backup.lon, spa_df.lon) # true
# Yep, should just be approximation error, so its fine

# Make sure env files are equal too, even if diff says they are
env_df_backup = CSV.read(
    joinpath("data", "proc", "distributions_env_full_backup.csv"), DataFrame
)
env_df_backup
env_df
isequal(env_df_backup, env_df) # false
isapprox(env_df_backup, env_df) # false ?!
isapprox(NaN, NaN) # oh probably just because of NaN
# whatever let's trust Git

## Export QC sites coordinates (for smaller scale analyses)
coords_qc = (left=-80.0, right=-55.0, bottom=45.0, top=63.0)
# Get site indices
spa_qc = filter(
    x ->
        (coords_qc.left - stride(wc_vars[1]; dims=2) <= x.lon < coords_qc.right) &&
            (coords_qc.bottom - stride(wc_vars[1]; dims=1) <= x.lat < coords_qc.top),
    spa_df,
)

# Test dimensions for Quebec coordinates
wc_vars_qc = SimpleSDMPredictor(WorldClim, BioClim, 1; coords_qc...)
wc_test_qc = SimpleSDMPredictor(WorldClim, BioClim, 1)[coords_qc]
isequal(size(wc_vars_qc), size(wc_test_qc)) # false
isapprox(longitudes(wc_vars_qc), unique(spa_qc.lon)) # true
isapprox(longitudes(wc_test_qc), unique(spa_qc.lon)) # false
isequal(length(wc_vars_qc.grid), nrow(spa_qc)) # equal
# So dimensions in the exported file should be ok
CSV.write(joinpath("data", "proc", "distributions_spa_qc.csv"), spa_qc; delim="\t")
# Git diff says there's difference
spa_qc_backup = CSV.read(
    joinpath("data", "proc", "distributions_spa_qc_backup.csv"), DataFrame
)
isapprox(spa_qc_backup, spa_qc) # true
# Again just an approximation error


using Raylib_jll

struct RayColor
    r::Cuchar
    g::Cuchar
    b::Cuchar
    a::Cuchar
end

const RAYWHITE = RayColor(245, 245, 245, 255)
const LIGHTGRAY = RayColor(200, 200, 200, 255)
const DARKGRAY = RayColor(80, 80, 80, 255)
const MAROON = RayColor(190, 33, 55, 255)
const RED = RayColor(230, 41, 55, 255)
const BLUE = RayColor(0, 121, 241, 255 )
const LIME = RayColor(0, 158, 47, 255)
const GOLD = RayColor(255, 203, 0, 255)
const SKYBLUE = RayColor(102, 191, 255, 255)
const BLACK = RayColor(0, 0, 0, 255)

struct RayVector3
    x::Cfloat
    y::Cfloat
    z::Cfloat
end

struct RayVector2
    x::Cfloat
    y::Cfloat
end

struct RayCamera3D
    position::RayVector3
    target::RayVector3
    up::RayVector3
    fovy::Cfloat
    projection::Cint
end

function main()
    screenWidth = 800
    screenHeight = 450

    @ccall libraylib.InitWindow(screenWidth::Cint, screenHeight::Cint,
                             "raylib [core] example - 3d camera first person"::Cstring)::Cvoid

    camera = Ref(RayCamera3D(
        RayVector3(4, 2, 4),
        RayVector3(0, 1.8, 0),
        RayVector3(0, 1, 0),
        60,
        0,
    ))

    heights = Vector{Cfloat}(undef, 20)
    positions = Vector{RayVector3}(undef, 20)
    colors = Vector{RayColor}(undef, 20)

    for i = 1:20
        heights[i] = @ccall libraylib.GetRandomValue(1::Cint, 12::Cint)::Cint
        positions[i] = RayVector3(
            @ccall(libraylib.GetRandomValue((-15)::Cint, 15::Cint)::Cint),
            heights[i] / 2,
            @ccall(libraylib.GetRandomValue((-15)::Cint, 15::Cint)::Cint),
        )
        colors[i] = RayColor(
            @ccall( libraylib.GetRandomValue(20::Cint, 255::Cint)::Cint),
            @ccall( libraylib.GetRandomValue(10::Cint, 55::Cint)::Cint),
            30,
            255,
        )
    end

    @ccall libraylib.SetCameraMode(camera[]::RayCamera3D, 3::Cint)::Cvoid
    @ccall libraylib.SetTargetFPS(60::Cint)::Cvoid

    while iszero(@ccall libraylib.WindowShouldClose()::Cint)
        @ccall libraylib.UpdateCamera(camera::Ref{RayCamera3D})::Cvoid

        @ccall libraylib.BeginDrawing()::Cvoid
        @ccall libraylib.ClearBackground(RAYWHITE::RayColor)::Cvoid

        @ccall libraylib.BeginMode3D(camera[]::RayCamera3D)::Cvoid

        @ccall libraylib.DrawPlane(
            RayVector3(0,0,0)::RayVector3,
            RayVector2(32, 32)::RayVector2,
            LIGHTGRAY::RayColor)::Cvoid

        @ccall libraylib.DrawCube(
            RayVector3(-16, 2.5, 0)::RayVector3, 1f0::Cfloat, 5f0::Cfloat, 32f0::Cfloat,
            BLUE::RayColor)::Cvoid
        @ccall libraylib.DrawCube(
            RayVector3(16, 2.5, 0)::RayVector3, 1f0::Cfloat, 5f0::Cfloat, 32f0::Cfloat,
            LIME::RayColor)::Cvoid
        @ccall libraylib.DrawCube(
            RayVector3(0, 2.5, 16)::RayVector3, 32f0::Cfloat, 5f0::Cfloat, 1f0::Cfloat,
            GOLD::RayColor)::Cvoid

        for i = 1:20
            @ccall libraylib.DrawCube(
                positions[i]::RayVector3, 2f0::Cfloat, heights[i]::Cfloat, 2f0::Cfloat,
                colors[i]::RayColor)::Cvoid
            @ccall libraylib.DrawCubeWires(positions[i]::RayVector3, 2f0::Cfloat, heights[i]::Cfloat, 2f0::Cfloat,
                                        MAROON::RayColor)::Cvoid

        end

        @ccall libraylib.EndMode3D()::Cvoid

        @ccall libraylib.DrawRectangle(10::Cint, 10::Cint, 220::Cint, 70::Cint,
                                    @ccall( libraylib.Fade(SKYBLUE::RayColor, 0.5::Cfloat)::RayColor )::RayColor)::Cvoid
        @ccall libraylib.DrawRectangleLines(10::Cint, 10::Cint, 220::Cint, 70::Cint, BLUE::RayColor)::Cvoid

        @ccall libraylib.DrawText("First person camera default controls:"::Cstring,
                               20::Cint, 20::Cint, 10::Cint, BLACK::RayColor)::Cvoid
        @ccall libraylib.DrawText("- Move with keys: W, A, S, D"::Cstring,
                               40::Cint, 40::Cint, 10::Cint, DARKGRAY::RayColor)::Cvoid
        @ccall libraylib.DrawText("- Mouse move to look around"::Cstring,
                               40::Cint, 60::Cint, 10::Cint, DARKGRAY::RayColor)::Cvoid

        @ccall libraylib.EndDrawing()::Cvoid
    end

    @ccall libraylib.CloseWindow()::Cvoid

end


export Population, Chromosome, imChromosome, Parameters, indiv_result_type, pop_result_type, elt_type
export CompositionalCircuit, CompGate, CompCircuit, Need, CC, Cc, Cg, Need

using Distributions
using DataFrames
using Random
using QuadGK
using CSV
using Statistics
using Distributed
using Printf
import Random.seed!
export MyInt, DIST_TYPE, Population, PopVect, IPopulation, IntRange, Circuit

if !@isdefined(MyInt)  # MyInt should be defined in CGP.jl
  const MyInt = UInt16
end
if !@isdefined(MyFunc)  # MyFunc should be defined in CGP.jl
  const MyFunc = UInt128
end
const elt_type = Union{Int64,MyInt,String,Float64}
const DIST_TYPE = Dict{Any,Float64}
#const IDIST_TYPE = Dict{Int64,Float64}
#const SDIST_TYPE = Dict{String,Float64}
const Population = Union{Vector{Int64},Vector{MyInt},Vector{String},Vector{Float64}}
const PopVect = Union{Vector{Vector{Int64}},Vector{Vector{MyInt}},Vector{Vector{String}},Vector{Vector{Float64}}}
#const IPopulation = Array{Int64,1}
#const MIPopulation = Array{MyInt,1}
#const SPopulation = Array{String,1}
const FPopulation = Array{Float64,1}
#const CPopulation = Vector{Chromosome}   # population of chromosomes
const IntRange = Union{Integer, Bool, AbstractRange{Int64}, AbstractRange{Bool}}
const FloatRange = Union{Float64,AbstractRange{Float64}}
const Goal =  Vector{MyInt}
const GoalList = Vector{Goal}
#Ones = 0x0f
Ones = 0xffff
#Ones = 0xffffffff

mutable struct Parameters
    mu::Integer
    lambda::Integer
    mutrate::Real
    targetfitness::Real
    numinputs::Integer
    numoutputs::Integer
    nodearity::Integer
    numinteriors::Integer   # Number of gates
    numlevelsback::Integer  # For LinCircuits, number of registers
end
abstract type Node end
mutable struct Func
    func::Function
    arity::Integer
    name::AbstractString
end
mutable struct InputNode <: Node
    index::Integer
    active::Bool
    cache::MyInt
end
mutable struct InteriorNode <: Node
    func::Func
    inputs::Vector{Integer}
    active::Bool
    cache::MyInt
end
mutable struct OutputNode <: Node
    input::Integer
end
PredType = Int64
mutable struct Chromosome
    params::Parameters
    inputs::Vector{InputNode}
    interiors::Vector{InteriorNode}
    outputs::Vector{OutputNode}
    fitness::Float64
    robustness::Union{Float64,PredType}
end
mutable struct LinCircuit
    # Here, numinteriors represents number of gates, numlevelsback represents number of computational registers
    # the First numoutputs registers are the output registers
    circuit_vects::Vector{Vector{MyInt}}
    params::Parameters
end
Circuit = Union{Chromosome,LinCircuit}

# indiv_result_type is used to parallelize population evolution in Indiv_evolution.jl
mutable struct indiv_result_type
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levelsback::Int64
  ngoals::Int64
  hamming_sel::Bool
  #avgfitness::Bool
  #robust_sel::Bool
  active_only::Bool
  maxsteps::Int64
  #gl_reps::Int64
  #fault_tol::Bool
  #fit_limit::Float64
  hdist::Float64
  steps::Int64
  same::Int64
  worse::Int64
  better::Int64
  nactive::Int64
  redundancy::Float64
  complexity::Float64
  degeneracy::Float64
  sdegeneracy::Float64
end

# indiv_result_type is used to parallelize population evolution in Indiv_evolution.jl
mutable struct geno_pheno_result_type
  goallist::Vector{Vector{MyInt}}
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levelsback::Int64
  ngoals::Int64
  hamming_sel::Bool
  active_only::Bool
  maxsteps::Int64
  gl_reps::Int64
  steps::Int64
  logsteps::Float64
  avgfit::Float64
  nactive::Int64
  redund::Float64
  complex::Float64
  gb_complex::Float64
  degen::Float64
  #gb_degen::Float64
  sdegen::Float64
  f_mutinf::Float64
  mutrobust::Float64
  evolvability::Float64
end

# env_result_type is used to parallelize population evolution in Env_evolution.jl
mutable struct env_result_type
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levelsback::Int64
  ngoals::Int64
  hamming_sel::Bool
  #robust_sel::Bool
  active_only::Bool
  maxsteps::Int64
  gl_reps::Int64
  num_flip_bits::Int64
  perturb_goal::Bool
  avgfitness::Bool
  perm_heuristic::Bool
  fit_limit::Float64
  steps::Int64
  avgfit::Float64   # average fitness of final chromosome
  same::Int64
  worse::Int64
  better::Int64
  nactive::Int64
  redundancy::Float64
  complexity::Float64
  degeneracy::Float64
  sdegeneracy::Float64
end

# pop_result_type is used to parallelize population evolution in Pop_evolution.jl
mutable struct pop_result_type
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levelsback::Int64
  ngoals::Int64
  popsize::Int64
  max_pop_gens::Int64
  max_indiv_steps::Int64
  hamming_sel::Bool
  robust_sel::Bool
  all_max_sel::Bool
  active_only::Bool
  maxfit::Float64
  maxrobust::Float64
  nactive::Int64
  redundancy::Float64
  complexity::Float64
  degeneracy::Float64
  sdegeneracy::Float64
end

# pop_result_type is used to parallelize population evolution in Inf_alleles.jl
mutable struct inf_alleles_result_type
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levsback::Int64
  gl::Vector{Goal}
  popsize::Int64
  max_pop_gens::Int64
  tourn_size::Int64
  func_evals::Int64
  fitness::Float64
  gen_finished::Int64
end

mutable struct evo_result_type
  goal::Goal
  nchromes::Int64
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levelsback::Int64
  nrepeats::Int64
  maxsteps::Int64
  all_count::Int64
  evolvable_count::Int64
  evo_diff_count::Int64
end

mutable struct evo_pairs_type
  source_g::Goal
  dest_g::Goal
  numinputs::Int64
  numoutputs::Int64
  numints::Int64
  levelsback::Int64
  maxsteps::Int64
  hamming_dist::Float64
  source_g_count::Int64
  dest_g_count::Int64
  steps::Int64
end
#=
abstract type CC end
mutable struct Cc <: CC
  circuits::Vector{CC}
  inputs::Vector{Int64}
  numoutputs::Int64
  Cc( circuits::Vector{CC}, inputs::Vector{Int64}, numoutputs::Int64 ) = ( x=new(); x.inputs=inputs; x.numoutput=numoutputs; x.circuits=x )
  Cc( cv::Vector{CC}, inputs::Vector{Int64}, numoutputs::Int64 ) = fCc( new(), cv, inputs, numoutputs )
end

function fCc( c::Cc, circuits::Vector{CC}, inputs::Vector{Int64}, numoutputs::Int64 )
  c.circuits=circuits
  c.inputs=inputs
  c.numoutputs=numoutputs
  return c
end

mutable struct Cg <: CC
  func::Func
  inputs::Vector{Int64}
  numoutputs::Int64
end

mutable struct Need
  goal::Goal
  circuit_index::Int64   # If nonzero, the index of the element of all_circuits that meets the need
  inputs::Vector{Int64}
end
=#
abstract type CC end

mutable struct Cc <: CC
  circuits::Vector{CC}
  inputs::Vector{Int64}
  numoutputs::Int64
  output::Goal    # cached output when applied to standard for length(inputs)
  key::String     # dictionary key
  cache_valid::Int64      # non-zero implies cache (which is output) is valid
  Cc( circuits::Vector{CC}, inputs::Vector{Int64}, numoutputs::Int64, output::Goal, key::String, cache_valid::Int64 ) = ( x=new(); x.inputs=inputs; x.numoutputs=numoutputs; x.circuits=circuits; x.output=output; x.key=key; x.cache_valid=x.cache_valid )
  Cc( cv::Vector{CC}, inputs::Vector{Int64}, numoutputs::Int64, output::Goal, key::String, cache_valid::Int64 ) = fCc( new(), cv, inputs, numoutputs, output::Goal, key::String, cache_valid::Int64 )
end

function fCc( c::Cc, circuits::Vector{CC}, inputs::Vector{Int64}, numoutputs::Int64, output::Goal, key::String, cache_valid::Int64 )
  c.circuits=circuits
  c.inputs=inputs
  c.numoutputs=numoutputs
  c.output = output  # use as a cache
  c.key = key
  c.cache_valid = cache_valid
  return c
end

mutable struct Cg <: CC
  func::Func
  inputs::Vector{Int64}
  numoutputs::Int64
  output::Goal
  key::String
end

mutable struct Need
  goal::Goal    # dictionary key
  inputs::Vector{Int64}
  circuit_key::String   # If not "", the key of the all_circuits dictionary that meets the need
  best_fitness::Float64
end


@testset "MultiResolutionHybrid" begin
    img = generate_surface(SynthSphere(50), 0.5)
    p,q = MultiResolutionHybrid()(img)
    @test_broken 1 == 2
end


struct NotImplementedOperation <: AbstractGroupOperation end

struct NotImplementedManifold <: Manifold{ℝ} end

struct NotImplementedGroupDecorator{M} <: AbstractDecoratorManifold{ℝ}
    manifold::M
end


module Train

export train!

using FluxGAN: GAN
using FluxGAN: image_grid

using Flux
using Flux.Losses: logitbinarycrossentropy
using EllipsisNotation
using Statistics
using Colors
using FileIO
using CUDA

######################
# training functions #
######################

# main training function - attempts to train on GPU

function train!(model::GAN, train_tensor::AbstractArray; kws...)
    if has_cuda_gpu()
        println("training GAN with GPU...")
        train_gpu!(model, train_tensor; kws...)
    else
        println("no GPU available :( training GAN with CPU...")
        train_cpu!(model, train_tensor; kws...)
    end
end

# gpu training function

function train_gpu!(model::GAN, train_tensor::AbstractArray;
                    iterations=1000, verbose=true, skip=50,
                    gif=false, gif_path="gifs", gif_layout=(5, 5), gif_fps=25,
                    gif_filename="", return_losses=false, kws...)

    d = model.hparams.latent_dim
    m = model.hparams.minibatch

    mean_gen_losses = Float32[]
    mean_dscr_losses = Float32[]

    gen_losses = Float32[]
    dscr_losses = Float32[]

    t0 = time_ns()

    if gif
        rows, cols = gif_layout
        h, w = model.hparams.img_size[1:2]
        gif_tensor = Array{Color}(undef, h * rows, w * cols, div(iterations, skip))
    end

    model = fmap(gpu, model)
    for i = 1:iterations
        if verbose && i % skip == 0
            mean_gen_loss = mean(gen_losses)
            mean_dscr_loss = mean(dscr_losses)
            push!(mean_gen_losses, mean_gen_loss)
            push!(mean_dscr_losses, mean_dscr_loss)
            training_message(i, iterations, t0, mean_gen_loss, mean_dscr_loss, skip; kws...)
            if gif
                gif_tensor[:,:,div(i, skip)] = image_grid(model, gif_layout; seed=true)
            end
            gen_losses = []
            dscr_losses = []
            t0 = time_ns()
        end
        for _ = 1:model.hparams.dscr_loops
            x_tensor = train_tensor[.., rand(1:size(train_tensor)[end], m)] |> gpu
            z_tensor = CUDA.randn(Float32, d, m)
            dscr_loss = train_discriminator!(model, x_tensor, z_tensor)
            push!(dscr_losses, dscr_loss)
        end
        z_tensor = CUDA.randn(Float32, d, m)
        gen_loss = train_generator!(model, z_tensor)
        push!(gen_losses, gen_loss)
    end
    model = fmap(cpu, model)

    if gif
        save(gif_path * "/" * gif_filename * "_n_$(iterations)_" *
             "grid_$(rows)_$(cols)" * ".gif",
             gif_tensor; fps=gif_fps)
    end

    if return_losses
        return mean_gen_losses, mean_dscr_losses
    end
end

# cpu training function

function train_cpu!(model::GAN, train_tensor::AbstractArray;
                    iterations=1000, verbose=true, skip=50,
                    gif=false, gif_path="images/gifs", gif_layout=(5, 5), gif_fps=25,
                    gif_filename="", kws...)

    d = model.hparams.latent_dim
    m = model.hparams.minibatch

    gen_loss = 0
    dscr_loss = 0

    t0 = time_ns()

    if gif
        gif_tensor = Array{Color}(undef, size(train_tensor)[1:2]...,
                                       *(gif_layout...),
                                       div(iterations, skip))
    end

    for i = 1:iterations
        if verbose && i % skip == 0
            training_message(i, iterations, t0, gen_loss, dscr_loss, skip; kws...)
            if gif
                gif_tensor[..,div(i, skip)] = image_grid_tensor(model,
                                                *(gif_layout...);
                                                seed=true)
            end
            gen_loss = 0
            dscr_loss = 0
            t0 = time_ns()
        end
        for _ = 1:model.hparams.dscr_loops
            x_tensor = train_tensor[.., rand(1:size(train_tensor)[end], m)]
            z_tensor = randn(Float32, d, m)
            dscr_loss += train_discriminator!(model, x_tensor, z_tensor)
        end
        z_tensor = randn(Float32, d, m)
        gen_loss += train_generator!(model, z_tensor)
    end

    if gif
        save(gif_path * "/" * gif_filename * "_n_$(iterations)_" *
             "grid_$(gif_layout[1])_$(gif_layout[2])" * ".gif",
             gif_tensor; fps=gif_fps)
    end
end

# discriminator train function

function train_discriminator!(model::GAN, xs::AbstractArray, zs::AbstractArray)
    ω = Flux.params(model.D)
    loss, ∇ = Flux.pullback(ω) do
        discriminator_loss(model.D(xs), model.D(model.G(zs)))
    end
    ∇ = ∇(1f0)
    Flux.update!(model.D_opt, ω, ∇)
    return loss
end

# generator train function

function train_generator!(model::GAN, zs::AbstractArray)
    θ = Flux.params(model.G)
    loss, ∇ = Flux.pullback(θ) do
        generator_loss(model.D(model.G(zs)))
    end
    ∇ = ∇(1f0)
    Flux.update!(model.G_opt, θ, ∇)
    return loss
end

##################
# loss functions #
##################

function discriminator_loss(real_outputs, fake_outputs)
    real_loss = logitbinarycrossentropy(real_outputs, 1)
    fake_loss = logitbinarycrossentropy(fake_outputs, 0)
    return real_loss + fake_loss
end

generator_loss(fake_outputs) = logitbinarycrossentropy(fake_outputs, 1)

# training message function

# TODO: add plotting functionality
#           * mean training loss
#           * maybe error bars

function training_message(i, iterations, t0, gen_loss, dscr_loss, skip)
    t = time_ns()
    println("\niteration: $i of ", iterations)
    println("avg seconds per loop   = ", Float32((t - t0) / skip)*1f-9)
    println("avg generator loss     = ", Float32(gen_loss))
    println("avg discriminator loss = ", Float32(dscr_loss))
end

end


module ARDESPOT

using POMDPs
using BeliefUpdaters
using Parameters
using CPUTime
using ParticleFilters
using D3Trees
using Random
using Printf
using POMDPModelTools

using BasicPOMCP # for ExceptionRethrow and NoDecision
import BasicPOMCP.default_action

import Random.rand

export
    DESPOTSolver,
    DESPOTPlanner,

    DESPOTRandomSource,
    MemorizingSource,
    MemorizingRNG,

    ScenarioBelief,
    previous_obs,

    default_action,
    NoGap,

    IndependentBounds,
    FullyObservableValueUB,
    DefaultPolicyLB,
    bounds,
    init_bounds,
    lbound,
    ubound,
    init_bound,

    ReportWhenUsed

# include("random.jl")
include("random_2.jl")

"""
    DESPOTSolver(<keyword arguments>)

Implementation of the ARDESPOT solver trying to closely match the pseudo code of:

http://bigbird.comp.nus.edu.sg/m2ap/wordpress/wp-content/uploads/2017/08/jair14.pdf

Each field may be set via keyword argument. The fields that correspond to algorithm
parameters match the definitions in the paper exactly.

# Fields
- `epsilon_0`
- `xi`
- `K`
- `D`
- `lambda`
- `T_max`
- `max_trials`
- `bounds`
- `default_action`
- `rng`
- `random_source`
- `bounds_warnings`
- `tree_in_info`

Further information can be found in the field docstrings (e.g.
`?DESPOTSolver.xi`)
"""
@with_kw mutable struct DESPOTSolver <: Solver
    "The target gap between the upper and the lower bound at the root of the partial DESPOT."
    epsilon_0::Float64                      = 0.0

    "The rate of target gap reduction."
    xi::Float64                             = 0.95

    "The number of sampled scenarios."
    K::Int                                  = 500

    "The maximum depth of the DESPOT."
    D::Int                                  = 90

    "Reguluarization constant."
    lambda::Float64                         = 0.01

    "The maximum online planning time per step."
    T_max::Float64                          = 1.0

    "The maximum number of trials of the planner."
    max_trials::Int                         = typemax(Int)

    "A representation for the upper and lower bound on the discounted value (e.g. `IndependentBounds`)."
    bounds::Any                             = IndependentBounds(-1e6, 1e6)

    """A default action to use if algorithm fails to provide an action because of an error.

    This can either be an action object, i.e. `default_action=1` if `actiontype(pomdp)==Int` or a function `f(pomdp, b, ex)` where b is the belief and ex is the exception that caused the planner to fail.
    """
    default_action::Any                     = ExceptionRethrow()

    "A random number generator for the internal sampling processes."
    rng::MersenneTwister                    = MersenneTwister(rand(UInt32))

    "A source for random numbers in scenario rollout"
    random_source::DESPOTRandomSource       = MemorizingSource(K, D, rng)

    "If true, sanity checks on the provided bounds are performed."
    bounds_warnings::Bool                   = true

    "If true, a reprenstation of the constructed DESPOT is returned by POMDPModelTools.action_info."
    tree_in_info::Bool                      = false
end

include("scenario_belief.jl")
include("default_policy_sim.jl")
include("bounds.jl")

struct DESPOTPlanner{P<:POMDP, B, RS<:DESPOTRandomSource, RNG<:AbstractRNG} <: Policy
    sol::DESPOTSolver
    pomdp::P
    bounds::B
    rs::RS
    rng::RNG
end

function DESPOTPlanner(sol::DESPOTSolver, pomdp::POMDP)
    bounds = init_bounds(sol.bounds, pomdp, sol)
    rng = deepcopy(sol.rng)
    rs = deepcopy(sol.random_source)
    Random.seed!(rs, rand(rng, UInt32))
    return DESPOTPlanner(deepcopy(sol), pomdp, bounds, rs, rng)
end

include("tree.jl")
include("planner.jl")
include("pomdps_glue.jl")

include("visualization.jl")
include("exceptions.jl")

end # module


using Random
using PlotUtils
using Makie
using AbstractPlotting
import Showoff
using Printf
using AbstractPlotting: Rect2D
import AbstractPlotting: IRect2D

const BBox = Rect2D{Float64}

left(rect::Rect2D) = minimum(rect)[1]
right(rect::Rect2D) = maximum(rect)[1]

bottom(rect::Rect2D) = minimum(rect)[2]
top(rect::Rect2D) = maximum(rect)[2]

abstract type Side end

struct Left <: Side end
struct Right <: Side end
struct Top <: Side end
struct Bottom <: Side end

Base.getindex(bbox::Rect2D, ::Left) = left(bbox)
Base.getindex(bbox::Rect2D, ::Right) = right(bbox)
Base.getindex(bbox::Rect2D, ::Bottom) = bottom(bbox)
Base.getindex(bbox::Rect2D, ::Top) = top(bbox)

mutable struct LayoutedAxis
    parent::Scene
    scene::Scene
    xlabel::Node{String}
    ylabel::Node{String}
    limits::Node{FRect2D} # these should steer the camera, not used yet
end

width(rect::Rect2D) = right(rect) - left(rect)
height(rect::Rect2D) = top(rect) - bottom(rect)

function BBox(left::Number, right::Number, top::Number, bottom::Number)
    mini = (left, bottom)
    maxi = (right, top)
    return BBox(mini, maxi .- mini)
end

function IRect2D(bbox::Rect2D)
    return IRect2D(
        round.(Int, minimum(bbox)),
        round.(Int, widths(bbox))
    )
end

struct RowCols{T <: Union{Number, Vector{Float64}}}
    lefts::T
    rights::T
    tops::T
    bottoms::T
end

function RowCols(ncols::Int, nrows::Int)
    return RowCols(
        zeros(ncols),
        zeros(ncols),
        zeros(nrows),
        zeros(nrows)
    )
end

Base.getindex(rowcols::RowCols, ::Left) = rowcols.lefts
Base.getindex(rowcols::RowCols, ::Right) = rowcols.rights
Base.getindex(rowcols::RowCols, ::Top) = rowcols.tops
Base.getindex(rowcols::RowCols, ::Bottom) = rowcols.bottoms

"""
    eachside(f)
Calls f over all sides (Left, Right, Top, Bottom), and creates a BBox from the result of f(side)
"""
function eachside(f)
    return BBox(map(f, (Left(), Right(), Top(), Bottom()))...)
end

"""
mapsides(
       f, first::Union{Rect2D, RowCols}, rest::Union{Rect2D, RowCols}...
   )::BBox
Maps f over all sides of the rectangle like arguments.
e.g.
```
mapsides(BBox(left, right, top, bottom)) do side::Side, side_val::Number
    return ...
end::BBox
```
"""
function mapsides(
        f, first::Union{Rect2D, RowCols}, rest::Union{Rect2D, RowCols}...
    )
    return eachside() do side
        f(side, getindex.((first, rest...), (side,))...)
    end
end

abstract type Alignable end

"""
Used to specify space that is occupied in a grid. Like 1:1|1:1 for the first square,
or 2:3|1:4 for a rect over the 2nd and 3rd row and the first four columns.
"""
struct Span
    rows::UnitRange{Int64}
    cols::UnitRange{Int64}
end

"""
An object that can be aligned that also specifies how much space it occupies in
a grid via its span.
"""
struct SpannedAlignable{T <: Alignable}
    al::T
    sp::Span
end
function side_indices(c::SpannedAlignable)
    return RowCols(
        c.sp.cols.start,
        c.sp.cols.stop,
        c.sp.rows.start,
        c.sp.rows.stop,
    )
end

"""
These functions tell whether an object in a grid touches the left, top, etc. border
of the grid. This means that it is relevant for the grid's own protrusion on that side.
"""
ismostin(sp::SpannedAlignable, grid, ::Left) = sp.sp.cols.start == 1
ismostin(sp::SpannedAlignable, grid, ::Right) = sp.sp.cols.stop == grid.ncols
ismostin(sp::SpannedAlignable, grid, ::Bottom) = sp.sp.rows.stop == grid.nrows
ismostin(sp::SpannedAlignable, grid, ::Top) = sp.sp.cols.start == 1

isleftmostin(sp::SpannedAlignable, grid) = ismostin(sp, grid, Left())
isrightmostin(sp::SpannedAlignable, grid) = ismostin(sp, grid, Right())
isbottommostin(sp::SpannedAlignable, grid) = ismostin(sp, grid, Bottom())
istopmostin(sp::SpannedAlignable, grid) = ismostin(sp, grid, Top())

struct SolvedAxisLayout <: Alignable
    inner::BBox
    outer::BBox
    axis::LayoutedAxis
end

struct AxisLayout <: Alignable
    decorations::BBox
    axis::LayoutedAxis
end

struct SolvedGridLayout <: Alignable
    bbox::BBox
    content::Vector{SpannedAlignable}
    nrows::Int
    ncols::Int
    grid::RowCols{Vector{Float64}}
end

struct GridLayout <: Alignable
    content::Vector{SpannedAlignable}
    nrows::Int
    ncols::Int
    colratios::Vector{Float64}
    rowratios::Vector{Float64}
    colgapfraction::Float64
    rowgapfraction::Float64
end

"""
All the protrusion functions calculate how much stuff "sticks out" of a layoutable object.
This is so that collisions are avoided, while what is actually aligned is the
"important" edges of the layout objects.
"""
leftprotrusion(x) = protrusion(x, Left())
rightprotrusion(x) = protrusion(x, Right())
bottomprotrusion(x) = protrusion(x, Bottom())
topprotrusion(x) = protrusion(x, Top())

protrusion(u::AxisLayout, side::Side) = u.decorations[side]
protrusion(sp::SpannedAlignable, side::Side) = protrusion(sp.al, side)

function protrusion(gl::GridLayout, side::Side)
    return mapreduce(max, gl.content, init = 0.0) do elem
        # we use only objects that stick out on this side
        # And from those we use the maximum protrusion
        ismostin(elem, gl, side) ? protrusion(elem, side) : 0.0
    end
end

protrusion(s::SolvedAxisLayout, ::Left) = left(s.inner) - left(s.outer)
function protrusion(s::SolvedAxisLayout, side::Side)
    return s.outer[side] - s.inner[side]
end

"""
This function solves a grid layout such that the "important lines" fit exactly
into a given bounding box. This means that the protrusions of all objects inside
the grid are not taken into account. This is needed if the grid is itself placed
inside another grid.
"""
function solve(gl::GridLayout, bbox::BBox)

    # first determine how big the protrusions on each side of all columns and rows are
    maxgrid = RowCols(gl.ncols, gl.nrows)
    # go through all the layout objects placed in the grid
    for c in gl.content
        idx_rect = side_indices(c)
        mapsides(idx_rect, maxgrid) do side, idx, grid
            grid[idx] = max(grid[idx], protrusion(c.al, side))
        end
    end
    # compute what size the gaps between rows and columns need to be
    colgaps = maxgrid.lefts[2:end] .+ maxgrid.rights[1:end-1]
    rowgaps = maxgrid.tops[2:end] .+ maxgrid.bottoms[1:end-1]

    # determine the biggest gap
    # using the biggest gap size for all gaps will make the layout more even, but one
    # could make this aspect customizable, because it might waste space
    maxcolgap = maximum(colgaps)
    maxrowgap = maximum(rowgaps)

    # determine the vertical and horizontal space needed just for the gaps
    # again, the gaps are what the protrusions stick into, so they are not actually "empty"
    # depending on what sticks out of the plots
    sumcolgaps = maxcolgap * (gl.ncols - 1)
    sumrowgaps = maxrowgap * (gl.nrows - 1)

    # compute what space remains for the inner parts of the plots
    remaininghorizontalspace = width(bbox) - sumcolgaps
    remainingverticalspace = height(bbox) - sumrowgaps

    # compute how much gap to add, in case e.g. labels are too close together
    # this is given as a fraction of the space used for the inner parts of the plots
    # so far, but maybe this should just be an absolute pixel value so it doesn't change
    # when resizing the window
    addedcolgap = gl.colgapfraction * remaininghorizontalspace
    addedrowgap = gl.rowgapfraction * remainingverticalspace

    # compute the actual space available for the rows and columns (plots without protrusions)
    spaceforcolumns = remaininghorizontalspace - addedcolgap * (gl.ncols - 1)
    spaceforrows = remainingverticalspace - addedrowgap * (gl.nrows - 1)

    # compute the column widths and row heights using the specified row and column ratios
    colwidths = gl.colratios ./ sum(gl.colratios) .* spaceforcolumns
    rowheights = gl.rowratios ./ sum(gl.rowratios) .* spaceforrows

    # this is the vertical / horizontal space between the inner lines of all plots
    colgap = maxcolgap + addedcolgap
    rowgap = maxrowgap + addedrowgap

    # compute the x values for all left and right column boundaries
    xleftcols = [left(bbox) + sum(colwidths[1:i-1]) + (i - 1) * colgap for i in 1:gl.ncols]
    xrightcols = xleftcols .+ colwidths

    # compute the y values for all top and bottom row boundaries
    ytoprows = [top(bbox) - sum(rowheights[1:i-1]) - (i - 1) * rowgap for i in 1:gl.nrows]
    ybottomrows = ytoprows .- rowheights

    # now we can solve the content thats inside the grid because we know where each
    # column and row is placed, how wide it is, etc.
    # note that what we did at the top was determine the protrusions of all grid content,
    # but we know the protrusions before we know how much space each plot actually has
    # because the protrusions should be static (like tick labels etc don't change size with the plot)

    gridboxes = RowCols(
        xleftcols, xrightcols,
        ytoprows, ybottomrows
    )
    solvedcontent = map(gl.content) do c
        idx_rect = side_indices(c)
        bbox_cell = mapsides(idx_rect, gridboxes) do side, idx, gridside
            gridside[idx]
        end
        solved = solve(c.al, bbox_cell)
        return SpannedAlignable(solved, c.sp)
    end
    # return a solved grid layout in which all objects are also solved layout objects
    return SolvedGridLayout(
        bbox, solvedcontent,
        gl.nrows, gl.ncols,
        gridboxes
    )
end

"""
This function solves a grid layout so that it fits exactly inside a bounding box.
Exactly means that the protrusions of all other objects inside this grid layout
also have to fit into the bounding box. This is needed if the grid is the outermost
object in the layout, the bounding box would then be the scene boundary.
"""
function outersolve(gl::GridLayout, bbox::BBox)
    maxgrid = RowCols(gl.ncols, gl.nrows)
    for c in gl.content
        idx_rect = side_indices(c)
        mapsides(idx_rect, maxgrid) do side, idx, grid
            grid[idx] = max(grid[idx], protrusion(c.al, side))
        end
    end

    topprot = maxgrid.tops[1]
    bottomprot = maxgrid.bottoms[end]
    leftprot = maxgrid.lefts[1]
    rightprot = maxgrid.rights[end]

    colgaps = maxgrid.lefts[2:end] .+ maxgrid.rights[1:end-1]
    rowgaps = maxgrid.tops[2:end] .+ maxgrid.bottoms[1:end-1]

    maxcolgap = gl.ncols <= 1 ? 0 : maximum(colgaps)
    maxrowgap = gl.nrows <= 1 ? 0 : maximum(rowgaps)

    sumcolgaps = maxcolgap * (gl.ncols - 1)
    sumrowgaps = maxrowgap * (gl.nrows - 1)

    remaininghorizontalspace = width(bbox) - sumcolgaps - leftprot - rightprot
    remainingverticalspace = height(bbox) - sumrowgaps - topprot - bottomprot

    addedcolgap = gl.colgapfraction * remaininghorizontalspace
    addedrowgap = gl.rowgapfraction * remainingverticalspace

    spaceforcolumns = remaininghorizontalspace - addedcolgap * (gl.ncols - 1)
    spaceforrows = remainingverticalspace - addedrowgap * (gl.nrows - 1)

    colwidths = gl.colratios ./ sum(gl.colratios) .* spaceforcolumns
    rowheights = gl.rowratios ./ sum(gl.rowratios) .* spaceforrows

    colgap = maxcolgap + addedcolgap
    rowgap = maxrowgap + addedrowgap

    xleftcols = map(1:gl.ncols) do i
        left(bbox) + leftprot + sum(colwidths[1:i-1]) + (i - 1) * colgap
    end
    xrightcols = xleftcols .+ colwidths

    ytoprows = map(1:gl.nrows) do i
        top(bbox) - topprot - sum(rowheights[1:i-1]) - (i - 1) * rowgap
    end

    ybottomrows = ytoprows .- rowheights
    gridboxes = RowCols(
        xleftcols, xrightcols,
        ytoprows, ybottomrows
    )
    solvedcontent = map(gl.content) do c
        idx_rect = side_indices(c)
        bbox_cell = mapsides(idx_rect, gridboxes) do side, idx, gridside
            gridside[idx]
        end
        solved = solve(c.al, bbox_cell)
        return SpannedAlignable(solved, c.sp)
    end
    # solvedcontent = solve.(gl.content)
    SolvedGridLayout(
        bbox, solvedcontent, gl.nrows, gl.ncols,
        gridboxes
    )
end

function solve(ua::AxisLayout, innerbbox)
    bbox = mapsides(innerbbox, ua.decorations) do side, iside, decside
        op = side isa Union{Left, Top} ? (-) : (+)
        return op(iside, decside)
    end
    SolvedAxisLayout(innerbbox, bbox, ua.axis)
end

const Indexables = Union{UnitRange, Int, Colon}

"""
This function allows indexing syntax to add a layout object to a grid.
You can do:

grid[1, 1] = obj
grid[1, :] = obj
grid[1:3, 2:5] = obj

and all combinations of the above
"""
function Base.setindex!(g, a::Alignable, rows::Indexables, cols::Indexables)
    if rows isa Int
        rows = rows:rows
    elseif rows isa Colon
        rows = 1:g.nrows
    end
    if cols isa Int
        cols = cols:cols
    elseif cols isa Colon
        cols = 1:g.ncols
    end

    if !((1 <= rows.start <= g.nrows) || (1 <= rows.stop <= g.nrows))
        error("invalid row span $rows for grid with $(g.nrows) rows")
    end
    if !((1 <= cols.start <= g.ncols) || (1 <= cols.stop <= g.ncols))
        error("invalid col span $cols for grid with $(g.ncols) columns")
    end

    push!(g.content, SpannedAlignable(a, Span(rows, cols)))
end

function axislines!(scene, rect)
    points = lift(rect) do r
        p1 = Point2(r.origin[1], r.origin[2] + r.widths[2])
        p2 = Point2(r.origin[1], r.origin[2])
        p3 = Point2(r.origin[1] + r.widths[1], r.origin[2])
        [p1, p2, p3]
    end
    lines!(scene, points, linewidth = 2, show_axis = false)
end

"""
A cheaper function that tries to come up with usable tick locations for a given value range
"""
function locateticks(xmin, xmax)
    # whats the distance?
    d = xmax - xmin
    # which order of magnitude is the distance in?
    ex = log10(d)
    # round the exponent to the closest one
    exrounded = round(ex)
    # this factor is used so we can find integer steps that then map back to
    # nice numbers in the given value range
    factor = 1 / 10 ^ (exrounded - 1)

    # minimum needs to be at least the lower value times the scaling factor
    xminf = ceil(xmin * factor)
    # maximum needs to be at most the higher value times the scaling factor
    xmaxf = floor(xmax * factor)

    # xminf and xmaxf are now both integers that are in an order of magnitude around ten steps apart
    df = xmaxf - xminf

    # step sizes we like
    steps = [5, 4, 2, 1]

    # from the highest to the lowest step size, choose the first that fits at least
    # two times between the end values (gives three ticks including the end value)
    for s in steps
        n, remainder = divrem(df, s)
        if n >= 2
            rang = 1:n
            # calculate the ticks by dividing with the factor from above
            ticks = [xminf; [xminf + x * s for x in rang]] ./ factor
            return ticks
        end
    end
end

function LayoutedAxis(parent::Scene)
    scene = Scene(parent, Node(IRect(0, 0, 100, 100)))
    limits = Node(FRect(0, 0, 100, 100))
    xlabel = Node("x label")
    ylabel = Node("y label")

    cam = cam2d!(scene)

    ticksnode = Node(Point2f0[])
    ticks = linesegments!(
        parent, ticksnode, linewidth = 2, show_axis = false
    )[end]

    # the algorithm from above seems to not give more than 7 ticks with the step sizes I chose
    nmaxticks = 7

    xticklabelnodes = [Node("0") for i in 1:nmaxticks]
    xticklabelposnodes = [Node(Point(0.0, 0.0)) for i in 1:nmaxticks]
    xticklabels = map(1:nmaxticks) do i
        text!(
            parent,
            xticklabelnodes[i],
            position = xticklabelposnodes[i],
            align = (:center, :top),
            textsize = 20,
            show_axis = false
        )[end]
    end

    yticklabelnodes = [Node("0") for i in 1:nmaxticks]
    yticklabelposnodes = [Node(Point(0.0, 0.0)) for i in 1:nmaxticks]
    yticklabels = map(1:nmaxticks) do i
        text!(
            parent,
            yticklabelnodes[i],
            position = yticklabelposnodes[i],
            align = (:center, :bottom),
            rotation = pi/2,
            textsize = 20,
            show_axis = false
        )[end]
    end

    on(cam.area) do a

        pxa = scene.px_area[]
        px_aspect = pxa.widths[1] / pxa.widths[2]

        # @printf("cam %.1f, %.1f, %.1f, %.1f\n", a.origin..., a.widths...)
        # @printf("pix %.1f, %.1f, %.1f, %.1f\n", pxa.origin..., pxa.widths...)

        width = px_aspect > 1 ? a.widths[1] / px_aspect : a.widths[1]
        xrange = (a.origin[1], a.origin[1] + width)

        if width == 0 || !isfinite(xrange[1]) || !isfinite(xrange[2])
            return
        end

        xtickvals = locateticks(xrange...)
        # xtickvals, vminbest, vmaxbest = optimize_ticks(xrange...)

        # this code here tries to transform between values given in pixels of the
        # scene and the camera area, but this is incorrect right now and everything
        # should just be determined by the now unused limits that are saved in the
        # LayoutedAxis object
        xfractions = (xtickvals .- xrange[1]) ./ width
        xrange_scene = (pxa.origin[1], pxa.origin[1] + pxa.widths[1])
        width_scene = xrange_scene[2] - xrange_scene[1]
        xticks_scene = xrange_scene[1] .+ width_scene .* xfractions
        ticksize = 10 # px
        y = pxa.origin[2]
        xtickstarts = [Point(x, y) for x in xticks_scene]
        xtickends = [t + Point(0.0, -ticksize) for t in xtickstarts]

        height = px_aspect < 1 ? a.widths[2] * px_aspect : a.widths[2]
        yrange = (a.origin[2], a.origin[2] + height)

        ytickvals = locateticks(yrange...)
        # ytickvals, vminbest, vmaxbest = optimize_ticks(yrange...)
        yfractions = (ytickvals .- yrange[1]) ./ height
        yrange_scene = (pxa.origin[2], pxa.origin[2] + pxa.widths[2])
        height_scene = yrange_scene[2] - yrange_scene[1]
        yticks_scene = yrange_scene[1] .+ height_scene .* yfractions
        ticksize = 10 # px
        x = pxa.origin[1]
        ytickstarts = [Point(x, y) for y in yticks_scene]
        ytickends = [t + Point(-ticksize, 0.0) for t in ytickstarts]

        # set and position tick labels
        xtickstrings = Showoff.showoff(xtickvals, :plain)
        nxticks = length(xtickvals)
        for i in 1:nmaxticks
            if i <= nxticks
                xticklabelnodes[i][] = xtickstrings[i]
                xticklabelposnodes[i][] = xtickends[i] + Point(0.0, -10.0)
                xticklabels[i].visible = true
            else
                xticklabels[i].visible = false
            end
        end

        ytickstrings = Showoff.showoff(ytickvals, :plain)
        nyticks = length(ytickvals)
        for i in 1:nmaxticks
            if i <= nyticks
                yticklabelnodes[i][] = ytickstrings[i]
                yticklabelposnodes[i][] = ytickends[i] + Point(-10.0, 0.0)
                yticklabels[i].visible = true
            else
                yticklabels[i].visible = false
            end
        end

        # set tick mark positions
        ticksnode[] = collect(Iterators.flatten(zip(
            [xtickstarts; ytickstarts],
            [xtickends; ytickends]
        )))
    end

    labelgap = 50

    xlabelpos = lift(scene.px_area) do a
        Point2(a.origin[1] + a.widths[1] / 2, a.origin[2] - labelgap)
    end

    ylabelpos = lift(scene.px_area) do a
        Point2(a.origin[1] - labelgap, a.origin[2] + a.widths[2] / 2)
    end

    tx = text!(
        parent, xlabel, textsize = 20, position = xlabelpos, show_axis = false
    )[end]
    tx.align = [0.5, 1]
    ty = text!(
        parent, ylabel, textsize = 20,
        position = ylabelpos, rotation = pi/2, show_axis = false
    )[end]
    ty.align = [0.5, 0]

    axislines!(parent, scene.px_area)

    LayoutedAxis(parent, scene, xlabel, ylabel, limits)
end

function applylayout(sg::SolvedGridLayout)
    for c in sg.content
        applylayout(c.al)
    end
end

function applylayout(sa::SolvedAxisLayout)
    sa.axis.scene.px_area[] = IRect2D(sa.inner)
end

function shrinkbymargin(rect, margin)
    IRect((rect.origin .+ margin), (rect.widths .- 2 .* margin))
end

begin
    scene = Scene(resolution=(600, 600));
    screen = display(scene)
    campixel!(scene);

    la1 = LayoutedAxis(scene)
    la2 = LayoutedAxis(scene)
    la3 = LayoutedAxis(scene)
    la4 = LayoutedAxis(scene)
    la5 = LayoutedAxis(scene)

    # lines!(la1.scene, rand(200, 2) .* 100, color=:black, show_axis=false)
    img = rand(100, 100)
    image!(la1.scene, img, show_axis=false)
    lines!(la2.scene, rand(200, 2) .* 100, color=:blue, show_axis=false)
    scatter!(la3.scene, rand(200, 2) .* 100, markersize=3, color=:red, show_axis=false)
    lines!(la4.scene, rand(200, 2) .* 100, color=:orange, show_axis=false)
    lines!(la5.scene, rand(200, 2) .* 100, color=:pink, show_axis=false)
    update!(scene)

    gl = GridLayout([], 2, 2, [1, 1], [1, 1], 0.01, 0.01)

    gl[2, 1:2] = AxisLayout(BBox(75, 0, 0, 75), la1)
    gl[1, 2] = AxisLayout(BBox(75, 0, 0, 75), la2)

    gl2 = GridLayout([], 2, 2, [0.8, 0.2], [0.2, 0.8], 0.01, 0.01)
    gl2[2, 1] = AxisLayout(BBox(75, 0, 0, 75), la3)
    gl2[1, 1] = AxisLayout(BBox(75, 0, 0, 75), la4)
    gl2[2, 2] = AxisLayout(BBox(75, 0, 0, 75), la5)

    gl[1, 1] = gl2

    sg = outersolve(gl, BBox(shrinkbymargin(pixelarea(scene)[], 30)))
    applylayout(sg)
    # when the scene is resized, apply the outersolve'd outermost grid layout
    # this recursively updates all layout objects that are contained in the grid
    on(scene.events.window_area) do area
        sg = outersolve(gl, BBox(shrinkbymargin(pixelarea(scene)[], 30)))
        applylayout(sg)
    end
    # @profiler wait(screen)
end


module SplitExplicit01

export OceanDGModel,
    OceanModel,
    Continuity3dModel,
    HorizontalModel,
    BarotropicModel,
    AbstractOceanProblem

#using Printf
using StaticArrays
using LinearAlgebra: I, dot, Diagonal

using ...VariableTemplates
using ...MPIStateArrays
using ...DGMethods: init_ode_state, basic_grid_info
using ...Mesh.Filters: CutoffFilter, apply!, ExponentialFilter
using ...Mesh.Grids:
    polynomialorders,
    dimensionality,
    dofs_per_element,
    VerticalDirection,
    HorizontalDirection,
    min_node_distance

using ...BalanceLaws
#import ...BalanceLaws: nodal_update_auxiliary_state!

using ...DGMethods.NumericalFluxes:
    NumericalFluxFirstOrder,
    NumericalFluxGradient,
    NumericalFluxSecondOrder,
    RusanovNumericalFlux,
    CentralNumericalFluxFirstOrder,
    CentralNumericalFluxGradient,
    CentralNumericalFluxSecondOrder

using ..Ocean: AbstractOceanProblem

import ...DGMethods.NumericalFluxes:
    update_penalty!, numerical_flux_second_order!, NumericalFluxFirstOrder

import ...DGMethods: LocalGeometry, DGModel, calculate_dt

import ...BalanceLaws:
    vars_state,
    flux_first_order!,
    flux_second_order!,
    source!,
    wavespeed,
    parameter_set,
    boundary_conditions,
    boundary_state!,
    update_auxiliary_state!,
    update_auxiliary_state_gradient!,
    compute_gradient_argument!,
    init_state_auxiliary!,
    init_state_prognostic!,
    compute_gradient_flux!,
    indefinite_stack_integral!,
    reverse_indefinite_stack_integral!,
    integral_load_auxiliary_state!,
    integral_set_auxiliary_state!,
    reverse_integral_load_auxiliary_state!,
    reverse_integral_set_auxiliary_state!

import ...SystemSolvers: BatchedGeneralizedMinimalResidual, linearsolve!

×(a::SVector, b::SVector) = StaticArrays.cross(a, b)
∘(a::SVector, b::SVector) = StaticArrays.dot(a, b)

abstract type AbstractOceanModel <: BalanceLaw end

function ocean_init_aux! end
function ocean_init_state! end
function ocean_model_boundary! end
function set_fast_for_stepping! end
function initialize_fast_state! end
function initialize_adjustment! end

include("SplitExplicitLSRK2nMethod.jl")
include("SplitExplicitLSRK3nMethod.jl")
include("OceanModel.jl")
include("Continuity3dModel.jl")
include("VerticalIntegralModel.jl")
include("BarotropicModel.jl")
include("IVDCModel.jl")
include("Communication.jl")
include("OceanBoundaryConditions.jl")

end


fetchvcat(J::judiVector) = J
fetchvcat(x) = vcat(x...)

function JUDI.time_modeling(model::Model, srcGeometry, srcData, recGeometry, recData, dm, srcnum::UnitRange{Int64}, op::Char, mode::Int64, options)

    # Broadcast common parameters
    _model = @bcast model
    _dm = isnothing(dm) ? dm : @bcast dm
    opts = make_opt([model, dm, op, srcGeometry])
    iter = make_parts(srcnum)
    # Run on azure
    results = @batchexec pmap(sx -> time_modeling_azure(_model, subsample(srcGeometry, sx), subsample(srcData, sx),
                                                        subsample(recGeometry, sx), subsample(recData, sx), _dm,
                                                        op, mode, subsample(options, sx)), iter) opts
    # Gather results
    if op=='F' || (op=='J' && mode==1)
        temp = fetch(results)
        argout1 = fetchvcat(temp)
    elseif op=='J' && mode==-1
        argout1 = fetchreduce(results; op=+, remote=true)
    else
        error("operation no defined")
    end
    return argout1
end

# FWI
function JUDI.fwi_objective(model::Model, source::judiVector, dObs::judiVector; options=JUDI.Options())
    # fwi_objective function for multiple sources. The function distributes the sources and the input data amongst the available workers.
    # Broadcast common parameters
    _model = @bcast model
    opts = make_opt([model, dObs, source])
    iter = make_parts(1:dObs.nsrc)
    results = @batchexec pmap(j -> fwi_objective_azure(_model, source[j], dObs[j], subsample(options, j)), iter) opts

    # Collect and reduce gradients
    obj, gradient = fetchreduce(results; op=+, remote=true)
    # first value corresponds to function value, the rest to the gradient
    return obj, gradient
end

# lsrtm

function JUDI.lsrtm_objective(model::Model, source::judiVector, dObs::judiVector, dm; options=JUDI.Options(), nlind=false)
    # lsrtm_objective function for multiple sources. The function distributes the sources and the input data amongst the available workers.

    # Broadcast common parameters
    _model = @bcast model
    _dm = isnothing(dm) ? dm : @bcast dm
    opts = make_opt([model, dObs, source, dm])
    iter = make_parts(1:dObs.nsrc)
    results = @batchexec pmap(j -> lsrtm_objective_azure(_model, source[j], dObs[j], _dm, subsample(options, j); nlind=nlind), iter) opts
    # Collect and reduce gradients
    obj, gradient = fetchreduce(results; op=+, remote=true)
    # first value corresponds to function value, the rest to the gradient
    return obj, gradient
end

# Make options for unique batch ids
make_opt(a::Array) = AzureClusterlessHPC.Options(;task_name="task_$(objectid(a))")

# Batchexec function
slicec(i::Integer, j::Integer) = i==j ? i : (i:j)
chunk(arr, n) = [arr[slicec(i, min(i + n - 1, length(arr)))] for i in 1:n:length(arr)]
make_parts(iter) = _njpi > 1 ? chunk(iter, _njpi) : iter