text:
import Printf: @sprintf
import Statistics: mean
#import Plots: plot
using RecipesBase

const ThreadRecord = Tuple{Int, Int, Bool, UInt64} #Jobnum, Thread ID, ?Start:Stop, Time

struct Job
    id    :: Int
    tid   :: Int
    start :: Int
    stop  :: Int
end

const ThreadLog = Dict{Int, Vector{Job}}

duration(job::Job) = job.stop - job.start
active(job::Job, t) = t >= job.start && t < job.stop

duration(jobs::Vector{Job}) = maximum(j.stop for j in jobs) - minimum(j.start for j in jobs) # Nonoptimized
jobcount(jobs::Vector{Job}, t) = length(filter(j -> active(j,t), jobs))
jobactive(jobs::Vector{Job}, t) = first(filter(j -> active(j,t), jobs))
gaptime(jobs::Vector{Job}) = sum(y.start - x.stop for (x,y) in zip(jobs[1:end-1], jobs[2:end]))

function _recordstolog!(log::ThreadLog, records::Vector{ThreadRecord}, t0 = 0)
    starts = Dict{Int, Float64}()
    for (i, (id, tid, start, t)) in enumerate(records)
        if start
            starts[id] = t
        else
            haskey(starts, id) || error("Stop encountered before start: record $i, job id $id")
            push!(get!(log, tid, Job[]), Job(id, tid, starts[id]-t0, t-t0))
        end
    end
    for jobs in values(log)
        sort!(jobs, by=x->x.id)
    end
    return log
end

_recordstolog(records, t0) = _recordstolog!(ThreadLog(), records, t0)

"""
    ThreadPools.dumplog(io, log)

"""

function dumplog(io::IO, log::ThreadLog)
    for job in Iterators.flatten(values(log))
        write(io, "$(job.id) $(job.tid) $(job.start) $(job.stop)\n")
    end
    close(io)
end

dumplog(fname::String, log::ThreadLog) = dumplog(open(fname, "w"), log)
dumplog(io, pool::AbstractThreadPool) = dumplog(io, pool.log)

"""
    ThreadPools.readlog(io) -> Dict of (thread # => job list)

Analyzes the output of a [`LoggedThreadPool`](@ref) and produces the history
of each job on each thread.

Each job in the job list is a struct of:
```
struct Job
    id    :: Int
    tid   :: Int
    start :: Float64
    stop  :: Float64
end
```
The default sorting order of the jobs in each thread are by stop time.  `io`
can either be an IO object or a filename.

# Example
```julia
julia> log = ThreadPools.readlog("mylog.txt")
Dict{Int64,Array{ThreadPools.Job,1}} with 3 entries:
  4 => ThreadPools.Job[Job(3, 4, 0.016, 0.327), Job(6, 4, 0.327, 0.928)]
  2 => ThreadPools.Job[Job(2, 2, 0.016, 0.233), Job(5, 2, 0.233, 0.749)]
  3 => ThreadPools.Job[Job(1, 3, 0.016, 0.139), Job(4, 3, 0.139, 0.546)]
```
"""
function readlog(io::IO)
    log = ThreadLog()
    job = 0
    tid = 0
    start = 0.0
    stop = 0.0
    for (i,line) in enumerate(readlines(io))
        try
            (a,b,c,d) = split(line)
            job   = parse(Int, a)
            tid   = parse(Int, b)
            start = parse(Int, c)
            stop  = parse(Int, d)
        catch
            error("Malformed log entry, line $i: $line")
        end
        push!(get!(log, tid, Job[]), Job(job, tid, start, stop))
    end
    close(io)
    return log
end

readlog(fname::String) = readlog(open(fname))

_pad(n) = string([" " for i in 1:n]...)

function _formatjob(log, tid, t, width)
    none = string(_pad(width-2), "- ")
    haskey(log, tid) || return none
    jobs = log[tid]
    count = jobcount(jobs, t)
    count > 0 || return none
    jobstr = string(jobactive(jobs, t).id)
    return string(_pad(width-2-length(jobstr)), count > 1 ? "*" : " ", jobstr, " ")
end

_tons(x) = Int(round(x/1e-9))

"""
    ThreadPools.showactivity([io, ]log, dt, t0=0, t1=Inf; nthreads=Threads.nthreads())

Produces a textual graph of the thread activity in the provided log.

The format of the output is

```julia
julia> ThreadPools.showactivity("mylog.txt", 0.1)
0.000   -   -   -   -
0.100   4   2   1   3
0.200   4   2   5   3
0.300   4   6   5   3
0.400   4   6   5   7
0.500   8   6   5   7
0.600   8   6   5   7
0.700   8   6   -   7
0.800   8   6   -   7
0.900   8   -   -   7
1.000   8   -   -   7
1.100   8   -   -   -
1.200   8   -   -   -
1.300   -   -   -   -
1.400   -   -   -   -
```
where the first column is time, and each column afterwards is the active job
id in each thread (threads 1:nthreads, left to right) at that point in
time.

If `io` is provided, the output will be written there.  `log` may be a log IO
object, or a filename to be opened and read.  `dt` is the time step
for each row, `t0` is the optional starting time, `t1` the optional stopping
time, and `nthreads` is the number of threads to print.
"""
function showactivity(io, log::ThreadLog, dt, t0=0, t1=Inf; nthreads=0)
    maxj = maximum(j.id for jobs in values(log) for j in jobs)
    width = length(string(maxj)) + 3

    maxt = maximum(duration(x) for x in values(log))
    dtns = _tons(dt)
    t0ns = _tons(t0)
    t1ns = min(t1==Inf ? typemax(Int64) : _tons(t1), (ceil(maxt/dtns)+2)*dtns)
    tns = (t0ns ÷ dtns)*dtns

    nthreads = nthreads == 0 ? Threads.nthreads() : nthreads

    # t = floor(t0/dt)*dt
    # tns  = Int(round(t/1e-9))
    # t0ns = Int(round(t0/1e-9))
    # t1ns = Int(round(t1/1e-9))
    # dtns = Int(round(dt/1e-9))
    # println((tns, t0ns, t1ns, dtns))
    while tns <= t1ns
        tstr = @sprintf "%0.3f" (tns*1e-9)
        println(io, "$tstr $(string([_formatjob(log, tid, tns, width) for tid in 1:nthreads]...))")
        tns += dtns
    end
end

showactivity(io, fname::String, dt, t0=0, t1=Inf; nthreads=0) = showactivity(io, readlog(fname), dt, t0, t1; nthreads=nthreads)
showactivity(log::ThreadLog, dt, t0=0, t1=Inf; nthreads=0) = showactivity(Base.stdout, log, dt, t0, t1; nthreads=nthreads)
showactivity(fname::String, dt, t0=0, t1=Inf; nthreads=0) = showactivity(Base.stdout, readlog(fname), dt, t0, t1; nthreads=nthreads)
showactivity(io, pool::AbstractThreadPool, dt, t0=0, t1=Inf; nthreads=0) = showactivity(io, pool.log, dt, t0, t1; nthreads=nthreads)
showactivity(pool::AbstractThreadPool, dt, t0=0, t1=Inf; nthreads=0) = showactivity(Base.stdout, pool.log, dt, t0, t1; nthreads=nthreads)

"""
    ThreadPools.showstats([io, ]log)

Produces a statistical analysis of the provided log.

# Example
```julia
julia> ThreadPools.showstats("mylog.txt")

    Total duration: 1.542 s
    Number of jobs: 8
    Average job duration: 0.462 s
    Minimum job duration: 0.111 s
    Maximum job duration: 0.82 s

    Thread 2: Duration 1.542 s, Gap time 0.0 s
    Thread 3: Duration 1.23 s, Gap time 0.0 s
    Thread 4: Duration 0.925 s, Gap time 0.0 s

```
"""
function showstats(io, log::Dict{Int, Vector{Job}})
    flat = collect(Iterators.flatten(jobs for jobs in values(log)))
    totduration = duration(flat)*1e-9
    durations = [duration(job)*1e-9 for job in flat]
    print(io, """

    Total duration: $(round(totduration; digits=3)) s
    Number of jobs: $(length(flat))
    Average job duration: $(round(mean(durations); digits=3)) s
    Minimum job duration: $(round(minimum(durations); digits=3)) s
    Maximum job duration: $(round(maximum(durations); digits=3)) s

""")
    for thrd in sort!(collect(keys(log)))
        dur = round(duration(log[thrd])*1e-9; digits=3)
        gap = round(gaptime(log[thrd])*1e-9; digits=3)
        println(io, "    Thread $thrd: Duration $dur s, Gap time $gap s")
    end
end

showstats(io, fname::String) = showstats(io, readlog(fname))
showstats(log::Dict{Int, Vector{Job}}) = showstats(Base.stdout, log)
showstats(fname::String) = showstats(Base.stdout, readlog(fname))

"""
    plot(log)
    plot(pool)

Produces a plot of the activity across the threads.  `log` can be
the output of `readlog` or a string representing the log file name.
"""
@recipe function f(log::ThreadLog)
    ids, threadids, starts, stops = _get_plot_data(log)

    seriescolor   --> [threadids threadids]'
    xguide  --> "Time [s]"
    yguide  --> "Job Index"
    legend  --> false

    [starts stops]',  [ids ids]'
end

@recipe function f(pool::AbstractThreadPool)
    log = pool.log
    ids, threadids, starts, stops = _get_plot_data(log)

    seriescolor   --> [threadids threadids]'
    xguide  --> "Time [s]"
    yguide  --> "Job Index"
    legend  --> false

    [starts stops]',  [ids ids]'
end

function _get_plot_data(log::ThreadLog)
    jobs = collect(Iterators.flatten(jobs for jobs in values(log)))
    t0 = minimum(j.start for j in jobs)
    sort!(jobs, by=j->j.id)
    ids = [j.id for j in jobs]
    starts = [(j.start-t0)/1e9 for j in jobs]
    stops = [(j.stop-t0)/1e9 for j in jobs]
    threadids = [j.tid for j in jobs]

    return ids, threadids, starts, stops
end


text:
#################### G and K Distribution ####################

immutable GK <: ContinuousUnivariateDistribution
  A::Float64
  B::Float64
  g::Float64
  k::Float64
  c::Float64

  function GK(A::Real, B::Real, g::Real, k::Real, c::Real)
    ## check args
    0.0 <= c < 1.0 || throw(Argumenterror("c must be 0 <= c < 1"))

    ## distribution
    new(A, B, g, k, c)
  end
  GK(A::Real, B::Real, g::Real, k::Real) = GK(A, B, g, k, 0.8)
end

## Parameters

location(d::GK) = d.A
scale(d::GK) = d.B
skewness(d::GK) = d.g
kurtosis(d::GK) = d.k
asymmetry(d::GK) = d.c

params(d::GK) = (d.A, d.B, d.g, d.k, d.c)

minimum(d::GK) = -Inf
maximum(d::GK) = Inf

insupport{T <: Real}(d::GK, x::AbstractVector{T}) = true

function quantile(d::GK, p::Float64)
  z = quantile(Normal(), p)
  z2gk(d, z)
end

function Qinv(d::GK, x::Real)
  z = gk2z(d, x)
  cdf(Normal(), z)
end

function validate(d::GK)
  f = z -> -Fgkz(d, z)
  optres = optimize(f, 0, 20 / abs(skewness(d)))
  max(optres.f_minimum, 1.0) <= 1 / asymmetry(d)  &&
    scale(d) > 0 && kurtosis(d) > -0.5
end

function pdf(d::GK, x::Real)
  exp(logpdf(d, x))
end

function logpdf(d::GK, x::Real)
  if !validate(d)
    return -Inf
  end

  z = gk2z(d, x)

  term0 = exp(-skewness(d) * z)
  term1 = log(scale(d)) + log(asymmetry(d)) + 0.5 * log(2 * pi) +
    z^2 / 2 + kurtosis(d) * log(1 + z^2)
  term2 = (1 / asymmetry(d) + (1 - term0) / (1 + term0))
  term3 = (1 + (2 * kurtosis(d) + 1) * z^2) / (1 + z^2)
  term4 = 2 * skewness(d) * z * term0 / (1 + term0)^2

  logQp = term1 + log(term2 * term3 + term4)

  return -logQp
end

function logpdf{T <: Real}(d::GK, x::AbstractVector{T})
  if !validate(d)
    return -Inf
  end

  value = 0.0
  for i in 1:length(x)
    z = gk2z(d, x[i])

    term0 = exp(-skewness(d) * z)
    term1 = log(scale(d)) + log(asymmetry(d)) + 0.5 * log(2 * pi) +
      z^2 / 2 + kurtosis(d) * log(1 + z^2)
    term2 = (1 / asymmetry(d) + (1 - term0) / (1 + term0))
    term3 = (1 + (2 * kurtosis(d) + 1) * z^2) / (1 + z^2)
    term4 = 2 * skewness(d) * z * term0 / (1 + term0)^2

    value -= term1 + log(term2 * term3 + term4)
  end
  return value
end

function fit_mle{T <: Real}(d::GK, x::AbstractVector{T})
  loglikelihood = theta ->
    - logpdf(GK(theta[1], theta[2], theta[3], theta[4]), x)

  A0 = median(x)
  B0 = diff(quantile(x, [0.25,0.75]))[1]
  g0 = mean((x - A0) .^ 3) / (var(x) ^ (3/2))
  k0 = mean((x - A0) .^ 4) / (sumabs2(x - A0) / length(x))^2 - 3

  optres = optimize(loglikelihood, [A0, B0, g0, k0])

  return GK(optres.minimum...)
end

function gk2z(d::GK, x::Real; max_expands::Integer = 100, expand::Integer = 10)
  y = (x - location(d)) / scale(d)
  if y == 0
    return 0
  elseif isfinite(y)
    d0 = GK(0, 1, skewness(d), kurtosis(d), asymmetry(d))
    inter = [-abs(y) / (1 - asymmetry(d)), abs(y)]
    iter = 0
    a, b = z2gk(d0, inter[1]) - y, z2gk(d0, inter[2]) - y
    while a * b > 0
      idx = findmin(abs([a,b]))[2]
      inter[idx] -= sign(a)*expand
      iter += 1
      iter < max_expands || break
    end
    return fzero(z -> z2gk(d0, z) - y, inter)
  else
    return y
  end
end

function Fgkz(d::GK, z::Float64)
  term1 = exp(-skewness(d) * z)
  term2 = 2 * skewness(d) * z * term1
  term3 = (1 + (2 * kurtosis(d) + 1) * z ^ 2) / (1 + z ^ 2)

  (1 - term1) / (1 + term1) + term2 / ((1 + term1) * term3)
end

function z2gk(d::GK, z::Float64)
  term1 = exp(-skewness(d) * z)
  term2 = (1.0 + asymmetry(d) * (1.0 - term1) / (1.0 + term1))
  term3 = (1.0 + z ^ 2) ^ kurtosis(d)
  location(d) + scale(d) * z * term2 * term3
end

ext: jl


text:
using DataFrames, Random

export partition_train_test
export train_test_split

"given a dataframe, split it into two random parts based on the at ratio"
function partition_train_test(data::DataFrame, at=0.66, seed=rand(1:1000))
    Random.seed!(seed)
    n = nrow(data)
    idx = shuffle(1:n)
    train_idx = view(idx, 1:floor(Int, at*n))
    test_idx = view(idx, (floor(Int, at*n)+1):n)
    data[train_idx,:], data[test_idx,:]
end

"""
_given a dataframe, split it to train and test groups_
#### parameters:
    data : DataFrame
        dataframe to split
    target_col : String or Symbol
        target variable
    train_share
#### returns: DataFrame
    a new dataframe with one hot encoded columns added
"""
function train_test_split(X::DataFrame, target_col::String; train_share=0.66, seed=rand(1:1000))
    train,test = partition_train_test(X, train_share, seed)
    X_train = select(train, Not(Symbol(target_col)))
    y_train = train[!, Symbol(target_col)]
    X_test = select(test, Not(Symbol(target_col)))
    y_test = test[!, Symbol(target_col)]
    return X_train, X_test, y_train, y_test
end


text:
using Flux
using TensorCast
using Zygote
using Zygote: @adjoint

## Utils
checkpoint(f, x...) = f(x...)
@adjoint checkpoint(f, x...) = f(x...), ȳ -> Zygote._pullback(f, x...)[2](ȳ)

## Attention

#=
    The attention module is made of 5 dense layers.
    The first 4 linearly project the input data and position data
    so they can be processed by later parts of the module.

    The last dense layer combines the attention calculations from the different
    heads into the module's final output
=#
function linear_projection(W, data)
    @reduce C[head, meaning, time, ins] :=
            sum(word) W[head, meaning, word] * data[word, time, ins]
    return C
end

function mask_and_softmax(scores, mask)
    # Mask
    if mask != nothing
        scores = scores .+ mask
    end

    # Softmax
    orig_sz = size(scores)
    scores = reshape(scores,
        orig_sz[1], orig_sz[2], orig_sz[3], orig_sz[4] * orig_sz[5])

    scores = softmax(scores, dims = 4)
    return reshape(scores,
        orig_sz[1], orig_sz[2], orig_sz[3], orig_sz[4], orig_sz[5])
end

function get_scores(my_queries, my_keys, my_pos, mask)
    # For every word in the query projection (i, j)
    # Add the key projection to the relevant section of the position projection
    # Then multiply by the (i, j)th query projection
    scores = []
    for i = 1:size(my_queries)[3]
        ins_scores = []
        for j = 1:size(my_queries)[4]
            Q = my_queries[:, :, i, j]
            K = my_keys
            P = my_pos[:,:, i:(i+size(my_queries)[3]-1),
                            j:(j+size(my_queries)[4]-1),]

            E = P .+ K
            @reduce C[head, q_time, q_ins, k_time, k_ins] :=
                sum(meaning) Q[head, meaning, q_time, q_ins] *
                                E[head, meaning, k_time, k_ins,]

            # Pass the information up from the for-loop
            #   in a way that doesn't break Flux.gradient()
            # I'm sure there's a better way, but I don't know it
            if ins_scores == []
                ins_scores = C
            else
                ins_scores = cat(ins_scores, C, dims = 3)
            end
        end
        # Pass the information up again
        if scores == []
            scores = ins_scores
        else
            scores = cat(scores, ins_scores, dims = 2)
        end
    end

    # Scale scores, and take the softmax over all query/key timesteps/instruments
    scores = scores ./ sqrt(sum(size(my_keys, 3)))

    return mask_and_softmax(scores, mask)
end

function get_scaled_values(scores, my_values)
    @reduce scaled_values[head, meaning, q_time, q_ins] :=
        sum(k_time, k_ins) my_values[head, meaning, k_time, k_ins] *
            scores[head, q_time, q_ins, k_time, k_ins]
end

struct Attention{
    S<:AbstractArray,
    T<:AbstractArray,
    U<:AbstractArray,
    V<:AbstractArray,
    W<:AbstractArray,
}
    W_Q::S # Dense layer for queries
    W_K::T # Dense layer for keys
    W_V::U # Dense layer for values
    W_P::V # Dense layer for positions
    W_H::W # Dense layer for heads
end

function Attention(
    word_size::Integer,
    latent_size::Integer;
    head_count = 8::Integer,
    initW = Flux.glorot_uniform,
)
    return Attention(
        initW(head_count, latent_size, word_size),
        initW(head_count, latent_size, word_size),
        initW(head_count, latent_size, word_size),
        initW(head_count, latent_size, word_size),
        initW(word_size, head_count * latent_size),
    )
end

Flux.@functor Attention

function (a::Attention)(
    query_data::AbstractArray,
    pos_data::AbstractArray;
    key_data = nothing::Union{Nothing,AbstractArray},
    mask = nothing::Union{Nothing,AbstractArray},
)
    W_Q, W_K, W_V, W_P, W_H = a.W_Q, a.W_K, a.W_V, a.W_P, a.W_H

    # For self-attention, keys and queries come from the same input
    if key_data == nothing
        key_data = query_data
    end

    # Get scaled values
    scaled_values = checkpoint(get_scaled_values,
                        checkpoint(get_scores,
                            linear_projection(W_Q, query_data),
                            linear_projection(W_K, key_data),
                            linear_projection(W_P, pos_data),
                            mask),
                        linear_projection(W_V, key_data))

    # Concatenate along heads
    dims = size(scaled_values)
    scaled_values =
        reshape(scaled_values, (dims[1] * dims[2], dims[3], dims[4]))

    # Transform back to input dimensions
    @reduce z[word, time, ins] :=
        sum(long_meaning) W_H[word, long_meaning] *
                        scaled_values[long_meaning, time, ins]
end

## Encoder

# An encoder is just a LayerNorm followed by an attention layer,
# With a residual connection and a relu on the end
struct Encoder
    norm::LayerNorm
    self_attn::Attention
end

function Encoder(head_count::Integer, word_size::Integer, latent_size::Integer)
    return Encoder(
        LayerNorm(word_size),
        Attention(word_size, latent_size, head_count = head_count),
    )
end

Flux.@functor Encoder

function (a::Encoder)(
    query_data::AbstractArray,
    pos_data::AbstractArray;
    mask = nothing::Union{Nothing,AbstractArray},
)
    norm, self_attn = a.norm, a.self_attn
    sublayer =
        relu.(query_data + self_attn(norm(query_data), pos_data, mask = mask))
end

## Decoder
# A decoder is just two encoders stacked on top of each other, with the
# second one getting it's key and value data from somewhere else.

struct Decoder
    norm1::LayerNorm
    self_attn::Attention
    norm2::LayerNorm
    cross_attn::Attention
end

function Decoder(head_count::Integer, word_size::Integer, latent_size::Integer)
    return Decoder(
        LayerNorm(word_size),
        Attention(word_size, latent_size, head_count = head_count),
        LayerNorm(word_size),
        Attention(word_size, latent_size, head_count = head_count),
    )
end

Flux.@functor Decoder

function (a::Decoder)(
    query_data::AbstractArray,
    key_data::AbstractArray,
    pos_data::AbstractArray;
    mask = nothing::Union{Nothing,AbstractArray},
)
    norm1, self_attn, norm2, cross_attn =
        a.norm1, a.self_attn, a.norm2, a.cross_attn

    sublayer1 =
        relu.(query_data + self_attn(norm1(query_data), pos_data, mask = mask))
    sublayer2 =
        relu.(
            sublayer1 +
            cross_attn(norm2(sublayer1), pos_data, key_data = norm2(key_data)),
        )
end

## Transformer
# The transformer just brings some number of encoders and decoders together,
# with the cross-attention in the decoders getting their key and value data
# from the final output of the encoder chain.
struct Transformer
    encoder::Union{Encoder,Chain}
    decoder::Union{Decoder,Chain}
end

function Transformer(
    N::Integer,
    head_count::Integer,
    word_size::Integer,
    latent_size::Integer,
)
    return Transformer(
        Chain([Encoder(head_count, word_size, latent_size) for i = 1:N]...),
        Chain([Decoder(head_count, word_size, latent_size) for i = 1:N]...),
    )
end

Flux.@functor Transformer

function (a::Transformer)(
    encoder_input,
    pos_data;
    decoder_input = nothing,
    mask = nothing,
)
    encoder, decoder = a.encoder, a.decoder

    if decoder_input == nothing
        decoder_input = encoder_input
    end

    encoded_data = encoder[1](encoder_input, pos_data, mask = mask)
    for i = 2:length(encoder)
        encoded_data = encoder[i](encoded_data, pos_data)
    end

    output = decoder[1](decoder_input, encoded_data, pos_data, mask = mask)
    for i = 2:length(decoder)
        output = decoder[i](output, encoded_data, pos_data)
    end

    return output
end


text:
using LinearAlgebra: BlasInt, checksquare
using LinearAlgebra.LAPACK: chkargsok

# po
## potrf
for (bname, fname,elty) in ((:cusolverDnSpotrf_bufferSize, :cusolverDnSpotrf, :Float32),
                            (:cusolverDnDpotrf_bufferSize, :cusolverDnDpotrf, :Float64),
                            (:cusolverDnCpotrf_bufferSize, :cusolverDnCpotrf, :ComplexF32),
                            (:cusolverDnZpotrf_bufferSize, :cusolverDnZpotrf, :ComplexF64))
    @eval begin
        function LinearAlgebra.LAPACK.potrf!(uplo::Char,
                        A::CuMatrix{$elty})
            cuuplo  = cublasfill(uplo)
            n       = checksquare(A)
            lda     = max(1, stride(A, 2))
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, cublasFillMode_t, Cint,
                          Ptr{$elty}, Cint, Ref{Cint}),
                         dense_handle(), cuuplo, n, A, lda, bufSize)

            buffer  = CuArray{$elty}(undef, bufSize[])
            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, cublasFillMode_t, Cint,
                          Ptr{$elty}, Cint, Ptr{$elty}, Cint, Ptr{Cint}),
                         dense_handle(), cuuplo, n, A, lda, buffer,
                         bufSize[], devinfo)
            info = BlasInt(_getindex(devinfo, 1))
            chkargsok(info)
            A, info
        end
    end
end

## potrs
for (fname,elty) in ((:cusolverDnSpotrs, :Float32),
                     (:cusolverDnDpotrs, :Float64),
                     (:cusolverDnCpotrs, :ComplexF32),
                     (:cusolverDnZpotrs, :ComplexF64))
    @eval begin
        function LinearAlgebra.LAPACK.potrs!(uplo::Char,
                        A::CuMatrix{$elty},
                        B::CuVecOrMat{$elty})
            cuuplo = cublasfill(uplo)
            n = checksquare(A)
            if size(B, 1) != n
                throw(DimensionMismatch("first dimension of B, $(size(B,1)), must match second dimension of A, $n"))
            end
            nrhs = size(B,2)
            lda  = max(1, stride(A, 2))
            ldb  = max(1, stride(B, 2))

            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, cublasFillMode_t, Cint, Cint,
                          Ptr{$elty}, Cint, Ptr{$elty}, Cint, Ptr{Cint}),
                         dense_handle(), cuuplo, n, nrhs, A, lda, B,
                         ldb, devinfo)
            info = _getindex(devinfo, 1)
            chkargsok(BlasInt(info))
            B
        end
    end
end

#getrf
for (bname, fname,elty) in ((:cusolverDnSgetrf_bufferSize, :cusolverDnSgetrf, :Float32),
                            (:cusolverDnDgetrf_bufferSize, :cusolverDnDgetrf, :Float64),
                            (:cusolverDnCgetrf_bufferSize, :cusolverDnCgetrf, :ComplexF32),
                            (:cusolverDnZgetrf_bufferSize, :cusolverDnZgetrf, :ComplexF64))
    @eval begin
        function getrf!(A::CuMatrix{$elty})
            m,n     = size(A)
            lda     = max(1, stride(A, 2))
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ptr{$elty}, Cint,
                          Ref{Cint}), dense_handle(), m, n, A, lda,
                         bufSize)

            buffer  = CuArray{$elty}(undef, bufSize[])
            devipiv = CuArray{Cint}(undef, min(m,n))
            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ptr{$elty},
                          Cint, Ptr{$elty}, Ptr{Cint}, Ptr{Cint}),
                         dense_handle(), m, n, A, lda, buffer,
                         devipiv, devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            elseif info > 0
                throw(LinearAlgebra.SingularException(info))
            end
            A, devipiv
        end
    end
end

#geqrf
for (bname, fname,elty) in ((:cusolverDnSgeqrf_bufferSize, :cusolverDnSgeqrf, :Float32),
                            (:cusolverDnDgeqrf_bufferSize, :cusolverDnDgeqrf, :Float64),
                            (:cusolverDnCgeqrf_bufferSize, :cusolverDnCgeqrf, :ComplexF32),
                            (:cusolverDnZgeqrf_bufferSize, :cusolverDnZgeqrf, :ComplexF64))
    @eval begin
        function geqrf!(A::CuMatrix{$elty})
            m, n    = size(A)
            lda     = max(1, stride(A, 2))
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)),libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ptr{$elty}, Cint,
                          Ref{Cint}), dense_handle(), m, n, A,
                         lda, bufSize)
            buffer  = CuArray{$elty}(undef, bufSize[])
            tau  = CuArray{$elty}(undef, min(m, n))
            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)),libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ptr{$elty},
                          Cint, Ptr{$elty}, Ptr{$elty}, Cint, Ptr{Cint}),
                         dense_handle(), m, n, A, lda, tau, buffer,
                         bufSize[], devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            end
            A, tau
        end
    end
end

#sytrf
for (bname, fname,elty) in ((:cusolverDnSsytrf_bufferSize, :cusolverDnSsytrf, :Float32),
                            (:cusolverDnDsytrf_bufferSize, :cusolverDnDsytrf, :Float64),
                            (:cusolverDnCsytrf_bufferSize, :cusolverDnCsytrf, :ComplexF32),
                            (:cusolverDnZsytrf_bufferSize, :cusolverDnZsytrf, :ComplexF64))
    @eval begin
        function sytrf!(uplo::Char,
                        A::CuMatrix{$elty})
            cuuplo = cublasfill(uplo)
            n      = size(A, 1)
            if size(A,2) != n
                throw(DimensionMismatch("SymTridiagonal matrix must be square!"))
            end
            lda = max(1, stride(A, 2))
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)),libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Ptr{$elty}, Cint,
                          Ref{Cint}), dense_handle(), n, A, lda,
                         bufSize)

            buffer  = CuArray{$elty}(undef, bufSize[])
            devipiv = CuArray{Cint}(undef, n)
            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)),libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, cublasFillMode_t, Cint,
                          Ptr{$elty}, Cint, Ptr{Cint}, Ptr{$elty}, Cint,
                          Ptr{Cint}), dense_handle(), cuuplo, n, A,
                         lda, devipiv, buffer, bufSize[], devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            elseif info > 0
                throw(LinearAlgebra.SingularException(info))
            end
            A, devipiv
        end
    end
end

#getrs
for (fname,elty) in ((:cusolverDnSgetrs, :Float32),
                     (:cusolverDnDgetrs, :Float64),
                     (:cusolverDnCgetrs, :ComplexF32),
                     (:cusolverDnZgetrs, :ComplexF64))
    @eval begin
        function getrs!(trans::Char,
                        A::CuMatrix{$elty},
                        ipiv::CuVector{Cint},
                        B::CuVecOrMat{$elty})
            cutrans = cublasop(trans)
            n = size(A, 1)
            if size(A, 2) != n
                throw(DimensionMismatch("LU factored matrix A must be square!"))
            end
            if size(B, 1) != n
                throw(DimensionMismatch("first dimension of B, $(size(B,1)), must match second dimension of A, $n"))
            end
            nrhs = size(B, 2)
            lda  = max(1, stride(A, 2))
            ldb  = max(1, stride(B, 2))

            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, cublasOperation_t, Cint, Cint,
                          Ptr{$elty}, Cint, Ptr{Cint}, Ptr{$elty}, Cint,
                          Ptr{Cint}), dense_handle(), cutrans, n, nrhs,
                         A, lda, ipiv, B, ldb, devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            end
            B
        end
    end
end

#ormqr
for (bname, fname, elty) in ((:cusolverDnSormqr_bufferSize, :cusolverDnSormqr, :Float32),
                             (:cusolverDnDormqr_bufferSize, :cusolverDnDormqr, :Float64),
                             (:cusolverDnCunmqr_bufferSize, :cusolverDnCunmqr, :ComplexF32),
                             (:cusolverDnZunmqr_bufferSize, :cusolverDnZunmqr, :ComplexF64))    @eval begin
        function ormqr!(side::Char,
                        trans::Char,
                        A::CuMatrix{$elty},
                        tau::CuVector{$elty},
                        C::CuVecOrMat{$elty})
            cutrans = cublasop(trans)
            cuside  = cublasside(side)
            if side == 'L'
                m   = size(A, 1)
                ldc = size(C, 1)
                n   = size(C, 2)
                if m > ldc
                    Ctemp = CuArray{$elty}(undef, m - ldc, n) .= 0
                    C = [C; Ctemp]
                    ldc = m
                end
                lda = m
            else
                m   = size(C, 1)
                n   = size(C, 2)
                ldc = m
                lda = n
            end
            k       = length(tau)
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)),libcusolver), cusolverStatus_t,
                          (cusolverDnHandle_t, cublasSideMode_t,
                           cublasOperation_t, Cint, Cint, Cint, Ptr{$elty},
                           Cint, Ptr{$elty}, Ptr{$elty}, Cint, Ref{Cint}),
                          dense_handle(), cuside,
                          cutrans, m, n, k, A,
                          lda, tau, C, ldc, bufSize)
            buffer  = CuArray{$elty}(undef, bufSize[])
            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)),libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, cublasSideMode_t,
                          cublasOperation_t, Cint, Cint, Cint, Ptr{$elty},
                          Cint, Ptr{$elty}, Ptr{$elty}, Cint, Ptr{$elty},
                          Cint, Ptr{Cint}),
                         dense_handle(), cuside,
                         cutrans, m, n, k, A, lda, tau, C, ldc, buffer,
                         bufSize[], devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            end
            side == 'L' ? C : C[:, 1:minimum(size(A))]
        end
    end
end

#orgqr
for (bname, fname, elty) in ((:cusolverDnSorgqr_bufferSize, :cusolverDnSorgqr, :Float32),
                             (:cusolverDnDorgqr_bufferSize, :cusolverDnDorgqr, :Float64),
                             (:cusolverDnCungqr_bufferSize, :cusolverDnCungqr, :ComplexF32),
                             (:cusolverDnZungqr_bufferSize, :cusolverDnZungqr, :ComplexF64))
    @eval begin
        function orgqr!(A::CuMatrix{$elty}, tau::CuVector{$elty})
            m = size(A , 1)
            n = min(m, size(A, 2))
            lda = max(1, stride(A, 2))
            k = length(tau)
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)), libcusolver), cusolverStatus_t,
                          (cusolverDnHandle_t, Cint, Cint, Cint, Ptr{$elty}, Cint,
                           Ptr{$elty}, Ref{Cint}),
                          dense_handle(), m, n, k, A, lda, tau, bufSize)
            buffer  = CuArray{$elty}(undef, bufSize[])
            devinfo = CuArray{Cint}(undef, 1)
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Cint, Ptr{$elty},
                          Cint, Ptr{$elty}, Ptr{$elty}, Cint, Ptr{Cint}),
                         dense_handle(), m, n, k, A,
                         lda, tau, buffer, bufSize[], devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            end
            if n < size(A, 2)
                A[:, 1:n]
            else
                A
            end
        end
    end
end

#gebrd
for (bname, fname, elty, relty) in ((:cusolverDnSgebrd_bufferSize, :cusolverDnSgebrd, :Float32, :Float32),
                                    (:cusolverDnDgebrd_bufferSize, :cusolverDnDgebrd, :Float64, :Float64),
                                    (:cusolverDnCgebrd_bufferSize, :cusolverDnCgebrd, :ComplexF32, :Float32),
                                    (:cusolverDnZgebrd_bufferSize, :cusolverDnZgebrd, :ComplexF64, :Float64))
    @eval begin
        function gebrd!(A::CuMatrix{$elty})
            m, n    = size(A)
            lda     = max(1, stride(A, 2))
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ref{Cint}),
                        dense_handle(), m, n, bufSize)
            buffer  = CuArray{$elty}(undef, bufSize[])
            devinfo = CuArray{Cint}(undef, 1)
            k       = min(m, n)
            D       = CuArray{$relty}(undef, k)
            E       = CuArray{$relty}(undef, k) .= 0
            TAUQ    = CuArray{$elty}(undef, k)
            TAUP    = CuArray{$elty}(undef, k)
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ptr{$elty},
                          Cint, Ptr{$relty}, Ptr{$relty}, Ptr{$elty},
                          Ptr{$elty}, Ptr{$elty}, Cint, Ptr{Cint}),
                         dense_handle(), m, n, A, lda, D, E, TAUQ,
                         TAUP, buffer, bufSize[], devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            end
            A, D, E, TAUQ, TAUP
        end
    end
end

for (bname, fname, elty, relty) in ((:cusolverDnSgesvd_bufferSize, :cusolverDnSgesvd, :Float32, :Float32),
                                    (:cusolverDnDgesvd_bufferSize, :cusolverDnDgesvd, :Float64, :Float64),
                                    (:cusolverDnCgesvd_bufferSize, :cusolverDnCgesvd, :ComplexF32, :Float32),
                                    (:cusolverDnZgesvd_bufferSize, :cusolverDnZgesvd, :ComplexF64, :Float64))
    @eval begin
        function gesvd!(jobu::Char,
                        jobvt::Char,
                        A::CuMatrix{$elty})
            m,n     = size(A)
            lda     = max(1, stride(A, 2))
            bufSize = Ref{Cint}(0)
            @check ccall(($(string(bname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cint, Cint, Ref{Cint}),
                        dense_handle(), m, n, bufSize)

            buffer  = CuArray{$elty}(undef, bufSize[])
            rbuffer = CuArray{$relty}(undef, 5 * min(m, n))
            devinfo = CuArray{Cint}(undef, 1)
            U       = CuArray{$elty}(undef, m, m)
            ldu     = max(1, stride(U, 2))
            S       = CuArray{$relty}(undef, min(m, n))
            Vt      = CuArray{$elty}(undef, n, n)
            ldvt    = max(1, stride(Vt, 2))
            @check ccall(($(string(fname)), libcusolver), cusolverStatus_t,
                         (cusolverDnHandle_t, Cuchar, Cuchar, Cint,
                          Cint, Ptr{$elty}, Cint, Ptr{$relty}, Ptr{$elty},
                          Cint, Ptr{$elty}, Cint, Ptr{$elty}, Cint,
                          Ptr{$relty}, Ptr{Cint}), dense_handle(),
                         jobu, jobvt, m, n, A, lda, S, U, ldu, Vt, ldvt,
                         buffer, bufSize[], rbuffer, devinfo)
            info = _getindex(devinfo, 1)
            if info < 0
                throw(ArgumentError("The $(info)th parameter is wrong"))
            end
            U, S, Vt
        end
    end
end


text:
using Flux
using Metalhead

function main()
    vgg = VGG19()
    Flux.testmode!(vgg, false)
    model = Chain(vgg.layers[1:end-2],
                  Dense(4096, 101),
                  softmax)
    Flux.testmode!(model,false)

    val_model = model |> cpu
    println(model)
    println(val_model)
    Flux.testmode!(val_model)
    println(model)
    println(val_model)
end

main()

max_stars_repo_name: terasakisatoshi/juliaExer

max_issues_repo_name: terasakisatoshi/juliaExer

max_forks_repo_name: terasakisatoshi/juliaExer


text:
module RelativeEstimationProject

using StaticArrays

include("relativedynamics.jl")
include("navigationsensor.jl")
include("estimationfilters.jl")

end # module

max_stars_repo_name: glemieux/RelativeEstimationProject

max_issues_repo_name: glemieux/RelativeEstimationProject

max_forks_repo_name: glemieux/RelativeEstimationProject


text:
using Documenter, InsarTimeseries

makedocs(
    modules = [InsarTimeseries],
    format = Documenter.HTML(),
    checkdocs = :exports,
    sitename = "InsarTimeseries.jl",
    pages = Any["index.md"],
)

deploydocs(repo = "github.com/scottstanie/InsarTimeseries.jl.git")


text:
const MinHeap = MutableBinaryHeap{Float32, DataStructures.FasterForward}

# Efficient inverse query for fenwick tree : adapted from https://codeforces.com/blog/entry/61364
function inverse_query(t::FenwickTree, v, N = length(t))
    tot, pos = 0, 0
    for i=floor(Int, log2(N)):-1:0
        new_pos = pos + 1 << i
        if new_pos <= N && tot + t.bi_tree[new_pos] < v
            tot += t.bi_tree[new_pos]
            pos = new_pos
        end
    end
    pos + 1
end

Base.getindex(t::FenwickTree, i::Int) = prefixsum(t, i) - prefixsum(t, i-1)

DataStructures.update!(t::FenwickTree, i, v) = inc!(t, i, v - t[i])

Base.zero(s::Type{Symbol}) = :zero # For initializing symbolic arrays

# construction of common mdp data
function mdp_data(S::T1, A::T2, capacity::Int, extras::Array{Symbol} = Symbol[]; ArrayType = Array, R = Float32, D = Bool, W = Float32) where {T1 <: AbstractSpace, T2 <: AbstractSpace}
    data = Dict{Symbol, ArrayType}(
        :s => ArrayType(fill(zero(type(S)), dim(S)..., capacity)),
        :a => ArrayType(fill(zero(type(A)), dim(A)..., capacity)),
        :sp => ArrayType(fill(zero(type(S)), dim(S)..., capacity)),
        :r => ArrayType(fill(zero(R), 1, capacity)),
        :done => ArrayType(fill(zero(D), 1, capacity))
        )
    for k in extras
        if k in [:return, :logprob, :xlogprob, :advantage, :cost, :cost_advantage, :cost_return]
            data[k] = ArrayType(fill(zero(R), 1, capacity))
        elseif k in [:weight]
            data[k] = ArrayType(fill(one(R), 1, capacity))
        elseif k in [:episode_end, :fail]
            data[k] = ArrayType(fill(false, 1, capacity))
        elseif k in [:t, :i]
            data[k] = ArrayType(fill(0, 1, capacity))
        elseif k in [:s0]
            data[k] = ArrayType(fill(zero(type(S)), dim(S)..., capacity))
        elseif k in [:x] # Disturbances for adversarial policy
            data[k] = ArrayType(fill(zero(type(A)), dim(A)..., capacity))
        else
            CRUX_WARNINGS && @warn "Unrecognized key: $k"
        end
    end
    data
end

## Prioritized replay parameters
@with_kw mutable struct PriorityParams
    minsort_priorities::MinHeap
    priorities::FenwickTree
    α::Float32 = 0.6
    β::Function = (i) -> 0.5f0
    max_priority::Float32 = 1.0
end

PriorityParams(N::Int; kwargs...) = PriorityParams(;minsort_priorities=MinHeap(fill(Inf32, N)), priorities=FenwickTree(fill(0f0, N)), kwargs...)
PriorityParams(N::Int, pp::PriorityParams) = PriorityParams(N, α=pp.α, β=pp.β, max_priority=pp.max_priority)

## Experience Buffer stuff
@with_kw mutable struct ExperienceBuffer{T <: AbstractArray}
    data::Dict{Symbol, T}
    elements::Int64
    next_ind::Int64
    indices::Vector{Int}
    priority_params::Union{Nothing, PriorityParams}
end

function ExperienceBuffer(data::Dict{Symbol, T}; elements=capacity(data), next_ind=1, indices=[], prioritized::Bool=false, priority_params::NamedTuple=(;)) where {T <: AbstractArray}
    pp = nothing
    if prioritized
        !haskey(data, :weight) && (data[:weight] = device(data)((ones(Float32, 1, capacity(data)))))
        pp = PriorityParams(capacity(data); priority_params...)
    end
    ExperienceBuffer(data, elements, next_ind, indices, pp)
end

function ExperienceBuffer(S::T1, A::T2, capacity::Int, extras::Array{Symbol}=Symbol[]; device=cpu,
                          prioritized=false, priority_params::NamedTuple=(;), R=Float32, D=Bool, W=Float32) where {T1 <: AbstractSpace, T2 <: AbstractSpace}
    Atype = device == gpu ? CuArray : Array
    data = mdp_data(S, A, capacity, extras, ArrayType=Atype, R=R, D=D, W=W)
    ExperienceBuffer(data, elements=0, prioritized=prioritized, priority_params=priority_params)
end

function buffer_like(b::ExperienceBuffer; capacity=capacity(b), device=device(b))
    data = Dict(k=>device(Array{eltype(v)}(undef, size(v)[1:end-1]..., capacity)) for (k,v) in b.data)
    ExperienceBuffer(data, 0, 1, Int[], isprioritized(b) ? PriorityParams(capacity, b.priority_params) : nothing)
end

function Flux.gpu(b::ExperienceBuffer)
    data = Dict{Symbol, CuArray}(k => v |> gpu for (k,v) in b.data)
    ExperienceBuffer(data, b.elements, b.next_ind, b.indices, b.priority_params)
end

function Flux.cpu(b::ExperienceBuffer)
    data = Dict{Symbol, Array}(k => v |> cpu for (k,v) in b.data)
    ExperienceBuffer(data, b.elements, b.next_ind, b.indices, b.priority_params)
end

function clear!(b::ExperienceBuffer)
    b.elements = 0
    b.next_ind = 1
    b.indices = Int[]
    isprioritized(b) && (b.priority_params = PriorityParams(capacity(b), b.priority_params))
    b
end

function Base.hcat(buffers::ExperienceBuffer...; kwargs...)
    data = Dict(k=>typeof(v)(undef, size(v)[1:end-1]...,0) for (k,v) in buffers[1].data)
    for b in buffers[1:end]
        @assert keys(data) == keys(b)
        for k in keys(b)
            data[k] = cat(data[k], b[k], dims=ndims(b[k]))
        end
    end
    ExperienceBuffer(data; kwargs...)
end

function Random.shuffle!(b::ExperienceBuffer)
    new_i = shuffle(1:length(b))
    for k in keys(b)
        b[k] .= bslice(b[k], new_i)
    end
    b
end

function split_batches(N, fracs)
    @assert sum(fracs) ≈ 1
    batches = floor.(Int, N .* fracs)
    batches[1] += N - sum(batches)
    batches
end

function Base.split(b::ExperienceBuffer, fracs)
    buffers = ExperienceBuffer[]
    start = 1
    for batch in split_batches(length(b), fracs)
        push!(buffers, ExperienceBuffer(minibatch(b, start:start+batch-1)))
        start += batch
    end
    buffers
end

function normalize!(b::ExperienceBuffer, S::AbstractSpace, A::AbstractSpace)
    haskey(b, :s) && (b[:s] .= tovec(b[:s], S))
    haskey(b, :sp) && (b[:sp] .= tovec(b[:sp], S))
    A isa ContinuousSpace && (b[:a] .= tovec(b[:a], A))
    b
end

function trim!(b::ExperienceBuffer, range)
    for k in keys(b)
        b.data[k] = bslice(b.data[k], range)
    end
    b.elements = min(b.elements, length(range))
    if b.next_ind > b.elements
        b.next_ind = 1
    end
    b
end

minibatch(b::ExperienceBuffer, indices) = Dict(k => bslice(b.data[k], indices) for k in keys(b))

Base.getindex(b::ExperienceBuffer, key::Symbol) = bslice(b.data[key], 1:b.elements)

Base.keys(b::ExperienceBuffer) = keys(b.data)

extra_columns(b::ExperienceBuffer) = collect(setdiff(keys(b), [:s, :a, :sp, :r, :done]))

Base.first(b::ExperienceBuffer) = first(b.data)

Base.haskey(b::ExperienceBuffer, k) = haskey(b.data, k)

Base.length(b::ExperienceBuffer) = b.elements

DataStructures.capacity(b::Union{Dict{Symbol, T}, ExperienceBuffer}) where {T <: AbstractArray} = size(first(b)[2])[end]

isprioritized(b::ExperienceBuffer) = !isnothing(b.priority_params)

device(b::ExperienceBuffer{CuArray}) = gpu
device(b::Dict{Symbol, CuArray}) = gpu
device(b::ExperienceBuffer{Array}) = cpu
device(b::Dict{Symbol, Array}) = cpu

function episodes(b::ExperienceBuffer)
    if haskey(b, :episode_end)
        ep_ends = findall(b[:episode_end][1,:])
        ep_starts = [1, ep_ends[1:end-1] .+ 1 ...]
    elseif haskey(b, :t)
        ep_starts = findall(b[:t][1,:] .== 1)
        ep_ends = [ep_starts[2:end] .- 1 ..., length(b)]
    else
        error("Need :episode_end flag or :t column to determine episodes")
    end
    zip(ep_starts, ep_ends)
end

function get_last_N_indices(b::ExperienceBuffer, N)
    # Make sure we don't exceed the number of elements actually stored in the buffer
    N = min(length(b), N)
    C = capacity(b)
    start_index = mod1(b.next_ind - N, C)
    mod1.(start_index:start_index + N - 1, C)
end

# Note: data can be a dictionary or an experience buffer
function Base.push!(b::ExperienceBuffer, data; ids = nothing)
    ids = isnothing(ids) ? UnitRange(1, size(data[first(keys(data))], 2)) : ids
    N, C = length(ids), capacity(b)
    I = mod1.(b.next_ind:b.next_ind + N - 1, C)
    for k in keys(b)
        if !haskey(data, k)
            # CRUX_WARNINGS && @warn "Pushed data does not contain $k"
            continue
        end
        v1 = bslice(b.data[k], I)
        v2 = collect(bslice(data[k], ids))
        @assert size(v1)[1:end-1] == size(v2)[1:end-1]
        copyto!(v1, v2)
    end
    isprioritized(b) && update_priorities!(b, I, b.priority_params.max_priority*ones(N))

    b.elements = min(C, b.elements + N)
    b.next_ind = mod1(b.next_ind + N, C)
    I
end

function update_priorities!(b, I::AbstractArray, v::AbstractArray)
    for i = 1:length(I)
        val = v[i] + eps(Float32)
        update!(b.priority_params.priorities, I[i], val^b.priority_params.α)
        update!(b.priority_params.minsort_priorities, I[i], val^b.priority_params.α)
        b.priority_params.max_priority = max(val, b.priority_params.max_priority)
    end
end

function Random.rand!(target::ExperienceBuffer, source::ExperienceBuffer...; i = 1, fracs = ones(length(source))./length(source))
    batches = split_batches(capacity(target), fracs)
    for (b, B) in zip(source, batches)
        B == 0 && continue
        isprioritized(b) ? prioritized_sample!(target, b, i=i, B=B) : uniform_sample!(target, b, B=B)
    end
end

function uniform_sample!(target::ExperienceBuffer, source::ExperienceBuffer; B = capacity(target))
    ids = rand(1:length(source), B)
    target.indices = ids
    push!(target, source, ids = ids)
end

# With guidance from https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py
function prioritized_sample!(target::ExperienceBuffer, source::ExperienceBuffer; i = 1, B = capacity(target))
    @assert haskey(source, :weight)
    N = length(source)
    prs = source.priority_params.priorities
    ptot = prefixsum(prs, N)
    Δp = ptot / B
    target.indices = [inverse_query(prs, (j + rand() - 1) * Δp, N-1) for j=1:B]
    pmin = first(source.priority_params.minsort_priorities) / ptot
    max_w = (pmin*N)^(-source.priority_params.β(i))
    source[:weight][1, target.indices] .= [(N * prs[id] / ptot)^source.priority_params.β(i) for id in target.indices] ./ max_w

    push!(target, source, ids = target.indices)
end


text:
# Heterozygous content
# --------------------

### The core heterozygosity planning method. All others eventually call this one.

"""
    plan_het(cb::ChromosomeBlueprint, pos::Int, alleles::Vector{DNA})

Plan heterozygosity between copies of a chromosome at position `pos`.

The `alleles` vector must contain a nucleotide for each copy of the
chromosome.

For example if you provided a vector of `[DNA_A, DNA_C]`, for the heterozygous
site you're defining at `pos`, the first copy of the chromosome will have an `A`,
and the second copy of the chromosome will have a `C`.

Creates a new chromosome blueprint, based on the input blueprint `cb`.

!!! tip
    Use the [`suggest_alleles`](@ref) function to help decide on a set of
    alleles to use.

!!! note
    In the new blueprint, the positions that were used to plan the heterozygosity,
    have been consumed, and cannot be used to plan any other subsequently added
    features.
"""
function plan_het(cb::ChromosomeBlueprint, pos::Int, alleles::Vector{DNA})
    if length(alleles) != ncopies(cb)
        throw(ArgumentError(string("alleles must have ", ncopies(cb), " entries")))
    end
    for a in alleles
        if isambiguous(a)
            throw(ArgumentError("alleles must not contain any ambiguous nucleotides"))
        end
    end
    if length(unique(alleles)) == 1
        throw(ArgumentError("alleles must contain at least two different nucleotides"))
    end
    checkavailability(cb, pos:pos)
    interplan = consume_regions(cb, pos:pos)
    newhet = HeterozygosityBlueprint(pos, alleles)
    return ChromosomeBlueprint(interplan, push!(hets(interplan), newhet))
end

### Allele suggestion methods.

"""
    suggest_alleles(groups::Vector{Int})

Suggest an allele pattern for a single heterozygous site.

The `groups` vector should have a number of
The nth value in the `groups` designates a group to the nth chromosome copy.

E.g. consider for a triploid, the vector `[2, 1, 2]` means that there are two
groups (group 1 and group 2), and that the second chromosome copy of the three
belongs to group 1, and the other two copies belong to group 2.

Chromosome copies that belong to the same group will get the same base at each
heterozygous site.

For the vector `[2, 1, 2]`, this method might return a suggested allele pattern
such as `[DNA_A, DNA_G, DNA_A]`.

!!! note
    Which base corresponds to which group number will be determined at random.
"""
function suggest_alleles(groups::Vector{Int})
    @assert all(g ≤ 4 for g in groups)
    possible = Random.shuffle!([DNA_A, DNA_C, DNA_G, DNA_T])
    return possible[groups]
end

"""
    suggest_alleles(groups::Vector{Int}...)

Suggest an allele pattern for a single heterozygous site.

Multiple integer vectors make up `groups`.

Each vector defines a group, and each value in a vector determines which copies
are allocated to that group.

E.g. consider for a triploid, two vectors of `[2] and [1, 3]` passed as arguments
to `groups`. They define two groups, the first vector contains a number
denoting the second chromosome copy, and he second vector contains numbers
representing the first and third chromosome copy.

For the vectors `[2]` and `[1, 3]`, this method might return a suggested allele
pattern such as `[DNA_A, DNA_G, DNA_A]`.

!!! note
    Which base corresponds to which group number will be determined at random.
"""
function suggest_alleles(groups::Vector{Int}...)
    l = sum(length(g) for g in groups)
    newgroups = Vector{Int}(undef, l)
    for (i, g) in enumerate(groups)
        for j in g
            newgroups[j] = i
        end
    end
    return suggest_alleles(newgroups)
end

"""
    suggest_alleles(ncopies::Int, ngroups::Int)

Suggest an allele pattern for a single heterozygous site.

By providing a number representing the number of copies of a chromosome,
and a number representing the number of groups (which at a minimum must be 2).
This method of `suggest_alleles` will randomly allocate the `ncopies` of the
chromosome into the `ngroups` groups.

For example, if you ask for a suggested allele pattern for 3 copies of a chromosome,
and 2 groups, this method might return a suggested pattern of `[DNA_G, DNA_A, DNA_G]`.

!!! note
    Which base corresponds to which group number will be determined at random.
"""
function suggest_alleles(ncopies::Int, ngroups::Int)
    @assert ngroups ≤ ncopies
    @assert ngroups > 1
    groups = Random.shuffle!(collect(Iterators.take(Iterators.cycle(Random.shuffle!(collect(Base.OneTo(ngroups)))), ncopies)))
    return suggest_alleles(groups)
end

"""
    suggest_alleles(npositions::Int, args...)

Suggest allele patterns for numerous heterozygous sites.

This method of `suggest_alleles` calls another `suggest_alleles` with `args...`
`npositions` times, and returns a vector of the results of each call.
"""
function suggest_alleles(npositions::Int, args...)
    return [suggest_alleles(args...) for _ in 1:npositions]
end

### Position and allele argument preparations, used to adapt the generic `plan_het` method's behaviour.

_prepare_het_positions(cb::ChromosomeBlueprint, pos::Int) = _prepare_het_positions(cb, suggest_regions(cb, 1, pos))
_prepare_het_positions(cb::ChromosomeBlueprint, pos::Float64) = _prepare_het_positions(cb, Int(floor(length(cb) * pos)))
_prepare_het_positions(cb::ChromosomeBlueprint, pos::Vector{Int}) = pos
function _prepare_het_positions(cb::ChromosomeBlueprint, pos::Vector{UnitRange{Int}})
    for p in pos
        @assert length(p) == 1
    end
    return [first(p) for p in pos]
end

_prepare_alleles(cb::ChromosomeBlueprint, npos::Int, x::Vector{Vector{DNA}}) = x
_prepare_alleles(cb::ChromosomeBlueprint, npos::Int, ngroups::Int) = suggest_alleles(npos, ncopies(cb), ngroups)
_prepare_alleles(cb::ChromosomeBlueprint, npos::Int, args...) = suggest_alleles(npos, args...)

"""
    plan_het(cb::ChromosomeBlueprint, pos, alle...)

A generic method of `plan_het`.

Creates a new chromosome blueprint, based on the input blueprint `cb`.

The `pos` argument determines which sites in a chromosome are made heterozygous,
and depending on the type of argument provided, the behaviour differs slightly:

- If `pos` is an integer, then that many available sites are selected at random
  to be heterozygous.

- If `pos` is a float, then it is treated as a proportion of the length of the,
  chromosome, which is converted into an integer, and that many sites are
  selected at random to be heterozygous.

- If `pos` is a vector of integers, it is trated as a list of positions the user
  wants to be heterozygous.

The `alle...` argument is a vararg argument which defines the how bases will be
allocated at each heterozygous site.

Depending on the type of the argument provided, the behaviour differs slightly,
based on the different methods of [`suggest_alleles`](@ref):

- If `alle...` is a single integer value. It is taken to mean the number of states
  each heterozygous site has. At a minumum this number must be two, as a site with
  only one state is not heterozygous by definition. For each site, which chromosome
  copies get which state is determined at random.

- If `alle...` is a vector of integers (one integer for every chromosome copy),
  then the nth value in the vector dictates the group the nth chromosome copy
  belongs to. For example for a triploid, the vector `[2, 1, 2]` means that
  there are two groups (group 1 and group 2), and that the second chromosome
  copy of the three belongs to group 1, and the other two copies belong to group
  2. Chromosome copies that belong to the same group will get the same base at
  each heterozygous site. In our example, copies 1 and 3 will get the same base,
  as they have been given the same group number.
  Chromosome copy 2 will end up with a different base. Which base corresponds to
  which group number will be determined at random.

- If `alle...` is multiple vectors of integers, then each vector constitutes a
  group, and each value in the vector determines which copies are in that group.
  It is alternative way to giving the same information as you can with a
  single vector (as above). E.g. consider the example of `[2, 1, 2]` from the
  previous point. This can be expressed as two vectors of `[2] and [1, 3]`:
  The first vector contains a number indicating the second chromosome copy, and
  the second vector contains numbers representing the first and third chromosome
  copy.

- If `alle...` is a vector of nucleotide vectors, then the nth vector of
  nucleotides determines which chromosome copy recieves which base at the nth
  heterozygous position. For example, for a triploid, `[[DNA_A, DNA_A, DNA_T], [DNA_G, DNA_C, DNA_C]]`
  would mean that at the first heterozygous site, the first and second copies
  of the chromosome would recieve an `DNA_A` base, and the third copy would
  recieve a `DNA_T` base. For the second heterozygous position, the first copy
  of the chromosome would recieve a DNA_G base, and the other two copies will
  recieve a `DNA_C` base.

(See also: [`plan_repetition`](@ref))
"""
function plan_het(cb::ChromosomeBlueprint, pos, alle...)
    positions = _prepare_het_positions(cb, pos)
    for p in positions
        if !(1 ≤ p ≤ length(cb))
            throw(ArgumentError(string("Position ", p, " is outside of the limits of the chromosome (", 1:length(cb), ')')))
        end
    end
    alleles = _prepare_alleles(cb, length(positions), alle...)

    @assert length(positions) == length(alleles)
    lastcb = cb
    for (p, a) in zip(positions, alleles)
        lastcb = plan_het(lastcb, p, a)
    end
    return lastcb
end


text:
module REPLMode

#using Markdown, UUIDs, Dates

import REPL
import REPL: LineEdit, REPLCompletions

import AKCalc

struct Option
end

struct Arg
end

struct Command
end

struct State
  gamedata::Union{AKCalc.GameData.Source, Nothing}
  playerdata::Union{AKCalc.PlayerData.Source, Nothing}
  statsdata::Union{AKCalc.StatsData.Source, Nothing}
end
State() = State(nothing, nothing, nothing)

function do_cmd(repl::REPL.AbstractREPL, input::String; do_rethrow=false)
  @assert false "Unimplemented"
end

function do_cmd!(repl::REPL.AbstractREPL, input::String; do_rethrow=false)
  @assert false "Unimplemented"
end

# Provide a string macro ak"cmd" that can be used in the same way
# as the REPLMode `ak> cmd`. Useful for testing and in environments
# where we do not have a REPL, e.g. IJulia.
struct MiniREPL <: REPL.AbstractREPL
  display::TextDisplay
  t::REPL.Terminals.TTYTerminal
end
function MiniREPL()
  env = get(ENV, "TERM", Sys.iswindows() ? "" : "dumb")
  MiniREPL(TextDisplay(stdout), REPL.Terminals.TTYTerminal(env, stdin, stdout, stderr))
end
REPL.REPLDisplay(repl::MiniREPL) = repl.display

const minirepl = Ref{MiniREPL}()

__init__() = minirepl[] = MiniREPL()

macro ak_str(str::String)
  :($(do_cmd)(minirepl[], $str; do_rethrow=true))
end

akstr(str::String) = do_cmd(minirepl[], str; do_rethrow=true)

function promptf()
  return "ak> "
end

function create_mode(repl::REPL.AbstractREPL, main::LineEdit.Prompt)
  ak_mode = LineEdit.Prompt(promptf;
    prompt_prefix = repl.options.hascolor ? Base.text_colors[:light_blue] : "",
    prompt_suffix = "",
    # complete = PkgCompletionProvider(),
    sticky = true)

  ak_mode.repl = repl
  hp = main.hist
  hp.mode_mapping[:ak] = ak_mode
  ak_mode.hist = hp

  search_prompt, skeymap = LineEdit.setup_search_keymap(hp)
  prefix_prompt, prefix_keymap = LineEdit.setup_prefix_keymap(hp, ak_mode)

  ak_mode.on_done = (s, buf, ok) -> begin
    ok || return REPL.transition(s, :abort)
    input = String(take!(buf))
    REPL.reset(repl)
    do_cmd(repl, input, do_rethrow=true)
    REPL.prepare_next(repl)
    REPL.reset_state(s)
    s.current_mode.sticky || REPL.transition(s, main)
  end

  mk = REPL.mode_keymap(main)

  b = Dict{Any,Any}[
    skeymap, mk, prefix_keymap, LineEdit.history_keymap,
    LineEdit.default_keymap, LineEdit.escape_defaults
  ]
  ak_mode.keymap_dict = LineEdit.keymap(b)

  return ak_mode
end

function repl_init(repl::REPL.AbstractREPL)
  main_mode = repl.interface.modes[1]
  ak_mode = create_mode(repl, main_mode)
  push!(repl.interface.modes, ak_mode)
  keymap = Dict{Any,Any}(
    '#' => function (s,args...)
      if isempty(s) || position(LineEdit.buffer(s)) == 0
        buf = copy(LineEdit.buffer(s))
        LineEdit.transition(s, ak_mode) do
          LineEdit.state(s, ak_mode).input_buffer = buf
        end
      else
        LineEdit.edit_insert(s, '#')
      end
    end
  )
  main_mode.keymap_dict = LineEdit.keymap_merge(main_mode.keymap_dict, keymap)
  return
end

end


text:
"""
$(TYPEDSIGNATURES)

Edit region numbers of grid cells via rectangular mask.
"""
function cellmask!(grid::ExtendableGrid,
                   maskmin::AbstractArray,
                   maskmax::AbstractArray,
                   ireg::Int;
                   tol=1.0e-10)
    xmaskmin=maskmin.-tol
    xmaskmax=maskmax.+tol
    ncells=num_cells(grid)
    cellnodes=grid[CellNodes]
    dim=dim_space(grid)
    cellregions=grid[CellRegions]
    coord=grid[Coordinates]
    for icell=1:ncells
        in_region=true
        for inode=1:num_targets(cellnodes,icell)
            ignode=cellnodes[inode,icell]
            for idim=1:dim
                if coord[idim,ignode]<xmaskmin[idim]
                    in_region=false
                elseif coord[idim,ignode]>xmaskmax[idim]
                    in_region=false
                end
            end
        end
        if in_region
            cellregions[icell]=ireg
        end
    end
    grid[NumCellRegions]=max(num_cellregions(grid),ireg)
end

"""
$(TYPEDSIGNATURES)

Edit region numbers of grid  boundary facets  via rectangular mask.
For 1D grids, inner boundaries can be added by this method.
"""
function bfacemask!(grid::ExtendableGrid,
                    maskmin::AbstractArray,
                    maskmax::AbstractArray,
                    ireg::Int;
                    tol=1.0e-10)

    xmaskmin=maskmin.-tol
    xmaskmax=maskmax.+tol

    nbfaces=num_bfaces(grid)
    bfacenodes=grid[BFaceNodes]
    Ti=eltype(bfacenodes)
    dim=dim_space(grid)
    bfaceregions=grid[BFaceRegions]
    new_bfacenodes=ElasticArray{Ti,2}(bfacenodes)
    coord=grid[Coordinates]

    function isbface(ix)
        for ibface=1:num_bfaces(grid)
            if bfacenodes[1,ibface]==ix
                return ibface
            end
        end
        return 0
    end

    function isbface(ix,iy)
        for ibface=1:num_bfaces(grid)
            if (bfacenodes[1,ibface] == ix && bfacenodes[2,ibface] == iy) ||
                (bfacenodes[1,ibface] == iy && bfacenodes[2,ibface] == ix)
                return ibface
            end
        end
        return 0
    end

    if dim_space(grid)==1
        for inode=1:num_nodes(grid)
            x=coord[1,inode]
            if x>xmaskmin[1] && x<xmaskmax[1]
                ibface=isbface(inode)
                if ibface>0
                    bfaceregions[ibface]=ireg
                else
                    push!(bfaceregions,ireg)
                    append!(new_bfacenodes,[inode])
                end
            end
        end
    else
        edgenodes=grid[EdgeNodes]
        for iedge=1:size(edgenodes,2)
            in_region=true
            for inode=1:num_targets(edgenodes,iedge)
                ignode=edgenodes[inode,iedge]
                for idim=1:dim_space(grid)
                    if coord[idim,ignode]<xmaskmin[idim]
                        in_region=false
                    elseif coord[idim,ignode]>xmaskmax[idim]
                        in_region=false
                    end
                end
            end
            if in_region
                ibface=isbface(edgenodes[1,iedge],edgenodes[2,iedge])
                if ibface>0
                    bfaceregions[ibface]=ireg
                else
                    push!(bfaceregions,ireg)
                    append!(new_bfacenodes,[edgenodes[1,iedge],edgenodes[2,iedge]])
                end
            end
        end
    end
    btype=grid[BFaceGeometries][1]
    grid[BFaceNodes]=Array{Ti,2}(new_bfacenodes)
    grid[BFaceGeometries]=VectorOfConstants(btype,length(bfaceregions))
    grid[NumBFaceRegions]=max(num_bfaceregions(grid),ireg)
    return grid
end


text:
using EasyPhys
using Base.Test

test_files = ["utilities.jl", "fitter.jl", "plot.jl"]

println("Running tests:")

for test in test_files
    println("Running tests on $(test)...")
    include(test)
end


text:
module Higher

include("rcd.jl")
include("rid.jl")

end

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
module Random
  export seed!
  function seed!(n)
    return MersenneTwister(n)
  end
end

max_stars_repo_name: tagordon/limbdark

max_issues_repo_name: tagordon/limbdark

max_forks_repo_name: tagordon/limbdark


text:
const ∞ = Inf

LightDirection() = Vector3(1, 1, 1) |> normalize

LightColor() = White


text:
module Cthulhu

using CodeTracking: definition
using InteractiveUtils

using Core: MethodInstance
const Compiler = Core.Compiler

Base.@kwdef mutable struct CthulhuConfig
    enable_highlighter::Bool = false
    highlighter::Cmd = `pygmentize -l`
    asm_syntax::Symbol = :att
    dead_code_elimination::Bool = true
    pretty_ast::Bool = false
end

"""
    Cthulhu.CONFIG

# Options
- `enable_highlighter::Bool`: Use command line `highlighter` to syntax highlight
  LLVM and native code.  Set to `true` if `highlighter` exists at the import
  time.
- `highlighter::Cmd`: A command line program that receives "llvm" or "asm" as
  an argument and the code as stdin.  Defaults to `$(CthulhuConfig().highlighter)`.
- `asm_syntax::Symbol`: Set the syntax of assembly code being used.
  Defaults to `$(CthulhuConfig().asm_syntax)`.
- `dead_code_elimination::Bool`: Enable dead-code elimination for high-level Julia IR.
  Defaults to `true`. DCE is known to be buggy and you may want to disable it if you
  encounter errors. Please report such bugs with a MWE to Julia or Cthulhu.
- `pretty_ast::Bool`: Use a pretty printer for the ast dump. Defaults to false.
"""
const CONFIG = CthulhuConfig()

include("callsite.jl")
include("reflection.jl")
include("ui.jl")
include("codeview.jl")

export descend, @descend, descend_code_typed, descend_code_warntype, @descend_code_typed, @descend_code_warntype

"""
    @descend_code_typed

Evaluates the arguments to the function or macro call, determines their
types, and calls `code_typed` on the resulting expression.
"""
macro descend_code_typed(ex0...)
    InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :descend_code_typed, ex0)
end

"""
    @descend_code_warntype

Evaluates the arguments to the function or macro call, determines their
types, and calls `code_warntype` on the resulting expression.
"""
macro descend_code_warntype(ex0...)
    InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :descend_code_warntype, ex0)
end

"""
    @descend

Shortcut for [`@descend_code_typed`](@ref).
"""
macro descend(ex0...)
    esc(:(@descend_code_typed($(ex0...))))
end

"""
    descend_code_typed(f, tt; kwargs...)

Given a function and a tuple-type, interactively explore the output of
`code_typed` by descending into `invoke` statements. Type enter to select an
`invoke` to descend into, select ↩  to ascend, and press q or control-c to
quit.

# Usage:
```julia
function foo()
    T = rand() > 0.5 ? Int64 : Float64
    sum(rand(T, 100))
end

descend_code_typed(foo, Tuple{})
```
"""
descend_code_typed(f, @nospecialize(tt); kwargs...) =
    _descend_with_error_handling(f, tt; iswarn=false, kwargs...)

"""
    descend_code_warntype(f, tt)

Given a function and a tuple-type, interactively explore the output of
`code_warntype` by descending into `invoke` statements. Type enter to select an
`invoke` to descend into, select ↩  to ascend, and press q or control-c to
quit.

# Usage:
```julia
function foo()
    T = rand() > 0.5 ? Int64 : Float64
    sum(rand(T, 100))
end

descend_code_warntype(foo, Tuple{})
```
"""
descend_code_warntype(f, @nospecialize(tt); kwargs...) =
    _descend_with_error_handling(f, tt; iswarn=true, kwargs...)

function _descend_with_error_handling(f, @nospecialize(tt); kwargs...)
    try
        _descend(f, tt; kwargs...)
    catch x
        if x isa InterruptException
            return nothing
        else
            rethrow(x)
        end
    end
    return nothing
end

"""
    descend

Shortcut for [`descend_code_typed`](@ref).
"""
const descend = descend_code_typed

##
# _descend is the main driver function.
# src/reflection.jl has the tools to discover methods
# src/ui.jl provides the user facing interface to which _descend responds
##
function _descend(mi::MethodInstance; iswarn::Bool, params=current_params(), optimize::Bool=true, kwargs...)
    debuginfo = true
    if :debuginfo in keys(kwargs)
        selected = kwargs[:debuginfo]
        # TODO: respect default
        debuginfo = selected == :source
    end

    display_CI = true
    while true
        debuginfo_key = debuginfo ? :source : :none
        (CI, rt, slottypes) = do_typeinf_slottypes(mi, optimize, params)
        preprocess_ci!(CI, mi, optimize, CONFIG)
        callsites = find_callsites(CI, mi, slottypes; params=params, kwargs...)

        display_CI && cthulu_typed(stdout, debuginfo_key, CI, rt, mi, iswarn)
        display_CI = true

        TerminalMenus.config(cursor = '•', scroll = :wrap)
        menu = CthulhuMenu(callsites)
        cid = request(menu)
        toggle = menu.toggle

        if toggle === nothing
            if cid == length(callsites) + 1
                break
            end
            if cid == -1
                throw(InterruptException())
            end
            callsite = callsites[cid]

            if callsite.info isa MultiCallInfo
                sub_callsites = map(ci->Callsite(callsite.id, ci), callsite.info.callinfos)
                if isempty(sub_callsites)
                    @warn "Expected multiple callsites, but found none. Please fill an issue with a reproducing example" callsite.info
                    continue
                end
                menu = CthulhuMenu(sub_callsites, sub_menu=true)
                cid = request(menu)
                if cid == length(sub_callsites) + 1
                    continue
                end
                if cid == -1
                    throw(InterruptException())
                end
                callsite = sub_callsites[cid]
            end

            if callsite.info isa GeneratedCallInfo || callsite.info isa FailedCallInfo
                @error "Callsite %$(callsite.id) failed to be extracted" callsite
            end

            # recurse
            next_mi = get_mi(callsite)
            if next_mi === nothing
                continue
            end

            _descend(next_mi; params=params, optimize=optimize,
                     iswarn=iswarn, debuginfo=debuginfo_key, kwargs...)

        elseif toggle === :warn
            iswarn ⊻= true
        elseif toggle === :optimize
            optimize ⊻= true
        elseif toggle === :debuginfo
            debuginfo ⊻= true
        elseif toggle === :highlighter
            CONFIG.enable_highlighter ⊻= true
            if CONFIG.enable_highlighter
                @info "Using syntax highlighter $(CONFIG.highlighter)"
            else
                @info "Turned off syntax highlighter for LLVM and native code."
            end
            display_CI = false
        elseif toggle === :dump_params
            @info "Dumping inference cache"
            Core.show(map(((i, x),) -> (i, x.result, x.linfo), enumerate(params.cache)))
            Core.println()
            display_CI = false
        else
            #Handle Standard alternative view, e.g. :native, :llvm
            view_cmd = get(codeviews, toggle, nothing)
            if view_cmd !== nothing
                view_cmd(stdout, mi, optimize, debuginfo, params, CONFIG)
                display_CI = false
            else
                error("Unknown option $toggle")
            end
        end
    end
end

function _descend(@nospecialize(F), @nospecialize(TT); params=current_params(), kwargs...)
    mi = first_method_instance(F, TT; params=params)
    _descend(mi; params=params, kwargs...)
end

end


text:
using Distributions

export acceptance_prob

abstract type InvMcmcMove end

acceptance_prob(move::InvMcmcMove) = move.counts[1]/move.counts[2]

function reset_counts!(mcmc::InvMcmcMove)
    mcmc.counts .= 0
end

# Any move must have the following
# prop_sample!(S_curr,S_prop,move) (returs log ratio)
# accept_reject!(S_curr,S_prop,model,move) (returns true false)
# counts field

# The most general enact accept/reject functions
# These can be specialised for a move if perhaps one can work out
# how to make the required changes more efficiently
function enact_accept!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    pointers::InteractionSequence{Int},
    move::InvMcmcMove
    )
    ε = length(S_prop) - length(S_curr)
    if ε > 0
        # We've increased dimension
        # Bring in pointers
        for i in 1:ε
            tmp = pop!(pointers)
            push!(S_curr,tmp)
        end
    else
        # We've decreased dimension
        # Send pointers back
        for i in 1:abs(ε)
            tmp = pop!(S_curr)
            push!(pointers, tmp)
        end
    end
    # Copy across from S_prop
    for (x,y) in zip(S_curr,S_prop)
        copy!(x, y)
    end
end

function enact_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    pointers::InteractionSequence{Int},
    move::InvMcmcMove
    )

    ε = length(S_prop) - length(S_curr)
    if ε > 0
        # We've increased dimension
        # Send back pointers
        for i in 1:ε
            tmp = pop!(S_prop)
            push!(pointers, tmp)
        end
    else
        # We've decreased dimension
        # Bring in pointers
        for i in 1:abs(ε)
            tmp = pop!(pointers)
            push!(S_prop, tmp)
        end
    end
    # Copy across from S_curr
    for (x,y) in zip(S_prop,S_curr)
        copy!(x, y)
    end
end

# Most general accept reject function


text:
# In total we can have trend, slope, seasonal, cycle and exogenous
# Maybe autoregressive
function validate_trend(trend::String)
    @assert trend in [
        "local level",
        "random walk",
        # "deterministic trend",
        # "linear deterministic trend",
        # "linear trend",
        "local linear trend",
        # "damped local lineaar trend",
        # "random walk with drift"
        "smooth trend"
    ]
    return true
end
function parse_trend(trend::String)
    validate_trend(trend)
    (has_irregular, has_trend, stochastic_trend,
    has_slope, stochastic_slope) = (false, false, false, false, false)
    if trend == "local level"
        (has_irregular, has_trend, stochastic_trend,
        has_slope, stochastic_slope) = (true, true, true, false, false)
    elseif trend == "random walk"
        (has_irregular, has_trend, stochastic_trend,
        has_slope, stochastic_slope) = (false, true, true, false, false)
    elseif trend == "local linear trend"
        (has_irregular, has_trend, stochastic_trend,
        has_slope, stochastic_slope) = (true, true, true, true, true)
    elseif trend == "smooth trend"
        (has_irregular, has_trend, stochastic_trend,
        has_slope, stochastic_slope) = (true, true, false, true, true)
    end
    # Validate trend booleans
    if !has_trend && stochastic_trend
        error("Invalid trend specification.")
    end
    if !has_trend && has_slope
        error("Invalid trend specification.")
    end
    if !has_slope && stochastic_slope
        error("Invalid trend specification.")
    end
    return (has_irregular, has_trend, stochastic_trend, has_slope, stochastic_slope)
end
function validate_seasonal(seasonal::String)
    if seasonal == "no"
        return true
    else
        spl = split(seasonal)
        # TODO better error messaage
        # Maybe a no string can be also valid
        @assert length(spl) == 2
        @assert spl[1] in ["deterministic", "stochastic"]
        return true
    end
end
function parse_seasonal(seasonal::String)
    validate_seasonal(seasonal)
    if seasonal == "no"
        has_seasonal, stochastic_seasonal, seasonal_freq = (false, false, 0)
        return has_seasonal, stochastic_seasonal, seasonal_freq
    else
        spl = split(seasonal)
        stochastic_seasonal = spl[1] == "stochastic"
        seasonal_freq = parse(Int, spl[2])
        has_seasonal = seasonal_freq == 0 ? false : true
        return has_seasonal, stochastic_seasonal, seasonal_freq
    end
end
function validate_cycle(cycle::String)
    if cycle == "no"
        return true
    else
        spl = split(cycle)
        # TODO better error messaage
        if length(spl) == 1
            @assert spl[1] in ["deterministic", "stochastic"]
            return true
        elseif length(spl) == 2
            @assert spl[1] in ["deterministic", "stochastic"]
            @assert spl[2] == "damped"
            return true
        end
    end
    return false
end
function parse_cycle(cycle::String)
    validate_cycle(cycle)
    if cycle == "no"
        has_cycle, stochastic_cycle, damped_cycle = (false, false, false)
        return has_cycle, stochastic_cycle, damped_cycle
    else
        spl = split(cycle)
        has_cycle = true
        # TODO better error messaage
        # Maybe a no string can be also valid
        if length(spl) == 1
            stochastic_cycle = spl[1] == "stochastic"
            return has_cycle, stochastic_cycle, false
        elseif length(spl) == 2
            stochastic_cycle = spl[1] == "stochastic"
            damped_cycle = spl[2] == "damped"
            return has_cycle, stochastic_cycle, damped_cycle
        end
    end
end

@doc raw"""
    UnobservedComponents(
        y::Vector{Fl};
        trend::String = "local level",
        seasonal::String = "no"
        cycle::String = "no"
    ) where Fl

An unobserved components model that can have trend/level, seasonal and cycle components.
Each component should be specified by strings, if the component is not desired in the model
a string with "no" can be passed as keyword argument.

These models take the general form

```math
\begin{gather*}
    \begin{aligned}
    y_t = \mu_t + \gamma_t + c_t + \varepsilon_t
    \end{aligned}
\end{gather*}
```
where ``y_t`` refers to the observation vector at time ``t``,
``\mu_t`` refers to the trend component, ``\gamma_t`` refers to the
seasonal component, ``c_t`` refers to the cycle, and
``\varepsilon_t`` is the irregular. The modeling details of these
components are given below.

### **Trend**

The trend component can be modeled in a lot of different ways, usually it is called level when
there is no slope component. The modelling options can be expressed as in the example `trend = "local level"`.

* Local Level
string: `"local level"`
```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  \mu_{t} + \varepsilon_{t} \quad \varepsilon_{t} \sim \mathcal{N}(0, \sigma^2_{\varepsilon})\\
        \mu_{t+1} &= \mu_{t} + \eta_{t} \quad \eta_{t} \sim \mathcal{N}(0, \sigma^2_{\eta})\\
    \end{aligned}
\end{gather*}
```

* Random Walk
string: `"random walk"`
```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  \mu_{t}\\
        \mu_{t+1} &= \mu_{t} + \eta_{t} \quad \eta_{t} \sim \mathcal{N}(0, \sigma^2_{\eta})\\
    \end{aligned}
\end{gather*}
```

* Local Linear Trend
string: `"local linear trend"`
```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  \mu_{t} + \gamma_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, \sigma^2_{\varepsilon})\\
        \mu_{t+1} &= \mu_{t} + \nu_{t} + \xi_{t} \quad &\xi_{t} \sim \mathcal{N}(0, \sigma^2_{\xi})\\
        \nu_{t+1} &= \nu_{t} + \zeta_{t} \quad &\zeta_{t} \sim \mathcal{N}(0, \sigma^2_{\zeta})\\
    \end{aligned}
\end{gather*}
```

* Smooth Trend
string: `"smooth trend"`
```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  \mu_{t} + \gamma_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, \sigma^2_{\varepsilon})\\
        \mu_{t+1} &= \mu_{t} + \nu_{t}\\
        \nu_{t+1} &= \nu_{t} + \zeta_{t} \quad &\zeta_{t} \sim \mathcal{N}(0, \sigma^2_{\zeta})\\
    \end{aligned}
\end{gather*}
```

**Seasonal**

The seasonal component is modeled as:

```math
\begin{gather*}
    \begin{aligned}
    \gamma_t = - \sum_{j=1}^{s-1} \gamma_{t+1-j} + \omega_t \quad \omega_t \sim N(0, \sigma^2_\omega)
    \end{aligned}
\end{gather*}
```

The periodicity (number of seasons) is s, and the defining character is
that (without the error term), the seasonal components sum to zero across
one complete cycle. The inclusion of an error term allows the seasonal
effects to vary over time. The modelling options can be expressed in terms
of `"deterministic"` or `"stochastic"` and the periodicity as a number in
the string, i.e., `seasonal = "stochastic 12"`.

**Cycle**

The cycle component is modeled as

```math
\begin{gather*}
    \begin{aligned}
        c_{t+1} &= \rho_c \left(c_{t} \cos(\lambda_c) + c_{t}^{*} \sin(\lambda_c)\right) \quad & \tilde\omega_{t} \sim \mathcal{N}(0, \sigma^2_{\tilde\omega})\\
        c_{t+1}^{*} &= \rho_c \left(-c_{t} \sin(\lambda_c) + c_{t}^{*} \sin(\lambda_c)\right) \quad &\tilde\omega^*_{t} \sim \mathcal{N}(0, \sigma^2_{\tilde\omega})\\
    \end{aligned}
\end{gather*}
```

The cyclical component is intended to capture cyclical effects at time frames much longer
than captured by the seasonal component. The parameter ``\lambda_c`` is the frequency of the cycle
and it is estimated via maximum likelihood. The inclusion of error terms allows the cycle
effects to vary over time. The modelling options can be expressed in terms
of `"deterministic"` or `"stochastic"` and the damping effect as a string, i.e.,
`cycle = "stochastic"`, `cycle = "deterministic"` or `cycle = "stochastic damped"`.

The UnobservedComponents model has some dedicated Plot Recipes, see [Visualization](@ref)

# References
 * Durbin, James, & Siem Jan Koopman.
   Time Series Analysis by State Space Methods: Second Edition.
   Oxford University Press, 2012
"""
mutable struct UnobservedComponents <: StateSpaceModel
    hyperparameters::HyperParameters
    trend::String
    seasonal::String
    cycle::String
    has_irregular::Bool
    has_trend::Bool
    stochastic_trend::Bool
    has_slope::Bool
    stochastic_slope::Bool
    # damped_slope::Bool TODO
    has_seasonal::Bool
    stochastic_seasonal::Bool
    seasonal_freq::Int
    has_cycle::Bool
    stochastic_cycle::Bool
    damped_cycle::Bool
    system::LinearUnivariateTimeInvariant
    results::Results

    function UnobservedComponents(y::Vector{Fl};
                                  trend::String = "local level",
                                  seasonal::String = "no", # for example "deterministic 3" or "stochastic 12"
                                  cycle::String = "no" # for example "deterministic damped" or "stochastic"
                                  ) where Fl

        (has_irregular, has_trend, stochastic_trend,
        has_slope, stochastic_slope) = parse_trend(trend)
        (has_seasonal, stochastic_seasonal, seasonal_freq) = parse_seasonal(seasonal)
        (has_cycle, stochastic_cycle, damped_cycle) = parse_cycle(cycle)
        # Define system matrices
        Z = build_Z(Fl, has_trend, has_slope, has_seasonal, seasonal_freq, has_cycle)
        T = build_T(Fl, has_trend, has_slope, has_seasonal, seasonal_freq, has_cycle)
        R = build_R(Fl, has_trend, has_slope, has_seasonal, seasonal_freq, has_cycle,
                        stochastic_trend, stochastic_slope, stochastic_seasonal, stochastic_cycle)
        d = zero(Fl)
        c = build_c(T)
        H = zero(Fl)
        Q = build_Q(R)

        system = LinearUnivariateTimeInvariant{Fl}(y, Z, T, R, d, c, H, Q)

        # Define hyperparameters names
        names = build_names(has_irregular, stochastic_trend, stochastic_slope,
                            stochastic_seasonal, has_cycle, stochastic_cycle, damped_cycle)
        hyperparameters = HyperParameters{Fl}(names)

        return new(hyperparameters, trend, seasonal, cycle,
                    has_irregular, has_trend, stochastic_trend,
                    has_slope, stochastic_slope, has_seasonal,
                    stochastic_seasonal, seasonal_freq,
                    has_cycle, stochastic_cycle, damped_cycle,
                    system, Results{Fl}())
    end
end

num_components(model::UnobservedComponents) = (model.has_trend + model.has_slope + model.has_seasonal +
                                                model.has_cycle)
function dict_components(model::UnobservedComponents)
    dict_components = OrderedDict{String, Int}()
    i = 1
    if model.has_trend
        dict_components["Trend"] = i
        i += 1
    end
    if model.has_slope
        dict_components["Slope"] = i
        i += 1
    end
    if model.has_seasonal
        dict_components["Seasonal"] = i
        i += model.seasonal_freq - 1
    end
    if model.has_cycle
        dict_components["Cycle"] = i
        i += model.has_cycle + 1
    end
    return dict_components
end

function build_Z(Fl::DataType,
                 has_trend::Bool,
                 has_slope::Bool,
                 has_seasonal::Bool,
                 seasonal_freq::Int,
                 has_cycle::Bool)
    Z = Fl[]
    if has_trend
        Z = vcat(Z, one(Fl))
    end
    if has_slope
        Z = vcat(Z, zero(Fl))
    end
    if has_seasonal
        Z = vcat(Z, [1; zeros(Fl, seasonal_freq - 2)])
    end
    if has_cycle
        Z = vcat(Z, [one(Fl); zero(Fl)])
    end
    return Z
end
function build_T(Fl::DataType,
                 has_trend::Bool,
                 has_slope::Bool,
                 has_seasonal::Bool,
                 seasonal_freq::Int,
                 has_cycle::Bool)
    # Caalculate how maany states for each component
    num_trend_states = has_trend + has_slope
    num_seasonal_states = has_seasonal ? seasonal_freq - 1 : 0
    num_cycle_states = 2 * has_cycle
    num_states = num_trend_states + num_seasonal_states + num_cycle_states
    T = zeros(Fl, num_states, num_states)
    if has_trend
        T[1, 1] = one(Fl)
    end
    if has_slope
        T[1, 2] = one(Fl)
        T[2, 2] = one(Fl)
    end
    if has_seasonal
        rows = num_trend_states + 1
        cols = num_trend_states + 1:num_trend_states + num_seasonal_states
        T[rows, cols] .= -one(Fl)
        rows = num_trend_states + 2:num_trend_states + num_seasonal_states
        cols = num_trend_states + 1:num_trend_states + num_seasonal_states - 1
        T[rows, cols] = Matrix{Fl}(I, seasonal_freq - 2, seasonal_freq - 2)
    end
    if has_cycle
        # Do nothing T should be filled with the cycle expression
    end
    return T
end
function build_R(Fl::DataType,
                 has_trend::Bool,
                 has_slope::Bool,
                 has_seasonal::Bool,
                 seasonal_freq::Int,
                 has_cycle::Bool,
                 stochastic_trend::Bool,
                 stochastic_slope::Bool,
                 stochastic_seasonal::Bool,
                 stochastic_cycle::Bool)
    # Assign some model properties
    num_trend_states = has_trend + has_slope
    num_stochastic_trend = stochastic_trend + stochastic_slope
    num_seasonal_states = has_seasonal ? seasonal_freq - 1 : 0
    num_cycle_states = 2 * has_cycle
    num_states = num_trend_states + num_seasonal_states + num_cycle_states
    num_stochastic_states = num_stochastic_trend + stochastic_seasonal + 2 * stochastic_cycle

    R = zeros(num_states, num_stochastic_states)
    for i in 1:num_stochastic_states
        R[i, i] = one(Fl)
    end
    return R
end
function build_c(T::Matrix{Fl}) where Fl
    m = size(T, 1)
    return zeros(Fl, m)
end
function build_Q(R::Matrix{Fl}) where Fl
    r = size(R, 2)
    return Matrix{Fl}(I, r, r)
end
function build_names(has_irregular::Bool,
                     stochastic_trend::Bool,
                     stochastic_slope::Bool,
                     stochastic_seasonal::Bool,
                     has_cycle::Bool,
                     stochastic_cycle::Bool,
                     damped_cycle::Bool)
    names = String[]
    if has_irregular
        push!(names, "sigma2_irregular")
    end
    if stochastic_trend
        push!(names, "sigma2_trend")
    end
    if stochastic_slope
        push!(names, "sigma2_slope")
    end
    if stochastic_seasonal
        push!(names, "sigma2_seasonal")
    end
    if has_cycle
        push!(names, "λ_cycle")
        if stochastic_cycle
            push!(names, "sigma2_cycle")
        end
        if damped_cycle
            push!(names, "ρ_cycle")
        end
    end
    return names
end

has_sigma2(str::String) = occursin("sigma2", str)
function sigma2_names(model::UnobservedComponents)
    return filter(has_sigma2, get_names(model))
end

function diag_indices(k::Int)
    idx = Int[]
    l = 1
    for i in 1:k, j in 1:k
        if i == j
            push!(idx, l)
        end
        l += 1
    end
    return idx
end

# Obligatory methods
function default_filter(model::UnobservedComponents)
    Fl = typeof_model_elements(model)
    steadystate_tol = Fl(1e-5)
    a1 = zeros(Fl, num_states(model))
    P1 = Fl(1e6) .* Matrix{Fl}(I, num_states(model), num_states(model))
    return UnivariateKalmanFilter(a1, P1, num_states(model), steadystate_tol)
end

function initial_hyperparameters!(model::UnobservedComponents)
    Fl = typeof_model_elements(model)
    observed_variance = variance_of_valid_observations(model.system.y)
    # TODO add heuristic for initial hyperparameters
    initial_hyperparameters = Dict{String,Fl}(get_names(model) .=> one(Fl))
    if model.has_irregular
        initial_hyperparameters["sigma2_irregular"] = observed_variance
    end
    if model.has_trend
        initial_hyperparameters["sigma2_trend"] = observed_variance
    end
    if model.has_cycle
        initial_hyperparameters["λ_cycle"] = Fl(2 * pi / 12)
        if model.damped_cycle
            initial_hyperparameters["ρ_cycle"] = Fl(0.7)
        end
    end
    set_initial_hyperparameters!(model, initial_hyperparameters)
    return model
end

function constrain_hyperparameters!(model::UnobservedComponents)
    Fl = typeof_model_elements(model)
    for variance_name in sigma2_names(model)
        constrain_variance!(model, variance_name)
    end
    if model.has_cycle
        # Durbin and Koopman (2012) comment possible values in their book pp. 48
        constrain_box!(model, "λ_cycle", Fl(2 * pi / 100), Fl(2 * pi / 1.5))
        if model.damped_cycle
            constrain_box!(model, "ρ_cycle", zero(Fl), Fl(1.2))
        end
    end
    return model
end

function unconstrain_hyperparameters!(model::UnobservedComponents)
    Fl = typeof_model_elements(model)
    for variance_name in sigma2_names(model)
        unconstrain_variance!(model, variance_name)
    end
    if model.has_cycle
        # Durbin and Koopman (2012) comment possible values in their book pp. 48
        unconstrain_box!(model, "λ_cycle", Fl(2 * pi / 100), Fl(2 * pi / 1.5))
        if model.damped_cycle
            unconstrain_box!(model, "ρ_cycle", zero(Fl), Fl(1.2))
        end
    end
    return model
end

function fill_model_system!(model::UnobservedComponents)
    # TODO (performance) maybe cache this indices
    # This creates some extra allocations
    diag_Q_idx = diag_indices(size(model.system.Q, 1))
    idx = 1

    if model.has_irregular
        model.system.H = get_constrained_value(model, "sigma2_irregular")
    end
    if model.stochastic_trend
        model.system.Q[diag_Q_idx[idx]] = get_constrained_value(model, "sigma2_trend")
        idx += 1
    end
    if model.stochastic_slope
        model.system.Q[diag_Q_idx[idx]] = get_constrained_value(model, "sigma2_slope")
        idx += 1
    end
    if model.stochastic_seasonal
        model.system.Q[diag_Q_idx[idx]] = get_constrained_value(model, "sigma2_seasonal")
        idx += 1
    end
    if model.stochastic_cycle
        model.system.Q[diag_Q_idx[idx]] = get_constrained_value(model, "sigma2_cycle")
        idx += 1
        model.system.Q[diag_Q_idx[idx]] = get_constrained_value(model, "sigma2_cycle")
        idx += 1
    end

    if model.has_cycle
        num_trend_states = model.has_trend + model.has_slope
        num_seasonal_states = model.has_seasonal ? model.seasonal_freq - 1 : 0
        num_trend_seasonal_states = num_trend_states + num_seasonal_states
        cycle_rows = num_trend_seasonal_states+1:num_trend_seasonal_states+2
        cycle_cols = cycle_rows
        if model.damped_cycle
            ρ = get_constrained_value(model, "ρ_cycle")
            model.system.T[cycle_rows[1], cycle_rows[1]] = ρ * cos(get_constrained_value(model, "λ_cycle"))
            model.system.T[cycle_rows[2], cycle_rows[1]] = ρ * -sin(get_constrained_value(model, "λ_cycle"))
            model.system.T[cycle_rows[1], cycle_rows[2]] = ρ * sin(get_constrained_value(model, "λ_cycle"))
            model.system.T[cycle_rows[2], cycle_rows[2]] = ρ * cos(get_constrained_value(model, "λ_cycle"))
        else
            model.system.T[cycle_rows[1], cycle_rows[1]] = cos(get_constrained_value(model, "λ_cycle"))
            model.system.T[cycle_rows[2], cycle_rows[1]] = -sin(get_constrained_value(model, "λ_cycle"))
            model.system.T[cycle_rows[1], cycle_rows[2]] = sin(get_constrained_value(model, "λ_cycle"))
            model.system.T[cycle_rows[2], cycle_rows[2]] = cos(get_constrained_value(model, "λ_cycle"))
        end
    end

    return model
end

function reinstantiate(model::UnobservedComponents, y::Vector{Fl}) where Fl
    return UnobservedComponents(y;
                                trend = model.trend,
                                seasonal = model.seasonal,
                                cycle = model.cycle)
end

has_exogenous(::UnobservedComponents) = false