text:
# these functions are related to calculating RV quantities
using UnitfulAstro, Unitful
using LinearAlgebra
using Optim, LineSearches
import GPLinearODEMaker; GLOM = GPLinearODEMaker

const n_kep_parms = 6

abstract type KeplerSignal end

"Convert Unitful units from one to another and strip the final units"
convert_and_strip_units(new_unit::Unitful.FreeUnits, quant::Quantity) = ustrip(uconvert(new_unit, quant))
convert_and_strip_units(new_unit::Unitful.FreeUnits, quant) = quant

"""
Calculate velocity semi-amplitude for RV signal of planets in m/s
adapted from eq. 12 of (http://exoplanets.astro.yale.edu/workshop/EPRV/Bibliography_files/Radial_Velocity.pdf)
"""
function velocity_semi_amplitude(
    m_star::Unitful.Mass,
    m_planet::Unitful.Mass,
    P::Unitful.Time;
    e::Real=0.,
    i::Real=π/2)

    # m_star = convert_and_strip_units(u"Msun", m_star)
    # m_planet = convert_and_strip_units(u"Msun", m_planet)
    # P = convert_and_strip_units(u"yr", P)
    assert_positive(P, m_star, m_planet)
    comb_mass = m_star + m_planet
    K_au_yr = (2 * π * sin(i) * m_planet / (sqrt((1-e)*(1+e)) * cbrt(comb_mass * comb_mass * P)))u"AU/yr"
    return uconvert(u"m/s", K_au_yr)
end

# unit_phase(t::Unitful.Time, P::Unitful.Time) = mod(t / P * 2 * π, 2 * π)
unit_phase(t::Unitful.Time, P::Unitful.Time) = t / P * 2 * π

"calculate mean anomaly"
mean_anomaly(t::Unitful.Time, P::Unitful.Time, M0::Real) = unit_phase(t, P) - M0
# mean_anomaly(t::Unitful.Time, P::Unitful.Time, M0::Real) = mod(unit_phase(t, P) - M0, 2 * π)
# mean_anomaly(t::Unitful.Time, P::Unitful.Time, M0::Real) = mod(t / P * 2 * π - M0, 2 * π)

"""   ecc_anom_init_guess_danby(M, ecc)
Initial guess for eccentric anomaly given mean anomaly (M) and eccentricity (ecc)
    Based on "The Solution of Kepler's Equations - Part Three"
    Danby, J. M. A. (1987) Journal: Celestial Mechanics, Volume 40, Issue 3-4, pp. 303-312  1987CeMec..40..303D
Julia function originally written by Eric Ford
"""
function ecc_anom_init_guess_danby(M::Real, e::Real)
    k = convert(typeof(e), 0.85)
    if M < zero(M); M += 2 * π end
    (M < π) ? M + k * e : M - k * e;
end

"""   update_ecc_anom_laguerre(E, M, ecc)
Update the current guess (E) for the solution to Kepler's equation given mean anomaly (M) and eccentricity (ecc)
    Based on "An Improved Algorithm due to Laguerre for the Solution of Kepler's Equation"
    Conway, B. A.  (1986) Celestial Mechanics, Volume 39, Issue 2, pp.199-211  1986CeMec..39..199C
Julia function originally written by Eric Ford
Basically looks for root of an nth order polynomial assuming that 1 root is
the one you are looking for (E) and the rest are evenly spaced
"""
function update_ecc_anom_laguerre(E::Real, M::Real, e::Real)
    # es = e * sin(E)
    # ec = e * cos(E)
    F = (E - e * sin(E)) - M
    Fp = one(E) - e * cos(E)
    # Fpp = es
    n = 5
    root = sqrt(abs((n - 1) * ((n - 1) * Fp * Fp - n * F * e * sin(E))))
    denom = Fp > zero(E) ? Fp + root : Fp - root
    return E - n * F / denom
end

"""
Loop to update the current estimate of the solution to Kepler's equation
Julia function originally written by Eric Ford
"""
function ecc_anomaly(
    t::Unitful.Time,
    P::Unitful.Time,
    M0::Number,
    e::Number;
    tol::Number=1e-8,
    max_its::Integer=200)

    # P = convert_and_strip_units(u"yr", P)
    # @assert (0 <= e <= 1) "eccentricity has to be between 0 and 1"
    M = mean_anomaly(t, P, M0)
    E = ecc_anom_init_guess_danby(M, e)
    for i in 1:max_its
       E_old = E
       E = update_ecc_anom_laguerre(E_old, M, e)
       if abs(E - E_old) < tol; break end
    end
    return E
end

# """
# Iterative solution for eccentric anomaly from mean anomaly
# from (http://www.csun.edu/~hcmth017/master/node16.html)
# could also implement third order method from
# (http://alpheratz.net/dynamics/twobody/KeplerIterations_summary.pdf)
# """
# function ecc_anomaly(t::Unitful.Time, P::Unitful.Time, e::Real)
#     @assert (0 <= e <= 1) "eccentricity has to be between 0 and 1"
#     M = mean_anomaly(t, P)
#     dif_thres = 1e-8
#     dif = 1
#     ecc_anom = copy(M)
#     while dif > dif_thres
#         ecc_anom_old = copy(ecc_anom)
#         ecc_anom -= (ecc_anom - e * sin(ecc_anom) - M) / (1 - e * cos(ecc_anom))  # first order method
#         dif = abs((ecc_anom - ecc_anom_old) / ecc_anom)
#     end
#     return ecc_anom
# end

"""
Solution for true anomaly from eccentric anomaly
from (https://en.wikipedia.org/wiki/True_anomaly#From_the_eccentric_anomaly)
"""
ϕ(
    E::Real,
    e::Real
    ) = 2 * atan( sqrt(1 + e)*sin(0.5*E), sqrt(1 - e) * cos(0.5*E) )

ϕ(
    t::Unitful.Time,
    P::Unitful.Time,
    M0::Real,
    e::Real
    ) = ϕ( ecc_anomaly(t, P, M0, e), e)

# """
# Calculate true anomaly for small e using equation of center approximating true
# anomaly from (https://en.wikipedia.org/wiki/Equation_of_the_center) O(e^8)
# """
# function ϕ_approx(t::Unitful.Time, P::Unitful.Time, e::Real)
#     @assert (0 <= e <= 1) "eccentricity has to be between 0 and 1"
#     M = mean_anomaly(t, P)
#
#     term_list = [eval_polynomial(e, [0, 2, 0, - 1/4, 0, 5/96, 0, 107/4608]),
#     eval_polynomial(e, [0, 0, 5/4, 0, -11/24, 0, 17/192, 0]),
#     eval_polynomial(e, [0, 0, 0, 13/12, 0, -43/64, 0, 95/512]),
#     eval_polynomial(e, [0, 0, 0, 0, 103/96, 0, -451/480, 0]),
#     eval_polynomial(e, [0, 0, 0, 0, 0, 1097/960, 0, -5957/4608]),
#     eval_polynomial(e, [0, 0, 0, 0, 0, 0, 1223/960, 0]),
#     eval_polynomial(e, [0, 0, 0, 0, 0, 0, 0, 47273/32256])]
#
#     sin_list = [sin(i * M) for i in 1:7]
#
#     return M + dot(term_list, sin_list)
#
# end
#
# ϕ_approx(t::Unitful.Time, P::Unitful.Time, e::Real) = ϕ_approx(t, convert_and_strip_units(u"yr", P), e)

"""
kepler_rv(t, K, P, M0, e, ω; γ)
Calculate radial velocity of star due to planet following Keplerian orbit at specified time.
Inputs:
- time
- Period
- eccentricity
- mean anomaly at time=0
- RV amplitude
- argument of periapsis
Optionally:
- velocity offset
"""
function kepler_rv(
    t::Unitful.Time,
    K::Unitful.Velocity,
    P::Unitful.Time,
    M0::Real,
    e_or_h::Real,
    ω_or_k::Real;
    kwargs...)

    return kepler_rv_ecc_anom(t, K, P, M0, e_or_h, ω_or_k; kwargs...)
    # return kepler_rv_true_anom(t, P, M0, e, K, ω; kwargs...)
end

function kepler_rv(
    t::Unitful.Time,
    m_star::Unitful.Mass,
    m_planet::Unitful.Mass,
    P::Unitful.Time,
    M0::Real,
    e::Real,
    ω::Real;
    i::Real=π/2,
    kwargs...)

    K = velocity_semi_amplitude(P, m_star, m_planet, e=e, i=i)
    return kepler_rv(t, K, P, M0, e, ω; kwargs...)
end

struct kep_signal <: KeplerSignal
    K::Unitful.Velocity
    P::Unitful.Time
    M0::Real  # ::Unitful.NoDims
    e::Real
    ω::Real  # ::Unitful.NoDims
    γ::Unitful.Velocity
    h::Real  # ::Unitful.NoDims
    k::Real  # ::Unitful.NoDims

    kep_signal(; K=0u"m/s", P=1u"d", M0=0, e_or_h=0, ω_or_k=0, γ=0u"m/s", use_hk::Bool=false) = kep_signal(K, P, M0, e_or_h, ω_or_k, γ; use_hk=use_hk)
    function kep_signal(K, P, M0, e_or_h, ω_or_k, γ; use_hk::Bool=false)
        if use_hk
            h = e_or_h
            k = ω_or_k
            e, ω = hk_2_eω(h, k)
        else
            e = e_or_h
            ω = ω_or_k
            h, k = eω_2_hk(e, ω)
        end
        return kep_signal(K, P, M0, e, ω, γ, h, k)
    end
    function kep_signal(
        K::Unitful.Velocity,
        P::Unitful.Time,
        M0::Real,
        e::Real,
        ω::Real,
        γ::Unitful.Velocity,
        h::Real,
        k::Real)

        # @assert 0 <= e < 1 "orbit needs to be bound"
        M0 = mod2pi(M0)
        ω = mod2pi(ω)
        # assert_positive(P)
        return new(K, P, M0, e, ω, γ, h, k)
    end
end
kepler_rv(t::Unitful.Time, ks::kep_signal) = kepler_rv(t, ks.K, ks.P, ks.M0, ks.e, ks.ω; γ=ks.γ)
(ks::KeplerSignal)(t::Unitful.Time) = kepler_rv(t, ks)

mutable struct kep_buffer{T1<:Real}
    ks::kep_signal
    rm_kep::Vector{T1}
    nprior::T1
end

# function nlogL_kep(
#     data::Vector{T},
#     times::Vector{T2} where T2<:Unitful.Time,
#     covariance::Union{Cholesky,Diagonal},
#     ks::kep_signal;
#     data_unit::Unitful.Velocity=1u"m/s")
#     return GLOM.nlogL(covariance, remove_kepler(data, times, ks_internal; data_unit=data_unit))
# end

function ∇nlogL_kep!(
    G::Vector{T},
    d::Vector{<:Integer},
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    ks::kep_signal,
    y::Vector{T};
    data_unit::Unitful.Velocity=1u"m/s",
    include_priors::Bool=true) where T<:Real

    α = covariance \ y
    for i in 1:length(G)
        d[i] = 1
        G[i] = GLOM.dnlogLdθ(remove_kepler(data, times, ks; data_unit=data_unit, d=d), α) - (include_priors * logprior_kepler(ks; d=d, use_hk=true))
        d[i] = 0
    end
end
∇nlogL_kep(data, times, covariance, ks; data_unit=data_unit, kwargs...) =
    ∇nlogL_kep(data, times, covariance, ks, remove_kepler(data, times, ks; data_unit=data_unit); data_unit=data_unit, kwargs...)
function ∇nlogL_kep(data, times, covariance, ks, y; G=zeros(typeof(data[1]), n_kep_parms), d=zeros(Int, n_kep_parms), kwargs...)
    ∇nlogL_kep!(G, d, data, times, covariance, ks, y; kwargs...)
    return G
end
∇nlogL_kep(glo_rv::GLO_RV, covariance, ks; kwargs...) =
    ∇nlogL_kep(glo_rv.GLO.y_obs, glo_rv.time, covariance, ks; data_unit=glo_rv.rv_factor, kwargs...)

function fit_kepler(
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    init_ks::kep_signal;
    data_unit::Unitful.Velocity=1u"m/s",
    print_stuff::Bool=true,
    include_priors::Bool=true,
    avoid_saddle::Bool=true
    ) where T<:Real

    # if P==0; return kep_signal_epicyclic(p, P, M0, e, ω, γ, coefficients) end
    assert_positive(init_ks.P)
    for i in 1:ndims(covariance)
        @assert size(covariance,i)==length(data) "covariance incompatible with data"
    end
    n_total_samp_points = length(data)
    n_samp_points = length(times)
    normalization = logdet(covariance) + n_total_samp_points * log(2 * π)
    K_u = unit(init_ks.K)
    P_u = unit(init_ks.P)
    γ_u = unit(init_ks.γ)

    current_x = ustrip.([init_ks.K, init_ks.P, init_ks.M0, init_ks.h, init_ks.k, init_ks.γ])
    buffer = kep_buffer(init_ks, remove_kepler(data, times, init_ks; data_unit=data_unit), -logprior_kepler(init_ks; use_hk=true))
    last_x = similar(current_x)

    function calculate_common!(x::Vector{T}, last_x::Vector{T}, buffer::kep_buffer) where {T<:Real}
        if x != last_x
            # copy!(last_x, x)
            last_x[:] = x
            buffer.ks = ks_from_vec(x, K_u, P_u, γ_u; use_hk=true)
            buffer.nprior = -logprior_kepler(buffer.ks; use_hk=true)
            buffer.nprior == Inf ? buffer.rm_kep[:] .= 0 : buffer.rm_kep = remove_kepler(data, times, buffer.ks; data_unit=data_unit)
        end
        # println("buffer: ", kep_parms_str(buffer.ks))
    end

    function f(x::Vector{T}, buffer::kep_buffer, last_x::Vector{T}) where {T<:Real}
        calculate_common!(x, last_x, buffer)
        # println(kep_parms_str(buffer.ks))
        # println(buffer.nprior)
        if buffer.nprior == Inf
            return buffer.nprior
        else
            return GLOM.nlogL(covariance, buffer.rm_kep, nlogL_normalization=normalization) + (include_priors * buffer.nprior)
        end
    end
    f(x) = f(x, buffer, last_x)

    d = zeros(Int, n_kep_parms)
    function g!(G::Vector{T}, x::Vector{T}, buffer::kep_buffer, last_x::Vector{T}) where {T<:Real}
        calculate_common!(x, last_x, buffer)
        if buffer.nprior == Inf
            G[:] .= 0
        else
            ∇nlogL_kep!(G, d, data, times, covariance, buffer.ks, buffer.rm_kep; data_unit=data_unit, include_priors=include_priors)
        end
    end
    g!(G, x) = g!(G, x, buffer, last_x)
    attempts = 0
    in_saddle = true
    while attempts < 120 && in_saddle
        attempts += 1
        if attempts > 1;
            if print_stuff; println("found saddle point. starting attempt $attempts with a perturbation") end
            current_x[1] = maximum([current_x[1] + centered_rand(; scale=0.3), 0.2])
            current_x[2] *= centered_rand(; scale=2e-2, center=1)
            current_x[3] = mod2pi(current_x[3] + centered_rand(; scale=4e-1))
            # e, ω = hk_2_eω(current_x[4], current_x[5])
            # e = maximum([minimum([e, 0.1]), 0])
            # ω = mod2pi(ω)
            current_x[4:5] .= eω_2_hk(0.05 * rand(), mod2pi(attempts * π / 3))
            current_x[6] += centered_rand(; scale=1e-1)
        end
        # println(current_x)
        result = optimize(f, g!, current_x, LBFGS(alphaguess=LineSearches.InitialStatic(alpha=3e-3))) # 27s
        current_x = copy(result.minimizer)
        ks = ks_from_vec(current_x, K_u, P_u, γ_u; use_hk=true)
        if print_stuff; println("fit attempt $attempts: "kep_parms_str(ks)) end
        # println(∇∇nlogL_kep(data, times, covariance, ks; data_unit=data_unit))
        new_det = det(∇∇nlogL_kep(data, times, covariance, ks; data_unit=data_unit, include_priors=include_priors))
        if print_stuff; println("determinant: ", new_det) end
        in_saddle = (new_det <= 0) && avoid_saddle
        if !in_saddle; return ks end
    end
    @error "no non-saddle point soln found"
end

"""
Simple radial velocity formula using true anomaly
adapted from eq. 11 of (http://exoplanets.astro.yale.edu/workshop/EPRV/Bibliography_files/Radial_Velocity.pdf)
"""
kepler_rv_true_anom(
    t::Unitful.Time,
    K::Unitful.Velocity,
    P::Unitful.Time,
    M0::Real,
    e::Real,
    ω::Real;
    γ::Unitful.Velocity=0u"m/s"
    ) = K * (e * cos(ω) + cos(ω + ϕ(t, P, M0, e))) + γ

function eω_2_hk(e, ω)
    return e * sin(ω), e * cos(ω)
end
function hk_2_eω(h, k)
    return sqrt(h * h + k * k), atan(h, k)
end

"""
Radial velocity formula directly using eccentric anomaly
Based on simplification of "An analytical solution for Kepler's problem"
Pál, András, Monthly Notices of the Royal Astronomical Society, 396, 3, 1737-1742.  2009MNRAS.396.1737P
see ηdot part of eq. 19
"""
function kepler_rv_ecc_anom(
    t::Unitful.Time,
    K::Unitful.Velocity,
    P::Unitful.Time,
    M0::Real,
    e_or_h::Real,
    ω_or_k::Real;
    γ::Unitful.Velocity=0u"m/s",
    use_hk::Bool=false)

    if use_hk
        h = e_or_h
        k = ω_or_k
        e, ω = hk_2_eω(h, k)
    else
        e = e_or_h
        ω = ω_or_k
        h, k = eω_2_hk(e, ω)
    end

    # k = e * cos(ω)
    E = ecc_anomaly(t, P, M0, e)
    j = sqrt((1-e)*(1+e))
    q = e * cos(E)
    # return K * j / (1 - q) * (cos(ω + E) - (1 - j) * cos(E) * cos(ω)) + γ  # equivalent
    return K * j / (1 - q) * (cos(ω + E) - k * q / (1 + j)) + γ
end

struct kep_signal_circ <: KeplerSignal
    K::Unitful.Velocity
    P::Unitful.Time
    M0minusω::Real  # ::Unitful.NoDims
    γ::Unitful.Velocity
    coefficients::Vector{T} where T<:Unitful.Velocity

    kep_signal_circ(; K=0u"m/s", P=0u"d", M0minusω=0, γ=0u"m/s") = kep_signal_circ(K, P, M0minusω, γ)
    kep_signal_circ(K, P, M0minusω, γ) = kep_signal_circ(K, P, M0minusω, γ, [K * cos(M0minusω), K * sin(M0minusω), γ])
    function kep_signal_circ(
        P::Unitful.Time,
        coefficients::Vector{T} where T<:Unitful.Velocity)

        return kep_signal_circ(
            sqrt(coefficients[1]^2 + coefficients[2]^2),
            P,
            mod2pi(atan(coefficients[2], coefficients[1])),
            coefficients[3],
            coefficients)
    end
    function kep_signal_circ(
        K::Unitful.Velocity,
        P::Unitful.Time,
        M0minusω::Real,  # ::Unitful.NoDims
        γ::Unitful.Velocity,
        coefficients::Vector{T} where T<:Unitful.Velocity)

        M0minusω = mod2pi(M0minusω)
        assert_positive(P)
        @assert length(coefficients) == 3 "wrong number of coefficients"
        return new(K, P, M0minusω, γ, coefficients)
    end
end
function kepler_rv_circ(t::Unitful.Time, ks::kep_signal_circ)
    phase = unit_phase.(t, ks.P)
    return (ks.coefficients[1] .* cos.(phase)) + (ks.coefficients[2] .* sin.(phase)) .+ ks.coefficients[3]
end
kepler_rv(t::Unitful.Time, ks::kep_signal_circ) = kepler_rv_circ(t, ks)

"""
A linear formulation of a Keplerian RV signal for low eccentricity by expanding about e=0
average difference is < 1% for e < 0.1 and < 10% for e < 0.35
for initial full orbit fits
coefficients[1] = K * cos(M0 - ω)
coefficients[2] = K * sin(M0 - ω)
coefficients[3] = e * K * cos(2 * M0 - ω)
coefficients[4] = e * K * sin(2 * M0 - ω)
coefficients[5] = γ
so
K = sqrt(coefficients[1]^2 + coefficients[2]^2)
e = sqrt(coefficients[3]^2 + coefficients[4]^2) / K
M0 = atan(coefficients[4], coefficients[3]) - atan(coefficients[2], coefficients[1])
ω = atan(coefficients[4], coefficients[3]) - 2 * atan(coefficients[2], coefficients[1])
"""
struct kep_signal_epicyclic <: KeplerSignal
    K::Unitful.Velocity
    P::Unitful.Time
    M0::Real
    e::Real
    ω::Real  # ::Unitful.NoDims
    γ::Unitful.Velocity
    h::Real
    k::Real  # ::Unitful.NoDims
    coefficients::Vector{T} where T<:Unitful.Velocity

    kep_signal_epicyclic(; K=0u"m/s", P=1u"d", M0=0, e_or_h=0, ω_or_k=0, γ=0u"m/s", use_hk::Bool=false) = kep_signal_epicyclic(K, P, M0, e_or_h, ω_or_k, γ; use_hk=use_hk)
    function kep_signal_epicyclic(K, P, M0, e_or_h, ω_or_k, γ; use_hk::Bool=false)
        if use_hk
            h = e_or_h
            k = ω_or_k
            e, ω = hk_2_eω(h, k)
        else
            e = e_or_h
            ω = ω_or_k
            h, k = eω_2_hk(e, ω)
        end
        return kep_signal_epicyclic(K, P, M0, e, ω, γ, h, k)
    end
    kep_signal_epicyclic(K, P, M0, e, ω, γ, h, k) = kep_signal_epicyclic(K, P, M0, e, ω, γ, h, k,
        [K * cos(M0 - ω), K * sin(M0 - ω), e * K * cos(2 * M0 - ω), e * K * sin(2 * M0 - ω), γ])
    function kep_signal_epicyclic(P::Unitful.Time, coefficients)
        K = sqrt(coefficients[1]^2 + coefficients[2]^2)
        e = sqrt(coefficients[3]^2 + coefficients[4]^2) / K
        # if e > 0.35; @warn "an orbit with this eccentricity would be very different from the output of kepler_rv_epicyclic" end
        M0 = mod2pi(atan(coefficients[4], coefficients[3]) - atan(coefficients[2], coefficients[1]))
        ω = mod2pi(atan(coefficients[4], coefficients[3]) - 2 * atan(coefficients[2], coefficients[1]))
        γ = coefficients[5]
        h, k = eω_2_hk(e, ω)
        return kep_signal_epicyclic(K, P, M0, e, ω, γ, h, k, coefficients)
    end
    function kep_signal_epicyclic(
        K::Unitful.Velocity,
        P::Unitful.Time,
        M0::Real,
        e::Real,
        ω::Real,  # ::Unitful.NoDims
        γ::Unitful.Velocity,
        h::Real,
        k::Real,
        coefficients::Vector{T} where T<:Unitful.Velocity)

        M0 = mod2pi(M0)
        ω = mod2pi(ω)
        assert_positive(P)
        @assert length(coefficients) == 5 "wrong number of coefficients"
        return new(K, P, M0, e, ω, γ, h, k, coefficients)
    end
end
function kepler_rv_epicyclic(t::Unitful.Time, ks::kep_signal_epicyclic)
    phase = unit_phase.(t, ks.P)
    return return (ks.coefficients[1] .* cos.(phase)) + (ks.coefficients[2] .* sin.(phase)) + (ks.coefficients[3] .* cos.(2 .* phase)) + (ks.coefficients[4] .* sin.(2 .* phase)) .+ ks.coefficients[5]
end
kepler_rv(t::Unitful.Time, ks::kep_signal_epicyclic) = kepler_rv_epicyclic(t, ks)
function fit_kepler_epicyclic(
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    P::Unitful.Time;
    data_unit::Unitful.Velocity=1u"m/s",
    print_stuff::Bool=true,
    avoid_saddle::Bool=true
    ) where T<:Real

    # if P==0; return kep_signal_epicyclic(p, P, M0, e, ω, γ, coefficients) end
    assert_positive(P)
    for i in 1:ndims(covariance)
        @assert size(covariance,i)==length(data) "covariance incompatible with data"
    end
    n_total_samp_points = length(data)
    n_samp_points = length(times)
    n_out = Int(n_total_samp_points / n_samp_points)
    phases = unit_phase.(times, P)
    design_matrix = zeros(n_total_samp_points, 5)
    design_matrix[1:n_out:end, :] = hcat(cos.(phases), sin.(phases), cos.(2 .* phases), sin.(2 .* phases), ones(length(times)))
    # n_total_samp_points > n_samp_points ? design_matrix = vcat(design_matrix, zeros(n_total_samp_points - n_samp_points, size(design_matrix, 2))) : design_matrix = design_matrix
    x = general_lst_sq(design_matrix, data; Σ=covariance)
    return kep_signal_epicyclic(P, x .* data_unit)
end
fit_kepler(data, times, covariance, ks::kep_signal_epicyclic; kwargs...) = fit_kepler_epicyclic(data, times, covariance, ks.P; kwargs...)

"""
A linear formulation of a Keplerian RV signal for a given eccentricity, period and initial mean anomaly
for fitting non-(e, P, M0) components
see Wright and Howard 2009 "EFFICIENT FITTING OF MULTI-PLANET KEPLERIAN MODELS TO RADIAL VELOCITY AND ASTROMETRY DATA" (https://arxiv.org/abs/0904.3725)
coefficients[1] = K * cos(ω)
coefficients[2] = -K * sin(ω)
coefficients[3] = γ + K * e * cos(ω)
so
K = sqrt(coefficients[1]^2 + coefficients[2]^2)
ω = atan(-coefficients[2], coefficients[1])
γ = coefficients[3] - K * e * cos(ω) = coefficients[3] - e * coefficients[1]
"""
struct kep_signal_wright <: KeplerSignal
    K::Unitful.Velocity
    P::Unitful.Time
    M0::Real
    e::Real
    ω::Real  # ::Unitful.NoDims
    γ::Unitful.Velocity
    h::Real
    k::Real
    coefficients::Vector{T} where T<:Unitful.Velocity
    kep_signal_wright(; K=0u"m/s", P=1u"d", M0=0, e_or_h=0, ω_or_k=0, γ=0u"m/s", use_hk::Bool=false) = kep_signal_wright(K, P, M0, e_or_h, ω_or_k, γ; use_hk=use_hk)
    function kep_signal_wright(K, P, M0, e_or_h, ω_or_k, γ; use_hk::Bool=false)
        if use_hk
            h = e_or_h
            k = ω_or_k
            e, ω = hk_2_eω(h, k)
        else
            e = e_or_h
            ω = ω_or_k
            h, k = eω_2_hk(e, ω)
        end
        return kep_signal_wright(K, P, M0, e, ω, γ, h, k)
    end
    kep_signal_wright(K, P, M0, e, ω, γ, h, k) = kep_signal_wright(K, P, M0, e, ω, γ, h, k,
        [K * cos(ω), -K * sin(ω), γ + K * e * cos(ω)])
    function kep_signal_wright(P::Unitful.Time, M0::Real, e::Real, coefficients::Vector{T} where T<:Unitful.Velocity)
        K = sqrt(coefficients[1]^2 + coefficients[2]^2)
        ω = mod2pi(atan(-coefficients[2], coefficients[1]))
        γ = coefficients[3] - e * coefficients[1]
        h, k = eω_2_hk(e, ω)
        return kep_signal_wright(K, P, M0, e, ω, γ, h, k, coefficients)
    end
    function kep_signal_wright(
        K::Unitful.Velocity,
        P::Unitful.Time,
        M0::Real,
        e::Real,
        ω::Real,  # ::Unitful.NoDims
        γ::Unitful.Velocity,
        h::Real,
        k::Real,
        coefficients::Vector{T} where T<:Unitful.Velocity)

        M0 = mod2pi(M0)
        ω = mod2pi(ω)
        assert_positive(P)
        @assert length(coefficients) == 3 "wrong number of coefficients"
        return new(K, P, M0, e, ω, γ, h, k, coefficients)
    end
end
function kepler_rv_wright(t::Unitful.Time, ks::kep_signal_wright)
    ϕ_t = ϕ(t, ks.P, ks.M0, ks.e)
    return (ks.coefficients[1] * cos.(ϕ_t)) + (ks.coefficients[2] * sin.(ϕ_t)) + ks.coefficients[3]
end
kepler_rv(t::Unitful.Time, ks::kep_signal_wright) = kepler_rv_wright(t, ks)
# kepler_rv(t::Unitful.Time, ks::kep_signal_wright) = kepler_rv(t, ks.K, ks.P, ks.M0, ks.e, ks.ω; γ=ks.γ)
function fit_kepler_wright_linear_step(
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    P::Unitful.Time,
    M0::Real,
    e::Real;
    data_unit::Unitful.Velocity=1u"m/s",
    return_extra::Bool=false
    ) where T<:Real

    # if P==0; return kep_signal_epicyclic(p, P, M0, e, ω, γ, coefficients) end
    # println(P, M0, e)
    assert_positive(P, e)
    for i in 1:ndims(covariance)
        @assert size(covariance,i)==length(data) "covariance incompatible with data"
    end
    n_total_samp_points = length(data)
    n_samp_points = length(times)
    n_out = Int(n_total_samp_points / n_samp_points)
    # println(e)
    if return_extra
        E_t = ecc_anomaly.(times, P, M0, e)
        ϕ_t = ϕ.(E_t, e)
    else
        ϕ_t = ϕ.(times, P, M0, e)
    end
    design_matrix = zeros(n_total_samp_points, 3)
    design_matrix[1:n_out:end, :] = hcat(cos.(ϕ_t), sin.(ϕ_t), ones(n_samp_points))
    # n_total_samp_points > n_samp_points ? design_matrix = vcat(design_matrix, zeros(n_total_samp_points - n_samp_points, size(design_matrix, 2))) : design_matrix = design_matrix
    if return_extra
        β, ϵ_int, ϵ_inv = general_lst_sq(design_matrix, data; Σ=covariance, return_ϵ_inv=true)
        return kep_signal_wright(P, M0, e, β .* data_unit), ϵ_int, ϵ_inv, design_matrix, E_t
    else
        return kep_signal_wright(P, M0, e, general_lst_sq(design_matrix, data; Σ=covariance) .* data_unit)
    end
end
fit_kepler_wright_linear_step(
    glo_rv::GLO_RV,
    covariance::Union{Cholesky,Diagonal} where T<:Real,
    P::Unitful.Time,
    M0::Real,
    e::Real;
    kwargs...) =
    fit_kepler_wright_linear_step(glo_rv.GLO.y_obs, glo_rv.time, covariance, P, M0, e; data_unit=glo_rv.rv_factor, kwargs...)

mutable struct kep_buffer_wright{T1<:Real}
    ks::kep_signal_wright
    rm_kep::Vector{T1}
    nprior::T1
    ϵ_int::Matrix{T1}
    ϵ_inv::Matrix{T1}
    design_matrix::Matrix{T1}
    E_t::Vector{T1}
end

function ∇nlogL_kep!(
    G::Vector{T},
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    buffer::kep_buffer_wright;
    data_unit::Unitful.Velocity=1u"m/s",
    hold_P::Bool=false,
    include_priors::Bool=true) where T<:Real

    @assert length(G) == (3 - hold_P)

    α_data = covariance \ data
    α_rm_kep = covariance \ buffer.rm_kep
    n_total_samp_points = length(data)
    n_samp_points = length(times)
    n_out = Int(n_total_samp_points / n_samp_points)
    e = buffer.ks.e
    mod_e2 = (1 + e) * (1 - e)
    sqrt_mod_e2 = sqrt(mod_e2)

    n_out = Int(length(data) / length(times))
    F = buffer.design_matrix
	sin_ϕ = F[1:n_out:end, 2]
	cos_ϕ = F[1:n_out:end, 1]
    β = buffer.ks.coefficients ./ data_unit

    dFdx_holder = zeros(n_total_samp_points, 3)
    function dudx!(dFdx::Matrix{T}, dϕdx::Vector{T}) where T<:Real
        dFdx[1:n_out:end, 1] = -sin_ϕ .* dϕdx
        dFdx[1:n_out:end, 2] = cos_ϕ .* dϕdx
        dβdx = (-buffer.ϵ_inv \ (((dFdx' * buffer.ϵ_int) + (F' * (covariance \ dFdx))) * (buffer.ϵ_inv \ (F' * α_data)))) + (buffer.ϵ_inv \ (dFdx' * α_data))
        return (dFdx * β) + (F * dβdx), dβdx
    end

    factor = 1 .- (e .* cos.(buffer.E_t))
    factor2 = factor .* factor
    dϕdM0 = -sqrt_mod_e2 ./ factor2
    dϕde = (mod_e2 .+ factor) .* sin.(buffer.E_t) ./ (sqrt_mod_e2 .* factor2)

    dudM0, dβdM0 = dudx!(dFdx_holder, dϕdM0)
    dude, dβde = dudx!(dFdx_holder, dϕde)

    G[end-1] = GLOM.dnlogLdθ(-dudM0, α_rm_kep)
    G[end] = GLOM.dnlogLdθ(-dude, α_rm_kep)

    if !hold_P
        t = convert_and_strip_units.(unit(buffer.ks.P), times)
        P = ustrip(buffer.ks.P)
        dϕdP = (-2 * π * sqrt_mod_e2 / P / P) .* t ./ factor2
        dudP, dβdP = dudx!(dFdx_holder, dϕdP)
        G[1] = GLOM.dnlogLdθ(-dudP, α_rm_kep)
        # println("dβdP:  ", dβdP)
        # println("dβdM0: ", dβdM0)
        # println("dβde:  ", dβde)
    end

    if include_priors
        kep_prior_G = logprior_kepler_tot(buffer.ks; d_tot=1, use_hk=false)

        K = buffer.ks.K / data_unit
        # K = sqrt(β[1]^2 + β[2]^2)
        function dPriordx(dβdx::Vector{T}, δPriorδx::T, dedx::Integer) where T<:Real
            dKdx = (β[1] * dβdx[1] + β[2] * dβdx[2]) / K * ustrip(data_unit)
            dωdx = (β[2] * dβdx[1] - β[1] * dβdx[2]) / (K * K)
            dγdx = (dβdx[3] - e * dβdx[1] - dedx * β[1]) * ustrip(data_unit)
            return kep_prior_G[1] * dKdx + kep_prior_G[5] * dωdx + kep_prior_G[6] * dγdx + δPriorδx
        end

        if !hold_P; G[1] -= dPriordx(dβdP, kep_prior_G[2], 0) end
        G[end-1] -= dPriordx(dβdM0, kep_prior_G[3], 0)
        G[end] -= dPriordx(dβde, kep_prior_G[4], 1)
    end
end
function ∇nlogL_kep(data, times, covariance, init_ks::kep_signal_wright; kwargs...)
    ks, ϵ_int, ϵ_inv, design_matrix, E_t = fit_kepler_wright_linear_step(data, times, covariance, init_ks.P, init_ks.M0, init_ks.e; data_unit=data_unit, return_extra=true)
    buffer = kep_buffer_wright(ks, remove_kepler(data, times, ks; data_unit=data_unit), -logprior_kepler(ks; use_hk=false), ϵ_int, ϵ_inv, design_matrix, E_t)
    G = zeros(3 - hold_P)
    ∇nlogL_kep!(G, data, times, covariance, buffer; kwargs...)
    return G
end

"You may have to mess wit the fit alpha"
function fit_kepler_wright(
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    init_ks::kep_signal_wright;
    data_unit::Unitful.Velocity=1u"m/s",
    print_stuff::Bool=true,
    hold_P::Bool=false,
    avoid_saddle::Bool=true,
    fit_alpha::Real=3e-3,
    include_priors::Bool=true
    ) where T<:Real

    optim_cb_local(x::OptimizationState) = optim_cb(x; print_stuff=print_stuff)

    # print("its this function")
    assert_positive(init_ks.P, init_ks.e)
    for i in 1:ndims(covariance)
        @assert size(covariance,i)==length(data) "covariance incompatible with data"
    end

    n_total_samp_points = length(data)
    n_samp_points = length(times)
    normalization = logdet(covariance) + n_total_samp_points * log(2 * π)
    P_u = unit(init_ks.P)

    ks, ϵ_int, ϵ_inv, design_matrix, E_t = fit_kepler_wright_linear_step(data, times, covariance, init_ks.P, init_ks.M0, init_ks.e; data_unit=data_unit, return_extra=true)
    if hold_P
        current_x = ustrip.([ks.M0, ks.e])
    else
        current_x = ustrip.([ks.P, ks.M0, ks.e])
    end
    buffer = kep_buffer_wright(ks, remove_kepler(data, times, ks; data_unit=data_unit), -logprior_kepler(ks; use_hk=false), ϵ_int, ϵ_inv, design_matrix, E_t)
    last_x = similar(current_x)

    function calculate_common!(x::Vector{T}, last_x::Vector{T}, buffer::kep_buffer_wright) where {T<:Real}
        if x != last_x
            # copy!(last_x, x)
            last_x[:] = x
            if (0 > x[end] || x[end] > 1)  # is 0 <= e <= 1 ?
                buffer.nprior = Inf
            else
                if hold_P
                    buffer.ks, buffer.ϵ_int, buffer.ϵ_inv, buffer.design_matrix, buffer.E_t = fit_kepler_wright_linear_step(data, times, covariance, init_ks.P, x[1], x[2]; data_unit=data_unit, return_extra=true)
                else
                    buffer.ks, buffer.ϵ_int, buffer.ϵ_inv, buffer.design_matrix, buffer.E_t = fit_kepler_wright_linear_step(data, times, covariance, x[1] * P_u, x[2], x[3]; data_unit=data_unit, return_extra=true)
                end
                buffer.nprior = -logprior_kepler(buffer.ks; use_hk=false)
            end
            buffer.nprior == Inf ? buffer.rm_kep = zeros(length(buffer.rm_kep)) : buffer.rm_kep = remove_kepler(data, times, buffer.ks; data_unit=data_unit)
        end
        # println("buffer: ", kep_parms_str(buffer.ks))
        # println(buffer.rm_kep[1:40:end])
    end

    function f(x::Vector{T}, buffer::kep_buffer_wright, last_x::Vector{T}) where {T<:Real}
        calculate_common!(x, last_x, buffer)
        if buffer.nprior == Inf
            return buffer.nprior
        else
            return GLOM.nlogL(covariance, buffer.rm_kep, nlogL_normalization=normalization) + (include_priors * buffer.nprior)
        end
    end
    f(x) = f(x, buffer, last_x)

    function g!(G::Vector{T}, x::Vector{T}, buffer::kep_buffer_wright, last_x::Vector{T}) where {T<:Real}
        calculate_common!(x, last_x, buffer)
        if buffer.nprior == Inf
            G[:] .= 0
        else
            ∇nlogL_kep!(G, data, times, covariance, buffer; data_unit=data_unit, hold_P=hold_P, include_priors=include_priors)
        end
    end
    g!(G, x) = g!(G, x, buffer, last_x)

    if avoid_saddle
        attempts = 0
        in_saddle = true
        while attempts < 10 && in_saddle
            attempts += 1
            if attempts > 1;
                if print_stuff; println("found saddle point. starting attempt $attempts with a perturbation") end
                if !hold_P; current_x[1] *= centered_rand(; scale=2e-2, center=1) end
                current_x[end-1] = mod2pi(current_x[end-1] + centered_rand(; scale=4e-1))
                current_x[end] = 0.05 * rand()
            end
            # println(current_x)
            result = optimize(f, g!, current_x, LBFGS(alphaguess=LineSearches.InitialStatic(alpha=fit_alpha)), Optim.Options(;callback=optim_cb_local)) # 27s
            current_x = copy(result.minimizer)
            ks = fit_kepler_wright_linear_step(data, times, covariance, buffer.ks.P, buffer.ks.M0, buffer.ks.e; data_unit=data_unit)
            if print_stuff; println("fit attempt $attempts: "kep_parms_str(ks)) end
            # println(∇∇nlogL_kep(data, times, covariance, ks; data_unit=data_unit))
            new_det = det(∇∇nlogL_kep(data, times, covariance, kep_signal(ks.K, ks.P, ks.M0, ks.e, ks.ω, ks.γ); data_unit=data_unit))
            if print_stuff; println("determinant: ", new_det) end
            in_saddle = new_det <= 0
            if !in_saddle; return ks end
        end
        @error "no non-saddle point soln found"
    else
        try
            result = optimize(f, g!, current_x, LBFGS(alphaguess=LineSearches.InitialStatic(alpha=fit_alpha)), Optim.Options(;callback=optim_cb_local)) # 27s
            return fit_kepler_wright_linear_step(data, times, covariance, buffer.ks.P, buffer.ks.M0, buffer.ks.e; data_unit=data_unit)
        catch
            # print("Kepler fit failed")
            return nothing
        end
    end

end
fit_kepler(data, times, covariance, init_ks::kep_signal_wright; kwargs...) =
    fit_kepler_wright(data, times, covariance, init_ks; kwargs...)

"Convert the solar phase information from SOAP 2.0 into days"
function convert_SOAP_phases(new_unit::Unitful.FreeUnits, phase::Real; P_rot::Unitful.Time=25.05u"d")
    # default P_rot is the solar rotation period used by SOAP 2.0 in days
    assert_positive(P_rot)
    return uconvert(new_unit, phase * P_rot)
end

function remove_kepler(
    data::Vector{<:Real},
    times::Vector{T2} where T2<:Unitful.Time,
    ks::KeplerSignal;
    data_unit::Unitful.Velocity=1u"m/s",
    d::Vector{<:Integer}=zeros(Int,n_kep_parms))

    validate_kepler_dorder(d)
    n_out = Int(length(data) / length(times))
    if findfirst(!iszero, d) == nothing
        y = copy(data)
        if ustrip(ks.K) != 0; y[1:n_out:end] -= uconvert.(unit(data_unit), ks.(times)) ./ data_unit end
    else  # this uses h and k not e and ω
        @assert typeof(ks) == kep_signal
        y = zeros(length(data))
        kep_deriv_simple(t) = kep_deriv(ks, t, d)
        y[1:n_out:end] -= ustrip.(kep_deriv_simple.(times)) ./ convert_and_strip_units(unit(ks.K), data_unit)
    end
    return y

end
remove_kepler(
    data::Vector{<:Unitful.Velocity},
    times::Vector{T2} where T2<:Unitful.Time,
    ks::KeplerSignal;
    kwargs...) =
    remove_kepler(ustrip.(data), times, ks; data_unit=1.0*unit(data[1]), kwargs...)
remove_kepler(
    glo_rv::GLO_RV,
    ks::KeplerSignal;
    kwargs...) =
    remove_kepler(glo_rv.GLO.y_obs, glo_rv.time, ks; data_unit=glo_rv.rv_factor, kwargs...)

fit_kepler(
    glo_rv::GLO_RV,
    covariance::Union{Cholesky,Diagonal},
    ks::KeplerSignal;
    kwargs...) where T<:Real =
    fit_kepler(glo_rv.GLO.y_obs, glo_rv.time, covariance, ks; data_unit=glo_rv.rv_factor, kwargs...)

kep_signal(ks::KeplerSignal) = kep_signal(ks.K, ks.P, ks.M0, ks.e, ks.ω, ks.γ)

kep_parms_str(ks::KeplerSignal) =
    "K: $(round(convert_and_strip_units(u"m/s", ks.K), digits=2))" * "m/s" * "  P: $(round(convert_and_strip_units(u"d", ks.P), digits=2))" * "d" * "  M0: $(round(ks.M0,digits=2))  e: $(round(ks.e,digits=2))  ω: $(round(ks.ω,digits=2)) γ: $(round(convert_and_strip_units(u"m/s", ks.γ), digits=2))" * "m/s"
kep_parms_str_short(ks::KeplerSignal) =
    "K: $(round(convert_and_strip_units(u"m/s", ks.K), digits=2))" * "m/s" * "  P: $(round(convert_and_strip_units(u"d", ks.P), digits=2))" * "d" * "  e: $(round(ks.e,digits=2))"

function ∇∇nlogL_kep(
    data::Vector{T},
    times::Vector{T2} where T2<:Unitful.Time,
    covariance::Union{Cholesky,Diagonal},
    ks::kep_signal;
    data_unit::Unitful.Velocity=1u"m/s",
    include_priors::Bool=false) where T<:Real

    y = remove_kepler(data, times, ks; data_unit=data_unit)
    α = covariance \ y
    H = zeros(n_kep_parms, n_kep_parms)
    for i in 1:n_kep_parms
        for j in 1:n_kep_parms
            if i <= j
                d = zeros(Int, n_kep_parms)
                d[j] += 1
                y2 = remove_kepler(data, times, ks; data_unit=data_unit, d=d)
                d[i] += 1
                y12 = remove_kepler(data, times, ks; data_unit=data_unit, d=d)
                d[j] -= 1
                y1 = remove_kepler(data, times, ks; data_unit=data_unit, d=d)
                H[i, j] = GLOM.d2nlogLdθ(y2, y12, α, covariance \ y1)
            end
        end
    end

    if include_priors
        H[:,:] -= logprior_kepler_tot(ks; d_tot=2, use_hk=true)
    end

    return Symmetric(H)
end

function ∇∇nlogL_GLOM_and_planet!(
    workspace::GLOM.nlogL_matrix_workspace,
    glo_rv::GLO_RV,
    total_hyperparameters::Vector{<:Real},
    ks::kep_signal;
    include_kepler_priors::Bool=false)

    GLOM.calculate_shared_∇nlogL_matrices!(workspace, glo_rv.GLO, total_hyperparameters)

    non_zero_inds = copy(glo_rv.GLO.non_zero_hyper_inds)
    n_hyper = length(non_zero_inds)
    full_H = zeros(n_kep_parms + n_hyper, n_kep_parms + n_hyper)
    full_H[1:n_hyper, 1:n_hyper] = GLOM.∇∇nlogL_GLOM(
        glo_rv.GLO, total_hyperparameters; Σ_obs=workspace.Σ_obs, y_obs=remove_kepler(glo_rv, ks))

    full_H[n_hyper+1:end,n_hyper+1:end] = ∇∇nlogL_kep(glo_rv.GLO.y_obs, glo_rv.time, workspace.Σ_obs, ks; data_unit=glo_rv.rv_factor, include_priors=include_kepler_priors)

    # TODO allow y and α to be passed to ∇∇nlogL_kep
    y = remove_kepler(glo_rv, ks)
    α = workspace.Σ_obs \ y
    for (i, nzind1) in enumerate(non_zero_inds)
        for j in 1:n_kep_parms
            d = zeros(Int, n_kep_parms)
            d[j] += 1
            y1 = remove_kepler(glo_rv, ks; d=d)
            full_H[i, j + n_hyper] = GLOM.d2nlogLdθ(y, y1, α, workspace.Σ_obs \ y1, workspace.βs[i])
        end
    end

    return Symmetric(full_H)
end
∇∇nlogL_GLOM_and_planet(glo_rv, total_hyperparameters, ks; kwargs...) =
    ∇∇nlogL_GLOM_and_planet!(
        GLOM.nlogL_matrix_workspace(glo_rv.GLO, total_hyperparameters),
        glo_rv, total_hyperparameters, ks; kwargs...)

function validate_kepler_dorder(d::Vector{<:Integer})
	@assert sum(d) < 3
	@assert minimum(d) == 0
	@assert length(d) == n_kep_parms
end

function validate_kepler_wright_dorder(d::Vector{<:Integer})
	@assert sum(d) < 2
	@assert minimum(d) == 0
	@assert length(d) == 3
end

function ks_from_vec(vector::Vector{<:Real}, K_u::Unitful.VelocityFreeUnits, P_u::Unitful.TimeFreeUnits, γ_u::Unitful.VelocityFreeUnits; use_hk::Bool=false)
    @assert length(vector) == n_kep_parms
    return kep_signal(vector[1] * K_u, vector[2] * P_u, vector[3], vector[4], vector[5], vector[6] * γ_u; use_hk=use_hk)
end


text:
"""
Construct adjacency map for a cartesian Nx × Ny lattice.
"""
function lattice_adjacency_map(Nx::Integer, Ny::Integer; pbc=true)
    L = Nx * Ny
    adjacency = zeros(Int, (L, L))
    for j in 0:Ny-1
        j_next = (j+1) % Ny
        for i in 0:Nx-1
            i_next = (i+1) % Nx
            # nearest neighbors
            if (pbc || i_next > 0) adjacency[j*Nx + i + 1, j*Nx + i_next + 1] = 1; end
            if (pbc || j_next > 0) adjacency[j*Nx + i + 1, j_next*Nx + i + 1] = 1; end
        end
    end
    adjacency = adjacency + transpose(adjacency)
    # only 0 or 1 entries
    return (adjacency .≠ 0)
end

"""
Construct horizontal adjacency map for a cartesian Nx × Ny lattice.
"""
function lattice_adjacency_map_horz(Nx::Integer, Ny::Integer; pbc=true)
    L = Nx * Ny
    adjacency = zeros(Int, (L, L))
    for j in 0:Ny-1
        for i in 0:Nx-1
            i_next = (i+1) % Nx
            # nearest neighbors along horizontal direction
            if (pbc || i_next > 0) adjacency[j*Nx + i + 1, j*Nx + i_next + 1] = 1; end
        end
    end
    adjacency = adjacency + transpose(adjacency)
    # only 0 or 1 entries
    return (adjacency .≠ 0)
end

"""
Construct vertical adjacency map for a cartesian Nx × Ny lattice.
"""
function lattice_adjacency_map_vert(Nx::Integer, Ny::Integer; pbc=true)
    L = Nx * Ny
    adjacency = zeros(Int, (L, L))
    for j in 0:Ny-1
        j_next = (j+1) % Ny
        for i in 0:Nx-1
            # nearest neighbors along vertical direction
            if (pbc || j_next > 0) adjacency[j*Nx + i + 1, j_next*Nx + i + 1] = 1; end
        end
    end
    adjacency = adjacency + transpose(adjacency)
    # only 0 or 1 entries
    return (adjacency .≠ 0)
end

function lattice_map_1d(Nx::Integer, Ny::Integer)
    L = Nx * Ny
    adjacency = zeros(Int, (L, L))
    for j in 1:L
        adjacency[j,j] = 1
    end
    # only 0 or 1 entries
    return (adjacency .≠ 0)
end

ext: jl

max_stars_repo_name: cmendl/density-matrix-bloch-qcircuit

max_issues_repo_name: cmendl/density-matrix-bloch-qcircuit

max_forks_repo_name: cmendl/density-matrix-bloch-qcircuit


text:
function life(rows, cols, some, generations)
    now = trues(rows*cols)
    for c=1:rows*cols
        now[c] = rand() < some
    end
    live(now,rows,generations)
end

function live(a,r,gen)
    if gen < 1
        return nothing
    end
    sleep(0.1)
    homeScreen()
    println("Generation ",lpad(gen,3,"0"))
    for c=1:length(a)
        if(a[c])
          print("o")
        else
          print(" ")
        end
        if mod(c,r)==0
            print("\n")
        end
    end
    b = falses(length(a))
    for c=1:length(a)
        neighbours = 1
        if c-1 >= 1
            neighbours += a[c-1]
        end
        if c+1 <= 1000
            neighbours += a[c+1]
        end
        if c-r-1 >= 1
            neighbours += a[c-r-1]
        end
        if c-r >= 1 && c-r <=1000
            neighbours += a[c-r]
        end
        if c-r+1 >= 1  && c-r+1 <=1000
            neighbours += a[c-r+1]
        end
        if c+r-1 >=1 && c+r-1 <= 1000
            neighbours += a[c+r-1]
        end
        if c+r <= 1000
            neighbours += a[c+r]
        end
        if c+51 <= 1000
            neighbours += a[c+51]
        end
        b[c] = a[c]
        if a[c]==0
            b[c] = (neighbours==3)
        else
          b[c] = (neighbours==2)
        end
    end
    gen = gen-1
    live(b,r,gen)
end

function homeScreen()
    print("\033[1;1H")
end

function clearScreen()
    run(`clear`)
end

sleep(0.1)
clearScreen()
life(50,20,0.6,200)

max_stars_repo_name: mukulv/gameOfLife

max_issues_repo_name: mukulv/gameOfLife

max_forks_repo_name: mukulv/gameOfLife


text:
"""A relative difference between prediction and observation.
The absolute value of the difference between the prediction and the
observation is divided by a reference value with the same units. This
reference scale should be chosen for each test such that normalization
produces directly comparable scores across tests. For example, if 5 volts
represents a medium size difference for TestA, and 10 seconds represents a
medium size difference for TestB, then 5 volts and 10 seconds should be
used for this reference scale in TestA and TestB, respectively. The
attribute `scale` can be passed to the compute method or set for the whole
class in advance. Otherwise, a scale of 1 (in the units of the
observation and prediction) will be used.
"""

#    _allowed_types = (float,)

#    _description = ('The relative difference between the prediction and the observation')
_best = 0.0  # A RelativeDifferenceScore of 0.0 is best
_worst = inf
scale = None

function _check_score(self, score):
        if score < 0.0:
            raise errors.InvalidScoreError(("RelativeDifferenceScore was initialized with "
                                            "a score of %f, but a RelativeDifferenceScore "
                                            "must be non-negative.") % score)
        end
end
function compute(cls, observation: Union[dict, float, int, pq.Quantity],
                 prediction: Union[dict, float, int, pq.Quantity],
                 key=None,
                 scale: Union[float, int, pq.Quantity, None] = None) -> 'RelativeDifferenceScore':
    """Compute the relative difference between the observation and a prediction.
    Returns:
        RelativeDifferenceScore: A relative difference between an observation and a prediction.
    """
    assert isinstance(observation, (dict, float, int, pq.Quantity))
    assert isinstance(prediction, (dict, float, int, pq.Quantity))

    obs, pred = cls.extract_means_or_values(observation, prediction,
                                            key=key)

    scale = scale or cls.scale or (obs/float(obs))
    assert type(obs) is type(scale)
    assert type(obs) is type(pred)

    if isinstance(obs, pq.Quantity):
        assert obs.units == pred.units, \
            "Prediction must have the same units as the observation"
        assert obs.units == scale.units, \
            "RelativeDifferenceScore.Scale must have the same units as the observation"
        pred = pred.rescale(obs.units)

    assert scale > 0, \
        "RelativeDifferenceScore.scale must be positive (not %g)" % scale
    value = np.abs(pred - obs) / scale
    value = utils.assert_dimensionless(value)
    return RelativeDifferenceScore(value)
end

function norm_score(self) -> float:
        """Return 1.0 for a ratio of 0.0, falling to 0.0 for extremely large values.
        Returns:
            float: The value of the norm score.
        """
        x = self.score
        return 1 / (1+x)
end
function __str__(self):
        return 'Relative Difference = %.2f' % self.score
end


text:
@testset "auxiliar" begin

    m = @agent 2 [x1,y1]::Local
    @test_nowarn AgentBasedModels.checkIsDeclared_(m,:x1)
    @test_nowarn AgentBasedModels.checkIsDeclared_(m,[:x1])
    @test_nowarn AgentBasedModels.checkIsDeclared_(m,[:x1,:y1])
    @test_throws ErrorException AgentBasedModels.checkIsDeclared_(m,:g)
    @test_throws ErrorException AgentBasedModels.checkIsDeclared_(m,[:g])
    @test_throws ErrorException AgentBasedModels.checkIsDeclared_(m,[:x1,:g])

    @test AgentBasedModels.subs_(:(x += 3*x),:x,:y) == :(y += 3*y)
    @test AgentBasedModels.subs_(:(x += 3*x),:x,:y,update=true) == :(y += 3*x)

    m = @agent(
        0,
        l::Local,
        g::Global,
        h::GlobalArray
    )
    p = AgentBasedModels.Program_(m);
    AgentBasedModels.updates_!(p)
    @test AgentBasedModels.vectorize_(m,:(l = 1), p) == :(localV[ic1_,1]=1)
    @test AgentBasedModels.vectorize_(m,:(l.i = 1), p) == :(localV[ic1_,1]=1)
    @test AgentBasedModels.vectorize_(m,:(l.j = 1), p) == :(localV[nnic2_,1]=1)
    @test AgentBasedModels.vectorize_(m,:(g = 1), p) == :(globalV[1]=1)
    @test AgentBasedModels.vectorize_(m,:(h[1,2] = 1), p) == :(h[1,2]=1)
    @test AgentBasedModels.vectorize_(m,:(id = 1), p) == :(identityV[ic1_,1]=1)
    @test AgentBasedModels.vectorize_(m,:(id.i = 1), p) == :(identityV[ic1_,1]=1)
    @test AgentBasedModels.vectorize_(m,:(id.j = 1), p) == :(identityV[nnic2_,1]=1)

    @test AgentBasedModels.vectorize_(m,:(l = l), p) == :(localV[ic1_,1]=localV[ic1_,1])
    @test AgentBasedModels.vectorize_(m,:(l.i = l.i), p) == :(localV[ic1_,1]=localV[ic1_,1])
    @test AgentBasedModels.vectorize_(m,:(l.j = l.j), p) == :(localV[nnic2_,1]=localV[nnic2_,1])
    @test AgentBasedModels.vectorize_(m,:(g = g), p) == :(globalV[1]=globalV[1])
    @test AgentBasedModels.vectorize_(m,:(h[1,2] = h[1,2]), p) == :(h[1,2]=h[1,2])
    @test AgentBasedModels.vectorize_(m,:(id = id), p) == :(identityV[ic1_,1]=identityV[ic1_,1])
    @test AgentBasedModels.vectorize_(m,:(id.i = id.i), p) == :(identityV[ic1_,1]=identityV[ic1_,1])
    @test AgentBasedModels.vectorize_(m,:(id.j = id.j), p) == :(identityV[nnic2_,1]=identityV[nnic2_,1])

    @test begin
        m = @agent(
            0,
            l::Local,
            g::Global,
            h::GlobalArray
        )

        prod([i in [:t,:N,:simulationBox,:loc_,:glob_,:id_,:h] for i in AgentBasedModels.agentArguments_(m)])
    end

    @test AgentBasedModels.cudaAdapt_(:(sin(e^2^x))) == :(CUDA.sin(Float32(e)^Float32(2)^x))
    @test AgentBasedModels.cudaAdapt_(:(zeros(e^2))) == :(CUDA.zeros(e^2))

    @test begin
        code = :(x .+= 1)
        f = AgentBasedModels.wrapInFunction_(:func,code)
        f = AgentBasedModels.subs_(f,:ARGS_,:x)
        eval(quote
            $f
            y = ones(10)
            func(y)
            y == 2*ones(10)
        end)
    end

    @test begin
        code = :(x[ic1_] += 1)
        code = AgentBasedModels.simpleFirstLoop_("cpu",code)
        f = AgentBasedModels.wrapInFunction_(:func,code)
        f = AgentBasedModels.subsArguments_(f,:ARGS_,[:x,:N])
        eval(quote
            $f
            N = 10
            y = ones(N)
            func(y,N)
            y == 2*ones(N)
        end)
    end

    @test begin
        code = :(x[ic1_] += 1)
        f = AgentBasedModels.simpleFirstLoopWrapInFunction_("cpu",:func,code)
        f = AgentBasedModels.subsArguments_(f,:ARGS_,[:x,:N])
        eval(quote
            $f
            N = 10
            y = ones(N)
            func(y,N)
            y == 2*ones(N)
        end)
    end

    if CUDA.has_cuda()
        @test begin
            code = :(x[ic1_] += 1)
            code = AgentBasedModels.simpleFirstLoop_("gpu",code)
            f = AgentBasedModels.wrapInFunction_(:func,code)
            f = AgentBasedModels.subsArguments_(f,:ARGS_,[:x,:N])
            eval(quote
                $f
                N = 10
                y = CUDA.ones(N)
                CUDA.@cuda func(y,N)
                Array(y) == 2*ones(N)
            end)
        end

        @test begin
            code = :(x[ic1_] += 1)
            f = AgentBasedModels.simpleFirstLoopWrapInFunction_("gpu", :func, code)
            f = AgentBasedModels.subsArguments_(f,:ARGS_,[:x,:N])
            eval(quote
                $f
                N = 10
                y = CUDA.ones(N)
                CUDA.@cuda func(y,N)
                Array(y) == 2*ones(N)
            end)
        end
    end

    @test_nowarn begin

        abm = @agent(
            1,

            [l1,l2]::Local,
            [g1,g2]::Global,
            [id1,id2]::Identity,
            [ga1,ga2]::GlobalArray,

            UpdateLocal =
            begin
                l1 += 1
                l2 = 3
                id1 = 1
            end,

            Equation =
            begin
                d(x) = a*dt + b*dW
            end,

            UpdateGlobal =
            begin
                g1 += 1
                g2 = 0
                ga1[1,2] += 1
            end
        )
        p = AgentBasedModels.Program_(abm);

        AgentBasedModels.updates_!(p)

        if p.update["Identity"] != Dict(:id1=>1) error() end
        if p.update["Global"] != Dict(:g1=>1,:g2=>2) error() end
        if p.update["GlobalArray"] != Dict(:ga1=>1) error() end
        if p.update["Local"] != Dict(:l2=>3,:l1=>2,:x=>1) error() end
        if p.update["Variables"] != Dict(:x=>1) error() end
    end

    @test AgentBasedModels.extract(:(f(x)*1+4),:f) == (:x)

end


text:
#' # MNIST: Knet.jl CNN
info("MNIST: Knet.jl CNN example") #jl-only

#' In this tutorial we will adapt the
#' [MNIST example](http://denizyuret.github.io/Knet.jl/latest/tutorial.html#Convolutional-neural-network-1)
#' from [Knet.jl](https://github.com/denizyuret/Knet.jl)
#' to utilize a custom augmentation pipeline.
#' In order to showcase the effect that image augmentation can
#' have on a neural network's ability to generalize, we will
#' limit the training set to just the first 500 images (of the
#' available 60,000!). For more information on the dataset see
#' [^MNIST1998].

#md #' !!! note
#md #'
#md #'     This tutorial is also available as a
#md #'     [Juypter](https://jupyter.org/) notebook. You can
#md #'     find a link to the Juypter version of this tutorial
#md #'     in the top right corner of this page.

#' ## Preparing the MNIST dataset
info("Preparing the MNIST dataset") #jl-only

#' In order to access, prepare, and visualize the MNIST images we
#' employ the help of three additional Julia packages. In the
#' interest of time and space we will not go into great detail
#' about their functionality. Feel free to click on their
#' respective names to find out more information about the
#' utility they can provide.
#'
#' - [MLDatasets.jl](https://github.com/JuliaML/MLDatasets.jl)
#'   has an MNIST submodule that offers a convenience interface
#'   to read the MNIST database.
#'
#' - [Images.jl](https://github.com/JuliaImages/Images.jl) will
#'   provide us with the necessary tools to process and display
#'   the image data in Julia / Juypter.
#'
#' - [MLDataUtils.jl](https://github.com/JuliaML/MLDataUtils.jl)
#'   implements a variety of functions to convert and partition
#'   Machine Learning datasets. This will help us prepare the
#'   MNIST data to be used with Knet.jl.

using Images, MLDatasets, MLDataUtils, Random
Random.seed!(42);
#md nothing # hide

#' As you may have seen previously in the
#' [elastic distortions tutorial](@ref elastic), the function
#' `MNIST.traintensor` returns the MNIST training images
#' corresponding to the given indices as a multi-dimensional
#' array. These images are stored in the native horizontal-major
#' memory layout as a single array. Because we specify that
#' the `eltype` of that array should be `Float32`, all the
#' individual values are scaled to be between `0.0` and `1.0`.
#' Also note, how the observations are laid out along the last
#' array dimension

@show summary(MNIST.traintensor(Float32, 1:500));
#md nothing # hide

#' The corresponding label of each image is stored as an integer
#' value between `0` and `9`. That means that if the label has
#' the value `3`, then the corresponding image is known to be a
#' handwritten "3". To show a more concrete example, the
#' following code reveals that the first training image denotes a
#' "5" and the second training image a "0" (etc).

@show summary(MNIST.trainlabels(1:500))
println("First eight labels: ", join(MNIST.trainlabels(1:8),", "))

#' For Knet we will require a slightly format for the images
#' and also the labels. More specifically, we add an additional
#' singleton dimension of length 1 to our image array. Think of
#' this as our single color channel (because MNIST images are gray).
#' Additionally we will convert our labels to proper 1-based indices.
#' This is because some functions provided by Knet expect the labels
#' to be in this format. We will do all this by creating a little
#' utility function that we will name `prepare_mnist`.

"""
    prepare_mnist(images, labels) -> (X, Y)

Change the dimension layout x1×x2×N of the given array
`images` to x1×x2×1×N and return the result as `X`.
The given integer vector `labels` is transformed into
an integer vector denoting 1-based class indices.
"""
function prepare_mnist(images, labels)
    X = reshape(images, (28, 28, 1, :))
    Y = convertlabel(LabelEnc.Indices{Int8}, labels, 0:9)
    X, Y
end
#md nothing # hide

#' With `prepare_mnist` defined, we can now use it in conjunction
#' with the functions in the `MLDatasets.MNIST` sub-module to load
#' and prepare our training set. Recall that for this tutorial only
#' the first 500 images of the training set will be used.

train_x, train_y = prepare_mnist(MNIST.traintensor(Float32, 1:500), MNIST.trainlabels(1:500))
@show summary(train_x) summary(train_y);
[MNIST.convert2image(train_x[:,:,1,i]) for i in 1:8]
#md tmp = hcat(ans...) # hide
#md save("mnist_knet_train.png",repeat(tmp, inner=(4,4))) # hide
#md nothing # hide

#md #' ![training images](mnist_knet_train.png)

#' Similarly, we use `MNIST.testtensor` and `MNIST.testlabels`
#' to load the full MNIST test set. We will utilize that data to
#' measure how well the network is able to generalize with and
#' without augmentation.

test_x, test_y = prepare_mnist(MNIST.testtensor(Float32), MNIST.testlabels())
@show summary(test_x) summary(test_y);
[MNIST.convert2image(test_x[:,:,1,i]) for i in 1:8]
#md tmp = hcat(ans...) # hide
#md save("mnist_knet_test.png",repeat(tmp, inner=(4,4))) # hide
#md nothing # hide

#md #' ![test images](mnist_knet_test.png)

#' ## Defining the Network
info("Defining the Network") #jl-only

#' With the dataset prepared, we can now define and instantiate
#' our neural network. To keep things simple, we will use the same
#' convolutional network as defined in the
#' [MNIST example](http://denizyuret.github.io/Knet.jl/latest/tutorial.html#Convolutional-neural-network-1)
#' of the Knet.jl package.

using Knet
#md nothing # hide

#' The first thing we will do is define the forward pass through
#' the network. This will effectively outline the computation
#' graph of the network architecture. Note how this does not
#' define some details, such as the number of neurons per layer.
#' We will define those later when initializing our
#' vector of weight arrays `w`.

"""
    forward(w, x) -> a

Compute the forward pass for the given minibatch `x` by using the
neural network parameters in `w`. The resulting (unnormalized)
activations of the last layer are returned as `a`.
"""
function forward(w, x)
    # conv1 (2x2 maxpool)
    a1 = pool(relu.(conv4(w[1], x)  .+ w[2]))
    # conv2 (2x2 maxpool)
    a2 = pool(relu.(conv4(w[3], a1) .+ w[4]))
    # dense1 (relu)
    a3 = relu.(w[5] * mat(a2) .+ w[6])
    # dense2 (identity)
    a4 = w[7] * a3 .+ w[8]
    return a4
end
#md nothing # hide

#' In order to be able to train our network we need to choose a
#' cost function. Because this is a classification problem we will
#' use the negative log-likelihood (provided by `Knet.nll`).
#' With the cost function defined we can the simply use the
#' higher-order function `grad` to create a new function `costgrad`
#' that computes us the corresponding gradients.

"""
    cost(w, x, y) -> AbstractFloat

Compute the per-instance negative log-likelihood for the data
in the minibatch `(x, y)` given the network with the current
parameters in `w`.
"""
cost(w, x, y) = nll(forward(w, x), y)
costgrad = grad(cost)
#md nothing # hide

#' Aside from the cost function that we need for training, we
#' would also like a more interpretable performance measurement.
#' In this tutorial we will use "accuracy" for its simplicity
#' and because we know that the class distribution for MNIST
#' is close to uniform.

"""
    acc(w, X, Y; [batchsize]) -> Float64

Compute the accuracy for the data in `(X,Y)` given the network
with the current parameters in `w`. The resulting value is
computed by iterating over the data in minibatches of size
`batchsize`.
"""
function acc(w, X, Y; batchsize = 100)
    sum = 0; count = 0
    for (x_cpu, y) in eachbatch((X, Y), maxsize = batchsize)
        x = KnetArray{Float32}(x_cpu)
        sum += Int(accuracy(forward(w,x), y, average = false))
        count += length(y)
    end
    return sum / count
end
#md nothing # hide

#' Before we can train or even just use our network, we need to
#' define how we initialize `w`, which is our the vector of
#' parameter arrays. The dimensions of these individual arrays
#' specify the filter sizes and number of neurons.
#' It can be helpful to compare the indices here with the indices
#' used in our `forward` function to see which array corresponds
#' to which computation node of our network.

function weights(atype = KnetArray{Float32})
    w = Array{Any}(undef, 8)
    # conv1
    w[1] = xavier(5,5,1,20)
    w[2] = zeros(1,1,20,1)
    # conv2
    w[3] = xavier(5,5,20,50)
    w[4] = zeros(1,1,50,1)
    # dense1
    w[5] = xavier(500,800)
    w[6] = zeros(500,1)
    # dense2
    w[7] = xavier(10,500)
    w[8] = zeros(10,1)
    return map(a->convert(atype,a), w)
end
#md nothing # hide

#' ## Training without Augmentation
info("Training baseline network without augmentation") #jl-only

#' In order to get an intuition for how useful augmentation can
#' be, we need a sensible baseline to compare to. To that end, we
#' will first train the network we just defined using only the
#' (unaltered) 500 training examples.

#' The package
#' [ValueHistories.jl](https://github.com/JuliaML/ValueHistories.jl)
#' will help us record the accuracy during the training process.
#' We will use those logs later to visualize the differences
#' between having augmentation or no augmentation.

using ValueHistories
using ProgressMeter #jl-only

#' To keep things simple, we will not overly optimize our
#' training function. Thus, we will be content with using a
#' closure. Because both, the baseline and the augmented version,
#' will share this "inefficiency", we should still get a decent
#' enough picture of their performance differences.

function train_baseline(; epochs = 500, batchsize = 100, lr = .03)
    w = weights()
    log = MVHistory()
    p = Progress(epochs, desc = "Baseline: ") #jl-only
    for epoch in 1:epochs
        for (batch_x_cpu, batch_y) in eachbatch((train_x ,train_y), batchsize)
            batch_x = KnetArray{Float32}(batch_x_cpu)
            g = costgrad(w, batch_x, batch_y)
            Knet.update!(w, g, lr = lr)
        end

        next!(p) #jl-only
        if (epoch % 5) == 0
            train = acc(w, train_x, train_y)
            test  = acc(w, test_x,  test_y)
            @trace log epoch train test
            msg = "epoch " * lpad(epoch,4) * ": train accuracy " * rpad(round(train, digits=3),5,"0") * ", test accuracy " * rpad(round(test,digits=3),5,"0")
            cancel(p, msg, :blue) #jl-only
#md             println(msg)
#jp             println(msg)
        end
    end
    finish!(p) #jl-only
    log
end
#md nothing # hide

#' Aside from the accuracy, we will also keep an eye on the
#' training time. In particular we would like to see if and how
#' the addition of augmentation causes our training time to
#' increase.

train_baseline(epochs=1) # warm-up
baseline_log = @time train_baseline(epochs=200);
#md nothing # hide

#' As we can see, the accuracy on the training set is around a
#' 100%, while the accuracy on the test set peaks around 90%. For
#' a mere 500 training examples, this isn't actually that bad of
#' a result.

#' ## Integrating Augmentor
info("Training network with augmentation") #jl-only

#' Now that we have a network architecture with a baseline to
#' compare to, let us finally see what it takes to add Augmentor
#' to our experiment. First, we need to include the package to
#' our experiment.

using Augmentor

#' The next step, and maybe the most human-hour consuming part of
#' adding image augmentation to a prediction problem, is to
#' design and select a sensible augmentation pipeline. Take a
#' look at the [elastic distortions tutorial](@ref elastic) for
#' an example of how to do just that.

#' For this example, we already choose a quite complicated but
#' promising augmentation pipeline for you. This pipeline was
#' designed to yield a large variation of effects as well as to
#' showcase how even deep pipelines are quite efficient in terms
#' of performance.

pl = Reshape(28,28) |>
     PermuteDims(2,1) |>
     ShearX(-5:5) * ShearY(-5:5) |>
     Rotate(-15:15) |>
     CropSize(28,28) |>
     Zoom(0.9:0.1:1.2) |>
     CacheImage() |>
     ElasticDistortion(10) |>
     PermuteDims(2,1) |>
     Reshape(28,28,1)
println(pl) #jl-only

#' Most of the used operations are quite self explanatory, but
#' there are some details about this pipeline worth pointing out
#' explicitly.

#' 1. We use the operation [`PermuteDims`](@ref) to convert the
#'    horizontal-major MNIST image to a julia-native
#'    vertical-major image. The vertical-major image is then
#'    processed and converted back to a horizontal-major array.
#'    We mainly do this here to showcase the option, but it is
#'    also to keep consistent with how the data is usually used
#'    in the literature. Alternatively, one could just work with
#'    the MNIST data in a vertical-major format all the way
#'    through without any issue.

#' 2. As counter-intuitive as it sounds, the operation
#'    [`CacheImage`](@ref) right before
#'    [`ElasticDistortion`](@ref) is actually used to improve
#'    performance. If we were to omit it, then the whole pipeline
#'    would be applied in one single pass. In this case, applying
#'    distortions on top of affine transformations lazily is in
#'    fact less efficient than using a temporary variable.

#' With the pipeline now defined, let us quickly peek at what
#' kind of effects we can achieve with it. In particular, lets
#' apply the pipeline multiple times to the first training image
#' and look at what kind of results it produces.

#jp [MNIST.convert2image(reshape(augment(train_x[:,:,:,1], pl), (28, 28))) for i in 1:2, j in 1:8]
#md [MNIST.convert2image(reshape(augment(train_x[:,:,:,1], pl), (28, 28))) for i in 1:8, j in 1:2]
#md tmp = vcat(hcat(ans[:,1]...), hcat(ans[:,2]...)) # hide
#md save("mnist_knet_aug.png",repeat(tmp, inner=(4,4))) # hide
#md nothing # hide

#md #' ![augmented samples](mnist_knet_aug.png)

#' As we can see, we can achieve a wide range of effects, from
#' more subtle to more pronounced. The important part is that all
#' examples are still clearly representative of the true label.

#' Next, we have to adapt the function `train_baseline` to make
#' use of our augmentation pipeline. To integrate Augmentor
#' efficiently, there are three necessary changes we have to
#' make.

#' 1. Preallocate a buffer with the same size and element type
#'    that each batch has.
#'
#'    ```
#'    batch_x_aug = zeros(Float32, 28, 28, 1, batchsize)
#'    ```

#' 2. Add a call to [`augmentbatch!`](@ref) in the inner loop of
#'    the batch iterator using our pipeline and buffer.
#'
#'    ```
#'    augmentbatch!(batch_x_aug, batch_x_org, pl)
#'    ```

#' 3. Replace `batch_x_org` with `batch_x_aug` in the constructor
#'    of `KnetArray`.
#'
#'    ```
#'    batch_x = KnetArray{Float32}(batch_x_aug)
#'    ```

#' Applying these changes to our `train_baseline` function
#' will give us something similar to the following function.
#' Note how all the other parts of the function remain exactly
#' the same as before.

function train_augmented(; epochs = 500, batchsize = 100, lr = .03)
    w = weights()
    log = MVHistory()
    p = Progress(epochs, desc = "Augmented: ") #jl-only
    batch_x_aug = zeros(Float32, size(train_x,1), size(train_x,2), 1, batchsize)
    for epoch in 1:epochs
        for (batch_x_cpu, batch_y) in eachbatch((train_x ,train_y), batchsize)
            augmentbatch!(CPUThreads(), batch_x_aug, batch_x_cpu, pl)
            batch_x = KnetArray{Float32}(batch_x_aug)
            g = costgrad(w, batch_x, batch_y)
            Knet.update!(w, g, lr = lr)
        end

        next!(p) #jl-only
        if (epoch % 5) == 0
            train = acc(w, train_x, train_y)
            test  = acc(w, test_x,  test_y)
            @trace log epoch train test
            msg = "epoch " * lpad(epoch,4) * ": train accuracy " * rpad(round(train, digits=3),5,"0") * ", test accuracy " * rpad(round(test,digits=3),5,"0")
            cancel(p, msg, :blue) #jl-only
#md             println(msg)
#jp             println(msg)
        end
    end
    finish!(p) #jl-only
    log
end
#md nothing # hide

#' You may have noticed in the code above that we also pass a
#' `CPUThreads()` as the first argument to [`augmentbatch!`](@ref).
#' This instructs Augmentor to process the images of the batch in
#' parallel using multi-threading. For this to work properly you
#' will need to set the environment variable `JULIA_NUM_THREADS`
#' to the number of threads you wish to use. You can check how
#' many threads are used with the function `Threads.nthreads()`

@show Threads.nthreads();
#md nothing # hide

#' Now that all pieces are in place, let us train our network
#' once more. We will use the same parameters except that now
#' instead of the original training images we will be using
#' randomly augmented images. This will cause every epoch to be
#' different.

train_augmented(epochs=1) # warm-up
augmented_log = @time train_augmented(epochs=200);
#md nothing # hide

#' As we can see, our network reaches far better results on our
#' testset than our baseline network did. However, we can also
#' see that the training took quite a bit longer than before.
#' This difference generally decreases as the complexity of the
#' utilized neural network increases. Yet another way to improve
#' performance (aside from simplifying the augmentation pipeline)
#' would be to increase the number of available threads.

#' ## Improving Performance
info("Improving Performance") #jl-only

#' One of the most effective ways to make the most out of the
#' available resources is to augment the next (couple) mini-batch
#' while the current minibatch is being processed on the GPU.
#' We can do this via julia's build in parallel computing
#' capabilities

#' First we need a worker process that will be responsible for
#' augmenting our dataset each epoch. This worker also needs
#' access to a couple of our packages

using Distributed
addprocs(1)
@everywhere using Augmentor, MLDataUtils
#md nothing # hide

#' Next, we replace the inner `eachbatch` loop with a more
#' complicated version using a `RemoteChannel` to exchange and
#' queue the augmented data.

function async_train_augmented(; epochs = 500, batchsize = 100, lr = .03)
    w = weights()
    log = MVHistory()
    p = Progress(epochs, desc = "Async Augmented: ") #jl-only
    for epoch in 1:epochs
        @sync begin
            local_ch = Channel{Tuple}(4) # prepare up to 4 minibatches in adavnce
            remote_ch = RemoteChannel(()->local_ch)
            @spawn begin
                # This block is executed on the worker process
                batch_x_aug = zeros(Float32, size(train_x,1), size(train_x,2), 1, batchsize)
                for (batch_x_cpu, batch_y) in eachbatch((train_x ,train_y), batchsize)
                    # we are still using multithreading
                    augmentbatch!(CPUThreads(), batch_x_aug, batch_x_cpu, pl)
                    put!(remote_ch, (batch_x_aug, batch_y))
                end
                close(remote_ch)
            end
            @async begin
                # This block is executed on the main process
                for (batch_x_aug, batch_y) in local_ch
                    batch_x = KnetArray{Float32}(batch_x_aug)
                    g = costgrad(w, batch_x, batch_y)
                    Knet.update!(w, g, lr = lr)
                end
            end
        end

        next!(p) #jl-only
        if (epoch % 5) == 0
            train = acc(w, train_x, train_y)
            test  = acc(w, test_x,  test_y)
            @trace log epoch train test
            msg = "epoch " * lpad(epoch,4) * ": train accuracy " * rpad(round(train,digits=3),5,"0") * ", test accuracy " * rpad(round(test,digits=3),5,"0")
            cancel(p, msg, :blue) #jl-only
#md             println(msg)
#jp             println(msg)
        end
    end
    finish!(p) #jl-only
    log
end
#md nothing # hide

#' Note that for this toy example the overhead of this approach
#' is greater than the benefit.

async_train_augmented(epochs=1) # warm-up
augmented_log = @time async_train_augmented(epochs=200);
#md nothing # hide

#' ## Visualizing the Results
info("Visualizing the Results") #jl-only

#' Before we end this tutorial, let us make use the
#' [Plots.jl](https://github.com/JuliaPlots/Plots.jl) package to
#' visualize and discuss the recorded training curves.
#' We will plot the accuracy curves of both networks side by side
#' in order to get a good feeling about their differences.

using Plots
#jp gr()
#md gr()
#md nothing # hide
unicodeplots() #jl-only

#+

#md default(bg_outside=colorant"#FFFFFF") # hide
plt = plot(
    plot(baseline_log,  title="Baseline",  ylim=(.5,1)),
    plot(augmented_log, title="Augmented", ylim=(.5,1)),
    size = (900, 400),
    xlab = "Epoch",
    ylab = "Accuracy",
    markersize = 1
)
#jp plt
#md png(plt, "mnist_knet_curves.png") # hide
#md nothing # hide

#md #' ![learning curves](mnist_knet_curves.png)

#' Note how the accuracy on the (unaltered) training set
#' increases faster for the baseline network than for the
#' augmented one. This is to be expected, since our augmented
#' network doesn't actually use the unaltered images for
#' training, and thus has not actually seen them. Given this
#' information, it is worth pointing out explicitly how the
#' accuracy on training set is still greater than on the test set
#' for the augmented network as well. This is also not a
#' surprise, given that the augmented images are likely more
#' similar to their original ones than to the test images.

#' For the baseline network, the accuracy on the test set
#' plateaus quite quickly (around 90%). For the augmented network
#' on the other hand, it the accuracy keeps increasing for quite
#' a while longer.

#' ## References
#'
#' [^MNIST1998]: LeCun, Yan, Corinna Cortes, Christopher J.C. Burges. ["The MNIST database of handwritten digits"](http://yann.lecun.com/exdb/mnist/) Website. 1998.


text:
@testset "AMP ACSF Descriptor" begin

@info("Testing AMP ACSF Descriptor for dc Si with cutoff=6.5 and Behler2011 parameters.")
using AMP, JuLIP, Test

at = bulk(:Si, cubic=true)
desc = amp_acsf(at)
#round.(desc[1], digit=7) == acsfpy
acsf_ref = [9.89642686254e+00, 6.33890586962e+00,
            4.09478636081e+00, 2.32524172255e+00,
            1.08689635509e+00, 4.18793801933e-01,
            5.49166080698e-02, 1.05677667385e-03,
            5.48207184853e+00, 1.46814261897e+01,
            2.33489680545e+00, 1.15342511467e+01,
            3.68647068294e+00, 9.83117086198e+00,
            1.59917860535e+00, 7.74387878439e+00,
            1.91527626092e+00, 5.07034523923e+00,
            8.57808585905e-01, 4.01287756422e+00,
            8.09806838507e-01, 2.12270653881e+00,
            3.78501769332e-01, 1.69140146963e+00,
            1.35537274890e-01, 1.22679889188e+00,
            2.17936760608e-03, 3.46608033973e-01,
            2.59711785058e-01, 6.72593672234e-01,
            1.27751356608e-01, 5.40633243783e-01,
            4.87646991566e-02, 4.01895952394e-01,
            5.18718040910e-04, 1.17867171366e-01,
            3.30515305924e-02, 8.43383109191e-02,
            1.72801778808e-02, 6.85669582075e-02,
            7.09381911607e-03, 5.25105488730e-02,
            5.61170353167e-05, 1.60425743126e-02,
            1.18713390808e-03, 3.00606908300e-03,
            6.40398948138e-04, 2.45933412306e-03,
            2.72582239093e-04, 1.91258708755e-03, 5.95430642794e-04]
acsf_now = vcat(desc[1,:]...)
println("AMP returns:",acsf_now)
println("Reference:",acsf_ref)
println("AMP desc size:",length(acsf_now))
println(@test acsf_now  ≈  acsf_ref)

desc = amp_acsf(at, Behler2011=false, cuttype="Tanhyper3")
acsf_now = vcat(desc[1,:]...)
println("AMP desc size:",length(acsf_now))
end


text:
# # How to Make a bar plot

# See [How to set up plots](3.0_set_up_plots.md) to get started

# ```julia
# using PowerGraphics
# const PG = PowerGraphics
# ```

# ### Make bar plots of results

# ```julia
# bar_plot(results)
# ```

# ![this one](plots-20/P__ThermalStandard_Bar.png)
# ![this one](plots-20/Example_Bar_Plot.png)

# ### Save the bar plots to a folder

# ```julia
# folder_path = joinpath(file_path, "plots_1")
# if !isdir(folder_path)
#    mkdir(folder_path)
# end
# bar_plot(results; save = folder_path)
# ```

# ![this one](plots-21/P__ThermalStandard_Bar.png)
# ![this one](plots-21/Example_saved_Bar_Plot.png)

# ### Show reserves in the bar plots

# ```julia
# bar_plot(results; reserves = true)
# ```

# ![this one](plots-22/P__ThermalStandard_Bar.png)
# ![this one](plots-22/Example_Bar_Plot_with_Reserves.png)

# ### Set different colors for the plots

# ```julia
# colors = [:pink :green :blue :magenta :black]
# bar_plot(results; seriescolor = colors)
# ```

# ![this one](plots-23/P__ThermalStandard_Bar.png)
# ![this one](plots-23/Example_Bar_Plot_with_Other_Colors.png)

# ### Set a title

# ```julia
# title = "Example of a Title"
# bar_plot(results; title = title)
# ```

# ![this one](plots-24/P__ThermalStandard_Bar.png)
# ![this one](plots-24/Example_of_a_Title.png)

# ### For saving the plot with the PlotlyJS backend, you can set a different format for saving
# ```julia
# bar_plot(results; save = path, format = "html")
# ```
# Default format for saving is png.
# Optional formats for saving include png, html, and svg.


text:
#=

Methods for offline testing

=#

function Debug(filepath::AbstractString, filetype::AbstractString; delay=0.0)

    if delay==0.0
        delay_b = false
        delay_t = 0.0
    else
        delay_b = true
        delay_t = delay
    end

    if filetype=="qq"

        a=matread(filepath)
        md=squeeze(a["data"]',2).*-1000
        d=Debug(true,"qq",filepath,md,1,floor(length(md)/SAMPLES_PER_DATA_BLOCK)*SAMPLES_PER_DATA_BLOCK,delay_b,delay_t)
    elseif filetype=="Intan"

        myheader = read_v_header(filepath)

        f=open(filepath, "r+")

        seekend(f)

        endpos = position(f)

        close(f)

        d=Debug(true,"Intan",filepath,zeros(Int16,1),10,endpos,delay_b,delay_t)
    end
    d
end

function readDataBlocks(rhd::RHD2000,s::Array{T,1}) where T<:Sorting
    fillFromOffline!(rhd)
    applySorting(rhd,s)
    true
end

function readDataBlocks(rhd::RHD2000,s::DArray{T,1,Array{T,1}}) where T<:Sorting
    fillFromOffline!(rhd)
    if rhd.cal<3
        offline_cal(s,rhd.v,rhd.buf,rhd.nums)
    else
        offline_sort(s,rhd.v,rhd.buf,rhd.nums)
    end
    true
end

function offline_cal(s,v,buf,nums)
    @sync for p in procs(s)
        @spawnat p begin
            cal!(localpart(s),v,buf,nums)
        end
    end
end

function offline_sort(s,v,buf,nums)
    @sync for p in procs(s)
        @spawnat p begin
            onlinesort!(localpart(s),v,buf,nums)
        end
    end
end

function fillFromOffline!(rhd::RHD2000)

    toff=rhd.time[end,1]

    if rhd.debug.m == "qq"

        for j=1:SAMPLES_PER_DATA_BLOCK
            for i=1:size(rhd.v,2)
                rhd.v[j,i]=round(Int16,rhd.debug.data[rhd.debug.ind])
            end
            rhd.debug.ind+=1
            rhd.time[j,1]=toff+j
        end
        if rhd.debug.ind>=rhd.debug.maxind
            rhd.debug.ind=1
        end
    elseif rhd.debug.m == "Intan"

        f=open(rhd.debug.filepath, "r+")

        seek(f,rhd.debug.ind)

        for i=1:size(rhd.v,2)
            for j=1:SAMPLES_PER_DATA_BLOCK
                rhd.v[j,i]=read(f,Int16)
            end
        end

        for j=1:SAMPLES_PER_DATA_BLOCK
            rhd.time[j,1]=toff+j
        end

        rhd.debug.ind = position(f)

        close(f)
        if rhd.debug.ind>=rhd.debug.maxind
            rhd.debug.ind=10
        end
    end
    #Filter
    for i=1:length(rhd.filts)
        for j=1:length(rhd.filts[i])
            apply_filter(rhd,rhd.filts[i][j],rhd.filts[i][j].chan)
        end
    end

    if rhd.debug.delay
        sleep(rhd.debug.delay_t)
    end
    nothing
end


text:
#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*--------------------------------------------------------------------
#* Author: Gerson Balbueno Bicca
#* $Id: DoublePipeIncr.mso								$
#*------------------------------------------------------------------
type DoublePipe_PressureDrop
	DoublePipe_PressureDrop()=begin
		N=outers.N
		Npoints=outers.Npoints
		new(
			length(Dict{Symbol,Any}(
				:Brief=>"Hydraulic Diameter of Pipe for Pressure Drop",
				:Lower=>1e-6
			)),
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of zones",
				:Default=>2
			)),
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of incremental points",
				:Default=>3
			)),
			fill(pressure(Dict{Symbol,Any}(
				:Brief=>"Incremental Local  Pressure",
				:Default=>1,
				:Lower=>1e-10,
				:Upper=>2e4,
				:DisplayUnit=>"kPa"
			)),(Npoints)),
			press_delta(Dict{Symbol,Any}(
				:Brief=>"Total Pressure Drop",
				:Default=>0.01,
				:Lower=>0,
				:DisplayUnit=>"kPa",
				:Symbol=>"\\Delta P"
			)),
			fill(press_delta(Dict{Symbol,Any}(
				:Brief=>"Incremental Pressure Drop for friction",
				:Default=>0.01,
				:Lower=>0,
				:DisplayUnit=>"kPa",
				:Symbol=>"\\Delta P_{fric}"
			)),(Npoints)),
			fill(fricfactor(Dict{Symbol,Any}(
				:Brief=>"Incremental Friction Factor",
				:Default=>0.05,
				:Lower=>1e-10,
				:Upper=>2000
			)),(N)),
			fill(positive(Dict{Symbol,Any}(
				:Brief=>"Incremental Reynolds Number",
				:Default=>100,
				:Lower=>1
			)),(N)),
			press_delta(Dict{Symbol,Any}(
				:Brief=>"Inlet Nozzle Pressure Drop",
				:Default=>0.01,
				:Lower=>0,
				:DisplayUnit=>"kPa"
			)),
			press_delta(Dict{Symbol,Any}(
				:Brief=>"Outlet Nozzle Pressure Drop",
				:Default=>0.01,
				:Lower=>0,
				:DisplayUnit=>"kPa"
			)),
			velocity(Dict{Symbol,Any}(
				:Brief=>"Inlet Nozzle Velocity",
				:Default=>1,
				:Upper=>1e5,
				:Lower=>0
			)),
			velocity(Dict{Symbol,Any}(
				:Brief=>"Outlet Nozzle Velocity",
				:Default=>1,
				:Upper=>1e5,
				:Lower=>0
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Inlet Nozzle rho-V^2",
				:Default=>1,
				:Upper=>1e6,
				:Unit=>"kg/s^2/m"
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Outlet Nozzle rho-V^2",
				:Default=>1,
				:Upper=>1e6,
				:Unit=>"kg/s^2/m"
			)),
			[:Dh,:N,:Npoints,],
			[:Plocal,:Pdrop,:Pd_fric,:fi,:Re,:Pdnozzle_in,:Pdnozzle_out,:Vnozzle_in,:Vnozzle_out,:RVsquare_in,:RVsquare_out,]
		)
	end
	Dh::length
	N::DanaInteger
	Npoints::DanaInteger
	Plocal::Array{pressure}
	Pdrop::press_delta
	Pd_fric::Array{press_delta}
	fi::Array{fricfactor}
	Re::Array{positive}
	Pdnozzle_in::press_delta
	Pdnozzle_out::press_delta
	Vnozzle_in::velocity
	Vnozzle_out::velocity
	RVsquare_in::positive
	RVsquare_out::positive
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export DoublePipe_PressureDrop
function atributes(in::DoublePipe_PressureDrop,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Brief]="to be documented"
	fields[:Info]="to be documented"
	drive!(fields,_)
	return fields
end
DoublePipe_PressureDrop(_::Dict{Symbol,Any})=begin
	newModel=DoublePipe_PressureDrop()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
Base.length(r::OneToRange) = _unsafe_length(r)

_unsafe_length(r::OneToRange) = Int(last(r) - zero(last(r)))

Base.length(::StepSRangeLen{T,Tr,Ts,R,S,L,F}) where {T,Tr,Ts,R,S,L,F} = L

Base.length(r::StepMRangeLen) = getfield(r, :len)

Base.length(::LinSRange{T,B,E,L,D}) where {T,B,E,L,D} = L

lendiv(::LinSRange{T,B,E,L,D}) where {T,B,E,L,D} = D

Base.length(r::LinMRange) = getfield(r, :len)

lendiv(r::LinMRange) = getfield(r, :lendiv)

function start_step_stop_to_length(::Type{T}, start, step, stop) where {T}
    if (start != stop) & ((step > zero(step)) != (stop > start))
        return 0
    else
        return Int(div((stop - start) + step, step))
    end
end

function start_step_stop_to_length(::Type{T}, start, step, stop) where {T<:Union{Int,UInt,Int64,UInt64,Int128,UInt128}}
    if (start != stop) & ((step > zero(step)) != (stop > start))
        return 0
    elseif step > 1
        return Int(div(unsigned(stop - start), step)) + 1
    elseif step < -1
        return Int(div(unsigned(start - stop), -step)) + 1
    elseif step > 0
        return Int(div(stop - start, step) + 1)
    else
        return Int(div(start - stop, -step) + 1)
    end
end

"""
    can_set_length(x) -> Bool

Returns `true` if type of `x` can have its length set independent of changing
its first or last position.
"""
can_set_length(::T) where {T} = can_set_length(T)
can_set_length(::Type{T}) where {T} = false
can_set_length(::Type{T}) where {T<:LinMRange} = true
can_set_length(::Type{T}) where {T<:StepMRangeLen} = true
can_set_length(::Type{T}) where {T<:StepMRange} = true
can_set_length(::Type{T}) where {T<:UnitMRange} = true
can_set_length(::Type{T}) where {T<:OneToMRange} = true

"""
    set_length!(x, len)

Change the length of `x` while maintaining it's first and last positions.
"""
set_length!(x::AbstractRange, val) = set_length!(x, Int(val))
function set_length!(r::LinMRange, len::Int)
    len >= 0 || throw(ArgumentError("set_length!($r, $len): negative length"))
    if len == 1
        r.start == r.stop || throw(ArgumentError("set_length!($r, $len): endpoints differ"))
        setfield!(r, :len, 1)
        setfield!(r, :lendiv, 1)
        return r
    end
    setfield!(r, :len, len)
    setfield!(r, :lendiv, max(len - 1, 1))
    return r
end
function set_length!(r::StepMRangeLen, len::Int)
    len >= 0 || throw(ArgumentError("length cannot be negative, got $len"))
    1 <= r.offset <= max(1,len) || throw(ArgumentError("StepMRangeLen: offset must be in [1,$len], got $offset"))
    setfield!(r, :len, len)
    return r
end
set_length!(x::OneToMRange{T}, len::T) where {T} = set_last!(x, len)
set_length!(x::UnitMRange{T}, len) where {T} = set_last!(x, T(first(x)+len-1))
function set_length!(r::StepMRange{T}, len) where {T}
    setfield!(r, :stop, convert(T, first(r) + step(r) * (len - 1)))
    return r
end

"""
    set_lendiv!(r, d)

Change the length of `x` while maintaining it's first and last positions.
"""
set_lendiv!(r::LinMRange, d) = set_lendiv!(r, Int(d))
function set_lendiv!(r::LinMRange, d::Int)
    d >= 0 || throw(ArgumentError("set_length!($r, $len): negative length"))
    #=
    we don't do this because on the off chance the user is intentionally
    setting lendiv we can't know if they want the length also set to 1 or 2
    if d == 1
        r.start == r.stop || throw(ArgumentError("set_length!($r, $len): endpoints differ"))
        setfield!(r, :len, 1)
        setfield!(r, :lendiv, 1)
        return r
    end
    =#
    setfield!(r, :len, d + 1)
    setfield!(r, :lendiv, d)
    return r
end


text:
import CanDecomp

A = rand(2, 3)
B = rand(5, 3)
C = rand(10, 3)
T_orig = CanDecomp.totensor(A, B, C)

Af = rand(size(A)...);
Bf = rand(size(B)...);
Cf = rand(size(C)...);

import StaticArrays
CanDecomp.candecomp!(StaticArrays.SVector(Af, Bf, Cf), T_orig, Val{:nnoptim}; regularization=1e-3, print_level=0, max_cd_iters=1000)

T_est = CanDecomp.totensor(Af, Bf, Cf);

import NTFk
import LinearAlgebra

@info("Norm $(LinearAlgebra.norm(T_est .- T_orig))")

import Cairo, Fontconfig

NTFk.plot2matrices(A, Af; progressbar=nothing)
NTFk.plot2matrices(B, Bf; progressbar=nothing)
NTFk.plot2matrices(C, Cf; progressbar=nothing)
NTFk.plotlefttensor(T_orig, T_est; progressbar=nothing)

max_stars_repo_name: jamesjcai/ScTenifold

max_issues_repo_name: jamesjcai/ScTenifold

max_forks_repo_name: jamesjcai/ScTenifold