#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions related to the associated Legendre functions.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Holmes, S. A. and W. E. Featherstone, 2002. A unified approach to the
#       Clenshaw summation and the recursive computation of very high degree and
#       order normalised associated Legendre functions. Journal of Geodesy,
#       76(5), pp. 279-299.
#
#       For more info.: http://mitgcm.org/~mlosch/geoidcookbook/node11.html
#
#   [2] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
#   [3] Schmidt, A (1917). Erdmagnetismus, Enzykl. Math. Wiss., 6, pp. 265–396.
#
#   [4] Winch, D. E., Ivers, D. J., Turner, J. P. R., Stening R. J (2005).
#       Geomagnetism and Schmidt quasi-normalization. Geophysical Journal
#       International, 160(2), pp. 487-504.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export denormalize_legendre, legendre!, legendre

export legendre_fully_normalized!, legendre_fully_normalized
export legendre_schmidt_quasi_normalized!, legendre_schmidt_quasi_normalized
export legendre_conventional!, legendre_conventional

"""
    legendre!([N,] P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the associated Legendre function `P_n,m[cos(ϕ)]`. The maximum degree and
order that will be computed are given by the parameters `n_max` and `m_max`. If
they are negative, then the dimensions of matrix `P` will be used.

The result will be stored at matrix `P`.

The optional parameter `N` can be used to select the normalization. The
following values are valid:

* `Val(:full)`: Compute the fully normalized associated Legendre function (see
  `legendre_fully_normalized!`).
* `Val(:schmidt)`: Compute the Schmidt quasi-normalized associated Legendre
  function (see `legendre_schmidt_quasi_normalized!`).
* `Val(:conv)`: Compute the conventional associated Legendre function (see
  `legendre_conventional!`).

If `N` is omitted, then the full normalization will be used.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

"""
legendre!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_fully_normalized!(P, float(ϕ), ph_term, n_max, m_max)

legendre!(::Val{:full}, P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_fully_normalized!(P, float(ϕ), ph_term, n_max, m_max)

legendre!(::Val{:schmidt}, P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_schmidt_quasi_normalized!(P, float(ϕ), ph_term, n_max, m_max)

legendre!(::Val{:conv}, P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_conventional!(P, float(ϕ), ph_term, n_max, m_max)

"""
    legendre([N,] ϕ::Number, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false)

Compute the associated Legendre function `P_n,m[cos(ϕ)]`. The maximum degree
that will be computed is `n_max` and the maximum order is `m_max`. Notice that
if `m_max` is higher than `n_max` or negative, than it is set to `n_max`.

The optional parameter `N` can be used to select the normalization. The
following values are valid:

* `Val(:full)`: Compute the fully normalized associated Legendre function (see
  `legendre_fully_normalized`).
* `Val(:schmidt)`: Compute the Schmidt quasi-normalized associated Legendre
  function (see `legendre_schmidt_quasi_normalized`).
* `Val(:conv)`: Compute the conventional associated Legendre function (see
  `legendre_conventional`).

If `N` is omitted, then the full normalization will be used (`Val(:full)`).

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

"""
legendre(ϕ::Number, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) =
    legendre_fully_normalized(float(ϕ), n_max, m_max, ph_term)

legendre(::Val{:full}, ϕ::Number, n_max::Integer, m_max::Integer = -1,
         ph_term::Bool = false) =
    legendre_fully_normalized(float(ϕ), n_max, m_max, ph_term)

legendre(::Val{:schmidt}, ϕ::Number, n_max::Integer, m_max::Integer = -1,
         ph_term::Bool = false) =
    legendre_schmidt_quasi_normalized(float(ϕ), n_max, m_max, ph_term)

legendre(::Val{:conv}, ϕ::Number, n_max::Integer, m_max::Integer = -1,
         ph_term::Bool = false) =
    legendre_conventional(float(ϕ), n_max, m_max, ph_term)

################################################################################
#                Fully Normalized Associated Legendre Functions
################################################################################

"""
    legendre_fully_normalized!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the fully normalized associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree and order that will be computed are given by the parameters
`n_max` and `m_max`. If they are negative, then the dimensions of matrix `P`
will be used:

    maximum degree -> number of rows
    maximum order  -> number of columns

The result will be stored at matrix `P`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Remarks

This algorithm was based on [1]. Our definition of fully normalized associated
Legendre function can be seen in [2, p. 546]. The conversion is obtained by:

                 _                     -
                |  (n-m)! . k . (2n+1)  |      k = 1 if m  = 0
    K_n,m = sqrt| --------------------- |,     k = 2 if m != 0
                |         (n+m)!        |
                 -                     -
    _
    P_n,m = P_n,m * K_n,m,

          _
    where P_n,m is the fully normalized Legendre associated function.

"""
function legendre_fully_normalized!(P::AbstractMatrix, ϕ::Number,
                                    ph_term::Bool = false, n_max::Integer = -1,
                                    m_max::Integer = -1)

    # Obtain the maximum degree and order that must be computed.
    n_max, m_max = _get_degree_and_order(P, n_max, m_max)

    # Auxiliary variables to improve code performance.
    c = cos(ϕ)
    s = sqrt(1-c^2)

    s_fact = !ph_term ? +s : -s

    @inbounds for n = 0:n_max
        # Starting values.
        if n == 0
            P[0+1,0+1] = 1
            continue
        elseif n == 1
            P[1+1,0+1] = +sqrt(3)*c

            if m_max > 0
                P[1+1,1+1] = +sqrt(3)*s_fact
            end

            continue
        end

        aux_n = (2n-1)*(2n+1)

        for m = 0:n

            if n == m
                P[n+1,n+1] = s_fact*sqrt( (2n+1)/(2n) )*P[n-1+1,n-1+1]
            else
                aux_nm = (n-m)*(n+m)
                a_nm   = sqrt( aux_n / aux_nm )
                b_nm   = sqrt( ( (2n+1)*(n+m-1)*(n-m-1) ) / ( aux_nm*(2n-3) ) )

                # We assume that the matrix is not initialized. Hence, we must
                # not access elements on the upper triangle.
                if m != n-1
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1] - b_nm*P[n-2+1,m+1]
                else
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1]
                end
            end

            # Check if the maximum desired order has been reached.
            m == m_max && break
        end
    end

    nothing
end

"""
    legendre_fully_normalized(ϕ::T, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) where T<:AbstractFloat

Compute the fully normalized associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree that will be computed is `n_max` and the maximum order is
`m_max`. Notice that if `m_max` is higher than `n_max` or negative, than it is
set to `n_max`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

# Remarks

This algorithm was based on [1]. Our definition of fully normalized associated
Legendre function can be seen in [2, p. 546]. The conversion is obtained by:

                 _                     -
                |  (n-m)! . k . (2n+1)  |      k = 1 if m  = 0
    K_n,m = sqrt| --------------------- |,     k = 2 if m != 0
                |         (n+m)!        |
                 -                     -
    _
    P_n,m = P_n,m * K_n,m,

          _
    where P_n,m is the fully normalized Legendre associated function.

"""
function legendre_fully_normalized(ϕ::T, n_max::Integer, m_max::Integer = -1,
                                   ph_term::Bool = false) where T<:AbstractFloat

    (n_max < 0) && throw(ArgumentError("n_max must be positive."))

    ( (m_max < 0) || (m_max > n_max) ) && (m_max = n_max)

    P = zeros(T, n_max+1, m_max+1)
    legendre_fully_normalized!(P, ϕ, ph_term)
    return P
end

################################################################################
#            Schmidt Quasi-Normalized Associated Legendre Functions
################################################################################

"""
    legendre_schmidt_quasi_normalized!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the Schmidt quasi-normalized associated Legendre function
`P_n,m[cos(ϕ)]` [3,4]. The maximum degree and order that will be computed are
given by the parameters `n_max` and `m_max`. If they are negative, then the
dimensions of matrix `P` will be used:

    maximum degree -> number of rows
    maximum order  -> number of columns

The result will be stored at matrix `P`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Remarks

This algorithm was based on [3,4]. The conversion is obtained by:

                 _           -
                |     (n-m)!  |    k = 1 if m  = 0
    K_n,m = sqrt| k. -------- |,   k = 2 if m != 0
                |     (n+m)!  |
                 -           -

    =
    P_n,m = P_n,m * K_n,m,

          =
    where P_n,m is the quasi-normalized normalized Legendre associated function.

"""
function legendre_schmidt_quasi_normalized!(P::AbstractMatrix, ϕ::Number,
                                            ph_term::Bool = false,
                                            n_max::Integer = -1,
                                            m_max::Integer = -1)

    # Obtain the maximum degree and order that must be computed.
    n_max, m_max = _get_degree_and_order(P, n_max, m_max)

    # Auxiliary variables to improve code performance.
    c = cos(ϕ)
    s = sqrt(1-c^2)

    s_fact = !ph_term ? +s : -s

    @inbounds for n = 0:n_max
        # Starting values.
        if n == 0
            P[0+1,0+1] = 1
            continue

        elseif n == 1
            P[1+1,0+1] = +c

            if m_max > 0
                P[1+1,1+1] = +s_fact
            end

            continue
        end

        aux_n = 2n-1 # -> sqrt( (2n-1)*(2n-1) )

        for m = 0:n

            if m == n
                P[n+1,n+1] = s_fact*sqrt( aux_n/(2n) )*P[n-1+1,n-1+1]
            else
                aux_nm = sqrt( (n-m)*(n+m) )
                a_nm   = aux_n / aux_nm
                b_nm   = sqrt( (n+m-1)*(n-m-1) ) / aux_nm

                # We assume that the matrix is not initialized. Hence, we must not
                # access elements on the upper triangle.
                if m != n-1
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1] - b_nm*P[n-2+1,m+1]
                else
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1]
                end
            end

            # Check if the maximum desired order has been reached.
            m == m_max && break
        end
    end

    nothing
end

"""
    legendre_schmidt_quasi_normalized(ϕ::T, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) where T<:AbstractFloat

Compute the Schmidt quasi-normalized associated Legendre function
`P_n,m[cos(ϕ)]`. The maximum degree that will be computed is `n_max` and the
maximum order is `m_max`. Notice that if `m_max` is higher than `n_max` or
negative, than it is set to `n_max`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

# Remarks

This algorithm was based on [3,4]. The conversion is obtained by:

                 _           -
                |     (n-m)!  |    k = 1 if m  = 0
    K_n,m = sqrt| k. -------- |,   k = 2 if m != 0
                |     (n+m)!  |
                 -           -

    =
    P_n,m = P_n,m * K_n,m,

          =
    where P_n,m is the quasi-normalized normalized Legendre associated function.

"""
function legendre_schmidt_quasi_normalized(ϕ::T, n_max::Integer,
                                           m_max::Integer = -1,
                                           ph_term::Bool = false) where T<:AbstractFloat

    (n_max < 0) && throw(ArgumentError("n_max must be positive."))

    ( (m_max < 0) || (m_max > n_max) ) && (m_max = n_max)

    P = zeros(T, n_max+1, m_max+1)
    legendre_schmidt_quasi_normalized!(P, ϕ, ph_term)
    return P
end

################################################################################
#                  Conventional Associated Legendre Function
################################################################################

"""
    legendre_conventional!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the conventional associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree and order that will be computed are given by the parameters
`n_max` and `m_max`. If they are negative, then the dimensions of matrix `P`
will be used:

    maximum degree -> number of rows
    maximum order  -> number of columns

The result will be stored at matrix `P`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

"""
function legendre_conventional!(P::AbstractMatrix, ϕ::Number,
                                ph_term::Bool = false, n_max::Integer = -1,
                                m_max::Integer = -1)

    # Obtain the maximum degree and order that must be computed.
    n_max, m_max = _get_degree_and_order(P, n_max, m_max)

    # Auxiliary variables to improve code performance.
    c = cos(ϕ)
    s = sqrt(1-c^2)

    s_fact = !ph_term ? +s : -s

    @inbounds for n = 0:n_max
        # Starting values.
        if n == 0
            P[0+1,0+1] = 1
            continue
        elseif n == 1
            P[1+1,0+1] = +c

            if m_max > 0
                P[1+1,1+1] = +s_fact
            end

            continue
        end

        aux_n = 2n-1 # -> sqrt( (2n-1)*(2n-1) )

        for m = 0:n

            if n == m
                P[n+1,n+1] = s_fact*aux_n*P[n-1+1,n-1+1]
            else
                aux_nm = n-m # -> sqrt( (n-m)*(n-m) )
                a_nm   = aux_n / aux_nm
                b_nm   = (n+m-1) / aux_nm # -> sqrt( (n+m-1)*(n+m-1) ) / aux_nm

                # We assume that the matrix is not initialized. Hence, we must
                # not access elements on the upper triangle.
                if m != n-1
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1] - b_nm*P[n-2+1,m+1]
                else
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1]
                end
            end

            # Check if the maximum desired order has been reached.
            m == m_max && break
        end
    end

    nothing
end

"""
    legendre_conventional(ϕ::T, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) where T<:AbstractFloat

Compute the conventional associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree that will be computed is `n_max` and the maximum order is
`m_max`. Notice that if `m_max` is higher than `n_max` or negative, than it is
set to `n_max`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

"""
function legendre_conventional(ϕ::T, n_max::Integer, m_max::Integer = -1,
                               ph_term::Bool = false) where T<:AbstractFloat

    (n_max < 0) && throw(ArgumentError("n_max must be positive."))

    ( (m_max < 0) || (m_max > n_max) ) && (m_max = n_max)

    P = zeros(T, n_max+1, m_max+1)
    legendre_conventional!(P, ϕ, ph_term)
    return P
end

################################################################################
#                                   Private
################################################################################

"""
    _get_degree_and_order(P, n_max, m_max)

Return the maximum degree and order to compute the Legendre associated functions
given the matrix `P` and the configuration values `n_max` and `m_max`.

"""
@inline function _get_degree_and_order(P, n_max, m_max)
    # Get the size of the matrix.
    (rows, cols) = size(P)

    # If the order or degree is less than 0, then the user wants to use all the
    # available memory.
    n_max < 0 && (n_max = rows - 1)
    m_max < 0 && (m_max = cols <= rows ? cols - 1 : n_max)

    # Make sure that the degree and order fits the matrix.
    n_max > rows - 1 && (n_max = rows - 1)
    ( (m_max > cols - 1) || (m_max > n_max) ) && (m_max = min(cols-1, n_max))

    return n_max, m_max
end


"""
    y = bquantize(x, nsd=3, abstol=eps, reltol=10*eps)

Bidirectionally quantize a n by 1 vector x to nsd signed digits, terminate early
if the error is less than the specified tolerances.

`y` is a 2-tuple of arrays with the following fields:
- `y`[1][i] is the quantized value in floating-point form
- `y`[2][i] is a 2-by-nsd (or less) matrix containing the powers of two
    (first row) and their signs (second row).

See also bunquantize.
"""
function bquantize(x, nsd=3, abstol=eps(Float64), reltol=10*eps(Float64))
    n = length(x)
    offset = -log2(0.75)

    y_val = zeros(n)
    y_csd = Array{Matrix{Int}}(undef, n)

    for i = 1:n
        xp = x[i]
        y_val[i] = 0.0
        for j = 1:nsd
            error = abs(y_val[i] - x[i])
            if error <= abstol || error <= abs(x[i])*reltol
                break
            end
            p = floor(Int, log2(abs(xp)) + offset)
            p2 = 2.0 .^ p
            sx = cmp(xp, 0)
            xp = xp - sx*p2
            y_val[i] = y_val[i] + sx*p2
            if isassigned(y_csd, i)
                y_csd[i] = hcat(y_csd[i], [p; sx])
            else
                y_csd[i] = reshape([p; sx], (2,1))
            end
        end
    end

    return y_val, y_csd
end


#### ODE Steppers for y'(t) = f(t, y(t)) #####

##### Secondary Methods ######

function Derivative(x, f, d=1, h=10.0^-8)
    """
    Estimate the derivative of f at x using the center
    finite difference technique with step size h.

    x and f can be vector valued, in which case
    d must be provided. d should be a list of only
    zeros and ones.

    For example, if x and f(x) are vectors of length four
    and we want the derivative with respect to the second
    component in x, then we call Derivative(x, f, [0, 1, 0, 0]).
    """
    (f(x+h*d) - f(x-h*d)) / (2*h)
end

function NewtonsMethod(x, f, error_tol=10.0^-10)
    """
    Estimate the root of f starting from x.

    x and f(x) are vector valued. Newton's method becomes

    \vec{x}_{n+1} = \vec{x}_{n} - inv(Jacobian) * \vec{f(\vec{x}_n)}

    I create the Jacobian by estimating the derivatives
    using a center finite difference method.
    """
    h, n = 10.0^-5, length(x)
    J = zeros(n, n)
    for j=1:n
        d = zeros(n)
        d[j] = 1
        J[:, j] = Derivative(x, f, d)
    end
    count = 0
    while norm(f(x)) > error_tol
        count += 1
        x -= inv(J) * f(x)
        if count > 50
            println("Newton's Method took unusually long so it was terminated prematurely.")
            return x
        end
    end
    x
end

##### End Secondary Methods #####

###### Single Steppers ######

## Step functions that do not include an error estimate ##
# Call these with Stepper(f, method, t0, y0, tf, h)

function EulerStep(f, t, y, h)
    h * f(t, y)
end

function ImprovedEulerStep(f, t, y, h)
    f0 = f(t, y)
    (h / 2.0) * (f0 + f(t + h, y + h * f0))
end

function BackwardEulerStep(f, t, y, h)
    """ Implicit Backward Euler Step """
    #Use Euler step as starting point for NewtonsMethod.
    NewtonsMethod(f(t, y), f1 -> h*f(t+h, y+f1) - f1)
end

function GL2Step(f, t, y, h)
    """ Implicit Gauss-Legendre Order 2 Method """
    #Use Euler step as starting point for NewtonsMethod.
    NewtonsMethod(f(t, y), f1 -> h*f(t+h/2.0, y+f1/2.0) - f1)
end

function MidpointStep(f, t, y, h)
    h * f(t + h / 2, y + (h / 2) * f(t, y))
end

function CrankNicholsonStep(f, t, y, h)
    """ Second order implicit """
    f1 = f(t, y)
    NewtonsMethod(f1, f2 -> (h/2.0) * (f1 + f(t+h, y+f2)) - f2)
end

function RK4Step(f, t, y, h)
    f1 = f(t, y)
    f2 = f(t + h / 2, y + (h / 2) * f1)
    f3 = f(t + h / 2, y + (h / 2) * f2)
    f4 = f(t + h, y + h * f3)
    (h / 6) * (f1 + f4 + 2 * (f2 + f3))
end

function Stepper(f, method, t0, y0, tf, h)
    """
    The step function `method` returns just the step,
    not an error estimate.
    """
    ts, ys = Float64[], Vector{Float64}[]
    while t0 <= tf
        push!(ts, t0)
        push!(ys, y0)
        y0 += method(f, t0, y0, h)
        t0 += h
    end
    ts, ys
end

## End step functions that do not include an error estimate ##

## Step functions that include an error estimate ##
# Call these with AdaptiveStepper(f, method, t0, y0, tf)

function SimpleErrorStep(f, t, y, h)
    """ Compare implicit Crank Nicholson to Midpoint to get error """
    step = CrankNicholsonStep(f, t, y, h)
    step, norm(step - MidpointStep(f, t, y, h))
end

function RKTrapStep(f, t, y, h)
    """ Implicit Runge-Kutta Trapazoidal Method """
    f1 = h * f(t, y)
    f2 = NewtonsMethod(f1, f2 -> h*f(t+h, y+f1/2.0+f2/2.0) - f2)

    soln = f1/2.0 + f2/2.0
    soln, norm(soln - f1)
end

function DopriStep(f, t, y, h)
    """ Dormand-Prince Method """
    f1 = h * f(t, y)
    f2 = h * f(t + h/5.0, y + f1/5.0)
    f3 = h * f(t + 3*h/10.0, y + 3*f1/40.0 + 9*f2/40.0)
    f4 = h * f(t + 4*h/5.0, y + 44*f1/45.0 - 56*f2/15.0 + 32*f3/9.0)
    f5 = h * f(t + 8*h/9.0, y + 19372*f1/6561.0 - 25360*f2/2187.0 + 64448*f3/6561.0 - 212*f4/729.0)
    f6 = h * f(t + h, y + 9017*f1/3168.0 - 355*f2/33.0 + 46732*f3/5247.0 + 49*f4/176.0 - 5103*f5/18656.0)
    soln = 35*f1/384.0 + 500*f3/1113.0 + 125*f4/192.0 - 2187*f5/6784.0 + 11*f6/84.0
    f7 = h * f(t + h, y + soln)
    alt = 5179*f1/57600.0 + 7571*f3/16695.0 + 393*f4/640.0 - 92097*f5/339200.0 + 187*f6/2100.0 + f7/40.0

    soln, norm(soln - alt)
end

function CarpStep(f, t, y, h)
    """ Cash-Karp Method """
    f1 = h * f(t, y)
    f2 = h * f(t + h/5.0, y + f1/5.0)
    f3 = h * f(t + 3*h/10, y + 3*f1/40.0 + 9*f2/40.0)
    f4 = h * f(t + 3*h/5.0, y + 3*f1/10.0 - 9*f2/10.0 + 6*f3/5.0)
    f5 = h * f(t + h, y - 11*f1/54.0 + 5*f2/2.0 - 70*f3/27.0 + 35*f4/27.0)
    f6 = h * f(t + 7*h/8.0, y + 1631*f1/55296 + 175*f2/512 + 575*f3/13824 + 44275*f4/110592.0 + 253*f5/4096.0)
    soln = 37*f1/378.0 + 250*f3/621.0 + 125*f4/594.0 + 512*f6/1771.0
    alt = 2825*f1/27648.0 + 18575*f3/48384.0 + 13525*f4/55296.0 + 277*f5/14336.0 + f6/4.0

    soln, norm(soln - alt)
end

function RKFStep(f, t, y, h)
    """ Runge-Kutta-Fehlberg Method """
    f1 = h * f(t, y)
    f2 = h * f(t + h/4.0, y + f1/4.0)
    f3 = h * f(t + 3*h/8.0, y + 3*f1/32.0 + 9*f2/32.0)
    f4 = h * f(t + 12*h/13.0, y + 1932*f1/2197.0 - 7200*f2/2197.0 + 7296*f3/2197.0)
    f5 = h * f(t + h, y + 439*f1/216.0 - 8*f2 + 3680*f3/513.0 - 845*f4/4104.0)
    f6 = h * f(t + h/2.0, y - 8*f1/27.0 + 2*f2 - 3544*f3/2565.0 + 1859*f4/4104.0 - 11*f5/40.0)
    soln = 16*f1/135.0 + 6656*f3/12825.0 + 28561*f4/56430.0 - 9*f5/50.0 + 2*f6/55.0
    alt = 25*f1/216.0 + 1408*f3/2565.0 + 2197*f4/4104.0 - f5/5.0

    soln, norm(soln - alt)
end

function AdaptiveStepper(f, method, t0, y0, tf, error_tol=10.0^-7)
    """
    The step function `method` must return the tuple,
    (step, error estimate).
    """
    ts, ys = Float64[t0], Vector{Float64}[y0]
    h, totalerror = 0.001, 0
    while t0 < tf
        if h < 10.0^-8 && h < tf - t0
            println("Step size effectively zero at t = ", t0)
            h = 0.001
        else
            if h > tf - t0
                h = tf - t0
            end
            m, error_est = method(f, t0, y0, h)
            if error_est > error_tol
                h *= .75
            else
                t0 += h
                y0 += m
                push!(ts, t0)
                push!(ys, y0)
                if error_est < error_tol / 10.0
                    h *= 1.2
                end
                totalerror += error_est
            end
        end
    end
    println("Estimated upper bound on total error of ", method, ": ", totalerror)
    ts, ys
end

## End step functions that include an error estimate ##

##### End Single Steppers #####

##### Multi Steppers ######
#Call these with MultiStepper(f, method, t0, y0, tf, h)

function AdamsBash2(f, t, y, h)
    """ Adams–Bashforth Order 2 explicit method """
    y0n = Array[y, y + RK4Step(f, t, y, h)]
    function g(t, yn)
        #yn = Array[y0, y1]
        y0, y1 = yn
        y2 = y1 + h * (1.5 * f(t + h, y1) - 0.5 * f(t, y0))
        Array[y1, y2]
    end
    g, y0n
end

function AdamsBash3(f, t, y, h)
    """ Adams–Bashforth Order 3 explicit method """
    y1 = y + RK4Step(f, t, y, h)
    y2 = y1 + RK4Step(f, t+h, y1, h)
    y0n = Array[y, y1, y2]
    function g(t, yn)
        #yn = Array[y0, y1, y2]
        y0, y1, y2 = yn
        y3 = y2 + h * (23.0 / 12.0 * f(t+2*h, y2) - 4.0 / 3.0 * f(t+h, y1) + 5.0 / 12.0 * f(t, y0))
        Array[y1, y2, y3]
    end
    g, y0n
end

function AdamsBash4(f, t, y, h)
    """ Adams–Bashforth Order 4 explicit method """
    y1 = y + RK4Step(f, t, y, h)
    y2 = y1 + RK4Step(f, t+h, y1, h)
    y3 = y2 + RK4Step(f, t+2*h, y2, h)
    y0n = Array[y, y1, y2, y3]
    function g(t, yn)
        #yn = Array[y0, y1, y2, y3]
        y0, y1, y2, y3 = yn
        y4 = y3+h * (55.0 / 24.0 * f(t+3*h, y3) - 59.0 / 24.0 * f(t+2*h, y2) + 37.0 / 24.0 * f(t+h, y1) - 3.0 / 8.0 * f(t, y0))
        Array[y1, y2, y3, y4]
    end
    g, y0n
end

function AdamsMoult2(f, t, y, h)
    """ Adams-Moulton Order 2 implicit method """
    y1 = y + RK4Step(f, t, y, h)
    y0n = Array[y, y1]
    function g(t, yn)
        y0, y1 = yn
        f0, f1 = f(t, y0), f(t+h, y1)
        y2 = NewtonsMethod(y1, y2 -> y1+h*(5*f(t+2*h, y2)/12.0 + 2*f1/3.0 - f0/12.0) - y2)
        Array[y1, y2]
    end
    g, y0n
end

function AdamsMoult3(f, t, y, h)
    """ Adams-Moulton Order 3 implicit method """
    y1 = y + RK4Step(f, t, y, h)
    y2 = y1 + RK4Step(f, t+h, y1, h)
    y0n = Array[y, y1, y2]
    function g(t, yn)
        y0, y1, y2 = yn
        f0, f1, f2 = f(t, y0), f(t+h, y1), f(t+2*h, y2)
        y3 = NewtonsMethod(y2, y3 -> y2+h*(3*f(t+3*h, y3)/8.0 + 19*f2/24.0 - 5*f1/24.0 + f0/24.0)-y3)
        Array[y1, y2, y3]
    end
    g, y0n
end

function AdamsMoult4(f, t, y, h)
    """ Adams-Moulton Order 4 implicit method """
    y1 = y + RK4Step(f, t, y, h)
    y2 = y1 + RK4Step(f, t+h, y1, h)
    y3 = y2 + RK4Step(f, t+2*h, y2, h)
    y0n = Array[y, y1, y2, y3]
    function g(t, yn)
        y0, y1, y2, y3 = yn
        f0, f1, f2, f3 = f(t, y0), f(t+h, y1), f(t+2*h, y2), f(t+3*h, y3)
        y4 = NewtonsMethod(y3, y4 -> y3+h*(251*f(t+4*h, y4)/720.0+646*f3/720.0 - 264*f2/720.0 + 106*f1/720.0 - 19*f0/720.0)-y4)
        Array[y1, y2, y3, y4]
    end
    g, y0n
end

function MultiStepper(f, method, t0, y0, tf, h)
    """
    method returns g(t, yn), y0n with y0n is
    an array  of previous values of y.
    """
    g, y0n = method(f, t0, y0, h)
    #y0n = Array[y0, y1, ..., yn]
    ts, ys = Float64[], Vector{Float64}[]
    while t0 <= tf + h
        push!(ts, t0)
        push!(ys, y0n[1])
        y0n = g(t0, y0n)
        t0 += h
    end
    ts, ys
end

##### End Multi Steppers #####

#### Test ####
function func(t, y)
    y1, ydot = y
    [ydot, -t*y1]
end

t, tmax, h = 0.0, 20, 0.01
y0 = [-1.0, 0.0]

#xs, ys = MultiStepper(func, AdamsMoult4, t, copy(y0), tmax, h)
xs, ys = Stepper(func, GL2Step, t, copy(y0), tmax, h)
#xs, ys = AdaptiveStepper(func, RKTrapStep, t, copy(y0), tmax, h)
figure(0)
plot(xs, [ys[i][1] for i=1:length(ys)], "o-", label="Position")
plot(xs, [ys[i][2] for i=1:length(ys)], "x--", label="Velocity")
legend()

xs, ys = AdaptiveStepper(func, RKFStep, t, copy(y0), tmax)
figure(1)
plot(xs, [ys[i][1] for i=1:length(ys)], "o-", label="Position")
plot(xs, [ys[i][2] for i=1:length(ys)], "x--", label="Velocity")
legend()


using DataFrames, Plots
pyplot()

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  function euler_ex71(f::Function, y0::Float64, x::FloatRange{Float64})
    y = Vector{Real}(length(x))
    y[1] = y0
    h = x.step/x.divisor
    for i in 2:length(x)
      k0 = h * f(x[i-1], y[i-1])
      k1 = h * f(x[i-1] + h, y[i-1] + k0)
      y[i] = y[i-1] + 1/2*(k0 + k1)
    end
    columns=[]
    push!(columns, collect(x))
    push!(columns, y)
    df = DataFrame(columns, [:x, :y])
  end

  f(x,y) = (x+y)/x
  fy(x) = x*(1+log(x/2))

  x1 = 2.0:0.25:2.5
  x2 = 2.0:0.1:2.5
  x3 = 1.8:0.1:2.6
  df1 = euler_ex71(f, 2.0, x1)
  df2 = euler_ex71(f, 2.0, x2)

  df1 |> display
  println()
  df2 |> display
  println()

  global p2

  p1 = plot(x3, fy.(x3), ylims=(1.8, 3.5))
  plot!(p1, df1[:x], df1[:y], line=(:dash, 1), marker=(:plus), color=:red)
  plot!(p1, df2[:x], df2[:y], line=(:dash, 1), marker=(:circle), color=:red)
  savefig("euler_ex7.2.png")
  plot(p1, layout=(1,1))
  gui()
end


## PART 4

println("\nPart 4");

bvals = [b-3:b+3];
println("\n\tEstimating σ for f1..f3 with bin around b=$b");

σf1jack = Float64[];
σf2jack = Float64[];
σf3jack = Float64[];

for bb in bvals
    jestb = jacknife(v, N2, bb);
    push!(σf1jack, jacknifeσ(f1, jestb));
    push!(σf2jack, jacknifeσ(f2, jestb));
    push!(σf3jack, jacknifeσ(f3, jestb));
end

σfjack = cat(2, σf1jack, σf2jack, σf3jack);

println("\n\tGenerating plots in σfvsb to show σf dependence on b.");
using Gadfly
σfvsb = Plot[];
for i in [1:3]
    push!(σfvsb, Gadfly.plot(
        x = bvals,
        y = σfjack[:,i],
        Geom.line,
        Geom.point,
        Guide.xlabel("Bin size b"),
        Guide.ylabel("Estimated Variance"),
        Guide.title("Estimated Variance of function f$i vs. Bin Size")
    ));
end

println("\n\tCompare σfjack to naive σ̂ from part 1:\n")
println("\t\tσ̂f1N2/σfjack[4,1] = $(σ̂f1N2/σfjack[4,1])");
println("\t\tσ̂f2N2/σfjack[4,2] = $(σ̂f2N2/σfjack[4,2])");
println("\t\tσ̂f3N2/σfjack[4,3] = $(σ̂f3N2/σfjack[4,3])");


using Minpack: hybrd
using Base.Test

@testset "analytic" begin

function residual(xvec, arg1, arg2)

    fvec = zeros(xvec)

    fvec[1] = xvec[1] - arg1
    fvec[2] = xvec[1] + xvec[2] - arg2

    return fvec
end

x0 = [1.0, 0.0]
args = (3.0, 1.0)
tol = 1e-8

x, f = hybrd(residual, x0, args, tol)

@test isapprox(x[1], 3.0; atol=1e-8)
@test isapprox(x[2], -2.0; atol=1e-8)
@test isapprox(f[1], 0.0; atol=1e-8)
@test isapprox(f[2], 0.0; atol=1e-8)

function residual2(xvec)

    fvec = zeros(xvec)

    fvec[1] = exp(-exp(-(xvec[1]+xvec[2]))) - xvec[2]*(1+xvec[1]^2)
    fvec[2] = xvec[1]*cos(xvec[2]) + xvec[2]*sin(xvec[1]) - 0.5

    return fvec
end

x0 = [0.0, 0.0]
args = ()
tol = 1e-8

x, f = hybrd(residual2, x0, args, tol)

@test isapprox(x[1], 0.3532; atol=1e-4)
@test isapprox(x[2], 0.6061; atol=1e-4)
@test isapprox(f[1], 0.0; atol=1e-8)
@test isapprox(f[2], 0.0; atol=1e-8)

end  # test set


# # [Open Circular Port Graph Examples](@id cpg_examples)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/CPG_examples.ipynb)

using AlgebraicDynamics.DWDDynam
using AlgebraicDynamics.CPortGraphDynam
using AlgebraicDynamics.CPortGraphDynam: draw, barbell, gridpath, grid, meshpath

using Catlab
using Catlab.WiringDiagrams
using Catlab.WiringDiagrams.CPortGraphs
using Catlab.Theories
using Catlab.CategoricalAlgebra

using OrdinaryDiffEq
using DynamicalSystems
using Plots, Plots.PlotMeasures

using PrettyTables

# ## SIR Epidemiology Model

# An SIR epidemiology model has three types of people: susceptible, infected, and recovered. When a susceptible person interacts with infected person, the susceptible person also becomes infected. Over time infected people  recover. Transition rates determine the how frequently susceptible people come into contact with infected people and how fast infected people recover. The system evolves according to the law of mass action.

# In a multi-city SIR model, each city has susceptible, infected, and recovered populations. To see the spread of the disease we will consider both susceptible and infected people moving between cities. To define a multi-city SIR model, we can compose multiple single-city SIR models using the composition syntax of open CPGs. The composition pattern will consist of three boxes each of which will be filled by a single-city SIR model. Ports expose the susceptible and infected populations of each city. One set of wires connect the susceptible and infected populations of cities 1 and 2. A second set of wires connect the susceptible and infected popuation of cities 2 and 3.

## Define the composition pattern
d₀ = OpenCPortGraph()
add_parts!(d₀, :Box, 1)
d₁ = barbell(2)
F = ACSetTransformation((Box=[2],), d₀, d₁)
G = ACSetTransformation((Box=[1],), d₀, d₁)
d₂ = apex(pushout(F,G))

## Define the primitive systems
β, μ, α₁, α₂ = 0.4, 0.4, 0.01, 0.01

sirfuncb = (u,x,p,t)->[-β*u[1]*u[2] - α₁*(u[1]-x[1]), # Ṡ
                        β*u[1]*u[2] - μ*u[2] - α₂*(u[2]-x[2]), #İ
                        μ*u[2] # Ṙ
                        ]
sirfuncm = (u,x,p,t)->[-β*u[1]*u[2] - α₁*(u[1]-(x[1]+x[3])/2),
                        β*u[1]*u[2] - μ*u[2] - α₂*(u[2]-(x[2]+x[4])/2),
                        μ*u[2]
                       ]

boundary  = ContinuousMachine{Float64}(2,3,sirfuncb, u->u[1:2])
middle    = ContinuousMachine{Float64}(4,3, sirfuncm, u->u[[1,2,1,2]])

## Compose
threecity = oapply(d₂, [boundary,middle,boundary])

# First, we will approximate the solution to the three city SIR model using Euler's method. The initial condition has 100 susceptible people in each city a single infected person in the first city. We show the infected populations in each city over time.

u0 = [100,1,0,100,0,0,100,0,0.0]

h = 0.01
nsteps = 100

threecity_approx = euler_approx(threecity, h)
traj = trajectory(threecity_approx, u0, [], nothing, nsteps)

map(traj) do u
    return (i1=u[2], i2=u[5], i3=u[8])
end |> pretty_table

# Next, we will solve the continuous system and plot the results. Over time the infected populations increase and the susceptible populations decrease. The delays in the plots illustrate how the disease spreads from city 1 to city 2 and then from city 2 to city 3.

## Solve and plot
tspan = (0.0, 1.0)

prob = ODEProblem(threecity, u0, tspan)
sol = solve(prob, Tsit5(); dtmax = 0.01)

plot(sol, lw=2, title = "SIR Epidemiology Model", bottom_margin=10mm, left_margin=10mm,
    label=["S" "I" "R"])

# ## Cellular automata

#  Circular port graphs are particularly useful for modeling systems where the composition pattern is given by a grid and where the dynamics are repetative. In the case of cellular automata the composition pattern is a row of $n$ cells each of which is connected to its two neighbors. The primitive systems are identical machines whose discrete dynamics are a specified rule. See [here](https://mathworld.wolfram.com/ElementaryCellularAutomaton.html) for a complete set of rules and the patterns they generate.

function Rule(k::Int)
    (left_neighbor, x, right_neighbor) ->
    Bool(digits(k, base=2, pad=8)[1 + right_neighbor + 2*x + 4*left_neighbor])
end

## Define the composition pattern
n = 100
row = apex(gridpath(n, 1))

## Define the primitive system which will be repeated for each cell
rule = DiscreteMachine{Bool}(2, 1, 2, (u, x, p, t)->Rule(p)(x[2], u[1], x[1]),
            u->[u[1], u[1]])

## Compose
automaton = oapply(row, rule)

## Solve and plot
u0 = zeros(Int, n); u0[Int(n/2)] = 1

rule_number = 126
traj = trajectory(automaton, u0, [0,0], rule_number, 100)
spy(Matrix(traj))


using Statistics, Random, Distributions

# returns the variable (or matrix), lagged p times,
# with the first p rows filled with ones (to avoid divide errors)
# remember to drop those rows before doing analysis
function lag(x,p::Int64)
	n = size(x,1)
	lagged_x = zeros(eltype(x),n,p)
	lagged_x = [ones(p); x[1:n-p]]
end

# lags of a vector from 1 to p in a p column array
function  lags(x,p)
	n = size(x,1)
	lagged_x = zeros(eltype(x),n,p)
	for i = 1:p
		lagged_x[:,i] = lag(x,i)
	end
    return lagged_x
end

# compute moving average using p most recent values, including current value
function ma(x, p)
    m = similar(x)
    for i = p:size(x,1)
        m[i] = mean(x[i-p+1:i])
    end
    return m
end

# auxiliary model: HAR-RV
# Corsi, Fulvio. "A simple approximate long-memory model
# of realized volatility." Journal of Financial Econometrics 7,
# no. 2 (2009): 174-196.
function HAR(y)
    ylags = lags(y,10)
    X = [ones(size(y,1)) ylags[:,1]  mean(ylags[:,1:4],dims=2) mean(ylags[:,1:10],dims=2)]
    # drop missings
    y = y[11:end]
    X = X[11:end,:]
    βhat = X\y
    σhat = std(y-X*βhat)
    vcat(βhat,σhat)
end

function aux_stat(y)
    y, m, s  = stnorm(abs.(y))
    # look for evidence of volatility clusters
    mm = ma(y,5)
    mm = mm[5:end]
    clusters = quantile(mm,0.75) -quantile(mm, 0.25)
    ϕ = HAR(y)
    vcat(m, s, clusters, ϕ)
end

# the dgp: simple discrete time stochastic volatility (SV) model
function SVmodel(σe, ρ, σu, n, shocks_u, shocks_e)
    burnin = size(shocks_u,1) - n
    hlag = 0.0
    h = ρ.*hlag .+ σu.*shocks_u[1] # figure out type
    y = σe.*exp(h./2.0).*shocks_e[1]
    ys = zeros(eltype(y),n)
    for t = 1:burnin+n
        h = ρ.*hlag .+ σu.*shocks_u[t]
        y = σe.*exp(h./2.0).*shocks_e[t]
        if t > burnin
            ys[t-burnin] = y
        end
        hlag = h
    end
    #plot(ys)
    sqrt(n)*aux_stat(ys)
end


#=
  From Joost-Pieter Katoen
  "Probabilistic Programming Quantitative Modeling for the Masses?"
  (MMB 2018 Conference, Erlangen)

  Slide 4:
  """
  How likely does a student end up with a bad mood after getting
  a bad grade for an easy exam, given that she is well prepared?
  """

Distributions of variable difficulty
easy       =>    6022  (0.602200)
hard       =>    3978  (0.397800)

Distributions of variable preparation
good       =>   10000  (1.000000)

Distributions of variable grade
good       =>    6821  (0.682100)
bad        =>    3179  (0.317900)

Distributions of variable mood
bad        =>    5004  (0.500400)
good       =>    4996  (0.499600)

  Cf ~/webppl/student_mood_after_exam.wppl

=#
using Turing, StatsPlots, Distributions
include("jl_utils.jl")

@model function student_mood_after_exam()
    bad = 1
    good = 2
    badGood = [bad,good]

    easy = 1
    hard = 2

    difficulty ~ Categorical([0.6,0.4]) # [easy,hard]

    preparation ~ Categorical([0.7,0.3]) # badGood

    grade ~
        (difficulty==easy && preparation==bad)  ? Categorical([0.95,0.05]) : # badGood
        (difficulty==easy && preparation==good) ? Categorical([0.5,0.5])   :
        (difficulty==hard && preparation==bad)  ? Categorical([0.6,0.4])   :
        (difficulty==hard && preparation==good) ? Categorical([0.05,0.95]) : Dirac(easy)

    mood ~ grade == bad ? # badGood
        Categorical([0.9,0.1]) :
        Categorical([0.3,0.7])

    # true ~ Dirac(difficulty == easy)
    true ~ Dirac(preparation == good)
    # true ~ Dirac(grade == bad)
    # true ~ Dirac(mood == bad)

end

model = student_mood_after_exam()

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

bad_good = ["bad","good"]
show_var_dist_pct(chns, :difficulty, ["easy","hard"])

println()
show_var_dist_pct(chns, :preparation, bad_good)

println()
show_var_dist_pct(chns, :grade, bad_good)

println()
show_var_dist_pct(chns, :mood, bad_good)


# indicator of the L2 norm ball with given radius

export IndBallL2

"""
**Indicator of a Euclidean ball**

    IndBallL2(r=1.0)

Returns the indicator function of the set
```math
S = \\{ x : \\|x\\| \\leq r \\},
```
where ``\\|\\cdot\\|`` is the ``L_2`` (Euclidean) norm. Parameter `r` must be positive.
"""
struct IndBallL2{R <: Real} <: ProximableFunction
    r::R
    function IndBallL2{R}(r::R) where {R <: Real}
        if r <= 0
            error("parameter r must be positive")
        else
            new(r)
        end
    end
end

is_convex(f::IndBallL2) = true
is_set(f::IndBallL2) = true

IndBallL2(r::R=1.0) where {R <: Real} = IndBallL2{R}(r)

function (f::IndBallL2)(x::AbstractArray{T}) where {R <: Real, T <: RealOrComplex{R}}
    if isapprox_le(norm(x), f.r, atol=eps(R), rtol=sqrt(eps(R)))
        return R(0)
    end
    return R(Inf)
end

function prox!(y::AbstractArray{T}, f::IndBallL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    scal = f.r/norm(x)
    if scal > 1
        y .= x
        return R(0)
    end
    for k in eachindex(x)
        y[k] = scal*x[k]
    end
    return R(0)
end

fun_name(f::IndBallL2) = "indicator of an L2 norm ball"
fun_dom(f::IndBallL2) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndBallL2) = "x ↦ 0 if ||x|| ⩽ r, +∞ otherwise"
fun_params(f::IndBallL2) = "r = $(f.r)"

function prox_naive(f::IndBallL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    normx = norm(x)
    if normx > f.r
        y = (f.r/normx)*x
    else
        y = x
    end
    return y, R(0)
end


# indicator of the ball of matrices with (at most) a given rank

using LinearAlgebra
using TSVD

export IndBallRank

"""
**Indicator of rank ball**

    IndBallRank(r=1)

Returns the indicator function of the set of matrices of rank at most `r`:
```math
S = \\{ X : \\mathrm{rank}(X) \\leq r \\},
```
Parameter `r` must be a positive integer.
"""
struct IndBallRank{I <: Integer} <: ProximableFunction
    r::I
    function IndBallRank{I}(r::I) where {I <: Integer}
        if r <= 0
            error("parameter r must be a positive integer")
        else
            new(r)
        end
    end
end

is_set(f::IndBallRank) = true
is_prox_accurate(f::IndBallRank) = false

IndBallRank(r::I=1) where {I <: Integer} = IndBallRank{I}(r)

function (f::IndBallRank)(x::AbstractArray{T, 2}) where {R <: Real, T <: RealOrComplex{R}}
    maxr = minimum(size(x))
    if maxr <= f.r return R(0) end
    U, S, V = tsvd(x, f.r+1)
    # the tolerance in the following line should be customizable
    if S[end]/S[1] <= 1e-7
        return R(0)
    end
    return R(Inf)
end

function prox!(y::AbstractMatrix{T}, f::IndBallRank, x::AbstractMatrix{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    maxr = minimum(size(x))
    if maxr <= f.r
        y .= x
        return R(0)
    end
    U, S, V = tsvd(x, f.r)
    # TODO: the order of the following matrix products should depend on the shape of x
    M = S .* V'
    mul!(y, U, M)
    return R(0)
end

fun_name(f::IndBallRank) = "indicator of the set of rank-r matrices"
fun_dom(f::IndBallRank) = "AbstractArray{Real,2}, AbstractArray{Complex,2}"
fun_expr(f::IndBallRank) = "x ↦ 0 if rank(x) ⩽ r, +∞ otherwise"
fun_params(f::IndBallRank) = "r = $(f.r)"

function prox_naive(f::IndBallRank, x::AbstractMatrix{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    maxr = minimum(size(x))
    if maxr <= f.r
        y = x
        return y, R(0)
    end
    F = svd(x)
    y = F.U[:,1:f.r]*(Diagonal(F.S[1:f.r])*F.V[:,1:f.r]')
    return y, R(0)
end


"""
    Ridge(alpha = 0.1)
Ridge Regression. `alpha` is the value multiplied by regularization term.

# Example
```jldoctest regression
julia> model = Ridge()
Ridge(Float64[], 0.1)

julia> fit!(model, x, t)
3-element Vector{Float64}:
 -0.5635468573581848
  2.1185952951687614
 40.334109796666425

julia> model(x)
20-element Vector{Float64}:
 175.12510514593038
 170.28763505842625
 149.54478779081344
 159.7466476176333
 165.4525001910219
 129.6935485937018
 164.79709438985097
 161.3621431448216
 170.17418141858434
 176.95192713546982
 174.8078461898064
 168.76930346791391
 171.09202561187362
 170.58861763111338
 155.04089855305028
 168.05151465675456
 149.76311329450505
 169.5183634524783
 141.7695072903308
 163.94744858842117
```
"""
mutable struct Ridge
    w::Array
    α::Float64
    Ridge(alpha = 0.1) = new(Array{Float64}(undef, 0), alpha)
end

function fit!(model::Ridge, x, t)
    check_size(x, t)
    x = expand(x)
    n = size(x, 1)
    _I = Matrix{Float64}(I, n, n)
    model.w = inv(x * x' .+ model.α * _I) * x * t
end

(model::Ridge)(x) = expand(x)' * model.w


#=
The number 3797 has an interesting property. Being prime itself, it is possible to continuously remove digits from left to right, and remain prime at each stage: 3797, 797, 97, and 7. Similarly we can work from right to left: 3797, 379, 37, and 3.

Find the sum of the only eleven primes that are both truncatable from left to right and right to left.

NOTE: 2, 3, 5, and 7 are not considered to be truncatable primes.
=#
include("projecteulerutils.jl")
using Primes

function istruncatableprime(a)
  da = reverse(digits(a))
  f1,f2 = copy(da), copy(da)
  for i in 1:length(da)-1
    deleteat!(f1,1)
    pop!(f2)
    if !isprime(nd(f1)) || !isprime(nd(f2))
      return false
    end
  end
  true
end

function calc()
  i,count,total = 10,0,0
  while count < 11
    if isprime(i) && istruncatableprime(i)
      count += 1
      total += i
    end
    i += 1
  end
  total
end
@time println(calc())


module solver

function euler(y0::AbstractArray{Float64}, f, a::Float64, b::Float64, h::Float64)::AbstractArray{Float64}
    k = ceil(Integer, (b - a) / h)
    t_range = range(a, stop=b, length=k)
    y = y0
    for t_n = t_range
        if t_n == a
            continue
        end
        y_n = y[:,end]
        y_next = y_n + h * f(t_n, y_n)
        y = [y y_next]
    end
    return y
end

function midpoint(y0::AbstractArray{Float64}, f, a::Float64, b::Float64, h::Float64)::AbstractArray{Float64}
    k = ceil(Integer, (b - a) / h)
    t_range = range(a, stop=b, length=k)
    y = y0
    for t_n = t_range
        if t_n == a
            continue
        end
        y_n = y[:,end]
        y_next = y_n + h * f(t_n + h / 2.0, y_n + h / 2.0 * f(t_n, y_n))
        y = [y y_next]
    end
    return y
end

end # module


const madsen = let res_init=zeros(3), jac_init=zeros(3,2), x_init=[3.0, 1.0]

    function res(x, r)
        r[1] = x[1] ^ 2 + x[2] ^ 2 + x[1] * x[2]
        r[2] = sin(x[1])
        r[3] = cos(x[2])
        return r
    end

    function jac(x, j)
        j[1,1] = 2.e0 * x[1] + x[2]
        j[1,2] = 2.e0*x[2] + x[1]
        j[2,1] = cos(x[1])
        j[2,2] = 0.e0
        j[3,1] = 0.e0
        j[3,2] = -sin(x[2])
        return j
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("madsen", res, jac,
                                                   res_init, jac_init, x_init;
                                                   scale=scale, verbose=verbose,
                                                   print_steps=print_steps)
end


module BasesAndSamples

using Combinatorics

export basis_monomial, basis_laguerre, basis_jacobi, basis_chebyshev, basis_gegenbauer
export sample_points_simplex, sample_points_padua, sample_points_rescaled_laguerre, sample_points_chebyshev, sample_points_chebyshev_mod

"""
    basis_monomial(d::Int, x...)

Generate the monomial basis in variables x... up to degree d (inclusive).
"""
function basis_monomial(d::Int, x...)
    n = length(x)
    if n > 1
        q = Vector{typeof(prod(one.(x)))}(undef, binomial(n+d,d))
    elseif n == 1
        q = Vector{typeof(one(x...))}(undef, binomial(n+d,d))
    end
    idx = 1
    for k=0:d
        for exponent in multiexponents(n,k)
            q[idx] = prod([x[i]^exponent[i] for i=1:length(exponent)])
            idx+=1
        end
    end
    return q
end

"""
    basis_laguerre(d::Integer, alpha, x)

Generate the Laguerre polynomials with parameter alpha up to degree d (inclusive). (https://en.wikipedia.org/wiki/Laguerre_polynomials)
"""
function basis_laguerre(d::Integer, alpha, x)
    v = Vector{typeof(one(alpha) * one(x))}(undef, 1 + d)
    v[1] = one(x)
    d == 0 && return v
    v[2] = 1 + alpha - x
    d == 1 && return v
    for l = 2:d
        v[l+1] = 1 // big(l) * ((2l - 1 + alpha - x) * v[l] - (l + alpha - 1) * v[l-1])
    end
    return v
end

"""
    basis_jacobi(d::Integer, alpha, beta, x, normalized = true)

Generate the Jacobi polynomials with parameters alpha and beta up to degree d (inclusive). (https://en.wikipedia.org/wiki/Jacobi_polynomials)
"""
function basis_jacobi(d::Integer, alpha, beta, x, normalized = true)
    q = Vector{typeof(one(alpha) * one(x))}(undef, d + 1)
    q[1] = one(x)
    d == 0 && return q
    q[2] = x # normalized
    if !normalized
        q[2] *= (alpha + 1)
    end
    d == 1 && return q
    for k = 2:d
        # what if alpha+beta = -n for some integer 1 <= n < 2d ?
        q[k+1] =
            (2 * k + alpha + beta - 1) /
            BigFloat(2k * (k + alpha + beta) * (2k + alpha + beta - 2)) *
            ((2 * k + alpha + beta) * (2k + alpha + beta - 2) * x + beta^2 - alpha^2) *
            q[k] + -2 * (k + alpha - 1) * (k + beta - 1) * (2 * k + alpha + beta) * q[k-1]
    end
    return q
end

"""
    basis_chebyshev(d::Int,x)

Generate a basis of chebyshev polynomials up to degree d (inclusive). (https://en.wikipedia.org/wiki/Chebyshev_polynomials)
"""
function basis_chebyshev(d::Int,x)
    v = Vector{typeof(one(x))}(undef, 1 + d)
    v[1] = one(x)
    d == 0 && return v
    v[2] = x
    d == 1 && return v
    for l = 2:d
        v[l+1] = 2x*v[l] - v[l-1]
    end
    return v
end

"""
    basis_gegenbauer(d, n, x)

Basis for the Gegenbauer polynomials in dimension n up to degree d.
 This is the Gegenbauer polynomial with parameter lambda = n/2-1,
 or the Jacobi polynomial with alpha = beta = (n-3)/2.
 Normalized to evaluate to 1 at 1.
 Taken from arxiv/2001.00256, ancillary files, SemidefiniteProgramming.jl.
"""
function basis_gegenbauer(d, n, x)
    v = Vector{typeof(one(x))}(undef, 1 + d)
    v[1] = one(x)
    d == 0 && return v
    v[2] = x
    d == 1 && return v
    for l = 2:d
        v[l+1] = (2l + n - 4) // (l + n - 3) * x * v[l] - (l - 1) // (l + n - 3) * v[l-1]
    end
    return v
end

"""
    sample_points_simplex(n, d)

Generate the rational sample points in the unit simplex with denominator d.
"""
function sample_points_simplex(n, d)
    #rational points in the unit simplex with denominator d
    #probably not very efficient, but I dont know how to do it better for general n.
    x = [zeros(BigFloat, n) for i = 1:binomial(n + d, d)] #need n+d choose d points for a unisolvent set, if symmetry is not used.
    idx = 1
    for I in CartesianIndices(ntuple(k -> 0:d, n)) #all tuples with elements in 0:d of length n
        if sum(Tuple(I)) <= d #in unit simplex
            x[idx] = [i / BigFloat(d) for i in Tuple(I)]
            idx += 1
        end
    end
    return x
end

"""
    sample_points_padua(d)

Generate the Padua points for degree d. (https://en.wikipedia.org/wiki/Padua_points)
"""
function sample_points_padua(d)
    #padua points:
    z = [Array{BigFloat}(undef, 2) for i = 1:binomial(2 + d, d)]
    z_idx = 1
    for j = 0:d
        delta_j = j % 2 == d % 2 == 1 ? 1 : 0
        mu_j = cospi(j / d)
        for k = 1:(div(d, 2)+1+delta_j)
            eta_k = j % 2 == 1 ? cospi((2 * k - 2) / (d + 1)) : cospi((2 * k - 1) / (d + 1))
            z[z_idx] = [mu_j, eta_k]
            z_idx += 1
        end
    end
    return z
end

"""
    sample_points_rescaled_laguerre(d)

Generate 'rescaled laguerre' points, as in SDPB.
"""
function sample_points_rescaled_laguerre(d)
    #as done in sdpb: ('rescaled Laguerre')
    # x[k] = sqrt(pi)/(64*(d+1)*log( 3- 2*sqrt(2))) * (-1+4*k)^2, with k=0:d
    constant = -sqrt(BigFloat(pi)) / (64 * (d + 1) * log(3 - 2 * sqrt(BigFloat(2))))
    x = zeros(BigFloat, d + 1)
    for k = 0:d
        x[k+1] = constant * (-1 + 4 * k)^2
    end
    return [[i] for i in x]
end

"""
    sample_points_chebyshev(d, a = -1, b = 1)

Generate the d+1 chebyshev points in [a,b]. (https://en.wikipedia.org/wiki/Chebyshev_nodes)
"""
function sample_points_chebyshev(d, a = -1, b = 1)
    #roots of chebyshev polynomials of the first kind, unisolvent for polynomials up to degree d
    points = [
        (a + b) / BigFloat(2) +
        (b - a) / BigFloat(2) * cos((2k - 1) / BigFloat(2(d + 1)) * BigFloat(pi)) for
        k = 1:d+1
    ]
    return [[i] for i in points]
end

"""
    sample_points_chebyshev_mod(d, a = -1, b = 1)

Generate the d+1 modified chebyshev points in [a,b], the chebyshev points divided by cos(pi/(2(d+1))).
"""
function sample_points_chebyshev_mod(d, a = -1, b = 1)
    #roots of chebyshev polynomials of the first kind, divided by cos(pi/2(d+1)) to get a lower lebesgue constant
    points = [
        (a + b) / BigFloat(2) +
        (b - a) / BigFloat(2) * cos((2k - 1) / BigFloat(2(d + 1)) * BigFloat(pi)) /
        cos(BigFloat(pi) / 2(d + 1)) for k = 1:d+1
    ]
    return [[i] for i in points]
end

end # end module


module Transformation

export Vec, Quat, Transf, dot, cross, magnitude, norm, ×, slerp, unit

"""
    Vec - three dimensional vector.

# Examples
```julia-repl
julia> Vec(1, 2, 3)
Vec(1, 2, 3)
```
"""
struct Vec
    x::Real
    y::Real
    z::Real
    function Vec(x::Real=0, y::Real=0, z::Real=0)::Vec
        new(x, y, z)
    end
end

function Base.:+(a::Vec, b::Vec)::Vec
    Vec(
        a.x + b.x,
        a.y + b.y,
        a.z + b.z
    )
end

function Base.:-(a::Vec, b::Vec)::Vec
    Vec(
        a.x - b.x,
        a.y - b.y,
        a.z - b.z
    )
end

function Base.:-(a::Vec)::Vec
    Vec(
        -a.x,
        -a.y,
        -a.z
    )
end

"""
    dot(a::Vec, b::Vec)::Vec

Compute the dot product of two vectors.
"""
function dot(a::Vec, b::Vec)
    a.x * b.x + a.y * b.y + a.z * b.z
end

"""
    cross(a::Vec, b::Vec)::Vec

Compute the cross product of two vectors.
"""
function cross(a::Vec, b::Vec)::Vec
    Vec(
        a.y * b.z - a.z * b.y,
        a.z * b.x - a.x * b.z,
        a.x * b.y - a.y * b.x
    )
end

"""
    ×(a::Vec, b::Vec)::Vec

Compute the cross product of two vectors.
"""
function ×(a::Vec, b::Vec)::Vec
    cross(a, b)
end

function Base.:*(a::Real, b::Vec)::Vec
    Vec(
        a * b.x,
        a * b.y,
        a * b.z
    )
end

function Base.:*(a::Vec, b::Real)::Vec
    Vec(
        a.x * b,
        a.y * b,
        a.z * b
    )
end

function Base.:/(a::Vec, b::Real)::Vec
    Vec(
        a.x / b,
        a.y / b,
        a.z / b
    )
end

function Base.isequal(a::Vec, b::Vec)::Bool
    a.x == b.x && a.y == b.y && a.z == b.z
end

"""
    norm(a::Vec)::Real

Compute the Euclidean norm of the given vector.
"""
function norm(a::Vec)::Real
    √(a.x^2 + a.y^2 + a.z^2)
end

"""
    unit(a::Vec)::Vec

Compute the unit vector colinear to the given one.
"""
function unit(a::Vec)::Vec
    a / norm(a)
end

"""
    round(a::Vec)::Vec

Compute the vector with it's elements rounder to nearest value.
"""
function Base.round(a::Vec)::Vec
    Vec(
        round(a.x),
        round(a.y),
        round(a.z)
    )
end

"""
    Quat - hypercomplex number.

Unit quaternions are used for rotation representation.

# Examples
```julia-repl
julia> # 120° counterclockwise rotation around the (1, 1, 1) axis
julia> Quat(2π/3, Vec(1, 1, 1))
Quat(0.5, 0.5, 0.5, 0.5)

julia> # 45° clockwise rotation around the (1, 0, 0) axis
julia> Quat(-π/4, Vec(1, 0, 0))
Quat(0.9238795325112867, -0.3826834323650898, -0.0, -0.0)
```
"""
struct Quat
    w::Real
    x::Real
    y::Real
    z::Real
    function Quat(w::Real=1, x::Real=0, y::Real=0, z::Real=0)::Quat
        new(w, x, y, z)
    end

    function Quat(angle::Real, axis::Vec)::Quat
        w = cos(angle / 2)
        xyz = unit(axis) * sin(angle / 2)
        Quat(w, xyz.x, xyz.y, xyz.z)
    end
end

function Base.:+(a::Quat, b::Quat)::Quat
    Quat(
        a.w + b.w,
        a.x + b.x,
        a.y + b.y,
        a.z + b.z
    )
end

function Base.:-(a::Quat, b::Quat)::Quat
    Quat(
        a.w - b.w,
        a.x - b.x,
        a.y - b.y,
        a.z - b.z
    )
end

function Base.:*(a::Quat, b::Quat)::Quat
    Quat(
        a.w * b.w - a.x * b.x - a.y * b.y - a.z * b.z,
        a.w * b.x + a.x * b.w + a.y * b.z - a.z * b.y,
        a.w * b.y - a.x * b.z + a.y * b.w + a.z * b.x,
        a.w * b.z + a.x * b.y - a.y * b.x + a.z * b.w
    )
end

function Base.:*(q::Quat, v::Vec)::Vec
    r = q * Quat(0, v.x, v.y, v.z) * conj(q)
    Vec(r.x, r.y, r.z)
end

function Base.:*(a::Quat, b::Real)::Quat
    Quat(
        a.w * b,
        a.x * b,
        a.y * b,
        a.z * b
    )
end

function Base.:*(a::Real, b::Quat)::Quat
    Quat(
        a * b.w,
        a * b.x,
        a * b.y,
        a * b.z
    )
end

function dot(a::Quat, b::Quat)::Real
    a.w * b.w + a.x * b.x + a.y * b.y + a.z * b.z
end

function Base.:/(a::Quat, b::Real)::Quat
    Quat(
        a.w / b,
        a.x / b,
        a.y / b,
        a.z / b
    )
end

"""
    conj(a::Quat)::Quat

Compute the quaternion conjugate.
If the quaternion is unit, it's conjugate represents the reverse rotation.

# Examples
```julia-repl
julia> q = Quat(π, Vec(1, 1, 1))
julia> q * conj(q)
Quat(1.0, 0.0, 0.0, 0.0)
```
"""
function Base.conj(a::Quat)::Quat
    Quat(
        a.w,
        -a.x,
        -a.y,
        -a.z
    )
end

"""
    norm(a::Quat)::Real

Compute the quaternion norm.
"""
function norm(a::Quat)::Real
    √(a.w^2 + a.x^2 + a.y^2 + a.z^2)
end

"""
    unit(a::Quat)::Quat

Compute the unit quaternion colinear to the given one.
"""
function unit(a::Quat)::Quat
    a / norm(a)
end

function slerp(a::Quat, b::Quat, t)::Quat
    d = dot(a, b)
    if d < 0
        b = -b
        d = -d
    end
    if d > 0.9995
        result = a + (a - b) * t
        return norm(result)
    end
    theta0 = acos(d)
    theta = theta0 * t
    sinTheta = sin(theta)
    sinTheta0 = sin(theta0)

    s1 = sinTheta / sinTheta0
    s0 = cos(theta) - d * s1
    return a * s0 + b * s1
end

"""
    Transf - spatial transformation representation.

The transformation consists of consecutive translation on vector `v` and rotation on `q` quaternion.
"""
struct Transf
    v::Vec
    q::Quat
    function Transf(v::Vec=Vec(), q::Quat=Quat())::Transf
        new(v, q)
    end
end

function Base.:+(a::Transf, b::Transf)::Transf
    Transf(
        a.v + a.q * b.v,
        a.q * b.q
    )
end

function Base.:+(a::Transf, b::Vec)::Vec
    a.v + a.q * b
end

"""
    conj(a::Transf)::Transf

Compute such transformation, that `a + conj(a)` equals to zero transformation.

# Examples
```julia-repl
julia> t = Transf(Vec(1, 2, 3), Quat(π/3, Vec(1, 1, 0)))
julia> t + conj(t)
Transf(Vec(0.0, 0.0, 0.0), Quat(1.0, 0.0, 0.0, 0.0))
```
"""
function Base.conj(a::Transf)::Transf
    q = conj(a.q)
    Transf(
        q * (-a.v),
        q
    )
end

end


"""
    poisson_eigenvalues(N, L, dim, ::Periodic)

Return the eigenvalues satisfying the discrete form of Poisson's equation
with periodic boundary conditions along the dimension `dim` with `N` grid
points and domain extent `L`.
"""
function poisson_eigenvalues(N, L, dim, ::Periodic)
    inds = reshape(1:N, reshaped_size(N, dim)...)
    return @. (2sin((inds - 1) * π / N) / (L / N))^2
end

"""
    poisson_eigenvalues(N, L, dim, ::Bounded)

Return the eigenvalues satisfying the discrete form of Poisson's equation
with staggered Neumann boundary conditions along the dimension `dim` with
`N` grid points and domain extent `L`.
"""
function poisson_eigenvalues(N, L, dim, ::Bounded)
    inds = reshape(1:N, reshaped_size(N, dim)...)
    return @. (2sin((inds - 1) * π / 2N) / (L / N))^2
end

"""
    poisson_eigenvalues(N, L, dim, ::Flat)

Return N-element array of `0.0` reshaped to three-dimensions.
This is also the first `poisson_eigenvalue` for `Bounded` and `Periodic` directions.
"""
poisson_eigenvalues(N, L, dim, ::Flat) = reshape(zeros(N), reshaped_size(N, dim)...)


input = 347991

# Part 1
function side_length(number:: Int)
  root = ceil(sqrt(number))
  if root % 2 == 0
    return root + 1
  else
    return root
  end
end

function get_position(number:: Int)
  slength = side_length(number)
  x = floor(slength / 2)

  side_centers = slength^2 - ((slength - 1) * collect(1:4)) - floor(slength / 2)
  y = minimum(abs.(side_centers - number))

  return x, y
end

(x, y) = abs.(get_position(input))
println("Part 1: $(x + y) steps from ($x, $y)")

# Part 2
mutable struct NegativeIndexSquareArray
  size::Int
  content::Array{Int, 3}

  NegativeIndexSquareArray(size) = new(size, zeros(size, size, 4))
end

function Base.getindex(array::NegativeIndexSquareArray, x, y)
  sx = sign(x) == 0? 1 : sign(x)
  sy = sign(y) == 0? 1 : sign(y)
  (x, y) = abs.([x, y])

  if sx == 1 && sy == 1
    array.content[x + 1, y + 1, 1]
  elseif sx == 1 && sy == -1
    array.content[x + 1, y, 2]
  elseif sx == -1 && sy == 1
    array.content[x, y + 1, 3]
  else # Both negative
    array.content[x, y, 4]
  end
end

function Base.setindex!(array::NegativeIndexSquareArray, value, x, y)
  sx = sign(x) == 0? 1 : sign(x)
  sy = sign(y) == 0? 1 : sign(y)
  (x, y) = abs.([x, y])

  if sx == 1 && sy == 1
    array.content[x + 1, y + 1, 1] = value
  elseif sx == 1 && sy == -1
    array.content[x + 1, y, 2] = value
  elseif sx == -1 && sy == 1
    array.content[x, y + 1, 3] = value
  else # Both negative
    array.content[x, y, 4] = value
  end
end

spiral = NegativeIndexSquareArray(10)
spiral[0, 0] = 1
(x, y) = [1, 0]
direction = im

while true
  spiral[x, y] = spiral[x - 1, y] + spiral[x - 1, y - 1] +
    spiral[x, y - 1] + spiral[x + 1, y - 1] + spiral[x + 1, y] +
    spiral[x + 1, y + 1] + spiral[x, y + 1] + spiral[x - 1, y + 1]

  if spiral[x, y] > input
    println("Part 2: $(spiral[x, y]) at ($x, $y)")
    break
  end

  left = direction * im
  if spiral[x + real(left), y + imag(left)] == 0
    direction = left
  end
  x += real(direction)
  y += imag(direction)
end


function bisect_left(A, p, r, v)
	i = p
	if p < r
	   q = floor(Int, (p + r) / 2)
	   if v <= A[q]
		   i = bisect_left(A, p, q, v)
	   else
		   i = bisect_left(A, q + 1, r, v)
	   end
	end
	return i
end

function bisect_right(A, p, r, v)
	i = p
	if p < r
	   q = floor(Int, (p + r) / 2)
	   if v < A[q]
		   i = bisect_right(A, p, q, v)
	   else
		   i = bisect_right(A, q + 1, r, v)
	   end
	end
	return i
end

function bisect_right(A, v)
	return bisect_right(A, 1, length(A) + 1, v)
end

function bisect_left(A, v)
	return bisect_left(A, 1, length(A) + 1, v)
end


# ---
# title: 312. Burst Balloons
# id: problem312
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Divide and Conquer, Dynamic Programming
# link: <https://leetcode.com/problems/burst-balloons/description/>
# hidden: true
# ---
#
# Given `n` balloons, indexed from `0` to `n-1`. Each balloon is painted with a
# number on it represented by array `nums`. You are asked to burst all the
# balloons. If the you burst balloon `i` you will get `nums[left] * nums[i] *
# nums[right]` coins. Here `left` and `right` are adjacent indices of `i`. After
# the burst, the `left` and `right` then becomes adjacent.
#
# Find the maximum coins you can collect by bursting the balloons wisely.
#
# **Note:**
#
#   * You may imagine `nums[-1] = nums[n] = 1`. They are not real therefore you can not burst them.
#   * 0 ≤ `n` ≤ 500, 0 ≤ `nums[i]` ≤ 100
#
# **Example:**
#
#
#
#     Input: [3,1,5,8]
#     Output: 167
#     Explanation:nums = [3,1,5,8] --> [3,5,8] -->   [3,8]   -->  [8]  --> []
#                  coins =  3*1*5      +  3*5*8    +  1*3*8      + 1*8*1   = 167
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


export newtoncg

"""
newtoncg(f,df,H,x)

Newton-CG method for solving min_x f(x)

"""
function newtoncg(f::Function,df::Function,H::Function,x::Vector;maxIter=20,atol=1e-8,out::Int=0,storeInterm::Bool=false,
	lineSearch::Function=(f,df,fk,dfk,xk,pk)->armijo(f,fk,dfk,xk,pk,maxIter=20),tolCG=1e-2,maxIterCG=30,P=d2f->identity)

    his = zeros(maxIter,6)
    X = (storeInterm) ? zeros(length(x),maxIter) : []
    i = 1; flag = -1; LL = []
    while i<=maxIter
        fk  = f(x)
        dfk = df(x)
        his[i,1:2] = [norm(fk) norm(dfk)]
        if storeInterm; X[:,i] = x; end;

        if(norm(dfk)<atol)
            flag = 0
            his = his[1:i,:]
            break
        end

        # get search direction
        d2f = H(x)
        PC  = P(d2f) # preconditioner
        pk,his[i,6],his[i,5],his[i,4] = cg(d2f,-dfk,out=-1,tol=tolCG,maxIter=maxIterCG,M=PC)
        if his[i,6]==-2 && norm(pk)==0
            pk = -dfk
        end
      # line search
        ak,his[i,3] = lineSearch(f,df,fk,dfk,x,pk)
        if his[i,3]==-1
            flag = -3
            his  = his[1:i,:]
            break
        end
      if out>0
            @printf "iter=%04d\t|f|=%1.2e\t|df|=%1.2e\tcgIter=%d\tLS=%1.4f\n" i his[i,1] his[i,2] his[i,4] ak
        end
        # update
        x  += ak*pk
        i+=1
    end
    i = min(maxIter,i)

    if out>=0
        if flag==-1
            println(@sprintf("newtoncg iterated maxIter (=%d) times but reached only atol of %1.2e instead of tol=%1.2e",i,his[i,2],atol))
        elseif flag==-3
            println(@sprintf("newtoncg stopped at iteration %d due to a line search fail.",i))
        elseif out>1
            println(@sprintf("newtoncg achieved desired atol of %1.2e at iteration %d.",atol,i))
        end
    end
    if storeInterm; X = X[:,1:i]; end
    return x,flag,his,X
end


# # Padua transform
# This demonstrates the Padua transform and inverse transform,
# explaining precisely the normalization and points

using FastTransforms

# We define the Padua points and extract Cartesian components:
N = 15
pts = paduapoints(N)
x = pts[:,1]
y = pts[:,2];

# We take the Padua transform of the function:
f = (x,y) -> exp(x + cos(y))
f̌ = paduatransform(f.(x , y));

# and use the coefficients to create an approximation to the function $f$:
f̃ = (x,y) -> begin
    j = 1
    ret = 0.0
    for n in 0:N, k in 0:n
        ret += f̌[j]*cos((n-k)*acos(x)) * cos(k*acos(y))
        j += 1
    end
    ret
end

# At a particular point, is the function well-approximated?
f̃(0.1,0.2) ≈ f(0.1,0.2)

# Does the inverse transform bring us back to the grid?
ipaduatransform(f̌) ≈ f̃.(x,y)


using LinearAlgebra
using Zygote: @adjoint

""" maximum eigenvalue of a symmetric square matrix """
function maxeig(symmat::Array{T,2})::T where {T}
    eig = eigen(Symmetric(symmat))
    max_idx = findmax(eig.values)[2]
    eigmax = eig.values[max_idx]
    return eigmax
end

@adjoint function maxeig(symmat::Array{T,2}) where {T}
    eig = eigen(Symmetric(symmat))
    max_idx = findmax(eig.values)[2]
    eigmax = eig.values[max_idx]
    vecmax = eig.vectors[:, max_idx]
    jac_adj = vecmax * (vecmax')
    return eigmax, c -> Tuple(jac_adj * c)
end

""" minimum eigenvalue of a symmetric square matrix """
function mineig(symmat::Array{T,2})::T where {T}
    eig = eigen(Symmetric(symmat))
    min_idx = findmin(eig.values)[2]
    eigmin = eig.values[min_idx]
    return eigmin
end

@adjoint function mineig(symmat::Array{T,2}) where {T}
    eig = eigen(Symmetric(symmat))
    min_idx = findmin(eig.values)[2]
    eigmin = eig.values[min_idx]
    vecmin = eig.vectors[:, min_idx]
    jac_adj = vecmin * (vecmin')
    return eigmax, c -> Tuple(jac_adj * c)
end

""" eigenvalue regularization """
function softplus_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return mean(softplus.(eig.values))
end

@adjoint function softplus_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return sum(softplus.(eig.values)),
    c -> Tuple(
        sum((
            sigmoid(eig.values[i]) * eig.vectors[:, i] * (eig.vectors[:, i]') for
            i = 1:size(symmat, 1)
        )) * c,
    )
end

function relu_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return sum(relu.(eig.values))
end

@adjoint function relu_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return sum(relu.(eig.values)),
    c -> Tuple(
        sum((
            0.5 * (sign(eig.values[i]) + 1) * eig.vectors[:, i] * (eig.vectors[:, i]')
            for i = 1:size(symmat, 1)
        )) * c,
    )
end

""" compute all eigenvalues """
function eigenvalues(symmat)
    eig = eigen(Symmetric(symmat))
    return eig.values
    # return eigvals(Symmetric(symmat))
end

@adjoint function eigenvalues(symmat)
    eig = eigen(Symmetric(symmat))
    return eig.values,
    c -> Tuple((eig.vectors[:, i] * (eig.vectors[:, i]')) * c for i = 1:size(symmat, 1))
    # vals = eigvals(Symmetric(symmat))
    # vecs = eigvecs(Symmetric(symmat))
    # return vals, c -> Tuple((vecs[:,i]*(vecs[:,i]'))*c for i=1:size(symmat,1))
end


include("prime.jl")
function  is_circular_prime(n::Int)
	if (!is_prime(n))
		return false
	end
	s=string(n)
	for i=1:length(s)
		if (!is_prime( parse(Int,s[i:length(s)]*s[1:i-1])))
			return false
		end

	end

	return true

end

function  compute()
	const N=1000000
	ans =sum( 1 for i =1:N if is_circular_prime(i))
	return ans

end

println(compute())


using LogExpFunctions, Test

@testset "xlogx & xlogy" begin
    @test iszero(xlogx(0))
    @test xlogx(2) ≈ 2.0 * log(2.0)

    @test iszero(xlogy(0, 1))
    @test xlogy(2, 3) ≈ 2.0 * log(3.0)
    @test xlogy(2, 3.0) ≈ 2.0 * log(3.0) # promoted
end

@testset "logistic & logit" begin
    @test logistic(2)        ≈ 1.0 / (1.0 + exp(-2.0))
    @test iszero(logit(0.5))
    @test logit(logistic(2)) ≈ 2.0
end

@testset "log1psq" begin
    @test iszero(log1psq(0.0))
    @test log1psq(1.0) ≈ log1p(1.0)
    @test log1psq(2.0) ≈ log1p(4.0)
    @test log1psq(Float16(0.1)) ≈ log1p(0.01)
end

# log1pexp, log1mexp, log2mexp & logexpm1

@testset "log1pexp" begin
    @test log1pexp(2.0)    ≈ log(1.0 + exp(2.0))
    @test log1pexp(-2.0)   ≈ log(1.0 + exp(-2.0))
    @test log1pexp(10000)  ≈ 10000.0
    @test log1pexp(-10000) ≈ 0.0

    @test log1pexp(2f0)      ≈ log(1f0 + exp(2f0))
    @test log1pexp(-2f0)     ≈ log(1f0 + exp(-2f0))
    @test log1pexp(10000f0)  ≈ 10000f0
    @test log1pexp(-10000f0) ≈ 0f0
end

@testset "log1mexp" begin
    @test log1mexp(-1.0)  ≈ log1p(- exp(-1.0))
    @test log1mexp(-10.0) ≈ log1p(- exp(-10.0))
end

@testset "log2mexp" begin
    @test log2mexp(0.0)  ≈ 0.0
    @test log2mexp(-1.0) ≈ log(2.0 - exp(-1.0))
end

@testset "logexpm1" begin
    @test logexpm1(2.0)            ≈  log(exp(2.0) - 1.0)
    @test logexpm1(log1pexp(2.0))  ≈  2.0
    @test logexpm1(log1pexp(-2.0)) ≈ -2.0

    @test logexpm1(2f0)            ≈  log(exp(2f0) - 1f0)
    @test logexpm1(log1pexp(2f0))  ≈  2f0
    @test logexpm1(log1pexp(-2f0)) ≈ -2f0
end

@testset "log1pmx" begin
    @test iszero(log1pmx(0.0))
    for x in [-0.65, -0.5, -0.3, -0.1, 0.4, 1.0, 2.0]
        z = BigFloat(x)
        @test log1pmx(x) ≈ Float64(log(1 + x) - x)
    end
end

@testset "logmxp1" begin
    @test iszero(logmxp1(1.0))
    for x in [0.1, 0.35, 0.5, 2.0, 3.0]
        z = BigFloat(x)
        @test logmxp1(x) ≈ Float64(log(z) - z + 1)
    end
end

@testset "logsumexp" begin
    @test logaddexp(2.0, 3.0)     ≈ log(exp(2.0) + exp(3.0))
    @test logaddexp(10002, 10003) ≈ 10000 + logaddexp(2.0, 3.0)
    @test logaddexp(2, 3.0)     ≈ log(exp(2.0) + exp(3.0)) # promotion

    @test logsumexp([1.0, 2.0, 3.0])          ≈ 3.40760596444438
    @test logsumexp((1.0, 2.0, 3.0))          ≈ 3.40760596444438
    @test logsumexp([1.0, 2.0, 3.0] .+ 1000.) ≈ 1003.40760596444438

    let cases = [([-Inf, -Inf], -Inf),   # correct handling of all -Inf
                 ([-Inf, -Inf32], -Inf), # promotion
                 ([-Inf32, -Inf32], -Inf32), # Float32
                 ([-Inf, Inf], Inf),
                 ([-Inf, 9.0], 9.0),
                 ([Inf, 9.0], Inf),
                 ([NaN, 9.0], NaN),  # NaN propagation
                 ([NaN, Inf], NaN),  # NaN propagation
                 ([NaN, -Inf], NaN), # NaN propagation
                 ([0, 0], log(2.0))] # non-float arguments
        for (arguments, result) in cases
            @test logaddexp(arguments...) ≡ result
            @test logsumexp(arguments) ≡ result
        end
    end
end

@testset "softmax" begin
    x = [1.0, 2.0, 3.0]
    r = exp.(x) ./ sum(exp.(x))
    @test softmax([1.0, 2.0, 3.0]) ≈ r
    softmax!(x)
    @test x ≈ r
end


### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ 1c630e40-628c-11ec-1793-05688dabf6c9
begin
	using Pkg; Pkg.activate(@__DIR__); Pkg.instantiate()
	Pkg.precompile()

	using PlutoUI
	using Statistics
	using StatsBase
	using RDatasets
	using Plots
	using StatsPlots
	using KernelDensity
	using Distributions
	using LinearAlgebra
	using HypothesisTests
	using PyCall
	using MLBase
end;

# ╔═╡ 7242a33b-ec37-48d6-afca-8c8387a45a83
PlutoUI.TableOfContents(aside=true, indent=true, depth=3)

# ╔═╡ dd2aa8b3-9da3-418d-bf31-74574483360d
html"""

"""

# ╔═╡ abc828f1-fd14-4329-8b47-f4135300f937
md"""
# 📊 Statistics

Having a solid understanding of statistics in data science allows us to understand our data better, and allows us to create a quantifiable evaluation of any future conclusions.

[YouTube Link](https://www.youtube.com/watch?v=AAGxWEJ_eWk)
"""

# ╔═╡ a8aae3d4-198c-482f-a972-05be31be39c1
md"""
## Load dataset

In this notebook, we will use eruption data on the faithful geyser (Yellowstone, USA). The data will contain wait times between every consecutive times the geyser goes off (in minutes) and the length of the eruptions (also in minutes).

Let's get the data first using the [RDatasets.jl](https://github.com/JuliaStats/RDatasets.jl) package...

**Note:** the dataset will be loaded as a `DataFrame`.
"""

# ╔═╡ 63420679-9b77-4412-ba1c-fc471da6b1e9
data = dataset("datasets", "faithful")

# ╔═╡ 9d56b2cc-c4c6-4ce2-8d9c-33174fe5bc33
md"""
We can use the `describe` function to get some stats from our data...
"""

# ╔═╡ 67ac4a1d-eb15-473f-a548-07425fd627d8
describe(data)

# ╔═╡ 5ee886cc-2730-40cf-8497-507dcdc08c5f
md"""
Let's call each column (feature) separately...
"""

# ╔═╡ 110fda28-1b77-44d6-8485-a4ec6e012011
begin
	eruption = data[!, :Eruptions]
	# eruption = data.Eruptions

	waittime = data[!, :Waiting]
	# waittime = data.Waiting
end;

# ╔═╡ 6e093fbb-1e3a-4131-b3c2-b5891e734fa7
md"""
Now, we can plot both variables in the same scatter plot...
"""

# ╔═╡ ed9da9c2-53e4-45d6-9ec1-02c8bf54c1a5
begin
	scatter(eruption, label="Eruptions")
	scatter!(waittime, label="Wait Time", xlabel="Index", ylabel="Time (minutes)")
end

# ╔═╡ 13137e30-cfde-459a-adb4-1ddc07e626e8
md"""
## Statistics plots

As you can see, this doesn't tell us much about the data... Let's try some statistical plots.
"""

# ╔═╡ ff9d63e7-808a-4465-a68d-659a222afbae
md"""
### Boxplot
"""

# ╔═╡ c3672771-248a-47f9-a019-2de03cdd79f4
boxplot(["Eruption Length"], eruption, legend=false,
		size=(200,400), whisker_width=1, ylabel="Time (minutes)")

# ╔═╡ 46b55947-428a-483f-a00f-a232e5329fbd
md"""
Statistical plots such as a boxplot (and a violin plot), can provide a much better understanding of the data. Here, we immediately see that the median time of each eruption is about 4 minutes.
"""

# ╔═╡ e0933f5c-50b5-47cf-bae8-7a45a219fefb
md"""
### Histogram

The next plot we will see is a histogram plot. Check its arguments below...

- `x`: AbstractVector of values to be binned.
- `bins`: Integer, NTuple{2,Integer}, AbstractVector or Symbol. Default is :auto (the Freedman-Diaconis rule). For histogram-types, defines the approximate number of bins to aim for, or the auto-binning algorithm to use (`:sturges`, `:sqrt`, `:rice`, `:scott` or `:fd`). For fine-grained control pass a Vector of break values, e.g. `range(minimum(x), stop = maximum(x), length = 25)`.
- `weights`: Vector of weights for the values in x, for weighted bin counts.
- `normalize`: Bool or Symbol. Histogram normalization mode. Possible values are: `false`/`:none` (no normalization, default), `true`/`:pdf` (normalize to a discrete Probability Density Function, where the total area of the bins is 1), `:probability` (bin heights sum to 1) and `:density` (the area of each bin, rather than the height, is equal to the counts - useful for uneven bin sizes).
- `bar_position`: Symbol. Choose from `:overlay` (default), `:stack`. (warning: May not be implemented fully).
- `bar_width`: nothing or Number. Width of bars in data coordinates. When nothing, chooses based on x (or y when orientation = `:h`).
- `bar_edges`: Bool. Align bars to edges (true), or centers (the default).
- `orientation`: Symbol. Horizontal or vertical orientation for bar types. Values `:h`, `:hor`, `:horizontal` correspond to horizontal (sideways, anchored to y-axis), and `:v`, `:vert`, and `:vertical` correspond to vertical (the default).

Let's plot eruption length data...
"""

# ╔═╡ 943c8862-b909-4750-96da-7bda62c5fc8c
histogram(eruption, bins=:sqrt, label="Eruption", color=:orange,
		  xlabel="Eruption Length (minutes)", ylabel="Frequency")

# ╔═╡ ed480ae6-3302-4d0e-b83e-53a9c395d5ce
md"""
### Kernel density estimates

Next, we will see how we can fit a kernel density estimation function to our data. We will make use of the [KernelDensity.jl](https://github.com/JuliaStats/KernelDensity.jl) package.

With Kernel density estimates, we can see how things are changing...
"""

# ╔═╡ ed58e41f-3f73-40c1-92ec-5f2845c62b97
ρ = kde(eruption)

# ╔═╡ 952506dc-6cd8-4ecd-ace8-b9d914560289
md"""
If we want the histogram and the kernel density graph to be aligned we need to remember that the "density contribution" of every point added to one of these histograms is `density * nb of elements * bin width`. We do this, since KDE are densities, while the y-axis is showing the frequency.

**Note:** Read more about kernel density estimates [here](https://en.wikipedia.org/wiki/Kernel_density_estimation).
"""

# ╔═╡ 3ea44987-7ba7-4df6-a82a-c11a98ea784d
# the range of density data (x-axis)
ρ.x

# ╔═╡ 31274ae3-20da-4f52-9153-6f6b4d4472cb
# density data (y-axis)
ρ.density

# ╔═╡ 53f3e80a-2c88-405f-9131-21d07b132389
# density * nb of elements * bin width
scaled_density = ρ.density .* length(eruption) .* 0.2

# ╔═╡ ccfe5310-cc85-457c-bfbc-9671062b9088
begin
	histogram(eruption, bins=:sqrt, label="Eruption", color=:lightblue,
		  	  xlabel="Eruption Length (minutes)", ylabel="Frequency")

	plot!(ρ.x, scaled_density, linewidth=1.5, color=:black, label="KDE fit")
end

# ╔═╡ cbf471e5-523c-4550-835b-5c45b9648d8c
md"""
## Probability distributions

First, we will take a look at one probablity distribution, namely the Normal Distribution and verify that it generates a bell curve.
"""

# ╔═╡ 914412a4-69d4-4d9f-b575-03cfa03a7873
# generating random normal data
normaldata = randn(100_000)

# ╔═╡ 86c9cf10-2ade-45de-8649-a8dd5b2ae586
begin
	# generating the KDE
	ρₙ = kde(normaldata)

	# scaling density estimates
	scaled_ρₙ = ρₙ.density .* length(normaldata) * 0.1
end

# ╔═╡ e137810d-9e3c-43c8-905c-4efc2213a130
begin
	# plotting histogram
	histogram(normaldata, color=:lightblue, label="Z ~ N(0,1)",
			  xlabel="Z", ylabel="Frequency")

	# plotting KDE
	plot!(ρₙ.x, scaled_ρₙ, linewidth=3, color=:black, label="KDE fit")
end

# ╔═╡ f6430d83-a92b-4c03-a07e-8a6146cc9e1c
md"""
Another way to generate the same plot is via using the [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) package and choosing the probability distribution you want, and then drawing random numbers from it.
"""

# ╔═╡ 4ffd9a7a-c62c-40a4-8ec1-4e2351a2a66a
md"""
### Normal distribution

As an example, we will use `N = Normal()` below to draw random numbers from a [Normal distribution](https://en.wikipedia.org/wiki/Normal_distribution). The default $\mu$ and $\sigma$ values are 0 and 1 (standard Normal distribution).
"""

# ╔═╡ 978e2f4c-b411-459c-8024-4f6205484f63
begin
	# drawing numbers ~N(0,1)
	N = Normal()
	gaussdata = rand(N, 100_000)

	# calculating KDE
	rhoN = kde(gaussdata)
	scaled_rhoN = rhoN.density .* length(gaussdata) .* 0.1
end

# ╔═╡ d1483a5e-0b7e-4317-adb0-8f024fad21eb
begin
	# plotting histogram
	histogram(gaussdata, color=:lightblue, label="Z ~ N(0,1)",
			  xlabel="Z", ylabel="Frequency")

	# plotting KDE
	plot!(rhoN.x, scaled_rhoN, linewidth=3, color=:black, label="KDE fit")
end

# ╔═╡ 62470113-95dc-4b92-9e1d-3118efa0bd08
md"""
### Binomial distribution

Now, let's draw random numbers from a [Binomial distribution](https://en.wikipedia.org/wiki/Binomial_distribution).
"""

# ╔═╡ 7749b6b5-cf01-4812-b1ce-6fe3204b0bcf
begin
	# drawing numbers ~B with n=40 trials
	B = Binomial(40)
	binomdata = rand(B, 100_000)

	# calculating KDE
	rhoB = kde(binomdata)
	scaled_rhoB = rhoB.density .* length(binomdata) .* 0.5
end

# ╔═╡ ff68a550-9cc6-416e-96b8-07de164e5fc4
begin
	# plotting histogram
	histogram(binomdata, color=:lightblue, label="Z ~ B(40,0.5)",
			  xlabel="Z", ylabel="Frequency")

	# plotting KDE
	plot!(rhoB.x, scaled_rhoB, linewidth=3, color=:black, label="KDE fit")
end

# ╔═╡ 693de4f8-88b1-476f-9319-64edafb049af
md"""
### Fit data to a distribution

Next, we will try to fit a given set of numbers to a parametric distribution. Thus, we are going to use the StatsBase.jl `fit` function.
"""

# ╔═╡ 7f9949d9-22e2-4d4e-92ec-6e52e01d10aa
begin
	# creating 1000 random values
	randomdata = rand(1000)
	# fitting random values to a Standard Normal distribution
	fitted_distrib = fit(Normal, randomdata)
	# drawing 1000 values from the new fitted distribution
	new_data = rand(fitted_distrib, 1000)
end

# ╔═╡ 98301aae-61ca-46ab-936a-fbb310854778
# note that `fitted_distrib` is a `Distribution`
fitted_distrib

# ╔═╡ f270c785-3a7a-4e43-9f21-b9a6e199383b
begin
	# plotting new data (Normal distribution)
	histogram(new_data, nbins=20, fillalpha=0.3, label="Normal Data")

	# plotting original data
	histogram!(randomdata, nbins=20, linecolor=:red, fillalpha=0.3,
			   label="Original Data")
end

# ╔═╡ 2d5e5740-a3c6-4448-b79f-f231a0ac9d96
md"""
We can try to fit eruptions data to a Normal distribution, although eruptions distribution does not seem to be Normal.
"""

# ╔═╡ ef651442-e846-4a5e-9f3b-0da63afeca16
begin
	erup_distrib = fit(Normal, eruption)
	normal_erup = rand(erup_distrib, 1000)
end

# ╔═╡ a5a54536-4b60-4a9c-b423-2804cce37382
erup_distrib

# ╔═╡ 25b15bda-9bb4-4ed9-90e4-c1f937051314
begin
	# plotting fitted data (Normal distribution)
	histogram(normal_erup, nbins=20, fillalpha=0.3, label="Normal Data")

	# plotting original data
	histogram!(eruption, nbins=20, linecolor=:red, fillalpha=0.3,
			   label="Original Data")
end

# ╔═╡ b82db811-d97d-4efd-81bc-a48780b23dcb
md"""
## Hypothesis testing

Next, we will perform hypothesis testing using the [HypothesisTests.jl](https://github.com/JuliaStats/HypothesisTests.jl) package.

Let's try to use the `OneSampleTTest` function to perform a one sample t-test of random sampled values.

The one-sample t-test is a statistical hypothesis test used to determine whether an unknown population mean is different from a specific value. In this case, the default value under $H_0$ is 0.
"""

# ╔═╡ af54608d-209f-4d16-9573-b4b8634d2acf
# generating 1000 random normal values
hypdata = randn(1000)

# ╔═╡ 867dab59-1c8d-44d9-9add-9f3999f1b86c
# performing one sample t-test
OneSampleTTest(hypdata)

# ╔═╡ 47c7f6df-0696-4b1e-b6c2-7eb8e782eb82
md"""
**Note:** as we can see above, we cannot reject the null hypothesis $H_0$. This was expected, since we drew sample values from a Standard Normal distribution, with $\mu=0$.

Now we can perform a one sample t-test of eruption data. We will probably be able to reject $H_0$, since the eruption length mean is close to 3.5 minutes...
"""

# ╔═╡ 6555c2a2-227f-419b-9b8b-d563767f3ea5
OneSampleTTest(eruption)

# ╔═╡ 78585d9d-5696-41dd-80f0-abab1585b008
md"""
## Correlation statistics

First of all, let's see visually if there is a correlation between `eruption` and `waittime` variables using a scatter plot...
"""

# ╔═╡ 650a1af5-b413-4794-bf51-77fa34cfb820
scatter(eruption, waittime, legend=false, color=:orange,
		xlabel="Eruption Length (minutes)", markersize=3,
		ylabel = "Time between Eruptions (minutes)")

# ╔═╡ 80db913d-faa8-43d5-9794-e17b02823cc5
md"""
These two variables may have a strong positive correlation. We can calculate both `Pearson` and `Spearman` correlation coefficients to check that.
"""

# ╔═╡ fcb84ac7-3bca-4284-8939-a55bd78a02fa
begin
	spearman_r = corspearman(eruption, waittime)
	"r(Spearman) = $spearman_r"
end

# ╔═╡ 960b8739-10d1-490d-ac49-e51527386161
begin
	pearson_r = cor(eruption, waittime)
	"r(Pearson) = $pearson_r"
end

# ╔═╡ 5bd50a69-4edb-435a-9263-e0b5d47851aa
md"""
Interesting! This means that the next time you visit Yellowstone National part ot see the faithful geysser and you have to wait for too long for it to go off, you will likely get a longer eruption!

Currently we are using the p-value of Spearman and Pearson correlation coefficients from Python. But you can follow the formula [here](https://stackoverflow.com/questions/53345724/how-to-use-julia-to-compute-the-pearson-correlation-coefficient-with-p-value) to implement your own.

Let's import the `scipy.stats` Python module.

**Note:** sometimes there are some issues getting Python and Julia to communicate as desired. We can explicity add `scipy.stats` using the [Conda.jl](https://github.com/JuliaPy/Conda.jl) package:

```julia
julia> using Conda
julia> Conda.add("scipy")
```
"""

# ╔═╡ d1ef6cd2-fda1-47f6-af12-8492ea70a539
scipy_stats = pyimport("scipy.stats");

# ╔═╡ c5f77715-02cb-437a-a45f-3f3a4dfba67b
begin
	rₛ, pvalueₛ = scipy_stats.spearmanr(eruption, waittime)
	"r(Spearman) = $(rₛ) || p-value = $(pvalueₛ)"
end

# ╔═╡ 615e7c40-bafe-475f-959d-29c5f5e845e9
begin
	rₚ, pvalueₚ = scipy_stats.pearsonr(eruption, waittime)
	"r(Pearson) = $(rₚ) || p-value = $(pvalueₚ)"
end

# ╔═╡ 32bc5a24-638e-43bc-a5c1-bd0a85db7849
md"""
## Confusion matrix & AUC

Finally, we will cover basic tools you will need such as confusion matrix and AUC scores. We use the [MLBase.jl](https://github.com/JuliaStats/MLBase.jl) package for that.

First, let's calculate a confusion matrix, given a target vector `y₁` and a prediction vector `ŷ₁`...
"""

# ╔═╡ 838a8a32-95f6-48c0-b87d-28a66ef431c5
begin
	# generating target and prediction vectors
	y₁ = [1, 1, 1, 1, 1, 1, 1, 2]
	yhat₁ = [1, 1, 2, 2, 1, 1, 1, 1]
end;

# ╔═╡ 68d7cd6c-bd76-406b-a3d7-4528ee9671fd
# compute default confusion matrix
C = confusmat(2, y₁, yhat₁)

# ╔═╡ 2bcf8647-368d-4b74-8e1a-4a4b1a577fdc
md"""
We may normalize cell values per class...
"""

# ╔═╡ 86ad3850-f631-42b2-8d19-bf20cba89c6b
C ./ sum(C, dims=2)

# ╔═╡ 38a2bb94-1d8f-4e1a-a07f-3bc1332301fd
md"""
Also, we may compute the accuracy (correct rate) from confusion matrix...

```math
accuracy = \frac{TP}{TP+TN+FP+FN}
```
"""

# ╔═╡ ea2c50c6-9404-4ba4-8fe7-4d1e1b3246ff
acc = sum(diag(C)) / length(y₁)

# ╔═╡ 70cbc1f9-171e-4664-87d1-6e73f0bc841c
md"""
Instead of calculating accuracy manually, we could use the `correctrate` function...
"""

# ╔═╡ 188c0219-8a54-4f28-95fd-9f2f525c12be
acc == correctrate(y₁, yhat₁)

# ╔═╡ 063d30ad-0ac4-4d50-b66b-9b06bf155fd8
md"""
Now, we can use the `roc` function to calculate ROC scores...
"""

# ╔═╡ f3cca2ef-17ec-44e0-996d-8a0cb2f2fcb1
begin
	# generating target and prediction vectors
	y₂ = [1, 1, 1, 1, 1, 1, 1, 0]
	yhat₂ = [1, 1, 0, 0, 1, 1, 1, 1]
end;

# ╔═╡ 9af05d0f-5daf-48f4-9949-79db70ae4bbc
ROC = MLBase.roc(y₂, yhat₂)

# ╔═╡ 92f98fe2-9154-4029-8282-45dee3f11d31
md"""
We may also use both `recall` and `precision` functions to compute these metrics...
"""

# ╔═╡ 0cfa5643-b65b-4374-a227-8f036ce32bbf
begin
	rec = recall(ROC)
	"Recall = $rec"
end

# ╔═╡ 93ea2a29-bac3-4629-bac5-5d03e8e45d80
begin
	prec = precision(ROC)
	"Precision = $prec"
end

# ╔═╡ Cell order:
# ╟─1c630e40-628c-11ec-1793-05688dabf6c9
# ╟─7242a33b-ec37-48d6-afca-8c8387a45a83
# ╟─dd2aa8b3-9da3-418d-bf31-74574483360d
# ╟─abc828f1-fd14-4329-8b47-f4135300f937
# ╟─a8aae3d4-198c-482f-a972-05be31be39c1
# ╠═63420679-9b77-4412-ba1c-fc471da6b1e9
# ╟─9d56b2cc-c4c6-4ce2-8d9c-33174fe5bc33
# ╠═67ac4a1d-eb15-473f-a548-07425fd627d8
# ╟─5ee886cc-2730-40cf-8497-507dcdc08c5f
# ╠═110fda28-1b77-44d6-8485-a4ec6e012011
# ╟─6e093fbb-1e3a-4131-b3c2-b5891e734fa7
# ╠═ed9da9c2-53e4-45d6-9ec1-02c8bf54c1a5
# ╟─13137e30-cfde-459a-adb4-1ddc07e626e8
# ╟─ff9d63e7-808a-4465-a68d-659a222afbae
# ╠═c3672771-248a-47f9-a019-2de03cdd79f4
# ╟─46b55947-428a-483f-a00f-a232e5329fbd
# ╟─e0933f5c-50b5-47cf-bae8-7a45a219fefb
# ╠═943c8862-b909-4750-96da-7bda62c5fc8c
# ╟─ed480ae6-3302-4d0e-b83e-53a9c395d5ce
# ╠═ed58e41f-3f73-40c1-92ec-5f2845c62b97
# ╟─952506dc-6cd8-4ecd-ace8-b9d914560289
# ╠═3ea44987-7ba7-4df6-a82a-c11a98ea784d
# ╠═31274ae3-20da-4f52-9153-6f6b4d4472cb
# ╠═53f3e80a-2c88-405f-9131-21d07b132389
# ╠═ccfe5310-cc85-457c-bfbc-9671062b9088
# ╟─cbf471e5-523c-4550-835b-5c45b9648d8c
# ╠═914412a4-69d4-4d9f-b575-03cfa03a7873
# ╠═86c9cf10-2ade-45de-8649-a8dd5b2ae586
# ╠═e137810d-9e3c-43c8-905c-4efc2213a130
# ╟─f6430d83-a92b-4c03-a07e-8a6146cc9e1c
# ╟─4ffd9a7a-c62c-40a4-8ec1-4e2351a2a66a
# ╠═978e2f4c-b411-459c-8024-4f6205484f63
# ╠═d1483a5e-0b7e-4317-adb0-8f024fad21eb
# ╟─62470113-95dc-4b92-9e1d-3118efa0bd08
# ╠═7749b6b5-cf01-4812-b1ce-6fe3204b0bcf
# ╠═ff68a550-9cc6-416e-96b8-07de164e5fc4
# ╟─693de4f8-88b1-476f-9319-64edafb049af
# ╠═7f9949d9-22e2-4d4e-92ec-6e52e01d10aa
# ╠═98301aae-61ca-46ab-936a-fbb310854778
# ╠═f270c785-3a7a-4e43-9f21-b9a6e199383b
# ╟─2d5e5740-a3c6-4448-b79f-f231a0ac9d96
# ╠═ef651442-e846-4a5e-9f3b-0da63afeca16
# ╠═a5a54536-4b60-4a9c-b423-2804cce37382
# ╠═25b15bda-9bb4-4ed9-90e4-c1f937051314
# ╟─b82db811-d97d-4efd-81bc-a48780b23dcb
# ╠═af54608d-209f-4d16-9573-b4b8634d2acf
# ╠═867dab59-1c8d-44d9-9add-9f3999f1b86c
# ╟─47c7f6df-0696-4b1e-b6c2-7eb8e782eb82
# ╠═6555c2a2-227f-419b-9b8b-d563767f3ea5
# ╟─78585d9d-5696-41dd-80f0-abab1585b008
# ╠═650a1af5-b413-4794-bf51-77fa34cfb820
# ╟─80db913d-faa8-43d5-9794-e17b02823cc5
# ╠═fcb84ac7-3bca-4284-8939-a55bd78a02fa
# ╠═960b8739-10d1-490d-ac49-e51527386161
# ╟─5bd50a69-4edb-435a-9263-e0b5d47851aa
# ╠═d1ef6cd2-fda1-47f6-af12-8492ea70a539
# ╠═c5f77715-02cb-437a-a45f-3f3a4dfba67b
# ╠═615e7c40-bafe-475f-959d-29c5f5e845e9
# ╟─32bc5a24-638e-43bc-a5c1-bd0a85db7849
# ╠═838a8a32-95f6-48c0-b87d-28a66ef431c5
# ╠═68d7cd6c-bd76-406b-a3d7-4528ee9671fd
# ╟─2bcf8647-368d-4b74-8e1a-4a4b1a577fdc
# ╠═86ad3850-f631-42b2-8d19-bf20cba89c6b
# ╟─38a2bb94-1d8f-4e1a-a07f-3bc1332301fd
# ╠═ea2c50c6-9404-4ba4-8fe7-4d1e1b3246ff
# ╟─70cbc1f9-171e-4664-87d1-6e73f0bc841c
# ╠═188c0219-8a54-4f28-95fd-9f2f525c12be
# ╟─063d30ad-0ac4-4d50-b66b-9b06bf155fd8
# ╠═f3cca2ef-17ec-44e0-996d-8a0cb2f2fcb1
# ╠═9af05d0f-5daf-48f4-9949-79db70ae4bbc
# ╟─92f98fe2-9154-4029-8282-45dee3f11d31
# ╠═0cfa5643-b65b-4374-a227-8f036ce32bbf
# ╠═93ea2a29-bac3-4629-bac5-5d03e8e45d80