text:
using ClusterManagers, DSGE, ModelConstructors, FileIO, Plots, Dates, OrderedCollections, Distributed
using Plots.PlotMeasures
GR.inline("pdf")
fn = dirname(@__FILE__)

# This script shows how to use exogenous regime-switching
# for forecasting as well as permanent and
# temporary alternative policies. The format
# of this script mirrors make_packet.jl,
# so users should familiarize themselves with that
# script first.

# In this script, we consider a case where
# in 1990:Q1, the standard deviations of
# all shocks are set to 0.9 of their values in
# the period 1959:Q3-1989:Q4, but by 2015:Q1,
# the standard deviations return to the pre-1990 values.
# Furthermore, we consider the effects of a permanent and temporary
# alternative policy of nominal GDP (NGDP) targeting.

# What do you want to do?
run_full_forecast = true  # Run full distribution forecast
temp_alt          = true  # Temporary NGDP targeting?
perm_alt          = true  # Permanent NGDP targeting?
make_plots        = true  # Make plots
save_plots        = false # Save plots to figurespath(m, "forecast")
add_workers       = false # Run in parallel
n_workers         = 10

# Initialize model objects and desired settings
custom_settings = Dict{Symbol, Setting}(:add_altpolicy_pgap => Setting(:add_altpolicy_pgap, true))
m = Model1002("ss10"; custom_settings = custom_settings)   # We will directly construct the matrix of parameter draws
m10 = Model1002("ss10") # b/c loading from a saved estimation file has not been fully implemented

for model in [m, m10]
    # Not strictly necessary to have the same settings, but it makes it easier for us
    model <= Setting(:sampling_method, :MH)
    usual_model_settings!(model, "181115"; cdvt = "181115", fcast_date = quartertodate("2018-Q4"))
    model <= Setting(:use_population_forecast, false)
    model <= Setting(:saveroot, joinpath("$(fn)", "..", "save"))
    model <= Setting(:dataroot, joinpath("$(fn)", "..", "save", "input_data"))
    model <= Setting(:date_forecast_end, quartertodate("2030-Q1"))
    model <= Setting(:forecast_horizons, DSGE.subtract_quarters(date_forecast_end(model), date_forecast_start(model)) + 1)
    model <= Setting(:forecast_block_size, 40)
    if get_setting(model, :sampling_method) == :SMC
        model <= Setting(:forecast_jstep, 1)
    end
end

forecast_string = ""

# Set up regime switching and forecast settings
m <= Setting(:n_regimes, 3, true, "reg", "")   # How many regimes?
m <= Setting(:regime_switching, true)          # Need to set that the model does have regime-switching
# m <= Setting(:regime_switching_ndraws, 100000) # b/c we don't have an estimation file, need this setting for dimension purposes
m <= Setting(:forecast_block_size, 5000)
m <= Setting(:use_parallel_workers, true)
baseline_regime_dates = Dict{Int, Date}(1 => date_presample_start(m), 2 => Date(1990, 3, 31),
                                        3 => Date(2015, 3, 31))
m <= Setting(:regime_dates, baseline_regime_dates)
setup_regime_switching_inds!(m)

# Some settings for alternative policy
if temp_alt || perm_alt
    # Define values for creating the alternative policy, for either case
    policy               = :ngdp
    replace_policy       = DSGE.ngdp_replace_eq_entries
    policy_eqcond        = DSGE.ngdp_eqcond
    policy_solve         = DSGE.ngdp_solve
    policy_forecast_init = DSGE.ngdp_forecast_init
    pol_str              = "NGDP"
end

# Set up parameters for regime-switching.
# See the full forecast code block for the construction
# of the matrix of parameter draws
overrides = forecast_input_file_overrides(m10)
overrides[:full] = "$(fn)/../test/reference/mhsave_vint=181115.h5"
modal_params = map(x -> x.value, m10.parameters) # The default parameters will be our "modal" parameters for this exercise
θ10 = load_draws(m10, :full)

para_regs = perm_alt ? (1:4) : (1:3) # if permanent altpolicy occurs, then we need to add a fourth regime
for i in para_regs
    adj = (i == 2) ? .9 : 1. # adjustment to the size of the standard deviation

    # For each parameter, we need to instruct it that a regime exists.
    # Note that the regimes MUST be entered in order of the regimes
    # because regime values are stored in an OrderedDict, which is sorted
    # by insertion order.
    ModelConstructors.set_regime_val!(m[:σ_g], i, adj * m10[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_b], i, adj * m10[:σ_b].value)
    ModelConstructors.set_regime_val!(m[:σ_μ], i, adj * m10[:σ_μ].value)
    ModelConstructors.set_regime_val!(m[:σ_ztil], i, adj * m10[:σ_ztil].value)
    ModelConstructors.set_regime_val!(m[:σ_λ_f], i, adj * m10[:σ_λ_f].value)
    ModelConstructors.set_regime_val!(m[:σ_λ_w], i, adj * m10[:σ_λ_w].value)
    ModelConstructors.set_regime_val!(m[:σ_r_m], i, adj * m10[:σ_r_m].value)
    ModelConstructors.set_regime_val!(m[:σ_σ_ω], i, adj * m10[:σ_σ_ω].value)
    ModelConstructors.set_regime_val!(m[:σ_μ_e], i, adj * m10[:σ_μ_e].value)
    ModelConstructors.set_regime_val!(m[:σ_γ], i, adj * m10[:σ_γ].value)
    ModelConstructors.set_regime_val!(m[:σ_π_star], i, adj * m10[:σ_π_star].value)
    ModelConstructors.set_regime_val!(m[:σ_lr], i, adj * m10[:σ_lr].value)
    ModelConstructors.set_regime_val!(m[:σ_z_p], i, adj * m10[:σ_z_p].value)
    ModelConstructors.set_regime_val!(m[:σ_tfp], i, adj * m10[:σ_tfp].value)
    ModelConstructors.set_regime_val!(m[:σ_gdpdef], i, adj * m10[:σ_gdpdef].value)
    ModelConstructors.set_regime_val!(m[:σ_corepce], i, adj * m10[:σ_corepce].value)
    ModelConstructors.set_regime_val!(m[:σ_gdp], i, adj * m10[:σ_gdp].value)
    ModelConstructors.set_regime_val!(m[:σ_gdi], i, adj * m10[:σ_gdi].value)

    for j = 1:DSGE.n_mon_anticipated_shocks(m)
        ModelConstructors.set_regime_val!(m[Symbol("σ_r_m$(j)")], i, adj * m10[Symbol("σ_r_m$(j)")])
    end
end

df = load_data(m; check_empty_columns = false)

# Full-distribution forecast
if run_full_forecast

    # Set up matrix of parameter draws
    m <= Setting(:forecast_jstep, 1)
    m <= Setting(:forecast_block_size, 100)
    param_mat = zeros(100, ModelConstructors.n_parameters_regime_switching(m)) # Need to use a special method to count the additional regimes
    n_pvec = n_parameters(m) # This is just length(m.parameters), so it counts the "unique" number of parameters, excluding regimes
    for i in 1:size(param_mat, 1)
        # Start w/regime 1 b/c all regime 1 parameters are grouped together
        param_mat[i, 1:n_pvec] = map(x -> (abs(x) < 1e-15) ? 0. : x, θ10[i, :])

        # Now do the other regimes, which grouped by parameter rather than by regime
        j = n_pvec # To help index for regime-switching parameters
        for para in m.parameters
            if !isempty(para.regimes)
                for (k, v) in para.regimes[:value] # If there exists multiple regimes
                    if k > 1
                        j += 1
                        drawv = θ10[i, m10.keys[para.key]] # Find the column for the parameter in θ10 using the parameter's key
                        param_mat[i, j] = (k  == 2) ? .9 * drawv : drawv # Adjust the value
                        if abs(param_mat[i, j]) < 1e-15 # Make sure essentially zero parameters are actually zero
                            param_mat[i, j] = 0.
                        end
                    end
                end
            end
        end
    end

    output_vars = [:histpseudo, :histobs, :histstdshocks,
                   :hist4qpseudo, :hist4qobs, :histutpseudo,
                   :forecastpseudo, :forecastobs, :forecastutpseudo,
                   :forecast4qpseudo, :forecast4qobs, :forecaststdshocks]

    if add_workers
        my_procs = addprocs(n_workers)
        @everywhere using DSGE, OrderedCollections
    end

    usual_model_forecast(m, :full, :none, output_vars,
                         forecast_string = forecast_string,
                         density_bands = [.5, .6, .68, .7, .8, .9],
                         check_empty_columns = false,
                         params = param_mat) # Need to pass in parameters directly, see ?forecast_one

    if add_workers
        rmprocs(my_procs)
    end

    # For comparison to the alternative policies. We want the raw output here, so we call
    # forecast_one_draw rataher than the wrappers usual_model_forecast or forecast_one
    # Note m is already calibrated to the "mode", so passing in parameters of same length
    # as m.parameters will just trigger the non-regime-switching updating of the parameters in m.
    # If the parameters of m were not already calibrated, then modal_params would need to be extended
    # to include the correct regime-switching values (rather than just map(x -> x.value, m10.parameter))
    fcast = DSGE.forecast_one_draw(m, :mode, :none, output_vars, modal_params,
                                   df; regime_switching = true, n_regimes = get_setting(m, :n_regimes))
end

if perm_alt
    # m <= Setting(:replace_eqcond, false)  # Make sure this setting is false
    m <= Setting(:pgap_type, policy)      # Use NGDP
    m <= Setting(:pgap_value, 12.0)       # parametrizing the NGDP rule
    m <= Setting(:regime_switching, true) # Regime-switching should still be on b/c historical regimes
    m <= Setting(:regime_dates, baseline_regime_dates) # Make sure to use historical regimes
    setup_regime_switching_inds!(m)

    m <= Setting(:gensys2, false) # Make sure we use the normal gensys algorithm

    setup_permanent_altpol!(m, AltPolicy(policy, policy_eqcond, policy_solve,
                                         forecast_init = policy_forecast_init)) # Set up permanent alternative policy
    fcast_permalt = DSGE.forecast_one_draw(m, :mode, :none, output_vars, modal_params,
                                           df; regime_switching = true, n_regimes = get_setting(m, :n_regimes))
end

if temp_alt
    # First, need to set up new set of regime dates to implement a temporary alternative policy (temp alt policy)
    temp_n_regimes = 16 # 3 historical regimes, 12 quarters or 3 years of the temp alt policy, and a return to normal in the last eregime
    temp_regime_dates = Dict{Int, Date}()
    temp_regime_dates[1] = date_presample_start(m)
    for (k, v) in baseline_regime_dates # Add the historical regimes
        temp_regime_dates[k] = v
    end
    end_date = Date(2018, 12, 31) + Dates.Year(3) + Dates.Month(3)
    for (i, date) in zip(4:temp_n_regimes,
                         Date(2018, 12, 31):Dates.Month(3):end_date) # add forecast dates
        temp_regime_dates[i] = date
    end
    m <= Setting(:regime_dates, temp_regime_dates)
    m <= Setting(:regime_switching, true)
    setup_regime_switching_inds!(m)

    # Add parameter values for additional regimes
    for i in 4:temp_n_regimes
        ModelConstructors.set_regime_val!(m[:σ_g], i, m10[:σ_g].value)
        ModelConstructors.set_regime_val!(m[:σ_b], i, m10[:σ_b].value)
        ModelConstructors.set_regime_val!(m[:σ_μ], i, m10[:σ_μ].value)
        ModelConstructors.set_regime_val!(m[:σ_ztil], i, m10[:σ_ztil].value)
        ModelConstructors.set_regime_val!(m[:σ_λ_f], i, m10[:σ_λ_f].value)
        ModelConstructors.set_regime_val!(m[:σ_λ_w], i, m10[:σ_λ_w].value)
        ModelConstructors.set_regime_val!(m[:σ_r_m], i, m10[:σ_r_m].value)
        ModelConstructors.set_regime_val!(m[:σ_σ_ω], i, m10[:σ_σ_ω].value)
        ModelConstructors.set_regime_val!(m[:σ_μ_e], i, m10[:σ_μ_e].value)
        ModelConstructors.set_regime_val!(m[:σ_γ], i, m10[:σ_γ].value)
        ModelConstructors.set_regime_val!(m[:σ_π_star], i, m10[:σ_π_star].value)
        ModelConstructors.set_regime_val!(m[:σ_lr], i, m10[:σ_lr].value)
        ModelConstructors.set_regime_val!(m[:σ_z_p], i, m10[:σ_z_p].value)
        ModelConstructors.set_regime_val!(m[:σ_tfp], i, m10[:σ_tfp].value)
        ModelConstructors.set_regime_val!(m[:σ_gdpdef], i, m10[:σ_gdpdef].value)
        ModelConstructors.set_regime_val!(m[:σ_corepce], i, m10[:σ_corepce].value)
        ModelConstructors.set_regime_val!(m[:σ_gdp], i, m10[:σ_gdp].value)
        ModelConstructors.set_regime_val!(m[:σ_gdi], i, m10[:σ_gdi].value)

        for j = 1:DSGE.n_mon_anticipated_shocks(m)
            ModelConstructors.set_regime_val!(m[Symbol("σ_r_m$(j)")], i, m10[Symbol("σ_r_m$(j)")].value)
        end
    end

    # Now set up settings for temp alt policy
    m <= Setting(:gensys2, true) # Temporary alternative policies use a special gensys algorithm
    m <= Setting(:replace_eqcond, true) # This new gensys algo replaces eqcond matrices, so this step is required
    regime_eqcond_info = Dict{Int, DSGE.EqcondEntry}() # Which eqcond to use in which periods
    for i in 4:(temp_n_regimes - 1)
        regime_eqcond_info[i] = DSGE.EqcondEntry(DSGE.ngdp()) # Use the alternative policy for these regimes!
    end
    regime_eqcond_info[temp_n_regimes] = DSGE.EqcondEntry(default_policy()) # Use the alternative policy for these regimes!
    m <= Setting(:regime_eqcond_info, regime_eqcond_info) # Add mapping of regimes to new eqcond matrices

    m <= Setting(:pgap_value, 12.0)  # parametrizing the NGDP rule
    m <= Setting(:pgap_type, policy)

    fcast_tempalt = DSGE.forecast_one_draw(m, :mode, :none, output_vars, modal_params,
                                         df; regime_switching = true, n_regimes = get_setting(m, :n_regimes))
end

if make_plots
    plots_dict = Dict()

    # Define some plotting functions
    function make_comp_plot(m::AbstractDSGEModel, plot_mat::OrderedDict{String, Dict{Symbol, Array{Float64}}},
                            obs::Symbol, product::Symbol,
                            dates_plot, data_plot,
                            fcast_date, nfcast;
                            styles::Vector{Symbol} = repeat([:solid], length(keys(plot_mat))),
                            colors::Vector{Symbol} = repeat([:red], length(keys(plot_mat))))
        series_ind = if product ==:forecastobs
            m.observables[obs]
        else
            m.pseudo_observables[obs]
        end

        if product == :forecastobs
            dates_plot = data_plot[!, :date]
            histobs_plot = obs == :obs_hours ? data_plot[!, obs] : 4 .* data_plot[!, obs]
            fcast_date_tmp = vcat(dates_plot[end - 1], fcast_date)
            dates_plot = map(x -> DSGE.prev_quarter(x), dates_plot)
            p = Plots.plot(dates_plot, histobs_plot, color = :black,
                           label = "Data", linewidth = 3, legend = :bottomright, left_margin = 20px)
        else
            fcast_date_tmp = fcast_date
            p = Plots.plot()
        end

        for (i, (key, val)) in enumerate(plot_mat)

            if obs == :obs_hours || product == :forecastpseudo
                forecast_plot = val[product][series_ind, :][1:nfcast]
            else
                forecast_plot = 4 .* val[product][series_ind, :][1:nfcast]
            end
            if product == :forecastobs
                forecast_plot = vcat(histobs_plot[end], forecast_plot)
            end

            plot!(fcast_date_tmp, forecast_plot, color = colors[i],
                  label = key, linewidth = 1, linestyle = styles[i])
            if obs == :obs_nominalrate
                ylims!((-2, 4))
            end
        end
        return p
    end

    function make_ngdp_plot(m::AbstractDSGEModel, plot_mat::OrderedDict{String, Dict{Symbol, Array{Float64}}},
                            dates_plot, data_plot,
                            fcast_date, nfcast;
                            styles::Vector{Symbol} = repeat([:solid], length(keys(plot_mat))),
                            colors::Vector{Symbol} = repeat([:red], length(keys(plot_mat))))
        fcast_permalt_date = map(x -> DSGE.prev_quarter(x), DSGE.get_quarter_ends(date_forecast_start(m),
                                                                                DSGE.quartertodate("2024-Q4")))
        nfcast_permalt     = length(fcast_permalt_date) - 1

        p = Plots.plot()
        for (i, (key, val)) in enumerate(plot_mat)
            plotobsgdp = vcat(df[end - 1, :obs_gdp], val[:forecastobs][m.observables[:obs_gdp], 1:nfcast_permalt])
            plotpi = vcat(df[end - 1, :obs_gdpdeflator],
                          val[:forecastobs][m.observables[:obs_gdpdeflator], 1:nfcast_permalt])
            ngdp_level = cumprod(exp.(plotobsgdp ./ 100.) + plotpi ./ 100.)
            ngdp_level ./= ngdp_level[1]
            ngdp_level .-= 1.
            ngdp_level .*= 100

            plot!(fcast_permalt_date[1:end], ngdp_level[1:end], color = colors[i], label = key, linewidth = 3, linestyle = styles[i])
        end
        return p
    end

    data_plot  = df[df[:, :date] .>= Date(2016, 3, 31), :]
    fcast_date = map(x -> DSGE.prev_quarter(x), DSGE.get_quarter_ends(date_forecast_start(m), date_forecast_end(m)))
    nfcast     = length(fcast_date)
    dates_plot = map(x -> DSGE.prev_quarter(x), data_plot[!, :date])

    for obs in [:obs_gdp, :obs_corepce, :obs_gdpdeflator, :obs_nominalrate, :obs_investment, :obs_consumption,
                :obs_spread, :obs_wages, :obs_hours]
        plots_dict[obs] = make_comp_plot(m, OrderedDict("Temp $(pol_str)" => fcast_tempalt,
                                                          "Baseline" => fcast,
                                                          "$(pol_str)" => fcast_permalt),
                                         obs, :forecastobs,
                                         dates_plot, data_plot, fcast_date, nfcast,
                                         colors = [:red, :blue, :purple])
    end

    plots_dict[:cum_NominalGDPLevel] = make_ngdp_plot(m, OrderedDict("Temp $(pol_str)" => fcast_tempalt,
                                                                       "Baseline" => fcast,
                                                                       "$(pol_str)" => fcast_permalt),
                                                      dates_plot, data_plot, fcast_date, nfcast,
                                                      colors = [:red, :blue, :purple])
end

if save_plots
    if !isdir(figurespath(m, "forecast"))
        mkdir(figurespath(m, "forecast"))
    end

    for (k, p) in plots_dict
        Plots.savefig(p, joinpath(figurespath(m, "forecast"), "$(pol_str)_$(string(k)).pdf"))
    end
end

nothing


text:
include("geneticAlgorithm.jl")
genesLength = 6
populationSize = 50
crossoverPoint = 3
mutationPercentage = 0.2
elitePercentage = 0.2
repeatSize = 3

population = generatePopulation(populationSize, genesLength)
calculatePopulationFitness!(population)
printPopulation(population)
popGen, repeatNum, population = geneticAlgorithm!(population, elitePercentage, crossoverPoint, mutationPercentage, repeatSize)
printPopulation(population)
print("Ukupan broj generacija je $popGen \n")
print("Broj ponavljanja poslednjeg najboljeg finess-a je $repeatNum \n")

ext: jl

max_stars_repo_name: pikacxe/Primenjeni-Algoritmi

max_issues_repo_name: pikacxe/Primenjeni-Algoritmi

max_forks_repo_name: pikacxe/Primenjeni-Algoritmi


text:
# NOTE: this is a quick and dirty implementation of the Stillinger-Weber model;
# in the future it will be good to do this more elegantly / more efficiently
# there is huge overhead in this code that can be
#   significantly improved in terms of performance, but for now we just
#   want a correct and readable code

# [Stillinger/Weber, PRB 1985]
# ---------------------------------------------------
# v2 = ϵ f2(r_ij / σ); f2(r) = A (B r^{-p} - r^{-q}) exp( (r-a)^{-1} )
# v3 = ϵ f3(ri/σ, rj/σ, rk/σ); f3 = h(rij, rij, Θjik) + ... + ...
# h(rij, rik, Θjik) = λ exp[ γ (rij-a)^{-1} + γ (rik-a)^{-1} ] * (cosΘjik+1/3)^2
#       >>>
# V2 = 0.5 * ϵ * A (B r^{-p} - r^{-q}) * exp( (r-a)^{-1} )
# V3 = √ϵ * λ exp[ γ (r-a)^{-1} ]
#
# Parameters from QUIP database:
# -------------------------------
# <per_pair_data atnum_i="14" atnum_j="14" AA="7.049556277" BB="0.6022245584"
         # p="4" q="0" a="1.80" sigma="2.0951" eps="2.1675" />
# <per_triplet_data atnum_c="14" atnum_j="14" atnum_k="14" lambda="21.0"
         # gamma="1.20" eps="2.1675" />

using ForwardDiff
using LinearAlgebra: dot
using JuLIP: JVecF

export StillingerWeber

"""
`bondangle(S1, S2) -> (dot(S1, S2) + 1.0/3.0)^2`

* not this assumes that `S1, S2` are normalised
* see `bondangle_d` for the derivative
"""
bondangle(S1, S2) = (dot(S1, S2) + 1.0/3.0)^2

"""
`b := bondangle(S1, S2)` then

`bondangle_d(S1, S2, r1, r2) -> b, db1, db2`

where `dbi` is the derivative of `b` w.r.t. `Ri` where `Si= Ri/ri`.
"""
function bondangle_d(S1, S2, r1, r2)
   d = dot(S1, S2)
   b1 = (1.0/r1) * S2 - (d/r1) * S1
   b2 = (1.0/r2) * S1 - (d/r2) * S2
   return (d+1.0/3.0)^2, 2.0*(d+1.0/3.0)*b1, 2.0*(d+1.0/3.0)*b2
end

function _ad_bondangle_(R)
   R1, R2 = R[1:3], R[4:6]
   r1, r2 = norm(R1), norm(R2)
   return bondangle(R1/r1, R2/r2)
end

# TODO: need a faster implementation of bondangle_dd
function bondangle_dd(R1, R2)
   R = [R1; R2]
   hh = ForwardDiff.hessian(_ad_bondangle_, R)
   h = zeros(JMatF, 2,2)
   h[1,1] = JMatF(hh[1:3,1:3])
   h[1,2] = JMatF(hh[1:3, 4:6])
   h[2,1] = JMatF(hh[4:6,1:3])
   h[2,2] = JMatF(hh[4:6, 4:6])
   return h
end

"""
Stillinger-Weber potential with parameters for Si.

Functional form and default parameters match the original SW potential
from [Stillinger/Weber, PRB 1985].

The `StillingerWeber` type can also by "abused" to generate arbitrary
bond-angle potentials of the form
   ∑_{i,j} V2(rij) + ∑_{i,j,k} V3(rij) V3(rik) (cos Θijk + 1/3)^2

Constructor admits the following key-word parameters:
`ϵ=2.1675, σ = 2.0951, A=7.049556277, B=0.6022245584,
                  p = 4, a = 1.8, λ=21.0, γ=1.20`

which enter the potential as follows:
```
V2(r) = 0.5 * ϵ * A * (B * (r/σ)^(-p) - 1.0) * exp(1.0 / (r/σ - a))
V3(r) = sqrt(ϵ * λ) * exp(γ / (r/σ - a))
```
"""
struct StillingerWeber{P1,P2} <: SitePotential
   V2::P1
   V3::P2
end

@pot StillingerWeber

cutoff(calc::StillingerWeber) = max(cutoff(calc.V2), cutoff(calc.V3))

function StillingerWeber(; brittle = false,
               ϵ=2.1675, σ = 2.0951, A=7.049556277, B=0.6022245584,
               p = 4, a = 1.8, λ = brittle ? 42.0 : 21.0, γ=1.20 )
   cutoff = a*σ-1e-2
   V2 = @analytic(r -> (ϵ*A) * (B*(r/σ)^(-p) - 1.0) * exp(1.0/(r/σ - a))) *
         HS(cutoff)
   V3 = @analytic(r -> sqrt(ϵ * λ) * exp( γ / (r/σ - a) )) * HS(cutoff)
   return StillingerWeber(V2, V3)
end

alloc_temp(V::StillingerWeber, N::Integer) =
      (  R = zeros(JVecF, N),
         S = zeros(JVecF, N),
         V3 = zeros(Float64, N)  )

function evaluate!(tmp, calc::StillingerWeber, R)
   Es = 0.0
   # three-body contributions
   for i = 1:length(R)
      r = norm(R[i])
      tmp.S[i] = R[i] / r
      tmp.V3[i] = calc.V3(r)
      Es += 0.5 * calc.V2(r)
   end
   for i1 = 1:(length(R)-1), i2 = (i1+1):length(R)
      Es += tmp.V3[i1] * tmp.V3[i2] * bondangle(tmp.S[i1], tmp.S[i2])
   end
   return Es
end

alloc_temp_d(V::StillingerWeber, N::Integer, T = Float64) =
      (  dV = zeros(JVec{T}, N),
         R = zeros(JVecF, N),
         r = zeros(T, N),
         S = zeros(JVec{T}, N),
         V3 = zeros(T, N),
         gV3 = zeros(JVec{T}, N)  )

alloc_temp_dd(V::StillingerWeber, N::Integer) =
      (  r = zeros(Float64, N),
         S = zeros(JVecF, N),
         V3 = zeros(Float64, N),
         gV3 = zeros(JVecF, N)  )

function evaluate_d!(dEs, tmp, calc::StillingerWeber, R::AbstractVector{<:JVec})
   for i = 1:length(R)
      tmp.r[i] = r = norm(R[i])
      tmp.S[i] = R[i] / r
      tmp.V3[i] = calc.V3(r)
      tmp.gV3[i] = grad(calc.V3, r, R[i])
      dEs[i] = 0.5 * grad(calc.V2, r, R[i])
   end
   for i1 = 1:(length(R)-1), i2 = (i1+1):length(R)
      a, b1, b2 = bondangle_d(tmp.S[i1], tmp.S[i2], tmp.r[i1], tmp.r[i2])
      dEs[i1] += (tmp.V3[i1] * tmp.V3[i2]) * b1 + (tmp.V3[i2] * a) * tmp.gV3[i1]
      dEs[i2] += (tmp.V3[i1] * tmp.V3[i2]) * b2 + (tmp.V3[i1] * a) * tmp.gV3[i2]
   end
   return dEs
end

function _ad_dV(V::StillingerWeber, R_dofs)
   R = vecs( reshape(R_dofs, 3, length(R_dofs) ÷ 3) )
   r = norm.(R)
   dV = zeros(eltype(R), length(R))
   tmpd = alloc_temp_d(V, length(R), eltype(R[1]))
   evaluate_d!(dV, tmpd, V, R)
   return mat(dV)[:]
end

function _ad_ddV!(hEs, V::StillingerWeber, R::AbstractVector{JVec{T}}) where {T}
   ddV = ForwardDiff.jacobian( Rdofs -> _ad_dV(V, Rdofs), mat(R)[:] )
   # convert into a block-format
   n = length(R)
   for i = 1:n, j = 1:n
      hEs[i, j] = ddV[ ((i-1)*3).+(1:3), ((j-1)*3).+(1:3) ]
   end
   return hEs
end

evaluate_dd!(hEs, tmp, V::StillingerWeber, R) = _ad_ddV!(hEs, V, R)

# function hess(V::StillingerWeber, r, R)
#    n = length(r)
#    hV = zeros(JMatF, n, n)
#
#    # two-body contributions
#    for (i, (r_i, R_i)) in enumerate(zip(r, R))
#       hV[i,i] += hess(V.V2, r_i, R_i)
#    end
#
#    # three-body terms
#    S = [ R1/r1 for (R1,r1) in zip(R, r) ]
#    V3 = [ V.V3(r1) for r1 in r ]
#    dV3 = [ grad(V.V3, r1, R1) for (r1, R1) in zip(r, R) ]
#    hV3 = [ hess(V.V3, r1, R1) for (r1, R1) in zip(r, R) ]
#
#    for i1 = 1:(length(r)-1), i2 = (i1+1):length(r)
#       # Es += V3[i1] * V3[i2] * bondangle(S[i1], S[i2])
#       # precompute quantities
#       ψ, Dψ_i1, Dψ_i2 = bondangle_d(S[i1], S[i2], r[i1], r[i2])
#       Hψ = bondangle_dd(R[i1], R[i2])  # <<<< this should be SLOW (AD)
#       # assemble local hessian contributions
#       hV[i1,i1] +=
#          hV3[i1] * V3[i2] * ψ       +   dV3[i1] * V3[i2] * Dψ_i1' +
#          Dψ_i1 * V3[i2] * dV3[i1]'  +   V3[i1] * V3[i2] * Hψ[1,1]
#       hV[i2,i2] +=
#          V3[i2] * hV3[i2] * ψ       +   V3[i1] * dV3[i2] * Dψ_i2' +
#          Dψ_i2 * V3[i1] * dV3[i2]'  +   V3[i1] * V3[i2] * Hψ[2,2]
#       hV[i1,i2] +=
#          dV3[i1] * dV3[i2]' * ψ     +   V3[i1] * Dψ_i1 * dV3[i2]' +
#          dV3[i1] * V3[i2] * Dψ_i2'  +   V3[i1] * V3[i2] * Hψ[1,2]
#       hV[i2, i1] +=
#          dV3[i2] * dV3[i1]' * ψ     +   V3[i1] * Dψ_i2 * dV3[i1]' +
#          dV3[i2] * V3[i1] * Dψ_i1'  +   V3[i1] * V3[i2] * Hψ[2,1]
#    end
#    return hV
# end

function precon!(hEs, tmp, V::StillingerWeber, R::AbstractVector{JVec{T}}, innerstab=0.0
                 ) where {T}
   n = length(R)
   r = tmp.r
   V3 = tmp.V3
   S = tmp.S

   # two-body contributions
   for (i, R1) in enumerate(R)
      r[i] = norm(R1)
      V3[i] = V.V3(r[i])
      S[i] = R1 / r[i]
      hEs[i,i] += precon!(nothing, V.V2, r[i], R1)
   end

   # three-body terms
   for i1 = 1:(n-1), i2 = (i1+1):n
      Θ = dot(S[i1], S[i2])
      dΘ1 = (T(1.0)/r[i1]) * S[i2] - (Θ/r[i1]) * S[i1]
      dΘ2 = (T(1.0)/r[i2]) * S[i1] - (Θ/r[i2]) * S[i2]
      # ψ = (Θ + 1/3)^2, ψ' = (Θ + 1/3), ψ'' = 2.0
      a = abs((V3[i1] * V3[i2] * T(2.0)))
      hEs[i1,i2] += a * dΘ1 * dΘ2'
      hEs[i1,i1] += a * dΘ1 * dΘ1'
      hEs[i2, i2] += a * dΘ2 * dΘ2'
      hEs[i2, i1] += a * dΘ2 * dΘ1'
   end

   return hEs
end


text:
using PyCall

@pyimport math

math.sin(math.pi / 2)


text:
module TestingEasyServers

using Test
using Sockets
using EasyServers
using Base: IOError

# a simple echo callback
function echo(sock::TCPSocket)
    req = readline(sock, keep=false)
    println(sock, "ok: ", req)
    req == "quit" ? :quit : :continue
end

const SendError = (VERSION < v"1.1" ? ArgumentError : IOError)

tryconnect(port::Integer; kwds...) =
    tryconnect(Sockets.localhost, port; kwds...)

function tryconnect(addr, port::Integer; secs::Real=5.0, verb::Bool=false)
    tmax = max(Float64(secs), 0.0)
    Δt = 0.001
    tsum = 0.0
    while true
        sock = try
            connect(addr, port)
        catch ex
            isa(ex, IOError) || rethrow(ex)
            nothing
        end
        if sock === nothing && tsum < tmax
            verb && println(
                stderr, "connection failed, $(tmax - tsum) seconds remaining")
            sleep(min(Δt, tmax - tsum))
            tsum += Δt # do not add actual delay to avoid rounding errors
            Δt *= 2
        else
            return sock
        end
    end
end

@testset "Simple server" begin
    # Create the server ourself but with an automatic port number.
    port, server = listenany(0)
    @async runserver(server, echo)
    sock = tryconnect(port; secs=10.0)
    @test isa(sock, TCPSocket)
    println(sock, "hello")
    @test readline(sock; keep=false) == "ok: hello"
    @test isopen(sock)
    println(sock, "quit")
    @test readline(sock; keep=false) == "ok: quit"
    @test readline(sock; keep=false) == ""
    @test_throws SendError println(sock, "are you there?")
    @test isopen(server) == false

    # Resuse the same port number.
    @async runserver(port, echo)
    sock = tryconnect(port; secs=10.0)
    @test isa(sock, TCPSocket)
    println(sock, "hello")
    @test readline(sock; keep=false) == "ok: hello"
    @test isopen(sock)
    println(sock, "quit")
    @test readline(sock; keep=false) == "ok: quit"
    @test readline(sock; keep=false) == ""
    @test_throws SendError println(sock, "are you there?")
end

end # module


text:
__precompile__(true)

"""
Main module for NQSIM, a proof of concept for a node-based version of QSIM, later potentially ported or integrated with MATSim
"""

module NQSIM
using DataStructures

q = Queue(Int64)
enqueue!(q, 1)
enqueue!(q, 2)
enqueue!(q, 3)
println("hello world", dequeue!(q), dequeue!(q), q)
end

max_stars_repo_name: muellermichel/NQSIM

max_issues_repo_name: muellermichel/NQSIM

max_forks_repo_name: muellermichel/NQSIM


text:
#------------------------------------------------------------------------------
# Author: Edward Tate <edward.tate@erlang-solutions.com>
#------------------------------------------------------------------------------

export simulate_seq

#------------------------------------------------------------------------------
# Types
#------------------------------------------------------------------------------

type Goal
    home_goal::Bool
    away_goal::Bool
    time::Int
end

type Match
    home_team::String
    away_team::String
    final_home_score::Int
    final_away_score::Int
    goals::Array{Goal}
end

#------------------------------------------------------------------------------
# Data Generation
#------------------------------------------------------------------------------

function sum_home_goals (goals)
    acc = 0
    for i = 1:length(goals)
        if goals[i].home_goal
            acc += 1
        end
    end
    acc
end

function sum_away_goals (goals)
    acc = 0
    for i = 1:length(goals)
        if goals[i].away_goal
            acc += 1
        end
    end
    acc
end

# The next outcome is determined by the probability of home score 'hp'
# and the probability of away scoring 'ap'.
# Both scoring is a possibility.

const home_strength = (2.27 * 0.36 * 1.37) / 90
const away_strength = (2.55 * 0.28) / 90

function next_outcome (time::Integer)
    # random number limit ?
    r = rand(Float32, 2)
    home_probability = (home_strength + (0.70 * time))
    away_probability = (away_strength + (0.65 * time))
    home_outcome = r[1] < home_probability
    away_outcome = r[2] < away_probability
    Goal(home_outcome, away_outcome, time)
end

function seq (typ, f, range)
    output::Array{typ} = Array(typ, length(range))
    for i in range
        output[i] = f(i)
    end
    output
end

function next_match (match_index)
    timeline_lim = 90 + rand(0:5)

    goals = seq(Goal, next_outcome, 1:timeline_lim)
    home_score = sum_home_goals(goals)
    away_score = sum_away_goals(goals)

    Match("", "", home_score, away_score, goals)
end

function simulate_seq (niter)
    out = seq(Match, next_match, 1:niter)

    total_home_score = 0
    total_away_score = 0
    for i in 1:niter
        total_home_score += out[i].final_home_score
        total_away_score += out[i].final_away_score
    end

    (uint32(total_home_score), uint32(total_away_score))
end

max_stars_repo_name: esl/FsFootball


text:
# Debug.

using Crayons
using TextWrap

using REPL

include("debug_recompile_trigger.jl") # The Makefile touches this file to force
                                      # recompilation. e.g. for ENV changes.

# Debug Level.

const DEFAULT_DEBUG_LEVEL = 0

"""
Set `ENV["JULIA_UNIX_IO_DEBUG_LEVEL"]` at compile time to control
the verbosity of debug messages.
Use `set_debug_level(n)` to change the debug level at runtime.
"""
DEBUG_LEVEL = parse(Int, get(ENV, "JULIA_UNIX_IO_DEBUG_LEVEL",
                                  "$DEFAULT_DEBUG_LEVEL"))

"""
### `set_debug_level` -- Control the verbosity of debug messages.

    set_debug_level(n)

`n` can't be higher than the compiled in `BASE_DEBUG_LEVEL`
(because the `@db` marcro reduces these messages to `:()` at compile time.)

Defaults to `ENV["JULIA_UNIX_IO_DEBUG_LEVEL"]`.
"""
function set_debug_level(n)
    global DEBUG_LEVEL
    @require n <= BASE_DEBUG_LEVEL
    DEBUG_LEVEL = n
end

const BASE_DEBUG_LEVEL = DEBUG_LEVEL

# Abbreviations.

db_t(t) = debug_time(t)
db_c(c,p) = constant_name(c;prefix=p)

# State.

"""
`task` is the task that produced the most recent log message.
`t0` is the `time()` the process started.
"""
mutable struct GlobalDebug
    task::Int
    t0::Float64
    lock::Base.ThreadSynchronizer
    GlobalDebug() = new(-1, Inf, Base.ThreadSynchronizer())
end

const global_debug = GlobalDebug()

function debug_init()
    global DEBUG_LEVEL
    global_debug.t0 = time()

    env_level = parse(Int, get(ENV, "JULIA_UNIX_IO_DEBUG_LEVEL", "0"))
    if env_level > DEBUG_LEVEL
        @error "ENV[\"JULIA_UNIX_IO_DEBUG_LEVEL\"] = $env_level, " *
               "but compiled-in DEBUG_LEVEL = $DEBUG_LEVEL.\n" *
               "Debug messages above $DEBUG_LEVEL will not be displayed " *
               "until the the UnixIO module is re-compiled."
    elseif env_level < DEBUG_LEVEL
        @info "ENV[\"JULIA_UNIX_IO_DEBUG_LEVEL\"] = $env_level -> DEBUG_LEVEL"
        DEBUG_LEVEL = env_level
    end
end

"""
`indent` is a stack of string pairs used to indent debug messages.
The first string of the pair contains the full call path.
The second string is padded with whitespace.

`bg_color` is used to render messages for this task.
"""
mutable struct TaskDebug
    indent::Vector{Tuple{String,String}}
    bg_color::Crayon
    prev::Tuple
    leader::String
    TaskDebug() = new([("","")], popfirst!(debug_background_colors), (), "")
end

"""
Get the `TaskDebug` struct for the current task.
"""
function task_local_debug()
    key = :UNIX_IO_DEBUG_STATE
    s = try
        task_local_storage(key)
    catch
        task_local_storage(key, TaskDebug())
    end
    s
end

const DEBUG_FILENAME_WIDTH=24
const DEBUG_BLANK_LINE_INDENT = 35 + DEBUG_FILENAME_WIDTH
const DEBUG_MIN_INDENT=6
"""
Indent this task's debug messages for new call to `function_name`.
"""
@noinline function debug_indent_push(function_name::String)
    function_name = function_name[1:min(length(function_name),
                                    DEBUG_MIN_INDENT)]
    pad = repeat(" ", length(function_name))
    v = task_local_debug().indent
    i, j = v[end]
    i = i == "" ? function_name : string(i, " │ ", function_name)
    j = j == "" ? pad : string(j, " │ ", pad)
    push!(v, (i, j))
end

"""
Revert to the indentation used before `debug_indent_push`.
"""
debug_indent_pop() = pop!(task_local_debug().indent)

# Printing.

"""
IO configuration for printing debug messages.
"""
debug_io_context(io=Base.stderr) =
    IOContext(io, :color => true,
                  :module => @__MODULE__,
                  :compact => true,
                  :limit => true,
                  :displaysize => (1,30))

"""
    debug_print(debug_level, lineno, message; [prefix= " │ "]

Print a debug `message` colored for `debug_level`.
Annotate the message with:
 - a timestamp,
 - `lineno`,
 - Thread and Task IDs.
"""
debug_print(n, l, v; kw...) = debug_print(n, l, "", v; kw...)

@noinline function debug_print(n::Int, lineno::String,
                               message::String, value="";
                               new_function="",
                               function_lineno="",
                               function_return="",
                               prefix::String=" │ ")
    @nospecialize
    DEBUG_LEVEL < n && return
    @ccall(jl_generating_output()::Cint) == 1 && return

    io = IOBuffer()
    ioc = debug_io_context(io)

    color = debug_level_colors[max(1,n)]
    bg_color = task_local_debug().bg_color

    # Has anything been logged since function entry?
    task_local = task_local_debug()
    function_indented = task_local.prev != (function_return, function_lineno)
    task_local.prev = (new_function, lineno)

@lock global_debug.lock begin

    # Create a Task ID and check for task-switch.
    task = Int(pointer_from_objref(current_task())) & 0xffff
    task_switched = task != global_debug.task
    global_debug.task = task

    # Use verbose indentation if there was a task-switch.
    indent = task_local.indent[end][2] # FIXMEtask_switched ? 1 : 2]

    if function_return != ""
        task_local.leader = ""
        if length(function_return) > DEBUG_MIN_INDENT
            if function_indented
                prefix = " ╰" *
                         repeat("─", length(function_return) - 1
                                     - DEBUG_MIN_INDENT) *
                         "▶ "
            else
                prefix = repeat(" ", length(function_return)
                                     - DEBUG_MIN_INDENT) *
                         " ╰▶ "
            end
        else
            prefix = " ╰▶ "
        end
    end

    if indent != ""
        indent *= prefix
    end

    thread_offset = repeat("   ", Threads.threadid()-1)

    if task_switched
        debug_write("\n")
    end

    if function_indented && task_local.leader != ""
        debug_write(task_local.leader)
        task_local.leader = ""
    end

    dbprint(ioc,
        bg_color,

        color, "[ UnixIO $n: ", inv(color),

        # Timestamp.
        lpad(debug_time(time()), 8), " ",

        # File and Line Number.
        lineno, " ",

        # Thread and Task IDs
        (task_switched ? (lpad(Threads.threadid(), 2), ".",
                          uppercase(string(task; base=16, pad=4)), " ")
                       : (repeat(" ", 8),))...,

        # Indented message.
        thread_offset, indent, message, value,

        "\e[0K", # Clear to end of line
        inv(bg_color),
        "\n")

    # Wrap long lines.
    width = displaysize(Base.stdout)[2] - 1
    if io.size > width
        line = String(take!(io))
#        id = task_local_debug().indent[end][2]
        pad = string(repeat(" ", DEBUG_BLANK_LINE_INDENT),
                     thread_offset, indent)
        if new_function != ""
            pad = string(pad, repeat(" ", length(new_function)), prefix)
        else
            pad *= "    "
        end
        if (width - length(pad)) < 16
            pad = ""
        end
        lines = string(wrap(line; width = width, subsequent_indent = pad),
                       "\n")
        debug_write(lines)
    else
        GC.@preserve io debug_write(pointer(io.data), io.size)
    end

# ╰╯ ╮
    if length(new_function) > DEBUG_MIN_INDENT
        io.size = 0
        io.ptr = 1
        print(ioc,
              bg_color,
              repeat(" ", DEBUG_BLANK_LINE_INDENT), thread_offset, indent,
              repeat(" ", DEBUG_MIN_INDENT),
              " ╭", repeat("─", length(new_function) - 1
                                - DEBUG_MIN_INDENT), "╯",
              "\e[0K", # Clear to end of line
              inv(bg_color),
              "\n")
        task_local.leader = String(take!(io))
    end
end # lock
end

"""
Write directly to STDERR.
"""
function debug_write(fd, p, l)
    while l > 0
        n = C.write(fd, p, l)
        if n > 0
            l -= n;
            p += n;
        elseif n == -1
            @assert errno() in (C.EAGAIN, C.EINTR)
        end
    end
    tcdrain(fd)
end
debug_write(p, l) = debug_write(Base.STDERR_NO, p, l)
debug_write(s::String) = GC.@preserve s debug_write(pointer(s), ncodeunits(s))

debug_time(t) = round(t - global_debug.t0; digits=4)

# Macro utilities.

"""
Filename: Line No. String
"""
function debug_lineno_str(source)
    f = basename(string(source.file))
    string(lpad(f,DEBUG_FILENAME_WIDTH), ":", rpad(source.line,4))
end

"""
String representation of function argument.
"""
function dbtiny(x; limit=16)
    s = dbshort(x)
    if length(s) > limit
        s = string(s[1:nextind(s,16)], "...")
    end
    return s
end

dbtiny(x::Cmd) = repr(x)

function dbshort(x)
    io = IOBuffer()
    dbshow(debug_io_context(io), x)
    String(take!(io))
end

"""
Function argument types (for arguments that have types).
"""
function argtypes(ex)
  f, args, kwargs = destructure_callex(ex)
  args = map(x->Base.is_expr(x, :(::)) ? (x.args...,) : missing, args)
  last.(skipmissing(args))
end

"""
Symbol for function argument.
"""
argsym(s::Symbol) = s
argsym(s::Expr) = s.head == :(::)        ? argsym(s.args[1]) :
                  s.head == :kw          ? argsym(s.args[1]) :
                  s.head == :...         ? argsym(s.args[1]) :
                  s.head == :parameters  ? missing           :
                                           " ?"

"""
Split body at end of "header" lines (e.g. `@nospecialize`).
"""
function body_split(body)
    line1(::LineNumberNode) = false
    line1(e::Expr) = e.args[1] != Symbol("@nospecialize")
    line1(::Any) = true
    i = findfirst(line1, body.args)
    while i > 2 && body.args[i-1] isa LineNumberNode
        i -= 1
    end
    head = body
    body = copy(head)
    deleteat!(head.args, i:length(head.args))
    deleteat!(body.args, 1:i-1)
    head, body
end

# Instrumentation macros.

"""
Like `Base.@lock`, but with logging when blocked waiting for lock.
"""
macro dblock(l, expr)
    if DEBUG_LEVEL < 0
        return esc(:(@lock $l $expr))
    end

    lineno = debug_lineno_str(__source__)

    quote
        l = $(esc(l))
        warn = islocked(l)
        warn && debug_print(1, $lineno, "🔒 Waiting $(string(l)) ⁉ ...")
        lock(l)
        warn && debug_print(1, $lineno, "🔓 Unlocked.")
        try
            $(esc(expr))
        finally
            unlock(l)
        end
    end
end

"""
Wrap the body of function `ex` wtih calls to `@debug_function_entry` and
`finally @debug_function_exit`.
(Do not use directly. Called by `@db`)
"""
macro debug_function(n::Int, ex::Expr, lineno::String)
    DEBUG_LEVEL < n && return esc(ex)

    @require ex.head == :function
    @require length(ex.args) == 2
    call, body = ex.args
    @require call.head in (:call, :where)
    @require body.head == :block
    if call.head != :call
        call = call.args[1]
    end
    @assert call.head == :call

    @info "@db Wrapping $call"

    # Split function Expr into parts.
    name = string(call.args[1])
    types = join(argtypes(call), ", ")
    if types != ""
        types = "($types)"
    end
    args = tuple((skipmissing(argsym(a) for a in call.args[2:end]))...)
    head, body = body_split(body)

    # Rewrite function body.
    body = quote

        _db_lineno::String = $lineno
        _db_fname::String = $name
        _db_level::Int = $n
        _db_returned::Bool = false
        _db_prefix::String = " │ "

        debug_print(
            _db_level,
            _db_lineno,
            string(_db_fname, " ┬", $types, " <- (", #┌─┐
                join(((dbtiny(_x) for _x in ($(args...),))...,), ", "),
            ")");
            new_function=_db_fname)

        debug_indent_push(_db_fname)

        _db_result = nothing
        try
            # Original function body.
            _db_result = $body

        catch _db_err
            _db_result = "⚠  $_db_err"
            rethrow(_db_err)
        finally
            if ! _db_returned
                debug_print(_db_level, _db_lineno,
                            _db_result == nothing ? "👍" : dbshort(_db_result);
                            function_lineno = _db_lineno,
                            function_return = _db_fname)
            end
            debug_indent_pop()
        end
        _db_result
    end
    append!(head.args, body.args)
    esc(ex)
end

"""
Print the return value before returning.
(Do not use directly. Called by `@db`)
"""
macro debug_return(n::Int, ex::Expr, message, lineno)
    DEBUG_LEVEL < n && return esc(ex)

    @require ex.head == :return
    v = gensym()
    message = something(message, :(dbshort($v)))
    esc(quote
        $v = $(ex.args[1])
        debug_print(_db_level, $lineno, $message; function_lineno = _db_lineno,
                                                  function_return = _db_fname)
        _db_returned = true
        return $v
    end)
end

const _db_prefix = " │ [ "

"""
Print a debug message.
(Do not use directly. Called by `@db`)
"""
macro debug_print(n::Int, ex, lineno)
    DEBUG_LEVEL < n && return :()
    esc(:(debug_print($n, $lineno, $ex; prefix=_db_prefix)))
end

"""
Print a debug message for `name` and `value`.
(Do not use directly. Called by `@db`)
"""
macro debug_show(n::Int, name, value, lineno)
    DEBUG_LEVEL < n && return :()
    esc(:(debug_print($n, $lineno, $name, $value; prefix=_db_prefix)))
end

"""
    @db [n] message

    @db [n] function name(args...)
        ...
        @db return x ["message"]
    end

`@db` prints debug messages.
The optional `n` parameter causes a message to be disabled
if `DEBUG_LEVEL < n`. Disabled messages have zero-cost.
The macro generates an empty expression `:()` for disabled messages.

When applied to a function definition, `@db` will print messages
for function entry and exit (both by `return` and by `throw`).

When applied to a return statement, `@db` will print the return
value (or an optional `message`) before returning.
"""
macro db(args...)

    db_level = args[1]
    if db_level isa Int
        args = args[2:end]
    else
        db_level =1
    end

    ex, message = (args..., nothing)

    lineno = debug_lineno_str(__source__)

    if ex isa Expr && ex.head == :return
        return esc(:(@debug_return $db_level $ex $message $lineno))
    elseif ex isa Expr && ex.head == :function
        return esc(:(@debug_function $db_level $ex $lineno))
    else
        if ex isa String || (ex isa Expr && ex.head == :string)
            return esc(:(@debug_print $db_level $ex $lineno))
        else
            return esc(Expr(
                :block,
                (:(@debug_show $db_level $(string(v, " = ")) $v $lineno)
                for v in args)...))
        end
    end

    @assert false
end

macro db_indent(n, f)
    DEBUG_LEVEL < n && return :()
    esc(:(debug_indent_push($f)))
end

macro db_unindent(n)
    DEBUG_LEVEL < n && return :()
    esc(:(debug_indent_pop()))
end

# Colors.

const debug_level_colors = [
    crayon"bold fg:202",
    crayon"bold fg:208",
    crayon"bold fg:220",
    crayon"bold fg:46",
    crayon"bold fg:51",
    crayon"bold fg:27"]

const debug_background_colors =
    Iterators.Stateful(
    Iterators.Cycle([
    crayon"bg:16",
    crayon"bg:18",
    crayon"bg:233",
    crayon"bg:234",
    crayon"bg:235",
    crayon"bg:236",
    crayon"bg:237",
    crayon"bg:238",
    crayon"bg:239",
    crayon"bg:240",
    crayon"bg:241",
    crayon"bg:242",
    crayon"bg:17",
    crayon"bg:18",
    crayon"bg:19",
    crayon"bg:20",
    crayon"bg:21",
    crayon"bg:22",
    crayon"bg:23",
    crayon"bg:24",
    crayon"bg:25",
    crayon"bg:58",
    crayon"bg:59",
    crayon"bg:60",]))

# Pretty Printing.

dbprint(io::IO, v::String) = print(io, v)
dbprint(io::IO, a) = print(io, a)

dbshow(io::IO, a) = show(io, a)

function dbprint(io::IO, a, args...)
    for a in (a, args...)
        if a isa String || a isa Crayon
            dbprint(io, a)
        else
            dbshow(io, a)
        end
    end
end

function dbshow(io::IO, v::Union{AbstractVector,AbstractSet})
    if get(io, :compact, false)
        print(IOContext(io, :typeinfo => typeof(v)), v)
    else
        print(io, v)
    end
end

function dbstring(args...)
    io = IOBuffer()
    ioc = debug_io_context(io)
    dbprint(ioc, args...)
    String(take!(io))
end

# Box Drawing.

"""
╔══╦══╗ ┌──┬──┐ ╭──╮ ┏━━┳━━┓    ▲     ■●◠ ◢◣ ░
║  ║  ║ │  │  │ │  │ ┃  ┃  ┃ ╌┈┉│┇┋╏ ▢□◯◡ ◥◤ ▒
╠══╬══╣ ├──┼──┤ │  │ ┣━━╋━━┫ ◀──┼──▶         ▓
║  ║  ║ │  │  │ │  │ ┃  ┃  ┃ ╍┅┄│┆┊╎  ▁▂▃▄▅▆▇█▉▊▋▌▍▎▏
╚══╩══╝ └──┴──┘ ╰──╯ ┗━━┻━━┛    ▼
"""

# Optional Exports.

module Debug

    for x in (Symbol("@db"),
              Symbol("@debug_print"),
              Symbol("@debug_function"),
              Symbol("@debug_return"),
              Symbol("@debug_show"),
              :dbshow,
              :dbtiny,
              :dbprint,
              :dbshort,
              :debug_print,
              :debug_indent_push,
              :debug_indent_pop,
              :_db_prefix)

       eval(:(import ..UnixIO:$x; export $x))
   end
end

# End of file: debug.jl


text:
#=
Created on Friday 4 December 2020
Last update: Tuesday 28 September 2021

@author: Bram De Jaegher
bram.de.jaegher@gmail.com
=#

module PlutoQuestions
  using Markdown
  using PlutoUI
  using Markdown: MD, Admonition

  export still_missing, keep_working, correct, not_defined, hint, fyi
  export Question, QuestionBlock, validate
  export NoDiff, Easy, Intermediate, Hard
  export @safe, @terminal

  include("styles.jl")
  include("admonition.jl")
  include("question.jl")
  include("macros.jl")
# Write your package code here.
end

max_stars_repo_name: Beramos/PlutoTeaching

max_issues_repo_name: Beramos/PlutoTeaching

max_forks_repo_name: Beramos/PlutoTeaching


text:
"""
Returns a set of all sub-expressions in an expression that look like \$some_expression
"""
function find_node_expressions(obj::Expr)
    node_expressions = Set()
    if isnodeexpression(obj)
        push!(node_expressions, obj)
    else
        for a in obj.args
            node_expressions = union(node_expressions, find_node_expressions(a))
        end
    end
    node_expressions
end

function find_node_expressions(x)
    # empty dict if x is not an Expr
    Set()
end

isnodeexpression(x) = false
function isnodeexpression(e::Expr)
    e.head == Symbol(:$) && length(e.args) == 1
end

"""
Replaces every subexpression that looks like a node expression with a substitute symbol stored in `exprdict`.
"""
function replace_node_expressions!(exp::Expr, exprdict)
    if isnodeexpression(exp)
        error("You can't @lift an expression that only consists of a single node.")
    else
        for (i, arg) in enumerate(exp.args)
            if isnodeexpression(arg)
                exp.args[i] = exprdict[arg]
            else
                replace_node_expressions!(arg, exprdict)
            end
        end
    end
    exp
end

replace_node_expressions!(x, exprdict) = begin end

"""
Replaces an expression with lift(argtuple -> expression, args...), where args
are all expressions inside the main one that begin with \$.

# Example:

x = Node(rand(100))
y = Node(rand(100))

## before
z = lift((x, y) -> x .+ y, x, y)

## after
z = @lift(\$x .+ \$y)

You can also use parentheses around an expression if that expression evaluates to a node.

```julia
nt = (x = Node(1), y = Node(2))
@lift(\$(nt.x) + \$(nt.y))
```
"""
macro lift(exp)

    node_expr_set = find_node_expressions(exp)

    if length(node_expr_set) == 0
        error("Did not find any expressions that looked like nodes.")
    end

    # store expressions with their substitute symbols, gensym them manually to be
    # able to escape the expression later
    node_expr_arg_dict = Dict(expr => gensym("arg$i") for (i, expr) in enumerate(node_expr_set))

    replace_node_expressions!(exp, node_expr_arg_dict)

    # keep an array for ordering
    node_expressions_array = collect(keys(node_expr_arg_dict))
    node_substitutes_array = [node_expr_arg_dict[expr] for expr in node_expressions_array]
    node_expressions_without_dollar = [n.args[1] for n in node_expressions_array]

    # the arguments to the lifted function
    argtuple = Expr(Symbol(:tuple), node_substitutes_array...)

    # the lifted function itself
    function_expression = Expr(Symbol(:->), argtuple, exp)

    # the full expression
    lift_expression = Expr(
        Symbol(:call),
        Symbol(:lift),
        esc(function_expression),
        esc.(node_expressions_without_dollar)...
    )

    lift_expression
end


text:
using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra

import OrdinaryDiffEq

@testset "schroedinger" begin

    Random.seed!(0)

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(op1::AbstractSuperOperator, op2::AbstractSuperOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))

    N = 3
    Ncutoff = 2
    tspan = [0.0:0.1:1.0;]

    Ω = [
        0 2 3
        2 0 1
        3 1 0
    ]

    ω = [1.0, 1.2, 1.5]

    basis_fock = FockBasis(Ncutoff)
    basis = tensor([basis_fock for i = 1:N]...)

    a = destroy(basis_fock)
    at = create(basis_fock)
    I = identityoperator(basis_fock)

    psi0 = tensor([coherentstate(basis_fock, i % Ncutoff) for i = 1:N]...)

    # Interaction picture
    Hrot = SparseOpType[]
    for i = 1:N, j = 1:N
        if i == j
            continue
        end
        h = embed(basis, [i, j], [a, Ω[i, j] * at])
        push!(Hrot, h)
    end
    Hrot = sum(Hrot)

    # Schroedinger picture
    function f(t, psi)
        H = SparseOpType[embed(basis, i, ω[i] * at * a) for i = 1:N]
        for i = 1:N, j = 1:N
            if i == j
                continue
            end
            h = embed(basis, [i, j], [a, exp(1im * (ω[i] - ω[j]) * t) * Ω[i, j] * at])
            push!(H, h)
        end
        sum(H)
    end

    tout, psi_rot_t = schroedinger(psi0, tspan, Hrot)
    tout, psi_t = schroedinger_dynamic(psi0, tspan, f)

    n_op = dense(at * a)
    for (i, t) in enumerate(tout)
        R = prod([embed(basis, i, exp(1im * ω[i] * t * n_op)) for i = 1:N])
        psi_rot = psi_rot_t[i]
        psi = psi_t[i]
        # @test abs(dagger(psi_rot)*R*psi) < 1e-5
        rho = dm(psi)
        rho_rot = dm(psi_rot)
        @test tracedistance(rho_rot, dense(R) * rho * dagger(dense(R))) < 1e-5
    end

    function fout(t, psi)
        deepcopy(psi)
    end
    t_fout, psi_fout = schroedinger(psi0, tspan, Hrot; fout = fout)
    @test t_fout == tout && psi_fout == psi_rot_t

    # test integration of propagator using 2 level system
    basis = SpinBasis(1 // 2)
    su = spinup(basis)
    u0 = identityoperator(basis)
    sx = sigmax(basis)
    sz = sigmaz(basis)

    # for the time dependent equation
    f2(t, psi) = sx * π
    tspan = [0:1.0;]
    t, u = schroedinger(u0, tspan, π * sx)

    # I think the tolerance on the differential equation is 1e-6, we expect the operator to be essentially the identity
    @test abs(expect(sz, u[end] * su)) - abs(expect(sz, u0 * su)) < 1e-6

    t, u = schroedinger_dynamic(u0, tspan, f2)
    @test abs(expect(sz, u[end] * su)) - abs(expect(sz, u0 * su)) < 1e-6

    mode1 = Mode(0.2, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 2, Energy)
    mol2 = Molecule([mode1], 2, Energy)
    agg = Aggregate([mol1, mol2])
    aggInds = getIndices(agg)
    aggIndsLen = length(aggInds)
    basis = GenericBasis([aggIndsLen])
    FCFact = getFranckCondonFactors(agg, aggInds)
    Ham = getAggHamiltonian(agg, aggInds, FCFact)

    ket0 = randstate(basis)
    ket0 = Ket(basis)
    for i = 1:aggIndsLen
        ket0.data[i] = (1.0 / aggIndsLen)^0.5
    end
    rho0 = dm(ket0)
    tspan = [0.0:0.1:1.0;]

    T, rho_t = schroedinger(
        ket0,
        tspan,
        Ham;
        reltol = 1e-10,
        abstol = 1e-10,
        alg = OrdinaryDiffEq.Tsit5(),
    )
    for t_i = 1:length(tspan)
        U_op = evolutionOperator(Ham, tspan[t_i])
        ket = U_op * ket0
        @test 1e-8 > D(ket, rho_t[t_i])
        # println(t_i, " ", D(ket.data, rho_t[t_i].data))
    end

end # testset


text:
using Distributed
if nprocs() == 1
    addprocs(3)
end

function surface_extraction(iterator)
  localvar = 2.5
  block_func = function (elm)
    # hypothetical (expensive) operation with elm
    # (in my code this is solving a block)
    localvar * elm
  end

  pmap(block_func, iterator)
end

surface_extraction_mock(1:10)


text:
using Documenter, TrackedDistributions

makedocs(;
    modules=[TrackedDistributions],
    format=Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    pages=[
        "Home" => "index.md",
        "API" => "api.md",
    ],
    repo="https://gitlab.invenia.ca/invenia/TrackedDistributions.jl/blob/{commit}{path}#L{line}",
    sitename="TrackedDistributions.jl",
    authors="Alex Robson",
    assets=[
        "assets/invenia.css",
        "assets/logo.png",
    ],
    strict = true,
    checkdocs = :exports,
)


text:
print_all(Lexers.RLexer, "r")

tokentest(
    Lexers.RLexer,
    "# Comments...",
    COMMENT_SINGLE => "# Comments...",
)

tokentest(
    Lexers.RLexer,
    "   # Comments...",
    TEXT => "   ",
    COMMENT_SINGLE => "# Comments...",
)

tokentest(
    Lexers.RLexer,
    "x <- 1",
    NAME => "x",
    TEXT => " ",
    OPERATOR => "<-",
    TEXT => " ",
    NUMBER => "1",
)

tokentest(
    Lexers.RLexer,
    "func(x, y)",
    NAME => "func",
    PUNCTUATION => "(",
    NAME => "x",
    PUNCTUATION => ",",
    TEXT => " ",
    NAME => "y",
    PUNCTUATION => ")",
)

tokentest(
    Lexers.RLexer,
    "c(1, 2)",
    KEYWORD_TYPE => "c",
    PUNCTUATION => "(",
    NUMBER => "1",
    PUNCTUATION => ",",
    TEXT => " ",
    NUMBER => "2",
    PUNCTUATION => ")",
)

tokentest(
    Lexers.RLexer,
    "-Inf + 1e9",
    OPERATOR => "-",
    KEYWORD_CONSTANT => "Inf",
    TEXT => " ",
    OPERATOR => "+",
    TEXT => " ",
    NUMBER => "1e9",
)

tokentest(
    Lexers.RLexer,
    "61.4e-12 * 1.2e2",
    NUMBER => "61.4e-12",
    TEXT => " ",
    OPERATOR => "*",
    TEXT => " ",
    NUMBER => "1.2e2",
)

tokentest(
    Lexers.RLexer,
    "4 %% 3",
    NUMBER => "4",
    TEXT => " ",
    OPERATOR => "%%",
    TEXT => " ",
    NUMBER => "3",
)

tokentest(
    Lexers.RLexer,
    "A %*% t(A)",
    NAME => "A",
    TEXT => " ",
    OPERATOR => "%*%",
    TEXT => " ",
    NAME => "t",
    PUNCTUATION => "(",
    NAME => "A",
    PUNCTUATION => ")",
)

tokentest(
    Lexers.RLexer,
    "FALSE && TRUE",
    KEYWORD_CONSTANT => "FALSE",
    TEXT => " ",
    OPERATOR => "&&",
    TEXT => " ",
    KEYWORD_CONSTANT => "TRUE",
)

tokentest(
    Lexers.RLexer,
    "' \" ' \" ' \" ' \\' ' \" \\\" \"",
    STRING => "' \" '",
    TEXT => " ",
    STRING => "\" ' \"",
    TEXT => " ",
    STRING => "' \\' '",
    TEXT => " ",
    STRING => "\" \\\" \"",
)

tokentest(
    Lexers.RLexer,
    "while (TRUE) {",
    KEYWORD_RESERVED => "while",
    TEXT => " ",
    PUNCTUATION => "(",
    KEYWORD_CONSTANT => "TRUE",
    PUNCTUATION => ")",
    TEXT => " ",
    PUNCTUATION => "{",
)

tokentest(
    Lexers.RLexer,
    "for (x in 1:10) {",
    KEYWORD_RESERVED => "for",
    TEXT => " ",
    PUNCTUATION => "(",
    NAME => "x",
    TEXT => " ",
    KEYWORD_RESERVED => "in",
    TEXT => " ",
    NUMBER => "1",
    OPERATOR => ":",
    NUMBER => "10",
    PUNCTUATION => ")",
    TEXT => " ",
    PUNCTUATION => "{",
)


text:
export build_dispatch
function build_dispatch(range, kernel_info, karg_memory, signal)
    debug_print("build_dispatch: Building $range KernelDispatchPacket for $(string(kernel_info.func))($(join(kernel_info.argtypes,',')))")
    dispatch_packet = KernelDispatchPacket(
        kernel_info.kernel_object,
        range)
    dispatch_packet.header.acquire_fence_scope = HSA.HSA_FENCE_SCOPE_SYSTEM
    dispatch_packet.header.release_fence_scope = HSA.HSA_FENCE_SCOPE_SYSTEM
    dispatch_packet.completion_signal = signal
    dispatch_packet.kernarg_address = karg_memory.ptr
    dispatch_packet.private_segment_size = kernel_info.private_size
    dispatch_packet.group_segment_size = kernel_info.group_size

    return dispatch_packet
end


text:
using Symbolics, Test

@variables t a x(t) y(t)
expr = [a*x - x*y,-3y + x*y]
@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget()) ==
    """
    real[] diffeqf(real t,real[] internal_var___u,real[] internal_var___p,real[] x_r,int[] x_i) {
      real internal_var___du[2];
      internal_var___du[1] = internal_var___p[1] * internal_var___u[1] + -1 * internal_var___u[1] * internal_var___u[2];
      internal_var___du[2] = internal_var___u[1] * internal_var___u[2] + -3 * internal_var___u[2];
      return internal_var___du;
    }
    """

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t]) ==
  """
  #include <math.h>
  void diffeqf(double* internal_var___du, double* internal_var___u, double* internal_var___p, double t) {
    internal_var___du[0] = internal_var___p[0] * internal_var___u[0] + -1 * internal_var___u[0] * internal_var___u[1];
    internal_var___du[1] = internal_var___u[0] * internal_var___u[1] + -3 * internal_var___u[1];
  }
  """

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget()) ==
  """
  diffeqf = @(t,internal_var___u) [
    internal_var___p(1) * internal_var___u(1) + -1 * internal_var___u(1) * internal_var___u(2);
    internal_var___u(1) * internal_var___u(2) + -3 * internal_var___u(2);
  ];
  """

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t]) ==
    Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                   lhsname=:internal_var___du,
                                   rhsnames=[:internal_var___u,:internal_var___p,:t])

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget()) ==
      Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget()) ==
      Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t]) ==
      Symbolics.build_function(expr,vcat(x,y),[a],t,target = Symbolics.CTarget(),
                                     lhsname=:internal_var___du,
                                     rhsnames=[:internal_var___u,:internal_var___p,:t])

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.StanTarget()) ==
      Symbolics.build_function(expr,vcat(x,y),[a],t,target = Symbolics.StanTarget())

@test Symbolics.build_function(expr,[x,y],[a],t,target = Symbolics.MATLABTarget()) ==
      Symbolics.build_function(expr,vcat(x,y),[a],t,target = Symbolics.MATLABTarget())

# Matrix CTarget test
let
    @variables x[1:4] y[1:4] z[1:4]
    expression = hcat(x,y,z) # Results in [x1 y1 z1; x2 y2 z2; ...]
    variables  = vcat(x,y,z) # Results in [x1, x2, x3, x4, y1, ...]
    cfunc      = build_function(expression, variables; target = Symbolics.CTarget(), expression = Val{true})

    # Generated function should be out[0] = in[0], out[1] = in[1], out[2] = in[2], etc.
    @test cfunc == "#include <math.h>\nvoid diffeqf(double* du, double* RHS1) {\n  du[0] = RHS1[0];\n  du[1] = RHS1[1];\n  du[2] = RHS1[2];\n  du[3] = RHS1[3];\n  du[4] = RHS1[4];\n  du[5] = RHS1[5];\n  du[6] = RHS1[6];\n  du[7] = RHS1[7];\n  du[8] = RHS1[8];\n  du[9] = RHS1[9];\n  du[10] = RHS1[10];\n  du[11] = RHS1[11];\n}\n"
end

# Scalar CTarget test
let
    @variables x y t
    expression = x + y + t
    cfunc = build_function(expression, [x], [y], t; target = Symbolics.CTarget(), expression = Val{true})

    # Generated function should be out[0] = in1[0] + in2[0] + in3
    @test cfunc == "#include <math.h>\nvoid diffeqf(double* du, double* RHS1, double* RHS2, double RHS3) {\n  du[0] = RHS3 + RHS1[0] + RHS2[0];\n}\n"
end

# Scalar CTarget test with scalar multiplication and powers
let
  @variables x y a t
  expression = x^2 + y^-1 + sin(a)^3.5 + 2t
  cfunc = build_function(expression, [x, y], [a], t; target = Symbolics.CTarget(), expression = Val{true})

  # Generated function should avoid scalar multiplication of the form "4t" (currently done by adding another "* 1") and other invalid C syntax
  @test cfunc == "#include <math.h>\nvoid diffeqf(double* du, double* RHS1, double* RHS2, double RHS3) {\n  du[0] = 2 * RHS3 * 1 + pow(RHS1[0], 2) + 1.0 / RHS1[1] + pow(sin(RHS2[0]), 3.5);\n}\n"
end

# Matrix StanTarget test
let
    @variables x[1:4] y[1:4] z[1:4]
    expression = hcat(x,y,z) # Results in [x1 y1 z1; x2 y2 z2; ...]
    variables  = vcat(x,y,z) # Results in [x1, x2, x3, x4, y1, ...]
    sfunc      = build_function(expression, vcat(x,y), z, []; target = Symbolics.StanTarget(), expression = Val{true})

    # Generated function should be out[1] = v[1], out[2] = v[2], ..., out[9] = p[9], etc.
    @test sfunc == """
    real[] diffeqf(real Any[],real[] internal_var___u,real[] internal_var___p,real[] x_r,int[] x_i) {
      real internal_var___du[12];
      internal_var___du[1] = internal_var___u[1];
      internal_var___du[2] = internal_var___u[2];
      internal_var___du[3] = internal_var___u[3];
      internal_var___du[4] = internal_var___u[4];
      internal_var___du[5] = internal_var___u[5];
      internal_var___du[6] = internal_var___u[6];
      internal_var___du[7] = internal_var___u[7];
      internal_var___du[8] = internal_var___u[8];
      internal_var___du[9] = internal_var___p[1];
      internal_var___du[10] = internal_var___p[2];
      internal_var___du[11] = internal_var___p[3];
      internal_var___du[12] = internal_var___p[4];
      return internal_var___du;
    }
    """
end

# Scalar StanTarget test
let
    @variables t x(t) y(t) z(t)
    expression = x + y + z
    sfunc = build_function(expression, vcat(x,y), [z], t; target = Symbolics.StanTarget(), expression = Val{true})

    # Generated function should be out[0] = v[0] + p[0] + t[0]
    @test sfunc == """
    real[] diffeqf(real t,real[] internal_var___u,real[] internal_var___p,real[] x_r,int[] x_i) {
      real internal_var___du[1];
      internal_var___du[1] = internal_var___u[1] + internal_var___u[2] + internal_var___p[1];
      return internal_var___du;
    }
    """
end

# Matrix MATLABTarget test
let
    @variables x[1:4] y[1:4] z[1:4]
    expression = hcat(x,y,z) # Results in [x1 y1 z1; x2 y2 z2; ...]
    variables  = vcat(x,y,z) # Results in [x1, x2, x3, x4, y1, ...]
    mfunc      = build_function(expression, vcat(x,y,z); target = Symbolics.MATLABTarget(), expression = Val{true})

    # Generated function should be of the same form as expression
    @test mfunc == "diffeqf = @(t,internal_var___u) [\n  internal_var___u(1), internal_var___u(5), internal_var___u(9);\n  internal_var___u(2), internal_var___u(6), internal_var___u(10);\n  internal_var___u(3), internal_var___u(7), internal_var___u(11);\n  internal_var___u(4), internal_var___u(8), internal_var___u(12);\n];\n"
end

# Scalar MATLABTarget test
let
    @variables x y z
    expression = x + y + z
    mfunc = build_function(expression, vcat(x,y,z); target = Symbolics.MATLABTarget(), expression = Val{true})

    # Generated function should be out[0] = v[0] + p[0] + t[0]
    @test mfunc == "diffeqf = @(t,internal_var___u) [\n  internal_var___u(1) + internal_var___u(2) + internal_var___u(3);\n];\n"
end


text:
using HetaSimulator, Plots
ENV["GKSwstype"] = "100"

model = load_jlmodel("../dist/julia/model.jl")

res = Scenario(model; tspan = (0,120)) |> sim

plotd = plot(res; vars = [:A, :B])
savefig(plotd, "julia-plot-1.png")

plotd = plot(res; vars = [:r1])
savefig(plotd, "julia-plot-2.png")

max_stars_repo_name: insysbio/heta-case-mini

max_issues_repo_name: insysbio/heta-case-mini

max_forks_repo_name: insysbio/heta-case-mini


text:
using StatsPlots

function plot_predictive(obstimes, X_true, Y_obs, chain)
    Xp = []
    for i in 1:10
        pred = predict(Y_obs, chain)
        push!(Xp, pred[2:end,6])
    end

    p = plot(obstimes, Xp; legend=false, color=:red, alpha=0.8)
    plot!(p, obstimes, X_true, color=:black, lw=3)
    scatter!(p, obstimes, Y_obs)
    return p
end

function plot_intermediate_samples(
    intermediate_samples,
    param_name,
    true_param,
    betas
)
    inter_samples = permutedims(hcat(intermediate_samples...), [2,1])

    num_samples = size(inter_samples, 1)
    p = plot(title="Intermediate samples for $(string(param_name))", legend=false)
    plot!(p, [true_param], color=:red, seriestype="vline")
    for i in 1:size(inter_samples, 2)
        ix = i > 10 ? (i-9)*10 : i
        beta = betas[ix]
		scatter!(
			p,
			map(x->x[param_name], inter_samples[:,i]),
			ones(num_samples)*beta,
			alpha=0.2,
            color=:black,
            xlims=[-0.1,0.1],
            ylims=[0,1.0]
		)
    end
    return p
end

function plot_intermediate_weights(intermediate_weights, betas; ixs=nothing)
    inter_weights = permutedims(hcat(intermediate_weights...), [2,1])

    num_samples = size(inter_weights, 1)
    p = plot(
        title="Intermediate weights",
        xlabel="betas",
        ylabel="log weight",
        legend=false
    )

    if isnothing(ixs)
        ixs = 1:size(inter_weights, 2)
    end
    @show size(inter_weights)
	for i in ixs
        ix = i > 10 ? (i-9)*10 : i
        beta = betas[ix]
		scatter!(
			p,
			ones(num_samples)*beta,
			inter_weights[:,i],
			alpha=0.2,
            color=:black
		)
    end
    return p
end

function plot_intermediate_acceptance(intermediate_acceptance, betas; ixs=nothing)
    inter_acceptance = permutedims(hcat(intermediate_acceptance...), [2,1])
    num_samples = size(inter_acceptance, 1)

    inter_means = mean(inter_acceptance, dims=1)[1,:]

    beta_ix = [i > 10 ? (i-9)*10 : i for i in 1:size(inter_acceptance, 2)]
    p = plot(
        beta_ix,
        inter_means,
        fillalpha=0.4,
        title="Intermediate Acceptance Rate",
        xlabel="betas",
        ylabel="acceptance rate",
        legend=false,
        color=:red,
        lw=2
    )

    if isnothing(ixs)
        ixs = 1:size(inter_acceptance, 2)
    end
	for i in ixs
        ix = i > 10 ? (i-9)*10 : i
		scatter!(
			p,
			ones(num_samples)*ix,
			inter_acceptance[:,i],
			alpha=0.2,
            color=:black
		)
    end
    return p
end

function plot_intermediate_ess(intermediate_weights, betas; ixs=nothing)
    inter_weights = permutedims(hcat(intermediate_weights...), [2,1])

    ess = exp.(
        2 * logsumexp(inter_weights; dims=1) - logsumexp(2 * inter_weights; dims=1)
    )[1,:]

    if isnothing(ixs)
        ixs = 1:size(inter_weights, 2)
    end
    dist_ixs = [i > 10 ? (i-9)*10 : i for i in ixs]
    betas = [betas[i] for i in dist_ixs]
    ess = ess[ixs]
    return plot(
        dist_ixs,
        ess,
        title="ESS for intermediate distributions",
        xlabel="beta index",
        ylabel="ESS",
        lw=2,
        color=:black,
        legend=false
    )
end

function plot_joint_dist(chain, beta_true, i_0_true)
    betas = get(chain, :β)[:β]
    i_0s = get(chain, :i₀)[:i₀]
    p = scatter(
        betas, i_0s, xlabel="β", ylabel="i₀", label="AnIS samples", color=:black
    )
    scatter!(
        p, [beta_true], [i_0_true], label="True Params", color=:red, marker=:x
    )
    return p
end

function plot_sir_samples(chain; color=:blue)
    betas = get(chain, :β)[:β]
    i_0s = get(chain, :i₀)[:i₀]
    p = scatter(
        betas,
        i_0s,
        xlabel=L"\beta",
        ylabel=L"i_0",
        alpha=0.5,
        legend=false,
        xlims=[0,2.3],
        xticks=[0.0, 0.5, 1.0, 1.5, 2.0],
        ylims=[0,60],
        fontfamily="serif-roman",
        framestyle=:semi,
        grid=false,
        color=color,
        thickness_scaling=1.5,
        markersize=2.5
    )
    return p
end

function plot_is_samples(samples, log_weights)
    return scatter(
        samples,
        log_weights,
        markersize=3,
        markeralpha=0.4,
        legend=false,
        xlabel="beta",
        ylabel="log weight",
        xlims=[0,2],
        ylims=[-40,0]
    )
end

function get_median_and_quantiles(errors)
    # errors: num_samples x num_chains
    qs = [quantile(errors[i,:], [0.25, 0.75]) for i in 1:size(errors,1)]
    qs = hcat(qs...)
    medians = median(errors, dims=2)[:,1]
    # medians: num_samples
    # qs: 2 x num_samples
    return medians, qs
end

function convergence_plot(
    problem_type,
    Z1_chn,
    Z2_chn,
    f,
    ground_truth;
    thinning=1,
    num_xlogspaced=nothing,
    Z1_negative_chn=nothing,
    mean_and_std=false
)
    # Shape: num_samples x num_chains
    Z1_ws = exp.(Array(Z1_chn[:log_weight]))
    Z2_ws = exp.(Array(Z2_chn[:log_weight]))

    num_samples = size(Z1_ws, 1)
    #resampled_ixs = StatsBase.sample(1:num_samples, pweights(Z2_ws), num_samples)
    #Z2_ws = Z2_ws[resampled_ixs]

    # NOTE: Usually we have to divide each estimate by the number of samples used
    # but because those "normalisations" cancel in our downstream calculations we
    # avoid them.
    Z1s_unnormalised = cumsum(Z1_ws, dims=1)
    Z2s_unnormalised = cumsum(Z2_ws, dims=1)

    if isnothing(Z1_negative_chn)
        tabi_ests = Z1s_unnormalised ./ Z2s_unnormalised
    else
        Z1_negative_ws = exp.(Array(Z1_negative_chn[:log_weight]))
        Z1_negatives_unnormalised = cumsum(Z1_negative_ws, dims=1)
        tabi_ests = (Z1s_unnormalised .- Z1_negatives_unnormalised) ./ Z2s_unnormalised
    end

    if problem_type == "sir"
        i₀ = Array(Z2_chn[:i₀])
        β = Array(Z2_chn[:β])
        anis_ests = Array{Float64,2}(undef, num_samples, size(Z2_ws, 2))
        for ix in CartesianIndices(Z2_ws)
            anis_ests[ix] = Z2_ws[ix] * f(i₀[ix], β[ix])
        end
        #i₀ = Array(Z2_chn[:i₀])[resampled_ixs]
        #β = Array(Z2_chn[:β])[resampled_ixs]
        #anis_ests = cumsum(map(zip(ones(num_samples), i₀, β)) do (w, i, b)
        #    w * f(i, b)
        #end)
        #anis_ests = anis_ests ./ (1:num_samples)
    elseif problem_type == "post_pred"
        xs = Z2_chn[:,namesingroup(Z2_chn, :x),:].value.data
        anis_ests = Array{Float64,2}(undef, num_samples, size(Z2_ws, 2))
        for ix in CartesianIndices(Z2_ws)
            sample_ix, chain_ix = ix[1], ix[2]
            anis_ests[ix] = Z2_ws[ix] * f(xs[sample_ix,:,chain_ix])
        end
    elseif problem_type == "seventh_moment"
        xs = Array(Z2_chn[:x])
        anis_ests = Array{Float64,2}(undef, num_samples, size(Z2_ws, 2))
        for ix in CartesianIndices(Z2_ws)
            sample_ix, chain_ix = ix[1], ix[2]
            anis_ests[ix] = Z2_ws[ix] * f(xs[sample_ix,chain_ix])
        end
    else
        error("Unknown problem type")
    end
    anis_ests = cumsum(anis_ests, dims=1)
    anis_ests = anis_ests ./ Z2s_unnormalised

    tabi_ests = tabi_ests[1:thinning:end,:]
    anis_ests = anis_ests[1:thinning:end,:]

    @show tabi_ests[end,:]
    @show anis_ests[end,:]
    tabi_rserror = relative_squared_error.(ground_truth, tabi_ests)
    anis_rserror = relative_squared_error.(ground_truth, anis_ests)
    @show tabi_rserror[end,:]
    @show anis_rserror[end,:]

    #ixs = 50:num_samples
    ixs = 1:thinning:num_samples
    taanis_ixs = if isnothing(Z1_negative_chn)
        2 * ixs # TAAnIS uses twice the number of samples
    else
        3 * ixs
    end

    data = [
        ("TAAnIS", taanis_ixs, tabi_rserror),
        ("AnIS", ixs, anis_rserror)
    ]
    error_plot = plot(
        yscale=:log10,
        xscale=:log10,
        xlabel="Number of Samples",
        ylabel="RSE",
        legend=false,
        thickness_scaling=1.5,
        fontfamily="serif-roman",
        framestyle=:semi,
        grid=false,
        xlims=(10,10^4),
        # ylims=(10^(-6.5),1),
        ylims=(10^(-8),1),
        xticks=[10, 100, 1_000, 10_000],
        yticks=[1, 0.01, 0.0001, 0.000001]
    )
    for (label, indexes, rserrors) in data
        if mean_and_std
            mds = vec(mean(rserrors, dims=2))
            qs = vec(std(rserrors, dims=2))
        else
            mds, qs = get_median_and_quantiles(rserrors)
        end
        means = vec(mean(rserrors, dims=2))
        mds, qs = get_median_and_quantiles(rserrors)

        if !isnothing(num_xlogspaced)
            log_ixs = unique(round.(Int, geomspace(1, num_samples, num_xlogspaced)))
            mds = mds[log_ixs]
            # qs = if mean_and_std
            #     qs[log_ixs]
            # else
            #     qs[:,log_ixs]
            # end
            qs = qs[:,log_ixs]
            means = means[log_ixs]
            indexes = Array(indexes)[log_ixs]
        end

        # Error bars differ when we use quantiles instead of standard deviation.
        # ribbon = if mean_and_std
        #     qs
        # else
        #     [(mds.-qs[1,:]), (qs[2,:]-mds)]
        # end
        ribbon = [(mds.-qs[1,:]), (qs[2,:]-mds)]
        plot!(
            error_plot,
            indexes,
            mds,#mds[ixs],
            ribbon=ribbon,
            label=label
        )
    end
    return error_plot, nothing

    ground_truth_plot = plot(
        taanis_ixs,
        tabi_ests[ixs],
        label="TAAnIS"
    )
    plot!(
        ground_truth_plot,
        ixs,
        anis_ests[ixs],
        label="AnIS"
    )
    plot!(
        ground_truth_plot,
        [1, 2*num_samples],
        repeat([ground_truth], 2),
        label="GT"
    )
    return error_plot, ground_truth_plot
end

function ess_batch(log_weights; thinning=1)
    # Input: (num_samples, num_chains)
    # Output: (num_samples/thinning, num_chains)

    num_samples = size(log_weights, 1)
    ess = Array{Float64}(undef, Int(num_samples/thinning), size(log_weights,2))

    ess_ix = 1
    for ix in 1:thinning:num_samples
        denominator = logsumexp(2 * log_weights[1:ix,:]; dims=1)
        numerator = 2 * logsumexp(log_weights[1:ix,:]; dims=1)
        ess[ess_ix,:] = exp.(numerator .- denominator)
        ess_ix += 1
    end
    return ess
end

function make_ess_plot(
    Z1_lws,
    Z2_lws,
    anis_lws_retargeted;
    thinning=1,
    anis_lws=nothing,
    yaxisscale=:log10,
    take_min=true,
    num_xlogspaced=nothing,
    legend=false
)
    num_samples = size(Z1_lws, 1)

    Z1_ess_cumsum = ess_batch(Z1_lws; thinning=thinning)
    Z2_ess_cumsum = ess_batch(Z2_lws; thinning=thinning)

    if isnothing(anis_lws)
        anis_lws = Z2_lws
    end
    anis_ess_cumsum = ess_batch(anis_lws; thinning=thinning)
    anis_ess_retargeted_cumsum = ess_batch(anis_lws_retargeted; thinning=thinning)

    ixs = Array(1:thinning:num_samples)
    # taanis_mds, taanis_qs = get_median_and_quantiles(taanis_min_ess)

    # ess_plot = plot(
    #     ixs*2,
    #     taanis_mds,
    #     ribbon=[(taanis_mds.-taanis_qs[1,:]), (taanis_qs[2,:].-taanis_mds)],
    #     xscale=:log10,
    #     yscale=yaxisscale,
    #     label="TAAnIS",
    #     legend=false,
    #     xlabel="Number of Samples",
    #     ylabel="ESS",
    #     xlims=(10,10^4),
    #     thickness_scaling=1.7
    # )
    # anis_ixs = Array(1:thinning:num_samples)
    # anis_mds, anis_qs = get_median_and_quantiles(anis_min_ess)
    # plot!(
    #     ess_plot,
    #     anis_ixs,
    #     anis_mds,
    #     ribbon=[(anis_mds.-anis_qs[1,:]), (anis_qs[2,:].-anis_mds)],
    #     label="AnIS"
    # )

    taanis_color = palette(:default)[1]
    anis_color = palette(:default)[2]
    ess_quantities = if take_min
        taanis_min_ess = min.(Z1_ess_cumsum, Z2_ess_cumsum)
        anis_min_ess = min.(anis_ess_retargeted_cumsum, anis_ess_cumsum)
        [
            ("TAAnIS", ixs*2, get_median_and_quantiles(taanis_min_ess), taanis_color, :solid),
            ("AnIS", ixs, get_median_and_quantiles(anis_min_ess), anis_color, :solid)
        ]
    else
        [
            ("Z1 (TAAnIS)", ixs*2, get_median_and_quantiles(Z1_ess_cumsum), taanis_color, :dash),
            ("Z2 (TAAnIS)", ixs*2, get_median_and_quantiles(Z2_ess_cumsum), taanis_color, :solid),
            ("AnIS", ixs, get_median_and_quantiles(anis_ess_cumsum), anis_color, :solid),
            ("AnIS retargeted", ixs, get_median_and_quantiles(anis_ess_retargeted_cumsum), anis_color, :dash)
        ]
    end

    name1, indices1, (medians1, quantiles1), cl1, ls1 = ess_quantities[1]
    if !isnothing(num_xlogspaced)
        log_ixs = unique(round.(Int, geomspace(1, num_samples, num_xlogspaced)))
        indices1 = indices1[log_ixs]
        medians1 = medians1[log_ixs]
        quantiles1 = quantiles1[:,log_ixs]
    end
    ess_plot = plot(
        indices1,
        medians1,
        ribbon=[(medians1.-quantiles1[1,:]), (quantiles1[2,:].-medians1)],
        xscale=:log10,
        yscale=yaxisscale,
        label=name1,
        color=cl1,
        linestyle=ls1,
        legend=legend,
        xlabel="Number of Samples",
        ylabel="ESS",
        xlims=(10,10^4),
        thickness_scaling=1.5,
        fontfamily="serif-roman",
        framestyle=:semi,
        grid=false
    )
    for (name, indices, (medians, quantiles), cl, ls) in ess_quantities[2:end]
        if !isnothing(num_xlogspaced)
            log_ixs = unique(round.(Int, geomspace(1, num_samples, num_xlogspaced)))
            indices = indices[log_ixs]
            medians = medians[log_ixs]
            quantiles = quantiles[:,log_ixs]
        end
        plot!(
            ess_plot,
            indices,
            medians,
            ribbon=[(medians.-quantiles[1,:]), (quantiles[2,:].-medians)],
            color=cl,
            linestyle=ls,
            label=name
        )
    end
    return ess_plot
end


text:
"""
    find_blocks(tokens, templates)

Given a list of tokens and a dictionary of block templates, find all blocks
matching templates. The blocks are sorted by order of appearance and inner
blocks are weeded out.
"""
function find_blocks(
            tokens::SubVector{Token};
            is_md::Bool=true
            )::Vector{Block} where T <: Dict{Symbol, BlockTemplate}

    blocks = Block[]
    isempty(tokens) && return blocks
    is_active = ones(Bool, length(tokens))

    # ------------------------------------------------------------------------
    if is_md
        # ======== PASS 0 ==============
        # raw blocks (??? ... ???)
        _find_blocks!(blocks, tokens, MD_PASS0, is_active)

        # ======== PASS 1 ==============
        # basically all container blocks
        # comment, raw html, raw latex, def blocks, code blocks
        # math blocks, div blocks, autolink, cu_brackets, h* blocks
        # and lxbegin/end
        _find_blocks!(blocks, tokens, MD_PASS1_TEMPLATES, is_active,
                      process_linereturn=true)

        sort!(blocks, by=from)

        # At this point we have the cu_brackets and begin/end in blocks
        # Form the begin...end environments and deactivate all tokens within
        _find_env_blocks!(blocks, tokens, is_active)

        # ======== PASS 2 ==============
        # brackets which may form a link: brackets and sq_brackets
        dt = _find_blocks!(blocks, tokens, MD_PASS2_TEMPLATES, is_active)
        form_links!(blocks)
        # here there may be brackets that are not part of links which
        # should have their content re-inspected
        @inbounds for b in filter(b_ -> b_.name == :BRACKETS, blocks)
            fromb  = from(b)
            tob    = to(b)
            retoks = Token[]
            for i in dt
                toki = tokens[i]
                fi   = from(toki)
                ti   = to(toki)
                # is the token in the scope ?
                fromb < fi && ti < tob && push!(retoks, toki)
            end
            # recurse
            append!(blocks, find_blocks(subv(retoks), is_md=true))
        end
        # discard leftover bracket blocks
        filter!(b -> b.name != :BRACKETS, blocks)

        # ======== PASS 3 ==============
        # remaining stuff e.g. emphasis tokens, lxnew* etc
        _find_blocks!(blocks, tokens, MD_PASS3_TEMPLATES, is_active)

    # ------------------------------------------------------------------------
    else
        # for HTML we barely do anything, a single pass is plenty enough
        _find_blocks!(blocks, tokens, HTML_TEMPLATES, is_active)
    end

    # remove blocks inside larger blocks (recursion)
    sort!(blocks, by=from)
    remove_inner!(blocks)

    # forming of double braces is done here to avoid clash with latex curly braces
    is_md && form_dbb!(blocks)

    return blocks
end

@inline find_blocks(t::Vector{Token}, a...) = find_blocks(subv(t), a...)

function _find_blocks!(
            blocks::Vector{Block},
            tokens::SubVector{Token},
            templates::Dict{Symbol, BlockTemplate},
            is_active::Vector{Bool}=ones(Bool, length(tokens));
            process_linereturn::Bool=false
            )::Vector{Int}
    #
    # keep track of what was deactivated, this is useful for md parsing
    # when discarding BRACKET tokens and re-enabling the tokens inside them;
    # only the tokens deactivated by it should be re-enabled.
    # so for instance:
    #   (abc _@@d *g* @@_ ef) --> first pass will deactivate `*`
    #   --> we should only re-enable `_`.
    #
    deactivated_tokens = Int[]

    isempty(templates) && return deactivated_tokens
    template_keys = keys(templates)
    n_tokens = length(tokens)
    @inbounds for i in eachindex(tokens)

        is_active[i] || continue
        opening = tokens[i].name

        if process_linereturn && opening in (:SOS, :LINE_RETURN)
            process_line_return!(blocks, tokens, i)
            continue
        elseif opening ∉ template_keys
            continue
        end

        template = templates[opening]
        closing  = template.closing
        nesting  = template.nesting

        if closing === NO_CLOSING
            push!(blocks, TokenBlock(tokens[i]))
            continue
        end

        # Try to find the closing token
        closing_index = -1
        open_depth    = 1
        for j in i+1:n_tokens
            # the tokens ahead might be inactive due to first pass
            is_active[j] || continue
            candidate = tokens[j].name
            # has to happen before opener to avoid ambiguity in emphasis tokens
            if candidate in closing
                open_depth -= 1
            elseif candidate == opening && nesting
                open_depth += 1
            end
            if open_depth == 0
                closing_index = j
                break
            end
        end

        if (closing_index == -1)
            # allow those to not be closed properly
            if opening ∈ CAN_BE_LEFT_OPEN
                continue
            end
            parser_exception(:BlockNotClosed, """
                An opening token '$(opening)' was found but not closed.
                """)
        end

        tokens_in_span = @view tokens[i:closing_index]
        new_block = Block(template.name, tokens_in_span)
        push!(blocks, new_block)

        # deactivate all tokens in the span of the block
        to_deactivate = i:closing_index
        is_active[to_deactivate] .= false
        append!(deactivated_tokens, collect(to_deactivate))
    end
    return deactivated_tokens
end

"""
    process_line_return!(blocks, tokens, i)

Process a line return followed by any number of white spaces and X. Depending
on `X`, it will lead to a different interpretation.

* a paragraph break (double line skip)
* a hrule (has --- or *** or ____ on the line)
* an item candidate (starts with * or + or ...)
* a table row candidate (startswith |)
* a blockquote (startswith >).

We disambiguate the different cases based on the two characters after the
whitespaces of the line return (the line return token captures `\n[ \t]*`).
"""
function process_line_return!(b::Vector{Block}, tv::SubVector{Token},
                              i::Int)::Nothing
    t = tv[i]
    if t.name == :SOS
        c = [first(t.ss), next_chars(t, 1)...]
    else
        c = next_chars(t, 2)
    end

    if isempty(c) || c[1] ∈ ('\n', EOS)
        # P_BREAK; if there's not two chars beyond `c` will be empty
        # otherwise if there's `\n` or `EOS` then it's a line skip
        push!(b, Block(:P_BREAK, t.ss))

    # ------------------------------------------------------------------------
    # Hrules
    # NOTE the line MUST start with a triple followed only by
    # the same character, whitespaces and the eventual line return.
    elseif c[1] == c[2] == '-'
        _hrule!(b, t, HR1_PAT)

    elseif c[1] == c[2] == '_'
        _hrule!(b, t, HR2_PAT)

    elseif c[1] == c[2] == '*'
        _hrule!(b, t, HR3_PAT)

    # ------------------------------------------------------------------------
    # List items
    # NOTE for an item candidate, the candidate might not capture
    # the full item if the full item is on several lines, this has to
    # be post-processed when assembling ITEM_x_CAND into lists.
    elseif c[1] in ('+', '-', '*') && c[2] in (' ', '\t')
        cand = until_next_line_return(t)
        ps   = parent_string(cand)
        push!(b, Block(:ITEM_U_CAND, subs(ps, from(t), to(cand))))

    elseif c[1] ∈ NUM_CHAR && c[2] in vcat(NUM_CHAR, ['.', ')'])
        cand = until_next_line_return(t)
        if match(OL_ITEM_PAT, cand) !== nothing
            ps   = parent_string(cand)
            push!(b, Block(:ITEM_O_CAND, subs(ps, from(t), to(cand))))
        end

    # ------------------------------------------------------------------------
    # Table Rows
    # NOTE we're stricter here than usual GFM, every row must start and end
    # with a pipe, every row must be on a single line.
    elseif c[1] == '|'
        # TABLE_ROW_CAND
        cand = until_next_line_return(t)
        if strip(cand)[end] == '|'
            ps = parent_string(cand)
            push!(b, Block(:TABLE_ROW_CAND, subs(ps, from(t), to(cand))))
        end

    # ------------------------------------------------------------------------
    # Blockquote
    elseif c[1] == '>'
        # Blockquote
        cand = until_next_line_return(t)
        ps   = parent_string(cand)
        push!(b, Block(:BLOCKQUOTE_LINE, subs(ps, from(t), to(cand))))
    end
    return
end

"""
    _hrule!(blocks, token, regex)

Helper function to match and process a hrule.
"""
function _hrule!(b, t, r)
    cand  = until_next_line_return(t)
    check = match(r, cand)
    ps    = parent_string(cand)
    isnothing(check) || push!(b, Block(:HRULE, subs(ps, from(t), to(cand))))
    return
end

"""
    form_links!(blocks)

Here we catch the following:

    * [A]     LINK_A   for html(A)
    * [A][B]  LINK_AR  for html(A)
    * [A](B)  LINK_AB  for html(A)
    * ![A]    IMG_A
    * ![A](B) IMG_AB
    * [A]: B  REF      (--> aggregate B, will need to distinguish later)

where 'A' is necessarily non empty, 'B' may be empty.

Note: currently we DO NOT support links with titles such as the following out
of simplicity:

* [A]: B C
* [A](B C)

this allows to not have to check whether B is a link and C is text. If the
user wants links with titles, they should create a command for it. We also do
not support link destinations between <...>.

Note: in the case of a LINK_A, we check around if the previous non whitespace
character and the next non whitespace character don't happen to be } {. In
that specific case, the link is
"""
function form_links!(blocks::Vector{Block})
    isempty(blocks) && return
    nblocks = length(blocks)
    remove  = Int[]
    i       = 1
    nb      = blocks[i]
    ps      = parent_string(nb)

    while i < nblocks
        b  = nb
        nb = blocks[i+1]

        if b.name == :SQ_BRACKETS
            pchar = previous_chars(b)
            nchar = next_chars(b)

            # NOTE: ![]: --> ![] takes precedence.
            # img: is it preceded by '!'?
            # ref: is it followed by ':'?
            # lnk: is the next char '('
            img = false
            if !isempty(pchar)
                img = pchar[1] == '!'
            end
            ref = false
            lab = false
            lar = false
            if !isempty(nchar)
                ref = !img && nchar[1] == ':'
                lab = nchar[1] == '(' && nb.name == :BRACKETS
                lar = nchar[1] == '[' && nb.name == :SQ_BRACKETS
            end
            lnk = lab | lar

            # ref ==> REF, stop
            #
            # img  & !lnk => IMG_A
            # img  & lab  => IMG_AB
            # img  & lar  => IMG_AR
            # !img & !lnk => LINK_A
            # !img & lab  => LINK_AB
            # !img & lar  => LINK_AR

            if ref
                # [...]: block
                # check if the block is at the start of line, otherwise discard
                ss = until_previous_line_return(b)
                if isempty(strip(ss))
                    blocks[i] = Block(:REF, subs(ps, from(b), next_index(b)))
                else
                    push!(remove, i)
                end
            else
                if img
                    if !lnk
                        blocks[i] = Block(:IMG_A, subs(ps, prev_index(b), to(b)))
                    elseif lab
                        blocks[i] = Block(:IMG_AB, subs(ps, prev_index(b), to(nb)))
                        push!(remove, i+1)
                    else
                        blocks[i] = Block(:IMG_AR, subs(ps, prev_index(b), to(nb)))
                        push!(remove, i+1)
                    end
                else
                    if !lnk
                        blocks[i] = Block(:LINK_A, subs(ps, from(b), to(b)))
                    elseif lab
                        blocks[i] = Block(:LINK_AB, subs(ps, from(b), to(nb)))
                        push!(remove, i+1)
                    else
                        blocks[i] = Block(:LINK_AR, subs(ps, from(b), to(nb)))
                        push!(remove, i+1)
                    end
                end
            end
        end
        i += 1
    end

    # check if the last block is maybe a standalone `(!)[...](:)`.
    i = nblocks
    b = blocks[i]
    if i ∉ remove && b.name == :SQ_BRACKETS
        pchar = previous_chars(b)
        nchar = next_chars(b)
        if isempty(pchar)
            img = false
        else
            img = pchar[1] == '!'
        end
        if isempty(nchar)
            ref = false
        else
            ref = !img && nchar[1] == ':'
        end

        if ref
            ss = until_previous_line_return(b)
            if isempty(strip(ss))
                blocks[i] = Block(:REF, subs(ps, from(b), next_index(b)))
            else
                push!(remove, i)
            end
        elseif img
            blocks[i] = Block(:IMG_A, subs(ps, prev_index(b), to(b)))
        else
            blocks[i] = Block(:LINK_A, subs(ps, from(b), to(b)))
        end
    end
    deleteat!(blocks, remove)
    return
end

"""
    remove_inner!(blocks)

Remove blocks which are part of larger blocks (these will get re-formed and
re-processed at an ulterior step).
"""
function remove_inner!(blocks::Vector{Block})
    isempty(blocks) && return
    n_blocks  = length(blocks)
    is_active = ones(Bool, n_blocks)
    @inbounds for i in eachindex(blocks)
        is_active[i] || continue
        to_current = to(blocks[i])
        next_outer = n_blocks + 1
        for j = i+1:n_blocks
            if from(blocks[j]) >= to_current
                next_outer = j
                break
            end
        end
        is_active[i+1:next_outer-1] .= false
    end
    deleteat!(blocks, [i for i in eachindex(blocks) if !is_active[i]])
    return
end

"""
    form_dbb!(blocks)

Find CU_BRACKETS blocks that start with `{{` and and with `}}` and mark them as :DBB.
"""
function form_dbb!(b::Vector{Block})
    @inbounds for i in eachindex(b)
        b[i].name === :CU_BRACKETS || continue
        ss = b[i].ss
        (startswith(ss, "{{") && endswith(ss, "}}")) || continue

        open  = Token(:DBB_OPEN, subs(ss, 1:2))
        li    = lastindex(ss)
        close = Token(:DBB_CLOSE, subs(ss, li-1:li))
        it    = @view b[i].inner_tokens[2:end-1]
        b[i]  = Block(:DBB, open => close, it)
    end
end

function _find_env_blocks!(
        blocks::Vector{Block},
        tokens::SubVector{Token},
        is_active::Vector{Bool}
        )::Nothing

    isempty(blocks) && return

    envs     = Block[]
    discard  = Int[]
    i        = 1
    n_blocks = length(blocks)
    curb     = blocks[i]

    @inbounds while i < n_blocks
        nxtb = blocks[i+1]
        j    = i

        if curb.name == :LX_BEGIN
            nxtb.name == :CU_BRACKETS || enverr()
            env_name  = content(nxtb) |> strip

            # look ahead trying to find the proper closing \end{...}
            open_depth    = 1
            closing_index = -1
            probe         = nxtb
            j            += 1

            while j < n_blocks && open_depth != 0
                cand  = probe.name
                probe = blocks[j + 1]

                closer = (cand == :LX_END) &&
                         (probe.name == :CU_BRACKETS) &&
                         (env_name == content(probe) |> strip)

                opener = (cand == :LX_BEGIN) &&
                         (probe.name == :CU_BRACKETS) &&
                         (env_name == content(probe) |> strip)

                if closer
                    open_depth -= 1
                elseif opener
                    open_depth += 1
                end

                j += 1
            end
            open_depth   != 0 && enverr("{$env_name}")
            closing_index = j

            # tokens in span (there is always at least LX_BEGIN and END)
            env_from = from(curb)
            env_to   = to(blocks[closing_index])

            toks_1   = findfirst(t -> from(t) >= env_from, tokens)
            toks_n   = findlast(t -> to(t) <= env_to, tokens)
            toks_r   = toks_1:toks_n

            # deactivate them all
            is_active[toks_r] .= false

            # mark all blocks in the range as to be discarded
            append!(discard, i:closing_index)

            # keep track of the block and its tokens
            b = Block(
                Symbol("ENV_$(env_name)"),
                tokens[toks_1] => tokens[toks_n],
                @view tokens[toks_r]
            )
            push!(envs, b)
        end

        curb = nxtb
        i    = j + 1
    end
    # discard all blocks within the env
    deleteat!(blocks, discard)
    append!(blocks, envs)
    return
end