#!/usr/bin/julia

###############################################################################
#
# Kpax3 - Bayesian cluster analysis of categorical data
#
# Copyright (c) 2016-2018 Alberto Pessia <alberto.pessia@gmail.com>
#
# Kpax3 is free software: you can redistribute it and/or modify it under the
# terms of the MIT License. Refer to LICENSE.md file for more info.
#
###############################################################################

###############################################################################
#
# CSV FILE SUPPORT
#
###############################################################################
#
# csv support is still in an early stage of development. The dataset can be
# analyzed without problems but the user API is not mature yet.
#
###############################################################################

###############################################################################
#
# INTRODUCTION
#
###############################################################################
#
# Welcome to the Kpax3 tutorial!
#
# In this tutorial you will learn how to:
# - load a generic categorical dataset from a csv file
# - setup the parameters
# - run the algorithm
# - save the results
# - plot the results
#
# If you didn't do it already, copy this file and the provided dataset to a
# location of your choice.
#
# The tutorial has been designed in such a way that you should just copy/paste
# the commands into your julia environment. Alternatively, it is also possible
# to run the script directly from the command line by typing:
#
# julia Kpax3_csv_tutorial.jl
#
# from within the directory where the file is located.
#
###############################################################################

###############################################################################
#
# NOTE FOR WINDOWS USERS
#
###############################################################################
#
# File paths should be written using the slash '/' character instead of the
# usual backslash '\' character. If you want to use the backslash, you have to
# escape it (i.e. use '\\'). Examples of valid file paths:
#
# "C:/Users/me/Documents/file.txt"
# "C:\\Users\\me\\Documents\\file.txt"
#
###############################################################################

###############################################################################
#
# LOAD THE PACKAGE
#
###############################################################################
#
# First of all, we need to load the package "Kpax3". If the package is not
# already installed in your system, please follow the installation instructions
# written in the README.md file.
#
# To load the package, open a julia session and run the command
using Kpax3;

###############################################################################
#
# INPUT/OUTPUT
#
###############################################################################
#
# The data we will use in this tutorial consist of 75 aligned protein genomes
# of ZIKA Virus (Human). Total length of each sequence is 3425 amino acids, but
# only 186 polymorphisms. Once converted to binary format, dataset will consist
# of 75 units and 382 binary variables.
#
# Data was retrieved from NCBI’s Virus Variation Resource database
# https://www.ncbi.nlm.nih.gov/genome/viruses/variation/Zika/
# on 2017-02-10 and aligned with MUSCLE
# http://www.drive5.com/muscle/
# and converted to a comma delimited file (csv) file.
#
# Edit the following variable with the path to the dataset
input_file = "path_to_zika_dataset.csv";

# Edit the following variable with the path to the output file.
# This file will contain intermediate output generated by Kpax3. Note that file
# extension is not necessary.
output_file = "path_to_output_file";

###############################################################################
#
# SETUP PARAMETERS
#
###############################################################################
#
# Kpax3 comes with default parameters values that should work in the majority
# of applications. At this point, you can already start analysing the dataset.
# Uncomment the following line if you do not want to customize your Kpax3 run
# and use default values:
#settings = KSettings(input_file, output_file);

# -----------------------------------------------------------------------------
# If any setting is omitted, its default value will be used.
#
# - Missing values
# First of all, we need to specify which characters are to be considered
# missing data.
#
# If the dataset is contained in a fasta file (genetic data), then the variable
# is called `miss` and the format is a UInt8 vector representing the
# characters:
#   Default value:
#     miss = UInt8['?', '*', '#', '-', 'b', 'j', 'x', 'z']
#
# If the dataset is contained in a csv file (generic categorical data), then
# the variable is called `misscsv` and the format is a String vector
# representing the values:
#  Default value:
#    misscsv = [""]
#
# - Ewens-Pitman distribution for the partition
# alpha is a Float64 variable constrained in (0, 1).
# theta is a Float64 variable constrained in (-alpha, infinity)
#
# Default values:
#   alpha = 0.5
#   theta = -0.25
#
# - Column status prior probabilities
# This is a non-negative Float64 vector of length 3. If it does not sum to 1,
# it is automatically normalized.
# Probabilities contained in this variable represent an initial guess for the
# proportion of sites possessing the following statuses:
# - uninformative (noise)
# - informative but not characteristic for any cluster (weak signal)
# - informative and characteristic for a cluster (strong signal)
#
# Default value:
#   gamma = [0.6; 0.35; 0.05]
#
# - Beta distribution hyper-parameter
# This is a Float64 variable and it's a bit tricky to explain and to properly
# setup. You can do it on your own, of course, but it requires some good
# understanding of the model. Read the Pessia and Corander (2017) paper to
# learn its usage. I suggest leaving it to its default value.
#
# Default value:
#   r = log(0.001) / log(0.95)
#
# - General settings
# Boolean variable 'verbose' to print diagnostics messages during run.
# Integer variable 'verbosestep' to print diagnostics messages every fixed
# number of steps.
#
# Default values:
#   verbose = false
#   verbosestep = 500
#
# - MCMC settings
# Integer variable 'T' to choose the length of the Markov Chain.
# Integer variable 'burnin' to choose the length of the burnin period.
# Integer variable 'tstep' to tune the thinning of the chain, i.e. samples will
# be collected every 'tstep' iterations.
# Float64 vector 'op' of length 3 with proportions of MCMC operators, i.e.
# [split-merge; gibbs; biased random walk]. To save time, we suggest using a
# very long chain with only split-merge and biased random walk. For particular
# datasets, gibbs sampler might be needed to reach convergence.
#
# Default values:
#   T = 100000
#   burnin = 10000
#   tstep = 1
#   op = [0.5; 0.0; 0.5]
#
# - Genetic Algorithm settings
# Integer variable 'popsize' to choose the population size.
# Integer variable 'maxiter' to choose the total number of iterations.
# Integer variable 'maxgap' to choose how many iterations without new
#   solutions we want to wait before stopping the search.
# Float64 variable 'xrate' for crossover rate.
# Float64 variable 'mrate' for mutation rate.
#
# Default values:
#   popsize = 20
#   maxiter = 20000
#   maxgap = 5000
#   xrate = 0.9
#   mrate = 0.005
# -----------------------------------------------------------------------------

# To save time in this tutorial, we will simulate a short chain. For a real
# data analysis, choose a longer chain in order to reduce the effect of
# autocorrelation between observations.
settings = KSettings(input_file,
                     output_file,
                     miss=UInt8['?', '*', '#', '-', 'b', 'j', 'x', 'z'],
                     misscsv=["", "?", "*", "#", "-", "B", "J", "X", "Z"],
                     alpha=0.5,
                     theta=-0.25,
                     gamma=[0.6; 0.35; 0.05],
                     r=log(0.001) / log(0.95),
                     verbose=true,
                     verbosestep=1000,
                     T=50000,
                     burnin=10000,
                     tstep=1,
                     op=[0.5; 0.0; 0.5],
                     popsize=20,
                     maxiter=20000,
                     maxgap=5000,
                     xrate=0.9,
                     mrate=0.005);

###############################################################################
#
# LOAD DATA
#
###############################################################################
#
# Loading data is straightforward.
zika_data = CategoricalData(settings);

###############################################################################
#
# INITIAL PARTITION
#
###############################################################################
#
# We need to start our search, or simulation, somewhere. Kpax3 offers two ways
# to input an initial partition of the units:
# 1) Provide a csv file with a known initial partition, for example a
#    solution from a previous run;
# 2) Heuristically find a good initial partition.
#
# Here, we will opt for the second option by using k-means. Choose a suitable
# initial number of clusters for your dataset.
using Clustering;
k = 10;
initial_partition = kmeans(convert(Array{Float64}, zika_data.data),
                           k).assignments;

# The csv file can be either a single column of integers (cluster indices)
# representing the partition, or a two-column file where the first column
# contains the units ids and the second column the cluster indices.
#initial_partition = "path_to_initial_partition.csv";

###############################################################################
#
# MAP ESTIMATE (OPTIMIZATION BY GENETIC ALGORITHM)
#
###############################################################################
#
# We will start by obtaining a point estimate equivalent to the maximum a
# posteriori (MAP). In general, MAP estimates provide good solutions to the
# problem of clustering genetic sequences. It is the recommended approach
# for huge datasets or if we are only interested in clustering the sequences.
#
# Keep in mind that the solution might be a local optimum, requiring multiple
# runs of the optimization algorithm to increase the chance of finding the
# global optimum.
#
# Fortunately for us, the genetic algorithm implemented by Kpax3 is good enough
# to get very close to the global optimum. Even if the solution is a local
# optimum, it will still be a good approximation to the exact solution.
map_solution = kpax3ga(zika_data, initial_partition, settings);

# Check the value of the log-posterior distribution (plus a constant)
# Compare different solutions by this value and pick the one with the maximum
# value.
using Printf;
@printf("Log-posterior value (plus a constant): %.4f\n", map_solution.logpp);

###############################################################################
#
# POSTERIOR DISTRIBUTION (MCMC)
#
###############################################################################
#
# To obtain a measure of uncertainty regarding the parameters of interest, we
# need to sample realizations of such parameters from the posterior
# distribution. We will do it by simulating a Markov Chain.
#
# NOTE
# This step is going to take a while.
#
# Depending on your dataset, you need to customize chain length and operators
# proportions. Obviously, the longer the chain (number of simulations) the
# better, at the cost of longer simulation times and increased size of hard
# drive memory required. If the dataset is too big for MCMC technique, MAP
# estimate (see previous section) is the recommended approach.
kpax3mcmc(zika_data, initial_partition, settings);

# Compute the estimate that minimize the posterior expected loss
mcmc_solution = kpax3estimate(zika_data, settings);

# If the MAP estimate is a local optimum, it might happen that the MCMC
# estimate will have a higher posterior probability
if mcmc_solution.logpp > map_solution.logpp
  # let's try again starting from a better solution. If no solution is found,
  # map_solution will be equal to mcmc_solution
  map_solution = kpax3ga(zika_data, mcmc_solution.R, settings);
  nothing;
end

###############################################################################
#
# SAVE RESULTS
#
###############################################################################
#
# We will save both the MAP estimate and the state that minimize the posterior
# expected loss. In general, MAP estimate is a better point estimate.
#
# Edit the path to the output files. No file extension needed.
map_output_file = "path_to_MAP_estimate_output_file";
mcmc_output_file = "path_to_MCMC_estimate_output_file";

writeresults(zika_data, map_solution, map_output_file, what=4, verbose=true);
writeresults(zika_data, mcmc_solution, mcmc_output_file, what=4, verbose=true);

# Output description:
# If what = 1:
#   filename_summary.txt        = Text file with the log posterior probability.
#                                 This can be used to compare different
#                                 parallel runs.
#   filename_partition.csv      = Best partition found by the algorithm.
#
# If what = 2 writes all the previous files plus:
#   filename_attributes.csv     = Complete columns classification matrix.
#
# If what = 3 writes all the previous files plus:
#   filename_characteristic.csv = Characteristic sites and their corresponding
#                                 values.
#
# If what = 4 writes all the previous files plus:
#   filename_dataset.txt        = Original dataset with highlighted clusters
#                                 and their corresponding characteristic
#                                 values. It is meaningful only if the data
#                                 values are single latin characters (a-zA-Z).
#
# To immediately check if the estimate is reliable, open
# "filename_characteristic.csv" and count how many characteristic sites have
# been found. If the ratio
# number of characteristic sites / total number of SNPs
# is greater than 0.1 and each cluster has at least a characteristic amino
# acid (the more the better), then the solution can be considered reliable.
#
# Suppose we are not satisfied with the current MAP solution and we think
# that units should be moved around, or a cluster should be split or two
# clusters merged. This could happen when metadata is available, so that we
# can immediately see if something is not right.
#
# After manually editing the partition file, we can load it again and check
# our hypothesis. No need to run the algorithm again. There is a function
# written specifically for this task.
#
# If needed, uncomment the following commands.
#my_partition = "path_to_manually_edited_partition.csv";
#my_solution = optimumstate(zika_data, my_partition, settings);
#@printf("Log-posterior value (plus a constant): %.4f\n", my_solution.logpp);
#
# You can save this solution the same way we did with MAP and MCMC estimates.

###############################################################################
#
# PLOT RESULTS
#
###############################################################################
#
# In the following, we will use the MAP estimate as the reference partition.
#
# Function signatures are shown in case you want to change default values.

# Choose the default plotting backend, such as GR. Install them with
#   ]add Plots GR
using Plots;
gr();

# For high quality figures, you can save all these plots to SVG format with the
# following command:
#   Plots.savefig(plotfunction(plot_parameters, size=(width, height),
#                 "path_to_figure.svg");
# where 'width' and 'height' are Float64 to be set according to your needs and
# 'plotfunction' is the corresponding desired plot.
#
# 'Plots.savefig' chooses file type automatically by the extension.

# Plot the complete dataset, highlighting clusters and polymorphic sites
# function kpax3plotd(x::AminoAcidData,
#                     state::AminoAcidState;
#                     clusterorder=:auto,
#                     clusterlabel=:auto)
kpax3plotd(zika_data, map_solution)

# Load posterior probabilities:
(k, pk) = readposteriork(output_file);
(id, P) = readposteriorP(output_file);
(site, aa, freq, C) = readposteriorC(output_file);

# Let's check how the Markov Chain performed by looking at trace and jump
# plots. We will do it separately for R (partition of the units) and C
# (partition of the columns).
#
# Compute the entropy of each sample and the average distance of entropies at
# different lags.
(entropy_R, avgd_R) = traceR(output_file, maxlag=50);
(entropy_C, avgd_C) = traceC(output_file, maxlag=50);

# function kpax3plottrace(entropy::Vector{Float64};
#                         maxlag=200,
#                         M=20000,
#                         main="")
kpax3plottrace(entropy_R)
kpax3plottrace(entropy_C)

# function kpax3plotdensity(entropy::Vector{Float64};
#                           maxlag=200,
#                           main="")
kpax3plotdensity(entropy_R)
kpax3plotdensity(entropy_C)

# function kpax3plotjump(avgd::Vector{Float64};
#                        main="")
kpax3plotjump(avgd_R)
kpax3plotjump(avgd_C)

# Posterior distribution of the number of clusters
# function kpax3plotk(k::Vector{Int},
#                     pk::Vector{Float64})
kpax3plotk(k, pk)

# Posterior distribution of adjacency matrix
# function kpax3plotp(R::Vector{Int},
#                     P::Matrix{Float64};
#                     clusterorder=:auto,
#                     clusterlabel=:auto)
kpax3plotp(map_solution.R, P)

# Posterior distribution of features types (not a very informative plot in this
# case because of few polymorphic sites compared to the length of the genome)
# function kpax3plotc(site::Vector{Int},
#                     freq::Vector{Float64},
#                     C::Matrix{Float64})
kpax3plotc(site, freq, C)


@testset "Persistence" begin
    (f, fio) = mktemp()
    close(fio)

    gx = PathGraph(5)
    for g in testgraphs(gx)
        mg = MetaGraph(g)
        set_prop!(mg, 1, 2, :weight, 0.2)
        set_prop!(mg, 2, 3, :color, :red)
        set_prop!(mg, 1, :name, "Alice")
        set_prop!(mg, 2, :name, "Bob")
        set_prop!(mg, :type, "My MetaGraph")

        @test savegraph(f, mg) == 1
        g2 = loadgraph(f, MGFormat())
        @test mg == g2
    end

    gx = PathDiGraph(5)
    for g in testdigraphs(gx)
        mg = MetaDiGraph(g)
        set_prop!(mg, 1, 2, :weight, 0.2)
        set_prop!(mg, 2, 3, :color, :red)
        set_prop!(mg, 1, :name, "Alice")
        set_prop!(mg, 2, :name, "Bob")
        set_prop!(mg, :type, "My MetaGraph")

        @test savegraph(f, mg) == 1
        g2 = loadgraph(f, MGFormat())
        @test mg == g2
    end
    rm(f)
end


# Fonction qu'on veut tester
# include("../NLPAtXmod.jl")
# importall NLPAtXmod

# On vérifie que le constructeur pour problème sans contrainte fonctionne
uncons_nlp_at_x = NLPAtX(zeros(10))

@test uncons_nlp_at_x.x == zeros(10)
@test uncons_nlp_at_x.fx == nothing
@test uncons_nlp_at_x.gx == nothing
@test uncons_nlp_at_x.Hx == nothing
@test uncons_nlp_at_x.mu == nothing
@test uncons_nlp_at_x.cx == nothing
@test uncons_nlp_at_x.Jx == nothing

@test uncons_nlp_at_x.lambda == zeros(0)
@test uncons_nlp_at_x.current_time == nothing
@test uncons_nlp_at_x.current_score == nothing
@test (!(uncons_nlp_at_x.evals == nothing))

# On vérifie que le constucteur pour problème avec contrainte fonctionne
cons_nlp_at_x = NLPAtX(zeros(10), zeros(10))

@test cons_nlp_at_x.x == zeros(10)
@test cons_nlp_at_x.fx == nothing
@test cons_nlp_at_x.gx == nothing
@test cons_nlp_at_x.Hx == nothing
@test cons_nlp_at_x.mu == nothing
@test cons_nlp_at_x.cx == nothing
@test cons_nlp_at_x.Jx == nothing
@test (false in (cons_nlp_at_x.lambda .== 0.0)) == false
@test cons_nlp_at_x.current_time == nothing
@test cons_nlp_at_x.current_score == nothing

# On vérifie que la fonction update! fonctionne
update!(uncons_nlp_at_x, x = ones(10), fx = 1.0, gx = ones(10))
update!(uncons_nlp_at_x, lambda = ones(10), current_time = 1.0)
update!(uncons_nlp_at_x, Hx = ones(10,10), mu = ones(10), cx = ones(10), Jx = ones(10,10))

@test (false in (uncons_nlp_at_x.x .== 1.0)) == false #assez bizarre comme test...
@test uncons_nlp_at_x.fx == 1.0
@test (false in (uncons_nlp_at_x.gx .== 1.0)) == false
@test (false in (uncons_nlp_at_x.Hx .== 1.0)) == false
@test uncons_nlp_at_x.mu == ones(10)
@test uncons_nlp_at_x.cx == ones(10)
@test (false in (uncons_nlp_at_x.Jx .== 1.0)) == false
@test (false in (uncons_nlp_at_x.lambda .== 1.0)) == false
@test uncons_nlp_at_x.current_time == 1.0
@test uncons_nlp_at_x.current_score == nothing

reinit!(uncons_nlp_at_x)
@test uncons_nlp_at_x.x == ones(10)
@test uncons_nlp_at_x.fx == nothing
reinit!(uncons_nlp_at_x, x = zeros(10))
@test uncons_nlp_at_x.x == zeros(10)
@test uncons_nlp_at_x.fx == nothing

nlp_64 = NLPAtX(ones(10))
nlp_64.x = ones(10)
nlp_64.fx = 1.0
nlp_64.gx = ones(10)

# nlp_32 = convert_nlp(Float32, nlp_64)
# @test typeof(nlp_32.x[1]) == Float32
# @test typeof(nlp_32.fx[1]) == Float32
# @test typeof(nlp_32.gx[1]) == Float32
# @test isnan(nlp_32.mu[1])
# @test isnan(nlp_32.current_time)
#
# @test typeof(nlp_64.x[1]) == Float64
# @test typeof(nlp_64.fx[1]) == Float64
# @test typeof(nlp_64.gx[1]) == Float64
# @test isnan(nlp_64.mu[1])
# @test isnan(nlp_64.current_time)

#Test the _size_check:
try
    NLPAtX(ones(5), gx = zeros(4))
    @test false
catch
    @test true
end
try
    NLPAtX(ones(5), mu = zeros(4))
    @test false
catch
    @test true
end
try
    NLPAtX(ones(5), Hx = zeros(4,4))
    @test false
catch
    @test true
end
try
    NLPAtX(ones(5), zeros(1), cx = zeros(0))
    @test false
catch
    @test true
end


using JuliaVariables
using MLStyle
import MacroTools

unwrap_scoped(ex) = @match ex begin
    Expr(:scoped, _, a) => unwrap_scoped(a)
    Expr(head, args...) => Expr(head, map(unwrap_scoped, args)...)
    a => a
end

globals(s::Symbol) = [s]
globals(x) = []

function globals(ex::Expr)
    branch(head, newargs) = union(newargs...)

    function leaf(v::JuliaVariables.Var)
        v.is_global ? [v.name] : Symbol[]
    end

    leaf(x) = []

    solved_ex = unwrap_scoped(solve_from_local!(simplify_ex(ex)))

    return foldall(leaf, branch)(solved_ex)
end

function graph(T, m::Model)
    g = T()

    mvars = variables(m)
    for v in mvars
        SimplePosets.add!(g, v)
    end

    for (v, expr) in pairs(m.vals)
        for p in globals(expr) ∩ variables(m)
            SimplePosets.add!(g, p, v)
        end
    end

    for (v, expr) in pairs(m.dists)
        for p in globals(expr) ∩ variables(m)
            SimplePosets.add!(g, p, v)
        end
    end

    g
end

export digraph
digraph(m::Model) = graph(SimpleDigraph{Symbol}, m)

export poset
poset(m::Model) = graph(SimplePoset{Symbol}, m)


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.7.0
#     language: julia
#     name: julia-1.7
# ---

# %%
struct Foo{A, B}
    a::A
    b::B
end

# %%
methods(Foo)

# %%
methods(Foo{Int, Int})

# %%
foo = Foo(1.2, 34)

# %%
mult(x::Foo{<:Number, <:Number}) = x.a * x.b

# %%
methods(mult)

# %%
mult(Foo(2, 3 + 4im))

# %%
mult(x::Foo{<:Integer, <:AbstractString}) = x.b ^ x.a

# %%
methods(mult)

# %%
mult(Foo(10, "hoge! "))

# %%


#
# Julia System Desciption Language
#
module Causal

using UUIDs
using DifferentialEquations
using Sundials
using LightGraphs
using DataStructures
using JLD2
using Plots
using ProgressMeter
using Logging
using LinearAlgebra
using Dates
using NLsolve
using Interpolations
using LibGit2
using DocStringExtensions
import GraphPlot.gplot
import FileIO: load
import Base: show, display, write, read, close, setproperty!, mv, cp, open,  istaskdone, istaskfailed,
    getindex, setindex!, size, isempty

include("utilities/utils.jl")
export equip

include("utilities/callback.jl")
export Callback,
    enable!,
    disable!,
    isenabled,
    applycallbacks

include("utilities/buffer.jl")
export BufferMode,
    LinearMode,
    CyclicMode,
    Buffer,
    Normal,
    Cyclic,
    Fifo,
    Lifo,
    write!,
    isfull,
    ishit,
    content,
    mode,
    snapshot,
    datalength,
    inbuf,
    outbuf

include("connections/link.jl")
export Link,
    launch,
    refresh!

include("connections/pin.jl")
export AbstractPin,
    Outpin,
    Inpin,
    connect!,
    disconnect!,
    isconnected,
    isbound

include("connections/port.jl")
export AbstractPort,
    Inport,
    Outport,
    datatype

include("components/componentsbase/hierarchy.jl")
export AbstractComponent,
    AbstractSource,
    AbstractSystem,
    AbstractSink,
    AbstractStaticSystem,
    AbstractDynamicSystem,
    AbstractSubSystem,
    AbstractMemory,
    AbstractDiscreteSystem,
    AbstractODESystem,
    AbstractRODESystem,
    AbstractDAESystem,
    AbstractSDESystem,
    AbstractDDESystem

include("components/componentsbase/macros.jl")
include("components/componentsbase/interpolant.jl")
export Interpolant

include("components/componentsbase/takestep.jl")
export readtime!,
    readstate,
    readinput!,
    writeoutput!,
    computeoutput,
    evolve!,
    takestep!,
    drive!,
    approve!

include("components/constants.jl")

include("components/sources/clock.jl")
export Clock,
    isrunning,
    ispaused,
    isoutoftime,
    set!,
    stop!,
    pause!

include("components/sources/generators.jl")
export @def_source,
    FunctionGenerator,
    SinewaveGenerator,
    DampedSinewaveGenerator,
    SquarewaveGenerator,
    TriangularwaveGenerator,
    ConstantGenerator,
    RampGenerator,
    StepGenerator,
    ExponentialGenerator,
    DampedExponentialGenerator

include("components/systems/staticsystems/staticsystems.jl")
export @def_static_system,
    StaticSystem,
    Adder,
    Multiplier,
    Gain,
    Terminator,
    Memory,
    Coupler,
    Differentiator

include("components/systems/dynamicalsystems/init.jl")

include("components/systems/dynamicalsystems/odesystems.jl")
export @def_ode_system,
    ODESystem,
    ContinuousLinearSystem,
    LorenzSystem,
    ChenSystem,
    ChuaSystem,
    RosslerSystem,
    VanderpolSystem,
    ForcedLorenzSystem,
    ForcedChenSystem,
    ForcedChuaSystem,
    ForcedRosslerSystem,
    ForcedVanderpolSystem,
    Integrator

include("components/systems/dynamicalsystems/discretesystems.jl")
export @def_discrete_system,
    DiscreteSystem,
    DiscreteLinearSystem,
    HenonSystem,
    LoziSystem,
    BogdanovSystem,
    GingerbreadmanSystem,
    LogisticSystem

include("components/systems/dynamicalsystems/sdesystems.jl")
export @def_sde_system,
    SDESystem,
    NoisyLorenzSystem,
    ForcedNoisyLorenzSystem

include("components/systems/dynamicalsystems/daesystems.jl")
export @def_dae_system,
    DAESystem,
    RobertsonSystem,
    PendulumSystem,
    RLCSystem

include("components/systems/dynamicalsystems/rodesystems.jl")
export @def_rode_system,
    RODESystem,
    MultiplicativeNoiseLinearSystem

include("components/systems/dynamicalsystems/ddesystems.jl")
export @def_dde_system,
    DDESystem,
    DelayFeedbackSystem

include("components/sinks/sinks.jl")
export @def_sink,
    Writer,
    Printer,
    Scope,
    write!,
    fwrite!,
    fread,
    update!

include("models/taskmanager.jl")
export TaskManager,
    checktaskmanager

include("models/simulation.jl")
export Simulation,
    setlogger,
    closelogger,
    report

include("models/model.jl")
export @defmodel,
    Model,
    inspect!,
    initialize!,
    run!,
    terminate!,
    simulate!,
    getloops,
    breakloop!,
    Node,
    Branch,
    addnode!,
    getnode,
    addbranch!,
    getbranch,
    deletebranch!,
    signalflow,
    troubleshoot

include("plugins/loadplugins.jl")
export AbstractPlugin,
    process,
    add,
    remove,
    enable,
    disable,
    check,
    @def_plugin

end  # module


abstract type DynamicInverterComponent <: DynamicComponent end
abstract type Converter <: DynamicInverterComponent end
abstract type DCSource <: DynamicInverterComponent end
abstract type Filter <: DynamicInverterComponent end
abstract type FrequencyEstimator <: DynamicInverterComponent end
abstract type InnerControl <: DynamicInverterComponent end

abstract type ActivePowerControl <: DeviceParameter end
abstract type ReactivePowerControl <: DeviceParameter end

"""
Parameters of a Outer-Loop controller using a virtual inertia with VSM for active power controller
and a reactive power droop controller.
# Conmutable structor
```julia
VirtualInertiaQDroop(A, R)
```
# Arguments
*  `A`::Float64 : Active power controller using virtual inertia with VSM
*  `R`::Float64 : Reactive power controller using reactive power droop
"""
mutable struct OuterControl{A <: ActivePowerControl, R <: ReactivePowerControl} <:
               DynamicInverterComponent
    active_power::A
    reactive_power::R
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int64
end

function OuterControl(
    active_power::A,
    reactive_power::R,
    ext = Dict{String, Any}(),
) where {A <: ActivePowerControl, R <: ReactivePowerControl}
    OuterControl(
        active_power,
        reactive_power,
        ext,
        vcat(active_power.states, reactive_power.states),
        active_power.n_states + reactive_power.n_states,
    )
end

function OuterControl(; active_power, reactive_power, ext = Dict{String, Any}())
    OuterControl(active_power, reactive_power, ext)
end

get_active_power(value::OuterControl) = value.active_power
get_reactive_power(value::OuterControl) = value.reactive_power
get_ext(value::OuterControl) = value.ext
get_states(value::OuterControl) = value.states
get_n_states(value::OuterControl) = value.n_states


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

@testset "graph.plainhypergraph" begin
    g = plainhypergraph(6, [Set([1, 2, 3, 4]), Set([1, 4, 5]), Set([6])])
    @test issetequal(neighbors(g, 4)[2], [1, 4, 5])
    @test issetequal(getedge(g, 1), 1:4)
    @test hasedge(g, 2, 3)
    @test !hasedge(g, 2, 5)
    @test_throws ErrorException nodeattr(g, 1)
    @test_throws ErrorException edgeattr(g, 1)
    @test_throws ErrorException edgeattr(g, 3, 4)

    @test issetequal(nodeset(g), 1:6)
    @test issetequal(edgeset(g), 1:3)
    @test_throws ErrorException nodeattrs(g)
    @test_throws ErrorException edgeattrs(g)

    @test nodecount(g) == 6
    @test edgecount(g) == 3
    @test neighborcount(g, 2) == 1
    @test degree(g, 3) == 1

    addnode!(g)
    @test degree(g, 7) == 0
    addedge!(g, Set([4, 6]))
    @test degree(g, 4) == 3
    @test_throws ErrorException nodeattrtype(g)
    @test_throws ErrorException edgeattrtype(g)

    emp = plainhypergraph()
    @test nodecount(emp) == 0

end # graph.plaingraph


export QLearner,
       GradientBanditLearner,
       MonteCarloLearner, MonteCarloExploringStartLearner, OffPolicyMonteCarloLearner,
       TDLearner, OffPolicyTDLearner, DoubleLearner, QLearner, DifferentialTDLearner, TDλReturnLearner,
       ReinforceLearner,  ReinforceBaselineLearner,
       DQN, DoubleDQN

include("gradient_bandit_learner.jl")
include("monte_carlo_learner.jl")
include("temporal_difference_learner.jl")
include("reinforce_learner.jl")
include("dqn.jl")


import .MCMCChains: Chains

# A basic chains constructor that works with the Transition struct we defined.
function bundle_samples(
    rng::AbstractRNG,
    model::DensityModel,
    s::Metropolis,
    N::Integer,
    ts::Vector{T},
    chain_type::Type{Chains};
    param_names=missing,
    kwargs...
) where {ModelType<:AbstractModel, T<:AbstractTransition}
    # Turn all the transitions into a vector-of-vectors.
    vals = copy(reduce(hcat,[vcat(t.params, t.lp) for t in ts])')

    # Check if we received any parameter names.
    if ismissing(param_names)
        param_names = ["Parameter $i" for i in 1:length(s.init_params)]
    else
        # Deepcopy to be thread safe.
        param_names = deepcopy(param_names)
    end

    # Add the log density field to the parameter names.
    push!(param_names, "lp")

    # Bundle everything up and return a Chains struct.
    return Chains(vals, param_names, (internals=["lp"],))
end


using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\kernelEvolutions.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\structs\\BasicStructs.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\utils\\CUDAGpuUtils.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\utils\\IterationUtils.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\utils\\ReductionUtils.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\utils\\MemoryUtils.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\distanceMetrics\\MeansMahalinobis.jl")

using ..BasicPreds, ..CUDAGpuUtils , ..MeansMahalinobis, ..IterationUtils,..ReductionUtils , ..MemoryUtils

function getExampleKernelArgs()
    nx=512 ; ny=512 ; nz=317
    #first we initialize the metrics on CPU so we will modify them easier
    goldBoolCPU= zeros(Float32,nx,ny,nz); #mimicks gold standard mask
    segmBoolCPU= ones(Float32,nx,ny,nz); #mimicks other  mask
    cartTrueGold =  CartesianIndices(zeros(3,3,5) ).+CartesianIndex(5,5,5);
    cartTrueSegm =  CartesianIndices(zeros(150,150,150) ).+CartesianIndex(100,100,100);

    goldBoolCPU[cartTrueGold].=Float32(1.0);
    segmBoolCPU[cartTrueSegm].=Float32(1.0);

    numberToLooFor= Float32(1.0);
    goldBoolGPU = CuArray(goldBoolCPU);
    segmBoolGPU = CuArray(segmBoolCPU);
    #we will fill it after we work with launch configuration
    loopXdim = UInt32(1);loopYdim = UInt32(1) ;loopZdim = UInt32(1) ;
    sizz = size(goldBoolCPU);maxX = UInt32(sizz[1]);maxY = UInt32(sizz[2]);maxZ = UInt32(sizz[3])
    threads = (32,20)
    blocks = 50
    loopXdim = UInt32(fld(maxX, threads[1]))
    loopYdim = UInt32(fld(maxY, threads[2]))
    loopZdim = UInt32(fld(maxZ,blocks ))

    totalCountGold= CuArray([0]);
    totalCountSegm= CuArray([0]);

    return (goldBoolGPU,segmBoolGPU,numberToLooFor
    ,loopYdim,loopXdim,loopZdim
    ,(maxX, maxY,maxZ)
    ,totalCountGold
    ,totalCountSegm)
end

# @testset "iter3d" begin

    args = getExampleKernelArgs()
    goldBoolGPU,segmBoolGPU,numberToLooFor,loopYdim,loopXdim,loopZdim,maxes,totalCountGold,totalCountSegm= args

    function testLoopKernelA(goldArr,segmArr
        ,numberToLooFor
        ,loopYdim::UInt32
        ,loopXdim::UInt32
        ,loopZdim::UInt32
        ,arrDims::Tuple{UInt32,UInt32,UInt32}
        ,totalCountGold,totalCountSegm)

        countGold::UInt32 = UInt32(0)
        countSegm::UInt32 = UInt32(0)
        shmemSum= @cuStaticSharedMem(UInt32, (32,2))
        clearSharedMemWarpLong(shmemSum, UInt8(2))

        @iter3d(arrDims,  loopXdim,loopYdim,loopZdim,if(  @inbounds(goldArr[x,y,z])  ==numberToLooFor)
            #updating variables needed to calculate means
            countGold+=UInt32(1)
        end)#if bool in arr

        sync_threads()

        @iter3dWithVal(segmArr,arrDims, loopXdim,loopYdim,loopZdim, if(  value  ==numberToLooFor)
            #updating variables needed to calculate means
            countSegm+=UInt32(1)
        end)#if bool in arr

        offsetIter = UInt8(1)
        @redWitAct(offsetIter,shmemSum,  countGold,+, countSegm,+  )
        @addAtomic(shmemSum,totalCountGold, totalCountSegm)

        return
    end

    @cuda threads=(32,20) blocks=50 testLoopKernelA(args...)

    @test totalCountGold[1]==3*3*5
    nx=512 ; ny=512 ; nz=317
    @test totalCountSegm[1]==(nx*ny*nz)

    using Revise, Parameters, Logging, Test
    using CUDA
    includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
    using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils, ..BasicStructs
    using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates
    using ..MainOverlap, ..RandIndex , ..ProbabilisticMetrics , ..VolumeMetric ,..InformationTheorhetic
    using ..CUDAAtomicUtils, ..TpfpfnKernel, ..InterClassCorrKernel

    dimx = 529
    dimy = 556
    dimz = 339
    # dimx = 333
    # dimy = 33
    # dimz = 33
    arrGold = CUDA.ones(UInt8,(dimx,dimy,dimz))
    arrAlgo =  CUDA.ones(UInt8,(dimx,dimy,dimz))
    sizz= size(arrGold)
    ##### load tests ...
    conf = ConfigurtationStruct(trues(12)...)
    numberToLooFor = UInt8(1)

    args,threads,blocks,metricsTuplGlobal= TpfpfnKernel.prepareForconfusionTableMetricsNoSliceWise(numberToLooFor  ,conf)

    # iterLoop= args[9]
    # pixPerSlice = Int64(args[10])
    singleSum= CUDA.zeros(UInt64,1)
    function testIter(goldGPU#goldBoolGPU
        , segmGPU#segmBoolGPU
        ,sliceMetricsTupl
        ,tp,tn,fp,fn#tp,tn,fp,fn
        ,arrDims,totalNumbOfVoxels,iterLoop,pixPerSlice
        ,numberToLooFor#numberToLooFor
        #,metricsTuplGlobal
        ,conf,singleSum)
        @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,begin
           CUDA.@atomic singleSum[1]+=1
            #CUDA.@cuprint "i $(i) \n"
        end)

    return
    end

    @cuda threads=threads blocks=blocks testIter(arrGold, arrAlgo,args..., singleSum) #args[8][3]  is number of slices ...

    @test Int64(singleSum[1]) ==dimx*dimy*dimz

    Int64(singleSum[1])-dimx*dimy*dimz

# end#testset


"""
    mutable struct Program_

Structure that contains all the pieces of code before they are finally assembled into the final evolve function.
"""
mutable struct Program_

    agent::Agent
    space::SimulationSpace

    declareVar::Expr
    declareF::Expr
    args::Array{Symbol,1}
    execInit::Expr
    execInloop::Expr
    execAfter::Expr
    returning::Expr

    update::Dict{String,Dict{Symbol,Int}}
end

function Program_(abm::Agent,space::SimulationSpace)
    return Program_(abm,space,quote end,quote end,Array{Symbol,1}([:t,:N,:dt]),
                    quote end,quote end,quote end,quote end,Dict{String,Dict{Symbol,Int}}())
end


ENV["PYCALL_JL_RUNTIME_PYTHON"] = Sys.which("python")
using Distributed
# using 6 cores
if 6 - Distributed.nprocs() > 0
    Distributed.addprocs(6 - Distributed.nprocs());
end
@everywhere using SACBP
using SACBP
using LinearAlgebra
using Plots
pyplot()
using ProgressMeter
using Printf
using Random
using JLD2
using FileIO

include(joinpath(@__DIR__, "../baseline_implementations/p_control.jl"))
include(joinpath(@__DIR__, "../baseline_implementations/mcts.jl"))
include(joinpath(@__DIR__, "../baseline_implementations/belief_ilqg.jl"))
include(joinpath(@__DIR__, "../baseline_implementations/tlqg.jl"))

function get_initial_true_state()
    px0_true = 1.5;    # True px(0)
    py0_true = 1.;     # True py(0)
    θ0_true = 11*pi/6.;# True theta(0) in (0,2*pi)
    vx0_true = 0.;     # True vx(0)
    vy0_true = 0.;     # True vy(0)
    ω0_true = 0.;      # True ω(0)
    m_true = 2.0;      # True mass
    a_true = 0.5;      # Height of the Plate.
    b_true = 1.0;      # Width of the Plate.
    J_true = 1/3*m_true*(a_true^2 + b_true^2); # Moment of Inertia
    rx_true = 0.25;    # Arm Length rx
    ry_true = 0.25;    # Arm Length ry
    μ_true = 4.0;      # Linear Friction Coefficient
    x_init = PhysManipulate2D(0.,[px0_true,py0_true],θ0_true,[vx0_true,vy0_true],ω0_true,m_true,J_true,[rx_true,ry_true],μ_true);
    return x_init, a_true, b_true
end

function get_initial_belief_state()
    μ_init = [4.,4.,pi/4,0.1,-0.1,pi/10,6.,2.,0.3,0.1,7.0];
    Σ_init = diagm([10.,10.,(pi/2)^2,2.,2.,(pi/4)^2,1.,1.0,1.0,1.0,5.0]);
    b_init = BelMvNormal(0.,μ_init,Σ_init);
    return b_init
end

function visualizeSimEnv(px::Real,py::Real,θ::Real,vx::Real,vy::Real,ω::Real,
                         a::Real,b::Real,rx::Real,ry::Real,dto=0.0;
                         plot_velocity=false)
    plt = plot();
    Vertices = [[-b/2,a/2],[-b/2,-a/2],[b/2,-a/2],[b/2,a/2]];
    xVertices = map(v -> v[1], Vertices);
    yVertices = map(v -> v[2], Vertices);
    VRotated = map(v -> [cos(θ) -sin(θ); sin(θ) cos(θ)]*v + [px, py], Vertices);
    xRotated = map(v -> v[1], VRotated);
    yRotated = map(v -> v[2], VRotated);
    Robot = [cos(θ) -sin(θ); sin(θ) cos(θ)]*[rx, ry] + [px, py];
    plot!(xRotated,yRotated,seriestype=:shape,color=:cyan,linecolor=:darkcyan,alpha=0.4,aspect_ratio=1.0,label="Object",xlabel="x [m]",ylabel="y [m]");
    scatter!((px,py),color=:darkcyan,markersize=5.0,label="C.M.")
    scatter!((Robot[1],Robot[2]),color=:red,markersize=10.0,label="Robot")
    if plot_velocity
        plot!([px,px+dto*vx],[py,py+dto*vy],line=:arrow,color=:green,width=2.0,label="Lin. Vel.")
        plot!([px + min(a,b)/4*cos(pi/2+ii) for ii = range(0.,stop=5*dto*ω,length=10)],
              [py + min(a,b)/4*sin(pi/2+ii) for ii = range(0.,stop=5*dto*ω,length=10)],line=:arrow,color=:orange,width=2.0,label="Ang. Vel.");
    end
    plot!(xVertices,yVertices,seriestype=:shape,color=:black,linecolor=:black,alpha=0.1,label="Goal",size=(500,400))
    xlims!((-1.,2.3))
    ylims!((-1.,1.5))
    return plt
end

function simulate(b_init::BelMvNormal{<:Real},
                  x_init::PhysManipulate2D{<:Real},
                  a_true::Real, b_true::Real,
                  controlMode::String, rng_num::Int64,
                  simModel::SimulateManipulate2D;
                  numSamples=10, # for sacbp only
                  u_param_min::Vector{Float64}=[-3.0, -3.0, -3.0],
                  u_param_max::Vector{Float64}=[3.0, 3.0, 3.0],
                  posGain::Real,
                  rotGain::Real,
                  Th::Float64=2.0, # planning horizon [s]
                  tf::Float64=20.0, # simulation horizon [s]
                  mcts_param::Union{Nothing, MCTSParams}=nothing,
                  ilqgPolicy::Union{Nothing, Dict{String,Any}}=nothing,
                  tlqg::Union{Nothing, TLQG_Planner_Manipulate2D}=nothing,
                  animate=true, plotfig=true, verbose=true)
    rng = MersenneTwister(rng_num)
    @assert controlMode in ["sacbp", "mcts", "ilqg", "tlqg", "tlqg_offline"];
    if animate
        # Animation Object.
        anim = Animation();
    end
    # Only for MCTS policies.
    if controlMode == "mcts"
        @assert !isnothing(mcts_param) "MCTS Parameters not given"
        bmdp = BelMDP(simModel, u_param_min, u_param_max, mcts_param.numActions)
        mctsSolverPControl = MCTS.DPWSolver(depth=Int64(Th/simModel.dto),
                                            n_iterations=mcts_param.numSamples,
                                            exploration_constant=mcts_param.expConst,
                                            k_action=mcts_param.kAct,
                                            alpha_action=mcts_param.αAct,
                                            k_state=mcts_param.kState,
                                            alpha_state=mcts_param.αState,
                                            rng=rng,
                                            estimate_value=RolloutEstimator(PControl(u_param_min, u_param_max, posGain, rotGain)));
        mctsPolicyPControl = solve(mctsSolverPControl,bmdp);
    end

    contTransModel = simModel.continuousTransModel
    discTransModel = simModel.discreteTransModel
    dtc, dto = simModel.dtc, simModel.dto
    Q, R, Cs, Cu = simModel.Q, simModel.R, simModel.Cs, simModel.Cu

    # Controller Compilation.
    U_default = [MControl2D(round(dtc*t,digits=5),zeros(3)) for t = 0:Int64(Th/dtc)-1];
    if controlMode == "sacbp" # SACBP (with PControl as nominal)
        @time U_default = pControlUpdate(simModel,b_init,U_default,u_param_min,u_param_max,posGain,rotGain);
        @time U_default = sacControlUpdate(simModel,b_init,U_default,u_param_min,u_param_max,numSamples,rng)[1];
    elseif controlMode == "mcts" # MCTS (with PControl as nominal)
        @time U_default = mctsControlUpdate(mctsPolicyPControl,b_init,U_default,dto)
    elseif controlMode == "ilqg" # Belief iLQG (initialized with PControl)
        @assert !isnothing(ilqgPolicy) "iLQG Policy file not given"
        @time U_default = ilqgControlUpdate(ilqgPolicy,b_init,U_default,u_param_min,u_param_max,dto,dto)
    elseif controlMode == "tlqg"
        @assert !isnothing(tlqg) "T-LQG Planner object not given"
        @time ~,U_default = tlqgControlUpdate(tlqg,b_init,U_default,verbose=verbose,online=true)
    elseif controlMode == "tlqg_offline"
        # offline == planner can take much longer time to compute NLP iterations (until convergence).
        # Only for collecting computation time statistics
        @assert !isnothing(tlqg) "T-LQG Planner object not given"
        @time ~,U_default = tlqgControlUpdate(tlqg,b_init,U_default,verbose=verbose,online=false)
    end

    println("Controller Compiled. Starting Simulation...")

    # Simulation.
    y_history = Vector{Float64}[];
    tcalc_true_history = Float64[];
    x_history = [x_init];
    b_history = [b_init];
    u_history = MControl2D{Float64}[];
    U_pool = copy(U_default);

    observeModel = ObserveModel_Manipulate2D();
    numSteps = Int64(tf/dtc);
    nextSampleTime = dto;
    @showprogress 1 for t = 1:numSteps
        if round(t*dtc,digits=5) > nextSampleTime
        # Belief & Control Update.
            x = x_history[end];
            if animate
                plt = visualizeSimEnv(x.pos[1],x.pos[2],x.θ,x.vel[1],x.vel[2],x.ω,a_true,b_true,x.r[1],x.r[2],dto,plot_velocity=true)
                time_str = @sprintf "%3.2f" x.t
                title!("Time: $(time_str) [s]")
                Plots.frame(anim);
            end
            u = u_history[end];
            # Observation.
            yVec = observe(observeModel,x,MControl2D(x.t,vec(u)),R,rng);
            push!(y_history,yVec);
            # Belief Update.
            b = pop!(b_history);
            b_new = trans(discTransModel,b,MControl2D(b.t,vec(u)),yVec,R);
            push!(b_history,b_new);
            if verbose
                println((b_new.t,"Belief Updated."));
            end
            # Control Update
            if controlMode == "sacbp";  # SACBP (with PControl)
                tcalc_true = @elapsed U_pool = pControlUpdate(simModel,b_new,U_pool,u_param_min,u_param_max,posGain,rotGain);
                tcalc_true += @elapsed U_pool,act_time_init,act_time_final,tcalc_simulated = sacControlUpdate(simModel,b_new,U_pool,u_param_min,u_param_max,numSamples,rng,offline=true);
                push!(tcalc_true_history,tcalc_true);
                if verbose
                    println((b.t,"Control Updated.","time_init=$(act_time_init)","time_final=$(act_time_final)"));
                end
            elseif controlMode == "mcts" # MCTS (with PControl)
                tcalc_true = @elapsed U_pool = mctsControlUpdate(mctsPolicyPControl,b_new,U_pool,dto);
                push!(tcalc_true_history,tcalc_true);
                if verbose
                    println((b.t,"Control Updated."));
                end
            elseif controlMode == "ilqg" # Belief iLQG (with PControl)
                tcalc_true = @elapsed U_pool = ilqgControlUpdate(ilqgPolicy,b_new,U_pool,u_param_min,u_param_max,dto,dto);
                push!(tcalc_true_history,tcalc_true);
                if verbose
                    println((b.t,"Control Updated."));
                end
            elseif controlMode == "tlqg" # T-LQG (with PControl)
                #println(b_new)
                tcalc_true, U_pool = tlqgControlUpdate(tlqg,b_new,U_pool,verbose=verbose,online=true)
                if !isnothing(tcalc_true)
                    push!(tcalc_true_history,tcalc_true);
                end
                if verbose
                    println((b.t,"Control Updated."));
                end
            elseif controlMode == "tlqg_offline" # T-LQG (with PControl), offline
                tcalc_true, U_pool = tlqgControlUpdate(tlqg,b_new,U_pool,verbose=verbose,online=false)
                if !isnothing(tcalc_true)
                    push!(tcalc_true_history,tcalc_true);
                end
                if verbose
                    println((b.t,"Control Updated."));
                end
            end
            # Control Execution.
            u = popfirst!(U_pool);
            push!(u_history,u);
            x_new = trans(TransModel_Manipulate2D(),x,u,dtc);
            push!(x_history,x_new);
            b_new = trans(contTransModel,b_new,MControl2D(b.t,vec(u)),dtc,Q);
            push!(b_history,b_new);

            # Append nominal control.
            u_nominal = MControl2D(round(b.t+Th,digits=5),zeros(3));
            push!(U_pool,u_nominal);
            nextSampleTime = round(nextSampleTime+dto,digits=5);
        else
            # No Control Update.
            b = b_history[end];
            # Control Execution.
            x = x_history[end];
            if animate && round(x.t/0.02, digits=5) == round(x.t/0.02)
                plt = visualizeSimEnv(x.pos[1],x.pos[2],x.θ,x.vel[1],x.vel[2],x.ω,a_true,b_true,x.r[1],x.r[2],dto,plot_velocity=true)
                time_str = @sprintf "%3.2f" x.t
                title!("Time: $(time_str) [s]")
                Plots.frame(anim);
            end
            u = popfirst!(U_pool);
            push!(u_history,u);
            x_new = trans(TransModel_Manipulate2D(),x,u,dtc);
            push!(x_history,x_new);
            b_new = trans(contTransModel,b,u,dtc,Q);
            push!(b_history,b_new);

            # Append nominal control.
            u_nominal = MControl2D(round(b.t+Th,digits=5),zeros(3));
            push!(U_pool,u_nominal);
        end
    end
    # Observation.
    x = x_history[end];
    if animate
        plt = visualizeSimEnv(x.pos[1],x.pos[2],x.θ,x.vel[1],x.vel[2],x.ω,a_true,b_true,x.r[1],x.r[2],dto,plot_velocity=true)
        time_str = @sprintf "%3.2f" x.t
        title!("Time: $(time_str) [s]")
        Plots.frame(anim);
    end
    u = u_history[end];
    # Observation.
    yVec = observe(observeModel,x,MControl2D(x.t,vec(u)),R,rng);
    push!(y_history,yVec);
    # Belief Update.
    b = pop!(b_history);
    b_new = trans(discTransModel,b,MControl2D(x.t,vec(u)),yVec,R);
    push!(b_history,b_new);
    if verbose
        println((b_new.t,"Belief Updated."));
    end
    log_dir = normpath(joinpath(@__DIR__, "../data",
                                "manipulation_with_unknown_parameters"));
    if animate
        gif(anim,joinpath(log_dir, "$(controlMode)_tf_$(tf)_rng_$(rng_num).gif"),fps=50)
        #gif(anim,joinpath(log_dir, "$(controlMode)_tf_$(tf)_rng_$(rng_num).gif"),fps=5)
    end
    return y_history,tcalc_true_history,b_history,x_history,u_history,U_pool
end

function simulate_main(controlMode::String; rng_num::Int64,
                       animate=true, plotfig=true, verbose=true)
    # PControl Parameters
    posGain = 1.0;
    rotGain = 0.5;

    # SAC Controller Parameters.
    dtc = 0.01;        # Euler Approximation Interval.
    dto = 0.2;         # Observation & Control Update Interval.
    Th = 2.0;          # Planning Horizon.
    dtexec = 0.2*dto;  # Control Execution Duration.
    Q = diagm([0.05, 0.05, 1.0*pi/180, 0.05, 0.05, 1.0*pi/180., 0., 0., 0., 0., 0.]); # state: [px, py, θ, vx, vy, ω, m, J, rx, ry, μ]
    R = diagm([0.1, 0.1, 5.0*pi/180, 0.1, 0.1, 5.0*pi/180., 0.2, 0.2, 10.0*pi/180.]) # observation: [px, py, θ, vx, vy, ω, ax, ay(linear acceleration), α(angular acceleration)]
    Cs = diagm([20.0, 20.0, 20.0, 15.0, 15.0, 15.0, 0., 0., 0., 0., 0.]); # State Cost Coefficient.
    Cu = diagm([0.1, 0.1, 0.1]);  # Control Cost Coefficient.
    u_param_min = [-3.,-3.,-3.];  # Lower Limit for PosControl
    u_param_max = [3.,3.,3.];     # Upper Limit for PosControl
    contTransModel = CPredictManipulate2D();
    discTransModel = DUpdateManipulate2D();
    costModel = CostManipulate2D();
    simModel = SimulateManipulate2D(contTransModel,discTransModel,costModel,dtc,dto,dtexec,Q,R,Cs,Cu);
    numSamples = 10;     # Number of Monte Carlo Samples for SACBP.

    # MCTS Parameters
    numActions = 7;             # Number of Actions per Coordinate.
    numSamples_mcts = 240;
    expConst = 100.0            # Exploration Constant.
    kAct = 8.0;                 # Action k.
    αAct = 0.2;                 # Action α.
    kState = 8.0;               # State k.
    αState = 0.2;               # State α.
    mcts_param = MCTSParams(numActions, numSamples_mcts, expConst, kAct, αAct, kState, αState);

    # iLQG Policy
    ilqg_policy_path = normpath(joinpath(@__DIR__, "../baseline_implementations",
                                         "ilqg_pcontrol_policy_20sec_2.jld2"));
    ilqgPolicy = FileIO.load(ilqg_policy_path);

    # Simulation Parameters
    tf = 20.;    # Simulation Time.

    # T-LQG Planner
    Th_tlqg = Th;
    kl_threshold = 25.0;
    tlqg = TLQG_Planner_Manipulate2D(Th_tlqg,dto,Q,R,Cs,Cu,u_param_min,u_param_max,posGain,rotGain,kl_threshold)

    x_init, a_true, b_true = get_initial_true_state();
    if plotfig
        frame = visualizeSimEnv(x_init.pos[1],x_init.pos[2],x_init.θ,
                                x_init.vel[1],x_init.vel[2],x_init.ω,
                                a_true, b_true, x_init.r[1], x_init.r[2])
        init_fig_path = normpath(joinpath(@__DIR__, "../data",
                                          "manipulation_with_unknown_parameters",
                                          "manipulation_initial.pdf"))
        savefig(init_fig_path);
    end

    b_init = get_initial_belief_state();

    y_history,tcalc_true_history,b_history,x_history,u_history,U_pool =
        simulate(b_init, x_init, a_true, b_true, controlMode, rng_num,
                 simModel, numSamples=numSamples,
                 u_param_min=u_param_min, u_param_max=u_param_max,
                 posGain=posGain, rotGain=rotGain,
                 Th=Th, tf=tf, mcts_param=mcts_param,
                 ilqgPolicy=ilqgPolicy, tlqg=tlqg,
                 animate=animate, plotfig=plotfig, verbose=verbose);
    println("Average control computation time: $(round(mean(tcalc_true_history), digits=3)) [s]")
    return y_history,tcalc_true_history,b_history,x_history,u_history,U_pool, Cs, Cu
end

if abspath(PROGRAM_FILE) == @__FILE__
    println("1. Simulation with SACBP Controller")
    simulate_main("sacbp", rng_num=202, animate=true, plotfig=true, verbose=false);
    println("2. Simulation with MCTS-DPW Controller")
    simulate_main("mcts", rng_num=202, animate=true, plotfig=true, verbose=false);
    println("3. Simulation with Belief iLQG Controller")
    simulate_main("ilqg", rng_num=202, animate=true, plotfig=true, verbose=false);
    println("4. Simulation with T-LQG Controller")
    simulate_main("tlqg", rng_num=202, animate=true, plotfig=true, verbose=false);
end


module Utils

using DataFrames
using DataStructures: OrderedDict
using Dates

using CodecZlib
using Serialization
using Printf

export dview, toposort, coldim, ncols, FlushedLogger
export serialize_gzip, deserialize_gzip

coldim(::AbstractDataFrame) = 2
coldim(::AbstractMatrix) = 1

"""
    ncols(d)

Return the number of columns, i.e., features of `d`.
"""
ncols(d) = size(d, coldim(d))

function dview(d::AbstractDataFrame, obs, cols)
    view(d, obs, cols)
end

dview(d::AbstractMatrix, obs, cols) = view(d, cols, obs)

"""
    isapprox(dfa::AbstractDataFrame, dfb::AbstractDataFrame)

Check if two `DataFrame` instances are approximatelly equal.
"""
function Base.isapprox(dfa::AbstractDataFrame, dfb::AbstractDataFrame)
    function isapproxifexists(a, b)
        try
            return a ≈ b
        catch e
            e isa MethodError && return a == b
            rethrow(e)
        end
    end
    names(dfa) == names(dfb) && eltypes(dfa) == eltypes(dfb) && size(dfa) == size(dfb) &&
        all(isapproxifexists(dfa[!,n], dfb[!,n]) for n in names(dfa))
end

"""
    toposort(nodes, predf, nodef = identity)

Sort graph nodes topologically.

Provide `node` an iterator over node objects, `predf` a function giving predecessor of a node and
a `nodef` function optionally applied to the node in the result. [See](https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search).

# Examples
```@meta
DocTestSetup = quote
    using DataStructures: OrderedDict
    using MLGoodies: toposort
end
```

```jldoctest
julia> nodes = OrderedDict(2 => [1], 5 => [1, 3, 4], 4 => [1, 2], 6 => [5], 1 => [], 3 => [1], 20 => [5, 10], 10 => [1]);

julia> toposort(nodes, v -> last(v), v -> first(v))
8-element Array{Any,1}:
  1
  2
  3
  4
  5
  6
 10
 20
```
"""
function toposort(nodes, predf, nodef = identity)
    ndict = OrderedDict(nodef(n) => predf(n) for n in nodes)

    tmarks = Set()
    pmarks = Set()
    l = []
    function visit(n)
        n in pmarks && return
        n in tmarks && error("not a DAG!")
        push!(tmarks, n)
        for p in ndict[n]
            visit(p)
        end
        delete!(tmarks, n)
        push!(l, n)
        push!(pmarks, n)
    end

    while length(pmarks) < length(nodes)
        for n in keys(ndict)
            n ∉ pmarks && visit(n)
        end
    end
    l
end

#-------------------------------------------------------------------------------
# Serialization
#-------------------------------------------------------------------------------
function serialize_gzip(fname, data)
    open(fname, "w") do f
        gf = GzipCompressorStream(f)
        Serialization.serialize(gf, data)
        close(gf)
    end
end

function deserialize_gzip(fname)
    open(fname) do f
        gf = GzipDecompressorStream(f)
        data = Serialization.deserialize(gf)
        close(gf)
        data
    end
end

end


#! /usr/bin/env julia

using CSV
using DataFrames
using MAT
using Statistics

@info "Read $(ARGS[1])"

file = matopen(ARGS[1]); # "goodPoints.mat"

good_points = read(file, "goodPoints");
good_bases = vec(read(file, "goodBases"));

close(file)

@info "Read $(ARGS[2])"

gene_per_base = CSV.read(ARGS[2], header=0); # "cell_barcode_names.csv"
gene_per_base = Dict("$(row[2])" => row[3] for row in eachrow(gene_per_base));

df_spatial = DataFrame(good_points, [:x, :y, :z]);
df_spatial[!, :gene] = [gene_per_base[b] for b in good_bases];

CSV.write("molecules.csv", df_spatial)

if length(ARGS) > 2
    import NPZ
    import Images

    @info "Read $(ARGS[3])"
    Images.save("segmentation.tiff", UInt16.(NPZ.npzread(ARGS[3])["labels"]))
end


# IIF testCompareVariableFactors repeated in IIF

using DistributedFactorGraphs
using IncrementalInference
using Test

@testset "testing compare functions for variables and factors..." begin

    fg = initfg()

    v1 = addVariable!(fg, :x0, ContinuousScalar)
    f1 = addFactor!(fg, [:x0;], Prior(Normal()))

    @test compareVariable(v1,v1)
    @test compareFactor(f1,f1)

    v2 = addVariable!(fg, :x1, ContinuousScalar)
    f2 = addFactor!(fg, [:x0;:x1], LinearRelative(Normal(2.0, 0.1)))

    fg2 = deepcopy(fg)

    @test !compareVariable(v1,v2)

    # not testing different factors in this way
    # @test !compareFactor(f1,f2)

    @test compareAllVariables(fg, fg)

    @test compareAllVariables(fg, fg2)

    @test compareSimilarVariables(fg, fg)
    @test compareSimilarVariables(fg, fg2)

    @test compareSimilarFactors(fg, fg)
    @test compareSimilarFactors(fg, fg2)

    @test compareFactorGraphs(fg, fg)
    @test compareFactorGraphs(fg, fg2)

    tree = solveTree!(fg)

    x1a = getVariable(fg, :x0)
    x1b = getVariable(fg2, :x0)

    @test !compareVariable(x1a, x1b, skipsamples=false)

    @test !compareSimilarVariables(fg, fg2, skipsamples=false)
    @test !compareSimilarFactors(fg, fg2, skipsamples=false)

    @test compareFactorGraphs(fg, fg)
    @test !compareFactorGraphs(fg, fg2, skipsamples=false)

    ensureAllInitialized!(fg2)

    @test compareSimilarVariables(fg, fg2, skipsamples=true, skip=Symbol[:initialized;:inferdim;:ppeDict;:solvedCount])
    # fg2 has been solved, so it should fail on the estimate dictionary
    @test !compareSimilarVariables(fg, fg2, skipsamples=true, skip=Symbol[:initialized;:inferdim])

    tree = buildTreeReset!(fg2)

    @test compareSimilarFactors(fg, fg2, skipsamples=true, skipcompute=true)

    @test !compareSimilarFactors(fg, fg2, skipsamples=true, skipcompute=false)

    @test compareFactorGraphs(fg, fg2, skipsamples=true, skipcompute=true, skip=[:initialized;:inferdim;:ppeDict; :solvedCount])

end

@testset "test subgraph functions..." begin

    fg = initfg()

    addVariable!(fg, :x0, ContinuousScalar)
    addFactor!(fg, [:x0;], Prior(Normal()))

    addVariable!(fg, :x1, ContinuousScalar)
    addFactor!(fg, [:x0;:x1], LinearRelative(Normal(2.0, 0.1)))

    addVariable!(fg, :x2, ContinuousScalar)
    addFactor!(fg, [:x1;:x2], LinearRelative(Normal(4.0, 0.1)))

    addVariable!(fg, :l1, ContinuousScalar)
    addFactor!(fg, [:x1;:l1], LinearRelative(Rayleigh()))

    sfg = buildSubgraph(LightDFG, fg, [:x0;:x1])

    @warn "FIXME This is NOT supposed to pass"
    @test_skip compareFactorGraphs(fg, sfg, skip=[:labelDict;:addHistory;:logpath])
    # drawGraph(sfg)

end


gw = GridWorld(sx=2, sy=2, rs=[GridWorldState(1,1)], rv=[10.0])

pvec = fill(GridWorldAction(:left), 5)

solver = VectorSolver(pvec)

p = solve(solver, gw)

for s in iterator(states(gw))
    @test action(p, s) == GridWorldAction(:left)
end

p2 = VectorPolicy(gw, pvec)
for s in iterator(states(gw))
    @test action(p2, s) == GridWorldAction(:left)
end


# https://docs.julialang.org/en/v1/base/collections/#Base.Pair
karty = [Pair(7,"7"),Pair(8,"8"),Pair(9,"9"),Pair(10,"10"),Pair(11,"dolník"),Pair(12,"horník"),Pair(13,"kráľ"),Pair(14,"eso")]
farby_kariet = ["zeleň", "žaluď", "červeň", "guľa"]

# https://discourse.julialang.org/t/cleanest-way-to-generate-all-combinations-of-n-arrays/20127
# help?> Iterators.product
vsetky_karty = collect(Iterators.product(karty, farby_kariet))

najhodnotnejsia_karta = vsetky_karty[1,1:4]


#############################################################################
# SCS.jl
# Wrapper around the SCS solver https://github.com/cvxgrp/scs
# See http://github.com/JuliaOpt/SCS.jl
#############################################################################
# test/options.jl
# Tests the ability to pass options
#############################################################################

using MathProgBase

# The normal test
A = [1.0 1.0 0.0 0.0 0.0;
     0.0 1.0 0.0 0.0 1.0;
     0.0 0.0 1.0 1.0 1.0]
collb = [0.0, 0.0, 0.0, 0.0, 0.0]
obj   = [3.0, 4.0, 4.0, 9.0, 5.0]
rowub = [ 5.0,  3.0,  9.0]
s = SCSSolver()
m = MathProgBase.ConicModel(s)
MathProgBase.loadproblem!(m, -obj, A, rowub, [(:NonNeg,1:3)],[(:NonNeg,1:5)])
MathProgBase.optimize!(m)

@test isapprox(MathProgBase.getobjval(m), -99.0, atol=1e-6, rtol=0.0)
@test !isapprox(MathProgBase.getobjval(m), -99.0, atol=1e-7, rtol=0.0)

# With eps = 1e-12, solution should be far more accurate
s = SCSSolver(eps=1e-12)
m = MathProgBase.ConicModel(s)
MathProgBase.loadproblem!(m, -obj, A, rowub, [(:NonNeg,1:3)],[(:NonNeg,1:5)])
MathProgBase.optimize!(m)
@test isapprox(MathProgBase.getobjval(m), -99.0, atol=1e-10, rtol=0.0)

# With a warmstart from the eps = 1e-8 solution, solution should be extremely accurate even after 1 iteration
SCS.addoption!(m, :warm_start, true)
SCS.addoption!(m, :max_iters, 1)
MathProgBase.optimize!(m)
@test isapprox(MathProgBase.getobjval(m), -99.0, atol=1e-10, rtol=0.0)

# Now let's do the same warmstart, but on a new instance of the same problem
primal_sol = m.primal_sol
dual_sol = m.dual_sol
slack = m.slack
s = SCSSolver(max_iters=1)
m = MathProgBase.ConicModel(s)
MathProgBase.loadproblem!(m, -obj, A, rowub, [(:NonNeg,1:3)],[(:NonNeg,1:5)])
MathProgBase.setwarmstart!(m, primal_sol; dual_sol = dual_sol, slack = slack)
MathProgBase.optimize!(m)
@test isapprox(MathProgBase.getobjval(m), -99.0, atol=1e-10, rtol=0.0)

# tests for incorrect options
s = SCSSolver(eps=1e-12, epps=1.0)
m = MathProgBase.ConicModel(s)
MathProgBase.loadproblem!(m, -obj, A, rowub, [(:NonNeg,1:3)],[(:NonNeg,1:5)])

@test_throws ArgumentError MathProgBase.optimize!(m)

err = try
    MathProgBase.optimize!(m)
catch ex
    ex
end
@test err.msg == "Unrecognized option passed to SCS: epps;
Recognized options are: linear_solver, normalize, scale, rho_x, max_iters, eps, alpha, cg_rate, verbose, warm_start and acceleration_lookback."

# tests for incorrect options
s = SCSSolver(linear_solver="AAA", eps=1e-12, epps=1.0)
m = MathProgBase.ConicModel(s)
MathProgBase.loadproblem!(m, -obj, A, rowub, [(:NonNeg,1:3)],[(:NonNeg,1:5)])

@test_throws ArgumentError MathProgBase.optimize!(m)

err = try
    MathProgBase.optimize!(m)
catch ex
    ex
end

@test err.msg == "Unrecognized linear_solver passed to SCS: AAA;\nRecognized options are: SCS.Direct, SCS.Indirect."


const idx_t = Cint
const real_t = Cfloat

const METIS_NOPTIONS = 40

## Return codes
const METIS_OK           = Cint(1)  # Returned normally
const METIS_ERROR_INPUT  = Cint(-2) # Returned due to erroneous inputs and/or options
const METIS_ERROR_MEMORY = Cint(-3) # Returned due to insufficient memory
const METIS_ERROR        = Cint(-4) # Some other errors

struct MetisError <: Exception code::Cint end
function Base.showerror(io::IO, me::MetisError)
    print(io, "MetisError: ")
    if me.code == Metis.METIS_ERROR_INPUT
        print(io, "input error")
    elseif me.code == Metis.METIS_ERROR_MEMORY
        print(io, "could not allocate the required memory")
    else
        print(io, "unknown error")
    end
    print(io, " (error code $(me.code)).")
end

## Operation type codes
const METIS_OP_PMETIS = Cint(0)
const METIS_OP_KMETIS = Cint(1)
const METIS_OP_OMETIS = Cint(2)

## (1-based) positions in options vector
const METIS_OPTION_PTYPE     = 1
const METIS_OPTION_OBJTYPE   = 2
const METIS_OPTION_CTYPE     = 3
const METIS_OPTION_IPTYPE    = 4
const METIS_OPTION_RTYPE     = 5
const METIS_OPTION_DBGLVL    = 6
const METIS_OPTION_NITER     = 7
const METIS_OPTION_NCUTS     = 8
const METIS_OPTION_SEED      = 9
const METIS_OPTION_NO2HOP    = 10
const METIS_OPTION_MINCONN   = 11
const METIS_OPTION_CONTIG    = 12
const METIS_OPTION_COMPRESS  = 13
const METIS_OPTION_CCORDER   = 14
const METIS_OPTION_PFACTOR   = 15
const METIS_OPTION_NSEPS     = 16
const METIS_OPTION_UFACTOR   = 17
const METIS_OPTION_NUMBERING = 18
const METIS_OPTION_HELP      = 19
const METIS_OPTION_TPWGTS    = 20
const METIS_OPTION_NCOMMON   = 21
const METIS_OPTION_NOOUTPUT  = 22
const METIS_OPTION_BALANCE   = 23
const METIS_OPTION_GTYPE     = 24
const METIS_OPTION_UBVEC     = 25

## Partitioning Schemes
const METIS_PTYPE_RB   = Cint(0)
const METIS_PTYPE_KWAY = Cint(1)

## Graph types for meshes
const METIS_GTYPE_DUAL  = Cint(0)
const METIS_GTYPE_NODAL = Cint(1)

## Coarsening Schemes
const METIS_CTYPE_RM   = Cint(0)
const METIS_CTYPE_SHEM = Cint(1)

## Initial partitioning schemes
const METIS_IPTYPE_GROW    = Cint(0)
const METIS_IPTYPE_RANDOM  = Cint(1)
const METIS_IPTYPE_EDGE    = Cint(2)
const METIS_IPTYPE_NODE    = Cint(3)
const METIS_IPTYPE_METISRB = Cint(4)

## Refinement schemes
const METIS_RTYPE_FM        = Cint(0)
const METIS_RTYPE_GREEDY    = Cint(1)
const METIS_RTYPE_SEP2SIDED = Cint(2)
const METIS_RTYPE_SEP1SIDED = Cint(3)

## Debug levels (bit positions)
const METIS_DBG_INFO       = Cint(1)    # Shows various diagnostic messages
const METIS_DBG_TIME       = Cint(2)    # Perform timing analysis
const METIS_DBG_COARSEN    = Cint(4)    # Show the coarsening progress
const METIS_DBG_REFINE     = Cint(8)    # Show the refinement progress
const METIS_DBG_IPART      = Cint(16)   # Show info on initial partitioning
const METIS_DBG_MOVEINFO   = Cint(32)   # Show info on vertex moves during refinement
const METIS_DBG_SEPINFO    = Cint(64)   # Show info on vertex moves during sep refinement
const METIS_DBG_CONNINFO   = Cint(128)  # Show info on minimization of subdomain connectivity
const METIS_DBG_CONTIGINFO = Cint(256)  # Show info on elimination of connected components
const METIS_DBG_MEMORY     = Cint(2048) # Show info related to wspace allocation

## Types of objectives
const METIS_OBJTYPE_CUT  = Cint(0)
const METIS_OBJTYPE_VOL  = Cint(1)
const METIS_OBJTYPE_NODE = Cint(2)

## Metis C API
function METIS_PartGraphRecursive(nvtxs, ncon, xadj, adjncy, vwgt, vsize, adjwgt, nparts,
                                  tpwgts, ubvec, options, edgecut, part)
    r = ccall((:METIS_PartGraphRecursive, libmetis), Cint,
              (Ref{idx_t}, Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
               Ptr{idx_t}, Ref{idx_t}, Ptr{real_t}, Ptr{real_t}, Ptr{idx_t}, Ptr{idx_t},
               Ptr{idx_t}),
              nvtxs, ncon, xadj, adjncy, vwgt, vsize, adjwgt, nparts, tpwgts, ubvec,
              options, edgecut, part)
    r == METIS_OK || throw(MetisError(r))
    return
end

function METIS_PartGraphKway(nvtxs, ncon, xadj, adjncy, vwgt, vsize, adjwgt, nparts,
                             tpwgts, ubvec, options, edgecut, part)
    r = ccall((:METIS_PartGraphKway, libmetis), Cint,
              (Ref{idx_t}, Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
               Ptr{idx_t}, Ref{idx_t}, Ptr{real_t}, Ptr{real_t}, Ptr{idx_t}, Ptr{idx_t},
               Ptr{idx_t}),
              nvtxs, ncon, xadj, adjncy, vwgt, vsize, adjwgt, nparts, tpwgts, ubvec,
              options, edgecut, part)
    r == METIS_OK || throw(MetisError(r))
    return
end

function METIS_MeshToDual(ne, nn, eptr, eind, ncommon, numflag, r_xadj, r_adjncy)
    r = ccall((:METIS_MeshToDual, libmetis), Cint,
              (Ref{idx_t}, Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ref{idx_t}, Ref{idx_t},
               Ptr{idx_t}, Ptr{idx_t}),
              ne, nn, eptr, eind, ncommon, numflag, r_xadj, r_adjncy)
    r == METIS_OK || throw(MetisError(r))
    return
end

# function METIS_MeshToNodal(ne, nn, eptr, eind, numflag, r_xadj, r_adjncy)
#     r = ccall((:METIS_MeshToNodal, libmetis), Cint,
#               (Ref{idx_t}, Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ref{idx_t}, Ptr{idx_t},
#                Ptr{idx_t}),
#               ne, nn, eptr, eind, numflag, r_xadj, r_adjncy)
#     r == METIS_OK || throw(MetisError(r))
#     return
# end

# function METIS_PartMeshNodal(ne, nn, eptr, eind, vwgt, vsize, nparts, tpwgts, options,
#                              objval, epart, npart)
#     r = ccall((:METIS_PartMeshNodal, libmetis), Cint,
#               (Ref{idx_t}, Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
#                Ref{idx_t}, Ptr{real_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}),
#               ne, nn, eptr, eind, vwgt, vsize, nparts, tpwgts, options, objval, epart,
#               npart)
#     r == METIS_OK || throw(MetisError(r))
#     return
# end

# function METIS_PartMeshDual(ne, nn, eptr, eind, vwgt, vsize, ncommon, nparts, tpwgts,
#                             options, objval, epart, npart)
#     r = ccall((:METIS_PartMeshDual, libmetis), Cint,
#               (Ref{idx_t}, Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
#                Ref{idx_t}, Ref{idx_t}, Ptr{real_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
#                Ptr{idx_t}),
#               ne, nn, eptr, eind, vwgt, vsize, ncommon, nparts, tpwgts, options, objval,
#               epart, npart)
#     r == METIS_OK || throw(MetisError(r))
#     return
# end

function METIS_NodeND(nvtxs, xadj, adjncy, vwgt, options, perm, iperm)
    r = ccall((:METIS_NodeND, libmetis), Cint,
              (Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
               Ptr{idx_t}),
              nvtxs, xadj, adjncy, vwgt, options, perm, iperm)
    r == METIS_OK || throw(MetisError(r))
    return
end

function METIS_Free(ptr)
    r = ccall((:METIS_Free, libmetis), Cint, (Ptr{idx_t},), ptr)
    r == METIS_OK || throw(MetisError(r))
    return
end

# function METIS_SetDefaultOptions(options)
#     r = ccall((:METIS_SetDefaultOptions, libmetis), Cint, (Ptr{idx_t},), options)
#     r == METIS_OK || throw(MetisError(r))
#     return
# end

# function METIS_NodeNDP(nvtxs, xadj, adjncy, vwgt, npes, options, perm, iperm, sizes)
#     r = ccall((:METIS_NodeNDP, libmetis), Cint,
#               (Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
#                Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}),
#               nvtxs, xadj, adjncy, vwgt, npes, options, perm, iperm, sizes)
#     r == METIS_OK || throw(MetisError(r))
#     return
# end

function METIS_ComputeVertexSeparator(nvtxs, xadj, adjncy, vwgt, options, sepsize, part)
    r = ccall((:METIS_ComputeVertexSeparator, libmetis), Cint,
              (Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
               Ptr{idx_t}),
              nvtxs, xadj, adjncy, vwgt, options, sepsize, part)
    r == METIS_OK || throw(MetisError(r))
    return
end

# function METIS_NodeRefine(nvtxs, xadj, vwgt, adjncy, where, hmarker, ubfactor)
#     r = ccall((:METIS_NodeRefine, libmetis), Cint,
#               (Ref{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t}, Ptr{idx_t},
#                Ptr{real_t}),
#               nvtxs, xadj, vwgt, adjncy, where, hmarker, ubfactor)
#     r == METIS_OK || throw(MetisError(r))
#     return
# end


"""
Main module for `Oceananigans.jl` -- a Julia software for fast, friendly, flexible,
data-driven, ocean-flavored fluid dynamics on CPUs and GPUs.
"""
module Oceananigans

if VERSION < v"1.6"
    error("This version of Oceananigans.jl requires Julia v1.6 or newer.")
end

export
    # Architectures
    CPU, GPU,

    # Logging
    OceananigansLogger,

    # Grids
    Center, Face,
    Periodic, Bounded, Flat,
    FullyConnected, LeftConnected, RightConnected,
    RectilinearGrid,
    LatitudeLongitudeGrid,
    ConformalCubedSphereFaceGrid,
    xnodes, ynodes, znodes, nodes,

    # Immersed boundaries
    ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom, ImmersedBoundaryCondition,

    # Advection schemes
    CenteredSecondOrder, CenteredFourthOrder, UpwindBiasedFirstOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder, WENO5,
    VectorInvariant, EnergyConservingScheme, EnstrophyConservingScheme,

    # Boundary conditions
    BoundaryCondition,
    FluxBoundaryCondition, ValueBoundaryCondition, GradientBoundaryCondition, OpenBoundaryCondition,
    FieldBoundaryConditions,

    # Fields and field manipulation
    Field, CenterField, XFaceField, YFaceField, ZFaceField,
    Average, Integral, Reduction, BackgroundField,
    interior, set!, compute!, regrid!, location,

    # Forcing functions
    Forcing, Relaxation, LinearTarget, GaussianMask, AdvectiveForcing,

    # Coriolis forces
    FPlane, ConstantCartesianCoriolis, BetaPlane, NonTraditionalBetaPlane,

    # BuoyancyModels and equations of state
    Buoyancy, BuoyancyTracer, SeawaterBuoyancy,
    LinearEquationOfState, TEOS10,
    BuoyancyField,

    # Surface wave Stokes drift via Craik-Leibovich equations
    UniformStokesDrift,

    # Turbulence closures
    VerticalScalarDiffusivity,
    HorizontalScalarDiffusivity,
    ScalarDiffusivity,
    VerticalScalarBiharmonicDiffusivity,
    HorizontalScalarBiharmonicDiffusivity,
    ScalarBiharmonicDiffusivity,
    SmagorinskyLilly,
    AnisotropicMinimumDissipation,
    ConvectiveAdjustmentVerticalDiffusivity,
    RiBasedVerticalDiffusivity,
    IsopycnalSkewSymmetricDiffusivity,
    FluxTapering,
    VerticallyImplicitTimeDiscretization,

    # Lagrangian particle tracking
    LagrangianParticles,

    # Models
    NonhydrostaticModel,
    HydrostaticFreeSurfaceModel,
    ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation,
    PressureField,
    fields,

    # Hydrostatic free surface model stuff
    VectorInvariant, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface,
    HydrostaticSphericalCoriolis,
    PrescribedVelocityFields,

    # Time stepping
    Clock, TimeStepWizard, time_step!,

    # Simulations
    Simulation, run!, Callback, iteration, stopwatch,
    iteration_limit_exceeded, stop_time_exceeded, wall_time_limit_exceeded,
    erroring_NaNChecker!,

    # Diagnostics
    StateChecker, CFL, AdvectiveCFL, DiffusiveCFL,

    # Output writers
    NetCDFOutputWriter, JLD2OutputWriter, Checkpointer,
    TimeInterval, IterationInterval, AveragedTimeInterval, SpecifiedTimes,
    AndSchedule, OrSchedule,

    # Output readers
    FieldTimeSeries, FieldDataset, InMemory, OnDisk,

    # Abstract operations
    ∂x, ∂y, ∂z, @at, KernelFunctionOperation,

    # MultiRegion and Cubed sphere
    MultiRegionGrid, XPartition,
    ConformalCubedSphereGrid,

    # Utils
    prettytime, apply_regionally!, construct_regionally, @apply_regionally, MultiRegionObject

using Printf
using Logging
using Statistics
using LinearAlgebra
using CUDA
using Adapt
using DocStringExtensions
using OffsetArrays
using FFTW
using JLD2
using NCDatasets

using Base: @propagate_inbounds
using Statistics: mean

import Base:
    +, -, *, /,
    size, length, eltype,
    iterate, similar, show,
    getindex, lastindex, setindex!,
    push!

#####
##### Abstract types
#####

"""
    AbstractModel

Abstract supertype for models.
"""
abstract type AbstractModel{TS} end

"""
    AbstractDiagnostic

Abstract supertype for diagnostics that compute information from the current
model state.
"""
abstract type AbstractDiagnostic end

"""
    AbstractOutputWriter

Abstract supertype for output writers that write data to disk.
"""
abstract type AbstractOutputWriter end

#####
##### Place-holder functions
#####

function run_diagnostic! end
function write_output! end
function location end
function instantiated_location end
function tupleit end
function fields end
function prognostic_fields end
function tracer_tendency_kernel_function end

#####
##### Include all the submodules
#####

# Basics
include("Architectures.jl")
include("Units.jl")
include("Grids/Grids.jl")
include("Utils/Utils.jl")
include("Logger.jl")
include("Operators/Operators.jl")
include("BoundaryConditions/BoundaryConditions.jl")
include("Fields/Fields.jl")
include("AbstractOperations/AbstractOperations.jl")
include("Advection/Advection.jl")
include("Solvers/Solvers.jl")
include("Distributed/Distributed.jl")

# Physics, time-stepping, and models
include("Coriolis/Coriolis.jl")
include("BuoyancyModels/BuoyancyModels.jl")
include("StokesDrift.jl")
include("TurbulenceClosures/TurbulenceClosures.jl")
include("LagrangianParticleTracking/LagrangianParticleTracking.jl")
include("Forcings/Forcings.jl")

include("ImmersedBoundaries/ImmersedBoundaries.jl")
include("TimeSteppers/TimeSteppers.jl")
include("Models/Models.jl")

# Output and Physics, time-stepping, and models
include("Diagnostics/Diagnostics.jl")
include("OutputWriters/OutputWriters.jl")
include("OutputReaders/OutputReaders.jl")
include("Simulations/Simulations.jl")

# Abstractions for distributed and multi-region models
include("MultiRegion/MultiRegion.jl")
include("CubedSpheres/CubedSpheres.jl")

#####
##### Needed so we can export names from sub-modules at the top-level
#####

using .Logger
using .Architectures
using .Utils
using .Advection
using .Grids
using .BoundaryConditions
using .Fields
using .Coriolis
using .BuoyancyModels
using .StokesDrift
using .TurbulenceClosures
using .LagrangianParticleTracking
using .Solvers
using .Forcings
using .ImmersedBoundaries
using .Distributed
using .Models
using .TimeSteppers
using .Diagnostics
using .OutputWriters
using .OutputReaders
using .Simulations
using .AbstractOperations
using .MultiRegion
using .CubedSpheres

function __init__()
    threads = Threads.nthreads()
    if threads > 1
        @info "Oceananigans will use $threads threads"

        # See: https://github.com/CliMA/Oceananigans.jl/issues/1113
        FFTW.set_num_threads(4*threads)
    end

    if CUDA.has_cuda()
        @debug "CUDA-enabled GPU(s) detected:"
        for (gpu, dev) in enumerate(CUDA.devices())
            @debug "$dev: $(CUDA.name(dev))"
        end

        CUDA.allowscalar(false)
    end
end

end # module


country() = executor(data["en"]["faker"]["address"]["country"])
city_prefix() = executor(data["en"]["faker"]["address"]["city_prefix"])
city_suffix()= executor(data["en"]["faker"]["address"]["city_suffix"])
street_prefix()= executor(data["en"]["faker"]["address"]["street_prefix"])
street_suffix()= executor(data["en"]["faker"]["address"]["street_suffix"])
secondary_address() = bothify(executor(data["en"]["faker"]["address"]["secondary_address"]))
state() = executor(data["en"]["faker"]["address"]["state"])
state_abbr() = executor(data["en"]["faker"]["address"]["state_abbr"])
municipality() = executor(data["en"]["faker"]["address"]["municipality"])
street_name() = executor(data["en"]["faker"]["address"]["street_name"])
street_address() = executor(data["en"]["faker"]["address"]["street_address"])
address() = ("$(street_address()) $(city()), $(state_abbr()) $(postcode())")
postcode() = bothify(executor(data["en"]["faker"]["address"]["postcode"]))
building_number() = bothify(executor(data["en"]["faker"]["address"]["building_number"]))
city() = executor(data["en"]["faker"]["address"]["city"])

function geo_coordinate(;center="None", radius=0.001)

        if center=="None"
            center=0
        end

        if center == 0
            return Float16(rand(-17.99 : 17.99))
        else
            return Float16(rand( center - radius : center + radius))
        end
end

latitude() = geo_coordinate() / 2.0
longitude() = geo_coordinate()/ 1.0


# Functions for plotting molecular dynamics simulation results

PLOTTING_TIME_UNIT = u"ps"
PLOTTING_TEMP_UNIT = u"K"
PLOTTING_ENERGY_UNIT = u"hartree"

# Extract (timestep, time) tuples at stride intervals from a simulation result
function plotting_time_range(result::MolecularDynamicsResult, stride::Integer)
    [(t, uconvert(PLOTTING_TIME_UNIT, time)) for (t, time) ∈ enumerate(get_time_range(result)) if (t - 1) % stride == 0]
end

"""
    plot_temperature(result::MolecularDynamicsResult, stride::Integer)::Plot

Plot the temperature of a `MolecularDynamicsResult` against time, sampling every `stride` points.
"""
function plot_temperature(result::MolecularDynamicsResult, stride::Integer)
    N = length(get_system(result))
    p = plot(
        title = "Temperature during Simulation [n = $(N)]",
        xlab = "Time",
        ylab = "Temperature",
    )
    plot_temperature!(p, result, stride, true)
end
"""
    plot_temperature!(p::Plot, result::MolecularDynamicsResult, stride::Integer, first_plot::Bool = false)::Plot

Plot the temperature of a `MolecularDynamicsResult` against time, sampling every `stride` points.
Adds the new line to an existing plot and update the legend only if it is the first plot.
If it is not the first plot, adds a vertical line to differentiate the segments of the simulation in the plot.
"""
function plot_temperature!(p::Plot, result::MolecularDynamicsResult, stride::Integer, first_plot::Bool = false)
    time_range = plotting_time_range(result, stride)
    if (!first_plot)
        vline!(
            p,
            [time_range[1][2]],
            label = false,
            color = :black,
            linestyle = :dot,
            lw = 2
        )
    end
    plot!(
        p,
        [time for (t, time) ∈ time_range],
        [uconvert(PLOTTING_TEMP_UNIT, temperature(result, t)) for (t, time) ∈ time_range],
        label = first_plot ? "Simulation Temperature" : nothing,
        color = 1,
    )
    reference_temp = uconvert(PLOTTING_TEMP_UNIT, reference_temperature(result))
    if (!ismissing(reference_temp))
        plot!(
            p,
            [time for (t, time) ∈ time_range],
            [reference_temp for (t, time) ∈ time_range],
            label = first_plot ? "Reference Temperature" : nothing,
            color = 2,
            linestyle = :dash,
            lw = 2
        )
    end
    p
end

"""
    plot_energy(result::MolecularDynamicsResult, stride::Integer)::Plot

Plot the kinetic, potential, and total energy of a `MolecularDynamicsResult` against time, sampling every `stride` points.
"""
function plot_energy(result::MolecularDynamicsResult, stride::Integer)
    N = length(get_system(result))
    p = plot(
        title = "Energy during Simulation [n = $(N)]",
        xlab = "Time",
        ylab = "Energy",
        legend = :right
    )
    plot_energy!(p, result, stride, true)
end
"""
    plot_energy!(p::Plot, result::MolecularDynamicsResult, stride::Integer, first_plot::Bool = false)::Plot

Plot the kinetic, potential, and total energy of a `MolecularDynamicsResult` against time, sampling every `stride` points.
Add the new lines to an existing plot and update the legend only if it is the first plot.
If it is not the first plot, add a vertical line to differentiate the segments of the simulation in the plot.
"""
function plot_energy!(p::Plot, result::MolecularDynamicsResult, stride::Integer, first_plot::Bool = false)
    time_range = plotting_time_range(result, stride)
    if (!first_plot)
        vline!(p, [time_range[1][2]], label = false, color = :black, linestyle = :dot, lw = 2)
    end
    plot!(
        p,
        [time for (t, time) ∈ time_range],
        [uconvert(PLOTTING_ENERGY_UNIT, kinetic_energy(result, t)) for (t, time) ∈ time_range],
        label = first_plot ? "Kinetic Energy" : nothing,
        color = 2
    )
    plot!(
        p,
        [time for (t, time) ∈ time_range],
        [uconvert(PLOTTING_ENERGY_UNIT, potential_energy(result, t)) for (t, time) ∈ time_range],
        label = first_plot ? "Potential Energy" : nothing,
        color = 1
    )
    plot!(
        p,
        [time for (t, time) ∈ time_range],
        [uconvert(PLOTTING_ENERGY_UNIT, total_energy(result, t)) for (t, time) ∈ time_range],
        label = first_plot ? "Total Energy" : nothing,
        color = 3
    )
end

"""
    plot_rdf(result::MolecularDynamicsResult, σ::Real, sample_fraction::Float64 = 1.0)::Plot

Plot the radial distribution function of a `MolecularDynamicsResult` sampling only the trailing `sample_fraction` of the timesteps.
Use σ (from Lennard Jones) as a normalization factor for the radius (assumed to be in atomic units).
"""
function plot_rdf(result::MolecularDynamicsResult, σ::Real, sample_fraction::Float64 = 1.0)
    plot_rdf(result, σ * LENGTH_UNIT, sample_fraction)
end
"""
    plot_rdf(result::MolecularDynamicsResult, σ::Unitful.Length, sample_fraction::Float64 = 1.0)::Plot

Plot the radial distribution function of a `MolecularDynamicsResult` sampling only the trailing `sample_fraction` of the timesteps.
Use σ (from Lennard Jones) as a normalization factor for the radius.
"""
function plot_rdf(result::MolecularDynamicsResult, σ::Unitful.Length, sample_fraction::Float64 = 1.0)
    @assert 0 < sample_fraction ≤ 1
    N = length(get_system(result))
    T = length(get_time_range(result)) - 1
    rs, grf = rdf(result, sample_fraction)
    @assert length(rs) == length(grf)
    plot(
        title = "Radial Distribution Function [n = $(N)] [T = $(T)]",
        xlab = "Distance r/σ",
        ylab = "Radial Distribution g(r)",
        legend = false
    )
    plot!(rs / austrip(σ), grf)
end


module NetworkModels

using Graphs
using Distributions
using Pipe
using LinearAlgebra

end


module PrettyPrint
export pprint, pprintln, pformat, is_simple, is_simple_t, is_atom_t, is_static_t, pp_impl, pp_impl_dynamic, pp_impl_static, MaxIndentExpected

struct PPPair{FieldName, FieldContent}
    name :: FieldName
    sep :: String
    content :: FieldContent
end

struct PPMetaSymbol{S}
    s :: S
end

Base.show(io::IO, p::PPMetaSymbol) = Base.print(io, p.s)

is_simple_t(a) = isprimitivetype(a)
is_simple_t(::Type{<:Number}) = true
is_simple_t(::Type{<:AbstractString}) = true
is_simple_t(::Type{<:AbstractChar}) = true
is_simple_t(::Type{<:Symbol}) = true
is_simple_t(::Type{<:Complex}) = true
is_simple_t(::Type{<:PPMetaSymbol}) = true
is_simple_t(::Type{Nothing}) = true
is_simple_t(::Type{PPPair{A1, A2}}) where {A1, A2} = is_simple_t(A1) && is_simple_t(A2)
is_simple(::T) where T = is_simple_t(T)

is_atom_t(_) = false
is_atom_t(::Type{<:Number}) = true
is_atom_t(::Type{<:AbstractString}) = true
is_atom_t(::Type{<:AbstractChar}) = true
is_atom_t(::Type{<:Symbol}) = true
is_atom_t(::Type{<:Complex}) = true
is_atom_t(::Type{<:PPMetaSymbol}) = true
is_atom_t(::Type{Nothing}) = true
is_atom(::T) where T = is_atom_t(T)

is_static_t(a) = isprimitivetype(a)
is_static_t(::Type{String}) = true
is_static_t(::Type{Symbol}) = true
is_static_t(::Type{<:Complex}) = true

const Indentation = Int
const MaxIndentExpected = Ref{Int}(30)

print_repeat(n::Int, io::IO, e) = for i = 1:n
    print(io, e)
end

function pprint_for_seq(io, left::String, right::String, seq :: Union{Vector, Tuple}, indent::Int)::Indentation
    print(io, left)
    n = length(seq)
    is_empty = n === 0
    if is_empty
        print(io, right)
        return indent + length(left) + length(right)
    end

    if all(is_simple(e) for e in seq)
        indent += length(left)
        indent′ = indent
        max_indent = MaxIndentExpected[]
        indent′ = pp_impl(io, seq[1], indent)
        if n == 1
            print(io, ",")
            indent′ += 2
        end
        for i in 2:n
            e = seq[i]
            print(io, ", ")

            if indent′ > max_indent
                print(io, '\n')
                indent′ = indent
                print_repeat(indent, io, ' ')
            else
                indent′ += 2
            end

            indent′ = pp_impl(io, e, indent′)
        end

        print(io, right)
        indent′ += length(right)

        return indent′
    end

    print(io, '\n')
    indent′ = indent + 2
    for e in seq
        print_repeat(indent′, io, ' ')
        pp_impl(io, e, indent′)
        print(io, ',')
        print(io, '\n')
    end
    print_repeat(indent, io, ' ')
    print(io, right)
    return indent + length(right)
end

"""
The implementation of `pprint` and `pformat`.

You can extend pretty print support via such boilerplate:
```
    function PrettyPrint.pp_impl(io, data::YourType, indent::Int)
        print(io, data)
    end
```
"""
function pp_impl(io, seq::Vector, indent::Indentation)::Indentation
    pprint_for_seq(io, "[", "]", seq, indent)
end

function pp_impl(io, seq::Tuple, indent::Indentation)::Indentation
    pprint_for_seq(io, "(", ")", seq, indent)
end

function pp_impl(io, seq::AbstractSet, indent::Indentation)::Indentation
    pprint_for_seq(io, "{", "}", collect(seq), indent)
end

function pp_impl(io, seq::AbstractDict, indent::Indentation)::Indentation
    pprint_for_seq(io, "{", "}", [PPPair(k, " : ", v) for (k, v) in seq], indent)
end

function pp_impl(io, p::PPPair{K, V}, indent)::Indentation where {K, V}
    indent = pp_impl(io, p.name, indent)
    print(io, p.sep)
    indent += length(p.sep)
    pp_impl(io, p.content, indent)
end

function pp_impl(io, p::Pair, indent::Indentation)::Indentation
    pp_impl(io, PPPair(p.first, " => ", p.second), indent)
end

pp_impl(io, p, indent)::Indentation =
    if is_static_t(typeof(p))
        pp_impl_static(io, p, indent)
    else
        pp_impl_dynamic(io, p, indent)
    end

@generated function pp_impl_static(io, data::T, indent)::Indentation where T
    if is_atom_t(T)
        return quote s = repr(data); print(io, s); length(s) + indent  end
    end
    fields = fieldnames(T)
    exp = Expr(:tuple)
    for field in fields
        push!(exp.args, Expr(:call, PPPair, PPMetaSymbol(field), "=", :(data.$field)))
    end
    left = string(T) * "("
    right = ")"
    :(pprint_for_seq(io, $left, $right, $exp, indent))
end

function pp_impl_dynamic(io, data, indent)::Indentation
    if is_atom(data)
        s = repr(data)
        print(io, s)
        return length(s) + indent
    end

    fields = propertynames(data)
    left = string(typeof(data)) * "("
    right = ")"
    seq = [PPPair(PPMetaSymbol(field), "=", getproperty(data, field)) for field in fields]
    pprint_for_seq(io, left, right, seq, indent)
end

"""
pretty print to IO.
"""
function pprint(io, data, indent)
    pp_impl(io, data, indent)
    nothing
end

function pprint(io, data)
    pprint(io, data, 0)
end

function pprint(data)
    pprint(stdout, data)
end

"""
return pretty formatted string
"""
function pformat(data)
    io = IOBuffer()
    pprint(io, data)
    String(take!(io))
end

pp_impl(io, s :: String, indent) = begin
    s = repr(s)
    print(io, s)
    length(s) + indent
end

pprintln(io, s) = begin
    pprint(io, s)
    print(io, '\n')
end

pprintln(s) = pprintln(stdout, s)

@deprecate pprint(io, data, indent, newline) pprint(io, data, indent)
@deprecate pp_impl(io, data, indent, newline) pp_impl(io, data, indent)

end # module


using DataStructures

"""
vl_graph_plot(
    df_nodes,
    df_edges,
    node_label       = true,
    node_labelsize   = 10,
    tooltip          = true,
    node_size        = 500,
    node_color       = "#9ecae9",
    node_sizefield   = nothing,
    node_colorfield  = nothing,
    node_colorscheme = "blues",
    node_opacity     = 1.0,
    edge_opacity     = 0.5,
    edge_weightfield = nothing,
    width            = 600,
    height           = 400
    )
Contains the VegaLite specification for producing the plot.
"""
function vl_graph_plot(
    graph_nodes,
    graph_edges;
    ew               = false,
    node_label       = true,
    node_labelsize   = 10,
    tooltip          = true,
    node_size        = 500,
    node_color       = "#9ecae9",
    node_sizefield   = nothing,
    node_colorfield  = nothing,
    node_sizefieldtype   = "q",
    node_colorfieldtype  = "q",
    node_colorscheme = "blues",
    node_opacity     = 1.0,
    edge_opacity     = 0.5,
    edge_weightfield = nothing,
    edge_weightfieldtype = "o",
    width            = 600,
    height           = 400
    )

  # Adjusting for extra fields
  if node_sizefield  != nothing
        graph_nodes[!,:node_sizefield] = node_sizefield
        node_sizefield = :node_sizefield
  end
  if node_colorfield  != nothing
        graph_nodes[!,:node_colorfield] = node_colorfield
        node_colorfield = :node_colorfield
  end

  if edge_weightfield  != nothing
        ew = true
        graph_edges[!,:ew] = edge_weightfield
        edge_weightfield = :edge_weightfield

  end

  # vegalite.jl plotting specification
  # v1 - plotting nodes
    v1 =@vlplot(
    data = graph_nodes,
    layer=[
    {
        mark={"type"=:circle,clip=false,
              opacity=node_opacity,
              size=node_size,
              color=node_color},
        # There is some bug with VegaLite
        # when trying to put [] aroung tooltip
        #  tooltip={field="keywords","type"="nominal"},
        x={"node_x:q",axis=nothing},
        y={"node_y:q",axis=nothing},
        size={field=node_sizefield,
              type=node_sizefieldtype,
              scale={rangeMin=100},
              legend=nothing},
        color={field=node_colorfield,
               type=node_colorfieldtype,
               scale={scheme=node_colorscheme}},
        width=width,
        height=height,
        selection={
            grid={
                    type=:interval,
                    bind=:scales,
                    }}
        },
    {
        mark={"type"=:text,clip=false,opacity=1,size=node_labelsize},
        x={"node_x:q",axis=nothing},
        y={"node_y:q",axis=nothing},
        text={"keywords","type"="nominal"}
    }]
    )

  # v2 - plotting edges
    v2 =@vlplot(
        data = graph_edges,
        layer=[
        {
          mark={"type"=:line,color="gray",clip=false,size=0.8,opacity=edge_opacity},
          encoding={
          x={"edges_x:q",axis=nothing},
          y={"edges_y:q",axis=nothing},
          detail={"pairs:o"}},
          width=width,
          height=height
        }
              ]
    );

    # Adjusting fields to avoid warnings when plotting
    if tooltip == true
      v1.layer[1]["encoding"]["tooltip"] = OrderedDict[OrderedDict("field"=>"keywords","type"=>"nominal")]
    end

    if ew == true
      v2.layer[1]["encoding"]["size"] = OrderedDict("field"=>"ew","type"=>edge_weightfieldtype,"legend"=>nothing)
    end

    if node_label == false
      deleteat!(v1.layer,2)
    end

    if node_sizefield  == nothing
        delete!(v1.layer[1]["encoding"],"size")
    end
    if node_colorfield == nothing
        delete!(v1.layer[1]["encoding"],"color")
    end

    graph_plot = @vlplot(view={stroke=nothing})+v2+v1
    return graph_plot
end


# convert_a3d

# Function that allows federation by setting the `output_format`  in the Stanmodel().

function convert_a3d(a3d_array, cnames, ::Val{:mcmcchains})
  pi = filter(p -> length(p) > 2 && p[end-1:end] == "__", cnames)
  p = filter(p -> !(p in  pi), cnames)

  MCMCChains.Chains(a3d_array,
    cnames,
    Dict(
      :parameters => p,
      :internals => pi
    )
  )
end