using StatsPlots
using Query
using DataFrames
using Distributions
using LaTeXStrings
include("load_data.jl")

using Plots: PX_PER_INCH
const LINEWIDTH = 6.75*PX_PER_INCH
const COLWIDTH = 3.25*PX_PER_INCH
height(width, ratio=16/9) = width / ratio

dpl = 6.75 * 100
dph = dpl / (16/9)

plot_kwargs = (
    size=(1000, 400),
    markersize=8,
    layout=(2, 4),
    tickfontsize=8,
    guidefontsize=9,
    titlefontsize=10,
    legendfontsize=8,
    legendtitlefontsize=9,
)

####################################################################################################
# Performance: Alg, Order, Smoothing, Sigma
####################################################################################################
for probname in unique(df.probname)
    _df = df |>
        @filter((_.steprule .== :standard)) |>
        @filter((_.probname .== probname)) |>
        @filter((_.L2 .< 1e2)) |>
        DataFrame

    xlims, ylims = Dict(
        "logistic" => ((5e0, 1e6), (1e-13, 1e1),),
        "lotkavolterra" => ((1e2, 3e6), (1e-12, 1e2)),
        "fitzhughnagumo" => ((1e2, 3e6), (1e-13, 1e2)),
        "rigidbody" => ((4e1, 1.5e6), (5e-14, 1e2)),
        "brusselator" => ((4e1, 2.5e6), (1e-12, 2e2),),
    )[probname]

    legendplot = plot(
        repeat([missing], 1, 10),
        label=["EKF-IWP1" "EKF-IWP2" "EKF-IWP3" "EKF-IWP4" "EKF-IWP5" "EKS-IWP1" "EKS-IWP2" "EKS-IWP3" "EKS-IWP4" "EKS-IWP5"],
        color=[1 2 3 4 5 1 2 3 4 5],
        linestyle=hcat(repeat([:solid], 5)..., repeat([:dash], 5)...),
        marker=hcat(repeat([:o], 5)..., repeat([:star], 5)...),
        grid=false, xlims=(20,3), showaxis=false,
        legend=:left,
    );

    # plots = []
    i = 0
    for alg in ["EKF0", "EKF1"],
        sigma in [:fixedMLE, :dynamic, :fixedMLEMV,]
        i+=1

        if sigma in (:fixedMLEMV,) && alg == "EKF1"
            sigma = :dynamicMV
            alg = "EKF0"
        end

        sigmaname = Dict(
            :fixedMLE => "fixed",
            :dynamic => "TV",
            :fixedMLEMV => "fixed-MV",
            :dynamicMV => "TV-MV",
        )[sigma]

        p = _df |>
            @filter((_.sigmarule .== sigma)) |>
            @filter((_.alg .== alg)) |>
            @df StatsPlots.plot(
                :num_evals, :L2,
                group=(:q_str, :smooth_str),
                color=:q,
                linestyle=[:solid, :dash, :dot, :dashdot][:smooth.+1],
                markershape=[:o, :utriangle][:smooth.+1],
                xscale=:log10, yscale=:log10,
                xlabel = (alg == "EKF1" || sigma in (:dynamicMV,)) ? L"\#Evaluations" : " ",
                ylabel = sigma == :fixedMLE ? "L2 Error" : "",
                title="EKF$(alg[end])/EKS$(alg[end]), Γ=$sigmaname",
                legend=false,
            );
        plot!(p, xlims=xlims, ylims=ylims,)
        savefig(p, "/home/nath/PhD-Projects/ode-filter-uq-paper/figures/$(probname)_performance_$i.tex")

    end
    l = @layout [
        [a b c
         d e f] g{0.13w}
    ]
    p = plot(plots..., legendplot;
         xlims=xlims, ylims=ylims,
         plot_kwargs...,
         layout=l,
         );
    pgfsave("figures/$(probname)_performance.pdf")
    savefig("figures/$(probname)_performance.pdf")
end

####################################################################################################
# Calibration
####################################################################################################
for probname in unique(df.probname)
    _df = df |>
        @filter((_.steprule .== :standard)) |>
        @filter((_.probname .== probname)) |>
        @filter((_.L2 .< 1e2)) |>
        DataFrame

    xlims, ylims = Dict(
        "logistic" => ((1e-13, 1e1), (1e-7, 1e3),),
        "lotkavolterra" => ((1e-12, 1e2), (5e-9, 1.1e3)),
        "fitzhughnagumo" => ((1e-13, 1e2), (1e-8, 1e4)),
        "rigidbody" => ((5e-14, 1e2), (8e-7, 2e4)),
        "brusselator" => ((1e-12, 2e2), (5e-9, 1e3),),
    )[probname]

    legendplot = plot(
        repeat([missing], 1, 12),
        label=["EKF-IWP1" "EKF-IWP2" "EKF-IWP3" "EKF-IWP4" "EKF-IWP5" "EKS-IWP1" "EKS-IWP2" "EKS-IWP3" "EKS-IWP4" "EKS-IWP5" "E[χ²(2)]" "99% interval"],
        color=[1 2 3 4 5 1 2 3 4 5 :black :black],
        linestyle=hcat(repeat([:solid], 5)..., repeat([:dash], 5)..., :dash, :dot),
        marker=hcat(repeat([:o], 5)..., repeat([:star], 5)..., :none, :none),
        grid=false, xlims=(20,3), showaxis=false,
        legend=:left,
    )

    plots = []
    for alg in ["EKF0", "EKF1"],
        sigma in [:fixedMLE, :dynamic, :fixedMLEMV]

        if sigma in (:fixedMLEMV,) && alg == "EKF1"
            sigma = :dynamicMV
            alg = "EKF0"
        end

        sigmaname = Dict(
            :fixedMLE => "fixed",
            :dynamic => "TV",
            :fixedMLEMV => "fixed-MV",
            :dynamicMV => "TV-MV",
        )[sigma]

        p = plot()
        _p = 1-0.99
        dist = Chisq(2)
        m, q_low, q_high = mean(dist), quantile(dist, 1-(_p/2)), quantile(dist, _p/2)
        hline!(p, [m], color=:black, linestyle=:dash, label="")
        hline!(p, [q_low, q_high], color=:black, linestyle=:dot, label="")

        _df |>
            @filter((_.sigmarule .== sigma)) |>
            @filter((_.alg .== alg)) |>
            @df StatsPlots.plot!(p,
                # :num_evals, :Χ²,
                :L2, :Χ²,
                group=(:q_str, :smooth_str),
                color=:q,
                linestyle=[:solid, :dash, :dot, :dashdot][:smooth.+1],
                markershape=[:o, :star][:smooth.+1],
                xscale=:log10, yscale=:log10,
                xlabel = (alg == "EKF1" || sigma == :dynamicMV) ? "L2 Error" : " ",
                ylabel = sigma == :fixedMLE ? "χ²-statistic" : "",
                title="EKF$(alg[end])/EKS$(alg[end]), Γ=$sigmaname",
                legend=false,
            )
        push!(plots, p)
    end
    l = @layout [
    [a b c
     d e f] g{0.13w}
    ]
    plot(plots..., legendplot;
         xlims=xlims, ylims=ylims,
         plot_kwargs...,
         layout=l,
         )
    savefig("figures/$(probname)_calibration.pdf")
end


using PyCall;

try
  pyimport("galgebra")
catch e
  run(PyCall.python_cmd(`-m pip install -e git+https://github.com/pygae/galgebra.git#egg=galgebra --user`))
  run(PyCall.python_cmd(`-m pip install sympy==1.3 --user`))
end


using Rust
using Base.Test
using Base.Libdl

# write your own tests here
@testset "Compile code string" begin
    rustlib = compile_code("""
        extern crate libc;
        use libc::uint32_t;

        #[no_mangle]
        pub extern fn addition(a: uint32_t, b: uint32_t) -> uint32_t {
           a + b
        }
    """)

    addition = dlsym(rustlib, :addition)
    @test addition != C_NULL
    @test UInt32(UInt32(2) + UInt32(10)) === ccall(addition, UInt32, (UInt32, UInt32), 2, 10)
end

@testset "Compile Julia function" begin
    @testset "Emission" begin
        function addition(a::UInt32, b::UInt32)
            a + b
        end

        name, emission = Rust.AST.generate_rust(addition, (UInt32, UInt32))
        @test isa(emission, String)
        @test name == "addition"
    end

    @testset "Compilation" begin
        function addition(a::UInt32, b::UInt32)
            a + b
        end

        addition = compile_function(addition, (UInt32, UInt32))

        @test addition != C_NULL
        @test UInt32(UInt32(2) + UInt32(10)) === ccall(addition, UInt32, (UInt32, UInt32), 2, 10)
    end

    @testset "Every function" begin
        function everything(a::UInt32, b::UInt32)
            (b >> a * b + div(b, a) * (a % b) - a | b) + (a & b + (~b $ a))
        end

        everything_rust = compile_function(everything, (UInt32, UInt32))

        @test everything_rust != C_NULL
        @test everything(UInt32(2), UInt32(10)) === ccall(everything_rust, UInt32, (UInt32, UInt32), 2, 10)
    end
end


using GLM

tests = ["glmFit"]

println("Running tests:")

for t in tests
    tfile = joinpath("test", "$(t).jl")
    println(" * $(tfile) ...")
    include(tfile)
end


using Pkg

ssh = true
if ssh
    git = "git@github.com:SebastianAment/"
else
    git = "https://github.com/SebastianAment/"
end

add(x) = Pkg.add(Pkg.PackageSpec(url = git * x * ".git"))

add("LazyInverse.jl")
add("LazyLinearAlgebra.jl")
add("LinearAlgebraExtensions.jl")
add("KroneckerProducts.jl")
add("WoodburyIdentity.jl")
add("CovarianceFunctions.jl")
add("OptimizationAlgorithms.jl")
add("BackgroundSubtraction.jl")
println("Trying to add GaussianDistributions.jl via ssh. If this fails, please add the package locally by typing \"]add path/to/GaussianDistributions.jl\" in the REPL.")
add("GaussianDistributions.jl") # needs ssh OR add path to local package


@forwardfun bfmi
@forwardfun geweke
@forwardfun ess
@forwardfun rhat
@forwardfun mcse


#environment.jl - gets the esizesize and fsizesize variables.
################################################################################

doc"""
`Unums.environment` outputs the current environment as the appropriate Unum type
"""
function environment()
  Unum{options[:env_ESS], options[:env_FSS]}
end

doc"""
  `Unums.setenvironment` sets the current environment as the appropriate Unum type
"""
function setenvironment(ESS, FSS)
  options[:env_ESS] = ESS
  options[:env_FSS] = FSS
  Unum{ESS, FSS}
end


srand(1)
M = 10_000_000; K = 100
svec1 = rand([string(rand(Char.(32:126), rand(1:8))...) for k in 1:M÷K], M)
# memcmp("abc","def",3)
memcmp(a::String, b::String, sz) = ccall(:memcmp, Int32, (Ptr{UInt8}, Ptr{UInt8}, Int32), pointer(a), pointer(b), sz)

@time sort!(svec1, lt=(x,y)->memcmp(x,y,8)<0, alg=QuickSort); # 2.7 seconds
issorted(svec1)


struct DummySystem <: AbstractSystem; end

"""
Generic struct for a system
"""
mutable struct System <: AbstractSystem
  dim::Int64
  t::Float64
  x::Vector{Float64}
  f::Vector{Float64}
  function System(dim::Int64,t::Float64,x::Vector{Float64},f::Vector{Float64})
    new(dim,t,x,f)
  end
end

function System(model::AbstractPotential)
  dim::Int64 = getdimensionality(model)
  x::Vector{Float64} = zeros(dim)
  f::Vector{Float64} = zeros(dim)
  System(dim,0.0,x,f)
end

"""
Generic struct for a system with a thermostat
"""
mutable struct ThermostattedSystem <: AbstractThermostattedSystem
  dim::Int64
  t::Float64
  x::Vector{Float64}
  f::Vector{Float64}
  thermostat::AbstractThermostat
  function ThermostattedSystem(dim::Int64,t::Float64,x::Vector{Float64},f::Vector{Float64},thermostat::AbstractThermostat)
    new(dim,t,x,f,thermostat)
  end
end

function ThermostattedSystem(model::AbstractPotential,thermostat::AbstractThermostat)
  dim::Int64 = getdimensionality(model)
  x::Vector{Float64} = zeros(dim)
  f::Vector{Float64} = zeros(dim)
  ThermostattedSystem(dim,0.0,x,f,thermostat)
end

#"""
#Struct for an overdamped system with a langevin thermostat
#"""
#mutable struct OverdampedLangevinSystem <: AbstractThermostattedSystem
#  dim::Int64
#  t::Float64
#  x::Vector{Float64}
#  f::Vector{Float64}
#  thermostat::AbstractLangevin
#  function OverdampedLangevinSystem(dim::Int64,t::Float64,x::Vector{Float64},f::Vector{Float64},thermostat::AbstractLangevin)
#    new(dim,t,x,f,thermostat)
#  end
#end
#
#function OverdampedLangevinSystem(model::AbstractPotential,thermostat::AbstractLangevin)
#  dim::Int64 = getdimensionality(model)
#  x::Vector{Float64} = zeros(dim)
#  f::Vector{Float64} = zeros(dim)
#  ThermostattedSystem(dim,0.0,x,f,thermostat)
#end

"""
Functions for interacting with the force
"""
function zeroforce!(system::AbstractSystem)
  system.f .= 0.0
end

function addforce!(dt::Float64,system::AbstractSystem,model::AbstractSinglePotential)
  system.x .+= dt .* model.f
end

function addforce!(dt::Float64,system::AbstractSystem,model::AbstractTrainablePotential)
  if model.condition(system)
    system.x .+= dt .* model.f
  end
end

# TODO maybe take this function out since it's a repeat of the one above
function addforce!(dt::Float64,system::ThermostattedSystem,model::AbstractSinglePotential)
  system.x .+= dt .* model.f ./ system.thermostat.gamma
#  system.x .+= dt .* model.f
end

function addforce!(dt::Float64,system::ThermostattedSystem,model::AbstractTrainablePotential)
#  println(system.t)
  if model.condition(system)
#    println("hey")
    system.x .+= dt .* model.f ./ system.thermostat.gamma
  end
end

# TODO do unit tests on Overdamped system before allowing this
#function addforce!(dt::Float64,system::OverdampedLangevinSystem,model::AbstractPotential)
#  system.x .+= dt .* model.f ./ system.thermostat.gamma
#end


module SMLMData

using Base
using StatsBase
using Distributions
using MAT
using DataFrames
using Images

include("structdefinitions.jl")
include("isolatesmld.jl")
include("catsmld.jl")
include("transformsmld.jl")
include("images.jl")
include("smite.jl")

end


# Represents a static order of an ITensor
@eval struct Order{N}
  (OrderT::Type{<:Order})() = $(Expr(:new, :OrderT))
end

@doc """
   Order{N}
A value type representing the order of an ITensor.
""" Order

"""
   Order(N) = Order{N}()
Create an instance of the value type Order representing
the order of an ITensor.
"""
Order(N) = Order{N}()

# Helpful if we want code to work generically
# for other Index-like types (such as IndexRange)
const IndexSet{IndexT<:Index} = Vector{IndexT}
const IndexTuple{IndexT<:Index} = Tuple{Vararg{IndexT}}

# Definition to help with generic code
const Indices{IndexT<:Index} = Union{Vector{IndexT},Tuple{Vararg{IndexT}}}

# Flatten combinations of tuples and vectors into a single collection
# of indices
tuple_vcat(t::Tuple) = t
tuple_vcat() = ()
tuple_vcat(t) = (t,)
tuple_vcat(a, args...) = (tuple_vcat(a)..., tuple_vcat(args...)...)

tuple_to_vector(t::Tuple) = collect(t)
tuple_to_vector(t) = t

function _narrow_eltype(v::Vector{T}) where {T}
  return convert(Vector{mapreduce(typeof, promote_type, v)}, v)
end
narrow_eltype(v::Vector{T}) where {T} = isconcretetype(T) ? v : _narrow_eltype(v)

push_or_append!(v, x::Union{Vector,Tuple}) = append!(v, x)
push_or_append!(v, x) = push!(v, x)

function _indices(is::Vector)
  isempty(is) && return is
  is_flat = Index[]
  for i in is
    push_or_append!(is_flat, i)
  end
  return narrow_eltype(is_flat)
end
indices(is::Vector{<:Index}) = narrow_eltype(is)

_indices(is::Tuple{Vararg{<:Index}}) = is
_indices(is::Tuple{Vararg{Union{<:Vector,<:Index}}}) = vcat(is...)
_indices(is::Tuple{Vararg{Union{<:Tuple,<:Index}}}) = tuple_vcat(is...)
_indices(is::Tuple{Vararg{Union{<:Tuple,<:Vector,<:Index}}}) = indices(tuple_to_vector.(is))
indices(is::Tuple{Vararg{<:Index}}) = is
function indices(is::Tuple)
  inds = _indices(is)
  if isempty(inds)
    # Otherwise it outputs `Any[]`, which breaks
    # some generic code like `dim`.
    return Index[]
  end
  return inds
end
indices(is::Union{<:Tuple,<:Vector,<:Index}...) = indices(is)
function indices(is::Vector)
  # This narrows the type. Also handles the empty case.
  all(i -> i isa Index, is) && return narrow_eltype(is)
  return _indices(is)
end

# To help with backwards compatibility
IndexSet(inds::IndexSet) = inds
IndexSet(inds::Indices) = collect(inds)
IndexSet(inds::Index...) = collect(inds)
IndexSet(f::Function, N::Int) = map(f, 1:N)
IndexSet(f::Function, ::Order{N}) where {N} = IndexSet(f, N)

Tuple(is::IndexSet) = _Tuple(is)
NTuple{N}(is::IndexSet) where {N} = _NTuple(Val(N), is)

"""
    not(inds::Union{IndexSet, Tuple{Vararg{<:Index}}})
    not(inds::Index...)
    !(inds::Union{IndexSet, Tuple{Vararg{<:Index}}})
    !(inds::Index...)

Represents the set of indices not in the specified
IndexSet, for use in pattern matching (i.e. when
searching for an index or priming/tagging specified
indices).
"""
not(is::Indices) = Not(is)
not(inds::Index...) = not(inds)
!(is::Indices) = not(is)
!(inds::Index...) = not(inds...)

# Convert to an Index if there is only one
# TODO: also define the function `only`
Index(is::Indices) = is[]

NDTensors.dims(is::IndexSet) = dim.(is)

"""
    dim(is::Indices)

Get the product of the dimensions of the indices
of the Indices (the total dimension of the space).
"""
NDTensors.dim(is::IndexSet) = Compat.mapreduce(dim, *, is; init=1)

"""
    dim(is::IndexSet, n::Int)

Get the dimension of the Index n of the Indices.
"""
NDTensors.dim(is::IndexSet, pos::Int) = dim(is[pos])

"""
    dag(is::Indices)

Return a new Indices with the indices daggered (flip
all of the arrow directions).
"""
dag(is::Indices) = map(i -> dag(i), is)

# TODO: move to NDTensors
NDTensors.dim(is::Tuple, pos::Integer) = dim(is[pos])

# TODO: this fixes an ambiguity error with base, move to NDTensors
NDTensors.similar(T::NDTensors.DenseTensor, inds::Tuple) = NDTensors._similar(T, inds)

# TODO: this is a weird definition, fix it
function NDTensors.similartype(
  ::Type{<:Tuple{Vararg{IndexT}}}, ::Type{Val{N}}
) where {IndexT,N}
  return NTuple{N,IndexT}
end

## # This is to help with some generic programming in the Tensor
## # code (it helps to construct an IndexSet(::NTuple{N,Index}) where the
## # only known thing for dispatch is a concrete type such
## # as IndexSet{4})
##
## #NDTensors.similartype(::Type{<:IndexSet},
## #                      ::Val{N}) where {N} = IndexSet
##
## #NDTensors.similartype(::Type{<:IndexSet},
## #                      ::Type{Val{N}}) where {N} = IndexSet

"""
    sim(is::Indices)

Make a new Indices with similar indices.

You can also use the broadcast version `sim.(is)`.
"""
sim(is::Indices) = map(i -> sim(i), is)

"""
    mindim(is::Indices)

Get the minimum dimension of the indices in the index set.

Returns 1 if the Indices is empty.
"""
function mindim(is::Indices)
  length(is) == 0 && (return 1)
  md = dim(is[1])
  for n in 2:length(is)
    md = min(md, dim(is[n]))
  end
  return md
end

"""
    maxdim(is::Indices)

Get the maximum dimension of the indices in the index set.

Returns 1 if the Indices is empty.
"""
function maxdim(is::Indices)
  length(is) == 0 && (return 1)
  md = dim(is[1])
  for n in 2:length(is)
    md = max(md, dim(is[n]))
  end
  return md
end

"""
    commontags(::Indices)

Return a TagSet of the tags that are common to all of the indices.
"""
commontags(is::Indices) = commontags(is...)

#
# Set operations
#

"""
    ==(is1::Indices, is2::Indices)

Indices equality (order dependent). For order
independent equality use `issetequal` or
`hassameinds`.
"""
function ==(A::Indices, B::Indices)
  length(A) ≠ length(B) && return false
  for (a, b) in zip(A, B)
    a ≠ b && return false
  end
  return true
end

"""
    ITensors.fmatch(pattern) -> ::Function

fmatch is an internal function that
creates a function that accepts an Index.
The function returns true if the Index matches
the provided pattern, and false otherwise.

For example:
```
i = Index(2, "s")
fmatch("s")(i) == true
```
"""
fmatch(is::Indices) = in(is)
fmatch(is::Index...) = fmatch(is)
fmatch(i::Index) = fmatch((i,))

fmatch(pl::Int) = hasplev(pl)

fmatch(tags::TagSet) = hastags(tags)
fmatch(tags::AbstractString) = fmatch(TagSet(tags))

fmatch(id::IDType) = hasid(id)

fmatch(n::Not) = !fmatch(parent(n))

# Function that always returns true
ftrue(::Any) = true

fmatch(::Nothing) = ftrue

"""
    ITensors.fmatch(; inds = nothing,
                      tags = nothing,
                      plev = nothing,
                      id = nothing) -> Function

An internal function that returns a function
that accepts an Index that checks if the
Index matches the provided conditions.
"""
function fmatch(; inds=nothing, tags=nothing, plev=nothing, id=nothing)
  return i -> fmatch(inds)(i) && fmatch(plev)(i) && fmatch(id)(i) && fmatch(tags)(i)
end

"""
    indmatch

Checks if the Index matches the provided conditions.
"""
indmatch(i::Index; kwargs...) = fmatch(; kwargs...)(i)

"""
    getfirst(is::Indices)

Return the first Index in the Indices. If the Indices
is empty, return `nothing`.
"""
function getfirst(is::Indices)
  length(is) == 0 && return nothing
  return first(is)
end

"""
    getfirst(f::Function, is::Indices)

Get the first Index matching the pattern function,
return `nothing` if not found.
"""
function getfirst(f::Function, is::Indices)
  for i in is
    f(i) && return i
  end
  return nothing
end

getfirst(is::Indices, args...; kwargs...) = getfirst(fmatch(args...; kwargs...), is)

Base.findall(is::Indices, args...; kwargs...) = findall(fmatch(args...; kwargs...), is)

# In general this isn't defined for Tuple but is
# defined for Vector
"""
    indexin(ais::Indices, bis::Indices)

For collections of Indices, returns the first location in
`bis` for each value in `ais`.
"""
function Base.indexin(ais::Indices, bis::Indices)
  return [findfirst(bis, ais[i]) for i in 1:length(ais)]
end

#function Base.indexin(a::Index, bis::Indices)
#  return [findfirst(bis, a)]
#end

findfirst(is::Indices, args...; kwargs...) = findfirst(fmatch(args...; kwargs...), is)

#
# Tagging functions
#

function prime(f::Function, is::Indices, args...)
  return map(i -> f(i) ? prime(i, args...) : i, is)
end

"""
    prime(A::Indices, plinc, ...)

Increase the prime level of the indices by the specified amount.
Filter which indices are primed using keyword arguments
tags, plev and id.
"""
function prime(is::Indices, plinc::Integer, args...; kwargs...)
  return prime(fmatch(args...; kwargs...), is, plinc)
end

prime(f::Function, is::Indices) = prime(f, is, 1)

prime(is::Indices, args...; kwargs...) = prime(is, 1, args...; kwargs...)

"""
    adjoint(is::Indices)

For is' notation.
"""
adjoint(is::Indices) = prime(is)

function setprime(f::Function, is::Indices, args...)
  return map(i -> f(i) ? setprime(i, args...) : i, is)
end

function setprime(is::Indices, plev::Integer, args...; kwargs...)
  return setprime(fmatch(args...; kwargs...), is, plev)
end

noprime(f::Function, is::Indices, args...) = setprime(is, 0, args...; kwargs...)

noprime(is::Indices, args...; kwargs...) = setprime(is, 0, args...; kwargs...)

function _swapprime(f::Function, i::Index, pl1pl2::Pair{Int,Int})
  pl1, pl2 = pl1pl2
  if f(i)
    if hasplev(i, pl1)
      return setprime(i, pl2)
    elseif hasplev(i, pl2)
      return setprime(i, pl1)
    end
    return i
  end
  return i
end

function swapprime(f::Function, is::Indices, pl1pl2::Pair{Int,Int})
  return map(i -> _swapprime(f, i, pl1pl2), is)
end

function swapprime(f::Function, is::Indices, pl1::Int, pl2::Int)
  return swapprime(f, is::Indices, pl1 => pl2)
end

function swapprime(is::Indices, pl1pl2::Pair{Int,Int}, args...; kwargs...)
  return swapprime(fmatch(args...; kwargs...), is, pl1pl2)
end

function swapprime(is::Indices, pl1::Int, pl2::Int, args...; kwargs...)
  return swapprime(fmatch(args...; kwargs...), is, pl1 => pl2)
end

replaceprime(f::Function, is::Indices, pl1::Int, pl2::Int) = replaceprime(f, is, pl1 => pl2)

function replaceprime(is::Indices, pl1::Int, pl2::Int, args...; kwargs...)
  return replaceprime(fmatch(args...; kwargs...), is, pl1 => pl2)
end

const mapprime = replaceprime

function _replaceprime(i::Index, rep_pls::Pair{Int,Int}...)
  for (pl1, pl2) in rep_pls
    hasplev(i, pl1) && return setprime(i, pl2)
  end
  return i
end

function replaceprime(f::Function, is::Indices, rep_pls::Pair{Int,Int}...)
  return map(i -> f(i) ? _replaceprime(i, rep_pls...) : i, is)
end

function replaceprime(is::Indices, rep_pls::Pair{Int,Int}...; kwargs...)
  return replaceprime(fmatch(; kwargs...), is, rep_pls...)
end

addtags(f::Function, is::Indices, args...) = map(i -> f(i) ? addtags(i, args...) : i, is)

function addtags(is::Indices, tags, args...; kwargs...)
  return addtags(fmatch(args...; kwargs...), is, tags)
end

settags(f::Function, is::Indices, args...) = map(i -> f(i) ? settags(i, args...) : i, is)

function settags(is::Indices, tags, args...; kwargs...)
  return settags(fmatch(args...; kwargs...), is, tags)
end

"""
    CartesianIndices(is::Indices)

Create a CartesianIndices iterator for an Indices.
"""
CartesianIndices(is::Indices) = CartesianIndices(_Tuple(dims(is)))

"""
    eachval(is::Index...)
    eachval(is::Tuple{Vararg{Index}})

Create an iterator whose values correspond to a
Cartesian indexing over the dimensions
of the provided `Index` objects.
"""
eachval(is::Index...) = eachval(is)
eachval(is::Tuple{Vararg{Index}}) = CartesianIndices(dims(is))

"""
    eachindval(is::Index...)
    eachindval(is::Tuple{Vararg{Index}})

Create an iterator whose values are Index=>value pairs
corresponding to a Cartesian indexing over the dimensions
of the provided `Index` objects.
# Example
```julia
i = Index(3; tags = "i")
j = Index(2; tags = "j")
T = randomITensor(j,i)
for iv in eachindval(i,j)
  @show T[iv...]
end
```
"""
eachindval(is::Index...) = eachindval(is)
eachindval(is::Tuple{Vararg{Index}}) = (is .=> Tuple(ns) for ns in eachval(is))

function removetags(f::Function, is::Indices, args...)
  return map(i -> f(i) ? removetags(i, args...) : i, is)
end

function removetags(is::Indices, tags, args...; kwargs...)
  return removetags(fmatch(args...; kwargs...), is, tags)
end

function _replacetags(i::Index, rep_ts::Pair...)
  for (tags1, tags2) in rep_ts
    hastags(i, tags1) && return replacetags(i, tags1, tags2)
  end
  return i
end

# XXX new syntax
# hastags(any, is, ts)
"""
    anyhastags(is::Indices, ts::Union{String, TagSet})
    hastags(is::Indices, ts::Union{String, TagSet})

Check if any of the indices in the Indices have the specified tags.
"""
anyhastags(is::Indices, ts) = any(i -> hastags(i, ts), is)

hastags(is::Indices, ts) = anyhastags(is, ts)

# XXX new syntax
# hastags(all, is, ts)
"""
    allhastags(is::Indices, ts::Union{String, TagSet})

Check if all of the indices in the Indices have the specified tags.
"""
allhastags(is::Indices, ts::String) = all(i -> hastags(i, ts), is)

# Version taking a list of Pairs
function replacetags(f::Function, is::Indices, rep_ts::Pair...)
  return map(i -> f(i) ? _replacetags(i, rep_ts...) : i, is)
end

function replacetags(is::Indices, rep_ts::Pair...; kwargs...)
  return replacetags(fmatch(; kwargs...), is, rep_ts...)
end

# Version taking two input TagSets/Strings
replacetags(f::Function, is::Indices, tags1, tags2) = replacetags(f, is, tags1 => tags2)

function replacetags(is::Indices, tags1, tags2, args...; kwargs...)
  return replacetags(fmatch(args...; kwargs...), is, tags1 => tags2)
end

function _swaptags(f::Function, i::Index, tags1, tags2)
  if f(i)
    if hastags(i, tags1)
      return replacetags(i, tags1, tags2)
    elseif hastags(i, tags2)
      return replacetags(i, tags2, tags1)
    end
    return i
  end
  return i
end

function swaptags(f::Function, is::Indices, tags1, tags2)
  return map(i -> _swaptags(f, i, tags1, tags2), is)
end

function swaptags(is::Indices, tags1, tags2, args...; kwargs...)
  return swaptags(fmatch(args...; kwargs...), is, tags1, tags2)
end

function swaptags(is::Indices, tags12::Pair, args...; kwargs...)
  return swaptags(is, first(tags12), last(tags12), args...; kwargs...)
end

function replaceinds(is::Indices, rep_inds::Pair{<:Index,<:Index}...)
  return replaceinds(is, zip(rep_inds...)...)
end

function replaceinds(is::Indices, rep_inds::Vector{<:Pair{<:Index,<:Index}})
  return replaceinds(is, rep_inds...)
end

function replaceinds(is::Indices, rep_inds::Tuple{Vararg{Pair{<:Index,<:Index}}})
  return replaceinds(is, rep_inds...)
end

function replaceinds(is::Indices, rep_inds::Pair)
  return replaceinds(is, Tuple(first(rep_inds)) .=> Tuple(last(rep_inds)))
end

# Check that the QNs are all the same
hassameflux(i1::Index, i2::Index) = (dim(i1) == dim(i2))

function replaceinds(is::Indices, inds1, inds2)
  is1 = inds1
  poss = indexin(is1, is)
  is_tuple = Tuple(is)
  for (j, pos) in enumerate(poss)
    isnothing(pos) && continue
    i1 = is_tuple[pos]
    i2 = inds2[j]
    i2 = setdir(i2, dir(i1))
    space(i1) ≠ space(i2) && error("Indices must have the same spaces to be replaced")
    is_tuple = setindex(is_tuple, i2, pos)
  end
  return (is_tuple)
end

replaceind(is::Indices, i1::Index, i2::Index) = replaceinds(is, (i1,), (i2,))

function replaceind(is::Indices, i1::Index, i2::Indices)
  length(i2) != 1 &&
    throw(ArgumentError("cannot use replaceind with an Indices of length $(length(i2))"))
  return replaceinds(is, (i1,), i2)
end

replaceind(is::Indices, rep_i::Pair{<:Index,<:Index}) = replaceinds(is, rep_i)

function swapinds(is::Indices, inds1, inds2)
  return replaceinds(is, (inds1..., inds2...), (inds2..., inds1...))
end

function swapinds(is::Indices, inds1::Index, inds2::Index)
  return swapinds(is, (inds1,), (inds2,))
end

function swapinds(is::Indices, inds12::Pair)
  return swapinds(is, first(inds12), last(inds12))
end

swapind(is::Indices, i1::Index, i2::Index) = swapinds(is, (i1,), (i2,))

removeqns(is::Indices) = is

# Permute is1 to be in the order of is2
# This is helpful when is1 and is2 have different directions, and
# you want is1 to have the same directions as is2
# TODO: replace this functionality with
#
# setdirs(is1::Indices, is2::Indices)
#
function permute(is1::Indices, is2::Indices)
  length(is1) != length(is2) && throw(
    ArgumentError(
      "length of first index set, $(length(is1)) does not match length of second index set, $(length(is2))",
    ),
  )
  perm = getperm(is1, is2)
  return is1[invperm(perm)]
end

#
# Helper functions for contracting ITensors
#

function compute_contraction_labels(Ais::Tuple, Bis::Tuple)
  have_qns = hasqns(Ais) && hasqns(Bis)
  NA = length(Ais)
  NB = length(Bis)
  Alabels = MVector{NA,Int}(ntuple(_ -> 0, Val(NA)))
  Blabels = MVector{NB,Int}(ntuple(_ -> 0, Val(NB)))

  ncont = 0
  for i in 1:NA, j in 1:NB
    Ais_i = @inbounds Ais[i]
    Bis_j = @inbounds Bis[j]
    if Ais_i == Bis_j
      if have_qns && (dir(Ais_i) ≠ -dir(Bis_j))
        error(
          "Attempting to contract IndexSet:\n\n$(Ais)\n\nwith IndexSet:\n\n$(Bis)\n\nQN indices must have opposite direction to contract, but indices:\n\n$(Ais_i)\n\nand:\n\n$(Bis_j)\n\ndo not have opposite directions.",
        )
      end
      Alabels[i] = Blabels[j] = -(1 + ncont)
      ncont += 1
    end
  end

  u = ncont
  for i in 1:NA
    if (Alabels[i] == 0)
      Alabels[i] = (u += 1)
    end
  end
  for j in 1:NB
    if (Blabels[j] == 0)
      Blabels[j] = (u += 1)
    end
  end

  return (Tuple(Alabels), Tuple(Blabels))
end

function compute_contraction_labels(Cis::Tuple, Ais::Tuple, Bis::Tuple)
  NA = length(Ais)
  NB = length(Bis)
  NC = length(Cis)
  Alabels, Blabels = compute_contraction_labels(Ais, Bis)
  Clabels = MVector{NC,Int}(ntuple(_ -> 0, Val(NC)))
  for i in 1:NC
    locA = findfirst(==(Cis[i]), Ais)
    if !isnothing(locA)
      if Alabels[locA] < 0
        error(
          "The noncommon indices of $Ais and $Bis must be the same as the indices $Cis."
        )
      end
      Clabels[i] = Alabels[locA]
    else
      locB = findfirst(==(Cis[i]), Bis)
      if isnothing(locB) || Blabels[locB] < 0
        error(
          "The noncommon indices of $Ais and $Bis must be the same as the indices $Cis."
        )
      end
      Clabels[i] = Blabels[locB]
    end
  end
  return (Tuple(Clabels), Alabels, Blabels)
end

#
# TupleTools
#

"""
    pop(is::Indices)

Return a new Indices with the last Index removed.
"""
pop(is::Indices) = (NDTensors.pop(Tuple(is)))

# Overload the unexported NDTensors version
NDTensors.pop(is::Indices) = pop(is)

# TODO: don't convert to Tuple
"""
    popfirst(is::Indices)

Return a new Indices with the first Index removed.
"""
popfirst(is::IndexSet) = (NDTensors.popfirst(Tuple(is)))

# Overload the unexported NDTensors version
NDTensors.popfirst(is::IndexSet) = popfirst(is)

"""
    push(is::Indices, i::Index)

Make a new Indices with the Index i inserted
at the end.
"""
push(is::IndexSet, i::Index) = NDTensors.push(is, i)

# Overload the unexported NDTensors version
NDTensors.push(is::IndexSet, i::Index) = push(is, i)

# TODO: deprecate in favor of `filterinds` (abuse of Base notation)
filter(is::Indices, args...; kwargs...) = filter(fmatch(args...; kwargs...), is)

# For ambiguity with Base.filter
filter(is::Indices, args::String; kwargs...) = filter(fmatch(args; kwargs...), is)

#
# QN functions
#

"""
    setdirs(is::Indices, dirs::Arrow...)

Return a new Indices with indices `setdir(is[i], dirs[i])`.
"""
function setdirs(is::Indices, dirs)
  return map(i -> setdir(is[i], dirs[i]), 1:length(is))
end

"""
    dir(is::Indices, i::Index)

Return the direction of the Index `i` in the Indices `is`.
"""
function dir(is1::Indices, i::Index)
  return dir(getfirst(is1, i))
end

"""
    dirs(is::Indices, inds)

Return a tuple of the directions of the indices `inds` in
the Indices `is`, in the order they are found in `inds`.
"""
function dirs(is1::Indices, inds)
  return map(i -> dir(is1, inds[i]), 1:length(inds))
end

"""
    dirs(is::Indices)

Return a tuple of the directions of the indices `is`.
"""
dirs(is::Indices) = dir.(is)

hasqns(is::Indices) = any(hasqns, is)

"""
    getperm(col1, col2)

Get the permutation that takes collection 2 to collection 1,
such that `col2[p] .== col1`.
"""
function getperm(s1, s2)
  N = length(s1)
  r = Vector{Int}(undef, N)
  return map!(i -> findfirst(==(s1[i]), s2), r, 1:length(s1))
end

# TODO: define directly for Vector
"""
    nblocks(::Indices, i::Int)

The number of blocks in the specified dimension.
"""
function NDTensors.nblocks(inds::IndexSet, i::Int)
  return nblocks(Tuple(inds), i)
end

# TODO: don't convert to Tuple
function NDTensors.nblocks(inds::IndexSet, is)
  return nblocks(Tuple(inds), is)
end

"""
    nblocks(::Indices)

A tuple of the number of blocks in each
dimension.
"""
NDTensors.nblocks(inds::Indices) = nblocks.(inds)

# TODO: is this needed?
function NDTensors.nblocks(inds::NTuple{N,<:Index}) where {N}
  return ntuple(i -> nblocks(inds, i), Val(N))
end

ndiagblocks(inds) = minimum(nblocks(inds))

"""
    flux(inds::Indices, block::Tuple{Vararg{Int}})

Get the flux of the specified block, for example:
```
i = Index(QN(0)=>2, QN(1)=>2)
is = (i, dag(i'))
flux(is, Block(1, 1)) == QN(0)
flux(is, Block(2, 1)) == QN(1)
flux(is, Block(1, 2)) == QN(-1)
flux(is, Block(2, 2)) == QN(0)
```
"""
function flux(inds::Indices, block::Block)
  qntot = QN()
  for n in 1:length(inds)
    ind = inds[n]
    qntot += flux(ind, Block(block[n]))
  end
  return qntot
end

"""
    flux(inds::Indices, I::Integer...)

Get the flux of the block that the specified
index falls in.
```
i = Index(QN(0)=>2, QN(1)=>2)
is = (i, dag(i'))
flux(is, 3, 1) == QN(1)
flux(is, 1, 2) == QN(0)
```
"""
flux(inds::Indices, vals::Integer...) = flux(inds, block(inds, vals...))

"""
    ITensors.block(inds::Indices, I::Integer...)

Get the block that the specified index falls in.

This is mostly an internal function, and the interface
is subject to change.

# Examples

```julia
i = Index(QN(0)=>2, QN(1)=>2)
is = (i, dag(i'))
ITensors.block(is, 3, 1) == (2,1)
ITensors.block(is, 1, 2) == (1,1)
```
"""
block(inds::Indices, vals::Integer...) = blockindex(inds, vals...)[2]

#show(io::IO, is::IndexSet) = show(io, MIME"text/plain"(), is)

#
# Read and write
#

function readcpp(io::IO, ::Type{<:Indices}; kwargs...)
  format = get(kwargs, :format, "v3")
  is = IndexSet()
  if format == "v3"
    size = read(io, Int)
    function readind(io, n)
      i = readcpp(io, Index; kwargs...)
      stride = read(io, UInt64)
      return i
    end
    is = IndexSet(n -> readind(io, n), size)
  else
    throw(ArgumentError("read IndexSet: format=$format not supported"))
  end
  return is
end

function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, is::Indices)
  g = create_group(parent, name)
  attributes(g)["type"] = "IndexSet"
  attributes(g)["version"] = 1
  N = length(is)
  write(g, "length", N)
  for n in 1:N
    write(g, "index_$n", is[n])
  end
end

function HDF5.read(
  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, T::Type{<:Indices}
)
  g = open_group(parent, name)
  if read(attributes(g)["type"]) != "IndexSet"
    error("HDF5 group or file does not contain IndexSet data")
  end
  N = read(g, "length")
  return T(n -> read(g, "index_$n", Index), N)
end


function build_robot!(vis::Visualizer, model::PointFootQuadruped;
	body_height=0.025,
	body_length=0.17,
	body_width=0.15,
	foot_radius=0.025,
	color_opacity=1.0,
	name::Symbol=:robot)

	body_mat = MeshPhongMaterial(color=RGBA(0.0, 0.0, 0.0, color_opacity))
	foot_mat = MeshPhongMaterial(color=RGBA(1.0, 165.0 / 255.0, 0.0, color_opacity))

	default_background!(vis)

	setobject!(vis[name][:body],
    	GeometryBasics.Rect(Vec(-body_length, -body_width, -body_height),Vec(2.0 * body_length, 2.0 * body_width, body_height)),
    	body_mat)

	foot1 = setobject!(vis[name][:foot1], GeometryBasics.Sphere(Point3f0(0),
        convert(Float32, foot_radius)),
        foot_mat)

	foot2 = setobject!(vis[name][:foot2], GeometryBasics.Sphere(Point3f0(0),
		convert(Float32, foot_radius)),
		foot_mat)

	foot3 = setobject!(vis[name][:foot3], GeometryBasics.Sphere(Point3f0(0),
		convert(Float32, foot_radius)),
		foot_mat)

	foot4 = setobject!(vis[name][:foot4], GeometryBasics.Sphere(Point3f0(0),
		convert(Float32, foot_radius)),
		foot_mat)
end

function set_robot!(vis::Visualizer, model::PointFootQuadruped, q::AbstractVector;
	p_shift=[0.0; 0.0; 0.025], name::Symbol=:robot)

	R = mrp_rotation_matrix(q[4:6])

	p_body  = q[0  .+ (1:3)] + p_shift
	p_foot1 = q[6  .+ (1:3)] + p_shift
	p_foot2 = q[9  .+ (1:3)] + p_shift
	p_foot3 = q[12 .+ (1:3)] + p_shift
	p_foot4 = q[15 .+ (1:3)] + p_shift

	settransform!(vis[name][:body], compose(Translation(p_body), LinearMap(mrp_rotation_matrix(q[4:6]))))
	settransform!(vis[name][:foot1], Translation(p_foot1))
	settransform!(vis[name][:foot2], Translation(p_foot2))
	settransform!(vis[name][:foot3], Translation(p_foot3))
	settransform!(vis[name][:foot4], Translation(p_foot4))

end

function visualize!(vis, model::PointFootQuadruped, q;
	Δt=0.1,
	body_height=0.025,
	body_length=0.17,
	body_width=0.15,
	foot_radius=0.025,
	color_opacity=1.0,
	fixed_camera=false,
	anim=MeshCat.Animation(convert(Int, floor(1.0 / Δt))),
	name::Symbol=:robot)

	build_robot!(vis, model,
		body_height=body_height,
		body_length=body_length,
		body_width=body_width,
		foot_radius=foot_radius,
		color_opacity=color_opacity,
		name=name)

	for (t, qt) in enumerate(q)
		MeshCat.atframe(anim, t) do
			set_robot!(vis, model, qt, name=name)
		end
	end

	MeshCat.setanimation!(vis, anim)

	return anim
end


# __BEGIN_LICENSE__
#
# WaveOptics.jl
#
# Copyright (c) 2015, Stanford University
#
# All rights reserved.
#
# Redistribution and use in source and binary forms for academic and other
# non-commercial purposes with or without modification, are permitted provided
# that the following conditions are met:
#
# * Redistributions of source code, including modified source code, must retain
#   the above copyright notice, this list of conditions and the following
#   disclaimer.
#
# * Redistributions in binary form or a modified form of the source code must
#   reproduce the above copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other materials provided with
#   the distribution.
#
# * Neither the name of The Leland Stanford Junior University, any of its
#   trademarks, the names of its employees, nor contributors to the source code
#   may be used to endorse or promote products derived from this software without
#   specific prior written permission.
#
# * Where a modified version of the source code is redistributed publicly in
#   source or binary forms, the modified source code must be published in a freely
#   accessible manner, or otherwise redistributed at no charge to anyone
#   requesting a copy of the modified source code, subject to the same terms as
#   this agreement.
#
# THIS SOFTWARE IS PROVIDED BY THE TRUSTEES OF THE LELAND STANFORD JUNIOR
# UNIVERSITY "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE LELAND STANFORD JUNIOR
# UNIVERSITY OR ITS TRUSTEES BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
# OF SUCH DAMAGE.
#
# __END_LICENSE__

"""
In order to avoid aliasing artifacts, the wave optics simulation must be
carried out at a resolution close to the diffraction limit. However, image
sensors sometimes undersample as compared to this limit, thus we may need to
downsample here to reach the resolution of the image sensor.

This utility function performs a downsampling by averaging together oversampling
x oversampling square blocks of pixels. This actually produces some aliasing,
but this aliasing should be present in the real optical system as well.

TODO: someday we could add pixel fill factor as a parameter to this function.
"""
function downsample_to_sensor_resolution(
    psf_image::AbstractArray{T,2},
    oversampling::Int64,
) where {T<:Real}
    if (oversampling == 1)
        return psf_image
    end

    # We must ensure that the image has an even number of pixels, so we pad
    # (on the bottom and right) with a row of black pixels if necessary to
    # ensure the image has an even number of pixels in each linear dimension.
    trim_r = size(psf_image, 1) % oversampling
    trim_c = size(psf_image, 2) % oversampling
    trimmed_image =
        view(psf_image, 1:size(psf_image, 1)-trim_r, 1:size(psf_image, 2)-trim_c)
    height, width = size(trimmed_image)

    result = zeros(T, (div(height, oversampling), div(width, oversampling)))
    for r = 1:oversampling
        for c = 1:oversampling
            result .= result .+ trimmed_image[r:oversampling:end, c:oversampling:end]
        end
    end

    return result
end

function downsample_to_sensor_resolution(
    psf_image::AbstractArray{T,3},
    oversampling::Int64,
) where {T<:Real}
    if (oversampling == 1)
        return psf_image
    end

    # We must ensure that the image has an even number of pixels, so we pad
    # (on the bottom and right) with a row of black pixels if necessary to
    # ensure the image has an even number of pixels in each linear dimension.
    I, J, K = size(psf_image)
    trim_r = I % oversampling
    trim_c = J % oversampling
    trimmed_image = view(psf_image, 1:I-trim_r, 1:J-trim_c, 1:K)
    height, width = size(trimmed_image)

    result = zeros(T, (div(height, oversampling), div(width, oversampling), K))
    @inbounds for r = 1:oversampling
        for c = 1:oversampling
            result .+= trimmed_image[r:oversampling:end, c:oversampling:end, :]
        end
    end

    return result
end


import Adapt
import .CUDA
using LinearAlgebra, NNlib

"""
    OneHotArray{T,L,N,M,I} <: AbstractArray{Bool,M}

These are constructed by [`onehot`](@ref) and [`onehotbatch`](@ref).
Parameter `I` is the type of the underlying storage, and `T` its eltype.
"""
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}

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

# 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]...]

function Base.showarg(io::IO, x::OneHotArray, toplevel)
    print(io, ndims(x) == 1 ? "OneHotVector(" : ndims(x) == 2 ? "OneHotMatrix(" : "OneHotArray(")
    Base.showarg(io, x.indices, false)
    print(io, ')')
    toplevel && print(io, " with eltype Bool")
    return nothing
end

# this is from /LinearAlgebra/src/diagonal.jl, official way to print the dots:
function Base.replace_in_print_matrix(x::OneHotLike, i::Integer, j::Integer, s::AbstractString)
    x[i,j] ? s : _isonehot(x) ? Base.replace_with_centered_mark(s) : s
end

# copy CuArray versions back before trying to print them:
Base.print_array(io::IO, X::OneHotLike{T, L, N, var"N+1", <:CuArray}) where {T, L, N, var"N+1"} =
  Base.print_array(io, cpu(X))
Base.print_array(io::IO, X::LinearAlgebra.AdjOrTrans{Bool, <:OneHotLike{T, L, N, var"N+1", <:CuArray}}) where {T, L, N, var"N+1"} =
  Base.print_array(io, cpu(X))

_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(x, labels, [default])

Return a `OneHotVector` which is roughly a sparse representation of `x .== labels`.

Instead of storing say `Vector{Bool}`, it stores the index of the first occurrence
of `x` in `labels`. If `x` is not found in labels, then it either returns `onehot(default, labels)`,
or gives an error if no default is given.

See also [`onehotbatch`](@ref) to apply this to many `x`s,
and [`onecold`](@ref) to reverse either of these, as well as to generalise `argmax`.

# Examples
```jldoctest
julia> β = Flux.onehot(:b, [:a, :b, :c])
3-element OneHotVector(::UInt32) with eltype Bool:
 ⋅
 1
 ⋅

julia> αβγ = (Flux.onehot(0, 0:2), β, Flux.onehot(:z, [:a, :b, :c], :c))  # uses default
(Bool[1, 0, 0], Bool[0, 1, 0], Bool[0, 0, 1])

julia> hcat(αβγ...)  # preserves sparsity
3×3 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
 1  ⋅  ⋅
 ⋅  1  ⋅
 ⋅  ⋅  1
```
"""
function onehot(x, labels)
  i = something(findfirst(isequal(x), labels), 0)
  i > 0 || error("Value $x is not in labels")
  OneHotVector{UInt32, length(labels)}(i)
end

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

"""
    onehotbatch(xs, labels, [default])

Returns a `OneHotMatrix` where `k`th column of the matrix is [`onehot(xs[k], labels)`](@ref onehot).
This is a sparse matrix, which stores just a `Vector{UInt32}` containing the indices of the
nonzero elements.

If one of the inputs in `xs` is not found in `labels`, that column is `onehot(default, labels)`
if `default` is given, else an error.

If `xs` has more dimensions, `M = ndims(xs) > 1`, then the result is an
`AbstractArray{Bool, M+1}` which is one-hot along the first dimension,
i.e. `result[:, k...] == onehot(xs[k...], labels)`.

# Examples
```jldoctest
julia> oh = Flux.onehotbatch(collect("abracadabra"), 'a':'e', 'e')
5×11 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
 1  ⋅  ⋅  1  ⋅  1  ⋅  1  ⋅  ⋅  1
 ⋅  1  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  1  ⋅  ⋅
 ⋅  ⋅  ⋅  ⋅  1  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅
 ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  1  ⋅  ⋅  ⋅  ⋅
 ⋅  ⋅  1  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  1  ⋅

julia> reshape(1:15, 3, 5) * oh  # this matrix multiplication is done efficiently
3×11 Matrix{Int64}:
 1  4  13  1  7  1  10  1  4  13  1
 2  5  14  2  8  2  11  2  5  14  2
 3  6  15  3  9  3  12  3  6  15  3
```
"""
onehotbatch(ls, labels, default...) = batch([onehot(l, labels, default...) for l in ls])

"""
    onecold(y::AbstractArray, labels = 1:size(y,1))

Roughly the inverse operation of [`onehot`](@ref) or [`onehotbatch`](@ref):
This finds the index of the largest element of `y`, or each column of `y`,
and looks them up in `labels`.

If `labels` are not specified, the default is integers `1:size(y,1)` --
the same operation as `argmax(y, dims=1)` but sometimes a different return type.

# Examples
```jldoctest
julia> Flux.onecold([false, true, false])
2

julia> Flux.onecold([0.3, 0.2, 0.5], [:a, :b, :c])
:c

julia> Flux.onecold([ 1  0  0  1  0  1  0  1  0  0  1
                      0  1  0  0  0  0  0  0  1  0  0
                      0  0  0  0  1  0  0  0  0  0  0
                      0  0  0  0  0  0  1  0  0  0  0
                      0  0  1  0  0  0  0  0  0  1  0 ], 'a':'e') |> String
"abeacadabea"
```
"""
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

function Base.:(*)(A::AbstractMatrix, B::OneHotLike{<:Any, L, 1}) 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 NNlib.gather(A, _indices(B))
end

function Base.:(*)(A::AbstractMatrix, B::Adjoint{Bool, <:OneHotMatrix})
  B_dim = length(_indices(parent(B)))
  size(A, 2) == B_dim || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $B_dim"))
  return NNlib.scatter(+, A, _indices(parent(B)), dstsize=(size(A,1), size(B,2)))
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


function minTest1()
  Rect3x2::Matrix{Float64} = [0.1 0.2; 0.3 0.4;0.5 0.6]
  Rect2x3::Matrix{Float64} = zeros(2,3)
  Rect2x3' .= Rect3x2
  println("Still Zeros: ", Rect2x3)
  Rect2x3 .= Rect3x2'
  println("Seems to Work: ", Rect2x3)
end

function foo(A, x, b=Vector{Float64}(length(x)))
  b .= A * x
  return b
end

println(foo([1 2; 3 4], [1,2]))
----------------
#this tests for a notice period negative bias to performance
#this version
function NNP_1D(HFLong::DataFrame; dependentVar::Symbol = :performance1D,
    titleText::String =  "Performance Near Withdawels: Impact of Notice Period")::String
    thresholdLS::RedemptionThresholds = RedemptionThresholds(SIG_REDEMPTION_THRESHOLDS)

    modelsNNP_1D::Vector{CTLM} = Vector{CTLM}()

    XNNP_1DSpecs::Vector{CTExpr} = Vector{CTExpr}()
    XNNP_1DNames::Vector{Vector{Symbol}}  = Vector{Vector{Symbol}}()

    push!(XNNP_1DSpecs, parse("noticePeriod"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod])

    push!(XNNP_1DSpecs, parse("noticePeriod + main_strategy"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod])

    push!(XNNP_1DSpecs, parse("noticePeriod + fRedemptionDate + main_strategy"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod])

    push!(XNNP_1DSpecs, parse("noticePeriod + fDate + main_strategy"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod])

    push!(XNNP_1DSpecs, parse("noticePeriod&negLFlows"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod_negLFlows])

    push!(XNNP_1DSpecs, parse("noticePeriod&negLFlows + main_strategy"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod_negLFlows])

    push!(XNNP_1DSpecs, parse("noticePeriod&negLFlows + fRedemptionDate + main_strategy"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod_negLFlows])

    push!(XNNP_1DSpecs, parse("noticePeriod&negLFlows + fDate + main_strategy"))
    push!(XNNP_1DNames, [:intercept; :noticePeriod_negLFlows])

    #push the specifications (use a common name to make a clean table)
    #=for s::Symbol ∈ thresholdLS.negLFlows

        push!(XNNP_1DSpecs, parse("noticePeriod&$s"))
        push!(XNNP_1DNames, [:intercept; :noticePeriod])

        push!(XNNP_1DSpecs, parse("noticePeriod&$s + fDate + main_strategy"))
        push!(XNNP_1DNames, [:intercept; :noticePeriod])

    end=#

    #run the model
    for iX ∈ 1:length(XNNP_1DSpecs)
        push!(modelsNNP_1D, CTLM(HFLong, XNNP_1DSpecs[iX],
            dependentVar, XNames=XNNP_1DNames[iX], YName = :performance))
    end

    ######IO Code for performance bias analysis

    descRowNamesNNP_1D::Vector{String} = ["Date F.E.",  "Redemp. Date F.E.","Strategy F.E.",#="Outflow \\% >", =# "N ('000s)"]

    #need to print the descriptive rows
    descContentNNP_1D::Vector{Vector{String}} =
        [Vector{String}(length(modelsNNP_1D)) for i∈ 1:length(descRowNamesNNP_1D)]

    #Determine the threshold levels (searches for the threshold level in each spec)
    thresholdVec::Vector{Int} = ((i::Int)->maximum(
        ((j::Int)->contains(XNNP_1DSpecs[i],"$(thresholdLS.negLFlows[j])")?
            Int(100.0*collect(SIG_REDEMPTION_THRESHOLDS)[j]):0).(1:length(thresholdLS.negLFlows))
        )).(1:length(XNNP_1DSpecs))

    for i ∈ 1:length(XNNP_1DSpecs)
        descContentNNP_1D[1][i] = "$(contains(XNNP_1DSpecs[i], "fDate")?"X":"")"
        descContentNNP_1D[2][i] = "$(contains(XNNP_1DSpecs[i], "fRedemptionDate")?"X":"")"
        descContentNNP_1D[3][i] = "$(contains(XNNP_1DSpecs[i], "main_strategy")?"X":"")"
        #descContentNNP_1D[3][i] = "$(contains(XNNP_1DSpecs[i], "negLFlows")?"X":"")"
        #descContentNNP_1D[4][i] = "$(thresholdVec[i])\\%"
        descContentNNP_1D[end][i] = "$(round(modelsNNP_1D[i].N/1000.0,1))"
    end

    #rows and variable names for the regression
    rowsNNP_1D::Vector{Symbol} = [:intercept; :noticePeriod; :noticePeriod_negLFlows]

    TableTextNNP_1D::String = texTable(modelsNNP_1D, getHomoskedΣ!, rowsNNP_1D,
        titleCaption = "Performance Flow-Bias Specifications",
        colNames = [["($i)" for i∈1:length(modelsNNP_1D)]],
        contentRowNames = ["intercept", "Notice Period", "Notice X Outflow"],
        descRowNames = descRowNamesNNP_1D,
        descContent = descContentNNP_1D,
        decimalDigits = 2,
        columnSepPt = -20,
        scaling = fill(100.0,length(rowsNNP_1D))::Vector{Float64},
        clearMem = USE_AGGRESSIVE_GC,
        caption = """NNP_1D specification compares performance prior to notice,
            with performance after notice. In some specifications, the notice
            period dummy is interacted with the amount of the outflow. Threshold
            refers to the minimum outflow included, thus some versions include
            only significant outflows. Data is from alive and dead hedge funds
            2000-2017. Interpret a value of 1.0 for the non-interacted
            specification as the discontinuity increasing performance by 1\\%.
            For the interacted specification, a 1.0 implies that outflows of
            10\\% corresponding to a 0.1\\% drop in monthly performance.""")

    return TableTextNNP_1D
end


using MosaicViews
using Test
using ImageCore, ColorVectorSpace
using OffsetArrays

# Because of the difference in axes types between paddedviews (Base.OneTo) and
# sym_paddedviews (UnitRange), the return type of `_padded_cat` isn't inferrable to a
# concrete type. But it is inferrable to a Union{A,B} where both A and B are concrete.
# While `@inferred(_padded_cat((A, B), ...))` would therefore fail, this is a close substitute.
function _checkinferred_paddedcat(V, A; kwargs...)
    vd = MosaicViews.valdim(first(A))
    RTs = Base.return_types(MosaicViews._padded_cat, (typeof(A), Bool, eltype(V), typeof(vd)))
    @test length(RTs) == 1
    RT = RTs[1]
    @test isconcretetype(RT) || (isa(RT, Union) && isconcretetype(RT.a) && isconcretetype(RT.b))
    return V
end
function checkinferred_mosaic(As::Tuple; kwargs...)
    V = mosaic(As; kwargs...)
    return _checkinferred_paddedcat(V, As)
end
function checkinferred_mosaic(A...; kwargs...)
    V = mosaic(A...; kwargs...)
    return _checkinferred_paddedcat(V, A)
end
function checkinferred_mosaic(As::AbstractVector{<:AbstractArray}; kwargs...)
    V = mosaic(As; kwargs...)
    return _checkinferred_paddedcat(V, (As...,); kwargs...)
end

@testset "MosaicView" begin
    @test_throws ArgumentError MosaicView(rand(2,2,2,2,2))

    @testset "1D input" begin
        A = [1,2,3]
        mva = @inferred MosaicView(A)
        @test parent(mva) === A
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (3, 1) # 1D vector is lifted to 2D matrix
        @test axes(mva) == (Base.OneTo(3), Base.OneTo(1))
        @test @inferred(getindex(mva, 1, 1)) === 1
        @test @inferred(getindex(mva, 2, 1)) === 2
    end

    @testset "2D input" begin
        A = [1 2;3 4]
        mva = @inferred MosaicView(A)
        @test parent(mva) === A
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (2, 2)
        @test axes(mva) == (Base.OneTo(2), Base.OneTo(2))
        @test @inferred(getindex(mva, 1, 1)) === 1
        @test @inferred(getindex(mva, 2, 1)) === 3
    end

    @testset "3D input" begin
        A = zeros(Int,2,2,2)
        A[:,:,1] = [1 2; 3 4]
        A[:,:,2] = [5 6; 7 8]
        mva = @inferred MosaicView(A)
        @test parent(mva) === A
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (4, 2)
        @test @inferred(getindex(mva,1,1)) === 1
        @test @inferred(getindex(mva,2,1)) === 3
        @test_throws BoundsError mva[0,1]
        @test_throws BoundsError mva[1,0]
        @test_throws BoundsError mva[1,3]
        @test_throws BoundsError mva[5,1]
        @test all(mva .== vcat(A[:,:,1],A[:,:,2]))
        # singleton dimension doesn't change anything
        @test mva == MosaicView(reshape(A,2,2,2,1))
    end

    @testset "4D input" begin
        A = zeros(Int,2,2,1,2)
        A[:,:,1,1] = [1 2; 3 4]
        A[:,:,1,2] = [5 6; 7 8]
        mva = @inferred MosaicView(A)
        @test parent(mva) === A
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (2, 4)
        @test @inferred(getindex(mva,1,1)) === 1
        @test @inferred(getindex(mva,2,1)) === 3
        @test_throws BoundsError mva[0,1]
        @test_throws BoundsError mva[1,0]
        @test_throws BoundsError mva[3,1]
        @test_throws BoundsError mva[1,5]
        @test all(mva .== hcat(A[:,:,1,1],A[:,:,1,2]))
        A = zeros(Int,2,2,2,3)
        A[:,:,1,1] = [1 2; 3 4]
        A[:,:,1,2] = [5 6; 7 8]
        A[:,:,1,3] = [9 10; 11 12]
        A[:,:,2,1] = [13 14; 15 16]
        A[:,:,2,2] = [17 18; 19 20]
        A[:,:,2,3] = [21 22; 23 24]
        mva = @inferred MosaicView(A)
        @test parent(mva) === A
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (4, 6)
        @test @inferred(getindex(mva,1,1)) === 1
        @test @inferred(getindex(mva,2,1)) === 3
        @test all(mva .== vcat(hcat(A[:,:,1,1],A[:,:,1,2],A[:,:,1,3]), hcat(A[:,:,2,1],A[:,:,2,2],A[:,:,2,3])))
    end
end

@testset "mosaicview" begin
    @testset "1D input" begin
        A1 = [1, 2, 3]
        A2 = [4, 5, 6]
        mva = checkinferred_mosaic(A1, A2)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A1)
        @test size(mva) == (6, 1)
        @test @inferred(getindex(mva, 1, 1)) == 1

        @test checkinferred_mosaic(A1, A2; nrow=1) == [1 4; 2 5; 3 6]
        @test mosaic(A1, A2; nrow=2) == reshape([1, 2, 3, 4, 5, 6], (6, 1))
        @test mosaic(A1, A2) == mosaic([A1, A2]) == mosaic((A1, A2))
    end

    @testset "2D input" begin
        A = [i*ones(Int, 2, 3) for i in 1:4]
        Ao = [i*ones(Int, 0i:1, 0:2) for i in 1:4]

        for B in (A, tuple(A...), Ao)
            @test_throws ArgumentError mosaic(B, nrow=0)
            @test_throws ArgumentError mosaic(B, ncol=0)
            @test_throws ArgumentError mosaic(B, nrow=1, ncol=1)

            mva = checkinferred_mosaic(B)
            @test mosaic(B...) == mva
            @test typeof(mva) <: MosaicView
            @test eltype(mva) == eltype(eltype(B))
            @test size(mva) == (8, 3)
            @test @inferred(getindex(mva,3,1)) === 2
            @test collect(mva) == [
                1  1  1
                1  1  1
                2  2  2
                2  2  2
                3  3  3
                3  3  3
                4  4  4
                4  4  4
            ]

            mva = checkinferred_mosaic(B, nrow=2)
            @test typeof(mva) <: MosaicView
            @test eltype(mva) == eltype(eltype(B))
            @test size(mva) == (4, 6)
            @test collect(mva) == [
             1  1  1  3  3  3
             1  1  1  3  3  3
             2  2  2  4  4  4
             2  2  2  4  4  4
            ]
        end

        @test mosaic(A...) == mosaic(A)
        @test mosaic(A..., nrow=2) == mosaic(A, nrow=2)
        @test mosaic(A..., nrow=2, rowmajor=true) == mosaic(A, nrow=2, rowmajor=true)

        A1 = reshape([1 2 3], (1, 3))
        A2 = reshape([4;5;6], (3, 1))
        @test checkinferred_mosaic([A1, A2]; center=false) == [
         1 2 3;
         0 0 0;
         0 0 0;
         4 0 0;
         5 0 0;
         6 0 0
        ]
        @test mosaic([A1, A2]; center=true) == [
         0 0 0;
         1 2 3;
         0 0 0;
         0 4 0;
         0 5 0;
         0 6 0
        ]
        @test mosaic([A1, A2]) == mosaic([A1, A2]; center=true)

        # same size but different axes
        A1 = fill(1, 1:2, 1:2)
        A2 = fill(2, 2:3, 2:3)
        @test collect(checkinferred_mosaic(A1, A2; center=true)) == [
            1 1 0;
            1 1 0;
            0 0 0;
            2 2 0;
            2 2 0;
            0 0 0;
        ]
        @test collect(mosaic(A1, A2; center=false)) == [
            1 1 0;
            1 1 0;
            0 0 0;
            0 0 0;
            0 2 2;
            0 2 2;
        ]

        @test mosaicview(A1) == A1 # a trivial case
    end

    @testset "3D input" begin
        A = [(k+1)*l-1 for i in 1:2, j in 1:3, k in 1:2, l in 1:2]
        B = reshape(A, 2, 3, :)
        @test_throws ArgumentError mosaicview(B, nrow=0)
        @test_throws ArgumentError mosaicview(B, ncol=0)
        @test_throws ArgumentError mosaicview(B, nrow=1, ncol=1)

        mva = checkinferred_mosaic(B)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(B)
        @test size(mva) == (8, 3)
        @test @inferred(getindex(mva,3,1)) === 2
        @test mva == [
            1  1  1
            1  1  1
            2  2  2
            2  2  2
            3  3  3
            3  3  3
            5  5  5
            5  5  5
        ]
        mva = checkinferred_mosaic(B, nrow=2)
        @test mva == MosaicView(A)
        @test typeof(mva) != typeof(MosaicView(A))
        @test parent(parent(mva)).data == B
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(B)
        @test size(mva) == (4, 6)

        @test mosaic(B, B) == mosaicview(cat(B, B; dims=4))
        @test mosaic(B, B, nrow=2) == mosaicview(cat(B, B; dims=4), nrow=2)
        @test mosaic(B, B, nrow=2, rowmajor=true) == mosaicview(cat(B, B; dims=4), nrow=2, rowmajor=true)
    end

    @testset "4D input" begin
        A = [(k+1)*l-1 for i in 1:2, j in 1:3, k in 1:2, l in 1:2]
        @test_throws ArgumentError mosaicview(A, nrow=0)
        @test_throws ArgumentError mosaicview(A, ncol=0)
        @test_throws ArgumentError mosaicview(A, nrow=1, ncol=1)

        mva = mosaicview(A)
        @test mva == MosaicView(A)
        @test typeof(mva) != typeof(MosaicView(A))
        @test parent(parent(mva)).data == reshape(A, 2, 3, :)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (4, 6)
        mva = mosaicview(A, npad=1)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (5, 7)
        @test mva == mosaicview(A, nrow=2, npad=1)
        @test mva == mosaicview(A, ncol=2, npad=1)
        @test @inferred(getindex(mva,3,1)) === 0
        @test @inferred(getindex(mva,2,5)) === 3
        @test mva == [
            1  1  1  0  3  3  3
            1  1  1  0  3  3  3
            0  0  0  0  0  0  0
            2  2  2  0  5  5  5
            2  2  2  0  5  5  5
        ]
        mva = mosaicview(A, ncol=3, npad=1)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (5, 11)
        @test mva == [
            1  1  1  0  3  3  3  0  0  0  0
            1  1  1  0  3  3  3  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0
            2  2  2  0  5  5  5  0  0  0  0
            2  2  2  0  5  5  5  0  0  0  0
        ]
        mva = mosaicview(A, rowmajor=true, ncol=3, npad=1)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (5, 11)
        @test @inferred(getindex(mva,3,1)) === 0
        @test @inferred(getindex(mva,2,5)) === 2
        @test mva == [
            1  1  1  0  2  2  2  0  3  3  3
            1  1  1  0  2  2  2  0  3  3  3
            0  0  0  0  0  0  0  0  0  0  0
            5  5  5  0  0  0  0  0  0  0  0
            5  5  5  0  0  0  0  0  0  0  0
        ]
        mva = mosaicview(A, rowmajor=true, ncol=3, npad=2)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (6, 13)
        @test mva == [
            1  1  1  0  0  2  2  2  0  0  3  3  3
            1  1  1  0  0  2  2  2  0  0  3  3  3
            0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0
            5  5  5  0  0  0  0  0  0  0  0  0  0
            5  5  5  0  0  0  0  0  0  0  0  0  0
        ]
        mva = mosaicview(A, fillvalue=-1.0, rowmajor=true, ncol=3, npad=1)
        @test typeof(mva) <: MosaicView
        @test eltype(mva) == eltype(A)
        @test size(mva) == (5, 11)
        @test mva == [
             1   1   1  -1   2   2   2  -1   3   3   3
             1   1   1  -1   2   2   2  -1   3   3   3
            -1  -1  -1  -1  -1  -1  -1  -1  -1  -1  -1
             5   5   5  -1  -1  -1  -1  -1  -1  -1  -1
             5   5   5  -1  -1  -1  -1  -1  -1  -1  -1
        ]

        @test mosaic(A, A) == mosaicview(cat(A, A; dims=5))
        @test mosaic(A, A, nrow=2) == mosaicview(cat(A, A; dims=4), nrow=2)
        @test mosaic(A, A, nrow=2, rowmajor=true) == mosaicview(cat(A, A; dims=4), nrow=2, rowmajor=true)
    end

    @testset "Colorant Array" begin
        A = rand(RGB{Float32}, 2, 3, 2, 2)
        mvaa = mosaicview(A)
        @test eltype(mvaa) == eltype(A)
        @test mvaa == @inferred(MosaicView(A))
        mvaa = mosaicview(A, rowmajor=true, ncol=3)
        @test eltype(mvaa) == eltype(A)
        @test @inferred(getindex(mvaa, 3, 4)) == RGB(0,0,0)
        mvaa = mosaicview(A, fillvalue=colorant"white", rowmajor=true, ncol=3)
        @test eltype(mvaa) == eltype(A)
        @test @inferred(getindex(mvaa, 3, 4)) == RGB(1,1,1)

        # this should work regardless they're of different size and color
        @test_nowarn mosaic(rand(RGB{Float32}, 4, 4),
                            rand(Gray{N0f8}, 5, 5))
    end

    # all arrays should have the same dimension
    @test_throws ArgumentError mosaic(ones(2), ones(1, 2))
    @test_throws ArgumentError mosaic((ones(2), ones(1, 2)))
    @test_throws ArgumentError mosaic([ones(2), ones(1, 2)])

    @testset "filltype" begin
        # always a concrete type
        A = checkinferred_mosaic(rand(N0f8, 4, 4), rand(Float64, 4, 4), rand(Float32, 4, 4))
        @test eltype(A) == Float64

        A = mosaic(Any[1 2 3; 4 5 6], rand(Float32, 4, 4))
        @test eltype(A) == Float32
        A = mosaic(rand(Float32, 4, 4), Any[1 2 3; 4 5 6])
        @test eltype(A) == Float32

        # FIXME:
        # A = checkinferred_mosaic(rand(Float64, 4, 4), Union{Missing, Float32}[1 2 3; 4 5 6])
        A = mosaic(rand(Float64, 4, 4), Union{Missing, Float32}[1 2 3; 4 5 6])
        @test eltype(A) == Union{Missing, Float64}
    end
end

@testset "deprecations" begin
    @info "deprecations are expected"
    # mosaicview -> mosaic deprecations
    A = [(k+1)*l-1 for i in 1:2, j in 1:3, k in 1:2, l in 1:2]
    mva_old = mosaicview(A, A, rowmajor=true, ncol=3, npad=1)
    mva_new = mosaic(A, A, rowmajor=true, ncol=3, npad=1)
    @test mva_old == mva_new
    mva_old = mosaicview(A, A, A, rowmajor=true, ncol=3, npad=1)
    mva_new = mosaic(A, A, A, rowmajor=true, ncol=3, npad=1)
    @test mva_old == mva_new
    mva_old = mosaicview([A, A], rowmajor=true, ncol=3, npad=1)
    mva_new = mosaic([A, A], rowmajor=true, ncol=3, npad=1)
    @test mva_old == mva_new
    mva_old = mosaicview((A, A), rowmajor=true, ncol=3, npad=1)
    mva_new = mosaic((A, A), rowmajor=true, ncol=3, npad=1)
    @test mva_old == mva_new
    # center keyword for `mosaicview` (still applies for `mosaic`)
    A = [(k+1)*l-1 for i in 1:2, j in 1:3, k in 1:2, l in 1:2]
    B = reshape(A, 2, 3, :)
    @test mosaicview(B, center=2) == mosaicview(B) # no op
    @test mosaicview(A, center=2) == mosaicview(A) # no op
end


import Base: show
import Oceananigans.Grids: short_show
using Oceananigans.Grids: show_domain

location_str(::Face) = "Face"
location_str(::Cell) = "Cell"
show_location(X, Y, Z) = "($(location_str(X())), $(location_str(Y())), $(location_str(Z())))"
show_location(field::AbstractField{X, Y, Z}) where {X, Y, Z} = show_location(X, Y, Z)

short_show(field::Field) = string("Field located at ", show_location(field))

show(io::IO, field::Field) =
    print(io, "$(short_show(field))\n",
          "├── data: $(typeof(field.data)), size: $(size(field.data))\n",
          "├── grid: $(short_show(field.grid))\n",
          "└── boundary conditions: $(short_show(field.boundary_conditions))")

short_show(array::OffsetArray{T, D, A}) where {T, D, A} = string("OffsetArray{$T, $D, $A}")


function fahrenheit_to_celsius(F)
    C = 0

    # Your code goes here: calculate the temperature in Celsius,
    # store in a variable (we called it C), and return it.

    return C
end


# To ensure that tests can run on travis we have to do a little
# hackadoodle here. The tests require a license file. We include
# a license file that is only valid for one day (the day when
# change is submitted).
# If there is no valid license file, we default to that file.

if haskey(ENV,"MOSEKLM_LICENSE_FILE")
    # that's nice
elseif haskey(ENV,"HOME")
    if isfile(joinpath(ENV["HOME"],"mosek","mosek.lic"))
        # our lucky day!
    else
        licfile = joinpath(@__DIR__,"..","test",".dontuse-probablyexpired.lic")
        println("Use license file: $licfile")
        import Mosek
        Mosek.putlicensepath(Mosek.msk_global_env,licfile)
    end
end

using MathOptInterfaceMosek

using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIB = MOI.Bridges

const optimizer = MosekOptimizer(QUIET = true)
# 1e-3 needed for rotatedsoc3 test
const config = MOIT.TestConfig(atol=1e-3, rtol=1e-3, query=false)

@testset "Continuous linear problems" begin
    # linear1 is failing because NumberOfConstraints does not take deletion into account
    # linear11 is failing because the following are not implemented:
    # * MOI.cantransformconstraint(instance, c2, MOI.LessThan(2.0))
    # * MOI.get(instance, MathOptInterface.ConstraintFunction())
    # linear13 is failing because it is FeasibilitySense
    MOIT.contlineartest(optimizer, config, ["linear1", "linear11", "linear13"])
end

# include("contquadratic.jl")
# @testset "Continuous quadratic problems" begin
#     # contquadratictest(GurobiSolver())
# end

@testset "Continuous conic problems" begin
    # lin1 and soc1 are failing because ListOfConstraints is not implemented
    MOIT.contconictest(MOIB.RootDet{Float64}(MOIB.GeoMean{Float64}(optimizer)), config, ["lin1v", "lin1f", "soc1v", "soc1f", "exp", "psds", "rootdets", "logdet"])
end

@testset "Mixed-integer linear problems" begin
    MOIT.intlineartest(optimizer, config, ["int2"])
end

#include("test_jump.jl")


using Random

"""

    sort_biosystem_compounds(compound_structs,
        biosystem_compounds,
        sortkey,
        zero_mass_behavior)

Return a vector of (compound,mass) sorted by increasing mass.
Always shuffles the order of the zero-mass compounds.

...
# Arguments
- `compound_structs::Compounds`: compounds with metadata
- `biosystem_compounds::IDs`: compounds to sort
- `sortkey::Symbol=:exact_mass`: how to sort compounds. Can also sort by `:mol_weight`
- `zero_mass_behavior::String="end"`: what to do with compounds that have no mass. Can also be `"random"`
...
"""
function sort_biosystem_compounds(
    compound_structs::Compounds,
    biosystem_compounds::IDs,
    sortkey::Symbol=:exact_mass,
    zero_mass_behavior::String="end")

    cpd_masses = [(i,getproperty(compound_structs[i],sortkey)) for i in biosystem_compounds if getproperty(compound_structs[i],sortkey)!=0]
    cpd_zero_masses = [(i,getproperty(compound_structs[i],sortkey)) for i in biosystem_compounds if getproperty(compound_structs[i],sortkey)==0]

    sort!(cpd_masses, by= x->x[2])

    if zero_mass_behavior=="end"
        return vcat(cpd_masses,shuffle(cpd_zero_masses))

    elseif zero_mass_behavior=="random"
        for tup in cpd_zero_masses
            insert!(cpd_masses,rand(1:length(cpd_masses)),tup) # this will permit insertion at longer indices as cpd_masses grows
        end
        return cpd_masses
    end
end

"""
    mix_it_up(tuples,beta,n_swaps)

Swap random tuples within `tuples`.
Based on formula from Handorf et al. 2008.
"""
function mix_it_up!(
    tuples::Vector{<:Tuple{String,Real}},
    beta::Real,
    n_swaps::Int,
    rng::AbstractRNG=Random.default_rng()) #where {R <: Real}

    for _ in 1:n_swaps
        i,j = rand(1:length(tuples),2) ## 2 random indices
        if swap_random(tuples[i],tuples[j],beta,rng) == true
            tuples[i], tuples[j] = tuples[j], tuples[i] ## changes them simultaneously
        end
    end
end

"""
    swap_random(ti,tj,beta)

Return `true` if tuples should be flipped.
"""
function swap_random(
    ti::Tuple{String,<:Real},
    tj::Tuple{String,<:Real},
    beta::Real,
    rng::AbstractRNG=Random.default_rng())

    p = getflipprobability(ti,tj,beta)

    rand(rng,Float64)<p ? true : false
end

"""
    getflipprobability(ti,tj,beta)

Get the probability of flipping tuples ti and tj.
Based on formula from Handorf et al. 2008.
"""
function getflipprobability(
    ti::Tuple{String,Real},
    tj::Tuple{String,Real},
    beta::Real)

    diff = ti[2] - tj[2]

    ## Determine probability of flipping (another way of writing below)
    # ti[1]==0 || tj[1]==0 ? p=.5 : diff>0 ? p = math.e**(-diff/float(beta)) : p=1.0

    if ti[2]==0 || tj[2]==0
        p = .5
    elseif diff > 0
        p = exp(-diff/beta)
    else
        p = 1.0
    end

    return p

end

"""
    randomizecompounds(biosystem_compounds,
        compound_structs,
        n_runs,
        n_swaps,
        beta,
        sortkey,
        zero_mass_behavior)

Return a vector of run results.
Each run result is a randomized list of biosystem_compounds.

...
# Arguments
- `biosystem_compounds::IDs`: compounds to sort
- `compound_structs::Compounds`: compounds with metadata
- `n_runs::Int`: number of runs
- `n_swaps::Int=1000`: number of compounds to within biosystem_compounds to swap when randomizing
- `beta::Real=20`: Mixing coefficient. See Handorf et al. 2008 for details.
- `sortkey::Symbol=:exact_mass`: how to sort compounds. Can also sort by `:mol_weight`
- `zero_mass_behavior::String="end"`: what to do with compounds that have no mass. Can also be `"random"`
...
"""
function randomizecompounds(
    biosystem_compounds::IDs,
    compound_structs::Compounds,
    n_runs::Int,
    n_swaps::Int=1000,
    beta::Real=20,
    sortkey::Symbol=:exact_mass,
    zero_mass_behavior::String="end",
    rng::AbstractRNG=Random.default_rng())

    randomized_cpd_lists = [Vector{String}(undef,length(biosystem_compounds)) for _ in 1:n_runs] # allocate output memory
    ## Generate seeds and RNG objects based on rng
    rng_seeds = [rand(rng,1:typemax(Int)) for i in 1:n_runs]
    rng_list = [MersenneTwister(i) for i in rng_seeds]
    for r in 1:n_runs
        cpds_masses = sort_biosystem_compounds(compound_structs, biosystem_compounds, sortkey, zero_mass_behavior)
        mix_it_up!(cpds_masses,beta,n_swaps,rng_list[r])
        randomized_cpd_lists[r] = [c[1] for c in cpds_masses]
    end
    return randomized_cpd_lists
end


using ChainRulesTestUtils
using FiniteDifferences
using ITensors
using Random

using ChainRulesCore: NoTangent

# TODO: maybe move this into ITensorChainRules module?
# These are useful definitions for testing code. The downside
# is adding dependencies on FiniteDifferences and ChainRulesTestUtils

#
# For ITensor compatibility with FiniteDifferences
#

function FiniteDifferences.to_vec(A::ITensor)
  # TODO: generalize to sparse tensors
  # TODO: define `itensor([1.0])` as well
  # as `itensor([1.0], ())` to help with generic code.
  function vec_to_ITensor(x)
    return isempty(inds(A)) ? ITensor(x[]) : itensor(x, inds(A))
  end
  return vec(array(A)), vec_to_ITensor
end

function FiniteDifferences.to_vec(x::Index)
  return (Bool[], _ -> x)
end

function FiniteDifferences.to_vec(x::Tuple{Vararg{Index}})
  return (Bool[], _ -> x)
end

function FiniteDifferences.to_vec(x::Pair{<:Tuple{Vararg{Index}},<:Tuple{Vararg{Index}}})
  return (Bool[], _ -> x)
end

function FiniteDifferences.rand_tangent(rng::AbstractRNG, A::ITensor)
  # TODO: generalize to sparse tensors
  return isempty(inds(A)) ? ITensor(randn(eltype(A))) : randomITensor(eltype(A), inds(A))
end

function FiniteDifferences.rand_tangent(rng::AbstractRNG, x::Index)
  return NoTangent()
end

function FiniteDifferences.rand_tangent(rng::AbstractRNG, x::Tuple{Vararg{Index}})
  return NoTangent()
end

function FiniteDifferences.rand_tangent(
  rng::AbstractRNG, x::Pair{<:Tuple{Vararg{Index}},<:Tuple{Vararg{Index}}}
)
  return NoTangent()
end

#
# For ITensor compatibility with ChainRulesTestUtils
#

function ChainRulesTestUtils.test_approx(
  actual::ITensor, expected::ITensor, msg=""; kwargs...
)
  ChainRulesTestUtils.@test_msg msg isapprox(actual, expected; kwargs...)
end

function ChainRulesTestUtils.test_approx(
  actual::ITensor, expected::Number, msg=""; kwargs...
)
  ChainRulesTestUtils.@test_msg msg isapprox(actual[], expected; kwargs...)
end

function ChainRulesTestUtils.test_approx(
  actual::Number, expected::ITensor, msg=""; kwargs...
)
  ChainRulesTestUtils.@test_msg msg isapprox(actual, expected[]; kwargs...)
end


function part1(filename::AbstractString)
    sum(length(line) - length(unescape_string(line[begin+1:end-1])) for line in eachline(filename))
end

function part2(filename::AbstractString)
    sum(length(escape_string(line)) - length(line) + 2 for line in eachline(filename))
end


@testset "Conversion from and to MATLAB model" begin
    filename = model_paths["iJO1366.mat"]

    mm = load_mat_model(filename)
    sm = convert(StandardModel, mm)
    mm2 = convert(MATModel, sm)

    @test Set(reactions(mm)) == Set(reactions(sm))
    @test Set(reactions(mm)) == Set(reactions(mm2))
end

@testset "MATModel generic interface" begin
    model = load_model(model_paths["e_coli_core.mat"])

    @test reaction_stoichiometry(model, "EX_ac_e") == Dict("ac_e" => -1)
    @test reaction_stoichiometry(model, 44) == Dict("ac_e" => -1)
end