text:
### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ a7871567-172a-42ec-b889-56c5440c99fa
begin
	using Revise
	using ProgressMeter
end

# ╔═╡ 7fa178d6-9f1c-11eb-3604-9777d8da228e
using Curiosity, Plots, PlutoUI, Statistics, Random

# ╔═╡ 60fa4e34-de7a-43c3-91d6-e4f60f1ca272
include("../experiments/tabular_tmaze.jl")

# ╔═╡ 3d63f4ea-a20a-438d-a359-35a7f84b789e
TTMU = Curiosity.TabularTMazeUtils

# ╔═╡ 6f96b22d-1bbe-4711-b4bd-b9b89130fead
default_args(ϵ) =
    Dict(
        # Behaviour Items
        "behaviour_eta" => 0.04,
        "behaviour_gamma" => 0.9,
        "behaviour_learner" => "GPI",
        "behaviour_update" => "TB",
        "behaviour_trace" => "AccumulatingTraces",
        "behaviour_opt" => "Auto",
        "behaviour_lambda" => 0.9,
        "behaviour_alpha_init" => 0.1,
        "exploration_param" => ϵ,
        "exploration_strategy" => "epsilon_greedy",

        # Demon Attributes
        "demon_alpha_init" => 0.1,
        "demon_eta" => 0.5,
        "demon_discounts" => 0.9,
        "demon_learner" => "SR",
        "demon_update" => "TB",
        "demon_policy_type" => "greedy_to_cumulant",
        "demon_opt" => "Auto",
        "demon_lambda" => 0.9,
        "demon_trace"=> "AccumulatingTraces",
        "demon_beta_m" => 0.9,
        "demon_beta_v" => 0.99,

        # Environment Config
        "constant_target"=> 1.0,
        "cumulant_schedule" => "DrifterDistractor",
        "distractor" => (1.0, 5.0),
        "drifter" => (1.0, sqrt(0.01)),
        "exploring_starts"=>false,

        # Agent and Logger
        "horde_type" => "regular",
        "intrinsic_reward" => "weight_change",
        "logger_keys" => [LoggerKey.TTMAZE_ERROR, LoggerKey.TTMAZE_UNIFORM_ERROR, LoggerKey.TTMAZE_OLD_ERROR],
        "save_dir" => "TabularTMazeExperiment",
        "seed" => 1,
        "steps" => 10000,
        "use_external_reward" => true,
        "logger_interval" => 100,
    )

# ╔═╡ b1573f0a-6622-4a40-b171-24f8afae3505
env = TabularTMaze(true, Curiosity.TMazeCumulantSchedules.Constant(1))

# ╔═╡ c655ab97-c20c-4e62-8435-533d50f47ade
begin
	start!(env)
	plot(env)
end

# ╔═╡ 037b8feb-c1de-4c5e-a8f4-fbd6b762673e
function get_policy_arrows(env,agent)
	#HARDCODED CAUSE IDK WHAT ELSE TO DO :(
	s_x = 7
	s_y = 9
	SIZE = 20
	OFFSET = 1*SIZE
	arrow_x = []
	arrow_y = []
	arrow_u = []
	arrow_v = []
	function generate_obs(state::Array{Int64})
		obs = zeros(Int64, 9, 7)
		obs[state[1],state[2]] = 1
		toReturn = findfirst(!iszero, obs[Curiosity.valid_state_mask()])
		return toReturn
	end
	for j ∈ 1:s_x
		for i ∈ 1:s_y
			sqr_i = ((i)*SIZE + 1):((i+1)*SIZE)
			sqr_j = ((j)*SIZE + 1):((j+1)*SIZE)

			k = generate_obs([i,j])
			if k != nothing # If it is not part of the 21 valid cells
				# @show k,i,j
				obs = zeros(5)
				obs[1] = k
				b_π = Curiosity.μ_π(agent, obs)
				# val, ACTION = findmax(b_π)

				# drawarray is bugged. Need to have it off. This is the correct transformation
				# when yflip=true
				# https://github.com/JuliaPlots/Plots.jl/issues/1498
				x_coord = OFFSET + (j-1) * SIZE + OFFSET/2
				y_coord = s_y * SIZE + 2*OFFSET - i * SIZE - OFFSET/2
				# yflip = false
				# x_coord = s_x * SIZE + 2 * OFFSET - (j) * SIZE - OFFSET/2
				# y_coord = OFFSET + (i-1) * SIZE + OFFSET/2

				max_prob,_ = findmax(b_π)
				for (ACTION,π) in enumerate(b_π)
					if π != max_prob
						continue
					end

					if ACTION == 1
						u,v = 0,1
					elseif ACTION == 2
						u,v, = 1,0
					elseif ACTION == 3
						u,v = 0,-1
					elseif ACTION == 4
						u,v = -1,0
					end
					push!(arrow_x,x_coord)
					push!(arrow_y,y_coord)
					push!(arrow_u, u*SIZE/2)
					push!(arrow_v,v*SIZE/2)
				end
			end
		end
	end
	return arrow_x,arrow_y,arrow_u,arrow_v
end

# ╔═╡ 77dae42d-3be5-43b0-93b0-611fa8c9f123
function main_experiment(parsed=default_args(0.1); progress=false, working=false, show_policy = false)

    num_steps = parsed["steps"]
    Random.seed!(parsed["seed"])

    cumulant_schedule = TTMU.get_cumulant_schedule(parsed)
    exploring_starts = parsed["exploring_starts"]
    env = TabularTMaze(exploring_starts, cumulant_schedule)

    agent = TabularTMazeExperiment.construct_agent(parsed)

    eps = 1
    max_num_steps = num_steps
    steps = Int[]

    if progress
        p = Progress(max_num_steps)
    end
	anim = Animation()
    while sum(steps) < max_num_steps
        cur_step = 0
        max_episode_steps = max_num_steps - sum(steps)
        tr, stp =
            run_episode!(env, agent, max_episode_steps) do (s, a, s_next, r, t)
                #This is a callback for every timestep where logger can go
                # agent is accesible in this scope
				plot(env, title=string("Eps: ", parsed["exploration_param"],
					" Step: ",cur_step,
					" Ep: ", eps,
					" DL: ", parsed["demon_learner"],
					" BL: ", parsed["behaviour_learner"]))
				if show_policy
					arrow_x, arrow_y, arrow_u, arrow_v = get_policy_arrows(env,agent)
					quiver!(arrow_x,arrow_y,quiver=(arrow_u,arrow_v))
				end
				frame(anim)
                cur_step+=1
            end
        push!(steps, stp)
        eps += 1

        if progress
            ProgressMeter.update!(p, sum(steps))
        end
    end
	anim
end

# ╔═╡ 465bd0cd-3204-41a2-a68c-f3e38cdd0b5f
anim = main_experiment(default_args(0.2), show_policy=true, progress=true)

# ╔═╡ 1bacedb7-8373-4d52-99b8-54e1e5460346
# let
# 	for ϵ ∈ [0.1, 0.2, 0.3, 0.4]
# 		anim = main_experiment(default_args(ϵ))
# 		mp4(anim, "GPI_SR_epsilon_$(ϵ).mp4")
# 	end
# end

# ╔═╡ Cell order:
# ╠═a7871567-172a-42ec-b889-56c5440c99fa
# ╠═7fa178d6-9f1c-11eb-3604-9777d8da228e
# ╠═3d63f4ea-a20a-438d-a359-35a7f84b789e
# ╠═60fa4e34-de7a-43c3-91d6-e4f60f1ca272
# ╠═6f96b22d-1bbe-4711-b4bd-b9b89130fead
# ╟─b1573f0a-6622-4a40-b171-24f8afae3505
# ╠═c655ab97-c20c-4e62-8435-533d50f47ade
# ╠═037b8feb-c1de-4c5e-a8f4-fbd6b762673e
# ╟─77dae42d-3be5-43b0-93b0-611fa8c9f123
# ╠═465bd0cd-3204-41a2-a68c-f3e38cdd0b5f
# ╠═1bacedb7-8373-4d52-99b8-54e1e5460346

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
abstract type AbstractModel end

struct Model <: AbstractModel

  function Model()

  end
end

macro model(expr::Expr)
  expr.head != :function && error("Expression is not a function definition!")
  func_def = splitdef(expr)
  dump(func_def)
  nothing
end

ext: jl


text:
# Autogenerated wrapper script for Cares_jll for x86_64-linux-gnu
export acountry, adig, ahost, libcares

JLLWrappers.@generate_wrapper_header("Cares")
JLLWrappers.@declare_executable_product(acountry)
JLLWrappers.@declare_executable_product(adig)
JLLWrappers.@declare_executable_product(ahost)
JLLWrappers.@declare_library_product(libcares, "libcares.so.2")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        acountry,
        "bin/acountry",
    )

    JLLWrappers.@init_executable_product(
        adig,
        "bin/adig",
    )

    JLLWrappers.@init_executable_product(
        ahost,
        "bin/ahost",
    )

    JLLWrappers.@init_library_product(
        libcares,
        "lib/libcares.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
"""
    min_radius(n::Integer, idx::Integer, data::Array{<:Real,2}, target::Array{<:Integer,1}, dist_func::Any)::Float64

Compute minimum raidus of hypersphere centered around each data sample so that each hypersphere contains
at least n samples from same class as the corresponding data sample as well as n samples from a different class.

Author: Jernej Vivod
"""
function min_radius(n::Integer, e_idx::Integer, data::Array{<:Real,2}, target::Array{<:Integer,1}, dist_func::Any)::Float64

    # Get masks for samples with same and different class.
    hit_neigh_mask = target .== target[e_idx]
    miss_neigh_mask = .!hit_neigh_mask

    # Get distances to samples from same and different classes.
    dists_same = dist_func(data[e_idx:e_idx, :], data[hit_neigh_mask, :])
    dists_other = dist_func(data[e_idx:e_idx, :], data[miss_neigh_mask, :])

    # Return computed radius.
    return max(sort(dists_same, dims=1)[n+1], sort(dists_other, dims=1)[n])
end

"""
    iterative_relief(data::Array{<:Real,2}, target::Array{<:Integer,1}, m::Integer=-1, min_incl::Integer=3,
                          max_iter::Integer=100, dist_func::Any=(e1, e2, w) -> sum(abs.(w.*(e1 .- e2)), dims=2);
                          f_type::String="continuous")::Array{Float64,1}

Compute feature weights using Iterative Relief algorithm.

---
# Reference:
- Bruce Draper, Carol Kaito, and Jose Bins. Iterative Relief. Proceedings
CVPR, IEEE Computer Society Conference on Computer Vision and
Pattern Recognition., 6:62 – 62, 2003.
"""
function iterative_relief(data::Array{<:Real,2}, target::Array{<:Integer,1}, m::Integer=-1, min_incl::Integer=3,
                          max_iter::Integer=100, dist_func::Any=(e1, e2, w) -> sum(abs.(w.*(e1 .- e2)), dims=2);
                          f_type::String="continuous")::Array{Float64,1}

    # Initialize distance weights.
    dist_weights = ones(Float64, 1, size(data, 2))

    # Initialize iteration counter, convergence indicator and
    # Array for storing feature weights from previous iteration.
    iter_count = 0
    convergence = false
    feature_weights_prev = zeros(Float64, size(data, 2))

    # Iterate until reached maximum iterations or convergence.
    @inbounds while iter_count < max_iter && !convergence

        # Increment iteration counter.
        iter_count += 1

        # Reset feature weights to zero and sample samples.
        feature_weights = zeros(Float64, size(data, 2))
        sample_idxs = StatsBase.sample(1:size(data, 1), if (m==-1) size(data,1) else m end, replace=false)

        # Set m if currently set to signal value -1.
        if (m == -1) m = size(data, 1) end

        # Go over sampled indices.
        @inbounds for idx in sample_idxs

            # Get next sampled sample.
            e = data[idx:idx, :]

            # Get minimum radius needed to include n samples from same class and n samples from different class.
            min_r = min_radius(min_incl, idx, data, target, (e1, e2) -> dist_func(e1, e2, dist_weights))

            # Compute hypersphere inclusions and distances to examples within the hypersphere.
            # Distances to examples from same class.

            dist_same_all = dist_func(data[(target .== target[idx]) .& (1:length(target) .!= idx), :], e, dist_weights)
            sel = dist_same_all .<= min_r
            dist_same = dist_same_all[sel]
            data_same = data[(target .== target[idx]) .& (1:length(target) .!= idx), :][vec(sel), :]

            # Distances to examples with different class.
            dist_other_all = dist_func(data[target .!= target[idx], :], e, dist_weights)
            sel = dist_other_all .<= min_r
            dist_other = dist_other_all[sel]
            data_other = (data[target .!= target[idx], :])[vec(sel), :]

            ### Weights Update ###

            w_miss = max.(0, 1 .- (dist_other.^2/min_r.^2))
            w_hit = max.(0, 1 .- (dist_same.^2/min_r.^2))

            if f_type == "continuous"
                # If features continuous.

                numerator1 = sum(abs.(e .- data_other) .* w_miss, dims=1)
                denominator1 = sum(w_miss) + eps(Float64)

                numerator2 = sum(abs.(e .- data_same) .* w_hit, dims=1)
                denominator2 = sum(w_hit) + eps(Float64)

                feature_weights .+= vec(numerator1 ./ denominator1 .- numerator2 ./ denominator2)

            elseif f_type == "discrete"
                # If features discrete.

                numerator1 = sum(Int64.(e .!= data_other) .* w_miss, dims=1)
                denominator1 = sum(w_miss) + eps(Float64)

                numerator2 = sum(Int64.(e .!= data_same) .* w_hit, dims=1)
                denominator2 = sum(w_hit) + eps(Float64)

                feature_weights .+= vec(numerator1 ./ denominator1 .- numerator2 ./ denominator2)

            else
                throw(DomainError(f_type, "f_type can only be equal to \"continuous\" or \"discrete\"."))
            end

            ######################

        end

        # Update distance weights by feature weights - use algorithm's own feature evaluations
        # to weight features when computing distances.
        dist_weights += permutedims(feature_weights)

        # Check convergence.
        if sum(abs.(feature_weights .- feature_weights_prev)) < 1.0e-3
            convergence = true
        end

        # Set current feature weights as previous feature weights.
        feature_weights_prev = feature_weights
    end

    # Return computed feature weights.
    return vec(dist_weights)

end


text:
using KMeansPNNSmoothing
using Test
using DelimitedFiles
using Random

# The a3 dataset was downloaded from here:
#
#  http://cs.uef.fi/sipu/datasets/
#
#  Clustering basic benchmark
#  P. Fränti and S. Sieranoja
#  K-means properties on six clustering benchmark datasets
#  Applied Intelligence, 48 (12), 4743-4759, December 2018
#  https://doi.org/10.1007/s10489-018-1238-7

a3 = Matrix(readdlm(joinpath(@__DIR__, "a3.txt"))')
a3 ./= maximum(a3)
n, m = size(a3)
k = 50

Random.seed!(472632723)

@testset "kmeans uniform" begin
    kmseeder = KMeansPNNSmoothing.KMUnif()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 7 < result.cost < 25
    @test result.exit_status == :converged
end

@testset "kmeans++" begin
    kmseeder = KMeansPNNSmoothing.KMPlusPlus()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPlusPlus{1}()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans maxmin" begin
    kmseeder = KMeansPNNSmoothing.KMMaxMin()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans PNN" begin
    kmseeder = KMeansPNNSmoothing.KMPNN()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans||" begin
    kmseeder = KMeansPNNSmoothing.KMScala()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans PNNS(UNIF)" begin
    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMUnif())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMUnif(); rlevel=2)
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMUnif(); rlevel=3)
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans PNNS([G]KM++)" begin
    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMPlusPlus{1}())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMPlusPlus{1}(); rlevel=2)
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMPlusPlus())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMPlusPlus(); rlevel=2)
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans PNNS(MAXMIN)" begin
    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMMaxMin())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMMaxMin(); rlevel=2)
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans PNNS(SCALA)" begin
    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMScala())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged

    kmseeder = KMeansPNNSmoothing.KMPNNS(KMeansPNNSmoothing.KMScala(); rlevel=2)
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans PNNSR" begin
    kmseeder = KMeansPNNSmoothing.KMPNNSR()
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans REFINE(UNIF)" begin
    kmseeder = KMeansPNNSmoothing.KMRefine(KMeansPNNSmoothing.KMUnif())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans REFINE(++)" begin
    kmseeder = KMeansPNNSmoothing.KMRefine(KMeansPNNSmoothing.KMPlusPlus{1}())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans REFINE(MAXMIN)" begin
    kmseeder = KMeansPNNSmoothing.KMRefine(KMeansPNNSmoothing.KMMaxMin())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

@testset "kmeans REFINE(SCALA)" begin
    kmseeder = KMeansPNNSmoothing.KMRefine(KMeansPNNSmoothing.KMScala())
    result = kmeans(a3, k; kmseeder, verbose=false)
    @test length(result.labels) == m
    @test all(∈(1:k), result.labels)
    @test size(result.centroids) == (2,k)
    @test 6.7 < result.cost < 11
    @test result.exit_status == :converged
end

max_stars_repo_name: carlobaldassi/KMeansPNNSmoothing


text:
"Metropolized Independent Sampling"
struct MIAlg <: SamplingAlgorithm end

"Metropolized Independent Sampling"
const MI = MIAlg()

isapproximate(::MIAlg) = true

"Sample `ω | y == true` with Metropolis Hasting"
function Base.rand(y::RandVar,
                   n,
                   alg::MIAlg,
                   ΩT::Type{OT};
                   cb = default_cbs(n),
                   hack = true) where {OT <: Ω}
  ω = ΩT()
  xω_, sb = applytrackerr(y, ω)
  plast = err(sb)
  qlast = 1.0
  ωsamples = ΩT[]
  accepted = 0
  for i = 1:n
    ω_ = ΩT()
    xω_, sb = applytrackerr(y, ω_)
    p_ = err(sb)
    ratio = p_ / plast
    if rand() < ratio
      ω = ω_
      plast = p_
      accepted += 1
    end
    push!(ωsamples, ω)
    cb((ω = ω, accepted = accepted, p = p_, i = i), IterEnd)
  end
  [applynotrackerr(y, ω_) for ω_ in ωsamples]
end


text:
function KapitzaPendulum(;l=1, g=9.8, ax=0, ay=0, omegax=0, omegay=0, phix=0, phiy=0)
    """
    A Kapitza Pendulum function defined by the parameters.

    :param l: length.
    :param g: acceleration due to gravity.

    :params ax, omegax, phix: the x coordinate of the pivot xp(t) = ax cos(omegax t + phix).
    :params ay, omegay, phiy: the y coordinate of the pivot yp(t) = ay cos(omegay t + phiy).

    :return: the function f where \vec{theta}'(t) = f(t, \vec{theta(t)}), x, y.
    """
    function f(t, theta)
        """
        For [theta'(t), theta''(t)] = f(t, [theta(t), theta'(t)]).

        :param t: time.
        :param theta: vector [theta(t), theta'(t)].

        :return: \vec{theta'(t)}.
        """
        th, thdot = theta
        [thdot, ax/l*omegax^2*cos(omegax*t+phix)*cos(th)-(g-ay*omegay^2*cos(omegay*t+phiy))*sin(th)/l]
    end
    function x(t, theta)
        """ x(t) = xp(t) + l sin(theta(t)) """
        ax*cos(omegax*t+phix)+l*sin(theta)
    end
    function y(t, theta)
        """ y(t) = yp(t) - y cos(theta(t)) """
        ay*cos(omegay*t+phiy)-l*cos(theta)
    end

    f, x, y
end

#Default arguments represents simple pendulum.
gfunc, x, y = KapitzaPendulum(ax=.2, omegax=100, ay=.2, omegay=100, phiy=pi/2.0)
t, tmax, h = 0.0, 2, 0.0001
y0 = [pi/4.0, 0.0]
ts, theta = MultiStepper(gfunc, AdamsMoult4, t, y0, tmax, h)

using PyPlot

# Theta plot
figure(0)
plot(ts, [theta[i][1] for i=1:length(theta)], "o-", label="Position")
plot(ts, [theta[i][2] for i=1:length(theta)], "x--", label="Velocity")
xlabel("time t")
ylabel("angle theta (radians)")
title("Kapitza Pendulum")
legend()

# Parametric plot
figure(1)
xaxis = [x(ts[i], theta[i][1]) for i=1:length(ts)]
yaxis = [y(ts[i], theta[i][1]) for i=1:length(ts)]
plot(xaxis, yaxis, "o-", label="Parametric Plot")
title("Kapitza Pendulum")
xlabel("x")
ylabel("y")
legend()

max_stars_repo_name: jiosue/QuadratureCandJulia

max_issues_repo_name: jiosue/QuadratureCandJulia

max_forks_repo_name: jiosue/QuadratureCandJulia


text:
export HilbertSpace

"""
    HilbertSpace{L, S, M}

`M` dimension Hilbert space of Site type `L` in with shape `S`.
"""
struct HilbertSpace{L, Shape, N, NSite}
    cache::Array{L, N}

    function HilbertSpace{L}(dims::Int...) where L
        Shape = Tuple{dims...}
        N = length(dims)
        NSite = prod(dims)
        new{L, Shape, N, NSite}(downs(L, dims))
    end
end

StaticArrays.Size(::Type{HilbertSpace{L, Shape}}) where {L, Shape} = Size(Shape)
Base.@pure Base.size(::Type{<:HilbertSpace{L, S}}) where {L, S} = tuple(S.parameters...)

function carrybit!(a::Array{L}) where L
    @inbounds for i in eachindex(a)
        if a[i] == up(L)
            a[i] = down(L)
        else
            a[i] = up(L)
            break
        end
    end
    a
end

Base.iterate(it::HilbertSpace{L}) where L = (fill!(it.cache, down(L)); iterate(it, 1))

function Base.iterate(it::HilbertSpace{L, S, N, NSite}, state) where {L, S, N, NSite}
    if_stop(it, state) && return nothing
    state == 1 && return (it.cache, state + 1)
    carrybit!(it.cache), state + 1
end

@generated function if_stop(it::HilbertSpace{L, S, N, NSite}, state) where {L, S, N, NSite}
    if sizeof(Int) * 8 > NSite
        :(state == $((1 << NSite) + 1))
    else
        quote
            flag = true
            for each in it.cache
                if each != up($L)
                    flag = false
                    break
                end
            end
            flag
        end
    end
end

Base.eltype(::HilbertSpace{L, S, N}) where {L, S, N} = Array{L, N}
Base.length(::HilbertSpace{L, S, N, NSite}) where {L, S, N, NSite} = 1 << NSite

function Base.collect(space::HilbertSpace)
    r = Vector{eltype(space)}(undef, length(space))
    @inbounds for (i, each) in enumerate(space)
        r[i] = copy(each)
    end
    r
end


text:
# # Replication of Mackowiak and Wiederholt (2009)

# This example replicates [Mackowiak and Wiederholt (2009)](https://www.aeaweb.org/articles?id=10.1257/aer.99.3.769) (henceforth MW) using the [DRIPs](https://github.com/afrouzi/DRIPs.jl) package.

# [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/afrouzi/DRIPs.jl/binder?filepath=examples) to run and modify the following code (no software is needed on the local machine).

# See [Afrouzi and Yang (2020)](http://www.afrouzi.com/dynamic_inattention.pdf) for background on the theory.

# ## Contents
# * [Setup](@ref mw2009_setup)
# * [Mapping the Problem to a DRIP](@ref mw2009_map_drip)
# * [Initialization](@ref mw2009_param)
# * [Functions](@ref mw2009_funcs)
#     * [Solving for the fixed point given $\omega$](@ref mw2009_fpointomega)
#     * [Solving for the optimal $\omega$](@ref mw2009_solveomega)
# * [Figures](@ref mw2009_figures)
#     * [Benchmark Economy](@ref mw2009_fig_benchmark)
#     * [Other values of real rigidity](@ref mw2009_robust_alpha)
#     * [Other values of information capacity](@ref mw2009_robust_kappa)
# * [Measuring Performance/Speed](@ref mw2009_performance)

# ## [Setup](@id mw2009_setup)

# The problem in MW is
# ```math
# \begin{aligned}
#     &\min_{\{\hat{\Delta}_{i,t},\hat{z}_{i,t}\}} \left\{ E\left[(\Delta_t-\hat{\Delta}_{i,t})^2\right]+\underset{\equiv \xi}{\underbrace{\left(\frac{\hat{\pi}_{14}}{\hat{\pi}_{11}}\right)}}^2E\left[(z_{i,t}-\hat{z}_{i,t})^2\right]\right\}, \\
#     s.t.\quad & \mathcal{I}(\{\Delta_t\};\{\hat{\Delta}_{i,t}\})+\mathcal{I}(\{z_{i,t}\};\{\hat{z}_{i,t}\})\leq \kappa, \\
#     & \{\Delta_t,\hat{\Delta}_{i,t}\} \perp \{z_{i,t},\hat{z}_{i,t}\}
# \end{aligned}
# ```
# where
# ```math
# \begin{aligned}
#     \Delta_t&\equiv p_t + \underset{\equiv 1-\alpha}{\underbrace{\left(\frac{|\hat{\pi}_{13}|}{|\hat{\pi}_{11}|}\right)}}(q_t-p_t) \\
#     p_t &= \int_0^1 \hat{\Delta}_{i,t}di \\
#     q_t &= \rho q_{t-1} + \nu_t, \nu_{q,t}\sim \mathcal{N}(0,\sigma_q^2) \\
#     z_{i,t}&= \rho z_{i,t-1} + \nu_{z,t}, \nu_{z,t}\sim \mathcal{N}(0,\sigma_z^2)
# \end{aligned}
# ```

# ## [Mapping the Problem to a DRIP](@id mw2009_map_drip)

# There are a few ways of translating the problem above to a `Drip` structure; however, the most efficient way, due to the independence assumption, is to write it as the sum of two DRIPs: one that solves the attention problem for the idiosyncratic shock, and one that solves the attention problem for the aggregate shock which also has endogenous feedback.

# Moreover, since the problem above has a fixed capacity, instead of a fixed cost of attention ($\omega$) as in DRIPs pacakge, we need to iterate over $\omega$'s to find the one that corresponds with $\kappa$.

# ## [Initialization](@id mw2009_param)
# Include the package::
using DRIPs;

# Assign parameters:

ρ  = 0.95;
σq = 0.01;
σz = 11.8*σq;
κ  = 3;
ξ  = 1;
α  = 1 - 0.15;

# Primitives of Drip:
using LinearAlgebra;
L  = 21; # length of trunction
A  = [zeros(1,L);[Matrix(I,L-1,L-1);zeros(1,L-1)]']; # MW truncate the state space with linear irfs of length 20
Qq = zeros(L,1); Qq[1]=σq;
Qz = zeros(L,1); Qz[1]=σz;
H  = zeros(L,1); H[1:21] = Array(1:-1/20:0);

# ## [Functions](@id mw2009_funcs)
# We start with a function that solves the aggregate problem with feedback for a given $\omega$.
# ### [Solving for the fixed point given $\omega$](@id mw2009_fpointomega)

function agg_drip(ω,A,Qq,        #primitives of drip except for H because H is endogenous
                  α,             #strategic complementarity
                  H;             #state space rep. of q
                  β     = 1,     #optional: discount factor, MW's parameterization implies β = 1
                  H0    = H,     #optional: initial guess for HΔ (H is the true solution when α=0)
                  maxit = 10000, #optional: max number of iterations for GE code
                  tol   = 1e-4,  #optional: tolerance for iterations
                  w     = 1)     #optional: update weight for RI
    errmin= 1;
    err   = 1;
    iter  = 0;
    L     = length(H);
    while (err > tol) & (iter < maxit)
            if iter == 0
                global agg  = Drip(ω,β,A,Qq,H0;w = w);
            else
                global agg  = Drip(ω,β,A,Qq,H0;Ω0 = agg.ss.Ω , Σ0 = agg.ss.Σ_1,w = w);
            end

            XFUN(jj) = ((I-agg.ss.K*agg.ss.Y')*agg.A)^jj * (agg.ss.K*agg.ss.Y') * (agg.A')^jj
            X = DRIPs.infinitesum(XFUN; maxit=200, start = 0);  #E[x⃗]=X×x⃗

            XpFUN(jj) = α^jj * X^(jj)
            Xp = DRIPs.infinitesum(XpFUN; maxit=200, start = 0);

            H1 = (1-α)*Xp'*H;

            err= 0.5*norm(H1-H0,2)/norm(H0)+0.5*err;
            if DRIPs.capacity(agg) < 1e-2 # perturb the initial guess if solution is the zero capacity one
                H0 = H0+rand(L).*(H-H0);
            else # store the solution if it has positive capacity
                H0 = H1;
                if err < errmin
                    global aggmin = agg;
                    errmin = err;
                end
            end
            iter += 1;
    end
    return(aggmin, errmin)
end;

# Now we need a function that iterates over $\omega$'s to find the one that corresponds to a given capacity for the MW problem.
# ### [Solving for the optimal $\omega$](@id mw2009_solveomega)

using Printf;
function MW(κ,α,A,Qq,Qz,Hq,Hz; #primitives of MW problem
            ω     = σq^2,      #optional: initial guess for ω
            tol   = 1e-3,      #optional: tolerance for κ
            maxit = 10000)     #optional: max iterations
    ωs    = [ω; 2*ω];
    caps  = [];
    iter  = 0;
    err   = 1;
    it    = 0;
    while (err > tol) & (iter < maxit)
        agg, errtemp = agg_drip(ω,A,Qq,α,H; H0 = rand(L),maxit=20,w=0.95);
        idi = Drip(ω,1,A,Qz,H,w = 0.9) ;
        cap = DRIPs.capacity(agg, unit = "bit") + DRIPs.capacity(idi, unit = "bit");
        x = ω/σq^2;
        @printf("ω = %.2fσq² for κ = %.2f \n",x,cap)
        push!(caps,cap);
        if it == 0
            ω = ωs[end];
        else
            slope = (caps[end]-caps[end-1])/(ωs[end]-ωs[end-1]);
            ω     = ω + (κ-caps[end])/slope;
            push!(ωs,ω);
        end
        err = abs(caps[end] - κ)/κ;
        it  += 1;
    end
    return(ω);
end;

# ## [Figures](@id mw2009_figures)
# Start with the benchmark calibration:
# ### [Benchmark Economy](@id mw2009_fig_benchmark)

ω = MW(3,α,A,Qq,Qz,H,H);

agg, err  = agg_drip(ω,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95);
idi       = Drip(ω,1,A,Qz,H,w = 0.9);

@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",DRIPs.capacity(agg),DRIPs.capacity(idi));

# Plot IRFs
iirfs  = irfs(idi, T = L)
airfs  = irfs(agg, T = L)

using Plots, LaTeXStrings; pyplot();
p1 = plot([iirfs.a[1,1,:],σz*H],
     label             = ["Perfect information" "Rational inattention"],
     marker            = [:square :circle],
     color             = [:gray25 :black],
     markercolor       = [false :black],
     markerstrokecolor = [:gray25 :black],
     markersize        = [7 3],
     xlabel            = "Periods",
     ylabel            = "Impulse responses to shocks \n of one standard deviation");
p2 = plot([σq*H,airfs.a[1,1,1:end],σq*agg.H],
     label             = ["Perfect information" "Rational inattention" L"$\Delta$ at the fixed point w/ RI"],
     marker            = [:circle :square :utriangle],
     color             = [:black :gray25 :gray50],
     markercolor       = [:black false false] ,
     markerstrokecolor = [:black :gray25 :gray50],
     markersize        = [3 7 7],
     xlabel            = "Periods",
     ylabel            = "Impulse responses to shocks \n of one standard deviation")

plot(p1,p2,
    layout     = (2,1),
    xlim       = (0,L+1),
    lw         = 1,
    legend     = :topright,
    legendfont = font(12),
    tickfont   = font(12),
    size       = (1000,550),
    xticks     = 1:2:21,
    framestyle = :box)

# ### [Other values of real rigidity ($\alpha$)](@id mw2009_robust_alpha)

# For $\alpha = 0.7$:

ω_α7 = MW(3,0.7,A,Qq,Qz,H,H);

agg_α7, err = agg_drip(ω_α7,A,Qq,0.7,H; H0 = rand(L), maxit = 100, w = 0.95);
idi_α7      = Drip(ω_α7,1,A,Qz,H,w = 0.9);

@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",
        DRIPs.capacity(agg_α7),DRIPs.capacity(idi_α7));

# For $\alpha = 0$:

ω_α0 = MW(3,0,A,Qq,Qz,H,H);

agg_α0, err = agg_drip(ω_α0,A,Qq,0,H; H0 = rand(L), maxit = 100, w = 0.95);
idi_α0      = Drip(ω_α0,1,A,Qz,H,w = 0.9);

@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",
        DRIPs.capacity(agg_α0),DRIPs.capacity(idi_α0));

# Plot IRFs:

airfs_α7  = irfs(agg_α7, T = L);
airfs_α0  = irfs(agg_α0, T = L);

plot(1:L,0.75*[σq*H,airfs.a[1,1,:],airfs_α7.a[1,1,:],airfs_α0.a[1,1,:]],
    label             = ["Perfect information" L"Rational inattention, benchmark economy, $\alpha = 0.85$" L"Lower degree of real rigidity, $\alpha = 0.7$" L"Lower degree of real rigidity, $\alpha = 0$"],
    marker            = [:circle :square :x :circle],
    color             = [:black :gray20 :gray40 :gray60],
    markercolor       = [:black false :gray40 false],
    markerstrokecolor = [:black :gray20 :gray40 :gray60],
    markersize        = [3 7 7 7],
    xlim              = (0,L+1),
    lw                = 1,
    xticks            = 1:2:21,
    legend            = :topright,
    legendfont        = font(11),
    tickfont          = font(11),
    size              = (1000,275),
    ylim              = (0,σq),
    framestyle        = :box,
    xlabel            = "Periods",
    ylabel            = "Impulse responses to shocks \n of one standard deviation")

# ### [Other values of information capacity ($\kappa$)](@id mw2009_robust_kappa)
# For $\kappa=4$:

ω_κ4 = MW(4,α,A,Qq,Qz,H,H);

agg_κ4, err = agg_drip(ω_κ4,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
idi_κ4      = Drip(ω_κ4,1,A,Qz,H,w = 0.9)

@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",
        DRIPs.capacity(agg_κ4),DRIPs.capacity(idi_κ4));

# For $\kappa=5$:

ω_κ5 = MW(5,α,A,Qq,Qz,H,H; ω = 0.1*σq^2);

agg_κ5, err = agg_drip(ω_κ5,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
idi_κ5      = Drip(ω_κ5,1,A,Qz,H,w = 0.9)

@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",
        DRIPs.capacity(agg_κ5),DRIPs.capacity(idi_κ5));

# Plot IRFs:

airfs_κ4  = irfs(agg_κ4, T = L);
airfs_κ5  = irfs(agg_κ5, T = L);

plot(1:L,0.75*[σq*H,airfs.a[1,1,:],airfs_κ4.a[1,1,:],airfs_κ5.a[1,1,:]],
    label             = ["Perfect information" L"Rational inattention, benchmark economy, $\kappa = 3$" L"Rational inattention $\kappa = 4$" L"Rational inattention $\kappa = 5$"],
    marker            = [:circle :square :x :circle],
    color             = [:black :gray20 :gray40 :gray60],
    markercolor       = [:black false :gray40 false],
    markerstrokecolor = [:black :gray20 :gray40 :gray60],
    markersize        = [3 7 7 7],
    xlim              = (0,L+1),
    lw                = 1,
    xticks            = 1:2:21,
    legend            = :topright,
    legendfont        = font(11),
    tickfont          = font(11),
    size              = (1000,275),
    ylim              = (0,σq),
    framestyle        = :box,
    xlabel            = "Periods",
    ylabel            = "Impulse responses to shocks \n of one standard deviation")

# ## [Measure Performance/Speed](@id mw2009_performance)
# ### Performance of the code for aggregate problem with feedback

# For random values of $\omega$ and benchmark values of other parameters:

using BenchmarkTools;
@benchmark agg_drip(ω,A,Qq,α,H; H0 = rand(L), maxit = 100, w = 0.95) setup = (ω = σq^2*5*rand())

# For $\alpha=0.7$ and random values of $\omega$:

@benchmark agg_drip(ω_α7,A,Qq,0.7,H; H0 = rand(L), maxit = 100, w = 0.95) setup = (ω = σq^2*5*rand())

# ### Performance of the code for idiosyncratic problem

@benchmark Drip(ω,1,A,Qz,H,w = 0.9) setup = (ω = σq^2*5*rand())


text:
module TodoFilter

using Redux

# actions
struct SetVisibilityFilter{T<:AbstractString} <: AbstractSyncAction
    filter::T
end

# state
struct State <: AbstractImmutableState
    filter::String
end

const SHOW_ALL = "show_all"
const SHOW_COMPLETED = "show_completed"
const SHOW_ACTIVE = "show_active"

# reducers
todo_filter(state::AbstractState, action::AbstractAction) = state
todo_filter(state::Vector{<:AbstractState}, action::AbstractAction) = state
todo_filter(s::State, a::SetVisibilityFilter) = State(a.filter)

end # module


text:
module MothModule

using ..Ahorn, Maple

@mapdef Entity "HonlyHelper/Moth" HonlyMoth(x::Integer, y::Integer)

const placements = Ahorn.PlacementDict(
    "Möth (HonlyHelper)" => Ahorn.EntityPlacement(
        HonlyMoth
    )
)

sprite = "objects/HonlyHelper/Bugs/Moth.png"

function Ahorn.selection(entity::HonlyMoth)
    x, y = Ahorn.position(entity)
    return Ahorn.Rectangle(x-2, y-2, 4, 4)
end

Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::HonlyMoth, room::Maple.Room) = Ahorn.drawSprite(ctx, sprite, 0, 0)

end

max_stars_repo_name: HollyMagala/HonlyHelper

max_issues_repo_name: HollyMagala/HonlyHelper

max_forks_repo_name: HollyMagala/HonlyHelper


text:
module PolytopesTests

##
using Gridap, Test

##
# Eliminating the NodeArray struct
##
# Checking Polytope and NFace APIs
D = 3

p = Polytope(1,1,1)

nf = nfaces(p)[9]

@test dim(p) == 3

@test extrusion(p).array == [1,1,1]

@test length(nfaces(p)) == 27

@test nf_nfs(p)[end] == [i for i in 1:27]

@test nf_dim(p)[end][1] == 1:8

@test dim(nf) == 1

@test space_dim(nf) == 3

@test anchor(nf).array == [0,0,0]

@test vertices_coordinates(p)[8].array == [1.0,1.0,1.0]

@test num_nfaces(p,1) == 12

@test extrusion(nface_ref_polytopes(p)[end]) == extrusion(p)

@test generate_admissible_permutations(p)[end] == [8,7,6,4]

@test length(equidistant_interior_nodes_coordinates(p,4)) == 27

@test equidistant_interior_nodes_coordinates(p,[2,2,2])[1].array == [0.5, 0.5, 0.5]

@test vertices_coordinates(p)[8].array == [1.0,1.0,1.0]

@test face_normals(p)[1][1].array ≈ [0.0,0.0,-1.0]

## Checking anisotropic order for n-cubes

D = 3

p = Polytope(1,1,1)

order = [2,4,3]

nodes = equidistant_interior_nodes_coordinates(p, order)

@test nodes[3].array ≈ [0.5,0.75,1.0/3.0]

##
# Adding outwards normals
D = 3
p = Polytope(1,1,1)
ns, f_os = Gridap.Polytopes.face_normals(p)
@test ns[1].array ≈ [0, 0 ,-1]
p = Polytope(1,2,2)
ns, f_os = Gridap.Polytopes.face_normals(p)
@test ns[1].array ≈ [0, 0 ,-1]
##
# Adding outwards normals
D = 2
p = Polytope(1,1)
ns, f_os = Gridap.Polytopes.face_normals(p)
@test ns[1].array ≈ [0,-1]
p = Polytope(1,2)
ns, f_os = Gridap.Polytopes.face_normals(p)
@test ns[1].array ≈ [0,-1]
##
# Adding outwards normals
D = 4
p = Polytope(1,1,1,1)
ns, f_os = Gridap.Polytopes.face_normals(p)
@test ns[1].array ≈ [0,0,0,-1]
p = Polytope(1,2,2,2)
ns, f_os = Gridap.Polytopes.face_normals(p)
@test ns[1].array ≈ [0,0,0,-1]
##
D = 3
p = Polytope(1,1,1)
nf_vs = Gridap.Polytopes._dimfrom_fs_dimto_fs(p,2,0)
vs = Gridap.Polytopes.vertices_coordinates(p)
@test length(nf_vs) == 6
@test nf_vs[end] == [2,4,6,8]
p = Polytope(Point(1,2,2))
nf_vs = Gridap.Polytopes._dimfrom_fs_dimto_fs(p,2,0)
@test length(nf_vs) == 4
@test nf_vs[end] == [2,3,4]
##

# using Gridap.Polytopes
p = Polytope(HEX_AXIS,HEX_AXIS)
for dim = 1:4
  p = Polytope(ones(Int,dim)...)
  _order = ones(Int,dim)
  for i in 1:5
    order = Tuple(_order*i)
    # nf = p.nfaces[end]
    vs = Gridap.Polytopes._interior_nodes_int_coords(p,order)
    @test length(vs) == max(0,(i-1)^dim)
  end
  vcs = Gridap.Polytopes.vertices_coordinates(p)
  @test length(vcs) == 2^dim
end
##
a = zeros(Int,4,4)
a[1,2] = 1
a[1,3] = 2
a[1,4] = 3
a[2,3] = 1
a[2,4] = 3
a[3,4] = 1
##
for dim = 1:4
  p = Polytope(2*ones(Int,dim)...)
  _order = ones(Int,dim)
  for i in 1:4
    order = Tuple(_order*i)
    # nf = p.nfaces[end]
    vs = Gridap.Polytopes._interior_nodes_int_coords(p,order)
    @test length(vs) == a[dim,i]
  end
  vcs = Gridap.Polytopes.vertices_coordinates(p)
  @test length(vcs) == dim+1
end
##
p = Polytope(HEX_AXIS, HEX_AXIS, TET_AXIS)
_order = ones(Int,3)
order = Tuple(_order*3)
# nf = p.nfaces[end]
vs = Gridap.Polytopes._interior_nodes_int_coords(p,order)
@test length(vs) == 1
vcs = Gridap.Polytopes.vertices_coordinates(p)
@test length(vcs) == 5
##
pl = Polytope(HEX_AXIS, TET_AXIS, HEX_AXIS)
_order = ones(Int,3)
order = Tuple(_order*3)
# p = pl.nfaces[end]
vs = Gridap.Polytopes._interior_nodes_int_coords(pl,order)
length(vs)
@test length(vs) == 2
vcs = Gridap.Polytopes.vertices_coordinates(pl)
@test length(vcs) == 6
##

# Developing the change of basis for all n-faces of a polytope
# 1. Given an n-face, determine all rigid-body permutations
# 2. Method that identifies the change of basis required to glue
# together two different orientations
# input: [gn1, ..., gnk] global vertex labels in the base n-face order
# input: [gn1', ..., gnk'] in the cell-based order of the n-face
# output: required change of basis
# Case 1: edges in 2D
# gid = [10,11]
# l2g = [10, 11]
# acceptable permutations = [ [1,2], [2,1] ]
# my_perm = 1
# gid = [11, 10]
# my_perm = 2

# Checking all topologies
##
D=3
# Cube
ext = (1,1,1)
polytope = Polytope(ext)
polytope.extrusion.array.data
x = Vector{Float64}(undef,D)
x .= polytope.nfaces[1].anchor.array
@test length(polytope.nfaces) == 27
@test num_nfaces(polytope,0) == 8
@test num_nfaces(polytope,1) == 12
@test num_nfaces(polytope,2) == 6
@test num_nfaces(polytope,3) == 1
# Pyramid
ext = (1,1,2)
polytope = Polytope(ext)
@test length(polytope.nfaces) == 19
# Prysm
ext = (1,2,1)
polytope = Polytope(ext)
@test length(polytope.nfaces) == 21
ext = Point(1,2,2)
polytope = Polytope(ext)
@test length(polytope.nfaces) == 15
##

# Test to check nodes on the closure of an nface of a polytope
##
D=3
orders=[2,3,2]
ext = Point(1,1,1)
polytope = Polytope(ext)
nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,orders)
cv1 = Gridap.CellValuesGallery.CellValueFromArray(polytope.nf_nfs) # cell to index
cv2 = Gridap.CellValuesGallery.CellValueFromArray(nfacenodes)
closurenfacenodes =  Gridap.CellValuesGallery.CellVectorByComposition(cv1,cv2)
@test length(closurenfacenodes[end-1])==12
@test closurenfacenodes[end-1][end-1]==33
##

# Similar test on the (open for dim > 0) nface
##
D=3
orders=[2,3,2]
ext = Point(1,1,1)
polytope = Polytope(ext)
nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,orders)
nfacenodes
@test length(nfacenodes[end-1])==2
@test nfacenodes[end-1][end-1]==33
##

# Test to check the node coordinates
##
D=3
orders=[2,3,4]
ext = Point(1,1,1)
polytope = Polytope(ext)
nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,orders)
@test length(nodes)==60
@test nodes[33] ≈ Point(0.5, 1.0/3.0, 0.0)
cv1 = Gridap.CellValuesGallery.CellValueFromArray(polytope.nf_nfs) # cell to index
cv2 = Gridap.CellValuesGallery.CellValueFromArray(nfacenodes)
closurenfacenodes =  Gridap.CellValuesGallery.CellVectorByComposition(cv1,cv2)
nfacenodes = closurenfacenodes[end-1]
coords = nodes[nfacenodes,:]
fco = i -> coords[i][1]
@test (prod(fco(i) for i=1:length(coords))==1)
##

# Checking _dimfrom_fs_dimto_fs
D = 3
p = Polytope(1,1,1)
nf_vs = Gridap.Polytopes._dimfrom_fs_dimto_fs(p,2,0)
@test length(nf_vs) == 6
@test nf_vs[end] == [2,4,6,8]
p = Polytope(Point(1,2,2))
nf_vs = Gridap.Polytopes._dimfrom_fs_dimto_fs(p,2,0)
@test length(nf_vs) == 4
@test nf_vs[end] == [2,3,4]

# Edge
##
p = Polytope(1,1,2)
perm_p = Gridap.Polytopes.generate_admissible_permutations(p)
@test length(perm_p) == 48
# @show perm_p
p = Polytope(1,2,2)
perm_p = Gridap.Polytopes.generate_admissible_permutations(p)
@test length(perm_p) == 24
##

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5
# D = 3
# p = Polytope(1,1,1)
# anc = p.nfaces[end].anchor
# ext = p.nfaces[end].ext
#
# p.nfaces
# # vertices coordinates
# vs_r = p.nf_dim[end][1]
# vs = p.nfaces[vs_r]
# num_vs = length(vs)
# vs_cs = Vector{Vector{Int}}(undef, num_vs)
# for (i_v,v) in enumerate(vs)
#   vs_cs[i_v] = v.anchor
# end
# perm = zeros(Int64, D*num_vs, D)
# c = 1
# for (i_v,v) in enumerate(vs)
#   global(c)
#   v_cs = vs_cs[i_v]
#   for idim = 1:D
#     perm[c, 1] = i_v
#     auxv = copy(v_cs)
#     v_cs[idim] == 0 ? auxv[idim] = 1 : auxv[idim] = 0
#     for k = 1:num_vs
#       if vs_cs[k] == auxv
#         perm[c, idim+1] = k
#         c += 1
#         exit
#       end
#     end
#   end
# end
# perm
# ##
#
#
#
# 1 2
# 2 1
#
# 1 -> 2 first axis
# 1 -> 3 second axis
#
# if placed in 2
# 2 -> 1 ok
# 2 -> 3 not possible
# 2 -> 4 ok
#
# we must do it with i,j !
#
# 1 = (0,0)
# 2 = (1,0)
# 3 = (0,1)
# 4 = (1,1)
#
# Thus, at 2
# we can choose the first axis to be e1 or e2
# If e1, since x2(1) = 1, +(-1,0) -> (0,0)
# Next e2, since x2(2) = 0, +(0,1) -> (1,1)
# It provides 2,1 and 2,4 all what is needed to determine orientations
# t
#
#
# 1 2 3 4
# 1 3 2 4
# 2 4 1 3
# 2 1 4 3
# 3 4 1 2
# 3 1 4 2
# 4 3 2 1
# 4 2 3 1
#
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

# Test to check the views of n-face set for a given n
##
# D=3
# ext = Point{D,Int}(1,1,1)
# polytope = Polytope(ext)
# for j=1:length(polytope.ext)+1
#   for i=1:length(polytope.nfaces[polytope.dimnfs[j]])
#     @test (sum(polytope.nfaces[polytope.dimnfs[j]][i].ext)==j-1)
#   end
# end
##
# using Gridap.Vtkio
# D = 3
# wedge = Polytope(1,2,1)
# @show wedge.nf_nfs
# writevtk(wedge,"wedge")
# pyramid = Polytope(1,1,2)
# writevtk(pyramid,"pyramid")
# hex = Polytope(1,1,1)
# writevtk(hex,"hex")
# tet = Polytope(1,2,2)
# writevtk(tet,"tet")

end # module


text:
struct AdjacencyMap{D,T,M} <: AbstractGraph{T}
    colptr::Vector{T}  # column i is in colptr[i]:(colptr[i+1]-1)
    rowval::Vector{T}  # row indices of stored values
    metadata::M

    AdjacencyMap{D,T,M}() where {D,T,M} = new{D::Bool,T,M}(T[1], T[], M())
    function AdjacencyMap{D,T,M}(n::Integer, m::M) where {D,T,M}
        (n < 0) && throw(ArgumentError("cannot have negative number of vertices"))
        _check_meta_length(0, m)
        return new{D::Bool,T,M}(fill(one(T), n + 1), T[], m)
    end
    function AdjacencyMap{D,T,M}(n::Integer) where {D,T,M}
        (n < 0) && throw(ArgumentError("cannot have negative number of vertices"))
        return new{D::Bool,T,M}(fill(one(T), n + 1), T[], M())
    end
    #= TODO appropriate way to bypass some checks
    function AdjacencyMap{D}(c::Vector{T}, r::Vector{T}, m::M) where {D,T,M}
        AdjacencyMap{D}(length(c) - 1, c, r, m)
    end
    =#
    function AdjacencyMap{D}(n::Integer, c::Vector{T}, r::Vector{T}, m::M) where {D,T,M}
        _check_type(n, T)
        _goodbuffers(n, c, r, m)
        return new{D::Bool,T,M}(c, r, m)
    end
    function AdjacencyMap{false,T,NoMetadata}(x::AbstractMatrix) where {N,T}
        n, m = size(x)
        n === m || throw(ArgumentError("Adjacency / distance matrices must be square"))
        g = new{false,T,NoMetadata}(fill(one(T), n+1), T[], no_metadata)
        @inbounds for j in OneTo(n)
            for k in OneTo(n)
                if ((j < k) && !iszero(x[j,k]))
                    throw(ArgumentError("Undirected AdjacencyMap can only be constructed from matrices where the lower triangular values are all zero"))
                end
                if !iszero(x[j,k])
                    add_edge!(x, j, k)
                end
            end
        end
        return g
    end
    function AdjacencyMap{true,T,NoMetadata}(x::AbstractMatrix) where {T}
        n, m = size(x)
        n === m || throw(ArgumentError("Adjacency / distance matrices must be square"))
        g = new{true,T,NoMetadata}(fill(one(T), n+1), T[], no_metadata)
        @inbounds for j in OneTo(n)
            for k in OneTo(n)
                if !iszero(x[j,k])
                    add_edge!(x, j, k)
                end
            end
        end
        return g
    end
    function AdjacencyMap{true,Ti,ListOfMeta{Tv,SUnitRange{1,1}}}(x::AbstractSparseMatrixCSC) where {Ti,Tv}
        n, m = size(x)
        n === m || throw(ArgumentError("Adjacency / distance matrices must be square"))
        new{true,Ti,ListOfMeta{Tv,SUnitRange{1,1}}}(
            copy(getcolptr(x)),
            copy(rowvals(x)),
            ListOfMeta(copy(nonzeros(x)), static(1):static(1))
        )
    end
    function AdjacencyMap{false,Ti,ListOfMeta{Tv,SUnitRange{1,1}}}(x::AbstractSparseMatrixCSC) where {Ti,Tv}
        n, m = size(x)
        n === m || throw(ArgumentError("Adjacency / distance matrices must be square"))
        # TODO should this be don by hand to allow for self loops
        if istriu(x)
            throw(ArgumentError("Undirected AdjacencyMap can only be constructed from matrices where the lower triangular values are all zero"))
        end
        new{true,Ti,ListOfMeta{Tv,SUnitRange{1,1}}}(
            copy(getcolptr(x)),
            copy(rowvals(x)),
            ListOfMeta(copy(nonzeros(x)), static(1):static(1))
        )
    end
end

const AdjMap{T,M} = AdjacencyMap{false,T,M}

const AdjDiMap{T,M} = AdjacencyMap{true,T,M}

const WAdjMap{Ti,Tv} = AdjMap{Ti,ListOfMeta{Tv,SUnitRange{1,1}}}

const WAdjDiMap{Ti,Tv} = AdjDiMap{Ti,ListOfMeta{Tv,SUnitRange{1,1}}}

WAdjMap{Ti}() where {Ti} = WAdjMap{Ti,Float64}()
WAdjDiMap{Ti}() where {Ti} = WAdjDiMap{Ti,Float64}()

WAdjMap() = WAdjMap{Int}()
WAdjDiMap() = WAdjDiMap{Int}()

WAdjMap{Ti}(n::Integer) where {Ti} = WAdjMap{Ti,Float64}(n)
WAdjDiMap{Ti}(n::Integer) where {Ti} = WAdjDiMap{Ti,Float64}(n)

WAdjMap(n::Integer) = WAdjMap{Int}(n)
WAdjDiMap(n::Integer) = WAdjDiMap{Int}(n)

weights(g::WAdjDiMap{Ti,Tv}) where {Ti,Tv} = MetaAdjView{Tv}(g, 1)
weights(g::WAdjMap{Ti,Tv}) where {Ti,Tv} = Symmetric(MetaAdjView{Tv}(g, 1), :U)

## empty
AdjDiMap{T}() where {T} = AdjDiMap{T,NoMetadata}()

AdjMap{T}() where {T} = AdjMap{T,NoMetadata}()

## from nvertices
AdjDiMap{T}(n::Integer) where {T} = AdjDiMap{T,NoMetadata}(n)
AdjDiMap(n::Integer) = AdjDiMap{typeof(n)}(n)

AdjMap{T}(n::Integer) where {T} = AdjMap{T,NoMetadata}(n)
AdjMap(n::Integer) = AdjMap{typeof(n)}(n)

# from matrix
AdjDiMap{T}(x::AbstractMatrix) where {T} = AdjDiMap{T,NoMetadata}(x)
AdjDiMap(x::AbstractMatrix) = AdjDiMap{Int}(x)
function AdjDiMap{Ti}(x::AbstractSparseMatrixCSC{Tv}) where {Tv,Ti}
    AdjDiMap{Ti,ListOfMeta{Tv,SUnitRange{1,1}}}(x)
end
function AdjDiMap(x::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
    AdjDiMap{Ti}(x::AbstractSparseMatrixCSC)
end

AdjMap{T}(x::AbstractMatrix) where {T} = AdjMap{T,NoMetadata}(x)
AdjMap(x::AbstractMatrix) = AdjMap{Int}(x)
function AdjMap{Ti}(x::AbstractSparseMatrixCSC{Tv}) where {Tv,Ti}
    AdjMap{Ti,ListOfMeta{Tv,SUnitRange{1,1}}}(x)
end
function AdjMap(x::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
    AdjMap{Ti}(x::AbstractSparseMatrixCSC)
end

vertices(g::AdjacencyMap) = OneTo(nvertices(g))

SparseArrays.getcolptr(g::AdjacencyMap) = getfield(g, :colptr)

SparseArrays.rowvals(g::AdjacencyMap) = getfield(g, :rowval)

Base.eltype(::AdjacencyMap{D,T,M}) where {D,T,M} = T
Base.eltype(::Type{<:AdjacencyMap{D,T,M}}) where {D,T,M} = T

@inline nvertices(g::AdjacencyMap) = length(getcolptr(g)) - 1

@inline nedges(g::AdjacencyMap) = Int(@inbounds(getcolptr(g)[nvertices(g) + 1])) - 1

function Base.size(g::AdjacencyMap, i::Int)
    if i === 1 || i === 2
        return nvertices(g)
    else
        return 1
    end
end
@inline function Base.size(g::AdjacencyMap)
    v = nvertices(g)
    return (v, v)
end

function Base.show(io::IO, @nospecialize(g::AdjacencyMap))
    dir = is_directed(g) ? "directed" : "undirected"
    print(io, "{$(nvertices(g)), $(nedges(g))} $dir simple $(eltype(g)) graph")
end

""" MetaAdjView """
struct MetaAdjView{Tv,D,Ti,M,K} <: AbstractSparseMatrixCSC{Tv,Ti}
    parent::AdjacencyMap{D,Ti,ListOfMeta{M,K}}
    index::Int

    function MetaAdjView{Tv}(p::AdjacencyMap{D,Ti,ListOfMeta{M,K}}, i::Int) where {Tv,D,Ti,M,K}
        new{Tv,D,Ti,M,K}(p, i)
    end
end

Base.parent(x::MetaAdjView) = getfield(x, :parent)

SparseArrays.getcolptr(x::MetaAdjView) = getcolptr(parent(x))

SparseArrays.rowvals(x::MetaAdjView) = rowvals(parent(x))

function SparseArrays.nonzeros(x::MetaAdjView{Tv,D,Ti,M,K}) where {Tv,D,Ti,M,K}
    ItemView{Tv,M,K}(metadata(parent(x)), getfield(x, :index))
end

nvertices(g::MetaAdjView) = nvertices(parent(g))

@inline function Base.size(g::MetaAdjView)
    v = nvertices(g)
    return (v, v)
end

Base.setindex!(x::MetaAdjView, val, i, j) = throw(MethodError(setindex!, (x, val, i, j)))

Base.setindex!(x::MetaAdjView, val, i) = throw(MethodError(setindex!, (x, val, i)))

function Metadata.metadata_keys(x::AdjacencyMap{D,T,M}) where {D,T,M<:ListOfMeta}
    Metadata.metadata_keys(metadata(x))
end
metadata(g::AdjacencyMap) = getfield(g, :metadata)

function metadata(x::AdjacencyMap, k::Int)
    @boundscheck (1 <= i <= length(metadata_keys(x))) && throw(BoundsError(x, i))
    return MetaAdjView{typeof(first(@inbounds(getfield(x, :metadata)[i])))}(x, k)
end

metadata_type(::Type{<:AdjacencyMap{D,T,M}}) where {D,T,M} = eltype(M)


text:
module Debugger

using ..Atom, MacroTools, ASTInterpreter, Lazy, Hiccup

include("stepper.jl")
include("entry.jl")

end # module


text:
module TestTree0

using Test
import Pkg: TOML
import TreeLas
import GraphIdx

include("data.jl")

function readtoml(fname::String; nan="-1.0")
    txt = read(open(fname), String)
    txt = replace(txt, "nan" => nan)
    TOML.parse(txt)
end

function instance(tree)
    y = Float64.(tree["y"])
    n = length(y)
    mu = GraphIdx.Weights(Float64.(get(tree, "mu", 1.0)))
    pi = Int.(tree["parent"]) .+ 1
    root::Int = get(tree, "root", -1) + 1
    if root <= 0
        root = GraphIdx.Tree.find_root(pi)
    end
    lam = Float64.(tree["lam"])
    if lam isa Array
        @assert findall(lam .< 0) == [root] "$(typeof(lam)): $lam, $root"
        lam[root] = NaN
    end
    lam = GraphIdx.Weights(lam)
    t = GraphIdx.Tree.RootedTree(root, pi)
    return y, t, lam, mu
end

@testset "tree0.toml                     " begin
    for (i, ti) in enumerate(readtoml(data_dir("test", "tree0.toml"))["tree"])
        @testset "tree0.toml[$i]" begin
            y, t, lam, mu = instance(ti)
            @test TreeLas.TreeDP.tree_dp(y, t, lam, mu) ≈ ti["x"]
        end
    end
end

@testset "tree0.1.toml                   " begin
    ti = readtoml(data_dir("test", "tree0.1.toml"))["tree"][1]
    y, t, lam, mu = instance(ti)
    x = TreeLas.TreeDP.tree_dp(y, t, lam, mu)
    @test all(isfinite.(x))
    @test minimum(y) <= minimum(x)
    @test maximum(x) <= maximum(y)
end

end

max_stars_repo_name: EQt/treelas

max_issues_repo_name: EQt/treelas

max_forks_repo_name: EQt/treelas


text:
import Base: ==, string, length, show, convert

export AbstractUnitProblem,
    LeasedProblem,
    LeasedProblemContsDrill,
    PerpetualProblem,
    sprime,
    ssprime,
    state,
    state_idx,
    end_ex1,
    end_ex0,
    end_lrn,
    end_inf

abstract type AbstractUnitProblem end

struct LeasedProblem <: AbstractUnitProblem
    dmax::Int
    Dmax::Int
    τ0max::Int
    τ1max::Int
    ext::Int
end

struct LeasedProblemContsDrill <: AbstractUnitProblem
    dmax::Int
    Dmax::Int
    τ0max::Int
    τ1max::Int
    ext::Int
end

struct PerpetualProblem <: AbstractUnitProblem
    dmax::Int
    Dmax::Int
end

PerpetualProblem(dmax,Dmax,τ0max,τ1max,ext) = PerpetualProblem(dmax,Dmax)

# see https://stackoverflow.com/questions/40160120/generic-constructors-for-subtypes-of-an-abstract-type
(::Type{T})(wp::AbstractUnitProblem) where {T<:AbstractUnitProblem} = T(_dmax(wp), _Dmax(wp), _τ0max(wp), _τ1max(wp), _ext(wp))
function (::Type{T})(wp::PerpetualProblem) where {T<:Union{LeasedProblem,LeasedProblemContsDrill}}
    @warn("PerpetualProblem does not have τ0max, τ1max, or ext information!!")
    T(_dmax(wp), _Dmax(wp), _τ0max(wp), _τ1max(wp), _ext(wp))
end

convert(::Type{T}, wp::AbstractUnitProblem) where {T<:AbstractUnitProblem} = T(wp)

_nstates_per_D(wp::AbstractUnitProblem) = 1
_nstates_per_D(wp::LeasedProblemContsDrill) = 2

_dmax(  wp::AbstractUnitProblem) = wp.dmax
_Dmax(  wp::AbstractUnitProblem) = wp.Dmax

_τ0max( wp::AbstractUnitProblem) = wp.τ0max
_τ1max( wp::AbstractUnitProblem) = wp.τ1max
_ext(   wp::AbstractUnitProblem) = wp.ext

_τ0max( wp::PerpetualProblem) = -1
_τ1max( wp::PerpetualProblem) = -1
_ext(   wp::PerpetualProblem) = 0

# ----------------------------------
# endpoints of each set of states
#   end_ex1 - Expiration of first lease term
#   end_ex0 - End of last lease term (expiration if present, or regular lease if not)
#   end_lrn - In learning stage, we have drilled _dmax(wp) wells
#   end_inf - Terminal state -- we've drilled everything
#   strt_ex - State where extension starts
# ----------------------------------
@inline end_ex1(wp::LeasedProblemContsDrill) = _τ1max(wp)+1
@inline end_ex0(wp::LeasedProblemContsDrill) = end_ex1(wp) + max(_τ0max(wp),_ext(wp))+1
@inline end_lrn(wp::LeasedProblemContsDrill) = end_ex0(wp) + _dmax(wp)+1
@inline end_inf(wp::LeasedProblemContsDrill) = end_lrn(wp) + 2*(_Dmax(wp)-1)+ 1
@inline strt_ex(wp::LeasedProblemContsDrill) = end_ex0(wp) - (_ext(wp)-1)     # t+d+2(D)+1 = (t+1) + (d+1) + 2*(D-1) + 1

@inline end_ex1(wp::LeasedProblem) = _τ1max(wp)+1
@inline end_ex0(wp::LeasedProblem) = end_ex1(wp) + max(_τ0max(wp),_ext(wp))+1
@inline end_lrn(wp::LeasedProblem) = end_ex0(wp) + _dmax(wp)+1
@inline end_inf(wp::LeasedProblem) = end_lrn(wp) + _Dmax(wp)
@inline strt_ex(wp::LeasedProblem) = end_ex0(wp) - (_ext(wp)-1)     # t+d+2(D)+1 = (t+1) + (d+1) + 2*(D-1) + 1

@inline end_ex1(wp::PerpetualProblem) = 0
@inline end_ex0(wp::PerpetualProblem) = 1
@inline end_lrn(wp::PerpetualProblem) = end_ex0(wp) + _dmax(wp)
@inline end_inf(wp::PerpetualProblem) = end_lrn(wp) + _Dmax(wp)
@inline strt_ex(wp::PerpetualProblem) = 1

# ----------------------------------
# length of states
# ----------------------------------

Base.length(wp::AbstractUnitProblem) = end_inf(wp)
_nS(    wp::AbstractUnitProblem) = length(wp)
_nSexp( wp::AbstractUnitProblem) = end_lrn(wp)

# ----------------------------------
# regions of state space where we have learning
# ----------------------------------

exploratory_terminal(wp::AbstractUnitProblem) = end_ex0(wp)+1
exploratory_learning(wp::AbstractUnitProblem) = end_ex0(wp)+2 : end_lrn(wp)

exploratory_terminal(wp::PerpetualProblem) = end_ex0(wp)
exploratory_learning(wp::PerpetualProblem) = end_ex0(wp)+1 : end_lrn(wp)

# ----------------------------------
# indices of where we are in things
# ----------------------------------

@inline ind_ex1(wp::AbstractUnitProblem) = end_ex1(wp)   : -1 : 1
@inline ind_ex0(wp::AbstractUnitProblem) = end_ex0(wp)   : -1 : end_ex1(wp)+1
@inline ind_exp(wp::AbstractUnitProblem) = end_ex0(wp)   : -1 : 1
@inline ind_lrn(wp::AbstractUnitProblem) = end_lrn(wp)   : -1 : end_ex0(wp)+1
@inline ind_inf(wp::AbstractUnitProblem) = end_inf(wp)-1 : -1 : end_lrn(wp)+1

@inline inf_fm_lrn(wp::AbstractUnitProblem) = (end_lrn(wp)+1) .+ _nstates_per_D(wp)*(0:_dmax(wp)-1)

# -------------------------------------------------------------------

function state_if_never_drilled(wp::AbstractUnitProblem, state0::Integer, t::Integer)

    if state0 <= end_ex1(wp)
        if state0 + t <= end_ex1(wp)
            return state0 + t
        elseif state0 + t <= end_ex1(wp) + _ext(wp)
            return (state0 + t) - (end_ex1(wp) - strt_ex(wp)) - 1
        else
            return end_ex0(wp)+1
        end

    elseif state0 <= end_ex0(wp)
        if state0 + t <= end_ex0(wp)
            return state0 + t
        else
            return end_ex0(wp)+1
        end

    else
        return end_ex0(wp)+1
    end
end

function state_if_never_drilled(wp::PerpetualProblem, state0::Integer, t::Integer)
    return 1
end

# ----------------------------------
# which state are we in?
# ----------------------------------

function state_idx(wp::LeasedProblemContsDrill, t1::Integer, t0::Integer, D::Integer, d1::Integer)::Int
    t1 >= 0 && t0 ∉ (0,_ext(wp)) && throw(error("Cannot be in primary + extension simultaneously"))
    t1 >= 0               && return end_ex1(wp) - t1
    t0 >= 0               && return end_ex0(wp) - t0
    D==0  && t0==-1       && return end_ex0(wp) + (1 + 0)
    D<_Dmax(wp) && t0==-1 && return end_lrn(wp) +  1 + 2*(D - 1) + (1-d1)
    D==_Dmax(wp)          && return end_lrn(wp) +  1 + 2*(_Dmax(wp)-1)   # drop last +1 since we have d1=1 at terminal
    throw(error("invalid state"))
end

function state_idx(wp::LeasedProblem, t1::Integer, t0::Integer, D::Integer, d1::Integer)::Int
    t1 >= 0 && t0 ∉ (0,_ext(wp)) && throw(error("Cannot be in primary + extension simultaneously"))
    t1 >= 0                && return end_ex1(wp) - t1
    t0 >= 0                && return end_ex0(wp) - t0
    D==0   && t0==-1       && return end_ex0(wp) + (1 + 0)
    D<=_Dmax(wp) && t0==-1 && return end_lrn(wp) + D
    throw(error("invalid state"))
end

function state_idx(wp::PerpetualProblem, t1::Integer, t0::Integer, D::Integer, d1::Integer)::Int
    D==0         && return end_ex0(wp)
    D<=_Dmax(wp) && return end_lrn(wp) + D
    throw(error("invalid state"))
end

# ----------------------------------
# information about state
# ----------------------------------

function _regime(wp::AbstractUnitProblem, s::Integer)::Symbol
    s <= 0 && throw(DomainError(s, "s <= 0"))
    s <= end_ex1(wp) && return :primary_WITH_extension
    s <= end_ex0(wp) && return :primary_or_extension
    s == end_ex0(wp) + 1 && return :expired
    s <= end_lrn(wp) && return :learn
    s <= end_inf(wp) && return :infill
    throw(DomainError())
end

function _horizon(wp::PerpetualProblem, sidx::Integer)::Symbol
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp)  && return :Infinite
    sidx <= end_lrn(wp)  && return :Learning
    sidx <  end_inf(wp)  && return :Infinite
    sidx == end_inf(wp)  && return :Terminal
    throw(DomainError())
end

function _horizon(wp::LeasedProblem, sidx::Integer)::Symbol
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp)   && return :Finite
    sidx == end_ex0(wp)+1 && return :Terminal
    sidx <= end_lrn(wp)   && return :Learning
    sidx <  end_inf(wp)   && return :Infinite
    sidx == end_inf(wp)   && return :Terminal
    throw(DomainError())
end

function _horizon(wp::LeasedProblemContsDrill, sidx::Integer)::Symbol
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp)   && return :Finite
    sidx == end_ex0(wp)+1 && return :Terminal
    sidx <= end_lrn(wp)   && return :Learning
    sidx <  end_inf(wp)   && return isodd(end_inf(wp)-sidx) ? :Infinite : :Finite
    sidx == end_inf(wp)   && return :Terminal
    throw(DomainError())
end

function s_of_D(wp::PerpetualProblem, D::Integer)
    D ∉ 0:_Dmax(wp) && throw(DomainError())
    D == 0          && return 1:end_ex0(wp)
    D <= _Dmax(wp)  && return end_lrn(wp) .+ (D:D)
end

function s_of_D(wp::LeasedProblem, D::Integer)
    D ∉ 0:_Dmax(wp) && throw(DomainError())
    D == 0          && return 1:end_ex0(wp)+1
    D <= _Dmax(wp)  && return end_lrn(wp) .+ (D:D)
end

function s_of_D(wp::LeasedProblemContsDrill, D::Integer)
    D ∉ 0:_Dmax(wp) && throw(DomainError())
    D == 0          && return 1:end_ex0(wp)+1
    D <= _Dmax(wp)  && return end_lrn(wp) .+ ((2*D-1):2*D)
end

function _D(wp::LeasedProblemContsDrill, sidx::Integer)::Int
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp) && return 0
    sidx <= end_lrn(wp) && return  sidx - end_ex0(wp) - 1
    sidx <= end_inf(wp) && return (sidx - end_lrn(wp) + isodd(sidx-end_lrn(wp)) ) / 2
    throw(DomainError())
end

function _D(wp::LeasedProblem, sidx::Integer)::Int
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp) && return 0
    sidx <= end_lrn(wp) && return sidx - end_ex0(wp) - 1
    sidx <= end_inf(wp) && return sidx - end_lrn(wp)
    throw(DomainError())
end

function _D(wp::PerpetualProblem, sidx::Integer)::Int
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp) && return 0
    sidx <= end_lrn(wp) && return sidx - end_ex0(wp)
    sidx <= end_inf(wp) && return sidx - end_lrn(wp)
    throw(DomainError())
end

_d1(wp::AbstractUnitProblem, sidx::Integer) = 0

function _d1(wp::LeasedProblemContsDrill, sidx::Integer)::Int
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp)+1 && return 0
    sidx <= end_lrn(wp)   && return 1
    sidx <  end_inf(wp)   && return isodd(sidx-end_lrn(wp))
    sidx == end_inf(wp)   && return 1
end

function _τ1(wp::AbstractUnitProblem, sidx::Integer)::Int
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex1(wp) && return end_ex1(wp)-sidx
    return -1
end

function _τ0(wp::AbstractUnitProblem, sidx::Integer)::Int
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex1(wp) && return _ext(wp)
    sidx <= end_ex0(wp) && return end_ex0(wp)-sidx
    return -1
end

_τ1(wp::PerpetualProblem, sidx::Integer) = -1
_τ0(wp::PerpetualProblem, sidx::Integer) = -1

"Deterministic state  (`τ1`, `τ0` `D`, `d1`)"
struct state
    τ1::Int  # Time remaining in initial term
    τ0::Int  # Time remaining in final term
    D::Int   # Wells drilled to date
    d1::Int  # Drilling last period == 1, no drilling last period == 0
end

# ------------- methods for state ----------

# Pretty print: https://docs.julialang.org/en/latest/manual/types.html#Custom-pretty-printing-1
string(s::state)::String = string((s.τ1, s.τ0, s.D, s.d1))
show(io::IO, s::state) = print(io, string(s))

function ==(s1::state, s2::state)
    s1.τ1 == s2.τ1    &&     s1.τ0 == s2.τ0   &&     s1.D == s2.D    &&    s1.d1 == s2.d1
end

state(wp::AbstractUnitProblem, s::Integer) = state(_τ1(wp,s), _τ0(wp,s), _D(wp,s), _d1(wp,s))

_d1(s::state) = s.d1
_τ1(s::state) = s.τ1
_τ0(s::state) = s.τ0
_D(s::state) = s.D
_sgnext(s::state) = s.τ1 == 0 && s.τ0 > 0
_τrem(s::state) = max(s.τ1,0) + max(s.τ0,0)
_τ11(s::state, wp::AbstractUnitProblem) = 2*_τrem(s)/maxlease(wp)-1
_Dgt0(wp::AbstractUnitProblem, sidx::Integer) = sidx > end_ex0(wp)+1
_Dgt0(s::state) = s.D > 0

stateinfo(wp::AbstractUnitProblem, st::state) = (_d1(st), _Dgt0(st), _sgnext(st), _τ11(st, wp),)

_sgnext(wp::AbstractUnitProblem, sidx::Integer) = sidx == end_ex1(wp)
function _sign_lease_extension(sidx::Integer,wp::AbstractUnitProblem)
    @warn "Deprecated. use `sgnext(wp,sidx)`"
    return _sgnext(wp,sidx)
end

function _τrem(wp::AbstractUnitProblem, sidx::Integer)
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex1(wp) && return end_ex1(wp)-sidx + _ext(wp)+1
    sidx <= end_ex0(wp) && return end_ex0(wp)-sidx
    return -1
end

_τrem(wp::PerpetualProblem,sidx::Integer) = -1

max_ext(wp::AbstractUnitProblem) = end_ex0(wp)-strt_ex(wp)+1

maxlease(wp::AbstractUnitProblem) = max( end_ex1(wp)+max_ext(wp), end_ex0(wp)-end_ex1(wp) ) - 1

_τ11(wp::AbstractUnitProblem, sidx::Integer) = 2*_τrem(wp,sidx)/maxlease(wp)-1

# ----------------------------------
# actions
# ----------------------------------

function _dmax(wp::AbstractUnitProblem, sidx::Integer)
    sidx <= 0 && throw(DomainError(sidx, "s <= 0"))
    sidx <= end_ex0(wp) && return _dmax(wp)
    sidx <= end_lrn(wp) && return 0
    sidx <  end_inf(wp) && return min((_Dmax(wp)-_D(wp,sidx)),_dmax(wp))
    sidx == end_inf(wp) && return 0
    throw(DomainError())
end

@inline actionspace(wp::AbstractUnitProblem, sidx::Integer) = 0:_dmax(wp,sidx)
@inline dp1space(   wp::AbstractUnitProblem, sidx::Integer) = actionspace(wp,sidx) .+ 1

# ----------------------------------
# states
# ----------------------------------

function state_space_vector(wp::LeasedProblemContsDrill)::Vector{state}
    _dmax(wp) <= _Dmax(wp) || throw(DomainError("dmax=$_dmax(wp), Dmax=$_Dmax(wp)"))
    # Primary term with extension
    exp1 = [state(τ1,_ext(wp),0,0) for τ1 in _τ1max(wp):-1:0]
    # Exploratory drilling (with terminal lease expiration)
    exp0 = [state(-1,τ0,0,0) for τ0 in _τ0max(wp):-1:-1]
    # Integrated (wrt information) infill
    inf_int = [state(-1,-1,D,1) for D in 1:_dmax(wp)]  # τmax+1 + (0:dmax) (note overlap!)
    # Infill drilling with immediately prior drilling
    infill  = [state(-1,-1,D,d1) for D in 1:_Dmax(wp)-1 for d1 in 1:-1:0] # τmax + 2 + dmax + (1:2*Dmax-2)
    # Infill terminal
    inf_term = [state(-1,-1,_Dmax(wp),1)] # τmax + 2 + dmax + 2*Dmax - 2 + 1  = τmax + dmax + 2*Dmax + 1
    return [exp1..., exp0..., inf_int..., infill..., inf_term...]
end

function state_space_vector(wp::LeasedProblem)::Vector{state}
    _dmax(wp) <= _Dmax(wp) || throw(DomainError("dmax=$_dmax(wp), Dmax=$_Dmax(wp)"))
    # Primary term with extension
    exp1 = [state(τ1,_ext(wp),0,0) for τ1 in _τ1max(wp):-1:0]
    # Exploratory drilling (with terminal lease expiration)
    exp0 = [state(-1,τ0,0,0) for τ0 in _τ0max(wp):-1:-1]
    # Integrated (wrt information) infill
    inf_int = [state(-1,-1,D,0) for D in 1:_dmax(wp)]  # τmax+1 + (0:dmax) (note overlap!)
    # Infill drilling with immediately prior drilling
    infill  = [state(-1,-1,D,0) for D in 1:_Dmax(wp)] # τmax + 2 + dmax + (1:2*Dmax-2)
    return [exp1..., exp0..., inf_int..., infill...]
end

function state_space_vector(wp::PerpetualProblem)::Vector{state}
    _dmax(wp) <= _Dmax(wp) || throw(DomainError("dmax=$_dmax(wp), Dmax=$_Dmax(wp)"))
    # Primary term with extension
    exp1 = [state(τ1,_ext(wp),0,0) for τ1 in _τ1max(wp):-1:0]
    # Exploratory drilling (with terminal lease expiration)
    exp0 = [state(-1,τ0,0,0) for τ0 in _τ0max(wp):-1:-1]
    # Integrated (wrt information) infill
    inf_int = [state(-1,-1,D,0) for D in 1:_dmax(wp)]
    # Infill drilling with immediately prior drilling
    infill  = [state(-1,-1,D,0) for D in 1:_Dmax(wp)]
    return [exp1..., exp0..., inf_int..., infill...]
end

# ----------------------------------
# sprime
# ----------------------------------

@inline function sprime(wp::LeasedProblemContsDrill, s::Integer, d::Integer)::Int
    if s <= 0
        throw(DomainError(s, "s <= 0"))
    elseif s == end_ex1(wp)
        d == 0 && return strt_ex(wp)
        d >  0 && return end_ex0(wp) + d + 1
    elseif s <= end_ex0(wp)
        d == 0 && return s+1
        d >  0 && return end_ex0(wp) + d + 1
    elseif s == end_ex0(wp)+1
        return s
    elseif s <= end_lrn(wp)
        return end_lrn(wp) + 2*(s-end_ex0(wp))-3  # was return endpts[3] + 2*d - 1. note d=(s-endpts[2]-1)
    elseif s < end_inf(wp)
        d == 0  && return s + isodd(s-end_lrn(wp))
        d >  0  && return s + 2*d - iseven(s-end_lrn(wp))
    elseif s == end_inf(wp)
        return s
    else
        throw(DomainError())
    end
end

@inline function sprime(wp::LeasedProblem, s::Integer, d::Integer)::Int
    if s <= 0
        throw(DomainError(s, "s <= 0"))
    elseif s == end_ex1(wp)
        d == 0 && return strt_ex(wp)
        d >  0 && return end_ex0(wp) + d + 1
    elseif s <= end_ex0(wp)
        d == 0 && return s+1
        d >  0 && return end_ex0(wp) + d + 1
    elseif s == end_ex0(wp)+1
        return s
    elseif s <= end_lrn(wp)
        return end_lrn(wp) + (s-end_ex0(wp)-1)
    elseif s < end_inf(wp)
        return s + d
    elseif s == end_inf(wp)
        return s
    else
        throw(DomainError())
    end
end

@inline function sprime(wp::PerpetualProblem, s::Integer, d::Integer)::Int
    if s <= 0
        throw(DomainError(s, "s <= 0"))
    elseif s == end_ex0(wp)
        d == 0 && return s
        d >  0 && return end_ex0(wp) + d
    elseif s <= end_lrn(wp)
        return end_lrn(wp) + (s-end_ex0(wp))
    elseif s < end_inf(wp)
        return s + d
    elseif s == end_inf(wp)
        return s
    else
        throw(DomainError())
    end
end

@inline sprimes(wp::AbstractUnitProblem, sidx::Integer) = (sprime(wp,sidx,d) for d in actionspace(wp,sidx))

"retrieve next state, skipping through learning if necessary."
function ssprime(wp::AbstractUnitProblem, s::Integer, d::Integer)::Int
    if s <= 0
        throw(DomainError(s, "s <= 0"))
    elseif s <= end_ex0(wp)
        return d == 0 ? sprime(wp,s,d) : sprime(wp, sprime(wp,s,d), d)  # because of LEARNING transition
    elseif s <= end_inf(wp)
        return sprime(wp,s,d)
    else
        throw(DomainError(s, "s > end_inf(wp)"))
    end
end

# ----------------------------------
# other functions
# ----------------------------------

# wp_info(wp::AbstractUnitProblem, s::Integer) = (dp1space(wp,s), sprimes(wp,s), _horizon(wp,s), state(wp,s), )
# action0(wp::AbstractUnitProblem,s::Integer) = sprime(wp, s, 0)
# stateinfo(wp::AbstractUnitProblem, s::Integer) = (_d1(wp,s), _Dgt0(wp,s), _sgnext(wp,s), _τ11(wp,s),)
# stateinfo(s::state) = (_d1(s), _Dgt0(s), _sgnext(s), _τ11(s),)

# infill_state_inds( wp::AbstractUnitProblem) = ind_inf(wp)
# explore_state_inds(wp::AbstractUnitProblem) = ind_exp(wp)

# end_pts(wp::AbstractUnitProblem) = (0, end_ex1(wp), end_ex0(wp), end_lrn(wp), end_inf(wp), strt_ex(wp),)
# _nSint( wp::AbstractUnitProblem) = end_lrn(wp) - end_ex0(wp)

# export
# ------------
# _dmax,
# _Dmax,
# _τ0max,
# _τ1max,
# _ext,
# strt_ex,
# _nS,
# _nSexp,
# ind_ex1,
# ind_ex0,
# ind_exp,
# ind_lrn,
# ind_inf,
# inf_fm_lrn,
# _horizon,
# _D,
# _d1,
# _τ1,
# _τ0,
# _Dgt0,
# _sgnext,
# _sign_lease_extension,
# _τrem,
# maxlease,
# _τ11,
# _dmax,
# actionspace,
# dp1space,
# state_space_vector,
# wp_info,
# action0


text:
# Testing gammatonegram response
s = read(open(joinpath(dirname(@__FILE__), "data", "sa2.raw")), Float64, 42701)
fs = 16000
(W, _) = AuditoryFilters.fft2gammatonemx(512, fs, 23, 1, 150, 4000)
G = gammatonegram(s, fs, 0.025, 0.010, 23, 150, 4000, 1)

G_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "G_matlab.csv")))
F_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "F_matlab.csv")))
W_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "W_matlab.csv")))

@test_approx_eq W W_matlab
@test_approx_eq G.frequencies F_matlab'
@test_approx_eq G.amplitude G_matlab


text:
@inline function typeparser(::Type{Bool}, source, pos, len, b, code, options::Options)
    x = false
    trues = options.trues
    falses = options.falses
    if trues === nothing
        if b == UInt8('t')
            pos += 1
            incr!(source)
            if eof(source, pos, len)
                code |= INVALID | EOF
                @goto done
            end
            b = peekbyte(source, pos)
            if b == UInt8('r')
                pos += 1
                incr!(source)
                if eof(source, pos, len)
                    code |= INVALID | EOF
                    @goto done
                end
                b = peekbyte(source, pos)
                if b == UInt8('u')
                    pos += 1
                    incr!(source)
                    if eof(source, pos, len)
                        code |= INVALID | EOF
                        @goto done
                    end
                    b = peekbyte(source, pos)
                    if b == UInt8('e')
                        pos += 1
                        incr!(source)
                        x = true
                        code |= OK
                        if eof(source, pos, len)
                            code |= EOF
                        end
                        @goto done
                    end
                end
            end
        else
            intx, intcode, intpos = typeparser(UInt8, source, pos, len, b, code, options)
            if ok(intcode) && intx < 0x02
                x = intx == 0x01 ? true : false
                code = intcode
                pos = intpos
                @goto done
            end
        end
    else
        if source isa AbstractVector{UInt8}
            startptr = pointer(source, pos)
            for (i, (ptr, ptrlen)) in enumerate(trues)
                if pos + ptrlen - 1 <= len
                    match = memcmp(startptr, ptr, ptrlen)
                    if match
                        x = true
                        code |= OK
                        pos = pos + ptrlen
                        if eof(source, pos, len)
                            code |= EOF
                        end
                        @goto done
                    end
                end
            end
        else # source isa IO
            for (i, (ptr, ptrlen)) in enumerate(trues)
                matched = match!(source, ptr, ptrlen)
                if matched
                    x = true
                    code |= OK
                    pos = pos + ptrlen
                    if eof(source, pos, len)
                        code |= EOF
                    end
                    @goto done
                end
            end
        end
    end
    if falses === nothing
        if b == UInt8('f')
            pos += 1
            incr!(source)
            if eof(source, pos, len)
                code |= INVALID | EOF
                @goto done
            end
            b = peekbyte(source, pos)
            if b == UInt8('a')
                pos += 1
                incr!(source)
                if eof(source, pos, len)
                    code |= INVALID | EOF
                    @goto done
                end
                b = peekbyte(source, pos)
                if b == UInt8('l')
                    pos += 1
                    incr!(source)
                    if eof(source, pos, len)
                        code |= INVALID | EOF
                        @goto done
                    end
                    b = peekbyte(source, pos)
                    if b == UInt8('s')
                        pos += 1
                        incr!(source)
                        if eof(source, pos, len)
                            code |= INVALID | EOF
                            @goto done
                        end
                        b = peekbyte(source, pos)
                        if b == UInt8('e')
                            pos += 1
                            incr!(source)
                            code |= OK
                            if eof(source, pos, len)
                                code |= EOF
                            end
                            @goto done
                        end
                    end
                end
            end
        end
    else
        if source isa AbstractVector{UInt8}
            startptr = pointer(source, pos)
            for (i, (ptr, ptrlen)) in enumerate(falses)
                if pos + ptrlen - 1 <= len
                    match = memcmp(startptr, ptr, ptrlen)
                    if match
                        x = false
                        code |= OK
                        pos = pos + ptrlen
                        if eof(source, pos, len)
                            code |= EOF
                        end
                        @goto done
                    end
                end
            end
        else # source isa IO
            for (i, (ptr, ptrlen)) in enumerate(falses)
                matched = match!(source, ptr, ptrlen)
                if matched
                    x = false
                    code |= OK
                    pos = pos + ptrlen
                    if eof(source, pos, len)
                        code |= EOF
                    end
                    @goto done
                end
            end
        end
    end
    fastseek!(source, pos - 1)
    code |= INVALID

@label done
    return x, code, pos
end


text:
"""
$(TYPEDEF)

A `Revolute` joint type allows rotation about a fixed axis.

The configuration vector for the `Revolute` joint type simply consists of the angle
of rotation about the specified axis. The velocity vector consists of the angular
rate, and is thus the time derivative of the configuration vector.
"""
struct Revolute{T} <: JointType{T}
    axis::SVector{3, T}
    rotation_from_z_aligned::RotMatrix3{T}

    function Revolute{T}(axis::AbstractVector) where {T}
        a = normalize(axis)
        new{T}(a, rotation_between(SVector(zero(T), zero(T), one(T)), SVector{3, T}(a)))
    end
end

"""
$(SIGNATURES)

Construct a new `Revolute` joint type, allowing rotation about `axis`
(expressed in the frame before the joint).
"""
Revolute(axis::AbstractVector{T}) where {T} = Revolute{T}(axis)

Base.show(io::IO, jt::Revolute) = print(io, "Revolute joint with axis $(jt.axis)")

function Random.rand(::Type{Revolute{T}}) where {T}
    axis = normalize(randn(SVector{3, T}))
    Revolute(axis)
end

RigidBodyDynamics.flip_direction(jt::Revolute) = Revolute(-jt.axis)

num_positions(::Type{<:Revolute}) = 1
num_velocities(::Type{<:Revolute}) = 1
has_fixed_subspaces(jt::Revolute) = true
isfloating(::Type{<:Revolute}) = false

@propagate_inbounds function set_configuration!(q::AbstractVector, joint::Joint{<:Any, <:Revolute}, θ::Number)
    @boundscheck check_num_positions(joint, q)
    @inbounds q[1] = θ
    q
end

@propagate_inbounds function set_velocity!(v::AbstractVector, joint::Joint{<:Any, <:Revolute}, θ̇::Number)
    check_num_velocities(joint, v)
    @inbounds v[1] = θ̇
    v
end

@propagate_inbounds function rand_configuration!(q::AbstractVector, ::Revolute)
    q[1] = randn()
    nothing
 end

@propagate_inbounds function joint_transform(jt::Revolute, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D, q::AbstractVector)
    aa = AngleAxis(q[1], jt.axis[1], jt.axis[2], jt.axis[3], false)
    Transform3D(frame_after, frame_before, convert(RotMatrix3{eltype(aa)}, aa))
end

@propagate_inbounds function joint_twist(jt::Revolute, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector, v::AbstractVector)
    angular = jt.axis * v[1]
    Twist(frame_after, frame_before, frame_after, angular, zero(angular))
end

@inline function bias_acceleration(jt::Revolute, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector, v::AbstractVector)
    S = promote_eltype(jt, q, v)
    zero(SpatialAcceleration{S}, frame_after, frame_before, frame_after)
end

@propagate_inbounds function joint_spatial_acceleration(jt::Revolute, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector, v::AbstractVector, vd::AbstractVector)
    S = promote_eltype(jt, q, v, vd)
    angular = convert(SVector{3, S}, jt.axis * vd[1])
    SpatialAcceleration(frame_after, frame_before, frame_after, angular, zero(angular))
end

@inline function motion_subspace(jt::Revolute, frame_after::CartesianFrame3D, frame_before::CartesianFrame3D,
        q::AbstractVector)
    S = promote_eltype(jt, q)
    angular = SMatrix{3, 1, S}(jt.axis)
    linear = zero(SMatrix{3, 1, S})
    GeometricJacobian(frame_after, frame_before, frame_after, angular, linear)
end

@inline function constraint_wrench_subspace(jt::Revolute, joint_transform::Transform3D)
    S = promote_eltype(jt, joint_transform)
    R = convert(RotMatrix3{S}, jt.rotation_from_z_aligned)
    Rcols12 = R[:, SVector(1, 2)]
    angular = hcat(Rcols12, zero(SMatrix{3, 3, S}))
    linear = hcat(zero(SMatrix{3, 2, S}), R)
    WrenchMatrix(joint_transform.from, angular, linear)
end

@propagate_inbounds function joint_torque!(τ::AbstractVector, jt::Revolute, q::AbstractVector, joint_wrench::Wrench)
    τ[1] = dot(angular(joint_wrench), jt.axis)
    nothing
end

@inline function velocity_to_configuration_derivative_jacobian(jt::Revolute, q::AbstractVector)
    T = promote_eltype(jt, q)
    @SMatrix([one(T)])
end

@inline function configuration_derivative_to_velocity_jacobian(jt::Revolute, q::AbstractVector)
    T = promote_eltype(jt, q)
    @SMatrix([one(T)])
end


text:
"""
    theoretical_memory_bandwidth(; device::CuDevice=CUDA.device(); verbose=true)
Estimates the theoretical maximal GPU memory bandwidth in GiB/s.
"""
function theoretical_memory_bandwidth(dev::CuDevice=CUDA.device(); verbose=true, io::IO=stdout)
    max_mem_clock_rate =
        CUDA.attribute(dev, CUDA.CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE) * 1000 # in Hz
    max_mem_bus_width =
        CUDA.attribute(dev, CUDA.CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH) / 8.0 # in bytes
    max_bw = 2.0 * max_mem_clock_rate * max_mem_bus_width * 2^(-30)
    if verbose
        printstyled(io, "Theoretical Maximal Memory Bandwidth (GiB/s):\n"; bold=true)
        print(io, " └ max: ")
        printstyled(io, round(max_bw; digits=1), "\n"; color=:green, bold=true)
    end
    return max_bw
end

"""
    memory_bandwidth([memsize; kwargs...])
Tries to estimate the peak memory bandwidth of a GPU in GiB/s by measuring the time
it takes to perform a memcpy of a certain amount of data (as specified by `memsize`).

**Keyword arguments:**
* `device` (default: `CUDA.device()`): CUDA device to be used.
* `dtype` (default: `Cchar`): element type of the vectors.
* `verbose` (default: `true`): toggle printing.
* `io` (default: `stdout`): set the stream where the results should be printed.

See also: [`memory_bandwidth_scaling`](@ref).
"""
function memory_bandwidth(
    memsize::UnitPrefixedBytes=GiB(0.5);
    dtype=Cchar,
    verbose=true,
    DtoDfactor=true,
    device=CUDA.device(),
    io::IO=stdout,
    kwargs...,
)
    device!(device) do
        N = Int(bytes(memsize) ÷ sizeof(dtype))
        mem_gpu = CUDA.rand(dtype, N)
        mem_gpu2 = CUDA.rand(dtype, N)

        # if verbose
        #     gpu = device(mem_gpu)
        #     println("Memsize: $(Base.format_bytes(sizeof(mem_gpu)))")
        #     println("GPU: ", gpu, " - ", name(gpu), "\n")
        # end

        return _perform_memcpy(
            mem_gpu, mem_gpu2; title="Memory", DtoDfactor, verbose, io=io, kwargs...
        )
    end
end

"""
    memory_bandwidth_scaling() -> datasizes, bandwidths
Measures the memory bandwidth (via [`memory_bandwidth`](@ref)) as a function of data size.
If `verbose=true` (default), displays a unicode plot. Returns the considered data sizes and GiB/s.
For further options, see [`memory_bandwidth`](@ref).
"""
function memory_bandwidth_scaling(;
    device=CUDA.device(), sizes=logspace(1, exp2(30), 10), verbose=true, io::IO=stdout, kwargs...
)
    bandwidths = zeros(length(sizes))
    for (i, s) in enumerate(sizes)
        bandwidths[i] = memory_bandwidth(B(s); device=device, verbose=false, kwargs...)
        clear_gpu_memory(device)
    end
    if verbose
        peak_val, idx = findmax(bandwidths)
        peak_size = sizes[idx]
        p = UnicodePlots.lineplot(
            sizes,
            bandwidths;
            xlabel="data size",
            ylabel="GiB/s",
            title=string(
                "Peak: ", round(peak_val; digits=2), " GiB/s (size = $(bytes(peak_size)))"
            ),
            xscale=:log2,
        )
        UnicodePlots.lineplot!(p, [peak_size, peak_size], [0.0, peak_val]; color=:red)
        println(io) # top margin
        println(io,p)
        println(io) # bottom margin
    end
    return sizes, bandwidths
end