#    Problem :
#    *********
#   The generalized Rosenbrock function.
#
#   Source: problem 5 in
#   S. Nash,
#   "Newton-type minimization via the Lanczos process",
#
#   SIF input: Nick Gould, Oct 1992.
#              minor correction by Ph. Shott, Jan 1995.
#
#   classification SUR2-AN-V-0
#
#   Number of variables
#
#
# Daniel Henderson, 08/2021

f = x -> begin
	N = lastindex(x)
    return 1.0 + 100sum((x[i] - x[i-1]^2)^2 for i = 2:N) + sum((x[i] - 1.0)^2 for i = 2:N)
end

g! = (g, x) -> begin
	N = lastindex(x)
	g[1] = -400 * x[1] * (x[2] - x[1]^2)
	α    =  200 * (x[2] - x[1]^2)
	for i in 2:N-1
		g[i] =  -400 * x[i] * (x[i+1] - x[i]^2) + 2 * (x[i] - 1.0) + α
		α    =   200 * (x[i+1] - x[i]^2)
	end
	g[N] = 200 * (x[N] - x[N-1]^2) + 2*(x[N] - 1.0)
	return g
end

fg! = (g, x) -> begin
	@warn "GENROSE fg! not implmented"
	# N = lastindex(x)
	# fx = 0.0
	# g[1] = -400 * x[1] * (x[2] - x[1]^2)
	# α    =  200 * (x[2] - x[1]^2)
	# for i in 2:N-1
	# 	g[i] =  -400 * x[i] * (x[i+1] - x[i]^2) + 2 * (x[i] - 1.0) + α
	# 	α    =   200 * (x[i+1] - x[i]^2)
	# 	fx  += (x[i] - x[i-1]^2)^2 + (x[i] - 1.0)^2
	# end
	# fx  += (x[N] - x[N-1]^2)^2 + (x[N] - 1.0)^2
	# g[N] = 200 * (x[N] - x[N-1]^2) + 2*(x[N] - 1.0)
	# return fx,  g
end

init = (n::Int=500) -> begin
	x0 = [j/(n+1) for j in 1:n]
	return n, x0
end

TestSet["GENROSE"] = UncProgram("GENROSE",  f, g!, fg!, init)


# Continuation of ex 7.13, Least Squares WRM

using NumericalMethodsforEngineers
using Test

@sym begin
  LagrangePolynomial(xi_, yi_) := Module([N, sum, num, den],
    begin
      N = Length(xi)
      sum = 0
      For( i=1, i <= N, Increment(i),
        begin
          num = 1
          den = 1
          For( j=1, j <= N, Increment(j),
            begin
              If(j != i, num *= (x-xi[j]))
              If(j != i, den *= (xi[i]-xi[j]))
            end
          )
          sum += yi[i] * num/den
        end
      ),
      Return(Simplify(sum))
    end
  )
end

@sym begin
  ClearAll(xi, yi, N1, Y, Ydotdot, C11, ytilde1, ytilde2)
  xi = [0, 1//2, 1]
  yi = [0, a, 1]
  Y(x_) := LagrangePolynomial(xi, yi)
  Ydotdot(x_) = D(Y(x), x, 2)
  R(x_) := Simplify(Expand(Ydotdot(x) - 3*x - 4*Y(x)))
  R(x_) = Simplify(Expand(R(x) ./ (a => -C1/4)))
  Rdot(x_) = D(R(x), C1)
  C11 = Solve(Integrate(R(x)*Rdot(x), [x, 0, 1]), C1)[1]
  SetJ(r, ToString(Simplify(R(x))))
  SetJ(C1, ToString(C11[1][2]))
  ytilde1 = Simplify(Y(x) ./ (a => -C1/4))
  ytilde2 = Simplify(ytilde1 ./ C11)
  SetJ(y, ToString(Simplify(Expand(ytilde2))));
end

println("\n\nExample 7.17: y'' = 3x + 4y, y(0)=0, y(1)=1")
println("by 1-point least squares Weighted Residual Method")
@sym Print("\nY(x) = $(Y(x))\n")
@sym Print("Y(x) = $(Simplify(Expand(Y(x) ./ (a => -C1/4))))\n")
@sym Println("R(x) = $(R(x))\n")
@sym Print("C1 = $(Solve(Integrate(R(x)*Rdot(x), [x, 0, 1]), C1)[1][1][2])\n")
@eval ytilde_1pt_leastsquares(x) = $(Meta.parse(y))
@eval rf_1pt_leastsquares(x, C1) = $(Meta.parse(r))
println()

rf_1pt_leastsquares_1(x) = rf_1pt_leastsquares(x, Meta.parse(C1))
println()

@test r == "4 + 2C1 + x + 4C1*x - 8x^2 - 4C1*x^2"
@test y == "(1/27)*x*(-8 + 35x)"
#@test (quadgk(rf_1pt_leastsquares_1, 0, 1))[1] < 5*eps()


function sma(ta::TimeArray, n::Int)
    tstamps = ta.timestamp[n:end]

    vals = zeros(size(ta.values,1) - (n-1), size(ta.values,2))
    for i in 1:size(ta.values,1) - (n-1)
        for j in 1:size(ta.values,2)
            vals[i,j] = mean(ta.values[i:i+(n-1),j])
        end
    end

    cname = String[]
    cols  = colnames(ta)
    for c in 1:length(cols)
        push!(cname, string(cols[c], "_sma_", n))
    end

    TimeArray(tstamps, vals, cname, ta.meta)
end

function ema(ta::TimeArray, n::Int; wilder=false)
    k = if wilder
        1 / n
    else
        2 / (n + 1)
    end

    tstamps = ta.timestamp[n:end]

    vals    =  zeros(size(ta.values,1), size(ta.values,2))
    # seed with first value with an sma value
    vals[n,:] = sma(ta, n).values[1,:]

    for i in n+1:size(ta.values,1)
        for j in 1:size(ta.values,2)
            vals[i,j] = ta.values[i,j] * k + vals[i-1, j] * (1-k)
        end
    end

    cname   = String[]
    cols    = colnames(ta)
    for c in 1:length(cols)
        push!(cname, string(cols[c], "_ema_", n))
    end

    TimeArray(tstamps, vals[n:length(ta),:], cname, ta.meta)
end

function kama(ta::TimeArray, n::Int=10, fn::Int=2, sn::Int=30)
    vola = moving(sum, abs.(ta .- lag(ta)), n)
    change = abs.(ta .- lag(ta, n))
    er = safediv.(change, vola)  # Efficiency Ratio

    # Smooth Constant
    sc = (er .* (2 / (fn + 1) - 2 / (sn + 1)) .+ 2 / (sn + 1)).^2

    cl = ta[n+1:end]
    vals = similar(Array{Float64}, indices(cl.values))
    # using simple moving average as initial kama
    pri_kama = mean(ta[1:n].values, 1)

    @assert length(cl) == length(sc)

    for idx ∈ 1:length(cl)
        vals[idx, :] =
            pri_kama =
            pri_kama .+ sc[idx].values .* (cl[idx].values .- pri_kama)
    end

    cols =
    if length(ta.colnames) == 1
        ["kama"]
    else
        ["$c\_kama" for c in ta.colnames]
    end

    TimeArray(cl.timestamp, vals, cols)
end

function env(ta::TimeArray, n::Int; e::Float64 = 0.1)
    tstamps = ta.timestamp[n:end]

    s = sma(ta, n)

    upper = s.values .* (1 + e)
    lower = s.values .* (1 - e)

    up_cname = string.(colnames(ta), "_env_$n", "_up")
    lw_cname = string.(colnames(ta), "_env_$n", "_low")

    u = TimeArray(tstamps, upper, up_cname, ta.meta)
    l = TimeArray(tstamps, lower, lw_cname, ta.meta)

    merge(l, u, :inner)
end

# Array dispatch for use by other algorithms

function sma(a::Array, n::Int)
    vals = zeros(size(a,1) - (n-1), size(a,2))

    for i in 1:size(a,1) - (n-1)
        for j in 1:size(a,2)
            vals[i,j] = mean(a[i:i+(n-1),j])
        end
    end

    vals
end

function ema(a::Array, n::Int; wilder=false)
    k = if wilder
        1 / n
    else
        2 / (n + 1)
    end

    vals = zeros(size(a,1), size(a,2))
    # seed with first value with an sma value
    vals[n,:] = sma(a, n)[1,:]

    for i in n+1:size(a,1)
        for j in 1:size(a,2)
            vals[i,j] = a[i,j] * k + vals[i-1, j] * (1-k)
        end
    end

    vals[n:end, :]
end

function env(a::AbstractArray, n::Int; e::Float64 = 0.1)
    s = sma(a, n)

    upper = @. s * (1 + e)
    lower = @. s * (1 - e)

    [lower upper]
end

doc"""
    sma(arr, n)

Simple Moving Average

```math
SMA = \frac{\sum_i^n{P_i}}{n}
```
"""
sma

doc"""
    ema(arr, n, wilder=false)

Exponemtial Moving Average

A.k.a. exponentially weighted moving average (EWMA)

```math
    \text{Let } k \text{denote the degree of weighting decrease}
```

If parameter `wilder` is `true`, ``k = \frac{1}{n}``,
else ``k = \frac{2}{n + 1}``.

```math
    EMA_t = k \times P_t + (1 - k) \times EMA_{t - 1}
```
"""
ema

doc"""

Kaufman's Adaptive Moving Average

**Arguments**:

- `n`: period

- `fn`: the fastest EMA constant

- `sn`: the slowest EMA constant

**Formula**:

```math
    \begin{align*}
        KAMA_t & = KAMA_{t-1} + SC \times (Price - KAMA_{t-1}) \\
        SC     & =
            (ER \times (\frac{2}{fn + 1} - \frac{2}{sn + 1}) + \frac{2}{sn + 1})^2 \\
        ER     & = \frac{Change}{Volatility} \\
        Change & = | Price - Price_{t-n} | \\
        Volatility & = \sum_{i}^{n} | Price_i - Price_{i-1} |
    \end{align*}
```
"""
kama

doc"""

    env(arr, n; e = 0.1)

Moving Average Envelope

```math
  \begin{align*}
    \text{Upper Envelope} & = \text{n period SMA } \times (1 + e) \\
    \text{Lower Envelope} & = \text{n period SMA } \times (1 - e)
  \end{align*}
```

**Arguments**

- `e`: the envelope, `0.1` implies the `10%` envelope.

**Reference**

- [TradingView](https://www.tradingview.com/wiki/Envelope_(ENV))
"""
env


##### Weighted var & std

## var
"""
    varm(x::AbstractArray, w::AbstractWeights, m, [dim]; corrected=false)

Compute the variance of a real-valued array `x` with a known mean `m`, optionally
over a dimension `dim`. Observations in `x` are weighted using weight vector `w`.
The uncorrected (when `corrected=false`) sample variance is defined as:
```math
\\frac{1}{\\sum{w}} \\sum_{i=1}^n {w_i\\left({x_i - m}\\right)^2 }
```
where ``n`` is the length of the input. The unbiased estimate (when `corrected=true`) of
the population variance is computed by replacing
``\\frac{1}{\\sum{w}}`` with a factor dependent on the type of weights used:
* `AnalyticWeights`: ``\\frac{1}{\\sum w - \\sum {w^2} / \\sum w}``
* `FrequencyWeights`: ``\\frac{1}{\\sum{w} - 1}``
* `ProbabilityWeights`: ``\\frac{n}{(n - 1) \\sum w}`` where ``n`` equals `count(!iszero, w)`
* `Weights`: `ArgumentError` (bias correction not supported)
"""
varm(v::RealArray, w::AbstractWeights, m::Real; corrected::DepBool=nothing) =
    _moment2(v, w, m; corrected=depcheck(:varm, corrected))

"""
    var(x::AbstractArray, w::AbstractWeights, [dim]; mean=nothing, corrected=false)

Compute the variance of a real-valued array `x`, optionally over a dimension `dim`.
Observations in `x` are weighted using weight vector `w`.
The uncorrected (when `corrected=false`) sample variance is defined as:
```math
\\frac{1}{\\sum{w}} \\sum_{i=1}^n {w_i\\left({x_i - μ}\\right)^2 }
```
where ``n`` is the length of the input and ``μ`` is the mean.
The unbiased estimate (when `corrected=true`) of the population variance is computed by
replacing ``\\frac{1}{\\sum{w}}`` with a factor dependent on the type of weights used:
* `AnalyticWeights`: ``\\frac{1}{\\sum w - \\sum {w^2} / \\sum w}``
* `FrequencyWeights`: ``\\frac{1}{\\sum{w} - 1}``
* `ProbabilityWeights`: ``\\frac{n}{(n - 1) \\sum w}`` where ``n`` equals `count(!iszero, w)`
* `Weights`: `ArgumentError` (bias correction not supported)
"""
function var(v::RealArray, w::AbstractWeights; mean=nothing,
                  corrected::DepBool=nothing)
    corrected = depcheck(:var, corrected)

    if mean == nothing
        varm(v, w, Statistics.mean(v, w); corrected=corrected)
    else
        varm(v, w, mean; corrected=corrected)
    end
end

## var along dim

function varm!(R::AbstractArray, A::RealArray, w::AbstractWeights, M::RealArray,
                    dim::Int; corrected::DepBool=nothing)
    corrected = depcheck(:varm!, corrected)
    rmul!(_wsum_centralize!(R, abs2, A, convert(Vector, w), M, dim, true),
          varcorrection(w, corrected))
end

function var!(R::AbstractArray, A::RealArray, w::AbstractWeights, dims::Int;
              mean=nothing, corrected::DepBool=nothing)
    corrected = depcheck(:var!, corrected)

    if mean == 0
        varm!(R, A, w, Base.reducedim_initarray(A, dims, 0, eltype(R)), dims;
                   corrected=corrected)
    elseif mean == nothing
        varm!(R, A, w, Statistics.mean(A, w, dims=dims), dims; corrected=corrected)
    else
        # check size of mean
        for i = 1:ndims(A)
            dA = size(A,i)
            dM = size(mean,i)
            if i == dims
                dM == 1 || throw(DimensionMismatch("Incorrect size of mean."))
            else
                dM == dA || throw(DimensionMismatch("Incorrect size of mean."))
            end
        end
        varm!(R, A, w, mean, dims; corrected=corrected)
    end
end

function varm(A::RealArray, w::AbstractWeights, M::RealArray, dim::Int;
                   corrected::DepBool=nothing)
    corrected = depcheck(:varm, corrected)
    varm!(similar(A, Float64, Base.reduced_indices(axes(A), dim)), A, w, M,
               dim; corrected=corrected)
end

function var(A::RealArray, w::AbstractWeights, dim::Int; mean=nothing,
                  corrected::DepBool=nothing)
    corrected = depcheck(:var, corrected)
    var!(similar(A, Float64, Base.reduced_indices(axes(A), dim)), A, w, dim;
         mean=mean, corrected=corrected)
end

## std
"""
    stdm(x::AbstractArray, w::AbstractWeights, m, [dim]; corrected=false)

Compute the standard deviation of a real-valued array `x` with a known mean `m`,
optionally over a dimension `dim`. Observations in `x` are weighted using weight vector `w`.
The uncorrected (when `corrected=false`) sample standard deviation is defined as:
```math
\\sqrt{\\frac{1}{\\sum{w}} \\sum_{i=1}^n {w_i\\left({x_i - m}\\right)^2 }}
```
where ``n`` is the length of the input. The unbiased estimate (when `corrected=true`) of the
population standard deviation is computed by replacing ``\\frac{1}{\\sum{w}}`` with a factor
dependent on the type of weights used:
* `AnalyticWeights`: ``\\frac{1}{\\sum w - \\sum {w^2} / \\sum w}``
* `FrequencyWeights`: ``\\frac{1}{\\sum{w} - 1}``
* `ProbabilityWeights`: ``\\frac{n}{(n - 1) \\sum w}`` where ``n`` equals `count(!iszero, w)`
* `Weights`: `ArgumentError` (bias correction not supported)
"""
stdm(v::RealArray, w::AbstractWeights, m::Real; corrected::DepBool=nothing) =
    sqrt(varm(v, w, m, corrected=depcheck(:stdm, corrected)))

"""
    std(x::AbstractArray, w::AbstractWeights, [dim]; mean=nothing, corrected=false)

Compute the standard deviation of a real-valued array `x`,
optionally over a dimension `dim`. Observations in `x` are weighted using weight vector `w`.
The uncorrected (when `corrected=false`) sample standard deviation is defined as:
```math
\\sqrt{\\frac{1}{\\sum{w}} \\sum_{i=1}^n {w_i\\left({x_i - μ}\\right)^2 }}
```
where ``n`` is the length of the input and ``μ`` is the mean.
The unbiased estimate (when `corrected=true`) of the population standard deviation is
computed by replacing ``\\frac{1}{\\sum{w}}`` with a factor dependent on the type of
weights used:
* `AnalyticWeights`: ``\\frac{1}{\\sum w - \\sum {w^2} / \\sum w}``
* `FrequencyWeights`: ``\\frac{1}{\\sum{w} - 1}``
* `ProbabilityWeights`: ``\\frac{n}{(n - 1) \\sum w}`` where ``n`` equals `count(!iszero, w)`
* `Weights`: `ArgumentError` (bias correction not supported)
"""
std(v::RealArray, w::AbstractWeights; mean=nothing, corrected::DepBool=nothing) =
    sqrt.(var(v, w; mean=mean, corrected=depcheck(:std, corrected)))

stdm(v::RealArray, m::RealArray, dim::Int; corrected::DepBool=nothing) =
    sqrt!(varm(v, m, dims=dim, corrected=depcheck(:stdm, corrected)))

stdm(v::RealArray, w::AbstractWeights, m::RealArray, dim::Int;
          corrected::DepBool=nothing) =
    sqrt.(varm(v, w, m, dim; corrected=depcheck(:stdm, corrected)))

std(v::RealArray, w::AbstractWeights, dim::Int; mean=nothing,
         corrected::DepBool=nothing) =
    sqrt.(var(v, w, dim; mean=mean, corrected=depcheck(:std, corrected)))

##### Fused statistics
"""
    mean_and_var(x, [w::AbstractWeights], [dim]; corrected=false) -> (mean, var)

Return the mean and standard deviation of collection `x`. If `x` is an `AbstractArray`,
`dim` can be specified as a tuple to compute statistics over these dimensions.
A weighting vector `w` can be specified to weight the estimates.
Finally, bias correction is be applied to the variance calculation if `corrected=true`.
See [`var`](@ref) documentation for more details.
"""
function mean_and_var(x; corrected::Bool=true)
    m = mean(x)
    v = varm(x, m; corrected=corrected)
    m, v
end

"""
    mean_and_std(x, [w::AbstractWeights], [dim]; corrected=false) -> (mean, std)

Return the mean and standard deviation of collection `x`. If `x` is an `AbstractArray`,
`dim` can be specified as a tuple to compute statistics over these dimensions.
A weighting vector `w` can be specified to weight the estimates.
Finally, bias correction is applied to the
standard deviation calculation if `corrected=true`.
See [`std`](@ref) documentation for more details.
"""
function mean_and_std(x; corrected::Bool=true)
    m = mean(x)
    s = stdm(x, m; corrected=corrected)
    m, s
end

function mean_and_var(x::RealArray, w::AbstractWeights; corrected::DepBool=nothing)
    m = mean(x, w)
    v = varm(x, w, m; corrected=depcheck(:mean_and_var, corrected))
    m, v
end
function mean_and_std(x::RealArray, w::AbstractWeights; corrected::DepBool=nothing)
    m = mean(x, w)
    s = stdm(x, w, m; corrected=depcheck(:mean_and_std, corrected))
    m, s
end

function mean_and_var(x::RealArray, dim::Int; corrected::Bool=true)
    m = mean(x, dims = dim)
    v = varm(x, m, dims = dim, corrected=corrected)
    m, v
end
function mean_and_std(x::RealArray, dim::Int; corrected::Bool=true)
    m = mean(x, dims = dim)
    s = stdm(x, m, dim; corrected=corrected)
    m, s
end

function mean_and_var(x::RealArray, w::AbstractWeights, dims::Int;
                      corrected::DepBool=nothing)
    m = mean(x, w, dims=dims)
    v = varm(x, w, m, dims; corrected=depcheck(:mean_and_var, corrected))
    m, v
end
function mean_and_std(x::RealArray, w::AbstractWeights, dims::Int;
                      corrected::DepBool=nothing)
    m = mean(x, w, dims=dims)
    s = stdm(x, w, m, dims; corrected=depcheck(:mean_and_std, corrected))
    m, s
end

##### General central moment
function _moment2(v::RealArray, m::Real; corrected=false)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += z * z
    end
    varcorrection(n, corrected) * s
end

function _moment2(v::RealArray, wv::AbstractWeights, m::Real; corrected=false)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += (z * z) * wv[i]
    end

    varcorrection(wv, corrected) * s
end

function _moment3(v::RealArray, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += z * z * z
    end
    s / n
end

function _moment3(v::RealArray, wv::AbstractWeights, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += (z * z * z) * wv[i]
    end
    s / sum(wv)
end

function _moment4(v::RealArray, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += abs2(z * z)
    end
    s / n
end

function _moment4(v::RealArray, wv::AbstractWeights, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += abs2(z * z) * wv[i]
    end
    s / sum(wv)
end

function _momentk(v::RealArray, k::Int, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += (z ^ k)
    end
    s / n
end

function _momentk(v::RealArray, k::Int, wv::AbstractWeights, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += (z ^ k) * wv[i]
    end
    s / sum(wv)
end

"""
    moment(v, k, [wv::AbstractWeights], m=mean(v))

Return the `k`th order central moment of a real-valued array `v`, optionally
specifying a weighting vector `wv` and a center `m`.
"""
function moment(v::RealArray, k::Int, m::Real)
    k == 2 ? _moment2(v, m) :
    k == 3 ? _moment3(v, m) :
    k == 4 ? _moment4(v, m) :
    _momentk(v, k, m)
end

function moment(v::RealArray, k::Int, wv::AbstractWeights, m::Real)
    k == 2 ? _moment2(v, wv, m) :
    k == 3 ? _moment3(v, wv, m) :
    k == 4 ? _moment4(v, wv, m) :
    _momentk(v, k, wv, m)
end

moment(v::RealArray, k::Int) = moment(v, k, mean(v))
function moment(v::RealArray, k::Int, wv::AbstractWeights)
    moment(v, k, wv, mean(v, wv))
end

##### Skewness and Kurtosis

# Skewness
# This is Type 1 definition according to Joanes and Gill (1998)
"""
    skewness(v, [wv::AbstractWeights], m=mean(v))

Compute the standardized skewness of a real-valued array `v`, optionally
specifying a weighting vector `wv` and a center `m`.
"""
function skewness(v::RealArray, m::Real)
    n = length(v)
    cm2 = 0.0   # empirical 2nd centered moment (variance)
    cm3 = 0.0   # empirical 3rd centered moment
    for i = 1:n
        @inbounds z = v[i] - m
        z2 = z * z

        cm2 += z2
        cm3 += z2 * z
    end
    cm3 /= n
    cm2 /= n
    return cm3 / sqrt(cm2 * cm2 * cm2)  # this is much faster than cm2^1.5
end

function skewness(v::RealArray, wv::AbstractWeights, m::Real)
    n = length(v)
    length(wv) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    cm2 = 0.0   # empirical 2nd centered moment (variance)
    cm3 = 0.0   # empirical 3rd centered moment

    @inbounds for i = 1:n
        x_i = v[i]
        w_i = wv[i]
        z = x_i - m
        z2w = z * z * w_i
        cm2 += z2w
        cm3 += z2w * z
    end
    sw = sum(wv)
    cm3 /= sw
    cm2 /= sw
    return cm3 / sqrt(cm2 * cm2 * cm2)  # this is much faster than cm2^1.5
end

skewness(v::RealArray) = skewness(v, mean(v))
skewness(v::RealArray, wv::AbstractWeights) = skewness(v, wv, mean(v, wv))

# (excessive) Kurtosis
# This is Type 1 definition according to Joanes and Gill (1998)
"""
    kurtosis(v, [wv::AbstractWeights], m=mean(v))

Compute the excess kurtosis of a real-valued array `v`, optionally
specifying a weighting vector `wv` and a center `m`.
"""
function kurtosis(v::RealArray, m::Real)
    n = length(v)
    cm2 = 0.0  # empirical 2nd centered moment (variance)
    cm4 = 0.0  # empirical 4th centered moment
    for i = 1:n
        @inbounds z = v[i] - m
        z2 = z * z
        cm2 += z2
        cm4 += z2 * z2
    end
    cm4 /= n
    cm2 /= n
    return (cm4 / (cm2 * cm2)) - 3.0
end

function kurtosis(v::RealArray, wv::AbstractWeights, m::Real)
    n = length(v)
    length(wv) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    cm2 = 0.0  # empirical 2nd centered moment (variance)
    cm4 = 0.0  # empirical 4th centered moment

    @inbounds for i = 1 : n
        x_i = v[i]
        w_i = wv[i]
        z = x_i - m
        z2 = z * z
        z2w = z2 * w_i
        cm2 += z2w
        cm4 += z2w * z2
    end
    sw = sum(wv)
    cm4 /= sw
    cm2 /= sw
    return (cm4 / (cm2 * cm2)) - 3.0
end

kurtosis(v::RealArray) = kurtosis(v, mean(v))
kurtosis(v::RealArray, wv::AbstractWeights) = kurtosis(v, wv, mean(v, wv))


mutable struct MoschopoulosParameters{T}
        θ₁::T
        C::T
        to_power::AbstractVector{T}
        γ::AbstractVector{T}
        ρ::T
        δ::AbstractVector{T}
end
function MoschopoulosParameters(α,θ)
    T = Base.promote_eltype(α,θ,[1.0])
    α = T.(α)
    θ = T.(θ)
    δ = T.([1])
    θ₁ = Base.minimum(θ)
    C = exp(sum(α .* log.(θ₁ ./ θ)))
    ρ = sum(α)
    to_power = (-θ₁ ./ θ) .+1
    γ = [sum(α .* to_power)] # gamma1
    return MoschopoulosParameters(θ₁,C,to_power,γ,ρ,δ)
end
Base.eltype(::MoschopoulosParameters{T}) where T = T

"""
    UnivariateGammaConvolution(α,θ)

Constructs a distribution that corresponds to the convolutions of Gamma(α[i],θ[i]) distributions.
The distribution can then be used through several methods,
following the Distributions.jl standard, to obtain pdf, cdf, random samples...

The pdf and cdf are handled by the Moshopoulos algorithm, and random samples by simply adding random gammas.
The code is type stable and handles any <:Real types, given by the parameters.

To fit the distribution, a loglikelyhood approach could be used. A more involved approach from Furman might be coded sometimes (but requires tanh-sinh integration and bigfloats...)

# Examples
```julia-repl
julia> dist = UnivariateGammaConvolution([1,0.5, 3.7],[4,2, 10]);
julia> sample = zeros(Float64,10);
julia> Random.rand!(dist,sample);
julia> pdf.((dist,),sample);
```
"""
struct UnivariateGammaConvolution{T<:Real} <: Distributions.ContinuousUnivariateDistribution where T
    α::AbstractVector{T}
    θ::AbstractVector{T}
    P::MoschopoulosParameters{T}
end

n(x::UnivariateGammaConvolution) = length(x.α)

function Base.show(io::IO, m::UnivariateGammaConvolution)
    println("Univariate Gamma Convolutions with parametrisation:")
    display([Text.(["α" "θ"]); [m.α m.θ]])
end

# Constructors :
UnivariateGammaConvolution(α::Real,θ::Real) = Distributions.Gamma(α,θ)
function UnivariateGammaConvolution(α::AbstractVector{T1},θ::AbstractVector{T2}) where {T1 <: Real, T2 <: Real}

    # Start by promoving alpha and theta to the same type :
    T = Base.promote_eltype(α,θ,[1.0]) # At least float.
    α = T.(α)
    θ = T.(θ)
    tol = eps(T) # Arbitrary fixed tolerence !

    order = sortperm(θ)
    n = length(θ)
    θ = θ[order]
    α = α[order]
    for i in 1:n
        if i < length(θ)
            for j in (i+1):n
                if j <= length(θ)
                    if abs(θ[i] - θ[j]) <= tol
                        new_α = α[i]+α[j]
                        new_θ = (α[i]*θ[i]+α[j]*θ[j])/(new_α)
                        α[i] = new_α
                        θ[i] = new_θ
                        deleteat!(α,j)
                        deleteat!(θ,j)
                        j = j-1
                    end
                end
            end
        end
    end
    to_be_kept = (α .> tol) .& (θ .> tol) .* (α.*θ .> tol)
    α = α[to_be_kept]
    θ = θ[to_be_kept]

    if length(α) == 1
        return Distributions.Gamma(α[1],θ[1])
    end
    return UnivariateGammaConvolution(α,θ,MoschopoulosParameters(α,θ))
end

#### Support
Distributions.@distr_support UnivariateGammaConvolution 0.0 Inf

#### Conversions
Base.eltype(d::UnivariateGammaConvolution) = typeof(d.α[1])
Base.convert(::Type{UnivariateGammaConvolution{T}}, d::UnivariateGammaConvolution{S}) where {T <: Real, S <: Real} = UnivariateGammaConvolution(T.(d.α), T.(d.θ))

#### parameters
shapes(d::UnivariateGammaConvolution) = d.α
scales(d::UnivariateGammaConvolution) = d.θ
rates(d::UnivariateGammaConvolution) = 1 / d.θ

#### Statistics
mean(d::UnivariateGammaConvolution) = sum(d.α .* d.θ)
var(d::UnivariateGammaConvolution) = sum(d.α .* d.θ .^ 2)

#### Sampling
function Base.rand(rng::Distributions.AbstractRNG,d::UnivariateGammaConvolution)
    sum(rand.(rng,Distributions.Gamma.(d.α,d.θ)))
end

#### Characteristic functions: pdf, cdf, mgf, cf

# Moshopoulos algorithm for pdf and cdf.
function MoschopoulosAlgo!(d::UnivariateGammaConvolution,x::Real, which)

    @assert(which in ["pdf","cdf"], "which should be etiher pdf or cdf")
    T = Base.eltype(d)
    atol = eps(T)
    rtol = T(0)
    entry_type = typeof(x)
    x = T(x)
    if x < T(0)
        return T(0)
    end
    k = 0
    out = T(0)
    while true
        if length(d.P.δ) < k+1
            # then compute the new deltas:
            n = length(d.P.δ)
            for k in n:(k+1)
                pushfirst!(d.P.γ, sum(d.α .* d.P.to_power .^ (k+1)))
                push!(d.P.δ,sum(d.P.γ[2:end] .* d.P.δ)/(k))
            end
        end
        dist = Distributions.Gamma(T(d.P.ρ + k),T(d.P.θ₁))
        if which == "pdf"
            step = d.P.δ[k+1] * Distributions.pdf(dist,x)
        elseif which == "cdf"
            step = d.P.δ[k+1] * Distributions.cdf(dist,x)
        end
        # if ((!isfinite(step)) & (x > T(0)))
        #     print("x = ", x)
        #     print(d)
        #     error("mince")
        # end
        @assert(!((!isfinite(step)) & (x > T(0))),"inf or nan append, the algorithm did not converge for x = $x")
        out += step
        if isapprox(step,T(0),atol=atol,rtol=rtol)
            break
        end
        k = k+1
    end
    out *= d.P.C
    return entry_type(out)
end

Distributions.pdf(d::UnivariateGammaConvolution,x::Real) = MoschopoulosAlgo!(d,x,"pdf")
Distributions.cdf(d::UnivariateGammaConvolution, x::Real) = MoschopoulosAlgo!(d,x,"cdf")
Distributions.logpdf(d::UnivariateGammaConvolution, x::Real) = log(pdf(d,x))
Distributions.mgf(d::UnivariateGammaConvolution, t::Real) = prod((1 - t .* d.θ) .^ (-d.α))
Distributions.cf(d::UnivariateGammaConvolution, t::Real) = prod((1 - (im * t) .* d.θ) .^ (-d.α))


using Distributions, LinearAlgebra, ForwardDiff, Seaborn, Pandas

# Evaluates target distribution π(q)
function π_q(q)
    q₁ = q[1]
    q₂ = q[2]
    in_bound = (-π/2 < q₁ < π/2) & (-π/2 < q₂ < π/2)
    return in_bound*(2/π * exp(-2*(q₁^2 + q₂^2)) + (sin(q₁*q₂)*sin(q₁)*cos(q₂))^2)
end

# Evaluates conditional probability π(p|q)
function π_pq(p,q)
    return pdf(MvNormal([0,0], 1.0*Matrix(I, 2, 2)), p)
end

# Evaluates Kinetic enery K(p,q)
function K(p,q)
    return -log(π_pq(p,q))
end

# Evaluates Potential energy V(q)
function V(q)
    return -log(π_q(q))
end

# Compute Hamiltonian at q,p
function H(p,q)
    return K(p,q) + V(q)
end

# Gets a trajectory of length L with n steps
function get_trajectory(q,p,L,n,dV_dq)

    # Step size
    ϵ = L/n

    q_trajectory = zeros(n,2)
    p_trajectory = zeros(n,2)

    q_trajectory[1,:] = q
    p_trajectory[1,:] = p

    for j=2:n
        ψ = p_trajectory[j-1,:] - ϵ/2 * dV_dq(q_trajectory[j-1,:])
        q_trajectory[j,:] = q_trajectory[j-1,:] + ϵ*ψ
        p_trajectory[j,:] = ψ - ϵ/2*dV_dq(q_trajectory[j,:])
    end

    return q_trajectory, p_trajectory
end

function hmc_alg(q₀, p₀)
    # Number of samples to return
    n = 25000

    q_samples = zeros(n,2)
    p_samples = zeros(n,2)
    q_samples[1,:] = q₀
    p_samples[1,:] = p₀

    trajectory_length = 1
    trajectory_partitions = 5

    # Gradient of V
    dV_dq(x) = ForwardDiff.gradient(V, x)

    for j=2:n

        q = q_samples[j-1,:]
        p = rand(MvNormal([0,0],1.0*Matrix(I, 2, 2)))
        q_proposal_trajectory, p_proposal_trajectory =
        get_trajectory(q, p, trajectory_length, trajectory_partitions, dV_dq)

        # Proposal
        q_L = q_proposal_trajectory[lastindex(q_proposal_trajectory)÷2,:]
        p_L = p_proposal_trajectory[lastindex(p_proposal_trajectory)÷2,:]

        # Compute acceptance probabilty
        α = min(1, exp(H(p,q) - H(-p_L,q_L)))

        if rand() < α
            q_samples[j,:] = q_L
            p_samples[j,:] = -p_L
        else
            q_samples[j,:] = q
            p_samples[j,:] = p
        end

    end

    return q_samples
end

samples = hmc_alg([0,0],[1,1])

using PyCall
using Plots
@pyimport seaborn as sns
pygui(true)
data = Pandas.DataFrame(Dict(:q₁=>samples[:,1], :q₂=>samples[:,2]))
rc("axes",facecolor = "#f7ece2")
rc("figure",facecolor = "#f7ece2")
rc("savefig",facecolor = "#f7ece2")
sns.jointplot(x = "q₁", y = "q₂", data = data, s = 1, alpha = 0.18)

using SpecialFunctions
q_values = -π/2:0.01:π/2

# Plot of q₁ marginal
function q₁_marginal(q)
    return √(2/π)*erf(π/√2)*exp(-2q^2) +
    sin(q)^2*(-π*q+π*q^3+sin(π*q))/(4*q^3-4*q)
end

q₁_values = [q₁_marginal(q) for q in q_values]

Plots.plot(q_values, q₁_values, legend = false, grid = false, xlabel = L"q_1", ylabel = L"\pi_{q_1}(q_1)", bg = RGB(247/255, 236/255, 226/255))

Plots.savefig("hmc_ex_2_marginal_1_plot.svg")

# Plot of q₂ marginal
function q₂_marginal(q)
    return √(2/π)*erf(π/√2)*exp(-2q^2) +
    1/4*cos(q)^2*(π + (1-2*q^2)*sin(π*q)/(q^3-q))
end

q₂_values = [q₂_marginal(q) for q in q_values]

Plots.plot(q_values, q₂_values, legend = false, grid = false, xlabel = L"q_2", ylabel = L"\pi_{q_2}(q_2)", bg = RGB(247/255, 236/255, 226/255))
Plots.savefig("hmc_ex_2_marginal_2_plot.svg")


# Lets use black-box optimization to do regression!
#
# To run this, just do:
#   julia regression_via_optimization.jl
# from a command line after having installed BlackBoxOptim.
#
# The basic idea is that a regression problem is a problem of deciding
# which set of coefficients (betas) minimize an objective function
# that compares the dependent variable y to a model based on the betas and the
# independent variables x.
#
# When implementing this, x will be a matrix d*n of floats where d is the dimension
# of the problem and n is the number of examples/cases. Thus y is an array/vector
# 1*n of floats and the coefficients is an array/vector (d+1)*1 of coefficients
# where the first is the intercept (beta0) while the rest are the coefficients.
#
# We get different types of regression depending on which objective function
# we select. Common to many of them is that they first calculate the deviances
# between the model and the actual values the model should predict:
function discrepancies(betas, x, y)
  y' .- ( betas[1] + sum(broadcast(*, betas[2:end], x), 1) )
end

# Given this setup we can now create an objective function for
# Ordinare Least Squares (OLS) regression. This is actually just the L2 norm:
function ols_regression_objective(beta, x, y)
  norm(discrepancies(beta, x, y), 2)
end

# And the L1 norm gives Least Absolute Deviations (LAD) (aka Robust regression) regression:
function lad_regression_objective(beta, x, y)
  norm(discrepancies(beta, x, y), 1)
end

# And we can do LASSO and Ridge regression by adding a penalty on large coefficients, but
# note that these also take a lambda constant:
function regularized_regression_objective(lambda, beta, x, y, p = 2, q = 1)
  norm(discrepancies(beta, x, y), p) + lambda * norm(beta[2:end], q)
end

# By selecting p = 2 and q = 1 we get LASSO regression:
function lasso_regression_objective(lambda, beta, x, y)
  regularized_regression_objective(lambda, beta, x, y, 2, 1)
end

# and by selecting p = 2 and q = 2 we get Ridge regression:
function ridge_regression_objective(lambda, beta, x, y)
  regularized_regression_objective(lambda, beta, x, y, 2, 2)
end

# Ok, we're ready to do some regression. We need some data, lets start with
# a way to generate data from models specified as julia functions. Lets start
# with a simple model of 3 vars: X1 + 2*X2 - X3
function m1(x)
  x[1,:] + 2*x[2,:] - x[3,:]
end

# and we generate random input for it:
x1 = rand(3, 100)
# and calc the outputs:
y1 = m1(x1)

# We can now search for the coefficients that minimize the OLS objective
# with a black-box optimization search like so (we allow coefficients to have
# a min value of -5.0 and a max value of 5.0, and the search is for 4
# coefficients, one intercept and three for each of the values of x):
using BlackBoxOptim
olsresult = bboptimize(b -> ols_regression_objective(b, x1, y1);
  SearchRange = (-10.0, 10.0), NumDimensions = 4, MaxSteps = 2e4)
@show (best_fitness(olsresult), best_candidate(olsresult))

# But the really nice thing is that we can easily consider other objectives such as the LAD:
ladresult = bboptimize(b -> lad_regression_objective(b, x1, y1);
  SearchRange = (-10.0, 10.0), NumDimensions = 4, MaxSteps = 2e4)
@show (best_fitness(ladresult), best_candidate(ladresult))

# For regularized regression we can optimize for different values of lambda so
# create a wrapper function that handles this:
function regularized_opt(lambda, func, x, y, dims, its = 2e4)
  bboptimize((b) -> func(lambda, b, x, y);
    SearchRange = (-10.0, 10.0), NumDimensions = dims, MaxSteps = its)
end

lassores1 = regularized_opt(1, lasso_regression_objective, x1, y1, 4)
lassores2 = regularized_opt(2, lasso_regression_objective, x1, y1, 4)
lassores3 = regularized_opt(3, lasso_regression_objective, x1, y1, 4)

ridgeres1 = regularized_opt(1, ridge_regression_objective, x1, y1, 4)
ridgeres2 = regularized_opt(2, ridge_regression_objective, x1, y1, 4)
ridgeres3 = regularized_opt(3, ridge_regression_objective, x1, y1, 4)

# Now lets create some support functions for printing models nicely.
linear_terms(num) = [@sprintf(" * X%d", i) for i in 1:num]
squared_terms(num) = [@sprintf(" * X%d^2", i) for i in 1:num]
linsq_terms(num) = vcat(linear_terms(num), squared_terms(num))

function sprint_predicted_model(fitresult, terms = nothing, skipIfLower = 1e-5)
  bestfit = best_candidate(fitresult)
  if terms == nothing
    terms = linear_terms(length(bestfit)-1)
  end
  if length(terms) < length(bestfit)
    terms = vcat([""], terms) # Put an empty term first which corresponds to intercept
  end

  signstr(value) = (value < 0.0) ? (@sprintf(" - %.3f", abs(value))) : (@sprintf(" + %.3f", value))

  elems = Any[]
  first_push = true
  for i in 1:length(bestfit)
    if abs(bestfit[i]) > skipIfLower
      str = join([signstr(bestfit[i]), terms[i]])
      if first_push
        # Strip away leading sign since this it the first push
        push!(elems, str[3:end])
        first_push = false
      else
        push!(elems, str)
      end
    end
  end

  join(elems)
end

# Let's try a model which involves a squared terms:
#   X1 + 4.13*X2*X2 - 3.14*X3
function m2(x)
  x[1,:] + (4.13 * x[2,:].^2) - (3.14 * x[3,:])
end

# and we generate random input for it:
x2 = rand(3, 100)
# and calc the outputs:
y2 = m2(x2)

# Before we regress we need to encode our beliefs about the general
# structure of the model. Let's say we believe there are squared terms but we
# do not know which ones. So we add one squared term per independent variable:
x2m = zeros(3+3, 100)
x2m[1:3,:] = x2
x2m[4,:] = x2[1,:].^2
x2m[5,:] = x2[2,:].^2
x2m[6,:] = x2[3,:].^2

# With this we can fit models:
m2_olsres = bboptimize((b) -> ols_regression_objective(b, x2m, y2);
  SearchRange = (-10.0, 10.0), NumDimensions = 7, MaxSteps = 5e4)
m2_ladres = bboptimize((b) -> lad_regression_objective(b, x2m, y2);
  SearchRange = (-10.0, 10.0), NumDimensions = 7, MaxSteps = 5e4)
m2_lassores1 = regularized_opt(1, lasso_regression_objective, x2m, y2, 7, 5e4)
m2_lassores2 = regularized_opt(2, lasso_regression_objective, x2m, y2, 7, 5e4)
m2_lassores3 = regularized_opt(3, lasso_regression_objective, x2m, y2, 7, 5e4)
m2_ridgeres1 = regularized_opt(1, ridge_regression_objective, x2m, y2, 7, 5e4)
m2_ridgeres2 = regularized_opt(2, ridge_regression_objective, x2m, y2, 7, 5e4)
m2_ridgeres3 = regularized_opt(3, ridge_regression_objective, x2m, y2, 7, 5e4)

# And now lets print our models nicely so user can see the results...

println("Model1 = 1.000 * X1 + 2.000 * X2 - 1.000 * X3")
println("OLS best fit: ", sprint_predicted_model(olsresult))
println("LAD best fit: ", sprint_predicted_model(ladresult))
println("LASSO best fit, lambda = 1: ", sprint_predicted_model(lassores1))
println("LASSO best fit, lambda = 2: ", sprint_predicted_model(lassores2))
println("LASSO best fit, lambda = 3: ", sprint_predicted_model(lassores3))
println("Ridge best fit, lambda = 1: ", sprint_predicted_model(ridgeres1))
println("Ridge best fit, lambda = 2: ", sprint_predicted_model(ridgeres2))
println("Ridge best fit, lambda = 3: ", sprint_predicted_model(ridgeres3))

terms = linsq_terms(3)
println("")
println("Model2 = 1.000 * X1 - 3.140 * X3 + 4.130 * X2^2 ")
println("OLS best fit: ", sprint_predicted_model(m2_olsres, terms))
println("LAD best fit: ", sprint_predicted_model(m2_ladres, terms))
println("LASSO best fit, lambda = 1: ", sprint_predicted_model(m2_lassores1, terms))
println("LASSO best fit, lambda = 2: ", sprint_predicted_model(m2_lassores2, terms))
println("LASSO best fit, lambda = 3: ", sprint_predicted_model(m2_lassores3, terms))
println("Ridge best fit, lambda = 1: ", sprint_predicted_model(m2_ridgeres1, terms))
println("Ridge best fit, lambda = 2: ", sprint_predicted_model(m2_ridgeres2, terms))
println("Ridge best fit, lambda = 3: ", sprint_predicted_model(m2_ridgeres3, terms))

# Conclusion: With black-box optimization you can easily fit regression models
# from very different paradigms without having to implement very much code.


"""
CORDIC (COordinate Rotation DIgital Computer), a.k.a. Volder's algorithm, for computing `sin` and
`cos` functions efficiently.
"""
module CORDIC

export cordic

# Plagiarizing @njuffa's answer from https://stackoverflow.com/a/23097989/.
#
# These coefficients are obtained by Remez algorithm rather than from Taylor series.
function arctan(x::Float64)::Float64
    # argument reduction:
    #
    # arctan (-x) = -arctan(x)
    # arctan (1/x) = 1/2 * pi - arctan (x), when x > 0
    z = abs(x)
    a = (z > 1.0) ? 1.0 / z : z

    # evaluate minimax polynomial approximation
    s = a * a  # a ^ 2
    q = s * s  # a ^ 4
    o = q * q  # a ^ 8

    # use Estrin's scheme for low-order terms
    p = fma(fma(fma(-0x1.53e1d2a25ff34p-16, s, 0x1.d3b63dbb65af4p-13),
                q,
                fma(-0x1.312788dde0801p-10, s, 0x1.f9690c82492dbp-9)),
            o,
            fma(fma(-0x1.2cf5aabc7cef3p-7, s, 0x1.162b0b2a3bfcep-6),
                q,
                fma(-0x1.a7256feb6fc5cp-6, s, 0x1.171560ce4a483p-5)))

    # use Horner's scheme for high-order terms
    p = fma(fma(fma(fma(fma(fma(fma(fma(fma(fma(fma(fma(p, s,
            -0x1.4f44d841450e1p-5), s,
            +0x1.7ee3d3f36bb94p-5), s,
            -0x1.ad32ae04a9fd1p-5), s,
            +0x1.e17813d66954fp-5), s,
            -0x1.11089ca9a5bcdp-4), s,
            +0x1.3b12b2db51738p-4), s,
            -0x1.745d022f8dc5cp-4), s,
            +0x1.c71c709dfe927p-4), s,
            -0x1.2492491fa1744p-3), s,
            +0x1.99999999840d2p-3), s,
            -0x1.555555555544cp-2) * s, a, a)
    # back substitution based on argument reduction
    r = (z > 1.0) ? (0x1.921fb54442d18p+0 - p) : p
    copysign(r, x)
end

ANGLES = @. arctan(exp2(-(0:27)))
K_VALUES = cumprod(@. 1.0 / sqrt(1.0 + exp2(-(0:2:50))))

"Returns a pair of `(sin(beta), cos(beta))` using CORDIC."
function cordic(beta::Float64, n_iterations::Integer=25)::Tuple{Float64, Float64}
    if beta < -pi / 2.0
        return .-cordic(beta + pi, n_iterations)
    elseif beta > pi / 2.0
        return .-cordic(beta - pi, n_iterations)
    end

    v = [1.0, 0.0]
    angle = ANGLES[1]

    i = 0
    while i < n_iterations
        σ = sign(beta)
        if σ == 0.0
            break
        end

        x = v[1] - σ * ldexp(v[2], -i)
        y = σ * ldexp(v[1], -i) + v[2]
        v = [x, y]
        beta -= σ * angle
        i += 1

        # +1 because the iterating variable starts from 0 (2 ^ 0, 2 ^ (-1), ...).
        if i + 1 > length(ANGLES)
            angle /= 2.0
        else
            angle = ANGLES[i + 1]
        end
    end

    if i > 0
        v *= K_VALUES[min(i, length(K_VALUES))]
    end
    (v[2], v[1])
end

const flatten = Iterators.flatten
function smoke()
    X = range(-pi*4, stop=pi*4, length=800)
    println("arctan error: ", max(@. abs(arctan(X) - atan(X))...))
    println("cordic error: ", max(abs.(flatten(cordic.(X)) .- flatten(sincos.(X)))...))
end

end  # module


###########
# This gives an example of Gauss–Seidel
#
# It is _much_ more efficient to use BlockBandedMatrices.jl
###########

using BandedMatrices, LazyArrays, LinearAlgebra

function finitedifference_2d(n)
    h = 1/n
    D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2
    D_xx = kron(D², Eye(n))
    D_yy = kron(Eye(n), D²)
    D_xx + D_yy
end

function gaussseidel(L, U, b, x=copy(b), M=5)
    for _=1:M
        @view(x[1:end-1]) .= Mul(U , @view(x[2:end]))
        x[end] = 0
        x .= b .- x
        x .= Ldiv(L, x)
    end
    x
end

n = 400
Δt = (1/n^2)/4; Δ = finitedifference_2d(n); A = I - Δt*Δ  # 160k x 160k discretization

L = LowerTriangular(A)
U = UpperTriangular(@view A[1:end-1,2:end])

b = randn(size(A,1));
@time u = A\b; # 6s
x = copy(b)
@time gaussseidel(L,U, b, x, 20) # 1.6s
norm(x - u) # 6*10^(-10)


function _lobatto_legendre_fast(s, T)
    c, b = FastGaussQuadrature.gausslobatto(s)
    shift!(b,c)
    QuadratureRule(2s-2, c, b, T)
end

"""
Lobatto-Legendre quadrature.
"""
function LobattoLegendreQuadrature(::Type{T}, s::Integer; IT=BigFloat, fast=false) where {T}
    if s == 1
        throw(ErrorException("Lobatto quadrature is not defined for one stage."))
    end

    if fast
        return _lobatto_legendre_fast(s, T)
    end

    D(k) = Polynomials.derivative(Polynomial(IT[0, 1, -1])^(k-1), k-2)
    P(k,x) = Polynomials.derivative(Polynomial(IT[-1, 0, 1])^k, k)(x) / factorial(k) / 2^k
    c = sort(IT.(Polynomials.roots(D(s)))); c[begin] = 0; c[end] = 1;
    b = [ 1 / ( s*(s-1) * P(s-1, 2c[i] - 1)^2 ) for i in 1:s ]
    return QuadratureRule(2s-2, c, b, T)
end

LobattoLegendreQuadrature(s; kwargs...) = LobattoLegendreQuadrature(Float64, s; kwargs...)


"""
    f₁(r, θ, ϕ, ψ)

Define a coordinate chart from 3-hyperspherical coordinates into Cartesian coordinates in ℝ⁴
with the given radius `r`, and angles: `θ`, `ϕ` and `ψ`. f₁: S³ ⟼ ℝ⁴.
"""
function f₁(r::Real, θ::Real, ϕ::Real, ψ::Real)
    x₁ = r * cos(θ)
    x₂ = r * sin(θ) * cos(ϕ)
    x₃ = r * sin(θ) * sin(ϕ) * cos(ψ)
    x₄ = r * sin(θ) * sin(ϕ) * sin(ψ)
    ℝ⁴(x₁, x₂, x₃, x₄)
end

"""
    f₂(α, ϕ₁, ϕ₂)

Define a coordinate chart from Hopf coordinates on 3-sphere into Cartesian coordinates in ℝ⁴
with the given angles `α`, `ϕ₁` and `ϕ₂`. f₂: S³ ⟼ ℝ⁴.
"""
function f₂(α::Real, ϕ₁::Real, ϕ₂::Real)
    @assert(0 < α < π / 2, "α ∈ (0, π/2)")
    @assert(0 ≤ ϕ₁ ≤ 2π, "ϕ₁ ∈ [0, 2π]")
    @assert(0 ≤ ϕ₂ ≤ 2π, "ϕ₂ ∈ [0, 2π]")
    x₁ = cos(ϕ₁) * sin(α)
    x₂ = sin(ϕ₁) * sin(α)
    x₃ = cos(ϕ₂) * cos(α)
    x₄ = sin(ϕ₂) * cos(α)
    ℝ⁴(x₁, x₂, x₃, x₄)
end

"""
    f₃(α, ϕ₁, ϕ₂)

Define a coordinate chart from Hopf coordinates on 3-sphere into Cartesian coordinates in ℝ⁴
with the given angles `α`, `ϕ₁` and `ϕ₂`. f₃: S³ ⟼ ℝ⁴.
"""
function f₃(α::Real, ϕ₁::Real, ϕ₂::Real)
    @assert(0 < α < π / 2, "α ∈ (0, π/2)")
    @assert(0 ≤ ϕ₁ ≤ 2π, "ϕ₁ ∈ [0, 2π]")
    @assert(0 ≤ ϕ₂ ≤ 4π, "ϕ₂ ∈ [0, 4π]")
    x₁ = cos((ϕ₁ + ϕ₂) / 2) * sin(α)
    x₂ = sin((ϕ₁ + ϕ₂) / 2) * sin(α)
    x₃ = cos((ϕ₂ - ϕ₁) / 2) * cos(α)
    x₄ = sin((ϕ₂ - ϕ₁) / 2) * cos(α)
    ℝ⁴(x₁, x₂, x₃, x₄)
end


clenshawcurtis_plan(μ) = length(μ) > 1 ? FFTW.plan_r2r!(μ, FFTW.REDFT00) : ones(μ)'

"""
Compute nodes and weights of the Clenshaw—Curtis quadrature rule with a Jacobi weight.
"""
clenshawcurtis{T<:AbstractFloat}(N::Int,α::T,β::T) = clenshawcurtis(N,α,β,clenshawcurtis_plan(zeros(T,N)))
clenshawcurtis{T<:AbstractFloat}(N::Int,α::T,β::T,plan) = T[cospi(k/(N-one(T))) for k=0:N-1],clenshawcurtisweights(N,α,β,plan)

"""
Compute weights of the Clenshaw—Curtis quadrature rule with a Jacobi weight.
"""
clenshawcurtisweights{T<:AbstractFloat}(N::Int,α::T,β::T) = clenshawcurtisweights(N,α,β,clenshawcurtis_plan(zeros(T,N)))
function clenshawcurtisweights{T<:AbstractFloat}(N::Int,α::T,β::T,plan)
    μ = chebyshevjacobimoments1(N,α,β)
    scale!(μ,inv(N-one(T)))
    plan*μ
    μ[1]/=2;μ[N]/=2
    return μ
end

# Chebyshev-T coefficients to values at Clenshaw-Curtis nodes

applyTN_plan(x) = length(x) > 1 ? FFTW.plan_r2r!(x, FFTW.REDFT00) : ones(x)'

applyTN!{T<:AbstractFloat}(x::Vector{T}) = applyTN!(x,applyTN_plan(x))
function applyTN!{T<:AbstractFloat}(x::Vector{T},plan)
    x[1] *= 2; x[end] *=2
    plan*x
    scale!(x,half(T))
end
applyTN{T<:AbstractFloat}(x::Vector{T},plan) = applyTN!(copy(x),plan)
applyTN{T<:AbstractFloat}(x::Vector{T}) = applyTN!(copy(x))

# Values at Clenshaw-Curtis nodes to Chebyshev-T coefficients

applyTNinv_plan(x) = length(x) > 1 ? FFTW.plan_r2r!(x, FFTW.REDFT00) : ones(x)'

applyTNinv!{T<:AbstractFloat}(x::Vector{T}) = applyTNinv!(x,applyTNinv_plan(x))
function applyTNinv!{T<:AbstractFloat}(x::Vector{T},plan)
    plan*x
    x[1] /= 2;x[end] /= 2
    scale!(x,inv(length(x)-one(T)))
end
applyTNinv{T<:AbstractFloat}(x::Vector{T},plan) = applyTNinv!(copy(x),plan)
applyTNinv{T<:AbstractFloat}(x::Vector{T}) = applyTNinv!(copy(x))

# sin(nθ) coefficients to values at Clenshaw-Curtis nodes except ±1

applyUN_plan(x) = length(x) > 0 ? FFTW.plan_r2r!(x, FFTW.RODFT00) : ones(x)'

applyUN!{T<:AbstractFloat}(x::AbstractVector{T}) = applyUN!(x,applyUN_plan(x))
function applyUN!{T<:AbstractFloat}(x::AbstractVector{T},plan)
    plan*x
    scale!(x,half(T))
end
applyUN{T<:AbstractFloat}(x::AbstractVector{T},plan) = applyUN!(copy(x),plan)
applyUN{T<:AbstractFloat}(x::AbstractVector{T}) = applyUN!(copy(x))


"""
**Inversedecay constructor**
```julia
    Inversedecay(; t0::Real=1, κ::Real=0.51)
```

Algorithm:
```math
\\Delta x_t = (t+t_0)^{-\\kappa}g_t
```

Simple learning rate respecting the [Robbins-Monro conditions](https://en.wikipedia.org/wiki/Stochastic_approximation#Robbins%E2%80%93Monro_algorithm) for κ ∈ (0.5,1]
"""
mutable struct Inversedecay <: Optimizer
    opt_type::String
    t::Int64
    t₀::Float64
    κ::Float64
end

function Inversedecay(; t0::Real=1, κ::Real=0.51)
    @assert t0 > 0.0 "t0 must be greater than 0"
    @assert (κ > 0.5 || κ <= 1.0) "κ argument is in (0.5,1]"

    Inversedecay("Inversedecay", 0, t0, κ)
end

params(opt::Inversedecay) = "t₀=$(opt.t₀), κ=$(opt.κ)"

function update(opt::Inversedecay, g_t::AbstractArray{T}) where {T<:Real}
    # update timestep
    opt.t += 1
    return (opt.t₀+opt.t)^(-opt.κ) * g_t
end


export randomInv, RandomVectorSketch

"""
Solve TAx=Tb with T a random projector
"""
function randomInv(
  stp::AbstractStopping;
  is_zero_start::Bool = false,
  k = 50,
  random_func = random_matrix_1,
  kwargs...,
)
  A, b = get_matrix(stp.pb), get_vector(stp.pb)
  m = size(A, 1)

  start!(stp)
  k = Int(floor(m / 8))
  T = random_func(k, m)
  xk = T * A \ T * b

  return update_and_stop!(stp, x = xk)
end

"""
Random vector sketch

Section 3.2 in Gower, R. M., & Richtárik, P. (2015).
Randomized iterative methods for linear systems.
SIAM Journal on Matrix Analysis and Applications, 36(4), 1660-1690.
"""
function RandomVectorSketch(
  stp::AbstractStopping;
  is_zero_start::Bool = false,
  random_func = random_matrix_1,
  kwargs...,
)
  A, b = get_matrix(stp.pb), get_vector(stp.pb)
  m, n = size(A)
  x0 = stp.current_state.x
  xk = x0
  stp.current_state.res = A * xk - b
  OK = start!(stp)

  while !OK
    s = vec(random_func(1, m)')
    As = A' * s
    xk = As == 0 ? x0 : x0 - dot(s, stp.current_state.res) / dot(As, As) * As

    update!(stp.current_state, x = xk)
    stp.current_state.res = A * xk - b
    OK = stop!(stp)
    x0 = xk
  end

  return stp
end


function stencil1d(m::Int)

    ic = Int(ceil(m/2.0))
    M = zeros((m, m))
    for i=1:m
        for j=1:m
            M[i, j] = (i - ic)^(j-1) / factorial(j-1)
        end
    end
    return inv(M)
end

function periodic_projection(N, m)

    js = [N-m+1:N...]
    append!(js, [1:N...], [1:m...])

    P = sparse([1:N...], [m+1:m+N...], ones(N), N, N+2*m)
    PT = sparse([1:N+2*m...], js, ones(N+2*m), N+2*m, N)

    return P, PT
end

function periodic_derivative_1d(n, m, N)

    Nb = N + 2*m
    D = sparse(zeros(Float64, (Nb, Nb)))
    stencil = stencil1d(m)[n+1,:]
    ic = Int(ceil(m/2.0))

    for i=1:m
        k = i - ic
        vec = stencil[i] * ones(Nb - abs(k))
        D += spdiagm(k => vec)
    end

    P, PT = periodic_projection(N, m)
    return P * D * PT
end

function periodic_derivative(nx, ny, mx, my, Nx, Ny, dx, dy)

    Dx = periodic_derivative_1d(nx, mx, Nx) / dx^nx
    Dy = periodic_derivative_1d(ny, my, Ny) / dy^ny
    return kron(Dy, Dx)
end


module firstTask

# Global variable
goldenRatio = (1 + sqrt(5)) / 2

"""
Function optimizes single variable function using golden section method (to find local minimum value in provided range)
"""
function goldenSection(f, a, b, tol = 1e-8)
  # Make sure that entered range is correct
  a = min(a, b)
  b = max(a, b)

  while abs(a - b) >= tol
    # Calculate next 2 points
    x1 = b - (b - a) / goldenRatio
    x2 = a + (b - a) / goldenRatio
    # Compare function values from calculated points and determine where minimum is located
    if f(x1) < f(x2)
      b = x2
    else
      a = x1
    end
  end
  # Get the approximate value and return it
  result = (a + b) / 2
  resultY = f(result)
  return (result, resultY)
end

export goldenSection

end # module


###############################################################################
## Logistic Problem
#    min_x  1/m * ∑_{i=1}^m log(1 + exp(-y_i dot(A_i, x))) + g(reg, x)
#
#    x : R^n
#    y : R^m observations, *values -1 and 1*
#    A : mxn matrix of samples

mutable struct LogisticPb{Tr, Tm, Txman} <: CompositeProblem
    A::Matrix{Float64}
    y::Vector{Float64}
    regularizer::Tr
    n::Int64
    x0::Txman
    M_x0::Tm
    λ₂::Float64
    function LogisticPb(
            A::Matrix{Float64},
            y::Vector{Float64},
            regularizer::Tr,
            n::Int64,
            x0::Txman,
            M_x0::Tm,
            λ₂::Float64
        ) where {Tr, Tm, Txman}
        @assert Set(y) ⊆ Set([-1.0, 1.0]) "Logistic rhs vector shoudl take values -1.0, 1.0, here: $(Set(y))."
        return new{Tr, Tm, Txman}(A, y, regularizer, n, x0, M_x0, λ₂)
    end
end

problem_dimension(pb::LogisticPb) = pb.n

"""
    logsig(t)

Compute the logarithm of sigmoid `-log(1+exp(-t))` with higher precision than plain
implementation.

Reference:
- F. Pedragosa's blog post http://fa.bianp.net/blog/2019/evaluate_logistic/
"""
@inline function logsig(t)
    if t < -33.3
        return t
    elseif t <= -18
        return t - exp(t)
    elseif t <= 37
        return -log1p(exp(-t))
    else
        return -exp(-t)
    end
end

## f
# 0th order
function f(pb::LogisticPb, x)
    m = size(pb.A, 1)

    Ax = pb.A * x
    fval = 0.0
    @inbounds @simd for i in 1:m
        fval -= logsig(pb.y[i] * Ax[i])
    end

    return fval / m + 0.5 * pb.λ₂ * norm(x, 2)^2
end

# 1st order
σ(x) = 1/(1+exp(-x))
function ∇f!(pb::LogisticPb, res, x)
    m = size(pb.A, 1)

    σyAx = pb.A * x
    σyAx .*= -pb.y
    σyAx .= σ.(σyAx)
    res .= transpose(pb.A) * (σyAx .* pb.y)
    res ./= -m
    res .+= pb.λ₂ .* x

    return res
end

# 2nd order
∇σ(x) = σ(x) * σ(-x)
function ∇²f_h!(pb::LogisticPb, res, x, h)
    m = size(pb.A, 1)

    yAx = -pb.y .* (pb.A * x)
    Ah = pb.A * h

    res .= transpose(pb.A) * (Ah .* ∇σ.(yAx))
    res ./= m
    res .+= pb.λ₂ .* h

    return res
end

# conditioning
function get_gradlips(pb::LogisticPb)
    m = size(pb.A, 1)
    return opnorm(pb.A)^2 / m + pb.λ₂
end
# get_μ_cvx(pb::LogisticPb) = (svdvals(pb.A)[end])^2


# General vector utilities that may have application outside of MkCell.

"Calculate angle between two vectors (in radians)"
ang(x,y) = acos(dot(x,y)/(norm(x)*norm(y)))

"""Get lattice a, b, c from cell"""
cell_abc(cell) = norm(cell[1,:]), norm(cell[2,:]), norm(cell[3,:])

"""Get cell angles α, β, γ (in degrees) from a cell."""
cell_angles(cell) = tuple((ang(cell[p[1],:], cell[p[2],:])*180.0/π for p in [(2,3), (1,3), (1,2)])...)


"""
`LinearKernel([ρ=1.0,[c=0.0]])`
The linear kernel is a Mercer kernel given by
```
    κ(x,y) = ρ²xᵀy + c
```
Where `c` is a real number
"""
struct LinearKernel{Tr, Tc<:Real} <: Kernel{Tr}
    transform::Tr
    c::Tc
end

function LinearKernel(ρ::T=1.0, c::Real=zero(T)) where {T<:Real}
    LinearKernel(ScaleTransform(ρ), c)
end

function LinearKernel(ρ::AbstractVector{T}, c::Real=zero(T)) where {T<:Real}
    LinearKernel(ARDTransform(ρ), c)
end

LinearKernel(t::Transform) = LinearKernel(t, 0.0)

params(k::LinearKernel) = (params(transform(k)),k.c)
opt_params(k::LinearKernel) = (opt_params(transform(k)),k.c)

@inline kappa(κ::LinearKernel, xᵀy::T) where {T<:Real} = xᵀy + κ.c

metric(::LinearKernel) = DotProduct()

"""
`PolynomialKernel([ρ=1.0[,d=2.0[,c=0.0]]])`
The polynomial kernel is a Mercer kernel given by
```
    κ(x,y) = (ρ²xᵀy + c)^d
```
Where `c` is a real number, and `d` is a shape parameter bigger than 1
"""
struct PolynomialKernel{Tr,Tc<:Real,Td<:Real} <: Kernel{Tr}
    transform::Tr
    d::Td
    c::Tc
    function PolynomialKernel{Tr, Tc, Td}(transform::Tr, d::Td, c::Tc) where {Tr<:Transform, Td<:Real, Tc<:Real}
        @check_args(PolynomialKernel, d, d >= one(Td), "d >= 1")
        return new{Tr, Td, Tc}(transform,d, c)
    end
end

function PolynomialKernel(ρ::Real=1.0, d::Td=2.0, c::Real=zero(Td)) where {Td<:Real}
    PolynomialKernel(ScaleTransform(ρ), d, c)
end

function PolynomialKernel(ρ::AbstractVector{T}, d::Real=2.0, c::Real=zero(T₁)) where {T<:Real}
    PolynomialKernel(ARDTransform(ρ), d, c)
end

function PolynomialKernel(t::Tr, d::Td=2.0, c::Tc=zero(eltype(Td))) where {Tr<:Transform, Td<:Real, Tc<:Real}
    PolynomialKernel{Tr, Tc, Td}(t, d, c)
end

params(k::PolynomialKernel) = (params(transform(k)),k.d,k.c)
opt_params(k::PolynomialKernel) = (opt_params(transform(k)),k.d,k.c)

@inline kappa(κ::PolynomialKernel, xᵀy::T) where {T<:Real} = (xᵀy + κ.c)^(κ.d)

metric(::PolynomialKernel) = DotProduct()


"""
Module: SymbolicDiff (Symbolic Operation for Arithmetic)
"""

"""
seval(f, dvar, env, cache)
Return the first derivative of expr f with respect to dvar
"""

function seval(f, dvar::Symbol)
    seval(f, dvar, globalenv, SymbolicCache())
end

function seval(f, dvar::Symbol, env::SymbolicEnv)
    seval(f, dvar, env, SymbolicCache())
end

function seval(f, dvar::Symbol, cache::SymbolicCache)
    seval(f, dvar, globalenv, cache)
end

###

function seval(f, dvar::SymbolicVariable{Tv}) where Tv
    seval(f, dvar.var)
end

function seval(f, dvar::SymbolicVariable{Tv}, env::SymbolicEnv) where Tv
    seval(f, dvar.var, env)
end

function seval(f, dvar::SymbolicVariable{Tv}, cache::SymbolicCache) where Tv
    seval(f, dvar.var, cache)
end

function seval(f, dvar::SymbolicVariable{Tv}, env::SymbolicEnv, cache::SymbolicCache) where Tv
    seval(f, dvar.var, env, cache)
end

###

function seval(f::SymbolicValue{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    Tv(0)
end

function seval(f::SymbolicVariable{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    f.var == dvar ? 1 : 0
end

function seval(f::AbstractNumberSymbolic{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    (dvar in f.params) || return 0
    get(cache, (f,dvar)) do
        retval = _eval(Val(f.op), f, dvar, env, cache)
        cache[(f,dvar)] = retval
    end
end

"""
_eval(::Val{xx}, dvar, f, env, cache)

Dispached function to evaluate the first derivative of f
"""

function _eval(::Val{:+}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    args = [seval(x, dvar, env, cache) for x = f.args]
    +(args...)
end

function _eval(::Val{:-}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    args = [seval(x, dvar, env, cache) for x = f.args]
    -(args...)
end

function _eval(::Val{:*}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    args = [seval(x, env, cache) for x = f.args]
    dargs = [seval(x, dvar, env, cache) for x = f.args]
    ret = dargs[1]
    s = args[1]
    for i = 2:length(args)
        ret *= args[i]
        ret += s * dargs[i]
        (i == length(args)) && break
        s *= args[i]
    end
    ret
end

function _eval(::Val{:/}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x,y = [seval(x, env, cache) for x = f.args]
    dx,dy = [seval(x, dvar, env, cache) for x = f.args]
    (dx * y - x * dy) / y^2
end

function _eval(::Val{:^}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x,y = [seval(x, env, cache) for x = f.args]
    dx,dy = [seval(x, dvar, env, cache) for x = f.args]
    x^(y-1) * (x * log(x) * dy + y * dx)
end

function _eval(::Val{:exp}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x, = [seval(x, env, cache) for x = f.args]
    dx, = [seval(x, dvar, env, cache) for x = f.args]
    exp(x) * dx
end

function _eval(::Val{:log}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x, = [seval(x, env, cache) for x = f.args]
    dx, = [seval(x, dvar, env, cache) for x = f.args]
    dx / x
end

function _eval(::Val{:sqrt}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x, = [seval(x, env, cache) for x = f.args]
    dx, = [seval(x, dvar, env, cache) for x = f.args]
    dx /(2 * sqrt(x))
end

function _eval(::Val{:sum}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    dx, = [seval(x, dvar, env, cache) for x = f.args]
    sum(dx)
end

function _eval(::Val{:dot}, f::SymbolicExpression{Tv}, dvar::Symbol, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x,y = [seval(x, env, cache) for x = f.args]
    dx,dy = [seval(x, dvar, env, cache) for x = f.args]
    dot(x,dy) + dot(dx,y)
end


# DBSCAN Clustering
#
#   References:
#
#       Martin Ester, Hans-peter Kriegel, Jörg S, and Xiaowei Xu
#       A density-based algorithm for discovering clusters
#       in large spatial databases with noise. 1996.
#

type DbscanResult <: ClusteringResult
    seeds::Vector{Int}          # starting points of clusters, size (k,)
    assignments::Vector{Int}    # assignments, size (n,)
    counts::Vector{Int}         # number of points in each cluster, size (k,)
end

## main algorithm

function dbscan{T<:Real}(D::DenseMatrix{T}, eps::Real, minpts::Int)
    # check arguments
    n = size(D, 1)
    size(D, 2) == n || error("D must be a square matrix.")
    n >= 2 || error("There must be at least two points.")
    eps > 0 || error("eps must be a positive real value.")
    minpts >= 1 || error("minpts must be a positive integer.")

    # invoke core algorithm
    _dbscan(D, convert(T, eps), minpts, 1:n)
end

function _dbscan{T<:Real}(D::DenseMatrix{T}, eps::T, minpts::Int, visitseq::AbstractVector{Int})
    n = size(D, 1)

    # prepare
    seeds = Int[]
    counts = Int[]
    assignments = zeros(Int, n)
    visited = zeros(Bool, n)
    k = 0

    # main loop
    for p in visitseq
        if assignments[p] == 0 && !visited[p]
            visited[p] = true
            nbs = _dbs_region_query(D, p, eps)
            if length(nbs) >= minpts
                k += 1
                cnt = _dbs_expand_cluster!(D, k, p, nbs, eps, minpts, assignments, visited)
                push!(seeds, p)
                push!(counts, cnt)
            end
        end
    end

    # make output
    return DbscanResult(seeds, assignments, counts)
end

## key steps

function _dbs_region_query{T<:Real}(D::DenseMatrix{T}, p::Int, eps::T)
    n = size(D,1)
    nbs = Int[]
    dists = view(D,:,p)
    for i = 1:n
        @inbounds if dists[i] < eps
            push!(nbs, i)
        end
    end
    return nbs::Vector{Int}
end

function _dbs_expand_cluster!{T<:Real}(D::DenseMatrix{T},           # distance matrix
                                       k::Int,                      # the index of current cluster
                                       p::Int,                      # the index of seeding point
                                       nbs::Vector{Int},            # eps-neighborhood of p
                                       eps::T,                      # radius of neighborhood
                                       minpts::Int,                 # minimum number of neighbors of a density point
                                       assignments::Vector{Int},    # assignment vector
                                       visited::Vector{Bool})       # visited indicators
    assignments[p] = k
    cnt = 1
    while !isempty(nbs)
        q = shift!(nbs)
        if !visited[q]
            visited[q] = true
            qnbs = _dbs_region_query(D, q, eps)
            if length(qnbs) > minpts
                for x in qnbs
                    if assignments[x] == 0
                        push!(nbs, x)
                    end
                end
            end
        end
        if assignments[q] == 0
            assignments[q] = k
            cnt += 1
        end
    end
    return cnt
end


# ---
# title: 458. Poor Pigs
# id: problem458
# author: Indigo
# date: 2021-02-02
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/poor-pigs/description/>
# hidden: true
# ---
#
# There are `buckets` buckets of liquid, where **exactly one** of the buckets is
# poisonous. To figure out which one is poisonous, you feed some number of
# (poor) pigs the liquid to see whether they will die or not. Unfortunately, you
# only have `minutesToTest` minutes to determine which bucket is poisonous.
#
# You can feed the pigs according to these steps:
#
#   1. Choose some live pigs to feed.
#   2. For each pig, choose which buckets to feed it. The pig will consume all the chosen buckets simultaneously and will take no time.
#   3. Wait for `minutesToDie` minutes. You may **not** feed any other pigs during this time.
#   4. After `minutesToDie` minutes have passed, any pigs that have been fed the poisonous bucket will die, and all others will survive.
#   5. Repeat this process until you run out of time.
#
# Given `buckets`, `minutesToDie`, and `minutesToTest`, return _the **minimum**
# number of pigs needed to figure out which bucket is poisonous within the
# allotted time_.
#
#
#
# **Example 1:**
#
#
#
#     Input: buckets = 1000, minutesToDie = 15, minutesToTest = 60
#     Output: 5
#
#
# **Example 2:**
#
#
#
#     Input: buckets = 4, minutesToDie = 15, minutesToTest = 15
#     Output: 2
#
#
# **Example 3:**
#
#
#
#     Input: buckets = 4, minutesToDie = 15, minutesToTest = 30
#     Output: 2
#
#
#
#
# **Constraints:**
#
#   * `1 <= buckets <= 1000`
#   * `1 <= minutesToDie <= minutesToTest <= 100`
#
#
## @lc code=start
using LeetCode

function poor_pigs(buckets::Int, minutesToDie::Int, minutesToTest::Int)
    state = minutesToTest ÷ minutesToDie + 1
    return ceil(Int, log(buckets) / log(state))
end
## @lc code=end


__precompile__()

module SparseRelaxation

export gauss_seidel, gauss_seidel!, weighted_jacobi!, weighted_jacobi

# TODO: how to improve locality
# order is different than it should be?
"""
    gauss_seidel(A::SparseMatrixCSC, x::Vector, b::Vector; backwards::Bool = false, iterations = 1)

Return a vector that is the result of applying one iteration of Gauss Seidel to
`Ax=b`. This iteration is equivalent to `L(b - Ux)` where `L` is the lower
triangular part of `A` and `U` is the strictly upper triangular part of `A`. If
`backwards` is true, then a backwards sweep is performed. `A` is assumed to be symmetric. If not, then `gauss_seidel(A)` is really Gauss-Seidel applied to `A^T`.
"""
function gauss_seidel{T}( A :: SparseMatrixCSC{T}
                        , x :: Vector{T}
                        , b :: Vector{T}
                        ; backwards :: Bool = false
                        , iterations :: Int = 1
                        ) :: Vector{T}
    x_new = copy(x)
    gauss_seidel!(A, x_new, b, backwards=backwards, iterations=iterations)
    x_new
end

"""
    gauss_seidel!(A::SparseMatrixCSC, x::Vector, b::Vector, backwards::Bool = false)

In place, more efficient version of `gauss_seidel(A, x, b)`.
"""
function gauss_seidel!{T}( A :: SparseMatrixCSC{T}
                         , x :: Vector{T}
                         , b :: Vector{T}
                         ; backwards :: Bool = false
                         , iterations :: Int = 1
                         ) :: Vector{T}
    rows = rowvals(A)
    vals = nonzeros(A)
    m, n = size(A)
    for col in 1:n
        colsum :: T = 0
        diag :: T = 0

        js = nzrange(A, col)
        # hack to avoid allocating
        start = backwards ? js.stop : js.start
        stop = backwards ? js.start : js.stop
        step = backwards ? -1 : 1
        @inbounds for j in start:step:stop
            row = rows[j]
            val = vals[j]
            if row != col
                colsum += val * x[row]
            else
                diag = val
            end
        end

        @inbounds x[col] = (b[col] - colsum) / diag
    end

    x
end

"""
In place version of `weighted_jacobi`.
"""
function weighted_jacobi!{T}( A :: SparseMatrixCSC{T}
                            , x :: Vector{T}
                            , b :: Vector{T}
                            ; weight :: T = 2/3
                            , iterations :: Int = 1
                            ) :: Vector{T}
    Di = spdiagm(map(x -> 1/x, diag(A)))
    R = A - spdiagm(diag(A))
    for i in 1:iterations
        x[:] = weight * Di * (b - R*x) + (1 - weight) * x
    end

    x
end

"""
    weighted_jacobi(A::SparseMatrixCSC, x::Vector, b::Vector; weight = 2/3, iterations = 1)

Return a vector that is the result of applying one iteration of Jacobi smoothing to
`Ax=b`. This iteration is equivalent to `wD(b - R*x) + (1-w)x` where `D` is the inverse diagonal and `R` is `A` without its diagonal.
"""
function weighted_jacobi{T}( A :: SparseMatrixCSC{T}
                           , x :: Vector{T}
                           , b :: Vector{T}
                           ; weight :: T = 2/3
                           , iterations :: Int = 1
                           ) :: Vector{T}
    x_new = copy(x)
    weighted_jacobi!(A, x_new, b, weight=weight, iterations=iterations)
    x_new
end

end


module FreudenthalTriangulation

using LinearAlgebra
using SparseArrays

export freudenthal_vertices, freudenthal_simplex, barycentric_coordinates, freudenthal_simplex_and_coords,
        freudenthal_simplex_and_coords!

# Construct the Freudenthal triangulation of the belief simplex

"""
    freudenthal_vertices(n::Int64, m::Int64)
Construct the list of Freudenthal vertices in an `n` dimensional space with grid resolution `m`.
The vertices are represented by a list of `n` dimensional vectors.
"""
function freudenthal_vertices(n::Int64, m::Int64)
    V = Vector{Int}[]
    v = Vector{Int}(undef, n)
    v[1] = m
    freudenthal_vertices!(V, v, 2)
    return V
end

function freudenthal_vertices!(V::Vector{Vector{Int64}}, v::Vector{Int64}, i::Int64)
    n = length(v)
    if i > n
        push!(V, copy(v))
        return
    end
    for k in 0 : v[i-1]
        v[i] = k
        freudenthal_vertices!(V, v, i+1)
    end
end

"""
    freudenthal_simplex(x::Vector{Float64})
Returns the list of vertices of the simplex of point `x` in the Freudenthal grid.
"""
function freudenthal_simplex(x::Vector{Float64})
    n = length(x)
    V = Vector{Vector{Int}}(undef, n+1)
    V[1] = floor.(Int, x)
    d = x - V[1]
    p = sortperm(d, rev=true)
    for i in 2 : n+1
        V[i] = copy(V[i-1])
        V[i][p[i-1]] += 1
    end
    return V
end

"""
    barycentric_coordinates(x::Vector{Int64}, V::Vector{Vector{Int64}})
Given a point `x` and its simplex `V` in the Freudenthal grid, returns the barycentric coordinates
of `x` in the grid. `V` must be in the same order as provided by the output of `freudenthal_simplex`
"""
function barycentric_coordinates(x, V)
    d = x - V[1]
    p = sortperm(d, rev=true)
    n = length(x)
    λ = Vector{Float64}(undef, n+1)
    λ[n+1] = d[p[n]]
    for i in n:-1:2
        λ[i] = d[p[i-1]] - d[p[i]]
    end
    λ[1] = 1.0 - sum(λ[2:end])
    return λ
end

"""
    freudenthal_coords(x::AbstractArray{Float64})
Given a point `x`, returns the simplex of the point `x` and the barycentric coordinates of `x` in the grid.
"""
function freudenthal_simplex_and_coords(x::Vector{Float64})
    V = freudenthal_simplex(x)
    return V, barycentric_coordinates(x, V)
end

"""
    freudenthal_simplex_and_coords!(x::AbstractArray{Float64}, V::Vector{Vector{Int64}}, λ::Vector{Float64})
Fills `V` and `λ` with the simplex points in the Freudenthal space and associated coordinates respectively.
"""
function freudenthal_simplex_and_coords!(x::AbstractArray{Float64}, V::Vector{Vector{Int64}}, λ::Vector{Float64})
    n = length(x)
    V[1] = floor.(Int, x)
    d = x - V[1]
    p = sortperm(d, rev=true)
    for i in 2 : n+1
        copyto!(V[i], V[i-1])
        V[i][p[i-1]] += 1
    end
    λ[n+1] = d[p[n]]
    for i in n:-1:2
        λ[i] = d[p[i-1]] - d[p[i]]
    end
    λ[1] = 1.0 - sum(λ[2:end])
    return V, λ
end

end # module


module ShearFunctions

using StaticArrays
using ..PhysicalVectors

export b1g_shear, b2g_shear, bxg_shears!

"""
    b1g_shear( eval_r::Vector2D, bob::Vector2D )

Calculate the ``B_{1g}`` from an edge dislocation for a given Burgers vector `bob` at site `eval_r`.

# Additional Information
* The edge dislocation is aligned along the ``z`` axis at the origin.
* This is in units of ``1/(1-σ)`` for a Poisson ratio ``σ``.

# Examples
```jldocstest
julia> b1g_shear( Vector2D(1.,1.), Vector2D(1.,0.) )
-0.07957747154594767
```
"""
function b1g_shear( eval_r::Vector2D, bob::Vector2D )
    return -1/π * ( eval_r.vec[1] * eval_r.vec[2] / magnitude2(eval_r) ) * (bob ⋅ eval_r) / magnitude2(eval_r)
end

"""
    b2g_shear( eval_r::Vector2D, bob::Vector2D )

Calculate the ``B_{2g}`` from an edge dislocation for a given Burgers vector `bob` at site `eval_r`.

# Additional Information
* The edge dislocation is aligned along the ``z`` axis at the origin.
* This is in units of ``1/(1-σ)`` for a Poisson ratio ``σ``.

# Examples
```jldocstest
julia> b2g_shear( Vector2D(1.,1.), Vector2D(1.,0.) )
0.15915494309189535
```
"""
function b2g_shear( bob::Vector2D, eval_r::Vector2D )
    return 1/(2*π) * ( ( eval_r.vec[1]^2 - eval_r.vec[2]^2 ) / magnitude2(eval_r) ) * (bob ⋅ eval_r) / magnitude2(eval_r)
end

"""
    bxg_shears!(eval_r::Vector2D, bob::Vector2D; diff::Function, source_r::Vector2D = Vector2D(0.,0.)) -> Tuple{T, T} where {T <: Real}

Calculate the two shears `(b1g, b2g)` from the ``B_{1g}`` and ``B_{2g}`` channels from and edge dislocation situated at the `source_r`
location and aligned along the ``z`` axis with a Burger's vector `bob`.

# Additional Information
* This function mutates both and `eval_r`. It should be used in cases where `eval_r` is a temporary vector.
* The keyword argument `diff <: Function` should be used in cases where `Vector2D` subtraction has a different definition than the normal one expected, for example with periodic boundary conditions.

# Examples
```jldocstest
julia> using StaticArrays

julia> A = MVector{2}(0., 0.)
2-element MVector{2, Float64} with indices SOneTo(2):
 0.0
 0.0

julia> bxg_shears!( A, Vector2D(1.,1.), Vector2D(1.,0.) )

julia> A
2-element MVector{2, Float64} with indices SOneTo(2):
 -0.07957747154594767
  0.15915494309189535
```
"""
function bxg_shears!( eval_r::Vector2D, bob::Vector2D; diff::Function = subtract!, source_r::Vector2D = Vector2D(0.,0.) )
    diff(eval_r, source_r)
    b1g = b1g_shear(eval_r, bob)
    b2g = b2g_shear(eval_r, bob)
    return (b1g, b2g)
end

end # module ShearFunctions


#
# 2D Position operators
#

function x_coord_matrix(p::Param)

    #Variables
    xmin        = p.xmin
    xmax        = p.xmax
    Nx          = p.xnodes

    ymin        = p.ymin
    ymax        = p.ymax
    Ny          = p.ynodes

     #Discretize space
     dx      = (xmax - xmin)/(Nx-1)
     x       = [xmin + (i-1)*dx for i in 1:Nx]

     dy      = (ymax - ymin)/(Ny-1)
     y       = [ymin + (i-1)*dy for i in 1:Ny]

     if(p.Boundaries == :periodic )
         dx      = (xmax - xmin)/(Nx)
         x       = [xmin + (i-1)*dx for i in 1:Nx]

         dy      = (ymax - ymin)/(Ny)
         y       = [ymin + (i-1)*dy for i in 1:Ny]
     end

    A = spzeros(Nx*Ny , Nx*Ny)

    index = 0

    for (ix , posx) in enumerate(x) , (iy , posy) in enumerate(y)

        index = (iy-1)*Nx+ix
        A[index , index] = posx

    end

    return A
end

function y_coord_matrix(p::Param)

    #Variables
    xmin        = p.xmin
    xmax        = p.xmax
    Nx          = p.xnodes

    ymin        = p.ymin
    ymax        = p.ymax
    Ny          = p.ynodes

    #Discretize space
    dx      = (xmax - xmin)/(Nx-1)
    x       = [xmin + (i-1)*dx for i in 1:Nx]

    dy      = (ymax - ymin)/(Ny-1)
    y       = [ymin + (i-1)*dy for i in 1:Ny]

    if(p.Boundaries == :periodic )
        dx      = (xmax - xmin)/(Nx)
        x       = [xmin + (i-1)*dx for i in 1:Nx]

        dy      = (ymax - ymin)/(Ny)
        y       = [ymin + (i-1)*dy for i in 1:Ny]
    end

    A = spzeros(Nx*Ny , Nx*Ny)

    index = 0

    for (ix , posx) in enumerate(x) , (iy , posy) in enumerate(y)

        index = (iy-1)*Nx+ix
        A[index , index] = posy

    end

    return A
end


# # Support Vector Machine
# We are showing how to solve a support vector machine problem with COSMO (and JuMP).
#
# ## Generating the Dataset
# We want to classify the points in this example dataset with $m = 100$ samples and $n = 2$ features:

using Distributions: MvNormal
using Plots, LinearAlgebra, SparseArrays, Random, Test

#-

## Generate dataset
rng = Random.MersenneTwister(123);
num_samples = 100;
Xpos = rand(rng, MvNormal([1.5, 1.5], 1.25), div(num_samples, 2))';
Xneg = rand(rng, MvNormal([-1.5, -1.5], 1.25), div(num_samples, 2))';
ypos = ones(div(num_samples, 2));
yneg = -ones(div(num_samples, 2));

#-
## Plot dataset
plot(Xpos[:, 1], Xpos[:, 2], color = :red, st=:scatter, markershape = :rect, label = "positive", xlabel = "x1", ylabel = "x2")
plot!(Xneg[:, 1], Xneg[:, 2], color = :blue, st=:scatter, markershape = :circle, label = "negative")

# with samples $(x_1, x_2, \ldots, x_m) \in \mathbb{R}^2$ and labels $y_i \in \{-1,1\}$.
#
# ## Solving SVM as a QP
# We want to compute the weights $w$ and bias term $b$ of the (soft-margin) SVM classifier:
#
# $$
# \begin{array}{ll}
#     \text{minimize}   & \|w\|^2 + \lambda \sum_{i=1}^m \text{max}(0, 1 - y_i(w^\top x_i  - b)),
# \end{array}
# $$
# where $\lambda$ is a hyperparameter. This problem can be solved as a quadratic program.
# We can rewrite above problem into an optimisation problem in primal form by introducing the auxiliary slack variables $t_i$:
#
# $$
# t_i = \text{max}(0, 1 - y_i(w^T x_i  - b)), \quad t_i \geq 0.
# $$

# This allows us to write the problems in standard QP format:
# $$
# \begin{array}{ll}
#     \text{minimize}   & \|w\|^2 + \lambda \sum_{i=1}^m t_i\\
#     \text{subject to} & y_i (w^\top x_i - b) \geq 1 - t_i, \quad \text{for } i = 1,\ldots, m\\
#                       & t_i \geq 0, \quad \text{for } i = 1,\ldots, m.
# \end{array}
# $$
# Next, we will remove the bias term $b$ by adding an initial feature $x_0 = -1$ to each sample (now: $n = 3$):

X = [-ones(num_samples) [Xpos; Xneg]];
y = [ypos; yneg];
m, n = size(X)

# ## Modelling in JuMP
# We can model this problem using `JuMP` and then hand it to `COSMO`:
using JuMP, COSMO
#-
λ = 1.0; # hyperparameter
model = JuMP.Model(optimizer_with_attributes(COSMO.Optimizer, "verbose" => true));

@variable(model, w[1:n]);
@variable(model, t[1:m] >= 0.);
@objective(model, Min, w' * w  + λ * ones(m)' * t);
@constraint(model, diagm(0 => y) * X * w .+ t .- 1 .>= 0);
status = JuMP.optimize!(model)
# The optimal weights $w = [w_0, w_1, w_2]^\top$ (where $w_0 = b$) are:
w_opt = JuMP.value.(w)

# ## Plotting the hyperplane
# The separating hyperplane is defined by $w^\top x - b = 0$. To plot the hyperplane, we calculate $x_2$ over a range of $x_1$ values:
# $$
# x_2 = (-w_1 x_1 - w_0) / w_2, \text{ where } w_0 = b.
# $$
x1 = -4:0.1:4;
x2 = (-w_opt[2] * x1  .- w_opt[1]) / w_opt[3]
plot!(x1, x2, label = "SVM separator", legend = :topleft)

# ## Modelling with COSMO
# The problem can also be solved by transforming it directly into `COSMO`'s problem format.
# Define `COSMO``s $x$-variable to be $x=[w, t]^\top$ and choose $P$, $q$, accordingly:
P = blockdiag(spdiagm(0 => ones(n)), spzeros(m, m));
q = [zeros(n); 0.5 * λ * ones(m)];

# Next we transform the first constraint $y_i (w^\top x_i - b) \geq 1 - t_i, \quad \text{for } i = 1,\ldots, m$ into
# `COSMO`'s constraint format: $Ax + b \in \mathcal{K}$.
A1 = [(spdiagm(0 => y) * X) spdiagm(0 => ones(m))];
b1 = -ones(m);
cs1 = COSMO.Constraint(A1, b1, COSMO.Nonnegatives);

# It remains to specify the constraint $t_i \geq 0, \quad \text{for } i = 1,\ldots, m$:
A2 = spdiagm(0 => ones(m));
b2 = zeros(m);
cs2 = COSMO.Constraint(A2, b2, COSMO.Nonnegatives, m+n, n+1:m+n);

# Create, assemble and solve the `COSMO.Model`:
model2 = COSMO.Model();
assemble!(model2, P, q, [cs1; cs2]);
result2 = COSMO.optimize!(model2);
w_opt2 = result2.x[1:3];
@test norm(w_opt2 - w_opt, Inf) < 1e-3


#
# Denoise an SPD Example with Cyclic Proximal Point applied to the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data as for the corresponding CP algorithm
#
using Manopt, Manifolds
using Images, CSV, DataFrames, LinearAlgebra, JLD2
#
# Settings
ExportResult = true
ExportOrig = true
ExportTable = true
asy_render_detail = 2
#
# Manifold and Data
f = artificial_SPD_image2(32)
pixelM = SymmetricPositiveDefinite(3)
resultsFolder = "examples/Total_Variation/SPD_TV/"
experimentName = "ImageCPPA"
if !isdir(resultsFolder)
    mkdir(resultsFolder)
end
if ExportOrig
    asymptote_export_SPD(
        resultsFolder * experimentName * "-orig.asy";
        data = f,
        scaleAxes = (7.5, 7.5, 7.5),
    )
    render_asymptote(resultsFolder * experimentName * "-orig.asy", render=asy_render_detail)
end
#
# Parameters
α = 6.0
maxIterations = 4000
#
# Build Problem for L2-TV
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
d = length(size(f))
rep(d) = (d > 1) ? [ones(Int, d)..., d] : d
fidelity(x) = 1 / 2 * distance(M, x, f)^2
Λ(x) = forward_logs(M, x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM), repeat(x, rep(d)...), Λ(x)), 1)
#
# Setup and Optimize
cost(x) = fidelity(x) + α * prior(x)
proxes = [(λ, x) -> prox_distance(M, λ, f, x, 2), (λ, x) -> prox_TV(M, α * λ, x, 1)]
x0 = f
@time o = cyclic_proximal_point(
    M,
    cost,
    proxes,
    x0;
    debug = [
        :Iteration,
        " | ",
        DebugProximalParameter(),
        " | ",
        :Change,
        " | ",
        :Cost,
        "\n",
        100,
        :Stop,
    ],
    record = [:Iteration, :Iterate, :Cost],
    stopping_criterion = StopAfterIteration(maxIterations),
    return_options = true,
)
y = get_solver_result(o)
yRec = get_record(o)
#
# Results
if ExportResult
    asymptote_export_SPD(
        resultsFolder *
        experimentName *
        "-result-$(maxIterations)-α$(replace(string(α), "." => "-")).asy";
        data = y,
        scaleAxes = (7.5, 7.5, 7.5),
    )
    render_asymptote(
        resultsFolder *
        experimentName *
        "-result-$(maxIterations)-α$(replace(string(α), "." => "-")).asy";
        render=asy_render_detail,
    )
end
if ExportTable
    A = cat([y[1] for y in yRec], [y[3] for y in yRec]; dims = 2)
    CSV.write(
        string(resultsFolder * experimentName * "ResultCost.csv"),
        DataFrame(A),
        writeheader = false,
    )
    save(
        resultsFolder * experimentName * "-CostValue.jld2",
        Dict("compareCostFunctionValue" => last(yRec)[3]),
    )
end


## Exercise 5-3
## Fermat’s Last Theorem says that there are no positive integers a, b, and c such that

## a^n+b^n=c^n
## for any value of n greater than 2.

## Write a function named checkfermat that takes four parameters—a, b, c and n—and checks to see if Fermat’s theorem holds. If n is greater than 2 and a^n + b^n == c^n the program should print, “Holy smokes, Fermat was wrong!” Otherwise the program should print, “No, that doesn’t work.”
println("Ans 1: ")

function check_fermats_theorem(a::Int, b::Int, c::Int, n::Int)
    if a^n + b^n == c^n && n > 2
        println("Holy smokes, Fermat was wrong!")
    else
        println("No, that doesn’t work.")
    end
end

check_fermats_theorem(3, 5, 3, 3)

## 2. Write a function that prompts the user to input values for a, b, c and n, converts them to integers, and uses checkfermat to check whether they violate Fermat’s theorem.
println("Ans 2: ")

function check_fermats_theorem2()
    print("Enter value for a: ")
    a = parse(Int64, readline())
    print("Enter value for b: ")
    b = parse(Int64, readline())
    print("Enter value for c: ")
    c = parse(Int64, readline())
    print("Enter value for n: ")
    n = parse(Int64, readline())

    check_fermats_theorem(a, b, c, n)
end

check_fermats_theorem2()

println("End.")


"""
Graph Laplacian of the line graph. With keyword `boundary = false`,
a graph Laplacian of the circle.
"""
linelaplacian(n; kwargs...) = linelaplacian(Float64,n; kwargs...)
function linelaplacian(T,n; boundary=true)
    A = SymTridiagonal(2ones(T,n), -ones(T,n-1))
    if boundary
        A[1,1] = A[end,end] = true
    end
    A
end

"""
Graph Laplacian of a `m×n` lattice.
"""
function gridlaplacian(T, m, n)
    S = sparse(T(0.0)I, n*m, n*m)
    linear = LinearIndices((1:m, 1:n))
    for i in 1:m
        for j in 1:n
            for (i2, j2) in ((i + 1, j), (i, j + 1))
                if i2 <= m && j2 <= n
                    S[linear[i, j], linear[i2, j2]] -= 1
                    S[linear[i2, j2], linear[i, j]] -= 1

                    S[linear[i, j], linear[i, j]] += 1
                    S[linear[i2, j2], linear[i2, j2]] += 1
                end
            end
        end
    end
    S
end

function gridderiv(T, m, n)
    S1 = sparse(T(0.0)*I, n*m, n*m)
    S2 = sparse(T(0.0)*I, n*m, n*m)
    S3 = sparse(T(0.0)*I, n*m, n*m)
    S4 = sparse(T(0.0)*I, n*m, n*m)

    linear = LinearIndices((1:m, 1:n))
    for i in 1:m
        for j in 1:n
            if i > 1
                S1[linear[i, j], linear[i, j]] -= 1
                S1[linear[i, j], linear[i - 1, j]] += 1
            end
            if j > 1
                S2[linear[i, j], linear[i, j]] -= 1
                S2[linear[i, j], linear[i, j - 1]] += +1
            end
            if i < m
                S3[linear[i, j], linear[i, j]] -= 1
                S3[linear[i, j], linear[i + 1, j]] += 1
            end
            if j < n
                S4[linear[i, j], linear[i, j]] -= 1
                S4[linear[i, j], linear[i, j + 1]] += +1
            end
        end
    end
    S1, S2, S3, S4
end

function boundary(A)
    m, n = size(A)
    B = zero(A)
    for i in 1:m
        for j in 1:n
            v = A[i,j]
            for (i2, j2) in ((i + 1, j), (i + 1, j +1 ), (i, j + 1))
                if 1 <= i2 <= m && 1<= j2 <= n
                    if v != A[i2, j2]
                        B[i,j] = 1.0
                        break
                    end
                end
            end
        end
    end
    B
end

function downop(T, m, n)
    S = sparse(T(0.0)I, (n÷2)*(m÷2), n*m)
    linearj = LinearIndices((1:m, 1:n))
    lineari = LinearIndices((1:m÷2, 1:n÷2))

    for i in 2:2:m
        for j in 2:2:n
            S[lineari[i÷2,j÷2], linearj[i, j]] = 1/4
            S[lineari[i÷2,j÷2], linearj[i-1, j]] = 1/4
            S[lineari[i÷2,j÷2], linearj[i, j-1]] = 1/4
            S[lineari[i÷2,j÷2], linearj[i-1, j-1]] = 1/4
        end
    end
    S
end
#=
m = n = 8
L = downop(F0, m, n); Λ = gridlaplacian(F0, m, n); Λ2 = gridlaplacian(F0, m÷2, n÷2)
heatmap([Matrix(8L*Λ*L') Matrix(Λ2) Matrix(8L*Λ*L' - Λ2)])
L = downop(F0, m, n); Λ = gridlaplacian(F0, m, n) + I/2; Λ2 = gridlaplacian(F0, m÷2, n÷2) + I
invm(x) = inv(Matrix(x))
F = Float64
L = downop(F, m, n); Λ = gridlaplacian(F, m, n)/2 + I/4; Λ2 = gridlaplacian(F, m÷2, n÷2) + I
heatmap([Matrix(L*invm(Λ)*L') Matrix(invm(Λ2)) Matrix(L*invm(Λ)*L' - invm(Λ2))])
mean(diag(L*invm(Λ)*L') - diag(invm(Λ2)))
.32007- .32008
=#


function minkowski_distance(set1::Array, set2::Array, p::Int64)

    @assert length(set1) == length(set2) "Sets don't have the same length."

    sum = zero(eltype(set1))

    for i in 1:length(set1)
        sum += abs(set2[i] - set1[i])^p
    end

    return sum^(1/p)
end


# N = 10 # 2520
N = 20 # 232792560
function isPrime(n::Int)
	for i in 2:ceil(sqrt(n))
		if n%i==0 && i!=n return false end
	end
	return true
end
primes = [x for x in 1:N if isPrime(x)]
print(primes)

function factorize(n::Int)
	f = Int[]
	while n > 1
		for i in 2:n
			if n%i==0
				n /= i
				push!(f, i)
				break
			end
		end
	end
	f
end
#println(factorize(2520))

p = Dict((i, 0) for i in primes)
for i in 1:N
	s = factorize(i)
	c = Dict((i,0) for i in primes)
	for j in s
		c[j] += 1
	end
	for (k,v) in c
		p[k] = max(p[k], v)
	end
end
println(p)
prod_p = 1
for (k, v) in p
	prod_p *= k^v
end
println(N, '\t', prod_p)


julia> sum([1,2,3,4,5].^2)
55

julia> sum([x^2 for x in [1,2,3,4,5]])
55

julia> mapreduce(x->x^2,+,[1:5])
55

julia> sum([x^2 for x in []])
0


using DifferentialEquations
using Plots
using Markdown

description = md"""

Solve the Lorenz System:
``dx/dt = σ(y-x),
  dy/dt = x(ρ-z) -y,
  dz/dt = xy - βz
``
where t∈[0,100] and u₀= [1.0, 0.0, 0.0]

"""
display(description)

# use in-place version (w/ ! symbol) so that
# the number of array allocations is minimal
function lorenz!(du, u, p, t)
    x,y,z = u
    σ, ρ, β = p
    du[1] = σ*(y-x)
    du[2] = x*(ρ-z) - y
    du[3] = x*y - β*z
end

# define IC's, parameters, and time span
u0 = [1.0; 0.0; 0.0]
p = [10.0; 28.0; 8/3]
tspan = (0.0, 100.0)

# solve the problem
prob = ODEProblem(lorenz!, u0, tspan, p)
sol = solve(prob)

# visualize the result
p1 = plot(sol, vars=(1, 2, 3)) # i.e. plot x,y,z
p2 = plot(sol, vars=(0,2))
plot(p1, p2)


function get_christoffel(x::TensorHead,TIT::TensorIndexType,metric::T ) where {T <: Tensor}
    """
    returns christoffel symbols as ``Γ_{ijk}``
    """
    @indices TIT i j k l
    g = metric

    simp(x) = contract_metric(x,TIT.metric)
    gg = diff(g(-i,-j),x(k))
    gg = simp(gg)

    h = (gg(-i,-j,-k) - gg(-j,-k,-i) + gg(-k,-i,-j))/2
    h = simp(h)

    return h(-i,-j,-k)
end

function get_riemann(x::TensorHead,h::S,TIT::TensorIndexType,metric::T) where {S <: Tensor, T <: Tensor}
    """
    compute Riemann curvature tensor as ``R_{ijkl}``
    """
    @indices TIT i j k l m n
    g = metric
    simp(x) = contract_metric(x,TIT.metric)
    dh = diff(h(-i,-j,-k),x(l))
    dh = simp(dh)

    hh = (TIT.metric(m,n)*h(-m,-i,-j))*h(-n,-k,-l)
    hh = simp(hh)

    Riemann = dh(-l,-i,-k,-j) - dh(-l,-j,-k,-i) + hh(-j,-k,-i,-l) - hh(-i,-k,-j,-l)
    return Riemann
end

"""
    Compute the Ricci tensor as ``R_{αβ}``
"""
function get_ricci(x::TensorHead, h::S, TIT::TensorIndexType, metric::T) where {S<:Tensor, T<:Tensor}
    @indices TIT i j k l
    g = metric
    Riemann = get_riemann(x, h, TIT, metric)
    Ricci = Riemann(-i, -j, -k, -l) * TIT.metric(i,k)
    return Ricci
end

"""
Compute the Ricci Scalar as ``R ≡ g^{μν}R_{μν}``
"""
function get_ricci_scalar(x::TensorHead, h::S, TIT::TensorIndexType, metric::T) where {S<:Tensor, T<:Tensor}
    @indices TIT i j
    g = metric
    Ricci = get_ricci(x, h, TIT, metric)
    R = Ricci(-i,-j)*TIT.metric(j,i)
    return R
end

"""
Compute the Einstein tensor as ``G^{μν} ≡ R^{μν} - \frac{1}{2}⋅ g^{μν}R
"""
function get_einstein(x::TensorHead, h::S, TIT::TensorIndexType, metric::T) where {S<:Tensor, T<:Tensor}
    @indices TIT i j
    g = metric
    R = get_ricci_scalar(x, h, TIT, metric)
    Ricci = get_ricci(x, h, TIT, metric)
    return Ricci(-i,-j) - 1/2*g(-i,-j)*R
end


function log(a::Real, x::Real, ε::Real)
    z = x
    t = 1
    y = 0
    while z > a || z < 1/a || t > ε
        if z > a
            z /= a
            y += t
        elseif z < 1/a
            z *= a
            y -= t
        else
            t /= 2
            z *= z
        end
    end
    return y
end


#softmax function: f(x)=exp(x)/sum(exp(x))

Fsoftmax(x::Array{Float64,2})=(exp(x)./sum(exp(x),1),[]) # TODO: better to subtract the max of x to make it numerically more stable

function Fsoftmax_inplace(value,auxvalue,x::Array{Float64,2})
    copy!(value,exp(x)./sum(exp(x),1)) # TODO: better to subtract the max of x to make it numerically more stable
end

function Dsoftmax(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::Array{Float64,2})
    axpy!(1.0,f_c.*(grad_c-repmat(sum(f_c.*grad_c,1),size(f_c,1),1)),grad_n)
end

if PROC=="GPU"

    function Fsoftmax(X::CudaArray)
        expX=CudaArray(Float64,size(X))
        exp!(X,expX)
        onr=CudaArray(Float64,(1,size(X,1))); fill!(onr,1.0);
        colsum=onr*expX
        onc=CudaArray(Float64,(size(X,1),1)); fill!(onc,1.0);
        out=CudaArray(Float64,size(X))
        vdiv!(1.0,expX,onc*colsum,out)
        free(expX); free(onr); free(colsum); free(onc)
        return (out,[]) # memory leak -- how to free out
    end

    function softmax!(X::CudaArray,out::CudaArray)
        expX=CudaArray(Float64,size(X))
        exp!(X,expX)
        onr=CudaArray(Float64,(1,size(X,1))); fill!(onr,1.0);
        colsum=onr*expX
        onc=CudaArray(Float64,(size(X,1),1)); fill!(onc,1.0);
        vdiv!(1.0,expX,onc*colsum,out)
        free(expX); free(onr); free(colsum); free(onc)
    end
    export softmax!

    function Fsoftmax_inplace(value,auxvalue,X::CudaArray)
        expX=CudaArray(Float64,size(X))
        exp!(X,expX)
        onr=CudaArray(Float64,(1,size(X,1))); fill!(onr,1.0);
        colsum=onr*expX
        onc=CudaArray(Float64,(size(X,1),1)); fill!(onc,1.0);
        vdiv!(1.0,expX,onc*colsum,value)
        free(expX); free(onr); free(colsum); free(onc)
    end

    function Dsoftmax(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray)
        tmp=CudaArray(Float64,size(f_c))
        vmult!(1.0,f_c,grad_c,tmp)
        onr=CudaArray(Float64,(1,size(tmp,1))); fill!(onr,1.0);
        colsum=onr*tmp
        onc=CudaArray(Float64,(size(tmp,1),1)); fill!(onc,1.0);
        copy!(tmp,grad_c)
        axpy!(-1.0,onc*colsum,tmp)
        vmult!(1.0,f_c,tmp,tmp)
        axpy!(1.0,tmp,grad_n)
        free(tmp); free(onr); free(onc);
    end

end

Derivative[Fsoftmax]=Dsoftmax
Inplace[Fsoftmax]=Fsoftmax_inplace

ADsoftmax(n)=ADnode(Fsoftmax,n)

softmax(n::ADnode)=ADnode(Fsoftmax,n)
softmax(x::Array)=exp(x)./sum(exp(x),1) # TODO: better to subtract the max of x to make it numerically more stableFsoftmax(x)

export softmax


module Utility

using LinearAlgebra
using Statistics, StatsBase

export sinkhorn

"""
    sinkhorn(A; r=[], c=[], maxit=1000, δ=1e-6, verbose=false)

Compute the row and column multiplicative factors that constrain marginals of `A` to unity.
Returns a boolean flag if algorithm converges within tolerance `δ` within `maxit` iterations.
If `r` is given and non-empty, row marginals will be constrained to user-supplied inputs.
If `c` is given and non-empty, column marginals will be constrained to user-supplied inputs.
"""
function sinkhorn(A; r=[], c=[], maxit=1000, δ=1e-6, verbose=false)
    if length(r) == 0
        r = size(A,2)
    end

    if length(c) == 0
        c = size(A,1)
    end

    x = ones(size(A,1))
    y = ones(size(A,2))
    for i ∈ 1:maxit
        δr = maximum(abs.(x.*(A*y)  .- r))
        δc = maximum(abs.(y.*(A'*x) .- c))

        if verbose
            @show minimum(A'*x), minimum(A*y)
            @show r, c
            @show i, δr, δc
        end

        (isnan(δr) || isnan(δc)) && return x, y, false
        (δr < δ && δc < δ) 		 && return x, y, true

        y = c ./ (A'*x)
        x = r ./ (A*y)
        if verbose
            @show mean(x), mean(y)
        end
    end

    return x, y, false
end

function sinkhorn(N, ϕ; maxit=1000, δ=1e-6, verbose=false)
    nrow = size(N,1)
    ncol = size(N,2)

    N² = N.^2
    u, v = ones(size(N,1)), ones(size(N,2))
    for i ∈ 1:maxit
        σ² = ((Diagonal(u.^2 ./ (1 .+ ϕ))*N²*Diagonal(v.^2)) .+ (Diagonal(u.*ϕ ./ (1 .+ ϕ))*N*Diagonal(v)))
        δr = maximum(abs.(sum(σ², dims=1) .- nrow))
        δc = maximum(abs.(sum(σ², dims=2) .- ncol))

        if verbose
            @show i, δr, δc
        end

        (isnan(δr) || isnan(δc)) && return u, v, false
        (δr < δ && δc < δ) 		 && return u, v, true

        A = vec(sum(Diagonal(u.^2 ./ (1 .+ ϕ)) * N², dims=1))
        B = vec(sum(Diagonal(u.*ϕ ./ (1 .+ ϕ)) * N , dims=1))

        v = @. (sqrt(B^2 + 4*A*nrow) - B) / (2*A)

        A = vec(sum(N² * Diagonal(v.^2),           dims=2))
        B = vec(sum(Diagonal(ϕ) * N * Diagonal(v), dims=2))

        u = @. (sqrt(B^2 + 4*A*ncol*(1+ϕ)) - B) / (2*A)
    end

    return u, v, false
end

end