#---------------------------------------------------------------------#
#This code computes the Legendre-Gauss points & weights
#which are the roots of the Legendre Polynomials.
#Written by F.X. Giraldo on 4/2008
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function legendre_gauss(P::Integer,DFloat)

    xgl=zeros(DFloat,P)
    wgl=zeros(DFloat,P)

    p=P-1; #Order of Polynomials [P is the number of points: p+1]
    ph=floor(typeof(P), (p+1)/2 )

    for i=1:ph
        x=cos( (2*i-1)*pi/(2*p+1) )
        for k=1:20
            (L0,L0_1,L0_2)=legendre_poly(p+1,x); #Compute the N+1 Legendre Polys

            #Get new Newton Iteration
            dx=-L0/L0_1
            x=x+dx
            if (abs(dx) < 1.0e-20)
                break
            end
        end
        xgl[p+2-i]=x
        wgl[p+2-i]=2/( (1-x^2)*L0_1^2 )
    end

    #Check for Zero Root
    if (p+1 != 2*ph)
        x=0
        (L0,L0_1,L0_2)=legendre_poly(p+1,x)
        xgl[ph+1]=x
        wgl[ph+1]=2/( (1-x^2)*L0_1^2 )
    end

    #Find remainder of roots via symmetry
    for i=1:ph
        xgl[i]=-xgl[p+2-i]
        wgl[i]=+wgl[p+2-i]
    end
    return xgl,wgl
end #function


using LinearAlgebra
using SparseArrays
using Plots

m = 10
x = LinRange(0,1,m+2) # easier to define eigs on [0,1]
xint = x[2:end-1]
h = x[2]-x[1]
A = (1/h^2) * spdiagm(0=>2*ones(m),-1=>-ones(m-1),1=>-ones(m-1))

λ,V = eigen(Matrix(A))
# scatter(xint,V[:,1])

function normalize_cols!(A)
    for i = 1:size(A,2)
        A[:,i] /= maximum(A[:,i])
    end
end
normalize_cols!(V)

λex = @. 2/h^2 * (1-cos((1:m)*pi*h))
Vex = hcat((sin.(k*pi*xint) for k = 1:m)...)
normalize_cols!(Vex)

# plot(xint,V[:,1] .- Vex[:,1])
# scatter(λ - λex)


@inline function checktail(tail::Symbol)
    if tail ∉ (:left, :right, :both)
        throw(ArgumentError("Unsupported tail specifier ($tail)"))
    end
end

"""
    logpvalue(nisect::Integer, na::Integer, nb::Integer, ntotal::Integer,
              [tail::Symbol = :right])

Log P-value for the two sets intersection.

`A` has `na` elemnts, B has `nb` elements,
they have `nisect` elements in common, and there are
`ntotal` elements in the "universe".

`tail` controls the null hypothesis:
* `:right` (default): by chance `A` and `B` would have ≥ elements in common
* `:left`: by chance `A` and `B` would have ≤ elements in common
* `:both`: by chance `A` and `B` would have either ≤ or ≥ elements in common, whichever is less probable
"""
function logpvalue(nisect::Integer, na::Integer, nb::Integer, ntotal::Integer,
                   tail::Symbol = :right)
    ((na >= 0) && (nb >= 0) && (ntotal >= 0)) ||
        throw(ArgumentError("Sets with negative number of elements"))
    ((na <= ntotal) && (nb <= ntotal)) ||
        throw(ArgumentError("Sets bigger than total number of elements"))
    # corner cases
    isect_max = min(na, nb)
    isect_min = max(0, na + nb - ntotal)
    if tail == :right
        (nisect > isect_max) && return -Inf
        (nisect <= isect_min) && return 0.0
    elseif tail == :left
        (nisect >= isect_max) && return 0.0
        (nisect < isect_min) && return -Inf
    elseif tail != :both
        # FIXME corner cases when tail == :both
        throw(ArgumentError("Unsupported tail specifier ($tail)"))
    end
    # normal cases
    distr = Distributions.Hypergeometric(na, ntotal - na, nb)
    if tail == :right
        return nisect < isect_max ? logccdf(distr, nisect-1) : logpdf(distr, nisect)
    elseif tail == :left
        return nisect > isect_min ? logcdf(distr, nisect) : logpdf(distr, nisect)
    elseif tail == :both
        return log(2.0) + min(logcdf(distr, nisect), logccdf(distr, nisect-1), log(0.5))
    else
        return NaN # noop
    end
end

# size of the sets intersection
# the sets are represented by the sorted sets of their tiles
function _isect_size(set1_tiles::AbstractVector{Int},
                     set2_tiles::AbstractVector{Int},
                     tile_sizes::Vector{Int})
    (length(set1_tiles)==0 || length(set2_tiles)==0) && return 0
    res = 0
    i = 1
    tile1 = set1_tiles[i] # current tile of the 1st set
    j = 1
    tile2 = set2_tiles[j] # current tile of the 2nd set
    @inbounds while true
        if tile1 < tile2
            i += 1
            (i <= length(set1_tiles)) || break
            tile1 = set1_tiles[i]
        elseif tile1 > tile2
            j += 1
            (j <= length(set2_tiles)) || break
            tile2 = set2_tiles[j]
        else # intersecting tile, adjust the overlap
            res += tile_sizes[tile1]
            i += 1
            j += 1
            (i <= length(set1_tiles) && j <= length(set2_tiles)) || break
            tile1 = set1_tiles[i]
            tile2 = set2_tiles[j]
        end
    end
    return res
end

# returns set×set score matrix
# a×b score is the Fisher's test P-value of their intersection
# a×a self-intersections are processed in the same way
# (might be used later in cover problem construction to penalize the use of
# multiple variables referencing the same set)
function _setXset_scores(tileXset::SparseMaskMatrix, total_size::Int,
                         set_sizes::AbstractVector{Int},
                         tile_sizes::AbstractVector{Int},
                         used_sets = Base.OneTo(size(tileXset, 2)))
    nsets = length(used_sets)
    res = zeros(Float64, nsets, nsets)
    # fill upper triangle
    @inbounds for i in 1:nsets
        set1_ix = used_sets[i]
        set1_size = set_sizes[set1_ix]
        set1_size == 0 && continue # skip empty sets
        set1_tiles = view(tileXset, :, set1_ix)
        for j in i:nsets
            set2_ix = used_sets[j]
            set2_size = set_sizes[set2_ix]
            set2_size == 0 && continue # skip empty sets
            set2_tiles = view(tileXset, :, set2_ix)
            # one-sided Fisher's P-value, right tail
            isect_size = _isect_size(set1_tiles, set2_tiles, tile_sizes)
            res[i, j] = min(logpvalue(isect_size, set1_size, set1_size + set2_size - isect_size, total_size),
                            logpvalue(isect_size, set2_size, set1_size + set2_size - isect_size, total_size))
        end
    end
    # symmetrize: copy upper triangle to lower
    @inbounds for i in 1:nsets
        for j in (i+1):nsets
            res[j, i] = res[i, j]
        end
    end
    return res
end

"""
    set_relevance(nset_observed::Integer, nset::Integer,
                  nobserved::Integer, ntotal::Integer) -> Float64

Calculates the relevance weight of the set that contains `nset` elements,
`nset_observed` of which were present (not necessarily enriched) in the data
that identified `nobserved` elements out of all known (`ntotal`).
It is used by [`SetMosaic`](@ref) to penalize the sets,
which could not be observed in the data (e.g. biological processes or pathways
that involve proteins not expressed by the cells used in the experiments).

While for [`MaskedSetMosaic`](@ref) it's recommended to use the IDs of data entities
(e.g. protein group IDs for proteomic data) to correctly count the set sizes
and estimate enrichment; `set_relevance()` should use the counts derived from
the original IDs of the annotation database (e.g. UniProt accession codes).
Otherwise it's not possible to correctly estimate the number of elements that
belong to the given annotated set, but were not observed in the data.

The returned value is the probability that no more than `nset_observed`
elements were observed at random.
"""
set_relevance(nset_observed::Integer, nset::Integer,
              nobserved::Integer, ntotal::Integer) =
    cdf(Hypergeometric(nset, ntotal - nset, nobserved), nset_observed)


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _recpol(x::T, y::T, degrees::Bool) where {T<:AbstractFloat}
    if degrees
        return hypot(x, y), rad2deg(atan(y, x))
    else
        return hypot(x, y), atan(y, x)
    end
end

"""
    recpol(x, y[, degrees=true]) -> radius, angle

### Purpose ###

Convert 2D rectangular coordinates to polar coordinates.

### Explanation ###

This is the partial inverse function of `polrec`.

### Arguments ###

* `x`: the abscissa coordinate of the point.  It may be a scalar or an array.
* `y`: the ordinate coordinate of the point.  It may be a scalar or an array
  of the same lenth as `x`.
* `degrees` (optional boolean keyword): if `true`, the output `angle` is given
 in degrees, otherwise in radians.  It defaults to `false`.

Mandatory arguments may also be passed as the 2-tuple `(x, y)`, so that it is
possible to execute `polrec(recpol(x, y))`.

### Output ###

A 2-tuple `(radius, angle)` with the polar coordinates of the input.  The
coordinate `angle` coordinate lies in the range \$[-\\pi, \\pi]\$ if
`degrees=false`, or \$[-180, 180]\$ when `degrees=true`.

If `x` and `y` are arrays, `radius` and `angle` are arrays of the same length as
`radius` and `angle`.

### Example ###

Calculate polar coordinates \$(r, \\varphi)\$ of point with rectangular
coordinates \$(x, y) = (2.24, -1.87)\$.

```jldoctest
julia> using AstroLib

julia> r, phi = recpol(2.24, -1.87)
(2.917961617293826, -0.6956158538564537)
```

Angle \$\\varphi\$ is given in radians.

"""
recpol(x::Real, y::Real; degrees::Bool=false) =
    _recpol(promote(float(x), float(y))..., degrees)

recpol(xy::Tuple{Real, Real}; degrees::Bool=false) =
    recpol(xy..., degrees=degrees)

function recpol(x::AbstractArray{X}, y::AbstractArray{Y};
                degrees::Bool=false) where {X<:Real, Y<:Real}
    @assert length(x) == length(y)
    typex = float(X)
    r = similar(x, typex)
    a = similar(x, typex)
    for i in eachindex(x)
        r[i], a[i] = recpol(x[i], y[i], degrees=degrees)
    end
    return r, a
end


# # Huber regression

# This example can be found here: <https://web.stanford.edu/~boyd/papers/pdf/cvx_applications.pdf>.
# Here we set `big_example = false` to only generate a small example which takes less time to run.
big_example = false
if big_example
    n = 300
    number_tests = 50
else
    n = 50
    number_tests = 10
end

# Generate data for Huber regression.
using Random
Random.seed!(1);
number_samples = round(Int, 1.5 * n);
beta_true = 5 * randn(n);
X = randn(n, number_samples);
Y = zeros(number_samples);
v = randn(number_samples);

#-

## Generate data for different values of p.
## Solve the resulting problems.
using Convex, SCS, Distributions
lsq_data = zeros(number_tests);
huber_data = zeros(number_tests);
prescient_data = zeros(number_tests);
p_vals = range(0, stop = 0.15, length = number_tests);
for i in 1:length(p_vals)
    p = p_vals[i]
    ## Generate the sign changes.
    factor = 2 * rand(Binomial(1, 1 - p), number_samples) .- 1
    Y = factor .* X' * beta_true + v

    ## Form and solve a standard regression problem.
    beta = Variable(n)
    fit = norm(beta - beta_true) / norm(beta_true)
    cost = norm(X' * beta - Y)
    prob = minimize(cost)
    solve!(prob, SCS.Optimizer; silent_solver = true)
    lsq_data[i] = evaluate(fit)

    ## Form and solve a prescient regression problem,
    ## i.e., where the sign changes are known.
    cost = norm(factor .* (X' * beta) - Y)
    solve!(minimize(cost), SCS.Optimizer; silent_solver = true)
    prescient_data[i] = evaluate(fit)

    ## Form and solve the Huber regression problem.
    cost = sum(huber(X' * beta - Y, 1))
    solve!(minimize(cost), SCS.Optimizer; silent_solver = true)
    huber_data[i] = evaluate(fit)
end

#-
using Plots

plot(p_vals, huber_data, label = "Huber", xlabel = "p", ylabel = "Fit")
plot!(p_vals, lsq_data, label = "Least squares")
plot!(p_vals, prescient_data, label = "Prescient")
#-

## Plot the relative reconstruction error for Huber and prescient regression,
## zooming in on smaller values of p.
indices = findall(p_vals .<= 0.08);
plot(p_vals[indices], huber_data[indices], label = "Huber")
plot!(p_vals[indices], prescient_data[indices], label = "Prescient")


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using QuadGK, Plots
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.3"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	sup = (-1,1)
	f1(x) = 3/4*(1-x^2)
	f2(x) = x < 0 ? x+1 : 1-x

	expect(f,support) = quadgk((x) -> x*f(x),support...)[1]
end

# ╔═╡ 3a837612-1f94-11eb-315f-0d9296c7ba90
begin
	x = -1:0.01:1
	plot(x, [f1.(x), f2.(x)], lab=["f1" "f2"])
end

# ╔═╡ aa0bdedc-1f93-11eb-12a7-f3ad5370c280
Text("Mean 1: $(expect(f1,sup)), Mean 2: $(expect(f2,sup))")

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.3"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═3a837612-1f94-11eb-315f-0d9296c7ba90
# ╠═aa0bdedc-1f93-11eb-12a7-f3ad5370c280
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


# # Direct and inverse FFTW

# Here we test the methods, in FFTW, to be used in our pseudo-spectral code. It is mostly a sanity check.

# We start by loading the packages we need:

using FFTW
using WGLMakie
using Test

# ## The spatial domain and its discretization

# We consider a square domain of sides $L = 2\pi$, for which the smallest wavenumber is $\kappa_0 = 2\pi/L$.

L = 2π
κ₀ = 2π/L
nothing

# We set the number $N$ of points for the mesh in each direction, yielding a mesh $(x_i, y_j)_{i, j = 1,\ldots ,N}$, with $x_N = y_N = L$, and steps $x_{i+1} - x_i = L/N$, and $y_{j+1} - y_j = L/N$. Due to the periodicity, we don't need to store the values corresponding to $i = j = 0$.

N = 128
x = y = (L/N):(L/N):L
nothing

# We may visualize the grid with a scatter plot, although if the mesh is too thin, we won't quite see the details. If using GLMakie or WGLMakie, one can zoom in for a detailed view.

fig, ax, plt = scatter(vec(x .* one.(y)'), vec(one.(x) .* y'))

fig

# ## A vorticity function for testing

# In order to test the methods from [FFTW.jl](http://www.fftw.org), we define a certain vorticity function and its derivatives

vort = sin.(one.(y) * x') .* cos.(3y * one.(x)')
dx_vort = cos.(one.(y) * x') .* cos.(3y * one.(x)')
dy_vort = - 3 * sin.(one.(y) * x') .* sin.(3y * one.(x)')
dd_vort = - sin.(one.(y) * x') .* cos.(3y * one.(x)') - 9 * sin.(one.(y) * x') .* cos.(3y * one.(x)')

nothing

# We may visualize the vorticity as a surface:

fig = Figure(backgroundcolor = RGBf(0.98, 0.98, 0.98))

ax = Axis3(fig[1, 1], xlabel = "x", ylabel = "y", zlabel = "vorticity level", title = "Vorticity graph")

surface!(ax, x, y, vort, colormap = :berlin)

fig

# Or, better yet, as a heatmap:

fig, ax, plt = heatmap(x, y, vort, colormap = :berlin)

ax = Axis(fig[1, 1], xlabel = "x", ylabel = "y", title = "Vorticity heatmap")

fig

# ## Testing the direct and inverse discrete Fourier transform

# We check going back and forth with the real fft and inverse real fft:

@testset "check composition `irfft ∘ rfft`" begin
    vort_hat = rfft(vort)
    vortback = irfft(vort_hat, N)
    @test vort ≈ vortback
end
nothing

# We can also visualize the excited modes:

κ_x = κ₀ * (0:1/N:1/2)
κ_y = κ₀ * (0:1/N:1-1/N)

vort_hat = rfft(vort)

fig, ax, plt = heatmap(κ_x, κ_y, abs.(vort_hat).^2, colormap = :berlin)

ax = Axis(fig[1, 1], xlabel = "k_x", ylabel = "k_y", title = "Enstrophy spectrum heatmap")

fig

# We have excited modes $(\pm 1, \pm 3)$. Due to the reality condition of the vorticity, the Fourier spectrum has an Hermitian symmetry around the origin. Using the real Fourier transform, only half of the modes need to be stored. In this case, only modes $(1, \pm 3)$ are retained. Moreover, the negative modes are shifted above. Due to the one-indexing of Julia, this means that modes $(\pm a, \pm b)$ are represented by indices $[a + 1, b + 1]$ and $[a + 1, N + 1 - b]$. In our case, the excited modes are associated with

vort_hat[4, 2]

# and

vort_hat[4, 128]

# These are the only excited modes, as we can check:

for i in 1:div(N, 2) + 1, j in 1:N
    if abs(vort_hat[i, j])^2 > eps()
        println("vort_hat[$i, $j] = $(vort_hat[i, j])")
    end
end

# ## Differential operators

# In order to check the derivatives in spectral space, we define the following operators. Actually, they are just vectors, since derivatives in spectral space act as diagonal operators. Hence, a straighforward Hadamard product suffices.

Dx_hat = im * κ₀ * [ifelse(k1 ≤ N/2 + 1, k1 - 1, k1 - 1 - N) for k2 in 1:N/2+1, k1 in 1:N]
Dy_hat = im * κ₀ * [k2 - 1 for k2 in 1:N/2+1, k1 in 1:N]

Delta_hat = - κ₀^2 * [
    ifelse(k1 ≤ N/2 + 1, (k1 - 1)^2 + (k2 - 1)^2, (k1 - 1 - N)^2 + (k2 - 1)^2)
    for k2 in 1:N/2+1, k1 in 1:N
]

Hu_hat = - Dy_hat ./ Delta_hat
Hu_hat[1, 1] = 0.0
Hv_hat = Dx_hat ./ Delta_hat
Hv_hat[1, 1] = 0.0

nothing

# ## Testing the spectral operators

# Now we are ready to test the differentiation in spectral space, with the operators just defined.

@testset "Check operators" begin
    vort_hat = rfft(vort)
    u_hat = Hu_hat .* vort_hat
    v_hat = Hv_hat .* vort_hat
    ## check derivative dx in spectral space:
    @test irfft(Dx_hat .* vort_hat, N) ≈ dx_vort
    ## check derivative dy in spectral space:
    @test irfft(Dy_hat .* vort_hat, N) ≈ dy_vort
    ## check Laplacian in spectral space:
    @test Delta_hat ≈ Dx_hat.^2 .+ Dy_hat.^2
    @test irfft(Delta_hat .* vort_hat, N) ≈ dd_vort
    ## check recovering velocity field from vorticity
    @test Dx_hat .* v_hat - Dy_hat .* u_hat ≈ vort_hat
end
nothing

# ## One-mode steady state

# For forced-periodic flows, when the forcing function contains a single Fourier mode, there is a corresponding steady state vorticity function corresponding also to that single-mode in Fourier space, only with a different amplitude.

# In order to construct such pairs of forcing mode / steady state, we set the viscosity of the flow, choose the mode to be forced and define the strength of that force:

ν = 1.0e-0 # viscosity
κ = (x = 1, y = 2) # forced mode
α = (re = 0.1, im = 0.05) # strength

nothing

# With these parameters, we find the curl `g_steady` of the forcing term and the vorticity `vort_steady` of the corresponding steady-state:

g_steady = ν * (
    2α.re * κ₀^4 * sum(abs2, κ)^2 * cos.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
    - 2α.im * κ₀^4 * sum(abs2, κ)^2 * sin.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
)

vort_steady = 2κ₀^2 * sum(abs2, κ) * (
    α.re * cos.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
    - α.im * sin.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
)
nothing

# Here are some visualizations of the curl of the forcing term and of the vorticity

heatmap(x, y, g_steady, xlabel="x", ylabel="y", title="Curl of the forcing term", titlefont=12)

heatmap(x, y, vort_steady, xlabel="x", ylabel="y", title="Steady state vorticity", titlefont=12)

nothing

# The discrete Fourier transform of this vector fields are given as

g_steady_hat = rfft(g_steady)

vort_steady_hat = rfft(vort_steady)

nothing

# ## Testing the steady state

# We are now ready to test the steadyness of this vorticity field:

@testset "Check single-mode stable steady state" begin
    u_steady_hat = Hu_hat .* vort_steady_hat
    v_steady_hat = Hv_hat .* vort_steady_hat

    u_steady = irfft(u_steady_hat, N)
    v_steady = irfft(v_steady_hat, N)

    vort_steady_back = irfft(vort_steady_hat, N)

    wu_steady_hat = rfft(vort_steady_back .* u_steady)
    wv_steady_hat = rfft(vort_steady_back .* v_steady)

    rhs_steady = g_steady_hat .+ Delta_hat .* vort_steady_hat .- Dx_hat .* wu_steady_hat .- Dy_hat .* wv_steady_hat

    vort_steady_sol_hat = - g_steady_hat ./ Delta_hat
    vort_steady_sol_hat[1, 1] = 0.0im
    ## Vanishing bilinear term on one-mode steady state
    @test maximum(abs.(Dx_hat .* wu_steady_hat .+ Dy_hat .* wv_steady_hat)) ≤ √eps()
    ## Vanishing RHS on steady state
    @test maximum(abs.(rhs_steady)) ≤ √eps()
    ## Vanishing linear Stokes on steady state
    @test maximum(abs.(g_steady_hat .+ Delta_hat .* vort_steady_hat)) ≤ √eps()

    ## Steady state equation
    @test g_steady_hat ≈ - Delta_hat .* vort_steady_hat ≈ - Delta_hat .* vort_steady_hat .- Dx_hat .* wu_steady_hat .- Dy_hat .* wv_steady_hat

    ## Steady state solution
    @test vort_steady_sol_hat ≈ vort_steady_hat
end
nothing

# All seems good.


using Bilevel
using ForwardDiff
using Bilevel: least_squares
using LinearAlgebra

n = 3
A = rand(n,n)

d = svd(A)
d.S[n] = d.S[n-1]
A = d.U*Diagonal(d.S)*d.V'

b = rand(n)
z0 = vcat(A[:],b)

function f(a)
    A = reshape(a[1:n^2],n,n) .* 2.
    b = a[n^2+1:end]
    x = least_squares(A,b)

    x
end

sol = f(z0)

# autodiff
J_auto = ForwardDiff.jacobian(f,z0)

# numerical
ϵ = sqrt(eps(1.))
J_num = zeros(size(J_auto))
for i = 1:length(z0)
    δ = zeros(length(z0))
    δ[i] = ϵ
    J_num[:,i] = (f(z0 + δ) .- sol)/ϵ
end

display(sol)
println("")
println("Auto")
display(J_auto)
println("")
display(maximum(abs.(J_auto)))
println("")
println("")
println("Finite diff")
display(J_num)
println("")
display(maximum(abs.(J_num)))
println("")
println("")

err = maximum(abs.(J_auto .- J_num))
display(err)
println("")


# ------ Smoothing functions often used to permit continuous differentiability -----

"""
    abs_smooth(x, delta_x)

Smooth out the absolute value function with a quadratic interval.
delta_x is the half width of the smoothing interval.
Typically usage is with gradient-based optimization.
"""
function abs_smooth(x, delta_x)

    if x < delta_x && x > -delta_x
        return x^2/(2*delta_x) + delta_x/2
    else
        return abs(x)
    end

end

# TODO AN: add smooth max/min with cubic splines


#=
The following iterative sequence is defined for the set of positive integers:

n → n/2 (n is even)
n → 3n + 1 (n is odd)

Using the rule above and starting with 13, we generate the following sequence:

13 → 40 → 20 → 10 → 5 → 16 → 8 → 4 → 2 → 1
It can be seen that this sequence (starting at 13 and finishing at 1) contains 10 terms. Although it has not been proved yet (Collatz Problem), it is thought that all starting numbers finish at 1.

Which starting number, under one million, produces the longest chain?

NOTE: Once the chain starts the terms are allowed to go above one million.
=#

function find_collatz(x::Int64)
  count = 0
  while x > 1
    x = iseven(x) ? div(x,2) : 3x+1
    count +=1
  end
  count
end

function calc()
  startnumber = 1
  longest = 1
  for i in 1:1_000_000
    c = find_collatz(i)
    if c > longest
      longest = c
      startnumber = i
    end
  end
  startnumber
end
@time println(calc())


using Test
include("../src/metric.jl")

@testset "Basic Loss Functions" begin

    @testset "MSE" begin
        @test mse(1, 1) == 0
        @test mse([1,1], [1,1]) == 0
        @test mse([1,1], [3,4]) == 6.5
    end

    @testset "MAE" begin
        @test mae(1, 1) == 0
        @test mae([1,1], [1,1]) == 0
        @test mae([1,1], [3,4]) == 2.5
    end

    @testset "RMSE" begin
        @test rmse(1, 1) == 0
        @test rmse([1,1], [1,1]) == 0
        @test rmse([1,1], [3,4]) == 2.5495097567963922
    end

    @testset "RMAE" begin
        @test rmae(1, 1) == 0
        @test rmae([1,1], [1,1]) == 0
        @test rmae([1,1], [3,4]) == 1.5811388300841898
    end

    @testset "RMSLE" begin
        @test rmsle(1, 1) == 0
        @test rmsle([1,1], [1,1]) == 0
        @test rmsle([100,101], [30,400]) == 1.2785158831395131
    end

    @testset "R2SCORE" begin
        @test r2score(1, 1) == 1
        @test r2score([1,1], [1,1]) == 1
        @test r2score([1,1], [3,4]) == -11.5
    end

    @testset "BINCROSSENTROPY" begin
        @test binarycrossentropy(1, 0.9) == 1
        @test binarycrossentropy([0,0,1], [0.2,0.3,0.8]) == 0.3861431206792365
    end

end


function bisect(f, a, b, tol)
    println("a = ", a, " b = ", b, " tol = ", tol)
    c = 0.0

    while (b - a) / 2 > tol
        c = (b + a) / 2
        if f(c) == 0
            return c
        end

        if f(a) * f(c) < 0
            b = c
        else
            a = c
        end
    end

    c = (b + a) / 2
    return c
end


@doc raw"""
    SquaredExponentialKernel([α=1])

The squared exponential kernel, or alternatively the Gaussian kernel, is identical to the
exponential kernel except that the Euclidean distance is squared:

```math
\kappa(\mathbf{x},\mathbf{y}) = \exp\left(-\alpha ||\mathbf{x} - \mathbf{y}||^2\right)
\qquad \alpha > 0
```

where ``\alpha`` is a scaling parameter of the squared Euclidean distance.
Just like the exponential kernel, the squared exponential kernel is an
isotropic Mercer kernel. The squared exponential kernel is more commonly known
as the radial basis kernel within machine learning communities.
"""
struct SquaredExponentialKernel{T<:AbstractFloat} <: MercerKernel{T}
    alpha::HyperParameter{T}
    SquaredExponentialKernel{T}(α::Real) where {T<:AbstractFloat} = new{T}(
        HyperParameter(convert(T,α), interval(OpenBound(zero(T)), nothing))
    )
end
SquaredExponentialKernel(α::T=1.0) where {T<:Real} = SquaredExponentialKernel{floattype(T)}(α)
GaussianKernel = SquaredExponentialKernel
RadialBasisKernel = SquaredExponentialKernel

@inline squaredexponentialkernel(z::T, α::T) where {T<:AbstractFloat} = exp(-α*z)

@inline basefunction(::SquaredExponentialKernel) = SquaredEuclidean()
@inline function kappa(κ::SquaredExponentialKernel{T}, z::T) where {T}
    squaredexponentialkernel(z, getvalue(κ.alpha))
end


using FastGaussQuadrature

abstract type AbstractQuadrature <: JalerkinType end;

#--- Quadrature

struct Quadrature{N} <: AbstractQuadrature
    """
    Quadrature rule with N nodes and weights of type T
    """
    nodes :: SVector{N, fp_precision()}
    weights :: SVector{N, fp_precision()}
end

@inline function get_nodes(q::Quadrature)
    q.nodes
end

@inline function get_weights(q::Quadrature)
    q.weights
end

@inline function quad(q::AbstractQuadrature, values::AbstractArray)
    """
    Compute integral of a function using a concrete quadrature rule,
    where 'values' are the values of the function on q.nodes
    """
    get_weights(q) ⋅ values
end

@inline function quad(q::AbstractQuadrature, f)
    """
    Compute integral of a function using a concrete quadrature formula
    """
    quad(q, f.(get_nodes(q)))
end

#--- Constructors for concrete quadrature rules

"""Gaussian quadrature rule on [-1,1] with n points"""
function GaussianQuadrature(n)
    nodes, weights = gausslegendre(n)
    Quadrature{n}(nodes, weights)
end

"""
Trapezoid quadrature rule on interval [-1, 1]
"""
function TrapezoidQuadrature()
    nodes = [-1, 1]
    weights = [1, 1]
    Quadrature{2}(nodes, weights)
end


function DX(delta, Nx, Ny)
    # This function creates a finite difference operator which represents a
    # first order x-derivitive.
    # INPUT
    # delta: resolution of the operator - scalar / a.u.
    # Nx: Number of elements in x-direction - integer / #
    # Ny: Number of elements in y-direction - integer / #
    # OUTPUT
    # operator:  (Nx*Ny X Nx*Ny)-Dim. FD-Dx operator - array / a.u.
    N = Nx*Ny
    next_vec = fill(1.0/(2.0*delta), N-1)
    if Ny > 1
        for ii in Nx:Nx:N-1 next_vec[ii] = 0.0 end
    end
    next_sparse_mat = sparse(2:N, 1:N-1, next_vec, N, N)
    operator = next_sparse_mat' - next_sparse_mat
    return operator
end

function DXX(delta, Nx, Ny)
    # This function creates a finite difference operator which represents a
    # second order x-derivitive.
    # INPUT
    # delta: resolution of the operator - scalar / a.u.
    # Nx: Number of elements in x-direction - integer / #
    # Ny: Number of elements in y-direction - integer / #
    # OUTPUT
    # operator:  (Nx*Ny X Nx*Ny)-Dim. FD-Dxx operator - array / a.u.
    N = Nx*Ny
    main_vec = fill(-2.0 /delta^2, N)
    next_vec = fill(1.0/delta^2, N-1)
    if Ny > 1
        for ii in Nx:Nx:(N-1) next_vec[ii] = 0.0 end
    end
    main_sparse_mat = sparse(1:N, 1:N, main_vec, N, N)
    next_sparse_mat = sparse(2:N, 1:(N-1), next_vec, N, N)
    operator = next_sparse_mat' + next_sparse_mat + main_sparse_mat
    return operator
end

function DY(delta, Nx, Ny)
    # This function creates a finite difference operator which represents a
    # first order y-derivitive.
    # INPUT
    # delta: resolution of the operator - scalar / a.u.
    # Nx: Number of elements in x-direction - integer / #
    # Ny: Number of elements in y-direction - integer / #
    # OUTPUT
    # operator:  (Nx*Ny X Nx*Ny)-Dim. FD-Dy operator - array / a.u.
    N = Nx*Ny
    next_vec = fill(1.0/(2*delta), (N-Ny))
    next_sparse_mat = sparse((1+Ny):N, 1:(N-Ny), next_vec, N, N)
    operator = next_sparse_mat' - next_sparse_mat
    return operator
end

function DYY(delta, Nx, Ny)
    # This function creates a finite difference operator which represents a
    # second order y-derivitive.
    # INPUT
    # delta: resolution of the operator - scalar / a.u.
    # Nx: Number of elements in x-direction - integer / #
    # Ny: Number of elements in y-direction - integer / #
    # OUTPUT
    # operator:  (Nx*Ny X Nx*Ny)-Dim. FD-Dyy operator - array / a.u.
    N = Nx*Ny
    main_vec = fill(-2.0/delta^2, N)
    next_vec = fill(1.0/delta^2, (N-Ny))
    main_sparse_mat = sparse(1:N, 1:N, main_vec, N, N)
    next_sparse_mat = sparse((1+Ny):N, 1:(N-Ny), next_vec, N, N)
    operator = next_sparse_mat + next_sparse_mat' + main_sparse_mat
    return operator
end


using ComplexSystems
using Plots

the_walkers = walker_set(5000,-5,5)
unshift!(the_walkers,walker(0.,0.,0))

the_walkers = dla(the_walkers,0.1,-5,5)

## Plotting
the_x = [w.x for w in the_walkers]
the_y = [w.y for w in the_walkers]
p = scatter(the_x,the_y,markersize=3.,
            leg=false,ticks=nothing,
            markerstrokewidth=0.,border=false)
title!("Diffusion Limit Aggregation")
png(p,"/figs/dla")

## Minkowski dimension
min_x=-5.
max_x=5.
min_y=-5.
max_y=5.

walkers_p = map(x -> Point(x.x,x.y),the_walkers);
Δs = []
Dimension = []
for Δ ∈ 0.01:0.05:0.96
    the_grid = make_square_grid(min_x,max_x,min_y,max_y,Δ,true);
    for punto ∈ walkers_p
        map(x -> inside_box(punto,x),the_grid);
    end
    non_empty_ones = sum([!x.empty for x ∈ the_grid]);
    println(non_empty_ones)
    dimension = log(non_empty_ones)/log(1/Δ)
    push!(Δs,Δ)
    push!(Dimension,dimension)
end

mini = round(Dimension[1];digits=3)
p = scatter(Δs,Dimension,legend=false)
title!("Dimension vs box size")
yaxis!("Dimension",:log10)
xaxis!("Length of square side",:log10)

png(p,"/figs/dimension")


module RandomLinearModels

using LinearAlgebra

function isstabled(A)
    all(x->abs(x)<=1, eigvals(A))
end

function isstablec(A)
    all(x->real(x)<=0, eigvals(A))
end

function controllability(A,B)
    n,m = size(B)
    R = zeros(n,n*m)
    Ak = zero(A) .+ I(n)
    for k = 1:n
        Rk = view(R, :, (k-1)*m .+ (1:m))
        Rk .= Ak*B
        Ak = Ak * A
    end
    return R
end

iscontrollable(A,B) = rank(controllability(A,B)) == size(A,1)

function genA(v, n, m)
    # Generate a random orthogonal matrix
    X = randn(n,n)
    Q = qr(X).Q
    A = Q*Diagonal(v)*Q'
end

genB(n,m) = randn(n,m)

function gendiscrete(n, m, tol::Real=1e-4)
    v = randn(n)
    v = v ./ (norm(v,Inf) + tol)     # make it (marginally) discrete stable
    A = genA(v, n, m)
    B = genB(n,m)
    A,B
end

function gencontinuous(n, m, tol::Real=1e-4)
    v = randn(n)
    v = v .- (maximum(v) + tol)     # make it (marginally) continuous stable
    A = genA(v, n, m)
    B = genB(n,m)
    A,B
end

function gencontrollable(n, m, type::Symbol=:discrete; tol::Real=1e-4, maxiter=20)
    cnt = 0
    while true
        if type == :discrete
            A,B = gendiscrete(n, m, tol)
        elseif type == :continuous
            A,B = gencontinuous(n, m, tol)
        end
        iscontrollable(A,B) && (return A,B)
        if cnt > maxiter
            throw(ErrorException("exceeded max number of attempt to find a controllable system."))
        end
        cnt += 1
    end
end

end


"""
bound(estimator, δ, n::Int)
Give the theoritical bound B such that for the estimator
    #TODO! correction math formula
```math
\\mathbb{P}( |X- \\mathbb{E}(X)|/\\sigma \\leq B) \\leq \\delta
```

"""
bound(δ::Real, n, M::EmpiricalMean) = sqrt(1 / (n * δ))
bound(δ::Real, n, M::MedianOfMean) = sqrt(32log(1 / δ) / n)
bound(δ::Real, n, M::TrimmedMean) = 19sqrt(2log(8 / δ) / n) # Typo in the review by Lugosie. It seems that the ouragous 19 does th job. It is probably improvable.
bound(δ::Real, n, M::Catoni) = sqrt(2log(2 / δ) / (n - 2log(2 / δ)))
bound(δ::Real, n, M::Huber) = 8sqrt(2log(4 / δ) / n)
""""
sqrt(2log(1 / δ) / n) # Multiply by a factor (1 + o(1)) where
o(1) = (1+O(sqrt(log(1/δ)/n)))*(1+log(log(1/δ))/log(1/δ)) goes to zeros with (log(1/δ)/n, δ)→ (0, 0)
see https://www.youtube.com/watch?v=Kr0Kl_sXsJM Q&A
"""
bound(δ::Real, n, M::LeeValiant) = 8sqrt(2log(4 / δ) / n)


"""
    rref(A::AbstractMatrix{T}) where {T<:Interval}

Computes the reduced row echelon form of the interval matrix `A` using maximum
mignitude as pivoting strategy.

### Examples

```jldoctest
julia> A = [2..4 -1..1; -1..1 2..4]
2×2 Matrix{Interval{Float64}}:
  [2, 4]  [-1, 1]
 [-1, 1]   [2, 4]

julia> rref(A)
2×2 Matrix{Interval{Float64}}:
 [2, 4]  [-1, 1]
 [0, 0]       [1.5, 4.5]
```
"""
function rref(A::AbstractMatrix{T}) where {T<:Interval}
    A1 = copy(A)
    return rref!(A1)
end

"""
    rref!(A::AbstractMatrix{T}) where {T<:Interval}

In-place version of [`rref`](@ref).
"""
function rref!(A::AbstractMatrix{T}) where {T<:Interval}
    m, n = size(A)
    minmn = min(m,n)
    @inbounds for k = 1:minmn

        if k < m
            # find maximum index
            migmax, kp = _findmax_mig(view(A, k:m, k))
            iszero(migmax) && throw(ArgumentError("Could not find a pivot with non-zero mignitude in column $k."))
            kp += k - 1

            # Swap rows k and kp if needed
            k != kp && _swap!(A, k, kp)
        end

        # Scale first column
        _scale!(A, k)

        # Update the rest
        _eliminate!(A, k)
    end
    return A
end

@inline function _findmax_mig(v)

    @inbounds begin
        migmax = mig(first(v))
        kp = firstindex(v)

        for (i, vi) in enumerate(v)
            migi = mig(vi)
            if migi > migmax
                kp = i
                migmax = migi
            end
        end
    end
    return migmax, kp
end

@inline function _swap!(A, k, kp)
    @inbounds for i = 1:size(A, 2)
        tmp = A[k,i]
        A[k,i] = A[kp,i]
        A[kp,i] = tmp
    end
end

@inline function _scale!(A, k)
    @inbounds begin
        Akkinv = inv(A[k,k])
        for i = k+1:size(A, 1)
            A[i,k] *= Akkinv
        end
    end
end

@inline function _eliminate!(A, k)
    m, n = size(A)
    @inbounds begin
        for j = k+1:n
            for i = k+1:m
                A[i,j] -= A[i,k]*A[k,j]
            end
        end

        for i = k+1:m
            A[i, k] = zero(eltype(A))
        end
    end
end


@testset "Test number of weights" begin
    for n = 1:5
        σ_parameters = MerweScaled(n=n, α=1.0, β=0.0, κ=3.0 - n)
        target_num = 2n + 1
        @test num_σ_points(σ_parameters) == target_num
        @test size(σ_parameters.m_w, 1) == target_num
        @test size(σ_parameters.Σ_w, 1) == target_num
    end
end

@testset "Test that weights are normalized" begin
    for n = 1:5
        σ_parameters = MerweScaled(n=n, α=1.0, β=0.0, κ=3.0 - n)
        @test sum(σ_parameters.m_w) - 1.0 < 1e-6
        @test sum(σ_parameters.Σ_w) - 1.0 < 1e-6
    end
end


# binocdf is the cdf for the Binomial distribution and binoccdf is its complement.
# Distributions.jl implements the cdf of its Binomial by calling out to Rmath which
# is a Julia package that wraps the corresponding functions used in R. Since we do not want to
# load the Distributions lib when starting Autotest we want a pure Julia implementation
# of binocdf. We base this code on the pure Ruby code for incomplete beta function available here:
#   http://rubydoc.info/gems/rubystats/0.2.3/Rubystats/SpecialMath#incomplete_beta-instance_method
#

# Cumulative Distribution Function evaluated at x of the Binomial(n, p) distribution:
function binocdf(x, n, p)
  x = ifloor(x)
  if x == n
    return 1.0
  else
    incompletebeta(1-p, n - x, x + 1)
  end
end

# Complementary cdf of Binomial distr func:
binoccdf(x, n, p) = 1.0 - binocdf(x, n, p)

# incompletebeta(x, a, b) evaluates incomplete beta function, here a, b > 0 and 0 <= x <= 1.
function incompletebeta(x, a, b)
  if x == 0
    return 0.0
  elseif x == 1
    return 1.0
  else
    beta_gam = exp( a * log(x) + b * log(1.0 - x) - log_beta(a, b) )
    if x < ((a+1) / (a+b+2))
      return (beta_gam * beta_fraction(x, a, b) / a)
    else
      return (1.0 - (beta_gam * beta_fraction(1.0 - x, b, a) / b))
    end
  end
end

const LOG_GAMMA_X_MAX_VALUE = 2.55e292

function log_beta(a, b)
  if (a <= 0.0) || (b <= 0.0) || (a + b) > LOG_GAMMA_X_MAX_VALUE
    return 0.0
  else
    return (lgamma(a) + lgamma(b) - lgamma(a + b))
  end
end

const XMININ = 2.23e-303
const PRECISION = 8.88e-016

function beta_fraction(x, a, b, MaxIterations = 200)
  c = 1.0
  sum_ab  = a + b
  a_plus  = a + 1.0
  a_minus = a - 1.0
  h = 1.0 - sum_ab * x / a_plus
  if abs(h) < XMININ
    h = XMININ
  end
  h     = 1.0 / h
  frac  = h
  m     = 1
  delta = 0.0

  while (m <= MaxIterations) && (abs(delta - 1.0) > PRECISION)
    m2 = 2 * m
    # even index for d
    d = m * (b - m) * x / ( (a_minus + m2) * (a + m2))
    h = 1.0 + d * h
    if abs(h) < XMININ
      h = XMININ
    end
    h = 1.0 / h
    c = 1.0 + d / c
    if abs(c) < XMININ
      c = XMININ
    end
    frac *= (h * c)
    # odd index for d
    d = -(a + m) * (sum_ab + m) * x / ((a + m2) * (a_plus + m2))
    h = 1.0 + d * h
    if abs(h) < XMININ
      h = XMININ
    end
    h = 1.0 / h
    c = 1.0 + d / c
    if abs(c) < XMININ
      c = XMININ
    end
    delta = h * c
    frac *= delta
    m += 1
  end

  return frac
end

# Compare to Distributions.jl
using Distributions

# call once to compile things
binocdf(1, 10, 0.3)
binocdf(1.0, 10.0, 0.3)
binocdf(2.0, 10, 0.3)

# Sample lower values with higher prob through a log transform.
# Often more realistic since people tend to use lower ranges/values
# frequently.
function logsample(min, max)
  shift = 1.0 - min
  logmin = log(min + shift)
  logmax = log(max + shift)
  exp(logmin + (logmax - logmin) * rand()) - shift
end

NumReps = 1e6
ourtimesum = 0.0
dtimesum = 0.0

for rep in 1:ifloor(NumReps)
  n = ifloor(logsample(2, 1e6))
  x = rand() * n # rand(0:n)
  p = rand()

  tic()
  ourres = binocdf(x, n, p)
  ourtimesum += toq()

  tic()
  dres = cdf(Binomial(n, p), x)
  dtimesum += toq()

  if !isapprox(ourres, dres)
    println("Failed for x = $x, n = $n, p = $p: our = $ourres, Distr.jl = $dres")
  end
end

println("Our average binocdf performance: ", ourtimesum / NumReps)
println("Distributions.jl average binocdf performance: ", dtimesum / NumReps)
println("Performance diff: ", round(ourtimesum / dtimesum, 2), "x")


"""
Returns the variance of data
"""
function Statistics.var(data1::VNamedD)
	σ=0.0
	μ=Statistics.mean(data1)
	n=0
	for iss in 1:length(data1)
		for ifield in names(data1[iss].d)[1]
			for ir = 1:data1[iss].n, it = 1:length(data1[iss].grid)
				n += 1
				σ += (data1[iss].d[ifield][it, ir]-μ)^2
			end
		end
	end
	return σ*inv(n)
end

function Statistics.mean(data1::VNamedD)
	n=0
	μ=0.0
	for iss in 1:length(data1)
		for ifield in names(data1[iss].d)[1]
			for ir = 1:data1[iss].n, it = 1:length(data1[iss].grid)
				n += 1
				μ += data1[iss].d[ifield][it, ir]
			end
		end
	end
	return μ*inv(n)
end

function addnoise!(dataN::VNamedD, data::VNamedD, SNR)

	σx=Statistics.var(data)

	σxN=sqrt(σx^2*inv(10^(SNR/10.)))

	# factor to be multiplied to each scalar
	α=sqrt(σxN)
	for iss in 1:length(data)
		for ifield in names(data[iss].d)[1]
			for ir = 1:data[iss].n, it = 1:length(data[iss].grid)
				dataN[iss].d[ifield][it, ir] = dataN[iss].d[ifield][it, ir] + α*Random.randn()
			end
		end
	end
end


using Pkg
Pkg.activate(pwd())

# # Optimal control
# ## Computing trajectories with no control

struct PMSM{T<:Real}
    ρ::T
    ω::T
    A::Matrix{T}
    invA::Matrix{T}

    function PMSM(ρ, ω)
        A = -ρ*[1 0; 0 1] -ω*[0 -1; 1 0]
        return new{eltype(A)}(ρ, ω, A, inv(A))
    end
end

#+

function expA(p::PMSM, t)
    ρ, ω = p.ρ, p.ω
    return exp(-ρ*t)*[cos(ω*t) sin(ω*t); -sin(ω*t) cos(ω*t)]
end

#+

ρ = 0.1
ω = 2
x0 = [0;-0.5]
q = [1;0]

ps = PMSM(ρ, ω)

# ### Exercise:
#
# Verify that the matrix exponential is computed correctly and that it is different from
# the elementwise exponential.
#
# **Hint**: The matrix exponential can also be computed directly by the `exp` function from
# the `LinearAlgebra` package.
#
# ---
# ### Solution:

# ---
#
# ### Exercise:
#
# Write two function `trajectory_fin_diff` and `trajectory_exact` which compute the
# trajectory. The first one should use the finite difference method to discretize the
# time, while the second one should use the closed-form formula.
#
# Plot both trajectories on time interval ``[0,10]`` with time discretization step
# ``\Delta t=0.01``. Since ``x(t)`` is a two-dimensional vector, plot each component
# on one axis.
#
# ---
# ### Solution:

# ---

ts = 0:0.0001:10

xs1 = trajectory_fin_diff(ps, x0, ts, q)
xs2 = trajectory_exact(ps, x0, ts, q)

plot(xs1[1,:], xs1[2,:], label="Finite differences");
plot!(xs2[1,:], xs2[2,:], label="True value")

# ## Solving the optimal control problem
# ### Exercise:
#
# Solve the optimal time for ``x_{\rm tar}= (0.25, -0.5)`` with the maximum voltage
# ``U_{\rm max} = 0.1``.
#
# **Hint**: To solve the equation above for ``t``, use the bisection method.

function bisection(f, a, b; tol=1e-6)
    fa = f(a)
    fb = f(b)
    fa == 0 && return a
    fb == 0 && return b
    fa*fb > 0 && error("Wrong initial values for bisection")
    while b-a > tol
        c = (a+b)/2
        fc = f(c)
        fc == 0 && return c
        if fa*fc > 0
            a = c
            fa = fc
        else
            b = c
            fb = fc
        end
    end
    return (a+b)/2
end

# ---
# ### Solution:

# ---

function trajectory_control(p::PMSM, x0, ts, q, U_max, p0)
    xs = zeros(length(x0), length(ts))

    for (i, t) in enumerate(ts)
        eAt  = expA(p, t)
        emAt = expA(p, -t)
        xs[:, i] = eAt*(x0 + p.invA * (I - emAt)*q + U_max/ρ*(exp(ρ*t) - 1)*p0)
    end
    return xs
end

#+

p0 = ps.ρ/(U_max*(exp(ps.ρ*τ)-1))*(expA(ps, -τ)*x_t - x0 - ps.invA*(I-expA(ps, -τ))*q)
p0 /= norm(p0)

#+

ts = range(0, τ; length=100)
traj = trajectory_control(ps, x0, ts, q, U_max, p0)

#+

plot(traj[1,:], traj[2,:], label="Optimal trajectory");
scatter!([x0[1]], [x0[2]], label="Starting point");
scatter!([x_t[1]], [x_t[2]], label="Target point")

#+

ts = 0:0.01:10
plt = plot();

for α = 0:π/4:2*π
    trj = trajectory_control(ps, x0, ts, q, U_max, [sin(α); cos(α)])
    plot!(plt, trj[1,:], trj[2,:], label="")
end

display(plt)


"""
    Berlage(; <keyword arguments>)

Create a Berlage wavelet.

# Arguments
- `dt::AbstractFloat=0.002`: sampling interval in secs.
- `f0::Real=20.0`: central frequency in Hz.
- `m::Real=2`: exponential parameter of Berlage wavelet.
- `alpha::Real=180.0`: alpha parameter of Berlage wavelet in rad/secs.
- `phi0::Real=0.0`: phase rotation in radians.

# Example
```julia
julia> w = Berlage(); plot(w);
```
**Reference**
* Aldridge, David F., 1990, The berlage wavelet: GEOPHYSICS, 55, 1508--1511.
"""
function Berlage(; dt::AbstractFloat=0.002, f0::Real=20.0, m::Real=2, alpha::Real=180.0,
                 phi0::Real=0.0)

    nw = floor(Int, 2.0/(f0*dt))
    t = dt*collect(0:1:nw-1)
    w = [ (t[i] ^ m) * exp(-alpha*t[i]) * cos(2*pi*f0*t[i] + phi0) for i in eachindex(t) ]

    w = w/maximum(w)

end


using ForwardDiff, Roots, Plots, LinearAlgebra, Optim,Statistics

function strong_backtracking(f, ∇, x, d; α=1, β=0.1, σ=0.4)
    y0, g0, y_prev, α_prev = f(x), ∇(x)⋅d, NaN, 0
    αlo, αhi = NaN, NaN
    # bracket phase
    while true
        y = f(x + α*d)
        if y > y0 + β*α*g0 || (!isnan(y_prev) && y ≥ y_prev)
            αlo, αhi = α_prev, α
            break
        end
        g = ∇(x + α*d)⋅d
        if abs(g) ≤ -σ*g0
            return α
        elseif g ≥ 0
            αlo, αhi = α, α_prev
            break
        end
        y_prev, α_prev, α = y, α, 2α
    end

    # zoom phase
    ylo = f(x + αlo*d)
    while true
        α = (αlo + αhi)/2
        y = f(x + α*d)
        if y > y0 + β*α*g0 || y ≥ ylo
            αhi = α
        else
            g = ∇(x + α*d)⋅d
            if abs(g) ≤ -σ*g0
                return α
            elseif g*(αhi - αlo) ≥ 0
                αhi = αlo
            end
            αlo = α
        end
    end
end

function bfgs(f,grad_f,x0;tol=1e-6,maxitr=10000)

    #init data
    ite = 0;x=x0;conv=false;err=0;
    xh = zeros(length(x0),maxitr+1);xh[:,1] = x0;

    #init bfgs parameters
    m = length(x)
    D = Matrix(1.0I,m,m)
    g = grad_f(x)

    #iteractions
    while ite < maxitr
        #line search
        d = -D*g

        #Inexact line search (strongWolf)
        alpha = strong_backtracking(f,grad_f,x,d;α=1, β=1e-4, σ=0.1)

        #update x
        s = alpha*d;
        x += s

        gl = grad_f(x)
        y = gl - g
        D = D - (s*y'*D + D*y*s')/(s'*y) + (1+(y'*D*y)/(s'*y))[1]*(s*s')/(s'*y)

        #update bfgs parameters
        g = gl

        #update ite and x history
        ite += 1
        xh[:,ite+1] = x

        #local minimum point
        err = norm(s);
        if err < tol ; conv=true; break; end
    end

    println(ite+1)
    #return x history. convergence and error
    return xh[:,1:ite+1], conv, err
end

#valores dos parametros
ka = 9.0
kb = 2.0
La = 10.0
Lb = 10.0

#Caso 1
#F1 = 2.0; F2 = 4.0;

#Caso 2
#F1 = 0.0 ;F2 = 4.0; #nao ta convergindo

#Caso 3
F1 = 0.0; F2 = 400.0;

#function
F(u) = 0.5e0 * ka * (sqrt(u[1] ^ 2 + (La - u[2]) ^ 2) - La) ^ 2 + 0.5e0 * kb * (sqrt(u[1] ^ 2 + (Lb + u[2]) ^ 2) - Lb) ^ 2 - F1 * u[1] - F2 * u[2];
grad_F(u) = ForwardDiff.gradient(F,u)

#tolerancia
tol = 1e-8

#numero maximo de iteracoes
iterac_max = 10000

#ponto inicial
#u0 = [-0.01,0.0]
u0 = [ 0.00,0.0]

#u*
@time u_est = bfgs(F,grad_F,u0;tol=tol,maxitr=iterac_max)
u_est = u_est[1][:,end]
println("u*: ",u_est)
println("f(u*): ",F(u_est))

#gradiente no ponto de minimo
println("∇f(u*): ",grad_F(u_est))

#verificar a hessiana no ponto de minimo
hess_F(u) = ForwardDiff.hessian(F,u)
println("H(u*): ",hess_F(u_est))
println("λ: ",eigvals(hess_F(u_est)))


export newton

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

Newton method for solving min_x f(x)

"""
function newton(f::Function,df::Function,H::Function,x::Vector;maxIter=20,atol=1e-8,out::Int=0,storeInterm::Bool=false)

    his = zeros(maxIter,2)
    X = (storeInterm) ? zeros(length(x),maxIter) : []
    i = 1; flag = -1; LL = []
    while i<=maxIter
        fk  = f(x)
        dfk = df(x)
        his[i,:] = [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)
        try
            LL    = cholfact(d2f)
        catch err
            if isa(err,Base.LinAlg.PosDefException)
                flag = -2
                his = his[1:i,:]
                break
            else
                throw(err)
            end
        end
        pk    = - (LL\dfk)
        if out>0
            @printf "iter=%04d\tf=%1.2e\t|df|=%1.2e\n" i his[i,1] his[i,2]
        end
        # update
        x  += pk
        i+=1
    end
    i = min(maxIter,i)

    if out>=0
        if flag==-1
            println(@sprintf("newton 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("newton stopped because the Hessian at iteration %d was not positive definite.",i))
        elseif out>1
            println(@sprintf("newton 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


export all_partitions
"""
`all_partitions(A::Set)` creates a `Set` containing all possible
partitions of the set `A`.

`all_partitions(n::Int)` creates the `Set` of all partitions of
the set `{1,2,...,n}`.

Both of these take an optional second argument `k` to specify that
only partitions with exactly `k` parts should be returned.
"""
function all_partitions(A::Set{T}) where T
    if length(A) < 2
        P = Partition(A)
        return Set([P])
    end
    # Set aside one element of A and recurse
    x = first(A)
    B = deepcopy(A)
    delete!(B,x)

    PB = all_partitions(B)
    PA = Set{Partition{T}}()  # place to hold partitions we create

    for P in PB
        # case 1: include x as a singleton
        P_parts = parts(P)
        push!(P_parts, Set([x]))
        Q = PartitionBuilder(P_parts)
        push!(PA,Q)

        # case 2: insert x into existing parts
        parts_list = collect(parts(P))
        np = length(parts_list)
        for k=1:np
            push!(parts_list[k],x)  # insert x into k'th part
            Q = PartitionBuilder(Set(parts_list)) # build the partition
            push!(PA,Q)
            delete!(parts_list[k],x) # take it back out
        end

    end
    return PA
end

function all_partitions(n::Int)
    if n < 0
        error("argument must be a nonnegative integer")
    end
    A = Set{Int}(collect(1:n))
    return all_partitions(A)
end

function all_partitions(A::Set{T},k::Int) where T
    if k<0
        error("Number of partitions (k) must be nonnegative")
    end
    n = length(A)
    PA = Set{Partition{T}}()
    if k>n  # no partitions possible
        return PA
    end
    if n==0 # special case when A is empty
        push!(PA, Partition(A))
        return PA
    end
    if k==0 # special case, no partitions
        return PA
    end

    x = first(A)
    B = deepcopy(A)
    delete!(B,x)

    # Step 1: partitions in which x is a singleton
    PB = all_partitions(B,k-1)
    for P in PB
        parts_set = parts(P)
        push!(parts_set, Set([x]))  # {x} as a single part
        Q = PartitionBuilder(parts_set)
        push!(PA,Q)
    end

    # Step 2: partitions in which x is with others
    PB = all_partitions(B,k)
    for P in PB
        for j=1:k  # each P has exactly k parts
            parts_list = collect(parts(P))
            push!(parts_list[j], x)  # add x to part j
            Q = PartitionBuilder(Set(parts_list))
            push!(PA,Q)
        end
    end

    return PA
end

function all_partitions(n::Int,k::Int)
    if n < 0
        error("argument must be a nonnegative integer")
    end
    A = Set{Int}(collect(1:n))
    return all_partitions(A,k)
end


# Newey-West covariance estimator
# from https://github.com/mcreel/Econometrics/blob/508aee681ca42ff1f361fd48cd64de6565ece221/src/NP/NeweyWest.jl
# under MIT licence https://github.com/mcreel/Econometrics/blob/508aee681ca42ff1f361fd48cd64de6565ece221/LICENSE
# and adapted

"""
    function newey_west(Z,nlags=0)

Returns the Newey-West estimator of the asymptotic variance matrix
INPUTS: Z, a nxk matrix with rows the vector zt'
        nlags, the number of lags
OUTPUTS: omegahat, the Newey-West estimator of the covariance matrix
"""
function newey_west(Z,nlags=0)
    n,k = size(Z)
    # de-mean the variables
    Z = Z .- mean(Z,dims=1)
    omegahat = Z'*Z/n       # sample variance
    # automatic lags?
    if nlags == 0
        nlags = max(1, round(Int, n^0.25))
    end
    # sample autocovariances
    for i = 1:nlags
       Zlag = @view(Z[1:n-i,:])
       ZZ = @view(Z[i+1:n,:])
       gamma = (ZZ'*Zlag)/n
       weight = 1.0 - (i/(nlags+1.0))
       omegahat += weight*(gamma + gamma')
    end
    return omegahat
end


"""
dₖ = solve_tCG(M, x, gradf_k, hessf_x_h, ν=1e-2, ϵ_residual = 1e-5)

Conjugate gradient method to solve the linear problem:
find d s.t. hessf_x_h(d) + gradf_k = 0

Note that we are looking for a direction d in the tangent space of M at x, and the kernel of
∇²f_M contains at least the orthogonal to this set. CG should naturally give a solution
living in the tangent space.
"""
function solve_tCG(M, x, gradfₖ, hessf_x_h; ν=1e-3, ϵ_residual = 1e-13, maxiter=1e5, safe_projection=false, printlev=0)
    norm_x(η) = norm(M, x, η)

    dₖ = zero_tangent_vector(M, x)
    rⱼ, rⱼ_prev = zero_tangent_vector(M, x), zero_tangent_vector(M, x)
    vⱼ, vⱼ_prev = zero_tangent_vector(M, x), zero_tangent_vector(M, x)

    d_type = :Unsolved
    j = 0

    (printlev>0) && @printf "\nj     norm(rⱼ)             norm(vⱼ)             ⟨vⱼ, hessf(x)[vⱼ]⟩   ν * norm(vⱼ)^2\n"
    while true
        # current residual, conjugated direction
        rⱼ = hessf_x_h(dₖ) + gradfₖ
        if safe_projection
            rⱼ = project(M, x, rⱼ)
        end

        vⱼ = - rⱼ
        if j ≥ 1
            βⱼ = norm_x(rⱼ)^2 / norm_x(rⱼ_prev)^2
            vⱼ += βⱼ * vⱼ_prev
        end

        hessf_x_vⱼ = hessf_x_h(vⱼ)
        (printlev>0) && @printf "%5i %.10e    %.10e    % .10e     %.10e\n" j norm_x(rⱼ) norm_x(vⱼ) inner(M, x, vⱼ, hessf_x_vⱼ) ν * norm_x(vⱼ)^2
        if norm_x(rⱼ) < ϵ_residual || inner(M, x, vⱼ, hessf_x_vⱼ) < ν * norm_x(vⱼ)^2 || j > maxiter
            ## Satisfying point obtained
            if j == 0
                dₖ = -gradfₖ
            end

            if norm_x(rⱼ) < ϵ_residual
                d_type = :Solved
                # printstyled("Exiting: ||rⱼ|| < ϵ : $(norm_x(rⱼ)) < $ϵ_residual\n", color=:red)
            elseif j > maxiter
                d_type = :MaxIter
                # printstyled("Exiting: j > maxiter : $j > $maxiter\n", color=:red)
            else
                d_type = :QuasiNegCurvature
                a = inner(M, x, vⱼ, hessf_x_vⱼ)
                b = ν * norm_x(vⱼ)^2
                # printstyled("Exiting: ⟨vⱼ, hessf_x(vⱼ)⟩ < ν * ||vⱼ||² : $a < $b\n", color=:red)
            end

            break
        end

        tⱼ = - inner(M, x, rⱼ, vⱼ) / inner(M, x, vⱼ, hessf_x_vⱼ)
        dₖ += tⱼ * vⱼ

        rⱼ_prev = deepcopy(rⱼ)
        vⱼ_prev = deepcopy(vⱼ)

        j += 1
    end

    return dₖ, j, d_type
end


# Get the summary of some numbers in fibonacci sequence.
v = [1,1]
s = 0
while (v[end] < 4000000)
    if v[end]%2==0 s+=v[end] end
    push!(v, v[end]+v[end-1])
end
println(s)
# 4613732


# ------------------------------------------------------------------------------------------
# # Orbit Diagrams, Surfaces of Section
# In this tutorial we will be using a `DynamicalSystem` instance to visualize aspects of the
# system.
#
# Topics:
# * `orbitdiagram`
# * `poincaresos`
# * `produce_orbitdiagram`
# * Docstrings
#
# # Orbit Diagram
# An "orbit diagram" is simply a plot that shows the long term behavior of a discrete system
# when a parameter is varied.
#
# This is exactly what the function `orbitdiagram` does:
# 1. Evolves the system for a transient amount of time.
# 2. Evolves & saves the output of the system for a chosen variable.
# 3. Changes/increments a parameter of the equations of motion.
# 4. Repeat steps 1-3 for all given parameter values!
#
# This approach is also related with the (much more advanced) [sensitivity
# analysis](http://docs.juliadiffeq.org/latest/analysis/sensitivity.html) from
# DifferentialEquations.jl
#
# ---
#
# Let's make the super-ultra-famous orbit diagram of the logistic map:
#
# $$x_{n+1} = rx_n(1-x_n)$$
# ------------------------------------------------------------------------------------------

using DynamicalSystems, PyPlot

logimap = Systems.logistic() # Systems module contains pre-defined well-known systems

# ------------------------------------------------------------------------------------------
# ---
#
# The call signature of `orbitdiagram` is:
#
# ```julia
# orbitdiagram(discrete_system, i, p_index, pvalues; n, Ttr, ...)
# ```
# * `i` is the index of the variable we want to save.
# * `p_index` is the index of the parameter we want to change.
# * `pvalues` are the values of the parameter that will change.
# * Keywords `Ttr` and `n` denote for how much transient time to evolve the system and how
# many states to save.
#
# ------------------------------------------------------------------------------------------

i = 1 # which variable to save (of course logistic map is 1D !)
n = 50 # how many values to save
Ttr = 5000 # transient iterations
p_index = 1
pvalues = 2:0.001:4  # parameter values
output = orbitdiagram(logimap, i, p_index, pvalues; n = n, Ttr = Ttr)
typeof(output)

# ------------------------------------------------------------------------------------------
# * The output is a vector of vectors. Each inner vector has length `n` and contains the
# values of the variable `i` at the given parameter value.
#
# Let's plot interactively!
# ------------------------------------------------------------------------------------------

function bf(pvalues, n, Ttr)
    logi = Systems.logistic()
    output = orbitdiagram(logi, 1, 1, pvalues; n = n, Ttr = Ttr)

    figure(figsize=(6,4))
    for (j, p) in enumerate(pvalues)
        plot(p .* ones(output[j]), output[j], linestyle = "None", # linestyle = None
        marker = "o", ms = 0.2, color = "black")
    end
    xlabel("\$r\$"); ylabel("\$x\$");
    xlim(pvalues[1], pvalues[end])
    # ylim(0, 1)
    return
end

bf(linspace(2.0, 4.0, 1000), 200, 2000)

bf(linspace(3.5, 3.6, 1000), 200, 2000)

# ------------------------------------------------------------------------------------------
# We can use `BenchmarkTools` to see how much time it takes to generate this orbit diagram
# ------------------------------------------------------------------------------------------

using BenchmarkTools
@btime output = orbitdiagram($logimap, $i, $p_index, $pvalues; n = $n, Ttr = $Ttr);
println("for total points: $(length(pvalues)*(Ttr+n)), out of which $(length(pvalues)*n) are saved")

# ------------------------------------------------------------------------------------------
# ---
#
# * `orbitdiagram` works with *any* discrete system! Check out the [documentation
# page](https://juliadynamics.github.io/DynamicalSystems.jl/latest/chaos/orbitdiagram/) for
# more!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Poincaré Surface of Section
# This is a technique to reduce a continuous system into a discrete map with 1 fewer
# dimension.
# The wikipedia entry on [Poincaré map](https://en.wikipedia.org/wiki/Poincar%C3%A9_map) has
# a lot of useful info, but the technique itself is very simple:
#
# 1. Define a hyperplane in the phase-space of the system.
# 2. Evolve the continuous system for long times. Each time the trajectory crosses this
# plane, record the state of the system.
# 3. Only crossings with a specific `direction` (either positive or negative) are allowed.
#
# And that's it! The recorded crossings are the Poincaré Surface of Section!
#
# ## Defining a hyperplane
# Let's say that our phase-space is $D$ dimensional. If the state of the system is
# $\mathbf{u} = (u_1, \ldots, u_D)$ then the equation for a hyperplane is
#
# $$
# a_1u_1 + \dots + a_Du_D = \mathbf{a}\cdot\mathbf{u}=b
# $$
# where $\mathbf{a}, b$ are the parameters that define the hyperplane.
#
# ---
#
# Here is the call signature for a function that does this:
#
# ```julia
# poincaresos(continuous_system, plane, tfinal = 100.0; direction = 1, ...)
# ```
# In code, `plane` can be either:
#
# * A `Tuple{Int, <: Number}`, like `(j, r)` : the hyperplane is defined as when the `j`
# variable of the system crosses the value `r`.
# * An `AbstractVector` of length `D+1`. The first `D` elements of the vector correspond to
# $\mathbf{a}$ while the last element is $b$. The hyperplane is defined with its formal
# equation.
#
# ---
#
# As an example, let's see a section of the Lorenz system:
# $$
# \begin{aligned}
# \dot{X} &= \sigma(Y-X) \\
# \dot{Y} &= -XZ + \rho X -Y \\
# \dot{Z} &= XY - \beta Z
# \end{aligned}
# $$
#
# ------------------------------------------------------------------------------------------

lor = Systems.lorenz()

tr = trajectory(lor, 100.0, dt = 0.01)
figure(figsize = (8,6))
plot3D(columns(tr)...);
xlabel("X"); ylabel("Y"); zlabel("Z");

psos = poincaresos(lor, (2, 0.0), 2000.0) # find where 2nd variable crosses 0.0

figure(figsize = (6,4))
plot(psos[:, 1], psos[:, 3], lw=0.0, marker ="o", ms = 1.0, color = "C1");
xlabel("X"); ylabel("Z");

# ------------------------------------------------------------------------------------------
# * We see that the surface of section is some kind of 1-dimensional object.
# * This is expected, because as we will show in the tutorial "Entropies & Dimensions" the
# Lorenz system (at least for the default parameters) lives in an almost 2-dimensional
# attractor.
#
# * This means that when you take a cut through this object, the result should be
# 1-dimensional!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# **Next, let's visualize the Poincaré Surface of Section in 3D**
# ------------------------------------------------------------------------------------------

"""
    meshgrid(x, y) -> X, Y
Create a meshgrid, such as that used in PyPlot's countour plots.
"""
function meshgrid(vx, vy)
    m, n = length(vy), length(vx)
    vx = reshape(vx, 1, n)
    vy = reshape(vy, m, 1)
    (repmat(vx, m, 1), repmat(vy, 1, n))
end

# ------------------------------------------------------------------------------------------
# We can get the attractor by calculating trajectories:
# ------------------------------------------------------------------------------------------

tr = trajectory(lor, 200.0, dt=0.01)
x, y, z = columns(tr);

# ------------------------------------------------------------------------------------------
# Next, let's make sure we can color points on our 3D plot based on where they are relative
# to a cut we'll make in the plot.
#
# To do this, we'll create a vector, `c`, that will store labels for the colors of all
# points:
# ------------------------------------------------------------------------------------------

c = Vector{String}(length(y))
for i in 1:length(y) # cut points: red
    if -0.1 < y[i] < 0.1
        c[i] = "C3"
    elseif y[i] < 0
        c[i] = "C0" # in front of cut: blue
    else
        c[i] = "C2" # behind cut: green
    end
end

figure(figsize = (8,6))
# First let's plot the attractor
plot3D(x,y,z, color = "black", lw = 0.2, alpha = 0.25)
scatter3D(x, y, z, color = c, s = 5.0)

# And then plot the PSOS plane:
x = [-20, 20]; z = [0, 40];
X, Z = meshgrid(x, z)
Y = zeros(X)

plot_surface(X, Y, Z, alpha = 0.25, color = "C1");
xlabel("X"); ylabel("Y"); zlabel("Z");

# ------------------------------------------------------------------------------------------
# # Producing an orbit diagram
#
# 1. The `orbitdiagram` function does not make sense for continuous systems. In order for
# `orbitdiagram` to have meaning one must have a map.
#
# 2. We can take advantage of the `poincaresos` function, and reduce a continuous system to
# a map.
# 3. Then, we can formally calculate an orbit diagram for a continuous system!
#
# In this example I will use the Shinriki oscillator, which displays a period-doubling route
# to chaos like the logistic map!
# ------------------------------------------------------------------------------------------

shinriki_voltage(V) = 2.295e-5*(exp(3.0038*V) - exp(-3.0038*V))
function shinriki_eom(u, p, t)
    R1 = p[1]

    du1 = (1/0.01)*(
    u[1]*(1/6.9 - 1/R1) - shinriki_voltage(u[1] - u[2]) - (u[1] - u[2])/14.5
    )

    du2 = (1/0.1)*(
    shinriki_voltage(u[1] - u[2]) + (u[1] - u[2])/14.5 - u[3]
    )

    du3 = (1/0.32)*(-u[3]*0.1 + u[2])
    return SVector{3}(du1, du2, du3)
end

# Be sure to give a mutable container for the parameter container!
shi = ContinuousDynamicalSystem(shinriki_eom, [-2, 0, 0.2], [22.0])

# ------------------------------------------------------------------------------------------
# Now we can use `shi`, our dynamical system, to calculate trajectory steps and then plot
# the result.
# ------------------------------------------------------------------------------------------

set_parameter!(shi, 1, 20.0) # parameter from 19 to 22
tr = trajectory(shi, 2000.0)
figure(figsize = (8,6))
plot3D(columns(tr[1000:end, :])..., color = "C3", alpha = 0.5, marker = "o", ms = 0.5);

# ------------------------------------------------------------------------------------------
# To get a feeling for the system, let's look at a couple `poicaresos`
# ------------------------------------------------------------------------------------------

figure(figsize=(8,4))
subplot(1,2,1)
# the function set_parameter! is useful here! (see docstring!)
R1 = 19.5
set_parameter!(shi, 1, R1)

# here I use `direction = -1` (makes more sense for the specific system)
psos = poincaresos(shi, (2, 0.0), 1000.0, Ttr = 100.0, direction = -1)

plot(psos[:, 1], psos[:, 3], lw=0.0, marker ="o", ms = 2.0, color = "C3");
xlabel("\$V_1\$"); ylabel("\$I_3\$")
title("R1 = $R1")

R1 = 21.0
set_parameter!(shi, 1, R1)
subplot(1,2,2)
psos = poincaresos(shi, (2, 0.0), 1000.0, Ttr = 100.0, direction = -1)
plot(psos[:, 1], psos[:, 3], lw=0.0, marker ="o", ms = 2.0, color = "C3");
xlabel("\$V_1\$");
title("R1 = $R1");

# ------------------------------------------------------------------------------------------
# Hm, this is interesting!
# 1. For some parameters the motion is clearly periodic (due to the distinct number of
# points)
# 2. But at other parameters the motion ***seems to be*** on a 2-dimensional manifold.
#
# **Don't be too quick to judge the second as chaotic though! In the next tutorial
# "Quantifying Chaos" I will show you tools to quantify chaotic behavior!**
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ---
#
# * We would like to see the behavior of the system while varying the parameter, like an
# orbit diagram.
#
# * We can do this by performing successive surfaces of section and then recording the
# values of a chosen variable!
#
# This procedure is bundled in the very convenient function `produce_orbitdiagram`:
#
# ```julia
# produce_orbitdiagram(continuous_system, plane, i, p_index, pvalues; kwargs...)
# ```
#
# The function produces successive `poincaresos` for the `plane`, and records the values of
# the `i`-th variable at the section.
#
# Most other arguments are simply propagated to `poincaresos`.
#
# *Lets use this function for the Shinriki oscillator example:*
#
# First we can store our input arguments for `produce_orbitdiagram` in a few input vars.
# ------------------------------------------------------------------------------------------

pvalues = linspace(19,22,201) # which parameter values to use
p_index = 1 # which parameter to change

i = 1                  # record 1st variable
plane = (2, 0.0)       # find PSOS of 2nd variable when crossing zero
tf = 200.0             # argument passed to poincaresos

# ------------------------------------------------------------------------------------------
# Now we're ready to call `produce_orbitdiagram` and plot the result!
#
# (NOTE: This computation takes some seconds, since it makes 201 PSOS!)
# ------------------------------------------------------------------------------------------

output = produce_orbitdiagram(shi, plane, i, p_index, pvalues;
                              # all keywords are passed to poincaresos:
                              tfinal = tf, Ttr = 200.0, direction = -1)

figure(figsize=(6,4))
for (j, p) in enumerate(pvalues)
    plot(p .* ones(output[j]), output[j], ls = "None",
    marker = "o", ms = 0.5, color = "black")
end
xlabel("\$R_1\$"); ylabel("\$V_1\$");

# ------------------------------------------------------------------------------------------
# ---
#
# ## Example with parameter-dependent plane
#
# Let's see one more case, using the Gissinger system:
#
# $$
# \begin{aligned}
# \dot{Q} &= \mu Q - VD \\
# \dot{D} &= -\nu D + VQ \\
# \dot{V} &= \Gamma -V + QD
# \end{aligned}
# $$
#
# ------------------------------------------------------------------------------------------

gis = Systems.gissinger([2.32865, 2.02514, 1.98312]) # initial condition necessary to see structure

set_parameter!(gis, 1, 0.112) # parameter from 0.1 to 0.15
tr = trajectory(gis, 1000.0)
figure(figsize = (8,6))
plot3D(columns(tr)...);

pvalues = linspace(0.1,0.15,201) # which parameter values to use
p_index = 1 # change 2nd parameter, which is μ

i = 2            # record i variable
j = 1            # find PSOS of j variable
offset = -0.75   # offset is important here

tf = 5000.0

output = produce_orbitdiagram(gis, (j, offset), i, p_index, pvalues;
                              # all keywords are passed to poincaresos:
                              tfinal = tf, Ttr = 2000.0)

figure(figsize=(8,6))
for (j, p) in enumerate(pvalues)
    length(output[j]) == 0 && continue
    plot(p .* ones(output[j]), output[j], ls = "None",
    marker = "o", ms = 0.5, color = "black", alpha = 0.5)
end
ylim(1,3)
xlabel("\$\\mu\$"); ylabel("\$D\$");

# ------------------------------------------------------------------------------------------
# ---
#
# 1. The above is quite cool, but for this specific system, the "optimal" hyperplane to use
# depends on the parameter `μ`.
#
# 2. In addition, it is not "optimal" to record a specific variable during the crossing, but
# instead a function of the variables.
#
# 3. For this approach, it is better to get the entire `poincaresos`s for each parameter,
# like I show in the following example:
# ------------------------------------------------------------------------------------------

# Define appropriate hyperplane for gissinger system
const ν = 0.1
const Γ = 0.9 # default parameters of the system

# I want hyperperplane defined by these two points:
Np(μ) = SVector{3}(sqrt(ν + Γ*sqrt(ν/μ)), -sqrt(μ + Γ*sqrt(μ/ν)), -sqrt(μ*ν))
Nm(μ) = SVector{3}(-sqrt(ν + Γ*sqrt(ν/μ)), sqrt(μ + Γ*sqrt(μ/ν)), -sqrt(μ*ν))

# Create hyperplane using normal vector to vector connecting points:
gis_plane(μ) = (d = (Np(μ) - Nm(μ)); [d[2], -d[1], 0, 0])

gis_plane(0.112)

μ = 0.12
set_parameter!(gis, 1, μ)
figure(figsize = (8,6))
psos = poincaresos(gis, gis_plane(μ), 2000.0, Ttr = 2000.0)
plot3D(columns(psos)..., marker = "o", ls = "None", ms = 2.0);

pvalues = linspace(0.1,0.14,201) # which parameter values to use

xs = Vector{Float64}[] # empty vector of vectors

for (i, μ) in enumerate(pvalues)
    set_parameter!(gis, 1, μ)
    N = Np(μ)
    psos = poincaresos(gis, gis_plane(μ), 4000.0, Ttr = 2000.0)
    push!(xs, [norm(N - k) for k in psos])
end

figure(figsize=(8,6))
for (j, p) in enumerate(pvalues)
    length(xs[j]) == 0 && continue
    plot(p .* ones(xs[j]), xs[j], ls = "None",
    marker = "o", ms = 0.2, color = "black", alpha = 0.5)
end
xlabel("\$\\mu\$"); ylabel("\$x\$");

figure(figsize=(8,6))
for (j, p) in enumerate(pvalues)
    length(xs[j]) == 0 && continue
    plot(p .* ones(xs[j]), xs[j], ls = "None",
    marker = "o", ms = 0.2, color = "black", alpha = 0.5)
end
ylim(2.7, 2.9)
xlabel("\$\\mu\$"); ylabel("\$x\$");

# ------------------------------------------------------------------------------------------
# # Docstrings
# ------------------------------------------------------------------------------------------

?orbitdiagram

?poincaresos

?produce_orbitdiagram

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/introduction-to-dynamicalsystems.jl/2. Orbit diagrams, Surfaces of Section.jl


function euclideanDistance(point1::Array, point2::Array)

    nVar = length(point1)
    acc = 0
    for i=1:nVar
        acc += (point1[i]-point2[i])^2
    end
    return sqrt(acc)
end

function pairwiseDistance(array1::Array, array2::Array, metric::Function)

    len = size(array1)[1]

    matrix = zeros(len, len)

    for i=1:len
        for j=i+1:len
            matrix[i,j] = euclideanDistance(array1[i,:], array2[j,:])
            matrix[j,i] = matrix[i,j]
        end
    end

    return matrix
end

function getKnearest(row::Array, k::Integer)

    Knearest = collect(1:k)
    len = length(row)

    for i=1:k
        for j=i+1:k
            if row[Knearest[i]] < row[Knearest[j]]
                Knearest[j], Knearest[i] = Knearest[i],  Knearest[j]
            end
        end
    end

    for i=k+1:len
        if row[i] < row[Knearest[1]]
            Knearest[1] = i
            for j=1:k-1
                if row[Knearest[j]] < row[Knearest[j+1]] || break
                    Knearest[j], Knearest[j+1] = Knearest[j+1],  Knearest[j]
                end
            end
        end
    end

    return Knearest
end

function randomIndexSelection(len::Integer, n::Integer, rng::Random.AbstractRNG)

    if n>len
        n = len;
    end

    indexes = Array{UInt32}(undef, n)

    i=1
    while i<=n
        inside = false
        indexes[i] = (rand(rng, UInt32)%len) + 1;
        for j=1:i
            if i!=j
                inside = indexes[i] == indexes[j]
                inside && break
            end
        end
        if !inside
            i+=1
        end
    end
    return indexes
end

function randomIndexSelection2(len::Integer, n::Integer, rng::Random.AbstractRNG)

    if n>len
        n = len;
    end

    indexes = collect(UInt32,1:len)
    selected = Array{UInt32}(undef, n)

    i=1
    for i=1:n
        inside = false
        index = (rand(rng, UInt32)%len) + 1;
        len = len - 1
        selected[i] = indexes[index]
        indexes[index] = indexes[end]
        pop!(indexes)

    end
    return selected
end

function findprevIndexes(predicate::Function, A::Array, i::Integer)
    previous = Array{Integer}(undef, 0)

    while (aux = findprev(predicate, A, i)) != nothing
        push!(previous, aux)
        i -= 1
    end

    return previous
end

function findnextIndexes(predicate::Function, A::Array, i::Integer)
    next = Array{Integer}(undef, 0)

    while (aux = findnext(predicate, A, i)) != nothing
        push!(next, aux)
        i += 1
    end

    return next
end


using DrWatson
@quickactivate "StatReth"

# %%
using StatsBase
using Distributions
using Plots

# %% 3.20 - 3.23
pdf.(Binomial(2, 0.7), 0:2)

rand(Binomial(2, 0.7), 1)

rand(Binomial(2, 0.7), 10)

N = 10_000
dummy_w = rand(Binomial(2, 0.7), N)
counts(dummy_w) ./ N |> display   # either
countmap(dummy_w)                 # or, depending on what you want

# %% 3.24, 3.25
dummy_w = rand(Binomial(9, 0.7), N)
histogram(dummy_w)

w = rand(Binomial(9, 0.6), N)

# %% 3.26
samples = let
    # %% 3.2 - 3.5
    n = 1000
    p_grid = range(0, 1, length = n)
    prob_p = ones(n)
    prob_data = @. pdf(Binomial(9, p_grid), 6)
    posterior = prob_data .* prob_p
    posterior ./= sum(posterior)
    weights = pweights(posterior)
    sample(p_grid, weights, 10_000)
end

# Version 1: Imagine this as having all the binomial distributions from the p of the
# samples added up. The random numbers are drawn from the resulting distribution.
mix = MixtureModel(Binomial.(9, samples))
w1 = rand(mix, 10_000)

# Version 2: From each p in the sample you make a binomial distribution and draw 1 value
# from it. Then you put them all into one big array.
w2 = rand.(Binomial.(9, samples), 1)
w2 = vcat(w2...)

# These are two different approaches but give the same result for lots of values drawn.
# However, if you are only drawing a few values (less than samples) the first version
# is preferrable since you can still use all the information in the samples.

histogram(w1, normalize = :probability, alpha = 0.6)
histogram!(w2, normalize = :probability, alpha = 0.6)


# Routines having to do with the normalizing function values on a grid with respect
# to a given measure.

normalizing_weights(grid, measure) = sqrt.(quadweights(grid, measure))

"""
Return a normalization for a sampling operator on a grid. The normalization ensures
a norm equivalence between the continuous norm of the function space associated
with the measure, and the discrete norm of the sampling vector.

In practice, this amounts to a diagonal scaling with the square roots of the
weights of a quadrature rule on the grid that is exact on the span of the dictionary
in the space.
"""
sampling_normalization(gb::GridBasis{T}, measure) where T =
	DiagonalOperator(gb, gb, sqrt.(quadweights(grid(gb), measure)))


module problem113

description = """
Non-bouncy numbers

Working from left-to-right if no digit is exceeded by the digit to its left it is called an increasing number; for example, 134468.
Similarly if no digit is exceeded by the digit to its right it is called a decreasing number; for example, 66420.
We shall call a positive integer that is neither increasing nor decreasing a "bouncy" number; for example, 155349.
As n increases, the proportion of bouncy numbers below n increases such that there are only 12951 numbers below one-million that are not bouncy and only 277032 non-bouncy numbers below 10^10.

How many numbers below a googol (10^100) are not bouncy?
"""

# Count the number of non-bouncy numbers having a given number of digits
# This is sufficent to solve the problem of non-bouncy's less than 10^100, numDigits = 100
# Non-bouncy count is increasing + decreasing - flat (eg, 111, 222, 333)
function non_bouncy_count(numdigits :: Integer)
  return increasing(numdigits) + decreasing(numdigits) - 9*numdigits
end

# count of increasing numbers with maximum number of digits
# Solve this by tabulating numdigits vs maxDigit in number, table[3, 7] is number of 3-digit increasing numbers having max digit=7
# A pattern appears where i,j = i-1,j + i,j-1
# First row and first column are all 1's
function increasing(numDigits :: Integer)
  table = zeros(Int, numDigits, 9)
  table[1,:] = 1
  table[:,1] = 1
  for i = 2:numDigits
    for j = 2:9
      table[i,j] = table[i-1, j] + table[i, j-1]
    end
  end
  return sum(table)
end

# Similar to increasing, tabulate the solution, but requires some care because 0 can appear in a decreasing number,
# but zero is not counted as a single digit decreasing number.
# In this function, table[4, 8] is number of 4-digit increasing numbers with min digit = 9 (offset by 1 to accomodate 0)
function decreasing(numDigits :: Integer)
  table = zeros(Int, numDigits, 10)
  table[1, :] = 1
  table[1, 1] = 0 # zero is not counted
  table[:,10] = 1
  for i = 2:numDigits
    for j = 9:-1:1
      table[i, j] = table[i-1, j] + table[i, j+1]
    end
  end
  return sum(table)
end

using Base.Test
@test increasing(1) == 9
@test increasing(2) == 45 + 9
@test decreasing(1) == 9
@test decreasing(2) == 54 + 9
@test non_bouncy_count(6) == 12951
@test non_bouncy_count(10) == 277032

end


using ComponentArrays
using DifferentialEquations
using UnPack: @unpack

tspan = (0.0, 20.0)

## Lorenz system
function lorenz!(D, u, p, t; f=0.0)
    @unpack σ, ρ, β = p
    @unpack x, y, z = u

    D.x = σ*(y - x)
    D.y = x*(ρ - z) - y - f
    D.z = x*y - β*z
    return nothing
end
function lorenz_jac!(D, u, p, t)
    @unpack σ, ρ, β = p
    @unpack x, y, z = u

    D[:x,:x] = -σ
    D[:x,:y] = σ

    D[:y,:x] = ρ
    D[:y,:y] = -1
    D[:y,:z] = -x

    D[:z,:x] = y
    D[:z,:y] = x
    D[:z,:z] = -β
    return nothing
end

lorenz_p = (σ=10.0, ρ=28.0, β=8/3)
lorenz_ic = ComponentArray(x=0.0, y=0.0, z=0.0)
lorenz_fun = ODEFunction(lorenz!, jac=lorenz_jac!)
lorenz_prob = ODEProblem(lorenz_fun, lorenz_ic, tspan, lorenz_p)

## Lotka-Volterra system
function lotka!(D, u, p, t; f=0.0)
    @unpack α, β, γ, δ = p
    @unpack x, y = u

    D.x =  α*x - β*x*y + f
    D.y = -γ*y + δ*x*y
    return nothing
end
function lotka_jac!(D, u, p, t)
    @unpack α, β, γ, δ = p
    @unpack x, y = u

    D[:x,:x] = α - β*y
    D[:x,:y] = -β*x

    D[:y,:x] = δ*y
    D[:y,:y] = -γ + δ*x
    return nothing
end

lotka_p = (α=2/3, β=4/3, γ=1.0, δ=1.0)
lotka_ic = ComponentArray(x=1.0, y=1.0)
lotka_fun = ODEFunction(lotka!, jac=lotka_jac!)
lotka_prob = ODEProblem(lotka_fun, lotka_ic, tspan, lotka_p)

## Composed Lorenz and Lotka-Volterra system
function composed!(D, u, p, t)
    c = p.c #coupling parameter
    @unpack lorenz, lotka = u

    lorenz!(D.lorenz, lorenz, p.lorenz, t, f=c*lotka.x)
    lotka!(D.lotka, lotka, p.lotka, t, f=c*lorenz.x)
    return nothing
end
function composed_jac!(D, u, p, t)
    c = p.c
    @unpack lorenz, lotka = u

    lorenz_jac!(@view(D[:lorenz,:lorenz]), lorenz, p.lorenz, t)
    lotka_jac!(@view(D[:lotka,:lotka]), lotka, p.lotka, t)

    @view(D[:lorenz,:lotka])[:y,:x] = -c
    @view(D[:lotka,:lorenz])[:x,:x] = c
    return nothing
end

comp_p = (lorenz=lorenz_p, lotka=lotka_p, c=0.01)
comp_ic = ComponentArray(lorenz=lorenz_ic, lotka=lotka_ic)
comp_fun = ODEFunction(composed!, jac=composed_jac!)
comp_prob = ODEProblem(comp_fun, comp_ic, tspan, comp_p)

## Solve problem
# We can solve the composed system...
comp_sol = solve(comp_prob, Rodas5())

# ...or we can unit test one of the component systems
lotka_sol = solve(lotka_prob, Rodas5())


## analytic.jl :special cases where eigenfunctions are known analytically

# (separable) exponential covariance function with p = 0.5
function compute_analytic(cov::CovarianceFunction{1,Exponential{T}} where {T},n::N where {N<:Integer},pts::V where {V<:AbstractVector})
	λ = cov.cov.λ
    ω = findroots(λ, n)
    ev = @. 2*λ/(λ^2*ω^2+1)
    n = @. sqrt(2)/2*sqrt(1/ω*(λ^2*ω^2*cos(ω)*sin(ω)+λ^2*ω^3-2*λ*ω*cos.(ω)^2-cos(ω)*sin(ω)+ω)+2*λ)
	ef = diagm(0 => 1.0/n)*( sin.(ω*pts') + λ*diagm(0 => ω)*cos.(ω*pts') )

	SpectralData(sqrt.(ev),ef')
end

# find all positive (>0) zeros of the transcendental function tan(ω) = 2*λ*ω/(λ^2*ω^2-1)
function findroots(λ::T, n::N) where {T <: AbstractFloat, N <: Integer}

	# define the transcendental function
	f(ω) = (λ^2*ω^2-1)*sin(ω)-2*λ*ω*cos(ω)

	# find range around singularity 1/λ
	left_point_of_range = (2*floor(1/(π*λ)-1/2))*π/2 # left odd multiple of π/2
	right_point_of_range = (2*ceil(1/(π*λ)-1/2)+1)*π/2 # right odd multiple of π/2

	# find roots before 1/λ, if any
	if left_point_of_range ≠ π/2
		roots = zeros(min(n,round(UInt64,floor(abs(1/λ/π-1/2)))))
		left_point = π/2
		right_point = 1*π
		for i = 1:length(roots)
			@inbounds roots[i] = bisect_root(f,left_point,right_point)[1]
			right_point = roots[i] + π
			left_point = left_point + π
		end
	else
		roots = zeros(0)
	end

	# find roots inside range around 1/λ
	( length(roots) ≥ n || floor(1/(π*λ)-1/2) < 0 ) ||
	push!(roots,bisect_root(f,left_point_of_range+eps(T),1/λ)[1]) # first intersection point
	( length(roots) ≥ n || ceil(1/(π*λ)-1/2) < 0 ) ||
	push!(roots,bisect_root(f,1/λ,right_point_of_range)[1]) # second intersection point

	# if the first root is zero, cut it off
	roots[1] == 0 ? shift!(roots) : [] # empty expression

	# find roots after 1/λ
	startindex = 1 + length(roots)
	if n-length(roots) > 0
		roots = [roots; zeros(n-length(roots))]
		left_point = (2*ceil(1/(π*λ)-1/2)+2)*π/2
		right_point = (2*ceil(1/(π*λ)-1/2)+3)*π/2
		for i = startindex:length(roots)
			@inbounds roots[i] = bisect_root(f,left_point,right_point)[1]
			right_point = roots[i] + π
			left_point = left_point + π
		end
	end
	return roots
end

# bissection method to find the zeros of a function in a particular interval [x1,x2]
function bisect_root(fn::Function, x1::T, x2::T) where {T<:Real}
	xm = middle(x1, x2)
	s1 = sign(fn(x1))
	s2 = sign(fn(x2))
	while x1 < xm < x2
		sm =  sign(fn(xm))

		if s1 != sm
			x2 = xm
			s2 = sm
		else
			x1 = xm
			s1 = sm
		end

		xm = middle(x1, x2)
	end

	return x1, x2
end

# helper function to find the real "mid point" of two given floating point numbers
function middle(x1::Float64, x2::Float64)
	# use the usual float rules for combining non-finite numbers
	if !isfinite(x1) || !isfinite(x2)
		return x1 + x2
	end

	# always return 0.0 when inputs have opposite sign
	if sign(x1) != sign(x2) && x1 != 0.0 && x2 != 0.0
		return 0.0
	end

	negate = x1 < 0.0 || x2 < 0.0

	x1_int = reinterpret(UInt64, abs(x1))
	x2_int = reinterpret(UInt64, abs(x2))
	unsigned = reinterpret(Float64, (x1_int + x2_int) >> 1)

	negate ? -unsigned : unsigned
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Julia 1.4.0
#     language: julia
#     name: julia-1.4
# ---

using Plots
using LaTeXStrings

# # Circle

# +
θ = 0:0.01:2π
x = @. cos(θ)
y = @. sin(θ)

p = plot(
    xlabel = L"x",
    ylabel = L"y",
    xlim = [-1, 1],
    ylim = [-1, 1],
    aspect_ratio = :equal,
)

plot!(p, x, y, label = L"x^2+y^2 = 1")
# -

# # Astroid

# +
θ = 0:0.01:2π
a = 2
x = @. a * cos(θ)^3
y = @. a * sin(θ)^3

eq = latexstring("x^{2/3}+y^{2/3} = $a^{2/3}")

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    xlim = [-2, 2],
    ylim = [-2, 2],
    aspect_ratio = :equal,
)

plot!(p, x, y, label = :false)
# -

# # Lissajous

# +
θ = 0:0.01:2π
x = @. sin(θ)
y = @. sin(2θ)

eq = latexstring("y^2=4x^2(1-x^2)")

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    xlim = [-1, 1],
    ylim = [-1, 1],
    aspect_ratio = :equal,
)

plot!(p, x, y, label=false)

# +
θ = -π/4:0.01:π/4

x = @. tan(θ)
y = @. cos(2θ)

eq = L"$y=\frac{1-x^2}{1+x^2}$"

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    xlim = [-1, 1],
    ylim = [0, 1],
    aspect_ratio = :equal,
)

plot!(p, x, y, label = :false)
# -

# # Archimedean Spiral

# +
θ = 0:0.01:2π
a = 2
r = @. a*θ
x = @. r * cos(θ)
y = @. r * sin(θ)

eq = L"r=a\theta"

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    aspect_ratio = :equal,
)

plot!(p, x, y, label = :false)
# -

# # Cardioid

# +
θ = 0:0.01:2π
a = 2
r = @. a*(1+cos(θ))
x = @. r * cos(θ)
y = @. r * sin(θ)

eq = L"r=a(1+\cos(\theta)"

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    aspect_ratio = :equal,
)

plot!(p, x, y, label = :false)
# -

# # Rose Curve

# ## n = 2

# +
θ = 0:0.01:2π
a = 2
n = 2
r = @. a*sin(n*θ)
x = @. r * cos(θ)
y = @. r * sin(θ)

eq = latexstring("r=a\\sin($n\\theta)")

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    aspect_ratio = :equal,
)

plot!(p, x, y, label = :false)
# -

# ## n = 3

# +
θ = 0:0.01:2π
a = 2
n = 3
r = @. a*sin(n*θ)
x = @. r * cos(θ)
y = @. r * sin(θ)

eq = latexstring("r=a\\sin($n\\theta)")

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    aspect_ratio = :equal,
)

plot!(p, x, y, label=false)
# -

# ## n = 4

# +
θ = 0:0.01:2π
a = 2
n = 4
r = @. a*sin(n*θ)
x = @. r * cos(θ)
y = @. r * sin(θ)

eq = latexstring("r=a\\sin($n\\theta)")

p = plot(
    title = eq,
    xlabel = L"x",
    ylabel = L"y",
    aspect_ratio = :equal,
)

plot!(p, x, y, label = :false)
# -

# ## Rose (gif)

# +
θ = 0:0.01:2π
a = 2
anim = @animate for n in 2:10
    r = @. a*sin(n*θ)
    x = @. r * cos(θ)
    y = @. r * sin(θ)

    eq = latexstring("r=a\\sin($n\\theta)")

    p = plot(
        title = eq,
        xlabel = L"x",
        ylabel = L"y",
        xlim=[-3,3],
        ylim=[-3,3],
        aspect_ratio = :equal,
    )

    plot!(p, x, y, label = :false)
end

anim;
# -

gif(anim, "roses.gif", fps=2)


using SymPy

N = 11

b = [symbols("b_$i", real=true) for i in 1:N]
B = diagm(b)

# 2nd order dissipation
D1 = zeros(Int, N, N)
for i in 1:N
    D1[i,i] = 1
end
for i in 1:N-1
    D1[i+1,i] = -1
end
D1[1,2] = 1
display(D1)

B1 = eye(Int, N); B1[1,1] = 0
display(B1)

display(D1' * B * D1)
display(D1' * B1 * D1)

# 4th order dissipation
D2 = zeros(Int, N, N)
for i in 1:N
    D2[i,i] = -2
end
for i in 1:N-1
    D2[i+1,i] = 1
    D2[i,i+1] = 1
end
boundary_width = 1
for i in 1:boundary_width
    D2[i,:] = D2[boundary_width+1,:]
    D2[end-i+1,:] = D2[end-boundary_width,:]
end
display(D2)

B2 = eye(Int, N); B2[1,1] = B2[end,end] = 0
display(B2)

display(D2' * B * D2)
display(D2' * B2 * D2)

# 6th order dissipation
D3 = zeros(Int, N, N)
for i in 1:N
    D3[i,i] = -3
end
for i in 1:N-1
    D3[i+1,i] = 3
    D3[i,i+1] = 1
end
for i in 1:N-2
    D3[i+2,i] = -1
end
boundary_width = 2
for i in 1:boundary_width
    D3[i,:] = D3[boundary_width+1,:]
end
for i in 1:boundary_width-1
    D3[end-i+1,:] = D3[end-boundary_width+1,:]
end
display(D3)

B3 = eye(Int, N); B3[1,1] = B3[2,2]= B3[end,end] = 0
display(B3)

display(D3' * B * D3)
display(D3' * B3 * D3)

# 8th order dissipation
D4 = zeros(Int, N, N)
for i in 1:N
    D4[i,i] = 6
end
for i in 1:N-1
    D4[i+1,i] = -4
    D4[i,i+1] = -4
end
for i in 1:N-2
    D4[i+2,i] = 1
    D4[i,i+2] = 1
end
boundary_width = 2
for i in 1:boundary_width
    D4[i,:] = D4[boundary_width+1,:]
    D4[end-i+1,:] = D4[end-boundary_width,:]
end
display(D4)

B4 = eye(Int, N); B4[1,1] = B4[2,2]= B4[end,end] = B4[end-1,end-1] = 0
display(B4)

display(D4' * B * D4)
display(D4' * B4 * D4)


using Plots

#graphs a cylindrical aprx with n planes, radius =r, height = h
function cylinder(n,r,h)
	p = plot(xlims=[-h,h],ylims=[-2*r,2*r],zlims=[-2*r,2*r],legend=false)
	x = range(0,h,length=10)
	alpha = range(0,2*pi,length=n)

	for theta in alpha
		if theta == alpha[end]
			break
		end
		y=r*cos(theta)*ones(length(x))
		z=r*sin(theta)*ones(length(x))
		plot!(x,y,z,seriescolor=:black)
	end
	display(p)
end

#graphs a function to be revolved, at angles theta_i in alpha
#use at least n=10,15
function PolarLines(a,b,n)

	#f is the function to revolve
	f(u) = 1+sin(u)cos(u)
	p = plot(legend=false)
	x=range(a,b,length=10)
	alpha = range(0,2*pi,length=n)
	for theta in alpha
		if theta == alpha[end]
			break
		end
		y=cos(theta)*f.(x)
		z=sin(theta)*f.(x)
		plot!(x,y,z,seriescolor=:black)
	end
	savefig(p,"fig.html")
	p
end

#graphs a function to be revolved, at vertical slices, and with PolarLines
#use at least n=10
function Rings(a,b,n,m)

	#f is the function to revolve
	f(u) = 1+sin(u)cos(u)
	q = plot([f,0],a,b)

	p = plot(legend=false)
	x=range(a,b,length=n)
	alpha = range(0,2*pi,length=m)

	for xi in x
		y = Float64[]
		z = Float64[]
		xp = xi*ones(m)
		for theta in alpha
			push!(y,cos(theta)*f(xi))
			push!(z,sin(theta)*f(xi))

		end
		plot!(p,xp,y,z,seriescolor=:black)

	end
	p1 = PolarLines(a,b,2*m)

	f = plot(
	p,p1,q,
	layout=@layout([a b;c]),
	legend=false,
	titlelocation=:left,
	bottom_margin=10px,
	window_title="Solid of revolution")
	savefig(f,"figures.html")
	f
end


# # Quadratic Program
#
# We want to solve the following quadratic program with decision variable `x`:
# $$
# \begin{array}{ll} \text{minimize} &  1/2 x^\top P x + q^\top x \\
# \text{subject to} &  l \leq A x \leq u
# \end{array}
# $$
# The problem can be solved with `COSMO` in the following way. Start by defining the problem data

using COSMO, SparseArrays, LinearAlgebra, Test

q = [1; 1.];
P = sparse([4. 1; 1 2]);
A = [1. 1; 1 0; 0 1];
l = [1.; 0; 0];
u = [1; 0.7; 0.7];

# First, we decide to solve the problem with two one-sided constraints using `COSMO.Nonnegatives` as the convex set:
Aa = [-A; A]
ba = [u; -l]
constraint1 = COSMO.Constraint(Aa, ba, COSMO.Nonnegatives);

# Next, we define the settings object, the model and then assemble everything:
settings = COSMO.Settings(verbose=true);
model = COSMO.Model();
assemble!(model, P, q, constraint1, settings = settings);
res = COSMO.optimize!(model);

# Alternatively, we can also use two-sided constraints with `COSMO.Box`:
constraint1 = COSMO.Constraint(A, zeros(3), COSMO.Box(l, u));

model = COSMO.Model();
assemble!(model, P, q, constraint1, settings = settings);
res_box = COSMO.optimize!(model);

# Let's check that the solution is correct:
@testset "QP Problem" begin
  @test norm(res.x[1:2] - [0.3; 0.7], Inf) < 1e-3
  @test norm(res_box.x[1:2] - [0.3; 0.7], Inf) < 1e-3
  @test abs(res.obj_val - 1.88) < 1e-3
  @test abs(res_box.obj_val - 1.88) < 1e-3
end
nothing


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
using Pkg, DrWatson

# ╔═╡ e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random
	Random.seed!(1)
end

# ╔═╡ ac3ceabe-1ee4-11eb-3935-313e96aafd62
md"## Listing 2.13"

# ╔═╡ e9fa266e-1eeb-11eb-2630-59930c474ed3
function montyHall(switchPolicy)
    prize, choice = rand(1:3), rand(1:3)
    if prize == choice
        revealed = rand(setdiff(1:3,choice))
    else
        revealed = rand(setdiff(1:3,[prize,choice]))
    end

    if switchPolicy
        choice = setdiff(1:3,[revealed,choice])[1]
    end
    return choice == prize
end

# ╔═╡ 9046a9c8-1f64-11eb-067f-6552cbd8919e
N = 10^6

# ╔═╡ 9046fdce-1f64-11eb-3888-651b7d2d17ac
Text("Success probability with policy I (stay): $(sum([montyHall(false) for _ in 1:N])/N)")

# ╔═╡ 90479df6-1f64-11eb-3e06-af4adc28a9d5
Text("Success probability with policy II (switch): $(sum([montyHall(true) for _ in 1:N])/N)")

# ╔═╡ ea086874-1eeb-11eb-2cb0-4186697d4294
md"## End of listing 2.13"

# ╔═╡ Cell order:
# ╟─ac3ceabe-1ee4-11eb-3935-313e96aafd62
# ╠═de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
# ╠═e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
# ╠═e9fa266e-1eeb-11eb-2630-59930c474ed3
# ╠═9046a9c8-1f64-11eb-067f-6552cbd8919e
# ╠═9046fdce-1f64-11eb-3888-651b7d2d17ac
# ╠═90479df6-1f64-11eb-3e06-af4adc28a9d5
# ╟─ea086874-1eeb-11eb-2cb0-4186697d4294


# ---
# title: 172. Factorial Trailing Zeroes
# id: problem172
# author: Tian Jun
# date: 2020-10-31
# 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

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


# Fractal Surface Interpolation demo

using Makie
using Random
using FractalTools
using GeometryBasics
using LinearAlgebra

# Get interpolation data
trig  = Triangle(
    Point(BigFloat(-1.), BigFloat(-1)),
    Point(BigFloat(1.), BigFloat(-1.)),
    Point(BigFloat(0.), BigFloat(2.))
)
f(x, y) = x^2 + y^2 + 1
pnts3d = getdata(f, trig, 100)

# Construct test points
tpnts3d = getdata(f, trig, 100)
tpnts = project(tpnts3d)

α = Node(0.001)
niter = Node(1)
err = lift(niter, α) do n, α
    itp = interpolate(pnts3d, α=α, maxiters=n)
    fval = map(pnt -> f(pnt...), tpnts)
    ival = map(pnt -> itp(pnt...), tpnts)
    val = abs.(fval - ival) ./ abs.(fval) * 100
    val
end

fig = Figure(title="Fractal Interpolation")
ax = Axis(fig[1,1], xlabel="npoints", ylabel="Relative Error (%)", title="Fractal Surface Interpolation Error Demo")
plt = lines!(ax, err)

gl = GridLayout()
gl2 = GridLayout()
gl[1, 1] = slider = Slider(fig, range=0.001 : 0.0001 : 0.01, startvalue=0.001, tellwidth=false)
gl2[1, 1] = button1 = Button(fig, label="Increment", tellwidth=false)
gl2[1, 2] = button2 = Button(fig, label="Decrement", tellwidth=false)
gl[2, 1] = gl2
gl[1, 2] = label1 = Label(fig, "alpha = $(slider.value[])")
gl[2, 2] = label2 = Label(fig, "niter = $(niter[])")
fig[2, 1] = gl

on(button1.clicks) do n
    niter[] += 1
    label2.text[] = "niter = " * string(niter[])
    autolimits!(ax)
end

on(button2.clicks) do n
    niter[] -= 1
    label2.text[] = "niter = " * string(niter[])
    autolimits!(ax)
end

on(slider.value) do val
    α[] = val
    label1.text[] = "alpha = " * string(val)
    autolimits!(ax)
end

display(fig)

# # For video recording.
# fps = 60
# record(fig.scene, "demo.mp4"; framerate = fps) do io
#     for i = 1 : 50 * fps
#         sleep(1/fps)
#         recordframe!(io)
#     end
# end


f(x,y) = 2x+3y


# https://projecteuler.net/problem=2
# 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.

result = 0
a = 1
b = 2

while a <= 4000000
    if a % 2 == 0
        result += a
    end
    a = a + b
    a, b = b, a
end

println(result)


# Based on the introductory discussion on
# https://www.maa.org/press/periodicals/loci/joma/the-sir-model-for-spread-of-disease-the-differential-equation-model

## Uncomment the following lines if used as a standalone
# using DifferentialEquations
# using Plots

function SIR(du,u,p,t)
    s,i,r = u
    b,k = p
    du[1] = -b*s*i
    du[2] = b*s*i - k*i
    du[3] = k*i
end

# TODO: make the following a doctest example
    # #initian conditions
    # p = [0.5/7900000.0,0.33]
    # u0 = [7900000.0,10.0,0.0]
    # tspan = (0.0,140.0)

    # #solve
    # sir = ODEProblem(SIR,u0,tspan,p)
    # sol = solve(sir)

    # #plot
    # plot(sol)


function gcd{T<:Integer}(a::T, b::T)
    neg = a < 0
    while b != 0
        t = b
        b = rem(a, b)
        a = t
    end
    g = abs(a)
    neg ? -g : g
end