# # Introduction
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__notebooks/01_Introduction.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__notebooks/01_Introduction.ipynb)
#
# Let's start with a classic:

message = "Hello World!"

#-

println(message)

# Julia is a general purpose programming language designed with scientific
# computing and mathematics in mind. For this reason, some Julia expressions
# look like mathematics.
#
# For example, let's write a Julia function to calculate
# a circle’s circumference $C$ as a function of the radius $r$:
#
# $C = 2 \pi r$

r = 10

C = 2π*r

# `π` is one of the
# [mathematical constants](https://docs.julialang.org/en/v1.1/base/numbers/#General-Number-Functions-and-Constants-1)
# defined in Julia. You can write `π` by typing `\pi` and pressing `<TAB>` in
# the REPL:

π

#-

pi

# The multiplication operator `*` is not needed when a literal number
# (e.g. `2`) is placed just before a variable or constant (e.g. `π`),
# see [*Numeric Literal Coefficients*](https://docs.julialang.org/en/v1.1/manual/integers-and-floating-point-numbers/#man-numeric-literal-coefficients-1).
# That makes polynomial expressions much cleaner:

f(x) = 1.56 + 2.24x + 3.47x^2

#-

f(0.04)

# This one-line function definition, i.e. `function_name(arguments) = body`, is
# the simplest we can use in Julia.

# #### Exercise 1
#
# Define an one-line function to calculate the length of the  hypotenuse of a
# right-angled triangle with catheti (legs) $a$ and $b$ as
# $\sqrt{a^{2} + b^{2}}$ using the `sqrt()` function:

## hypotenuse(...

# If your solution is correct, the following test should pass without errors.

using Test

# The `Test` module exports the `@test` macro, useful for writing unit tests of
# the code. We are going to use `≈`, `\approx<TAB>`, (`a ≈ b` or
# `isapprox(a, b)`) test if two values can be equal taking into account
# floating point errors. :

@test hypotenuse(2, 5) ≈ hypot(2, 5)

# In real life, you can use the function `hypot` to avoid
# [underflow and overflow problems](https://www.johndcook.com/blog/2010/06/02/whats-so-hard-about-finding-a-hypotenuse/).

# ## Control Flow
#
# ### `if`

distance = hypot(1.0, 3.0, 2.5)

#-

if distance <= 2.08  # Å
    "Disulphide"
elseif distance <= 6  # Å
    "Contact"
else
    "Not interacting"
end

# Each condition should be a boolean variable (`true` or `false`) or an
# expression that evaluates to a boolean value.

# ### Short-circuit evaluation

condition = true
condition && println("It's true!")

# #### Exercise 2
#
# Can you write the equivalent expression using `if` ?

## if ...

# ### Ternary operator

condition = true
result = condition ? "👍" : "👎"

# This is equivalent to:

result = if condition
        "👍"
    else
        "👎"
end

#-

result

# ## `for` loops

dna = "ATGCAT"  # dna is a string

for base in dna
    println(base)  # base is a character
end

# `for base = dna` and `for base ∈ dna` are also a possible notations, i.e.
# using `=` or `∈` instead of `in`:

for base ∈ dna
    println(base)
end

# Julia has a nice unicode support and unicode characters can be used in
# variables names, operators, etc. You can write `∈` in the REPL by typing `\in`
# and pressing the `<TAB>` key.

# ## Functions
#
# ### One-line syntax

"Return the number of codons in the sequence."
codon_number(seq) = div(length(seq), 3) # This is a Julia comment

# A strings before the function declaration is used as docstring. You can access
# the function documentation by typing `?` in the REPL and the name of the
# function, e.g.: `?codon_number`

## Try to look at the documentation of condon_number

#-

codon_number(dna)

# ### Multi-line syntax

# Functions with longer bodies are defined using the `function` keyword:

"Returns the GC number in a given sequence."
function count_gc(string)
    count = 0
    for char in string
        if char == 'C' || char == 'G'  # || is the short-circuiting or.
            ## ' ' (instead of " ") is used to define a character.
            count += 1  # i.e. count = count + 1
        end
    end
    count  # i.e. return count
end

#-

count_gc(dna)

# By default, a function returns the result of the last evaluated expression,
# `count` in the previous example. Otherwise, the keyword `return` should be
# used:

"""
`is_dna` iterates a string and returns:
- `false` if a character different from **A, C, T or G** is found
- `true` otherwise
"""
function is_dna(string)
    for char in string
        if char != 'A' && char != 'C' && char != 'T' && char != 'G'
            return false
        end
    end
    true
end

#-

is_dna("ACHL")

# Here, the `return` keyword is used inside the loop to return `false` as soon
# as an incorrect character is found.
#
# `"""` defines multiline strings. In this example, the multiline string is the
# documentation string of the function. Markdown syntax can be used in
# docstrings.

#md ?is_dna
#nb ?is_dna

# #### Exercise 3
#
# Write a function that returns `true` if the given string (argument) has at
# least one `N` using the keywords `function`, `for`, `if` and `return`.

## has_n(...

# `@testset` is useful to aggregate tests for a particular functionality:

using Test
@testset "Exercise 3" begin
    @test has_n("ACTGN")
    @test !has_n("ACTG")
end

# #### Material
#
# You can find more complete information about
# [functions](https://docs.julialang.org/en/v1.1/manual/functions/)
# and [control flow](https://docs.julialang.org/en/v1.1/manual/control-flow/)
# statements in the Julia manual.


md"""
$$
  x[n] = A \cos(\omega n T + \phi) = A \cos(2\pi f n T + \phi)
$$
- $A$: amplitude
- $\omega$: angular frequency (`radians/sec`)
- $f = \omega / 2\pi$: frequency in Hertz (`cycles/sec`)
- $\phi$: initial phase (`radians`)
- $n$: time index
- $t = nT = n / f_s$, where $f_s$ is the sampling frequency

"""
begin
  import Pkg
  Pkg.activate(mktempdir())
  Pkg.add([
      "AudioDisplay",
      "PyPlot",
  ])
end

begin
  A = .8
  f = 1000    # i.e. 1000 cycles/sec
  ϕ = pi/2
  fs = 44100  # i.e. 44100 samples/sec
  t = range(-.002, .002, step=1/fs)
  x = A .* cos.(2πf.*t .+ ϕ)
  size(x)
end

A * cos.(2π*f.*t .+ ϕ)


function getDivGrad(n1,n2,n3)
	# D = getDivGrad(n1,n2,n3)
	# builds 3D divergence operator
	D1 = kron(speye(n3),kron(speye(n2),ddx(n1)))
	D2 = kron(speye(n3),kron(ddx(n2),speye(n1)))
	D3 = kron(ddx(n3),kron(speye(n2),speye(n1)))

	Div = [D1 D2 D3]
	return Div*Div'
end

function ddx(n)
# generate 1D finite difference on staggered grid
	return d = spdiags(ones(n)*[-1 1],[0,1],n,n+1)
end

function spdiags(B,d,m,n)
# A = spdiags(B,d,m,n)
# creates a sparse matrix from its diagonals
	d = d[:]
	p = length(d)
	len = zeros(p+1,1)
	for k = 1:p
		len[k+1] = round(Int,len[k]+length(max(1,1-d[k]):min(m,n-d[k])))
	end
	a = zeros(round(Int,len[p+1]),3)
	for k = 1:p
		# Append new d[k]-th diagonal to compact form
		i = max(1,1-d[k]):min(m,n-d[k])
		a[(round(Int,len[k])+1):round(Int,len[k+1]),:] = [i i+d[k] B[i+(m>=n)*d[k],k]]
	end

	A = sparse(round(Int,a[:,1]),round(Int,a[:,2]),a[:,3],m,n)
	return A
end


# Newey-West covariance estimator
function NeweyWest(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
=#
    n,k = size(Z)
    # de-mean the variables
    Z = Z .- mean(Z,1)
    omegahat = Z'*Z/n # sample variance
    # automatic lags?
    if nlags == 0
        nlags = max(0, Int(round(n^0.25)))
    end
    # sample autocovariances
    for i = 1:nlags
       Zlag = Z[1:n-i,:]
       ZZ = Z[i+1:n,:]
       gamma = (ZZ'*Zlag)/n
       weight = 1.0 - (i/(nlags+1.0))
       omegahat += weight*(gamma + gamma')
    end
    return omegahat
end


using ApproxFun, OrdinaryDiffEq
using LinearAlgebra
using DiffEqOperators

function makesim(sts::Union{STS{:KuramotoSivashinsky}, STS{:ksiva}})
    @unpack ic, T, S, Δt, N = sts
    L = try
            L = sts.p.L
        catch e
            L = 22
        end
    # prepare initial condition
    if ic isa Union{Nothing, Integer}
        # here we integrate a pre-determined initial condition
        # for a random amount of time
        rng = MersenneTwister(isnothing(ic) ? 42 : ic)
        S += Δt*rand(rng, 10:1000)
        x = L*(0:N-1)/N
        u0 = @. cos(x) + 0.1*sin(x/8) + 0.01*cos((2π/L)*x)
    elseif ic isa AbstractVector{<:Real}
        u0 = ic
        @assert length(u0) == N
    else
        error(ICERROR)
    end
    saveat = S:Δt:(T+S)
    t, u = kuramoto_sivashinsky(u0, L, saveat)
    t = 0:Δt:T
    return @strdict u t sts
end

"""
    kuramoto_sivashinsky(u, L, saveat)

Solve the Kuramoto Sivashinsky system given by the PDE
```math
u_t = -u*u_x - u_{xx} - u_{xxxx}
```
with periodic boundary conditions and initial condition `u`.

`L` is the physical spatial length (with `length(u)` points sampled
equidistantly in it).

`saveat` dictates at which timepoints to save the system (which is solved
using `OrdinaryDiffEq`).
"""
function kuramoto_sivashinsky(u, L, saveat)
    n = length(u)                  # number of gridpoints
    F = Fourier(0..L)
    T = ApproxFun.plan_transform(F, n)
    Ti = ApproxFun.plan_itransform(F, n)

    #Linear Part
    D  = (Derivative(F) → F)[1:n,1:n]
    D2 = Derivative(F,2)[1:n,1:n]
    D4 = Derivative(F,4)[1:n,1:n]
    A = DiffEqArrayOperator(Diagonal(-D2-D4))

    #Nonlinear Part
    ks = KSFunctor(zeros(n))

    params = (D, T, Ti)
    prob = SplitODEProblem(A, ks, T*u, (0.0, saveat[end]), params)
    sol = solve(prob, ETDRK4(), dt=Float64(saveat.step), saveat=saveat,
                save_start = false, save_end = false)
    # @assert sol.t == saveat
    return sol.t, map(u -> Ti*u, sol.u)
end

struct KSFunctor{T}; tmp::Vector{T}; end
function (ks::KSFunctor)(du,u,p,t)
    D, T, Ti = p
    mul!(du,Ti,u)
    @. du = -1/2*du^2
    mul!(ks.tmp, T, du)
    mul!(du, D, ks.tmp)
end


export SG_Filter, maxDerivativeOrder, polynomialOrder, apply_SG_filter, apply_SG_filter2D

import Base: filter, length

function _Vandermonde(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)::Array{T,2}

    @assert halfWidth>=0
    @assert degree>=0

    x=T[i for i in -halfWidth:halfWidth]

    n = degree+1
    m = length(x)

    V = Array{T}(undef,m, n)

    for i = 1:m
        V[i,1] = T(1)
    end
    for j = 2:n
        for i = 1:m
            V[i,j] = x[i] * V[i,j-1]
        end
    end

    return V
end

# ================================================================

struct SG_Filter{T<:AbstractFloat,N}
    _filter_set::Array{LinearFilter_DefaultCentered{T,N},1}
end

"""
    function SG_Filter(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)

Creates a `SG_Filter` structure used to store Savitzky-Golay filters.

* filter length is 2*`halfWidth`+1
* polynomial degree is `degree`, which defines `maxDerivativeOrder`

You can apply these filters using the
* `apply_SG_filter`
* `apply_SG_filter2D`
functions.

Example:

```jldoctest
julia> sg = SG_Filter(halfWidth=5,degree=3);

julia> maxDerivativeOrder(sg)
3

julia> length(sg)
11

julia> filter(sg,derivativeOrder=2)
Filter(r=-5:5,c=[0.03497, 0.01399, -0.002331, -0.01399, -0.02098, -0.02331, -0.02098, -0.01399, -0.002331, 0.01399, 0.03497])

```
"""
function SG_Filter(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)::SG_Filter
    @assert degree>=0
    @assert halfWidth>=1
    @assert 2*halfWidth>degree

    V=_Vandermonde(T,halfWidth=halfWidth,degree=degree)
    Q,R=qr(V)
    # breaking change in Julia V1.0,
    # see https://github.com/JuliaLang/julia/issues/27397
    #
    # before Q was a "plain" matrix, now stored in compact form
    #
    # SG=R\Q'
    #
    # must be replaced by
    #
    # Q=Q*Matrix(I, size(V))
    # SG=R\Q'
    #
    Q=Q*Matrix(I, size(V))
    SG=R\Q'

    n_filter,n_coef = size(SG)

    buffer=Array{LinearFilter_DefaultCentered{T,n_coef},1}()

    for i in 1:n_filter
        SG[i,:]*=factorial(i-1)
        push!(buffer,LinearFilter(SG[i,:]))
    end

# Returns filters set
    return SG_Filter{T,n_coef}(buffer)
end

# ================================================================

"""
    function filter(sg::SG_Filter{T,N};derivativeOrder::Int=0)

Returns the filter to be used to compute the smoothed derivatives of order `derivativeOrder`.

See: `SG_Filter`
"""
function filter(sg::SG_Filter{T,N};derivativeOrder::Int=0) where {T<:AbstractFloat,N}
    @assert 0<= derivativeOrder <= maxDerivativeOrder(sg)
    return sg._filter_set[derivativeOrder+1]
end

"""
    function length(sg::SG_Filter{T,N})

Returns filter length, this is an odd number.

See: `SG_Filter`
"""
Base.length(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = length(filter(sg))

"""
    function maxDerivativeOrder(sg::SG_Filter{T,N})

Maximum order of the smoothed derivatives we can compute using `sg` filters.

See: `SG_Filter`
"""
maxDerivativeOrder(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = size(sg._filter_set,1)-1

"""
    function polynomialOrder(sg::SG_Filter{T,N})

Returns the degree of the polynomial used to construct the
Savitzky-Golay filters. This is mainly a 'convenience' function, as it
is equivalent to `maxDerivativeOrder`

See: `SG_Filter`
"""
polynomialOrder(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = maxDerivativeOrder(sg)

"""
    function apply_SG_filter(signal::Array{T,1},
                             sg::SG_Filter{T};
                             derivativeOrder::Int=0,
                             left_BE::Type{<:BoundaryExtension}=ConstantBE,
                             right_BE::Type{<:BoundaryExtension}=ConstantBE)

Applies an 1D Savitzky-Golay and returns the smoothed signal.
"""
function apply_SG_filter(signal::AbstractArray{T,1},
                         sg::SG_Filter{T};
                         derivativeOrder::Int=0,
                         left_BE::Type{<:BoundaryExtension}=ConstantBE,
                         right_BE::Type{<:BoundaryExtension}=ConstantBE) where {T<:AbstractFloat}

    return directCrossCorrelation(filter(sg,derivativeOrder=derivativeOrder),
                                  signal,
                                  left_BE,
                                  right_BE)
end

"""
    function apply_SG_filter2D(signal::Array{T,2},
                               sg_I::SG_Filter{T},
                               sg_J::SG_Filter{T};
                               derivativeOrder_I::Int=0,
                               derivativeOrder_J::Int=0,
                               min_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                               max_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                               min_J_BE::Type{<:BoundaryExtension}=ConstantBE,
                               max_J_BE::Type{<:BoundaryExtension}=ConstantBE)

Applies an 2D Savitzky-Golay and returns the smoothed signal.
"""
function apply_SG_filter2D(signal::AbstractArray{T,2},
                           sg_I::SG_Filter{T},
                           sg_J::SG_Filter{T};
                           derivativeOrder_I::Int=0,
                           derivativeOrder_J::Int=0,
                           min_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                           max_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                           min_J_BE::Type{<:BoundaryExtension}=ConstantBE,
                           max_J_BE::Type{<:BoundaryExtension}=ConstantBE) where {T<:AbstractFloat}

    return directCrossCorrelation2D(filter(sg_I,derivativeOrder=derivativeOrder_I),
                                    filter(sg_J,derivativeOrder=derivativeOrder_J),
                                    signal,
                                    min_I_BE,
                                    max_I_BE,
                                    min_J_BE,
                                    max_J_BE)
end


function randindex(c1,c2)
    # rand_index - calculates Rand Indices to compare two partitions
    # (AR, RI, MI, HI) = rand(c1,c2), where c1,c2 are vectors listing the
    # class membership, returns the "Hubert & Arabie adjusted Rand index".
    # (AR, RI, MI, HI) = rand(c1,c2) returns the adjusted Rand index,
    # the unadjusted Rand index, "Mirkin's" index and "Hubert's" index.
    #
    # See L. Hubert and P. Arabie (1985) "Comparing Partitions" Journal of
    # Classification 2:193-218

    c = counts(c1,c2,(1:maximum(c1),1:maximum(c2))) # form contingency matrix

    n = round(Int,sum(c))
    nis = sum(sum(c,2).^2)        # sum of squares of sums of rows
    njs = sum(sum(c,1).^2)        # sum of squares of sums of columns

    t1 = binomial(n,2)            # total number of pairs of entities
    t2 = sum(c.^2)                # sum over rows & columnns of nij^2
    t3 = .5*(nis+njs)

    # Expected index (for adjustment)
    nc = (n*(n^2+1)-(n+1)*nis-(n+1)*njs+2*(nis*njs)/n)/(2*(n-1))

    A = t1+t2-t3;        # no. agreements
    D = -t2+t3;          # no. disagreements

    if t1 == nc
        # avoid division by zero; if k=1, define Rand = 0
        ARI = 0
    else
        # adjusted Rand - Hubert & Arabie 1985
        ARI = (A-nc)/(t1-nc)
    end

    RI = A/t1            # Rand 1971      # Probability of agreement
    MI = D/t1            # Mirkin 1970    # p(disagreement)
    HI = (A-D)/t1        # Hubert 1977    # p(agree)-p(disagree)

    return (ARI, RI, MI, HI)
end

randindex(R::ClusteringResult, c0::AbstractVector{Int}) = randindex(assignments(R), c0)
randindex(R1::ClusteringResult, R2::ClusteringResult) = randindex(assignments(R2), assignments(R1))


#=
The following code is modified from
https://github.com/billmclean/GaussQuadrature.jl with the original license:

> The MIT License (MIT)

> Copyright (c) 2013 billmclean

> Permission is hereby granted, free of charge, to any person obtaining a copy
> of this software and associated documentation files (the "Software"), to deal
> in the Software without restriction, including without limitation the rights
> to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
> copies of the Software, and to permit persons to whom the Software is
> furnished to do so, subject to the following conditions:

> The above copyright notice and this permission notice shall be included in
> all copies or substantial portions of the Software.

> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
> FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
> AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
> LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
> OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
> SOFTWARE.
=#

# October 2013 by Bill McLean, School of Maths and Stats,
# The University of New South Wales.
#
# Based on earlier Fortran codes
#
# gaussq.f original version 20 Jan 1975 from Stanford
# gaussq.f modified 21 Dec by Eric Grosse
# gaussquad.f95 Nov 2005 by Bill Mclean
#
# This module provides functions to compute the abscissae x[j] and
# weights w[j] for the classical Gauss quadrature rules, including
# the Radau and Lobatto variants.  Thus, the sum
#
#           n
#           ∑  w[j] f(x[j])
#          j=1
#
# approximates
#
#           hi
#          ∫  f(x) w(x) dx
#          lo
#
# where the weight function w(x) and interval lo < x < hi are as shown
# in the table below.
#
# Name                      Interval     Weight Function
#
# Legendre                 -1 < x < 1          1
# Chebyshev (first kind)   -1 < x < 1     1 / sqrt(1-x²)
# Chebyshev (second kind)  -1 < x < 1       sqrt(1-x²)
# Jacobi                   -1 < x < 1     (1-x)ᵅ (1+x)ᵝ
# Laguerre                  0 < x < ∞     xᵅ exp(-x)
# Hermite                  -∞ < x < ∞      exp(-x²)
#
# In addition to these classical rules, the module generates Gauss rules
# for logarithmic weights of the form
#
#    w(x) = x^ρ log(1/x)   for 0 < x < 1.
#
# For the Jacobi and Laguerre rules we require α > -1 and
# β > -1, so that the weight function is integrable.  Likewise, for
# log weight we require ρ > -1.
#
# Use the endpoint argument to include one or both of the end points
# of the interval of integration as an abscissa in the quadrature
# rule, as follows.
#
# endpoint = NeitherEndPoint()   Default      lo < x[j] < hi, j = 1:n.
# endpoint = LeftEndPoint()      Left Radau   lo = x[1] < x[j] < hi, j = 2:n.
# endpoint = RightEndPoint()     Right Radau  lo < x[j] < x[n] = hi, j = 1:n-1.
# endpoint = BothEndPoint()      Lobatto      lo = x[1] < x[j] < x[n] = hi, j = 2:n-1.
#
#
# The code uses the Golub and Welsch algorithm, in which the abscissae
# x[j] are the eigenvalues of a symmetric tridiagonal matrix whose
# entries depend on the coefficients in the 3-term recurrence relation
# for the othonormal polynomials generated by the weighted inner product.
#
# References:
#
#   1.  Golub, G. H., and Welsch, J. H., Calculation of Gaussian
#       quadrature rules, Mathematics of Computation 23 (April,
#       1969), pp. 221-230.
#   2.  Golub, G. H., Some modified matrix eigenvalue problems,
#       Siam Review 15 (april, 1973), pp. 318-334 (section 7).
#   3.  Stroud and Secrest, Gaussian Quadrature Formulas, Prentice-
#       Hall, Englewood Cliffs, N.J., 1966.

# Enumeration type used to specify which endpoints of the integration
# interval should be included amongst the quadrature points: neither,
# left, right, or both.

abstract type EndPoint end
struct NeitherEndPoint <: EndPoint end
struct LeftEndPoint <: EndPoint end
struct RightEndPoint <: EndPoint end
struct BothEndPoint <: EndPoint end

"""
    x, w = legendregauss(T, n, endpoint::EndPoint=Bennu.NeitherEndPoint())

Returns points `x` and weights `w` for the `n`-point Gauss-Legendre rule
for the interval `-1 < x < 1` with weight function `w(x) = 1`.

Use `endpoint=LeftEndPoint()`, `RightEndPoint() ` or `BothEndPoints()` for the
left Radau, right Radau, or Lobatto rules, respectively.
"""
function legendregauss(::Type{T}, n, endpoint=NeitherEndPoint()) where {T}
    @assert n ≥ 1
    a, b = legendrecoefficients(T, n)
    return gaussrule(-one(T), one(T), a, b, endpoint)
end

"""
    x, w = legendregauss(n, endpoint::EndPoint=Bennu.NeitherEndPoint())

Convenience function with type `T = Float64`:
"""
legendregauss(n, endpoint::EndPoint=NeitherEndPoint()) = legendregauss(Float64, n, endpoint)

"""
    x, w = legendregausslobatto(T, n)

Returns points `x` and weights `w` for the `n`-point Legendre-Gauss-Lobatto rule
for the interval `-1 ≤ x ≤ 1` with weight function `w(x) = 1`.
"""
function legendregausslobatto(::Type{T}, n) where {T}
    return legendregauss(T, n, BothEndPoint())
end

"""
    x, w = legendregausslobatto(n)

Convenience function with type `T = Float64`:
"""
legendregausslobatto(n) = legendregausslobatto(Float64, n)

function legendrecoefficients(::Type{T}, n) where {T}
    a = zeros(T, n)
    b = zeros(T, n + 1)
    b[1] = sqrt(convert(T, 2))
    for k in 2:(n + 1)
        b[k] = (k - 1) / sqrt(convert(T, (2k - 1) * (2k - 3)))
    end
    return a, b
end

"""
    x, w = gaussrule(lo, hi, a, b, endpoint, maxiterations=100)

Generates the points `x` and weights `w` for a Gauss rule with weight
function `w(x)` on the interval `lo < x < hi`.

The arrays `a` and `b` hold the coefficients (as given, for instance, by
`legendrecoefficients`) in the three-term recurrence relation for the monic
orthogonal polynomials `p(0,x)`, `p(1,x)`, `p(2,x)`, ... , that is,

    p(k, x) = (x-a[k]) p(k-1, x) - b[k]² p(k-2, x),    k ≥ 1,

where `p(0, x) = 1` and, by convention, `p(-1, x) = 0` with

              hi
    b[1]^2 = ∫  w(x) dx.
             lo

Thus, `p(k, x) = xᵏ + lower degree terms` and

     hi
    ∫  p(j, x) p(k, x) w(x) dx = 0 if j ≠ k.
    lo
"""
function gaussrule(lo, hi, a, b, endpoint::EndPoint, maxiterations=100)
    T = promote_type(typeof(lo), typeof(hi), eltype(a), eltype(b))
    n = length(a)
    @assert length(b) == n + 1
    if endpoint isa LeftEndPoint
        if n == 1
            a[1] = lo
        else
            a[n] = tridiagonalshiftsolve(n, lo, a, b) * b[n]^2 + lo
        end
    elseif endpoint isa RightEndPoint
        if n == 1
            a[1] = hi
        else
            a[n] = tridiagonalshiftsolve(n, hi, a, b) * b[n]^2 + hi
        end
    elseif endpoint isa BothEndPoint
        if n == 1
            error("Must have at least two points for both ends.")
        end
        g = tridiagonalshiftsolve(n, lo, a, b)
        t1 = (hi - lo) / (g - tridiagonalshiftsolve(n, hi, a, b))
        b[n] = sqrt(t1)
        a[n] = lo + g * t1
    end
    w = zero(a)
    tridiagonaleigenproblem!(a, b, w, maxiterations)
    for i in 1:n
        w[i] = (b[1] * w[i])^2
    end
    idx = sortperm(a)
    # Ensure end point values are exact.
    if endpoint isa LeftEndPoint || endpoint isa BothEndPoint
        a[idx[1]] = lo
    elseif endpoint isa RightEndPoint || endpoint isa BothEndPoint
        a[idx[n]] = hi
    end
    return a[idx], w[idx]
end

function tridiagonalshiftsolve(n, shift, a, b) where {T}
    #
    # Perform elimination to find the nth component s = delta[n]
    # of the solution to the nxn linear system
    #
    #     ( J_n - shift I_n ) delta = e_n,
    #
    # where J_n is the symmetric tridiagonal matrix with diagonal
    # entries a[i] and off-diagonal entries b[i], and e_n is the nth
    # standard basis vector.
    #
    t = a[1] - shift
    for i in 2:(n - 1)
        t = a[i] - shift - b[i]^2 / t
    end
    return one(t) / t
end

function tridiagonaleigenproblem!(d, e, z, maxiterations)
    #
    # Finds the eigenvalues and first components of the normalised
    # eigenvectors of a symmetric tridiagonal matrix by the implicit
    # QL method.
    #
    # d[i]   On entry, holds the ith diagonal entry of the matrix.
    #        On exit, holds the ith eigenvalue.
    #
    # e[i]   On entry, holds the [i,i-1] entry of the matrix for
    #        i = 2, 3, ..., n.  (The value of e[1] is not used.)
    #        On exit, e is overwritten.
    #
    # z[i]   On exit, holds the first component of the ith normalised
    #        eigenvector associated with d[i].
    #
    # maxiterations The maximum number of QL iterations.
    #
    # Martin and Wilkinson, Numer. Math. 12: 377-383 (1968).
    # Dubrulle, Numer. Math. 15: 450 (1970).
    # Handbook for Automatic Computation, Vol ii, Linear Algebra,
    #        pp. 241-248, 1971.
    #
    # This is a modified version of the Eispack routine imtql2.
    #
    T = promote_type(eltype(d), eltype(e), eltype(z))
    n = length(z)
    z[1] = one(T)
    z[2:n] .= zero(T)
    e[n + 1] = zero(T)

    if n == 1 # Nothing to do for a 1x1 matrix.
        return
    end
    for l in 1:n
        for j in 1:maxiterations
            # Look for small off-diagonal elements.
            m = n
            for i in l:(n - 1)
                if abs(e[i + 1]) <= eps(T) * (abs(d[i]) + abs(d[i + 1]))
                    m = i
                    break
                end
            end
            p = d[l]
            if m == l
                continue
            end
            if j == maxiterations
                msg = string("No convergence after ", j, " iterations",
                             " (try increasing maxiterations)")
                error(msg)
            end
            # Form shift
            g = (d[l + 1] - p) / (2 * e[l + 1])
            r = hypot(g, one(T))
            g = d[m] - p + e[l + 1] / (g + copysign(r, g))
            s = one(T)
            c = one(T)
            p = zero(T)
            for i in (m - 1):-1:l
                f = s * e[i + 1]
                b = c * e[i + 1]
                if abs(f) < abs(g)
                    s = f / g
                    r = hypot(s, one(T))
                    e[i + 2] = g * r
                    c = one(T) / r
                    s *= c
                else
                    c = g / f
                    r = hypot(c, one(T))
                    e[i + 2] = f * r
                    s = one(T) / r
                    c *= s
                end
                g = d[i + 1] - p
                r = (d[i] - g) * s + 2 * c * b
                p = s * r
                d[i + 1] = g + p
                g = c * r - b
                # Form first component of vector.
                f = z[i + 1]
                z[i + 1] = s * z[i] + c * f
                z[i] = c * z[i] - s * f
            end
            d[l] -= p
            e[l + 1] = g
            e[m + 1] = zero(T)
        end
    end
    return
end


using ProgressBars, Plots, StatsBase, Statistics, LaTeXStrings, JLD

function RandomWalker2D(TotalSteps)
    choices = ((1,0),(0,1),(-1,0),(0,-1))
    return reduce(.+,rand(choices,TotalSteps))
end

function RGyration(ensemble)
    MassCenter = (mean(getindex.(ensemble,1)),mean(getindex.(ensemble,2)))
    return √(sum(pos -> reduce(.+, (pos .- MassCenter).^2), ensemble)/length(ensemble))
end

TotalSteps = 40
RWNumbers = 10000000

Data = [[RandomWalker2D(TS) for i ∈ 1:RWNumbers] for TS ∈ ProgressBar(1:TotalSteps)]

RgList = [RGyration(Data[i]) for i ∈ 1:TotalSteps]
save("../../Data/Q1/Q1-Rg.jld", "RgList", RgList)

plot(1:TotalSteps,sqrt.(1:TotalSteps), label = L"Y = \sqrt{X}", c = :black, linestyle = :dash)
scatter!(RgList, label = L"Data\ Point", framestyle = :box, c = :steelblue, legend = 140)
plot!(xlabel = L"Time", ylabel = L"{R_g}_{(t)}",
    title = L"Radius\ of\ Gyration\ of\ RW\ per\ Time\ for\ 10^6 \ Particles")
savefig("../../Figs/Q1/Q1-Rg(t).pdf")


#=
Popis
  Funkcia vypocita pre kazdy vrchol hodnotu jeho closeness centrality.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozena vlastnost vrchola :: String
  normalize : parameter pre aplikovanie normalizacie :: Bool

Navratova hodnota
  closeness : slovnik obsahujuci mapovanie vrcholov a k nim prisluchajucich hodnot
           closeness centrality (index_vrchola => hodnota_closeness_centrality) :: Dict{Int64, Float64}()

Poznamky
  Hodnota centrality moze byt normalizovana nastavenim parametra normalize na True/False.
  Vypocet centrality je upraveny aby velkost grafu neovplyvnovala hodnotu centrality.
=#

function centrality_closeness(g::GenericGraph; dist_key::AbstractString="weight", normalize::Bool=true)

  epi = AttributeEdgePropertyInspector{Float64}(dist_key)
  closeness = Dict{Int64, Float64}()

  for v in Graphs.vertices(g)

      distances = filter(x -> x != Inf, dijkstra_shortest_paths(g, epi, [v]).dists)
      total_dist = sum(distances)

      if total_dist > 0 && num_vertices(g) > 1

        closeness[v.index] = (length(distances) - 1.0) / total_dist

        if normalize

          norm_coefficient = (length(distances)-1.0) / ( num_vertices(g) - 1 )
          closeness[v.index] *= norm_coefficient

        end

      end

  end

  return closeness

end


@doc raw"""
    SquaredEuclidean

The squared Euclidean function is defined by:

```math
f(\mathbf{x}, \mathbf{y}) = (\mathbf{x} - \mathbf{y})^{\intercal}(\mathbf{x} - \mathbf{y})
```
"""
struct SquaredEuclidean <: Metric end

@inline base_aggregate(::SquaredEuclidean, s::T, x::T, y::T) where {T} = s + (x-y)^2

@inline isstationary(::SquaredEuclidean) = true
@inline isisotropic(::SquaredEuclidean)  = true


"""
    The Tower of Hanoi

# Brief:
    The Tower of Hanoi is a mathematical game or puzzle consisting of three rods and a number
    of disks of various diameters, which can slide onto any rod.
    The puzzle begins with the disks stacked on one rod in order of decreasing size,
    the smallest at the top, thus approximating a conical shape.
    The objective of the puzzle is to move the entire stack to the last rod,
    obeying the following rules:
        1. Only one disk may be moved at a time.
        2. Each move consists of taking the upper disk from one of the stacks and placing it on
           top of another stack or on an empty rod.
        3. No disk may be placed on top of a disk that is smaller than it.

# Complexity: O(2^n)

# Functions
    - solveUtil(tower1, tower2, tower3, n::Int, solution::Array{Pair}) - The recursive function which finds the solution and saves it to
                                                                         the array of pairs where each pair represents one move
    - solve(tower1, tower2, tower3,n::Int) - The main function for finding the set of steps to move n rings from tower 1 to tower 3, using tower 2
    - printSolution(solution::Array{Pair}) - Prints the solution given as an array of pairs by printing each pair ( step ) in a different line

# Reference:- [Wikipedia](https://en.wikipedia.org/wiki/Tower_of_Hanoi)

# Contributed by:- [Nikola Mircic](https://github.com/Nikola-Mircic)
"""

module Hanoi
    function solveUtil(tower1, tower2, tower3, n::Int, solution::Array{Pair})
        if n==1
            push!(solution, tower1 => tower3) # There is only one ring, so just move it from tower1 to tower3
        else
            solveUtil(tower1, tower3, tower2, n-1, solution) # Move n-1 rings from tower1 to tower2 to free the bottom ring
            push!(solution, tower1 => tower3) # Move the bottom ring from tower1 to tower3
            solveUtil(tower2, tower1, tower3, n-1, solution) # Now move those n-1 rings from tower2 to tower3
        end
    end

    function solve(tower1, tower2, tower3,n::Int)
        solution = Array{Pair}(undef, 0)
        solveUtil(tower1, tower2, tower3, n, solution)
        return solution
    end

    function printSolution(solution::Array{Pair})
        for step in solution
            println(step)
        end
    end
end


"""
    permutation(n, r)

    returns the number of ways to choose r items from n items
    with order and without repetition

# Arguments:
- `n`: Positive integers of items to choose from
- 'r': Positive integers of items to choose

Contributed By:- [Mayur Dahibhate](https://github.com/mayurdahibhate)
"""

function permutation(n, r)
    function factorial(n)
        fact = 1

        if n == 0 || n == 1
            return fact
        end

        for i in 1:n
            fact = fact * i
        end

        return fact
    end

    perm = factorial(n) / factorial(n - r)

    return convert(Int64, perm)
end


using QuadGK: quadgk

export fourier_coefficient, FourierSeries, eval_series, rotating_frame_series

"""
    fourier_coefficient(f::Function, frequency::Real, harmonic::Int)

Calculate the fourier coefficient `1/T ∫_0^T f(τ) e^(-2πi * frequency * harmonic * τ) dτ`
of a periodic function `f`.

# args
* `f`: the function of one variable.
* `frequency`: the frequency of periodicity of `f`.
* `harmonic`: the desired harmonic.

# returns
The fourier coefficient.
"""
function fourier_coefficient(f::Function, frequency::Real, harmonic::Int)
    T = 1/frequency
    g(t) = f(t) * exp(-2π * 1im * frequency * harmonic * t)
    integral, _ = quadgk(g, 0.0, T, atol=sqrt(eps(T)))
    return integral/T
end

"""
Representation of a Fourier series `f(t) = ∑ Aₙ e^(iωnt)`
"""
struct FourierSeries
    frequency::Real
    terms::Dict{Int, <:Number} # harmonic => coefficient
end

"""
    FourierSeries(f::Function, frequency::Real, harmonics::AbstractVector{Int})

Construct a FourierSeries from a periodic function `f` on given harmonics.

# args
* `f`: the periodic function.
* `frequency`: the frequency of periodicity of f.
* `harmonics`: a vector of harmonics to compute.

# returns
A FourierSeries.
"""
function FourierSeries(f::Function, frequency::Real, harmonics::AbstractVector{Int})
    terms = Dict(h => fourier_coefficient(f, frequency, h) for h in harmonics)
    return FourierSeries(frequency, terms)
end

"""
    eval_series(fs::FourierSeries, times::AbstractVector{<:Real})

Evaluate a FourierSeries on given times.

# args
* `fs`: a FourierSeries.
* `times`: a vector of times.

# returns
A vector of values.
"""
function eval_series(fs::FourierSeries, times::AbstractVector{<:Real})
    return reduce(+, coeff * exp.(2π * 1im * fs.frequency * harmonic * times) for (harmonic, coeff) in fs.terms)
end

"""
    rotating_frame_series(fs::FourierSeries, harmonics::AbstractVector{Int})

Compute the fourier series for `exp(i∫_0^t f(τ) dτ)` where `f` is the function given by
the FourierSeries `fs`. This function ignores a constant term in `fs` since that maps onto
an overall exponential factor which is not part of a FourierSeries.

# args
* `fs`: a FourierSeries.
* `harmonics`: the desired harmonics.

# returns
The FourierSeries.
"""
function rotating_frame_series(fs::FourierSeries, harmonics::AbstractVector{Int})
    iω = 2π * 1im * fs.frequency
    function exp_of_integral(t::Real)
        int = 0.0
        for (h, c) in fs.terms
            if h != 0
                int += c * (exp(iω * h * t) - 1)/(iω * h)
            end
        end
        return exp(1im * int)
    end
    return FourierSeries(exp_of_integral, fs.frequency, harmonics)
end


"""
Finds roots of 4th order polynomials of the form
```math
ax⁴ + bx³ + cx² + dx + e = 0
```
based on the Fast Quartic and Cubic Solver
https://github.com/NKrvavica/fqs
by Nino Krvavica.
"""
module QuarticSolver

function solvequadratic(a0, b0, c0)
    inva0 = 1/a0
    a, b = b0*inva0, c0*inva0

    a₀ = -0.5*a
    Δ = a₀^2 - b
    sqrtΔ = sqrt(Δ)

    r1 = a₀ - sqrtΔ
    r2 = a₀ + sqrtΔ

    return r1, r2
end

function solvecubic(a0, b0, c0, d0)
    inva0 = 1/a0
    a, b, c = b0*inva0, c0*inva0, d0*inva0

    third = 1/3
    thirda = a*third
    thirda² = thirda^2
    sqrt3 = sqrt(3)

    f = third*b - thirda²
    g = thirda*(2*thirda² - b) + c
    h = 0.25*g^2 + f^3  # discriminant (`Δ` or `d` in some papers)

    if f == g == h == 0  # this check seems a little strong
        r1 = -cbrt(c)
        return r1, r1, r1
    elseif h isa Complex || h <= 0   # casus irreducibilis
        j = sqrt(-f)
        k = acos(-0.5*g/j^3)
        m = cos(third*k)
        n = sqrt3*sin(third*k)
        r1 = 2*j*m - thirda
        r2 = -j*(m + n) - thirda
        r3 = -j*(m - n) - thirda
        return r1, r2, r3
    else
        sqrth = sqrt(h)
        S = cbrt(-0.5*g + sqrth)
        U = cbrt(-0.5*g - sqrth)
        SplusU = S + U
        SminusU = S - U
        tmp = SminusU*sqrt3*0.5im
        r1 = SplusU - thirda
        r2 = -0.5*SplusU - thirda + tmp
        r3 = -0.5*SplusU - thirda - tmp
        return r1, r2, r3
    end
end

function solvequartic(a0, b0, c0, d0, e0)
    inva0 = 1/a0
    a, b, c, d = b0*inva0, c0*inva0, d0*inva0, e0*inva0

    a₀ = 0.25*a
    a₀² = a₀^2

    # Subsidiary cubic equation
    p = 3*a₀² - 0.5*b
    q = a*a₀² - b*a₀ + 0.5*c
    r = 3*a₀²^2 - b*a₀² + c*a₀ - d

    # One root of cubic
    z0, _, _ = solvecubic(1, p, r, p*r-0.5*q^2)
    z0 = complex(z0)

    s = sqrt(2*p + 2*z0)
    s == 0 ? t = z0^2 + r : t = -q/s

    r0, r1 = solvequadratic(1, s, z0 + t)
    r2, r3 = solvequadratic(1, -s, z0 - t)

    return r0-a₀, r1-a₀, r2-a₀, r3-a₀
end

end # module


using DifferentialEquations
using SimpleDiffEq
using Tables
using DataFrames
using StatsPlots
using BenchmarkTools

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

δt = 0.1
tmax = 40.0
tspan = (0.0,tmax)

u0 = [990.0,10.0,0.0]; # S,I,R

p = [0.05,10.0,0.25]; # β,c,γ

prob_ode = ODEProblem(sir_ode!, u0, tspan, p);

sol_ode = solve(prob_ode, dt = δt);

df_ode = DataFrame(Tables.table(sol_ode'))
rename!(df_ode,["S","I","R"])
df_ode[!,:t] = sol_ode.t;

@df df_ode plot(:t,
    [:S :I :R],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number")

@benchmark solve(prob_ode, dt = δt);


using LinearAlgebra, Distances
export Diffusion
export hamming_dist, jaccard_dist, diffusion_dist

"""
Compute the Hamming distance between two adjacency matrices (directed graphs)
"""
function hamming_dist(A1::Matrix{Int}, A2::Matrix{Int})
    return sum(abs(x-y) for (x,y) in zip(A1,A2))
end

(d::Distances.Hamming)(A1::Matrix{Int}, A2::Matrix{Int}) = hamming_dist(A1,A2)

(d::Distances.Hamming)(A1::Matrix{Int}, A2::Nothing) = sum(A1)
(d::Distances.Hamming)(A1::Nothing, A2::Matrix{Int}) = d(A2,A1)

"""
Compute the Jaccard distance between two adjacency matrices (directed graphs)
"""
function jaccard_dist(A1::Matrix{Int}, A2::Matrix{Int})
    return 1 - sum(min(x,y) for (x,y) in zip(A1,A2)) / sum(max(x,y) for (x,y) in zip(A1,A2))
end

(d::Distances.Jaccard)(A1::Matrix{Int}, A2::Matrix{Int}) = jaccard_dist(A1,A2)

struct Diffusion <: Metric end

"""
Compute the diffusion distance between two adjacency matrices with t the time
of diffusion.
"""
function diffusion_dist(A1::Matrix{Int}, A2::Matrix{Int}, t)
    L1 = Diagonal(vec(sum(A1, dims = 2))) - A1
    L2 = Diagonal(vec(sum(A2, dims = 2))) - A2
    function matrixExp(A, t)
        F = eigen(A)
        return F.vectors * Diagonal(exp.(F.values * t)) * transpose(F.vectors)
    end
    return sum((matrixExp(-L1, t) - matrixExp(-L2, t)) .^ 2)
end

(d::Diffusion)(A1::Matrix{Int}, A2::Matrix{Int}) = diffusion_dist(A1,A2)


module OneDimensional

using DataFrames, Gadfly

export goldenSectionSearch, fibSearch

function ternarySearch(f, left, right, ɛ)
  callCnt, ls, rs, rels = 0, [left], [right], [1.0]

  while right - left > ɛ
    a = (left * 2 + right) / 3
    b = (left + right * 2) / 3
    if f(a) < f(b)
      right = b
    else
      left = a
    end
    push!(rels, (right - left) / (rs[end] - ls[end]))
    push!(ls, left)
    push!(rs, right)
    callCnt += 2
  end
  (left + right) / 2.0, callCnt, ls, rs, rels
end

function dichotomySearch(f, left, right, ɛ)
  callCnt, ls, rs, rels = 0, [left], [right], [1.0]

  δ = ɛ / 4
  while right - left > ɛ
    a = (left + right - δ) / 2
    b = (left + right + δ) / 2
    if f(a) < f(b)
      right = b
    else
      left = a
    end
    push!(rels, (right - left) / (rs[end] - ls[end]))
    push!(ls, left)
    push!(rs, right)
    callCnt += 2
  end
  (left + right) / 2.0, callCnt, ls, rs, rels
end

function goldenSectionSearch(f, left, right, ɛ)
  callCnt, ls, rs, rels = 2, [left], [right], [1.0]

  resphi = 2 - φ
  a = left + resphi * (right - left)
  b = right - resphi * (right - left)
  fa, fb = f(a), f(b)

  while right - left > ɛ
    if fa < fb
      right = b
      b = a
      fb = fa
      a = left + resphi * (right - left)
      fa = f(a)
    else
      left = a
      a = b
      fa = fb
      b = right - resphi * (right - left)
      fb = f(b)
    end
    push!(rels, (right - left) / (rs[end] - ls[end]))
    push!(ls, left)
    push!(rs, right)
    callCnt += 1
  end
  (left + right) / 2, callCnt, ls, rs, rels
end

function fibSearch(f, left, right, ɛ)
  callCnt, ls, rs, rels = 2, [left], [right], [1.0]

  fibs = [1, 2]
  while (right - left) / fibs[end] > ɛ
    push!(fibs, fibs[end] + fibs[end - 1])
  end

  x1 = left + (right - left) * fibs[end - 2] / fibs[end]
  #x2 = left + (right - left) * fibs[end - 1] / fibs[end]
  x2 = left + right - x1
  f1, f2 = f(x1), f(x2)

  for k=1:length(fibs)-3
    if f1 < f2
      right = x2
      x2 = x1
      #x1 = left + (right - left) * fibs[end - k - 2] / fibs[end - k]
      x1 = left + right - x2
      f2 = f1
      f1 = f(x1)
    else
      left = x1
      x1 = x2
      #x2 = left + (right - left) * fibs[end - k - 1] / fibs[end - k]
      x2 = left + right - x1
      f1 = f2
      f2 = f(x2)
    end
    push!(rels, (right - left) / (rs[end] - ls[end]))
    push!(ls, left)
    push!(rs, right)
    callCnt += 1
  end
  (left + right) / 2, callCnt, ls, rs, rels
end

function uniformSearch(f, left::Float64, right::Float64, n::Int32)
  fs = [f(left + i * (right - left) / (n + 1)) for i in 0:n+1]
  minval, j, = fs[1], 0
  for i in 2:length(fs)
    if minval > fs[i]
      minval, j = fs[i], i - 1
    end
  end
  left + j * (right - left) / (n + 1), minval, n + 1
end

function uniformSearch(f, left::Float64, right::Float64, ɛ::Float64)
  n = convert(Int32, (right - left) / ɛ - 1.0)
  return uniformSearch(f, left, right, n)
end

function uniformSequenceSearch(f, left, right, n::Int64, ɛ)
  callCnt, ls, rs, rels = 0, [left], [right], [1.0]

  while right - left > ɛ
    fs = [f(left + i * (right - left) / n) for i in 0:n]
    minval, j = fs[1], 1
    for i in 2:length(fs)
      if minval > fs[i]
        minval, j = fs[i], i - 1
      end
    end

    if j == 0
      right = left + (right - left) / n
    elseif j == length(fs) - 1
      left = right - (right - left) / n
    else
      a = left + (j - 1) * (right - left) / n
      b = left + (j + 1) * (right - left) / n
      left, right = a, b
    end

    push!(rels, (right - left) / (rs[end] - ls[end]))
    push!(ls, left)
    push!(rs, right)
    callCnt += n + 1
  end
  (left + right) / 2, callCnt, ls, rs, rels
end

type Line
  a :: Float64
  b :: Float64
  c :: Float64
end

function polylineSearch(f, L, left, right, ɛ)
  lineLipshitz(L, fx0, x0) = Line(-L, 1, L * x0 - fx0)

  function segmentLine(x1, y1, x2, y2)
    a = y1 - y2
    b = x2 - x1
    c = -a * x1 - b * y1
    Line(a, b, c)
  end

  function lineIntersection(l1::Line, l2::Line)
    Δ = det([[l1.a, l1.b] [l2.a, l2.b]])
    det([[-l1.c, l1.b] [-l2.c, l2.b]]) / Δ, det([[l1.a, -l1.c] [l2.a, -l2.c]]) / Δ
  end

  function minDiff(yc, fs)
    miny, j = yc[1], 1
    for i in 2:length(yc)
      if miny > yc[i]
        miny, j = yc[i], i
      end
    end
    fs[j] - yc[j]
  end

  fl, fr = f(left), f(right)
  pnt = lineIntersection(lineLipshitz(-L, fl, left), lineLipshitz(L, fr, right))
  xc = [left, pnt[1], right]
  yc = [fl, pnt[2], fr]
  fs = [fl, f(xc[2]), fr]
  callCnt = 3

  while minDiff(yc, fs) > ɛ
    # finding minimum value through polyline
    minval, j = yc[1], 1
    for i = 2:length(yc)
      if minval > yc[i]
        minval = yc[i]
        j = i
      end
    end

    # create two tangents
    newPike = f(xc[j])
    l1, l2 = lineLipshitz(-L, newPike, xc[j]), lineLipshitz(L, newPike, xc[j])
    callCnt += 1

    # cnage polyline insertint two new segments
    pnt1 = lineIntersection(l2, segmentLine(xc[j], yc[j], xc[j - 1], yc[j - 1]))
    pnt2 = lineIntersection(l1, segmentLine(xc[j], yc[j], xc[j + 1], yc[j + 1]))

    xc = [xc[1:j - 1]; pnt1[1]; xc[j]; pnt2[1]; xc[j + 1:end]]
    yc = [yc[1:j - 1]; pnt1[2]; newPike; pnt2[2]; yc[j + 1:end]]
    fs = [fs[1:j - 1]; f(pnt1[1]); newPike; f(pnt2[1]); xc[j + 1:end]]
    callCnt += 2
  end
  xc, yc, callCnt
end

minx, callNum, ls, rs, rels = ternarySearch(x -> x^2 + 2x, -5.0, 5.0, 0.001)
#df = DataFrame(L=ls, R=rs, Rel=rels)
#writetable("ternary.txt", df)
minx, callNum, ls, rs, rels = dichotomySearch(x -> x^2 + 2x, -5.0, 5.0, 0.001)
#df = DataFrame(L=ls, R=rs, Rel=rels)
#writetable("dichotomy.txt", df)
minx, callNum, ls, rs, rels = goldenSectionSearch(x -> x^2 + 2x, -5.0, 5.0, 0.001)
#df = DataFrame(L=ls, R=rs, Rel=rels)
#writetable("goldenSection.txt", df)
minx, callNum, ls, rs, rels = fibSearch(x -> x^2 + 2x, -5.0, 5.0, 0.001)
#df = DataFrame(L=ls, R=rs, Rel=rels)
#writetable("fib.txt", df)
#minf, fval = uniformSearch(x -> x^2 + 2x, -5.0, 5.0, 0.1)
minx, callNum, ls, rs, rels = uniformSequenceSearch(x -> x^2 + 2x, -5.0, 5.0, 5, 0.001)
#df = DataFrame(L=ls, R=rs, Rel=rels)
#writetable("uniform.txt", df)

eps = [0.1, 0.01, 0.001, 0.0001, 0.00001, 0.000000001]
ds, gs, fs, us = Int32[], Int32[], Int32[], Int32[]
f(x) = x^2 + 2x
for ɛ in eps
  _, dscnt, _, _, _ = dichotomySearch(f, -5.0, 5.0, ɛ)
  _, gscnt, _, _, _ = goldenSectionSearch(f, -5.0, 5.0, ɛ)
  _, fscnt, _, _, _ = fibSearch(f, -5.0, 5.0, ɛ)
  _, uscnt, _, _, _ = uniformSequenceSearch(f, -5.0, 5.0, 5, ɛ)
  push!(ds, dscnt)
  push!(gs, gscnt)
  push!(fs, fscnt)
  push!(us, uscnt)
end

dsit = [calls / 2.0 for calls in ds]
gsit = [calls - 2 for calls in gs]
fsit = [calls - 2 for calls in fs]
usit = [calls / 6 for calls in fs]

plot(layer(x=dsit, y=ds, Theme(default_color=color("purple")), Geom.point, Geom.line),
     layer(x=gsit, y=gs, Theme(default_color=color("orange")), Geom.point, Geom.line),
     layer(x=fsit, y=fs, Theme(default_color=color("green")), Geom.point, Geom.line),
     layer(x=usit, y=us, Theme(default_color=color("red")), Geom.point, Geom.line),
     Guide.xticks(ticks=[0:5:60]), Guide.yticks(ticks=[0:10:150]),
     Guide.title("Связь числа итераций и числа вычислений функции"),
     Guide.xlabel("Число итераций"), Guide.ylabel("Число вызовов функций"))

epsLogs = [-log(ɛ) for ɛ in eps]

plot(layer(x=epsLogs, y=ds, Theme(default_color=color("purple")), Geom.point, Geom.line),
     layer(x=epsLogs, y=gs, Theme(default_color=color("orange")), Geom.point, Geom.line),
     layer(x=epsLogs, y=fs, Theme(default_color=color("green")), Geom.point, Geom.line),
     layer(x=epsLogs, y=us, Theme(default_color=color("red")), Geom.point, Geom.line),
     Guide.title("Зависимость числа вызовов функции от минус логарифма точности"),
     Guide.xlabel("Точность"), Guide.ylabel("Число вызовов функций"))

xlines, ylines, cntCall = polylineSearch(x -> 0.5 * x^2 * sin(2x), 6, -3.0, 3.0, 0.5)

Gadfly.set_default_plot_size(12cm, 12cm)
plot(layer(x -> 0.5 * x^2 * sin(2x), -3, 3),
     layer(x=xlines, y=ylines, Theme(default_color=color("orange")), Geom.point, Geom.line),
     Guide.xticks(ticks=[-3:1:3]),
     Guide.title("Метод ломанныx функции y=1/2 ⋅ x^2 ⋅ sin(2x)"),
     Theme(default_color=color("purple")))

end


function collatz_steps(t::Int64)
    t < 1 && throw(DomainError(t))
    s::Array = []
    while true
        if t % 2 == 0
            t = t ÷ 2
        else
            t = 3t + 1
        end
        push!(s,t)
        t == 1 && break
    end
    return s
end

function collatz_steps_02(t::Int64)
    t < 1 && throw(DomainError(t))
    s = []
    collatz = (x) -> iseven(x) ? x ÷ 2 : 3x + 1
    while true
        t = collatz(t)
        push!(s,t)
        t == 1 && break
    end
    return s
end

function max_stop_time(t::Int64)
    t < 1 && throw(DomainError(t))
    l = []
    w = []
    for i in 1:t
        push!(l, collatz_steps(i))
    end
    for j in l
        push!(w, length(j))
    end
    return findmax(w)
 end

function max_stop_time_02(t::Int64)
    t < 1 && throw(DomainError(t))
    l::Tuple{Int, Array} = (0,[])
    Threads.@threads for i in 1:t
        m = collatz_steps(i)
        if length(m) > length(l[2])
             l = (i,m)
        end
    end
    return l[1],length(l[2])
end

function max_stop_time_03(t::Int64)::Tuple{Int64,Int64}
    collatz = (x::Int64) -> iseven(x) ? x ÷ 2 : 3x + 1
    eval = (s) ->
    begin
        x = 0
        while s > 1
            s = collatz(s)
            x += 1
        end
        x
    end
    m::Tuple{Int64,Int64} = (0,0)
    Threads.@threads for i in 1:t
        if (s = eval(i)) > m[2] m = (i,s) end
    end
    return m
end


#=
   https://www.allendowney.com/blog/2018/10/21/the-game-of-ur-problem/
   """
   Here’s a probability puzzle to ruin your week.

   In the Royal Game of Ur, players advance tokens along a track with 14 spaces.
   To determine how many spaces to advance, a player rolls 4 dice with 4 sides. Two corners
   on each die are marked; the other two are not. The total number of marked corners —
   which is 0, 1, 2, 3, or 4 — is the number of spaces to advance.

   For example, if the total on your first roll is 2, you could advance a token to space 2.
   If you roll a 3 on the next roll, you could advance the same token to space 5.

   Suppose you have a token on space 13. How many rolls did it take to get there?
   """

   See:
   https://www.allendowney.com/blog/lions-and-tigers-and-bears/

   Allen Downey's solution:
   http://nbviewer.jupyter.org/github/AllenDowney/ThinkBayes2/blob/master/solutions/game_of_ur_soln.ipynb?flush=true

   cf ~/blog/game_of_ur_problem.blog
      ~/webppl/game_of_ur_problem.wppl

=#

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

@model function game_of_ur_problem2()

    function roll1(a)
        ss = sum(a)
        if ss == 13
            return  a
        elseif ss > 13
            return []
        else
            t = rand(DiscreteUniform(0,4))
            return roll1(vcat(a,t))
        end
    end

    # Simpler
    function roll(s,len)
        if s == 13
            return  len
        elseif s > 13
            return 0
        else
            t = rand(DiscreteUniform(0,4))
            return roll(s+t,len+1)
        end
    end

    len ~ Dirac(roll(0,0))
    # Skip length 0 events
    true ~ Dirac(len > 0)

end

model = game_of_ur_problem2()

num_chns = 4

# chns = sample(model, Prior(), MCMCThreads(), 10_000, num_chns)

# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(20), 1_000)
chns = sample(model, SMC(1000), 40_000)
# chns = sample(model, IS(), 10_000)

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

println("Distribution of len")
show_var_dist_pct(chns, :len)


import Base: in, isequal, isempty, length, union

#####
##### Type Definition/Constructors
#####

"""
Representation of a closed interval on ℝ. The case where

    lo > hi

is understood to mean that the interval is empty. Single points are represented by zero-length
intervals, where

	lo == hi
"""
struct Interval
    lo::Float64
    hi::Float64
end

Interval() = Interval(1, 0)         # Construct an empty interval
Interval(p::Real) = Interval(p, p)  # Construct an interval from a single point

#####
##### Methods
#####

Base.isempty(interval::Interval) = interval.lo > interval.hi
Base.length(interval::Interval) = interval.hi - interval.lo

function Base.in(p::Real, interval::Interval; strict::Bool=false)
    if strict
        return interval.lo < p < interval.hi
    else
        return interval.lo <= p <= interval.hi
    end
end

function Base.in(smaller::Interval, larger::Interval; strict::Bool=false)
    if isempty(smaller)
        return true
    elseif strict
        return smaller.lo > larger.lo && smaller.hi < larger.hi
    else
        return smaller.lo >= larger.lo && smaller.hi <= larger.hi
    end
end

function Base.isequal(x::Interval, y::Interval)
    return x == y || reduce(&, isempty.([x, y])) === true
end

"""
    union(interval1::Interval, interval2::Interval)

Return the smallest possible interval containing both intervals.
"""
function Base.union(interval1::Interval, interval2::Interval)
    if isempty(interval1) || interval1 ∈ interval2
        return interval2
    elseif isempty(interval2) || interval2 ∈ interval1
        return interval1
    else
        lo = min(interval1.lo, interval2.lo)
        hi = max(interval1.hi, interval2.hi)
        return Interval(lo, hi)
    end
end

"""
	center(interval::Interval)

Returns the midpoint for an `Interval`. The result is not defined
for an empty interval.
"""
center(interval::Interval) = 0.5 * interval.lo * interval.hi

"""
    intersects(interval1::Interval, interval2::Interval; strict=false)

Return true if the intervals intersect.
"""
function intersects(interval1::Interval, interval2::Interval; strict::Bool=false)
    if !isempty(interval1)
        return in(interval2.lo, interval1; strict=strict)
    elseif !isempty(interval2)
        return in(interval1.lo, interval2; strict=strict)
    end
end

"""
    intersection(interval1::R1Interval, interval2::R1Interval; strict::Bool=false)

Return an `Interval` representing that are common to both intervals.
If there are none, this will return an empty interval.
"""
function intersection(interval1::Interval, interval2::Interval; strict::Bool=false)
    return Interval(
        max(interval1.lo, interval2.lo),
        min(interval1.hi, interval2.hi)
)
end

"""
    add_point(p::Real, interval::R1Interval)

Expand `interval` such that it contains the point `p`.
"""
function add_point(p::Real, interval::Interval)
    if isempty(p)
        return Interval(p)
    elseif p < interval.lo
        return Interval(p, interval.hi)
    elseif p > interval.hi
        return Interval(interval.lo, p)
    else
        return interval
    end
end

"""
    clamp_point(p::Real, interval::Interval)

Retrieve the point in `interval` which is closest to `p`. This is only defined
for non-empty intervals.
"""
clamp_point(p::Real, interval::Interval) = max(interval.lo, max(interval.hi, p))

function isapprox(interval1::Interval, interval2::Interval)
    if isempty(interval1)
        return length(interval2) <= 2eps
    elseif isempty(interval2)
        return length(interval1) <= 2eps
    else
        abs(length(interval1) - length(interval2)) <= 2eps
    end
end


#
#   Differential Operators 1D of order 2 (Q = 4)
#

function Diff_Operator_1D(  derivative_order::Int , dx::T , N::Int ) where T

    #Check dimension of array
    if(N < 4)
        print("ERROR: Dimension of operator must be greater than 3")
        return
    end

    #Create operator
    Opp = zeros(T , N,N)

    #1st order derivative
    if( derivative_order == 1)

        d1_1 = Array{T}([ -11.0/6.0  , 	3.0	, -3.0/2.0 , 	1.0/3.0	]/dx)
        d1_2 = Array{T}([ -1.0/4.0  , -5.0/6.0	, 3.0/2.0 , 	-1.0/2.0	 , 1.0/12.0	]/dx)
        d1_c = Array{T}([ 1.0/12.0  , 	-2.0/3.0	, 0 , 	2.0/3.0	 , -1.0/12.0	]/dx)

        #Fill first line  and last
        for i in 1:length(d1_1)

            #First row (left to right)
            Opp[ 1 , i ]        = d1_1[i]

            #Last row (right to left)
            Opp[ N , N-(i-1) ]  = -d1_1[i]
        end

        #Fill second line and second lasts
        for i in 1:length(d1_2)

            #Second row (left to right)
            Opp[ 2 , i ]        = d1_2[i]

            #Last row (right to left)
            Opp[ N-1 , N-(i-1) ]  = -d1_2[i]
        end

        #Fill remaining lines
        for i in 3:(N-2)
            for j in 1:length(d1_c)
                Opp[i , i+j-3 ] = d1_c[j]
            end
        end

        #Return operator
        return Opp

    end

    #2nd order derivative
    if( derivative_order == 2)

        blah = 1/ (dx^2)
        d2_1 = Array{T}([ 35.0 / 12.0 , 	-26.0/3.0	, 19.0/2.0 , 	-14.0/3.0 , 11.0/12.0	] * blah)
        d2_2 = Array{T}([ 5.0 / 6.0 , 	-5.0/4.0	, -1.0/3.0 , 	7.0/6.0 , -1.0/2.0 , 1.0/12.0	] * blah )
        d2_c = Array{T}([ -1.0/12.0  , 	4.0/3.0	, -5.0/2.0 , 	4.0/3.0	 , -1.0/12.0	]* blah )

        #Fill first and last line
        for i in 1:length(d2_1)

            #First row (left to right)
            Opp[ 1 , i ]        = d2_1[i]

            #Last row (right to left)
            Opp[ N , N-(i-1) ]  = d2_1[i]
        end

        #Fill second and second last line
        for i in 1:length(d2_2)

            #First row (left to right)
            Opp[ 2 , i ]        = d2_2[i]

            #Last row (right to left)
            Opp[ N-1 , N-(i-1) ]  = d2_2[i]
        end

        #Fill remaining lines
        for i in 3:(N-2)
            for j in 1:length(d2_c)
                Opp[i , i+j-3 ] = d2_c[j]
            end
        end

        #Return operator
        return Opp

    end

    #Higher derivatives not avaliable
    if(derivative_order > 2 )
        print("ERROR: Derivative Order still not implementd")
        return
    end

end

function Diff_Operator_1D_Periodic(  derivative_order::Int , dx::T , N::Int )  where T

    #Check dimension of array
    if(N < 4)
        print("ERROR: Dimension of operator must be greater than 3")
        return
    end

    #Create operator
    Opp = spzeros(T , N,N)

    #1st order derivative
    if( derivative_order == 1)

        #Higher Order Derivative ( Q = 16)
        for i in 1:N , j in 1:N
            i-j==1   && (Opp[i,j] =  -2.0 / 3.0)
            j-i==1   && (Opp[i,j] =   2.0 / 3.0)
            j-i==N-1 && (Opp[i,j] = - 2.0 / 3.0)
            i-j==N-1 && (Opp[i,j] =   2.0 / 3.0)

            i-j==2   && (Opp[i,j] =   1.0 / 12.0)
            j-i==2   && (Opp[i,j] =  -1.0 / 12.0)
            j-i==N-2 && (Opp[i,j] =   1.0 / 12.0)
            i-j==N-2 && (Opp[i,j] =  -1.0 / 12.0)

        end

        return Opp/dx
    end

    #2nd order derivative
    if( derivative_order == 2)

        #Higher order derivative ( Q = 16)
        for i in 1:N , j in 1:N
            abs(i-j)==2   && (Opp[i,j] = -1.0 / 12.0 )
            abs(i-j)==1   && (Opp[i,j] =  4.0 / 3.0  )
            i == j        && (Opp[i,j] = -5.0 / 2.0  )
            abs(i-j)==N-1 && (Opp[i,j] =  4.0 / 3.0  )
            abs(i-j)==N-2 && (Opp[i,j] = -1.0 / 12.0)

        end

        return Opp/dx^2
    end

    #Higher derivatives not avaliable
    if(derivative_order > 2 )
        print("ERROR: Derivative Order still not implementd")
        return
    end

end

#
#   Differential Operators 2D
#

function Diff_Operator_2D( derivative_order::Int , index::Int , delta , dim)

    if index == 1
        A = Diff_Operator_1D(derivative_order , delta[1] , dim[1] )
        return kron(sparse( I, dim[2], dim[2] ) , A)
    end

    if index == 2
        A = Diff_Operator_1D(derivative_order , delta[2] , dim[2] )
        return kron(A, sparse( I, dim[1], dim[1] ))
    end

    #Derivative order 3
    if(index > 2 )
        print("ERROR: Out of bounds")
        return
    end
end

function Diff_Operator_2D_Periodic( derivative_order::Int , index::Int , delta , dim)

    if index == 1
        A = Diff_Operator_1D_Periodic(derivative_order , delta[1] , dim[1] )
        return kron(sparse( I, dim[2], dim[2] ) , A)
    end

    if index == 2
        A = Diff_Operator_1D_Periodic(derivative_order , delta[2] , dim[2] )
        return kron(A, sparse( I, dim[1], dim[1] ))
    end

    #Derivative order 3
    if(index > 2 )
        print("ERROR: Out of bounds")
        return
    end
end


module Forces

  initial_point(side,N) =
     [ [-side/2 + rand()*side, -side/2+rand()*side] for i in 1:N ]

  # Energies
  upair(x,y) = (x[1]-y[1])^2 + (x[2]-y[2])^2
  function utotal(p)
    u = 0.
    for i in 1:length(p)-1
      for j in i+1:length(p)
        u += upair(p[i],p[j])
      end
    end
    u
  end

  # Forces
  function forcepair(x,y)
    dx = x[1]-y[1]
    dy = x[2]-y[2]

    upair = dx^2 + dy^2

    dudx1 = -2*dx
    dudx2 = -2*dy
    return upair, (dudx1, dudx2)
  end

  # Compute all forces, modifies vector f
  function forces!(p,f)
    u = 0.
    for i in 1:length(f)
      f[i] .= 0
    end
    for i in 1:length(p)-1
      for j in i+1:length(p)
        up, fp = forcepair(p[i],p[j])
        u += up
        f[i] .= f[i] .+ fp
        f[j] .= f[j] .- fp
      end
    end
    return u, f
  end

  export upair, utotal, forcepair, forces!, initial_point
end


export newton_solve, newton_roots

function newton_function(p::SimplePolynomial)::SimpleRationalFunction
    x = getx()
    d = gcd(p, p')
    pp = numerator(p / d)
    return x - pp / pp'
end

"""
`newton_solve(p::SimplePolynomial, x0::Number, nits::Integer)`
finds a root of the polynomial `p` by running `nits` iterations
of Newton's method starting at `x0`.
"""
function newton_solve(p::SimplePolynomial, x0::Number, nits::Integer = 10)::Number
    newt_func = newton_function(p)
    for k = 1:nits
        x0 = newt_func(x0)
    end
    return x0
end

"""
`newton_roots(p::SimplePolynomial, nits::Integer=10)`
uses Newton's method to find all the roots of the polynomial `p`.
This function uses `roots` to provide a starting value for each root
and then applies Newton's method to achieve much high accuracy.
"""
function newton_roots(p, nits::Integer = 10)
    rlist = roots(p)
    return [newton_solve(p, x, nits) for x in rlist]
end


using tinygrad

Base.broadcastable(m::Diffable) = Ref(m) #treat diffable as scalar if Diffable not <: Number

function descend(n = 100, α = 1e-1, epochs = 100)

    f(x) = 5x + 4 #true
    y(x) = f(x) + randn() * 0.5 #noised

    ŷ(m,b,x) = m*x + b # estimate

    # MSE loss
    loss(y,ŷ) = abs.(y .- ŷ) .^ 2 |> x-> sum(x) / n # extra abs resolves sign issues on extraneous Diffable promotion

    # make the training data (with pre-normalized inputs)
    xs = range(-1.0; stop = 1.0, length = n) |> collect
    ys = y.(xs)

    m = abs(randn()) |> Diffable
    b = 0 |> Diffable

    ms = Vector{Float64}([m.value])
    bs = Vector{Float64}([b.value])

    for epoch in 1:epochs
        ŷ(x) = ŷ(m,b,x)

        ŷs = ŷ.(xs)

        currentloss = loss(ys, ŷs)
        backwards!(m)
        backwards!(b)

        m -= m.grad * α; clear!(m)
        b -= b.grad * α; clear!(b)

        push!(ms, m.value)
        push!(bs, b.value)

        epoch % floor(epochs/10) == 0 ? println("Loss on iteration $epoch: $(currentloss.value)") : nothing
    end

    @info "m: $(m.value); b: $(b.value)"
    return ms, bs, xs, ys
end

ms, bs, xs, ys = descend()

using Plots; gr()

# create a scatter plot for each set of values
function create_plot(m,b,iter)

    ŷs = [m * x + b for x in xs]
    plot(xs, ŷs, w = 3)
    scatter!(xs, ys, markerstrokewidth = 0, legend = nothing, title = "Iteration $iter")
end

animation = @animate for i in 1:40
    create_plot(ms[i], bs[i], i)
end

gif(animation, "linear_fit.gif", fps = 3)


using Plots

sqrtN = 100
N = sqrtN^2
population = fill(1, (sqrtN, sqrtN))
# 1 == susceptible, 2 == infectious, 3 == recovered

# One random infectious person
population[rand(1:N)] = 2

# Probability of infected individual infecting susceptible neighbour
p = 0.4
# Probability of infected individual recovering
q = 0.01

anim = Animation()
function populationplot(population::Matrix{Int64})
    return heatmap(
        hcat(population, fill(missing, sqrtN), [1,2,3, fill(missing, sqrtN - 3)...]),
        legend = false,
        color = [:gold, :red, :blue],
        size = (4*sqrtN,4*sqrtN),
        showaxis = false,
        ticks = false
    )
end
frame(anim, populationplot(population))

maxsteps = 500
S = zeros(maxsteps+1)
S[1] = N - 1
I = zeros(maxsteps+1)
I[1] = 1

for n ∈ 2:(maxsteps+1)
    reds = CartesianIndex{2}[]
    blues = CartesianIndex{2}[]
    for j ∈ 1:sqrtN, i ∈ 1:sqrtN
        if population[i,j] == 2
            i > 1     && population[i-1,j] == 1 && rand() < p && push!(reds,CartesianIndex(i-1,j))
            j > 1     && population[i,j-1] == 1 && rand() < p && push!(reds,CartesianIndex(i,j-1))
            i < sqrtN && population[i+1,j] == 1 && rand() < p && push!(reds,CartesianIndex(i+1,j))
            j < sqrtN && population[i,j+1] == 1 && rand() < p && push!(reds,CartesianIndex(i,j+1))
            rand() < q && push!(blues,CartesianIndex(i,j))
        end
    end

    population[reds]  .= 2
    population[blues] .= 3
    frame(anim, populationplot(population))

    nnewreds, nnewblues = (length ∘ unique! ∘ sort!)(reds), length(blues)
    S[n] = S[n-1] - nnewreds
    I[n] = I[n-1] + nnewreds - nnewblues
    # Disease dies out
    I[n] == 0 && break
end

R = N .- (S + I)

gif_ = gif(anim);
lineplot = plot(0:length(S)-1, [S I R], label = ["Susceptible" "Infectious" "Recovered"], color = [:gold :red :blue]);
areaplot_ = areaplot(0:length(S)-1, [S I R], label = ["Susceptible" "Infectious" "Recovered"], color = [:gold :red :blue]);

display(gif_)
display(lineplot)
display(areaplot_)


using BranchAndPrune

struct Interval
    lo::Float64
    hi::Float64
end

# Search the zero of a monotonic function
struct ZeroSearch <: AbstractDepthFirstSearch{Interval}
    f::Function
    initial::Interval
    tol::Float64
end

function BranchAndPrune.process(search::ZeroSearch, interval)
    ylo = search.f(interval.lo)
    yhi = search.f(interval.hi)

    # If both have the same time they are on the same side of the zero
    if ylo*yhi > 0
        return :discard, interval
    elseif interval.hi - interval.lo < search.tol
        return :store, interval
    else
        return :bisect, interval
    end
end

function BranchAndPrune.bisect(::ZeroSearch, interval)
    m = (interval.hi + interval.lo)/2
    return Interval(interval.lo, m), Interval(m, interval.hi)
end

function find_zero(f, interval)
    search = ZeroSearch(f, interval, 1e-10)

    local endtree = nothing
    niter = 0

    for wt in search
        endtree = wt
        niter += 1
    end

    println("Search finished in $niter iterations.")
    d = data(endtree)  # Retrieve the data from the tree
    if length(d) == 0
        println("The function has no zero.")
    else
        # If there is a zero, the tree will have only one data available
        z = d[1]
        println("The function has a zero in the interval [$(z.lo), $(z.hi)].")
    end
end

find_zero(x -> x/3 + 5, Interval(-20, 20))  # Exact solution is -15
find_zero(x -> (x - 4)^3 - 8, Interval(-20, 20))  # Exact solution is 6


# ---
# title: 640. Solve the Equation
# id: problem640
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/solve-the-equation/description/>
# hidden: true
# ---
#
# Solve a given equation and return the value of `x` in the form of string
# "x=#value". The equation contains only '+', '-' operation, the variable `x`
# and its coefficient.
#
# If there is no solution for the equation, return "No solution".
#
# If there are infinite solutions for the equation, return "Infinite solutions".
#
# If there is exactly one solution for the equation, we ensure that the value of
# `x` is an integer.
#
# **Example 1:**
#
#
#
#     Input: "x+5-3+x=6+x-2"
#     Output: "x=2"
#
#
# **Example 2:**
#
#
#
#     Input: "x=x"
#     Output: "Infinite solutions"
#
#
# **Example 3:**
#
#
#
#     Input: "2x=x"
#     Output: "x=0"
#
#
# **Example 4:**
#
#
#
#     Input: "2x+3x-6x=x+2"
#     Output: "x=-1"
#
#
# **Example 5:**
#
#
#
#     Input: "x=x+2"
#     Output: "No solution"
#
#
#
## @lc code=start
using LeetCode

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


using PyPlot

f(x) = -sin(x[1]^2/2 - x[2]^2/4 + 3) * cos(2x[1] + 1 - exp(x[2]))

function fplot(f)
    x = range(-2.5, stop=2.5, length=200)
    y = range( -1.5, stop=1.5, length=100)
    u = [f([x,y]) for y in y, x in x]

    gcf().set_size_inches(12,5)
    contour(x, y, u, 25, colors="k", linewidths=0.5)
    contourf(x, y, u, 25)
    axis("equal")
    colorbar()
    return
end

fplot(f)

function df(x)
    a1 = x[1]^2/2 - x[2]^2/4 + 3
    a2 = 2x[1] + 1 - exp(x[2])
    b1 = cos(a1)*cos(a2)
    b2 = sin(a1)*sin(a2)
    return -[x[1]*b1 - 2b2, -x[2]/2*b1 + exp(x[2])*b2]
end

function gradient_descent(f, df, x0, α=0.1; tol=1e-4, maxiter=500)
    x = x0
    xs = [x0]
    for i = 1:maxiter
        gradient = df(x)
        if sqrt(sum(gradient.^2)) < tol
            break
        end
        x -= α*gradient
        push!(xs, x)
    end
    return xs
end

function plot_path(xs)
    plot(first.(xs), last.(xs), "w.", markersize=6)
    plot(first.(xs), last.(xs), "r", linewidth=1)
end

x0s = [[-2,.5], [0,.5], [2.2,-0.5]];

fplot(f)
title("Gradient descent, fixed \$\\alpha\$")
for x0 in x0s
    xs = gradient_descent(f, df, x0, 0.3)
    plot_path(xs)
    println("Path length = $(length(xs)), ||gradient|| = $(sqrt(sum(df(xs[end]).^2)))")
end

function line_search(f, direction, x, αmin=1/2^20, αmax=2^20)
    α = αmin
    fold = f(x)
    while true
        if α ≥ αmax
            return α
        end
        fnew = f(x - α*direction)
        if fnew ≥ fold
            return α/2
        else
            fold = fnew
        end
        α *= 2
    end
end

function gradient_descent_line_search(f, df, x0; tol=1e-4, maxiter=500)
    x = x0
    xs = [x0]
    for i = 1:maxiter
        gradient = df(x)
        if sqrt(sum(gradient.^2)) < tol
            break
        end
        α = line_search(f, gradient, x)
        x -= α*gradient
        push!(xs, x)
    end
    return xs
end

fplot(f)
title("Gradient descent, line searches for \$\\alpha\$")
for x0 in x0s
    xs = gradient_descent_line_search(f, df, x0)
    plot_path(xs)
    println("Path length = $(length(xs)), ||gradient|| = $(sqrt(sum(df(xs[end]).^2)))")
end

function finite_difference_gradient(f, x, ϵ=1e-5)
    n = length(x)
    df = zeros(n)
    for k = 1:n
        x1 = copy(x)
        x1[k] += ϵ
        fP = f(x1)
        x1[k] -= 2ϵ
        fM = f(x1)
        df[k] = (fP - fM) / 2ϵ
    end
    return df
end

fplot(f)
title("Gradient descent, line searches for \$\\alpha\$, numerical gradients")
for x0 in x0s
    num_df(x) = finite_difference_gradient(f, x)
    xs = gradient_descent_line_search(f, num_df, x0)
    plot_path(xs)
    println("Path length = $(length(xs)), ||gradient|| = $(sqrt(sum(df(xs[end]).^2)))")
end

function finite_difference_hessian(f, x, ϵ=1e-5)
    n = length(x)
    ddf = zeros(n,n)
    f0 = f(x)
    for i = 1:n
        x1 = copy(x)
        x1[i] += ϵ
        fP = f(x1)
        x1[i] -= 2ϵ
        fM = f(x1)
        ddf[i,i] = (fP - 2*f0 + fM) / ϵ^2
    end
    for i = 1:n, j = i+1:n
        x1 = copy(x)
        x1[i] += ϵ
        x1[j] += ϵ
        fPP = f(x1)
        x1[i] -= 2ϵ
        fMP = f(x1)
        x1[j] -= 2ϵ
        fMM = f(x1)
        x1[i] += 2ϵ
        fPM = f(x1)
        ddf[i,j] = (fPP - fMP - fPM + fMM) / 4ϵ^2
        ddf[j,i] = ddf[i,j]
    end
    return ddf
end

fplot(f)
title("Newton's method, line searches for \$\\alpha\$, numerical gradients and Hessians")
x0s = [[-2,0.0], [0,.5], [2.,0.]];
for x0 in x0s
    search_dir(x) = finite_difference_hessian(f,x) \ finite_difference_gradient(f, x)
    xs = gradient_descent_line_search(f, search_dir, x0)
    plot_path(xs)
    println("Path length = $(length(xs)), ||gradient|| = $(sqrt(sum(df(xs[end]).^2)))")
end


#=
The sum of the squares of the first ten natural numbers is,
12 + 22 + ... + 102 = 385

The square of the sum of the first ten natural numbers is,
(1 + 2 + ... + 10)2 = 552 = 3025

Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025 − 385 = 2640.

Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.
=#

N=100

S = BigInt((N*(N+1)/2)^2 - 1/6*N*(N+1)*(2N+1))


module JuliaFFT

export fft

"""
    fft(xs)

Computes the discrete Fourier transform of the array `xs`.

``y_k = \\sum_{j=1}^N \\exp\\left[ - 2 \\pi i \\frac{(j-1)(k-1)}{N} \\right] x_j``

The implementation is fully generic; as long as the input types have
exponentiation, addition, and multiplication (by complex numbers) defined, the
algorithm will work.  Not particularly fast (about 2ms for FFT of 1024 element
array on my modern-ish MacBook).

Currently only works for 1D arrays.

"""
function fft(xs)
    n = size(xs,1)

    if n == 1
        return xs
    else
        xso = xs[1:2:end]
        xse = xs[2:2:end]

        ffo = fft(xso)
        ffe = fft(xse)

        runity = exp.(-2im*pi*range(0, stop=n/2-1)/n)

        return vcat(ffo .+ runity .* ffe, ffo .- runity.*ffe)
    end
end

end # module


# Simple SIR Model for COVID-19 Modeling
# Michele Fornino
# 3/23/2020
#
# Based on Atkeson (2020)

# Housekeeping
using Plots
using DifferentialEquations

# Define System of Differential Equations
function SIRE!(du, u, p, t)
  # Order of variables in vectors u and du:
  # u[1] S(usceptible)
  # u[2] E(xposed)
  # u[3] I(nfected)
  # u[4] R(recovered/dead)
  #
  # Order of parameters in vector p:
  # p[1] γ
  # p[2] σ

  # Define auxiliary variable
  S, E, I, R = u
  γ, σ, R0 = p
  N = sum(u)
  β = R0 * γ
  du[1] = - β * S / N * I
  du[2] = β * S / N * I - σ * E
  du[3] = σ * E - γ * I
  du[4] = γ * I
end

# Parameters
γ = 1.0/18.0
σ = 1.0/5.2
R0 = 2.5
p = [γ, σ, R0]

I0 = 1.0e-7
E0 = 4.0 * I0
S0 = 1.0 - E0 - I0
R0 = 0.0

u0 = [S0, E0, I0, R0]

T = 18.0 * 30.0
tspan = (0.0, T)

# Define and solve problem
prob = ODEProblem(SIRE!, u0, tspan, p)
sol = solve(prob)


using MultipathChannel
using Plots
using Random
using StatsBase
using LaTeXStrings
using Statistics

Random.seed!(1)

# Rayleigh Distribution
# ---------- Mathmatical Model ----------
σ = 1/√2
normDistRange = -4σ:0.1:4σ
x = normDistRange
normDist = 1/√(2π*σ^2) * exp.(-x.^2/(2σ^2))

rayleighDistRange = 0:0.05:5σ
x = rayleighDistRange
rayleighDist = x./σ^2 .* exp.(-x.^2/(2σ^2))

# ---------- Rayleigh Fading ----------
channelmodel = "rayleigh"
maxdoppler = 0
samplingrate = 15.36e6
numiters = 50000
amp = Vector{Complex{Float64}}(undef, numiters)
for i = 1:numiters
    chan = initChan(maxdoppler, samplingrate, channelmodel)
    amp[i] = applyChan([1], chan, initialtime=rand()*10000)[1]
end

# Create Histograms of in-phase, quadrature, and envelope
HistoRxI = fit(Histogram, real.(amp), normDistRange)
HistoRxINorm = StatsBase.normalize(HistoRxI, mode=:pdf)
HistoRxQ = fit(Histogram, imag.(amp), normDistRange)
HistoRxQNorm = StatsBase.normalize(HistoRxQ, mode=:pdf)
HistoRxEnv = fit(Histogram, abs.(amp), rayleighDistRange)
HistoRxEnvNorm = StatsBase.normalize(HistoRxEnv, mode=:pdf)

# Show a graph of Histograms
p1 = plot(HistoRxINorm, title="In-phase", legend=false); plot!(normDistRange, normDist); xlabel!("Amplitude"); ylabel!("Power Distribution Function")
p2 = plot(HistoRxQNorm, title="Quadrature", legend=false); plot!(normDistRange, normDist); xlabel!("Amplitude"); ylabel!("Power Distribution Function")
p3 = plot(HistoRxEnvNorm, title="Envelope", legend=false); plot!(rayleighDistRange, rayleighDist); xlabel!("Amplitude"); ylabel!("Power Distribution Function")
plot(p1,p2,p3, layout=(1,3))
savefig("RayleighAmplitude.png")


import Statistics
import StatsBase

"""
    r2_score(ytrue, ypred)

Computes coefficient of determination. Higher values are better. Best value
is 1.
"""
function r2_score(
        ytrue::AbstractVector{<:Real},
        ypred::AbstractVector{<:Real},
        )
    if length(ytrue) != length(ypred)
        error("length(ytrue) != length(ypred)")
    end
    if length(ytrue) == 0
        error("length(ytrue) == 0")
    end
    # ybar = mean of the true y values
    ybar = Statistics.mean(ytrue)
    # SStot = total sum of squares
    SStot = sum( (ytrue .- ybar).^2 )
    # SSres = sum of squares of residuals
    residuals = ytrue .- ypred
    SSres = sum( residuals.^2 )
    R2 = 1 - SSres/SStot
    return R2
end


######### StanVariational Bernoulli example  ###########

using StanVariational

bernoulli_model = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";

bernoulli_data = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

# Keep tmpdir across multiple runs to prevent re-compilation
tmpdir = joinpath(@__DIR__, "tmp")

stanmodel = VariationalModel(
  "bernoulli", bernoulli_model; tmpdir = tmpdir)

(sample_file, log_file) = stan_variational(stanmodel; data=bernoulli_data)

if sample_file !== Nothing

  (chns, cnames) = read_variational(stanmodel)

  # Show the same output in DataFrame format
  sdf = read_summary(stanmodel)
  display(sdf)
  println()

  # Retrieve mean value of theta from the summary
  sdf[:theta, :mean]

end


using LinearAlgebra         # 引入线性代数库
x=zeros(10)                 # 10 个元素的零矢量，列矢量
a=ones(3)                   # 3 个元素的矢量，全填充 1.0，列矢量
H=[1.0 2.0 3.0              # 直接构造矩阵
   4.0 5.0 6.0
   7.0 8.0 9.0]
H[2,1]                      # 4.0
ndims(H)                    # H 的维度，2
eltype(H)                   # H 中元素的类型，Float64
length(H)                   # H 中元素数目，9
size(H)                     # H 每一维度的大小所构成的元组，(3,3)
sum(H, dims=2)              # 对第2维求和，相似地，该用法可以用于以下函数： mean, std,
                            # prod, minimum, maximum, any, all
                            # 使用统计函数前需要 using Statistics
H[:,1]                      # 矩阵的第1列
H[1,:]                      # 矩阵的第1行（仍是列矢量）
Hc=H; Hc[1,1]=100.0; H[1,1] # 100.0，数组的赋值是按引用复制
Hc=copy(H)                  # 这样才是真的复制一新的矩阵
H*a                         # 矩阵乘法
exp(H)                      # 矩阵的函数
a*a                         # 出错的矩阵乘法
                            # 列矢量和列矢量的乘法行列不对应无法进行
a⋅a                         # 但是可以做点积
                            # ⋅ 用 \cdot<Tab> 可以打出来（在 REPL、Jupyter 中）
                            # 或者借助插件（见前文开发环境部分）
a×a                         # 还可以做叉积
sum(a)
H'                          # 共轭转置，相当于 adjoint(H)
typeof(H')                  # 注意共轭转置是惰性求值
                            # 它并不直接分配一块内存并写入计算的结果
                            # 这对计算效率有很大的帮助
G=rand(1000,1000); G=G+G'   # 构造一个实对称阵
eigvals(G)                  # 这样的小矩阵，我们可以直接求本征值
eigvecs(G)                  # 本征矢


import SCS
import Convex
import InsarTimeseries

import PyPlot
plt = PyPlot

huber_fit = InsarTimeseries.huber_fit
huber = InsarTimeseries.huber

obj_l1(z) = norm(z, 1)
obj_l2(z) = norm(z, 2)
obj_huber(z) = sum(huber(z)) / 2

function l1_fit(A, b)
    x = Convex.Variable(size(A, 2))
    problem = Convex.minimize(norm(A * x - b, 1))
    Convex.solve!(problem, SCS.SCSSolver(verbose = false))
    return x.value
end

# Fit a line with L1, huber, and L2
function compare_l1_huber(m = 500, m_true = 1, b_true = -5, sigma = 3; with_outlier = true)
    # slope of line is m_true, y intercept is b_true
    x_true = collect(range(0, 10, length = m))
    data_true = (m_true .* x_true) .+ b_true
    @show m_true, b_true
    # b = x_true .+ (sigma .* randn(m))  # Gaussian noise
    b = data_true .+ (2sigma .* rand(m) .- sigma)  # Uniform noise [-sigma, sigma]
    if with_outlier
        extra = maximum(b) - minimum(b)
        b[1] += 2extra
        b[end] -= 2extra
    end

    n = 1
    A = hcat(x_true, ones(m))

    huber_sol = huber_fit(A, b)
    l1_sol = l1_fit(A, b)
    l2_sol = A \ b

    plt.figure()
    plt.plot(x_true, data_true, label = "truth")
    plt.scatter(x_true, b, label = "data")
    plt.plot(x_true, A * huber_sol, label = "Huber")
    plt.plot(x_true, A * l1_sol, label = "L1")
    plt.plot(x_true, A * l2_sol, label = "L2")
    plt.legend()
    plt.show(block = false)

    res_huber = A * huber_sol .- b
    res_l1 = A * l1_sol .- b
    res_l2 = A * l2_sol .- b
    return huber_sol, l1_sol, l2_sol, huber_sol - l1_sol, res_huber, res_l1, res_l2
end

function plot_resid(rh, r1, r2, bins = nothing)
    vm = max(maximum(abs.(rh)), maximum(abs.(r1)))
    fig, axes = plt.subplots(1, 3)
    n1, _, _ = axes[1].hist(rh, range = (-vm, vm), bins = bins)
    n2, _, _ = axes[2].hist(r1, range = (-vm, vm), bins = bins)
    n3, _, _ = axes[3].hist(r2, range = (-vm, vm), bins = bins)
    axes[1].set_title("Huber residuals")
    axes[2].set_title("L1 (CVX) residuals")
    axes[3].set_title("L2 least sq.")
    nmax = max(maximum(n1), maximum(n2), maximum(n3))
    axes[1].set_ylim(0, nmax + 2)
    axes[2].set_ylim(0, nmax + 2)
    axes[3].set_ylim(0, nmax + 2)
    return fig, axes
end

points = 40
sh, s1, s2, dif, rh, r1, r2 = compare_l1_huber(points, with_outlier = true)

nbins = 35
fig, axes = plot_resid(rh, r1, r2, nbins)


# Arvind Prasadan
# 2019 January
# DMF Package
# Evaluates the Autocorrelation Function

"""
	autocorrelation(x; circular = false, bias_correct = false)

Evaluates the Autocorrelation Function of a signal

# Arguments
- `x`: signal (vector)
- `circular`: Boolean whether to use the circular inner product or not
- `bias_correct`: Boolean whether to use n - k instead of n in normalization for a lag of k (theoretically undesirable); ignored if circular is true

# Outputs
- `gamma`: Autocorrelation function (lags 0 to n - 1)
"""

function autocorrelation(x; circular = false, bias_correct = false)

    x = x[:]
    n = length(x)

    x = x .- mean(x)

    gamma = zeros(n)
    if false == circular
        for k = 0:1:(n - 1)
            gamma[k + 1] = x[1:1:(n - k)]' * x[(1 + k):1:n]

            if true == bias_correct
                gamma[k + 1] = gamma[k + 1] / (n - k)
            end
        end

        if false == bias_correct
            gamma = gamma / n
        end
    elseif true == circular
        for k = 0:1:(n - 1)
            gamma[k + 1] = circshift(x, k)' * x
        end

        gamma = gamma / n
    end

    return gamma
end


#
# Simulation of the SIRD epidemics model
#
# Mechanism:  U: Susceptible ; S: Sick ; D: Dead ; I : Immune
#             U + S -> 2S (kc: contamination)
#             S -> D (kd: death)
#             S -> I (ki: immunization)
#             I -> U (kiu: loss of immunity)
#

# Diferential equations
dUdt(U,S,I,k) = -k[1]*U*S + k[4]*I
dSdt(U,S,k) = k[1]*U*S - (k[2]+k[3])*S
dDdt(S,k) = k[2]*S
dIdt(S,I,k) = k[3]*S - k[4]*I

# error
sird_error(U,S,D,I) = abs(1. - (U + S + D + I))

# A method that will be used if the trajectory structure was not created outside
function sird(input :: SIRDInput)
  traj = SIRDTraj(input)
  sird!(input,traj)
  return traj
end

# We need this layer for the optimization procedure to be able to receive a variable (modifies traj)
function sird!(input :: SIRDInput, traj :: SIRDTraj)
  k = [ input.kc, input.kd, input.ki, input.kiu ]
  sird!(k, input, traj)
end

# Function that performs the simulation, modifies the data in traj structure (modifies traj)
function sird!(k :: Vector{Float64}, input :: SIRDInput, traj :: SIRDTraj)

  # Number of steps
  dt = input.dt
  nsteps = round(Int64,input.tmax/dt)
  if nsteps%input.nsave != 0
    error(" In SIRD (tmax/dt) must be a multiple of nsave")
  end
  isave = round(Int64,nsteps/input.nsave)

  # To clear out the code
  U = traj.U
  S = traj.S
  D = traj.D
  I = traj.I
  time = traj.time

  # Initial populations (p for "previous" to the first step)
  Sp = input.Si
  Dp = input.Di
  Ip = input.Ii
  if Sp + Dp + Ip > 1.
    error("Sum of sick, dead, and immune cannot be greater than 1.")
  end
  Up = 1. - (Sp+Dp+Ip)

  # Running simulation
  nlast = nsteps
  nsaved = 0
  for i in 1:nsteps
    # Update populations (c for current)
    Uc = Up + dUdt(Up,Sp,Ip,k)*dt
    Sc = Sp + dSdt(Up,Sp,k)*dt
    Dc = Dp + dDdt(Sp,k)*dt
    Ic = Ip + dIdt(Sp,Ip,k)*dt
    # Save data if required at this step
    if i%isave == 0
      nsaved = nsaved + 1
      U[nsaved] = Uc
      S[nsaved] = Sc
      D[nsaved] = Dc
      I[nsaved] = Ic
      time[nsaved] = i*dt
    end
    # check if epidemics ended
    if Sc < input.tol
      nlast = i
      break
    end
    # check if integration is going fine
    err = sird_error(Uc,Sc,Dc,Ic)
    if err > input.err
      if input.print
        println(input.print," SIRD: Integration error too large at step: ", i," error = ",err)
      end
      nlast = i
      break
    end
    # Update previous values to current ones
    Up = Uc
    Sp = Sc
    Dp = Dc
    Ip = Ic
  end

  if input.print
    println(" Epidemics ended at step = ", nlast)
    println(" Last step saved = ", nsaved, " time = ", time[nsaved])
  end

  # Filling up vector up to the end, if the trajectory ended before
  if nsaved < input.nsave
    if nsaved == 0 # the simulation ended before the first point was saved
      for i in 1:input.nsave
        U[i] = Up
        S[i] = Sp
        D[i] = Dp
        I[i] = Ip
        time[i] = i*isave*dt
      end
    else
      for i in nsaved+1:input.nsave
        U[i] = U[nsaved]
        S[i] = S[nsaved]
        D[i] = D[nsaved]
        I[i] = I[nsaved]
        time[i] = i*isave*dt
      end
    end
  end

end


using LinearAlgebra
"""
Quasi-Newton Methods for Solving Systems of Nonlinear Equations
@author Oscar Veliz
"""

"""
    F(X)

System of nonlinear equations on vector X ([x;y]).
Overwrite for your own testing, currently set to:
x^2 - y - 1
x - y^2 + 1

# Example
```julia-repl
julia> F([1.0; 2.0])
[-2.0, -2.0]
```
"""
F(X) = [X[1]^2 - X[2] - 1; X[1] - X[2]^2 + 1]

"""
    Secant(X, eps=10^-6)

Wolfe-Bittner Secant Method for the system of nonlinear equations in [`F`](@ref).
Requires `d+1` starting vectors in X.
# Example
```julia-repl
julia> Secant(transpose([1.0 1.0;1.0 2.0;1.5 2.0]))
2-element Array{Float64,1}:
 1.6180339736070495
 1.6180339507485015
```
"""
function Secant(X, eps::Float64 = 10^-6)
    n = size(X,1)
    FN = mapslices(F, X; dims=1)
    FX = FN[:,end]
    #println(X) #uncomment to see starting points
    while norm(FX) > eps
        A = [ones(1,n+1); FN]
        b = zeros(n+1,1)
        b[1,1] = 1
        p = A \ b
        Xbar = X * p
        #print(Xbar) #uncomment to see each itr
        X = hcat(X, Xbar)
        X = X[1:end,2:end]
        FX = F(Xbar)
        FN = hcat(FN,FX)
        FN = FN[1:end,2:end]
    end
    X[:,end]
end

"""
    steffJ(X, FX = F(X), n = size(X,1))

Traub-Steffensen approximation of the Jacobian.
Requires starting vector X.
FX is optional in case F(X) was already computed for efficiency.
# Example
```julia-repl
julia> steffJ([1.5;2.0])
2×2 Array{Float64,2}:
 2.25  -1.0
 1.0   -2.5
```
"""
function steffJ(X,FX = F(X),n = size(X,1))
    H = Diagonal(FX)
    J = X #placeholder to set dims
    for i = 1:n
        J = hcat(J, F(X + H[:,i]) - FX)
    end
    J = J[1:end, 2:end] * inv(H) #removes placeholder before mult
end

"""
    Steffensen(X, eps=10^-6)

Traub-Steffensen Method for a system of nonlinear equations in [`F`](@ref).
Requires starting vector X.
# Example
```julia-repl
julia> Steffensen([1.5;2.0])
2-element Array{Float64,1}:
 1.6180339736216265
 1.6180339387805316
```
"""
function Steffensen(X, eps::Float64 = 10^-6)
    n = size(X,1)
    #println(X) #uncomment to see itr
    FX = F(X)
    while(norm(FX) > eps)
        J = steffJ(X,FX,n)
        X = X - J \ FX
        #println(X) #uncomment to see itr
        FX = F(X)
    end
    X
end

"""
    fdJ(X, FX = F(X), h = 10^-4, n = size(X,1))

(Forward) Finite Difference approximation of the Jacobian.
Requires starting vector X. Optional h and n.
FX is optional in case F(X) was already computed for efficiency.
# Example
```julia-repl
julia> fdJ([1.5;2.0])
2×2 Array{Float64,2}:
 3.0001  -1.0
 1.0     -4.0001
```
"""
function fdJ(X,FX = F(X),h::Float64 = 10^-4,n = size(X,1))
    H = h * I(n)
    J = X #placeholder to set dims
    for i = 1:n
        J = hcat(J, F(X + H[:,i]) - FX)
    end
    J = h^-1 * J[1:end, 2:end] #removes placeholder and divide by h
end

"""
    FiniteDiff(X, h = 10^-4, eps = 10^-6)

Generalized Finite-Difference Method for a system of nonlinear equations in [`F`](@ref).
Requires starting vector X. Optional h and eps.
# Example
```julia-repl
julia> FiniteDiff([1.5;2.0])
2-element Array{Float64,1}:
 1.6180340699246842
 1.6180341416981485
```
"""
function FiniteDiff(X,h::Float64 = 10^-4, eps::Float64 = 10^-6)
    n = size(X,1)
    #println(X) #uncomment to see itr
    FX = F(X)
    while(norm(FX) > eps)
        J = fdJ(X,FX,h,n)
        X = X - J \ FX
        #println(X) #uncomment to see itr
        FX = F(X)
    end
    X
end

"""
    Broyden(X, eps = 10^-6)

"Good" Broyden Method for a system of nonlinear equations in [`F`](@ref).
Requires starting vector X. Optional eps. Initial J approximated by differencing.
# Example
```julia-repl
julia> Broyden([1.5;2.0])
2-element Array{Float64,1}:
 1.6180340721138766
 1.6180340507162783
```
"""
function Broyden(X, eps::Float64 = 10^-6)
    #println(X) #uncomment to see itr
    FX = F(X)
    invJ = inv(fdJ(X,FX))
    while(norm(FX) > eps)
        Xold = X
        X = X - invJ*FX
        #println(X) #uncomment to see itr
        deltaX = X - Xold
        FXold = FX
        FX = F(X)
        deltaF = FX - FXold
        trans = transpose(deltaX)
        invJ = invJ + (deltaX - invJ*deltaF)/(trans*invJ*deltaF)*trans*invJ
    end
    X
end

"""
Main
"""

println(Secant(transpose([1.0 1.0;1.0 2.0;1.5 2.0])))
println(Steffensen([1.5;2.0]))
println(FiniteDiff([1.5;2.0]))
println(Broyden([1.5;2.0]))


fun(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

function fun_grad!(g, x)
g[1] = -2.0 * (1.0 - x[1]) - 400.0
* (x[2] - x[1]^2) * x[1]
g[2] = 200.0 * (x[2] - x[1]^2)
end

function fun_hess!(h, x)
h[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
h[1, 2] = -400.0 * x[1]
h[2, 1] = -400.0 * x[1]
h[2, 2] = 200.0
end;

x0 = [0.0, 0.0]
df = TwiceDifferentiable(fun, fun_grad!, fun_hess!, x0)

lx = [-0.5, -0.5]; ux = [0.5, 0.5]
dfc = TwiceDifferentiableConstraints(lx, ux)

res = optimize(df, dfc, x0, IPNewton())

ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())

lx = fill(-Inf, 2); ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())

lx = Float64[]; ux = Float64[]
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())

con_c!(c, x) = (c[1] = x[1]^2 + x[2]^2; c)
function con_jacobian!(J, x)
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    J
end
function con_h!(h, x, λ)
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
end;

lx = Float64[]; ux = Float64[]
lc = [-Inf]; uc = [0.5^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())

lc = [0.1^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())

function con2_c!(c, x)
    c[1] = x[1]^2 + x[2]^2     ## First constraint
    c[2] = x[2]*sin(x[1])-x[1] ## Second constraint
    c
end
function con2_jacobian!(J, x)
    # First constraint
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    # Second constraint
    J[2,1] = x[2]*cos(x[1])-1.0
    J[2,2] = sin(x[1])
    J
end
function con2_h!(h, x, λ)
    # First constraint
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
    # Second constraint
    h[1,1] += λ[2]*x[2]*-sin(x[1])
    h[1,2] += λ[2]*cos(x[1])
    # Symmetrize h
    h[2,1]  = h[1,2]
    h
end;

x0 = [0.25, 0.25]
lc = [-Inf, 0.0]; uc = [0.5^2, 0.0]
dfc = TwiceDifferentiableConstraints(con2_c!, con2_jacobian!, con2_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())

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


using KernelAbstractions

function averageK(K, i, j, k)
    Kw = harmonicMean(K[i - 1, j,     k    ],     K[i,     j,     k    ])
    Ke = harmonicMean(K[i,     j,     k    ],     K[i + 1, j,     k    ])
    Ks = harmonicMean(K[i,     j - 1, k    ],     K[i,     j,     k    ])
    Kn = harmonicMean(K[i,     j,     k    ],     K[i,     j + 1, k    ])
    Kb = harmonicMean(K[i,     j,     k - 1],     K[i,     j,     k    ])
    Kt = harmonicMean(K[i,     j,     k    ],     K[i,     j,     k + 1])
    return Kw, Ke, Ks, Kn, Kb, Kt
end

function harmonicMean(v1, v2)
    return 2.0 * v1 * v2 / (v1 + v2)
end

@kernel function pressureKernel!(source, h, hⁿ⁺¹, Δx, Δy, Δz, Δt, K, SS)
    i, j, k = @index(Global, NTuple)

    @inbounds begin
        F = K[i,j,k] * (
            (h[i + 1,j,k] + h[i - 1,j,k] - 2.0 * h[i,j,k]) / (Δx * Δx) +
            (h[i,j + 1,k] + h[i,j - 1,k] - 2.0 * h[i,j,k]) / (Δy * Δy) +
            (h[i,j,k + 1] + h[i,j,k - 1] - 2.0 * h[i,j,k]) / (Δz * Δz)
            ) + source[i,j,k] / (Δx * Δy)

        hⁿ⁺¹[i,j,k] = h[i,j,k] + Δt * F * (1.0 / SS)
    end

end

@inline function pressureKernelAveraged!(i, j, k, source, h, u, v, w, hⁿ⁺¹, uⁿ⁺¹, vⁿ⁺¹, wⁿ⁺¹, Δx, Δy, Δz, Δt, K, SS)

    Kw, Ke, Ks, Kn, Kb, Kt = averageK(K, i, j, k)

    F = (1.0 / Δx * Δx) * (Ke * (h[i + 1,j  ,k  ] - h[i,j,k]) - Kw * (h[i,j,k] - h[i - 1,j  ,k  ])) +
        (1.0 / Δy * Δy) * (Kn * (h[i  ,j + 1,k  ] - h[i,j,k]) - Ks * (h[i,j,k] - h[i  ,j - 1,k  ])) +
        (1.0 / Δz * Δz) * (Kt * (h[i  ,j  ,k + 1] - h[i,j,k]) - Kb * (h[i,j,k] - h[i  ,j  ,k - 1])) +
        source[i,j,k] / (Δx * Δy)

    hⁿ⁺¹[i,j,k] = h[i,j,k] + Δt * F * (1.0 / SS)
end


using Plots, Distributions, Random; pyplot()
Random.seed!(0)

n = 2000
data = rand(Normal(),n)
l, m = minimum(data), maximum(data)

delta = 0.3;
bins = [(x,x+delta) for x in l:delta:m-delta]
if last(bins)[2] < m
    push!(bins,(last(bins)[2],m))
end
L = length(bins)

inBin(x,j) = first(bins[j]) <= x && x < last(bins[j])
sizeBin(j) = last(bins[j]) - first(bins[j])
f(j) = sum([inBin(x,j)  for x in data])/n
h(x) = sum([f(j)/sizeBin(j) * inBin(x,j) for j in 1:L])

xGrid = -4:0.01:4
histogram(data,normed=true, bins=L,
    label="Built-in histogram",
    c=:blue, la=0, alpha=0.6)
plot!(xGrid,h.(xGrid), lw=3, c=:red, label="Manual histogram",
    xlabel="x",ylabel="Frequency")
plot!(xGrid,pdf.(Normal(),xGrid),label="True PDF",
    lw=3, c=:green, xlims=(-4,4), ylims=(0,0.5))


abstract AbstractCovarianceFunction

abstract StationaryCovarianceFunction <: AbstractCovarianceFunction

Base.call(f::StationaryCovarianceFunction,x,y) = f(sqrt(sumabs2(x-y)))

doc"""
A Gaussian covariance function is of the form

C(x_1,x_2) = σ^2 e^{-\|x_1-x_2\|^2/2L^2}

where L is the bandwidth of the Gaussian kernel.
"""
type GaussianCovarianceFunction <: StationaryCovarianceFunction
    L::Real
    σ::Real
end

Base.call(f::GaussianCovarianceFunction,h) = f.σ^2*exp(-h^2/(2*f.L^2))

doc"""
A Matérn covariance function takes the form

C(x_1,x_2) = σ^2 \frac{2^{1-ν}}{Γ(ν)}(\sqrt{2ν}\|x_1-x_2\|/ρ)^ν K_ν(\sqrt{2ν}\|x_1-x_2\|/ρ)

where ν gives the order of the modified Bessel function of the second kind, K_ν, ρ is the bandwidth/correlation length and σ is the total covariance.
"""
type MatérnCovarianceFunction <: StationaryCovarianceFunction
    ν::Real
    ρ::Real
    σ::Real
end

Base.call(f::MatérnCovarianceFunction,h) = (h==0?f.σ^2:f.σ^2*2.0^(1-f.ν)/gamma(f.ν)*(sqrt(2*f.ν)*h/f.ρ)^f.ν*besselk(f.ν,sqrt(2*f.ν)*h/f.ρ))

doc"""
Generates a covariance matrix from a stationary covariance function
and a matrix of distances as might be generated by Distances.jl's
`pairwise` function
"""
covariancematrix(f::StationaryCovarianceFunction,D::Matrix) = map(f,D)


using LinearAlgebra
using Statistics

include("algo_util.jl")

abstract type FirstOrder <: DescentDirectionMethod end
function solve(M::FirstOrder, f, ∇f, x, max_iters; num_eval_termination=true)
    init!(M, f, ∇f, x)
    x_hist = [x]
    while true
        x = step!(M, f, ∇f, x)
        push!(x_hist, x)
        if num_eval_termination && (count(∇f) >= max_iters/2) # 2 calls per iteration
            break
        end
    end
    return x, x_hist
end

# Vanilla Gradient Descent
Base.@kwdef mutable struct GradientDescent <: FirstOrder
    α = 1e-3
end
init!(M::GradientDescent, f, ∇, x) = M
function step!(M::GradientDescent, f, ∇f, x)
    α = M.α
    return x - α * ∇f(x)
end

# GD + Momentum
Base.@kwdef mutable struct GDMomentum <: FirstOrder
    α = 1e-3
    β = 0.9
    v = nothing
end
function init!(M::GDMomentum, f, ∇f, x)
    M.v = zeros(length(x))
    return M
end
function step!(M::GDMomentum, f, ∇f, x)
    α, β, = M.α, M.β
    M.v[:] = β * M.v - α * ∇f(x)
    return x + M.v
end

# GD + Nesterov Momentum
Base.@kwdef mutable struct GDMomentumNesterov <: FirstOrder
    α = 1e-3
    β = 0.9
    v = nothing
end
function init!(M::GDMomentumNesterov, f, ∇f, x)
    M.v = zeros(length(x))
    return M
end
function step!(M::GDMomentumNesterov, f, ∇f, x)
    α, β = M.α, M.β
    M.v[:] = β * M.v - α * ∇f(x + β * M.v)
    return x + M.v
end

# Adam
Base.@kwdef mutable struct Adam <: FirstOrder
    α = 1e-3
    β1 = 0.9
    β2 = 0.999
    ϵ = 1e-8
    k = 0
    v = nothing
    s = nothing
end
function init!(M::Adam, f, ∇f, x)
    M.v = zeros(length(x))
    M.s = zeros(length(x))
    return M
end
function step!(M::Adam, f, ∇f, x)
    α, β1, β2, ϵ = M.α, M.β1, M.β2, M.ϵ
    g = ∇f(x)
    M.v[:] = β1 * M.v + (1 - β1) * g
    M.s[:] = β2 * M.s + (1 - β2) * g .* g
    M.k += 1
    v̂ = M.v ./ (1 - β1 ^ M.k)
    ŝ = M.s ./ (1 - β2 ^ M.k)
    return x - (α * v̂) ./ (sqrt.(ŝ) .+ ϵ)
end

# Gradient Descent + Backtracking Line Search
Base.@kwdef mutable struct GDApproxLineSearch <: FirstOrder
    α = 3e-3
    approx_line_search = backtracking_line_search
end
init!(M::GDApproxLineSearch, f, ∇, x) = M
function step!(M::GDApproxLineSearch, f, ∇f, x)
    g = ∇f(x)
    α = M.approx_line_search(f, ∇f, x, -g, M.α)
    return x - α * g
end