text:
using HDF5
using FastGaussQuadrature
using LaTeXStrings
using LinearAlgebra
using Particles
using Random
using ReducedBasisMethods
using SparseArrays
using Statistics

fpath = "../runs/BoT_Np5e4_k_010_050_np_10_T25_projections.h5"

params = read_sampling_parameters(fpath)

μₜᵣₐᵢₙ = h5read(fpath, "parameters/mu_train")

IP = IntegratorParameters(fpath)

poisson = PoissonSolverPBSplines(fpath)

Ψ = h5read(fpath, "projections/Psi")
Ψₑ = h5read(fpath, "projections/Psi_e")
Πₑ = sparse(h5read(fpath, "projections/Pi_e"))

# Reference draw
P₀ = ParticleList(h5read(fpath, "initial_condition/x_0"),
                  h5read(fpath, "initial_condition/v_0"),
                  h5read(fpath, "initial_condition/w") )

nₜₑₛₜ = 10
κₜₑₛₜ_ₘᵢₙ = 0.1
κₜₑₛₜ_ₘₐₓ = 0.5

μₜₑₛₜ = zeros(nₜₑₛₜ, 5)
for i in 1:nₜₑₛₜ
    μₜₑₛₜ[i,:] = [κₜₑₛₜ_ₘᵢₙ, params.ε, params.a, params.v₀, params.σ]
end

λ = 0
for i in 1:nₜₑₛₜ
    if nₜₑₛₜ > 1
        μₜₑₛₜ[i,1] = rand(1)[1]*(κₜₑₛₜ_ₘₐₓ - κₜₑₛₜ_ₘᵢₙ) + κₜₑₛₜ_ₘᵢₙ
#         μₜₑₛₜ[i,1] = (1-λ)*κₜₑₛₜ_ₘᵢₙ + λ*κₜₑₛₜ_ₘₐₓ
#         λ += 1/(nₜₑₛₜ-1)
    end
end

μₜₑₛₜ = μₜₑₛₜ[sortperm(μₜₑₛₜ[:, 1]), :]

GC.gc()

IPₜₑₛₜ = IntegratorParameters(IP.dt, IP.nₜ, IP.nₜ+1, IP.nₕ, IP.nₚ, nₜₑₛₜ)
ICₜₑₛₜ = IntegratorCache(IPₜₑₛₜ)

@time Rₜₑₛₜ = ReducedBasisMethods.integrate_vp(P₀, μₜₑₛₜ, params, poisson, IPₜₑₛₜ, ICₜₑₛₜ;
                                              given_phi = false, save = true)
# Xₜₑₛₜ = Rₜₑₛₜ.X
# Vₜₑₛₜ = Rₜₑₛₜ.V
# Φₜₑₛₜ = Rₜₑₛₜ.Φ;

Φₜₑₛₜ = copy(Rₜₑₛₜ.Φ)

# no saving
@time ReducedBasisMethods.integrate_vp(P₀, μₜₑₛₜ, params, poisson, IPₜₑₛₜ, ICₜₑₛₜ;
                                        given_phi = false, save = false)

# Reduced Model
k = size(Ψ)[2]
kₑ = size(Ψₑ)[2]

RIC = ReducedIntegratorCache(IPₜₑₛₜ, k, kₑ)

ΨᵀPₑ = Ψ' * Ψₑ * inv(Πₑ' * Ψₑ)
ΠₑᵀΨ = Πₑ' * Ψ

@time Rᵣₘ = reduced_integrate_vp(P₀, Ψ, ΨᵀPₑ, ΠₑᵀΨ, μₜₑₛₜ, params, poisson, IPₜₑₛₜ, RIC;
                                   DEIM=true, given_phi = false, save = true)
# Xᵣₘ = Ψ * Rᵣₘ.Zₓ
# Vᵣₘ = Ψ * Rᵣₘ.Zᵥ
# Φᵣₘ = Rᵣₘ.Φ;

# no saving
@time reduced_integrate_vp(P₀, Ψ, ΨᵀPₑ, ΠₑᵀΨ, μₜₑₛₜ, params, poisson, IPₜₑₛₜ, RIC;
                            DEIM=true, given_phi = false, save=false)

# Saving
h5save("../runs/BoT_Np5e4_k_010_050_np_10_T25_DEIM.h5", IPₜₑₛₜ, poisson, params, μₜᵣₐᵢₙ, μₜₑₛₜ, Rₜₑₛₜ, Rᵣₘ, Ψ);

println(norm(Rᵣₘ.Φ - Φₜₑₛₜ))
println(norm(Rₜₑₛₜ.Φ - Φₜₑₛₜ))


text:
function solve_re{T<:FloatingPoint}(model::Blanchard_Kahn_Form{T},cutoff::T)

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  a     = copy(model.a)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  n = nx+ny

  (s,q) = schur(complex(a))
  q = q'  # So now q*a*q' = s

  # Reorder the eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,q) = reorder_complex_schur(s,q)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(s[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  q11 = q[1:nx,1:nx]
  q12 = q[1:nx,(nx+1):n]
  q21 = q[(nx+1):n,1:nx]
  q22 = q[(nx+1):n,(nx+1):n]
  s11 = s[1:nx,1:nx]

  p = real((q11-(q12/q22)*q21)\s11*(q11-(q12/q22)*q21))
  k = c[1:nx,:]
  f = -real(q22\q21)

  soln = Blanchard_Kahn_Soln(p,k,f,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Blanchard_Kahn_Form{T},cutoff::T,tol::T)

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  a     = copy(model.a)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  n = nx+ny

  s = eigvals(a)

  grc = 0
  for i = 1:n

    if norm(s[i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium

  a11 = a[1:nx,1:nx]
  a12 = a[1:nx,(nx+1):n]
  a21 = a[(nx+1):n,1:nx]
  a22 = a[(nx+1):n,(nx+1):n]

  f0 = -a22\a21
  f = Array(T,ny,nx)
  len = Inf
  while len > tol

    f = (a22-f0*a12)\(f0*a11-a21)
    len = maxabs(f0-f)
    f0 = copy(f)

  end

  p = a11+a12*f
  k = c[1:nx,:]

  soln = Blanchard_Kahn_Soln(p,k,f,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Klein_Form{T},cutoff::T)

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  a     = copy(model.a)
  b     = copy(model.b)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  n = nx+ny

  (s,t,q,z) = schur(complex(a),complex(b))
  q = q'  # So now q*a*z = s and q*b*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  z11 = z[1:nx,1:nx]
  z21 = z[(nx+1):n,1:nx]
  t11 = t[1:nx,1:nx]
  s11 = s[1:nx,1:nx]

  p = real((z11/t11)*(s11/z11))
  k = c[1:nx,:]
  f = real(z21/z11)

  soln = Klein_Soln(p,k,f,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Structural_Form{T},cutoff::T)

  a     = copy(model.a)
  a1    = copy(model.a1)
  b     = copy(model.b)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  nx = size(a1,1)
  ny = size(b,1)

  n = nx+ny

  a = [zeros(nx,nx) eye(ny); -a1 a]
  b = [eye(nx) zeros(nx,ny); zeros(ny,nx) b]
  c = [zeros(nx,size(c,2)); -c]

  (s,t,q,z) = schur(complex(a),complex(b))
  q = q'  # So now q*a*z = s and q*b*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  z11 = z[1:nx,1:nx]
  z12 = z[1:nx,(nx+1):n]
  z21 = z[(nx+1):n,1:nx]
  z22 = z[(nx+1):n,(nx+1):n]
  t11 = t[1:nx,1:nx]
  s11 = s[1:nx,1:nx]
  s12 = s[1:nx,(nx+1):n]
  s22 = s[(nx+1):n,(nx+1):n]

  qc  = -q*c
  qcs = qc[1:nx,:]       # The part of qc that is in the stable block
  qcu = qc[(nx+1):n,:]   # The part of qc that is in the unstable block

  p = real((z11/t11)*(s11/z11))
  k = real((z11/t11)*(s12/s22)*qcu-(z11/t11)*(s11/z11)*(z12/s22)*qcu-(z11/t11)*qcs)

  soln = Structural_Soln(p,k,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Structural_Form{T},cutoff::T,tol::T)

  a     = copy(model.a)
  a1    = copy(model.a1)
  b     = copy(model.b)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  p0 = copy(a1)
  p  = copy(p0)

  len = Inf
  while len > tol

    p = (a-b*p0)\a1
    len = maxabs(p-p0)
    p0 = copy(p)

  end
  k = (a-b*p0)\c

  nx = size(a1,1)
  ny = size(b,1)
  aa = [zeros(nx,nx) eye(ny); -a1 a]
  bb = [eye(nx) zeros(nx,ny); zeros(ny,nx) b]

  lambda = eigvals(bb,aa)

  grc = 0
  for i = 1:(nx+ny)

    if norm(lambda[i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  soln = Structural_Soln(p,k,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Sims_Form{T},cutoff::T)

  # This function is modeled on a routine written by
  # Christopher A. Sims (http://www.princeton.edu/~sims/)

  gamma0 = copy(model.gamma0)
  gamma1 = copy(model.gamma1)
  c      = copy(model.c)
  psi    = copy(model.psi)
  pi     = copy(model.pi)
  sigma  = copy(model.sigma)

  n  = size(gamma0,1)

  (s,t,q,z) = schur(complex(gamma1),complex(gamma0))
  q = q'  # So now q*gamma1*z = s and q*gamma0*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  # Define some ingredients that are needed for any solution

  q1 = q[1:(n-grc),:]
  q2 = q[(n-grc+1):n,:]
  z1 = z[:,1:(n-grc)]'
  z2 = z[:,(n-grc+1):n]'
  t2 = t[(n-grc+1):n,(n-grc+1):n]
  s2 = s[(n-grc+1):n,(n-grc+1):n]
  q2pi  = q2*pi
  q2psi = q2*psi

  # Check for the existance of a solution

  (pi_rows,pi_cols)   = size(pi)
  (psi_rows,psi_cols) = size(psi)

  if grc == 0
    q2pi    = zeros(0,pi_cols)
    uq2pi   = zeros(0,0)
    dq2pi   = zeros(0,0)
    vq2pi   = zeros(pi_cols,0)

    q2psi   = zeros(0,psi_cols)
    uq2psi  = zeros(0,0)
    dq2psi  = zeros(0,0)
    vq2psi  = zeros(psi_cols,0)

    non_zero_index_psi = []
  else
    (uq2pi,dq2pi,vq2pi) = svd(q2pi)
    non_zero_index_pi = find(dq2pi .> 2*eps(T)*norm(q2pi))
    uq2pi = uq2pi[:,non_zero_index_pi]
    vq2pi = vq2pi[:,non_zero_index_pi]
    dq2pi = diagm(dq2pi[non_zero_index_pi])

    (uq2psi,dq2psi,vq2psi) = svd(q2psi)
    non_zero_index_psi = find(dq2psi .> 2*eps(T)*norm(q2psi))
    uq2psi = uq2psi[:,non_zero_index_psi]
    vq2psi = vq2psi[:,non_zero_index_psi]
    dq2psi = diagm(dq2psi[non_zero_index_psi])
  end

  if isempty(non_zero_index_psi)
    iid_exist = true
  else
    iid_exist = (norm(uq2psi - uq2pi*uq2pi'*uq2psi) < 2*eps(T)*norm(uq2psi))
  end

  if iid_exist == false
    general_exist = false
  else
    t2iq2psi = t2\q2psi
    p        = t2\s2
    h        = t2iq2psi
    for i = 1:n

      h = [p*h t2iq2psi]

    end
    (uh,dh,vh) = svd(h)
    non_zero_index_dh = find(dh .> 2*eps(T)*norm(h))
    uh = uh[:,non_zero_index_dh]
    vh = vh[:,non_zero_index_dh]
    dh = diagm(dh[non_zero_index_dh])
    general_exist = (norm(uh-uq2pi*uq2pi'*uh) < 2*eps(T)*norm(uh,1))
  end

  # Check for uniqueness

  if grc == n
    q1pi  = zeros(0,pi_cols)
    uq1pi = zeros(0,0)
    vq1pi = zeros(pi_cols,0)
    dq1pi = zeros(0,0)
  else
    q1pi = q1*pi
    (uq1pi,dq1pi,vq1pi) = svd(q1pi)
    non_zero_index_dq1pi = find(dq1pi .> 2*eps(T)*norm(q1pi))
    uq1pi = uq1pi[:,non_zero_index_dq1pi]
    vq1pi = vq1pi[:,non_zero_index_dq1pi]
    dq1pi = diagm(dq1pi[non_zero_index_dq1pi])
  end

  if isempty(vq1pi)
    unique = true
  else
    resid = vq1pi-vq2pi*vq2pi'vq1pi
    (u1,d1,v1) = svd(resid)
    unique = ((sumabs(d1) > 2*eps(T)*n) != true)
  end
  tmat = [eye(n-grc) -(uq2pi*(dq2pi\vq2pi')*vq1pi*dq1pi*uq1pi')']
  g0   = [tmat*t; zeros(grc,(n-grc)) eye(grc)]
  g1   = [tmat*s; zeros(grc,n)]

  g0i    = inv(g0)
  g1     = real(z*g0i*g1*z')
  c      = real(z*g0i*[tmat*q*c; (t[(n-grc+1):n,(n-grc+1):n]-s[(n-grc+1):n,(n-grc+1):n])\q2*c])
  impact = real(z*g0i*[tmat*q*psi; zeros(grc,size(psi,2))])
  fmat   = real(s[(n-grc+1):n,(n-grc+1):n]\t[(n-grc+1):n,(n-grc+1):n])
  fwt    = -real(s[(n-grc+1):n,(n-grc+1):n]\q2*psi)
  ywt    = real(z*g0i[:,(n-grc+1):n])

  soln = Sims_Soln(g1,c,impact,ywt,fmat,fwt,sigma,grc,iid_exist,general_exist,unique)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Gomme_Klein_Form,cutoff::T)

  nx     = copy(model.nx)
  ny     = copy(model.ny)
  deriv1 = copy(model.derivs)
  deriv2 = copy(model.hessians)
  eta    = copy(model.eta)
  sigma  = copy(model.sigma)

  n = nx+ny

  # Reorder the second derivatives to put them in the order that Gomme and Klein use
  # Probably do away with this in a subsequent update

  hes = [deriv2[(n+1):2*n,:]; deriv2[1:n,:]]
  for i = 2:n

    hes = [hes; deriv2[2*n*(i-1)+n+1:2*n*i,:]; deriv2[2*n*(i-1)+1:2*n*(i-1)+n,:]]

  end
  hes = [hes[:,(n+1):2*n] hes[:,1:n]]

  # Construct partitioned first derivative matrices

  fx  = deriv1[1:n,1:nx]
  fy  = deriv1[1:n,(nx+1):n]
  fxp = deriv1[1:n,(n+1):(n+nx)]
  fyp = deriv1[1:n,(n+nx+1):(2*n)]

  # Solve the linearized first-order system

  a =  [fx fy]
  b = -[fxp fyp]
  c =  [eta; zeros(ny,size(eta,2))]

  first_order_model = Klein_Form(nx,ny,a,b,c,model.sigma)
  first_order_soln  = solve_re(first_order_model,cutoff)

  hx = first_order_soln.p
  gx = first_order_soln.f
  grc = first_order_soln.grc

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Set up the LP problem needed to construct the matrices on the second-order states

  m = [hx; gx*hx; eye(nx); gx]
  q = kron(eye(n),m')*hes*m
  b1 = kron(fxp,eye(nx))
  b2 = kron(fyp,eye(nx))
  b4 = kron(fy,eye(nx))
  c1 = kron(eye(ny),hx')
  c2 = kron(gx,eye(nx))

  qq = [kron(eye(nx),b4)+kron(hx',b2*c1) kron(eye(nx),b1+b2*c2)]
  xx = -qq\vec(q)
  gxx = reshape(xx[1:nx^2*ny],nx*ny,nx)
  hxx = reshape(xx[nx^2*ny+1:n*nx*nx],nx^2,nx)

  # Set up the LP problem needed to construct the intercepts, which contain the volatility effects

  qq = [fxp+fyp*gx fyp+fy]
  nn = [eye(nx); gx; zeros(n,nx)]
  q = fyp*tracem(kron(eye(ny),eta*(sigma')*eta')*gxx)+tracem(kron(eye(n),nn')*hes*nn*eta*(sigma')*eta')
  ss = -qq\q
  ssh = ss[1:nx]
  ssg = ss[nx+1:n]

  soln = Gomme_Klein_Soln(ssh,hx,hxx,ssg,gx,gxx,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Lombardo_Sutherland_Form,cutoff::T)

  nx     = copy(model.nx)
  ny     = copy(model.ny)
  deriv1 = copy(model.derivs)
  deriv2 = copy(model.hessians)
  eta    = copy(model.eta)
  sigma  = copy(model.sigma)

  n  = nx+ny
  ns = size(eta,2)

  # Solve the first order system

  a1 = -deriv1[1:n,(n+1):2*n]
  a2 = deriv1[1:n,1:n]
  c  = [eta; zeros(ny,size(eta,2))]

  (s,t,q,z) = schur(complex(a2),complex(a1))
  q = q'  # So now q*a*z = s and q*b*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  z11 = z[1:nx,1:nx]
  z12 = z[1:nx,nx+1:n]
  z21 = z[(nx+1):n,1:nx]
  z22 = z[(nx+1):n,(nx+1):n]

  t11 = t[1:nx,1:nx]
  t12 = t[1:nx,nx+1:n]
  t22 = t[nx+1:n,nx+1:n]

  s11 = s[1:nx,1:nx]
  s12 = s[1:nx,nx+1:n]
  s22 = s[nx+1:n,nx+1:n]

  phi   = real((z11/t11)*(s11/z11))
  gamma = c[1:nx,:]
  omega = [eye(nx); real(z21/z11)]  # Defines [x; y] as a function of [x]

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Set up the matrices for the second order terms

  for i = 1:n

    for  j = 1:2*n

      deriv2[(i-1)*2*n+j,j] = deriv2[(i-1)*2*n+j,j]/2

    end

    deriv2[(i-1)*2*n+1:i*2*n,:] = 2*deriv2[(i-1)*2*n+1:i*2*n,:]

  end

  v2vh_n = vec_to_vech(n)

  a4 = v2vh_n*vec(deriv2[1:n,1:n])
  a5 = v2vh_n*vec(deriv2[(n+1):2*n,(n+1):2*n])
  for i = 2:n

    a4 = [a4 v2vh_n*vec(deriv2[(i-1)*2*n+1:(i-1)*2*n+n,1:n])]
    a5 = [a5 v2vh_n*vec(deriv2[(i-1)*2*n+n+1:i*2*n,(n+1):2*n])]

  end

  a4 = a4'
  a5 = a5'

  vh2v_nx = vech_to_vec(nx)
  v2vh_nx = vec_to_vech(nx)
  vh2v_ns = vech_to_vec(ns)
  v2vh_ns = vec_to_vech(ns)
  p       = permutation(nx,ns)

  r           = v2vh_n*kron(omega,omega)*vh2v_nx            # n*(n+1)/2   times nx*(nx+1)/2
  phi_tilda   = v2vh_nx*kron(phi,phi)*vh2v_nx               # nx*(nx+1)/2 times nx*(nx+1)/2
  gamma_tilda = v2vh_nx*kron(gamma,gamma)*vh2v_ns           # nx*(nx+1)/2 times ns*(ns+1)/2
  psi_tilda   = v2vh_nx*(kron(phi,gamma)+kron(gamma,phi)*p) # nx*(nx+1)/2 times nx*ns

  g = a4*r+a5*r*phi_tilda                                   # n times nx*(nx+1)/2
  h = a5*r*gamma_tilda                                      # n times ns*(ns+1)/2

  sig = v2vh_ns*vec(sigma)                                  # ns*(ns+1)/2 times ns*ns

  # The next step is to resolve the system using the q and z matrices obtained earlier

  g = q*g
  h = q*h

  gs = g[1:nx,:]
  gu = g[nx+1:n,:]

  hs = h[1:nx,:]
  hu = h[nx+1:n,:]

  m = reshape((kron(phi_tilda',t22)-kron(eye(int(nx*(nx+1)/2)),s22))\vec(gu),ny,int(nx*(nx+1)/2))

  hx  = real(z11*(t11\s11)/z11)
  hxx = real(-(z11*(t11\s11)/z11)*z12*m + z11*(t11\(s12*m-t12*m*phi_tilda+gs)) + z12*m*phi_tilda)

  gx  = real(z21/z11)
  gxx = real((z22-z21*(z11\z12))*m)

  vbar = (eye(int(nx*(nx+1)/2))-phi_tilda)\gamma_tilda
  m3   = (eye(ny)-s22\t22)\(s22\(gu*vbar+hu))
  m2   = reshape((eye(int(nx*(nx+1)/2)*ny)-kron(phi_tilda',s22\t22))\(vec(s22\gu)),ny,int(nx*(nx+1)/2))

  p4 = -real(z22'\(m3-m2*vbar))
  p3 = -z22'\m2
  p2 = -z22'\z12'

  r4 = (t11*z21'+t12*z22')*p4
  r3 = (t11*z21'+t12*z22')*p3
  r2 = t11*z11'+t12*z12'+(t11*z21'+t12*z22')*p2

  d4 = (s11*z21'+s12*z22')*p4+hs

  ssg = p4*sig
  ssh = real(r2\(d4-r4-r3*gamma_tilda))*sig

  soln = Lombardo_Sutherland_Soln(ssh,hx,hxx,ssg,gx,gxx,phi_tilda,gamma_tilda,psi_tilda,sigma,grc,soln_type)

  return soln

end

ext: jl

max_stars_repo_name: dingliumath/SolveDSGE

max_issues_repo_name: dingliumath/SolveDSGE

max_forks_repo_name: dingliumath/SolveDSGE


text:
module SMMWeightMatrix

using Statistics
using Distributions
using HypothesisTests
using Random
using LinearAlgebra

include("summary_stats.jl")
include("weight_matrix.jl")

export WeightMatrix, block_bootstrap_estimator, select_moments

end # module


text:
type Quotes
    name::String
    symbols::String
    properties::Array{Tuple{String,Any},1}
end

type HistoricalData
    symbol::String
    dates::Array{Dates.Date,1}
    open::Array{Float64,1}
    high::Array{Float64,1}
    low::Array{Float64,1}
    close::Array{Float64,1}
    volume::Array{Float64,1}
    adjusted_close::Array{Float64,1}
end


text:
using LoopVectorization, Test
T = Float32

function test_broadcast(::Type{T}) where {T}
  M, N = 37, 47
  @show T, @__LINE__
  R = T <: Integer ? (T(-100):T(100)) : T
  a = rand(R, 99, 99, 99)
  b = rand(R, 99, 99, 1)
  bl = LowDimArray{(true, true, false)}(b)
  @test size(bl) == size(b)
  @test LoopVectorization.ArrayInterface.size(bl) ===
        (size(b, 1), size(b, 2), LoopVectorization.StaticInt(1))

  br = reshape(b, (99, 99))
  c1 = a .+ b
  c2 = @turbo inline = false a .+ bl
  @test c1 ≈ c2
  fill!(c2, 99999)
  @turbo inline = false c2 .= a .+ br
  @test c1 ≈ c2
  fill!(c2, 99999)
  @turbo inline = false c2 .= a .+ b
  @test c1 ≈ c2
  br = reshape(b, (99, 1, 99))
  bl = LowDimArray{(true, false, true)}(br)
  @test size(bl) == size(br)
  @test LoopVectorization.ArrayInterface.size(bl) ===
        (size(br, 1), LoopVectorization.StaticInt(1), size(br, 3))
  @. c1 = a + br
  fill!(c2, 99999)
  @turbo inline = false @. c2 = a + bl
  @test c1 ≈ c2
  fill!(c2, 99999)
  @turbo inline = false @. c2 = a + br
  @test c1 ≈ c2
  br = reshape(b, (1, 99, 99))
  bl = LowDimArray{(false,)}(br)
  @test size(bl) == size(br)
  @test LoopVectorization.ArrayInterface.size(bl) ===
        (LoopVectorization.StaticInt(1), size(br, 2), size(br, 3))
  @. c1 = a + br
  fill!(c2, 99999)
  @test c1 ≈ @turbo inline = false @. c2 = a + bl
  # @test c1 ≈ c2
  br = reshape(rand(R, 99), (1, 99, 1))
  bl = LowDimArray{(false,)}(br)
  @test size(bl) == size(br)
  @. c1 = a + br
  fill!(c2, 99999)
  @turbo inline = false @. c2 = a + bl
  @test c1 ≈ c2

  if T <: Integer
    xs = rand(-T(100):T(100), M)
  else
    xs = rand(T, M)
  end
  max_ = maximum(xs, dims = 1)
  @test (@turbo inline = false exp.(xs .- LowDimArray{(false,)}(max_))) ≈
        exp.(xs .- LowDimArray{(false,)}(max_))
  @test size(LowDimArray{(false,)}(max_)) == size(max_)

  if T === Int32
    a = rand(T(1):T(100), 73, 1)
    @test sqrt.(Float32.(a)) ≈ @turbo inline = false sqrt.(a)
  elseif T === Int64
    a = rand(T(1):T(100), 73, 1)
    @test sqrt.(a) ≈ @tturbo inline = false sqrt.(a)
  else
    a = rand(T, 73, 1)
    @test sqrt.(a) ≈ @turbo inline = false sqrt.(a)
  end

  a = rand(R, M)
  B = rand(R, M, N)
  c = rand(R, N)
  c′ = c'
  d1 = @. a + B * c′
  d2 = @tturbo inline = false @. a + B * c′
  @test d1 ≈ d2

  @. d1 = a + B * c′
  @turbo inline = false @. d2 = a + B * c′
  @test d1 ≈ d2

  d3 = a .+ B * c
  d4 = @turbo inline = false a .+ B *ˡ c
  @test d3 ≈ d4

  fill!(d3, -1000.0)
  fill!(d4, 91000.0)

  d3 .= a .+ B * c
  @turbo inline = false d4 .= a .+ B *ˡ c
  @test d3 ≈ d4

  fill!(d4, 91000.0)
  @turbo inline = false @. d4 = a + B *ˡ c
  @test d3 ≈ d4

  M, K, N = 77, 83, 57
  A = rand(R, M, K)
  B = rand(R, K, N)
  C = rand(R, M, N)
  At = copy(A')
  D1 = C .+ A * B
  D2 = @tturbo inline = false C .+ A .*ˡ B
  @test D1 ≈ D2
  if RUN_SLOW_TESTS
    fill!(D2, -999999)
    D2 = @turbo inline = false C .+ At' *ˡ B
    @test D1 ≈ D2
    fill!(D2, -999999)
    @test A * B ≈ (@turbo inline = false @. D2 = A *ˡ B)
    D1 .= view(C, 1, :)' .+ A * B
    fill!(D2, -999999)
    @turbo inline = false D2 .= view(C, 1, :)' .+ A .*ˡ B
    @test D1 ≈ D2
    C3d = rand(R, 3, M, N)
    D1 .= view(C3d, 1, :, :) .+ A * B
    fill!(D2, -999999)
    @tturbo inline = false D2 .= view(C3d, 1, :, :) .+ A .*ˡ B
    @test D1 ≈ D2
  end
  D1 .= 9999
  @turbo inline = false D2 .= 9999
  @test D1 == D2
  D1 .= -99999
  @tturbo inline = false D2' .= -99999
  @test D1 == D2

  b = rand(T, K)
  x = rand(R, N)
  D1 .= C .+ A * (b .+ x')
  @tturbo inline = false @. D2 = C + A *ˡ (b + x')
  @test D1 ≈ D2
  D2 = @turbo inline = false @. C + A *ˡ (b + x')
  @test D1 ≈ D2
  if T === Int64
    xd = rand(-1_000_000_000_000:1_000_000_000_000, 89)
  elseif T === Int32
    xd = rand(-Int32(10_000_000):Int32(10_000_000), 89)
  else
    xd = rand(R, 89)
  end
  yd = rand(R, 89)
  yd[yd.==0] .= 77
  @test xd .÷ yd == @turbo xd .÷ yd

  if T <: Union{Float32,Float64}
    D3 = cos.(B')
    D4 = @turbo inline = false cos.(B')
    @test D3 ≈ D4

    fill!(D3, -1e3)
    fill!(D4, 9e9)
    Bt = transpose(B)
    @. D3 = exp(Bt)
    @tturbo inline = false @. D4 = exp(Bt)
    @test D3 ≈ D4

    D1 = similar(B)
    D2 = similar(B)
    D1t = transpose(D1)
    D2t = transpose(D2)
    @. D1t = exp(Bt)
    @turbo inline = false @. D2t = exp(Bt)
    @test D1t ≈ D2t

    fill!(D1, -1e3)
    fill!(D2, 9e9)
    @. D1' = exp(Bt)
    lset = @tturbo inline = false @. D2' = exp(Bt)

    @test D1 ≈ D2

    a = rand(137)
    b1 = @turbo inline = false @. 3 * a + sin(a) + sqrt(a)
    b2 = @. 3 * a + sin(a) + sqrt(a)
    @test b1 ≈ b2
    three = 3
    fill!(b1, -9999)
    @tturbo inline = false @. b1 = three * a + sin(a) + sqrt(a)
    @test b1 ≈ b2

    C = rand(100, 10, 10)
    D1 = C .^ 0.3
    D2 = @tturbo inline = false C .^ 0.3
    @test D1 ≈ D2
    @. D1 = C^2
    @turbo inline = false @. D2 = C^2
    @test D1 ≈ D2
    @turbo view(C, 1:100, 1:10, 1:10) .= 0
    @test all(==(0), C)
    @turbo view(B, axes(B, 1), axes(B, 2))' .= 2
    @test all(==(2), B)
  end
end

@testset "broadcast_float" begin
  for T ∈ (Float32, Float64, Int32, Int64)
    @time test_broadcast(T)
  end
end


text:
# # Classification: Sparse Variational Approximation for Non-Conjugate Likelihoods with Optim's L-BFGS
#
#
# This example demonstrates how to carry out non-conjugate Gaussian process
# inference using the stochastic variational Gaussian process (SVGP) model. For
# a basic introduction to the functionality of this library, please refer to the
# [User Guide](@ref).
#
# ## Setup

using ApproximateGPs
using ParameterHandling
using Zygote
using Distributions
using LinearAlgebra
using Optim

using Plots
default(; legend=:outertopright, size=(700, 400))

using Random
Random.seed!(1234);

# ## Generate some training data
#
# For our binary classification model, we will use the standard approach of a
# latent GP with a Bernoulli likelihood. This results in a generative model that
# we can use to produce some training data.
#
# First, we define the underlying latent GP
# ```math
# f \sim \mathcal{GP}(0, k(\cdot, \cdot'))
# ```
# and sample a function `f`.

k_true = [30.0, 1.5]
kernel_true = k_true[1] * (SqExponentialKernel() ∘ ScaleTransform(k_true[2]))

jitter = 1e-8  # for numeric stability
lgp = LatentGP(GP(kernel_true), BernoulliLikelihood(), jitter)
x_true = 0:0.02:6
f_true, y_true = rand(lgp(x_true))

plot(x_true, f_true; seriescolor="red", label="")  # Plot the sampled function

# Then, the output of this sampled function is pushed through a logistic sigmoid
# `μ = σ(f)` to constrain the output to `[0, 1]`.

μ = mean.(lgp.lik.(f_true))
plot(x_true, μ; seriescolor="red", label="")

# Finally, the outputs `y` of the process are sampled from a Bernoulli
# distribution with mean `μ`. We're only interested in the outputs at a subset
# of inputs `x`, so we first pick some random input locations and then find the
# corresponding values for `y`.

N = 30  # The number of training points
mask = sample(1:length(x_true), N; replace=false, ordered=true)  # Subsample some input locations
x, y = x_true[mask], y_true[mask]

scatter(x, y; label="Sampled outputs")
plot!(x_true, mean.(lgp.lik.(f_true)); seriescolor="red", label="True mean")

# ## Creating an SVGP
#
# Now that we have some data sampled from a generative model, we can try to recover the
# true generative function with an SVGP classification model.
#
# For this, we shall use a mixture of
# [ParameterHandling.jl](https://github.com/invenia/ParameterHandling.jl) to
# deal with our constrained parameters and
# [Optim.jl](https://julianlsolvers.github.io/Optim.jl/stable/) to perform
# optimimisation.
#
# The required parameters for the SVGP are - the kernel hyperparameters `k`, the
# inducing inputs `z` and the mean and covariance of the variational
# distribution `q`; given by `m` and `A` respectively. ParameterHandling
# provides an elegant way to deal with the constraints on these parameters,
# since `k` must be positive and `A` must be positive definite. For more
# details, see the
# [ParameterHandling.jl](https://github.com/invenia/ParameterHandling.jl)
# readme.

# Initialise the parameters

M = 15  # number of inducing points
raw_initial_params = (
    k=(var=positive(rand()), precision=positive(rand())),
    z=bounded.(range(0.1, 5.9; length=M), 0.0, 6.0),  # constrain z to simplify optimisation
    m=zeros(M),
    A=positive_definite(Matrix{Float64}(I, M, M)),
);

# `flatten` takes the `NamedTuple` of parameters and returns a flat vector of
# `Float64` - along with a function `unflatten` to reconstruct the `NamedTuple`
# from a flat vector. `value` takes each parameter in the `NamedTuple` and
# applies the necessary transformation to return the constrained value which can
# then be used to construct the SVGP model. `unpack` therefore takes a flat,
# unconstrained `Vector{Float64}` and returns a `NamedTuple` of constrained
# parameter values.

flat_init_params, unflatten = ParameterHandling.flatten(raw_initial_params)
unpack = ParameterHandling.value ∘ unflatten;

# Now, we define a function to build everything needed for an SVGP model from
# the constrained parameters. The two necessary components are the `LatentGP`
# which we are trying to approximate and the
# `SparseVariationalApproximation` struct. This struct takes as arguments
# the inducing points `fz`, and the variational posterior distribution `q`.
# These elements can then be passed to the loss function (the `elbo`) along with
# the data `x` and `y`.

lik = BernoulliLikelihood()
jitter = 1e-3  # added to aid numerical stability

function build_SVGP(params::NamedTuple)
    kernel = params.k.var * (SqExponentialKernel() ∘ ScaleTransform(params.k.precision))
    f = LatentGP(GP(kernel), lik, jitter)
    q = MvNormal(params.m, params.A)
    fz = f(params.z).fx
    return SparseVariationalApproximation(fz, q), f
end

function loss(params::NamedTuple)
    svgp, f = build_SVGP(params)
    fx = f(x)
    return -elbo(svgp, fx, y)
end;

# Optimise the parameters using LBFGS.

opt = optimize(
    loss ∘ unpack,
    θ -> only(Zygote.gradient(loss ∘ unpack, θ)),
    flat_init_params,
    LBFGS(;
        alphaguess=Optim.LineSearches.InitialStatic(; scaled=true),
        linesearch=Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(; iterations=4_000);
    inplace=false,
)

# Finally, build the optimised SVGP model, and sample some functions to see if
# they are close to the true function.

final_params = unpack(opt.minimizer)

svgp_opt, f_opt = build_SVGP(final_params)
post_opt = posterior(svgp_opt)
l_post_opt = LatentGP(post_opt, BernoulliLikelihood(), jitter)

post_f_samples = rand(l_post_opt.f(x_true, 1e-6), 20)
post_μ_samples = mean.(l_post_opt.lik.(post_f_samples))

plt = plot(x_true, post_μ_samples; seriescolor="red", linealpha=0.2, label="")
scatter!(plt, x, y; seriescolor="blue", label="Data points")
vline!(final_params.z; label="Pseudo-points")
plot!(
    x_true, mean.(lgp.lik.(f_true)); seriescolor="green", linewidth=3, label="True function"
)


text:
using OMEinsum

N = 100
ca = randn(Float32, N, N)
ct = randn(Float32, N, N, N)
cq = randn(Float32, N, N, N, N)

function bfunc_star_cpu()
    ein"ji,kl,li->jkl"(ca,ca,ca)
end

function bfunc_t3_cpu()
    ein"ji,kli,li->jkl"(ca,ct,ca)
end

function bfunc_psum_cpu()
    ein"iikl->"(cq)
end

using BenchmarkTools
display(@benchmark bfunc_star_cpu() seconds = 1)
display(@benchmark bfunc_t3_cpu() seconds = 1)
display(@benchmark bfunc_psum_cpu() seconds = 1)


text:
@testset "994.rotting-oranges.jl" begin
    @test oranges_rotting([[2, 1, 1], [0, 1, 1], [1, 0, 1]]) == -1
    @test oranges_rotting([[2, 1, 1], [1, 1, 0], [0, 1, 1]]) == 4
    @test oranges_rotting([[0, 2]]) == 0
end


text:
# ---
# title: 78. Subsets
# id: problem78
# author: Jerry Ling
# date: 2021-12-01
# difficulty: Medium
# categories: Array, Backtracking, Bit Manipulation
# link: <https://leetcode.com/problems/subsets/description/>
# hidden: true
# ---
#
# Given an integer array `nums`, return _all possible subsets (the power set)_.
#
# The solution set must not contain duplicate subsets.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3]
#     Output: [[],[1],[2],[1,2],[3],[1,3],[2,3],[1,2,3]]
#
#
# **Example 2:**
#
#
#
#     Input: nums = [0]
#     Output: [[],[0]]
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 10`
#   * `-10 <= nums[i] <= 10`
#
#
## @lc code=start
using LeetCode

function subsets(arr::Vector{T}) where T
    res = [T[]]
    for e in arr
        for i in eachindex(res)
            push!(res, vcat(res[i], e))
        end
    end
    res
end

## add your code here:
## @lc code=end


text:
# CSCI5573Final/components.jl
# Licensed under the MIT License. See LICENSE.md file in the project root for
# full license information.

# ----------------------------- Begin Processor -----------------------------
mutable struct Processor
    task::Union{_Task, Nothing}
    Multiplier::Float64
    resource::Resource
    queue::Vector{Any}
    Processor(Multiplier, resource) = new(nothing, Multiplier, resource, Vector{Int64}())
end

busy(processor::Processor) = length(processor.queue) == 0

function TotalWork(processor::Processor)
    total = 0
    for ID in processor.queue
        total = total + TASKS[ID].Cost
    end
    return total * processor.Multiplier |> round |> Int64
end

function SendToProcessor(processor::Processor, ID::Int64, time::Int64)
    push!(processor.queue, ID)
    @with_resource processor.resource begin
        if !(ID in processor.queue)
            return
        end

        filter!(e->e!==ID, processor.queue)

        processor.task = TASKS[ID]

        for i in 1:time
            task_work = Int64(round(1/processor.Multiplier))
            working!(processor.task, task_work)

            work(CLOCK_CYCLE)

            if isDone(processor.task)
                break
            end
        end

        if isDone(processor.task)
            push!(TERMINATED, processor.task.ID)
        else
            enqueue!(READY_QUEUE, processor.task.ID)
        end
    end
    processor.task = nothing
end

# ----------------------------- End Processor -----------------------------

# ----------------------------- Begin Routing -----------------------------

@process Enqueuer() begin
    local Incoming::Vector{Int64} = []

    for i in unique(collect(keys(TASKS)))
        push!(Incoming, i)
    end

    active_tasks = []

    push!(TERMINATED, 0)
    filter!(e->e!==0, Incoming)

    while(true)
        launch = []

        # Clean up finished tasks and enqueue newly ready tasks
        for ID in copy(TERMINATED)
            for child in copy(TASKS[ID].Children)
                remove_dependency!(TASKS[child], ID)

                if length(TASKS[child].Dependencies) === 0 && child in Incoming
                    push!(launch, child)
                end
            end

            if verbose
                println("Killing $ID")
            end

            filter!(e->e!==ID, active_tasks)
            filter!(e->e!==ID, TERMINATED)
            filter!(e->e!==ID, Incoming)
        end

        wait(CLOCK_CYCLE)

        filter!(e->e in Incoming, launch)

        if verbose
            if !isempty(launch)
                println("launching:", launch)
            end
        end

        for task in copy(launch)
            if verbose
                println("Enqueuing $task")
            end
            enqueue!(READY_QUEUE, task)
            filter!(e->e!==task, Incoming)
            push!(active_tasks, task)
        end

        if verbose
            println("Active Tasks", active_tasks)
            println("Ready Queue", READY_QUEUE.data)
        end

        if length(Incoming) === 0 & length(active_tasks) === 1 & length(TERMINATED) === 0
            global COMPLETE = true
        end

        wait(CLOCK_CYCLE)
    end
end

@process Dispatcher(ID::Int64, processor::Int64, process_time::Int64, comm_time::Int64) begin
    process_store!(current_process(), :process_task, ID)
    local this_task = TASKS[ID]

    if this_task.Type === :END
        if verbose
            println("End Queued at $(current_time())")
        end
        global COMPLETE
        while(!COMPLETE)
            wait(1)
        end

        global FINISH_TIME = "$(current_time())"
        println("Simulation finished in $(FINISH_TIME)")
        stop_simulation()
    end

    if this_task.Type === :COMPUTATION
        SendToProcessor(PROCESSORS[processor], ID, process_time+comm_time)
    end
end
# ----------------------------- End Routing -----------------------------

# Just a little sneaky fella down here
# She manages the results destination.
SCHEDULER = ""


text:
# # [SIR example with event scheduling](@id sir_scheduling)
# Sean L. Wu (@slwu89), 2021-1-9

using Individual.Sampling
using Individual.SchemaBase
using Individual.SchemaEvents

using Catlab.Present, Catlab.CSetDataStructures, Catlab.Theories, Catlab.CategoricalAlgebra, Catlab.Graphics, Catlab.Graphs
using Plots, GraphViz

# ## Introduction
# This tutorial shows how to simulate the SIR model using the event scheduling schema available in Individual.jl. The stochastic
# process being simulated is exactly the same as the basic [Markov SIR model tutorial](@ref sir_basic), and all parameters
# are identical, so please see that tutorial for reference if needed.

# The schema looks like this:

to_graphviz(TheorySchedulingIBM)

# The schema expands on the basic `TheoryIBM` for Markov models; the relationship between states and people is the same.
# There is a new object, `Event`, which in epidemiological models might contain recovery, hospitalization, death, etc.
# There are morphisms into the objects `EventLabel` and `EventListener`. `EventLabel` is simply a string for writing
# self-documenting models. The `EventListener` stores functions called "event listeners" which are called with the
# set of persons scheduled for that event on that timestep, and may queue state updates, and schedule or cancel other events.

# The set of scheduled (queued) events is also an object `Scheduled`, with morphisms into the objects of events and persons, and a delay attribute.
# For each queued event these tell us which event will occur, who it will happen to, and after how many timesteps it (the event listener) should fire.
# This event scheduling mechanism allows events which occur after a non-Geometric (i.e. non-Markovian) delay, such as fixed delays
# or any other distribution.

# If we had additional combinatorial data describing each person (e.g. discrete age bin for each person), we could make another attribute and a morphism from people to that attribute.
# Additional atomic data for each person (e.g. neutralizing antibody titre) would be an attribute of people. In this way we can simulate
# a general class of individual based models relevant to epidemiology, ecology, and the social sciences.

# ## Parameters

N = 1000
I0 = 8
S0 = N - I0
Δt = 0.1
tmax = 100
steps = Int(tmax/Δt)
γ = 1/10 # recovery rate
R0 = 2.5
β = R0 * γ # R0 for corresponding ODEs

initial_states = fill("S", N)
initial_states[rand(1:N, I0)] .= "I"
state_labels = ["S", "I", "R"];

# ## Model object

# The "SchedulingIBM" (Individual Based Model with Scheduling) schema needs several type parameters.
# The first is the same as the "IBM" schema, associating a unique name to each categorical state.
# The second is the delay, the number of time steps after which a scheduled event will fire.
# The third is the name for each event. The last is for the listeners associated with each event.

const SIR = SchedulingIBM{String, Int64, String, Vector{Function}}()
initialize_states(SIR, initial_states, state_labels);

# ## Processes

# The infection process is the same as the basic SIR model.

function infection_process(t::Int)
    I = npeople(SIR, "I")
    N = npeople(SIR)
    λ = β * I/N
    S = get_index_state(SIR, "S")
    S = bernoulli_sample(S, λ, Δt)
    queue_state_update(SIR, S, "I")
end

# The recovery process is queues future recovery events. We first find who is infected,
# and then take the set difference of those persons with those who are already scheduled for
# recovery, which is just the set of persons who need a recovery scheduled.

function recovery_process(t::Int)

    I = get_index_state(SIR, "I")
    already_scheduled = get_scheduled(SIR, "Recovery")
    to_schedule = setdiff(I, already_scheduled)

    if length(to_schedule) > 0
        rec_times = delay_geom_sample(length(to_schedule), γ, Δt)
        schedule_event(SIR, to_schedule, rec_times, "Recovery")
    end
end

# ## Events

# The event listener associated with recovery is quite simple, just updating the state to R.
# We create an event with the label "Recovery" and a single listener, and add it to the model.

function recovery_listener(target, t::Int)
    queue_state_update(SIR, target, "R")
end

add_event(SIR, "Recovery", recovery_listener);

# ## Simulation

# We use `render_process` to create a rendering (output) process and
# a matrix giving state counts by time step. Then we draw a trajectory and plot the results.
# We need to use the `simulation_loop` from the module for models with event scheduling.

state_out, render_process = render_states(SIR, steps)

SchemaEvents.simulation_loop(SIR, [infection_process, recovery_process, render_process], steps)

plot(
    (1:steps) * Δt,
    state_out,
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number"
)


text:
export trouve_intervalleA
function trouve_intervalleA(h :: AbstractNLPModel;
                            t₀ :: Float64 = h.meta.x0[1],
                            tₘ :: Float64 = 100.0,
                            direction :: String  =  "Nwt",
                            tol :: Float64 = 1e-10,
                            verbose :: Bool = false,
                            kwargs...)
        # println("on est dans trouve_intervalleA")
        tim1 = t₀
        ti = (tim1 + tₘ) / 2

        i = 0

        h₀ = obj(h, [t₀])[1]; dh₀ = grad(h, [t₀])[1]

        him1 = obj(h, [tim1])[1]; dhim1 = grad(h, [tim1])[1]
        hi = obj(h, [ti])[1]; dhi = grad(h, [ti])[1]

        # println("on est avant les verbose")

        verbose &&
            @printf("iter tim1        dhim1        him1         ti        ")
        verbose &&
            @printf("dhi        hi\n")
        verbose &&
            @printf("%4d %7.2e %7.2e  %7.2e  %7.2e  %7.2e  %7.2e \n",
                    i,  tim1, dhim1, him1, ti, dhi, hi)
        # println("avant le while")
        while i < 50
          # print_with_color(:geen, "dans le while i = $i \n")
          him1 = hi
          hi = obj(h, [ti])[1]
          if (hi > h₀) ||((hi  > him1) && (i > 1))
            (topt, iter) = zoom_generic(h, tim1, ti;
                                        direction = direction,
                                        verbose = verbose)
            gopt = grad(h, [topt])[1]
            status = :NotSolved
            (abs(gopt) < 1e-07) && (status = :Optimal)
            (i >= 50) && (status = :Tired)
            tired = i > 50
            optimal = abs(gopt) < 1e-07
            return (topt, obj(h, [topt])[1], norm(gopt, Inf), i, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
          end

          dhim1 = dhi
          dhi = grad(h, [ti])[1]
          if (abs(dhi) <= tol)
            iter = i
            topt = ti
            gopt = grad(h, [topt])[1]
            status = :NotSolved
            (abs(gopt) < 1e-07) && (status = :Optimal)
            (i >= 50) && (status = :Tired)
            tired = i > 50
            optimal = abs(gopt) < 1e-07
            return (topt, obj(h, [topt])[1], norm(gopt, Inf), i, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
          end

          if (dhi >= 0.0)
            (topt, iter) = zoom_generic(h, ti, tim1,
                                        verbose = verbose,
                                        direction = direction)
            gopt = grad(h, [topt])[1]
            status = :NotSolved
            (abs(gopt) < 1e-07) && (status = :Optimal)
            (i >= 50) && (status = :Tired)
            tired = i > 50
            optimal = abs(gopt) < 1e-07
            return (topt, obj(h, [topt])[1], norm(gopt, Inf), i, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)

          end

          tim1 = ti
          ti = (tim1 + tₘ) / 2

          i += 1

          verbose &&
            @printf("%4d %7.2e %7.2e  %7.2e  %7.2e  %7.2e  %7.2e \n",
                    i,  tim1, dhim1, him1, ti, dhi, hi)
        end

        return (ti, obj(h, [ti])[1], norm(grad(h, [ti])[1], Inf), i, false, true, :stalled, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
end


text:
function JuMP.solve_time(bm::BilevelModel)
    return bm.solve_time
end
function build_time(bm::BilevelModel)
    return bm.build_time
end

function JuMP.set_optimizer_attribute(bm::BilevelModel, name::String, value)
    _check_solver(bm)
    return JuMP.set_optimizer_attribute(bm, MOI.RawParameter(name), value)
end
function JuMP.set_optimizer_attribute(
    bm::BilevelModel, attr::MOI.AbstractOptimizerAttribute, value
)
    _check_solver(bm)
    return MOI.set(bm.solver, attr, value)
end
function JuMP.set_optimizer_attributes(bm::BilevelModel, pairs::Pair...)
    for (name, value) in pairs
        JuMP.set_optimizer_attribute(bm.solver, name, value)
    end
end

function JuMP.get_optimizer_attribute(bm::BilevelModel, name::String)
    _check_solver(bm)
    return JuMP.get_optimizer_attribute(bm.solver, MOI.RawParameter(name))
end
function JuMP.get_optimizer_attribute(
    bm::BilevelModel, attr::MOI.AbstractOptimizerAttribute
)
    _check_solver(bm)
    return MOI.get(bm.solver, attr)
end

function JuMP.set_silent(bm::BilevelModel)
    _check_solver(bm)
    return MOI.set(bm.solver, MOI.Silent(), true)
end
function JuMP.unset_silent(bm::BilevelModel)
    _check_solver(bm)
    return MOI.set(bm.solver, MOI.Silent(), false)
end

function JuMP.set_time_limit_sec(bm::BilevelModel, limit)
    _check_solver(bm)
    return MOI.set(bm.solver, MOI.TimeLimitSec(), limit)
end
function JuMP.unset_time_limit_sec(bm::BilevelModel)
    _check_solver(bm)
    return MOI.set(bm.solver, MOI.TimeLimitSec(), nothing)
end
function JuMP.time_limit_sec(bm::BilevelModel)
    _check_solver(bm)
    return MOI.get(bm.solver, MOI.TimeLimitSec())
end

function JuMP.simplex_iterations(bm::BilevelModel)
    _check_solver(bm)
    return MOI.get(bm.solver, MOI.SimplexIterations())
end

function JuMP.barrier_iterations(bm::BilevelModel)
    _check_solver(bm)
    return MOI.get(bm.solver, MOI.BarrierIterations())
end

function JuMP.node_count(bm::BilevelModel)
    _check_solver(bm)
    return MOI.get(bm.solver, MOI.NodeCount())
end

function JuMP.result_count(bm::BilevelModel)::Int
    _check_solver(bm)
    return MOI.get(bm.solver, MOI.ResultCount())
end

function set_copy_names(bm::BilevelModel)
    bm.copy_names = true
    return nothing
end
function unset_copy_names(bm::BilevelModel)
    bm.copy_names = false
    return nothing
end
function get_copy_names(bm::BilevelModel)
    return bm.copy_names
end
function set_pass_start(bm::BilevelModel)
    bm.pass_start = true
    return nothing
end
function unset_pass_start(bm::BilevelModel)
    bm.pass_start = false
    return nothing
end
function get_pass_start(bm::BilevelModel)
    return bm.pass_start
end


text:
module ImputationAlgamest

export locf, locf!, lerp, qlerp,
       findmissings, findpresents, removemissings,
       findnothings, findsomethings, removenothings,
       findNaNs, findnonNaNs, removeNaNs,
       isnothing, issomething

const IntFloat = Union{Signed, AbstractFloat}

const Ints   = Union{Int32, Int64}
const Floats = Union{Float32, Float64}

const MaybeInts    = Union{Missing, Ints}
const MaybeFloats  = Union{Missing, Floats}

const MaybeNothing = Union{Missing, Nothing}

include("EndpointRanges.jl")
using .EndpointRanges

include("indicies.jl")  #  index non-values
include("locf.jl")      #  last observation carrys forward
include("lerp.jl")      #  simple and weighted interpolations
include("getset.jl")    #  datavector's access and restorage smoothified

end # ImputationAlgamest


text:
using Test
using Random
using StatisticalVerificationOfIsingEnergies

@testset "estimation minimum" begin

    x = [0.1, 0.1, 0.1, 0.2, 0.2, 0.3]
    @test estimate_ground_state_energy(x, 0.34) ≈ -0.06106446065424451
end

@testset "error calculation" begin
    x = [0.1, 0.1, 0.1, 0.2, 0.2, 0.3]
    @test squared_error(x, 0.34) ≈ 0.2614025705829984
end

@testset "bootstrap histogram of minimas" begin
    Random.seed!(1234)
    x = [1., 1., 2., 2., 3., 4., 5.]
    y = bootstrap_hists_of_mins(x, 0.19, 2)
    @test y ≈ [-6.276354333015176, -6.893114986285872]
end

@testset "p-value" begin
    Random.seed!(1234)
    x = [0.1, 1., 1.2, 2., 2.2, 3., 3.1, 3.5];
    @test bootstrap_get_pvalue(x, 0.19) ≈ 0.655
end

@testset "estimate temparature" begin
    x = [1. ,1.5, 2., 3., 4.]
    @test estiamte_temperature(x, 0.) ≈ 0.8958957699200034

end


text:
export batch_find_untargeted_attack

using DelimitedFiles, Dates

@enum SolveRerunOption never = 1 always = 2 resolve_ambiguous_cases = 3 refine_insecure_cases = 4 retarget_infeasible_cases =
    5

struct BatchRunParameters
    nn::NeuralNet
    pp::PerturbationFamily
    norm_order::Real
    tolerance::Real
end

# BatchRunParameters are used for result folder names
function Base.show(io::IO, t::BatchRunParameters)
    print(io, "$(t.nn.UUID)__$(t.pp)__$(t.norm_order)__$(t.tolerance)")
end

function mkpath_if_not_present(path::String)
    !isdir(path) ? mkpath(path) : nothing
end

function extract_results_for_save(d::Dict)::Dict
    m = d[:Model]
    r = Dict()
    r[:SolveTime] = d[:SolveTime]
    r[:ObjectiveBound] = getobjbound(m)
    r[:ObjectiveValue] = getobjectivevalue(m)
    r[:TargetIndexes] = d[:TargetIndexes]
    r[:SolveStatus] = d[:SolveStatus]
    r[:PredictedIndex] = d[:PredictedIndex]
    r[:TighteningApproach] = d[:TighteningApproach]
    r[:TotalTime] = d[:TotalTime]
    if !isnan(r[:ObjectiveValue])
        r[:PerturbationValue] = d[:Perturbation] |> getvalue
        r[:PerturbedInputValue] = d[:PerturbedInput] |> getvalue
    end
    return r
end

function is_infeasible(s::Symbol)::Bool
    s == :Infeasible || s == :InfeasibleOrUnbounded
end

function get_uuid()::String
    Dates.format(now(), "yyyy-mm-ddTHH.MM.SS.sss")
end

function generate_csv_summary_line(
    sample_number::Integer,
    results_file_relative_path::String,
    r::Dict,
)
    [
        sample_number,
        results_file_relative_path,
        r[:PredictedIndex],
        r[:TargetIndexes],
        r[:SolveTime],
        r[:SolveStatus],
        r[:SolveStatus] |> is_infeasible,
        r[:ObjectiveValue],
        r[:ObjectiveBound],
        r[:TighteningApproach],
        r[:TotalTime],
    ] .|> string
end

function generate_csv_summary_line_optimal(sample_number::Integer, d::Dict)
    @assert(d[:PredictedIndex] in d[:TargetIndexes])
    [
        sample_number,
        "",
        d[:PredictedIndex],
        d[:TargetIndexes],
        0,
        :Optimal,
        false,
        0,
        0,
        :NA,
        d[:TotalTime],
    ] .|> string
end

function create_summary_file_if_not_present(summary_file_path::String)
    if !isfile(summary_file_path)
        summary_header_line = [
            "SampleNumber",
            "ResultRelativePath",
            "PredictedIndex",
            "TargetIndexes",
            "SolveTime",
            "SolveStatus",
            "IsInfeasible",
            "ObjectiveValue",
            "ObjectiveBound",
            "TighteningApproach",
            "TotalTime",
        ]

        open(summary_file_path, "w") do file
            writedlm(file, [summary_header_line], ',')
        end
    end
end

function verify_target_indices(
    target_indices::AbstractArray{<:Integer},
    dataset::MIPVerify.LabelledDataset,
)
    num_samples = MIPVerify.num_samples(dataset)
    @assert(minimum(target_indices) >= 1, "Target sample indexes must be 1 or larger.")
    @assert(
        maximum(target_indices) <= num_samples,
        "Target sample indexes must be no larger than the total number of samples $(num_samples)."
    )
end

function initialize_batch_solve(
    save_path::String,
    nn::NeuralNet,
    pp::MIPVerify.PerturbationFamily,
    norm_order::Real,
    tolerance::Real,
)::Tuple{String,String,String,DataFrames.DataFrame}

    results_dir = "run_results"
    summary_file_name = "summary.csv"

    batch_run_parameters = MIPVerify.BatchRunParameters(nn, pp, norm_order, tolerance)
    main_path = joinpath(save_path, batch_run_parameters |> string)

    main_path |> mkpath_if_not_present
    joinpath(main_path, results_dir) |> mkpath_if_not_present

    summary_file_path = joinpath(main_path, summary_file_name)
    summary_file_path |> create_summary_file_if_not_present

    dt = CSV.read(summary_file_path)
    return (results_dir, main_path, summary_file_path, dt)
end

function save_to_disk(
    sample_number::Integer,
    main_path::String,
    results_dir::String,
    summary_file_path::String,
    d::Dict,
    solve_if_predicted_in_targeted::Bool,
)
    if !(d[:PredictedIndex] in d[:TargetIndexes]) || solve_if_predicted_in_targeted
        r = extract_results_for_save(d)
        results_file_uuid = get_uuid()
        results_file_relative_path = joinpath(results_dir, "$(results_file_uuid).mat")
        results_file_path = joinpath(main_path, results_file_relative_path)

        matwrite(results_file_path, Dict(ascii(string(k)) => v for (k, v) in r))
        summary_line = generate_csv_summary_line(sample_number, results_file_relative_path, r)
    else
        summary_line = generate_csv_summary_line_optimal(sample_number, d)
    end

    open(summary_file_path, "a") do file
        writedlm(file, [summary_line], ',')
    end
end

"""
$(SIGNATURES)

Determines whether to run a solve on a sample depending on the `solve_rerun_option` by
looking up information on the most recent completed solve recorded in `summary_dt`
matching `sample_number`.

`summary_dt` is expected to be a `DataFrame` with columns `:SampleNumber`, `:SolveStatus`,
and `:ObjectiveValue`.

Behavior for different choices of `solve_rerun_option`:
+ `never`: `true` if and only if there is no previous completed solve.
+ `always`: `true` always.
+ `resolve_ambiguous_cases`: `true` if there is no previous completed solve, or if the
    most recent completed solve a) did not find a counter-example BUT b) the optimization
    was not demosntrated to be infeasible.
+ `refine_insecure_cases`: `true` if there is no previous completed solve, or if the most
    recent complete solve a) did find a counter-example BUT b) we did not reach a
    provably optimal solution.
"""
function run_on_sample_for_untargeted_attack(
    sample_number::Integer,
    summary_dt::DataFrame,
    solve_rerun_option::MIPVerify.SolveRerunOption,
)::Bool
    previous_solves = filter(row -> row[:SampleNumber] == sample_number, summary_dt)
    if size(previous_solves)[1] == 0
        return true
    end
    # We now know that previous_solves has at least one element.

    if solve_rerun_option == MIPVerify.never
        return !(sample_number in summary_dt[!, :SampleNumber])
    elseif solve_rerun_option == MIPVerify.always
        return true
    elseif solve_rerun_option == MIPVerify.resolve_ambiguous_cases
        last_solve_status = previous_solves[end, :SolveStatus]
        last_objective_value = previous_solves[end, :ObjectiveValue]
        return (last_solve_status == "UserLimit") && (last_objective_value |> isnan)
    elseif solve_rerun_option == MIPVerify.refine_insecure_cases
        last_solve_status = previous_solves[end, :SolveStatus]
        last_objective_value = previous_solves[end, :ObjectiveValue]
        return !(last_solve_status == "Optimal") && !(last_objective_value |> isnan)
    else
        throw(DomainError("SolveRerunOption $(solve_rerun_option) unknown."))
    end
end

"""
$(SIGNATURES)

Runs [`find_adversarial_example`](@ref) for the specified neural network `nn` and `dataset`
for samples identified by the `target_indices`, with the target labels for each sample set
to the complement of the true label.

It creates a named directory in `save_path`, with the name summarizing
  1) the name of the network in `nn`,
  2) the perturbation family `pp`,
  3) the `norm_order`
  4) the `tolerance`.

Within this directory, a summary of all the results is stored in `summary.csv`, and
results from individual runs are stored in the subfolder `run_results`.

This functioned is designed so that it can be interrupted and restarted cleanly; it relies
on the `summary.csv` file to determine what the results of previous runs are (so modifying
this file manually can lead to unexpected behavior.)

If the summary file already contains a result for a given target index, the
`solve_rerun_option` determines whether we rerun [`find_adversarial_example`](@ref) for this
particular index.

`main_solver` specifies the solver used to solve the MIP problem once it has been built.

# Named Arguments:
+ `save_path`: Directory where results will be saved. Defaults to current directory.
+ `pp, norm_order, tolerance, rebuild, tightening_algorithm, tightening_solver, cache_model,
  solve_if_predicted_in_targeted` are passed
  through to [`find_adversarial_example`](@ref) and have the same default values;
  see documentation for that function for more details.
+ `solve_rerun_option::MIPVerify.SolveRerunOption`: Options are
  `never`, `always`, `resolve_ambiguous_cases`, and `refine_insecure_cases`.
  See [`run_on_sample_for_untargeted_attack`](@ref) for more details.
"""
function batch_find_untargeted_attack(
    nn::NeuralNet,
    dataset::MIPVerify.LabelledDataset,
    target_indices::AbstractArray{<:Integer},
    main_solver::MathProgBase.SolverInterface.AbstractMathProgSolver;
    save_path::String = ".",
    solve_rerun_option::MIPVerify.SolveRerunOption = MIPVerify.never,
    pp::MIPVerify.PerturbationFamily = MIPVerify.UnrestrictedPerturbationFamily(),
    norm_order::Real = 1,
    tolerance::Real = 0.0,
    rebuild = false,
    tightening_algorithm::MIPVerify.TighteningAlgorithm = DEFAULT_TIGHTENING_ALGORITHM,
    tightening_solver::MathProgBase.SolverInterface.AbstractMathProgSolver = MIPVerify.get_default_tightening_solver(
        main_solver,
    ),
    cache_model = true,
    solve_if_predicted_in_targeted = true,
    adversarial_example_objective::AdversarialExampleObjective = closest,
)::Nothing

    verify_target_indices(target_indices, dataset)
    (results_dir, main_path, summary_file_path, dt) =
        initialize_batch_solve(save_path, nn, pp, norm_order, tolerance)

    for sample_number in target_indices
        should_run = run_on_sample_for_untargeted_attack(sample_number, dt, solve_rerun_option)
        if should_run
            # TODO (vtjeng): change function signature for get_image and get_label
            Memento.info(MIPVerify.LOGGER, "Working on index $(sample_number)")
            input = MIPVerify.get_image(dataset.images, sample_number)
            true_one_indexed_label = MIPVerify.get_label(dataset.labels, sample_number) + 1
            d = find_adversarial_example(
                nn,
                input,
                true_one_indexed_label,
                main_solver,
                invert_target_selection = true,
                pp = pp,
                norm_order = norm_order,
                tolerance = tolerance,
                rebuild = rebuild,
                tightening_algorithm = tightening_algorithm,
                tightening_solver = tightening_solver,
                cache_model = cache_model,
                solve_if_predicted_in_targeted = solve_if_predicted_in_targeted,
                adversarial_example_objective = adversarial_example_objective,
            )

            save_to_disk(
                sample_number,
                main_path,
                results_dir,
                summary_file_path,
                d,
                solve_if_predicted_in_targeted,
            )
        end
    end
    return nothing
end

"""
$(SIGNATURES)

Determines whether to run a solve on a sample depending on the `solve_rerun_option` by
looking up information on the most recent completed solve recorded in `summary_dt`
matching `sample_number`.

`summary_dt` is expected to be a `DataFrame` with columns `:SampleNumber`, `:TargetIndexes`, `:SolveStatus`,
and `:ObjectiveValue`.
"""
function run_on_sample_for_targeted_attack(
    sample_number::Integer,
    target_label::Integer,
    summary_dt::DataFrame,
    solve_rerun_option::MIPVerify.SolveRerunOption,
)::Bool
    match_sample_number = summary_dt[!, :SampleNumber] .== sample_number
    match_target_label = summary_dt[!, :TargetIndexes] .== "[$(target_label)]"
    match = match_sample_number .& match_target_label
    previous_solves = summary_dt[match, :]
    if size(previous_solves)[1] == 0
        return true
    end
    # We now know that previous_solves has at least one element.

    if solve_rerun_option == MIPVerify.never
        return !(sample_number in summary_dt[!, :SampleNumber])
    elseif solve_rerun_option == MIPVerify.always
        return true
    elseif solve_rerun_option == MIPVerify.retarget_infeasible_cases
        last_solve_status = previous_solves[end, :SolveStatus]
        return (last_solve_status == "Infeasible")
    elseif solve_rerun_option == MIPVerify.resolve_ambiguous_cases
        last_solve_status = previous_solves[end, :SolveStatus]
        last_objective_value = previous_solves[end, :ObjectiveValue]
        return (last_solve_status == "UserLimit") && (last_objective_value |> isnan)
    elseif solve_rerun_option == MIPVerify.refine_insecure_cases
        last_solve_status = previous_solves[end, :SolveStatus]
        last_objective_value = previous_solves[end, :ObjectiveValue]
        return !(last_solve_status == "Optimal") && !(last_objective_value |> isnan)
    else
        throw(DomainError("SolveRerunOption $(solve_rerun_option) unknown."))
    end
end

"""
$(SIGNATURES)

Runs [`find_adversarial_example`](@ref) for the specified neural network `nn` and `dataset`
for samples identified by the `target_indices`, with each of the target labels in `target_labels`
individually targeted.

Otherwise same parameters as [`batch_find_untargeted_attack`](@ref).
"""
function batch_find_targeted_attack(
    nn::NeuralNet,
    dataset::MIPVerify.LabelledDataset,
    target_indices::AbstractArray{<:Integer},
    main_solver::MathProgBase.SolverInterface.AbstractMathProgSolver;
    save_path::String = ".",
    solve_rerun_option::MIPVerify.SolveRerunOption = MIPVerify.never,
    target_labels::AbstractArray{<:Integer} = [],
    pp::MIPVerify.PerturbationFamily = MIPVerify.UnrestrictedPerturbationFamily(),
    norm_order::Real = 1,
    tolerance::Real = 0.0,
    rebuild = false,
    tightening_algorithm::MIPVerify.TighteningAlgorithm = DEFAULT_TIGHTENING_ALGORITHM,
    tightening_solver::MathProgBase.SolverInterface.AbstractMathProgSolver = MIPVerify.get_default_tightening_solver(
        main_solver,
    ),
    cache_model = true,
    solve_if_predicted_in_targeted = true,
)::Nothing
    results_dir = "run_results"
    summary_file_name = "summary.csv"

    verify_target_indices(target_indices, dataset)
    (results_dir, main_path, summary_file_path, dt) =
        initialize_batch_solve(save_path, nn, pp, norm_order, tolerance)

    for sample_number in target_indices
        for target_label in target_labels
            should_run = run_on_sample_for_targeted_attack(
                sample_number,
                target_label,
                dt,
                solve_rerun_option,
            )
            if should_run
                input = MIPVerify.get_image(dataset.images, sample_number)
                true_one_indexed_label = MIPVerify.get_label(dataset.labels, sample_number) + 1
                if true_one_indexed_label == target_label
                    continue
                end

                Memento.info(
                    MIPVerify.LOGGER,
                    "Working on index $(sample_number), with true_label $(true_one_indexed_label) and target_label $(target_label)",
                )

                d = find_adversarial_example(
                    nn,
                    input,
                    target_label,
                    main_solver,
                    invert_target_selection = false,
                    pp = pp,
                    norm_order = norm_order,
                    tolerance = tolerance,
                    rebuild = rebuild,
                    tightening_algorithm = tightening_algorithm,
                    tightening_solver = tightening_solver,
                    cache_model = cache_model,
                    solve_if_predicted_in_targeted = solve_if_predicted_in_targeted,
                )

                save_to_disk(
                    sample_number,
                    main_path,
                    results_dir,
                    summary_file_path,
                    d,
                    solve_if_predicted_in_targeted,
                )
            end
        end
    end
    return nothing
end

function batch_build_model(
    nn::NeuralNet,
    dataset::MIPVerify.LabelledDataset,
    target_indices::AbstractArray{<:Integer},
    tightening_solver::MathProgBase.SolverInterface.AbstractMathProgSolver;
    pp::MIPVerify.PerturbationFamily = MIPVerify.UnrestrictedPerturbationFamily(),
    tightening_algorithm::MIPVerify.TighteningAlgorithm = DEFAULT_TIGHTENING_ALGORITHM,
)::Nothing

    verify_target_indices(target_indices, dataset)

    for sample_number in target_indices
        Memento.info(MIPVerify.LOGGER, "Working on index $(sample_number)")
        input = MIPVerify.get_image(dataset.images, sample_number)
        build_reusable_model_uncached(nn, input, pp, tightening_solver, tightening_algorithm)
    end
    return nothing
end


text:
using LinearAlgebra
using Random
using StatsFuns: softplus
using Roots # for inverse

################################################################################
#                            Planar and Radial Flows                           #
#             Ref: Variational Inference with Normalizing Flows,               #
#               D. Rezende, S. Mohamed(2015) arXiv:1505.05770                  #
################################################################################

###############
# PlanarLayer #
###############

# TODO: add docstring

# FIXME: using `TrackedArray` for the parameters, we end up with
# nested tracked structures; don't want this.
mutable struct PlanarLayer{T1,T2} <: Bijector{1}
    w::T1
    u::T1
    b::T2
end

function get_u_hat(u, w)
    # To preserve invertibility
    return (
        u + (planar_flow_m(w' * u) - w' * u)[1]
        * w / (norm(w[:,1],2) ^ 2)
    )   # from A.1
end

function PlanarLayer(dims::Int, container=Array)
    w = container(randn(dims, 1))
    u = container(randn(dims, 1))
    b = container(randn(1))
    return PlanarLayer(w, u, b)
end

planar_flow_m(x) = -1 .+ softplus.(x)   # for planar flow from A.1
dtanh(x) = 1 .- (tanh.(x)) .^ 2         # for planar flow
ψ(z, w, b) = dtanh(w' * z .+ b) .* w    # for planar flow from eq(11)

# An internal version of transform that returns intermediate variables
function _transform(flow::PlanarLayer, z)
    u_hat = get_u_hat(flow.u, flow.w)
    transformed = z + u_hat * tanh.(flow.w' * z .+ flow.b) # from eq(10)
    return (transformed=transformed, u_hat=u_hat)
end

(b::PlanarLayer)(z) = _transform(b, z).transformed

function _forward(flow::PlanarLayer, z)
    transformed, u_hat = _transform(flow, z)
    # Compute log_det_jacobian
    psi = ψ(z, flow.w, flow.b)
    log_det_jacobian = log.(abs.(1.0 .+ psi' * u_hat))          # from eq(12)
    return (rv=transformed, logabsdetjac=vec(log_det_jacobian)) # from eq(10)
end

forward(flow::PlanarLayer, z::AbstractMatrix) = _forward(flow, z)

function forward(flow::PlanarLayer, z::AbstractVector{<: Real})
    res = _forward(flow, z)
    return (rv=vec(res.rv), logabsdetjac=res.logabsdetjac[1])
end

function (ib::Inversed{<: PlanarLayer})(y::AbstractMatrix{<: Real})
    flow = ib.orig
    u_hat = get_u_hat(flow.u, flow.w)
    # Define the objective functional; implemented with reference from A.1
    f(y) = alpha -> (flow.w' * y)[1] - alpha - (flow.w' * u_hat)[1] * tanh(alpha+flow.b[1])
    # Run solver
    alphas_ = [find_zero(f(y[:,i:i]), 0.0, Order16()) for i in 1:size(y, 2)]
    alphas = alphas_'
    z_para = (flow.w ./ norm(flow.w,2)) * alphas
    z_per = y - z_para - u_hat * tanh.(flow.w' * z_para .+ flow.b)

    return z_para + z_per
end

function (ib::Inversed{<: PlanarLayer})(y::AbstractVector{<: Real})
    return vec(ib(reshape(y, (length(y), 1))))
end

logabsdetjac(flow::PlanarLayer, x) = forward(flow, x).logabsdetjac
isclosedform(b::Inversed{<:PlanarLayer}) = false


text:
"Describes a component of growth media"
struct MediaComponent
    "metabolite index in a COBRA model"
    metabolite::Int
    "index of corresponding exchange reaction"
    exchange_reaction::Int
    "Maximum uptake rate"
    V::Float64
    "concentration of the metabolite in the media"
    concentration::Float64

    function MediaComponent(i::Integer, ex::Integer, V::Real, c::Real)
        @assert i > 0 && ex > 0 && V ≥ 0 && c ≥ 0
        new(i, ex, V, c)
    end
end

function MediaComponent(cobra_model::COBRA.LPproblem, metabolite::Integer, V::Real, c::Real)
    rxns = exchange_reaction(cobra_model, metabolite)
    @assert length(rxns) == 1
    MediaComponent(metabolite, first(rxns), V, c)
end

"A growth media is a collection of media components."
Media = Vector{MediaComponent}

"""
Sets uptake bounds: u ≤ min(V, c/ξ)
"""
function set_bounds_ξ(cobra_model::COBRA.LPproblem, media::Media, ξ::Real)
    @assert 0 ≤ ξ < Inf
    for x in media
        if x.c == 0
            cobra_model.ub[x.exchange_reaction] = 0
        elseif ξ == 0
            cobra_model.ub[x.exchange_reaction] = x.V
        else
            cobra_model.ub[x.exchange_reaction] = min(x.V, x.c / ξ)
        end
    end
end


text:
"""
    getMonthNum(t, f)

Returns month index in simulation
"""
function getMonthNum(t::Int, f::Dict)
    return 1 + (Dates.month(f[:timestamp][t])
        - Dates.month(f[:timestamp][1])
        + 12 * (Dates.year(f[:timestamp][t])
            - Dates.year(f[:timestamp][1])))
end

"""
    findDate(date, dates)

Returns index of first occurence of date in dates.
"""
function findDate(date::Date,
                  dates::Array{Date,1})
    candidates = find(x->x==date, dates)
    if length(candidates) == 0
        error("Could not find date " * string(date) * " in dataset")
    end
    if length(candidates) > 24
        error("More than 24 timestamps for date " * string(date) * "...")
    end
    if length(candidates) < 24
        warn("Less than 24 timestamps for date " * string(date) * " - OK?")
    end
    return candidates[1]
end

function save_data(sesi::JuMP.Model, f::Dict, p::Dict,
    folderNm::String)

    T = f[:T];

    # write results to csv
    output = DataFrame()

    for col in [:gHeater, :eHeater, :hwHeater, :eHRC,
        :hwHRC, :cwHRC, :eChiller, :cwChiller,
        :eGrid, :gGrid, :umHot, :umCold]
        output[col] = getvalue(sesi[col])
    end
    output[:elecCEF] = (output[:eHRC] + output[:eHeater]
                        + output[:eChiller])
    output[:cwStorage] = getvalue(sesi[:cwTank])[1:T-1];
    output[:hwStorage] = getvalue(sesi[:hwTank])[1:T-1];
    #output[:elecSolar] = getvalue(getvariable(sesi, :elecSolar));

    # Go back to UTC for storage
    output[:Date] = f[:timestamp] - Dates.Hour(p[:utc_offset])

    output[:elecCampus] = f[:loadKWH]
    output[:cwLoads] = f[:loadCW]
    output[:hwLoads] = f[:loadHW]
    output[:elecPrice] = f[:elecPrice]

    # rename some variables
    output[:hwHeater] = output[:hwHeater]
    output[:cwChiller] = output[:cwChiller]
    output[:cwHRC] = output[:cwHRC]
    output[:hwHRC] = output[:hwHRC]

    output[:elecGrid] = output[:eGrid]
    output[:gasGrid] = output[:gGrid]
    output[:elecHRC] = output[:eHRC]
    output[:elecChiller] = output[:eChiller]
    output[:elecHeater] = output[:eHeater]

    # make sure we have the write order for output
    cols = [:Date, :elecCEF,:elecCampus,:cwLoads, :hwLoads, :hwStorage,
    :cwStorage, :gHeater, :elecHeater, :hwHeater, :elecHRC, :hwHRC, :cwHRC,
    :elecChiller, :cwChiller, :elecPrice, :elecGrid, :gasGrid, :umHot, :umCold]

    if p[:ev]
        output[:evSOC] = sum(getvalue(sesi[:LevelOfCharge])[:,1:T-1])
        output[:elecEV_in] = getvalue(sesi[:elecEV_in])
        output[:elecEV_out] = getvalue(sesi[:elecEV_out])
        append!(cols, [:evSOC, :elecEV_in, :elecEV_out])
    end

    mkpath(folderNm)
    CSV.write(joinpath(folderNm, "data.csv"), output[cols])

    # write params and forecast to json file
    out_f = f
     # convert to UTC for storage
    if p[:cbp]
        out_f[:timestamp_cbp] = [ts - Dates.Hour(p[:utc_offset])
            for ts in out_f[:timestamp_cbp]]
        out_f[:cbp_bid] = getvalue(sesi[:cbp_bid]) # also store this here
    end
    forecast_json = json(out_f);
    open(joinpath(folderNm, "forecast.json"), "w") do fw
        write(fw, forecast_json)
    end
    params_json = json(p);
    open(joinpath(folderNm, "params.json"), "w") do fw
        write(fw, params_json)
    end
end

"""
    campus_bill(sesi, f, verb)

Calculate campus bill, after an optimization run. This is usually (slightly)
different from the optimization objective (that could include e.g. optimization
constraints). Note that this function also serves as a sanity check - it is
calculated after the optimization, just based on the resulting schedule
"""

function campus_bill(sesi::JuMP.Model, f::Dict, p::Dict; verb::Int=0)

    T = f[:T]
    M = f[:M]
    # startT = f[:startT]
    # endT = f[:endT]
    eGrid = getvalue(sesi[:eGrid])[1:T-1]
    gGrid = getvalue(sesi[:gGrid])[1:T-1]
    umCold = getvalue(sesi[:umCold])[1:T-1]
    umHot = getvalue(sesi[:umHot])[1:T-1]

    elecCEP = 0.
    for var in [:eChiller, :eHeater, :eHRC]
        if haskey(sesi.objDict, var)
            elecCEP = elecCEP + getvalue(sesi[var])[1:T-1]
        end
    end

    elecSolar = 0.;
    if haskey(sesi.objDict, :elecSolar)
        elecSolar = getvalue(sesi[:elecSolar])[1:T-1]
    end

    if p[:sw]
        swHRC = getvalue(sesi[:swHRC])
        swChiller = getvalue(sesi[:swChiller])
        swHeater = getvalue(sesi[:swHeater])
    end

    if p[:ev]
        evDischarge = getvalue(sesi[:evDischarge])
    end

    # bill dictionary contains detailed bill
    bill = Dict();

    bill["objective"] = 0.# recreate objective to make debugging easier

    # note: need to recalculate peak variables "by hand" because they will not
    # always exist in the JuMP model
    peak = zeros(M);
    for t = 1:T-1
        m = getMonthNum(t, f)
        peak[m] = max(peak[m], eGrid[t])
    end

    bill["demandCharge"] = sum(peak .* f[:demandCharge])
    bill["eEnergyCost"] = sum(eGrid .* f[:elecPrice])
    bill["gEnergyCost"] = sum(gGrid .* f[:gasPrice])

    bill["objective"] = bill["gEnergyCost"] + bill["eEnergyCost"]
    if p[:peak]
        bill["objective"] += bill["demandCharge"]
    end

    bill["agg_dollar"] = (bill["eEnergyCost"] + bill["demandCharge"]
                            + bill["gEnergyCost"])

    if haskey(f, :carbonIntensityGrid)
        bill["agg_carbon"] = (sum(eGrid .* f[:carbonIntensityGrid])
            + sum(gGrid * f[:carbonIntensityGas]))
        #+ sum(elecSolar * f[:carbonIntensitySolar])
        bill["CEP_carbon"] = (sum(elecCEP .* f[:carbonIntensityGrid])
        + sum(gGrid * f[:carbonIntensityGas]))
    end

    if p[:carbon]
        bill["objective"] += f[:carbonPrice] * bill["agg_carbon"]
    end

    if p[:cbp]
        cbp_bid = getvalue(sesi[:cbp_bid])
        delivery_ratio = 1 # for now
        bill["cbp"] = - sum(f[:cbp_price][m] * delivery_ratio *
                    cbp_bid[m] for m in 1:M)
        bill["unmet"] = sum(getvalue(sesi[:cbp_unmet]))*f[:cbp_unmet_penalty]
        bill["agg_dollar"] += bill["cbp"]
        bill["objective"] += bill["cbp"]
	    bill["objective"] += bill["unmet"]
    end

    if p[:sw]
        bill["swCost"] = f[:swPenalty] * (swHRC'swHRC + swChiller'swChiller
            + swHeater'swHeater)
        bill["objective"] += bill["swCost"]
    end

    if p[:ev]
        bill["ev"] = sum(evDischarge * p[:usageCostBatteryEV])
        bill["objective"] += bill["ev"]
    end

    elecRateInternal = (sum(f[:elecPrice] .* eGrid)
        + sum(peak .* f[:demandCharge]))/sum(eGrid)
    if p[:cbp]
        elecRateInternal += bill["cbp"]/sum(eGrid)
    end
    bill["elecRateInternal"] = elecRateInternal
    bill["CEP_dollar"] = (elecRateInternal * sum(elecCEP)
        + sum(gGrid .* f[:gasPrice]));

    # unmet loads
    if p[:umLoadsL2]
        bill["umLoads"] = p[:umHot]*umHot'umHot+p[:umCold]*umCold'umCold
    else
        bill["umLoads"] = p[:umHot] * sum(umHot) + p[:umCold] * sum(umCold)
    end

    bill["objective"] += bill["umLoads"]

    # sanity
    if abs(bill["objective"]-getvalue(getobjective(sesi)))>1e-3
        warn("Objective is not properly reconstructed - Debugging this now")
        sw = p[:swPenalty] * (swHRC'swHRC + swChiller'swChiller
            + swHeater'swHeater)
        dc = sum(getvalue(sesi[:peak][m]) * f[:demandCharge][m] for m = 1:M)
        carbon = 0.
        if p[:carbon]
            carbon = f[:carbonPrice] * (
                sum(f[:carbonIntensityGrid][t] * getvalue(sesi[:eGrid][t])
                    for t = 1:T-1)
                + f[:carbonIntensityGas] * sum(getvalue(sesi[:gGrid][t])
                    for t = 1:T-1))
        end
        energy = sum(getvalue(sesi[:eGrid][t]) * f[:elecPrice][t]
            + getvalue(sesi[:gGrid][t]) * f[:gasPrice] for t = 1:T-1)
        cbp = 0.
        cbp_unmet = 0.
        if p[:cbp]
            deliv_ratio = 1.
            cbp = - sum(f[:cbp_price][m] * deliv_ratio
                * getvalue(sesi[:cbp_bid][m]) for m in 1:M)
            cbp_unmet = sum(getvalue(sesi[:cbp_unmet]))*p[:cbp_unmet_penalty]
        end

        if p[:umLoadsL2]
            umLoads = (p[:umHot]*getvalue(sesi[:umHot])'getvalue(sesi[:umHot])
                       +p[:umCold]*getvalue(sesi[:umCold])'getvalue(sesi[:umCold]))
       else
            umLoads = (p[:umHot] * sum(getvalue(sesi[:umHot]))
                       + p[:umCold] * sum(getvalue(sesi[:umCold])))
        end
        obj = sw + dc + carbon + energy + cbp + cbp_unmet + umLoads

        println("Switching")
        println(abs(sw - bill["swCost"]))
        println("Demand Charge")
        println(abs(dc - bill["demandCharge"]))
        if p[:carbon]
            println("carbon cost")
            println(abs(carbon - f[:carbonPrice] * bill["agg_carbon"]))
        end
        println("energy bill")
        println(abs(energy - bill["eEnergyCost"] - bill["gEnergyCost"]))
        if haskey(sesi.objDict, :cbp_bid)
            println("cbp")
            println(abs(cbp - bill["cbp"]))
            println("cbp_unmet")
            println(abs(cbp_unmet - bill["unmet"]))
        end
        println("objective")
        println(abs(bill["objective"]-getvalue(getobjective(sesi))))
    end

    # verbose output
    if verb > 0
        @printf("Printing hourly equivalent rates\n")
        @printf("elec demand:\t\t%.4f \$/kWh\n",
            bill["demandCharge"]/sum(eGrid))
        @printf("electricity:\t\t%.4f \$/kWh\n",
            bill["eEnergyCost"]/sum(eGrid))
        if haskey(sesi.objDict, :cbp_bid)
            @printf("cbp:\t\t\t%.4f \$/kWh\n", bill["cbp"]/sum(eGrid))
        end
        @printf("elec internal:\t\t%.6f \$/kWh\n", bill["elecRateInternal"])

        @printf("gas:\t\t\t%.4f\$/mmbtu\n", bill["gEnergyCost"]/sum(gGrid))
    end
    if verb > 1
        @printf("Printing aggregate costs:\n")
        @printf("elec demand:\t\t%.4f \$\n",
            bill["demandCharge"])
        @printf("electricity:\t\t%.4f \$\n",
            bill["eEnergyCost"])
        if haskey(sesi.objDict, :cbp_bid)
            @printf("cbp:\t\t\t%.4f \$\n", bill["cbp"])
        end
        @printf("gas:\t\t\t%.4f\$\n", bill["gEnergyCost"])
        @printf("total:\t\t\t%.4f\$\n", bill["agg_dollar"])
        @printf("obj:\t\t\t%.4f\$\n", bill["objective"])
    end
    if verb > 2
        print("Energy Cost (M\$): ")
        println(bill["agg_dollar"]/1e6)
        print("Energy Cost CEP (M\$): ")
        println(bill["CEP_dollar"]/1e6)
        @printf("elec rate internal: %.2f\n", bill["elecRateInternal"])
        if haskey(bill, "agg_carbon")
            print("Carbon (1,000 tons): ")
            println(bill["agg_carbon"]/1e3)
            print("Carbon for CEP (1,000 tons): ")
            println(bill["CEP_carbon"]/1e3)
        end
        print("Internal rate for electricity(\$): ")
        println(elecRateInternal)
    end

    if verb > 3
        print("Peaks(MW): ")
        println(peak/1e3)
        print("Demand charge (M\$): ")
        println(sum(peak .* f[:demandCharge])/1e6)
        print("Electricity (M\$): ")
        println(sum(f[:elecPrice] .* eGrid)/1e6)
        print("Gas (1,000\$): ")
        println(sum(f[:gasPrice] .* gGrid)/1e3)
    end

    return bill
end

function compare_bills(bbill::Dict, nbill::Dict)
    comp = Dict()
    for it in unique(vcat(collect(keys(bbill)), collect(keys(nbill))))
        bb = 0.
        nb = 0.
        if haskey(bbill, it)
            bb = bbill[it]
        end
        if haskey(nbill, it)
            nb = nbill[it]
        end
        comp[it] = bb - nb
    end
    comp
end

max_stars_repo_name: jdechalendar/sesi

max_issues_repo_name: jdechalendar/sesi

max_forks_repo_name: jdechalendar/sesi


text:
@deprecate ordered isordered

@deprecate compact compress
@deprecate uncompact decompress

@deprecate CategoricalArray(::Type{T}, dims::NTuple{N,Int}; ordered=false) where {T, N} CategoricalArray{T}(dims, ordered=ordered)
@deprecate CategoricalArray(::Type{T}, dims::Int...; ordered=false) where {T} CategoricalArray{T}(dims, ordered=ordered)

@deprecate CategoricalVector(::Type{T}, m::Integer; ordered=false) where {T} CategoricalVector{T}(m, ordered=ordered)

@deprecate CategoricalMatrix(::Type{T}, m::Int, n::Int; ordered=false) where {T} CategoricalMatrix{T}(m, n, ordered=ordered)

@deprecate CategoricalArray(dims::Int...; ordered=false) CategoricalArray(undef, dims...; ordered=ordered)

@deprecate CategoricalArray{T, N, R}(dims::NTuple{N,Int}; ordered=false) where {T, N, R} CategoricalArray{T, N, R}(undef, dims; ordered=ordered)

@deprecate CategoricalArray{T, N}(dims::NTuple{N,Int}; ordered=false) where {T, N} CategoricalArray{T, N}(undef, dims; ordered=ordered)
@deprecate CategoricalArray{T}(dims::NTuple{N,Int}; ordered=false) where {T, N} CategoricalArray{T}(undef, dims; ordered=ordered)
@deprecate CategoricalArray{T, 1}(m::Int; ordered=false) where {T} CategoricalArray{T, 1}(undef, m; ordered=ordered)
@deprecate CategoricalArray{T, 2}(m::Int, n::Int; ordered=false) where {T} CategoricalArray{T, 2}(undef, m, n; ordered=ordered)
@deprecate CategoricalArray{T, 1, R}(m::Int; ordered=false) where {T, R} CategoricalArray{T, 1, R}(undef, m; ordered=ordered)
# R <: Integer is required to prevent default constructor from being called instead
@deprecate CategoricalArray{T, 2, R}(m::Int, n::Int; ordered=false) where {T, R <: Integer} CategoricalArray{T, 2, R}(undef, m, n; ordered=ordered)
@deprecate CategoricalArray{T, 3, R}(m::Int, n::Int, o::Int; ordered=false) where {T, R} CategoricalArray{T, 3, R}(undef, m, n, o; ordered=ordered)
@deprecate CategoricalArray{T}(m::Int; ordered=false) where {T} CategoricalArray{T}(undef, m; ordered=ordered)
@deprecate CategoricalArray{T}(m::Int, n::Int; ordered=false) where {T} CategoricalArray{T}(undef, m, n; ordered=ordered)
@deprecate CategoricalArray{T}(m::Int, n::Int, o::Int; ordered=false) where {T} CategoricalArray{T}(undef, m, n, o; ordered=ordered)

@deprecate CategoricalVector(m::Integer; ordered=false) CategoricalVector(undef, m; ordered=ordered)
@deprecate CategoricalVector{T}(m::Int; ordered=false) where {T} CategoricalVector{T}(undef, m; ordered=ordered)

@deprecate CategoricalMatrix(m::Int, n::Int; ordered=false) CategoricalMatrix(undef, m, n; ordered=ordered)
@deprecate CategoricalMatrix{T}(m::Int, n::Int; ordered=false) where {T} CategoricalMatrix{T}(undef, m::Int, n::Int; ordered=ordered)

Base.@deprecate_binding CategoricalString CategoricalValue{String, R} where R
Base.@deprecate_binding CatValue Union{CategoricalValue{<:Any, R}, CategoricalValue{String, R}} where R

@deprecate iscatvalue(x::Type) x <: CategoricalValue && x !== Union{}
@deprecate iscatvalue(x::Any) typeof(x) <: CategoricalValue

@deprecate(catvaluetype(::Type{T}, ::Type{R}) where {T >: Missing, R},
    CategoricalValue{CategoricalArrays.leveltype(Base.nonmissingtype(T)), R})
@deprecate(catvaluetype(::Type{<:CategoricalValue{T}}, ::Type{R}) where {T >: Missing, R},
    CategoricalValue{CategoricalArrays.leveltype(Base.nonmissingtype(T)), R})

# AbstractString interface for CategoricalString
import Base: eltype, length, lastindex, sizeof, nextind, prevind,
    iterate, getindex, codeunit, ascii, isvalid, match, collect, reverse,
    ncodeunits, isempty, firstindex, lpad, rpad, occursin, startswith, endswith,
    repeat, split, rsplit, strip, lstrip, rstrip, eachmatch,
    uppercase, lowercase, titlecase, uppercasefirst, lowercasefirst,
    chop, chomp, escape_string, textwidth, isascii, *, ^, findfirst, findlast, keys, replace
# NOTE: drop dependency on Unicode when removing these deprecations
import Unicode: normalize, graphemes
@deprecate eltype(x::CategoricalValue{String}) eltype(get(x))
@deprecate eltype(T::Type{<:CategoricalValue{String}}) eltype(CategoricalArrays.leveltype(T))
@deprecate length(x::CategoricalValue{String}) length(String(x))
@deprecate lastindex(x::CategoricalValue{String}) lastindex(String(x))
@deprecate sizeof(x::CategoricalValue{String}) sizeof(String(x))
@deprecate nextind(x::CategoricalValue{String}, i::Int, j::Int=1) nextind(String(x), i, j)
@deprecate prevind(x::CategoricalValue{String}, i::Int, j::Int=1) prevind(String(x), i, j)
@deprecate iterate(x::CategoricalValue{String}) iterate(String(x))
@deprecate iterate(x::CategoricalValue{String}, i::Int) iterate(String(x), i)
@deprecate getindex(x::CategoricalValue{String}, i::Int) getindex(String(x), i)
@deprecate codeunit(x::CategoricalValue{String}, i::Integer) codeunit(String(x), i)
@deprecate ascii(x::CategoricalValue{String}) ascii(String(x))
@deprecate isvalid(x::CategoricalValue{String}) isvalid(String(x))
@deprecate isvalid(x::CategoricalValue{String}, i::Integer) isvalid(String(x), i)
@deprecate match(r::Regex, s::CategoricalValue{String},
      idx::Integer=firstindex(s), add_opts::UInt32=UInt32(0); kwargs...) match(r, String(s), idx, add_opts; kwargs...)
@deprecate collect(x::CategoricalValue{String}) collect(String(x))
@deprecate reverse(x::CategoricalValue{String}) reverse(String(x))
@deprecate ncodeunits(x::CategoricalValue{String}) ncodeunits(String(x))
@deprecate keys(x::CategoricalValue{String}) keys(String(x))

@deprecate isempty(x::CategoricalValue{String}) isempty(String(x))
@deprecate firstindex(x::CategoricalValue{String}) firstindex(String(x))
@deprecate normalize(x::CategoricalValue{String}, s::Symbol) normalize(String(x), s) false
@deprecate normalize(x::CategoricalValue{String}; kwargs...) normalize(String(x); kwargs...) false
@deprecate graphemes(x::CategoricalValue{String}) graphemes(String(x)) false
@deprecate length(x::CategoricalValue{String}, i::Int, j::Int) length(String(x), i, j)
@deprecate repeat(x::CategoricalValue{String}, i::Integer) repeat(String(x), i)
@deprecate eachmatch(r::Regex, x::CategoricalValue{String}; overlap=false) eachmatch(r, String(x), overlap=overlap)
@deprecate lpad(x::CategoricalValue{String}, n::Integer, c::Union{AbstractChar,AbstractString}=' ') lpad(String(x), n)
@deprecate rpad(x::CategoricalValue{String}, n::Integer, c::Union{AbstractChar,AbstractString}=' ') rpad(String(x), n)
@deprecate occursin(x::CategoricalValue{String}, y::AbstractString) occursin(String(x), y)
@deprecate occursin(x::AbstractString, y::CategoricalValue{String}) occursin(x, String(y))
@deprecate occursin(x::Regex, y::CategoricalValue{String}) occursin(x, String(y))
@deprecate occursin(x::AbstractChar, y::CategoricalValue{String}) occursin(x, String(y))
@deprecate startswith(x::CategoricalValue{String}, y::AbstractString) startswith(String(x), y)
@deprecate startswith(x::AbstractString, y::CategoricalValue{String}) startswith(x, String(y))
@deprecate endswith(x::CategoricalValue{String}, y::AbstractString) endswith(String(x), y)
@deprecate endswith(x::AbstractString, y::CategoricalValue{String}) endswith(x, String(y))
@deprecate split(x::CategoricalValue{String}; kwargs...) split(String(x); kwargs...)
@deprecate split(x::CategoricalValue{String}, dlm; kwargs...) split(String(x), dlm; kwargs...)
@deprecate rsplit(x::CategoricalValue{String},; kwargs...) rsplit(String(x); kwargs...)
@deprecate rsplit(x::CategoricalValue{String}, dlm; kwargs...) rsplit(String(x), dlm; kwargs...)
@deprecate strip(x::CategoricalValue{String}) strip(String(x))
@deprecate lstrip(x::CategoricalValue{String}) lstrip(String(x))
@deprecate rstrip(x::CategoricalValue{String}) rstrip(String(x))
@deprecate lstrip(x::CategoricalValue{String}, chars::Base.Chars) lstrip(String(x), chars)
@deprecate rstrip(x::CategoricalValue{String}, chars::Base.Chars) rstrip(String(x), chars)
@deprecate strip(x::CategoricalValue{String}, chars::Base.Chars) strip(String(x), chars)
@deprecate strip(pred, x::CategoricalValue{String}) strip(pred, String(x))
@deprecate lstrip(pred, x::CategoricalValue{String}) lstrip(pred, String(x))
@deprecate rstrip(pred, x::CategoricalValue{String}) rstrip(pred, String(x))
@deprecate lowercase(x::CategoricalValue{String}) lowercase(String(x))
@deprecate uppercase(x::CategoricalValue{String}) uppercase(String(x))
@deprecate lowercasefirst(x::CategoricalValue{String}) lowercasefirst(String(x))
@deprecate uppercasefirst(x::CategoricalValue{String}) uppercasefirst(String(x))
@deprecate titlecase(x::CategoricalValue{String}) titlecase(String(x))
@deprecate chop(x::CategoricalValue{String}; kwargs...) chop(String(x); kwargs...)
@deprecate chomp(x::CategoricalValue{String}) chomp(String(x))
@deprecate textwidth(x::CategoricalValue{String}) textwidth(String(x))
@deprecate isascii(x::CategoricalValue{String}) isascii(String(x))
@deprecate escape_string(x::CategoricalValue{String}) escape_string(String(x))
@deprecate *(x::CategoricalValue{String}, y::AbstractString) String(x) * y
@deprecate *(x::AbstractString, y::CategoricalValue{String}) x * String(y)
@deprecate *(x::CategoricalValue{String}, y::CategoricalValue{String}) String(x) * String(y)
@deprecate ^(x::CategoricalValue{String}, n::Integer) String(x)^n
@deprecate findfirst(needle::AbstractString, haystack::CategoricalValue{String}) findfirst(needle, String(haystack))
@deprecate findlast(needle::AbstractString, haystack::CategoricalValue{String}) findlast(needle, String(haystack))
@deprecate findfirst(needle::Function, haystack::CategoricalValue{String}) findfirst(needle, String(haystack))
@deprecate findlast(needle::Function, haystack::CategoricalValue{String}) findlast(needle, String(haystack))
@deprecate findfirst(needle::Base.Fix2, haystack::CategoricalValue{String}) findfirst(needle, String(haystack))
@deprecate findlast(needle::Base.Fix2, haystack::CategoricalValue{String}) findlast(needle, String(haystack))
@deprecate replace(x::CategoricalValue{String}, old_new::Pair...; kwargs...) replace(String(x), old_new...; kwargs...)


text:
using WRS
using Test

@testset "WRS" begin

    include("test_pb2gen.jl")

end


text:
using HDF5,JLD,KUparser,KUnet,Base.Test

info("Load data")
@date @load "conll07.tst.jld4"
ncpu = 10
nbatch = 10
ndeps = length(deprel)
pt = ArcEager13
p0 = pt(1,ndeps)
s1 = corpus[1]

ft = Array[["s0w"], ["s0p"], ["s0w", "s0p"]]
@date (p1,x1,y1)=oparse(pt, s1, ndeps, ft)

SFarray = Array(Any, length(KUparser.SFhash))
for k in keys(KUparser.SFhash)
    v = KUparser.SFhash[k]
    SFarray[v] = k
end

for j=1:size(x1,2)
    println(SFarray[find(x1[:,j])])
end


text:
module TestsUKBB

using Test
using CSV
using DataFrames
using TMLEEpistasis
using CategoricalArrays
using TOML
using Serialization
using TMLE
using JLD2
using MLJBase

include("helper_fns.jl")

function test_base_serialization(tmle_reports, n_expected; phenotype_id=1)
    tmle_report = tmle_reports["$(phenotype_id)_1"]
    @test tmle_report isa TMLE.TMLEReport
    test_queries((tmle_report.query,),
        (Query(name="QUERY_1", case=(RSID_10 = "AG", RSID_100 = "AG"), control=(RSID_10 = "GG", RSID_100 = "GG")),)
    )
    @test size(tmle_report.influence_curve, 1) == n_expected
    @test tmle_report.estimate isa Real
    @test tmle_report.initial_estimate isa Real

    # Check second tmle_report
    tmle_report = tmle_reports["$(phenotype_id)_2"]
    @test tmle_report isa TMLE.TMLEReport
    test_queries((tmle_report.query,),
        (Query(name="QUERY_2", case=(RSID_10 = "AG", RSID_100 = "AA"), control=(RSID_10 = "GG", RSID_100 = "GG")),)
    )
    @test size(tmle_report.influence_curve, 1) == n_expected
    @test tmle_report.estimate isa Real
    @test tmle_report.initial_estimate isa Real
end

@testset "Test actualqueries" begin
    genotypes = DataFrame(
        rs1 = ["AC", "CC", "AC", "CC", missing],
        rs2 = ["CG", "CG", "GG", "GG", "GG"]
    )
    queries = [
    Query(
        case=(rs1 = "AC", rs2 = "CG"),
        control=(rs1 = "CC", rs2 = "GG"),
        name="QUERY_1"),
    Query(
        case=(rs1 = "AC", rs2 = "CG"),
        control=(rs1 = "CC", rs2 = "CC"),
        name="QUERY_2")
    ]
    @test TMLEEpistasis.genotypes_combinations(queries[1]) ==
        DataFrame(
            rs1 = ["CC", "AC", "CC", "AC"],
            rs2 = ["GG", "GG", "CG", "CG"]
        )
    @test TMLEEpistasis.genotypes_combinations(queries[2]) ==
        DataFrame(
            rs1 = ["CC", "AC", "CC", "AC"],
            rs2 = ["CC", "CC", "CG", "CG"]
        )

    actual_queries = TMLEEpistasis.actualqueries(genotypes, queries)
    @test actual_queries == queries[1:1]
end

@testset "Test preprocess" begin
    n = 10
    base_sample_ids = 1:n

    confounders = DataFrame(rand(n, 3), :auto)
    confounders.sex = [0, 1, 1, 1, 0, 0 , 0, 1, 1, 0]
    confounders.SAMPLE_ID = base_sample_ids

    genotypes = DataFrame(
        SAMPLE_ID = base_sample_ids,
        RSID_1    = ["AC", "AC", "AC", "AC", "AC", "CC", "CC", "CC", "CC", "CC"],
        RSID_2    = ["AG", "AA", "AA", "AA", "AA", "AG", "AG", "AA", "AA", "AA"]
        )

    queries = [TMLE.Query(case=(RSID_2="AA", RSID_1="CC"), control=(RSID_2="AG", RSID_1="AC"))]
    phenotypes = DataFrame(
        SAMPLE_ID = base_sample_ids,
        Y1 = rand(n),
        Y2 = rand(n)
    )
    T, W, y, sample_ids = TMLEEpistasis.preprocess(genotypes, confounders, phenotypes, Real, queries)

    # Check the order of columns has been reversed
    @test T == genotypes[!, [:RSID_2, :RSID_1]]
    @test W == confounders[!, Not("SAMPLE_ID")]
    @test all(eltype(W[!, col]) === Float64 for col in names(W))
    @test y == phenotypes[!, Not("SAMPLE_ID")]
    @test sample_ids ==
            Dict(
                "Y1" => string.(base_sample_ids),
                "Y2" => string.(base_sample_ids),
            )
    # Checking joining on sample_ids
    genotypes = DataFrame(
        SAMPLE_ID = base_sample_ids,
        RSID_1    = ["AC", "AC", missing, "AC", "AC", "CC", "CC", "CC", "CC", "CC"],
        RSID_2    = ["AG", "AG", "AA", "AA", "AA", "AG", "AG", "AA", "AA", "AA"]
        )
    phenotypes = DataFrame(
        SAMPLE_ID = base_sample_ids,
        Y1 = [0, 1, 0, 1, 0, 1, 0, 0, 1, missing],
    )
    T, W, y, sample_ids = TMLEEpistasis.preprocess(genotypes, confounders, phenotypes, Bool, queries)

    @test size(T) == (9, 2)
    # For y the missing value hasn't been dropped yet, it will be dropped during TMLE
    @test size(y) == (9, 1)
    @test size(W) == (9, 4)
    @test sample_ids == Dict("Y1" => ["1", "2", "4", "5", "6", "7", "8", "9"])
end

@testset "Test UKBBVariantRun with continuous target" begin
    # Only one continuous phenotype / machines not saved / no adaptive cv
    estimatorfile = joinpath("config", "tmle_config.toml")
    build_query_file()
    parsed_args = Dict(
        "genotypes" => genotypesfile,
        "phenotypes" => continuous_phenotypefile,
        "confounders" => confoundersfile,
        "queries" => queryfile,
        "estimator" => estimatorfile,
        "out" => "RSID_10_RSID_100.hdf5",
        "phenotypes-list" => phenotypelist_file,
        "verbosity" => 0,
        "adaptive-cv" => false,
        "save-full" => false,
        "target-type" => "Real",
    )

    UKBBVariantRun(parsed_args)
    # Essential results
    file = jldopen(parsed_args["out"])
    # Missing data
    n_expected = 486
    @test size(file["SAMPLE_IDS"]["CONTINUOUS_1"], 1) == n_expected
    test_base_serialization(file["TMLEREPORTS"], n_expected)
    # QUERY_3 contains a missing genotype so is not actually computed
    @test keys(file["TMLEREPORTS"]) == ["1_1", "1_2"]
    close(file)

    # Clean
    rm(parsed_args["out"])
    rm(parsed_args["queries"])
end

@testset "Test UKBBVariantRun with binary targets" begin
    estimatorfile = joinpath("config", "tmle_config.toml")
    build_query_file()
    parsed_args = Dict(
        "genotypes" => genotypesfile,
        "phenotypes" => binary_phenotypefile,
        "confounders" => confoundersfile,
        "queries" => queryfile,
        "estimator" => estimatorfile,
        "out" => "RSID_10_RSID_100.hdf5",
        "phenotypes-list" => nothing,
        "verbosity" => 0,
        "adaptive-cv" => true,
        "save-full" => true,
        "target-type" => "Bool",
    )

    UKBBVariantRun(parsed_args)

    # Essential results
    file = jldopen(parsed_args["out"])

    @test size(file["SAMPLE_IDS"]["BINARY_1"], 1) == 487
    @test size(file["SAMPLE_IDS"]["BINARY_2"], 1) == 485

    tmlereports = file["TMLEREPORTS"]
    # QUERY_3 contains a missing genotype so is not actually computed
    @test keys(tmlereports) == ["1_1", "1_2", "2_1", "2_2"]
    test_base_serialization(tmlereports, 487, phenotype_id=1)
    test_base_serialization(tmlereports, 485, phenotype_id=2)

    machines = file["MACHINES"]
    Gmach = machines["G"]
    @test length(report(Gmach).cv_report) == 3
    Qmach₁ = machines["Q_1"]
    @test length(report(Qmach₁).cv_report) == 4
    Qmach₂ = machines["Q_2"]
    @test length(report(Qmach₂).cv_report) == 4
    # Adaptive CV
    @test size(report(Qmach₁).cv_report.HALClassifier_1.per_fold[1], 1) == 20

    # Clean
    rm(parsed_args["out"])
    rm(parsed_args["queries"])
end

end;

true


text:
"""
select a subset of groups and swap one poorly particle from each selected groups
* `de`: differential evolution object
* `groups`: groups of particles
"""
function migration!(de, groups)
    # select groups for migration
    sub_group = select_groups(de, groups)
    # select particles within groups for migration
    p_idx,particles = select_particles(sub_group)
    # swap the particles so that p1->pn, p2 -> p1,..., pn -> pn-1
    shift_particles!(sub_group, p_idx, particles)
    return nothing
end

"""
select a subset of groups for migration and return their indices
* `de`: differential evolution object
* `groups`: groups of particles
"""
function select_groups(de, groups)
    N = rand(2:de.n_groups)
    sub_group = sample(groups, N, replace=false)
    return sub_group
end

"""
Select particles from groups for migration. Returns particle index and particles.
* `group_idx`: indices of groups for migration
* `groups`: groups of particles
"""
function select_particles(sub_group)
    Ng = length(sub_group)
    p_idx = fill(0, Ng)
    particles = Vector{eltype(sub_group[1])}(undef,Ng)
    for (i,g) in enumerate(sub_group)
        p_idx[i],particles[i] = select_particle(g)
    end
    return p_idx,particles
end

"""
Select particle from a single chain inversely proportional to its weight
* `group`: a group of particles
"""
function select_particle(group)
    w = map(x -> x.weight, group)
    θ = exp.(-w)/sum(exp.(-w))
    idx = sample(1:length(group), Weights(θ))
    # if numberical error occurs, select the worst particle index (lower is worse)
    any(isnan, θ) ? idx = findmin(w)[2] : nothing
    return idx,group[idx]
end

"""
Swap the particles so that p1->pn, p2 -> p1,..., pn -> pn-1 where
pi is the particle belonging to the ith group
* `group_idx`: indices of groups for migration
* `p_idx`: particle index
* `particles`: particle objects representing position in parameter space
* `groups`: groups of particles
"""
function shift_particles!(sub_group, p_idx, particles)
    # perform a circular shift
    particles = circshift(particles, 1)
    # assign shifted particles to the new chain
    for (g,j,p) in zip(sub_group, p_idx, particles)
        g[j] = p
    end
end