using CalculusWithJulia
using Test

## test package
@testset "test packages" begin

    ## Roots
    @test fzero(sin, 3, 4)  ≈ pi
    @test fzero(sin, 3.0)  ≈ pi

    ## ForwardDiff
    f(x) = sin(x)
    @test f'(2)  ≈ cos(2)
    @test f''(2)  ≈ -sin(2)

    ## QuadGK
    @test quadgk(sin, 0, pi)[1]  ≈ 2
    @test abs(quadgk(x -> x^2, 0, 1)[2]) <= 1e-14

    ## HQuadrature
    @test hquadrature(x -> 2, (0.0, 0.0), (2pi, 1pi))[1]  ≈ 2 * (2pi) * (1pi)

    ## SymPy
    @test limit(x -> sin(x)/x, 0) == 1
    @test integrate(sin,  0, pi) == 2

end

@testset "multidimensional" begin

    x = [[1,2,3], [4,5,6]]
    @test unzip(x)[1] == [1, 4]
    @test unzip(x)[2] == [2, 5]
    @test unzip(x)[3] == [3, 6]

    xs = [1,2,3]
    ys = [1,2]
    M = ((x,y) -> (x,y)).(xs, ys')
    @test unzip(M, recursive=true)  == ([[1,2,3], [1,2,3]], [[1,1,1], [2,2,2]])

    @vars x y z real=true

    f(x,y,z) = x*y*z
    f(v) = f(v...)
    @test gradient(f(x,y,z), [x,y,z]) == [y*z, x*z, x*y]
    @test ∇(f)([1,2,3]) ≈  [2*3, 1*3, 1*2]

    F(x,y,z) = [x,y,z]
    F(v) = F(v...)
    @test divergence(F(x,y,z), [x,y,z]) == 3
    @test  (∇⋅F)([1,0,0]) ≈ 3

    F(x,y,z) = [-y, x, 0]
    F(v) = F(v...)
    @test curl(F(x,y,z), [x,y,z]) == [0,0,2]
    @test all((∇ × F)([1,2,3]) .≈ [0.0, 0.0, 2.0])

end

@testset "derivatives" begin

    f(x) = x^2
    @test secant(f, 0, 1)(1/2) ≈ 1/2
    @test tangent(f, 1/2)(1) ≈ f(1/2) + 2*(1/2)*(1-1/2)

    @test f'(1) ≈ 2
    @test f''(1)  ≈ 2
    @test iszero(f'''(1))

    r(t) = [t, t^2, t^3]
    @test all(r'(1) .≈ [1.0, 2*1, 3*1^2])
end

@testset "integration" begin

    @test quadgk(sin, 0, pi)[1]  ≈ 2
    @test hcubature(xy -> sin(xy[1])*sin(xy[2]), (0,0), (pi,pi))[1] ≈ 4
    @test riemann(sin, 0, pi, 10_000)  ≈ 2
    @test fubini((x,y) -> 1, (x->-sqrt(1-x^2), x->sqrt(1-x^2)), (-1,1)) ≈ pi

end

@testset "plot-utils" begin

    if isinteractive()

        pyplot()
        f(x) = 1/x

        plotif(f, f, -1, 1)
        plotif(f, f', -1, 1)
        trimplot(f, -1, 1)
        signchart(f, -1, 1)

        rr(theta) = cos(theta)
        plot_polar(rr, 0, 2pi)

        rr(t) = [sin(t), cos(t)]
        plot_parametric_curve(rr, 0, 2pi)
        arrow!(rr(2pi), rr'(2pi))

        rr(t) = [sin(t), cos(t), t]
        plot_parametric_curve(rr, 0, 2pi)
        arrow!(rr(2pi), rr'(2pi))

        F(u, v) = [u*cos(v), u*sin(v)]
        us = range(0, 1, length=20)
        vs = range(0, 2pi, length=20)
        plot(legend=false)
        plot!(unzip(F.(us, vs'))...) # constant angle, draws rays
        plot!(unzip(F.(us', vs))...) # constant radii, draws circles

        pyplot() # surface doesn't work with gr()
        Phi(u,v) = [u*cos(v), u*sin(v), u]
        plot_parametric_surface(Phi, xlims=(0, 1), ylims=(0,2pi))

        V(x,y) = [x, x-y]
        vectorfieldplot(V, xlims=(-2,2))#, nx=8, ny=8) # ylim=(-5,5) plotly not happy, gr is

        V(x,y,z) = [-y, x, 0]
        vectorfieldplot3d(V, nx=4, ny=4, nz=3)

        a,b = 1, 3
        F(x,y,z) =  (x^2+((1+b)*y)^2+z^2-1)^3-x^2*z^3-a*y^2*z^3
        CalculusWithJulia.plot_implicit_surface(F, xlims=(-2,2),ylims=(-1,1),zlims=(-1,2))
    end

end


@testset "material" begin

@testset "tensorize" begin
    @test MaxwellBase.tensorize(3) === SMatrix{3,3,ComplexF64,9}(3,0,0, 0,3,0, 0,0,3)
    @test MaxwellBase.tensorize([2,4,6]) === SMatrix{3,3,ComplexF64,9}(2,0,0, 0,4,0, 0,0,6)
    @test MaxwellBase.tensorize([1 4 7; 2 5 8; 3 6 9]) === SMatrix{3,3,ComplexF64,9}(1,2,3,4,5,6,7,8,9)
end

@testset "Material" begin
    m_nn = Material("(num,num)", ε = 2, μ = 3)
    m_vv = Material("(vec,vec)", ε = [2.1,2.2,2.3], μ = [3.1,3.2,3.3])
    m_mm = Material("(mat,mat)", ε = [2 2 2; 2 2 2; 2 2 2], μ = [3 3 3; 3 3 3; 3 3 3])

    @test matparam(m_nn,EE) == diagm(0=>[2,2,2]) && matparam(m_nn,HH) == diagm(0=>[3,3,3]) && string(m_nn) == "(num,num)"
    @test matparam(m_vv,EE) == diagm(0=>[2.1,2.2,2.3]) && matparam(m_vv,HH) == diagm(0=>[3.1,3.2,3.3]) && string(m_vv) == "(vec,vec)"
    @test matparam(m_mm,EE) == fill(2, (3,3)) && matparam(m_mm,HH) == fill(3, (3,3)) && string(m_mm) == "(mat,mat)"
end

@testset "τ_trans" begin
    # Check if τ_trans for a 2×2 matrix is τ_trans for the 3×3 version of the matrix.  The
    # 3×3 version needs to have the original 2×2 matrix as the top-left 2×2 block, and the
    # third row and column has the (3,3) entry as the only nonzero entry.
    ε2 = @SMatrix rand(ComplexF64, 2, 2)
    ε3temp = zeros(ComplexF64, 3, 3)
    ε3temp[1:2,1:2] .= ε2
    ε3temp[3,3] = rand(ComplexF64)
    ε3 = MaxwellBase.S²ComplexF3(ε3temp)

    τ2 = MaxwellBase.τ_trans(ε2)
    τ3 = MaxwellBase.τ_trans(ε3)
    @test τ3[1:2,1:2] == τ2
end

end  # @testset "material"


# Example 2.5 from A. Law (2013): Simulation Modeling and Analysis, 5ed

using EventSimulation
using Distributions   # Exponential distribution
using StatsBase       # weighted mean
using DataFrames      # nicer ouptut of results

mutable struct Job
    tin::Float64  # time when job entered the system
    tout::Float64 # time when job left the system
    ttrue::Float64 # time needed without cycling
    remaining::Float64 # time remaining to finish the job
end

mutable struct SharedComputer <: AbstractState
    a::Exponential{Float64}    # arrival time distribution
    p::Exponential{Float64}    # processing time distribution
    τ::Float64                 # switching time
    q::Float64                 # processing quantum time
    n::Int                     # number of terminals
    reps::Int                  # number of finished jobs to collect
    job_queue::Vector{Job}     # queue of jobs waiting for processing
    finished_jobs::Vector{Job} # list of finished jobs
    job::Union{Job, Nothing}  # current job on computer
    # this structure holds information for reporting
    # for each perior between events it holds the length of queue,
    # if the computer were busy and the length of time period
    state_chunks::Vector{Tuple{Int, Bool, Float64}}
    SharedComputer(a, p, τ, q, n, reps) =
        new(Exponential(a), Exponential(p), τ, q, n, reps,
            Job[], Job[], nothing, Tuple{Int, Bool, Float64}[])
end

function arrival(s)
    ttrue = rand(s.state.p)
    j = Job(s.now, NaN, ttrue, ttrue)
    isnothing(s.state.job) && register!(s, start_compute, 0.0)
    push!(s.state.job_queue, j)
end

function start_compute(s)
    j = popfirst!(s.state.job_queue)
    s.state.job = j
    processed = j.remaining < s.state.q ? j.remaining : s.state.q
    j.remaining -= processed
    register!(s, stop_compute, processed + s.state.τ)
end

function stop_compute(s)
    j = s.state.job
    s.state.job = nothing
    if j.remaining == 0
        j.tout = s.now
        push!(s.state.finished_jobs, j)
        register!(s, arrival, rand(s.state.a))
    else
        push!(s.state.job_queue, j)
    end
    isempty(s.state.job_queue) || register!(s, start_compute, 0.0)
    length(s.state.finished_jobs) == s.state.reps && terminate!(s)
end

monitor(s, Δ) =
    push!(s.state.state_chunks,
          (length(s.state.job_queue), !isnothing(s.state.job), Δ))

function run(a, p, τ, q, n, reps)
    sc = SharedComputer(a, p, τ, q, n, reps)
    s = Scheduler(sc, Float64, monitor)
    for _ in 1:n
        register!(s, arrival, rand(sc.a))
    end
    go!(s)
    chunks = s.state.state_chunks
    (n=n,
     response=mean(j.tout - j.tin for j in s.state.finished_jobs),
     queue=mean(getindex.(chunks, 1), Weights(getindex.(chunks, 3))),
     CPU=mean(getindex.(chunks, 2), Weights(getindex.(chunks, 3))))
end

DataFrame([run(25.0, 0.8, 0.015, 0.1, n, 1000) for n in 10:10:80])


"""
    logtp_equal_rates_equal_time(i, j)

Evaluate the logarithm of a transition probability of a simple birth and death
process when the death rate `μ` is equal to the birth rate `λ` and the time `t`
is equal to `1 / λ`. It is a very unlikely event (let's say impossible) but it
simplifies into a special case.
"""
@inline function logtp_equal_rates_equal_time(
  i::I,
  j::I
) where {
  I <: Integer
}
  lgamma(i + j) - lgamma(i) - lgamma(j + 1) - (i + j) * log(2)
end

"""
    logtp_equal_rates(i, j, t, λ)

Evaluate the logarithm of a transition probability of a simple birth and death
process.
"""
@inline function logtp_equal_rates(
  i::I,
  j::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = λ * t

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + (i + j) * log(θ / (1 + θ)) +
  log_hypergeometric(a, b, logexpm1(-2 * log(θ)))
end

"""
    logtp_equal_rates_alternating(i, j, t, λ)

Evaluate the logarithm of a transition probability of a simple birth and death
process when terms are alternating in sign.
"""
@inline function logtp_equal_rates_alternating(
  i::I,
  j::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = λ * t

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + (i + j) * log(θ / (1 + θ)) +
  log_meixner_ortho_poly(a, b, 2 * log(θ))
end

"""
    logtp_equal_rates_equal_time_extinction(i)

Evaluate the logarithm of a transition probability of a simple birth and death
process when the death rate `μ` is equal to the birth rate `λ`, time `t` is
equal to `1 / λ`, and 'j' is equal to zero. It is a very unlikely event
(let's say impossible) but it simplifies into a special case.
"""
@inline function logtp_equal_rates_equal_time_extinction(
  i::I
) where {
  I <: Integer
}
  - i * log(2)
end

"""
    logtp_equal_rates_extinction(i, t, λ)

Evaluate the logarithm of the probability of extinction for a simple birth and
death process when the death rate `μ` is equal to the birth rate `λ`.
"""
@inline function logtp_equal_rates_extinction(
  i::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  i * log((λ * t) / (1 + λ * t))
end

"""
    logtp_pure_birth(i, j, t, λ)

Evaluate the logarithm of the transition probability of a pure-birth process.
"""
@inline function logtp_pure_birth(
  i::I,
  j::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  lgamma(j) + (j - i) * logexpm1(λ * t) -
  (lgamma(i) + lgamma(j - i + 1) + j * λ * t)
end

"""
    logtp_pure_death(i, j, t, μ)

Evaluate the logarithm of the transition probability of a pure-death process.
"""
@inline function logtp_pure_death(
  i::I,
  j::I,
  t::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  lgamma(i + 1) + (i - j) * logexpm1(μ * t) -
  (lgamma(j + 1) + lgamma(i - j + 1) + i * μ * t)
end

"""
    logtp_equal_time(i, j, λ, μ)

Evaluate the logarithm of a transition probability of a simple birth and death
process when the time `t` is equal to `log(λ / μ) / (λ - μ)`. It is a very
unlikely event (let's say impossible) but it simplifies into a special case.
"""
@inline function logtp_equal_time(
  i::I,
  j::I,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  lgamma(i + j) + i * log(μ) + j * log(λ) -
  (lgamma(i) + lgamma(j + 1) + (i + j) * log(λ + μ))
end

"""
    logtp_sum(i, j, t, λ, μ)

Evaluate the logarithm of a transition probability of a simple birth and death
process.
"""
@inline function logtp_sum(
  i::I,
  j::I,
  t::F,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = (λ - μ) * t
  ω = log(λ / μ)

  x = if λ < μ
    logexpm1(θ - ω) + log1mexp(θ + ω) - 2 * log1mexp(θ)
  else
    logexpm1(θ + ω) + log1mexp(θ - ω) - 2 * logexpm1(θ)
  end

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + j * ω +
  (i + j) * log(expm1(θ) / expm1(θ + ω)) +
  log_hypergeometric(a, b, x)
end

"""
    logtp_sum_alternating(i, j, t, λ, μ)

Evaluate the logarithm of a transition probability of a simple birth and death
process when terms are alternating in sign.
"""
@inline function logtp_sum_alternating(
  i::I,
  j::I,
  t::F,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = (λ - μ) * t
  ω = log(λ / μ)

  x = log(2 * λ * μ) +
      logexpm1(log1pexp(2 * θ) - (θ + log(F(2)))) -
      log((λ - μ)^2)

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + j * ω +
  (i + j) * log(expm1(θ) / expm1(θ + ω)) +
  log_meixner_ortho_poly(a, b, x)
end

"""
    logtp_equal_time_extinction(i, t, λ, μ)

Evaluate the logarithm of the probability of extinction for a simple birth and
death process when `t = log(λ / μ) / (λ - μ)`.
"""
@inline function logtp_equal_time_extinction(
  i::I,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  i * log(μ / (λ + μ))
end

"""
    logtp_extinction(i, t, λ, μ)

Evaluate the logarithm of the probability of extinction for a simple birth and
death process.
"""
@inline function logtp_extinction(
  i::I,
  t::F,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  i * log(expm1((λ - μ) * t) / expm1(log(λ / μ) + (λ - μ) * t))
end

"""
    trans_prob(i, j, t, λ, μ; log_value)

Evaluate the transition probability of a simple birth and death process, i.e.
the probability of moving from `i` to `j` in `t` time when the birth rate is
equal to `λ` and the death rate is equal to `μ`.

The minimum amount of units `i` allowed at time 0 is 1 but `j` is allowed to be
0. Time `t`, birth rate `λ`, and death rate `μ` cannot be negative.

Let ``α = (μ e^{(λ - μ) t} - μ) / (λ e^{(λ - μ) t} - μ)`` and
``β = (λ e^{(λ - μ) t} - λ) / (λ e^{(λ - μ) t} - μ)``. Transition probability
is equal to (Bailey, 1964):
``\\sum_{h = 0}^{\\min(i, j)} \\binom{i}{h} \\binom{i + j - h - 1}{i - 1}
α^{i - h} β^{j - h} (1 - α - β)^{h}``.

# References:

Bailey, N. T. J. (1964). The elements of stochastic processes with applications
to the natural sciences. Wiley, New York, NY, USA. ISBN 0-471-04165-3.
"""
function trans_prob(
  i::I,
  j::I,
  t::Real,
  η::Vector{R};
  log_value::Bool=true
) where {
  I <: Integer,
  R <: Real
}
  if i < 1
    msg = "Initial population size 'i' must be greater than zero."
    throw(DomainError(i, msg))
  end

  if j < 0
    msg = "Final population size 'j' must be greater than or equal to zero."
    throw(DomainError(j, msg))
  end

  if t < 0
    msg = "Time 't' must be greater than or equal to zero."
    throw(DomainError(t, msg))
  end

  if η[1] < 0
    msg = "Birth rate 'η[1]' must be greater than or equal to zero."
    throw(DomainError(η[1], msg))
  end

  if η[2] < 0
    msg = "Death rate 'η[2]' must be greater than or equal to zero."
    throw(DomainError(η[2], msg))
  end

  # promote values to the appropriate float
  F = float(R)
  t, λ, μ = F(t), F(η[1]), F(η[2])
  ϵ = floatmin(F)

  log_trans_prob::F = if (t < ϵ) || ((λ < ϵ) && (μ < ϵ))
    if i == j
      0
    else
      -Inf
    end
  elseif λ ≈ μ
    ξ = 1 / λ

    if j > 0
      if t ≈ ξ
        logtp_equal_rates_equal_time(i, j)
      elseif t < ξ
        logtp_equal_rates(i, j, t, λ)
      else
        logtp_equal_rates_alternating(i, j, t, λ)
      end
    else
      if t ≈ ξ
        logtp_equal_rates_equal_time_extinction(i)
      else
        logtp_equal_rates_extinction(i, t, λ)
      end
    end
  elseif μ < ϵ
    if j >= i
      logtp_pure_birth(i, j, t, λ)
    else
      -Inf
    end
  elseif λ < ϵ
    if j <= i
      logtp_pure_death(i, j, t, μ)
    else
      -Inf
    end
  else
    ξ = log(λ / μ) / (λ - μ)

    if j > 0
      if t ≈ ξ
        logtp_equal_time(i, j, λ, μ)
      elseif t < ξ
        logtp_sum(i, j, t, λ, μ)
      else
        logtp_sum_alternating(i, j, t, λ, μ)
      end
    else
      if t ≈ ξ
        logtp_equal_time_extinction(i, λ, μ)
      else
        logtp_extinction(i, t, λ, μ)
      end
    end
  end

  if log_value
    log_trans_prob
  else
    exp(log_trans_prob)
  end
end


export NormalNormalTransformer

@doc raw"""
    NormalNormalTransformer(normal_sampler, new_params, old_params=(0.0, 1.0))

Generate normal random variables from another normal sampler.

Transform samples from a normal distribution sampler into samples from a normal
distribution with different parameters. This uses the result that if ``Z`` is a
standard normal random variable (``Z \sim \text{Normal}(0, 1)``) then
``X =  \sigma Z + \mu`` has distribution ``\text{Normal}(\mu, \sigma^2)``.

# Arguments
* `normal_sampler::sampler`: a sampler producing normally distributed samples.
* `new_params::Tuple{Real, Real}`: a tuple (mu, sigma2) giving the mean and
    variance of the target distribution
* `old_params::Tuple{Real, Real}=(0.0, 1.0)`: a tuple (mu, sigma2) giving the
    mean and variance of the normal sampler. Defaults to a standard normal
    distribution.

# Examples
```julia
# Transforming ten standard normal random variables into N(1, 2)
s = BoxMullerTransformSampler()
t = NormalNormalTransformer(s, (1, 2))
transformed_samples = sample(t, 10)
```
"""
struct NormalNormalTransformer <: Sampler
    normal_sampler::Sampler
    new_params::NamedTuple{(:mu, :sigma2), Tuple{Real, Real}}
    old_params::NamedTuple{(:mu, :sigma2), Tuple{Real, Real}}
    function NormalNormalTransformer(
        normal_sampler::Sampler,
        new_params::Tuple{Real, Real},
        old_params::Tuple{Real, Real}=(0.0, 1.0)
    )
        # Name parameter tuples
        new_params = (mu = new_params[1], sigma2 = new_params[2])
        old_params = (mu = old_params[1], sigma2 = old_params[2])
        # Validation
        if !(old_params[:sigma2] > 0 && new_params[:sigma2] > 0)
            error("variance must be positive")
        end
        new(normal_sampler, new_params, old_params)
    end
end

@doc raw"""
    sample(s::NormalNormalTransformer, N)

Transform normal samples into normal samples with new parameters.

# Arguments
* `s::NormalNormalTransformer`: a normal-normal transformer.
* `N::Integer`: the number of samples to generate.

# Returns
* `samples::Array{Float64}`: a 2-D array in which each column is a sample.
"""
function sample(s::NormalNormalTransformer, N::Integer)
    old_samples = sample(s.normal_sampler, N)
    std_samples = (old_samples .- s.old_params[:mu]) ./
        sqrt(s.old_params[:sigma2])
    new_samples = sqrt(s.new_params[:sigma2]) .* std_samples .+
        s.new_params[:mu]

    return new_samples
end


using Primes

# These primes are good example primes, as they have p = ±4 (mod 9), so taking cube roots is easy in F_p^2
plist = [5, 13, 23, 31, 41, 59, 67]
bigplist = [big(2)^107 - 1, big(2)^521 - 1]

# The prime used in the OG SIDH paper. Does not have above property^
OGp = big(2)^63*big(3)^41*11-1

# If p = 2 (mod 3) we have an easy supersingular curve
# E = Affine_EC_TwiHes(F, inv(4*one(F)), zero(F))
# Or twisted to
# E = Affine_EC_TwiHes(F, sqrt(F(b))*inv(4*one(F)), zero(F))
# This is isomorphic to the Weierstrass curve y^2 = x^3 + 1

# Similarly (slightly worse though) If p = 3 (mod 4) we have an easy supersingular curve
# y^2 = x^3 - x
# Translating this to twisted Hessian form will be future work

# If p = 1 (mod 12), just run away, dont look back.

#Scratch all this, much easier than that...
#If p = 2 (mod 3), then 0 is a supersingular j-invariant
#If p = 3 (mod 4), then 1728 is a supersingular j-invariant

#Example:
OGp = big(2)^63*big(3)^41*11-1
F = Fpsquare(OGp)
E = TwiHesEC(F, inv(4*one(F)), zero(F))
#w = F([1850222081870264162797200520152908562431, 2673953603825912900316077551548884232036])
#MAYBE: all points of order two are of the form (X : 1 : 1)
#Want to find roots of x^3-a = 0, where a = 3700444163740528325594401040305817124855
#One lies in F_p, the other two are the roots of x^2 + 3700444163740528325594401040305817124861*x + 4
#Roots are r1 = 3700444163740528325594401040305817124862 + 1647463043911297475037754062791951339209*i
#and r2 = 3700444163740528325594401040305817124862 + 2052981119829230850556646977513865785654*i
#r1 = F([3700444163740528325594401040305817124862, 1647463043911297475037754062791951339209])
#P = E(-r1, one(F), one(F))
#E0 = Isogeny(E,P).codomain
#c = F([3398654792751657841130893587213062545659])*w
#Q = E0(zero(F), -w, one(F))
#E1 = Isogeny(E0, Q).codomain

#Simple example with isogeny over F_p
F = Zmod(31)
E = TwiHesEC(F, F(1), F(7))
listP = allPointsNaive(E)
phi = Isogeny(E, E(F(13), F(1), F(1)))

#Simple example with isogeny over F_p^2
F = Fpsquare(31)
E = TwiHesEC(F, F([4, 5]), F([1,2]))
listP = allPointsNaive(E)
phi = Isogeny(E, E(F([9,19]), F([1]), F([1])))


"""
    Module Basis1D
General math tools for 1D bases used by DG
"""

module Basis1D

export gauss_lobatto_quad, gauss_quad
export jacobiP, grad_jacobiP
export vandermonde1D, grad_vandermonde1D
export chebyshev
export dPhi1D

using LinearAlgebra
using SpecialFunctions

"""
    gauss_lobatto_quad(degree, α=0, β=0)
Initialize Gauss-Legendre-Lobatto (GLL) quadrature nodes and weights (x,w)
to be able to exactly integrate (α,β) Jacobi polynomial of given degree
"""
function gauss_lobatto_quad(degree, α=0, β=0)
    # Need n points to integrate polynomials of degree 2n-3
    n = ceil(Int, (degree+3)/2)
    p = n-1
    if (α!=0) && (β!=0)
        error("currently only allows alpha == beta == 0")
    end
    x = zeros(p+1)
    w = zeros(p+1)
    if p == 0
        x[1] = 0
        w[1] = 2
    elseif p == 1
        x[1] = -1.0
        x[2] = 1.0
        w[1] = 1.0
        w[2] = 1.0
    else
        xint, w = gauss_quad(2*p-3, α+1, β+1)
        x = [-1 transpose(xint) 1]

        V = vandermonde1D(p,x)
        w = vec(sum(inv(V*V'),dims=2))
    end
    return x[:],w[:]
end

"""
    gauss_quad(degree, α, β)
Initialize Gaussian quadrature nodes and weights (x,w) to be able to exactly
integrate Jacobi Polynomial (α,β) of given degree
"""
function gauss_quad(degree, α=0, β=0)
    # Need n points to integrate polynomials of degree 2n-1
    n = ceil(Int, (degree+1)/2)
    p = n-1
    if p == 0
        x = [-(α-β)/(α+β+2)]
        w = [2]
        return x, w
    end

    J = zeros(p+1, p+1)
    h₁ = @. 2*(0:p)+α+β
    J = diagm(0 => (@. -1/2*(α^2-β^2)/(h₁+2)/h₁), 1 => (@. 2/(h₁[1:p]+2)*sqrt((1:p)*((1:p)+α+β)*((1:p)+α)*((1:p)+β)/(h₁[1:p]+1)/(h₁[1:p]+3))))
    if α+β<10*eps()
        J[1,1] = 0.0
    end
    J = J + transpose(J) # Finalize symmetric tridiagonal Jacobi matrix

    x, V = eigen(J)
    w = @. transpose(V[1,:])^2*2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    return x[:], w[:]
end

"""
    grad_jacobiP(p, r, α, β)
Evaluate the derivative of Jacobi Polynomial (α, β) of order p at nodes r
"""
function grad_jacobiP(p, r, α, β)
    dP = zeros(length(r))
    if p != 0
        dP = sqrt(p*(p+α+β+1))*jacobiP(p-1,r,α+1,β+1)
    end
    return dP
end

"""
    jacobiP(p, x, α, β)
Evaluate Jacobi Polynomial (α, β) of order p at points x
"""
function jacobiP(p, x, α, β)
    xp = x
    if size(xp, 2) == 1
        xp = transpose(xp)
    end

    PL = zeros(p+1,length(xp))
    γ₀ = 2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    PL[1,:] .= 1.0/sqrt(γ₀)
    if p == 0
        P = transpose(PL)
        return P
    end
    γ₁ = (α+1)*(β+1)/(α+β+3)*γ₀
    PL[2,:] = ((α+β+2).*xp/2 .+ (α-β)/2)/sqrt(γ₁)
    if p == 1
        P = PL[p+1,:]
        return P
    end

    aold = 2/(2+α+β)*sqrt((α+1)*(β+1)/(α+β+3))

    for i = 1:p-1
        h₁ = 2i+α+β
        anew = 2/(h₁+2)*sqrt((i+1)*(i+1+α+β)*(i+1+α)*(i+1+β)/(h₁+1)/(h₁+3))
        bnew = -(α^2-β^2)/h₁/(h₁+2)
        PL[i+2,:] = 1/anew*(-aold*transpose(PL[i,:]).+(xp.-bnew).*transpose(PL[i+1,:]))
        aold = anew
    end

    P = PL[p+1,:]
    return P;
end

"""
    vandermonde1D(p, r)
Initialize the 1D Vandermonde matrix of order p Legendre polynomials at nodes r
"""
function vandermonde1D(p, r)
    V1D = zeros(length(r), p+1)
    for j = 1:p+1
        V1D[:,j] = jacobiP(j-1, r[:], 0, 0)
    end
    return V1D
end

"""
    grad_vandermonde1D(p, r)
Initialize the 1D gradient Vandermonde matrix of order p Legendre polynomials at nodes r
"""
function grad_vandermonde1D(p, r)
    V1D = zeros(length(r), p+1)
    for j = 1:p+1
        V1D[:,j] = grad_jacobiP(j-1, r[:], 0, 0)
    end
    return V1D
end

"""
    chebyshev(p)
Return the 1D Chebyshev nodes of order p on [-1,1]
"""
function chebyshev(p)
    cheby = cos.((p:-1:0)*π/p)
end

"""
    interpolation_matrix1D(x_from, x_to)
Compute an interpolation matrix from a set of 1D points to another set of 1D points.
Assumes that points interpolating from provide enough accuracy (aka - they are
well spaced out and of high enough order), and define an interval. Points
interpolating onto can be of any size, but must be defined on this same interval.
Interpolation matrix ∈ ℜ^(size of x_to x size of x_from)
"""
function interpolation_matrix1D(x_from, x_to)
    # Create nodal representation of reference bases
    order = size(x_from,1) - 1 # Assumes order = (size of x_from) - 1
    n_from = size(x_from,1)
    l_from = vandermonde1D(order, x_from)

    eye = diagm(0=>ones(n_from))
    V = reshape(l_from, n_from,n_from)
    coeffs_phi = V \ eye

    # Compute reference bases on the output points
    l_to = vandermonde1D(order, x_to)
    n_to = size(l_to,1)
    V_to = reshape(l_to, n_to,n_from)

    # Construct interpolation matrix
    Interp1D = V_to*coeffs_phi
end

"""
    dPhi1D(x_from, x_to)
Compute the gradient of the basis functions defined by 1D points on another
set of 1D points.

Assumes that points interpolating from provide enough accuracy (aka - they are
well spaced out and of high enough order), and define an interval. Points
interpolating onto can be of any size, but must be defined on this same interval.
dPhi ∈ ℜ^(size of x_to × size of x_from × 1)
"""
function dPhi1D(x_from, x_to)
    dim = 1
    # Create nodal representation of reference bases
    order = size(x_from,1) - 1 # Assumes order = (size of x_from) - 1
    n_from = size(x_from,1)
    l_from = vandermonde1D(order, x_from)

    eye = diagm(0=>ones(n_from))
    V = reshape(l_from, n_from,n_from)
    coeffs_phi = V \ eye

    # Compute derivative of reference bases on the output points
    dl_to = grad_vandermonde1D(order, x_to)
    n_to = size(dl_to,1)
    dV = reshape(dl_to, n_to,n_from,dim)
    # Construct gradient of phi = dV*coeffs_phi
    dPhi_to = Array{Float64,3}(undef, n_to,n_from,dim)
    for l = 1:dim
        dPhi_to[:,:,l] = dV[:,:,l]*coeffs_phi
    end
    return dPhi_to
end

end


function integrate(f, a, b, n)
	h = (b - a) / n
	x = zeros(1, 3)
	x[1] = f(a) + f(b)
	for i ∈ 1:(n - 1)
		X = a + i*h
		if iseven(i)
			x[3] += f(X)
		else
			x[2] += f(X)
		end
	end
	return h*(x[1] + 2*x[3] + 4*x[2])/3
end

f(x) = x
println(integrate(f, 0, 3, 100))


module SymRCM

# (C) 2020, Petr Krysl

export symrcm

using SparseArrays

"""
    adjgraph(A)

Compute the adjacency graph from a sparse matrix.

The sparse matrix `A` is assumed to be symmetric.
The results will be wrong if it isn't.

- `sortbydeg`: Should the neighbor lists be sorted by column degree? The default is
  `true`, but often results of very similar quality are obtained when this is
  set to `false` and the lists are not sorted. The second option is much
  faster, as the sorting is expensive.
"""
function adjgraph(A::SparseMatrixCSC; sortbydeg = true)
    colptr = A.colptr
    rowval = A.rowval
    ncols = length(colptr)-1
    neighbors = Vector{Vector{eltype(colptr)}}(undef, ncols)
    cdeg = diff(colptr) # the degree is colptr[j+1]-colptr[j]
    for j in 1:ncols
        cstart = colptr[j]
        jdeg = cdeg[j]
        neighbors[j] = [rowval[cstart+m-1] for m in 1:jdeg]
    end
    # All of these sorts can be done in parallel,  they are totally independent.
    # The question is when to switch over to parallel execution so as to
    # amortize the cost of starting up threads.
    if sortbydeg
        for j in 1:ncols
            sort!(neighbors[j], by = j -> cdeg[j])
        end
    end
    return neighbors
end

"""
    nodedegrees(adjgr::Vector{Vector{Int}})

Compute the degrees of the nodes in the adjacency graph.

conn = [9 1 8 4;
       1 3 2 8;
       8 2 7 5;
       2 6 7 7];
nfens = 9;
adjgr = adjgraph(conn, nfens)
nodedegrees(adjgr)

julia> degrees = node_degrees(adjgr)
9-element Array{Int64,1}:
 5
 6
 3
 3
 3
 2
 4
 7
 3
"""
function nodedegrees(adjgr::Vector{Vector{T}}) where {T}
    degrees = fill(0, length(adjgr))
    for k = 1:length(adjgr)
        degrees[k] = length(adjgr[k])
    end
    return degrees
end

"""
    symrcm(adjgr::Vector{Vector{T}}, degrees::Vector{T}) where {T}

Reverse Cuthill-McKee node-renumbering algorithm.
"""
function symrcm(adjgr::Vector{Vector{T}}, degrees::Vector{T}) where {T}
    @assert length(adjgr) == length(degrees)
    # Initialization
    n = length(adjgr)
    ndegperm = sortperm(degrees) # sorted nodal degrees
    inR = fill(false, n) # Is a node in the result list?
    inQ = fill(false, n) # Is a node in the queue?
    R = T[]
    sizehint!(R, n)
    Q = T[] # Node queue
    sizehint!(Q, n)
    while true
        P = zero(T) # Find the next node to start from
        while !isempty(ndegperm)
            i = popfirst!(ndegperm)
            if !inR[i]
                P = i
                break
            end
        end
        if P == zero(T)
            break # That was the last node
        end
        # Now we have a node to start from: put it into the result list
        push!(R, P); inR[P] = true
        empty!(Q) # empty the queue
        append!(Q, adjgr[P]); inQ[adjgr[P]] .= true # put adjacent nodes in queue
        while length(Q) >= 1
            C = popfirst!(Q) # child to put into the result list
            inQ[C] = false # make note: it is not in the queue anymore
            if !inR[C]
                push!(R, C); inR[C] = true
            end
            for i in adjgr[C] # add all adjacent nodes into the queue
                if (!inR[i]) && (!inQ[i]) # contingent on not being in result/queue
                    push!(Q, i); inQ[i] = true
                end
            end
        end
    end
    return reverse(R) # reverse the result list
end

"""
    symrcm(A::SparseMatrixCSC; sortbydeg = true)

Reverse Cuthill-McKee node-renumbering algorithm.

Compute the adjacency graph from a sparse matrix. The sparse matrix `A` is
assumed to be symmetric. The results will be wrong if it isn't.

- `sortbydeg`: Should the neighbor lists be sorted by column degree? The default is
  `true`, but often results of very similar quality are obtained when this is
  set to `false` and the lists are not sorted. The second option can be much
  faster, as the sorting is expensive when the neighbor lists are long.
"""
function symrcm(A::SparseMatrixCSC; sortbydeg = true)
    ag = adjgraph(A; sortbydeg = sortbydeg)
    nd = nodedegrees(ag)
    return symrcm(ag, nd)
end

end # module


using JuMP
using Clp

m = Model(Clp.Optimizer)

@variable(m, x₁ ≥ 0)
@variable(m, x₂ ≥ 0)
@variable(m, x₃ ≥ 0)

@objective(m, Max, 2x₁  + 4x₂ + 3x₃)
@constraint(m, 3x₁  + 4x₂ + 2x₃ ≤ 60)
@constraint(m, 2x₁  + x₂ + 2x₃ ≤ 40)
@constraint(m, x₁  + 3x₂ + 2x₃ ≤ 80)

print(m)

optimize!(m)

status = termination_status(m)

println("Solution status: ", status)

println("Objective value: ", objective_value(m))
println("x₁ = ", value(x₁))
println("x₂ = ", value(x₂))
println("x₃ = ", value(x₃))


using Distributions, LaTeXStrings, Random, Plots
Random.seed!(1234)

# Simulate n coin flips with a bias of 0.6
n = 1000
flips = rand(n) .< 0.6
p = sum(flips)
q = n - p

# Create plot of Beta(α,β) density
α = 10
β = 5

x_values = 0:0.001:1

# Get prior and posterior values
prior = [pdf(Beta(α,β),x) for x in x_values]
posterior = [pdf(Beta(α+p,β+q),x) for x in x_values]

# Create plots
plot(x_values, posterior, seriestype = :line, label = "posterior",
 yaxis = L"f(\theta|y)", xaxis = L"\theta", grid = false)

plot!(x_values, prior, seriestype = :line, label = "prior", grid = false)

posterior_quantile = cdf(Beta(α+p,β+q),0.65) - cdf(Beta(α+p,β+q),0.55)
confidence_interval = [quantile(Beta(α+p,β+q),0.025), quantile(Beta(α+p,β+q),0.975)]
bayes_mean = mean(Beta(α+p,β+q))

plot(x_values, normal_values)
savefig("beta_prior_coin_ex_2.svg")


using DrWatson
@quickactivate "StatReth"

# %%
using Distributions
using Plots

# %% 4.6
w = 6
n = 9
p_grid = range(0, 1, length = 100)
posterior = @. pdf(Binomial(n, p_grid), w) * pdf(Uniform(0, 1), p_grid)
posterior ./= sum(posterior)

plot(p_grid, posterior)


# Counting Inversions
#
# This implementation assumes all the values in the collection are distinct.

function count_inversions_naive(a::AbstractVector)::Int
    cnt = 0
    n = length(a)
    for i in 1:(n - 1)
        cnt += sum(a[i] .> @view a[(i + 1):n])
    end
    cnt
end

"Merge sorts `a` and returns its number of inversions."
function count_inversions_mergesort!(a::AbstractVector)::Int
    aux = zeros(eltype(a), length(a))
    ms_inv_r!(a, aux)
end

function ms_inv_r!(a::AbstractVector, aux::AbstractVector)::Int
    n = length(a)
    if n <= 1
        return 0
    end

    m = div(n, 2)
    cnt = @views ms_inv_r!(a[1:m], aux[1:m]) +
                 ms_inv_r!(a[(m + 1):n], aux[(m + 1):n])

    i = 1
    j = m + 1
    for k in 1:n
        if i > m
            aux[k] = a[j]
            j += 1
        elseif j > n
            aux[k] = a[i]
            i += 1
        elseif a[i] < a[j]
            aux[k] = a[i]
            i += 1
        else
            aux[k] = a[j]
            j += 1
            cnt += m - i + 1
        end
    end
    a .= aux

    cnt
end

# Plargarized from https://stackoverflow.com/a/23201616
function count_inversions_fenwick(a::AbstractVector)::Int
    n = length(a)
    cnt = 0
    count_tree = zeros(Int, n)
    ranks = Dict(v => i for (i, v) in enumerate(sort(a)))
    for x in Iterators.reverse(a)
        r = ranks[x]
        i = r - 1
        while i > 0
            cnt += count_tree[i]
            i -= i & (-i)
        end

        i = r
        while i <= n
            count_tree[i] += 1
            i += i & (-i)
        end
    end
    cnt
end

#===== Unit Tests =====#

using Test
using Random: seed!, shuffle!

@testset "Count inversions" begin
    n = 1000
    seed!()
    a = collect(1:n)
    shuffle!(a)

    n_inversions_naive = count_inversions_naive(a)
    n_inversions_fenwick = count_inversions_fenwick(a)
    n_inversions_mergesort = count_inversions_mergesort!(a)
    @test n_inversions_naive == n_inversions_mergesort == n_inversions_fenwick
    @test issorted(a)
end


using StatsBase
using Plots

Parameters = [[[0 0 ; 0 0.16], [0 ; 0]],
        [[0.85 0.04 ; -0.04 0.85], [0 ; 1.6]],
        [[0.20 -0.26 ; 0.23 0.22], [0 ; 1.6]],
        [[-0.15 0.28 ; 0.26 0.24], [0 ; 0.44]]]

f1(vector, Pars) = Pars[1][1] * vector + Pars[1][2]
f2(vector, Pars) = Pars[2][1] * vector + Pars[2][2]
f3(vector, Pars) = Pars[3][1] * vector + Pars[3][2]
f4(vector, Pars) = Pars[4][1] * vector + Pars[4][2]

function RandomBarnsleyFern(P, PointNums)
    Functions = [f1, f2, f3, f4]
    X = [0.0 for i = 1:PointNums]
    Y = [0.0 for i = 1:PointNums]
    for i in 1:PointNums
        point = [rand() ; rand()]
        for j in 1:P
            s = sample([1,2,3,4], Weights([0.01,0.85,0.07,0.07]))
            point = Functions[s](point, Parameters)
        end
        X[i]= point[1]
        Y[i]= point[2]
    end
    return X, Y
end

for i in [10, 100]
    for j in [10000,1000000]
        x,y = RandomBarnsleyFern(i, j)

        scatter(x,y,markersize = 10^4/j, legend = false, border=:none, color =:green,dpi=200, label = "Points number = $j Operations number = $i")
        savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet1\\Figs\\Q6\\RandomBarnsleyFern-p$(i)num$(j).png")
    end
end


# ## Quiver plots

#md # ![FILE_NAME.png](images/FILE_NAME.png)

n = 7
f(x,y) = 1 / hypot(x, y)
x = repeat(-3:(2*3)/n:3, 1, n) |> vec
y = repeat(-3:(2*3)/n:3, 1, n)' |> vec
vx = f.(x,y) .* cos.(atan.(y,x)) |> vec
vy = f.(x,y) .* sin.(atan.(y,x)) |> vec

quiver(x, y, quiver=(vx, vy), aspect_ratio=:equal)


"""
Return the rotation angle around the
reference axis for the first vector to be
aligned to the second vector.

Arguments
----------
:Type{ProjectedRotationAngle}
    must be type ``ProjectedRotationAngle``

input:AbstractArray
    the vector to be rotated

target:AbstractArray
    the target vector to be aligned to

ref_axis:AbstractArray
    the reference axis around which input vector will be rotated

tol_near_zero=1e-15:AbstractFloat
    (optional) tolerance for the length of the input vectors to be close to zero
"""
function gage(::Type{ProjectedRotationAngle}, input::AbstractArray, target::AbstractArray,
    ref_axis::AbstractArray, tol_near_zero::AbstractFloat=1e-15)

    @debug @assert length(input) == length(target)

    ref_axis_magnitude = norm(ref_axis, 2)
    input_magnitude = norm(input, 2)
    target_magnitude = norm(target, 2)

    if ref_axis_magnitude < tol_near_zero # this means the input vector is the origin
        return 0. # no need to rotate
    else
        unit_ref_axis = ref_axis / ref_axis_magnitude
    end

    if input_magnitude < tol_near_zero # this means the input vector is the origin
        return 0. # no need to rotate
    else
        unit_input = input / input_magnitude
    end

    if target_magnitude < tol_near_zero # this means the target vector is the origin
        return 0. # no need to rotate
    else
        unit_target = target / target_magnitude
    end

    # get the projection along a direction normal to the ref_axis
    function normal_projection(input::AbstractArray, unit_ref_axis::AbstractArray)
        return input - dot(input, unit_ref_axis) * unit_ref_axis
    end

    return gage(RotationAngle, normal_projection(unit_input, unit_ref_axis),
        normal_projection(unit_target, unit_ref_axis))
end


module DiffSolve

export eulersolve, eulercromer, eulermid, verlet, velverlet, rk4

"""
Solve a system of first order ODEs dx/dt = f(t, x(t)) using the Euler method.
x(t) is a vector.

# Arguments
- `init`: initial x vector
- `func`: the function dx/dt = f(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
vector of all t and matrix containing x at all t (the (i, j) element in the matrix
represents the jth variable in the x vector x[j] at the ith time step t[i]).
"""
function eulersolve(func::Function, init::Vector{Float64}, start::Float64, stop::Float64;
        step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Matrix{Float64}(undef, length(ts), length(init))
    xs[1, :] = init

    x = init
    for (i, t) in enumerate(@view ts[1:end-1])
        x += step * func(t, x)
        xs[i + 1, :] = x
    end

    return ts, xs
end

"""
Solve a system of first order ODEs dx/dt = f(t, x(t)) using the 4th order Runge–Kutta
method (a.k.a. RK4, or Classic Runge–Kutta). x(t) is a vector.

# Arguments
- `init`: initial x vector
- `func`: the function dx/dt = f(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
vector of all t and matrix containing x at all t (the (i, j) element in the matrix
represents the jth variable in the x vector x[j] at the ith time step t[i]).
"""
function rk4(func::Function, init::Vector{Float64}, start::Float64, stop::Float64;
        step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Matrix{Float64}(undef, length(ts), length(init))
    xs[1, :] = init

    x = init
    for (i, t) in enumerate(@view ts[1:end-1])
        k1 = func(t, x)
        k2 = func(t + step / 2, x .+ step / 2 * k1)
        k3 = func(t + step / 2, x .+ step / 2 * k2)
        k4 = func(t + step, x .+ step * k3)

        x += step / 6 * (k1 + 2k2 + 2k3 + k4)
        xs[i + 1, :] = x
    end

    return ts, xs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using the Euler–Cromer method
(a.k.a. the semi-implicit euler method, or the Newton–Størmer–Verlet).

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors.
"""
function eulercromer(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))
    xs[1] = initx
    vs[1] = initv

    x = initx
    v = initv
    for (i, t) in enumerate(@view ts[1:end-1])
        v += step * acceleration(t, x)
        x += step * v

        xs[i + 1] = x
        vs[i + 1] = v
    end

    return ts, xs, vs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using the Midpoint method
(a modified euler method).

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors. Due to the way this method works, the v vector is the
velocity at midpoints between the time steps.
"""
function eulermid(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))
    xs[1] = initx
    vs[1] = initv - acceleration(start, initx) * step / 2

    x = initx
    v = vs[1]
    for (i, t) in enumerate(@view ts[1:end-1])
        v += step * acceleration(t, x)
        x += step * v

        xs[i + 1] = x
        vs[i + 1] = v
    end

    return ts, xs, vs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using Verlet integration
(a.k.a. the Störmer–Verlet method).

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors.
"""
function verlet(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))

    xs[1] = initx
    vs[1] = initv
    xs[2] = initx + initv * step + acceleration(start, initx) * step * step / 2

    for i in 2:length(xs)-1
        xs[i + 1] = 2 * xs[i] - xs[i - 1] + acceleration(ts[i], xs[i]) * step * step
        vs[i + 1] = (xs[i + 1] - xs[i]) / step
    end

    return ts, xs, vs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using the velocity Verlet method

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors.
"""
function velverlet(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))

    xs[1] = initx
    vs[1] = initv

    for i in 1:length(ts)-1
        accbefore = acceleration(ts[i], xs[i])
        xs[i + 1] = xs[i] + vs[i] * step + accbefore * step * step / 2
        accafter = acceleration(ts[i + 1], xs[i + 1])
        vs[i + 1] = vs[i] + (accbefore + accafter) * step / 2
    end

    return ts, xs, vs
end

end


using SpecialFunctions
using Base.MathConstants
using DataFrames, CSV
using Printf

"""
    gamma_derivatives()

Derive a table of derivatives of the Gamma function at 1, i.e., Gamma^{(m)}(a)

This code is used to create the results, which are included in the main code via a lookup table.

"""
function gamma_derivatives(k::Integer, x::Float64)
    #    general form from: http://erikerlandson.github.io/blog/2016/06/15/computing-derivatives-of-the-gamma-function/
    return SpecialFunctions.gamma(x) * D(k,0,x)
end

function D(k::Integer, n::Integer, x::Float64)
    # general form from: http://erikerlandson.github.io/blog/2016/06/15/computing-derivatives-of-the-gamma-function/
    if n<0
        throw(DomainError(n))
    elseif k<0
        throw(DomainError(k))
    end
    if k==0
        return 1
    elseif k==1
        return SpecialFunctions.polygamma(n,x)
    else
        total = D(k-1,n+1,x)
        for j=0:n
            total += binomial(n,j) * D(1,j,x) * D(k-1,n-j,x)
        end
        return total
    end
end

function g1(t::Int) # Crandall,2012, p.17
    # t derivate of Gamma function at 1
    PG21 = -2.4041138063191902 # SpecialFunctions.polygamma(2,1)
    PG41 = -24.8862661234409   # SpecialFunctions.polygamma(4,1)
    if t==0
        return 1.0
    elseif t==1
        return -γ
    elseif t==2
        # return γ^2 - γ + pi^2/6, Crandall seems to be wrong on this
        return γ^2 + pi^2/6  # from Mathematica
    elseif t==3
        # return -2*γ^3 + 9*γ^2 - (π^2+6)*γ + 3*π^2/2 - 4*SpecialFunctions.zeta(3), Crandall seems to be wrong on this
        return -γ^3  - π^2*γ/2 + PG21 # from Mathematica
    elseif t==4
        return γ^4 + γ^2*π^2 + 3*π^4/20 - 4*γ*PG21 # from Mathematica
    elseif t==5
        return -γ^5 - (20/12)*γ^3*π^2 - (9/12)*γ*π^4 + (10*γ^2 + (20/12)*π^2)*PG21 + PG41 # from Mathematica
    else
        return NaN
    end
end

function g2(t::Int) # http://erikerlandson.github.io/blog/2016/06/15/computing-derivatives-of-the-gamma-function/
    # t derivate of Gamma function at 1
    return gamma_derivatives(t,1.0)
end

# #### output a table
# K = 14
# k = collect(0:K)
# G1 = zeros(Float64, K+1)
# G2 = zeros(Float64, K+1)
# for i=1:K+1
#     G1[i] = g1(i-1)
#     G2[i] = g2(i-1)
#     @printf("G1[i] = %.20f, G2[i] = %.20f\n", G1[i], G2[i])
# end

# df = DataFrame( k=k, G1=G1, G2=G2 )
# CSV.write("../data/gamma_derivatives.csv", df)


#############################################################################
#############################################################################
#
# This file contains units tests for polynomial operations
#
#############################################################################
#############################################################################

"""
Test product of polynomials.
"""
function prod_test_poly(;N::Int = 10, N_prods::Int = 20, seed::Int = 0)
    Random.seed!(seed)
    for _ in 1:N
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        prod = p1*p2
        @assert leading(prod) == leading(p1)*leading(p2)
    end

    for _ in 1:N
        p_base = Polynomial(Term(1,0))
        for _ in 1:N_prods
            p = rand(Polynomial)
            prod = p_base*p
            @assert leading(prod) == leading(p_base)*leading(p)
            p_base = prod
        end
    end
    println("prod_test_poly - PASSED")
end

"""
Test product of polynomials modulo P.
"""
function prod_test_poly_mod_p(;N::Int = 10, N_prods::Int = 20, seed::Int = 0)
    Random.seed!(seed)
    sample_primes_array = [3,5,7,11,13,17,19]
    for _ in 1:N
        rand_prime = rand(sample_primes_array)
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        p1mod = PolynomialModP(p1, rand_prime)
        p2mod = PolynomialModP(p2, rand_prime)
        prod = p1mod*p2mod
        @assert mod(leading(prod), rand_prime) == leading(mod(p1mod.terms*p2mod.terms, rand_prime))
    end

    for _ in 1:N
        rand_prime = rand(sample_primes_array)
        p_base = PolynomialModP(Polynomial(Term(1,0)), rand_prime)
        for _ in 1:N_prods
            p = PolynomialModP(rand(Polynomial), rand_prime)
            prod = p_base*p
            prod = PolynomialModP(prod, rand_prime)
            @assert leading(prod.terms) == leading(mod(p_base.terms*p.terms, rand_prime))
            p_base = prod
        end
    end
    println("prod_test_poly_mod_p - PASSED")
end

"""
Test power of polynomial/integer mod P
"""
function power_test_poly(;N::Int = 10, seed::Int = 0)
    sample_primes_array = [3,5,7,11,13,17,19]
    Random.seed!(seed)
    x = x_poly()
    print("Proceeding with power tests ~ [")
    for i in 1:N
        exponent = rand(10:20)
        rand_prime = rand(sample_primes_array)
        p1 = rand(Polynomial)
        out = 1
        for i in 1:exponent
            out = mod(out*p1, rand_prime)
        end
        if ^(p1, exponent, rand_prime) - out != 0
            println("--- TEST FAILED ---")
            break
        end
        print("-")
    end
    print("]\n")
    println("power_test_poly - COMPLETED\n")
end

"""
Test power of polynomial mod P
"""
function power_test_poly_modp(;N::Int = 10, seed::Int = 0)
    sample_primes_array = [3,5,7,11,13,17,19]
    Random.seed!(seed)
    x = x_poly()
    print("Proceeding with power tests using PolynomialModP data type ~ [")
    for i in 1:N
        exponent = rand(10:20)
        rand_prime = rand(sample_primes_array)
        p1 = rand(Polynomial)
        p1mod = PolynomialModP(p1, rand_prime)
        out = 1
        for i in 1:exponent
            out = mod(out*p1mod.terms, p1mod.prime)
        end
        if ^(p1mod, exponent) - out != 0
            println("--- TEST FAILED ---")
            break
        end
        print("-")
    end
    print("]\n")
    println("power_test_poly_modp - COMPLETED\n")
end

"""
Benchmark multiplication
"""
function prod_benchmark(;N::Int = 50, seed::Int = 0)
    println("Begin prod_benchmark: Demonstrating performance of polynomial multiplication algorithm")
    Random.seed!(seed)
    p1 = 0
    p2 = 0
    x = x_poly()
    for i in 1:N
        p1 += rand(1:N)*x^(i)
        p2 += rand(1:N)*x^(i)
        if i%10 == 0
            print("Time taken to multiply two polynomials with $i terms:")
            @time prod = p1*p2
            @assert leading(prod) == leading(p1)*leading(p2)
        end
    end
    println("prod_benchmark COMPLETED\n")
end

function prod_modP_benchmark(;N::Int = 50, seed::Int = 0)
    println("Begin prod_modP_benchmark: Demonstrating performance of polynomial multiplication algorithm in modulo P")
    Random.seed!(seed)
    sample_primes_array = [3,5,7,11,13,17,19]
    p1 = 0
    p2 = 0
    x = x_poly()
    for i in 1:N
        rand_prime = rand(sample_primes_array)
        p1 += rand(1:N)*x^(i)
        p2 += rand(1:N)*x^(i)
        p1mod = PolynomialModP(p1, rand_prime)
        p2mod = PolynomialModP(p2, rand_prime)
        if i%10 == 0
            print("Time taken to multiply two polynomials with $i terms:")
            @time prod = p1mod*p2mod
            @assert leading(prod) == leading(mod((mod(p1mod.terms, rand_prime) * mod(p2mod.terms, rand_prime)), rand_prime))
        end
    end
    println("prod_modP_benchmark COMPLETED\n")
end

function crt_benchmark(;N::Int = 50, seed::Int = 0)
    println("Begin crt_benchmark: Demonstrating performance of polynomial multiplication using the Chinese Remainder Theorem")
    Random.seed!(seed)
    sample_primes_array = [3,5,7,11,13,17,19]
    p1 = 0
    p2 = 0
    x = x_poly()
    for i in 1:N
        x = x_poly()
        rand_prime = rand(sample_primes_array)
        p1 += rand(1:N)*x^(i)
        p2 += rand(1:N)*x^(i)
        if i%10 == 0
            print("Time taken to multiply two polynomials with $i terms:")
            @time prod = mult_poly_with_crt(p1, p2)
            @assert prod == p1 * p2 "Error: CRT does not net p1*p2"
        end
    end
    println("crt_benchmark COMPLETED\n")
end

"""
Test derivative of polynomials (as well as product).
"""
function prod_derivative_test_poly(;N::Int = 10^3,  seed::Int = 0)
    Random.seed!(seed)
    for _ in 1:N
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        p1d = derivative(p1)
        p2d = derivative(p2)
        @assert (p1d*p2) + (p1*p2d) == derivative(p1*p2)
    end
    println("prod_derivative_test_poly - PASSED")
end

"""
Test derivative of polynomials modulo P (as well as product).
"""
function prod_derivative_test_poly_mod_p(;N::Int = 10^3,  seed::Int = 0)
    Random.seed!(seed)
    sample_primes_array = [3,5,7,11,13,17,19]
    for _ in 1:N
        rand_prime = rand(sample_primes_array)
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        p1mod = PolynomialModP(p1, rand_prime)
        p2mod = PolynomialModP(p2, rand_prime)
        p1d = derivative(p1)
        p2d = derivative(p2)
        p1mod_d = PolynomialModP(p1d, rand_prime)
        p2mod_d = PolynomialModP(p2d, rand_prime)
        @assert mod((p1d*p2) + (p1*p2d), rand_prime) == mod(derivative(p1mod.terms*p2mod.terms), rand_prime)
    end
    println("prod_derivative_test_poly_mod_p - PASSED")
end

"""
Test division of polynomials modulo p.
"""
## REMOVE
function division_test_poly(;prime::Int = 101, N::Int = 10^4, seed::Int = 0)
    Random.seed!(seed)
    for _ in 1:N
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        p_prod = p1*p2
        q, r = Polynomial(), Polynomial()
        try
            q, r = divide(p_prod, p2)(prime)
            if (q, r) == (nothing,nothing)
                println("Unlucky prime: $p1 is reduced to $(p1 % prime) modulo $prime")
                continue
            end
        catch e
            if typeof(e) == DivideError
                @assert mod(p2, prime) == 0
            else
                throw(e)
            end
        end
        @assert iszero( mod(q*p2+r - p_prod, prime) )
    end
    println("division_test_poly - PASSED")
end

"""
Test the extended euclid algorithm for polynomials modulo p.
"""
function ext_euclid_test_poly(;prime::Int=101, N::Int = 10^3, seed::Int = 0)
    Random.seed!(seed)
    for _ in 1:N
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        g, s, t = extended_euclid_alg(p1, p2, prime)
        @assert mod(s*p1 + t*p2 - g, prime) == 0
    end
    println("ext_euclid_test_poly - PASSED")
end

"""
Test division of polynomials modulo p using PolynomialModP data type.
"""
function division_test_poly_mod_p(;prime::Int = 101, N::Int = 10^3, seed::Int = 0)
    Random.seed!(seed)
    for _ in 1:N
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        p1mod = PolynomialModP(p1, prime)
        p2mod = PolynomialModP(p2, prime)
        p_prod = PolynomialModP(p1mod.terms*p2mod.terms, prime)
        q, r = Polynomial(), Polynomial()
        try
            q, r = divide(p_prod, p2mod)
            if (q, r) == (nothing,nothing)
                println("Unlucky prime: $p1 is reduced to $(p1 % prime) modulo $prime")
                continue
            end
        catch e
            if typeof(e) == DivideError
                @assert mod(p2mod.terms, prime) == 0
            else
                throw(e)
            end
        end
        @assert iszero(mod(q*p2mod.terms+r - p_prod.terms, prime))
    end
    println("division_test_poly_mod_p - PASSED")
end

"""
Test the extended euclid algorithm for polynomials modulo p.
"""
function ext_euclid_test_poly_mod_p(;prime::Int=101, N::Int = 10^3, seed::Int = 0)
    Random.seed!(seed)
    for _ in 1:N
        p1 = rand(Polynomial)
        p2 = rand(Polynomial)
        p1mod = PolynomialModP(p1, prime)
        p2mod = PolynomialModP(p2, prime)
        g, s, t = extended_euclid_alg(p1mod, p2mod)
        @assert mod(s*p1mod.terms + t*p2mod.terms - g, prime) == 0
    end
    println("ext_euclid_test_poly_mod_p - PASSED")
end


# collatz conjenture

#Jl Plots

function collatz_next(n)
    if n % 2 == 0
        return n ÷ 2
    else
        return 3n+1
    end
end

function collatz_path(n)
    steps = Vector{Int}()
    while n != 1
        push!(steps, n)
        n = collatz_next(n)
    end
    push!(steps, 1)
    return steps
end

using Plots

x = 1:10000
y = length.(collatz_path.(x))
Plots.scatter(x, y, label="Stopping time", legend=:topleft, size=(800,600))

x = 2 .^ (0:13)
y = length.(collatz_path.(x))
Plots.scatter!(x, y, color=:red, label="for powers of two")


module DTWDTfunctions
export test, computeErrorFunction, accumulateErrorFunction, backtrackDistanceFunction, computeDTWerror

"""
 USAGE: err = computeErrorFunction( u1, u0, nSample, lag )

 INPUT:

   u1      = trace that we want to warp; size = (nsamp,1)

   u0      = reference trace to compare with: size = (nsamp,1)

   nSample = numer of points to compare in the traces

   lag     = maximum lag in sample number to search

   norm    = 'L2' or 'L1' (default is 'L2')

 OUTPUT:

    err = the 2D error function; size = (nsamp,2*lag+1)

 The error function is equation 1 in Hale, 2013.

 You could umcomment the L1 norm and comment the L2 norm if you want on Line 29

 Original by Di Yang
 Last modified by Dylan Mikesell (25 Feb. 2015)
"""
function computeErrorFunction(u1::Union{Array{Float32,1},Array{Float64,1}}, u0::Union{Array{Float32,1},Array{Float64,1}}, nSample::Int, lag::Int; norm::String="L2")

    if lag >= nSample
        error("computeErrorFunction:lagProblem ","lag must be smaller than nSample");
    end

    #using JLD2
    #@load "../exampledata/sineShiftData.jld2"
    #nSample = length(u0)
    #lag = 80
    #norm = "L2"

    # Allocate error function variable
    err = zeros(Float64, nSample, 2 * lag + 1 );

    #--------------------------------------------------------------------------
    # initial error calculation
    for ll = -lag:lag # loop over lags

        thisLag = ll + lag + 1;

        for ii = 1:nSample # loop over samples

            if ( ii + ll >= 1 && ii + ll <= nSample ) # skip corners for now, we will come back to these

                diff = u1[ii] - u0[ii + ll]; # sample difference

                if norm == "L2"
                        err[ii, thisLag] = diff^2; # difference squared error
                elseif norm == "L1"
                        err[ii, thisLag] = abs(diff); # absolute value errors
                else
                    error("norm type is not defined.")
                end

            end

        end

    end

    #--------------------------------------------------------------------------
    # Now fix corners with constant extrapolation
    for ll = -lag:lag # loop over lags

        thisLag = ll + lag + 1;

        for ii = 1:nSample # loop over samples

            if ( ii + ll < 1 ) # lower left corner (negative lag, early time)

                err[ii, thisLag] = err[-ll + 1, thisLag];

            elseif ( ii + ll > nSample ) # upper right corner (positive lag, late time)

                err[ii, thisLag] = err[nSample - ll, thisLag];

            end

        end

    end

    return err
end

"""

 USAGE: d = accumulation_diw_mod( dir, err, nSample, lag, b )

 INPUT:

   dir = accumulation direction ( dir > 0 = forward in time, dir <= 0 = backward in time)

   err = the 2D error function; size = (nsamp,2*lag+1)

   nSample = numer of points to compare in the traces

   lag = maximum lag in sample number to search

   b = strain limit (integer value >= 1)

 OUTPUT:

    d = the 2D distance function; size = (nsamp,2*lag+1)

 The function is equation 6 in Hale, 2013.

 Original by Di Yang Last modified by Dylan Mikesell (25 Feb. 2015)
"""
function accumulateErrorFunction(dir::Int, err::Array{Float64,2}, nSample::Int, lag::Int, b::Int)

    #nSample = length(u0)
    #lag = 80
    #dir = -1

    nLag = (2 * lag ) + 1; # number of lags from [ -lag : +lag ]

    # allocate distance matrix
    d = zeros(Float64, nSample, nLag);

    #--------------------------------------------------------------------------
    # Setup indices based on forward or backward accumulation direction
    #--------------------------------------------------------------------------
    if dir > 0            # FORWARD
        iBegin = 1;       # start index
        iEnd   = nSample; # end index
        iInc   = 1;       # increment
    else                  # BACKWARD
        iBegin = nSample; # start index
        iEnd   = 1;       # stop index
        iInc   = -1;      # increment
    end
    #--------------------------------------------------------------------------
    # Loop through all times ii in forward or backward direction

    for ii = iBegin:iInc:iEnd

        # min/max to account for the edges/boundaries
        ji = max(1, min(nSample, ii - iInc ));     # i-1 index
        jb = max(1, min(nSample, ii - iInc * b )); # i-b index

        # loop through all lags l
        for ll = 1:nLag

            # -----------------------------------------------------------------
            # check limits on lag indices
            lMinus1 = ll - 1; # lag at l-1

            if lMinus1 < 1  # check lag index is greater than 1
                lMinus1 = 1; # make lag = first lag
            end

            lPlus1 = ll + 1; # lag at l+1

            if lPlus1 > nLag # check lag index less than max lag
    #             lPlus1 = nLag - 1; # D.Y. version
                lPlus1 = nLag; # D.M. version
            end
            # -----------------------------------------------------------------

            # get distance at lags (ll-1, ll, ll+1)
            distLminus1 = d[jb, lMinus1]; # minus:  d( i-b, j-1 )
            distL       = d[ji, ll];      # actual: d( i-1, j   )
            distLplus1  = d[jb, lPlus1];  # plus:   d( i-b, j+1 )

            if ji != jb # equation 10 in Hale (2013)
                for kb = ji:-iInc:jb+iInc # sum errors over i-1:i-b+1
                    distLminus1 = distLminus1 + err[kb, lMinus1];
                    distLplus1  = distLplus1  + err[kb, lPlus1];
                end
            end

            # equation 6 (if b=1) or 10 (if b>1) in Hale (2013) after treating boundaries
            d[ii, ll] = err[ii, ll] + min(distLminus1, distL, distLplus1);
        end
    end
    return d
end

"""

 USAGE: stbar = backtrackDistanceFunction( dir, d, err, lmin, b )

 INPUT:

   dir   = side to start minimization ( dir > 0 = front, dir <= 0 =  back)

   d     = the 2D distance function; size = (nsamp,2*lag+1)

   err   = the 2D error function; size = (nsamp,2*lag+1)

   lmin  = minimum lag to search over

   b     = strain limit (integer value >= 1)

 OUTPUT:

   stbar = vector of integer shifts subject to |u(i)-u(i-1)| <= 1/b

 The function is equation 2 in Hale, 2013.

 Original by Di Yang Last modified by Dylan Mikesell (19 Dec. 2014)

"""
function backtrackDistanceFunction(dir::Int, d::Array{Float64,2}, err::Array{Float64,2}, lmin::Int, b::Int)

    #d = dist
    #dir = -1
    #lmin = -maxLag
    #b = 1

    nSample = size(d,1); # number of samples
    nLag    = size(d,2); # number of lags
    stbar   = zeros(Int64, nSample); # allocate

    #--------------------------------------------------------------------------
    # Setup indices based on forward or backward accumulation direction
    #--------------------------------------------------------------------------
    if dir > 0            # FORWARD
        iBegin = 1;       # start index
        iEnd   = nSample; # end index
        iInc   = 1;       # increment
    else                  # BACKWARD
        iBegin = nSample; # start index
        iEnd   = 1;       # stop index
        iInc   = -1;      # increment
    end
    #--------------------------------------------------------------------------
    # start from the end (front or back)
    ll0 = argmin(d[iBegin,:]); # find minimum accumulated distance at front or back depending on 'dir'
    stbar[iBegin] = ll0 + lmin - 1; # absolute value of integer shift
    #--------------------------------------------------------------------------
    # move through all time samples in forward or backward direction
    ii = iBegin;

    while ii != iEnd
        if ii == iBegin
            ll = ll0;
        else
            ll = ll_next
        end

        # min/max for edges/boundaries
        ji = max( 1, min(nSample, ii + iInc) );
        jb = max( 1, min(nSample, ii + iInc * b) );

        # -----------------------------------------------------------------
        # check limits on lag indices

        lMinus1 = ll - 1; # lag at l-1

        if lMinus1 < 1 # check lag index is greater than 1
            lMinus1 = 1; # make lag = first lag
        end

        lPlus1 = ll + 1; # lag at l+1

        if lPlus1 > nLag # check lag index less than max lag
            lPlus1 = nLag; # D.M. and D.Y. version
        end
        # -----------------------------------------------------------------
        # get distance at lags (ll-1, ll, ll+1)
        distLminus1 = d[jb, lMinus1]; # minus:  d( i-b, j-1 )
        distL       = d[ji, ll];      # actual: d( i-1, j   )
        distLplus1  = d[jb, lPlus1];  # plus:   d( i-b, j+1 )

        if ji != jb # equation 10 in Hale (2013)
            for kb = ji:iInc:jb-iInc # sum errors over i-1:i-b+1
                distLminus1 = distLminus1 + err[kb, lMinus1];
                distLplus1  = distLplus1  + err[kb, lPlus1];
            end
        end

        dl = min(distLminus1, distL, distLplus1); # update minimum distance to previous sample

        if ( dl != distL ) # then ll != ll and we check forward and backward
            if ( dl == distLminus1 )
                global ll_next = lMinus1;
            else # ( dl == lPlus1 )
                global ll_next = lPlus1;
            end
        else
            ll_next = ll
        end

        # assume ii = ii - 1
        ii += iInc; # previous time sample

        stbar[ii] = ll_next + lmin - 1; # absolute integer of lag
        # now move to correct time index, if smoothing difference over many
        # time samples using 'b'

        if ( ll_next == lMinus1 || ll_next == lPlus1 ) # check edges to see about b values
            if ( ji != jb ) # if b>1 then need to move more steps
                for kb = ji:iInc:jb - iInc
                    ii = ii + iInc; # move from i-1:i-b-1
                    stbar[ii] = ll_next + lmin - 1; # constant lag over that time
                end
            end
        end
    end

    return stbar

    #------------------------#
    # while loop is tricky with Julia
    # we modified ll to ll_next
    # this is the original source from dylanmikesell
    #------------------------#
    # ii = iBegin;
    # while (ii ~= iEnd)
    #
    #     % min/max for edges/boundaries
    #     ji = max( 1, min( [ nSample, ii + iInc ] ) );
    #     jb = max( 1, min( [ nSample, ii + iInc * b ] ) );
    #
    #     % -----------------------------------------------------------------
    #     % check limits on lag indices
    #     lMinus1 = ll - 1; % lag at l-1
    #
    #     if lMinus1 < 1 % check lag index is greater than 1
    #         lMinus1 = 1; % make lag = first lag
    #     end
    #
    #     lPlus1 = ll + 1; % lag at l+1
    #
    #     if lPlus1 > nLag % check lag index less than max lag
    #         lPlus1 = nLag; % D.M. and D.Y. version
    #     end
    #     % -----------------------------------------------------------------
    #
    #     % get distance at lags (ll-1, ll, ll+1)
    #     distLminus1 = d( jb, lMinus1 ); % minus:  d( i-b, j-1 )
    #     distL       = d( ji, ll );      % actual: d( i-1, j   )
    #     distLplus1  = d( jb, lPlus1 );  % plus:   d( i-b, j+1 )
    #
    #     if (ji ~= jb) % equation 10 in Hale (2013)
    #         for kb = ji : iInc : jb - iInc % sum errors over i-1:i-b+1
    #             distLminus1 = distLminus1 + err( kb, lMinus1 );
    #             distLplus1  = distLplus1  + err( kb, lPlus1  );
    #         end
    #     end
    #
    #     dl = min( [ distLminus1, distL, distLplus1 ] ); % update minimum distance to previous sample
    #
    #     if ( dl ~= distL ) % then ll ~= ll and we check forward and backward
    #         if ( dl == distLminus1 )
    #             ll = lMinus1;
    #         else % ( dl == lPlus1 )
    #             ll = lPlus1;
    #         end
    #     end
    #
    #     % assume ii = ii - 1
    #     ii = ii + iInc; % previous time sample
    #
    #     stbar(ii) = ll + lmin - 1; % absolute integer of lag
    #
    #     % now move to correct time index, if smoothing difference over many
    #     % time samples using 'b'
    #     if ( ll == lMinus1 || ll == lPlus1 ) % check edges to see about b values
    #         if ( ji ~= jb ) % if b>1 then need to move more steps
    #             for kb = ji : iInc : jb - iInc
    #                 ii = ii + iInc; % move from i-1:i-b-1
    #                 stbar(ii) = ll + lmin - 1; % constant lag over that time
    #             end
    #         end
    #     end
    # end

end

"""

 Compute the accumulated error along the warping path for Dynamic
 Time Warping.

 USAGE: function error = computeDTWerror( Aerr, u, lag0 )

 INPUT:

   Aerr = error MATRIX (equation 13 in Hale, 2013)

   u    = warping function (samples) VECTOR

   lag0 = value of maximum lag (samples) SCALAR

 Written by Dylan Mikesell Last modified: 25 February 2015
"""
function computeDTWerror(Aerr::Array{Float64,2}, u::Array{Int64,1}, lag0::Int)

    npts = length(u);

    if size(Aerr,1) != npts
        #println("Funny things with dimensions of error matrix: check inputs.");
        #Aerr = transpose(Aerr);
        error("Funny things with dimensions of error matrix: check inputs.")
    end

    error = 0; # initialize

    # accumulate error
    for ii = 1:npts
        idx = lag0 + 1 + u[ii]; # index of lag
        error += Aerr[ii,idx]; # sum error
    end

    return error

end

end


# indicator of the zero cone

"""
  IndZero()

Returns the indicator function of the zero point, or "zero cone", i.e.,
  `g(x) = 0 if x = 0, +∞ otherwise`
"""

immutable IndZero <: IndicatorConvexCone end

function (f::IndZero){T <: RealOrComplex}(x::AbstractArray{T})
  for k in eachindex(x)
    if x[k] != zero(T)
      return Inf
    end
  end
  return 0.0
end

function prox!{T <: RealOrComplex}(f::IndZero, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  for k in eachindex(x)
    y[k] = zero(T)
  end
  return 0.0
end

fun_name(f::IndZero) = "indicator of the zero cone"
fun_dom(f::IndZero) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndZero) = "x ↦ 0 if all(x = 0), +∞ otherwise"
fun_params(f::IndZero) = "none"

function prox_naive(f::IndZero, x::AbstractArray, gamma::Real=1.0)
  return zero(x), 0.0
end


using Distributed

## Function to get the number of points inside a circle
function in_circle_points(n)
    points_in = 0 ## Points inside the circle
    for i = 1:n
        x, y = rand(), rand() ## Randomly generate x,y btwn 0 & 1
        ## Count if generated value is within radius of the circle
        ## int the unit square
        points_in += (x*x + y*y) <= 1
    end
    ## Return the number of points in
    return points_in
end

## Run code in parallel across n cores
function parallel_pi(n)
    ## Create the number of workers specified by the
    p = nworkers()
    ## Combine results via addition (+)
    n_in = @distributed (+) for i=1:p
        in_circle_points(n / p)
    end
    return 4 * n_in / n ## Approximation of pi
end

## Run code serially on 1 worker
function serial_pi(n)
    n_in = in_circle_points(n)
    return 4 * n_in / n
end


@testset "Gauss–Chebyshev" begin

    n = 10

    @testset "x.^2" begin
        x, w = gausschebyshev(n,1)
        @test dot(x.^2,w) ≈ pi/2
        x, w = gausschebyshev(n,2)
        @test dot(x.^2,w) ≈ pi/8
        x, w = gausschebyshev(n,3)
        @test dot(x.^2,w) ≈ pi/2
        x, w = gausschebyshev(n,4)
        @test dot(x.^2,w) ≈ pi/2
    end

    @testset "x^3" begin
        x, w = gausschebyshev(n,1)
        @test abs(dot(x.^3,w))<1e-15
        x, w = gausschebyshev(n,2)
        @test abs(dot(x.^3,w))<1e-15
        x, w = gausschebyshev(n,3)
        @test dot(x.^3,w) ≈ 3*pi/8
        x, w = gausschebyshev(n,4)
        @test dot(x.^3,w) ≈ -3*pi/8
    end
end


#= Fornberg algorithm
This implements the Fornberg (1988) algorithm (https://doi.org/10.1090/S0025-5718-1988-0935077-0)
and hermite-based finite difference Fornberg(2020) algorithm (https://doi.org/10.1093/imanum/draa006)
to obtain Finite Difference weights over arbitrary points to arbitrary order.
Inputs:
        order: The derivative order for which we need the coefficients
        x0   : The point in the array 'x' for which we need the coefficients
        x    : A dummy array with relative coordinates, e.g., central differences
               need coordinates centred at 0 while those at boundaries need
               coordinates starting from 0 to the end point
        dfdx : optional argument to consider weights of the first-derivative of function or not
                if    dfdx == false (default kwarg), implies Fornberg(1988)
                      dfdx == true,                  implies Fornberg(2020)
    Outputs:
        if dfdx == false (default kwarg),   _C : weights to approximate derivative of required order using function values only.
                                 else,   _D,_E : weights to approximate derivative of required order using function and its first-
                                                 derivative values respectively.
=#

function calculate_weights(order::Int, x0::T, x::AbstractVector; dfdx::Bool = false) where T<:Real
    N = length(x)
    @assert order < N "Not enough points for the requested order."
    M = order
    c1 = one(T)
    c4 = x[1] - x0
    C = zeros(T, N, M+1)
    C[1,1] = 1
    @inbounds for i in 1 : N-1
        i1 = i + 1
        mn = min(i, M)
        c2 = one(T)
        c5 = c4
        c4 = x[i1] - x0
        for j in 0 : i-1
            j1 = j + 1
            c3 = x[i1] - x[j1]
            c2 *= c3
            if j == i-1
                for s in mn : -1 : 1
                    s1 = s + 1
                    C[i1,s1] = c1*(s*C[i,s] - c5*C[i,s1]) / c2
                end
                C[i1,1] = -c1*c5*C[i,1] / c2
           end
            for s in mn : -1 : 1
                s1 = s + 1
                C[j1,s1] = (c4*C[j1,s1] - s*C[j1,s]) / c3
            end
            C[j1,1] = c4 * C[j1,1] / c3
        end
        c1 = c2
    end
    #=
        This is to fix the problem of numerical instability which occurs when the sum of the stencil_coefficients is not
        exactly 0.
        https://scicomp.stackexchange.com/questions/11249/numerical-derivative-and-finite-difference-coefficients-any-update-of-the-fornb
        Stack Overflow answer on this issue.
        http://epubs.siam.org/doi/pdf/10.1137/S0036144596322507 - Modified Fornberg Algorithm
    =#
    _C = C[:,end]
    if order != 0
        _C[div(N,2)+1] -= sum(_C)
    end
    if dfdx == false
        return _C
    else
        A = x .- x';
        s = sum(1 ./ (A + I(N)), dims = 1) .- 1;
        cp = factorial.(0:M);
        cc = C./cp'
        c̃ = zeros(N, M+2);
        for k in 1:M+1
           c̃[:,k+1] = sum(cc[:,1:k].*cc[:,k:-1:1], dims = 2);
        end
        E = c̃[:,1:M+1] - (x .- x0).*c̃[:,2:M+2];
        D = c̃[:,2:M+2] + 2*E.*s';
        D = D.*cp';
        E = E.*cp';

        _D = D[:,end];   _E = E[:,end]
        return _D, _E
    end
end


"""
Stationary vector with iterative methods
"""

"""
stguess(Q::MatT, ::Type{Tv} = Float64)::Vector{Tv}

Get a vector which is guessed as the stationary vector of CTMC.
This is used as the initial vector for the iterative methods.
"""

function stguess(Q::MatT, ::Type{Tv} = Float64)::Vector{Tv} where {Tv,MatT}
    m, n = size(Q)
    @assert m == n
    result = Vector{Tv}(undef, n)
    for (i,x) in enumerate(spdiag(Q))
        result[i] = 1/x
    end
    result ./= sum(result)
end

"""
stgs(Q::SparseCSC{Tv,Ti}, x0::Vector{Tv}=stguess(Q,Tv); maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6))

Get a stationary vector of CTMC.

Parameters:
- Q: CTMC Kernal
- x0: Initial vector for iteration
- maxiter: The maximum number of iteration. The algorithm stops when the number of iteration becomes maxiter.
- steps: The number of steps to check the convergence
- rtol: the tolerance error. When the relative errors of two successive vectors with steps attains rtol, the algorithm stops.
Return value:
A tuple of
- x: stationary vector
- conv: A boolean whether the algorithm converges or not
- iter: The number of iterations
- rerror: The relative error when the algorithm stops
"""

function stgs(Q::SparseMatrixCSC{Tv,Ti}; x0::Vector{Tv}=stguess(Q,Tv),
        maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6)) where {Tv,Ti}
    stgs(SparseCSC(Q), x0=x0, maxiter=maxiter, steps=steps, rtol=rtol)
end

function stgs(Q::SparseCSC{Tv,Ti}; x0::Vector{Tv}=stguess(Q,Tv),
        maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6)) where {Tv,Ti}
    m, n = size(Q)
    @assert m == n
    b = zeros(Tv, n)
    x = copy(x0)
    iter = 0
    conv = false
    rerror::Tv = 0
    prevx = similar(x)
    while true
        prevx .= x
        for i in 1:steps
            gsstep!(x, Q, b)
            x ./= sum(x)
        end
        # rerror = maximum(abs.((x - prevx) ./ x))
        rerror = maximum(abs.(x - prevx)) / maximum(x)
        iter += steps
        if rerror < rtol
            conv = true
            break
        end
        if iter >= maxiter
            break
        end
    end
    return x, conv, iter, rerror
end

"""
stpower(P::AbstractMatrix{Tv}, x0::Vector{Tv}=stguess(Q,Tv); maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6))

Get a stationary vector of DTMC with power method.

Parameters:
- P: The transition probability matrix for DTMC
- x0: Initial vector for iteration
- maxiter: The maximum number of iteration. The algorithm stops when the number of iteration becomes maxiter.
- steps: The number of steps to check the convergence
- rtol: the tolerance error. When the relative errors of two successive vectors with steps attains rtol, the algorithm stops.
Return value:
A tuple of
- x: stationary vector
- conv: A boolean whether the algorithm converges or not
- iter: The number of iterations
- rerror: The relative error when the algorithm stops
"""

function stpower(P::AbstractMatrix{Tv}; x0::Vector{Tv}=stguess(P,Tv),
    maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6)) where {Tv}
    m, n = size(P)
    @assert m == n
    Pdash = P'
    x = copy(x0)
    iter = 0
    conv = false
    rerror::Tv = 0
    prevx = similar(x)
    while true
        prevx .= x
        for i in 1:steps
            x = Pdash * x
            x ./= sum(x)
        end
        # rerror = maximum(abs.((x - prevx) ./ x))
        rerror = maximum(abs.(x - prevx)) / maximum(x)
        iter += steps
        if rerror < rtol
            conv = true
            break
        end
        if iter >= maxiter
            break
        end
    end
    return x, conv, iter, rerror
end


using MLJ
import RDatasets: dataset
import DataFrames: DataFrame, select, Not, describe
using Random

data = dataset("datasets", "USArrests")
names(data)

describe(data, :mean, :std)

X = select(data, Not(:State))
X = coerce(X, :UrbanPop=>Continuous, :Assault=>Continuous);

PCA = @load PCA pkg=MultivariateStats

pca_mdl = PCA(pratio=1)
pca = machine(pca_mdl, X)
fit!(pca)
PCA
W = MLJ.transform(pca, X);

schema(W).names

r = report(pca)
cumsum(r.principalvars ./ r.tvar)

data = dataset("ISLR", "OJ")

feature_names = [
    :PriceCH, :PriceMM, :DiscCH, :DiscMM, :SalePriceMM, :SalePriceCH,
    :PriceDiff, :PctDiscMM, :PctDiscCH
]

X = select(data, feature_names);

Random.seed!(1515)

SPCA = Pipeline(
    Standardizer(),
    PCA(pratio=1-1e-4)
)

spca = machine(SPCA, X)
fit!(spca)
W = MLJ.transform(spca, X)
names(W)

rpca = collect(values(report(spca).report_given_machine))[2]
cs = cumsum(rpca.principalvars ./ rpca.tvar)

using PyPlot

figure(figsize=(8,6))

PyPlot.bar(1:length(cs), cs)
plot(1:length(cs), cs, color="red", marker="o")

xlabel("Number of PCA features", fontsize=14)
ylabel("Ratio of explained variance", fontsize=14)

Random.seed!(1515)

KMeans = @load KMeans pkg=Clustering
SPCA2 = Pipeline(
    Standardizer(),
    PCA(),
    KMeans(k=3)
)

spca2 = machine(SPCA2, X)
fit!(spca2)

assignments = collect(values(report(spca2).report_given_machine))[3].assignments
mask1 = assignments .== 1
mask2 = assignments .== 2
mask3 = assignments .== 3;

using PyPlot

figure(figsize=(8, 6))
for (m, c) in zip((mask1, mask2, mask3), ("red", "green", "blue"))
    plot(W[m, 1], W[m, 2], ls="none", marker=".", markersize=10, color=c)
end

xlabel("PCA-1", fontsize=13)
ylabel("PCA-2", fontsize=13)
legend(["Group 1", "Group 2", "Group 3"], fontsize=13)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


module BracketingTests

import ..Algopt # ensure we are using the correct Algopt

using Test
using Algopt.Bracketing:
search, find_bracket, search_bracket,
Bracket,
GoldenSection

using Algopt.TestFunctions: Rosenbrock

quad2 = x->(x - 2)^2

ros = Rosenbrock(a = 1, b = 5)

@testset "Algopt.Bracketing.search(::GoldenSection)" begin
    @test 2 ≈ search(GoldenSection(), quad2)
    @test 2 ≈ search(GoldenSection(), quad2, -1e6)
    @test !(2 ≈ search(GoldenSection(max_iter = 2), quad2))

    @test abs(0 - search(GoldenSection(), x->x'x)) <= eps()

    # Line search of α on 2D Rosenbrock function
    ros_line = α->ros.f([10, 10] + α * [-1, -1])
    @test 9 ≈ search(GoldenSection(), ros_line, 0)
end

@testset "Algopt.Bracketing.find_bracket" begin
    bracket = find_bracket(quad2, 0)
    @test bracket.left < 2
    @test bracket.right > 2

    bracket = find_bracket(quad2, 10)
    @test bracket.left < 2
    @test bracket.right > 2

    bracket = find_bracket(quad2, -1e6)
    @test bracket.left < 2
    @test bracket.right > 2
end

@testset "Algopt.Bracketing.search_bracket" begin
    bracket = search_bracket(GoldenSection(), quad2, Bracket(0, 10))
    @test bracket.left ≈ 2
    @test bracket.right ≈ 2

    bracket = search_bracket(GoldenSection(), quad2, Bracket(-1e6, 1e6))
    @test bracket.left ≈ 2
    @test bracket.right ≈ 2

    bracket = search_bracket(GoldenSection(max_iter = 2), quad2, Bracket(-1e6, 1e6))
    @test !(bracket.left ≈ 2)
    @test !(bracket.right ≈ 2)
end

end # module


# project x onto matrix A's column space, given A is full column rank
function Proj(x,A)
    y = try
        pinv(A' * A) * (A' * x)
    catch err
        if err != Nothing
            A' * x
        end
    end
    return A * y
end

function wavelet_perp_Matrix(w,A)
    N, k = size(A)
    B = zeros(N, k - 1)
    if k == 1
        return
    end
    tmp = A'*w
    if tmp[1] > 1e-6
        M = Matrix{Float64}(I, k-1, k-1)
        return A*vcat((-tmp[2:end] ./ tmp[1])', M)
    end
    if tmp[end] > 1e-6
        M = Matrix{Float64}(I, k-1, k-1)
        return A*vcat(M,(-tmp[1:end-1] ./ tmp[end])')
    end
    ind = findall(tmp .== maximum(tmp))[1]
    rest_ind = setdiff([i for i in 1:k], ind)
    M = Matrix{Float64}(I, k-1, k-1)
    B = A * vcat(vcat(M[1:ind-1,:], (-tmp[rest_ind] ./ tmp[ind])'), M[ind:end,:])
    return B
end

# Gram-Schmidt Process Orthogonalization
function gram_schmidt(A; tol = 1e-12)
    # Input: matirx A
    # Output: orthogonalization matrix of A's column vectors

    # Convert the matrix to a list of column vectors
    a = [A[:,i] for i in 1:size(A,2)]

    # Start Gram-Schmidt process
    q = []
    complement_dim = 0
    for i = 1:length(a)
        qtilde = a[i]
        for j = 1:(i - 1)
            qtilde -= (q[j]' * a[i]) * q[j]
        end
        if norm(qtilde) < tol
            complement_dim = size(A,2) - i + 1
            break
        end
        push!(q, qtilde / norm(qtilde))
    end
    Q = zeros(size(A,1), length(q))
    for i = 1:length(q)
        Q[:,i] .= q[i]
    end
    return Q, complement_dim
end

#
"""
    mgslp(A::Matrix{Float64}; tol::Float64 = 1e-12, p::Float64 = 1.0)

Modified Gram-Schmidt Process Orthogonalization with ℓᵖ pivoting algorithm (MGSLp)

# Input Arguments
- `A::Matrix{Float64}`: whose column vectors are to be orthogonalized.

# Output Argument
- `A::Matrix{Float64}`: orthogonalization matrix of A's column vectors based on ℓᵖ pivoting.
"""
function mgslp(A::Matrix{Float64}; tol::Float64 = 1e-12, p::Float64 = 1.0)
    n = size(A, 2)
    # Convert the matrix to a list of column vectors
    a = [A[:, i] for i in 1:n]

    # Start modified Gram-Schmidt process
    q = []
    v = copy(a)
    vec_lp_norm = [norm(v[i], p) for i in 1:n]
    complement_dim = 0
    for i = 1:n
        # Pivoting based on minimum ℓᵖ-norm
        idx = findmin(vec_lp_norm)[2]
        v[i], v[idx + i - 1] = v[idx + i - 1], v[i]
        # Check the linear dependency
        if norm(v[i]) < tol
            complement_dim = n - i + 1
            break
        end
        qtilde = v[i] / norm(v[i], 2)
        vec_lp_norm = []
        for j = (i + 1):n
            v[j] .-= (qtilde' * v[j]) * qtilde
            push!(vec_lp_norm, norm(v[j], p))
        end
        push!(q, qtilde)
    end
    Q = zeros(size(A,1), length(q))
    for i = 1:length(q)
        Q[:,i] .= q[i]
    end
    return Q, complement_dim
end


export hindmarsh_rose

"""
Equations of motion for the Hindmarsh-Rose system.
"""
function eom_hindmarsh_rose(u, p, t)
    a, b, c, d, r, s, xᵣ, I = (p...,)
    x, y, z = (u...,)

	ϕ = -a*x^3 + b*x^2
	ψ = c - d*x^2
    dx = y + ϕ - z + I
	dy = ψ - y
	dz = r*(s*(x - xᵣ) - z)
    return SVector{3}(dx, dy, dz)
end

"""
	hindmarsh_rose(u₀, p)

Initialise a Hindmarsh-Rose system, which is a model of neuronal
spiking.

```math
\\begin{aligned}
\\dfrac{dx}{dt} &= y + \\phi(x) - z + I
\\dfrac{dy}{dt} &= \\psi(x) - y
\\dfrac{dz}{dt} &= r[s(x - x_R) - z],
\\end{aligned}
```
where

```math
\\begin{aligned}
\\phi(x) &= -ax^3+bx^2
\\psi(x) &= c - dx^2
\\end{aligned}
```

If parameters other than the defaults are to be used, they must be
provided as a vector `[a, b, c, d, r, s, xᵣ, I]`.
"""
function hindmarsh_rose(u₀, p)
    ContinuousDynamicalSystem(eom_hindmarsh_rose, u₀, p)
end
hindmarsh_rose(;u₀ = rand(3), a = 1, b = 3, c = 1, d = 5, r = 1e-3, s = 4, xᵣ = - 8/5, I = -8) =
    hindmarsh_rose(u₀, [a, b, c, d, r, s, xᵣ, I])


#=
proximal.jl

    Evaluate the proximal operator for a variety of commonly used functions in
    optimization, data science, machine learning, and econometrics

@author: Quint Wiersma <q.wiersma@vu.nl>

@date: 2022/01/13
=#

"""
	soft_thresh(x, λ)

Compute soft thresholding operator with scaling parameter `λ` at `x`, proximal
operator of ``ℓ₁``-norm.

#### Arguments
  - `x::Real`	: input
  - `λ::Real`	: scaling parameter

#### Returns
  - `y::Real`	: soft thresholded value
"""
function soft_thresh(x::Real, λ::Real)
	# Soft thresholding
	y= sign(x)*max(abs(x) - λ, zero(λ))

	return y
end

"""
	block_soft_thresh(x, λ)

Compute block soft thresholding operator with scaling parameter `λ` at `x`,
proximal operator of the ``ℓ₂``-norm.

#### Arguments
  - `x::AbstractVector`	: input (n x 1)
  - `λ::Real`			: scaling parameter

#### Returns
  - `y::AbstractVector`	: block soft thresholded value (n x 1)
"""
function block_soft_thresh(x::AbstractVector, λ::Real)
    y= similar(x)
	block_soft_thresh!(y, x, λ)

	return y
end

"""
	block_soft_thresh!(y, x, λ)

Compute block soft thresholding operator with scaling parameter `λ` at `x`,
proximal operator of the ``ℓ₂``-norm, storing the results in `y`. See also
`block_soft_thresh`.
"""
function block_soft_thresh!(y::AbstractVector, x::AbstractVector, λ::Real)
	T= eltype(x)
	# Scaling
	τ= max(one(λ) - λ*inv(norm(x) + eps(T)), zero(λ))

	# Block soft thresholding
	@. y= τ*x

	return nothing
end

"""
	block_soft_thresh!(x, λ)

Compute block soft thresholding operator with scaling parameter `λ` at `x`,
proximal operator of the ``ℓ₂``-norm, overwriting `x`. See also
`block_soft_thresh`.
"""
function block_soft_thresh!(x::AbstractVector, λ::Real)
	T= eltype(x)
	# Scaling
	τ= max(one(λ) - λ*inv(norm(x) + eps(T)), zero(λ))

	# Block soft thresholding
	@. x= τ*x

	return nothing
end

"""
	shrinkage(x, λ)

Compute shrinkage operator with scaling parameter `λ` at `x`, proximal operator
of the squared ℓ₂-norm (ridge).

#### Arguments
  - `x::Real`	: input
  - `λ::Real`	: scaling parameter

#### Returns
  - `y::Real`	: shrunken value
"""
function shrinkage(x::Real, λ::Real)
	# Scaling
	τ= one(λ)*inv(one(λ) + λ)

	# Shrinkage
	y= τ*x

	return y
end

"""
    shrinkage!(x, λ, fac, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b` using a factorization `fac` of ``I + λA``, overwriting
`x`. See also `shrinkage`.
"""
function shrinkage!(x::AbstractVector, λ::Real, fac::Factorization, b::AbstractVector)
    # x - λb
    x.-= λ .* b
    # (I + λA)⁻¹(x - λb)
    ldiv!(fac, x)

    return nothing
end

"""
    shrinkage(x, λ, fac, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b` using a factorization `fac` of ``I + λA``.

#### Arguments
  - `x::AbstractVector` : input
  - `λ::Real`           : scaling parameter
  - `fac::Factorization`: factorization of ``I + λA``
  - `b::AbstractVector` : linear coefficients

#### Returns
  - `y::AbstractVector` : shrunken values
"""
function shrinkage(x::AbstractVector, λ::Real, fac::Factorization, b::AbstractVector)
    y= similar(x)
    shrinkage!(y, λ, fac, b)

    return y
end

"""
    shrinkage!(x, λ, A, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b`, overwriting `x`. See also `shrinkage`.
"""
function shrinkage!(x::AbstractVector, λ::Real, A::AbstractMatrix, b::AbstractVector)
    # scaling I + λA
    S= λ .* A
    @inbounds @fastmath for i in axes(A,1)
        S[i,i]+= one(eltype(A))
    end

    # factorization
    C= cholesky!(Hermitian(S))

    shrinkage!(x, λ, C, b)

    return nothing
end

"""
    shrinkage(x, λ, A, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b`.

#### Arguments
  - `x::AbstractVector` : input
  - `λ::Real`           : scaling parameter
  - `A::AbstractMatrix` : quadratic coefficients
  - `b::AbstractVector` : linear coefficients

#### Returns
  - `y::AbstractVector` : shrunken values
"""
function shrinkage(x::AbstractVector, λ::Real, A::AbstractMatrix, b::AbstractVector)
    y= similar(x)
    shrinkage!(y, λ, A, b)

    return y
end

function smooth!(x::AbstractVector, λ::Real, f::Function, ∇f!::Function, x0::AbstractVector)
    # adjust objective function and gradient
    # f(y) + (1/2λ)‖y - x‖₂²
    g(y::AbstractVector)= f(y) + inv(λ + λ) * ( sum(abs2, y) + sum(abs2, x) - 2 * dot(y, x) )
    # ∇f(y) + (1/λ)(y - x)
    ∇g!(∇g::AbstractVector, y::AbstractVector)= begin
                                                    # gradient
                                                    ∇f!(∇g, y)
                                                    # quadratic term
                                                    ∇g.+= inv(λ) .* (y .- x)
                                                end

    # solve using L-BFGS
    res= optimize(g, ∇g!, x0, LBFGS(), Optim.Options(g_tol = 1e-4))
    x.= Optim.minimizer(res)
    # store current estimate to warm start next iteration
    x0.= x

    return nothing
end


# This file calculates the UL Decomposition of a Toeplitz operator,
# via the Wiener–Hopf factorization.  If the symbol of
# the Toeplitz operator is G(z), we find a Wiener–Hopf factorization
#
# 	G(z) = φ_+(z)*φ_-(z)
#
# where φ_± are analytic in the interior/exterior of the unit
# circle.  Then we have the UL decomposition
#
#   T[G]=T[φ_+]*T[φ_-]
#
# The Wiener–Hopf factorization is found by rephrasing as a
# Riemann–Hilbert problem:
#
#   φ_+(z)=φ_-(z)^(-1)*G(z)
#
# and writing
#
#   φ_-(z)^(-1) = 1 + C_-V(z) and φ_+(z)=1+C_+V(z)=1+V+C_-V(z)
#
# where C_± are the left/right limits of the Cauchy operator.
# This leads to a singular integral equation
#
#   V + C_-V*(1-G) = G - 1
#
# This approach also extends to the matrix case.

using ApproxFun, SingularIntegralEquations

# Scalar case

G=Fun(z->2+cos(z+1/z),Circle()) # the symbol of the Toeplitz operator
T=ToeplitzOperator(G)

C  = Cauchy(-1)
V  = (I+(1-G)*C)\(G-1)

Φmi = 1+C*V
Φp = V+Φmi

L  = ToeplitzOperator(1/Φmi)
U  = ToeplitzOperator(Φp)

norm((T-U*L)[1:10,1:10])  # check the accuracy

# Matrix case

G=Fun(z->[-1 -3; -3 -1]/z +
         [ 2  2;  1 -3] +
         [ 2 -1;  1  2]*z,Circle())
T=ToeplitzOperator(G)

C  = Cauchy(-1)
V  = V=(I+(I-G)*C)\(G-I)

Φmi = I+C*V
Φp = V+Φmi

L  = ToeplitzOperator(inv(Φmi))
U  = ToeplitzOperator(Φp)

norm((T-U*L)[1:10,1:10])  # check the accuracy


# Three-state Bioreactor
id = :PerezGalvan2017e

pL = [6.45; 0.46; 1.05]
pU = [6.55; 0.47; 1.1]
PerezGalvan2017e_x0(p::Vector{T}) where T = T[p[1]; 5.0*one(T); 15.0*one(T)]
function PerezGalvan2017e_f!(du::Vector{T}, u::Vector{T}, p::Vector{T}, t) where T
    x3m = 50.0; x2f = 20.0; D = 0.202; Y = 0.4; β = 0.2; α = 0.5
    μ = p[2]*(1.0 - u[3]/x3m)*u[2]/(p[3] + u[2])
    du[1] = (μ - D)*u[1]
    du[2] = D*(x2f - u[2]) - μ*u[1]/Y
    du[3] = D*u[3] + (α*μ + β)*u[1]
    return
end
tspan = (0.0, 7.7)

prob = ODERelaxProb(PerezGalvan2017e_f!, tspan, PerezGalvan2017e_x0, pL, pU)
url = "https://www.sciencedirect.com/science/article/abs/pii/S0098135417300923"
source = "PerezGalvan2017"
desc = "Three-state Bioreactor"


"""
    eom_chuacircuit_nscroll_sine(u, p, t) -> SVector{3}

Equations of motion for n-scroll chaotic attractors from an adjusted Chua
system [1].

# References
1. Tang, Wallace KS, et al. "Generation of n-scroll attractors via
sine function." IEEE Transactions on Circuits and Systems I:
Fundamental Theory and Applications 48.11 (2001): 1369-1372.

"""
function eom_chuacircuit_nscroll_sine(u, p, t)
    α, β, γ, a, b, c, σx, σy, σz = (p...,)
    x, y, z = (u...,)

    n::Int = c + 1
    if x >= 2*a*c
        fx = (b*pi/2*a)*(x - 2*a*c)
    elseif -2*a*c < x < 2*a*c
        d = ifelse(isodd(n), pi, 0)
        fx = -b*sin((pi*x/2*a) + d)
    elseif x <= -2*a*c
        fx = (b*pi/2*a)*(x + 2*a*c)
    end

    ηx = σx == 0 ? 0 : rand(Normal(0, σx))
    ηy = σy == 0 ? 0 : rand(Normal(0, σy))
    ηz = σz == 0 ? 0 : rand(Normal(0, σz))

    dx = α*(y - fx) + ηx
    dy = x - y + z + ηy
    dz = -β*y - γ*z + ηz
    return SVector{3}(dx, dy, dz)
end

"""
    chua_nscroll_sinefunc(u₀, α, β, γ, a, b, c::Int)

Generate n-scroll chaotic attractors from an adjusted Chua system [1].

# References
1. Tang, Wallace KS, et al. "Generation of n-scroll attractors via
sine function." IEEE Transactions on Circuits and Systems I:
Fundamental Theory and Applications 48.11 (2001): 1369-1372.
"""
function chuacircuit_nscroll_sine(u₀, α, β, γ, a, b, c::Int, σx, σy, σz)
    p = [α, β, γ, a, b, c, σx, σy, σz]
    ContinuousDynamicalSystem(eom_chuacircuit_nscroll_sine, u₀, p)
end
chuacircuit_nscroll_sine(;u₀ = [0.0, 0.0, 0.28695],
        α = 10.814, β = 14, γ = 0, a = 1.3, b = 0.11, c = 2,
        σx = 0.0, σy = 0.0, σz = 0.0) =
    chuacircuit_nscroll_sine(u₀, α, β, γ, a, b, c, σx, σy, σz)


using Test
using FlightMechanics

ψ, θ, ϕ = 0.5, 0.05, 0.2  # rad
p, q, r = 1.0, 0.2, 0.1  # rad/s

# Test regression
ear = body_angular_velocity_to_euler_angles_rates(p, q, r, θ, ϕ)
pqr = euler_angles_rates_to_body_angular_velocity(ear..., θ, ϕ)
@test isapprox(pqr, [p, q, r])

# Test against trimmer function
ψ_dot = 0.2
pqr1 = euler_angles_rates_to_body_angular_velocity(ψ_dot, 0.0, 0.0, θ, ϕ)
pqr2 = turn_rate_angular_velocity(ψ_dot, θ, ϕ)
@test isapprox(pqr1, pqr2)

# Test body_angular_velocity_to_quaternion_rates
quat = euler2quaternion(ψ, θ, ϕ)
qdot = body_angular_velocity_to_quaternion_rates(p, q, r, quat...)

Ω_inv = inv(
    [0 -p -q -r;
     p  0  r -q;
     q -r  0  p;
     r  q -p  0]
)
@test isapprox(Ω_inv * qdot * 2, quat)

# Test uvw_to_tasαβ
exp_tas = 100  # m/s
exp_α = 0.01  # rad
exp_β = 0.005  # rad
u, v, w = wind2body(tas, 0, 0, exp_α, exp_β)
tas, α, β = uvw_to_tasαβ(u, v, w)
@test isapprox(tas, exp_tas)
@test isapprox(α, exp_α)
@test isapprox(β, exp_β, atol=1e-7)

# Test tas_α_β_dot_from_uvw_dot
u = 50.0
v = 1.0
w = 5.0
u_dot = 5.0
v_dot = 0.5
w_dot = 2.0
tas_dot, α_dot, β_dot = tas_α_β_dot_from_uvw_dot(u, v, w, u_dot, v_dot, w_dot)

# Test alpha_dot against formula and tas_dot
tas = √(u*u + v*v + w*w)
@test isapprox(α_dot, (w_dot * u - w * u_dot) / (u*u + w*w))
@test isapprox(tas_dot, (u * u_dot + v * v_dot + w * w_dot) / tas)
# Test β_dot against Stevens implementation (function uses Morelli)
@test isapprox(β_dot, (v_dot * tas - v * tas_dot) / (tas * sqrt(u*u + w*w)))

# tas_α_β_dot_from_uvw_dot <--> uvw_dot_from_tas_α_β
tas, α, β = uvw_to_tasαβ(u, v, w)
u_dot_, v_dot_, w_dot_ = uvw_dot_from_tas_α_β(tas, α, β, tas_dot, α_dot, β_dot)
@test isapprox([u_dot, v_dot, w_dot], [u_dot_, v_dot_, w_dot_], atol=1e-4)


using Flux, DiffEqFlux, DifferentialEquations, Plots, StatsBase, RecursiveArrayTools

const u0 = Float32[2.; 0.]
const datasize = 1000
const batchsize = 30
const batchtime = 10
const tspan = (0.0f0,25f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
true_sol = solve(prob,Tsit5(),saveat=t)
ode_data = Array(true_sol)

function get_batch()
    batch_0s = sample(1:datasize-batchtime, batchsize, replace=false,ordered=false)
    batch_idxs = range.(batch_0s,batch_0s.+(batchtime-1))
    batch_ts = [t[i] for i in batch_idxs]
    batch_u0 = [true_sol[u0] for u0 in batch_0s]
    batch_u = [true_sol[:,idxs] for idxs in batch_idxs]
    return batch_u0, batch_u, batch_ts
end
batch_u0, batch_u, batch_ts = get_batch();

dudt = Chain(x -> x.^3,
             Dense(2,50,tanh),
             Dense(50,2))
ps = Flux.params(dudt)
function n_ode(batch_u0, batch_t)
    neural_ode(dudt,batch_u0,tspan,Tsit5(),
               saveat=batch_t, reltol=1e-7,abstol=1e-9)
end
#= n_ode = x->neural_ode(dudt,x,tspan,Tsit5(), =#
                      #= saveat=t,reltol=1e-7,abstol=1e-9) =#

n_ode(batch_u0, batch_ts)

pred = n_ode(u0,t) # Get the prediction using the correct initial condition
scatter(t,ode_data[1,:],label="data")
scatter!(t,Flux.data(pred[1,:]),label="prediction")

function predict_n_ode(batch_u0,batch_t)
  n_ode(batch_u0, batch_t)
end
loss_n_ode(batch_u0, batch_u,batch_t) = sum(abs2,batch_u .- predict_n_ode(batch_u0,batch_t))

data = Iterators.repeated(get_batch(), 1000)
opt = ADAM(0.1)
cb = function () #callback function to observe training
    bu0,bu,bt = get_batch()
    display(loss_n_ode(bu0,bu,bt))
  # plot current prediction against data
  #= cur_pred = Flux.data(predict_n_ode(bu0,bt)) =#
  #= pl = scatter(bt,ode_data[1,:],label="data") =#
  #= scatter!(pl,bt,cur_pred[1,:],label="prediction") =#
  #= display(plot(pl)) =#
end

# Display the ODE with the initial parameter values.
cb()

Flux.train!(loss_n_ode, ps, data, opt, cb = cb)


using Statistics, MultivariateStats, RDatasets, LinearAlgebra, Plots; pyplot()

data = dataset("datasets", "iris")
data = data[:,[:SepalLength,:SepalWidth,:PetalLength,:PetalWidth]]
x = convert(Array{Float64,2},data)'

model = fit(PCA, x, maxoutdim=4, pratio = 0.999)
M = projection(model)

function manualProjection(x)
    covMat = cov(x')
    ev = eigvals(covMat)
    eigOrder = sortperm(eigvals(covMat),rev=true)
    eigvecs(covMat)[:,eigOrder]
end

println("Manual vs. package: ",maximum(abs.(M-manualProjection(x))))

pcVar = principalvars(model) ./ tvar(model)
cumVar = cumsum(pcVar)
pcDat = M[:,1:2]'*x

p1 = plot(pcVar, c=:blue, label="Variance due to PC")
	plot!(1:length(cumVar), cumVar, label="Cumulative Variance", c=:red,
        xlabel="Principle component",ylabel="Variance Proportion",ylims=(0,1))
p2 = scatter(pcDat[1,:],pcDat[2,:], c=:blue, xlabel="PC 1", ylabel="PC 2",
            msw=0, legend=:none)
plot(p1, p2, size=(800,400))