text:
"""
    struct HamiltonianMC <: MCMCAlgorithm

Hamiltonian Monte Carlo sampling algorithm.

The gradient of the target density is computed via auto-differentiation.

* Note: The fields of `HamiltonianMC` are still subject to change, and not
yet part of stable public API!*

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct HamiltonianMC <: MCMCAlgorithm
    metric::HMCMetric = DiagEuclideanMetric()
    gradient::Module = ForwardDiff
    integrator::HMCIntegrator = LeapfrogIntegrator()
    proposal::HMCProposal = NUTS()
    adaptor::HMCAdaptor = StanHMCAdaptor()
end

export HamiltonianMC

bat_default(::Type{MCMCSampling}, ::Val{:trafo}, mcalg::HamiltonianMC) = PriorToGaussian()

bat_default(::Type{MCMCSampling}, ::Val{:nsteps}, mcalg::HamiltonianMC, trafo::AbstractDensityTransformTarget, nchains::Integer) = 10^4

get_mcmc_tuning(algorithm::HamiltonianMC) = MCMCNoOpTuning()

# MCMCIterator subtype for HamiltonianMC
mutable struct AHMCIterator{
    AL<:HamiltonianMC,
    D<:AbstractDensity,
    R<:AbstractRNG,
    PR<:RNGPartition,
    SV<:DensitySampleVector,
    I<:AdvancedHMC.AbstractIntegrator,
    P<:AdvancedHMC.AbstractProposal,
    A<:AdvancedHMC.AbstractAdaptor,
    H<:AdvancedHMC.Hamiltonian,
} <: MCMCIterator
    algorithm::AL
    density::D
    rng::R
    rngpart_cycle::PR
    info::MCMCIteratorInfo
    samples::SV
    nsamples::Int64
    stepno::Int64
    hamiltonian::H
    transition::AdvancedHMC.Transition
    integrator::I
    proposal::P
    adaptor::A
end

function AHMCIterator(
    rng::AbstractRNG,
    algorithm::MCMCAlgorithm,
    density::AbstractDensity,
    info::MCMCIteratorInfo,
    x_init::AbstractVector{P},
) where {P<:Real}
    stepno::Int64 = 0

    npar = totalndof(density)

    params_vec = Vector{P}(undef, npar)
    params_vec .= x_init
    !(params_vec in var_bounds(density)) && throw(ArgumentError("Parameter(s) out of bounds"))

    log_posterior_value = eval_logval(density, params_vec, strict = true)

    T = typeof(log_posterior_value)
    W = Float64 # ToDo: Support other sample weight types

    sample_info = AHMCSampleID(info.id, info.cycle, 1, CURRENT_SAMPLE, 0.0, 0, false, 0.0)
    current_sample = DensitySample(params_vec, log_posterior_value, one(W), sample_info, nothing)
    samples = DensitySampleVector{Vector{P},T,W,AHMCSampleID,Nothing}(undef, 0, npar)
    push!(samples, current_sample)

    nsamples::Int64 = 0

    rngpart_cycle = RNGPartition(rng, 0:(typemax(Int16) - 2))

    metric = AHMCMetric(algorithm.metric, npar)
    logval_posterior(v) = eval_logval(density, v)

    hamiltonian = AdvancedHMC.Hamiltonian(metric, logval_posterior, algorithm.gradient)
    hamiltonian, t = AdvancedHMC.sample_init(rng, hamiltonian, params_vec)

    algorithm.integrator.step_size == 0.0 ? algorithm.integrator.step_size = AdvancedHMC.find_good_stepsize(hamiltonian, params_vec) : nothing
    ahmc_integrator = AHMCIntegrator(algorithm.integrator)

    ahmc_proposal = AHMCProposal(algorithm.proposal, ahmc_integrator)
    ahmc_adaptor = AHMCAdaptor(algorithm.adaptor, metric, ahmc_integrator)

    chain = AHMCIterator(
        algorithm,
        density,
        rng,
        rngpart_cycle,
        info,
        samples,
        nsamples,
        stepno,
        hamiltonian,
        t,
        ahmc_integrator,
        ahmc_proposal,
        ahmc_adaptor
    )

    reset_rng_counters!(chain)

    chain
end

function MCMCIterator(
    rng::AbstractRNG,
    algorithm::HamiltonianMC,
    density::AbstractDensity,
    chainid::Integer,
    startpos::AbstractVector{<:Real}
)
    cycle = 0
    tuned = false
    converged = false
    info = MCMCIteratorInfo(chainid, cycle, tuned, converged)
    AHMCIterator(rng, algorithm, density, info, startpos)
end

@inline _current_sample_idx(chain::AHMCIterator) = firstindex(chain.samples)
@inline _proposed_sample_idx(chain::AHMCIterator) = lastindex(chain.samples)

getalgorithm(chain::AHMCIterator) = chain.algorithm

getdensity(chain::AHMCIterator) = chain.density

getrng(chain::AHMCIterator) = chain.rng

mcmc_info(chain::AHMCIterator) = chain.info

nsteps(chain::AHMCIterator) = chain.stepno

nsamples(chain::AHMCIterator) = chain.nsamples

current_sample(chain::AHMCIterator) = chain.samples[_current_sample_idx(chain)]

sample_type(chain::AHMCIterator) = eltype(chain.samples)

function reset_rng_counters!(chain::AHMCIterator)
    set_rng!(chain.rng, chain.rngpart_cycle, chain.info.cycle)
    rngpart_step = RNGPartition(chain.rng, 0:(typemax(Int32) - 2))
    set_rng!(chain.rng, rngpart_step, chain.stepno)
    nothing
end

function samples_available(chain::AHMCIterator)
    i = _current_sample_idx(chain::AHMCIterator)
    chain.samples.info.sampletype[i] == ACCEPTED_SAMPLE
end

function get_samples!(appendable, chain::AHMCIterator, nonzero_weights::Bool)::typeof(appendable)
    if samples_available(chain)
        samples = chain.samples

        for i in eachindex(samples)
            st = samples.info.sampletype[i]
            if (
                (st == ACCEPTED_SAMPLE || st == REJECTED_SAMPLE) &&
                (samples.weight[i] > 0 || !nonzero_weights)
            )
                push!(appendable, samples[i])
            end
        end
    end
    appendable
end

function next_cycle!(chain::AHMCIterator)
    _cleanup_samples(chain)

    chain.info = MCMCIteratorInfo(chain.info, cycle = chain.info.cycle + 1)
    chain.nsamples = 0
    chain.stepno = 0

    reset_rng_counters!(chain)

    resize!(chain.samples, 1)

    i = _proposed_sample_idx(chain)
    @assert chain.samples.info[i].sampletype == CURRENT_SAMPLE
    chain.samples.weight[i] = 1

    t_stat = chain.transition.stat

    chain.samples.info[i] = AHMCSampleID(
        chain.info.id, chain.info.cycle, chain.stepno, CURRENT_SAMPLE,
        t_stat.hamiltonian_energy, t_stat.tree_depth,
        t_stat.numerical_error, t_stat.step_size
    )

    chain
end

function _cleanup_samples(chain::AHMCIterator)
    samples = chain.samples
    current = _current_sample_idx(chain)
    proposed = _proposed_sample_idx(chain)
    if (current != proposed) && samples.info.sampletype[proposed] == CURRENT_SAMPLE
        # Proposal was accepted in the last step
        @assert samples.info.sampletype[current] == ACCEPTED_SAMPLE
        samples.v[current] .= samples.v[proposed]
        samples.logd[current] = samples.logd[proposed]
        samples.weight[current] = samples.weight[proposed]
        samples.info[current] = samples.info[proposed]

        resize!(samples, 1)
    end
end

function mcmc_step!(chain::AHMCIterator)
    _cleanup_samples(chain)

    samples = chain.samples
    algorithm = getalgorithm(chain)

    chain.stepno += 1
    reset_rng_counters!(chain)

    rng = getrng(chain)
    density = getdensity(chain)

    # Grow samples vector by one:
    resize!(samples, size(samples, 1) + 1)
    samples.info[lastindex(samples)] = AHMCSampleID(
        chain.info.id, chain.info.cycle, chain.stepno, PROPOSED_SAMPLE,
        0.0, 0, false, 0.0
    )

    current = _current_sample_idx(chain)
    proposed = _proposed_sample_idx(chain)
    @assert current != proposed

    current_params = samples.v[current]
    proposed_params = samples.v[proposed]

    # Propose new variate:
    samples.weight[proposed] = 0

    ahmc_step!(rng, algorithm, chain, proposed_params, current_params)
    tstat = chain.transition.stat

    current_log_posterior = samples.logd[current]
    T = typeof(current_log_posterior)

    # Evaluate prior and likelihood with proposed variate:
    proposed_log_posterior = eval_logval(density, proposed_params, strict = false)

    samples.logd[proposed] = proposed_log_posterior

    accepted = current_params != proposed_params

    if accepted
        samples.info.sampletype[current] = ACCEPTED_SAMPLE
        samples.info.sampletype[proposed] = CURRENT_SAMPLE
        chain.nsamples += 1

        samples.info.hamiltonian_energy[proposed] = tstat.hamiltonian_energy
        samples.info.tree_depth[proposed] = tstat.tree_depth
        samples.info.divergent[proposed] = tstat.numerical_error
        samples.info.step_size[proposed] = tstat.step_size
    else
        samples.info.sampletype[proposed] = REJECTED_SAMPLE
    end

    delta_w_current, w_proposed = if accepted
        (0, 1)
    else
        (1, 0)
    end

    samples.weight[current] += delta_w_current
    samples.weight[proposed] = w_proposed

    nothing
end

function ahmc_step!(rng, algorithm, chain, proposed_params, current_params)

    chain.transition = AdvancedHMC.step(rng, chain.hamiltonian, chain.proposal, chain.transition.z)

    tstat = AdvancedHMC.stat(chain.transition)

    if typeof(algorithm.adaptor) <: StanHMCAdaptor
        i = chain.adaptor.state.i
        n_adapts = algorithm.adaptor.n_adapts
    else
        i, n_adapts = chain.info.converged ? (3, 2) : (1, 1)
    end

    chain.hamiltonian, chain.proposal, isadapted = AdvancedHMC.adapt!(
        chain.hamiltonian,
        chain.proposal,
        chain.adaptor,
        Int(i),
        Int(n_adapts),
        chain.transition.z.θ,
        tstat.acceptance_rate
    )

    if i == n_adapts
        chain.info = MCMCIteratorInfo(chain.info, tuned = isadapted)
    end

    tstat = merge(tstat, (is_adapt=isadapted,))

    proposed_params[:] = chain.transition.z.θ
    nothing
end


text:
using ITCC
using Base.Test

# write your own tests here
p = [0.05 0.05 0.05 0.0 0.0 0.0
        0.05 0.05 0.05 0.0 0.0 0.0
        0.0 0.0 0.0 0.05 0.05 0.05
        0.0 0.0 0.0 0.05 0.05 0.05
        0.04 0.04 0.0 0.04 0.04 0.04
        0.04 0.04 0.04 0.0 0.04 0.04]

k = 3
l = 2
n = 5
convergeThresh = 0.001

cX = [3 1 2 2 3 3]
cY = [1 1 2 1 2 2]

result = itcc(p, k, l, n, convergeThresh, cX, cY);

cX_expected = [1 1 2 2 3 3]
cY_expected = [1 1 1 2 2 2]

q_expected = [0.05400000000000001 0.05400000000000001 0.04200000000000001 0.0 0.0 0.0
 0.05400000000000001 0.05400000000000001 0.04200000000000001 0.0 0.0 0.0
 0.0 0.0 0.0 0.04200000000000001 0.05400000000000001 0.05400000000000001
 0.0 0.0 0.0 0.04200000000000001 0.05400000000000001 0.05400000000000001
 0.03600000000000002 0.03600000000000002 0.02800000000000001 0.02800000000000001 0.03600000000000002 0.03600000000000002
 0.03600000000000002 0.03600000000000002 0.02800000000000001 0.02800000000000001 0.03600000000000002 0.03600000000000002]

p_clust_expected = [0.3 0.0
				0.0 0.3
 				0.2 0.2]

kl_expected = 0.066335621500106

converged_expected = true

num_iters_expected = 3

@test result.cX == cX_expected
@test result.cY == cY_expected
@test isapprox(result.q, q_expected)
@test isapprox(result.p_clust, p_clust_expected)
@test isapprox(result.kl, kl_expected)
@test result.converged == converged_expected
@test result.num_iters == num_iters_expected

ext: jl

max_stars_repo_name: slongwell/ITCC

max_issues_repo_name: slongwell/ITCC

max_forks_repo_name: slongwell/ITCC


text:
mpSim, isOkay = initialiseFromDB( "db/simDB", showInfo=false )
run( mpSim )
@test now( mpSim ) == 300.0


text:
# Very basic setup, useful for testing
using DFTK
using Plots

a = 10.26  # Silicon lattice constant in Bohr
lattice = a / 2 * [[0 1 1.];
                   [1 0 1.];
                   [1 1 0.]]
Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
atoms = [Si => [ones(3)/8, -ones(3)/8]]

model = model_LDA(lattice, atoms)
kgrid = [4, 4, 4]  # k-point grid (Regular Monkhorst-Pack grid)
Ecut = 15           # kinetic energy cutoff in Hartree
basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)

scfres = self_consistent_field(basis, tol=1e-8)


text:
#' ---
#' title: Maximum cut problem, linear and semidefinite optimization formulations
#' ---

#' **Originally Contributed by**: Mathieu Besançon

using LinearAlgebra
using SparseArrays
using Test
import Random

using JuMP
using GLPK
using LightGraphs
using GraphPlot
import SCS
using Colors: @colorant_str

#' ## The Maximum Cut problem

# TODO explain problem

#' The Integer Linear Formulation

#' $$
#' \begin{align*}
#' \max_{x,z} & \sum_{\forall (i,j) \in E} w_{i,j} z_{ij} \\
#' \text{s.t. } & z_{ij} \leq x_i + x_j \\
#' & z_ij \leq 2 - (x_i + x_j)\\
#' & z_{ij} \in \{0,1\} \forall (i,j) \in E \\
#' & x_{i} \in \{0,1\} \forall i \in V
#' \end{align*}
#' $$

function compute_cut_value(g::LightGraphs.AbstractGraph, w::AbstractMatrix, vertex_subset)
    return sum(w[i,j] for i in vertices(g) for j in neighbors(g, i) if in(i, vertex_subset)  != in(j, vertex_subset))
end

#' The following example graph is taken from https://en.wikipedia.org/wiki/Graph_theory#/media/File:6n-graf.svg
#'

g = SimpleGraph(6)
add_edge!(g, 1, 2)
add_edge!(g, 1, 5)
add_edge!(g, 2, 3)
add_edge!(g, 2, 5)
add_edge!(g, 3, 4)
add_edge!(g, 4, 5)
add_edge!(g, 4, 6);

GraphPlot.gplot(g, nodelabel=1:6)

#' Creating a unit weight matrix

w = spzeros(6, 6)

for e in edges(g)
   (i, j) = Tuple(e)
   w[i,j] = 1
end

#' Computing value for a solution

@test compute_cut_value(g, w, (2, 5, 6)) ≈ 5
@show compute_cut_value(g, w, (2, 5, 6));

#' Building the linear optimization model

n = nv(g)

linear_max_cut = Model(GLPK.Optimizer)

@variable(linear_max_cut, z[1:n,1:n], Bin)
@variable(linear_max_cut, x[1:n], Bin)
@constraint(linear_max_cut, [i = 1:n, j = 1:n; (i, j) in edges(g)], z[i,j] <= x[i] + x[j])
@constraint(linear_max_cut, [i = 1:n, j = 1:n; (i, j) in edges(g)], z[i,j] <= 2 - (x[i] + x[j]))

@objective(linear_max_cut, Max, dot(w, z))

optimize!(linear_max_cut)

@test objective_value(linear_max_cut) ≈ 6.0
@show objective_value(linear_max_cut);

x_linear = value.(x)
@show x_linear;
@show [(i,j) for i in 1:n-1 for j in i+1:n if JuMP.value.(z)[i,j] > 0.5];

#' Visualizing a solution

nodecolor = [colorant"lightseagreen", colorant"orange"]
all_node_colors = [nodecolor[round(Int, xi + 1)]  for xi in x_linear]
GraphPlot.gplot(g, nodelabel=1:6, nodefillc=all_node_colors)

#' ### Semidefinite formulation

#' The maximum cut problem can also be formulated as a quadratic optimization model.
#' If the solution vector is composed of $\{-1, 1\}$ for belonging to the subset or not.
#' Then, the product $x_i x_j$ is $-1$ if i and j are in different subsets, and 1 otherwise.
#' This implies $1 - x_i x_j$ is $0$ if i and j are in the same subset, and 2 otherwise.
#' Since each pair is counted twice with (i,j) and (j,i), we need to divide it by 4:
#' $\frac{w_{ij}}{4} (1 - x_i x_j)$ is $0$ if i and j are in the same subset, and $z_{ij}$ otherwise.

#' The resulting model is given by:
#' $$
#' \begin{align*}
#' \max_{x} & \frac{w_{ij}}{4} (1 - x_i x_j) \\
#' \text{s.t. } & x_i \in \{-1, 1\} \forall i \in V
#' \end{align*}
#' $$

#' Introducing a matrix variable $Y_{ij} = x_i x_j$,
#' we can equivalently reformulate the above quatratic problem as follows:
#' $$
#' \begin{align*}
#' \max_{x,Y} & \frac{w_{ij}}{4} (1 - Y_{ij}) \\
#' \text{s.t. } & Y_{ii} = 1 \forall i \in V \\
#'              & Y = x x^T
#' \end{align*}
#' $$

#' The constraint $Y = x x^T$ requires $Y$ to be a Positive Semidefinite matrix
#' of rank one. This optimization problem is fully equivalent to the initial
#' quadratic problem and hard to solve.
#' A possible relaxation is to remove the rank-1 constraint, yielding the following problem:
#' $$
#' \begin{align*}
#' \max_{Y} & \frac{w_{ij}}{4} (1 - Y_{ij}) \\
#' \text{s.t. } & Y_{ii} = 1 \forall i \in V \\
#'              & Y \in \mathcal{S}_n^+
#' \end{align*}
#' $$

#' with $\mathcal{S}_n^+$ the set of positive semidefinite matrices.
#' Such model can be implemented in JuMP and solved using an underlying solver
#' which supports semidefinite optimization.

sdp_max_cut = Model(optimizer_with_attributes(SCS.Optimizer, "verbose" => 0))

@variable(sdp_max_cut, Y[1:n,1:n] in PSDCone())
@constraint(sdp_max_cut, [i = 1:n], Y[i,i] == 1)
@objective(sdp_max_cut, Max, 1/4 * sum(w[i,j] * (1 - Y[i,j]) for i in 1:n for j in 1:n))

optimize!(sdp_max_cut)

@test objective_value(linear_max_cut) >= 6.0
@show objective_value(linear_max_cut);

#' Re-computing the vector:
#' $Y$ approximates $x x^T$, we can compute the estimate by randomized rounding.
#' Goemans, M. X., & Williamson, D. P. (1995).
#' Improved approximation algorithms for maximum cut and satisfiability problems using semidefinite programming.
#' Journal of the ACM (JACM), 42(6), 1115-1145.

F = svd(value.(Y))
U = F.U * Diagonal(sqrt.(F.S))

Random.seed!(33)

#' Generating a normal random vector: the vector
#' needs to be uniformly generated on a sphere centered at 0.

x = randn(size(U, 2))
xhat = sign.(U * x) .> 0

@show collect(zip(round.(Int, x_linear), xhat))
@test LinearAlgebra.norm1(round.(Int, x_linear) .- xhat) == 0 || LinearAlgebra.norm1(round.(Int, x_linear) .- xhat) == n

#' ### Representing the SDP approximated solution

nodecolor = [colorant"lightseagreen", colorant"orange"]
all_node_colors = [nodecolor[xi + 1] for xi in xhat]
GraphPlot.gplot(g, nodelabel=1:6, nodefillc=all_node_colors)


text:
module PlotsPyPlot

include("init.jl")
include("backend.jl")
include("output.jl")

__init__() = pyplot()

end # module


text:
using HDF5, Base.Test
include("../util.jl")

# Read in an input matrix and reference eigenvalues/eigenvectors
h5 = h5open(joinpath(dirname(@__FILE__), "..", "reference", "eig.h5"),"r")
matrix = read(h5,"matrix_in")
ref_U = read(h5,"U")
ref_S_diag = read(h5, "S_diag")
close(h5)

# Eigendecomposition with only 1 worker
addprocs(1)
S_diag, U = eig(matrix)

# Test for equality
@test_matrix_approx_eq_eps ref_S_diag S_diag 1e-6 1e-2
@test_matrix_approx_eq_eps ref_U U 1e-6 1e-2
# TODO calibrate above for when use_parallel_workers ≡ true

nothing


text:
"""
```
BruSan{T} <: AbstractNLCTFPModel{T}
```

The `BruSan` type defines the structure of a standard "Brunnermeier-Sannikov" style model
augmented to feature jumps.

### Fields

#### Parameters
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model
  parameters.

* `keys::OrderedDict{Symbol,Int}`: Maps human-readable names for all model
  parameters to their indices in `parameters`.

#### Inputs to Measurement and Equilibrium Functional Equations

The following fields are dictionaries that map human-readable names to indices.

* `state_variables::OrderedDict{Symbol,Int}`: Maps each state variable to its dimension on a Cartesian grid

* `endogenous_variables::OrderedDict{Symbol,Int}`: Maps endogenous variables
    calculated during the solution of functional equations to an index

* `exogenous_shocks::OrderedDict{Symbol,Int}`: Maps each shock to a column in
  the measurement equation

* `observables::OrderedDict{Symbol,Int}`: Maps each observable to a row in the
  model's measurement equation matrices.

* `pseudo_observables::OrderedDict{Symbol,Int}`: Maps each pseudo-observable to
  a row in the model's pseudo-measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"an_schorfheide\", cached
  here for filepath computation.

* `subspec::String`: The model subspecification number, indicating that some
  parameters from the original model spec (\"ss0\") are initialized
  differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect computation
  without changing the economic or mathematical setup of the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If `true`,
  settings from `m.test_settings` are used in place of those in `m.settings`.

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

* `pseudo_observable_mappings::OrderedDict{Symbol,PseudoObservable}`: A
  dictionary that stores names and transformations to/from model units. See
  `PseudoObservable` for further details.
"""
mutable struct BruSan{T} <: AbstractNLCTFPModel{T}
    parameters::ParameterVector{T}                         # vector of all time-invariant model parameters
    keys::OrderedDict{Symbol,Int}                          # human-readable names for all the model
                                                           # parameters and steady-states
    state_variables::OrderedDict{Symbol,Int}               # dimension number of state variable

    functional_variables::OrderedDict{Symbol,Int}
    derivatives::OrderedDict{Symbol,Int}
    endogenous_variables::OrderedDict{Symbol,Int}
    exogenous_shocks::OrderedDict{Symbol,Int}
    observables::OrderedDict{Symbol,Int}
    pseudo_observables::OrderedDict{Symbol,Int}

    spec::String                                           # Model specification number (eg "m990")
    subspec::String                                        # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                         # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}                    # Settings/flags for testing mode
    rng::MersenneTwister                                   # Random number generator
    testing::Bool                                          # Whether we are in testing mode or not

    observable_mappings::OrderedDict{Symbol, Observable}
    pseudo_observable_mappings::OrderedDict{Symbol, PseudoObservable}
end

description(m::BruSan) = "Julia implementation of model defined in 'Public Liquidity and Financial Crises' (2020) by Wenhao Li: BruSan, $(m.subspec)"

"""
`init_model_indices!(m::BruSan)`

Arguments:
`m:: BruSan`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and observables.
"""
function init_model_indices!(m::BruSan)
    # State_Variables
    state_variables = collect([:η])

    # Exogenous shocks
    exogenous_shocks = collect([:K_sh]) # capital shock K

    # Variables for functional equations
    functional_variables = collect([:q, :vₑ, :vₕ])

    # Derivatives of variables
    init_derivatives!(m, OrderedDict{Symbol, Vector{Int}}(:q => standard_derivs(1), :vₑ => standard_derivs(1), :vₕ => standard_derivs(1)),
                      state_variables)
    derivatives = keys(get_derivatives(m))

    # Endogenous variables
    endogenous_variables = collect([:φ_e, :φ_h, :σ_q, :μ_q, :σ_η, :μ_η, :ςₑ, :ςₕ, :ι, :Φ, :μ_K, :dr_f,
                                    :σ_vₑ, :σ_vₕ, :μ_vₑ, :μ_vₕ])

    # Observables
    observables = keys(m.observable_mappings)

    # Pseudo-observables
    pseudo_observables = keys(m.pseudo_observable_mappings)

    for (i,k) in enumerate(state_variables);        m.state_variables[k]      = i end
    for (i,k) in enumerate(exogenous_shocks);       m.exogenous_shocks[k]     = i end
    for (i,k) in enumerate(endogenous_variables);   m.endogenous_variables[k] = i end
    for (i,k) in enumerate(functional_variables);   m.functional_variables[k] = i end
    for (i,k) in enumerate(derivatives);            m.derivatives[k]          = i end
    for (i,k) in enumerate(observables);            m.observables[k]          = i end
    for (i,k) in enumerate(pseudo_observables);     m.pseudo_observables[k]   = i end
end

function BruSan(subspec::String = "ss0";
                custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                testing = false)

    # Model-specific specifications
    spec               = split(basename(@__FILE__),'.')[1]
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # initialize empty model
    m = BruSan{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), OrderedDict{Symbol,Int}(),

            # model indices
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}(),
            OrderedDict{Symbol,PseudoObservable}())

    # Set settings
    model_settings!(m)

    # default_test_settings!(m)
    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Set observable and pseudo-observable transformations
    # init_observable_mappings!(m)
    # init_pseudo_observable_mappings!(m)

    # Initialize parameters
    init_parameters!(m)

    init_model_indices!(m)
    init_subspec!(m)

    return m
end

"""
```
init_parameters!(m::BruSan)
```
Initializes the model's parameters.
"""
function init_parameters!(m::BruSan)
    # Technology
    m <= parameter(:aₑ, 1., (0., Inf), (0., 1e3), ModelConstructors.Untransformed(),
                   Uniform(0., 1e3), fixed = false, description = "Banker's productivity")
    m <= parameter(:aₕ, 0.7, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household's productivity")
    m <= parameter(:δ, 0.1, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.),
                   fixed = false, description = "Depreciation rate of capital")
    m <= parameter(:σ, 0.033, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Volatility of capital shocks")
    m <= parameter(:χ₁, 1., (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Adjustment cost parameter 1 of internal investment function")
    m <= parameter(:χ₂, 1., (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Adjustment cost parameter 2 of internal investment function")

    # Preferences
    m <= parameter(:ρₑ, 0.04, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Expert discount rate")
    m <= parameter(:ρₕ, 0.04, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household discount rate")
    m <= parameter(:νₑ, 0.0, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Expert death rate")
    m <= parameter(:νₕ, 0.0, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household death rate")
    m <= parameter(:γₑ, 1.0, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Expert risk aversion coefficient")
    m <= parameter(:γₕ, 1.0, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household risk aversion coefficient")
    m <= parameter(:ψₑ, 1.0, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Expert inverse Frisch elasticity")
    m <= parameter(:ψₕ, 1.0, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household inverse Frisch elasticity")

    # Stationary distribution
    m <= parameter(:τ, 0.0, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Redistributive wealth tax on experts")
end

"""
```
model_settings!(m::BruSan)
```
creates the model's settings.
"""
function model_settings!(m::BruSan)

    # Numerical settings for grid
    m <= Setting(:N, 150, "Grid size")
    m <= Setting(:stategrid_method, :exponential, "Type of grid to construct for state variables")
    m <= Setting(:stategrid_dimensions, OrderedDict{Symbol, Tuple{Float64, Float64, Int}}(:η => (1e-3, 1. - 1e-3, get_setting(m, :N))),
                 "Information about the dimensions of the state space")
    m <= Setting(:stategrid_splice, 0.1, "The grid is constructed in two parts. This value is where the first half stops.")

    # Numerical settings for no jump equilibrium
    m <= Setting(:boundary_conditions, OrderedDict{Symbol, Vector{Float64}}(:q => [0.; 0.]),
                 "Boundary conditions for differential equations.")
    m <= Setting(:ode_reltol, 1e-4, "Relative tolerance for ODE integration")
    m <= Setting(:ode_abstol, 1e-12, "Absolute tolerance for ODE integration")
    m <= Setting(:ode_integrator, Tsit5(), "Numerical ODE integrator for no-jump solution")
    m <= Setting(:backup_ode_integrators, [DP5(), RK4(), Euler()], "Back up numerical ODE integrators for no-jump solution")
    m <= Setting(:interpolant, Gridded(Linear()), "Interpolation method for value functions")
    m <= Setting(:max_q₀_perturb, 1.05, "Maximum perturbation for the initial value of q at η = 0")

    # Numerical settings for functional iteration
    m <= Setting(:tol, 1e-4, "Tolerance for functional iteration")
    m <= Setting(:learning_rate, 0.4, "Learning rate for the update of functional variables after each loop")
    m <= Setting(:max_iter, 40, "Maximum number of loops during functional iteration")
    m <= Setting(:error_method, :total_error, "Method for calculating error at the end of each loop during functional iteration")
    m <= Setting(:q₀_perturb, 0.006, "Perturbation of boundary condition for q at 1e-3")
    m <= Setting(:κp_grid, Vector{Float64}(undef, 0),
                 "Vector of guesses for κp used during iteration for new κp and xK within each functional loop")
    m <= Setting(:p_interpolant, Gridded(Linear()), "Interpolant method for fitted p.")
    m <= Setting(:xK_interpolant, Gridded(Linear()), "Interpolant method for xK after solving xK for each guess of κp.")
    m <= Setting(:xg_interpolant, Gridded(Linear()), "Interpolant method for xg after solving xK and κp.")
    m <= Setting(:κp_interpolant, Gridded(Linear()), "Interpolant method for κp after solving xK and κp.")
    m <= Setting(:Q̂_interpolant, Gridded(Linear()), "Interpolant method for Q̂.")
    m <= Setting(:inside_iteration_nlsolve_tol, 1e-6, "Tolerance for nlsolve in the inside iteration")
    m <= Setting(:inside_iteration_nlsolve_max_iter, 400, "Maximum number of iterations for nlsolve for the inside iteration")
    m <= Setting(:xg_tol, 2e-3, "Lower but permissible tolerance for xg fixed point in the inside iteration")
    m <= Setting(:yg_tol, 1e-5, "Lowest permissible value for yg")
    m <= Setting(:p_tol, 1e-8, "Tolerance when calculating p after completing an inside iteration")
    m <= Setting(:firesale_bound, 0.99, "Upper bound for size of firesale jumps that do not wipe out a banker's net worth")
    m <= Setting(:firesale_interpolant, Gridded(Linear()), "Interpolant method for the firesale value.")
    m <= Setting(:N_GH, 10, "Number of nodes for Gauss-Hermite quadrature.")
    m <= Setting(:Q̂_tol, 1e-5, "Tolerance for approximation of Q̂ via Gauss-Hermite quadrature")
    m <= Setting(:Q̂_max_it, 1000, "Maximum number of iterations for approximation of Q̂.")
    m <= Setting(:dt, 1. / 12., "Simulation interval for Q̂ approximation as a fraction of a 1 year")
end


text:
# ## Raibert policy
ContactImplicitMPC.@with_kw mutable struct RaibertOptions
    live_plotting::Bool=false # Use the live plotting tool to debug
end

mutable struct Raibert <: ContactImplicitMPC.Policy
    kr_c::Real
    kr_p::Real
    kr_v_stance::Real
    kr_v_flight::Real
    kθ_c::Real
    kθ_p::Real
    kθ_v::Real
    u::AbstractVector
    contact::Bool
    v0::Real
    Tstance::Real
    Tflight::Real
    q0::AbstractVector
    q1::AbstractVector
    opts::RaibertOptions
end

function raibert_policy(model::ContactImplicitMPC.Hopper2D;
    v0=0.5, Tstance=0.13, Tflight=0.62, h = 0.1,
    kr_c =  8e1,
    kr_p = -1e3,
    kr_v_stance = -1e-2,
    kr_v_flight = -1e1,
    kθ_c =  0.0,
    kθ_p = -3e1,
    kθ_v = -1e1,
    mpc_opts = RaibertOptions(),
    )

    u  = zeros(model.nu)
    q0 = zeros(model.nq)
    q1 = zeros(model.nq)
    contact = false

    Raibert(kr_c, kr_p, kr_v_stance, kr_v_flight, kθ_c, kθ_p, kθ_v,
        u, contact, v0, Tstance, Tflight, copy(q0), copy(q1), mpc_opts)
end

function ContactImplicitMPC.policy(p::Raibert, x, traj, t)
    # Initialization
    h = traj.h
    p.q0 = copy(traj.q[t])
    p.q1 = copy(x)

    # Detect contact
    if any(traj.γ[max(1,t-1)] .> 1.5e-2)
        p.contact = true
    else
        p.contact = false
    end

    # Velocities
    qv = (p.q1 - p.q0)/h
    θv = qv[3]
    rv = qv[4]

    # References
    θ1 = p.q1[3]
    r1 = p.q1[4]
    rref = 0.5
    # td = touchdown
    θtd  = asin(p.v0*Tstance/(2*rref))/2

    # Gains
    kr_c = p.kr_c
    kr_p = p.kr_p
    kr_v_stance = p.kr_v_stance
    kr_v_flight = p.kr_v_flight

    kθ_c = p.kθ_c
    kθ_p = p.kθ_p
    kθ_v = p.kθ_v

    if p.contact
        # regulate around θtd using same gain kθ_p but taking into account
        # that stance time is much lower than flight time
        p.u[1] = kθ_c + kθ_p*(θ1 + θtd) * p.Tflight/p.Tstance
        p.u[2] = kr_c + kr_p*(r1 - rref) + kr_v_stance*rv
    else
        # regulate around θtd
        p.u[1] = kθ_p*(θ1 - θtd)  + kθ_v*θv
        p.u[2] = kr_p*(r1 - rref) + kr_v_flight*rv
    end

    return p.u .* h
end


text:
#============================================

Graph the amount of time required to insert a new element into
a Skiplist.

============================================#

pushfirst!(LOAD_PATH, joinpath(@__DIR__, ".."))

using BenchmarkTools
using Logging
using PyPlot
using Skiplists

pygui(true)

list_sizes = (100_000, 120_000, 140_000, 160_000)

# Helper functions

function generate_list(N)
    list = SkipList{Int64}()
    for ii = 1:N
        insert!(list, ii)
    end
    list
end

randint(N) = abs(mod(rand(Int64), N))

# Plotting script

results = []

for N in list_sizes
    push!(results, @benchmark(
        insert!(list, ii),
        setup=(list = generate_list($N); ii = randint($N);),
        seconds=10,
       ))
    @info "Obtained benchmarks for N = $N"
end

@info "Benchmarks completed"

rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["font.size"] = 15

times = [median(r.times) for r in results]
plot(list_sizes, times)
xlabel("List size")
ylabel("Insertion time (μs)")
show()


text:
###############################################################################
#Package of functions related to the Butterfly relaxation
#
#Citation:
#Dussault, J. P., Haddou, M., Kadrani, A., & Migot, T. (2017).
#How to Compute a M-stationary point of the MPCC.
###############################################################################

"""
psi(x) : relaxation function
"""
function psi(x :: Union{Float64,Vector},
             r :: Float64,
             s :: Float64,
             t :: Float64)

 return s .+ t.*theta(x .- s, r)
end

function invpsi(x :: Union{Float64,Vector},
                r :: Float64,
                s :: Float64,
                t :: Float64)

 #inverse not defined for t = 0.
 return t != 0 ? s .+ invtheta((x.-s)./t, r) : -Inf
end

function dpsi(x :: Union{Float64,Vector},
              r :: Float64,
              s :: Float64,
              t :: Float64)

 return t .* dtheta(x .- s, r)
end

function ddpsi(x :: Union{Float64,Vector},
               r :: Float64,
               s :: Float64,
               t :: Float64)

 return t .* ddtheta(x .- s,r)
end

"""
phi(x,r,s,t) : évalue la fonction de relaxation
de la contrainte de complémentarité en (yg,yh) in (ncc,ncc)

output : Array of size 1 x ncc
"""
function phi(yg :: Union{Float64,Vector},
             yh :: Union{Float64,Vector},
             r  :: Float64,
             s  :: Float64,
             t  :: Float64)

 return (yg - psi(yh, r, s, t)) .* (yh - psi(yg, r, s, t))
end
"""
phi(x,r,s,t) : évalue la fonction de relaxation
 de la contrainte de complémentarité en x in n+2ncc

output : Array of size 1 x ncc
"""
function phi(x   :: Vector,
             ncc :: Int64,
             r   :: Float64,
             s   :: Float64,
             t   :: Float64)

 n = length(x)-2*ncc

 yg = x[n+1:n+ncc]
 yh = x[n+ncc+1:n+2*ncc]

 return phi(yg, yh, r, s, t)
end

"""
dphi(x,r,s,t) : évalue la fonction de relaxation de la contrainte de complémentarité en (yg,yh) in (ncc,ncc)

output : Array of size 1 x 2ncc
"""
function dphi(yg :: Union{Float64,Vector},
              yh :: Union{Float64,Vector},
              r  :: Float64,
              s  :: Float64,
              t  :: Float64)

 dg = (yh - psi(yg,r,s,t)) - dpsi(yg,r,s,t).*(yg - psi(yh,r,s,t))
 dh = (yg - psi(yh,r,s,t)) - dpsi(yh,r,s,t).*(yh - psi(yg,r,s,t))

 return [dg;dh]
end

"""
dphi(yg,yh,r,s,t) : évalue la fonction de relaxation de la contrainte de complémentarité en x in n+2ncc

output : Array of size 1 x (n +2ncc)
"""
function dphi(x   :: Vector,
              ncc :: Int64,
              r   :: Float64,
              s   :: Float64,
              t   :: Float64)

 n = length(x)-2*ncc

 yg = x[n+1:n+ncc]
 yh = x[n+ncc+1:n+2*ncc]

 return [zeros(n);dphi(yg,yh,r,s,t)]
end

"""
ddphi(x,r,s,t) : évalue la fonction de relaxation de la contrainte de complémentarité en (yg,yh) in (ncc,ncc)

output : Array of size 2ncc x 2ncc
"""
function ddphi(yg :: Union{Float64,Vector},
               yh :: Union{Float64,Vector},
               r  :: Float64,
               s  :: Float64,
               t  :: Float64)

 dgg = -2*dpsi(yg,r,s,t) - ddpsi(yg,r,s,t).*(yg - psi(yh,r,s,t))
 dhh = -2*dpsi(yh,r,s,t) - ddpsi(yh,r,s,t).*(yh - psi(yg,r,s,t))
 dgh = 1 .+ dpsi(yg,r,s,t).*dpsi(yh,r,s,t)

 return dgg, dgh, dhh
end

"""
alpha_max : see ThetaFct.alpha_theta_max

output : Array of size 1 x 2
"""
function alpha_max(x  :: Union{Float64,Vector},
                   dx :: Union{Float64,Vector},
                   y  :: Union{Float64,Vector},
                   dy :: Union{Float64,Vector},
                   r  :: Float64,
                   s  :: Float64,
                   t  :: Float64)

 return alpha_theta_max(x, dx, y, dy, r, s, t)
end


text:
using Indicators
using PyPlot
using Base.Dates

# Generate some toy sample data
srand(1)
n = 250
op = 100.0 + cumsum(randn(n))
hi = op + rand(n)
lo = op - rand(n)
cl = 100.0 + cumsum(randn(n))
for i = 1:n
	if cl[i] > hi[i]
		cl[i] = hi[i]
	elseif cl[i] < lo[i]
		cl[i] = lo[i]
	end
end
ohlc = [op hi lo cl]
hlc = [hi lo cl]
hl = [hi lo]
t = collect(today():Day(1):today()+Day(n-1))

# Overlays
subplot(411)
plot(t, cl, lw=2, c="k", label="Random Walk")
grid(ls="-", c=[0.8,0.8,0.8])
plot(t, sma(cl,n=40), c=[1,0.5,0], label="SMA (40)")
plot(t, ema(cl,n=10), c=[0,1,1], label="EMA (10)")
plot(t, wma(cl,n=20), c=[1,0,1], label="WMA (20)")
plot(t, psar(hl),   "bo", label="Parabolic SAR")
legend(loc="best", frameon=false)

# MACD
subplot(412)
plot(t, macd(cl)[:,1], label="MACD", c=[1,0.5,1])
plot(t, macd(cl)[:,2], label="Signal", c=[0.5,0.25,0.5])
bar(t, macd(cl)[:,3], align="center", label="Histogram", color=[0,0.5,0.5], alpha=0.25)
plot([t[1],t[end]], [0,0], ls="--", c=[0.5,0.5,0.5])
grid(ls="-", c=[0.8,0.8,0.8])
legend(loc="best", frameon=false)

# RSI
subplot(413)
plot(t, rsi(cl), c=[0.5,0.5,0], label="RSI")
grid(ls="-", c=[0.8,0.8,0.8])
plot([t[1],t[end]], [30,30], c="g")
plot([t[1],t[end]], [70,70], c="r")
legend(loc="best", frameon=false)

# ADX
subplot(414)
plot(t, adx(hlc)[:,1], "g-", label="DI+")
plot(t, adx(hlc)[:,2], "r-", label="DI-")
plot(t, adx(hlc)[:,3], c=[0,0,1], lw=2, label="ADX")
grid(ls="-", c=[0.8,0.8,0.8])
legend(loc="best", frameon=false)

tight_layout()


text:
push!(LOAD_PATH, "../src/")
ENV["GKS_WSTYPE"]=100
using FdeSolver, Documenter

makedocs(
         format=Documenter.HTML(;
         canonical="https://github.com/JuliaTurkuDataScience/FdeSolver.jl"
         ),
         sitename = "FdeSolver.jl",
         modules  = [FdeSolver],
         pages=[
                "Home" => "intro.md"
                "Manual" => "index.md"
                "Examples" => "examples.md"
               ])

deploydocs(;
        repo="github.com/JuliaTurkuDataScience/FdeSolver.jl",
)


text:
"""
```julia
AGeom=Vector{AGeomss,1}
```
An acquisition geometry, bundled into the mutable `AGeom` type,
has to be specified in order to
create any `Expt` variable.
Acquisition has supersources, sources and receivers.
For `AGeom` is a mutable type to define source-receiver geometry.
Each supersource has `ns` sources that are
injected (or active) simultaneously, and
a set of `nr` receivers to
record.
`AGeomss` is a subtype for each supersource, therefore:
This package provides tools to easily create commonly-used `AGeom` instances,
after deciding on a spatial grid `mgrid` for the experiment.

## Indexing and Fields
Let `ageom` be an instance of this type, then fields
can be accessed using:

* `ageom[i]` : acquisition for ith supersource
* `ageom[i].s[p]` : position coordinates of sources where `p ∈ [:z, :x, :y]`
* `ageom[i].r[p]` : position coordinates of receivers
* `ageom[i].ns` : number of sources
* `ageom[i].nr` : number of receivers

```julia
ageom=AGeom(mgrid, attrib, SSrcs(10), Recs(20))
```
Return pre-defined acquisition geometries (in this case with 10 supersources and 20 receivers), based on `attrib`, on an input mesh `mgrid`.
The sources and receivers positions are same for all supersources, and are not placed on the edges of the mesh.
Choose `attrib::Symbol`
* `=:xwell` cross-well acquisition;
* `=:surf` surface acquisition;
* `=:vsp` vertical seismic profiling;
* `=:rvsp`  reverse vertical seismic profiling;
* `=:downhole` downhole sources and receivers.

TODO: For 3D, just mean(ygrid) is used, need to work on more exotic gallery of 3D acquisitions.
"""
AGeom = Array{AGeomss,1}

function Array{AGeomss,1}(
    mgrid::AbstractArray{T},
    attrib::Symbol,
    ss::SSrcs = SSrcs(1),
    r::Recs = Recs(10),
) where {T<:StepRangeLen}
    N = length(mgrid)
    @assert N ∈ [2, 3]
    # ouput a coordinate given a 1-D grid usign weights a and b
    getp(m, a = 0.5) = a * m[1] + (1 - a) * m[end]
    function getp(m::AbstractArray{T}, a = fill(0.5, length(m))) where {T<:StepRangeLen}
        @assert length(m) == length(a)
        @assert all(a .> 0)
        @assert all(a .< 1)
        return [getp(mgrid[i], a[i]) for i = 1:length(m)]
    end
    ageom = AGeom(mgrid, ss, Srcs(1), r) # just use one source per supersource
    if (attrib == :xwell)
        sp1 = [0.9, 0.9]
        sp2 = [0.1, 0.9]
        rp1 = [0.9, 0.1]
        rp2 = [0.1, 0.1]
    elseif (attrib == :surf)
        sp1 = [0.95, 0.9]
        sp2 = [0.95, 0.1]
        rp1 = [0.95, 0.9]
        rp2 = [0.95, 0.1]
    elseif (attrib == :vsp)
        sp1 = [0.95, 0.9]
        sp2 = [0.95, 0.1]
        rp1 = [0.75, 0.9]
        rp2 = [0.1, 0.9]
    elseif (attrib == :rvsp)
        sp1 = [0.75, 0.9]
        sp2 = [0.1, 0.9]
        rp1 = [0.95, 0.9]
        rp2 = [0.95, 0.1]
    elseif (attrib == :downhole)
        sp1 = [0.75, 0.75]
        sp2 = [0.9, 0.75]
        rp1 = [0.75, 0.75]
        rp2 = [0.9, 0.75]
    elseif (attrib == :microseismic)
        sp1 = [rand(Uniform(0.1,0.2)), rand(Uniform(0.7,0.9))]
        sp2 = [rand(Uniform(0.1,0.2)), rand(Uniform(0.7,0.9))]
        rp1 = [0.95, 0.9]
        rp2 = [0.95, 0.1]
    else
        error("invalid attrib")
    end
    if (N == 2)
        update!(ageom, SSrcs(), getp(mgrid, sp1), getp(mgrid, sp2))
        update!(ageom, Recs(), getp(mgrid, rp1), getp(mgrid, rp2))
    else # just use mean y for 3D
        update!(
            ageom,
            SSrcs(),
            getp(mgrid, [sp1[1], 0.5, sp1[2]]),
            getp(mgrid, [sp2[1], 0.5, sp2[2]]),
        )
        update!(
            ageom,
            Recs(),
            getp(mgrid, [rp1[1], 0.5, rp1[2]]),
            getp(mgrid, [rp2[1], 0.5, rp2[2]]),
        )
    end
    return ageom
end


text:
using Test
using JSON3
using Pinecone

APIKEY = ENV["PINECONE_API_KEY"]
CLOUDENV="us-west1-gcp"

context = Pinecone.init(APIKEY, CLOUDENV)

INDEX = "filter-example"
NAMESPACE = "mynamespace"
index = Pinecone.Index(INDEX);

@testset "Test describe_index_stats()" begin
    context = Pinecone.init(APIKEY, CLOUDENV)
    index = PineconeIndex(INDEX)
    result = Pinecone.describe_index_stats(context, index)
    @test result !== nothing
    @test typeof(result) == String
    parsed = JSON3.read(result)
    @test haskey(parsed, "namespaces") == true
    @test haskey(parsed, "dimension") == true
 end


text:
function _posterior(
    latent::Latent1D,
    parameters::Parameters2PL,
    response::Response,
    w::PolyaGammaSample,
)
    sigma2 = (parameters.a^2) * w.val
    sigma2 = 1 / (sigma2 + (1 / Distributions.var(latent.prior)))
    mu =
        sigma2 * (
            parameters.a * (parameters.a * parameters.b * w.val + (response.val - 0.5)) +
            (latent.prior.μ / Distributions.var(latent.prior))
        )
    return Distributions.Normal(mu, sqrt(sigma2))
end

function posterior(
    examinee::AbstractExaminee,
    item::AbstractItem,
    response::Response,
    w::PolyaGammaSample
)
    return _posterior(examinee.latent, item.parameters, response, w)
end

function __posterior(
    latent::Latent1D,
    parameters::Vector{Parameters2PL},
    responses_val::Vector{Float64},
    W_val::Vector{Float64},
)
    prior = latent.prior
    sigma2 = mapreduce((p, w_val) -> (p.a^2) * w_val, +, parameters, W_val)
    sigma2 = 1 / (sigma2 + (1 / Distributions.var(prior)))
    mu =
        sigma2 * (
            mapreduce(
                (p, w_val, r_val) -> p.a * (p.a * p.b * w_val +
                                    #(get_responses_by_item_id(i.id, responses_e)[1].val - 0.5)
                                    (r_val - 0.5)),
                +,
                parameters,
                W_val,
                responses_val,
            ) + (prior.μ / Distributions.var(prior))
        )
    return Distributions.Normal(mu, sqrt(sigma2))::Distributions.ContinuousDistribution
end

function _posterior(
    latent::Latent1D,
    parameters::Vector{Parameters2PL},
    responses::Vector{Response},
    W::Vector{PolyaGammaSample},
)
    prior = latent.prior
    sigma2 = mapreduce((i, w) -> (i.a^2) * w.val, +, parameters, W)
    sigma2 = 1 / (sigma2 + (1 / Distributions.var(prior)))
    mu =
        sigma2 * (
            mapreduce(
                (i, w, r) -> i.a * (i.a * i.b * w.val +
                                    #(get_responses_by_item_id(i.id, responses_e)[1].val - 0.5)
                                    (r.val - 0.5)),
                +,
                parameters,
                W,
                responses,
            ) + (prior.μ / Distributions.var(prior))
        )
    return Distributions.Normal(mu, sqrt(sigma2))::Distributions.ContinuousDistribution
end

function posterior(
    examinee::AbstractExaminee,
    items::Vector{<:AbstractItem}, #only items answered by examinee sorted by i.idx
    responses::Vector{<:AbstractResponse}, #only responses of examinee sorted by i.idx
    W::Vector{PolyaGammaSample};
    already_sorted = false,
    )
    if !already_sorted
        sort!(items, by = i -> i.idx)
        sort!(responses, by = r -> r.item_idx)
        already_sorted = true
    end
    return _posterior(examinee.latent, map(i -> i.parameters, items), responses, W)
end

function update_posterior!(
    examinee::AbstractExaminee,
    items::Vector{<:AbstractItem},
    responses::Vector{<:AbstractResponse},
    W::Vector{PolyaGammaSample};
    already_sorted = false,
    )
    if !already_sorted
        sort!(items, by = i -> i.idx)
        sort!(responses, by = r -> r.item_idx)
        already_sorted = true
    end
    examinee.latent.posterior = posterior(examinee, items, responses, W; already_sorted = already_sorted)
end


text:
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Computing D_KL
# The function compute D_KL between 2 Dirichlet dists
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_D_KL(α::Array{Float64,1}, β::Array{Float64,1})
    α_0 = sum(α)
    β_0 = sum(β)
    D1 = + lgamma(α_0) - sum(lgamma.(α))
    D2 = - lgamma(β_0) + sum(lgamma.(β))
    D3 = sum( ( α .- β ) .* ( digamma.(α) .- digamma(α_0) ) )
    D = D1 + D2 + D3
    return D
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Computing surprise
# The function compute surprise of observing x at time t
# {s(t-1) -> x(t)} -> S(t)
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_surprise(Agent, t::Int)
    x = Agent.Input.obs[t]
    s = Agent.State[t-1]
    α = Agent.Π[:, s+1, t-1]
    ϵ = Agent.Param.ϵ

    if Agent.HyperParam.surprise_form == "BF"
        p_x = α[x+1]/(sum(α))
        S = 1/p_x
        return S

    elseif Agent.HyperParam.surprise_form == "Sh"
        p_x = α[x+1]/(sum(α))
        S = - log(p_x)
        return S

    elseif Agent.HyperParam.surprise_form == "CC"
        β = ones(size(α)) .* ϵ
        β[x+1] = ϵ + 1
        S = Func_compute_D_KL(α, β)
        return S
    end

    error("Surprise format is not known!")

end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Computing gamma
# The function compute gamma of corresponding surprise at time t
# S(t) -> gamma(t)
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_gamma(Agent, t::Int)
    p_c = Agent.Param.p_change

    if Agent.HyperParam.gamma_form == "fix"
        gamma=p_c
        return gamma

    elseif Agent.HyperParam.gamma_form == "satur"
        S = Agent.Surprise[t]
        m = p_c / (1 - p_c)
        gamma = m * S / (1 + m*S)
        return gamma
    end

    error("Gamma format is not known!")

end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Computing t_since_change
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_t_since_change(Agent, t::Int)
    gamma = Agent.Gamma[t]
    t_since_change = Agent.t_since_change[t-1]
    return ( (1-gamma)*t_since_change + 1 )
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update Counter
# The function update counter after observing x at time t
# x(t) -> Counter(:,t)
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_update_counter(Agent, t::Int)
    x = Agent.Input.obs[t]
    counter = Agent.Counter[:,t-1]
    counter = Agent.Param.count_decay .* counter
    counter[x+1] += 1
    return counter
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update belief
# Having x(t) and gamma(t), the belief is updated
# x(t) and gamma(t)  -> Π(:,:,t)
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_update_belief(Agent, t::Int)
    x = Agent.Input.obs[t]
    s = Agent.State[t-1]

    Π = Agent.Π[:, :, t-1]
    α = Π[:, s+1]

    γ = Agent.Gamma[t]
    ϵ = Agent.Param.ϵ

    if Agent.HyperParam.update_form == "VarSMiLe"
        α = ( (1 - γ) .* α ) .+ ( γ * ϵ )
        α[x+1] += 1
        Π[:, s+1] = α
        return Π

    elseif Agent.HyperParam.update_form == "Leaky"
        γ = Agent.Param.p_change
        α = ( (1 - γ) .* α )
        α[x+1] += 1
        Π[:, s+1] = α
        return Π
    end

    error("Update format is not known!")
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update Novelty
# After updating belief and counter, novely vector is computed
# Counter(:,t) or Π(:,:,t) -> Novelty(:,t)
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_update_novelty(Agent, t::Int)
    if Agent.HyperParam.novelty_form == "log-count"
        P_obs = (Agent.Counter[:,t] .+ 1) ./ (sum( Agent.Counter[:,t] .+ 1) )
        Novelty = - log.(P_obs)
        return Novelty

    elseif Agent.HyperParam.novelty_form == "trap-det"
        Novelty = -1 .* (Agent.Counter[:,t] .> Agent.Param.C_thresh)
        return Novelty

    elseif Agent.HyperParam.novelty_form == "trap-det-2"
        Count = Agent.Counter[:,t]
        Novelty = zeros(size(Count))
        ind_Count = sortperm(Count[:],rev=true)[1:Int(Agent.Param.C_thresh)]
        Novelty[ind_Count[Count[ind_Count].>0]] .= -1
        return Novelty
    end

    error("Novelty format is not known!")
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Compute the estimation of the transition probability
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_P_trans(Agent, t::Int)
    Π = Agent.Π[:,:,t]
    P_Tran = zeros(size(Π))

    for j = 1:(size(Π)[2])
        α0 = sum(Π[:,j])
        P_Tran[:,j] = Π[:,j]./α0;
    end

    return P_Tran
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update "State", defined as the combination of real state and action
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_update_State(Agent, t::Int)
    State = Agent.Input.obs[t] * Agent.Param.num_act + Agent.Input.act[t]
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update Q and U Values for model based
# After updating belief and novelt, new Q and U values are computed either for
# novelty or reward
# Π(:,:,t) and Novelty(:,t) -> Q(:,t) and U(:,t)
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_update_prior_sweep(Agent, t::Int, RorN::String)
    if RorN == "Reward"
        U = Agent.U_R[:,t-1]
        Q = zeros(size(Agent.Q_R_MB[:,t-1]))
        Reward = Agent.HyperParam.reward_func[:]
        Reward = repeat(Reward,1,length(Q))
        if Agent.Param.γ_RL_R_MB == -1
            γ = Agent.Param.γ_RL_R
        else
            γ = Agent.Param.γ_RL_R_MB
        end

    elseif RorN == "Novelty"
        U = Agent.U_N[:,t-1]
        Q = zeros(size(Agent.Q_N_MB[:,t-1]))
        Reward = Agent.Novelty[:,t]
        Reward = repeat(Reward,1,length(Q))
        if Agent.Param.γ_RL_N_MB == -1
            γ = Agent.Param.γ_RL_N
        else
            γ = Agent.Param.γ_RL_N_MB
        end

    elseif RorN == "MB_Curiosity"
        U = Agent.U_N[:,t-1]
        Q = zeros(size(Agent.Q_N_MB[:,t-1]))
        Reward = - log.(Func_compute_P_trans(Agent, t))
        if Agent.Param.γ_RL_N_MB == -1
            γ = Agent.Param.γ_RL_N
        else
            γ = Agent.Param.γ_RL_N_MB
        end

    else
        error("RorN format is not known!")
    end

    P_Tran = Func_compute_P_trans(Agent, t)
    V = zeros(size(U))
    p = zeros(size(U))

    A = Agent.Param.num_act

    # Calculating Q
    for i = 1:length(Q)
        Q[i] = sum( (Reward[:,i] .+ (γ .* U) ) .* P_Tran[:,i] )
    end

    # Calculating V
    for i = 1:length(U)
        s_i = (i-1)*A
        V[i] = findmax( Q[ s_i .+ (1:A) ] )[1]
        p[i] = abs( V[i] - U[i] )
    end

    for i_cycle = 1:Agent.Param.N_prior_sweep
        X_i = findmax(p)[2]
        dV = V[X_i] - U[X_i]
        U[X_i] = V[X_i]

        for i = 1:length(Q)
            Q[i] = Q[i] + ( γ * dV * P_Tran[X_i, i] )
        end

        for i = 1:length(U)
            s_i = (i-1)*A
            V[i] = findmax( Q[ s_i .+ (1:A) ] )[1]
            p[i] = abs( V[i] - U[i] )
        end
    end

    return Q, U
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Compute δ, either RPE or NPE
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_δ(Agent, t, RorN)
    if RorN == "Reward"
        Q = Agent.Q_R_MF[:,t-1]
        Reward = Agent.HyperParam.reward_func[:]
        Reward = repeat(Reward,1,length(Q))
        γ = Agent.Param.γ_RL_R

    elseif RorN == "Novelty"
        Q = Agent.Q_N_MF[:,t-1]
        Reward = Agent.Novelty[:,t]                     # WARNING
        Reward = repeat(Reward,1,length(Q))
        γ = Agent.Param.γ_RL_N

    elseif RorN == "MB_Curiosity"
        Q = Agent.Q_N_MF[:,t-1]
        Reward = - log.(Func_compute_P_trans(Agent, t))
        γ = Agent.Param.γ_RL_N

    else
        error("RorN format is not known!")
    end

    state_old = Agent.State[t-1]                        # (s_{t-1},a_{t-1})
    obs_new = Agent.Input.obs[t]                        #  s_{t}

    A = Agent.Param.num_act
    Q_old = Q[ state_old + 1 ]                          # Q(s_{t-1},a_{t-1})
    V_new = findmax( Q[ (obs_new*A) .+ (1:A) ] )[1]     # max_{a'} Q(s_{t},a')
    R_new = Reward[obs_new + 1, state_old + 1]          # R_{t}

    δ = R_new + γ*V_new - Q_old

    return δ
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update Eligibility Trace E
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_ElTrace(Agent, t, RorN)
    if RorN == "Reward"
        E = Agent.E_R[:,t-1]
        λ = Agent.Param.λ_elig_R
        γ = Agent.Param.γ_RL_R

    elseif (RorN == "Novelty") || (RorN == "MB_Curiosity")
        E = Agent.E_N[:,t-1]
        λ = Agent.Param.λ_elig_N
        γ = Agent.Param.γ_RL_N

    else
        error("RorN format is not known!")
    end

    state_old = Agent.State[t-1]                        # (s_{t-1},a_{t-1})
    E = (λ*γ) * E
    E[state_old+1] = 1

    return E
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Update Q values for model-free
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_update_Q_MF(Agent, t, RorN)
    if RorN == "Reward"
        Q = Agent.Q_R_MF[:,t-1]
        δ = Agent.δ_R_Seq[t]
        E = Agent.E_R[:,t]

    elseif (RorN == "Novelty") || (RorN == "MB_Curiosity")
        Q = Agent.Q_N_MF[:,t-1]
        δ = Agent.δ_N_Seq[t]
        E = Agent.E_N[:,t]

    else
        error("RorN format is not known!")
    end

    γ_S = Agent.Gamma[t]
    if Agent.HyperParam.MF_surprise_modul == "Adapt"
        α_0 = Agent.Param.α
        Δα = Agent.Param.Δα
        α = α_0 + Δα*γ_S
    elseif Agent.HyperParam.MF_surprise_modul == "DecayingAdapt"
        α_0 = Agent.Param.α
        Δα = Agent.Param.Δα
        t_since_change = Agent.t_since_change[t]
        α = α_0/(1 + Δα*t_since_change)
    else
        error("MF_surprise_modul format is not known!")
    end

    Q = Q .+ ((α*δ) * E)

    return Q
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Computing Action Probability
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
function Func_compute_P_action(Agent, t::Int; Only_Reward=false)
    if Only_Reward
        Q_MB = Agent.Q_R_MB[:,t] .+ ( Agent.Param.β_ratio_N2R_exploit .* Agent.Q_N_MB[:,t] )
        Q_MF = Agent.Q_R_MF[:,t] .+ ( Agent.Param.β_ratio_N2R_exploit .* Agent.Q_N_MF[:,t] )
    else
        Q_MB = Agent.Q_R_MB[:,t] .+ ( Agent.Param.β_ratio_N2R .* Agent.Q_N_MB[:,t] )
        Q_MF = Agent.Q_R_MF[:,t] .+ ( Agent.Param.β_ratio_N2R .* Agent.Q_N_MF[:,t] )
    end

    if Agent.First_Epi
        Q = (Agent.Param.w_MF_to_MB_1 * Agent.Param.w_MF_to_MB_Scaler) .* Q_MF .+
            (1-Agent.Param.w_MF_to_MB_1) .* Q_MB
    else
        Q = (Agent.Param.w_MF_to_MB_2 * Agent.Param.w_MF_to_MB_Scaler) .* Q_MF .+
            (1-Agent.Param.w_MF_to_MB_2) .* Q_MB
    end

    Q = Q[ (Agent.Input.obs[t]*Agent.Param.num_act) .+ (1:Agent.Param.num_act) ]
    Q = (Q .- findmax(Q)[1]) .* Agent.Param.β_R
    P_Action = exp.(Q)
    P_Action = P_Action ./ sum(P_Action)

    return P_Action
end


text:
import CUDA

# Repeat cudnn calls that fail due to memory issues
macro cudnn_retry(ex)
    quote
        res = CUDA.CUDNN.@retry_reclaim x->(x ∈ (CUDA.CUDNN.CUDNN_STATUS_ALLOC_FAILED, CUDA.CUDNN.CUDNN_STATUS_EXECUTION_FAILED)) begin
            $ex
        end
        if res != CUDA.CUDNN.CUDNN_STATUS_SUCCESS
            CUDA.CUDNN.throw_api_error(res)
        end
    end |> esc
end


text:
abstract type AbstractIterState{T<:Real} end
####Fixed point methods
#contains various state variables for running the solver
mutable struct IterFPState{TΘ⁺<:VolumePartsIter,T} <: AbstractIterState{T}
  Θ⁺::TΘ⁺
  λ⁺::T
  ftol::T

  griddelta::T
  gridstart::T
  gridend::T
  remainingpregriddelta::Vector{T}
  currentgriddelta::T

  starttime::Float64
  besttime::Float64
  maxelapsedsincebest::Float64
  maxelapsed::Float64

  iter::Int
  bestiter::Int
  maxitersincebest::Int
  maxiter::Int

  lastreporttime::Float64
  stop::Union{Nothing,Symbol}
end

#updates the problem state given the solution for an iteration
function updatefpstate!(Ψ::IterFPState{TΘ⁺,T}, Θ::VolumePartsIter, λ::T) where {TΘ⁺,T}
  Ψ.iter+=1
  t=time()
  if abs(Ψ.λ⁺ - λ) > Ψ.ftol
    Ψ.besttime=t
    Ψ.bestiter=Ψ.iter
  end

  if λ < Ψ.λ⁺ #keep track of the best solution, even if it doesn't meet ftol criteria
    Ψ.Θ⁺.Π .= Θ.Π
    Ψ.λ⁺ = λ
  end

  #adjusts the grid delta if a pre-conditioning glide path was provided
  if length(Ψ.remainingpregriddelta) > 0
    Ψ.currentgriddelta = popfirst!(Ψ.remainingpregriddelta)
    if length(Ψ.remainingpregriddelta) == 0
      Ψ.currentgriddelta = Ψ.griddelta #the last point should be the current grid
      @info "Grid delta pre-conditioning glide path complete"
    end
    Ψ.besttime=t
    Ψ.bestiter=Ψ.iter
    return nothing #this short-circuits the stopping criteria
  end

  #run through the stopping criteria
  elapsed::Int = round(time() - Ψ.starttime)
  elapsedsincebest::Int = round(time() - Ψ.besttime)
  itersincebest::Int = Ψ.iter - Ψ.bestiter
  if  elapsed > Ψ.maxelapsed
    @info "STOP: max time $elapsed/$(Ψ.maxelapsed) reached"
    Ψ.stop=:maxtime
  elseif elapsedsincebest>Ψ.maxelapsedsincebest
    @info "STOP: max time with minimal improvement $elapsedsincebest/$(Ψ.maxelapsedsincebest) reached"
    Ψ.stop=:maxtimesincebest
  elseif Ψ.iter ≥ Ψ.maxiter
    @info "STOP: max iter $(Ψ.iter) reached."
    Ψ.stop=:maxiter
  elseif itersincebest  ≥ Ψ.maxitersincebest
    @info "STOP: max iter since best $itersincebest reached."
    Ψ.stop=:maxitersincebest
  end

  return nothing
end

#contains some defaults for creating teh initial solver state
IterFPState(Θ::VolumePartsIter, ::Type{T};
  Θ⁺::VolumePartsIter = deepcopy(Θ),
  λ⁺::T=T(Inf),
  ftol::T=max(T(PARAM[:fpftol]), zero(T)),

  griddelta::T = PARAM[:fpgriddelta] |> T,
  gridstart::T = PARAM[:fpgridstart] |> T,
  gridend::T = PARAM[:fpgridend] |> T,
  remainingpregriddelta::Vector{T}= T.(PARAM[:fppregriddelta]),
  currentgriddelta::T= (
    length(remainingpregriddelta) > 0 ? popfirst!(remainingpregriddelta) : griddelta) |> T,

  starttime::Float64=time(),
  besttime::Float64=time(),
  maxelapsedsincebest::Float64=PARAM[:fpmaxelapsedsincebest],
  maxelapsed::Float64=PARAM[:fpmaxelapsed],

  iter::Int=0,
  bestiter::Int=0,
  maxitersincebest::Int=PARAM[:fpmaxitersincebest],
  maxiter::Int=PARAM[:fpmaxiter],

  lastreporttime::Float64=time(),
  stop::Union{Nothing,Symbol} = nothing,
  ) where T = IterFPState(Θ⁺, λ⁺, ftol,
    griddelta, gridstart, gridend, remainingpregriddelta, currentgriddelta,
    starttime, besttime, maxelapsedsincebest, maxelapsed,
    iter, bestiter, maxitersincebest, maxiter,
    lastreporttime, stop)

#reports the current status of the solver
function status(Ψ::IterFPState, λ; minreporttime::Float64 = PARAM[:iterminreporttime])
  (time()<minreporttime+Ψ.lastreporttime) && return nothing
  Ψ.lastreporttime=time()

  elapsed::Int = round(time() - Ψ.starttime)
  elapsedsincebest::Int = round(time() - Ψ.besttime)

  itersincebest::Int = Ψ.iter - Ψ.bestiter

  println("iter=$(Ψ.iter)/$(Ψ.maxiter), itersincebest=$itersincebest/$(Ψ.maxitersincebest)," *
    "loss=$λ, bestloss=$(Ψ.λ⁺), " *
    "time=$elapsed/$(Ψ.maxelapsed), " *
    "timesincebest = $elapsedsincebest/$(round(Ψ.maxelapsedsincebest)), " *
    "griddelta=$(Ψ.currentgriddelta)")

  return nothing
end

function solveA!(iterloss, Θ::VolumePartsIter,
    Xv::AbstractXYIter{TM,TV,T}, _Xv::XYIterAlloc{TM,TV,T}, RegressionType::Val, ::Val{:fixedpoint}
    ) where {TM,TV,T}

  Ψ::IterFPState = IterFPState(Θ, T)

  try
    while true #main solving loop

      #update the growths and values
      λ::T = updateA!(iterloss, Θ,Xv,_Xv,RegressionType, ΔG=Ψ.currentgriddelta)

      #update the problem state
      updatefpstate!(Ψ,Θ,λ)

      #check the stopping criteria and report status
      status(Ψ, λ)
      (!(Ψ.stop === nothing)) && break

    end
  catch err
    @warn "Aborted solver with error $err"
    Ψ.stop = :aborted
  finally #do this as a finallizer in case there is an error

    #ensure that we capture the best λ
    Θ.Π .= Ψ.Θ⁺.Π

    b = IOBuffer()
    elapsed::Int = round(time() - Ψ.starttime)
    elapsedsincebest::Int = round(time() - Ψ.besttime)
    itersincebest::Int = Ψ.iter - Ψ.bestiter
    write(b, "***Results info:\n" *
      "stop_reason: $(Ψ.stop)\n" *
      "iterations: $(Ψ.iter)/$(Ψ.maxiter)\n" *
      #"start_time: $(Ψ.starttime)\n" *
      "elapsed_at_completion: $(elapsed)/$(Ψ.maxelapsed)\n" *
      "elapsed since best: ~$(elapsedsincebest)/$(Ψ.maxelapsedsincebest)\n" *
      "iter since best: ~$(itersincebest)/$(Ψ.maxitersincebest)\n" *
      "Best resid: $(Ψ.λ⁺)\n")

    write(b, "\n\n***Other Simulation parameters:\n" *
      "regression method: $(RegressionType)\n" *
      "report interval: $(PARAM[:iterminreporttime])\n" *
      "Types (Matrix, Vector, T): $TM, $TV, $T\n" *
      "grid (excludes preconditioning) start:delta:end: $(Ψ.gridstart):$(Ψ.griddelta):$(Ψ.gridend)\n" *
      "currentgriddelta: $(Ψ.currentgriddelta)\n")

    if length(PARAM[:iterpregriddelta]) > 0
      write(b,
      "precondition max: $(maximum(PARAM[:iterpregriddelta]))\n" *
      "precondition length: $(length(PARAM[:iterpregriddelta]))\n"
      )
    end
    #compute the full jacobian with AD

    #∇Θver = gradient(()->loss(Θver,X,y),Flux.params(Θver))
    #rmsver = sqrt(mean(vec∇(Θver, ∇Θver).^2))

    #write(b, "\n\nlst squares gradient rms: $rmsver\n")

    iterinfo = String(take!(b))
    println(iterinfo)
    open("$(PARAM[:outputpath])\\$(iterresultsname(Ψ.λ⁺))_info.txt", "w+") do f
      write(f, iterinfo)
    end

    return Ψ.λ⁺
  end

  @assert false #should never hit this point due to the finalizer
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
export with_terminal, Dump, Show, Print

import Suppressor: @color_output, @capture_out, @capture_err
import Logging:  ConsoleLogger, with_logger
import Markdown: htmlesc
import Base: show

struct WithTerminalOutput
    stdout::String
    stderr::String
    value::Any
end

function show(io::IO, mime::MIME"text/html", with_terminal::WithTerminalOutput)
	show(io, mime, HTML("""
		$(terminal_css)

$(htmlesc(with_terminal.stdout))

            $(isempty(with_terminal.stderr) ? "" : "
" * htmlesc(with_terminal.stderr) * "
")

    """))
end

const terminal_css = """

"""

"""
Run the function, and capture all messages to `stdout`. The result will be a small terminal displaying the captured text.

This allows you to to see the messages from `println`, `dump`, `Pkg.status`, etc.

Example:

```julia
with_terminal() do
    x=1+1
    println(x)
    @warn "Oopsie!"
end
```

```julia
with_terminal(dump, [1,2,[3,4]])
```

See also [PlutoUI.Dump](@ref).

"""
function with_terminal(f::Function, args...; kwargs...)
    local spam_out, spam_err, value
	@color_output false begin
		spam_out = @capture_out begin
            spam_err = @capture_err begin
	            with_logger(ConsoleLogger(stdout)) do
                    value = f(args...; kwargs...)
                end
			end
		end
  end
	WithTerminalOutput(spam_out, spam_err, value)
end

"""
    Dump(x; maxdepth=8)

Every part of the representation of a value. The depth of the output is truncated at maxdepth.

This is a variant of [`Base.dump`](@ref) that returns the representation directly, instead of printing it to stdout.

See also: [`Print`](@ref) and [`with_terminal`](@ref).
"""
function Dump(x; maxdepth=8)
	sprint() do io
		dump(io, x; maxdepth=maxdepth)
	end |> Text
end

"""
    Show(mime::MIME, data)

An object that can be rendered as the `mime` MIME type, by writing `data` to the IO stream. For use in environments with rich output support. Read more about [`Base.show`](@ref).

# Examples

```julia
Show(MIME"text/html"(), "I can be <b>rendered</b> as <em>HTML</em>!")
```

```julia
Show(MIME"image/png"(), read("dog.png"))
```

`Base.showable` and `Base.show` are defined for a `Show`.

```julia
s = Show(MIME"text/latex"(), "\\\\frac{hello}{world}")

showable(MIME"text/latex"(), s) == true

repr(MIME"text/latex"(), s) == "\\\\frac{hello}{world}"
```

"""
struct Show
    mime::MIME
    data
end

Base.showable(m::MIME, x::Show) = x.mime == m
Base.show(io::IO, ::MIME, x::Show) = write(io, x.data)

"""
    Print(xs...)

The text that would be printed when calling `print(xs...)`. Use `string(xs...)` if you want to use the result as a `String`.

See also: [`Dump`](@ref) and [`with_terminal`](@ref).
"""
Print(xs...) = Text(string(xs...))


text:
# Currently unused.
const grid_definition = 512

"""
$(SIGNATURES)

This func returns a square mesh, centered on zero, with (width * height) nodes
"""
function squareMesh(width::Int, height::Int)
    nodeList = Vector{Point3D}(undef, height * width)
    nodeArray = Matrix{Point3D}(undef, width, height)
    count = 1
    midpoint = width / 2
    for y in 1:height, x in 1:width
        newPoint = Point3D(x, y, 0, x, y)
        nodeList[count] = newPoint
        nodeArray[x, y] = newPoint
        count += 1
    end

    triangles = Vector{Triangle}(undef, (width - 1) * (height - 1) * 2)
    count = 1
    for y = 1:(height - 1)
        for x = 1:(width - 1)
            # here x and y establish the column of squares we're in
            index_ul = (y - 1) * width + x
            index_ur = index_ul + 1

            index_ll = y * width + x
            index_lr = index_ll + 1

            triangles[count] = Triangle(index_ul, index_ll, index_ur)
            count += 1
            triangles[count] = Triangle(index_lr, index_ur, index_ll)
            count += 1
        end
    end

    return Mesh(nodeList, nodeArray, triangles, width, height)
end

"""
$(SIGNATURES)

Warning: not guaranteed to work in 3D?
"""
function centroid(mesh::Mesh, index::Int)
    triangle = mesh.triangles[index]
    p1 = mesh.nodes[triangle.pt1]
    p2 = mesh.nodes[triangle.pt2]
    p3 = mesh.nodes[triangle.pt3]

    return centroid(p1, p2, p3)
end

"""
$(SIGNATURES)

Given 3 points and 3 velocities, calculate the `t` required to bring the area of that triangle to zero
"""
function findT(
    p1::Point3D,
    p2::Point3D,
    p3::Point3D,
    dp1::Point3D,
    dp2::Point3D,
    dp3::Point3D,
)
    x1 = p2.x - p1.x
    y1 = p2.y - p1.y

    x2 = p3.x - p1.x
    y2 = p3.y - p1.y

    u1 = dp2.x - dp1.x
    v1 = dp2.y - dp1.y

    u2 = dp3.x - dp1.x
    v2 = dp3.y - dp1.y

    a = u1 * v2 - u2 * v1
    b = x1 * v1 + y2 * u1 - x2 * v1 - y1 * u2
    c = x1 * y2 - x2 * y1

    if a != 0
        quotient = b^2 - 4a * c
        if quotient >= 0
            d = sqrt(quotient)
            return (-b - d) / 2a, (-b + d) / 2a
        else
            return -123.0, -123.0
        end
    else

        # cool, there just isn't any dependence on t^2, but there is still on t!
        return -c / b, -c / b
    end
end

"""
$(SIGNATURES)

This function saves the mesh object in obj format.
"""
function saveObj!(mesh::Mesh, filename::String; scale=1.0, scalez=1.0, reverse=false, flipxy=false)
    # This function saves the mesh object in stl format
    open(filename, "w") do io
        for vertex in mesh.nodes
            if flipxy
                println(io, "v ", vertex.y * scale, " ", vertex.x * scale, " ", vertex.z * scalez)
            else
                println(io, "v ", vertex.x * scale, " ", vertex.y * scale, " ", vertex.z * scalez)
            end
        end

        for face in mesh.triangles
            if reverse
                println(io, "f ", face.pt3, " ", face.pt2, " ", face.pt1)
            else
                println(io, "f ", face.pt1, " ", face.pt2, " ", face.pt3)
            end
        end

        println(io, "dims ", mesh.width, " ", mesh.height)
    end
end

"""
$(SIGNATURES)
"""
function Obj2Mesh(filename)
    lines = readlines(filename)

    vertexLines = [l for l in lines if startswith(l, "v")]
    nodeList = Vector{Point3D}(undef, size(vertexLines))
    count = 1

    for line in vertexLines
        elements = split(line, " ")
        x = parse(Float64, elements[2])
        y = parse(Float64, elements[3])
        z = parse(Float64, elements[4]) * 10
        pt = Point3D(x, y, z, 0, 0)
        nodeList[count] = pt
        count += 1
    end

    faceLines = [l for l in lines if startswith(l, "f")]
    triangles = Vector{Triangle}(undef, size(faceLines))
    for line in faceLines
        elements = split(line, " ")
        triangle = Triangle(
            parse(Int64, elements[2]),
            parse(Int64, elements[3]),
            parse(Int64, elements[4]),
        )
    end

    dimsLines = [l for l in lines if startswith(l, "dims")]
    elements = split(dimsLines[1], " ")

    return Mesh(nodeList, triangles, parse(Int64, elements[2]), parse(Int64, elements[3]))
end

"""
$(SIGNATURES)
"""
function ∇(f::Matrix{Float64})
    width, height = size(f)

    ∇fᵤ = zeros(Float64, width, height)   # the right edge will be filled with zeros
    ∇fᵥ = zeros(Float64, width, height)   # the buttom edge will be filled with zeros

    for x in 1:width, y in 1:height
        ∇fᵤ[x, y] = (x == width ? 0 : f[x + 1, y] - f[x, y])
        ∇fᵥ[x, y] = (y == height ? 0 : f[x, y + 1] - f[x, y])
    end

    return ∇fᵤ, ∇fᵥ
end

"""
$(SIGNATURES)
"""
function getPixelArea(mesh::Mesh)
    # A Mesh is a grid of 3D points. The X and Y coordinates are not necessarily aligned or square
    # The Z coordinate represents the value. brightness is just proportional to area.
    # pixelAreas = Matrix{Float64}(undef, mesh.width - 1, mesh.height - 1)
    pixelAreas = zeros(mesh.width - 1, mesh.height - 1)

    for x in 1:mesh.width - 1, y in 1:mesh.height - 1
        upperLeft = mesh.nodeArray[x, y]
        upperRight = mesh.nodeArray[x + 1, y]

        lowerLeft = mesh.nodeArray[x, y + 1]
        lowerRight = mesh.nodeArray[x + 1, y + 1]

        #=
        *------*
        |    / |
        |   /  |
        |  /   |
        | /    |
        *------* =#
        pixelAreas[x, y] =
            triangle_area(lowerLeft, upperRight, upperLeft) +
            triangle_area(lowerLeft, lowerRight, upperRight)
    end

    return pixelAreas
end

"""
$(SIGNATURES)
"""
function relax!(matrix::Matrix{Float64}, D::Matrix{Float64})
    # This function implements successive over relaxation for a matrix and its associated error matrix
    # There is a hardcoded assumption of Neumann boundary conditions--that the derivative across the
    # boundary must be zero in all cases. See:
    # https://math.stackexchange.com/questions/3790299/how-to-iteratively-solve-poissons-equation-with-no-boundary-conditions
    # sz = size(matrix)
    # width = sz[1]
    # height = sz[2]
    width, height = size(matrix)
    # ω = 2 / (1 + π / width)
    ω = 1.99
    # println("OMEGA $(ω)")
    max_update = 0
    for y = 1:height
        for x = 1:width
            val = matrix[x, y]

            if x == 1 && y == 1
                # Top left corner
                val_down = matrix[x, y + 1]
                val_right = matrix[x + 1, y]
                delta = ω / 2 * (val_down + val_right - 2 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            elseif x == 1 && y == height
                # Bottom left corner
                val_up = matrix[x, y - 1]
                val_right = matrix[x + 1, y]
                delta = ω / 2 * (val_up + val_right - 2 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            elseif x == width && y == 1
                # Top right corner
                val_down = matrix[x, y + 1]
                val_left = matrix[x - 1, y]
                delta = ω / 2 * (val_down + val_left - 2 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            elseif x == width && y == height
                # Bottom right corner
                val_up = matrix[x, y - 1]
                val_left = matrix[x - 1, y]
                delta = ω / 2 * (val_up + val_left - 2 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta

            elseif x == 1
                # Along the left edge, but not the top or buttom corner
                val_up = matrix[x, y - 1]
                val_down = matrix[x, y + 1]
                val_right = matrix[x + 1, y]
                delta = ω / 3 * (val_up + val_down + val_right - 3 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            elseif x == width
                # Along the right edge, but not the top or buttom corner
                val_up = matrix[x, y - 1]
                val_down = matrix[x, y + 1]
                val_left = matrix[x - 1, y]
                delta = ω / 3 * (val_up + val_down + val_left - 3 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            elseif y == 1
                # Along the top edge, but not the left or right corner
                val_down = matrix[x, y + 1]
                val_left = matrix[x - 1, y]
                val_right = matrix[x + 1, y]
                delta = ω / 3 * (val_down + val_left + val_right - 3 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            elseif y == height
                # Along the bottom edge, but not the left or right corner
                val_up = matrix[x, y - 1]
                val_left = matrix[x - 1, y]
                val_right = matrix[x + 1, y]
                delta = ω / 3 * (val_up + val_left + val_right - 3 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            else
                # The normal case, in the middle of the mesh!
                val_up = matrix[x, y - 1]
                val_down = matrix[x, y + 1]
                val_left = matrix[x - 1, y]
                val_right = matrix[x + 1, y]

                # The new way
                # ∇x₁ =

                # The old way
                delta = ω / 4 * (val_up + val_down + val_left + val_right - 4 * val - D[x, y])
                if abs(delta) > max_update
                    max_update = abs(delta)
                end
                matrix[x, y] += delta
            end
            # node.z = .25 * (node_up.z + node_down.z + node_left.z + node_right.z) # simple averaging
            # node.z += ω/4 * (node_up.z + node_down.z + node_left.z + node_right.z - 4 * node.z)

            # matrix[x, y] += ω/4 * (val_up + val_down + val_left + val_right - 4 * val - D[x, y])
        end
    end

    max_update

    # for y = 1:height
    #     for x = 1:width
    #         val = matrix[x, y]
    #     end
    # end
end

"""
$(SIGNATURES)

This function will take a `grid_definition x grid_definition` matrix and returns a `grid_definition x grid_definition` mesh.

It currently takes the size of the matrix passed as argument.
"""
function matrix_to_mesh(matrix::Matrix{Float64})
    w, h = size(matrix)

    retval = squareMesh(w, h)
    for x in 1:w, y in 1:h
        index = (y - 1) * retval.width + x
        node = retval.nodes[index]
        node.z = matrix[x, y]
    end

    return retval
end

"""
$(SIGNATURES)
"""
function marchMesh!(mesh::Mesh, ϕ::Matrix{Float64})
    ∇ϕᵤ, ∇ϕᵥ = ∇(ϕ)

    imgWidth, imgHeight = size(ϕ)   # should be grid_definitionxgrid_definition

    # For each point in the mesh we need to figure out its velocity
    velocities = Matrix{Point3D}(undef, mesh.width, mesh.height)
    for x in 1:mesh.width, y in 1:mesh.height
        # XY coordinates in the mesh ARE XY coordinates in the image. The mesh just needs an extra row and column
        # at the bottom right edge so that the triangles can be closed

        if x == mesh.width
            u = 0
        else
            u = (y == mesh.height ? ∇ϕᵤ[x, y - 1] : ∇ϕᵤ[x, y])
        end

        if y == mesh.height
            v = 0
        else
            v = (x == mesh.width ? ∇ϕᵥ[x - 1, y] : ∇ϕᵥ[x, y])
        end

        velocities[x, y] = Point3D(-u, -v, 0, 0, 0)
    end

    min_t = 10000
    triangleCount = 1
    for triangle in mesh.triangles
        p1 = mesh.nodes[triangle.pt1]
        p2 = mesh.nodes[triangle.pt2]
        p3 = mesh.nodes[triangle.pt3]

        v1 = velocities[p1.ix, p1.iy]
        v2 = velocities[p2.ix, p2.iy]
        v3 = velocities[p3.ix, p3.iy]

        t1, t2 = findT(p1, p2, p3, v1, v2, v3)

        if 0 < t1 < min_t
            min_t = t1
        end

        if 0 < t2 < min_t
            min_t = t2
        end
        triangleCount += 1
    end

    println("Overall min_t:", min_t)
    δ = min_t / 2

    for point in mesh.nodes
        v = velocities[point.ix, point.iy]
        point.x = v.x * δ + point.x
        point.y = v.y * δ + point.y
    end

    # saveObj(mesh, "gateau.obj")
end

    """
$(SIGNATURES)
"""
function quantifyLoss!(D, suffix, img)
    println("Loss:")
    println("Minimum: $(minimum(D))")
    println("Maximum: $(maximum(D))")
    println("SUM: $(sum(D))")

    blue = zeros(size(D))
    blue[D .> 0] = D[D .> 0]
    red = zeros(size(D))
    red[D .< 0] = -D[D .< 0]
    green = zeros(size(D))

    println(size(blue))
    println(size(red))
    println(size(green))

    rgbImg = RGB.(red, green, blue)'
    save("./examples/loss_$(suffix).png", map(clamp01nan, rgbImg))
end

"""
$(SIGNATURES)
"""
function oneIteration(meshy, img, suffix)
    # Remember meshy is (will be) `grid_definition x grid_definition` just like the image
    # `grid_definition x grid_definition`, so LJ is `grid_definition x grid_definition`.
    LJ = getPixelArea(meshy)
    D = Float64.(LJ - img)
    # Shift D to ensure its sum is zero
    D .-= sum(D) / (512 * 512)

    # Save the loss image as a png
    println(minimum(D))
    println(maximum(D))
    quantifyLoss!(D, suffix, img)

    # These lines are just for plotting the quiver of L, which I needed for the blog post
    # ∇Lᵤ, ∇Lᵥ = ∇(D)
    # plotVAsQuiver(∇Lᵤ, ∇Lᵥ, stride=10, scale=10, max_length=200)
    # println("okay")
    # return

    width, height = size(img)

    ϕ = zeros(width, height)

    println("Building Phi")
    for i = 1:10000
        max_update = relax!(ϕ, D)

        if isnan(max_update)
            println("MAX UPDATE WAS NaN. CANNOT BUILD PHI")
            return
        end

        if i % 500 == 0
            println(max_update)
        end

        if max_update < 0.00001
            println("Convergence reached at step $(i) with max_update of $(max_update)")
            break
        end
    end

    # saveObj!(
    #     matrix_to_mesh(ϕ * 0.02),
    #     "./examples/phi_$(suffix).obj",
    #     reverse=false,
    #     flipxy=true,
    # )
    # plotAsQuiver(ϕ * -1.0, stride=30, scale=1.0, max_length=200, flipxy=true, reversex=false, reversey=false)
    # saveObj(matrix_to_mesh(D * 10), "D_$(suffix).obj")

    # Now we need to march the x,y locations in our mesh according to this gradient!
    marchMesh!(meshy, ϕ)
# saveObj!(meshy, "./examples/mesh_$(suffix).obj", flipxy=true)
end

"""
$(SIGNATURES)
"""
function setHeights!(mesh, heights, heightScale=1.0, heightOffset=10)
    width, height = size(heights)
    for y = 1:height
        for x = 1:width
            mesh.nodeArray[x, y].z = heights[x, y] * heightScale + heightOffset
            if x == 100 && y == 100
                println("Example heights: $(heights[x, y])  and  $(heights[x, y] * heightScale) and $(heights[x, y] * heightScale + heightOffset)")
            end
        end
    end

    # get the side edge
    for y = 1:height
        mesh.nodeArray[width + 1, y].z = mesh.nodeArray[width, y].z
    end

    # get the bottom edge
    for x = 1:width + 1
        mesh.nodeArray[x, height + 1].z = mesh.nodeArray[x, height].z
    end

    # # get the pesky corner!
# mesh.nodeArray[width + 1, height + 1].z = mesh.nodeArray[width, height].z
end

"""
$(SIGNATURES)
"""
function solidify(inputMesh, offset=100)
    width = inputMesh.width
    height = inputMesh.height
    totalNodes = width * height * 2
    nodeList = Vector{Point3D}(undef, totalNodes)
    nodeArrayTop = Matrix{Point3D}(undef, width, height)
    nodeArrayBottom = Matrix{Point3D}(undef, width, height)

    # imagine a 4x4 image. 4 * 2 + 2 * 2 = 12
    numEdgeNodes = width * 2 + (height - 2) * 2

    numTrianglesTop = (width - 1) * (height - 1) * 2
    numTrianglesBottom = numTrianglesTop
    numTrianglesEdges = numEdgeNodes * 2

    totalTriangles = numTrianglesBottom + numTrianglesTop + numTrianglesEdges

    println("Specs: $(width)  $(height)  $(totalNodes)  $(numEdgeNodes)  $(numTrianglesBottom) $(totalTriangles)")

    # Build the bottom surface
    count = 1
    for y = 1:height
        for x = 1:width
            newPoint = Point3D(x, y, -offset, x, y)
            nodeList[count] = newPoint
            nodeArrayBottom[x, y] = newPoint
            count += 1
        end
    end

    # Copy in the top surface
    for y = 1:height
        for x = 1:width
            node = inputMesh.nodeArray[x, y]
            copiedPoint = Point3D(node.x, node.y, node.z, node.ix, node.iy)
            if node.ix != x
                println("OH NO POINTS NOT MATCHED $(x) vs $(node.ix)")
            end
            if node.iy != y
                println("OH NO POINTS NOT MATCHED $(y) vs $(node.iy)")
            end

            nodeList[count] = copiedPoint
            nodeArrayTop[x, y] = copiedPoint
            count += 1
        end
    end

    println("We now have $(count - 1) valid nodes")

    triangles = Vector{Triangle}(undef, totalTriangles)
    # Build the triangles for the bottom surface
    count = 1
    for y = 1:(height - 1)
        for x = 1:(width - 1)
          # here x and y establish the column of squares we're in
            index_ul = (y - 1) * width + x
            index_ur = index_ul + 1

            index_ll = y * width + x
            index_lr = index_ll + 1

            triangles[count] = Triangle(index_ul, index_ll, index_ur)
            count += 1
            triangles[count] = Triangle(index_lr, index_ur, index_ll)
            count += 1
        end
    end

    println("We've filled up $(count - 1) triangles")
    if count != numTrianglesBottom + 1
        println("Hmm aren't count and triangles bottom equal? $(count) vs $(numTrianglesBottom + 1)")
        end

    # Build the triangles for the top surface
    for y = 1:(height - 1)
        for x = 1:(width - 1)
          # here x and y establish the column of squares we're in
            index_ul = (y - 1) * width + x + totalNodes / 2
            index_ur = index_ul + 1

            index_ll = y * width + x + totalNodes / 2
            index_lr = index_ll + 1

            triangles[count] = Triangle(index_ul, index_ur, index_ll)
            count += 1
            triangles[count] = Triangle(index_lr, index_ll, index_ur)
            count += 1
        end
    end

    println("We've filled up $(count - 1) triangles")

    # Build the triangles to close the mesh
    x = 1
    for y = 1:(height - 1)
        ll = (y - 1) * width + x
        ul = ll + totalNodes / 2
        lr = y * width + x
        ur = lr + totalNodes / 2
        triangles[count] = Triangle(ll, ul, ur)
        count += 1
    triangles[count] = Triangle(ur, lr, ll)
        count += 1
    end

    x = width
    for y = 1:(height - 1)
        ll = (y - 1) * width + x
        ul = ll + totalNodes / 2
        lr = y * width + x
        ur = lr + totalNodes / 2
        triangles[count] = Triangle(ll, ur, ul)
        count += 1
        triangles[count] = Triangle(ur, ll, lr)
        count += 1
    end

    y = 1
    for x = 2:width
        ll = (y - 1) * width + x
        ul = ll + totalNodes / 2
        lr = (y - 1) * width + (x - 1)
        ur = lr + totalNodes / 2
        triangles[count] = Triangle(ll, ul, ur)
        count += 1
        triangles[count] = Triangle(ur, lr, ll)
        count += 1
    end

    y = height
    for x = 2:width
        ll = (y - 1) * width + x
        ul = ll + totalNodes / 2
        lr = (y - 1) * width + (x - 1)
        ur = lr + totalNodes / 2
        triangles[count] = Triangle(ll, ur, ul)
        count += 1
        triangles[count] = Triangle(ur, ll, lr)
        count += 1
    end

Mesh(nodeList, nodeArrayBottom, triangles, width, height)
end

"""
$(SIGNATURES)
"""
function findSurface(mesh, image, f, imgWidth)
    width, height = size(image)

    # imgWidth = .1 # m
    # f = 1.0  # m
    H = f
    metersPerPixel = imgWidth / width
    println(metersPerPixel)

    # η = 1.49
    n₂ = 1
    n₁ = 1.49
    Nx = zeros(width + 1, height + 1)
    Ny = zeros(width + 1, height + 1)

    for j = 1:height
        for i = 1:width
            node = mesh.nodeArray[i, j]
            dx = (node.ix - node.x) * metersPerPixel
            dy = (node.iy - node.y) * metersPerPixel

            little_h = node.z * metersPerPixel
            H_minus_h = H - little_h
            dz = H_minus_h

            # k = η * sqrt(dx * dx + dy * dy + H_minus_h * H_minus_h) - H_minus_h
            # Nx[i, j] = 1/k * dx
            # Ny[i, j] = 1/k * dy
            Ny[i, j] = tan(atan(dy / dz) / (n₁ - 1))
            Nx[i, j] = tan(atan(dx / dz) / (n₁ - 1))

        end
    end

    divergence = zeros(width, height)
    # We need to find the divergence of the Vector field described by Nx and Ny

    for j = 1:height
        for i = 1:width
            δx = (Nx[i + 1, j] - Nx[i, j])
            δy = (Ny[i, j + 1] - Ny[i, j])
            divergence[i, j] = δx + δy
        end
    end
    println("Have all the divergences")
    println("Divergence sum: $(sum(divergence))")
    divergence .-= sum(divergence) / (width * height)

    h = zeros(width, height)
    max_update = 0
    for i = 1:10000
        max_update = relax!(h, divergence)

        if i % 500 == 0
            println(max_update)
        end
        if max_update < 0.00001
            println("Convergence reached at step $(i) with max_update of $(max_update)")
            break
        end
    end

    # println("HEIGHT STATS")
    # h .-= sum(h) / (width * height)
    # println(minimum(h))
    # println(maximum(h))
    # println(sum(h))
    # println("-----------")

    # saveObj!(matrix_to_mesh(h * 10), "examples/heightmap.obj")
    h, metersPerPixel
end

"""
$(SIGNATURES)
"""
function testSquareMesh!()
    mesh = squareMesh(100, 50)

    println(mesh.nodeArray[1, 1])
    println(mesh.nodes[1])

    mesh.nodeArray[1, 1].x = 8
    println(mesh.nodeArray[1, 1])
    println(mesh.nodes[1])

    mesh.nodes[1].y += 12
println(mesh.nodeArray[1, 1])
    println(mesh.nodes[1])

end

"""
$(SIGNATURES)
"""
function testSolidify!()
    println("Testing solidification")
    width = 100
    height = 100
    origMesh = squareMesh(width, height)

    for y in 1:height, x in 1:width
        x2 = (x - width / 2) / width
        y2 = (y - height / 2) / height
        value = x2 * x2 + y2 * y2
        origMesh.nodeArray[x, y].z = 15 - value * 25
    end

    saveObj!(origMesh, "./examples/testSolidify.obj")
    solidMesh = solidify(origMesh, 0)
    saveObj!(solidMesh, "./examples/testSolidify2.obj")
end

"""
$(SIGNATURES)
"""
function plotAsQuiver(
    g;
    stride=4,
    scale=300,
    max_length=2,
    flipxy=false,
    reversey=false,
    reversex=false,
)
    h, w = size(g)
    xs = Float64[]
    ys = Float64[]
    us = Float64[]
    vs = Float64[]

    for x in 1:stride:w, y in 1:stride:h
        reversex ? push!(xs, x) : push!(xs, -x)
        reversey ? push!(ys, -y) : push!(ys, y)

        p1 = g[y, x]
        u = (g[y, x + 1] - g[y, x]) * scale
        v = (g[y + 1, x] - g[y, x]) * scale

        u = -u

        reversey && (v = -v)
        reversex && (u = -u)

        # println(u, v)
        u >= 0 ? push!(us, min(u, max_length)) : push!(us, max(u, -max_length))
        v >= 0 ? push!(vs, min(v, max_length)) : push!(vs, max(v, -max_length))
    end

    q =
        flipxy ? quiver(ys, xs, quiver=(vs, us), aspect_ratio=:equal) :
        quiver(xs, ys, quiver=(us, vs), aspect_ratio=:equal)

    display(q)
    readline()
end

"""
$(SIGNATURES)
"""
function plotVAsQuiver(vx, vy; stride=4, scale=300, max_length=2)
    h, w = size(vx)

    xs = Float64[]
    ys = Float64[]
    us = Float64[]
    vs = Float64[]

    for x in 1:stride:w, y in 1:stride:h
        push!(xs, x)
        push!(ys, h - y)

        u = max(vx[x, y], 0.001)
        v = max(vy[x, y], 0.001)

        push!(us, u)
        push!(vs, v)
        # println(u, ": ", v)
    end

    # readline()
    q = quiver(xs, ys, quiver=(us, vs), aspect_ratio=:equal)
    display(q)
    readline()
end

"""
$(SIGNATURES)
"""
function engineer_caustics(img)
    img = Gray.(img)
    img2 = permutedims(img) * 1.0
    width, height = size(img2)

    # meshy is the same size as the image
    meshy = squareMesh(width + 1, height + 1)

    # We need to boost the brightness of the image so that its sum and the sum of the area are equal
    mesh_sum = width * height
    image_sum = sum(img2)
    boost_ratio = mesh_sum / image_sum

    # img3 is `grid_definition x grid_definition` and is normalised to the same (sort of) _energy_ as the
    # original image.
    img3 = img2 .* boost_ratio

    oneIteration(meshy, img3, "it1")
    oneIteration(meshy, img3, "it2")
    oneIteration(meshy, img3, "it3")
    oneIteration(meshy, img3, "it4")
    # oneIteration(meshy, img3, "it5")
    # oneIteration(meshy, img3, "it6")

    artifactSize = 0.1  # meters
    focalLength = 0.2 # meters
    h, metersPerPixel = findSurface(meshy, img3, focalLength, artifactSize)

    setHeights!(meshy, h, 1.0, 10)

    solidMesh = solidify(meshy)
    saveObj!(
        solidMesh,
        "./examples/original_image.obj",
        scale=1 / 512 * artifactSize,
        scalez=1 / 512.0 * artifactSize,
    )

    return meshy, img3
end

"""
$(SIGNATURES)
"""
function main()
    @assert size(ARGS) == (1,) "Intented usage is: julia create_mesh.jl image.png"

    img = load(ARGS[1])
    return engineer_caustics(img)
end

"""
$(SIGNATURES)
"""
main()

max_stars_repo_name: Mehul-Bagaria/causticsEngineering

max_issues_repo_name: Mehul-Bagaria/causticsEngineering

max_forks_repo_name: Mehul-Bagaria/causticsEngineering


text:
"""
A package of useful optical propagation methods and simple 2-dimensional light field data types.
"""
module OpticalPropagation

using FFTW
using Unitful
using RecipesBase

import Base:
    +, -, *, /, \

export
    AbstractLightField2D,
    MonoLightField2D,
    angularspectrum,
    fresnel2,
    fresnel1

include("lightfield2d.jl")
include("angularspectrum.jl")
include("fresnel.jl")
include("utils.jl")

end # module