text:
#these map model types to model functions
fluxmodels = Dict{Symbol, Function}(
  :levelmle => fluxlevelmlemodel,
  :levellsq => fluxlevellsqmodel,
)

fluxresultsname(λ⁺) = "results_$(PARAM[:crspdatalabel])" *
  "_$(PARAM[:fluxopt])" *
  "_$(PARAM[:fluxmodel])" *
  "_$(PARAM[:fluxpredictiontype])" *
  "_$(PARAM[:fluxresultssuffix])" *
  "_$λ⁺"

#primary function for estimating the measure using flux
function formfluxregression(panel::DataFrame, ms::MeasureSpec;
  fluxmodel::Function = fluxmodels[PARAM[:fluxmodel]])

  issorted(panel, [:permno, :date]) || error("panel must be sorted")

  zs::ZSpec = condition!(panel, ms)

  #testing
  if PARAM[:testmeasure]
    testvolumepartsindiceslsq()
    testvolumepartsindicesmle()
  end

  if PARAM[:refreshmeasure]
    λ⁺, results::DataFrame = fluxmodel(panel, zs)
    destandardize!(results, zs)

    #this will be the name of the results file
    resultname::String = fluxresultsname(λ⁺)

    results |> CSV.write("$(PARAM[:outputpath])\\$(resultname).csv")

  end

end

#pre-run on optimimization to find optimal learning rate
function findηmax(optgen, ::Type{T}=Float32;
  loss::Function = error("loss function is required"),
  Θ::AbstractVolumeParts = error("original VolumeParts model is required"),
  X::NamedTuple= error("X is required"),
  y = error("y is required"),
  ηsearchmax::T = T(PARAM[:fluxsearchmax]),
  ηstart::T = T(PARAM[:fluxsearchstart]),
  Δη::T = T(PARAM[:fluxsearchdelta]),
  searchiter::Int=PARAM[:fluxsearchiter],
  reportinterval::NInt = nothing) where T<:Real

  Θ₀ = copytrainable(Θ)  #the best model
  Π = Flux.params(Θ) #formatted parameters

  #record η here
  ηmax::T=ηstart
  λmax::T = -T(Inf)
  η::T=ηstart
  λmin::T = T(Inf)

  #set up an iterator so we can use the built-in training loop
  repdat = Iterators.repeated((X,y), searchiter)

  #callback for reporting
  reportcb() = println("η: $η\tηmax: $ηmax\t ηsupermax: $ηsearchmax λmin: $λmin")
  throttledreportcb = reportinterval === nothing ? reportcb : throttle(reportcb, reportinterval)

  @info "Beginning search for ηmax"
  while η ≤ ηsearchmax
    opt = optgen(η)
    Flux.train!(loss, Π, repdat, opt)
      #cb=throttle(()->println("Searching...loss=", loss(ws, RLws,V)), reportinterval))
    λ::T = loss(X,y)

    if λ < λmin
      λmin = λ
      ηmax = η
    end

    if λ > λmax
      λmax = λ
    end
    throttledreportcb()

    if λ === T(Inf)
      @warn "infinite loss. Assuming ηmax ($ηmax) < η($η) and breaking."
      break
    end

    #reset and repeat
    updatexfromy!(Θ, Θ₀)
    η *= exp(Δη)
  end

  updatexfromy!(Θ, Θ₀)
  ((λmax - λmin)/(abs(λmin) + abs(λmax)) ≤ 0.01) && error("
    λmin ($λmin) nearly equal to λmax ($λmax), consider expanding search bounds")
  (λmin === T(Inf)) && error("no ηmax found!!")

  @info "ηmax search complete."
  reportcb()
  return ηmax
end

function solvevol!(optgen, ::Type{T}=Float32;
    loss::Function = error("loss function is required"),
    Θ::AbstractVolumeParts = error("original vol parts model is required"),
    X::NamedTuple = error("X is required"),
    y = error("y is required"),
    debugcb::Function = ()->nothing,
    maxiter = PARAM[:fluxmaxiter],
    maxiternoimprovement = PARAM[:fluxmaxiternoimprovement],
    reportinterval = PARAM[:fluxreportinterval],
    ηsupermindenom = T(PARAM[:fluxlearningsupermindenom]),
    Δη⁺ = T(PARAM[:fluxlearningdeltaup]),
    Δη⁻ = T(PARAM[:fluxlearningdeltadown]),
    cyclemaxdenom = T(PARAM[:fluxcyclemaxdenom]),
    cyclemindenom = T(PARAM[:fluxcyclemindenom]),
    cycledeltadenom = T(PARAM[:fluxcycledeltadenom]),
    cycles = PARAM[:fluxcycles],
    ∇rmstol = PARAM[:fluxgradrmstol],
    stopnoimprovement = PARAM[:fluxstopnoimprovement],
    stopsmallgradient = PARAM[:fluxstopsmallgradient]
  ) where T<:Real

  #get the data
  Π = Flux.params(Θ) #formatted parameters
  Θ⁺ = copytrainable(Θ)  #the best model
  Θtemp = copytrainable(Θ)  #a temporary storage location

  PARAM[:fluxopt] == :adadelta && error("udpate not supported with :adadelta")

  #initialize the current state
  local ∞::T = T(Inf)
  local λ ::T= ∞
  local λ⁺::T = ∞ #the best loss

  local itersincebest::Int = 0
  local iter::Int = 0

  #cycle parameters
  local ηmax::T = findηmax(optgen, T, loss=loss, Θ=Θ, X=X, y=y, reportinterval=reportinterval)
  ηmax = ηmax / cyclemaxdenom
  local ηmin::T = ηmax / cyclemindenom

  local ηsupermin::T = ηmin / ηsupermindenom
  (ηmax ≤ ηmin) && error("ηmax $ηmax ≤ ηmin $ηmin")
  local Δcycle::T = ηmin / cycledeltadenom

  #learning parameters
  η::T = ηmin
  local η⁺::T = η #learnign rate at best
  local ηupmult::T = exp(Δη⁺) #multiply η by this when we decrease the learning rate
  local ηdownmult::T = exp(-Δη⁻) #multiply η by this when we reduce the learning rate
  local isbestλ::Bool

  #this will yield an updated optimizer given a learning rate
  #updateopt(η) = (η, optgen(η))
  @inline function updateopt!(η, opt)
    opt.eta = η
  end
  updateopt(η) = (η, optgen(η))

  #initialize the optimizer
  opt = optgen(η)
  opt⁺ = deepcopy(opt)

  #main gradient function
  get∇Θ() = gradient(Π) do
      λ = loss(X,y) #store the loss so that we don't have to calculate it again

      if λ < λ⁺ #check if this is the best we have
        λ⁺ = λ
        η⁺ = η
        itersincebest = 0
        isbestλ = true
      else
        itersincebest += 1
        isbestλ=false
      end

      return λ
    end

  #Report progress to the user
  function reportcb(∇Θ)

    #get the gradients for monitoring
    rmsnow = ∇rms(Θ, ∇Θ)
    updatexfromy!(Θtemp, Θ) #temp=Θ
    updatexfromy!(Θ, Θ⁺) #Θ=Θ⁺
    ∇Θ = get∇Θ()
    rmsbest = ∇rms(Θ, ∇Θ)
    updatexfromy!(Θ, Θtemp) #Θ=temp

    println("loss: $λ\tbestloss: $(λ⁺)\tη:$η\tcurrent∇rms: $(rmsnow)\tbest∇rms: $(rmsbest)\t",
      iter==0 ? "" : "iter: $iter/$maxiter\titersincebest: $itersincebest/$maxiternoimprovement")

  end
  throttledreportcb = reportinterval === nothing ? reportcb : throttle(reportcb, reportinterval)

  #three exit conditions: 1) maxiter 2) no improvement for some time 3) an error
  @info "beginning cycle conditioning. ηmax=$ηmax, ηmin=$ηmin, η=$η"
  @progress for cycle ∈ 0:cycles
    @info "beginning cycle $cycle of $(cycles). Est. $((ηmax-ηmin)/Δcycle*2.0) iter/cycle"
    (cycle==cycles) && @info "YAY! ITS THE FINAL CYCLE"
    while η ≤ ηmax
      ∇Θ = get∇Θ()
      (λ==∞) && debugcb() #infinity means something went wrong
      Flux.update!(opt, Π, ∇Θ)
      if isbestλ
        updatexfromy!(Θ⁺, Θ)
        opt₊ = deepcopy(opt)
      end
      η, opt = updateopt(η + Δcycle)
      #η = updateopt!(η + Δcycle, opt)
      throttledreportcb(∇Θ) #keep user appraised
    end
    while η ≥ ηmin
      ∇Θ = get∇Θ()
      (λ==∞) && debugcb() #infinity means something went wrong
      Flux.update!(opt, Π, ∇Θ)
      if isbestλ
        updatexfromy!(Θ⁺, Θ)
        opt₊ = deepcopy(opt)
      end
      η, opt = updateopt(η - Δcycle)
      #η = updateopt!(η - Δcycle, opt)
      throttledreportcb(∇Θ) #keep user appraised
    end
  end

  @info "Final cycle complete. Beginning final convex optimization from best point"
  updatexfromy!(Θ, Θ⁺)
  itersincebest = 0
  isbestλ = false

  opt = opt⁺
  @progress for i ∈ 1:maxiter
    iter = i #for reporting purposes
    ∇Θ = get∇Θ()

    isbestλ && updatexfromy!(Θ⁺, Θ)
    if λ==∞ #reset and lower learning rate
      debugcb()  #this is an error scenario where learning rate can't decrease
    elseif isbestλ && (η ≤ ηmax) #if its the best answer, note model and accelerate learning
      #η, opt = updateopt(η * ηupmult)
      η = updateopt!(η * ηupmult, opt)
    elseif (!isbestλ) && η ≥ ηsupermin #if its not, decellerate learning
      #η, opt = updateopt(η * ηdownmult)
      η = updateopt!(η * ηdownmult, opt)
    end

    #run through the stopping criteria
    if stopnoimprovement && (itersincebest ≥ maxiternoimprovement) #this is the main stopping criteria
      @info "No improvement in $maxiternoimprovement iterations."
      reportcb(∇Θ) #run the callback once more
      break
    elseif stopsmallgradient && ∇rms(Θ, ∇Θ) < ∇rmstol
      @info "GOOD NEWS: ∇rms ($(∇rms(Θ, ∇Θ))) < tol ($∇rmstol)."
      reportcb(∇Θ) #run the callback once more
      break
    elseif iter == maxiter
      @warn "max iter reached, returning best result"
      reportcb(∇Θ)
    else
      throttledreportcb(∇Θ) #keep user appraised
      Flux.update!(opt, Π, ∇Θ)
    end
  end

  ∇Θ = get∇Θ()
  printτ = :τ ∈ propertynames(Θ) #check if τ (risk or precision) is used in the optimization process
  open("$(PARAM[:outputpath])\\$(fluxresultsname(λ⁺))_info.txt", "w+") do f
    println(f, "******Loss: loss: $λ\tbestloss: $(λ⁺)\tη:$η\tgradrms: $(∇rms(Θ, ∇Θ))",
    "\n******Current A:\n", ∇Θ[Θ.A]',
    "\n******Current V₀:\n", ∇Θ[Θ.V₀],
    (printτ ? "\n******Current τ:\n$(display(∇Θ[Θ.τ]'))" : ""),
    "\n++++++Switch to best: ", isnothing(updatexfromy!(Θ,Θ⁺)),
    "\n++++++Best A:\n", ∇Θ[Θ.A]',
    "\n++++++Best V₀:\n", ∇Θ[Θ.V₀],
    (printτ ? "\n++++++Best τ:\n$(display(∇Θ[Θ.τ]'))" : ""),
    )
  end
  return λ⁺
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
# AUTO GENERATED FILE - DO NOT EDIT

export viasp_viaspdash

"""
    viasp_viaspdash(;kwargs...)

A ViaspDash component.

Keyword arguments:
- `id` (String; optional): The ID used to identify this component in Dash callbacks.
"""
function viasp_viaspdash(; kwargs...)
        available_props = Symbol[:id]
        wild_props = Symbol[]
        return Component("viasp_viaspdash", "ViaspDash", "viasp_dash", available_props, wild_props; kwargs...)
end

ext: jl

max_stars_repo_name: glaserL/viasp

max_issues_repo_name: glaserL/viasp

max_forks_repo_name: glaserL/viasp


text:
# This file was generated, do not modify it. # hide
auto = dataset("ISLR", "Auto")
first(auto, 3)


text:
include("predict.jl")
include("assume.jl")
include("observe.jl")
include("noparam.jl")
include("beta-binomial.jl")
include("resample.jl")
include("importance_sampling.jl")

include("test_clonetask.jl")
include("test_tarray.jl")

# include("test_tarray2.jl")
include("test_particlecontainer.jl")


text:
@nospecialize
# TODO: make show code a bit more DRY

struct TransducerFolder
    xform::Transducer
end

function __foldl__(rf, val, xff::TransducerFolder)
    @nospecialize
    xf = _normalize(xff.xform)
    while xf isa Composition
        val = @next(rf, val, xf.outer)
        xf = xf.inner
    end
    val = @next(rf, val, xf)
    return complete(rf, val)
end

function print_arrow(io, pre, post)
    print(io, pre)
    printstyled(io, "|>"; color=:light_black, bold=true)
    print(io, post)
end

print_spaced_arrow(io) = print_arrow(io, ' ', ' ')

function Base.summary(io::IO, xform::Transducer)
    xff = TransducerFolder(xform)
    n = mapfoldl(Count(), right, xff)
    f = mapfoldl(Take(1), right, xff)
    l = mapfoldl(TakeLast(1), right, xff)
    @assert n > 0
    if n == 1
        print(io, nameof(typeof(xform)))
    else
        print(io, nameof(typeof(f)))
        print_spaced_arrow(io)
        if n > 2
            print(io, "(", n - 2, " transducers...)")
            print_spaced_arrow(io)
        end
        print(io, nameof(typeof(l)))
    end
    return
end

is_anonymous(n) = n == :anonymous || startswith(string(n), '#')

_is_default_arg(xf, name::Symbol, value) = is_default_arg(xf, Val(name), value)
is_default_arg(xf, name, value) = false
is_default_arg(xf::Scan, ::Val{:init}, value) = value === Init(xf.f)
is_default_arg(xf::GroupBy, ::Val{:init}, value) = value isa DefaultInit

_name_of_transducer_type(xf) = nameof(typeof(xf))

function show_transducer(io, mime, xf)
    @nospecialize io, xf
    printstyled(io, _name_of_transducer_type(xf); color=:cyan)
    printstyled(io, '('; color=:yellow)
    show_args(io, mime, xf)
    printstyled(io, ')'; color=:yellow)
end

function show_args(io, mime, xf)
    @nospecialize io, xf
    indent = get(io, :indent, "") * ' ' ^ 4
    multiline = mime !== nothing && any(propertynames(xf)) do n
        getproperty(xf, n) isa Composition
    end
    isfirst = true
    for n in propertynames(xf)
        v = getproperty(xf, n)
        _is_default_arg(xf, n, v) && continue
        if isfirst
            isfirst = false
            if multiline
                println(io)
                print(io, indent)
            end
        else
            if multiline
                println(io, ",")
                print(io, indent)
            else
                print(io, ", ")
            end
        end
        if v isa Function
            n = nameof(v)
            if is_anonymous(n)
                print(io, parentmodule(v), ".λ❓")
            else
                print(io, n)
            end
        else
            if multiline
                io′ = io
            else
                io′ = IOContext(io, :compact => true)
            end
            io′ = IOContext(io′, :indent => indent, :first_indent => "")
            _show_field(io′, mime, v)
        end
    end
    if multiline
        println(io)
        print(io, get(io, :indent, ""))
    end
end

_show_field(io, ::Nothing, v) = show(io, v)
_show_field(io, mime, v) = show(io, mime, v)
_show_field(io, ::MIME"text/plain", v::Array) = show(io, v)

show_args(io, mime, xf::Union{OfType{T}, NotA{T}}) where T = show(io, T)

function _show_impl(io, mime, xform::Transducer)
    @nospecialize io, xform
    indent = get(io, :indent, "")
    first_indent = get(io, :first_indent, indent)
    next_indent = indent * ' ' ^ 4
    foreach(FlagFirst(), TransducerFolder(xform)) do (isfirst, xf)
        if mime === nothing
            inner_indent = ""
            if !isfirst
                print_spaced_arrow(io)
            end
        else
            if isfirst
                print(io, first_indent)
                inner_indent = indent
            else
                print_arrow(io, ' ', '\n')
                print(io, next_indent)
                print(io)
                inner_indent = next_indent
            end
        end
        show_transducer(IOContext(io,
                                  :indent => inner_indent,
                                  :first_indent => inner_indent),
                        mime,
                        xf)
    end
    return
end

Base.show(io::IO, m::MIME"text/plain", xf::Transducer) = _show_impl(io, m, xf)

function Base.show(io::IO, xf::Transducer)
    if get(io, :limit, false)
        summary(io, xf)
    else
        _show_impl(io, nothing, xf)
    end
    return
end

Base.show(io::IO, m::MIME"text/plain", rf::AbstractReduction) =
    _show_impl(io, m, rf)

function _show_impl(io, mime, x)
    indent = get(io, :indent, "")
    first_indent = get(io, :first_indent, indent)
    print(io, first_indent)
    _show_noindent(io, mime, x)
end

_show_noindent(io, mime, x) = show(io, mime, x)

function _show_noindent(io, mime, x::Function)
    m = parentmodule(x)
    if m !== Base
        show(io, mime, m)
        print(io, '.')
    end
    print(io, nameof(x))
end

function _show_impl(io, mime, rf::AbstractReduction)
    indent = get(io, :indent, "")
    first_indent = get(io, :first_indent, indent)
    next_indent = indent * ' ' ^ 4
    next_io = IOContext(io,
                        :indent => next_indent,
                        :first_indent => next_indent)

    print(io, first_indent)
    print(io, Base.typename(typeof(rf)).name)
    println(io, '(')
    _show_multiline_args(next_io, mime, rf)
    print(io, ")")
    return
end

function _show_multiline_args(io, mime, rf::Reduction)
    _show_impl(io, mime, xform(rf))
    println(io, ",")
    _show_impl(io, mime, inner(rf))
end

function _show_multiline_args(io, mime, rf::BottomRF)
    _show_impl(io, mime, inner(rf))
end

function _show_multiline_args(io, mime, rf::Splitter)
    _show_impl(io, mime, inner(rf))
end

function _show_multiline_args(io, mime, rf::Joiner)
    _show_impl(io, mime, inner(rf))
end

@specialize


text:
@doc raw"""
    y = prox_distance(M,λ,f,x [,p=2])

compute the proximal map $\operatorname{prox}_{\lambda\varphi}$ with
parameter λ of $\varphi(x) = \frac{1}{p}d_{\mathcal M}^p(f,x)$.

# Input
* `M` – a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold) $\mathcal M$
* `λ` – the prox parameter
* `f` – a point $f ∈ \mathcal M$ (the data)
* `x` – the argument of the proximal map

# Optional argument
* `p` – (`2`) exponent of the distance.

# Ouput
* `y` – the result of the proximal map of $\varphi$
"""
function prox_distance(M::Manifold, λ, f, x, p::Int=2)
  d = distance(M,f,x)
  if p==2
    t =  λ/(1+λ);
  elseif p==1
    if λ < d
      t = λ/d;
    else
      t = 1.;
    end
  else
      throw(ErrorException(
        "Proximal Map of distance(M,f,x) not implemented for p=$(p) (requires p=1 or 2)"
      ))
  end
  return exp(M,x,log(M,x,f),t);
end
@doc raw"""
    (y1,y2) = prox_TV(M,λ,(x1,x2) [,p=1])

Compute the proximal map $\operatorname{prox}_{\lambda\varphi}$ of
$\varphi(x,y) = d_{\mathcal M}^p(x,y)$ with
parameter `λ`.

# Input
* `M` – a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold)
* `λ` – a real value, parameter of the proximal map
* `(x1,x2)` – a tuple of two points,

# Optional
(default is given in brackets)
* `p` – (1) exponent of the distance of the TV term

# Ouput
* `(y1,y2)` – resulting tuple of points of the
  $\operatorname{prox}_{\lambda\varphi}($ `(x1,x2)` $)$
"""
function prox_TV(M::mT,λ::Number, pointTuple::Tuple{T,T}, p::Int=1) where {mT <: Manifold,T}
  x1 = pointTuple[1];
  x2 = pointTuple[2];
  d = distance(M,x1,x2);
  if p==1
    t = min(0.5, λ/d);
  elseif p==2
    t = λ/(1+2*λ);
  else
    throw(ErrorException(
      "Proximal Map of TV(M,x1,x2,p) not implemented for p=$(p) (requires p=1 or 2)"
    ))
  end
  return (  exp(M, x1, log(M, x1, x2), t), exp(M, x2, log(M, x2, x1), t)  );
end
@doc raw"""
    ξ = prox_TV(M,λ,x [,p=1])

compute the proximal maps $\operatorname{prox}_{\lambda\varphi}$ of
all forward differences orrucirng in the power manifold array, i.e.
$\varphi(xi,xj) = d_{\mathcal M}^p(xi,xj)$ with `xi` and `xj` are array
elemets of `x` and `j = i+e_k`, where `e_k` is the $k$th unitvector.
The parameter `λ` is the prox parameter.

# Input
* `M` – a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold)
* `λ` – a real value, parameter of the proximal map
* `x` – a point.

# Optional
(default is given in brackets)
* `p` – (1) exponent of the distance of the TV term

# Ouput
* `y` – resulting  point containinf with all mentioned proximal
  points evaluated (in a cylic order).
"""
function prox_TV(M::PowerManifold{𝔽,N,T}, λ, x, p::Int=1) where {𝔽,N<:Manifold,T}
    power_size = power_dimensions(M)
    R = CartesianIndices(Tuple(power_size))
    d = length(power_size)
    maxInd = Tuple(last(R))
    y = copy(x)
    for k in 1:d # for all directions
        ek = CartesianIndex(ntuple(i  ->  (i==k) ? 1 : 0, d) ) #k th unit vector
        for l in 0:1
            for i in R # iterate over all pixel
                if (i[k] % 2) == l
                    I = [i.I...] # array of index
                    J = I .+ 1 .* (1:d .== k) #i + e_k is j
                    if all( J .<= maxInd ) # is this neighbor in range?
                        j = CartesianIndex(J...) # neigbbor index as Cartesian Index
                        (y[i],y[j]) = prox_TV( M.manifold,λ,(y[i],y[j]),p) # Compute TV on these
                    end
                end
            end # i in R
        end # even odd
    end # directions
    return y
end
@doc raw"""
    ξ = prox_parallel_TV(M,λ,x [,p=1])

compute the proximal maps $\operatorname{prox}_{\lambda\varphi}$ of
all forward differences orrucirng in the power manifold array, i.e.
$\varphi(xi,xj) = d_{\mathcal M}^p(xi,xj)$ with `xi` and `xj` are array
elemets of `x` and `j = i+e_k`, where `e_k` is the $k$th unitvector.
The parameter `λ` is the prox parameter.

# Input
* `M`     – a `PowerManifold` manifold
* `λ`     – a real value, parameter of the proximal map
* `x`     – a point

# Optional
(default is given in brackets)
* `p` – (`1`) exponent of the distance of the TV term

# Ouput
* `y`  – resulting Array of points with all mentioned proximal
  points evaluated (in a parallel within the arrays elements).

*See also* [`prox_TV`](@ref)
"""
function prox_parallel_TV(M::PowerManifold, λ, x::Array{T,1}, p::Int=1) where {T}
  R = CartesianIndices(x[1])
  d = ndims(x[1])
  if length(x) != 2*d
    throw(ErrorException("The number of inputs from the array ($(length(x))) has to be twice the data dimensions ($(d))."))
  end
  maxInd = Tuple(last(R))
  # create an array for even/odd splitted proxes along every dimension
  y = reshape(deepcopy(x),d,2)
  x = reshape(x,d,2)
  for k in 1:d # for all directions
    ek = CartesianIndex(ntuple(i  ->  (i==k) ? 1 : 0, d) ) #k th unit vector
    for l in 0:1 # even odd
      for i in R # iterate over all pixel
        if (i[k] % 2) == l
          I = [i.I...] # array of index
          J = I .+ 1 .* (1:d .== k) #i + e_k is j
          if all( J .<= maxInd ) # is this neighbor in range?
            j = CartesianIndex(J...) # neigbbor index as Cartesian Index
            # parallel means we apply each (direction even/odd) to a seperate copy of the data.
            (y[k,l+1][i],y[k,l+1][j]) = prox_TV( M.manifold,λ,(x[k,l+1][i],x[k,l+1][j]),p) # Compute TV on these
          end
        end
      end # i in R
    end # even odd
  end # directions
  return y[:] # return as onedimensional array
end
@doc raw"""
    (y1,y2,y3) = prox_TV2(M,λ,(x1,x2,x3),[p=1], kwargs...)

Compute the proximal map $\operatorname{prox}_{\lambda\varphi}$ of
$\varphi(x_1,x_2,x_3) = d_{\mathcal M}^p(c(x_1,x_3),x_2)$ with
parameter `λ`>0, where $c(x,z)$ denotes the mid point of a shortest
geodesic from `x1` to `x3` that is closest to `x2`.

# Input

* `M`          – a manifold
* `λ`          – a real value, parameter of the proximal map
* `(x1,x2,x3)` – a tuple of three points

* `p` – (`1`) exponent of the distance of the TV term

# Optional
`kwargs...` – parameters for the internal [`subgradient_method`](@ref)
    (if `M` is neither `Euclidean` nor `Circle`, since for these a closed form
    is given)

# Output
* `(y1,y2,y3)` – resulting tuple of points of the proximal map
"""
function prox_TV2(M::Manifold,λ,pointTuple::Tuple{T,T,T},p::Int=1;
  stopping_criterion::StoppingCriterion = StopAfterIteration(5),
  kwargs...) where {T}
  if p != 1
    throw(ErrorException(
      "Proximal Map of TV2(M,λ,pT,p) not implemented for p=$(p) (requires p=1) on general manifolds."
    ))
  end
  PowX = [pointTuple...]
  PowM = PowerManifold(M, NestedPowerRepresentation(), 3)
  xInit = PowX
  F(x) = 1/2*distance(PowM,PowX,x)^2 + λ*costTV2(PowM,x)
  ∂F(x) = log(PowM,x,PowX) + λ*∇TV2(PowM,x)
  xR = subgradient_method(PowM,F,∂F,xInit;stopping_criterion=stopping_criterion, kwargs...)
  return (xR...,)
end
function prox_TV2(M::Circle,λ,pointTuple::Tuple{T,T,T},p::Int=1) where {T}
  w = [1., -2. ,1. ]
  x = [pointTuple...]
  if p==1 # Theorem 3.5 in Bergmann, Laus, Steidl, Weinmann, 2014.
    m = min( λ, abs(  sym_rem( sum( x .* w  ) ) )/(dot(w,w))   )
    s = sign( sym_rem(sum(x .* w)) )
    return Tuple(  sym_rem.( x  .-  m .* s .* w ) )
  elseif p==2 # Theorem 3.6 ibd.
    t = λ * sym_rem( sum( x .* w ) ) / (1 + λ*dot(w,w) )
    return Tuple( sym_rem.( x - t.*w )  )
  else
    throw(ErrorException(
      "Proximal Map of TV2(Circle,λ,pT,p) not implemented for p=$(p) (requires p=1 or 2)"
    ))
  end
end
function prox_TV2(M::Euclidean,λ,pointTuple::Tuple{T,T,T},p::Int=1) where {T}
  w = [1., -2. ,1. ]
  x = [pointTuple...]
  if p==1 # Example 3.2 in Bergmann, Laus, Steidl, Weinmann, 2014.
    m = min.(Ref(λ),  abs.( x .* w  ) / (dot(w,w))   )
    s = sign.( sum(x .* w) )
    return x  .-  m .* s .* w
  elseif p==2 # Theorem 3.6 ibd.
    t = λ * sum( x .* w ) / (1 + λ*dot(w,w) )
    return x - t.*w
  else
    throw(ErrorException(
      "Proximal Map of TV2(Euclidean,λ,pT,p) not implemented for p=$(p) (requires p=1 or 2)"
    ))
  end
end
@doc raw"""
    ξ = prox_TV2(M,λ,x,[p])

compute the proximal maps $\operatorname{prox}_{\lambda\varphi}$ of
all centered second order differences orrucirng in the power manifold array, i.e.
$\varphi(x_k,x_i,x_j) = d_2(x_k,x_i.x_j)$, where $k,j$ are backward and forward
neighbors (along any dimension in the array of `x`).
The parameter `λ` is the prox parameter.

# Input
* `M` – a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold)
* `λ` – a real value, parameter of the proximal map
* `x` – a points.

# Optional
(default is given in brackets)
* `p` – (`1`) exponent of the distance of the TV term

# Ouput
* `y` – resulting point with all mentioned proximal points
  evaluated (in a cylic order).
"""
function prox_TV2(M::PowerManifold{N,T}, λ, x, p::Int=1) where {N,T}
  power_size = power_dimensions(M)
  R = CartesianIndices(power_size)
  d = length(size(x))
  minInd = [first(R).I...]
  maxInd = [last(R).I...]
  y = copy(x)
  for k in 1:d # for all directions
    for l in 0:1
      for i in R # iterate over all pixel
        if (i[k] % 3) == l
          I = [i.I...] # array of index
          JForward = I .+ 1 .* (1:d .== k) #i + e_k
          JBackward = I .+ 1 .* (1:d .== k) # i - e_k
          if all( JForward .<= maxInd ) && all( JBackward .>= minInd)
            jForward = CartesianIndex{d}(JForward...) # neigbbor index as Cartesian Index
            jBackward = CartesianIndex{d}(JForward...) # neigbbor index as Cartesian Index
            (y[jBackward], y[i], y[jForward]) =
              prox_TV2( M.manifold, λ, (y[jBackward], y[i], y[jForward]),p) # Compute TV on these
          end
        end # if mod 3
      end # i in R
    end # for mod 3
  end # directions
  return y
end
@doc raw"""
    prox_collaborative_TV(M,λ,x [,p=2,q=1])

compute the prox of the collaborative TV prox for x on the `PowerManifold`
manifold, i.e. of the function

```math
F^q(x) = \sum_{i ∈ \mathcal G}
  \Bigl( \sum_{j ∈ \mathcal I_i}
    \sum_{k=1^d} \lVert X_{i,j}\rVert_x^p\Bigr)^\frac{q/p},
```

where $\mathcal G$ is the set of indices for $x ∈ \mathcal M$ and $\mathcal I_i$
is the set of its forward neighbors.
This is adopted from the paper by Duran, Möller, Sbert, Cremers:
_Collaborative Total Variation: A General Framework for Vectorial TV Models_
(arxiv: [1508.01308](https://arxiv.org/abs/1508.01308)), where the most inner
norm is not on a manifold but on a vector space, see their Example 3 for
details.
"""
function prox_collaborative_TV(N::PowerManifold, λ, x, Ξ,p=2.,q=1.)
  # Ξ = forward_logs(M,x)
  pdims = power_dimensions(N)
  if length(pdims) == 1
    d = 1
    s = 1
    iRep = (1,)
  else
    d = pdims[end]
    s = length(pdims)-1
    if s != d
      throw( ErrorException( "the last dimension ($(d)) has to be equal to the number of the previous ones ($(s)) but its not." ))
    end
    iRep = (Integer.(ones(d))...,d)
  end
  if q==1 # Example 3 case 2
    if p==1
      normΞ = norm.(Ref(N.manifold), x, Ξ)
      return  max.(normΞ .- λ, 0.) ./ ( (normΞ .== 0) .+ normΞ )  .*  Ξ
    elseif p==2 # Example 3 case 3
      norms = sqrt.( sum( norm.(Ref(N.manifold),x,Ξ).^2, dims=d+1))
      if length(iRep) > 1
        norms = repeat(norms,inner=iRep)
      end
      # if the norm is zero add 1 to avoid division by zero, also then the
      # nominator is already (max(-λ,0) = 0) so it stays zero then
      return  max.(norms .- λ, 0.) ./ ( (norms .== 0) .+ norms )  .*  Ξ
    else
      throw( ErrorException("The case p=$p, q=$q is not yet implemented"))
    end
  elseif q==Inf
    if p==2
      norms = sqrt.( sum( norm.(Ref(N.manifold),x,Ξ).^2, dims=d+1))
      if length(iRep) > 1
        norms = repeat(norms,inner=iRep)
      end
    elseif p==1
      norms = sum( norm.(Ref(N.manifold), x, Ξ), dims=d+1)
      if length(iRep) > 1
        norms = repeat(norms,inner=iRep)
      end
    elseif p==Inf
      norms = norm.(Ref(N.manifold),x,Ξ)
    else
      throw( ErrorException("The case p=$p, q=$q is not yet implemented"))
    end
    return (λ .* Ξ) ./ max.(Ref(λ), norms)
  end # end q
  throw( ErrorException("The case p=$p, q=$q is not yet implemented"))
end
prox_collaborative_TV(N::PowerManifold, λ, x, Ξ, p::Int, q::Float64=1.) = prox_collaborative_TV(N,λ,x,Ξ,Float64(p),q)
prox_collaborative_TV(N::PowerManifold, λ, x, Ξ, p::Float64, q::Int) = prox_collaborative_TV(N,λ,x,Ξ,p,Float64(q))
prox_collaborative_TV(N::PowerManifold, λ, x, Ξ, p::Int, q::Int) = prox_collaborative_TV(N,λ,x,Ξ,Float64(p),Float64(q))


text:
struct GenericDistribution{T<:Real} <: Prior
    pivots::Vector{T}
    weights::Vector{T}
    pdf::Vector{T}
    pdf_function::Function
    low::T
    high::T
end

function GenericDistribution(f::Function; low::Real = 0, high::Real = 1)::GenericDistribution{Float64}

    eps          = 1e-3
    low, high    = max(low, eps), min(high, 1 - eps)
    p, ω         = gauss_legendre_25(low, high)
    z            = quadgk( p -> f(p), low, high; atol = 1e-5 )[1]
    pdf_function = p -> ( (p < low) | (p > high) ) ? 0.0 : f(p) / z
    pdf          = pdf_function.(p)
    zz           = (high - low)/2 * sum(pdf .* ω)
    pdf          = pdf ./ zz # normalize
    return GenericDistribution{Float64}(
        convert(Vector{Float64}, p),
        convert(Vector{Float64}, ω),
        convert(Vector{Float64}, pdf),
        pdf_function,
        convert(Float64, low),
        convert(Float64, high)
    )
end

function pdf(p::TR, prior::GenericDistribution{TR})::TR where {TR<:Real}

    return prior.pdf_function(p)
end

function cdf(p::TR, prior::GenericDistribution{TR})::TR where {TR<:Real}

    return max(0.0, min(1.0, quadgk(prior.pdf_function, 0, p; atol = 1e-5)[1]))
end

function expectation(f::Function, prior::GenericDistribution{TR})::TR where
    {TR<:Real}

    return (prior.high - prior.low)/2 * sum( prior.pdf .* f.(prior.pivots) .* prior.weights )
end

function mean(prior::GenericDistribution{TR})::TR where
    {TR<:Real}

    return expectation(p -> p, prior)
end

function condition(prior::GenericDistribution{TR}; low::TR = prior.low, high::TR = prior.high)::GenericDistribution{TR} where
    {TR<:Real}

    return GenericDistribution{TR}(TD; low = max(low, prior.low), high = min(high, prior.high))
end

function update(prior::GenericDistribution{TR}, x::TI, n::TI)::GenericDistribution{TR} where
    {TR<:Real,TI<:Integer}

    return GenericDistribution{TR}( p -> pdf(prior, p) * pmf(x, n, p); low = prior.low, high = prior.high )
end

function string(prior::GenericDistribution)
    eps          = 1e-3
    low, high    = max(prior.low, eps), min(prior.high, 1 - eps)
    if (prior.low > low) | (prior.high < high)
        @sprintf "GenericDistribution|[%.2f,%.2f](%s)" prior.low prior.high prior.design
    else
        "GenericDistribution|[0, 1]"
    end
end


text:
import LinearAlgebra

mutable struct CLArray{T, N} <: CLObject
    ctx::Context
    queue::CmdQueue
    buffer::Buffer{T}
    size::NTuple{N, Int}
end

const CLMatrix{T} = CLArray{T,2}
const CLVector{T} = CLArray{T,1}

## constructors

function CLArray(buf::Buffer{T}, queue::CmdQueue, sz::Tuple{Vararg{Int}}) where T
    ctx = context(buf)
    CLArray(ctx, queue, buf, sz)
end

function CLArray(queue::CmdQueue,
                 flags::Tuple{Vararg{Symbol}},
                 hostarray::AbstractArray{T,N}) where {T, N}
    ctx = context(queue)
    buf = Buffer(T, ctx, flags, hostbuf=hostarray)
    sz = size(hostarray)
    CLArray(ctx, queue, buf, sz)
end

CLArray(queue::CmdQueue, hostarray::AbstractArray{T,N};
        flags=(:rw, :copy)) where {T, N} = CLArray(queue, (:rw, :copy), hostarray)

Base.copy(A::CLArray; ctx=A.ctx, queue=A.queue,
          buffer=A.buffer, size=A.size) = CLArray(ctx, queue, buffer, size)

function Base.deepcopy(A::CLArray{T,N}) where {T, N}
    new_buf = Buffer(T, A.ctx, prod(A.size))
    copy!(A.queue, new_buf, A.buffer)
    return CLArray(A.ctx, A.queue, new_buf, A.size)
end

"""
Create in device memory array of type `t` and size `dims` filled by value `x`.
"""
function Base.fill(::Type{T}, q::CmdQueue, x::T, dims...) where T
    ctx = info(q, :context)
    v = opencl_version(ctx)
    if v.major == 1 && v.minor >= 2
        buf = Buffer(T, ctx, prod(dims))
        fill!(q, buf, x)
    else
        buf = Buffer(T, ctx, (:rw, :copy), prod(dims), hostbuf=fill(x, dims))
    end
    return CLArray(buf, q, dims)
end

Base.zeros(::Type{T}, q::CmdQueue, dims...) where {T} = fill(T, q, T(0), dims...)
Base.zeros(q::CmdQueue, dims...) = fill(Float64, q, Float64(0), dims...)
Base.ones(::Type{T}, q::CmdQueue, dims...) where {T} = fill(T, q, T(1), dims...)
Base.ones(q::CmdQueue, dims...) = fill(Float64, q, Float64(1), dims...)

## core functions

buffer(A::CLArray) = A.buffer
Base.pointer(A::CLArray) = A.buffer.id
context(A::CLArray) = context(A.buffer)
queue(A::CLArray) = A.queue
Base.eltype(A::CLArray{T, N}) where {T, N} = T
Base.size(A::CLArray) = A.size
Base.size(A::CLArray, dim::Integer) = A.size[dim]
Base.ndims(A::CLArray) = length(size(A))
Base.length(A::CLArray) = prod(size(A))
Base.:(==)(A:: CLArray, B:: CLArray) =
    buffer(A) == buffer(B) && size(A) == size(B)
Base.reshape(A::CLArray, dims...) = begin
    @assert prod(dims) == prod(size(A))
    return copy(A, size=dims)
end

## show

Base.show(io::IO, A::CLArray{T,N}) where {T, N} =
    print(io, "CLArray{$T,$N}($(buffer(A)),$(size(A)))")

## to_host

function to_host(A::CLArray{T,N}; queue=A.queue) where {T, N}
    hA = Array{T}(undef, size(A)...)
    copy!(queue, hA, buffer(A))
    return hA
end

## other array operations

const TRANSPOSE_PROGRAM_PATH = joinpath(dirname(@__FILE__), "kernels/transpose.cl")

function max_block_size(queue::CmdQueue, h::Int, w::Int)
    dev = info(queue, :device)
    dim1, dim2 = info(dev, :max_work_item_size)[1:2]
    wgsize = info(dev, :max_work_group_size)
    wglimit = floor(Int, sqrt(wgsize))
    return gcd(dim1, dim2, h, w, wglimit)
end

"""
Transpose CLMatrix A, write result to a preallicated CLMatrix B
"""
function LinearAlgebra.transpose!(B::CLMatrix{Float32}, A::CLMatrix{Float32};
                                  queue=A.queue)
    block_size = max_block_size(queue, size(A, 1), size(A, 2))
    ctx = context(A)
    kernel = get_kernel(ctx, TRANSPOSE_PROGRAM_PATH, "transpose",
                          block_size=block_size)
    h, w = size(A)
    lmem = LocalMem(Float32, block_size * (block_size + 1))
    set_args!(kernel, buffer(B), buffer(A), UInt32(h), UInt32(w), lmem)
    return enqueue_kernel(queue, kernel, (h, w), (block_size, block_size))
end

"""Transpose CLMatrix A"""
function LinearAlgebra.transpose(A::CLMatrix{Float32};
                        queue=A.queue)
    B = zeros(Float32, queue, reverse(size(A))...)
    ev = LinearAlgebra.transpose!(B, A, queue=queue)
    wait(ev)
    return B
end

"""Transpose CLMatrix A, write result to a preallicated CLMatrix B"""
function LinearAlgebra.transpose!(B::CLMatrix{Float64}, A::CLMatrix{Float64};
                         queue=A.queue)
    block_size = max_block_size(queue, size(A, 1), size(A, 2))
    ctx = context(A)
    kernel = get_kernel(ctx, TRANSPOSE_PROGRAM_PATH, "transpose_double",
                          block_size=block_size)
    h, w = size(A)
    # lmem = LocalMem(Float64, block_size * (block_size + 1))
    lmem = LocalMem(Float64, block_size * block_size)
    set_args!(kernel, buffer(B), buffer(A), UInt32(h), UInt32(w), lmem)
    return enqueue_kernel(queue, kernel, (h, w), (block_size, block_size))
end

"""Transpose CLMatrix A"""
function LinearAlgebra.transpose(A::CLMatrix{Float64};
                        queue=A.queue)
    B = zeros(Float64, queue, reverse(size(A))...)
    ev = LinearAlgebra.transpose!(B, A, queue=queue)
    wait(ev)
    return B
end


text:
@testset "Utilities" begin
  @test_throws ArgumentError readgeotable("doesnotexist.csv")

  grid = boundgrid(data1D, (100,))
  @test coordinates(grid, 1) == [0.]
  @test coordinates(grid, 100) == [100.]
end


text:
export MTIterativeSampler

mutable struct MTIterativeSampler{N,P,IC,MIC,EC,S,SC,V} <: AbstractIterativeSampler
    net::N
    problem::P
    itercache::IC
    mitercache::MIC
    sampled_values::EC

    sampler::S
    sampler_cache::SC
    v::V
end

"""
    MTIterativeSampler(net, sampler, algorithm, prob)

Creates a multithreaded iterative sampler.
"""
function MTIterativeSampler(net,
                            sampl,
                            prob,
                            algo=prob)
    sampl = multithread(sampl)

    itercache      = SamplingCache(algo, prob, net)
    mitercache     = [SamplingCache(algo, prob, net) for i=1:Threads.nthreads()]
    v              = state(prob, net)
    sampler_cache  = init_sampler!(sampl, net, v)

    evaluated_vals = EvaluatedNetwork(algo, net)

    MTIterativeSampler(net, prob, itercache, mitercache, evaluated_vals, sampl, sampler_cache, v)
end

function sample!(is::MTIterativeSampler)
    # Init the samplers on each thread
    init_sampler!(is.sampler, is.net, is.v, is.sampler_cache)

    # Sample on every thread
    Threads.@threads for i=1:Threads.nthreads()
        s = sampler_list(is.sampler)[i]
        n = is.sampler_cache.nets[i]
        v = is.sampler_cache.σs[i]
        _ic = is.mitercache[i]
        zero!(_ic)
        while true
            if !(s isa FullSumSampler)
                sample_network!(_ic, is.problem, n, v)
            else
                sample_network_wholespace!(_ic, is.problem, n, v)
            end
            !samplenext!(v, s, n, is.sampler_cache.caches[i]) && break
        end
    end

    # Add results of all threads
    zero!(is.itercache);
    for ic=is.mitercache
        NeuralQuantum.add!(is.itercache, ic)
    end
    evaluation_post_sampling!(is.sampled_values, is.itercache)
    return is.sampled_values
end

Base.show(io::IO, is::MTIterativeSampler) = print(io,
    "MTIterativeSampler for parallel sampling with $(Threads.nthreads()) threads:"*
    "\n\tnet\t\t: $(is.net)"*
    "\n\tproblem\t: $(is.problem)"*
    "\n\tsampler\t: $(is.sampler)")


text:
"""
    StatisticalSubstitutionModel

Subtype of `StatsBase.StatisticalModel`, to fit discrete data to a model
of trait substitution along a network.
See [`fitdiscrete`](@ref) to fit a trait substitution model to discrete data.
It returns an object of type `StatisticalSubstitutionModel`, to which standard
functions can be applied, like `loglikelihood(object)`, `aic(object)` etc.
"""
mutable struct StatisticalSubstitutionModel <: StatsBase.StatisticalModel
    model::SubstitutionModel
    ratemodel::RateVariationAcrossSites #allows rates to vary according to gamma
    """stationary for NASM models (log(1/4)), uniform in all other cases"""
    prioratroot::Vector{Float64}
    net::HybridNetwork
    """ data: trait[i] for leaf with n.number = i
        type Int: for indices of trait labels in getlabels(model)
        allows missing, but not half-ambiguous and not multi-state"""
    trait::Vector{Vector{Union{Missing, Int}}}
    "number of columns in the data: should be equal to `length(trait[i])` for all i"
    nsites::Int
    "vector of weights, one for each column, taken as 1s if not provided"
    siteweight::Union{Nothing, Vector{Float64}}
    "total sum of site weights"
    totalsiteweight::Float64
    """
    log of transition probabilities: where e goes from X to Y
    logtrans[i,j,e,r] = P{Y=j|X=i} where i=start_state, j=end_state, e=edge.number, r = rate (if using RateVariationAcrossSites)
    all displayed trees use same edge numbers and same logtrans as in full network
    """
    loglik::Union{Missings.Missing, Float64}
    """
    log of transition probabilities.
    size: k,k, net.numEdges, r where k=nstates(model)
    """
    logtrans::Array{Float64,4}
    # type based on extracting displayed trees
    displayedtree::Vector{HybridNetwork}
    """
    prior log tree weight: log product of γ's.
    In fit!: priorltw = `inheritanceWeight.(trees)`
    (which returns missing for any negative γ and would cause an error here)
    """
    priorltw::Vector{Float64}
    """
    partial log-likelihoods for active trait(s) given a particular displayed tree
    and a particular rate category, at indices [i, n.number or e.number]:
    - forward likelihood: log P{data below node n (in tree t) given state i at n}
    - direct likelihood:  log P{data below edge e (in tree t) given i at parent of e}
    - backward likelihood:log P{data at all non-descendants of n (in t) and state i at n}
    sizes: k, net.numNodes or net.numEdges.
    """
    # reset for each trait and rate
    forwardlik::Array{Float64,2} # size: k, net.numNodes
    directlik::Array{Float64,2}  # size: k, net.numEdges
    backwardlik::Array{Float64,2}# size: k, net.numNodes
    "log-likelihood of site k"
    _sitecache::Array{Float64,1} # size: nsites
    "log-likelihood of ith displayed tree t, given rate category j, of site k"
    _loglikcache::Array{Float64, 3} # size: nsites, nrates, ntrees

    "inner (default) constructor: from model, rate model, network, trait and site weights"
    function StatisticalSubstitutionModel(model::SubstitutionModel,
            ratemodel::RateVariationAcrossSites,
            net::HybridNetwork, trait::AbstractVector,
            siteweight=nothing::Union{Nothing, Vector{Float64}},
            maxhybrid=length(net.hybrid)::Int)
        length(trait) > 0 || error("no trait data!")
        nsites = length(trait[1])
        siteweight === nothing || length(siteweight) == nsites ||
            error("siteweight must be of same length as the number of traits")
        totalsiteweight = (siteweight === nothing ? float(nsites) : sum(siteweight))
        k = nstates(model)
        if typeof(model) <: NASM
            prioratroot = log.(stationary(model)) #refactor to save in obj
        else # other trait models
            prioratroot = [-log(k) for i in 1:k] # uniform prior at root
        end
        # T = eltype(getlabels(model))
        # extract displayed trees
        trees = displayedTrees(net, 0.0; nofuse=true)
        for tree in trees
            preorder!(tree) # no need to call directEdges! before: already done on net
        end
        ntrees = 2^maxhybrid
        ntrees >= length(trees) ||
            error("""maxhybrid is too low.
                    Call using maxhybrid >= current number of hybrids""")
        # log tree weights: sum log(γ) over edges, for each displayed tree
        priorltw = inheritanceWeight.(trees)
        all(!ismissing, priorltw) ||
          error("one or more inheritance γ's are missing or negative. fix using setGamma!(network, edge)")
        maxedges = length(net.edge) + 3*(maxhybrid-length(net.hybrid))
        maxnodes = length(net.node) + 2*(maxhybrid-length(net.hybrid))
        logtrans   = zeros(Float64, k,k, maxedges, length(ratemodel.ratemultiplier))
        forwardlik = zeros(Float64, k, maxnodes)
        directlik  = zeros(Float64, k, maxedges)
        backwardlik= zeros(Float64, k, maxnodes)
        _sitecache = Vector{Float64}(undef, nsites)
        _loglikcache = zeros(Float64, nsites, length(ratemodel.ratemultiplier), ntrees)
        new(deepcopy(model), deepcopy(ratemodel), prioratroot,
            net, trait, nsites, siteweight, totalsiteweight, missing, # missing log likelihood
            logtrans, trees,
            priorltw, forwardlik, directlik, backwardlik,_sitecache,_loglikcache)
    end
end
const SSM = StatisticalSubstitutionModel

# fasta constructor: from net, fasta filename, modsymbol, and maxhybrid
# Works for DNA in fasta format. Probably need different versions for
# different kinds of data (snp, amino acids). Similar to fitdiscrete()
function StatisticalSubstitutionModel(net::HybridNetwork, fastafile::String,
        modsymbol::Symbol, ratecategories=1::Int,
        maxhybrid=length(net.hybrid)::Int)
    for e in net.edge # check for missing or inappropriate γ values
        if e.hybrid
            e.gamma > 0.0 && continue
            setGamma!(e, (e.isMajor ? 0.6 : 0.4)) # to maintain isMajor as is
        else
            e.gamma == 1.0 || error("tree edge number $(e.number) has γ not 1.0")
        end
    end
    data, siteweights = readfastatodna(fastafile, true)
    model = defaultsubstitutionmodel(net, modsymbol, data, siteweights)
    ratemodel = RateVariationAcrossSites(1.0, ratecategories)
    dat2 = traitlabels2indices(view(data, :, 2:size(data,2)), model)
    o, net = check_matchtaxonnames!(data[:,1], dat2, net) # calls resetNodeNumbers, which calls preorder!
    trait = dat2[o]
    obj = StatisticalSubstitutionModel(model, ratemodel, net, trait, siteweights,
                           maxhybrid)
end

StatsBase.loglikelihood(obj::SSM) = obj.loglik
StatsBase.islinear(SSM) = false
StatsBase.dof(obj::SSM) = nparams(obj.model) + nparams(obj.ratemodel)
function Base.show(io::IO, obj::SSM)
    disp =  "$(typeof(obj)):\n"
    disp *= string(obj.model)
    disp *= "$(obj.nsites) traits, $(length(obj.trait)) species\n"
    if obj.ratemodel.ncat != 1
        disp *= "variable rates across sites ~ discretized gamma with\n alpha=$(obj.ratemodel.alpha)"
        disp *= "\n $(obj.ratemodel.ncat) categories"
        disp *= "\n rate multipliers: $(round.(obj.ratemodel.ratemultiplier, digits=5))\n"
    end
    disp *= "on a network with $(obj.net.numHybrids) reticulations"
    if !ismissing(obj.loglik)
        disp *= "\nlog-likelihood: $(round(obj.loglik, digits=5))"
    end
    print(io, disp)
end
# nobs: nsites * nspecies, minus any missing, but ignores correlation between species
# fixit: extend the StatsBase methods
# coef, coefnames, coeftable, confint,
# deviance (not from loglik in StatsBase), nulldeviance (from intercept only. here?)
# nullloglikelihood (from intercept only. here?)
# score (gradient of loglik)
# informationmatrix: Fisher by default, observed info matrix if expected=false
# stderror, vcov, params
# weights

"""
    fitdiscrete(net, model, tipdata)
    fitdiscrete(net, model, RateVariationAcrossSites, tipdata)
    fitdiscrete(net, model, species, traits)
    fitdiscrete(net, model, RateVariationAcrossSites, species, traits)
    fitdiscrete(net, model, dnadata, dnapatternweights)
    fitdiscrete(net, model, RateVariationAcrossSites, dnadata, dnapatternweights)
    fitdiscrete(net, modSymbol, species, traits)
    fitdiscrete(net, modSymbol, dnadata, dnapatternweights)

Calculate the maximum likelihood (ML) score of a network or tree given
one or more discrete characters at the tips. Along each edge, transitions
are modelled with a continous time Markov `model`, whose parameters are
estimated (by maximizing the likelihood). At each hybrid node,
the trait is assumed to be inherited from either of the two immediate
parents according to the parents' average genetic contributions
(inheritance γ). The model ignores incomplete lineage sorting.
The algorithm extracts all trees displayed in the network.

Data can given in one of the following:
- `tipdata`: dictionary taxon => state label, for a single trait.
- `tipdata`: data frame for a single trait, in which case the taxon names
  are to appear in column 1 or in a column named "taxon" or "species", and
  trait *labels* are to appear in column 2 or in a column named "trait".
  Here, trait labels should be as they appear in `getlabels(model)`.
- `species`: vector of strings, and `traits`: DataFrame of traits,
  with rows in the order corresponding to the order of species names.
  Again, trait labels should be as they appear in `getlabels(model)`.
  All traits are assumed to follow the same model, with same parameters.
- `dnadata`: the first part of the output of readfastatodna,
  a dataframe of BioSequence DNA sequences, with taxon in column 1 and
  a column for each site.
- `dnapatternweights`: the second part of the output of readfastatodna,
  an array of weights, one weights for each of the site columns.
  The length of the weight is equal to nsites.
  If using dnapatternweights, must provide dnadata.
- RateVariationAcrossSites: model for rate variation (optional)

Optional arguments (default):
- `optimizeQ` (true): should model rate parameters be fixed,
  or should they be optimized?
- `optimizeRVAS` (true): should the model optimize the α parameter
  for the variability of rates across sites?
- `NLoptMethod` (`:LN_COBYLA`, derivative-free) for the optimization algorithm.
  For other options, see the
  [NLopt](https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/).
- tolerance values to control when the optimization is stopped:
  `ftolRel` (1e-12), `ftolAbs` (1e-10) on the likelihood, and
  `xtolRel` (1e-10), `xtolAbs` (1e-10) on the model parameters.
- bounds for the alpha parameter of the Gamma distribution of
  rates across sites: `alphamin=0.05`, `alphamax=50`.
- `verbose` (false): if true, more information is output.

**limitation**: the root is assumed to be in each displayed tree,
unless the prior at the root is taken to be the stationary distribution.
This leads to a wrong likelihood if the root had a hybrid child edge.
fixit in the future, with new option in `deleteleaf!` and to extract displayed trees.

# examples:

```jldoctest fitDiscrete_block
julia> net = readTopology("(((A:2.0,(B:1.0)#H1:0.1::0.9):1.5,(C:0.6,#H1:1.0::0.1):1.0):0.5,D:2.0);");

julia> m1 = BinaryTraitSubstitutionModel([0.1, 0.1], ["lo", "hi"]);

julia> using DataFrames

julia> dat = DataFrame(species=["C","A","B","D"], trait=["hi","lo","lo","hi"]);

julia> fit1 = fitdiscrete(net, m1, dat)
PhyloNetworks.StatisticalSubstitutionModel:
Binary Trait Substitution Model:
rate lo→hi α=0.27222
rate hi→lo β=0.34981
1 traits, 4 species
on a network with 1 reticulations
log-likelihood: -2.7277

julia> tips = Dict("A" => "lo", "B" => "lo", "C" => "hi", "D" => "hi");

julia> fit2 = fitdiscrete(net, m1, tips; xtolRel=1e-16, xtolAbs=1e-16, ftolRel=1e-16)
PhyloNetworks.StatisticalSubstitutionModel:
Binary Trait Substitution Model:
rate lo→hi α=0.27222
rate hi→lo β=0.34981
1 traits, 4 species
on a network with 1 reticulations
log-likelihood: -2.7277
```

Note that a copy of the network is stored in the fitted object,
but the internal representation of the network may be different in
`fit1.net` and in the original network `net`:

```jldoctest fitDiscrete_block
julia> net = readTopology("(sp1:3.0,(sp2:2.0,(sp3:1.0,sp4:1.0):1.0):1.0);");

julia> using BioSymbols

julia> tips = Dict("sp1" => BioSymbols.DNA_A, "sp2" => BioSymbols.DNA_A, "sp3" => BioSymbols.DNA_G, "sp4" => BioSymbols.DNA_G);

julia> mJC69 = JC69([0.25], false);

julia> fitJC69 = fitdiscrete(net, mJC69, tips)
PhyloNetworks.StatisticalSubstitutionModel:
Jukes and Cantor 69 Substitution Model,
absolute rate version
off-diagonal rates equal to 0.29233/3.
rate matrix Q:
               A       C       G       T
       A       *  0.0974  0.0974  0.0974
       C  0.0974       *  0.0974  0.0974
       G  0.0974  0.0974       *  0.0974
       T  0.0974  0.0974  0.0974       *
1 traits, 4 species
on a network with 0 reticulations
log-likelihood: -4.99274

julia> rv = RateVariationAcrossSites()
Rate Variation Across Sites using Discretized Gamma Model
alpha: 1.0
categories for Gamma discretization: 4
ratemultiplier: [0.14578, 0.51313, 1.07083, 2.27025]

julia> fitdiscrete(net, mJC69, rv, tips; optimizeQ=false, optimizeRVAS=false)
PhyloNetworks.StatisticalSubstitutionModel:
Jukes and Cantor 69 Substitution Model,
absolute rate version
off-diagonal rates equal to 0.25/3.
rate matrix Q:
               A       C       G       T
       A       *  0.0833  0.0833  0.0833
       C  0.0833       *  0.0833  0.0833
       G  0.0833  0.0833       *  0.0833
       T  0.0833  0.0833  0.0833       *
1 traits, 4 species
variable rates across sites ~ discretized gamma with
 alpha=1.0
 4 categories
 rate multipliers: [0.14578, 0.51313, 1.07083, 2.27025]
on a network with 0 reticulations
log-likelihood: -5.2568

```
fixit: add option to allow users to specify root prior,
using either equal frequencies or stationary frequencies for trait models.
"""
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    tips::Dict; kwargs...) #tips::Dict no ratemodel version
    ratemodel = RateVariationAcrossSites(1.0, 1)
    fitdiscrete(net, model, ratemodel, tips; kwargs...)
end

#tips::Dict version with ratemodel
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel, ratemodel::RateVariationAcrossSites,
    tips::Dict; kwargs...)
        species = String[]
    dat = Vector{Int}[] # indices of trait labels
    for (k,v) in tips
        !ismissing(v) || continue
        push!(species, k)
        vi = findfirst(isequal(v), getlabels(model))
        vi !== nothing || error("trait $v not found in model")
        push!(dat, [vi])
    end
    o, net = check_matchtaxonnames!(species, dat, net) # dat[o] would make a shallow copy only
    StatsBase.fit(StatisticalSubstitutionModel, net, model, ratemodel, view(dat, o); kwargs...)
end

#dat::DataFrame, no rate model version
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    dat::DataFrame; kwargs...)
    ratemodel = RateVariationAcrossSites(1.0, 1)
    fitdiscrete(net, model, ratemodel, dat; kwargs...)
end

#dat::DataFrame with rate model version
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    ratemodel::RateVariationAcrossSites, dat::DataFrame; kwargs...)
    i = findfirst(isequal(:taxon), DataFrames.names(dat))
    if i===nothing i = findfirst(isequal(:species), DataFrames.names(dat)); end
    if i===nothing i=1; end # first column if no column "taxon" or "species"
    j = findfirst(isequal(:trait), DataFrames.names(dat))
    if j===nothing j=2; end
    if i==j
        error("""expecting taxon names in column 'taxon', or 'species' or
        column 1, and trait values in column 'trait' or column 2.""")
    end
    species = dat[:,i]    # modified in place later
    dat = traitlabels2indices(dat[!,j], model)   # vec of vec, indices
    o, net = check_matchtaxonnames!(species, dat, net)
    StatsBase.fit(StatisticalSubstitutionModel, net, model, ratemodel, view(dat, o); kwargs...)
end

#species, dat version, no ratemodel
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    species::Array{String}, dat::DataFrame; kwargs...)
    ratemodel = RateVariationAcrossSites(1.0, 1)
    fitdiscrete(net, model, ratemodel, species, dat; kwargs...)
end

#species, dat version with ratemodel
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    ratemodel::RateVariationAcrossSites, species::Array{String},
    dat::DataFrame; kwargs...)
    dat2 = traitlabels2indices(dat, model) # vec of vec, indices
    o, net = check_matchtaxonnames!(copy(species), dat2, net)
    StatsBase.fit(StatisticalSubstitutionModel, net, model, ratemodel, view(dat2, o); kwargs...)
end

#wrapper: species, dat version with model symbol
function fitdiscrete(net::HybridNetwork, modSymbol::Symbol,
    species::Array{String}, dat::DataFrame, rvSymbol=:noRV::Symbol; kwargs...)
    rate = startingrate(net)
    labels = learnlabels(modSymbol, dat)
    if modSymbol == :JC69
        model = JC69([1.0], true) # 1.0 instead of rate because relative version
    elseif modSymbol == :HKY85
        model = HKY85([1.0], # transition/transversion rate ratio
                       empiricalDNAfrequencies(dat, repeat([1.], inner=size(dat, 2))),
                       true)
    elseif modSymbol == :ERSM
        model = EqualRatesSubstitutionModel(length(labels), rate, labels);
    elseif modSymbol == :BTSM
        model = BinaryTraitSubstitutionModel([rate, rate], labels)
    elseif modSymbol == :TBTSM
        model = TwoBinaryTraitSubstitutionModel([rate, rate, rate, rate, rate, rate, rate, rate], labels)
    else
        error("model $modSymbol is unknown or not implemented yet")
    end

    if rvSymbol == :RV
        rvas = RateVariationAcrossSites(1.0, 4)
    else
        rvas = RateVariationAcrossSites(1.0, 1)
    end
    fitdiscrete(net, model, rvas, species, dat; kwargs...)
end

#dnadata with dnapatternweights version, no ratemodel
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    dnadata::DataFrame, dnapatternweights::Array{Float64}; kwargs...)
    ratemodel = RateVariationAcrossSites(1.0, 1)
    fitdiscrete(net, model, ratemodel, dnadata, dnapatternweights; kwargs...)
end

#dnadata with dnapatternweights version with ratemodel
function fitdiscrete(net::HybridNetwork, model::SubstitutionModel,
    ratemodel::RateVariationAcrossSites,dnadata::DataFrame,
    dnapatternweights::Array{Float64}; kwargs...)
    dat2 = traitlabels2indices(dnadata[!,2:end], model)
    o, net = check_matchtaxonnames!(dnadata[:,1], dat2, net)
    kwargs = (:siteweights => dnapatternweights, kwargs...)
    StatsBase.fit(StatisticalSubstitutionModel, net, model, ratemodel, view(dat2, o);
        kwargs...)
end

#wrapper for dna data
function fitdiscrete(net::HybridNetwork, modSymbol::Symbol, dnadata::DataFrame,
    dnapatternweights::Array{Float64}, rvSymbol=:noRV::Symbol; kwargs...)
    rate = startingrate(net)
    if modSymbol == :JC69
        model = JC69([1.0], true)  # 1.0 instead of rate because relative version
    elseif modSymbol == :HKY85
        model = HKY85([1.0], # transition/transversion rate ratio
                      empiricalDNAfrequencies(view(dnadata, :, 2:size(dnadata,2)), dnapatternweights),
                      true)
    elseif modSymbol == :ERSM
        model = EqualRatesSubstitutionModel(4, rate, [BioSymbols.DNA_A, BioSymbols.DNA_C, BioSymbols.DNA_G, BioSymbols.DNA_T]);
    elseif modSymbol == :BTSM
        error("Binary Trait Substitution Model supports only two trait states, but dna data has four states.")
    elseif modSymbol == :TBTSM
        error("Two Binary Trait Substitution Model does not support dna data: it supports two sets of potentially correlated two trait states.")
    else
        error("model $modSymbol is unknown or not implemented yet")
    end

    if rvSymbol == :RV
        rvas = RateVariationAcrossSites(1.0, 4)
    else
        rvas = RateVariationAcrossSites(1.0, 1)
    end
    fitdiscrete(net, model, rvas, dnadata, dnapatternweights; kwargs...)
end

"""
    fit(StatisticalSubstitutionModel, net, model, traits; kwargs...)
    fit!(StatisticalSubstitutionModel; kwargs...)

Internal function called by [`fitdiscrete`](@ref): with same key word arguments `kwargs`.
But dangerous: `traits` should be a vector of vectors as for [`fitdiscrete`](@ref)
**but** here `traits` need to contain the *indices* of trait values corresponding
to the indices in `getlabels(model)`, and species should appear in `traits` in the
order corresponding to the node numbers in `net`.
See [`traitlabels2indices`](@ref) to convert trait labels to trait indices.

**Warning**: does *not* perform checks. [`fitdiscrete`](@ref) calls this function
after doing checks, preordering nodes in the network, making sure nodes have
consecutive numbers, species are matched between data and network etc.
"""
function StatsBase.fit(::Type{SSM}, net::HybridNetwork, model::SubstitutionModel,
    ratemodel::RateVariationAcrossSites, trait::AbstractVector; kwargs...)

    sw = nothing
    if haskey(kwargs, :siteweights)
        sw = kwargs[:siteweights]
        kwargs = filter(p -> p.first != :siteweights, kwargs)
    end
    obj = StatisticalSubstitutionModel(model, ratemodel, net, trait, sw)
    fit!(obj; kwargs...)
end

function fit!(obj::SSM; optimizeQ=true::Bool, optimizeRVAS=true::Bool,
    closeoptim=false::Bool, verbose=false::Bool, maxeval=1000::Int,
    ftolRel=fRelBL::Float64, ftolAbs=fAbsBL::Float64,
    xtolRel=xRelBL::Float64, xtolAbs=xAbsBL::Float64,
    alphamin=alphaRASmin, alphamax=alphaRASmax)
    all(x -> x >= 0.0, [e.length for e in obj.net.edge]) || error("branch lengths should be >= 0")
    all(x -> x >= 0.0, [e.gamma for e in obj.net.edge]) || error("gammas should be >= 0")
    if optimizeQ && nparams(obj.model) <1
        @debug "The Q matrix for this model is fixed, so there are no rate parameters to optimize. optimizeQ will be set to false."
        optimizeQ = false
    end
    if optimizeRVAS && (nparams(obj.ratemodel) == 0)
        @debug "Rate model has one rate category, so there are no parameters to optimize. optimizeRVAS will be set to false."
        optimizeRVAS = false
    end
    if !optimizeQ && !optimizeRVAS
        discrete_corelikelihood!(obj)
        verbose && println("loglik = $(loglikelihood(obj)) under fixed parameters, no optimization")
        return obj
    end
    if optimizeQ
        function loglikfun(x::Vector{Float64}, grad::Vector{Float64}) # modifies obj
            setrates!(obj.model, x)
            res = discrete_corelikelihood!(obj)
            verbose && println("loglik: $res, model rates: $x")
            length(grad) == 0 || error("gradient not implemented")
            return res
        end
        # optimize Q under fixed RVAS parameters
        # set-up optimization object for Q
        NLoptMethod=:LN_COBYLA # no gradient # :LN_COBYLA for (non)linear constraits, :LN_BOBYQA for bound constraints
        nparQ = nparams(obj.model)
        optQ = NLopt.Opt(NLoptMethod, nparQ)
        NLopt.ftol_rel!(optQ,ftolRel) # relative criterion
        NLopt.ftol_abs!(optQ,ftolAbs) # absolute criterion
        NLopt.xtol_rel!(optQ,xtolRel)
        NLopt.xtol_abs!(optQ,xtolAbs)
        NLopt.maxeval!(optQ, maxeval) # max number of iterations
        # NLopt.maxtime!(optQ, t::Real)
        NLopt.lower_bounds!(optQ, zeros(Float64, nparQ))
        if typeof(obj.model) == HKY85 # set an upper bound on kappa values
            NLopt.upper_bounds!(optQ, fill(kappamax,nparQ))
        end
        NLopt.max_objective!(optQ, loglikfun)
        fmax, xmax, ret = NLopt.optimize(optQ, obj.model.rate)
        setrates!(obj.model, xmax)
        obj.loglik = fmax
        verbose && println("got $(round(fmax, digits=5)) at $(round.(xmax, digits=5)) after $(optQ.numevals) iterations (return code $(ret))")
    end
    if optimizeRVAS
        function loglikfunRVAS(alpha::Vector{Float64}, grad::Vector{Float64})
            setalpha!(obj.ratemodel, alpha[1])
            res = discrete_corelikelihood!(obj)
            verbose && println("loglik: $res, rate variation model shape parameter alpha: $(alpha[1])")
            length(grad) == 0 || error("gradient not implemented")
            return res
        end
        # set-up optimization object for RVAS parameter
        NLoptMethod=:LN_COBYLA # no gradient
        # :LN_COBYLA for (non)linear constraits, :LN_BOBYQA for bound constraints
        nparRVAS = nparams(obj.ratemodel)
        optRVAS = NLopt.Opt(NLoptMethod, nparRVAS)
        NLopt.ftol_rel!(optRVAS,ftolRel) # relative criterion
        NLopt.ftol_abs!(optRVAS,ftolAbs) # absolute criterion
        NLopt.xtol_rel!(optRVAS,xtolRel)
        NLopt.xtol_abs!(optRVAS,xtolAbs)
        NLopt.maxeval!(optRVAS,1000) # max number of iterations
        # NLopt.maxtime!(optRVAS, t::Real)
        NLopt.lower_bounds!(optRVAS, fill(alphamin, (nparRVAS,)) ) # for 0 as lower bound: zeros(Float64, nparRVAS)
        NLopt.upper_bounds!(optRVAS, fill(alphamax, (nparRVAS,)) ) # delete to remove upper bound
        NLopt.max_objective!(optRVAS, loglikfunRVAS)
        fmax, xmax, ret = NLopt.optimize(optRVAS, [obj.ratemodel.alpha]) # optimization here!
        setalpha!(obj.ratemodel, xmax[1])
        obj.loglik = fmax
        verbose && println("RVAS: got $(round(fmax, digits=5)) at $(round.(xmax, digits=5)) after $(optRVAS.numevals) iterations (return code $(ret))")
    end
    if optimizeQ && optimizeRVAS && closeoptim
        # optimize Q under fixed RVAS parameters: a second time
        fmax, xmax, ret = NLopt.optimize(optQ, obj.model.rate)
        setrates!(obj.model, xmax)
        obj.loglik = fmax
        verbose && println("got $(round(fmax, digits=5)) at $(round.(xmax, digits=5)) after $(optQ.numevals) iterations (return code $(ret))")
        # optimize RVAS under fixed Q: a second time
        fmax, xmax, ret = NLopt.optimize(optRVAS, [obj.ratemodel.alpha])
        setalpha!(obj.ratemodel, xmax[1])
        obj.loglik = fmax
        verbose && println("RVAS: got $(round(fmax, digits=5)) at $(round.(xmax, digits=5)) after $(optRVAS.numevals) iterations (return code $(ret))")
    end
    return obj
end

"""
    update_logtrans(obj::SSM)

Initialize and update `obj.logtrans`, the log transition probabilities
along each edge in the full network.
They are re-used for each displayed tree, which is why edges are not fused
around degree-2 nodes when extracting displayed trees.
"""
function update_logtrans(obj::SSM)
    rates = obj.ratemodel.ratemultiplier
    k = nstates(obj.model)
    Ptmp = MMatrix{k,k,Float64}(undef) # memory to be re-used
    for edge in obj.net.edge # update logtrans: same for all displayed trees, all traits
        enum = edge.number
        len = edge.length
        for i = 1:length(rates)
            obj.logtrans[:,:,enum, i] .= log.(P!(Ptmp, obj.model, len * rates[i]))
        end
    end
end

"""
    update_logtrans(obj::SSM, edge::Edge)

Update the log-transition probabilities associates to one particular `edge`
in the network.
"""
function update_logtrans(obj::SSM, edge::Edge)
    rates = obj.ratemodel.ratemultiplier
    enum = edge.number
    len = edge.length
    for i in 1:length(rates)
        pmat = view(obj.logtrans, :,:,enum,i)
        @inbounds pmat .= log.(P!(pmat, obj.model, len * rates[i]))
    end
end

"""
    discrete_corelikelihood!(obj::StatisticalSubstitutionModel;
                             whichtrait::AbstractVector{Int} = 1:obj.nsites)

Calculate the likelihood and update `obj.loglik` for discrete characters on a network,
calling [`discrete_corelikelihood_trait!`](@ref).
The algorithm extracts all displayed trees and weighs the likelihood under all these trees.
The object's partial likelihoods are updated:
- forward and direct partial likelihoods are re-used, one trait at a time,
- overall likelihoods on each displayed tree, given each rate category and for
  each given site/trait: are cached in `_loglikcache`.
"""
function discrete_corelikelihood!(obj::SSM; whichtrait::AbstractVector{Int} = 1:obj.nsites)
    # fill _loglikcache
    nr = length(obj.ratemodel.ratemultiplier)
    nt = length(obj.displayedtree)
    update_logtrans(obj)
    for t in 1:nt
        for ri in 1:nr
            for ci in whichtrait
                obj._loglikcache[ci,ri,t] = discrete_corelikelihood_trait!(obj,t,ci,ri)
                # note: -Inf expected if 2 tips have different states, but separated by path of total length 0.0
            end
        end
    end
    # aggregate over trees and rates
    for ti in 1:nt
        obj._loglikcache[:,:,ti] .+= obj.priorltw[ti]
    end
    # obj._loglikcache .-= log(nr)
    # done below in 1 instead nsites x nrates x ntrees calculations, but _loglikcache not modified
    for ci in whichtrait
        obj._sitecache[ci] = logsumexp(view(obj._loglikcache, ci,:,1:nt))
    end
    if obj.siteweight !== nothing
        obj._sitecache .*= obj.siteweight
    end
    loglik = sum(obj._sitecache)
    loglik -= log(nr) * obj.totalsiteweight
    obj.loglik = loglik
    return loglik
end

"""
    discrete_corelikelihood_trait!(obj::SSM, t::Integer, ci::Integer, ri::Integer)

Return the likelihood for tree `t`, trait (character/site) index `ci` and rate category `ri`.
Update & modify the forward & directional log-likelihoods `obj.forwardlik`
and `obj.directlik`, which are indexed by [state, node_number or edge_number].
Used by [`discrete_corelikelihood!`](@ref).

**Preconditions**: `obj.logtrans` updated, edges directed, nodes/edges preordered
"""
function discrete_corelikelihood_trait!(obj::SSM, t::Integer, ci::Integer, ri::Integer)
    forwardlik = obj.forwardlik
    directlik  = obj.directlik
    tree = obj.displayedtree[t]
    k = nstates(obj.model)   # also = size(logtrans,1) if not RateVariationAcrossSites
    fill!(forwardlik, 0.0) # re-initialize for each trait, each iteration
    fill!(directlik,  0.0)
    loglik = 0.
    for ni in reverse(1:length(tree.nodes_changed)) # post-order
        n = tree.nodes_changed[ni]
        nnum = n.number # same n.number across trees for a given node
        if n.leaf # need forwardlik initialized at 0: keep at 0 = log(1) if no data
            state = obj.trait[nnum][ci] # here: data assumed in a row n.number
            if !ismissing(state)
                for i in 1:k
                    forwardlik[i,nnum] = -Inf64 # log(0) = -Inf if i != observed state
                end
                forwardlik[state, nnum] = 0.
            end
        else # forward likelihood = product of direct likelihood over all children edges
            for e in n.edge
                n == getParent(e) || continue # to next edge if n is not parent of e
                forwardlik[:,nnum] .+= view(directlik, :,e.number)
            end
        end
        if ni==1 # root is first index in nodes changed
            loglik = logsumexp(obj.prioratroot + view(forwardlik, :,nnum)) # log P{data for ci | tree t, rate ri}
            break # out of loop over nodes
        end
        # if we keep going, n is not the root
        # calculate direct likelihood on the parent edge of n
        for e in n.edge
            if n == getChild(e)
                lt = view(obj.logtrans, :,:,e.number, ri)
                for i in 1:k # state at parent node
                    directlik[i,e.number] = logsumexp(view(lt,i,:) + view(forwardlik,:,nnum))
                end
                break # we visited the parent edge: break out of for loop
            end
        end #loop over edges
    end # of loop over nodes
    return loglik
end

"""
    posterior_logtreeweight(obj::SSM, trait = 1)

Return an array A such that A[t] = log of P(tree `t` and trait `trait`)
if a single `trait` is requested, or A[i,t]= log of P(tree `t` and trait `i`)
if `trait` is a vector or range (e.g. `trait = 1:obj.nsites`).
These probabilities are conditional on the model parameters in `obj`.

**Precondition**: `_loglikcache` updated by [`discrete_corelikelihood!`](@ref)

# examples

```jldoctest
julia> net = readTopology("(((A:2.0,(B:1.0)#H1:0.1::0.9):1.5,(C:0.6,#H1:1.0::0.1):1.0):0.5,D:2.0);");

julia> m1 = BinaryTraitSubstitutionModel([0.1, 0.1], ["lo", "hi"]); # arbitrary rates

julia> using DataFrames

julia> dat = DataFrame(species=["C","A","B","D"], trait=["hi","lo","lo","hi"]);

julia> fit = fitdiscrete(net, m1, dat); # optimized rates: α=0.27 and β=0.35

julia> pltw = PhyloNetworks.posterior_logtreeweight(fit)
2-element Array{Float64,1}:
 -0.08356519024776699
 -2.523619878044531

julia> exp.(pltw) # posterior trees probabilities (sum up to 1)
2-element Array{Float64,1}:
 0.9198311206979973
 0.08016887930200293

julia> round.(exp.(fit.priorltw), digits=4) # the prior tree probabilities are similar here (tiny data set!)
2-element Array{Float64,1}:
 0.9
 0.1
```
"""
function posterior_logtreeweight(obj::SSM, trait = 1)
    # ts[site,tree] = log P(data and tree) at site, integrated over rates
    d = length(size(trait)) # 0 if single trait, 1 if vector of several traits
    ts = dropdims(mapslices(logsumexp, view(obj._loglikcache, trait,:,:),
                            dims=d+1); dims=1)
    if d>0 ts = permutedims(ts); end # now: ts[tree] or ts[tree,site]
    siteliks = mapslices(logsumexp, ts, dims=1) # 1 x ntraits array (or 1-element vector)
    # -log(nr) missing from both ts and siteliks, but cancels out next
    ts .-= siteliks
    return ts
end

"""
    traitlabels2indices(data, model::SubstitutionModel)

Check that the character states in `data` are compatible with (i.e. subset of)
the trait labels in `model`. All columns are used.
`data` can be a DataFrame or a Matrix (multiple traits), or a Vector (one trait).
Return a vector of vectors (one per species) with integer entries,
where each state (label) is replaced by its index in `model`.
For DNA data, any ambiguous site is treated as missing.
"""
function traitlabels2indices(data::AbstractVector, model::SubstitutionModel)
    A = Vector{Vector{Union{Missings.Missing,Int}}}(undef, 0) # indices of trait labels
    labs = getlabels(model)
    for l in data
        vi = missing
        if !ismissing(l)
            vi = findfirst(isequal(l), getlabels(model)) # value index in model labels
            vi !== nothing || error("trait $l not found in model")
        end
        push!(A, [vi])
    end
    return A
end

function traitlabels2indices(data::Union{AbstractMatrix,AbstractDataFrame},
                             model::SubstitutionModel)
    A = Vector{Vector{Union{Missings.Missing,Int}}}(undef, 0) # indices of trait labels
    labs = getlabels(model)
    isDNA = typeof(labs) == Array{DNA,1}
    ntraits = size(data,2) #starting with spcies column
    for i in 1:size(data,1) # go row by row (which is species by species)
        V = Vector{Union{Missings.Missing,Int}}(undef, ntraits)
        for j in 1:ntraits
            vi = missing # value index
            @inbounds l = data[i,j] # value label
            if !isDNA && ismissing(l)
                vi = missing
            elseif isDNA #else if DNA
                if typeof(l) == String #takes string and converts to a Char so that we can convert to DNA
                    l = Vector{Char}(l)[1]
                end
                l = convert(DNA, l)
            end
            if !ismissing(l)
                vi = findfirst(isequal(l), labs)
                if vi === nothing
                    # ideally, handle ambiguous DNA types optimally
                    if isDNA #&& BioSymbols.isambiguous(l)
                        vi = missing
                    else
                        error("trait $l not found in model")
                    end
                end
            end
            V[j] = vi
        end
        push!(A, V)
    end
    return A
end

"""
    check_matchtaxonnames!(species, data, net)

Modify `species` and `dat` by removing the species (rows) absent from the network.
Return a new network (`net` is *not* modified) with tips matching those in species:
if some species in `net` have no data, these species are pruned from the network.
The network also has its node names reset, such that leaves have nodes have
consecutive numbers starting at 1, with leaves first.
Used by [`fitdiscrete`](@ref) to build a new [`StatisticalSubstitutionModel`](@ref).
"""
function check_matchtaxonnames!(species::AbstractVector, dat::AbstractVector, net::HybridNetwork)
    # 1. basic checks for dimensions and types
    eltt = eltype(dat)
    @assert eltt <: AbstractVector "traits should be a vector of vectors"
    @assert nonmissingtype(eltype(eltt)) <: Integer "traits should be integers (label indices)"
    @assert !isempty(dat) "empty data vector!"
    ntraits = length(dat[1])
    for d in dat
        @assert length(d)==ntraits "all species should have the same number of traits"
    end
    @assert length(dat) == length(species) "need as many species as rows in trait data"
    # 2. match taxon labels between data and network
    netlab = tipLabels(net)
    ind2notinnet = findall(x -> x ∉ netlab, species) # species not in network
    deleteat!(species, ind2notinnet)
    deleteat!(dat,     ind2notinnet)
    nvalues = [sum(.!ismissing.(d)) for d in dat] # species with completely missing data
    indmissing = findall(x -> x==0, nvalues)
    deleteat!(species, indmissing)
    deleteat!(dat,     indmissing)
    indnotindat = findall(x -> x ∉ species, netlab) # species not in data
    net = deepcopy(net)
    if !isempty(indnotindat)
        @warn "the network contains taxa with no data: those will be pruned"
        for i in indnotindat
            deleteleaf!(net, netlab[i])
        end
    end
    # 3. calculate order of rows to have species with node.number i on ith row
    resetNodeNumbers!(net; checkPreorder=true, type=:ape) # tip species: now with numbers 1:n
    resetEdgeNumbers!(net, false) # to use edge as indices: 1:numEdges
    netlab = [n.name for n in sort(net.leaf, by = x -> x.number)]
    nspecies = length(netlab)
    o = Vector{Int}(undef, nspecies)
    for i in 1:nspecies
        @inbounds o[i] = something(findfirst(isequal(netlab[i]), species),0)
    end
    count(!iszero, o) == nspecies || # number of non-zeros should be total size
        error("weird: even after pruning, species in network have no data")
    return (o,net)
end

"""
    ancestralStateReconstruction(obj::SSM, trait::Integer = 1)

Estimate the marginal probability of ancestral states for discrete character
number `trait` (first trait by default).
The parameters of the [`StatisticalSubstitutionModel`](@ref) object `obj`
must first be fitted using [`fitdiscrete`](@ref), and ancestral state reconstruction
is conditional on the estimated parameters. If these parameters were estimated
using all traits, they are used as is, to do ancestral state reconstruction of the
particular `trait` of interest.

**output**: data frame with a first column for the node numbers, a second column for
the node labels, and a column for each possible state: the entries in these columns
give the marginal probability that a given node has a given state.

warnings:
- node numbers and node labels refer to those in `obj.net`, which might
  have a different internal representation of nodes than the original network
  used to build `obj`.
- `obj` is modified: its likelihood fields (forward, directional & backward)
  are updated to make sure that they correspond to the current parameter values
  in `obj.model`, and to the `trait` of interest.

limitations: the following are not checked.
- Assumes that every node in the large network is also present
  (with descendant leaves) in each displayed tree.
  This is not true if the network is not tree-child...
- Assumes that the root is also in each displayed tree, which
  may not be the case if the root had a hybrid child edge.

See also [`posterior_logtreeweight`](@ref) and
[`discrete_backwardlikelihood_trait!`](@ref) to update `obj.backwardlik`.

# examples

```jldoctest
julia> net = readTopology("(((A:2.0,(B:1.0)#H1:0.1::0.9):1.5,(C:0.6,#H1:1.0::0.1):1.0):0.5,D:2.0);");

julia> m1 = BinaryTraitSubstitutionModel([0.1, 0.1], ["lo", "hi"]);

julia> using DataFrames

julia> dat = DataFrame(species=["C","A","B","D"], trait=["hi","lo","lo","hi"]);

julia> fit1 = fitdiscrete(net, m1, dat);

julia> asr = ancestralStateReconstruction(fit1)
9×4 DataFrames.DataFrame
│ Row │ nodenumber │ nodelabel │ lo       │ hi       │
│     │ Int64      │ String    │ Float64  │ Float64  │
├─────┼────────────┼───────────┼──────────┼──────────┤
│ 1   │ 1          │ A         │ 1.0      │ 0.0      │
│ 2   │ 2          │ B         │ 1.0      │ 0.0      │
│ 3   │ 3          │ C         │ 0.0      │ 1.0      │
│ 4   │ 4          │ D         │ 0.0      │ 1.0      │
│ 5   │ 5          │ 5         │ 0.286018 │ 0.713982 │
│ 6   │ 6          │ 6         │ 0.319454 │ 0.680546 │
│ 7   │ 7          │ 7         │ 0.168548 │ 0.831452 │
│ 8   │ 8          │ 8         │ 0.767361 │ 0.232639 │
│ 9   │ 9          │ H1        │ 0.782777 │ 0.217223 │

julia> using PhyloPlots

julia> plot(fit1.net, :R, nodeLabel = asr[!,[:nodenumber, :lo]], tipOffset=0.2); # pp for "lo" state
```
"""
function ancestralStateReconstruction(obj::SSM, trait::Integer = 1)
    # posterior probability of state i at node n: proportional to
    # sum_{tree t, rate r} exp( ltw[t] + backwardll[i,n] given t,r + forwardll[i,n] given t,r ) / nr
    trait <= obj.nsites || error("trait $trait is larger than the number of traits in the data")
    nnodes = length(obj.net.node) # may be smaller than 2nd size of bkd or frd
    update_logtrans(obj)
    bkd = view(obj.backwardlik, :, 1:nnodes)
    frd = view(obj.forwardlik, :, 1:nnodes)
    ltw = obj.priorltw
    res = similar(bkd) # first: hold the cumulative logsumexp of bkd + frd + ltw
    fill!(res, -Inf64)
    nr = length(obj.ratemodel.ratemultiplier)
    for t in 1:length(obj.displayedtree)
        ltprior = ltw[t]
        for ri in 1:nr
            # update forward & directional likelihoods
            discrete_corelikelihood_trait!(obj,t,trait,ri)
            # update backward likelihoods
            discrete_backwardlikelihood_trait!(obj,t,ri)
            # P{state i at node n} ∝ bkd[i,n] * frd[i,n] given tree & rate:
            # res = logaddexp(res, ltw[t] + bkd + frd)  ---   -log(nr) will cancel out
            broadcast!(logaddexp, res, res, ltprior .+ bkd + frd)
        end
    end
    # normalize the results at each node: p_i / sum(p_j over all states j)
    traitloglik = logsumexp(res[:,1]) # sum p_j at node 1 or at any node = loglikelihood + log(nr)
    res .= exp.(res .- traitloglik)
    nodestringlabels = Vector{String}(undef, nnodes)
    for n in obj.net.node
        nodestringlabels[n.number] = (n.name == "" ? string(n.number) : n.name)
    end
    dat = DataFrame(transpose(res), Symbol.(getlabels(obj.model)))
    insertcols!(dat, 1, :nodenumber => collect(1:nnodes), makeunique=true)
    insertcols!(dat, 2, :nodelabel  => nodestringlabels,  makeunique=true)
    return dat
end

"""
    discrete_backwardlikelihood_trait!(obj::SSM, tree::Integer, ri::Integer)

Update and return the backward likelihood (last argument `backwardlik`)
assuming rate category `ri` and tree index `tree`,
using current forward and backwards likelihoods in `obj`:
these depend on the trait (or site) given to the last call to
`discrete_corelikelihood_trait!`.
Used by `ancestralStateReconstruction`.

**warning**: assume correct transition probabilities.
"""
function discrete_backwardlikelihood_trait!(obj::SSM, t::Integer, ri::Integer)
    backwardlik = obj.backwardlik
    directlik  = obj.directlik
    tree = obj.displayedtree[t]
    k = nstates(obj.model)
    fill!(backwardlik, 0.0) # re-initialize for each trait, each iteration
    bkwtmp = Vector{Float64}(undef, k) # to hold bkw lik without parent edge transition
    if typeof(obj.model) == NASM
        logprior = log.(stationary(obj.model))
    else #trait models
        logprior = [-log(k) for i in 1:k] # uniform prior at root
    end
    for ni in 1:length(tree.nodes_changed) # pre-order traversal to calculate backwardlik
        n = tree.nodes_changed[ni]
        nnum = n.number
        if ni == 1 # n is the root
            backwardlik[:,nnum] = logprior
        else
            pe = getMajorParentEdge(n)
            pn = getParent(pe)
            bkwtmp[:] = backwardlik[:,pn.number] # use bktmp's original memory
            for se in pn.edge
                if se != pe && pn == getParent(se) # then se is sister edge to pe
                    bkwtmp .+= view(directlik, :,se.number)
                end
            end
            lt = view(obj.logtrans, :,:,pe.number,ri)
            for j in 1:k # state at node n
                backwardlik[j,nnum] = logsumexp(bkwtmp + view(lt,:,j))
            end
        end
    end
    return backwardlik
end

"""
    learnlabels(model::Symbol, dat::DataFrame)

Return unique non-missing values in `dat`, and check that these labels
can be used to construct of substitution model of type `model`.

# examples:

```jldoctest
julia> using DataFrames

julia> dat = DataFrame(trait1 = ["A", "C", "A", missing]); # 4×1 DataFrame

julia> PhyloNetworks.learnlabels(:BTSM, dat)
2-element Array{String,1}:
 "A"
 "C"

julia> PhyloNetworks.learnlabels(:JC69, dat)
2-element Array{String,1}:
 "A"
 "C"
```
"""
function learnlabels(modSymbol::Symbol, dat::AbstractDataFrame)
    labels = mapreduce(x -> unique(skipmissing(x)), union, eachcol(dat))
    if modSymbol == :BTSM
        length(labels) == 2 || error("Binary Trait Substitution Model supports traits with two states. These data have do not have two states.")
    elseif modSymbol == :TBTSM
        unique(skipmissing(dat[!,1])) == 2 && unique(skipmissing(dat[!,2]) == 2) ||
          error("Two Binary Trait Substitution Model supports two traits with two states each.")
    elseif modSymbol in [:HKY85, :JC69]
        typeof(labels) == Array{DNA,1} ||
        (typeof(labels) == Array{Char,1}   &&       all(in.(uppercase.(labels), "-ABCDGHKMNRSTVWY"))) ||
        (typeof(labels) == Array{String,1} && all(occursin.(uppercase.(labels), "-ABCDGHKMNRSTVWY"))) ||
        # "ACGT" would dissallow ambiguous sites
          error("$modSymbol requires that trait data are dna bases A, C, G, and T")
    end
    return labels
end

"""
    startingrate(net)

Estimate an evolutionary rate appropriate for the branch lengths in the network,
which should be a good starting value before optimization in `fitdiscrete`,
assuming approximately 1 change across the entire tree.
If all edge lengths are missing, set starting rate to 1/(number of taxa).
"""
function startingrate(net::HybridNetwork)
    totaledgelength = 0.0
    for e in net.edge
        if e.length > 0.0
            totaledgelength += e.length
        end
    end
    if totaledgelength == 0.0 # as when all edge lengths are missing
        totaledgelength = net.numTaxa
    end
    return 1.0/totaledgelength
end

"""
    defaultsubstitutionmodel(network, modsymbol::Symbol, data::DataFrame,
                  siteweights::Vector)

Return a statistical substitution model (SSM) with appropriate state labels
and a rate appropriate for the branch lengths in `net`
(see [`startingrate`](@ref)).
The `data` frame must have the actual trait/site data in columns 2 and up,
as when the species names are in column 1.
For DNA data, the relative rate model is returned, with a
stationary distribution equal to the empirical frequencies.
"""
function defaultsubstitutionmodel(net::HybridNetwork, modsymbol::Symbol, data::DataFrame,
        siteweights=repeat([1.], inner=size(data,2))::AbstractVector)
    rate = startingrate(net)
    actualdat = view(data, :, 2:size(data,2))
    labels = learnlabels(modsymbol, actualdat)
    if modsymbol == :JC69
        return JC69([1.0], true) # 1.0 instead of rate because relative version
    elseif modsymbol == :HKY85 # transition/transversion rate ratio
        return HKY85([1.0], empiricalDNAfrequencies(actualdat, siteweights), true)
    elseif modsymbol == :ERSM
        return EqualRatesSubstitutionModel(length(labels), rate, labels);
    elseif modsymbol == :BTSM
        return BinaryTraitSubstitutionModel([rate, rate], labels)
    elseif modsymbol == :TBTSM
        return TwoBinaryTraitSubstitutionModel([rate, rate, rate, rate, rate,
                                                rate, rate, rate], labels)
    else
        error("model $modsymbol is unknown or not implemented yet")
    end
end

# fixit: new type for two (dependent) binary traits
# need new algorithm for model at hybrid nodes: displayed trees aren't enough


text:
abstract type AbstractMetric end
struct WuPalmer <: AbstractMetric end
struct SimilarityDepth <: AbstractMetric end

function similarity(w1::AbstractString, w2::AbstractString)
    phrase1 = basic_stem.(split(w1, ' '))
    phrase2 = basic_stem.(split(w2, ' '))
    max(_similarity_of_stemmed_words(basic_stem(w1), basic_stem(w2)),
        mean(_similarity_of_stemmed_words(a, b) for a in phrase1 for b in phrase2))
end

function _similarity_of_stemmed_words(w1::AbstractString, w2::AbstractString)
    max(similarity(WuPalmer(), w1, w2),
        similarity(SimilarityDepth(), w1, w2))
end

function similarity(::WuPalmer, w1::AbstractString, w2::AbstractString)
    if w1 == w2
        return 1.0
    end
    if !(w1 in keys(CACHE.paths)) || !(w2 in keys(CACHE.paths))
        return 0.0
    end
    s = 0.0
    for p1 in CACHE.paths[w1]
        for p2 in CACHE.paths[w2]
            s = max(s, similarity(WuPalmer(), p1, p2))
        end
    end
    return s
end

function similarity(::WuPalmer, p1::BasicPath, p2::BasicPath)
    2 * common_ancestor_depth(p1, p2) / (length(p1) + length(p2))
end

function common_ancestor_depth(p1::BasicPath, p2::BasicPath)
    max_depth = min(length(p1), length(p2))
    for i in 1:max_depth
        if p1[i] != p2[i]
            return i - 1
        end
    end
    return max_depth
end

function similarity(::SimilarityDepth, w1::AbstractString, w2::AbstractString)
    if w1 == w2
        return 1.0
    end
    s = 0.0
    if !(w1 in keys(CACHE.similarity_groups)) || !(w2 in keys(CACHE.synsets))
        return s
    end
    for groups in CACHE.similarity_groups[w1]
        for synset in CACHE.synsets[w2]
            s = max(s, similarity(SimilarityDepth(), groups, synset))
        end
    end
    s
end

const ADJECTIVE_EQUIVALENT_DEPTH = 10

function similarity(::SimilarityDepth, groups::BasicSimilarityGroups, synset::BasicSynset)
    d = similarity_distance(groups, synset)
    if d === nothing
        return 0.0
    end
    2 * (ADJECTIVE_EQUIVALENT_DEPTH - d) / (2 * ADJECTIVE_EQUIVALENT_DEPTH)
end

function similarity_distance(groups::BasicSimilarityGroups, synset::BasicSynset)
    for i in 1:length(groups)
        if synset in groups[i]
            return i - 1
        end
    end
    nothing
end