text:
@doc raw"""
    HyperbolicDiffusionEquations2D

The linear hyperbolic diffusion equations in two space dimensions.
A description of this system can be found in Sec. 2.5 of the book "I Do Like CFD, Too: Vol 1".
The book is freely available at http://www.cfdbooks.com/ and further analysis can be found in
the paper by Nishikawa [DOI: 10.1016/j.jcp.2007.07.029](https://doi.org/10.1016/j.jcp.2007.07.029)
"""
struct HyperbolicDiffusionEquations2D{RealT<:Real} <: AbstractHyperbolicDiffusionEquations{2, 3}
  Lr::RealT     # reference length scale
  inv_Tr::RealT # inverse of the reference time scale
  nu::RealT     # diffusion constant
end

function HyperbolicDiffusionEquations2D(; nu=1.0, Lr=inv(2pi))
  Tr = Lr^2 / nu
  HyperbolicDiffusionEquations2D(promote(Lr, inv(Tr), nu)...)
end

varnames(::typeof(cons2cons), ::HyperbolicDiffusionEquations2D) = ("phi", "q1", "q2")
varnames(::typeof(cons2prim), ::HyperbolicDiffusionEquations2D) = ("phi", "q1", "q2")
default_analysis_errors(::HyperbolicDiffusionEquations2D)     = (:l2_error, :linf_error, :residual)

@inline function residual_steady_state(du, ::HyperbolicDiffusionEquations2D)
  abs(du[1])
end

# Set initial conditions at physical location `x` for pseudo-time `t`
function initial_condition_poisson_periodic(x, t, equations::HyperbolicDiffusionEquations2D)
  # elliptic equation: -νΔϕ = f
  # depending on initial constant state, c, for phi this converges to the solution ϕ + c
  if iszero(t)
    phi = 0.0
    q1  = 0.0
    q2  = 0.0
  else
    phi = sin(2.0*pi*x[1])*sin(2.0*pi*x[2])
    q1  = 2*pi*cos(2.0*pi*x[1])*sin(2.0*pi*x[2])
    q2  = 2*pi*sin(2.0*pi*x[1])*cos(2.0*pi*x[2])
  end
  return SVector(phi, q1, q2)
end

@inline function source_terms_poisson_periodic(u, x, t, equations::HyperbolicDiffusionEquations2D)
  # elliptic equation: -νΔϕ = f
  # analytical solution: phi = sin(2πx)*sin(2πy) and f = -8νπ^2 sin(2πx)*sin(2πy)
  @unpack inv_Tr = equations
  C = -8 * equations.nu * pi^2

  x1, x2 = x
  tmp1 = sinpi(2 * x1)
  tmp2 = sinpi(2 * x2)
  du1 = -C*tmp1*tmp2
  du2 = -inv_Tr * u[2]
  du3 = -inv_Tr * u[3]

  return SVector(du1, du2, du3)
end

@inline function initial_condition_poisson_nonperiodic(x, t, equations::HyperbolicDiffusionEquations2D)
  # elliptic equation: -ν Δϕ = f in Ω, u = g on ∂Ω
  if iszero(t)
    T = eltype(x)
    phi = one(T)
    q1  = one(T)
    q2  = one(T)
  else
    # TODO: sincospi
    sinpi_x1,  cospi_x1  = sincos(pi*x[1])
    sinpi_2x2, cospi_2x2 = sincos(pi*2*x[2])
    phi =  2 *      cospi_x1 * sinpi_2x2 + 2 # ϕ
    q1  = -2 * pi * sinpi_x1 * sinpi_2x2     # ϕ_x
    q2  =  4 * pi * cospi_x1 * cospi_2x2     # ϕ_y
  end
  return SVector(phi, q1, q2)
end

@inline function source_terms_poisson_nonperiodic(u, x, t, equations::HyperbolicDiffusionEquations2D)
  # elliptic equation: -ν Δϕ = f in Ω, u = g on ∂Ω
  # analytical solution: ϕ = 2cos(πx)sin(2πy) + 2 and f = 10π^2cos(πx)sin(2πy)
  @unpack inv_Tr = equations

  # TODO: sincospi
  x1, x2 = x
  du1 = 10 * pi^2 * cospi(x1) * sinpi(2 * x2)
  du2 = -inv_Tr * u[2]
  du3 = -inv_Tr * u[3]

  return SVector(du1, du2, du3)
end

@inline function boundary_condition_poisson_nonperiodic(u_inner, orientation, direction, x, t,
                                                        surface_flux_function,
                                                        equations::HyperbolicDiffusionEquations2D)
  # elliptic equation: -ν Δϕ = f in Ω, u = g on ∂Ω
  u_boundary = initial_condition_poisson_nonperiodic(x, one(t), equations)

  # Calculate boundary flux
  if direction in (2, 4) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
  end

  return flux
end

@inline function initial_condition_harmonic_nonperiodic(x, t, equations::HyperbolicDiffusionEquations2D)
  # elliptic equation: -ν Δϕ = 0 in Ω, u = g on ∂Ω
  if t == 0.0
    phi = 1.0
    q1  = 1.0
    q2  = 1.0
  else
    C   = inv(sinh(pi))
    # TODO: sincospi
    sinpi_x1, cospi_x1 = sincos(pi*x[1])
    sinpi_x2, cospi_x2 = sincos(pi*x[2])
    sinh_pix1 = sinh(pi*x[1])
    cosh_pix1 = cosh(pi*x[1])
    sinh_pix2 = sinh(pi*x[2])
    cosh_pix2 = cosh(pi*x[2])
    phi = C *      (sinh_pix1 * sinpi_x2 + sinh_pix2 * sinpi_x1)
    q1  = C * pi * (cosh_pix1 * sinpi_x2 + sinh_pix2 * cospi_x1)
    q2  = C * pi * (sinh_pix1 * cospi_x2 + cosh_pix2 * sinpi_x1)
  end
  return SVector(phi, q1, q2)
end

@inline function source_terms_harmonic(u, x, t, equations::HyperbolicDiffusionEquations2D)
  # harmonic solution ϕ = (sinh(πx)sin(πy) + sinh(πy)sin(πx))/sinh(π), so f = 0
  @unpack inv_Tr = equations
  phi, q1, q2 = u

  du2 = -inv_Tr * q1
  du3 = -inv_Tr * q2

  return SVector(0, du2, du3)
end

"""
    initial_condition_eoc_test_coupled_euler_gravity(x, t, equations::HyperbolicDiffusionEquations2D)

Setup used for convergence tests of the Euler equations with self-gravity used in
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
in combination with [`source_terms_harmonic`](@ref).
"""
function initial_condition_eoc_test_coupled_euler_gravity(x, t, equations::HyperbolicDiffusionEquations2D)

  # Determine phi_x, phi_y
  G = 1.0 # gravitational constant
  C = -2.0*G/pi
  A = 0.1 # perturbation coefficient must match Euler setup
  rho1 = A * sin(pi * (x[1] + x[2] - t))
  # intialize with ansatz of gravity potential
  phi = C * rho1
  q1  = C * A * pi * cos(pi*(x[1] + x[2] - t)) # = gravity acceleration in x-direction
  q2  = q1                                     # = gravity acceleration in y-direction

  return SVector(phi, q1, q2)
end

"""
    initial_condition_sedov_self_gravity(x, t, equations::HyperbolicDiffusionEquations2D)

Adaptation of the Sedov blast wave with self-gravity taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
based on
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
Should be used together with [`boundary_condition_sedov_self_gravity`](@ref).
"""
function initial_condition_sedov_self_gravity(x, t, equations::HyperbolicDiffusionEquations2D)
  # for now just use constant initial condition for sedov blast wave (can likely be improved)
  phi = 0.0
  q1  = 0.0
  q2  = 0.0
  return SVector(phi, q1, q2)
end

"""
    boundary_condition_sedov_self_gravity(u_inner, orientation, direction, x, t,
                                          surface_flux_function,
                                          equations::HyperbolicDiffusionEquations2D)

Adaptation of the Sedov blast wave with self-gravity taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
based on
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
Should be used together with [`initial_condition_sedov_self_gravity`](@ref).
"""
function boundary_condition_sedov_self_gravity(u_inner, orientation, direction, x, t,
                                                surface_flux_function,
                                                equations::HyperbolicDiffusionEquations2D)
  u_boundary = initial_condition_sedov_self_gravity(x, t, equations)

  # Calculate boundary flux
  if direction in (2, 4) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
  end

  return flux
end

# Calculate 1D flux in for a single point
@inline function flux(u, orientation::Integer, equations::HyperbolicDiffusionEquations2D)
  phi, q1, q2 = u
  @unpack inv_Tr = equations

  if orientation == 1
    f1 = -equations.nu*q1
    f2 = -phi * inv_Tr
    f3 = zero(phi)
  else
    f1 = -equations.nu*q2
    f2 = zero(phi)
    f3 = -phi * inv_Tr
  end

  return SVector(f1, f2, f3)
end

# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation
@inline function max_abs_speed_naive(u_ll, u_rr, orientation::Integer, equations::HyperbolicDiffusionEquations2D)
  λ_max = sqrt(equations.nu * equations.inv_Tr)
end

@inline function flux_godunov(u_ll, u_rr, orientation::Integer, equations::HyperbolicDiffusionEquations2D)
  # Obtain left and right fluxes
  phi_ll, p_ll, q_ll = u_ll
  phi_rr, p_rr, q_rr = u_rr
  f_ll = flux(u_ll, orientation, equations)
  f_rr = flux(u_rr, orientation, equations)

  # this is an optimized version of the application of the upwind dissipation matrix:
  #   dissipation = 0.5*R_n*|Λ|*inv(R_n)[[u]]
  λ_max = sqrt(equations.nu * equations.inv_Tr)
  f1 = 1/2 * (f_ll[1] + f_rr[1]) - 1/2 * λ_max * (phi_rr - phi_ll)
  if orientation == 1 # x-direction
    f2 = 1/2 * (f_ll[2] + f_rr[2]) - 1/2 * λ_max * (p_rr - p_ll)
    f3 = 1/2 * (f_ll[3] + f_rr[3])
  else # y-direction
    f2 = 1/2 * (f_ll[2] + f_rr[2])
    f3 = 1/2 * (f_ll[3] + f_rr[3]) - 1/2 * λ_max * (q_rr - q_ll)
  end

  return SVector(f1, f2, f3)
end

@inline have_constant_speed(::HyperbolicDiffusionEquations2D) = Val(true)

@inline function max_abs_speeds(eq::HyperbolicDiffusionEquations2D)
  λ = sqrt(eq.nu * eq.inv_Tr)
  return λ, λ
end

# Convert conservative variables to primitive
@inline cons2prim(u, equations::HyperbolicDiffusionEquations2D) = u

# Convert conservative variables to entropy found in I Do Like CFD, Too, Vol. 1
@inline function cons2entropy(u, equations::HyperbolicDiffusionEquations2D)
  phi, q1, q2 = u
  w1 = phi
  w2 = equations.Lr^2 * q1
  w3 = equations.Lr^2 * q2

  return SVector(w1, w2, w3)
end

# Calculate entropy for a conservative state `u` (here: same as total energy)
@inline entropy(u, equations::HyperbolicDiffusionEquations2D) = energy_total(u, equations)

# Calculate total energy for a conservative state `u`
@inline function energy_total(u, equations::HyperbolicDiffusionEquations2D)
  # energy function as found in equations (2.5.12) in the book "I Do Like CFD, Vol. 1"
  phi, q1, q2 = u
  return 0.5 * (phi^2 + equations.Lr^2 * (q1^2 + q2^2))
end


text:
using MacroTools, InteractiveUtils, SpecialFunctions

struct Variable
  name::Symbol
  number::Int
end

Symbol(x::Variable) = Symbol(x.name, x.number)

Base.show(io::IO, x::Variable) = print(io, ":(", x.name, x.number, ")")

Base.print(io::IO, x::Variable) = Base.show_unquoted(io, x, 0, -1)
Base.show_unquoted(io::IO, x::Variable, ::Int, ::Int) =
  print(io, x.name, x.number)

struct Wengert
  variable::Symbol
  instructions::Vector{Any}
end

Wengert(; variable = :y) = Wengert(variable, [])

Base.keys(w::Wengert) = (Variable(w.variable, i) for i = 1:length(w.instructions))
Base.lastindex(w::Wengert) = Variable(w.variable, length(w.instructions))

Base.getindex(w::Wengert, v::Variable) = w.instructions[v.number]

function Base.show(io::IO, w::Wengert)
  println(io, "Wengert List")
  for (i, x) in enumerate(w.instructions)
    print(io, Variable(w.variable, i), " = ")
    Base.println(io, x)
  end
end

Base.push!(w::Wengert, x) = x

function Base.push!(w::Wengert, x::Expr)
  isexpr(x, :block) && return pushblock!(w, x)
  x = Expr(x.head, map(x -> x isa Expr ? push!(w, x) : x, x.args)...)
  push!(w.instructions, x)
  return lastindex(w)
end

function pushblock!(w::Wengert, x)
  bs = Dict()
  rename(ex) = Expr(ex.head, map(x -> get(bs, x, x), ex.args)...)
  for arg in MacroTools.striplines(x).args
    if @capture(arg, x_ = y_)
      bs[x] = push!(w, rename(y))
    else
      push!(w, rename(arg))
    end
  end
  return Variable(w.variable, length(w.instructions))
end

function Wengert(ex; variable = :y)
  w = Wengert(variable = variable)
  push!(w, ex)
  return w
end

function Expr(w::Wengert)
  cs = Dict()
  for x in w.instructions
    x isa Expr || continue
    for v in x.args
      v isa Variable || continue
      cs[v] = get(cs, v, 0) + 1
    end
  end
  bs = Dict()
  rename(ex::Expr) = Expr(ex.head, map(x -> get(bs, x, x), ex.args)...)
  rename(x) = x
  ex = :(;)
  for v in keys(w)
    if get(cs, v, 0) > 1
      push!(ex.args, :($(Symbol(v)) = $(rename(w[v]))))
      bs[v] = Symbol(v)
    else
      bs[v] = rename(w[v])
    end
  end
  push!(ex.args, rename(bs[lastindex(w)]))
  return unblock(ex)
end

addm(a, b) = a == 0 ? b : b == 0 ? a : :($a + $b)
mulm(a, b) = 0 in (a, b) ? 0 : a == 1 ? b : b == 1 ? a : :($a * $b)
mulm(a, b, c...) = mulm(mulm(a, b), c...)

function derive(w::Wengert, x; out = w)
  ds = Dict()
  ds[x] = 1
  d(x) = get(ds, x, 0)
  for v in keys(w)
    ex = w[v]
    Δ = @capture(ex, a_ + b_) ? addm(d(a), d(b)) :
        @capture(ex, a_ * b_) ? addm(mulm(a, d(b)), mulm(b, d(a))) :
        @capture(ex, a_^n_Number) ? mulm(d(a),n,:($a^$(n-1))) :
        @capture(ex, a_ / b_) ? :($(mulm(b, d(a))) - $(mulm(a, d(b))) / $b^2) :
        @capture(ex, sin(a_)) ? mulm(:(cos($a)), d(a)) :
        @capture(ex, cos(a_)) ? mulm(:(-sin($a)), d(a)) :
        @capture(ex, exp(a_)) ? mulm(v, d(a)) :
        @capture(ex, log(a_)) ? mulm(:(1/$a), d(a)) :
        error("$ex is not differentiable")
    ds[v] = push!(out, Δ)
  end
  return out
end

function derive_r(w::Wengert, x)
  ds = Dict()
  d(x) = get(ds, x, 0)
  d(x, Δ) = ds[x] = haskey(ds, x) ? addm(ds[x],Δ) : Δ
  d(lastindex(w), 1)
  for v in reverse(collect(keys(w)))
    ex = w[v]
    Δ = d(v)
    if @capture(ex, a_ + b_)
      d(a, Δ)
      d(b, Δ)
    elseif @capture(ex, a_ * b_)
      d(a, push!(w, mulm(Δ, b)))
      d(b, push!(w, mulm(Δ, a)))
    elseif @capture(ex, a_^n_Number)
      d(a, mulm(Δ, n, :($a^$(n-1))))
    elseif @capture(ex, a_ / b_)
      d(a, push!(w, mulm(Δ, b)))
      d(b, push!(w, :(-$(mulm(Δ, a))/$b^2)))
    else
      error("$ex is not differentiable")
    end
  end
  push!(w, d(x))
  return w
end

ext: jl


text:
# SDE.jl

"""
    abstract type AbstactSDE <: Any

subtype: SDE.

A supertype for implementing SDE types.
dX_t=θ(X_t)dt + σ(X_t)dW_t.
"""
abstract type AbstractSDE end

"""
    struct SDE{S, T} <: AbstractSDE where {S<:Function, T<:SDEModel}
      b::S
      model::T
    end

# Constructors
```julia
SDE(b::S, model::T) where {S<:Function, T<:SDEModel}
SDE(model::T, b::S) where {S<:Function, T<:SDEModel}
```

Implements a stochastic differential equation dX_t=b(X_t)dt+σ(X_t)dW_t on time
interval [0,model.endtime], with W_t a Brownian motion. The drift function b is
functions R→R. The variance is given by model.σ (either a function or a Float64).
Beginvalue X_0=model.beginvalue and model.Δ>0 is the precision with which we
discretise.

# Example

```julia
model = SDEModel(x->2+sin(x), 0.0,10.0,0.001)
sde = SDE(sin, model)
```
"""
struct SDE{S, T} <: AbstractSDE where {S<:Function, T<:SDEModel}
  b::S
  model::T
end

# Constructor
SDE(model::T, b::S) where {S<:Function, T<:SDEModel} = SDE(b, model)

function calculatenextsamplevalue(prevXval::Float64, b::Function,
    model::SDEModel{Float64}, BMincrement::Float64)
  return prevXval + b(prevXval)*model.Δ + model.σ*BMincrement
end

function calculatenextsamplevalue(prevXval::Float64, b::Function,
    model::T, BMincrement::Float64) where T<:SDEModel{S} where S<:Function
  return prevXval + b(prevXval)*model.Δ + model.σ(prevXval)*BMincrement
end

"""
  rand(sde::SDE)

Returns a SamplePath object which represents a sample path from an
AbstractSDE subtype sde. From time 0.0 to time sde.model.endtime, discretised
with precision sde.model.Δ.

# Examples

```julia
model = SDEModel(1.0,0.0,10.0,0.01)
sde = SDE(sin, model)
X = rand(sde)
```
"""
function rand(sde::SDE)::SamplePath
  timeinterval = 0.0:sde.model.Δ:sde.model.endtime
  lengthoftimeinterval = length(timeinterval)
  BMincrements = sqrt(sde.model.Δ) .* randn(lengthoftimeinterval)
  samplevalues = Array{Float64}(lengthoftimeinterval)
  prevXval = samplevalues[1] = sde.model.beginvalue
  for k in 2:lengthoftimeinterval
    samplevalues[k] = calculatenextsamplevalue(prevXval, sde.b, sde.model,
        BMincrements[k])
    prevXval = samplevalues[k]
  end
  return SamplePath(timeinterval, samplevalues)
end


text:
export
    AtomBijection, calculateRMSD, assignByName, assignCAlphaAtoms, assignBackboneAtoms

"""
    AtomBijection
`AtomBijection` is a bijective Map of `Atom`s. Read usage at [Bijection's documentation](https://github.com/scheinerman/Bijections.jl#using-a-bijection)
"""
const AtomBijection = Bijection{Atom,Atom}

"""
    calculateRMSD(bijec::AtomBijection)
Gets the [Root-mean-square-deviation](https://en.wikipedia.org/wiki/Root-mean-square_deviation)
of the atom in `bijec`.
"""
function calculateRMSD(bijec::AtomBijection)
    sum_of_squares::Float64 = 0
    for (left_at, right_at) in bijec
        sum_of_squares += sqeuclidean(left_at.position_,right_at.position_)
    end
    sum_of_squares = sqrt(sum_of_squares / length(bijec))
    return sum_of_squares
    #= julia alternative
    return sqrt(sum(map(sqeuclidean, keys(bijec), values(bijec)))) / length(bijec))
    =#
end

"""
Maps atoms in `A` to atoms of equal name in `B`.
"""
function assignByName(A::T,  B::S, limit_to_selection::Bool) where {T,S<:AbstractAtom}
    result = AtomBijection()
    A_names = Dict{String, Atom}([(atom.name_, atom) for atom in collectAtoms(A) ])
    for atom in collectAtoms(B)
        if haskey(A_names,atom.name_) &&
                    ( !limit_to_selection || isSelected(atom) || isSelected(A_names[atom.name_]))
             result[A_names[atom.name_]] = atom
        end
        delete!(A_names, atom.name_)
    end
    return result
end

"""
    assignCAlphaAtoms(A::Composite,  B::Composite, ::Bool)
Maps C-alpha atoms in `A` to C-alpha atoms in `B`.
"""
function assignCAlphaAtoms(A::AbstractComposite,  B::AbstractComposite, limit_to_selection::Bool) where {T,S<:AbstractAtom}
    result = AtomBijection()
    res_list_A= collectResidues(A)
    res_list_B= collectResidues(B)
    for (res_A,res_B) in zip(res_list_A,res_list_B)
        at_A = getAtom(res_A,"CA")
        at_B = getAtom(res_B,"CA")
        if !isnothing(at_A) && !isnothing(at_B) &&
                                    (!limit_to_selection || isSelected(at_a) || isSelected(at_B))
            result[at_A] = at_B
        end
    end
    return result
end

"""
    assignBackboneAtoms(A::Composite, B::Composite, limit_to_selection::Bool)
Maps the backbone atoms in `A` to the backbone atoms in `B`.
"""
assignBackboneAtoms(A::AbstractComposite, B::AbstractComposite, limit_to_selection::Bool) where {T,S<:AbstractAtom} = begin
    result = AtomBijection()
    res_list_A= collectResidues(A)
    res_list_B= collectResidues(B)
    for (res_A,res_B) in zip(res_list_A,res_list_B)
        backbone_atoms = String["CA", "C", "H", "O", "N"]
        for at_name in backbone_atoms
            at_A = getAtom(res_A,at_name)
            at_B = getAtom(res_B,at_name)

            if !isnothing(at_A) && !isnothing(at_B) &&
                                    (!limit_to_selection || isSelected(at_a) || isSelected(at_B))
                result[at_A] = at_B
            end
        end
    end
    return result
end


text:
using MomentClosure
using MomentClosure: expected_coeff
using ModelingToolkit: value, simplify
using SymbolicUtils: expand
using Catalyst
using Test

@parameters b
b = b.val
@test expected_coeff(2, 3) == 8
@test isequal(expected_coeff(b, 3), b*(1 + (6b) + (6(b^2))))
expr = ((b*((1 + b)^-1)) + ((b^6)*((1 + b)^-6)) + (57(b^2)*((1 + b)^-2)) + (302(b^3)*((1 + b)^-3)) + (302(b^4)*((1 + b)^-4)) + (57(b^5)*((1 + b)^-5)))*((1 + b)^6)
@test isequal(expected_coeff(b, 6), expr)

@parameters t, c₁, c₂, c₃, c₄, Ω
@variables X(t), Y(t)

S_mat = [ 1 -1  1 -1;
         -1  1  0  0]
a = [c₁*X*Y*(X-1)/Ω^2, c₂*X, c₃*Ω, c₄*X]
rn = ReactionSystemMod(t, [X, Y], [c₁, c₂, c₃, c₄, Ω], a, S_mat)

sys = generate_central_moment_eqs(rn, 2, 4)
expr1 = sys.odes.eqs[2].rhs
@test isequal(MomentClosure.Differential(t)(sys.μ[1,0]), sys.odes.eqs[1].lhs)
μ = sys.μ
M = sys.M
# simplify(value()) is needed due to diffs between ModelingToolkit Num type and SymbolicUtils types
expr2 = c₂*μ[1,0] + c₁*M[1,1]*(Ω^-2) + c₁*μ[0,1]*μ[1,0]*Ω^-2 - c₁*M[2,1]*Ω^-2 -
        2*c₁*M[1,1]*μ[1,0]*Ω^-2 - c₁*M[2,0]*μ[0,1]*Ω^-2 - c₁*μ[0,1]*Ω^-2*μ[1,0]^2
expr2 = simplify(value.(expr2))
@test isequal(expand(expr1), expr2)

sys = generate_central_moment_eqs(rn, 2)
expr1 = sys.odes.eqs[2].rhs
@test isequal(MomentClosure.Differential(t)(sys.μ[1,0]), sys.odes.eqs[1].lhs)
@test isequal(expand(expr1), expr2)

sys = generate_raw_moment_eqs(rn, 2)
μ = sys.μ
expr1 = sys.odes.eqs[4].rhs
expr2 = c₂*μ[2,0] + c₁*μ[1,1]/Ω^2 - c₁*μ[3,1]/Ω^2 -c₂*(μ[1,0] + μ[1,1]) +
        c₃*μ[0,1]*Ω - c₄*μ[1,1] - c₁*μ[1,2]/Ω^2 + c₁*μ[2,2]/Ω^2
expr2 = simplify(value.(expr2))
@test isequal(expand(expr1), expand(expr2))

# corner case - a linear propensity with rate coefficient 1
rn = @reaction_network begin
    1, X → 0
end
sys = generate_raw_moment_eqs(rn, 2)
μ = sys.μ
@test isequal(MomentClosure.Differential(t)(μ[(1,)]) ~ -μ[(1,)], sys.odes.eqs[1])

# time-dependent propensity tests

rn = @reaction_network begin
        (X^1.5), X ⇒ Y
end c₁
@test_throws ErrorException generate_raw_moment_eqs(rn, 2)

rn = @reaction_network begin
        (c₁*Y^3*X^1.5), X ⇒ Y
end c₁
@test_throws ErrorException generate_raw_moment_eqs(rn, 2)

rn = @reaction_network begin
        (c₁*sin(exp(X)+Y)), X ⇒ Y
end c₁
@test_throws ErrorException generate_raw_moment_eqs(rn, 2)

rn = @reaction_network begin
        (c₁*sin(exp(t))+Y^(3+1)), X ⇒ Y
end c₁
eqs = generate_raw_moment_eqs(rn, 2)
@test eqs.q_order == 5


text:
# This set of functions is based entirely on the excellent
# scikit-learn package's method of calculating an ROC's area
# under the curve (AUC). The functions below produce the
# same output as scikit-learn.

function _binary_clf_curve(y_true, y_score)
    y_true = y_true .== 1       # make y_true a boolean vector
    desc_score_indices = sortperm(y_score, rev = true)

    y_score = y_score[desc_score_indices]
    y_true = y_true[desc_score_indices]

    distinct_value_indices = find(diff(y_score))
    threshold_idxs = push!(distinct_value_indices, length(y_score))

    tps = cumsum(y_true)[threshold_idxs]
    fps = threshold_idxs - tps
    return (fps, tps, y_score[threshold_idxs])
end

# y = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1, 1, 1]
# y_score = [0.3, 0.2, 0.3, 0.23, 0.5, 0.34, 0.45, 0.54, 0.6, 0.7, 0.8, 0.65, 0.5, 0.4, 0.3, 0.2, 0.6, 0.7, 0.5, 0.2, 0.1, 0.7, 0.2, 0.7, 0.4]

# _binary_clf_curve(y, y_score)

function roc_curve(y_true, y_score)
    fps, tps, thresholds = _binary_clf_curve(y_true, y_score)
    fpr = fps/fps[end]
    tpr = tps/tps[end]
    return (fpr, tpr, thresholds)
end

# roc_curve(y, y_score)

"""
Given 2 vectors, `x` and `y`, this function returns the indices that
sort the elements by `x`, with `y` breaking ties. See the example below.

julia> a = [2, 1, 3, 2]
julia> b = [3, 4, 1, 0]
julia> order = sortperm2(a, b)
4-element Array{Int64,1}:
 2
 4
 1
 3

julia> hcat(a[order], b[order]
4×2 Array{Int64,2}:
 1  4
 2  0
 2  3
 3  1
 """
function sortperm2(x, y; rev = false)
    n = length(x)
    no_ties = n == length(Set(x))
    if no_ties
        res = sortperm(x, rev = rev)
    else
        ord1 = sortperm(x, rev = rev)
        x_sorted = x[ord1]
        i = 1
        while i < n
            println(i)
            # println("x_i is $(x_sorted[i]) and x_(i+1) is $(x_sorted[i+1])")
            if x_sorted[i] == x_sorted[i+1]
                if rev && y[ord1][i] < y[ord1][i+1]
                    println("(1.) Switching $(y[ord1][i]) with $(y[ord1][i+1])")
                    ord1[i], ord1[i+1] = ord1[i+1], ord1[i]
                    i = i > 1 ? i - 1 : i
                    continue
                elseif !rev && y[ord1][i] > y[ord1][i+1]
                    println("(2.) Switching $(y[ord1][i]) with $(y[ord1][i+1])")
                    ord1[i], ord1[i+1] = ord1[i+1], ord1[i]
                    i = i > 1 ? i - 1 : i
                    continue
                end
            end
            i += 1
        end
        res = ord1
    end
    res
end

# a = [1, 5, 1, 4, 3, 4, 4, 3, 1, 4, 5, 3, 5]
# b = [9, 4, 0, 4, 0, 2, 1, 2, 1, 3, 2, 1, 1]
#
# ord = sortperm2(a, b, rev = true)
# hcat(a[ord], b[ord])
#
# ord = sortperm2(a, b)
# hcat(a[ord], b[ord])

function trapsum(y, x)
    d = diff(x)
    res = sum((d .* (y[2:end] + y[1:(end-1)]) ./ 2.0))
    res
end

# trapsum([1, 2, 3], [4, 6, 8])     # 8.0

function auc(x, y, reorder = false)
    direction = 1
    if reorder
        order = sortperm2(x, y)
        x, y = x[order], y[order]
    else
        dx = diff(x)
        if any(dx .<= 0)
            if all(dx .<= 0)
                direction = -1
            else
                error("Reordering is not turned on, and the x array is not increasing: $x")
            end
        end
    end
    area = direction * trapsum(y, x)
    return area
end

function roc_auc_score(y_true, y_score)
    if length(Set(y_true)) == 1
        warn("Only one class present in y_true.\n
              The AUC is not defined in that case; returning -Inf.")
        res = -Inf
    else
        fpr, tpr, thresholds = roc_curve(y_true, y_score)
        res = auc(fpr, tpr, true)
    end
    res
end

# roc_auc_score(y, y_score)

max_stars_repo_name: paulstey/boosting

max_issues_repo_name: paulstey/boosting

max_forks_repo_name: paulstey/boosting


text:
# This is a copy of Knet/data/imdb.jl and Knet/tutorial/07.imdb.ipynb cleaned up

# Based on https://github.com/fchollet/keras/raw/master/keras/datasets/imdb.py
# Also see https://github.com/fchollet/keras/raw/master/examples/imdb_lstm.py
# Also see https://github.com/ilkarman/DeepLearningFrameworks/raw/master/common/utils.py

using Pkg; for p in ("PyCall","JSON","JLD2"); haskey(Pkg.installed(),p) || Pkg.add(p); end
using PyCall,JSON,JLD2,Random
np = pyimport("numpy")

"""

    imdb()

Load the IMDB Movie reviews sentiment classification dataset from
https://keras.io/datasets and return (xtrn,ytrn,xtst,ytst,dict) tuple.

# Keyword Arguments:
- url=https://s3.amazonaws.com/text-datasets: where to download the data (imdb.npz) from.
- dir=Pkg.dir("Knet/data"): where to cache the data.
- maxval=nothing: max number of token values to include. Words are ranked by how often they occur (in the training set) and only the most frequent words are kept. nothing means keep all, equivalent to maxval = vocabSize + pad + stoken.
- maxlen=nothing: truncate sequences after this length. nothing means do not truncate.
- seed=0: random seed for sample shuffling. Use system seed if 0.
- pad=true: whether to pad short sequences (padding is done at the beginning of sequences). pad_token = maxval.
- stoken=true: whether to add a start token to the beginning of each sequence. start_token = maxval - pad.
- oov=true: whether to replace words >= oov_token with oov_token (the alternative is to skip them). oov_token = maxval - pad - stoken.

"""
function imdb(;
              url = "https://s3.amazonaws.com/text-datasets",
              dir = "./", # joinpath(@__DIR__, "imdb"),
              data="imdb.npz",
              dict="imdb_word_index.json",
              jld2="imdb.jld2",
              maxval=nothing,
              maxlen=nothing,
              seed=0, oov=true, stoken=true, pad=true
              )
    global _imdb_xtrn,_imdb_ytrn,_imdb_xtst,_imdb_ytst,_imdb_dict
    if !(@isdefined _imdb_xtrn)
        isdir(dir) || mkpath(dir)
        jld2path = joinpath(dir,jld2)
        if !isfile(jld2path)
            @info("Downloading IMDB...")
            datapath = joinpath(dir,data)
            dictpath = joinpath(dir,dict)
            isfile(datapath) || download("$url/$data",datapath)
            isfile(dictpath) || download("$url/$dict",dictpath)
            d = np.load(datapath)
            _imdb_xtrn = map(a->np.asarray(a,dtype=np.int32), get(d, "x_train"))
            _imdb_ytrn = Array{Int8}(get(d, "y_train") .+ 1)
            _imdb_xtst = map(a->np.asarray(a,dtype=np.int32), get(d, "x_test"))
            _imdb_ytst = Array{Int8}(get(d, "y_test") .+ 1)
            _imdb_dict = Dict{String,Int32}(JSON.parsefile(dictpath))
            JLD2.@save jld2path _imdb_xtrn _imdb_ytrn _imdb_xtst _imdb_ytst _imdb_dict
            #rm(datapath)
            #rm(dictpath)
        end
        @info("Loading IMDB...")
        JLD2.@load jld2path _imdb_xtrn _imdb_ytrn _imdb_xtst _imdb_ytst _imdb_dict
    end
    if seed != 0; Random.seed!(seed); end
    xs = [_imdb_xtrn;_imdb_xtst]
    if maxlen == nothing; maxlen = maximum(map(length,xs)); end
    if maxval == nothing; maxval = maximum(map(maximum,xs)) + pad + stoken; end
    if pad; pad_token = maxval; maxval -= 1; end
    if stoken; start_token = maxval; maxval -= 1; end
    if oov; oov_token = maxval; end
    function _imdb_helper(x,y)
        rp = randperm(length(x))
        newy = y[rp]
        newx = similar(x)
        for i in 1:length(x)
            xi = x[rp[i]]
            if oov
                xi = map(w->(w<=oov_token ? w : oov_token), xi)
            else
                xi = filter(w->(w<=oov_token), xi)
            end
            if stoken
                xi = [ start_token; xi ]
            end
            if length(xi) > maxlen
                xi = xi[end-maxlen+1:end]
            end
            if pad && length(xi) < maxlen
                xi = append!(repeat([pad_token], maxlen-length(xi)), xi)
            end
            newx[i] = xi
        end
        newx,newy
    end
    xtrn,ytrn = _imdb_helper(_imdb_xtrn,_imdb_ytrn)
    xtst,ytst = _imdb_helper(_imdb_xtst,_imdb_ytst)
    return xtrn,ytrn,xtst,ytst,_imdb_dict
end

### From 07.imdb.ipynb

using Pkg
for p in ("Knet",) #"ProgressMeter")
    haskey(Pkg.installed(),p) || Pkg.add(p)
end

EPOCHS=3          # Number of training epochs
BATCHSIZE=64      # Number of instances in a minibatch
EMBEDSIZE=125     # Word embedding size
NUMHIDDEN=100     # Hidden layer size
MAXLEN=150        # maximum size of the word sequence, pad shorter sequences, truncate longer ones
VOCABSIZE=30000   # maximum vocabulary size, keep the most frequent 30K, map the rest to UNK token
NUMCLASS=2        # number of output classes
DROPOUT=0.0       # Dropout rate
LR=0.001          # Learning rate
BETA_1=0.9        # Adam optimization parameter
BETA_2=0.999      # Adam optimization parameter
EPS=1e-08         # Adam optimization parameter

using Knet # : Knet
#ENV["COLUMNS"]=92                     # column width for array printing
#include(Knet.dir("data","imdb.jl"))   # defines imdb loader

#@doc imdb

#@time
(xtrn,ytrn,xtst,ytst,imdbdict)=imdb(maxlen=MAXLEN,maxval=VOCABSIZE);

#summary.((xtrn,ytrn,xtst,ytst,imdbdict))

# Words are encoded with integers
#rand(xtrn)'

# Each word sequence is padded or truncated to length 150
#length.(xtrn)'

# Define a function that can print the actual words:
imdbvocab = Array{String}(undef,length(imdbdict))
for (k,v) in imdbdict; imdbvocab[v]=k; end
imdbvocab[VOCABSIZE-2:VOCABSIZE] = ["<unk>","<s>","<pad>"]
function reviewstring(x,y=0)
    x = x[x.!=VOCABSIZE] # remove pads
    """$(("Sample","Negative","Positive")[y+1]) review:\n$(join(imdbvocab[x]," "))"""
end

# Hit Ctrl-Enter to see random reviews:
#r = rand(1:length(xtrn))
#println(reviewstring(xtrn[r],ytrn[r]))

# Here are the labels: 1=negative, 2=positive
#ytrn'

using Knet: param, dropout, RNN

struct SequenceClassifier; input; rnn; output; pdrop; end

SequenceClassifier(input::Int, embed::Int, hidden::Int, output::Int; pdrop=0) =
    SequenceClassifier(param(embed,input), RNN(embed,hidden,rnnType=:gru), param(output,hidden), pdrop)

function (sc::SequenceClassifier)(input)
    embed = sc.input[:, permutedims(hcat(input...))]
    embed = dropout(embed,sc.pdrop)
    hidden = sc.rnn(embed)
    hidden = dropout(hidden,sc.pdrop)
    return sc.output * hidden[:,:,end]
end

using Knet: minibatch
dtrn = minibatch(xtrn,ytrn,BATCHSIZE;shuffle=true)
dtst = minibatch(xtst,ytst,BATCHSIZE)
#length.((dtrn,dtst))

# For running experiments
#using Knet: train!, Adam
#import ProgressMeter

function trainresults(file,model)
    # if (print("Train from scratch? ");readline()[1]=='y')
    #     updates = 0; prog = ProgressMeter.Progress(EPOCHS * length(dtrn))
    #     function callback(J)
    #         ProgressMeter.update!(prog, updates)
    #         return (updates += 1) <= prog.n
    #     end
    #     opt = Adam(lr=LR, beta1=BETA_1, beta2=BETA_2, eps=EPS)
    #     train!(model, dtrn; callback=callback, optimizer=opt, pdrop=DROPOUT)
    #     Knet.gc()
    #     Knet.save(file,"model",model)
    # else
        isfile(file) || download("http://people.csail.mit.edu/deniz/models/tutorial/$file",file)
        model = Knet.load(file,"model")
    # end
    return model
end

#using Knet: nll, accuracy
model = SequenceClassifier(VOCABSIZE,EMBEDSIZE,NUMHIDDEN,NUMCLASS)
#nll(model,dtrn), nll(model,dtst), accuracy(model,dtrn), accuracy(model,dtst)

model = trainresults("imdbmodel113.jld2",model);

# 33s (0.059155148f0, 0.3877507f0, 0.9846153846153847, 0.8583733974358975)
#nll(model,dtrn), nll(model,dtst), accuracy(model,dtrn), accuracy(model,dtst)

predictstring(x)="\nPrediction: " * ("Negative","Positive")[argmax(Array(vec(model([x]))))]
UNK = VOCABSIZE-2
str2ids(s::String)=[(i=get(imdbdict,w,UNK); i>=UNK ? UNK : i) for w in split(lowercase(s))]

# Here we can see predictions for random reviews from the test set; hit Ctrl-Enter to sample:
# r = rand(1:length(xtst))
# println(reviewstring(xtst[r],ytst[r]))
# println(predictstring(xtst[r]))

# Here the user can enter their own reviews and classify them:
# println(predictstring(str2ids(readline(stdin))))


text:
function value(obj::ExtendedUnivariateDifferentiableMatrixFunction, arg0::Derivative)
    return jcall(obj, "value", Vector{Vector{Derivative}}, (Derivative,), arg0)
end

function value(obj::UnivariateDifferentiableMatrixFunction, arg0::DerivativeStructure)
    return jcall(obj, "value", Vector{Vector{DerivativeStructure}}, (DerivativeStructure,), arg0)
end

function value(obj::UnivariateMatrixFunction, arg0::jdouble)
    return jcall(obj, "value", Vector{Vector{jdouble}}, (jdouble,), arg0)
end


text:
export zerossiz, onessiz, randsiz, randnsiz
export shzeros, shones, shrand, shrandn
export shzerossiz, shonessiz, shrandsiz, shrandnsiz
export zeroel, oneel
export +, *, repeat, nop, id,  istrue, isfalse, not, or, and, @dict
export plus, minus, times, divby, square, power
export any, anyequal, all, allequal, unequal
export inside
export minimum, maximum

#######################################
##  zerossiz, onessiz, randsiz, randnsiz

zerossiz(a::AbstractArray, T::Type = Float64) = zeros(T, a...)
onessiz(a::AbstractArray, T::Type = Float64) = ones(T, a...)
randsiz(a::AbstractArray, T::Type = Float64) = rand(T, a...)
randnsiz(a::AbstractArray) = randn(a...)
randnsiz(a::AbstractArray, T::Type) = (r = Array{T,length(a)r}(a...); fillrandn(r))

shzerossiz(a::AbstractArray, T::Type = Float64) = shzeros(T, a...)
shonessiz(a::AbstractArray, T::Type = Float64) = shones(T, a...)
shrandsiz(a::AbstractArray, T::Type = Float64) = shrand(T, a...)
shrandnsiz(a::AbstractArray, T::Type = Float64) = shrandn(a...)

fillrand(a::SharedArray{T}) where {T} = (for i = 1:length(a) a[i] = rand(T) end; a)
fillrandn(a::SharedArray{T}) where {T} = (for i = 1:length(a) a[i] = randn() end; a)
shzeros(a...) = SharedArray(Float64, a...)
shzeros(T::Type, a...) = SharedArray{T,length(a)}(a...)
shones(a...) = fill!(SharedArray{Float64,length(a)}(a...), one(Float64))
shones(T::Type, a...) = fill!(SharedArray{T,length(a)}(a...), one(T))
shrand(a...) = fillrand(SharedArray{Float64,length(a)}(a...))
shrand(T::Type, a...) = fillrand(SharedArray{T,length(a)}(a...))
shrandn(a...) = fillrandn(SharedArray{Float64,length(a)}(a...))
shrandn(T::Type, a...) = fillrandn(SharedArray{T,length(a)}(a...))

#######################################
## zeroel, oneel

zeroel(a) = zero(eltype(a))
oneel(a) = one(eltype(a))

macro dict(a...)
    esc(Expr(:call, :Dict, [Meta.parse(":$x => $x") for x in a]...))
end

import Base.*
*(a::Union{Char,AbstractString}...) = string(a...)

import Base.repeat
# repeat(a::AbstractArray{T,1}, n::Integer) where T = flatten(map(unstack(1:n), x->a))
repeat(a::AbstractArray{T,2}, n::Integer) where T = flatten(map(unstack(1:n), x->a))
repeat(a, n) = flatten(map(unstack(1:n), x->a))

nop(a...) = return
id(a) = a
id(a...) = a

istrue(a::Bool) = a
istrue(f::Function) = istrue(f())
isfalse(a) = !(istrue(a))

#######################################
##  not, or, and

not = !
or(a,b) =  a || b
and(a,b) = a && b

#######################################
###  any, anyequal, all, allequal

import Base.any
function any(a::AbstractArray, f::Function)
    v = trytoview(a,1)
    for i = 1:len(a)
        if f(trytoview(a,i,v))
            return true
        end
    end
    return false
end

anyequal(a::AbstractArray, x) = return any(a, y -> isequal(y, x))

import Base.all
function all(a::AbstractArray, f::Function)
    v = trytoview(a,1)
    for i = 1:len(a)
        if !f(trytoview(a,i,v))
            return false
        end
    end
    return true
end

allequal(a::AbstractArray, x) = all(a, y -> isequal(y, x))

#######################################
##  internal data management

newarraysize(a::Number,n::Int) = (n,)
function newarraysize(a,n::Int)
    s = siz(a)
    if s[end] == 1
        s = s[1:end-1]
    end
    tuple(s...,n)
end

arraylike(a::T, n::Int, array = nothing) where {T<:AbstractString} = Array{T,1}(undef, n)
function arraylike(a::T, n::Int, array = nothing) where T
    # if !(a == nothing || (array != nothing && !(eltype(array) <: Number)))
    if T <: AbstractArray
        s = newarraysize(a,n)
        r = Array{eltype(a),length(s)}(undef, s...)
    else
        return Vector{T}(undef,n)
    end
end

function sharraylike(a, n::Int)
    s = newarraysize(a,n)
    SharedArray{eltype(a),length(s)}(s...)
end

plus(a,b) = a.+b
minus(a,b) = a.-b
times(a,b) = a.*b
divby(a,b) = a./b
unequal(a,b) = !(isequal(a,b))
square(a) = a.*a
power(a,n) = a.^n

inside = in

import Base: minimum, maximum
minimum(::Type{T} = Float64) where {T<:AbstractFloat} = -floatmax(T)
maximum(::Type{T} = Float64) where {T<:AbstractFloat} =  floatmax(T)
minimum(::Type{T}) where {T<:Number} = typemin(T)
maximum(::Type{T}) where {T<:Number} = typemax(T)


text:
"""
    wireframe(x, y, z)
    wireframe(positions)
    wireframe(mesh)

Draws a wireframe, either interpreted as a surface or as a mesh.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Wireframe) do scene
    default_theme(scene, LineSegments)
end

function convert_arguments(::Type{<: Wireframe}, x::AbstractVector, y::AbstractVector, z::AbstractMatrix)
    (ngrid(x, y)..., z)
end

xvector(x::AbstractVector, len) = x
xvector(x::ClosedInterval, len) = range(minimum(x), stop=maximum(x), length=len)
xvector(x::AbstractMatrix, len) = x

yvector(x, len) = xvector(x, len)'
yvector(x::AbstractMatrix, len) = x

function plot!(plot::Wireframe{<: Tuple{<: Any, <: Any, <: AbstractMatrix}})
    points_faces = lift(plot[1:3]...) do x, y, z
        M, N = size(z)
        points = vec(Point3f0.(xvector(x, M), yvector(y, N), z))
        # Connect the vetices with faces, as one would use for a 2D Rectangle
        # grid with M,N grid points
        faces = decompose(LineFace{GLIndex}, Tesselation(Rect2D(0, 0, 1, 1), (M, N)))
        connect(points, faces)
    end
    linesegments!(plot, Attributes(plot), points_faces)
end

function plot!(plot::Wireframe{Tuple{T}}) where T
    points = lift(plot[1]) do g
        # get the point representation of the geometry
        indices = decompose(LineFace{GLIndex}, g)
        points = decompose(Point3f0, g)
        return connect(points, indices)
    end
    linesegments!(plot, Attributes(plot), points)
end


text:
#
# some stuff for generating the incomes module
# output will need hand-editing.
#
n = 0
items = []
for line in eachline("etc/incomes.txt")
    global n, items
    m = match(r" *(.*?) *=.*", line )
    if m !== nothing
        # print(m[1])
        s = m[1]
        su = uppercase(s)
        if s[1] !== '#'
            push!(items,(s,su))
        end
    end
end
println("\n\n# declarations  ----------------")
n = size(items)[1]
for i in 1:n
    println("    const $(items[i][2]) = $i")
end

println("\n\n# names ----------------")
for i in 1:n
    censored = titlecase(replace(items[i][1], r"[_]" => " "))
    print("    elseif i == $(items[i][2])
            return \"$censored\"
    ")
end

println("\n\n# mappings ----------------")
for i in 1:n
    k = items[i][2]
    j = items[i][1]
    if match( r"SPARE.*", k ) === nothing
        print( "
        if haskey(incd, Definitions.$j )
            out[$k] = incd[Definitions.$j]
        end")
    end
end

println("\n\n# exports ----------------")
for i in 1:n
    k = items[i][2]
    println("export $k")
end

max_stars_repo_name: grahamstark/ScottishTaxBenefitModel

max_issues_repo_name: grahamstark/ScottishTaxBenefitModel

max_forks_repo_name: grahamstark/ScottishTaxBenefitModel


text:
"""
   SIGNAL TO NOISE MEASUREMENT
"""
function SNRatio( corr::Array{T,N}, method::S, width::I ) where {T<:Real,N}

    max2       = 0.0;
    h, w, d    = size(corr)
    idx1, max1 = firstPeak( corr )

    if     method == "peak2peak"
        idx2, max2 = secondPeak( corr, width, idx1 )
    elseif method == "peak2mean"
        max2 = mean(corr)
    else
        throw(ErrorException("Error: use 'peak2peak' or 'peak2mean'.\n"))
    end

    # if max2 == 0, sig2noise = inf. Check this at return
    sig2noise = max2 != 0.0 ? max1 / max2 : 0.0
    return sig2Noise;
end


text:
export local_basis

function local_basis(ψ::CuMPS, i::Int)
    d = physical_dim(ψ, i)
    ret = CUDA.zeros(Int, d)
    @inline function kernel(ret)
        state = (blockIdx().x-1) * blockDim().x + threadIdx().x
            if state <= length(ret)
                if state == 1
                    ret[state] = -1
                else
                    ret[state] = state - 1
                end
            end
        return
    end
    threads, blocks = cudiv(length(ret))
    CUDA.@cuda threads=threads blocks=blocks kernel(ret)
    ret
end

LinearAlgebra.I(ψ::CuMPS, i::Int) = CuMatrix(I(physical_dim(ψ, i)))


text:
module ColorDiffCharts

using Colors

struct ColorDiffChartSVG <: Main.SVG
    buf::IOBuffer
end

const c = Lab.((Colors.JULIA_LOGO_COLORS.green,
                Colors.JULIA_LOGO_COLORS.red,
                Colors.JULIA_LOGO_COLORS.purple))

function ColorDiffChartSVG(metric::Colors.DifferenceMetric)
    io = IOBuffer()

    id = nameof(typeof(metric))

    d12 = colordiff(c[1], c[2], metric=metric)
    d23 = colordiff(c[2], c[3], metric=metric)
    d31 = colordiff(c[3], c[1], metric=metric)

    dn = 2*d31^2 - d23^2 + 2*d12^2
    ds0 = -(d31-d23-d12) * (d31-d23+d12) * (d31+d23-d12) * (d31+d23+d12)
    # `ds0` should be non-negative in metric systems.
    # However, DE_CMC is not metric but quasimetric.
    ds = max(ds0, 0)

    x = [0,
         -sqrt(ds/dn)/2,
          sqrt(ds/dn)/2]

    y = [-sqrt(dn) / 3,
         (( -d31^2 - d23^2 + 5*d12^2) * sqrt(dn)) / 6dn,
         ((5*d31^2 - d23^2 -   d12^2) * sqrt(dn)) / 6dn]

    # verification
    if ds0 >= 0
        d12 ≈ hypot(x[1] - x[2], y[1] - y[2]) || error("distance mismatch: d12")
        d23 ≈ hypot(x[2] - x[3], y[2] - y[3]) || error("distance mismatch: d23")
        d31 ≈ hypot(x[3] - x[1], y[3] - y[1]) || error("distance mismatch: d31")
    end

    scale = 7.5 / abs(y[1])

    sx = scale .* x
    sy = scale .* y
    r = scale * 20

    mx12, my12 = (sx[1] + sx[2]) * 0.5, (sy[1] + sy[2]) * 0.5
    mx23, my23 = (sx[2] + sx[3]) * 0.5, (sy[2] + sy[3]) * 0.5
    mx31, my31 = (sx[3] + sx[1]) * 0.5, (sy[3] + sy[1]) * 0.5

    a = atan(sy[2] - sy[3], sx[2] - sx[3])
    dy23, dx23 = 8 .* sincos(a)

    simplify(x) = replace(string(round(x, sigdigits=3)), r"\.0+$"=>"")
    rd12, rd23, rd31 = simplify.((d12, d23, d31))

    write(io,
        """
        <svg xmlns="http://www.w3.org/2000/svg" version="1.1"
             viewBox="-15 -12 30 30" width="30mm" height="30mm"
             stroke="none" stroke-linejoin="round"
             style="display:inline; margin-left:2em; margin-bottom:2em">
        <defs>
          <marker id="marker_s_$id" orient="auto" style="overflow:visible">
            <path d="M 7,-30 l -7,4 7,4 M 0,-26 h 9 M 0,-3 v -27"
                  style="fill:none;stroke:currentColor;stroke-opacity:0.7;" />
          </marker>
          <marker id="marker_e_$id" orient="auto" style="overflow:visible">
            <path d="M -7,-30 l 7,4 -7,4 M 0,-26 h -9 M 0,-3 v -27"
                  style="fill:none;stroke:currentColor;stroke-opacity:0.7;" />
          </marker>
        </defs>
        <g>
          <circle cx="$(sx[3])" cy="$(sy[3])" fill="#$(hex(c[3]))" r="$r" />
          <circle cx="$(sx[2])" cy="$(sy[2])" fill="#$(hex(c[2]))" r="$r" />
          <circle cx="$(sx[1])" cy="$(sy[1])" fill="#$(hex(c[1]))" r="$r" />
        </g>
        <g style="stroke:currentColor;stroke-width:0.2;stroke-opacity:0.4;
           marker-start:url(#marker_s_$id);marker-end:url(#marker_e_$id)">
          <path d="M$(sx[3]),$(sy[3]) L$(sx[2]),$(sy[2])"/>
          <path d="M$(sx[2]),$(sy[2]) L$(sx[1]),$(sy[1])"/>
          <path d="M$(sx[1]),$(sy[1]) L$(sx[3]),$(sy[3])"/>
        </g>
        <g fill="currentColor" style="font-size:2.8px;">
          <text x="$(mx12-3)" y="$my12" text-anchor="end">$rd12</text>
          <text x="$(mx23+dy23)" y="$(my23-dx23+1)" text-anchor="middle">$rd23</text>
          <text x="$(mx31+3)" y="$my31" text-anchor="start">$rd31</text>
          <text x="-14" y="17" style="font-size:2.5px;">$id</text>
        </g>
        </svg>""")
    ColorDiffChartSVG(io)
end

end


text:
struct FTL
    L :: Array{Float64, 1};

    function FTL(K)
        new(zeros(K));
    end
end

function act(h::FTL)
    u = h.L .== minimum(h.L);
    u ./ sum(u);
end

function incur!(h::FTL, l)
    h.L .+= l;
end

mutable struct AdaHedge
    L :: Array{Float64, 1};
    Δ :: Float64;

    function AdaHedge(K)
        new(zeros(K), 0.01);
    end
end

function act(h::AdaHedge)
    η = log(length(h.L))/h.Δ;
    u = exp.(-η.*(h.L .- minimum(h.L)));
    u ./ sum(u);
end

function incur!(h::AdaHedge, l)
    u = act(h);
    η = log(length(h.L))/h.Δ;
    h.L .+= l;
    m = minimum(l) - 1/η*log(u'exp.(-η.*(l .- minimum(l))));
    # println(η, " ", u, " ", u'l, " ", -1/η*log(u'exp.(-η.*l)));
    @assert isfinite(m) && u'l ≥ m - 1e-7 "hedge loss $(u'l), Δ=$(h.Δ), η=$η, mix loss $m, u=$u, L=$(h.L)";
    h.Δ  += u'l - m;
end

# Menard calls this "Lazy Mirror Descent" As perhaps the main feature
# here is the mixing step, we call it "Fixed Share".  This is the
# version with slowly decreasing O(1/sqrt(t)) learnign and switching
# rates.  The switching is used to force exploration
mutable struct FixedShare
    w :: Array{Float64, 1};
    t; # time
    S; # loss range: S ≥ \max_k l_k - \min_k l_k
       # TODO: adapt to it (like AdaHedge?)

    function FixedShare(K; S=1)
        new(fill(1/K,K), 0, S);
    end
end

function act(h::FixedShare)
    h.w
end

function incur!(h::FixedShare, l)
    K = length(l);
    h.t += 1;

    # Exponential Weights update
    η = sqrt(log(K)/h.t)/h.S;   # decaying learning rate
    h.w .*= exp.(-η*(l .- minimum(l)));
    h.w ./= sum(h.w);

    # Fixed Share (to unifor prior)
    γ = 1/(4*sqrt(h.t)); # decaying switching rate
    h.w .= (1-γ).*h.w .+ γ/K;
end


text:
module NonUniformRandomVariateGeneration

using Random
import Random.GLOBAL_RNG

include("binomial.jl")
include("multinomial.jl")
include("poisson.jl")
include("gamma.jl")
include("beta.jl")
include("uniform.jl")
include("categorical.jl")

export sampleBinomial, sampleMultinomial, sampleMultinomial!, sampleGamma,
  sampleBeta, samplePoisson, sampleSortedUniforms!, sampleSortedUniforms,
  sampleCategorical, sampleCategoricalSorted, sampleCategoricalSorted!

end


text:
function build_uri(config, path)
    # API basic
    # http://username:password@dev.staffjoy.com/path
    # Note that our api keys are username and there is no password, so
    # http://apikey:@dev.staffjoy.com/path

    return string(config["protocol"], "://", config["api_key"], ":@", config["host"], path)
end

function days(organization)
    DAYS = [
        "monday",
        "tuesday",
        "wednesday",
        "thursday",
        "friday",
        "saturday",
        "sunday",
    ]

    start_day = organization["day_week_starts"]
    start_index = findfirst(DAYS, start_day)

    if start_index == 0
        Logging.err( "Unable to match organization start day")
        error()
    end

    output = Any[]
    for i in start_index:(start_index + length(DAYS) - 1)
        push!(
            output,
            # Stupid 1-indexing
            DAYS[(i - 1) % length(DAYS)+ 1]
        )
    end
    return output
end

max_stars_repo_name: shdann/autoshift

max_issues_repo_name: shdann/autoshift

max_forks_repo_name: shdann/autoshift


text:
#= Custom currency creation =#

"""
Add a new currency to the (global) currency list and return a unit for that
currency. Prefer the `@usingcustomcurrency` macro, which leads to more clear
code, whenever possible. This function takes three arguments: the symbol for the
currency, a string description of the currency (following the conventions
outlined in the documentation for `currencyinfo`), and an exponent representing
the number of decimal points to describe the minor currency unit in terms of the
major currency unit. Conventionally, the symbol used to describe custom
currencies should consist of only lowercase letters.

    btc = newcurrency!(:btc, "Bitcoin", 8)  # 1.00000000 BTC
"""
function newcurrency!(T::Symbol, name::AbstractString, expt::Int)
    ISO4217[T] = (expt, String(name), 0)
    Monetary(T)
end

"""
Add a new currency to the (global) currency list and assign a variable in the
local namespace to that currency's unit. Provide three arguments: an identifier
for the currency, a string description of the currency (following the
conventions outlined in the documentation for `currencyinfo`), and an exponent
representing the number of decimal points to describe the minor currency unit in
terms of the major currency unit. Conventionally, the identifer used to describe
custom currencies should consist of only lowercase letters.

    @usingcustomcurrency btc "Bitcoin" 8
    10btc  # 10.00000000 btc
"""
macro usingcustomcurrency(symb, name, exponent)
    quote
        $symb = newcurrency!($(Expr(:quote, symb)), $name, $exponent)
        nothing
    end |> esc
end


text:
using QuantumGate
using LinearAlgebra

zerostate = [1; 0]
onestate = [0; 1]

onestate ⊗ onestate ⊗ zerostate


text:
using ONNX, Flux, ProtoBuf
include("ops_tests.jl")

#Test Constant
main_test("$ONNX_TEST_PATH/test_constant", read_output("$ONNX_TEST_PATH/test_constant"))

#Test BatchNorm epsilon
main_test("$ONNX_TEST_PATH/test_batchnorm_epsilon", read_output("$ONNX_TEST_PATH/test_batchnorm_epsilon"),
    read_input("$ONNX_TEST_PATH/test_batchnorm_epsilon")[1] ,
    read_input("$ONNX_TEST_PATH/test_batchnorm_epsilon")[2],
    read_input("$ONNX_TEST_PATH/test_batchnorm_epsilon")[3],
    read_input("$ONNX_TEST_PATH/test_batchnorm_epsilon")[4],
    read_input("$ONNX_TEST_PATH/test_batchnorm_epsilon")[5])
"""
# Test BatchNorm example
main_test("$ONNX_TEST_PATH/test_batchnorm_example",
    read_output("$ONNX_TEST_PATH/test_batchnorm_example"),
    read_input("$ONNX_TEST_PATH/test_batchnorm_example")[1] ,
    read_input("$ONNX_TEST_PATH/test_batchnorm_example")[2],
    read_input("$ONNX_TEST_PATH/test_batchnorm_example")[3],
    read_input("$ONNX_TEST_PATH/test_batchnorm_example")[4],
    read_input("$ONNX_TEST_PATH/test_batchnorm_example")[5])
"""


text:
function to_str(input)
    out = []
    for i in input
        push!(out, parse(Int, i))
    end
    return out
end

function process_input(adapters)
    sort!(adapters)

    prepend!(adapters, 0)
    append!(adapters, adapters[end] + 3)

    println(is_valid(adapters))
    count_jumps(adapters)
    # count_combs(adapters)
end

function jumps(adapters)
    jumps = []
    a = adapters[1]
    for i in adapters[2:length(adapters)]
        push!(jumps, i-a)
        a = i
    end
    return jumps
end

# I was stuck on part 2, figured this out thanks to
# dynamic programming tip from  Jonathan Paulson
# https://youtu.be/cE88K2kFZn0
comb_dict = Dict()
function num_comb(i)
    println(adapters)
    if i == length(adapters)
        return 1
    end

    if haskey(comb_dict, i)
        return comb_dict[i]
    end

    n = 0
    for j = (i+1):length(adapters)
        if adapters[j] - adapters[i] <= 3
            n += num_comb(j)
        end
    end
    comb_dict[i] = n
    return n
end

# adapters = [16, 10, 15, 5, 1, 11, 7, 19, 6, 12, 4]
# adapters = to_str(readlines("./test_input.txt"))
adapters = to_str(readlines("./adapters.txt"))

process_input(adapters)
println("Number of combinations: ", num_comb(1))


text:
module InvertedIndices

export InvertedIndex, Not

"""
    InvertedIndex(idx)
    Not(idx)

Construct an inverted index, selecting all indices not in the passed `idx`.

Upon indexing into an array, the `InvertedIndex` behaves like a 1-dimensional
collection of the indices of the array that are not in `idx`. Bounds are
checked to ensure that all indices in `idx` are within the bounds of the array
— even though they are skipped. The `InvertedIndex` behaves like a
1-dimensional collection of its inverted indices. If `idx` spans multiple
dimensions (like a multidimensional logical mask or `CartesianIndex`), then the
inverted index will similarly span multiple dimensions.
"""
struct InvertedIndex{T}
    skip::T
end
const Not = InvertedIndex

# This is a little tricky because trues isn't indices-aware, but we also don't
# want to use fill(true) in the 1-indexed case since we want to favor BitArrays
@inline trues(tup::Tuple{Vararg{Base.OneTo}}) = Base.trues(map(Base.unsafe_length, tup))
@inline trues(tup::Tuple{Vararg{Base.AbstractUnitRange}}) = fill(true, tup)

@inline function Base.to_indices(A, inds, I::Tuple{InvertedIndex, Vararg{Any}})
    v = trues(spanned_indices(inds, I))
    v[I[1].skip] = false
    to_indices(A, inds, (v, Base.tail(I)...))
end

# Determining the indices that the InvertedIndex spans is tricky due to partial
# linear indexing. Lean on `Base.uncolon` until the deprecation goes through.
@inline spanned_indices(inds, I::Tuple{InvertedIndex,Vararg{Any}}) = (Base.uncolon(inds, (:, Base.tail(I)...)).indices,)

NIdx{N} = Union{CartesianIndex{N}, AbstractArray{CartesianIndex{N}}, AbstractArray{Bool,N}}
@inline spanned_indices(inds, I::Tuple{InvertedIndex{<:NIdx{0}},Vararg{Any}}) = ()
@inline spanned_indices(inds, I::Tuple{InvertedIndex{<:NIdx{1}},Vararg{Any}}) = (Base.uncolon(inds, (:, Base.tail(I)...)).indices,)
@inline function spanned_indices{N}(inds, I::Tuple{InvertedIndex{<:NIdx{N}},Vararg{Any}})
    heads, tails = Base.IteratorsMD.split(inds, Val{N})
    (Base.front(heads)..., Base.uncolon((heads[end], tails...), (:, Base.tail(I)...)).indices)
end

end # module


text:
# Autogenerated wrapper script for CLIME_jll for aarch64-linux-musl
export lime_contents, lime_extract_record, lime_extract_type, lime_pack, lime_unpack

JLLWrappers.@generate_wrapper_header("CLIME")
JLLWrappers.@declare_executable_product(lime_contents)
JLLWrappers.@declare_executable_product(lime_extract_record)
JLLWrappers.@declare_executable_product(lime_extract_type)
JLLWrappers.@declare_executable_product(lime_pack)
JLLWrappers.@declare_executable_product(lime_unpack)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        lime_contents,
        "bin/lime_contents",
    )

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

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

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

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

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
export RandomPolicy

"Policy that selects actions uniformly at random."
struct RandomPolicy{R <: AbstractRNG} <: PolicySolution
    domain::Domain
    rng::R
end

RandomPolicy(domain) = RandomPolicy(domain, Random.GLOBAL_RNG)

get_action(sol::RandomPolicy, state::State) =
    rand_action(sol, state)
rand_action(sol::RandomPolicy, state::State) =
    rand(sol.rng, available(state, sol.domain))


text:
function ImGui_ImplGlfw_ErrorCallback(code::Cint, description::Ptr{Cchar})::Cvoid
    @error "GLFW ERROR: code $code msg: $(unsafe_string(description))"
    return nothing
end

function ImGui_ImplGlfw_MouseButtonCallback(window::GLFW.Window, button, action, mods)
    io::Ptr{ImGuiIO} = igGetIO()
    ctx = unsafe_pointer_to_objref(unsafe_load(io.UserData))
    if ctx.PrevUserCallbackMousebutton != C_NULL
        ccall(ctx.PrevUserCallbackMousebutton, Cvoid, (Ptr{Cvoid}, Cint, Cint, Cint), window.handle, button, action, mods)
    end

    button = Cint(button)
    if action == GLFW.PRESS && button ≥ 0 && button < length(ctx.MouseJustPressed)
        ctx.MouseJustPressed[button+1] = true
    end

    return nothing
end

function ImGui_ImplGlfw_ScrollCallback(window::GLFW.Window, xoffset, yoffset)
    io::Ptr{ImGuiIO} = igGetIO()
    ctx = unsafe_pointer_to_objref(unsafe_load(io.UserData))
    if ctx.PrevUserCallbackScroll != C_NULL
        ccall(ctx.PrevUserCallbackScroll, Cvoid, (Ptr{Cvoid}, Cdouble, Cdouble), window.handle, xoffset, yoffset)
    end

    io.MouseWheelH = unsafe_load(io.MouseWheelH) + Cfloat(xoffset)
    io.MouseWheel = unsafe_load(io.MouseWheel) + Cfloat(yoffset)

    return nothing
end

function ImGui_ImplGlfw_KeyCallback(window::GLFW.Window, key, scancode, action, mods)
    io::Ptr{ImGuiIO} = igGetIO()
    ctx = unsafe_pointer_to_objref(unsafe_load(io.UserData))
    if ctx.PrevUserCallbackKey != C_NULL
        ccall(ctx.PrevUserCallbackKey, Cvoid, (Ptr{Cvoid}, Cint, Cint, Cint, Cint), window.handle, key, scancode, action, mods)
    end

    key = Cint(key)
    if key ≥ 0 && key < length(unsafe_load(io.KeysDown))
        if action == GLFW.PRESS
            c_set!(io.KeysDown, key, true)
            ctx.KeyOwnerWindows[key+1] = window
        end
        if action == GLFW.RELEASE
            c_set!(io.KeysDown, key, false)
            ctx.KeyOwnerWindows[key+1] = GLFW.Window(C_NULL)
        end
    end

    # modifiers are not reliable across systems
    io.KeyCtrl = c_get(io.KeysDown, GLFW.KEY_LEFT_CONTROL) || c_get(io.KeysDown, GLFW.KEY_RIGHT_CONTROL)
    io.KeyShift = c_get(io.KeysDown, GLFW.KEY_LEFT_SHIFT) || c_get(io.KeysDown, GLFW.KEY_RIGHT_SHIFT)
    io.KeyAlt = c_get(io.KeysDown, GLFW.KEY_LEFT_ALT) || c_get(io.KeysDown, GLFW.KEY_RIGHT_ALT)
    if Sys.iswindows()
        io.KeySuper = false
    else
        io.KeySuper = c_get(io.KeysDown, GLFW.KEY_LEFT_SUPER) || c_get(io.KeysDown, GLFW.KEY_RIGHT_SUPER)
    end

    return nothing
end

function ImGui_ImplGlfw_CharCallback(window::GLFW.Window, x)
    io::Ptr{ImGuiIO} = igGetIO()
    ctx = unsafe_pointer_to_objref(unsafe_load(io.UserData))
    if ctx.PrevUserCallbackChar != C_NULL
        ccall(ctx.PrevUserCallbackChar, Cvoid, (Ptr{Cvoid}, Cuint), window, x)
    end

    0 < Cuint(x) < 0x10000 && ImGuiIO_AddInputCharacter(io, x)

    return nothing
end

function ImGui_ImplGlfw_MonitorCallback(monitor::Ptr{Cvoid}, x::Cint)
    io::Ptr{ImGuiIO} = igGetIO()
    ctx = unsafe_pointer_to_objref(unsafe_load(io.UserData))
    ctx.WantUpdateMonitors = true
    return nothing
end

function ImGui_ImplGlfw_WindowCloseCallback(window::GLFW.Window)
    viewport::Ptr{ImGuiViewport} = igFindViewportByPlatformHandle(window.handle)
    if viewport != C_NULL
        viewport.PlatformRequestClose = true
    end
    return nothing
end

function ImGui_ImplGlfw_WindowPosCallback(window::GLFW.Window, x, y)
    viewport::Ptr{ImGuiViewport} = igFindViewportByPlatformHandle(window.handle)
    if viewport != C_NULL
        data::Ptr{ImGuiViewportDataGlfw} = unsafe_load(viewport.PlatformUserData)
        if data != C_NULL
            ignore_event = igGetFrameCount() ≤ (unsafe_load(data.IgnoreWindowPosEventFrame) + 1)
            ignore_event && return nothing
        end
        viewport.PlatformRequestMove = true
    end
    return nothing
end

function ImGui_ImplGlfw_WindowSizeCallback(window::GLFW.Window, x, y)
    viewport::Ptr{ImGuiViewport} = igFindViewportByPlatformHandle(window.handle)
    if viewport != C_NULL
        data::Ptr{ImGuiViewportDataGlfw} = unsafe_load(viewport.PlatformUserData)
        if data != C_NULL
            ignore_event = igGetFrameCount() ≤ (unsafe_load(data.IgnoreWindowSizeEventFrame) + 1)
            ignore_event && return nothing
        end
        viewport.PlatformRequestResize = true
    end
    return nothing
end


text:
import KnetLayers: _batchSizes2indices, _pack_sequence, arrtype
#####
##### Abstract Type Definitions
#####

"""
    Model

    Abstract type for all the models in Morse Package
    see `subtypes(Model)`
"""
abstract type Model;end;

"""
    Sequential <: Model

    Abstract type for all sequential models in Morse
    see `subtypes(Sequential)`
"""
abstract type Sequential <: Model;end;

function predict(M::Model, sentence::String; v::Vocabulary, p::Parser, maxL=20)
    encoded_ios = map(w->EncodedIO(collect(w),v=v,mask=v.specialTokens.unk,len=maxL), split(sentence))
    d = miniBatch(encoded_ios, v)
    preds,_ = predict(M, d; v=v)
    return map(enumerate(d.encodedIOs)) do (i,encodedIO)
        if isa(M,Sequential)
           @inbounds pred  = StringAnalysis(encodedIO.chars, preds[i,:], v=v)
        else
           @inbounds pred  = StringAnalysis(encodedIO.chars, preds[i], v=v)
        end
        makeFormat(pred,p)
    end
end

"""
    Discriminative <: Model

    Abstract type for all discriminative models in Morse
    see `subtypes(Discriminative)`
"""
abstract type Discriminative <: Model;end;

#####
##### Primitives
#####

"""
    WordEncoder(config::Dict,v::Vocabulary)
    (M::WordEncoder)(input::Vector{Int}, batchSizes::Vector{Int},
                     unsortingIndices::Vector{Int}; training=false)

    The `WordEncoder` described in the paper.
"""
struct WordEncoder
    embed::Embed
    dropout::Dropout
    encoder::LSTM
end

WordEncoder(o::Dict,v::Vocabulary) = WordEncoder(
    Embed(input=length(v.chars), output=o[:embedSizes][1]),
    Dropout(p=o[:dropouts]),
    LSTM(input=o[:embedSizes][1], hidden=o[:hiddenSizes][1], seed=o[:seed]) |> fgbias!
)

# function (M::WordEncoder)(inputs::Vector{Int}, sizes::Vector{Int}, indices::Vector{Int}; training=false)
#     out = M.encoder(M.dropout(M.embed(inputs), enable=training), batchSizes=sizes, hy=true, cy=true)
#     return out.hidden[:,:,end][:,indices], out.memory[:,:,end][:,indices]
# end

function (M::WordEncoder)(inputs::Vector{Int}, sizes::Vector{Int}, indices::Vector{Int}; training=false)
    if value(M.encoder.params) isa KnetArray
        out = M.encoder(M.dropout(M.embed(inputs), enable=training), batchSizes=sizes, hy=true, cy=true)
        return out.hidden[:,:,end][:,indices], out.memory[:,:,end][:,indices]
    else
        hs=[]; cs=[];
        for ind in _batchSizes2indices(sizes)
            out = M.encoder(reshape(M.embed(inputs[ind]),size(M.embed.weight,1),1,length(ind)); hy=true, cy=true)
            push!(hs, reshape(out.hidden,size(out.hidden,1),1))
            push!(cs, reshape(out.memory,size(out.memory,1),1))
        end
        return hcat(hs[indices]...), hcat(cs[indices]...)
    end
end

"""
    ContextEncoder(config::Dict,v::Vocabulary)
    (M::ContextEncoder)(hiddens; training=false)

    The `ContextEncoder` described in the paper.
"""
struct ContextEncoder
    encoder::LSTM
    reducer::Dense
    dropout::Dropout
end

ContextEncoder(o::Dict,v::Vocabulary) = ContextEncoder(
    LSTM(input=o[:hiddenSizes][1], hidden=o[:hiddenSizes][2], bidirectional=true, seed=o[:seed]) |> fgbias!,
    Dense(input=2o[:hiddenSizes][2],output=o[:hiddenSizes][3], activation=ReLU()),
    Dropout(p=o[:dropouts])
)

function (M::ContextEncoder)(hiddens; training=false)
    input = reshape(hiddens,size(hiddens,1),1,size(hiddens,2))
    y = M.encoder(M.dropout(input; enable=training)).y
    h = M.reducer(reshape(y,size(y,1),size(y,3)))
    return h, zero(value(h)) #TO-DO: zeros
end

"""
    OutputEncoder(config::Dict,v::Vocabulary)
    (M::OutputEncoder)(tagEmbeddings::Vector)
    (M::OutputEncoder)(emb, h, c)

    The `OutputEncoder` described in the paper.
"""
struct OutputEncoder
    encoder::LSTM
    dropout::Dropout
    previous::Int
    hidden
end

OutputEncoder(o::Dict,v::Vocabulary) = OutputEncoder(
    LSTM(input=o[:embedSizes][2], hidden=o[:hiddenSizes][4], seed=o[:seed]) |> fgbias!,
    Dropout(p=o[:dropouts]),
    o[:previous],
    arrtype(zeros(o[:hiddenSizes][4],1,1))
)

function (M::OutputEncoder)(emb, h, c)
    y = M.encoder(emb, h, c; hy=true, cy=true)
    return y.hidden, y.memory
end

function (M::OutputEncoder)(tagEmbeddings::Vector)
    hos, cos = [], [];
    H,_,_ = size(M.hidden) # assume one layer output encoder
    for i=1:length(tagEmbeddings)
        ho, co = zero(M.hidden), zero(M.hidden) #TO-DO zeros
        for j = max(i-M.previous,1):i-1
             out = M.encoder(M.dropout(tagEmbeddings[j],enable=true), ho, co; hy=true, cy=true)
             ho, co = out.hidden, out.memory
        end
        push!(hos,reshape(ho,H,1))
        push!(cos,reshape(co,H,1))
    end
    return hcat(hos...), hcat(cos...)
end

"""
    Decoder(config::Dict,v::Vocabulary)
    (M::Decoder)(input, hiddens1, hiddens2; training=false)

    The two layer `Decoder` described in the paper.
"""
struct Decoder
    L1::LSTM
    L2::LSTM
    dropout::Dropout
end

function make3D(x)
    if ndims(x) < 3
        return reshape(x,size(x)...,1)
    end
    return x
end

function (M::Decoder)(input, hiddens1, hiddens2; training=false)
    out1   = M.L1(M.dropout(input; enable=training), make3D.(hiddens1)...)
    out2   = M.L2(M.dropout(out1.y; enable=training), make3D.(hiddens2)...)
    return (out1.hidden, out1.memory),
           (M.dropout(out2.hidden; enable=training), out2.memory)
end

Decoder(o::Dict,v::Vocabulary) = Decoder(
    LSTM(input=o[:embedSizes][2], hidden=o[:hiddenSizes][3], seed=o[:seed]) |> fgbias!,
    LSTM(input=o[:hiddenSizes][3], hidden=o[:hiddenSizes][3], seed=o[:seed]) |> fgbias!,
    Dropout(p=o[:dropouts])
)

"""
    recurrentLoss(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)

    Calculate recurrent negative likelihood for sequential models
"""
function recurrentLoss(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)
    total = 0.0
    input = fill(startIndex, size(d.seqOutputs,1))
    for i=1:size(d.seqOutputs,2)
        hiddens1, hiddens2 = M.decoder(M.outputEmbed(input), hiddens1, hiddens2; training=true)
        h2 = hiddens2[1]
        scores = M.output(reshape(h2,size(h2,1),size(h2,2)))
        # loss calculation
        total += M.loss(scores, d.seqOutputs[:,i] .* d.masks[:,i]; average=false)
        input  = d.seqOutputs[:,i]
    end
    return total
end

"""
    recurrentPredict(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)

    Make prediciton with recurrent decoder for sequential models except MorseModel due to its `OutputEncoder`
"""
function recurrentPredict(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)
    wordNumber, timeSteps = size(d.seqOutputs)
    preds = zeros(Int, wordNumber, timeSteps)
    input = fill(startIndex, wordNumber)
    total = 0.0
    for i=1:timeSteps
        hiddens1, hiddens2 = M.decoder(M.outputEmbed(input), hiddens1, hiddens2; training=false)
        h2 = hiddens2[1]
        scores = M.output(reshape(h2,size(h2,1),size(h2,2)))
        # prediction
        input = vec(getindex.(argmax(convert(Array,scores),dims=1),1))
        preds[:,i] = input
        #loss
        total += M.loss(scores, d.seqOutputs[:,i] .* d.masks[:,i]; average=false)
    end
    return preds, total
end

#####
##### Sequential Models
#####

"""
    MorseModel <: Sequential
    MorseModel(config::Dict,v::Vocabulary)
    loss(M::MorseModel, d::SentenceBatch; v::Vocabulary)
    predict(M::MorseModel, d::SentenceBatch; v::Vocabulary)

    The `MorseModel` described in the paper
"""
struct MorseModel <: Sequential
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEncoder::OutputEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

MorseModel(o::Dict,v::Vocabulary) = MorseModel(
    WordEncoder(o,v), ContextEncoder(o,v), OutputEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::MorseModel, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    # output encoder
    esize = size(M.outputEmbed.weight,1);
    tagEmbeddings = [reshape(M.outputEmbed(d.seqOutputs[i,range]),esize,1,length(range)) for (i,range)=enumerate(d.tagRange)]
    hiddens3 = M.outputEncoder(tagEmbeddings)
    # Decoder
    hiddens2 = hiddens2 .+ hiddens3
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow)/sum(d.masks)
end

function predict(M::MorseModel, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    # Decoder
    total = 0.0
    wordNumber, timeSteps = size(d.seqOutputs)
    preds = zeros(Int, wordNumber, timeSteps)
    hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))

    outEmbeddings = [];
    for w=1:wordNumber
        dh2 = (hiddens2[1][:,w:w] .+ hiddens3[1], hiddens2[2][:,w:w] .+ hiddens3[2])
        dh1 = (hiddens1[1][:,w:w], hiddens1[2][:,w:w])

        input = v.specialIndices.bow
        outEmbedding = [];
        continueStoring = true

        for t=1:timeSteps
            rnninput = M.outputEmbed([input])
            if continueStoring && input == v.specialIndices.eow
                continueStoring = false
            elseif continueStoring && input ∉ v.specialIndices && length(v.tags[input]) > 1
                push!(outEmbedding,rnninput)
            end
            dh1, dh2 = M.decoder(rnninput, dh1, dh2; training=false)
            h2 = dh2[1]
            scores = M.output(reshape(h2, size(h2,1), size(h2,2)))
            #prediction
            input = argmax(convert(Array,scores))[1]
            preds[w,t] = input
            # loss calculation
            total += M.loss(scores, d.seqOutputs[w:w,t] .* d.masks[w:w,t]; average=false)
        end

        if length(outEmbedding) != 0

            embedding = hcat(outEmbedding...)
            push!(outEmbeddings, reshape(embedding, size(embedding,1), 1, size(embedding,2)))

            if length(outEmbeddings) > M.outputEncoder.previous
                popfirst!(outEmbeddings);
            end

            if length(outEmbeddings) > 0
                hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))
                for embs in outEmbeddings
                    hiddens3 = M.outputEncoder(embs, hiddens3[1], hiddens3[2])
                end
            end
        end
    end
    return preds, total
end

(M::MorseModel)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)

"""
    S2S <: Sequential
    S2S(o::Dict,v::Vocabulary)
    loss(M::S2S, d::SentenceBatch; v::Vocabulary)
    predict(M::S2S, d::SentenceBatch; v::Vocabulary)

    Ablation of `MorseModel` where `ContextEncoder` and `OutputEncoder` removed.
"""
struct S2S <: Sequential
    wordEncoder::WordEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

S2S(o::Dict,v::Vocabulary) = S2S(
    WordEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::S2S, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    recurrentLoss(M, d, (nothing, nothing), hiddens2; startIndex=v.specialIndices.bow) / sum(d.masks)
end

function predict(M::S2S, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    recurrentPredict(M, d, (nothing, nothing), hiddens2; startIndex=v.specialIndices.bow)
end

(M::S2S)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)
"""
    S2SContext <: Sequential
    S2SContext(config::Dict,v::Vocabulary)
    loss(M::S2SContext, d::SentenceBatch; v::Vocabulary)
    predict(M::S2SContext, d::SentenceBatch; v::Vocabulary)

    Ablation of the `MorseModel` where `OutputEncoder` removed
"""
struct S2SContext <: Sequential
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

S2SContext(o::Dict,v::Vocabulary) = S2SContext(
    WordEncoder(o,v),
    ContextEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::S2SContext, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow) / sum(d.masks)
end

function predict(M::S2SContext, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    recurrentPredict(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow)
end

(M::S2SContext)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)
#####
##### Discriminative Models
#####

"""
    Classifier <: Discriminative
    Classifier(config::Dict,v::Vocabulary)
    loss(M::Classifier, d::SentenceBatch; v::Vocabulary)
    predict(M::Classifier, d::SentenceBatch; v::Vocabulary)

    Ablation of MorseModel where `OutputEncoder` removed and `Decoder` replaced with a classifer.
"""
struct Classifier <: Discriminative
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    output::Linear
    loss::CrossEntropyLoss
end

Classifier(o::Dict,v::Vocabulary) = Classifier(
    WordEncoder(o,v),
    ContextEncoder(o,v),
    Linear(input=o[:hiddenSizes][3]+o[:hiddenSizes][1], output=length(v.comptags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::Classifier, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    scores = M.output(vcat(hiddens1[1],hiddens2[1]))
    M.loss(scores,d.compositeOutputs)
end

function predict(M::Classifier, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    scores   = M.output(vcat(hiddens1[1],hiddens2[1]))
    preds    = vec(getindex.(argmax(convert(Array,scores),dims=1),1))
    loss     = M.loss(scores,d.compositeOutputs)
    return preds,loss
end

(M::Classifier)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)
#####
##### Disambiguator Models
#####

abstract type Disambiguator <: Sequential; end;

struct S2SContexDis <: Disambiguator
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

S2SContexDis(o::Dict,v::Vocabulary) = S2SContexDis(
    WordEncoder(o,v),
    ContextEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::S2SContexDis, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow) / sum(d.masks)
end

function predict(M::S2SContexDis, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    recurrentPredict(M, d, hiddens1, hiddens2; v=v)
end

"""
    recurrentPredict(M::Disambiguator, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)

    Make prediciton with recurrent decoder for disambiguation models except MorseDis
"""
function recurrentPredict(M::Disambiguator, d::SentenceBatch, hiddens1, hiddens2; v::Vocabulary)

    preds = Vector{Int}[]
    total = 0.0

    for (w,encodedIO) in enumerate(d.encodedIOs)

        min_loss = typemax(Float64)
        min_loss_seq = Int[];

        for (i,analysis) in enumerate(encodedIO.analyses)

            h1 = (hiddens1[1][:,w:w],hiddens1[2][:,w:w])
            h2 = (hiddens2[1][:,w:w],hiddens2[2][:,w:w])

            seq = [analysis.lemma; analysis.tags; v.specialIndices.eow]

            wtotal = 0.0
            input = v.specialIndices.bow

            for o in seq
                h1, h2 = M.decoder(M.outputEmbed([input]), h1, h2; training=false)
                hh2 = h2[1]
                scores = M.output(reshape(hh2,size(hh2,1),size(hh2,2)))
                wtotal += M.loss(scores,[o])
                input   = o
            end

            loss = wtotal
            if loss < min_loss
                min_loss = loss; min_loss_seq=seq;
            end
        end
        push!(preds,min_loss_seq)
        total+=min_loss
    end

    return PadSequenceArray(preds; pad=v.specialIndices.mask),total
    #TO-DO: total causes incorrect loos in evaluate.jl
end

"""
    MorseDis <: Disambiguator
    MorseDis(config::Dict,v::Vocabulary)
    loss(M::MorseDis, d::SentenceBatch; v::Vocabulary)
    predict(M::MorseDis, d::SentenceBatch; v::Vocabulary)
"""
struct MorseDis <: Disambiguator
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEncoder::OutputEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

MorseDis(o::Dict,v::Vocabulary) = MorseDis(
    WordEncoder(o,v), ContextEncoder(o,v), OutputEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::MorseDis, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    # output encoder
    esize = size(M.outputEmbed.weight,1);
    tagEmbeddings = [reshape(M.outputEmbed(d.seqOutputs[i,range]),esize,1,length(range)) for (i,range)=enumerate(d.tagRange)]
    hiddens3 = M.outputEncoder(tagEmbeddings)
    # Decoder
    hiddens2 = hiddens2 .+ hiddens3
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow)/sum(d.masks)
end

function  predict(M::MorseDis, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    # Output Encoder
    hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))
    # Decoder

    outEmbeddings = [];
    preds = Vector{Int}[]
    total = 0.0

    for (w,encodedIO) in enumerate(d.encodedIOs)

        min_loss = typemax(Float64)
        min_loss_seq::Union{Nothing,Vector{Int}} = nothing;
        min_outEmbed::Union{Nothing,Vector} = nothing;

        for (i,analysis) in enumerate(encodedIO.analyses)

            h2 = (hiddens2[1][:,w:w] .+ hiddens3[1], hiddens2[2][:,w:w] .+ hiddens3[2])
            h1 = (hiddens1[1][:,w:w],hiddens1[2][:,w:w])

            wtotal = 0.0
            input = v.specialIndices.bow
            outEmbedding = [];
            seq = [analysis.lemma; analysis.tags; v.specialIndices.eow]

            for o in seq
                rnninput = M.outputEmbed([input])
                 if input ∉ v.specialIndices && length(v.tags[input])>1
                    push!(outEmbedding,rnninput)
                end
                h1, h2 = M.decoder(rnninput, h1, h2; training=false)
                hh2 = h2[1]
                scores = M.output(reshape(hh2,size(hh2,1),size(hh2,2)))
                wtotal += M.loss(scores,[o])
                input   = o
            end

            loss = wtotal
            if loss < min_loss
                min_loss = loss; min_loss_seq=seq; min_outEmbed=outEmbedding;
            end
        end

        push!(preds,min_loss_seq)
        total += min_loss

        if length(min_outEmbed) != 0

            embedding = hcat(min_outEmbed...)
            push!(outEmbeddings, reshape(embedding, size(embedding,1), 1, size(embedding,2)))

            if length(outEmbeddings) > M.outputEncoder.previous
                popfirst!(outEmbeddings);
            end

            if length(outEmbeddings) > 0
                hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))
                for embs in outEmbeddings
                    hiddens3 = M.outputEncoder(embs, hiddens3[1], hiddens3[2])
                end
            end
        end
    end
    return PadSequenceArray(preds; pad=v.specialIndices.mask),total
end

#####
##### Utils
#####
function fgbias!(m::LSTM)
    for inputSource in (:i,:h), layer in (:10,:11)
        bias = get(m.gatesview, Symbol(:b,inputSource,:f,layer), nothing)
        if bias !== nothing
            bias .= 0.5
        end
    end
    return m
end

function transfer!(target, tv::Vocabulary, source, sv::Vocabulary)
    for f in (:wordEncoder, :decoder, :contextEncoder, :outputEmbed, :output, :outputEncoder)
        if isdefined(target,f) && isdefined(source,f)
            transfer!(getproperty(target,f), tv, getproperty(source,f), sv)
        end
    end
end

function transfer!(target::Embed, tv::Vocabulary, source::Embed, sv::Vocabulary)
    transfer!(target.weight, tv.tags, source.weight, sv.tags; axis=2)
end

function transfer!(target::Linear, tv::Vocabulary, source::Linear, sv::Vocabulary)
    transfer!(target.mult.weight, tv.tags, source.mult.weight, sv.tags; axis=1)
    transfer!(target.bias, tv.tags, source.bias, sv.tags; axis=1)
end

function transfer!(target, tv::IndexedDict, source, sv::IndexedDict; axis=1)
    for (k,i) in sv.toIndex
        if (j=get(tv,k,nothing)) !== nothing
            if ndims(target) == 2
                if axis == 1
                    target[j,:] = source[i,:]
                else
                    target[:,j] = source[:,i]
                end
            else
                 target[j] = source[i]
            end
        end
    end
end

function transfer!(source::Decoder, sv::Vocabulary, target::Decoder, tv::Vocabulary)
    transfer!(source.L1,target.L1)
    transfer!(source.L2,target.L2)
end

function transfer!(source::ContextEncoder,sv::Vocabulary,target::ContextEncoder,tv::Vocabulary)
    transfer!(source.encoder,target.encoder)
    transfer!(source.reducer,target.reducer)
end

function transfer!(source::OutputEncoder,sv::Vocabulary,target::OutputEncoder,tv::Vocabulary)
    transfer!(source.encoder,target.encoder)
end

function transfer!(target::WordEncoder, tv::Vocabulary, source::WordEncoder, sv::Vocabulary)
    transfer!(target.embed.weight, tv.chars, source.embed.weight, sv.chars; axis=2)
    transfer!(target.encoder, source.encoder)
end

function transfer!(target::Linear, source::Linear)
    transfer!(target.mult,source.mult)
    transfer!(target.bias,source.bias)
end
transfer!(target::Multiply, source::Multiply) = transfer!(target.weight,source.weight)
transfer!(target::Dense, source::Dense) = transfer!(target.linear,source.linear)
transfer!(target::LSTM, source::LSTM) = transfer!(target.params,source.params)
transfer!(target::Param, source::Param) = copyto!(value(target),value(source))