mutable struct BasicContMuvParameterNState{N<:Real} <: ParameterNState{Continuous, Multivariate}
  value::Matrix{N}
  loglikelihood::Vector{N}
  logprior::Vector{N}
  logtarget::Vector{N}
  gradloglikelihood::Matrix{N}
  gradlogprior::Matrix{N}
  gradlogtarget::Matrix{N}
  tensorloglikelihood::Array{N, 3}
  tensorlogprior::Array{N, 3}
  tensorlogtarget::Array{N, 3}
  dtensorloglikelihood::Array{N, 4}
  dtensorlogprior::Array{N, 4}
  dtensorlogtarget::Array{N, 4}
  diagnosticvalues::Matrix
  size::Integer
  sizesquared::Integer
  sizecubed::Integer
  monitor::Vector{Bool}
  n::Integer
  diagnostickeys::Vector{Symbol}

  function BasicContMuvParameterNState{N}(
    size::Integer,
    n::Integer,
    monitor::Vector{Bool}=[true; fill(false, 12)],
    diagnostickeys::Vector{Symbol}=Symbol[],
    ::Type{N}=Float64,
    diagnosticvalues::Matrix=Array{Any}(length(diagnostickeys), isempty(diagnostickeys) ? 0 : n)
  ) where N<:Real
    instance = new()

    fnames = fieldnames(BasicContMuvParameterNState)
    for i in 2:4
      l = (monitor[i] == false ? zero(Integer) : n)
      setfield!(instance, fnames[i], Array{N}(l))
    end
    for i in (1, 5, 6, 7)
      s, l = (monitor[i] == false ? (zero(Integer), zero(Integer)) : (size, n))
      setfield!(instance, fnames[i], Array{N}(s, l))
    end
    for i in 8:10
      s, l = (monitor[i] == false ? (zero(Integer), zero(Integer)) : (size, n))
      setfield!(instance, fnames[i], Array{N}(s, s, l))
    end
    for i in 11:13
      s, l = (monitor[i] == false ? (zero(Integer), zero(Integer)) : (size, n))
      setfield!(instance, fnames[i], Array{N}(s, s, s, l))
    end

    instance.diagnosticvalues = diagnosticvalues
    instance.size = size
    instance.sizesquared = instance.size^2
    instance.sizecubed = instance.size^3
    instance.monitor = monitor
    instance.n = n
    instance.diagnostickeys = diagnostickeys

    instance
  end
end

BasicContMuvParameterNState(
  size::Integer,
  n::Integer,
  monitor::Vector{Bool}=[true; fill(false, 12)],
  diagnostickeys::Vector{Symbol}=Symbol[],
  ::Type{N}=Float64,
  diagnosticvalues::Matrix=Array{Any}(length(diagnostickeys), isempty(diagnostickeys) ? 0 : n)
) where N<:Real =
  BasicContMuvParameterNState{N}(size, n, monitor, diagnostickeys, N, diagnosticvalues)

function BasicContMuvParameterNState(
  size::Integer,
  n::Integer,
  monitor::Vector{Symbol},
  diagnostickeys::Vector{Symbol}=Symbol[],
  ::Type{N}=Float64,
  diagnosticvalues::Matrix=Array{Any}(length(diagnostickeys), isempty(diagnostickeys) ? 0 : n)
) where N<:Real
  fnames = fieldnames(BasicContMuvParameterNState)
  BasicContMuvParameterNState(
    size, n, [fnames[i] in monitor ? true : false for i in 1:13], diagnostickeys, N, diagnosticvalues
  )
end

const ContMuvMarkovChain = BasicContMuvParameterNState

function copy!(nstate::BasicContMuvParameterNState, state::BasicContMuvParameterState, i::Integer)
  fnames = fieldnames(BasicContMuvParameterNState)

  for j in 2:4
    if nstate.monitor[j]
      getfield(nstate, fnames[j])[i] = getfield(state, fnames[j])
    end
  end

  for j in (1, 5, 6, 7)
    if nstate.monitor[j]
      getfield(nstate, fnames[j])[1+(i-1)*state.size:i*state.size] = getfield(state, fnames[j])
    end
  end

  for j in 8:10
    if nstate.monitor[j]
      getfield(nstate, fnames[j])[1+(i-1)*nstate.sizesquared:i*nstate.sizesquared] = getfield(state, fnames[j])
    end
  end

  for j in 11:13
    if nstate.monitor[j]
      getfield(nstate, fnames[j])[1+(i-1)*nstate.sizecubed:i*nstate.sizecubed] = getfield(state, fnames[j])
    end
  end

  if !isempty(nstate.diagnosticvalues)
    nstate.diagnosticvalues[:, i] = state.diagnosticvalues
  end
end

eltype{N<:Real}(::Type{BasicContMuvParameterNState{N}}) = N
eltype{N<:Real}(::BasicContMuvParameterNState{N}) = N

==(z::S, w::S) where {S<:BasicContMuvParameterNState} = reduce(&, [getfield(z, n) == getfield(w, n) for n in fieldnames(S)[1:17]])

isequal(z::S, w::S) where {S<:BasicContMuvParameterNState} =
  reduce(&, [isequal(getfield(z, n), getfield(w, n)) for n in fieldnames(S)[1:17]])

function show(io::IO, nstate::BasicContMuvParameterNState{N}) where N<:Real
  fnames = fieldnames(BasicContMuvParameterNState)
  fbool = map(n -> !isempty(getfield(nstate, n)), fnames[1:13])
  indentation = "  "

  println(io, "BasicContMuvParameterNState:")

  println(io, indentation*"eltype: $(eltype(nstate))")
  println(io, indentation*"state size = $(nstate.size)")
  println(io, indentation*"number of states = $(nstate.n)")

  print(io, indentation*"monitored components:")
  if !any(fbool)
    println(io, " none")
  else
    print(io, "\n")
    for i in 1:13
      if fbool[i]
        println(io, string(indentation^2, fnames[i]))
      end
    end
  end

  print(io, indentation*"diagnostics:")
  if isempty(nstate.diagnostickeys)
    print(io, " none")
  else
    for k in nstate.diagnostickeys
      print(io, "\n")
      print(io, string(indentation^2, k))
    end
  end
end


#####
##### AlphaZero.jl
##### Jonathan Laurent, Carnegie Mellon University (2019)
#####

module AlphaZero

export MCTS, MinMax, GameInterface, GI, Report, Network, Benchmark
export Params, SelfPlayParams, LearningParams, ArenaParams
export MctsParams, MemAnalysisParams
export Env, train!, learning!, self_play!, get_experience
export AbstractGame, AbstractPlayer, interactive!
export MctsPlayer, RandomPlayer, EpsilonGreedyPlayer, NetworkPlayer, Human
export SamplesWeighingPolicy, CONSTANT_WEIGHT, LOG_WEIGHT, LINEAR_WEIGHT
export ColorPolicy, ALTERNATE_COLORS, BASELINE_WHITE, CONTENDER_WHITE
export AbstractNetwork, OptimiserSpec, Nesterov, CyclicNesterov, Adam
export SimpleNet, SimpleNetHP, ResNet, ResNetHP
export AbstractSchedule, PLSchedule, StepSchedule
export Session, resume!, save, play_game, run_new_benchmark
export Explorer, explore

include("util.jl")
import .Util
using .Util: Option, @unimplemented

include("game.jl")
using .GameInterface
const GI = GameInterface

include("mcts.jl")
import .MCTS

include("networks/network.jl")
using .Network

include("ui/log.jl")
using .Log

import Plots
import Colors
import JSON2
import JSON3

using Formatting
using Crayons
using Colors: @colorant_str
using ProgressMeter
using Base: @kwdef
using Serialization: serialize, deserialize
using DataStructures: Stack, CircularBuffer
using Distributions: Categorical, Dirichlet
using Statistics: mean

include("schedule.jl")
include("params.jl")
include("report.jl")
include("memory.jl")
include("learning.jl")
include("play.jl")
include("training.jl")

include("minmax.jl")
import .MinMax

include("benchmark.jl")
using .Benchmark

# We support Flux and Knet
const USE_KNET = true

if USE_KNET
  @eval begin
    include("networks/knet.jl")
    using .KNets
  end
else
  @eval begin
    include("networks/flux.jl")
    using .FluxNets
  end
end

include("ui/explorer.jl")
include("ui/plots.jl")
include("ui/json.jl")
include("ui/session.jl")

end


using Base: @propagate_inbounds

using OffsetArrays
using JLD2

using Oceananigans.Architectures
using Oceananigans.Grids
using Oceananigans.Fields

using Oceananigans.Grids: topology, total_size, interior_parent_indices
using Oceananigans.Fields: show_location

import Oceananigans: short_show
import Oceananigans.Fields: Field, set!, interior

struct FieldTimeSeries{X, Y, Z, K, A, T, D, G, B, χ} <: AbstractDataField{X, Y, Z, A, G, T, 4}
                   data :: D
           architecture :: A
                   grid :: G
    boundary_conditions :: B
                  times :: χ

    function FieldTimeSeries{X, Y, Z, K}(data::D, arch::A, grid::G, bcs::B, times::χ) where {X, Y, Z, K, D, A, G, B, χ}
        T = eltype(grid)
        return new{X, Y, Z, K, A, T, D, G, B, χ}(data, arch, grid, bcs, times)
    end
end

#####
##### Constructors
#####

"""
    FieldTimeSeries{LX, LY, LZ}([architecture = CPU()], grid, times, boundary_conditions=nothing)

Return `FieldTimeSeries` at location `(LX, LY, LZ)`, on `grid`, at `times`, with
`boundary_conditions`, and initialized with zeros of `eltype(grid)`.
"""
function FieldTimeSeries{LX, LY, LZ}(architecture, grid, times, boundary_conditions=nothing) where {LX, LY, LZ}
    location = (LX, LY, LZ)
    Nt = length(times)
    data_size = total_size(location, grid)
    raw_data = zeros(architecture, grid, data_size..., Nt)
    data = offset_data(raw_data, grid, location)
    return FieldTimeSeries{LX, LY, LZ, InMemory}(data, architecture, grid, boundary_conditions, times)
end

# CPU() default
FieldTimeSeries{LX, LY, LZ}(grid::AbstractGrid, times, bcs=nothing) where {LX, LY, LZ} =
    FieldTimeSeries{LX, LY, LZ}(CPU(), grid, times, bcs)

"""
    FieldTimeSeries(path, name;
                    architecture = CPU(),
                    backend = InMemory(),
                    grid = nothing,
                    iterations = nothing,
                    times = nothing)

Returns a `FieldTimeSeries` for the field `name` describing a field's time history from a JLD2 file
located at `path`.

Keyword arguments
=================

- `architecture`: The architecture on which to store time series data. `CPU()` by default.

- `backend`: `InMemory()` to load data into a 4D array or `OnDisk()` to lazily load data from disk
             when indexing into `FieldTimeSeries`.

- `grid`: A grid to associated with data, in the case that the native grid
          was not serialized properly.

- `iterations`: Iterations to load. Defaults to all iterations found in the file.

- `times`: Save times to load, as determined through an approximate floating point
           comparison to recorded save times. Defaults to times associated with `iterations`.
           Takes precedence over `iterations` if `times` is specified.
"""
FieldTimeSeries(path, name; architecture=CPU(), backend=InMemory(), kwargs...) =
    FieldTimeSeries(path, name, architecture, backend; kwargs...)

#####
##### InMemory time serieses
#####

const InMemoryFieldTimeSeries{X, Y, Z} = FieldTimeSeries{X, Y, Z, InMemory}

struct UnspecifiedBoundaryConditions end

function FieldTimeSeries(path, name, architecture, backend::InMemory;
                         grid = nothing,
                         location = nothing,
                         boundary_conditions = UnspecifiedBoundaryConditions(),
                         iterations = nothing,
                         times = nothing)

    file = jldopen(path)

    # Non-defaults
    isnothing(grid)       && (grid       =  file["serialized/grid"])
    isnothing(iterations) && (iterations =  parse.(Int, keys(file["timeseries/t"])))
    isnothing(times)      && (times      =  [file["timeseries/t/$i"] for i in iterations])
    isnothing(location)   && (location   =  file["timeseries/$name/serialized/location"])

    if boundary_conditions isa UnspecifiedBoundaryConditions
        boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"]
    end

    LX, LY, LZ = location

    time_series = FieldTimeSeries{LX, LY, LZ}(architecture, grid, times, boundary_conditions)

    set!(time_series, path, name)

    return time_series
end

Base.getindex(fts::InMemoryFieldTimeSeries{LX, LY, LZ}, n::Int) where {LX, LY, LZ} =
    Field(LX, LY, LZ, fts.architecture, fts.grid, fts.boundary_conditions, view(fts.data, :, :, :, n))

#####
##### set!
#####

"""
    Field(path::String, name::String, iter; architecture=GPU(), grid=nothing)

Load a Field saved in JLD2 file at `path`, with `name` and at `iter`ation.
`architecture = CPU()` by default, and `grid` is loaded from `path` if not specified.
"""
function Field(location, path::String, name::String, iter;
               architecture = CPU(),
               grid = nothing,
               boundary_conditions = nothing)

    file = jldopen(path)

    raw_data = arch_array(architecture, file["timeseries/$name/$iter"])
    isnothing(grid) && (grid = file["serialized/grid"])

    close(file)

    try
        data = offset_data(raw_data, grid, location)
        return Field(location, architecture, grid, boundary_conditions, data)
    catch
        field = Field(location, architecture, grid, boundary_conditions)
        interior(field) .= raw_data
        return field
    end
end

function set!(time_series::InMemoryFieldTimeSeries, path::String, name::String)

    file = jldopen(path)
    file_iterations = parse.(Int, keys(file["timeseries/t"]))
    file_times = [file["timeseries/t/$i"] for i in file_iterations]
    close(file)

    for (n, time) in enumerate(time_series.times)
        file_index = findfirst(t -> t ≈ time, file_times)
        file_iter = file_iterations[file_index]

        field_n = Field(location(time_series), path, name, file_iter,
                        boundary_conditions = time_series.boundary_conditions,
                        grid = time_series.grid)

        set!(time_series[n], field_n)
    end

    return nothing
end

function set!(time_series::FieldTimeSeries, fields_vector::AbstractVector{<:AbstractField})
    raw_data = parent(time_series.data)
    ArrayType = array_type(time_series.architecture)

    file = jldopen(path)

    for (n, field) in enumerate(fields_vector)
        raw_data[:, :, :, n] .= parent(field)
    end

    close(file)

    return nothing
end

# TODO: this is a bit of type-piracy (with respect to the Oceananigans.Fields module)...
# is there a better way?

# FieldTimeSeries[i] returns ViewField
const ViewField = Field{<:Any, <:Any, <:Any, <:Any, <:SubArray}

using OffsetArrays: IdOffsetRange

parent_indices(idx::Int) = idx
parent_indices(idx::Base.Slice{<:IdOffsetRange}) = Colon()

# Is this too surprising?
Base.parent(vf::ViewField) = view(parent(parent(vf.data)), parent_indices.(vf.data.indices)...)

"Returns a view of `f` that excludes halo points."
@inline interior(f::FieldTimeSeries{X, Y, Z}) where {X, Y, Z} =
    view(parent(f.data),
         interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx),
         interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy),
         interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz),
         :)

#####
##### OnDisk time serieses
#####

struct OnDiskData
    path :: String
    name :: String
end

function FieldTimeSeries(path, name, architecture, backend::OnDisk; grid=nothing)
    file = jldopen(path)

    if isnothing(grid)
        grid = file["serialized/grid"]
    end

    iterations = parse.(Int, keys(file["timeseries/t"]))
    times = [file["timeseries/t/$i"] for i in iterations]

    data = OnDiskData(path, name)
    LX, LY, LZ = file["timeseries/$name/serialized/location"]
    bcs = file["timeseries/$name/serialized/boundary_conditions"]

    close(file)

    return FieldTimeSeries{LX, LY, LZ, OnDisk}(data, architecture, grid, bcs, times)
end

# For creating an empty `FieldTimeSeries`.
function FieldTimeSeries(grid, location, times; architecture=CPU(), ArrayType=array_type(architecture), name="", filepath="", bcs=nothing)
    LX, LY, LZ = location

    Nt = length(times)
    data_size = total_size(location, grid)

    raw_data = zeros(data_size..., Nt) |> ArrayType
    data = offset_data(raw_data, grid, location)

    return FieldTimeSeries{LX, LY, LZ}(InMemory(), data, architecture, grid, bcs, times, name, filepath, 4)
end

#####
##### Methods
#####

# Include the time dimension.
@inline Base.size(fts::FieldTimeSeries) = (size(location(fts), fts.grid)..., length(fts.times))

@propagate_inbounds Base.getindex(f::FieldTimeSeries{X, Y, Z, InMemory}, i, j, k, n) where {X, Y, Z} = f.data[i, j, k, n]

function Base.getindex(fts::FieldTimeSeries{X, Y, Z, OnDisk}, n::Int) where {X, Y, Z}
    # Load data
    file = jldopen(fts.data.path)
    iter = keys(file["timeseries/t"])[n]
    raw_data = file["timeseries/$(fts.data.name)/$iter"] |> array_type(fts.architecture)
    close(file)

    # Wrap Field
    loc = (X, Y, Z)
    field_data = offset_data(raw_data, fts.grid, loc)

    return Field(loc..., fts.architecture, fts.grid, fts.boundary_conditions, field_data)
end

Base.setindex!(fts::FieldTimeSeries, val, inds...) = Base.setindex!(fts.data, val, inds...)

#####
##### Show methods
#####

backend_str(::InMemory) = "InMemory"
backend_str(::OnDisk) = "OnDisk"

#####
##### show
#####

short_show(fts::FieldTimeSeries{X, Y, Z, K}) where {X, Y, Z, K} =
    string("$(join(size(fts), "×")) FieldTimeSeries{$(backend_str(K()))} located at $(show_location(fts))")

Base.show(io::IO, fts::FieldTimeSeries{X, Y, Z, K, A}) where {X, Y, Z, K, A} =
    print(io, "$(short_show(fts))\n",
          "├── architecture: $A\n",
          "└── grid: $(short_show(fts.grid))")


# todo change
__precompile__(false)

module Bayesian_EMR_MSM

using CmdStan
using LinearAlgebra
using Distributions
using Random
using DataFrames
using CSV

include("timeseries/timeseries.jl")
include("estimations/estimations.jl")
include("predictions/predictions.jl")

end # module


function trailingmeasure!(F::TF, crsp::DataFrame, calendarinterval::DatePeriod,
  minfraction::Float64;
  Ftarget::Symbol = error("Ftarget is required"),
  Fmeasure::Symbol = error("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing,)::Nothing where TF

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{String} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (error("$Fnew already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || error("dataframe not sorted by [:permno, :date]")

  #find the minimum date in the entire dataset
  alldates = unique(crsp.date) |> sort!
  initialminpoints = round(
    searchsortedfirst(alldates, alldates[1] + calendarinterval)*minfraction,RoundUp) |> Int

  #now compute the min number of points for each date
  function getminpoints(i)
    (alldates[i] ≤ alldates[1] + calendarinterval) && return initialminpoints
    lowerdateindex = searchsortedfirst(alldates, alldates[i] + Day(1) - calendarinterval)
    return round((i - lowerdateindex + 1) * minfraction,RoundUp) |> Int
  end
  minpointindex = Dict{Date, Int}(
    alldates[i] => getminpoints(i) for i ∈ initialminpoints:length(alldates))
  minstartdate = alldates[initialminpoints]
  @assert (alldates[alldates .≥ minstartdate] .|> (d)->haskey(minpointindex, d)) |> all

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #when we are looking at eom values only
  @inline getj₀(d::Date, dates::AbstractVector{Date}) = searchsortedfirst(
    dates, d + Day(1) - calendarinterval)

  #the below is a low-effort performance booster- don't check arrays below this
  absminpoints::Int = minimum(values(minpointindex) |> collect)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = size(scrsp, 1)
    (absminpoints < nrows) || continue #if there are enough  entries to compute the trailing returns

    #a hack to only compute calcualtions at end of month
    for j ∈ absminpoints:nrows
      applyconditional && (!scrsp[j,Fconditional]) && continue #if we are only computing end of month values
      currentdate = scrsp[j, :date]
      (currentdate < minstartdate) && continue
      j₀::Int = getj₀(currentdate, scrsp.date)
      N::Int = j-j₀ +1
      #if the returns are not stale and complete enough, record them
      (N < minpointindex[currentdate]) && continue
      scrsp[j, Fmeasure] = F(scrsp[j₀:j , Ftarget])
    end
  end

  return nothing
end

#applies a function to a vector of trailing variables
#this version uses a rolling calendar interval with a set number
#of minimum points
function trailingmeasure!(F::Function, crsp::DataFrame, calendarinterval::DatePeriod,
  minpoints::Int;
  Ftarget::Symbol = error("Ftarget is required"),
  Fmeasure::Symbol = error("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing)::Nothing

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{String} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (error("$Fmeasure already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || error("dataframe not sorted by [:permno, :date]")

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #when we are looking at eom values only
  @inline getj₀(d::Date, dates::AbstractVector{Date}) = searchsortedfirst(
    dates, d + Day(1) - calendarinterval)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = size(scrsp, 1)
    (minpoints ≤ nrows) || continue #if there are enough  entries to compute the trailing returns

    #a hack to only compute calcualtions at end of month
    for j ∈ minpoints:nrows
      applyconditional && (!scrsp[j,Fconditional]) && continue #if we are only computing end of month values
      j₀::Int = getj₀(scrsp[j, :date], scrsp.date)
      N::Int = j-j₀ +1
      #if the returns are not stale and complete enough, record them
      (N < minpoints) && continue
      scrsp[j, Fmeasure] = F(scrsp[j₀:j , Ftarget])
    end
  end

  return nothing
end

#trailing index measure
function stricttrailingmeasure!(F::TF, crsp::DataFrame, calendarinterval::DatePeriod;
  Ftarget::Symbol = error("Ftarget is required"),
  Fmeasure::Symbol = error("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing,)::Nothing where TF

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{String} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (error("$Fnew already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || error("dataframe not sorted by [:permno, :date]")

  #find the lagged dates in the entire dataset
  alldates = unique(crsp.date) |> sort!
  alldatesidx = Dict{Date, Int}(alldates[i] => i for i ∈ 1:length(alldates))
  initialminpoints = searchsortedfirst(alldates, alldates[1] + calendarinterval) |> Int

  #now compute the lagged dates
  getlaggedpoint(i) = alldates[searchsortedfirst(alldates, alldates[i] - calendarinterval)]

  laggeddateidx = Dict{Date, Date}(
    alldates[i] => getlaggedpoint(i) for i ∈ initialminpoints:length(alldates))
  minstartdate = alldates[initialminpoints]
  @assert (alldates[alldates .≥ minstartdate] .|> (d)->haskey(laggeddateidx, d)) |> all

  #low effort check to only look at rows after the minimum days across all intervals
  minintervaldays = minimum(((k,v)->alldatesidx[k]-alldatesidx[v]
    ).(keys(laggeddateidx), values(laggeddateidx)))

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = length(scrsp.date)

    #if there aren't enough dates within a permno to create lagged dates, skip the permno
    (firstdate, lastdate) = (scrsp[1,:date], scrsp[end,:date])
    if ((nrows < minintervaldays) ||
        (lastdate < minstartdate) ||
        (laggeddateidx[lastdate] < firstdate))

      continue
    end

    #a hack to only compute calcualtions at end of month
    for j ∈ minintervaldays:length(scrsp.date)
      #if we are only computing end of month values, make sure we are at a month end
      applyconditional && (!scrsp[j,Fconditional]) && continue
      currentdate = scrsp[j, :date]

       #skip if currentdate is before the first entry in the lookup table
      currentdate < minstartdate && continue

      #now check if the lagged date is a valid index entry for this permno
      targetlaggeddate = laggeddateidx[currentdate]
      j₀ = searchsortedfirst(scrsp.date, targetlaggeddate)
      actuallaggeddate = scrsp[j₀, :date]
      (actuallaggeddate ≠ targetlaggeddate) && continue
      scrsp[j, Fmeasure] = F(scrsp[j₀:j , Ftarget])
    end
  end

  return nothing
end

#applies a function to a vector of trailing variables
#this version uses a rolling window computed by the number of entries as opposed to
#one computed by the dates
function trailingmeasure!(F::TF, crsp::DataFrame, windowsize::Int;
  Ftarget::Symbol = throw("Ftarget is required"),
  Fmeasure::Symbol = throw("Fmeasurename is required"),
  Fconditional::Union{Nothing, Symbol} = nothing,
  maxcalendarinterval::DatePeriod = throw("maxcalendarinterval is required")
  )::Nothing where TF

  applyconditional::Bool = (!isnothing(Fconditional))
  crspnames::Vector{Symbol} = names(crsp)

  #check and make sure the target doesn't already have data in it (or somethign is probably wrong)
  (sum((ismissing).(crsp[!,Fmeasure])) == size(crsp,1)) ||  (throw("$Fnew already contains data"))

  #make sure everything is sorted
  (issorted(crsp, [:permno, :date])) || throw("dataframe not sorted by [:permno, :date]")

  #split and apply by permno
  requiredfields::Vector{Symbol} = [Ftarget; applyconditional ? Fconditional : []]
  completecrsp::SubDataFrame = view(crsp, completecases(crsp, requiredfields), :)
  scrsps::GroupedDataFrame = groupby(completecrsp, :permno)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    nrows = size(scrsp, 1)
    (nrows<windowsize) || continue #if there are enough  entries to compute the trailing returns

    for j ∈ windowsize:nrows
      applyconditional && (!scrsp[j,Fconditional]) && continue #if we are only computing end of month values
      staledate::Date = scrsp[j, :date] - maxcalendarinterval
      j₀::Int = j-windowsize+1

      (scrsp[j₀, :date] ≤ staledate)  && continue #elliminates stale windows
      scrsp[j, Fmeasurename] = F(scrsp[j₀:j, Ftarget])
    end
  end

  return nothing
end


module PreliminarySelection
    using GLMNet
    using ..ModelSelection

    export lasso!, lasso, lassoselection, PRELIMINARYSELECTION_EXTRAKEY, _lasso!, _lasso

    include("const.jl")
    include("strings.jl")
    include("utils.jl")
    include("core.jl")
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    PointwiseLearn(model)

A learning solver that converts geospatial data to a tabular format
with features (and possibly labels) for each point, and then solves
the problem with classical statistical learning `model`.

## Parameters

* `model` - Learning model implementing the `MLJModelInterface.jl`.

## References

* Hoffimann et al. 2020. [Geostatistical Learning: Challenges and Opportunities]
  (https://arxiv.org/abs/2102.08791)
"""
struct PointwiseLearn{M} <: LearningSolver
  model::M
end

function solve(problem::LearningProblem, solver::PointwiseLearn)
  sdata = sourcedata(problem)
  tdata = targetdata(problem)
  ptask = task(problem)
  model = solver.model

  # assert model is compatible with task
  @assert iscompatible(model, ptask) "$model is not compatible with $ptask"

  # learn model on source data
  lmodel = learn(ptask, sdata, model)

  # apply model to target data
  perform(ptask, tdata, lmodel)
end

# -----------
# IO METHODS
# -----------

function Base.show(io::IO, ::PointwiseLearn)
  print(io, "PointwiseLearn")
end

function Base.show(io::IO, ::MIME"text/plain", solver::PointwiseLearn)
  println(io, solver)
  print(io, "  └─model ⇨ ")
  show(IOContext(io, :compact => true), solver.model)
end


###############################################################################
#
#   Groups.jl : Nemo groups
#
###############################################################################

include("Rings.jl")

################################################################################
#
#  Element types for instances of groups
#
################################################################################

elem_type(::T) where {T <: Group} = elem_type(T)


using CausalityTools
using TransferEntropy, PerronFrobenius, StateSpaceReconstruction, Simplices
using CrossMappings
using TimeseriesSurrogates
using Documenter, DocumenterMarkdown
using PyCall, Conda
Conda.add("scipy")
using Plots
using DynamicalSystems
using Distributions
using StaticArrays
using Statistics
using StatsBase

ENV["GKSwstype"] = "100"

PAGES = [
    "index.md",
    "Discretization" => "discretization.md",

    #"StateSpaceReconstruction" => [
        #"Delay embeddings" => "embed.md",
        #"Rectangular partitioning" => "partitioning_rectangular.md",
        #"Triangulation partitioning" => "partitioning_triangulation.md"
    #],

    "PerronFrobenius" => [
        "Estimator documentation" => [
            "Transfer operator grid docs" => "transferoperator/transferoperator_grid_docs.md",
        #    "Transfer operator triangulation docs" => "transferoperator/transferoperator_triang_docs.md"
        ],
        #"Examples" => [
        #    "Transfer operator grid example" => "transferoperator/transferoperator_grid_example.md",
        #    "Transfer operator triangulation" => "transferoperator/transferoperator_triang_example.md"
        #],

        "Invariant measure" => "invariantmeasure/invariantmeasure_docs.md",#[
            #"Documentation" => ,
            #"Examples" => "invariantmeasure/invariantmeasure_example.md"
        #]
    ],

    "Causality algorithms" => [
        "TransferEntropy" => [
            #"Common interface" => "transferentropy/commoninterface_TE.md",
            "kNN estimator" => "transferentropy/estimator_TE_kNN.md",
            "Transfer operator grid estimator" => "transferentropy/estimator_TE_transferoperator_grid.md",
            "Visitation frequency estimator" => "transferentropy/estimator_TE_visitfreq.md",
            "Wrapper" => "transferentropy/wrapper_TE.md"
        ],
        "Cross mappings" => [
            "Convergent cross mapping" => "crossmappings/crossmapping.md"
        ]
    ],

    "Surrogate data" => [
        "The method of surrogate data" => "surrogates/surrogates_overview.md",
        "Docs" => [
            "randomshuffle" => "surrogates/randomshuffle_docs.md",
            "randomamplitudes" => "surrogates/randomamplitudes_docs.md",
            "randomphases" => "surrogates/randomphases_docs.md",
            "aaft" => "surrogates/aaft_docs.md",
            "iaaft" => "surrogates/iaaft_docs.md"
        ],
        "Examples" => [
            "randomshuffle" => "surrogates/randomshuffle_example.md",
            "randomamplitudes" => "surrogates/randomamplitudes_example.md",
            "randomphases" => "surrogates/randomphases_example.md",
            "aaft" => "surrogates/aaft_example.md",
            "iaaft" => "surrogates/iaaft_example.md"
        ]
    ],

    #"Workflow" => "workflow.md",

    #"Interop with DynamicalSystems.jl" => [
    #    "Example 1" => "interop_dynamicalsystems_infoflow1.md"
    #],

    #"Examples" => [
    #    "TE partitioning schemes"=> "transferentropy/examples_TE_different_partitionings.md"
    #],

    "Example systems" => [
		"ar1" => "example_systems/ar1.md",
        "henon2" => "example_systems/henon2.md",
        "logistic2" => "example_systems/logistic2.md",
        "logistic3" => "example_systems/logistic3.md",
		"linearmap3d_nonlinearcoupling" => "example_systems/linearmap3d_nonlinearcoupling.md",
		"verdes" => "example_systems/verdes.md"
    ]
]

makedocs(
    sitename = "CausalityTools.jl documentation",
    modules = [CausalityTools,
                TransferEntropy,
                PerronFrobenius,
                StateSpaceReconstruction,
                TimeseriesSurrogates,
                CrossMappings],
    format = :markdown,
    pages = PAGES
)

if !Sys.iswindows()
    deploydocs(
        deps   = Deps.pip("mkdocs==0.17.5", "mkdocs-material==2.9.4",
        "python-markdown-math", "pygments", "pymdown-extensions"),
        repo   = "github.com/kahaaga/CausalityTools.jl.git",
        target = "site",
        make = () -> run(`mkdocs build`)
    )
end


"""
    Consistent(predictions::AbstractVector)

Create an object that can be used for resampling `predictions` and corresponding targets.

The set of predictions has to be provided as a vector consisting of
- probabilities,
- vectors of probabilities summing up to 1,
- or distributions that can be sampled from.

Consistency resampling can be performed with `rand` and `rand!`, as described in the
documentation of [`Random`](https://docs.julialang.org/en/v1/stdlib/Random/).

# Examples

```jldoctest
julia> predictions = rand(100);

julia> consistent = Consistent(predictions);

julia> rand(consistent) isa Tuple{Float64,Bool}
true

julia> first(rand(consistent)) in predictions
true
```

!!! note
    One can use `Random.Sampler` to create a sampler that is optimized for the
    generation of multiple samples, e.g., by building an alias table:
    ```jldoctest
    julia> using Random: Sampler, GLOBAL_RNG

    julia> sampler = Sampler(GLOBAL_RNG, Consistent(rand(100)));

    julia> rand(sampler) isa Tuple{Float64,Bool}
    true
    ```
    However, for small number of samples the setup cost can outweigh the
    improved sampling efficiency.

!!! note
    Multiple samples are returned as an array of tuples of predictions and corresponding
    targets. If you prefer a tuple of arrays of predictions and targets instead, you can use
    a package such as
    [StructArrays](https://github.com/JuliaArrays/StructArrays.jl) and perform in-place
    sampling:
    ```jldoctest
    julia> using StructArrays, Random

    julia> predictions = Vector{Float64}(undef, 100);

    julia> targets = Vector{Bool}(undef, 100);

    julia> rand!(StructVector((predictions, targets)), Consistent(rand(100)));
    ```

# References

Bröcker, J. and Smith, L.A., 2007. [Increasing the reliability of reliability diagrams](https://doi.org/10.1175/WAF993.1). Weather and forecasting, 22(3), pp. 651-661.
"""
struct Consistent{T,S,P<:AbstractVector{T}}
    predictions::P

    function Consistent{T,S}(predictions::AbstractVector{T}) where {T,S}
        return new{T,S,typeof(predictions)}(predictions)
    end
end

function Consistent(predictions::AbstractVector{<:Real})
    isempty(predictions) && error("at least one prediction is required")
    isprobabilities = all(predictions) do p
        zero(p) <= p <= one(p)
    end
    isprobabilities || error("predictions have to be probabilities")
    return Consistent{eltype(predictions),Bool}(predictions)
end

function Consistent(predictions::AbstractVector{<:AbstractVector{<:Real}})
    isempty(predictions) && error("at least one prediction is required")
    isnormalized = all(predictions) do p
        sum_p = sum(p)
        return sum_p ≈ one(sum_p)
    end
    isnormalized || error("predictions have to be normalized")
    return Consistent{eltype(predictions),Int}(predictions)
end

# fallback: sample one target to compute the type of targets
# `eltype` does not work reliably for determining e.g. the type of
# `rand(d::Distributions.Distribution)`
function Consistent(predictions::AbstractVector)
    isempty(predictions) && error("at least one prediction is required")
    S = typeof(rand(first(predictions)))
    return Consistent{eltype(predictions),S}(predictions)
end

Base.eltype(::Type{<:Consistent{T,S}}) where {T,S} = Tuple{T,S}


"""
    search_cv(estimator, X, y, splits, params; [verbosity = 1])

Exhaustive search over the specified sets of parameter values for the `estimator` with features
data `X` and label `y`. The iterable `splits` provides vectors of indices for the training dataset.
The remaining indices are used to create the validation dataset.

Return an array with a tuple for each set of parameters value, where the first entry is a set of
parameter values and the second entry the cross-validation outcome of those values. This outcome is
a dictionary with an entry for the validation dataset and, if the parameter `is_training_metric` is
set in the `estimator`, an entry for the training dataset. Each entry of the dictionary is
another dictionary with an entry for each validation metric in the `estimator`. Each of these
entries is an array that holds the validation metric's value for each dataset, at the last valid
iteration.

# Arguments
* `estimator::LGBMEstimator`: the estimator to be fit.
* `X::Matrix{TX<:Real}`: the features data.
* `y::Vector{Ty<:Real}`: the labels.
* `splits`: the iterable providing arrays of indices for the training dataset.
* `params`: the iterable providing dictionaries of pairs of parameters (Symbols) and values to
    configure the `estimator` with.
* `verbosity::Integer`: keyword argument that controls LightGBM's verbosity. `< 0` for fatal logs
    only, `0` includes warning logs, `1` includes info logs, and `> 1` includes debug logs.
"""
function search_cv(estimator::LGBMEstimator, X::Matrix{TX}, y::Vector{Ty},
                                      splits, params; verbosity::Integer = 1) where {TX<:Real,Ty<:Real}
    n_params = length(params)
    results = Array{Tuple{Dict{Symbol,Any},Dict{String,Dict{String,Vector{Float64}}}}}(undef,n_params)
    for (search_idx, search_params) in enumerate(params)
        log_info(verbosity, "\nSearch: ", search_idx, "\n", search_params, "\n")

        search_estimator = deepcopy(estimator)
        foreach(param -> setfield!(search_estimator, param[1], param[2]), search_params)
        search_results = cv(search_estimator, X, y, deepcopy(splits),
                            verbosity = ifelse(verbosity == 1, 0, verbosity))

        results[search_idx] = (search_params, search_results)
        cv_logsummary(search_results, verbosity)
    end
    log_info(verbosity, "\nSearch finished\n")

    return results
end


using .CUDA

@inline function evaluate_dist(f, ref, frame, px, py, qx, qy, rₚx, rₚy)
    val = Distances.eval_start(f, ref, frame)
    @inbounds for ox in -rₚx:rₚx, oy in -rₚy:rₚy
        val = Distances.eval_reduce(f, val, Distances.eval_op(f, frame[px+ox, py+oy], ref[qx+ox, qy+oy]))
    end
    return Distances.eval_end(f, val)
end

function best_match_fullsearch_kernel!(R, ref, frame, f, rₚ, Δₛ, rₛ)
    M, N = size(frame)
    rₚx, rₚy = rₚ.I
    Δₛx, Δₛy = Δₛ.I
    rₛx, rₛy = rₛ.I

    px_start = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    py_start = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    Δx = gridDim().x * blockDim().x
    Δy = gridDim().y * blockDim().y
    px_end = M-rₚx
    py_end = N-rₚy

    if px_start <= rₚx || py_start <= rₚy
        return nothing
    end

    # TODO:
    # we can further share the patch ref[px_start-rₚx:px_end+rₚx, py_start-rₚy:py_end+rₚy]
    # using shared memory so as to reduce global memory access.

    # NOTE:
    # For better performance, loop along the row order. This could
    # produce results that different to CPU versions on constant fields.
    @inbounds for px in px_start:Δx:px_end, py in py_start:Δy:py_end
        qx_range = max(rₚx+1, px - rₛx):Δₛx:min(px_end, px + rₛx)
        qy_range = max(rₚy+1, py - rₛy):Δₛy:min(py_end, py + rₛy)
        min_val, min_pos_x, min_pos_y = eltype(frame)(Inf), first(qx_range), first(qy_range)
        @inbounds for qx in qx_range, qy in qy_range
            val = evaluate_dist(f, ref, frame, px, py, qx, qy, rₚx, rₚy)
            if val < min_val
                min_val = val
                min_pos_x = qx
                min_pos_y = qy
            end
        end
        R[px, py] = CartesianIndex(min_pos_x, min_pos_y)
    end
    return nothing
end

# modified from https://github.com/JuliaStats/Distances.jl/blob/988c92b8b2b6d8a28e1b8aea336f572025ada2f2/src/metrics.jl#L198
# some metrics are removed because they are not supported yet: RogersTanimoto
DistancesMetrics = Union{Euclidean,SqEuclidean,PeriodicEuclidean,Chebyshev,Cityblock,TotalVariation,Minkowski,Hamming,Jaccard,CosineDist,ChiSqDist,KLDivergence,RenyiDivergence,BrayCurtis,JSDivergence,SpanNormDist,GenKLDivergence}

function best_match(
        S::FullSearch{F},
        ref::CuArray{<:Any,2},
        frame::CuArray{<:Any,2}; offset=false) where F<:DistancesMetrics
    rₚ, Δₛ, rₛ = S.patch_radius, S.search_stride, S.search_radius
    size_check(ref, frame)

    # CuArray{<:RGB} is not supported
    T = Float32
    ref = CUDA.adapt(CuArray{T, 2}, ref)
    frame = CUDA.adapt(CuArray{T, 2}, frame)

    threads = (16, 16)
    blocks = ceil.(Int, size(ref)./threads)

    R_frame = CartesianIndices(size(frame))
    R_frame = first(R_frame)+rₚ:last(R_frame)-rₚ

    cu_matches = OffsetArray(CuArray(fill(CartesianIndex(0, 0), size(R_frame))), rₚ.I)
    CUDA.@sync begin
        @cuda threads=threads blocks=blocks best_match_fullsearch_kernel!(cu_matches, ref, frame, S.f, rₚ, Δₛ, rₛ)
    end

    if offset
        m = parent(cu_matches)
        @. m = m - R_frame
    end
    return cu_matches
end

function multi_match_fullsearch_kernel!(R, ref, frame, f, rₚ, Δₛ, rₛ, K, sizeof_vals)
    T = eltype(frame)
    M, N = size(frame)
    rₚx, rₚy = rₚ.I
    Δₛx, Δₛy = Δₛ.I
    rₛx, rₛy = rₛ.I

    px_start = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    py_start = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    Δx = gridDim().x * blockDim().x
    Δy = gridDim().y * blockDim().y
    px_end = M-rₚx
    py_end = N-rₚy

    if px_start <= rₚx || py_start <= rₚy
        return nothing
    end

    threads = (blockDim().x, blockDim().y)
    vals = @cuDynamicSharedMem T (threads..., K)
    inds = @cuDynamicSharedMem Int16 (threads..., K) offset=sizeof_vals

    idx = (threadIdx().x, threadIdx().y)
    vals = view(vals, idx..., :)
    inds = view(inds, idx..., :)

    # NOTE:
    # For better performance, loop along the row order. This could
    # produce results that different to CPU versions on constant fields.
    @inbounds for px in px_start:Δx:px_end, py in py_start:Δy:py_end
        qx_range = max(rₚx+1, px - rₛx):Δₛx:min(px_end, px + rₛx)
        qy_range = max(rₚy+1, py - rₛy):Δₛy:min(py_end, py + rₛy)
        Rq = CartesianIndices((qx_range, qy_range))

        # Partial sort the smallest K values `vals` and their associated indices `inds`
        # using 1-pass insertion sort with binary search.
        # This is significantly faster than the max-heap sort provided in `src/utils.jl`
        fill!(vals, typemax(eltype(vals)))
        fill!(inds, 0)
        for i = 1:length(Rq)
            qx, qy = Rq[i].I
            x = evaluate_dist(f, ref, frame, px, py, qx, qy, rₚx, rₚy)
            # 1-pass insertion sort
            idx = _insertion_index(x, vals)
            idx == 0 && continue
            # bump the right-hand side and then insert
            for j in K:-1:idx+1
                vals[j] = vals[j-1]
                inds[j] = inds[j-1]
            end
            vals[idx] = x
            inds[idx] = i
        end

        @inbounds for i in 1:K
            R[px, py, i] = Rq[inds[i]]
        end
    end
    return nothing
end

function _insertion_index(x, vals)
    # find the first index such that x < vals[k]
    l, r = 1, length(vals)
    i = (l+r)÷2
    i_old = 1
    @inbounds while i_old != i
        if x < vals[i]
            r = i
        else
            l = i
        end
        i_old = i
        i = (l+r)÷2
    end
    return @inbounds x < vals[l] ? l : x < vals[r] ? r : 0
end

function multi_match(
        S::FullSearch{F},
        ref::CuArray{<:Any, 2},
        frame::CuArray{<:Any, 2}; num_patches, offset=false) where F<:DistancesMetrics

    rₚ, Δₛ, rₛ = S.patch_radius, S.search_stride, S.search_radius
    K = num_patches
    size_check(ref, frame)

    # CuArray{<:RGB} is not supported
    T = Float16
    Tind = Int16
    ref = CUDA.adapt(CuArray{T, 2}, ref)
    frame = CUDA.adapt(CuArray{T, 2}, frame)

    # FIXME: this currently doesn't work with large num_patches
    threads = (8, 4)
    blocks = ceil.(Int, size(frame)./threads)

    function get_shmem(threads, K)
        N = prod(threads)
        sizeof_vals = N * K * sizeof(T)
        sizeof_inds = N * K * sizeof(Tind)
        return sizeof_vals, sizeof_inds
    end

    shmem_sizes = get_shmem(threads, K)
    shmem = sum(shmem_sizes)

    R_frame = CartesianIndices(size(frame))
    R_frame = first(R_frame)+rₚ:last(R_frame)-rₚ
    cu_matches = OffsetArray(CUDA.fill(CartesianIndex(0, 0), size(R_frame)..., K), (rₚ.I..., 0));

    CUDA.@sync begin
        @cuda threads=threads blocks=blocks shmem=shmem multi_match_fullsearch_kernel!(cu_matches, ref, frame, S.f, rₚ, Δₛ, rₛ, K, shmem_sizes[1])
    end

    if offset
        m = parent(cu_matches)
        @. m = m - R_frame
    end
    return cu_matches
end


x = asindex(1:2:10, :x)
y = asindex(1:5, :y)

@testset "vcat" begin
    @testset "vcat_keys" begin
        for (x, y, xykeys, xyvals, xyname, xyindex) in ()
            @test vcat_keys(x, y) == xykeys
            @test vcat_keys(keys(x), y) == xykeys
            @test vcat_keys(x, keys(x)) == xykeys
            @test vcat_keys(keys(x), keys(y)) == xykeys

            @test vcat_values(x, y) == xyvals
            @test vcat_values(values(x), y) == xyvals
            @test vcat_values(x, values(x)) == xyvals
            @test vcat_values(values(x), values(y)) == xyvals

            @test vcat_names(x, y) == xyname
            @test vcat_names(dimnames(x), y) == xyname
            @test vcat_names(x, dimnames(x)) == xyname
            @test vcat_names(dimnames(x), dimnames(y)) == xyname

            @test vcat(x, y) == xyname
        end
    end
end


#!/usr/bin/env julia

using Weber
using Lazy: @> # see https://github.com/MikeInnes/Lazy.jl

version = v"0.0.1"
sid,skip = @read_args("Gap Detection ($version).")

#===============================================================================
Experiment Settings
===============================================================================#
# 2AFC gap detection. Each stimulus consists of a marker---a band-pass noise---
# and masker---a notch noise of the same width. Each stimulus either does or
# does not have a gap in the middle. The task is to indicate which of 2 stimuli
# has this gap.

const atten_dB = 30 # adjust to calibrate sound levels
const n_trials_per_block = 60
const n_blocks = 6

const marker_center_freq = 200Hz
const marker_width_octaves = 1
const marker_SNR = 20

const gap_ramp_ms = 0.5ms
const feedback_delay = 700ms

const visual_delay = 300ms

const SOA = 900ms # "onset" asynchrony of the no-gap and gap interval
const trial_spacing = 500ms # how long to wait at the start and end of trials
const max_delta = 400ms
const first_delta = 50ms
const big_step = 5ms
const little_step = 1ms

const adapter = levitt_adapter(first_delta=first_delta/max_delta,down=3,up=1,
                               big=big_step / first_delta,
                               little=little_step / first_delta,
                               min_delta=0,max_delta=1)

experiment = Experiment(
  skip=skip,
  moment_resolution=5ms,
  columns = [
    :sid => sid,
    :condition => "example",
    :version => version,
    :block
  ]
)

isresponse(e) = iskeydown(e,key"p") || iskeydown(e,key"q")

#==============================================================================
Stimulus and Trial Generation
===============================================================================#

low_freq = marker_center_freq*2^(-marker_width_octaves/2)
high_freq = marker_center_freq*2^(marker_width_octaves/2)

# marker: the band-pass noise which will have a gap (or not)
marker() = @> begin
  noise(trial_spacing+SOA+trial_spacing)
  bandpass(low_freq,high_freq)
  ramp
  attenuate(atten_dB)
end

# masker: an infinite stream of noise above and below the marker
masker() = @> begin
  noise(trial_spacing+SOA+trial_spacing)
  bandstop(low_freq,high_freq)
  ramp
  attenuate(atten_dB+marker_SNR)
end

# gap_noise: generates the gap
function gap_noise(gap_first,adapter)
  @assert max_delta*delta(adapter) > 2gap_ramp_ms || delta(adapter) == 0
  if delta(adapter) == 0 # no gap
    mix(masker(),marker)
  elseif gap_first
    @> begin
      envelope(1,trial_spacing)
      fadeto(silence(max_delta*delta(adapter)),gap_ramp_ms)
      fadeto(envelope(1,SOA+trial_spacing-max_delta*delta(adapter)),gap_ramp_ms)
      mult(marker())
      mix(masker())
    end
  else # gap second
    @> begin
      envelope(1,trial_spacing+SOA)
      fadeto(silence(max_delta*delta(adapter)),gap_ramp_ms)
      fadeto(envelope(1,trial_spacing-max_delta*delta(adapter)))
      mult(marker())
      mix(masker())
    end
  end
end

# creates a single, 2-interval forced-choice trial

function gap2AFC(adapter;keys...)
  gap_first = rand(Bool)
  resp = response(adapter,key"q" => "gap_first",key"p" => "gap_second",
                  correct=(gap_first? "gap_first" : "gap_second");
                  keys...)
  stim() = gap_noise(gap_first,adapter)

  [moment(feedback_delay,play,stim),
   show_cross(),
   moment(trial_spacing,display,"Interval 1"),
   moment(SOA,display,"Interval 2"),
   moment(trial_spacing+visual_delay,display,
          "Was there a gap in the first [Q] or second interval [P]?"),
   resp,await_response(isresponse)]
end

#===============================================================================
Experiment Creation
===============================================================================#

setup(experiment) do
  # play a test tone, to verify sound levels
  addbreak(moment(record,"start"),
           moment(250ms,play,@> tone(1kHz,1s) ramp attenuate(atten_dB)),
           moment(1s))

  addbreak(instruct("""

      On each trial, you will hear a gap in a noise at one of two
  possible time intervals. Indicate which of the two intervals had a gap. Hit
  'Q' if the first interval had a gap, and 'P' if the second had a gap.

  """),moment(display,colorant"gray"))

  for block in 1:n_blocks
    @addtrials let adapter = adapter
      for trial in 1:n_trials_per_block
        addtrial(gap2AFC(adapter,block=block))
      end

      # define this string during experiment setup
      # when we know what block we're on...
      heading = "Block $block of $n_blocks"

      function threshold_report()
        mean,sd = estimate(adapter)
        thresh = usconvert(ms,round(mean,3)*max_delta)
        thresh_sd = uconvert(ms,round(sd,3)*max_delta)

        # define this string during run time when we know
        # what the threshold estimate is.
        "$heading: Threshold $thresh (SD: $thresh_sd)\n"*
        "Hit spacebar to continue..."
      end

      addbreak(moment(display,threshold_report,clean_whitespace=false),
               await_response(iskeydown(key":space:")))
    end
  end
end

run(experiment)


# using SuperResolution

mutable struct EvalSR
    _selectIDX::Array{Tuple{Int64,Int64},1}
    _modus::Int64
    wModel::Array{SuperResolution.ModelSR,1}
    RefModel::Array{SuperResolution.ModelSR,1}
    distKL::Array{SuperResolution.DistanceSR,1}
    wkldSR::Array{SuperResolution.clusWKLD,1}
    gapSR::Array{SuperResolution.clusGAP,1}

    max_mean_idx::Array{Int64,1}
    max_median_idx::Array{Int64,1}
    max_max_idx::Array{Int64,1}
    max_minDelta_idx::Array{Int64,1}

    diff_WK_max_idx::Array{Int64,1}
    diff_GAP_max_idx::Array{Int64,1}

    N_observations::Array{Int64,1}
    K_ground_truth::Array{Int64,1}

    K_KL_estimated_by_mean::Array{Int64,1}
    K_KL_estimated_by_median::Array{Int64,1}
    K_KL_estimated_by_max::Array{Int64,1}
    K_KL_estimated_by_minDelta::Array{Int64,1}
    K_KL_estimated_by_95quantile::Array{Int64,1}

    K_GAP_estimated::Array{Int64,1}

    corrDetecRate_KL_by_mean::Float64
    corrDetecRate_KL_by_median::Float64
    corrDetecRate_KL_by_max::Float64
    corrDetecRate_KL_by_minDelta::Float64
    corrDetecRate_KL_by_95quantile::Float64

    corrDetecRate_GAP::Float64

    VALID_DF::DataFrame
    error_DF::DataFrame

    locID_DF::Array{DataFrame,1}
    evo_DF::Array{DataFrame,1}

    function EvalSR(SEG::SuperResolution.SegmentSR,RESTR_SEG::Array{Tuple{Int64,Int64},1},modINPUT=5)#RESTR_SEG::Int64)
        obj=new()
        obj._selectIDX=RESTR_SEG
        obj._modus=modINPUT
        print("*** modINPUT = $(modINPUT) ***\n")
        NUM_SEG=length(RESTR_SEG)
        #obj._selectIDX=RESTR_SEG[1:RESTR_SEG]
        obj.wModel=Array{SuperResolution.ModelSR,1}(undef,NUM_SEG)
        obj.RefModel=Array{SuperResolution.ModelSR,1}(undef,NUM_SEG)
        obj.distKL=Array{SuperResolution.DistanceSR,1}(undef,NUM_SEG)
        obj.wkldSR=Array{SuperResolution.clusWKLD,1}(undef,NUM_SEG)
        obj.gapSR=Array{SuperResolution.clusGAP,1}(undef,NUM_SEG)

        obj.diff_WK_max_idx=Array{Int64,1}(undef,NUM_SEG)
        obj.diff_GAP_max_idx=Array{Int64,1}(undef,NUM_SEG)

        obj.K_ground_truth=Array{Int64,1}(undef,NUM_SEG)
        obj.N_observations=Array{Int64,1}(undef,NUM_SEG)

        select_mode=[   :K_KL_estimated_by_mean,                    #   modINPUT=1
                        :K_KL_estimated_by_median,                  #   modINPUT=2
                        :K_KL_estimated_by_max,                     #   modINPUT=3
                        :K_KL_estimated_by_minDelta,                #   modINPUT=4
                        :K_KL_estimated_by_95quantile]              #   modINPUT=5

        @time(for i=1:NUM_SEG
                print("\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n")
                print(">>> Starting Segment $(i)/$(NUM_SEG) with index $(SEG.idx[i])\n")
                print("\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n")
                print("(wModel:)\n")
                obj.wModel[i]=SuperResolution.ModelSR(SEG,RESTR_SEG[i],1)     #for segment create all models  WKLD
                # print("(RefModel:)\n")
                #obj.RefModel[i]=SuperResolution.ModelSR(SEG,RESTR_SEG[i],2)   #for segment create all models KLD
                # obj.RefModel[i]=obj.wModel[i]   #temporary workaround to save computation time

                obj.wModel[i].wkld[1:(end-1)].=obj.wModel[i].wkld[2:end]
                obj.wModel[i].weight[1:(end-1)].=obj.wModel[i].weight[2:end]
                obj.wModel[i].geom[1:(end-1)].=obj.wModel[i].geom[2:end]

                obj.wkldSR[i]=SuperResolution.clusWKLD(obj.wModel[i],1)    #cluster WKLD
                # obj.gapSR[i]=SuperResolution.clusGAP(obj.wModel[i],obj.RefModel[i],1)  #Gap statistic
        end)
        obj.K_ground_truth=map(i->SEG.K[CartesianIndex(RESTR_SEG[i])],1:NUM_SEG)
        obj.N_observations=map(i->SEG.N[CartesianIndex(RESTR_SEG[i])],1:NUM_SEG)

        obj.max_mean_idx=map(i->obj.wkldSR[i].max_mean_idx[1],1:length(obj.wkldSR))
        obj.max_median_idx=map(i->obj.wkldSR[i].max_median_idx[1],1:length(obj.wkldSR))
        obj.max_max_idx=map(i->obj.wkldSR[i].max_max_idx[1],1:length(obj.wkldSR))
        obj.max_minDelta_idx=map(i->obj.wkldSR[i].max_minDelta_idx[1],1:length(obj.wkldSR))
        # obj.diff_WK_max_idx=map(i->obj.gapSR[i].diff_WK_max_idx[1],1:length(obj.gapSR))

        #obj.K_KL_estimated_by_mean=map(i->obj.wModel[i].K[CartesianIndex(RESTR_SEG[i])],1:NUM_SEG)
        obj.K_KL_estimated_by_mean=map(i->size(obj.wModel[i].wkld[obj.max_mean_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_median=map(i->size(obj.wModel[i].wkld[obj.max_median_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_max=map(i->size(obj.wModel[i].wkld[obj.max_max_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_minDelta=map(i->size(obj.wModel[i].wkld[obj.max_minDelta_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_95quantile=map(i->length( obj.wkldSR[i].wkldMinDelta[ .!(obj.wkldSR[i].wkldMinDelta .< quantile(obj.wkldSR[i].wkldMinDelta, 0.95))] )+1,1:length(obj.wkldSR))  # +1 to account for indexing

        # obj.K_KL_estimated_by_95quantile=map(i->length( obj.wkldSR[i].wkldMin[ .!(obj.wkldSR[i].wkldMin .< quantile(obj.wkldSR[i].wkldMin, 0.95))] )+1,1:length(obj.wkldSR))  # +2 to account for indexing and +1 to account for floor quantile to closest existing wlkd value lower than the quantile
        # obj.K_GAP_estimated=map(i->size(obj.wModel[i].wkld[obj.diff_WK_max_idx[i]],1),1:NUM_SEG)
        #obj.corrDetecRate_GAP=sum(abs.(obj.K_ground_truth-obj.diff_WK_max_idx))/sum(obj.K_ground_truth)
        obj.corrDetecRate_KL_by_mean=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_mean))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_median=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_median))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_max=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_max))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_minDelta=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_minDelta))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_95quantile=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_95quantile))./obj.K_ground_truth))
        # obj.corrDetecRate_GAP=mean(exp.((-abs.(obj.K_ground_truth-obj.K_GAP_estimated))./obj.K_ground_truth))
        obj.VALID_DF=DataFrame( idx=obj._selectIDX,
                                N_observations=obj.N_observations,
                                K_ground_truth=obj.K_ground_truth,
                                K_KL_estimated_by_mean=obj.K_KL_estimated_by_mean,
                                K_KL_estimated_by_median=obj.K_KL_estimated_by_median,
                                K_KL_estimated_by_max=obj.K_KL_estimated_by_max,
                                K_KL_estimated_by_minDelta=obj.K_KL_estimated_by_minDelta,
                                K_KL_estimated_by_95quantile=obj.K_KL_estimated_by_95quantile,
                                corrDetecRate_KL_by_mean=obj.corrDetecRate_KL_by_mean,
                                corrDetecRate_KL_by_median=obj.corrDetecRate_KL_by_median,
                                corrDetecRate_KL_by_max=obj.corrDetecRate_KL_by_max,
                                corrDetecRate_KL_by_minDelta=obj.corrDetecRate_KL_by_minDelta,
                                corrDetecRate_KL_by_95quantile=obj.corrDetecRate_KL_by_95quantile,
                                dev_by_mean=(obj.K_ground_truth-obj.K_KL_estimated_by_mean),
                                dev_by_median=(obj.K_ground_truth-obj.K_KL_estimated_by_median),
                                dev_by_max=(obj.K_ground_truth-obj.K_KL_estimated_by_max),
                                dev_by_minDelta=(obj.K_ground_truth-obj.K_KL_estimated_by_minDelta),
                                dev_by_95quantile=(obj.K_ground_truth-obj.K_KL_estimated_by_95quantile) )

        print("\n<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n")
        print("<<< K_ground_truth $(obj.K_ground_truth)\n")
        print("<<< N_observations $(obj.N_observations)\n")
        print("------------------------------------------------\n")
        # print("<<< K_KL_estimated_by_mean $(obj.K_KL_estimated_by_mean)\n")
        # print("<<< KL(mean) correct detection rate $(obj.corrDetecRate_KL_by_mean)\n")
        # print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_mean)\n")
        # print("------------------------------------------------\n")
        print("<<< K_KL_estimated_by_median $(obj.K_KL_estimated_by_median)\n")
        print("<<< KL(median) correct detection rate $(obj.corrDetecRate_KL_by_median)\n")
        print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_median)\n")
        print("------------------------------------------------\n")
        # print("<<< K_KL_estimated_by_max $(obj.K_KL_estimated_by_max)\n")
        # print("<<< KL(max) correct detection rate $(obj.corrDetecRate_KL_by_max)\n")
        # print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_max)\n")
        # print("------------------------------------------------\n")
        print("<<< K_KL_estimated_by_minDelta $(obj.K_KL_estimated_by_minDelta)\n")
        print("<<< KL(minDelta) correct detection rate $(obj.corrDetecRate_KL_by_minDelta)\n")
        print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_minDelta)\n")
        print("------------------------------------------------\n")
        print("<<< K_KL_estimated_by_95quantile $(obj.K_KL_estimated_by_95quantile)\n")
        print("<<< KL(95quantile) correct detection rate $(obj.corrDetecRate_KL_by_95quantile)\n")
        print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_95quantile)\n")
        print("------------------------------------------------\n")
        # print("<<< K_GAP_estimated $(obj.K_GAP_estimated)\n")
        # print("<<< GAP correct detection rate $(obj.corrDetecRate_GAP)\n")
        # print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_GAP_estimated)\n")
        # print("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n")
        obj.locID_DF=map(k->DataFrame(  mol_ID=obj.wModel[k].mol_ID,
                                        est_ID=obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1],
                                        x=obj.wModel[k].loc[1][:,1],
                                        y=obj.wModel[k].loc[1][:,2],
                                        # xC=vcat(map(i->repeat([obj.wModel[k].loc[end-obj.K_KL_estimated_by_95quantile[k]+1][i,:][1]],sum(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1].==unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])[i])),1:length(unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])))...),
                                        # yC=vcat(map(i->repeat([obj.wModel[k].loc[end-obj.K_KL_estimated_by_95quantile[k]+1][i,:][2]],sum(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1].==unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])[i])),1:length(unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])))...),
                                        x_clus_ID=vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...),
                                        y_clus_ID=vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...),
                                        x_mol_ID=vcat(map(i->mean(obj.wModel[k].loc[1][:,1][obj.wModel[k].mol_ID.==i]),obj.wModel[k].mol_ID)...),
                                        y_mol_ID=vcat(map(i->mean(obj.wModel[k].loc[1][:,2][obj.wModel[k].mol_ID.==i]),obj.wModel[k].mol_ID)...),
                                        intra_dist_clus_ID=map(j->pairwise(Euclidean(), transpose(hcat( obj.wModel[k].loc[1][:,1],obj.wModel[k].loc[1][:,2])),
                                                                                transpose(hcat( vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...),
                                                                                vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...))),
                                                        dims=2)[j,j],1:length(obj.wModel[k].loc[1][:,1])),
                                        ),1:length(obj._selectIDX))

        obj.evo_DF=map(k->DataFrame(    K=map(i->length(unique(obj.wModel[k].member[i])),1:length(obj.wModel[k].member)),
                                        median_intra_dist_clus_ID=map(l->
                                                                    median(map(j->
                                                                        pairwise(Euclidean(), transpose(hcat( obj.wModel[k].loc[1][:,1],obj.wModel[k].loc[1][:,2])),
                                                                                transpose(hcat( vcat(map(i->obj.wModel[k].loc[l][:,1][(1:length(obj.wModel[k].loc[l][:,1]))[obj.wModel[k].clus_ID[l].==i]],obj.wModel[k].member[l])...),
                                                                                vcat(map(i->obj.wModel[k].loc[l][:,2][(1:length(obj.wModel[k].loc[l][:,2]))[obj.wModel[k].clus_ID[l].==i]],obj.wModel[k].member[l])...))),
                                                                                dims=2)[j,j],
                                                                            1:length(obj.wModel[k].loc[1][:,1]))),
                                                                    1:length(obj.wModel[k].loc))
                                        ),1:length(obj._selectIDX))

        obj.error_DF=DataFrame(     idx=obj._selectIDX,
                                    clus_vs_mol_MSE=map(k->(1/length(obj.locID_DF[k].x_clus_ID)) * sum( (obj.locID_DF[k].x_clus_ID-obj.locID_DF[k].x_mol_ID).^2 + (obj.locID_DF[k].y_clus_ID-obj.locID_DF[k].y_mol_ID).^2 ),1:length(obj.locID_DF)),
                                    clus_vs_mol_MAE=map(k->(1/length(obj.locID_DF[k].x_clus_ID)) * sum(sqrt.( (obj.locID_DF[k].x_clus_ID-obj.locID_DF[k].x_mol_ID).^2 + (obj.locID_DF[k].y_clus_ID-obj.locID_DF[k].y_mol_ID).^2 )),1:length(obj.locID_DF))
        )
        print("------------------------------------------------\n")
        print("<<< mean clus_vs_mol_MSE $(mean(obj.error_DF.clus_vs_mol_MSE))\n")
        print("<<< median clus_vs_mol_MSE $(median(obj.error_DF.clus_vs_mol_MSE))\n")
        print("<<< sample standard deviation $(Statistics.std(obj.error_DF.clus_vs_mol_MSE))\n")
        print("------------------------------------------------\n")
        return obj
    end
end
export EvalSR

# map(i-> EVAL.wModel[1].loc[end-EVAL.K_KL_estimated_by_95quantile[1]+1][:,1][collect(1:length(EVAL.wModel[1].loc[end-EVAL.K_KL_estimated_by_95quantile[1]+1][:,1]))[EVAL.wModel[1].clus_ID[end-EVAL.K_KL_estimated_by_95quantile[1]+1].==i]],EVAL.wModel[1].member[end-EVAL.K_KL_estimated_by_95quantile[1]+1])


using Flux: @treelike

using ..Basic
using ..Basic: onehot

export Gpt, lmloss
export load_gpt_pretrain

struct Gpt
    pe::PositionEmbedding
    ts::Chain
    drop::Dropout
end

@treelike Gpt

"""
    Gpt(size::Int, head::Int, ps::Int, layer::Int;
        max_len::Int = 512, trainable = true, act = gelu, pdrop = 0.1)
    Gpt(size::Int, head::Int, hs::Int, ps::Int, layer::Int;
        max_len::Int = 512, trainable = true, act = gelu, pdrop = 0.1)

the Generative Pretrain model.
"""
function Gpt(size::Int, head::Int, ps::Int, layer::Int;
             max_len::Int = 512, trainable = true, act = gelu, pdrop = 0.1)
    rem(size, head) != 0 && error("size not divisible by head")
    Gpt(size, head, div(size, head), ps, layer; max_len=max_len, trainable=trainable, act=act, pdrop=pdrop)
end

function Gpt(size::Int, head::Int, hs::Int, ps::Int, layer::Int;
             max_len::Int = 512, trainable = true, act = gelu, pdrop = 0.1)
    Gpt(PositionEmbedding(size, max_len; trainable=trainable),
        Chain([Transformer(size, head, hs, ps; future=false, act=act, pdrop=pdrop) for i = 1:layer]...),
        Dropout(pdrop))
end

function (gpt::Gpt)(x::T, mask=nothing)::T where T
    pe = gpt.pe(x)
    e = x .+ pe
    e = gpt.drop(e)::T
    t = gpt.ts(e)::T
    t = mask === nothing ? t : t .* mask
    t #size(t) == (size, seq_len, batch)
end

"""
    lmloss(embed, onehot, encoding, mask)

compute the language modeling loss for Gpt, onehot is the onehot array of the origin
input sentence. encoding the output of Gpt, mask is the mask between input sentences.

"""
lmloss(embed::Embed, o::OneHotArray, t::AbstractArray{T}, mask) where T = lmloss(embed, tofloat(T, o), t, mask)
function lmloss(embed::Embed, et, t::AbstractArray{T, N}, mask) where {T,N}
    if N == 3
        t = t[:, 1:end-1, :]
        s = size(t)
        sim = logsoftmax(transpose(embed.embedding) * reshape(t, s[1], :)) #(vocab, seq_len*batch)
        sim = reshape(sim, :, s[2], s[3])
        return logcrossentropy(et[:, 2:end, :], sim, mask[:, 2:end, :])
    elseif N == 2
        t = t[:, 1:end-1]
        sim = logsoftmax(transpose(embed.embedding) * t)
        return logcrossentropy(et[:, 2:end], sim, mask[:, 2:end])
    end
end

function lmloss(gpt::Gpt, embed::Embed, x)
    e, mask = embed(x)
    t = gpt(e, mask)
    lmloss(embed, onehot(embed, x), t, mask)
end

function Base.show(io::IO, gpt::Gpt)
    hs = div(size(gpt.ts[1].mh.iqproj.W)[1], gpt.ts[1].mh.head)
    h, ps = size(gpt.ts[1].pw.dout.W)

    print(io, "Gpt(")
    print(io, "layers=$(length(gpt.ts.layers)), ")
    print(io, "head=$(gpt.ts[1].mh.head), ")
    print(io, "head_size=$(hs), ")
    print(io, "pwffn_size=$(ps), ")
    print(io, "size=$(h))")
end


function make_number(data, start, nbytes)
    sum = 0
    for ii = 1:nbytes
        sum *= 256
        sum += data[start+ii]
    end
    return sum
end

function get_xref_record(data, start, w)
    v = Vector{Int}()
    for tw in w
        n = make_number(data, start, tw)
        start += tw
        push!(v, n)
    end
    return v
end

createIfObjectStream(stm::CosStream) =
    get(stm, cn"Type") === cn"ObjStm" ? CosObjectStream(stm) : stm

function read_xref_stream(xrefstm::IDD{CosStream},
                          xref::Dict{CosIndirectObjectRef, CosObjectLoc})
    @assert get(xrefstm, cn"Type") === cn"XRef"
    size = get(xrefstm, cn"Size")
    @assert size !== CosNull

    w = get(xrefstm, cn"W")
    @assert w != CosNull
    @assert length(w) == 3

    index = get(xrefstm, cn"Index")

    index === CosNull &&
        (index = CosArray(CosObject[CosInt(0), size]))

    # The xref stream may be accessed later.
    # There is no point encrypting this data
    # Ideal will be to remove the filter.
    cosStreamRemoveFilters(xrefstm)

    input = get(xrefstm)
    data = read(input)
    util_close(input)
    datasize = length(data)

    w_n = get(w,true) #This size is 3
    recsize = sum(w_n)

    lenidx = length(index)
    @assert rem(lenidx,2) == 0
    idx_int=get(index,true)

    it = 0 #iterator for data
    for i = 1:div(lenidx,2)
        for j = 0:idx_int[2*i] - 1
            oid = idx_int[2*i-1]+j
            record = get_xref_record(data, it, w_n)
            @assert length(record) == 3
            @assert record[1] in 0:2

            loc = (record[1] == 1) ? record[2] :
                  (record[1] == 2) ? record[3] : 0
            stm = (record[1] == 2) ? CosIndirectObjectRef(record[2],0) : CosNull
            ref = (record[1] == 1) ? CosIndirectObjectRef(oid, record[3]) :
                  (record[1] == 2) ? CosIndirectObjectRef(oid, 0) :
                                     CosIndirectObjectRef(0,0)
            it += recsize
            record[1] != 0 && !haskey(xref,ref) && (xref[ref] = CosObjectLoc(loc,stm))
            it >= datasize && break
        end
        it >= datasize && break
    end
    return xref
end

function read_object_info_from_stm(stm::CosStream,
                                   oids::Vector{Int},
                                   oloc::Vector{Int},
                                   n::Int, first::Int)
    filename = get(stm, cn"F")
    io = util_open(String(filename),"r")
    try
        for i = 1:n
            oids[i] = convert(Int, parse_unsignednumber(io))
            oloc[i] = convert(Int, parse_unsignednumber(io)) + first
        end
    finally
        util_close(io)
    end
end

cosObjectStreamGetObject(stm::CosIndirectObject{CosObjectStream},
                         ref::CosIndirectObjectRef, loc::Int) =
                             cosObjectStreamGetObject(stm.obj, ref, loc)

function cosObjectStreamGetObject(stm::CosObjectStream,
                                  ref::CosIndirectObjectRef,
                                  loc::Int)
    objtuple = get(ref)
    stm.oids[loc+1] != objtuple[1] && return CosNull
    dirobj = cosObjectStreamGetObject(stm, CosNull, loc)
    return CosIndirectObject(objtuple[1], objtuple[2], dirobj)
end

function cosObjectStreamGetObject(stm::CosObjectStream,
                                  ref::CosNullType,
                                  loc::Int)
    filename = get(stm, cn"F")
    io = util_open(String(filename),"r")
    ps = io
    try
        seek(ps, stm.oloc[loc+1])
        obj = parse_value(ps)
        return obj
    finally
        util_close(ps)
    end
end


"""
    cudnnRNNForward(w, x; hiddenSize, o...)
    cudnnRNNForward!(y, w, x; hiddenSize, o...)
    cudnnRNNForward(w, x, d::cudnnRNNDescriptor; o...)
    cudnnRNNForward!(y, w, x, d::cudnnRNNDescriptor; o...)

Apply the RNN specified with weights `w` and configuration given by `d` or keyword options
to input `x`.

Keyword arguments for hidden input/output:
* `hx=nothing`: initialize the hidden vector if specified (by default initialized to 0).
* `cx=nothing`: initialize the cell vector (only in LSTMs) if specified (by default initialized to 0).
* `hy=nothing`: return the final hidden vector in `hy` if set to `Ref{Any}()`.
* `cy=nothing`: return the final cell vector in `cy` (only in LSTMs) if set to `Ref{Any}()`.

Keyword arguments specifying the RNN when `d::cudnnRNNDescriptor` is not given:
* `hiddenSize::Integer`: hidden vector size, which must be supplied when `d` is not given
* `algo::cudnnRNNAlgo_t = CUDNN_RNN_ALGO_STANDARD`: RNN algo (`CUDNN_RNN_ALGO_STANDARD`, `CUDNN_RNN_ALGO_PERSIST_STATIC`, or `CUDNN_RNN_ALGO_PERSIST_DYNAMIC`).
* `cellMode::cudnnRNNMode_t = CUDNN_LSTM`: Specifies the RNN cell type in the entire model (`CUDNN_RNN_RELU`, `CUDNN_RNN_TANH`, `CUDNN_LSTM`, `CUDNN_GRU`).
* `biasMode::cudnnRNNBiasMode_t = CUDNN_RNN_DOUBLE_BIAS`: Sets the number of bias vectors (`CUDNN_RNN_NO_BIAS`, `CUDNN_RNN_SINGLE_INP_BIAS`, `CUDNN_RNN_SINGLE_REC_BIAS`, `CUDNN_RNN_DOUBLE_BIAS`). The two single bias settings are functionally the same for `RELU`, `TANH` and `LSTM` cell types. For differences in `GRU` cells, see the description of `CUDNN_GRU` in cudnn docs.
* `dirMode::cudnnDirectionMode_t = CUDNN_UNIDIRECTIONAL`: Specifies the recurrence pattern: `CUDNN_UNIDIRECTIONAL` or `CUDNN_BIDIRECTIONAL`. In bidirectional RNNs, the hidden states passed between physical layers are concatenations of forward and backward hidden states.
* `inputMode::cudnnRNNInputMode_t = CUDNN_LINEAR_INPUT`: Specifies how the input to the RNN model is processed by the first layer. When `inputMode` is `CUDNN_LINEAR_INPUT`, original input vectors of size `inputSize` are multiplied by the weight matrix to obtain vectors of `hiddenSize`. When `inputMode` is `CUDNN_SKIP_INPUT`, the original input vectors to the first layer are used as is without multiplying them by the weight matrix.
* `mathPrec::DataType = eltype(x)`: This parameter is used to control the compute math precision in the RNN model. For `Float16` input/output can be `Float16` or `Float32`, for `Float32` or `Float64` input/output, must match the input/output type.
* `mathType::cudnnMathType_t = math_mode()`: Sets the preferred option to use NVIDIA Tensor Cores accelerators on Volta (SM 7.0) or higher GPUs. When `dataType` is `CUDNN_DATA_HALF`, the `mathType` parameter can be `CUDNN_DEFAULT_MATH` or `CUDNN_TENSOR_OP_MATH`. The `ALLOW_CONVERSION` setting is treated the same `CUDNN_TENSOR_OP_MATH` for this data type. When `dataType` is `CUDNN_DATA_FLOAT`, the `mathType` parameter can be `CUDNN_DEFAULT_MATH` or `CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION`. When the latter settings are used, original weights and intermediate results will be down-converted to `CUDNN_DATA_HALF` before they are used in another recursive iteration. When dataType is `CUDNN_DATA_DOUBLE`, the `mathType` parameter can be `CUDNN_DEFAULT_MATH`.
* `inputSize::Integer = size(x,1)`: Size of the input vector in the RNN model. When the `inputMode=CUDNN_SKIP_INPUT`, the `inputSize` should match the `hiddenSize` value.
* `projSize::Integer = hiddenSize`: The size of the LSTM cell output after the recurrent projection. This value should not be larger than `hiddenSize`. It is legal to set `projSize` equal to `hiddenSize`, however, in this case, the recurrent projection feature is disabled. The recurrent projection is an additional matrix multiplication in the LSTM cell to project hidden state vectors `ht` into smaller vectors `rt = Wr * ht`, where `Wr` is a rectangular matrix with `projSize` rows and hiddenSize columns. When the recurrent projection is enabled, the output of the LSTM cell (both to the next layer and unrolled in-time) is `rt` instead of `ht`. The recurrent projection can be enabled for LSTM cells and `CUDNN_RNN_ALGO_STANDARD` only.
* `numLayers::Integer = 1`: Number of stacked, physical layers in the deep RNN model. When `dirMode= CUDNN_BIDIRECTIONAL`, the physical layer consists of two pseudo-layers corresponding to forward and backward directions.
* `dropout::Real = 0`: When non-zero, dropout operation will be applied between physical layers. A single layer network will have no dropout applied. Dropout is used in the training mode only.
* `auxFlags::Integer = CUDNN_RNN_PADDED_IO_ENABLED`: Miscellaneous switches that do not require additional numerical values to configure the corresponding feature. In future cuDNN releases, this parameter will be used to extend the RNN functionality without adding new API functions (applicable options should be bitwise OR-ed). Currently, this parameter is used to enable or disable padded input/output (`CUDNN_RNN_PADDED_IO_DISABLED`, `CUDNN_RNN_PADDED_IO_ENABLED`). When the padded I/O is enabled, layouts `CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED` and `CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED` are permitted in RNN data descriptors.

Other keyword arguments:
* `layout::cudnnRNNDataLayout_t = CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED`: The memory layout of the RNN data tensor. Options are `CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED`: Data layout is padded, with outer stride from one time-step to the next; `CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_PACKED`: The sequence length is sorted and packed as in the basic RNN API; `CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED`: Data layout is padded, with outer stride from one batch to the next.
* `seqLengthArray::Vector{Cint} = nothing`: An integer array with `batchSize` number of elements. Describes the length (number of time-steps) of each sequence. Each element in `seqLengthArray` must be greater than or equal to 0 but less than or equal to `maxSeqLength`. In the packed layout, the elements should be sorted in descending order, similar to the layout required by the non-extended RNN compute functions. The default value `nothing` assumes uniform `seqLength`s, no padding.
* `devSeqLengths::CuVector{Cint} = nothing`: Device copy of `seqLengthArray`
* `fwdMode::cudnnForwardMode_t = CUDNN_FWD_MODE_INFERENCE`: set to `CUDNN_FWD_MODE_TRAINING` when training
"""
cudnnRNNForward, cudnnRNNForward!

# Public methods
cudnnRNNForward(w, x; hiddenSize, o...)     = cudnnRNNForwardWithDefaults(w, x; hiddenSize, o...)
cudnnRNNForward!(y, w, x; hiddenSize, o...) = cudnnRNNForwardWithDefaults(w, x; y, hiddenSize, o...)
cudnnRNNForward(w, x, d::cudnnRNNDescriptor; o...)     = cudnnRNNForwardWithDefaults(w, x; rnnDesc=d, o...)
cudnnRNNForward!(y, w, x, d::cudnnRNNDescriptor; o...) = cudnnRNNForwardWithDefaults(w, x; y, rnnDesc=d, o...)

# Private method
function cudnnRNNForwardWithDefaults(
    w, x;

    # input hidden vectors
    hx = nothing,
    cx = nothing,

    # output buffers
    y = nothing,
    hy = nothing,
    cy = nothing,

    # rnnDescriptor parameters
    # TODO: look into GetClip, SetClip
    algo::cudnnRNNAlgo_t = CUDNN_RNN_ALGO_STANDARD,
    cellMode::cudnnRNNMode_t = CUDNN_LSTM,
    biasMode::cudnnRNNBiasMode_t = CUDNN_RNN_DOUBLE_BIAS,
    dirMode::cudnnDirectionMode_t = CUDNN_UNIDIRECTIONAL,
    inputMode::cudnnRNNInputMode_t = CUDNN_LINEAR_INPUT,
    dataType::DataType = eltype(x),
    mathPrec::DataType = dataType, # has to match dataType with one extra possibility dt=Float16 => mp=Float16|Float32
    mathType::cudnnMathType_t = math_mode(),
    inputSize::Integer = size(x,1),
    hiddenSize::Integer = 0,
    projSize::Integer = hiddenSize,
    numLayers::Integer = 1,
    dropout::Real = 0,
    auxFlags::Integer = CUDNN_RNN_PADDED_IO_ENABLED, # When the padded I/O is enabled, layouts CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED and CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED are permitted in RNN data descriptors.

    # rnnDescriptor
    rnnDesc::cudnnRNNDescriptor = cudnnRNNDescriptor(algo, cellMode, biasMode, dirMode, inputMode, cudnnDataType(dataType), cudnnDataType(mathPrec), mathType, Int32(inputSize), checkHidden(hiddenSize), Int32(projSize), Int32(numLayers), cudnnDropoutDescriptor(Cfloat(dropout)), UInt32(auxFlags)),

    # rnnData parameters:
    layout::cudnnRNNDataLayout_t = CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED, # padded [X,B,T] array
    seqLengthArray::Union{Nothing,Vector{Cint}} = nothing,         # assume no padding by default
    paddingFill::Ptr{Cvoid} = C_NULL,

    # rnnForward parameters
    fwdMode::cudnnForwardMode_t = CUDNN_FWD_MODE_INFERENCE, # set to CUDNN_FWD_MODE_TRAINING when training
    devSeqLengths::Union{Nothing,CuArray{Cint,1}} = nothing,
    reserveSpace::Union{CuArray,Nothing} = nothing,
    workspace::Union{CuArray,Nothing} = nothing,

    # gradient buffers: layer designers may want to preallocate, so leave them as kwargs
    dw = Ref{Any}(nothing),
    dx = Ref{Any}(nothing),
    dhx = Ref{Any}(nothing),
    dcx = Ref{Any}(nothing),
)
    # Verify all inputs: they should be compatible with rnnDesc (in case it is supplied), not necessarily with kwargs:
    rd = cudnnGetRNNDescriptor_v8(rnnDesc)
    @assert rd.hiddenSize > 0  "hiddenSize > 0 must be provided"
    @assert cudnnDataType(eltype(x)) == rd.dataType "Input x type not compatible with RNN"
    @assert size(x,1) == rd.inputSize               "Input x size not compatible with RNN"
    ydims = (rd.projSize << (rd.dirMode === CUDNN_BIDIRECTIONAL), size(x)[2:end]...)
    if y !== nothing
        @assert cudnnDataType(eltype(y)) == rd.dataType  "Output y type not compatible with RNN"
        @assert size(y) == ydims                         "Output y size not compatible with RNN or input x"
    else
        y = similar(x, ydims)
    end
    if layout === CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED
        if seqLengthArray === nothing; seqLengthArray = fill(Cint(size(x,3)), size(x,2)); end
        @assert size(x,2) == length(seqLengthArray)  "Input x batchsize not compatible with seqLengthArray"
        @assert size(x,3) >= maximum(seqLengthArray) "Input x seqLength not compatible with seqLengthArray"
        xDesc = cudnnRNNDataDescriptor(rd.dataType, layout, Cint(size(x,3)), Cint(size(x,2)), Cint(size(x,1)), seqLengthArray, paddingFill)
        yDesc = cudnnRNNDataDescriptor(rd.dataType, layout, Cint(size(y,3)), Cint(size(y,2)), Cint(size(y,1)), seqLengthArray, paddingFill)
    elseif layout === CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED
        if seqLengthArray === nothing; seqLengthArray = fill(Cint(size(x,2)), size(x,3)); end
        @assert size(x,3) == length(seqLengthArray)  "Input x batchsize not compatible with seqLengthArray"
        @assert size(x,2) >= maximum(seqLengthArray) "Input x seqLength not compatible with seqLengthArray"
        xDesc = cudnnRNNDataDescriptor(rd.dataType, layout, Cint(size(x,2)), Cint(size(x,3)), Cint(size(x,1)), seqLengthArray, paddingFill)
        yDesc = cudnnRNNDataDescriptor(rd.dataType, layout, Cint(size(y,2)), Cint(size(y,3)), Cint(size(y,1)), seqLengthArray, paddingFill)
    elseif layout === CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_PACKED
        if seqLengthArray === nothing; seqLengthArray = fill(Cint(size(x,3)), size(x,2)); end
        @assert length(x)÷size(x,1) == sum(seqLengthArray) "Input x size not compatible with seqLengthArray"
        xDesc = cudnnRNNDataDescriptor(rd.dataType, layout, maximum(seqLengthArray), Cint(length(seqLengthArray)), Cint(size(x,1)), seqLengthArray, paddingFill)
        yDesc = cudnnRNNDataDescriptor(rd.dataType, layout, maximum(seqLengthArray), Cint(length(seqLengthArray)), Cint(size(y,1)), seqLengthArray, paddingFill)
    else
        error("Unknown layout $layout")
    end
    if devSeqLengths === nothing; devSeqLengths = CuArray(seqLengthArray); end

    hdims = (rd.projSize, length(seqLengthArray), rd.numLayers << (rd.dirMode === CUDNN_BIDIRECTIONAL))
    cdims = (rd.hiddenSize, length(seqLengthArray), rd.numLayers << (rd.dirMode === CUDNN_BIDIRECTIONAL))
    if hx !== nothing
        @assert cudnnDataType(eltype(hx)) == rd.dataType "Hidden hx type not compatible with RNN"
        @assert size(hx) == hdims                        "Hidden hx size not compatible with RNN"
    end
    if hy !== nothing
        @assert hy isa Ref{Any}
        if isassigned(hy) && hy[] !== nothing
            @assert cudnnDataType(eltype(hy[])) == rd.dataType "Hidden hy type not compatible with RNN"
            @assert size(hy[]) == hdims                        "Hidden hy size not compatible with RNN"
        else
            hy[] = similar(y, hdims)
        end
    end
    if rd.cellMode === CUDNN_LSTM
        if cx !== nothing
            @assert cudnnDataType(eltype(cx)) == rd.dataType "Hidden cx type not compatible with RNN"
            @assert size(cx) == cdims                        "Hidden cx size not compatible with RNN"
        end
        if cy !== nothing
            @assert cy isa Ref{Any}
            if isassigned(cy) && cy[] !== nothing
                @assert cudnnDataType(eltype(cy[])) == rd.dataType "Hidden cy type not compatible with RNN"
                @assert size(cy[]) == cdims                        "Hidden cy size not compatible with RNN"
            else
                cy[] = similar(y, cdims)
            end
        end
    end
    hDesc = cudnnTensorDescriptor(CUDNN_TENSOR_NCHW, rd.dataType, Cint(3), Cint[reverse(hdims)...])
    cDesc = cudnnTensorDescriptor(CUDNN_TENSOR_NCHW, rd.dataType, Cint(3), Cint[reverse(cdims)...])

    weightSpaceSize = cudnnRNNWeightSpaceSize(rnnDesc)
    @assert sizeof(w) >= weightSpaceSize "RNN weights should be at least $weightSpaceSize bytes."

    # Backward called separately on each variable. We will calculate all gradients on first call. Use `dready` to avoid subsequent calls.
    dready = Ref{Bool}(false)   # this will be turned to `true` by the first backward call.

    y_h_c = cudnnRNNForwardAD(w, x, hx, cx; rnnDesc, fwdMode, devSeqLengths, xDesc, yDesc, y, hDesc, hy=(hy isa Ref ? hy[] : hy), cDesc, cy=(cy isa Ref ? cy[] : cy), workspace, reserveSpace, dw, dx, dhx, dcx, dready)
    if hy isa Ref; hy[] = y_h_c[2]; end
    if cy isa Ref && rd.cellMode === CUDNN_LSTM; cy[] = y_h_c[3]; end
    return y_h_c[1]             # only return y; hy and cy can be accessed through keyword arguments. They still need to be in AutoGrad return value to be included in gradient calc.
end

# AD method

function cudnnRNNForwardAD(w, x, hx, cx; rnnDesc, fwdMode, devSeqLengths, xDesc, yDesc, y, hDesc, hy, cDesc, cy, workspace, reserveSpace, dw, dx, dhx, dcx, dready)
    (workspaceSize, reserveSpaceSize) = cudnnRNNTempSpaceSizes(rnnDesc, fwdMode, xDesc)
    if reserveSpaceSize > 0 && reserveSpace === nothing; reserveSpace = cudnnTempSpace(reserveSpaceSize); end
    @assert sizeof(reserveSpace) >= reserveSpaceSize  "reserveSpace should be at least $reserveSpaceSize bytes"
    # Cannot use @workspace here because it is shared between forw and back calls
    if workspaceSize > 0 && workspace === nothing; workspace = cudnnTempSpace(workspaceSize); end
    @assert sizeof(workspace) >= workspaceSize  "workspace should be at least $workspaceSize bytes"
    cudnnRNNForward(handle(), rnnDesc, fwdMode, devSeqLengths, xDesc, x, yDesc, y, hDesc, something(hx, CU_NULL), something(hy, CU_NULL), cDesc, something(cx, CU_NULL), something(cy, CU_NULL), sizeof(w), w, sizeof(workspace), something(workspace, CU_NULL), sizeof(reserveSpace), something(reserveSpace, CU_NULL))
    return (y, hy, cy)
end

# Helper methods

function cudnnRNNWeightSpaceSize(rnnDesc::cudnnRNNDescriptor)
    ws = Csize_t[0]
    cudnnGetRNNWeightSpaceSize(handle(), rnnDesc, ws)
    ws[1]
end

function cudnnRNNTempSpaceSizes(rnnDesc::cudnnRNNDescriptor, fwdMode::cudnnForwardMode_t, xDesc::cudnnRNNDataDescriptor)
    ws = Csize_t[0]; rs = Csize_t[0]
    cudnnGetRNNTempSpaceSizes(handle(), rnnDesc, fwdMode, xDesc, ws, rs)
    ws[1], rs[1]
end

function cudnnGetRNNDescriptor_v8(rnnDesc::cudnnRNNDescriptor)
    (algo, cellMode, biasMode, dirMode, inputMode, dataType, mathPrec, mathType, inputSize, hiddenSize, projSize, numLayers, dropout, auxFlags) = (Ref{cudnnRNNAlgo_t}(), Ref{cudnnRNNMode_t}(), Ref{cudnnRNNBiasMode_t}(), Ref{cudnnDirectionMode_t}(), Ref{cudnnRNNInputMode_t}(), Ref{cudnnDataType_t}(), Ref{cudnnDataType_t}(), Ref{cudnnMathType_t}(), Ref{Int32}(), Ref{Int32}(), Ref{Int32}(), Ref{Int32}(), Ref{Ptr{Nothing}}(), Ref{UInt32}())
    cudnnGetRNNDescriptor_v8(rnnDesc, algo, cellMode, biasMode, dirMode, inputMode, dataType, mathPrec, mathType, inputSize, hiddenSize, projSize, numLayers, dropout, auxFlags)
    (algo, cellMode, biasMode, dirMode, inputMode, dataType, mathPrec, mathType, inputSize, hiddenSize, projSize, numLayers, dropout, auxFlags) = (algo[], cellMode[], biasMode[], dirMode[], inputMode[], dataType[], mathPrec[], mathType[], inputSize[], hiddenSize[], projSize[], numLayers[], dropout[], auxFlags[])
    (; rnnDesc, algo, cellMode, biasMode, dirMode, inputMode, dataType, mathPrec, mathType, inputSize, hiddenSize, projSize, numLayers, dropout, auxFlags)
end

checkHidden(h) = (h > 0 ? Int32(h) : error("hiddenSize > 0 is required"))

"""
    cudnnGetRNNWeightParams(w, d::cudnnRNNDescriptor)
    cudnnGetRNNWeightParams(w; hiddenSize, o...)

Return an array of weight matrices and bias vectors of an RNN specified by `d` or keyword
options as views into `w`. The keyword arguments and defaults in the second form are the
same as those in cudnnRNNForward specifying the RNN.

In the returned array `a[1,l,p]` and `a[2,l,p]` give the weight matrix and bias vector for
the l'th layer and p'th parameter or `nothing` if the specified matrix/vector does not
exist. Note that the matrices should be transposed for left multiplication, e.g. `a[1,l,p]'
* x`

The `l` index refers to the pseudo-layer number. In uni-directional RNNs, a pseudo-layer is
the same as a physical layer (pseudoLayer=1 is the RNN input layer, pseudoLayer=2 is the
first hidden layer). In bi-directional RNNs, there are twice as many pseudo-layers in
comparison to physical layers:

    pseudoLayer=1 refers to the forward direction sub-layer of the physical input layer
    pseudoLayer=2 refers to the backward direction sub-layer of the physical input layer
    pseudoLayer=3 is the forward direction sub-layer of the first hidden layer, and so on

The `p` index refers to the weight matrix or bias vector linear ID index.

If cellMode in rnnDesc was set to CUDNN_RNN_RELU or CUDNN_RNN_TANH:

    Value 1 references the weight matrix or bias vector used in conjunction with the input from the previous layer or input to the RNN model.
    Value 2 references the weight matrix or bias vector used in conjunction with the hidden state from the previous time step or the initial hidden state.

If cellMode in rnnDesc was set to CUDNN_LSTM:

    Values 1, 2, 3 and 4 reference weight matrices or bias vectors used in conjunction with the input from the previous layer or input to the RNN model.
    Values 5, 6, 7 and 8 reference weight matrices or bias vectors used in conjunction with the hidden state from the previous time step or the initial hidden state.
    Value 9 corresponds to the projection matrix, if enabled (there is no bias in this operation).

Values and their LSTM gates:

    Values 1 and 5 correspond to the input gate.
    Values 2 and 6 correspond to the forget gate.
    Values 3 and 7 correspond to the new cell state calculations with hyperbolic tangent.
    Values 4 and 8 correspond to the output gate.

If cellMode in rnnDesc was set to CUDNN_GRU:

    Values 1, 2 and 3 reference weight matrices or bias vectors used in conjunction with the input from the previous layer or input to the RNN model.
    Values 4, 5 and 6 reference weight matrices or bias vectors used in conjunction with the hidden state from the previous time step or the initial hidden state.

Values and their GRU gates:

    Values 1 and 4 correspond to the reset gate.
    Values 2 and 5 reference to the update gate.
    Values 3 and 6 correspond to the new hidden state calculations with hyperbolic tangent.

"""
function cudnnGetRNNWeightParams(
    w;
    hiddenSize::Integer,
    inputSize::Integer = hiddenSize,
    projSize::Integer = hiddenSize,
    algo::cudnnRNNAlgo_t = CUDNN_RNN_ALGO_STANDARD,
    cellMode::cudnnRNNMode_t = CUDNN_LSTM,
    biasMode::cudnnRNNBiasMode_t = CUDNN_RNN_DOUBLE_BIAS,
    dirMode::cudnnDirectionMode_t = CUDNN_UNIDIRECTIONAL,
    inputMode::cudnnRNNInputMode_t = CUDNN_LINEAR_INPUT,
    dataType::DataType = Float32,
    mathPrec::DataType = dataType,
    mathType::cudnnMathType_t = math_mode(),
    numLayers::Integer = 1,
    dropout::Real = 0,
    auxFlags::Integer = CUDNN_RNN_PADDED_IO_ENABLED,
)
    cudnnGetRNNWeightParams(w, cudnnRNNDescriptor(algo, cellMode, biasMode, dirMode, inputMode, cudnnDataType(dataType), cudnnDataType(mathPrec), mathType, Int32(inputSize), checkHidden(hiddenSize), Int32(projSize), Int32(numLayers), cudnnDropoutDescriptor(Cfloat(dropout)), UInt32(auxFlags)))
end

function cudnnGetRNNWeightParams(w, rnnDesc::cudnnRNNDescriptor)
    d = cudnnGetRNNDescriptor_v8(rnnDesc)
    T = juliaDataType(d.dataType)
    weightSpace = reinterpret(T, w)
    nlayers = d.numLayers << (d.dirMode === CUDNN_BIDIRECTIONAL)
    nparams = (d.cellMode === CUDNN_RNN_RELU || d.cellMode === CUDNN_RNN_TANH ? 2 :
               d.cellMode === CUDNN_LSTM ? 9 : d.cellMode === CUDNN_GRU  ? 6 :
               error("Unknown cellMode $(d.cellMode)"))
    a = Array{Any}(undef, 2, nlayers, nparams)
    p = Ref{Ptr{Cvoid}}(0)
    cudnnCreateTensorDescriptor(p); mDesc = cudnnTensorDescriptor(p[])
    cudnnCreateTensorDescriptor(p); bDesc = cudnnTensorDescriptor(p[])
    mAddr = Ref{CuPtr{Cvoid}}(0)
    bAddr = Ref{CuPtr{Cvoid}}(0)
    for l in 1:nlayers, p in 1:nparams
        cudnnGetRNNWeightParams(handle(), rnnDesc, l-1, sizeof(weightSpace), weightSpace, p-1, mDesc, mAddr, bDesc, bAddr)
        mT,mD,mS = cudnnGetTensorDescriptor(mDesc)
        bT,bD,bS = cudnnGetTensorDescriptor(bDesc)
        @assert mT === bT === T
        if mAddr[] === CU_NULL
            a[1,l,p] = nothing
        else
            m0 = (mAddr[] - pointer(weightSpace)) ÷ sizeof(T) |> Int
            a[1,l,p] = reshape(view(weightSpace, (m0+1):(m0+prod(mD))), (mD[1],mD[2]))
        end
        if bAddr[] === CU_NULL
            a[2,l,p] = nothing
        else
            b0 = (bAddr[] - pointer(weightSpace)) ÷ sizeof(T) |> Int
            a[2,l,p] = view(weightSpace, (b0+1):(b0+prod(bD)))
        end
    end
    cudnnDestroyTensorDescriptor.((mDesc,bDesc))
    return a
end


####
# NOT FUNCTIONAL
#TODO: Convert this to v2 model - This will not work for v2
####

import ..GPU

import ..YOLO: loadWeights!

export loadWeights!

"""
    flipkernel(x)

Flips given kernel
"""
flipkernel(x) = x[end:-1:1, end:-1:1, :, :]

"""
    updateconv!(c, gamma, mean, variance)

Applies batch-normalization to given convolutional layer
"""
function updateconv!(c, gamma, mean, variance)
    gamma4 = reshape(gamma, 1, 1, 1, :)
    variance4 = reshape(variance, 1, 1, 1, :)
    gamma3 = reshape(gamma, 1, 1, :, 1)
    mean3 = reshape(mean, 1, 1, :, 1)
    variance3 = reshape(variance, 1, 1, :, 1)
    #update c.w
    a = convert(Float32, 0.001)
    c.w = (c.w .* gamma4) ./ sqrt.(variance4 .+ a)
    #update c.b
    c.b = c.b .- (gamma3 .* mean3 ./ sqrt.(variance3 .+ a))
    return c.w, c.b
end

"""
    loadWeights!(model::YOLO.v2_tiny.v2Chain, weights_filepath::String)

Loads layers' weights from given weights file path.
"""
function loadWeights!(model::v2Chain, settings::Settings)
    outlength = settings.cell_bboxes * (5 + settings.num_classes)

    open(settings.weights_filepath, "r") do io
        readconstants!(io)
        #First Conv layer
        loadconv!(model.layers[1], io, 3, 3, 3, 16)
        #Second Conv layer
        loadconv!(model.layers[3], io, 3, 3, 16, 32)
        #Third Conv layer
        loadconv!(model.layers[5], io, 3, 3, 32, 64)
        #4th Conv layer
        loadconv!(model.layers[7], io, 3, 3, 64, 128)
        #5th Conv layer
        loadconv!(model.layers[9], io, 3, 3, 128, 256)
        #6th Conv layer
        loadconv!(model.layers[11], io, 3, 3, 256, 512)
        #YoloPad
        model.layers[13].w[1, 1, 1, 1] = 1
        #7th Conv layer
        loadconv!(model.layers[14], io, 3, 3, 512, 1024)
        #8th Conv layer
        loadconv!(model.layers[15], io, 3, 3, 1024, 1024)
        #last layer
        read!(io, model.layers[16].b)
        toRead = Array{Float32}(undef, 1024 * outlength)
        read!(io, toRead)
        toRead = reshape(toRead, 1, 1, 1024, outlength)
        model.layers[16].w = permutedims(toRead, [2, 1, 3, 4])
        model.layers[16].w = flipkernel(model.layers[16].w)

        if GPU >= 0
            model.layers[16].w = KnetArray(model.layers[16].w)
            model.layers[16].b = KnetArray(model.layers[16].b)
            model.layers[13].w = KnetArray(model.layers[13].w)
        end
    end
end

"""
    loadconv!(c, io, d1, d2, d3, d4)

Loads the io to given convolutional layer and updates it by batch-normalization
"""
function loadconv!(c, io::IOStream, d1, d2, d3, d4)
    read!(io, c.b)
    gamma = Array{Float32}(undef, d4)
    mean = Array{Float32}(undef, d4)
    variance = Array{Float32}(undef, d4)
    read!(io, gamma)
    read!(io, mean)
    read!(io, variance)
    toRead = Array{Float32}(undef, d4 * d3 * d2 * d1)
    read!(io, toRead)
    toRead = reshape(toRead, d1, d2, d3, d4)
    c.w = permutedims(toRead, [2, 1, 3, 4])
    c.w, c.b = updateconv!(c, gamma, mean, variance)
    c.w = flipkernel(c.w)
    if GPU >= 0
        c.w = KnetArray(c.w)
        c.b = KnetArray(c.b)
    end
end

"""
    readconstants!(io)

Read constant and unnecessary numbers from the io stream
"""
function readconstants!(io::IOStream)
    major = read(io, Int32)
    minor = read(io, Int32)
    revision = read(io, Int32)
    iseen = read(io, Int32)
end