text:
#   Main Module of the Diagonalization.jl Package for Julia language
#   v 0.2.1 - last update 25th of June 2021.

#   MIT License
#   Copyright (c) 2019-2021,
#   Marco Congedo, CNRS, Grenoble, France:
#   https://sites.google.com/site/marcocongedo/home

# TODO
# - svd full or not?

module Diagonalizations

using LinearAlgebra, Statistics, StatsBase, CovarianceEstimation,
      Base.Threads, PosDefManifold

# Special instructions and variables
BLAS.set_num_threads(Sys.CPU_THREADS)

import Base:size, length, eltype, ==, ≈, ≠, ≉

export
   SCM, LShr, LShrLW, NShrLW, # constants for CovarianceEstimator types
   LinearFilters,
   LinearFilter, LF,          # Structs and its alias
   # methods for the above Structs
   size, cut, length, eltype, ==, ≈, ≠, ≉,
   eig,                       # utilities
   nonDiagonality, nonD,
   spForm,
   genDataMatrix,
   pca,                       # linear filters
   whitening,
   csp,
   cstp,
   mca,
   cca,

   ajd,
   gmca,
   gcca,
   majd,

   _cov # temp

# Consts
const titleFont     = "\x1b[35m"
const separatorFont = "\x1b[95m"
const defaultFont   = "\x1b[0m"
const greyFont      = "\x1b[90m"
const 📌            = titleFont*"Diagonalizations.jl"*defaultFont

const SCM=SimpleCovariance()
const LShr=LinearShrinkage
const LShrLW=LShr(ConstantCorrelation())
const NShrLW=AnalyticalNonlinearShrinkage()

# Types
○          = nothing
SorH       = Union{Symmetric, Hermitian}
VecSorH    = Union{Array{Symmetric, 1}, Array{Hermitian, 1}}
SorHo      = Union{SorH, Nothing}
Vec        = Vector{T} where T<:Union{Real, Complex}
Veco       = Union{Vector{T}, Nothing} where T<:Union{Real, Complex}
Mat        = Matrix{T} where T<:Union{Real, Complex}
Mato       = Union{Matrix{T}, Nothing} where T<:Union{Real, Complex}
VecMat     = Vector{Matrix{T}} where T<:Union{Real, Complex}
VecMato    = Union{Vector{Matrix{T}}, Nothing} where T<:Union{Real, Complex}
VecVecMat  = Vector{Vector{Matrix{T}}} where T<:Union{Real, Complex}
Diagonalo  = Union{Diagonal, Nothing}
Float64o   = Union{Float64, Nothing}
Stringo    = Union{String, Nothing}
Tmean      = Union{Int64, AbstractVector, AbstractMatrix, Nothing}
Into       = Union{Int64, Nothing}
Tw         = Union{StatsBase.AbstractWeights, Nothing}
Twf        = Union{Tw, Function}
TeVar      = Union{Float64, Int64}
TeVaro     = Union{TeVar, Nothing}

abstract type LinearFilters end

# ===============================================================#
# Structure for all kinds of linear filter
struct LinearFilter <: LinearFilters
   F     :: AbstractArray         # filter
   iF    :: AbstractArray         # left-inverse of the filter
   D     :: Diagonalo   # Diag. Matrix of ev associated with the filter
   eVar  :: Float64o    # explained variance (0, 1]
   ev    :: Veco        # vector of all eigenvalues
   arev  :: Veco        # vector of all accumulated regularized eigenvalues
   name  :: String      # name of the filter

   # Universal constructor of (simple) filters
   function LinearFilter(F::Mat, iF::Mat, D::Diagonalo=○, eVar::Float64o=○,
               ev::Veco=○, arev::Veco=○, name::Stringo=○, check::Bool=true)

     ℹ = 📌*": Simple Linear Filter construction: "
     size(F) ≠ reverse(size(iF))         && throw(DimensionMismatch(ℹ*"
       the size of matrices `F`($(size(F))) and `iF`($(size(iF))) must be one the reverse of the other"))
     (check && iF*F≉ I)                  && throw(ArgumentError(ℹ*"
       matrix `iF` is not a left-inverse of matrix `F`"))
     (D≠○ && size(F, 2) ≠ size(D, 2))    && throw(DimensionMismatch(ℹ*"
       the dimensions of `D`($(size(D, 2))) must be equal to the number of columns of matrix `F`($(size(F, 2)))"))
     (ev≠○ && length(ev)>size(F, 1)) && throw(DimensionMismatch(ℹ*"
       the length of `ev`($(length(ev))) cannot exceed the number of rows of matrix `F`($(size(F, 1)))"))
     (arev≠○ && size(F, 1) < length(arev)) && throw(DimensionMismatch(ℹ*"
       the length of `arev`($(length(arev))) cannot exceed the number of rows of matrix `F`($(size(F, 1)))"))

     new(F, iF, D, eVar, ev, arev, name===○ ? "Linear Filter" : name)
   end

   # minimal constructor given only the filter matrix
   LinearFilter(F::Mat, name::Stringo=○) =
      LinearFilter(F, pinv(F), ○, ○, ○, ○, name===○ ? "Linear Filter" : name, false)

   # Universal constructor of joint linear filters
   function LinearFilter(F::VecMat, iF::VecMat,
                         D::Diagonalo, eVar::Float64o=○, ev::Veco=○,
                         arev::Veco=○, name::Stringo=○, check::Bool=true)

     ℹ = 📌*": Joint Linear Filter construction: "
     length(F) ≠ length(iF) && throw(DimensionMismatch(ℹ*"
       The number of matrices in `F`($(length(F))) and `iF`($(length(iF))) must be the same "))
     for i=1:length(F)
       size(F[i]) == reverse(size(iF[i])) && continue
       size(F[i]) ≠ reverse(size(iF[i]))  && throw(DimensionMismatch(ℹ*"
       the size of matrices `F[$i]`($(size(F[i]))) and `iF[$i]`($(size(iF[i]))) must be one the reverse of the other"))
       # BREAK!
     end
     if check for i=1:length(F) iF*F≉ I   && throw(DimensionMismatch(ℹ*"
       matrix `iF[$i]` is not a left-inverse of matrix `F[$i]`"))
     end end
     n2=minimum(size(f, 2) for f ∈ F)
     (D≠○ && n2 ≠ size(D, 2)) && throw(DimensionMismatch(ℹ*"
       the dimension of `D`($(size(D, 2))) must be equal to the minimum number of columns of the matrices in `F`($n2)"))
     n1=minimum(size(f, 1) for f ∈ F)
     (ev≠○ && length(ev)>n1) && throw(DimensionMismatch(ℹ*"
       the length of `ev`($(length(ev))) cannot exceed the minimum number of rows of the matrices in `F`($n1)"))
     (arev≠○ && length(arev)>n1) && throw(DimensionMismatch(ℹ*"
       the length of `arev` ($(length(arev))) cannot exceed the minimum number of rows of the matrices in `F`($n1)"))
     new(F, iF, D, eVar, ev, arev, name===○ ? "Joint Linear Filter" : name)
   end

   # minimal constructor given only the filters matrix
   LinearFilter(F::VecMat, name::Stringo=○) =
      LinearFilter(F, [pinv(X) for X ∈ F],
                   ○, ○, ○, ○, name===○ ? "Linear Filter" : name, false)
end # Struct LinearFilter

LF=LinearFilter # alias

size(f::LF) = f.F isa Matrix ? size(f.F) : (size(f.F[i]) for i=1:length(f.F))

length(f::LF) = f.F isa Matrix ? 1 : length(f.F)

eltype(f::LF) = f.F isa Matrix ? eltype(f.F) : eltype(f.F[1])

function ==(f::LF, g::LF)
   ((f.F isa Matrix) ≠ (g.F isa Matrix)) && return false
   length(f.F) ≠ length(g.F) && return false
   arevOK=(f.arev≠○ && g.arev≠○) ? f.arev≈g.arev   : true
   evOK  =(f.ev≠○ && g.ev≠○)     ? f.ev≈g.ev       : true
   eVarOK=(f.eVar≠○ && g.eVar≠○) ? f.eVar≈g.eVar   : true
   DOK   =(f.D≠○ && g.D≠○)       ? f.D≈g.D         : true
   if f.F isa Matrix
      FOK   = ((spForm(f.iF*g.F)+spForm(g.iF*f.F))/2)<0.01
   else
      FOK   = (mean(spForm(f.iF[i]*g.F[i]) for i=1:length(g.F)) +
               mean(spForm(g.iF[i]*f.F[i]) for i=1:length(f.F)))/2 <0.01
   end
   f.name==g.name && arevOK && evOK && eVarOK && DOK && FOK
end
≈(f::LF, g::LF)= ==(f, g)

≠(f::LF, g::LF)= !==(f, g)
≉(f::LF, g::LF)= !≈(f, g)

# Given a filter, create a reduced filter with dimension p
#### TOUT DOUX: allow passing an eVar value
function cut(f::LinearFilter, p::Int64)
   p=clamp(p, 1, length(f.ev))
   h=1:p
   if f.F isa Matrix
      LF(f.F[:, h], f.iF[h, :], f.D===○ ? ○ : Diagonal(diag(f.D)[h]),
         f.eVar===○ ? ○ : f.eVar[p], f.ev, f.arev, f.name*"($p)", false)
   else
      LF([X[:, h] for X ∈ f.F], [X[h, :] for X ∈ f.iF],
          f.D===○ ? ○ : Diagonal(diag(f.D)[h]), f.eVar===○ ? ○ : f.eVar[p],
          f.ev, f.arev, f.name*"($p)", false)
   end
end
# ===============================================================#

include("tools.jl")
include("pca.jl")
include("csp.jl")
include("cca.jl")
include("../src/optim/common.jl")
include("../src/optim/JoB.jl")
include("../src/optim/LogLike.jl")
include("../src/optim/Jade.jl")
include("../src/optim/Gajd.jl")
include("../src/optim/QnLogLike.jl")
include("ajd.jl")
include("gcca.jl")

# override the Base.size method
## size(f::LinearFilters) = size(f.p) ## TODO field p in the LF struct
# ++++++++++++++++++++  Show override  +++++++++++++++++++ # (REPL output)
function Base.show(io::IO, ::MIME{Symbol("text/plain")}, f::LinearFilter)
  if f.F isa Matrix
     f.name===○ ? println(io, titleFont, "\n⋱ Linear Filter") :
                  println(io, titleFont, "\n⋱ "*f.name)
     println(io, separatorFont, "⤔   ⤔    ⤔     ⤔      ⤔")
     println(io, separatorFont,".F    ", defaultFont, "matrix{$(eltype(f.F))}$((size(f.F)))")
     println(io, separatorFont,".iF   ", defaultFont, "matrix{$(eltype(f.iF))}$(size(f.iF))")
  else
     f.name===○ ? println(io, titleFont, "\n⋱ Joint Linear Filter") :
                  println(io, titleFont, "\n⋱ "*f.name)
     println(io, separatorFont, "⤕   ⤕    ⤕     ⤕      ⤕")
     println(io, separatorFont,".F    ", defaultFont, "$(length(f.F))-vector of matrices{$(eltype(f.F[1]))}$((size(f.F[1])))")
     println(io, separatorFont,".iF   ", defaultFont, "$(length(f.iF))-vector of matrices{$(eltype(f.iF[1]))}$((size(f.iF[1])))")
  end
  f.D===○ ?    nothing : println(io, separatorFont,".D    ", defaultFont, "diagonal{$(eltype(f.D))}($(size(f.D)))")
  f.eVar===○ ? nothing : println(io, separatorFont,".eVar ", defaultFont, "explained variance=$(round(f.eVar, digits=3))")
  f.ev===○ ?   nothing : println(io, separatorFont,".ev   ", defaultFont, "vector{$(eltype(f.ev))}($(length(f.ev)))")
  f.arev===○ ? nothing : println(io, separatorFont,".arev ", defaultFont, "vector{$(eltype(f.arev))}($(length(f.arev)))")
  println(io, greyFont,"LEGEND:")
  println(io, greyFont,"F=filter, i=inverse, ev=all eigenvalues")
  println(io, greyFont,"D=Diagonal eigenvalues associated with F")
  println(io, greyFont,"ar=accumulated regularized ev")
end

##

println("\n⭐ "," Welcome to the ", 📌, " package", " ⭐\n")
@info " "
println(" Your Machine `",gethostname(),"` (",Sys.MACHINE, ")")
println(" runs on kernel ",Sys.KERNEL," with word size ",Sys.WORD_SIZE,".")
println(" CPU  Threads: ", Sys.CPU_THREADS)
println(" Base.Threads: ", "$(Threads.nthreads())")
println(" BLAS Threads: ", "$(Sys.CPU_THREADS)", "\n")

end # module


text:
function TimeStampedAngularCoordinates(arg0::AbsoluteDate, arg1::FieldRotation)
    return TimeStampedAngularCoordinates((AbsoluteDate, FieldRotation), arg0, arg1)
end

function TimeStampedAngularCoordinates(arg0::AbsoluteDate, arg1::PVCoordinates, arg2::PVCoordinates)
    return TimeStampedAngularCoordinates((AbsoluteDate, PVCoordinates, PVCoordinates), arg0, arg1, arg2)
end

function TimeStampedAngularCoordinates(arg0::AbsoluteDate, arg1::PVCoordinates, arg2::PVCoordinates, arg3::PVCoordinates, arg4::PVCoordinates, arg5::jdouble)
    return TimeStampedAngularCoordinates((AbsoluteDate, PVCoordinates, PVCoordinates, PVCoordinates, PVCoordinates, jdouble), arg0, arg1, arg2, arg3, arg4, arg5)
end

function TimeStampedAngularCoordinates(arg0::AbsoluteDate, arg1::Rotation, arg2::Vector3D, arg3::Vector3D)
    return TimeStampedAngularCoordinates((AbsoluteDate, Rotation, Vector3D, Vector3D), arg0, arg1, arg2, arg3)
end

function add_offset(obj::TimeStampedAngularCoordinates, arg0::AngularCoordinates)
    return jcall(obj, "addOffset", TimeStampedAngularCoordinates, (AngularCoordinates,), arg0)
end

function get_date(obj::TimeStampedAngularCoordinates)
    return jcall(obj, "getDate", AbsoluteDate, ())
end

function interpolate(::Type{TimeStampedAngularCoordinates}, arg0::AbsoluteDate, arg1::AngularDerivativesFilter, arg2::Collection)
    return jcall(TimeStampedAngularCoordinates, "interpolate", TimeStampedAngularCoordinates, (AbsoluteDate, AngularDerivativesFilter, Collection), arg0, arg1, arg2)
end

function revert(obj::TimeStampedAngularCoordinates)
    return jcall(obj, "revert", TimeStampedAngularCoordinates, ())
end

function shifted_by(obj::TimeStampedAngularCoordinates, arg0::jdouble)
    return jcall(obj, "shiftedBy", TimeStampedAngularCoordinates, (jdouble,), arg0)
end

function subtract_offset(obj::TimeStampedAngularCoordinates, arg0::AngularCoordinates)
    return jcall(obj, "subtractOffset", TimeStampedAngularCoordinates, (AngularCoordinates,), arg0)
end

ext: jl


text:
Base.one(T::Type{Valid{N,ES}}) where {N,ES} = Valid(one(Vnum{N,ES}), one(Vnum{N,ES}))
Base.zero(T::Type{Valid{N,ES}}) where {N,ES} = Valid(zero(Vnum{N,ES}), zero(Vnum{N,ES}))

abstract type ∅; end
abstract type ℝ; end
abstract type ℝp; end

convert(T::Type{Valid{N, ES}}, ::Type{∅}) where {N,ES}  = Valid(zero(     Vnum{N,ES}), maxneg( Vnum{N,ES}))
convert(T::Type{Valid{N, ES}}, ::Type{ℝ}) where {N,ES}  = Valid(-maxpos( Vnum{N,ES}), maxpos(      Vnum{N,ES}))
convert(T::Type{Valid{N, ES}}, ::Type{ℝp}) where {N,ES} = Valid(Vnum{N,ES}(Inf),       maxpos(      Vnum{N,ES}))

abstract type __plusstar; end
Base.:+(::typeof(Base.:*)) = __plusstar
abstract type __minusstar; end
Base.:-(::typeof(Base.:*)) = __minusstar
abstract type __positives; end
abstract type __negatives; end

ℝ(::Type{__plusstar}) =      __plusstar
ℝ(::Type{__minusstar}) =     __minusstar
ℝ(::typeof(+)) =             __positives
ℝ(::typeof(-)) =             __negatives

convert(T::Type{Valid{N,ES}}, ::Type{__plusstar}) where {N,ES} =  Valid(minpos(Vnum{N, ES}),  maxpos(Vnum{N, ES}))
convert(T::Type{Valid{N,ES}}, ::Type{__minusstar}) where {N,ES} = Valid(-maxpos(Vnum{N, ES}), maxneg(Vnum{N, ES}))
convert(T::Type{Valid{N,ES}}, ::Type{__positives}) where {N,ES} = Valid(zero(Vnum{N, ES}),          maxpos(Vnum{N, ES}))
convert(T::Type{Valid{N,ES}}, ::Type{__negatives}) where {N,ES} = Valid(-maxpos(Vnum{N, ES}),         zero(Vnum{N, ES}))

export ∅, ℝ, ℝp

max_issues_repo_name: MohHizzani/SigmoidNumbers

max_forks_repo_name: MohHizzani/SigmoidNumbers


text:
export hmm_loglik

import ..HMMModel

using mlpack._Internal.io

import mlpack_jll
const hmm_loglikLibrary = mlpack_jll.libmlpack_julia_hmm_loglik

# Call the C binding of the mlpack hmm_loglik binding.
function hmm_loglik_mlpackMain()
  success = ccall((:hmm_loglik, hmm_loglikLibrary), Bool, ())
  if !success
    # Throw an exception---false means there was a C++ exception.
    throw(ErrorException("mlpack binding error; see output"))
  end
end

" Internal module to hold utility functions. "
module hmm_loglik_internal
  import ..hmm_loglikLibrary

import ...HMMModel

# Get the value of a model pointer parameter of type HMMModel.
function IOGetParamHMMModel(paramName::String)::HMMModel
  HMMModel(ccall((:IO_GetParamHMMModelPtr, hmm_loglikLibrary), Ptr{Nothing}, (Cstring,), paramName))
end

# Set the value of a model pointer parameter of type HMMModel.
function IOSetParamHMMModel(paramName::String, model::HMMModel)
  ccall((:IO_SetParamHMMModelPtr, hmm_loglikLibrary), Nothing, (Cstring, Ptr{Nothing}), paramName, model.ptr)
end

# Serialize a model to the given stream.
function serializeHMMModel(stream::IO, model::HMMModel)
  buf_len = UInt[0]
  buf_ptr = ccall((:SerializeHMMModelPtr, hmm_loglikLibrary), Ptr{UInt8}, (Ptr{Nothing}, Ptr{UInt}), model.ptr, Base.pointer(buf_len))
  buf = Base.unsafe_wrap(Vector{UInt8}, buf_ptr, buf_len[1]; own=true)
  write(stream, buf)
end
# Deserialize a model from the given stream.
function deserializeHMMModel(stream::IO)::HMMModel
  buffer = read(stream)
  HMMModel(ccall((:DeserializeHMMModelPtr, hmm_loglikLibrary), Ptr{Nothing}, (Ptr{UInt8}, UInt), Base.pointer(buffer), length(buffer)))
end
end # module

"""
    hmm_loglik(input, input_model; [verbose])

This utility takes an already-trained HMM, specified with the `input_model`
parameter, and evaluates the log-likelihood of a sequence of observations, given
with the `input` parameter.  The computed log-likelihood is given as output.

For example, to compute the log-likelihood of the sequence `seq` with the
pre-trained HMM `hmm`, the following command may be used:

```julia
julia> using CSV
julia> seq = CSV.read("seq.csv")
julia> _ = hmm_loglik(seq, hmm)
```

# Arguments

 - `input::Array{Float64, 2}`: File containing observations,
 - `input_model::HMMModel`: File containing HMM.
 - `verbose::Bool`: Display informational messages and the full list of
      parameters and timers at the end of execution.  Default value `false`.

# Return values

 - `log_likelihood::Float64`: Log-likelihood of the sequence.  Default
      value `0`.

"""
function hmm_loglik(input,
                    input_model::HMMModel;
                    verbose::Union{Bool, Missing} = missing,
                    points_are_rows::Bool = true)
  # Force the symbols to load.
  ccall((:loadSymbols, hmm_loglikLibrary), Nothing, ());

  IORestoreSettings("Hidden Markov Model (HMM) Sequence Log-Likelihood")

  # Process each input argument before calling mlpackMain().
  IOSetParamMat("input", input, points_are_rows)
  hmm_loglik_internal.IOSetParamHMMModel("input_model", convert(HMMModel, input_model))
  if verbose !== nothing && verbose === true
    IOEnableVerbose()
  else
    IODisableVerbose()
  end

  IOSetPassed("log_likelihood")
  # Call the program.
  hmm_loglik_mlpackMain()

  return IOGetParamDouble("log_likelihood")
end


text:
# # High angular resolution diffusion imaging
#
# In this example we will compute the signal using the same gradient sequence in many
# different directions.

# ## Building a biological model
#
# We start by loading SpinDoctor and a Makie plotting backend.

# LSP indexing solution                                                          #src
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983 #src
if isdefined(@__MODULE__, :LanguageServer)                                       #src
    include("../src/SpinDoctor.jl")                                              #src
    using .SpinDoctor                                                            #src
end                                                                              #src

using SpinDoctor
using LinearAlgebra

if haskey(ENV, "GITHUB_ACTIONS")
    using CairoMakie
else
    using GLMakie
end

# Here we create a recipe for five stacked plates with isotropic diffusion tensors. They
# should allow for free diffusion in the horizontal direction, but a rather restricted
# vertical diffusion with the permeable membranes.

ncell = 5
setup = PlateSetup(;
    name = "Plates",
    width = 50.0,
    depth = 50.0,
    heights = fill(5.0, ncell),
    bend = 0.0,
    twist = 0.0,
)
coeffs = coefficients(
    setup;
    D = [0.002 * I(3) for _ = 1:ncell],
    T₂ = fill(Inf, ncell),
    ρ = fill(1.0, ncell),
    κ = (; interfaces = fill(1e-4, ncell - 1), boundaries = fill(0.0, ncell)),
    γ = 2.67513e-4,
)

# We then proceed to build the geometry and finite element mesh.

mesh, = create_geometry(setup; recreate = true)
plot_mesh(mesh)

# The mesh looks good, so we may then proceed to assemble the biological model and the
# associated finite element matrices.

model = Model(; mesh, coeffs...)
matrices = assemble_matrices(model);

# We may also compute some useful quantities, including a scalar diffusion coefficient from
# the diffusion tensors.

volumes = get_cmpt_volumes(model.mesh)
D_avg = 1 / 3 * tr.(model.D)' * volumes / sum(volumes)
ncompartment = length(model.mesh.points)

# The gradient pulse sequence will be a PGSE with both vertical and horizontal components.
# This allows for both restricted vertical diffusion and almost unrestricted horizontal
# diffusion. The different approaches should hopefully confirm this behaviour.

directions = unitsphere(200)
profile = PGSE(2500.0, 4000.0)
b = 1000
g = √(b / int_F²(profile)) / model.γ
gradients = [ScalarGradient(d, profile, g) for d ∈ eachcol(directions)]

# We may solve the BTPDE for each gradient.

btpde = BTPDE(; model, matrices)
solver = IntervalConstantSolver(; timestep = 10.0)
ξ, = solve_multigrad(btpde, gradients, solver)

# The signals are computed from the magnetization field through quadrature.

signal = [abs(compute_signal(matrices.M, ξ)) for ξ ∈ ξ]

# We may plot the directionalized signal.

plot_hardi(directions, signal)

# The signal attenuates the most in the horizontal direction, as that is where diffusion is
# restricted the least.


text:
_pgf_array(a) = "{$(join(a, ','))}"

const COLORS = ["cyan", "red", "orange", "lime", "magenta", "lightgray", "darkgray"]

function pgf_boxplot(lqs, hqs, mds, mns, mxs, labels; h="10cm", w="6cm")
	s = """
	\\begin{tikzpicture}
	\\begin{axis}[
		height=$h,
		width=$w,
		ymin=0, ymax=$(length(labels)+1),
		xmin=0, xmax=$(length(labels)+1),
		ytick=$(_pgf_array(1:length(labels))),
		yticklabels=$(_pgf_array(labels))
		]
	"""

	for (lq, hq, md, mn, mx) in zip(lqs, hqs, mds, mns, mxs)
		s *= """
		\\addplot [boxplot prepared={
		lower whisker=$mn, lower quartile=$lq,
		median=$md, upper quartile=$hq,
		upper whisker=$mx},
		] coordinates {};
		"""
	end
	s *= """
	\\end{axis}
	\\end{tikzpicture}
	"""
	s
end

function pgf_boxplot_grouped(lqs, hqs, mds, mns, mxs, labels, groups; h="10cm", w="6cm")
	s = """
	\\begin{tikzpicture}
	\\begin{axis}[
		height=$h,
		width=$w,
		ymin=0, ymax=$(length(labels)+1),
		xmin=0, xmax=$(length(labels)+1),
		ytick=$(_pgf_array(1:length(labels))),
		yticklabels=$(_pgf_array(labels))
		]
	"""

	for (lq, hq, md, mn, mx, g) in zip(lqs, hqs, mds, mns, mxs, groups)
		s *= """
		\\addplot [$(g%2==0 ? "fill=lightgray" : "fill=gray"), boxplot prepared={
		lower whisker=$mn, lower quartile=$lq,
		median=$md, upper quartile=$hq,
		upper whisker=$mx},
		] coordinates {};
		"""
	end
	s *= """
	\\end{axis}
	\\end{tikzpicture}
	"""
	s
end

function pgf_boxplot_grouped_color(lqs, hqs, mds, mns, mxs, labels, groups; h="10cm", w="6cm")
	s = """
	\\begin{tikzpicture}
	\\begin{axis}[
		height=$h,
		width=$w,
		ymin=0, ymax=$(length(labels)+1),
		xmin=0, xmax=$(length(labels)+1),
		ytick=$(_pgf_array(1:length(labels))),
		yticklabels=$(_pgf_array(labels))
		]
	"""

	for (lq, hq, md, mn, mx, g) in zip(lqs, hqs, mds, mns, mxs, groups)
		s *= """
		\\addplot [fill=$(COLORS[g]), boxplot prepared={
		lower whisker=$mn, lower quartile=$lq,
		median=$md, upper quartile=$hq,
		upper whisker=$mx},
		] coordinates {};
		"""
	end
	s *= """
	\\end{axis}
	\\end{tikzpicture}
	"""
	s
end

function pgf_boxplot_grouped_colorpos(lqs, hqs, mds, mns, mxs, labels, groups, mgroups; h="10cm", w="6cm")
	yticks = [1]
	for (cg, ng) in zip(mgroups[1:end-1], mgroups[2:end-1])
		if ng == cg # if next item belongs to a group with the same id move by one otherwise skip row
			push!(yticks, yticks[end] + 1)
		else
			push!(yticks, yticks[end] + 2)
		end
	end
	push!(yticks, yticks[end] + 1)

	s = """
	\\begin{tikzpicture}
	\\begin{axis}[
		height=$h,
		width=$w,
		ymin=0, ymax=$(yticks[end] + 1),
		xmin=0, xmax=$(length(labels) + 1),
		ytick=$(_pgf_array(yticks)),
		yticklabels=$(_pgf_array(labels))
		]
	"""

	for (lq, hq, md, mn, mx, g, y) in zip(lqs, hqs, mds, mns, mxs, groups, yticks)
		s *= """
		\\addplot [fill=$(COLORS[g]), boxplot prepared={
		draw position=$(y),
		lower whisker=$mn, lower quartile=$lq,
		median=$md, upper quartile=$hq,
		upper whisker=$mx},
		] coordinates {};
		"""
	end
	s *= """
	\\end{axis}
	\\end{tikzpicture}
	"""
	s
end


text:
module UnitfulRecipes

using RecipesBase
using Unitful: Quantity, unit, ustrip, Unitful, dimension
export @P_str

#==========
Main recipe
==========#

@recipe function f(::Type{T}, x::T) where T <: AbstractArray{<:Quantity}
    axisletter = plotattributes[:letter]   # x, y, or z
    fixaxis!(plotattributes, x, axisletter)
end

function fixaxis!(attr, x, axisletter)
    # Attribute keys
    axislabel = Symbol(axisletter, :guide) # xguide, yguide, zguide
    axislims = Symbol(axisletter, :lims)   # xlims, ylims, zlims
    axisunit = Symbol(axisletter, :unit)   # xunit, yunit, zunit
    axis = Symbol(axisletter, :axis)       # xaxis, yaxis, zaxis
    # Get the unit
    u = pop!(attr, axisunit, unit(eltype(x)))
    # If the subplot already exists, get unit from its axis label
    sp = get(attr, :subplot, 1)
    if sp ≤ length(attr[:plot_object])
        label = attr[:plot_object][sp][axis][:guide]
        if label isa UnitfulString
            u = label.unit
        end
    end
    # Fix the attributes: labels, lims, marker/line stuff, etc.
    append_unit_if_needed!(attr, axislabel, u)
    fixlims!(attr, axislims, u)
    fixmarkercolor!(attr)
    fixmarkersize!(attr)
    fixlinecolor!(attr)
    # Strip the unit
    ustrip.(u, x)
end

# Recipe for (x::AVec, y::AVec, z::Surface) types
const AVec = AbstractVector
const AMat{T} = AbstractArray{T,2} where T
@recipe function f(x::AVec, y::AVec, z::AMat{T}) where T <: Quantity
    u = get(plotattributes, :zunit, unit(eltype(z)))
    z = fixaxis!(plotattributes, z, :z)
    append_unit_if_needed!(plotattributes, :colorbar_title, u)
    x, y, z
end

# Recipe for vectors of vectors
@recipe function f(::Type{T}, x::T) where T <: AbstractVector{<:AbstractVector{<:Quantity}}
    axisletter = plotattributes[:letter]   # x, y, or z
    [fixaxis!(plotattributes, x, axisletter) for x in x]
end

# Recipes for functions
@recipe function f(f::Function, x::T) where T <: AVec{<:Quantity}
    x, f.(x)
end
@recipe function f(x::T, f::Function) where T <: AVec{<:Quantity}
    x, f.(x)
end
@recipe function f(x::T, y::AVec, f::Function) where T <: AVec{<:Quantity}
    x, y, f.(x',y)
end
@recipe function f(x::AVec, y::T, f::Function) where T <: AVec{<:Quantity}
    x, y, f.(x',y)
end
@recipe function f(x::T1, y::T2, f::Function) where {T1<:AVec{<:Quantity}, T2<:AVec{<:Quantity}}
    x, y, f.(x',y)
end

#===============
Attribute fixing
===============#

# Markers / lines
function fixmarkercolor!(attr)
    u = ustripattribute!(attr, :marker_z)
    fixlims!(attr, :clims, u)
    u == Unitful.NoUnits || append_unit_if_needed!(attr, :colorbar_title, u)
end
fixmarkersize!(attr) = ustripattribute!(attr, :markersize)
fixlinecolor!(attr) = ustripattribute!(attr, :line_z)

# Lims
function fixlims!(attr, key, u)
    if haskey(attr, key)
        lims = attr[key]
        if lims isa NTuple{2, Quantity}
            attr[key] = ustrip.(u, lims)
        end
    end
end

# strip unit from attribute[key]
function ustripattribute!(attr, key)
    if haskey(attr, key)
        v = attr[key]
        u = unit(eltype(v))
        attr[key] = ustrip.(u, v)
        return u
    else
        return Unitful.NoUnits
    end
end

#=======================================
Label string containing unit information
=======================================#

abstract type AbstractProtectedString <: AbstractString end
struct ProtectedString{S} <: AbstractProtectedString
    content::S
end
struct UnitfulString{S,U} <: AbstractProtectedString
    content::S
    unit::U
end
# Minimum required AbstractString interface to work with Plots
const S = AbstractProtectedString
Base.iterate(n::S) = iterate(n.content)
Base.iterate(n::S, i::Integer) = iterate(n.content, i)
Base.codeunit(n::S) = codeunit(n.content)
Base.ncodeunits(n::S) = ncodeunits(n.content)
Base.isvalid(n::S, i::Integer) = isvalid(n.content, i)
Base.pointer(n::S) = pointer(n.content)
Base.pointer(n::S, i::Integer) = pointer(n.content, i)
"""
    P_str(s)

Creates a string that will be Protected from recipe passes.

Example:
```julia
julia> plot(rand(10)*u"m", xlabel=P"This label is protected")

julia> plot(rand(10)*u"m", xlabel=P"This label is not")
```
"""
macro P_str(s)
    return ProtectedString(s)
end

#=====================================
Append unit to labels when appropriate
=====================================#

function append_unit_if_needed!(attr, key, u::Unitful.Units)
    label = get(attr, key, nothing)
    append_unit_if_needed!(attr, key, label, u)
end
# dispatch on the type of `label`
append_unit_if_needed!(attr, key, label::ProtectedString, u) = nothing
append_unit_if_needed!(attr, key, label::UnitfulString, u) = nothing
function append_unit_if_needed!(attr, key, label::Nothing, u)
    attr[key] = UnitfulString(string(u), u)
end
function append_unit_if_needed!(attr, key, label::S, u) where {S <: AbstractString}
    if label ≠ ""
        attr[key] = UnitfulString(S(string(label, " (", u, ")")), u)
    end
end

end # module


text:
"""
    DetMDP(; ns = 10^4, na = 10)

Returns a random deterministic MDP.
"""
function DetMDP(; ns = 10^4, na = 10)
    mdp = MDP(ns, na, init = "deterministic")
    mdp.reward = mdp.reward .* (mdp.reward .< -1.5)
    mdp
end

"""
    DetTreeMDP(; na = 4, depth = 5)

Returns a treeMDP with random rewards at the leaf nodes.
"""
function DetTreeMDP(; na = 4, depth = 5)
    mdp = treeMDP(na, depth, init = "deterministic")
end

"""
    DetTreeMDPwithinrew(; na = 4, depth = 5)

Returns a treeMDP with random rewards.
"""
function DetTreeMDPwithinrew(; args...)
    mdp = DetTreeMDP(; args...)
    nonterminals = findall(x -> x == 0, mdp.isterminal)
    mdp.reward[:, nonterminals] = -rand(ENV_RNG, mdp.na, length(nonterminals))
    mdp
end

"""
    StochTreeMDP(; na = 4, depth = 4, bf = 2)

Returns a random stochastic treeMDP with branching factor `bf`.
"""
function StochTreeMDP(; na = 4, depth = 4, bf = 2)
    mdp = treeMDP(na, depth, init = "random", branchingfactor = bf)
    mdp
end

"""
    StochMDP(; na = 10, ns = 50) = MDP(ns, na)
"""
StochMDP(; na = 10, ns = 50) = MDP(ns, na)

"""
    AbsorbingDetMDP(;ns = 10^3, na = 10)

Returns a random deterministic absorbing MDP
"""
function AbsorbingDetMDP(;ns = 10^3, na = 10)
    mdp = MDP(ns, na, init = "deterministic")
    mdp.reward .= mdp.reward .* (mdp.reward .< -.5)
    mdp.initialstates = 1:div(ns, 100)
    reset!(mdp)
    setterminalstates!(mdp, ns - div(ns, 100) + 1:ns)
    mdp
end


text:
module DescriptiveWrapper

using JavaCall

const AbstractStorelessUnivariateStatistic = @jimport org.hipparchus.stat.descriptive.AbstractStorelessUnivariateStatistic
const AbstractUnivariateStatistic = @jimport org.hipparchus.stat.descriptive.AbstractUnivariateStatistic
const AggregatableStatistic = @jimport org.hipparchus.stat.descriptive.AggregatableStatistic
const Class = @jimport java.lang.Class
const DescriptiveStatistics = @jimport org.hipparchus.stat.descriptive.DescriptiveStatistics
const DoubleConsumer = @jimport java.util.function.DoubleConsumer
const Iterable = @jimport java.lang.Iterable
const JString = @jimport java.lang.String
const MultivariateSummaryStatistics = @jimport org.hipparchus.stat.descriptive.MultivariateSummaryStatistics
const Object = @jimport java.lang.Object
const RealMatrix = @jimport org.hipparchus.linear.RealMatrix
const StatisticalMultivariateSummary = @jimport org.hipparchus.stat.descriptive.StatisticalMultivariateSummary
const StatisticalSummary = @jimport org.hipparchus.stat.descriptive.StatisticalSummary
const StatisticalSummaryValues = @jimport org.hipparchus.stat.descriptive.StatisticalSummaryValues
const StorelessMultivariateStatistic = @jimport org.hipparchus.stat.descriptive.StorelessMultivariateStatistic
const StorelessUnivariateStatistic = @jimport org.hipparchus.stat.descriptive.StorelessUnivariateStatistic
const StreamingStatistics = @jimport org.hipparchus.stat.descriptive.StreamingStatistics
const StreamingStatistics_StreamingStatisticsBuilder = @jimport org.hipparchus.stat.descriptive.StreamingStatistics$StreamingStatisticsBuilder
const UnivariateStatistic = @jimport org.hipparchus.stat.descriptive.UnivariateStatistic
const WeightedEvaluation = @jimport org.hipparchus.stat.descriptive.WeightedEvaluation

export AbstractStorelessUnivariateStatistic
export AbstractUnivariateStatistic
export AggregatableStatistic
export DescriptiveStatistics
export MultivariateSummaryStatistics
export StatisticalMultivariateSummary
export StatisticalSummary
export StatisticalSummaryValues
export StorelessMultivariateStatistic
export StorelessUnivariateStatistic
export StreamingStatistics
export UnivariateStatistic
export WeightedEvaluation
export accept
export add_value
export aggregate
export and_then
export apply
export builder
export clear
export copy
export equals
export evaluate
export get_class
export get_covariance
export get_data
export get_dimension
export get_element
export get_geometric_mean
export get_kurtosis
export get_max
export get_mean
export get_median
export get_min
export get_n
export get_percentile
export get_population_variance
export get_quadratic_mean
export get_result
export get_second_moment
export get_skewness
export get_sorted_values
export get_standard_deviation
export get_sum
export get_sum_log
export get_sum_of_logs
export get_sum_of_squares
export get_sum_sq
export get_summary
export get_values
export get_variance
export get_window_size
export hash_code
export increment
export increment_all
export notify
export notify_all
export remove_most_recent_value
export replace_most_recent_value
export set_data
export set_window_size
export to_string
export wait

include("MomentWrapper/MomentWrapper.jl")
include("RankWrapper/RankWrapper.jl")
include("SummaryWrapper/SummaryWrapper.jl")
include("VectorWrapper/VectorWrapper.jl")
include("abstract_storeless_univariate_statistic.jl")
include("abstract_univariate_statistic.jl")
include("aggregatable_statistic.jl")
include("descriptive_statistics.jl")
include("multivariate_summary_statistics.jl")
include("statistical_multivariate_summary.jl")
include("statistical_summary.jl")
include("statistical_summary_values.jl")
include("storeless_multivariate_statistic.jl")
include("storeless_univariate_statistic.jl")
include("streaming_statistics.jl")
include("univariate_statistic.jl")
include("weighted_evaluation.jl")

end


text:
using Test

@timed_testset "batch_processing_helpers/" begin
    include("batch_processing_helpers/unit.jl")
    include("batch_processing_helpers/integration.jl")
end


text:
using Compat
using Dates
using Infinity
using Infinity.Utils
using Random
using Test
using TimeZones: ZonedDateTime

@testset "Infinity" begin
    include("utils.jl")
    include("infinite.jl")
    include("infextendedreal.jl")
    include("infextendedtime.jl")
end


text:
using Geomechanics
using Test

@testset "Geomechanics.jl" begin
    # Write your own tests here.
end


text:
using Oceananigans.Utils: prettytime

struct ForwardMap{V,D,P}
    model_time_series::V
    observations::D
    parameters::P
end

mutable struct ModelTimeSeries{U, V, B, E}
    u :: U
    v :: V
    b :: B
    e :: E
end

ModelTimeSeries(grid, targets) = ModelTimeSeries([XFaceField(grid) for i = targets],
                                                 [YFaceField(grid) for i = targets],
                                                 [CenterField(grid) for i = targets],
                                                 [CenterField(grid) for i = targets])

function model_time_series(simulation, observations, parameters)

    model = simulation.model

    # Sometimes observations will have fewer data objects than the model, so pad observations with redundant objects
    redundant_data = [observations[1] for _ in 1:(batch_size(model) - length(observations))]
    full_observations = [observations; redundant_data]

    all_targets = getproperty.(full_observations, :targets)
    max_simulation_length = maximum(length.(all_targets))
    starts = getindex.(all_targets, 1)

    initialize_forward_run!(model, full_observations, parameters, starts)

    outputs = [ModelTimeSeries(data.grid, data.targets) for data in observations]

    u = get_model_field(model, :u)
    v = get_model_field(model, :v)
    b = get_model_field(model, :b)
    e = get_model_field(model, :e)

    # this should be improved
    all_lengths = length.(getproperty.(observations, :t))
    longest_sim = observations[argmax(all_lengths)]
    t = longest_sim.t # times starting from zero

    start_time = time_ns()

    for save_index in 1:max_simulation_length
        simulation.stop_time = t[save_index]
        run!(simulation)

        u = model.velocities.u
        v = model.velocities.v
        b = model.tracers.b
        e = model.tracers.e

        capture_model_state!(outputs, save_index, observations, u, v, b, e)
    end

    end_time = time_ns()
    elapsed_time = (end_time - start_time) * 1e-9

    @info "The forward run took $(prettytime(elapsed_time))"

    return ForwardMap(outputs, observations, parameters)
end

function capture_model_state!(outputs, save_index, observations, u, v, b, e)
    for (data_index, data) in enumerate(observations)
        output = outputs[data_index]
        if save_index <= length(data.targets)
            u_snapshot = parent(output.u[save_index])
            v_snapshot = parent(output.v[save_index])
            b_snapshot = parent(output.b[save_index])
            e_snapshot = parent(output.e[save_index])

            copyto!(u_snapshot, view(parent(u), 1, data_index, :))
            copyto!(v_snapshot, view(parent(v), 1, data_index, :))
            copyto!(b_snapshot, view(parent(b), 1, data_index, :))
            copyto!(e_snapshot, view(parent(e), 1, data_index, :))
        end
    end

    return nothing
end


text:
rewardmatrix(p::DDP{nStateVars,nChoiceVars,C}) where
    {nStateVars,nChoiceVars,C<:Int} =
    _rewardmatrix(p.rewardfunc, p.tStateVectors, getchoicevars(p.tStateVectors, p.tChoiceVectors))

rewardmatrix(p::DDP{nStateVars,nChoiceVars,C}) where
    {nStateVars,nChoiceVars,C<:Vector{Float64}} =
    _rewardmatrix(p.rewardfunc, p.tStateVectors, p.tChoiceVectors)

""" Return a nChoices x nStates reward matrix"""
function _rewardmatrix(rewardfunc, tStateVectors, tChoiceVectors)
    # output dimension is nChoices x nStates

    nChoices = prod(length.(tChoiceVectors))
    nStates = prod(length.(tStateVectors))
    mReward = zeros(Float64, nChoices, nStates)

    iterator_states = Iterators.product(tStateVectors...)

    iterator_choices = getiterator(tChoiceVectors)

    for (i, states) in enumerate(iterator_states)
        for (j, choices) in enumerate(iterator_choices)
           # @show choices
           mReward[j,i] = rewardfunc(states, choices)
           # error("stop")
        end
    end

    return mReward
end

rewardmatrix_partial(p::DDP) = _rewardmatrix_partial(p.options.rewardfunc_partial,
    p.tStateVectors)
# rewardmatrix_partial(p::DDP) = _rewardmatrix_partial(p.options.rewardfunc_partial,
#     getchoicevars(p.tStateVectors, p.tChoiceVectors),
#     getnonchoicevars(p.tStateVectors, p.tChoiceVectors))

function _rewardmatrix_partial(rewardfunc_partial, tStateVectors)
    # output dimension is vector of length nStates

    # julia does not know dim of endog state vars from this expression
    nStates::Int64 = prod(length.(tStateVectors))

    mReward = zeros(Float64, nStates)

    iterator = Iterators.product(tStateVectors...)

    for (i, states) in enumerate(iterator)
       mReward[i] = rewardfunc_partial(states)
    end

    return mReward

end

"""Return a nEndogStates x nExogStates reward matrix."""
# function _rewardmatrix_partial(rewardfunc_partial,
#     tEndogStateVectors, tExogStateVectors)
#     # output dimension is nEndogStates x nExogStates
#
#     # julia does not know dim of endog state vars from this expression
#     nEndogStates::Int64 = prod(length.(tEndogStateVectors))
#     nExogStates::Int64 = prod(length.(tExogStateVectors))
#
#     # mStates::Array{Float64,2} = gridmake(tStateVectors...)
#
#     mReward = zeros(Float64, nEndogStates, nExogStates)
#
#     iterator_endogstates = Iterators.product(tEndogStateVectors...)
#     iterator_exogstates = Iterators.product(tExogStateVectors...)
#
#     for (i, exogstates) in enumerate(iterator_exogstates)
#         for (j, endogstates) in enumerate(iterator_endogstates)
#            mReward[j,i] = rewardfunc_partial((endogstates..., exogstates...))
#         end
#     end
#
#     return mReward
#
# end


text:
export mat,
       liou,
       unitalproj,
       choi_liou_involution,
       swap_involution,
       choi2liou,
       choi2kraus,
       choi2stinespring,
       kraus2choi,
       kraus2liou,
       kraus2stinespring,
       liou2choi,
       liou2kraus,
       liou2stinespring,
       liou2pauliliou,
       pauliliou2liou,
       dissipator,
       hamiltonian,
       depol,
       istp,
       iscp,
       ischannel,
       isunital,
       isliouvillian,
       nearestu

mat( v::Vector, r=round(Int,sqrt(length(v))), c=round(Int,sqrt(length(v))) ) = reshape( v, r, c )
liou{T<:AbstractMatrix}( left::T, right::T ) = kron( transpose(right), left )
liou{T<:AbstractMatrix}( m::T ) = kron( conj(m), m )

function choi_liou_involution( r::Matrix )
  d = round(Int, sqrt(size(r,1)) )
  rl = reshape( r, (d, d, d, d) )
  rl = permutedims( rl, [1,3,2,4] )
  reshape( rl, size(r) )
end

function swap_involution( r::Matrix )
  d = round(Int, sqrt(size(r,1)) )
  rl = reshape( r, (d, d, d, d) )
  rl = permutedims( rl, [3,4,1,2] )
  reshape( rl, size(r) )
end

function choi2liou( r::Matrix  )
  sqrt(size(r,1))*choi_liou_involution( r )
end

function liou2choi( r::Matrix )
  choi_liou_involution( r )/sqrt(size(r,1))
end

function choi2kraus{T}( r::Matrix{T}  )
  (vals,vecs) = eig( sqrt(size(r,1))*r )
  #vals = eigvals( sqrt(size(r,1))*r )
  kraus_ops = Matrix{T}[]
  for i in 1:length(vals)
    push!(kraus_ops, sqrt(vals[i])*mat(vecs[:,i]))
  end
  kraus_ops
end

function choi2stinespring{T}( r::Matrix{T}  )
  (vals,vecs) = eig( Hermitian(sqrt(size(r,1))*r) ) # we are assuming Hermiticity-preserving maps
  #vals = eigvals( sqrt(size(r,1))*r )
  A_ops = Matrix{T}[]
  B_ops = Matrix{T}[]
  d = length(vals)
  for i in 1:length(vals)
    push!(A_ops, kron(sqrt(abs(vals[i]))*mat(vecs[:,i]),ket(i-1,d)))
    push!(B_ops, sign(vals[i])*A_ops[end])
  end
  return sum(A_ops),sum(B_ops)
end

function liou2stinespring{T}( r::Matrix{T} )
  return r |> liou2choi |> choi2stinespring
end

function kraus2liou{T}( k::Vector{Matrix{T}} )
  l = zeros(T,map(x->x^2,size(k[1])))
  for i in 1:length(k)
    l = l + liou(k[i],k[i]')
  end
  l
end

function liou2kraus( l::Matrix )
  choi2kraus( liou2choi( l ) )
end

function kraus2choi{T}( k::Vector{Matrix{T}} )
  c = zeros(T,map(x->x^2,size(k[1])))
  for i in 1:length(k)
    c = vec(k[i])*vec(k[i])'
  end
  c/sqrt(size(c,1))
end

# TODO: Add support for sparse matrices
function dissipator( a::Matrix )
  liou(a,a') - 1/2 * liou(a'*a, eye(a)) - 1/2 * liou(eye(a),a'*a)
end

function hamiltonian( h::Matrix )
  -1im * ( liou(h,eye(h)) - liou(eye(h),h) )
end

function pauliliou2liou( m::Matrix )
  if size(m,1) != size(m,2)
    error("Only square matrices supported")
  elseif size(m,1) != 4^(floor(log2(size(m,1))/2))
    error("Only matrices with dimension 4^n supported.")
  end
  dsq = size(m,1)
  res = zeros(Complex128,size(m))
  n = round(Int,log(2,dsq)/2)
  for (i,pi) in enumerate(allpaulis(n))
    for (j,pj) in enumerate(allpaulis(n))
      res += m[i,j] * vec(complex(pi)) * vec(complex(pj))' / sqrt(dsq)
    end
  end
  res
end

function liou2pauliliou{T}( m::Matrix{T} )
  if size(m,1) != size(m,2)
    error("Only square matrices supported")
  elseif size(m,1) != 4^(floor(log2(size(m,1))/2))
    error("Only matrices with dimension 4^n supported.")
  end
  dsq = size(m,1)
  res = zeros(Complex128,size(m))
  n = round(Int,log(2,dsq)/2)
  for (i,pi) in enumerate(allpaulis(n))
    for (j,pj) in enumerate(allpaulis(n))
      res[i,j] += trace( m * vec(complex(pi)) * vec(complex(pj))' / sqrt(dsq) )
    end
  end
  res
end

"""
`depol(d, p=1.0)`

`depol` return a superoperator that replaces the input with a maximally
mixed state with probability p, and leaves it unchanged with probability (1-p).
"""
function depol( d::Int, p=1.0 )
  choi2liou( p * eye(d^2)/d^2 + (1-p) * projector(_max_entangled_state(d)) )
end

"""
`unitalproj(m)`

Given a superoperator `j` in a Liouville representation, `unitalproj`
extracts the closest superoperator (in Frobenius norm) that is
unital. The result may not be completely positive.
"""
function unitalproj{T}( m::Matrix{T} )
  d2 = size(m,1)
  d  = round(Int,sqrt(d2))
  id = projector(normalize(vec(eye(d))))
  id*m*id + (I-id)*m*(I-id)
end

"""
ishermitian(m; tol)

Checks if a matrix is Hermitian.
"""
function ishermitian(m; tol=0.0)
    ah = (m-m')/2
    tol = tol==0.0 ? 1e2*eps(abs(one(eltype(m)))) : tol
    norm(ah,Inf)<tol
end

"""
iscp(m; tol)

Checks if the liouville representation of a map (in the natural, computational basis) is completely positive.
"""
function iscp(m; tol=0.0)
    ispossemidef(liou2choi(m))
end

"""
istp(m; tol)

Checks if the Liouville representation of a map is trace preserving (TP).
"""
function istp(m; tol=0.0)
    tol = tol==0.0 ? 1e2*eps(abs(one(eltype(m)))) : tol
    dsq = size(m,1)
    d = round(Int,sqrt(dsq))
    norm(m'*vec(eye(d))-vec(eye(d)),Inf) < tol
end

"""
ischannel(m; tol)

Checks if the Liouville representation of a map is completely positive (CP) and trace preserving (TP).
"""
function ischannel(m; tol=0.0)
    iscp(m,tol=tol) && istp(m,tol=tol)
end

"""
isunital(m; tol)

Checks the conditions for unitality of a map in Liouville representation.
"""
function isunital(m; tol=0.0)
    tol = tol==0.0 ? 1e2*eps(abs(one(eltype(m)))) : tol
    dsq = size(m,1)
    d = round(Int,sqrt(dsq))
    norm(m*vec(eye(d))-vec(eye(d)),Inf) < tol
end

"""
isliouvillian(m; tol)

Checks the conditions for a physical Liouvillian matrix (CPTP map generator)
"""
function isliouvillian(m;tol=0.0)
    tol = tol==0.0 ? 1e2*eps(abs(one(eltype(m)))) : tol

    mΓ = choi_liou_involution(m)
    d = round(Int,sqrt(size(m,1)))
    ω = sum([kron(ket(i,d),ket(i,d)) for i in 0:d-1])/sqrt(d)
    Πω = projector(ω)
    Πω⊥ = eye(d^2)-Πω

    return ishermitian(mΓ,tol=tol) && norm(ω'*m,Inf)<tol && ispossemidef(Πω⊥*mΓ*Πω⊥,tol=tol)
end

"""
liouvillian_violation(m; tol)

Computes violation of conditions for a physical Liouvillian matrix (CPTP map generator)
"""
function liouvillian_violation(m;tol=0.0)
    tol = tol==0.0 ? 1e2*eps(abs(one(eltype(m)))) : tol

    mΓ = choi_liou_involution(m)
    d = round(Int,sqrt(size(m,1)))
    ω = sum([kron(ket(i,d),ket(i,d)) for i in 0:d-1])/sqrt(d)
    Πω = projector(ω)
    Πω⊥ = eye(d^2)-Πω

    return (ahermpart(mΓ), ω'*m, Πω⊥*mΓ*Πω⊥)
end

"""
nearestu(l)

Computes the unitary CP map closest to a given CP map in an interferometric sense.
See D. Oi, [Phys. Rev. Lett. 91, 067902 (2003)](http://journals.aps.org/prl/abstract/10.1103/PhysRevLett.91.067902)
"""
function nearestu(l)
    c = liou2choi(l)
    vals,vecs = eig(Hermitian(c))
    imax = indmax(vals)
    Λ = mat(vecs[:,imax])
    U,Σ,V = svd(Λ)
    W = U*V'
    return kron(conj(W),W)
end


text:
function get_2Ddata(data, Nx, Ny)
    beam2d = reshape(data["copol"], (Ny, Nx))'
    theta = deg2rad.(reshape(data["theta"], (Ny, Nx)))'
    phi = deg2rad.(reshape(data["phi"], (Ny, Nx)))'
    x = rad2deg.(theta .* cos.( phi ))
    y = rad2deg.(theta .* sin.( phi ))
    return (beam2d, x, y)
end

function symmetrizer(psi_in_ipix, beam2d, Nx, Ny)
    beam_reshape = [beam2d'[1:Nx+1,:] rot180(beam2d'[Nx+1:end,:])]

    SIZE = size(beam_reshape)
    beam_sym = zeros(SIZE[1], SIZE[2])
    n = length(psi_in_ipix)

    for i in 1:n
        rot_i = Int(round(rad2deg(2psi_in_ipix[i])))
        #println(rot_i)
        beam_sym .+= circshift(beam_reshape, (0, rot_i))
    end

    beam_sym = beam_sym ./ n
    beam_reconst_sym = [beam_sym[:,1:Nx] ;reverse(rot180(beam_sym[1:end-1,Nx+1:Ny-1]), dims=2)]'

    return beam_reconst_sym
end

@. dBi(x) =10log10(x)

function gen_beammap(res::Resolution, beam2d, target_pixel)
    Nx, Ny = size(beam2d)
    beam_reshape = rot180([ beam2d'[1:Nx+1,:] reverse(rot180(beam2d'[Nx+1:end,:]), dims=2) ])
    itp = Interpolations.interpolate(beam_reshape, BSpline(Linear()))

    itp = extrapolate(itp, Flat())

    target_ang = pix2angRing(res, target_pixel)
    ipix = beam_pointor(res, target_ang[1], target_ang[2])

    beam_map = zeros(res.numOfPixels)
    @inbounds @threads for i in 1:res.numOfPixels
        ang = @views rad2deg.(pix2angRing(res, ipix[i]))
        beam_map[i] = itp(ang[1]*2, ang[2]*2)
    end
    return beam_map
end

function quaternion(theta, rotate_axis)
    V = @views @SVector [cos(theta/2), rotate_axis[1]*sin(theta/2), rotate_axis[2]*sin(theta/2), rotate_axis[3]*sin(theta/2)]
    return Quaternion(V)
end

function beam_pointor(res::Resolution, theta, phi)
    transposed_ipix = zeros(Int64, res.numOfPixels)
    qy = @views quaternion(theta, [0,-1,0])
    qz = @views quaternion(phi, [0,0,-1])
    @inbounds @threads for i in 1:res.numOfPixels
        exang = pix2angRing(res, i)
        vec = @views ang2vec(exang[1], exang[2])
        q_vec = Quaternion(0, vec[1], vec[2], vec[3])
        #Q = qy*qz
        P = qy*qz * q_vec / qz/qy
        ang = @views vec2ang(P.q1, P.q2, P.q3)
        transposed_ipix[i] = ang2pixRing(res, ang[1], ang[2])
    end
    return transposed_ipix
end

mutable struct AlmPair{T<:Matrix{ComplexF64}}
    alm::T
    almconj::T
end

function gen_Blm(res::Resolution, beammap)
    lmax = 3res.nside - 1
    beammap_pol = zeros(3, res.numOfPixels)
    beammap_pol[2,:] .= beammap .* (res.numOfPixels ./ (4π*sum(beammap)))
    Blm = hp.map2alm(beammap_pol, lmax = lmax, datapath = "/home/cmb/yusuket/program/MapData/healpy-data/.", use_pixel_weights=true)
    return AlmPair(Blm, conj.(Blm))
end

function gen_GaussBeammap(base, fwhm, peak_vec)
    sigma = fwhm/(2.0*sqrt(2.0*log(2.0)))
    r = @. sqrt((base[1] - peak_vec[1])^2 + (base[2] - peak_vec[2])^2 + (base[3] - peak_vec[3])^2)
    theta = zeros(length(r))
    @inbounds @threads for i in eachindex(r)
        r[i] = ifelse(r[i]>2.0, 2.0, r[i])
        theta[i] = 2.0*asin(r[i]/2.0)
    end
    return @. exp(-theta^2/(2.0*sigma^2))/((sigma*sqrt(2*π)))/((sigma*sqrt(2*π)))
end


text:
#TODO: update return_timings=false to a TypeFlag()
abstract type MatchingFlag end
struct returnTimings <: MatchingFlag end
struct noTimings <: MatchingFlag end

#match1 gets priority over match2
function combine_matchings(match1::Dict{Int,Int},match2::Dict{Int,Int})

    new_matching = copy(match1)
    for (i,j) in match2
        if !haskey(new_matching,i)
            new_matching[i] = j
        end
    end
    return new_matching
end

function degree_based_matching(A::SparseMatrixCSC{T,Int},B::SparseMatrixCSC{T,Int}) where T

	m = size(A,1)
	n = size(B,1)

    noise_factor = .01
	d_A = A*ones(m) .+ noise_factor*rand(m)
	d_B = B*ones(n) .+ noise_factor*rand(n)

	p_A = sort(1:m,by=i->d_A[i]) #breaks ties with some variation
	p_B = sort(1:n,by=j->d_B[j])

	return p_A,p_B

end

function low_rank_matching(U::Array{T,2},V::Array{T,2}) where T
    n,d1 = size(U)
    m,d2 = size(V)
    @assert d1 == d2

    matching_weights = zeros(T,d1)
    matchings = Array{Dict{Int,Int},1}(undef,d1)

    for i = 1:d1
        matchings[i],matching_weights[i] = rank_one_matching(U[:,i],V[:,i])
    end

    D = zeros(d1,d2)
    for j in 1:d2
        for i in 1:d1
            w = 0.0
            for (v_i,v_j) in matchings[j]
                edge_w = U[v_i,i]*V[v_j,i]
                if edge_w > 0
                    w += edge_w
                end
            end
            D[i,j] = matching_weights[i]/w
        end
    end

    d_j = maximum(D,dims= 1)
    opt_j = argmin(d_j).I[2]

    return matchings[opt_j], matching_weights[opt_j]

end

function rank_one_matching(u::Array{T,1},v::Array{T,1}) where T

    #get maximum matching by rearrangement theorem
    u_perm = sortperm(u,rev=true)
    v_perm = sortperm(v,rev=true)

    matching_weight = zero(T )
    Match_mapping = Dict{Int,Int}()
    for (i,j) in zip(u_perm,v_perm)

        #only match positive to positive and negative to negative
        if (u[i] > 0 && v[j] > 0 ) || (u[i] < 0 && v[j] < 0)
            Match_mapping[i] = j
            matching_weight += u[i]*v[j]
        end
    end
    return Match_mapping, matching_weight

end

function search_Krylov_space(A::ThirdOrderSymTensor,B::ThirdOrderSymTensor,U::Array{Float64,2},V::Array{Float64,2};returnScoringTimings::T=noTimings()) where {T <:MatchingFlag}

    best_score = -1
    best_i = -1
    best_j = -1
    best_matching = Dict{Int,Int}()

    Triangle_check = Dict{Array{Int,1},Int}()
    A_unique_nnz = length(A.values)
    B_unique_nnz = length(B.values)

    if A_unique_nnz > B_unique_nnz
        for i in 1:A_unique_nnz
            Triangle_check[A.indices[:,i]] = 1
        end
        Input_tensor = B
    else
        for i in 1:B_unique_nnz
            Triangle_check[B.indices[:,i]] = 1
        end
        Input_tensor = A
    end

    scoring_time =0.0
    for i in 1:size(U,2)
       for j in 1:size(V,2)
            A_to_B_mapping, _ = rank_one_matching(U[:,i],V[:,j])
            B_to_A_mapping, _ = rank_one_matching(V[:,j],U[:,i])
            if typeof(returnScoringTimings) === returnTimings
                if A_unique_nnz > B_unique_nnz
                    (matched_tris,gaped_tris),t = @timed TAME_score(Triangle_check,Input_tensor,B_to_A_mapping)
                else
                    (matched_tris,gaped_tris),t = @timed TAME_score(Triangle_check,Input_tensor,A_to_B_mapping)
                end
                scoring_time += t
            else # noTimings
                if A_unique_nnz > B_unique_nnz
                    matched_tris,gaped_tris = TAME_score(Triangle_check,Input_tensor,B_to_A_mapping)
                else
                    matched_tris,gaped_tris = TAME_score(Triangle_check,Input_tensor,A_to_B_mapping)
                end
            end

            if matched_tris > best_score
                best_matching = A_to_B_mapping
                best_score = matched_tris
                best_i = i
                best_j = j
            end
        end
    end
    if typeof(returnScoringTimings) === returnTimings
        return best_score, best_i, best_j, best_matching, scoring_time
    else
        return best_score, best_i, best_j, best_matching
    end
end

function search_Krylov_space(A::SymTensorUnweighted{S},B::SymTensorUnweighted{S},U::Array{Float64,2},V::Array{Float64,2};returnScoringTimings::T=noTimings())where {S <: Motif,T <:MatchingFlag}

    A_order,A_unique_nnz = size(A.indices)
    B_order,B_unique_nnz = size(B.indices)
    #@assert A_order == B_order

    best_score = -1
    best_i = -1
    best_j = -1
    best_matching = Dict{Int,Int}()

    motif_check = Dict{Array{Int,1},Int}()

    if A_unique_nnz > B_unique_nnz
        for i in 1:A_unique_nnz
            motif_check[A.indices[:,i]] = 1
        end
        Input_tensor = B
    else
        for i in 1:B_unique_nnz
            motif_check[B.indices[:,i]] = 1
        end
        Input_tensor = A
    end

    scoring_time =0.0
    for i in 1:size(U,2)
       for j in 1:size(V,2)
            A_to_B_mapping, _ = rank_one_matching(U[:,i],V[:,j])
            B_to_A_mapping, _ = rank_one_matching(V[:,j],U[:,i])
            if typeof(returnScoringTimings) === returnTimings
                if A_unique_nnz > B_unique_nnz
                    (matched_motifs,gaped_motifs),t = @timed TAME_score(motif_check,Input_tensor,B_to_A_mapping)
                else
                    (matched_motifs,gaped_motifs),t = @timed TAME_score(motif_check,Input_tensor,A_to_B_mapping)
                end
                scoring_time += t
            else
                if A_unique_nnz > B_unique_nnz
                    matched_motifs,gaped_motifs = TAME_score(motif_check,Input_tensor,B_to_A_mapping)
                else
                    matched_motifs,gaped_motifs = TAME_score(motif_check,Input_tensor,A_to_B_mapping)
                end
            end

            if matched_motifs > best_score
                best_matching = A_to_B_mapping
                best_score = matched_motifs
                best_i = i
                best_j = j
            end
        end
    end
    if typeof(returnScoringTimings) === returnTimings
        return best_score, best_i, best_j, best_matching, scoring_time
    else
        return best_score, best_i, best_j, best_matching
    end
end

#=
function search_Krylov_space(A::Array{SymTensorUnweighted,1},B::Array{SymTensorUnweighted,1},U::Array{Float64,2},V::Array{Float64,2})

    A_order,A_unique_nnz = size(A.indices)
    B_order,B_unique_nnz = size(B.indices)
    #@assert A_order == B_order

    best_score = -1
    best_i = -1
    best_j = -1
    best_matching = Dict{Int,Int}()

    motif_check = Dict{Array{Int,1},Int}()

    if A_unique_nnz > B_unique_nnz
        for i in 1:A_unique_nnz
            motif_check[A.indices[:,i]] = 1
        end
        Input_tensor = B
    else
        for i in 1:B_unique_nnz
            motif_check[B.indices[:,i]] = 1
        end
        Input_tensor = A
    end

    for i in 1:size(U,2)
       for j in 1:size(V,2)

            if A_unique_nnz > B_unique_nnz
                matched_motifs,gaped_motifs,matching = TAME_score(motif_check,Input_tensor,V[:,j],U[:,i])
            else
                matched_motifs,gaped_motifs,matching = TAME_score(motif_check,Input_tensor,U[:,i],V[:,j])
            end

            if matched_motifs > best_score
                best_matching = matching
                best_score = matched_motifs
                best_i = i
                best_j = j
            end
        end
    end
    return best_score, best_i, best_j, best_matching
end
=#

function search_Krylov_space(A::Array{SymTensorUnweighted{S},1},B::Array{SymTensorUnweighted{S},1},U::Array{Float64,2},V::Array{Float64,2};returnScoringTimings::T=noTimings()) where {S <: Motif,T <:MatchingFlag}

    @assert length(A) == length(B)
    for i = 1:length(A) #ensure orders are the same at each i
        @assert size(A[i].indices,1) == size(B[i].indices,1)
    end

    #create sets of edges
    larger_edge_sets = [ ]

    for (A_motifs,B_motifs) in zip(A,B)
        if size(A_motifs.indices,2) > size(B_motifs.indices,2)  #larger set of edges gets
            #push!(larger_edge_sets,Set(eachcol(A_motifs.indices)))
            push!(larger_edge_sets,Dict([(Array(col),1) for col in eachcol(A_motifs.indices)]))
        else
            push!(larger_edge_sets,Dict([(Array(col),1) for col in eachcol(B_motifs.indices)]))
        end
    end

    best_matching_score = -1
    best_matching_idx = (-1,-1)
    best_matched_motifs::Array{Int,1} = []
    best_mapping::Dict{Int,Int} = Dict()

    max_iter = size(U,2) - 1 #skip the first column of U and V
    scoring_timing = 0.0

    for i=1:max_iter
        for j=1:max_iter
            A_to_B_mapping, _ = rank_one_matching(U[:,i],V[:,j])
            B_to_A_mapping, _ = rank_one_matching(V[:,j],U[:,i])

            #TODO: handle mapping inversion

            #matched_motifs, matching_score = motif_matching_counts(A,B,mapping)
            matched_motifs = zeros(Int,length(A))
            if typeof(returnScoringTimings) === returnTimings

                for i =1:length(A)
                    if size(A[i].indices,2) > size(B[i].indices,2)
                        (matched_motifs[i],_),t = @timed TAME_score(larger_edge_sets[i], B[i], B_to_A_mapping)
                    else
                        (matched_motifs[i],_),t = @timed TAME_score(larger_edge_sets[i], A[i], A_to_B_mapping)
                    end
                    scoring_timing += t
                end
            else
                for i =1:length(A)
                    if size(A[i].indices,2) > size(B[i].indices,2)
                        matched_motifs[i],_ = TAME_score(larger_edge_sets[i], B[i], B_to_A_mapping)
                    else
                        matched_motifs[i],_ = TAME_score(larger_edge_sets[i], A[i], A_to_B_mapping)
                    end
                end
            end
            #matching_score = motif_matching_count(B[end], A[end], mapping)

            matching_score = 0
            for i= 1:length(A)
                matching_score += binomial(size(A[i].indices,1),2)*matched_motifs[i]
            end

            if matching_score > best_matching_score
                best_matching_score = matching_score
                best_matching_idx = (i,j)
                best_matched_motifs = copy(matched_motifs)
                best_mapping = A_to_B_mapping
            end

        end
    end

    best_i, best_j = best_matching_idx
    if typeof(returnScoringTimings) === returnTimings
        return best_matching_score, best_matched_motifs,  best_i, best_j, best_mapping, scoring_timing
    else
        return best_matching_score, best_matched_motifs,  best_i, best_j, best_mapping
    end
end

#TODO: TEST
#used when we don't want to recreate the triangle matching dictionary multiple times
function TAME_score(Triangle_Dict::Dict{Array{Int,1},Int},Input_tensor,
                    u::Array{Float64,1},v::Array{Float64,1})

    Match_mapping, _ = rank_one_matching(u,v)
    return TAME_score(Triangle_Dict,Input_tensor,Match_mapping)...,Match_mapping

end

function TAME_score(A,B,U::Array{Float64,2},V::Array{Float64,2};return_timings::T=noTimings()) where {T <: MatchingFlag}

    if typeof(return_timings) === returnTimings
        (Match_mapping, _), matching_time = @timed low_rank_matching(U,V)
        (triangle_count, gaped_triangles), scoring_time = @timed TAME_score(A,B,Match_mapping)
        return triangle_count, gaped_triangles,Match_mapping, matching_time, scoring_time
    else
        Match_mapping,weight = low_rank_matching(U,V)
        return TAME_score(A,B,Match_mapping)...,Match_mapping
    end

end

function TAME_score(A::ThirdOrderSymTensor, B::ThirdOrderSymTensor, X::SparseMatrixCSC{Float64,Int64};return_timings::T=noTimings()) where {T <: MatchingFlag}

    if typeof(return_timings) === returnTimings
        x ,bipartite_matching_time = @timed bipartite_matching(X)
        matching_dict = Dict(enumerate(x.match))
        (triangle_count, gaped_triangles), scoring_time = @timed TAME_score(A,B,matching_dict)
        return triangle_count, gaped_triangles, matching_dict, bipartite_matching_time, scoring_time
    else
        x = bipartite_matching(X)
        matching_dict = Dict(enumerate(x.match))
        return TAME_score(A,B,matching_dict)..., matching_dict
    end

end

function TAME_score(A::Union{ThirdOrderSymTensor,Array{SymTensorUnweighted{S},1}},
                    B::Union{ThirdOrderSymTensor,SymTensorUnweighted{S},Array{SymTensorUnweighted{S},1}},X::Array{Float64,2};return_timings::T=noTimings(),kwargs...) where {T <: MatchingFlag,S <: Motif}

    if typeof(return_timings) === returnTimings
        (_,_,matching,_) ,matching_time = @timed bipartite_matching_primal_dual(X;kwargs...)
        matching_dict = matching_array_to_dict(matching)
        (triangle_count, gaped_triangles), scoring_time = @timed TAME_score(A,B,matching_dict)
        return triangle_count, gaped_triangles, matching_dict, matching_time, scoring_time
    else
        _,_,matching,_ = bipartite_matching_primal_dual(X;kwargs...)
        matching_dict = matching_array_to_dict(matching)
        return TAME_score(A,B,matching_dict)...,matching_dict
    end
end

function TAME_score(A::SymTensorUnweighted{S},B::SymTensorUnweighted{S},X::Array{Float64,2};return_timings::T=noTimings(),kwargs...) where {T <: MatchingFlag,S <: Motif}

    if typeof(return_timings) === returnTimings
        (_,_,matching,_) ,matching_time = @timed bipartite_matching_primal_dual(X;kwargs...)
        matching_dict = matching_array_to_dict(matching)
        (triangle_count, gaped_triangles), scoring_time = @timed TAME_score(A,B,matching_dict)
        return triangle_count, gaped_triangles, matching_dict, matching_time, scoring_time
    else
        _,_,matching,_ = bipartite_matching_primal_dual(X;kwargs...)
        matching_dict = matching_array_to_dict(matching)
        #return TAME_score(A,B,matching_array_to_dict(matching))
        return TAME_score(A,B,matching_dict)...,matching_dict
        #BUG?
    end
end

function TAME_score(A::ThirdOrderSymTensor,B::ThirdOrderSymTensor,x::Array{Float64,1};return_timings::T=noTimings(),kwargs...) where {T <: MatchingFlag}

    X = reshape(x,A.n,B.n)
    if typeof(return_timings) === returnTimings
        (_,_,matching,_) ,matching_time = @timed bipartite_matching_primal_dual(X;kwargs...)
        matching_dict = matching_array_to_dict(matching)
        (triangle_count, gaped_triangles), scoring_time = @timed TAME_score(A,B,matching_dict)
        return triangle_count, gaped_triangles, matching_dict, matching_time, scoring_time
    else
        _,_,matching,_ = bipartite_matching_primal_dual(X;kwargs...)
        matching_dict = matching_array_to_dict(matching)
        return TAME_score(A,B,matching_dict)...,matching_dict
    end

end

"""
   Match_mapping is expected to map V_A -> V_B
"""

function TAME_score(A::ThirdOrderSymTensor,B::ThirdOrderSymTensor,Match_mapping::Dict{Int,Int})
    # matching oriented A -> B

    match_len = length(Match_mapping)

    Triangle_check = Dict{Array{Int,1},Int}()
    gaped_triangles = 0
    triangle_count = 0

    if size(A.indices,2) > size(B.indices,2)

        for i in 1:size(A.indices,2)
            Triangle_check[A.indices[:,i]] = 1
        end

        #invert to map V_B indices to V_A
        Match_mapping = Dict(value => key for (key, value) in Match_mapping)

        for i in 1:size(B.indices,2)
            v_i,v_j,v_k = B.indices[:,i]

            matched_triangle =
              sort([get(Match_mapping,v_i,-1),get(Match_mapping,v_j,-1),get(Match_mapping,v_k,-1)])
            match = get(Triangle_check,matched_triangle,0)
            if match == 1
                triangle_count += 1
            else
                gaped_triangles += 1
            end
        end

    else
        for i in 1:size(B.indices,2)
            Triangle_check[B.indices[:,i]] = 1
        end

        for i in 1:size(A.indices,2)
            v_i,v_j,v_k = A.indices[:,i]
            matched_triangle =
               sort([get(Match_mapping,v_i,-1), get(Match_mapping,v_j,-1),get(Match_mapping,v_k,-1)])

            match = get(Triangle_check,matched_triangle,0)
            if match == 1
                triangle_count += 1
            else
                gaped_triangles += 1
            end
        end
    end

   # triangles, triangle_weight = count_triangles(sub_A,sub_B)

    return triangle_count, gaped_triangles

end

#TODO: may remove the array mapping routines
function TAME_score(A::SymTensorUnweighted{S},B::SymTensorUnweighted{S}, mapping::Array{Int,1}) where {S <: Motif}

    # matching oriented B -> A
    @assert size(A.indices,1) == size(B.indices,1)

    if size(B.indices,2) > size(A.indices,2)
        inverted_mapping = -ones(Int, maximum(mapping))

        for (v,u) in enumerate(mapping)
            if u != -1
                inverted_mapping[u] = v
            end
        end

        #return (-1,"ERROR:mapping needs to be inverted")
        return TAME_score(B, A,inverted_mapping)
    end

    order =  size(A.indices,1)
    A_motifs = Set(eachcol(A.indices))

    matched_motifs = 0
    missed_motifs = 0

    for idx =1:size(B.indices,2)
        edge = B.indices[:,idx]
        new_edge = [mapping[i] for i in edge]
        sort!(new_edge)

        if new_edge in A_motifs
            matched_motifs += 1
        else
            missed_motifs += 1
        end

    end
    return matched_motifs, matched_motifs

end

function TAME_score(A::Array{SymTensorUnweighted{S},1},B::Array{SymTensorUnweighted{S},1},A_to_B_mapping::Dict{Int,Int}) where {S <: Motif}
    # matching oriented A -> B

    @assert length(A) == length(B)
    for i = 1:length(A) #ensure orders are the same at each i
        @assert size(A[i].indices,1) == size(B[i].indices,1)
    end

    B_to_A_mapping::Dict{Int,Int} = Dict([(j,i) for (i,j) in A_to_B_mapping])

    #create sets of edges
    larger_edge_sets = [ ]

    for (A_motifs,B_motifs) in zip(A,B)
        if size(A_motifs.indices,2) > size(B_motifs.indices,2)  #larger set of edges gets
            #push!(larger_edge_sets,Set(eachcol(A_motifs.indices)))
            push!(larger_edge_sets,Dict([(Array(col),1) for col in eachcol(A_motifs.indices)]))
        else
            push!(larger_edge_sets,Dict([(Array(col),1) for col in eachcol(B_motifs.indices)]))
        end
    end

    best_matching_score = -1
    best_matched_motifs::Array{Int,1} = []

    matched_motifs = zeros(Int,length(A))
    for i =1:length(A)
        if size(A[i].indices,2) > size(B[i].indices,2)
            matched_motifs[i],_ = TAME_score(larger_edge_sets[i], B[i], B_to_A_mapping)
        else
            matched_motifs[i],_ = TAME_score(larger_edge_sets[i], A[i], A_to_B_mapping)
        end
    end
    #matching_score = motif_matching_count(B[end], A[end], mapping)

    matching_score = 0
    for i= 1:length(A)
        #println(binomial(size(A[i],1),2))
        #println(matched_motifs[i])
        matching_score += binomial(size(A[i].indices,1),2)*matched_motifs[i]
    end

    return matching_score, matched_motifs
end

function TAME_score(A::SymTensorUnweighted{Clique},B::SymTensorUnweighted{Clique}, mapping::Dict{Int,Int})
    # matching oriented A -> B

    @assert size(A.indices,1) == size(B.indices,1)

    #=
    if size(B.indices,2) > size(A.indices,2)
        println("inverting")
        inverted_mapping = Dict{Int,Int}()

        for (v,u) in mapping
            if u != -1
                inverted_mapping[u] = v
            end
        end

        #return (-1,"ERROR:mapping needs to be inverted")
        return TAME_score(B, A,inverted_mapping)
    end
    =#

    order =  size(A.indices,1)
    B_motifs = Set(eachcol(B.indices))

    matched_motifs = 0
    gaped_motifs = 0

    for idx =1:size(A.indices,2)
        edge = A.indices[:,idx]
        new_edge = sort([get(mapping,i,-1) for i in edge])
        if new_edge in B_motifs
            matched_motifs += 1
        else
            gaped_motifs += 1
        end

    end
    return matched_motifs, gaped_motifs

end

function TAME_score(A::SymTensorUnweighted{Cycle},B::SymTensorUnweighted{Cycle}, mapping::Dict{Int,Int})
    # matching oriented B -> A

    @assert size(A.indices,1) == size(B.indices,1)

    if size(B.indices,2) > size(A.indices,2)
        inverted_mapping = Dict{Int,Int}()

        for (v,u) in mapping
            if u != -1
                inverted_mapping[u] = v
            end
        end
        return TAME_score(B, A,inverted_mapping)
    end

    order =  size(A.indices,1)
    A_cycle_hashes = Set{Float64}()
	for i= 1:size(A.indices,2)
		#A.indices are assumed to have the smallest index a the start of the column
		push!(A_cycle_hashes,DistributedTensorConstruction.cycle_hash(A.indices[:,i]))
	end

    matched_motifs = 0
    missed_motifs = 0

	l = size(B.indices,1)
    for idx =1:size(B.indices,2)
        edge = B.indices[:,idx]
        new_edge = [get(mapping,i,-1) for i in edge]

        new_edge = [new_edge[(i+argmin(new_edge)+2)%l+1] for i in 1:l]
			#cycles the edge so smallest index is first
		hash = DistributedTensorConstruction.cycle_hash(new_edge)

        if hash in A_cycle_hashes
            matched_motifs += 1
        else
            missed_motifs += 1
        end

    end

    return matched_motifs, missed_motifs
end

function TAME_score(Triangle_Dict::Dict{Array{Int,1},Int},Input_tensor::ThirdOrderSymTensor,
                    Match_mapping::Union{Dict{Int,Int},Vector{Int}})

    triangle_count = 0
    gaped_triangles = 0

    for i in 1:length(Input_tensor.values)
        v_i,v_j,v_k = Input_tensor.indices[:,i]

        matched_triangle =
          sort([get(Match_mapping,v_i,-1),get(Match_mapping,v_j,-1),get(Match_mapping,v_k,-1)])

        match = get(Triangle_Dict,matched_triangle,0)
        if match == 1
            triangle_count += 1
        else
            gaped_triangles += 1
        end
    end
    return triangle_count, gaped_triangles
end

#TODO: may remove, see above

function TAME_score(A_motifs::Dict{Array{Int,1},Int}, B::SymTensorUnweighted{S}, mapping::Array{Int,1}) where {S <: Motif}

    order =  size(B.indices,1)
    #A_motifs = Set(eachcol(Indices_A))

    matched_motifs = 0
    missed_motifs = 0

    for idx =1:size(B.indices,2)
        edge = B.indices[:,idx]
        new_edge = [mapping[i] for i in edge]
        sort!(new_edge)

        if new_edge in A_motifs
            matched_motifs += 1
        else
            missed_motifs += 1
        end

    end
    return matched_motifs, missed_motifs

end

function TAME_score(A_motifs::Dict{Array{Int,1},Int}, B::SymTensorUnweighted{S}, mapping::Dict{Int,Int}) where {S <: Motif}

    order =  size(B.indices,1)
    #A_motifs = Set(eachcol(Indices_A))

    matched_motifs = 0
    gaped_motifs = 0

    for idx =1:size(B.indices,2)
        edge = B.indices[:,idx]
        new_edge = [get(mapping,i,-1) for i in edge]
        sort!(new_edge)

        if haskey(A_motifs,new_edge)
            matched_motifs += 1
        else
            gaped_motifs += 1
        end

    end
    return matched_motifs, gaped_motifs

end

function edges_matched(A::SparseMatrixCSC{T,Int}, B::SparseMatrixCSC{T,Int}, mapping::Dict{Int,Int}) where T

    edges_matched = 0
    edges_missed = 0

    B_nnz = Set(zip(findnz(B)[1:2]...))

    for (i,j,_) in zip(findnz(A)...)

        ip = get(mapping,i,-1)
        jp = get(mapping,j,-1)

        if (ip,jp) in B_nnz
            edges_matched += 1
        else
            edges_missed += 1
        end

    end

    return edges_matched, edges_missed
end

function matching_array_to_dict(matching::Vector{Int})

    matching_dict = Dict{Int,Int}()
    for (i,j) in enumerate(matching)
        if i != -1 && j != -1 #-1 indicates unmatched
            matching_dict[i] = j
        end
    end
    return matching_dict
end

"""------------------------------------------------------------------------------
  A function for solving the a maximum bipartite matching problem with a dense
  m x n input X.

  Inputs
  ------
  * X - (Matrix{Float64}):
    The weights of the input bipartite graph.
  * tol - (Float):
    The tolerance to determine whether or not the flow through an edge is at
    capacity.
  * 'normalize_weights' - (Bool):
    Whether or not to normalize the maximum edge weight to be 1.0, implictly
    alters the tolerance to which the problem is solved.

  Outputs
  -------
  * val - (Float):
    The value of the maximum weighted matching.
  * noute - (Int):
    The cardinality of the matching found.
  * match1 - (Array{Int,1}):
    The mapping linking left hand side nodes to the right hand side. The ith
    entry maps node i to a right hand side node.
  * match2 - (Array{Int,1}):
    The mapping linking right hand side nodes to the right hand side. The jth
    entry maps node j to a left hand side node. The last n entries in the array
    represent dummy nodes which are used to account non-matches which may arise
    from negative weights.
------------------------------------------------------------------------------"""
function bipartite_matching_primal_dual(X::Union{Matrix{T},Adjoint{T,Matrix{T}}};primalDualTol::Float64=1e-8,
                                       normalize_weights::Bool=false) where T
    #to get the access pattern right, we must match the right hand side to the left hand side.

	m,n = size(X)
	#@assert m >= n  #error occurs when m < n
    if m < n
        val, noute, B_to_A, _ = bipartite_matching_primal_dual(X';primalDualTol,normalize_weights)
        A_to_B = -ones(Int,m) # negative 1 to make no mistake that node is unmatched
        for (j,i) in enumerate(B_to_A)
            if i != -1
                A_to_B[i] = j
            end
        end
        return val, noute, A_to_B, B_to_A
    end

    # variables used for the primal-dual algorithm
    # normalize ai values # updated on 2-19-2019
    if normalize_weights
        X ./= maximum(abs.(X))
    end

	#initialize variables
    alpha=zeros(Float64,n)
    bt=zeros(Float64,m+n)#beta
    match1 = zeros(Int64,n)
    match2 = zeros(Int64,n+m)
    queue=zeros(Int64,m)
    t=zeros(Int64,m+n)
    tmod = zeros(Int64,m+n)
    ntmod=0

    # initialize the primal and dual variables

	for j = 1:n
		for i=1:m
			if X[i,j] > alpha[j]
			   alpha[j]=X[i,j]
			end
		end
    end

    # dual variables (bt) are initialized to 0 already
    # match1 and match2 are both 0, which indicates no matches

    j=1
    @inbounds while j<=n

        for i=1:ntmod
            t[tmod[i]]=0
        end
        ntmod=0
        # add i to the stack
        head=1
        tail=1
        queue[head]=j
        while head <= tail && match1[j]==0 #queue empty + i is unmatched

 		 	k=queue[head]
			#println("begining of queue loop")

            for i=1:m+1 #iterate over column k

				if i == m+1 #check the dummy node
					i = k+m
				end
				if i == k+m #dummy nodes don't have weight
					if 0.0 < alpha[k] + bt[i] - primalDualTol
						continue
					end
                elseif X[i,k] < alpha[k] + bt[i] - primalDualTol
					continue
                end # skip if tight

                if t[i]==0
                    tail=tail+1 #put the potential match in the queue
                    if tail <= m
                        queue[tail]=match2[i]
					end
                    t[i]=k  #try vertex k for vertex j
                    ntmod=ntmod+1
                    tmod[ntmod]=i
                    if match2[i]<1  #if i is unmatched
                        while i>0   #unfurl out to get an augmented path
                            match2[i]=t[i]
                            k=t[i]
                            temp=match1[k]
                            match1[k]=i
							i=temp
						end
                        break
                    end
                end
            end
            head=head+1

        end

        #if node j was unable to be matched, update flows and search for new augmenting path
		if match1[j] < 1
            theta=Inf
            for i=1:head-1
                t1=queue[i]
                for t2=1:m
                    if t[t2] == 0 && alpha[t1] + bt[t2] - X[t2,t1] < theta
                        theta = alpha[t1] + bt[t2] - X[t2,t1]
                    end
                end
				#check t1's dummy node
				if t[t1 + m] == 0 && alpha[t1] + bt[t1 + m] < theta
                        theta = alpha[t1] + bt[t1 + m]
				end
            end

            for i=1:head-1
                alpha[queue[i]] -= theta
            end
            for i=1:ntmod
                bt[tmod[i]] += theta
            end

            continue
        end

        j=j+1
    end

	#count
    val=0.0
    for j=1:n
        for i=1:m
            if i==match1[j]
                val=val+X[i,j]
            end
        end
    end

	#count how many are properly matched
    noute = 0
    for j=1:n
        if match1[j]<=m
            noute=noute+1
        end
	end

    A_to_B = -ones(Int,m)
    B_to_A = match1
    for (j,i) in enumerate(B_to_A)
        A_to_B[i] = j
    end

    return val, noute, A_to_B, B_to_A
end

#
# Local Search Methods
#

function seq_similarity(U,V,matching)

	sim = 0.0
	for (i,ip) in matching
		sim += dot(U[i,:],V[ip,:])
	end
	return sim
end

function TAME_score(A_motifs::Set{Vector{Int}},B::SymTensorUnweighted{Clique},
						B_edge_mask::Set{Int64},U::Matrix{T},V::Matrix{T},
						mapping::Dict{Int,Int}) where T

	#@assert maximum(keys(mapping)) <= B.n  # mapping oriented from B -> A

    #top_sim = 0
	motifs_matched = 0

    for idx in B_edge_mask

        edge = B.indices[:,idx]

        new_edge = [get(mapping,i,-1) for i in edge]
        sort!(new_edge)

		#_,jp,kp = edge
		#_,j,k = new_edge

        if new_edge in A_motifs

			#top_sim += 1 + max(dot(U[j,:],V[jp,:]) + dot(U[k,:],V[kp,:]),
		    #					   dot(U[j,:],V[kp,:]) + dot(U[k,:],V[jp,:]))
			motifs_matched += 1

        end

    end

    return top_sim, motifs_matched

end

function TAME_score(A_motifs::Set{Vector{Int}},B::SymTensorUnweighted{Clique},
					B_edge_mask::Set{Int64}, mapping::Dict{Int,Int}) where T

	motifs_matched = 0

    for idx in B_edge_mask

        edge = B.indices[:,idx]

        new_edge = [get(mapping,i,-1) for i in edge]
        sort!(new_edge)

        if new_edge in A_motifs

			motifs_matched += 1

        end

    end
    return motifs_matched

end

function TAME_score(A_motifs::Set{Vector{Int}},B::ThirdOrderSymTensor,
					B_edge_mask::Set{Int64}, mapping::Dict{Int,Int}) where T

	motifs_matched = 0

    for idx in B_edge_mask

        edge = B.indices[:,idx]

        new_edge = [get(mapping,i,-1) for i in edge]
        sort!(new_edge)

        if new_edge in A_motifs

			motifs_matched += 1

        end

    end

    return motifs_matched

end

max_stars_repo_name: chuckcol/LambdaTAME

max_issues_repo_name: chuckcol/LambdaTAME

max_forks_repo_name: chuckcol/LambdaTAME


text:
for f in (:real, :imag, :conj)
    @eval ($f)(a::STaylor1{N,T}) where {N,T<:Number} = STaylor1{N,T}(($f).(a.coeffs))
end

adjoint(a::STaylor1) = conj(a)
isinf(a::STaylor1) = any(isinf.(a.coeffs))
isnan(a::STaylor1) = any(isnan.(a.coeffs))

function abs(a::STaylor1{N,T}) where {N,T<:Real}
    if a[0] > zero(T)
        return a
    elseif a[0] < zero(T)
        return -a
    else
        throw(ArgumentError(
        """The 0th order Taylor1 coefficient must be non-zero
        (abs(x) is not differentiable at x=0)."""))
    end
end

abs2(a::STaylor1{N,T}) where {N,T<:Real} = a^2

deg2rad(z::STaylor1{N, T}) where {N, T<:AbstractFloat} = z * (convert(T, pi) / 180)
deg2rad(z::STaylor1{N, T}) where {N, T<:Real} = z * (convert(float(T), pi) / 180)

rad2deg(z::STaylor1{N, T}) where {N, T<:AbstractFloat} = z * (180 / convert(T, pi))
rad2deg(z::STaylor1{N, T}) where {N, T<:Real} = z * (180 / convert(float(T), pi))

function mod(a::STaylor1{N,T}, x::T) where {N, T<:Real}
    return STaylor1{N,T}(ntuple(i -> i == 1 ? mod(constant_term(a), x) : a.coeffs[i], Val(N)))
end

function mod(a::STaylor1{N,T}, x::S) where {N, T<:Real, S<:Real}
    R = promote_type(T, S)
    a = convert(STaylor1{N,R}, a)
    return mod(a, convert(R, x))
end

function rem(a::STaylor1{N,T}, x::T) where {N, T<:Real}
    return STaylor1{N,T}(ntuple(i -> i == 1 ? rem(constant_term(a), x) : a.coeffs[i], Val(N)))
end

function rem(a::STaylor1{N,T}, x::S) where {N, T<:Real, S<:Real}
    R = promote_type(T, S)
    a = convert(STaylor1{N,R}, a)
    return rem(a, convert(R, x))
end

function mod2pi(a::STaylor1{N,T}) where {N, T<:Real}
    return STaylor1{N,T}(ntuple(i -> i == 1 ? mod2pi(constant_term(a)) : a.coeffs[i], Val(N)))
end


text:
module FactorPostProcessing

export svd_logs

using LinearAlgebra, Statistics, LinearAlgebra.BLAS
using BeastUtils.Logs, BeastUtils.MatrixUtils, BeastUtils.ESS

const DEBUG = false

function my_show(x, s::String)
    println("$s:")
    display(x)
    println()
end

function process_log(log_path::String, header::Union{Regex, String}, k::Int,
                    p::Int;
                    fac_header::String = "factors.",
                    rotate_factors::Bool = false,
                    find_best::Bool = true,
                    do_svd::Bool = true,
                    rotation_cols::Vector{Int} = zeros(Int, k),
                    transposed::Bool = false,
                    relevant_cols::Vector{Int} = collect(1:p),
                    relevant_rows::Vector{Int} = collect(1:k))

    cols, data = get_log(log_path, burnin = 0.0)
    L_inds = findall(x -> startswith(x, header), cols)
    @show k
    @show p
    @assert length(L_inds) == k * p

    f_inds = rotate_factors ? findall(x -> startswith(x, fac_header), cols) : Int[]
    n_taxa = div(length(f_inds), k)
    @assert length(f_inds) == n_taxa * k

    L_cols = cols[L_inds]
    L_data = data[:, L_inds]

    f_cols = cols[f_inds]
    f_data = data[:, f_inds]

    n = size(data, 1)

    d_storage = zeros(n, k)
    v_storage = zeros(n, k * p)
    l_storage = zeros(n, k * p)
    f_storage = zeros(n, n_taxa * k)

    ke = length(relevant_rows)
    pe = length(relevant_cols)

    do_svd = do_svd && ke != 0

    other_rows = setdiff(1:k, relevant_rows)
    other_cols = setdiff(1:k, relevant_cols)

    L_full = zeros(k, p)
    Ft_full = zeros(k, n_taxa)
    UtFt_full = zeros(k, n_taxa)

    L = zeros(ke, pe)
    Ft = zeros(ke, n_taxa)
    UtFt = zeros(ke, n_taxa)

    fac_inds = vec([(i - 1) * k + row for row in relevant_rows, i in 1:n_taxa])

    S = zeros(ke)
    Vt = zeros(ke, pe)

    for i = 1:n
        fill_L!(L_full, L_data, i, transposed)
        for j = 1:pe
            for l = 1:ke
                L[l, j] = L_full[relevant_rows[l], relevant_cols[j]] # TODO: add inbounds
            end
        end

        if DEBUG
            my_show(L_full, "L_full")
            my_show(L, "L")
        end

        if do_svd
            s = svd(L)
            S .= s.S
            Vt .= s.Vt
        else
            for i = 1:ke
                S[i] = norm(@view L[i, :])
                for j = 1:pe
                    Vt[i, j] = L[i, j] / S[i]
                end
            end
        end

        if do_svd
            for j = 2:ke
                @assert S[j - 1] >= S[j]
            end
        end
        # rotate_svd!(s, cols = rotation_cols)
        if DEBUG
            L2 = U * Diagonal(S) * Vt
            @assert maximum(abs.(L2 - L)) < 1e-12
        end

        for j = 1:ke
            d_storage[i, relevant_rows[j]] = S[j]
            offset = (relevant_rows[j] - 1) * p
            for l = 1:pe
                col = relevant_cols[l]
                ind = offset + col
                v_storage[i, ind] = Vt[j, l]
                l_storage[i, ind] = Vt[j, l] * S[j]
            end
        end
        for j = 1:(k - ke)
            row = other_rows[j]
            l = @view L_full[row, :]
            nrm = norm(l)
            d_storage[i, row] = nrm

            rng = ((row - 1) * p + 1):(row * p)
            l_storage[i, rng] .= l
            if nrm != 0
                v_storage[i, rng] .= l / nrm # TODO: memory efficiency
            end
        end

        #
        # d_storage[i, :] .= s.S
        # v_storage[i, :] .= vec(s.Vt')
        #
        # l_storage[i, :] .= vec(s.Vt' * Diagonal(s.S))

        if rotate_factors
            U = Matrix(Diagonal(ones(ke)))
            if do_svd
                U = s.U
            end
            copyto!(Ft, @view f_data[i, fac_inds])
            gemm!('T', 'N', 1.0, U, Ft, 0.0, UtFt)
            for j = 1:n_taxa
                offset = (j - 1) * k
                for l = 1:ke
                    f_storage[i, offset + relevant_rows[l]] = UtFt[l, j]
                end
            end
            for j = 1:n_taxa
                offset = (j - 1) * k
                for l = 1:(k - ke)
                    ind = offset + other_rows[l]
                    f_storage[i, ind] = f_data[i, ind]
                end
            end
        end

        # if i > 2
        #     error("break")
        # end
    end

    force_rotation = false
    if any(rotation_cols .!= 0)
        force_rotation = true
    end
    if !force_rotation
        if !find_best
            fill!(rotation_cols, 1)
        else
            rotation_cols .= find_rotation_cols(v_storage, k, p)
        end
    end
    reflect!(v_storage, l_storage, f_storage, rotation_cols, k, p)
    return d_storage, v_storage, l_storage, f_storage
end

function fill_L!(L::Matrix{Float64}, data::Matrix{Float64}, row::Int,
                 transposed::Bool)
    k, p = size(L)
    if transposed
        for i = 1:p
            offset = k * (i - 1)
            for j = 1:k
                L[j, i] = data[row, offset + j]
            end
        end
    else
        for i = 1:k
            offset = p * (i - 1)
            for j = 1:p
                L[i, j] = data[row, offset + j]
            end
        end
    end
end

function rotate_svd!(s::SVD{Float64,Float64,Array{Float64,2}}; cols::Vector{Int} = ones(Int, length(s.S)))
    k = length(s.S)
    rev = ones(k)
    for i = 1:k
        if s.Vt[i, cols[i]] < 0.0
            rev[i] = -1.0
        end
    end
    Q = Diagonal(rev)
    s.U .= s.U * Q #TODO: make memory efficient
    s.Vt .= Q * s.Vt #TODO: see above
end

function find_rotation_cols(data::Matrix{Float64}, k::Int, p::Int)
    abs_mat = abs.(data)
    means = vec(mean(abs_mat, dims = 1))
    vars = vec(var(abs_mat, dims = 1))

    zs = means ./ sqrt.(vars)
    nan_inds = findnans(zs)
    zs[nan_inds] .= 0.0
    cols = zeros(Int, k)
    for i = 1:k
        l = ((i - 1) * p) + 1
        u = i * p
        zmax = findmax(@view zs[l:u])
        cols[i] = zmax[2]
    end
    return cols
end

function process_loadings(log_path::String, header::Union{Regex, String}, k::Int, p::Int; find_best::Bool = true)
    cols, data = get_log(log_path, header, burnin = 0.0)

    @assert length(cols) == k * p

    rotation_cols = ones(Int, k)
    if find_best
        rotation_cols .= find_rotation_cols(data, k, p)
    end
    reflect!(data, rotation_cols, k, p)
    return data, cols
end

function reflect!(data::Matrix{Float64}, loadings::Matrix{Float64},
                    factors::Matrix{Float64}, cols::Vector{Int}, k::Int, p::Int)
    @show cols
    @assert length(cols) == k
    n, kp = size(data)
    @assert kp == k * p

    @assert size(factors, 1) == n
    n_taxa = div(size(factors, 2), k)
    @assert n_taxa * k == size(factors, 2)

    for i = 1:k
        offset = (i - 1) * p
        check_ind = offset + cols[i]
        fac_inds = i:k:(n_taxa * k)
        for j = 1:n
            if data[j, check_ind] < 0.0
                for k = (offset + 1):(i * p)
                    data[j, k] = -data[j, k]
                    loadings[j, k] = -loadings[j, k]
                end
                for ind in fac_inds
                    factors[j, ind] = -factors[j, ind]
                end
            end
        end
    end
    return data
end

function svd_logs(path::String, new_path::String, k::Int, p::Int;
        cols::Vector{Int} = zeros(Int, k),
        Lid::String = "L",
        fid::String = "factors.",
        rotate_factors::Bool = false,
        transposed::Bool = false,
        do_svd::Bool = true,
        relevant_rows::Vector{Int} = collect(1:k),
        relevant_cols::Vector{Int} = collect(1:p))

    d, v, l, f = process_log(path, Lid, k, p, rotation_cols = cols,
                                transposed = transposed,
                                fac_header = fid,
                                rotate_factors = rotate_factors,
                                relevant_rows = relevant_rows,
                                relevant_cols = relevant_cols,
                                do_svd = do_svd)
    d_labels = ["sv$i" for i = 1:k]
    v_labels = vec(["V$i$j" for j = 1:p, i = 1:k])
    l_labels = vec(["L$i$j" for j = 1:p, i = 1:k])
    cols = get_cols(path)
    f_inds = rotate_factors ? findall(x -> startswith(x, fid), cols) : Int[]
    f_labels = string.(cols[f_inds])

    states = get_log(path, "state", burnin = 0.0)[2]

    labels = ["state"; d_labels; v_labels; l_labels; f_labels]
    data = [states d v l f]

    make_log(new_path, data, labels, includes_states = true)
end

function svd_logs(path::String, new_path::String;
    prec_start = "factorPrecision",
    Lid::String = "L",
    fid::String = "factors.",
    rotate_factors::Bool = false,
    transposed::Bool = false)

    cols = Logs.get_cols(path)
    p = length(findall(x -> startswith(x, prec_start), cols))
    kp = length(findall(x -> startswith(x, Lid), cols))
    k, r = divrem(kp, p)
    if r != 0
        error("Unable to determine the number of factors and traits.")
    end

    return svd_logs(path, new_path, k, p, Lid = Lid, fid = fid,
            rotate_factors = rotate_factors, transposed = transposed)
end

# k = 2
# p = 100
#
# path = joinpath(Directories.desktop, "simFactor_N100_P100_K2_gibbs.log")
# new_path = joinpath(Directories.desktop, "gibbs.log")
#
# svd_logs(path, new_path, k, p)

end


text: plot(randn(10, 3))


text:
#######################################################################################
#                                   Load Packages
#######################################################################################
cd(@__DIR__)
using Pkg
Pkg.activate("../..")
using Revise, StatsPlots, ACTRModels, Distributions, FFTDists, DifferentialEvolutionMCMC
include("Semantic_FFT_Model.jl")
include("model_functions.jl")
include("../../utilities/plotting.jl")
Random.seed!(373401)
#######################################################################################
#                                   Generate Data
#######################################################################################
blc = 1.5
δ = 1.0
parms = (noise = true, τ = 0.0, s = 0.2, mmp = true)
stimuli = get_stimuli()
n_reps = 10
data = map(s -> simulate(parms, s, n_reps; blc, δ), stimuli)
#######################################################################################
#                                    Define Model
#######################################################################################
priors = (
    blc = (Normal(1.5, 1),),
    δ = (Truncated(Normal(1.0, .5), 0.0, Inf),)
  )
bounds = ((-Inf,Inf),(eps(),Inf))
#######################################################################################
#                                 Estimate Parameters
#######################################################################################
model = DEModel(parms; priors, model=loglike, data)
de = DE(;bounds, burnin=1000, priors, n_groups=2, Np=4)
n_iter = 2000
chains = sample(model, de, MCMCThreads(), n_iter, progress=true)
#######################################################################################
#                                  Posterior Correlations
#######################################################################################
cors = cor(chains)
#######################################################################################
#                                         Plot
#######################################################################################
pyplot()
posteriors = plot(chains, seriestype=:pooleddensity, grid=false, titlefont=font(10),
     xaxis=font(8), yaxis=font(8), color=:grey, size=(300,250))
plot!(title = "", xlabel = "blc", subplot = 1)
plot!(title = "", xlabel = "δ", subplot = 2)
#######################################################################################
#                                  Posterior Predictive
#######################################################################################
rt_preds(s) = posterior_predictive(x -> simulate(parms, s, n_reps; x...), chains, 1000)
temp_preds = map(s -> rt_preds(s), stimuli)
preds = merge.(temp_preds)
posterior_plots = Plots.Plot[]
for (pred, stimulus) in zip(preds, stimuli)
    prob_yes = length(pred.yes)/(length(pred.yes) + length(pred.no))
    object = stimulus.object
    category = stimulus.category
    hist = histogram(layout=(1,2), xlims=(0,2.5),  ylims=(0,3.5), title="$object-$category",
        grid=false, titlefont=font(6), xaxis=font(6), yaxis=font(6), xlabel="Yes RT",
        xticks = 0:2, yticks=0:3)

    if !isempty(pred.yes)
        histogram!(hist, pred.yes, xlabel="Yes RT", norm=true, grid=false, color=:grey, leg=false,
            size=(300,250), subplot=1)
        hist[1][1][:y] *= prob_yes
    end

    prob_no = 1 - prob_yes
    object = stimulus.object
    category = stimulus.category
    histogram!(hist, pred.no, xlabel="No RT", norm=true, grid=false, color=:grey, leg=false,
        size=(300,250), subplot=2)
    hist[2][1][:y] *= prob_no
    push!(posterior_plots, hist)
end
posterior_plot = plot(posterior_plots..., layout=(2,2), size=(330,300))
savefig(posteriors, "Semantic_FFT_Posterior.eps")
savefig(posterior_plot, "Semantic_FFT_Postpred.eps")