using LinearAlgebra: eigen, cholesky, diag, diagm, UpperTriangular, LowerTriangular
using Statistics: mean

"""
  Functions for mvrnorm, covriance to correlation conversion (cov2cor)
  and vice versa cor2cov.
"""
function generateRandomMatrix(p::Int64, n::Int64)
  x = rand(n, p)
  return x' * x
end

"""
  Function to convert covariance to correlation matrix
"""
function cov2cor(E::Array{T, 2}) where {T <: AbstractFloat}
  p = size(E)[1]
  d::Array{T, 1} = diag(E)
  d .^= -0.5
  C = (d * d') .* E
  C = 0.5.*(UpperTriangular(C) + LowerTriangular(C)')
  C = C + C'
  for i in 1:p
    C[i, i] = T(1)
  end
  return C
end

"""
  Function to convert correlation to covariance matrix
  d is the variance vector of the covariance matrix
"""
function cor2cov(C::Array{T, 2}, d::Array{T, 1}) where {T <: AbstractFloat}
  p = size(C)[1]
  d1 = d .^0.5
  E = (d1 * d1') .* C
  E = 0.5.*(UpperTriangular(E) + LowerTriangular(E)')
  E = E + E'
  for i in 1:p
    E[i, i] /= 2
  end
  return E
end

"""
  Function that carries samples from a multivariate normal
  distribution when supplied with the number of samples (n),
  the mean vector (mu), and the covariance matrix (sigma).

  Example:
  mvrnorm(100, zeros(10), cov2cor(generateRandomMatrix(10, 10000)))
"""
function mvrnorm(n::Int64, mu::Array{T, 1}, sigma::Array{T, 2}) where {T <: AbstractFloat}
  p = size(sigma)[1]
  A = cholesky(sigma).L
  output = zeros(T, (n, p))
  for i = 1:n
    output[i, :] = mu + A * randn(p)
  end
  return output
end

#=======================================================================================================#
"""
  Random Number Generator for Beta (and Uniform) Distribution
"""
abstract type AbstractSampleDistribution end
struct BetaSampleDistribution{T} <: AbstractSampleDistribution
  alpha::T
  beta::T
  function BetaSampleDistribution(alpha::T, beta::T) where {T <: AbstractFloat}
    return new{T}(alpha, beta)
  end
  function BetaSampleDistribution(alpha::T, beta::T) where {T <: Integer}
    return BetaSampleDistribution(Float64(alpha), Float64(beta))
  end
end

function calcSample(ualpha::T, vbeta::T) where {T <: AbstractFloat}
  return ualpha/(ualpha + vbeta)
end
function condition(ualpha::T, vbeta::T) where {T <: AbstractFloat}
  return (ualpha + vbeta) > T(1)
end

"""
  Version 1 of the random sample from beta distribution
  Reference: C. P. Robert, G. Casella,
            Monte Carlo Statistical Methods, Example 2.11 p44
"""
function sample(distrib::BetaSampleDistribution{T}, shape::Tuple{Vararg{Int64}}) where {T <: AbstractFloat}
  n = prod(shape)
  U = rand(T, n); V = rand(T, n)
  Y::Array{T} = zeros(T, shape)
  ialpha = 1/distrib.alpha; ibeta = 1/distrib.beta
  for i in 1:n
    u::T = U[i]; v::T = V[i]
    ualpha = u .^ialpha; vbeta = v .^ibeta;
    while condition(ualpha, vbeta)
      u = rand(T, 1)[1]; v = rand(T, 1)[1]
      ualpha = u .^ialpha; vbeta = v .^ibeta;
    end
    Y[i] = calcSample(ualpha, vbeta)
  end
  return Y
end

function sample(distrib::BetaSampleDistribution{T}, shape::Int64) where {T <: AbstractFloat}
  return sample(distrib, (shape,))
end

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

struct UniformSampleDistribution{T} <: AbstractSampleDistribution
  min::T
  max::T
  function UniformSampleDistribution(min::T, max::T) where {T <: AbstractFloat}
    @assert(min < max, "Minimum value is not less than maximum value")
    return new{T}(min, max)
  end
  function UniformSampleDistribution(min::T, max::T) where {T <: Integer}
    return UniformSampleDistribution(Float64(min), Float64(max))
  end
end

function sample(distrib::UniformSampleDistribution{T}, shape::Tuple{Vararg{Int64}}) where {T <: AbstractFloat}
  rsample::Array{T} = rand(T, shape)
  return (rsample .* (distrib.max - distrib.min)) .+ distrib.min
end
function sample(distrib::UniformSampleDistribution{T}, shape::Int64) where {T <: AbstractFloat}
  return sample(distrib, (shape,))
end

function Base.min(x::Array{T}) where {T <: AbstractFloat}
  n = length(x)
  @assert(n > 0, "Length of array is zero.")
  if n == 1
    return x[1]
  end
  ret::T = x[1]
  for i in 2:n
    ret = ret > x[i] ? x[i] : ret
  end
  return ret
end

function Base.max(x::Array{T}) where {T <: AbstractFloat}
  n = length(x)
  @assert(n > 0, "Length of array is zero.")
  if n == 1
    return x[1]
  end
  ret::T = x[1]
  for i in 2:n
    ret = ret < x[i] ? x[i] : ret
  end
  return ret
end

function Base.range(x::Array{T}) where {T <: AbstractFloat}
  return [min(x), max(x)]
end

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

"""
  Identity Matrix
"""
function I(::Type{T}, p::Int64) where {T <: Number}
  x = zeros(T, (p, p))
  for i in 1:p
    x[i, i] = T(1)
  end
  return x
end
"""
  Convenience overload
"""
function I(p::Int64)
  return I(Float64, p)
end

"""
  Types for generating random correlation matrices
"""
abstract type AbstractRandomCorrelationMatrix end

struct BetaGenerator <: AbstractRandomCorrelationMatrix end
struct OnionGenerator <: AbstractRandomCorrelationMatrix end
struct UniformGenerator <: AbstractRandomCorrelationMatrix end
struct VineGenerator <: AbstractRandomCorrelationMatrix end

"""
  Vine method for generating a random correlation matrix
  Source:
    https://stats.stackexchange.com/questions/2746/how-to-efficiently-generate-random-positive-semidefinite-correlation-matrices/125017#125017

    Check it with to make sure the solutions are correct:
    Journal of Multivariate Analysis 100 (2009) 1989–2001
    Generating random correlation matrices based on vines and extended
    onion method. Daniel Lewandowski(a)*, Dorota Kurowicka (a),
    Harry Joe (b).

  randomCorrelationMatrix(VineGenerator(), 10, 2.0)
"""
function randomCorrelationMatrix(::VineGenerator, d::Int64, eta::T) where {T <: AbstractFloat}

  beta = eta + (d - 1)/2
  P = zeros(T, (d, d))
  S = I(T, d)

  for k in 1:(d - 1)
    beta = beta - (1/2)
    distrib = BetaSampleDistribution(beta, beta)

    for i in (k + 1):d
      P[k, i] = sample(distrib, 1)[1]
      P[k, i] = (P[k, i] - 0.5)*2
      p = P[k, i]

      for l in (k - 1):-1:1
        p = p * sqrt((1 - P[l, i]^2)*(1 - P[l, k]^2)) + P[l, i]*P[l, k]
      end

      S[k, i] = p
      S[i, k] = p
    end
  end

  return S
end

"""
  Random correlation matrix using the Onion Generator
"""
function randomCorrelationMatrix(::OnionGenerator, d::Int64, eta::T) where {T <: AbstractFloat}
  beta::T = eta + (d - 2)/2
  distrib = BetaSampleDistribution(beta, beta)
  u = sample(distrib, 1)[1]
  r = I(T, 2)
  r[1, 2] = r[2, 1] = 2*u - 1
  for k in 2:(d - 1)
    beta -= T(1/2)
    distrib = BetaSampleDistribution(T(k/2), beta)
    y = sample(distrib, 1)[1]
    U = rand(T, k)
    w = sqrt(y) .* U

    ev = eigen(r)
    A = ev.vectors * diagm(abs.(ev.values).^(0.5)) * ev.vectors'
    z = A * w
    r = [r z; z' T(1)]
  end
  return r
end

"""
  Random correlation matrix by sampling from Beta Distribution
  # Example:
  randomCorrelationMatrix(BetaGenerator(), 10, (1.0, 1.0))
"""
function randomCorrelationMatrix(::BetaGenerator, d::Int64, (alpha, beta)::Tuple{T, T}) where {T <: AbstractFloat}
  distrib = BetaSampleDistribution(alpha, beta)
  r = sample(distrib, (d, d))
  # Change range to (-1, 1)
  r .= (r .* 2) .+ T(-1)
  # Symmetry
  r .= T(0.5) .* (r + r')
  for i in 1:d
    r[i, i] = T(1)
  end
  ev = eigen(r)
  r = ev.vectors * diagm(sort(abs.(ev.values))) * ev.vectors'
  maxR = max(r)
  r ./= maxR
  r .= T(0.5) .* (r + r')
  for i in 1:d
    r[i, i] = T(1)
  end
  return r
end

"""
  Random correlation matrix by sampling from Uniform Distribution
  # Example:
  randomCorrelationMatrix(Float64, UniformGenerator(), 10)
"""
function randomCorrelationMatrix(::Type{T}, ::UniformGenerator, d::Int64) where {T <: AbstractFloat}
  distrib = UniformSampleDistribution(T(-1), T(1))
  r = sample(distrib, (d, d))
  r .= T(0.5) .* (r + r')
  for i in 1:d
    r[i, i] = T(1)
  end
  ev = eigen(r)
  r = ev.vectors * diagm(sort(abs.(ev.values))) * ev.vectors'
  maxR = max(abs.(r)...)
  r ./= maxR
  r .= T(0.5) .* (r + r')
  for i in 1:d
    r[i, i] = T(1)
  end
  return r
end

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

"""
  Function to simulate X and eta, p = number of parameters
  (including intercept), n = number of samples.

  # Example:
    using Random: seed!
    seed!(0)
    simulateData(Float64, 10, 1000)
"""
function simulateData(::Type{T}, p::Int64, n::Int64, delta::T = T(0)) where {T <: AbstractFloat}

  corr = randomCorrelationMatrix(T, UniformGenerator(), p)
  mu = zeros(T, p)
  X = mvrnorm(n, mu, corr)

  #= The intercept =#
  X[:, 1] .= T(1)

  b = zeros(T, p)
  idist = UniformSampleDistribution(T(0), T(0.3))
  b[1] = sample(idist, 1)[1]
  if length(b) > 1
    distrib = UniformSampleDistribution(T(-0.1), T(0.1))
    b[2:p] = sample(distrib, p - 1)
  end

  eta = X*b
  sd = 0.5*abs(mean(eta))
  eta = delta .+ eta .+ sd .* randn(n)

  return (X = X, eta = eta)
end


function LGnodes(alpha::Real, N::Int)
    # generate nodes of Laguerre-Gauss quadrature nodes.
    alpha > -1 || error("alpha <= -1 not allowed.")
    j = 1:N |> collect
    b = j .* (j .+ alpha)
    # +- diagonal
    dl = -sqrt.(b)
    j = 0:N |> collect
    # main diagonal
    d = 2.0 * j .+ alpha .+ 1
    A = SymTridiagonal(d, dl)
    return eigvals(A)
end

function LGRnodes(N::Int)
    # Laguerre-Gauss-Radau quadrature nodes.
    xi = LGnodes(1.0, N - 1)
    x = [0; xi]
    return x
end

function LGR(N::Int)
    # get nodes and weights for LGR.
    x = LGRnodes(N)
    LN = eval_laguerre_function.(N, x)
    # equation 7.37 in Shen, Tang & Wang.
    w = 1 ./ ( (N + 1) * LN.^2)
    return x, w
end

function eval_laguerre_function(n::Int, x::Real)
    # evaluate Laguerre function using the recurrence relation
    L0 = exp(-x / 2)
    L1 = (1 - x) * exp(-x / 2)
    if n == 0
        return L0
    end
    if n == 1
        return L1
    end
    # n >= 2 case
    Lnew = 0.0
    for i = 2:n
        Lnew = (2 - 1 / i - x / i) * L1 - (1 - 1 / i) * L0
        L0 = L1
        L1 = Lnew
    end
    return Lnew
end

function eval_LaguerreFunction(n::Int, x::Real)
    # method in Shen, SIAM J NUMER ANAL, Vol 38, No 4, pp. 1113-1133
    # STABLE AND EFFICIENT SPECTRAL METHODS IN UNBOUNDED
    # DOMAINS USING LAGUERRE FUNCTIONS
    signL0 = 1
    # TODO: need to indentify the error in the equations cited.
    # d0 = 1
    # signL1 = sign(1 - x)
    # d1 = abs(1 - x) * exp(-x / 2)
    # if n == 0
    #     return exp(-x / 2)
    # end
    # if n == 1
    #     return signL1 * d1
    # end
    # for i = 2:n
    #     # common term
    #     t = (2 * i - 1 - x) / i * signL1 * d1 * (d1 / d0)^(i - 2) - (1 - 1 / i) * signL0
    #     signLi = sign(t)
    #     di = d0^(1. - 2 / i) * abs(t)^(1.0 / n)
    #     # update old variables.
    #     signL0 = signL1
    #     signL1 = signLi
    #     d0 = d1
    #     d1  = di
    #     if i == n
    #         return signLi * di^n
    #     end
    # end

end

function LGRquad(f::Function, N::Int)
    # numerical integration using Laguerre-Gauss-Radau quadrature.
    # from 0 to infinity.
    x, w = LGR(N)
    return sum(f.(x) .* w)
end

function laguerre_transform(f::Function, N::Int)
    # interpolation: represent f in a series expansion
    # f = sum_{n=0}^N a_n Lnhat(x)
    # where Lnhat is the Laguerre function.
    # an = \sum_{j=0}^N f(x_j)Lnh(x_j)what_j
    x, w = LGR(N)
    an = zeros(N + 1)
    for i = 1:N + 1
        Ln = eval_laguerre_function.(i - 1, x)
        an[i] = sum(f.(x) .* Ln .* w)
    end
    return an
end

function inverse_laguerre_transform(an::AbstractVector, x::AbstractArray)
    # take an array of coefficients and reconstruct the function.
    N = length(an) - 1
    # f = sum_{n=0}^N an Lnh(x)
    f = zeros(size(x))

    for (a, n) in zip(an, 0:N)
        Lnh = eval_laguerre_function.(n, x)
        f += a * Lnh
    end
    return f
end


using NumericalMethodsforEngineers

a = [16. 4. 8.; 4. 5. -4.; 8. -4. 22.]
b = [16., 18., -22.]

a = sparse(a)

f = cholesky(a)
f |> display

y = f.L \ b
c = f.U \ y
println("Solution Vector: \n $c")

@test c == a \ b


u, v, w = ℝ³(rand(3)), ℝ³(rand(3)), ℝ³(rand(3))
a, b = ℝ³(rand(3)), ℝ³(rand(3))
zero = ℝ³([0.0; 0.0; 0.0])
α, β = rand(2)

## Abstract Vector Space Tests ##

@test isapprox(u + (v + w), (u + v) + w) # 1. associativity of addition
@test isapprox(u + v, v + u) # 2. commutativity of addition
@test isapprox(u + zero, u) # 3. the zero vector
@test isapprox(u - u, zero) # 4. the inverse element
@test isapprox(α * (u + v), α * u + α * v) # 5. distributivity Ι
@test isapprox((α + β) * u, α * u + β * u) # 6. distributivity ΙΙ
@test isapprox(α * (β * u), (α * β) * u) # 7. associativity of scalar multiplication
@test isapprox(1u, u) # 8. the unit scalar 1

## Product Space Tests ##

@test isapprox(dot(u, v), dot(v, u)) # symmetric
@test isapprox(dot(u, α * a + β * b), α * dot(u, a) + β * dot(u, b)) # linear
@test dot(u, u) ≥ 0 && isapprox(dot(u, zero), 0) # positive semidefinite
@test size(outer(u, v)) == (length(u), length(v))

## isapprox

segments = 5
array1 = Array{ℝ³,2}(undef, segments, segments)
for i in 1:segments
    for j in 1:segments
        array1[i, j] = ℝ³(rand(3))
    end
end

array2 = map(x -> 2x, array1)

@test isapprox(array1, array1, atol = TOLERANCE)
@test isapprox(array1, array2, atol = TOLERANCE) == false


export rayleigh_criterion

"""
    rayleigh_criterion(focal_length, diameter; λ=500e-9)

Calculates the resolution of a microscope according to Rayleigh.
`focal_length` is the focal length of the lens.
`λ` is the wavelength and `diameter` is the diameter of the aperture of the lens.
See [Wikipedia](https://en.wikipedia.org/wiki/Angular_resolution#The_Rayleigh_criterion).

 # Examples
```julia-repl
julia> rayleigh_criterion(100, 200)
3.05e-7

julia> rayleigh_criterion(100, 200, λ=100e-9)
6.099999999999999e-8
```
"""
function rayleigh_criterion(focal_length, diameter; λ=500e-9)
    return rayleigh_criterion(calc_NA(focal_length, diameter), λ=λ)
end

"""
    rayleigh_criterion(NA; λ=500e-9)

Calculates the resolution of a microscope according to Rayleigh.
`λ` is the wavelength and `NA` the numerical aperture of the system.
See [Wikipedia](https://en.wikipedia.org/wiki/Angular_resolution#The_Rayleigh_criterion).

 # Examples
```julia-repl
julia> rayleigh_criterion(1.0)
3.05e-7

julia> rayleigh_criterion(1.0, λ=100e-9)
6.099999999999999e-8
```
"""
function rayleigh_criterion(NA; λ=500e-9)
    return 0.61 * λ / NA
end


using PyCall
using PyPlot
using SpecialFunctions
# uncomment for non-jupyter use:
#pygui(true)

### FUNCTIONS, used for several subtasks:
# %%

# the Bernoulli PDF:
function bernoulli_pdf(θ, y)
    if 0 <= θ <= 1 && (y==0 || y==1)    # Only admissible values
        return θ^y * (1-θ)^(1-y)
    else
        return 0
    end
end

# Likelihood fct of a (series of) coin tosses with bias:
function likelihood_fct(θ,y)
    # Make array s.t. multiple theta values can be input:
    likelihood = Array{Float64}(undef,length(θ))
    # Do that for every theta value requested:
    for ii in 1:length(θ)
        # Elementwise application of the beronoulli_pdf over outcome vector y, then MULTIPLY together:
        likelihood[ii] = prod(bernoulli_pdf.(θ[ii], y))
    end
    return likelihood
end

#%%

# Log-Bernoulli PDF:
function log_bernoulli_pdf(θ,y)
    if 0 < θ < 1 && (y==0 || y==1)    # Only admissible values
        return y * log(θ) + (1-y) * log(1-θ)
    else
        return -Inf
    end
end

# Log-Likelihood fct of a (series of) coin tosses with bias:
function log_likelihood_fct(θ, y)
    # Make array s.t. multiple theta values can be input:
    likelihood = Array{Float64}(undef,length(θ))
    # Do that for every theta value requested:
    for ii in 1:length(θ)
        # Elementwise application of the log_beronoulli_pdf over outcome vector y, then SUM together:
        likelihood[ii] = sum(log_bernoulli_pdf.(θ[ii],y))
    end
    return likelihood
end

#%%

### Choose and run the subtask:

###########
### Task 2
# Given the parameter, plot outcome probabilities
#%%
θ = 0.5
outcomes = [0, 1]

figure()
bar(outcomes, bernoulli_pdf.(θ,outcomes))
xticks([0, 1])
xlabel("y")
ylabel("θ")
plt.show()

#%%

θ = 0.25

figure()
bar(outcomes, bernoulli_pdf.(θ,outcomes))
xticks([0, 1])
xlabel("y")
ylabel("θ")
plt.show()

#%%

###########
### Task 3
# Given the outcome, plot parameter likelihoods
#%%
θ = 0:0.01:1;
outcome = 1;

figure()
plot(likelihood_fct.(θ,outcome))
plt.show()

#%%
outcome = 0

figure()
plot(likelihood_fct.(θ,outcome))
plt.show()
#%%

###########
### Task 4
# Evaluate likelihood fct after n flips

#%%
θ = 0.5

# a)
println("a) Regular Likelihood:")

for i in 1:3
    n = 10 * 100^(i-1)

    outcome = (rand(n) .>= 0.5) * 1

    print("L(y_1,y_2,...,y_$n | θ = $θ) = ")
    println(likelihood_fct(θ,outcome))
end

println("\nWith rising n, the likelihood -> 0")

#%%
# b) and c)
println("\nb) log-likelihoods, l, and L = exp(l):")

for i in 1:4
    n = 10 * 100^(i-1)

    outcome = (rand(n) .>= 0.5) * 1
    l = log_likelihood_fct(θ,outcome)
    L = exp.(l)

    print("l(y_1,y_2,...,y_$n | θ = $θ) = ")
    println(l)
    print("L(y_1,y_2,...,y_$n | θ = $θ) = ")
    println(L)
end

println("\nYes, n > 100000 could be evaluated without over-/underflow in logarithmic space.")

#%%
# d)
println("\nd) Plotting the likelihood function for a series")
θ = 0:0.01:1;

outcomeSeries = [[1], [1 1], [1 1 0 1]]

for thisOutcome in outcomeSeries
    figure()
    subplot(1,2,1)
    plot(θ,likelihood_fct(θ,thisOutcome),label="L($thisOutcome|θ)")
    grid()
    legend()
    subplot(1,2,2)
    plot(θ,log_likelihood_fct(θ,thisOutcome),label="l($thisOutcome|θ)")
    grid()
    legend()
    plt.show()
end
#%%

###########
### Task 5
# Implement the beta function and create plots

### i)
# Logarithmic version of the beta-distribution
function Log_Beta(θ,a::Any,b::Any)
    #using SpecialFunctions
    if 0 < θ < 1
        # Numerator
        num = log(θ)*(a-1) + log(1-θ)*(b-1)
        # The denominator (Beta-function) can be integrated or made up of Gamma-functions
        # Since we only treat integer values here, Γ(a) = (a-1)!
        #den = log(factorial(a-1) * factorial(b-1) / factorial(a+b-1))
        den = log(SpecialFunctions.beta(a,b))
        return num - den
    else
        return -Inf
    end
end

#%%
### ii) sub-tasks 1-3
function Log_posterior(θ,a,b,y)
    heads = sum(y)
    tails = length(y) - heads
    posterior = Array{Float64}(undef,length(θ))
    # Log addition, the brackets are for converting into a single float:
    for ii in 1:length(θ)
        posterior[ii] = (Log_Beta(θ[ii],a,b) .+ log_likelihood_fct(θ[ii], y) .-
                    log(SpecialFunctions.beta(a+heads,b+tails)))[1];
                    #log(factorial(a+heads-1) * factorial(b+tails-1) / factorial(a+heads+b+tails-1)))[1];
    end
    return posterior
end

#%%

println("\nd) Plotting the likelihood function for a series")
θ = 0:0.001:1;

#%%
# Flat prior:
a = 1;
b = 1;

figure()
subplot(1,2,1)
plot(θ,exp.(Log_Beta.(θ,a,b)), label="beta(θ,$a,$b)")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_Beta.(θ,a,b), label="log(beta(θ,$a,$b))")
grid()
legend()
plt.show()

# %%
y = [1]
a = 1
b = 1

figure()
subplot(1,2,1)
plot(θ,exp.(Log_posterior(θ,a,b,y)), label="posterior for $y")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_posterior(θ,a,b,y), label="log-posterior for $y")
grid()
legend()
plt.show()

# Alternatively directly by Beta:
heads = sum(y)
tails = length(y) - heads

figure()
subplot(1,2,1)
plot(θ,exp.(Log_Beta.(θ,a+heads,b+tails)), label="beta(θ,$(a+heads),$(b+tails))")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_Beta.(θ,a+heads,b+tails), label="beta(θ,$(a+heads),$(b+tails))")
grid()
legend()
plt.show()

# %%
y = [1 1]
a = 1
b = 1

figure()
subplot(1,2,1)
plot(θ,exp.(Log_posterior(θ,a,b,y)), label="posterior for $y")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_posterior(θ,a,b,y), label="log-posterior for $y")
grid()
legend()
plt.show()

# Alternatively directly by Beta:
heads = sum(y)
tails = length(y) - heads

figure()
subplot(1,2,1)
plot(θ,exp.(Log_Beta.(θ,a+heads,b+tails)), label="beta(θ,$(a+heads),$(b+tails))")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_Beta.(θ,a+heads,b+tails), label="beta(θ,$(a+heads),$(b+tails))")
grid()
legend()
plt.show()

# %%
y = [1 1 0 1]
a = 1
b = 1

figure()
subplot(1,2,1)
plot(θ,exp.(Log_posterior(θ,a,b,y)), label="posterior for $y")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_posterior(θ,a,b,y), label="log-posterior for $y")
grid()
legend()
plt.show()

# Alternatively directly by Beta:
heads = sum(y)
tails = length(y) - heads

figure()
subplot(1,2,1)
plot(θ,exp.(Log_Beta.(θ,a+heads,b+tails)), label="beta(θ,$(a+heads),$(b+tails))")
grid()
legend()
subplot(1,2,2)
plot(θ,Log_Beta.(θ,a+heads,b+tails), label="beta(θ,$(a+heads),$(b+tails))")
grid()
legend()
plt.show()

#%%

println("The differences are a scaling in the non-log curves and thus
a shifting in the log-curves.
The differences arise because before we only used the LIKELIHOOD function
before, but now we use a normalised POSTERIOR.")

# %%
### iii) More informative prior and reproduce Fig 6.4

# First Column
θ = 0:.01:1
a = 100
b = 100

y = [ones(1,17) zeros(1,3)]

subplot(3,1,1)
plot(θ,exp.(Log_Beta.(θ,a,b)))
grid()
subplot(3,1,2)
plot(θ,exp.(log_likelihood_fct(θ,y)))
grid()
subplot(3,1,3)
plot(θ,exp.(Log_posterior(θ,a,b,y)))
grid()
plt.show()
#%%

# Second Column
θ = 0:.01:1
a = 18.25
b = 6.75

y = [ones(1,17) zeros(1,3)]

subplot(3,1,1)
plot(θ,exp.(Log_Beta.(θ,a,b)))
grid()
subplot(3,1,2)
plot(θ,exp.(log_likelihood_fct(θ,y)))
grid()
subplot(3,1,3)
plot(θ,exp.(Log_posterior(θ,a,b,y)))
grid()
plt.show()

#%%

# Third Column
θ = 0:.001:1
a = 1
b = 1

y = [ones(1,17) zeros(1,3)]

subplot(3,1,1)
plot(θ,exp.(Log_Beta.(θ,a,b)))
grid()
subplot(3,1,2)
plot(θ,exp.(log_likelihood_fct(θ,y)))
grid()
subplot(3,1,3)
plot(θ,exp.(Log_posterior(θ,a,b,y)))
grid()
plt.show()


# # Wasserstein non-negative matrix factorisation
#
# In this example, we implement Wasserstein non-negative matrix factorisation (W-NMF)
# following the paper [^RCP16] by Rolet et al.
#
# [^RCP16]: Rolet, Antoine, Marco Cuturi, and Gabriel Peyré. ["Fast dictionary learning with a smoothed Wasserstein loss."](https://marcocuturi.net/Papers/rolet16fast.pdf) Artificial Intelligence and Statistics. PMLR, 2016.
#
# ## Introduction
#
# Matrix decomposition is a classical problem in machine learning. Given a $m \times n$ matrix $X$ representing $n$ repeated $m$-dimensional observations, one may seek matrices $D, \Lambda$ of appropriate dimensions such that
# ```math
# X \approx D \Lambda.
# ```
# For a target rank $k < \min \{ m, n \}$, i.e. $D \in \mathbb{R}^{m \times k}, \Lambda \in \mathbb{R}^{k \times n}$, this problem can be thought of as seeking a low-dimensional linear representation $D \Lambda$ of the potentially high-dimensional dataset $X$.
#
# Lee and Seung [^LS99] observed that the data matrix $X$ is non-negative in many practical applications, and that naturally one may want the factor matrices $D, \Lambda$ to be also non-negative.
#
# [^LS99]: Lee, Daniel D., and H. Sebastian Seung. "Learning the parts of objects by non-negative matrix factorization." Nature 401.6755 (1999): 788-791.
#
# For a given $m \times n$ matrix $X \ge 0$, finding the factor matrices $D \in \mathbb{R}^{m \times k}, \Lambda \in \mathbb{R}^{k \times n}$ such that $X \approx D \Lambda$ with $D \ge 0, \Lambda \ge 0$ is known as the rank-$k$ non-negative matrix factorization (NMF) problem. Typically, such an approximate representation is sought by solving a minimisation problem
# ```math
# \min_{D \in \mathbb{R}^{m \times k}, \Lambda \in \mathbb{R}^{k \times n}, D \ge 0, \Lambda \ge 0} \Phi(X, D \Lambda),
# ```
# where $(X, Y) \mapsto \Psi(X, Y)$ is a loss function defined on matrices. Commonly used choices of $\Phi$ include the squared Frobenius loss $\Phi(X, Y) = \| X - Y \|_F^2$ and Kullback-Leibler divergence $\Phi(X, Y) = \sum_{ij} X_{ij} \log(X_{ij}/Y_{ij})$.
#
# Both these loss functions (and many other common choices) decompose elementwise in their arguments, that is, they can be written as $\Phi(X, Y) = \sum_{ij} f(X_{ij}, Y_{ij})$ for some function $f$ acting on scalars.
#
# Rolet et al. note that pointwise loss functions cannot account for spatial correlations in datasets with underlying geometry, and propose to use entropy-regularised optimal transport as a loss function that is sensitive to the spatial arrangement of the data.
# They argue that for datasets such as images, optimal transport is a more natural choice of loss function, and that it achieves superior performance.
#
# In particular, suppose that the columns of $X$ encode image data, and that $C \in \mathbb{R}^{m \times m}$ encodes the squared Euclidean distances on the imaging domain. The problem that Rolet et al. pose is
# ```math
# \min_{D \in \mathbb{R}^{m \times k}_{\ge 0}, \Lambda \in \mathbb{R}^{k \times n}_{\ge 0}} \sum_{i = 1}^{n} \operatorname{OT}_{\varepsilon}(X_i, (D\Lambda)_i) + \rho_1 E(D) + \rho_2 E(\Lambda),
# ```
# where $\operatorname{OT}_{\varepsilon}(\alpha, \beta)$ is the entropy-regularised optimal transport loss between two probability distributions $\alpha, \beta$ for a cost $C$, and $E$ is an entropy barrier function (a smooth approximation to the non-negativity constraint),
# ```math
# E(A) = \sum_{ij} (A_{ij} \log(A_{ij}) - 1).
# ```
# The parameters $\rho_1, \rho_2$ control how "sharp" the non-negativity constraints are. As $\rho_1, \rho_2 \to 0$, the smoothed constraint approaches the hard non-negativity constraint. Finally, $\varepsilon$ controls the regularisation level for the optimal transport loss.
#
# This example shows how this method can be implemented using functions from the `Dual` submodule of OptimalTransport.jl.
#
# ## Set up prerequisites
#
# Load packages that will be required later.
#
using OptimalTransport
import OptimalTransport.Dual: Dual
using MLDatasets: MLDatasets
using StatsBase
using Plots;
default(; palette=:Set1_3)
using LogExpFunctions
using NNlib: NNlib
using LinearAlgebra
using Distances
using Base.Iterators
using NMF
using Optim
#
# Define `simplex_norm!`, which normalises `x` to sum to `1` along `dims`.
#
function simplex_norm!(x; dims=1)
    return x .= x ./ sum(x; dims=dims)
end
#
# ## Implementation
#
# We now implement Wasserstein-NMF. Rolet et al. split the original non-convex problem into a pair of convex problems, one for $D$ and one for $\Lambda$.
# ```math
# \begin{aligned}
# &\min_{D \in \mathbb{R}^{m \times k}} \sum_{i = 1}^{n} \operatorname{OT}_{\varepsilon}(X_i, (D\Lambda)_i) + \rho_2 E(D), \\
# &\min_{\Lambda \in \mathbb{R}^{k \times n}} \sum_{i = 1}^{n} \operatorname{OT}_{\varepsilon}(X_i, (D\Lambda)_i) + \rho_1 E(\Lambda).
# \end{aligned}
# ```
# For each of these problems, the dual problem can be derived (see Section 3.3 of Rolet et al. for details on how this is done). These turn out to be
# ```math
# \begin{aligned}
# &\min_{G \in \mathbb{R}^{m \times n}} \sum_{i = 1}^n \operatorname{OT}_{\varepsilon}^*(X_i, G_i) + \rho_2 \sum_{i = 1}^{n} E^*(-(G \Lambda^\top)_i / \rho_2) \\
# &\min_{G \in \mathbb{R}^{m \times n}} \sum_{i = 1}^n \operatorname{OT}_{\varepsilon}^*(X_i, G_i) + \rho_1 \sum_{i = 1}^{n} E^*(-(D^\top G)_i / \rho_1).
# \end{aligned}
# ```
# The semi-dual of entropy-regularised optimal transport loss, $\operatorname{OT}_{\varepsilon}$, is implemented as `Dual.ot_entropic_semidual` and its gradient can be computed by `Dual.ot_entropic_semidual_grad`. $E^*$ turns out to be `logsumexp`, for which we implement a wrapper function:
#
function E_star(x; dims=1)
    return logsumexp(x; dims=dims)
end;
#
# The gradient of `logsumexp` is `softmax`, so we define also its gradient:
#
function E_star_grad(x; dims=1)
    return NNlib.softmax(x; dims=1)
end;
#
# Thus, for each problem we may define the dual objective and its gradient. We note that `ot_entropic_semidual(_grad)` automatically broadcasts along columns of its input. There is therefore no need to make multiple calls to the function, thus allowing for more efficient evaluation.
#
function dual_obj_weights(X, K, ε, D, G, ρ1)
    return sum(Dual.ot_entropic_semidual(X, G, ε, K)) + ρ1 * sum(E_star(-D' * G / ρ1))
end
function dual_obj_weights_grad!(∇, X, K, ε, D, G, ρ1)
    return ∇ .= Dual.ot_entropic_semidual_grad(X, G, ε, K) - D * E_star_grad(-D' * G / ρ1)
end
function dual_obj_dict(X, K, ε, Λ, G, ρ2)
    return sum(Dual.ot_entropic_semidual(X, G, ε, K)) + ρ2 * sum(E_star(-G * Λ' / ρ2))
end
function dual_obj_dict_grad!(∇, X, K, ε, Λ, G, ρ2)
    return ∇ .= Dual.ot_entropic_semidual_grad(X, G, ε, K) - E_star_grad(-G * Λ' / ρ2) * Λ
end;
#
# The only remaining part of Wasserstein-NMF to implement is the conversion at optimality from the dual variable $G$ to the primal variables $D, \Lambda$. From the results of Theorems 3 and 4 in Rolet et al., we have for $\Lambda$:
# ```math
# \Lambda_i = \operatorname{softmax}((-D^\top G)_i / \rho_1),
# ```
# and we have for $D$:
# ```math
# D_i = \operatorname{softmax}((-G \Lambda^\top)_i / \rho_2).
# ```
#
function getprimal_weights(D, G, ρ1)
    return NNlib.softmax(-D' * G / ρ1; dims=1)
end
function getprimal_dict(Λ, G, ρ2)
    return NNlib.softmax(-G * Λ' / ρ2; dims=1)
end;
#
# We can now implement functions `solve_weights` and `solve_dict` that solve the respective dual problems for the next iterates of `Λ` and `D`.
#
function solve_weights(X, K, ε, D, ρ1; alg, options)
    opt = optimize(
        g -> dual_obj_weights(X, K, ε, D, g, ρ1),
        (∇, g) -> dual_obj_weights_grad!(∇, X, K, ε, D, g, ρ1),
        zero.(X),
        alg,
        options,
    )
    return getprimal_weights(D, Optim.minimizer(opt), ρ1)
end
function solve_dict(X, K, ε, Λ, ρ2; alg, options)
    opt = optimize(
        g -> dual_obj_dict(X, K, ε, Λ, g, ρ2),
        (∇, g) -> dual_obj_dict_grad!(∇, X, K, ε, Λ, g, ρ2),
        zero.(X),
        alg,
        options,
    )
    return getprimal_dict(Λ, Optim.minimizer(opt), ρ2)
end;

#
# ## Example: noisy univariate Gaussians
#
# We set up each observation as a mixture of 3 Gaussians with means sampled from `N(6, σ), N(0, σ), N(-6, σ)` respectively, and mixture weights sampled uniformly from `[0, 1]`. The resulting mixture model is normalised to sum to 1 on the discrete domain `coord`.
#
f(x, μ, σ) = exp.(-(x .- μ) .^ 2)
coord = range(-12, 12; length=100)
N = 100
σ = 1
X = hcat(
    [
        rand() * f(coord, σ * randn() + 6, 1) +
        rand() * f(coord, σ * randn(), 1) +
        rand() * f(coord, σ * randn() - 6, 1) for _ in 1:N
    ]...,
)
X = simplex_norm!(X);
#
# We visualise the  observations.
#
plot(coord, X; alpha=0.1, color=:blue, title="Input data", legend=nothing)
#
# We can apply NMF with a squared Frobenius loss using the NMF.jl package. We seek `k = 3` components.  This performs poorly, since the pointwise nature of the loss function cannot handle the translational noise in the data.
#
k = 3
result = nnmf(X, k; alg=:multmse)
plot(coord, result.W; title="NMF with Frobenius loss", palette=:Set1_3)
#
# We can now set up a cost matrix corresponding to the domain `coord`.
#
C = pairwise(SqEuclidean(), coord)
C = C / mean(C);
# Specify parameters
ε = 0.025
ρ1, ρ2 = (5e-2, 5e-2);
# Compute Gibbs kernel
K = exp.(-C / ε);
#
# Now we use a random initialisation, where columns of `D` and `Λ` are normalised to sum to 1.
#
D = rand(size(X, 1), k) # dictionary
simplex_norm!(D; dims=1) # norm columnwise
Λ = rand(k, size(X, 2)) # weights
simplex_norm!(Λ; dims=1); # norm rowwise
#
# We now run 10 iterations of Wasserstein-NMF.
#
n_iter = 10
for iter in 1:n_iter
    @info "Wasserstein-NMF: iteration $iter"
    D .= solve_dict(
        X,
        K,
        ε,
        Λ,
        ρ2;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
    Λ .= solve_weights(
        X,
        K,
        ε,
        D,
        ρ1;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
end
#
# We observe that Wasserstein-NMF learns atoms (columns of $D$) corresponding to the three Gaussians used to generate the input data.
#
plot(coord, D; title="NMF with Wasserstein loss", palette=:Set1_3)

# ## Example: image data (MNIST)
# Here we will download MNIST dataset using MLDatasets.jl and downscale each image to 14x14 to allow for faster runtime (since we are running on CPU).
#
sizex, sizey = 28, 28
factor = 2 # downscale factor
Σ = hcat([sum(I(sizex)[:, i:(i + factor - 1)]; dims=2) for i in 1:factor:sizex]...)
sizex, sizey = sizex ÷ factor, sizey ÷ factor
N = 256
x, y = MLDatasets.MNIST.traindata(Float64, sample(1:60_000, N; replace=false))
x = permutedims(x, (2, 1, 3))
x = cat([Σ' * x[:, :, i] * Σ for i in 1:N]...; dims=3)
X = simplex_norm!(reshape(x, (sizex * sizey, :)));
#
# The columns of `X` now correspond to images "flattened" as vectors. We can preview a few images.
#
M = 25
plot(
    [
        heatmap(
            reshape(X[:, i], sizex, sizey);
            legend=:none,
            axis=nothing,
            showaxis=false,
            aspect_ratio=:equal,
            c=:Blues,
            yflip=true,
        ) for i in 1:M
    ]...;
    layout=(5, M ÷ 5),
    plot_title="Input images",
)

# Now we can set up `coord`, cost matrix `C`, and specify parameters.
#
coord = reshape(collect(product(1:sizex, 1:sizey)), :)
C = pairwise(SqEuclidean(), coord)
C = C / mean(C);
ε = 0.0025
ρ1, ρ2 = (5e-3, 5e-3);
# We compute the Gibbs kernel from `C`:
K = exp.(-C / ε);
#
# Let us aim to learn `k = 25` atoms.
k = 25;
#
# Initialise again randomly
#
D = rand(size(X, 1), k) # dictionary
simplex_norm!(D; dims=1) # norm columnwise
Λ = rand(k, size(X, 2)) # weights
simplex_norm!(Λ; dims=1); # norm rowwise
#
# We now run 15 iterations of Wasserstein-NMF.
#
n_iter = 15
for iter in 1:n_iter
    @info "Wasserstein-NMF: iteration $iter"
    D .= solve_dict(
        X,
        K,
        ε,
        Λ,
        ρ2;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
    Λ .= solve_weights(
        X,
        K,
        ε,
        D,
        ρ1;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
end
#
# We can inspect the atoms learned (columns of `D`):
#

plot(
    [
        heatmap(
            reshape(D[:, i], sizex, sizey);
            legend=:none,
            axis=nothing,
            showaxis=false,
            aspect_ratio=:equal,
            c=:Blues,
            yflip=true,
        ) for i in 1:k
    ]...;
    layout=(5, k ÷ 5),
    plot_title="Learned atoms",
)

#
# Finally, we can look at the images constructed by the low-rank approximation `DΛ` and compare to the original images that we previewed earlier.
#
X_hat = D * Λ
plot(
    [
        heatmap(
            reshape(X_hat[:, i], sizex, sizey);
            legend=:none,
            axis=nothing,
            showaxis=false,
            aspect_ratio=:equal,
            c=:Blues,
            yflip=true,
        ) for i in 1:M
    ]...;
    layout=(5, M ÷ 5),
    plot_title="Low rank approximation",
)


using JuMP
using GLPK
using GraphPlot
using LightGraphs
using Colors
using LinearAlgebra

G = [
    0 1 0 0 0 0;
    1 0 1 1 0 0;
    0 1 0 0 1 1;
    0 1 0 0 1 0;
    0 0 1 1 0 0;
    0 0 1 0 0 0;
]

g = SimpleGraph(G)

gplot(g)

vertex_cover = Model(GLPK.Optimizer)

@variable(vertex_cover, y[1:nv(g)], Bin)
@constraint(vertex_cover, [i = 1:nv(g), j = 1:nv(g); G[i,j] == 1], y[i] + y[j] >= 1)
@objective(vertex_cover, Min, sum(y))

optimize!(vertex_cover)
@show value.(y);

membership = round.(Int, value.(y)) # Change to Int
membership = membership + ones(Int, nv(g)) # Make the color groups one indexed
nodecolor = [colorant"red", colorant"blue"] # Blue to represent vertices in the cover
nodefillc = nodecolor[membership]
gplot(g, nodefillc = nodefillc)

G = [
    0 1 0 0 0 0 0 0 0 1 0 ;
    1 0 1 0 0 0 0 0 0 0 1;
    0 1 0 1 0 1 0 0 0 0 0;
    0 0 1 0 1 0 0 0 0 0 0;
    0 0 0 1 0 1 0 0 0 0 0;
    0 0 1 0 1 0 1 0 0 0 0;
    0 0 0 0 0 1 0 1 0 0 0;
    0 0 0 0 0 0 1 0 1 0 1;
    0 0 0 0 0 0 0 1 0 1 1;
    1 0 0 0 0 0 0 0 1 0 1;
    0 1 0 0 0 0 0 1 1 1 0;
]

g = SimpleGraph(G)

gplot(g)

dominating_set = Model(GLPK.Optimizer)

@variable(dominating_set, x[1:nv(g)], Bin)
@constraint(dominating_set, [i = 1:nv(g)], dot(G[i,:], x) >= 1)
@objective(dominating_set, Min, sum(x))

optimize!(dominating_set)
@show value.(x);

membership = convert(Array{Int},value.(x)) # Change to Int
membership = membership + ones(Int, nv(g)) # Make the color groups one indexed
nodecolor = [colorant"red", colorant"blue"] # Blue to represent vertices in the set
nodefillc = nodecolor[membership]
gplot(g, nodefillc = nodefillc)

G = [
    0 0 0 0 1 0 0 0;
    0 0 0 0 0 1 0 0;
    0 0 0 0 0 0 1 0;
    0 0 0 0 0 0 0 1;
    1 0 0 0 0 1 0 1;
    0 1 0 0 1 0 1 0;
    0 0 1 0 0 1 0 1;
    0 0 0 1 1 0 1 0;
]

g = SimpleGraph(G)

gplot(g)

matching = Model(GLPK.Optimizer)

@variable(matching, m[i = 1:nv(g), j = 1:nv(g)], Bin)
@constraint(matching, [i = 1:nv(g)], sum(m[i,:]) <= 1)
@constraint(matching, [i = 1:nv(g), j = 1:nv(g); G[i,j] == 0], m[i,j] == 0)
@constraint(matching, [i = 1:nv(g), j = 1:nv(g)], m[i,j] == m[j,i])
@objective(matching, Max, sum(m))

optimize!(matching)
@show value.(m);

G = [
    0 1 0 0 1 1 0 0 0 0;
    1 0 1 0 0 0 1 0 0 0;
    0 1 0 1 0 0 0 1 0 0;
    0 0 1 0 1 0 0 0 1 0;
    1 0 0 1 0 0 0 0 0 1;
    1 0 0 0 0 0 1 0 0 1;
    0 1 0 0 0 1 0 1 0 0;
    0 0 1 0 0 0 1 0 1 0;
    0 0 0 1 0 0 0 1 0 1;
    0 0 0 0 1 1 0 0 1 0;
]

g = SimpleGraph(G)

gplot(g)

k = nv(g)

k_colouring = Model(GLPK.Optimizer)

@variable(k_colouring, z[1:k], Bin)
@variable(k_colouring, c[1:nv(g),1:k], Bin)
@constraint(k_colouring, [i = 1:nv(g)], sum(c[i,:]) == 1)
@constraint(k_colouring, [i = 1:nv(g), j = 1:nv(g), l = 1:k; G[i,j] == 1], c[i,l] + c[j,l] <= 1)
@constraint(k_colouring, [i = 1:nv(g), l = 1:k], c[i,l] <= z[l])

@objective(k_colouring, Min, sum(z))

optimize!(k_colouring)
@show value.(z);
@show value.(c);

c = value.(c)
membership = zeros(nv(g))
for i in 1:nv(g)
    for j in 1:k
        if c[i,j] == 1
            membership[i] = j
            break
        end
    end
end
membership = convert(Array{Int},membership)

nodecolor = distinguishable_colors(nv(g), colorant"green")
nodefillc = nodecolor[membership]
gplot(g, nodefillc = nodefillc)


using OrdinaryDiffEq
using LabelledArrays

function f(out,du,u,p,t)
  out.x = - 0.04u.x              + 1e4*u.y*u.z - du.x
  out.y = + 0.04u.x - 3e7*u.y^2 - 1e4*u.y*u.z - du.y
  out.z = u.x + u.y + u.z - 1.0
end

u₀ = LVector(x=1.0, y=0.0, z=0.0)
du₀ = LVector(x=-0.04, y=0.04, z=0.0)
tspan = (0.0,100000.0)

differential_vars = LVector(x=true, y=true, z=false)
prob = DAEProblem(f,du₀,u₀,tspan,differential_vars=differential_vars)

sol = solve(prob, DImplicitEuler())

function f1(du,u,p,t)
   du.x .= -1 .* u.x .* u.y .* p[1]
   du.y .= -1 .* u.y .* p[2]
end
const n = 1000
u_0 = @LArray  fill(1000.0,2*n) (x = (1:n),y = (n+1:2*n))
p = [0.1,0.1]
prob1 = ODEProblem(f1,u_0,(0,100.0),p)
sol = solve(prob1, Rodas5())
sol = solve(prob1, Rodas5(autodiff=false))


#=
Given an undirected graph represented as an adjacency matrix and an integer k, write a function to determine whether each vertex in the graph can be colored such that no two adjacent vertices share the same color using at most k colors.
=#

using Test
include("Solutions/problem56_color_graph_k_colors.jl")

graph = [0 1 1 0;1 0 0 0;1 0 0 1;0 0 1 0]

@test colorable(graph, 1) == false
@test colorable(graph, 2)
@test colorable(graph, 3)

graph = [0 1 1 0 1;1 0 0 0 1;1 0 0 1 0;0 0 1 0 1;1 1 0 1 0]

@test colorable(graph, 1) == false
@test colorable(graph, 2) == false
@test colorable(graph, 3)


"""

```
sphereradiusfromshellvolume(volume,step)
```

Computes the radius that corresponds to a spherical shell of
a given volume.

"""
function sphereradiusfromshellvolume(volume,step)

  fourthirdsofpi = 4*pi/3
  if 3*step*volume - pi*step^4 <= 0.
    return 0.
  end
  rmin = (sqrt(3*pi)*sqrt(3*step*volume-pi*step^4)-3*pi*step^2)/(6*pi*step)
  return ( 0.5*( volume/fourthirdsofpi + 2*rmin^3 ) )^(1/3)

end


# load diabetes.data

using LinearAlgebra
using CSV
using DataFrames
using Plots

Diabete_Data = CSV.read("diabetes.csv",DataFrame)

n,d = size(Diabete_Data)
d = d-1

#println(Diabete_Data)

# compute the optimal theta

X = zeros(n,d)
y = zeros(n,1)

for i=1:d
    X[:,i] = Diabete_Data[:,i]
end
#X = [ones(n,1)  X[:,3]]
X = [ones(n,1) X]
y = Diabete_Data[:,end]

theta_opt = X\y

# achieved loss
MSE = norm(X*theta_opt-y,2)^2/n

println("theta_opt: ", theta_opt)
println("RMSE: ",sqrt(MSE))

# plot
figure()
plot(y,y,"k")
plot(y,X*theta_opt,"o",alpha=0.4)
xlabel(L"$y$")
ylabel(L"predicted $y$")
title("prediction accuracy")
axis("square")
grid("on")

figure()
hist(y-X*theta_opt, bins=10, density=true)
title("prediction error")
grid("on")

# put JHK data into your predictor

#age
#sex
#bmi
#map	   mean arterial pressure
#s1  tc :   total cholesterol
#s2  ldl    low density lipoprotein
#s3  hdl    high density lipoprotein
#s4  tch
#s5  ltg
#s6  glu
#        age  sex  bmi    map   tc     ldl    hdl     tch   ltg   glu
X_JHK = [41   1    18.3   90    171    80.0   74.9    2     4.75  90.0]
X_JHK = [1 X_JHK]

y_JHK = X_JHK*theta_opt

# plot

figure()
plot(y,y,"k")
plot(y,X*theta_opt,"o",alpha=0.4)
plot(y_JHK,y_JHK,"ro",alpha=0.4)
xlabel(L"$y$")
ylabel(L"predicted $y$")
axis("square")
grid("on")


"""
        game(decisionMatrix; verbose = false)

    Solve the zero-sum game.

# Arguments:
 - `decisionMatrix::Array{Float,2}`: n × m matrix of gain values for row player.
 - `verbose::Bool`: Logical value indicating whether the optimization logs are shown or not. Default is false.

# Description
game() solves zero-sum games using a gain matrix designed for the row player. The game matrix has n rows
and m columns for n and m strategies of row and column players, respectively.

# Output
- `::GameResult`: GameResult object that holds mixed strategy probabilities for row player and the value of game.
If a pure strategy exists for the row player, than the propabilities vector is a one-hat vector.

# Examples
```julia-repl
julia> # Rock & Paper & Scissors game
julia> mat = [0 -1 1; 1 0 -1; -1 1 0]
3×3 Array{Int64,2}:
  0  -1   1
  1   0  -1
 -1   1   0

julia> dm = makeDecisionMatrix(mat)
3×3 DataFrame
 Row │ Crt1     Crt2     Crt3
     │ Float64  Float64  Float64
─────┼───────────────────────────
   1 │     0.0     -1.0      1.0
   2 │     1.0      0.0     -1.0
   3 │    -1.0      1.0      0.0

julia> result = game(dm);

julia> result.row_player_probabilities
3-element Array{Float64,1}:
 0.3333333333333333
 0.33333333333333337
 0.3333333333333333

julia> result.value
0.0

```

# References
Zhou, Hai-Jun. "The rock–paper–scissors game." Contemporary Physics 57.2 (2016): 151-163.
"""
    function game(decisionMatrix::DataFrame; verbose::Bool=false)::GameResult

    newDecisionMatrix = similar(decisionMatrix)

    nrow, ncol = size(decisionMatrix)

    minmat = minimum(decisionMatrix)
    modified = false

    if minmat < 0
        modified = true
        newDecisionMatrix = decisionMatrix .- minmat
    end

    # dm = convert(Matrix, newDecisionMatrix)
    dm = Matrix(newDecisionMatrix)

    model = Model(GLPK.Optimizer);
    MOI.set(model, MOI.Silent(), !verbose)

    @variable(model, x[1:nrow])
    @objective(model, Min, sum(x))

    for i in 1:nrow
        @constraint(model, sum(x[1:nrow] .* dm[:, i]) >= 1)
    end

    for i in 1:nrow
        @constraint(model, x[i] >= 0)
    end

    optimize!(model);

    values = JuMP.value.(x)

    gamevalue = 1 / objective_value(model)
    p = gamevalue * values

    if modified
        gamevalue -= abs(minmat)
    end

result = GameResult(
        p,
        gamevalue
    )

    return result

end


"""
    ways_place_containers(x,m)

Given containers of sizes x, and m balls, enumerate all the ways to place the balls
into the containers without violating the constraints.

Lets illustrate this function through an example. Suppose we want to find
all the ways to place 4 balls into 4 containers with capacities [2,1,3,4]

We define f(m,e) as the number of solutions of placing m balls into the
e leftmost containers. Then we could build a table, vertically m and horizontally e.

The base cases are
m  e|   1   2   3   4
--------------------------
0   |   1   1   1   1
1   |   1
2   |   1
3   |   0
4   |   0

And then we could fill in the rest of the table using the recursive formula
f(m,e)=∑f(m-i,e-1)  where i=0:min(m,c[e]), where c[e] is the capacity of the e^th container.

m  e|   1   2   3   4
--------------------------
0   |   1   1   1   1
1   |   1   2   3   4
2   |   1   2   5   9
3   |   0   1   6   15
4   |   0   0   5   20

Then we could use this table to build the actual solutions with another builder function.
"""
function ways_place_containers(x,m)
    #assume x=[x_1,x_2,x_3...x_n] respresents n containers each with
    #positive capacity x_i. Now we want to place m balls into these containers
    #without breaking the capacity
    #this solution will use a dynamic programming approach.
    #First, find the number of solutions, then allocate an array, then
    #build that table using the solution table and a builder function.

    #base cases
    solution_table = zeros(Int64,m+1,length(x))
    solution_table[1,1:end] .= 1
    for r = 2:m+1
        solution_table[r,1] = r-1 <= x[1]
    end

    #build solution table
    for c=2:length(x)
        for r=2:m+1
            solution_table[r,c] = sum(solution_table[r-min(r-1,x[c]):r,c-1])
        end
    end

    ways = zeros(Int64,solution_table[end,end],length(x))
    table_builder!(ways,solution_table,x,1,length(x),m)
    return ways
end

function table_builder!(output,solution_table,container_sizes,row,col,balls_left)
    (balls_left <= 0) && return
    r1 = row
    for balls_to_place = 0:min(balls_left,container_sizes[col])
        if col==1
            output[row,col] = balls_left
            continue
        end
        ways = solution_table[balls_left + 1 - balls_to_place,col-1]
        ways <= 0 && continue
        output[r1:r1+ways-1,col] .= balls_to_place
        table_builder!(output,solution_table,container_sizes,r1,col-1,balls_left - balls_to_place)

        r1 = r1 +ways
    end
end

"""
    canonical_placement(x,y)

Given containers of size x, and number of balls to place into each container y,
return the placement as a tuple of 1s and 0s.

Example:

canonical_placement([2,3,4],[1,1,1]) returns

(1,0,1,0,0,1,0,0,0)

Which one could think of as 10|100|1000

"""
function canonical_placement(x,y)
    #assume x=[x_1,x_2,x_3...x_n] respresents n containers each with positive capacity x_i.
    #assume y=[y_1,y_2,y_3...y_n] means placed 0<=y_i<=x_i balls into container i
    #returns one way this could have happened.
    #For example, x=[1,2,6],y=[1,1,2] could be 1|10|010100 or 1|01|110000
    tmp = zeros(Int,sum(x))
    ind = 1
    for i = 1:length(x)
        tmp[ind:ind+y[i]-1] .= 1
        ind += x[i]
    end
    return tmp
end

function count_occurrences(factor)
    #assume factor is a sorted tuple like (0,0,0,1,1,3,3)
    #returns two arrays, the first one keys, the second one
    #number of occurrences, like for this case,
    #[0,1,3],[3,2,2]
    x = [factor[1]]
    y = [1]
    for i=2:length(factor)
        if factor[i] == factor[i-1]
            y[end] += 1
        else
            push!(x,factor[i])
            push!(y,1)
        end
    end
    return x,y
end


module problem123

description = """
Prime square remainders

Let pn be the nth prime: 2, 3, 5, 7, 11, ..., and let r be the remainder when (pn−1)^n + (pn+1)^n is divided by pn^2.
For example, when n = 3, p3 = 5, and 4^3 + 6^3 = 280 ≡ 5 mod 25.
The least value of n for which the remainder first exceeds 10^9 is 7037.

Find the least value of n for which the remainder first exceeds 10^10.
"""

# Computes a^b % n
# Definition from Wikipedia on Modular exponentiation: http://en.wikipedia.org/wiki/Modular_exponentiation
# BigInt is required, as a*a will overflow for large a and n
function powmod(a :: BigInt, b :: Integer, n :: Integer)
  x = BigInt(1)
  a = a % n
  while b > 0
    if (b % 2) == 1
      x = (x * a) % n
    end
    b = b >> 1
    a = (a * a) % n
  end
  return x
end

function solve(rmin = 10^10, maxprime = 10^6)
  ps = primes(maxprime)
  for n = 1:length(ps)
    p = BigInt(ps[n])
    psq = p * p
    if psq < rmin
      continue
    end
    r = (powmod(p-1, n, psq) + powmod(p+1, n, psq)) % psq
    if r > rmin
      return n
    end
  end
  error("No solution for rmin=$rmin maxprime=$maxprime")
end

using Base.Test

@test (powmod(BigInt(4), 3, 25) + powmod(BigInt(6), 3, 25)) % 25 == 5
@test solve(10^9) == 7037

end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Scale(; low=0.25, high=0.75)

The scale transform of `x` is the value `(x .- xl) ./ (xh .- xl))`
where `xl = quantile(x, low)` and `xh = quantile(x, high)`.
"""
struct Scale <: Colwise
  low::Float64
  high::Float64

  function Scale(low, high)
    @assert 0 ≤ low ≤ high ≤ 1 "invalid quantiles"
    new(low, high)
  end
end

Scale(; low=0.25, high=0.75) = Scale(low, high)

assertions(::Type{Scale}) = [assert_continuous]

isrevertible(::Type{Scale}) = true

function colcache(transform::Scale, x)
  levels = (transform.low, transform.high)
  xl, xh = quantile(x, levels)
  xl == xh && ((xl, xh) = (zero(xl), one(xh)))
  (xl=xl, xh=xh)
end

colapply(::Scale, x, c)  = @. (x - c.xl) / (c.xh - c.xl)

colrevert(::Scale, y, c) = @. (c.xh - c.xl) * y + c.xl

"""
    MinMax()

The transform that is equivalent to `Scale(low=0, high=1)`.
"""
MinMax() = Scale(low=0.0, high=1.0)

"""
    Interquartile()

The transform that is equivalent to `Scale(low=0.25, high=0.75)`.
"""
Interquartile() = Scale(low=0.25, high=0.75)


export sample_banana, log_pdf_banana

"""
    sample_banana(N; μ = 0.0, σ = 2.0, bananicity = 0.2)

Generate `N` samples [x₁; x₂] of the Banana distribution defined as:
    x₁ ∼ N(μ, σ²),
    ϵ  ∼ N(0, 1),
    x₂ ∼ bananicity × (x₁ - σ²) + ϵ.
"""
function sample_banana(N; μ = 0.0, σ = 2.0, bananicity = 0.2)
    X = zeros(2,N)
    for i=1:N
        X[1,i] = μ + σ*randn()
        X[2,i] = bananicity*(X[1,i]^2 - σ^2) + randn()
    end
    return X
end

# log_pdf_banana(X) = log_pdf(X[1]) + log_pdf(X[2] - X[1]^2)
# log_pdf_banana(X; a = 0.0, b = 1.0, c = 1.0) = log_pdf((X[1] -a)/b) + log_pdf(X[2] - c*X[1]^2)

"""
    log_pdf_banana(X; μ = 0.0, σ = 2.0, bananicity = 0.2)

Compute the logarithm of the porbability density function (pdf) of the Banana distribution at `X`.
See documentation of sample_banana for details on the Banana distribution.
"""
log_pdf_banana(X::AbstractVector; μ = 0.0, σ = 2.0, bananicity = 0.2) = log_pdf((X[1]-μ)/σ) + log_pdf(X[2] - bananicity*(X[1]^2 - σ^2))


@testset "_unit_ball_volume" begin
    @test MinimumVolumeEllipsoids._unit_ball_volume(0) == 1.0
    @test MinimumVolumeEllipsoids._unit_ball_volume(1) ≈ 2.0 atol = 1e-10
    @test MinimumVolumeEllipsoids._unit_ball_volume(2) ≈ π atol = 1e-10
    @test MinimumVolumeEllipsoids._unit_ball_volume(3) ≈ 4π / 3 atol = 1e-10
    @test MinimumVolumeEllipsoids._unit_ball_volume(4) ≈ π^2 / 2 atol = 1e-10
    @test MinimumVolumeEllipsoids._unit_ball_volume(5) ≈ 8π^2 / 15 atol = 1e-10
end

@testset "volume" begin
    @test volume(Ellipsoid(PDMat(2 * Matrix(I, 2, 2)), zeros(2))) ≈ π atol = 1e-10
    @test volume(Ellipsoid(PDMat(3 * Matrix(I, 3, 3)), zeros(3))) ≈ 4π / 3 atol = 1e-10
    @test volume(Ellipsoid(PDMat(4 * Matrix(I, 4, 4)), zeros(4))) ≈ π^2 / 2 atol = 1e-10
end


using LinearAlgebra

# Base transition matrix
T = [-1 1 0 0; 0 -1 1 0; 0 0 -1 1; 1 0 0 -1]

Λ, V = eigen(T)
Λ, W = eigen(T')

# Right eigenvectors
V⁴ = round.(V[:, 1] ./ V[1, 1])
V³ = round.(V[:, 2] ./ V[1, 2])
V² = round.(V[:, 3] ./ V[1, 3])
V¹ = round.(V[:, 4] ./ V[1, 4])

# Left eigenvectors
W⁴ = round.(W[:, 1] ./ W[1, 1])
W³ = round.(W[:, 2] ./ W[1, 2])
W² = round.(W[:, 3] ./ W[1, 3])
W¹ = round.(W[:, 4] ./ W[1, 4])

# Flip Order (take into account later)
W[:, 1] .= W¹
W[:, 2] .= W²
W[:, 3] .= W³
W[:, 4] .= W⁴
inv(W) * 4

# Quick Check (right eigenvector)
T * V - V * Diagonal(Λ)
maximum(abs.(T * V - V * Diagonal(Λ))) ≤ eps(10.0)
# Quick Check (left eigenvector)
T' * W - W * Diagonal(reverse(Λ))    # since we flipped the order so that index 1 corresponds to the lowest eigenvalue
W' * T - Diagonal(reverse(Λ))' * W'  # same thing, recall that ' means adjoint, thus takes the conjucate
maximum(abs.(T' * W - W * Diagonal(reverse(Λ)))) ≤ eps(10.0)

# The following transformation is what is used
S = real.(V)
S[:, 1] .= real.(W¹)
S[:, 2] .= real.(0.5 * (W² + W³))
S[:, 3] .= real.(0.5 * im * (W² - W³))
S[:, 4] .= real.(W⁴)

display(S)
display(inv(S))
display(inv(S) * T * S)

# The following partitions the magnitude of
# the anti-symmetric component and the symmetric component
# of the transition matrix
sT = (T + T') / 2
aT = (T - T') / 2
λ = 1.0
ω = 1.0
inv(S) * (λ * sT + ω * aT) * S

## check kronecker products corresponding to random phase
kroneckered_phase = kron(sT, I + 0 * sT) + kron(I + 0 * sT, sT)
Λ, V = eigen(kroneckered_phase)

## check out higher dimensional one
idmat = I + 0 * sT # the identity matrix
kroneckered_phase = kron(sT, idmat, idmat) + kron(idmat, sT, idmat) + kron(idmat, idmat, sT)
Λ, V = eigen(kroneckered_phase)


"""
    Chebyshev

A functor struct representing the basis of Chebyshev polynomials.
"""
struct Chebyshev end

"""
    chebyshev_basis(k,x)

Evaluate the `k`-th Chebyshev basis function at `x`.
"""
function chebyshev_basis(k,x)
    if abs(x) < 1
        cos(k * acos(x))
    elseif x >= 1
        cosh(k * acosh(x))
    else
        (-1)^k * cosh(k * acosh(-x))
    end
end

chebyshev_basis(k,x::AbstractVector) = chebyshev_basis.(k,x)

(::Chebyshev)(k,x) = chebyshev_basis(k,x)


@testset "Matrix" begin

    @testset "Matrix: Determinant" begin
        M1 = [1 0; 2 2]
        M2 = rand(3,3)
        M3 = rand(4,4)

        @test determinant(M1) == det(M1)
        @test round(determinant(M2),digits = 4) == round(det(M2),digits = 4)
        @test round(determinant(M3),digits = 4) == round(det(M3),digits = 4)
    end

    @testset "Matrix: LU Decompose" begin
        mat = [
            2  -1 -2;
            -4   6  3;
            -4  -2  8
        ]

        L = [
            1  0 0;
            -2  1 0;
            -2 -1 1
        ]
        U = [
            2 -1 -2;
            0  4 -1;
            0  0  3
        ]

        @test lu_decompose(mat) == (L,U)
    end

    @testset "Matrix: rotation matrix" begin
        theta = pi/6
        a = [1;0]
        b = [0;1]
        @test rotation_matrix(theta)*a == [cos(theta);sin(theta)]
        @test rotation_matrix(theta)*b == [-sin(theta);cos(theta)]
    end

    @testset "Matrix: Gauss-Jordan Elimination" begin
        M1 = [1 2 3; 4 5 6];
        M2 = [1 2 3; 4 8 12];
        M3 = [1 2; 4 8];
        M4 = [1 4 8  1  4
              4 5 6  8  11
              1 3 2  4  8
              4 5 67 23 0];
        M5 = Float64[1 4 8  1  4
                     4 5 6  8  11
                     1 3 2  4  8
                     4 5 67 23 0];
        M6 = [1 4 8
              4 0 6
              1 3 2];

        R3 =  [ 1.0  0.0  0.0  0.0  -1.10637
                0.0  1.0  0.0  0.0   1.89743
                0.0  0.0  1.0  0.0  -0.444913
                0.0  0.0  0.0  1.0   1.07598];

        @test gauss_jordan(M1) == Float64[1 0 -1; 0 1 2]
        @test_throws AssertionError gauss_jordan(M2)
        @test_throws AssertionError gauss_jordan(M3)
        @test isapprox(gauss_jordan(M4), R3, atol=1e-5)
        @test isapprox(gauss_jordan(M5), R3, atol=1e-5)
        @test isapprox(gauss_jordan(M6), Float64[1 0 0; 0 1 0; -0 -0  1], atol=1e-5)
    end

end


module SonarEquations

using SpecialFunctions

"""
	SE = signal_excess_passive(SL, TL, NL, DI, DT)

Calculates the signal excess in a passive sonar scenario for source level `SL` (dB), transmission loss `TL` (dB), noise level `NL` (dB), directivity index `DI` (dB) and detection threshold `DT` (dB).
"""
function signal_excess_passive(SL, TL, NL, DI, DT)
	SE = SL - TL - NL + DI - DT
end

"""
	detection_index_gaussian(p_dtc, p_fal)

Calculates the detection index for a non-fluctuating signal in a Gaussian-distributed noise environment for a probability of detection `p_dtc` and probability of false alarm `p_fal`.
"""
function detection_index_gaussian(p_dtc, p_fal)
	if p_fal ≥ p_dtc
		d = 0
	else
		# d = 2(erfcinv(2p_dtc) - erfcinv(2p_fal))^2
		d = 2(erfinv(1 - 2p_dtc) - erfinv(1 - 2p_fal))^2
	end
	return d
end
# function detection_index_gaussian(p_dtc, p_fal)
# 	d = 2(erfcinv(2p_dtc) - erfcinv(2p_fal))^2
# 	# d = 2(erfinv(1 - 2p_dtc) - erfinv(1 - 2p_fal))^2
# end

# function detection_index_gaussian(SL, TL, NL, B, t)
# 	d = B*t*((SL - TL)/(B*NL))^2
# end
function detection_index_gaussian(SL, TL, NL, B, t)
	if TL ≥ SL
		d = 0
	else
		d = B*t*((SL - TL)/NL)^2
	end
	return d
end

"""

"""
function detection_index_exponential(p_dtc, p_fal)
	d = (log(p_fal)/log(p_dtc) - 1)^2
end

"""

"""
function detection_index_rayleigh(p_dtc, p_fal)
	d = π/(4 - π)*(sqrt(log(p_fal)/log(p_dtc)) - 1)^2
end

"""

"""
function detection_threshold(d, B, t)
	DT = 5log10(d*B/t)
end

"""
	probability_of_detection_gaussian(d, p_fal)

Calculates the probability of detection for a non-fluctuating signal in a Gaussian-distributed noise environment for detection index `d` and probability of false alarm `p_fal`.
"""
function probability_of_detection_gaussian(d, p_fal)
	p_dtc = erfc(erfcinv(2p_fal) - sqrt(d/2))/2
end

"""

"""
function probability_of_detection_exponential(d, p_fal)
	p_dtc = exp(log(p_fal)/(sqrt(d) + 1))
end

"""

"""
function probability_of_detection_rayleigh(d, p_fal)
	p_dtc = exp(log(p_fal)/(sqrt((4 - π)/π*d) + 1)^2)
end

end


# This file includes utilities

export Q, dbtoval, valtodb, energy

"""
    $SIGNATURES

Computes the energy of the discrete time signal `s` sampled with `ts` (defaults to 1) seconds.
"""
energy(s::AbstractVector, ts=1.) = sum(abs.(s).^2) * ts

"""
    $SIGNATURES

Converts `γ` from `dB` value to its real value.

# Example
```julia
julia> dbtoval(2)  # γ = 2 dB
1.5848931924611136
```
"""
dbtoval(γ) = 10^(γ / 10)

"""
    $SIGNATURES

Converts `val` tı dB scale.
"""
valtodb(val) = 10 * log10(val)

"""
    $SIGNATURES

Q-function defined as.
```math
    Q(x) = \\int_{x}^{\\infty} exp(-\\dfrac{x^2}{2}) dx
```
"""
Q(x) = 1 / 2 * erfc(x / sqrt(2))


@testset "Univariate functions" begin
    # bspline0 from Daisy benchmarks
    f = x -> (1 - x)^3 / 6.0
    dom = Interval(-4.5, -0.3)
    xref = Interval(0.36616666666666675, 27.729166666666668)

    for solver in available_solvers
        x = enclose(f, dom, solver)
        rleft, rright = relative_precision(x, xref)
        if solver isa TaylorModelsEnclosure
            # for this example, TaylorModels cannot tightly approximate the lhs
            @test rleft ≤ 30 && rright ≤ 1e-10
        elseif solver isa SumOfSquaresEnclosure
            # solver returns negative -1.06259e-06
            @test rleft ≥ -1e-5 && rright ≤ 1e-5
        else
            @test rleft ≤ 1e-5 && rright ≤ 1e-5
        end
    end
end

@testset "Univariate example the Quickstart Guide" begin
    f(x) = -x^3/6 + 5x
    dom = Interval(1, 4)

    x = enclose(f, dom, NaturalEnclosure())
    xref = Interval(-5.66667, 19.8334)
    rleft, rright = relative_precision(x, xref)
    @test rleft ≤ 1e-5 && rright ≤ 1e-5

    x = enclose(f, dom, TaylorModelsEnclosure(order=4))
    xref = Interval(-4.2783, 12.7084)
    rleft, rright = relative_precision(x, xref)
    @test rleft ≤ 1e-5 && rright ≤ 1e-5

    x = enclose(f, dom, MooreSkelboeEnclosure())
    xref = Interval(4.83299, 10.5448)
    rleft, rright = relative_precision(x, xref)
    @test rleft ≤ 1e-5 && rright ≤ 1e-5

    x = enclose(f, dom, SumOfSquaresEnclosure(backend=SDPA.Optimizer))
    xref = Interval(4.8333, 10.541)
    rleft, rright = relative_precision(x, xref)
    @test rleft ≤ 1e-5 && rright ≤ 1e-5
end

@testset "Test univariate polynomial input" begin
    @polyvar x
    p = -x^3/6 + 5x
    dom = Interval(1, 4)

    x = enclose(p, dom)
    xref = Interval(-5.66667, 19.8334)
    rleft, rright = relative_precision(x, xref)
    @test rleft ≤ 1e-5 && rright ≤ 1e-5
end


using CompScienceMeshes
using LinearAlgebra

o = point(3,4,2)
a = point(4,2,4)
b = point(2,8,2)
c = point(0,4,6)

tet = simplex(o,a,b,c)
qw = quadpoints(tet,4)
print(volume(tet),"==",abs(dot(a-o,cross(b-o,c-o)))/6,"==",abs(sum(w for (p,w) in qw)),"\n")

o,x,y,z = CompScienceMeshes.euclidianbasis(3)

tet = simplex(o,x,2y,3z)
volume(tet)

qw = quadpoints(tet,4)
length(qw)

print(abs(sum(w for (p,w) in qw)-volume(tet)) < 0.001,"\n")
print(sum(w*sin(p[1])*p[2]^2*p[3]^3 for (p,w) in qw))
#https://www.wolframalpha.com/input/?i=int_0^1+int_0^(2-2x)+int_0^(3-3x-1.5y)+sin(x)y^2z^3+dz+dy+dx


using MultivariateStats
using LinearAlgebra
using Test
import Statistics: mean, cov
import Random

@testset "Whitening" begin

    Random.seed!(34568)

    ## data

    d = 3
    n = 5

    X = rand(d, n)
    mval = vec(mean(X, dims=2))
    C = cov(X, dims=2)
    C0 = copy(C)

    emax = maximum(eigvals(C))
    rc = 0.01
    Cr = C + (emax * rc) * Matrix(I, d, d)

    ## cov_whitening

    cf = cholesky(Hermitian(C, :U))
    W = cov_whitening(cf)
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    cf = cholesky(Hermitian(C, :L))
    W = cov_whitening(cf)
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    W = cov_whitening(C)
    @test C == C0
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    W = cov_whitening(C, 0.0)
    @test C == C0
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    W = cov_whitening(C, rc)
    @test istriu(W)
    @test W'Cr * W ≈ Matrix(I, d, d)

    # Whitening from data

    f = fit(Whitening, X)
    W = f.W
    @test isa(f, Whitening{Float64})
    @test mean(f) === f.mean
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)
    @test transform(f, X) ≈ W' * (X .- f.mean)

    f = fit(Whitening, X; regcoef=rc)
    W = f.W
    @test W'Cr * W ≈ Matrix(I, d, d)

    f = fit(Whitening, X; mean=mval)
    W = f.W
    @test W'C * W ≈ Matrix(I, d, d)

    f = fit(Whitening, X; mean=0)
    Cx = (X * X') / (n - 1)
    W = f.W
    @test W'Cx * W ≈ Matrix(I, d, d)

    # invsqrtm

    R = invsqrtm(C)
    @test C == C0
    @test R ≈ inv(sqrt(C))

    # mixing types
    X = rand(Float64, 5, 10)
    XX = convert.(Float32, X)

    M = fit(Whitening, X)
    MM = fit(Whitening, XX)

    # mixing types should not error
    transform(M, XX)
    transform(MM, X)

    # type consistency
    @test eltype(mean(M)) == Float64
    @test eltype(mean(MM)) == Float32
end


__precompile__(true)
module FourierSeries

    using FFTW
    export
        fourierSeriesStepReal,
        fourierSeriesSampledReal,
        fourierSeriesSynthesisReal,
        repeatPeriodically

    @doc raw"""
    # Function call

    `fourierSeriesStepReal(t, u, T, hMax)`

    # Description

    Calculates the real Fourier coefficients of a step function `u(t)` based on
    the data pairs `t` and `u` according to the following variables:

    ```math
    \mathtt{a[k]} =
    \frac{2}{\mathtt{T}} \int_{\mathtt{t[1]}}^{\mathtt{t[1]+T}}
    \mathtt{u(t)} \cos(k\omega \mathtt{t}) \operatorname{d}\mathtt{t} \\
    ```

    ```math
    \mathtt{b[k]} =
    \frac{2}{\mathtt{T}} \int_{\mathtt{t[1]}}^{\mathtt{t[1]+T}}
    \mathtt{u(t)} \sin(k\omega \mathtt{t}) \operatorname{d}\mathtt{t}
    ```

    # Function arguments

    `t` Time vector instances where a step occurs

    `u` Data vector of which Fourier coefficients shall be determined of, i.e.,
    `u[k]` is the data value from the right at `t[k]`

    `T` Time period of `u`

    `hMax` Maximum harmonic number to be determined

    # Return values

    `h` Vector of harmonic numbers start at, i.e., `h[1] = 0` (dc value of `u`),
    `h[hMax+1] = hMax`

    `f` Frequency vector associated with harmonic numbers, i.e. `f = h / T`

    `a` Cosine coefficients of Fourier series, where `a[1]` = dc value of `u`

    `b` Sine coefficients of Fourier series, where `b[1] = 0`

    # Examples

    Copy and paste code:

    ```julia
    # Fourier approximation of square wave form
    ts = [0; 1]                 # Sampled time data
    us = [-1; 1]                # Step function of square wave
    T = 2                       # Period
    hMax = 7                    # Maximum harmonic number
    (h, f, a, b) = fourierSeriesStepReal(ts, us, 2, hMax)
    (t, u)=fourierSeriesSynthesisReal(f, a, b)
    using PyPlot
    # Extend (t, u) by one element to right periodically
    (tx, ux) = repeatPeriodically(ts, us, right=1)
    step(tx, ux, where="post")  # Square wave function
    plot(t, u)                  # Fourier approximation up to hMax = 7
    ```
    """
    function fourierSeriesStepReal(t,u,T,hMax)
        # Check if t and have equal lengths
        if length(t) != length(u)
            error("module FourierSeries: function fourierSeriesStepReal:\n
    Vectors t and u have different lengths")
        end
        # Check if vector u is NOT of type Complex
        if u[1] isa Complex
            error("module FourierSeries: function fourierSeriesStepReal:\n
    The analyzed functions `u` must not be of Type ::Complex")
        end
        # Determine length of function u
        N = length(u)
        # Initialization of real result vectors
        a = zeros(hMax+1)
        b = zeros(hMax+1)
        # Cycle through loop to determine coefficients
        i = collect(1:N)
        # Time vector, extended by period T
        τ = cat(t, [T+t[1]], dims = 1)
        # DC value
        a[1] = sum(u .* (τ[i.+1] - τ[i])) / T
        b[1] = 0
        global a, b, τ, i, T
        for k in collect(1:hMax)
            a[k+1] = sum(u.*(+sin.(k * τ[i.+1] * 2 * pi / T)
                            .-sin.(k * τ[i] * 2 * pi / T))) / (k * pi)
            b[k+1] = sum(u.*(-cos.(k * τ[i.+1] * 2 * pi / T)
                            .+cos.(k * τ[i] * 2 * pi / T))) / (k * pi)
        end
        # Number of harmonics
        h = collect(0:hMax)
        # Frequencies
        f = h / T
        return (h,f,a,b)
    end

    """
    # Function call

    `fourierSeriesSampledReal(t,u,T,hMax)`

    # Description

    Calculates the real Fourier coefficients of a sampled function `u(t)` based
    on the data pairs `t` and `u`

    # Function arguments

    `t` Vector of sample times

    `u` Data vector of which Fourier coefficients shall be determined of, i.e.,
    `u[k]` is the sampled data value associated with `t[k]`

    `T` Time period of `u`

    `hMax` Maximum harmonic number to be determined

    # Return values

    `h` Vector of harmonic numbers start at, i.e., `h[1] = 0` (dc value of `u`),
    `h[hMax+1] = hMax`

    `f` Frequency vector associated with harmonic numbers, i.e. `f = h / T`

    `a` Cosine coefficients of Fourier series, where `a[1]` = dc value of `u`

    `b` Sine coefficients of Fourier series, where `b[1] = 0`

    # Examples

    Copy and paste code:

    ```julia
    # Fourier approximation of square wave form
    ts = [ 0; 1; 2; 3; 4; 5; 6; 7; 8; 9]    # Sampled time data
    us = [-1;-1;-1;-1;-1; 1; 1; 1; 1; 1]    # Step function of square wave
    T = 10                      # Period
    hMax = 7                    # Maximum harmonic number
    (h, f, a, b) = fourierSeriesSampledReal(ts, us, hMax)
    (t, u) = fourierSeriesSynthesisReal(f, a, b)
    using PyPlot
    # Extend (t, u) by one element to right periodically
    (tx, ux) = repeatPeriodically(ts, us, right=1)
    plot(tx, ux, marker="o")    # Square wave function
    plot(t, u)                  # Fourier approximation up to hMax = 7
    ```
    """
    function fourierSeriesSampledReal(t, u, hMax::Int64=typemax(Int64))
        # Check if vector u is NOT of type Complex
        if u[1] isa Complex
            error("module FourierSeries: function fourierSeriesSampledReal:\n
    The analyzed functions `u` must not be of Type ::Complex")
        end
        # Determine length of function u
        N = length(u)
        # Determine period of function u
        T = t[end] - t[1] + t[2] - t[1]
        # Apply fast Fourier transform
        c = 2 * fft(u) / N
        # Correct DC value by factor 1/2
        c[1] = c[1] / 2;
        if mod(N,2) == 1
            # Number of sampled data, N, is odd
            hMaxMax = div(N-1, 2)
        else
            # Number of sampled data, N, is even
            hMaxMa x =div(N, 2)
        end
        # Number of harmonics
        h = collect(0:hMaxMax)
        # Frequencies
        f = h / T
        # Real and imaginary coefficients
        a = +real(c[1:min(hMax,hMaxMax) + 1])
        b = -imag(c[1:min(hMax,hMaxMax) + 1])
        # Return vectors
        return (h, f, a, b)
    end

    function fourierSeriesSynthesisReal(f, a, b;hMax=length(a)-1, N=1000, t0=0)
        # Check if vectors f, a and b have equal lengths
        if length(f) != length(a)
            error("module Fourier: function fourierSeriesSynthesisReal:\n
    Vectors `f` and `a` have different lengths")
        end
        if length(a) != length(b)
            error("module Fourier: function fourierSeriesSynthesisReal:\n
    Vectors `a` and `b` have different lengths")
        end
        # Determine Period of synthesized function
        T = 1 / f[2]
        # Initialization of synthesis function f
        u = fill(a[1],N)
        # time samples of synthesis function, considering start time t0
        t = collect(0:N-1) / N * T + fill(t0, N)
        # hMax may either be a scalar of vector
        if hMax isa Array
            # If hMax is an array, then synthesize only harmonic numbers
            # indicated by hMax
            kRange = hMax
        else
            # If hMax is a scalar, then treat hMax as the maximum harmonic
            # number, which may not exceed length(a)-1, as a[1] equals the dc
            # component (harmonic 0) and a[hMax-1] represents harmonic number
            # hMax
            kRange = collect(1:min(hMax, length(a)-1))
        end
        # Calculate superosition
        for k in kRange
            u = u + a[k+1] * cos.(k * t * 2 * pi / T)
                  + b[k+1] * sin.(k * t * 2 * pi / T)
        end
        return (t, u)
    end

    function repeatPeriodically(t, u; left = 0, right = 1)
        # Check if vectors t and u have equal lengths
        if length(t) != length(u)
            error("module Fourier: function repeatPeriodically:\n
    Vectors `t` and `u` have different lengths")
        end
        # Determine length of arrays t and u
        N = length(t)
        # Determine period of time axis
        T = t[end] - t[1] + t[2] - t[1]
        # Determine how many times the vectors t and u shall be repeated
        outerRight = Int(floor((right-1)/N) + 1)
        outerLeft  = Int(floor((left -1)/N) + 1)
        # Repeat vector t by full periods and consider period T
        tx = repeat(t, outer = outerRight+outerLeft + 1) +
             repeat(collect(-outerLeft: outerRight), inner = N) * T
        # Repeat vector u by full periods
        ux = repeat(u, outer = outerLeft + outerRight + 1)
        # Limit output vectors elements
        tx = tx[outerLeft * N - left + 1:(outerLeft * N) + N + right]
        ux = ux[outerLeft * N - left + 1:(outerLeft * N) + N + right]
        return (tx, ux)
    end

end


#==============================================================================
    Code for solving the Hamiltonian Jacboi Bellman for
	   an basic model: a neoclassical growth model: ρV(k) = max_{c} U(c) + V'(k)[F(k)-δk - c]
	   Where s(k) = F(k) - δk - c(k) and c(k) = U'^{-1}(V'(k))

	Translated Julia code from Matlab code by Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm

    Updated to Julia 1.0.0
==============================================================================#

using Distributions, Plots, SparseArrays, LinearAlgebra

σ= 2.0 #
ρ = 0.05 #the discount rate
δ = 0.05 # the depreciation rate
A = 1.0
α= 0.3

k_ss = (α*A/(ρ+δ))^(1/(1-α))

H = 10000
k_min = 0.001*k_ss
k_max = 2.0*k_ss

k = LinRange(k_min, k_max, H)
k = convert(Array, k) # create grid for a values
dk = (k_max-k_min)/(H-1)

Δ = 1000

maxit = 10
ε = 10e-6

dVf, dVb = [zeros(H,1) for i in 1:6]

#initial guess for V
v0 = (A.*k.^α).^(1-σ)/(1-σ)/ρ
v= v0

dist=[]

for n=1:maxit
	V=v

    # forward difference
	dVf[1:H-1] = (V[2:H]-V[1:H-1])/dk
	dVf[H] = (A.*k_max.^α - δ.*k_max)^(-σ)

	# backward difference
	dVb[2:H] = (V[2:H]-V[1:H-1])/dk
	dVb[1] = (A.*k_min.^α - δ.*k_min)^(-σ) # the boundary condition

	I_concave = dVb .> dVf

    # consumption and savings with forward difference
    cf = dVf.^(-1/σ)
    μ_f = A.*k.^α - δ.*k -cf

    # consumption and savings with backward difference
    cb = dVb.^(-1/σ)
    μ_b = A.*k.^α - δ.*k -cb

	c0 = A.*k.^α - δ.*k
    dV0 = c0.^(-σ)

    # Now to make a choice between forward and backward difference
    If = μ_f .> 0
    Ib = μ_b .< 0
    I0 = (1.0.-If-Ib)

    global dV_Upwind= dVf.*If + dVb.*Ib + dV0.*I0

    global c = dV_Upwind.^(-1/σ)
    u = c.^(1-σ)/(1-σ)

    # create the transition matrix
    X = -min.(μ_b,0)/dk
    Y = -max.(μ_f,0)/dk + min.(μ_b,0)/dk
    Z = max.(μ_f,0)/dk

    AA = sparse(Diagonal((Y[:]))) + [zeros(1,H); sparse(Diagonal((X[2:H]))) zeros(H-1,1)] + [zeros(H-1,1) sparse(Diagonal((Z[1:H-1]))); zeros(1,H)]
    B = (ρ + 1/Δ)*sparse(I,H,H) - AA
    b = u + V./Δ

    V = B\b
    V_change = V-v
    global v= V

	push!(dist,findmax(abs.(V_change))[1])
	if dist[n] .< ε
		println("Value Function Converged Iteration=")
		println(n)
		break
	end
end

plot(dist, grid=false,
		xlabel="Iteration", ylabel="||V^{n+1} - V^n||",
		ylims=(-0.001,0.030),
		legend=false, title="")
png("Convergence")

v_err = c.^(1-σ)/(1-σ) + dV_Upwind.*(A.*k.^α - δ.*k -c) - ρ.*v

plot(k, v_err, grid=false,
		xlabel="k", ylabel="Error in the HJB equation",
		xlims=(k_min,k_max),
		legend=false, title="")
png("HJB_error")

plot(k, v, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k_min,k_max),
		legend=false, title="")
png("Value_function_vs_k")

plot(k, c, grid=false,
		xlabel="k", ylabel="c(k)",
		xlims=(k_min,k_max),
		legend=false, title="")
png("c(k)_vs_k")

# approximation at the borrowing constraint
k_dot = (A.*k.^α - δ.*k -c)

plot(k, k_dot, grid=false,
		xlabel="k", ylabel="s(k)",
		xlims=(k_min,k_max), title="", label="s(k)", legend=false)
plot!(k, zeros(H,1), label="", line=:dash)
png("stateconstraint")


function thomas(a::Vector{Float64}, b::Vector{Float64}, c::Vector{Float64},
                d::Vector{Float64}, n::Int64)

    x = copy(d)
    c_prime = copy(c)

    # Setting initial elements
    c_prime[1] /= b[1]
    x[1] /= b[1]

    for i = 2:n
        # Scale factor is for c_prime and x
        scale = 1.0 / (b[i] - c_prime[i-1]*a[i])
        c_prime[i] *= scale
        x[i] = (x[i] - a[i] * x[i-1]) * scale
    end

    # Back-substitution
    for i = n-1:-1:1
        x[i] -= (c_prime[i] * x[i+1])
    end

    return x

end

function main()
    a = [0.0, 2.0, 3.0]
    b = [1.0, 3.0, 6.0]
    c = [4.0, 5.0, 0.0]
    d = [7.0, 5.0, 3.0]

    println(
        """The system
        $(join((b[1], c[1], "",   "|", d[1]), "\t"))
        $(join((a[2], b[2], c[2], "|", d[2]), "\t"))
        $(join(("",   a[3], b[3], "|", d[3]), "\t"))
        Has the solution:"""
    )

    soln = thomas(a, b, c, d, 3)

    println(soln)
end

main()


function translation(x::T,y::T,z::T)::Array{T,2} where {T<:Real}
    return [1 0 0 x; 0 1 0 y; 0 0 1 z; 0 0 0 1]
end

function translation(x::Real,y::Real,z::Real)
    return translation(promote(x,y,z)...)
end

function scaling(x::T,y::T,z::T)::Array{T,2} where {T<:Real}
    return [x 0 0 0; 0 y 0 0; 0 0 z 0; 0 0 0 1]
end

function scaling(x::Real,y::Real,z::Real)
    return scaling(promote(x,y,z)...)
end

function rotation_x(angle::T)::Array{T,2} where {T<:Real}
    return [1 0 0 0; 0 cos(angle) -sin(angle) 0; 0 sin(angle) cos(angle) 0; 0 0 0 1]
end

function rotation_y(angle::T)::Array{T,2} where {T<:Real}
    return [cos(angle) 0 sin(angle) 0; 0 1 0 0; -sin(angle) 0 cos(angle) 0; 0 0 0 1]
end

function rotation_z(angle::T)::Array{T,2} where {T<:Real}
    return [cos(angle) -sin(angle) 0 0; sin(angle) cos(angle) 0 0; 0 0 1 0; 0 0 0 1]
end

function shearing(x_y::T,x_z::T,y_x::T,y_z::T,z_x::T,z_y::T)::Array{T,2} where {T<:Real}
    return [1 x_y x_z 0; y_x 1 y_z 0; z_x z_y 1 0; 0 0 0 1]
end


"""Julia program to print the n'th term in the Golomb sequence
Golomb sequence is a non-decreasing integer sequence where n'th
term is  equal to the number of times n appears in the sequence"""

function golomb_sequence(n)
    dp = zeros(Int, n + 1)
    dp[1] = 1

    for i in 2:n
        dp[i] = 1 + dp[i - dp[dp[i - 1]]]
    end
    return dp[n]
end

print("Enter the value of n?, where you need the n'th number in the golomb sequence. ")
n = readline()
n = parse(Int, n)
if (n <= 0)
    println("The given value of n is invalid.")
    exit()
end
res = golomb_sequence(n)
println("The $n'th term in the golomb sequence is $res.")

"""
Time Complexity: O(n), where 'n' is the given number
Space Complexity: O(n)

SAMPLE INPUT AND OUTPUT

SAMPLE 1
Enter the value of n?, where you need the n'th number in the golomb sequence. 5
The 5'th term in the golomb sequence is 3.

SAMPLE 2
Enter the value of n?, where you need the n'th number in the golomb sequence. 867
The 867'th term in the golomb sequence is 79.

"""


using  LinearAlgebra
using Plots

function secant_method(f′, x0, x1, ϵ)
    g0 = f′(x0)
    Δ = Inf
    while abs(Δ) > ϵ
        g1 = f′(x1)
        Δ = (x1 - x0)/(g1 - g0)*g1
        x0, x1, g0 = x1, x1 - Δ, g1
    end
    return x1
end

f = x->sin(10x)+cos(3x)
df(x) = 10*cos(10*x)-3*sin(3*x)

#dois pontos iniciais
x0 = 3.4
x1 = 3.8

#convergencia
convergence = 1e-8

xs = secant_method(df,x0,x1,convergence)
xs,f(xs),df(xs)

plot(f,0.,8.,lw=2,draw_arrow="true")
#scatter!([a,c],[f(a),f(c)])
scatter!([x0,x1],[f(x0),f(x1)])
scatter!([xs],[f(xs)])


"""`semimajoraxis(x, y, z, u, v, w, m1, m2)`

Calculates semimajor axis for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function semimajoraxis(x, y, z, u, v, w, m1, m2)
    r = sqrt(x^2+y^2+z^2)
    vsq = u^2+v^2+w^2
    M = m1+m2
    return 1/( (2/r)-(vsq/M) )
end

"""`eccentricity(x, y, z, u, v, w, m1, m2)`

Calculates eccentricity for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function eccentricity(x, y, z, u, v, w, m1, m2)
    hsq = (y*w-z*v)^2 + (z*u-x*w)^2 + (x*v-y*u)^2
    a = semimajoraxis(x, y, z, u, v, w, m1, m2)
    M = m1+m2
    quotient = hsq/( M*a )
    return sqrt(1 - quotient)
end

"""`inclination(x, y, z, u, v, w)`

Calculates inclination for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies."""
function inclination(x, y, z, u, v, w)
    hz = x*v-y*u
    h = sqrt( (y*w-z*v)^2 + (z*u-x*w)^2 + (x*v-y*u)^2 )
    return acos(hz/h)
end

"""`aei(x, y, z, u, v, w, m1, m2)`

Returns semimajor axis `a`, eccentricity `e` and inclination `I` for the two-body
problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)`vectors
between two bodies with masses `m1` and `m2`."""
function aei(x, y, z, u, v, w, m1, m2)
    r = sqrt(x^2+y^2+z^2)
    vsq = u^2+v^2+w^2
    M = m1+m2
    a = 1/( (2/r)-(vsq/M) )
    hsq = (y*w-z*v)^2 + (z*u-x*w)^2 + (x*v-y*u)^2
    h = sqrt(hsq)
    quotient = hsq/( M*a )
    e = sqrt(1 - quotient)
    hz = x*v-y*u
    inc = acos(hz/h)
    return a, e, inc
end

"""`ae(x, y, z, u, v, w)`

Returns semimajor axis `a` and eccentricity `e` for the two-body problem defined
by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two
bodies with masses `m1` and `m2`."""
function ae(x, y, z, u, v, w, m1, m2)
    r = sqrt(x^2+y^2+z^2)
    vsq = u^2+v^2+w^2
    M = m1+m2
    a = 1/( (2/r)-(vsq/M) )
    hsq = (y*w-z*v)^2 + (z*u-x*w)^2 + (x*v-y*u)^2
    quotient = hsq/( M*a )
    e = sqrt(1 - quotient)
    return a, e
end

"""`longascnode(x, y, z, u, v, w)`

Calculates longitude of ascending node for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies."""
function longascnode(
    x, y, z,
    u, v, w
    )

    #the longitude of ascending node is computed as
    #the angle between x-axis and the vector n = (-hy,hx,0)
    #where hx, hy, are resp., the x and y comps. of ang mom vector per unit mass, h

    res = atan( y*w-z*v, x*w-z*u)

    if constant_term(res) ≥ zero(constant_term(res))
        return res
    else
        return res+2pi
    end

end

"""`argperi(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) argument of pericentre for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function argperi(
    x, y, z,
    u, v, w,
    m1, m2
    )

    #n = (z-axis unit vector)×h = (-hy, hx, 0)
    n = [x*w-z*u, y*w-z*v, zero(x)]
    e = lrlvec(x,y,z,u,v,w,m1,m2) #cartesian comps. of Laplace-Runge-Lenz vector
    n = n/sqrt(n[1]^2+n[2]^2+n[3]^2)
    e = e/sqrt(e[1]^2+e[2]^2+e[3]^2)
    cosω = dot(n, e)

    if constant_term(e[3]) >= zero(constant_term(x))
        return acos(cosω)
    else
        return 2pi-acos(cosω)
    end

end

"""`longperi(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) longitude of pericentre for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function longperi(
    x, y, z,
    u, v, w,
    m1, m2
    )

    return argperi(x, y, z, u, v, w, m1, m2)+longascnode(x, y, z, u, v, w)
end

"""`trueanomaly(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) true anomaly for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function trueanomaly(
    x, y, z,
    u, v, w,
    m1, m2
    )

    R = sqrt(x^2+y^2+z^2)
    V2 = u^2+v^2+w^2
    V = sqrt(V2)

    a = 1.0/( (2.0/R)-(V2/(m1+m2)) )

    h2 = (y*w-z*v)^2 + (z*u-x*w)^2 + (x*v-y*u)^2
    h = sqrt(h2)

    quotient=h2/( (m1+m2)*a )

    e = sqrt(1.0 - quotient)

    Rdot = (x*u+y*v+z*w)/R

    sin_f = Rdot*a*(1.0-e^2)/( e*h )
    cos_f = (1.0/e)*( -1.0+a*(1.0-e^2)/R )

    res = atan(sin_f, cos_f)

    if res ≥ zero(res)
        return res
    else
        return res+2pi
    end
end

"""`ecanomaly(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) eccentric anomaly for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function ecanomaly(
    x, y, z,
    u, v, w,
    m1, m2
    )

    a = semimajoraxis(x, y, z, u, v, w, m1, m2)
    e = eccentricity(x, y, z, u, v, w, m1, m2)
    R = sqrt(x^2+y^2+z^2)
    Rdot = (x*u+y*v+z*w)/R
    mu = m1+m2

    cosE = (a-R)/(a*e)
    sinE = Rdot*R/( e*sqrt(mu*a) )

    res = atan(sinE, cosE)

    return res

end

"""`meananomaly(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) mean anomaly for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function meananomaly(
    x, y, z,
    u, v, w,
    m1, m2
    )

    E = ecanomaly(x, y, z, u, v, w, m1, m2)
    e = eccentricity(x, y, z, u, v, w, m1, m2)

    return E-e*sin(E)
end

"""`timeperipass(t, x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) time of pericentre passage, at time `t`, for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function timeperipass(
    t, x, y, z,
    u, v, w,
    m1, m2
    )

    me_an = meananomaly(x, y, z, u, v, w, m1, m2) #mean anomaly
    a = semimajoraxis(x, y, z, u, v, w, m1, m2)
    me_mo = meanmotion(m1+m2, a) #mean motion

    return t-me_an/me_mo
end

"""`rungelenzx(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) x-component of the Laplace-Runge-Lenz vector for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function rungelenzx(
    x, y, z,
    u, v, w,
    m1, m2
    )

    # v = [u, v, w]
    # h = [y.*w-z.*v, z.*u-x.*w, x.*v-y.*u]
    mu = m1+m2
    # r = [x, y, z]
    rmag = sqrt(x^2+y^2+z^2)

    lrl_x = ( -(z*u-x*w)*w+(x*v-y*u)*v )/(mu)-x/rmag

    return lrl_x
end

"""`rungelenzy(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) y-component of the Laplace-Runge-Lenz vector for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function rungelenzy(
    x, y, z,
    u, v, w,
    m1, m2
    )

    # v = [u, v, w]
    # h = [y.*w-z.*v, z.*u-x.*w, x.*v-y.*u]
    mu = m1+m2
    # r = [x, y, z]
    rmag = sqrt(x^2+y^2+z^2)

    lrl_y = ( -(x*v-y*u)*u+(y*w-z*v)*w )/(mu)-y/rmag

    return lrl_y
end

"""`rungelenzz(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) z-component of the Laplace-Runge-Lenz vector for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function rungelenzz(
    x, y, z,
    u, v, w,
    m1, m2
    )

    # v = [u, v, w]
    # h = [y.*w-z.*v, z.*u-x.*w, x.*v-y.*u]
    mu = m1+m2
    # r = [x, y, z]
    rmag = sqrt(x^2+y^2+z^2)

    lrl_z = ( -(y*w-z*v)*v+(z*u-x*w)*u )/(mu)-z/rmag

    return lrl_z
end

"""`rungelenzmag(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) magnitude of the Laplace-Runge-Lenz vector for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function rungelenzmag(
    x, y, z,
    u, v, w,
    m1, m2
    )

    # v = [u, v, w]
    # h = [y.*w-z.*v, z.*u-x.*w, x.*v-y.*u]
    mu = m1+m2
    # r = [x, y, z]
    rmag = sqrt(x^2+y^2+z^2)

    lrl_x = ( -(z*u-x*w)*w+(x*v-y*u)*v )/(mu)-x/rmag
    lrl_y = ( -(x*v-y*u)*u+(y*w-z*v)*w )/(mu)-y/rmag
    lrl_z = ( -(y*w-z*v)*v+(z*u-x*w)*u )/(mu)-z/rmag

    return sqrt(lrl_x^2+lrl_y^2+lrl_z^2)
end

"""`lrlvec(x, y, z, u, v, w, m1, m2)`

Calculates the instantaneous (osculating) cartesian components of the Laplace-Runge-Lenz vector for the two
body problem defined by the relative position `(x,y,z)` and velocity `(u,v,w)` vectors between two bodies
with masses `m1` and `m2`."""
function lrlvec(
    x, y, z,
    u, v, w,
    m1, m2
    )

    mu = m1+m2
    rmag = sqrt(x^2+y^2+z^2)

    lrl_x = ( -(z*u-x*w)*w+(x*v-y*u)*v )/(mu)-x/rmag
    lrl_y = ( -(x*v-y*u)*u+(y*w-z*v)*w )/(mu)-y/rmag
    lrl_z = ( -(y*w-z*v)*v+(z*u-x*w)*u )/(mu)-z/rmag

    return [lrl_x, lrl_y, lrl_z]
end

# compute eccentric anomaly (E) from eccentricity (e) and mean anomaly (M)
function eccentricanomaly(e::T, M::T) where {T <: Number}
    M0 = mod2pi(M)
    E0 = M0 + sign(sin(M0))*0.85*e #Murray-Dermotts' initial estimate
    # successive approximations via Newtons' method
    for i in 0:4
        #TODO: implement modified Newton's method for Kepler's equation (Murray-Dermott)
        Eans = E0 - (E0-e*sin(E0)-M0)/(1.0-e*cos(E0))
        E0 = Eans
    end
    return E0
end

# compute true anomaly (f) from eccentricity (e) and eccentric anomaly (E)
function trueanomaly(e,E)
    Enew = mod2pi(E)
    return 2.0*atan(  sqrt((1.0+e)/(1.0-e))*tan(Enew/2)  )
end

# compute true anomaly from eccentricity and mean anomaly
function meanan2truean(e,M)
    return trueanomaly(e, eccentricanomaly(e, M))
end

# get mean motion from mass parameter (mu) and semimajor axis (a)
function meanmotion(mu,a)
    return sqrt(mu/(a^3))
end

# get mean anomaly from mean motion (n), time (t) and time of pericenter passage (taup)
function meananomaly(n, t, taup)
    return n*(t-taup)
end

# compute true anomaly from time, a, e, mu and taup
function time2truean(a, e, mu, t, taup)
    return meanan2truean(e, meananomaly(meanmotion(mu, a), t, taup))
end


using Plots, LaTeXStrings, Statistics,  Distributions,StatsPlots, JLD

N=1000000
σList=[1, 2, 3, 4, 5  ]
function GaussianD(σ, UniformD)
    x1=rand(UniformD, round(Int,N/2))
    x2=rand(UniformD, round(Int,N/2))
    ρ=σ * .√(2 * log.(1 ./ (1 .- x1)))
    θ=2π * x2
    y1=ρ.*cos.(θ)
    y2=ρ.*sin.(θ)
    return [y1 ; y2]
end
Data=[GaussianD(σList[i], rand(N)) for i in 1:5]
save("../../chapter6/6.4/samples.jld","Samples", Data)
colors=[:rosybrown, :turquoise,:khaki1,:lightblue1, :lightpink2]
plot()
for i in 1:length(σList)
        histogram!(Data[i],bins=70,linecolor=:gray, color=colors[i],normalize = true, label=L"Histogram\ of\ Samples")
        plot!(Normal(mean(Data[i]),std(Data[i])), c = :black, label= L"\sigma = %$(round(std(Data[i]),digits = 2)))")
end
plot!(xlabel=L"",title=L"Normalized\ Distribution\ of\ Samples")
# plot(plots..., size = (1000,1000), plot_title = "Normalized Distribution of Samples", dpi=600)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter6\\6.4\\Figs\\N_Distribution.png")


# Astronomy example from Hairer et al 1992 p.129 (see also Szebehely
# 1967): the orbit of a satellite in planet-moon system.  This system
# has know periodic solutions, so called “Arenstorf orbits”
# (Arenstorf, 1963); one of which is used as a test-case.
#
# This is a non-stiff problem.  Adaptive solvers will fare much better
# as there are three regions where the satellite experiences very
# strong acceleration interspersed with regions of very small
# acceleration.

export threebody

threebody = let
    tcname = :threebody
    T = Float64 # the datatype used
    Tarr = Matrix
    dof = 4 # degrees of freedom
    dae=0  # index of DAE, ==0 for ODE
    # stiffness of system, one of the three constants
    stiffness = [nonstiff, mildlystiff, stiff][1]

    # parameters
    μ = 0.012277471
    μ′ = 1-μ

    function fn!(t,y,dydt)
        # the ode function

        y1, y2, y1′, y2′ = y[:] # this allocates memory... but less than below
        D1 = ((y1+μ )^2 + y2^2)^(3/2)
        D2 = ((y1-μ′)^2 + y2^2)^(3/2)
        y1″ = y1 + 2*y2′ - μ′*(y1+μ)/D1 - μ*(y1-μ′)/D2
        y2″ = y2 - 2*y1′ - μ′*y2/D1     - μ*y2/D2
        dydt[:] = [y1′, y2′, y1″, y2″] # this allocates memory...
        return nothing
    end
    fn!( ; T_::Type=T, dof_=dof) = zeros(T_,dof_)

    jac! = nothing
    mass! = nothing

    ic = T[0.994, 0.0, 0.0, -2.00158510637908252240537862224] # vector of initial conditions
    tspan = T[0, 17.0652165601579625588917206249] # integration interval
    refsol = ic # reference solution at tspan[2]
    refsolinds = trues(dof)
    scd_absinds = ic.==0.0
    tc = TestCaseExplicit{tcname, T, Tarr}(
                             stiffness,
                             dae,
                             dof,
                             fn!,
                             jac!,
                             mass!,
                             ic,
                             tspan,
                             refsol,
                             refsolinds,
                             scd_absinds)

    # put into the right buckets
    tc_all[tcname] = tc
    tc_nonstiff[tcname] = tc
end


##
##  i n t e r p o l a t e . j l  Interpolation
##

function interp1d{T<:Real}(xi::Array{T,1}, yi::Array{T,1}, x::T;
                           method::Symbol = :cubic)
    n = length(xi)
    if length(yi) != n
        error("Vectors 'xi', 'yi' must be of same length.")
    end
    if n == 1
       error("At least two data points must be provided.")
    end
    if any(xi[1:end-1] .> xi[2:end])
        error("Elements in 'xi' must be strictly increasing.")
    end

    local k::Int = 0
    if x < xi[1]; k = 0
    elseif x >= xi[end]; k = n
    else
        k = maximum(find(xi[1:(end-1)] .<= x))
    end

    if method == :constant
        if k == 0; k = 1; end
        return yi[k]

    elseif method == :nearest
        if k == 0; k = 1
        elseif k < n
            if x >= (xi[k] + xi[k+1]) / 2.0
                k += 1
            end
        end
        return yi[k]

    elseif method == :linear
        if k == 0; k = 1
        elseif k == n; k = n-1; end
        s = x - xi[k]
        delta = (yi[k+1] - yi[k]) / (xi[k+1] - xi[k])
        return yi[k] + (x - xi[k])*delta

    elseif method == :spline
        error("Method $(method) is not yet implemented.")

    elseif method == :cubic
        return pchip(xi, yi, x)

    else
        error("Unknown interpolation method: $(method)")
    end
end

function interp1d{T<:Real}(xi::Array{T,1}, yi::Array{T,1}, x::Array{T,1};
                           method::Symbol = :cubic)
    local m::Int = length(x)
    y = zeros(m)
    for i = 1:m
        y[i] = interp1d(xi, yi, x[i], method = method)
    end
    return y
end

#-- Moler: Piecewise Cubic Hermitean Interpolation Polynomials
#
function pchip{T<:Real}(xi::Array{T,1}, yi::Array{T,1}, x::T)
    n = length(xi)
    if length(yi) != n
        error("Vectors 'xi', 'yi' must be of same length.")
    end
    if any(xi[1:end-1] .> xi[2:end])
        error("Elements in 'xi' must be strictly increasing.")
    end

    # First derivatives
    h = diff(xi)
    delta = diff(yi) ./ h
    d = pchip_slopes(h, delta)

    # Piecewise polynomial coefficients
    a = (3.0*delta .- 2.0*d[1:(n-1)] .- d[2:n]) ./ h
    b = (d[1:(n-1)] .- 2.0*delta .+ d[2:n]) ./ h.^2

    # Find subinterval index k so that xi[k] <= x < xi[k+1]
    if x < xi[1]
        k = 1
    else
        k = maximum(find(xi[1:(n-1)] .<= x))
    end

    # Evaluate interpolant
    s = x - xi[k]
    v = yi[k] + s*(d[k] + s*(a[k] + s*b[k]))

    return v
end

function pchip_slopes(h, delta)

    # Slopes at interior points
    local n = length(h) + 1
    d = zeros(n)
    k = find(sign(delta[1:(n-2)]) .* sign(delta[2:(n-1)]) .> 0.0) .+ 1
    w1 = 2.0*h[k] .+ h[k.-1]
    w2 = h[k] .+ 2.0*h[k.-1]
    d[k] = (w1+w2) ./ (w1./delta[k.-1] + w2./delta[k])

    # Slopes at endpoints
    d[1] = pchip_end(h[1], h[2], delta[1], delta[2])
    d[n] = pchip_end(h[n-1], h[n-2], delta[n-1], delta[n-2])

    return d
end

function pchip_end(h1, h2, del1, del2)

    # Noncentered, shape-preserving, three-point formula.
    d = ((2.0*h1 + h2)*del1 - h1*del2) / (h1 + h2)
    if sign(d) != sign(del1)
        d = 0
    elseif sign(del1) != sign(del2)  &&  abs(d) > abs(3*del1)
        d = 3.0*del1
    end

    return d
end


using  LinearAlgebra
using Plots

function newtons_method(∇f, H, x, ϵ, k_max)
    k, Δ = 1, Inf
    while abs(Δ) > ϵ && k ≤ k_max
        Δ =  ∇f(x)/abs(H(x))
        x -= Δ
        k += 1
    end
    return x
end

#definicao da funcao
f = x-> -(0.4/((x^2+1))^(1/2))+((x^2+1)^(1/2))*(1-0.4/(x^2+1))-x    #multiplicada por -1 para encontrar o minimo
g = x-> -(-(0.4/((x^2+1))^(1/2))+((x^2+1)^(1/2))*(1-0.4/(x^2+1))-x) #funcao original

#definicao da primeira derivada
df(x) = 1.2*x/((x^2+1)^(3/2))+(1-0.4/(x^2+1))*x/((x^2+1)^(1/2))-1

#definicao da segunda derivada
ddf(x) = (-0.6*x^2+1.8)/((x^2+1)^(5/2))

#converge para o ponto de minimo
initial_point = 1       #ponto inicial
convergence = 1e-8      #convergencia
maximum_iterations = 10 #numero maximo de iteracoes
xs = newtons_method(df,ddf,initial_point,convergence,maximum_iterations)
#xs,f(xs),ddf(xs),df(xs)

println("initial_point = ", initial_point)
println("x* = ", xs)
println("f(x*) = ", g(xs))

plot(g,min(initial_point,xs)-1.,max(initial_point,xs)+1.,lw=2,draw_arrow="true",label="Function")
scatter!([initial_point],[g(initial_point)])
scatter!([xs],[g(xs)])


function chebyshev_polynomial(order::S,x::T) where {T<:Number,S<:Integer}

  polynomial    = Array{T}(undef,1,order+1)
  polynomial[1] = one(T)

  for i = 2:order+1
    if i == 2
      polynomial[i] = x
    else
      polynomial[i] = 2*x*polynomial[i-1]-polynomial[i-2]
    end
  end

  return polynomial

end

function chebyshev_polynomial(order::S,x::AbstractArray{T,1}) where {T<:Number,S<:Integer}

  polynomial      = Array{T}(undef,length(x),order+1)
  polynomial[:,1] = ones(T,length(x))

  for i = 2:order+1
    for j = 1:length(x)
      if i == 2
        polynomial[j,i] = x[j]
      else
        polynomial[j,i] = 2*x[j]*polynomial[j,i-1]-polynomial[j,i-2]
      end
    end
  end

  return polynomial

end

function derivative_of_chebyshev_polynomial(order::S,x::T) where {T<:Number,S<:Integer}

  polynomial    = Array{T}(undef,1,order+1)
  poly_deriv    = Array{T}(undef,1,order+1)
  polynomial[1] = one(T)
  poly_deriv[1] = zero(T)

  for i = 2:order+1
    if i == 2
      polynomial[i] = x
      poly_deriv[i] = one(T)
    else
      polynomial[i] = 2*x*polynomial[i-1]-polynomial[i-2]
      poly_deriv[i] = 2*polynomial[i-1]+2*x*poly_deriv[i-1]-poly_deriv[i-2]
    end
  end

  return poly_deriv

end


# # Numerical Stability

using Distributions
using HMMBase
using PyPlot
using Seaborn

rc("axes", xmargin = 0) # hide
set_style("whitegrid")  # hide

# Let's consider the case of a Normal distribution with null variance,
# such a case can appear during maximum likelihood estimation if only
# one observation is associated to a state:

A = [0.99 0.01; 0.1 0.9]
B = [Normal(0, 1), Normal(10.5, 0)]
hmm = HMM(A, B)
#-

y = rand(hmm, 500)
figure(figsize = (9,2)) # hide
plot(y)
gcf() # hide
#-

# The likelihood of a Normal distribution with null variance goes to infinity for `y = μ`,
# as there is a division by zero in the density function:
println(extrema(likelihoods(hmm, y)))
println(extrema(likelihoods(hmm, y, logl = true)))

# To avoid propagating these non-finite quantities (for example in the forward-backward algorithm),
# you can use the `robust` option:
println(extrema(likelihoods(hmm, y, robust = true)))
println(extrema(likelihoods(hmm, y, logl = true, robust = true)))

# This truncates `+Inf` to the largest Float64, and `-Inf` to the smallest Float64:
prevfloat(Inf), nextfloat(-Inf)

# In the log. case we use `log(prevfloat(Inf))` to avoid overflows when taking
# the exp. of the log-likelihood.