using Distributed
addprocs(15)
@everywhere include("rijke_tangent_state_estimation.jl")
filenames = Array{String,1}(undef, nworkers())
for i = 1:nworkers()
	filenames[i] = string("../data/rijke_exp9_", string(i), ".jld")
end
pmap(assimilate_parameter_and_trajectory, filenames)


function fix_text!(ax, default=:small)
    sc = ax.scene[OldAxis]
    if default == :big
        sc[:names, :textsize] = (25, 25, 25)
        sc[:ticks, :textsize][] = (15, 20, 20)
        sc[:frame, :linewidth][] = (180, 180, 180)
    else
        sc[:names, :textsize] = (10, 10, 10)
        sc[:ticks, :textsize][] = (10, 10, 10)
        sc[:frame, :linewidth][] = (180, 180, 180)
    end
end

function default3d_plot!(fig, f::ScalarField, cmap; kwargs...)
    if !haskey(kwargs, :levels)
        kwargs = Iterators.flatten((kwargs, (:levels=>6,)))
    end

    contour(fig[1,1], f,
        alpha = 0.3, transparency = true, colormap=cmap; kwargs...)
end

function default3d_plot!(fig, f::ScalarVariable, cmap; kwargs...)
    size = get(kwargs, :size3d, 35)

    scattervariable(fig[1,1], f, size=size)
end

function section_plot!(ax, f::ScalarField, slice_dir, idx_value; kwargs...)
    dim = dir_to_idx(slice_dir)
    levels = get(kwargs, :levels, 6)
    f_section = @lift(slice(f, dim, $idx_value))
    labels = string.(filter(i->i≠slice_dir, (:x,:y,:z)))

    plt = contour!(ax, f_section, levels=levels)

    ax.xlabel = labels[1]
    ax.ylabel = labels[2]
    ax.aspect = DataAspect()

    return plt
end

function section_plot!(ax, f::ScalarVariable, slice_dir, idx_value; kwargs...)
    dim = dir_to_idx(slice_dir)
    ϵ = get(kwargs, :ϵ, 0.2)
    grid = getdomain(f)
    sort!(f, dim)

    f_section = @lift slice(f, dim, $idx_value, ϵ)
    labels = string.(filter(i->i≠slice_dir, (:x,:y,:z)))

    plt = scattervariable!(ax, f_section, size=2)

    ax.xlabel = labels[1]
    ax.ylabel = labels[2]
    ax.aspect = DataAspect()

    return plt
end

function add_legend!(fig, section_plt, ::ScalarField; kwargs...)
    levels = get(kwargs, :levels, 6)
    legend = Colorbar(fig, section_plt,
        label = "Field values",
        width = Relative(3/4),
        height = 10,
        vertical = false,
        tellheight = true,
        colormap = cgrad(:viridis, levels, categorical = true))

    fig[2,1:2] = legend
end

function add_legend!(fig, section_plt, ::ScalarVariable; kwargs...)
    # legend = Colorbar(fig, section_plt,
    #     label = "Variable values",
    #     width = Relative(3/4),
    #     height = 10,
    #     vertical = false,
    #     tellheight = true,
    #     colormap = cgrad(:viridis))

    # fig[2,1:2] = legend
end

function slider(f::ScalarField, slice_dir; kwargs...)
    dim = dir_to_idx(slice_dir)
    grid = getdomain(f)
    slider_domain = grid[dim]
    sl = Slider(axes(slider_domain, 1))
    sl_label = @lift round(slider_domain[$(sl.value)], digits=2)

    return sl, sl_label
end

function slider(f::ScalarVariable, slice_dir; kwargs...)
    dim = dir_to_idx(slice_dir)
    grid = getdomain(f)
    slider_domain = grid[dim]
    sl = Slider(1:2*10^4:length(slider_domain))
    @debug sl.value
    sl_label = @lift round(grid[dim][$(sl.value)], digits=2)

    return sl, sl_label
end

function section_widget(f, slice_dir=:x; kwargs...)
    fig = Figure(resolution=(800, 600))
    f_approx = approximate_field(f)
    grid = getdomain(f_approx)
    dim = dir_to_idx(slice_dir)
    cmap = RGBAf0.(to_colormap(:viridis, 50), 1.0)
    cmap[1:2] .= RGBAf0(0,0,0,0)

    lvl_slider = Slider(1:20)
    # Contour plot
    ax, plt = default3d_plot!(fig, f_approx, cmap; kwargs...)
    fix_text!(ax)

    # Slider
    sl, sl_label = slider(f_approx, slice_dir; kwargs...)

    section_ax = Axis(fig[1,2])
    # Section through the field
    section_plt = section_plot!(section_ax, f_approx, slice_dir, sl.value; kwargs...)

    add_legend!(fig, section_plt, f; kwargs...)

    # Transparent plane

    plane_origin = [minimum(grid)...]
    ws = maximum(grid) .- plane_origin
    ws[dim] = 0

    p = lift(sl.value) do idx
        autolimits!(section_ax)
        update!(fig.scene)

        new_origin = grid[dim][idx]
        plane_origin[dim] = new_origin
        origin = Point3f0(plane_origin...)
        widths = Point3f0(ws...)
        @debug "Plane with origin at $origin"
        FRect3D(origin, widths)
    end

    a = RGBAf0(0,0,0,0)
    c = RGBAf0(0.2, 0.2, 1.0, 1.0)
    img = AbstractPlotting.ImagePattern([c a; a c])
    mesh!(fig[1, 1], p; color = RGBAf0(0.2, 0.2, 1.0, 0.3), transparency=true)

    # Build app
    app = JSServe.App() do
        dom = md"""
        $(fig.scene)
        $(DOM.div("Slice at: ", sl_label, sl))
        """

        DOM.div(JSServe.TailwindCSS, JSServe.Styling, dom)
    end

end


using HypothesisTests, Test

mutable struct TestTest <: HypothesisTests.HypothesisTest end

@testset "Common" begin
@test_throws DimensionMismatch HypothesisTests.check_same_length([1], [])
@test_throws ArgumentError HypothesisTests.check_alpha(0.0)

result = HypothesisTests.population_param_of_interest(TestTest())
@test result[1] == "not implemented yet"
@test isnan(result[2])
end


import Pandas

@testset "summarystats" begin
    rng = MersenneTwister(42)
    nchains, ndraws = 4, 10
    idata = convert_to_inference_data(Dict(
        "a" => randn(rng, nchains, ndraws),
        "b" => randn(rng, nchains, ndraws, 3, 4),
    ))

    @test summarystats(idata) isa Pandas.DataFrame
    @test summarystats(idata; fmt = "wide") isa Pandas.DataFrame
    @test summarystats(idata; fmt = "long") isa Pandas.DataFrame
    s = summarystats(idata)
    @test "a" ∈ Pandas.index(s)
    @test "b" ∉ Pandas.index(s)
    @test "b[0,0]" ∉ Pandas.index(s)
    @test "b[1,1]" ∈ Pandas.index(s)
    @test "b[0,0]" ∈ Pandas.index(summarystats(idata; index_origin = 0))

    s2 = summarystats(idata; fmt = "xarray")
    @test s2 isa ArviZ.Dataset
end

@testset "ArviZ.summary" begin
    rng = MersenneTwister(42)
    nchains, ndraws = 4, 10
    data = Dict(
        "a" => randn(rng, nchains, ndraws),
        "b" => randn(rng, nchains, ndraws, 3, 4),
    )

    @test ArviZ.summary(data) isa Pandas.DataFrame
end


function f(x)
a=x*5
b=a+3
c=a-4
d=b/c
end


module TestTimescaleDB

using Test
using TSML

function test_timescaledb()
  tsdb = TimescaleDB()
  mpipeline = Pipeline(Dict(
         :transformers => [tsdb]
     )
  )
  fit!(mpipeline)
  respipe = transform!(mpipeline)
  @test sum(respipe.Value) ≈ 560997.0
  @test nrow(respipe) == 48349
  @test ncol(respipe) == 2
end
@testset "TimescaleDB" begin
  test_timescaledb()
end

end


#!/bin/julia

using DelimitedFiles

#=
FAZER

- Criar funçao para gerar amostras com algumas palavras
- Criar um módulo de limpeza de texto chamado `Limpeza.jl` (copiar as Regex do projeto tfidf-simples)
- Criar ferramenta que descobre qual o idioma de uma frase com base na lista de bigramas coletados.

=#

home = pwd()
homecorpus = home * "/" * "corpus"

function get_linguas()::Array{String,1}
	lis = readdir(homecorpus)
	return lis
end

function get_textos(lingua::String)::Array{String,1}
	return readdir(homecorpus * "/" * lingua)
end

function ler(lingua::String, texto::String)::String
	arquivo = homecorpus * "/" * lingua * "/" * texto
	return read(arquivo, String)
end

function lerlinhas(lingua::String, texto::String)::Array{String,1}
	arquivo = homecorpus * "/" * lingua * "/" * texto
	return readlines(arquivo)
end

function pushdict!(dict::Dict, item::Any)
	if (item in keys(dict))
		dict[item] += 1
	else
		push!(dict, item => 1)
	end
end

function bigramas!(tokens::Dict, texto::String)::Dict{Any,Any}

	words = split(lowercase(texto))

	for w in words

		#   | margem |        meio        | margem |
		#   |    1   |        0..n        |   1    |
		#
		# a:     [[a_,1,1]]
		# de:    [[de,1,1]]
		# com:   [[co,1,0], [om,0,1]]
		# para:  [[pa,1,0], [ar,0,0], [ra,0,1]]
		# arara: [[ar,1,0], [ra,0,0], [ar,0,0], [ra,0,1]]
		chars = split(w, "")
		len = length(chars)
		margin = 1 # margen direita e esquerda

		if len == 1
			token = Any[]
			push!(token, chars[1])
			push!(token, nothing)
			push!(token, 1)
			push!(token, 1)
			pushdict!(tokens, token)
			continue
		end

		i = 1
		while i < len
			token = Any[]
			push!(token, chars[i])
			push!(token, chars[i+1])
			if i <= margin # margem direita
				push!(token, i)
			else
				push!(token, 0)
			end
			if i >= len - margin # margem esquerda
				push!(token, len - i)
			else
				push!(token, 0)
			end
			pushdict!(tokens, token)
			i += 1
		end
	end

	return tokens
end

function main()
	linguas = Dict{String, Dict}(map(x -> Pair(x, Dict()), get_linguas()))
	for lingua in keys(linguas)
		tokens = Dict()
		textos = get_textos(lingua)
		for texto in textos
			bigramas!(tokens, ler(lingua, texto))
		end
		linguas[lingua] = tokens
		println(lingua, ": ", length(tokens))
	end

	# CSV: [lingua, countsum]
	L = map(k -> [k, sum(values(linguas[k])) ], collect(keys(linguas)))
	writedlm( "linguas.csv",  L, ',')

	# CSV: [char1, char2, margin1, margin2, df, idf]
	d = Dict()
	foreach(x -> foreach(y -> pushdict!(d, y), keys(x[2])), collect(linguas))
	B = map(k -> vcat(k[1], k[2], log(length(linguas)/k[2])), collect(d))
	writedlm("bigramas.csv",  B, ',')

end

main()


# AUTO GENERATED FILE - DO NOT EDIT

export vtk_polydata

"""
    vtk_polydata(;kwargs...)
    vtk_polydata(children::Any;kwargs...)
    vtk_polydata(children_maker::Function;kwargs...)

A PolyData component.
PolyData is exposing a vtkPolyData to a downstream filter
It takes the following set of properties:
  - points: [x, y, z, x, y, z, ...],
  - verts: [cellSize, pointId0, pointId1, ..., cellSize, pointId0, ...]
  - lines: [cellSize, pointId0, pointId1, ..., cellSize, pointId0, ...]
  - polys: [cellSize, pointId0, pointId1, ..., cellSize, pointId0, ...]
  - strips: [cellSize, pointId0, pointId1, ..., cellSize, pointId0, ...]
Cell connectivity helper property:
  - connectivity: 'manual', // [manual, points, triangles, strips]
Keyword arguments:
- `children` (Array of a list of or a singular dash component, string or numbers | a list of or a singular dash component, string or number; optional)
- `id` (String; optional): The ID used to identify this component.
- `port` (Real; optional): downstream connection port
- `points` (Array of Reals; optional): xyz coordinates
- `verts` (Array of Reals; optional): verts cells
- `lines` (Array of Reals; optional): lines cells
- `polys` (Array of Reals; optional): polys cells
- `strips` (Array of Reals; optional): strips cells
- `connectivity` (String; optional): Type of connectivity `manual` or implicit such as `points`, `triangles`, `strips`
"""
function vtk_polydata(; kwargs...)
        available_props = Symbol[:children, :id, :port, :points, :verts, :lines, :polys, :strips, :connectivity]
        wild_props = Symbol[]
        return Component("vtk_polydata", "PolyData", "dash_vtk", available_props, wild_props; kwargs...)
end

vtk_polydata(children::Any; kwargs...) = vtk_polydata(;kwargs..., children = children)
vtk_polydata(children_maker::Function; kwargs...) = vtk_polydata(children_maker(); kwargs...)


"""
Provides a definition of a FULL method for constructing Koza trees.
"""
type KozaFullBuilderDefinition <: KozaBuilderDefinition
  """
  The minimum depth of trees created according to this method.
  """
  min_depth::Int

  """
  The maximum depth of trees created according to this method.
  """
  max_depth::Int

  KozaFullBuilderDefinition() = new(1, 8)
  KozaFullBuilderDefinition(min::Int, max::Int) = new(min, max)
end

"""
Implements the Koza FULL initialisation method.
"""
type KozaFullBuilder{T <: KozaTree} <: KozaBuilder
  """
  The minimum depth of trees created according to this method.
  """
  min_depth::Int

  """
  The maximum depth of trees created according to this method.
  """
  max_depth::Int

  KozaFullBuilder() = new(1, 8)
  KozaFullBuilder(min::Int, max::Int) = new(min, max)
end

"""
Composes a Koza Full builder, for a given sub-type of Koza Tree, using a
provided definition.
"""
compose!(d::KozaFullBuilderDefinition, t::Type) =
  KozaFullBuilder{t}(d.min_depth, d.max_depth)

"""
Implements the Koza FULL initialisation method, where some `d` is selected
between a specified minimum and maximum depth, and used to generate a full tree
with a depth equal to `d`.

**Properties:**

* `min::Int`, the minimum tree depth.
* `max::Int`, the maximum tree depth.
"""
full() = KozaFullBuilderDefinition()
full(min::Int, max::Int) = KozaFullBuilderDefinition(min, max)
function full(f::Function)
  def = KozaFullBuilderDefinition()
  f(def)
  def
end

"""
Constructs a Koza Tree belonging to a given representation using the FULL
method, according to a set of provided specifics.
"""
function build{T}(b::KozaFullBuilder{T}, r::KozaTreeRepresentation)
  t = T()
  t.root = build(KozaFullBuilder, r, t, 1, rand(b.min_depth:b.max_depth))
  return t
end

build(::Type{KozaFullBuilder},
  r::KozaTreeRepresentation,
  p::KozaParent,
  d::Integer,
  md::Integer
) =
  if d == md
    return sample_terminal(r, p)
  else
    n = sample_non_terminal(r, p)
    for i in 1:arity(n)
      n.children[i] = build(KozaFullBuilder, r, n, d+1, md)
    end
    return n
  end


module Bridges

using VecMathOptInterface
const MOI = VecMathOptInterface
const MOIU = MOI.Utilities

const CI = MOI.ConstraintIndex

include("bridge.jl")
include("set_map.jl")
include("bridge_optimizer.jl")

include("Variable/Variable.jl")
include("Constraint/Constraint.jl")
include("Objective/Objective.jl")

include("flip_sign.jl")
include("soc_rsoc.jl")

include("lazy_bridge_optimizer.jl")
include("debug.jl")

"""
    full_bridge_optimizer(model::MOI.ModelLike, ::Type{T}) where {T}

Returns a [`LazyBridgeOptimizer`](@ref) bridging `model` for every bridge
defined in this package (see below for the few exceptions) and for the
coefficient type `T` in addition to the bridges in the list returned by
`MOI.get(model, MOI.Bridges.ListOfNonstandardBridges{T}())`.

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

!!! note
    The following bridges are not added by `full_bridge_optimizer` except if
    they are in the list returned by `MOI.get(model, MOI.Bridges.ListOfNonstandardBridges{T}())`
    (see the docstrings of the corresponding bridge for the reason they are not
    added):
    * [`Constraint.SOCtoNonConvexQuadBridge`](@ref),
      [`Constraint.RSOCtoNonConvexQuadBridge`](@ref) and
      [`Constraint.SOCtoPSDBridge`](@ref).
    * The subtypes of [`Constraint.AbstractToIntervalBridge`](@ref) (i.e.
      [`Constraint.GreaterToIntervalBridge`](@ref) and
      [`Constraint.LessToIntervalBridge`](@ref)) if `T` is not a subtype of
      `AbstractFloat`.
"""
function full_bridge_optimizer(model::MOI.ModelLike, ::Type{T}) where {T}
    bridged_model = LazyBridgeOptimizer(model)
    for BT in MOI.get(model, ListOfNonstandardBridges{T}())
        add_bridge(bridged_model, BT)
    end
    Variable.add_all_bridges(bridged_model, T)
    Constraint.add_all_bridges(bridged_model, T)
    Objective.add_all_bridges(bridged_model, T)
    return bridged_model
end

"""
    ListOfNonstandardBridges{T}() <: MOI.AbstractOptimizerAttribute

Any optimizer can be wrapped in a [`LazyBridgeOptimizer`](@ref) using
[`full_bridge_optimizer`](@ref). However, by default [`LazyBridgeOptimizer`](@ref)
uses a limited set of bridges that are:

  1. implemented in `MOI.Bridges`
  2. generally applicable for all optimizers.

For some optimizers however, it is useful to add additional bridges, such as
those that are implemented in external packages (e.g., within the solver package
itself) or only apply in certain circumstances (e.g.,
[`Constraint.SOCtoNonConvexQuadBridge`](@ref)).

Such optimizers should implement the `ListOfNonstandardBridges` attribute to
return a vector of bridge types that are added by [`full_bridge_optimizer`](@ref)
in addition to the list of default bridges.

Note that optimizers implementing `ListOfNonstandardBridges` may require
package-specific functions or sets to be used if the non-standard bridges
are not added. Therefore, you are recommended to use
`model = MOI.instantiate(Package.Optimizer; with_bridge_type = T)` instead of
`model = MOI.instantiate(Package.Optimizer)`. See
[`VecMathOptInterface.instantiate`](@ref).

## Examples

### An optimizer using a non-default bridge in `MOI.Bridges`

Solvers supporting [`MOI.ScalarQuadraticFunction`](@ref) can support
[`MOI.SecondOrderCone`](@ref) and [`MOI.RotatedSecondOrderCone`](@ref) by
defining:
```julia
function MOI.get(::MyQuadraticOptimizer, ::ListOfNonstandardBridges{Float64})
    return Type[
        MOI.Bridges.Constraint.SOCtoNonConvexQuadBridge{Float64},
        MOI.Bridges.Constraint.RSOCtoNonConvexQuadBridge{Float64},
    ]
end
```

### An optimizer defining an internal bridge

Suppose an optimizer can exploit specific structure of a constraint, e.g., it
can exploit the structure of the matrix `A` in the linear system of equations
`A * x = b`.

The optimizer can define the function:
```julia
struct MatrixAffineFunction{T} <: MOI.AbstractVectorFunction
    A::SomeStructuredMatrixType{T}
    b::Vector{T}
end
```
and then a bridge
```julia
struct MatrixAffineFunctionBridge{T} <: MOI.Constraint.AbstractBridge
    # ...
end
# ...
```
from `VectorAffineFunction{T}` to the `MatrixAffineFunction`. Finally, it
defines:
```julia
function MOI.get(::Optimizer{T}, ::ListOfNonstandardBridges{T}) where {T}
    return Type[MatrixAffineFunctionBridge{T}]
end
```
"""
struct ListOfNonstandardBridges{T} <: MOI.AbstractOptimizerAttribute end

attribute_value_type(::ListOfNonstandardBridges) = Vector{Type}

MOI.is_copyable(::ListOfNonstandardBridges) = false

MOI.get_fallback(model::MOI.ModelLike, ::ListOfNonstandardBridges) = Type[]

print_num_bridges(io::IO, ::Variable.EmptyMap) = nothing

print_num_bridges(io::IO, ::Constraint.EmptyMap) = nothing

print_num_bridges(io::IO, ::Objective.EmptyMap) = nothing

function print_num_bridges(io::IO, B::Variable.Map)
    s(n) = n == 1 ? "" : "s"
    indent = " "^get(io, :indent, 0)
    n = length(B)
    print(io, "\n$(indent)with $(n) variable bridge$(s(n))")
    return
end

function print_num_bridges(io::IO, B::Constraint.Map)
    s(n) = n == 1 ? "" : "s"
    indent = " "^get(io, :indent, 0)
    n = length(B)
    print(io, "\n$(indent)with $(n) constraint bridge$(s(n))")
    return
end

function print_num_bridges(io::IO, B::Objective.Map)
    s(n) = n == 1 ? "" : "s"
    indent = " "^get(io, :indent, 0)
    n = length(B)
    print(io, "\n$(indent)with $(n) objective bridge$(s(n))")
    return
end

function Base.show(io::IO, B::AbstractBridgeOptimizer)
    MOIU.print_with_acronym(io, summary(B))
    print_num_bridges(io, Variable.bridges(B))
    print_num_bridges(io, Constraint.bridges(B))
    print_num_bridges(io, Objective.bridges(B))
    if :model in propertynames(B)
        indent = " "^get(io, :indent, 0)
        print(io, "\n$(indent)with inner model ")
        show(IOContext(io, :indent => get(io, :indent, 0) + 2), B.model)
    end
    return
end

end


import Adapt
import .CUDA

struct OneHotArray{T<:Integer, L, N, var"N+1", I<:Union{T, AbstractArray{T, N}}} <: AbstractArray{Bool, var"N+1"}
  indices::I
end
OneHotArray{T, L, N, I}(indices) where {T, L, N, I} = OneHotArray{T, L, N, N+1, I}(indices)
OneHotArray(indices::T, L::Integer) where {T<:Integer} = OneHotArray{T, L, 0, 1, T}(indices)
OneHotArray(indices::I, L::Integer) where {T, N, I<:AbstractArray{T, N}} = OneHotArray{T, L, N, N+1, I}(indices)

_indices(x::OneHotArray) = x.indices
_indices(x::Base.ReshapedArray{<: Any, <: Any, <: OneHotArray}) =
  reshape(parent(x).indices, x.dims[2:end])

const OneHotVector{T, L} = OneHotArray{T, L, 0, 1, T}
const OneHotMatrix{T, L, I} = OneHotArray{T, L, 1, 2, I}

OneHotVector(idx, L) = OneHotArray(idx, L)
OneHotMatrix(indices, L) = OneHotArray(indices, L)

function _show_elements(x::OneHotArray)
  xbool = convert(Array{Bool}, cpu(x))
  xrepr = join(split(repr(MIME("text/plain"), xbool; context= :limit => true), "\n")[2:end], "\n")

  return xrepr
end

function Base.show(io::IO, ::MIME"text/plain", x::OneHotArray{<:Any, L, <:Any, N, I}) where {L, N, I}
  join(io, string.(size(x)), "×")
  print(io, " Flux.OneHotArray{")
  join(io, string.([L, N, I]), ",")
  println(io, "}:")
  print(io, _show_elements(x))
end
function Base.show(io::IO, ::MIME"text/plain", x::OneHotVector{T, L}) where {T, L}
  print(io, string.(length(x)))
  print(io, "-element Flux.OneHotVector{")
  join(io, string.([L, T]), ",")
  println(io, "}:")
  print(io, _show_elements(x))
end

# use this type so reshaped arrays hit fast paths
# e.g. argmax
const OneHotLike{T, L, N, var"N+1", I} =
  Union{OneHotArray{T, L, N, var"N+1", I},
        Base.ReshapedArray{Bool, var"N+1", <:OneHotArray{T, L, <:Any, <:Any, I}}}

_isonehot(x::OneHotArray) = true
_isonehot(x::Base.ReshapedArray{<:Any, <:Any, <:OneHotArray{<:Any, L}}) where L = (size(x, 1) == L)

Base.size(x::OneHotArray{<:Any, L}) where L = (Int(L), size(x.indices)...)

_onehotindex(x, i) = (x == i)

Base.getindex(x::OneHotVector, i::Integer) = _onehotindex(x.indices, i)
Base.getindex(x::OneHotVector{T, L}, ::Colon) where {T, L} = x

Base.getindex(x::OneHotArray, i::Integer, I...) = _onehotindex.(x.indices[I...], i)
Base.getindex(x::OneHotArray{<:Any, L}, ::Colon, I...) where L = OneHotArray(x.indices[I...], L)
Base.getindex(x::OneHotArray{<:Any, <:Any, <:Any, N}, ::Vararg{Colon, N}) where N = x
Base.getindex(x::OneHotArray, I::CartesianIndex{N}) where N = x[I[1], Tuple(I)[2:N]...]

_onehot_bool_type(x::OneHotLike{<:Any, <:Any, <:Any, N, <:Union{Integer, AbstractArray}}) where N = Array{Bool, N}
_onehot_bool_type(x::OneHotLike{<:Any, <:Any, <:Any, N, <:CuArray}) where N = CuArray{Bool, N}

function Base.cat(x::OneHotLike{<:Any, L}, xs::OneHotLike{<:Any, L}...; dims::Int) where L
  if isone(dims) || any(x -> !_isonehot(x), (x, xs...))
    return cat(map(x -> convert(_onehot_bool_type(x), x), (x, xs...))...; dims = dims)
  else
    return OneHotArray(cat(_indices(x), _indices.(xs)...; dims = dims - 1), L)
  end
end

Base.hcat(x::OneHotLike, xs::OneHotLike...) = cat(x, xs...; dims = 2)
Base.vcat(x::OneHotLike, xs::OneHotLike...) = cat(x, xs...; dims = 1)

# optimized concatenation for matrices and vectors of same parameters
Base.hcat(x::T, xs::T...) where {L, T <: OneHotLike{<:Any, L, <:Any, 2}} =
  OneHotMatrix(reduce(vcat, _indices.(xs); init = _indices(x)), L)
Base.hcat(x::T, xs::T...) where {L, T <: OneHotLike{<:Any, L, <:Any, 1}} =
  OneHotMatrix(reduce(vcat, _indices.(xs); init = _indices(x)), L)

batch(xs::AbstractArray{<:OneHotVector{<:Any, L}}) where L = OneHotArray(_indices.(xs), L)

Adapt.adapt_structure(T, x::OneHotArray{<:Any, L}) where L = OneHotArray(adapt(T, _indices(x)), L)

Base.BroadcastStyle(::Type{<:OneHotArray{<: Any, <: Any, <: Any, N, <: CuArray}}) where N = CUDA.CuArrayStyle{N}()

Base.map(f, x::OneHotLike) = Base.broadcast(f, x)

Base.argmax(x::OneHotLike; dims = Colon()) =
  (_isonehot(x) && dims == 1) ?
    reshape(CartesianIndex.(_indices(x), CartesianIndices(_indices(x))), 1, size(_indices(x))...) :
    invoke(argmax, Tuple{AbstractArray}, x; dims = dims)

"""
    onehot(l, labels[, unk])

Return a `OneHotVector` where only first occourence of `l` in `labels` is `1` and
all other elements are `0`.

If `l` is not found in labels and  `unk` is present, the function returns
`onehot(unk, labels)`; otherwise the function raises an error.

# Examples
```jldoctest
julia> Flux.onehot(:b, [:a, :b, :c])
3-element Flux.OneHotVector{3,UInt32}:
 0
 1
 0

julia> Flux.onehot(:c, [:a, :b, :c])
3-element Flux.OneHotVector{3,UInt32}:
 0
 0
 1
```
"""
function onehot(l, labels)
  i = something(findfirst(isequal(l), labels), 0)
  i > 0 || error("Value $l is not in labels")
  OneHotVector{UInt32, length(labels)}(i)
end

function onehot(l, labels, unk)
  i = something(findfirst(isequal(l), labels), 0)
  i > 0 || return onehot(unk, labels)
  OneHotVector{UInt32, length(labels)}(i)
end

"""
    onehotbatch(ls, labels[, unk...])

Return a `OneHotMatrix` where `k`th column of the matrix is `onehot(ls[k], labels)`.

If one of the input labels `ls` is not found in `labels` and `unk` is given,
return [`onehot(unk, labels)`](@ref) ; otherwise the function will raise an error.

# Examples
```jldoctest
julia> Flux.onehotbatch([:b, :a, :b], [:a, :b, :c])
3×3 Flux.OneHotArray{3,2,Vector{UInt32}}:
 0  1  0
 1  0  1
 0  0  0
```
"""
onehotbatch(ls, labels, unk...) = batch([onehot(l, labels, unk...) for l in ls])

"""
    onecold(y[, labels = 1:length(y)])

Inverse operations of [`onehot`](@ref).

# Examples
```jldoctest
julia> Flux.onecold([true, false, false], [:a, :b, :c])
:a

julia> Flux.onecold([0.3, 0.2, 0.5], [:a, :b, :c])
:c
```
"""
onecold(y::AbstractVector, labels = 1:length(y)) = labels[argmax(y)]
function onecold(y::AbstractArray, labels = 1:size(y, 1))
  indices = _fast_argmax(y)
  xs = isbits(labels) ? indices : collect(indices) # non-bit type cannot be handled by CUDA

  return map(xi -> labels[xi[1]], xs)
end

_fast_argmax(x::AbstractArray) = dropdims(argmax(x; dims = 1); dims = 1)
function _fast_argmax(x::OneHotLike)
  if _isonehot(x)
    return _indices(x)
  else
    return _fast_argmax(convert(_onehot_bool_type(x), x))
  end
end

@nograd OneHotArray, onecold, onehot, onehotbatch

function Base.:(*)(A::AbstractMatrix, B::OneHotLike{<:Any, L}) where L
  _isonehot(B) || return invoke(*, Tuple{AbstractMatrix, AbstractMatrix}, A, B)
  size(A, 2) == L || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
  return A[:, onecold(B)]
end
for wrapper in [:Adjoint, :Transpose]
  @eval begin
    function Base.:*(A::$wrapper{<:Any, <:AbstractMatrix{T}}, b::OneHotVector{<:Any, L}) where {L, T}
      size(A, 2) == L ||
        throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))

      return A[:, onecold(b)]
    end

    function Base.:*(A::$wrapper{<:Number, <:AbstractVector{T}}, b::OneHotVector{<:Any, L}) where {L, T}
      size(A, 2) == L ||
        throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))

      return A[onecold(b)]
    end
  end
end


# demo
# ====

# fibonacci
# ---------

fib(n) = n ≤ 2 ? n : fib(n-1) + fib(n-2)

fib(1000)   # never terminates in ordinal execution
fib(m)      # undef var
fib("1000") # obvious type error

# language features
# -----------------

# user-defined types
struct Ty{T}
    fld::T
end

function foo(a)
    v = Ty(a)
    return bar(v)
end

# macros will be expanded
@inline bar(n::T)     where {T<:Number} = n < 0 ? zero(T) : one(T)
@inline bar(v::Ty{T}) where {T<:Number} = bar(v.fdl) # typo "fdl"
@inline bar(v::Ty)                      = bar(convert(Number, v.fld))

foo(1.2)
foo("1") # `String` can't be converted to `Number`


function SeisProcessHeaders(in,out,functions,parameters;group="some",key=[],ntrace=1000000,update_tracenum=true)

	if (group=="all")
		h1 = SeisReadHeaders(in,group=group,key=key,itrace=1,ntrace=ntrace)
		for ifunc = 1 : length(functions)
			f = functions[ifunc]
			p = parameters[ifunc]
			h2 = f(h1;p...)
			h1 = copy(h2)
		end
		SeisWriteHeaders(out,h1)
	else
		itrace_in = 1
		itrace_out = 1
		NX = GetNumTraces(in)
		while itrace_in <= NX
			nx = NX - itrace_in + 1
			ntrace = nx > ntrace ? ntrace : nx
			h1 = SeisReadHeaders(in,group=group,key=key,itrace=itrace_in,ntrace=ntrace)
			num_traces_in = length(h1)
			for ifunc = 1 : length(functions)
				f = functions[ifunc]
				p = parameters[ifunc]
				h2 = f(h1;p...)
				h1 = copy(h2)
			end
			num_traces_out = length(h1)
			SeisWriteHeaders(out,h1,itrace=itrace_out,update_tracenum=update_tracenum)
			itrace_in += num_traces_in
			itrace_out += num_traces_out
		end
	end
end


#
# This file is part of the Simulate.jl Julia package, MIT license
#
# Paul Bayer, 2019
#
# This is a Julia package for discrete event simulation
#

"""
    𝐶

`𝐶` (𝐶 = `\\itC`+`tab`) is the default simulation clock. If you do one
simulation at a time, you can use it for time keeping.

# Examples

```jldoctest
julia> using Simulate

julia> reset!(𝐶)
"clock reset to t₀=0.0, sampling rate Δt=0.0."

julia> 𝐶  # default clock
Clock thread 1 (+ 0 ac): state=Simulate.Idle(), t=0.0 , Δt=0.0 , prc:0
  scheduled ev:0, cev:0, sampl:0
```
"""
const 𝐶 = Clock()

# ----------------------------------------------------
# setting and getting clock parameters
# ----------------------------------------------------

"""
    setUnit!(clk::Clock, new::FreeUnits)

set a clock to a new time unit in `Unitful`. If necessary convert
current clock times to the new unit.

# Arguments
- `clk::Clock`
- `new::FreeUnits`: new is one of `ms`, `s`, `minute` or `hr` or another Unitful
    `Time` unit.

# Examples
```jldoctest
julia> using Simulate, Unitful

julia> import Unitful: Time, s, minute, hr

julia> c = Clock(t0=60)     # setup a new clock with t0=60
Clock thread 1 (+ 0 ac): state=Simulate.Undefined(), t=60.0 , Δt=0.0 , prc:0
  scheduled ev:0, cev:0, sampl:0

julia> tau(c) # current time is 60.0 NoUnits
60.0

julia> setUnit!(c, s)       # set clock unit to Unitful.s
60.0 s

julia> tau(c) # current time is now 60.0 s
60.0 s

julia> setUnit!(c, minute)  # set clock unit to Unitful.minute
1.0 minute

julia> tau(c)               # current time is now 1.0 minute
1.0 minute

julia> isa(tau(c), Time)
true

julia> uconvert(s, tau(c))  # ... which can be converted to other time units
60.0 s

julia> tau(c).val           # it has a value of 1.0
1.0

julia> c.time               # internal clock time is set to 1.0 (a Float64)
1.0

julia> c.unit               # internal clock unit is set to Unitful.minute
minute
```
"""
function setUnit!(clk::Clock, new::FreeUnits)
    if isa(1new, Time)
        if clk.unit == new
            println("clock is already set to $new")
        elseif clk.unit == NoUnits
            clk.unit = new
        else
            old = clk.unit
            clk.unit = new
            fac = uconvert(new, 1*old).val
            clk.time *= fac
            clk.end_time *= fac
            clk.tev *= fac
            clk.Δt *= fac
            clk.tn *= fac
        end
    else
        clk.unit = NoUnits
    end
    tau(clk)
end

"""
```
tau(clk::Clock=𝐶)
```
Return the current simulation time.

# Examples

```jldoctest
julia> using Simulate

julia> reset!(𝐶)
"clock reset to t₀=0.0, sampling rate Δt=0.0."
julia> tau() # gives the central time
0.0
```
"""
tau(clk::Clock=𝐶) = clk.unit == NoUnits ? clk.time : clk.time*clk.unit

"""
```
sync!(clk::Clock, to::Clock=𝐶)
```
Force a synchronization of two clocks. Change all registered times of
`clk` accordingly. Convert or force clk.unit to to.unit.
"""
function sync!(clk::Clock, to::Clock=𝐶)
    if (clk.unit == NoUnits) | (clk.unit == to.unit)
        fac = 1
    elseif to.unit == NoUnits
        println(stderr, "Warning: deleted time unit without conversion")
        fac = 1
    else
        fac = uconvert(to.unit, 1clk.unit).val
    end
    Δt = to.time - clk.time*fac
    clk.time = clk.time*fac + Δt
    clk.unit = to.unit
    clk.tn  = clk.tn*fac + Δt
    clk.tev  = clk.tev*fac + Δt
    clk.end_time = clk.end_time*fac + Δt
    clk.Δt = to.Δt
    evq = PriorityQueue{DiscreteEvent,Float64}()
    for (ev, t) ∈ pairs(clk.sc.events)
        evq[ev] = t*fac + Δt
    end
    clk.sc.events = evq
    clk
end

"""
```
reset!(clk::Clock, Δt::Number=0; t0::Number=0, hard::Bool=true, unit=NoUnits)
```
reset a clock

# Arguments
- `clk::Clock`
- `Δt::Number=0`: time increment
- `t0::Float64=0` or `t0::Time`: start time
- `hard::Bool=true`: time is reset, all scheduled events and sampling are
    deleted. If hard=false, then only time is reset, event and
    sampling times are adjusted accordingly.
- `unit=NoUnits`: the Time unit for the clock after reset. If a `Δt::Time` is
    given, its Time unit goes into the clock Time unit. If only t0::Time is given,
    its Time unit goes into the clock time unit.

# Examples

```jldoctest
julia> using Simulate, Unitful

julia> import Unitful: s

julia> c = Clock(1s, t0=60s)
Clock thread 1 (+ 0 ac): state=Simulate.Undefined(), t=60.0 s, Δt=1.0 s, prc:0
  scheduled ev:0, cev:0, sampl:0

julia> reset!(c)
"clock reset to t₀=0.0, sampling rate Δt=0.0."

julia> c
Clock thread 1 (+ 0 ac): state=Simulate.Idle(), t=0.0 , Δt=0.0 , prc:0
  scheduled ev:0, cev:0, sampl:0
```
"""
function reset!(clk::Clock, Δt::Number=0;
                t0::Number=0, hard::Bool=true, unit=NoUnits)
    if  isa(1unit, Time)
        Δt = isa(Δt, Time) ? uconvert(unit, Δt).val : Δt
        t0 = isa(t0, Time) ? uconvert(unit, t0).val : t0
    elseif isa(Δt, Time)
        unit = Unitful.unit(1Δt)
        Δt = Δt.val
        t0 = isa(t0, Time) ? uconvert(unit, t0).val : t0
    elseif isa(t0, Time)
        unit = Unitful.unit(t0)
        t0 = t0.val
    else
        nothing
    end
    if hard
        clk.state = Idle()
        clk.time = t0
        clk.unit = unit
        clk.tn = t0
        clk.tev = t0
        clk.end_time = t0
        clk.evcount = 0
        clk.scount = 0
        clk.Δt = Δt
        clk.sc.events = PriorityQueue{DiscreteEvent,Float64}()
        clk.sc.cevents = DiscreteCond[]
        clk.processes = Dict{Any, Prc}()
        clk.sc.samples = Sample[]
    else
        sync!(clk, Clock(Δt, t0=t0, unit=unit))
    end
    "clock reset to t₀=$(float(t0*unit)), sampling rate Δt=$(float(Δt*unit))."
end

"""
    tadjust(clk::Clock, t::Unitful.Time) :: Float64

Adjust/convert `t` given according to clock settings and return a Float64 value.
"""
function tadjust(clk::Clock, t::Unitful.Time) :: Float64
    if clk.unit == NoUnits
        println(stderr, "Warning: clock has no time unit, ignoring units")
        return t.val
    else
        return uconvert(clk.unit, t).val
    end
end

"""
```
sample_time!([clk::Clock], Δt::Number)
```
set the clock's sample rate starting from now (`tau(clk)`).

# Arguments
- `clk::Clock`: if not supplied, set the sample rate on 𝐶,
- `Δt::Number`: sample rate, time interval for sampling
"""
function sample_time!(clk::Clock, Δt::Number)
    clk.Δt = Δt isa Unitful.Time ? tadjust(clk, Δt) : Δt
    clk.tn = clk.time + clk.Δt
end
sample_time!(Δt::Number) = sample_time!(𝐶, Δt)

"""
```
periodic!([clk::Clock], ex::Union{Expr, Fun}, Δt::Number=clk.Δt;
        scope::Module=Main, spawn=false)
```
Register a function or expression for periodic execution at the clock`s sample rate.

# Arguments
- `clk::Clock`: if not supplied, it samples on 𝐶,
- `ex::Union{Expr, Fun}`: an expression or function,
- `Δt::Number=clk.Δt`: set the clock's sampling rate, if no Δt is given, it takes
    the current sampling rate, if that is 0, it calculates one,
- `scope::Module=Main`: optional, an evaluation scope for a given expression.
"""
function periodic!(clk::Clock, ex::Union{Expr, Fun}, Δt::Number=clk.Δt;
                 scope::Module=Main, spawn=false)
    clk.Δt = Δt == 0 ? scale(clk.end_time - clk.time)/100 : Δt
    assign(clk, Sample(ex, scope), spawn ? spawnid(clk) : 0)
end
periodic!(ex::Union{Expr, Fun}, Δt::Number=𝐶.Δt; kw...) = periodic!(𝐶, ex, Δt; kw...)

"Is a Clock busy?"
busy(clk::Clock) = clk.state == Busy()

# ----------------------------------------------------
# step! transition functions for Clocks
# ----------------------------------------------------

"""
    step!(clk::Clock, ::Undefined, ::Init)

initialize a clock.
"""
function step!(clk::Clock, ::Undefined, ::Init)
    clk.state = Idle()
end

"""
    step!(clk::Clock, ::Undefined, σ::Union{Step,Run})

if uninitialized, initialize and then Step or Run.
"""
function step!(clk::Clock, ::Undefined, σ::Union{Step,Run})
    step!(clk, clk.state, Init(0))
    step!(clk, clk.state, σ)
end

"""
    setTimes(clk::Clock)

set clock times for next event or sampling action. The internal clock times
`clk.tev` and `clk.tn` must always be set to be at least `clk.time`.
"""
function setTimes(clk::Clock)
    if !isempty(clk.sc.events)
        clk.tev = nextevtime(clk)
        clk.tn = clk.Δt > 0 ? clk.time + clk.Δt : clk.time
    else
        clk.tn = clk.Δt > 0 ? clk.time + clk.Δt : clk.time
        clk.tev = clk.tn
    end
end

"""
    step!(c::Clock, ::Union{Idle,Busy,Halted}, ::Step)

step forward to next tick or scheduled event.

At a tick evaluate 1) all sampling functions or expressions, 2) all conditional
events, then 3) if an event is encountered, trigger the event.

The internal clock times `c.tev` and `c.tn` are always at least `c.time`.
"""
step!(c::Clock, ::Union{Idle,Halted}, ::Step) = do_step!(c)

"""
    do_run(c::Clock, Δt::Float64)

Run a clock for a time Δt.
"""
function do_run!(c::Clock, Δt::Float64)
    c.end_time = c.time + Δt
    c.evcount = 0
    c.scount = 0
    setTimes(c)
    while any(i->(c.time < i ≤ c.end_time), (c.tn, c.tev))
        do_step!(c)
        if c.state == Halted()
            return c.end_time
        end
    end
    c.time = c.end_time
end

"""
    step!(clk::Clock, ::Idle, σ::Run)

Run a simulation for a given duration.

The duration is given with `Run(duration)`. Call scheduled events and evaluate
sampling expressions at each tick in that timeframe.
"""
function step!(clk::Clock, ::Idle, σ::Run)
    tend = do_run!(clk, σ.duration)
    if clk.state == Halted()
        return
    end
    # catch remaining events scheduled for the end_time
    while (length(clk.sc.events) ≥ 1) && (nextevtime(clk) ≤ tend + Base.eps(tend)*10)
        step!(clk, clk.state, Step())
        tend = nextfloat(tend)
    end
    clk.time = clk.end_time

    "run! finished with $(clk.evcount) clock events, $(clk.scount) sample steps, simulation time: $(clk.time)"
end

"""
    step!(clk::Clock, ::Busy, ::Stop)

Stop the clock.
"""
function step!(clk::Clock, ::Busy, ::Stop)
    clk.state = Halted()
    "Halted after $(clk.evcount) events, simulation time: $(clk.time)"
end

"""
    step!(clk::Clock, ::Halted, ::Resume)

Resume a halted clock.
"""
function step!(clk::Clock, ::Halted, ::Resume)
    clk.state = Idle()
    step!(clk, clk.state, Run(clk.end_time - clk.time))
end

"""
    step!(clk::Clock, q::ClockState, σ::ClockEvent)

catch all step!-function.
"""
function step!(clk::Clock, q::ClockState, σ::ClockEvent)
    println(stderr, "Warning: undefined transition ",
            "$(typeof(clk)), ::$(typeof(q)), ::$(typeof(σ)))\n",
            "maybe, you should reset! the clock!")
end

# --------------------------------------------------------
# Clock user interface functions
# --------------------------------------------------------

"""
    run!(clk::Clock, duration::Number)

Run a simulation for a given duration. Call scheduled events and evaluate
sampling expressions at each tick in that timeframe.
"""
function run!(clk::Clock, duration::Number)
    duration = duration isa Unitful.Time ? tadjust(clk, duration) : duration
    step!(clk, clk.state, Run(duration))
end

"""
    incr!(clk::Clock)

Take one simulation step, execute the next tick or event.
"""
incr!(clk::Clock) = step!(clk, clk.state, Step())

"""
    stop!(clk::Clock)

Stop a running simulation.
"""
stop!(clk::Clock) = step!(clk, clk.state, Stop())

"""
    resume!(clk::Clock)

Resume a halted simulation.
"""
resume!(clk::Clock) = step!(clk, clk.state, Resume())

"""
    init!(clk::Clock)

initialize a clock.
"""
init!(clk::Clock) = step!(clk, clk.state, Init(""))


include("../imports.jl")

@testset ExtendedTestSet "projective" begin
    image = Image(rand(RGB, 32, 32))
    tfm = Project(Translation(20, 10))

    @test_nowarn apply(tfm, image)
    timage = apply(tfm, image)
    @test timage.bounds.rs == (21:52, 11:42)

    @test_nowarn apply!(timage, tfm, image)

end


using DataDeps

register(DataDep(
  "blog-feedback",
  "https://archive.ics.uci.edu/ml/datasets/BlogFeedback",
  "https://archive.ics.uci.edu/ml/machine-learning-databases/00304/BlogFeedback.zip",
  "1ba74e5ad920f7cd037502b2968581cc695146a226a73eff52fe8ad875ed4bcf",
  post_fetch_method=(path -> begin
    unpack(path)
    new_path = joinpath(dirname(path), "blogData_train.csv")
    mv(new_path, path)

    UCIData.process_dataset(path,
        target_index=281,
        feature_indices=1:280,
        categoric_indices=63:262,
    )
  end),
))


"""
    QXZoo.jl - implementation of circuit generator for
    different quantum algorithms.
"""
module QXZoo

using Reexport

include("gates/gate_ops.jl")
include("gates/default_gates.jl")
include("gates/gate_map.jl")
include("gates/composite_gates.jl")

include("circuits/circuit.jl")
include("algorithms/Algorithms.jl")

@reexport using QXZoo.GateMap
@reexport using QXZoo.GateOps
@reexport using QXZoo.DefaultGates
@reexport using QXZoo.CompositeGates
@reexport using QXZoo.Circuit

# Initialise default gates into cache
GateMap.init_cache()

end


"""
    FIRITools

`FIRITools` is a collection of convenience functions for working with
Faraday-International Reference Ionosphere (FIRI) model profiles.

The underlying FIRI-2018 model `DATA`, `HEADER`, and `ALTITUDE`'s can be accessed
as e.g. `FIRITools.HEADER`. The unit of `ALTITUDE` is meters and `DATA` is electrons per
cubic meter.

Only the function `firi` is exported, but also see `FIRITools.quantile` and
`FIRITools.extrapolate`.

# References

Friedrich, M., Pock, C., & Torkar, K. (2018).
FIRI-2018, an updated empirical model of the lower ionosphere.
Journal of Geophysical Research: Space Physics, 123, 6737-6751.
[doi: 10.1029/2018JA025437](https://doi.org/10.1029/2018JA025437)
"""
module FIRITools

using Artifacts, Statistics, DelimitedFiles
using TypedTables
using Interpolations

export firi

const MIN_DENSITY = 1e-4

const DF = readdlm(joinpath(artifact"firi", "firi2018.csv"), ',')

function parseheader()
    df = Table(
        code=convert(Vector{String}, DF[2,2:end]),
        month=convert(Vector{Int}, DF[3,2:end]),
        chi=convert(Vector{Int}, DF[5,2:end]),  # deg
        lat=convert(Vector{Int}, DF[6,2:end]),  # deg
        f10_7=convert(Vector{Int}, DF[7,2:end])
    )

    return df
end

const HEADER = parseheader()

# Start altitude at 60 km even though table has down to 55 km (60 km is stated lower limit)
const DATA = convert(Matrix{Float64}, DF[14:end-2, 2:end])  # last 2 rows have Missing
const ALTITUDE = convert(Vector{Int}, DF[14:end-2, 1])*1000  # convert to m
@assert length(ALTITUDE) == size(DATA, 1) "ALTITUDE does not match DATA"

"""
    values(s)

Return the FIRI model values of parameter `s`.
"""
function values(s)
    if s in ("f10_7", :f10_7)
        return [75, 130, 200]  # == unique(HEADER.f10_7)
    elseif s in ("chi", :chi)
        return [0, 30, 45, 60, 75, 80, 85, 90, 95, 100, 130]
    elseif s in ("lat", :lat)
        return [0, 15, 30, 45, 60]
    elseif s in ("month", :month)
        return [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
    else
        throw(ArgumentError("$s is not an FIRI model parameter"))
    end
end

"""
    firi(;chi=(0, 130), lat=(0, 60), f10_7=(75, 200), month=(1, 12))

Return the average FIRI profile across model parameters selected by the keyword arguments.
"""
function firi(;chi=(0, 130), lat=(0, 60), f10_7=(75, 200), month=(1, 12))
    profiles = selectprofiles(chi=chi, lat=lat, f10_7=f10_7, month=month)
    profile = mean(profiles, dims=2)
    replace!(x -> x<MIN_DENSITY ? MIN_DENSITY : x, profile)
    return dropdims(profile, dims=2)
end

"""
    firi(chi, lat; f10_7=(75, 200), month=(1, 12))

Return the average FIRI profile across the model parameters selected by the keyword
arguments and interpolated at solar zenith angle `chi` and latitude `lat`.
"""
function firi(chi, lat; f10_7=(75, 200), month=(1, 12))
    profiles = selectprofiles(chi, lat, f10_7=f10_7, month=month)
    profile = mean(profiles, dims=2)
    replace!(x -> x<MIN_DENSITY ? MIN_DENSITY : x, profile)
    return dropdims(profile, dims=2)
end

"""
    selectprofiles(;chi=(0, 130), lat=(0, 60), f10_7=(75, 200), month=(1, 12))

Return matrix of selected model profiles.
"""
function selectprofiles(;chi=(0, 130), lat=(0, 60), f10_7=(75, 200), month=(1, 12))
    mask = trues(length(HEADER))

    mask .&= select(HEADER.chi, chi) .&
             select(HEADER.lat, lat) .&
             select(HEADER.f10_7, f10_7) .&
             select(HEADER.month, month)

    return DATA[:,mask]
end

"""
    selectprofiles(chi, lat; f10_7=(75, 200), month=(1, 12))

Return matrix of selected model profiles interpolated at solar zenith angle `chi` and
latitude `lat`.
"""
function selectprofiles(chi, lat; f10_7=(75, 200), month=(1, 12))
    issorted(unique(HEADER.chi)) || throw(DomainError("chi column of HEADER is not sorted"))
    issorted(unique(HEADER.lat)) || throw(DomainError("lat column of HEADER is not sorted"))

    if chi > 130
        @warn "`chi` greater than 130° uses `chi = 130°`" maxlog=3
        chi = oftype(chi, 130)
    end

    if chi < 0 || lat < 0
        throw(ArgumentError("`chi` or `lat` below 0° is not supported. See Friedrich et al., “FIRI-2018”."))
    end

    mask = trues(length(HEADER))

    mask .&= select(HEADER.f10_7, f10_7) .&
             select(HEADER.month, month)

    chibnds = twoclosest(HEADER.chi, chi)
    latbnds = twoclosest(HEADER.lat, lat)

    # Mask all but the 2 closest values of chi and lat to the target chi and lat
    mask .&= .!isnothing.(indexin(HEADER.chi, chibnds)) .&
             .!isnothing.(indexin(HEADER.lat, latbnds))

    N = count(mask)
    maskidxs = findall(mask)

    # We need to explicitly handle 1D vs 2D interpolation (if lat, or chi don't need interp)
    if isequal(chibnds...) && isequal(latbnds...)
        profiles = DATA[:, maskidxs]
    elseif isequal(chibnds...)
        profiles = Matrix{Float64}(undef, length(ALTITUDE), N÷2)

        j = 1
        for i = 1:2:N
            for a in axes(DATA, 1)  # altitude
                d = DATA[a,maskidxs[[i, i+1]]]
                itp = LinearInterpolation(latbnds, d, extrapolation_bc=Line())
                profiles[a,j] = itp(lat)
            end
            j += 1
        end
    elseif isequal(latbnds...)
        profiles = Matrix{Float64}(undef, length(ALTITUDE), N÷2)

        j = 1
        for i = 1:2:N
            for a in axes(DATA, 1)  # altitude
                d = DATA[a,maskidxs[[i, i+1]]]
                itp = LinearInterpolation(chibnds, d, extrapolation_bc=Line())
                profiles[a,j] = itp(chi)
            end
            j += 1
        end
    else
        # Although not a general solution, we know the format of Chi guarantees that occurences
        # of the two closest values will be in pairs and that both Chi and Lat are sorted
        # so every 4 indices in the mask form a group that covers
        # [(chi_low, lat_low) (chi_high, lat_low);
        #  (chi_low, lat_high) (chi_high, lat_high)]
        profiles = Matrix{Float64}(undef, length(ALTITUDE), N÷4)

        j = 1
        for i = 1:4:N
            for a in axes(DATA, 1)  # altitude
                d = [DATA[a,maskidxs[i]] DATA[a,maskidxs[i+2]];
                     DATA[a,maskidxs[i+1]] DATA[a,maskidxs[i+3]]]
                itp = LinearInterpolation((chibnds, latbnds), d, extrapolation_bc=Line())
                profiles[a,j] = itp(chi, lat)
            end
            j += 1
        end
    end

    return profiles
end

"""
    twoclosest(s, v)

Return:
    - the two closest values of vector `s` to `v` if `v` is outside of the range of
        values in `s` (extrapolating)
    - or the closest value on each side of `v` in column `s` if interpolating
    - or `[v, v]` if `v` is exactly in `s`.

The two values are sorted.
"""
function twoclosest(s, v)
    uvals = unique(s)

    if v in uvals
        return [v, v]
    elseif v >= maximum(uvals) || v <= minimum(uvals)
        # We're extrapolating or at a min/max
        I = sortperm(abs.(v .- uvals))
        return sort([uvals[I[1]], uvals[I[2]]])  # 2 closest matches
    else
        # We're interpolating, make sure we return the closest on each side of `v`
        udiffs = v .- uvals
        v1 = minimum(x->x > 0 ? x : Inf, udiffs)  # closest val with positive diff
        v2 = -minimum(x->x < 0 ? abs(x) : Inf, udiffs)  # closest val with negative diff

        I1 = findfirst(isequal(v1), udiffs)
        I2 = findfirst(isequal(v2), udiffs)
        return sort([uvals[I1], uvals[I2]])
    end
end

"""
    select(h, x::Tuple)

Return `x[1] .<= h .<= x[2]`.
"""
function select(h, x::T) where T<:Tuple
    first(x) .<= h .<= last(x)
end

"""
    select(h, x)

Return `x .== h` if `h` is in `x`, otherwise `nothing`.
"""
function select(h, x::T) where T<:Number
    x in h || @warn "$x is not in model parameters. Looking for interpolating form of `selectprofiles`?"

    return x .== h
end

"""
    quantile(p; chi=(0, 130), lat=(0, 60), f10_7=(75, 200), month=(1, 12))

Compute the quantile(s) `p` at each `ALTITUDE` with data columns selected by the keyword
arguments.
"""
function quantile(p; chi=(0, 130), lat=(0, 60), f10_7=(75, 200), month=(1, 12))
    profiles = selectprofiles(chi=chi, lat=lat, f10_7=f10_7, month=month)
    profile = _quantile(profiles, p)

    replace!(x -> x<MIN_DENSITY ? MIN_DENSITY : x, profile)

    if p isa Number
        # make `profile` a Vector for consistency with `firi`
        profile = dropdims(profile, dims=2)
    end

    return profile
end

"""
    quantile(chi, lat, p; f10_7=(75, 200), month=(1, 12))

Compute the quantile(s) `p` at each `ALTITUDE` with data columns selected by keyword
arguments and interpolated to solar zenith angle `chi` and latitude `lat`.
"""
function quantile(chi, lat, p; f10_7=(75, 200), month=(1, 12))
    profiles = selectprofiles(chi, lat, f10_7=f10_7, month=month)
    profile = _quantile(profiles, p)

    replace!(x -> x<MIN_DENSITY ? MIN_DENSITY : x, profile)

    if p isa Number
        # make `profile` a Vector for consistency with `firi`
        profile = dropdims(profile, dims=2)
    end

    return profile
end

function _quantile(profiles, p)
    profile = Array{Float64, 2}(undef, length(ALTITUDE), length(p))
    for j in eachindex(p)
        i = 1
        for r in eachrow(profiles)
            profile[i,j] = Statistics.quantile(r, p[j])
            i += 1
        end
    end
    return profile
end

"""
    extrapolate([z,] profile, newz)

Return `profile`, originally defined at sorted altitudes `z`, exponentially interpolated
and extrapolated to altitudes of `newz`.

If `z` is not specified, it is assumed that `profile` is defined at `FIRITools.ALTITUDE`.
"""
function extrapolate(z::AbstractRange, profile::AbstractVector, newz::AbstractRange)
    length(z) == length(profile) ||
        throw(ArgumentError("`z` and `profile` must be equal lengths."))

    itp = LinearInterpolation(z, log.(profile), extrapolation_bc=Line())
    itpprofile = itp(newz)

    return exp.(itpprofile)
end
extrapolate(profile, newz) = extrapolate(first(ALTITUDE):1000:last(ALTITUDE), profile, newz)

function extrapolate(z::AbstractRange, profile::AbstractMatrix, newz::AbstractRange)
    expprofile = Vector{eltype(profile)}()
    nrows = 0
    for i in axes(profile, 2)
        p = extrapolate(z, view(profile,:,i), newz)
        append!(expprofile, p)
        nrows = length(p)
    end
    return reshape(expprofile, nrows, size(profile, 2))
end

# NOTE: at this time, models() isn't very important because only 1 model is supported
"""
    models()

Return the available models.
"""
models() = first.(splitext.(filter(endswith(".csv"), readdir(artifact"firi"))))

end  # module


@testset "1297.maximum-number-of-occurrences-of-a-substring.jl" begin
    @test max_freq("aababcaab", 2, 3, 4) == 2
    @test max_freq("aaaa", 1, 3, 3) == 2
    @test max_freq("aabcabcab", 2, 2, 3) == 3

    @test max_freq("abcde", 2, 3, 3) == 0
end


mutable struct DEStats
  nf::Int
  nw::Int
  nsolve::Int
  naccept::Int
  nreject::Int
  maxeig::Float64
end

DEStats() = DEStats(ntuple(_->0, 5)..., 0.0)

function Base.show(io::IO, s::DEStats)
  println(io, summary(s))
  println(io, "Number of function evaluations: $(s.nf)")
  println(io, "Number of W matrix evaluations: $(s.nw)")
  println(io, "Number of linear solves:        $(s.nsolve)")
  println(io, "Number of accepted steps:       $(s.naccept)")
  print(io,   "Number of rejected steps:       $(s.nreject)")
  iszero(s.maxeig) || print(io, "\nMaximum eigenvalue recorded:    $(s.maxeig)")
end


"""
    plotabm(model::ABM{<: ContinuousSpace}; ac, as, am, kwargs...)
    plotabm(model::ABM{<: DiscreteSpace}; ac, as, am, kwargs...)

Plot the `model` as a `scatter`-plot, by configuring the agent shape, color and size
via the keywords `ac, as, am`.
These keywords can be constants, or they can be functions, each accepting an agent
and outputting a valid value for color/shape/size.

The keyword `scheduler = model.scheduler` decides the plotting order of agents
(which matters only if there is overlap).

The keyword `offset` is a function with argument `offest(a::Agent)`.
It targets scenarios where multiple agents existin within
a grid cell as it adds an offset (same type as `agent.pos`) to the plotted agent position.

All other keywords are propagated into `Plots.scatter` and the plot is returned.
"""
function plotabm(
    model::ABM{<:Union{GridSpace,ContinuousSpace}};
    ac = "#765db4",
    as = 10,
    am = :circle,
    scheduler = model.scheduler,
    offset = nothing,
    kwargs...,
)

    ids = scheduler(model)
    colors = typeof(ac) <: Function ? [ac(model[i]) for i in ids] : ac
    sizes = typeof(as) <: Function ? [as(model[i]) for i in ids] : as
    markers = typeof(am) <: Function ? [am(model[i]) for i in ids] : am
    if offset == nothing
        pos = [model[i].pos for i in ids]
    else
        pos = [model[i].pos .+ offset(model[i]) for i in ids]
    end

    scatter(
        pos;
        markercolor = colors,
        markersize = sizes,
        markershapes = markers,
        label = "",
        markerstrokewidth = 0.5,
        markerstrokecolor = :black,
        kwargs...,
    )
end


export isafter

using Dates

function isafter(str::AbstractString, date::AbstractString=Dates.format(Dates.now(), "Y-mm-ddTHH:MM:SS"))::Bool
    comparison = Dates.DateTime(date)
    original = Dates.DateTime(str)

    return original > comparison
end


using Chain
using DataFrames
using DataFramesMeta
using Dates
using StatsBase

"""
Takes a DataFrame of Twitter data and a date manipulation function to group on and calculates volume

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Returns
-  `DataFrame`: Containing date function output and n volume
"""
function get_volume_by(twit_data::DataFrame, date_func::Function)
    volume_by = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        combine(groupby(_, :temp), nrow)
    end

    rename!(volume_by, :nrow => "n")
    rename!(volume_by, :temp => Symbol(date_func))

    return volume_by
end

"""
Takes a DataFrame of Twitter data and a date manipulation function to group on and calculates cumulative volume

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Returns
-  `DataFrame`: Containing date function output and n cumulative volume
"""
function get_cumulative_volume_by(twit_data::DataFrame, date_func::Function)
    volume_by = get_volume_by(twit_data, date_func)
    volume_by = @transform(volume_by, :n = cumsum(:n))

    return volume_by
end

"""
Takes a DataFrame of Twitter data and a date manipulation function to group on and calculates
number of quotes and retweets

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function
- `manip_func::Function`: Column manipulation function (default mean)

# Returns
-  `DataFrame`: Containing date function output, nquotes and nretweets
"""
function get_interactions_by(twit_data::DataFrame, date_func::Function, manip_func::Function = mean)
    interactions = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        @by(:temp, :nquotes = manip_func(:quote_count), :nretweets = manip_func(:retweet_count), :nreplies = manip_func(:reply_count))
    end

    rename!(interactions, :temp => Symbol(date_func))

    return interactions
end

"""
Takes a DataFrame of Twitter data and extracts mention data

# Arguments
- `twit_data::DataFrame`: Twitter data

# Returns
-  `DataFrame`: Containing original twitter data plus mentions and Nmentions columns
"""
function add_mentions(twit_data::DataFrame)
    twit_data.mentions = get_mentions.(lowercase.(twit_data.text))
    twit_data.n_mentions = length.(twit_data.mentions)

    return twit_data
end

"""
Takes a DataFrame of Twitter data and a date manipulation function to group on and calculates
number of mentions

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function
- `manip_func::Function`: Column manipulation function (default mean)

# Returns
-  `DataFrame`: Containing date function output and number of mentions
"""
function get_mentions_by(twit_data::DataFrame, date_func::Function, manip_func::Function = mean)
    twit_data = add_mentions(twit_data)

    mentions_by = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        @by(:temp, :nmentions = manip_func(:n_mentions))
    end

    rename!(mentions_by, :temp => Symbol(date_func))

    return mentions_by
end

"""
Takes a DataFrame of Twitter data and plots the 'ratio' or number retweets against the number of quotes

# Arguments
- `twit_data::DataFrame`: Twitter data

# Output
A plot of retweets against quotes
"""
function generate_ratio_plot(twit_data::DataFrame)
    ratio_df = data(twit_data) *
        mapping(:quote_count => "N Quotes",:retweet_count => "N Retweets")
    return ratio_df
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the volume of tweets grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of tweet volume by date function grouping
"""
function generate_volume_by_plot(twit_data::DataFrame, date_func::Function)
    volume = data(get_volume_by(twit_data, date_func)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :n => "Number of posts")

    return volume
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the cumulative volume of tweets grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of cumulative tweet volume by date function grouping
"""
function generate_cumulative_volume_by_plot(twit_data::DataFrame, date_func::Function)
    cumulative_volume = data(get_cumulative_volume_by(twit_data, date_func)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :n => "Cumulative number of posts")

    return cumulative_volume
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the average number of mentions per tweet grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of average mentions by date function grouping
"""
function generate_avg_mentions_by_plot(twit_data::DataFrame, date_func::Function)
    mentions = data(get_mentions_by(twit_data, date_func, mean)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nmentions => "Average number of mentions per tweet")

    return mentions
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the total number of mentions grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of total mentions by date function grouping
"""
function generate_total_mentions_by_plot(twit_data::DataFrame, date_func::Function)
    mentions = data(get_mentions_by(twit_data, date_func, sum)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nmentions => "Total number of mentions")

    return mentions
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the average number of retweets per tweet grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of average retweets by date function grouping
"""
function generate_avg_retweets_by_plot(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, mean), Symbol(date_func), :nretweets)
    mentions = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nretweets => "Average number of retweets per tweet")

    return mentions
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the total number of grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of total retweets by date function grouping
"""
function generate_total_retweets_by_plot(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nretweets)
    mentions = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nretweets => "Total number of retweets")

    return mentions
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the average number of quotes per tweet grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of average quote tweets by date function grouping
"""
function generate_avg_quotes_by_plot(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, mean), Symbol(date_func), :nquotes)
    mentions = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nquotes => "Average number of quotes per tweet")

    return mentions
end

"""
Takes a DataFrame of Twitter data and date manipulation function and produces a plot
of the total number of quotes grouped by the date manipulation function

# Arguments
- `twit_data::DataFrame`: Twitter data
- `date_func::Function`: Date manipulation function

# Output
A plot of total quote tweets by date function grouping
"""
function generate_total_quotes_by_plot(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nquotes)
    mentions = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nquotes => "Total number of quotes")

    return mentions
end

function get_unique_users(twit_data::DataFrame, date_func::Function)
    users = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        @by(:temp, :nusers = length(unique(:author_id)))
    end

    rename!(users, :temp => Symbol(date_func))

    return users
end

function generate_unique_users(twit_data::DataFrame, date_func::Function)
    df = get_unique_users(twit_data, date_func)
    users = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nusers => "N unique users")

    return users
end

function get_quote_volume(twit_data::DataFrame, date_func::Function)
    quotes = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        @by(:temp, :nquotes = count(i -> (i=="quoted" || "quoted" in split(strip(i,['[', ']']),", ")), :sourcetweet_type))
    end

    rename!(quotes, :temp => Symbol(date_func))

    return quotes
end

function generate_quote_volume(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nquotes)
    quotes = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nquotes => "N quote tweets")

    return quotes
end

function get_retweet_volume(twit_data::DataFrame, date_func::Function)
    df = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        @by(:temp, :nretweets = count(i -> (i=="retweeted" || "retweeted" in split(strip(i,['[', ']']),", ")), :sourcetweet_type))
    end

    rename!(df, :temp => Symbol(date_func))

    return df
end

function generate_retweet_volume(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nretweets)
    quotes = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nretweets => "N quote tweets")

    return quotes
end

function get_reply_volume(twit_data::DataFrame, date_func::Function)
    df = @chain twit_data begin
        @transform(:temp = date_func.(:date))
        @by(:temp, :nreplies = count(i -> (i!="NA"), :in_reply_to_user_id))
    end

    rename!(df, :temp => Symbol(date_func))

    return df
end

function generate_reply_volume(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nreplies)
    quotes = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nreplies => "N reply tweets")

    return quotes
end

function get_unique_user_rate(twit_data::DataFrame, date_func::Function)
    df = @chain twit_data begin
        @transform(:temp1 = minute.(:date), :temp2 = hour.(:date), :temp3 = dayofyear.(:date))
        @by([:temp1, :temp2, :temp3], :nunique = length(unique(:author_id)))
        @by(:temp3, :nuniqueAvg = mean(:nunique))
    end

    rename!(df, :temp3 => Symbol(date_func))

    return df
end

function generate_unique_user_rate(twit_data::DataFrame, date_func::Function)
    df = get_unique_user_rate(twit_data, date_func)
    quotes = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nuniqueAvg => "Unique user tweet rate")

    return quotes
end

function get_quote_rate(twit_data::DataFrame, date_func::Function)
    is_quote(type) = type == "quoted" || "quoted" in split(strip(type,['[', ']']),", ")
    twit_data = filter(:sourcetweet_type => is_quote, twit_data)
    df = @chain twit_data begin
        @transform(:temp1 = minute.(:date), :temp2 = hour.(:date), :temp3 = dayofyear.(:date))
        @by([:temp1, :temp2, :temp3], :nunique = length(unique(:author_id)))
        @by(:temp3, :nuniqueAvg = mean(:nunique))
    end

    rename!(df, :temp3 => Symbol(date_func))

    return df
end

function generate_quote_rate(twit_data::DataFrame, date_func::Function)
    df = get_quote_rate(twit_data, date_func)
    quotes = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nuniqueAvg => "Quote tweet rate")

    return quotes
end

function get_retweet_rate(twit_data::DataFrame, date_func::Function)
    is_retweet(type) = type == "retweeted" || "retweeted" in split(strip(type,['[', ']']),", ")
    twit_data = filter(:sourcetweet_type => is_retweet, twit_data)
    df = @chain twit_data begin
        @transform(:temp1 = minute.(:date), :temp2 = hour.(:date), :temp3 = dayofyear.(:date))
        @by([:temp1, :temp2, :temp3], :nunique = length(unique(:author_id)))
        @by(:temp3, :nuniqueAvg = mean(:nunique))
    end

    rename!(df, :temp3 => Symbol(date_func))

    return df
end

function generate_retweet_rate(twit_data::DataFrame, date_func::Function)
    df = get_retweet_rate(twit_data, date_func)
    retweets = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nuniqueAvg => "Retweet rate")

    return retweets
end

function get_reply_rate(twit_data::DataFrame, date_func::Function)
    is_reply(type) = type != "NA"
    twit_data = filter(:in_reply_to_user_id => is_reply, twit_data)
    df = @chain twit_data begin
        @transform(:temp1 = minute.(:date), :temp2 = hour.(:date), :temp3 = dayofyear.(:date))
        @by([:temp1, :temp2, :temp3], :nunique = length(unique(:author_id)))
        @by(:temp3, :nuniqueAvg = mean(:nunique))
    end

    rename!(df, :temp3 => Symbol(date_func))

    return df
end

function generate_reply_rate(twit_data::DataFrame, date_func::Function)
    df = get_reply_rate(twit_data, date_func)
    replys = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nuniqueAvg => "reply tweet rate")

    return replys
end

function get_quote_rate_2(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nquotes)
    df = @chain df begin
        @transform(:nquotes = :nquotes/1440)
    end
    return df
end

function generate_quote_rate_2(twit_data::DataFrame, date_func::Function)
    df = get_quote_rate_2(twit_data, date_func)
    replys = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nquotes => "quote tweet rate")

    return replys
end

function get_retweet_rate_2(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nretweets)
    df = @chain df begin
        @transform(:nretweets = :nretweets/1440)
    end
    return df
end

function generate_retweet_rate_2(twit_data::DataFrame, date_func::Function)
    df = get_retweet_rate_2(twit_data, date_func)
    replys = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nretweets => "retweet rate")

    return replys
end

function get_reply_rate_2(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nreplies)
    df = @chain df begin
        @transform(:nreplies = :nreplies/1440)
    end
    print(df)
    return df
end

function generate_reply_rate_2(twit_data::DataFrame, date_func::Function)
    df = get_reply_rate_2(twit_data, date_func)
    replys = data(df) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :nreplies => "reply rate")

    return replys
end

function get_cum_quote_volume(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nquotes)
    df = @transform(df, :n = cumsum(:nquotes))
    return df
end

function generate_cum_quote_volume(twit_data::DataFrame, date_func::Function)
    cumulative_volume = data(get_cum_quote_volume(twit_data, date_func)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :n => "Cumulative number of quote tweets")

    return cumulative_volume
end

function get_cum_retweet_volume(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nretweets)
    df = @transform(df, :n = cumsum(:nretweets))
    return df
end

function generate_cum_retweet_volume(twit_data::DataFrame, date_func::Function)
    cumulative_volume = data(get_cum_retweet_volume(twit_data, date_func)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :n => "Cumulative number of retweets")

    return cumulative_volume
end

function get_cum_reply_volume(twit_data::DataFrame, date_func::Function)
    df = select(get_interactions_by(twit_data, date_func, sum), Symbol(date_func), :nreplies)
    df = @transform(df, :n = cumsum(:nreplies))
    return df
end

function generate_cum_reply_volume(twit_data::DataFrame, date_func::Function)
    cumulative_volume = data(get_cum_reply_volume(twit_data, date_func)) *
        visual(Lines) *
        mapping(Symbol(date_func) => string(date_func), :n => "Cumulative number of replies")

    return cumulative_volume
end

function wrap_in_makie(out_of_AoG, annotations, use_mad = false, fig = Figure())
    ax1 = Axis(fig[1,1],
        xlabel = annotations[1],
        ylabel = annotations[2],
        title = annotations[3]
        )

    draw!(ax1, out_of_AoG)

    if use_mad
        mean_data = mean(out_of_AoG.data[2])
        mad_data = mad(out_of_AoG.data[2])

        hlines!(ax1, mean_data)
        hspan!(ax1, mean_data-mad_data , mean_data+mad_data, color = (:blue, 0.2))
    end

    return fig
end

function cumulative_sorting(df, column)
    @chain df begin
        @subset(:true_type .!= "retweeted")
        groupby(:author_id)
        combine(column => sum => :tot)
        @orderby(-:tot)
        @transform(_,
        :cumulative = cumsum(:tot),
        :rownumber = rownumber.(eachrow(_)))
    end
end

function graph2(df, date_func)
    fig = Figure()

    g1 = wrap_in_makie(generate_quote_volume(df, date_func), ["Day of year", "N quotes", "Daily quote volume"], false, fig[1,1])
    g2 = wrap_in_makie(generate_retweet_volume(df, date_func), ["Day of year", "N retweets", "Daily retweet volume"], false, fig[1,2])
    g3 = wrap_in_makie(generate_reply_volume(df, date_func), ["Day of year", "N replies", "Daily reply volume"], false, fig[2,1])

    return fig
end

function graph4(df, date_func)
    fig = Figure()
    g1 = wrap_in_makie(generate_quote_rate_2(df, date_func), ["Day of year", "Avg n quotes per min", "Daily quote rate"], true, fig[1,1])
    g2 = wrap_in_makie(generate_retweet_rate_2(df, date_func), ["Day of year", "Avg n retweets per min", "Daily retweet rate"], true, fig[1,2])
    g3 = wrap_in_makie(generate_reply_rate_2(df, date_func), ["Day of year", "Avg n replies per min", "Daily reply rate"], true, fig[2,1])

    return fig
end

function graph5(df, date_func)
    fig = Figure()
    g1 = wrap_in_makie(generate_cum_quote_volume(df, date_func), ["Day of year", "Cumulative n quotes", "Daily quote volume"], false, fig[1,1])
    g2 = wrap_in_makie(generate_cum_retweet_volume(df, date_func), ["Day of year", "Cumulative n retweets", "Daily retweet volume"], false, fig[1,2])
    g3 = wrap_in_makie(generate_cum_reply_volume(df, date_func), ["Day of year", "Cumulative n replies", "Daily reply volume"], false, fig[2,1])

    return fig
end


import Base:
    ==,
    hash

##### state, action, and observation spaces and related functions #####

struct RockSampleState
    index::Int64
end
RockSampleState() = RockSampleState(-1)

struct RockSampleAction
    index::Int64
end
RockSampleAction() = RockSampleAction(-1)

struct RockSampleObs
    index::Int64
end
RockSampleObs() = RockSampleObs(-1)

## spaces
struct RockSampleStateSpace
    min_index::Int64
    max_index::Int64
end

struct RockSampleActionSpace
    min_index::Int64
    max_index::Int64
end

struct RockSampleObsSpace
    min_index::Int64
    max_index::Int64
end

## space iterator types
struct RockSampleStateIterator
    min_index::Int64
    max_index::Int64
end

struct RockSampleActionIterator
    min_index::Int64
    max_index::Int64
end

struct RockSampleObsIterator
    min_index::Int64
    max_index::Int64
end

# iterator enabling functions
Base.start(::RockSampleStateIterator)   = RockSampleState(0)
Base.start(::RockSampleActionIterator)  = RockSampleAction(0)
Base.start(::RockSampleObsIterator)     = RockSampleObs(0)

Base.next(::RockSampleStateIterator, state::RockSampleState)        =
    (state, RockSampleState(state.index+1))
Base.next(::RockSampleActionIterator, action::RockSampleAction)     =
    (action, RockSampleAction(action.index+1))
Base.next(::RockSampleObsIterator, obs::RockSampleObs)              =
    (obs, RockSampleObs(obs.index+1))

Base.done(iter::RockSampleStateIterator, state::RockSampleState)    =
    (state.index > iter.max_index)
Base.done(iter::RockSampleActionIterator, action::RockSampleAction) =
    (action.index > iter.max_index)
Base.done(iter::RockSampleObsIterator, obs::RockSampleObs)          =
    (obs.index > iter.max_index)

## iterator creation functions
POMDPs.iterator(space::RockSampleStateSpace)                        =
    RockSampleStateIterator(space.min_index, space.max_index)
POMDPs.iterator(space::RockSampleActionSpace)                       =
    RockSampleActionIterator(space.min_index, space.max_index)
POMDPs.iterator(space::RockSampleObsSpace)                          =
    RockSampleObsIterator(space.min_index, space.max_index)

==(x::RockSampleState, y::RockSampleState)      = (x.index == y.index)
==(x::RockSampleAction, y::RockSampleAction)    = (x.index == y.index)
==(x::RockSampleObs, y::RockSampleObs)          = (x.index == y.index)

Base.hash(x::RockSampleState,  h::Int64 = 0)    = x.index
Base.hash(x::RockSampleAction, h::Int64 = 0)    = x.index
Base.hash(x::RockSampleObs, h::Int64 = 0)       = x.index

######## RockSample type definition ########
mutable struct RockSample <: POMDPs.POMDP{RockSampleState, RockSampleAction, RockSampleObs}
    #problem parameters
    grid_size::Int64
    n_rocks::Int64
    rand_max::Int64
    model_seed ::UInt32 # random seed to construct arbitrary size scenarios
    belief_seed::UInt32 # random seed used in initial belief construction,
                        # if needed

    #problem properties
    n_cells::Int64
    n_actions::Int64
    n_states::Int64
    n_observations::Int64
    robot_start_cell::Int64
    half_eff_distance::Int64
    discount::Float64

    #internal variables and structures
    rock_set_start::Int64
    rock_at_cell::Array{Int64,1}
    cell_to_coords::Array{Vector{Int64},1}
    observation_effectiveness::Array{Float64,2}
    rocks::Array{Int64,1}
    T::Array{RockSampleState,2}
    R::Array{DESPOTReward,2}
    actions::Array{Int64,1} # needed for large problems

    #observation aliases
    #TODO: think how to best convert to const (or just wait for struct fields to be implemented...)
    BAD_OBS::Int64
    GOOD_OBS::Int64
    NONE_OBS::Int64
    TERMINAL_OBS::Int64

    # the const value for 'seed' is meant to provide compatibility with the C++ version of DESPOT
    function RockSample(grid_size::Int64 = 4,
                        n_rocks::Int64 = 4;
                        rand_max::Int64 = 2^31-1,
                        belief_seed::UInt32 = convert(UInt32, 479),
                        model_seed::UInt32  = convert(UInt32, 476),
                        discount::Float64 = 0.95)

          this = new()
          # problem parameters
          this.grid_size = grid_size
          this.n_rocks = n_rocks
          this.rand_max = rand_max
          this.belief_seed = belief_seed
          this.model_seed  = model_seed
          this.discount = discount

          # problem properties
          this.n_cells = grid_size*grid_size
          this.n_actions = n_rocks + 5
          this.n_states = (grid_size*grid_size+1)*(1 << n_rocks)
          this.n_observations = 4
          this.robot_start_cell = 1
          this.half_eff_distance = 20

          #internal variables and structures
          this.rock_set_start = 0
          this.rock_at_cell = Array{Int64}(this.n_cells)
          this.cell_to_coords = Array{Vector{Int64}}(this.n_cells)
          this.observation_effectiveness = Array{Float64}(this.n_cells, this.n_cells)
          this.rocks = Array{Int64}(this.n_rocks)                       # locations
          this.T = Array{RockSampleState}(this.n_states, this.n_actions)
          this.R = Array{Float64}(this.n_states, this.n_actions)
          this.actions = collect(1:n_rocks + 5) # default ordering
          this.BAD_OBS      = 0
          this.GOOD_OBS     = 1
          this.NONE_OBS     = 2
          this.TERMINAL_OBS = 3

          init_problem(this)
          return this
     end
end

POMDPs.state_index(pomdp::RockSample, state::RockSampleState)   = state.index
POMDPs.action_index(pomdp::RockSample, action::RockSampleAction) = action.index
POMDPs.obs_index(pomdp::RockSample, obs::RockSampleObs)       = obs.index

## distribution types
mutable struct RockSampleTransitionDistribution
    pomdp::RockSample
    state::RockSampleState
    action::RockSampleAction
end

mutable struct RockSampleObsDistribution
    pomdp::RockSample
    state::RockSampleState
    action::RockSampleAction
    next_state::RockSampleState
    debug::Int64 #TODO: consider removing

    #TODO: consider removing, not really needed except for debugging
    function RockSampleObsDistribution(pomdp::RockSample,
                                               state::RockSampleState,
                                               action::RockSampleAction,
                                               next_state::RockSampleState,
                                               debug::Int64 = 0)
        this = new()
        this.pomdp = pomdp
        this.state = state
        this.action = action
        this.next_state = next_state
        this.debug = debug

        return this
    end
end

start_state(pomdp::RockSample) =
    RockSampleState(make_state_index(pomdp, pomdp.robot_start_cell, pomdp.rock_set_start))

# Creates a default belief structure to store the problem's initial belief
DESPOT.create_belief(pomdp::RockSample)         =
    ParticleBelief(Array{Particle{RockSampleState}}(0))

create_transition_distribution(pomdp::RockSample)    =
    RockSampleTransitionDistribution(pomdp, RockSampleState(-1), RockSampleAction(-1))
create_observation_distribution(pomdp::RockSample)   =
    RockSampleObsDistribution(pomdp,
                              RockSampleState(-1),
                              RockSampleAction(-1),
                              RockSampleState(-1))

# default initial state distribution represented as a set of generic particles
function POMDPs.initial_state_distribution(pomdp::RockSample)
    n_states = 1 << pomdp.n_rocks
    pool = Array{POMDPToolbox.Particle{RockSampleState}}(0)

    p = 1.0/n_states
    for k = 0:n_states-1 #TODO: can make faster, potentially
        push!(pool, POMDPToolbox.Particle{RockSampleState}(RockSampleState(make_state_index(pomdp, pomdp.robot_start_cell, k)), p))
    end

    return POMDPToolbox.ParticleDistribution(pool)
end

## accessor functions
POMDPs.n_states(pomdp::RockSample)          = pomdp.n_states
POMDPs.n_actions(pomdp::RockSample)         = pomdp.n_actions
POMDPs.n_observations(pomdp::RockSample)    = pomdp.n_observations

# 0-based indexing in the following functions
POMDPs.states(pomdp::RockSample)    = RockSampleStateSpace(0, pomdp.n_states-1)
POMDPs.actions(pomdp::RockSample)   = RockSampleActionSpace(0, pomdp.n_actions-1)
POMDPs.observations(pomdp::RockSample) = RockSampleActionSpace(0, pomdp.n_observations-1)

POMDPs.discount(pomdp::RockSample) = pomdp.discount

## RockSample initialization
function init_4_4(pomdp::RockSample)
  pomdp.rocks[1] = cell_num(pomdp,0,2) # rocks is an array
  pomdp.rocks[2] = cell_num(pomdp,2,2)
  pomdp.rocks[3] = cell_num(pomdp,3,2)
  pomdp.rocks[4] = cell_num(pomdp,3,3)
  pomdp.robot_start_cell = cell_num(pomdp,2,0)
end

function init_7_8(pomdp::RockSample)
  pomdp.rocks[1] = cell_num(pomdp,0,1)
  pomdp.rocks[2] = cell_num(pomdp,1,5)
  pomdp.rocks[3] = cell_num(pomdp,2,2)
  pomdp.rocks[4] = cell_num(pomdp,2,3)
  pomdp.rocks[5] = cell_num(pomdp,3,6)
  pomdp.rocks[6] = cell_num(pomdp,5,0)
  pomdp.rocks[7] = cell_num(pomdp,5,3)
  pomdp.rocks[8] = cell_num(pomdp,6,2)
  pomdp.robot_start_cell = cell_num(pomdp,3,0)
end

function init_11_11(pomdp::RockSample)
  pomdp.rocks[1] = cell_num(pomdp,7,0)
  pomdp.rocks[2] = cell_num(pomdp,3,0)
  pomdp.rocks[3] = cell_num(pomdp,2,1)
  pomdp.rocks[4] = cell_num(pomdp,6,2)
  pomdp.rocks[5] = cell_num(pomdp,7,3)
  pomdp.rocks[6] = cell_num(pomdp,2,3)
  pomdp.rocks[7] = cell_num(pomdp,7,4)
  pomdp.rocks[8] = cell_num(pomdp,2,5)
  pomdp.rocks[9] = cell_num(pomdp,9,6)
  pomdp.rocks[10] = cell_num(pomdp,7,9)
  pomdp.rocks[11] = cell_num(pomdp,1,9)
  pomdp.robot_start_cell = cell_num(pomdp,5,0)
end

function init_15_15(pomdp::RockSample)
  pomdp.rocks[1] = cell_num(pomdp,7,0)
  pomdp.rocks[2] = cell_num(pomdp,3,0)
  pomdp.rocks[3] = cell_num(pomdp,2,1)
  pomdp.rocks[4] = cell_num(pomdp,6,2)
  pomdp.rocks[5] = cell_num(pomdp,7,3)
  pomdp.rocks[6] = cell_num(pomdp,2,3)
  pomdp.rocks[7] = cell_num(pomdp,7,4)
  pomdp.rocks[8] = cell_num(pomdp,2,5)
  pomdp.rocks[9] = cell_num(pomdp,9,6)
  pomdp.rocks[10] = cell_num(pomdp,7,9)
  pomdp.rocks[11] = cell_num(pomdp,1,9)
  pomdp.rocks[12] = cell_num(pomdp,8,11)
  pomdp.rocks[13] = cell_num(pomdp,10,13)
  pomdp.rocks[14] = cell_num(pomdp,9,14)
  pomdp.rocks[15] = cell_num(pomdp,2,12)
  pomdp.robot_start_cell = cell_num(pomdp,5,0)
end

function init_general(pomdp::RockSample, seed::Array{UInt32,1})

    rockIndex::Int64 = 1 # rocks is an array
    if !is_linux()
        srand(seed[1])
    end

    while rockIndex <= pomdp.n_rocks
        if is_linux()
            cell = ccall((:rand_r, "libc"), Int, (Ptr{Cuint},), seed) % pomdp.n_cells
        else
            cell = Base.rand(0:pomdp.rand_max) % pomdp.n_cells
        end

        if findfirst(pomdp.rocks, cell) == 0
            pomdp.rocks[rockIndex] = cell
            rockIndex += 1
        end
    end
    pomdp.robot_start_cell = cell_num(pomdp, round(Integer, pomdp.grid_size/2), 0)
end

function init_problem(pomdp::RockSample)

    pomdp.rocks = Array{Int64}(pomdp.n_rocks)
    seed = Cuint[convert(UInt32, pomdp.model_seed)]

    if pomdp.grid_size == 4 && pomdp.n_rocks == 4
        init_4_4(pomdp)
    elseif pomdp.grid_size == 7 && pomdp.n_rocks == 8
        init_7_8(pomdp)
    elseif pomdp.grid_size == 11 && pomdp.n_rocks == 11
        init_11_11(pomdp)
    elseif pomdp.grid_size == 15 && pomdp.n_rocks == 15
        init_15_15(pomdp)
    else
        init_general(pomdp, seed)
    end

    # Compute rock set start
    pomdp.rock_set_start = 0

    if !is_linux()
        srand(seed[1])
    end

    for i in 0 : pomdp.n_rocks-1
        if is_linux()
            rand_num = ccall((:rand_r, "libc"), Int, (Ptr{Cuint},), seed)
        else #Windows, etc
            rand_num = Base.rand(0:pomdp.rand_max)
        end

        if (rand_num & 1) == 1
            pomdp.rock_set_start |= (1 << i)
        end
    end

    # Fill in cellToCoord and init rock_at_cell mappings
    fill!(pomdp.rock_at_cell, -1)

    for i in 0 : pomdp.n_cells-1
        pomdp.cell_to_coords[i+1] = [trunc(Integer, i/pomdp.grid_size), i % pomdp.grid_size]
    end

    for i in 0 : pomdp.n_rocks-1
        pomdp.rock_at_cell[pomdp.rocks[i+1]+1] = i # rock_at_cell and rocks are arrays
    end

    # T and R - ALL INDICES BELOW ARE OFFSET BY +1 (for 1-based array indexing)
    for cell in 0 : pomdp.n_cells-1
        for rock_set = 0:(1 << pomdp.n_rocks)-1
        s_index = make_state_index(pomdp, cell, rock_set)

            #initialize transition and rewards with default values
            for a_index in 0:pomdp.n_actions-1
                pomdp.T[s_index+1,a_index+1] =
                    RockSampleState(s_index)
                pomdp.R[s_index+1,a_index+1] = 0.
            end

            row, col = pomdp.cell_to_coords[cell+1]
            # North
            if row == 0
                pomdp.T[s_index+1,1] =
                    RockSampleState(s_index)
                pomdp.R[s_index+1,1] = -100.
            else
                pomdp.T[s_index+1,1] =
                    RockSampleState(make_state_index(pomdp, cell_num(pomdp,row-1,col), rock_set))
                pomdp.R[s_index+1,1] = 0
            end

            # South
            if row == pomdp.grid_size-1
                pomdp.T[s_index+1,2] =
                    RockSampleState(s_index)
                pomdp.R[s_index+1,2] = -100.
            else
                pomdp.T[s_index+1,2] =
                    RockSampleState(make_state_index(pomdp, cell_num(pomdp,row+1,col), rock_set))
                pomdp.R[s_index+1,2] = 0.
            end

            # East
            if col == pomdp.grid_size-1
                pomdp.T[s_index+1,3] =
                    RockSampleState(make_state_index(pomdp, pomdp.n_cells, rock_set))
                pomdp.R[s_index+1,3] = 10.
            else
                pomdp.T[s_index+1,3] =
                    RockSampleState(make_state_index(pomdp, cell_num(pomdp,row,col+1), rock_set))
                pomdp.R[s_index+1,3] = 0.
            end

            # West
            if col == 0
                pomdp.T[s_index+1, 4] = RockSampleState(s_index)
                pomdp.R[s_index+1, 4] = -100.
            else
                pomdp.T[s_index+1, 4] =
                    RockSampleState(make_state_index(pomdp, cell_num(pomdp,row,col-1), rock_set))
                pomdp.R[s_index+1, 4] = 0.
            end

            # Sample
            rock = pomdp.rock_at_cell[cell+1] # array
            if rock != -1
                if rock_status(rock, rock_set)
                    pomdp.T[s_index+1, 5] =
                        RockSampleState(make_state_index(pomdp, cell, sample_rock_set(rock, rock_set)))
                    pomdp.R[s_index+1, 5] = +10.
                else
                    pomdp.T[s_index+1, 5] = RockSampleState(s_index)
                    pomdp.R[s_index+1, 5] = -10.
                end
            else
                pomdp.T[s_index+1, 5] = RockSampleState(s_index)
                pomdp.R[s_index+1, 5] = -100.
            end

            # Check
            for a_index in 5:pomdp.n_actions-1
                pomdp.T[s_index+1, a_index+1] = RockSampleState(s_index)
                pomdp.R[s_index+1, a_index+1] = 0.
            end
        end
    end

    # Terminal states
    for k = 0:(1 << pomdp.n_rocks)-1
        s_index = make_state_index(pomdp, pomdp.n_cells, k);
        for a_index in 0:pomdp.n_actions-1
            pomdp.T[s_index+1, a_index+1] = RockSampleState(s_index)
            pomdp.R[s_index+1, a_index+1] = 0.
        end
    end

    # precompute observation effectiveness table
    for i in 0 : pomdp.n_cells-1
        for j in 0 : pomdp.n_cells-1
        agent = pomdp.cell_to_coords[i+1]
        other = pomdp.cell_to_coords[j+1]
        dist = sqrt((agent[1] - other[1])^2 + (agent[2]-other[2])^2)
        pomdp.observation_effectiveness[i+1,j+1] =
            (1 + 2^(-dist / pomdp.half_eff_distance)) * 0.5 # Array indexing starts from 1.
                                                    # Remember to subtract one to go back
        end
    end
end

## utility functions

# True for good rock, false for bad rock, x can be a rock set or state index
rock_status(rock::Int64, x::Int64) = (((x >>> rock) & 1) == 1 ? true : false)

cell_num(pomdp::RockSample, row::Int64, col::Int64) = row * pomdp.grid_size + col

make_state_index(pomdp::RockSample, cell::Int64, rock_set::Int64) =
    convert(Int64, (cell << pomdp.n_rocks) + rock_set)

POMDPs.reward(pomdp::RockSample, s::RockSampleState, a::RockSampleAction) =
    pomdp.R[s.index+1, a.index+1]

function POMDPs.transition(
                    pomdp::RockSample,
                    state::RockSampleState,
                    action::RockSampleAction)

    distribution = create_transition_distribution(pomdp)
    distribution.pomdp = pomdp
    distribution.state = state
    distribution.action = action

    return distribution
end

function POMDPs.observation(
                    pomdp::RockSample,
                    state::RockSampleState,
                    action::RockSampleAction,
                    next_state::RockSampleState)

    distribution = create_observation_distribution(pomdp)
    distribution.pomdp = pomdp
    distribution.state = next_state
    distribution.action = action
    distribution.next_state = next_state

    return distribution
end

function POMDPs.rand(
                    rng::AbstractRNG,
                    distribution::RockSampleTransitionDistribution,
                    sample=nothing)

    return RockSampleState(
        distribution.pomdp.T[distribution.state.index+1, distribution.action.index+1].index)
end

function POMDPs.rand(
                    rng::AbstractRNG,
                    distribution::RockSampleObsDistribution,
                    sample=nothing)

    # generate a new random number regardless of whether it's used below or not
    rand_num::Array{Float64} = Array{Float64}(1)
    rand!(rng, rand_num)

    if (distribution.action.index < 5)
        obs = isterminal(distribution.pomdp, distribution.next_state) ?
                    RockSampleObs(distribution.pomdp.TERMINAL_OBS) :
                    RockSampleObs(distribution.pomdp.NONE_OBS) # rs.T is an array
    else
        rock_cell = distribution.pomdp.rocks[distribution.action.index - 4] # would be [action-5] with 0-based indexing
        agent_cell = cell_of(distribution.pomdp, distribution.state)
        eff = distribution.pomdp.observation_effectiveness[agent_cell+1, rock_cell+1]

        if (rand_num[1] <= eff) == rock_status(distribution.action.index - 5, distribution.state.index)
            obs = RockSampleObs(distribution.pomdp.GOOD_OBS)
        else
            obs = RockSampleObs(distribution.pomdp.BAD_OBS)
        end
    end

    return obs
end

function POMDPs.rand(
                    rng::AbstractRNG,
                    state_space::RockSampleStateSpace,
                    sample=nothing)

    if is_linux()
        random_number = ccall((:rand_r, "libc"), Int, (Ptr{Cuint},), rng.seed) / rng.rand_max
    else #Windows, etc
        srand(seed)
        random_number = Base.rand()
    end

    return RockSampleState(floor(
           random_number*(state_space.max_index-state_space.min_index)))
end

function POMDPs.pdf(distribution::RockSampleObsDistribution,
                    obs::RockSampleObs)
  # Terminal state should match terminal obs
  if isterminal(distribution.pomdp, distribution.next_state)
      if obs.index == distribution.pomdp.TERMINAL_OBS
          return 1.
      else
          return 0.
      end
  end

  if (distribution.action.index < 5)
      if obs.index == distribution.pomdp.NONE_OBS
          return 1.
      else
          return 0.
      end
  end

  if ((obs.index != distribution.pomdp.GOOD_OBS) && (obs.index != distribution.pomdp.BAD_OBS))
    return 0.
  end

  rock::Int64       = distribution.action.index - 5
  rockCell::Int64   = distribution.pomdp.rocks[rock+1]
  agentCell::Int64  = cell_of(distribution.pomdp, distribution.next_state)

  eff::Float64 = distribution.pomdp.observation_effectiveness[agentCell+1, rockCell+1]

  rstatus::Bool = rock_status(rock, distribution.next_state.index)
  if ((obs.index == distribution.pomdp.GOOD_OBS) && (rstatus == true)) ||
     ((obs.index == distribution.pomdp.BAD_OBS) && (rstatus == false))
    return eff
  else
    return 1. - eff
  end
end

POMDPs.isterminal(pomdp::RockSample, s::RockSampleState)        =
    cell_of(pomdp, s) == pomdp.n_cells

POMDPs.isterminal_obs(pomdp::RockSample, obs::RockSampleObs)    =
    obs.index == pomdp.TERMINAL_OBS

# Which cell the agent is in
cell_of(pomdp::RockSample, s::RockSampleState) = (s.index >>> pomdp.n_rocks)

# The rock set after sampling a rock from it
sample_rock_set(rock::Int64, rock_set::Int64) = (rock_set & ~(1 << rock))

# The set of rocks in the state
rock_set_of(pomdp::RockSample, s::RockSampleState) =
    s.index & ((1 << pomdp.n_rocks)-1)

function show_state(pomdp::RockSample, s::RockSampleState)
  ac = cell_of(pomdp, s)
  for i in 0:pomdp.grid_size-1
    for j in 0:pomdp.grid_size-1
      if ac == cell_num(pomdp,i,j)
        if pomdp.rock_at_cell[ac+1] == -1 # array
          print("R ")
        elseif rock_status(pomdp.rock_at_cell[ac+1], rock_set_of(pomdp,s))
          print("G ")
        else
          print("B ")
        end # if rock_at_cell[ac] == -1
        continue
      end # if ac == cell_num(i, j)
      if pomdp.rock_at_cell[cell_num(pomdp,i,j)+1] == -1
        print(". ")
      elseif (rock_status(pomdp.rock_at_cell[cell_num(pomdp,i,j)+1], rock_set_of(pomdp,s)))
        print("1 ")
      else
        print("0 ")
      end # if rock_at_cell[cell_num(i,j)] == -1

    end # for j in 1:grid_size
    println("")
  end # i in 1:grid_size
end

function show_obs(pomdp::RockSample, obs::RockSampleObs)
    if obs.index == pomdp.NONE_OBS
        println("NONE")
    elseif obs.index == pomdp.GOOD_OBS
        println("GOOD")
    elseif obs.index == pomdp.BAD_OBS
        println("BAD")
    elseif obs.index == pomdp.TERMINAL_OBS
        println("TERMINAL")
    else
        println("UNKNOWN")
    end
end


module TestDay15

import AOC2021.Day15.part1, AOC2021.Day15.part2
using Test

test_input = """
1163751742
1381373672
2136511328
3694931569
7463417111
1319128137
1359912421
3125421639
1293138521
2311944581"""

@testset "Day 15" begin
    @testset "part 1" begin
        @test part1(test_input) == 40
    end
    @testset "part 2" begin
        @test part2(test_input) == 315
    end
end

end # module


function set_qp1test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [0.571429, 0.428572, 0.857143],
        constraint_primal = [4.0, 1.0],
        variable_dual = [0.0, 0.0, 0.0],
        constraint_dual = [0.714286, 0.857144]
    )
    return
end

function set_qp2test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [0.571429, 0.428572, 0.857143],
        constraint_primal = [4.0, 1.0],
        variable_dual = [0.0, 0.0, 0.0],
        constraint_dual = [0.714286, 0.857144]
    )
    # SOLVE 2
    LQOI.set_solution!(solver,
        variable_primal = [0.571429, 0.428572, 0.857143],
        constraint_primal = [4.0, 1.0],
        variable_dual = [-0.0, -0.0, -0.0],
        constraint_dual = [-1.42857, -1.71429]
    )
    return
end

function set_qp3test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [0.250002, 0.749998],
        constraint_primal = [1.0],
        variable_dual = [6.97141e-6, 4.06505e-8],
        constraint_dual = [2.75]
    )
    # SOLVE 2
    LQOI.set_solution!(solver,
        variable_primal = [1.0, 0.0],
        constraint_primal = [1.0],
        variable_dual = [-0.0, -1.0],
        constraint_dual = [2.0]
    )
    return
end

function set_qcp1test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [0.500003, 1.75],
        constraint_primal = [1.25, 2.25],
        quadratic_primal = [2.0],
        variable_dual = [-0.0, -0.0],
        constraint_dual = [4.48497e-8, 5.48002e-8],
        quadratic_dual = [1.0]
    )
    return
end

function set_qcp2test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [1.41421],
        constraint_primal = Float64[],
        quadratic_primal = [2.0],
        variable_dual = [-0.0],
        constraint_dual = Float64[],
        quadratic_dual = [0.353553]
    )
    return
end

function set_qcp3test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [1.41421],
        constraint_primal = Float64[],
        quadratic_primal = [2.0],
        variable_dual = [0.0],
        constraint_dual = Float64[],
        quadratic_dual = [-0.353553]
    )
    return
end

function set_socp1test_solutions!(solver)
    # SOLVE 1
    LQOI.set_solution!(solver,
        variable_primal = [0.5, 0.5, 0.707107],
        constraint_primal = [1.0],
        quadratic_primal = [-9.55708e-13],
        variable_dual = [-3.72382e-10, -3.72383e-10, 2.4911e-10],
        constraint_dual = [0.707107],
        quadratic_dual = [-1.41421]
    )
    return
end