text:
# This file is a part of Julia. License is MIT: http://julialang.org/license
typealias ComplexOrReal{T} Union{T,Complex{T}}

gamma(x::Float64) = nan_dom_err(ccall((:tgamma,libm),  Float64, (Float64,), x), x)
gamma(x::Float32) = nan_dom_err(ccall((:tgammaf,libm),  Float32, (Float32,), x), x)

"""
    gamma(x)

Compute the gamma function of `x`.
"""
gamma(x::Real) = gamma(float(x))

function lgamma_r(x::Float64)
    signp = Array{Int32}(1)
    y = ccall((:lgamma_r,libm),  Float64, (Float64, Ptr{Int32}), x, signp)
    return y, signp[1]
end
function lgamma_r(x::Float32)
    signp = Array{Int32}(1)
    y = ccall((:lgammaf_r,libm),  Float32, (Float32, Ptr{Int32}), x, signp)
    return y, signp[1]
end
lgamma_r(x::Real) = lgamma_r(float(x))
lgamma_r(x::Number) = lgamma(x), 1 # lgamma does not take abs for non-real x
"`lgamma_r(x)`: return L,s such that `gamma(x) = s * exp(L)`" lgamma_r

"""
    lfact(x)

Compute the logarithmic factorial of `x`
"""
lfact(x::Real) = (x<=1 ? zero(float(x)) : lgamma(x+one(x)))

"""
    lgamma(x)

Compute the logarithm of the absolute value of [`gamma`](@ref) for
`Real` `x`, while for `Complex` `x` it computes the
principal branch cut of the logarithm of `gamma(x)` (defined for negative `real(x)`
by analytic continuation from positive `real(x)`).
"""
function lgamma end

# asymptotic series for log(gamma(z)), valid for sufficiently large real(z) or |imag(z)|
@inline function lgamma_asymptotic(z::Complex{Float64})
    zinv = inv(z)
    t = zinv*zinv
    # coefficients are bernoulli[2:n+1] .// (2*(1:n).*(2*(1:n) - 1))
    return (z-0.5)*log(z) - z + 9.1893853320467274178032927e-01 + # <-- log(2pi)/2
       zinv*@evalpoly(t, 8.3333333333333333333333368e-02,-2.7777777777777777777777776e-03,
                         7.9365079365079365079365075e-04,-5.9523809523809523809523806e-04,
                         8.4175084175084175084175104e-04,-1.9175269175269175269175262e-03,
                         6.4102564102564102564102561e-03,-2.9550653594771241830065352e-02)
end

# Compute the logΓ(z) function using a combination of the asymptotic series,
# the Taylor series around z=1 and z=2, the reflection formula, and the shift formula.
# Many details of these techniques are discussed in D. E. G. Hare,
# "Computing the principal branch of log-Gamma," J. Algorithms 25, pp. 221-236 (1997),
# and similar techniques are used (in a somewhat different way) by the
# SciPy loggamma function.  The key identities are also described
# at http://functions.wolfram.com/GammaBetaErf/LogGamma/
function lgamma(z::Complex{Float64})
    x = real(z)
    y = imag(z)
    yabs = abs(y)
    if !isfinite(x) || !isfinite(y) # Inf or NaN
        if isinf(x) && isfinite(y)
            return Complex(x, x > 0 ? (y == 0 ? y : copysign(Inf, y)) : copysign(Inf, -y))
        elseif isfinite(x) && isinf(y)
            return Complex(-Inf, y)
        else
            return Complex(NaN, NaN)
        end
    elseif x > 7 || yabs > 7
        # use the Stirling asymptotic series for sufficiently large x or |y|
        return lgamma_asymptotic(z)
    elseif x < 0.1 # use reflection formula to transform to x > 0
        if x == 0 && y == 0 # return Inf with the correct imaginary part for z == 0
            return Complex(Inf, signbit(x) ? copysign(oftype(x, pi), -y) : -y)
        end
        # the 2pi * floor(...) stuff is to choose the correct branch cut for log(sinpi(z))
        return Complex(1.1447298858494001741434262, # log(pi)
                       copysign(6.2831853071795864769252842, y) # 2pi
                       * floor(0.5*x+0.25)) -
               log(sinpi(z)) - lgamma(1-z)
    elseif abs(x - 1) + yabs < 0.1
        # taylor series around zero at z=1
        # ... coefficients are [-eulergamma; [(-1)^k * zeta(k)/k for k in 2:15]]
        w = Complex(x - 1, y)
        return w * @evalpoly(w, -5.7721566490153286060651188e-01,8.2246703342411321823620794e-01,
                                -4.0068563438653142846657956e-01,2.705808084277845478790009e-01,
                                -2.0738555102867398526627303e-01,1.6955717699740818995241986e-01,
                                -1.4404989676884611811997107e-01,1.2550966952474304242233559e-01,
                                -1.1133426586956469049087244e-01,1.000994575127818085337147e-01,
                                -9.0954017145829042232609344e-02,8.3353840546109004024886499e-02,
                                -7.6932516411352191472827157e-02,7.1432946295361336059232779e-02,
                                -6.6668705882420468032903454e-02)
    elseif abs(x - 2) + yabs < 0.1
        # taylor series around zero at z=2
        # ... coefficients are [1-eulergamma; [(-1)^k * (zeta(k)-1)/k for k in 2:12]]
        w = Complex(x - 2, y)
        return w * @evalpoly(w, 4.2278433509846713939348812e-01,3.2246703342411321823620794e-01,
                               -6.7352301053198095133246196e-02,2.0580808427784547879000897e-02,
                               -7.3855510286739852662729527e-03,2.8905103307415232857531201e-03,
                               -1.1927539117032609771139825e-03,5.0966952474304242233558822e-04,
                               -2.2315475845357937976132853e-04,9.9457512781808533714662972e-05,
                               -4.4926236738133141700224489e-05,2.0507212775670691553131246e-05)
    end
    # use recurrence relation lgamma(z) = lgamma(z+1) - log(z)
    # to shift to x > 7 for asymptotic series
    shiftprod = Complex(x,yabs)
    x += 1
    sb = false # == signbit(imag(shiftprod)) == signbit(yabs)
    # To use log(product of shifts) rather than sum(logs of shifts),
    # we need to keep track of the number of + to - sign flips in
    # imag(shiftprod), as described in Hare (1997), proposition 2.2.
    signflips = 0
    while x <= 7
        shiftprod *= Complex(x,yabs)
        sb′ = signbit(imag(shiftprod))
        signflips += sb′ & (sb′ != sb)
        sb = sb′
        x += 1
    end
    shift = log(shiftprod)
    if signbit(y) # if y is negative, conjugate the shift
        shift = Complex(real(shift), signflips*-6.2831853071795864769252842 - imag(shift))
    else
        shift = Complex(real(shift), imag(shift) + signflips*6.2831853071795864769252842)
    end
    return lgamma_asymptotic(Complex(x,y)) - shift
end
lgamma{T<:Union{Integer,Rational}}(z::Complex{T}) = lgamma(float(z))
lgamma{T<:Union{Float32,Float16}}(z::Complex{T}) = Complex{T}(lgamma(Complex{Float64}(z)))

gamma(z::Complex) = exp(lgamma(z))

# Bernoulli numbers B_{2k}, using tabulated numerators and denominators from
# the online encyclopedia of integer sequences.  (They actually have data
# up to k=249, but we stop here at k=20.)  Used for generating the polygamma
# (Stirling series) coefficients below.
#   const A000367 = map(BigInt, split("1,1,-1,1,-1,5,-691,7,-3617,43867,-174611,854513,-236364091,8553103,-23749461029,8615841276005,-7709321041217,2577687858367,-26315271553053477373,2929993913841559,-261082718496449122051", ","))
#   const A002445 = [1,6,30,42,30,66,2730,6,510,798,330,138,2730,6,870,14322,510,6,1919190,6,13530]
#   const bernoulli = A000367 .// A002445 # even-index Bernoulli numbers

"""
    digamma(x)

Compute the digamma function of `x` (the logarithmic derivative of [`gamma(x)`](@ref)).
"""
function digamma(z::ComplexOrReal{Float64})
    # Based on eq. (12), without looking at the accompanying source
    # code, of: K. S. Kölbig, "Programs for computing the logarithm of
    # the gamma function, and the digamma function, for complex
    # argument," Computer Phys. Commun.  vol. 4, pp. 221–226 (1972).
    x = real(z)
    if x <= 0 # reflection formula
        ψ = -π * cot(π*z)
        z = 1 - z
        x = real(z)
    else
        ψ = zero(z)
    end
    if x < 7
        # shift using recurrence formula
        n = 7 - floor(Int,x)
        for ν = 1:n-1
            ψ -= inv(z + ν)
        end
        ψ -= inv(z)
        z += n
    end
    t = inv(z)
    ψ += log(z) - 0.5*t
    t *= t # 1/z^2
    # the coefficients here are Float64(bernoulli[2:9] .// (2*(1:8)))
    ψ -= t * @evalpoly(t,0.08333333333333333,-0.008333333333333333,0.003968253968253968,-0.004166666666666667,0.007575757575757576,-0.021092796092796094,0.08333333333333333,-0.4432598039215686)
end

"""
    trigamma(x)

Compute the trigamma function of `x` (the logarithmic second derivative of [`gamma(x)`](@ref)).
"""
function trigamma(z::ComplexOrReal{Float64})
    # via the derivative of the Kölbig digamma formulation
    x = real(z)
    if x <= 0 # reflection formula
        return (π * csc(π*z))^2 - trigamma(1 - z)
    end
    ψ = zero(z)
    if x < 8
        # shift using recurrence formula
        n = 8 - floor(Int,x)
        ψ += inv(z)^2
        for ν = 1:n-1
            ψ += inv(z + ν)^2
        end
        z += n
    end
    t = inv(z)
    w = t * t # 1/z^2
    ψ += t + 0.5*w
    # the coefficients here are Float64(bernoulli[2:9])
    ψ += t*w * @evalpoly(w,0.16666666666666666,-0.03333333333333333,0.023809523809523808,-0.03333333333333333,0.07575757575757576,-0.2531135531135531,1.1666666666666667,-7.092156862745098)
end

signflip(m::Number, z) = (-1+0im)^m * z
signflip(m::Integer, z) = iseven(m) ? z : -z

# (-1)^m d^m/dz^m cot(z) = p_m(cot z), where p_m is a polynomial
# that satisfies the recurrence p_{m+1}(x) = p_m′(x) * (1 + x^2).
# Note that p_m(x) has only even powers of x if m is odd, and
# only odd powers of x if m is even.  Therefore, we write p_m(x)
# as p_m(x) = q_m(x^2) m! for m odd and p_m(x) = x q_m(x^2) m! if m is even.
# Hence the polynomials q_m(y) satisfy the recurrence:
#     m odd:  q_{m+1}(y) = 2 q_m′(y) * (1+y) / (m+1)
#    m even:  q_{m+1}(y) = [q_m(y) + 2 y q_m′(y)] * (1+y) / (m+1)
# This function computes the coefficients of the polynomial q_m(y),
# returning an array of the coefficients of 1, y, y^2, ...,
function cotderiv_q(m::Int)
    m < 0 && throw(ArgumentError("$m < 0 not allowed"))
    m == 0 && return [1.0]
    m == 1 && return [1.0, 1.0]
    q₋ = cotderiv_q(m-1)
    d = length(q₋) - 1 # degree of q₋
    if isodd(m-1)
        q = Array{Float64}(length(q₋))
        q[end] = d * q₋[end] * 2/m
        for i = 1:length(q)-1
            q[i] = ((i-1)*q₋[i] + i*q₋[i+1]) * 2/m
        end
    else # iseven(m-1)
        q = Array{Float64}(length(q₋) + 1)
        q[1] = q₋[1] / m
        q[end] = (1 + 2d) * q₋[end] / m
        for i = 2:length(q)-1
            q[i] = ((1 + 2(i-1))*q₋[i] + (1 + 2(i-2))*q₋[i-1]) / m
        end
    end
    return q
end

# precompute a table of cot derivative polynomials
const cotderiv_Q = [cotderiv_q(m) for m in 1:100]

# Evaluate (-1)^m d^m/dz^m [π cot(πz)] / m!.  If m is small, we
# use the explicit derivative formula (a polynomial in cot(πz));
# if m is large, we use the derivative of Euler's harmonic series:
#          π cot(πz) = ∑ inv(z + n)
function cotderiv(m::Integer, z)
    isinf(imag(z)) && return zero(z)
    if m <= 0
        m == 0 && return π * cot(π*z)
        throw(DomainError())
    end
    if m <= length(cotderiv_Q)
        q = cotderiv_Q[m]
        x = cot(π*z)
        y = x*x
        s = q[1] + q[2] * y
        t = y
        # evaluate q(y) using Horner (TODO: Knuth for complex z?)
        for i = 3:length(q)
            t *= y
            s += q[i] * t
        end
        return π^(m+1) * (isodd(m) ? s : x*s)
    else # m is large, series derivative should converge quickly
        p = m+1
        z -= round(real(z))
        s = inv(z^p)
        n = 1
        sₒ = zero(s)
        while s != sₒ
            sₒ = s
            a = (z+n)^p
            b = (z-n)^p
            s += (a + b) / (a * b)
            n += 1
        end
        return s
    end
end

# Helper macro for polygamma(m, z):
#   Evaluate p[1]*c[1] + x*p[2]*c[2] + x^2*p[3]*c[3] + ...
#   where c[1] = m + 1
#         c[k] = c[k-1] * (2k+m-1)*(2k+m-2) / ((2k-1)*(2k-2)) = c[k-1] * d[k]
#         i.e. d[k] = c[k]/c[k-1] = (2k+m-1)*(2k+m-2) / ((2k-1)*(2k-2))
#   by a modified version of Horner's rule:
#      c[1] * (p[1] + d[2]*x * (p[2] + d[3]*x * (p[3] + ...))).
# The entries of p must be literal constants and there must be > 1 of them.
macro pg_horner(x, m, p...)
    k = length(p)
    me = esc(m)
    xe = esc(x)
    ex = :(($me + $(2k-1)) * ($me + $(2k-2)) * $(p[end]/((2k-1)*(2k-2))))
    for k = length(p)-1:-1:2
        cdiv = 1 / ((2k-1)*(2k-2))
        ex = :(($cdiv * ($me + $(2k-1)) * ($me + $(2k-2))) *
               ($(p[k]) + $xe * $ex))
    end
    :(($me + 1) * ($(p[1]) + $xe * $ex))
end

# compute oftype(x, y)^p efficiently, choosing the correct branch cut
pow_oftype(x, y, p) = oftype(x, y)^p
pow_oftype(x::Complex, y::Real, p::Complex) = oftype(x, y^p)
function pow_oftype(x::Complex, y::Real, p::Real)
    if p >= 0
        # note: this will never be called for y < 0,
        # which would throw an error for non-integer p here
        return oftype(x, y^p)
    else
        yp = y^-p # use real power for efficiency
        return oftype(x, Complex(yp, -zero(yp))) # get correct sign of zero!
    end
end

# Generalized zeta function, which is related to polygamma
# (at least for integer m > 0 and real(z) > 0) by:
#    polygamma(m, z) = (-1)^(m+1) * gamma(m+1) * zeta(m+1, z).
# Our algorithm for the polygamma is just the m-th derivative
# of our digamma approximation, and this derivative process yields
# a function of the form
#          (-1)^(m) * gamma(m+1) * (something)
# So identifying the (something) with the -zeta function, we get
# the zeta function for free and might as well export it, especially
# since this is a common generalization of the Riemann zeta function
# (which Julia already exports).   Note that this geneneralization
# is equivalent to Mathematica's Zeta[s,z], and is equivalent to the
# Hurwitz zeta function for real(z) > 0.

"""
    zeta(s, z)

Generalized zeta function ``\\zeta(s, z)``, defined
by the sum ``\\sum_{k=0}^\\infty ((k+z)^2)^{-s/2}``, where
any term with ``k+z=0`` is excluded.  For ``\\Re z > 0``,
this definition is equivalent to the Hurwitz zeta function
``\\sum_{k=0}^\\infty (k+z)^{-s}``.   For ``z=1``, it yields
the Riemann zeta function ``\\zeta(s)``.
"""
function zeta(s::ComplexOrReal{Float64}, z::ComplexOrReal{Float64})
    ζ = zero(promote_type(typeof(s), typeof(z)))

    (z == 1 || z == 0) && return oftype(ζ, zeta(s))
    s == 2 && return oftype(ζ, trigamma(z))

    x = real(z)

    # annoying s = Inf or NaN cases:
    if !isfinite(s)
        (isnan(s) || isnan(z)) && return (s*z)^2 # type-stable NaN+Nan*im
        if real(s) == Inf
            z == 1 && return one(ζ)
            if x > 1 || (x >= 0.5 ? abs(z) > 1 : abs(z - round(x)) > 1)
                return zero(ζ) # distance to poles is > 1
            end
            x > 0 && imag(z) == 0 && imag(s) == 0 && return oftype(ζ, Inf)
        end
        throw(DomainError()) # nothing clever to return
    end
    if isnan(x)
        if imag(z)==0 && imag(s)==0
            return oftype(ζ, x)
        else
            return oftype(ζ, Complex(x,x))
        end
    end

    m = s - 1

    # Algorithm is just the m-th derivative of digamma formula above,
    # with a modified cutoff of the final asymptotic expansion.

    # Note: we multiply by -(-1)^m m! in polygamma below, so this factor is
    #       pulled out of all of our derivatives.

    cutoff = 7 + real(m) + abs(imag(m)) # TODO: this cutoff is too conservative?
    if x < cutoff
        # shift using recurrence formula
        xf = floor(x)
        nx = Int(xf)
        n = ceil(Int,cutoff - nx)
        minus_s = -s
        if nx < 0 # x < 0
            # need to use (-z)^(-s) recurrence to be correct for real z < 0
            # [the general form of the recurrence term is (z^2)^(-s/2)]
            minus_z = -z
            ζ += pow_oftype(ζ, minus_z, minus_s) # ν = 0 term
            if xf != z
                ζ += pow_oftype(ζ, z - nx, minus_s)
                # real(z - nx) > 0, so use correct branch cut
                # otherwise, if xf==z, then the definition skips this term
            end
            # do loop in different order, depending on the sign of s,
            # so that we are looping from largest to smallest summands and
            # can halt the loop early if possible; see issue #15946
            # FIXME: still slow for small m, large Im(z)
            if real(s) > 0
                for ν in -nx-1:-1:1
                    ζₒ= ζ
                    ζ += pow_oftype(ζ, minus_z - ν, minus_s)
                    ζ == ζₒ && break # prevent long loop for large -x > 0
                end
            else
                for ν in 1:-nx-1
                    ζₒ= ζ
                    ζ += pow_oftype(ζ, minus_z - ν, minus_s)
                    ζ == ζₒ && break # prevent long loop for large -x > 0
                end
            end
        else # x ≥ 0 && z != 0
            ζ += pow_oftype(ζ, z, minus_s)
        end
        # loop order depends on sign of s, as above
        if real(s) > 0
            for ν in max(1,1-nx):n-1
                ζₒ= ζ
                ζ += pow_oftype(ζ, z + ν, minus_s)
                ζ == ζₒ && break # prevent long loop for large m
            end
        else
            for ν in n-1:-1:max(1,1-nx)
                ζₒ= ζ
                ζ += pow_oftype(ζ, z + ν, minus_s)
                ζ == ζₒ && break # prevent long loop for large m
            end
        end
        z += n
    end

    t = inv(z)
    w = isa(t, Real) ? conj(oftype(ζ, t))^m : oftype(ζ, t)^m
    ζ += w * (inv(m) + 0.5*t)

    t *= t # 1/z^2
    ζ += w*t * @pg_horner(t,m,0.08333333333333333,-0.008333333333333333,0.003968253968253968,-0.004166666666666667,0.007575757575757576,-0.021092796092796094,0.08333333333333333,-0.4432598039215686,3.0539543302701198)

    return ζ
end

"""
    polygamma(m, x)

Compute the polygamma function of order `m` of argument `x` (the `(m+1)`th derivative of the
logarithm of [`gamma(x)`](@ref))
"""
function polygamma(m::Integer, z::ComplexOrReal{Float64})
    m == 0 && return digamma(z)
    m == 1 && return trigamma(z)

    # In principle, we could support non-integer m here, but the
    # extension to complex m seems to be non-unique, the obvious
    # extension (just plugging in a complex m below) does not seem to
    # be the one used by Mathematica or Maple, and sources do not
    # agree on what the "right" extension is (e.g. Mathematica & Maple
    # disagree).   So, at least for now, we require integer m.

    # real(m) < 0 is valid, but I don't think our asymptotic expansion
    # here works in this case.  m < 0 polygamma is called a
    # "negapolygamma" function in the literature, and there are
    # multiple possible definitions arising from different integration
    # constants. We throw a DomainError() since the definition is unclear.
    real(m) < 0 && throw(DomainError())

    s = Float64(m+1)
    # It is safe to convert any integer (including `BigInt`) to a float here
    # as underflow occurs before precision issues.
    if real(z) <= 0 # reflection formula
        (zeta(s, 1-z) + signflip(m, cotderiv(m,z))) * (-gamma(s))
    else
        signflip(m, zeta(s,z) * (-gamma(s)))
    end
end

"""
    invdigamma(x)

Compute the inverse [`digamma`](:func:`digamma`) function of `x`.
"""
function invdigamma(y::Float64)
    # Implementation of fixed point algorithm described in
    # "Estimating a Dirichlet distribution" by Thomas P. Minka, 2000

    # Closed form initial estimates
    if y >= -2.22
        x_old = exp(y) + 0.5
        x_new = x_old
    else
        x_old = -1.0 / (y - digamma(1.0))
        x_new = x_old
    end

    # Fixed point algorithm
    delta = Inf
    iteration = 0
    while delta > 1e-12 && iteration < 25
        iteration += 1
        x_new = x_old - (digamma(x_old) - y) / trigamma(x_old)
        delta = abs(x_new - x_old)
        x_old = x_new
    end

    return x_new
end

"""
    beta(x, y)

Euler integral of the first kind
``\\operatorname{B}(x,y) = \\Gamma(x)\\Gamma(y)/\\Gamma(x+y)``.
"""
function beta(x::Number, w::Number)
    yx, sx = lgamma_r(x)
    yw, sw = lgamma_r(w)
    yxw, sxw = lgamma_r(x+w)
    return exp(yx + yw - yxw) * (sx*sw*sxw)
end

"""
    lbeta(x, y)

Natural logarithm of the absolute value of the [`beta`](@ref)
function ``\\log(|\\operatorname{B}(x,y)|)``.
"""
lbeta(x::Number, w::Number) = lgamma(x)+lgamma(w)-lgamma(x+w)

"""
    zeta(s)

Riemann zeta function ``\\zeta(s)``.
"""
function zeta(s::ComplexOrReal{Float64})
    # Riemann zeta function; algorithm is based on specializing the Hurwitz
    # zeta function above for z==1.

    # blows up to ±Inf, but get correct sign of imaginary zero
    s == 1 && return NaN + zero(s) * imag(s)

    if !isfinite(s) # annoying NaN and Inf cases
        isnan(s) && return imag(s) == 0 ? s : s*s
        if isfinite(imag(s))
            real(s) > 0 && return 1.0 - zero(s)*imag(s)
            imag(s) == 0 && return NaN + zero(s)
        end
        return NaN*zero(s) # NaN + NaN*im
    elseif real(s) < 0.5
        if abs(real(s)) + abs(imag(s)) < 1e-3 # Taylor series for small |s|
            return @evalpoly(s, -0.5,
                             -0.918938533204672741780329736405617639861,
                             -1.0031782279542924256050500133649802190,
                             -1.00078519447704240796017680222772921424,
                             -0.9998792995005711649578008136558752359121)
        end
        return zeta(1 - s) * gamma(1 - s) * sinpi(s*0.5) * (2π)^s / π
    end

    m = s - 1

    # shift using recurrence formula:
    #   n is a semi-empirical cutoff for the Stirling series, based
    #   on the error term ~ (|m|/n)^18 / n^real(m)
    n = ceil(Int,6 + 0.7*abs(imag(s-1))^inv(1 + real(m)*0.05))
    ζ = one(s)
    for ν = 2:n
        ζₒ= ζ
        ζ += inv(ν)^s
        ζ == ζₒ && break # prevent long loop for large m
    end
    z = 1 + n
    t = inv(z)
    w = t^m
    ζ += w * (inv(m) + 0.5*t)

    t *= t # 1/z^2
    ζ += w*t * @pg_horner(t,m,0.08333333333333333,-0.008333333333333333,0.003968253968253968,-0.004166666666666667,0.007575757575757576,-0.021092796092796094,0.08333333333333333,-0.4432598039215686,3.0539543302701198)

    return ζ
end

"""
    eta(x)

Dirichlet eta function ``\\eta(s) = \\sum^\\infty_{n=1}(-1)^{n-1}/n^{s}``.
"""
function eta(z::ComplexOrReal{Float64})
    δz = 1 - z
    if abs(real(δz)) + abs(imag(δz)) < 7e-3 # Taylor expand around z==1
        return 0.6931471805599453094172321214581765 *
               @evalpoly(δz,
                         1.0,
                         -0.23064207462156020589789602935331414700440,
                         -0.047156357547388879740146103148112380421254,
                         -0.002263576552598880778433550956278702759143568,
                         0.001081837223249910136105931217561387128141157)
    else
        return -zeta(z) * expm1(0.6931471805599453094172321214581765*δz)
    end
end

# Converting types that we can convert, and not ones we can not
# Float16, and Float32 and their Complex equivalents can be converted to Float64
# and results converted back.
# Otherwise, we need to make things use their own `float` converting methods
# and in those cases, we do not convert back either as we assume
# they also implement their own versions of the functions, with the correct return types.
# This is the case for BitIntegers (which become `Float64` when `float`ed).
# Otherwise, if they do not implement their version of the functions we
# manually throw a `MethodError`.
# This case occurs, when calling `float` on a type does not change its type,
# as it is already a `float`, and would have hit own method, if one had existed.

# If we really cared about single precision, we could make a faster
# Float32 version by truncating the Stirling series at a smaller cutoff.
# and if we really cared about half precision, we could make a faster
# Float16 version, by using a precomputed table look-up.

for T in (Float16, Float32, Float64)
    @eval f64(x::Complex{$T}) = Complex128(x)
    @eval f64(x::$T) = Float64(x)
end

for f in (:digamma, :trigamma, :zeta, :eta, :invdigamma)
    @eval begin
        function $f(z::Union{ComplexOrReal{Float16}, ComplexOrReal{Float32}})
            oftype(z, $f(f64(z)))
        end

        function $f(z::Number)
            x = float(z)
            typeof(x) === typeof(z) && throw(MethodError($f, (z,)))
            # There is nothing to fallback to, as this didn't change the argument types
            $f(x)
        end
    end
end

for T1 in (Float16, Float32, Float64), T2 in (Float16, Float32, Float64)
    (T1 == T2 == Float64) && continue # Avoid redefining base definition

    @eval function zeta(s::ComplexOrReal{$T1}, z::ComplexOrReal{$T2})
        ζ = zeta(f64(s), f64(z))
        convert(promote_type(typeof(s), typeof(z)),  ζ)
    end
end

function zeta(s::Number, z::Number)
    t = float(s)
    x = float(z)
    if typeof(t) === typeof(s) && typeof(x) === typeof(z)
        # There is nothing to fallback to, since this didn't work
        throw(MethodError(zeta,(s,z)))
    end
    zeta(t, x)
end

function polygamma(m::Integer, z::Union{ComplexOrReal{Float16}, ComplexOrReal{Float32}})
    oftype(z, polygamma(m, f64(z)))
end

function polygamma(m::Integer, z::Number)
    x = float(z)
    typeof(x) === typeof(z) && throw(MethodError(polygamma, (m,z)))
    # There is nothing to fallback to, since this didn't work
    polygamma(m, x)
end

max_stars_repo_name: oxinabox/julia

max_issues_repo_name: oxinabox/julia

max_forks_repo_name: oxinabox/julia


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

if !isdefined(@__MODULE__, Symbol("@verify_error"))
    macro verify_error(arg)
        arg isa String && return esc(:(print && println(stderr, $arg)))
        (arg isa Expr && arg.head === :string) || error("verify_error macro expected a string expression")
        pushfirst!(arg.args, GlobalRef(Core, :stderr))
        pushfirst!(arg.args, :println)
        arg.head = :call
        return esc(arg)
    end
end

function check_op(ir::IRCode, domtree::DomTree, @nospecialize(op), use_bb::Int, use_idx::Int, print::Bool)
    if isa(op, SSAValue)
        if op.id > length(ir.stmts)
            def_bb = block_for_inst(ir.cfg, ir.new_nodes.info[op.id - length(ir.stmts)].pos)
        else
            def_bb = block_for_inst(ir.cfg, op.id)
        end
        if (def_bb == use_bb)
            if op.id > length(ir.stmts)
                @assert ir.new_nodes.info[op.id - length(ir.stmts)].pos <= use_idx
            else
                if op.id >= use_idx
                    @verify_error "Def ($(op.id)) does not dominate use ($(use_idx)) in same BB"
                    error("")
                end
            end
        else
            if !dominates(domtree, def_bb, use_bb) && !(bb_unreachable(domtree, def_bb) && bb_unreachable(domtree, use_bb))
                # At the moment, we allow GC preserve tokens outside the standard domination notion
                #@Base.show ir
                @verify_error "Basic Block $def_bb does not dominate block $use_bb (tried to use value $(op.id))"
                error("")
            end
        end
    elseif isa(op, GlobalRef)
        if !isdefined(op.mod, op.name) || !isconst(op.mod, op.name)
            @verify_error "Unbound GlobalRef not allowed in value position"
            error("")
        end
    elseif isa(op, Union{OldSSAValue, NewSSAValue})
        #@Base.show ir
        @verify_error "Left over SSA marker"
        error("")
    elseif isa(op, Union{SlotNumber, TypedSlot})
        @verify_error "Left over slot detected in converted IR"
        error("")
    end
end

function count_int(val::Int, arr::Vector{Int})
    n = 0
    for x in arr
        if x === val
            n += 1
        end
    end
    n
end

function verify_ir(ir::IRCode, print::Bool=true)
    # For now require compact IR
    # @assert isempty(ir.new_nodes)
    # Verify CFG
    last_end = 0
    # Verify statements
    domtree = construct_domtree(ir.cfg.blocks)
    for (idx, block) in pairs(ir.cfg.blocks)
        if first(block.stmts) != last_end + 1
            #ranges = [(idx,first(bb.stmts),last(bb.stmts)) for (idx, bb) in pairs(ir.cfg.blocks)]
            @verify_error "First statement of BB $idx ($(first(block.stmts))) does not match end of previous ($last_end)"
            error("")
        end
        last_end = last(block.stmts)
        terminator = ir.stmts[last_end][:inst]

        bb_unreachable(domtree, idx) && continue
        for p in block.preds
            p == 0 && continue
            c = count_int(idx, ir.cfg.blocks[p].succs)
            if c == 0
                @verify_error "Predecessor $p of block $idx not in successor list"
                error("")
            elseif c == 2
                if count_int(p, block.preds) != 2
                    @verify_error "Double edge from $p to $idx not correctly accounted"
                    error("")
                end
            end
        end
        if isa(terminator, ReturnNode)
            if !isempty(block.succs)
                @verify_error "Block $idx ends in return or unreachable, but has successors"
                error("")
            end
        elseif isa(terminator, GotoNode)
            if length(block.succs) != 1 || block.succs[1] != terminator.label
                @verify_error "Block $idx successors ($(block.succs)), does not match GotoNode terminator"
                error("")
            end
        elseif isa(terminator, GotoIfNot)
            if terminator.dest == idx + 1
                @verify_error "Block $idx terminator forms a double edge to block $(idx+1)"
                error("")
            end
            if length(block.succs) != 2 || (block.succs != [terminator.dest, idx+1] && block.succs != [idx+1, terminator.dest])
                @verify_error "Block $idx successors ($(block.succs)), does not match GotoIfNot terminator"
                error("")
            end
        elseif isexpr(terminator, :enter)
            @label enter_check
            if length(block.succs) != 2 || (block.succs != Int[terminator.args[1], idx+1] && block.succs != Int[idx+1, terminator.args[1]])
                @verify_error "Block $idx successors ($(block.succs)), does not match :enter terminator"
                error("")
            end
        else
            if length(block.succs) != 1 || block.succs[1] != idx + 1
                # As a special case, we allow extra statements in the BB of an :enter
                # statement, until we can do proper CFG manipulations during compaction.
                for idx in first(block.stmts):last(block.stmts)
                    stmt = ir.stmts[idx][:inst]
                    if isexpr(stmt, :enter)
                        terminator = stmt
                        @goto enter_check
                    end
                    isa(stmt, PhiNode) || break
                end
                @verify_error "Block $idx successors ($(block.succs)), does not match fall-through terminator ($terminator)"
                error("")
            end
        end
        for s in block.succs
            if !(idx in ir.cfg.blocks[s].preds)
                #@Base.show ir.cfg
                #@Base.show ir
                #@Base.show ir.argtypes
                @verify_error "Successor $s of block $idx not in predecessor list"
                error("")
            end
        end
    end
    for (bb, idx) in bbidxiter(ir)
        # We allow invalid IR in dead code to avoid passes having to detect when
        # they're generating dead code.
        bb_unreachable(domtree, bb) && continue
        stmt = ir.stmts[idx][:inst]
        stmt === nothing && continue
        if isa(stmt, PhiNode)
            @assert length(stmt.edges) == length(stmt.values)
            for i = 1:length(stmt.edges)
                edge = stmt.edges[i]
                if !(edge == 0 && bb == 1) && !(edge in ir.cfg.blocks[bb].preds)
                    #@Base.show ir.argtypes
                    #@Base.show ir
                    @verify_error "Edge $edge of φ node $idx not in predecessor list"
                    error("")
                end
                edge == 0 && continue
                isassigned(stmt.values, i) || continue
                val = stmt.values[i]
                phiT = ir.stmts[idx][:type]
                if isa(val, SSAValue)
                    if !(types(ir)[val] ⊑ phiT)
                        #@verify_error """
                        #    PhiNode $idx, has operand $(val.id), whose type is not a sub lattice element.
                        #    PhiNode type was $phiT
                        #    Value type was $(ir.stmts[val.id][:type])
                        #"""
                        #error("")
                    end
                elseif isa(val, GlobalRef) || isa(val, Expr)
                    @verify_error "GlobalRefs and Exprs are not allowed as PhiNode values"
                    error("")
                end
                check_op(ir, domtree, val, Int(edge), last(ir.cfg.blocks[stmt.edges[i]].stmts)+1, print)
            end
        elseif isa(stmt, PhiCNode)
            for i = 1:length(stmt.values)
                val = stmt.values[i]
                if !isa(val, SSAValue)
                    @verify_error "Operand $i of PhiC node $idx must be an SSA Value."
                    error("")
                end
                if !isa(ir[val], UpsilonNode)
                    @verify_error "Operand $i of PhiC node $idx must reference an Upsilon node."
                    error("")
                end
            end
        else
            if isa(stmt, Expr) || isa(stmt, ReturnNode) # TODO: make sure everything has line info
                if !(stmt isa ReturnNode && !isdefined(stmt, :val)) # not actually a return node, but an unreachable marker
                    if ir.stmts[idx][:line] <= 0
                        #@verify_error "Missing line number information for statement $idx of $ir"
                    end
                end
            end
            if isa(stmt, Expr)
                if stmt.head === :(=)
                    if stmt.args[1] isa SSAValue
                        @verify_error "SSAValue as assignment LHS"
                        error("")
                    end
                    if stmt.args[2] isa GlobalRef
                        # undefined GlobalRef as assignment RHS is OK
                        continue
                    end
                elseif stmt.head === :gc_preserve_end
                    # We allow gc_preserve_end tokens to span across try/catch
                    # blocks, which isn't allowed for regular SSA values, so
                    # we skip the validation below.
                    continue
                elseif stmt.head === :isdefined && length(stmt.args) == 1 && stmt.args[1] isa GlobalRef
                    # a GlobalRef isdefined check does not evaluate its argument
                    continue
                end
            end
            for op in userefs(stmt)
                op = op[]
                check_op(ir, domtree, op, bb, idx, print)
            end
        end
    end
end

function verify_linetable(linetable::Vector{LineInfoNode}, print::Bool=true)
    for i in 1:length(linetable)
        line = linetable[i]
        if i <= line.inlined_at
            @verify_error "Misordered linetable"
            error("")
        end
    end
end

ext: jl

max_stars_repo_name: jd-foster/julia

max_issues_repo_name: jd-foster/julia

max_forks_repo_name: jd-foster/julia


text:
using StatsBase
using RDatasets
iris_dataframe = dataset("datasets", "iris")
head(iris_dataframe)
typeof(iris_dataframe[1,:Species])
describe(iris_dataframe)

max_stars_repo_path: Chapter 3/3.jl

max_stars_repo_name: PacktPublishing/Julia-for-data-science

max_issues_repo_path: Chapter 3/3.jl

max_issues_repo_name: PacktPublishing/Julia-for-data-science

max_forks_repo_path: Chapter 3/3.jl

max_forks_repo_name: PacktPublishing/Julia-for-data-science


text:
make_ygrid(y, ngrid=100) = range(minimum(y), maximum(y), length=ngrid)

_ecdf(xs::AbstractArray{<:Real}, ys::AbstractArray{<:Real}) = ecdf(xs).(ys)

function postmean_cdf(chain::AbstractVector, ygrid; exponentiate=false,
                      N=10000)
  pp = posterior_predictive(chain)
  B = length(pp.C)
  if exponentiate
    cdf_C = mean(map(p -> _ecdf(exp.(rand(p, N)), ygrid), pp.C))
    cdf_T = mean(map(p -> _ecdf(exp.(rand(p, N)), ygrid), pp.T))
  else
    cdf_C = mean(map(p -> _ecdf(rand(p, N), ygrid), pp.C))
    cdf_T = mean(map(p -> _ecdf(rand(p, N), ygrid), pp.T))
  end
  return (C=cdf_C, T=cdf_T)
end

function posterior_predictive(chain::AbstractVector)
  function pp(c::NamedTuple, i::Symbol)
    comps = SkewT.(c.mu, c.sigma, c.nu, c.phi)
    MixtureModel(comps, grab(c, :eta, i))
  end
  return (C=pp.(chain, :C), T=pp.(chain, :T))
end

function compute_post_density(chain::AbstractVector, ygrid)
  function compute_pdf(c::NamedTuple, i::Symbol)
    comps = SkewT.(c.mu, c.sigma, c.nu, c.phi)
    return pdf.(MixtureModel(comps, grab(c, :eta, i)), ygrid)
  end
  return (C=compute_pdf.(chain, :C), T=compute_pdf.(chain, :T))
end

function printsummary(chain, metrics; digits=3)
  if :loglike in keys(metrics)
    println("mean loglike: ", mean(metrics[:loglike]))
    println("DIC: ", MCMC.dic(metrics[:loglike]))
  end

  sanitize(s) = round(mean(s), digits=digits)
  sanitizevec(v) = round.(mean(v), digits=digits)

  function print_uni_param(sym)
    param = getindex.(chain, sym)
    println("mean $(sym): ", sanitize(param))
  end

  function print_mv_param(sym)
    param = getindex.(chain, sym)
    println("mean $(sym): ", sanitizevec(param))
  end

  foreach(print_mv_param, [:etaC, :etaT, :mu, :sigma, :nu, :phi])
  foreach(print_uni_param, [:tau])
end

function plot_post_density!(chain, ygrid; plot_mean=true, alpha=0.3)
  pdfC, pdfT = compute_post_density(chain, ygrid)
  pdfC = hcat(pdfC...)
  pdfT = hcat(pdfT...)

  pdfC_mean = mean(pdfC, dims=2)
  pdfC_lower = vec(quantiles(pdfC, 0.025, dims=2))
  pdfC_upper = vec(quantiles(pdfC, 0.975, dims=2))

  pdfT_mean = mean(pdfT, dims=2)
  pdfT_lower = vec(quantiles(pdfT, 0.025, dims=2))
  pdfT_upper = vec(quantiles(pdfT, 0.975, dims=2))

  p = let
    plot!(ygrid, pdfC_lower, fillrange=pdfC_upper, alpha=alpha, color=:blue,
         label=nothing)
    plot_mean && plot!(ygrid, pdfC_mean, color=:blue, label=nothing)
    plot!(ygrid, pdfT_lower, fillrange=pdfT_upper, alpha=alpha, color=:red,
          label=nothing)
    plot_mean && plot!(ygrid, pdfT_mean, color=:red, label=nothing)
    ylabel!("density")
  end

  return p
end

function plot_simtruth(dC, dT, ygrid; kwargs...)
  plot!(ygrid, pdf.(dC, ygrid), color=:blue, ls=:dot; kwargs...)
  plot!(ygrid, pdf.(dT, ygrid), color=:red, ls=:dot; kwargs...)
  ylabel!("density")
end

# Add methods for computing and plotting CDF of F_i
function compute_Fi_tilde_cdf(gtilde::Gtilde, chain::AbstractVector,
                              pzero::Pzero; gridsize::Int=100,
                              exponentiate=false, N::Integer=10000,
                              max_quantile=0.999)
  gammaC_mean = mean(infer_Pzero1(pzero, 10000).distC)
  gammaT_mean = mean(infer_Pzero1(pzero, 10000).distT)

  y = [gtilde.yC; gtilde.yT]
  if exponentiate
    qmax = exp(quantile(y, max_quantile))
    ygrid = range(0, qmax, length=gridsize)
  else
    ygrid = make_ygrid(y, gridsize)
  end
  cdfs = postmean_cdf(chain, ygrid, exponentiate=exponentiate, N=N)

  cdfC = cdfs.C * (1 - gammaC_mean) .+ gammaC_mean
  cdfT = cdfs.T * (1 - gammaT_mean) .+ gammaT_mean

  # Compute area between curves.
  area = sum(abs.(cdfC - cdfT) * (ygrid[2] - ygrid[1]))

  if exponentiate
    return (C=cdfC, T=cdfT, ygrid=ygrid, area=area, qmax=qmax)
  else
    return (C=cdfC, T=cdfT, ygrid=ygrid, area=area)
  end
end
function plot_Fi_tilde_cdf!(gtilde::Gtilde, chain::AbstractVector,
                            pzero::Pzero; gridsize::Int=200,
                            exponentiate=false, N=10000, max_quantile=0.999)
  cdfs = compute_Fi_tilde_cdf(gtilde, chain, pzero, gridsize=gridsize,
                              exponentiate=exponentiate, N=N,
                              max_quantile=max_quantile)
  plot!(cdfs.ygrid, cdfs.C, lw=2, color=:blue, label=nothing)
  plot!(cdfs.ygrid, cdfs.T, lw=2, color=:red, label=nothing)

  println("Area between curves: $(cdfs.area) (exponentiate=$(exponentiate))")
  if exponentiate
    println("Normalized area between Fi CDFs: $(cdfs.area / cdfs.qmax)")
  end
end

function compute_Fi_tilde_cdf_truth(simdata::NamedTuple) end
function plot_Fi_tilde_cdf_truth!(simdata::NamedTuple) end

"""
Return the average number of small clusters in the two samples.
"""
function count_small_clusters(chain::AbstractVector; p::Real=0.01)
  etaC = getindex.(chain, :etaC)
  etaT = getindex.(chain, :etaT)
  nsmall_clust = map(eta -> sum(eta .< p), etaC + etaT)
  return mean(nsmall_clust)
end

function hellinger(chain::AbstractVector, pzero::Pzero, nsamps::Integer;
                   verbose::Int=2)
  B = length(chain)
  gamma_post = infer_Pzero1(pzero, 1)

  function hell(c)
    gammaC = rand(gamma_post.distC)
    gammaT = rand(gamma_post.distT)
    st = SkewT.(c.mu, c.sigma, c.nu, c.phi)
    GC = MixtureModel(st, c.etaC)
    GT = MixtureModel(st, c.etaT)
    FC = HurdleModel(-Inf, GC, gammaC)
    FT = HurdleModel(-Inf, GT, gammaT)
    H2 = MCMC.hellinger2(FC, FT, nsamps)
    if !(0 < H2 < 1)
      verbose >= 2 && println("WARNING: H2 = $(H2). Clamping to be in [0, 1].")
      H2 = clamp(H2, 0, 1)
    end
    verbose >= 1 && flush(stdout)
    return sqrt(H2)
  end

  iters = verbose >= 1 ? MCMC.ProgressBars.ProgressBar(chain) : chain
  return map(hell, iters)
end


text:
## Exports
module Expect
export ExpectProc, interact, expect!, with_timeout!, raw!, sendeof
export ExpectTimeout, ExpectEOF

## Imports
import Base.Libc: strerror
import Base: Process, TTY, wait, wait_readnb, wait_readbyte
import Base: kill, process_running, process_exited, success
import Base: write, print, println, flush, eof, close
import Base: read, readbytes!, readuntil
import Base: isopen, bytesavailable, readavailable

## UNIX/tty support lib
@static if Sys.isunix()
    const LIBEXJL = joinpath(@__DIR__, "../deps/libexjl.so")

    function _set_cloexec(fd::RawFD)
        ccall((:exjl_set_cloexec, LIBEXJL), Cint, (Cint,), fd)
    end

    function _sendeof(tty::TTY)
        flush(tty) # flush any pending buffer in jl/uv
        ccall((:exjl_sendeof, LIBEXJL), Cint, (Ptr{Nothing},), tty)
    end
end

## Types
struct ExpectTimeout <: Exception end
struct ExpectEOF <: Exception end

mutable struct ExpectProc <: IO
    proc::Process
    timeout::Real
    encode::Function
    decode::Function
    in_stream::IO
    out_stream::IO
    before
    match
    buffer::Vector{UInt8}

    function ExpectProc(cmd::Cmd, timeout::Real; env::Base.EnvDict=ENV, encoding="utf8", pty=true)
        # TODO: only utf8 is currently supported
        @assert encoding == "utf8"
        encode = x->collect(transcode(UInt8, x))
        decode = x->transcode(String, copy(x))

        in_stream, out_stream, proc = _spawn(cmd, env, pty)
        new(proc, timeout, encode, decode,
            in_stream, out_stream,
            nothing, nothing, [])
    end
end

## Support functions
function raw!(tty::TTY, raw::Bool)
    # UV_TTY_MODE_IO (cfmakeraw) is only available with libuv 1.0 and not
    # directly supported by jl_tty_set_mode (JL_TTY_MODE_RAW still performs NL
    # conversion).
    UV_TTY_MODE_NORMAL = 0
    UV_TTY_MODE_IO = 2
    mode = raw ? UV_TTY_MODE_IO : UV_TTY_MODE_NORMAL
    ret = ccall(:uv_tty_set_mode, Cint, (Ptr{Nothing},Cint), tty.handle, mode)
    ret == 0
end

function raw!(proc::ExpectProc, raw::Bool)
    if !isa(proc.out_stream, TTY)
        # pipes are always raw
        return raw
    else
        raw!(proc.out_stream, raw)
    end
end

function _spawn(cmd::Cmd, env::Base.EnvDict, pty::Bool)
    env = copy(ENV)
    env["TERM"] = "dumb"
    setenv(cmd, env)

    if pty && Sys.isunix()
        O_RDWR = Base.Filesystem.JL_O_RDWR
        O_NOCTTY = Base.Filesystem.JL_O_NOCTTY

        fdm = RawFD(ccall(:posix_openpt, Cint, (Cint,), O_RDWR|O_NOCTTY))
        fdm == RawFD(-1) && error("openpt failed: $(strerror())")
        ttym = TTY(fdm)
        in_stream = out_stream = ttym

        rc = _set_cloexec(fdm)
        rc != 0 && error("ioctl failed: $(strerror())")

        rc = ccall(:grantpt, Cint, (Cint,), fdm)
        rc != 0 && error("grantpt failed: $(strerror())")

        rc = ccall(:unlockpt, Cint, (Cint,), fdm)
        rc != 0 && error("unlockpt failed: $(strerror())")

        pts = ccall(:ptsname, Ptr{UInt8}, (Cint,), fdm)
        pts == C_NULL && error("ptsname failed: $(strerror())")

        fds = RawFD(ccall(:open, Cint, (Ptr{UInt8}, Cint), pts, O_RDWR|O_NOCTTY))
        fds == RawFD(-1) && error("open failed: $(strerror())")
        if raw!(out_stream, true) == false
            ccall(:close, Cint, (Cint,), fds)
            error("raw! failed: $(strerror())")
        end

        local proc::Process
        try
            proc = run(cmd, (fds, fds, fds), wait=false)
        catch ex
            ccall(:close, Cint, (Cint,), fds)
            close(ttym)
            rethrow(ex)
        end

        @async begin
            wait(proc)
            close(ttym)
            ccall(:close, Cint, (Cint,), fds)
        end

        Base.start_reading(in_stream)
    else
        pipe = Pipe()
        proc = open(cmd, "r", pipe)
        out_stream = Base.pipe_writer(pipe)
        in_stream = Base.pipe_reader(proc)
    end

    return (in_stream, out_stream, proc)
end

# Process management
kill(proc::ExpectProc, signum::Integer=15) = kill(proc.proc, signum)
wait(proc::ExpectProc) = wait(proc.proc)
process_running(proc::ExpectProc) = process_running(proc.proc)
process_exited(proc::ExpectProc) = process_exited(proc.proc)
success(proc::ExpectProc) = success(proc.proc)

# Writing functions
flush(proc::ExpectProc) = flush(proc.out_stream)
close(proc::ExpectProc) = close(proc.out_stream)

function sendeof(proc::ExpectProc)
    if typeof(proc.out_stream) <: TTY
        return _sendeof(proc.out_stream) == 0
    else
        close(proc.out_stream)
        return true
    end
end

write(proc::ExpectProc, buf::Vector{UInt8}) = write(proc.out_stream, buf)
write(proc::ExpectProc, buf::AbstractString) = write(proc, proc.encode(buf))
print(proc::ExpectProc, x::AbstractString) = write(proc, x)
println(proc::ExpectProc, x::AbstractString) = write(proc, string(x, "\n"))
# resolve method ambiguity
const ContiguousString = Union{String,SubString{String}}
write(proc::ExpectProc, buf::ContiguousString) = write(proc, proc.encode(buf))
print(proc::ExpectProc, x::ContiguousString) = write(proc, x)
println(proc::ExpectProc, x::ContiguousString) = write(proc, string(x, "\n"))

# Reading functions
function _timed_wait(func::Function, proc::ExpectProc; timeout::Real=proc.timeout)
    if isinf(timeout)
        return func()
    end
    thunk = current_task()
    timer = Timer(timeout) do _
        Base.throwto(thunk, ExpectTimeout())
    end
    local ret
    try
        ret = func()
    finally
        close(timer)
    end
    return ret
end

function eof(proc::ExpectProc; timeout::Real=proc.timeout)
    _timed_wait(proc; timeout=timeout) do
        eof(proc.in_stream)
    end
end

function wait_readnb(proc::ExpectProc, nb::Int; timeout::Real=proc.timeout)
    _timed_wait(proc; timeout=timeout) do
        wait_readnb(proc.in_stream, nb)
    end
end

function wait_readbyte(proc::ExpectProc, c::UInt8; timeout::Real=proc.timeout)
    _timed_wait(proc; timeout=timeout) do
        wait_readbyte(proc.in_stream, c)
    end
end

read(proc::ExpectProc, ::Type{UInt8}; timeout::Real=proc.timeout) =
    _timed_wait(proc; timeout=timeout) do
        read(proc.in_stream, UInt8)
    end

readbytes!(proc::ExpectProc, b::AbstractVector{UInt8}, nb=length(b); timeout::Real=proc.timeout) =
    _timed_wait(proc; timeout=timeout) do
        readbytes!(proc.in_stream, b, nb)
    end

readuntil(proc::ExpectProc, delim::AbstractString; timeout::Real=proc.timeout, keep::Bool=false) =
    _timed_wait(proc; timeout=timeout) do
        readuntil(proc.in_stream, delim, keep=keep)
    end

isopen(proc::ExpectProc) = isopen(proc.in_stream)
bytesavailable(proc::ExpectProc) = bytesavailable(proc.in_stream)
readavailable(proc::ExpectProc) = readavailable(proc.in_stream)

# Expect
function _expect_search(buf::AbstractString, str::AbstractString)
    pos = findfirst(str, buf)
    return pos == nothing ? nothing : (buf[pos], pos)
end

function _expect_search(buf::AbstractString, regex::Regex)
    m = match(regex, buf)
    return m == nothing ? nothing : (m.match, m.offset:(m.offset+length(m.match)-1))
end

function _expect_search(buf::AbstractString, vec::Vector)
    for idx=1:length(vec)
        ret = _expect_search(buf, vec[idx])
        if ret != nothing
            return idx, ret[1], ret[2]
        end
    end
    return nothing
end

function expect!(proc::ExpectProc, vec; timeout::Real=proc.timeout)
    proc.match = nothing
    pos = 0:-1
    idx = 0
    while true
        if bytesavailable(proc.in_stream) > 0
            proc.buffer = vcat(proc.buffer, readavailable(proc.in_stream))
        end
        if length(proc.buffer) > 0
            buffer = try proc.decode(proc.buffer); finally; end
            if buffer != nothing
                ret = _expect_search(buffer, vec)
                if ret != nothing
                    idx, proc.match, pos = ret
                    break
                end
            end
        end
        if !isopen(proc.in_stream)
            throw(ExpectEOF())
        end
        wait_readnb(proc, 1; timeout=timeout)
    end
    proc.before = proc.decode(proc.buffer[1:pos[1]-1])
    proc.buffer = proc.buffer[pos[end]+1:end]
    return idx
end

function expect!(proc::ExpectProc, regex::Regex; timeout::Real=proc.timeout)
    expect!(proc, [regex]; timeout=timeout)
    proc.before
end

function expect!(proc::ExpectProc, str::AbstractString; timeout::Real=proc.timeout)
    # TODO: this is worth implementing efficiently
    expect!(proc, [str]; timeout=timeout)
    proc.before
end

# Helpers
function with_timeout!(func::Function, proc::ExpectProc, timeout::Real)
    orig = proc.timeout
    proc.timeout = timeout
    local ret
    try
        ret = func()
    finally
        proc.timeout = orig
    end
    return ret
end

function interact(func::Function, cmd::Cmd, args...; kwargs...)
    proc = ExpectProc(cmd, args...; kwargs...)
    try
        func(proc)
    catch err
        kill(proc)
        rethrow(err)
    finally
        close(proc)
    end

    # wait+kill performed by success itself
    success(proc)
end

end


text:
abstract type LateralDriverModel{A} <: DriverModel{A} end

# TODO: Why does this struct live in the `lateral_driver_models.jl` file
"""
	ProportionalLaneTracker

A controller that executes the lane change decision made by the `lane change models`

# Constructors
	ProportionalLaneTracker(;σ::Float64 = NaN,kp::Float64 = 3.0,kd::Float64 = 2.0)

# Fields
- `a::Float64 = NaN` predicted acceleration
- `σ::Float64 = NaN` optional stdev on top of the model, set to zero or NaN for deterministic behavior
- `kp::Float64 = 3.0` proportional constant for lane tracking
- `kd::Float64 = 2.0` derivative constant for lane tracking
"""
@with_kw mutable struct ProportionalLaneTracker <: LateralDriverModel{Float64}
    a::Float64 = NaN # predicted acceleration
    σ::Float64 = NaN # optional stdev on top of the model, set to zero or NaN for deterministic behavior
    kp::Float64 = 3.0 # proportional constant for lane tracking
    kd::Float64 = 2.0 # derivative constant for lane tracking
end

function track_lateral!(model::ProportionalLaneTracker, laneoffset::Float64, lateral_speed::Float64)
    model.a = -laneoffset*model.kp - lateral_speed*model.kd
    model
end

function observe!(model::ProportionalLaneTracker, scene::Scene{Entity{S, D, I}}, roadway::Roadway, egoid::I) where {S, D, I}

    ego_index = findfirst(egoid, scene)
    veh_ego = scene[ego_index]
    t = posf(veh_ego.state).t # lane offset
    dt = velf(veh_ego.state).t # rate of change of lane offset
    model.a = -t*model.kp - dt*model.kd

    model
end
function Base.rand(rng::AbstractRNG, model::ProportionalLaneTracker)
    if isnan(model.σ) || model.σ ≤ 0.0
        model.a
    else
        rand(rng, Normal(model.a, model.σ))
    end
end
function Distributions.pdf(model::ProportionalLaneTracker, a_lat::Float64)
    if isnan(model.σ) || model.σ ≤ 0.0
        Inf
    else
        pdf(Normal(model.a, model.σ), a_lat)
    end
end
function Distributions.logpdf(model::ProportionalLaneTracker, a_lat::Float64)
    if isnan(model.σ) || model.σ ≤ 0.0
        Inf
    else
        logpdf(Normal(model.a, model.σ), a_lat)
    end
end


text:
include("filtering.jl");
## filter per treatment
list = ["PreVehicle",
    "Altanserin",
    "SB242084",
    "Way_100135",
    "Methysergide",
    "Citalopram"]
s = filter(r->r.Treatment in list &&
    r.Trial < 61,
    streaks)

## number of pokes before leaving
Df = combine(groupby(s,[:Phase,:Treatment])) do dd
    ecdf(dd,:Num_pokes; mode = :conf)
end
Drug_colors!(Df)
gd = groupby(Df,:Phase)
tp = [@df subdf plot(:Xaxis, :Mean,
    xlabel = :Phase[1],
    xticks = 0:5:30,
    group = :Treatment,
    linecolor = :color,
    legend = false,
    ribbon = :ERR,
    linewidth = 2,
    fillalpha = 0.3) for subdf in gd]
ord = [5,3,4,2,1]
pp =  plot(tp[ord]...)
note = plot(xlims = (0,0.5), ylims = (0,0.5), annotations = (0.25,0.25,
    Plots.text("mean plus 95% C.I.", :center)),
    border = :none)
plot(vcat(tp[ord],note)...)
savefig(joinpath(figs_loc,"Fig3/CumNumPokes.pdf"))

## Signed rank test Num Pokes
df1 = combine(groupby(s,:Phase)) do dd
    subdf = unstack(dd,:Treatment,:Num_pokes)
    current_drug = Symbol(subdf[1,:Phase])
    rename!(subdf, current_drug => :Drug)
end
df2 = combine(groupby(df1,[:Phase])) do dd
    wilcoxon(dd,:Drug, :PreVehicle; f = x -> mean(skipmissing(x)))
end
Drug_colors!(df2)
plot_wilcoxon(df2)
savefig(joinpath(figs_loc,"Fig3/WilcoxonNumPokes.pdf"))
## Traveltime analysis
df1 = combine(groupby(s,:Phase)) do dd
    subdf = unstack(dd,:Treatment,:Travel_to)
    current_drug = Symbol(subdf[1,:Phase])
    rename!(subdf, current_drug => :Drug)
end
df2 = combine(groupby(df1,[:Phase])) do dd
    wilcoxon(dd,:Drug, :PreVehicle; f = x -> mean(skipmissing(x)))
end
Drug_colors!(df2)
plot_wilcoxon(df2)
savefig(joinpath(figs_loc,"Fig3/WilcoxonTravel.pdf"))
##


text:
module APTDecoder

if Sys.islinux()
    # must load first
    # https://github.com/JuliaIO/ImageMagick.jl/issues/142#issuecomment-433691895
    using ImageMagick
end
using Images
using Printf
using Dates
using Statistics
using LinearAlgebra
using SatelliteToolbox
using PyPlot
using CodecZlib
using RemoteFiles
import DSP
using FileIO
using GeoDatasets

scans_per_seconds = 2

# https://web.archive.org/web/20190814072342/https://noaa-apt.mbernardi.com.ar/how-it-works.html
# frequency ratio is 5/4
sync_frequency = (1040., # channel A
                  832.)   # channel B

include("GeoMapping.jl")
include("georef.jl")
include("decode.jl")
include("plot.jl")

end # module

max_stars_repo_name: Alexander-Barth/APTDecoder

max_issues_repo_name: Alexander-Barth/APTDecoder

max_forks_repo_name: Alexander-Barth/APTDecoder


text:
export label_components!

"""
label_components!(lattice::Lattice)

labeling lattice site
"""
function label_components!(sqlattice::Square)
    if sqlattice.lattice_config.neighbortype == "nn"
        sqlattice.lattice_properties.labeled_lattice_sites = label_components(sqlattice.lattice_config.lattice_sites)
    else # next nearest neighbor
        sqlattice.lattice_properties.labeled_lattice_sites = label_components(sqlattice.lattice_config.lattice_sites, trues(3,3))
    end

    sqlattice.lattice_properties.islabeled = true
    sqlattice.lattice_properties.nclusters = maximum(sqlattice.lattice_properties.labeled_lattice_sites)

    return nothing
end

function label_components!(trilattice::Triangular)
    site = trilattice.lattice_config.lattice_sites
    row, col = size(site)
    labelnum = 0
    labeled_site = - Int.(site)
    searchlist = Vector{Vector{Int}}()

    function checkneighbor!(i,j)
        if labeled_site[i,j] == -1
            # x o o neighbor
            # o o o
            # o o x
            if j < col && labeled_site[i, j+1] == -1; push!(searchlist, [i, j+1]); end
            if 1 < j && labeled_site[i, j-1] == -1; push!(searchlist, [i, j-1]); end
            if i < row && labeled_site[i+1, j] == -1; push!(searchlist, [i+1, j]); end
            if 1 < i && labeled_site[i-1, j] == -1; push!(searchlist, [i-1, j]); end
            if 1 < i && j < col && labeled_site[i-1, j+1] == -1; push!(searchlist, [i-1, j+1]); end
            if i < row && 1 < j && labeled_site[i+1, j-1] == -1; push!(searchlist, [i+1, j-1]); end
        end
    end

    @simd for j in 1:col
        @simd for i in 1:row
            if labeled_site[i,j] == -1
                labelnum += 1
                checkneighbor!(i,j)
                labeled_site[i,j] = labelnum

                while !isempty(searchlist)
                    tmppos = pop!(searchlist)
                    checkneighbor!(tmppos...)
                    labeled_site[tmppos...] = labelnum
                end
            end
        end
    end

    trilattice.lattice_properties.nclusters = labelnum
    trilattice.lattice_properties.islabeled = true
    trilattice.lattice_properties.labeled_lattice_sites = labeled_site

    return nothing
end

function label_components!(hclattice::Honeycomb)
    site = hclattice.lattice_config.lattice_sites
    row, col = size(site)
    labelnum = 0
    labeled_site = - Int.(site)
    searchlist = Vector{Vector{Int}}()

    function checkneighbor!(i,j)
        if labeled_site[i,j] == -1
            if j < col && labeled_site[i, j+1] == -1; push!(searchlist, [i, j+1]); end # right
            if 1 < j && labeled_site[i, j-1] == -1; push!(searchlist, [i, j-1]); end # left

            if iseven(i+j)
                if 1 < i && labeled_site[i-1, j] == -1; push!(searchlist, [i-1, j]); end # above
            else
                if i < row && labeled_site[i+1, j] == -1; push!(searchlist, [i+1, j]); end # down
            end
        end
    end

    @simd for j in 1:col
        @simd for i in 1:row
            if labeled_site[i,j] == -1
                labelnum += 1
                checkneighbor!(i,j)
                labeled_site[i,j] = labelnum

                while !isempty(searchlist)
                    tmppos = pop!(searchlist)
                    checkneighbor!(tmppos...)
                    labeled_site[tmppos...] = labelnum
                end
            end
        end
    end

    hclattice.lattice_properties.nclusters = labelnum
    hclattice.lattice_properties.islabeled = true
    hclattice.lattice_properties.labeled_lattice_sites = labeled_site

    return nothing
end


text:
module grid
using LinearAlgebra
export Spacegrid,grid_n,grid_k,grid_xy_square
struct Spacegrid
    x::Array{Float64,2}
    y::Array{Float64,2}
end

function grid_n(Nx,Ny)
    #reciprocal space coordinate indices
    ny=[c  for r in -Nx:Nx, c in -Ny:Ny]
    nx=[r  for r in -Nx:Nx, c in -Ny:Ny]
    nx=vec(nx)
    ny=vec(ny)
    #difference matrix for 2dft
    dnx=[nx[a]-nx[b] for a in 1:length(nx),b in 1:length(nx)]
    dny=[ny[a]-ny[b] for a in 1:length(ny),b in 1:length(ny)]
    return nx,ny,dnx,dny
end

function grid_k(nx,ny,theta,alpha,lambda,ax,ay,epsilon_ref)
    #straightforward
    #all k vectors are generally normalized to k0 here
    k0=2*pi/real(lambda)
    #The incoming wave, transformed from spherical coordinates to normalized cartesian coordinates, |kin|=1
    kin=[sin(theta*pi/180)*cos(alpha*pi/180),sin(theta*pi/180)*sin(alpha*pi/180),cos(theta*pi/180)]*real(sqrt(epsilon_ref))
    #the spatial vectors are the sum of the incoming vector and the reciprocal lattice vectors
    kx=kin[1].+(2*pi/ax)*nx/k0;
    ky=kin[2].+(2*pi/ay)*ny/k0;
    #need matrix for later computations
    Kx=Diagonal(kx)
    Ky=Diagonal(ky)
    return k0,Kx,Ky,kin
end

function grid_xy_square(acc)
    #creating a square meshgrid
    x=[r  for r in -acc/2+.5:acc/2-.5, c in -acc/2+.5:acc/2-.5]/acc
    y=[c  for r in -acc/2+.5:acc/2-.5, c in -acc/2+.5:acc/2-.5]/acc
    return Spacegrid(x,y)
end

end


text:
using PyPlot, AcousticWave
nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
v  = 2500. * ones(nz, nx);
v[101:end, 1:end] = 3000.;
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# initialize source term
f0 = 30.0; ot = [0.0]; pos_src = [5 101]
src = InitSources(pos_src, f0, ot, dt);
# receiver location
irx = collect(1:2:nx); irz = 5*ones(Int64, length(irx));
pos_rec = convert(Array{Int64,2}, hcat(irz, irx))
# position of receiver
shot  = MultiStepForward(pos_rec, src[1], fidMtx, tmax=1.0)
SeisPlot(shot.d, pclip=95, cmap="seismic", name="/Users/wenlei/Desktop/shot.pdf")
dp = dpDirect(shot, v[1], src[1], dx, dt)
shot.d = dp .* shot.d
SeisPlot(shot.d, cmap="seismic", name="/Users/wenlei/Desktop/shot_rd.pdf")

# compute source side wavefield with smooth velocity model
v1 = modSmooth(v, 10)
fidMtx1 = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v1);
pathsrc = "/Users/wenlei/Desktop/src.bin"
MultiStepForward(pathsrc, src, fidMtx1, tmax=1.0)
pathout  = "/Users/wenlei/Desktop/srcWave"
waveAnim(pathsrc, pathout, interval=100)

# compute receiver side wavefield with smooth velocity model
pathrec = "/Users/wenlei/Desktop/rec.bin"
MultiStepAdjoint(pathrec, shot, fidMtx1)
animRec = "/Users/wenlei/Desktop/recWave"
waveAnim(pathrec, animRec, interval=100)
(Ivz, Ivx, Ipz, Ipx) = RTMimaging(pathsrc, pathrec, image_type="all");
SeisPlot(Ipz, cmap="seismic")
Ip = RTMimaging(pathsrc, pathrec, image_type="p");
SeisPlot(Ip, cmap="seismic")

# ==============================================================================
using PyPlot, AcousticWave
nz = 101; nx = 101; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
v  = 2500. * ones(nz, nx);
v[101:end, 1:end] = 3000.;
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
m31 = fidMtx.MpzBvz * fidMtx.MvzBvz;
m33 = fidMtx.MpzBvz * fidMtx.MvzBp + fidMtx.MpzBpz;
m34 = fidMtx.MpzBvz * fidMtx.MvzBp;

m42 = fidMtx.MpxBvx * fidMtx.MvxBvx;
m43 = fidMtx.MpxBvx * fidMtx.MvxBp;
m44 = fidMtx.MpxBvx * fidMtx.MvxBp + fidMtx.MpxBpx;

N = (nz+2*ext)*(nx+2*ext);
M = vcat(hcat(fidMtx.MvzBvz, spzeros(N,N) , fidMtx.MvzBp, fidMtx.MvzBp),
         hcat(spzeros(N,N) , fidMtx.MvxBvx, fidMtx.MvxBp, fidMtx.MvxBp),
         hcat(m31          , spzeros(N,N) , m33         , m34         ),
         hcat(spzeros(N,N) , m42          , m43         , m44         ))

# ==============================================================================
# using PyPlot, AcousticWave
# nz = 101; nx = 101; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
# v  = 3000. * ones(nz, nx);
# v[101:end, 1:end] = 3000.;
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# fidMtx1= InitFidMtx_back(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
#
# function OneStepBackward(spt2::SnapShot, spt1::SnapShot, fidMtx::FidMtx)
#     spt2.it = spt1.it - 1
#     spt2.pz = fidMtx1.MpzBpz * spt1.pz + fidMtx1.MpzBvz * spt1.vz;
#     spt2.px = fidMtx1.MpxBpx * spt1.px + fidMtx1.MpxBvx * spt1.vx;
#     spt2.vz = fidMtx1.MvzBvz * spt1.vz + fidMtx1.MvzBp * (spt2.pz+spt2.px);
#     spt2.vx = fidMtx1.MvxBvx * spt1.vx + fidMtx1.MvxBp * (spt2.pz+spt2.px);
#     return nothing
# end
#
# spt1 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, 2);
# spt2 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, 2);
# spt3 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, 2);
# spt1.vz = rand(length(spt1.vz)); spt1.vx = rand(length(spt1.vx));
# spt1.pz = rand(length(spt1.pz)); spt1.px = rand(length(spt1.px));
# OneStepBackward(spt2, spt1, fidMtx);
# OneStepAdjoint!(spt3, spt1, fidMtx);
#
# p2 = spt2.pz + spt2.px
# p3 = spt3.pz + spt3.px
#
# tmp = spt3-spt2

max_stars_repo_name: WneleiGao/AcousticWaveLab

max_issues_repo_name: WneleiGao/AcousticWaveLab

max_forks_repo_name: WneleiGao/AcousticWaveLab


text:
const Data{M, T} = NamedTuple{(:poly, :sigtail, :sigratio),
                              Tuple{Polynomial{M, T}, Polynomial{M, T}, M}}
const SigTable{I, M, T} = Dict{Tuple{I, M}, Data{M, T}}
const GenData{I} = NamedTuple{(:position, :att_key, :tag),
                              Tuple{I, I, Symbol}}
gendata(i::I, j::I, tg::Symbol) where I = (position = i, att_key = j, tag = tg)

mutable struct SigPolynomialΓ{I, M, T, MΓ<:Context{M}, TΓ<:Context{T}, PΓ<:PolynomialΓ{M, T, MΓ, TΓ}, S}<:Context{Tuple{I, M}}
    po::PΓ
    tbl::SigTable{I, M, T}
    ord_indices::Dict{I, GenData{I}}
end

pos_type(::SigPolynomialΓ{I}) where {I} = I
mon_type(::SigPolynomialΓ{I, M}) where {I, M} = M
coeff_type(::SigPolynomialΓ{I, M, T}) where {I, M, T} = T
mod_order(::SigPolynomialΓ{I, M, T, MΓ, TΓ, PΓ, S}) where {I, M, T, MΓ, TΓ, PΓ, S} = S

gendata(ctx::SigPolynomialΓ{I}) where I = (position = zero(I), att_key = zero(I), tag = :f)

function pos(ctx::SigPolynomialΓ{I, M},
             p::Tuple{I, M}) where {I, M}

    iszero(p[1]) && return zero(I)
    ctx.ord_indices[p[1]][:position]
end

function gettag(ctx::SigPolynomialΓ{I, M},
                p::Tuple{I, M}) where {I, M}

    iszero(p[1]) && return :nil
    ctx.ord_indices[p[1]][:tag]
end

getattkey(ctx::SigPolynomialΓ{I, M}, p::Tuple{I, M}) where {I, M} = ctx.ord_indices[p[1]][:att_key]

function new_gen!(ctx::SigPolynomialΓ{I, M, T},
                  info_hashmap::Dict{I, Info},
                  posit::I,
                  attached_to::I,
                  tagg::Symbol,
                  pol::Polynomial{M, T}) where {I, M, T}

    # register new generator
    posit_key = register!(ctx, pol, info_hashmap)

    # rebuild ord_indices
    new_pos!(ctx, posit_key, posit, attached_to, tagg)
end

function register!(ctx::SigPolynomialΓ{I, M, T},
                   pol::Polynomial{M, T},
                   info_hashmap::Dict{I, Dict{String, Int}}) where {I, M, T}

    posit_key = maximum(keys(ctx.ord_indices)) + one(I)
    ctx((posit_key, one(ctx.po.mo)), pol)
    info_hashmap[posit_key] = new_info()
    ctx.ord_indices[posit_key] = gendata(ctx)
    return posit_key
end

function new_pos!(ctx::SigPolynomialΓ{I},
                  posit_key::I,
                  posit::I,
                  attached_to::I,
                  tagg::Symbol) where I

    delete!(ctx.ord_indices, posit_key)
    new_dict_arr = [(k, i >= posit ? gendata(i + one(I), j, tg) : gendata(i, j, tg))
                    for (k, (i, j, tg)) in ctx.ord_indices]
    push!(new_dict_arr, (posit_key, gendata(posit, attached_to, tagg)))
    ctx.ord_indices = Dict(new_dict_arr)
end

function new_pos!(ctx::SigPolynomialΓ{I}, posit_key::I, posit::I) where I
    new_pos!(ctx, posit_key, posit, ctx.ord_indices[posit_key][:att_key],
             ctx.ord_indices[posit_key][:tag])
end

# find maximal index
function maxpos(ctx::SigPolynomialΓ{I, M}) where {I, M}

    maximum(v -> v[:position], values(ctx.ord_indices))
end

# return original generator of higher index than pos if it exists
function f_left(ctx::SigPolynomialΓ{I, M}, pos::I) where {I, M}

    result = zero(I)
    for (i, v) in ctx.ord_indices
        if v[:tag] == :f && v[:position] > pos
            if iszero(result) || v[:position] < ctx.ord_indices[result][:position]
                result = i
            end
        end
    end
    return result
end

function idxsigpolynomialctx(coefficients,
                             ngens;
                             monomials=nothing,
                             index_type=UInt32,
                             mask_type=UInt32,
                             pos_type=UInt32,
                             mod_order=:POT,
                             deg_bound = 0,
                             kwargs...)
    if isnothing(monomials)
        moctx = ixmonomialctx(; indices=index_type, mask_type=mask_type, deg_bound=deg_bound, kwargs...)
    end
    po = polynomialctx(coefficients, monomials = moctx)
    tbl = SigTable{pos_type, index_type, eltype(coefficients)}()
    ord_indices = Dict([(pos_type(i), gendata(pos_type(i), zero(pos_type(i)), :f)) for i in 1:ngens])
    SigPolynomialΓ{pos_type, eltype(moctx), eltype(coefficients),
                   typeof(moctx), typeof(coefficients), typeof(po), mod_order}(po, tbl, ord_indices)
end

# registration functions

function (ctx::SigPolynomialΓ{I, M, T})(sig::Tuple{I, M},
                                        pol::Polynomial{M, T},
                                        sigtail::Polynomial{M, T}) where {I, MO, M, T}
    if iszero(pol)
        ratio = one(ctx.po.mo)
    else
        ratio = div(ctx.po.mo, sig[2], leadingmonomial(pol))
    end
    val = Data{M, T}((pol, sigtail, ratio))
    try
        ctx.tbl[sig] = val
    catch
        insert!(ctx.tbl, sig, val)
    end
end

(ctx::SigPolynomialΓ{I, M, T})(sig::Tuple{I, M}, pol::Polynomial{M, T}) where {I, M, T} = ctx(sig, pol, zero(pol))

Base.getindex(ctx::SigPolynomialΓ{I, M}, sig::Tuple{I, M}) where {I, M} = getindex(ctx.tbl, sig)

# get functions

# WARNING: if sig::Tuple{J, MO} where J != I or MO != M then this will convert sig to a Tuple{I, M}

@inline function (ctx::SigPolynomialΓ{I, M})(sig::Tuple{I, M}) where {I, M, MO}
    ctx.tbl[sig]
    # get(ctx.tbl, sig) do
    #     error("Nothing registered under the signature $(sign)")
    # end
end

function (ctx::SigPolynomialΓ{I, M, T})(m::M, sig::Tuple{I, M}; orig_elem = false) where {I, M, T}
    key = (sig[1], mul(ctx.po.mo, m, sig[2]))
    if !(orig_elem)
        get(ctx.tbl, key) do
            val = ctx.tbl[sig]
            Data{M, T}((mul(ctx.po, val[:poly], m), mul(ctx.po, val[:sigtail], m), val[:sigratio]))
        end
    end
    val = ctx.tbl[sig]
    Data{M, T}((mul(ctx.po, val[:poly], m), mul(ctx.po, val[:sigtail], m), val[:sigratio]))
end

# get projection to highest index

function project(sig::Tuple{I, M},
                 sigtail::Polynomial{M, T}) where {I, M, T}

    Polynomial{M, T}(vcat(sig[2], sigtail.mo), vcat(one(T), sigtail.co))
end

function project(ctx::SigPolynomialΓ{I, M},
                 sig::Tuple{I, M}) where {I, M}
    project(sig, ctx(sig)[:sigtail])
end

# forwarding of functions on polynomials/monomials

function mul(ctx::SigPolynomialΓ{I, M}, m::M, sig::Tuple{I, M}) where {I, M}
    (sig[1], mul(ctx.po.mo, m, sig[2]))
end

function divides(ctx::SigPolynomialΓ{I, M}, s1::Tuple{I, M}, s2::Tuple{I, M}) where {I, M}
    s1[1] == s2[1] && divides(ctx.po.mo, s1[2], s2[2])
end

@inline leadingmonomial(ctx::SigPolynomialΓ{I, M}, sig::Tuple{I, M}) where {I, M} = leadingmonomial(ctx(sig)[:poly])

@inline leadingmonomial(ctx::SigPolynomialΓ{I, M}, m::M, sig::Tuple{I, M}) where {I, M} = leadingmonomial(ctx(m, sig)[:poly])

# sorting

@inline @generated function lt(ctx::SigPolynomialΓ{I, M, T, MΓ, TΓ, PΓ, S},
                               a::Tuple{I, M},
                               b::Tuple{I, M}) where {I, M, T, MΓ, TΓ, PΓ, S}

    if S == :POT
        quote
            if a[1] == b[1]
                return lt(ctx.po.mo, a[2], b[2])
            end
            return ctx.ord_indices[a[1]][:position] < ctx.ord_indices[b[1]][:position]
        end
    elseif S == :TOP
        quote
            if a[2] == b[2]
                return ctx.ord_indices[a[1]][:position] < ctx.ord_indices[b[1]][:position]
            end
            return lt(ctx.po.mo, a[2], b[2])
        end
    end
end

# Abstract Algebra

(ctx::SigPolynomialΓ)(i, m::AA.MPolyElem) = (pos_type(ctx)(i), ctx.po.mo(m))
(ctx::SigPolynomialΓ{I, M})(sig::Tuple{I, M}, p::AA.MPolyElem) where {I, M} = ctx(sig, ctx.po(p))

function (R :: AA.MPolyRing)(ctx::SigPolynomialΓ{I, M},
                             sig::Tuple{I, M}) where {I, M}

    R(ctx.po, ctx(sig)[:poly])
end


text:
export DataFrame, get, update!, clear!

"""
    DataFrame(values::CellRangeValues)::Union{Nothing,DataFrame}

Creates a DataFrame from spreadsheet range values.  The first row is converted to column names.  All other rows are converted to string values.

# Arguments
- `values::CellRangeValues`: cell values
"""
DataFrame(values::CellRangeValues)::Union{Nothing,DataFrame} = values.values == nothing ? nothing : DataFrame([values.values[1,i]=>values.values[2:end,i] for i in 1:size(values.values,2)]...)

"""
Converts sheet values to a string matrix.
"""
function _matrix(inputs::Array{String,2})::Array{String,2}
    return inputs
end

"""
Converts sheet values to a string matrix.
"""
function _matrix(inputs)::Array{String,2}
    m = fill("", length(inputs), mapreduce(length, max, inputs))

    for (i,r) in enumerate(inputs)
        for (j,v) in enumerate(r)
            m[i,j] = v
        end
    end

    return m
end

"""
    get(client::GoogleSheetsClient, range::CellRange)::CellRangeValues

Gets a range of cell values from a spreadsheet.

# Arguments
- `client::GoogleSheetsClient`: client
- `range::CellRange`: cell range
"""
function Base.get(client::GoogleSheetsClient, range::CellRange)::CellRangeValues
    result = gsheet_api_sheet_get(client; spreadsheetId=range.spreadsheet.id, majorDimension="ROWS", range=range.range)
    return CellRangeValues(CellRange(range.spreadsheet, result["range"]), haskey(result, "values") ? _matrix(result["values"]) : nothing, result["majorDimension"])
end

"""
    get(client::GoogleSheetsClient, ranges::CellRanges)::Vector{CellRangeValues}

Gets multiple ranges of cell values from a spreadsheet.

# Arguments
- `client::GoogleSheetsClient`: client
- `ranges::CellRanges`: cell ranges
"""
function Base.get(client::GoogleSheetsClient, ranges::CellRanges)::Vector{CellRangeValues}
    result = gsheet_api_sheet_batchget(client; spreadsheetId=ranges.spreadsheet.id, majorDimension="ROWS", ranges=ranges.ranges)
    return [CellRangeValues(CellRange(ranges.spreadsheet, r["range"]), haskey(r, "values") ? _matrix(r["values"]) : nothing, r["majorDimension"]) for r in result["valueRanges"] ]
end

"""
    update!(client::GoogleSheetsClient, range::CellRange, values::Array{<:Any,2}; raw::Bool=false)::UpdateSummary

Updates a range of cell values in a spreadsheet.

# Arguments
- `client::GoogleSheetsClient`: client for interacting with Google Sheets.
- `range::CellRange`: cell range to modify.
- `values::Array{<:Any,2}`: values to place in the spreadsheet.
- `raw::Bool=false`: true treats values as raw, unparsed values and and are simply
    inserted as a string.  false treats values exactly as if they were entered into
    the Google Sheets UI, for example "=A1+B1" is a formula.
"""
function update!(client::GoogleSheetsClient, range::CellRange, values::Array{<:Any,2}; raw::Bool=false)::UpdateSummary
    # There are serialization problems if the values are not Array{Any,2} or Array{String,2}
    values = Any[values[i,j] for i in 1:size(values,1), j in 1:size(values,2)]

    body = Dict(
        "values" => values,
        "majorDimension" => "ROWS",
    )

    result = gsheet_api_sheet_update(client; spreadsheetId=range.spreadsheet.id, range=range.range, valueInputOption= raw ? "RAW" : "USER_ENTERED", body=body)
    return UpdateSummary(CellRange(range.spreadsheet, result["updatedRange"]), result["updatedColumns"], result["updatedRows"], result["updatedCells"])
end

"""
    update!(client::GoogleSheetsClient, range::CellRange, df::DataFrame; kwargs...)::UpdateSummary

Updates a range of cell values in a spreadsheet.

# Arguments
- `client::GoogleSheetsClient`: client for interacting with Google Sheets.
- `range::CellRange`: cell range to modify.
- `df::DataFrame`: dataframe of values.
- `kwargs...`: keyword arguments.
"""
function update!(client::GoogleSheetsClient, range::CellRange, df::DataFrame; kwargs...)::UpdateSummary

    function df_to_string(df::DataFrame)::Matrix{String}
        rst = fill("", nrow(df)+1, ncol(df))

        for (j,n) in enumerate(names(df))
            rst[1, j] = n
        end

        for i in 1:nrow(df), j in 1:ncol(df)
            rst[i+1,j] = "$(df[i,j])"
        end

        return rst
    end

    return update!(client, range, df_to_string(df); kwargs...)
end

"""
    clear!(client::GoogleSheetsClient, range::CellRange)::UpdateSummary

Clears a range of cell values in a spreadsheet.

# Arguments
- `client::GoogleSheetsClient`: client
- `range::CellRange`: cell range
"""
function clear!(client::GoogleSheetsClient, range::CellRange)::UpdateSummary
    v = get(client, range)
    rng = v.range
    vls = v.values

    if isnothing(values)
        throw(ErrorException("No data found: range=$range"))
    end

    vls .= ""
    return update!(client, CellRange(range.spreadsheet, rng.range), vls)
end


text:
# This is a project for the computational physics course Fall 2018 at the Perimeter Institute
# The goal of this project is to get familiar with Julia language features and practice them

module MyProject_1

using Plots
using DataFrames
using Dates
using StatsBase
##############################################################################################################################################

CITY_DATA = Dict("Chicago" => "chicago.csv", "New York City" => "new_york_city.csv", "Washington" => "washington.csv")

export load_data
function load_data(city::String, skipmissing = true)
	# load data file into a dataframe
	# By default, delete the rows with missing values
	df = readtable(CITY_DATA[city])
	if skipmissing
		dropmissing!(df)
	end
	# convert the Start Time column to datetime
	dateformat = "y-m-d H:M:S"
	df[:Start_Time] = map(x->Date(x,dateformat), df[:Start_Time])
	df[:End_Time] = map(x->Date(x,dateformat), df[:End_Time])
	df
end

export time_stats
function time_stats(df::DataFrame)
	# Displays statistics on the most frequent times of travel
	println("\nCalculating The Most Frequent Times of Travel...\n")
	@time begin
		# display the most common month
		df[:month] = map(x->Dates.month(x), df[:Start_Time])
		popular_month = mode(df[:month])
		println("The most popular month is:$popular_month")
		# display the most common day of week
		df[:day] = map(x->Dates.dayofweek(x), df[:Start_Time])
		popular_day = mode(df[:day])
		println("The most popular day of week is:$popular_day")
	end
end

export station_stats
function station_stats(df::DataFrame)
	# Displays statistics on the most popular stations and trip
	println("\nCalculating The Most Popular Stations and Trip...\n")
	@time begin
	    #display most commonly used start station
	    popular_start_station = mode(df[:Start_Station])
	    println("The most popular start station is:$popular_start_station")
	    # display most commonly used end station
	    popular_end_station = mode(df[:End_Station])
	    println("The most popular end station is:$popular_end_station")
	    # display most frequent combination of start station and end station trip
	    trip = "From" .* df[:Start_Station] .* "to" .* df[:End_Station]
	    popular_trip = mode(trip)
	    println("The most popular combination of start and end stations is:$popular_trip")
	end
end

export trip_duration_stats
function trip_duration_stats(df::DataFrame)
	# Displays statistics on the total and average trip duration
	println("\nCalculating Trip Duration...\n")
	@time begin
	N = nrow(df)
	# display total travel time
	total_travel_time = sum(df[:Trip_Duration]) / 60
	println("The total travel time (in minutes) is:$total_travel_time")
	# display mean travel time
	mean_travel_time = total_travel_time / N
	println("The mean travel time (in minutes) is:$mean_travel_time")
	end
end

export user_stats
function user_stats(df::DataFrame)
	# Displays statistics on bikeshare users
	println("\nCalculating User Stats...\n")
	@time begin
		# display counts of user types
		user_types = countmap(dropmissing!(df)[:User_Type])
		@show user_types
		# display counts of gender
		if :Gender in names(df)
			gender = countmap(dropmissing!(df)[:Gender])
			@show gender
		else
			println("There is no data for gender")
		end
		# display earliest, most recent, and most common year of birth
		if :Birth_Year in names(df)
			earliest = minimum(df[:Birth_Year])
			most_recent = maximum(df[:Birth_Year])
			popular = mode(df[:Birth_Year])
			println("The earliest year of birth is $earliest \nThe most recent year of birth is $most_recent \n The most common year of birth is $popular")
		else
			println("There is no data for birth year")
		end
	end
end

export show_results
function show_results(df::DataFrame)
	time_stats(df)
	station_stats(df)
	trip_duration_stats(df)
	user_stats(df)
end

# Practice some Plots
export practice_plots
function practice_plots(df::DataFrame)
	plot(2017 .- df[df.Gender .== "Female", :][:Birth_Year],
	df[df.Gender .== "Female", :][:Trip_Duration] ./60,
	seriestype=:scatter,
	color = "red",
  	label = "Female",
  	title = "Trip Duration vs Age",
  	xlabel = "Age",
 	ylabel = "Trip Duration(min)",
  	xlims = (0,120),
  	ylims = (0,1500),
 	legend =:topright)

	plot!(2017 .- df[df.Gender .== "Male", :][:Birth_Year],
	df[df.Gender .== "Male", :][:Trip_Duration] ./60,
	seriestype=:scatter,
	color = "blue",
  	label = "Male",
 	title = "Trip Duration vs Age",
  	xlabel = "Age",
  	ylabel = "Trip Duration(min)",
 	xlims = (0,120),
	ylims = (0,1500),
	legend =:topright)
end
##############################################################################################################################################

end # module


text:
using FrankWolfe
using ProgressMeter
using Arpack
using Plots
using DoubleFloats
using ReverseDiff

using JSON
using LaTeXStrings

results = JSON.Parser.parsefile(joinpath(@__DIR__, "polynomial_result.json"))

iteration_list = [
    [x[1] + 1 for x in results["trajectory_arr_lafw"]],
    [x[1] + 1 for x in results["trajectory_arr_bcg"]],
    collect(eachindex(results["function_values_gd"])),
]
time_list = [
    [x[5] for x in results["trajectory_arr_lafw"]],
    [x[5] for x in results["trajectory_arr_bcg"]],
    results["gd_times"],
]
primal_list = [
    [x[2] - results["ref_primal_value"] for x in results["trajectory_arr_lafw"]],
    [x[2] - results["ref_primal_value"] for x in results["trajectory_arr_bcg"]],
    [x - results["ref_primal_value"] for x in results["function_values_gd"]],
]
test_list = [
    [x[6] for x in results["trajectory_arr_lafw"]],
    [x[6] for x in results["trajectory_arr_bcg"]],
    results["function_values_test_gd"],
]
label = [L"\textrm{L-AFW}", L"\textrm{BCG}", L"\textrm{GD}"]
coefficient_error_values = [
    [x[7] for x in results["trajectory_arr_lafw"]],
    [x[7] for x in results["trajectory_arr_bcg"]],
    results["coefficient_error_gd"],
]

FrankWolfe.plot_results(
    [primal_list, primal_list, test_list, test_list],
    [iteration_list, time_list, iteration_list, time_list],
    label,
    [L"\textrm{Iteration}", L"\textrm{Time}", L"\textrm{Iteration}", L"\textrm{Time}"],
    [L"\textrm{Primal Gap}", L"\textrm{Primal Gap}", L"\textrm{Test loss}", L"\textrm{Test loss}"],
    xscalelog=[:log, :identity, :log, :identity],
    legend_position=[:bottomleft, nothing, nothing, nothing],
    filename="polynomial_result.pdf",
)