struct LSSSchur{wBType,wEType,BType,EType}
  wBinv::wBType
  wEinv::wEType
  B::BType
  E::EType
end

struct LSSSensitivityFunction{iip,F,A,J,JP,S,PJ,UF,PF,JC,PJC,Alg,fc,JM,pJM,MM,CV,
     PGPU,PGPP,CONFU,CONGP,DG} <: DiffEqBase.AbstractODEFunction{iip}
  f::F
  analytic::A
  jac::J
  jac_prototype::JP
  sparsity::S
  paramjac::PJ
  uf::UF
  pf::PF
  J::JM
  pJ::pJM
  jac_config::JC
  paramjac_config::PJC
  alg::Alg
  numparams::Int
  numindvar::Int
  f_cache::fc
  mass_matrix::MM
  colorvec::CV
  pgpu::PGPU
  pgpp::PGPP
  pgpu_config::CONFU
  pgpp_config::CONGP
  dg_val::DG
end

function LSSSensitivityFunction(sensealg,f,analytic,jac,jac_prototype,sparsity,paramjac,u0,
                                    alg,p,f_cache,mm,
                                    colorvec,tspan,g,dg)

  uf = DiffEqBase.UJacobianWrapper(f,tspan[1],p)
  pf = DiffEqBase.ParamJacobianWrapper(f,tspan[1],copy(u0))

  if DiffEqBase.has_jac(f)
    jac_config = nothing
  else
    jac_config = build_jac_config(sensealg,uf,u0)
  end

  if DiffEqBase.has_paramjac(f)
    paramjac_config = nothing
  else
    paramjac_config = build_param_jac_config(sensealg,pf,u0,p)
  end
  numparams = length(p)
  numindvar = length(u0)
  J = Matrix{eltype(u0)}(undef,numindvar,numindvar)
  pJ = Matrix{eltype(u0)}(undef,numindvar,numparams) # number of funcs size

  # compute gradients of objective
  if dg != nothing
    pgpu = nothing
    pgpp = nothing
    pgpu_config = nothing
    pgpp_config = nothing
    if dg isa Tuple && length(dg) == 2
      dg_val = (similar(u0, numindvar),similar(u0, numparams))
      dg_val[1] .= false
      dg_val[2] .= false
    else
      dg_val = similar(u0, numindvar) # number of funcs size
      dg_val .= false
    end
  else
    pgpu = UGradientWrapper(g,tspan[1],p) # ∂g∂u
    pgpp = ParamGradientWrapper(g,tspan[1],u0) #∂g∂p
    pgpu_config = build_grad_config(sensealg,pgpu,u0,tspan[1])
    pgpp_config = build_grad_config(sensealg,pgpp,p,tspan[1])
    dg_val = (similar(u0, numindvar),similar(u0, numparams))
    dg_val[1] .= false
    dg_val[2] .= false
  end

  LSSSensitivityFunction{isinplace(f),typeof(f),typeof(analytic),
                             typeof(jac),typeof(jac_prototype),typeof(sparsity),
                             typeof(paramjac),
                             typeof(uf),
                             typeof(pf),typeof(jac_config),
                             typeof(paramjac_config),typeof(alg),
                             typeof(f_cache),
                             typeof(J),typeof(pJ),typeof(mm),typeof(f.colorvec),
                             typeof(pgpu),typeof(pgpp),typeof(pgpu_config),typeof(pgpp_config),typeof(dg_val)}(
                             f,analytic,jac,jac_prototype,sparsity,paramjac,uf,pf,J,pJ,
                             jac_config,paramjac_config,alg,
                             numparams,numindvar,f_cache,mm,colorvec,
                             pgpu,pgpp,pgpu_config,pgpp_config,dg_val)
end

struct ForwardLSSProblem{A,C,solType,dtType,umidType,dudtType,SType,Ftype,bType,ηType,wType,vType,windowType,
    ΔtType,G0,G,DG,resType}
  sensealg::A
  diffcache::C
  sol::solType
  dt::dtType
  umid::umidType
  dudt::dudtType
  S::SType
  F::Ftype
  b::bType
  η::ηType
  w::wType
  v::vType
  window::windowType
  Δt::ΔtType
  Nt::Int
  g0::G0
  g::G
  dg::DG
  res::resType
end

function ForwardLSSProblem(sol, sensealg::ForwardLSS, g, dg = nothing;
                            kwargs...)

  @unpack f, p, u0, tspan = sol.prob
  isinplace = DiffEqBase.isinplace(f)

  p == nothing && error("You must have parameters to use parameter sensitivity calculations!")
  !(sol.u isa AbstractVector) && error("`u` has to be an AbstractVector.")

  sense = LSSSensitivityFunction(sensealg,f,f.analytic,f.jac,
                                     f.jac_prototype,f.sparsity,f.paramjac,
                                     u0,sensealg,
                                     p,similar(u0),f.mass_matrix,
                                     f.colorvec,
                                     tspan,g,dg)

  @unpack numparams, numindvar = sense
  Nt = length(sol.t)
  Ndt = Nt-one(Nt)

  # pre-allocate variables
  dt = similar(sol.t, Ndt)
  umid = Matrix{eltype(u0)}(undef,numindvar,Ndt)
  dudt = Matrix{eltype(u0)}(undef,numindvar,Ndt)
  # compute their values
  discretize_ref_trajectory!(dt, umid, dudt, sol, Ndt)

  S = LSSSchur(dt,u0,numindvar,Nt,Ndt,sensealg.alpha)

  if sensealg.alpha isa Number
    η = similar(dt,Ndt)
    window = nothing
    g0 = g(u0,p,tspan[1])
  else
    η = nothing
    window = similar(dt,Nt)
    g0 = nothing
  end

  b = Matrix{eltype(u0)}(undef,numindvar*Ndt,numparams)
  w = similar(dt,numindvar*Ndt)
  v = similar(dt,numindvar*Nt)

  Δt = tspan[2] - tspan[1]
  wB!(S,Δt,Nt,numindvar,dt)
  wE!(S,Δt,dt,sensealg.alpha)

  B!(S,dt,umid,sense,sensealg)
  E!(S,dudt,sensealg.alpha)
  F = SchurLU(S)

  res = similar(u0, numparams)

  ForwardLSSProblem{typeof(sensealg),typeof(sense),typeof(sol),typeof(dt),
    typeof(umid),typeof(dudt),
    typeof(S),typeof(F),typeof(b),typeof(η),typeof(w),typeof(v),typeof(window),typeof(Δt),
    typeof(g0),typeof(g),typeof(dg),typeof(res)}(sensealg,sense,sol,dt,umid,dudt,S,F,b,η,w,v,
    window,Δt,Nt,g0,g,dg,res)
end

function LSSSchur(dt,u0,numindvar,Nt,Ndt,alpha)
  wBinv = similar(dt,numindvar*Nt)
  if alpha isa Number
    wEinv = similar(dt,Ndt)
    E = Matrix{eltype(u0)}(undef,numindvar*Ndt,Ndt)
  else
    wEinv = nothing
    E = nothing
  end
  B = Matrix{eltype(u0)}(undef,numindvar*Ndt,numindvar*Nt)

  LSSSchur(wBinv,wEinv,B,E)
end

# compute discretized reference trajectory
function discretize_ref_trajectory!(dt, umid, dudt, sol, Ndt)
  for i=1:Ndt
    tr = sol.t[i+1]
    tl = sol.t[i]
    ur = sol.u[i+1]
    ul = sol.u[i]
    dt[i] = tr-tl
    copyto!((@view umid[:,i]), (ur + ul)/2)
    copyto!((@view dudt[:,i]), (ur - ul)/dt[i])
  end
  return nothing
end

function wB!(S::LSSSchur,Δt,Nt,numindvar,dt)
  @unpack wBinv = S
  fill!(wBinv, one(Δt))
  dim = numindvar * Nt
  tmp = @view wBinv[1:numindvar]
  tmp ./= dt[1]
  tmp = @view wBinv[dim-2:end]
  tmp ./= dt[end]
  for indx = 2:Nt-1
    tmp = @view wBinv[(indx-1)*numindvar+1:indx*numindvar]
    tmp ./= (dt[indx]+dt[indx-1])
  end

  wBinv .*= 2*Δt
  return nothing
end

wE!(S::LSSSchur,Δt,dt,alpha::Union{CosWindowing,Cos2Windowing}) = nothing

function wE!(S::LSSSchur,Δt,dt,alpha)
  @unpack wEinv = S
  @. wEinv = Δt/(alpha^2*dt)
  return nothing
end

function B!(S::LSSSchur,dt,umid,sense,sensealg)
  @unpack B = S
  @unpack f,J,uf,numindvar,f_cache,jac_config = sense

  fill!(B, zero(eltype(J)))

  for (i,u) in enumerate(eachcol(umid))
    if DiffEqBase.has_jac(f)
      f.jac(J,u,uf.p,uf.t) # Calculate the Jacobian into J
    else
      jacobian!(J, uf, u, f_cache, sensealg, jac_config)
    end
    B0 = @view B[(i-1)*numindvar+1:i*numindvar,i*numindvar+1:(i+1)*numindvar]
    B1 = @view B[(i-1)*numindvar+1:i*numindvar,(i-1)*numindvar+1:i*numindvar]
    B0 .+= I/dt[i] - J/2
    B1 .+= -I/dt[i] -J/2
  end
  return nothing
end

E!(S::LSSSchur,dudt,alpha::Union{CosWindowing,Cos2Windowing}) = nothing

function E!(S::LSSSchur,dudt,alpha)
  @unpack E = S
  numindvar, Ndt = size(dudt)
  for i=1:Ndt
    tmp = @view E[(i-1)*numindvar+1:i*numindvar,i]
    copyto!(tmp, (@view dudt[:,i]))
  end
  return nothing
end

# compute Schur
function SchurLU(S::LSSSchur)
  @unpack B, E, wBinv, wEinv = S
  Smat = B*Diagonal(wBinv)*B'
  (wEinv !== nothing) && (Smat .+= E*Diagonal(wEinv)*E')
  F = lu!(Smat)
  return F
end

function b!(b, prob::ForwardLSSProblem)
  @unpack diffcache, umid, sensealg = prob
  @unpack f, f_cache, pJ, pf, paramjac_config, uf, numindvar = diffcache

  for (i,u) in enumerate(eachcol(umid))
    if DiffEqBase.has_paramjac(f)
      f.paramjac(pJ, u, uf.p, pf.t)
    else
      pf.u = u
      jacobian!(pJ, pf, uf.p, f_cache, sensealg, paramjac_config)
    end
    tmp = @view b[(i-1)*numindvar+1:i*numindvar,:]
    copyto!(tmp, pJ)
  end
  return nothing
end

function shadow_forward(prob::ForwardLSSProblem; t0skip=zero(prob.Δt), t1skip=zero(prob.Δt))
  shadow_forward(prob,prob.sensealg,prob.sensealg.alpha,t0skip,t1skip)
end

function shadow_forward(prob::ForwardLSSProblem,sensealg::ForwardLSS,alpha::Number,t0skip,t1skip)
  @unpack sol, S, F, window, Δt, diffcache, b, w, v, η, res, g, g0, dg, umid = prob
  @unpack wBinv, wEinv, B, E = S
  @unpack dg_val, pgpu, pgpu_config, pgpp, pgpp_config, numparams, numindvar, uf = diffcache

  n0 = searchsortedfirst(sol.t, sol.t[1]+t0skip)
  n1 = searchsortedfirst(sol.t, sol.t[end]-t1skip)

  b!(b,prob)

  ures = @view sol.u[n0:n1]
  umidres = @view umid[:,n0:n1-1]

  # reset
  res .*=false

  for i=1:numparams
    #running average
    g0 *= false
    bpar = @view b[:,i]
    w .= F\bpar
    v .= Diagonal(wBinv)*(B'*w)
    η .= Diagonal(wEinv)*(E'*w)

    ηres = @view η[n0:n1-1]

    for (j,u) in enumerate(ures)
      vtmp = @view v[(n0+j-2)*numindvar+1:(n0+j-1)*numindvar]
      #  final gradient result for ith parameter
      if dg===nothing
        if dg_val isa Tuple
          DiffEqSensitivity.gradient!(dg_val[1], pgpu, u, sensealg,pgpu_config)
          DiffEqSensitivity.gradient!(dg_val[2], pgpp, uf.p, sensealg,pgpp_config)
          res[i] += dot(dg_val[1],vtmp)
          res[i] += dg_val[2][i]
        else
          DiffEqSensitivity.gradient!(dg_val, pgpu, u, sensealg,pgpu_config)
          res[i] += dot(dg_val,vtmp)
        end
      else
        if dg_val isa Tuple
          dg[1](dg_val[1],u,uf.p,nothing,n0+j-1)
          dg[2](dg_val[2],u,uf.p,nothing,n0+j-1)
          res[i] -= dot(dg_val[1],vtmp)
          res[i] -= dg_val[2][i]
        else
          dg(dg_val,u,uf.p,nothing,n0+j-1)
          res[i] -= dot(dg_val,vtmp)
        end
      end
    end
    # mean value
    res[i] = res[i]/(n1-n0+1)

    for (j,u) in enumerate(eachcol(umidres))
      # compute objective
      gtmp = g(u,uf.p,nothing)
      g0 += gtmp
      res[i] -= ηres[j]*gtmp/(n1-n0)
    end
    res[i] = res[i] + sum(ηres)*g0/(n1-n0)^2

  end
  return res
end

function shadow_forward(prob::ForwardLSSProblem,sensealg::ForwardLSS,alpha::CosWindowing,t0skip,t1skip)
  @unpack sol, S, F, window, Δt, diffcache, b, w, v, dg, res = prob
  @unpack wBinv, B = S
  @unpack dg_val, pgpu, pgpu_config, pgpp, pgpp_config, numparams, numindvar, uf = diffcache

  b!(b,prob)

  # windowing (cos)
  @. window = (sol.t-sol.t[1])*convert(eltype(Δt),2*pi/Δt)
  @. window = one(eltype(window)) - cos(window)
  window ./= sum(window)

  res .*= false

  for i=1:numparams
    bpar = @view b[:,i]
    w .= F\bpar
    v .= Diagonal(wBinv)*(B'*w)
    for (j,u) in enumerate(sol.u)
      vtmp = @view v[(j-1)*numindvar+1:j*numindvar]
      #  final gradient result for ith parameter
      if dg===nothing
        if dg_val isa Tuple
          DiffEqSensitivity.gradient!(dg_val[1], pgpu, u, sensealg,pgpu_config)
          DiffEqSensitivity.gradient!(dg_val[2], pgpp, uf.p, sensealg,pgpp_config)
          res[i] += dot(dg_val[1],vtmp)*window[j]
          res[i] += dg_val[2][i]*window[j]
        else
          DiffEqSensitivity.gradient!(dg_val, pgpu, u, sensealg,pgpu_config)
          res[i] += dot(dg_val,vtmp)*window[j]
        end
      else
        if dg_val isa Tuple
          dg[1](dg_val[1],u,uf.p,nothing,j)
          dg[2](dg_val[2],u,uf.p,nothing,j)
          res[i] -= dot(dg_val[1],vtmp)*window[j]
          res[i] -= dg_val[2][i]*window[j]
        else
          dg(dg_val,u,uf.p,nothing,j)
          res[i] -= dot(dg_val,vtmp)*window[j]
        end
      end
    end
  end
  return res
end

function shadow_forward(prob::ForwardLSSProblem,sensealg::ForwardLSS,alpha::Cos2Windowing,t0skip,t1skip)
    @unpack sol, S, F, window, Δt, diffcache, b, w, v, dg, res = prob
    @unpack wBinv, B = S
    @unpack dg_val, pgpu, pgpu_config, pgpp, pgpp_config, numparams, numindvar, uf = diffcache

    b!(b,prob)

    res .*= false

    # windowing cos2
    @. window = (sol.t-sol.t[1])*convert(eltype(Δt),2*pi/Δt)
    @. window = (one(eltype(window)) - cos(window))^2
    window ./= sum(window)

    for i=1:numparams
      bpar = @view b[:,i]
      w .= F\bpar
      v .= Diagonal(wBinv)*(B'*w)
      for (j,u) in enumerate(sol.u)
        vtmp = @view v[(j-1)*numindvar+1:j*numindvar]
        #  final gradient result for ith parameter
        if dg===nothing
          if dg_val isa Tuple
            DiffEqSensitivity.gradient!(dg_val[1], pgpu, u, sensealg,pgpu_config)
            DiffEqSensitivity.gradient!(dg_val[2], pgpp, uf.p, sensealg,pgpp_config)
            res[i] += dot(dg_val[1],vtmp)*window[j]
            res[i] += dg_val[2][i]*window[j]
          else
            DiffEqSensitivity.gradient!(dg_val, pgpu, u, sensealg,pgpu_config)
            res[i] += dot(dg_val,vtmp)*window[j]
          end
        else
          if dg_val isa Tuple
            dg[1](dg_val[1],u,uf.p,nothing,j)
            dg[2](dg_val[2],u,uf.p,nothing,j)
            res[i] -= dot(dg_val[1],vtmp)*window[j]
            res[i] -= dg_val[2][i]*window[j]
          else
            dg(dg_val,u,uf.p,nothing,j)
            res[i] -= dot(dg_val,vtmp)*window[j]
          end
        end
      end
    end
    return res
end

struct AdjointLSSProblem{A,C,solType,dtType,umidType,dudtType,SType,FType,hType,bType,wType,
    ΔtType,G0,G,DG,resType}
  sensealg::A
  diffcache::C
  sol::solType
  dt::dtType
  umid::umidType
  dudt::dudtType
  S::SType
  F::FType
  h::hType
  b::bType
  wa::wType
  Δt::ΔtType
  Nt::Int
  g0::G0
  g::G
  dg::DG
  res::resType
end

function AdjointLSSProblem(sol, sensealg::AdjointLSS, g, dg = nothing;
                            kwargs...)

  @unpack f, p, u0, tspan = sol.prob
  isinplace = DiffEqBase.isinplace(f)

  p == nothing && error("You must have parameters to use parameter sensitivity calculations!")
  !(sol.u isa AbstractVector) && error("`u` has to be an AbstractVector.")

  sense = LSSSensitivityFunction(sensealg,f,f.analytic,f.jac,
                                     f.jac_prototype,f.sparsity,f.paramjac,
                                     u0,sensealg,
                                     p,similar(u0),f.mass_matrix,
                                     f.colorvec,
                                     tspan,g,dg)

  @unpack numparams, numindvar = sense
  Nt = length(sol.t)
  Ndt = Nt-one(Nt)

  # pre-allocate variables
  dt = similar(sol.t, Ndt)
  umid = Matrix{eltype(u0)}(undef,numindvar,Ndt)
  dudt = Matrix{eltype(u0)}(undef,numindvar,Ndt)
  # compute their values
  discretize_ref_trajectory!(dt, umid, dudt, sol, Ndt)

  S = LSSSchur(dt,u0,numindvar,Nt,Ndt,sensealg.alpha)

  if sensealg.alpha isa Number
    g0 = g(u0,p,tspan[1])
  else
    g0 = nothing
  end

  b = Vector{eltype(u0)}(undef,numindvar*Ndt)
  h = Vector{eltype(u0)}(undef,Ndt)
  wa = similar(dt,numindvar*Ndt)

  Δt = tspan[2] - tspan[1]
  wB!(S,Δt,Nt,numindvar,dt)
  wE!(S,Δt,dt,sensealg.alpha)

  B!(S,dt,umid,sense,sensealg)
  E!(S,dudt,sensealg.alpha)
  F = SchurLU(S)
  wBcorrect!(S,sol,g,Nt,sense,sensealg,dg)

  h!(h,g0,g,umid,p,S.wEinv)

  res = similar(u0, numparams)

  AdjointLSSProblem{typeof(sensealg),typeof(sense),typeof(sol),typeof(dt),
    typeof(umid),typeof(dudt),
    typeof(S),typeof(F),typeof(h),typeof(b),typeof(wa),typeof(Δt),
    typeof(g0),typeof(g),typeof(dg),typeof(res)}(sensealg,sense,sol,dt,umid,dudt,S,F,h,b,wa,
    Δt,Nt,g0,g,dg,res)
end

function h!(h,g0,g,u,p,wEinv)

  for (j,uj) in enumerate(eachcol(u))
    # compute objective
    h[j] = g(uj,p,nothing)
  end
  h .= -(h .- mean(h)) / (size(u)[2])

  @. h = wEinv*h

  return nothing
end

function wBcorrect!(S,sol,g,Nt,sense,sensealg,dg)
  @unpack dg_val, pgpu, pgpu_config, numparams, numindvar, uf = sense
  @unpack wBinv = S

  for (i,u) in enumerate(sol.u)
    _wBinv = @view wBinv[(i-1)*numindvar+1:i*numindvar]
    if dg === nothing
      if dg_val isa Tuple
        DiffEqSensitivity.gradient!(dg_val[1], pgpu, u, sensealg, pgpu_config)
        @. _wBinv = _wBinv*dg_val[1]/Nt
      else
        DiffEqSensitivity.gradient!(dg_val, pgpu, u, sensealg, pgpu_config)
        @. _wBinv = _wBinv*dg_val/Nt
      end
    else
      if dg_val isa Tuple
        dg[1](dg_val[1],u,uf.p,nothing,i)
        @. _wBinv = -_wBinv*dg_val[1]/Nt
      else
        dg(dg_val,u,uf.p,nothing,i)
        @. _wBinv = -_wBinv*dg_val/Nt
      end
    end
  end
  return nothing
end

function shadow_adjoint(prob::AdjointLSSProblem; t0skip=zero(prob.Δt), t1skip=zero(prob.Δt))
  shadow_adjoint(prob,prob.sensealg,prob.sensealg.alpha,t0skip,t1skip)
end

function shadow_adjoint(prob::AdjointLSSProblem,sensealg::AdjointLSS,alpha::Number,t0skip,t1skip)
  @unpack sol, S, F, Δt, diffcache, h, b, wa, res, g, g0, dg, umid = prob
  @unpack wBinv, B, E = S
  @unpack dg_val, pgpp, pgpp_config, numparams, numindvar, uf, f, f_cache, pJ, pf, paramjac_config = diffcache

  b .= E*h + B*wBinv
  wa .= F\b

  n0 = searchsortedfirst(sol.t, sol.t[1]+t0skip)
  n1 = searchsortedfirst(sol.t, sol.t[end]-t1skip)

  umidres = @view umid[:,n0:n1-1]
  wares = @view wa[(n0-1)*numindvar+1:(n1-1)*numindvar]

  # reset
  res .*= false

  if dg_val isa Tuple
    for (j,u) in enumerate(eachcol(umidres))
      if dg === nothing
        DiffEqSensitivity.gradient!(dg_val[2], pgpp, uf.p, sensealg, pgpp_config)
        @. res += dg_val[2]
      else
        dg[2](dg_val[2],u,uf.p,nothing,n0+j-1)
        @. res -= dg_val[2]
      end
    end
    res ./= (size(umidres)[2])
  end

  for (j,u) in enumerate(eachcol(umidres))
    _wares = @view wares[(j-1)*numindvar+1:j*numindvar]
    if DiffEqBase.has_paramjac(f)
      f.paramjac(pJ, u, uf.p, pf.t)
    else
      pf.u = u
      jacobian!(pJ, pf, uf.p, f_cache, sensealg, paramjac_config)
    end

    res .+= pJ'*_wares
  end

  return res
end

check_for_g(sensealg::Union{ForwardLSS,AdjointLSS},g)=((sensealg.alpha isa Number && g===nothing) && error("Time dilation needs explicit knowledge of g. Either pass `g` as a kwarg or use ForwardLSS/AdjointLSS with windowing."))


function save_ivfadc_index(filename::AbstractString,
                           ivfadc::IVFADCIndex{U,I,Dc,Dr,T,Q}
                          ) where {U,I,Dc,Dr,T,Q<:NaiveQuantizer}
    open(filename, "w") do fid
        # Initialize all variables needed to write
        nrows, nclusters = size(ivfadc.coarse_quantizer.vectors)
        n = length(ivfadc)
        m = length(ivfadc.residual_quantizer.codebooks)
        k = ivfadc.residual_quantizer.k
        d = size(ivfadc.residual_quantizer.codebooks[1].vectors, 1)

        # Get types
        coarsequanttype_str = "NaiveQuantizer"
        quanttype_str = string(typeof(ivfadc.residual_quantizer.quantization))
        quanteltype_str = string(U)
        indextype_str = string(I)
        disttypecoarse_str = string(Dc)
        disttyperesidual_str = string(Dr)
        origeltype_str = string(T)

        # Start writing text information
        println(fid, "$nrows $nclusters")
        println(fid, "$n $m $k $d")
        println(fid, "$coarsequanttype_str")
        println(fid, "$quanttype_str")
        println(fid, "$quanteltype_str")
        println(fid, "$indextype_str")
        println(fid, "$disttypecoarse_str")
        println(fid, "$disttyperesidual_str")
        println(fid, "$origeltype_str")

        # Write coarse quantizer
        _write_naive_coarse_quantizer(fid, ivfadc.coarse_quantizer)

        # Write residual quantizer
        _write_residual_quantizer(fid, ivfadc.residual_quantizer, (nrows, nclusters))

        # Write inverse index
        _write_inverse_index(fid, ivfadc.inverse_index, (nrows, nclusters))
    end
end

_write_naive_coarse_quantizer(fid::IO, quantizer) = begin
    nclusters = size(quantizer.vectors, 2)
    for i in 1:nclusters
        write(fid, quantizer.vectors[:,i])
    end
end

_write_residual_quantizer(fid::IO, quantizer, dims) = begin
    nrows, nclusters = dims
    d = size(quantizer.codebooks[1].vectors, 1)
    m = length(quantizer.codebooks)
    for i in 1:m  # codebooks
        write(fid, quantizer.codebooks[i].codes)
        for j in 1:d
            write(fid, quantizer.codebooks[i].vectors[j,:])
        end
    end
    for i in 1:nrows  # rotation matrix
        write(fid, quantizer.rot[:,i])
    end
end

_write_inverse_index(fid::IO, inverse_index, dims) = begin
    nrows, nclusters = dims
    for i in 1:nclusters
        clsize = length(inverse_index[i].idxs)
        write(fid, clsize)
        write(fid, inverse_index[i].idxs)
        for j in 1:clsize
            write(fid, inverse_index[i].codes[j])
        end
    end
end

# Generate a WordVectors object from binary file
function load_ivfadc_index(filename::AbstractString)
    open(filename, "r") do fid
        nrows, nclusters = map(x -> parse(Int, x), split(readline(fid), ' '))
        n, m, k, d = map(x -> parse(Int, x), split(readline(fid), ' '))
        CQ = _read_type_from_line(readline(fid))
        Q = _read_type_from_line(readline(fid))
        U = eval(Symbol(readline(fid)))
        I = eval(Symbol(readline(fid)))
        Dc = _read_type_from_line(readline(fid))
        Dr = _read_type_from_line(readline(fid))
        T = eval(Symbol(readline(fid)))

        # Read coarse quantizer
        coarse_quantizer = _load_coarse_quantizer(CQ, fid;
                                dims=(nrows, nclusters),
                                disttype=Dc,
                                vectype=T)

        # Read residual quantizer
        cbooks = Vector{CodeBook{U,T}}(undef, m)
        codes_length = sizeof(U) * k
        codevecs_length = sizeof(T) * k
        for i = 1:m
            codes = collect(reinterpret(U, read(fid, codes_length)))
            vectors = Matrix{T}(undef, d, k)
            for j in 1:d
                vectors[j,:] = collect(reinterpret(T, read(fid, codevecs_length)))
            end
            cbooks[i] = CodeBook(codes, vectors)
        end
        binary_length = sizeof(T) * nrows
        rotmat = Matrix{T}(undef, nrows, nrows)
        for i in 1:nrows
            rotmat[:,i] = collect(reinterpret(T, read(fid, binary_length)))
        end
        residual_quantizer = ArrayQuantizer(Q(), (nrows, n), cbooks, k, Dr(), rotmat)

        # Read inverse index
        inverse_index = IVFADC.InvertedIndex{I,U}(undef, nclusters)
        encoded_length = sizeof(U) * m
        for i in 1:nclusters
            clsize = reinterpret(Int, read(fid, sizeof(Int)))[1]
            idxs_length = sizeof(I) * clsize
            idxs = collect(reinterpret(I, read(fid, idxs_length)))
            codes = Vector{Vector{U}}(undef, clsize)
            for j in 1:clsize
                codes[j] = collect(reinterpret(U, read(fid, encoded_length)))
            end
            inverse_index[i] = InvertedList(idxs, codes)
        end
        return IVFADCIndex(coarse_quantizer, residual_quantizer, inverse_index)
    end
end

_read_type_from_line(line) = begin
    if occursin('.', line)
        _module, _type = String.(split(line, "."))
    else
        _module, _type = string(@__MODULE__), line
    end
    return eval(Expr(:., Symbol(_module), QuoteNode(Symbol(_type))))
end

_load_coarse_quantizer(::Type{CQ},
                       fid::IO;
                       dims=(0,0),
                       disttype::Type{Dc}=typeof(DEFAULT_COARSE_DISTANCE),
                       vectype::Type{T}=AbstractFloat
                      ) where {CQ<:NaiveQuantizer,Dc,T} = begin
    nrows, nclusters = dims
    data = Matrix{T}(undef, nrows, nclusters)
    binary_length = sizeof(T) * nrows
    for i in 1:nclusters
        data[:,i] = collect(reinterpret(T, read(fid, binary_length)))
    end
    coarse_quantizer = IVFADC.NaiveQuantizer{Dc,T}(data)
end

function save_ivfadc_index(filename::AbstractString,
                           ivfadc::IVFADCIndex{U,I,Dc,Dr,T,Q}
                          ) where {U,I,Dc,Dr,T,Q<:HNSWQuantizer}
    open(filename, "w") do fid
        # Initialize all variables needed to write
        nrows = length(ivfadc.coarse_quantizer.hnsw.data[1])
        nclusters = length(ivfadc.coarse_quantizer.hnsw.data)
        n = length(ivfadc)
        m = length(ivfadc.residual_quantizer.codebooks)
        k = ivfadc.residual_quantizer.k
        d = size(ivfadc.residual_quantizer.codebooks[1].vectors, 1)

        # Get types
        coarsequanttype_str = "HNSWQuantizer"
        quanttype_str = string(typeof(ivfadc.residual_quantizer.quantization))
        quanteltype_str = string(U)
        indextype_str = string(I)
        disttypecoarse_str = string(Dc)
        disttyperesidual_str = string(Dr)
        origeltype_str = string(T)

        # Start writing text information
        println(fid, "$nrows $nclusters")
        println(fid, "$n $m $k $d")
        println(fid, "$coarsequanttype_str")
        println(fid, "$quanttype_str")
        println(fid, "$quanteltype_str")
        println(fid, "$indextype_str")
        println(fid, "$disttypecoarse_str")
        println(fid, "$disttyperesidual_str")
        println(fid, "$origeltype_str")

        # Write coarse quantizer
        _write_hnsw_coarse_quantizer(fid, ivfadc.coarse_quantizer.hnsw)

        # Write residual quantizer
        _write_residual_quantizer(fid, ivfadc.residual_quantizer, (nrows, nclusters))

        # Write inverse index
        _write_inverse_index(fid, ivfadc.inverse_index, (nrows, nclusters))
    end
end

_write_hnsw_coarse_quantizer(fid::IO,
                             hnsw::HierarchicalNSW{T,F,V,M}
                            ) where {T,F,V,M} = begin
    println(fid, string(T))
    nclusters= length(hnsw.data)

    # Start writing
    write(fid, hnsw.lgraph.M0)          # .lgraph
    write(fid, hnsw.lgraph.M)
    write(fid, hnsw.lgraph.m_L)
    nlinkedlists = length(hnsw.lgraph.linklist)
    write(fid, nlinkedlists)
    for i in 1:nlinkedlists
        listlength = length(hnsw.lgraph.linklist[i])
        write(fid, listlength)
        write(fid, hnsw.lgraph.linklist[i])
    end
    write(fid, hnsw.added)              # .added
    for i in 1:nclusters                # .data
        write(fid, hnsw.data[i])
    end
    write(fid, hnsw.ep)                 # .ep
    write(fid, hnsw.entry_level)        # .entry_level
    write(fid, hnsw.vlp.num_elements)   # .vlp
    nvlps = length(hnsw.vlp.pool)
    write(fid, nvlps)
    for i in 1:nvlps
        write(fid, hnsw.vlp.pool[i].visited_value)  # UInt8
        write(fid, length(hnsw.vlp.pool[i].list))
        write(fid, hnsw.vlp.pool[i].list)
    end
    # metric can be skipped             # .metric
    write(fid, hnsw.efConstruction)     # .efConstruction
    write(fid, hnsw.ef)                 # .ef
end

_load_coarse_quantizer(::Type{CQ},
                       fid::IO;
                       dims=(0,0),
                       disttype::Type{Dc}=typeof(DEFAULT_COARSE_DISTANCE),
                       vectype::Type{F}=AbstractFloat
                      ) where {CQ<:HNSWQuantizer,Dc,F} = begin
    # We already have F (float type) and M (metric type = Dc)
    T = eval(Symbol(readline(fid)))
    V = Vector{Vector{F}}
    intsize = sizeof(Int)
    nrows, nclusters = dims

    # .lgraph
    M0 = reinterpret(Int, read(fid, intsize))[1]
    M = reinterpret(Int, read(fid, intsize))[1]
    m_L = reinterpret(Float64, read(fid, 8))[1]  # a Float64
    nlinkedlists = reinterpret(Int, read(fid, intsize))[1]
    linklist = Vector{Vector{T}}(undef, nlinkedlists)
    for i in 1:nlinkedlists
        listlength = reinterpret(Int, read(fid, intsize))[1]
        linklist[i] = collect(reinterpret(T, read(fid, sizeof(T) * listlength)))
    end
    lgraph = HNSW.LayeredGraph{T}(linklist, M0, M, m_L)

    # .added
    added = collect(reinterpret(Bool, read(fid, nclusters * sizeof(Bool))))

    # .data
    data = V(undef, nclusters)
    for i in 1:nclusters
        data[i] = collect(reinterpret(F, read(fid, nrows * sizeof(F))))
    end

    # .ep
    ep = reinterpret(T, read(fid, sizeof(T)))[1]

    # .entry_level
    entry_level = reinterpret(Int, read(fid, intsize))[1]

    # .vlp
    num_elements = reinterpret(Int, read(fid, intsize))[1]
    nvlps = reinterpret(Int, read(fid, intsize))[1]
    pool = Vector{HNSW.VisitedList}(undef, nvlps)
    for i in 1:nvlps
        visited_value = read(fid, 1)[1]
        len = reinterpret(Int, read(fid, intsize))[1]
        list = read(fid, len)
        pool[i] = HNSW.VisitedList(list, visited_value)
    end
    vlp = HNSW.VisitedListPool(pool, num_elements)

    # .metric need not be written

    # .efConstruction
    efConstruction = reinterpret(Int, read(fid, intsize))[1]

    # .ef
    ef = reinterpret(Int, read(fid, intsize))[1]
    hnsw = HNSW.HierarchicalNSW{T,F,V,Dc}(lgraph, added, data,
                ep, entry_level, vlp, Dc(), efConstruction, ef)
    return HNSWQuantizer{T,V,Dc,F}(hnsw)
end


using Compat, FastGroupBy
import Base: Forward
import SortingAlgorithms: RADIX_MASK, RADIX_SIZE, uint_mapping
using BenchmarkTools

using SortingAlgorithms, Base.Order, Compat, IndexedTables, DataFrames
import Base: ht_keyindex, rehash!, _setindex!
import SortingAlgorithms: uint_mapping, RADIX_SIZE, RADIX_MASK
import DataFrames: DataFrame, AbstractDataFrame
import IndexedTables: IndexedTable, column
import PooledArrays.PooledArray
import CategoricalArrays.CategoricalArray
import SplitApplyCombine.groupreduce

using StatsBase

const N = 1_000_000
const K = 100
by = rand(Int32(1):Int32(N/K), N);
val = similar(by);
T = Int

using DataStructures

# define a heap that invovles the byval and S
struct ByVal{T<:Integer, S <: Number}
    index::Integer
    by::T
    val::S
end

import Base.isless
isless(x::ByVal,y::ByVal) = x.index < y.index

"WAAAAAAAYYY too slow uses less memory by building a heap"
function sumby_lessmem_heap{T, S<:Number}(by::AbstractVector{T},  val::AbstractVector{S})::Dict{T,S}
    o = Forward
    lo = 1
    hi = length(by)

    if lo >= hi;  return by;  end

    # Make sure we're sorting a bits type
    if !isbits(T)
      error("the by vector must be of bits types (e.g. Float64, Integer) got $T")
    end

    # Init
    iters = ceil(Integer, sizeof(T)*8/RADIX_SIZE)
    bin = zeros(UInt32, 2^RADIX_SIZE, iters)

    # Histogram for each element, radix
    for i = lo:hi
        v = uint_mapping(o, by[i])
        for j = 1:iters
            idx = @compat(Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK)) + 1
            @inbounds bin[idx,j] += 1
        end
    end

    # Sort!
    len = hi-lo+1
    for j = 1:iters
        # Unroll first data iteration, check for degenerate case
        v = uint_mapping(o, by[hi])
        idx = @compat(Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK)) + 1

        # are all values the same at this radix?
        if bin[idx,j] == len;  continue;  end

        cbin = cumsum(bin[:,j])
        ci = cbin[idx]

        # ByValHeap -- keeps the data in a heap
        bvheap = binary_maxheap(ByVal)

        # if ci == hi then do nothing
        if ci < hi
          push!(bvheap, ByVal(ci, by[ci], val[ci]))
          by[ci] = by[hi]
          val[ci] = val[hi]
        end
        cbin[idx] -= 1

        loopto = lo-1
        i = hi - 1 # i is used to keep track of where we are up to

        # Finish the loop...
        #   @inbounds for i in hi-1:-1:lo
        while i > 0
          if isempty(bvheap)
              loopto = lo-1
          else
              loopto = top(bvheap).index
          end
          while i > loopto
              v = uint_mapping(o, by[i])
              idx = @compat(Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK)) + 1
              ci = cbin[idx]
              if ci > i # it's ok to just put it in place
                  by[ci] = by[i]
                  val[ci] = val[i]
              elseif ci < i # keep the value in a heap for retrieval
                  push!(bvheap, ByVal(ci, by[ci], val[ci]))
                  by[ci] = by[i]
                  val[ci] = val[i]
                  # make the loop end sooner
                  if ci > loopto; loopto = ci; end
              end

              cbin[idx] -= 1
              i -= 1
          end
          # if it got here without being 0 then it must be up to a heap value
          if i != 0
              tmp = pop!(bvheap)
              v = uint_mapping(o, tmp.by)
              idx = @compat(Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK)) + 1
              ci = cbin[idx]
              if ci > i # it's ok to just put it in place
                  by[ci] = tmp.by
                  val[ci] = tmp.val
              elseif ci < i # keep the value in a heap for retrieval
                  push!(bvheap, ByVal(ci, by[ci], val[ci]))
                  by[ci] = tmp.by
                  val[ci] = tmp.val
              end

              cbin[idx] -= 1
              i -= 1
          end
        end
    end

    sumby_contiguous(by, val)
end

sumby_lessmem_heap([1:10...],[1:10]

function abc()
  by1, val1 = copy(by), copy(val)
  @elapsed sumby_lessmem_heap(by1, val1)
end

function def()
  by1, val1 = copy(by), copy(val)
  @elapsed sumby(by1, val1)
end

a1 = mean([abc() for i = 1:5][2:end])
b1 = mean([def() for i = 1:5][2:end])
1 - a1/b1

function ghi()
  by1, val1 = copy(by), copy(val)
  @elapsed sumby_lessmem_chain(by1, val1)
end
c1 = mean([ghi() for i = 1:5][2:end])
1 - c1/b1


# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

function MOI.features_available(d::NLPEvaluator)
    # Check if we are missing any hessians for user-defined multivariate
    # operators, in which case we need to disable :Hess and :HessVec.
    d.disable_2ndorder = any(
        op -> op.∇²f === nothing,
        d.data.operators.registered_multivariate_operators,
    )
    if d.disable_2ndorder
        return [:Grad, :Jac, :JacVec]
    end
    return [:Grad, :Jac, :JacVec, :Hess, :HessVec]
end

function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
    # Check that we support the features requested by the user.
    available_features = MOI.features_available(d)
    for feature in requested_features
        if !(feature in available_features)
            error("Unsupported feature $feature")
        end
    end
    moi_index_to_consecutive_index =
        Dict(x => i for (i, x) in enumerate(d.ordered_variables))
    N = length(moi_index_to_consecutive_index)
    #
    largest_user_input_dimension = 1
    for op in d.data.operators.registered_multivariate_operators
        largest_user_input_dimension = max(largest_user_input_dimension, op.N)
    end
    d.objective = nothing
    d.user_output_buffer = zeros(largest_user_input_dimension)
    d.jac_storage = zeros(max(N, largest_user_input_dimension))
    d.constraints = _FunctionStorage[]
    d.last_x = fill(NaN, N)
    d.want_hess = :Hess in requested_features
    want_hess_storage = (:HessVec in requested_features) || d.want_hess
    coloring_storage = Coloring.IndexedSet(N)
    max_expr_length = 0
    #
    main_expressions = [c.expression.nodes for (_, c) in d.data.constraints]
    if d.data.objective !== nothing
        pushfirst!(main_expressions, d.data.objective.nodes)
    end
    d.subexpression_order, individual_order = _order_subexpressions(
        main_expressions,
        [expr.nodes for expr in d.data.expressions],
    )
    num_subexpressions = length(d.data.expressions)
    d.subexpression_linearity = Vector{Linearity}(undef, num_subexpressions)
    subexpression_variables = Vector{Vector{Int}}(undef, num_subexpressions)
    subexpression_edgelist =
        Vector{Set{Tuple{Int,Int}}}(undef, num_subexpressions)
    d.subexpressions = Vector{_SubexpressionStorage}(undef, num_subexpressions)
    d.subexpression_forward_values = zeros(num_subexpressions)
    d.subexpression_reverse_values = zeros(num_subexpressions)
    for k in d.subexpression_order
        # Only load expressions which actually are used
        d.subexpression_forward_values[k] = NaN
        subex = _SubexpressionStorage(
            d.data.expressions[k],
            d.subexpression_linearity,
            moi_index_to_consecutive_index,
        )
        d.subexpressions[k] = subex
        d.subexpression_linearity[k] = subex.linearity
        if d.want_hess
            empty!(coloring_storage)
            _compute_gradient_sparsity!(coloring_storage, subex.nodes)
            # union with all dependent expressions
            for idx in _list_subexpressions(subex.nodes)
                union!(coloring_storage, subexpression_variables[idx])
            end
            subexpression_variables[k] = collect(coloring_storage)
            empty!(coloring_storage)
            linearity = _classify_linearity(
                subex.nodes,
                subex.adj,
                d.subexpression_linearity,
            )
            edgelist = _compute_hessian_sparsity(
                subex.nodes,
                subex.adj,
                linearity,
                coloring_storage,
                subexpression_edgelist,
                subexpression_variables,
            )
            subexpression_edgelist[k] = edgelist
        end
    end
    max_chunk = 1
    if d.data.objective !== nothing
        d.objective = _FunctionStorage(
            main_expressions[1],
            d.data.objective.values,
            N,
            coloring_storage,
            d.want_hess,
            d.subexpressions,
            individual_order[1],
            d.subexpression_linearity,
            subexpression_edgelist,
            subexpression_variables,
            moi_index_to_consecutive_index,
        )
        max_expr_length = max(max_expr_length, length(d.objective.nodes))
        max_chunk = max(max_chunk, size(d.objective.seed_matrix, 2))
    end
    for (k, (_, constraint)) in enumerate(d.data.constraints)
        idx = d.data.objective !== nothing ? k + 1 : k
        push!(
            d.constraints,
            _FunctionStorage(
                main_expressions[idx],
                constraint.expression.values,
                N,
                coloring_storage,
                d.want_hess,
                d.subexpressions,
                individual_order[idx],
                d.subexpression_linearity,
                subexpression_edgelist,
                subexpression_variables,
                moi_index_to_consecutive_index,
            ),
        )
        max_expr_length = max(max_expr_length, length(d.constraints[end].nodes))
        max_chunk = max(max_chunk, size(d.constraints[end].seed_matrix, 2))
    end
    # 10 is hardcoded upper bound to avoid excess memory allocation
    max_chunk = min(max_chunk, 10)
    if d.want_hess || want_hess_storage
        d.input_ϵ = zeros(max_chunk * N)
        d.output_ϵ = zeros(max_chunk * N)
        #
        len = max_chunk * max_expr_length
        d.forward_storage_ϵ = zeros(len)
        d.partials_storage_ϵ = zeros(len)
        d.reverse_storage_ϵ = zeros(len)
        #
        len = max_chunk * length(d.subexpressions)
        d.subexpression_forward_values_ϵ = zeros(len)
        d.subexpression_reverse_values_ϵ = zeros(len)
        #
        for k in d.subexpression_order
            len = max_chunk * length(d.subexpressions[k].nodes)
            d.subexpressions[k].forward_storage_ϵ = zeros(len)
            d.subexpressions[k].partials_storage_ϵ = zeros(len)
            d.subexpressions[k].reverse_storage_ϵ = zeros(len)
        end
        d.max_chunk = max_chunk
        if d.want_hess
            _compute_hessian_lagrangian_structure(d)
        end
    end
    return
end

function MOI.eval_objective(d::NLPEvaluator, x)
    if d.objective === nothing
        error("No nonlinear objective.")
    end
    _reverse_mode(d, x)
    return d.objective.forward_storage[1]
end

function MOI.eval_objective_gradient(d::NLPEvaluator, g, x)
    if d.objective === nothing
        error("No nonlinear objective.")
    end
    _reverse_mode(d, x)
    fill!(g, 0.0)
    _extract_reverse_pass(g, d, d.objective)
    return
end

function MOI.eval_constraint(d::NLPEvaluator, g, x)
    _reverse_mode(d, x)
    for i in 1:length(d.constraints)
        g[i] = d.constraints[i].forward_storage[1]
    end
    return
end

function MOI.jacobian_structure(d::NLPEvaluator)
    J = Tuple{Int64,Int64}[]
    for (row, constraint) in enumerate(d.constraints)
        for col in constraint.grad_sparsity
            push!(J, (row, col))
        end
    end
    return J
end

function MOI.eval_constraint_jacobian(d::NLPEvaluator, J, x)
    _reverse_mode(d, x)
    fill!(J, 0.0)
    offset = 0
    for ex in d.constraints
        for i in ex.grad_sparsity
            d.jac_storage[i] = 0.0
        end
        _extract_reverse_pass(d.jac_storage, d, ex)
        for (k, idx) in enumerate(ex.grad_sparsity)
            J[offset+k] = d.jac_storage[idx]
        end
        offset += length(ex.grad_sparsity)
    end
    return
end

function MOI.eval_constraint_jacobian_product(d::NLPEvaluator, y, x, w)
    fill!(y, 0.0)
    J_struct = MOI.jacobian_structure(d)
    nnz = length(J_struct)
    J = zeros(nnz)
    MOI.eval_constraint_jacobian(d, J, x)
    for (k, (i, j)) in enumerate(J_struct)
        y[i] += J[k] * w[j]
    end
    return
end

function MOI.eval_constraint_jacobian_transpose_product(
    d::NLPEvaluator,
    y::AbstractVector{Float64},
    x::AbstractVector{Float64},
    w::AbstractVector{Float64},
)
    fill!(y, 0.0)
    J_struct = MOI.jacobian_structure(d)
    nnz = length(J_struct)
    J = zeros(nnz)
    MOI.eval_constraint_jacobian(d, J, x)
    for (k, (i, j)) in enumerate(J_struct)
        y[j] += J[k] * w[i]
    end
    return y
end

function MOI.hessian_lagrangian_structure(d::NLPEvaluator)
    @assert d.want_hess
    return d.hessian_sparsity
end

function _compute_hessian_lagrangian_structure(d::NLPEvaluator)
    d.hessian_sparsity = Tuple{Int64,Int64}[]
    if d.objective !== nothing
        append!(d.hessian_sparsity, zip(d.objective.hess_I, d.objective.hess_J))
    end
    for c in d.constraints
        append!(d.hessian_sparsity, zip(c.hess_I, c.hess_J))
    end
    return
end

function MOI.eval_hessian_lagrangian(d::NLPEvaluator, H, x, σ, μ)
    @assert d.want_hess
    _reverse_mode(d, x)
    fill!(d.input_ϵ, 0.0)
    offset = 0
    if d.objective !== nothing
        offset += _eval_hessian(d, d.objective, H, σ, offset)::Int
    end
    for (i, ex) in enumerate(d.constraints)
        offset += _eval_hessian(d, ex, H, μ[i], offset)::Int
    end
    return
end

function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
    _reverse_mode(d, x)
    fill!(h, 0.0)
    T = ForwardDiff.Partials{1,Float64}
    input_ϵ = reinterpret(T, d.input_ϵ)
    output_ϵ = reinterpret(T, d.output_ϵ)
    for i in 1:length(x)
        input_ϵ[i] = ForwardDiff.Partials((v[i],))
    end
    # forward evaluate all subexpressions once
    subexpr_forward_values_ϵ = reinterpret(T, d.subexpression_forward_values_ϵ)
    for i in d.subexpression_order
        subexpr = d.subexpressions[i]
        subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(
            d,
            subexpr,
            reinterpret(T, subexpr.forward_storage_ϵ),
            reinterpret(T, subexpr.partials_storage_ϵ),
            input_ϵ,
            subexpr_forward_values_ϵ,
            d.data.operators,
        )
    end
    # we only need to do one reverse pass through the subexpressions as well
    subexpr_reverse_values_ϵ = reinterpret(T, d.subexpression_reverse_values_ϵ)
    fill!(subexpr_reverse_values_ϵ, zero(T))
    fill!(d.subexpression_reverse_values, 0.0)
    fill!(d.reverse_storage_ϵ, 0.0)
    fill!(output_ϵ, zero(T))
    if d.objective !== nothing
        _forward_eval_ϵ(
            d,
            d.objective,
            reinterpret(T, d.forward_storage_ϵ),
            reinterpret(T, d.partials_storage_ϵ),
            input_ϵ,
            subexpr_forward_values_ϵ,
            d.data.operators,
        )
        _reverse_eval_ϵ(
            output_ϵ,
            d.objective,
            reinterpret(T, d.reverse_storage_ϵ),
            reinterpret(T, d.partials_storage_ϵ),
            d.subexpression_reverse_values,
            subexpr_reverse_values_ϵ,
            σ,
            zero(T),
        )
    end
    for (i, con) in enumerate(d.constraints)
        _forward_eval_ϵ(
            d,
            con,
            reinterpret(T, d.forward_storage_ϵ),
            reinterpret(T, d.partials_storage_ϵ),
            input_ϵ,
            subexpr_forward_values_ϵ,
            d.data.operators,
        )
        _reverse_eval_ϵ(
            output_ϵ,
            con,
            reinterpret(T, d.reverse_storage_ϵ),
            reinterpret(T, d.partials_storage_ϵ),
            d.subexpression_reverse_values,
            subexpr_reverse_values_ϵ,
            μ[i],
            zero(T),
        )
    end
    for i in length(d.subexpression_order):-1:1
        j = d.subexpression_order[i]
        subexpr = d.subexpressions[j]
        _reverse_eval_ϵ(
            output_ϵ,
            subexpr,
            reinterpret(T, subexpr.reverse_storage_ϵ),
            reinterpret(T, subexpr.partials_storage_ϵ),
            d.subexpression_reverse_values,
            subexpr_reverse_values_ϵ,
            d.subexpression_reverse_values[j],
            subexpr_reverse_values_ϵ[j],
        )
    end
    for i in 1:length(x)
        h[i] += output_ϵ[i].values[1]
    end
    return
end


import Knet: update!, train!

abstract type AbstractTrainable end
abstract type AbstractLayer <: AbstractTrainable end
# TODO: rename StaticLayer -> AbstractStaticLayer and StochasticLayer -> AbstractStochasticLayer
abstract type StaticLayer <: AbstractLayer end
abstract type StochasticLayer <: AbstractLayer end

"""
Initialise optimizer for each trainable parameters.
"""
function initoptim!(model::AbstractTrainable, otype; args...)
    for p in Knet.params(model)
        p.opt = otype(; args...)
    end
end

"""
Update all parameters in `ps` by back-propgating `losstape`.
"""
function update!(ps::Array, losstape::AutoGrad.Tape)
    for p in ps
        g = AutoGrad.grad(losstape, p)
        @assert p.opt != nothing "$p has no optimizer set up"
        update!(AutoGrad.value(p), g, p.opt)
    end
end

"""
Dynamic Adam - an Adam that can initialise first and second momentum during optimization.
This is used to deal dynamic models.
"""
mutable struct DynamicAdam
    adam::Knet.Adam
    isreset::Bool
end

DynamicAdam(isreset=true; kwargs...) = DynamicAdam(Knet.Adam(; kwargs...), isreset)

for T in (Array{Float32},Array{Float64},Knet.KnetArray{Float32},Knet.KnetArray{Float64}); @eval begin
    function update!(w::$T, g::$T, p::DynamicAdam)
        if !(p.adam.fstm === nothing) && length(g) > length(p.adam.fstm)
            F = eltype($T)
            if p.isreset
                p.adam.fstm = $T(zeros(F, size(g)...))
                p.adam.scndm = $T(zeros(F, size(g)...))
            else
                # Vector
                if length(size(g)) == 1
                    d_pad = length(g) - length(p.adam.fstm)
                    p.adam.fstm = vcat([p.adam.fstm, $T(zeros(F, d_pad))]...)
                    p.adam.scndm = vcat([p.adam.scndm, $T(zeros(F, d_pad))]...)
                # Matrix
                else
                    # Expand rows
                    if size(g, 1) > size(p.adam.fstm, 1)
                        d_pad = size(g, 1) - size(p.adam.fstm, 1)
                        p.adam.fstm = vcat([p.adam.fstm, $T(zeros(F, d_pad, size(g, 2)))]...)
                        p.adam.scndm = vcat([p.adam.scndm, $T(zeros(F, d_pad, size(g, 2)))]...)
                    # Expand cols
                    else
                        d_pad = size(g, 2) - size(p.adam.fstm, 2)
                        p.adam.fstm = hcat([p.adam.fstm, $T(zeros(F, size(g, 1), d_pad))]...)
                        p.adam.scndm = hcat([p.adam.scndm, $T(zeros(F, size(g, 1), d_pad))]...)
                    end
                end
            end
        end
        update!(w, g, p.adam)
    end
end; end

"""
Update all parameters of `model` by back-propgating `losstape`.
"""
update!(model::AbstractTrainable, losstape::AutoGrad.Tape) = update!(Knet.params(model), losstape)

"""
Compute the number of parameters.
"""
function numparams(model::AbstractTrainable)
    n = 0
    for p in Knet.params(model)
        n += prod(size(p))
    end
    return n
end

"""
Neural models.

The difference between a layer and a model is that for models, loss are defined.

All `NeuralModel` should be compatible with the following signature.
```julia
Base.show(nm::NeuralModel)
(nm::NeuralModel)(data, training::Val{true})    # loss
(nm::NeuralModel)(data, training::Val{false})   # evaluate metric
```
"""
abstract type NeuralModel <: AbstractTrainable end
Base.show(nm::NeuralModel) = error("Method not implemented")

"""
Train the model on a dataset.

For each batch in the dataset, do:
1. Compute the gradient of loss;
2. Update model parameters using the optimizers set for each parameters.

It returns batch averaged loss in the end.
"""
function train!(model::NeuralModel, dataloader; kwargs...)
    loss_list = []
    # Pop `:epoch` out from from `kwargs` if exist
    local epoch = nothing
    if :epoch in keys(kwargs)
        epoch = kwargs[:epoch]
        kwargs = Dict(kwargs...)
        # NOTE: the line below is uncommented because `:epoch` is useful for annealing
        #       so we need to deal with possible unwanted keyword agurments manually
        # pop!(kwargs, :epoch)
    end
    for (i, data_batch) in enumerate(dataloader)
        # Calculate the iteration till now if given the epoch
        if epoch != nothing
            kwargs[:iter] = length(dataloader) * (epoch - 1) + i
        end
        losstape = Knet.@diff model(data_batch, Val(:true); kwargs...)
        update!(model, losstape)
        push!(loss_list, Knet.value(losstape))
    end
    return mean(loss_list)
end

"""
Evaluate the model on a dataset.

It returns batch averaged metric value in the end.
"""
function evaluate(model::NeuralModel, dataloader; kargs...)
    loss_list = []
    for data_batch in dataloader
        push!(loss_list, model(data_batch, Val(:false); kargs...)[1])
    end
    return mean(loss_list)
end

export initoptim!, update!, DynamicAdam, numparams
export AbstractTrainable, AbstractLayer, StaticLayer, StochasticLayer
export NeuralModel, train!, evaluate

include("activations.jl")
export log1pexp, logexpm1, softplus, invsoftplus, leaky_relu
include("layers.jl")
export Dense, LazyDense, DynamicIn, DynamicOut, Chain
include("stochastic_layers.jl")
# TODO: somehow make below automatic
export GaussianDense, GaussianLazyDense, GaussianDynamicIn, GaussianDynamicOut
export GaussianLogVarDense, GaussianLogVarLazyDense, GaussianLogVarDynamicIn, GaussianLogVarDynamicOut
export KumaraswamyDense, KumaraswamyLazyDense, KumaraswamyDynamicIn, KumaraswamyDynamicOut
export BernoulliDense, BernoulliLazyDense, BernoulliDynamicIn, BernoulliDynamicOut
export BernoulliLogitDense, BernoulliLogitLazyDense, BernoulliLogitDynamicIn, BernoulliLogitDynamicOut
export GumbelBernoulliLogitDense, GumbelBernoulliLogitLazyDense, GumbelBernoulliLogitDynamicIn, GumbelBernoulliLogitDynamicOut
include("sbc.jl")
export MeanFieldSBC, StructuredSBC

include("rho.jl")
export Rho, updategrad!


#= Custom calibration configuration file. =#

using Distributions
using StatsBase
using LinearAlgebra
using CalibrateEDMF
using CalibrateEDMF.ModelTypes
using CalibrateEDMF.DistributionUtils
using CalibrateEDMF.ReferenceModels
using CalibrateEDMF.ReferenceStats
using CalibrateEDMF.LESUtils
using CalibrateEDMF.TurbulenceConvectionUtils
const src_dir = dirname(pathof(CalibrateEDMF))
using CalibrateEDMF.HelperFuncs
# Import EKP modules
using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions
using JLD2

namelist_args = [
    ("time_stepping", "t_max", 2.0 * 3600),
    ("time_stepping", "dt_max", 30.0),
    ("time_stepping", "dt_min", 20.0),
    ("stats_io", "frequency", 120.0),
    ("grid", "dz", 150.0),
    ("grid", "nz", 20),
]

# Cases defined as structs for quick access to default configs
struct Bomex end

function get_config()
    config = Dict()
    # Flags for saving output data
    config["output"] = get_output_config()
    # Define regularization of inverse problem
    config["regularization"] = get_regularization_config()
    # Define reference used in the inverse problem
    config["reference"] = get_reference_config(Bomex())
    # Define the parameter priors
    config["prior"] = get_prior_config()
    # Define the kalman process
    config["process"] = get_process_config()
    # Define the SCM static configuration
    config["scm"] = get_scm_config()
    return config
end

function get_output_config()
    config = Dict()
    config["outdir_root"] = mktempdir()
    config["overwrite_scm_file"] = false # Flag for overwritting SCM input file
    return config
end

function get_regularization_config()
    config = Dict()
    config["perform_PCA"] = true # Performs PCA on data
    config["variance_loss"] = 1.0e-2 # Variance truncation level in PCA
    config["normalize"] = true  # whether to normalize data by pooled variance
    config["tikhonov_mode"] = "relative" # Tikhonov regularization
    config["tikhonov_noise"] = 1.0e-4 # Tikhonov regularization
    config["dim_scaling"] = false # Dimensional scaling of the loss
    return config
end

function get_process_config()
    config = Dict()
    config["N_iter"] = 2
    config["N_ens"] = 5
    config["algorithm"] = "Inversion" # "Sampler", "Unscented"
    return config
end

function get_reference_config(::Bomex)
    config = Dict()
    config["case_name"] = ["Bomex"]
    # Flag to indicate whether reference data is from a perfect model (i.e. SCM instead of LES)
    config["y_reference_type"] = SCM()
    config["Σ_reference_type"] = SCM()
    config["y_names"] = [["thetal_mean", "qt_mean"]]
    ref_root_dir = mktempdir()
    config["y_dir"] = [joinpath(ref_root_dir, "Output.Bomex.ref")]
    config["scm_suffix"] = ["scm"]
    config["scm_parent_dir"] = [ref_root_dir]
    config["t_start"] = [0.0]
    config["t_end"] = [2.0 * 3600]
    config["namelist_args"] = repeat([namelist_args], 2)
    return config
end

function get_prior_config()
    config = Dict()
    config["constraints"] =
        Dict("entrainment_factor" => [bounded(0.0, 0.5)], "detrainment_factor" => [bounded(0.0, 0.5)])
    config["unconstrained_σ"] = 0.5
    return config
end

function get_scm_config()
    config = Dict()
    config["namelist_args"] = nothing
    return config
end


module ParametricOptInterface

using MathOptInterface
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const POI = ParametricOptInterface

"""
    Parameter(Float64)

The `Parameter` structure stores the numerical value associated to a given parameter.
# Example:
```julia-repl
julia> ParametricOptInterface.Parameter(5)
ParametricOptInterface.Parameter(5)
```
"""
struct Parameter <: MOI.AbstractScalarSet
    val::Float64
end

"""
    ParametricOptimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer

Declares a `ParametricOptimizer`, which allows the handling of parameters in a optimization model.
# Example:
```julia-repl
julia> ParametricOptInterface.ParametricOptimizer(GLPK.Optimizer())
ParametricOptInterface.ParametricOptimizer{Float64,GLPK.Optimizer}
```
"""
mutable struct ParametricOptimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer
    optimizer::OT
    parameters::Dict{MOI.VariableIndex, T}
    updated_parameters::Dict{MOI.VariableIndex, T}
    variables::Dict{MOI.VariableIndex, MOI.VariableIndex}
    last_index_added::Int
    affine_constraint_cache::Dict{MOI.ConstraintIndex, Array{MOI.ScalarAffineTerm{Float64},1}}
    quadratic_constraint_cache_pv::Dict{MOI.ConstraintIndex, Array{MOI.ScalarQuadraticTerm{Float64},1}}
    quadratic_constraint_cache_pp::Dict{MOI.ConstraintIndex, Array{MOI.ScalarQuadraticTerm{Float64},1}}
    quadratic_constraint_cache_pc::Dict{MOI.ConstraintIndex, Array{MOI.ScalarAffineTerm{Float64},1}}
    quadratic_constraint_variables_associated_to_parameters_cache::Dict{MOI.ConstraintIndex, Array{MOI.ScalarAffineTerm{Float64},1}}
    quadratic_added_cache::Dict{MOI.ConstraintIndex, MOI.ConstraintIndex}
    last_quad_add_added::Int
    affine_objective_cache::Array{MOI.ScalarAffineTerm{T},1}
    quadratic_objective_cache_pv::Array{MOI.ScalarQuadraticTerm{T},1}
    quadratic_objective_cache_pp::Array{MOI.ScalarQuadraticTerm{T},1}
    quadratic_objective_cache_pc::Array{MOI.ScalarAffineTerm{T},1}
    quadratic_objective_variables_associated_to_parameters_cache::Array{MOI.ScalarAffineTerm{T},1}
    function ParametricOptimizer(optimizer::OT) where OT
        new{Float64, OT}(
            optimizer,
            Dict{MOI.VariableIndex, Float64}(),
            Dict{MOI.VariableIndex, Float64}(),
            Dict{MOI.VariableIndex, MOI.VariableIndex}(),
            0,
            Dict{MOI.ConstraintIndex, Array{MOI.ScalarAffineTerm{Float64},1}}(),
            Dict{MOI.ConstraintIndex, Array{MOI.ScalarQuadraticTerm{Float64},1}}(),
            Dict{MOI.ConstraintIndex, Array{MOI.ScalarQuadraticTerm{Float64},1}}(),
            Dict{MOI.ConstraintIndex, Array{MOI.ScalarAffineTerm{Float64},1}}(),
            Dict{MOI.ConstraintIndex, Array{MOI.ScalarAffineTerm{Float64},1}}(),
            Dict{MOI.ConstraintIndex, MOI.ConstraintIndex}(),
            0,
            Array{MOI.ScalarAffineTerm{Float64},1}(),
            Array{MOI.ScalarQuadraticTerm{Float64},1}(),
            Array{MOI.ScalarQuadraticTerm{Float64},1}(),
            Array{MOI.ScalarAffineTerm{Float64},1}(),
            Array{MOI.ScalarAffineTerm{Float64},1}()
        )
    end
end

function MOI.add_variable(model::ParametricOptimizer)
    model.last_index_added += 1
    v_p = MOI.VariableIndex(model.last_index_added)
    v = MOI.add_variable(model.optimizer)
    model.variables[v_p] = v
    return v_p
end

"""
    MOI.add_constrained_variable(model::ParametricOptimizer, set::Parameter)

Adds a parameter, that is, a variable parameterized to the value provided in `set`, to the `model` specified.
The `model` must be a `ParametricOptInterface.ParametricOptimizer` model.
Returns the MOI.VariableIndex of the parameterized variable and the MOI.ConstraintIndex associated.

#Example:
```julia-repl
julia> w, cw = MOI.add_constrained_variable(optimizer, POI.Parameter(0))
(MathOptInterface.VariableIndex(1), MathOptInterface.ConstraintIndex{MathOptInterface.SingleVariable,ParametricOptInterface.Parameter}(1))
```
"""
function MOI.add_constrained_variable(model::ParametricOptimizer, set::Parameter)
    model.last_index_added += 1
    p = MOI.VariableIndex(model.last_index_added)
    model.parameters[p] = set.val
    return p, MOI.ConstraintIndex{MOI.SingleVariable, typeof(set)}(model.last_index_added)
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.SingleVariable, set::MOI.AbstractScalarSet)
    if haskey(model.parameters, f.variable)
        error("Cannot constrain a parameter")
   elseif haskey(model.variables, f.variable)
       return MOI.add_constraint(model.optimizer, f, set)
   else
       error("Variable not in the model")
   end
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.ScalarAffineFunction{T}, set::MOI.AbstractScalarSet) where T
    if !any(haskey(model.parameters, f.terms[i].variable_index) for i = 1:length(f.terms))
        return MOI.add_constraint(model.optimizer, f, set)
    else
        vars = MOI.ScalarAffineTerm{T}[]
        params = MOI.ScalarAffineTerm{T}[]

        for i in f.terms
            if haskey(model.variables, i.variable_index)
                push!(vars, i)
            elseif haskey(model.parameters, i.variable_index)
                push!(params, i)
            else
                error("Constraint uses a variable that is not in the model")
            end
        end
        param_constant = 0
        for j in params
            param_constant += j.coefficient * model.parameters[j.variable_index]
        end
        fvar = MOI.ScalarAffineFunction(vars, f.constant + param_constant)
        ci = MOIU.normalize_and_add_constraint(model.optimizer, fvar, set)
        model.affine_constraint_cache[ci] = params
        return ci
    end
end

function update_constant!(s::MOI.LessThan{T}, val) where T
    MOI.LessThan{T}(s.upper - val)
end

function update_constant!(s::MOI.GreaterThan{T}, val) where T
    MOI.GreaterThan{T}(s.lower - val)
end

function update_constant!(s::MOI.EqualTo{T}, val) where T
    MOI.EqualTo{T}(s.value - val)
end

function MOI.get(model::ParametricOptimizer, attr::MOI.VariablePrimal, v::MOI.VariableIndex)
    if haskey(model.parameters, v)
         return model.parameters[v]
    elseif haskey(model.variables, v)
        return MOI.get(model.optimizer, attr, model.variables[v])
    else
        error("Variable not in the model")
    end
end

"""
    MOI.set(model::ParametricOptimizer, ::MOI.ConstraintSet, cp::MOI.ConstraintIndex{MOI.SingleVariable, P}, set::P) where {P <: Parameter}

Sets the parameter to a given value, using its `MOI.ConstraintIndex` as reference.

#Example:
```julia-repl
julia> MOI.set(optimizer, MOI.ConstraintSet(), cw, POI.Parameter(2.0))
2.0
```
"""
function MOI.set(model::ParametricOptimizer, ::MOI.ConstraintSet, cp::MOI.ConstraintIndex{MOI.SingleVariable, P}, set::P) where {P <: Parameter}
    p = MOI.VariableIndex(cp.value)
    if haskey(model.parameters, p)
        return model.updated_parameters[p] = set.val
    else
        error("Parameter not in the model")
    end
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveFunction{F}, f::F) where {F <: MOI.ScalarAffineFunction{T}} where T

    # clear previously defined objetive function cache
    model.affine_objective_cache = Array{MOI.ScalarAffineTerm{T},1}()
    model.quadratic_objective_cache_pv = Array{MOI.ScalarQuadraticTerm{T},1}()
    model.quadratic_objective_cache_pp = Array{MOI.ScalarQuadraticTerm{T},1}()
    model.quadratic_objective_cache_pc = Array{MOI.ScalarQuadraticTerm{T},1}()

    if !any(haskey(model.parameters, f.terms[i].variable_index) for i = 1:length(f.terms))
        MOI.set(model.optimizer, attr, f)
        return
    else
        vars = MOI.ScalarAffineTerm{T}[]
        params = MOI.ScalarAffineTerm{T}[]
        for i in f.terms
            if haskey(model.variables, i.variable_index)
                push!(vars, i)
            elseif haskey(model.parameters, i.variable_index)
                push!(params, i)
            else
                error("Objective function uses a variable that is not in the model")
            end
        end
        param_constant = 0
        for j in params
            param_constant += j.coefficient * model.parameters[j.variable_index]
        end
        fvar = MOI.ScalarAffineFunction(vars, f.constant + param_constant)
        co = MOI.set(model.optimizer, attr, fvar)
        model.affine_objective_cache = params
        return
    end
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveFunction{F}, v::F) where {F <: MOI.SingleVariable}
    if haskey(model.parameters, v)
        error("Cannot use a parameter as objective function alone")
    elseif haskey(model.variables, v)
        return MOI.set(model.optimizer, attr, MOI.SingleVariable(model.variables[v.variable]))
    else
        error("Variable not in the model")
    end
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    return MOI.set(model.optimizer, attr, sense)
end

function MOI.get(model::ParametricOptimizer, attr::T) where {
                                                                T <: Union{
                                                                    MOI.TerminationStatus,
                                                                    MOI.ObjectiveValue,
                                                                    MOI.PrimalStatus
                                                                    }
                                                            }
    return MOI.get(model.optimizer, attr)
end

function MOI.set(model::ParametricOptimizer, attr::T, bool::Bool) where {T <: MOI.Silent}
    MOI.set(model.optimizer, T(), bool)
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.VectorOfVariables, set::MOI.AbstractVectorSet) where T
    if any(haskey(model.parameters, f.variables) for i = 1:length(f.variables))
        error("VectorOfVariables does not allow parameters")
    else
        return MOI.add_constraint(model.optimizer, f, set)
    end
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.VectorAffineFunction{T}, set::MOI.AbstractVectorSet) where T
    if any(haskey(model.parameters, f.terms[i].variable_index) for i = 1:length(f.terms))
        error("VectorAffineFunction does not allow parameters")
    else
        return MOI.add_constraint(model.optimizer, f, set)
    end
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.ScalarQuadraticFunction{T}, set::MOI.AbstractScalarSet) where T

    if (!any(haskey(model.parameters, f.affine_terms[i].variable_index) for i = 1:length(f.affine_terms)) &&
        !any(haskey(model.parameters, f.quadratic_terms[j].variable_index_1) for j = 1:length(f.quadratic_terms)) &&
        !any(haskey(model.parameters, f.quadratic_terms[j].variable_index_2) for j = 1:length(f.quadratic_terms)))

        return MOI.add_constraint(model.optimizer, f, set)

    else

        quad_params = MOI.ScalarQuadraticTerm{T}[] #outside declaration so it has default value
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[] #outside declaration so it has default value; parameter as variable_index_1
        aux_variables_associated_to_parameters =  MOI.VariableIndex[] #outside declaration so it has default value

        if any(haskey(model.parameters, f.quadratic_terms[i].variable_index_1) for i = 1:length(f.quadratic_terms)) ||
            any(haskey(model.parameters, f.quadratic_terms[i].variable_index_2) for i = 1:length(f.quadratic_terms))

            quad_terms = MOI.ScalarQuadraticTerm{T}[]

            for i in f.quadratic_terms
                if haskey(model.variables, i.variable_index_1) && haskey(model.variables, i.variable_index_2)
                    push!(quad_terms, i) # if there are only variables, it remains a quadratic term

                elseif haskey(model.parameters, i.variable_index_1) && haskey(model.variables, i.variable_index_2)
                    # This is the case when i.variable_index_1 is a parameter and i.variable_index_2 is a variable.
                    # Thus, it creates an affine term. Convention: param as 1, var as 2
                    aux = MOI.ScalarQuadraticTerm(i.coefficient, i.variable_index_1, i.variable_index_2)
                    push!(quad_aff_vars, aux)
                    push!(aux_variables_associated_to_parameters, i.variable_index_2)

                elseif haskey(model.variables, i.variable_index_1) && haskey(model.parameters, i.variable_index_2)
                    # Check convention defined above
                    aux = MOI.ScalarQuadraticTerm(i.coefficient, i.variable_index_2, i.variable_index_1)
                    push!(quad_aff_vars, aux)
                    push!(aux_variables_associated_to_parameters, i.variable_index_1)

                elseif haskey(model.parameters, i.variable_index_1) && haskey(model.parameters, i.variable_index_2)
                    # This is the case where both variable_index_1,2 are actually parameters
                    aux = MOI.ScalarQuadraticTerm(i.coefficient, i.variable_index_1, i.variable_index_2)
                    push!(quad_params, aux)

                else
                    error("Constraint uses a variable that is not in the model")
                end
            end

        else
            quad_terms = f.quadratic_terms
        end

        aff_params = MOI.ScalarAffineTerm{T}[] #outside declaration so it has default value
        variables_associated_to_parameters =  MOI.ScalarAffineTerm{T}[] #outside declaration so it has default value

        if any(haskey(model.parameters, f.affine_terms[i].variable_index) for i = 1:length(f.affine_terms))

            aff_vars = MOI.ScalarAffineTerm{T}[]

            for i in f.affine_terms
                if haskey(model.variables, i.variable_index)
                    push!(aff_vars, i)
                    if i.variable_index in aux_variables_associated_to_parameters
                        push!(variables_associated_to_parameters, i)
                    end
                elseif haskey(model.parameters, i.variable_index)
                    push!(aff_params, i)
                else
                    error("Constraint uses a variable that is not in the model")
                end
            end

        else
            aff_vars = f.affine_terms
        end

        aff_terms = MOI.ScalarAffineTerm{T}[]

        for i in quad_aff_vars
            aux = MOI.ScalarAffineTerm(i.coefficient * model.parameters[i.variable_index_1], i.variable_index_2)
            push!(aff_terms, aux)
        end

        aff_terms = vcat(aff_terms, aff_vars)

        const_term = f.constant

        for j in aff_params
            const_term += j.coefficient * model.parameters[j.variable_index]
        end

        for j in quad_params
            const_term += j.coefficient * model.parameters[j.variable_index_1] * model.parameters[j.variable_index_2]
        end

        f_quad = if !isempty(quad_terms)
            MOI.ScalarQuadraticFunction(
                        aff_terms,
                        quad_terms,
                        const_term
                    )
        else
            MOI.ScalarAffineFunction(
                        aff_terms,
                        const_term
                    )
        end

        model.last_quad_add_added += 1
        ci = MOIU.normalize_and_add_constraint(model.optimizer, f_quad, set)
        new_ci = MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, typeof(set)}(model.last_quad_add_added)
        model.quadratic_added_cache[new_ci] = ci

        if !isempty(quad_aff_vars)
            model.quadratic_constraint_cache_pv[new_ci] = quad_aff_vars
        end

        if !isempty(quad_params)
            model.quadratic_constraint_cache_pp[new_ci] = quad_params
        end

        if !isempty(aff_params)
            model.quadratic_constraint_cache_pc[new_ci] = aff_params
        end

        if !isempty(variables_associated_to_parameters)
            model.quadratic_constraint_variables_associated_to_parameters_cache[new_ci] = variables_associated_to_parameters
        end

        return ci
    end

end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveFunction{F}, f::F) where {F <: MOI.ScalarQuadraticFunction{T}} where T

    # clear previously defined objetive function cache
    model.affine_objective_cache = Array{MOI.ScalarAffineTerm{T},1}()
    model.quadratic_objective_cache_pv = Array{MOI.ScalarQuadraticTerm{T},1}()
    model.quadratic_objective_cache_pp = Array{MOI.ScalarQuadraticTerm{T},1}()
    model.quadratic_objective_cache_pc = Array{MOI.ScalarQuadraticTerm{T},1}()

    if (!any(haskey(model.parameters, f.affine_terms[i].variable_index) for i = 1:length(f.affine_terms)) &&
        !any(haskey(model.parameters, f.quadratic_terms[j].variable_index_1) for j = 1:length(f.quadratic_terms)) &&
        !any(haskey(model.parameters, f.quadratic_terms[j].variable_index_2) for j = 1:length(f.quadratic_terms)))

        MOI.set(model.optimizer, attr, f)
        return
    else

        quad_params = MOI.ScalarQuadraticTerm{T}[] #outside declaration so it has default value
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[] #outside declaration so it has default value; parameter as variable_index_1
        aux_variables_associated_to_parameters =  MOI.VariableIndex[] #outside declaration so it has default value

        if (any(haskey(model.parameters, f.quadratic_terms[i].variable_index_1) for i = 1:length(f.quadratic_terms)) ||
             any(haskey(model.parameters, f.quadratic_terms[i].variable_index_2) for i = 1:length(f.quadratic_terms)))

            quad_terms = MOI.ScalarQuadraticTerm{T}[]

            for i in f.quadratic_terms
                if haskey(model.variables, i.variable_index_1) && haskey(model.variables, i.variable_index_2)
                    push!(quad_terms, i) # if there are only variables, it remains a quadratic term

                elseif haskey(model.parameters, i.variable_index_1) && haskey(model.variables, i.variable_index_2)
                    # This is the case when i.variable_index_1 is a parameter and i.variable_index_2 is a variable.
                    # Thus, it creates an affine term. Convention: param as 1, var as 2
                    aux = MOI.ScalarQuadraticTerm(i.coefficient, i.variable_index_1, i.variable_index_2)
                    push!(quad_aff_vars, aux)
                    push!(aux_variables_associated_to_parameters, i.variable_index_2)

                elseif haskey(model.variables, i.variable_index_1) && haskey(model.parameters, i.variable_index_2)
                    # Check convention defined above
                    aux = MOI.ScalarQuadraticTerm(i.coefficient, i.variable_index_2, i.variable_index_1)
                    push!(quad_aff_vars, aux)
                    push!(aux_variables_associated_to_parameters, i.variable_index_1)

                elseif haskey(model.parameters, i.variable_index_1) && haskey(model.parameters, i.variable_index_2)
                    # This is the case where both variable_index_1,2 are actually parameters
                    aux = MOI.ScalarQuadraticTerm(i.coefficient, i.variable_index_1, i.variable_index_2)
                    push!(quad_params, aux)

                else
                    error("Constraint uses a variable that is not in the model")
                end
            end

        else
            quad_terms = f.quadratic_terms
        end

        aff_params = MOI.ScalarAffineTerm{T}[] #outside declaration so it has default value
        variables_associated_to_parameters =  MOI.ScalarAffineTerm{T}[] #outside declaration so it has default value

        if any(haskey(model.parameters, f.affine_terms[i].variable_index) for i = 1:length(f.affine_terms))

            aff_vars = MOI.ScalarAffineTerm{T}[]

            for i in f.affine_terms
                if haskey(model.variables, i.variable_index)
                    push!(aff_vars, i)
                    if i.variable_index in aux_variables_associated_to_parameters
                        push!(variables_associated_to_parameters, i)
                    end
                elseif haskey(model.parameters, i.variable_index)
                    push!(aff_params, i)
                else
                    error("Constraint uses a variable that is not in the model")
                end
            end

        else
            aff_vars = f.affine_terms
        end

        aff_terms = MOI.ScalarAffineTerm{T}[]

        for i in quad_aff_vars
            aux = MOI.ScalarAffineTerm(i.coefficient * model.parameters[i.variable_index_1], i.variable_index_2)
            push!(aff_terms, aux)
        end

        aff_terms = vcat(aff_terms, aff_vars)

        const_term = f.constant

        for j in aff_params
            const_term += j.coefficient * model.parameters[j.variable_index]
        end

        for j in quad_params
            const_term += j.coefficient * model.parameters[j.variable_index_1] * model.parameters[j.variable_index_2]
        end

        f_quad = if !isempty(quad_terms)
            MOI.ScalarQuadraticFunction(
                        aff_terms,
                        quad_terms,
                        const_term
                    )
        else
            MOI.ScalarAffineFunction(
                        aff_terms,
                        const_term
                    )
        end

        MOI.set(model.optimizer, attr, f_quad)

        model.quadratic_objective_cache_pv = quad_aff_vars
        model.quadratic_objective_cache_pp = quad_params
        model.quadratic_objective_cache_pc = aff_params
        model.quadratic_objective_variables_associated_to_parameters_cache = variables_associated_to_parameters

        return
    end
end

function MOI.optimize!(model::ParametricOptimizer)
    if !isempty(model.updated_parameters)

        for (ci, fparam) in model.affine_constraint_cache
            param_constant = 0
            for j in fparam
                if haskey(model.updated_parameters, j.variable_index)
                    param_old = model.parameters[j.variable_index]
                    param_new = model.updated_parameters[j.variable_index]
                    aux = param_new - param_old
                    param_constant += j.coefficient * aux
                end
            end
            if param_constant != 0
                set = MOI.get(model.optimizer, MOI.ConstraintSet(), ci)
                set = POI.update_constant!(set, param_constant)
                MOI.set(model.optimizer, MOI.ConstraintSet(), ci, set)
            end
        end

        if !isempty(model.affine_objective_cache)
            objective_constant = 0
            for j in model.affine_objective_cache
                if haskey(model.updated_parameters, j.variable_index)
                    param_old = model.parameters[j.variable_index]
                    param_new = model.updated_parameters[j.variable_index]
                    aux = param_new - param_old
                    objective_constant += j.coefficient * aux
                end
            end
            if objective_constant != 0
                F = MOI.get(model.optimizer, MOI.ObjectiveFunctionType())
                f = MOI.get(model.optimizer, MOI.ObjectiveFunction{F}())
                fvar = MOI.ScalarAffineFunction(f.terms, f.constant + objective_constant)
                MOI.set(model.optimizer,MOI.ObjectiveFunction{F}(), fvar)
            end
        end

        for (ci, fparam) in model.quadratic_constraint_cache_pc
            param_constant = 0
            for j in fparam
                if haskey(model.updated_parameters, j.variable_index)
                    param_old = model.parameters[j.variable_index]
                    param_new = model.updated_parameters[j.variable_index]
                    aux = param_new - param_old
                    param_constant += j.coefficient * aux
                end
            end
            if param_constant != 0
                set = MOI.get(model.optimizer, MOI.ConstraintSet(), model.quadratic_added_cache[ci])
                set = POI.update_constant!(set, param_constant)
                MOI.set(model.optimizer, MOI.ConstraintSet(), model.quadratic_added_cache[ci], set)
            end
        end

        if !isempty(model.quadratic_objective_cache_pc)
            objective_constant = 0
            for j in model.quadratic_objective_cache_pc
                if haskey(model.updated_parameters, j.variable_index)
                    param_old = model.parameters[j.variable_index]
                    param_new = model.updated_parameters[j.variable_index]
                    aux = param_new - param_old
                    objective_constant += j.coefficient * aux
                end
            end
            if objective_constant != 0
                F = MOI.get(model.optimizer, MOI.ObjectiveFunctionType())
                f = MOI.get(model.optimizer, MOI.ObjectiveFunction{F}())
                fvar = if F == MathOptInterface.ScalarAffineFunction{Float64}
                    MOI.ScalarAffineFunction(f.terms, f.constant + objective_constant)
                else
                    MOI.ScalarQuadraticFunction(f.affine_terms, f.quadratic_terms, f.constant + objective_constant)
                end
                MOI.set(model.optimizer,MOI.ObjectiveFunction{F}(), fvar)
            end
        end

        for (ci, fparam) in model.quadratic_constraint_cache_pp
            param_constant = 0
            for j in fparam
                if haskey(model.updated_parameters, j.variable_index_1) && haskey(model.updated_parameters, j.variable_index_2)
                    param_new_1 = model.updated_parameters[j.variable_index_1]
                    param_new_2 = model.updated_parameters[j.variable_index_2]
                    param_old_1 = model.parameters[j.variable_index_1]
                    param_old_2 = model.parameters[j.variable_index_2]
                    param_constant += j.coefficient * ((param_new_1 * param_new_2)-(param_old_1 * param_old_2))
                elseif haskey(model.updated_parameters, j.variable_index_1)
                    param_new_1 = model.updated_parameters[j.variable_index_1]
                    param_old_1 = model.parameters[j.variable_index_1]
                    param_old_2 = model.parameters[j.variable_index_2]
                    param_constant += j.coefficient * param_old_2 * (param_new_1 - param_old_1)
                elseif haskey(model.updated_parameters, j.variable_index_2)
                    param_new_2 = model.updated_parameters[j.variable_index_2]
                    param_old_1 = model.parameters[j.variable_index_1]
                    param_old_2 = model.parameters[j.variable_index_2]
                    param_constant += j.coefficient * param_old_1 * (param_new_2 - param_old_2)
                end
            end
            if param_constant != 0
                set = MOI.get(model.optimizer, MOI.ConstraintSet(), model.quadratic_added_cache[ci])
                set = POI.update_constant!(set, param_constant)
                MOI.set(model.optimizer, MOI.ConstraintSet(), model.quadratic_added_cache[ci], set)
            end
        end

        if !isempty(model.quadratic_objective_cache_pp)
            objective_constant = 0
            for j in model.quadratic_objective_cache_pp
                if haskey(model.updated_parameters, j.variable_index_1) && haskey(model.updated_parameters, j.variable_index_2)
                    param_new_1 = model.updated_parameters[j.variable_index_1]
                    param_new_2 = model.updated_parameters[j.variable_index_2]
                    param_old_1 = model.parameters[j.variable_index_1]
                    param_old_2 = model.parameters[j.variable_index_2]
                    objective_constant += j.coefficient * ((param_new_1 * param_new_2)-(param_old_1 * param_old_2))
                elseif haskey(model.updated_parameters, j.variable_index_1)
                    param_new_1 = model.updated_parameters[j.variable_index_1]
                    param_old_1 = model.parameters[j.variable_index_1]
                    param_old_2 = model.parameters[j.variable_index_2]
                    objective_constant += j.coefficient * param_old_2 * (param_new_1 - param_old_1)
                elseif haskey(model.updated_parameters, j.variable_index_2)
                    param_new_2 = model.updated_parameters[j.variable_index_2]
                    param_old_1 = model.parameters[j.variable_index_1]
                    param_old_2 = model.parameters[j.variable_index_2]
                    objective_constant += j.coefficient * param_old_1 * (param_new_2 - param_old_2)
                end
            end
            if objective_constant != 0
                F = MOI.get(model.optimizer, MOI.ObjectiveFunctionType())
                f = MOI.get(model.optimizer, MOI.ObjectiveFunction{F}())
                fvar = if F == MathOptInterface.ScalarAffineFunction{Float64}
                    MOI.ScalarAffineFunction(f.terms, f.constant + objective_constant)
                else
                    MOI.ScalarQuadraticFunction(f.affine_terms, f.quadratic_terms, f.constant + objective_constant)
                end
                MOI.set(model.optimizer,MOI.ObjectiveFunction{F}(), fvar)
            end
        end

        constraint_aux_dict = Dict{Any,Any}()

        for (ci, fparam) in model.quadratic_constraint_cache_pv
            for j in fparam
                if haskey(model.updated_parameters, j.variable_index_1)
                    coef = j.coefficient
                    param_new = model.updated_parameters[j.variable_index_1]
                    if haskey(constraint_aux_dict, (ci, j.variable_index_2))
                        constraint_aux_dict[(ci, j.variable_index_2)] += param_new*coef
                    else
                        constraint_aux_dict[(ci, j.variable_index_2)] = param_new*coef
                    end
                end
            end
        end

        for (ci, fparam) in model.quadratic_constraint_variables_associated_to_parameters_cache
            for j in fparam.terms
                coef = j.coefficient
                if haskey(constraint_aux_dict, (ci, j.variable_index))#
                    constraint_aux_dict[(ci, j.variable_index)] += coef
                else
                    constraint_aux_dict[(ci, j.variable_index)] = coef
                end
            end
        end

        for (key, value) in constraint_aux_dict
            MOI.modify(model.optimizer, key[1], MOI.ScalarCoefficientChange(key[2], value))
        end

        objective_aux_dict = Dict{Any,Any}()

        if !isempty(model.quadratic_objective_cache_pv)
            for j in model.quadratic_objective_cache_pv
                if haskey(model.updated_parameters, j.variable_index_1)
                    coef = j.coefficient
                    param_new = model.updated_parameters[j.variable_index_1]
                    if haskey(objective_aux_dict, (j.variable_index_2))
                        objective_aux_dict[(j.variable_index_2)] += param_new*coef
                    else
                        objective_aux_dict[(j.variable_index_2)] = param_new*coef
                    end
                end
            end
        end

        for j in model.quadratic_objective_variables_associated_to_parameters_cache
            coef = j.coefficient
            if haskey(objective_aux_dict, j.variable_index)
                objective_aux_dict[j.variable_index] += coef
            else
                objective_aux_dict[j.variable_index] = coef
            end
        end

        F_pv = MOI.get(model.optimizer, MOI.ObjectiveFunctionType())

        for (key, value) in objective_aux_dict
            MOI.modify(model.optimizer, MOI.ObjectiveFunction{F_pv}(), MOI.ScalarCoefficientChange(key, value))
        end

        for (i, val) in model.updated_parameters
            model.parameters[i] = val
        end
        empty!(model.updated_parameters)
    end
    MOI.optimize!(model.optimizer)
end

end # module


export Combiner

# TODO: Have combiner store the locations
# of the uncombined and combined indices
# This can generalize to a Combiner that combines
# multiple set of indices, e.g. (i,j),(k,l) -> (a,b)
struct Combiner <: TensorStorage{Number}
  perm::Vector{Int}
  comb::Vector{Int}
  Combiner(perm::Vector{Int},comb::Vector{Int}) = new(perm,comb)
end

Combiner() = Combiner(Int[],Int[])

data(::Combiner) = error("Combiner storage has no data")

blockperm(C::Combiner) = C.perm
blockcomb(C::Combiner) = C.comb

Base.eltype(::Type{<:Combiner}) = Number

Base.eltype(::StoreT) where {StoreT<:Combiner} = eltype(StoreT)

Base.promote_rule(::Type{<:Combiner},
                  StorageT::Type{<:Dense}) = StorageT

#
# CombinerTensor (Tensor using Combiner storage)
#

const CombinerTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:Combiner}

combinedindex(T::CombinerTensor) = inds(T)[1]
uncombinedinds(T::CombinerTensor) = popfirst(inds(T))

blockperm(C::CombinerTensor) = blockperm(store(C))
blockcomb(C::CombinerTensor) = blockcomb(store(C))

Base.conj(T::CombinerTensor; always_copy = false) = T

function contraction_output(::TensorT1,
                            ::TensorT2,
                            indsR::IndsR) where {TensorT1<:CombinerTensor,
                                                 TensorT2<:DenseTensor,
                                                 IndsR}
  TensorR = contraction_output_type(TensorT1,TensorT2,IndsR)
  return similar(TensorR,indsR)
end

function contraction_output(T1::TensorT1,
                            T2::TensorT2,
                            indsR) where {TensorT1<:DenseTensor,
                                          TensorT2<:CombinerTensor}
  return contraction_output(T2,T1,indsR)
end

function contract!!(R::Tensor{<:Number,NR},
                    labelsR::NTuple{NR},
                    T1::CombinerTensor{<:Number,N1},
                    labelsT1::NTuple{N1},
                    T2::Tensor{<:Number,N2},
                    labelsT2::NTuple{N2}) where {NR,N1,N2}
  if N1 ≤ 1
    #println("identity")
    return R
  elseif N1 + N2 == NR
    error("Cannot perform outer product involving a combiner")
  elseif count_common(labelsT1,labelsT2) == 1 && N1 == 2
    # This is the case of index replacement
    ui = setdiff(labelsT1, labelsT2)[]
    newind = inds(T1)[findfirst(==(ui),labelsT1)]
    cpos1,cpos2 = intersect_positions(labelsT1,labelsT2)
    storeR = copy(store(T2))
    indsR = setindex(inds(T2),newind,cpos2)
    return tensor(storeR,indsR)
  elseif count_common(labelsT1,labelsT2) == 1 && length(inds(T1)) != 2
    # This is the case of uncombining
    cpos1,cpos2 = intersect_positions(labelsT1,labelsT2)
    storeR = copy(store(T2))
    indsC = deleteat(inds(T1),cpos1)
    indsR = insertat(inds(T2),indsC,cpos2)
    return tensor(storeR,indsR)
  elseif is_combiner(labelsT1,labelsT2)
    # This is the case of combining
    Alabels,Blabels = labelsT2,labelsT1
    final_labels    = contract_labels(Blabels, Alabels)
    final_labels_n  = contract_labels(labelsT1,labelsT2)
    indsR = inds(R)
    if final_labels != final_labels_n
      perm  = getperm(final_labels_n, final_labels)
      indsR = permute(inds(R), perm)
      labelsR = permute(labelsR, perm)
    end
    cpos1,cposR = intersect_positions(labelsT1,labelsR)
    labels_comb = deleteat(labelsT1,cpos1)
    vlR = [labelsR...]
    for (ii, li) in enumerate(labels_comb)
      insert!(vlR, cposR+ii, li)
    end
    deleteat!(vlR, cposR)
    labels_perm = tuple(vlR...)
    perm = getperm(labels_perm,labelsT2)
    T2p = reshape(R,permute(inds(T2),perm))
    permutedims!(T2p,T2,perm)
    R = reshape(T2p,indsR)
  end
  return R
end

function contract!!(R::Tensor{<:Number,NR},
                    labelsR::NTuple{NR},
                    T1::Tensor{<:Number,N1},
                    labelsT1::NTuple{N1},
                    T2::CombinerTensor{<:Number,N2},
                    labelsT2::NTuple{N2}) where {NR,N1,N2}
  return contract!!(R,labelsR,T2,labelsT2,T1,labelsT1)
end

function Base.show(io::IO,
                   mime::MIME"text/plain",
                   S::Combiner)
  println(io, "Permutation of blocks: ", S.perm)
  println(io, "Combination of blocks: ", S.comb)
end

function Base.show(io::IO,
                   mime::MIME"text/plain",
                   T::CombinerTensor)
  summary(io, T)
  println(io)
  show(io, mime, store(T))
end


using RecursiveArrayTools: ArrayPartition
using ComponentArrays: ComponentVector
using ..RingPolymers: RingPolymers

get_positions(u::ComponentVector) = u.r
get_velocities(u::ComponentVector) = u.v
get_flat_positions(u::ComponentVector) = @view get_positions(u)[:]
get_flat_velocities(u::ComponentVector) = @view get_velocities(u)[:]

# These must be kept since the symplectic integrators (VelocityVerlet) in DiffEq use ArrayPartitions.
get_positions(u::ArrayPartition) = u.x[2]
get_velocities(u::ArrayPartition) = u.x[1]

function get_quantum_subsystem end

get_mapping_positions(u::ComponentVector) = u.qmap
get_mapping_momenta(u::ComponentVector) = u.pmap
get_mapping_positions(u::ComponentVector, i) = @view get_mapping_positions(u)[:,i]
get_mapping_momenta(u::ComponentVector, i) = @view get_mapping_momenta(u)[:,i]


# using OpenAIGym
# import OpenAIGym
using Flux
using JuliaRL
using Random
import LearnBase.IntervalSet

# rand(rng, int_set::)
Base.rand(rng::AbstractRNG, s::IntervalSet{T}) where T <: AbstractVector = Float64[rand(rng) * (s.hi[i] - s.lo[i]) + s.lo[i] for i=1:length(s)]

function test_gym(num_episodes, num_steps, to_render::Bool=false)

    rng = MersenneTwister(1)
    env = GymEnv(:BipedalWalker, :v2; seed=1)
    println(typeof(env))
    for i ∈ 1:num_episodes
        T = 0
        R = 0.0
        action_set = get_actions(env)
        st = start!(env)
        for i in 1:num_steps
            action = rand(rng, action_set)
            step!(env, action)
            if to_render
                render(env)
            end
            R += get_reward(env)
            T += 1
        end

        @info("Episode $i finished after $T steps. Total reward: $R")
    end

    # You can explicitly close the connection for the python environment here. But will be taken care of by the garbage collector eventually.
    close(env)

end


"""
The basic voting function when every agent's got an equal voting power
and the candidate is being compared to the registry mean
"""
function simpleVote(benchmark, candidate, agents)
    rejectAgents = []
    proAgents = []
    quorum = length(agents)
    for agent in agents
        push!(evaluateCandidateByAgent(candidate, agent) >= benchmark ? proAgents : rejectAgents, agent)
    end
    [length(proAgents) > quorum/2, proAgents, rejectAgents]
end

"""
The voting function when every agent's got an equal voting power
but only agents who hold some tokens are eligible to participate
"""
function binaryTokenVote(benchmark, candidate, agents)
    rejectAgents = []
    proAgents = []
    quorum = 0
    for agent in agents
        if agent.balance > 0
            if (evaluateCandidateByAgent(candidate, agent) <= benchmark)
                push!(rejectAgents, agent)
            else
                push!(proAgents, agent)
            end
            quorum += 1
        end
    end
    [length(proAgents) > quorum/2, proAgents, rejectAgents]
end

"""
The voting function when every agent's got a voting power
proportional to the amount of tokens under possesion
"""
function proRataTokenVote(benchmark, candidate, agents)
    rejectAgents = []
    proAgents = []
    proTokens = 0
    quorum = 0
    for agent in agents
        if agent.balance > 0
            if (evaluateCandidateByAgent(candidate, agent) <= benchmark)
                push!(rejectAgents, agent)
            else
                proTokens += agent.balance
                push!(proAgents, agent)
            end
            quorum += agent.balance
        end
    end
    [proTokens > quorum/2, proAgents, rejectAgents]
end


#****************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#****************************************************************************

function getpotstr(pot)

if pot==0
    sp0 = "(T *u, T *rho, T *mu, T *eta, int *kappa, int nrho, int nmu, int neta, int nkappa, int ng)\n";
    sp1 = "(double *, double *, double *, double *, int*, int, int, int, int, int);\n";
    sp2 = "(float *, float *, float *, float *, int *, int, int, int, int, int);\n";
    sp3 = "(u, rho, mu, eta, kappa, nrho, nmu, neta, nkappa, ng);\n";
elseif pot==1
    sp0 = "(T *u, T *xi, T *qi, int *ti, int *ai, T *mu, T *eta, int *kappa, int dim, int ncq, int nmu, int neta, int nkappa, int ng)\n";
    sp1 = "(double *, double *, double *, int *, int *, double *, double *, int*, int, int, int, int, int, int);\n";
    sp2 = "(float *, float *, float *, int *, int *, float *, float *, int *, int, int, int, int, int, int);\n";
    sp3 = "(u, xi, qi, ti, ai, mu, eta, kappa, dim, ncq, nmu, neta, nkappa, ng);\n";
elseif pot==2
    sp0 = "(T *u, T *xij, T *qi, T *qj, int *ti, int *tj, int *ai, int *aj, T *mu, T *eta, int *kappa, int dim, int ncq, int nmu, int neta, int nkappa, int ng)\n";
    sp1 = "(double *, double *, double *, double *, int *, int *, int *, int *, double *, double *, int*, int, int, int, int, int, int);\n";
    sp2 = "(float *, float *, float *, float *, int *, int *, int *, int *, float *, float *, int *, int, int, int, int, int, int);\n";
    sp3 = "(u, xij, qi, qj, ti, tj, ai, aj, mu, eta, kappa, dim, ncq, nmu, neta, nkappa, ng);\n";
elseif pot==3
    sp0 = "(T *u, T *xij, T *xik, T *qi, T *qj, T *qk, int *ti, int *tj, int *tk, int *ai, int *aj, int *ak, T *mu, T *eta, int *kappa, int dim, int ncq, int nmu, int neta, int nkappa, int ng)\n";
    sp1 = "(double *, double *, double *, double *, double *, double *, int *, int *, int *, int *, int *, int *, double *, double *, int*, int, int, int, int, int, int);\n";
    sp2 = "(float *, float *, float *, float *, float *, float *, int *, int *, int *, int *, int *, int *, float *, float *, int *, int, int, int, int, int, int);\n";
    sp3 = "(u, xij, xik, qi, qj, qk, ti, tj, tk, ai, aj, ak, mu, eta, kappa, dim, ncq, nmu, neta, nkappa, ng);\n";
elseif pot==4
    sp0 = "(T *u, T *xij, T *xik, T *xil, T *qi, T *qj, T *qk, T *ql, int *ti, int *tj, int *tk, int *tl, int *ai, int *aj, int *ak, int *al, T *mu, T *eta, int *kappa, int dim, int ncq, int nmu, int neta, int nkappa, int ng)\n";
    sp1 = "(double *, double *, double *, double *, double *, double *, double *, double *, int *, int *, int *, int *, int *, int *, int *, int *, double *, double *, int*, int, int, int, int, int, int);\n";
    sp2 = "(float *, float *, float *, float *, float *, float *, float *, float *, int *, int *, int *, int *, int *, int *, int *, int *, float *, float *, int *, int, int, int, int, int, int);\n";
    sp3 = "(u, xij, xik, xil, qi, qj, qk, ql, ti, tj, tk, tl, ai, aj, ak, al, mu, eta, kappa, dim, ncq, nmu, neta, nkappa, ng);\n";
end

return sp0, sp1, sp2, sp3

end