using VortexDistributions, Distributions, Statistics, BenchmarkTools

include("../util.jl")

##
N = 2^4
Lx = 200; Ly = Lx
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi4 = Torus(psi0,x,y)

##
N=2^5
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi5 = Torus(psi0,x,y)

##
N=2^6
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi6 = Torus(psi0,x,y)

##
N=2^7
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi7 = Torus(psi0,x,y)

##
N=2^8
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi8 = Torus(psi0,x,y)

##
N=2^9
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi9 = Torus(psi0,x,y)

##
N=2^10
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi10 = Torus(psi0,x,y)

##
N=2^11
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi11 = Torus(psi0,x,y)

##
N=2^12
x = LinRange(-Lx/2,Lx/2, N+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, N+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi12 = Torus(psi0,x,y)

##
rand_vorts = rand_pointvortex(2, psi)

vortex!(psi4, rand_vorts)
vortex!(psi5, rand_vorts)
vortex!(psi6, rand_vorts)
vortex!(psi7, rand_vorts)
vortex!(psi8, rand_vorts)
vortex!(psi9, rand_vorts)
vortex!(psi10, rand_vorts)
vortex!(psi11, rand_vorts)
vortex!(psi12, rand_vorts)

bench4 = @benchmark naivePlaquette(psi4.ψ, [psi4.x, psi4.y], 0)
bench5 = @benchmark naivePlaquette(psi5.ψ, [psi5.x, psi5.y], 0)
bench6 = @benchmark naivePlaquette(psi6.ψ, [psi6.x, psi6.y], 0)
bench7 = @benchmark naivePlaquette(psi7.ψ, [psi7.x, psi7.y], 0)
bench8 = @benchmark naivePlaquette(psi8.ψ, [psi8.x, psi8.y], 0)
bench9 = @benchmark naivePlaquette(psi9.ψ, [psi9.x, psi9.y], 0)
bench10 = @benchmark naivePlaquette(psi10.ψ, [psi10.x, psi10.y], 0)
bench11 = @benchmark naivePlaquette(psi11.ψ, [psi11.x, psi11.y], 0)
bench12 = @benchmark naivePlaquette(psi12.ψ, [psi12.x, psi12.y], 0)

bench_all = [bench4, bench5, bench6, bench7, bench8, bench9, bench10, bench11, bench12]
N_all = [2^x for x=4:12]

using Plots, LaTeXStrings
bench_all_mean = [mean(bench_all[i].times) for i in 1:length(bench_all)]
plot(log.(2, N_all), log.(bench_all_mean), marker='*', xlab = L"n = 2^x", ylab = L"time = ln(10^9 s)", title = "log-log plaquette n x n grid", label = "naive")

fench4 = @benchmark remove_vortices_edge(findvortices_jumps(psi4), psi4)
fench5 = @benchmark remove_vortices_edge(findvortices_jumps(psi5), psi5)
fench6 = @benchmark remove_vortices_edge(findvortices_jumps(psi6), psi6)
fench7 = @benchmark remove_vortices_edge(findvortices_jumps(psi7), psi7)
fench8 = @benchmark remove_vortices_edge(findvortices_jumps(psi8), psi8)
fench9 = @benchmark remove_vortices_edge(findvortices_jumps(psi9), psi9)
fench10 = @benchmark remove_vortices_edge(findvortices_jumps(psi10), psi10)
fench11 = @benchmark remove_vortices_edge(findvortices_jumps(psi11), psi11)
fench12 = @benchmark remove_vortices_edge(findvortices_jumps(psi12), psi12)

fench_all = [fench4, fench5, fench6, fench7, fench8, fench9, fench10, fench11, fench12]
fench_all_mean = [mean(fench_all[i].times) for i in 1:length(fench_all)]

plot!(log.(2, N_all), log.(fench_all_mean), marker='*', label = "optimised")


############################################################################################
"""
$(TYPEDEF)

Utility struct to help assign boundaries to parameter - keeps scalar parameter constrained.

# Fields
$(TYPEDFIELDS)
"""
struct Constrained{T<:Real} <: AbstractConstraint
    min::T
    max::T
    function Constrained(min::S, max::T, ϵ = 1e-6) where {S<:AbstractFloat, T<:AbstractFloat}
        ArgCheck.@argcheck min < max
        vals = promote(min, max)
        return new{eltype(vals)}(vals[1] + ϵ, vals[2] - ϵ)
    end
end
#=
#!NOTE: Bijectors seems to constrain Int to Float -> so we will remove this option
function Constrained(min::I, max::I) where {I<:Integer}
    ArgCheck.@argcheck min < max
    return Constrained{I}(min, max)
end
=#

############################################################################################
function construct_transform(info::Constrained, val)
    transform = Bijectors.TruncatedBijector(info.min, info.max)
    return transform, Bijectors.inverse(transform)
end

############################################################################################
function _check(
    _rng::Random.AbstractRNG,
    constraint::Constrained,
    val::R,
) where {R<:Real}
    ArgCheck.@argcheck typeof(val) == typeof(constraint.min) == typeof(constraint.max) "Type of Constrained boundaries must match type of val"
    ArgCheck.@argcheck constraint.min < val < constraint.max "val must be between boundaries"
    return true
end

############################################################################################
# Constrained constraint -> just flatten
function construct_flatten(
    output::Type{T},
    flattentype::F,
    unflattentype::U,
    constraint::Constrained,
    x
) where {T<:AbstractFloat,F<:FlattenTypes,U<:UnflattenTypes}
    return construct_flatten(T, flattentype, unflattentype, x)
end

############################################################################################
#=
!NOTE: For this constraint, we use a Bijector as Transformer, so we do not need to add a new functors
1. MyTransformer <: AbstractTransformer, MyInverseTransformer <: AbstractTransformer
2. define a function construct_transform(MyConstraint, val) -> MyTransformer, MyInverseTransformer
3. overload unconstrain and log_abs_det_jac on MyTransformer.
4. overload constrain on MyInverseTransformer.
=#

############################################################################################
#Export
export
    Constrained,
    construct_flatten,
    construct_transform,
    constrain,
    unconstrain,
    log_abs_det_jac


"""
    msmplot(T; pi_i=nothing, x=nothing, y=nothing, filename=nothing,
            edgewidth_scale=10.0, arrow_scale=0.1, nodesize=0.5, fontsize=10, names=[], dpi=100)

Visualize the graphical structure of the given Markov state model parameters.
`T` is a transition probability matrix whose elements Tij represents the probablity of transition from state i to j.
`T` should satisfy the detailed balance condition. `pi_i` is a vector whose elements are equilibrium probabilities of states.
`x` and `y` are X and Y coordinates of states, respectively.

# Examples
```julia-repl
julia> nstate = 5
julia> T, pi_i = msmtransitionmatrix(rand(nstate, nstate))
julia> x = rand(nstate); y = rand(nstate)
julia> msmplot(T, pi_i=pi_i, x=x, y=y)
```
"""
function msmplot(T; pi_i=nothing, x=nothing, y=nothing, filename=nothing,
                 edgewidth_scale=3.0, arrow_scale=0.0001, nodesize=0.5, fontsize=10, names=[], dpi=100)
  n = size(T, 1)
  g = MetaDiGraph(n, 0.0)
  U = typeof(T[1, 1])

  for i = 1:n
    for j = 1:n
      if !iszero(T[i, j])
        Graphs.SimpleGraphs.add_edge!(g, i, j)
      end
     end
  end

  for i = 1:n
    for j = 1:n
      if !iszero(T[i, j])
        set_prop!(g, Edge(i, j), :weight, T[i, j])
      end
    end
  end

  if isnothing(pi_i)
    pi_i = ones(1, n)./n
    for i = 1:(n*100)
      pi_i .= pi_i * T
    end
    pi_i = pi_i./sum(pi_i)
    pi_i = pi_i[:]
  end

  if !isnothing(x)
    std_x = std(x)
    std_y = std(y)
    max_xy = maximum([maximum(x),maximum(y)])
    xx = (x[:].-mean(x))./std_x .- std_x
    yy = (y[:].-mean(y))./std_y .- std_y
  else
    xx = nothing
    yy = nothing
  end

  gr()
#  pyplot()
  p = graphplot(g, edge_width=(s,d,w) -> get_prop(g, Edge(s,d), :weight)*edgewidth_scale,
            #arrow=arrow(:closed, :head, 0.1, 0.1),
            arrow=arrow(:closed, :head, arrow_scale, arrow_scale),
            #markersize = markersize,
            node_weights = pi_i[:],
            #node_weights = 1:n,
            #markercolor = range(colorant"yellow", stop=colorant"red", length=n),
            markercolor = :white,
            names = names,
            nodesize = nodesize,
            fontsize = fontsize,
            fontcolor = :white,
            nodeshape = :circle,
            linecolor = :darkgrey,
            #shorten=0.95,
            x = xx,
            y = yy,
            curves = true,
            curvature_scalar = 0.8,
            dpi = dpi
            )

  if !isnothing(filename)
    savefig(p, filename)
  end

  return p
end

function msmsample(p)
  p_cum = cumsum(p) ./ sum(p)
  r = rand()
  for i = 1:length(p_cum)
    if r <= p_cum[i]
      return i
    end
  end
end

"""
    msmgenerate(nframe::Int, T, pi_i)
                states = zeros(typeof(nframe), nframe)

Randomly samples a state trajectory from the the given transition matrix, and equilibrium probabilities of states.

# Examples
```julia-repl
julia> T, pi_i = msmtransitionmatrix(C)
julia> states = msmgenerate(1000, T, pi_i)
```
"""
function msmgenerate(nframe::Int, T, pi_i)
  states = zeros(typeof(nframe), nframe)

  states[1] = msmsample(pi_i)
  for iframe = 2:nframe
    states[iframe] = msmsample(T[states[iframe-1], :])
  end
  return states
end

"""
    msmgenerate(nframe, T, pi_i, emission) -> states, observations

Randomly samples a state trajectory and observations from the given transition matrix, equilibrium probabilities of states, and emissions.

# Examples
```julia-repl
julia> T, pi_i = msmtransitionmatrix(C)
julia> states, observations = msmgenerate(1000, T, pi_i, emission)
```
"""
function msmgenerate(nframe::Int, T, pi_i, emission)
  states = zeros(typeof(nframe), nframe)
  observations = zeros(typeof(nframe), nframe)

  states[1] = msmsample(pi_i)
  observations[1] = msmsample(emission[states[1], :])
  for iframe = 2:nframe
    states[iframe] = msmsample(T[states[iframe-1], :])
    observations[iframe] = msmsample(emission[states[iframe], :])
  end
  return states, observations
end

"""
    msmcountmatrix(indexOfCluster; tau=1) -> C::Matrix

Transition count matrix from a sinlgle binned trajectory or a set of binned trajectories.
`indexOfCluster` is a vector or a set of vectors containing binned trajectorie(s).
Returns count matrix for transitions from state i to state j in a time step of tau.

# Examples
```julia-repl
julia> ta = mdload("ak.dcd")
julia> X = compute_distancemap(ta["atomname CA"])
julia> F = clusterkcenters(X)
julia> c = msmcountmatrix(F.indexOfCluster, tau=10)
```
"""
function msmcountmatrix(indexOfCluster; tau=1)
  if eltype(indexOfCluster) <: Number
    indexOfCluster2 = [indexOfCluster]
  else
    indexOfCluster2 = indexOfCluster
  end

  nstate = maximum(maximum(indexOfCluster2))
  ntraj = length(indexOfCluster2)
  U = typeof(indexOfCluster2[1][1])
  c = spzeros(U, nstate, nstate)
  for itraj = 1:ntraj
    nframe = length(indexOfCluster2[itraj])

    index_from = 1:(nframe-tau)
    index_to   = (1+tau):nframe
    indexOfCluster_from = indexOfCluster2[itraj][index_from]
    indexOfCluster_to   = indexOfCluster2[itraj][index_to]

    s = ones(U, length(indexOfCluster_from))
    c_1 = sparse(indexOfCluster_from, indexOfCluster_to, s, nstate, nstate)
    c .= c .+ c_1
  end

  return Array(c)
end

"""
    msmtransitionmatrix(C; TOLERANCE=10^(-4), verbose=true) -> T::Matrix, p::Vector

Estimate the transition probability matrix from count matrix `C`.
Detailed balance is implicitly imposed in the estimation.

Returns the transition probability matrix and the equilibrium probabilities of states.

# Examples
```julia-repl
julia> ta = mdload("ak.dcd")
julia> X = compute_distancemap(ta["atomname CA"])
julia> F = clusterkcenters(X)
julia> C = msmcountmatrix(F.indexOfCluster, tau=10)
julia> T, p = msmtransitionmatrix(C)
```
# References
```
This routines uses the reversible maximum likelihood estimator described in
K. A. Beauchamp, G. R. Bowman, T. J. Lane, L. Maibaum, I. S. Haque, and V. S. Pande,
MSMBuilder2: Modeling Conformational Dynamics on the Picosecond to Millisecond Scale,
J. Chem. Theory Comput. 7, 3412 (2011).
```
"""
function msmtransitionmatrix(C; TOLERANCE=10^(-4), verbose=true)
  c = Matrix{Float64}(C)
  nstate = size(c, 1)

  c_sym = c + c'
  x = c_sym

  c_rs = sum(c, dims=2)
  x_rs = sum(x, dims=2)

  logL_old = 2.0*TOLERANCE
  logL = 0.0
  count_iteration = 0

  x = zeros(Float64, nstate, nstate)
  x_new = zeros(Float64, nstate, nstate)

  while abs(logL_old - logL) > TOLERANCE
    count_iteration = count_iteration + 1
    logL_old = logL

    # fixed-point method
    for i = 1:nstate
      for j = 1:nstate
        denom = (c_rs[i]/x_rs[i]) + (c_rs[j]/x_rs[j])
        x_new[i, j] = (c[i, j] + c[j, i]) / denom
      end
    end

    # update
    x_rs .= sum(x_new, dims=2)
    x .= x_new
    logL = 0.0
    for i = 1:nstate
      for j = 1:nstate
        if !iszero(x[i, j]) & !iszero(x_rs[i])
          logL = logL + c[i, j] * log(x[i, j] / x_rs[i]);
        end
      end
    end

    if verbose & (mod(count_iteration, 10) == 0)
      Printf.@printf("%d iteration  LogLikelihood = %8.5e  delta = %8.5e  tolerance = %8.5e\n", count_iteration, logL, abs(logL_old-logL), TOLERANCE)
    end
  end

  pi_i = x_rs./sum(x_rs)
  pi_i = pi_i[:]
  t = x ./ x_rs

  return t, pi_i
end

function msmforward(data_list, T, pi_i, emission)
    ndata = length(data_list)
    nstate = length(T[1, :])
    logL = zeros(Float64, ndata)
    alpha_list = []
    factor_list = []
    for idata = 1:ndata
        data = data_list[idata]
        nframe = length(data)
        alpha  = zeros(Float64, (nframe, nstate))
        factor = zeros(Float64, nframe)
        alpha[1, :] = pi_i.*emission[:, data[1]]
        factor[1] = sum(alpha[1, :])
        alpha[1, :] = alpha[1, :]./factor[1]
        for iframe = 2:nframe
            alpha[iframe, :] = sum(alpha[iframe-1, :] .* T, dims=1)' .* emission[:, data[iframe]]
            factor[iframe] = sum(alpha[iframe, :])
            alpha[iframe, :] = alpha[iframe, :]./factor[iframe]
        end
        logL[idata] = sum(log.(factor))
        push!(alpha_list, alpha)
        push!(factor_list, factor)
    end
    logL, alpha_list, factor_list
end

function msmforward_missing(data_list, T, pi_i, emission)
  ndata = length(data_list)
  nstate = length(T[1, :])
  logL = zeros(Float64, ndata)
  alpha_list = []
  factor_list = []
  for idata = 1:ndata
      data = data_list[idata]
      nframe = length(data)
      alpha  = zeros(Float64, (nframe, nstate))
      factor = zeros(Float64, nframe)
      if data[1] === missing
        alpha[1, :] = pi_i
      else
        alpha[1, :] = pi_i.*emission[:, data[1]]
      end
      factor[1] = sum(alpha[1, :])
      alpha[1, :] = alpha[1, :]./factor[1]
      for iframe = 2:nframe
          if data[iframe] === missing
            alpha[iframe, :] = sum(alpha[iframe-1, :] .* T, dims=1)'
          else
            alpha[iframe, :] = sum(alpha[iframe-1, :] .* T, dims=1)' .* emission[:, data[iframe]]
          end
          factor[iframe] = sum(alpha[iframe, :])
          alpha[iframe, :] = alpha[iframe, :]./factor[iframe]
      end
      logL[idata] = sum(log.(factor))
      push!(alpha_list, alpha)
      push!(factor_list, factor)
  end
  logL, alpha_list, factor_list
end

function msmbackward(data_list, factor_list, T, pi_i, emission)
    ndata = length(data_list)
    nstate = length(T[1, :])
    logL = zeros(Float64, ndata)
    beta_list = []
    for idata = 1:ndata
        data   = data_list[idata]
        factor = factor_list[idata]
        nframe = length(data)
        beta   = zeros(Float64, (nframe, nstate))
        beta[nframe, :] .= 1.0
        for iframe = (nframe-1):-1:1
            beta[iframe, :] = sum((T .* (emission[:, data[iframe+1]] .* beta[iframe+1, :])'), dims=2) ./ factor[iframe+1]
        end
        logL[idata] = sum(log.(factor))
        push!(beta_list, beta)
    end
    logL, beta_list
end

function msmbackward_missing(data_list, factor_list, T, pi_i, emission)
    ndata = length(data_list)
    nstate = length(T[1, :])
    logL = zeros(Float64, ndata)
    beta_list = []
    for idata = 1:ndata
        data   = data_list[idata]
        factor = factor_list[idata]
        nframe = length(data)
        beta   = zeros(Float64, (nframe, nstate))
        beta[nframe, :] .= 1.0
        for iframe = (nframe-1):-1:1
            if data[iframe+1] === missing
                beta[iframe, :] = sum((T .* beta[iframe+1, :]'), dims=2) ./ factor[iframe+1]
            else
                beta[iframe, :] = sum((T .* (emission[:, data[iframe+1]] .* beta[iframe+1, :])'), dims=2) ./ factor[iframe+1]
            end
        end
        logL[idata] = sum(log.(factor))
        push!(beta_list, beta)
    end
    logL, beta_list
end

"""
    msmbaumwelch(observations, T_init, p_init, emission_init; TOLERANCE=10.0^(-4), MAXITERATION=Inf64) -> T::Matrix, p::Vector, emission::Matrix

Baum-Welch algorithm estimates the most probable transition probabilities from the given observation data.
In this function, detailed balance is implicitly imposed in the estimation,
so the equilibrium probabilities can be determined from the estimated transition probabilities.
Also, unlike the original Baum-Welch algorithm, the emission is NOT estimated in this function,
because the emission probabilites are usually known a priori in cases of molecular experiments.
`observations` is a set of observed vectors. `T_init`, `p_init` are initial transition probabilities and equilibrium probabilities, respectively.
`emission_init` is a matrix whose rows correspond to states, and columns correspond to observations.

Returns the transition probability matrix, the equilibrium probabilities of states,
and the emission probabilities (though the emission does not change).

# Examples
```julia-repl
julia> ta = mdload("ak.dcd")
julia> X = compute_distancemap(ta["atomname CA"])
julia> F = clusterkcenters(X)
julia> C = msmcountmatrix(F.indexOfCluster, tau=10)
julia> T, p = msmtransitionmatrix(C)
```
# References
```
The algorithm of this routines is based on the descriptions in PRML book by C. Bishop.
```
"""
function msmbaumwelch(data_list, T0, pi_i0, emission0; TOLERANCE=10.0^(-4), MAXITERATION=Inf64)
    ## setup
    check_convergence = Inf64
    count_iteration = 0
    logL_old = 1.0
    #if not isinstance(data_list, list):
    #    data_list = [data_list]
    ndata = length(data_list)
    nobs = length(emission0[1, :])
    nstate = length(T0[1, :])
    T = similar(T0)
    emission = similar(emission0)
    pi_i = similar(pi_i0)
    while (check_convergence > TOLERANCE) & (count_iteration <= MAXITERATION)
        ## E-step
        logL, alpha_list, factor_list = msmforward(data_list, T0, pi_i0, emission0)
        #print("1"); println(logL)
        logL2, beta_list = msmbackward(data_list, factor_list, T0, pi_i0, emission0)
        #print("2"); println(logL2)
        log_alpha_list = []
        for a in alpha_list
            push!(log_alpha_list, log.(a))
        end
        log_beta_list = []
        for b in beta_list
            push!(log_beta_list, log.(b))
        end
        log_T0 = log.(T0)
        log_emission0 = log.(emission0)

        ## M-step
        # pi
        # pi = np.zeros(nstate, dtype=np.float64)
        # log_gamma_list = []
        # for idata in range(ndata):
        #     log_gamma_list.append(log_alpha_list[idata] + log_beta_list[idata])
        #     pi = pi + np.exp(log_gamma_list[idata][0, :])
        # pi = pi/np.sum(pi)
        # pi_i = pi_i0
        # emission
        # emission = np.zeros((nstate, nobs), dtype=np.float64)
        # for idata in range(ndata):
        #     data = data_list[idata]
        #     for istate in range(nstate):
        #         for iobs in range(nobs):
        #             id = (data == iobs)
        #             if np.any(id):
        #                 emission[istate, iobs] = emission[istate, iobs] + np.sum(np.exp(log_gamma_list[idata][id, istate]))
        # emission[np.isnan(emission)] = 0.0
        # emission = emission / np.sum(emission, axis=1)[:, None]
        emission = emission0
        # T
        T = zeros(Float64, (nstate, nstate))
        for idata = 1:ndata
          data = data_list[idata]
          nframe = length(data)
          for iframe = 2:nframe
            #log_xi = bsxfun(@plus, log_alpha{idata}(iframe-1, :)', log_beta{idata}(iframe, :));
            log_xi = log_alpha_list[idata][iframe-1, :] .+ log_beta_list[idata][iframe, :]'
            #T = T .+ exp(bsxfun(@plus, log_xi, log_emission0(:, data(iframe))') + log_T0)./factor{idata}(iframe);
            T = T .+ exp.((log_xi .+ log_emission0[:, data[iframe]]') .+ log_T0) ./ factor_list[idata][iframe]
          end
        end
        #T[np.isnan(T)] = 0.0
        T = T ./ sum(T, dims=2)

        ## reversible T
        T, pi_i = msmtransitionmatrix(T, TOLERANCE=10^(-8), verbose=false)

        ## Check convergence
        count_iteration += 1
        logL = sum(logL)
        check_convergence = abs(logL_old - logL)
        if mod(count_iteration, 100) == 0
            Printf.@printf("%d iteration LogLikelihood = %e  delta = %e  tolerance = %e\n" , count_iteration, logL, check_convergence, TOLERANCE)
        end
        logL_old = logL
        pi_i0 = pi_i
        emission0 = emission
        T0 = T
    end
    T, pi_i, emission
end

function msmbaumwelch_missing(data_list, T0, pi_i0, emission0; TOLERANCE=10.0^(-4), MAXITERATION=Inf64)
  ## setup
  check_convergence = Inf64
  count_iteration = 0
  logL_old = 1.0
  #if not isinstance(data_list, list):
  #    data_list = [data_list]
  ndata = length(data_list)
  nobs = length(emission0[1, :])
  nstate = length(T0[1, :])
  T = similar(T0)
  emission = similar(emission0)
  pi_i = similar(pi_i0)
  while (check_convergence > TOLERANCE) & (count_iteration <= MAXITERATION)
      ## E-step
      logL, alpha_list, factor_list = msmforward_missing(data_list, T0, pi_i0, emission0)
      #print("1"); println(logL)
      logL2, beta_list = msmbackward_missing(data_list, factor_list, T0, pi_i0, emission0)
      #print("2"); println(logL2)
      log_alpha_list = []
      for a in alpha_list
          push!(log_alpha_list, log.(a))
      end
      log_beta_list = []
      for b in beta_list
          push!(log_beta_list, log.(b))
      end
      log_T0 = log.(T0)
      log_emission0 = log.(emission0)

      ## M-step
      # pi
      # pi = np.zeros(nstate, dtype=np.float64)
      # log_gamma_list = []
      # for idata in range(ndata):
      #     log_gamma_list.append(log_alpha_list[idata] + log_beta_list[idata])
      #     pi = pi + np.exp(log_gamma_list[idata][0, :])
      # pi = pi/np.sum(pi)
      # pi_i = pi_i0
      # emission
      # emission = np.zeros((nstate, nobs), dtype=np.float64)
      # for idata in range(ndata):
      #     data = data_list[idata]
      #     for istate in range(nstate):
      #         for iobs in range(nobs):
      #             id = (data == iobs)
      #             if np.any(id):
      #                 emission[istate, iobs] = emission[istate, iobs] + np.sum(np.exp(log_gamma_list[idata][id, istate]))
      # emission[np.isnan(emission)] = 0.0
      # emission = emission / np.sum(emission, axis=1)[:, None]
      emission = emission0
      # T
      T = zeros(Float64, (nstate, nstate))
      for idata = 1:ndata
        data = data_list[idata]
        nframe = length(data)
        for iframe = 2:nframe
          log_xi = log_alpha_list[idata][iframe-1, :] .+ log_beta_list[idata][iframe, :]'
          if data[iframe] === missing
            T = T .+ exp.(log_xi .+ log_T0) ./ factor_list[idata][iframe]
          else
            T = T .+ exp.((log_xi .+ log_emission0[:, data[iframe]]') .+ log_T0) ./ factor_list[idata][iframe]
          end
        end
      end
      #T[np.isnan(T)] = 0.0
      T = T ./ sum(T, dims=2)

      ## reversible T
      T, pi_i = msmtransitionmatrix(T, TOLERANCE=10^(-8), verbose=false)

      ## Check convergence
      count_iteration += 1
      logL = sum(logL)
      check_convergence = abs(logL_old - logL)
      if mod(count_iteration, 100) == 0
          Printf.@printf("%d iteration LogLikelihood = %e  delta = %e  tolerance = %e\n" , count_iteration, logL, check_convergence, TOLERANCE)
      end
      logL_old = logL
      pi_i0 = pi_i
      emission0 = emission
      T0 = T
  end
  T, pi_i, emission
end

"""
    msmviterbi(observation, T, p, emission) -> states::Vector

Viterbi algorithm estimates the most probable hidden state sequence from the observation data.
`observations` is a set of observed vectors. `T`, `p` are the transition probabilities and equilibrium probabilities, respectively.
`emission` is a matrix whose rows correspond to states, and columns correspond to observations.

Returns the transition probability matrix, the equilibrium probabilities of states,
and the emission probabilities (though the emission does not change).

# Examples
```julia-repl
julia> nframe = 1000
julia> states, observations = msmgenerate(nframe, T, pi_i, emission)
julia> states_estimated = msmviterbi(T, pi_i, emission, observation)
```
# References
```
The algorithm of this routines is based on the descriptions in PRML book by C. Bishop.
```
"""
function msmviterbi(observation, T, pi_i, emission)
    T2 = T .+ 10^(-100)
    T2 .= T2 ./ sum(T2, dims=2)
    pi_i2 = pi_i .+ 10^(-100)
    pi_i2 = pi_i2 ./sum(pi_i2)
    emission2 = emission .+ 10^(-100)
    emission2 .= emission2 ./ sum(emission2, dims=1)

    nframe = size(observation, 1)
    nstate = size(T2, 1)
    P = zeros(eltype(T), nstate, nframe)
    I = zeros(eltype(T), nstate, nframe)
    state_estimated = zeros(eltype(observation), nframe)

    # initialization
    if observation[1] === missing
      P[:, 1] .= log.(pi_i2)
    else
      P[:, 1] .= log.(pi_i2) .+ log.(emission2[:, observation[1]])
    end
    I[:, 1] .= zeros(eltype(T2), nstate)

    # argmax forward
    Z = zeros(eltype(T2), nstate, nstate)
    for t = 2:nframe
        Z .= P[:, (t-1):(t-1)] .+ log.(T2)
        #if observation[t] === missing
        #  Z .= P[:, (t-1):(t-1)] .+ log.(T2)
        #else
        #  Z .= P[:, (t-1):(t-1)] .+ log.(T2) .+ log.(emission2[:, observation[t]:observation[t]]')
        #end
        I[:, t] .= getindex.(argmax(Z, dims=1), 1)[:]
        if observation[t] === missing
          P[:, t] .= maximum(Z, dims=1)[:]
        else
          P[:, t] .= maximum(Z, dims=1)[:] .+ log.(emission[:, observation[t]])
        end
    end

    # termination
    P_star = maximum(P[:, nframe])
    state_estimated[nframe] = argmax(P[:, nframe])
    #@show P[:, nframe]

    # decoding
    for t = (nframe-1):-1:1
        state_estimated[t] = I[state_estimated[t+1], t+1]
    end

    return state_estimated
end

function msmimpliedtime(indexOfCluster, tau)
  n = length(tau)
  k = min(10, n)
  nstate = Int(maximum(indexOfCluster))
  #implied_timescale = zeros(Float64, n, k)
  implied_timescale = Matrix{Union{Float64, Missing}}(undef, n, nstate-1)
  for i in 1:n
    t = tau[i]
    #@printf "calculating tau = %d\n" t
    C = msmcountmatrix(indexOfCluster, tau=t)
    # TODO fixme; tarjan algorithm to remove non-ergodic graphs
    T, p = msmtransitionmatrix(C, TOLERANCE=10^(-8), verbose=false)
    F = eigen(T, sortby = x -> -real(x))
    #@show F.values
    for istate = 2:nstate
      r = real(F.values[istate])
      if r >  0.0
        implied_timescale[i, istate-1] = - t / log(r)
      else
        implied_timescale[i, istate-1] = missing
      end
    end
  end
  return implied_timescale
end


using PtFEM

old = pwd()
ProjDir = joinpath(dirname(@__FILE__), "..", "examples", "4 Static Equilibrium")
#cd(ProjDir)

include(joinpath(ProjDir, "p4.1.1.jl"))

#=
data = Dict(
  # Rod(nels, np_types, nip, finite_element(nod, nodof))
  :struc_el => Rod(4, 1, 1, Line(2, 1)),
  :properties => [1.0e5;],
  :x_coords => range(0, stop=1, length=5),
  :support => [(1, [0])],
  :loaded_nodes => [(1,[-0.625]),(2,[-1.25]),(3,[-1.25]),(4,[-1.25]),(5,[-0.625])]
)
data |> display
println()

@time m = FE4_1(data)
println()

println("Displacements:")
m.displacements |> display
println()

println("Actions:")
m.actions |> display
println()
=#

@show m.loads
println()

m1 = fromSkyline(m.kv, m.kdiag)
@show m1
println()

@show m.kdiag
println()

@show m.kv
println()

ms = skyline2sparse(m.kv, m.kdiag)
@show ms
println()

loads = zeros(m.neq+1)
if :loaded_nodes in keys(data)
  for i in 1:size(data[:loaded_nodes], 1)
    loads[m.nf[:, data[:loaded_nodes][i][1]]+1] = data[:loaded_nodes][i][2]
  end
end

@show loads
println()

@show m1 \ loads[2:end]
println()

@show ms \ loads[2:end]
println()
cd(old)


"""
Package: ForecastPlots

    ccf(x1::{AbstractVector,DataFrame},
        x2::{AbstractVector,DataFrame};
        type = "cor",
        lag = Integer(ceil(10*log10(length(x1)))),
        alpha = (0.95,0.99);
        plot = true)

Compute the cross-correlation or cros-covariance of two univariate series with an optional plot

The results are normalized to preserve homoscedasticity. The distribution used to normalize the data is an approximation of a Fisher Transformation via a Normal Distribution. There is a plot recipe for the returned object, if the type is `cor` the plot will also show confidence intervals for the given alpha values.

If, for a given integer `k`, `x2` repeats `x1` values such that x1[t] = x2[t+k] for all `i` then high correlation value will be placed *at the right from the center* in the results. That is, this convention will be represented in the plots as `x1_t = x2_{t+k} -> _____0__k__` meaning x2 behavior can be predicted by x1 in k units.

# Arguments
- `x1`: Vector or uni-dimensional DataFrame of data.
- `x2`: Vector or uni-dimensional DataFrame of data.
- `type`: Valid values are "cor" for correlation (default) and "cov" for convariance.
- `lag`: Maximum number of lags.
- `alpha`: A tuple with two thresholds (t1,t2) with t1 <= t2 to plot confidence intervals. The default values are 0.95 and 0.99.
- `plot`: When true displays a plot with the results.

# Returns

Cros-correlation vector with an optional plot

# Examples
```julia-repl
x1 = rand(100);
x2 = circshift(x1,6);
ccf(x1, x2; type="cor");
ccf(x1, x2; type="cor", plot = false)
41-element Vector{Float64}:
  0.09861519494432992
 -0.04384418688776631
 -0.1900182513825246
  [...]
```
"""
function ccf(x1::AbstractVector{<:A},
             x2::AbstractVector{<:B};
             type::String = "cor",
             lag::Integer = Integer(ceil(10*log10(length(x1)))),
             alpha::Tuple{AbstractFloat,AbstractFloat} = (0.95,0.99),
             plot::Bool = true,
             kw...) where {A<:Real, B<:Real}

    N = length(x1)
    @assert N == length(x2) "Vectors should be of equal size"
    @assert N >= 4 "Vectors should have at least four values"
    @assert isnothing(findfirst(isnan,x1)) "No missing values allowed"
    @assert isnothing(findfirst(isnan,x2)) "No missing values allowed"
    @assert isnothing(findfirst(ismissing,x1)) "No missing values allowed"
    @assert isnothing(findfirst(ismissing,x2)) "No missing values allowed"
    @assert type in  ["cor","cov"] "The options for `type` are `cor` and `cov`"
    @assert 1 <= lag <= N-4
    @assert length(alpha) == 2
    @assert 0.0 < alpha[1] < alpha[2] < 1.0

    # Promote to same Float type
    x1,x2,_ = Base.promote(x1,x2,[Base.promote_op(/,A,B)(1)])

    # auto-ccf
    auto = x1 == x2

    ft = (type == "cor" ? Statistics.cor : Statistics.cov)

    x = []

    for i in 1:lag
        push!(x,ft(x1[i+1:N],x2[1:(N-i)]))
    end

    if auto
        kx = collect(N-1:-1:N-lag)
    else
        x = vcat(ft(x1,x2),x) # adding central value (lag=0)
        x = reverse(x)
        for i in 1:lag
            push!(x,ft(x2[i+1:N],x1[1:(N-i)]))
        end
        kx = vcat(collect(N-lag:N-1), N, collect(N-1:-1:N-lag))
    end

    # ccf normalized with the standard error for the
    # Normal distsribution of an approximation for
    # the Fisher transformation
    function fse(v::Real)::AbstractFloat
        1/sqrt(v-3)
    end
    k = fse.(kx)/fse(N)
    cres = x ./ k

    call = "ccf("* (auto ? "x" : "x1, x2") *
        "; type=\""*type*
        "\", lag="*string(lag)*
        ", alpha="*string(alpha)*")"

    # Confidence Intervals
    a1 = alpha[1]
    a2 = alpha[2]
    z1 = Distributions.quantile(Normal(), a1 + (1-a1)/2)
    ci1 = z1*fse(N)
    z2 = Distributions.quantile(Normal(), a2 + (1-a2)/2)
    ci2 = z2*fse(N)
    ci = (ci1,ci2)

    plot && display(cf_gr((cres, N=N, type, lag, alpha, ci, auto); kw...))

    cres

end

function cf_gr(args; kw...)

    # plot configuration

    titlefont = font(8, "Courier")
    xtickfont = font(3, "Courier")
    ytickfont = font(3, "Courier")
    yguidefont = font(5, "Courier")
    xguidefont = font(5, "Courier")

    N = args.N
    if split(args.type,'_')[1] == "pacf"
        type = (args.type) == "pacf_step" ? "Stepwise PACF" : "Real PACF"
        type = (args.type) == "pacf_stepwise-real" ? "Stepwise and Real PACF   |s|r|" : type
        auto_prefix = ""
    else
        type = (args.type == "cor") ? "Correlation" : "Covariance"
        auto_prefix = args.auto ? "Auto-" : "Cross-"
    end

    yguide = auto_prefix*type
    xguide = "Lag"

    # Center ticks
    lr = length(args.cres)
    gap = div(lr,11)
    mp = div(lr+1,2)
    rs = mp:gap:lr
    ls = mp-gap:-gap:1

    if args.auto
        xticks = 1:gap:lr
    else
        xt = vcat(reverse(collect(ls)),collect(rs))
        xticks = (xt, xt .- xt[div(end+1,2)])
    end

    plot(collect(1:size(args.cres)[1]+1) .- ((ndims(args.cres) == 2) ? .1 : 0),
         ndims(args.cres) == 2 ? args.cres[:,1] : args.cres;
         seriestype = :sticks, linewidth = 100/(args.lag+1),
         legend = :none, kw...)

    if ndims(args.cres) == 2

        plot!(collect(1:size(args.cres)[1]+1) .+ .1,
              args.cres[:,2];
              seriestype = :sticks, linewidth = 100/(args.lag+1))

    end

    if (args.type == "cor") | (split(args.type,'_')[1] == "pacf")
        lg = args.lag
        a1 = args.alpha[1]; c1 = args.ci[1]
        a2 = args.alpha[2]; c2 = args.ci[2]
        dc = c2-c1

        plot!([-c1,c1]; seriestype = :hline, seriescolor = :black, linealpha = 0.5, linestyle = :dash)

        ax = args.auto ? lg : 2*lg + 1
        annotations =
            [(ax,c1+dc/8,string("CI %",a1*100),font(3, "Courier")),
             (ax,c2+dc/8,string("CI %",a2*100),font(3, "Courier"))]

        plot!([-c2,c2]; seriestype = :hline, seriescolor = :black,
              linealpha = 0.5, linestyle = :dot, annotations)

    end

end


using RecipesBase
import ColorTypes.HSV

@recipe function f(b::UniformRBFE)
    Nv = size(b.μ,1)
    dist = b.μ[end]-b.μ[1]
    v = range(b.μ[1]-0.1dist, b.μ[end]+0.1dist, length=200)
    seriestype := :line
    color_palette --> [HSV(i,1,0.7) for i in range(0, 255, length=Nv)]
    xguide --> "Scheduling signal \$v\$"
    title --> "Basis function expansion"
    label --> ""
    v, b(v)
end

@recipe function f(rbf::MultiUniformRBFE, style=:default)
    # color_palette --> [HSV(i,1,0.7) for i in range(0, 255, length=Nv)]
    xguide --> "Scheduling signal \$v_1\$"
    yguide --> "Scheduling signal \$v_2\$"
    title --> "Basis function expansion"
    label --> ""
    seriestype --> :surface

    c       = rbf.μ
    minb    = minimum(c,dims=2)
    maxb    = maximum(c,dims=2)
    dist    = maxb-minb
    Npoints = 50
    v = [range(mi, ma, length=Npoints) for (mi,ma) in zip(minb,maxb)]
    vg = meshgrid(v...)
    # v  = plotattributes[:seriestype] == :surface ? vg : v # not really sure what this is for
    if style == :default
        bg = map(vg...) do v1,v2
            maximum(rbf([v1,v2]))
        end
        (v..., bg)
    elseif style == :full
        colorbar --> false
        seriesalpha --> 0.5
        bg = map(vg...) do v1,v2
            rbf([v1,v2])
        end
        bg = cat(bg...,dims=3)
        bg = reshape(bg, size(bg,1), Npoints, Npoints)
        for i = 1:size(bg,1)
            @series (v..., bg[i,:,:])
        end
    end

end


"""
Row-by-row sparsification using a quantile threshold
"""
function sparsify_signal!(U_guess; quantile_threshold=0.95)
    sparsify(x, eps) = abs(x) > eps ? x : 0;
    for i in 1:size(U_guess, 1)
        tmp = abs.(U_guess[i, :])
        epsilon = quantile(tmp[tmp.>0.0], quantile_threshold)
        for i2 in 1:size(U_guess,2)
            U_guess[i,i2] = sparsify(U_guess[i,i2], epsilon)
        end
    end
end

"""
Row-by-row sparsification using a quantile threshold
"""
function sparsify_signal(U_guess; quantile_threshold=0.95)
    for i in 1:size(U_guess, 1)
        tmp = abs.(U_guess[i, :])
        epsilon = quantile(tmp[tmp.>0.0], quantile_threshold)
        ind = abs.(U_guess[i, :]) .< epsilon
        U_guess[i, ind] .= 0.0
    end
    return U_guess
end

"""Keeps only largest 'n' terms"""
function keep_n_terms(X, n)
    ind = sortperm(vec(abs.(X)), rev=true)
    new_X = zeros(size(X))
    new_X[ind[1:n]] = X[ind[1:n]]
    return new_X
end

"""
Row-by-row sparsification using a hard threshold
"""
function threshold_signal!(U_guess; hard_threshold=1e-5)
    sparsify(x, eps) = abs(x) > eps ? x : 0;
    for i in 1:size(U_guess, 1)
        tmp = abs.(U_guess[i, :])
        for i2 in 1:size(U_guess,2)
            U_guess[i,i2] = sparsify(U_guess[i,i2], hard_threshold)
        end
    end
end

# """
# Simple implementation of sequential least squares thresholding for a linear
#     regression model
# """
# function sparse_regression(X::Matrix, y::Matrix;
#                            num_iter=2,
#                            quantile_threshold=0.1,
#                            hard_threshold=nothing,
#                            num_terms=nothing)
#     if size(y,1) > 1
#         A = zeros(size(y,1), size(X,1))
#         # For the 2d predictor case, just loop over rows
#         for i in 1:size(y,1)
#             if num_terms !== nothing
#                 # Constrain the number of terms row by row
#                 if length(num_terms)>1
#                     n = num_terms[i]
#                 else
#                     n = num_terms
#                 end
#             else
#                 n = nothing
#             end
#             A[i, :] = sparse_regression(X, y[i:i,:],
#                             num_iter=num_iter,
#                             quantile_threshold=quantile_threshold,
#                             hard_threshold=hard_threshold,
#                             num_terms=n)
#         end
#         return A
#     end
#     # If this is only one row, initialize with L2
#     A = y/X
#     # sparsify_signal!(A, quantile_threshold=quantile_threshold)
#     # Redo, removing the rows in X that were zeroed out
#     for i in 1:num_iter
#         if hard_threshold == nothing && num_terms == nothing
#             sparsify_signal!(A, quantile_threshold=quantile_threshold)
#         elseif num_terms == nothing
#             threshold_signal!(A, hard_threshold=hard_threshold)
#         else
#             A = keep_n_terms(A, num_terms)
#         end
#         which_rows = vec(abs.(A) .> 0)
#         A_subset = y/(X[which_rows, :])
#         A[which_rows] = A_subset
#         A[.!which_rows] .= 0.0 # Retain the previous zeros
#     end
#
#     return A
# end

"""
Removes outliers, by default 3 times the median
"""
function remove_outliers(dat::Vector, metric=median, max_val=3)
    m = 3*median(abs.(dat))
    return dat[abs.(dat) .< m]
end

"""
Replaces outliers (by default 3 times the median) with 'replace_val'

function replace_outliers(dat::Vector,
            replace_val=NaN, metric=median, max_val=3)
"""
function replace_outliers(dat::Vector,
            replace_val=NaN, metric=median, max_val=3)
    new_dat = copy(dat)
    m = 3*median(abs.(dat))
    for i in 1:length(dat)
        if dat[i] > m
            new_dat[i] = replace_val
        end
    end
    return new_dat
end

# export sparse_regression, sparsify_signal, threshold_signal!,
#         keep_n_terms
export sparsify_signal, threshold_signal!,
        keep_n_terms, remove_outliers, replace_outliers


"""
The way the Louvain code works, you start with n clusters (singletons), and as nodes
move, many of these become empty. At certain points in the algorithm, it is
helpful to renumber the cluster IDs, so that they go from 1 to M where
M = number of non-empty clusters.

e.g. c = [2 3 5 3 9] ---> c = [1 2 3 2 4]

This version does this with both the Clusters array of arrays AND
the cluster indicator vector c:

c[i] = (integer) cluster ID that node i belongs to
Clusters[j] = (integer array) node IDs for nodes in cluster j
"""
function renumber(c::Vector{Int64},Clusters::Vector{Vector{Int}})

    n = length(c)
    map = sort(unique(c))     # map from new cluster ID to old cluster ID
    old2new = Dict()    # map from old cluster ID to new cluster ID
    for i = 1:length(map)
        old2new[map[i]] = i
    end
    cnew = zeros(Int64,n)

    Clusters = Clusters[map]

    # Rename the clusters
    for i = 1:n
        # newClus = findfirst(x->x == c[i],map)
        newClus = old2new[c[i]]
        cnew[i] = newClus
        push!(Clusters[newClus],i)
    end

    return cnew, Clusters

end

"""
See above function, this function doesn't care about the Clusters array
"""
function renumber(c::Vector{Int64})

    n = length(c)
    map = sort(unique(c))     # map from new cluster ID to old cluster ID
    old2new = Dict()          # map from old cluster ID to new cluster ID
    for i = 1:length(map)
        old2new[map[i]] = i
    end
    cnew = zeros(Int64,n)

    # Rename the clusters
    for i = 1:n
        cnew[i] = old2new[c[i]]
    end

    return cnew

end

# These are some useful functions for converting back and forth between
# different ways to store hypergraphs.
function NeighborList(He2n::SparseArrays.SparseMatrixCSC{Float64,Int64},Hn2e::SparseArrays.SparseMatrixCSC{Float64,Int64})
    """
    NeighborList: given hypergraph (edge x node) indicence and (node x edge)
    incidence, return a list of neighbors for each node.
    """
    m,n = size(He2n)
    rp = He2n.rowval
    ci = He2n.colptr
    Neighbs = Vector{Vector{Int64}}()
    for i = 1:n
        # Get list of edges that node i is in
        Edges = rp[ci[i]:ci[i+1]-1]

        # build a set of neighboring nodes
        NeighbSet = Vector{Int64}()
        for j = 1:length(Edges)
            ej = Edges[j]               # For each node ej adjacent to I...
            nodes = getnodes(Hn2e,ej)   # ...get nodes in edge ej
            append!(NeighbSet,nodes)
        end
        push!(Neighbs, sort(setdiff(unique(NeighbSet),i)))
    end

    return Neighbs
end

function NeighborList(node2edge::Vector{Vector{Int64}},edge2node::Vector{Vector{Int64}})
    """
    NeighborList: given edge to node list and node to edge list, compute
    neighbors
    """
    n = length(node2edge)
    m = length(edge2node)
    Neighbs = Vector{Vector{Int64}}()
    for i = 1:n
        # Get list of edges that node i is in
        Edges = node2edge[i]

        # build a set of neighboring nodes
        NeighbSet = Vector{Int64}()
        for j = 1:length(Edges)
            ej = Edges[j]               # For each node ej adjacent to I...
            nodes = edge2node[ej]       # ...get nodes in edge ej
            append!(NeighbSet,nodes)
        end
        push!(Neighbs, sort(setdiff(unique(NeighbSet),i)))
    end

    return Neighbs
end

function getedges(He2n::SparseArrays.SparseMatrixCSC{Float64,Int64},I::Int64)
    """
    Given the edge-by-node incidence matrix He2n, return the set of hyperedges
    that node I is in.
    """
    first = He2n.colptr[I]
    last = He2n.colptr[I+1]-1
    edges = He2n.rowval[first:last]

    return edges
end

function getnodes(Hn2e::SparseArrays.SparseMatrixCSC{Float64,Int64},J::Int64)
    """
    Given the node-by-edge incidence matrix Hn2e, return the set of nodes
    contained in hyperedge J
    """
    first = Hn2e.colptr[J]
    last = Hn2e.colptr[J+1]-1
    nodes = Hn2e.rowval[first:last]
    mult = Hn2e.nzval[first:last]
    edge = Vector{Int64}()
    # need to adjust for multiplicities
    for t = 1:length(nodes)
        node = nodes[t]
        for k = 1:mult[t]
            push!(edge,node)
        end
    end

    return edge
end

function incidence2elist(H::SparseArrays.SparseMatrixCSC{Float64,Int64},nodelist::Bool=false)
    """
    This converts a hypergraph in incidence matrix form to hyperedge list form.
    Incidence matrix form:  H[e,u] = 1  iff node u is in hyperedge e
    Edgelist form: Hyperedges[j] = array of nodes in hyperedge j

    nodelist == true: means you actually want to get a map from node IDs to hyperedge ids
    """
    if ~nodelist
        # unless you want the node2edge map, transpose first
        H = SparseArrays.sparse(H')
    end
    rp = H.rowval
    ci = H.colptr
    nz = H.nzval
    Hyperedges = Vector{Vector{Int64}}()
    n,m = size(H)

    for i = 1:m
        startedge = ci[i]
        endedge = ci[i+1]-1
        nodes = rp[startedge:endedge]
        mult = nz[startedge:endedge]
        edge = Vector{Int64}()
        # need to adjust for multiplicities
        for t = 1:length(nodes)
            node = nodes[t]
            for k = 1:mult[t]
                push!(edge,node)
            end
        end

        push!(Hyperedges,edge)
    end
    return Hyperedges
end

function elist2incidence(Hyperedges::Vector{Vector{Int64}}, N::Int64)
    """
    Converts a hyperedge list into a binary incidence matrix for the hypergraph.

    This is the exact inverse of incidence2elist
    """
    U = Vector{Int64}()
    E = Vector{Int64}()
    M = length(Hyperedges)
    for enum = 1:length(Hyperedges)
        e = Hyperedges[enum]
        for node in e
            push!(U,node)
            push!(E,enum)
        end
    end

    H = SparseArrays.sparse(E,U,ones(length(U)),M,N)
    return H
end

function hyperedge_formatting(H::hypergraph)
    """
    Convert a hypergraph from dictionary-based format (type 'hypergraph')
    to Hyperedge list with weights vector
    """
    Hyperedges = Vector{Vector{Int64}}()
    weights = Vector{Float64}()
    for key in keys(H.E)
        ed = collect(keys(H.E[key]))
        wt = collect(values(H.E[key]))
        append!(weights,wt)
        append!(Hyperedges,ed)
    end
    return Hyperedges, weights
end

function hypergraph2incidence(H::hypergraph)
    """
    Convert dictionary hypergraph format to binary edge-by-node incidence matrix
    """
    Hyperedges = Vector{Vector{Int64}}()
    weights = Vector{Float64}()
    for key in keys(H.E)
        ed = collect(keys(H.E[key]))
        wt = collect(values(H.E[key]))
        append!(weights,wt)
        append!(Hyperedges,ed)
    end
    N = length(H.D)
    He2n = elist2incidence(Hyperedges,N)
    return He2n, weights
end

function EdgeMap(H::hypergraph)
    """
    Given a hypergraph, extract a map from node indices to hyperedge indices
    """
    n = length(H.D)
    Hyp, weights = hyperedge_formatting(H)
    Hin = elist2incidence(Hyp,n)

    return incidence2elist(Hin,true)
end


__precompile__()

module HemirealNumbers

import Base: +, -, *, /, \, ^, abs, abs2, conj, convert, isfinite, promote_op, promote_rule, real, show, zero
import Base.LinAlg: A_mul_Bt, A_mul_Bc, At_mul_B, Ac_mul_B, At_mul_Bt, Ac_mul_Bc

export PureHemi, Hemireal, μ, ν, mu, nu

immutable PureHemi{T<:Real} <: Number
    m::T
    n::T
end
PureHemi(m::Real, n::Real) = PureHemi(promote(m,n)...)

immutable Hemireal{T<:Real} <: Number
    r::T
    h::PureHemi{T}
end

Hemireal(r::Real, m::Real, n::Real) = Hemireal(r, PureHemi(m, n))
function Hemireal{R,H}(r::R, h::PureHemi{H})
    T = promote_type(R,H)
    Hemireal{T}(r, PureHemi{T}(h.m, h.n))
end
Hemireal{R}(h::PureHemi{R}) = Hemireal(zero(R), h)
Hemireal{R}(r::R) = Hemireal(r, zero(PureHemi{R}))

const μ = PureHemi(true,false)
const ν = PureHemi(false,true)

## PureHemi implementation
convert{R}(::Type{PureHemi{R}}, x::PureHemi) = PureHemi{R}(x.m, x.n)
convert{R}(::Type{PureHemi{R}}, x::Real) = x == 0 ? zero(PureHemi{R}) : throw(DomainError())  # error("Non-zero reals cannot be converted to pure-hemi numbers")

(-)(x::PureHemi) = PureHemi(-x.m, -x.n)

(+)(x::PureHemi, y::PureHemi) = PureHemi(x.m+y.m, x.n+y.n)
(-)(x::PureHemi, y::PureHemi) = PureHemi(x.m-y.m, x.n-y.n)
(*)(x::PureHemi, y::PureHemi) = x.m*y.n + x.n*y.m
(*)(c::Bool, x::PureHemi)     = PureHemi(c*x.m, c*x.n)
(*)(c::Real, x::PureHemi)     = PureHemi(c*x.m, c*x.n)
(*)(x::PureHemi, c::Real)     = c*x
(/)(x::PureHemi, c::Real)     = PureHemi(x.m/c, x.n/c)
(\)(c::Real, x::PureHemi)     = x/c
function (/)(x::PureHemi, y::PureHemi)
    @assert x.m*y.n == x.n*y.m
    y.n == 0 ? x.m/y.m : x.n/y.n
end
(\)(x::PureHemi, y::PureHemi) = y/x
(^)(x::PureHemi, p::Integer)  = (x*x)^(p/2)  # note x*x is much faster than x^2
(^)(x::PureHemi, p::Rational) = (x*x)^(p/2)
(^)(x::PureHemi, p::Real)     = (x*x)^(p/2)

# Symmetric division (there are other solutions y to y*x = c)
(/)(c::Real, x::PureHemi) = PureHemi(c/(2*x.n), c/(2*x.m))
(\)(x::PureHemi, c::Real) = c/x

real{R}(::Type{PureHemi{R}}) = R
real{R}(x::PureHemi{R}) = zero(R)
conj(x::PureHemi) = x
mu(x::PureHemi) = x.m
nu(x::PureHemi) = x.n

zero{R}(::Type{PureHemi{R}}) = PureHemi{R}(0, 0)
zero{R}(::PureHemi{R}) = zero(PureHemi{R})

isfinite(x::PureHemi) = isfinite(x.m) && isfinite(x.n)
# I'm not entirely certain whether the next two should even be defined.
# (note that abs2(x) != x*x)
abs2(x::PureHemi) = x.m*x.m + x.n*x.n
abs(x::PureHemi) = sqrt(abs2(x))

## Hemireal implementation
convert{R}(::Type{Hemireal{R}}, x::Hemireal) = Hemireal{R}(x.r, x.h)
convert{R}(::Type{Hemireal{R}}, x::PureHemi) = Hemireal{R}(0, x)
convert{R}(::Type{Hemireal{R}}, x::Real) = Hemireal{R}(x, zero(PureHemi{R}))
convert{R}(::Type{Hemireal{R}}, r::Real, m::Real, n::Real) = Hemireal{R}(r, PureHemi{R}(m, n))

(-)(x::Hemireal) = Hemireal(-x.r, -x.h)

(+)(x::Hemireal, y::Hemireal) = Hemireal(x.r+y.r, x.h+y.h)
(-)(x::Hemireal, y::Hemireal) = Hemireal(x.r-y.r, x.h-y.h)
(*)(x::Hemireal, y::Hemireal) = Hemireal(x.r*y.r + x.h*y.h, x.r*y.h + x.h*y.r)

(+)(x::Hemireal, y::PureHemi) = x+Hemireal(y)
(+)(x::PureHemi, y::Hemireal) = y+x
(-)(x::Hemireal, y::PureHemi) = x-Hemireal(y)
(-)(x::PureHemi, y::Hemireal) = Hemireal(x)-y
(*)(x::Hemireal, y::PureHemi) = x*Hemireal(y)
(*)(x::PureHemi, y::Hemireal) = y*x

(+)(x::Hemireal, y::Real) = x+Hemireal(y)
(+)(x::Real, y::Hemireal) = y+x
(-)(x::Hemireal, y::Real) = x-Hemireal(y)
(-)(x::Real, y::Hemireal) = Hemireal(x)-y
(*)(c::Bool, x::Hemireal) = Hemireal(c*x.r, c*x.h)
(*)(c::Real, x::Hemireal) = Hemireal(c*x.r, c*x.h)
(*)(x::Hemireal, c::Real) = c*x
(/)(x::Hemireal, c::Real) = Hemireal(x.r/c, x.h/c)
(\)(c::Real, x::Hemireal) = x/c

real{R}(::Type{Hemireal{R}}) = R
real{R}(x::Hemireal{R}) = x.r
conj(x::Hemireal) = x
mu(x::Hemireal) = mu(x.h)
nu(x::Hemireal) = nu(x.h)

zero{R}(::Type{Hemireal{R}}) = Hemireal{R}(0, PureHemi{R}(0, 0))
zero{R}(::Hemireal{R}) = zero(Hemireal{R})

isfinite(x::Hemireal) = isfinite(x.r) && isfinite(x.h)
# ?
abs2(x::Hemireal) = x.r*x.r + abs2(x.h)
abs(x::Hemireal) = sqrt(abs2(x))

for f in (:mu, :nu)
    @eval begin
        function $f(X::AbstractArray)
            out = Array(real(eltype(X)), size(X))
            for I in eachindex(X)
                out[I] = $f(X[I])
            end
            out
        end
    end
end

(*){H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVector{H2}) = A_mul_B!(Array(promote_type(real(H1),real(H2)), size(A,1)), A, B)
(*){H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedMatrix{H2}) = A_mul_B!(Array(promote_type(real(H1),real(H2)), size(A,1), size(B,2)), A, B)
A_mul_Bt{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVecOrMat{H2}) = A_mul_Bt!(Array(promote_type(real(H1),real(H2)), size(A,1), size(B,1)), A, B)
A_mul_Bc{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVecOrMat{H2}) = A_mul_Bt!(Array(promote_type(real(H1),real(H2)), size(A,1), size(B,1)), A, B)
At_mul_B{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVector{H2}) = At_mul_B!(Array(promote_type(real(H1),real(H2)), size(A,2)), A, B)
At_mul_B{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedMatrix{H2}) = At_mul_B!(Array(promote_type(real(H1),real(H2)), size(A,2), size(B,2)), A, B)
Ac_mul_B{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVector{H2}) = At_mul_B!(Array(promote_type(real(H1),real(H2)), size(A,2)), A, B)
Ac_mul_B{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedMatrix{H2}) = At_mul_B!(Array(promote_type(real(H1),real(H2)), size(A,2), size(B,2)), A, B)
At_mul_Bt{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVecOrMat{H2}) = At_mul_Bt!(Array(promote_type(real(H1),real(H2)), size(A,2), size(B,1)), A, B)
Ac_mul_Bc{H1<:PureHemi,H2<:PureHemi}(A::StridedVecOrMat{H1}, B::StridedVecOrMat{H2}) = At_mul_Bt!(Array(promote_type(real(H1),real(H2)), size(A,2), size(B,1)), A, B)

promote_op{H1,H2}(::Base.MulFun, ::Type{PureHemi{H1}}, ::Type{PureHemi{H2}}) = promote_type(H1,H2)
promote_op{H1,H2}(::Base.DotMulFun, ::Type{PureHemi{H1}}, ::Type{PureHemi{H2}}) = promote_type(H1,H2)
promote_op{R<:Real,H}(::Base.MulFun, ::Type{R}, ::Type{PureHemi{H}}) = PureHemi{promote_type(R,H)}
promote_op{R<:Real,H}(::Base.DotMulFun, ::Type{R}, ::Type{PureHemi{H}}) = PureHemi{promote_type(R,H)}
promote_op{R<:Real,H}(::Base.MulFun, ::Type{PureHemi{H}}, ::Type{R}) = promote_op(Base.MulFun(), R, PureHemi{H})
promote_op{R<:Real,H}(::Base.DotMulFun, ::Type{PureHemi{H}}, ::Type{R}) = promote_op(Base.MulFun(), R, PureHemi{H})

promote_rule{H}(::Type{Bool}, ::Type{PureHemi{H}}) = Hemireal{promote_type(Bool,H)}
promote_rule{R,H}(::Type{Irrational{R}}, ::Type{PureHemi{H}}) = Hemireal{promote_type(Irrational{R},H)}
promote_rule{R<:Real,H}(::Type{R}, ::Type{PureHemi{H}}) = Hemireal{promote_type(R,H)}
promote_rule{H1,H2}(::Type{PureHemi{H1}}, ::Type{PureHemi{H2}}) = PureHemi{promote_type(H1,H2)}
promote_rule{H}(::Type{Bool}, ::Type{Hemireal{H}}) = Hemireal{promote_type(Bool,H)}
promote_rule{R,H}(::Type{Irrational{R}}, ::Type{Hemireal{H}}) = Hemireal{promote_type(Irrational{R},H)}
promote_rule{R<:Real,H}(::Type{R}, ::Type{Hemireal{H}}) = Hemireal{promote_type(R,H)}
promote_rule{H1,H2}(::Type{PureHemi{H1}}, ::Type{Hemireal{H2}}) = Hemireal{promote_type(H1,H2)}
promote_rule{H1,H2}(::Type{Hemireal{H1}}, ::Type{Hemireal{H2}}) = Hemireal{promote_type(H1,H2)}

show(io::IO, x::PureHemi) = x.n >= 0 ? print(io, x.m, "μ + ", x.n, 'ν') : print(io, x.m, "μ - ", abs(x.n), 'ν')
show(io::IO, x::Hemireal) = x.h.m >= 0 ? print(io, x.r, " + ", x.h) : print(io, x.r, " - ", PureHemi(-x.h.m, x.h.n))

end # module


using ChainRulesCore
using LinearAlgebra
using KrylovKit

@doc raw"
    num_grad(f, K::Real; [δ = 1e-5])

return the numerical gradient of `f` at `K` calculated with
`(f(K+δ/2) - f(K-δ/2))/δ`

# example

```jldoctest; setup = :(using TensorNetworkAD)
julia> TensorNetworkAD.num_grad(x -> x * x, 3) ≈ 6
true
```
"
function num_grad(f, K; δ::Real=1e-5)
    if eltype(K) == ComplexF64
        (f(K + δ / 2) - f(K - δ / 2)) / δ +
            (f(K + δ / 2 * 1.0im) - f(K - δ / 2 * 1.0im)) / δ * 1.0im
    else
        (f(K + δ / 2) - f(K - δ / 2)) / δ
    end
end

@doc raw"
    num_grad(f, K::AbstractArray; [δ = 1e-5])

return the numerical gradient of `f` for each element of `K`.

# example

```jldoctest; setup = :(using TensorNetworkAD, LinearAlgebra)
julia> TensorNetworkAD.num_grad(tr, rand(2,2)) ≈ I
true
```
"
function num_grad(f, a::AbstractArray; δ::Real=1e-5)
    b = Array(copy(a))
    df = map(CartesianIndices(b)) do i
        foo = x -> (ac = copy(b); ac[i] = x; f(_arraytype(a)(ac)))
        num_grad(foo, b[i], δ=δ)
    end
    return _arraytype(a)(df)
    # map(CartesianIndices(a)) do i
    #     foo = x -> (ac = copy(a); ac[i] = x; f(ac))
    #     num_grad(foo, a[i], δ=δ)
    # end
end

# patch since it's currently broken otherwise
function ChainRulesCore.rrule(::typeof(Base.typed_hvcat), ::Type{T}, rows::Tuple{Vararg{Int}}, xs::S...) where {T,S}
    y = Base.typed_hvcat(T, rows, xs...)
    function back(ȳ)
        return NoTangent(), NoTangent(), NoTangent(), permutedims(ȳ)...
    end
    return y, back
end

# improves performance compared to default implementation, also avoids errors
# with some complex arrays
function ChainRulesCore.rrule(::typeof(LinearAlgebra.norm), A::AbstractArray)
    n = norm(A)
    function back(Δ)
        return NoTangent(), Δ .* A ./ (n + eps(0f0)), NoTangent()
    end
    return n, back
end

"""
    ChainRulesCore.rrule(::typeof(leftenv), Au::AbstractArray{T}, Ad::AbstractArray{T}, M::AbstractArray{T}, FL::AbstractArray{T}; kwargs...) where {T}

```
           ┌──   Au  ──┐
           │     │     │
dM    = - FL ──    ──  ξl
           │     │     │
           └──   Ad  ──┘

           ┌──       ──┐
           │     │     │
dAu  = -  FL ──  M ──  ξl
           │     │     │
           └──   Ad  ──┘

           ┌──   Au  ──┐
           │     │     │
dAd  = -  FL  ── M ──  ξl
           │     │     │
           └──       ──┘
```
"""

function ChainRulesCore.rrule(::typeof(leftenv), Au::AbstractArray{T}, Ad::AbstractArray{T}, M::AbstractArray{T}, FL::AbstractArray{T}; kwargs...) where {T}
    λl, FL = leftenv(Au, Ad, M, FL)
    # @show λl
    function back((dλ, dFL))
        dFL -= Array(ein"abc,abc ->"(conj(FL), dFL))[] * FL
        ξl, info = linsolve(FR -> ein"((abc,ceh),dgeb),fgh -> adf"(Au, FR, M, Ad), conj(dFL), -λl, 1; maxiter = 1)
        @assert info.converged==1
        # errL = ein"abc,cba ->"(FL, ξl)[]
        # abs(errL) > 1e-1 && throw("FL and ξl aren't orthometric. err = $(errL)")
        dAu = -ein"((adf,fgh),dgeb),ceh -> abc"(FL, Ad, M, ξl)
        dAd = -ein"((adf,abc),dgeb),ceh -> fgh"(FL, Au, M, ξl)
        dM = -ein"(adf,abc),(fgh,ceh) -> dgeb"(FL, Au, Ad, ξl)
        return NoTangent(), conj(dAu), conj(dAd), conj(dM), NoTangent()...
    end
    return (λl, FL), back
end

"""
    function ChainRulesCore.rrule(::typeof(norm_FL), Au::AbstractArray{T}, Ad::AbstractArray{T}, FL::AbstractArray{T}; kwargs...) where {T}

```
           ┌──       ──┐
           │     │     │
dAu  = -  FL ────┼──── ξl
           │     │     │
           └──   Ad  ──┘

           ┌──   Au  ──┐
           │     │     │
dAd  = -  FL  ───┼───  ξl
           │     │     │
           └──       ──┘
```
"""

function ChainRulesCore.rrule(::typeof(norm_FL), Au::AbstractArray{T}, Ad::AbstractArray{T}, FL::AbstractArray{T}; kwargs...) where {T}
    λl, FL = norm_FL(Au, Ad, FL)
    function back((dλ, dFL))
        dFL -= Array(ein"ab,ab ->"(conj(FL), dFL))[] * FL
        ξl, info = linsolve(FR -> ein"(be,acb), dce -> ad"(FR,Au,Ad), conj(dFL), -λl, 1; maxiter = 1)
        # @assert info.converged==1
        # errL = ein"abc,cba ->"(FL, ξl)[]
        # abs(errL) > 1e-1 && throw("FL and ξl aren't orthometric. err = $(errL)")
        dAu = -ein"(ad,dce), be -> acb"(FL, Ad, ξl)
        dAd = -ein"(ad,acb), be -> dce"(FL, Au, ξl)
        return NoTangent(), conj(dAu), conj(dAd), NoTangent()...
    end
    return (λl, FL), back
end


input = readlines()

numbers = parse.(Int, split(input[1], ","))

function check_victory(m, drawn)
    if maximum(sum(drawn; dims=1)) == 5 || maximum(sum(drawn; dims=2)) == 5
        (@. m * !drawn) |> sum
    else
        0
    end
end

function solve()
    best = []
    for i in 1:100
        m = [
            parse.(Int, split(input[3 + (i-1)*6]));;
            parse.(Int, split(input[4 + (i-1)*6]));;
            parse.(Int, split(input[5 + (i-1)*6]));;
            parse.(Int, split(input[6 + (i-1)*6]));;
            parse.(Int, split(input[7 + (i-1)*6]))
        ]

        drawn = falses(size(m))

        for (j, n) in enumerate(numbers)
            pos = findnext(==(n), m, CartesianIndex(1, 1))
            if !isnothing(pos)
                drawn[pos] = true
            end
            score = check_victory(m, drawn)
            if score > 0
                score *= n
                push!(best, (j, score))
                break
            end
        end
    end
    sort!(best)
    best[begin], best[end]
end

# Part 1 and 2
solve() .|> println


using IterativeSolvers
using SparseArrays
using Random
using Printf

#########################
# Method Implementation #
#########################

@inline function omega(t, s)
    angle = sqrt(2.)/2
    ns = norm(s)
    nt = norm(t)
    ts = dot(t,s)
    rho = abs(ts/(nt*ns))
    om = ts/(nt*nt)
    if rho < angle
        om = om*convert(typeof(om),angle)/rho
    end
    om
end

#s = 8, , ,
function sidrs!(X0, A, C::T; s::Number = 4, tol=missing, droptol=(eps(real(eltype(C)))), maxiter=size(A, 2), verbose::Bool=false) where {T}
	if ismissing(tol)
		tol = 10sqrt(eps(real(eltype(C))))
	end
	verbose && @show tol
	X = X0

    verbose && @printf("=== idrs ===\n%4s\t%7s\n","iter","resnorm")
    R = C - A(X)
    normR = norm(R)
	iter = 1

    if normR <= tol # Initial guess is a good enough solution
        return X
    end

    Z = zero(C)
	d = size(C,1)
    P = T[Diagonal(rand(d) .- 0.5) + sprandn(d,d, 1/d) for k in 1:s]
    U = T[copy(C) for k in 1:s]
    G = T[copy(C) for k in 1:s] #
    Q = copy(Z)
    V = copy(Z)

    M = Matrix{eltype(C)}(I,s,s)
    f = zeros(eltype(C),s)
    c = zeros(eltype(C),s)

    om::eltype(C) = 1
    while normR > tol && iter ≤ maxiter
        for i in 1:s
            f[i] = dot(P[i], R)
        end
        for k in 1:s

            # Solve small system and make v orthogonal to P

            c = LowerTriangular(M[k:s,k:s])\f[k:s]
            V = c[1] * G[k]
            Q = c[1] * U[k]

            for i = k+1:s
                V += c[i-k+1] * G[i]
                Q += c[i-k+1] * U[i]
            end

            # Compute new U[:,k] and G[:,k], G[:,k] is in space G_j
            V = R - V

            U[k] = Q + om * V
            G[k] = A(U[k])

            # Bi-orthogonalise the new basis vectors

            for i in 1:k-1
                alpha = dot(P[i],G[k])/M[i,i]
                G[k] -= alpha * G[i]
                U[k] -= alpha * U[i]
            end

            # New column of M = P'*G  (first k-1 entries are zero)

            for i in k:s
                M[i,k] = dot(P[i],G[k])
            end

            #  Make r orthogonal to q_i, i = 1..k

            beta = f[k]/M[k,k]
            R -= beta * G[k]
            X += beta * U[k]

			droptol!(X, droptol)

            normR = norm(R)
            verbose && @printf("%3d\t%1.2e\t%1.2e\n",iter,normR,sparsity(X))
            if normR < tol || iter == maxiter
                iter == maxiter && warn("not converged")
                break
            end
            if k < s
                f[k+1:s] .-=  beta*M[k+1:s,k]
            end
            iter += 1
        end # k in 1:s

		#@show sparsity(R), sparsity(X), sparsity(Q)

        # Now we have sufficient vectors in G_j to compute residual in G_j+1
        # Note: r is already perpendicular to P so v = r
        copyto!(V, R)
        Q = A(V)
        om = omega(Q, R)
        R -= om .* Q
        X += om .* V
		droptol!(X, droptol)

        normR = norm(R)

        iter += 1
    end
    verbose && print("\n")
    X
end

slyap_op(B) = X -> B*X + X*B'
linl_op(B) = X -> B*X
linr_op(B) = X -> X*B
slyapunov(B, C; tol=missing, droptol = 1e-7) = sidrs!(copy(C), slyap_op(B), -C; maxiter=prod(size(B)), verbose=true, tol=tol, droptol = droptol)
slsolve(B, Y; tol=missing, droptol = 1e-7) = sidrs!(copy(Y), linl_op(B), Y; maxiter=prod(size(B)), verbose=true, tol=tol, droptol = droptol)
srsolve(B, Y; tol=missing, droptol = 1e-7) = sidrs!(copy(Y), linr_op(B), Y; maxiter=prod(size(B)), verbose=true, tol=tol, droptol = droptol)

"""
Compute lyaponov equation using the Kronecker product for testing.
"""
kronlyap(B,C) = -reshape(inv(kron(B, Diagonal(I,size(B,1))) + kron(Diagonal(I,size(B,1)), B))*vec(C), size(C))


using Gen
using PyCall
using GenRedner: get_depth_renderer

plt = pyimport("matplotlib.pyplot")

const depth_renderer = get_depth_renderer(num_samples=2, print_timing=true)

function construct_simple_scene()

    vertices = vcat(

        # first foreground triangle (closest to camera)
        [-0.25, -0.25, 1.0]', # 1
        [0.25, -0.25, 1.0]', # 2
        [0.0, 0.25, 1.0]', # 3

        # second foreground triangle
        [-1, -1, 1.5]', # 4
        [1, -1, 1.5]', # 5
        [0.0, 1, 1.5]', # 6

        # two triangles forming a background
        [-2, -2, 2.0]', # 7
        [-2, 2, 2.0]', # 8
        [2, 2, 2.0]', # 9
        [2, -2, 2.0]' # 10
    )

    indices = vcat(
        [1, 2, 3]', # first triangle
        [4, 5, 6]', # second triangle
        [7, 8, 9]', # background triangle 1
        [9, 10, 7]' # background triangle 2
    )

    return (vertices, indices)
end

function construct_corrupted_scene()
    (vertices, indices) = construct_simple_scene()
    corrupted_vertices = copy(vertices)

    # offset one vertex of the front triangle by 0.1 units
    corrupted_vertices[1,:] += [-0.1, 0.0, 0.0]
    return corrupted_vertices, indices
end

# get scene and corrupted scene
vertices, indices = construct_simple_scene()
corrupted_vertices, corrupted_indices = construct_corrupted_scene()

# focal lengths
fx = 48.0
fy = 48.0

# principal point
cx = 32.0
cy = 24.0

# image resolution
width = 64
height = 48

intrinsics = [fx, fy, cx, cy]
dims = [width, height]

function forward_example()

    # run forward three times (to show variability) and plot using matplotlib
    plt.figure(figsize=(18, 27))

    for replicate in [1, 2, 3]

        trace = simulate(depth_renderer, (vertices, indices, intrinsics, dims))
        points = get_retval(trace)

        plt.subplot(3, 2, (replicate-1) * 2 + 1)
        plt.imshow(points[:,:,3], cmap="Greys", vmin=0, vmax=2.0, origin="lower")
        plt.title("depth image")

        plt.subplot(3, 2, (replicate-1) * 2 + 2, projection="3d")
        ax = plt.gca()
        ax.scatter(points[:,:,1], points[:,:,2], points[:,:,3], s=2)
        ax.scatter([0.0], [0.0], [0.0], color="red", s=10, marker="o")
        ax.quiver([0.0], [0.0], [0.0], [1.0], [0.0], [0.0], length=0.5, color="red")
        ax.quiver([0.0], [0.0], [0.0], [0.0], [1.0], [0.0], length=0.5, color="green")
        ax.quiver([0.0], [0.0], [0.0], [0.0], [0.0], [1.0], length=0.5, color="blue")
        ax.set_xlim([-2, 2])
        ax.set_ylim([-2, 2])
        ax.set_zlim([-0, 2])
        plt.title("point cloud")

    end

    plt.tight_layout()
    plt.savefig("forward.png")

end

forward_example()

struct NoiseModel <: Gen.Distribution{Matrix{Float64}} end
const noise_model = NoiseModel()

function Gen.logpdf(::NoiseModel,
        x::Matrix{Float64},
        mu::Matrix{Float64})
    diffs = (x .- mu)
    return -sum(diffs .^ 2)
end

Gen.has_argument_grads(::NoiseModel) = (true,)
Gen.has_output_grad(::NoiseModel) = false

function Gen.logpdf_grad(::NoiseModel,
        x::Matrix{Float64},
        mu::Matrix{Float64})
    diffs = x .- mu
    mu_grad = 2 * diffs
    return (nothing, mu_grad)
end

@gen function model()

    # prior on two foreground triangles
    foreground_vertices_list = []
    for i in 1:6
        v = ({:vertices => i} ~ mvnormal([0.0, 0.0, 1.0], [1000.0 0.0 0.0; 0.0 1000.0 0.0; 0.0 0.0 1000.0]))
        push!(foreground_vertices_list, v)
    end
    foreground_vertices_mat = copy(transpose(hcat(foreground_vertices_list...)))
    foreground_indices = vcat(
        [1, 2, 3]', # first triangle
        [4, 5, 6]', # second triangle
    )

    # fixed background
    background_vertices = vcat(
        # two triangles forming a background
        [-2, -2, 2.0]', # 7
        [-2, 2, 2.0]', # 8
        [2, 2, 2.0]', # 9
        [2, -2, 2.0]' # 10
    )
    background_indices = vcat(
        [7, 8, 9]', # background triangle 1
        [9, 10, 7]' # background triangle 2
    )

    points ~ depth_renderer(
        vcat(foreground_vertices_mat, background_vertices),
        vcat(foreground_indices, background_indices),
        intrinsics, dims)

    depths = points[:,:,3]
    obs ~ noise_model(depths)
end

function visualize_trace(fname, trace)
    plt.clf()
    points = trace[:points]
    plt.figure(figsize=(27, 8))

    plt.subplot(1, 5, 1)
    plt.title("observed depth image")
    plt.imshow(trace[:obs], cmap="Greys", vmin=0, vmax=2.0, origin="lower")

    plt.subplot(1, 5, 2)
    plt.title("reconstructed depth image")
    plt.imshow(points[:,:,3], cmap="Greys", vmin=0.0, vmax=2.0, origin="lower")

    plt.subplot(1, 5, 3)
    plt.title("observed .- latent diff")
    plt.imshow(trace[:obs] .- points[:,:,3], cmap="RdBu", vmin=-1.0, vmax=1.0, origin="lower")

    plt.subplot(1, 5, 4)
    plt.title("gradient with respect to depth")
    plt.imshow(logpdf_grad(noise_model, trace[:obs], points[:,:,3])[2], cmap="RdBu", vmin=-1.0, vmax=1.0, origin="lower")

    plt.subplot(1, 5, 5, projection="3d")
    plt.title("latent point cloud")
    ax = plt.gca()
    ax.scatter(points[:,:,1], points[:,:,2], points[:,:,3], s=2)
    ax.scatter([0.0], [0.0], [0.0], color="red", s=10, marker="o")
    ax.quiver([0.0], [0.0], [0.0], [1.0], [0.0], [0.0], length=0.5, color="red")
    ax.quiver([0.0], [0.0], [0.0], [0.0], [1.0], [0.0], length=0.5, color="green")
    ax.quiver([0.0], [0.0], [0.0], [0.0], [0.0], [1.0], length=0.5, color="blue")
    ax.set_xlim([-2, 2])
    ax.set_ylim([-2, 2])
    ax.set_zlim([0, 2])

    for t in 1:2
        x1, y1, z1 = trace[:vertices => indices[t,1]]
        x2, y2, z2 = trace[:vertices => indices[t,2]]
        x3, y3, z3 = trace[:vertices => indices[t,3]]
        ax.plot([x1, x2], [y1, y2], [z1, z2], color="black", linewidth=2)
        ax.plot([x3, x2], [y3, y2], [z3, z2], color="black")
        ax.plot([x1, x3], [y1, y3], [z1, z3], color="black")
    end

    (_, _, choice_grads) = choice_gradients(trace, select(:vertices))
    for v in 1:6
        x, y, z = trace[:vertices => v]
        grad = choice_grads[:vertices => v]
        grad = grad / 10
        xstop, ystop, zstop = [x, y, z] .+ grad
        ax.plot([x, xstop], [y, ystop], [z, zstop], color="purple")
    end

    plt.tight_layout()
    plt.savefig(fname)
end

using Printf: @sprintf
using Random: seed!

function do_inference()
    seed!(1) # note: this does not affect Redner's random seed

    num_latent_vertices = 6

    observed = depth_renderer(vertices, indices, intrinsics, dims)[:,:,3]
    constraints = choicemap()
    constraints[:obs] = observed

    # initialize to foreground vertices that are slightly wrong
    for i in 1:6
        constraints[:vertices => i] = corrupted_vertices[i,:]
    end
    trace, = generate(model, (), constraints)
    frame = 0

    fname = @sprintf("iter_%03d.png", frame); frame += 1
    visualize_trace(fname, trace)
    iters = 500
    for iter in 1:iters
        println("iter: $iter")

        # stochastic gradient ascent for MAP estimation
        (_, _, grads) = choice_gradients(trace, select((:vertices => v for v in 1:num_latent_vertices)...))
        constraints = choicemap()
        for v in 1:num_latent_vertices
            constraints[:vertices => v] = trace[:vertices => v] .+ (0.000001) * grads[:vertices => v]
        end

        trace, _ = update(trace, (), (), constraints)
        if iter % 10 == 0
            fname = @sprintf("iter_%03d.png", frame); frame += 1
            visualize_trace(fname, trace)
        end
        plt.clf()

       println("iter; $iter of $iters")#,  acc3: $acc3")

    end
    fname = @sprintf("iter_%03d.png", frame); frame += 1
    visualize_trace(fname, trace)
end

do_inference()


"""
Maximal ratio of circumcircle radius and length of shortest edge before  a triangle
is considered bad.
"""
beta = sqrt(2.)

minArea = 1e-5

"""
    compute_circumcenter_radius(mesh::Mesh, p::VertexIndex, r::VertexIndex, q::VertexIndex)

Computes the position and radius of the triangle circumcenter.

# Returns
* cx, cy, rs with ``c_x``, ``c_y`` the position and ``r_s`` the radius square of the circum center.
"""
function compute_circumcenter_radius(mesh::Mesh, p::VertexIndex, r::VertexIndex, q::VertexIndex)
  local pr = mesh.tesselation.vertices[p]
  local qr = mesh.tesselation.vertices[q]
  local rr = mesh.tesselation.vertices[r]
  local ax = pr[1]
  local ay = pr[2]
  local bx = qr[1]
  local by = qr[2]
  local cx = rr[1]
  local cy = rr[2]
  local sx = (ax^2.0*(by-cy) + ay^2.0*(by-cy) + bx^2.0*(cy-ay) + by^2.0*(cy-ay) + cx^2.0*(ay-by) + cy^2.0*(ay-by))/2.0
  local sy = (ax^2.0*(cx-bx) + ay^2.0*(cx-bx) + bx^2.0*(ax-cx) + by^2.0*(ax-cx) + cx^2.0*(bx-ax) + cy^2.0*(bx-ax))/2.0
  local a = ax*(by - cy) + bx * (cy-ay) + cx*(ay-by)
  local b = ax^2.0 * (bx*cy - by * cx) + ay^2.0 * (bx*cy - by*cx) + bx^2.0*(ay*cx-ax*cy) + by^2.0*(ay*cx-ax*cy) + cx^2.0*(ax*by-ay*bx) + cy^2.0*(ax*by-ay*bx)
  local rs = ((b/a + (sx*sx + sy*sy)/(a*a)))
  return sx/a, sy/a, rs
end

"""
    compute_offcenter(mesh:Mesh, p::VertexIndex, q::VertexIndex, r::VertexIndex)

Computes position of offcenter[^1] for given triangle ``p - q - r``.

# Remarks
* ``p - q`` must be the shortest edge of the triangle, i.e. the smallest angle is located at ``r``.

# References
[^1] Alper Üngör (2009). Off-centers: A new type of Steiner points for computing size-optimal quality-guaranteed Delaunay triangulations, Computational Geometry, *42* (2)

![SteinerPoint](SteinerPoint.svg)
"""
function compute_offcenter(mesh::Mesh, p::VertexIndex, q::VertexIndex, r::VertexIndex)
  local pr = mesh.tesselation.vertices[p]
  local qr = mesh.tesselation.vertices[q]
  local rr = mesh.tesselation.vertices[r]

  # position of bisection start of pq and circum center
  local bsr = 0.5 * (pr + qr)
  ccx, ccy, ccrs = compute_circumcenter_radius(mesh, p, q, r)

  # length of pq and radius of offcenter
  local lpq = norm(pr - qr)
  local ocr = beta * lpq * 0.999999

  # distance edge midpoint to center of offcircle ``c_2``
  local dobsoc = sqrt(ocr*ocr - (lpq/2.)*(lpq/2.) )

  # distance edge midpoint to center of outcircle
  local dbscc = norm(bsr - [ccx, ccy])
  if (dbscc <= dobsoc)
    # we take the out circle center
    return [ccx, ccy]
  else
    # we take off center

    # bisection: pq-midpoint
    local c2 = bsr + dobsoc * normalize([ccx, ccy] - bsr)
    local offcenter = c2 + ocr * normalize([ccx, ccy] - bsr)
    return offcenter
  end
end

"""
    vertex_encroaches_segment(mesh::Mesh, ei::Int, vertex::Point)

Checks if a potential vertex encroaches the segment denoted by `ei`.

# Remarks
* If edge `ei` is not a segment, i.e. not marked as a boundary, it always
  returns `false`
"""
function vertex_encroaches_segment(mesh::Mesh, ei::Int, vertex::Point)
    if (!isconstraint(mesh.tesselation, ei))
      return false
    end

    vio = org(mesh.tesselation, ei)
    vid = dest(mesh.tesselation, ei)
    v1 = mesh.tesselation.vertices[vio]
    v2 = mesh.tesselation.vertices[vid]
    mp = 0.5*(v1 + v2)
    dr = vertex - mp # from midpoint to new vertex
    dx = v1 - mp # from midpoint to end point

    return (dr[1]^2 + dr[2]^2) <= (dx[1]^2 + dx[2]^2)
  end

""""
    get_shortest_edge(mesh::Mesh, ai::VertexIndex, bi::VertexIndex, ci::VertexIndex)

Given a triangle with CCW-ordered vertices `ai` - `bi` - `ci`, it computes the
edge length and returns a sorted tuple `p` - `q` - `r` such that the triangle
``p - q - r`` has its shortest edge between `p` and `q`.
"""
function get_shortest_edge(mesh::Mesh, ai::VertexIndex, bi::VertexIndex, ci::VertexIndex)
  # nomenclature as in Bronstein et al..
  xa = mesh.tesselation.vertices[ai]
  xb = mesh.tesselation.vertices[bi]
  xc = mesh.tesselation.vertices[ci]
  vab = (xb - xa) # edge a
  vac = (xc - xa) # edge b
  vbc = (xb - xc) # edge a
  c = norm(vab)
  b = norm(vac)
  a = norm(vbc)
  #cosa = (b'*b + c'*c - a'*a)/(2. * b' * c)
  #cosb = (a'*a + c'*c - b'*b)/(2. * c' * a)
  #cosc = (a'*a + b'*b - c'*c)/(2. * a' * b)
  ordered = [(a, ai, bi, ci), (b, bi, ci, ai), (c, ci, ai, bi)]

  # order in place
  if (ordered[1][1] > ordered[2][1])
    ordered[1], ordered[2] = ordered[2], ordered[1]
  end
  if (ordered[2][1] > ordered[3][1])
    ordered[2], ordered[3] = ordered[3], ordered[2]
  end
  if (ordered[1][1] > ordered[2][1])
    ordered[1], ordered[2] = ordered[2], ordered[1]
  end

  return (ordered[1][2], ordered[1][3], ordered[1][4])
end

 """
    function refine_triangle(mesh::Mesh, tri::Triangle)

Refines the triangle located to the left of edge `ei` by adding a new vertex
at the offcenter or the midpoint of encroached edges.

# Remarks
* If a new vertex at the off-center of the triangle does not encroach any segments
  it is inserted at the off-center.
* Otherwise, a new vertex is inserted at the midpoint of each edge that the
  off-center would encroach on.
  try
"""
function refine_triangle(mesh::Mesh, tri::Triangle)

  # get ordered vertices
  p, q, r = get_shortest_edge(mesh, tri.vertices[1], tri.vertices[2], tri.vertices[3])

  # get offcenter
  offcenter = compute_offcenter(mesh, p, q, r)

  # get indices of all encroached segments
  indices = [i for i in 1:4:length(mesh.tesselation.edges)]
  encroachedIndices = filter(i->vertex_encroaches_segment(mesh, i, offcenter), indices)

  if isempty(encroachedIndices)
    # insert off-cente  r
    # we need to push the *unscaled point* so we push it straight to tesselation
    push_scaled!(mesh, offcenter')
  else
    # insert mid points
    for ei in encroachedIndices
      # get midpoint and insert it
      oi = org(mesh.tesselation, ei)
      di = dest(mesh.tesselation, ei)
      mp = 0.5 * (mesh.tesselation.vertices[oi] + mesh.tesselation.vertices[di])

      # we need to push the *unscaled point* so we push it straight to tesselation
      newVertexIndex = length(mesh.tesselation.vertices) + 1
      push_scaled!(mesh, mp')

      # we need to fix the face locations for the new faces.. we do that by
      # iterating over the star of the new vertex, starting at the segment
      startEdge = mesh.tesselation.vertexCache[newVertexIndex]
      while (!isconstraint(mesh.tesselation, startEdge))
        startEdge = onext(mesh.tesselation, startEdge)
      end
      faceIsInside = mesh.tesselation.edges[startEdge].constraintType == OnRightSide
      curEdge = startEdge
      while (true)
        # mark left face as inside/outside
        mesh.faceLocation[lface(mesh.tesselation, curEdge)] = faceIsInside
        curEdge = onext(mesh.tesselation, curEdge)
        if (curEdge == startEdge)
          break
        elseif (isconstraint(mesh.tesselation, curEdge))
          faceIsInside = !faceIsInside
        end
      end
    end
  end
end

"""
    check_triangle(mesh::Mesh, tri::Triangle)

Checks if the triangle satisfies the quality measure ``r/\bar{pq} \le \beta``,
with ``r`` the radius of the circumcircle and ``\bar{pq}`` the length of its
shortest edge.

# Returns
* `true`, if the triangle satisfies the quality measure.
"""
function check_triangle(mesh::Mesh, tri::Triangle)
  # TODO: we really need to get rid of the double computation of offcenter etc.
  #       maybe put a nullable in Triangle?

  # check if the area is below limit
  ra = mesh.tesselation.vertices[tri.vertices[1]]
  rb = mesh.tesselation.vertices[tri.vertices[2]]
  rc = mesh.tesselation.vertices[tri.vertices[3]]

  xa, ya = ra[1], ra[2]
  xb, yb = rb[1], rb[2]
  xc, yc = rc[1], rc[2]

  area = 0.5*abs((xa-xc)*(yb-ya)-(xa-xb)*(yc-ya))
  if area < minArea
    return true
  end

  # get ordered vertices and circumcircle
  p, q, r = get_shortest_edge(mesh, tri.vertices[1], tri.vertices[2], tri.vertices[3])
  sx, sy, rs = DelaunayMeshes.compute_circumcenter_radius(mesh, 1, 2, 3)

  lpq = norm(mesh.tesselation.vertices[p] - mesh.tesselation.vertices[q])
  return sqrt(rs)/lpq <= DelaunayMeshes.beta
end

"""
    refine_valid_triangles!(mesh::

Gets the first triangle that does not satisfy the quality measure or `null`.
"""
function refine_valid_triangles!(mesh::Mesh)
  foundBadTriangle = true
  counter = 0
  while foundBadTriangle && counter < 10000
    counter = counter + 1
    foundBadTriangle = false

    # get the first valid but bad triangle
    itState = start(mesh.tesselation)
    while (!done(mesh.tesselation, itState))
      tri, itState = next(mesh.tesselation, itState)
      if (mesh.faceLocation[tri.face] && !check_triangle(mesh, tri))

        # the triangle is bad and valid - we need to refine it
        refine_triangle(mesh, tri)
        foundBadTriangle = true

        # break out to outer loop
        break
      end
    end
  end
end


using Flux; Flux.CuArrays.allowscalar(false)
using Flux: params, throttle, Optimise.@epochs
using SparseGaussianProcesses
import Plots

gp = SparseGaussianProcess(SquaredExponentialKernel(1), inducing_points = PseudoDataInducingPoints)

f(x) = 2 * sin.(x |> vec) + randn(length(x))/10
x = reshape(-5:0.1:5, (1,:)) |> collect
y = f(x)' |> collect
gp.inducing_points(reshape(-5:1:5,1,:), gp.kernel)

(gp,x,y) = gpu.((gp,x,y))

rand!(gp.prior_basis, gp.kernel; num_features=67)
rand!(gp; num_samples=17)
loss(gp,x,y)|>sum

opt = ADAM(0.01)
dataset = Iterators.repeated((x,y), 200)

Flux.train!((x,y) -> loss(gp,x,y)|>sum, params(gp), dataset, opt; cb = throttle(() -> @show(loss(gp,x,y)|>sum), 0.1))

(gp,x,y) = cpu.((gp,x,y))

plot_gp_intervals(gp, x, y)


#### Code and functions specific to the Yodzis-Innes metabolic model
#### "Slow organisms exhibit sudden population disappearances in a reddened world" by Greyson-Gaito, Gellner, & McCann.

@with_kw mutable struct YodInnScalePar
    aₜ = 54.9
    aᵣ = 34.3
    rat = 10^-1
    δ = 0.55
    fⱼ = 0.99
    aⱼ = 89.2
    k = 3.0
    R₀ = 0.8
    ε = 0.1
    x = ( aₜ / aᵣ ) * rat^0.25
    y = fⱼ * aⱼ / aₜ
end

function yod_inn_scale_II!(du, u, p, t)
    @unpack aₜ, aᵣ, y, δ, k, R₀, ε, rat = p
    fₑ = ( (rat * ε) ^ 0.25 ) / ( rat ^ 0.25 )
    xₛ = ( aₜ / aᵣ ) * (rat * ε)^0.25
    R, C = u
    du[1] = R * (1 - R / k) -  ( xₛ * y / ( (1 - δ) * fₑ ) )  * R * C  / (R₀ + R)
    du[2] = ( xₛ * y * R * C  / (R₀ + R) ) - xₛ * C
    return
end #setting up the Yodzis-Innes model

function yod_inn_scale_II(u, par)
    du = similar(u)
    yod_inn_scale_II!(du, u, par, 0.0)
    return du
end #setting up the Yodzis-Innes model

function eq_yodinn_II(p)
    @unpack x, y, δ, k, R₀ = p
    eq_II_R = R₀ / (y - 1)
    eq_II_C = ( ( 1 - δ ) / x ) * eq_II_R * (1 - eq_II_R / k)
    return vcat(eq_II_R, eq_II_C)
end #calculating the interior equilibrium of the Yodzis-Innes model

function con_iso_yodinn(p)
    @unpack R₀, y = p
    R₀/(y - 1)
end #calculating the consumer isocline of the Yodzis-Innes model

function res_iso_yodinn(R, p)
    @unpack x, y, δ, R₀, k  = p
    (R*(R*δ - R + R₀*δ - R₀) + k*(-R*δ + R - R₀*δ + R₀))/(k*x*y)
end #function to help calculate the resource isocline of the Yodzis-Innes model

function iso_plot_YodInn(resrange, par)
    data = [res_iso_yodinn(R, par) for R in resrange]
    plot(collect(resrange), data)
    return vlines(con_iso_yodinn(par), 0, 0.5, colors="orange")
end #plots the isoclines of the Yodzis-Innes model

function YodInn_pert(ep, R0, freq, r, seed, tsend, tvals)
    Random.seed!(seed)
    par = YodInnScalePar()
    par.ε = ep
    par.R₀ = R0
    noise = noise_creation(r, tsend / freq)
    count = 1
    u0 = [eq_yodinn_II(par)[1], eq_yodinn_II(par)[2] + noise[1]]
    tspan = (0.0, tsend)

    function pert_cb2(integrator)
        count += 1
        if isapprox(integrator.u[2], 0.00000000; atol = 1e-8)
            integrator.u[2] = 0.00000000
        else
            integrator.u[2] = maximum([integrator.u[2] + noise[count], 0.0])
        end
    end

    cb = PeriodicCallback(pert_cb2, freq, initial_affect = false) #as of may 29th - initial_affect does not actually do the affect on the first time point
    prob = ODEProblem(yod_inn_scale_II!, u0, tspan, par)
    sol = DifferentialEquations.solve(prob, callback = cb, reltol = 1e-8)
    return solend = sol(tvals)
end #function to numerically solve the Yodzis-Innes consumer-resource model with added noise to the consumer

function pert_YodInn_phase_plot(ep, R0, freq, r, seed, tsend, tvals)
    sol = YodInn_pert(ep, R0, freq, r, seed, tsend, tvals)
    plot(sol[1, :], sol[2, :], color = "green")
    xlabel("Resource")
    return ylabel("Consumer")
end #plot a phase space time series for the Yodzis-Innes model

### Below are functions to help calculate the feasible R₀ values, the isocline functions, and the Hopf bifurcation values for the Yodzis-Innes model
# #Feasible R₀
# function feasibleR0(p)
#     @unpack y, k = p
#     lower = ( (y-1)/(y+1) ) * k
#     upper = (y-1) * k
#     return vcat(lower, upper)
# end

# feasibleR0(YodInnScalePar()) #0.699849 1.82557

## isoclines
# using SymPy
# @vars R C
# @vars x y R₀ k δ

# f(R, C) = R * (1 - R / k) - ( x * y / (1 - δ) ) * R * C / ( R₀ + R )
# g(R ,C) = x * y * R * C / (R + R₀) - x * C

# SymPy.solve(f(R,C),R)
# SymPy.solve(f(R,C),C)
# SymPy.solve(g(R,C),R)
# SymPy.solve(g(R,C),C)

# #find maximum of resource isocline
# using Calculus
# Calculus.simplify(Calculus.differentiate("(R*(R*δ - R + R₀*δ - R₀) + k*(-R*δ + R - R₀*δ + R₀))/(k*x*y)", :R))
# @vars R C
# @vars x y R₀ k δ a
# r(R) = (((((((R * δ - R) + R₀ * δ) - R₀) + R * (δ - 1)) + k * (1 + -δ)) * (k * x * y)) / (k * x * y) ^ 2)

# SymPy.solve(r(R), R)

# function hopf_YodInn(p)
#     @unpack k, R₀, δ, x, y= p
#     hopf_R = -R₀/2 + k/2
#     hopf_C = (hopf_R*(hopf_R*δ - hopf_R + R₀*δ - R₀) + k*(-hopf_R*δ + hopf_R - R₀*δ + R₀))/(k*x*y)
#     return vcat(hopf_R,hopf_C)
# end

# hopf_YodInn(YodInnScalePar()) #R=1.1, C=6.651085519687272


# ---
# title: 877. Stone Game
# id: problem877
# author: Indigo
# date: 2021-06-16
# difficulty: Medium
# categories: Math, Dynamic Programming, Minimax
# link: <https://leetcode.com/problems/stone-game/description/>
# hidden: true
# ---
#
# Alex and Lee play a game with piles of stones.  There are an even number of
# piles **arranged in a row** , and each pile has a positive integer number of
# stones `piles[i]`.
#
# The objective of the game is to end with the most stones.  The total number of
# stones is odd, so there are no ties.
#
# Alex and Lee take turns, with Alex starting first.  Each turn, a player takes
# the entire pile of stones from either the beginning or the end of the row.
# This continues until there are no more piles left, at which point the person
# with the most stones wins.
#
# Assuming Alex and Lee play optimally, return `True` if and only if Alex wins
# the game.
#
#
#
# **Example 1:**
#
#
#
#     Input: piles = [5,3,4,5]
#     Output: true
#     Explanation:
#     Alex starts first, and can only take the first 5 or the last 5.
#     Say he takes the first 5, so that the row becomes [3, 4, 5].
#     If Lee takes 3, then the board is [4, 5], and Alex takes 5 to win with 10 points.
#     If Lee takes the last 5, then the board is [3, 4], and Alex takes 4 to win with 9 points.
#     This demonstrated that taking the first 5 was a winning move for Alex, so we return true.
#
#
#
#
# **Constraints:**
#
#   * `2 <= piles.length <= 500`
#   * `piles.length` is even.
#   * `1 <= piles[i] <= 500`
#   * `sum(piles)` is odd.
#
#
## @lc code=start
using LeetCode
"""
function stone_game(piles::Vector{Int})
    len = length(piles)
    dp = Matrix{Int}(undef, len, len)
    for i in 1:len
        dp[i, i] = piles[i]
    end
    for i in (len - 1):-1:1, j in (i + 1):len
        dp[i, j] = max(piles[i] - dp[i + 1, j], piles[j] - dp[i, j - 1])
    end
    return dp[1, end] > 0
end
"""
stone_game(::Vector{Int}) = true
## @lc code=end


#!/usr/bin/env julia

module NRI

import ..ReconstructionEvaluations
const RE = ReconstructionEvaluations;

"""

   nri( om::SparseMatrixCSC; correct=true )

Implementation of the Neural Reconstruction Index. Incorporates
the correction for false positive values if "correct" is true.
"""
function nri( om::SparseMatrixCSC; correct=true )

  TP, segTP = compute_TPs( om )
  FP, segFP = compute_FPs( om; correct=correct )
  FN, segFN = compute_FNs( om )

  NRI = nri( TP, FP, FN )
  segNRI = nri( segTP, segFP, segFN )

  segNRIw = nri_weight( segTP, segFP, segFN )

  NRI, segNRI, segNRIw
end

function nri( om::Array; correct=true )
  nri( sparse(om); correct=correct )
end

function nri( TP, FP, FN )
  2*TP ./ (2*TP + FP + FN)
end

function nri( TP::SparseVector, FP::SparseVector, FN::SparseVector )

  inds = find(TP); #if a segment has no TP's, segNRI=0
  res = sparsevec(Int[],Float64[],length(TP))

  for i in inds res[i] = (2*TP[i]) / (2*TP[i] + FP[i] + FN[i]) end

  res
end

function nri( table1::AbstractArray, table2::AbstractArray; correct=true )
  count_table, A_to_inds, B_to_inds = RE.build_count_table(table1,table2)

  NRI, segNRI, segNRIw = nri( count_table; correct=correct )

  res_segNRI  = sparsevec( Int[], Float64[], maximum(keys(A_to_inds))+1 )
  res_segNRIw = sparsevec( Int[], Float64[], maximum(keys(A_to_inds))+1 )
  for (k,v) in A_to_inds
      res_segNRI[k] = segNRI[v]
      res_segNRIw[k] = segNRIw[v]
  end

  NRI, res_segNRI, res_segNRIw
end

function nri( fname1::AbstractString, fname2::AbstractString )

  @assert isfile(fname1) && isfile(fname2)

  table1, table2 = RE.load_edges(fname1), RE.load_edges(fname2)

  nri( table1, table2 )
end

"""

    compute_TPs( om, offset=true )

Counts the number of true positive paths within each
segment in a network specified by a sparse overlap matrix,
and summarizes their total.

Offset indicates whether segment ids are offset by 1 to include
synapse "insertions" and deletions. This feature is not fully implemented
yet.
"""
function compute_TPs( om, offset=true )

  rs, cs = findn(om); vs = nonzeros(om);

  TPs = choosetwo(vs);
  segTPs = sparsevec(Int[],Int[],size(om,1));

  oT = one(eltype(rs));
  for (r,c,tp) in zip(rs,cs,TPs)

    if r == oT continue end
    if c == oT continue end

    segTPs[r] += tp
  end

  sum(segTPs), segTPs
end

"""

    compute_FPs( om, offset=true; correct=true )

Counts the number of false positive paths within each
segment in a network specified by a sparse overlap matrix,
and summarizes their total.

offset indicates whether segment ids are offset by 1 to include
synapse "insertions" and deletions. This feature is not fully implemented
yet.

correct indicates whether to fully penalize false positives (unlike within
the first draft of MATLAB nri code). The false setting is meant to be
used as a comparison
"""
function compute_FPs( om, offset=true; correct=true )

  rs, cs = findn(om); vs = nonzeros(om);

  sum_over_cols = sparsevec(Int[],Int[],size(om,2))
  for (c,v) in zip(cs, vs) sum_over_cols[c] += v end

  #needs to be a float to account for attributing
  # mergers to gt segments
  segFPs = sparsevec(Int[],Float64[],size(om,1))

  oT = one(eltype(rs))
  for (r,c,v) in zip(rs,cs,vs)

    if c == oT continue end

    other_count = sum_over_cols[c] - v
    segFPs[r] += (v*other_count) / 2

    #fully accounting for synapse insertions (if using the correct vers)
    if r == oT && correct segFPs[r] += choosetwo(v) end
  end

  sum(segFPs), segFPs
end

"""

    compute_FNs( om, offset=true )

Counts the number of false negative paths within each
segment in a network specified by a sparse overlap matrix,
and summarizes their total.

offset indicates whether segment ids are offset by 1 to include
synapse "insertions" and deletions. This feature is not fully implemented
yet.

correct indicates whether to fully penalize false positives (unlike within
the first draft of MATLAB nri code). The false setting is meant to be
used as a comparison
"""
function compute_FNs( om, offset=true )

  rs, cs = findn(om); vs = nonzeros(om);

  sum_over_rows = sparsevec(Int[],Int[],size(om,1))
  for (r,v) in zip(rs, vs) sum_over_rows[r] += v end

  segFNs = sparsevec(Int[],Float64[],size(om,1))

  oT = one(eltype(rs))
  for (r,c,v) in zip(rs,cs,vs)

    if r == oT continue end

    other_count = sum_over_rows[r] - v
    segFNs[r] += (v*other_count) / 2

    #accounting for synapse deletions
    if c == oT segFNs[r] += choosetwo(v) end
  end

  sum(segFNs), segFNs
end

"""

    nri_weight( segTP, segFP, segFN )

Finds the weight for each segment s.t. the full network
NRI is a weighted sum of each individual NRI values. This value
is equal to the sum of "positives" (TP+FP) and "negatives" (TP+FN)
for each segment.
"""
function nri_weight( segTP, segFP, segFN )

  total_weight = 2*sum(segTP) + sum(segFP) + sum(segFN)

  is = Set(find(segTP)); union!(is, find(segFP)); union!(is, find(segFN));

  res = sparsevec(Int[],Float64[],length(segTP))

  for i in is res[i] = (2*segTP[i] + segFP[i] + segFN[i]) / total_weight end

  res
end

"""

    nri_by_class( seg_nri, seg_nriw, class_map )
"""
function nri_by_class( seg_nri, seg_nriw, class_map )

  is = find(seg_nriw)

  class_nri = Dict( k => 0. for k in unique(values(class_map)) )
  class_w   = Dict( k => 0. for k in unique(values(class_map)) )

  for i in is
    if !haskey(class_map,i) continue end

    class_w[class_map[i]] += seg_nriw[i]
  end

  for i in is
    if !haskey(class_map,i) continue end
    class = class_map[i]

    class_nri[class] += (seg_nri[i] * seg_nriw[i]) / class_w[class]
  end

  class_nri, class_w
end

#Utility fns
choosetwo(x) = x * ((x-1)/2)
choosetwo!(x) = for i in eachindex(x) x[i] *= (x[i]-1)/2 end
choosetwo(x::Array) = [ v*((v-1)/2) for v in x ]

end #module end


module Graphtypes

export
    Graph,
    SimpleGraph,
    StructureGraph,

    simple_graph,
    structure_graph,

    connected,
    count_connections,
    list_outgoing,
    total_edges,
    total_vertices,
    list_edges,
    connect!,
    disconnect!,
    reverse!

abstract Graph{directed, adjacent}

type SimpleGraph <: Graph
    adjacent::Matrix{Int64} #adjacency matrix
end

type StructureGraph <: Graph
    #undirected, no duplicate edges
    adjacent::Matrix{Bool}
end

function simple_graph (v_count::Int64)
    g = SimpleGraph(zeros(Int64, v_count, v_count))
    return g
end

function structure_graph (v_count::Int64)
    g = StructureGraph(zeros(Bool, v_count, v_count))
    return g
end

function list_edges(g::SimpleGraph)
    size = total_vertices(g)
    edges = Array(Int64, 0, 2)
    for i in 1:size
        for j in 1:size
            for k in 1:count_connections(g,i,j)
                edges = [edges, [i j]]
    end end end
    return edges
end

function list_edges(g::StructureGraph)
    size = total_vertices(g)
    edges = Array(Int64, 0, 2)
    for i in 1:size
        for j in i:size
            if g.adjacent[i,j]
                edges = [edges, [i j]]
    end end end
    return edges
end

function list_outgoing(g::SimpleGraph, from::Int64)
    return find( x->(x != 0), g.adjacent[from, :])
end

function list_outgoing(g::StructureGraph, from::Int64)
    return sort(unique([    find( x->(x != 0), g.adjacent[from, :]),
                    find( x->(x != 0), g.adjacent[:, from])]))
end

function connected(g::Graph, from::Int64, to::Int64)
    return g.adjacent[from, to] != 0
end

function count_connections(g::Graph, from::Int64, to::Int64)
    return g.adjacent[from, to]
end

function total_edges(g::SimpleGraph)
    return sum(g.adjacent)
end

function total_edges(g::StructureGraph)
    return sum(g.adjacent)/2
end

function total_vertices(g::Graph)
    return size(g.adjacent)[2]
end

function connect!(g::SimpleGraph, from::Int64, to::Int64)
    g.adjacent[from, to] += 1
end

function connect!(g::StructureGraph, from::Int64, to::Int64)
    g.adjacent[from, to] = true
    g.adjacent[to, from] = true
end

function disconnect!(g::SimpleGraph, from::Int64, to::Int64)
    if g.adjacent[from, to] > 0
        g.adjacent[from, to] -= 1
        return true
    else
        return false
    end
end

function disconnect!(g::StructureGraph, from::Int64, to::Int64)
    if g.adjacent[from, to]
        g.adjacent[from, to] = false
        return true
    else
        return false
    end
end
function reverse!(g::Graph, i1::Int64, i2::Int64)
    if disconnect!(g,i1,i2)
        connect!(g,i2,i1)
        return true
    end
    return false
end

end