text:
##########################################
#            CUSTOM TYPES                #
##########################################

struct PrimalDualVars{T}
    X::VectorTrajectory{T}
    U::VectorTrajectory{T}
    λ::Matrix{Vector{T}}
end

# function PrimalDualVars(prob::Problem)
#     n,m,N = size(prob)
#     X = prob.X
#     U = prob.U
#
#     p = num_constraints(prob)
#     y_part = dual_partition(n,m,p,N)
#     λ = [ zeros(y) for y in y_part, i=1:1]
#     PrimalDualVars(X,U,λ)
# end
#
# function copy(V::PrimalDualVars)
#     PrimalDualVars(deepcopy(V.X), deepcopy(V.U), deepcopy(V.λ))
# end
#
struct KKTFactors{T}
    E::MatrixTrajectory{T}
    F::MatrixTrajectory{T}
    G::Vector{Cholesky{T,Matrix{T}}}
    K::MatrixTrajectory{T}
    L::MatrixTrajectory{T}
    M::MatrixTrajectory{T}
    H::Vector{Cholesky{T,Matrix{T}}}
end
#
# function KKTFactors(n,m,p_active,N)
#     E = [zeros(n,n) for p in p_active]
#     F = [zeros(n,p) for p in p_active]
#     G = [cholesky(Matrix(I,n,n)) for p in p_active]
#
#     K = [zeros(p,n) for p in p_active]
#     L = [zeros(p,p) for p in p_active]
#     M = [zeros(p,n) for p in p_active]
#     H = [cholesky(Matrix(I,p,p)) for p in p_active]
#     KKTFactors(E,F,G,K,L,M,H)
# end
#
#
#
struct SequentialNewtonSolver{T} <: DirectSolver{T}
    opts::ProjectedNewtonSolverOptions{T}
    stats::Dict{Symbol,Any}
    V::PrimalDualVars{T}
    V_::PrimalDualVars{T}
    δx::VectorTrajectory{T}
    δu::VectorTrajectory{T}
    δλ::Matrix{Vector{T}}
    r::Matrix{Vector{T}}
    L::KKTFactors{T}

    Q::Vector{Expansion{T,Diagonal{T,Vector{T}},Diagonal{T,Vector{T}}}}
    Qinv::Vector{Diagonal{T,Vector{T}}}
    Rinv::Vector{Diagonal{T,Vector{T}}}
    fVal::Vector{Vector{T}}
    ∇F::Vector{PartedArray{T,2,Matrix{T},P}} where P
    C::PartedVecTrajectory{T}
    ∇C::Vector{PartedArray{T,2,Matrix{T},P} where P}
    active_set::Vector{Vector{Bool}}
    p::Vector{Int}
end
#
# function SequentialNewtonSolver(prob::Problem{T}, opts::ProjectedNewtonSolverOptions{T}) where T
#     n,m,N = size(prob)
#
#     NN = N*n + (N-1)*m
#     p = num_constraints(prob)
#     pcum = insert!(cumsum(p),1,0)
#     P = sum(p) + N*n
#
#     # Partitions
#     part_a = (primals=1:NN, duals=NN+1:NN+P) #, ν=NN .+ (1:N*n), λ=NN + N*n .+ (1:sum(p)))
#     part_f = create_partition2(prob.model)
#     y_part = dual_partition(n,m,p,N)
#
#     # Options
#     stats = Dict{Symbol,Any}()
#
#     # PrimalDuals Variables
#     V = PrimalDualVars(copy(prob))
#     V_ = copy(V)
#
#     # Deviations
#     δx = [zeros(n) for k = 1:N]
#     δu = [zeros(m) for k = 1:N-1]
#     δλ = [zeros(y) for y in y_part, i=1:1]
#     r  = [zeros(y) for y in y_part, i=1:1]
#
#     # KKT Factors
#     L = KKTFactors(n,m,p,N)
#
#     # Costs
#     Q = [Expansion{T,Diagonal,Diagonal}(n,m) for k = 1:N]
#     Qinv = [Diagonal(zeros(T,n,n)) for k = 1:N]
#     Rinv = [Diagonal(zeros(T,m,m)) for k = 1:N-1]
#
#     # Constraints
#     fVal = [zeros(n) for k = 1:N]
#     ∇F = [PartedMatrix(zeros(n,n+m+1),part_f) for k = 1:N]
#
#     constraints = prob.constraints
#     C = [PartedVector(con) for con in constraints.C]
#     ∇C = [PartedMatrix(con,n,m) for con in constraints.C]
#     C[N] = PartedVector(constraints[N],:terminal)
#     ∇C[N] = PartedMatrix(constraints[N],n,m,:terminal)
#
#     # Active Set
#     active_set = [ones(Bool,pk) for pk in p]
#
#     SequentialNewtonSolver(opts, stats, V, copy(V), δx, δu, δλ, r, L, Q, Qinv, Rinv, fVal, ∇F, C, ∇C, active_set, p)
# end
#
# ##########################################
# #        TYPE UTIL FUNCTIONS             #
# ##########################################
#
# function size(solver::SequentialNewtonSolver)
#     n,m = length(solver.δx[1]), length(solver.δu[1])
#     N = length(solver.Q)
#     return n,m,N
# end
#
# function num_active_constraints(solver::SequentialNewtonSolver)
#     n,m,N = size(solver)
#     sum(sum.(solver.active_set)) + N*n
# end
#
# """
# Construct 2N partition of dual variables (n,n,p1,n,p2,n,p3,...,n,pN-1,pN)
# """
# function dual_partition(n,m,p,N)
#     y_part = ones(Int,2,N-1)*n
#     y_part[2,:] = p[1:end-1]
#     y_part = vec(y_part)
#     insert!(y_part,1,3)
#     push!(y_part, p[N])
#     return y_part
# end
#
# function dual_partition(solver::SequentialNewtonSolver)
#     n,m,N = size(solver)
#     p_active = sum.(solver.active_set)
#     dual_partition(n,m,p_active,N)
# end
#
# function update!(prob::Problem, solver::SequentialNewtonSolver, V=solver.V, active_set=true)
#     dynamics_constraints!(prob, solver, V)
#     dynamics_jacobian!(prob, solver, V)
#     update_constraints!(prob, solver, V)
#     constraint_jacobian!(prob, solver, V)
#     if active_set
#         active_set!(prob, solver)
#     end
#     cost_expansion!(prob, solver, V)
#     invert_hessian!(prob, solver)
#     nothing
# end
#
# function active_duals(V::PrimalDualVars{T}, a)::Vector{SubArray{T,1,Vector{T},Tuple{Vector{Int}},false}} where T
#     N = length(V.X)
#     n = length(V.X[1])
#     dyn = collect(1:n)
#     λ = [begin
#             if i == 1 || iseven(i)
#                 view(V.λ[i],dyn)
#             elseif isodd(i) || i == 2N
#                 k = (i-1)÷2
#                 view(V.λ[i],a[k])
#             end
#         end for i = 1:2N
#         ]
#     return λ
# end
#
# """
# Take the 2-norm of the residual: [g + Y'λ; y]
# """
# function res2(solver::SequentialNewtonSolver, V)
#     N = length(solver.Q)
#
#     # Calculate g + Y'λ
#     λa = active_duals(V, solver.active_set)
#     δx,δu = jac_T_mult(solver, λa)
#     for k = 1:N-1
#         δx[k] += solver.Q[k].x
#         δu[k] += solver.Q[k].u
#     end
#     δx[N] += solver.Q[N].x
#     # return δx, δu, y
#     y = active_constraints(solver)
#     return sqrt(norm(δx)^2 + norm(δu)^2 + norm(y)^2)
# end
#
# function packZ(X,U)
#     N = length(X)
#     z = [[x;u] for (x,u) in zip(X[1:end-1],U)]
#     push!(z,X[N])
#     vcat(z...)
# end
#
#
# #########################################
# #         CONSTRAINT FUNCTIONS          #
# #########################################
#
# function dynamics_jacobian!(prob::Problem, ∇F, X, U)
#     n,m,N = size(prob)
#     ∇F[1].xx .= Diagonal(I,n)
#     dt = get_dt_traj(prob,U)
#     for k = 1:N-1
#         jacobian!(∇F[k+1], prob.model, X[k], U[k], dt[k])
#     end
# end
# dynamics_jacobian!(prob::Problem, solver::SequentialNewtonSolver, V=solver.V) =
#     dynamics_jacobian!(prob, solver.∇F, V.X, V.U)
#
#
# function active_set!(prob::Problem, solver::SequentialNewtonSolver)
#     n,m,N = size(prob)
#     for k = 1:N
#         active_set!(solver.active_set[k], solver.C[k], solver.opts.active_set_tolerance)
#     end
# end
#
# """
# Get the 2N vector of active constraints
# """
# function active_constraints!(y, solver::SequentialNewtonSolver)
#     N = length(solver.Q)
#     y[1] = solver.fVal[1]
#     a = solver.active_set
#
#     for k = 1:N-1
#         y[2k] = solver.fVal[k+1]
#         y[2k+1] = solver.C[k][a[k]]
#     end
#     y[end] = solver.C[N][a[N]]
#     nothing
# end
# function active_constraints(solver::SequentialNewtonSolver)
#     y_part = dual_partition(solver)
#     y = [zeros(y) for y in y_part]
#     active_constraints!(y, solver)
#     return y
# end
#
#
# #########################################
# #       COST EXPANSION FUNCTIONS        #
# #########################################
#
# function cost_expansion!(prob::Problem, solver::SequentialNewtonSolver, V=solver.V)
#     N = prob.N
#     X,U, dt = V.X, V.U, get_dt_traj(prob,V.U)
#     for k = 1:N-1
#         cost_expansion!(solver.Q[k], prob.obj[k], X[k], U[k], dt[k])
#     end
#     cost_expansion!(solver.Q[N], prob.obj[N], X[N])
# end
#
# function invert_hessian!(prob::Problem, solver::SequentialNewtonSolver)
#     N = prob.N
#     for k = 1:N-1
#         solver.Qinv[k] = inv(solver.Q[k].xx)
#         solver.Rinv[k] = inv(solver.Q[k].uu)
#     end
#     solver.Qinv[N] = inv(solver.Q[N].xx)
# end
#
#
#
# #########################################
# #     CORE LINEAR ALGEBRA FUNCTIONS     #
# #########################################
#
# """
# Calculate r = Y*z, store result in r, where z is split into x and u
# """
# function jac_mult(solver::SequentialNewtonSolver, x, u)
#     y_part = dual_partition(solver)
#     r = [zeros(y) for y in y_part]
#
#     n,m,N = size(solver)
#
#     # Extract variables from solver
#     Q = solver.Q
#     Qinv = solver.Qinv
#     Rinv = solver.Rinv
#     a = solver.active_set
#     ∇F = view(solver.∇F,2:N)
#     fVal = view(solver.fVal,2:N)
#     ∇C = solver.∇C
#     c = solver.C
#
#     xi,ui = 1:n, n .+ (1:m)
#     C = [∇C[k][a[k],xi] for k = 1:N]
#     D = [∇C[k][a[k],ui] for k = 1:N-1]
#
#     # Calculate residual
#     r[1] = x[1]
#     for k = 1:N-1
#         r[2k] = (∇F[k].xx*x[k] + ∇F[k].xu*u[k] - x[k+1])
#         r[2k+1] = C[k]*x[k] + D[k]*u[k]
#     end
#     r[end] = C[N]*x[N]
#     return r
# end
#
# """
# Calculate x,u = Y'λ
# """
# function jac_T_mult(solver::SequentialNewtonSolver, λ)
#     n,m,N = size(solver)
#     x = [zeros(n) for k = 1:N]
#     u = [zeros(m) for k = 1:N-1]
#
#     Q = solver.Q
#     ∇F = view(solver.∇F,2:N)
#     ∇C = solver.∇C
#     a = solver.active_set
#
#     xi,ui = 1:n, n .+ (1:m)
#     C = [∇C[k][a[k],xi] for k = 1:N]
#     D = [∇C[k][a[k],ui] for k = 1:N-1]
#
#     x[1] = λ[1] + ∇F[1].xx'λ[2] + C[1]'λ[3]
#     u[1] =        ∇F[1].xu'λ[2] + D[1]'λ[3]
#     for k = 2:N-1
#         i = 2k
#         x[k] = -λ[i-2] + ∇F[k].xx'λ[i] + C[k]'λ[i+1]
#         u[k] =           ∇F[k].xu'λ[i] + D[k]'λ[i+1]
#     end
#     x[end] = -λ[end-2] + C[N]'λ[end]
#     return x,u
# end
#
# """
# Compute the cholesky factorization of (Y*Hinv*Y')
# where Y is the constraint jacobian and Hinv is the inverse of the cost Hessian
# """
# function calc_factors!(solver::SequentialNewtonSolver, Qinv=solver.Qinv, Rinv=solver.Rinv)
#     n,m,N = size(solver)
#     Nb = 2N  # number of blocks
#     p_active = sum.(solver.active_set)
#     Pa = sum(p_active)
#     y_parts = dual_partition(solver)
#
#     # Extract arrays
#     S_factors = solver.L
#     E = S_factors.E
#     F = S_factors.F
#     G = S_factors.G
#
#     K = S_factors.K
#     L = S_factors.L
#     M = S_factors.M
#     H = S_factors.H
#
#
#     # Extract variables from solver
#     a = solver.active_set
#     Q = solver.Q
#     fVal = view(solver.fVal, 2:N)
#     ∇F = view(solver.∇F, 2:N)
#     c = solver.C
#     ∇C = solver.∇C
#
#     # Get Active Jacobians
#     xi,ui = 1:n, n .+ (1:m)
#     C = [∇C[k][a[k],1:n] for k = 1:N]
#     D = [∇C[k][a[k],n+1:n+m] for k = 1:N-1]
#
#     # Initial condition
#     if Qinv[1] isa UniformScaling
#         G0 = cholesky_reg(Array(Diagonal(I,n)))
#     else
#         G0 = cholesky_reg(Array(Qinv[1]))
#     end
#
#     F[1] = ∇F[1].xx*G0.L
#     _G = ∇F[1].xu*Rinv[1]*∇F[1].xu' + Qinv[2]
#     G[1] = cholesky(Symmetric(_G))
#
#     L[1] = C[1]*G0.L
#     M[1] = D[1]*Rinv[1]*∇F[1].xu'/(G[1].U)
#     H[1] = cholesky(Symmetric(D[1]*Rinv[1]*D[1]' - M[1]*M[1]'))
#
#     G_ = G[1]
#     M_ = M[1]
#     H_ = H[1]
#     i = 4
#     for k = 2:N-1
#         E[k] = -∇F[k].xx*Qinv[k]/G_.U
#         F[k] = -E[k]*M_'/H_.U
#         G[k] = cholesky_reg(Symmetric(∇F[k].xx*Qinv[k]*∇F[k].xx' + ∇F[k].xu*Rinv[k]*∇F[k].xu' + Qinv[k+1] - E[k]*E[k]' - F[k]*F[k]'))
#         G[k].info != 0 ? println("failed G cholesky at k = $k") : nothing
#         i += 1
#
#         K[k] = -C[k]*Qinv[k]/G_.U
#         L[k] = -K[k]*M_'/H_.U
#         M[k] = (C[k]*Qinv[k]*( solver.Q[k].xx - G_.U\(I - (M_'/H_.U) * (H_.L\M_) )/G_.L )*Qinv[k]*∇F[k].xx' + D[k]*Rinv[k]*∇F[k].xu')/G[k].U
#         H[k] = cholesky_reg(C[k]*Qinv[k]*C[k]' + D[k]*Rinv[k]*D[k]' - K[k]*K[k]' - L[k]*L[k]' - M[k]*M[k]')
#         H[k].info != 0 ? println("failed H cholesky at k = $k") : nothing
#         i += 1
#
#         G_, M_, H_ = G[k], M[k], H[k]
#     end
#
#     # Terminal
#     E[N] = -C[N]*Qinv[N]/G_.U
#     F[N] = -E[N]*M_'/H_.U
#     G[N] = cholesky_reg(Symmetric(C[N]*Qinv[N]*C[N]' - E[N]*E[N]' - F[N]*F[N]'))
#     G[N].info != 0 ? println("failed G cholesky at k = N") : nothing
#
#     # Append G0 onto the end
#     push!(G, G0)
#
#     return nothing
# end
#
# function cholesky_reg(A::AbstractMatrix)
#     C = cholesky(A, check=true)
#     if C.info != 0 && false
#         E = eigen(A)
#         v = min(minimum(E.values),-1e-2)
#         cholesky(A - 2I*v)
#     end
#     return C
# end
#
#
# """
# Solve the system S*δλ = r
#     where S = (Y*Hinv*Y') and is already factored and stored in a `KKTFactors` type
# """
# function solve_cholesky(solver::SequentialNewtonSolver, r)
#     n,m,N = size(solver)
#     Nb = 2N  # number of blocks
#     p_active = sum.(solver.active_set)
#     Pa = sum(p_active)
#     y_part = dual_partition(solver)
#
#
#     # Init arrays
#     S_factors = solver.L
#     E = S_factors.E
#     F = S_factors.F
#     G = S_factors.G
#
#     K = S_factors.K
#     L = S_factors.L
#     M = S_factors.M
#     H = S_factors.H
#
#
#     λ = [zeros(y) for y in y_part]
#     λ_ = deepcopy(λ)
#
#     # Extract variables from solver
#     a = solver.active_set
#     Q = solver.Q
#     fVal = view(solver.fVal, 2:N)
#     ∇F = view(solver.∇F, 2:N)
#     c = solver.C
#     ∇C = solver.∇C
#
#     # Initial condition
#     p = p_active[1]
#     G0 = G[end]
#     λ_[1] = G0.L\r[1]
#     λ_[2] = G[1].L\(r[2] - F[1]*λ_[1])
#     λ_[3] = H[1].L\(r[3] - M[1]*λ_[2] - L[1]*λ_[1])
#
#     i = 4
#     for k = 2:N-1
#         λ_[i] = G[k].L\(r[i] - F[k]*λ_[i-1] - E[k]*λ_[i-2])
#         i += 1
#
#         λ_[i] = H[k].L\(r[i] - M[k]*λ_[i-1] - L[k]*λ_[i-2] - K[k]*λ_[i-3])
#         i += 1
#     end
#
#     # Terminal
#     λ_[i] = G[N].L\(r[i] - F[N]*λ_[i-1] - E[N]*λ_[i-2])
#
#
#     # BACK SUBSTITUTION
#     λ[Nb] = G[N].U\λ_[Nb]
#     λ[Nb-1] = H[N-1].U\(λ_[Nb-1] - F[N]'λ[Nb])
#     λ[Nb-2] = G[N-1].U\(λ_[Nb-2] - M[N-1]'λ[Nb-1] - E[N]'λ[Nb])
#
#     i = Nb-3
#     for k = N-2:-1:1
#         # println("\n Time step $k")
#         λ[i] = H[k].U\(λ_[i] - F[k+1]'λ[i+1] - L[k+1]'λ[i+2])
#         i -= 1
#         λ[i] = G[k].U\(λ_[i] - M[k]'λ[i+1] - E[k+1]'λ[i+2] - K[k+1]'λ[i+3])
#         i -= 1
#     end
#     λ[1] = G0.U\(λ_[1] - F[1]'λ[2] - L[1]'λ[3])
#     return λ
# end
#
#
#
# ######################
# #   UPDATE METHODS   #
# ######################
#
# function _update_primals!(V::PrimalDualVars, δx, δu)
#     N = length(V.X)
#     for k = 1:N-1
#         V.X[k] += δx[k]
#         V.U[k] += δu[k]
#     end
#     V.X[N] += δx[N]
#     return nothing
# end
#
# function +(V::PrimalDualVars{T}, δz::NTuple{2,Vector{Vector{T}}}) where T
#     V_ = copy(V)
#     _update_primals!(V_, δz[1], δz[2])
#     return V_
# end
#
# function _update_duals!(V::PrimalDualVars, δλ, active_set)
#     N = length(V.X)
#     V.λ[1] += δλ[1]
#     for k = 1:N-1
#         a = active_set[k]
#         V.λ[2k] += δλ[2k]
#         λk = view(V.λ[2k+1], a)
#         λk .+= δλ[2k+1]
#     end
#     λN = view(V.λ[2N], active_set[N])
#     λN .+= δλ[2N]
# end
#
# function +(V::PrimalDualVars, δλ::Tuple{Vector{Vector{T}},Vector{Vector{Bool}}}) where T
#     δλ,a = δλ
#     V_ = copy(V)
#     _update_duals!(V_, δλ, a)
#     return V_
# end
#
# function +(V::PrimalDualVars,
#         δV::Tuple{NTuple{3,Vector{Vector{T}}},Vector{Vector{Bool}}}) where T
#     δV,a = δV
#     δx,δu,δλ = δV
#     V_ = copy(V)
#     _update_primals!(V_, δx, δu)
#     _update_duals!(V_, δλ, a)
#     return V_
# end
# # function solve_cholesky(prob::Problem, solver::SequentialNewtonSolver, r,
# #         Qinv=solver.Qinv, Rinv=solver.Rinv)
# #     n,m,N = size(prob)
# #     Nb = 2N  # number of blocks
# #     p_active = sum.(solver.active_set)
# #     Pa = sum(p_active)
# #
# #     x,u = SVector(1:n...), SVector(1:m...)
# #
# #     # Init arrays
# #     S_factors = solver.L
# #     E = S_factors.E
# #     F = S_factors.F
# #     G = S_factors.G
# #
# #     K = S_factors.K
# #     L = S_factors.L
# #     M = S_factors.M
# #     H = S_factors.H
# #
# #     r = solver.r
# #     λ = solver.δλ
# #     λ_ = deepcopy(λ)
# #
# #     # Extract variables from solver
# #     a = solver.active_set
# #     Q = solver.Q
# #     fVal = view(solver.fVal, 2:N)
# #     ∇F = view(solver.∇F, 2:N)
# #     c = solver.C
# #     ∇C = solver.∇C
# #
# #     # Initial condition
# #     p = p_active[1]
# #     if Qinv[1] isa UniformScaling
# #         G0 = cholesky(Diagonal(I,n))
# #     else
# #         G0 = cholesky(Qinv[1])
# #     end
# #     λ_[1] = G0.L\r[1]
# #
# #     F[1] = ∇F[1].xx*G0.L
# #     _G = ∇F[1].xu*Rinv[1]*∇F[1].xu' + Qinv[2]
# #     G[1] = cholesky(Symmetric(_G))
# #     λ_[2] = G[1].L\(r[2] - F[1]*λ_[1])
# #
# #     C = ∇C[1].x[a[1],x]
# #     D = ∇C[1].u[a[1],u]
# #     L[1] = C*G0.L
# #     M[1] = D*Rinv[1]*∇F[1].xu'/(G[1].U)
# #     H[1] = cholesky(Symmetric(D*Rinv[1]*D' - M[1]*M[1]'))
# #     λ_[3] = H[1].L\(r[3] - M[1]*λ_[2] - L[1]*λ_[1])
# #
# #     G_ = G[1]
# #     M_ = M[1]
# #     H_ = H[1]
# #     i = 4
# #     for k = 2:N-1
# #         # println("\n Time Step $k")
# #         p = p_active[k]
# #         p_ = p_active[k-1]
# #
# #         C = (∇C[k].x[solver.active_set[k],:])
# #         D = (∇C[k].u[solver.active_set[k],:])
# #
# #         E[k] = -∇F[k].xx*Qinv[k]/G_.U
# #         F[k] = -E[k]*M_'/H_.U
# #         G[k] = cholesky(Symmetric(∇F[k].xx*Qinv[k]*∇F[k].xx' + ∇F[k].xu*Rinv[k]*∇F[k].xu' + Qinv[k+1] - E[k]*E[k]' - F[k]*F[k]'))
# #         λ_[i] = G[k].L\(r[i] - F[k]*λ_[i-1] - E[k]*λ_[i-2])
# #         i += 1
# #
# #         K[k] = -C*Qinv[k]/G_.U
# #         L[k] = -K[k]*M_'/H_.U
# #         M[k] = (C*Qinv[k]*( solver.Q[k].xx - G_.U\(I - (M_'/H_.U) * (H_.L\M_) )/G_.L )*Qinv[k]*∇F[k].xx' + D*Rinv[k]*∇F[k].xu')/G[k].U
# #         H[k] = cholesky(C*Qinv[k]*C' + D*Rinv[k]*D' - K[k]*K[k]' - L[k]*L[k]' - M[k]*M[k]')
# #         λ_[i] = H[k].L\(r[i] - M[k]*λ_[i-1] - L[k]*λ_[i-2] - K[k]*λ_[i-3])
# #         i += 1
# #
# #         G_, M_, H_ = G[k], M[k], H[k]
# #     end
# #
# #     # Terminal
# #     p = p_active[N]
# #     p_ = p_active[N-1]
# #
# #     C = ∇C[N].x[a[N],:]
# #     D = ∇C[N].u[a[N],:]
# #
# #     E[N] = -C*Qinv[N]/G_.U
# #     F[N] = -E[N]*M_'/H_.U
# #     G[N] = cholesky(Symmetric(C*Qinv[N]*C' - E[N]*E[N]' - F[N]*F[N]'))
# #     λ_[i] = G[N].L\(r[i] - F[N]*λ_[i-1] - E[N]*λ_[i-2])
# #
# #     # return λ_
# #
# #     # BACK SUBSTITUTION
# #     λ[Nb] = G[N].U\λ_[Nb]
# #     λ[Nb-1] = H[N-1].U\(λ_[Nb-1] - F[N]'λ[Nb])
# #     λ[Nb-2] = G[N-1].U\(λ_[Nb-2] - M[N-1]'λ[Nb-1] - E[N]'λ[Nb])
# #
# #     i = Nb-3
# #     for k = N-2:-1:1
# #         # println("\n Time step $k")
# #         λ[i] = H[k].U\(λ_[i] - F[k+1]'λ[i+1] - L[k+1]'λ[i+2])
# #         i -= 1
# #         λ[i] = G[k].U\(λ_[i] - M[k]'λ[i+1] - E[k+1]'λ[i+2] - K[k+1]'λ[i+3])
# #         i -= 1
# #     end
# #     λ[1] = G0.U\(λ_[1] - F[1]'λ[2] - L[1]'λ[3])
# #
# #     return λ, λ_, r
# #
# # end
# #
# # function solve_cholesky_static(prob::Problem, solver::SequentialNewtonSolver, vals)
# #     n,m,N = size(prob)
# #     Nb = 2N  # number of blocks
# #     p_active = sum.(solver.active_set)
# #     Pa = sum(p_active)
# #
# #     x,u = SVector(1:n...), SVector(1:m...)
# #
# #     # Init arrays
# #     E = vals.E
# #     F = vals.F
# #     G = vals.G
# #
# #     K = vals.K
# #     L = vals.L
# #     M = vals.M
# #     H = vals.H
# #
# #     r = vals.r
# #     λ_ = vals.λ_
# #     λ = vals.λ
# #
# #     # Extract variables from solver
# #     a = solver.active_set
# #     a = vals.a
# #     Q = solver.Q
# #     Qinv = solver.Qinv
# #     Rinv = solver.Rinv
# #     fVal = view(solver.fVal, 2:N)
# #     ∇F = view(solver.∇F, 2:N)
# #     c = solver.C
# #     ∇C = solver.∇C
# #
# #     # Calculate residual
# #     r[1] = solver.fVal[1] - Qinv[1]*Q[1].x
# #     for k = 1:N-1
# #         C = ∇C[k].x[a[k],x]
# #         D = ∇C[k].u[a[k],u]
# #         r[2k] = fVal[k] - (∇F[k].xx*Qinv[k]*Q[k].x + ∇F[k].xu*Rinv[k]*Q[k].u - Qinv[k+1]*Q[k+1].x)
# #         r[2k+1] = c[k][a[k]] - (C*Qinv[k]*Q[k].x + D*Rinv[k]*Q[k].u)
# #     end
# #     C = ∇C[N].x[a[N],:]
# #     r[end] = c[N][a[N]] - C*Qinv[N]*Q[N].x
# #
# #
# #     # Initial condition
# #     p = p_active[1]
# #     G0 = cholesky(Qinv[1])
# #     λ_[1] = G0.L\r[1]
# #
# #     F[1] = SMatrix{n,n}(∇F[1].xx*G0.L)
# #     _G = SMatrix{n,n}(∇F[1].xu*Rinv[1]*∇F[1].xu' + Qinv[2])
# #     G[1] = cholesky(Symmetric(_G))
# #     λ_[2] = G[1].L\(r[2] - F[1]*λ_[1])
# #
# #     C = ∇C[1].x[a[1],x]
# #     D = ∇C[1].u[a[1],u]
# #     L[1] = SMatrix{p,n}(C*G0.L)
# #     M[1] = SMatrix{p,n}(D*Rinv[1]*∇F[1].xu'/(G[1].U))
# #     _H = D*Rinv[1]*D' - M[1]*M[1]'
# #     H[1] = cholesky(Symmetric(_H))
# #     λ_[3] = H[1].L\(r[3] - M[1]*λ_[2] - L[1]*λ_[1])
# #
# #     G_ = G[1]
# #     M_ = M[1]
# #     H_ = H[1]
# #     i = 4
# #     for k = 2:N-1
# #         # println("\n Time Step $k")
# #         p = p_active[k]
# #         p_ = p_active[k-1]
# #
# #         C = (∇C[k].x[a[k],x])
# #         D = (∇C[k].u[a[k],u])
# #         # C = (∇C[k].x[solver.active_set[k],:])
# #         # D = (∇C[k].u[solver.active_set[k],:])
# #
# #         E[k] = -∇F[k].xx*Qinv[k]/G_.U
# #         if p_ > 0
# #             F[k] = SMatrix{n,p_}(-E[k]*M_'/H_.U)
# #             G[k] = cholesky(Symmetric(∇F[k].xx*Qinv[k]*∇F[k].xx' + ∇F[k].xu*Rinv[k]*∇F[k].xu' + Qinv[k+1] - E[k]*E[k]' - F[k]*F[k]'))
# #             λ_[i] = G[k].L\(r[i] - F[k]*λ_[i-1] - E[k]*λ_[i-2])
# #         else
# #             F[k] = @SMatrix zeros(n,0)
# #             G[k] = cholesky(Symmetric(∇F[k].xx*Qinv[k]*∇F[k].xx' + ∇F[k].xu*Rinv[k]*∇F[k].xu' + Qinv[k+1] - E[k]*E[k]'))
# #             λ_[i] = G[k].L\(r[i] - E[k]*λ_[i-2])
# #         end
# #         i += 1
# #
# #         if p > 0 #|| true
# #             K[k] = SMatrix{p,n}(-C*Qinv[k]/G[k].U)
# #             Q = solver.Q[k].xx
# #             if p_ > 0
# #                 L[k] = SMatrix{p,p_}(-K[k]*M_'/H_.U)
# #                 _M = (C*Qinv[k]*( Q - G_.U\(I - (M_'/H_.U) * (H_.L\M_) )/G_.L )*Qinv[k]*∇F[k].xx' + D*Rinv[k]*∇F[k].xu')/G[k].U
# #                 M[k] = SMatrix{p,n}(_M)
# #                 H[k] = cholesky(C*Qinv[k]*C' + D*Rinv[k]*D' - K[k]*K[k]' - L[k]*L[k]' - M[k]*M[k]')
# #                 λ_[i] = H[k].L\(r[i] - M[k]*λ_[i-1] - L[k]*λ_[i-2] - K[k]*λ_[i-3])
# #             else
# #                 L[k] = @SMatrix zeros(p,0)
# #                 _M = (C*Qinv[k]*( Q - G_.U\(I)/G_.L )*Qinv[k]*∇F[k].xx' + D*Rinv[k]*∇F[k].xu')/G[k].U
# #                 M[k] = SMatrix{p,n}(_M)
# #                 H[k] = cholesky(C*Qinv[k]*C' + D*Rinv[k]*D' - K[k]*K[k]' - M[k]*M[k]')
# #                 λ_[i] = H[k].L\(r[i] - M[k]*λ_[i-1] - K[k]*λ_[i-3])
# #             end
# #         else
# #             L[k] = @SMatrix zeros(p_active[k],p_active[k-1])
# #         end
# #         i += 1
# #
# #         G_, M_, H_ = G[k], M[k], H[k]
# #     end
# #
# #     # Terminal
# #     p = p_active[N]
# #     p_ = p_active[N-1]
# #
# #     C = ∇C[N].x[a[N],:]
# #     D = ∇C[N].u[a[N],:]
# #
# #     E[N] = -C*Qinv[N]/G_.U
# #     F[N] = SMatrix{p,p_}(-E[N]*M_'/H_.U)
# #     G[N] = cholesky(Symmetric(C*Qinv[N]*C' - E[N]*E[N]' - F[N]*F[N]'))
# #     λ_[i] = G[N].L\(r[i] - F[N]*λ_[i-1] - E[N]*λ_[i-2])
# #
# #     # return λ_
# #
# #     # BACK SUBSTITUTION
# #     λ[Nb] = G[N].U\λ_[Nb]
# #     λ[Nb-1] = H[N-1].U\(λ_[Nb-1] - F[N]'λ[Nb])
# #     λ[Nb-2] = G[N-1].U\(λ_[Nb-2] - M[N-1]'λ[Nb-1] - E[N]'λ[Nb])
# #
# #     i = Nb-3
# #     for k = N-2:-1:1
# #         # println("\n Time step $k")
# #         if p_active[k] > 0
# #             if p_active[k+1] > 0
# #                 λ[i] = H[k].U\(λ_[i] - F[k+1]'λ[i+1] - L[k+1]'λ[i+2])
# #             else
# #                 λ[i] = H[k].U\(λ_[i] - F[k+1]'λ[i+1])
# #             end
# #         end
# #         i -= 1
# #         if p_active[k] > 0
# #             if p_active[k+1] > 0
# #                 λ[i] = G[k].U\(λ_[i] - M[k]'λ[i+1] - E[k+1]'λ[i+2] - K[k+1]'λ[i+3])
# #             else
# #                 λ[i] = G[k].U\(λ_[i] - M[k]'λ[i+1] - E[k+1]'λ[i+2])
# #             end
# #         else
# #             λ[i] = G[k].U\(λ_[i] - E[k+1]'λ[i+2] - K[k+1]'λ[i+3])
# #         end
# #         i -= 1
# #     end
# #     λ[1] = G0.U\(λ_[1] - F[1]'λ[2] - L[1]'λ[3])
# #
# #     return λ, λ_, r
# #
# # end
# #
# #
# # function set_active_set!(solver::SequentialNewtonSolver)
# #     N = length(solver.δx)
# #     p = sum.(solver.active_set)
# #     y_part = ones(Int,2,N-1)*n
# #     y_part[2,:] = p[1:end-1]
# #     y_part = vec(y_part)
# #     insert!(y_part,1,3)
# #     push!(y_part, p[N])
# #     for i = 1:2N
# #         solver.δλ[i] = zeros(y_part[i])
# #         solver.r[i] = zeros(y_part[i])
# #     end
# # end
# #

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
# one batch: vector of matrices with shape (features,batchsize)
# [rand(features,batchsize) for 1:seq_len]

function get_dl(; seq_len=32, batchsize=16)
    filepath = joinpath(@__DIR__, "half-cheetah-data")
    data_dir = filepath

    obs_size = size(npzread("$(data_dir)/trace_0000.npy"),2) # 17

    all_files = readdir(data_dir)
    all_files = ["$(data_dir)/$(f)" for f in all_files]
    train_files = all_files[16:25]
    test_files = all_files[6:15]
    valid_files = all_files[1:5]

    # goal:
    # size(first(train_dl)[1]) == [(17,16) for 1:32]
    # x_data == [ [(17,16) for 1:32] for 1:60 ]

    # sequences = (32,17,100)
    # sequences = [(32,17) for 1:100]
    #

    train_x, train_y = _load_files(train_files, seq_len)
    test_x, test_y = _load_files(test_files, seq_len)
    valid_x, valid_y = _load_files(valid_files, seq_len)

    # (970×32×17)
    train_x = permutedims(train_x,(3,1,2))
    train_y = permutedims(train_y,(3,1,2))
    test_x = permutedims(test_x,(3,1,2))
    test_y = permutedims(test_y,(3,1,2))
    valid_x = permutedims(valid_x,(3,1,2))
    valid_y = permutedims(valid_y,(3,1,2))
    # (17x970x32)

    train_x_new = train_y_new = Vector{Matrix{Float32}}[]
    test_x_new = test_y_new = Vector{Matrix{Float32}}[]
    valid_x_new = valid_y_new = Vector{Matrix{Float32}}[]
    # train_x_new = train_y_new = []
    # test_x_new = test_y_new = []
    # valid_x_new = valid_y_new = []
    for i in 1:batchsize:size(train_x,2)-batchsize-1
        push!(train_x_new, Flux.unstack(train_x[:,i:i+batchsize-1,:],3))
        push!(train_y_new, Flux.unstack(train_y[:,i+1:i+batchsize,:],3))
        # push!(train_x_new, train_x[:,i:i+batchsize-1,:])
        # push!(train_y_new, train_y[:,i+1:i+batchsize,:])
    end
    for i in 1:batchsize:size(test_x,2)-batchsize-1
        push!(test_x_new, Flux.unstack(test_x[:,i:i+batchsize-1,:],3))
        push!(test_y_new, Flux.unstack(test_y[:,i+1:i+batchsize,:],3))
        # push!(test_x_new, test_x[:,i:i+batchsize-1,:])
        # push!(test_y_new, test_y[:,i+1:i+batchsize,:])
    end
    for i in 1:batchsize:size(valid_x,2)-batchsize-1
        push!(valid_x_new, Flux.unstack(valid_x[:,i:i+batchsize-1,:],3))
        push!(valid_y_new, Flux.unstack(valid_y[:,i+1:i+batchsize,:],3))
        # push!(valid_x_new, valid_x[:,i:i+batchsize-1,:])
        # push!(valid_y_new, valid_y[:,i+1:i+batchsize,:])
    end

    # train_x_new = Flux.stack(train_x_new,4)
    # train_y_new = Flux.stack(train_y_new,4)
    # test_x_new = Flux.stack(test_x_new,4)
    # test_y_new = Flux.stack(test_y_new,4)
    # valid_x_new = Flux.stack(valid_x_new,4)
    # valid_y_new = Flux.stack(valid_y_new,4)

    #train_x = Flux.unstack(train_x,3)

    # train_dl = Flux.Data.DataLoader((train_x,train_y),batchsize=1)
    # test_dl = Flux.Data.DataLoader((test_x,test_y),batchsize=batchsize)
    # valid_dl = Flux.Data.DataLoader((valid_x,valid_y),batchsize=batchsize)
    # train_dl = Flux.Data.DataLoader((train_x_new,train_y_new),batchsize=1)
    # test_dl = Flux.Data.DataLoader((test_x_new,test_y_new),batchsize=1)
    # valid_dl = Flux.Data.DataLoader((valid_x_new,valid_y_new),batchsize=1)

    # @show size(train_x_new)
    # @show size(train_x_new[1])
    # @show size(train_x_new[1][1])
    #
    # @show size(train_x_new[2])
    # @show size(train_x_new[2][1])

    train_dl = zip(train_x_new, test_y_new) #|> f32
    test_dl = zip(test_x_new, test_y_new) #|> f32
    valid_dl = zip(valid_x_new, valid_y_new) #|> f32

    # train_dl = (train_x_new, test_y_new)
    # test_dl = (test_x_new, test_y_new)
    # valid_dl = (valid_x_new, valid_y_new)

    return train_dl, test_dl, valid_dl
end

function _load_files(files, seq_len)
    all_x = []
    all_y = []
    for f in files
        arr = Float32.(npzread(f))
        x, y = _cut_in_sequences(arr, seq_len, 10)
        push!(all_x, x...)
        push!(all_y, y...)
    end

    return Flux.stack(all_x,1), Flux.stack(all_y,1)
    #return all_x, all_y
end

function _cut_in_sequences(x, seq_len, inc=1)
    seq_x = []
    seq_y = []
    for s in 1:inc:size(x,1)-seq_len-1
        i_s = s
        i_e = i_s + seq_len - 1
        push!(seq_x, x[i_s:i_e,:])
        push!(seq_y, x[i_s+1:i_e+1,:])
    end
    return seq_x, seq_y
end

# train_dl = get_dl()[1][1]
# x = train_dl[1]

ext: jl


text:
    using AutoDiffSource
using Base.Test

function checkdiff_inferred(f, δf, x0...)
    x = [x0...]
    y0 = f(x...)
    @assert length(y0) == 1 "Scalar functions only"
    y, ∇f = Test.@inferred δf(x...)
    @assert isapprox(y0, y) "Return values do not match"
    @assert typeof(y0) === typeof(y) "Return type doesn't match"
    ∂x = Test.@inferred ∇f(1.0)
    if isa(∂x, Tuple)
        @assert typeof(x0) === typeof(∂x) "Gradient type doesn't match: $(typeof(x0)) vs $(typeof(∂x))"
    else
        @assert typeof(x0) === typeof((∂x,)) "Gradient type doesn't match: : $(typeof(x0)) vs $(typeof((∂x,)))"
    end
    checkgrad(f, x, ∂x)
end

# Test example
sigmoid(x) = 1 / (1 + exp(-x))
function autoencoder(We1, We2, Wd, b1, b2, input)
    firstLayer = sigmoid.(We1 * input + b1)
    encodedInput = sigmoid.(We2 * firstLayer + b2)
    reconstructedInput = sigmoid.(Wd * encodedInput)
    return reconstructedInput
end
@δ function autoencoderError(We1, We2, Wd, b1, b2, input)
    reconstructedInput = autoencoder(We1, We2, Wd, b1, b2, input)
    return sum((input - reconstructedInput).^2)
end
@assert checkdiff_inferred(autoencoderError, δautoencoderError, randn(3,3), randn(3,3), rand(3,3), randn(3), randn(3), randn(3))

# Array indexing
@δ rosenbrock(x, y) = sum(100*(y-x.^2).^2 + (1-x).^2)
@δ function rosenbrock(x)
    l = length(x)
    rosenbrock(x[1:l-1], x[2:l])
end
@assert checkdiff_inferred(rosenbrock, δrosenbrock, randn(3))

# diffentiate a 3rd party function without source code, recursively
@δ rosen2(x::Vector{Float64}, y::Vector{Float64}) = sum(100*(y-x.^2).^2 + (1-x).^2)
function rosen(x::Vector{Float64})
    l = length(x)
    rosen2(x[1:l-1], x[2:l])
end
@δ rosen
@assert checkdiff_inferred(rosen, δrosen, randn(3))

# check basic use
@δ f(x, y) = (x + y) * y
@test checkdiff(f, δf, 2., 3.)

# check numerical constants
@δ function f2(x, y::AbstractFloat)
    srand(1)
    z = 4.0*3.2*2.5x - y^2 + rand(size(x)) + randn(size(x))
    sum(z) / y + randn() + rand()
end
@test checkdiff_inferred(f2, δf2, rand(5), rand())
@test checkdiff_inferred(f2, δf2, rand(3,3), rand())

# test broadcast
@δ f3(x) = sum(abs.(x))
@test checkdiff_inferred(f3, δf3, rand(5)-0.5)

# test multi argument functions and reuse
@δ f4(x, y) =  x * y, x - y
@δ function f5(x)
    (a, b) = f4(x, x+3)
    a * x + 4b
end
@test checkdiff_inferred(f5, δf5, rand()-0.5)

# test external constants
const f6r = rand(5)
@δ f6(x) = sum(f6r .* x.^2)
@test checkdiff_inferred(f6, δf6, rand(5))

# test ...
@δ function f7(x...)
    a, b = x
    return a * b
end
@test checkdiff_inferred(f7, δf7, rand(2)...)

# test ...
const f8_const = rand(13)
@δ function f8(a, x)
    b, c, d, e, f, g, h, i, j, k, l, m, n = x
    a * b + c * d - e * f + g * h - i * j / k * l + m * n
end
@δ test_f8(x) = f8(x, f8_const)
@test checkdiff_inferred(test_f8, δtest_f8, rand())

# test matrix multiply
@δ f9(x, y) = sum(x * y)
@test checkdiff_inferred(f9, δf9, rand(3), rand(3)')
@test checkdiff_inferred(f9, δf9, rand(3)', rand(3))
#@test checkdiff_inferred(f9, δf9, rand(3)', rand(3, 3))
@test checkdiff_inferred(f9, δf9, rand(3, 3), rand(3))
@test checkdiff_inferred(f9, δf9, rand(3, 3), rand(3, 3))

# test sequence of plus and times
@δ f10(x, y, z) = (x + y + z) * x * y * z
@test checkdiff_inferred(f10, δf10, rand(), rand(), rand())

# test multiple return values
@δ f11(x, y) = (x*y, y/x)
@δ function f12(x, y)
    return x+y, y-x
end
@δ function f13(x, y)
    a, b = f11(y, y)
    c, d = f12(x, b)
    return a+b+c+d+x
end
@δ function f14(x)
    f13(x, 2.3) + x
end
@test checkdiff(f14, δf14, rand())

# check constants
@δ function f15(x, y::Float64)
    srand(1)
    z = x + rand()
    return z+y
end
@δ function f16(x)
    f15(x, 2.3) + x
end
@test checkdiff_inferred(f16, δf16, rand())

# test tuple constants
@δ function f17(x)
    z = x.^3 + x
    state_const = (z, x.^2)
    state_const, x.^3
end
@δ function test_f17(x)
    state_const, y = f17(x)
    sum(y)
end
@test checkdiff_inferred(test_f17, δtest_f17, rand())
@test checkdiff_inferred(test_f17, δtest_f17, rand(3))

# single return symbol
@δ function f18(x)
    y = x + x
    return y
end
@test checkdiff_inferred(f18, δf18, rand())

# single return symbol
@δ function f19(x)
    y = x + x
    y
end
@test checkdiff_inferred(f19, δf19, rand())

# reassignment inside
@δ function f20(x)
    x = x + x
    x = x * x
    x
end
@test checkdiff_inferred(f20, δf20, rand())

# test broadcast
@eval @δ f21(x, y) = mean(f13.(x, y))
@eval @test checkdiff_inferred(f21, δf21, rand(), rand())
@eval @test checkdiff_inferred(f21, δf21, rand(3), rand(3, 2))
@eval @test checkdiff_inferred(f21, δf21, rand(3, 2), rand(3))
@eval @test checkdiff_inferred(f21, δf21, rand(5), rand())
@eval @test checkdiff_inferred(f21, δf21, rand(3, 2), rand())
@eval @test checkdiff_inferred(f21, δf21, rand(), rand(5))
@eval @test checkdiff_inferred(f21, δf21, rand(), rand(3, 2))

# (scalar, scalar), (scalar, const), (const, scalar), (const, const)
for o in [:+, :-, :*, :/, :^, :min, :max]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = $o(x, y) + $o(1.0, 3.0)
    @eval @test checkdiff_inferred($t, $δt, rand(), rand())
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = $o(x, 2.)
    @eval @test checkdiff_inferred($t, $δt, rand())
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = $o(2., x)
    @eval @test checkdiff_inferred($t, $δt, rand())
end

# (scalar)
for o in [:abs, :sum, :sqrt, :exp, :log, :-, :sign, :log1p, :expm1,
          :sin, :cos, :tan, :sinh, :cosh, :tanh, :asin, :acos, :atan,
          :round, :floor, :ceil, :trunc, :mod2pi, :maximum, :minimum, :transpose,
          :erf, :erfc, :gamma, :lgamma]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = $o(x) + x
    @eval @test checkdiff_inferred($t, $δt, rand())
end

# (vector, vector), (matrix, matrix), (const, *), (*, const), (const, const)
# (vector, matrix), (matrix, vector), (vector, scalar), (matrix, scalar), (scalar, vector), (scalar, matrix)
for o in [:.+, :.-, :.*, :./, :.^]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = sum($o(x, y)+$o(1.0, 3.0))
    @eval @test checkdiff_inferred($t, $δt, rand(5), rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(x, 3.))
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(3., x))
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = sum($o(x, y))
    @eval @test checkdiff_inferred($t, $δt, rand(), rand())
    @eval @test checkdiff_inferred($t, $δt, rand(3), rand(3, 2))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand(3))
    @eval @test checkdiff_inferred($t, $δt, rand(5), rand())
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand())
    @eval @test checkdiff_inferred($t, $δt, rand(), rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(), rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(x, ones(1, 2)))
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
    @eval @test checkdiff_inferred($t, $δt, rand(1, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(5, 1))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(ones(1, 2), x))
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
    @eval @test checkdiff_inferred($t, $δt, rand(1, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(5, 1))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = sum($o(x, y))
    @eval @test checkdiff_inferred($t, $δt, rand(1, 1), rand(1, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 1), rand(3, 2))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand(3, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(5, 1), rand(1, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand(1, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(1, 1), rand(5, 1))
    @eval @test checkdiff_inferred($t, $δt, rand(1, 1), rand(3, 2))
end

# (scalar), (vector), (matrix)
for o in [:abs, :sqrt, :exp, :log, :sign, :log1p, :expm1,
          :sin, :cos, :tan, :sinh, :cosh, :tanh, :asin, :acos, :atan,
          :round, :floor, :ceil, :trunc, :mod2pi,
          :erf, :erfc, :gamma, :lgamma]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o.(x) + x)
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
end

# (vector, scalar), (matrix, scalar), (scalar, vector), (scalar, matrix), (const, *), (*, const)
for o in [:+, :-, :*]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = sum($o(x, y))
    @eval @test checkdiff_inferred($t, $δt, rand(5), rand())
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand())
    @eval @test checkdiff_inferred($t, $δt, rand(), rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(), rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(x, 4.))
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(5., x))
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
end

# (vector, scalar), (matrix, scalar), (*, const)
for o in [:/]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = sum($o(x, y))
    @eval @test checkdiff_inferred($t, $δt, rand(5), rand())
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand())

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(x, 5.))
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
end

for o in [:dot]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = $o(x, y)
    @eval @test checkdiff_inferred($t, $δt, rand(5), rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand(3, 2))
end

for o in [:transpose, :maximum, :minimum]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o(x))
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
end

# (vector, vector), (matrix, matrix), (const, *), (*, const)
for o in [:min, :max]
    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x, y) = sum($o.(x, y))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2), rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o.(x, 3.))
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))

    t = gensym(o)
    δt = Symbol("δ$t")
    @eval @δ $t(x) = sum($o.(3., x))
    @eval @test checkdiff_inferred($t, $δt, rand())
    @eval @test checkdiff_inferred($t, $δt, rand(5))
    @eval @test checkdiff_inferred($t, $δt, rand(3, 2))
end

# https://github.com/gaika/AutoDiffSource.jl/issues/15
@δ f_15(x, y) = (x + 2y) * y^2
v, ∇f_25 = δf_15(3,2)
@test v == 28.0 && ∇f_25() == (4.0, 36.0)


text:
using Chron
using Test, Statistics, StatsBase

include("testUtilities.jl")
include("testDist.jl")
include("testStratOnly.jl")
include("testRadiocarbon.jl")
include("testCoupled.jl")


text:
module VoronoiCells

import GeometryBasics
import VoronoiDelaunay
import VoronoiDelaunay: getx, gety

using RecipesBase

export
    Rectangle,
    Tessellation,

    left,
    lower,
    right,
    upper,
    voronoiarea,
    voronoicells

include("Points.jl")
include("Rectangle.jl")
include("Clipping.jl")
include("Sort.jl")
include("Cells.jl")
include("Plot.jl")
include("Area.jl")

const BoundingBox = Rectangle(
    VoronoiDelaunay.Point2D(VoronoiDelaunay.min_coord, VoronoiDelaunay.min_coord),
    VoronoiDelaunay.Point2D(VoronoiDelaunay.max_coord, VoronoiDelaunay.max_coord)
)

end # module


text:
"""
    MaxEntropyGraphs

Julia module for working with maximum entropy graphs
"""
module MaxEntropyGraphs

import Base: show
#import Graphs: degree
#import LinearAlgebra: diagind
import StatsBase: countmap

using Graphs
# using SimpleWeightedGraphs

using NLsolve
using IndirectArrays
using LoopVectorization
# MODELS
export UBCM, DBCM
export solve! #, UBCMCompact, solve!

# METRICS
#export degree

include("models.jl")
#include("metrics.jl")
end # module


text:
function rbf(γ::Array, X₁::Matrix, X₂::Matrix)::Matrix
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      K = Array{eltype(γ)}(undef, N₁, N₂)
      for j in 1:N₂
            for i in 1:N₁
                  sqd = zero(eltype(γ))
                  for k in 1:F
                        sqd += γ[k]*(X₁[k,i] - X₂[k,j])^2
                  end
                  K[i,j] = exp(-0.5*sqd)
            end
      end
      return K
end

function ∇ₓrbf(γ::Array, X₁::Matrix, X₂::Matrix)::Array  # derivative on X1
    F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
    ∇ₓK = Array{eltype(γ)}(undef, F, N₁, N₂)
    for j in 1:N₂
          for i in 1:N₁
                sqd = zero(eltype(γ))
                for m in 1:F
                      sqd += γ[m]*(X₁[m,i]-X₂[m,j])^2
                end
                for l in 1:F
                      ∇ₓK[l,i,j] = -γ[l]*(X₁[l,i]-X₂[l,j])*exp(-0.5*sqd)
                end
          end
    end
    return ∇ₓK
end

function linear(γ::Array, X₁::Matrix, X₂::Matrix)::Array
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      K = Array{eltype(γ)}(undef, N₁, N₂)
      for j in 1:N₂
            for i in 1:N₁
                  dot_prod = zero(eltype(γ))
                  for m in 1:F
                        dot_prod += γ[m]*X₁[m,i]*X₂[m,j]
                  end
                  K[i,j] = dot_prod + γ[F+1]
            end
      end
      return K
end

function ∇ₓlinear(γ::Array, X₁::Matrix, X₂::Matrix)::Array
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      ∇ₓK = Array{eltype(γ)}(undef, F, N₁, N₂)
      for j in 1:N₂
            for i in 1:N₁
                  for l in 1:F
                        ∇ₓK[l,i,j] = γ[l]*X₂[l,j]
                  end
            end
      end
      return ∇ₓK
end

function mixed_effect_kernel(θ::Array, X₁::Matrix, X₂::Matrix)::Matrix
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      K = Array{eltype(θ)}(undef, N₁, N₂)
      M = ones(F, F)
      D = diagm(ones(F))
      for j in 1:N₂
            for i in 1:N₁
                  sqd = zero(eltype(θ))
                  for k in 1:F
                        sqd += θ[k]*(X₁[k,i] - X₂[k,j])^2
                  end
                  K[i,j] = exp(-0.5*sqd)
            end
      end
      𝐾 = kron(K, θ[F+1]*M+(1-θ[F+1])*D)
      return 𝐾
end

max_stars_repo_name: HamletWantToCode/DiffGPFlux

max_issues_repo_name: HamletWantToCode/DiffGPFlux


text:
using CSV
using DataFrames
using Tables

filename = "example.csv"
file_path  = joinpath(@__DIR__, filename) #ensures we are using the correct directory regardless of how file is run

df = CSV.read(file_path, DataFrame)
display(df)


text:
using DiscreteAdjoint, OrdinaryDiffEq, ForwardDiff
using Test

@testset "DiscreteAdjoint.jl" begin

    function f(du, u, p, t)
        du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
        du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
    end
    p = [1.5,1.0,3.0,1.0]; tspan = (0.0, 10.0); u0 = [1.0,1.0];
    prob = ODEProblem(f, u0, tspan, p)
    sol = solve(prob, Tsit5(), u0=u0, p=p, abstol=1e-10, reltol=1e-10, tstops=0:0.1:10.0)

    dg(out,u,p,t,i) = out .= 1
    t = 0:0.1:10.0

    dp_fd, du0_fd = discrete_adjoint(sol, dg, t; autojacvec=ForwardDiffVJP())
    # dp = [8.30536  -159.484  75.2035  -339.195]
    # du0 = [-39.1275  -8.78778]

    dp_rd, du0_rd = discrete_adjoint(sol, dg, t; autojacvec=ReverseDiffVJP())
    # dp = [8.30536  -159.484  75.2035  -339.195]
    # du0 = [-39.1275  -8.78778]

    dp_rdc, du0_rdc = discrete_adjoint(sol, dg, t; autojacvec=ReverseDiffVJP(true))
    # dp = [8.30536  -159.484  75.2035  -339.195]
    # du0 = [-39.1275  -8.78778]

    dp_z, du0_z = discrete_adjoint(sol, dg, t; autojacvec=ZygoteVJP())
    # dp = [8.30536  -159.484  75.2035  -339.195]
    # du0 = [-39.1275  -8.78778]

    # forward AD solution
    function sum_of_solution(x)
        _prob = remake(prob, u0=x[1:2], p=x[3:end])
        sum(solve(_prob, Tsit5(), abstol=1e-10, reltol=1e-10, saveat=0.1))
    end

    dx = ForwardDiff.gradient(sum_of_solution,[u0;p])
    # dx = [-39.1275  -8.78778  8.30536  -159.484  75.2035  -339.195]

    @test isapprox(du0_fd, dx[1:2])
    @test isapprox(dp_fd, dx[3:end])

    @test isapprox(du0_rd, dx[1:2])
    @test isapprox(dp_rd, dx[3:end])

    @test isapprox(du0_rdc, dx[1:2])
    @test isapprox(dp_rdc, dx[3:end])

    @test isapprox(du0_z, dx[1:2])
    @test isapprox(dp_z, dx[3:end])

end


text:
macro cast(ex)
    esc(codegen_ast_cast(__module__, QuoteNode(__source__), ex))
end

macro main(ex)
    esc(codegen_entry(__module__, QuoteNode(__source__), ex))
end

macro main()
    esc(codegen_entry(__module__, QuoteNode(__source__)))
end

function codegen_ast_cast(m::Module, line, ex)
    if ex isa Symbol
        casted = codegen_ast_cast_module(m, line, ex)
    elseif Meta.isexpr(ex, :module)
        casted = codegen_ast_cast_module(m, line, ex)
    elseif is_function(ex)
        casted = codegen_ast_cast_function(m, line, ex)
    else
        error("unkown expression: $ex, expect module name or function definition")
    end

    return quote
        $(codegen_casted_commands(m))
        $casted
    end
end

function is_argument_with_type(ex)
    ex isa Expr || return false
    return Meta.isexpr(ex, :(::))
end

function is_vararg_with_type(ex)
    ex isa Expr || return false
    Meta.isexpr(ex, :(...)) || return false
    return is_argument_with_type(ex.args[1])
end

function is_optional_argument_with_type(ex)
    ex isa Expr || return false
    Meta.isexpr(ex, :kw) || return false
    return is_argument_with_type(ex.args[1])
end

function split_leaf_command(fn::JLFunction)
    # use ::<type> as hint if there is no docstring
    args = map(fn.args) do each
        if each isa Symbol
            xcall(Comonicon, :JLArgument; name = QuoteNode(each))
        elseif is_argument_with_type(each) # :($name::$type)
            xcall(
                Comonicon,
                :JLArgument;
                name = QuoteNode(each.args[1]),
                type = wrap_type(fn, each.args[2]),
            )
        elseif is_vararg_with_type(each) # :($name::$type...)
            name = each.args[1].args[1]
            type = each.args[1].args[2]
            xcall(
                Comonicon,
                :JLArgument;
                name = QuoteNode(name),
                type = wrap_type(fn, type),
                require = false,
                vararg = true,
            )
        elseif is_optional_argument_with_type(each) # Expr(:kw, :($name::$type), value)
            name = each.args[1].args[1]
            type = each.args[1].args[2]
            value = each.args[2]
            xcall(
                Comonicon,
                :JLArgument;
                name = QuoteNode(name),
                type = wrap_type(fn, type),
                require = false,
                default = string(value),
            )
        elseif Meta.isexpr(each, :...) # :($name...)
            xcall(Comonicon, :JLArgument; name = QuoteNode(each.args[1]), require = false, vararg = true)
        elseif Meta.isexpr(each, :kw) && each.args[1] isa Symbol # Expr(:kw, name::Symbol, value)
            xcall(
                Comonicon,
                :JLArgument;
                name = QuoteNode(each.args[1]),
                require = false,
                default = string(each.args[2]),
            )
        else
            throw(Meta.ParseError("invalid syntax: $each"))
        end
    end

    flags, options = [], []
    if !isnothing(fn.kwargs)
        for each in fn.kwargs
            Meta.isexpr(each, :kw) ||
                error("options should have default values or make it a positional argument")
            expr = each.args[1]
            value = each.args[2]
            if expr isa Symbol # Expr(:kw, name::Symbol, value)
                push!(options, xcall(Comonicon, :JLOption, QuoteNode(expr), Any, string(value)))
            elseif Meta.isexpr(expr, :(::))
                name = expr.args[1]
                type = expr.args[2]

                if type === :Bool || type === Bool
                    value == false || error(
                        "Boolean options must use false as " *
                        "default value, and will be parsed as flags. got $name",
                    )
                    push!(flags, xcall(Comonicon, :JLFlag, QuoteNode(name)))
                else
                    push!(
                        options,
                        xcall(
                            Comonicon,
                            :JLOption,
                            QuoteNode(name),
                            type,
                            string(value) * "::" * string(type),
                        ),
                    )
                end
            end
        end
    end
    args = Expr(:ref, :($Comonicon.JLArgument), args...)
    options = Expr(:ref, :($Comonicon.JLOption), options...)
    flags = Expr(:ref, :($Comonicon.JLFlag), flags...)
    return args, options, flags
end

function wrap_type(def::JLFunction, type)
    def.whereparams === nothing && return type
    return Expr(:where, type, def.whereparams...)
end

function codegen_ast_cast_function(m::Module, @nospecialize(line), ex::Expr)
    fn = JLFunction(ex)
    args, options, flags = split_leaf_command(fn)

    @gensym cmd
    name = default_name(fn.name)
    return quote
        $ex
        Core.@__doc__ $(fn.name)
        $cmd = $Comonicon.cast($(fn.name), $name, $args, $options, $flags, $line)
        $Comonicon.set_cmd!($m.CASTED_COMMANDS, $cmd, $name)
    end
end

function codegen_ast_cast_module(m::Module, line, ex)
    name = name_only(ex)
    @gensym cmd
    cmd_name = default_name(name)
    return quote
        $(Expr(:toplevel, ex))
        Core.@__doc__ $name
        $cmd = $Comonicon.cast($name, $cmd_name, $line)
        $Comonicon.set_cmd!($m.CASTED_COMMANDS, $cmd, $cmd_name)
    end
end

function codegen_create_entry(m::Module, line, @nospecialize(ex))
    @gensym cmd entry
    quote
        $(codegen_entry_cmd(m::Module, line, cmd, ex))
        $entry = $ComoniconTypes.Entry($cmd, $(get_version(m)), $line)
        $Comonicon.set_cmd!($m.CASTED_COMMANDS, $entry, "main")
        $m.eval($ComoniconTargetExpr.emit_expr($entry))
    end
end

function codegen_entry(m::Module, line, @nospecialize(ex = nothing))
    if m === Main
        codegen_script_entry(m, line, ex)
    else
        codegen_project_entry(m, line, ex)
    end
end

function codegen_script_entry(m::Module, line, @nospecialize(ex))
    quote
        $(codegen_create_entry(m, line, ex))
        command_main()
    end
end

function codegen_project_entry(m::Module, line, @nospecialize(ex))
    quote
        $(codegen_create_entry(m, line, ex))

        # entry point for apps
        function julia_main()::Cint
            try
                return command_main()
            catch
                Base.invokelatest(Base.display_error, Base.catch_stack())
                return 1
            end
        end

        """
            comonicon_install(;kwargs...)
        Install the CLI manually. This will use the default configuration in `Comonicon.toml`,
        if it exists. See also [`comonicon_build`](@ref). For more detailed reference, please
        refer to [Comonicon documentation](https://docs.comonicon.org).
        """
        comonicon_install(; kwargs...) = $ComoniconBuilder.install($m; kwargs...)

        """
            comonicon_install_path([quiet=false])
        Install the `PATH` and `FPATH` to your shell configuration file. You can use `comonicon_install_path(true)`,
        to skip interactive prompt.
        For more detailed reference, please refer to [Comonicon documentation](https://docs.comonicon.org).
        """
        comonicon_install_path() = $ComoniconBuilder.install_env_path()

        precompile(Tuple{typeof($m.command_main),Array{String,1}})
    end
end

function codegen_entry_cmd(m::Module, line, cmd, ex)
    if isnothing(ex)
        return codegen_multiple_main_entry(m, line, cmd)
    else
        return codegen_single_main_entry(m, line, cmd, ex)
    end
end

function codegen_multiple_main_entry(m::Module, line, cmd)
    if has_comonicon_toml(m)
        options = ComoniconOptions.read_options(m)
        name = options.name
    else
        name = default_name(nameof(m))
    end

    @gensym doc
    return quote
        Core.@__doc__ const COMMAND_ENTRY_DOC_STUB = nothing
        $doc = @doc(COMMAND_ENTRY_DOC_STUB)
        if $Comonicon.has_docstring($doc)
            $doc = $Comonicon.read_description($doc)
        else
            $doc = nothing
        end

        $cmd = $ComoniconTypes.NodeCommand($name, copy($m.CASTED_COMMANDS), $doc, $line)
    end
end

function codegen_single_main_entry(m::Module, line, cmd, ex)
    fn = JLFunction(ex)
    args, options, flags = split_leaf_command(fn)
    name = default_name(fn.name)
    return quote
        $(codegen_casted_commands(m))
        $ex
        Core.@__doc__ $(fn.name)
        $cmd = $Comonicon.cast($(fn.name), $name, $args, $options, $flags, $line)
    end
end

function codegen_casted_commands(m::Module)
    isdefined(m, :CASTED_COMMANDS) && return
    return :(const CASTED_COMMANDS = Dict{String,Any}())
end

function cast(m::Module, name::String = default_name(m), line = LineNumberNode(0))
    isdefined(m, :CASTED_COMMANDS) || error("module $m does not contain any @cast commands")
    NodeCommand(name, copy(m.CASTED_COMMANDS), split_docstring(m), line)
end

function cast(
    f::Function,
    name::String,
    args::Vector{JLArgument} = JLArgument[],
    options::Vector{JLOption} = JLOption[],
    flags::Vector{JLFlag} = JLFlag[],
    line = LineNumberNode(0),
)
    doc = split_docstring(f)::JLMD
    args, vararg = cast_args(doc, args, line)
    flags = cast_flags(doc, flags, line)
    options = cast_options(doc, options, line)
    return LeafCommand(
        f,
        name,
        args,
        count(x -> x.require, args),
        vararg,
        flags,
        options,
        doc.desc,
        line,
    )
end

function cast_args(doc::JLMD, args::Vector{JLArgument}, line)
    args = map(args) do each
        Argument(
            string(each.name),
            each.type,
            each.vararg,
            each.require,
            each.default,
            Description(get(doc.arguments, string(each.name), nothing)),
            line,
        )
    end

    if last(args).vararg
        return args[1:end-1], last(args)
    else
        return args, nothing
    end
end

function cast_flags(doc::JLMD, flags::Vector{JLFlag}, line)
    cmd_flags = Dict{String,Flag}()
    for each in flags
        name = replace(string(each.name), '_' => '-')
        if haskey(doc.flags, name)
            doc_flag = doc.flags[name]::JLMDFlag
            cmd_flags[name] =
                flg = Flag(;
                    sym = each.name,
                    name = name,
                    short = doc_flag.short,
                    description = Description(doc_flag.desc),
                    line = line,
                )

            if doc_flag.short
                cmd_flags[name[1:1]] = flg
            end
        else
            cmd_flags[name] = Flag(; sym = each.name, name = name, line = line)
        end
    end
    return cmd_flags
end

function cast_options(doc::JLMD, options::Vector{JLOption}, line)
    cmd_options = Dict{String,Option}()
    for each in options
        name = replace(string(each.name), '_' => '-')
        if haskey(doc.options, name)
            option = doc.options[name]::JLMDOption
            cmd_options[name] =
                opt = Option(;
                    sym = each.name,
                    name = name,
                    hint = option.hint, # use user defined hint
                    type = each.type,
                    short = option.short,
                    description = option.desc,
                    line = line,
                )

            if option.short
                cmd_options[name[1:1]] = opt
            end
        else
            cmd_options[name] = Option(;
                sym = each.name,
                name = name,
                hint = each.hint,
                type = each.type,
                line = line,
            )
        end
    end
    return cmd_options
end


text:
# Inner state of a random walk Metropolis algorithm in R^d
struct RWMState{d,FT<:AbstractFloat, T <: AbstractVector{FT},
    FunT <: Function,  RngT <: Random.AbstractRNG}
    rng::RngT # The random number generator
    q!::FunT  # The proposal function
    x::T      # Current state
    u::T      # Draw from q!
    y::T      # Proposal state
end

# Convenience constructors:
"""
    s = RWMState(x0, [rng, [q!]])

Constructor for RWMState (random walk Metropolis state).

# Arguments
- `x0`: The initial state vector
- `rng`: Random number generator; default `Random.GLOBAL_RNG`.
- `q!`: Symmetric proposal distribution; default `randn!`.
  Called by `q!(rng, u)` which puts a draw to vector `u`

If `s` is `RWMState`, then `s.x` is the current state. New proposal
may be drawn to `s.y` by calling `draw!`, and the proposal is
accepted by calling `accept!(s).
"""

function RWMState(x0::AbstractVector{FT}, rng::AbstractRNG=Random.GLOBAL_RNG,
    q::Function=randn!) where {FT <: AbstractFloat}
    state = RWMState{length(x0),FT,typeof(similar(x0)),typeof(q),typeof(rng)}(rng,
    q,  similar(x0), similar(x0), similar(x0))
    copyto!(state.x, x0)
    state
end

# Draw u from proposal
@inline function draw_u!(s::RWMState) :: Nothing
    s.q!(s.rng, s.u)
    nothing
end

"""
   draw!(s::RWMState, sc::AbstractFloat)

Draw proposal from `s.x` to `s.y` using scaling `sc`.
"""
@inline function draw!(s::RWMState, sc::AbstractFloat) :: Nothing
    draw_u!(s)
    s.y .= sc.*s.u .+ s.x
    nothing
end

# Draw proposal using scalar multiplier
@inline function draw!(s::RWMState{1}, sc::AbstractFloat) :: Nothing
    draw_u!(s)
    s.y[] = sc*s.u[] + s.x[]
    nothing
end

"""
   draw!(s::RWMState, L::Cholesky, sc::AbstractFloat)

Draw proposal from `s.x` to `s.y` using scaling `sc*L.L`, with default `sc=1.0`.
"""
@inline function draw!(s::RWMState{d,FT}, L::Cholesky{FT}, sc::FT=one(FT)) where {
    d, FT <: AbstractFloat}
    draw_u!(s)
    if L.uplo == 'L'
        lowerTriInplaceMultiplyAdd!(s.y, sc, L.factors, s.u, one(FT), s.x)
    else
        mul!(s.y, L.L, s.u)
        s.y .= sc .* s.y .+ s.x
    end
    nothing
end

"""
   accept!(s::RWMState)

Accept the proposal, that is, set copy `s.y` to `s.x`.
"""
@inline function accept!(sr::RWMState) :: Nothing
    sr.x .= sr.y
    nothing
end

@inline function accept!(sr::RWMState{1}) :: Nothing
    sr.x[] = sr.y[]
    nothing
end


text:
using CSV, DataFrames, StatsPlots
using StanSample, Statistics

ProjDir = @__DIR__

rcd = CSV.read(joinpath(ProjDir, "RedCardData.csv"), DataFrame;
  missingstring = "NA")
rcd = dropmissing(rcd, :rater1)
rcd[!, :rater1] = Float64.(rcd.rater1)

stan_logistic_0 ="
data {
  int<lower=0> N;
  int<lower=0> n_redcards[N];
  int<lower=0> n_games[N];
  vector[N] rating;
}
parameters {
  vector[2] beta;
}
model {
  beta[1] ~ normal(0, 10);
  beta[2] ~ normal(0, 1);
  n_redcards ~ binomial_logit(n_games, beta[1] + beta[2] * rating);
}
";

stan_logistic_1 = "
functions {
  real partial_sum_lpmf(int[] slice_n_redcards,
                        int start, int end,
                        int[] n_games,
                        vector rating,
                        vector beta) {
    return binomial_logit_lupmf(slice_n_redcards |
                               n_games[start:end],
                               beta[1] + beta[2] * rating[start:end]);
  }
}
data {
  int<lower=0> N;
  int<lower=0> n_redcards[N];
  int<lower=0> n_games[N];
  vector[N] rating;
  int<lower=1> grainsize;
}
parameters {
  vector[2] beta;
}
model {

  beta[1] ~ normal(0, 10);
  beta[2] ~ normal(0, 1);
  target += reduce_sum(partial_sum_lupmf, n_redcards, grainsize,
                       n_games, rating, beta);
}
";

isdir(joinpath(ProjDir, "tmp")) && rm(joinpath(ProjDir, "tmp"), recursive=true);
if !isdir("tmp")
  tmp = mkdir(joinpath(ProjDir, "tmp"))
else
  tmp = joinpath(ProjDir, "tmp")
end;

logistic_0 = SampleModel("logistic0", stan_logistic_0, tmp)
logistic_1 = SampleModel("logistic1", stan_logistic_1, tmp)

data = Dict(
    :N => size(rcd, 1),
    :grainsize => 1,
    :n_redcards => rcd.redCards,
    :n_games => rcd.games,
    :rating => rcd.rater1
)

println("\nUsing $(Threads.nthreads()) Julia threads.\n")

println("Timing of logistic_0 (4 Julia chains):")
@time rc_0 = stan_sample(logistic_0; data);

if success(rc_0)
    dfs_0 = read_summary(logistic_0)
    dfs_0[8:9, [1,2,4,8,9,10]] |> display
    println()
end

println("Timing of logistic_1 (4 Julia chains):")
@time rc_1 = stan_sample(logistic_1; data);

if success(rc_1)
    dfs_1 = read_summary(logistic_1)
    dfs_1[8:9, [1,2,4,8,9,10]] |> display
    println()
end

println("Timing of logistic_1 (4 C++ chains, 4 num_threads):")
@time rc_2 = stan_sample(logistic_1; data,
  use_cpp_chains=true, num_threads=4);

if success(rc_2)
    dfs_2 = read_summary(logistic_1)
    dfs_2[8:9, [1,2,4,8,9,10]] |> display
    println()
end

println("Timing of logistic_1 (4 C++ chains, 8 num_threads):")
@time rc_3 = stan_sample(logistic_1; data,
  use_cpp_chains=true, num_threads=8);

if success(rc_3)
    dfs_3 = read_summary(logistic_1)
    dfs_3[8:9, [1,2,4,8,9,10]] |> display
    println()
end

println("Timing of logistic_1 (2 C++ chains, 2 Julia chains, 8 num_threads):")
@time rc_4 = stan_sample(logistic_1; data,
  use_cpp_chains=true, check_num_chains=false,
  num_threads=8, num_cpp_chains=2, num_julia_chains=2);

if success(rc_4)
    dfs_4 = read_samples(logistic_1, :dataframe)
    size(dfs_4) |> display
    mean(Array(dfs_4), dims=1) |> display
    read_summary(logistic_1, true) |> display
    println()
end


text:
############################################################################################
"""
$(TYPEDEF)
Default Configuration for MALA sampler.

# Fields
$(TYPEDFIELDS)
"""
struct ConfigMALA <: AbstractConfiguration
    "Target Acceptance Rate"
    δ::Float64
    "Default size for tuning iterations in each cycle."
    window::ConfigTuningWindow
    function ConfigMALA(
        δ::Float64,
        window::ConfigTuningWindow
    )
        @argcheck 0.0 < δ <= 1.0 "Acceptance rate not bounded between 0 and 1"
        return new(
            δ,
            window
        )
    end
end

"""
$(SIGNATURES)
Initialize Mala custom configurations.

# Examples
```julia
```

"""
function init(
    ::Type{AbstractConfiguration},
    mcmc::Type{MALA},
    objective::Objective,
    proposalconfig::ConfigProposal;
    ## Target Acceptance Rate
    δ=0.574,
    ## MCMC Phase tune variables based on standard HMC best practice target acceptance rate of 0.80
    window= ConfigTuningWindow(
        [1, 5, 1],
        Int.(floor.([75, 50, 25] * (0.80 / 0.574))),
        [Warmup(), Adaptionˢˡᵒʷ(), Adaptionᶠᵃˢᵗ(), Exploration()],
        )
)
    return ConfigMALA(
        δ,
        window
    )
end

############################################################################################
#export
export ConfigMALA


text:
###
# Hcat
####
hcat(A::AbstractQuasiArray...) = apply(hcat, A...)

ApplyStyle(::typeof(hcat), ::Type{<:AbstractQuasiArray}...) = LazyQuasiArrayApplyStyle()

axes(f::ApplyQuasiMatrix{<:Any,typeof(hcat)}) = (axes(f.args[1],1), oneto(sum(size.(f.args,2))))

_vecormat_getindex(A::AbstractQuasiMatrix,k,ξ) = A[k,ξ]
_vecormat_getindex(A::AbstractQuasiVector,k,ξ) = A[k]
@inline LazyArrays._view_hcat(a::AbstractQuasiVector, kr, jr::Colon) = view(a, kr)
# TODO: generalize
@inline LazyArrays._view_hcat(a::AbstractQuasiVector, kr::Number, jr) = a[kr]

function getindex(f::ApplyQuasiMatrix{T,typeof(hcat)}, k::Number, j::Number) where T
    ξ = j
    for A in f.args
        n = size(A,2)
        ξ ≤ n && return T(_vecormat_getindex(A,k,ξ))::T
        ξ -= n
    end
    throw(BoundsError(f, (k,j)))
end

\(A::AbstractQuasiArray, H::ApplyQuasiMatrix{<:Any,typeof(hcat)}) = hcat((Ref(A) .\ H.args)...)

function ==(A::ApplyQuasiMatrix{<:Any,typeof(hcat)}, B::ApplyQuasiMatrix{<:Any,typeof(hcat)})
    axes.(A.args,2) == axes.(B.args,2) && return all(A.args .== B.args)
    Base.invoke(==, Tuple{AbstractQuasiArray,AbstractQuasiArray}, A, B)
end

"""
    UnionVcat

is an analogue of `Vcat` that takes the union of the axes.
If there is overlap, it uses the first in order.
"""
struct UnionVcat{T,N,Args<:Tuple} <: AbstractQuasiArray{T,N}
    args::Args
end

function UnionVcat{T,2}(A...) where T
    ncols = axes(A[1], 2)
    for j = 2:length(A)
        if axes(A[j], 2) != ncols
            throw(ArgumentError("axes of columns of each array must match (got $(map(x->axes(x,2), A)))"))
        end
    end
    UnionVcat{T,2,typeof(A)}(A)
end

UnionVcat{T,1}(A...) where T = UnionVcat{T,1,typeof(A)}(A)

UnionVcat{T}(a::AbstractQuasiOrVector...) where T = UnionVcat{T,1}(a...)
UnionVcat{T}(a::AbstractQuasiOrMatrix...) where T = UnionVcat{T,2}(a...)
UnionVcat(a...) = UnionVcat{mapreduce(eltype,promote_type,a)}(a...)

axes(A::UnionVcat{<:Any,1}) = (union(axes.(A.args,1)...),)
axes(A::UnionVcat{<:Any,2}) = (union(axes.(A.args,1)...),axes(A.args[1],2))

function _getindex(::Type{IND}, A::UnionVcat{T,1}, (x,)::IND) where {IND,T}
    for a in A.args
        x in axes(a,1) && return convert(T,a[x])::T
    end
    throw(BoundsError(A, I))
end

function _getindex(::Type{IND}, A::UnionVcat{T,2}, (x,j)::IND) where {IND,T}
    for a in A.args
        x in axes(a,1) && return convert(T,a[x,j])::T
    end
    throw(BoundsError(A, I))
end


text:
import Base: eltype

"""
All networks in the package belong to the `AbstractEcologicalNetwork` type. They
all have a field `A` to represent interactions as a *matrix*, and a number of
fields for species. See the documentation for `AbstractBipartiteNetwork` and
`AbstractUnipartiteNetwork`, as well as `AllowedSpeciesTypes` for the allowed
types for species.

Note that *all* species in a network (including both levels of a bipartite
network) *must* have the same type. For example, `["a", :b, "c"]` is not a valid
array of species, as not all its elements have the same type.
"""
abstract type AbstractEcologicalNetwork end

"""
This abstract type groups all unipartite networks, regardless of the type of
information. Unipartite networks have *a single* field for species, named `S`,
which has the same number of elements as the size of the matrix.

Any unipartite network can be declared (we'll use the example of a binary
network) either using `UnipartiteNetwork(A, S)` (assuming `A` is a matrix of
interactions and `S` is a vector of species names), or `UnipartiteNetwork(A)`,
in which case the species will be named automatically.
"""
abstract type AbstractUnipartiteNetwork <: AbstractEcologicalNetwork end

"""
This abstract type groups all bipartite networks, regardless of the type of
information. Bipartite networks have *two* fields for species, named `T` (for
top, corresponding to matrix *rows*), and `B` (for bottom, matrix *columns*).

Any bipartite network can be declared (we'll use the example of a binary
network) either using `BipartiteNetwork(A, T, B)` (assuming `A` is a matrix of
interactions and `T` and `B` are vectors of species names for the top and bottom
level), or `BipartiteNetwork(A)`, in which case the species will be named
automatically.
"""
abstract type AbstractBipartiteNetwork <: AbstractEcologicalNetwork end

"""
A bipartite deterministic network is a matrix of boolean values.
"""
mutable struct BipartiteNetwork{Bool, T} <: AbstractBipartiteNetwork
  A::Matrix{Bool}
  T::Vector{T}
  B::Vector{T}
  function BipartiteNetwork{Bool, NT}(A::M, T::Vector{NT}, B::Vector{NT}) where {M<:AbstractMatrix{Bool}, NT}
    check_bipartiteness(A, T, B)
    new{Bool,NT}(A, T, B)
  end
end

"""
An unipartite deterministic network is a matrix of boolean values.
"""
mutable struct UnipartiteNetwork{Bool, T} <: AbstractUnipartiteNetwork
  A::Matrix{Bool}
  S::Vector{T}
  function UnipartiteNetwork{Bool, NT}(A::M, S::Vector{NT}) where {M<:AbstractMatrix{Bool}, NT}
    check_unipartiteness(A, S)
    new{Bool,NT}(A, S)
  end
end

"""
A bipartite probabilistic network is a matrix of floating point numbers, all of
which must be between 0 and 1.
"""
mutable struct BipartiteProbabilisticNetwork{IT<:AbstractFloat, NT} <: AbstractBipartiteNetwork
  A::Matrix{IT}
  T::Vector{NT}
  B::Vector{NT}
  function BipartiteProbabilisticNetwork{IT, NT}(A::Matrix{IT}, T::Vector{NT}, B::Vector{NT}) where {IT<:AbstractFloat, NT}
    check_bipartiteness(A, T, B)
    check_probability_values(A)
    new{IT,NT}(A, T, B)
  end
end

"""
A bipartite quantitative network is matrix of numbers. It is assumed that the
interaction strength are *positive*.
"""
mutable struct BipartiteQuantitativeNetwork{IT<:Number, NT} <: AbstractBipartiteNetwork
  A::Matrix{IT}
  T::Vector{NT}
  B::Vector{NT}
  function BipartiteQuantitativeNetwork{IT, NT}(A::Matrix{IT}, T::Vector{NT}, B::Vector{NT}) where {IT<:Number, NT}
    check_bipartiteness(A, T, B)
    new{IT,NT}(A, T, B)
  end
end

"""
A unipartite probabilistic network is a square matrix of floating point numbers,
all of which must be between 0 and 1.
"""
mutable struct UnipartiteProbabilisticNetwork{IT<:AbstractFloat, NT} <: AbstractUnipartiteNetwork
  A::Matrix{IT}
  S::Vector{NT}
  function UnipartiteProbabilisticNetwork{IT, NT}(A::Matrix{IT}, S::Vector{NT}) where {IT<:AbstractFloat,NT}
    check_unipartiteness(A, S)
    check_probability_values(A)
    new{IT,NT}(A, S)
  end
end

"""
A unipartite quantitative network is a square matrix of numbers.
"""
mutable struct UnipartiteQuantitativeNetwork{IT<:Number, NT} <: AbstractUnipartiteNetwork
  A::Matrix{IT}
  S::Vector{NT}
  function UnipartiteQuantitativeNetwork{IT, NT}(A::Matrix{IT}, S::Vector{NT}) where {IT<:Number,NT}
    check_unipartiteness(A, S)
    new{IT,NT}(A, S)
  end
end

"""
This is a union type for both Bipartite and Unipartite probabilistic networks.
Probabilistic networks are represented as arrays of floating point values ∈
[0;1].
"""
ProbabilisticNetwork = Union{BipartiteProbabilisticNetwork, UnipartiteProbabilisticNetwork}

"""
This is a union type for both Bipartite and Unipartite deterministic networks.
All networks from these class have adjacency matrices represented as arrays of
Boolean values.
"""
BinaryNetwork = Union{BipartiteNetwork, UnipartiteNetwork}

"""
This is a union type for both unipartite and bipartite quantitative networks.
All networks of this type have adjancency matrices as two-dimensional arrays of
numbers.
"""
QuantitativeNetwork = Union{BipartiteQuantitativeNetwork, UnipartiteQuantitativeNetwork}

"""
All non-probabilistic networks
"""
DeterministicNetwork = Union{BinaryNetwork, QuantitativeNetwork}

all_types = [
    :BipartiteNetwork, :UnipartiteNetwork,
    :BipartiteProbabilisticNetwork, :UnipartiteProbabilisticNetwork,
    :BipartiteQuantitativeNetwork, :UnipartiteQuantitativeNetwork
    ]
for T in all_types
  @eval begin
    """
        eltype(N::$($T){IT,ST}) where {IT,ST}

    Returns a tuple with the type of the interactions, and the type of the species.
    """
    function eltype(N::$T{IT,ST}) where {IT,ST}
      return (IT,ST)
    end
  end
end


text:
module _vscodeserver

function remlineinfo!(x)
    if isa(x, Expr)
        if x.head == :macrocall && x.args[2] != nothing
            id = findall(map(x -> (isa(x, Expr) && x.head == :line) || (isdefined(:LineNumberNode) && x isa LineNumberNode), x.args))
            deleteat!(x.args, id)
            for j in x.args
                remlineinfo!(j)
            end
            insert!(x.args, 2, nothing)
        else
            id = findall(map(x -> (isa(x, Expr) && x.head == :line) || (isdefined(:LineNumberNode) && x isa LineNumberNode), x.args))
            deleteat!(x.args, id)
            for j in x.args
                remlineinfo!(j)
            end
        end
    end
    x
end

using REPL, Sockets, Base64, Pkg
import Base: display, redisplay
global active_module = :Main

struct InlineDisplay <: AbstractDisplay end

pid = Base.ARGS[1]

function change_module(newmodule::String, print_change = true)
    global active_module
    smods = Symbol.(split(newmodule, "."))

    val = Main
    for i = 1:length(smods)
        if isdefined(val, smods[i])
            val = getfield(val, smods[i])
        else
            println("Could not find module $newmodule")
            return
        end
    end
    expr = Meta.parse(newmodule)
    active_module = expr
    repl = Base.active_repl
    main_mode = repl.interface.modes[1]
    main_mode.prompt = string(newmodule,"> ")
    main_mode.on_done = REPL.respond(repl,main_mode; pass_empty = false) do line
        if !isempty(line)
            ex = Meta.parse(line)
            if ex isa Expr && ex.head == :module
                ret = :( Base.eval($expr, Expr(:(=), :ans, Expr(:toplevel, Meta.parse($line)))) )
            else
                ret = :( Core.eval($expr, Expr(:(=), :ans, Meta.parse($line))) )
            end
        else
            ret = :(  )
        end
        out = connect(to_vscode)
        write(out, string("repl/variables,", getVariables(), "\n"))
        close(out)
        return ret
    end
    print(" \r ")
    print_change && println("Changed root module to $expr")
    printstyled(string(newmodule,"> "), bold = true, color = :green)
end

function get_available_modules(m=Main, out = Module[])
    for n in names(m, all = true, imported = true)
        if isdefined(m, n) && getfield(m, n) isa Module  && !(getfield(m, n) in out)
            M = getfield(m, n)
            push!(out, M)
            get_available_modules(M, out)
        end
    end
    out
end

function getVariables()
    M = @__MODULE__
    variables = []
    msg = ""
    for n in names(M)
        !isdefined(M, n) && continue
        x = getfield(M, n)
        t = typeof(x)
        msg = string(msg, ";", n, "::", t)
    end
    return msg
end

function generate_pipe_name(name)
    if Sys.iswindows()
        "\\\\.\\pipe\\vscode-language-julia-$name-$pid"
    elseif Sys.isunix()
        joinpath(tempdir(), "vscode-language-julia-$name-$pid")
    end
end

!(Sys.isunix() || Sys.iswindows()) && error("Unknown operating system.")

global_lock_socket_name = generate_pipe_name("terminal")
from_vscode = generate_pipe_name("torepl")
to_vscode = generate_pipe_name("fromrepl")

if issocket(from_vscode)
    rm(from_vscode)
end

@async begin
    server = listen(from_vscode)
    while true
        sock = accept(server)
        cmd = readline(sock)
        !startswith(cmd, "repl/") && continue
        text = readuntil(sock, "repl/endMessage")[1:end-15]
        if cmd == "repl/include"
            cmod = Core.eval(active_module)
            ex = Expr(:call, :include, strip(text, '\n'))
            cmod.eval(ex)
        elseif cmd == "repl/getVariables"
            out = connect(to_vscode)
            write(out, getVariables())
            close(out)
        end
    end
end

conn = connect(global_lock_socket_name)

function display(d::InlineDisplay, ::MIME{Symbol("image/png")}, x)
    payload = stringmime(MIME("image/png"), x)
    print(conn, "image/png", ":", sizeof(payload), ";")
    print(conn, payload)
end

displayable(d::InlineDisplay, ::MIME{Symbol("image/png")}) = true

function display(d::InlineDisplay, ::MIME{Symbol("image/svg+xml")}, x)
    payload = stringmime(MIME("image/svg+xml"), x)
    print(conn, "image/svg+xml", ":", sizeof(payload), ";")
    print(conn, payload)
end

displayable(d::InlineDisplay, ::MIME{Symbol("image/svg+xml")}) = true

function display(d::InlineDisplay, ::MIME{Symbol("text/html")}, x)
    payload = stringmime(MIME("text/html"), x)
    print(conn, "text/html", ":", sizeof(payload), ";")
    print(conn, payload)
end

displayable(d::InlineDisplay, ::MIME{Symbol("text/html")}) = true

function display(d::InlineDisplay, ::MIME{Symbol("juliavscode/html")}, x)
    payload = stringmime(MIME("juliavscode/html"), x)
    print(conn, "juliavscode/html", ":", sizeof(payload), ";")
    print(conn, payload)
end

Base.Multimedia.istextmime(::MIME{Symbol("juliavscode/html")}) = true

displayable(d::InlineDisplay, ::MIME{Symbol("juliavscode/html")}) = true

function display(d::InlineDisplay, ::MIME{Symbol("application/vnd.vegalite.v2+json")}, x)
    payload = stringmime(MIME("application/vnd.vegalite.v2+json"), x)
    print(conn, "application/vnd.vegalite.v2+json", ":", sizeof(payload), ";")
    print(conn, payload)
end

displayable(d::InlineDisplay, ::MIME{Symbol("application/vnd.vegalite.v2+json")}) = true

function display(d::InlineDisplay, ::MIME{Symbol("application/vnd.plotly.v1+json")}, x)
    payload = stringmime(MIME("application/vnd.plotly.v1+json"), x)
    print(conn, "application/vnd.plotly.v1+json", ":", sizeof(payload), ";")
    print(conn, payload)
end

displayable(d::InlineDisplay, ::MIME{Symbol("application/vnd.plotly.v1+json")}) = true

function display(d::InlineDisplay, x)
    if showable("application/vnd.vegalite.v2+json", x)
        display(d,"application/vnd.vegalite.v2+json", x)
    elseif showable("application/vnd.plotly.v1+json", x)
        display(d,"application/vnd.plotly.v1+json", x)
    elseif showable("juliavscode/html", x)
        display(d,"juliavscode/html", x)
    # elseif showable("text/html", x)
    #     display(d,"text/html", x)
    elseif showable("image/svg+xml", x)
        display(d,"image/svg+xml", x)
    elseif showable("image/png", x)
        display(d,"image/png", x)
    else
        throw(MethodError(display,(d,x)))
    end

end
if length(Base.ARGS) >= 3 && Base.ARGS[3] == "true"
    atreplinit(i->Base.Multimedia.pushdisplay(InlineDisplay()))
end

# Load revise?
load_revise = Base.ARGS[2] == "true" && (VERSION < v"1.1" ? haskey(Pkg.Types.Context().env.manifest, "Revise") : haskey(Pkg.Types.Context().env.project.deps, "Revise"))

end

if _vscodeserver.load_revise
    @eval using Revise
    Revise.async_steal_repl_backend()
end


text:
using LAP_julia: normalize_to_zero_one, create_sparse_flow_from_sparse, create_sparse_flow_from_full, vec_len
using TimerOutputs

function gen_flow(img_size)

    a, b, c = randn(), randn(), randn()
    f(z) = a + b*z + c*z^2 # + d*z^3 + e*z^4

    X = ones(img_size[1]) * collect(range(0,1,length=img_size[2]))'
    Y = collect(range(0,1,length=img_size[1])) * ones(img_size[2])'
    A = Y .* im + X

    B = f.(A)
    max_len = maximum(LAP_julia.vec_len.(B))
    B = B./max_len
    return B
end

B = gen_flow((12,12)); showflow(B)

println(B)

f(1 * im)

imgshow(rand(12,123))

using TimerOutputs, LAP_julia

img, imgw, flow = gen_init(:chess, :quad, flow_args=[20])
img, imgw, flow = gen_init(:lena, :quad, flow_args=[20])

TimerOutputs.enable_debug_timings(LAP_julia)
timer = TimerOutput("Registration");
@timeit timer "polyfilter lap" begin
    flow_est, source_reg = pflap(img, imgw, display=false, timer=timer)
end
print_timer(timer)

showflow(flow_est)
showflow(flow-flow_est)

b = zeros(123,123)
@benchmark a = zeros(123,123)
@benchmark a = similar(b)
@benchmark a = Array{Float64}(undef, 123, 123)
@benchmark a = Array{Float32}(undef, 123, 123)

similar(b, (234, 234))

## fix vis methods

img = gen_chess()
img = gen_lena()
flow = ones(size(img)) .* 20 .+ ones(size(img)) .* 10 .* im
imgw = warp_img(img, -real(flow), -imag(flow))
showflow(flow)
imgshowflow(imgw, flow)
imgshow(img)

##
img, imgw, flow = gen_init()

showflow(flow)
imgshow(imgw)

## speedtest jl

img, imgw, flow = gen_init(:lena, :uniform, flow_args=[-1 - 4im, 15]);
## SPARSE LAP
timer=TimerOutput("reg alg: sparse lap");
method_kwargs = Dict(:timer => timer, :point_count => 200, :spacing => 15)
# method_kwargs = Dict(:timer => timer, :point_count => 35, :spacing => 35, :flow_interpolation_method => :rbf)
flow_est, source_reg, timer, results, (estim_at_inds, inds) = test_registration_alg(sparse_lap, img, imgw, flow, [25, (51,51)], method_kwargs, timer=timer)

norm_err = normalize_to_zero_one(abs.(err))
sparse_estim_flow = create_sparse_flow_from_sparse(estim_at_inds, inds, size(flow_est))
addpoints(inds, labels=map(x -> string(round(x, digits=2)), norm_err))
sparse_truth_flow = create_sparse_flow_from_full(flow, inds)
showflow(sparse_truth_flow.-sparse_estim_flow)

ncc(err, vec_len.((map(ind -> flow[ind], inds) .- estim_at_inds)))
addpoints(inds, labels=map(x -> string(round(x, digits=2)), vec_len.((map(ind -> flow[ind], inds) .- estim_at_inds))))

to = timer["reg alg: sparse lap"]["sparse lap"]["filtering"]

TimerOutputs.tottime(to)/10e8

showflow(flow_est-flow)

showflow(flow_est)
showflow(flow)
imgoverlay(img, source_reg)
imgoverlay(img, imgw)

## LAP
img, imgw, flow = gen_init(:lena, :tiled, flow_args=[15, 30])
timer=TimerOutput("reg alg: lap");
flow_est, source_reg, timer, results = test_registration_alg(lap, img, imgw, flow, [25, (51,51)], Dict(:timer => timer), timer=timer);

img, imgw = gen_anhir()
flow_est, source_reg = sparse_pflap(img, imgw)
showflow(flow_est)
imgshow(source_reg)
imgshow(img)
imgshow(imgw)

sections = ["reg alg: lap", "lap", "calculate flow"];
TimerOutputs.time(get_timer(old_timer, sections))/TimerOutputs.time(get_timer(new_timer, sections))

showflow(flow_est)
showflow(flow)

## PF LAP
timer=TimerOutput("pf lap");
method_kwargs =Dict(:timer => timer, :display => false, :max_repeats => 1)
flow_est, source_reg, timer, results = test_registration_alg(pflap, img, imgw, flow, method_kwargs=method_kwargs, timer=timer);

@code_warntype pflap(img, imgw, display = false)

showflow(flow_est)
showflow(flow)
showflow(flow - flow_est, figtitle="flow - old flow")
## SPARSE PF LAP
img, imgw, flow = gen_init();
img, imgw, flow = gen_init(:spaghetti, :uniform, flow_args=[-1 - 4im, 15]);

timer=TimerOutput("sparse pf lap");
method_kwargs = Dict(:timer => timer, :display => false, :max_repeats => 1, :point_count => 500, :spacing => 10)
flow_est, source_reg, timer, results = test_registration_alg(sparse_pflap, img, imgw, flow, method_kwargs=method_kwargs, timer=timer)

showflow(flow_est)
showflow(flow)
showflow(flow-flow_est)

imgoverlay(img, source_reg)
imgshow(img - source_reg)
imgshow(img)
imgshow(imgw)
imgshow(source_reg)

timer["sparse pf lap"]["single filter pyramid level"]

window_sizes = cat(collect(1:51), collect(61:10:101), collect(141:40:381), dims=1)
limit = 70
map(x -> x <= limit ? x : 0, window_sizes)
img_size = 50
img = gen_init(:chess, chess_args=[25, img_size/25])

##
filter_size = 2*filter_half_size + 1
filter_half_size = 2
k = (-filter_half_size:filter_half_size)

basis = collect(reshape(1:length(k)^2, filter_size, filter_size))

basis * k

basis' * k

##
filter_half_size = 15
filter_size = 2*filter_half_size + 1
k = (-filter_half_size:filter_half_size)

# Get the displacement vector field from the filters

basis = similar(img, (filter_size, filter_size, filter_num))
forward_ker, backward_ker = prepare_gaussian_filters(filter_half_size)
for k in 1:3
    basis[:, :, k] .= parent(broadcast(*, forward_ker[k]...))
end

basis

image_size = size(img)
all_coeffs = rand((image_size[1]^2), 3)

u1_top = zeros(image_size);
u1_bot = zeros(image_size);
u2_top = zeros(image_size);
u2_bot = zeros(image_size);

for n in 1:filter_num
    @views u1_top[:] .= u1_top[:] .- sum(transpose(basis[:, :, n]) * k) .* all_coeffs[:, n];
    @views u1_bot[:] .= u1_bot[:] .+ sum(basis[:, :, n]) .* all_coeffs[:, n];

    @views u2_top[:] .= u2_top[:] .- sum(basis[:, :, n] * k) .* all_coeffs[:, n];
    @views u2_bot[:] .= u2_bot[:] .+ sum(basis[:, :, n]) .* all_coeffs[:, n];
end

u_est = 2 .* ((im .* u1_top ./ u1_bot) .+ (u2_top ./ u2_bot));

sum(transpose(basis[:, :, 2]) * k)
sum(basis[:, :, 1] * k)
sum(basis[:, :, 2])

## use the fact that some of the sums are 0

u1_top = zeros(image_size);
u1_bot = zeros(image_size);
u2_top = zeros(image_size);
u2_bot = zeros(image_size);

for n in 1:filter_num
    if n == 1
        @views u1_top[:] .= u1_top[:] .- 0 .* all_coeffs[:, n];
        @views u1_bot[:] .= u1_bot[:] .+ 1 .* all_coeffs[:, n];

        @views u2_top[:] .= u2_top[:] .- 0 .* all_coeffs[:, n];
        @views u2_bot[:] .= u2_bot[:] .+ 1 .* all_coeffs[:, n];
    elseif n == 2
        @views u1_top[:] .= u1_top[:] .- 0 .* all_coeffs[:, n];
        @views u1_bot[:] .= u1_bot[:] .+ 0 .* all_coeffs[:, n];

        @views u2_top[:] .= u2_top[:] .- 1 .* all_coeffs[:, n];
        @views u2_bot[:] .= u2_bot[:] .+ 0 .* all_coeffs[:, n];
    elseif n == 3
        @views u1_top[:] .= u1_top[:] .- 1 .* all_coeffs[:, n];
        @views u1_bot[:] .= u1_bot[:] .+ 0 .* all_coeffs[:, n];

        @views u2_top[:] .= u2_top[:] .- 0 .* all_coeffs[:, n];
        @views u2_bot[:] .= u2_bot[:] .+ 0 .* all_coeffs[:, n];
    end
end

u_est_2 = 2 .* ((im .* u1_top ./ u1_bot) .+ (u2_top ./ u2_bot));

n = 3
sum(transpose(basis[:, :, n]) * k)
sum(basis[:, :, n])
sum(basis[:, :, n] * k)
sum(basis[:, :, n])

isapprox(u_est, u_est_2)

showflow(u_est - u_est_2)

## coeffs
B = ones(123, 3)

B[[1, 3, 12,12, 14], :] .= 12

dif = ones(123,123,3)

fhs = 2
shift = CartesianIndex(fhs, fhs)

ind = CartesianIndex(5,5)
dif[ind-shift:ind+shift, 2] .= 2

dif[ind, :]

@benchmark Dict()

function foo(img1, dict=Dict(:sone => img1))
    println(typeof(img1))
    println(typeof(dict))
end

foo(img)
img1 = img

imgfilt = zeros(size(img))
σ = 25
gaus = gaussian(σ, 2 * σ + 1)

@benchmark out = imfilter(img1, kernelfactors((gaus, gaus)), "symmetric")

@benchmark imfilter!(img1, img1, kernelfactors((gaus, gaus)), "symmetric")

out == img1
out == imgfilt

using LAP_julia: gen_uniform_flow

uni_flow = gen_uniform_flow((200, 200), 1+1im, 15)

## fair comparison matlab vs this

img, imgw, flow = gen_init(:lena, :uniform, flow_args=[-1 - 4im, 15]);

timer=TimerOutput("reg alg: lap");
flow_est, source_reg, timer, results = test_registration_alg(lap, img, imgw, flow, [25, (51,51)], Dict(:timer => timer), timer=timer);

div_coef = zeros(12,12)
div_coef[3:4,3:7] .= 1
mask = falses(12,12)
mask[2:6, 3:4] .= true
@benchmark newly_inpainted = ((div_coef .!= 0) .& (mask .== 0))
@benchmark newly_inpainted = ((div_coef .!= 0) .* (mask .== 0))

mask
div_coef

source = imgw
target = img

source_hist = adjust_histogram(source, Matching(targetimg = target))

@benchmark out = adjust_histogram($source, Matching(targetimg = $target))
@benchmark adjust_histogram!($source, Matching(targetimg = $target)) # takes 6ms on average

@benchmark (target, source) = LAP_julia.pad_images($target, $source)

@benchmark level_count = floor(Int64, log2(minimum(size(target))/8)+1)+1

@benchmark mask =

using LAP_julia: max_displacement

function foo()
    for k in 1:120
        println("iter: ", k)
        img, imgw = gen_anhir()
        try
            flow_est, source_reg = sparse_pflap(img, imgw)
        catch e
            println(e)
            return img, imgw
        end
        if max_displacement(flow_est) > 200
            return img, imgw
        end
    end
    println("all fine")
end

img, imgw = foo()

img, imgw = gen_anhir()

flow = gen_tiled_flow(size(img))
imgw = warp_img(img, real(flow), imag(flow))

imgshow(img, origin_bot=true, figtitle="target")
imgshow(imgw, origin_bot=true, figtitle="source")

max_repeats = 3;
flow_est, source_reg, figs = sparse_pflap(img, imgw, display=true, max_repeats=max_repeats)

using LAP_julia
assess_flow_quality(flow_est, flow)

@manipulate for k in slider(0:max_repeats*size(figs,1), value=1, label="k")
    if k == 0
        imgshow(imgw, origin_bot=true, figtitle="source")
    else
        reshape(permuteddimsview(figs, [2, 1, 3]), (:, 4))[k, 2]
    end
end

reshape(figs, (:, 5))

figs[1,1,2]

imgshow(source_reg, origin_bot=true, figtitle="source_reg")
imgshow(img, origin_bot=true, figtitle="target")

for k in 1:120
    try
        flow_est, source_reg, figs = sparse_pflap(img, imgw, display=true)
    catch e
        println(e)
        break
    end
end

img = gen_chess(10,10)
inds = find_edge_points(img, point_count=150, spacing=8)

imgshow(img); addpoints(inds)

##

a = zeros(Float16, 123,123);
c = ones(Float16, 5,5);
@benchmark imfilter(a, c)
b = zeros(Float64, 123,123);
d = ones(Float64, 5,5);
@benchmark imfilter(b, d)

##

img, imgw, flow = gen_init()

showflow(flow)

small_img = imresize(img, ratio = 0.5)
small_imgw = imresize(imgw, ratio = 0.5)

## small version
@btime small_flow_est, small_source_reg, rest = sparse_lap(small_img, small_imgw, 15, point_count = 500, spacing = 5)
imgshow(small_source_reg)
showflow(small_flow_est)

@btime enlarged_small_flow = imresize(small_flow_est, ratio = 2) .* 2

showflow(enlarged_small_flow.-flow)

## normal version
@btime flow_est, source_reg, rest = sparse_lap(img, imgw, 30)
imgshow(source_reg)
showflow(flow_est)

showflow(flow_est.-flow)

## classic lap
flow_est_lap, source_reg_lap = lap(img, imgw, 30)

showflow(flow_est_lap.-flow)

mutated_flow = imresize(small_flow, ratio = 3/2)

showflow(mutated_flow)
showflow(flow)
showflow(flow.-mutated_flow)

##
using CSV, FileIO

base_path = "/Users/MrTrololord/Documents/anhir/";

loc_table = CSV.read(joinpath(base_path, "location_table.csv"))
train_rows = loc_table[loc_table[:status] .== "training", :]

random_row = train_rows[rand(1:size(train_rows, 1)), :]
target_path = random_row[Symbol("Target image")]

img = load(joinpath(base_path, "dataset", target_path))
imgg = Gray.(img)

imfilter(img, ones(13,13))

eltype(img)
eltype(imgg)

sizeof(eltype(img))
sizeof(eltype(imgg))

##
target = img

level_count = floor(Int64, log2(minimum((1024,1024))/8)+1)+1

minimum(size(target))/1024

half_size_pyramid = Int64.(2 .^ range(level_count-1, stop=0, length=level_count))

1024/8

128/16

## optable

macro optable(expr)
    if expr.args[1] == :(=>) && expr.args[2] isa Int
        n = expr.args[2]
        nexpr = expr.args[3]
        name = string(nexpr.args[1])
        descr = string(nexpr)
        :(optable($(esc(nexpr)), $name, $descr, $n))
    elseif expr.args[1] == :(=>) && expr.args[2] isa String
        name = expr.args[2]
        nexpr = expr.args[3]
        descr = string(nexpr)
        :(optable($(esc(nexpr)), $name, $descr))
    else
        name = string(expr.args[1])
        descr = string(expr)
        :(optable($(esc(expr)), $name, $descr))
    end
end

function optable(op, name, descr)
    fname = joinpath("..", "assets", string(name, ".png"))
    i = 2
    while isfile(fname)
        fname = joinpath("..", "assets", string(name, i, ".png"))
        i = i + 1
    end
    out = augment(pattern, op)
    save(fname, out)
    header = length(descr) < 20 ? "Output for `$descr`" : "`$descr`"
    tbl = string(
        "Input | $header\n",
        "------|--------\n",
        "![input](../assets/testpattern.png) | ![output]($fname)\n"
    )
    Markdown.parse(tbl)
end

function optable(img1, img2, basename1, basename2, descr1, descr2)
    # save fig 1 as next in line
    imgshow(img1); i = 1;
    fname1 = joinpath("..", "..", @__DIR__, "assets", string(basename1, i, ".png"))
    while isfile(fname1)
        i = i + 1
        fname1 = joinpath("..", "..", @__DIR__, "assets", string(basename1, i, ".png"))
    end
    savefig(fname1)
    # save fig 2 as next in line
    imgshow(img2); j = 1;
    fname2 = joinpath("..", "..", @__DIR__, "assets", string(basename2, j, ".png"))
    while isfile(fname2)
        i = i + 1
        fname2 = joinpath("..", "..", @__DIR__, "assets", string(basename2, j, ".png"))
    end
    savefig(fname2)
    tbl = string(
        "$descr1 | $descr2\n",
        "------|--------\n",
        "![input]($fname1) | ![output]($fname2)\n"
    )
    Markdown.parse(tbl)
end

optable(img, imgw, "img", "imgw", "target", "source")

function homografy_flow(flow_size, max_magnitude = 10)

    a, b, c = randn(), randn(), randn()
    α, β, γ = randn(ComplexF64), randn(ComplexF64), randn(ComplexF64)

    u(x,y)=(α*x + β*y + γ)/(a*x + b*y + c)

    X = ones(flow_size[1]) * collect(range(0,1,length=flow_size[2]))'
    Y = collect(range(0,1,length=flow_size[1])) * ones(flow_size[2])'
    A = Y .* im + X

    B = u.(X,Y)

    # max_len = maximum(LAP_julia.vec_len.(B))
    # B = B .* (max_magnitude/max_len)

end

rand_uniform(a, b) = rand() * (b-a) + a
using Distributions
mu = 0    #The mean of the truncated Normal
sigma = 1 #The standard deviation of the truncated Normal
lb = 0    #The truncation lower bound
ub = 1    #The truncation upper bound
d = Truncated(Normal(mu, sigma), lb, ub)  #Construct the distribution type
x = rand(Normal(mu, sigma), 100) #Simulate 100 obs from the truncated Normal
rand(Normal(mu, 0.2))

function rand_H(flow_size)
    a = flow_size[1] * 0.4
    H = [rand(Normal(1, 0.4)) rand(Normal(mu, 0.2)) rand_uniform(-a, a);
         rand(Normal(mu, 0.2)) rand(Normal(1, 0.4)) rand_uniform(-a, a);
         rand(Normal(mu, 0.0002)) rand(Normal(mu, 0.0002)) 1.]
    @info H
    return H
end

function gen_homo_flow(flow_size, max_magnitude=10)
    H = rand_H(flow_size)
    flow = LAP_julia.make_flow_from_H(H, flow_size)
    max_len = maximum(LAP_julia.vec_len.(flow))
    flow = flow .* (max_magnitude/max_len)
    @assert median(LAP_julia.vec_len.(flow)) > (0.1*max_magnitude) H
    return flow
end

flow = gen_homo_flow((123,123), 3)
showflow(flow)

img = gen_spaghetti(size(flow))
imgw = warp_img(img, -real.(flow), -imag.(flow))

method = sparse_pflap
timer = TimerOutput(string(method));
method_kwargs = Dict(:display => false, :timer => timer, :match_source_histogram => true, :spacing => 8)
flow_est, source_reg, timer, time_in_secs = test_registration_alg(method, img, imgw, flow.*(-1), timer=method_kwargs[:timer], method_kwargs=method_kwargs);

source_reg_zeros = warp_img(img, -real.(flow_est), -imag.(flow_est), border_strat=:zeros)
showflow(flow.*(-1))
showflow(flow_est)

# (Mat_<float>(3,3) << (1-rng.uniform(-0.05f, 0.05f)),
# (rng.uniform(-0.03f, 0.03f)), (rng.uniform(10.f, 20.f)),
# (rng.uniform(-0.03f, 0.03f)), (1-rng.uniform(-0.05f, 0.05f)),(rng.uniform(10.f, 20.f)),
# (rng.uniform(0.0001f, 0.0003f)), (rng.uniform(0.0001f, 0.0003f)), 1.f);
img, imgw, flow = gen_init();

flow = homografy_flow(size(img))
showflow(flow)
showflow(flow[215:230, 100:110], skip_count = 0)

mean(flow)
median(real.(flow))
median(imag.(flow))
maximum(real.(abs.(flow)))
maximum(imag.(abs.(flow)))

findall(real.(abs.(flow)) .== maximum(real.(abs.(flow))))

imgshow(imag.(flow), origin_bot=true)

flow1 = homografy_flow(size(img))
showflow(flow1)

imgw = warp_img(img, flow)

imgshow(img)
imgshow(imgw)

timer=TimerOutput("sparse pf lap");
method_kwargs = Dict(:timer => timer, :display => false, :max_repeats => 1, :point_count => 500, :spacing => 10)
flow_est, source_reg, timer, results = test_registration_alg(sparse_pflap, img, imgw, flow1, method_kwargs=method_kwargs, timer=timer)

showflow(flow_est)

figs[1][1,1,4]

imgshow(img)

showflow(flow1)
showflow(flow_est-flow1)

timer=TimerOutput("pf lap");
method_kwargs =Dict(:timer => timer, :display => false, :max_repeats => 1)
flow_est, source_reg, timer, results = test_registration_alg(pflap, img, imgw, flow1, method_kwargs=method_kwargs, timer=timer);


text:
using Sockets: IPv6
is_ch_type(::Val{:IPv6})  = true
result_type(::Val{:IPv6})  = Vector{IPv6}

function read_col_data(sock::ClickHouseSock, num_rows::VarUInt, ::Val{:IPv6})
    tmp = Vector{IPv6}(undef, UInt64(num_rows))
    Base.read!(sock.io, tmp)
    return tmp
end

function write_col_data(sock::ClickHouseSock,
                        data::AbstractVector{IPv6}, ::Val{:IPv6})

    Base.write(sock.io, data)
end


text:
using KrylovMethods
using Base.Test
include("getDivGrad.jl")

println("=== Testing SSOR for real matrix ===")
# small full system

# CG: test sparse Laplacian
A = getDivGrad(32,32,32)
rhs = randn(size(A,1))
tol = 1e-2

# tests with A being matrix
x1,flag1,relres1,iter1,resvec1       = ssor(A,rhs,tol=tol,maxIter=100)
x2,flag2,relres2,iter2,resvec2       = ssor(A,rhs,tol=tol,maxIter=100,omega=2/3,out=1)
x3,flag3,relres3,iter3,resvec3     = ssor(A,rhs,tol=tol,maxIter=100,omega=1/2,out=1,storeInterm=true)

# test relative residuals
@test norm(A*x1-rhs)/norm(rhs) <= tol
@test norm(A*x2-rhs)/norm(rhs) <= tol
@test norm(A*x3[:,end]-rhs)/norm(rhs) <= tol

println("=== Testing SSOR as Preconditioner ===")
omega = 1.2;
d = omega./diag(A);
x = zeros(length(rhs)) # pre allocation for the preconditioner result.
PC = r -> (x[:]=0.0; return ssorPrecTrans!(A,x,r,d));
y = KrylovMethods.cg(A,rhs,tol=1e-12,maxIter=200,M=PC,out=1)[1]
@test norm(A*y-rhs)/norm(rhs) <= 1e-12

println("=== Testing SSOR for complex matrix ===")
# small full system

# CG: test sparse Laplacian
A = getDivGrad(32,32,32)
A += 1im*speye(size(A,1))
rhs = randn(size(A,1)) + 1im * randn(size(A,1))
tol = 1e-2

#
x0,flag0,relres0,iter0,resvec0       = ssor(A,rhs,tol=tol,maxIter=3,out=2)
@test flag0==-1
@test iter0==3
@test length(resvec0)==3

# tests with A being matrix
x1,flag1,relres1,iter1,resvec1       = ssor(A,rhs,tol=tol,maxIter=100)
x2,flag2,relres2,iter2,resvec2       = ssor(A,rhs,tol=tol,maxIter=100,omega=2/3,out=1)
x3,flag3,relres3,iter3,resvec3       = ssor(A,rhs,tol=tol,maxIter=100,omega=1/2,out=1)

# test relative residuals
@test norm(A*x1-rhs)/norm(rhs) <= tol
@test norm(A*x2-rhs)/norm(rhs) <= tol
@test all(diff(resvec1).<0)
@test all(diff(resvec2).<0)
@test all(diff(resvec3).<0)

println("=== SSOR : All tests passed. ====")