# fitting noisy data to an exponential model
model(xpts, p) = p[1]*exp(-xpts.*p[2])

# some example data
srand(12345)
xpts = linspace(0,10,20)
data = model(xpts, [1.0 2.0]) + 0.01*randn(length(xpts))

beta, r, J = curve_fit(model, xpts, data, [0.5, 0.5])
@assert norm(beta - [1.0, 2.0]) < 0.05

# can also get error estimates on the fit parameters
errors = estimate_errors(beta, r, J)
@assert norm(errors - [0.016, 0.075]) < 0.01


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ ec794b0c-67b1-11eb-23e7-459641028308
begin
	using Markdown
	using InteractiveUtils
	using Test
	using LinearAlgebra
end

# ╔═╡ 47f1758c-67b1-11eb-0a89-27d2595e5912
md"""

## Task \# 8

Compute the norms and condition numbers

#### Norm

```math
||A||^2 = max \frac{||Ax||^2}{||x||^2} = max \frac{ x^T A^T A x }{ x^T x } = \lambda_{max}(A^T A) = \sigma^{2}_{max}
```

```math
||A|| = \sqrt{\sigma^{2}_{max}} = \sigma_{max}
```

```math
||A^{-1}|| = \frac{1}{\sigma_{min}}
```

"""

# ╔═╡ 9aa05ff6-67b5-11eb-0b36-95986372b105
function normCalculator(A::Array{Int})
	AᵀA = A' * A

	A_eigenvals, A_eigenvec = eigen(AᵀA)

	A_normFromEigenSquared = √max(A_eigenvals...)

	A_norm = opnorm(A)

	return (A_normFromEigenSquared, A_norm)
end

# ╔═╡ 4533fdcc-67b3-11eb-30f8-851ef7fc43c9
begin
	A₁ = [
		1 7
		1 1
	]

	A₁normFromEigenSquared, A₁norm = normCalculator(A₁)

	@test round.(A₁normFromEigenSquared, digits=5) == round.(A₁norm, digits=5)

	A₁normFromEigenSquared, A₁norm, cond(A₁)
end

# ╔═╡ bf65578a-67b8-11eb-2627-9139089fcf68
md"""
#### Cond is Inf, because of `` || A_2^{-1} || = \sigma_{min} = 0 ``
"""

# ╔═╡ 1c2ed1ac-67b5-11eb-1e77-67e96dbd2802
begin
	A₂ = [
		1 1
		0 0
	]

	A₂normFromEigenSquared, A₂norm = normCalculator(A₂)

	@test round.(A₂normFromEigenSquared, digits=5) == round.(A₂norm, digits=5)

	A₂normFromEigenSquared, A₂norm, cond(A₂)
end

# ╔═╡ 1568a204-67b9-11eb-3d55-d72928803e5e
md"""
#### Cond is 1, because of `` || A_2 || = || A_2^{-1} || ``
"""

# ╔═╡ ed7d86f6-67b3-11eb-273c-43770e6f299c
begin
	A₃ = [
		1 1
		-1 1
	]

	A₃normFromEigenSquared, A₃norm = normCalculator(A₃)

	@test round.(A₃normFromEigenSquared, digits=5) == round.(A₃norm, digits=5)

	A₃normFromEigenSquared, A₃norm, cond(A₃)
end

# ╔═╡ c3dc9efe-67bd-11eb-0224-af9d70653b73
md"""
### Task \# 13
Estimate the condition number
"""

# ╔═╡ e21e095c-67bd-11eb-12e4-e92b18b0b223
begin
	A₄ = [
		1 1
		1 1.0001
	]
	cond(A₄)
end

# ╔═╡ 67060602-67c2-11eb-046e-dfc437a3dc81
md"""
### Task \#17
Fibonacci matrix and the SVD
"""

# ╔═╡ 832167be-67c2-11eb-1693-ab8ea1a74995
begin
	F = [
		1 1
		1 0
	]
	svd(F)
end

# ╔═╡ 8d7f475a-67c3-11eb-243b-53c5d68004f1
md"""
### Task \#18
"""

# ╔═╡ 9a4b6df6-67c3-11eb-27d6-2546dd023fb4
begin
	K = [
		1 2
		2 5
	]
	lu(K)
end

# ╔═╡ Cell order:
# ╠═ec794b0c-67b1-11eb-23e7-459641028308
# ╟─47f1758c-67b1-11eb-0a89-27d2595e5912
# ╠═9aa05ff6-67b5-11eb-0b36-95986372b105
# ╠═4533fdcc-67b3-11eb-30f8-851ef7fc43c9
# ╟─bf65578a-67b8-11eb-2627-9139089fcf68
# ╠═1c2ed1ac-67b5-11eb-1e77-67e96dbd2802
# ╟─1568a204-67b9-11eb-3d55-d72928803e5e
# ╠═ed7d86f6-67b3-11eb-273c-43770e6f299c
# ╟─c3dc9efe-67bd-11eb-0224-af9d70653b73
# ╠═e21e095c-67bd-11eb-12e4-e92b18b0b223
# ╟─67060602-67c2-11eb-046e-dfc437a3dc81
# ╠═832167be-67c2-11eb-1693-ab8ea1a74995
# ╟─8d7f475a-67c3-11eb-243b-53c5d68004f1
# ╠═9a4b6df6-67c3-11eb-27d6-2546dd023fb4


# plots the contours of the OLS objective function,
# with the true and estimated coefficients
# generate data
using Plots, Econometrics
n = 20
x = [ones(n) randn(n)]
beta = randn(2)
y = x*beta+randn(n)
# ols estimate and objective function
betahat = x\y
s = (a,b)-> (1/n)*sum((y - x*[a,b]).^2)
# plot the contours and the points
closeall()
b1 = range(-4.0,stop=4.0,length=100)
b2 = range(-4.0,stop=4.0,length=100)
p = contour(b1,b2,(b1,b2)->s(b1,b2),fill=true, c=:viridis)
scatter!([beta[1]],[beta[2]], markersize=10, label="true")
scatter!([betahat[1]],[betahat[2]], markersize=10, label="estimated")
# do OLS by simulated annealing, and plot the path
lb = -4*ones(2,1)
ub = -lb
thetastart = fill(-3.0, 2, 1)
obj = theta-> (1/n)*sum((y - x*theta).^2)[1,1]
results = samin(obj, thetastart, lb, ub, verbosity=2)
path = results[4]
plot!(path[:,4], path[:,5])


using Plots; pyplot()
using DeSolver
using LinearAlgebra
using LaTeXStrings

# Number of collocation points
N = 50
h= pi/N

# create function
function f(u,t)
    du = zeros(length(u))
    for i in 1:length(u)
        if i ==1
            du[1] = (-2u[1] + 2u[2])/h^2
        elseif i == length(u)
            du[length(u)] = (2u[length(u)-1] - 2u[length(u)])/h^2
        else
            du[i] = (u[i-1] - 2u[i] + u[i+1])/h^2
        end
    end
    return du
end

xs=range(0, stop=pi, length=N)
u0=sin.(xs)

soltest1=odesolve(f, u0, (0.,10.), 1, method="RK4")
soltest2=odesolve(f, u0, (0.,10.), .1, method="RK4")
soltest3=odesolve(f, u0, (0.,10.), .01, method="RK4")


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

"""
    ln_mean(x, y)

Compute the logarithmic mean

    ln_mean(x, y) = (y - x) / (log(y) - log(x)) = (y - x) / log(y / x)

Problem: The formula above has a removable singularity at `x == y`. Thus,
some care must be taken to implement it correctly without problems or loss
of accuracy when `x ≈ y`. Here, we use the approach proposed by
Ismail and Roe (2009).
Set ξ = y / x. Then, we have

    (y - x) / log(y / x) = (x + y) / log(ξ) * (ξ - 1) / (ξ + 1)

Set f = (ξ - 1) / (ξ + 1) = (y - x) / (x + y). Then, we use the expansion

    log(ξ) = 2 * f * (1 + f^2 / 3 + f^4 / 5 + f^6 / 7) + O(ξ^9)

Inserting the first few terms of this expansion yields

    (y - x) / log(ξ) ≈ (x + y) * f / (2 * f * (1 + f^2 / 3 + f^4 / 5 + f^6 / 7))
                     = (x + y) / (2 + 2/3 * f^2 + 2/5 * f^4 + 2/7 * f^6)

Since divisions are usually more expensive on modern hardware than
multiplications (Agner Fog), we try to avoid computing two divisions. Thus,
we use

    f^2 = (y - x)^2 / (x + y)^2
        = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y)

Given ε = 1.0e-4, we use the following algorithm.

    if f^2 < ε
      # use the expansion above
    else
      # use the direct formula (y - x) / log(y / x)
    end

# References
- Ismail, Roe (2009).
  Affordable, entropy-consistent Euler flux functions II: Entropy production at shocks.
  [DOI: 10.1016/j.jcp.2009.04.021](https://doi.org/10.1016/j.jcp.2009.04.021)
- Agner Fog.
  Lists of instruction latencies, throughputs and micro-operation breakdowns
  for Intel, AMD, and VIA CPUs.
  https://www.agner.org/optimize/instruction_tables.pdf
"""
@inline function ln_mean(x, y)
  epsilon_f2 = 1.0e-4
  f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
  if f2 < epsilon_f2
    return (x + y) / @evalpoly(f2, 2, 2/3, 2/5, 2/7)
  else
    return (y - x) / log(y / x)
  end
end

"""
    inv_ln_mean(x, y)

Compute the inverse `1 / ln_mean(x, y)` of the logarithmic mean
[`ln_mean`](@ref).

This function may be used to increase performance where the inverse of the
logarithmic mean is needed, by replacing a (slow) division by a (fast)
multiplication.
"""
@inline function inv_ln_mean(x, y)
  epsilon_f2 = 1.0e-4
  f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
  if f2 < epsilon_f2
    return @evalpoly(f2, 2, 2/3, 2/5, 2/7) / (x + y)
  else
    return log(y / x) / (y - x)
  end
end

# `Base.max` and `Base.min` perform additional checks for signed zeros and `NaN`s
# which are not present in comparable functions in Fortran/C++. For example,
# ```julia
# julia> @code_native debuginfo=:none syntax=:intel max(1.0, 2.0)
#         .text
#         vmovq   rcx, xmm1
#         vmovq   rax, xmm0
#         vcmpordsd       xmm2, xmm0, xmm0
#         vblendvpd       xmm2, xmm0, xmm1, xmm2
#         vcmpordsd       xmm3, xmm1, xmm1
#         vblendvpd       xmm3, xmm1, xmm0, xmm3
#         vmovapd xmm4, xmm2
#         test    rcx, rcx
#         jns     L45
#         vmovapd xmm4, xmm3
# L45:
#         test    rax, rax
#         js      L54
#         vmovapd xmm4, xmm3
# L54:
#         vcmpltsd        xmm0, xmm0, xmm1
#         vblendvpd       xmm0, xmm4, xmm2, xmm0
#         ret
#         nop     word ptr cs:[rax + rax]
#
# julia> @code_native debuginfo=:none syntax=:intel min(1.0, 2.0)
#         .text
#         vmovq   rcx, xmm1
#         vmovq   rax, xmm0
#         vcmpordsd       xmm2, xmm0, xmm0
#         vblendvpd       xmm2, xmm0, xmm1, xmm2
#         vcmpordsd       xmm3, xmm1, xmm1
#         vblendvpd       xmm3, xmm1, xmm0, xmm3
#         vmovapd xmm4, xmm2
#         test    rcx, rcx
#         jns     L58
#         test    rax, rax
#         js      L67
# L46:
#         vcmpltsd        xmm0, xmm1, xmm0
#         vblendvpd       xmm0, xmm4, xmm2, xmm0
#         ret
# L58:
#         vmovapd xmm4, xmm3
#         test    rax, rax
#         jns     L46
# L67:
#         vmovapd xmm4, xmm3
#         vcmpltsd        xmm0, xmm1, xmm0
#         vblendvpd       xmm0, xmm4, xmm2, xmm0
#         ret
#         nop     word ptr cs:[rax + rax]
#         nop     dword ptr [rax]
# ```
# In contrast, we get the much simpler and faster version
# ```julia
# julia> @code_native debuginfo=:none syntax=:intel Base.FastMath.max_fast(1.0, 2.0)
#         .text
#         vmaxsd  xmm0, xmm1, xmm0
#         ret
#         nop     word ptr cs:[rax + rax]
#
# julia> @code_native debuginfo=:none syntax=:intel Base.FastMath.min_fast(1.0, 2.0)
#         .text
#         vminsd  xmm0, xmm0, xmm1
#         ret
#         nop     word ptr cs:[rax + rax]
# ```
# when using `@fastmath`, which we also get from
# [Fortran](https://godbolt.org/z/Yrsa1js7P)
# or [C++](https://godbolt.org/z/674G7Pccv).
"""
    max(x, y, ...)

Return the maximum of the arguments. See also the `maximum` function to take
the maximum element from a collection.

This version in Trixi.jl is semantically equivalent to `Base.max` but may
be implemented differently. In particular, it may avoid potentially expensive
checks necessary in the presence of `NaN`s (or signed zeros).

# Examples

julia> max(2, 5, 1)
5
"""
@inline max(args...) = @fastmath max(args...)

"""
    min(x, y, ...)

Return the minimum of the arguments. See also the `minimum` function to take
the minimum element from a collection.

This version in Trixi.jl is semantically equivalent to `Base.min` but may
be implemented differently. In particular, it may avoid potentially expensive
checks necessary in the presence of `NaN`s (or signed zeros).

# Examples

julia> min(2, 5, 1)
1
"""
@inline min(args...) = @fastmath min(args...)

end # @muladd


#Algorithm 12.2
#Heat equation with backwards-difference
#Boundary conditions u(0, t) = u(l, t) = 0, 0 < t < T

using Printf

function heatbd(f, L, T, α, m, N)
	w = zeros(Float64, m)
	l = zeros(Float64, m)
	u = zeros(Float64, m)
	z = zeros(Float64, m)
	h = L/m
	k = T/N
	λ = k*α^2/h^2
	for i ∈ 1:(m - 1)
		w[i] = f(i * h)
	end
	l[1] = 1 + 2λ
	u[1] = -λ/l[1]
	for i ∈ 2:(m - 2)
		l[i] = 1 + 2λ + λ*u[i - 1]
		u[i] = -λ/l[i]
	end
	l[m - 1] = 1 + 2λ + λ*u[m - 2]
	for j ∈ 1:N
		t = j*k
		z[1] = w[1] / l[1]
		for i ∈ 2:(m - 1)
			z[i] = (w[i] + λ*z[i - 1])/l[i]
		end
		w[m - 1] = z[m - 1]
		for i ∈ (m - 2):1
			w[i] = z[i][t] - u[i]w[i+1]
		end
		for i ∈ 1:(m - 1)
			x = i*h
			println(x, w[i])
			x = i * h
			@printf("%3d %11.8f %14.8f\n", i, x, w[i])
		end
	end
end

function distrib(x)
	return sin(π*x)
end

heatbd(distrib, 1, 10, 0.5, 100, 100)


# ---
# title: 172. Factorial Trailing Zeroes
# id: problem172
# author: zhwang
# date: 2022-03-25
# difficulty: Easy
# categories: Math
# link: <https://leetcode.com/problems/factorial-trailing-zeroes/description/>
# hidden: true
# ---
#
# Given an integer `n`, return _the number of trailing zeroes in`n!`_.
#
# **Follow up:** Could you write a solution that works in logarithmic time
# complexity?
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 3
#     Output: 0
#     Explanation:  3! = 6, no trailing zero.
#
#
# **Example 2:**
#
#
#
#     Input: n = 5
#     Output: 1
#     Explanation:  5! = 120, one trailing zero.
#
#
# **Example 3:**
#
#
#
#     Input: n = 0
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 104`
#
#
## @lc code=start
using LeetCode

## simulate Mathematica's function NestWhileList
function nest_while_list(f::Function, val::T, chk::Function)::Vector{T} where T
    res = [val]
    while chk(val)
        val = f(val)
        push!(res, val)
    end
    res
end

trailing_zeroes(n::Int) = sum(nest_while_list(i->i÷5, n÷5, >(1)))
## equivalent to the following
## function trailing_zeroes(n::Int)::Int
##     res = 0
##     while n >=5
##         n ÷= 5
##         res += n
##     end
##     res
## end

## @lc code=end

## @lc test=start
@testset "172.factorial-trailing-zeroes.jl" begin
    @test all(trailing_zeroes(i) == 0 for i in 0:4)
    @test trailing_zeroes(5) == 1
    @test trailing_zeroes(1999) == 496
    @test trailing_zeroes(2000) == trailing_zeroes(2001) == 499
end
## @lc test=end

## @lc info=start
# link: [solution to 172]()
## @lc info=end


"""
    pinv_M, rnull_optional, lnull_optional = pinv_qr(M::AbstractArray)

Computes Moore-Penrose pseudoinverse of `M` using a rank-revealing QR
factorization and optionally returns an orthonormal basis for right
and left null-space. This can be significantly faster than using
SVD for computing the pseudoinverse.  Moreover, `A` can be sparse or
dense.

# Examples
```julia-repl
julia> pinv_qr([3 8 17; 2 1 4; 8 3 21])
([-0.10344827586206914 1.344827586206895 -0.17241379310344793;0.1149425287356321 0.8390804597701144 -0.2528735632183907;0.022988505747126506 -0.6321839080459762 0.1494252873563217], Matrix{Float64}(undef, 3, 0), Matrix{Float64}(undef, 3, 0))
```
"""
function pinv_qr(M::AbstractArray{T}) where {T}
    m = size(M, 1)
    n = size(M, 2)

    F = qr(M)
    Q = F.Q*I  # so Q is allways m x m, with Matrix(F.Q) we get a m x n Matrix for m >= n
    R = Matrix(F.R)

    # compute rank
    abs_diag_R = abs.(diag(R))
    tol = 100 * minimum(size(M)) * maximum(abs_diag_R) * eps(T)  # MAGIC CONSTANT
    rnk = sum(tol .< abs_diag_R)

    # compute pseudoinverse
    M_pinv = Matrix{T}(undef, n, m)

    S = R[1:rnk, 1:rnk] \ R[1:rnk, rnk+1:end]
    W = R[1:rnk, 1:rnk] \ Q[:, 1:rnk]'
    X = (S' * S + I) \ (S' * W)
    M_pinv[rnk+1:end, :] = X
    M_pinv[1:rnk, :] = W - S*X

    # computing null spaces
    right_null_size = n - rnk
    left_null_size = m - rnk

    # right null space
    if right_null_size > 0
        rnull, = qr(Matrix([S; -I(right_null_size)]))
        rnull = Matrix(rnull)
    else
        rnull = zeros(T, n, 0)
    end

    # left null space
    if left_null_size > 0
        lnull = Q[: ,end - (left_null_size - 1): end]
    else
        lnull = zeros(T, m, 0)
    end

    return M_pinv, rnull, lnull
end

function pinv_qr(M::AbstractArray{<:Int})
    pinv_qr(float(M))
end


using ImplicitGraphs, Permutations

"""
    iTransposition(n::Int, adjacent::Bool=true)

Create an `ImplicitGraph` whose vertices are all permutations of `1:n`
in which two vertices are adjacent if they differ by a transpostion.

* If `adjacent` is true, then only consider transpositions of the form `(a,a+1)`
* Otherwise, consider all possible transpositions of the form `(a,b)` where `1≤a<b≤n`.
"""
function iTransposition(n::Int, adjacent::Bool = true)
    @assert n > 0 "Require n>0, got n=$n"

    vcheck(v::Permutation) = length(v) == n

    function outs(v::Permutation)
        if adjacent
            return [v * Transposition(n, i, i + 1) for i = 1:n-1]
        end
        return [v * Transposition(n, i, j) for i = 1:n for j = 1:n if i ≠ j]
    end

    return ImplicitGraph{Permutation}(vcheck, outs)
end

"""
    trans_string(p::Permutation)

We assume that `p` is a transposition.
"""
function trans_string(p::Permutation)::String
    n = length(p)
    FP = fixed_points(p)
    @assert n == length(FP) + 2 "Permutation must be a transposition"
    elts = sort(setdiff(1:n, FP))
    a, b = elts
    return "($a,$b)"
end

"""
    pscore(p::Permutation)::Int

Measures the extent to which `p` is different from the identity permutation
using this formula:

`sum(abs(p[k]-k) for k=1:length(p))`
"""
function pscore(p::Permutation)::Int
    n = length(p)
    sum(abs(i - p(i)) for i = 1:n)
end

"""
    crazy_trans_product(p::Permutation, adjacent::Bool=true)::String

Express `p` as the product of transpositions. With `adjacent=true` the transpositions
are all of the form `(a,a+1)`; otherwise, all possible transpositions are allowed.

Result is expressed as a `String`.
"""
function crazy_trans_product(p::Permutation, adjacent::Bool = true)::String
    n = length(p)
    G = iTransposition(n, adjacent)
    P = reverse(guided_path_finder(G, p, Permutation(n), score = pscore))

    nP = length(P)

    if nP == 1
        return "()"
    end

    Q = [P[k]' * P[k+1] for k = 1:nP-1]

    prod(trans_string(Q[j]) for j = 1:length(Q))
end


#=
Each new term in the Fibonacci sequence is generated by adding the previous two terms. By starting with 1 and 2, the first 10 terms will be:

1, 2, 3, 5, 8, 13, 21, 34, 55, 89, ...

By considering the terms in the Fibonacci sequence whose values do not exceed four million, find the sum of the even-valued terms.
=#

function calc()
  fibs = [1, 1]
  while fibs[end] <= 4_000_000
    push!(fibs, fibs[end] + fibs[end-1])
  end
  filter(x -> iseven(x) && x < 4_000_000, fibs) |> sum |> println
end
@time calc()


# some utility functions
const Γ = gamma

function generalized_binomial(α::T, n::S) where {T,S <: Integer}
    R = Base.promote_op(/,T,S)
    out = one(R)
    for i in n:-1:1
        out *= α/i
        α -= 1
    end
    out
end

## memoized version of the above
@memoize function generalized_binomial′(a, k::Int)
    k < 0 && return NaN * one(a)
    k == 0 && return one(a)/1
    generalized_binomial′(a,k-1) * (a-k+1)/k
end

## x⁽ⁱ⁾ is falling (x)⋅(x-1)⋯(x-n+1)
function Pochhammer(::Val{:falling}, z::T,n::Int) where {T}
    iszero(n) && return one(T)
    out = z
    for i in 1:n-1
        z -= 1
        out *= z
    end
    out
end

## (x)ᵢ is rising (x)⋅(x+1)⋯(x+n-1) = Γ(z+n)/Γ(z)
function Pochhammer(::Val{:rising}, z::T,n::Int) where {T}
    iszero(n) && return one(T)
    out = z
    for i in 1:n-1
        z += 1
        out *= z
    end
    out
end
# default is (x)ᵢ
Pochhammer(z,n::Int) = Pochhammer(Val(:rising),  z, n)

# (x)_n/n!
Pochhammer_factorial(z, n) = Pochhammer_factorial(Val(:rising), z, n)

function Pochhammer_factorial(::Val{:rising}, z, n)
    iszero(n) && return one(z)/1
    prod((z+i)/(n-i) for i in 0:n-1)
end

## Alternative to HypergeometricFunctions.jl so that
## variables can be inserted into the numerator
"""
    pFq(as, bs, z; [maxevals=1000])

Compute the generalized [hypergeometric](https://en.wikipedia.org/wiki/Generalized_hypergeometric_function) function. The `HypergeometricFunction.jl` package would normally be used for such calculations, but this one accepts polynomial values for the `as` and `z` values.

# Example

From [mathworld](https://mathworld.wolfram.com/HypergeometricFunction.html)

```jldoctest
julia> using Polynomials, SpecialPolynomials

julia> import SpecialPolynomials: pFq, Pochhammer

julia> pFq((1/3,2/3), 5/6, 27/32) ≈ 8/5
true

julia> pFq([1/4, 1/2], [3/4], 80/81; maxevals=2000) ≈ 9/5
true

julia> x = variable()
Polynomial(x)

julia> n = 5
5

julia> pFq((-n,n+1), 1, (1-x)/2) ≈ basis(Legendre,n)(x)
true

julia> α, β, n = 1/2, 1/2, 5;

julia> Pochhammer(α+1,n)/factorial(n) * pFq((-n,n+α+β+1), α+1, (1-x)/2) ≈ basis(Jacobi{α,β}, n)(x)
true
```
"""
function pFq(as, bs, z; maxevals=1000)
    n = 1
    acc = trm = one(z)

    p,q = length(as), length(bs)

    # element in `as` is negative integer or 0 && converges
    # p ≤ q && converges
    # p = q + 1 |x| < 1 && converge
    # p > q + 1 && diverges
    while n < maxevals
        a = isempty(as) ? 1 : prod(as)
        b = isempty(bs) ? 1 : prod(bs)

        iszero(a) && return acc
        iszero(b) && return Inf  # check on b

        trm *= a /b * z  /n
        acc += trm

        as = plus_1(as)

        bs = plus_1(bs)
        n += 1
    end

    if (p > q + 1 || (p == q+1 && abs(z) < 1))
        return acc
    else
        NaN*acc
    end
end
## tuples
@inline plus_1(as) = map(x->x+1, as)

# use HypergeometricFunction.mFn if possible
function pFq(as::Tuple{A,B}, bs::Tuple{C}, z::AbstractFloat) where {A,B,C}
    HypergeometricFunctions._₂F₁(as[1],as[2],bs[1],z)
end
function pFq(as::Tuple{A,B,C}, bs::Tuple{D,E}, z::AbstractFloat) where {A,B,C,D,E}
    HypergeometricFunctions._₃F₂(as[1],as[2],as[3],bs[1],bs[2],z)
end

## Try to speed up  quadgk by not specializing on F
mutable struct Wrapper
    F
end
(F::Wrapper)(x) = F.F(x)

_quadgk(f, a, b) = quadgk(Wrapper(f), a, b)[1]
const ∫ = _quadgk

checked_div(a, b) = (iszero(a) && iszero(b)) ? a : a/b


#=
  Cookie problem.

  From Think Bayes, page 3
  """
  Suppose there are two bowls of cookies.
  Bowl 1 contains 30 vanilla cookies and 10 chocolate cookies.
  Bowl 2 contains 20 of each.

  Now suppose you choose one of the bowls at random and, without looking,
  select a cookie at random. The cookie is vanilla.
  What is the probability that it came from Bowl 1?
  """

  Distributions of variable bowl (num:0)
  1.00000 =>    5991  (0.599100)
  2.00000 =>    4009  (0.400900)

  Cf  ~/webppl/cookie_problem.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function cookie_problem()
    vanilla = 1
    chocolate = 2

    bowl1 = 1
    bowl2 = 2

    bowl ~ DiscreteUniform(bowl1,bowl2)
    cookie ~ (bowl == bowl1) ? Categorical(simplex([30,10])) : # [vanilla,chocolate]
                               Categorical(simplex([20,20]))

    true ~ Dirac(cookie == vanilla)

end

model = cookie_problem()
num_chains = 4

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

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

show_var_dist_pct(chns, :bowl)


using DifferentialEquations
using Plots
using Markdown
using LinearAlgebra
using StaticArrays

description=md"""

solve the matrix system
``u' = Au``
for matrix A and i.c. u₀

"""
display(description)

# use immutable arrays that are stack-allocated
A = @SMatrix [1.0 0 0 -5
              4 -2 4 -3
              -4 0 0 1
              5 -2 2 3]

u0 = rand(4, 2)

tspan = (0.0, 1.0)

f(du, u, p, t) = mul!(du, A, u)  # in-place version of matrix multiply

prob = ODEProblem(f, u0, tspan)
sol = solve(prob)
plot(sol)


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.7.0-DEV
#     language: julia
#     name: julia-1.7
# ---

# %%
using Plots
plotly()
f(x, y) = sin(x) * sin(y) + (x - 2y)^2/50
x = range(-π, π; length=401)
y = range(-π, π; length=401)
surface(x, y, f; xlim=extrema(x), ylim=extrema(y), camera=(30, 45), size=(500, 500), colorbar=false)

# %%
fmin, pt = findmin(p -> f(p...), Iterators.product(x, y))

# %% tags=[]
surface(x, y, f; xlim=extrema(x), ylim=extrema(y), camera=(30, 45), size=(500, 500), colorbar=false)
scatter!([pt[1]], [pt[2]], [f(pt...)]; label="", color=:cyan, ms=2)

# %%
f(p) = sin(p[1]) * sin(p[2]) + (p[1] - 2p[2])^2/50
x = y = range(-π, π; length=401)
findmin(f, Iterators.product(x, y))

# %%
f(p) = sin(p[1]) * sin(p[2]) + (p[1] - 2p[2])^2/50
x = y = range(-π, π, 1000)
findmin(f, Iterators.product(x, y))

# %%
using Plots
plotly()
x = y = range(-3, 3, length=100)
f(x, y) = exp(-x^2-y^2)
z = f.(x', y)
surface(x, y, z; colorbar=false, size=(600, 600), color=:jet)

# %%
using Plots
plotly()
x = y = range(-3, 3, length=100)
z = exp.(.-x'.^ 2 .- y.^ 2)
surface(x, y, z; colorbar=false, size=(600, 600), color=:jet)

# %%
using Plots
plotly()
x = y = range(-3, 3, length=100)
z = @. exp(-x'^2 .- y^2)
surface(x, y, z; colorbar=false, size=(600, 600), color=:jet)

# %%

0004/v1.7 findmin and Plots plotly example.jl


#=

  Linear regression

  This is actually a WebPPL example (example/linear_regression.wppl) but translated via my BLOG model.

  xs = [0,1,2,3]
  ys = [0,1,4,6]

  mu should be ~ 2 and b ~ 0.
  y4: y for x=4 should be 2*4+0 = 8
  y5: y for x=5 should be 2*5+0 = 10

  Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

           m    2.0134    0.3960     0.0040    0.0088   1980.8480    1.0001      445.5349
           b   -0.2678    0.7118     0.0071    0.0135   2785.0095    0.9999      626.4079
       sigma    0.8463    0.4970     0.0050    0.0155    929.0047    0.9999      208.9529
          y4    7.7715    1.5278     0.0153    0.0314   2240.9116    1.0001      504.0287
          y5    9.7916    1.8747     0.0187    0.0401   2060.8303    1.0000      463.5246

  With 5 data points (x=[0,1,2,3,4] and y=[0,1,4,6,8]) with got a much better result:
  Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

           m    2.0000    0.0001     0.0000    0.0000    684.0973    0.9999      170.6830
           b    0.0000    0.0004     0.0000    0.0000   1140.3205    1.0005      284.5111
       sigma    0.0002    0.0003     0.0000    0.0000    111.5388    1.0486       27.8290
          y4    8.0000    0.0005     0.0000    0.0000    713.2871    1.0046      177.9658
          y5   10.0000    0.0006     0.0000    0.0000    553.0126    1.0016      137.9772

  cf ~/blog/linear_regression2.blog
     ~/webppl/linear_regression2.wppl
=#

using Turing
include("jl_utils.jl")

@model function linear_regression2(xs,ys)
    n = length(xs)

    m ~ Normal(0,2)
    b ~ Normal(0,2)
    sigma ~ Gamma(1,1)

    xs ~ filldist(Uniform(0,14),n)

    for i in 1:n
        ys[i] ~ Normal(m*xs[i] + b, sigma)
    end

    # Prediction for x=4 and x=5
    y4 ~ Normal(m*4 + b, sigma)
    y5 ~ Normal(m*5 + b, sigma)

end

xs = [0,1,2,3]
ys = [0,1,4,6]
n = 4
xs = [i for i in 0:n]
ys = [2*x for x in xs]
println("xs: ", xs)
println("ys: ", ys)
model = linear_regression2(xs,ys)

# chns = sample(model, Prior(), 100_000)
# chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 100_000)
chns = sample(model, NUTS(), 10_000)

display(chns)


"""
    gram_schmidt(A::Array[, ret_coef_mat::Bool])

Return orthogonalized basis.
Optionally returns the Gram-Schmidt coefficient matrix using the ret_coef_mat parameter.

# Example

```julia
julia> x = [
        [ 3  -1   5]
        [-5   2  -1]
        [-3   9   2]
    ]

julia> gram_schmidt(x)
3×3 Array{Float64,2}:
  3.0      -1.0       5.0
 -3.11429   1.37143   2.14286
  2.68728   6.5689   -0.298587

```
"""
function gram_schmidt{T<:AbstractFloat}(A::Array{T, 2}; ret_coef_mat::Bool = false)
    # Assume the vectors are given as row vectors of X
    # Maybe redo later since vectors in Julia are stored in column-major order.

    X = transpose(copy(A))
    X_star = zeros(X)

    # Coefficient matrix
    μ = zeros(X)

    n_cols = size(X)[2]
    for i in 1:n_cols
        X_star[:, i] = X[:, i]
        for j in 1:i-1
            t = dot(X[:, i], X_star[:, j]) / dot(X_star[:, j], X_star[:, j])
            μ[j, i] = t
            X_star[:, i] -= (t * X_star[:, j])
        end
        μ[i,i] = 1
    end

    if ret_coef_mat
        transpose(X_star), transpose(μ)
    else
        transpose(X_star)
    end
end

function gram_schmidt{T<:Integer}(A::Array{T, 2}; ret_coef_mat::Bool = false)
    gram_schmidt(float(A), ret_coef_mat=ret_coef_mat)
end


using Statistics, StatsBase, Random, LinearAlgebra
Random.seed!(1)

function cmHitTime()
    catIndex, mouseIndex, t = 1, 5, 0
        while catIndex != mouseIndex
            catIndex += catIndex == 1 ? 1 : rand([-1,1])
            mouseIndex += mouseIndex == 5 ? -1 : rand([-1,1])
            t += 1
        end
    return t
end
function mcTraj(P,initState,T,stopState=0)
    n = size(P)[1]
    state = initState
    traj = [state]
    for t in 1:T-1
        state = sample(1:n,weights(P[state,:]))
        push!(traj,state)
        if state == stopState
            break
        end
    end
    return traj
end
N = 10^6
P = [   0   1   0   0   0;
        1/4 0   1/4 1/4 1/4;
        0   1/2 0   0   1/2;
        0   1/2 0   0   1/2;
        0   0   0   0   1]

theor = [1 0 0 0] * (inv(I - P[1:4,1:4])*ones(4))
est1 = mean([cmHitTime() for _ in 1:N])
est2 = mean([length(mcTraj(P,1,10^6,5))-1 for _ in 1:N])

P[5,:] = [1 0 0 0 0]
pi5 = sum(mcTraj(P,1,N) .== 5)/N
est3 = 1/pi5 - 1

println("Theoretical: ", theor)
println("Estimate 1: ",est1)
println("Estimate 2: ",est2)
println("Estimate 3: ",est3)


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Plots
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.01"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	lambda  = 1/4.5
	expDist = Exponential(1/lambda)
	n, N    = 10, 10^6

	means     = Array{Float64}(undef, N)
	variances = Array{Float64}(undef, N)

	for i in 1:N
		data = rand(expDist,n)
		means[i] = mean(data)
		variances[i] = var(data)
	end
end

# ╔═╡ 78fd898a-2527-11eb-048f-3be84c5d8649
("Actual mean: ",mean(expDist),
		"Mean of sample means: ",mean(means))

# ╔═╡ 78fdc0d2-2527-11eb-3fe8-5738d4af3806
("Actual variance: ",var(expDist),
	"Mean of sample variances: ",mean(variances))

# ╔═╡ 78fe5bf8-2527-11eb-2c19-1d2e7d1277af
begin
	stephist(means, bins=200, c=:blue, normed=true,
		label="Histogram of Sample Means")
	stephist!(variances, bins=600, c=:red, normed=true,
		label="Histogram of Sample Variances", xlims=(0,40), ylims=(0,0.4),
			xlabel = "Statistic value", ylabel = "Density")
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.01"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═78fd898a-2527-11eb-048f-3be84c5d8649
# ╠═78fdc0d2-2527-11eb-3fe8-5738d4af3806
# ╠═78fe5bf8-2527-11eb-2c19-1d2e7d1277af
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


# coordinates.jl
# Maps having to do with coordinate transforms.

"""
A Cartesion to Polar map. First dimension is interpreted as radial distance, second as an angle.
The unit circle is mapped to a square [-1,1]x[-1,1]

"""
struct CartToPolarMap{T} <: AbstractMap{SVector{2,T}, SVector{2,T}}
end

(m::CartToPolarMap)(x) = applymap(m, x)

applymap(map::CartToPolarMap, x::SVector{2,T}) where {T} = SVector{2,T}(sqrt(x[1]^2+x[2]^2)*2-1, atan(x[2],x[1])/pi)

inv(map::CartToPolarMap{T}) where {T} = PolarToCartMap{T}()

"""
A Polar to Cartesian map. The angle is mapped to the second dimension, radius to the first.
A square [-1,1]x[-1,1] is mapped to the unit circle
"""
struct PolarToCartMap{T} <: AbstractMap{SVector{2,T}, SVector{2,T}}
end

(m::PolarToCartMap)(x) = applymap(m, x)

applymap(map::PolarToCartMap, x::SVector{2,T}) where {T} = SVector{2,T}((x[1]+1)/2*cos(pi*x[2]), (x[1]+1)/2*sin(pi*x[2]))

inv(map::PolarToCartMap{T}) where {T} = CartToPolarMap{T}()


# Cost functions
"""
    mae(ŷ, y)

Return the mean of absolute error `sum(abs.(ŷ .- y)) / length(y)`
"""
mae(ŷ, y) = sum(abs.(ŷ .- y)) * 1 // length(y)

"""
    mse(ŷ, y)

Return the mean squared error `sum((ŷ .- y).^2) / length(y)`.
"""
mse(ŷ, y) = sum((ŷ .- y).^2) * 1 // length(y)

"""
    msle(ŷ, y; ϵ=eps(eltype(ŷ)))

Returns the mean of the squared logarithmic errors `sum((log.(ŷ .+ ϵ) .- log.(y .+ ϵ)).^2) / length(y)`.
The `ϵ` term provides numerical stability.

This error penalizes an under-predicted estimate greater than an over-predicted estimate.
"""
msle(ŷ, y; ϵ=eps(eltype(ŷ))) = sum((log.(ŷ .+ ϵ) .- log.(y .+ ϵ)).^2) * 1 // length(y)

"""
    huber_loss(ŷ, y; δ=1.0)

Computes the mean of the Huber loss given the prediction `ŷ` and true values `y`. By default, δ is set to 1.0.

                    | 0.5*|ŷ - y|,   for |ŷ - y| <= δ
      Hubber loss = |
                    |  δ*(|ŷ - y| - 0.5*δ),  otherwise

[`Huber Loss`](https://en.wikipedia.org/wiki/Huber_loss).
"""
function huber_loss(ŷ, y;  δ=eltype(ŷ)(1))
   abs_error = abs.(ŷ .- y)
   temp = abs_error .<  δ
   x = eltype(ŷ)(0.5)
   hub_loss = sum(((abs_error.^2) .* temp) .* x .+ δ*(abs_error .- x*δ) .* (1 .- temp)) * 1 // length(y)
end

function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Nothing)
  return -sum(y .* log.(ŷ)) * 1 // size(y, 2)
end

function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Number)
  return -sum(y .* log.(ŷ)) .* weight * 1 // size(y, 2)
end

function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::AbstractVector)
  return -sum(y .* log.(ŷ) .* weight) * 1 // size(y, 2)
end

"""
    crossentropy(ŷ, y; weight=1)

Return the crossentropy computed as `-sum(y .* log.(ŷ) .* weight) / size(y, 2)`.

See also [`logitcrossentropy`](@ref), [`binarycrossentropy`](@ref).
"""
crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight=nothing) = _crossentropy(ŷ, y, weight)

"""
    logitcrossentropy(ŷ, y; weight=1)

Return the crossentropy computed after a [softmax](@ref) operation:

  -sum(y .* logsoftmax(ŷ) .* weight) / size(y, 2)

See also [`crossentropy`](@ref), [`binarycrossentropy`](@ref).
"""
function logitcrossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight = 1)
  return -sum(y .* logsoftmax(ŷ) .* weight) * 1 // size(y, 2)
end

"""
    binarycrossentropy(ŷ, y; ϵ=eps(ŷ))

Return `-y*log(ŷ + ϵ) - (1-y)*log(1-ŷ + ϵ)`. The ϵ term provides numerical stability.

Typically, the prediction `ŷ` is given by the output of a [`sigmoid`](@ref) activation.
"""
binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)

# Re-definition to fix interaction with CuArrays.
CuArrays.@cufunc binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)

"""
    logitbinarycrossentropy(ŷ, y)

`logitbinarycrossentropy(ŷ, y)` is mathematically equivalent to `binarycrossentropy(σ(ŷ), y)`
but it is more numerically stable.

See also [`binarycrossentropy`](@ref), [`sigmoid`](@ref), [`logsigmoid`](@ref).
"""
logitbinarycrossentropy(ŷ, y) = (1 - y)*ŷ - logσ(ŷ)

# Re-definition to fix interaction with CuArrays.
CuArrays.@cufunc logitbinarycrossentropy(ŷ, y) = (1 - y)*ŷ - logσ(ŷ)

"""
    normalise(x; dims=1)

Normalises `x` to mean 0 and standard deviation 1, across the dimensions given by `dims`. Defaults to normalising over columns.

```julia-repl
julia> a = reshape(collect(1:9), 3, 3)
3×3 Array{Int64,2}:
  1  4  7
  2  5  8
  3  6  9

julia> normalise(a)
3×3 Array{Float64,2}:
  -1.22474  -1.22474  -1.22474
  0.0       0.0       0.0
  1.22474   1.22474   1.22474

julia> normalise(a, dims=2)
3×3 Array{Float64,2}:
  -1.22474  0.0  1.22474
  -1.22474  0.0  1.22474
  -1.22474  0.0  1.22474
```
"""
function normalise(x::AbstractArray; dims=1)
  μ′ = mean(x, dims = dims)
  σ′ = std(x, dims = dims, mean = μ′, corrected=false)
  return (x .- μ′) ./ σ′
end

"""
    kldivergence(ŷ, y)

KLDivergence is a measure of how much one probability distribution is different from the other.
It is always non-negative and zero only when both the distributions are equal everywhere.

[KL Divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).
"""
function kldivergence(ŷ, y)
  entropy = sum(y .* log.(y)) * 1 //size(y,2)
  cross_entropy = crossentropy(ŷ, y)
  return entropy + cross_entropy
end

"""
    poisson(ŷ, y)

Poisson loss function is a measure of how the predicted distribution diverges from the expected distribution.
Returns `sum(ŷ .- y .* log.(ŷ)) / size(y, 2)`

[Poisson Loss](https://peltarion.com/knowledge-center/documentation/modeling-view/build-an-ai-model/loss-functions/poisson).
"""
poisson(ŷ, y) = sum(ŷ .- y .* log.(ŷ)) * 1 // size(y,2)

"""
    hinge(ŷ, y)

Measures the loss given the prediction `ŷ` and true labels `y` (containing 1 or -1).
Returns `sum((max.(0, 1 .- ŷ .* y))) / size(y, 2)`

[Hinge Loss](https://en.wikipedia.org/wiki/Hinge_loss)
See also [`squared_hinge`](@ref).
"""
hinge(ŷ, y) = sum(max.(0, 1 .-  ŷ .* y)) * 1 // size(y, 2)

"""
    squared_hinge(ŷ, y)

Computes squared hinge loss given the prediction `ŷ` and true labels `y` (conatining 1 or -1).
Returns `sum((max.(0, 1 .- ŷ .* y)).^2) / size(y, 2)`

See also [`hinge`](@ref).
"""
squared_hinge(ŷ, y) = sum((max.(0, 1 .- ŷ .* y)).^2) * 1 // size(y, 2)

"""
    dice_coeff_loss(ŷ, y; smooth=1)

Loss function used in Image Segmentation. Calculates loss based on dice coefficient. Similar to F1_score.
Returns `1 - 2*sum(|ŷ .* y| + smooth) / (sum(ŷ.^2) + sum(y.^2) + smooth)`

[V-Net: Fully Convolutional Neural Networks forVolumetric Medical Image Segmentation](https://arxiv.org/pdf/1606.04797v1.pdf)
"""
dice_coeff_loss(ŷ, y; smooth=eltype(ŷ)(1.0)) = 1 - (2*sum(y .* ŷ) + smooth) / (sum(y.^2) + sum(ŷ.^2) + smooth)

"""
    tversky_loss(ŷ, y; β=0.7)

Used with imbalanced data to give more weightage to False negatives.
Larger β weigh recall higher than precision (by placing more emphasis on false negatives)
Returns `1 - sum(|y .* ŷ| + 1) / (sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1)`

[Tversky loss function for image segmentation using 3D fully convolutional deep networks](https://arxiv.org/pdf/1706.05721.pdf)
"""
tversky_loss(ŷ, y; β=eltype(ŷ)(0.7)) = 1 - (sum(y .* ŷ) + 1) / (sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1)

"""
    flatten(x::AbstractArray)

Transforms (w,h,c,b)-shaped input into (w x h x c,b)-shaped output,
by linearizing all values for each element in the batch.
"""
function flatten(x::AbstractArray)
  return reshape(x, :, size(x)[end])
end


#=
# 3. Non-uniqueness and discrete circulation

There is no unique solution to the two-dimensional potential flow problem in the presence of an impenetrable surface, since one can choose any value for the uniform value $\mathfrak{s}_0$ and still enforce the no-penetration condition. Equivalently, we can specify any circulation about the body and still enforce this condition. For later use, let us write this uniform surface streamfunction as $\mathfrak{s}_0 = s_0 \mathfrak{1}$, where $s_0$ is a single scalar value and $\mathfrak{1}$ is a vector of ones on all surface points. The discrete circulation $\Gamma_b$ about the body is given by the sum of the bound vortex sheet data and can be written compactly as

$\Gamma_b = \mathfrak{1}^{T} \mathfrak{f}.$

In `GridPotentialFlow.jl`, the user cannot specify $s_0$ directly. Instead, the package solves the system

$\begin{bmatrix}
\mathsf{L} & \mathsf{R} & 0\\
\mathsf{E} &  0 & \mathfrak{1}\\
0 & \mathfrak{1}^{T} & 0\\
\end{bmatrix} \begin{pmatrix} \mathsf{s} \\ \mathfrak{f} \\ s_0 \end{pmatrix} =
\begin{pmatrix} -\mathsf{w} \\ \mathfrak{s}'_b \\ \Gamma_b \end{pmatrix}.$

and the circulation can be specified using the `Γ` keyword for each body. If it is not specified, it is set to zero. The return value of `solvesystem` has a field `ψ₀`, whose entries represent the value for $s_0$ for each body, which are in this case the Lagrange multipliers to enforce the specified circulation.
=#

# We now illustrate how to specify the circulation by simulating the flow around two cylinders that have different circulations.

#md # ```@setup 3.-Non-uniqueness-and-discrete-circulation
#md # using GridPotentialFlow
#md # using Plots
#md # ```
#!md using GridPotentialFlow
#!md using Plots
Δx = 0.01
Lx = 2.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx);

# The two cylinders are positioned using the tools from `RigidTransform` from `RigidBodyTools.jl` and the circulation of the first one is set to a non-zero value.
R = Lx/8;
b_left = PotentialFlowBody(Circle(R,Δx),Γ=1.0)
T = RigidTransform((-Lx/4,0.0),0.0)
T(b_left);
b_right = PotentialFlowBody(Circle(R,Δx))
T = RigidTransform((Lx/4,0.0),0.0)
T(b_right);

# A plot of the streamfunction then clearly shows the non-zero circulation about the left cylinder.
model = VortexModel(g,bodies=[b_left,b_right])
s = streamfunction(model)
plot(s,g)
plot!(b_left,fillcolor=:black,fillrange=0,fillalpha=0.25,linecolor=:black,linewidth=2)
plot!(b_right,fillcolor=:black,fillrange=0,fillalpha=0.25,linecolor=:black,linewidth=2)

#=
When the previously introduced saddle point system

$\begin{bmatrix}
\mathsf{L} & \mathsf{R} \\
\mathsf{E} &  0
\end{bmatrix} \begin{pmatrix} \mathsf{s} \\ \mathfrak{f} \end{pmatrix} =
\begin{pmatrix} -\mathsf{w} \\ \mathfrak{s}'_b - \mathfrak{s}_0 \end{pmatrix}.$

is solved, the resulting circulation is

$\Gamma_{b} = \mathfrak{1}^{T} \mathsf{S}^{-1}\left(\mathfrak{s}_{b}^{\prime}+\mathsf{EL}^{-1} \mathsf{w}\right)-s_{0} \mathfrak{1}^{T} \mathsf{S}^{-1} \mathfrak{1},$

where $\mathsf{S} = -\mathsf{EL}^{-1}\mathsf{R}$ is the Schur complement. The scalar factor $\mathfrak{1}^{T} \mathsf{S}^{-1} \mathfrak{1}$ in this expression is a property of the set of points and their immersion into the Cartesian grid. Part of this factor, $\mathsf{S}^{-1} \mathfrak{1}$, represents the bound vortex sheet strength associated with a uniform, unit-strength streamfunction on the surface. This sheet has a particularly important role in some of the discussion to follow, so we will denote its strength by $\mathfrak{f}_0$:

$\mathfrak{f}_0 \equiv \mathsf{S}^{-1} \mathfrak{1}.$

We demonstrate below the distribution of $\mathfrak{f}_0$ for elliptical cylinders with different aspect ratios.
=#

Δx = 0.004
Lx = 2.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx);

# Let us consider four cylinders with aspect ratios 1, 2, 3, and infinity. An ellipse with infinite aspect ratio is a flat plate, so we construct that case separately.
a = Lx/4 # semi-major axis
AR_list = [1.0,2.0,3.0]
ellipses = [Ellipse(a,a/AR,Δx) for AR in AR_list]
bodies = PotentialFlowBody[]
push!(bodies,[PotentialFlowBody(SplinedBody(hcat(e.x,e.y),3*cellsize(g))) for e in ellipses]...)
push!(bodies,PotentialFlowBody(Plate(2*a,3*Δx)));
colors = [:blue,:red,:green,:black]
plot()
for i in 1:length(bodies)
    plot!(bodies[i],fillrange=0,fillalpha=0.0,linecolor=colors[i],linewidth=2)
end
plot!(xlim=xlim,ylim=ylim,xlabel="x",ylabel="y")

# To compute $\mathfrak{f}_0$, we have to access the internal fields of the associated `system` of our vortexmodel.
f₀_list = []
for i in 1:length(bodies)
    Δs = dlength(bodies[i]);
    model = VortexModel(g,bodies=[bodies[i]]);
    ones = ScalarData(length(bodies[i]))
    ones .= 1.0
    f₀ = model.system.ibp.Sfact\ones
    push!(f₀_list,f₀)
end

# Because the flat plate is an open body, we need to mirror and append its contents to itself to compare it to the ellipses. Furthermore, its bound vortex sheet strength should be divided by two because it represents $\gamma = \gamma^{+} + \gamma^{-}$
append!(f₀_list[end],reverse!(f₀_list[end]));
Δs_plate = dlengthmid(bodies[end]);
append!(Δs_plate,reverse!(Δs_plate));
plot(ylim=(-2,0),xlabel="p/N",title="f₀ for ellipses of different aspect ratios")
for i in 1:length(bodies)-1
    plot!((1:length(f₀_list[i]))/length(f₀_list[i]),f₀_list[i]./dlength(bodies[i]),linecolor=colors[i],label="AR=$(AR_list[i])")
end
plot!((1:length(f₀_list[end]))/length(f₀_list[end]),0.5*f₀_list[end]./Δs_plate,linecolor=colors[end],label="AR=inf")


# # Train Kernel Parameters

# Here we show a few ways to train (optimize) the kernel (hyper)parameters at the example of kernel-based regression using KernelFunctions.jl.
# All options are functionally identical, but differ a little in readability, dependencies, and computational cost.

# We load KernelFunctions and some other packages. Note that while we use `Zygote` for automatic differentiation and `Flux.optimise` for optimization, you should be able to replace them with your favourite autodiff framework or optimizer.

using KernelFunctions
using LinearAlgebra
using Distributions
using Plots
using BenchmarkTools
using Flux
using Flux: Optimise
using Zygote
using Random: seed!
seed!(42);

# ## Data Generation
# We generate a toy dataset in 1 dimension:

xmin, xmax = -3, 3  # Bounds of the data
N = 50 # Number of samples
x_train = rand(Uniform(xmin, xmax), N)  # sample the inputs
σ = 0.1
y_train = sinc.(x_train) + randn(N) * σ  # evaluate a function and add some noise
x_test = range(xmin - 0.1, xmax + 0.1; length=300)
nothing #hide

# Plot the data

scatter(x_train, y_train; label="data")
plot!(x_test, sinc; label="true function")

# ## Manual Approach
# The first option is to rebuild the parametrized kernel from a vector of parameters
# in each evaluation of the cost function. This is similar to the approach taken in
# [Stheno.jl](https://github.com/JuliaGaussianProcesses/Stheno.jl).

# To train the kernel parameters via [Zygote.jl](https://github.com/FluxML/Zygote.jl),
# we need to create a function creating a kernel from an array.
# A simple way to ensure that the kernel parameters are positive
# is to optimize over the logarithm of the parameters.

function kernelcall(θ)
    return (exp(θ[1]) * SqExponentialKernel() + exp(θ[2]) * Matern32Kernel()) ∘
           ScaleTransform(exp(θ[3]))
end
nothing #hide

# From theory we know the prediction for a test set x given
# the kernel parameters and normalization constant:

function f(x, x_train, y_train, θ)
    k = kernelcall(θ[1:3])
    return kernelmatrix(k, x, x_train) *
           ((kernelmatrix(k, x_train) + exp(θ[4]) * I) \ y_train)
end
nothing #hide

# Let's look at our prediction.
# With starting parameters `p0` (picked so we get the right local
# minimum for demonstration) we get:

p0 = [1.1, 0.1, 0.01, 0.001]
θ = log.(p0)
ŷ = f(x_test, x_train, y_train, θ)
scatter(x_train, y_train; label="data")
plot!(x_test, sinc; label="true function")
plot!(x_test, ŷ; label="prediction")

# We define the following loss:

function loss(θ)
    ŷ = f(x_train, x_train, y_train, θ)
    return norm(y_train - ŷ) + exp(θ[4]) * norm(ŷ)
end
nothing #hide

# The loss with our starting point:

loss(θ)

# Computational cost for one step:

@benchmark let
    θ = log.(p0)
    opt = Optimise.ADAGrad(0.5)
    grads = only((Zygote.gradient(loss, θ)))
    Optimise.update!(opt, θ, grads)
end

# ### Training the model

# Setting an initial value and initializing the optimizer:
θ = log.(p0) # Initial vector
opt = Optimise.ADAGrad(0.5)
nothing #hide

# Optimize

anim = Animation()
for i in 1:15
    grads = only((Zygote.gradient(loss, θ)))
    Optimise.update!(opt, θ, grads)
    scatter(
        x_train, y_train; lab="data", title="i = $(i), Loss = $(round(loss(θ), digits = 4))"
    )
    plot!(x_test, sinc; lab="true function")
    plot!(x_test, f(x_test, x_train, y_train, θ); lab="Prediction", lw=3.0)
    frame(anim)
end
gif(anim, "train-kernel-param.gif"; show_msg=false, fps=15);
nothing; #hide

# ![](train-kernel-param.gif)

# Final loss
loss(θ)

# ## Using ParameterHandling.jl
# Alternatively, we can use the [ParameterHandling.jl](https://github.com/invenia/ParameterHandling.jl) package
# to handle the requirement that all kernel parameters should be positive.
# The package also allows arbitrarily nesting named tuples that make the parameters
# more human readable, without having to remember their position in a flat vector.

using ParameterHandling

raw_initial_θ = (
    k1=positive(1.1), k2=positive(0.1), k3=positive(0.01), noise_var=positive(0.001)
)

flat_θ, unflatten = ParameterHandling.value_flatten(raw_initial_θ)
flat_θ #hide

# We define a few relevant functions and note that compared to the previous `kernelcall` function, we do not need explicit `exp`s.

function kernelcall(θ)
    return (θ.k1 * SqExponentialKernel() + θ.k2 * Matern32Kernel()) ∘ ScaleTransform(θ.k3)
end
nothing #hide

function f(x, x_train, y_train, θ)
    k = kernelcall(θ)
    return kernelmatrix(k, x, x_train) *
           ((kernelmatrix(k, x_train) + θ.noise_var * I) \ y_train)
end
nothing #hide

function loss(θ)
    ŷ = f(x_train, x_train, y_train, θ)
    return norm(y_train - ŷ) + θ.noise_var * norm(ŷ)
end
nothing #hide

initial_θ = ParameterHandling.value(raw_initial_θ)
nothing #hide

# The loss at the initial parameter values:

(loss ∘ unflatten)(flat_θ)

# Cost per step

@benchmark let
    θ = flat_θ[:]
    opt = Optimise.ADAGrad(0.5)
    grads = (Zygote.gradient(loss ∘ unflatten, θ))[1]
    Optimise.update!(opt, θ, grads)
end

# ### Training the model

# Optimize

opt = Optimise.ADAGrad(0.5)
for i in 1:15
    grads = (Zygote.gradient(loss ∘ unflatten, flat_θ))[1]
    Optimise.update!(opt, flat_θ, grads)
end
nothing #hide

# Final loss

(loss ∘ unflatten)(flat_θ)

# ## Flux.destructure
# If we don't want to write an explicit function to construct the kernel, we can alternatively use the `Flux.destructure` function.
# Again, we need to ensure that the parameters are positive. Note that the `exp` function is now part of the loss function, instead of part of the kernel construction.

# We could also use ParameterHandling.jl here.
# To do so, one would remove the `exp`s from the loss function below and call `loss ∘ unflatten` as above.

θ = [1.1, 0.1, 0.01, 0.001]

kernel = (θ[1] * SqExponentialKernel() + θ[2] * Matern32Kernel()) ∘ ScaleTransform(θ[3])

params, kernelc = Flux.destructure(kernel);

# This returns the trainable `params` of the kernel and a function to reconstruct the kernel.
kernelc(params)

# From theory we know the prediction for a test set x given
# the kernel parameters and normalization constant

function f(x, x_train, y_train, θ)
    k = kernelc(θ[1:3])
    return kernelmatrix(k, x, x_train) * ((kernelmatrix(k, x_train) + (θ[4]) * I) \ y_train)
end
nothing #hide

function loss(θ)
    ŷ = f(x_train, x_train, y_train, exp.(θ))
    return norm(y_train - ŷ) + exp(θ[4]) * norm(ŷ)
end
nothing #hide

# Cost for one step

@benchmark let θt = θ[:], optt = Optimise.ADAGrad(0.5)
    grads = only((Zygote.gradient(loss, θt)))
    Optimise.update!(optt, θt, grads)
end

# ### Training the model

# The loss at our initial parameter values:
θ = log.([1.1, 0.1, 0.01, 0.001]) # Initial vector
loss(θ)

# Initialize optimizer

opt = Optimise.ADAGrad(0.5)
nothing #hide

# Optimize

for i in 1:15
    grads = only((Zygote.gradient(loss, θ)))
    Optimise.update!(opt, θ, grads)
end
nothing #hide

# Final loss
loss(θ)


# https://turing.ml/dev/tutorials/07-poisson-regression/

# Import Turing, Distributions and DataFrames
using Turing, Distributions, DataFrames, Distributed

# Import MCMCChain, Plots, and StatsPlots for visualizations and diagnostics.
using MCMCChains, Plots, StatsPlots

# Set a seed for reproducibility.
using Random
Random.seed!(99);

# Turn off progress monitor.
Turing.setprogress!(false)

# The true Poisson parameter (mean) of toy data
theta_noalcohol_meds = 1    # no alcohol, took medicine
theta_alcohol_meds = 3      # alcohol, took medicine
theta_noalcohol_nomeds = 6  # no alcohol, no medicine
theta_alcohol_nomeds = 36   # alcohol, no medicine

# no of samples for each of the above cases
q = 100

# Generate data from different Poisson distributions
noalcohol_meds = Poisson(theta_noalcohol_meds)
alcohol_meds = Poisson(theta_alcohol_meds)
noalcohol_nomeds = Poisson(theta_noalcohol_nomeds)
alcohol_nomeds = Poisson(theta_alcohol_nomeds)

nsneeze_data = vcat(rand(noalcohol_meds, q), rand(alcohol_meds, q), rand(noalcohol_nomeds, q), rand(alcohol_nomeds, q))
alcohol_data = vcat(zeros(q), ones(q), zeros(q), ones(q))
meds_data = vcat(zeros(q), zeros(q), ones(q), ones(q))

df = DataFrame(nsneeze = nsneeze_data, alcohol_taken = alchol_data, nomeds_taken = meds_data, product_alcohol_meds = meds_data.*alcohol_data)
df[sample(1:nrow(df), 5, replace = false), :]

# Data Plotting
p1 = Plots.histogram(df[(df[:,:alcohol_taken] .== 0) .& (df[:,:nomeds_taken] .== 0), 1], title = "no_alcohol+meds")
p2 = Plots.histogram((df[(df[:,:alcohol_taken] .== 1) .& (df[:,:nomeds_taken] .== 0), 1]), title = "alcohol+meds")
p3 = Plots.histogram((df[(df[:,:alcohol_taken] .== 0) .& (df[:,:nomeds_taken] .== 1), 1]), title = "no_alcohol+no_meds")
p4 = Plots.histogram((df[(df[:,:alcohol_taken] .== 1) .& (df[:,:nomeds_taken] .== 1), 1]), title = "alcohol+no_meds")
plot(p1, p2, p3, p4, layout = (2, 2), legend = false)

# Convert the DataFrame object to matrices.
data = Matrix(df[:, [:alcohol_taken, :nomeds_taken, :product_alcohol_meds]])
data_labels = df[:, :nsneeze]
data

# Rescale our matrices to help Turing sampler
# in initializing the parameter estimates(normalizing).
data = (data .- mean(data, dims=1)) ./ std(data, dims=1)

@model poisson_regression(x, y, n, σ²) = begin
    b0 ~ Normal(0, σ²)
    b1 ~ Normal(0, σ²)
    b2 ~ Normal(0, σ²)
    b3 ~ Normal(0, σ²)
    for i = 1:n
        theta = b0 + b1*x[i, 1] + b2*x[i,2] + b3*x[i,3]
        y[i] ~ Poisson(exp(theta))
    end
end;

# Retrieve the number of observations.
n, _ = size(data)

# Sample using NUTS.
num_chains = 4
m = poisson_regression(data, data_labels, n, 10)
chain = sample(m, NUTS(200, 0.65), MCMCThreads(), 2_500, num_chains; discard_adapt=false)

# check whether our chains have converged
gelmandiag(chain)

# Taking the first chain
c1 = chain[:,:,1]

# Calculating the exponentiated means
b0_exp = exp(mean(c1[:b0]))
b1_exp = exp(mean(c1[:b1]))
b2_exp = exp(mean(c1[:b2]))
b3_exp = exp(mean(c1[:b3]))

print("The exponent of the meaned values of the weights (or coefficients are): \n")
print("b0: ", b0_exp, " \n", "b1: ", b1_exp, " \n", "b2: ", b2_exp, " \n", "b3: ", b3_exp, " \n")
print("The posterior distributions obtained after sampling can be visualised as :\n")

plot(chain)

# removing warmup samples

# Note the standard deviation before removing the warmup samples
describe(chain)

# Removing the first 200 values of the chains.
chains_new = chain[201:2500,:,:]
describe(chains_new)

plot(chains_new)


# Conjugate

export Conjugate

"""
**Convex conjugate**

    Conjugate(f)

Returns the convex conjugate (also known as Fenchel conjugate, or Fenchel-Legendre transform) of function `f`, that is
```math
f^*(x) = \\sup_y \\{ \\langle y, x \\rangle - f(y) \\}.
```
"""
struct Conjugate{T <: ProximableFunction} <: ProximableFunction
  f::T
  function Conjugate{T}(f::T) where {T<: ProximableFunction}
    if is_convex(f) == false
      error("`f` must be convex")
    end
    new(f)
  end
end

is_prox_accurate(f::Conjugate) = is_prox_accurate(f.f)
is_convex(f::Conjugate) = true
is_cone(f::Conjugate) = is_cone(f.f) && is_convex(f.f)
is_smooth(f::Conjugate) = is_strongly_convex(f.f)
is_strongly_convex(f::Conjugate) = is_smooth(f.f)
is_quadratic(f::Conjugate) = is_strongly_convex(f.f) && is_generalized_quadratic(f.f)
is_generalized_quadratic(f::Conjugate) = is_quadratic(f.f)

fun_dom(f::Conjugate) = fun_dom(f.f)

Conjugate(f::T) where {T <: ProximableFunction} = Conjugate{T}(f)

# only prox! is provided here, call method would require being able to compute
# an element of the subdifferential of the conjugate

function prox!(y::AbstractArray{R}, g::Conjugate, x::AbstractArray{R}, gamma::R=one(R)) where R <: Real
  # Moreau identity
  v = prox!(y, g.f, x/gamma, one(R)/gamma)
  if is_set(g)
    v = zero(R)
  else
    v = dot(x,y) - gamma*dot(y,y) - v
  end
  for k in eachindex(y)
    y[k] = x[k] - gamma*y[k]
  end
  return v
end

# complex case, need to cast inner products to real

function prox!(y::AbstractArray{Complex{T}}, g::Conjugate, x::AbstractArray{Complex{T}}, gamma::R=one(R)) where {R <: Real, T <: RealOrComplex{R}}
  v = prox!(y, g.f, x/gamma, one(R)/gamma)
  if is_set(g)
    v = zero(R)
  else
    v = real(dot(x,y)) - gamma*real(dot(y,y)) - v
  end
  for k in eachindex(y)
    y[k] = x[k] - gamma*y[k]
  end
  return v
end

# naive implementation

function prox_naive(g::Conjugate, x::AbstractArray{T}, gamma::Real=1.0) where T <: RealOrComplex
  y, v = prox_naive(g.f, x/gamma, 1.0/gamma)
  return x - gamma*y, real(dot(x,y)) - gamma*real(dot(y,y)) - v
end


"""
    diffusion_matrix(domain, n)

Assemble discrete diffusion matrix.
"""
function diffusion_matrix end

function diffusion_matrix(domain::ClosedIntervalDomain, n)
    Δx = 1 // n * (domain.right - domain.left)
    inds = [-1, 0, 1]
    stencil = [1, -2, 1] / Δx^2
    diags = [i => fill(s, n + 1 - abs(i)) for (i, s) ∈ zip(inds, stencil)]
    D = spdiagm(diags...)
    D[1, 1] = stencil[1] + stencil[2]
    D[end, end] = stencil[end-1] + stencil[end]
    D
end

function diffusion_matrix(domain::PeriodicIntervalDomain, n)
    Δx = 1 // n * (domain.right - domain.left)
    inds = [-1, 0, 1]
    stencil = [1, -2, 1] / Δx^2
    diags = [i => fill(s, n - abs(i)) for (i, s) ∈ zip(inds, stencil)]
    D = spdiagm(diags...)
    D[1, end] = stencil[1]
    D[end, 1] = stencil[end]
    D
end


# this file depends on jacobi.jl from JuliaOptics/OpticsPolynomials that
# should be found in the same directory

export zernike
export zernike_series
export zernike_sum
export zernike_norm
export zernike_nm_to_fringe
export zernike_nm_to_ansi_j
export zernike_ansi_j_to_nm
export zernike_noll_to_nm
export zernike_fringe_to_nm
export zernike_zero_separation

export apply_weights

"""
    apply_weights(data, weights)

Given data of arbitrary dimensionality with multiple realizations or modes along
the final dimension, weight the final dimension by weights and return the sum.

Will have one fewer dimension than the input data.
"""
function apply_weights(data, weights)
    shape = size(data)
    prod = reshape(data, (:, shape[end])) * weights
    return reshape(prod, shape[1: end-1])
end

function kronecker(i,j)
    return i==j ? 1 : 0
end

"""
    zernike_norm(n, m)

Norm of Zernike polynomial of radial order n, azimuthal order m.

The norm is the average squared distance to zero.  By multiplying a zernike
value by the norm, the term is given unit stdev or RMS.
"""
function zernike_norm(n::Integer, m::Integer)
    return √(2 * (n+1)) / (1 + kronecker(m, 0))
end

"""
    zernike_nm_to_fringe(n, m)

Map (n,m) ANSI indices to a single fringe index.
"""
function zernike_nm_to_fringe(n::Integer, m::Integer)
    term1 = (1 + (n + abs(m))/2)^2
    term2 = 2 * abs(m)
    term3 = (1 + sign(m)) / 2
    return Int(trunc(term1 - term2 - term3)) + 1
end

"""
    zernike_nm_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_ansi_j_to_nm`](@ref) (reciprocal of this function)
"""
function zernike_nm_to_ansi_j(n::Integer, m::Integer)
    return (n * (n + 2) + m) ÷ 2
end

"""
    zernike_ansi_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_nm_to_ansi_j`](@ref) (reciprocal of this function)
"""
function zernike_ansi_j_to_nm(j::Integer)
    n = (-3 + √(9 + 8j))÷2
    m = 2j - n * (n + 2)
    return Int(n), Int(m)
end

"""
    zernike_noll_to_nm(j)

Map j Noll index to ANSI (n,m) indices.
"""
function zernike_noll_to_nm(j::Integer)
    n = ceil(Int, (-1 + √(1 + 8j))/2 - 1)
    if n == 0
        m = 0
    else
        nseries = (n+1) * (n+2) ÷ 2
        residual = j - nseries

        if isodd(j)
            sign = -1
        else
            sign = 1
        end

        if isodd(n)
            ms = [1,1]
        else
            ms = [0]
        end

        for i=0:(n÷2)-1
            push!(ms, ms[end]+2)
            push!(ms, ms[end])
        end

        m = ms[end+residual] * sign
    end
    return n, Int(m)
end

"""
    zernike_fringe_to_nm(j)

Map j Fringe index to ANSI (n,m) indices.
"""
function zernike_fringe_to_nm(j::Integer)
    m_n = 2 * (ceil(Int, √j) - 1)
    g_s = (m_n ÷ 2)^2 + 1
    n = m_n ÷ 2 + (j-g_s)÷2
    m = (m_n - n) * (1 - mod(j-g_s, 2) * 2)
    return n, m
end

"""
    zernike_zero_separation(n)

Minimum zero separation of Zernike polynomial of radial order n.  Useful for
computing sample count requirements.
"""
function zernike_zero_separation(n::Integer)
    return 1 / n^2
end

"""
    zernike(n, m, ρ, θ[; norm])

Zernike polynomial of radial order n and azimuthal order m, evaluated at the
point (ρ, θ).  No normalization is required of (ρ, θ), though the polynomials
are orthogonal only over the unit disk.

norm is a boolean flag indicating whether the result should be orthonormalized
(scaled to unit RMS) or not.

The zernike polynomials' radial basis is a special case of the Jacobi
polynomials under the transformation n_jacobi = (n-m)/2, α=0, β=|m|, x=2ρ^2-1.

See also: [`zernike_series`](@ref), [`zernike_sum`](@ref).
"""
function zernike(n::Integer, m::Integer, ρ, θ; norm::Bool=true)
    x = 2 .* ρ.^2 .- 1
    am = abs(m)
    n_j = (n - am) ÷ 2
    out = jacobi(n_j, 0, am, x)
    if m != 0
        if m < 0
            out .*= (ρ.^am .* sin.(m.*θ))
        else
            out .*= (ρ.^am .* cos.(m.*θ))
        end

    end
	if norm
		out .*= zernike_norm(n,m)
	end
    return out
end

"""
    zernike_series(nms, ρ, θ[; norm])

Compute a series of Zernike polynomials of orders (n,m).
Returns an array with shape (size(ρ)..., length(nms)).
That is, the _final_ dimension contains the modes and the first dimension(s)
are spatial.

See also: [`zernike`](@ref), [`zernike_sum`](@ref)
"""
function zernike_series(nms, ρ, θ; norm::Bool=true)
    # 1.  Calculate a lookup table of Jacobi polynomials for the unique elements combinations of $|m|$ and $n_j$.
    # 2.  Calculate a lookup table of $\rho^{|m|}$ for each unique $|m|$
    # 3.  Calculate a lookup table of $\sin{m\theta}$ and $\cos{m\theta}$ for each unique $m$.
    # 4.  For each input $n, m$ look up the appropriate elements and compute the Zernike polynomial.
	x = 2 .* ρ.^2 .- 1
    n_j_list = [((n - abs(m)) ÷ 2, abs(m)) for (n,m) in nms]
    abs_m_type = typeof(n_j_list[1][2])
    dtype = eltype(ρ)
    jacobi_sequences_mnj = Dict{abs_m_type,Array{typeof(n_j_list[1][1]), 1}}()
	for (nj, am) in n_j_list
		a = get!(jacobi_sequences_mnj, am) do
			Array{abs_m_type, 1}()
		end
		push!(a, nj)
	end
	for key in keys(jacobi_sequences_mnj)
		# sort may not be needed.  Inherently sorted
		# but not sure if unique guarantees not shuffling
		# input
		sort!(unique!(jacobi_sequences_mnj[key]))
	end
    # jacobi sequences_mnj now maps |m| => max n_j
    # this is step 1 at the top of the function
    jacobi_sequences = Dict{abs_m_type,Array{dtype,ndims(ρ)+1}}()
    for k in keys(jacobi_sequences_mnj)
        njs = jacobi_sequences_mnj[k]
        jacobi_sequences[k] = jacobi_series(njs, 0., dtype(k), x)
    end

    # jacobi_sequences now maps |m| => arrays with n_j along the last dimension
    # powers_of_rho is now ρ^|m| for each unique |m|
    plane_idx = ndims(ρ)+1
    sines = Dict{abs_m_type, typeof(ρ)}()
    cosines = Dict{abs_m_type, typeof(ρ)}()
    powers_of_rho = Dict{abs_m_type, typeof(ρ)}()
    out = Array{eltype(ρ), ndims(ρ)+1}(undef, size(ρ)..., length(nms))
    for i in eachindex(nms)
        n, m = nms[i]
        absm = abs(m)

        out_view = selectdim(out, plane_idx, i)
		n_j = (n-abs(m))÷2
		jac_plane = first(searchsorted(jacobi_sequences_mnj[absm], n_j))
        jacobi_piece = selectdim(jacobi_sequences[absm], plane_idx, jac_plane)
        if m == 0
            if norm
                norm_val = zernike_norm(n, m)
                out_view .= norm_val .* jacobi_piece
            else
                out_view .= jacobi_piece
            end
        else
            if m < 0
				# these evaluate the sin/cosine only as needed,
				# the do block is only called if get! would error
                azpiece = get!(sines, m) do
					sin.(m .* θ)
				end
            else
                azpiece = get!(cosines, m) do
					cos.(m .* θ)
				end
            end
            radialpiece = get!(powers_of_rho, absm) do
				ρ .^ absm
			end

            if norm
                norm_val = zernike_norm(n, m)
                out_view .= norm_val .* jacobi_piece .* radialpiece .* azpiece
            else
                out_view .= jacobi_piece .* radialpiece .* azpiece
            end
        end
    end
    return out
end

"""
    zernike_sum(nms, weights, ρ, θ[; norm])

Compute a sum of Zernike polynomials of orders (n,m) weighted by weights.

See also: [`zernike`](@ref), [`zernike_series`](@ref)
"""
function zernike_sum(nms, weights, ρ, θ, norm::Bool=true)
    data = zernike_series(nms, ρ, θ, norm=norm)
    return apply_weights(data, weights)
end


# bessel_minus_struve.jl

module bessel_minus_struve
export BesselI_minus_StruveL

using ArbNumerics

"""
Implementation of I[n, z] - L[-n, z] where I[n, z] is the modified Bessel function of the first kind, and L[-n, z] is the modified Struve function. n is an integer ≥ 0 and the order of the functions, and z∈ℜ>0 a positive real argument. Notice the order are opposite signs.

The expansion of the BesselI and StruveL functions are:

    I[n, z] = ∑_{k=0}^∞ (z/2)^{2k+n} / (Γ(k+1)⋅Γ(k+n+1))
    L[-n, z] = ∑_{k=0}^∞ (z/2)^{2k-n+1} / (Γ(k+3/2)⋅Γ(k-n+3/2))

with Γ(x) the Gamma function. The values of I[n, z] and L[-n, z] can be incredibly close, so ArbReal types are used as they allow for accurate arbitrary precision floats so that the difference, I[n, z] - L[-n, z], can be distinguished. ArbReal types are faster than BigFloat types, and the error of a float can be tracked using ball(::ArbReal)[2].

Unfortunately, as ArbReal(bits = 64) * ArbReal(bits = 128) = ArbReal(bits = 64), if the starting precision is too small to accurately produce I[n, z] - L[-n, z], then the entire summation has to be restarted at a higher precision otherwise rounding errors will be too large, which slows the algorithm sometimes. Therefore, if the precision is found to be too low, the precision is just doubled (rather than increased by some set amount) to reach a required precision quickly.

The general structure of this arbitrary precision summation algorithm is used elsewhere too since you just change the arguments and term appropriately.
"""

function BesselI_minus_StruveL(n, z, a; prec = 64) # z > 0 & n >= 0

    # Initialise precision of ArbReal to prec.
    p = prec
    setextrabits(0)
    setprecision(ArbReal, p) # ArbReal(bit = 64 + 8) has same precision as Float64 (which is why we add 8 bits).

    n = ArbReal("$n")
    z = ArbReal("$z")
    a = ArbReal("$a")

    k = ArbReal("0.0")
    result = ArbReal("0.0")
    term = ArbReal("1.0")
    err = eps(result)  # Machine accuracy of specified precision = prec.

    while abs(midpoint(term)) > err * abs(midpoint(result))
        term = ArbReal((abs(z) * a / 2)^(2 * k + n + 1) / (ArbNumerics.gamma(k + 1) * ArbNumerics.gamma(k + n + 2)) - (abs(z) * a / 2)^(2 * k - n) / (ArbNumerics.gamma(k + 3//2) * ArbNumerics.gamma(k - n + 1//2)))
        result += term
        # println("term: k = ", k, "\nterm value: ", ball(term), "\ncumulant result: ", ball(result), "\n")
        k += 1
        # Double precision if rounding error in result exceeds accuracy specified by prec.
        while radius(result) > err * abs(midpoint(result))
            p *= 2
            setprecision(ArbReal, p)
            # variables have to be re-parsed into the higher precision ArbReal type.
            # println("Not precise enough. Error = ", radius(result), " > ", ArbReal("$err"), ". Increasing precision to ", p, " bits.\n")
            n = ArbReal("$n")
            z = ArbReal("$z")
            a = ArbReal("$a")
            k = ArbReal("0")
            result = ArbReal("0.0")
            term = ArbReal("1.0")
        end
    end
    # println("z: ", ArbReal(z, bits = prec), ". Final result: ", ArbReal(result, bits = prec))
    ArbReal(result * √π * ArbNumerics.gamma(-n - 1/2) * (abs(z) / 2)^n * z / 2, bits = prec) # return to specified precision.
end
# @time bsi = BesselI_minus_StruveL(3, 55, 1)
# @show(bsi)

# β = ArbReal("2.0")
# α = ArbReal("7.0")
# v = ArbReal("5.8")
# w = ArbReal("1.6")
# R = ArbReal((v^2 - w^2) / (w^2 * v))
# a = ArbReal(sqrt(β^2 / 4 + R * β * coth(β * v / 2)))
# z = ArbReal("90.61")
# n = ArbReal("12.0")
# c = BesselI_minus_StruveL(n, z, a; prec = 24)

end # end of module


# # 110: 1D Reaction Diffusion equation with two species
# ([source code](SOURCE_URL))
#
# Solve the nonlinear coupled reaction diffusion problem
#
# ```math
# -\nabla (0.01+2u_2)\nabla u_1 + u_1u_2= 0.0001(0.01+x)
# ```
#
# ```math
# -\nabla (0.01+2u_1)\nabla u_2 - u_1u_2 = 0.0001(1.01-x)
# ```
#
#
# in $\Omega=(0,1)$ with boundary condition $u_1(0)=1$, $u_2(0)=0$ and $u_1(1)=1$, $u_2(1)=1$.
#

module Example110_ReactionDiffusion1D_TwoSpecies

using Printf
using VoronoiFVM
using ExtendableGrids
using GridVisualize

function main(;n=100,Plotter=nothing,verbose=false,unknown_storage=:sparse)
    h=1/n
    grid=VoronoiFVM.Grid(collect(0:h:1))

    eps::Vector{Float64}=[1.0,1.0]

    physics=VoronoiFVM.Physics(num_species=2,

                               reaction=function(f,u,node)
                               f[1]=u[1]*u[2]
                               f[2]=-u[1]*u[2]
                               end,

                               flux=function(f,u0,edge)
                               u=unknowns(edge,u0)

                               nspecies=2
                               f[1]=eps[1]*(u[1,1]-u[1,2])*(0.01+u[2,1]+u[2,2])
                               f[2]=eps[2]*(u[2,1]-u[2,2])*(0.01+u[1,1]+u[1,2])
                               end,

                               source=function(f,node)
                               f[1]=1.0e-4*(0.01+node[1])
                               f[2]=1.0e-4*(0.01+1.0-node[1])
                               end,

                               storage=function(f,u,node)
                               f[1]=u[1]
                               f[2]=u[2]
                               end
                               )

    sys=VoronoiFVM.System(grid,physics,unknown_storage=unknown_storage)

    enable_species!(sys,1,[1])
    enable_species!(sys,2,[1])

    boundary_dirichlet!(sys,1,1,1.0)
    boundary_dirichlet!(sys,1,2,0.0)

    boundary_dirichlet!(sys,2,1,1.0)
    boundary_dirichlet!(sys,2,2,0.0)

    inival=unknowns(sys)
    U=unknowns(sys)
    inival.=0

    control=VoronoiFVM.NewtonControl()
    control.verbose=verbose
    control.damp_initial=0.1
    u5=0
    p=GridVisualizer(Plotter=Plotter,layout=(2,1))
    for xeps in [1.0,0.5,0.25,0.1,0.05,0.025,0.01]
        eps=[xeps,xeps]
        solve!(U,inival,sys,control=control)
        inival.=U
        scalarplot!(p[1,1],grid,U[1,:],clear=true,title="U1, eps=$(xeps)")
        scalarplot!(p[2,1],grid,U[2,:],clear=true,title="U2, eps=$(xeps)",reveal=true)
        sleep(0.2)
        u5=U[5]
    end
    return u5
end

function test()
    testval=0.7117546972922056
    main(unknown_storage=:sparse) ≈ testval && main(unknown_storage=:dense) ≈ testval
end
end


# Calcul des distances a vol d'oiseau entre 2 coordonnees GPS
# X. Gandibleux - Octobre 2019

using Printf
using DelimitedFiles

# Donnees dans un fichier au format CSV séparé par des virgules avec :
# - ligne 1 : header
# - colonne 1 : nom du lieu
# - colonne 2 : position x du lieu
# - colonne 3 : position y du lieu
# - colonne 4 : latitude du lieu
# - colonne 5 : longitude du lieu
# donnees extraites de : https://data.education.gouv.fr
m = readdlm("GPSlyceesNantes.csv", ',')

# Fonction de conversion degre vers radian
function deg_rad(v)
    return pi/180*v
end

# Calcul du distancier
lieu = String[]
for i in 2:size(m,1)
    push!(lieu, m[i,1])
end

dist = zeros(size(m,1)-1, size(m,1)-1)
for i in 2:size(m,1)-1
    for j in i+1:size(m,1)

        LatitudeA  = m[i,4]
        LongitudeA = m[i,5]
        LatitudeB  = m[j,4]
        LongitudeB = m[j,5]

        # formule des sinus : http://villemin.gerard.free.fr/aGeograp/Distance.htm
        # https://www.coordonnees-gps.fr/conversion-coordonnees-gps

        d = acos(
             sin(deg_rad(LatitudeA)) * sin(deg_rad(LatitudeB))
             +
             cos(deg_rad(LatitudeA)) * cos(deg_rad(LatitudeB)) * cos(deg_rad(LongitudeB) - deg_rad(LongitudeA))
            ) * 6371

        @printf("%2d - %2d = %6.3f km | %s <-> %s \n",i-1, j-1, d, m[i,1], m[j,1])

        dist[i-1,j-1] = d
        dist[j-1,i-1] = d
    end
end

println("distancier symétrique crée : ", length(lieu), " lieux")


using Distributions

points = [1, 14, 2, 3, 4, 7, 9, 15, 22, 4]
function GMM(points::Array; k::Int64 = 2, tol::Float64 = .001)

    prior = 1 / k
    n = length(points)
    @assert n > k "Need more data points than latent classes!"

    post = repmat([prior], n, k)
    post = hcat(post, points)

    # initial random latent gaussians
    sd = sqrt(var(points))
    gauss = [Distributions.Normal(rand(points), sd) for i in 1:k]

    est_mu(w, points) = sum(w .* points) / sum(w)
    est_var(w, mu, points) = sum(w .* (points - mu).^2) / sum(w)

    prev = post[:, 1]
    converged = false

    while !converged

        # for each point, compute the posterior of belonging to each latent gaussian
        for i in 1:n
            p = post[i, end]
            normalizer = sum([pdf(norm, p) for norm in gauss]) * prior

            for j in 1:k
                post[i, j] = (pdf(gauss[j], p) * prior) / normalizer
            end
        end

        diff = sum(abs(prev - post[:, 1]))

        if diff < tol
            converged = true
        else
            prev = post[:, 1]
        end

        mus = [est_mu(post[:, i], points) for i in 1:k]
        vars = [est_var(post[:, i], mus[i], points) for i in 1:k]
        gauss = [Distributions.Normal(mus[i], sqrt(vars[i])) for i in 1:k]
    end

    post
end


# GaussQuadrature.jl
mutable struct DNq
    n::Array{Float64, 2}
    w::Array{Float64, 1}
end

"""
    initializenode(ndims::Int, quadpts::Int)
Initialize the nodes and weights of Gauss-Hermite quadrature (Multidimentional version).
Generate `ndims` dimensional quadrature nodes and weights.

# Arguments

# Optional arguments

- `dentype`
    - `:DN` Discrete normal
    - `:GH` Gauss-Hermite quadrature

# Example
```julia
ghn = initializenode(3, 21)
```
"""
function initializenode(ndims, quadpts; θrange = (-4, 4))
    nodes = range(θrange[1], stop = θrange[2], length = quadpts)
    weights = pdf.(Normal(), nodes) ./ sum(pdf.(Normal(), nodes))
    ghn = Dict(zip(1:quadpts, nodes))
    ghw = Dict(zip(1:quadpts, weights))
    # https://stackoverflow.com/questions/54256481/converting-array-of-cartesianindex-to-2d-matrix-in-julia
    as_ints(a::AbstractArray{CartesianIndex{L}}) where L = reshape(reinterpret(Int, a), (L, size(a)...))
    idx = convert(Matrix, as_ints(vec(CartesianIndices(Array{Float64}(undef, Tuple(fill(quadpts, ndims))))))')
    n = [ghn[idx[i, j]] for i in axes(idx, 1), j in axes(idx, 2)]
    w = prod([ghw[idx[i, j]] for i in axes(idx, 1), j in axes(idx, 2)], dims = 2)[:]
    return DNq(n, w)
end


# # Logistic regression
using DataFrames
using Plots
using RDatasets
using Convex
using SCS

#-

## we'll use iris data
## predict whether the iris species is versicolor using the sepal length and width and petal length and width
iris = dataset("datasets", "iris")
## outcome variable: +1 for versicolor, -1 otherwise
Y = [species == "versicolor" ? 1.0 : -1.0 for species in iris.Species]
## create data matrix with one column for each feature (first column corresponds to offset)
X = hcat(ones(size(iris, 1)), iris.SepalLength, iris.SepalWidth, iris.PetalLength, iris.PetalWidth);

#-

## solve the logistic regression problem
n, p = size(X)
beta = Variable(p)
problem = minimize(logisticloss(-Y.*(X*beta)))

solve!(problem, SCSSolver(verbose=false))

#-

## let's see how well the model fits
using Plots
logistic(x::Real) = inv(exp(-x) + one(x))
perm = sortperm(vec(X*beta.value))
plot(1:n, (Y[perm] .+ 1)/2, st=:scatter)
plot!(1:n, logistic.(X*beta.value)[perm])


function generatemonomials(x,m,;option=0)
    mons = [];
    if option == 0
        mons = monomials(x,[1:m;], mon -> (degree(mon) == m));
    elseif 1 <= option <= n
        mons = monomials(x,[1:m;], mon -> ((degree(mon) == m)&&(degree(mon,x[option])<=1)));
    else
        throw(DomainError(option,"option must be non-negative and <=$(n)"))
    end
    return mons
end

function differentiatepow(u,var::MP.AbstractVariable;power=1)
    res = u;
    if power>1
        res = differentiate(differentiatepow(u,var,power=power-1),var);
    elseif power==1
        res = differentiate(u,var);
    elseif power < 0
        throw(DomainError(power,"power must be non-negative"));
    end
    return res;
end

function dalpha(u,alpha,vars)
    if length(alpha) != length(vars)
        throw(DomainError(alpha,"alpha must be a valid multi-index for the variables x"))
    end
    res = u;
    for i in 1:1:length(alpha)
        res = differentiatepow(res,vars[i],power=alpha[i]);
    end
    return res;
end

function polyderivative(u,poly,vars)
    pterms = terms(poly);
    pcoefs = map(coefficient,pterms);
    pmons = map(monomial,pterms);
    fixedexp = z -> exponents(prod(vars.^0)*z);
    ppows = map(fixedexp,pmons);
    Dvec = map(z->dalpha(u,z,vars),ppows);
    if length(pcoefs)!= length(Dvec)
        throw(ErrorException("poly was not well defined"))
    end
    return sum(pcoefs[i]*Dvec[i] for i in 1:1:length(Dvec));
end

function laplacian(u,vars,;power=1)
    res = u;
    if power > 1
        res = polyderivative(laplacian(u,vars,power=power-1),sum(vars.*vars),vars);
    elseif power == 1
        res = polyderivative(u,sum(vars .* vars),vars);
    elseif power < 0
        throw(DomainError(power,"power must be non-negative"));
    end
    return res;
end

function normpartialder(hvec,var::MP.AbstractVariable,vars,;power=1)
    res = hvec;
    if power>1
        der = normpartialder(hvec,var,vars,power=power-1)
        res = (der[1]-2,der[1]*var*der[2]+sum(vars.*vars)*differentiate(der[2],var));
    elseif power==1
        res = (hvec[1]-2,hvec[1]*var*hvec[2]+sum(vars.*vars)*differentiate(hvec[2],var));
    elseif power < 0
        throw(DomainError(power,"power must be non-negative"));
    end
    return res;
    return
end

function normdalpha(hvec,alpha,vars)
    if length(alpha) != length(vars)
        throw(DomainError(alpha,"alpha must be a valid multi-index for the variables x"))
    end
    res = hvec;
    for i in 1:1:length(alpha)
        res = normpartialder(res,vars[i],vars,power=alpha[i]);
    end
    return res;
end

function normpolyderivative(hvec,poly,vars)
    pterms = terms(poly);
    pcoefs = map(coefficient,pterms);
    pmons = map(monomial,pterms);
    fixedexp = z -> exponents(prod(vars.^0)*z);
    ppows = map(fixedexp,pmons);
    Dvec = map(z->normdalpha(hvec,z,x),ppows);
    if length(pcoefs)!= length(Dvec)
        throw(ErrorException("poly was not well defined"))
    end
    return Dvec
end

function generatebasissphere(m, vars)
    n = length(vars)
    genpoly = (2-n,1);
    monos = generatemonomials(vars,m,option=1);
    fixedexp = z -> exponents(prod(vars.^0)*z);
    mpows = map(fixedexp,monos)
    return map(m->normdalpha(genpoly,m,vars)[2],mpows)
end

function dcoefficient(n,m,i,j)
    d = 0;
    if i == j
        d = 1.0/(2.0^j*factorial(j)*prod([(2*m + n - 2*j - 2*l) for l in 1:j]));
    elseif i < j
        d = -dcoefficient(n,m,i,j-1)/(2*(j-i)*(2*m+n-2-2*i-2*j));
    end
    return d;
end

function harmonicdecomposition(u,vars)
    norm = sum(vars .* vars);
    n =length(vars);
    m = maxdegree(u);
    if mindegree(u) != m
        throw(DomainError(u,"polynomial must be homogeneous"));
    end
    return [ (2*i,sum([dcoefficient(n,m,i,j)*norm^(j-i)*laplacian(u,vars,power=j) for j in i:floor(Int,m/2)])) for i in 0:floor(Int,m/2) ];
end


function sieve(n :: Int)
    isprime = trues(n)
    isprime[1] = false
    for p in 2:n
        if isprime[p]
            j = p * p
            if j > n
                return findall(isprime)
            else
                for k in j:p:n
                  isprime[k] = false
                end
            end
        end
    end
end


# # Holomorphic integration on the unit disk
# In this example, we explore integration of a harmonic function:
# ```math
#   f(x,y) = \frac{x^2-y^2+1}{(x^2-y^2+1)^2+(2xy+1)^2},
# ```
# over the unit disk. In this case, we know from complex analysis that the
# integral of a holomorphic function is equal to $\pi \times f(0,0)$.
# We analyze the function on an $N\times M$ tensor product grid defined by:
# ```math
# \begin{aligned}
# r_n & = \cos\left[(n+\tfrac{1}{2})\pi/2N\right],\quad{\rm for}\quad 0\le n < N,\quad{\rm and}\\
# \theta_m & = 2\pi m/M,\quad{\rm for}\quad 0\le m < M;
# \end{aligned}
# ```
# we convert the function samples to Chebyshev×Fourier coefficients using
# `plan_disk_analysis`; and finally, we transform the Chebyshev×Fourier
# coefficients to Zernike polynomial coefficients using `plan_disk2cxf`.
#
# For the storage pattern of the arrays, please consult the
# [documentation](https://MikaelSlevinsky.github.io/FastTransforms).

using FastTransforms, LinearAlgebra, Plots
const GENFIGS = joinpath(pkgdir(FastTransforms), "docs/src/generated")
!isdir(GENFIGS) && mkdir(GENFIGS)
plotlyjs()

# Our function $f$ on the disk:
f = (x,y) -> (x^2-y^2+1)/((x^2-y^2+1)^2+(2x*y+1)^2)

# The Zernike polynomial degree:
N = 15
M = 4N-3

# The radial grid:
r = [sinpi((N-n-0.5)/(2N)) for n in 0:N-1]

# The angular grid (mod $\pi$):
θ = (0:M-1)*2/M

# On the mapped tensor product grid, our function samples are:
F = [f(r*cospi(θ), r*sinpi(θ)) for r in r, θ in θ]

# We superpose a surface plot of $f$ on top of the grid:
X = [r*cospi(θ) for r in r, θ in θ]
Y = [r*sinpi(θ) for r in r, θ in θ]
scatter3d(vec(X), vec(Y), vec(0F); markersize=0.75, markercolor=:red)
surface!(X, Y, F; legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "zernike.html"))
###```@raw html
###<object type="text/html" data="../zernike.html" style="width:100%;height:400px;"></object>
###```

# We precompute a (generalized) Zernike--Chebyshev×Fourier plan:
α, β = 0, 0
P = plan_disk2cxf(F, α, β)

# And an FFTW Chebyshev×Fourier analysis plan on the disk:
PA = plan_disk_analysis(F)

# Its Zernike coefficients are:
U = P\(PA*F)

# The Zernike coefficients are useful for integration. The integral of $f(x,y)$
# over the disk should be $\pi/2$ by harmonicity. The coefficient of $Z_{0,0}$
# multiplied by `√π` is:
U[1, 1]*sqrt(π)

# Using an orthonormal basis, the integral of $[f(x,y)]^2$ over the disk is
# approximately the square of the 2-norm of the coefficients:
norm(U)^2, π/(2*sqrt(2))*log1p(sqrt(2))

# But there's more! Next, we repeat the experiment using the Dunkl-Xu
# orthonormal polynomials supported on the rectangularized disk.
N = 2N
M = N

# We analyze the function on an $N\times M$ mapped tensor product $xy$-grid defined by:
# ```math
# \begin{aligned}
# x_n & = \cos\left(\frac{2n+1}{2N}\pi\right) = \sin\left(\frac{N-2n-1}{2N}\pi\right),\quad {\rm for} \quad 0 \le n < N,\quad{\rm and}\\
# z_m & = \cos\left(\frac{2m+1}{2M}\pi\right) = \sin\left(\frac{M-2m-1}{2M}\pi\right),\quad {\rm for} \quad 0 \le m < M,\\
# y_{n,m} & = \sqrt{1-x_n^2}z_m.
# \end{aligned}
# ```
# Slightly more accuracy can be expected by using an auxiliary array:
# ```math
#   w_n = \sin\left(\frac{2n+1}{2N}\pi\right),\quad {\rm for} \quad 0 \le n < N,
# ```
# so that $y_{n,m} = w_nz_m$.
#
# The x grid
w = [sinpi((n+0.5)/N) for n in 0:N-1]
x = [sinpi((N-2n-1)/(2N)) for n in 0:N-1]

# The z grid
z = [sinpi((M-2m-1)/(2M)) for m in 0:M-1]

# On the mapped tensor product grid, our function samples are:
F = [f(x[n], w[n]*z) for n in 1:N, z in z]

# We superpose a surface plot of $f$ on top of the grid:
X = [x for x in x, z in z]
Y = [w*z for w in w, z in z]
scatter3d(vec(X), vec(Y), vec(0F); markersize=0.75, markercolor=:green)
surface!(X, Y, F; legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "dunklxu.html"))
###```@raw html
###<object type="text/html" data="../dunklxu.html" style="width:100%;height:400px;"></object>
###```

# We precompute a Dunkl-Xu--Chebyshev plan:
P = plan_rectdisk2cheb(F, β)

# And an FFTW Chebyshev² analysis plan on the rectangularized disk:
PA = plan_rectdisk_analysis(F)

# Its Dunkl-Xu coefficients are:
U = P\(PA*F)

# The Dunkl-Xu coefficients are useful for integration. The integral of $f(x,y)$
# over the disk should be $\pi/2$ by harmonicity. The coefficient of $P_{0,0}$
# multiplied by `√π` is:
U[1, 1]*sqrt(π)

# Using an orthonormal basis, the integral of $[f(x,y)]^2$ over the disk is
# approximately the square of the 2-norm of the coefficients:
norm(U)^2, π/(2*sqrt(2))*log1p(sqrt(2))


# Parts of this code were taken / derived from Graphs.jl. See LICENSE for
# licensing details.

type FloydWarshallState{T}<:AbstractPathState
    dists::Matrix{T}
    parents::Matrix{Int}
end

doc"""Uses the [Floyd-Warshall algorithm](http://en.wikipedia.org/wiki/Floyd–Warshall_algorithm)
to compute shortest paths between all pairs of vertices in graph `g`. Returns a
`FloydWarshallState` with relevant traversal information, each is a
vertex-indexed vector of vectors containing the metric for each vertex in the
graph.

Note that this algorithm may return a large amount of data (it will allocate
on the order of $\mathcal{O}(nv^2)$).
"""
function floyd_warshall_shortest_paths{T}(
    g::SimpleGraph,
    distmx::AbstractArray{T, 2} = DefaultDistance()
)

    n_v = nv(g)
    dists = fill(typemax(T), (n_v,n_v))
    parents = zeros(Int, (n_v,n_v))

    # fws = FloydWarshallState(Matrix{T}(), Matrix{Int}())
    for v in 1:n_v
        dists[v,v] = zero(T)
    end
    undirected = !is_directed(g)
    for e in edges(g)
        u = src(e)
        v = dst(e)

        d = distmx[u,v]

        dists[u,v] = min(d, dists[u,v])
        parents[u,v] = u
        if undirected
            dists[v,u] = min(d, dists[v,u])
            parents[v,u] = v
        end
    end
    for w in vertices(g), u in vertices(g), v in vertices(g)
        if dists[u,w] == typemax(T) || dists[w,v] == typemax(T)
            ans = typemax(T)
        else
            ans = dists[u,w] + dists[w,v]
        end
        if dists[u,v] > ans
            dists[u,v] = dists[u,w] + dists[w,v]
            parents[u,v] = parents[w,v]
        end
    end
    fws = FloydWarshallState(dists, parents)
    # for r in 1:size(parents,1)    # row by row
    #     push!(fws.parents, vec(parents[r,:]))
    # end
    # for r in 1:size(dists,1)
    #     push!(fws.dists, vec(dists[r,:]))
    # end

    return fws
end

function enumerate_paths(s::FloydWarshallState, v::Integer)
    pathinfo = s.parents[v,:]
    paths = Vector{Int}[]
    for i in 1:length(pathinfo)
        if i == v
            push!(paths, Vector{Int}())
        else
            path = Vector{Int}()
            currpathindex = i
            while currpathindex != 0
                push!(path,currpathindex)
                currpathindex = pathinfo[currpathindex]
            end
            push!(paths, reverse(path))
        end
    end
    return paths
end

enumerate_paths(s::FloydWarshallState) = [enumerate_paths(s, v) for v in 1:size(s.parents,1)]
enumerate_paths(st::FloydWarshallState, s::Integer, d::Integer) = enumerate_paths(st, s)[d]


export modnewton

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

Newton method with Hessian shift for solving min_x f(x)

"""
function modnewton(f::Function,df::Function,H::Function,x::Vector;maxIter=20,atol=1e-8,out::Int=0,
	storeInterm::Bool=false,beta=1e-1,
		lineSearch::Function=(f,df,fk,dfk,xk,pk)->armijo(f,fk,dfk,xk,pk,maxIter=10))

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

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

        # get search direction
        d2f = H(x)
        tau = (minimum(diag(d2f)) > 0) ? 0 :  -minimum(diag(d2f))+beta
        nfac = 0;
        for k=1:20
            try
                nfac +=1
                LL    = cholfact(d2f + tau*speye(length(x)))
                break;
            catch err
                if isa(err,Base.LinAlg.PosDefException) && k<20
                    tau = max(2*tau,beta)
                elseif isa(err,Base.LinAlg.PosDefException)
                    flag = -2
                    his[i,4] = nfac
                    his = his[1:i,:]
                    break; break;
                else
                    throw(err)
                end
            end
        end
        his[i,4] = nfac
        pk    = - (LL\dfk)

        # line search
		ak,his[i,3] = lineSearch(f,df,fk,dfk,x,pk)
        if his[i,3]==-1
            flag = -3
            his  = his[1:i,:]
            break
        end
        if out>0
            @printf "iter=%04d\t|f|=%1.2e\t|df|=%1.2e\tLS=%d\n" i his[i,1] his[i,2] his[i,3]
        end
        # update x and H
        x  += ak*pk

        i+=1
    end
    i = min(maxIter,i)

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


#!/usr/bin/env julia

#=
# data should be sorted and works by divide and conquer
# run-time complexity O(log n)
=#

function binary_search(search_for, data_array, low=1, high=-1)
  if high == -1
    high = length(data_array)
  end
  if high < low
    return -1
  end
  mid = Int(floor(low + (high - low) / 2))
  if data_array[mid] == search_for
    return mid
  end
  if data_array[mid] > search_for
    return binary_search(search_for, data_array, low, mid-1)
  else
    return binary_search(search_for, data_array, mid+1, high)
  end
end

     ## 1   2   3   4   5   6   7   8   9   10
data = [10, 11, 12, 22, 24, 27, 34, 47, 56, 67]

@time println("at index: ", binary_search(47, data))
@time println("at index: ", binary_search(12, data))
@time println("at index: ", binary_search(7, data))
@time println("at index: ", binary_search(67, data))
@time println("at index: ", binary_search(34, data))


module BifurcationLib

# a demonstration of bifurcation using the logistic map

# equation: x_n+1 = r x_n (1-x_n)

# how does equilibrium change as the parameter r changes?

using Plots
# using ArgParse

### utility functions:

# rounds x to an aribtrary number of decimal places
function decimal_round(x::AbstractFloat, num_decimal_places::Unsigned)::AbstractFloat
  rounding_factor = 10^num_decimal_places
  round(rounding_factor * x)/rounding_factor
end

### Logistic map stuff

# gives logistic map of a value x and parameter (growth rate) r
logistic_map(x, r) = r*x*(1-x)

# gives many results of logistic map, starting with nth iteration and returning
# the results of m iterations after the nth
function logistic_map_series(x, r, n, m)
  nth_x = x
  for i = 1:n
    nth_x = logistic_map(nth_x, r)
  end

  mth_x = nth_x
  series = []
  for i = 1:m
    mth_x = logistic_map(mth_x, r)
    append!(series, [mth_x])
  end
  series
end

### Plotting things

# r0: first r to use

# r_inc: amount to increment r by each time

# r_final: final r (if r_final-r0 is not divisibible by r_inc then the actual
# final r will be slightly less than r_final)

function generate_points(x0, r0, n, m, r_inc, r_final)
  x = []
  y = []

  for r = r0:r_inc:r_final
    this_x = fill(r, (m))
    this_y = logistic_map_series(x0, r, n, m)
    append!(x, this_x)
    append!(y, this_y)
  end

  x, y
end

function plot_points(x, y, n, m)
  # set backend
  gr()

  # create plot
  plot(x, y,
    seriestype = :scatter,
    title = "Values of recursive iterations of logistic map for various r",
    xlabel = "r",
    ylabel = string("next ", m, " values after first ", n, " iterations"),
    dpi = 1000,
    markersize = 1,
    markeralpha = 0.5,
    # markeralpha = 0.1,
    # markerstrokestyle = :dot,
    # markerstrokecolor = :blue,
    markerstrokewidth = 0,
    markerfillcolor = :blue,
  )
end

function generate_and_plot(x0, r0, n, m, r_inc, r_final)
  println("generating points for plot with the following parameters:")
  println("x0      = ", x0)
  println("r0      = ", r0)
  println("n       = ", n)
  println("m       = ", m)
  println("r_inc   = ", r_inc)
  println("r_final = ", r_final)
  println()

  println("total number of poins to generate: ",
    floor( (r_final-r0) * r_inc ) * (n + m)
  )
  println()

  println("generating...")
  x, y = generate_points(x0, r0, n, m, r_inc, r_final)

  println("points generated. Plotting...")
  p = plot_points(x, y, n, m)
  # display(p)
  filename = string(
      "x0=", x0, ",",
      "r0=", r0, ",",
      "n=", n, ",",
      "m=", m, ",",
      "r_inc=", r_inc, ",",
      "r_final=", r_final, ".png")
  println("saving figure to ", filename)
  savefig(p, filename)
  println("done.")
  # p
end

function example_plot()
  x0 = 0.5
  r0 = 0.0
  n = 100
  m = 100
  r_inc = 0.01
  r_final = 4.0 # seems like anything higher than this diverges quickly to -inf

  generate_and_plot(x0, r0, n, m, r_inc, r_final)
end

end # module


using Distributions

# ρ(x,y) in integrand
ρ(x,y) = (0 < x < 1) && (0 < y < 2) ? 1/2 : 0

#=
Implementation of random walk sampler Metropolis-Hastings.
Inputs:
- n: The number of states to return
- b: The standard deviation of the normal distribution used in the proposal
     distribution
Outputs:
- Returns an array of length n denoting states S_i in Metropolis-Hastings alg.
=#
function metropolis_hastings(n, b)

    # Samples
    samples = zeros(n,2)
    samples[1,:] = rand(2)

    # Conditional Q distribution
    d = Normal(0,b)

    acceptances = 0

    for i=2:n
        # Propose Y
        ξx = rand(d)
        ξy = rand(d)
        X = samples[i-1,1] + ξx
        Y = samples[i-1,2] + ξy

        # Acceptance probability
        α = min(1, ρ(X,Y)/ρ(samples[i-1,1],samples[i-1,2]))

        # Next state
        if rand() < α
            samples[i,1] = X
            samples[i,2] = Y
        else
            samples[i,:] = samples[i-1,:]
        end
    end

    return samples
end

# Obtain 10000 samples with b = 1
S = metropolis_hastings(10000,1)

# f(x) in integrand
f(x,y) = exp(sin(x*y))

# Estimate integral
M = 1000
int_estimate = 2*mean([f(x,y) for (x,y) in zip(S[M:end,1], S[M:end,2])])


"""
    calc_nchk(n::Integer,k::Integer)

Calculates binomial coefficient: n choose k
"""
function calc_nchk(n::Integer,k::Integer)
    accum::Int = 1
    for i in 1:k
        accum = accum * (n-k+i) ÷ i
    end
    return accum
end

"""
    get_nchk(n::Integer,k::Integer)

Looks up binomial coefficient from a precomputed table: n choose k
"""
@inline function get_nchk(n,k)
    return binom_coeff[n+1,k+1]
end


"""
pde_lin_test(n)

Test the linear operators.
"""
function pde_lin_test(n)
    h = 1.0 / (n + 1)
    x = collect(h:h:1.0-h)
    fdata = fishinit(n)
    z = rand(n, n)
    L2d = Lap2d(n)
    DX = Dx2d(n)
    DY = Dy2d(n)
    lapok = lap_test(n, x, z, L2d, fdata)
    discok = disc_test(n, x, L2d, DX, DY, fdata)
    lapok && discok
end

"""
disc_test()
Have I broken the discretizations?
"""
function disc_test(n, x, L2d, DX, DY, fdata)
    n2 = n * n
    ue = solexact(x)
    ux = dxexact(x)
    uy = dyexact(x)
    D2u = l2dexact(x)
    ue1 = reshape(ue, (n2,))
    uex1 = reshape(ux, (n2,))
    uey1 = reshape(uy, (n2,))
    ued21 = reshape(D2u, (n2,))
    # test DX
    dx21 = DX * ue1
    dxerr = norm(dx21 - uex1, Inf)
    # test DY
    dy21 = DY * ue1
    dyerr = norm(dy21 - uey1, Inf)
    # test Laplacian
    du21 = L2d * ue1
    d2err = norm(du21 - ued21, Inf)
    pass = (d2err < 0.75) && (dxerr < 0.1) && (dyerr < 1.e-12)
    pass || println("Discretization test fails")
    return pass
end

"""
lap_test()

Does the FFT invert the discrete Laplacian?
Does the discrete Laplacian pass a simple eigenvalue test.
"""
function lap_test(n, x, z, L2d, fdata)
    randok = rand_test(z, n, L2d, fdata)
    eigok = eig_test(n, x, fdata)
    pass = randok && eigok
    return pass
end

function rand_test(z, n, L2d, fdata)
    n2 = n * n
    z1 = reshape(z, (n2,))
    y1 = L2d * z1
    y = reshape(y1, (n, n))
    mz = fish2d(y, fdata)
    q = reshape(mz, (n2,))
    pass = (norm(q - z1, Inf) < 1.e-12)
    pass || println("rand_test fails, norm =", norm(q - z1))
    return pass
end

function eig_test(n, x, fdata)
    lambda = pi * pi * 5
    efunx = sin.(pi * x)
    efuny = sin.(2 * pi * x)
    efunu = efunx * efuny'
    vfun = fish2d(efunu, fdata)
    pass = (norm(lambda * vfun - efunu, Inf) < 1.e-2)
    pass || println("eig test fails")
    return pass
end


"""
Cross-entropy loss as defined in PyTorch:

http://pytorch.org/docs/master/nn.html#torch.nn.CrossEntropyLoss

This method works with one-hot-encoded inputs
"""
function cross_entropy_loss(ŷ::AbstractMatrix, y::AbstractMatrix)
    sublosses = -sum(y .* ŷ, 1) .+ log.(sum(exp.(ŷ), 1))
    return mean(sublosses)
end

@require CuArrays begin

    # have to split sublosses into 2 subexpressions since otherwise CUDAnative failes with
    # 'ERROR: Broadcast output type Any is not concrete'
    function cross_entropy_loss(ŷ::CuArray{T,2}, y::CuArray{T,2}) where T
        w1 = -sum(y .* ŷ, 1)
        w2 = sum(exp.(ŷ), 1)
        sublosses = w1 .+ CUDAnative.log.(w2)
        return mean(sublosses)
    end

end


# Huber loss function

"""
  HuberLoss(rho::Real=1.0, mu::Real=1.0)

Returns the function `g(x) = (mu/2)||x||² if ||x|| ⩽ rho, and rho*mu*(||x||-rho/2) otherwise`.
"""

immutable HuberLoss{R <: Real} <: ProximableFunction
  rho::R
  mu::R
  function HuberLoss(rho::R, mu::R)
    if rho <= 0.0 || mu <= 0.0
      error("parameters rho and mu must be positive")
    else
      new(rho, mu)
    end
  end
end

HuberLoss{R <: Real}(rho::R=1.0, mu::R=1.0) = HuberLoss{R}(rho, mu)

function (f::HuberLoss){T <: Union{Real, Complex}}(x::AbstractArray{T})
  normx = vecnorm(x)
  if normx <= f.rho
    return (f.mu/2)*normx^2
  else
    return f.rho*f.mu*(normx-f.rho/2)
  end
end

function prox!{T <: Union{Real, Complex}}(f::HuberLoss, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  normx = vecnorm(x)
  mugam = f.mu*gamma
  scal = (1-min(mugam/(1+mugam), mugam*f.rho/(normx)))
  for k in eachindex(y)
    y[k] = scal*x[k]
  end
  normy = scal*normx
  if normy <= f.rho
    return (f.mu/2)*normy^2
  else
    return f.rho*f.mu*(normy-f.rho/2)
  end
end

fun_name(f::HuberLoss) = "Huber loss"
fun_dom(f::HuberLoss) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::HuberLoss) = "x ↦ (μ/2)||x||² if ||x||⩽ρ, μρ(||x||-ρ/2) otherwise"
fun_params(f::HuberLoss) = string("ρ = $(f.rho), μ = $(f.mu)")

function prox_naive{T <: Union{Real, Complex}}(f::HuberLoss, x::AbstractArray{T}, gamma::Real=1.0)
  y = (1-min(f.mu*gamma/(1+f.mu*gamma), f.mu*gamma*f.rho/(vecnorm(x))))*x
  if vecnorm(y) <= f.rho
    return y, (f.mu/2)*vecnorm(y)^2
  else
    return y, f.rho*f.mu*(vecnorm(y)-f.rho/2)
  end
end