module TestVsSVector
include("preamble.jl")

using Bifurcations: resolved_points
using Bifurcations.Continuations: find_errors, print_errors
using Bifurcations: examples
using Bifurcations.Examples.Calcium
using Bifurcations.Examples.PredatorPrey

use_array(ex) =
    ex.make_prob(
        f = generalized_f(ex),
        u0 = Array(ex.u0),
    )

EXAMPLES = [
    # name => (
    #     prob1_svec,   # SVECtor
    #     prob1_stda,   # STanDard Array
    #     sweep_ids,    # sweeps to be tested
    # )
    :Calcium => (
        Calcium.prob,
        use_array(Calcium),
        :,
        ),
    # PredatorPrey* tests fail at sweeps[4].u[8].  Probably this kind
    # of test is not a good idea after all.  But let's keep it as-is
    # for now to "freeze" the implementation.
    # TODO: Find out how PredatorPrey* tests diverge.
    :PredatorPreyImmutableState => (
        PredatorPrey.prob,
        use_array(PredatorPrey),
        1:3,
    ),
    :PredatorPreyMutableState => (
        PredatorPrey.prob,
        PredatorPrey.make_prob(u0=Array(PredatorPrey.u0)),
        1:3,
    ),
]

@testset "example $name" for (name, probs) in EXAMPLES
    prob1_svec, prob1_stda, sweep_ids = probs
    solver1_svec = init(prob1_svec)
    solver1_stda = init(prob1_stda)

    solve!(solver1_svec)
    solve!(solver1_stda)
    sweeps_svec = solver1_svec.super.sol.sweeps
    sweeps_stda = solver1_stda.super.sol.sweeps

    @test length(sweeps_svec) == length(sweeps_stda)
    for (i, (sweep_svec, sweep_stda)) in enumerate(zip(sweeps_svec[sweep_ids],
                                                       sweeps_stda[sweep_ids]))
        @test length(sweep_svec.u) == length(sweep_stda.u)
        for (j, (u_svec, u_stda)) in enumerate(zip(sweep_svec.u,
                                                   sweep_stda.u))
            # Manually abort on the first test failure:
            if ! (u_svec ≈ u_stda)
                println("sweeps[$i].u[$j]")
                @test u_svec ≈ u_stda
                break
            else
                @test u_svec ≈ u_stda
            end
        end
    end
end

end  # module


using Fed.Config: BaseConfig, VanillaConfig, QuantizedConfig,
    QDiffConfig, QDiffStaticConfig, GDConfig, GDStaticConfig
using SHA
using CRC32c

"""
    hash_crc32(data)

Hash data with CRC32 and returns the hash as a `Vector{UInt8}`.
"""
function hash_crc32(data)
    hash = Vector{UInt8}(undef, 4)
    data = crc32c(data)
    hash[1] = (data & 0x000000ff)
    hash[2] = (data & 0x0000ff00) >> 8
    hash[3] = (data & 0x00ff0000) >> 16
    hash[4] = (data & 0xff000000) >> 24
    return hash
end

# config
const SERVERURL = "http://127.0.0.1:8080"
const NUM_COMM_ROUNDS = 100
const FRACTION_CLIENTS = 0.1f0
const NUM_TOTAL_CLIENTS = 100

const base_config = BaseConfig(
    SERVERURL,
    NUM_COMM_ROUNDS,
    FRACTION_CLIENTS,
    NUM_TOTAL_CLIENTS
)

newVanillaConfig() = VanillaConfig(base_config)
newQuantizedConfig(::Type{T}) where T <: Unsigned = QuantizedConfig{T}(base_config)
newQDiffConfig(::Type{T}, is_client::Bool) where T <: Unsigned =
    QDiffConfig{T}(base_config, 256, is_client)

newQDiffStaticConfig(::Type{T}, is_client::Bool) where T <: Unsigned =
    QDiffStaticConfig{T}(base_config, 256, is_client)

newGDConfig(::Type{T}, port::Int, is_client::Bool) where T <: Unsigned =
    GDConfig{T}(
        base_config,
        256,
        sha1,
        # hash_crc32,
        round(T, 0.125 * sizeof(T) * 8), # % of each weight which goes in the basis
        "127.0.0.1",
        port,
        is_client
    )

newGDStaticConfig(::Type{T}, port::Int, is_client::Bool) where T <: Unsigned =
    GDStaticConfig{T}(
        base_config,
        256,
        # sha1,
        hash_crc32,
        round(T, 0.6 * sizeof(T) * 8), # % of each weight which goes in the basis
        "127.0.0.1",
        port,
        is_client
    )

"""
    newconfig()

Returns a new current configuration.
"""
# newconfig(_, _) = newVanillaConfig()
# newconfig(_, _) = newQuantizedConfig(UInt8)
# newconfig(_, is_client) = newQDiffConfig(UInt8, is_client)
# newconfig(_, is_client) = newQDiffStaticConfig(UInt16, is_client)
newconfig(port, is_client) = newGDConfig(UInt8, port, is_client)
# newconfig(port, is_client) = newGDStaticConfig(UInt8, port, is_client)

"""
    newmodel()

Returns a new current model.
"""
newmodel() = initialize_mlp(28*28, 200, 10)


const JULIAN_EPOCH = TTEpoch(AstroDates.JULIAN_EPOCH, AstroDates.H12)
const J2000_EPOCH = TTEpoch(AstroDates.J2000_EPOCH, AstroDates.H12)
const MODIFIED_JULIAN_EPOCH = TTEpoch(AstroDates.MODIFIED_JULIAN_EPOCH, AstroDates.H00)
const FIFTIES_EPOCH = TTEpoch(AstroDates.FIFTIES_EPOCH, AstroDates.H00)
const CCSDS_EPOCH = TTEpoch(AstroDates.CCSDS_EPOCH, AstroDates.H00)
const GALILEO_EPOCH = TTEpoch(AstroDates.GALILEO_EPOCH, AstroDates.H00)
const GPS_EPOCH = TTEpoch(AstroDates.GPS_EPOCH, AstroDates.H00)
const UNIX_EPOCH = TAIEpoch(AstroDates.UNIX_EPOCH, Time(0, 0, 10.0))

const PAST_INFINITY = TAIEpoch(UNIX_EPOCH, -Inf)
const FUTURE_INFINITY = TAIEpoch(UNIX_EPOCH, Inf)

const EPOCH_77 = TAIEpoch(1977, 1, 1)


mpSim = ManpowerSimulation( "sim" )

# Setting Attributes.
attribute1 = Attribute( "level" )
setPossibleAttributeValues!( attribute1, ["Junior", "Senior", "Master"] )
setInitialAttributeValues!( attribute1, Dict( "Junior" => 1.0 ) )
attribute2 = Attribute( "branch" )
setPossibleAttributeValues!( attribute2, ["A", "B", "reserve", "none"] )
setInitialAttributeValues!( attribute2, Dict( "none" => 1.0 ) )
attribute3 = Attribute( "isCareer" )
setPossibleAttributeValues!( attribute3, ["yes", "no"] )
setInitialAttributeValues!( attribute3, Dict( "no" => 1.0 ) )
setSimulationAttributes!( mpSim, [attribute1, attribute2, attribute3] )

# Setting Base Nodes.
node1 = BaseNode( "A junior" )
setNodeRequirements!( node1, ("level", "Junior"), ("branch", "A"),
    ("isCareer", "no") )
node2 = BaseNode( "B junior" )
setNodeRequirements!( node2, ("level", "Junior"), ("branch", "B"),
    ("isCareer", "no") )
node3 = BaseNode( "Reserve junior" )
setNodeRequirements!( node3, ("level", "Junior"), ("branch", "reserve"),
    ("isCareer", "no") )
node4 = BaseNode( "A senior" )
setNodeRequirements!( node4, ("level", "Senior"), ("branch", "A"),
    ("isCareer", "yes") )
node5 = BaseNode( "B senior" )
setNodeRequirements!( node5, ("level", "Senior"), ("branch", "B"),
    ("isCareer", "yes") )
node6 = BaseNode( "Master" )
setNodeRequirements!( node6, ("level", "Master"), ("branch", "none"),
    ("isCareer", "yes") )
setNodeTarget!.( [node1, node2, node3, node4, node5, node6],
    [-1, -1, -1, 30, 30, 25] )
setSimulationBaseNodes!( mpSim, [node1, node2, node3, node4, node5, node6] )
setSimulationBaseNodeOrder!( mpSim, Dict( "Master" => 1, "A senior" => 2,
    "B senior" => 2, "A junior" => 3, "B junior" => 3 ) )

# Setting Compound Nodes.
cnode1 = CompoundNode( "Branch A" )
setCompoundNodeComponents!( cnode1, ["A junior", "A senior"] )
cnode2 = CompoundNode( "Branch B" )
setCompoundNodeComponents!( cnode2, ["B junior", "B senior"] )
cnode3 = CompoundNode( "Junior" )
setCompoundNodeComponents!( cnode3,
    ["A junior", "B junior", "Reserve junior"] )
cnode4 = CompoundNode( "Career" )
setCompoundNodeComponents!( cnode4, ["A senior", "B senior", "Master"] )
cnode5 = CompoundNode( "Empty" )
setSimulationCompoundNodes!( mpSim,
    [cnode1, cnode2, cnode3, cnode4 , cnode5] )

# Setting recruitment.
rec1 = Recruitment( "EW" )
setRecruitmentTarget!( rec1, "A junior" )
rec2 = Recruitment( "EW" )
setRecruitmentTarget!( rec2, "B junior" )
rec3 = Recruitment( "EW" )
setRecruitmentTarget!( rec3, "Reserve junior" )

setRecruitmentSchedule!.( [rec1, rec2, rec3], 12 )
setRecruitmentFixed!.( [rec1, rec2, rec3], 10 )
setRecruitmentAgeFixed!.( [rec1, rec2, rec3], 240 )
addSimulationRecruitment!( mpSim, rec1, rec2, rec3 )

# Setting THROUGH Transitions.
ttrans1 = Transition( "Promotion", "A junior", "A senior" )
addTransitionCondition!( ttrans1, MPcondition( "time in node", >=, 24 ) )
ttrans2 = Transition( "Promotion", "B junior", "B senior" )
addTransitionCondition!( ttrans2, MPcondition( "time in node", >=, 24 ) )
ttrans3 = Transition( "Promotion", "A senior", "Master" )
addTransitionCondition!( ttrans3, MPcondition( "time in node", >=, 36 ) )
ttrans4 = Transition( "Promotion", "B senior", "Master" )
addTransitionCondition!( ttrans4, MPcondition( "time in node", >=, 36 ) )
ttrans5 = Transition( "Reserve", "Reserve junior", "A senior" )
addTransitionCondition!( ttrans5, MPcondition( "time in node", >=, 24 ) )
ttrans6 = Transition( "Reserve", "Reserve junior", "B senior" )
addTransitionCondition!( ttrans6, MPcondition( "time in node", >=, 24 ) )

setTransitionSchedule!.( [ttrans1, ttrans2, ttrans3, ttrans4, ttrans5,
    ttrans6], 12 )
setTransitionMaxAttempts!.( [ttrans1, ttrans2, ttrans3, ttrans4, ttrans5,
    ttrans6], 1 )
setTransitionFluxLimits!.( [ttrans1, ttrans2, ttrans3, ttrans4, ttrans5,
    ttrans6], [8, 8, 0, 0, 0, 0], [8, 8, 5, 5, 4, 4] )
addSimulationTransition!( mpSim, ttrans1, ttrans2, ttrans3, ttrans4, ttrans5,
    ttrans6 )

# Setting OUT Transitions.
otrans1 = Transition( "B-", "A junior" )
addTransitionCondition!( otrans1, MPcondition( "time in node", >=, 24 ) )
otrans2 = Transition( "B-", "B junior" )
addTransitionCondition!( otrans2, MPcondition( "time in node", >=, 24 ) )
otrans3 = Transition( "B-", "Reserve junior" )
addTransitionCondition!( otrans3, MPcondition( "time in node", >=, 24 ) )
otrans4 = Transition( "PE", "A senior" )
addTransitionCondition!( otrans4, MPcondition( "tenure", >=, 120 ) )
otrans5 = Transition( "PE", "B senior" )
addTransitionCondition!( otrans5, MPcondition( "tenure", >=, 120 ) )
otrans6 = Transition( "PE", "Master" )
addTransitionCondition!( otrans6, MPcondition( "tenure", >=, 120 ) )

setTransitionMaxAttempts!.( [otrans1, otrans2, otrans3, otrans4, otrans5,
    otrans6], 1 )
setTransitionFluxLimits!.( [otrans1, otrans2, otrans3, otrans4, otrans5,
    otrans6], 0, -1 )
addSimulationTransition!( mpSim, otrans1, otrans2, otrans3, otrans4, otrans5,
    otrans6 )
setSimulationTransitionTypeOrder!( mpSim, Dict( "PE" => 1, "Promotion" => 2,
    "Reserve" => 3, "B-" => 4 ) )

# Miscellaneous configs.
setSimulationLength!( mpSim, 300 )

@test verifySimulation!( mpSim )

@testset "function saveSimulationConfiguration" begin
    setSimulationDatabase!( mpSim, joinpath( "reports", "simDBcopy" ) )
    saveSimulationConfiguration( mpSim )
    setSimulationDatabase!( mpSim, joinpath( "reports", "simDB" ) )
end  # @test "function saveSimulationConfiguration"


@testset "shapes" begin
    inputsize = 20
    poollength = 10
    datalength = 100
    x = rand(Float32, inputsize, poollength, 1, datalength)
    m = TPALSTM(inputsize, 10, poollength, 2, 32, 1)
    @test size(m(x)) == (1, datalength)
    if Flux.CUDA.functional()
        @test size(gpu(m)(gpu(x))) == (1, datalength)
    end
    FluxArchitectures.initialize_bias!(m)
    @test all([all(m.lstm.chain.chain[i].cell.b .== 1) for i in 1:length(m.lstm.chain.chain)])
end

@testset "misc" begin
    @test repr(TPALSTM(10, 10, 10, 2, 32, 1)) == "TPALSTM(10, 10, 10, 2)"
    @test repr(TPALSTM(10, 10, 10, 2, 20, 4)) == "TPALSTM(10, 10, 10, 2, 20, 4)"
end


using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = page_model()
    include("../src/components/MarketDamagesBurke.jl")

    marketdamages = addmarketdamagesburke(m)

    set_param!(m, :MarketDamagesBurke, :rtl_realizedtemperature, readpagedata(m, "test/validationdata/$valdir/rtl_realizedtemperature.csv"))
    set_param!(m, :MarketDamagesBurke, :rcons_per_cap_SLRRemainConsumption, readpagedata(m,"test/validationdata/$valdir/rcons_per_cap_SLRRemainConsumption.csv"))
    set_param!(m, :MarketDamagesBurke, :rgdp_per_cap_SLRRemainGDP, readpagedata(m,"test/validationdata/$valdir/rgdp_per_cap_SLRRemainGDP.csv"))
    set_param!(m, :MarketDamagesBurke, :isatg_impactfxnsaturation, 28.333333333333336)

    p = load_parameters(m)
    p["y_year_0"] = 2015.
    p["y_year"] = Mimi.dim_keys(m.md, :time)
    set_leftover_params!(m, p)

    run(m)

    rcons_per_cap = m[:MarketDamagesBurke, :rcons_per_cap_MarketRemainConsumption]
    rcons_per_cap_compare = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_MarketRemainConsumption.csv")
    @test rcons_per_cap ≈ rcons_per_cap_compare rtol=1e-1

    rgdp_per_cap = m[:MarketDamagesBurke, :rgdp_per_cap_MarketRemainGDP]
    rgdp_per_cap_compare = readpagedata(m, "test/validationdata/$valdir/rgdp_per_cap_MarketRemainGDP.csv")
    @test rgdp_per_cap ≈ rgdp_per_cap_compare rtol=1e-2
end


export PDChar
export readpdftxt

using CodecZlib
import Base.Unicode

mutable struct PDChar
    page::Int
    str::String
    fcoord::Rectangle
    gcoord::Rectangle
end

istext(c::PDChar) = !isdraw(c)
isdraw(c::PDChar) = c.str[1] == '[' && c.str[end] == ']'

function Base.string(c::PDChar)
    join([t.page, t.str, t.coord...], "\t")
end

function readpdftxt(filepath::String)
    if endswith(filepath,".gz")
        lines = open(s -> readlines(GzipDecompressorStream(s)), filepath)
    else
        lines = open(readlines, filepath)
    end
    chars = PDChar[]
    for line in lines
        isempty(line) && continue
        items = split(line, '\t')
        page = parse(Int, items[1])
        str = Unicode.normalize(items[2], :NFKC)
        if str[1] == '[' && str[end] == ']'
            fcoord = Rectangle()
            gcoord = Rectangle()
        else
            coord = map(x -> parse(Float64,x), split(items[3]," "))
            fcoord = Rectangle(coord...)
            coord = map(x -> parse(Float64,x), split(items[4]," "))
            gcoord = Rectangle(coord...)
        end
        c = PDChar(page, str, fcoord, gcoord)
        push!(chars, c)
    end
    chars
end

function saveimages(inpath::String; options=[])
    jar = downloadjar()
    command = `java -classpath $jar paperai.pdfextract.ImageExtractor $inpath $options`
    run(command)
end

function download_jar()
    version = "0.3.1"
    jarfile = normpath(joinpath(@__DIR__,"../deps/pdfextract-$version.jar"))
    if !isfile(jarfile)
        url = "https://github.com/paperai/pdfextract/releases/download/v$version/pdfextract-$version.jar"
        println("Downloading $url")
        download(url, jarfile)
    end
    jarfile
end

function readpdf_v024(filename::String)
    @assert endswith(filename, ".pdf")
    jar = download_jar()
    pdfstr = readstring(`java -classpath $jar PDFExtractor $filename`)
    chars = PDChar[]
    lines = split(pdfstr, "\n")
    for line in lines
        isempty(line) && continue
        items = split(line, '\t')
        page = parse(Int, items[2])
        _type = items[3]
        str = normalize_string(items[4], :NFKC)
        if length(items) > 4
            coord = map(x -> parse(Float64,x), split(items[5]," "))
        else
            coord = Float64[]
        end
        c = PDChar(page, _type, str, coord)
        push!(chars, c)
    end
    chars
end


"""This module defines classes of the package."""
# Author: Henri Gérard <hgerard.proy@gmail.com>
# License: MIT

abstract type RegressionModel end

# Define an object to
struct LinearRegression <: RegressionModel
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function LinearRegression()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

struct LogisticRegression <: RegressionModel
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function LogisticRegression()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

abstract type ParallelParam end

struct Sequential <: ParallelParam
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function Sequential()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

struct Parallel <: ParallelParam
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function Parallel()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

abstract type AlgoParams end

struct OptParams <: AlgoParams
    itmax::Int
    stability::Float64
    learning_rate::Float64
    verbosity::Int
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function OptParams(itmax, stability, learning_rate; verbosity = 0)
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new(itmax, stability, learning_rate, verbosity)
    end
end

struct ProjParams <: AlgoParams
    ITER_MAX::Int
    precision::Float64
    sample::Int
    para_proj::ParallelParam
    para_inter::ParallelParam
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function ProjParams(ITER_MAX, precision, sample; para_proj=Sequential(), para_inter=Sequential())
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new(ITER_MAX, precision, sample, para_proj, para_inter)
    end
end

abstract type AmbiguitySet end

struct DivergenceSet <: AmbiguitySet
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function DivergenceSet()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

struct WassersteinSet <: AmbiguitySet
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function WassersteinSet()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

abstract type GeneralConstraint end

struct PositiveConstraint <: GeneralConstraint
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function PositiveConstraint()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

struct DROConstraint <: GeneralConstraint
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function DROConstraint()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

struct EntropicConstraint <: GeneralConstraint
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function EntropicConstraint()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

abstract type DivergenceConstraint<:GeneralConstraint end

struct KLConstraint <: DivergenceConstraint
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function KLConstraint()
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        return new()
    end
end

struct RobustModel
    descent_direction::Array{Float64,1}
    I0::UnitRange{Int64}
    name::GeneralConstraint
    regressionModel::RegressionModel
    """
    Aggregator of projection

    Parameters
    ----------
    i::Int

    robustModel::RobustModel

    Returns
    -------

    """
    function RobustModel(N, nb_features, ϵ, ambiguity, regressionModel)
        """
        Aggregator of projection

        Parameters
        ----------
        i::Int

        robustModel::RobustModel

        Returns
        -------

        """
        if ambiguity == "KLdivergence"
            descent_direction = [ϵ; 1; zeros(nb_features); ones(N)/N]
            I0 = 1:N
            dim = N+1
            name = KLConstraint()
        elseif ambiguity == "wasserstein"
            descent_direction = [ϵ; zeros(nb_features); ones(N)/N]
            I0 = 1:N^2
            dim = N^2+1
            name = DROConstraint()
        elseif ambiguity == "entropic"
            descent_direction = [1; zeros(nb_features); ones(N)/N]
            I0 = 1:N
            dim = N+1
            name = EntropicConstraint()
        end
        return new(descent_direction, I0, name, regressionModel)
    end
end


# Directions
const directions = [(-1, 0),(1, 0),(0, 1),(0, -1)]
# For readability
PICK_FOOD_1 = 5
PLACE_FOOD_1 = 6
# Perform a move action
function move(env::NatureEnv, player::Int, dir::Int)
    new_pos = env.players[player].pos .+ directions[dir]
    outofbounds(env, new_pos) && return
    env.players[player].pos = new_pos
end
# Perform a food action
function food(env::NatureEnv, p::Int, idx::Int)
    # [ pick1, place1, pick2, place2, ... pickf, placef]
    pick  = idx % 2 == 1
    place = idx % 2 == 0
    food_type = floor(Int, ((idx - 1) / 2)) + 1
    fframe   = env.food_frames[food_type]
    player   = env.players[p]
    num_food = fframe[player.pos...]
    place_amount = 0.5

    if pick && num_food > 0
        player.food_counts = player.food_counts .+ Tuple(num_food * onehot(food_type, env.food_types))
        for (placed_player, placed_count) in env.place_record[player.pos][food_type]
            if placed_player != p
                env.exchanges[food_type] += placed_count
            end
        end
        fframe[player.pos...] = 0
        env.place_record[player.pos][food_type] = small_dd_builder()    elseif place && player.food_counts[food_type] >= place_amount
        player.food_counts = Tuple(player.food_counts .-  (place_amount .* onehot(food_type, env.food_types)))
        fframe[player.pos...] += place_amount
        env.place_record[player.pos][food_type][p] = env.place_record[player.pos][food_type][p] + place_amount
    end
end

function comm(env::NatureEnv, p::Int, symbol::Int)
    push!(env.comms[env.step], (env.players[p].pos..., p, symbol))
end


mutable struct HistoryODEIntegrator{algType,IIP,uType,tType,tdirType,ksEltype,SolType,CacheType} <: AbstractODEIntegrator{algType,IIP,uType,tType}
  sol::SolType
  u::uType
  k::ksEltype
  t::tType
  dt::tType
  uprev::uType
  tprev::tType
  alg::algType
  dtcache::tType
  tdir::tdirType
  saveiter::Int
  saveiter_dense::Int
  cache::CacheType
end

function (integrator::HistoryODEIntegrator)(t, deriv::Type=Val{0}; idxs=nothing)
    OrdinaryDiffEq.current_interpolant(t, integrator, idxs, deriv)
end

function (integrator::HistoryODEIntegrator)(
    val::AbstractArray, t::Union{Number,AbstractArray}, deriv::Type=Val{0}; idxs=nothing
)
    OrdinaryDiffEq.current_interpolant!(val, t, integrator, idxs, deriv)
end

mutable struct DDEIntegrator{algType,IIP,uType,tType,P,eigenType,tTypeNoUnits,tdirType,
                             ksEltype,SolType,F,CacheType,IType,FP,O,dAbsType,dRelType,H,
                             tstopsType,discType,FSALType,EventErrorType,CallbackCacheType} <: AbstractDDEIntegrator{algType,IIP,uType,tType}
    sol::SolType
    u::uType
    k::ksEltype
    t::tType
    dt::tType
    f::F
    p::P
    uprev::uType
    uprev2::uType
    tprev::tType
    prev_idx::Int
    prev2_idx::Int
    fpsolver::FP
    order_discontinuity_t0::Int
    "Discontinuities tracked by callback."
    tracked_discontinuities::Vector{Discontinuity{tType,Int}}
    discontinuity_interp_points::Int
    discontinuity_abstol::dAbsType
    discontinuity_reltol::dRelType
    "Future time stops for propagated discontinuities."
    tstops_propagated::tstopsType
    "Future propagated discontinuities."
    d_discontinuities_propagated::discType
    alg::algType
    dtcache::tType
    dtchangeable::Bool
    dtpropose::tType
    tdir::tdirType
    eigen_est::eigenType
    EEst::tTypeNoUnits
    qold::tTypeNoUnits
    q11::tTypeNoUnits
    erracc::tTypeNoUnits
    dtacc::tType
    success_iter::Int
    iter::Int
    saveiter::Int
    saveiter_dense::Int
    cache::CacheType
    callback_cache::CallbackCacheType
    kshortsize::Int
    force_stepfail::Bool
    last_stepfail::Bool
    just_hit_tstop::Bool
    event_last_time::Int
    vector_event_last_time::Int
    last_event_error::EventErrorType
    accept_step::Bool
    isout::Bool
    reeval_fsal::Bool
    u_modified::Bool
    opts::O
    destats::DDEStats
    history::H
    integrator::IType
    fsalfirst::FSALType
    fsallast::FSALType

    # incomplete initialization without fsalfirst and fsallast
    function DDEIntegrator{algType,IIP,uType,tType,P,eigenType,tTypeNoUnits,
                           tdirType,ksEltype,SolType,F,CacheType,IType,FP,
                           O,dAbsType,dRelType,H,tstopsType,discType,
                           FSALType,EventErrorType,CallbackCacheType}(
                               sol,u,k,t,dt,f,p,uprev,uprev2,tprev,prev_idx,prev2_idx,
                               fpsolver,order_discontinuity_t0,tracked_discontinuities,
                               discontinuity_interp_points,discontinuity_abstol,discontinuity_reltol,
                               tstops_propagated, d_discontinuities_propagated,
                               alg,dtcache,dtchangeable,dtpropose,tdir,eigen_est,EEst,qold,
                               q11,erracc,dtacc,success_iter,iter,saveiter,saveiter_dense,
                               cache,callback_cache,kshortsize,force_stepfail,last_stepfail,
                               just_hit_tstop,event_last_time,vector_event_last_time,last_event_error,
                               accept_step,isout,reeval_fsal,u_modified,opts,destats,
                               history,integrator) where
      {algType,IIP,uType,tType,P,eigenType,tTypeNoUnits,tdirType,ksEltype,SolType,F,
       CacheType,IType,FP,O,dAbsType,dRelType,H,tstopsType,discType,
       FSALType,EventErrorType,CallbackCacheType}

      new{algType,IIP,uType,tType,P,eigenType,tTypeNoUnits,tdirType,ksEltype,SolType,F,
          CacheType,IType,FP,O,dAbsType,dRelType,H,tstopsType,discType,FSALType,
          EventErrorType,CallbackCacheType}(
            sol,u,k,t,dt,f,p,uprev,uprev2,tprev,prev_idx,prev2_idx,fpsolver,
            order_discontinuity_t0,tracked_discontinuities,discontinuity_interp_points,
            discontinuity_abstol,discontinuity_reltol,tstops_propagated,
            d_discontinuities_propagated,alg,dtcache,dtchangeable,dtpropose,tdir,
            eigen_est,EEst,qold,q11,erracc,dtacc,success_iter,iter,saveiter,saveiter_dense,
            cache,callback_cache,kshortsize,force_stepfail,last_stepfail,just_hit_tstop,
            event_last_time,vector_event_last_time,last_event_error,accept_step,isout,
            reeval_fsal,u_modified,opts,destats,history,integrator)
    end
end

function (integrator::DDEIntegrator)(t, deriv::Type=Val{0}; idxs=nothing)
    OrdinaryDiffEq.current_interpolant(t, integrator, idxs, deriv)
end

function (integrator::DDEIntegrator)(val::AbstractArray, t::Union{Number,AbstractArray},
                                     deriv::Type=Val{0}; idxs=nothing)
    OrdinaryDiffEq.current_interpolant!(val, t, integrator, idxs, deriv)
end


struct SDDEProblem{uType,tType,lType,lType2,isinplace,P,NP,F,G,H,K,ND} <:
                          AbstractSDDEProblem{uType,tType,lType,isinplace,ND}
  f::F
  g::G
  u0::uType
  h::H
  tspan::tType
  p::P
  noise::NP
  constant_lags::lType
  dependent_lags::lType2
  kwargs::K
  noise_rate_prototype::ND
  seed::UInt64
  neutral::Bool
  order_discontinuity_t0::Rational{Int}

  @add_kwonly function SDDEProblem{iip}(f::AbstractSDDEFunction{iip}, g, u0, h, tspan, p = NullParameters();
                                       noise_rate_prototype = nothing, noise= nothing, seed = UInt64(0),
                                       constant_lags = (), dependent_lags = (),
                                       neutral = f.mass_matrix !== I && det(f.mass_matrix) != 1,
                                       order_discontinuity_t0 = 0//1,
                                       kwargs...) where {iip}
    _tspan = promote_tspan(tspan)
    new{typeof(u0),typeof(_tspan),typeof(constant_lags),typeof(dependent_lags),isinplace(f),
        typeof(p),typeof(noise),typeof(f),typeof(g),typeof(h),typeof(kwargs),typeof(noise_rate_prototype)}(
        f, g, u0, h, _tspan, p, noise, constant_lags, dependent_lags, kwargs, noise_rate_prototype, seed, neutral, order_discontinuity_t0)
  end

  function SDDEProblem{iip}(f::AbstractSDDEFunction{iip}, g , h, tspan::Tuple, p = NullParameters();
                           order_discontinuity_t0 = 1//1, kwargs...) where iip
    SDDEProblem{iip}(f, g, h(p, first(tspan)), h, tspan, p;
                    order_discontinuity_t0 = max(1//1, order_discontinuity_t0), kwargs...)
  end

  function SDDEProblem{iip}(f,g, args...; kwargs...) where iip
    SDDEProblem{iip}(convert(SDDEFunction{iip}, f,g),g, args...; kwargs...)
  end
end

SDDEProblem(f, g, args...; kwargs...) =
  SDDEProblem(convert(SDDEFunction, f, g), g, args...; kwargs...)

SDDEProblem(f::AbstractSDDEFunction, args...; kwargs...) =
  SDDEProblem{isinplace(f)}(f, args...; kwargs...)


"""
Ledgers

This package provides support for general financial ledgers

See README.md for the full documentation

Copyright 2019-2020, Eric Forgy, Scott P. Jones and other contributors

Licensed under MIT License, see LICENSE.md
"""
module Ledgers

using UUIDs, StructArrays, AbstractTrees
import Instruments
using Instruments: Position
# Although importing Instruments is enough to extend the methods below, we still import
# them explicitly so we can use them like `Ledgers.position` instead of `Instruments.position`.
# Unfortunately, `Base` also exports `position`.
import Instruments: instrument, currency, symbol, position, amount

export Identifier, AccountId, AbstractAccount, AccountWrapper, AccountNumber
export LedgerAccount, Ledger, Entry, Account, AccountGroup, ledgeraccount, add_account, ledger
export id, balance, credit!, debit!, post!, add_entry!
# Although the methods below are exported by Instruments, we export them explicitly
# because we do not expect users to use `Instruments` directly.
export instrument, currency, symbol, position, amount

include("accounts.jl")

abstract type AbstractLedgerAccount{P} <: AbstractAccount{P} end

struct Entry{P,A<:AbstractLedgerAccount{P}}
    debit::A
    credit::A
    amount::P
end

struct LedgerId <: Identifier
    value::UUID
end

LedgerId() = LedgerId(uuid4())

struct Ledger{P <: Position,A<:AbstractLedgerAccount{P}}
    id::LedgerId
    accounts::Dict{AccountId,A}
    entries::Vector{Entry{P}}
end

function Ledger(accounts::AbstractVector{<:AbstractLedgerAccount{P}}, entries::AbstractVector{Entry{P}}; id=LedgerId()) where {P <: Position}
    ledger = Ledger(P)
    add_account!.(Ref(ledger), accounts)
    add_entry!.(Ref(ledger), entries)
    return ledger
end

Base.iterate(ledger::Ledger, state=1) = iterate(ledger.entries, state)

Ledger(::Type{P}) where {P <: Position} = Ledger{P,LedgerAccount{P}}(LedgerId(),Dict{AccountId,LedgerAccount{P}}(), Vector{Entry{P}}())

mutable struct LedgerAccount{P} <: AbstractLedgerAccount{P}
    id::AccountId
    balance::P
    ledger::Ledger{P}
end

function LedgerAccount(balance::P, ledger) where {P <: Position}
    acc = LedgerAccount{P}(AccountId(), balance, ledger)
    add_account!(ledger, acc)
    acc
end

LedgerAccount(::Type{P}, ledger) where {P <: Position} =
    LedgerAccount(P(0), ledger)

ledgeraccount(acc::LedgerAccount) = acc
ledgeraccount(acc::AbstractAccount) = ledgeraccount(acc.account)

ledger(acc::LedgerAccount) = acc.ledger
ledger(acc::AbstractAccount) = ledger(ledgeraccount(acc))

balance(acc::LedgerAccount) = acc.balance
balance(acc::AbstractAccount) = balance(ledgeraccount(acc))
id(acc::LedgerAccount) = acc.id

function add_account!(ledger::Ledger{P}, acc::AbstractAccount{P}) where {P <: Position}
    acc = ledgeraccount(acc)
    if id(acc) ∈ keys(ledger.accounts)
        warn("Account already in ledger.")
        return
    end
    ledger.accounts[acc.id] = acc
end

function add_accounts!(ledger::Ledger{P}, ag::AccountGroup) where P
    for acc in ag
        add_account!(ledger, acc)
    end
    ledger
end

function add_entry!(ledger::Ledger, entry::Entry)
    (id(entry.credit) ∈ keys(ledger.accounts)
        && id(entry.debit) ∈ keys(ledger.accounts)) || error("Unknown account in entry.")
    push!(ledger.entries, entry)
    post!(entry)
end

Instruments.symbol(::Type{A}) where {P, A<:AbstractAccount{P}} = symbol(P)

Instruments.currency(::Type{A}) where {P, A<:AbstractAccount{P}} = currency(P)

Instruments.instrument(::Type{A}) where {P, A<:AbstractAccount{P}} = instrument(P)

Instruments.position(::Type{A}) where {P, A<:AbstractAccount{P}} = P

Instruments.amount(acc::AbstractAccount) = amount(balance(acc))

debit!(acc::LedgerAccount, amt::Position) = (acc.balance -= amt)

credit!(acc::LedgerAccount, amt::Position) = (acc.balance += amt)

function post!(entry::Entry)
    debit!(ledgeraccount(entry.debit), entry.amount)
    credit!(ledgeraccount(entry.credit), entry.amount)
end

Base.getindex(ledger::Ledger, ix) = ledger.entries[ix]

Base.getindex(ledger::Ledger, id::AccountId) =
    ledger.accounts[id]

Base.getindex(grp::AccountGroup, id::AccountId) =
    grp.accounts[id]

# struct EntityId <: Identifier
#     value::UUID
# end

# EntityId() = EntityId(uuid4())

# struct Entity
#     id::EntityId
#     name::String
#     ledgers::Dict{Type{<:Position},Ledger{<:AbstractAccount}}
# end

# const chartofaccounts = Dict{String,AccountGroup{<:Cash}}()

# iscontra(a::AccountGroup) = !isequal(a,a.parent) && !isequal(a.parent,getledger(a)) && !isequal(a.parent.isdebit,a.isdebit)

# function loadchart(ledgername,ledgercode,csvfile)
#     data, headers = readdlm(csvfile,',',String,header=true)
#     nrow,ncol = size(data)

#     ledger = add(AccountGroup(ledgername,ledgercode))
#     for i = 1:nrow
#         row = data[i,:]
#         number = row[1]
#         name = row[2]
#         parent = chartofaccounts[row[3]]
#         isdebit = isequal(row[4],"Debit")
#         add(AccountGroup(parent,name,number,isdebit))
#     end
#     return ledger
# end

# function trim(a::AccountGroup,newparent::AccountGroup=AccountGroup(a.parent.name,a.parent.number))
#     newaccount = isequal(a,a.parent) ? newparent : AccountGroup(newparent,a.name,a.number,a.isdebit,a.balance)
#     for subaccount in a.accounts
#         balance(subaccount).a_print_account

function _print_account(io::IO, acc)
    iobuff = IOBuffer()
    isempty(acc.number.value) || print(iobuff,"[$(acc.number)] ")
    print(iobuff,acc.name,": ")
    acc.isdebit ? print(iobuff,balance(acc)) : print(iobuff,-balance(acc))
    print(io,String(take!(iobuff)))
end

Base.show(io::IO, id::Identifier) = print(io, id.value)

Base.show(io::IO, number::AccountNumber) = print(io, number.value)

Base.show(io::IO, acc::LedgerAccount) = print(io, "$(string(id(acc))): $(balance(acc))")

Base.show(io::IO, acc::AbstractAccount) = _print_account(io, acc)

Base.show(io::IO, acc::AccountGroup) = print_tree(io, acc)

Base.show(io::IO, entry::Entry) = print_tree(io, entry)

Base.show(io::IO, ledger::Ledger) = print_tree(io, ledger)

Base.show(io::IO, a::Vector{<:AbstractAccount}) = print_tree(io, a)

AbstractTrees.printnode(io::IO, acc::Ledger{P}) where {P <: Position} =
    print(io, "$(symbol(P)) Ledger: [$(acc.id)]")

AbstractTrees.children(acc::Ledger) = collect(values(acc.accounts))

AbstractTrees.printnode(io::IO, acc::AccountGroup) = _print_account(io, acc)

AbstractTrees.children(acc::AccountGroup) = vcat(collect(values(acc.subgroups)), collect(values(acc.accounts)))

AbstractTrees.printnode(io::IO, b::AbstractVector{<:AbstractAccount}) =
    isempty(b) ? print(io, "Accounts: None") : print(io, "Accounts:")

AbstractTrees.children(b::AbstractVector{<:AbstractAccount}) = b

AbstractTrees.children(entry::Entry) = [entry.debit, entry.credit]

AbstractTrees.printnode(io::IO, ::Entry) = print(io, "Entry:")

end # module Ledgers


# by lazarusA # HIDE
# using GLMakie # HIDE
using GLMakie, NetCDF
using YAXArrays.Datasets: open_mfdataset
GLMakie.activate!() # HIDE
let
    dataset = open_mfdataset("./data/ETOPO1_halfdegree.nc")
    lon = dataset.lon.values
    lat = dataset.lat.values
    data = dataset.ETOPO1avg.data[:,:]

    fig = Figure(resolution = (1250,700),fontsize = 25, backgroundcolor = :black)
    ax = Axis3(fig, aspect = (1,0.6,0.1),azimuth = -0.65π, elevation = 0.225π)
    pltobj=surface!(ax,lon,lat, data, colormap = :starrynight, backlight = 1.0f0,
        colorrange =(-6500,5500))
    cbar  =Colorbar(fig,pltobj,label = "ETOPO1 [m]",ticklabelsize = 20,
        tickalign =1, vertical = true, ticksize=15, labelcolor = :white,
        ticklabelcolor = :white, tickcolor = :grey90, height = Relative(0.5))
    hidespines!(ax)
    hidedecorations!(ax)
    fig[1,1] = ax
    fig[1,2] = cbar
    fig
    save(joinpath(@__DIR__, "output", "topography.png"), fig, px_per_unit = 2.0) # HIDE
end

using Pkg # HIDE
Pkg.status(["GLMakie","NetCDF", "YAXArrays"]) # HIDE


using Test
using Dorothy

@DorothyAll()

@testset "PDB" begin

	@testset "Read" begin
		m1 = readf("$(Dorothy.datapath)/1BTL.pdb")
		@test length(m1) == 2236
		(a,b,c), (α,β,γ) = pbcgeometry(m1.header.cell)
		@test [a,b,c] ≈ [43.1, 64.4, 91.2]
		@test m1.header.title ==
				"CRYSTAL STRUCTURE OF ESCHERICHIA COLI TEM1 BETA-LACTAMASE" *
				" AT 1.8 ANGSTROMS RESOLUTION"
		@test m1.resids[1] == 26
		@test m1.resnames[1] == "HIS"
		@test m1.names[1] == "N"
		@test m1.ids[1] == 1
		@test m1.chainids[1] == "A"
		@test m1.occupancies[1] == 1.0
		@test m1.bfactors[1] == 14.51
		@test m1.R[end] ≈ [-12.176, -11.865, 48.715]
		@test ! haskey(m1, :V)
	end

	@testset "Write" begin
		m1 = readf("$(Dorothy.datapath)/1BTL.pdb")
		io = IOBuffer()
		write(specify(io, "structure/x-pdb"), m1)
		seekstart(io)
		m2 = read(specify(io, "structure/x-pdb"))
		@test length(m2) == 2236
		(a,b,c), (α,β,γ) = pbcgeometry(m2.header.cell)
		@test [a,b,c] ≈ [43.1, 64.4, 91.2]
		@test m2.header.title ==
				"CRYSTAL STRUCTURE OF ESCHERICHIA COLI TEM1 BETA-LACTAMASE" *
				" AT 1.8 ANGSTROMS RESOLUTION"
		@test m2.resids[1] == 26
		@test m2.resnames[1] == "HIS"
		@test m2.names[1] == "N"
		@test m2.ids[1] == 1
		@test m2.chainids[1] == "A"
		@test m2.occupancies[1] == 1.0
		@test m2.bfactors[1] == 14.51
		@test m2.R[end] ≈ [-12.176, -11.865, 48.715]
		@test ! haskey(m2, :V)
	end

	@testset "Convert" begin
		m1 = readf("$(Dorothy.datapath)/1BTL.pdb")
		io = IOBuffer()
		write(specify(io, "structure/x-gro"), m1)
		seekstart(io)
		m2 = read(specify(io, "structure/x-gro"))
		@test length(m2) == 2236
		(a,b,c), (α,β,γ) = pbcgeometry(m2.header.cell)
		@test [a,b,c] ≈ [43.1, 64.4, 91.2]
		@test m2.header.title ==
				"CRYSTAL STRUCTURE OF ESCHERICHIA COLI TEM1 BETA-LACTAMASE" *
				" AT 1.8 ANGSTROMS RESOLUTION"
		@test m2.resids[1] == 26
		@test m2.resnames[1] == "HIS"
		@test m2.names[1] == "N"
		@test m2.ids[1] == 1
		@test m2.R[end] ≈ [-12.18, -11.87, 48.72]
		@test ! haskey(m2, :V)
	end

end # @testset


module TetMesh

using StaticArrays: MVector, @MVector, SVector
using WriteVTK: vtk_grid, MeshCell, VTKCellTypes
using Lazy: @as
using Base: setindex

struct NodeInfo
    coords::MVector{2, Float64}
    adjacent_nodes::Vector{Int}
end

struct ElementInfo{NodesPerFace, InternalNodes}
    control_nodes::MVector{3, Int}
    faces::MVector{3, Int}
    face_nodes::MVector{3, MVector{NodesPerFace, Int}}
    internal_nodes::MVector{InternalNodes, Int}
    adjacent_elements::Vector{Int}
end

struct FaceDetails{NodesPerFace}
    number::Int
    element::Int
    nodes::MVector{NodesPerFace, Int}
end

struct BoundaryRepresentation{NodesPerFace}
    nodes::Vector{Int}
    faces::Vector{FaceDetails{NodesPerFace}}
end

struct Mesh{NodesPerFace, InternalNodes, BoundaryNodesPerFace}
    nodes::Vector{NodeInfo}
    elements::Vector{ElementInfo{NodesPerFace, InternalNodes}}
    boundaries::Vector{BoundaryRepresentation{BoundaryNodesPerFace}}
    boundary_tags::Vector{Vector{String}}
end

function deserialize(stream::IO)
    max_element_adjacencies = read(stream, Int)
    max_node_adjacencies = read(stream, Int)
    NodesPerFace = read(stream, Int)
    InternalNodes = read(stream, Int)

    # Read nodes
    count = read(stream, Int)
    nodes = [deserialize(stream, NodeInfo) for i in 1:count]

    # Read elements
    count = read(stream, Int)
    elements = [
        deserialize(stream, ElementInfo{NodesPerFace, InternalNodes})
        for i in 1:count
    ]

    # Read boundaries
    count = read(stream, Int)
    boundaries = [
        deserialize(stream, BoundaryRepresentation{NodesPerFace+2})
        for i in 1:count
    ]

    # Read tags
    count = read(stream, Int)
    tags = [deserialize(stream, Vector{String}) for i in 1:count]

    Mesh{NodesPerFace, InternalNodes, NodesPerFace+2}(nodes, elements, boundaries, tags)
end

"""
Returns three items:
  - points: An array of the coordinates of vertices of triangles in the mesh;
  - elements: An array of WriteVTK.MeshCell with indices of the vertices of
    triangles in `points`;
  - node_map: `node_map[i]` maps a node `i` in the triangle mesh to a node index
    in the full (possibly higher-order) mesh.
"""
function points_and_cells(msh::Mesh)
    node_indices = zeros(Int, length(msh.nodes))
    coords = Vector{SVector{2, Float64}}(undef, length(msh.nodes))
    vertices = NTuple{3, Int}[]
    node_map = Int[]

    node_index = 1
    for element in msh.elements
        verts = (0, 0, 0)
        i = 1
        for node in (element.control_nodes .+ 1)
            if node_indices[node] == 0
                node_indices[node] = node_index
                coords[node_index] = msh.nodes[node].coords
                node_index += 1
                push!(node_map, node)
            end
            verts = setindex(verts, node_indices[node], i)
            i += 1
        end
        push!(vertices, verts)
    end

    nnodes = node_index - 1
    resize!(coords, nnodes)

    (@as v coords reinterpret(Float64, v) reshape(v, 2, :)),
    [MeshCell(VTKCellTypes.VTK_TRIANGLE, verts) for verts in vertices],
    node_map
end

function initialize_vtk(target, msh)
    points, cells, node_map = points_and_cells(msh)
    vtk_grid(target, points, cells), node_map
end

function add_scalar_data(target, msh::Mesh, node_map::Vector{Int}, data::AbstractVector, name)
    @assert length(data) == length(msh.nodes)
    target[name] = @view data[node_map]
    target
end

function add_vector_data(target, msh::Mesh, node_map::Vector{Int}, data::AbstractMatrix, name)
    @assert size(data, 2) == length(msh.nodes)
    target[name] = @view data[:, node_map]
    target
end

function deserialize(stream::IO, ::Type{NodeInfo})
    coords = deserialize(stream, MVector{2, Float64})
    adjacent = deserialize(stream, Vector{Int})
    NodeInfo(coords, adjacent)
end

function deserialize(stream::IO, ::Type{MVector{N, T}}) where {N, T}
    x = @MVector zeros(T, N)
    read!(stream, x)
end

function deserialize(stream::IO, ::Type{Vector{T}}) where T
    sz = read(stream, Int)
    x = Vector{T}(undef, sz)
    read!(stream, x)
end

function deserialize(stream::IO, ::Type{ElementInfo{NodesPerFace, InternalNodes}}) where {NodesPerFace, InternalNodes}
    control_nodes = deserialize(stream, MVector{3, Int})
    faces = deserialize(stream, MVector{3, Int})
    face_nodes = @MVector [deserialize(stream, MVector{NodesPerFace, Int}) for i in 1:3]
    internal_nodes = deserialize(stream, MVector{InternalNodes, Int})
    adjacent_elements = deserialize(stream, Vector{Int})
    ElementInfo(control_nodes, faces, face_nodes, internal_nodes, adjacent_elements)
end

function deserialize(stream::IO, ::Type{BoundaryRepresentation{N}}) where N
    nodes = deserialize(stream, Vector{Int})
    faces = Vector{FaceDetails{N}}(undef, read(stream, Int))
    for i in eachindex(faces)
        faces[i] = deserialize(stream, FaceDetails{N})
    end
    BoundaryRepresentation{N}(nodes, faces)
end

function deserialize(stream::IO, ::Type{FaceDetails{N}}) where N
    number = read(stream, Int)
    element = read(stream, Int)
    nodes = deserialize(stream, MVector{N, Int})
    FaceDetails{N}(number, element, nodes)
end

function deserialize(stream::IO, ::Type{Vector{String}})
    count = read(stream, Int)
    [deserialize(stream, String) for i in 1:count]
end

function deserialize(stream::IO, ::Type{String})
    count = read(stream, Int)
    String(read!(stream, Vector{UInt8}(undef, count)))
end

end # module TetMesh


export EulerStepper

mutable struct EulerStepper{N,EOS,K,MD,M,ED,E,D,B} <: Stepper{N}
    eos::EOS
    positions::Vector{SVectorF{N}}
    velocities::Velocities{N}
    mass_component::M
    energy_component::E
    diffusion_component::D
    boundary_component::B
    pressures::Pressures
    time_derivatives::TimeDerivatives{N,MD,ED}

    function EulerStepper(
        positions::AbstractVector{SVectorF{N}},
        velocities::Velocities{N},
        mass_component::M,
        energy_component::E,
        diffusion_component::D,
        boundary_component::B,
        eos::EOS,
    ) where {
        N,
        EOS<:EquationOfState,
        K<:Kernel{N},
        MD,
        M<:MassComponent{MD,K},
        ED,
        E<:EnergyComponent{ED},
        D<:DiffusionComponent,
        B<:BoundaryComponent{N},
    }
        pressures = Pressures(mass_component, energy_component, eos)
        time_derivatives =
            TimeDerivatives(velocities, mass_component, energy_component)
        updatefluid!(
            time_derivatives,
            mass_component,
            energy_component,
            diffusion_component,
            pressures,
            velocities,
            positions,
            eos,
            boundary_component,
        )
        new{N,EOS,K,MD,M,ED,E,D,B}(
            eos,
            positions,
            velocities,
            mass_component,
            energy_component,
            diffusion_component,
            boundary_component,
            pressures,
            time_derivatives,
        )
    end
end

Base.show(
    io::IO,
    ::Type{EulerStepper{N,EOS,K,MD,M,ED,E,D,B}},
) where {N,EOS,K,MD,M,ED,E,D,B} = print(io, "EulerStepper")

function Base.show(
    io::IO,
    ::EulerStepper{N,EOS,K,MD,M,ED,E,D,B},
) where {N,EOS,K,MD,M,ED,E,D,B}
    print(
        io,
        """
        EulerStepper{
            $(N)D
            $(EOS)
            $(K)
            $(M)
            $(E)
            $(D)
            $(B)
        }""",
    )
end

get_positions(stepper::EulerStepper) = stepper.positions
get_velocities(stepper::EulerStepper) = get_velocities(stepper.velocities)
get_mass_densities(stepper::EulerStepper) =
    get_mass_densities(stepper.mass_component)
get_kernel_widths(stepper::EulerStepper) =
    get_kernel_widths(stepper.mass_component)
get_specific_energies(stepper::EulerStepper) =
    get_specific_energies(stepper.energy_component)
get_pressures(stepper::EulerStepper) = get_pressures(stepper.pressures)

function step!(stepper::EulerStepper, Δt::Number)
    evolvevariables!(
        stepper.velocities,
        stepper.mass_component,
        stepper.energy_component,
        stepper.time_derivatives,
        Δt,
    )
    evolvepositions!(stepper.positions, stepper.velocities, Δt)
    evolveboundaries!(stepper.boundary_component, Δt)
    updatefluid!(
        stepper.time_derivatives,
        stepper.mass_component,
        stepper.energy_component,
        stepper.diffusion_component,
        stepper.pressures,
        stepper.velocities,
        stepper.positions,
        stepper.eos,
        stepper.boundary_component,
    )
end


using KernelMethods.KMap: fftraversal, size_criterion, dnet
import KernelMethods.KMap: centroid!
using SimilaritySearch
using NearNeighborGraph
using TextModel

export codebook, FFTraversal, ApproxKDCentroids

abstract type Clustering end

struct FFTraversal <: Clustering
    numcenters::Int
    k::Int
end

# Approximate k distance centroids
struct ApproxKDCentroids <: Clustering
    numcenters::Int
    k::Int
    maxiter::Int
    tol::Float64
    recall::Float64
end

ApproxKDCentroids(numcenters, k; maxiter=10, tol=0.001, recall=0.9) = ApproxKDCentroids(numcenters, k, maxiter, tol, recall)

function create_distance(::Type{VBOW})
    #cosine_distance
    angle_distance
end

function create_distance(::Type{DenseCosine{Float32}})
    #CosineDistance()
    AngleDistance()
end

function create_distance(::Type{Vector{F}}) where {F <: AbstractFloat}
    L2Distance()
end

"""
Creates an nearest neighbor index to speedup `codebook`
"""
function create_index(db, recall)
    T = typeof(db[1])
    # return Sequential(db, create_distance(T))
    index = LocalSearchIndex(T, create_distance(T), search=BeamSearch(), neighborhood=LogSatNeighborhood(), recall=recall)
    NearNeighborGraph.fit!(index, db)
    index
end

"""
Creates a sampling of the dataset using the farthest first algorithm
(an approximation of the kcenter problem)

"""
function codebook(algo::FFTraversal, X::AbstractVector{T}) where {T <: Union{DenseCosine{Float32},VBOW}}
    n = length(X)
    rlist = Int[]
    # dmax = 0.0
    function callback(pivID, _dmax)
        push!(rlist, pivID)
        # dmax = _dmax
    end

    xcodes = KnnResult[KnnResult(algo.k) for i in 1:n]
    function callbackdist(pivID, objID, d)
        push!(xcodes[objID], pivID, d)
    end

    fftraversal(callback, X, CosineDistance(), size_criterion(algo.numcenters), callbackdist)
    sort!(rlist)

    centers = X[rlist]
    H = Dict(r => i for (i, r) in enumerate(rlist))
    codes = [[H[p.objID] for p in xcodes[i]] for i in 1:n]
    centers, codes, rlist
end

function associate_centroids_and_score(algo, C, X, codes, distances)
    index = create_index(C, algo.recall)

    Threads.@threads for objID in 1:length(X)
        res = search(index, X[objID], KnnResult(algo.k))
        codes[objID] = [p.objID for p in res]
        distances[objID] = last(res).dist
    end

    mean(distances)
end

const _empty_int_array = Int[]
"""
Computes a clustering based on random sampling and an approximate nearest neighbor index, the procedure is iterative like the kmeans algorithm but instead of use the nearest neighbor `codebook` uses `k` nearest neighbors to converge and an approximate nn algorithm.

"""
function codebook(algo::ApproxKDCentroids, X::AbstractVector{T}) where {T <: Union{DenseCosine{Float32},VBOW}}
    n = length(X)
    rlist = rand(1:n, algo.numcenters) |> Set |> collect   # initial selection
    codes = [_empty_int_array for i in 1:n]
    C = X[rlist]
    iter = 0
    distances = zeros(Float64, n)
    scores = [typemax(Float64), associate_centroids_and_score(algo, C, X, codes, distances)]
    @assert !isnan(scores[end]) "ERROR invalid score $scores"
    while iter < algo.maxiter && abs(scores[end-1] - scores[end]) > algo.tol
        iter += 1
        info("*** starting iteration: $iter; scores: $scores ***")
        invindex = [Int[] for i in 1:algo.numcenters]
        for (objID, plist) in enumerate(codes)
            for refID in plist
                push!(invindex[refID], objID)
            end
        end

        info("*** computing centroids ***")
        Threads.@threads for i in 1:length(invindex)
            plist = invindex[i]
            # C[i] can be empty because we could be using approximate search
            if length(plist) > 0
                C[i] = centroid!(X[plist])
            end
        end

        info("*** computing $(algo.k) nearest references ***")
        s = associate_centroids_and_score(algo, C, X, codes, distances)

        push!(scores, s)
        @assert !isnan(scores[end]) "ERROR invalid score $scores"
        info("*** new score with $(algo.k) references: $scores ***")
    end

    info("*** finished computation of $(algo.k) references, scores: $scores ***")
    C, codes, rlist
end


using Caesar, RoME
using TransformUtils, Rotations, CoordinateTransformations
using PyCall, PyLCM
using LibBSON
# using Colors, FixedPointNumbers #needed if pushing to DrakeVisualizer

# generate python bindings for lcmtypes and add them to PYTHONPATH
@show lcmtpath = joinpath(dirname(@__FILE__),"lcmtypes")
run(`lcm-gen -p --ppath $(lcmtpath) $(lcmtpath)/hauv_submap_points_t.lcm`)
println("Adding lcmtypes dir to Python path: $(lcmtpath)")
unshift!(PyVector(pyimport("sys")["path"]),lcmtpath)
@pyimport hauv

function lcmpointcloudmsg!(vc, slaml::SLAMWrapper,
                          msgdata  )
  msg = hauv.submap_points_t[:decode](msgdata)

  vert = getVert(slaml.fg, slaml.lastposesym, api=IncrementalInference.dlapi) # fetch from database

  # 2d arrays of points and colors (from LCM data into arrays{arrays})
  ptarr = [[pt[1], pt[2], pt[3]] for pt in msg[:points_local]]
  carr = [[UInt8(c.data[1]),UInt8(c.data[2]),UInt8(c.data[3])] for c in msg[:colors]]

  # push to mongo (using BSON as a quick fix)
  # (for deserialization, see src/DirectorVisService.jl:cachepointclouds!)
  serialized_point_cloud = BSONObject(Dict("pointcloud" => ptarr))
  appendvertbigdata!(slaml.fg, vert, "BSONpointcloud", string(serialized_point_cloud).data)
  serialized_colors = BSONObject(Dict("colors" => carr))
  appendvertbigdata!(slaml.fg, vert, "BSONcolors", string(serialized_colors).data)

  # push to DrakeVisualizer
  #=
  position = Translation(msg[:pos]...)
  orientation = Rotations.Quat(msg[:orientation]...)
  Tf = position ∘ LinearMap(orientation)
  vsym = Symbol(vert.label)
  setgeometry!(vc[:submaps][vsym], Triad())
  settransform!(vc[:submaps][vsym], Tf)

  pointcloud = PointCloud(ptarr)
  pointcloud.channels[:rgb] = map(x -> RGB(reinterpret(N0f8, x)...), carr)
  setgeometry!(vc[:submaps][vsym][:pc], pointcloud)
  =#

  nothing
end

# this function is on notice!
function setupSLAMinDB(;cloudGraph=nothing, addrdict=nothing)
  if cloudGraph != nothing
    return SLAMWrapper(initfg(sessionname=addrdict["session"], cloudgraph=cloudGraph), nothing, 0)
  else
    return SLAMWrapper(RoME.initfg(), nothing, 0)
  end
end

# prepare the factor graph with just one node
# (will prompt on stdin for db credentials)
cloudGraph, addrdict = standardcloudgraphsetup()

slam = setupSLAMinDB(cloudGraph=cloudGraph, addrdict=addrdict)
slam.fg = identitypose6fg(fg=slam.fg)

function listenone(slaml)
  lc = LCM()

  # lcm callback closure
  pchdl = (channel, msg_data) -> lcmpointcloudmsg!( slaml, msg_data )

  subscribe(lc, "SUBMAPS", pchdl)

  # run this thing for one message only....
  handle(lc)
end

listenone(slam)

# Juno.breakpoint(@__FILE__(), 52)


##############################################################################
##
## LSMR GPU
##
##############################################################################
using CuArrays.CUDAnative
import CuArrays: allowscalar
allowscalar(false)

##############################################################################
##
## Conversion FixedEffect between CPU and GPU
##
##############################################################################

# https://github.com/JuliaGPU/CuArrays.jl/issues/306
cuzeros(T::Type, n::Integer) = fill!(CuVector{T}(undef, n), zero(T))
function CuArrays.cu(T::Type, fe::FixedEffect)
	refs = CuArray(fe.refs)
	interaction = cu(T, fe.interaction)
	FixedEffect{typeof(refs), typeof(interaction)}(refs, interaction, fe.n)
end
CuArrays.cu(T::Type, w::Union{Fill, Ones, Zeros}) = fill!(CuVector{T}(undef, length(w)), w[1])
CuArrays.cu(T::Type, w::AbstractVector) = CuVector{T}(w)

##############################################################################
##
## FixedEffectLinearMap on the GPU (code by Paul Schrimpf)
##
##############################################################################

const N_THREADS = 256

function FixedEffectLinearMap{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:lsmr_gpu}}) where {T}
	fes = [cu(T, fe) for fe in fes]
	sqrtw = cu(T, sqrt.(values(weights)))
	colnorm = [colnorm!(cuzeros(T, fe.n), fe.refs, fe.interaction, sqrtw) for fe in fes]
	caches = [cache!(cuzeros(T, length(sqrtw)), fe.interaction, sqrtw, scale, fe.refs) for (fe, scale) in zip(fes, colnorm)]
	return FixedEffectLinearMap{T}(fes, sqrtw, colnorm, caches)
end

function colnorm!(fecoef::CuVector, refs::CuVector, y::CuVector, sqrtw::CuVector)
	nblocks = cld(length(refs), N_THREADS)
	@cuda threads=N_THREADS blocks=nblocks colnorm!_kernel!(fecoef, refs, y, sqrtw)
	fecoef .= sqrt.(fecoef)
end

function colnorm!_kernel!(fecoef, refs, y, sqrtw)
	index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
	stride = blockDim().x * gridDim().x
	@inbounds for i = index:stride:length(y)
		CuArrays.CUDAnative.atomic_add!(pointer(fecoef, refs[i]), abs2(y[i] * sqrtw[i]))
	end
end

function cache!(y::CuVector, interaction::CuVector , sqrtw::CuVector, fecoef::CuVector, refs::CuVector)
	nblocks = cld(length(y), N_THREADS)
	@cuda threads=N_THREADS blocks=nblocks cache_kernel!(y, interaction, sqrtw, fecoef, refs)
	return y
end

function cache_kernel!(y, interaction, sqrtw, fecoef, refs)
	index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
	stride = blockDim().x * gridDim().x
	@inbounds for i = index:stride:length(y)
		y[i] += interaction[i] * sqrtw[i] / fecoef[refs[i]]
	end
end

function mean!(fecoef::CuVector, refs::CuVector, α::Number, y::CuVector, cache::CuVector)
	nblocks = cld(length(y), N_THREADS)
	@cuda threads=N_THREADS blocks=nblocks mean_kernel!(fecoef, refs, α, y, cache)
end

function mean_kernel!(fecoef, refs, α, y, cache)
	index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
	stride = blockDim().x * gridDim().x
	@inbounds for i = index:stride:length(y)
		CuArrays.CUDAnative.atomic_add!(pointer(fecoef, refs[i]), α * y[i] * cache[i])
	end
end

function demean!(y::CuVector, α::Number, fecoef::CuVector, refs::CuVector, cache::CuVector)
	nblocks = cld(length(y), N_THREADS)
	@cuda threads=N_THREADS blocks=nblocks demean_kernel!(y, α, fecoef, refs, cache)
end

function demean_kernel!(y, α, fecoef, refs, cache)
	index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
	stride = blockDim().x * gridDim().x
	@inbounds for i = index:stride:length(y)
		y[i] += α * fecoef[refs[i]] * cache[i]
	end
end

##############################################################################
##
## Implement AbstractFixedEffectSolver interface
##
##############################################################################
struct FixedEffectSolverLSMRGPU{T} <: AbstractFixedEffectSolver{T}
	m::FixedEffectLinearMap{T}
	b::CuVector{T}
	r::CuVector{T}
	x::FixedEffectCoefficients{T}
	v::FixedEffectCoefficients{T}
	h::FixedEffectCoefficients{T}
	hbar::FixedEffectCoefficients{T}
	tmp::Vector{T} # used to convert AbstractVector to Vector{T}
	fes::Vector{<:FixedEffect}
end

function AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:lsmr_gpu}}) where {T}
	m = FixedEffectLinearMap{T}(fes, weights, Val{:lsmr_gpu})
	b = cuzeros(T, length(weights))
	r = cuzeros(T, length(weights))
	x = FixedEffectCoefficients([cuzeros(T, fe.n) for fe in fes])
	v = FixedEffectCoefficients([cuzeros(T, fe.n) for fe in fes])
	h = FixedEffectCoefficients([cuzeros(T, fe.n) for fe in fes])
	hbar = FixedEffectCoefficients([cuzeros(T, fe.n) for fe in fes])
	tmp = zeros(T, length(weights))
	FixedEffectSolverLSMRGPU(m, b, r, x, v, h, hbar, tmp, fes)
end

function solve_residuals!(r::AbstractVector, feM::FixedEffectSolverLSMRGPU{T}; tol::Real = sqrt(eps(T)), maxiter::Integer = 100_000) where {T}
	copyto!(feM.tmp, r)
	copyto!(feM.r, feM.tmp)
	feM.r .*=  feM.m.sqrtw
	copyto!(feM.b, feM.r)
	fill!(feM.x, 0.0)
	x, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar; atol = tol, btol = tol, maxiter = maxiter)
	mul!(feM.r, feM.m, feM.x, -1.0, 1.0)
	feM.r ./=  feM.m.sqrtw
	copyto!(feM.tmp, feM.r)
	copyto!(r, feM.tmp)
	return r, div(ch.mvps, 2), ch.isconverged
end

function solve_residuals!(X::AbstractMatrix, feM::FixedEffectSolverLSMRGPU; kwargs...)
	iterations = Int[]
	convergeds = Bool[]
	for j in 1:size(X, 2)
		_, iteration, converged = solve_residuals!(view(X, :, j), feM; kwargs...)
		push!(iterations, iteration)
		push!(convergeds, converged)
	end
	return X, iterations, convergeds
end

function solve_coefficients!(r::AbstractVector, feM::FixedEffectSolverLSMRGPU{T}; tol::Real = sqrt(eps(T)), maxiter::Integer = 100_000) where {T}
	copyto!(feM.tmp, r)
	copyto!(feM.b, feM.tmp)
	feM.b .*= feM.m.sqrtw
	fill!(feM.x, 0.0)
	x, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar; atol = tol, btol = tol, maxiter = maxiter)
	for (x, scale) in zip(feM.x.x, feM.m.colnorm)
		x ./=  scale
	end
	x = Vector{eltype(r)}[collect(x) for x in feM.x.x]
	full(normalize!(x, feM.fes; tol = tol, maxiter = maxiter), feM.fes), div(ch.mvps, 2), ch.isconverged
end


using Oceananigans: AbstractModel, AbstractOutputWriter, AbstractDiagnostic

using Oceananigans.Architectures: AbstractArchitecture, CPU
using Oceananigans.Advection: CenteredSecondOrder
using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions

using Oceananigans.BoundaryConditions: UVelocityBoundaryConditions,
                                       VVelocityBoundaryConditions,
                                       TracerBoundaryConditions

using Oceananigans.Fields: Field, tracernames, TracerFields, XFaceField, YFaceField, CenterField
using Oceananigans.Forcings: model_forcing
using Oceananigans.Grids: with_halo, topology, inflate_halo_size, halo_size
using Oceananigans.TimeSteppers: Clock, TimeStepper
using Oceananigans.TurbulenceClosures: with_tracers, DiffusivityFields
using Oceananigans.Utils: tupleit

function ShallowWaterTendencyFields(arch, grid, tracer_names)

    uh = XFaceField(arch, grid, UVelocityBoundaryConditions(grid))
    vh = YFaceField(arch, grid, VVelocityBoundaryConditions(grid))
    h  = CenterField(arch,  grid, TracerBoundaryConditions(grid))
    tracers = TracerFields(tracer_names, arch, grid)

    return merge((uh=uh, vh=vh, h=h), tracers)
end

function ShallowWaterSolutionFields(arch, grid, bcs)

    uh_bcs = :uh ∈ keys(bcs) ? bcs.uh : UVelocityBoundaryConditions(grid)
    vh_bcs = :vh ∈ keys(bcs) ? bcs.vh : VVelocityBoundaryConditions(grid)
    h_bcs  = :h  ∈ keys(bcs) ? bcs.h  : TracerBoundaryConditions(grid)

    uh = XFaceField(arch, grid, uh_bcs)
    vh = YFaceField(arch, grid, vh_bcs)
    h = CenterField(arch, grid, h_bcs)

    return (uh=uh, vh=vh, h=h)
end

struct ShallowWaterModel{G, A<:AbstractArchitecture, T, V, R, F, E, B, Q, C, K, TS} <: AbstractModel{TS}

                          grid :: G         # Grid of physical points on which `Model` is solved
                  architecture :: A         # Computer `Architecture` on which `Model` is run
                         clock :: Clock{T}  # Tracks iteration number and simulation time of `Model`
    gravitational_acceleration :: T         # Gravitational acceleration, full, or reduced
                     advection :: V         # Advection scheme for velocities _and_ tracers
                      coriolis :: R         # Set of parameters for the background rotation rate of `Model`
                       forcing :: F         # Container for forcing functions defined by the user
                       closure :: E         # Diffusive 'turbulence closure' for all model fields
                    bathymetry :: B         # Bathymetry/Topography for the model
                      solution :: Q         # Container for transports `uh`, `vh`, and height `h`
                       tracers :: C         # Container for tracer fields
                 diffusivities :: K         # Container for turbulent diffusivities
                   timestepper :: TS        # Object containing timestepper fields and parameters

end

function ShallowWaterModel(;
                           grid,
                           gravitational_acceleration,
  architecture::AbstractArchitecture = CPU(),
                               clock = Clock{eltype(grid)}(0, 0, 1),
                           advection = UpwindBiasedFifthOrder(),
                            coriolis = nothing,
                 forcing::NamedTuple = NamedTuple(),
                             closure = nothing,
                          bathymetry = nothing,
                            solution = nothing,
                             tracers = (),
                       diffusivities = nothing,
     boundary_conditions::NamedTuple = NamedTuple(),
                 timestepper::Symbol = :RungeKutta3)

    tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)

    Hx, Hy, Hz = inflate_halo_size(grid.Hx, grid.Hy, grid.Hz, topology(grid), advection, closure)
    grid = with_halo((Hx, Hy, Hz), grid)

    model_field_names = (:uh, :vh, :h, tracers...)
    boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, model_field_names)

    solution = ShallowWaterSolutionFields(architecture, grid, boundary_conditions)
    tracers  = TracerFields(tracers, architecture, grid, boundary_conditions)
    diffusivities = DiffusivityFields(diffusivities, architecture, grid,
                                      tracernames(tracers), boundary_conditions, closure)

    # Instantiate timestepper if not already instantiated
    timestepper = TimeStepper(timestepper, architecture, grid, tracernames(tracers);
                              Gⁿ = ShallowWaterTendencyFields(architecture, grid, tracernames(tracers)),
                              G⁻ = ShallowWaterTendencyFields(architecture, grid, tracernames(tracers)))

    # Regularize forcing and closure for model tracer and velocity fields.
    model_fields = merge(solution, tracers)
    forcing = model_forcing(model_fields; forcing...)
    closure = with_tracers(tracernames(tracers), closure)

    return ShallowWaterModel(grid,
                             architecture,
                             clock,
                             eltype(grid)(gravitational_acceleration),
                             advection,
                             coriolis,
                             forcing,
                             closure,
                             bathymetry,
                             solution,
                             tracers,
                             diffusivities,
                             timestepper)
end


module ColumnwiseLUSolver

export ManyColumnLU, SingleColumnLU

using ..Mesh.Grids
using ..Mesh.Grids: VerticalDirection
using ..Mesh.Topologies
using ..DGmethods
using ..DGmethods: BalanceLaw, DGModel, num_state, num_diffusive
using ..LinearSolvers
const LS = LinearSolvers
using ..MPIStateArrays
using LinearAlgebra
using GPUifyLoops

include("ColumnwiseLUSolver_kernels.jl")

abstract type AbstractColumnLUSolver <: AbstractLinearSolver  end

"""
    ManyColumnLU()

This solver is used for systems that are block diagonal where each block is
associated with a column of the mesh.  The systems are solved using a
non-pivoted LU factorization.
"""
struct ManyColumnLU <: AbstractColumnLUSolver  end

"""
    SingleColumnLU()

This solver is used for systems that are block diagonal where each block is
associated with a column of the mesh.  Moreover, each block is assumed to be
the same.  The systems are solved using a non-pivoted LU factorization.
"""
struct SingleColumnLU <: AbstractColumnLUSolver  end

struct ColumnwiseLU{F, AT}
  f::F
  A::AT
end

function LS.prefactorize(op, solver::AbstractColumnLUSolver, Q, args...)
  dg = op.f!

  # TODO: can we get away with just passing the grid?
  A = banded_matrix(op, dg, similar(Q), similar(Q), args...;
                    single_column = typeof(solver) <: SingleColumnLU)

  band_lu!(A, dg)

  ColumnwiseLU(dg, A)
end

function LS.linearsolve!(clu::ColumnwiseLU{F}, ::AbstractColumnLUSolver,
                         Q, Qrhs, args...) where {F <: DGModel}
  dg = clu.f
  A = clu.A
  Q .= Qrhs

  band_forward!(Q, A, dg)
  band_back!(Q, A, dg)
end

"""
    band_lu!(A, dg::DGModel)

"""
function band_lu!(A, dg::DGModel)
  bl = dg.balancelaw
  grid = dg.grid
  topology = grid.topology
  @assert isstacked(topology)
  @assert typeof(dg.direction) <: VerticalDirection

  FT = eltype(A)
  device = typeof(A) <: Array ? CPU() : CUDA()

  nstate = num_state(bl, FT)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqj = dimensionality(grid) == 2 ? 1 : Nq

  eband = num_diffusive(bl, FT) == 0 ? 1 : 2

  nrealelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nrealelem, nvertelem)

  threads = (Nq, Nqj)
  blocks = nhorzelem

  if ndims(A) == 2
    # single column case
    #
    # TODO Would it be faster to copy the matrix to the host and factorize it
    # there?
    threads = (1, 1)
    blocks = 1
    A = reshape(A, 1, 1, size(A)..., 1)
  end

  @launch(device, threads=threads, blocks=blocks,
          band_lu_knl!(A, Val(Nq), Val(threads[1]), Val(threads[2]),
                       Val(nstate), Val(nvertelem), Val(blocks),
                       Val(eband)))
end

function band_forward!(Q, A, dg::DGModel)
  bl = dg.balancelaw
  grid = dg.grid
  topology = grid.topology
  @assert isstacked(topology)
  @assert typeof(dg.direction) <: VerticalDirection

  FT = eltype(A)
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  nstate = num_state(bl, FT)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqj = dimensionality(grid) == 2 ? 1 : Nq

  eband = num_diffusive(bl, FT) == 0 ? 1 : 2

  nrealelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nrealelem, nvertelem)

  @launch(device, threads=(Nq, Nqj), blocks=nhorzelem,
          band_forward_knl!(Q.data, A, Val(Nq), Val(Nqj), Val(nstate),
                            Val(nvertelem), Val(nhorzelem), Val(eband)))
end

function band_back!(Q, A, dg::DGModel)
  bl = dg.balancelaw
  grid = dg.grid
  topology = grid.topology
  @assert isstacked(topology)
  @assert typeof(dg.direction) <: VerticalDirection

  FT = eltype(A)
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  nstate = num_state(bl, FT)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqj = dimensionality(grid) == 2 ? 1 : Nq

  eband = num_diffusive(bl, FT) == 0 ? 1 : 2

  nrealelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nrealelem, nvertelem)

  @launch(device, threads=(Nq, Nqj), blocks=nhorzelem,
          band_back_knl!(Q.data, A, Val(Nq), Val(Nqj), Val(nstate),
                         Val(nvertelem), Val(nhorzelem), Val(eband)))
end

"""
    banded_matrix(dg::DGModel, [Q::MPIStateArray, dQ::MPIStateArray,
                  single_column=false])

Forms the banded matrices for each the column operator defined by the `DGModel`
dg.  If `single_column=false` then a banded matrix is stored for each column and
if `single_column=true` only the banded matrix associated with the first column
of the first element is stored. The bandwidth of the DG column banded matrix is
`p = q = (polynomialorder + 1) * nstate * nvertelem - 1` with `p` and `q` being
the upper and lower bandwidths.

The banded matrices are stored in the LAPACK band storage format
<https://www.netlib.org/lapack/lug/node124.html>.

The banded matrices are returned as an arrays where the array type matches that
of `Q`. If `single_column=false` then the returned array has 5 dimensions, which
are:
- first horizontal column index
- second horizontal column index
- band index (-q:p)
- vertical DOF index with state `s`, vertical DOF index `k`, and vertical
  element `ev` mapping to `s + nstate * (k - 1) + nstate * nvertelem * (ev - 1)`
- horizontal element index

If the `single_column=true` then the returned array has 2 dimensions which are
the band index and the vertical DOF index.
"""
function banded_matrix(dg::DGModel, Q::MPIStateArray = MPIStateArray(dg),
                       dQ::MPIStateArray = MPIStateArray(dg);
                       single_column = false)
  banded_matrix((dQ, Q) -> dg(dQ, Q, nothing, 0; increment=false),
                dg, Q, dQ; single_column = single_column)
end

"""
    banded_matrix(f!::Function, dg::DGModel,
                  Q::MPIStateArray = MPIStateArray(dg),
                  dQ::MPIStateArray = MPIStateArray(dg), args...;
                  single_column = false, args...)

Forms the banded matrices for each the column operator defined by the linear
operator `f!` which is assumed to have the same banded structure as the
`DGModel` dg.  If `single_column=false` then a banded matrix is stored for each
column and if `single_column=true` only the banded matrix associated with the
first column of the first element is stored. The bandwidth of the DG column
banded matrix is `p = q = (polynomialorder + 1) * nstate * nvertelem - 1` with
`p` and `q` being the upper and lower bandwidths.

The banded matrices are stored in the LAPACK band storage format
<https://www.netlib.org/lapack/lug/node124.html>.

The banded matrices are returned as an arrays where the array type matches that
of `Q`. If `single_column=false` then the returned array has 5 dimensions, which
are:
- first horizontal column index
- second horizontal column index
- band index (-q:p)
- vertical DOF index with state `s`, vertical DOF index `k`, and vertical
  element `ev` mapping to `s + nstate * (k - 1) + nstate * nvertelem * (ev - 1)`
- horizontal element index

If the `single_column=true` then the returned array has 2 dimensions which are
the band index and the vertical DOF index.

Here `args` are passed to `f!`.
"""
function banded_matrix(f!, dg::DGModel,
                       Q::MPIStateArray = MPIStateArray(dg),
                       dQ::MPIStateArray = MPIStateArray(dg),
                       args...; single_column = false)
  bl = dg.balancelaw
  grid = dg.grid
  topology = grid.topology
  @assert isstacked(topology)
  @assert typeof(dg.direction) <: VerticalDirection

  FT = eltype(Q.data)
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  nstate = num_state(bl, FT)
  N = polynomialorder(grid)
  Nq = N + 1

  # p is lower bandwidth
  # q is upper bandwidth
  eband = num_diffusive(bl, FT) == 0 ? 1 : 2
  p = q = nstate * Nq * eband - 1

  nrealelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nrealelem, nvertelem)

  dim = dimensionality(grid)

  Nqj = dim == 2 ? 1 : Nq

  # first horizontal DOF index
  # second horizontal DOF index
  # band index -q:p
  # vertical DOF index
  # horizontal element index
  A = if single_column
    similar(Q.data, p + q + 1, Nq * nstate * nvertelem)
  else
    similar(Q.data, Nq, Nqj, p + q + 1, Nq * nstate * nvertelem,
            nhorzelem)
  end
  fill!(A, zero(FT))

  # loop through all DOFs in a column and compute the matrix column
  for ev = 1:nvertelem
    for s = 1:nstate
      for k = 1:Nq
        # Set a single 1 per column and rest 0
        @launch(device, threads=(Nq, Nqj, Nq), blocks=(nvertelem, nhorzelem),
                knl_set_banded_data!(bl, Val(dim), Val(N), Val(nvertelem),
                                     Q.data, k, s, ev, 1:nhorzelem,
                                     1:nvertelem))

        # Get the matrix column
        f!(dQ, Q, args...)

        # Store the banded matrix
        @launch(device, threads=(Nq, Nqj, Nq),
                blocks=(2 * eband + 1, nhorzelem),
                knl_set_banded_matrix!(bl, Val(dim), Val(N), Val(nvertelem),
                                       Val(p), Val(q), Val(eband+1),
                                       A, dQ.data, k, s, ev, 1:nhorzelem,
                                       -eband:eband))
      end
    end
  end
  A
end

"""
    banded_matrix_vector_product!(dg::DGModel, A, dQ::MPIStateArray,
                                  Q::MPIStateArray)

Compute a matrix vector product `dQ = A * Q` where `A` is assumed to be a matrix
created using the `banded_matrix` function.

This function is primarily for testing purposes.
"""
function banded_matrix_vector_product!(dg::DGModel, A, dQ::MPIStateArray,
                                       Q::MPIStateArray)
  bl = dg.balancelaw
  grid = dg.grid
  topology = grid.topology
  @assert isstacked(topology)
  @assert typeof(dg.direction) <: VerticalDirection

  FT = eltype(Q.data)
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  eband = num_diffusive(bl, FT) == 0 ? 1 : 2
  nstate = num_state(bl, FT)
  N = polynomialorder(grid)
  Nq = N + 1
  p = q = nstate * Nq * eband - 1

  nrealelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nrealelem, nvertelem)

  dim = dimensionality(grid)

  Nqj = dim == 2 ? 1 : Nq

  @launch(device, threads=(Nq, Nqj, Nq),
          blocks=(nvertelem, nhorzelem),
          knl_banded_matrix_vector_product!(bl, Val(dim), Val(N),
                                            Val(nvertelem), Val(p), Val(q),
                                            dQ.data, A, Q.data, 1:nhorzelem,
                                            1:nvertelem))
end

end


using Test
using DeepCorrections
using Random
using POMDPs
using POMDPModels
using POMDPSimulators
using Flux
using DeepQLearning
using DiscreteValueIteration

test_env_path = joinpath(dirname(pathof(DeepQLearning)),"..", "test", "test_env.jl")
include(test_env_path)

function evaluate(mdp, policy, rng, n_ep=100, max_steps=100)
    avg_r = 0.
    sim = RolloutSimulator(rng=rng, max_steps=max_steps)
    for i=1:n_ep
        DeepQLearning.resetstate!(policy)
        avg_r += simulate(sim, mdp, policy)
    end
    return avg_r/=n_ep
end

@testset begin "Dummy Corrections"
    # set random seed for test block
    Random.seed!(2)
    rng = MersenneTwister(2)
    mdp = TestMDP((5,5), 4, 6)
    model = Chain(x->flattenbatch(x), Dense(100, 8, tanh), Dense(8, 4))
    dqn_solver = DeepQLearningSolver(qnetwork = model, max_steps=20000, learning_rate=0.005,
                                    eval_freq=2000,num_ep_eval=100,
                                    log_freq = 1000,
                                    double_q = false, dueling=true, prioritized_replay=false,
                                    rng=rng)
    solver = DeepCorrectionSolver(dqn = dqn_solver)
    corr_pol = solve(solver, mdp)
    r_basic = evaluate(mdp, corr_pol, rng)
    @test r_basic >= 1.0
end

@testset begin "Perfect Qlo"
    # set random seed for test block
    Random.seed!(2)
    rng = MersenneTwister(2)
    # value table
    mdp = SimpleGridWorld()
    vi_pol = solve(ValueIterationSolver(), mdp)
    function vi_values(mdp::SimpleGridWorld, s::AbstractArray)
        s_gw = convert_s(GWPos, s, mdp)
        si = stateindex(mdp, s_gw)
        return vi_pol.qmat[si, :]
    end
    model = Chain(x->flattenbatch(x), Dense(2, 32, relu), Dense(32, n_actions(mdp)))
    dqn_solver = DeepQLearningSolver(qnetwork = model, prioritized_replay=true, max_steps=40_000, learning_rate=0.01,
                                     log_freq=1000, eval_freq = 2000,
                                     eps_fraction = 0.1, eps_end = 0.01,
                                     double_q=true, dueling=true, rng=rng)
    solver = DeepCorrectionSolver(dqn = dqn_solver,
                                  lowfi_values = vi_values,
                                  correction = additive_correction)

    corr_policy = solve(solver, mdp)
    r_gw =  evaluate(mdp, corr_policy, rng)
    @test r_gw >= 0.
end

@testset begin "Use Policy for Lowfi Values"
        # set random seed for test block
    Random.seed!(2)
    rng = MersenneTwister(2)
    # value table
    mdp = SimpleGridWorld()
    vi_pol = solve(ValueIterationSolver(), mdp)
    POMDPs.stateindex(mdp::SimpleGridWorld, s::AbstractArray) = stateindex(mdp, convert(Vector{Int64}, s))
    model = Chain(x->flattenbatch(x), Dense(2, 32, relu), Dense(32, n_actions(mdp)))
    dqn_solver = DeepQLearningSolver(qnetwork = model, prioritized_replay=true, max_steps=40_000, learning_rate=0.01,
                                     log_freq=1000, eval_freq = 2000,
                                     eps_fraction = 0.1, eps_end = 0.01,
                                     double_q=true, dueling=true, rng=rng)
    solver = DeepCorrectionSolver(dqn = dqn_solver,
                                  lowfi_values = vi_pol,
                                  correction = additive_correction)

    corr_policy = solve(solver, mdp)
    r_gw =  evaluate(mdp, corr_policy, rng)
    @test r_gw >= 0.
end


abstract type ExponentialFamily end
function bregmanDivergence()
end
function update_parameters!()
end
function dimension()
end
function sourceParameters()
end
function sourceVector()
end
function source_to_expectation()
end
include("WatsonDistribution.jl")
include("MultivariateGaussian.jl")


@testset "basic integration" begin
    m = SquareRooter(4)
    report = IC.train!(m, Step(2), InvalidValue(), NumberLimit(3); verbosity=0);
    @test report[1] == (Step(2), (new_iterations = 6,))
    @test report[2] == (InvalidValue(), (done=false, log=""))
    report[3] == (NumberLimit(3),
                  (done=true,
                   log="Stop triggered by NumberLimit(3) "*
                   "stopping criterion. "))

    m = SquareRooter(4)
    @test_logs((:info, r"final loss"),
               (:info, r"final training loss"),
               (:info, r"Stop triggered by Num"),
               IC.train!(m, Step(2),  InvalidValue(), NumberLimit(3)));
    @test_logs((:info, r"Using these controls"),
               (:info, r"Stepping model for 2 more iterations"),
               (:info, r"Stepping model for 2 more iterations"),
               (:info, r"Stepping model for 2 more iterations"),
               (:info, r"final loss"),
               (:info, r"final training loss"),
               (:info, r"total control cycles\: 3"),
               (:info, r"A total of 6 iterations added"),
               (:info, r"Stop triggered by NumberLimit"),
               IC.train!(m, Step(2),
                         InvalidValue(),
                         NumberLimit(3);
                         verbosity=2));
end

struct WrappedSquareRooter
    s::SquareRooter
end

@testset "expose and Callback(..., raw=true)" begin
    IC.train!(w::WrappedSquareRooter, n) = IC.train!(w.s, n)
    IC.loss(w::WrappedSquareRooter) = IC.loss(w.s)
    IC.train!(w::WrappedSquareRooter, n) = IC.train!(w.s, n)
    IC.expose(w::WrappedSquareRooter) = w.s

    wrapper = WrappedSquareRooter(SquareRooter(4))

    roots = Any[]
    wrappers = Any[]

    g(w) = push!(wrappers, deepcopy(w))

    c1 = Callback(s->push!(roots, s.root))
    c2 = Callback(g, raw=true)

    IC.train!(wrapper, Step(1), c1, c2, NumberLimit(2), verbosity=0)

    @test map(w -> w.s.root, wrappers) ≈ roots
end

@testset "skip integration" begin
    m = SquareRooter(4)
    numbers = Int[]
    IC.train!(m,
              Step(1),
              IterationControl.skip(
                  WithNumberDo(x->push!(numbers, x)), predicate=3),
              NumberLimit(10), verbosity=0)
    @test numbers == [3, 6, 9]
end

@testset "integration test related to #38" begin
    model = IterationControl.SquareRooter(4)
    @test_logs((:info, r"number"),
               (:info, r"number"),
               (:info, r"final loss"),
               (:info, r"final training loss"),
               (:info, r"Stop triggered by"),
               IC.train!(model,
                         Step(1),
                         Threshold(2.1),
                         WithNumberDo(),
                         IterationControl.skip(WithLossDo(), predicate=3)))
end

@testset "integration test with PQ" begin
    # get complete training losses:
    model = SquareRooter(1e40)
    IC.train!(model, 70)
    tlosses = IC.training_losses(model)

    # get complete losses:
    model = SquareRooter(1e40)
    losses = map(1:70) do _
        IC.train!(model, 1)
        IC.loss(model)
    end

    # train with PQ criterion:
    K = 5
    model = SquareRooter(1e40)
    pq = PQ(k=K)
    states = Any[]
    f(s) = push!(states, s)
    controls = [Step(1), NumberLimit(70), IC.with_state_do(pq, f=f)]
    IC.train!(model, controls..., verbosity=0)

    # check internal state of PQ makes sense, as far as `loss` and
    # `training_losses` are concerned:
    @test map(s->s.loss, states) == losses
    @test map(enumerate(states)) do (j, s)
        L = min(K, j) # what length of PQ training_losses should be
        t= s.loss == losses[j] &&
            s.training_losses == reverse(tlosses[j-L+1:j])
    end |> all
end

# https://github.com/JuliaAI/MLJIteration.jl/issues/36
@testset "integration test related to MLJIteration.jl #36" begin
    model = SquareRooter(NaN)
    report =
        IC.train!(model, Step(1), InvalidValue(), NumberLimit(3), verbosity=0)
    @test report[1][2].new_iterations == 1
    @test occursin("Stopping", report[2][2].log)
end