text:
"""
`module Lsq`

This sub-modules contains code to convert a database `LsqDB` into
linear LSQ systems, apply suitable weights, solve the resulting systems
and return an IP (e.g. an NBodyIP).

The ordering of observations (entries of Y, row-indices of Psi) is given by
the loop ordering
```
for ct in configtypes
   for dt in datatypes  # observation types
      for dat in db.data_groups[ct]
      # or: db.kron_groups[ct][dt][:,:,...]
```
E.g., if the configtypes are "E", "F", then
```
Y = [E, E, E, E...., F, F, F, F, ...]
```
"""
module Lsq

import JuLIP, IPFitting, StatsBase

using StaticArrays
using JuLIP: vecs
using JuLIP: AbstractCalculator, Atoms
using JuLIP.Potentials: OneBody
using JuLIP.MLIPs: SumIP
using IPFitting: Dat, LsqDB, weighthook, observations,
                      observation, hasobservation, eval_obs, vec_obs
using IPFitting.Data: configtype
using IPFitting.DB: dbpath, _nconfigs, matrows

using LinearAlgebra: lmul!, Diagonal, qr, qr!, cond, norm, svd, I, pinv, mul!, cholesky
using InteractiveUtils: versioninfo
using LowRankApprox
using JuLIP.Utils
using JuLIP: JVecF
using GLMNet
using Roots
using ACE
using IterativeSolvers
using Optim
using Mmap
#using GenSPGL
#using SGDOptim

const Err = IPFitting.Errors

export lsqfit, onb

"""
`collect_observations(db::LsqDB, weights::Dict, Vref) -> Y, W, Icfg

Construct vectors of observations and weights from the lsq system specified
by `db`. The observations do not have the weights applied yet.

Weights are specified as follows:
```
   wh = weighthook(obskey, cfg)
   w = w1 * wh
```
where `w1` is obtained from the `weights` dictionary.
```
weights = Dict( "default" => Dict( "E" => 30.0, "F" => 1.0, "V" => 0.3 ),
                "bulk" => Dict( "E" => 100.0, "F" = 20.0, "V" => 0.0 ),
                ... )
```
To compute the weights, the code first looks for the `configtype` entry. If
it doesn't exist, then it reverts to the `default` entry. The weights obtained
from the `weights` dictionary is then multiplied by a

- The above procedure can be overwritten by an entry in `cfg.info["W"]`;
in this case the re-weighting via the `weighthook` is ignored as well.
- If `configtype(cfg) ∈ weights["ignore"]` then this configuration is skipped.
- If `obskey ∈ weights["ignore"]` then this observation is skipped
"""
function collect_observations(db::LsqDB,
                              weights::Dict,
                              Vref)#, scal_wgs = false )

   nrows = size(db.Ψ, 1)  # total number of observations if we collect
                          # everything in the database
   Y = zeros(Float64, nrows)
   W = zeros(Float64, nrows)
   Icfg = zeros(Int, nrows)
   for (obskey, dat, icfg) in observations(db)  # obskey ∈ {"E","F",...}; d::Dat
      if configtype(dat) in weights["ignore"] || obskey in weights["ignore"]
         continue
      end

      irows = matrows(dat, obskey)
      ct = configtype(dat)

      # ----- Observation ------
      obs = observation(obskey, dat)
      # If Vref != nothing the it is a calculator and we can subtract
      # its value => this allows us to fit from a reference potential
      if haskey(dat.info, "Vref")
         obs = obs - dat.info["Vref"][obskey]
      elseif Vref != nothing
         obs = obs - vec_obs(obskey, eval_obs(obskey, Vref, dat))
      end
      Y[irows] .= obs
      Icfg[irows] .= icfg
      # ------ Weights --------
      # modify the weights from extra information in the dat structure
      W[irows] .= _get_weights(weights,
                               weighthook(obskey, dat),
                               dat, obskey, obs)#, scal_wgs)
   end
   return Y, W, Icfg
end

"""
see documentation in `collect_observations`
"""
function _get_weights(weights, wh, dat, obskey, o)
   # initialise the weight to 0.0. If this isn't overwritten then it means
   # we will ignore this observation.
   w = 0.0

   # if there is a "W" entry in dat.info then this means all defaults
   # are over-written and we use this weight directly as is
   if haskey(dat.info, "W")
      # now check that this observation type (E, F, V) exists in dat.D["W"]
      # if yes return the corresponding weight, if not we ignore that
      # observation.
      if haskey(dat.info["W"], obskey)
         w = dat.info["W"][obskey]
      end
   else
      # if no "W" dict exists extract the correct weights from
      # the `weights` dictionary
      # first check what the current configtype has a separate entry
      # and otherwise use the default
      ct = configtype(dat)
      if haskey(weights, ct)
         cfgkey = ct
      else
         cfgkey = "default"
      end
      # now check whether there is an entry for the current observation  type
      if haskey(weights[cfgkey], obskey)
         w = weights[cfgkey][obskey] * wh
      end
   end

   # transform the weights into a vector (if necessary) and return
   if length(w) == 1
      return w * ones(length(o))
   elseif length(w) == length(o)
      return w
   end
   error("_get_weights: length(w) is neither 1 nor length(o)?!?!?")
end

function _regularise!(Ψ::Matrix{T}, Y::Vector{T}, basis, regularisers;
                      verbose=false, Ibasis = :) where {T}
   # TODO: incorporate Ibasis again !
   # assemble the regularisers
   Ψreg = Matrix{Float64}(undef, 0, size(Ψ, 2))
   Yreg = Float64[]
   for reg in regularisers
      if reg isa Matrix
         P = reg
         q = zeros(size(P, 1))
      else
         P, q = Matrix(reg, basis; verbose=verbose)
      end
      Ψreg = vcat(Ψreg, P)
      Yreg = vcat(Yreg, q)
   end
   # check they all have the correct size
   return vcat(Ψ, Ψreg), vcat(Y, Yreg)
end
#using DelimitedFiles

function _forceprecon(db, Y, weights)
   if !haskey(weights, "precon")
      return db.Ψ, Y
   end
   if isempty(weights["precon"])
      return db.Ψ, Y
   end

   Ψ = copy(db.Ψ)

   for (obskey, dat, icfg) in observations(db)  # obskey ∈ {"E","F",...}; d::Dat
      if obskey != "F"; continue; end
      ct = configtype(dat)
      if !haskey(weights["precon"], ct) && !haskey(weights["precon"], "default"); continue; end  ###### TOBE changed if we want to precondition only certain config_types

      #preconfun = ["precon"][ct]
      if haskey(weights["precon"], ct)
         P = weights["precon"][ct](dat.at)
      elseif haskey(weights["precon"], "default")
         P = weights["precon"]["default"](dat.at)
      end
      #   E(R + U) - E(R) - DE(R).U ~ U' D^2E(R) U ~ U' P U -> sqrt(U' P U)
      # DE(R+U) - DE(R) ~ D^2E(R) U
      # instead : DE(R+U) - DE(R) ~ P U
      #         || P^{-1/2} (DE(R+U) - DE(R)) || ~ || P^{1/2} U || = sqrt(U' P U)
      #         || P^{-1} (DE(R+U) - DE(R)) || ~ || U ||  => PROBABLY NOT THIS!!!!
      # all we need is
      #     sqrt(P)^2 = P
      #    || sqrt(P)^{-1} F ||^2 = F' sqrt{P}^{-2} F = F' P^{-1} F
      # Option 1: sqrt(P) = Q sqrt(Λ) Q'
      # Option 2: sqrt(P) = chol(P).L
      if !haskey(weights["precon"],"_solver")
         solver = "chol"
      else
         solver = weights["precon"]["_solver"]
      end
      if solver == "chol"
         rtP = cholesky(P, Val(false)).L
      elseif solver == "sqrt"
         rtP = sqrt(P)
      elseif solver == "id"
         rtP = cholesky(P)   # rtP is just P now, but factorized
      else
         error("unknown root solver")
      end

      irows = matrows(dat, obskey)
      Y[irows] = rtP \ Y[irows]
      Ψ[irows, :] = rtP \ Ψ[irows, :]
   end
   @info("USING FORCE PRECONDITIONER")
   return Ψ, Y
end

"""
`get_lsq_system(db::LsqDB; kwargs...) -> Ψ, Y`

Assemble the least squares system + rhs (observations). The `kwargs` can be used
to select a subset of the available data or basis set, and to adjust the
weights by config_type and observation type. See `lsqfit` for a list
of keyword arguments. Allowed/required kwargs are `verbose, weights,
E0, Vref, Ibasis, Itrain, regularisers`.
"""
@noinline function get_lsq_system(db::LsqDB;
                                  verbose=true,
                                  Ibasis = :,
                                  Itrain = :,
                                  E0 = nothing,
                                  Vref = OneBody(E0),
                                  weights = nothing,
                                  #scal_wgs = false,
                                  regularisers = [])
   weights = _fix_weights!(weights)

   # get the observations vector and the weights vector
   # the Vref potential is subtracted from the observations
   Y, W, Icfg = collect_observations(db, weights, Vref)#, scal_wgs)

   # check for NaNs
   any(isnan, Y) && error("NaN detected in observations vector")
   any(isnan, W) && error("NaN detected in weights vector")

   # NON-DIAGONAL ADDITIONAL WEIGHTS
   #  If the user supplies a force preconditioner, then we premultiply
   #  Ψ and Y with another weight function that modifies everything
   Ψ, Y = _forceprecon(db, Y, weights)

   # convert : into a vector or list
   if Itrain == : ; Itrain = 1:length(Y); end

   # get the right slice of the big fat huge enourmous LSQ matrix
   # the columns are just the basis functions
   Icols = Ibasis
   # the rows are those which have (a) a non-zero weight
   # and (b) the configuration index is part of the training set
   Irows = intersect( findall(W .!= 0) |> sort,
                      findall(in.(Icfg, Ref(Itrain))) )

   # we make this a view but make sure to copy it before applying the weights
   Ψ = @view Ψ[Irows, Icols]
   Y = Y[Irows]
   W = W[Irows]

   if any(isnan, Ψ)
      @error("discovered NaNs in LSQ system matrix")
      naninds = findall(isnan.(Ψ))
      n = min(length(naninds), 10)
      @show length(naninds)
      @show naninds[1:n]
      error("discovered NaNs in LSQ system matrix")
   end

   # regularise
   if !isempty(regularisers)
      Ψ, Y = _regularise!(Ψ, Y, db.basis, regularisers; verbose=verbose, Ibasis=Ibasis)
      # adjust the weight-vector to the new system size
      append!(W, ones(length(Y)-length(W)))
   else
      Ψ = collect(Ψ)
   end

   # now rescale Y and Ψ according to W => Y_W, Ψ_W; then the two systems
   #   \| Y_W - Ψ_W c \| -> min  and \| W (Y - Ψ*c)^T \| -> min
   # are equivalent
   Y = Diagonal(W) * Y
   lmul!(Diagonal(W), Ψ)

   # this should be it ...
   return Ψ, Y
end

_show_free_mem() =
      @info("Free Memory: ≈ $(round(Sys.free_memory()*1e-9, digits=2)) GB")

"""
`lsqfit(db; kwargs...) -> IP, errs`

Given the pre-computed least-squares system `db` (cf `LsqDB`) setup a least
squares system, compute the solution and return an interatomic potential
`IP` as well as the associated error estimates `errs`.

## Keyword Arguments:

* `E0` (required unless `Vref` is provided) : energy of the `OneBody` term
* `Vref` : a reference potential from which to start the fit.
* `weights` (required unless 'sigmas' is provided) : A dictionary specifying
the weights for different types of configurations, see `?collect_observations`
for how this is specified.
* `sigmas` (required unless 'weights' is provided) : A dictionary specifying
the sigmas for different types of configurations, see `?collect_observations`
for how this is specified.
* `Ibasis` : indices of basis functions to be used in the fit, default is `:`
* `verbose` : true or false
* `solver` : -experimental, still need to  write the docs for this-
* `regularisers` : a list of regularisers to be added to the lsq functional.
Each regulariser `R` can be of an arbtirary type but this type must implement the
conversion to matrix
```
Matrix(R, basis; verbose={true,false})
```
If `Areg = Matrix(R, basis)` then this corresponds to adding
`|| Areg * x ||²` to the least squares functional.
* `deldb = false` : experimental - if `true` then the lsq matrix in the `db` is
deleted after assembling the weighted lsq system
* `asmerrs = false` : experimental - if `true` then the lsq matrix is used to
assemble the fitting errors.
* `saveqr = nothing` : experimental - if `saveqr` is a Dict then the factors
of the QR factorisation are stored in it for future use outside of this
function

## Return types

* `IP`: whatever `JuLIP.MLIPs.combine` returns
* `fitinfo::Dict`: adds all sort of information about the fit
"""

function do_qr!(db, Ψ, Y)
   @info("Performing QR decomposition")
   GC.gc();
   qrΨ = qr!(Ψ)

   saveqr = Dict()
   saveqr["Q"] = Matrix(qrΨ.Q)
   saveqr["R"] = Matrix(qrΨ.R)
   saveqr["Y"] = Y

   db = LsqDB(db.basis, db.configs, db.Ψ, db.dbpath, saveqr)

   flush_qr(db)

   return db
end

@noinline function lsqfit(db::LsqDB;
                solver=(:qr, ), verbose=true,
                Ibasis = :,
                Itrain = :,
                Itest = nothing,
                E0 = 0.0,
                Vref = nothing, # OneBody(E0),
                weights = nothing,
                sigmas = nothing,
                regularisers = [],
                deldb = false,
                asmerrs = false,
                lasso = false,
                saveqr = nothing,
                #scal_wgs = false,
                kwargs...)

   if !isnothing(weights) && !isnothing(sigmas)
      @error("cannot specify both weights and sigmas")
   elseif !isnothing(sigmas)
      weights = sigmas
      for (config, config_weights) in weights
         for (key, weight) in config_weights
            config_weights[key] = 1.0/config_weights[key]
         end
      end
   end
   weights = _fix_weights!(weights)
   Jbasis = ((Ibasis == Colon()) ? (1:length(db.basis)) : Ibasis)

   verbose && @info("assemble lsq system")
   verbose && _show_free_mem()
   Ψ, Y = get_lsq_system(db; verbose=verbose, Vref=Vref,
                             Ibasis=Ibasis, Itrain = Itrain,
                             weights = weights,
                             regularisers = regularisers)
                             #scal_wgs = scal_wgs)
   verbose && _show_free_mem()
   if deldb
      @info("Deleting database - `db` can no longer be saved to disk")
      db.Ψ = Matrix{Float64}(undef, 0,0)
      db.dbpath = ""
   end
   GC.gc()
   verbose && _show_free_mem()

   κ, p_1, int_order = 0.0, 0.0, 0.0
   if (solver[1] == :qr) || (solver == :qr)
      verbose && @info("solve $(size(Ψ)) LSQ system using QR factorisation")
      qrΨ = qr!(Ψ)
      κ = cond(qrΨ.R)
      GC.gc();
      verbose && @info("cond(R) = $(cond(qrΨ.R))")
      c = qrΨ \ Y
      rel_rms = norm(qrΨ.Q * (qrΨ.R * c) - Y) / norm(Y)

      if saveqr isa Dict
         saveqr["Q"] = qrΨ.Q
         saveqr["R"] = qrΨ.R
         saveqr["Y"] = Y
      end

      qrΨ = nothing

   elseif solver[1] == :svd
      verbose && @info("solve $(size(Ψ)) LSQ system using SVD factorisation")
      ndiscard = solver[2]
      F = svd(Ψ)
      c = F.V[:,1:(end-ndiscard)] * (Diagonal(F.S[1:(end-ndiscard)]) \ (F.U' * Y)[1:(end-ndiscard)])
      rel_rms = norm(Ψ * c - Y) / norm(Y)

   elseif solver[1] == :rrqr
      verbose && @info("solve $(size(Ψ)) LSQ system using Rank-Revealing QR factorisation")
      qrΨ = pqrfact(Ψ, rtol=solver[2])
      verbose && @info("cond(R) = $(cond(qrΨ.R))")
      c = qrΨ \ Y
      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :rid
      r = solver[2]
      verbose && @info("solve $(size(Ψ)) LSQ system using Ridge Regression [r = $(r)] ")
      qrΨ = qr!(Ψ)
      Nb = size(qrΨ.R, 1)
      y = (Y' * Matrix(qrΨ.Q))[1:Nb]
      η0 = norm(Y - Matrix(qrΨ.Q) * y)

      τ = r * η0
      Γ = Matrix(I, Nb, Nb)

      c = reglsq(Γ = Γ, R = Matrix(qrΨ.R), y=y, τ= τ, η0 = η0 );

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :lap
      rscal = solver[2][1]
      r = solver[2][2]
      verbose && @info("solve $(size(Ψ)) LSQ system using Laplacian Regularisation [r = $(r)] ")

      qrΨ = qr!(Ψ)
      Nb = size(qrΨ.R, 1)
      y = (Y' * Matrix(qrΨ.Q))[1:Nb]
      η0 = norm(Y - Matrix(qrΨ.Q) * y)

      τ = r * η0

      s = ACE.scaling(db.basis.BB[2], rscal)
      l = append!(ones(length(db.basis.BB[1])), s)

      Γ = Diagonal(l)

      c = reglsq(Γ = Γ, R = Matrix(qrΨ.R), y=y, τ= τ, η0 = η0 );

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :rrqr_lap
      r_tol = solver[2]
      rlap = solver[3]

      s = ACE.scaling(db.basis.BB[2], 2)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(rlap .* l)

      Ψreg = vcat(Ψ, Γ)

      qrΨreg = pqrfact(Ψreg, rtol=r_tol)
      Yreg = vcat(Y, zeros(length(l)))

      c = qrΨreg \ Yreg

      rel_rms = norm(qrΨreg * c - Yreg) / norm(Yreg)
   elseif solver[1] == :rrqr_lap_rescale
      r_tol = solver[2]
      rlap = solver[3]

      s = ACE.scaling(db.basis.BB[2], 2)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(rlap .* l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      qrΨreg = pqrfact(Ψreg, rtol=r_tol)
      c = D_inv * (qrΨreg \  Y)
      rel_rms = norm(Ψreg * c - Y) / norm(Y)
   elseif solver[1] == :lap_elastic_net_rel
      rlap_scal = solver[2][1]
      rtol = solver[2][2]
      etol = solver[2][3]

      @info("rlap_scal = $(rlap_scal), rrqr_tol=$(rtol), e_tol = $(etol)")

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      non_zero_ind = [j for (j,i) in enumerate(Ψ[1,:]) if sum(i) != 0]
      zero_ind = [j for (j,i) in enumerate(Ψ[1,:]) if sum(i) == 0]

      Ψreg_red = Ψreg[:, setdiff(1:end, zero_ind)]

      qrΨ = pqrfact(Ψreg_red, rtol=rtol)
      cred = qrΨ \ Y

      cred_big = zeros(length(Ψreg[1,:]))

      for (i,k) in enumerate(non_zero_ind)
        cred_big[k] = cred[i]
      end

      c = D_inv * cred_big

      rel_rms0 = norm(Ψ * c - Y) / norm(Y)

      @info("rel_rms0 = $(rel_rms0)")
      @info("etol*rel_rms0 = $(etol*rel_rms0)")

      function _f(Ψreg, Y, α; etol=1e-5, rtol=1e-9, return_solution=false)
          cv = glmnet(Ψreg, Y, alpha=α)
          theta = cv.betas[:, end]

          non_zero_ind = findall(x -> x != 0.0, theta)
          zero_ind = findall(x -> x == 0.0, theta)

          Ψreg_red = Ψreg[:, setdiff(1:end, zero_ind)]

          qrΨ = pqrfact(Ψreg_red, rtol=rtol)
          cred = qrΨ \ Y

          cred_big = zeros(length(Ψreg[1,:]))

          for (i,k) in enumerate(non_zero_ind)
              cred_big[k] = cred[i]
          end

          c = D_inv * cred_big

          rel_rms = norm(Ψ * c - Y) / norm(Y)
          @info("rel_rms=$(rel_rms) and α=$(α), keeping $(length(non_zero_ind)) basis functions ($(round(length(non_zero_ind)/length(theta), digits=2)*100)%)")

          if return_solution
              return c
          else
              return rel_rms - etol
          end
      end

      α = find_zero(α -> _f(Ψreg, Y, α, etol=etol*rel_rms0, rtol=rtol), (0,1), Roots.Bisection(), xatol=1E-6)#atol=etol/2) #atol=0.5
      @info("α found! α=$(α)")

      c = _f(Ψreg, Y, α, return_solution=true)
      ##
      s_1 = ACE.scaling(db.basis.BB[2], 1)
      p_1 = append!(ones(length(db.basis.BB[1])), s_1)
      #int_order = db.basis.BB[2].pibasis.inner[1].orders
      ##
   elseif solver[1] == :lap_rrqr
      rlap_scal = solver[2][1]
      rtol = solver[2][2]
      reduce = false
      if length(solver[2]) == 3 && solver[2][3] == true
         reduce = true
      end

      @info("rlap_scal = $(rlap_scal), rrqr_tol=$(rtol)")

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      if reduce
         Ψreg = convert.(Float32, Ψreg)
         Y = convert.(Float32, Y)
      end

      qrΨ = pqrfact!(Ψreg, rtol=rtol)
      cred = qrΨ \ Y

      c = D_inv * cred

      rel_rms = norm(Ψ * c - Y) / norm(Y)
      ##
      s_1 = ACE.scaling(db.basis.BB[2], 1)
      p_1 = append!(ones(length(db.basis.BB[1])), s_1)
      #int_order = db.basis.BB[2].pibasis.inner[1].orders
      ##
   elseif solver[1] == :lap_tik
      rlap_scal = solver[2][1]
      r_tik = solver[2][2]

      @info("rlap_scal = $(rlap_scal), r_tik=$(r_tik)")

      non_zero_ind = [j for (j,i) in enumerate(Ψ[1,:]) if sum(i) != 0]
      zero_ind = [j for (j,i) in enumerate(Ψ[1,:]) if sum(i) == 0]

      nbasis_full = length(Ψ[1,:])
      Ψ = Ψ[:, setdiff(1:end, zero_ind)]
      nbasis = length(Ψ[1,:])
      nobs = length(Ψ[:,1])

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l[:, setdiff(1:end, zero_ind)])

      D_inv = pinv(Γ)
      Ψred = similar(Ψ)
      mul!(Ψred, Ψ, D_inv)
      Ψ = nothing

      Ψreg = zeros(nobs + nbasis, nbasis)

      for (i,row) in enumerate(eachrow(Ψred))
         Ψreg[i,:] = row
      end

      for i in 1:nbasis
         Ψreg[nobs+i,i] = r_tik
      end
      Ψred = nothing

      Yreg = vcat(Y, zeros(nbasis))

      GC.gc()
      qrΨreg = qr!(Ψreg)
      cred = qrΨreg \ Yreg

      cred_small = D_inv * cred
      c = zeros(nbasis_full)

      for (i,k) in enumerate(non_zero_ind)
         c[k] = cred_small[i]
      end

      rel_rms = 0.0
      ##
      s_1 = ACE.scaling(db.basis.BB[2], 1)
      p_1 = append!(ones(length(db.basis.BB[1])), s_1)
      #int_order = db.basis.BB[2].pibasis.inner[1].orders
      ##
   elseif solver[1] == :elastic_net_lsqr
      damp = solver[2][1]
      atol = solver[2][2]
      α = solver[2][3]

      @info("damp=$(damp), lsqr_atol=$(atol), alpha=$(α) | ELNET LSQR")

      cv = glmnet(Ψ, Y, alpha=α)
      theta = cv.betas[:,end]

      non_zero_ind = findall(x -> x != 0.0, theta)
      zero_ind = findall(x -> x == 0.0, theta)

      @info("keeping $(length(non_zero_ind)) basis functions ($(round(length(non_zero_ind)/length(theta), digits=2)*100)%)")

      Ψred = Ψ[:, setdiff(1:end, zero_ind)]

      cred, lsqrinfo = lsqr(Ψred, Y, damp=damp, atol=atol, log=true)
      println(lsqrinfo)

      c = zeros(length(Ψ[1,:]))

      for (i,k) in enumerate(non_zero_ind)
        c[k] = cred[i]
      end

      rel_rms = norm(Ψ * c - Y) / norm(Y)
      #int_order = db.basis.BB[2].pibasis.inner[1].orders
      ##
   elseif solver[1] == :elastic_net_lap
      α = solver[2][1]
      rlap_scal = solver[2][2]
      r = solver[2][3]
      cv = glmnet(Ψ, Y, alpha=α)
      theta = cv.betas[:,end]

      non_zero_ind = findall(x -> x != 0.0, theta)
      zero_ind = findall(x -> x == 0.0, theta)

      @info("keeping $(length(non_zero_ind)) basis functions ($(round(length(non_zero_ind)/length(theta), digits=2)*100)%)")

      Ψred = Ψ[:, setdiff(1:end, zero_ind)]

      qrΨred = qr!(Ψred)
      Nb = size(qrΨred.R, 1)
      y = (Y' * Matrix(qrΨred.Q))[1:Nb]
      η0 = norm(Y - Matrix(qrΨred.Q) * y)

      τ = r * η0

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      lred = [l[i] for i in non_zero_ind]
      Γ = Diagonal(lred)

      cred = reglsq(Γ = Γ, R = Matrix(qrΨred.R), y=y, τ= τ, η0 = η0 );

      c = zeros(length(Ψ[1,:]))

      for (i,k) in enumerate(non_zero_ind)
        c[k] = cred[i]
      end

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :lap_elastic_net_lap
      α = solver[2][1]
      rlap_scal = solver[2][2]
      r = solver[2][3]

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      cv = glmnet(Ψreg, Y, alpha=α)
      theta = cv.betas[:,end]

      non_zero_ind = findall(x -> x != 0.0, theta)
      zero_ind = findall(x -> x == 0.0, theta)

      Ψred = Ψreg[:, setdiff(1:end, zero_ind)]

      qrΨred = qr!(Ψred)
      Nb = size(qrΨred.R, 1)
      y = (Y' * Matrix(qrΨred.Q))[1:Nb]
      η0 = norm(Y - Matrix(qrΨred.Q) * y)

      τ = r * η0

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      lred = [l[i] for i in non_zero_ind]
      Γ = Diagonal(lred)

      cred = reglsq(Γ = Γ, R = Matrix(qrΨred.R), y=y, τ= τ, η0 = η0 );

      cred_big = zeros(length(Ψ[1,:]))

      for (i,k) in enumerate(non_zero_ind)
        cred_big[k] = cred[i]
      end

      c = D_inv * cred_big

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :gd_reg
      λ = solver[2][1]
      rlap_scal = solver[2][2]

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      function error(c, y, Ψ, λ)
         ŷ = Ψ*c
         loss = mean((y .- ŷ).^2) + λ*norm(c)
         return loss
     end

      res = optimize(c -> error(c, Y, Ψreg, λ), zeros(length(Ψ[1,:])), GradientDescent())

      creg = Optim.minimizer(res)

      c = D_inv * creg

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :lbfgs_reg
      λ = solver[2][1]
      rlap_scal = solver[2][2]

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      function _error(c, y, Ψ, λ)
         ŷ = Ψ*c
         loss = mean((y .- ŷ).^2) + λ*norm(c)
         return loss
      end

      res = optimize(c -> _error(c, Y, Ψreg, λ), zeros(length(Ψ[1,:])), LBFGS())

      creg = Optim.minimizer(res)

      c = D_inv * creg
      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :itlsq
      damp = solver[2][1]
      rlap_scal = solver[2][2]
      atol = solver[2][3]
      a2b = solver[2][4]
      if length(solver[2]) == 5
         maxiter, c_init = solver[2][5]
         @info("Using a given approximate solution c, maxiter=$(maxiter)")
      else
         c_init = zeros(length(db.Ψ[1,:]))
         maxiter=100000
      end
      @info("damp=$(damp), rlap_scal=$(rlap_scal), lsqr_atol=$(atol), a2b=$(a2b)")

      s = ACE.scaling(db.basis.BB[2], rlap_scal; a2b = a2b)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      mul!(Ψ,Ψ,D_inv)

      creg, lsqrinfo = lsqr!(c_init, Ψ, Y, damp=damp, atol=atol, maxiter=maxiter, log=true)
      println(lsqrinfo)

      c = D_inv * creg

      rel_rms = norm(Ψ * creg - Y) / norm(Y)
   elseif solver[1] == :itlsq_lap2b
      damp = solver[2][1]
      rlap_scal = solver[2][2]
      atol = solver[2][3]
      @info("damp=$(damp), rlap_scal=$(rlap_scal), lsqr_atol=$(atol) | scaling 2B ONLY")

      s = ACE.scaling(db.basis.BB[1], rlap_scal)#; a2b = a2b)
      l = append!(s, ones(length(db.basis.BB[2])))
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      mul!(Ψ,Ψ,D_inv)

      creg, lsqrinfo = lsqr(Ψ, Y, damp=damp, atol=atol, log=true)
      println(lsqrinfo)

      c = D_inv * creg

      rel_rms = norm(Ψ * creg - Y) / norm(Y)
   elseif solver[1] == :lap2b_rid
      r = solver[2][1]
      rlap_scal = solver[2][2]
      @info("r=$(r), rlap_scal=$(rlap_scal) | scaling 2B ONLY")

      s = ACE.scaling(db.basis.BB[1], rlap_scal)#; a2b = a2b)
      l = append!(s, ones(length(db.basis.BB[2])))
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      mul!(Ψ,Ψ,D_inv)

      qrΨ = qr!(Ψ)
      Nb = size(qrΨ.R, 1)
      y = (Y' * Matrix(qrΨ.Q))[1:Nb]
      η0 = norm(Y - Matrix(qrΨ.Q) * y)

      τ = r * η0
      Γ = Matrix(I, length(l), length(l))

      creg = reglsq(Γ = Γ, R = Matrix(qrΨ.R), y=y, τ= τ, η0 = η0 );

      c = D_inv * creg

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :lsqr
      damp = solver[2][1]
      atol = solver[2][2]
      @info("damp=$(damp), lsqr_atol=$(atol) | SIMPLE LSQR")

      c, lsqrinfo = lsqr(Ψ, Y, damp=damp, atol=atol, log=true)
      println(lsqrinfo)

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :itlsq_Nw
      damp = solver[2][1]
      rlap_scal = solver[2][2]
      atol = solver[2][3]
      Nweight = solver[2][4]
      @info("damp=$(damp), rlap_scal=$(rlap_scal), lsqr_atol=$(atol)")

      V2scal = Nweight["2B"] .* ACE.scaling(db.basis.BB[1], rlap_scal)
      NBscal = ACE.scaling(db.basis.BB[2], rlap_scal)
      RegDiag = vcat(V2scal, NBscal)
      Γ = Diagonal(RegDiag)

      D_inv = pinv(Γ)
      mul!(Ψ,Ψ,D_inv)

      creg, lsqrinfo = lsqr(Ψ, Y, damp=damp, atol=atol, log=true)
      println(lsqrinfo)

      c = D_inv * creg

      rel_rms = norm(Ψ * creg - Y) / norm(Y)
   elseif solver[1] == :elastic_net_itlsq
      α = solver[2][1]
      damp = solver[2][2]
      rlap_scal = solver[2][3]
      atol = solver[2][4]
      a2b = solver[2][5]
      @info("α=$(α), damp=$(damp), rlap_scal=$(rlap_scal), lsqr_atol=$(atol), a2b=$(a2b)")

      s = ACE.scaling(db.basis.BB[2], rlap_scal; a2b = a2b)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      mul!(Ψ,Ψ,D_inv)

      cv = glmnet(Ψ, Y, alpha=α)
      theta = cv.betas[:,end]

      non_zero_ind = findall(x -> x != 0.0, theta)
      zero_ind = findall(x -> x == 0.0, theta)

      @info("α=$(α), keeping $(length(non_zero_ind)) basis functions ($(round(length(non_zero_ind)/length(theta), digits=2)*100)%)")
      Ψ = Ψ[:, setdiff(1:end, zero_ind)]

      cred, lsqrinfo = lsqr(Ψ, Y, damp=damp, atol=atol, log=true)
      println(lsqrinfo)

      cred_big = zeros(length(l))

      for (i,k) in enumerate(non_zero_ind)
        cred_big[k] = cred[i]
      end

      c = D_inv * cred_big

      rel_rms = norm(Ψ * cred - Y) / norm(Y)
      #rel_rms = 0.0
      rel_rms = norm(Ψ * cred - Y) / norm(Y)
   elseif solver[1] == :no2b_itlsq
      damp = solver[2][1]
      rlap_scal = solver[2][2]
      atol = solver[2][3]
      @info("Not keeping ACE 2B, damp=$(damp), rlap_scal=$(rlap_scal), lsqr_atol=$(atol)")

      len_pair = length(db.basis.BB[1])
      len_ace = length(db.basis.BB[2])

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(len_pair), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      mul!(Ψ,Ψ,D_inv)

      I_ord2 = findall(get_orders(db.basis.BB[2]) .== 1);

      zero_ind = I_ord2 .+ len_pair
      non_zero_ind = []
      for i in 1:(len_pair+len_ace)
          if i ∉ zero_ind
              append!(non_zero_ind, i)
          end
      end

      #theta = append!(ones(length(db.basis.BB[1])), cs)

      #non_zero_ind = findall(x -> x != 0.0, theta)
      #zero_ind = findall(x -> x == 0.0, theta)

      @info("keeping $(length(non_zero_ind)) basis functions ($(round(length(non_zero_ind)/(len_pair+len_ace), digits=2)*100)%)")
      Ψ = Ψ[:, setdiff(1:end, zero_ind)]

      cred, lsqrinfo = lsqr(Ψ, Y, damp=damp, atol=atol, log=true)
      println(lsqrinfo)

      cred_big = zeros(length(l))

      for (i,k) in enumerate(non_zero_ind)
        cred_big[k] = cred[i]
      end

      c = D_inv * cred_big

      rel_rms = norm(Ψ * cred - Y) / norm(Y)
      #rel_rms = 0.0
      rel_rms = norm(Ψ * cred - Y) / norm(Y)
   elseif solver[1] == :elastic_net_lap_manual
      α = solver[2][1]
      rscal = solver[2][2]
      lap_pen = solver[2][3]

      cv = glmnet(Ψ, Y, alpha=α)
      theta = cv.betas[:,end]

      non_zero_ind = findall(x -> x != 0.0, theta)
      zero_ind = findall(x -> x == 0.0, theta)

      @info("keeping $(length(non_zero_ind)) basis functions ($(round(length(non_zero_ind)/length(theta), digits=2)*100)%)")

      Ψred = Ψ[:, setdiff(1:end, zero_ind)]

      s = ACE.scaling(db.basis.BB[2], rscal)
      l = append!(ones(length(db.basis.BB[1])), s)

      lred = [l[i] for i in non_zero_ind]
      Γ = Diagonal(lap_pen .* lred)

      Ψreg = vcat(Ψred, collect(Γ))
      Yreg = vcat(Y, zeros(length(lred)))

      #qrΨreg = qr!(Ψreg)
      cred = Ψreg \ Yreg

      c = zeros(length(Ψ[1,:]))

      for (i,k) in enumerate(non_zero_ind)
        c[k] = cred[i]
      end

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   elseif solver[1] == :spgl_lap_rrqr
      it = solver[2][1]
      rlap_scal = solver[2][2]
      rtol = solver[2][3]

      s = ACE.scaling(db.basis.BB[2], rlap_scal)
      l = append!(ones(length(db.basis.BB[1])), s)
      Γ = Diagonal(l)

      D_inv = pinv(Γ)
      Ψreg = Ψ * D_inv

      opts = spgOptions(iterations = it)
      theta, r, g, info = spgl1(Ψ, Y, sigma = 1e-5, options=opts)

      non_zero_ind = findall(x -> x != 0.0, theta)
      zero_ind = findall(x -> x == 0.0, theta)

      Ψred = Ψreg[:, setdiff(1:end, zero_ind)]

      qrΨ = pqrfact!(Ψred, rtol=rtol)
      cred = qrΨ \ Y

      cred_big = zeros(length(Ψ[1,:]))

      @show length(cred_big)
      @show length(zero_ind)

      for (i,k) in enumerate(non_zero_ind)
        cred_big[k] = cred[i]
      end

      c = D_inv * cred_big

      rel_rms = norm(Ψ * c - Y) / norm(Y)
   else
      error("unknown `solver` in `lsqfit`")
   end

   # delete the lsq system and gc again
   Ψ = nothing
   GC.gc()
   verbose && _show_free_mem()

   if verbose
      @info("Relative RMSE on training set: $rel_rms")
   end

   IP = JuLIP.MLIPs.combine(db.basis, c)
   if Vref != nothing
      IP = SumIP(Vref, IP)
   end

   infodict = asm_fitinfo(db, IP, c, Ibasis, weights,
                          Vref, solver, E0, regularisers, verbose,
                          Itrain, Itest, asmerrs)
   infodict["kappa"] = κ
   infodict["p_1"] = p_1
   #infodict["int_order"] = int_order
   GC.gc()
   return IP, infodict
end

function asm_fitinfo(db, IP, c, Ibasis, weights,
                     Vref, solver, E0, regularisers, verbose,
                     Itrain = :, Itest = nothing, asmerrs=true)
   if Ibasis isa Colon
      Jbasis = collect(1:length(db.basis))
   else
      Jbasis = Ibasis
   end

   cfgtypes = setdiff(unique(configtype.(db.configs)), weights["ignore"])

   # compute errors TODO: still need to fix this!
   if asmerrs
      verbose && @info("Assemble errors table")
      @warn("new error implementation... redo this part please ")
      errs = Err.lsqerrors(db, c, Jbasis;
               cfgtypes=cfgtypes, Vref=Vref, Icfg=Itrain)
      if Itest != nothing
         errtest = Err.lsqerrors(db, c, Jbasis;
                  cfgtypes=cfgtypes, Vref=Vref, Icfg=Itest)
      else
         errtest = Dict()
      end
   end
   # --------------------------------------------------------------------
   # ASSEMBLE INFO DICT
   # --------------------------------------------------------------------
   verbose && @info("Assemble Information about the fit")

   # Julia Version Info
   iob = IOBuffer()
   versioninfo(iob)
   juliainfo = String(take!(iob))

   infodict = Dict("solver" => String(solver[1]),
                   "E0"     => E0,
                   "Ibasis" => Vector{Int}(Jbasis),
                   "c"      => c,
                   "dbpath" => dbpath(db),
                   "weights" => weights,
                   "regularisers"  => Dict.(regularisers),
                   "juliaversion"  => juliainfo,
                   "IPFitting_version" => "na", # get_pkg_info("IPFitting"),
                  )
   # TODO: fix IPFitting_version retrieval

   if asmerrs
      infodict["errors"] = errs
      infodict["errtest"] = errtest
   end
   # --------------------------------------------------------------------
   return infodict
end

import Pkg

# function get_pkg_info(pkg::AbstractString)
#     pkgs = [Pkg.API.check_package_name(pkg)]
#     ctx = Pkg.Types.Context()
#     Pkg.API.project_resolve!(ctx.env, pkgs)
#     Pkg.API.project_deps_resolve!(ctx.env, pkgs)
#     Pkg.API.manifest_resolve!(ctx.env, pkgs)
#     Pkg.API.ensure_resolved(ctx.env, pkgs)
#     i = Pkg.Display.status(ctx, pkgs)[1]
#     return Dict("name" => i.name,
#                 "uuid" => string(i.uuid),
#                 "ver" => string(i.new.ver))
# end

_fix_weights!(::Nothing) = _fix_weights!(Dict{String, Any}())

function _fix_weights!(weights::Dict)
   if !haskey(weights, "ignore")
      weights = Dict{Any,Any}(weights)
      weights["ignore"] = String[]
   end
   if !haskey(weights, "default")
      weights["default"] = Dict("E" => 1.0, "F" => 1.0, "V" => 1.0)
   end
   return weights
end

"""
`function reglsq(; kwargs...)`

Solves the problem
```
   min || Γ * x || subj to || A * x - z || ≦ tol
```

### Arguments:
* `Γ, R, z, tol` : specify the problem to be solved
* Instead of `tol` one may pass the kwargs `η0, τ`, which will set `tol = sqrt(τ^2-η0^2)`
* `abstol = 1e-2, reltol = 1e-2` : termination tolerance
* `λinit = 1e-3, λmax = 1e2` : search parameters
"""
function reglsq( ;
                Γ = nothing,
                R = nothing,
                y = nothing,
                η0 = nothing,
                τ = nothing,
                tol = sqrt(τ^2 - η0^2),
                abstol = 1e-5,
                reltol = 1e-2,
                λinit = 1e-3,
                λmax = 1e2,
                verbose = true)
   @assert size(Γ, 2) == size(R, 2)
   @assert length(y) == size(R, 1)
   Qc, Rc = qr(collect(Γ))
   Y = [y; zeros(size(Γ, 1))]

   verbose && @info("`reglsq` : solve regularised least squares")
   x = nothing

   function _solve(λ)
      x = qr([ R; λ * Rc ]) \ Y
      return x, norm(R * x - y)
   end

   # Bracketing
   λ1 = 0.0
   λ2 = λinit
   while true
      x, η = _solve(λ2)
      if η > tol
         break
      end
      if λ2 > λmax
         @error("""reglsq reached λmax = $λmax in the bracketing stage.
                  The returned coeffcients are in the interior of the constraint
                  set ||Rx-z|| ≦ tol. If stronger regularisation is required,
                  then please increase the λmax parameter.""")
         return x
      end

      λ1 = λ2
      λ2 *= 10
   end

   verbose && @info("found bracket, starting bisection")

   # Bisection
   while true
      λ = 0.5 * (λ1 + λ2)
      x, η = _solve(λ)
      if tol / (1+reltol) <= η <= tol * (1+reltol)
         # success
         break
      end
      if η > tol
         λ2 = λ
      else
         λ1 = λ
      end
      if abs(λ1 - λ2) < abstol
         @error("""bracket has become too small; returning current solutions,
                   if it is unsatisfactory, then please reduce the
                   `abstol` parameter.""", λ1, λ2, abstol)
         return x
      end
   end

   verbose && @info("found a solution")
   return x
end

function reglsq(lsqdb::LsqDB, saveqr::Dict, Γ; tol_fact = 1.2, verbose=true,
                E0 = nothing, Vref = OneBody(E0))
   R = saveqr["R"]
   Y = saveqr["Y"]
   Nb = size(R, 1)
   y = (Y' * saveqr["Q"])[1:Nb]
   _rmse = norm(Y - saveqr["Q"] * y)
   creg = IPFitting.Lsq.reglsq(; Γ = Γ, R=R, y=y,
                            η0 = _rmse, τ = tol_fact * _rmse,
                            verbose=verbose)
   verbose && @info("|creg| = $(norm(creg)); |Γ creg| = $(norm(Γ*creg))")
   return JuLIP.MLIPs.SumIP(Vref,
                            JuLIP.MLIPs.combine(lsqdb.basis, creg))
end

end

# @noinline function onb(db::LsqDB;
#                          solver=(:qr, ), verbose=true,
#                          Ibasis = :,
#                          Itrain = :,
#                          E0 = nothing,
#                          Vref = OneBody(E0),
#                          weights = nothing,
#                          regularisers = [],
#                          combineIP = nothing,
#                          kwargs...)
#
#    verbose && @info("assemble lsq system")
#    Ψ, _ = get_lsq_system(db; verbose=verbose, Vref=Vref,
#                              Ibasis=Ibasis, Itrain=Itrain,
#                              weights = weights,
#                              regularisers = regularisers,
#                              kwargs...)
#    @assert solver[1] == :qr
#    verbose && @info("QR-factorize Ψ, size=$(size(Ψ))")
#    qrΨ = qr(Ψ)
#    verbose && @info("cond(R) = $(cond(qrΨ.R))")
#    Rinv = pinv(qrΨ.R)
#    basis = db.basis[Ibasis]
#    onb = []
#    for n = 1:size(Rinv, 2)
#       push!(onb, combineIP(basis, Rinv[:,n]))
#    end
#    return [b for b in onb]
# end
#
# rlap = solver[2][2]
# p = solver[2][3]
#
# scale = true
# if length(solver[2]) == 4 && solver[2][4] == false
#    scale = false
# end
#
# s_p = scaling(db.basis.BB[2], p)
# p_p = append!(ones(length(db.basis.BB[1])), s_p)
# Γ = collect(Diagonal(rlap .* p_p))
#
# if scale == true
#    D_inv = pinv(Γ)
#    Ψreg = Ψ * D_inv
#    Yreg = Y
# else
#    Ψreg = vcat(Ψ, Γ)
#    Yreg = vcat(Y, ones(length(Ψ[1,:])))
# end
#
# if startswith(String(solver[1]), "lap_elastic_net_perc")
#    perc_ob = solver[2][1]
#    function _f(α, perc_ob, Ψreg, Yreg)
#       cv = glmnet(Ψreg, Yreg, alpha=α)
#       theta = cv.betas[:,end]
#
#       zero_ind = findall(x -> x == 0.0, theta)
#       non_zero_ind = findall(x -> x != 0.0, theta)
#
#       perc = round(length(non_zero_ind)/length(theta) * 100.0, digits=2)
#       @info("α=$(α) keeps $(perc)% of total basis functions")
#
#       return length(non_zero_ind)/length(theta) - perc_ob
#    end
#
#    α = find_zero(α -> _f(α, perc_ob, Ψreg, Yreg), (0,1), Roots.Bisection(), atol=0.01)
#    @info("α found! α=$(α)")
# else
#    α = solver[2][1]
# end
#
# cv = glmnet(Ψreg, Yreg, alpha=α)
# theta = cv.betas[:,end]
#
# zero_ind = findall(x -> x == 0.0, theta)
# non_zero_ind = findall(x -> x != 0.0, theta)
#
# perc = round(length(non_zero_ind)/length(db.basis)*100, digits=2)
# @info("Reduced LSQ Problem using $(perc)% of total basis functions [$(length(non_zero_ind))]")
#
# if scale == true
#    Ψred = Ψreg[:, setdiff(1:end, zero_ind)]
# else
#    Ψred = Ψ[:, setdiff(1:end, zero_ind)]
# end
#
# if endswith(String(solver[1]), "rrqr")
#    rtol = solver[3]
#    @info("Performing RRQR [$(rtol)]")
#    qrΨred = pqrfact(Ψred, rtol=rtol)
#    cred = qrΨred \ Y
# else
#    @info("Performing QR decomposition")
#    cred = Ψred \ Y
# end
#
# c_scal = zeros(length(db.basis))
#
# for (i,k) in enumerate(non_zero_ind)
#    c_scal[k] = cred[i]
# end
#
# if scale == true
#    c = D_inv * c_scal
# else
#    c = c_scal
# end
#
# if any(iszero, c[1:length(db.basis.BB[1])])
#    @info("Warning: culled a 2B function!! [TODO]")
# end
#
# rel_rms = norm(Ψ * c - Y) / norm(Y)
#
# ##
# s_1 = scaling(db.basis.BB[2], 1)
# p_1 = append!(ones(length(db.basis.BB[1])), s_1)
# ##


text:
abstract type GTR <: NASM end

const _πA(mod::GTR) = mod.πA
const _πC(mod::GTR) = mod.πC
const _πG(mod::GTR) = mod.πG
const _πT(mod::GTR) = mod.πT

ext: jl


text:
using Combinatorics
using CausalInference
using Test
using Random

Random.seed!(2)

@testset "combinations" begin
    for n in 1:10
        r = shuffle(1:n)
        for k in 0:n+1
            w = rand(1:n)

            c1 = collect(copy(c) for c in combinations(r, k) if r[w] ∉ c)
            comb = CausalInference.combinations_without(r, k, w)
            c2 = collect(copy(c) for c in comb)

            @test length(c1) == length(comb)
            @test sort(c1) == sort(c2)
        end
    end
end


text:
# Sparsity promoting LSRTM w/ the linearized Bregman method
# Module contains main algorithm + preconditioners
# Author: Philipp Witte
# Date: December 2016
#

module LeastSquaresMigration
import Base.norm

using JOLI
using PyPlot
using JLD
using ..TimeModeling

# Files
include("linbreg_lsrtm.jl")
include("model_preconditioners.jl")
include("data_topmute.jl")
include("auxiliary_functions.jl")

end

max_stars_repo_name: SINBADconsortium/opesciSLIM

max_issues_repo_name: SINBADconsortium/opesciSLIM

max_forks_repo_name: SINBADconsortium/opesciSLIM


text:
# TODO: Add backward testing after new-gradcheck in AutoGrad
@testset "primitive" begin

    @testset "Multiply" begin

        m = Multiply(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype)
        B = rand(1:10)

        x1   = _test_atype(zeros(_INPUTdim, B))
        y1_1 = m(x1)
        y1_2 = _test_atype(zeros(_OUTPUTdim, B))
        @test y1_1 == y1_2

        x2   = _test_atype(randn(_INPUTdim, B))
        y2_1 = m(x2)
        y2_2 = m.w * x2
        @test y2_1 == y2_2
    end

    @testset "Embed" begin
        x1 = rand(1:_INPUTdim, 5)
        m  = Embed(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype)
        y1 = m(x1)
        y2 = hcat([m.w[:,i] for i in x1]...)

        @test y1 == y2
    end

    @testset "BatchMul" begin
        m  = BatchMul(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype)
        B  = rand(1:20); T = rand(1:4)

        x1 = _test_atype(randn(_INPUTdim, B, T))
        y1 = m(x1, flatten=true)
        # put batched instances into the columns
        x1_flat = hcat([ x1[:,:,i] for i  in 1:T ]...)
        y2 = m.w * x1_flat
        @test y1 == y2

        x2   = _test_atype(randn(_INPUTdim, B, T, 3))
        y2_1 = m(x2, flatten=true)
        x2_flat = reshape(x2, _INPUTdim, B*T*3)
        y2_2 = m.w * x2_flat
        @test y2_2 == y2_1
    end

    @testset "Linear" begin
        m = Linear(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype)
        x1 = _test_atype(zeros(_INPUTdim, 1))
        y1 = m(x1)
        y2 = _test_atype(zeros(_OUTPUTdim, 1))
        @test y1 == y2

        m = Linear(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype, binit=ones, winit=randn)
        y1_2 = m(x1)
        y2_2 = _test_atype(ones(_OUTPUTdim, 1))
        @test y1_2 == y2_2
    end

    @testset "Dense" begin
        m = Dense(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype, binit=ones, winit=randn)
        x1 = _test_atype(zeros(_INPUTdim, 1))
        y1 = m(x1)
        y2 = _test_atype(ones(_OUTPUTdim, 1))
        @test y1 == y2

        m = Dense(input=_INPUTdim, output=_OUTPUTdim, atype=_test_atype, binit=ones, winit=randn, activation=relu)
        x1 = _test_atype(zeros(_INPUTdim, 1))
        y1 = m(x1)
        y2 = _test_atype(ones(_OUTPUTdim, 1))
        @test y1 == y2
    end
end


text:
using PyPlot
stress_y =  200.0

function vm_upper(a, c)
    vals = f(a[1], a[2], c)
    vm(vals[1], vals[2], 200)
end
vm(a,b) = sqrt(a^2 - a*b + b^2) - stress_y
f(m,c) = [600*cos(c) 600*sin(c)].*m
x_vals = []
max_iter = 100
y_vals = []
for i=0:0.1:(2*pi+0.3)
    wf(x) = f(x, i)
    t = 0.01
    step = 2
    merkki = -1
    s11, s22 = wf(t)
    ii = 0
    while (abs(vm(s11, s22)) > 1e-7) && ii < max_iter
        val = vm(s11, s22)
        if sign(val) != merkki
            merkki *= -1
            step *= -0.5
        end
        t += step
        s11, s22 = wf(t)
        ii += 1
    end
    push!(x_vals, s11)
    push!(y_vals, s22)
end

plot(x_vals, y_vals)
PyPlot.title("von Mises yield surface in 2D")
PyPlot.xlabel("Eig Stress 1")
PyPlot.ylabel("Eig Stress 2")
grid()
show()


text:
"""
	Modify a cluster by adding or removing vertices by picking at each step
	the vertex that has the maximum value of (Deg_external - Deg_Internal).
	Each vertex can be added in/removed only once.
"""
function refineCut(G::SparseMatrixCSC{Tv,Ti}, s::Array{Int64,1}) where {Tv,Ti}

	n = max(G.n, G.m)

	# inset - keep track of vertices in the low conductance set
	inset = zeros(Bool, n)
	for u in s
		inset[u] = true
	end

	# pq_in - costs are maintained to add elements to the set
	# pq_out - costs are maintained to remove elements from the set
	pq_in = DataStructures.PriorityQueue(Base.Order.Reverse)
	pq_out = DataStructures.PriorityQueue(Base.Order.Forward)

	for u in 1:n
		cost = 0

		for i in 1:deg(G,u)
			v = nbri(G,u,i)
			if inset[v]
				cost = cost + weighti(G,u,i)
			else
				cost = cost - weighti(G,u,i)
			end
		end

		if inset[u] == false
			pq_in[u] = cost
		end
		if inset[u] == true
			pq_out[u] = cost
		end
	end

	improve = true
	tries = 0
	while improve && tries < n
		improve = false

		# try to add a new vertex into the set
		if !isempty(pq_in)
			u,contrib = DataStructures.peek(pq_in)

			if contrib >= 0
				improve = true
				DataStructures.dequeue!(pq_in)
				inset[u] = true

				for i in 1:deg(G,u)
					v = nbri(G,u,i)

					if v in keys(pq_in)
						pq_in[v] = pq_in[v] + weighti(G,u,i)
					end
					if v in keys(pq_out)
						pq_out[v] = pq_out[v] + weighti(G,u,i)
					end
				end
			end
		end

		# try to remove a vertex from the set
		if !isempty(pq_out)
			u,contrib = DataStructures.peek(pq_out)
			if contrib < 0
				improve = true
				DataStructures.dequeue!(pq_out)
				inset[u] = false

				for i in 1:deg(G,u)
					v = nbri(G,u,i)
					if v in keys(pq_in)
						pq_in[v] = pq_in[v] - weighti(G,u,i)
					end
					if v in keys(pq_out)
						pq_out[v] = pq_out[v] - weighti(G,u,i)
					end
				end
			end
		end
	end

	news = (Int64)[]
	for i in 1:n
		if inset[i]
			push!(news, i)
		end
	end

	return news
end

"""
	Modify a cluster by passing through all the vertices exactly once and
	adding/removing them based on the value of (Deg_external - Deg_Internal).
"""
function dumbRefineCut(G::SparseMatrixCSC{Tv,Ti}, s::Array{Int64,1}) where {Tv,Ti}

	n = max(G.n, G.m)

	news = BitSet(s)

	for v in 1:n
		nrbsA = 0
		nrbsnotA = 0
		for i in 1:deg(G, v)
			u = nbri(G, v, i)
			if u in news
				nrbsA = nrbsA + weighti(G, v, i)
			else
				nrbsnotA = nrbsnotA + weighti(G, v, i)
			end
		end

		if nrbsA >= nrbsnotA
			push!(news, v)
		else
			if v in news
				pop!(news, v)
			end
		end
	end

	return collect(news)

end


text:
function _pathway_from(
  id::Int64,
  network::TransmissionNetwork;
  depth::Real=Inf,
  include_id::Bool=true)

  local path
  if network.external[id] || any(network.internal[:, id])
    path = Int64[]
    push!(path, id)
    path_lengths = [0]
    while (length(path) > path_lengths[end]) & (depth >= length(path_lengths))
      push!(path_lengths, length(path))
      for j in path[(path_lengths[end-1]+1):path_lengths[end]]
        append!(path, findall(network.internal[j, :]))
      end
    end
    if include_id == false
      popfirst!(path)
    end
  else
    path = nothing
  end
  @logmsg LogLevel(-5000) "Transmission pathway from Individual $id: $path"
  return path
end

function _pathway_from(
  id::Nothing,
  network::TransmissionNetwork;
  depth::Real=Inf,
  include_id::Bool=true)

  path = Union{Nothing, Int64}[nothing]
  append!(path, findall(network.external))
  path_lengths = [0; 1]
  while (length(path) > path_lengths[end]) & (depth >= length(path_lengths)-1)
    push!(path_lengths, length(path))
    for j in path[(path_lengths[end-1]+1):path_lengths[end]]
      append!(path, findall(network.internal[j, :]))
    end
  end
  if include_id == false
    popfirst!(path)
  end
  @logmsg LogLevel(-5000) "Transmission pathway from external source: $path"
  return path
end

function _pathway_to(
  id::Int64,
  network::TransmissionNetwork;
  depth::Real=Inf,
  include_id::Bool=true)

  local path
  if network.external[id] || any(network.internal[:, id])
    path = Union{Nothing, Int64}[id]
    next_id = findfirst(network.internal[:, id])
    while (length(path) < depth) && (next_id !== nothing)
      push!(path, next_id)
      next_id = findfirst(network.internal[:, path[end]])
    end
    push!(path, next_id)
    if include_id == false
      popfirst!(path)
    end
  else
    path = nothing
  end
  @logmsg LogLevel(-5000) "Transmission pathway to Individual $id: $path"
  return path
end

function _source(id::Int64, network::TransmissionNetwork)
  source = findfirst(network.internal[:, id])
  @logmsg LogLevel(-5000) "Transmission source lookup" Individual = id Source = source
  return source
end

function generate(
  ::Type{Tree},
  events::Events{S},
  obs::Vector{Float64},
  network::TransmissionNetwork) where {
  S <: DiseaseStateSequence}

  if length(obs) != individuals(events)
    throw(ErrorException("Infection observation vector does not match number of individuals in population"))
  end

  # Initialization
  expos = 0
  tree = Tree()
  addnode!(tree) # Root/MRCA node
  local event_times
  local parent_node
  local branch_length

  if State_E in S
    event_times = [events.exposure obs]
  else
    event_times = [events.infection obs]
  end
  event_nodes = Array{Union{Nothing, Int64}, 2}(fill(nothing, size(event_times)))
  event_order = sortperm(event_times[:])
  event_lookup = CartesianIndices(size(event_times))
  @logmsg LogLevel(-5000) "event_nodes, event_order, and event_lookup created" ArraySize = size(event_times)

  pathways = [_pathway_from(i, network, include_id = true) for i = 1:individuals(events)]

  # Determine significant transmissions (results in a observation in the future)
  significant_transmissions = (event_times[:, 1] .== -Inf) .| [pathways[i] !== nothing && any(event_times[pathways[i], 2] .>= event_times[i, 1]) for i = 1:individuals(events)]
  @logmsg LogLevel(-5000) "Significant transmitters" significant_transmissions

  start_time = minimum(events)

  # Iterate through all events to build tree
  for i = eachindex(event_order)

    event_time = event_times[event_order[i]]

    # Stop when no remaining valid event times
    isnan(event_time) && break

    # Determine the individual and event type
    event_individual, event_type = Tuple(event_lookup[event_order[i]])

    @debug "Building phylogenetic tree..." Event = i Individual = event_individual Type = event_type == 1 ? "Transmission" : "Observation" Time = event_time Index = event_order[i]

    # For purposes of tree generation, replace -Inf times (initial condition coding) with start_time
    if event_time == -Inf
      @logmsg LogLevel(-5000) "Setting -Inf event time to $(start_time)"
      event_time = start_time
    end

    # For transmission events...
    if event_type == 1

      # Add transmission event to tree only if it is significant
      if significant_transmissions[event_individual]

        # For significant external transmission events...
        if network.external[event_individual]
          @logmsg LogLevel(-5000) "Event $i is a significant external transmission"
          parent_node = 1
          expos += 1
          if expos > 1
            @warn "Current support is intended for a single intial event with no external exposures" ExternalExposures = expos
          end
          branch_length = event_time - start_time
          @logmsg LogLevel(-5000) "Set parent node to the root" ParentNode = parent_node BranchLength = branch_length
          branch!(tree, parent_node, branch_length)

          # Record tree and node ID
          event_nodes[event_order[i]] = length(tree.nodes)

        # For other significant transmission events...
        else

          # Determine source of transmission
          source = _source(event_individual, network)

          # If initial condition, no source
          if source === nothing
            @warn "Shouldn't be!"

          # Previous significant transmissions from this source of transmission
          else
            @logmsg LogLevel(-5000) "Event $i is a significant internal transmission" Source = source
            source_prior_transmissions = findall(network.internal[source, :][:] .& (event_times[:, 1] .< event_time) .& significant_transmissions)

            # For when the source of transmission has yet to be detected...
            if isnan(event_times[source, 2]) || event_times[source, 2] > event_times[event_individual, 1]
              @logmsg LogLevel(-5000) "Individual $source, the transmission source for individual $event_individual had not yet been detected"
              # When there have been prior transmissions from this undetected source
              if length(source_prior_transmissions) > 0
                @logmsg LogLevel(-5000) "Individual $source has had earlier transmissions" PriorTransmissionIds = source_prior_transmissions PriorTransmissionTimes = event_times[source_prior_transmissions, 1] PriorTransmissioNodes = event_nodes[source_prior_transmissions, 1]
                parent_node = event_nodes[source_prior_transmissions[argmax(event_times[source_prior_transmissions, 1])], 1]
                branch_length = event_time - maximum(event_times[source_prior_transmissions, 1])

              # When there have not been any prior transmissions from this undetected source
              else
                parent_node = event_nodes[source, 1]
                branch_length = event_time - max(event_times[source, 1], start_time)
                @logmsg LogLevel(-5000) "Individual $source has had no earlier transmissions" Source = source ParentNode1 = event_nodes[source, 1] ParentNode2 = parent_node BranchLength = branch_length EventTime = event_time ParentNodeTime = event_times[source, 1]
              end

            # For when the source of transmission has been detected...
            else
              @logmsg LogLevel(-5000) "Individual $source, the transmission source for individual $event_individual had been already detected"
              # And detection of source of transmission is most recent, relevant event...
              if length(source_prior_transmissions) == 0 || all(event_times[source, 2] .> event_times[source_prior_transmissions, 1])
                @logmsg LogLevel(-5000) "The detection of individual $source is the recentmost relevant event to transmission of Individual $event_individual"
                # Determine parent node of the observation leaf node
                parent_node = parentnode(tree, event_nodes[source, 2])
                branch_length = event_time - max(event_times[source, 2], start_time)

              # Or some prior transmission is most recent, relevant event...
              else
                @logmsg LogLevel(-5000) "Individual $source has had earlier transmissions, after detection" PriorTransmissionIds = source_prior_transmissions PriorTransmissionTimes = event_times[source_prior_transmissions, 1] PriorTransmissioNodes = event_nodes[source_prior_transmissions, 1]
                parent_node = event_nodes[source_prior_transmissions[argmax(event_times[source_prior_transmissions, 1])], 1]
                branch_length = event_time - maximum(event_times[source_prior_transmissions, 1])
              end
            end
            # Add branch and node
            @logmsg LogLevel(-5000) "branch! information" Individual = event_individual Source = source ParentNode = parent_node Node = length(tree.nodes) + 1 BranchLength = branch_length
            branch!(tree, parent_node, branch_length)

            # Record node id
            event_nodes[event_order[i]] = length(tree.nodes)
          end
        end
      else
        @logmsg LogLevel(-5000) "Event $i is an insignificant transmission"
      end

    # Infection observation event
    elseif event_type == 2
      # Determine if individual has had significant transmissions prior to being observed as infected
      prior_transmissions = findall(network.internal[event_individual, :][:] .&
                            (event_times[:, 1] .< event_time) .&
                            significant_transmissions)

      if length(prior_transmissions) > 0
        # Parent node will be the internal node associated with the final pre-observation transmission
        @logmsg LogLevel(-5000) "Individual $event_individual transmitted to others prior to their detection" PriorTransmissionIds = prior_transmissions PriorTransmissionTimes = event_times[prior_transmissions, 1] PriorTransmissioNodes = event_nodes[prior_transmissions, 1]
        parent_id = prior_transmissions[argmax(event_times[prior_transmissions, 1])]
        parent_node = event_nodes[parent_id, 1]
        branch_length = event_time - event_times[parent_id, 1]
      else
        # Individual has not transmitted to others prior to being observed as infected
        @logmsg LogLevel(-5000) "Individual $event_individual has not transmitted to others prior to their detection"
        parent_node = event_nodes[event_individual, 1]
        branch_length = event_time - max(event_times[event_individual, 1], start_time)
      end

      if any(network.internal[event_individual, :][:] .& (event_times[:, 1] .>= event_time) .& significant_transmissions)
        # Individual goes on to have significant transmissions
        # Internal node for observation
        @logmsg LogLevel(-5000) "Individual $event_individual has significant future transmissions, add internal node for observation" ParentNode = parent_node BranchLength = branch_length
        branch!(tree, parent_node, branch_length)
        # Add zero-length branch so observation event will be a leaf node
        @logmsg LogLevel(-5000) "Add a zero length branch to ensure observation of individual $event_individual is a leaf node" ParentNode = length(tree.nodes) BranchLength = 0.0
        branch!(tree, length(tree.nodes), 0.0)

      else
        # Individual does not go one to have any significant transmissions
        # so will surely be a leaf node
        @logmsg LogLevel(-5000) "Individual $event_individual does not have significant future transmissions - their infection observation represents a leaf node" ParentNode = parent_node BranchLength = branch_length
        branch!(tree, parent_node, branch_length)
      end
      # Record leaf node id
      event_nodes[event_order[i]] = length(tree.nodes)
    end
  end
  return tree, event_nodes[:,2]
end

function generate(
  ::Type{Tree},
  events::Events{S},
  obs::EventObservations{S, M},
  network::TransmissionNetwork) where {
  S <: DiseaseStateSequence,
  M <:ILM}
  return generate(Tree, events, obs.infection, network)
end


text:
using MittagLefflerFunctions
using PyPlot

f(x) = exp(x)

nmax = 10
M = nmax + 2
a = chebyshev_coefs(Float64, f, nmax, M)

x = range(-1, 1, length=201)
y = Float64[ chebyshev_sum(a, xval) for xval in x ]

figure(1)
plot(x, y-f.(x))
grid(true)


text:
#struct Vertex { union { struct { _float x, y, z; }; _float pos[3]; }; };

""" TODO """
type ShapeCoordinates
	#std::string& name
	positions::AVec3f
	normals::AVec3f
	texcoords::AVec2f
	indicies::AIndex
	origin::Vec3f
	scale::Vec3f

	ShapeCoordinates() = new(AVec3f(),AVec3f(),AVec1f(),AIndex(),Vec3f(),Vec3f())
	ShapeCoordinates(positions,normals,texcoords,indicies,origin,scale) = new(positions,normals,texcoords,indicies,origin,scale)
end

""" TODO """
type VertexIndex
	v::Integer
	vt::Integer
	vn::Integer

	VertexIndex() = new(0,0,0)
	VertexIndex(a::Integer) = new(a,a,a)
	VertexIndex(v::Integer, vn::Integer, vt::Integer) = new(v,vt,vn)
end

AVertexIndex = Array{VertexIndex,1}
AAVertexIndex = Array{AVertexIndex,1}
AGroupIndex = Array{Array{Array{Integer,1},1},1}

""" TODO """
function clear(sc::ShapeCoordinates)
	sc.positions = AVec3f()
	sc.normals = AVec3f()
	sc.texcoords = AVec1f()
	sc.indicies = AIndex()
	sc.origin = Vec3f()
	sc.scale = Vec3f()
end

# for std::map
""" TODO """
function isSmaller(a::VertexIndex, b::VertexIndex)
	if a.v != b.v return (a.v < b.v) end
	if a.vn != b.vn return (a.vn < b.vn) end
	if a.vt != b.vt return (a.vt < b.vt) end
	false
end

""" TODO """
isSpace(c::Char) = (c == ' ') || (c == '\t')

""" TODO """
isNewLine(c::Char) = (c == '\r') || (c == '\n') || (c == '\0')

""" TODO """
function getOrigin(min_BBOX::Vec3f, max_BBOX::Vec3f, scale::Number)
	(max_BBOX + min_BBOX) * (scale * 0.5f0)
end

""" TODO """
function getScale(min_BBOX::Vec3f, max_BBOX::Vec3f)
	size = max_BBOX - min_BBOX
	println("size: ",size)
	scale_value = size.x > size.z ? size.x : size.z
	println("1scale_value: ",scale_value)
	scale_value = scale_value > size.y ? scale_value : size.y
	println("2scale_value: ",scale_value)
	scale_value = abs(scale_value)
	println("3scale_value: ",scale_value)
	if scale_value > 1 scale_value = 1 / scale_value
	#if scale_value > 0 && scale_value < 1 scale_value = scale_value
	else scale_value = 1
	end
	scale_value
end

#function parseval(T::DataType, str::String, start=1)
#	t = search(str, r"[\+\-]?[0-9]+(\.[0-9]+)?", start)
#
##	v = 0
#	try
#		v = parse(T, str[t.start:t.stop])
##	catch(e)
#	end
#
#	(v, t.stop)
#end

""" TODO """
function parseVal(T::DataType, str::AbstractString, default=T(0))
	try
		return parse(T, str)
	catch(e)
	end
	default
end

#for m in eachmatch(Regex("([0-9]+) ([0-9]+) ([0-9]+)"),"1 2 3s") println(m.captures) end
""" TODO """
function parseArray(T::DataType, str::AbstractString)
	a = Array{T,1}()
	m = matchall(Regex("(^|(?<=\\s))[+-]?([0-9]+[.])?[0-9]+([eE][-+]?[0-9]+)?((?=\\s)|\$)"),str) # only valid: (space?)(number)(space?)
	#for x in m push!(a, parseval(T, x)) end
	map(x->parseVal(T,x),m)
end

""" TODO """
function parseVec(linenr::Number, str::AbstractString, count::Number)
	if count < 1 error("count < 1") end

	a = parseArray(Float32, str)
	len = length(a)

	if len != count	warn(linenr, ": missmatch count: required=$count != found=$(length(a)) ", str)	end

	# on count missmatch: just read until max(count) reachted
	v = zeros(Float32,count)
	for i=1:(len>=count ? count : len) v[i]=a[i] end

	result = nothing
	if count == 1 result = Vec1f(v)
	elseif count == 2 result = Vec2f(v)
	elseif count == 3 result = Vec3f(v)
	elseif count == 4 result = Vec4f(v)
	else result = Vec{count, Float32}(v)
	end
	result
end

# Make index zero-base, and also support relative index.
""" TODO """
function parseIndex(str::AbstractString, n::Integer)
	#a = parseArray(Integer, str)
	i = parseVal(Int, str, -1) #length(a)>0 ? a[1] : 0
	#if idx >= 0 i = idx
	#else i = n + idx # negative value = relative
	#end
	i
end

""" TODO """
has(str::AbstractString, index::Integer, c::Char) = index>0 && index<=length(str) && str[index] == c

# Parse triples: i, i/j/k, i//k, i/j
""" TODO """
function parseTriple(line::AbstractString, vsize::Integer, vnsize::Integer, vtsize::Integer)
	vi = VertexIndex()

	tmp = line
	vi.v = parseIndex(tmp, vsize)
	token = search(tmp, r"[/\s]").start
	if !has(line,token,'/') return vi end
	token+=1

	# i//k
	if has(line,token,'/')
		token+=1
		tmp = line[token:end]
		vi.vn = parseIndex(tmp, vnsize)
		token = search(tmp, r"/\s").start
		return vi
	end

	# i/j/k or i/j
	tmp = line[token:end]
	vi.vt = parseIndex(tmp, vtsize)
	token = search(tmp, r"[/\s]").start

	if !has(line,token,'/') return vi
	else token+=1  # skip '/'
	end

	tmp = line[token:end]
	vi.vn = parseIndex(tmp, vnsize)
	token = search(tmp, r"[/\s]").start
	vi
end

""" TODO """
function updateVertex2(vi::VertexIndex, vCache::Dict{VertexIndex, Integer}, coord::ShapeCoordinates, cache::ShapeCoordinates)
	if haskey(vCache, vi) return vCache[vi] end

	if vi.v > 0 && vi.v > length(coord.positions)
		error("$(vi.v) > $(length(coord.positions))")
		return 0
	end

	v = coord.positions[vi.v]
	println("push position ", v)
	push!(cache.positions, v) #MatrixMath.scale(v,coord.scale) - coord.origin)

	if vi.vn > 0
		v = coord.normals[vi.vn]
		println("push normal ", v)
		push!(cache.normals, v)
	end

	if vi.vt > 0
		v = coord.texcoords[vi.vt]
		println("push texcoord ", v)
		push!(cache.texcoords, v)
	end

	idx = length(cache.positions)
	vCache[vi] = idx
	idx
end

const MinUInt32 = UInt32(0)
const MaxUInt32 = UInt32(2^32 - 1)

""" TODO """
correctVec(v::Vec3f, tmp::ShapeCoordinates) = MatrixMath.scale(v, tmp.scale) #MatrixMath.scale(v, tmp.scale) - tmp.origin

""" TODO """
function parseAll(
	vCache::Dict{VertexIndex,Integer},
	groups::AGroupIndex,
	#const int material_id,
	#const std::string &name,
	#bool clearCache,
	tmp::ShapeCoordinates,
	cache::ShapeCoordinates)

	clear(cache) # remove old

	cache.positions = map(x->correctVec(x, tmp),tmp.positions)
	cache.normals = tmp.normals
	cache.texcoords = tmp.texcoords

	for gs in groups
		for g in gs
			for index in g
				index += -1 #fix: index in file starts with 1 (not 0)
				index = index < MinUInt32 || index > MaxUInt32 ? MaxUInt32 : index
				push!(cache.indicies, UInt32(index))
				break # take only vertex indicies
			end
		end
	end

	# Flatten vertices and indices
	#=
	if length(group) < 1
		warn("group: $(length(group))")
		return false
	end

	println("indicies group $(length(group))")
	for i=1:length(group)
		face = group[i]

		i0 = face[1]
		i1 = VertexIndex()
		i2 = face[2]

		npolys = length(face)

		# Polygon -> triangle fan conversion
		for k=3:npolys
			i1 = i2
			i2 = face[k]

			v0 = updateVertex(i0, vCache, sc, cache)
			v1 = updateVertex(i1, vCache, sc, cache)
			v2 = updateVertex(i2, vCache, sc, cache)

			println("push index ", v0); push!(cache.indicies, v0)
			println("push index ", v1); push!(cache.indicies, v1)
			println("push index ", v2); push!(cache.indicies, v2)
		end
	end
	=#

	true
end

""" TODO """
function parseIndicies(tmp_vi::AAVertexIndex,tmp::ShapeCoordinates, cache::ShapeCoordinates)

	clear(cache) # remove old

	ig = 0
	iCount = length(tmp_vi)
	vCount = length(tmp.positions)
	vnCount = length(tmp.normals)
	vtCount = length(tmp.texcoords)

	if !iCount warn("No index coordinates! $tmp")	end

	if !vCount
		error("No vertex coordinates! $tmp")
		return false
	end

	# For each vertex of each triangle
	for it in tmp_vi
		ig+=1
		for jt in it
			vi = jt
			iv = vi.v
			ivn = vi.vn
			ivt = vi.vt

			if iv > 0
				if iv <= vCount
					push!(cache.positions, tmp.positions[iv])
				else
					error("Missmatch count of vertex coordinates! $ig ($iv <= 0 or >= $vCount).")
					return false
				end
			end

			if ivn > 0
				if ivn <= vnCount
					push!(cache.normals, tmp.normals[ivn])
				else
					error("Missmatch count of normal coordinates! $ig ($ivn <= 0 or >= $vnCount).")
					return false
				end
			end

			if ivt > 0
				if ivt <= vtCount
					push!(cache.texcoords, tmp.texcoords[ivt])
				else
					error("Missmatch count of texture coordinates! $ig ($ivt <= 0 or >= $vtCount).")
					return false
				end
			end
		end
	end

	vCount = length(cache.positions)
	vnCount = length(cache.normals)
	vtCount = length(cache.texcoords)

	if !vCount
		error("No vertex coordinates!")
		return false
	end

	if vnCount > 0 && vCount != vnCount
		error("Missmatch count of vertex and normal coordinates! ($vCount != $vnCount).")
		return false
	end

	if vtCount > 0 && vCount != vtCount
		error("Missmatch count of vertex and texture coordinates! ($vCount != $vtCount).")
		return false
	end

	true
end

""" TODO """
function loadOBJ(path::String, mesh::Mesh)
	println("Load OBJ file $path...");

	tmp_v = AVec3f()
	tmp_vn = AVec3f()
	tmp_vt = AVec2f()
	tmp_g = AGroupIndex()

	vertexCache = Dict{VertexIndex,Integer}()

	min_BBOX = zeros(Vec3f)
	max_BBOX = zeros(Vec3f)

	linenr = 0
	warn(string("read ", path))
	open(path) do file
		content=readstring(file)
		content=Base.replace(content, "\r", "")
		content=Base.replace(content, "\\\n", "") #remove broken lines
		#content=string(content,'\0')
		open(x -> println(x, content), string(path, ".cleaned"), "w+")
		for line in split(content,r"\n")
			linenr += 1
			line = lstrip(line)
			len = length(line) #+remove r
			id = len

			if len == 0 || line == "" continue end

			token = search(line, r"\s+").start
			if token <= 0
				#error(string(linenr, ": invalid token = ", token))
				warn(string(linenr, ": no separator -> skip"))
				continue
			end

			#println(string(linenr, ": search..."))

			first = line[1]
			second = len > 1 ? line[2] : ""
			third = len > 2 ? line[3] : ""

			if first == '\0' || first == '#' continue # skip line

			elseif first == 'v' && isSpace(second)
				part = line[token:end]
				#warn(string(linenr, ": ", line))
				v = parseVec(linenr, part, 3)

				# set bounding box
				if v.x < min_BBOX.x min_BBOX.x = v.x end
				if v.y < min_BBOX.y min_BBOX.y = v.y end
				if v.z < min_BBOX.z min_BBOX.z = v.z end
				if v.x > max_BBOX.x max_BBOX.x = v.x end
				if v.y > max_BBOX.y max_BBOX.y = v.y end
				if v.z > max_BBOX.z max_BBOX.z = v.z end

				push!(tmp_v,v)
				continue

			elseif first == 'v' && second == 'n' && isSpace(third)
				part = line[token:end]
				#warn(string(linenr, ": ", line))
				v = parseVec(linenr, part, 3)
				#push!(tmp_vn, v)
				continue

			elseif first == 'v' && second == 't' && isSpace(third)
				part = line[token:end]
				#warn(string(linenr, ": ", line))
				v = parseVec(linenr, part, 2)
				#push!(tmp_vt,v)
				continue

			elseif first == 'f' && isSpace(second)

				token += 1
				part = line[token:end]
				#warn(string(linenr, ": ", line))

				group = map((p)->map(function(x)
					v=parseVal(UInt32,x, -1)
					#println(linenr, ": index ", part, " => ", v)
					v
				end,split(p, r"/")),split(part, r"\s+"))

				#if j>0
				#	println("=> ", vi)
				#	push!(group,vi)
				#end

				#max = length(a) #search(line[token:end],r"[\n]").start
				#for i=1:max
				#	part = line[token:end]
				#	vi = parseTriple(part, Integer(round(length(tmp_v) / 3)), Integer(round(length(tmp_vn) / 3)), Integer(round(length(tmp_vt) / 2)))
				#	println(linenr, ": index ", part, " => ", vi)
				#	push!(group,vi)
				#	n = search(part, r"[ \t]").start
				#	token += n
				#end

				push!(tmp_g, group)
			end
		end
	end

	scale_value = getScale(min_BBOX, max_BBOX)
	println("scale_value: ",scale_value)
	tmp_sc = ShapeCoordinates( tmp_v, tmp_vn, tmp_vt, AIndex(), getOrigin(min_BBOX, max_BBOX, scale_value), Vec3f(scale_value) )
	cache = ShapeCoordinates()
	result = parseAll(vertexCache, tmp_g, tmp_sc, cache)

	if result
		#println("positions $(length(cache.positions)): ", cache.positions)
		#println("normals $(length(cache.normals)): ", cache.normals)
		#println("texcoords $(length(cache.texcoords)): ", cache.texcoords)
		#println("indicies $(length(cache.indicies)). ", cache.indicies)

		MeshManager.update(mesh, :MODEL, ()->(:TRIANGLES, cache.positions,cache.normals,cache.texcoords,cache.indicies))
		#shape.material_ids.push_back(material_id);
		#shape.name = name;
		#if (clearCache)	vCache.clear();
		println("Finsihed loading OBJ file $path.")
	end

	result
end

max_stars_repo_name: Gilga/JGE

max_issues_repo_name: Gilga/JGE

max_forks_repo_name: Gilga/JGE


text:
__precompile__()

module OrdinaryDiffEq

  using DiffEqBase

  using Compat

  import DiffEqBase: solve, solve!, init, step!, build_solution, initialize!

  using Parameters, GenericSVD, ForwardDiff, InplaceOps, RecursiveArrayTools,
        NLsolve, Juno, Calculus, Roots, DataStructures, Iterators

  import Base: linspace

  import Base: start, next, done, eltype

  import ForwardDiff.Dual

  import DiffEqBase: resize!,deleteat!,addat!,full_cache,user_cache,u_cache,du_cache,
                     resize_non_user_cache!,deleteat_non_user_cache!,addat_non_user_cache!,
                     terminate!,get_du, get_dt,get_proposed_dt,set_proposed_dt!,
                     u_modified!,savevalues!,add_tstop!,add_saveat!,set_reltol!,
                     set_abstol!

  #Constants

  const TEST_FLOPS_CUTOFF = 1e10

  include("misc_utils.jl")
  include("algorithms.jl")
  include("alg_utils.jl")
  include("caches.jl")
  include("integrators/unrolled_tableaus.jl")
  include("integrators/type.jl")
  include("integrators/integrator_utils.jl")
  include("integrators/fixed_timestep_integrators.jl")
  include("integrators/explicit_rk_integrator.jl")
  include("integrators/low_order_rk_integrators.jl")
  include("integrators/high_order_rk_integrators.jl")
  include("integrators/verner_rk_integrators.jl")
  include("integrators/feagin_rk_integrators.jl")
  include("integrators/implicit_integrators.jl")
  include("integrators/rosenbrock_integrators.jl")
  include("integrators/threaded_rk_integrators.jl")
  include("integrators/integrator_interface.jl")
  include("integrators/composite_integrator.jl")
  include("iterator_interface.jl")
  include("constants.jl")
  include("callbacks.jl")
  include("composite_solution.jl")
  include("solve.jl")
  include("initdt.jl")
  include("interp_func.jl")
  include("dense/generic_dense.jl")
  include("dense/interpolants.jl")
  include("dense/stiff_addsteps.jl")
  include("dense/low_order_rk_addsteps.jl")
  include("dense/verner_addsteps.jl")
  include("dense/high_order_rk_addsteps.jl")
  include("derivative_wrappers.jl")

  #General Functions
  export solve, solve!, init, step!

  #Callback Necessary
  export ode_addsteps!, ode_interpolant,
        terminate!, savevalues!, copyat_or_push!, isfsal

  export constructDP5, constructVern6, constructDP8, constructDormandPrince, constructFeagin10,
        constructFeagin12, constructFeagin14

  # Reexport the Alg Types

  export OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqCompositeAlgorithm,
        Discrete, FunctionMap, Euler, Midpoint, RK4, ExplicitRK, BS3, BS5, DP5, DP5Threaded, Tsit5,
        DP8, Vern6, Vern7, Vern8, TanYam7, TsitPap8, Vern9, ImplicitEuler,
        Trapezoid, Rosenbrock23, Rosenbrock32, Feagin10, Feagin12, Feagin14,
        CompositeAlgorithm
end # module


text:
struct CopositiveInnerBridge{T, S} <: MOIB.Variable.AbstractBridge
    matrix_variables::Vector{MOI.VariableIndex}
    matrix_constraint::MOI.ConstraintIndex{MOI.VectorOfVariables, S}
    nonneg_variables::Vector{MOI.VariableIndex}
    nonneg_constraint::MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.Nonnegatives}
end

function MOIB.Variable.bridge_constrained_variable(
    ::Type{CopositiveInnerBridge{T, S}},
    model::MOI.ModelLike, set::SOS.CopositiveInnerCone) where {T, S}

    side_dimension = MOI.side_dimension(set.psd_inner)
    num_off_diag = MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(side_dimension - 1))
    return CopositiveInnerBridge{T, S}(
        MOI.add_constrained_variables(model, set.psd_inner)...,
        MOI.add_constrained_variables(model, MOI.Nonnegatives(num_off_diag))...
    )
end

function MOIB.Variable.supports_constrained_variable(
    ::Type{<:CopositiveInnerBridge}, ::Type{<:SOS.CopositiveInnerCone})
    return true
end
function MOIB.added_constrained_variable_types(::Type{CopositiveInnerBridge{T, S}}) where {T, S}
    return [(S,), (MOI.Nonnegatives,)]
end
function MOIB.added_constraint_types(::Type{<:CopositiveInnerBridge})
    return Tuple{DataType, DataType}[]
end
function MOIB.Variable.concrete_bridge_type(
    ::Type{<:CopositiveInnerBridge{T}}, ::Type{SOS.CopositiveInnerCone{S}}) where {T, S}
    return CopositiveInnerBridge{T, S}
end

# Attributes, Bridge acting as a model
function MOI.get(bridge::CopositiveInnerBridge,
                 ::MOI.NumberOfVariables)
    return length(bridge.matrix_variables) + length(bridge.nonneg_variables)
end
function MOI.get(bridge::CopositiveInnerBridge,
                 ::MOI.ListOfVariableIndices)
    return Iterators.flatten((bridge.matrix_variables, bridge.nonneg_variables))
end
function MOI.get(bridge::CopositiveInnerBridge{T, S},
                 ::MOI.NumberOfConstraints{MOI.VectorOfVariables, S}) where {T, S}
    return 1
end
function MOI.get(bridge::CopositiveInnerBridge{T, S},
                 ::MOI.ListOfConstraintIndices{MOI.VectorOfVariables, S}) where {T, S}
    return [bridge.matrix_constraint]
end
function MOI.get(bridge::CopositiveInnerBridge,
                 ::MOI.NumberOfConstraints{MOI.VectorOfVariables, MOI.Nonnegatives})
    return 1
end
function MOI.get(bridge::CopositiveInnerBridge,
                 ::MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.Nonnegatives})
    return [bridge.nonneg_constraint]
end

# Indices
function MOI.delete(model::MOI.ModelLike, bridge::CopositiveInnerBridge)
    MOI.delete(model, bridge.matrix_variables)
    MOI.delete(model, bridge.nonneg_variables)
end

# Attributes, Bridge acting as a constraint

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::CopositiveInnerBridge)
    return SOS.CopositiveInnerCone(MOI.get(model, attr, bridge.matrix_constraint))
end

# TODO ConstraintPrimal, ConstraintDual

# See https://jump.dev/MathOptInterface.jl/v0.9.1/apireference/#MathOptInterface.AbstractSymmetricMatrixSetTriangle
function matrix_indices(k)
    j = div(1 + isqrt(8k - 7), 2)
    i = k - div((j - 1) * j, 2)
    return i, j
end
# Vector index for the vectorization of the triangular part.
function vector_index(i, j)
    return div((j - 1) * j, 2) + i
end
# Vector index for the vectorization of the off-diagonal triangular part.
function offdiag_vector_index(i, j)
    if i < j
        return vector_index(i, j - 1)
    else
        throw(ArgumentError())
    end
end

function MOI.get(model::MOI.ModelLike, attr::MOI.VariablePrimal,
                 bridge::CopositiveInnerBridge, i::MOIB.IndexInVector)
    value = MOI.get(model, attr, bridge.matrix_variables[i.value])
    row, col = matrix_indices(i.value)
    if row != col
        value += MOI.get(model, attr, bridge.nonneg_variables[offdiag_vector_index(i, j)])
    end
    return value
end

function MOIB.bridged_function(bridge::CopositiveInnerBridge{T},
                               i::MOIB.IndexInVector) where T
    func = convert(MOI.ScalarAffineFunction{T},
                   bridge.matrix_variables[i.value])
    row, col = matrix_indices(i.value)
    if row != col
        func = MOIU.operate!(+, T, func,
            bridge.nonneg_variables[vector_index(row, col - 1)])
    end
    return func
end
function MOIB.Variable.unbridged_map(
    bridge::CopositiveInnerBridge{T},
    vi::MOI.VariableIndex, i::MOIB.IndexInVector) where T

    F = MOI.ScalarAffineFunction{T}
    func = convert(F, vi)
    map = bridge.matrix_variables[i.value] => func
    row, col = matrix_indices(i.value)
    if row == col
        return (map,)
    else
        nneg = bridge.nonneg_variables[vector_index(row, col - 1)]
        return (map, nneg => zero(F))
    end
end

max_stars_repo_name: blegat/SumOfSquares

max_issues_repo_name: blegat/SumOfSquares

max_forks_repo_name: blegat/SumOfSquares


text:
include("KalmanFilter.jl")

export EM_step
export EM_convergence

"""
    This function reestimates parameters based on the Estimation Maximization (EM) algorithm. This is a two-step procedure:
    (1) E-step: the expectation of the log-likelihood is calculated using previous parameter estimates.
    (2) M-step: Parameters are re-estimated through the maximisation of the log-likelihood (maximise result from (1)).
    See "Maximum likelihood estimation of factor models on data sets with arbitrary pattern of missing data" for details about parameter derivation (Banbura & Modugno, 2010).

    parameters:
        y: Array
            data (standardized and ready to be used)
        A: Array
            transition matrix
        C: Array
            measurement matrix
        Q: Array
            covariance for transition equation residuals
        R: Array
            covariance for measurement equation residuals
        Z0: Array
            initial values of state
        V0: Array
            initial value of factor covariance matrix
        p: Int
            number of lags in transition equation
        blocks: Array
            block structure for each series
        R_mat: Array
            estimation structure for quarterly variables (i.e. "tent")
        q: Array
            constraints on loadings
        nM: Int
            number of monthly variables
        monthly_quarterly_array: Array
            indices for quarterly variables
    returns: Dict
        A_new => updated transition matrix
        C_new => updated measurement matrix
        Q_new => updated covariance matrix for residuals for transition matrix
        R_new => updated covariance matrix for residuals of measurement equation
        Z0 => initial value of state
        V0 => initial value of factor covariance matrix
        loglik => log likelihood
"""
function EM_step(y_est; A, C, Q, R, Z0, V0, p, blocks, R_mat, q, nM, monthly_quarterly_array)
    n = size(y_est)[1]
    n_obs = size(y_est)[2]
    pC = size(R_mat)[2]
    ppC = maximum([p, pC])
    n_blocks = size(blocks)[2]
    nM = sum(.!monthly_quarterly_array)
    y = copy(y_est)

    output = kalman_filter(y_est, A, C, Q, R, Z0, V0)
    Zsmooth = output[:Zsmooth]; Vsmooth = output[:Vsmooth]; VVsmooth = output[:VVsmooth]; loglik = output[:loglik]

    ### MAXIMIZATION STEP (TRANSITION EQUATION). See Banbura & Modugno (2010) for details.
    A_new = copy(A)
    Q_new = copy(Q)
    V0_new = copy(V0)

    # 2A. Update factor parameters individually
    for i in 1:n_blocks # Loop for each block: factors are uncorrelated
        # Setup indexing
        p1 = (i - 1) * ppC
        b_subset = (p1 + 1):(p1 + p) # Subset blocks: Helps for subsetting Zsmooth, Vsmooth
        t_start = p1 + 1 # Transition matrix factor idx start
        t_end = p1 + ppC # Transition matrix factor idx end

        # Estimate factor portion of Q, A. Note: EZZ, EZZ_BB, EZZ_FB are parts of equations 6 and 8 in BM 2010
        # E[f_t * f_t' | Omega_T]
        EZZ = (Zsmooth[b_subset, 2:end] * transpose(Zsmooth[b_subset, 2:end])) |> x->
            x + reshape(sum(Vsmooth[b_subset, b_subset, 2:end], dims=3), size(x)...)
        # E[f_{t-1} * f_{t-1}' | Omega_T]
        EZZ_BB = (Zsmooth[b_subset, 1:(end-1)] * transpose(Zsmooth[b_subset, 1:(end-1)])) |> x->
            x + reshape(sum(Vsmooth[b_subset, b_subset, 1:(end-1)], dims=3), size(x)...)
        # E[f_t * f_{t-1}' | Omega_T]
        EZZ_FB = (Zsmooth[b_subset, 2:end] * transpose(Zsmooth[b_subset, 1:(end-1)])) |> x->
            x  + reshape(sum(VVsmooth[b_subset,b_subset,:], dims=3), size(x)...)
        # Select transition matrix/covariance matrix for block i
        Ai = A[t_start:t_end, t_start:t_end]
        Qi = Q[t_start:t_end, t_start:t_end]
        Ai[1,1:p] = EZZ_FB[1,1:p] * inv(EZZ_BB[1:p, 1:p]) # Equation 6: Estimate VAR(p) for factor
        Qi[1,1] = ((EZZ[1,1] .- Ai[1, 1:p] * transpose(EZZ_FB[1, 1:p])) / n_obs)[1] # Equation 8: Covariance matrix of residuals of VAR
        # Place updated results in output matrix
        A_new[t_start:t_end,t_start:t_end] = copy(Ai)
        Q_new[t_start:t_end,t_start:t_end] = copy(Qi)
        V0_new[t_start:t_end,t_start:t_end] = Vsmooth[t_start:t_end, t_start:t_end,1]
    end

    # 2B. Update parameters for idiosyncratic component
    rp1 = n_blocks * ppC # Col size of factor portion
    t_start = rp1 + 1 # Start of idiosyncratic component index
    i_subset = t_start:(rp1 + nM) # Gives indices for monthly idiosyncratic component values

    #% The three equations below estimate the idiosyncratic component (for eqns 6, 8 BM 2010)
    # E[f_t * f_t' | \Omega_T]
    EZZ = (diag(Zsmooth[t_start:end, 2:end] * transpose(Zsmooth[t_start:end, 2:end])) +
        diag(sum(Vsmooth[t_start:end, t_start:end, 2:end], dims = 3) |> x-> reshape(x, size(x)[1:2]...))) |>
        diagm
    # E[f_{t-1} * f_{t-1}' | \Omega_T]
    EZZ_BB = (diag(Zsmooth[t_start:end, 1:(end-1)] * transpose(Zsmooth[t_start:end, 1:(end-1)])) +
        diag(sum(Vsmooth[t_start:end, t_start:end, 1:(end-1)], dims = 3) |> x-> reshape(x, size(x)[1:2]...))) |>
        diagm
    # E[f_t * f_{t-1}' | \Omega_T]
    EZZ_FB = (diag(Zsmooth[t_start:end, 2:end] * transpose(Zsmooth[t_start:end, 1:(end-1)])) +
        diag(sum(VVsmooth[t_start:end, t_start:end, :], dims = 3) |> x-> reshape(x, size(x)[1:2]...))) |>
        diagm

    Ai = EZZ_FB * diagm(transpose(1 / diag(EZZ_BB))) # Equation 6
    Qi = (EZZ - Ai * transpose(EZZ_FB)) / n_obs # Equation 8

    # Place updated results in output matrix
    A_new[i_subset, i_subset] = Ai[1:nM, 1:nM]
    Q_new[i_subset, i_subset] = Qi[1:nM, 1:nM]
    V0_new[i_subset, i_subset] = diagm(diag(Vsmooth[i_subset, i_subset, 1]))

    # 3. Maximization step (measurement equation)
    Z0 = Zsmooth[:,1]
    # Set missing data series values to 0
    y[ismissing.(y)] .= 0
    C_new = copy(C) # loadings
    bl = unique(blocks) # Gives unique loadings
    n_bl = size(bl)[1] # Number of unique loadings
    # Initialize indices: these later help with subsetting
    bl_idxM = [] # Indicator for monthly factor loadings
    bl_idxQ = [] # Indicator for quarterly factor loadings
    R_con = [] # Block diagonal matrix giving monthly-quarterly aggreg scheme
    q_con = []
    # Loop through each block
    for i in 1:n_blocks
        push!(bl_idxQ, repeat(bl[:,i], 1, ppC))
        push!(bl_idxM, hcat(repeat(bl[:,i], 1, 1), zeros(n_bl, ppC - 1)))
        push!(R_con, R_mat)
        push!(q_con, zeros(size(R_mat)[1], 1))
    end
    bl_idxQ = hcat(bl_idxQ...) .== 1.0
    bl_idxM = hcat(bl_idxM...) .== 1.0
    R_con = blockdiag(sparse.(R_con)...) |> collect
    q_con = vcat(q_con...)

    # Indicator for monthly/quarterly blocks in measurement matrix
    index_freq_M = .!monthly_quarterly_array[1:nM] # Gives 1 for monthly series
    n_index_M = length(index_freq_M) # Number of monthly series
    c_index_freq = cumsum(.!monthly_quarterly_array) # Cumulative number of monthly series

    for i in 1:n_bl # Loop through unique loadings
        bl_i = bl[i,:]
        rs = sum(bl_i) # Total num of blocks loaded, or number of quarterly variables?
        idx_i = findall(x->x == true, [x == bl_i for x in eachrow(blocks)]) # Indices for bl_i (which variables in block)
        idx_iM = idx_i[idx_i .<= nM] # Only monthly series
        n_i = length(idx_iM) # Number of monthly series
        # Initialize sums in equation 13 of BGR 2010
        denom = zeros(Int(n_i * rs), Int(n_i * rs))
        nom = zeros(Int(n_i), Int(rs))

        # Stores monthly indicies. These are done for input robustness
        index_freq_i = index_freq_M[idx_iM]
        index_freq_ii = c_index_freq[idx_iM]
        index_freq_ii = index_freq_ii[index_freq_i]

        # Update monthly variables: loop through each period
        for t in 1:n_obs
            Wt = diagm(.!ismissing.(y_est)[idx_iM,t]) * 1 # Gives selection matrix (TRUE for nonmissing values)
            global bl_idxM_ext = vcat(bl_idxM[i,:], repeat([false], size(Zsmooth)[1] - length(bl_idxM)))
            # E[f_t * t_t' | Omega_T]
            denom = denom + kron(Zsmooth[bl_idxM_ext, t + 1] * transpose(Zsmooth[bl_idxM_ext, t + 1]) + Vsmooth[bl_idxM_ext, bl_idxM_ext, t + 1], Wt)
            # E[y_t * f_t' | Omega_T]
            nom = nom + y[idx_iM, t] *
                transpose(Zsmooth[bl_idxM_ext, t + 1]) -
                Wt[:,index_freq_i] * (Zsmooth[rp1 .+ index_freq_ii, t + 1] *
                transpose(Zsmooth[bl_idxM_ext, t + 1]) +
                Vsmooth[rp1 .+ index_freq_ii, bl_idxM_ext, t + 1])
        end
        vec_C = inv(denom) * reshape(nom, prod(collect(size(nom)), dims=1)[1]) # Eqn 13 BGR 2010
        C_new[idx_iM, bl_idxM_ext] = reshape(vec_C, Int(n_i), Int(rs)) # Place updated monthly results in output matrix
        idx_iQ = idx_i[idx_i .> nM] # Index for quarterly series
        rps = rs * ppC
        # Monthly-quarterly aggregation scheme
        R_con_i = R_con[:,bl_idxQ[i,:]]
        q_con_i = copy(q_con)
        no_c = [!all(x->Int(x)==0,row) for row in eachrow(R_con_i)] # any non-zero values in row
        R_con_i = R_con_i[no_c,:]
        q_con_i = q_con_i[no_c,:]

        # Loop through quarterly series in loading, this parallels monthly code. This has to be checked for working with multiple quarterly variables.
        for j in idx_iQ
            # Initialization
            denom = zeros(Int(rps), Int(rps))
            nom = zeros(1, Int(rps))
            idx_jQ = j - nM # Ordinal position of quarterly variable
            # Location of factor structure corresponding to quarterly variable residuals
            index_freq_jQ = (rp1 + n_index_M + 5 * (idx_jQ - 1) + 1):(rp1 + n_index_M + 5 * idx_jQ)

            # Place quarterly values in output matrix
            V0_new[index_freq_jQ, index_freq_jQ] = Vsmooth[index_freq_jQ, index_freq_jQ, 1]
            A_new[index_freq_jQ[1], index_freq_jQ[1]] = Ai[index_freq_jQ[1] - rp1, index_freq_jQ[1] - rp1]
            Q_new[index_freq_jQ[1], index_freq_jQ[1]] = Qi[index_freq_jQ[1] - rp1, index_freq_jQ[1] - rp1]
            # Update quarterly variables: loop through each period
            for t in 1:n_obs
                Wt = diagm(.!ismissing.(y_est)[j:j,t]) * 1 # Selection matrix for quarterly values, double check this used to be indx_iQ, had to change to work with multiple quarterly variables.
                # Intermediate steps in BGR equation 13
                global bl_idxQ_ext = vcat(bl_idxQ[i,:], repeat([false], size(Zsmooth)[1] - size(bl_idxQ)[2]))
                denom = denom + kron(Zsmooth[bl_idxQ_ext, t + 1] * transpose(Zsmooth[bl_idxQ_ext, t + 1]) +  Vsmooth[bl_idxQ_ext, bl_idxQ_ext, t + 1], Wt)
                nom = nom + y[j,t] * transpose(Zsmooth[bl_idxQ_ext, t + 1])
                nom = nom - Wt * ([1 2 3 2 1] * Zsmooth[index_freq_jQ, t + 1] * transpose(Zsmooth[bl_idxQ_ext, t + 1]) + [1 2 3 2 1] * Vsmooth[index_freq_jQ, bl_idxQ_ext, t + 1])
            end
            C_i = inv(denom) * transpose(nom)
            # BGR equation 13
            C_i_constr = C_i - inv(denom) * transpose(R_con_i) * inv(R_con_i * inv(denom) * transpose(R_con_i)) * (R_con_i * C_i - q_con_i)
            C_new[j, bl_idxQ_ext] = C_i_constr # Place updated values in output structure
        end
    end

    # 3B. Update covariance of residuales for measurement equation
    # Initialize covariance of residuals of observation equation
    R_new = zeros(Int(n), Int(n))
    for t in 1:n_obs
        Wt = diagm(.!ismissing.(y_est)[:,t]) * 1.0 # Selection matrix
        # BGR equation 15
        R_new = R_new + (y[:,t] - Wt * C_new * Zsmooth[:,t + 1]) * transpose(y[:,t] - Wt * C_new * Zsmooth[:,t + 1]) +
            Wt * C_new * Vsmooth[:,:,t + 1] * transpose(C_new) * Wt + (collect(I(n)) - Wt) * R * (collect(I(n)) - Wt)
    end
    RR = diag(R_new)
    RR[.!monthly_quarterly_array] .= 1e-4 # Ensure non-zero measurement error, see Doz, Giannone, Reichlin (2012) for reference
    RR[(nM + 1):end] .= 1e-4
    R_new = diagm(RR)

    # final return
    return Dict(
        :A_new => A_new,
        :C_new => C_new,
        :Q_new => Q_new,
        :R_new => R_new,
        :Z0 => Z0,
        :V0 => V0,
        :loglik => loglik
        )
end

"""
  This function checks whether EM has converged. Convergence occurs if the slope of the log-likelihood function falls below 'threshold' (i.e. f(t) - f(t-1)| / avg < threshold) where avg = (|f(t)| + |f(t-1)|)/2 and f(t) is log likelihood at iteration t.
      This stopping criterion is from Numerical Recipes in C (pg. 423). With MAP estimation (using priors), the likelihood can decrease even if the mode of the posterior increases.
  (1) E-step: the expectation of the log-likelihood is calculated using previous parameter estimates.
  (2) M-step: Parameters are re-estimated through the maximisation of the log-likelihood (maximise result from (1)).
  See "Maximum likelihood estimation of factor models on data sets with arbitrary pattern of missing data" for details about parameter derivation (Banbura & Modugno, 2010).

  parameters:
      loglik: Float
          log-likelihood from current EM iteration
      prev_loglik: Float
          log-likelihood from previous EM iteration
      threshold: Float
          convergence threshhold
  returns: Dict
      converged => 1 if convergence criteria satistifed, 0 otherwise
      decrease => 1 if log likelihood decreased, 0 otherwise
"""
function EM_convergence(loglik, prev_loglik, threshold)::Dict{Symbol, Int64}
    # initialization
    converged = 0
    decrease = 0

    if (loglik - prev_loglik) < -1e-3 # allows for imprecision
        decrease = 1
    end
    # check convergence criteria
    delta = abs(loglik - prev_loglik)
    avg_loglik = (abs(loglik) + abs(prev_loglik) + eps()) / 2

    if ((delta / avg_loglik) < threshold)
        converged = 1
    end

    return Dict(
        :converged => converged,
        :decrease => decrease
    )
end


text:
using Raylib_jll

struct RayColor
    r::Cuchar
    g::Cuchar
    b::Cuchar
    a::Cuchar
end

const RAYWHITE = RayColor(245, 245, 245, 255)
const LIGHTGRAY = RayColor(200, 200, 200, 255)
const DARKGRAY = RayColor(80, 80, 80, 255)
const MAROON = RayColor(190, 33, 55, 255)
const RED = RayColor(230, 41, 55, 255)
const BLUE = RayColor(0, 121, 241, 255 )
const LIME = RayColor(0, 158, 47, 255)
const GOLD = RayColor(255, 203, 0, 255)
const SKYBLUE = RayColor(102, 191, 255, 255)
const BLACK = RayColor(0, 0, 0, 255)

struct RayVector3
    x::Cfloat
    y::Cfloat
    z::Cfloat
end

struct RayVector2
    x::Cfloat
    y::Cfloat
end

struct RayCamera3D
    position::RayVector3
    target::RayVector3
    up::RayVector3
    fovy::Cfloat
    projection::Cint
end

function main()
    screenWidth = 800
    screenHeight = 450

    @ccall libraylib.InitWindow(screenWidth::Cint, screenHeight::Cint,
                             "raylib [core] example - 3d camera first person"::Cstring)::Cvoid

    camera = Ref(RayCamera3D(
        RayVector3(4, 2, 4),
        RayVector3(0, 1.8, 0),
        RayVector3(0, 1, 0),
        60,
        0,
    ))

    heights = Vector{Cfloat}(undef, 20)
    positions = Vector{RayVector3}(undef, 20)
    colors = Vector{RayColor}(undef, 20)

    for i = 1:20
        heights[i] = @ccall libraylib.GetRandomValue(1::Cint, 12::Cint)::Cint
        positions[i] = RayVector3(
            @ccall(libraylib.GetRandomValue((-15)::Cint, 15::Cint)::Cint),
            heights[i] / 2,
            @ccall(libraylib.GetRandomValue((-15)::Cint, 15::Cint)::Cint),
        )
        colors[i] = RayColor(
            @ccall( libraylib.GetRandomValue(20::Cint, 255::Cint)::Cint),
            @ccall( libraylib.GetRandomValue(10::Cint, 55::Cint)::Cint),
            30,
            255,
        )
    end

    @ccall libraylib.SetCameraMode(camera[]::RayCamera3D, 3::Cint)::Cvoid
    @ccall libraylib.SetTargetFPS(60::Cint)::Cvoid

    while iszero(@ccall libraylib.WindowShouldClose()::Cint)
        @ccall libraylib.UpdateCamera(camera::Ref{RayCamera3D})::Cvoid

        @ccall libraylib.BeginDrawing()::Cvoid
        @ccall libraylib.ClearBackground(RAYWHITE::RayColor)::Cvoid

        @ccall libraylib.BeginMode3D(camera[]::RayCamera3D)::Cvoid

        @ccall libraylib.DrawPlane(
            RayVector3(0,0,0)::RayVector3,
            RayVector2(32, 32)::RayVector2,
            LIGHTGRAY::RayColor)::Cvoid

        @ccall libraylib.DrawCube(
            RayVector3(-16, 2.5, 0)::RayVector3, 1f0::Cfloat, 5f0::Cfloat, 32f0::Cfloat,
            BLUE::RayColor)::Cvoid
        @ccall libraylib.DrawCube(
            RayVector3(16, 2.5, 0)::RayVector3, 1f0::Cfloat, 5f0::Cfloat, 32f0::Cfloat,
            LIME::RayColor)::Cvoid
        @ccall libraylib.DrawCube(
            RayVector3(0, 2.5, 16)::RayVector3, 32f0::Cfloat, 5f0::Cfloat, 1f0::Cfloat,
            GOLD::RayColor)::Cvoid

        for i = 1:20
            @ccall libraylib.DrawCube(
                positions[i]::RayVector3, 2f0::Cfloat, heights[i]::Cfloat, 2f0::Cfloat,
                colors[i]::RayColor)::Cvoid
            @ccall libraylib.DrawCubeWires(positions[i]::RayVector3, 2f0::Cfloat, heights[i]::Cfloat, 2f0::Cfloat,
                                        MAROON::RayColor)::Cvoid

        end

        @ccall libraylib.EndMode3D()::Cvoid

        @ccall libraylib.DrawRectangle(10::Cint, 10::Cint, 220::Cint, 70::Cint,
                                    @ccall( libraylib.Fade(SKYBLUE::RayColor, 0.5::Cfloat)::RayColor )::RayColor)::Cvoid
        @ccall libraylib.DrawRectangleLines(10::Cint, 10::Cint, 220::Cint, 70::Cint, BLUE::RayColor)::Cvoid

        @ccall libraylib.DrawText("First person camera default controls:"::Cstring,
                               20::Cint, 20::Cint, 10::Cint, BLACK::RayColor)::Cvoid
        @ccall libraylib.DrawText("- Move with keys: W, A, S, D"::Cstring,
                               40::Cint, 40::Cint, 10::Cint, DARKGRAY::RayColor)::Cvoid
        @ccall libraylib.DrawText("- Mouse move to look around"::Cstring,
                               40::Cint, 60::Cint, 10::Cint, DARKGRAY::RayColor)::Cvoid

        @ccall libraylib.EndDrawing()::Cvoid
    end

    @ccall libraylib.CloseWindow()::Cvoid

end


text:
"""
    cm()

Returns the path to `Rfam.cm` file containing the covariance models of all the families.
"""
function cm()
    local_path = joinpath(cache, "Rfam.cm")
    if !isfile(local_path)
        @info "Downloading Rfam.cm to $local_path ..."
        download(rfam_baseurl * "/Rfam.cm.gz", "$local_path.gz"; timeout = Inf)
        gunzip(local_path * ".gz")
    end
    return local_path
end

"""
    clanin()

Returns the path to `Rfam.clanin`.
"""
function clanin()
    local_path = joinpath(cache, "Rfam.clanin")
    if !isfile(local_path)
        @info "Downloading Rfam.clanin to $local_path ..."
        download("$rfam_baseurl/Rfam.clanin", "$local_path"; timeout = Inf)
    end
    return local_path
end

"""
    seed()

Returns the path to `Rfam.seed`.
"""
function seed()
    local_path = joinpath(cache, "Rfam.seed")
    if !isfile(local_path)
        @info "Downloading Rfam.seed to $local_path ..."
        download("$rfam_baseurl/Rfam.seed.gz", "$local_path.gz"; timeout = Inf)
        gunzip("$local_path.gz")
    end
    return local_path
end