text:
using .Networks

using Dates
using DataStructures
using JuMP
using Printf
using Parameters

@with_kw mutable struct RSOAssessmentModel <: AbstractModelContainer
    model::Model = Model()

    # Uncertainty model
    #bus,ts
    uncertain_load::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    #limitableUnit,ts
    uncertain_prod::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()

    # Market model
    #unit, ts
    p_injected::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    #pilotable, ts
    b_in::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    b_marg::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    b_out::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()

    #ts : rho (subproblem variable)
    overflow::SortedDict{DateTime,VariableRef} =
        SortedDict{DateTime,VariableRef}()

    # duals : branch,ts (subproblem)
    pos_flow_limit_duals::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()
    pos_flow_limit_indicators::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()
    neg_flow_limit_duals::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()
    neg_flow_limit_indicators::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()

    #sum_bus ptdf[branch,bus] * (prod_bus - conso_bus)
    signed_flow::SortedDict{Tuple{String,DateTime},AffExpr} =
        SortedDict{Tuple{String,DateTime},AffExpr}()

    #objective
    full_obj::AffExpr = AffExpr(0.)
end

function add_load_uncertainties_vars!(model_container::RSOAssessmentModel,
                                    network, TS, assessment_uncertainties)
    for bus in Networks.get_buses(network)
        bus_id = Networks.get_id(bus)
        lb_l = get_assessment_uncertainties_lb(assessment_uncertainties, bus_id)
        ub_l = get_assessment_uncertainties_ub(assessment_uncertainties, bus_id)
        for ts in TS
            name =  @sprintf("uncertain_load[%s,%s]", bus_id, ts);
            model_container.uncertain_load[bus_id,ts] = @variable(model_container.model, base_name = name,
                                                                lower_bound = lb_l, upper_bound = ub_l)
        end
    end

    return model_container.uncertain_load
end
function add_prod_uncertainties_vars!(model_container::RSOAssessmentModel,
                                    network, TS, assessment_uncertainties, tso_actions)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        lb_l = get_assessment_uncertainties_lb(assessment_uncertainties, gen_id)
        for ts in TS
            # this allows to set p_inj[gen_id] == uncertain_prod without needing to express min(uncertain, limit) in the model
            ub_l = min(get_assessment_uncertainties_ub(assessment_uncertainties, gen_id),
                    get_limitation(tso_actions, gen_id, ts),
                    Networks.get_p_max(limitable_gen))
            name =  @sprintf("uncertain_prod[%s,%s]", gen_id, ts);
            model_container.uncertain_prod[gen_id, ts] = @variable(model_container.model, base_name=name,
                                                                lower_bound = lb_l, upper_bound = ub_l)
        end
    end

    return model_container.uncertain_prod
end
function add_uncertainties_vars!(model_container::RSOAssessmentModel,
                                network, TS, assessment_uncertainties, tso_actions)
    add_load_uncertainties_vars!(model_container,
                                network, TS, assessment_uncertainties)
    add_prod_uncertainties_vars!(model_container,
                                network, TS, assessment_uncertainties, tso_actions)
    return model_container
end

function add_limitable_prod_vars!(model_container::RSOAssessmentModel, network, TS)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        p_max = Networks.get_p_max(limitable_gen)
        for ts in TS
            name =  @sprintf("p_injected[%s,%s]", gen_id, ts)
            model_container.p_injected[gen_id,ts] =
                @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=p_max)
        end
    end

    return model_container.p_injected
end

function add_use_limitables_constraints!(model_container::RSOAssessmentModel, network, TS, tso_actions)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        for ts in TS
            limit_l = get_limitation(tso_actions, gen_id, ts)
            @assert(upper_bound(model_container.uncertain_prod[gen_id, ts]) <= limit_l)
            @constraint(model_container.model,
                        model_container.p_injected[gen_id,ts] == model_container.uncertain_prod[gen_id, ts])
        end
    end
    return model_container
end

function add_pilotable_prod_vars!(model_container_p::RSOAssessmentModel, network, TS, tso_actions)
    p_injected_l = model_container_p.p_injected
    b_in_l = model_container_p.b_in
    b_marg_l = model_container_p.b_marg
    b_out_l = model_container_p.b_out

    for pilotable_gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
        gen_id = Networks.get_id(pilotable_gen)
        for ts in TS
            pmin_l, pmax_l = safeget_imposition(tso_actions, gen_id, ts)
            name_l =  @sprintf("b_in[%s,%s]", gen_id, ts);
            b_in_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("b_marginal[%s,%s]", gen_id, ts);
            b_marg_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("b_out[%s,%s]", gen_id, ts);
            b_out_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("p_injected[%s,%s]", gen_id, ts);
            p_injected_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l,
                                                lower_bound=0., upper_bound = pmax_l);

            @constraint(model_container_p.model,
                        p_injected_l[gen_id,ts] >= ( pmax_l * b_in_l[gen_id,ts] +  pmin_l * b_marg_l[gen_id,ts] ));
            @constraint(model_container_p.model,
                        p_injected_l[gen_id,ts] <= (1 - b_out_l[gen_id,ts]) * pmax_l);
            @constraint(model_container_p.model,
                        b_out_l[gen_id,ts] + b_marg_l[gen_id,ts] + b_in_l[gen_id,ts] == 1.);
        end
    end

    return model_container_p
end

function add_market_vars!(model_container::RSOAssessmentModel, network, TS, tso_actions)
    add_limitable_prod_vars!(model_container, network, TS)
    add_pilotable_prod_vars!(model_container, network, TS, tso_actions)
end

function add_cheapest_prod_constraints!(model_container::RSOAssessmentModel, network, TS)
    pilotables = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    for (index, gen_1) in enumerate(pilotables)
        cost_1 = Networks.get_prop_cost(gen_1)
        gen_id_1 = Networks.get_id(gen_1)
        for gen_2 in pilotables[index:end]
            cost_2 = Networks.get_prop_cost(gen_2)
            if cost_1 < cost_2
                gen_id_2 = Networks.get_id(gen_2)
                for ts in TS
                    @constraint(model_container.model, model_container.b_in[gen_id_1,ts] >= model_container.b_in[gen_id_2,ts] );
                end
            end
        end
    end

    return model_container
end

function add_eod_constraint!(model_container_p::RSOAssessmentModel, network, TS)

    for ts in TS

        supply_l = AffExpr(0)
        for gen in Networks.get_generators(network)
            gen_id = Networks.get_id(gen)
            supply_l += model_container_p.p_injected[gen_id, ts];
        end

        demand_l = AffExpr(0.)
        for bus in Networks.get_buses(network)
            bus_id = Networks.get_id(bus)
            demand_l += model_container_p.uncertain_load[bus_id, ts];
        end

        @constraint(model_container_p.model, supply_l == demand_l);
    end

    return model_container_p
end

function add_overflow_var!(model_container::RSOAssessmentModel, TS, upper_bound)
    for ts in TS
        name = @sprintf("overflow[%s]", ts)
        model_container.overflow[ts] =
            @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=upper_bound)
    end
    return model_container.overflow
end

function create_objective!(model_container::RSOAssessmentModel, network, TS)::AffExpr
    model_container.full_obj = AffExpr(0.)
    add_coeffxsum_cost!(model_container.full_obj, model_container.overflow, 1.)
    return model_container.full_obj
end

function add_uncertainties_constraints!(model_container_l::RSOAssessmentModel, network, TS, assessment_uncertainties, configs)
    #FIXME constraints on uncertainties go here !
    #e.g | u[bus1] - u[bus2] | < 5
    #e.g ( 1 - u[bus1] / ub[limitable_1] ) <  ( 1 - u[limitable_1] / ub[limitable_1] )
end

function create_flow_limit_duals_and_indicators(model_container::RSOAssessmentModel, network, TS)
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)
        for ts in TS
            name=@sprintf("pos_flow_duals[%s,%s]", branch_id, ts)
            model_container.pos_flow_limit_duals[branch_id, ts] =
                @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=1.)
            name=@sprintf("pos_flow_indicators[%s,%s]", branch_id, ts)
            model_container.pos_flow_limit_indicators[branch_id, ts] =
                @variable(model_container.model, base_name=name, binary=true)
            name=@sprintf("neg_flow_duals[%s,%s]", branch_id, ts)
            model_container.neg_flow_limit_duals[branch_id, ts] =
                @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=1.)
            name=@sprintf("neg_flow_indicators[%s,%s]", branch_id, ts)
            model_container.neg_flow_limit_indicators[branch_id, ts] =
                @variable(model_container.model, base_name=name, binary=true)
        end
    end

    return model_container
end

function create_flow_expressions(model_container::RSOAssessmentModel, network, TS)
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)

        for ts in TS

            branch_flow_expr = AffExpr(0.)
            for bus in Networks.get_buses(network)
                bus_id = Networks.get_id(bus)
                ptdf = Networks.safeget_ptdf_elt(network, branch_id, bus_id)

                # + injections
                for gen in Networks.get_generators(bus)
                    gen_id = Networks.get_id(gen)
                    branch_flow_expr += ptdf * model_container.p_injected[gen_id,ts]
                end

                # - loads
                branch_flow_expr -= ptdf * model_container.uncertain_load[bus_id,ts]
            end

            model_container.signed_flow[branch_id, ts] = branch_flow_expr
        end
    end
end

function add_flow_constraints(model_container::RSOAssessmentModel, network, TS)
    model_l = model_container.model
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)
        flow_limit_l = Networks.get_limit(branch)

        for ts in TS
            overflow_l = model_container.overflow[ts]
            signed_flow_l = model_container.signed_flow[branch_id, ts]

            name = @sprintf("c_pos_flow_limit[%s,%s]",branch_id,ts)
            @constraint(model_l, signed_flow_l <= flow_limit_l + overflow_l , base_name=name)
            name = @sprintf("c_neg_flow_limit[%s,%s]",branch_id,ts)
            @constraint(model_l, -signed_flow_l <= flow_limit_l + overflow_l , base_name=name)

        end
    end
end

function add_kkt_stationarity_constraints(model_container::RSOAssessmentModel, network, TS)
    for ts in TS
        duals_sum_l = AffExpr(0.)
        for branch_id in Networks.get_id.(Networks.get_branches(network))
            add_to_expression!(duals_sum_l,
                model_container.pos_flow_limit_duals[branch_id,ts] + model_container.neg_flow_limit_duals[branch_id,ts])
        end
        @constraint(model_container.model, duals_sum_l == 1.)
    end

    return model_container
end

function add_kkt_complementarity_constraints(model_container, network, TS)
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)
        flow_limit_l = Networks.get_limit(branch)

        for ts in TS
            # pos_flow_limit_duals
            kkt_var = model_container.pos_flow_limit_duals[branch_id, ts]
            cstr_expr = - model_container.signed_flow[branch_id, ts] + flow_limit_l + model_container.overflow[ts]
            b_indicator = model_container.pos_flow_limit_indicators[branch_id, ts]
            ub_cstr = compute_ub(cstr_expr)
            ub_kkt = upper_bound(kkt_var)
            formulate_complementarity_constraints!(model_container.model,
                                                    kkt_var, cstr_expr, b_indicator, ub_kkt, ub_cstr)
            # neg_flow_limit_duals
            kkt_var = model_container.neg_flow_limit_duals[branch_id, ts]
            cstr_expr = model_container.signed_flow[branch_id, ts] + flow_limit_l + model_container.overflow[ts]
            b_indicator = model_container.neg_flow_limit_indicators[branch_id, ts]
            ub_cstr = compute_ub(cstr_expr)
            ub_kkt = upper_bound(kkt_var)
            formulate_complementarity_constraints!(model_container.model,
                                                    kkt_var, cstr_expr, b_indicator, ub_kkt, ub_cstr)
        end
    end

    return model_container
end

function formulate_overflow_subproblem(model_container::RSOAssessmentModel, network, TS)
    create_flow_limit_duals_and_indicators(model_container, network, TS)
    create_flow_expressions(model_container, network, TS)
    add_flow_constraints(model_container, network, TS)
    add_kkt_stationarity_constraints(model_container, network, TS)
    add_kkt_complementarity_constraints(model_container, network, TS)

    return model_container
end

function formulate_rso_assessment(network, TS, assessment_uncertainties, tso_actions, configs)
    model_container_l = RSOAssessmentModel()

    add_uncertainties_vars!(model_container_l, network, TS, assessment_uncertainties, tso_actions)
    add_market_vars!(model_container_l, network, TS, tso_actions)
    add_overflow_var!(model_container_l, TS, configs.BIG_M)

    add_use_limitables_constraints!(model_container_l, network, TS, tso_actions)
    add_cheapest_prod_constraints!(model_container_l, network, TS)
    add_eod_constraint!(model_container_l, network, TS)
    add_uncertainties_constraints!(model_container_l, network, TS, assessment_uncertainties, configs)

    formulate_overflow_subproblem(model_container_l, network, TS)

    create_objective!(model_container_l, network, TS)
    @objective(get_model(model_container_l), Max, model_container_l.full_obj)

    return model_container_l
end

function is_validated(model_container::RSOAssessmentModel)
    return value(model_container.full_obj) < 1e-09
end

function has_positive_slack(model_container::RSOAssessmentModel)::Bool
    return false # ( has_positive_value(model_container.lower.lol_model.p_loss_of_load)
           #  || has_positive_value(model_container.lower.lol_model.p_cut_prod) )
end

function get_assessment_uncertainties_lb(assessment_uncertainties, bus_or_limitable::String)
    return assessment_uncertainties[bus_or_limitable][1]
end
function get_assessment_uncertainties_ub(assessment_uncertainties, bus_or_limitable::String)
    return assessment_uncertainties[bus_or_limitable][2]
end

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf


text:
using GeometryTypes, Test
function f(G, x)
    fx = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

    if !(G == nothing)
        G1 = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
        G2 = 200.0 * (x[2] - x[1]^2)
        gx = Point(G1,G2)

        return fx, gx
    else
        return fx
    end
end
f_obj = OnceDiffed(f)
minimize(f_obj, Point(3.0,3.0), GradientDescent(Inverse()), OptOptions())
minimize(f_obj, Point(3.0,3.0), BFGS(Inverse()), OptOptions())
minimize(f_obj, Point(3.0,3.0), DFP(Inverse()), OptOptions())
minimize(f_obj, Point(3.0,3.0), SR1(Inverse()), OptOptions())

minimize(f_obj, Point(3.0,3.0), GradientDescent(Direct()), OptOptions())
minimize(f_obj, Point(3.0,3.0), BFGS(Direct()), OptOptions())
minimize(f_obj, Point(3.0,3.0), DFP(Direct()), OptOptions())
minimize(f_obj, Point(3.0,3.0), SR1(Direct()), OptOptions())

ext: jl


text:
using PkgSRA
using Plots, Random, Distributions, Interpolations
pyplot()
Random.seed!(11)
using BSON: @save
include("../scripts/paper_settings.jl");
# include("../../utils/sindy_turing_utils.jl")
# include("../../utils/sindy_statistics_utils.jl")
# include("../../utils/main_algorithm_utils.jl")

################################################################################
#####
##### SIR example
#####
include(EXAMPLE_FOLDERNAME*"example_sir.jl")

# Define the multivariate forcing function
num_ctr = 3;
    U_starts = rand(1, num_ctr) .* tspan[2]/2
    U_widths = 0.6;
    amplitude = 100.0
my_U_func_time2(t) = U_func_time(t, u0,
                        U_widths, U_starts,
                        F_dim=1,
                        amplitude=amplitude)

# Get data
sol = solve_sir_system(U_func_time=my_U_func_time2)
dat = Array(sol)
plot(ts,dat[1,:], label="S");
    plot!(ts,dat[2,:], label="I");
    plot!(ts,dat[3,:], label="R")
numerical_grad = numerical_derivative(dat, ts)
true_grad = core_dyn_true(dat)
plot(ts, true_grad[2,:]);plot!(ts, numerical_grad[2,:])

## Also get baseline true/ideal cases
# Uncontrolled
# dat_raw = Array(solve_sir_system())
# numerical_grad_raw = numerical_derivative(dat_raw, ts)
# true_grad_raw = zeros(size(dat))
# for i in 1:size(dat,2)
#     true_grad_raw[:,i] = sir_system(dat_raw[:,i], Float64.(p), [0])
# end

# True control signal
U_true = zeros(size(dat))
for (i, t) in enumerate(ts)
    U_true[:,i] = my_U_func_time2(t)
end
plot(ts,dat[1,:], label="S");
    plot!(ts,dat[2,:], label="I");
    plot!(ts,dat[3,:], label="R")
    plot!(ts,U_true', label="Control")

###
### Fit a naive model
###
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);

# Upgrade: Use AIC
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_unctr,best_criterion,all_criteria,all_models) =
    sindyc_ensemble(dat, numerical_grad, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,100),
                    use_clustering_minimization=true)
println("Naive SINDy model:")
    print_equations(sindy_unctr, digits=5)
println("True equations are:")
    print_equations(core_dyn_true, digits=5)
scatter(sum.(val_list), all_criteria)
    title!("AIC for various sparsities (naive model)")
    xlabel!("Number of nonzero terms")

# Integrate this (poor) model
# Callback aborts if it blows up
condition(u,t,integrator) = any(abs.(u).>1e4)
cb = DiscreteCallback(condition, terminate!)
prob_unctr = ODEProblem(sindy_unctr, u0, tspan, [0], callback=cb)
sindy_dat_unctr = Array(solve(prob_unctr, Tsit5(), saveat=ts));

# let d = sindy_dat_unctr, d2 = dat
#     i = 1;
#     plot(d[i,:], label="Naive model");#plot!(d[2,:]);plot!(d[3,:])
#     plot!(d2[i,:], label="Data");#plot!(d2[2,:]);plot!(d2[3,:])
# end

###
### Calculate the residuals
###
# NOTE: do not use Turing, for speed
sindy_grad1 = sindy_unctr(dat, 0)
residual = numerical_grad .- sindy_grad1
noise_guess = 2*abs(median(residual[1,:])) # Really should use Turing here

plot(residual[3,:], label="Residual")
    hline!([noise_guess], label="Noise line");
    hline!([-noise_guess], label="Noise line")
    # title!("Residual")

ctr_guess = process_residual(residual, noise_guess)

###
### Subsample based on residual
###
accepted_ind = subsample_using_residual(residual,
            noise_guess, min_length=4)

num_pts = 2000
start_ind = 1
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1]
dat_sub = dat[:,subsample_ind]
grad_sub = numerical_grad[:,subsample_ind]

# Plot the subsampled points
i = 2
    plot(dat[i,:])
    scatter!(subsample_ind, dat_sub[i,:], color=:blue)
    plot!(U_true[i,:])
    title!("Subsampled points for index $i")

# Any control signal in the subset?
# U_sub = U_true[:,subsample_ind]
# plot(U_sub')
#     title!("Control signals in the subsampled dataset")

# val_list = calc_permutations(6,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_sub,best_criterion,all_criteria,all_models,best_index) =
    sindyc_ensemble(dat_sub, grad_sub, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,2*noise_guess),
                    use_clustering_minimization=true)
println("Best index is $best_index:")
    print_equations(sindy_sub, digits=5)
println("True equations are:")
    print_equations(core_dyn_true, digits=5)
scatter(sum.(val_list), all_criteria)
    title!("AIC for various sparsities")
    xlabel!("Number of nonzero terms")

# Plot a function for the subsampled points
# f(S, I, R) = -0.00027*I.*I .- 0.0002*I.*R
# g(S, I, R) = -0.2*I
# i = 2
#     plot(g(dat[1,:], dat[2,:], dat[3,:]), label="True")
#     scatter!(subsample_ind, g(dat_sub[1,:], dat_sub[2,:], dat_sub[3,:]))
#     plot!(f(dat[1,:], dat[2,:], dat[3,:]), label="Found")
#     title!("Subsampled points for index $i")


text:
export RBM

struct RBM{VT,MT,F} <: KetNeuralNetwork
    a::VT
    b::VT
    W::MT

    f::F
end
#@functor RBM

functor(x::RBM) = (a=x.a, b=x.b, W=x.W), y -> RBM(y..., x.f)

"""
    RBMSplit([T=Complex{STD_REAL_PREC}], N, α, f=af_softplus, [initW, initb])

Constructs a Restricted Bolzmann Machine to encode a wavefunction,
with weights of type `T`, `N` input neurons and `αN` hidden neurons.
This is the Neural Quantum State (NQS) Ansatz of [1].

`N` must match the size of the lattice.

By default the activation function is a sigmoid. You can also use logcosh by
passing as an additional parameter `af_logcosh`.

The initial parameters of the neurons are initialized with a rescaled normal
distribution of width 0.01 for the coupling matrix and 0.05 for the local
biases. The default initializers can be overriden by specifying

initW=(dims...)->rescaled_normal(T, 0.01, dims...)
initb=(dims...)->rescaled_normal(T, 0.05, dims...)
inita=(dims...)->rescaled_normal(T, 0.01, dims...)

Refs:
    https://arxiv.org/abs/1606.02318
"""
RBM(in, α, args...) = RBM(ComplexF32, in, α, args...)
RBM(T::Type, hilb::AbstractHilbert, args...) =
    RBM(T, nsites(hilb), args...)
RBM(T::Type, in::Int, α, σ::Function=af_softplus,
    initW=(dims...)->rescaled_normal(T, 0.01, dims...),
    initb=(dims...)->rescaled_normal(T, 0.01, dims...),
    inita=(dims...)->rescaled_normal(T, 0.01, dims...)) =
    RBM(inita(in), initb(convert(Int,α*in)),
        initW(convert(Int,α*in), in), σ)

out_type(net::RBM{VT,MT,F}) where {VT,MT,F} = eltype(VT)
is_analytic(net::RBM) = true

(net::RBM)(σ::AbstractVector) = transpose(net.a)*σ .+ sum(net.f.(net.b .+ net.W*σ))
(net::RBM)(σ::AbstractMatrix) = transpose(net.a)*σ .+ sum(net.f.(net.b .+ net.W*σ), dims=1)

function Base.show(io::IO, m::RBM{T,VT,F}) where {T,VT,F}
    print(io, "RBM($(eltype(VT)), n=$(length(m.a)), α=$(length(m.b)/length(m.a)), f=$F)")
end

# Cached version
mutable struct RBMCache{VT} <: NNCache{RBM}
    θ::VT
    logℒθ::VT
    σ::VT
    valid::Bool # = false
end

cache(net::RBM) =
    RBMCache(similar(net.b),
             similar(net.b),
             similar(net.a),
             false)

function (net::RBM)(c::RBMCache, σ_r)
    T=eltype(net.W)
    θ = c.θ
    logℒθ = c.logℒθ

    #θ .= net.b .+ net.W * σ
    σ = copyto!(c.σ, σ_r)
    copyto!(θ, net.b)
    mul!(θ, net.W, σ, T(1.0), T(1.0))
    #BLAS.gemv!('N', T(1.0), net.W, σ, T(1.0), θ)

    logℒθ .= net.f.(θ)
    logψ = dot(σ,net.a) + sum(logℒθ)
    return logψ
end

function logψ_and_∇logψ!(∇logψ, net::RBM, c::RBMCache, σ_r)
    T=eltype(net.W)
    θ = c.θ
    logℒθ = c.logℒθ

    #θ .= net.b .+ net.W * σ
    σ = copyto!(c.σ, σ_r)
    copyto!(θ, net.b)
    mul!(θ, net.W, σ, T(1.0), T(1.0))
    #BLAS.gemv!('N', T(1.0), net.W, σ, T(1.0), θ)

    logℒθ .= net.f.(θ)
    logψ = dot(σ,net.a) + sum(logℒθ)

    ∇logψ.a .= σ
    ∇logψ.b .= fwd_der.(net.f, θ)
    ∇logψ.W .= ∇logψ.b  .* transpose(σ)
    return logψ
end


text:
struct LocallyConstant <: AbstractRegression

end

"""
    Apply the analog method on catalog of historical data
    to generate forecasts.
"""
function predict(LocallyConstant, xf::Array{Float64,2}, weight)

    # weighted mean and covariance
    xf_mean = vec(sum(xf .* weights', dims = 2))
    Exf = xf .- xf_mean

    Symmetric(1.0 ./ (1.0 .- sum(weights .^ 2)) .* (Exf .* weights') * Exf')

end


text:
using MultiDimDictionaries

d = MultiDimDictionary(([(1, 1), (2, 1), (1, 2)]), [1, 2, 3])

# Get elements

@show d[1, 1] == 1
@show d[2, 1] == 2
@show d[1, 2] == 3

@show d[(1, 1)] == 1
@show d[(2, 1)] == 2
@show d[(1, 2)] == 3

@show d[LinearIndex(1)] == 1
@show d[LinearIndex(2)] == 2
@show d[LinearIndex(3)] == 3

# Error
# @show d[(1, 1)] == 1
# @show d[(2, 1)] == 2
# @show d[(1, 2)] == 3

# Set elements

@show d[1, 1] = 4
@show d[(2, 1)] = 5
@show d[LinearIndex(3)] = 6

@show d[1, 1] == 4
@show d[2, 1] == 5
@show d[1, 2] == 6

insert!(d, (2, 2), 12)

@show d[2, 2] == 12

delete!(d, (1, 1))

@show !haskey(d, (1, 1))
@show haskey(d, (1, 2))
@show haskey(d, (2, 1))
@show haskey(d, (2, 2))

delete!(d, 1, 2)

@show !haskey(d, 1, 1)
@show !haskey(d, 1, 2)
@show haskey(d, 2, 1)
@show haskey(d, 2, 2)

d2 = MultiDimDictionary([(3, 2), (2, 3)], [32, 23])
d12 = merge(d, d2)

@show d12 isa MultiDimDictionary
@show !haskey(d12, 1, 2)
@show haskey(d12, 2, 1)
@show haskey(d12, 3, 2)

d1 = MultiDimDictionary()
d1[1] = 1
d1[2] = 2

@show d1[1] == 1
@show d1[2] == 2

d2 = MultiDimDictionary()
d2[1] = 3
d2[2] = 4

@show d2[1] == 3
@show d2[2] == 4

dv = [d1; d2]

@show dv[1] == 1
@show dv[2] == 2
@show dv[3] == 3
@show dv[4] == 4

dh = hcat(d1, d2)

@show dh[1, 1] == 1
@show dh[2, 1] == 2
@show dh[1, 2] == 3
@show dh[2, 2] == 4

dh = hcat(d1, d2; new_dim_keys=("X", "Y"))

@show dh[1, "X"] == 1
@show dh[2, "X"] == 2
@show dh[1, "Y"] == 3
@show dh[2, "Y"] == 4

dh_X = dh[:, "X"]

@show dh_X[1, "X"] == 1
@show dh_X[2, "X"] == 2
@show !isassigned(dh_X, 1, "Y")
@show !isassigned(dh_X, 2, "Y")

dh_Y = dh[:, "Y"]

@show dh_Y[1, "Y"] == 3
@show dh_Y[2, "Y"] == 4

dh_Z = dh[:, "Z"]

@show !isassigned(dh_Z, 1, "X")
@show !isassigned(dh_Z, 2, "X")
@show !isassigned(dh_Z, 1, "Y")
@show !isassigned(dh_Z, 2, "Y")

d = d1 ⊔ d2

@show d[1, 1] == 1
@show d[1, 2] == 2
@show d[2, 1] == 3
@show d[2, 2] == 4

d = disjoint_union(d1, d2; new_dim_keys=("X", "Y"))

@show d["X", 1] == 1
@show d["X", 2] == 2
@show d["Y", 1] == 3
@show d["Y", 2] == 4

d_slice = d[[("X", 1), ("Y", 2)]]

@show d_slice["X", 1] == 1
@show !isassigned(d_slice, "X", 2)
@show !isassigned(d_slice, "Y", 1)
@show d_slice["Y", 2] == 4

d2 = (d ⊔ d)[1, :]

@show d2[1, "X", 1] == 1
@show d2[1, "X", 2] == 2
@show d2[1, "Y", 1] == 3
@show d2[1, "Y", 2] == 4


text:
function getemptymaze(dimx, dimy)
    maze = ones(Int, dimx, dimy)
    maze[1,:] .= maze[end,:] .= 0
    maze[:, 1] .= maze[:, end] .= 0
    maze
end

function setwall!(maze, startpos, endpos)
    dimx, dimy = startpos - endpos
    if dimx == 0
        maze[startpos[1], startpos[2]:endpos[2]] .= 0
    else
        maze[startpos[1]:endpos[1], startpos[2]] .= 0
    end
end

function indto2d(maze, pos)
    dimx = size(maze, 1)
    [rem(pos, dimx), div(pos, dimx) + 1]
end
function posto1d(maze, pos)
    dimx = size(maze, 1)
    (pos[2] - 1) * dimx + pos[1]
end

function checkpos(maze, pos)
    count = 0
    for dx in -1:1
        for dy in -1:1
            count += maze[(pos + [dx, dy])...] == 0
        end
    end
    count
end

function addrandomwall!(maze)
    startpos = rand(ENV_RNG, findall(x -> x != 0, maze[:]))
    startpos = indto2d(maze, startpos)
    starttouch = checkpos(maze, startpos)
    if starttouch > 0
        return 0
    end
    endx, endy = startpos
    if rand(ENV_RNG, 0:1) == 0 # horizontal
        while checkpos(maze, [endx, startpos[2]]) == 0
            endx += 1
        end
        if maze[endx + 1, startpos[2]] == 1 &&
            maze[endx + 1, startpos[2] + 1] ==
            maze[endx + 1, startpos[2] - 1] == 0
            endx -= 1
        end
    else
        while checkpos(maze, [startpos[1], endy]) == 0
            endy += 1
        end
        if maze[startpos[1], endy + 1] == 1 &&
            maze[startpos[1] + 1, endy + 1] ==
            maze[startpos[1] - 1, endy + 1] == 0
            endx -= 1
        end
    end
    setwall!(maze, startpos, [endx, endy])
    return 1
end

function setTandR!(d)
    for s in d.statefrommaze[findall(x -> x != 0, reshape(d.maze, :))]
        setTandR!(d, s)
    end
end
function setTandR!(d, s)
    T = d.mdp.trans_probs
    R = d.mdp.reward
    statefrommaze = d.statefrommaze
    goals = d.goals
    ns = d.mdp.observationspace.n
    maze = d.maze
    idx_goals = findfirst(x -> x == s, goals)
    if idx_goals !== nothing
        R.value[s] = d.goalrewards[idx_goals]
    end
    pos = indto2d(maze, d.mazefromstate[s])
    for (aind, a) in enumerate(([0, 1], [1, 0], [0, -1], [-1, 0]))
        nextpos = maze[(pos + a)...] == 0 ? pos : pos + a
        if d.neighbourstateweight > 0
            positions = []
            push!(positions, nextpos)
            weights = [1.]
            for dir in ([0, 1], [1, 0], [0, -1], [-1, 0])
                if maze[(nextpos + dir)...] != 0
                    push!(positions, nextpos + dir)
                    push!(weights, d.neighbourstateweight)
                end
            end
            states = map(p -> statefrommaze[posto1d(maze, p)], positions)
            weights /= sum(weights)
            T[aind, s] = sparsevec(states, weights, ns)
        else
            nexts = statefrommaze[posto1d(maze, nextpos)]
            T[aind, s] = sparsevec([nexts], [1.], ns)
        end
    end
end

function isinsideframe(maze, i)
    nx, ny = size(maze)
    i > nx && i < nx * (ny - 1) && i % nx != 0 && i % nx != 1
end
function n_effective(n, f, list)
    N = n === nothing ? div(length(list), Int(1/f)) : n
    min(N, length(list))
end
function breaksomewalls!(m; f = 1/50, n = nothing, rng = ENV_RNG)
    zeros = Int[]
    for i in 1:length(m)
        m[i] == 0 && isinsideframe(m, i) && push!(zeros, i)
    end
    pos = sample(rng, zeros, n_effective(n, f, zeros), replace = false)
    m[pos] .= 1
    m
end
function addobstacles!(m; f = 1/100, n = nothing, rng = ENV_RNG)
    nz = findall(x -> x == 1, reshape(m, :))
    pos = sample(rng, nz, n_effective(n, f, nz), replace = false)
    m[pos] .= 0
    m
end
"""
    struct DiscreteMaze
        mdp::MDP
        maze::Array{Int, 2}
        goals::Array{Int, 1}
        statefrommaze::Array{Int, 1}
        mazefromstate::Array{Int, 1}
"""
struct DiscreteMaze{T}
    mdp::T
    maze::Array{Int, 2}
    goals::Array{Int, 1}
    goalrewards::Array{Float64, 1}
    statefrommaze::Array{Int, 1}
    mazefromstate::Array{Int, 1}
    neighbourstateweight::Float64
end
"""
    DiscreteMaze(; nx = 40, ny = 40, nwalls = div(nx*ny, 10), ngoals = 1,
                   goalrewards = 1, stepcost = 0, stochastic = false,
                   neighbourstateweight = .05)

Returns a `DiscreteMaze` of width `nx` and height `ny` with `nwalls` walls and
`ngoals` goal locations with reward `goalreward` (a list of different rewards
for the different goal states or constant reward for all goals), cost of moving
`stepcost` (reward = -`stepcost`); if `stochastic = true` the actions lead with
a certain probability to a neighbouring state, where `neighbourstateweight`
controls this probability.
"""
function DiscreteMaze(; nx = 40, ny = 40, nwalls = div(nx*ny, 10), kwargs...)
    m = getemptymaze(nx, ny)
    [addrandomwall!(m) for _ in 1:nwalls]
    breaksomewalls!(m)
    DiscreteMaze(m; kwargs...)
end
function DiscreteMaze(maze; ngoals = 1, goalrewards = 1., stepcost = 0,
                      stochastic = false,
                      neighbourstateweight = stochastic ? .05 : 0.,
                      compressed = true)
    na = 4
    nzpos = findall(x -> x != 0, reshape(maze, :))
    statefrommaze = compressed ? cumsum(reshape(maze, :)) : collect(1:length(maze))
    mazefromstate = compressed ? nzpos : collect(1:length(maze))
    legalstates = statefrommaze[nzpos]
    ns = length(mazefromstate)
    T = Array{SparseVector{Float64,Int}}(undef, na, ns)
    goals = sort(sample(ENV_RNG, legalstates, ngoals, replace = false))
    R = DeterministicNextStateReward(fill(-stepcost, ns))
    isterminal = zeros(Int, ns); isterminal[goals] .= 1
    isinitial = setdiff(legalstates, goals)
    res = DiscreteMaze(MDP(DiscreteSpace(ns, 1),
                           DiscreteSpace(na, 1),
                           rand(ENV_RNG, legalstates),
                           T, R,
                           isinitial, isterminal),
                       maze,
                       goals,
                       typeof(goalrewards) <: Number ? fill(goalrewards, ngoals) :
                                                       goalrewards,
                       statefrommaze,
                       mazefromstate,
                       neighbourstateweight)
    setTandR!(res)
    res
end

interact!(env::DiscreteMaze, a) = interact!(env.mdp, a)
reset!(env::DiscreteMaze) = reset!(env.mdp)
getstate(env::DiscreteMaze) = getstate(env.mdp)
actionspace(env::DiscreteMaze) = actionspace(env.mdp)

function plotenv(env::DiscreteMaze)
    goals = env.goals
    mazefromstate = env.mazefromstate
    m = deepcopy(env.maze)
    m[mazefromstate[goals]] .= 3
    m[mazefromstate[env.mdp.state]] = 2
    imshow(m, colormap = 21, size = (400, 400))
end


text:
# cudnnReduceTensorOp_t
const CUDNN_REDUCE_TENSOR_ADD = Cint(0)
const CUDNN_REDUCE_TENSOR_MUL = Cint(1)
const CUDNN_REDUCE_TENSOR_MIN = Cint(2)
const CUDNN_REDUCE_TENSOR_MAX = Cint(3)
const CUDNN_REDUCE_TENSOR_AMAX = Cint(4)
const CUDNN_REDUCE_TENSOR_AVG = Cint(5)
const CUDNN_REDUCE_TENSOR_NORM1 = Cint(6)
const CUDNN_REDUCE_TENSOR_NORM2 = Cint(7)
const CUDNN_REDUCE_TENSOR_MUL_NO_ZEROS = Cint(8)

# cudnnReduceTensorIndices_t
const CUDNN_REDUCE_TENSOR_NO_INDICES = Cint(0)
const CUDNN_REDUCE_TENSOR_FLATTENED_INDICES = Cint(1)

mutable struct ReduceTensorDesc
    ptr::Cptr

    function ReduceTensorDesc(::Type{T}, op::Cint) where T
        ref = Ref{Cptr}()
        @cudnn :cudnnCreateReduceTensorDescriptor (Ptr{Cptr},) ref
        desc = new(ref[])
        finalizer(desc) do x
            @cudnn :cudnnDestroyReduceTensorDescriptor (Cptr,) x.ptr
        end

        ind = op == CUDNN_REDUCE_TENSOR_MIN || op == CUDNN_REDUCE_TENSOR_MAX ?
            CUDNN_REDUCE_TENSOR_FLATTENED_INDICES :
            CUDNN_REDUCE_TENSOR_NO_INDICES
        @cudnn(:cudnnSetReduceTensorDescriptor,
            (Cptr,Cint,Cint,Cint,Cint,Cint),
            desc, op, datatype(T), CUDNN_NOT_PROPAGATE_NAN, ind, CUDNN.CUDNN_32BIT_INDICES)
        desc
    end
end

Base.cconvert(::Type{Cptr}, desc::ReduceTensorDesc) = desc.ptr

const DICT_ReduceTensorDesc = Dict()

function reduce(A::CuArray{T}, dim, op) where T
    if size(A,dim) == 1 # CUDNN_STATUS_BAD_PARAM
        C = A
        indices = zeros(CuArray{Cint}, length(A)÷size(A,dim))
        return C, indices
    end

    h = gethandle()
    reducedesc = get!(DICT_ReduceTensorDesc, (T,op)) do
        ReduceTensorDesc(T, op)
    end
    adesc = TensorDesc(A, 4)
    cdims = Int[size(A)...]
    cdims[dim] = 1
    C = similar(A, cdims...)
    cdesc = TensorDesc(C, 4)

    ref = Ref{Csize_t}()
    @cudnn(:cudnnGetReductionIndicesSize,
        (Cptr,Cptr,Cptr,Cptr,Ptr{Csize_t}),
        h, reducedesc, adesc, cdesc, ref)
    indices = ref[] == 0 ? C_NULL : CuArray{Cint}(Int(ref[])÷sizeof(Cint))

    ref = Ref{Csize_t}()
    @cudnn(:cudnnGetReductionWorkspaceSize,
        (Cptr,Cptr,Cptr,Cptr,Ptr{Csize_t}),
        h, reducedesc, adesc, cdesc, ref)
    workspace = CuArray{UInt8}(Int(ref[]))

    @cudnn(:cudnnReduceTensor,
        (Cptr,Cptr,Cptr,Csize_t,Cptr,Csize_t,
        Cptr,Cptr,Cptr,Cptr,Cptr,Cptr),
        h, reducedesc, indices, length(indices)*sizeof(Cint), workspace, length(workspace),
        T[1], adesc, A, T[0], cdesc, C)
    C, indices
end


text:
const SWIPL_LIB = "libswipl"

"""
    Flags
#ifdef PL_KERNEL
#define PL_Q_DEBUG              0x0001  /* = TRUE for backward compatibility */
#endif
#define PL_Q_NORMAL             0x0002  /* normal usage */
#define PL_Q_NODEBUG            0x0004  /* use this one */
#define PL_Q_CATCH_EXCEPTION    0x0008  /* handle exceptions in C */
#define PL_Q_PASS_EXCEPTION     0x0010  /* pass to parent environment */
#define PL_Q_ALLOW_YIELD        0x0020  /* Support I_YIELD */
#define PL_Q_EXT_STATUS         0x0040  /* Return extended status */
#ifdef PL_KERNEL
#define PL_Q_DETERMINISTIC      0x0100  /* call was deterministic */
#endif
"""

const PL_Q_DEBUG = 0x0001
const PL_Q_NORMAL = 0x0002
const PL_Q_NODEBUG = 0x0004
const PL_Q_CATCH_EXCEPTION = 0x0008
const PL_Q_PASS_EXCEPTION = 0x0010
const PL_Q_ALLOW_YIELD = 0x0020
const PL_Q_EXT_STATUS = 0x0040

"""
    Term types
#define PL_VARIABLE      (1)            /* nothing */
#define PL_ATOM          (2)            /* const char * */
#define PL_INTEGER       (3)            /* int */
#define PL_RATIONAL      (4)            /* rational number */
#define PL_FLOAT         (5)            /* double */
#define PL_STRING        (6)            /* const char * */
#define PL_TERM          (7)
#define PL_NIL           (8)            /* The constant [] */
#define PL_BLOB          (9)            /* non-atom blob */
#define PL_LIST_PAIR     (10)           /* [_|_] term */
                                        /* PL_unify_term() */
#define PL_FUNCTOR       (11)           /* functor_t, arg ... */
#define PL_LIST          (12)           /* length, arg ... */
#define PL_CHARS         (13)           /* const char * */
#define PL_POINTER       (14)           /* void * */
"""

const PL_VARIABLE = 1
const PL_ATOM = 2
const PL_INTEGER = 3
const PL_FLOAT = 5
const PL_STRING = 6
const PL_TERM = 7
const PL_NIL = 8
const PL_LIST_PAIR = 10
const PL_FUNCTOR = 11

"""
Foreign functions Flags

#define PL_FA_NOTRACE           (0x01)  /* foreign cannot be traced */
#define PL_FA_TRANSPARENT       (0x02)  /* foreign is module transparent */
#define PL_FA_NONDETERMINISTIC  (0x04)  /* foreign is non-deterministic */
#define PL_FA_VARARGS           (0x08)  /* call using t0, ac, ctx */
#define PL_FA_META              (0x40)  /* Additional meta-argument spec */
#define PL_FA_SIG_ATOMIC        (0x80)  /* Internal: do not dispatch signals */
"""

const PL_DETERMINISTIC = 0
const PL_FA_META = 0x40
const PL_FA_TRANSPARENT = 0x02
const PL_FA_NONDETERMINISTIC = 0x04
const PL_FA_NOTRACE = 0x01
const PL_FA_VARARGS = 0x08

"""
 Functions for starting and shutting down the SWIPL instance
"""
function PL_start(exec_path::String)
    argv = [exec_path, "-q"]
    ccall((:PL_initialise, SWIPL_LIB), Cint, (Cint, Ptr{Ptr{UInt8}}), length(argv), argv)
end

PL_start() = PL_start("/usr/local/bin/swipl")

function PL_is_initialised(exec_path::String)
    argv = [exec_path, "-q"]
    ccall((:PL_is_initialised, SWIPL_LIB), Cint, (Cint, Ptr{Ptr{UInt8}}), length(argv), argv)
end

PL_is_initialised() = PL_is_initialised()

function PL_stop(status::Int64)
    ccall((:PL_halt, SWIPL_LIB), Cint, (Cint,), convert(Cint, status))
end

PL_stop() = PL_stop(0)

function PL_cleanup(status::Int64)
    ccall((:PL_cleanup, SWIPL_LIB), Cint, (Cint,), status)
end

PL_cleanup() = PL_cleanup(0)

"""
    Atoms and functors
"""

function PL_new_atom(name::String)
    name_to_use = collect([convert(UInt8, c) for c in name])
    ccall((:PL_new_atom, SWIPL_LIB), Cint, (Cstring,), name) #Cstring(pointer(name)))
end

function PL_atom_chars(term::Cint)
    r = ccall((:PL_atom_chars, SWIPL_LIB), Cstring, (Cint,), term)
    unsafe_string(r)
end

function PL_new_functor(atom::Cint, arity::Int64)
    ccall((:PL_new_functor, SWIPL_LIB), Cint, (Cint, Cint), atom, convert(Cint, arity))
end

function PL_functor_name(functor::Cint)
    ccall((:PL_functor_name, SWIPL_LIB), Cint, (Cint,), functor)
end

function PL_functor_arity(functor::Cint)
    ccall((:PL_functor_arity, SWIPL_LIB), Cint, (Cint,), functor)
end

"""
    Testing types of terms
"""

function PL_term_type(term::Cint)
    ccall((:PL_term_type, SWIPL_LIB), Cint, (Cint,), term)
end

function PL_is_variable(term::Cint)
    ccall((:PL_is_variable, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_ground(term::Cint)
    ccall((:PL_is_ground, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_atom(term::Cint)
    ccall((:PL_is_atom, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_string(term::Cint)
    ccall((:PL_is_string, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_integer(term::Cint)
    ccall((:PL_is_integer, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_float(term::Cint)
    ccall((:PL_is_float, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_callable(term::Cint)
    ccall((:PL_is_callable, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_compound(term::Cint)
    ccall((:PL_is_compound, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_list(term::Cint)
    ccall((:PL_is_list, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_pair(term::Cint)
    ccall((:PL_is_pair, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_atomic(term::Cint)
    ccall((:PL_is_atomic, SWIPL_LIB), Cint, (Cint,), term) > 0
end

function PL_is_number(term::Cint)
    ccall((:PL_is_number, SWIPL_LIB), Cint, (Cint,), term) > 0
end

"""
    Getting data from terms
"""

function PL_get_atom(term::Cint)
    atm = Ref{Cint}(0)
    ccall((:PL_get_atom, SWIPL_LIB), Cint, (Cint, Ref{Cint}), term, atm)
    atm[]
end

function PL_get_chars(term::Cint)
    atm = Ref{Cstring}(C_NULL)
    ccall((:PL_get_chars, SWIPL_LIB), Cint, (Cint, Ref{Cstring}, Cint), term, atm, convert(Cint, 0x00000002))
    unsafe_string(atm[])
end

function PL_get_integer(term::Cint)
    atm = Ref{Cint}(0)
    ccall((:PL_get_integer, SWIPL_LIB), Cint, (Cint, Ref{Cint}), term, atm)
    atm[]
end

function PL_get_bool(term::Cint)
    atm = Ref{Cint}(0)
    ccall((:PL_get_bool, SWIPL_LIB), Cint, (Cint, Ref{Cint}), term, atm)
    atm[]
end

function PL_get_atom_chars(term::Cint)
    tmp = Ref{Cstring}(C_NULL)
    rc = ccall((:PL_get_atom_chars, SWIPL_LIB), Cint, (Cint, Ref{Cstring}), term, tmp)
    unsafe_string(tmp[])
end

function PL_get_float(term::Cint)
    val = Ref{Cfloat}(0)
    ccall((:PL_get_float, SWIPL_LIB), Cint, (Cint, Ref{Cfloat}), term, val)
    val[]
end

function PL_get_arity_name(term::Cint)
    name = Ref{Cint}(0)
    ar = Ref{Csize_t}(0)
    ccall((:PL_get_arity_name, SWIPL_LIB), Cint, (Cint, Ref{Cint}, Ref{Csize_t}), term, name, ar)
    (name[], ar[])
end

function PL_get_compound_name_arity(term::Cint)
    name = Ref{Cint}(0)
    ar = Ref{Csize_t}(0)
    ccall((:PL_get_compound_name_arity, SWIPL_LIB), Cint, (Cint, Ref{Cint}, Ref{Csize_t}), term, name, ar)
    (name[], ar[])
end

function PL_get_arg(index::Csize_t, term::Cint, term_result::Cint)
    ccall((:PL_get_arg, SWIPL_LIB), Cint, (Csize_t, Cint, Cint), index, term, term_result)
end

function PL_get_list(list::Cint, head::Cint, tail::Cint)
    ccall((:PL_get_list, SWIPL_LIB), Cint, (Cint, Cint, Cint), list, head, tail)
end

"""
    Constructing terms
"""
function PL_new_term_ref()
    ccall((:PL_new_term_ref, SWIPL_LIB), Cint, ())
end

function PL_new_term_refs(arity::Int64)
    ccall((:PL_new_term_refs, SWIPL_LIB), Cint, (Cint,), arity)
end

function PL_put_variable(term::Cint)
    ccall((:PL_put_variable, SWIPL_LIB), Cvoid, (Cint,), term)
end

function PL_put_atom(term::Cint, atom::Cint)
    ccall((:PL_put_atom, SWIPL_LIB), Cvoid, (Cint, Cint), term, atom)
end

function PL_put_bool(term::Cint, atom::Cint)
    ccall((:PL_put_bool, SWIPL_LIB), Cvoid, (Cint, Cint), term, atom)
end

function PL_put_atom_chars(term::Cint, name::String)
    ccall((:PL_put_atom_chars, SWIPL_LIB), Cint, (Cint, Cstring), term, name)
end

function PL_put_string_chars(term::Cint, name::String)
    ccall((:PL_put_string_chars, SWIPL_LIB), Cint, (Cint, Cstring), term, name)
end

function PL_put_integer(term::Cint, val::Int64)
    ccall((:PL_put_integer, SWIPL_LIB), Cint, (Cint, Clong), term, val)
end

function PL_put_float(term::Cint, val::Float64)
    ccall((:PL_put_float, SWIPL_LIB), Cint, (Cint, Cdouble), term, val)
end

function PL_put_functor(term::Cint, functor::Cint)
    ccall((:PL_put_functor, SWIPL_LIB), Cint, (Cint, Cint), term, functor)
end

function PL_put_term(term::Cint, functor::Cint)
    ccall((:PL_put_term, SWIPL_LIB), Cint, (Cint, Cint), term, functor)
end

function PL_put_list(term::Cint)
    ccall((:PL_put_list, SWIPL_LIB), Cint, (Cint,), term)
end

function PL_put_nil(term::Cint)
    ccall((:PL_put_nil, SWIPL_LIB), Cint, (Cint,), term)
end

function PL_copy_term_ref(term::Cint)
    ccall((:PL_copy_term_ref, SWIPL_LIB), Cint, (Cint,), term)
end

function PL_cons_functor(term::Cint, functor::Cint, compound_term::Cint)
    ccall((:PL_cons_functor_v, SWIPL_LIB), Cint, (Cint, Cint, Cint), term, functor, compound_term)
end

function PL_cons_list(term::Cint, head::Cint, tail::Cint)
    ccall((:PL_cons_list, SWIPL_LIB), Cint, (Cint, Cint, Cint), term, head, tail)
end

"""
    Unification
"""

function PL_unify(term1::Cint, term2::Cint)
    ccall((:PL_unify, SWIPL_LIB), Cint, (Cint, Cint), term1, term2)
end

function PL_unify_atom(term1::Cint, term2::Cint)
    ccall((:PL_unify, SWIPL_LIB), Cint, (Cint, Cint), term1, term2)
end

function PL_unify_bool(term::Cint, val::Bool)
    ccall((:PL_unify_book, SWIPL_LIB), Cint, (Cint, Cint), term, val)
end

function PL_unify_atom_chars(term::Cint, val::String)
    ccall((:PL_unify_atom_chars, SWIPL_LIB), Cint, (Cint, Cstring), term, val)
end

function PL_unify_string_chars(term::Cint, val::String)
    ccall((:PL_unify_string_chars, SWIPL_LIB), Cvoid, (Cint, Cstring), term, val)
end

function PL_unify_integer(term::Cint, val::Int64)
    ccall((:PL_unify_integer, SWIPL_LIB), Cint, (Cint, Cint), term, val)
end

function PL_unify_float(term::Cint, val::Float64)
    ccall((:PL_unify_float, SWIPL_LIB), Cint, (Cint, Cdouble), term, val)
end

function PL_unify_functor(term::Cint, functor::Cint)
    ccall((:PL_unify_functor, SWIPL_LIB), Cint, (Cint, Cint), term, functor)
end

function PL_unify_compound(term::Cint, functor::Cint)
    ccall((:PL_unify_compound, SWIPL_LIB), Cint, (Cint, Cint), term, functor)
end

function PL_unify_list(term::Cint, head::Cint, tail::Cint)
    ccall((:PL_unify_list, SWIPL_LIB), Cint, (Cint, Cint, Cint), term, head, tail)
end

function PL_unify_nil(term::Cint)
    ccall((:PL_unify_nil, SWIPL_LIB), Cint, (Cint, ), term)
end

function PL_unify_arg(index::Int64, term::Cint, arg_term::Cint)
    ccall((:PL_unify_arg, SWIPL_LIB), Cint, (Cint, Cint, Cint), index, term, arg_term)
end

"""
    Calling Prolog
"""

function PL_predicate(predicate_name::String, arity::Int64)
    ccall((:PL_predicate, SWIPL_LIB), Ptr{Cvoid}, (Cstring, Cint, Cstring), predicate_name, arity, C_NULL)
end

function PL_open_query(predicate_t::Ptr{Cvoid}, term::Cint, flag::UInt16)
    ccall((:PL_open_query, SWIPL_LIB), Cint, (Ptr{Cvoid}, Cint, Ptr{Cvoid}, Cint), C_NULL, flag, predicate_t, term)
end

PL_open_query(predicate_t::Ptr{Cvoid}, term::Cint) = PL_open_query(predicate_t, term, PL_Q_NORMAL)

function PL_next_solution(qid_t::Cint)
    ccall((:PL_next_solution, SWIPL_LIB), Cint, (Cint, ), qid_t)
end

function PL_cut_query(qid_t::Cint)
    ccall((:PL_cut_query, SWIPL_LIB), Cint, (Cint, ), qid_t)
end

function PL_close_query(qid_t::Cint)
    ccall((:PL_close_query, SWIPL_LIB), Cint, (Cint,), qid_t)
end

"""
    Foreign function interface
"""
function PL_register_foreign(name::String, arity::Int64, fnc::Ptr{Cvoid}, flag::Int64)
    ccall((:PL_register_foreign, SWIPL_LIB), Cint, (Cstring, Cint, Ptr{Cvoid}, Cint), name, arity, fnc, flag)
end


text:
using SubstitutionTilings.NumFields

@NumFields.simple_number_field_concrete Qρ  [-36, 0, 8, 0] ρ

@NumFields.simple_number_field_lazy Qζ [-1, 1, -1, 1] ζ
Base.promote_rule(::Type{Qζ}, ::Type{<:Integer}) = Qζ

function func_macro(x::Qζ)
    return reduce!(x*x*x*x*x - x*x + x)
end
@benchmark func_macro(data) setup=(data=Qζ(rand(1:1000,4)))
# median 749 ns microseconds mutable

@NumFields.simple_number_field_concrete Qζ2 [-1, 1, -1, 1] ζ2
function Base.show(io::IO, x::Qζ2)
    result = "("
    some_shown = false
    for (i, coeff) in enumerate(x.coeffs)
        if (coeff != 0)
            if some_shown
                result = result*" + "
            else
                some_shown = true
            end
            result = result*string(coeff)
            result = result*"ζ2^"
            result = result*string(i-1)
        end
    end
    result = result*")/"
    result = result*string(x.denom)
    print(io, result)
end

Base.promote_rule(::Type{Qζ2}, ::Type{<:Integer}) = Qζ2
function func_concrete(x)
    return x*x*x*x*x - x*x + x
end
@benchmark func_concrete(data) setup=(data=Qζ2(rand(1:1000,4)))
# median 1.7 micros

function arith(x)
    return div(x[1] * x[2] - x[4], x[3]^5)
end


text:
"""
This module provides utilities to build neural networks with Flux,
along with a library of standard architectures.
"""
module FluxNets

export SimpleNet, SimpleNetHP, ResNet, ResNetHP

using ..Network
using Base: @kwdef
import ..GameInterface, ..Util, ..CyclicSchedule

using CUDA

import Flux
import Functors

CUDA.allowscalar(false)
array_on_gpu(::Array) = false
array_on_gpu(::CuArray) = true
array_on_gpu(arr) = error("Usupported array type: ", typeof(arr))

using Flux: relu, softmax, flatten
using Flux: Chain, Dense, Conv, BatchNorm, SkipConnection
import Zygote

#####
##### Flux Networks
#####

"""
    FluxNetwork{Game} <: AbstractNetwork{Game}

Abstract type for neural networks implemented using the _Flux_ framework.

The `regularized_params_` function must be overrided for all layers containing
parameters that are subject to regularization.

Provided that the above holds, `FluxNetwork` implements the full
network interface with the following exceptions:
[`Network.HyperParams`](@ref), [`Network.hyperparams`](@ref),
[`Network.forward`](@ref) and [`Network.on_gpu`](@ref).
"""
abstract type FluxNetwork{Game} <: AbstractNetwork{Game} end

function Base.copy(nn::Net) where Net <: FluxNetwork
  #new = Net(Network.hyperparams(nn))
  #Flux.loadparams!(new, Flux.params(nn))
  #return new
  return Base.deepcopy(nn)
end

Network.to_cpu(nn::FluxNetwork) = Flux.cpu(nn)

function Network.to_gpu(nn::FluxNetwork)
  CUDA.allowscalar(false)
  return Flux.gpu(nn)
end

function Network.set_test_mode!(nn::FluxNetwork, mode)
  Flux.testmode!(nn, mode)
end

Network.convert_input(nn::FluxNetwork, x) =
  Network.on_gpu(nn) ? Flux.gpu(x) : x

Network.convert_output(nn::FluxNetwork, x) = Flux.cpu(x)

Network.params(nn::FluxNetwork) = Flux.params(nn)

# This should be included in Flux
function lossgrads(f, args...)
  val, back = Zygote.pullback(f, args...)
  grad = back(Zygote.sensitivity(val))
  return val, grad
end

function Network.train!(callback, nn::FluxNetwork, opt::Adam, loss, data, n)
  optimiser = Flux.ADAM(opt.lr)
  params = Flux.params(nn)
  for (i, d) in enumerate(data)
    l, grads = lossgrads(params) do
      loss(d...)
    end
    Flux.update!(optimiser, params, grads)
    callback(i, l)
  end
end

function Network.train!(
    callback, nn::FluxNetwork, opt::CyclicNesterov, loss, data, n)
  lr = CyclicSchedule(
    opt.lr_base,
    opt.lr_high,
    opt.lr_low, n=n)
  momentum = CyclicSchedule(
    opt.momentum_high,
    opt.momentum_low,
    opt.momentum_high, n=n)
  optimiser = Flux.Nesterov(opt.lr_low, opt.momentum_high)
  params = Flux.params(nn)
  for (i, d) in enumerate(data)
    l, grads = lossgrads(params) do
      loss(d...)
    end
    Flux.update!(optimiser, params, grads)
    optimiser.eta = lr[i]
    optimiser.rho = momentum[i]
    callback(i, l)
  end
end

regularized_params_(l) = []
regularized_params_(l::Flux.Dense) = [l.W]
regularized_params_(l::Flux.Conv) = [l.weight]

# Reimplementation of what used to be Flux.prefor, does not visit leaves
function foreach_flux_node(f::Function, x, seen = IdDict())
  Functors.isleaf(x) && return
  haskey(seen, x) && return
  seen[x] = true
  f(x)
  for child in Flux.trainable(x)
    foreach_flux_node(f, child, seen)
  end
end

function Network.regularized_params(net::FluxNetwork)
  ps = Flux.Params()
  foreach_flux_node(net) do p
    for r in regularized_params_(p)
      any(x -> x === r, ps) || push!(ps, r)
    end
  end
  return ps
end

function Network.gc(::FluxNetwork)
  GC.gc(true)
  # CUDA.reclaim()
end

#####
##### Common functions between two-head neural networks
#####

"""
    TwoHeadNetwork{Game} <: FluxNetwork{G}

An abstract type for two-head neural networks implemented with Flux.

Subtypes are assumed to have fields
`hyper`, `common`, `vhead` and `phead`. Based on those, an implementation
is provided for [`Network.hyperparams`](@ref), [`Network.forward`](@ref) and
[`Network.on_gpu`](@ref), leaving only [`Network.HyperParams`](@ref) to
be implemented.
"""
abstract type TwoHeadNetwork{G} <: FluxNetwork{G} end

function Network.forward(nn::TwoHeadNetwork, board)
  c = nn.common(board)
  v = nn.vhead(c)
  p = nn.phead(c)
  return (p, v)
end

# Flux.@functor does not work do to Network being parametric
function Flux.functor(nn::Net) where Net <: TwoHeadNetwork
  children = (nn.common, nn.vhead, nn.phead)
  constructor = cs -> Net(nn.hyper, cs...)
  return (children, constructor)
end

Network.hyperparams(nn::TwoHeadNetwork) = nn.hyper

Network.on_gpu(nn::TwoHeadNetwork) = array_on_gpu(nn.vhead[end].b)

#####
##### Include networks library
#####

include("architectures/simplenet.jl")
include("architectures/resnet.jl")

end


text:
module SQLREPL

import LibPQ
import DataFrames
import ReplMaker
import REPL

# This uses some implementation details from SQLStrings
using SQLStrings

function match_magic_syntax(str)
    m = match(r"(\\d) *(.*)", str)
    if !isnothing(m)
        return (m[1], m[2])
    else
        return nothing
    end
end

# One imperfect way to allow multi-line editing
function valid_input_checker(prompt_state)
    cmdstr = String(take!(copy(REPL.LineEdit.buffer(prompt_state))))
    magic = match_magic_syntax(cmdstr)
    if !isnothing(magic)
        return !isempty(magic[2])
    end
    length(findall(")", cmdstr)) == length(findall("(", cmdstr))
end

function libpq_eval(conn, str)
    magic = match_magic_syntax(str)
    if !isnothing(magic)
        if magic[1] == "\\d"
            table_name = magic[2]
            query = sql```
                SELECT
                    column_name,
                    data_type,
                    character_maximum_length,
                    column_default,
                    is_nullable
                FROM INFORMATION_SCHEMA.COLUMNS
                WHERE table_name = $table_name
            ```
            # TODO: Allow a pattern here and use `LIKE` with table_name?
        else
            error("Unknown magic command $(magic[1])")
        end
    else
        # Trick: We macro expand @sql_cmd here in order to resolve symbols from
        # SQLStrings without requiring `using SQLStrings` in `Main`.
        query = macroexpand(@__MODULE__, :(@sql_cmd $str))
    end
    return quote
        $LibPQ.execute($conn, $query) |> $DataFrames.DataFrame
    end
end

"""
    SQLREPL.connect(conn; start_key=')')

Connect a PostgresSQL to a Julia REPL mode activated with `start_key`.  The
connection `conn` can either be a Postgres connection string or a
`LibPQ.Connection` object.
"""
function connect(conn::LibPQ.Connection; start_key=')')
    ReplMaker.initrepl(s->libpq_eval(conn, s),
                       valid_input_checker = valid_input_checker,
                       prompt_text="SQL> ",
                       prompt_color = :blue,
                       start_key=start_key,
                       mode_name="SQL")
end

function connect(connstr::AbstractString; kws...)
    connect(LibPQ.Connection(connstr); kws...)
end

end


text:
module IDX

# The following description of the IDX format comes from
#
#    http://yann.lecun.com/exdb/mnist/
#
# which is also the source of the MNIST data.
#
#    THE IDX FILE FORMAT
#
#    The IDX file format is a simple format for vectors and multidimensional matrices of various numerical types.  The basic
#    format is
#
#       magic number
#       size in dimension 0
#       size in dimension 1
#       size in dimension 2
#       .....
#       size in dimension N
#       data
#
#    The magic number is an integer (MSB first). The first 2 bytes are always 0.
#
#    The third byte codes the type of the data:
#
#       0x08: unsigned byte
#       0x09: signed byte
#       0x0B: short (2 bytes)
#       0x0C: int (4 bytes)
#       0x0D: float (4 bytes)
#       0x0E: double (8 bytes)
#
#    The 4-th byte codes the number of dimensions of the vector/matrix: 1 for vectors, 2 for matrices....
#
#    The sizes in each dimension are 4-byte integers (MSB first, high endian, like in most non-Intel processors).

# ----------------------------------------
# EXPORTED INTERFACE
# ----------------------------------------

export load

# ----------------------------------------
# CONSTANTS
# ----------------------------------------

const type_constructors = [ UInt8,
                            Int8,
                            Int16, # Actually, this one is unused. I'm really not sure why ... I haven't found proper
                                   # documentation for the format other than the website reference quoted above.
                            Int16,
                            Int32,
                            Float32,
                            Float64 ]

# ----------------------------------------
# FUNCTIONS
# ----------------------------------------

function items(handle, constructor, sizes)
    data = reshape(convert(Array{constructor}, read(handle, prod(sizes) * sizeof(constructor))), Tuple(reverse(sizes)))

    return ( sizes, constructor, data )
end

function load(file_name)
    handle = open(file_name, "r")

    ( _, _, constructor, dimension ) = magic_number(handle)

    return items(handle, constructor, sizes(handle, dimension))
end

function magic_number(handle; constructor = nothing)
    magic_number = bswap(read(handle, UInt32))

    ( _, _, type_index, dimension ) = ( UInt8(magic_number & 0xFF000000 >> 0x18),
                                        UInt8(magic_number & 0x00FF0000 >> 0x10),
                                        UInt8(magic_number & 0x0000FF00 >> 0x08),
                                        UInt8(magic_number & 0x000000FF >> 0x00) )

    if type_index < 0x01 || type_index > 0x07 || type_index == 0x04
        throw(DomainError(type_index, "Invalid magic number ... the type index \"$type_index\" is out of range."))
    end

    return ( 0x00, 0x00, constructor == nothing ? type_constructors[type_index] : constructor, dimension )
end

function sizes(handle, dimension)
    return map(_::Int64 -> bswap(read(handle, UInt32)), 1:dimension)
end

end

max_stars_repo_name: jlegare/IDX

max_issues_repo_name: jlegare/IDX

max_forks_repo_name: jlegare/IDX


text:
using JLD, HDF5, GZip

function sub_read(w2::Array{Float64, 2}, co1::Int64, co2::Int64, cur_rows::Array{Float64, 2})
	w2[1 + co1 * 21: (co1 + 1) * 21, 1 + co2 * 21: (co2 + 1) * 21] = cur_rows[1:end, 1:end]
	w2[1 + co2 * 21: (co2 + 1) * 21, 1 + co1 * 21: (co1 + 1) * 21] = cur_rows[1:end, 1:end]'
end

function read_raw(raw_file, out_file)
	in_file = open(raw_file)
	raw_read = readlines(in_file)
	close(in_file)

	ncol = findfirst(n -> n[1] == '#', raw_read) - 1

	w1 = zeros(ncol, 20)
	count = 0
	for line in raw_read[1:ncol]
		count += 1
		w1[count, :] = map(xx -> parse(Float64, xx), split(strip(line), "\t"))
	end

	w2 = zeros(ncol * 21, ncol * 21)
	co1 = 0
	co2 = 0
	row = 0
	cur_rows = zeros(21, 21)
	for line in raw_read[(ncol + 1):(end-1)]
		if line[1] == '#'
			if !(co1 == 0 && co2 == 0)
				sub_read(w2, co1, co2, cur_rows)
			end

			row = 0
			pos = split(strip(line))
			cur_rows *= 0.0 #reset
			co1 = parse(Int64, pos[2])
			co2 = parse(Int64, pos[3])
		else
			row += 1
			cur_rows[row, :] = map(mm -> parse(Float64, mm), split(strip(line), "\t"))
		end
	end

	w1_x = hcat(w1, zeros(ncol, 1))'
	w1_alt = reshape(w1_x, (21 * ncol, 1))

	save(out_file, "w1", Float64.(w1_alt), "w2", w2)

end

function main()
	raw_file = ARGS[end - 1]
	jld_file = ARGS[end]
	if isfile(raw_file)
		if endswith(raw_file, ".raw")
			if endswith(jld_file, ".jld")
				read_raw(raw_file, jld_file)
			else
				println("Second file should end in .jld")
			end
		else
			println("First file should end in .raw")
		end
	else
		println("Incorrect input, check if raw file exists.")
	end
end

main()


text:
using ArgParse

function readInput(filename::String)
    open(filename) do file
        data = read(file, String)
        line, image = map(strip, split(data, "\n\n"))
        line = map(x -> x == "#", split(line, ""))

        rows = split(image, "\n")
        matrix = falses(length(rows), length(rows))
        for (i, row) in enumerate(rows)
            matrix[i, :] = map(x -> x == "#", split(row, ""))
        end

        return line, matrix
    end
end

function step(line, image, pad_value; pad=2)
    function pad_false(image, pad)
        image_padded = falses(size(image) .+ pad*2)
        image_padded[1+pad:end-pad, 1+pad:end-pad] .= image
        return image_padded
    end

    function pad_true(image, pad)
        image_padded = trues(size(image) .+ pad*2)
        image_padded[1+pad:end-pad, 1+pad:end-pad] .= image
        return image_padded
    end

    image_padded = pad_value ? pad_true(image, pad) : pad_false(image, pad)
    image_next = pad_value != line[1] ? trues(size(image_padded)) : falses(size(image_padded))

    kernel = 1 .<< (8:-1:0)
    for x in 2:size(image_next, 1) - 1
        for y in 2:size(image_next, 2) - 1
            view = image_padded[x - 1:x + 1, y - 1:y + 1]
            # row first
            index = sum(kernel .* vec(transpose(view)))
            image_next[x, y] = line[1 + index]
        end
    end

    # keep image size as small as possible
    pivot = image_next[1, 1]
    while all(image_next[pad + 1, :] .== pivot) && all(image_next[end-pad, :] .== pivot) &&
        all(image_next[:, pad + 1] .== pivot) && all(image_next[:, end-pad] .== pivot)
        image_next = image_next[2:end-1, 2:end-1]
    end

    return image_next
end

function multiple_steps(line, image; steps::Int, pad_odd::Bool = false, pad_even::Bool = false)
    for i in 1:steps
        image = step(line, image, i % 2 == 1 ? pad_odd : pad_even)
    end
    return image, sum(image)
end

function parse_cli()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--example"
            help = "Calculate example"
            action = :store_true
    end
    return parse_args(s)
end

args = parse_cli();
if args["example"]
    line, image = readInput("example.txt");

    image2, sum_image2 = multiple_steps(line, image, steps=2);
    @assert sum_image2 == 35;

    image50, sum_image50 = multiple_steps(line, image2, steps=48);
    @assert sum_image50 == 3351;
end

line, image = readInput("input.txt");

image2, sum_image2 = multiple_steps(line, image, steps=2, pad_even=true);
println(sum_image2);

image50, sum_image50 = multiple_steps(line, image2, steps=48, pad_even=true);
println(sum_image50);

# using BenchmarkTools
# @btime multiple_steps(line, image, steps=50, pad_even=true);


text:
@noinline function currentstate(cf::LinearStateSpace, ::Type{Val{false}})
    idx = cf.t[]::Int
    (cf.f.a[idx], cf.f.P[idx])
end

@noinline function currentstate(cf::LinearStateSpace, ::Type{Val{true}})
    idx = cf.t[]::Int
    (cf.f.a[idx], cf.f.P[idx], cf.f.Pinf[idx], cf.f.Pstar[idx])
end

# function updatelik!(ss, v, F, F⁻¹, ispd)
#     p, m, r = size(ss)
#     s = ispd ? 0.5*(logdet(F)) : s = 0.5*(logdet(F) + v'F⁻¹*v)
#     ss.loglik[1] += -p/2*log(2π) - first(s)
# end

function updatelik!(ss, v, F, F⁻¹, Fflag)
    p, m, r = size(ss)
    s = Fflag == 1 ? 0.5*(logdet(F)) : 0.5*(logdet(F) + v'F⁻¹*v)
    ss.loglik[1] += -p/2*log(2π) - first(s)
end

function updatelik!(ss, v, F, F⁻¹)
    p, m, r = size(ss)
    s = 0.5*(logdet(F) + v'F⁻¹*v)
    ss.loglik[1] += -p/2*log(2π) - first(s)
end

function loglik(v, F, F⁻¹)
     - 0.5*(logdet(F) + v'F⁻¹*v)
end

function loglik(v, F, F⁻¹, flag)
    s = flag == 1 ? 0.5*(logdet(F)) : 0.5*(logdet(F) + v'F*v)
    - first(s)
end

isdiffuse(Pinf::StaticMatrix) = any(Pinf.data.!=0)
isdiffuse(Pinf::AbstractFloat) = Pinf>0
isdiffuse(Pinf::AbstractMatrix) = any(Pinf.!=0)

numstates(a::AbstractArray) = size(a,1)
numstates(a::AbstractFloat) = 1

nummeasur(a::AbstractArray) = size(a,1)
nummeasur(a::AbstractFloat) = 1

numr(a::AbstractArray) = size(a,2)
numr(a::AbstractFloat) = 1

Base.size(cf::LinearStateSpace) = size(cf.p)

Base.size(p::KFParms{P}) where P<:AbstractFloat = (1,1,1)
Base.size(p::KFParms{P}) where P<:AbstractArray = (size(p.Z)..., size(p.Q, 1))

testposdef(x::AbstractArray) = isposdef(x)
testposdef(x::StaticMatrix) = all(map(u-> u[1]>0, eig(x)))

function safeinverse(x)
    flag = any(abs.(x) .> 1e-07) ? 1 : 0
    if flag > 0
        try
            (flag, inv(x))
        catch mesg
            throw(DomainError)
        end
    else
        (flag, x)
    end
end

toarray(b) = b
toarray(b::Real) = [b]

function checkfilterdataconsistency(ss::LinearStateSpace, y::AbstractVector)
    T = length(y)
    p, m, r = size(ss.p)
    @assert p==1 "Inconsistent dimension. Y must be a (Tx$p) Matrix{T}."
    (T, p, m, r, ss.t[])
end

function checkfilterdataconsistency(ss::LinearStateSpace, Y::AbstractArray)
    q, T = size(Y)
    p, m, r = size(ss.p)
    msg = "a ($q x T) Matrix"
    @assert p==q "Inconsistent dimension. Y must be "*msg
    (T, p, m, r, ss.t[])
end

function resizecontainer!(ss::LinearStateSpace, n::Int64)
    resize!(ss.f.a.data, n)
    resize!(ss.f.P.data, n)
    resize!(ss.c.F.data, n)
    resize!(ss.c.F⁻¹.data, n)
    resize!(ss.c.Y.data, n)
end

function Masked(a::NTuple)
    rnan = map(y->findall(x -> isnan.(x), y), a)
    nnan = map(y->findall(x -> !isnan.(x), y), a)
    freeitr = map(x -> enumerate(x), rnan)
    fixeditr = map(x -> enumerate(x), nnan)
    Masked(rnan, nnan, freeitr, fixeditr, a, [1])
end

prototypical(m::Masked) = zeros(sum(map(i->length(i), m.freeitr)))

function OptimSSM(m::Masked, Z, CHH, T, R, CQQ, a1, P1, P1inf, Pstar, y)
    thetatmp = prototypical(m)
    jcfg = ForwardDiff.GradientConfig(nothing, thetatmp, ForwardDiff.Chunk(thetatmp))
    hcfg = ForwardDiff.HessianConfig(nothing, thetatmp, ForwardDiff.Chunk(thetatmp))
    OptimSSM(m, jcfg, hcfg, Z, CHH, R, T, CQQ,
             convert(Vector,a1),
             convert(Matrix, P1),
             convert(Matrix, P1inf),
             convert(Matrix, Pstar), y)
end

allfixed(x::Masked) = maximum(map(i->length(i), x.freeitr))==0 ? true : false

function remask!(m::Masked, b::NTuple, theta)
    count = m.count
    map(b, m.freeitr) do y,x
        for (i,j) in x
            y[j] = theta[count[1]];
            count[1] += 1
        end
    end
    count[1] = 1
    map(b, m.fixeditr, m.parent) do y, x, z
        for (i,j) in x
            @inbounds y[j] = z[j]
        end
    end
    nothing
end

function islowertriangular(x::AbstractMatrix{T}) where T
    r, c = size(x)
    out = true
    for j in 2:c
        for i in 1:j-1
            out = x[i,j] == zero(T)
            if !out
                break
            end
        end
    end
    out
end

@noinline function setupdiffusematrices(a1::T, P1::M, P1inf::M) where {T<:AbstractVector, M<:AbstractMatrix}
    idx = findall(diag(P1inf) .> 0)
    aa = copy(convert(Vector, a1))
    PP = copy(convert(Matrix, P1))
    RR = Matrix(1.0*I, size(P1, 1), size(P1, 1))

    map(idx) do i
        PP[i,:] .= 0.0
        PP[:,i] .= 0.0
        RR[:,i] .= 0.0
        RR[i,:] .= 0.0
        aa[i]   = 0.0
    end

    a1′ = convert(typeof(a1), aa)
    P1′ = convert(typeof(P1), PP)
    Pstar = convert(typeof(P1), RR*RR')
    (a1′::T, P1′::M, Pstar::M)
end

function solve_discrete_lyapunov(A::AbstractMatrix{T},
    B::AbstractMatrix{T}, max_it::Int=50) where T<:Real
    # TODO: Implement Bartels-Stewardt
    n = size(A, 2)
    alpha0 = reshape([A;], n, n)
    gamma0 = reshape([B;], n, n)

    alpha1 = similar(alpha0)
    gamma1 = similar(gamma0)
    gamma  = similar(gamma0)
    diff = 5
    n_its = 1

    while diff > 1e-15
        A_mul_B!(alpha1, alpha0, alpha0)
        A_mul_B!(gamma1, alpha0, gamma0)
        A_mul_Bt!(gamma, gamma1, alpha0)
        gamma1 .= gamma0 .+ gamma
        #gamma1 = gamma0 + alpha0*gamma0*alpha0'
        diff = maximum(abs, gamma)
        copy!(alpha0, alpha1)
        copy!(gamma0, gamma1)
        n_its += 1
        if n_its > max_it
        error("Exceeded maximum iterations, check input matrices")
        end
    end

    return gamma1
end

function Base.show(io::IO, ::MIME"text/plain", ss::LinearStateSpace)
    id = size(ss)
    print(io, "Linear State Space model, $id")
end


text:
module SimulationSummary

using Statistics
using Serialization
using Printf
using Parameters
using TOML
using CSV

using ..SimulationParams: SimulationParameters
using ..Models: AbstractBond, SSHBond, Bond
using ..Models: AbstractModel, SSHModel, HolsteinModel, write_phonons!
using ..Models: write_phonons!, write_M_matrix!
using ..Measurements: measure_κ
using ..MuFinder: MuTuner, update_μ!, estimate_μ
using ..Utilities: reshaped

export write_simulation_summary!

"""
Write results of simulation to simulation summary file.
"""
function write_simulation_summary!(datafolder::String,Nbins::Int=10)

    # unpack last checkpoint from simulation
    chkpnt = deserialize(joinpath(datafolder,"checkpoint.jls"))
    @unpack model, μ_tuner, sim_stats = chkpnt

    # data foldername
    foldername = splitdir(datafolder)[end]

    # read in config file
    files    = readdir(datafolder)
    tomls    = findall(f -> endswith(f, r"\.toml|\.TOML"), files)
    tomlfile = joinpath(datafolder, files[tomls[1]])
    input    = TOML.parsefile(tomlfile)

    # write final phonon config to file
    write_phonons!(model,joinpath(datafolder,"$(foldername)_config.out"))

    # write M matrix to file
    if haskey(input["simulation"],"write_M_matrix")
        if input["simulation"]["write_M_matrix"]
            write_M_matrix!(model, joinpath(datafolder,"$(foldername)_matrix.out"))
        end
    end

    # construct filename for summary file
    filename = joinpath(datafolder, "$(foldername)_summary.out")

    # open summary file
    open(filename,"w") do fout

        # copy contents of input file to output file
        write(fout,"#########################\n")
        write(fout,"## INPUT FILE CONTENTS ##\n")
        write(fout,"#########################\n\n")

        TOML.print(fout, input)
        write(fout,"\n")

        # write bond defintions to summary file
        write(fout,"######################\n")
        write(fout,"## BOND DEFINITIONS ##\n")
        write(fout,"######################\n\n")

        write_bond_definitions!(fout,model)

        # write phonon definitions
        write(fout,"########################\n")
        write(fout,"## PHONON DEFINITIONS ##\n")
        write(fout,"########################\n\n")

        write_phonon_definitions!(fout,model)

        # write chemical potential defintions
        write(fout,"#########################\n")
        write(fout,"## CHEMICAL POTENTIALS ##\n")
        write(fout,"#########################\n\n")

        write_chemical_potential_definitions!(fout,model)

        # write simulation statistics to file
        write(fout, "#####################","\n")
        write(fout, "## SIMULATION INFO ##","\n")
        write(fout, "#####################","\n","\n")

        total_time = sim_stats["simulation_time"] + sim_stats["measurement_time"] + sim_stats["write_time"]
        @printf fout "Total Time (min)        = %.8f\n" total_time
        @printf fout "Simulation Time (min)   = %.8f\n" sim_stats["simulation_time"]
        @printf fout "Measurement Time (min)  = %.8f\n" sim_stats["measurement_time"]
        @printf fout "Write Time (min)        = %.8f\n" sim_stats["write_time"]
        @printf fout "Iterative Solver Steps  = %.8f\n" sim_stats["iters"]
        @printf fout "Acceptance Rate         = %.8f\n" sim_stats["acceptance_rate"]
        @printf fout "Reflect Acceptance Rate = %.8f\n" sim_stats["reflect_acceptance_rate"]
        @printf fout "Swap Acceptance Rate    = %.8f\n" sim_stats["swap_acceptance_rate"]

        # write global measurements to file
        write(fout,"\n")
        write(fout, "#########################","\n")
        write(fout, "## GLOBAL MEASUREMENTS ##","\n")
        write(fout, "#########################","\n","\n")

        write_global_measurements!(fout,model,μ_tuner,datafolder,Nbins)

        # write global measurements to file
        write(fout, "\n")
        write(fout, "##########################","\n")
        write(fout, "## ON-SITE MEASUREMENTS ##","\n")
        write(fout, "##########################","\n","\n")

        write_onsite_measurements!(fout,model,μ_tuner,datafolder,Nbins)

        # write global measurements to file
        write(fout,"\n")
        write(fout, "#############################","\n")
        write(fout, "## INTER-SITE MEASUREMENTS ##","\n")
        write(fout, "#############################","\n","\n")

        write_intersite_measurements!(fout,model,datafolder,Nbins)

        # write susceptibilities
        write(fout,"\n")
        write(fout, "######################","\n")
        write(fout, "## SUSCEPTIBILITIES ##","\n")
        write(fout, "######################","\n","\n")

        write_susceptibilities!(fout,model,datafolder,Nbins)

        # write on-site correlations
        write(fout,"\n")
        write(fout, "##################","\n")
        write(fout, "## CORRELATIONS ##","\n")
        write(fout, "##################","\n","\n")

        write_correlations!(fout,model,datafolder,Nbins)
    end

    return nothing
end

######################
## PRIVATE FUNCTION ##
######################

"""
Write bond defintions to file.
"""
function write_bond_definitions!(fout,model::SSHModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    for (id, bond) in enumerate(bonds)

        write(fout,"Bond ID       = ", string(id), "\n")
        write(fout,"name          = ",bond.name, "\n")
        write(fout,"t_avg         = ", string(bond.t),  "\n")
        write(fout,"t_std         = ", string(bond.σt), "\n")
        write(fout,"Initial Orbit = ", string(bond.o₁), "\n")
        write(fout,"Final Orbit   = ", string(bond.o₂), "\n")
        write(fout,"Displacement  = ", string(bond.v),  "\n", "\n")
    end

    return nothing
end

function write_bond_definitions!(fout,model::HolsteinModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    for (id, bond) in enumerate(bonds)

        write(fout,"Bond ID       = ", string(id), "\n")
        write(fout,"t_avg         = ", string(bond.t),  "\n")
        write(fout,"t_std         = ", string(bond.σt), "\n")
        write(fout,"Initial Orbit = ", string(bond.o₁), "\n")
        write(fout,"Final Orbit   = ", string(bond.o₂), "\n")
        write(fout,"Displacement  = ", string(bond.v),  "\n", "\n")
    end

    return nothing
end

"""
Write Phonon definitions to file.
"""
function write_phonon_definitions!(fout,model::SSHModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    id = 0
    for bond in bonds

        if bond.has_phonon
            id += 1
            write(fout,"SSH Phonon ID = ", string(id), "\n")
            write(fout,"name = ",bond.name, "\n")
            write(fout,"t_avg = ", string(bond.t),  "\n")
            write(fout,"t_std = ", string(bond.σt), "\n")
            write(fout,"alpha_avg = ", string(bond.α),  "\n")
            write(fout,"alpha_std = ", string(bond.σα), "\n")
            write(fout,"alpha2_avg = ", string(bond.α₂),  "\n")
            write(fout,"alpha2_std = ", string(bond.σα₂), "\n")
            write(fout,"omega_avg = ", string(bond.ω),  "\n")
            write(fout,"omega_std = ", string(bond.σω), "\n")
            write(fout,"omega4_avg = ", string(bond.ω₄),  "\n")
            write(fout,"omega4_std = ", string(bond.σω₄), "\n")
            write(fout,"Initial Orbit = ", string(bond.o₁), "\n")
            write(fout,"Final Orbit   = ", string(bond.o₂), "\n")
            write(fout,"Displacement  = ", string(bond.v),  "\n", "\n")
        end
    end

    return nothing
end

function write_phonon_definitions!(fout,model::HolsteinModel{T1,T2,T3}) where {T1,T2,T3}

    L₁ = model.lattice.L1
    L₂ = model.lattice.L2
    L₃ = model.lattice.L3
    n  = model.lattice.unit_cell.norbits

    ω  = reshaped(model.ω,(n,L₁,L₂,L₃))
    ω₄ = reshaped(model.ω₄,(n,L₁,L₂,L₃))
    λ  = reshaped(model.λ,(n,L₁,L₂,L₃))
    λ₂ = reshaped(model.λ₂,(n,L₁,L₂,L₃))

    for o in 1:n

        write(fout,"Orbit = ", string(o), "\n")

        vals  = @view ω[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Omega_avg   = ", string(avg), "\n")
        write(fout,"Omega_std   = ", string(stdev), "\n")

        vals  = @view ω₄[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Omega4_avg  = ", string(avg), "\n")
        write(fout,"Omega4_std  = ", string(stdev), "\n")

        vals  = @view λ[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Lambda_avg  = ", string(avg), "\n")
        write(fout,"Lambda_std  = ", string(stdev), "\n")

        vals  = @view λ₂[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Lambda2_avg = ", string(avg), "\n")
        write(fout,"Lambda2_std = ", string(stdev), "\n","\n")
    end

    return nothing
end

"""
Write chemical potential defintions to file.
"""
function write_chemical_potential_definitions!(fout,model::AbstractModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    L₁ = model.lattice.L1
    L₂ = model.lattice.L2
    L₃ = model.lattice.L3
    n  = model.lattice.unit_cell.norbits

    μ  = reshaped(model.μ,(n,L₁,L₂,L₃))

    for o in 1:n

        vals  = @view μ[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Orbit  = ", string(o), "\n")
        write(fout,"Mu_avg = ", string(avg), "\n")
        write(fout,"Mu_std = ", string(stdev), "\n", "\n")
    end

    return nothing
end

"""
Write global measurements to file.
"""
function write_global_measurements!(fout, model::AbstractModel{T1,T2,T3}, μ_tuner::MuTuner{T1}, datafolder::String, Nbins::Int) where {T1,T2,T3}

    # file statistics will be written to
    statsfn = joinpath(datafolder, "global_measurements_stats.out")

    # data folder
    datafolder = joinpath(datafolder,"global_measurements_f")

    # get files in data folder
    files = readdir(datafolder,join=true)

    # number of data files
    Nmeas = length(files)

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas%Nbins==0

    # get measurements
    measurements = String[]
    open(files[1],"r") do fin
        for line in readlines(fin)
            atoms = split(line)
            push!( measurements , atoms[1] )
        end
    end

    # container to hold binned data
    binned_data = Dict(string(k)=>zeros(T1,Nbins) for k in measurements)

    # container to contain statistics
    stats = Dict(string(k)=>zeros(T1,2) for k in measurements)

    # iterate over files in data folder
    for (i,file) in enumerate(files)
        # open file
        open(file,"r") do fin
            # iterate over each line in data file
            for line in eachline(fin)
                # split the line at white space
                atoms = split(line)
                # get the bin
                bin   = div(i-1,N)+1
                # record measurement
                binned_data[atoms[1]][bin] += parse(T1,atoms[2]) / N
            end
        end
    end

    # calculate average and error bar for each measurement
    for measurement in keys(binned_data)

        # deal with chemical potential seperately due to μ-tuning algorithm
        if measurement=="mu"
            estimate_μ(μ_tuner)
            stats[measurement][1] = μ_tuner.μ_avg
            stats[measurement][2] = μ_tuner.μ_err
        else
            avg, err = mean_and_error(binned_data[measurement])
            stats[measurement][1] = avg
            stats[measurement][2] = err
        end
    end

    # calculate the compressibility
    β   = model.β
    N   = model.Nsites
    n   = stats["density"][1]
    Δn  = stats["density"][2]
    N²  = stats["Nsqr"][1]
    ΔN² = stats["Nsqr"][2]
    κ, Δκ = measure_κ(β,N,N²,ΔN²,n,Δn)
    stats["compressibility"] = [κ,Δκ]

    # write averaged statistics to file
    open(statsfn,"w") do statsout

        # header
        header = "global_meas avg error\n"
        write(fout,    header)
        write(statsout,header)

        # iterate over measurements
        for measurement in keys(stats)

            # calculate mean and standard deviation of mean
            avg = stats[measurement][1]
            err = stats[measurement][2]

            # write statistics to file
            @printf fout     "%s %.8f %.8f\n" measurement avg err
            @printf statsout "%s %.8f %.8f\n" measurement avg err
        end
    end

    return nothing
end

"""
Write on-site measreuments to file.
"""
function write_onsite_measurements!(fout,model::AbstractModel{T1,T2,T3},μ_tuner::MuTuner{T1},datafolder::String,Nbins::Int) where {T1,T2,T3}

    # number of site per unit cell
    nₛ = model.lattice.unit_cell.norbits::Int

    # data folder
    folder = joinpath(datafolder, "onsite_measurements_f")

    # get files containing data
    files = readdir(folder,join=true)

    # number of data files
    Nmeas = length(files)

    # file statistics will be written to
    statsfn = joinpath(datafolder, "onsite_measurements_stats.out")

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas%Nbins==0

    # get measurements
    measurements = String[]
    open(files[1],"r") do fin
        header = readline(fin)
        for line in readlines(fin)
            push!( measurements , split(line)[1] )
        end
    end
    measurements = unique(measurements)

    # container to hold binned data
    binned_data = Dict(string(k)=>zeros(T1,Nbins,nₛ) for k in measurements)

    # container to contain statistics
    stats = Dict(string(k)=>zeros(T1,2,nₛ) for k in measurements)

    # iterate over data files
    for (i,file) in enumerate(files)
        bin = div(i-1,N) + 1
        open(file,"r") do fin
            header = readline(fin)
            for line in eachline(fin)
                atoms = split(line)
                meas  = atoms[1]
                orbit = parse(Int,atoms[2])
                val   = parse(T1,atoms[3])
                binned_data[meas][bin,orbit] += val / N
            end
        end
    end

    # iterating over measurements
    for measurement in keys(stats)

        # deal with chemical potential seperately due to μ-tuning algorithm
        if measurement=="mu"
            μ  = reshaped(model.μ,(nₛ,div(model.Nsites,nₛ)))
            μ″ = mean(μ)
            μ′ = μ_tuner.μ_avg
            Δμ = μ_tuner.μ_err
            dμ = μ′-μ″
            # calculating average and error for on-site energy of each orbital
            for orbit in 1:nₛ
                μ₀ = @view μ[orbit,:]
                stats[measurement][1,orbit] = mean(μ₀) + dμ
                stats[measurement][2,orbit] = Δμ
            end
        else
            # iterate over orbits
            for orbit in 1:nₛ

                # calculating average and error of measurement
                data     = @view binned_data[measurement][:,orbit]
                avg, err = mean_and_error(data)
                stats[measurement][1,orbit] = avg
                stats[measurement][2,orbit] = err
            end
        end
    end

    # write averaged statistics to file
    open(statsfn,"w") do statsout

        # construct header
        header = "measurement orbit avg error\n"
        write(fout,    header)
        write(statsout,header)

        # iterate over measurements
        for measurement in keys(stats)

            # iterate over orbits
            for orbit in 1:nₛ

                # write out measurement
                avg  = stats[measurement][1,orbit]
                err  = stats[measurement][2,orbit]
                @printf fout     "%s %d %.6f %.6f\n" measurement orbit avg err
                @printf statsout "%s %d %.6f %.6f\n" measurement orbit avg err
            end
        end
    end

    return nothing
end

"""
Write on-site measreuments to file.
"""
function write_intersite_measurements!(fout,model::AbstractModel{T1,T2,T3},datafolder::String,Nbins::Int) where {T1,T2,T3}

    # number
    nᵥ = model.nbonds

    # data folder
    folder = joinpath(datafolder, "intersite_measurements_f")

    # get files containing data
    files = readdir(folder,join=true)

    # number of data files
    Nmeas = length(files)

    # file statistics will be written to
    statsfn = joinpath(datafolder, "intersite_measurements_stats.out")

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas%Nbins==0

    # get measurements
    measurements = String[]
    open(files[1],"r") do fin
        header = readline(fin)
        for line in readlines(fin)
            push!( measurements , split(line)[1] )
        end
    end
    measurements = unique(measurements)

    # container to hold binned data
    binned_data = Dict(string(k)=>zeros(T1,Nbins,nᵥ) for k in measurements)

    # container to contain statistics
    stats = Dict(string(k)=>zeros(T1,2,nᵥ) for k in measurements)

    # iterate over data files
    for (i,file) in enumerate(files)
        bin = div(i-1,N) + 1
        open(file,"r") do fin
            header = readline(fin)
            for line in eachline(fin)
                atoms = split(line)
                meas  = atoms[1]
                bond  = parse(Int,atoms[2])
                val   = parse(T1,atoms[3])
                binned_data[meas][bin,bond] += val / N
            end
        end
    end

    # iterating over measurements
    for measurement in keys(stats)

        # iterate over orbits
        for vector in 1:nᵥ

            # calculating average and error of measurement
            data     = @view binned_data[measurement][:,vector]
            avg, err = mean_and_error(data)
            stats[measurement][1,vector] = avg
            stats[measurement][2,vector] = err
        end
    end

    # write averaged statistics to file
    open(statsfn,"w") do statsout

        # construct header
        header = "measurement bond avg error\n"
        write(fout,    header)
        write(statsout,header)

        # iterate over measurements
        for measurement in keys(stats)

            # iterate over orbits
            for b in 1:nᵥ

                # write out measurement
                avg  = stats[measurement][1,b]
                err  = stats[measurement][2,b]
                @printf fout     "%s %d %.6f %.6f\n" measurement b avg err
                @printf statsout "%s %d %.6f %.6f\n" measurement b avg err
            end
        end
    end

    return nothing
end

"""
Write correlations to file.
"""
function write_correlations!(fout,model::AbstractModel,datafolder::String,Nbins::Int)

    # get contents of data folder
    contents = readdir(datafolder)
    # iterate of names in data folder
    for name in contents
        # if name corresponds to folder
        if endswith(name,"_f")
            # if name of directory starts with capital letter
            if isuppercase(name[1])
                # split name by underscore
                atoms = split(name,"_")
                # if a correlation function
                if !endswith(atoms[1],"Susc")
                    # name of measurement
                    measurement = string(atoms[1])
                    # position or momentum space correlation function
                    space = string(atoms[2])
                    # process data
                    write_correlation!(fout,model,measurement,space,Nbins,datafolder)
                end
            end
        end
    end

    return nothing
end

"""
Write a correlation to file.
"""
function write_correlation!(fout,model::AbstractModel{T},measurement::String,space::String,Nbins::Int,datafolder::String) where {T<:AbstractFloat}

    # data folder
    folder = joinpath(datafolder,"$(measurement)_$(space)_f")
    files  = readdir(folder,join=true)

    # read info in from key file
    Lₜ      = 1
    pairs   = Vector{Tuple{Int,Int}}()
    header  = ""
    keyfile = joinpath(folder,"$(measurement)_$(space)_key.out")
    open(keyfile,"r") do fin
        # construct header for stats file, determining whether to label first two colmns
        # "orbit1 orbit2" or "bond1 bond2" depending on contents of key file
        keyheader = readline(fin)
        atoms     = split(keyheader)
        if space == "position"
            header = "$(atoms[2]) $(atoms[3]) r3 r2 r1 tau $(measurement) error\n"
        else
            header = "$(atoms[2]) $(atoms[3]) k3 k2 k1 tau $(measurement) error\n"
        end
        # get unique pairs of orbitals/bonds the correlation function was measured for
        # and determine if measurement was time dependent
        for line in eachline(fin)
            atoms = split(line)
            p     = ( parse(Int,atoms[2]) , parse(Int,atoms[3]) )
            if !(p in pairs)
                push!(pairs,p)
            end
            τ  = parse(Int,atoms[7])
            Lₜ = max(τ+1,Lₜ)
        end
    end

    # get number of pairs of orbitals/bonds
    nₚ = length(pairs)

    # get lattice size
    L₁ = model.lattice.L1::Int
    L₂ = model.lattice.L2::Int
    L₃ = model.lattice.L3::Int
    V  = L₁*L₂*L₃*nₚ*Lₜ

    # number of measurements
    Nmeas = length(files)-1

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas % Nbins == 0

    # file statistics will be written to
    statsfn = joinpath(datafolder, "$(measurement)_$(space)_stats.out")

    # declare container for binned data
    container  = zeros(T,Nbins,Lₜ,L₁,L₂,L₃,nₚ)
    container′ = reshaped(container,Nbins,V)

    # iterate over number of files
    for m in 1:Nmeas
        # get bin number
        bin = div(m-1,N) + 1
        # iterate of rows in current data file
        for row in CSV.Rows(files[m], header=1, types=[Int,T], delim=" ", reusebuffer=true)
            # calculate binned averages
            container′[bin,row[1]] += row[2] / N
        end
    end

    # write stats file
    open(statsfn,"w") do statsout
        # write header
        write(fout,     header)
        write(statsout, header)
        # counter to indicate what chunk of data file to read
        n = 0
        # iterate of all position or momentum space variable
        for p in 1:nₚ
            n₁ = pairs[p][1]
            n₂ = pairs[p][2]
            for l₃ in 0:L₃-1
                for l₂ in 0:L₂-1
                    for l₁ in 0:L₁-1
                        # write averaged stats to file
                        for τ in 0:Lₜ-1
                            data     = @view container[:,τ+1,l₁+1,l₂+1,l₃+1,p]
                            avg, err = mean_and_error(data)
                            @printf fout     "%d %d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ τ avg err
                            @printf statsout "%d %d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ τ avg err
                        end
                        # increment to next chunk of file
                        n += 1
                    end
                end
            end
        end
    end

    write(fout,"\n")

    return nothing
end

"""
Write susceptibilities to file.
"""
function write_susceptibilities!(fout,model::AbstractModel, datafolder::String,Nbins::Int)

    # get contents of data folder
    contents = readdir(datafolder)
    # iterate of names in data folder
    for name in contents
        # if name corresponds to folder
        if endswith(name,"_f")
            # if name of directory starts with capital letter
            if isuppercase(name[1])
                # split name by underscore
                atoms = split(name,"_")
                # if a correlation function
                if endswith(atoms[1],"Susc")
                    # name of measurement
                    measurement = string(atoms[1])
                    # position or momentum space correlation function
                    space = string(atoms[2])
                    # process data
                    write_susceptibility!(fout,model,measurement,space,Nbins,datafolder)
                end
            end
        end
    end

    return nothing
end

"""
Write a susceptibility to file.
"""
function write_susceptibility!(fout,model::AbstractModel{T},measurement::String,space::String,Nbins::Int,datafolder::String) where {T<:AbstractFloat}

    # data folder
    folder = joinpath(datafolder,"$(measurement)_$(space)_f")
    files  = readdir(folder,join=true)

    # read info in from key file]
    pairs   = Vector{Tuple{Int,Int}}()
    header  = ""
    keyfile = joinpath(folder,"$(measurement)_$(space)_key.out")
    open(keyfile,"r") do fin
        # construct header for stats file, determining whether to label first two colmns
        # "orbit1 orbit2" or "bond1 bond2" depending on contents of key file
        keyheader = readline(fin)
        atoms     = split(keyheader)
        if space == "position"
            header = "$(atoms[2]) $(atoms[3]) r3 r2 r1 $(measurement) error\n"
        else
            header = "$(atoms[2]) $(atoms[3]) k3 k2 k1 $(measurement) error\n"
        end
        # get unique pairs of orbitals/bonds the correlation function was measured for
        for line in eachline(fin)
            atoms = split(line)
            p     = ( parse(Int,atoms[2]) , parse(Int,atoms[3]) )
            if !(p in pairs)
                push!(pairs,p)
            end
        end
    end

    # get number of pairs of orbitals/bonds
    nₚ = length(pairs)

    # get lattice size
    L₁ = model.lattice.L1::Int
    L₂ = model.lattice.L2::Int
    L₃ = model.lattice.L3::Int
    V  = L₁*L₂*L₃*nₚ

    # number of measurements
    Nmeas = length(files)-1

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas % Nbins == 0

    # file statistics will be written to
    statsfn = joinpath(datafolder, "$(measurement)_$(space)_stats.out")

    # container to hold binned data
    container  = zeros(T,Nbins,L₁,L₂,L₃,nₚ)
    container′ = reshaped(container,Nbins,V)

    # iterate over files
    for m in 1:Nmeas
        # calculate bin
        bin = div(m-1,N) + 1
        # iterate over rows of csv file
        for row in CSV.Rows(files[m], types=[Int,T], reusebuffer=true, delim=' ')
            index = row[1]
            val   = row[2]
            container′[bin,index] += val / N
        end
    end

    # write avg and error of measurements to file
    open(statsfn,"w") do statsout
        write(statsout,header)
        write(fout,    header)
        for p in 1:nₚ
            n₁ = pairs[p][1]
            n₂ = pairs[p][2]
            for l₃ in 0:L₃-1
                for l₂ in 0:L₂-1
                    for l₁ in 0:L₁-1
                        vals     = @view container[:,l₁+1,l₂+1,l₃+1,p]
                        avg, err = mean_and_error(vals)
                        @printf fout     "%d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ avg err
                        @printf statsout "%d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ avg err
                    end
                end
            end
        end
    end

    write(fout,"\n")

    return nothing
end

"""
Calculate mean and standard deviation of the mean of a set of data.
"""
function mean_and_error(v::AbstractArray{T})::Tuple{T,T} where {T<:Number}

    N   = length(v)
    avg = mean(v)
    err = std(v,corrected=true,mean=avg)/sqrt(N)
    return avg, err
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/ElPhDynamics

max_forks_repo_name: cohensbw/ElPhDynamics