text:
using LinearAlgebra: norm, transpose

"""
	tree_approximation!(newtree::Tree, path::Function, nIterations::Int64, p::Int64=2, r::Int64=2)

Returns a valuated probability scenario tree approximating the input stochastic process.

Args:
- *newtree* - Tree with a certain branching structure,

- *path* - function generating samples from the stochastic process to be approximated,

- *nIterations* - number of iterations for stochastic approximation procedure,

- *p* - choice of norm (default p = 2 (Euclidean distance)), and,

- *r* - transportation distance parameter
"""
function tree_approximation!(newtree::Tree, path::Function, nIterations::Int64, p::Int64=2, r::Int64=2)
    leaf, omegas, probaLeaf = leaves(newtree)      # leaves, indexes and probabilities of the leaves of the tree
    dm = size(newtree.state, 2)                    # dm = dimension of the states of the nodes of the tree.
    T = height(newtree)                            # height of the tree = number of stages - 1
    n = length(leaf)                               # number of leaves = no of omegas
    d = zeros(Float64, dm, length(leaf))
    samplepath = zeros(Float64, T+1, dm)           # T + 1 = the number of stages in the tree.
    probaLeaf = zero(probaLeaf)
    probaNode = nodes(newtree)                                # all nodes of the tree
    path_to_leaves = [root(newtree, i) for i in leaf]         # all the paths from root to the leaves
    path_to_all_nodes = [root(newtree, j) for j in probaNode] # all paths to other nodes
    for k = 1 : nIterations
	if (rem(k,10)==0)
		print("Progress: $(round(k/nIterations*100,digits=2))%   \r")
		flush(stdout)
	end
        critical = max(0.0, 0.2 * sqrt(k) - 0.1 * n)
        #tmp = findall(xi -> xi <= critical, probaLeaf)
        tmp = Int64[inx for (inx, ppf) in enumerate(probaLeaf) if ppf <= critical]
        samplepath .= path()  # a new trajectory to update the values on the nodes
        #The following part addresses the critical probabilities of the tree so that we don't loose the branches
        if !isempty(tmp) && !iszero(tmp)
            probaNode = zero(probaNode)
            probaNode[leaf] = probaLeaf
            for i = leaf
                while newtree.parent[i] > 0
                    probaNode[newtree.parent[i]] = probaNode[newtree.parent[i]] + probaNode[i]
                    i = newtree.parent[i]
                end
            end
            for tmpi = tmp
                rt = path_to_leaves[tmpi]
                #tmpi = findall(pnt -> pnt <= critical, probaNode[rt])
                tmpi = Int64[ind for (ind, pnt) in enumerate(probaNode[rt]) if pnt <= critical]
                newtree.state[rt[tmpi],:] .= samplepath[tmpi,:]
            end
        end
        #To the step  of STOCHASTIC COMPUTATIONS
        endleaf = 0 #start from the root
        for t = 1 : T+1
            tmpleaves = newtree.children[endleaf + 1]
            disttemp = Inf #or fill(Inf,dm)
            for i = tmpleaves
                dist = norm(view(samplepath, 1 : t) - view(newtree.state, path_to_all_nodes[i]), p)
                if dist < disttemp
                    disttemp = dist
                    endleaf = i
                end
            end
        end
        #istar = findall(lf -> lf == endleaf, leaf)
        istar = Int64[idx for (idx, lf) in enumerate(leaf) if lf == endleaf]
        probaLeaf[istar] .+= 1.0                                                            #counter  of probabilities
        StPath = path_to_leaves[endleaf - (leaf[1] - 1)]
        delta = newtree.state[StPath,:] - samplepath
        d[:,istar] .+= norm(delta, p).^(r)
        delta .=  r .* norm(delta, p).^(r - p) .* abs.(delta)^(p - 1) .* sign.(delta)
        ak = 1.0 ./ (30.0 .+ probaLeaf[istar]) #.^ 0.75        # step size function - sequence for convergence
        newtree.state[StPath,:] = newtree.state[StPath,:] - delta .* ak
    end
    probabilities  = map(plf -> plf / sum(probaLeaf), probaLeaf) #divide every element by the sum of all elements
    t_dist = (d * hcat(probabilities) / nIterations) .^ (1 / r)
    newtree.name = "$(newtree.name) with d=$(t_dist) at $(nIterations) iterations"
    newtree.probability .= build_probabilities!(newtree, hcat(probabilities)) #build the probabilities of this tree
    return newtree
end


text:
# UNION clause.

mutable struct UnionClause <: AbstractSQLClause
    over::Union{SQLClause, Nothing}
    all::Bool
    args::Vector{SQLClause}

    UnionClause(;
               over = nothing,
               all = false,
               args) =
        new(over, all, args)
end

UnionClause(args...; over = nothing, all = false) =
    UnionClause(over = over, all = all, args = SQLClause[args...])

"""
    UNION(; over = nothing, all = false, args)
    UNION(args...; over = nothing, all = false)

A `UNION` clause.

# Examples

```jldoctest
julia> c = FROM(:measurement) |>
           SELECT(:person_id, :date => :measurement_date) |>
           UNION(all = true,
                 FROM(:observation) |>
                 SELECT(:person_id, :date => :observation_date));

julia> print(render(c))
SELECT
  "person_id",
  "measurement_date" AS "date"
FROM "measurement"
UNION ALL
SELECT
  "person_id",
  "observation_date" AS "date"
FROM "observation"
```
"""
UNION(args...; kws...) =
    UnionClause(args...; kws...) |> SQLClause

dissect(scr::Symbol, ::typeof(UNION), pats::Vector{Any}) =
    dissect(scr, UnionClause, pats)

function PrettyPrinting.quoteof(c::UnionClause, ctx::QuoteContext)
    ex = Expr(:call, nameof(UNION))
    if c.all !== false
        push!(ex.args, Expr(:kw, :all, c.all))
    end
    if isempty(c.args)
        push!(ex.args, Expr(:kw, :args, Expr(:vect)))
    else
        append!(ex.args, quoteof(c.args, ctx))
    end
    if c.over !== nothing
        ex = Expr(:call, :|>, quoteof(c.over, ctx), ex)
    end
    ex
end

rebase(c::UnionClause, c′) =
    UnionClause(over = rebase(c.over, c′), args = c.args, all = c.all)

ext: jl


text:
"""
    @do block :while condition

Control flow statement that executes in a **local** scope a `block` of code at least once,
and then repeatedly executes the `block` or not, depending on a given boolean `condition`
after the `:while` symbol at the end of the `block`.`.

# Examples
```julia
julia> Pkg.clone("https://github.com/Ismael-VC/DoWhile.jl.git")

julia> using DoWhile

julia> i = 0
0

julia> @do begin
           @show i
           i += 1
       end :while i < 5
i = 0
i = 1
i = 2
i = 3
i = 4

julia> @do (
           @show i;
           i += 1
       ) :while i < 5
i = 5

julia> i
6

julia> @do (@show i; i += 1) :while i ≤ 10
i = 6
i = 7
i = 8
i = 9
i = 10

julia> @do begin @show i; i += 1 end :while i ≤ 10
i = 11

julia> i
12
```

# Help
```julia
help?> DoWhile

?help?> @do

julia> @doc DoWhile

julia> @doc @do
```
"""
module DoWhile

export @do

let
    do_while = quote
        macro fo(block, sym, cond)
            s = sym.args[1]
            s != :while && error("@do expected :while symbol, got :$s")

            quote
                let
                    $(esc(block))
                    while $(esc(cond))
                        $(esc(block))
                    end
                end
            end
        end
    end

    do_while.args[2].args[1].args[1] = :do

    eval(do_while)
end

@doc (@doc DoWhile) :(@do)

end # module


text:
struct HierarchicalPeriodicHMM{F,T} <: AbstractHMM{F}
    a::Vector{T}
    A::Array{T,3}
    B::Array{<:Distribution{F},4}
    HierarchicalPeriodicHMM{F,T}(a, A, B) where {F,T} = assert_hierarchicalperiodichmm(a, A, B) && new(a, A, B)
end

HierarchicalPeriodicHMM(
    a::AbstractVector{T},
    A::AbstractArray{T,3},
    B::AbstractArray{<:Distribution{F},4},
) where {F,T} = HierarchicalPeriodicHMM{F,T}(a, A, B)

HierarchicalPeriodicHMM(A::AbstractArray{T,3}, B::AbstractArray{<:Distribution{F},4}) where {F,T} =
    HierarchicalPeriodicHMM{F,T}(ones(size(A, 1)) ./ size(A, 1), A, B)

function assert_hierarchicalperiodichmm(a::AbstractVector, A::AbstractArray{T,3} where {T}, B::AbstractArray{<:Distribution,4})
    @argcheck isprobvec(a)
    @argcheck istransmats(A)
    @argcheck size(A, 3) == size(B, 2) ArgumentError("Period length must be the same for transition matrix and distribution")
    @argcheck length(a) == size(A, 1) == size(B, 1) ArgumentError("Number of transition rates must match length of chain")
    return true
end

function assert_hierarchicalperiodichmm(hmm::HierarchicalPeriodicHMM)
    assert_periodichmm(hmm.a, hmm.A, hmm.B)
end

size(hmm::HierarchicalPeriodicHMM, dim = :) = (size(hmm.B, 1), size(hmm.B, 3), size(hmm.B, 2), size(hmm.B, 4))[dim]
# K, D, T, size_memory
copy(hmm::HierarchicalPeriodicHMM) = HierarchicalPeriodicHMM(copy(hmm.a), copy(hmm.A), copy(hmm.B))

function rand_test(
    rng::AbstractRNG,
    hmm::HierarchicalPeriodicHMM,
    n2t::AbstractArray{Int},
    useless;
    init = rand(rng, Categorical(hmm.a)),
    seq = false,
    kwargs...
)
    T = size(hmm.B, 2)
    N = length(n2t)

    z = Vector{Int}(undef, N)
    (T >= 1) && (z[1] = init)
    for n = 2:N
        tm1 = n2t[n-1] # periodic t-1
        z[n] = rand(rng, Categorical(hmm.A[z[n-1], :, tm1]))
    end
    y = rand(rng, hmm, n2t, z, useless; kwargs...)
    seq ? (z, y) : y
end

function rand_test(
    rng::AbstractRNG,
    hmm::HierarchicalPeriodicHMM,
    n2t::AbstractArray{Int},
    z::AbstractVector{<:Integer},
    yini;
    size_memory = size(hmm, 4)
)
    T = size(hmm, 2)
    max_size_memory = maximum(size_memory)
    y = Matrix{Bool}(undef, length(z), T)
    memory = Int.(log.(size_memory) / log(2))
    D = size(y, 2)

    @argcheck D == size(hmm, 2)
    @argcheck length(n2t) == length(z)

    p = zeros(D)
    for j = 1:D
        if memory[j] > 0
            # One could do some specialized for each value of memory e.g. for memory = 1, we have simply previous_day_cat = y[n-1,:].+1
            N_ini = length(yini[1:memory[j], j])
            y[1:N_ini, j] = yini[1:memory[j], j]

            for n in eachindex(z)[N_ini+1:end]
                t = n2t[n] # periodic t
                previous_day_cat = bin2digit([y[n-m, j] for m = 1:memory[j]])

                p = succprob(hmm.B[z[n], t, j, previous_day_cat])
                y[n, j] = rand(Bernoulli(p))
            end
        else
            for n in eachindex(z)[1:end]
                t = n2t[n] # periodic t
                p = succprob(hmm.B[z[n], t, j, 1])
                y[n, j] = rand(Bernoulli(p))
            end
        end
    end
    y
end

function rand(
    rng::AbstractRNG,
    hmm::HierarchicalPeriodicHMM,
    n2t::AbstractArray{Int},
    useless;
    init = rand(rng, Categorical(hmm.a)),
    seq = false
)
    T = size(hmm.B, 2)
    N = length(n2t)

    z = Vector{Int}(undef, N)
    (T >= 1) && (z[1] = init)
    for n = 2:N
        tm1 = n2t[n-1] # periodic t-1
        z[n] = rand(rng, Categorical(hmm.A[z[n-1], :, tm1]))
    end
    y = rand(rng, hmm, n2t, z, useless)
    seq ? (z, y) : y
end

function rand(
    rng::AbstractRNG,
    hmm::HierarchicalPeriodicHMM,
    n2t::AbstractArray{Int},
    z::AbstractVector{<:Integer},
    yini
)
    T = size(hmm, 2)
    y = Matrix{Bool}(undef, length(z), T)
    memory = Int(log(size(hmm, 4)) / log(2))
    D = size(y, 2)

    @argcheck D == size(hmm, 2)
    @argcheck length(n2t) == length(z)

    p = zeros(D)
    if memory > 0
        # One could do some specialized for each value of memory e.g. for memory = 1, we have simply previous_day_cat = y[n-1,:].+1
        N_ini = size(yini, 1)
        @argcheck N_ini == memory # Did we gave the correct number of initial conditions
        size(yini, 1) == D
        y[1:N_ini, :] = yini
        previous_day_cat = zeros(Int, D)
        for n in eachindex(z)[N_ini+1:end]
            t = n2t[n] # periodic t
            previous_day_cat[:] = bin2digit.(eachcol([y[n-m, j] for m = 1:memory, j = 1:D]))
            p[:] = succprob.(hmm.B[CartesianIndex.(z[n], t, 1:D, previous_day_cat)])
            y[n, :] = rand(Product(Bernoulli.(p)))
        end
    else
        for n in eachindex(z)[1:end]
            t = n2t[n] # periodic t
            p[:] = succprob.(hmm.B[CartesianIndex.(z[n], t, 1:D, 1)])
            y[n, :] = rand(Product(Bernoulli.(p)))
        end
    end
    y
end

rand(hmm::HierarchicalPeriodicHMM, n2t::AbstractArray{Int}, useless; kwargs...) = rand(GLOBAL_RNG, hmm, n2t, useless; kwargs...)

rand(hmm::HierarchicalPeriodicHMM, n2t::AbstractArray{Int}, z::AbstractVector{<:Integer}, useless; kwargs...) = rand(GLOBAL_RNG, hmm, n2t, z, useless; kwargs...)

rand_test(hmm::HierarchicalPeriodicHMM, n2t::AbstractArray{Int}, useless; kwargs...) = rand(GLOBAL_RNG, hmm, n2t, useless; kwargs...)

rand_test(hmm::HierarchicalPeriodicHMM, n2t::AbstractArray{Int}, z::AbstractVector{<:Integer}, useless; kwargs...) = rand(GLOBAL_RNG, hmm, n2t, z, useless; kwargs...)

function likelihoods!(L::AbstractMatrix, hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}, lag_cat::Matrix{Int})
    N, K, T, D = size(observations, 1), size(hmm, 1), size(hmm, 3), size(hmm, 2)
    @argcheck size(L) == (N, K)

    @inbounds for i in OneTo(K), n in OneTo(N)
        t = n2t[n] # periodic t
        L[n, i] = pdf(product_distribution(hmm.B[CartesianIndex.(i, t, 1:D, lag_cat[n, :])]), observations[n, :])
    end
end

function loglikelihoods!(LL::AbstractMatrix, hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}, lag_cat::Matrix{Int})
    N, K, T, D = size(observations, 1), size(hmm, 1), size(hmm, 3), size(hmm, 2)
    @argcheck size(LL) == (N, K)

    @inbounds for i in OneTo(K), n in OneTo(N)
        t = n2t[n] # periodic t
        LL[n, i] = logpdf(product_distribution(hmm.B[CartesianIndex.(i, t, 1:D, lag_cat[n, :])]), observations[n, :])
    end
end

####

# # In-place update of the observations distributions.
#
# quikest and most generic version -> to use
function update_B!(B::AbstractArray{T,4} where {T}, γ::AbstractMatrix, observations, estimator, idx_tj::Matrix{Vector{Vector{Int}}})
    @argcheck size(γ, 1) == size(observations, 1)
    @argcheck size(γ, 2) == size(B, 1)
    N = size(γ, 1)
    K = size(B, 1)
    T = size(B, 2)
    D = size(B, 3)
    size_memory = size(B, 4)
    ## For periodicHMM only the n observations corresponding to B(t) are used to update B(t)

    @inbounds for t in OneTo(T)
        for i in OneTo(K)
            for j = 1:D
                for m = 1:size_memory
                    if sum(γ[idx_tj[t, j][m], i]) > 0
                        B[i, t, j, m] = fit_mle(Bernoulli, observations[idx_tj[t, j][m], j], γ[idx_tj[t, j][m], i])
                    else
                        B[i, t, j, m] = Bernoulli(eps())
                    end
                end
            end
        end
    end
end

function fit_mle(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}; init = :none, kwargs...)
    hmm = copy(hmm)

    if init == :kmeans
        kmeans_init!(hmm, observations, display = get(kwargs, :display, :none))
    end

    history = fit_mle!(hmm, observations, n2t; kwargs...)
    hmm, history
end

function fit_mle!(
    hmm::HierarchicalPeriodicHMM,
    observations,
    n2t::AbstractArray{Int}
    ;
    display = :none,
    maxiter = 100,
    tol = 1e-3,
    robust = false,
    estimator = fit_mle)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K, T, size_memory = size(observations, 1), size(hmm, 1), size(hmm, 3), size(hmm, 4)
    @argcheck T == size(hmm.B, 2)
    history = EMHistory(false, 0, [])

    # n2t = date # dayofyear_Leap.(date)

    # Allocate memory for in-place updates
    c = zeros(N)
    α = zeros(N, K)
    β = zeros(N, K)
    γ = zeros(N, K)
    ξ = zeros(N, K, K)
    LL = zeros(N, K)

    # assign category for observation depending in the past observations
    memory = Int(log(size_memory) / log(2))
    lag_cat = conditional_to(observations, memory)
    idx_tj = idx_observation_of_past_cat(lag_cat, n2t, T, K, size_memory)

    loglikelihoods!(LL, hmm, observations, n2t, lag_cat)
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

    forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
    backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
    posteriors!(γ, α, β)

    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter
        update_a!(hmm.a, α, β)
        update_A!(hmm.A, ξ, α, β, LL, n2t)
        update_B!(hmm.B, γ, observations, estimator, idx_tj)

        # Ensure the "connected-ness" of the states,
        # this prevents case where there is no transitions
        # between two extremely likely observations.
        robust && (hmm.A .+= eps())

        @check isprobvec(hmm.a)
        @check istransmats(hmm.A)

        # loglikelihoods!(LL, hmm, observations, n2t)
        loglikelihoods!(LL, hmm, observations, n2t, lag_cat)

        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

        forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
        backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
        posteriors!(γ, α, β)

        logtotp = sum(c)
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    history
end

## For AbstractHMM + n2t + lag_cat

function loglikelihoods(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1])
    (logl !== nothing) && deprecate_kwargs("logl")
    N, K, size_memory = size(observations, 1), size(hmm, 1), size(hmm, 4)
    LL = Matrix{Float64}(undef, N, K)

    memory = Int(log(size_memory) / log(2))
    lag_cat = conditional_to(observations, memory; past = past)

    loglikelihoods!(LL, hmm, observations, n2t, lag_cat)
    if robust
        replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    end
    LL
end

function forward(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1])
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations, n2t; robust = robust, past = past)
    forward(hmm.a, hmm.A, LL, n2t)
end

function backward(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1])
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations; robust = robust, past = past)
    backward(hmm.a, hmm.A, LL, n2t)
end

function posteriors(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1])
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations, n2t; robust = robust, past = past)
    posteriors(hmm.a, hmm.A, LL, n2t)
end

function viterbi(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1])
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations, n2t; robust = robust, past = past)
    viterbi(hmm.a, hmm.A, LL, n2t)
end


text:
function pilot_sampling(num_PilotSamples)
    shift_array_SpinUp = Vector{Float64}()
    shift_array_SpinDn = Vector{Float64}()
    output_file = open("result","a+")
    for N = 1 : num_PilotSamples
        eff_spec_SpinUp, eff_spec_SpinDn = matrix_product()
        Z_up, Z_dn = recursion_partition_function(eff_spec_SpinUp, eff_spec_SpinDn)
        write(output_file, string(Z_up[num_SpinUp] * Z_dn[num_SpinDn]), "\n")
    end
    close(output_file)
end

"Functions for Importance Sampling"

"function roulette_wheel_generator(γ, num_sites)
    roulette_wheel = Vector{Float64}()
    cumulative_array = Vector{Float64}()
    W = 0 # stands for total weight
    for i = 0 : num_sites
        selection_length = exp( -2 * i * γ + 2 * (num_sites - i) * γ)
        W += selection_length
        push!(roulette_wheel, selection_length)
        push!(cumulative_array, W)
    end
    roulette_wheel /= W
    cumulative_array /= W
    return W, roulette_wheel, cumulative_array
end

function roulette_selection(r, W, cumulative_array) # Binary Search
    L = length(cumulative_array)
    i_index, f_index = 1, L
    p_index = floor(Int64, (i_index + f_index) / 2) # pointer index
    while true
        if cumulative_array[p_index] <= r <= cumulative_array[p_index + 1]
            break
        elseif p_index == 1 && r <= cumulative_array[p_index]
            return p_index
        end
        if r < cumulative_array[p_index]
            f_index = p_index
            p_index = floor(Int64, (i_index + f_index) / 2)
        else
            i_index = p_index
            p_index = floor(Int64, (i_index + f_index) / 2)
        end
    end
    p_index + 1
end"

function σlistGenerator(num_sites)
    σlist = Array{Int64,1}[]
    for i = 0 : num_sites
        σarray = 2 * vcat(zeros(Int16, i), ones(Int16, num_sites - i)) .- 1
        push!(σlist, σarray)
    end
    σlist
end

function σfieldGenerator(k, σlist, num_sites) # i takes value from 0 to num_sites
    σfield = Random.shuffle(σlist[k + 1])
    WeightPerStep = (num_sites + 1) * binomial(num_sites, k) / 2 ^ num_sites
    σfield, WeightPerStep
end

max_stars_repo_name: TongSericus/CanonicalAFQMC

max_issues_repo_name: TongSericus/CanonicalAFQMC

max_forks_repo_name: TongSericus/CanonicalAFQMC


text:
@testset "WriteBuffer" begin
    buf = BSONWriteBuffer()
    @test length(buf) == 0
    @test size(buf) == (0,)
    @test sizeof(buf) == 0
    push!(buf, 0x1)
    @test length(buf) == 1
    @test size(buf) == (1,)
    @test sizeof(buf) == 1
    @test GC.@preserve buf unsafe_load(pointer(buf)) == 0x1
    resize!(buf, 5)
    @test length(buf) == 5
    @test buf[1] == 0x1
    @test_throws BoundsError buf[6]
    empty!(buf)
    @test length(buf) == 0

    buf = BSONWriteBuffer()
    sizehint!(buf, 10)
    @test length(buf) == 0
    @test length(buf.data) == 10
    resize!(buf, 10)
    @test length(buf) == 10
    @test length(buf.data) == 10
    push!(buf, 0x2)
    @test length(buf) == 11
    @test length(buf.data) > 11
end


text:
using StaticArrays

const MAT_INFO_KEY = Int

struct MatHist
    h::SVector{2,Int}
end

function Base.length(h::MatHist)
    l = 0
    for a in h.h
        iszero(a) && break
        l += 1
    end
    return l
end

struct IIEMatrixGame{T} <: Game{MatHist, MAT_INFO_KEY}
    R::Matrix{NTuple{2,T}}
end

IIEMatrixGame(g::MatrixGame) = IIEMatrixGame(g.R)

IIEMatrixGame() = IIEMatrixGame([
    (0,0) (-1,1) (1,-1);
    (1,-1) (0,0) (-1,1);
    (-1,1) (1,-1) (0,0)
])

CounterfactualRegret.initialhist(::IIEMatrixGame) = MatHist(SA[0,0])

CounterfactualRegret.isterminal(::IIEMatrixGame, h::MatHist) = length(h) > 1

function CounterfactualRegret.utility(game::IIEMatrixGame, i::Int, h::MatHist)
    length(h) > 1 ? game.R[h.h...][i] : 0
end

CounterfactualRegret.player(::IIEMatrixGame, h::MatHist) = length(h)+1

CounterfactualRegret.player(::IIEMatrixGame, k::MAT_INFO_KEY) = k+1

function CounterfactualRegret.next_hist(::IIEMatrixGame, h::MatHist, a)
    l = length(h)
    return MatHist(setindex(h.h, a, l+1))
end

CounterfactualRegret.infokey(::IIEMatrixGame, h) = length(h)

function CounterfactualRegret.actions(game::IIEMatrixGame, h::MatHist)
    iszero(length(h)) ? (1:size(game.R,1)) : (1:size(game.R,2))
end

## extras

function Base.print(io::IO, solver::AbstractCFRSolver{K,G}) where {K,G<:IIEMatrixGame}
    println(io)
    for (k,v) in solver.I
        σ = copy(v.s)
        σ ./= sum(σ)
        σ = round.(σ, digits=3)
        println(io, "Player: $(k) \t σ: $σ")
    end
end

function cumulative_strategies(hist::Vector{Vector{Float64}})
    Lσ = length(hist[1])
    mat = Matrix{Float64}(undef, length(hist), Lσ)
    σ = zeros(Float64, Lσ)

    for (i,σ_i) in enumerate(hist)
        σ = σ + (σ_i - σ)/i
        mat[i,:] .= σ
    end
    return mat
end

@recipe function f(sol::AbstractCFRSolver{K,G}) where {K,G <: IIEMatrixGame}
    layout --> 2
    link := :both
    framestyle := [:axes :axes]

    xlabel := "Training Steps"

    L1 = length(sol.I[0].σ)
    labels1 = Matrix{String}(undef, 1, L1)
    for i in eachindex(labels1); labels1[i] = L"a_{%$(i)}"; end

    @series begin
        subplot := 1
        ylabel := "Strategy"
        title := "Player 1"
        labels := labels1
        reduce(hcat,sol.I[0].hist)'
    end

    L2 = length(sol.I[1].σ)
    labels2 = Matrix{String}(undef, 1, L2)
    for i in eachindex(labels2); labels2[i] = L"a_{%$(i)}"; end

    @series begin
        subplot := 2
        title := "Player 2"
        labels := labels2
        reduce(hcat,sol.I[1].hist)'
    end
end

max_stars_repo_name: WhiffleFish/HelloCFR

max_issues_repo_name: WhiffleFish/HelloCFR

max_forks_repo_name: WhiffleFish/HelloCFR


text:
function parse_input(dirname::String)
    f = open(dirname)
    lines = readlines()
    @assert lines[1] == "problemData"
    chainlength = split(lines[2], ',')
    @assert chainlength[1] == "maxChainLength"
    if chainlength[2] == "Infinity"
        max_chain_length = Inf64
    else
        max_chain_length = Int64(chainlength[2])
    end
    close(f)
    return max_chain_length
end


text:
#include("MCMCplot.jl"); traceplot("MCMC_samples_residual_variance.txt","plotly",4); savefig("plot.png");
using DelimitedFiles,Plots,Plots.PlotMeasures,StatsPlots

function traceplot(file,backend="plotly",nplots=4)
#catch errors when no backends are installed
    if backend == "pyplot"
        pyplot(size=(300*nplots,200*nplots))
    elseif backend == "plotly"
        plotly(size=(300*nplots,200*nplots))
    end
    mychain,mylabel=readdlm(file,',',header=true)
    if nplots > length(mylabel)
        nplots=length(mylabel)
    end
    mychain = mychain[:,1:nplots]
    mylabel = mylabel[1:nplots]

    steps = 1:size(mychain,1)
    h1=plot(mychain, layout=(nplots,1),title= reshape(mylabel,1,length(mylabel)),
         label="",title_location=:left,titlefont=12,xlabel="iterations",ylabel = "residual variance",ytickfont=font(8))
    # add horizontal line
    hline!(h1,cumsum(mychain,dims=1)./steps,layout=(nplots,1),label="",color=:red,linewidth=0.1)
    #density plot
    h2=density(mychain,layout= (nplots,1),label="",orientation = :horizontal,xlim =[0,1.5],ylim = [-10,25])
    plot(h1, h2, link = :y, layout = grid(1,2,widths = [0.7,0.3]))
end

max_stars_repo_name: QTL-rocks/Plots

max_issues_repo_name: QTL-rocks/Plots

max_forks_repo_name: QTL-rocks/Plots


text:
function minimize(objective, x0, scheme::QuasiNewton, B0=nothing, options=OptOptions())
    minimize(objective, x0, (scheme, BackTracking()), B0, options)
end
function minimize(objective::T1, x0, approach::Tuple{<:Any, <:LineSearch}, B0=nothing,
                  options::OptOptions=OptOptions(),
                  linesearch::T2 = BackTracking()
                  ) where {T1, T2}

    x, fx, ∇fx, z, fz, ∇fz, B = prepare_variables(objective, approach, x0, copy(x0), B0)

    # first iteration
    z, fz, ∇fz, B, is_converged = iterate(x, fx, ∇fx, z, fz, ∇fz, B, approach, objective, options)

    iter = 0
    while iter <= options.maxiter && !is_converged
        iter += 1

        # take a step and update approximation
        z, fz, ∇fz, B, is_converged = iterate(x, fx, ∇fx, z, fz, ∇fz, B, approach, objective, options, false)
    end

    return z, fz, ∇fz, iter
end

function iterate(x, fx, ∇fx, z, fz, ∇fz, B, approach, objective, options, is_first=nothing)
    # split up the approach into the hessian approximation scheme and line search
    scheme, linesearch = approach

    # Move nexts into currs
    fx = fz
    x = copy(z)
    ∇fx = copy(∇fz)

    # Update current gradient and calculate the search direction
    d = find_direction(B, ∇fx, scheme) # solve Bd = -∇f

    # # Perform line search along d
    α, f_α, ls_success = linesearch(objective, d, x, fx, ∇fx, 1.0)

    # # Calculate final step vector and update the state
    s = @. α * d
    z = @. x + s

    # Update approximation
    fz, ∇fz, B = update_obj(objective, d, s, ∇fx, z, ∇fz, B, scheme, is_first)

    # Check for gradient convergence
    is_converged = converged(z, ∇fz, options.g_tol)

    return z, fz, ∇fz, B, is_converged
end


text:
# parameters
mass_bin_name = ARGS[1]
tslice = "t1"
path_wavelist = "src"
path_to_working_folder = "data"
list_of_files = ARGS[2:end]

println("Starting with ARGS = ", ARGS)
######################################################
using DelimitedFiles
using LinearAlgebra

push!(LOAD_PATH,"src")
using SDMHelper
using FittingPWALikelihood
using amplitudes_compass
using PWAHelper

BmatMC = read_cmatrix(
    joinpath(path_to_working_folder,"integrmat_$(mass_bin_name)_$(tslice)_mc.txt"));

# read wavelist throw waves below threshol
wavelist = get_wavelist(joinpath(path_wavelist,"wavelist_formated.txt");
     path_to_thresholds=joinpath(path_wavelist,"thresholds_formated.txt"),
     M3pi=Meta.parse(mass_bin_name[1:4])/1000)

const noϵ =  [1] # flat wave
const posϵ = [i for (i,ϵ) in enumerate(wavelist[:,6]) if ϵ=="+"]
const negϵ = [i for (i,ϵ) in enumerate(wavelist[:,6]) if ϵ=="-"]

# Model description
const ModelBlocks = [noϵ, posϵ, negϵ, negϵ]

# load precalculated data array
const PsiRD = read_precalc_basis(
    joinpath(path_to_working_folder,"functions_$(mass_bin_name)_$(tslice)_rd.bin"));

# get functions to calculate LLH, derivative and hessian
LLH, GRAD, LLH_and_GRAD!, HES = createLLHandGRAD(PsiRD, BmatMC, ModelBlocks);

println("---> Start calculating LLHs")
llhs = [LLH(vcat(readdlm(path_and_filename)...)) for path_and_filename in list_of_files]
sorted_list = sort(collect(zip(list_of_files, llhs)), by=x->x[2])
writedlm(joinpath(path_to_working_folder,"llh_attmpts_$(mass_bin_name)_$(tslice).txt"), sorted_list)

###########################################################################
println("---> Trying to improve minimum")
# fit the best minimum again
## normalize B
Bscale = [BmatMC[i,i] for i in 1:size(BmatMC,1)];
for i=1:size(BmatMC,1), j=1:size(BmatMC,2)
    BmatMC[i,j] *= 1/sqrt(Bscale[i]*Bscale[j])
end
## normalize Psi
for i in 1:size(PsiRD,2)
    PsiRD[:,i] .*= 1.0/sqrt(Bscale[i])
end
## get functions to calculate LLH, derivative and hessian
LLH, GRAD, LLH_and_GRAD!, HES = createLLHandGRAD(PsiRD, BmatMC, ModelBlocks);
## load the best parameters
best_parameters = vcat(readdlm(sorted_list[1][1])...)
parscale = abs.(extnd(sqrt.(Bscale), get_parameter_map(ModelBlocks, size(BmatMC,1))))
best_parameters_scalled = best_parameters .* parscale
## fit
@time minpars = minimize(LLH, LLH_and_GRAD!;
    algorithm = :LD_LBFGS, verbose=1, starting_pars=best_parameters_scalled,
    llhtolerance = 1e-6)

max_reldiff_of_parameters = max(abs.(1 .- minpars ./ best_parameters_scalled)...)
@show max_reldiff_of_parameters
(max_reldiff_of_parameters > 0.1) && warn("Too big improvement in the fit while fine tuning!")

@time minpars2 = minimize(LLH, LLH_and_GRAD!;
    algorithm = :LD_SLSQP, verbose=1, starting_pars=minpars,
    llhtolerance = 1e-6)

max_reldiff_of_parameters2 = max(abs.(1 .- minpars2 ./ minpars)...)
@show max_reldiff_of_parameters2
(max_reldiff_of_parameters2 > 0.1) && warn("Too big improvement in the fit while fine tuning!")

###########################################################################
println("---> Calculating Hessian")
# get derivative
hes = HES(minpars2)
singular = (det(hes) ≈ 0.0); singular && warn("Hessian matrix is singular!")
invhes = !(singular) ? inv(hes) : fill(0.0,size(hes)...)
invhes_scaled_back = invhes ./ transpose(parscale) ./ parscale
writedlm(joinpath(path_to_working_folder,"invhes_$(splitdir(sorted_list[1][1])[end])"), invhes_scaled_back)
println("Done")


text:
using LightXML

let docstr = """<Students>
      <Student Name="April" Gender="F" DateOfBirth="1989-01-02" />
      <Student Name="Bob" Gender="M"  DateOfBirth="1990-03-04" />
      <Student Name="Chad" Gender="M"  DateOfBirth="1991-05-06" />
      <Student Name="Dave" Gender="M"  DateOfBirth="1992-07-08">
        <Pet Type="dog" Name="Rover" />
      </Student>
      <Student DateOfBirth="1993-09-10" Gender="F" Name="Émily" />
    </Students>"""

    doc = parse_string(docstr)
    xroot = root(doc)
    for elem in xroot["Student"]
        println(attribute(elem, "Name"))
    end
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
Type representing any object drawable on image
"""
abstract type Drawable end

"""
    p = Point(x,y)
    p = Point(c)

A `Drawable` point on the image
"""
struct Point <: Drawable
    x::Int
    y::Int
end

abstract type Line <: Drawable end
abstract type Circle <: Drawable end

"""
    line = LineTwoPoints(p1, p2)

A `Drawable` infinite length line passing through the two points
`p1` and `p2`.
"""
struct LineTwoPoints <: Line
    p1::Point
    p2::Point
end

"""
    line = LineNormal(ρ, θ)

A `Drawable` infinte length line having perpendicular length `ρ` from
origin and angle `θ` between the perpendicular and x-axis

"""
struct LineNormal{T<:Real, U<:Real} <: Line
    ρ::T
    θ::U
end

"""
    circle = CircleThreePoints(p1, p2, p3)

A `Drawable` circle passing through points `p1`, `p2` and `p3`
"""
struct CircleThreePoints <: Circle
    p1::Point
    p2::Point
    p3::Point
end

"""
    circle = CirclePointRadius(center, ρ)

A `Drawable` circle having center `center` and radius `ρ`
"""
struct CirclePointRadius{T<:Real} <: Circle
    center::Point
    ρ::T
end

"""
    ls = LineSegment(p1, p2)

A `Drawable` finite length line between `p1` and `p2`
"""
struct LineSegment <: Drawable
    p1::Point
    p2::Point
end

"""
    path = Path([point])

A `Drawable` sequence of line segments connecting consecutive pairs
of points in `[point]`.
!!! note
    This will create a non-closed path. For a closed path, see `Polygon`
"""
struct Path <: Drawable
    vertices::Vector{Point}
end

"""
    ellipse = Ellipse(center, ρx, ρy)

A `Drawable` ellipse with center `center` and parameters `ρx` and `ρy`

"""
struct Ellipse{T<:Real, U<:Real} <: Drawable
    center::Point
    ρx::T
    ρy::U
end

"""
    polygon = Polygon([vertex])

A `Drawable` polygon i.e. a closed path created by joining the
consecutive points in `[vertex]` along with the first and last point.
!!! note
    This will create a closed path. For a non-closed path, see `Path`
"""
struct Polygon <: Drawable
    vertices::Vector{Point}
end

"""
    rp = RegularPolygon(center, side_count, side_length, θ)

A `Drawable` regular polygon.

#Arguments
* `center::Point` : the center of the polygon
* `side_count::Int` : number of sides of the polygon
* `side_length::Real` : length of each side
* `θ::Real` : orientation of the polygon w.r.t x-axis (in radians)

"""
struct RegularPolygon{T<:Real, U<:Real} <: Drawable
    center::Point
    side_count::Int
    side_length::T
    θ::U
end

"""
    cross = Cross(c, range::UnitRange{Int})
A `Drawable` cross passing through the point `c` with arms ranging across `range`.
"""
struct Cross <: Drawable
    c::Point
    range::UnitRange{Int}
end

"""
    img = draw!(img, drawable, color)
    img = draw!(img, drawable)

Draws `drawable` on `img` using color `color` which
defaults to `oneunit(eltype(img))`
"""
draw!(img::AbstractArray{T,2}, object::Drawable) where {T<:Colorant} = draw!(img, object, oneunit(T))

"""
    img = draw!(img, [drawable], [color])
    img = draw!(img, [drawable] ,color)
    img = draw!(img, [drawable])

Draws all objects in `[drawable]` in the given order on `img` using
corresponding colors from `[color]` which defaults to `oneunit(eltype(img))`
If only a single color `color` is specified then all objects will be
colored with that color.
"""
function draw!(img::AbstractArray{T,2}, objects::AbstractVector{U}, colors::AbstractVector{V}) where {T<:Colorant, U<:Drawable, V<:Colorant}
    colors = copy(colors)
    while length(colors) < length(objects)
        push!(colors, oneunit(T))
    end
    foreach((object, color) -> draw!(img, object, color), objects, colors)
    img
end

draw!(img::AbstractArray{T,2}, objects::AbstractVector{U}, color::T = oneunit(T)) where {T<:Colorant, U<:Drawable} =
    draw!(img, objects, [color for i in 1:length(objects)])

"""
    img_new = draw(img, drawable, color)
    img_new = draw(img, [drawable], [color])

Draws the `drawable` object on a copy of image `img` using color
`color`. Can also draw multiple `Drawable` objects when passed
as a `AbstractVector{Drawable}` with corresponding colors in `[color]`
"""
draw(img::AbstractArray{T,2}, args...) where {T<:Colorant} = draw!(copy(img), args...)

Point(τ::Tuple{Int, Int}) = Point(τ...)
Point(p::CartesianIndex) = Point(p[2], p[1])

function draw!(img::AbstractArray{T,2}, point::Point, color::T) where T<:Colorant
    drawifinbounds!(img, point, color)
end

"""

    img_new = drawifinbounds!(img, y, x, color)
    img_new = drawifinbounds!(img, Point, color)
    img_new = drawifinbounds!(img, CartesianIndex, color)

Draws a single point after checkbounds() for coordinate in the image.
Color Defaults to oneunit(T)

"""

drawifinbounds!(img::AbstractArray{T,2}, p::Point, color::T = oneunit(T)) where {T<:Colorant} = drawifinbounds!(img, p.y, p.x, color)
drawifinbounds!(img::AbstractArray{T,2}, p::CartesianIndex{2}, color::T = oneunit(T)) where {T<:Colorant} = drawifinbounds!(img, Point(p), color)

function drawifinbounds!(img::AbstractArray{T,2}, y::Int, x::Int, color::T) where {T<:Colorant}
    if checkbounds(Bool, img, y, x) img[y, x] = color end
    img
end


text:
using MultivariateStats
using PLSRegressor

const defdir = PLSRegressor.dir("datasets")

function gethousingdata(dir, filename)
    url = "https://archive.ics.uci.edu/ml/machine-learning-databases/housing/housing.data"
    mkpath(dir)
    path = download(url, "$(defdir)/$filename")
end

function loaddata(test=0.1)
    filename = "housing.data"
    file = "$(defdir)/$filename"

    isfile("$(defdir)/$filename") || gethousingdata(defdir, filename)

    data = readdlm(file)

    nfeatures = size(data)[2] - 1
    target_idx = size(data)[2]

    x = data[:, 1:nfeatures]
    y = data[:, target_idx:target_idx]
    if test == 0
        xtrn = xtst = x
        ytrn = ytst = y
    else
        r = randperm(size(x,1))          # trn/tst split
        n = round(Int, (1-test) * size(x,1))
        xtrn=x[r[1:n], :]
        ytrn=y[r[1:n], :]
        xtst=x[r[n+1:end], :]
        ytst=y[r[n+1:end], :]
    end
    (xtrn, [ytrn...], xtst, [ytst...])
end

(xtrn, ytrn, xtst, ytst) = loaddata()

model    = PLSRegressor.fit(xtrn, ytrn, nfactors = 3)
pred     = PLSRegressor.predict(model, xtst)

println("[PLS] mae error :", mean(abs.(ytst .- pred)))

# linear least squares from MultiVariateStats
sol = llsq(xtrn, ytrn)
a, b = sol[1:end-1], sol[end]
yp = xtst * a + b
println("[LLS] mae error :",mean(abs.(ytst .- yp)))

### if you want to save or load model use this
#PLSRegressor.save(model,filename="/tmp/pls_model.jld",modelname="pls_model")
#model = PLSRegressor.load(filename="/tmp/pls_model.jld",modelname="pls_model")


text:
# # A model with homeowners and renters

using Revise
using Test, Aiyagari
using QuantEcon, Parameters, Interpolations
using Plots

u(c; γ) = c^(1-γ) / (1-γ)

function u(c,h; ξ=0.8159, ρ=map(s -> (s-1)/s, 0.13), γ=2.0)
  C = (ξ * h^ρ + (1-ξ) * c^ρ)^(1/ρ)
  u(C, γ=γ)
end

# Exogenous states (incomes)
z_grid = [0.5; 1.0; 1.5]
z_prob = [0.7 0.15 0.15;
          0.2 0.6 0.2;
          0.15 0.15 0.7]
z_MC = MarkovChain(z_prob, z_grid, :z)

# Moving shocks
move_grid = Symbol[:just_moved, :normal, :move]
move_grid = [1, 2, 3]
move_prob = [0.7 0.3 0.0;
             0.0 0.99 0.01;
             1.0 0.0 0.0]
#move_prob = ones(3,3)/3

move_MC = MarkovChain(move_prob, move_grid, :move)

# itp_scheme = BSpline(Cubic(Line(OnGrid())))
# a_grid = LinRange(0.0, 0.7, 50)
#
# val = u.(a_grid .+ permutedims([exo.z for exo in exo1.grid]), γ=2.0)
#
# 𝔼V = extrapolated_𝔼V(a_grid, itp_scheme, val, exo1, 3, Aiyagari.Conditional(:move))
#
# 𝔼V([1.0, 2.0, 0.7])

mutable struct HousingAS{T1,T2,T3,T4} <: AggregateState
  r::T1
  p::T2 # house price
  ρ::T3 # rent
  dist::T4 # the distribution over idiosynchratic states
end

function HousingAS(r, p, a_grid, exo, param; ρ=p * (param.δ + r))
  dist_proto = zeros((length(a_grid), length(exo)))
  HousingAS(r, p, ρ, dist_proto)
end

include("housing-simple-nlopt.jl")
include("housing-nlopt.jl")
include("renting-nlopt.jl")
#include("housing-simple-jump.jl")
r = 0.29

exo = ExogenousStateSpace([z_MC])
exo = ExogenousStateSpace([z_MC, move_MC])

# ## Forever home owners
r_own = 0.29
a_grid_own = LinRange(0.0, 0.7, 50)
h_grid_own = LinRange(eps(), 2, 10)
endo = EndogenousStateSpace((w=a_grid_own, h=h_grid_own))
param_own  = (β = 0.7, θ = 0.9, δ = 0.1, h_thres = eps())
agg_state_own = HousingAS(r_own, 2.2, a_grid_own, exo, param_own)

out_C = solve_bellman(a_grid_own, exo, agg_state_own, param_own, Owner(Aiyagari.Conditional(:move)), tol=1e-5)
out_UC = solve_bellman(a_grid_own, exo, agg_state_own, param_own, Owner(Aiyagari.Unconditional()), tol=1e-5)
out_U = solve_bellman(endo, exo, agg_state_own, param_own, Owner(Aiyagari.Unconditional()), tol=2e-5)

@testset "conditional vs unconditional" begin
  itp_scheme = BSpline(Cubic(Line(OnGrid())))
  V_C = extrapolated_𝔼V(a_grid_own, itp_scheme, out_C.val, exo, 1, Aiyagari.Conditional(:move))
  V_UC = extrapolated_𝔼V(a_grid_own, itp_scheme, out_C.val, exo, 1, Aiyagari.Unconditional())
  V_UC_old = extrapolated_𝔼V(a_grid_own, itp_scheme, out_U.val, exo_old, 1, Aiyagari.Unconditional())

  V_C_p(a) = ForwardDiff.derivative(V_C, a)
  V_UC_p(a) = ForwardDiff.derivative(V_UC, a)
  V_UC_old_p(a) = ForwardDiff.derivative(V_UC_old, a)

  @test all(V_C.(a_grid_own) .≈ V_UC.(a_grid_own))
  @test all(V_C_p.(a_grid_own) .≈ V_UC_p.(a_grid_own))

  @show maximum(abs, V_UC.(a_grid_own) .- V_UC_old.(a_grid_own))
  @show maximum(abs, V_UC_p.(a_grid_own) .- V_UC_old_p.(a_grid_own))

  using Distributions
  a_rand = rand(Uniform(0, 0.7), 500)
  a_low = rand(Uniform(-5, 0), 500)
  a_high = rand(Uniform(0.7, 5), 500)

  @test all(V_C.(a_rand) .≈ V_UC.(a_rand))
  @test all(V_C.(a_low) .≈ V_UC.(a_low))
  @test all(V_C.(a_high) .≈ V_UC.(a_high))
end

using DelimitedFiles

@testset "regression test simple housing" begin
  #writedlm("test/matrices/housing_simple_nlopt_value.txt", val)
  value_test = readdlm("test/matrices/housing_simple_nlopt_value.txt")

  for i in 1:length(z_grid)
    Δ = maximum(abs, value_test[:,i] .- Aiyagari.get_cond_𝔼V(val, exo, i, :z => i))
    @test abs(Δ) < 1e-6
  end
end

#@show all(value_test .≈ val) || maximum(abs, value_test .- val)

hh = reshape(out_U.policies_full.h, (size(endo)..., length(exo)))

plot(a_grid_own, hh, title="house size", xlab="wealth", legend=:false)
plot(h_grid_own, reshape(permutedims(hh, (2,1,3)), (10,150)), title="house size", xlab="wealth", legend=:false, alpha=0.3)
 ylims!(0,2)
#

#plot(a_grid_own, policies_full.m, title="mortgage", xlab="wealth")

dist = stationary_distribution(z_MC, a_grid_own, policies_full.w_next)

#writedlm("test/matrices/housing_simple_nlopt_dist.txt", dist)
dist_test = readdlm("test/matrices/housing_simple_nlopt_dist.txt")
all(dist_test .≈ dist) || maximum(abs, dist_test .- dist)

plot(a_grid_own, dist, xlab="wealth" )

# ## Forever renters
r_rent = 0.15

a_grid_rent = LinRange(-0.05, 1.5, 50)
param_rent = (β = 0.7, θ = 0.9, δ = 0.1, h_thres = Inf)
agg_state_rent = HousingAS(r_rent, 2.2, a_grid_rent, exo, param_rent, ρ=2.2 * (param_rent.δ + r))

@unpack val, policy, policies_full = solve_bellman(a_grid_rent, exo, agg_state_rent, param_rent, Renter(Aiyagari.Unconditional()), maxiter=100)

plot(a_grid_rent, Aiyagari.get_cond_𝔼V(val, exo, 1, :z => 1))
 plot!(a_grid_rent, Aiyagari.get_cond_𝔼V(val, exo, 2, :z => 2))
 plot!(a_grid_rent, Aiyagari.get_cond_𝔼V(val, exo, 3, :z => 3))

plot(a_grid_rent, val)

plot(a_grid_rent, policies_full.w_next, title="house size", xlab="wealth", legend=:topleft)

# ### Stationary distribution

dist = stationary_distribution(exo.mc, a_grid_rent, policies_full.w_next)
plot(a_grid_rent, dist, xlab="wealth" )

# ## Own big, rent small

r = 0.10
a_grid = LinRange(0.0, 1.0, 40)
param_both = (β = 0.7, θ = 0.9, δ = 0.1, h_thres = 0.75)
param = [param_own, param_rent]
agg_state_both = HousingAS(r, 2.2, a_grid, exo, param_both, ρ= 1.07 * 2.2 * (param_both.δ + r))

 @unpack val, policy, policies_full, owner = solve_bellman(a_grid, exo, agg_state_both, param, OwnOrRent(), maxiter=70)

   plot(a_grid, owner, title="Who owns?")

# using StructArrays
# move = StructArray(exo.grid).move
# move_long = repeat(permutedims(move), 40, 1)

w_next_all = policies_full[1].w_next .* owner .+ policies_full[2].w_next .* .!owner
 h_all = policies_full[1].h .* owner .+ policies_full[2].h .* .!owner
 c_all = policies_full[1].c .* owner .+ policies_full[2].c .* .!owner

 scatter(a_grid, h_all .* (owner .== 1), color=:blue, legend=:false, title="House size", alpha=0.3, markerstrokewidth=0, xlab="wealth")
 scatter!(a_grid, h_all .* (owner .== 0), color=:red, alpha=0.3, markerstrokewidth=0)
 ylims!(0.3, 2.25)
 title!("no moving costs (blue == own, red == rent)")

savefig("no-moving-costs.png")
hline!([param_both.h_thres], color=:gray, label="", linestyle=:dash)

plot(a_grid, w_next_all, legend=false, title="wealth")

#

#
# plot(a_grid, w_next_all, legend=:left, title="cash-at-hand")
#   #plot!(a_grid, a_grid)
#   #plot!(a_grid[[1;end]], a_grid[[end;end]], legend=false)
#
# using Plots
# using DelimitedFiles
# #writedlm("test/matrices/housing_simple_nlopt_value.txt", val)
# value_test = readdlm("test/matrices/housing_simple_nlopt_value.txt")
#
# @test all(val .== value_test)
#
# using Plots
# plot(val)
#
# scatter(a_grid, policies_full.w_next)
# scatter(a_grid, policies_full[1].h)
# scatter!(a_grid, policies_full[2].h)

# scatter(a_grid, policies_full[1].m .* owner)
# scatter(a_grid, c_all)
#
# all(policies_full.conv)

# ## Stationary distribution

dist = stationary_distribution(exo.mc, a_grid, w_next_all)
plot(a_grid, dist, xlab="wealth" )

# #writedlm("test/matrices/housing_simple_nlopt_dist.txt", dist)
# dist_test = readdlm("test/matrices/housing_simple_nlopt_dist.txt")
# @test all(dist .== dist_test)
#
# using StatsBase
# mean(vec(policies_full.m), Weights(vec(dist)))
# mean(vec(policies_full.h), Weights(vec(dist)))
# #926 μs
#


text:
function test(args::Dict)

  is_sorted = args["--sorted"]
  is_sorted |= args["-s"]

  test(is_sorted)

end

function test(is_sorted=false)

  scratch_file = "$(dirname(@__FILE__))"
  scratch_file *= "/../../../"
  scratch_file *= "tmp/scratch.jl"

  open(scratch_file, "w") do f
    write(f, "tests_are_sorted = $is_sorted \n")
  end

  target_name = bump()

  Pkg.test(target_name)

  open(scratch_file, "w") do f
    write(f, "")
  end

end


text:
# generally, Documenter.jl assumes markdown files ends with `.md`
const markdown_exts = [".md",]
const markdown_footer = raw"""

---

*This page was generated using [DemoCards.jl](https://github.com/johnnychen94/DemoCards.jl).*

"""

"""
    struct MarkdownDemoCard <: AbstractDemoCard
    MarkdownDemoCard(path::String)

Constructs a markdown-format demo card from existing markdown file `path`.

# Fields

Besides `path`, this struct has some other fields:

* `path`: path to the source markdown file
* `cover`: path to the cover image
* `id`: cross-reference id
* `title`: one-line description of the demo card
* `author`: author(s) of this demo.
* `date`: the update date of this demo.
* `description`: multi-line description of the demo card
* `hidden`: whether this card is shown in the generated index page

# Configuration

You can pass additional information by adding a YAML front matter to the markdown file.
Supported items are:

* `cover`: an URL or a relative path to the cover image. If not specified, it will use the first available image link, or all-white image if there's no image links.
* `description`: a multi-line description to this file, will be displayed when the demo card is hovered. By default it uses `title`.
* `id`: specify the `id` tag for cross-references. By default it's infered from the filename, e.g., `simple_demo` from `simple demo.md`.
* `title`: one-line description to this file, will be displayed under the cover image. By default, it's the name of the file (without extension).
* `author`: author name. If there are multiple authors, split them with semicolon `;`.
* `date`: any string contents that can be passed to `Dates.DateTime`. For example, `2020-09-13`.
* `hidden`: whether this card is shown in the layout of index page. The default value is `false`.

An example of the front matter:

```text
---
title: passing extra information
cover: cover.png
id: non_ambiguious_id
author: Jane Doe; John Roe
date: 2020-01-31
description: this demo shows how you can pass extra demo information to DemoCards package. All these are optional.
hidden: false
---
```

See also: [`JuliaDemoCard`](@ref DemoCards.JuliaDemoCard), [`DemoSection`](@ref DemoCards.DemoSection), [`DemoPage`](@ref DemoCards.DemoPage)
"""
mutable struct MarkdownDemoCard <: AbstractDemoCard
    path::String
    cover::Union{String, Nothing}
    id::String
    title::String
    description::String
    author::String
    date::DateTime
    hidden::Bool
end

function MarkdownDemoCard(path::String)::MarkdownDemoCard
    # first consturct an incomplete democard, and then load the config
    card = MarkdownDemoCard(path, "", "", "", "", "", DateTime(0), false)

    config = parse(card)
    card.cover = load_config(card, "cover"; config=config)
    card.title = load_config(card, "title"; config=config)
    card.date = load_config(card, "date"; config=config)
    card.author = load_config(card, "author"; config=config)

    # Unlike JuliaDemoCard, Markdown card doesn't accept `julia` compat field. This is because we
    # generally don't know the markdown processing backend. It might be Documenter, but who knows.
    # More generally, any badges can just be manually added by demo writter, if they want.
    # `date` and `author` fields are added just for convinience.

    # default id requires a title
    card.id    = load_config(card, "id"; config=config)
    # default description requires a title
    card.description = load_config(card, "description"; config=config)
    card.hidden = load_config(card, "hidden"; config=config)
    return card
end

"""
    save_democards(card_dir::String, card::MarkdownDemoCard)

process the original markdown file and save it.

The processing pipeline is:

1. strip the front matter
2. insert a level-1 title and id
"""
function save_democards(card_dir::String,
                        card::MarkdownDemoCard;
                        credit,
                        kwargs...)
    isdir(card_dir) || mkpath(card_dir)

    markdown_path = joinpath(card_dir, basename(card))

    _, _, body = split_frontmatter(read(card.path, String))
    config = parse(Val(:Markdown), body)
    need_header = !haskey(config, "title")
    # @ref syntax: https://juliadocs.github.io/Documenter.jl/stable/man/syntax/#@ref-link-1
    header = need_header ? "# [$(card.title)](@id $(card.id))\n" : "\n"

    footer = credit ? markdown_footer : "\n"
    write(markdown_path, header, make_badges(card)*"\n\n", body, footer)
end


text:
"""
	stroboscopicmap(ds::ContinuousDynamicalSystem, T; kwargs...)  → smap

Return a map (integrator) that produces iterations over a period `T` of the `ds`,
known as a stroboscopic map. See [Integrator API](@ref) for handling integrators.

See also [`poincaremap`](@ref).

## Keyword Arguments
* `u0`: initial state
* `diffeq` is a `NamedTuple` (or `Dict`) of keyword arguments propagated into
  `init` of DifferentialEquations.jl.

## Example
```julia
f = 0.27; ω = 0.1
ds = Systems.duffing(zeros(2); ω, f, d = 0.15, β = -1)
smap = stroboscopicmap(ds, 2π/ω; diffeq = (;reltol = 1e-8))
reinit!(smap, [1.0, 1.0])
u = step!(smap)
u = step!(smap, 4) # step 4 iterations forward
```
"""
function stroboscopicmap(ds::CDS, T; u0 = get_state(ds), diffeq = NamedTuple())
	integ = integrator(ds, u0; diffeq)
	return StroboscopicMap{typeof(integ), dimension(ds), typeof(T)}(integ, T)
end
struct StroboscopicMap{I, D, F} <: GeneralizedDynamicalSystem
	integ::I
	T::F
end
isdiscretetime(::StroboscopicMap) = true
DelayEmbeddings.dimension(::StroboscopicMap{I, D}) where {I, D} = D

integrator(p::StroboscopicMap) = p
function step!(smap::StroboscopicMap)
	step!(smap.integ, smap.T, true)
	return smap.integ.u
end
function step!(smap::StroboscopicMap, n::Int)
	for k in 1:n; step!(smap.integ, smap.T, true); end
	return smap.integ.u
end
function reinit!(smap::StroboscopicMap, u0)
	reinit!(smap.integ, u0)
	return
end
function get_state(smap::StroboscopicMap)
	return smap.integ.u
end

function Base.show(io::IO, smap::StroboscopicMap)
    println(io, "Iterator of the stroboscopic map")
    println(io,  rpad(" rule f: ", 14),     DynamicalSystemsBase.eomstring(smap.integ.f.f))
    println(io,  rpad(" Period: ", 14),     smap.T)
end

current_time(smap::StroboscopicMap) = current_time(smap.integ)
function (smap::StroboscopicMap)(t)
	if t == current_time(smap)
		return get_state(smap)
	else
		error("Can't extrapolate discrete systems!")
	end
end
integrator(pinteg::StroboscopicMap, args...; kwargs...) = pinteg


text:
using JFVM
using Base.Test

# write your own tests here
JFVM_test()
@test 1==1