##################################################
## Classical derivative-based, iterative, root-finding algorithms
##
## If ri = f^(i-1)/f^(i), then these have an update step `xn - delta` where:
##
## * Newton: delta = r1  # order 2 if simple root (multiplicity 1)
## * Halley: delta = 2*r2/(2r2 - r1) * r1 # order 3 if simple root
## * Schroder: delta = r2  / (r2 - r1) * r1  # order 2
## * Thukral(3): delta =  (-2*r3)*(r2 - r1)/(r1^2 - 3*r1*r3 + 2*r2*r3) * r1 # order 3
## * Thukral(4): delta =  3*r1*r2*r4*(r1^2 - 3*r1*r3 + 2*r2*r3)/(-r1^3*r2 + 4*r1^2*r2*r4 + 3*r1^2*r3*r4 - 12*r1*r2*r3*r4 + 6*r2^2*r3*r4) # order 4
##
## The latter two come from
## [Thukral](http://article.sapub.org/10.5923.j.ajcam.20170702.05.html). They are not implemented.

## Newton
abstract type AbstractNewtonLikeMethod <: AbstractUnivariateZeroMethod end
fn_argout(::AbstractNewtonLikeMethod) = 2
struct Newton <: AbstractNewtonLikeMethod end
"""

    Roots.Newton()

Implements Newton's [method](http://tinyurl.com/b4d7vls):
`xᵢ₊₁ =  xᵢ - f(xᵢ)/f'(xᵢ)`.  This is a quadratically convergent method requiring
one derivative. Two function calls per step.

Example
```
find_zero((sin,cos), 3.0, Roots.Newton())
```

If function evaluations are expensive one can pass in a function which returns (f, f/f') as follows

```
find_zero(x -> (sin(x), sin(x)/cos(x)), 3.0, Roots.Newton())
```

This can be advantageous if the derivative is easily computed from the
value of f, but otherwise would be expensive to compute.

The error, `eᵢ = xᵢ - α`, can be expressed as `eᵢ₊₁ = f[xᵢ,xᵢ,α]/(2f[xᵢ,xᵢ])eᵢ²` (Sidi, Unified treatment of regula falsi, Newton-Raphson, secant, and Steffensen methods for nonlinear equations).

"""
Newton

# we store x0,x1,fx0,fx1 **and** Δ = fx1/f'(x1)
struct NewtonState{T,S} <: AbstractUnivariateZeroState{T,S}
    xn1::T
    xn0::T
    Δ::T
    fxn1::S
    fxn0::S
end

function init_state(M::Newton, F::Callable_Function, x)
    x₀ = float(first(x))
    fx₀, Δ = F(x₀)
    x₁ = x₀ - Δ
    state = init_state(M, F, x₀, x₁, fx₀, fx₀)
end

# compute fx₁, Δ
function init_state(::Newton, F, x₀, x₁, fx₀, fx₁)
    fx₁, Δ = F(x₁)
    NewtonState(x₁, x₀, Δ, fx₁, fx₀)
end

initial_fncalls(M::Newton) = 2

function update_state(method::Newton, F, o::NewtonState{T,S}, options, l=NullTracks()) where {T, S}

    xn0, xn1 = o.xn0, o.xn1
    fxn0, fxn1 = o.fxn0, o.fxn1
    Δ = o.Δ

    if isissue(Δ)
        log_message(l, "Issue with `f/f′'")
        return o, true
    end

    xn0, xn1 = xn1, xn1-Δ
    fxn0 = fxn1
    fxn1, Δ = F(xn1)
    incfn(l,2)

    @set! o.xn0 = xn0
    @set! o.xn1 = xn1
    @set! o.Δ = Δ
    @set! o.fxn0 = fxn0
    @set! o.fxn1 = fxn1

    return o, false

end

"""
    Roots.newton(f, fp, x0; kwargs...)

Implementation of Newton's method: `xᵢ₊₁ =  xᵢ - f(xᵢ)/f'(xᵢ)`.

Arguments:

* `f::Function` -- function to find zero of

* `fp::Function` -- the derivative of `f`.

* `x0::Number` -- initial guess. For Newton's method this may be complex.

With the `FowardDiff` package derivatives may be computed automatically. For example,  defining
`D(f) = x -> ForwardDiff.derivative(f, float(x))` allows `D(f)` to be used for the first derivative.

Keyword arguments are passed to `find_zero` using the `Roots.Newton()` method.

See also `Roots.newton((f,fp), x0)` and `Roots.newton(fΔf, x0)` for simpler implementations.

"""
newton(f, fp, x0; kwargs...) = find_zero((f, fp), x0, Newton(); kwargs...)

## Halley
abstract type AbstractHalleyLikeMethod <: AbstractUnivariateZeroMethod end
fn_argout(::AbstractHalleyLikeMethod) = 3

"""
    Roots.Halley()

Implements Halley's [method](http://tinyurl.com/yd83eytb),
`xᵢ₊₁ = xᵢ - (f/f')(xᵢ) * (1 - (f/f')(xᵢ) * (f''/f')(xᵢ) * 1/2)^(-1)`
This method is cubically converging, but requires more function calls per step (3) than
other methods.

Example
```
find_zero((sin, cos, x->-sin(x)), 3.0, Roots.Halley())
```

If function evaluations are expensive one can pass in a function which
returns `(f, f/f',f'/f'')` as follows

```
find_zero(x -> (sin(x), sin(x)/cos(x), -cos(x)/sin(x)), 3.0, Roots.Halley())
```

This can be advantageous if the derivatives are easily computed from
the computation for f, but otherwise would be expensive to compute separately.

The error, `eᵢ = xᵢ - α`, satisfies
`eᵢ₊₁ ≈ -(2f'⋅f''' -3⋅(f'')²)/(12⋅(f'')²) ⋅ eᵢ³` (all evaluated at `α`).

"""
struct Halley <: AbstractHalleyLikeMethod
end

struct HalleyState{T,S} <: AbstractUnivariateZeroState{T,S}
    xn1::T
    xn0::T
    Δ::T
    ΔΔ::T
    fxn1::S
    fxn0::S
end

# we compute one step here to get x₁
function init_state(M::AbstractHalleyLikeMethod, F::Callable_Function, x)
    x₀ = float(first(x))
    fx₀, (Δ, ΔΔ) = F(x₀)
    x₁ = x₀ - 2ΔΔ/(2ΔΔ - Δ) * Δ
    state = init_state(M, F, x₀, x₁, fx₀, fx₀)
end

function init_state(::AbstractHalleyLikeMethod, F, x₀, x₁, fx₀, fx₁)
    fx₁, (Δ, ΔΔ) = F(x₁)
    HalleyState(x₁, x₀, Δ, ΔΔ, fx₁, fx₀)
end

initial_fncalls(M::AbstractHalleyLikeMethod) = 2*3

function update_state(method::Halley, F, o::HalleyState{T,S}, options::UnivariateZeroOptions, l=NullTracks()) where {T,S}
    xn = o.xn1
    fxn = o.fxn1
    r1, r2 = o.Δ, o.ΔΔ

    Δ =  2*r2/(2r2 - r1) * r1
    if isissue(Δ)
        log_message(l, "Issue with computing `Δ`")
        return (o, true)
    end

    xn1::T = xn - Δ
    fxn1::S, (r1::T, r2::T) = F(xn1)
    incfn(l,3)

    @set! o.xn0 = xn
    @set! o.xn1 = xn1
    @set! o.fxn0 = fxn
    @set! o.fxn1 = fxn1
    @set! o.Δ = r1
    @set! o.ΔΔ = r2

    return o, false

end

"""
    Roots.halley(f, fp, fpp, x0; kwargs...)

Implementation of Halley's method (cf `?Roots.Halley()`).

Arguments:

* `f::Function` -- function to find zero of

* `fp::Function` -- derivative of `f`.

* `fpp:Function` -- second derivative of `f`.

* `x0::Number` -- initial guess

With the `FowardDiff` package derivatives may be computed automatically. For example,  defining
`D(f) = x -> ForwardDiff.derivative(f, float(x))` allows `D(f)` and `D(D(f))` to be used for the first and second
derivatives, respectively.

Keyword arguments are passed to `find_zero` using the `Roots.Halley()` method.

"""
halley(f, fp, fpp, x0; kwargs...) = find_zero((f, fp, fpp), x0, Halley(); kwargs...)


using DifferentialEquations
using IterableTables, DataFrames
using NamedTuples
using Gadfly

function F(du,u,p,t)
    S_H, E_H, I_H, R_H, S_V, E_V, I_V = u

    # host dynamics
    host_infection = (p.β*S_H*I_V)/p.N_H
    host_mortality = p.μ_H .* u[1:4] # include S_H, so easier to remove mortality
    host_births = sum(host_mortality)
    host_progression = p.σ_H*E_H
    recovery = p.λ*I_H

    du[1] = -host_infection + host_births
    du[2] = host_infection - host_progression
    du[3] = host_progression - recovery
    du[4] = recovery
    du[1:4] -= host_mortality

    # vector dynamics
    vec_infection = (p.β*S_V*I_H)/p.N_H
    vec_mortality = p.μ_V .* u[5:7] # include S_V, so easier to remove mortality
    vec_births = sum(vec_mortality)
    vec_progression = p.σ_V*E_V

    du[5] = -vec_infection + vec_births
    du[6] = vec_infection - vec_progression
    du[7] = vec_progression
    du[5:7] -= vec_mortality

end

# nb: in >= Julia v0.7, can eliminate this import
#  and the @NT syntax
u0 = [
    S_H=100.0,   E_H=0.0, I_H=1.0, R_H=0.0,
    S_V=10000.0, E_V=0.0, I_V=0.0
]
p = @NT(
  μ_H=1/365, μ_V=1/30, σ_H=1/3, σ_V=1/7, λ=1/14,
  β=0.05, N_H = sum(u0[1:4])
)
tspan = (0.0, 365.0)
prob = ODEProblem(F, u0, tspan, p)
sol = @time solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8,saveat=linspace(0,365,365*10+1))

df = DataFrame(sol)
rename!(df, Dict(:timestamp => :t,
  :value1 => :S_H, :value2 => :E_H, :value3 => :I_H, :value4 => :R_H,
  :value5 => :S_V, :value6 => :E_V, :value7 => :I_V
))
mlt = melt(df,:t) # convert results into long format for plotting
mlt[:host] = contains.(string.(mlt[:variable]),"H"); # tag which entries are host vs vector
df

# Plot
fig1a = plot(mlt[mlt[:host] .== true,:], x=:t, y=:value, color=:variable, Geom.line)
fig1b = plot(mlt[mlt[:host] .!= true,:], x=:t, y=:value, color=:variable, Geom.line)
vstack(fig1a,fig1b)


HOME_DIR = dirname(@__DIR__)
PACKAGE_DIR = dirname(HOME_DIR)
pushfirst!(LOAD_PATH, PACKAGE_DIR)
using MPCHelper

using ControlSystems
using LinearAlgebra
using LazySets

using Test

using Plots

## Initialization ---------------------------------------------------------------
# Linear 2D System
A = [1 -2; 0 1]
B= Matrix([1 1]')
R = 0.1*I
Q = Matrix(I,size(A))/10
K = dlqr(A,B,Q,R)
A_stab = A - B*K; eigvals(A_stab)
P = dlyap(A_stab, Q)
l_x = [3,3]
l_u = 1
z = [0.0 ;0.0]; v = 0.0

# # Linear 3D System
# A = [0.8 -2 3; 0 0.5 0; 0 0 0.3]
# B= [1 0; 0 1; 0 0.4]
# R = 0.1*Matrix(I,size(B,2), size(B,2))
# Q = 0.1*Matrix(I,size(A))
# K = dlqr(A,B,Q,R)
# A_stab = A - B*K; eigvals(A_stab)
# P = dlyap(A_stab, Q)
# l_x = [3,3,3]
# l_u = [2,2]
# z = [0.0 ;0.0; 0.0]; v = [0.0, 0.0]

# Main Part
α = 2
E = get_Ellipsoid(P, α)
l(x,u) = dot(l_x,x) + dot(l_u,u)
γ = l_x + K'*l_u
ωstar = ω_opt(γ, P, α)
hΩ = hΩ_opt(γ, P, α)
# Scale γ -> vec_γ, such that vec_γ'*γ the same cost as ω*γ
vec_γ = γ*hΩ/norm(γ)^2
@test ωstar'*γ ≈ vec_γ'*γ
# lmax(z,v) for Ellispoid defined by P, α
lmax(z,v) = lmax_mask(z,v, P, α, l_x, l_u, K)
# lmax(z,v) for Ellispoid defined by Ellipsoid E (Type LazySet)
lmax_lazy(z,v) = lmax_mask(z,v, E, l_x, l_u, K)
lmax_opt = lmax(z,v)
lmax2_opt = lmax_lazy(z,v)
@test hΩ ≈ lmax_opt
x_opt = z .+ ωstar
u_opt = vec(v .+ K*ωstar)
@test l(x_opt, u_opt) ≈ lmax_opt

## Plotting (only for 2D systems)
grid_res = -5:0.2:5
plot(E,1e-3, alpha=0.5, lab="Omega", size=(500,500))
plot_vec!(ωstar, line=(3,:blue), lab="w_opt")
plot_vec!(vec_γ, line=(2,:black, :dash), lab="lx + K'lu")
# plot_vec!([1,1], line=(2,:black, :dash), lab="lx + K'lu")
plot_lin_contour!(l_x, grid_res; color=:bluesreds, line=(1,:dash), colorbar=false)
plot_lin_contour!(γ, grid_res; color=:bluesreds, line=(1))#, contour_labels = true)
l_opt(x) = (-γ[1]*x + hΩ) / γ[2] # γ'x = hα
plot!(grid_res, l_opt.(grid_res), line=(:red), lab="lopt")
xlims!(-5,5)
ylims!(-5,5)

## PLotting Surfaces
grid_res = 0:0.1:10
surface(grid_res,grid_res, (x,y) -> lmax([x,y],0) )
surface!(grid_res,grid_res, (x,y) -> l([x,y],0))

## Solve OP for lmax and compare to analytical solution
using JuMP
using Ipopt
import LazySets.Approximations.overapproximate
n = size(A,2)

function lmax_optimized(z,v)
    HPolyrep_epsilon = overapproximate(E, 1e-5) # approximate ellipse by polytope
    model = Model(with_optimizer(Ipopt.Optimizer))
    JuMP.register(model, :l, 2, l, autodiff=true )
    @variable(model, ω[1:n])
    constrain_variable_by_set2!(model, ω, HPolyrep_epsilon)
    @objective(model, Max, l(z+ω, K*ω .+ v))
    # print(model)
    optimize!(model)
    opt_val = objective_value(model)
    opt_ω = value.(ω)
    @show opt_ω, opt_val
    return opt_val
end

lmax_optimized(z,v) - lmax(z,v)
@test abs(lmax_optimized(z,v) - lmax(z,v)) <= 1e-4


using MathepiaModels
using DifferentialEquations
using Plots
using Test
const DE = DifferentialEquations
@testset "SIRbasic" begin
    u_0 = [1000, 0.1, 0]
    p_data = (Λ = 0, d = 0, α = 0, N = 1000, β = 0.2, γ = 0.1)
    tspan_data = (0.0, 100.0)
    prob_data = DE.ODEProblem(SIRbasic, u_0, tspan_data, p_data)
    data_solve = DE.solve(prob_data, Tsit5(), abstol = 1e-12, reltol = 1e-12, saveat = 0.1)
    data_withoutnois = Array(data_solve)
    data = data_withoutnois #+ Float32(2e-1)*randn(eltype(data_withoutnois), size(data_withoutnois))
    Plots.plot(data_solve.t, data[1, :], label = "Train S")
    Plots.plot!(data_solve.t, data[2, :], label = "Train I")
    @info "SIRbasic succeed"
end

@testset "SEIRbasic" begin
    u_0 = [1000, 0, 1, 0]
    p_data = (Λ = 0, d = 0, α = 0, N = 1000, β = 0.2, γ = 0.1, σ=1/5.2, ke = 0.3)
    tspan_data = (0.0, 100.0)
    prob_data = DE.ODEProblem(SEIRbasic, u_0, tspan_data, p_data)
    data_solve = DE.solve(prob_data, Tsit5(), abstol = 1e-12, reltol = 1e-12, saveat = 0.1)
    data_withoutnois = Array(data_solve)
    data = data_withoutnois #+ Float32(2e-1)*randn(eltype(data_withoutnois), size(data_withoutnois))
    Plots.plot(data_solve.t, data[1, :], label = "Train S")
    Plots.plot!(data_solve.t, data[2, :], label = "Train E")
    @info "SEIRbasic succeed"
end

@testset "SEIARbasic" begin
    u_0 = [1000, 0, 1, 0, 0]
    p_data = (Λ = 0, d = 0, α = 0, N = 1000, β = 0.2, γ = 0.1, σ=1/5.2, ke = 0.3, ka =0.3, ρ =0.3, γA = 0.1)
    tspan_data = (0.0, 100.0)
    prob_data = DE.ODEProblem(SEIARbasic, u_0, tspan_data, p_data)
    data_solve = DE.solve(prob_data, Tsit5(), abstol = 1e-12, reltol = 1e-12, saveat = 0.1)
    data_withoutnois = Array(data_solve)
    data = data_withoutnois #+ Float32(2e-1)*randn(eltype(data_withoutnois), size(data_withoutnois))
    @show Plots.plot(data_solve.t, data[1, :], label = "Train S")
    Plots.plot!(data_solve.t, data[2, :], label = "Train E")
    @info "SEIARbasic succeed"
end

@testset "SISbasic" begin
    u_0 = [1000,  1]
    p_data = (N = 1000, β = 0.2, γ = 0.1)
    tspan_data = (0.0, 1000.0)
    prob_data = DE.ODEProblem(SISbasic, u_0, tspan_data, p_data)
    data_solve = DE.solve(prob_data, Tsit5(), abstol = 1e-12, reltol = 1e-12, saveat = 0.1)
    data_withoutnois = Array(data_solve)
    data = data_withoutnois #+ Float32(2e-1)*randn(eltype(data_withoutnois), size(data_withoutnois))
    @show Plots.plot(data_solve.t, data[1, :], label = "Train S")
    Plots.plot!(data_solve.t, data[2, :], label = "Train E")
    @info "SIS succeed"
end


# f(x)=abs(x)

function Fabs(x)
    return (abs(x),nothing)
end

function Fabs_inplace(value::Array,auxvalue,x) # inplace
    copy!(value,abs(x))
end

function Dabs(derivativeIDX,f_c,faux_c,grad_c,grad_n,x)
    axpy!(1.0,sign(x).*grad_c,grad_n)
end

if PROC=="GPU"
    function Fabs_inplace(value::CudaArray,auxvalue,x::CudaArray) # inplace
        fill!(value,0.0)
        abs!(x,value)
    end

    function Dabs(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray)
        xsigny_update!(grad_c,x,grad_n)
    end
end

Derivative[Fabs]=Dabs # Define dictionary lookup
Inplace[Fabs]=Fabs_inplace

import Base.abs

abs(n::ADnode)=ADnode(Fabs,n)

###abs(A::ADtrans)=transpose(abs(node[A.parent])) # abs(A')=(abs(A))' TODO:check

export abs


"""
Problem 1:

Given an input array of integers, count the number of values which increase relatively
to their precedent.

Problem 2:
Same, but using the convolution of the signal with a 3-position rectangular window.

"""

using DSP

# Read input
f = open("./inputs/day1.txt", "r")
data = map( x -> parse(Int,x), eachline(f))

# Problem
function countIncreasing(v::Vector{Int}) :: Int
  counter = 0
  for i in 2:length(v)
    if v[i] > v[i-1]
      counter += 1
    end
  end
  return counter
end

countIncreasingSums(v) = countIncreasing(conv(v,[1,1,1])[3:end-2])

println("Problem 1: ")
println(countIncreasing(data))
println("Problem 2: ")
println(countIncreasingSums(data))


module stest
using LinearAlgebra
using BenchmarkTools
using ..sparse

function test_create_sparse_vec()
    data = [10, 20, 1e-11, -30, 1e-5, 50, 1e-50]
    index = [1, 10, 15, 20]
    # vec = sparse.SparseVector(data, index)

    vec2 = sparse.create_sparse(data)

end

function test_create_sparse_mat()

    n_row = 5
    n_col = 3

    data        = [10,  20, 30, 40, 50, 60, 70, 80, 90]
    row_index   = [1,   2,  3,  4,  5,  1,  1,  4,  5, n_row]
    col_index   = [1,   1,  1,  2,  2,  2,  3,  3,  3, n_col]

    spmat = sparse.SparseMatrixCOO(data, row_index, col_index)

end

function test_add()

    data_x  = [3.3, 4, 5, 6, 10, 4,  5,  2,  3]
    index_x = [3,   4, 6, 8, 12, 15, 16, 18, 20, 25]

    sp_x = sparse.SparseVector(data_x, index_x)

    data_y  = [5.0, 2, 3, 5, 2, 3,  5,  2,  3]
    index_y = [4,   5, 6, 7, 8, 14, 15, 16, 18, 25]

    sp_y = sparse.SparseVector(data_y, index_y)

    sp_z = sparse.add(sp_x, sp_y)

    show(sp_z)

end

function test_matrix_vector_mult()

    n_row = 5
    n_col = 3

    data        = [10,  20, 30, 40, 50, 60, 70, 80, 90]
    row_index   = [1,   2,  3,  4,  5,  1,  1,  4,  5, n_row]
    col_index   = [1,   1,  1,  2,  2,  2,  3,  3,  3, n_col]

    spmat = sparse.SparseMatrixCOO(data, row_index, col_index)

    fvec = [10.0, 5, 6]

    res = spmat * fvec # sparse.matrix_vector_prod(spmat, fvec)
    show(res)

    fmat = sparse.create_full(spmat)
    res_theo = fmat * fvec
    show(res_theo)

    show(res_theo - res)

end

function create_matrix(n_dim = 8)

    A = zeros(Float64, n_dim, n_dim)

    main_diag = [1; 4 * ones(n_dim-2); 4]
    side_diag = -ones(n_dim-1)
    sideside_diag = -ones(n_dim-3)

    main_diag_idx = diagind(A, 0)
    side_m1_diag_idx = diagind(A, 1)
    side_p1_diag_idx = diagind(A, -1)

    side_m3_diag_idx = diagind(A, 3)
    side_p3_diag_idx = diagind(A, -3)

    A[ main_diag_idx ] = main_diag
    A[ side_m1_diag_idx ] = side_diag
    A[ side_p1_diag_idx ] = side_diag
    A[ side_m3_diag_idx ] = sideside_diag
    A[ side_p3_diag_idx ] = sideside_diag

    return A
end

function test_csc()
    n_dim = 8
    A = create_matrix(n_dim)
    A[:,1] .= 0.0
    A[:,4] .= 0.0

    spmat = sparse.create_sparse_csc(A)

    Af = sparse.create_full(spmat)

    vec = rand(n_dim) # ones(n_dim)

    res = sparse.matrix_vector_prod(spmat, vec)

    res_theo = A * vec

    println( "$(norm(res - res_theo))" )

end

function test_matmult(spmat, vec)
    res = sparse.matrix_vector_prod(spmat, vec)
end

function benchmark_matmult()

    n_dim = 25000 # 25_000
    A = create_matrix(n_dim)
    spmat = sparse.create_sparse_csc(A)
    vec = rand(n_dim)

    b = @benchmark test_matmult($spmat, $vec)
    println("=== sparse matrix vec mult === ")
    show(
        stdout,
        MIME("text/plain"),
        b)
    println("")
    println("")

    b = @benchmark $A * $vec
    println("=== default full matrix vec mult === ")
    show(
        stdout,
        MIME("text/plain"),
        b)

    # res_sp = test_matmult(spmat, vec)
    # res_fu = A * vec;

    # println("$(norm(res_sp - res_fu))")

end

n_dim = 8
A = create_matrix(n_dim)
A[7, 6] = 3.14
A[8,7] = 2.6
spmat = sparse.create_sparse_csc(A)

(L, d, U) = sparse.matrix_split(spmat)

Lf = sparse.create_full(L)
Uf = sparse.create_full(U)
Df = diagm(d)

println(" $(norm( Lf + Uf + Df - A )) ")

# stest.L[8,6]

end # stest


#
# Denoise an S2-valued image Example with Cyclic Proximal Point applied to the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data as for the corresponding CP algorithm
#
using Manopt, Manifolds
using Images, CSV, DataFrames, LinearAlgebra, JLD2
#
# Settings
experiment_name = "S2_Whirl_CPPA"
results_folder = joinpath(@__DIR__, "S2_TV")
export_result = true
export_orig = true
export_table = true
export_function_value = true
asymptote_render_detail = 2
!isdir(results_folder) && mkdir(results_folder)
#
# Manifold & Data
f = artificial_S2_whirl_image(64)
pixelM = Sphere(2);

if export_orig
    orig_filename = joinpath(results_folder, experiment_name * "-orig.asy")
    asymptote_export_S2_data(orig_filename; data=f)
    render_asymptote(orig_filename; render=asymptote_render_detail)
end
#
# Parameters
α = 1.5
maxIterations = 4000
#
# Build Problem for L2-TV
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
d = length(size(f))
iRep = [Integer.(ones(d))..., d]
fidelity(M, x) = 1 / 2 * distance(M, x, f)^2
Λ(M, x) = forward_logs(M, x) # on T_xN
prior(M, x) = norm(norm.(Ref(pixelM), repeat(x, iRep...), Λ(M, x)), 1)
#
# Setup and Optimize
cost(M, x) = fidelity(M, x) + α * prior(M, x)
proxes = ((M, λ, x) -> prox_distance(M, λ, f, x, 2), (M, λ, x) -> prox_TV(M, α * λ, x, 1))
x0 = f
@time o = cyclic_proximal_point(
    M,
    cost,
    proxes,
    x0;
    debug=[
        :Iteration,
        " | ",
        DebugProximalParameter(),
        " | ",
        :Change,
        " | ",
        :Cost,
        "\n",
        100,
        :Stop,
    ],
    record=[:Iteration, :Iterate, :Cost],
    stopping_criterion=StopAfterIteration(maxIterations),
    λ=i -> π / (2 * i),
    return_options=true,
)
y = get_solver_result(o)
yRec = get_record(o)
#
# Results
if export_result
    result_filename = joinpath(
        results_folder,
        experiment_name * "-result-$(maxIterations)-α$(replace(string(α), "." => "-")).asy",
    )
    asymptote_export_S2_data(result_filename; data=y)
    render_asymptote(result_filename; render=asymptote_render_detail)
end
if export_table
    A = cat([y[1] for y in yRec], [y[3] for y in yRec]; dims=2)
    table_filename = joinpath(results_folder, experiment_name * "-recorded-cost.csv")
    CSV.write(table_filename, DataFrame(A); writeheader=false)
end
if export_function_value
    fctval_filename = joinpath(results_folder, experiment_name * "-cost.jld2")
    values = Dict("cost_function_value" => last(yRec)[3], "iterations" => length(yRec) - 1)
    save(fctval_filename, values)
end


export poisson, poisson!

 # get the maximum intensity of an image
mymax(X::AbstractArray{<:RGB}) = max_rgb(X)
mymax(X::AbstractArray{<:Gray}) = gray(maximum(X))
mymax(X::AbstractArray) = maximum(X)

 # noise function for poisson
function noise_f(x, scaling, max_intens)
    # scale image to max_intensity and apply poisson noise
    # after Poisson noise scale it back
    return pois_rand(x * scaling / max_intens) * max_intens / scaling
end

f_pois(scaling, max_intens) = x -> scaling == nothing ? pois_rand(x) : noise_f(x, scaling, max_intens)

comb_pois(x, n) = n

poisson(X::AbstractArray, scaling=nothing; clip=false) = poisson!(copy(X), scaling, clip=clip)
function poisson!(X::AbstractArray, scaling=nothing; clip=false)
    max_intens = convert(Float64, mymax(X))
    return apply_noise!(comb_pois, f_pois(scaling, max_intens), X, clip)
end

"""
    poisson(X; scaling=nothing, clip=false)

Returns the array `X` affected by Poisson noise.
At every position the Poisson noise affects the intensity individually
and the values at the positions represent the expected value of the Poisson
Distribution.
Since Poisson Noise due to discrete events you should
provide the optional argument `scaling`. This `scaling` connects
the highest value of the array with the discrete number of events.
The highest value will be then scaled and the poisson noise is applied
Afterwards we scale the whole array back so that the initial intensity
is preserved but with applied Poisson noise.
`clip` is a keyword argument. If given, it clips the values to [0, 1]
"""
poisson


### PARALLELISATION:
# start julia with "julia -p auto" or "julia -p 4" to use multiple processors
using Revise
using Random
using SpecialFunctions
using PyCall
using PyPlot
using Plots
using StatsPlots
using CmdStan

#%%

# If run from command line externally:
#projDir = dirname(@__FILE__)

# If run from Jupyter/Hydrogen, maybe change to suit you:
projDir= "/home/johhub/Desktop/ABDA/A4"
#projDir= "/lhome/johhub/Desktop/ABDA/A4"
tmpDir = projDir*"/tmp"

################################################################

function runStanMC(y,a,b,burnIn,N,projDir,tmpDir,keepchains=false,noOfChains=4)
    # REQUIRES MODULES:
    # CmdStan
    #
    # INPUT:
    # y --> observed data
    # burnIn --> number of bur in samples
    # N --> total number of sample points
    # projDir --> directory where the analysis should be conducted, and plots should be saved
    # tmpDir --> directory where the temp files can be stored
    # keepchains (optional) --> keep all vectors of the various chains, false
    #                           per defaul, i.e. all chain are thrown in one vector
    #
    # OUTPUT:
    # x_s --> sampled x values for all chains (Nx4 array)

    ### 1) Define the Stan model, note that variable names must
    # correspond to observed data later
    # dichtonomous model
    dichtModel = "
    data {
    int<lower=0> J; // number of flips in the observed data, setting the minimum amount of data under lower
    int<lower=0,upper=1> y[J];  // coin flips
    }
    parameters {
    real<lower=0,upper=1> theta; // prob of getting a head
    }
    transformed parameters { // no transformed variables to use
    }
    model {
    theta ~ beta($a, $b);         // prior distribution for theta
    y ~ bernoulli(theta);       // likelihood, note that stan will create the posterior automatically.
    }
    ";

    ### 2) Data for the Stan model, note that variable names must
    # correspond to defined model in stan-code
    observedData = Dict("J" => length(y), "y" => y);

    ### 3) Chain specs
    myModel = Stanmodel(
                    Sample(save_warmup=true,
                           num_warmup=burnIn,
                           num_samples=N,
                           thin=1),   # thin: Period between saved samples
                    name = "flipping",
                    model = dichtModel,
                    printsummary = false,
                    tmpdir = tmpDir,
                    nchains = noOfChains);    # number of chains (4 default)

    ### 4) Run CmdStan:
    # rc:     return code (0 if all is fine)
    # chn:    chain results
    # cnames: vector of variable names
    rc, chn, cnames = stan(myModel,
                           observedData,
                           projDir,
                           diagnostics = false,
                           CmdStanDir = CMDSTAN_HOME);

    ### Results without burn-in
    #chns = chn[(n_burnin+1):end, :, :]
    # the 7th (out of 8) coulmn in the .value method are the actual sampled points
    # the 3rd dimension are the different chains: 1,2,3 --> chain 1,2,3

    x_s = 1* chn.value[(burnIn+1):end,7,:]  # the 1* converts the "Axis Array" in the chain into an Array{Float64,2}

    if keepchains == false
        x_s = reshape(x_s,:);   # Throw all the chain samples in a single vetor
    end

    return x_s
end

#%%

# Chain inputs:
N = 10^6
N_burnin = 1000

#%%

# First Column of Figure 6.4:
a = 100 # prior specs
b = 100 # prior specs
y = [ones(Int64,17); zeros(Int64,3)];    # make an Array{Intt64,1} array

# to receive the prior curve, just set the data to empty array []:
chain = runStanMC([],a,b,N_burnin,N,projDir,tmpDir);
histogram(chain, bins=100, normalize=:pdf, label="prior")
density!(chain,linewidth=3,label="prior")

# to receive the likelihood curve, just set the prior to a=1, b=1 (flat):
chain = runStanMC(y,1,1,N_burnin,N,projDir,tmpDir);
histogram!(chain, bins=100, normalize=:pdf, label="likelihood")
density!(chain,linewidth=3,label="likelihood")

posterior1 = runStanMC(y,a,b,N_burnin,N,projDir,tmpDir);
histogram!(posterior1, bins=100, normalize=:pdf, label="posterior")
density!(posterior1,linewidth=3,label="posterior")

Plots.savefig("A4-A-C1_Stan.pdf")

#%%
# Second Column
a = 18.25
b = 6.75
y = [ones(Int64,17); zeros(Int64,3)];    # make an Array{Intt64,1} array

# to receive the prior curve, just set the data to empty array []:
chain = runStanMC([],a,b,N_burnin,N,projDir,tmpDir);
histogram(chain, bins=100, normalize=:pdf, label="prior")
density!(chain,linewidth=3,label="prior")

# to receive the likelihood curve, just set the prior to a=1, b=1 (flat):
chain = runStanMC(y,1,1,N_burnin,N,projDir,tmpDir);
histogram!(chain, bins=100, normalize=:pdf, label="likelihood")
density!(chain,linewidth=3,label="likelihood")

posterior2 = runStanMC(y,a,b,N_burnin,N,projDir,tmpDir);
histogram!(posterior2, bins=100, normalize=:pdf, label="posterior")
density!(posterior2,linewidth=3,label="posterior")

Plots.savefig("A4-A-C2_Stan.pdf")

#%%
# Third Column
a = 1
b = 1
y = [ones(Int64,17); zeros(Int64,3)];    # make an Array{Intt64,1} array

# to receive the prior curve, just set the data to empty array []:
chain = runStanMC([],a,b,N_burnin,N,projDir,tmpDir);
histogram(chain, bins=100, normalize=:pdf, label="prior")
density!(chain,linewidth=3,label="prior")

# to receive the likelihood curve, just set the prior to a=1, b=1 (flat):
chain = runStanMC(y,1,1,N_burnin,N,projDir,tmpDir);
histogram!(chain, bins=100, normalize=:pdf, label="likelihood")
density!(chain,linewidth=3,label="likelihood")

posterior3 = runStanMC(y,a,b,N_burnin,N,projDir,tmpDir);
histogram!(posterior3, bins=100, normalize=:pdf, label="posterior")
density!(posterior3,linewidth=3,label="posterior")

Plots.savefig("A4-A-C3_Stan.pdf")

#%%

density(posterior1, label="a=100, b=100")
density!(posterior2, label="a=18.25, b=6.75")
density!(posterior3, label="a=1, b=1")
Plots.savefig("A4-A-all_Stan.pdf")


#=
Pentagonal numbers are generated by the formula, Pn=n(3n−1)/2. The first ten pentagonal numbers are:

1, 5, 12, 22, 35, 51, 70, 92, 117, 145, ...

It can be seen that P4 + P7 = 22 + 70 = 92 = P8. However, their difference, 70 − 22 = 48, is not pentagonal.

Find the pair of pentagonal numbers, Pj and Pk, for which their sum and difference are pentagonal and D = |Pk − Pj| is minimised; what is the value of D?
=#
function pent(x)
  div(x*(3x-1),2)
end

function ispent(x)
  n = (sqrt(24x+1)+1)/6
  return n == round(Int,n) && n > 0
end

function calc()
  i = 0
  pentagonals = Int64[]
  while true
    i += 1
    p = pent(i)
    for j in pentagonals
      diff = abs(p-j)
      if ispent(diff)
        intsum = p + j
        if ispent(intsum)
          return diff
        end
      end
    end
    push!(pentagonals, p)
  end
end
@time println(calc())


#----------------- Get the A matrix
function getDivGrad(n1,n2,n3)
    # the Divergence
    D1 = kron(speye(n3),kron(speye(n2),ddx(n1)))
    D2 = kron(speye(n3),kron(ddx(n2),speye(n1)))
    D3 = kron(ddx(n3),kron(speye(n2),speye(n1)))
    # DIV from faces to cell-centers

    Div = [D1 D2 D3]

    return Div*Div'
end
#----------------- 1D finite difference on staggered grid

# generate 1D derivatives
ddx(n)=spdiags(ones(n)*[-1 1],[0,1],n,n+1)

#------------- Build a diagonal matrix
function spdiags(B,d,m,n)
    # spdiags(B,d,m,n)
    # creates a sparse matrix from its diagonals
    d = d[:]
    p = length(d)

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

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


# using LinearAlgebra
# include("sphere.jl")
##

"""
This function compute the associated Legendre polynomial P(l,m,x)
    for m >= 0, l >= m, and abs(x) <= 1
"""
function Plm(l::Int, m::Int, x::AbstractFloat)
    @assert 0 <= m <= l
    pmm = 1.0
    if m > 0
        somx2 = sqrt((1.0 - x) * (1.0 + x))
        fact = 1.0
        for i = 1:m
            pmm *= (-fact) * somx2
            fact += 2.0
        end
    end

    if l == m
        return pmm
    end
    pmmp1 = x * (2.0 * m + 1.0) * pmm
    if (l == (m + 1))
        return pmmp1
    end

    pll = 0.0

    for ll = (m+2):l
        pll = ((2.0 * ll - 1.0) * x * pmmp1 - (ll + m - 1.0) * pmm) / (ll - m)
        pmm = pmmp1
        pmmp1 = pll
    end

    return pll
end
"""
This function compute the normalized associated Legendre polynomial P(l,m,x)
    for m >= 0, l >= m, and abs(x) <= 1. It is alias  of GSL sf_legendre_Plm
"""
function sph_Plm(l::Int, m::Int, x::AbstractFloat)
    @assert 0 <= m <= l
    return sf_legendre_sphPlm(l, m, x)
end
"""
This function compute all of the normalized associated Legendre polynomial P(l,m)
    for m >= 0, l >= m, and abs(x) <= 1. It is alias  of GSL sf_legendre_array.
    This implementation include the Condon-Shortley phase factor.
"""
function Plm_array(l_max::Int, x::AbstractFloat; normal = :SPHARM)
    normalization = Dict(
        :SPHARM => GSL_SF_LEGENDRE_SPHARM,
        :NONE => GSL_SF_LEGENDRE_NONE,
        :SCHMIDT => GSL_SF_LEGENDRE_SCHMIDT,
        :FULL => GSL_SF_LEGENDRE_FULL,
    )
    return sf_legendre_array_e(normalization[normal], l_max, x,1)
end
"""
This function returns the index of Plm_array corresponding to P(l,m)
"""
function Plm_array_index(l::Int, m::Int)
    @assert 0 <= m <= l
    return sf_legendre_array_index(l,m) + 1
end
"""
This function returns the the general real spherical harmonics function Y(l,m) at
    pol angle and az_angle. l ∈ (0,..,n) and ∈ (-l,..l) m are the
    order and degree of spherical harmonics respectively.
    Pol_angle ∈ (0,..pi) and az_angle ∈ (0,2pi) are the polar angle and azimuthal angle
    coordinate.
"""
function real_spharm(l::Int, m::Int, pol_angle, az_angle)
    sqrt2 = sqrt(2.0)
    if m == 0
        return sph_Plm(l, m, cos(pol_angle))
    elseif m > 0
        return sqrt2 * cos(m * az_angle) * sph_Plm(l, m, cos(pol_angle))
    elseif m < 0
        return sqrt2 * sin(-m * az_angle) * sph_Plm(l, -m, cos(pol_angle))
    end
end
"""
This function returns all of general real spherical harmonics function Y(l,m)
    with order l.
    Pol_angle ∈ (0,..pi) and az_angle ∈ (0,2pi) are the polar angle and azimuthal angle
    coordinate.
"""
function real_spharm_array(l_max::Int, pol_angle, az_angle)
    sqrt2 = sqrt(2.0)
    len_array = (l_max + 1)^2
    plm_array = Plm_array(l_max, cos(pol_angle); normal = :SPHARM)
    spharm_array = Vector{Float64}(undef, len_array)
    for j = 1:len_array
        l, m = j2lm(j)
        idx = Plm_array_index(l, abs(m))
        if m == 0
            spharm_array[j] = plm_array[idx]
        elseif m > 0
            spharm_array[j] = sqrt2 * cos(m * az_angle) * plm_array[idx]
        elseif m < 0
            spharm_array[j] = sqrt2 * sin(-m * az_angle) * plm_array[idx]
        end
    end
    return spharm_array
end

"""
This function returns the order l and degree m of Y(l,m) from linear
    index j of spharm_array
"""
function j2lm(j)
    j = j - 1
    l = floor(Int, sqrt(j))
    m = j - l * (l + 1)
    return l, m
end
"""
This function returns the linear index j of spharm_array
    from the order l and degree m of Y(l,m)
"""
function lm2j(l, m)
    return l^2 + l + m + 1
end

"""
This function returns the The matrix of general real spherical
    harmonics function Y(l,m). The row index i  corresponds to points x_i and
    column index j obtained from j =  l^2 + l + m +1 where l and m are the order
    and degree of spherical harmonics respectively.
"""
function spharm_mat(l_max, pol_angle::T, az_angle::T) where {T<:AbstractVector}
    n_row = length(pol_angle)
    n_column = (l_max + 1)^2
    mat = Matrix{eltype(pol_angle)}(undef, n_row, n_column)
    for i = 1:n_row
        mat[i, :] .= real_spharm_array(l_max, pol_angle[i], az_angle[i])
    end
    return mat
end

"""
This function returns the The matrix of coefficient of SH function Y(l,m). The
    row index i  corresponds to points x_i and column index j obtained from j =  l^2 + l + m +1 where l and m are the order
    and degree of spherical harmonics respectively.
"""
function spharm_coefs(l_max::Int, x, pol_angle,az_angle)
    mat = spharm_mat(l_max, pol_angle, az_angle)
    coefs = mat \ x
    return coefs
end


using Plots
using Polynomials

function main()

    # legendre_poly()
    spherical_harmonics()

end

function legendre_poly()

    n = 5
    xlin = range(-1, 1, length=100)

    plot(xaxis=("x", (-1,1)), yaxis=("y", (-1,1)), title="Legendre Polynomials",
         framestyle=:origin, legend=:outertopright, dpi=300)
    for i in 0:n

        pn = Pn(i)
        plot!(xlin, pn.(xlin), label="P"*string(i))

    end

    savefig("lp.png")

end

function Pn(n)
    p0 = Polynomial([-1, 0, 1])

    an = 1/(2^n * factorial(n))

    dp = (p0)^n
    for i in 1:n

        dp = derivative(dp)

    end

    return an*dp
end

function spherical_harmonics()

    theta = range(0, 2π, length=100)
    phi = range(0, π/2, length=100)

    n = 0
    m = 0

    plot3d(theta, phi, Ynm.(theta, phi, n, m))

end

function Ynm(θ, ϕ, n, m)

    pn = Pn(n)
    dp = pn
    pm = abs(m)

    for i in 0:pm
        dp = derivative(dp)
    end

    ptheta = (-1)^pm * sin(θ)^pm * dp(cos(θ))

    if m >= 0
        return ptheta * cos(m*ϕ)
    else
        return ptheta * sin(pm*ϕ)
    end

end

main()


# Problem 1
# If we list all the natural numbers below 10 that are multiples of 3 or 5,
# we get 3, 5, 6 and 9. The sum of these multiples is 23.
# Find the sum of all the multiples of 3 or 5 below 1000.

multiple_3_5(x) = ((x % 3) == 0) || ((x % 5) == 0)

∑ = 0
for i in 1:1000-1
    global ∑ += multiple_3_5(i) ? i : 0
end
print(∑)


"""
    SigmoidDerivative(x)

Derivative of the Sigmoid activation function.

    Usage:
```julia-repl
SigmoidDerivative(x)
```
Arguments:
- `x` : Input value.
"""
function SigmoidDerivative(x::Float64)
    x₁ = x * (1.0 - x)
    return x₁
end


using SpecialFunctions

@doc raw"""
Evaluates the Boys function ``Fₘ(x) = ∫₀¹ t²ᵐ \exp(-xt²) dt``.
`IND==0` gives around 14 digits, `IND==1` around 6 and `IND==3` around 3.
"""
function boys(m::Int, x::T, IND=0) where T <: Real
    x < 0 && throw(DomainError(x, "`x` must be nonnegative."))
    if x > 0
        # gamma_inc returns tuple of incomplete Γ function ratios,
        # we only care about the first returned result
        Px, _ = gamma_inc(m + 1 / T(2), x, IND)
        1 / 2x^(m + 1 / T(2)) * gamma(m + 1 / T(2)) * Px
    else
        1 / T(2m + 1)
    end
end

function doublefactorial(n::Integer)
    n < 0 && throw(DomainError(n, "`n` must be nonnegative."))
    start = iseven(n) ? 2 : 1
    f::typeof(n*n) = 1
    for i::typeof(n*n) = start:2:n
        f *= i
    end
    return f
end


#! /usr/bin/julia

# Rosetta Code, Arithmetic-geometric mean/Calculate Pi

function agm_step{T<:FloatingPoint}(x::T, y::T)
    (0.5*(x + y), sqrt(x*y))
end

function approx_pi_step{T<:FloatingPoint,U<:Integer}(x::T, y::T, z::T, n::U)
    (a, g) = agm_step(x, y)
    k = n + 1
    s = z + 2^(k+1)*(a^2 - g^2)
    return (a, g, s, k)
end

function approx_pi{T<:FloatingPoint}(a::T, g::T, s::T)
    4a^2/(1 - s)
end

prec = 512
set_bigfloat_precision(prec)
println("Approximating pi using ", prec, "-bit floats.")
println("   k     Error  Result")
a = big(1.0)
g = a/sqrt(big(2.0))
s = big(0.0)
k = 0
oldagpi = big(0.0)
for i in 1:100
    (a, g, s, k) = approx_pi_step(a, g, s, k)
    agpi = approx_pi(a, g, s)
    2eps(agpi) < abs(agpi-oldagpi) || break
    oldagpi = agpi
    err = pi - agpi
    print(@sprintf("  %2d ", i))
    print(@sprintf(" %9.1e", err))
    println(@sprintf(" %.60e", agpi))
end


include("../optimizers/linesearch.jl")
include("../utils/functions.jl")
using Gadfly

function optimize_linear_constraints(f, A, f_g, f_h, yinit=0)
    """
    Convert constrained optimization problem min f(x) s.t. Ax = b
        to unconstrained problem min g(y) = f(B*y + x0), where x0 is a solution to Ax = b
        and B is a basis for the kernel of A.

    Defaults to using Newton's line search algorithm.

    f: constrained objective
    A: constraint matrix
    f_g: gradient of f
    f_h: hessian of f
    yinit: init y with a vector of all yinit values
    """

    m, n = size(A)
    b = zeros(m)
    x0 = A \ b
    B = nullspace(A)

    y0 = ones(size(B)[2]) .* yinit

    g(y) = f(B*y + x0)
    g_g(y) = B' * f_g(B*y + x0)
    g_h(y) = B' * f_h(B*y + x0) * B

    yvals = line_search(g, y0, g_g, g_h, "newton", 5000)
    xvals = [B*y + x0 for y in yvals]
    lambdas = [A' \ f_g(x) for x in xvals]

    return yvals, xvals, lambdas
end

function generate_constraint_matrix(n)
    A = ones(2,n)
    for i in 1:n
        if i % 2 == 1
            A[2,i] = -1
        end
    end
    return A
end

A = generate_constraint_matrix(30)
b = zeros(2)
x0 = A \ b; B = nullspace(A)
f = cute
f_g = cute_g
f_h = cute_h
g(y) = f(B*y + x0)
g_g(y) = B' * f_g(B*y + x0)
g_h(y) = B' * f_h(B*y + x0) * B

yvals, xvals, lambdas = optimize_linear_constraints(f, A, f_g, f_h, 0);

niters = length(xvals)
fx = [f(x) for x in xvals]
Gadfly.plot(x=1:niters, y=fx, Geom.line,
    Guide.xlabel("iteration"), Guide.ylabel("f(x)"), Guide.title("Value of function"),
    Scale.x_log10, Scale.y_log10)

grads = [norm(g_g(y),2) for y in yvals]
ratios = grads[2:niters,:]./grads[1:niters-1,:]
Gadfly.plot(x=1:niters-1, y=ratios, Geom.line,
Guide.xlabel("iteration"), Guide.ylabel("f_g(x)"), Guide.title("Gradient norm ratios"),
    Scale.x_log10, Scale.y_log10)

Gadfly.plot(x=1:niters, y=grads, Geom.line,
Guide.xlabel("iteration"), Guide.ylabel("f_g(x)"), Guide.title("Gradient norm"),
    Scale.x_log10, Scale.y_log10)

# Lagrange multipliers
lambdas[end]

A = generate_constraint_matrix(100)
x0 = A \ b; B = nullspace(A)
f = cute
f_g = cute_g
f_h = cute_h
g(y) = f(B*y + x0)
g_g(y) = B' * f_g(B*y + x0)
g_h(y) = B' * f_h(B*y + x0) * B

yvals, xvals, lambdas = optimize_linear_constraints(f, A, f_g, f_h, 0);
# Lagrange multipliers
println(lambdas[end])
# Final value
println(f(xvals[end]))

niters = length(xvals)
fx = [f(x) for x in xvals]
#gy = [g(y) for y in yvals]

Gadfly.plot(x=1:niters, y=fx, Geom.line,
    Guide.xlabel("iteration"), Guide.ylabel("f(x)"), Guide.title("Value of function"),
    Scale.x_log10, Scale.y_log10)

grads = [norm(g_g(y),2) for y in yvals]
ratios = grads[2:niters,:]./grads[1:niters-1,:]
Gadfly.plot(x=1:niters-1, y=ratios, Geom.line,
Guide.xlabel("iteration"), Guide.ylabel("f_g(x)"), Guide.title("Gradient norm ratios"),
    Scale.x_log10, Scale.y_log10)

Gadfly.plot(x=1:niters, y=grads, Geom.line,
Guide.xlabel("iteration"), Guide.ylabel("f_g(x)"), Guide.title("Gradient norm"),
    Scale.x_log10, Scale.y_log10)


abstract type AbstractUtility end

#
# Separable utility
#

# Consumption utility

@doc doc"""
Type used to evaluate log utility. Log utility takes the form

u(c) = \log(c)

Additionally, this code assumes that if c < 1e-10 then

u(c) = log(1e-10) + 1e10*(c - 1e-10)

"""
struct LogUtility <: AbstractUtility
    ξ::Float64
end

LogUtility() = LogUtility(1.0)

(u::LogUtility)(c::Float64) =
    c > 1e-10 ? u.ξ*log(c) : u.ξ*(log(1e-10) + 1e10*(c - 1e-10))
derivative(u::LogUtility, c::Float64) =
    c > 1e-10 ? u.ξ / c : u.ξ*1e10

"""
Type used to evaluate CRRA utility. CRRA utility takes the form

u(c) = ξ c^(1 - γ) / (1 - γ)

Additionally, this code assumes that if c < 1e-10 then

u(c) = ξ (1e-10^(1 - γ) / (1 - γ) + 1e-10^(-γ) * (c - 1e-10))
"""
struct CRRAUtility <: AbstractUtility
    γ::Float64
    ξ::Float64

    function CRRAUtility(γ, ξ=1.0)
        if abs(γ - 1.0) < 1e-8
            error("Your value for γ is very close to 1... Consider using LogUtility")
        end

        return new(γ, ξ)
    end
end

(u::CRRAUtility)(c::Float64) =
    c > 1e-10 ?
           u.ξ * (c^(1.0 - u.γ) - 1.0) / (1.0 - u.γ) :
           u.ξ * ((1e-10^(1.0 - u.γ) - 1.0) / (1.0 - u.γ) + 1e-10^(-u.γ)*(c - 1e-10))
derivative(u::CRRAUtility, c::Float64) =
    c > 1e-10 ? u.ξ * c^(-u.γ) : u.ξ*1e-10^(-u.γ)

# Labor Utility

"""
Type used to evaluate constant Frisch elasticity (CFE) utility. CFE
utility takes the form

v(l) = ξ l^(1 + 1/ϕ) / (1 + 1/ϕ)

Additionally, this code assumes that if l < 1e-10 then

v(l) = ξ (1e-10^(1 + 1/ϕ) / (1 + 1/ϕ) - 1e-10^(1/ϕ) * (1e-10 - l))
"""
struct CFEUtility <: AbstractUtility
    ϕ::Float64
    ξ::Float64

    function CFEUtility(ϕ, ξ=1.0)
        if abs(ϕ - 1.0) < 1e-8
            error("Your value for ϕ is very close to 1... Consider using LogUtility")
        end

        return new(ϕ, ξ)
    end
end

(u::CFEUtility)(l::Float64) =
    l > 1e-10 ?
           -u.ξ * l^(1.0 + 1.0/u.ϕ)/(1.0 + 1.0/u.ϕ) :
           -u.ξ * (1e-10^(1.0 + 1.0/u.ϕ)/(1.0 + 1.0/u.ϕ) + 1e-10^(1.0/u.ϕ) * (l - 1e-10))
derivative(u::CFEUtility, l::Float64) =
    l > 1e-10 ? -u.ξ * l^(1.0/u.ϕ) : -u.ξ * 1e-10^(1.0/u.ϕ)

"""
Type used to evaluate elliptical utility function. Elliptical utility takes form

v(l) = b (1 - l^μ)^(1 / μ)
"""
struct EllipticalUtility <: AbstractUtility
    b::Float64
    μ::Float64
end

# These defaults are pulled straight from Evans Phillips 2017
EllipticalUtility(;b=0.5223, μ=2.2926) = EllipticalUtility(b, μ)

(u::EllipticalUtility)(l::Float64) =
    u.b * (1.0 - l^u.μ)^(1.0 / u.μ)
derivative(u::EllipticalUtility, l::Float64) =
    -u.b * (1.0 - l^u.μ)^(1.0/u.μ - 1.0) * l^(u.μ - 1.0)


# Julia translation of http://nbviewer.jupyter.org/github/barbagroup/AeroPython/blob/master/lessons/04_Lesson04_vortex.ipynb
# Lesson 3 doublet

using PyPlot
using Distributions

close("all")
meshgrid(x,y) = (repmat(x',length(y),1),repmat(y,1,length(x)))

N = 50                                # number of points in each direction
x_start, x_end = -2.0, 2.0            # boundaries in the x-direction
y_start, y_end = -1.0, 1.0            # boundaries in the y-direction
x = linspace(x_start, x_end, N)    # creates a 1D-array with the x-coordinates
y = linspace(y_start, y_end, N)    # creates a 1D-array with the y-coordinates

X,Y=meshgrid(x,y)

gamma = 5.0

x_vortex, y_vortex = 0.0, 0.0   # location of the source

function get_velocity_vortex(strength, xv, yv, X, Y):
      """Returns the velocity field generated by a vortex.

      Arguments
      ---------
      strength -- strength of the vortex.
      xv, yv -- coordinates of the vortex.
      X, Y -- mesh grid.
      """
      u = + strength/(2*pi).*(Y-yv)./((X-xv).^2+(Y-yv).^2)
      v = - strength/(2*pi).*(X-xv)./((X-xv).^2+(Y-yv).^2)

      return u, v
    end

function get_stream_function_vortex(strength, xv, yv, X, Y):
    """Returns the stream-function generated by a vortex.

    Arguments
    ---------
    strength -- strength of the vortex.
    xv, yv -- coordinates of the vortex.
    X, Y -- mesh grid.
    """
    psi = strength/(4*pi).*log((X-xv).^2+(Y-yv).^2)

    return psi
  end

  # computes the velocity field on the mesh grid
  u_vortex, v_vortex = get_velocity_vortex(gamma, x_vortex, y_vortex, X, Y)

  # computes the stream-function on the mesh grid
  psi_vortex = get_stream_function_vortex(gamma, x_vortex, y_vortex, X, Y)

  # plots the streamlines
  size = 10
  PyPlot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
  PyPlot.grid(true)
  PyPlot.xlabel("x", fontsize=16)
  PyPlot.ylabel("y", fontsize=16)
  PyPlot.xlim(x_start, x_end)
  PyPlot.ylim(y_start, y_end)
  PyPlot.streamplot(X, Y, u_vortex, v_vortex, density=2, linewidth=1, arrowsize=1, arrowstyle="->")
  PyPlot.scatter(x_vortex, y_vortex, color="#CD2305", s=80, marker="o")

strength_sink = -1.0            # strength of the sink
x_sink, y_sink = 0.0, 0.0       # location of the sink

function get_velocity_sink(strength, xs, ys, X, Y):
    """Returns the velocity field generated by a sink.

    Arguments
    ---------
    strength -- strength of the sink.
    xs, ys -- coordinates of the sink.
    X, Y -- mesh grid.
    """
    u = strength/(2*pi).*(X-xs)./((X-xs).^2+(Y-ys).^2)
    v = strength/(2*pi).*(Y-ys)./((X-xs).^2+(Y-ys).^2)

    return u, v
  end

function get_stream_function_sink(strength, xs, ys, X, Y):
      """Returns the stream-function generated by a sink.

      Arguments
      ---------
      strength -- strength of the sink.
      xs, ys -- coordinates of the sink.
      X, Y -- mesh grid.
      """
      psi = strength/(2*pi)*atan2((Y-ys), (X-xs))

      return psi
    end

    # computes the velocity field on the mesh grid
u_sink, v_sink = get_velocity_sink(strength_sink, x_sink, y_sink, X, Y)

# computes the stream-function on the mesh grid
psi_sink = get_stream_function_sink(strength_sink, x_sink, y_sink, X, Y)

# superposition of the sink and the vortex
u = u_vortex + u_sink
v = v_vortex + v_sink
psi = psi_vortex + psi_sink

# plots the streamlines
size = 10
PyPlot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
PyPlot.xlabel("x", fontsize=16)
PyPlot.ylabel("y", fontsize=16)
PyPlot.xlim(x_start, x_end)
PyPlot.ylim(y_start, y_end)
PyPlot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=1, arrowstyle="->")
PyPlot.scatter(x_vortex, y_vortex, color="#CD2305", s=80, marker="o");


######### CmdStan optimize program example  ###########

using Compat, CmdStan, Test, Statistics

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

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

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

  global stanmodel, rc, optim, cnames
  stanmodel = Stanmodel(Optimize(), name="bernoulli", model=bernoulli);

  rc, optim, cnames = stan(stanmodel, bernoullidata, ProjDir, CmdStanDir=CMDSTAN_HOME);

  if rc == 0
    println()
    display(optim)
    println()
    println("Test round.(mean(optim[1][\"theta\"]), digits=1) ≈ 0.3")
    @test round.(mean(optim["theta"]), digits=1) ≈ 0.3
  end
end # cd


using JuMP
using Ipopt
using CSV, DataFrames, DataStructures, MathOptFormat
using Statistics,Distributions, StatsFuns, StatsBase
using PyPlot

function critical_fractile(dist_type)
    #closed form of solution for newsvendor for normal distribution
    # dist_type - str, "norm" or "exp" for normal or exponential distribution, respectively

    b = 1 #underage cost
    h = 1 #overage cost

    if dist_type == "norm"
        mean = 50
        var = 50

        beta = b/(b+h)

        b_inv = norminvcdf(beta) #inverse of the cdf for normal distribution

        return mean + b_inv * var

    elseif dist_type == "exp"
        # plugged and chugged with cdf of exponential function to get warm start value (solver kept converging to point of infeasibility)
        cdf_func(x) = cdf(Exponential(50),x)

        # print(cdf_func(34.65))

        model3 = Model(with_optimizer(Ipopt.Optimizer,print_level=0))

        @variable(model3, x_var, start=34) #use warm start to avoid converge to infeasible point (see above)

        JuMP.register(model3, :cdf_func, 1, cdf_func; autodiff = true)  #to use cdf function in constraint

        @NLconstraint(model3,con1, cdf_func(x_var) >= b/(b+h))

        @objective(model3,Min,x_var)
        optimize!(model3)

        return value(x_var)

    else
        return nothing
    end
end

function scarf(df)
    #closed form scarf solution
    # df - dataframe, data of observed demand

    b = 1 #underage cost
    h = 1 #overage cost

    mu_hat = mean(df[:,"Column1"]) #sample mean
    sigma_sq_hat = var(df[:,"Column1"]) #sample variance

    return mu_hat + sigma_sq_hat/2 * (sqrt(b / h) - sqrt(h / b))
end

function SAA(df)
    # compute SAA solution (empirical risk minimization)
    # df - dataframe, data of observed demand

    b = 1 #underage cost
    h = 1 #overage cost

    model = Model(with_optimizer(Ipopt.Optimizer,print_level=0))

    data_size = length(df[:,"Column1"])

    @variable(model, x_var) #stocking quantity
    @variable(model, y_vars[1:data_size]) #variable to linearize max(d_i - x, x - d_i) in objective function

    @constraint(model,con1[i = 1:data_size], y_vars[i] >= x_var - df[:,"Column1"][i])
    @constraint(model, con2[i = 1:data_size], y_vars[i] >= df[:,"Column1"][i] - x_var )

    @NLobjective(model,Min, 1 / data_size * sum(y_vars[i] for i = 1:data_size))

    optimize!(model)

    return value(x_var)
end


using DifferentialEquations
using ModelingToolkit
using OrdinaryDiffEq
using DataFrames
using DataFrames
using StatsPlots
using BenchmarkTools

@parameters t β c γ
@variables S(t) I(t) R(t)
@derivatives D'~t
N=S+I+R # This is recognized as a derived variable
eqs = [D(S) ~ -β*c*I/N*S,
       D(I) ~ β*c*I/N*S-γ*I,
       D(R) ~ γ*I];

sys = ODESystem(eqs)
sys = ode_order_lowering(sys);

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

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

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

prob_ode = ODEProblem(sys,u0,tspan,p;jac=true)

sol_ode = solve(prob_ode);

df_ode = DataFrame(sol_ode(t)')
df_ode[!,:t] = t;

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

@benchmark solve(prob_ode)

include(joinpath(@__DIR__,"tutorials","appendix.jl"))
appendix()


function generate_primes(target)
    primes = [2]
    i = 3

    while length(primes) != target
        square_root = sqrt(i)
        for prime in primes

            if prime > square_root
                push!(primes, i)
                break
            end

            if i % prime == 0
                break
            end

        end
        i += 2
    end

    return primes[end]
end

println(generate_primes(10001))


@testset "qraddcol with β = 0" begin
    m = 100
    A = randn(m,0)
    R = zeros(0,0)

    for i in 1:m
        a = randn(m)
        R = qraddcol(A, R, a)
        A = [A a]
        @test  norm( R'*R - A'*A ) <1e-5
    end
end

@testset "qraddcol with β > 0" begin
    m = 100
    A = randn(m,0)
    R = zeros(0,0)
    β = 0.1
    for i in 1:m
        a = randn(m)
        B = [A a]
        R = qraddcol(A, R, a, β)
        @test norm( R'*R - B'*B - β^2*I ) <1e-5
    end
end

@testset "qraddrow" begin
    m, n = 3, 3
    A = randn(m,m)
    Q, R = qr(A)
    for i in 1:100
        a = randn(m)'
        A = [A; a]
        R = qraddrow(R, a)
        @test norm( R'R - A'*A ) <1e-5
    end

end

@testset "qrdelcol" begin
    m = 100
    A = randn(m,m)
    Q, R = qr(A)
    for i in 100:-1:1
        k = rand(1:i)
        A = A[:,1:i .!= k]
        R = qrdelcol(R, k)
        @test norm( R'*R - A'*A ) <1e-5#--> less_than(1e-5)
    end
end


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 9af04825-842b-4794-a317-6f1b78a64fb5
begin
	using Statistics, CSV, DataFrames
	using ForwardDiff
end

# ╔═╡ be7baa4f-6871-4d74-b4d2-1c2534cdfa17
# Row Orientation
begin
	using Random
	Random.seed!(0)

	struct Layer_dense
		w::Matrix{Float64}
		b::Matrix{Float64}
		σ::Function

		function Layer_dense(n_in::Int, n_out::Int; σ::Function = identity)
			w = 0.01 .* randn(n_out, n_in)
			b = reshape(zeros(n_out), (1, n_out))
			return new(w, b, σ)
		end
	end

	function (L::Layer_dense)(input)
		W, b, σ = L.w, L.b, L.σ
		σ( (input * W') .+ b )
	end
end

# ╔═╡ c4ea1432-abb4-11eb-1cf5-edac0735d67d
md"# LayerDense"

# ╔═╡ 45d31083-883b-4e84-913e-c040c6fdfd5f
md"""
>No futuro, podemos pensar em novas formas de fazer esse objeto. Por exemplo, passando a matriz W como input construtor do layer.
"""

# ╔═╡ dc76be1b-292e-4545-887e-1e62faf8e1b6
md"""
### Struct by Column Orientation
"""

# ╔═╡ fe237f5f-8f90-429a-8ddc-e5933b8fd808
begin
	Random.seed!(0)

	struct LayerDense{T<:Real, F<:Function}
		W::Matrix{T}
		b::Vector{T}
		σ::F

		function LayerDense(W::Matrix{T},
							b::Vector{T},
							σ::Function = identity) where {T<:Real}
			@assert size(W, 1) == length(b)

			new{T, typeof(σ)}(W, b, σ)
		end
	end

	function LayerDense(W::Matrix{T},
						b::Vector{S},
						σ::Function = identity) where {T<:Real, S<:Real}
		R = promote_type(T, S)
		return LayerDense(Matrix{R}(W), Vector{R}(b), σ)
	end

	function LayerDense(W::Matrix{T}, σ::Function = identity) where {T<:Real}
		b = zeros(T, size(W, 1))
		return LayerDense(W, b, σ)
	end

	function LayerDense(n_in::Int, n_out::Int, σ::Function = identity)
		W = 0.01 * randn(n_out, n_in)
		b = zeros(n_out)
		return LayerDense(W, b, σ)
	end

	function (L::LayerDense)(input::AbstractVector{<:Real})
		W, b, σ = L.W, L.b, L.σ
		return σ(W * input + b)
	end

	function (L::LayerDense)(input::AbstractMatrix{<:Real})
		W, b, σ = L.W, L.b, L.σ
		return σ(W * input .+ b)
	end
end

# ╔═╡ 4e6d165a-9707-4029-af4a-2356b0fe9c41
md"#### Example using the struct LayerDense"

# ╔═╡ 5d3ae0e6-7c5d-48b7-9845-d204a1df9934
W = rand(3, 4)

# ╔═╡ e85bc0f7-65ab-4a41-a4ba-e63c8b7c8184
B = LayerDense(W) # Layer de 4 inputs e 3 outputs

# ╔═╡ 142bb8f9-4188-46d2-94a2-8407d633c496
begin
	obs01 = [2, 3, 4, 5]
	obs02 = [2, 3, 4, 5]
	input = [obs01 obs02] # Each observation must be on a colunm
end

# ╔═╡ 2377f014-0bd0-457b-b9e9-3ace1ad7c331
B(input) # Each column is an output

# ╔═╡ c99c7f4f-7e62-4036-ba08-4478e5480d19
md"#### Adding sigmoid functions"

# ╔═╡ 96ce3aa7-9261-4e60-93ce-4d9321657d85
md"""
Função **Logística**: \
$$\;\;\;\;\;$$ $$\sigma(x) = {1 \over 1 + e^{-x}}$$

Na implementação via código é necessário fazer uma adaptação que evita problemas numéricos de *Overflow* resultando na seguinte função: \
$$\sigma(x) = \left\{
  \begin{array}{lr}
    {1 \over 1 + e^{-|x|}}, & x \ge 0\\
    {e^{-|x|} \over 1 + e^{-|x|}}, & x < 0
  \end{array}
\right.$$ \

Para isso, utilizamos o módulo de $$x$$ na expressão $$e^{-|x|}$$ que possui a seguinte propriedade: \
$$e^{-|x|} = \left\{
  \begin{array}{lr}
    e^{-x}, & x \ge 0\\
    e^x, & x < 0
  \end{array}
\right.$$

Com o modulo, quando x é positivo ou zero ($$x ≥ 0$$) a adaptação resulta na própria função logística, mas quando x é negativo são necessários os seguintes passos: \
$$Se \space x \ge 0:$$  $$\;\;\;\;\;$$  $$\sigma(x) = {1 \over 1 + e^{-|x|}}$$ \
$$Se \space x < 0:$$ $$\;\;\;\;\;$$ $$1 = e^{-x} ⋅ e^x → {1 \over 1 + e^{-x}} = {e^{-x} ⋅ e^x \over e^{-x} ⋅ e^x + e^{-x}} = {e^{-x} ⋅ e^x \over e^{-x}(e^x + 1)} = {e^x \over 1 + e^x} = {e^{-|x|} \over 1 + e^{-|x|}}$$
"""

# ╔═╡ 58265283-a36a-4c83-9101-3387165da021
begin
	function σ(x::Real)
		t = exp(-abs(x))
		return ifelse(x ≥ 0, inv(1 + t), t / (1 + t))
	end

	σ(x::AbstractVecOrMat{T}) where {T<:Real} = σ.(x)

	const sigmoid = σ
end

# ╔═╡ fc1ddc84-62dc-4af1-9b9e-9d37f581a252
md"""
#### ReLU function
$$ReLU(x) = \left\{
  \begin{array}{lr}
    x, & x \ge 0\\
    0, & x < 0
  \end{array}
\right.$$
$$Or$$
$$ReLU(x) = max(0, x)$$
"""

# ╔═╡ 45bafada-34ec-4c41-b0a2-000f35ebffef
begin
	relu(x::Real) = max(zero(x), x) # zero function keep the type of x
	relu(x::AbstractVecOrMat{<:Real}) = relu.(x)
end

# ╔═╡ c85cd202-7223-4c46-80ef-afd750f28ead
md"""
$$softplus(x) = ln(1 + e^x)$$
$$For \space x > 0:$$
$$1 = e^{-x} ⋅ e^x → ln(1 + e^x) = ln(e^{-x} ⋅ e^x + e^x) = ln(e^x(e^{-x} + 1)) = ln(e^x) + ln(1 + e^{-x})$$
$$softplus(x) = x + ln(1 + e^{-x})$$
"""

# ╔═╡ d6fce490-1498-469d-8ecc-55d0a71d1565
begin
	#log1p(x) is the seme as log(1 + x)
	softplus(x::Real) = ifelse(x > 0, x + log(1 + exp(-x)), log(1 + exp(x)))
	softplus(x::AbstractVecOrMat{<:Real}) = softplus.(x)
end

# ╔═╡ 04cadf6a-4fd1-4b5f-aa4d-140ea6e5cff2
function softmax(x::Vector{T}) where {T<:Real}
	m = maximum(x)
	exp_val = exp.(x .- m)
	s = sum(exp_val)
	return exp_val ./ s
end

# ╔═╡ a7898ffd-1acb-4220-bbc8-3b428b95db85
function softmax(x::Matrix{T}) where {T<:Real}
	m = maximum(x, dims = 1)
	exp_val = exp.(x .- m)
	s = sum(exp_val, dims = 1)
	return exp_val ./ s
end

# ╔═╡ c116a1e4-fd4d-446f-bdea-02301f43b2c0
md"##### Example with sigmoid functions"

# ╔═╡ d0c91a1c-8798-4985-978b-4c47d1abc212
C = LayerDense(2, 3, σ)

# ╔═╡ 3de0841e-f0f3-436c-a891-dd4af32d2af6
# C(randn(2, 3))
C([[1, 4] [2, 3] [3, 5]])

# ╔═╡ 963f6276-2440-4bfe-ae0f-539d8bfae0a2
D = LayerDense(2, 3, softmax)

# ╔═╡ 56191b0c-fe49-4ed4-9ca7-2e22e7b028ba
md"""
#### Cost Function
"""

# ╔═╡ 7b5a6859-a9e1-48b1-a0fd-1cde06249d03
begin
	function xlogy(x::Real, y::Real)
		result = x * log(y)
		ifelse(iszero(x) && !isnan(y), zero(result), result)
	end

	function crossentropy(ŷ::AbstractVecOrMat{<:Real},
						  y::AbstractVecOrMat{<:Real};
						  dims::Int = 1,
						  agg::Function = mean)
    	agg(.-sum(xlogy.(y, ŷ), dims = dims))
	end
end

# ╔═╡ 4b070348-fdbc-46cd-85f0-29b64432ae21
begin
	outp2 = D([1 4; 2 3; 3 5]')
	target2 = [0, 1, 0]
	crossentropy(outp2, target2)
end

# ╔═╡ e3ec0333-14ad-42a8-9cd3-4e3937dc5ff0
begin
	outp1 = [0.7 0.1 0.2
			 0.1 0.5 0.4
	   		 0.02 0.9 0.08]'
	target1 = [1 0 0
			   0 1 0
		  	   0 1 0]'
	crossentropy(outp1, target1)
end

# ╔═╡ 8f3cbc01-53ba-4ff0-8245-564f8fee53f2
md"##### Importing data from nnfs python package"

# ╔═╡ 099f5fa8-6687-4c77-9f61-119d75377bab
data = CSV.read("spiral_data.csv", DataFrame)

# ╔═╡ b5981a68-16e7-451f-9e13-d74b2b78a958
md"#### Chain type, params(), gradient()"

# ╔═╡ d2fa0422-7f39-4365-8f3b-bfa6997ae0ae
begin
	Random.seed!(1998)

	L1 = LayerDense(10, 5, σ)
	L2 = LayerDense(5, 2, softmax)

	output = L2(L1(rand(10)))
end

# ╔═╡ 72f68243-8ae1-465c-854d-a76c07c5e346
function params(layers)
	p = []
	for layer in layers
		push!(p, layer.W[:], layer.b )
	end
	return vcat(p...)
end

# ╔═╡ e84cff98-1bea-42cb-8e8d-6f5ef08226f4
begin
	Θ = params([L1, L2])

	U = (
		size(L1.W, 1)*size(L1.W, 2), length(L1.b),
		size(L2.W, 1)*size(L2.W, 2), length(L2.b),
	)
	U = cumsum(U)

	L1w = reshape( Θ[ 1:U[1] ], size(L1.W, 1), size(L1.W, 2))

	L1b = Θ[ U[1]+1:U[2] ]

	L2w = reshape( Θ[ U[2]+1:U[3] ], size(L2.W, 1), size(L2.W, 2))

	L2b = Θ[ U[3]+1:U[4] ]

	( length(Θ), U, (L1w, L1b, L2w, L2b) )
end

# ╔═╡ a5a511f7-38f0-4074-9192-7f11d611bf42
begin

	struct Chain2
		Layers::Vector{LayerDense{Float64}}
		Wl::Vector{Int}
		Wc::Vector{Int}
		b::Vector{Int}

		function Chain(L...)
			Wl = Int[]
			Wc = Int[]
			b = Int[]
			for (i, l) in enumerate(L)
				Wl[i], Wc[i] = size(l.W)
				b[i] = length(l.b)
			end
			return new{}(L, Wl, Wc, b)
		end

	end

end

# ╔═╡ 154e4560-0769-429e-89bd-3b0392209d6d
typeof( (1.3,2.12) )

# ╔═╡ 74ccd8c6-f513-4178-a41f-7b70835516e7
begin
	struct Chain{T<:Tuple}
		layers::T
		Chain(layers...) = new{typeof(layers)}(layers)
	end

	function (L::Chain)(input::AbstractVecOrMat{<:Real})
		return ∘(reverse(L.layers)...)(input)
	end
end

# ╔═╡ 4504eaff-0c86-4606-a410-c94a8e8a653e
begin
	#L1 = LayerDense(10, 5, σ)
	#L2 = LayerDense(5, 2, softmax)

	C1 = Chain(L1, L2)
	typeof(C1.layers)
end

# ╔═╡ ac678c79-a804-44e0-80d8-15a40dff327d
begin
	p = []
	for layer in C1.layers
		push!(p, layer.W[:], layer.b )
	end
	vcat(p...)
end

# ╔═╡ 4b019add-7eb4-496a-b0ef-3c3eee9d85bd
begin
	C1(rand(10,12))
end

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[compat]
CSV = "~0.8.5"
DataFrames = "~1.2.2"
ForwardDiff = "~0.10.19"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[CSV]]
deps = ["Dates", "Mmap", "Parsers", "PooledArrays", "SentinelArrays", "Tables", "Unicode"]
git-tree-sha1 = "b83aa3f513be680454437a0eee21001607e5d983"
uuid = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
version = "0.8.5"

[[ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "bdc0937269321858ab2a4f288486cb258b9a0af7"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.3.0"

[[CommonSubexpressions]]
deps = ["MacroTools", "Test"]
git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7"
uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
version = "0.3.0"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "727e463cfebd0c7b999bbf3e9e7e16f254b94193"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.34.0"

[[CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "ee400abb2298bd13bfc3df1c412ed228061a2385"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.7.0"

[[DataFrames]]
deps = ["Compat", "DataAPI", "Future", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrettyTables", "Printf", "REPL", "Reexport", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"]
git-tree-sha1 = "d785f42445b63fc86caa08bb9a9351008be9b765"
uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
version = "1.2.2"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[DiffResults]]
deps = ["StaticArrays"]
git-tree-sha1 = "c18e98cba888c6c25d1c3b048e4b3380ca956805"
uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
version = "1.0.3"

[[DiffRules]]
deps = ["NaNMath", "Random", "SpecialFunctions"]
git-tree-sha1 = "3ed8fa7178a10d1cd0f1ca524f249ba6937490c0"
uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
version = "1.3.0"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "a32185f5428d3986f47c2ab78b1f216d5e6cc96f"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.5"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[ForwardDiff]]
deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "NaNMath", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
git-tree-sha1 = "b5e930ac60b613ef3406da6d4f42c35d8dc51419"
uuid = "f6369f11-7733-5829-9624-2563aa707210"
version = "0.10.19"

[[Future]]
deps = ["Random"]
uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[InvertedIndices]]
deps = ["Test"]
git-tree-sha1 = "15732c475062348b0165684ffe28e85ea8396afc"
uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
version = "1.0.0"

[[IrrationalConstants]]
git-tree-sha1 = "f76424439413893a832026ca355fe273e93bce94"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.0"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[LogExpFunctions]]
deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "3d682c07e6dd250ed082f883dc88aee7996bf2cc"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.0"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "0fb723cd8c45858c22169b2e42269e53271a6df7"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.7"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "2ca267b08821e86c5ef4376cffed98a46c2cb205"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.1"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[NaNMath]]
git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
version = "0.3.5"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "bfd7d8c7fd87f04543810d9cbd3995972236ba1b"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "1.1.2"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PooledArrays]]
deps = ["DataAPI", "Future"]
git-tree-sha1 = "a193d6ad9c45ada72c14b731a318bedd3c2f00cf"
uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
version = "1.3.0"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "0d1245a357cc61c8cd61934c07447aa569ff22e6"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.1.0"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[SentinelArrays]]
deps = ["Dates", "Random"]
git-tree-sha1 = "54f37736d8934a12a200edea2f9206b03bdf3159"
uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
version = "1.3.7"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "b3363d7460f7d098ca0912c69b082f75625d7508"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.0.1"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[SpecialFunctions]]
deps = ["ChainRulesCore", "LogExpFunctions", "OpenSpecFun_jll"]
git-tree-sha1 = "a322a9493e49c5f3a10b50df3aedaf1cdb3244b7"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "1.6.1"

[[StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "3240808c6d463ac46f1c1cd7638375cd22abbccb"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "1.2.12"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "d0c690d37c73aeb5ca063056283fde5585a41710"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.5.0"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─c4ea1432-abb4-11eb-1cf5-edac0735d67d
# ╟─45d31083-883b-4e84-913e-c040c6fdfd5f
# ╟─dc76be1b-292e-4545-887e-1e62faf8e1b6
# ╠═fe237f5f-8f90-429a-8ddc-e5933b8fd808
# ╟─4e6d165a-9707-4029-af4a-2356b0fe9c41
# ╠═5d3ae0e6-7c5d-48b7-9845-d204a1df9934
# ╠═e85bc0f7-65ab-4a41-a4ba-e63c8b7c8184
# ╠═142bb8f9-4188-46d2-94a2-8407d633c496
# ╠═2377f014-0bd0-457b-b9e9-3ace1ad7c331
# ╟─c99c7f4f-7e62-4036-ba08-4478e5480d19
# ╟─96ce3aa7-9261-4e60-93ce-4d9321657d85
# ╠═58265283-a36a-4c83-9101-3387165da021
# ╟─fc1ddc84-62dc-4af1-9b9e-9d37f581a252
# ╠═45bafada-34ec-4c41-b0a2-000f35ebffef
# ╟─c85cd202-7223-4c46-80ef-afd750f28ead
# ╠═d6fce490-1498-469d-8ecc-55d0a71d1565
# ╠═04cadf6a-4fd1-4b5f-aa4d-140ea6e5cff2
# ╠═a7898ffd-1acb-4220-bbc8-3b428b95db85
# ╟─c116a1e4-fd4d-446f-bdea-02301f43b2c0
# ╠═d0c91a1c-8798-4985-978b-4c47d1abc212
# ╠═3de0841e-f0f3-436c-a891-dd4af32d2af6
# ╠═963f6276-2440-4bfe-ae0f-539d8bfae0a2
# ╟─56191b0c-fe49-4ed4-9ca7-2e22e7b028ba
# ╠═7b5a6859-a9e1-48b1-a0fd-1cde06249d03
# ╠═4b070348-fdbc-46cd-85f0-29b64432ae21
# ╠═e3ec0333-14ad-42a8-9cd3-4e3937dc5ff0
# ╟─8f3cbc01-53ba-4ff0-8245-564f8fee53f2
# ╟─099f5fa8-6687-4c77-9f61-119d75377bab
# ╠═9af04825-842b-4794-a317-6f1b78a64fb5
# ╟─b5981a68-16e7-451f-9e13-d74b2b78a958
# ╠═d2fa0422-7f39-4365-8f3b-bfa6997ae0ae
# ╠═ac678c79-a804-44e0-80d8-15a40dff327d
# ╠═72f68243-8ae1-465c-854d-a76c07c5e346
# ╠═e84cff98-1bea-42cb-8e8d-6f5ef08226f4
# ╠═a5a511f7-38f0-4074-9192-7f11d611bf42
# ╠═154e4560-0769-429e-89bd-3b0392209d6d
# ╠═74ccd8c6-f513-4178-a41f-7b70835516e7
# ╠═4504eaff-0c86-4606-a410-c94a8e8a653e
# ╠═4b019add-7eb4-496a-b0ef-3c3eee9d85bd
# ╟─be7baa4f-6871-4d74-b4d2-1c2534cdfa17
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


function defaultfcn(z)
    f = 1e9
    ϵᵣ = 5 - 2im
    μᵣ = 1 - 2im
    d = 1e-2
    c = 3e8
    ω = 2π*f
    k₀ = ω/c
    cc = ϵᵣ^2*(k₀*d)^2*(ϵᵣ*μᵣ - 1)
    w = ϵᵣ^2*z^2 + z^2*tan(z)^2 - cc
end

# Analysis parameters
xb = -2  # real part begin
xe = 2  # real part end
yb = -2  # imag part begin
ye = 2  # imag part end
r = 0.1999  # initial mesh step

origcoords = rectangulardomain(complex(xb, yb), complex(xe, ye), r)

# matlab results from https://github.com/PioKow/GRPF for comparison
matlab_zroots = [-1.624715288135189 + 0.182095877702038im,
                 -1.520192978034417 - 0.173670452237129im,
                 -0.515113098919392 + 0.507111597359180im,
                  0.515113098795215 - 0.507111597284675im,
                  1.520192978034417 + 0.173670452237129im,
                  1.624715288135189 - 0.182095877702037im]

matlab_zpoles = [-1.570796326699632 - 0.000000000206961im,
                  1.570796326699632 - 0.000000000206961im]

zroots, zpoles = grpf(defaultfcn, origcoords)

@test length(zroots) == 6
@test length(zpoles) == 2

@test approxmatch(zroots, matlab_zroots)
@test approxmatch(zpoles, matlab_zpoles)

pzroots, pzpoles, quadrants, phasediffs, tess, g2f = grpf(defaultfcn, origcoords, PlotData())

@test approxmatch(pzroots, matlab_zroots)
@test approxmatch(pzpoles, matlab_zpoles)

tmpr, tmpp = grpf(defaultfcn, origcoords, GRPFParams(8000, 1e-9))

@test approxmatch(tmpr, zroots)
@test approxmatch(tmpp, zpoles)

tmpr, tmpp, quadrants, phasediffs, tess, g2f = grpf(defaultfcn, origcoords, PlotData(), GRPFParams(8000, 1e-9))

@test approxmatch(tmpr, zroots)
@test approxmatch(tmpp, zpoles)

tmpr, tmpp = grpf(defaultfcn, origcoords, GRPFParams(8000, 1e-9, true))

@test approxmatch(tmpr, zroots)
@test approxmatch(tmpp, zpoles)

# Test with very low maxnodes
maxnodes = 10
params = GRPFParams(100, maxnodes, 3, maxnodes-1, 1e-9, false)
@test_logs (:warn,"GRPFParams.maxnodes reached") grpf(defaultfcn, origcoords, PlotData(), params)

# Test with big origcoords
xb = big(xb)  # real part begin
xe = big(xe)  # real part end
yb = big(yb)  # imag part begin
ye = big(ye)  # imag part end
r = big(r)  # initial mesh step

origcoords = rectangulardomain(complex(xb, yb), complex(xe, ye), r)

bzroots, bzpoles = grpf(defaultfcn, origcoords)

@test all(isa.(bzroots, Complex{BigFloat}))
@test all(isa.(bzpoles, Complex{BigFloat}))

@test approxmatch(bzroots, zroots)
@test approxmatch(bzpoles, zpoles)

# Test region with no roots or poles
xb, xe = -1, 1
yb, ye = -2, -1
origcoords = rectangulardomain(complex(xb, yb), complex(xe, ye), r)

nozroots, nozpoles = grpf(defaultfcn, origcoords)

@test length(nozroots) == 0
@test length(nozpoles) == 0

nozroots, nozpoles, quadrants, phasediffs, tess, g2f = grpf(defaultfcn, origcoords, PlotData())

@test length(nozroots) == 0
@test length(nozpoles) == 0

xb, xe = big"-1", big"1"
yb, ye = big"-2", big"-1"
r = big"0.2"

origcoords = rectangulardomain(complex(xb, yb), complex(xe, ye), r)

bnozroots, bnozpoles = grpf(defaultfcn, origcoords)

@test eltype(bnozroots) == Complex{BigFloat}
@test eltype(bnozpoles) == Complex{BigFloat}

@test length(bnozroots) == 0
@test length(bnozpoles) == 0


#A deck with 52 cards - 4 suites and 13 ranks - A,2,...,10,J,Q,K
deck = [(suite,r) for suite in ["♡","♢", "♣", "♠"]  for r in range(1,13)]

function simulate_pair(n = 100, verbose=false)
    pairs = 0
    for i in 1:n
        hand = shuffle(deck)[1:2]
        if length(unique([r for (s,r) in hand])) == 1
                if verbose
                    println(hand)
                end
                pairs += 1
        end
    end
    100.*pairs/n
end
good_pairs(len_deck = 52) = (len_deck*3)/2
all_pairs(len_deck = 52) = binomial(len_deck, 2)

function simulate_disjoint_triplet(n = 100, verbose=false)
    triplets = 0
    for i in 1:n
        hand = shuffle(deck)[1:3]
        if length(unique([r for (s,r) in hand])) == 3
                if verbose
                    println(hand)
                end
                triplets += 1
        end
    end
    100.*triplets/n
end
disjoint_triplets(len_deck = 52) = (len_deck*(len_deck-4)*(len_deck-8))/factorial(3)
all_triplets(len_deck = 52) = binomial(len_deck, 3)

function simulate_one_pair_hand(n = 100, verbose=false)
    pairs = 0
    for i in 1:n
        hand = shuffle(deck)[1:5]
        if length(unique([r for (s,r) in hand])) == 4
                if verbose
                    println(hand)
                end
                pairs += 1
        end
    end
    100.*pairs/n
end
prob_one_pair_hand = 100.* good_pairs(52) * disjoint_triplets(48) / binomial(52,5)


export tanhActivation, tanhActivation!

"""
 hyperbolic tan activation A = tanh(Y)

 Input:

   Y - array of features

 Optional Input:

   doDerivative - flag for computing derivative, set via varargin
                  Ex: tanhActivation(Y,true);

 Output:

  A  - activation
  dA - derivatives
"""
function tanhActivation(Y::Array{T,2},doDerivative::Bool=false) where {T <: Number}

    A = tanh.(Y)
    dA = zeros(A)
    if doDerivative
         dA .= one(T) .- A.^2
    end
    return A, dA
end

function tanhActivation!(A::Array{T,2},dA=[],doDerivative::Bool=false) where {T <: Number}

    A .= tanh.(A)
    if doDerivative
        if isempty(dA)
            dA = one(T) .- A.^2
        else
            dA .= one(T) .- A.^2
        end
    end
    return A, dA
end


# # Minimizing the sum of the k-largest λ
#
# We show how to find the sum of absolute value of the `k` largest eigenvalues of a symmetric matrix $A \in \mathbb{S}^n$. This problem can be solved as a semidefinite program. The primal and dual forms are stated in *Alizadeh* \[1\]:
# $$
# \begin{array}{llll} \text{maximize} &   \text{Tr}(AY) - \text{Tr}(AW) &    \text{minimize} &  kz + Tr(U) + Tr(V)            \\
# \text{subject to} &  \text{Tr}(Y + W) = k                             &   \text{subject to} &   zI + V - A \succeq 0            \\
#                   &  0 \preceq  Y \preceq I                           &  &                    zI + U + A \succeq 0             \\
#                   &  0 \preceq  W \preceq I                           & &                     U, V \succeq 0,
# \end{array}
# $$
# where $Y, W$ are the variables of the primal and $U, V$ are the variables of the dual problem.
#-
using LinearAlgebra, JuMP, COSMO, Random
rng = Random.MersenneTwister(212)

n = 10
A = 5 .* randn(rng, 10, 10)
A = Symmetric(A, :U)

# We are interested in minimizing the sum of absolute values of the `k=3` largest eigenvalues. Let's formulate the problem in `JuMP` with `COSMO` as the backend solver:

k = 3
model = JuMP.Model(optimizer_with_attributes(COSMO.Optimizer, "verbose" => true));
@variable(model, Y[1:n, 1:n], PSD);
@variable(model, W[1:n, 1:n], PSD);

@objective(model, Max, tr(A * Y) - tr(A * W));
@constraint(model, tr(Y + W) == k);
@constraint(model, Symmetric(I - Y) in PSDCone());
@constraint(model, Symmetric(I - W) in PSDCone());
status = JuMP.optimize!(model)

#-
opt_objective = JuMP.objective_value(model)

# Now, we can check the solution by computing the sum of the absolute value of the 3-largest eigenvalues:
k_λ_abs = sum(sort(abs.(eigen(A).values), rev = true)[1:k])

# ### Solve the dual
#
# Alternatively, we can solve the dual problem:
#-
model = JuMP.Model(with_optimizer(COSMO.Optimizer, verbose=true));
@variable(model, V[1:n, 1:n], PSD);
@variable(model, U[1:n, 1:n], PSD);
@variable(model, z);

@objective(model, Min, k * z + tr(V) + tr(U));
@constraint(model, Symmetric(z .* diagm(0 => ones(n)) + V - A) in PSDCone());
@constraint(model, Symmetric(z .* diagm(0 => ones(n)) + U + A) in PSDCone());
status = JuMP.optimize!(model)

#-
opt_objective = JuMP.objective_value(model)

# This gives the same result.
# ## Problem with A as variable
#
# Above problems are mostly helpful for illustrative purpose. It is obviously easier to find the sum of the k-largest eigenvalues by simply computing the eigenvalues of $A$. However, above results become useful if finding $A$ itself is part of the problem. For example, assume we want to find a valid matrix $A$ under the constraints: $C\, \text{vec}(A) = b$ with the minimum sum of absolute values of the k-largest eigenvalues. We can then solve the equivalent problem:
# $$
# \begin{array}{ll} \text{minimize} &  kz + Tr(U) + Tr(V)     \\
#  \text{subject to} &   C \text{vec}(A) = b \\
#                    & zI + V - A \succeq 0            \\
#                    &    zI + U + A \succeq 0             \\
#                    &   U, V \succeq 0.
# \end{array}
# $$
#
# ## References
# [1] Alizadeh - Interior point methods in semidefinite programming with applications to combinatorial optimization (1995)


using CompScienceMeshes
using SauterSchwabQuadrature

pI = point(1,5,3)
pII = point(2,5,3)
pIII = point(7,1,0)

Sourcechart = simplex(pI, pII, pIII)
Testchart = simplex(pII, pIII, pI)

Accuracy = 12

function integrand(x,y)
			return(((x-pI)'*(y-pII))*exp(-im*1*norm(x-y))/(4pi*norm(x-y)))
end

result = sauterschwabintegral(Sourcechart, Testchart, integrand, Accuracy, Accuracy)
println(result)

#=For those who want to test the sauterschwab_nonparameterized() function,
may uncomment the following five lines=#

#sourcechart = simplex(pI,pII,pIII)
#testchart = simplex(pI,pII,pIII)
#cf = CommonFace(Accuracy)
#result2 = sauterschwab_nonparameterized(sourcechart, testchart, integrand, cf)
#println(result2)

#=sourcechart = testchart, hence the required condition is fulfilled.
The user may also compare the two results and see that both are equal=#


function normalize(array, target_range)
    target_min, target_max = target_range
    orig_min, orig_max = minimum(array), maximum(array)
    normalized_array = target_min .+ ((target_max-target_min).*(array.-orig_min))./(orig_max-orig_min)
    return normalized_array
end
a = [1,2,3,4]
print(normalize(a, [-1,1]))


## Nerlove model by OLS (fminunc)
# and fmincon, imposing that factor shares are in (0,1) and sum to 1
using DelimitedFiles, Econometrics
## prepare the data
cd(@__DIR__)
data = readdlm("../Data/nerlove.data")
data = log.(data[:,2:end])
n = size(data,1)
y = data[:,1]
x = [ones(n,1) data[:,2:5]]

## bounds and restriction
lb = [-1e6, -1e6, 0., 0., 0.0]
ub = [1e6, 1e6, 1., 1., 1.]
R = [0. 0. 1. 1. 1.]
r = [1.0]

## define the objective function and start value
obj = theta -> (y-x*theta)'*(y-x*theta)
startval = (ub+lb)/2.0

## OLS
thetahat, objvalue = fminunc(obj, startval)
println("the OLS estimates: obj. value: ", round(objvalue,digits=5))
prettyprint(thetahat)

## restricted LS
thetahat, objvalue_r, flag = fmincon(obj, startval, R, r, lb, ub) # both lower and upper bounds
println("the restricted LS estimates: obj. value: ", round(objvalue_r,digits=5))
prettyprint(thetahat)

println("Exercise: you should construct a qF test using the unrestricted and")
println("restricted objective function values")


function hailstonelength(n::Integer)
    len = 1
    while n > 1
        n = ifelse(iseven(n), n ÷ 2, 3n + 1)
        len += 1
    end
    return len
end

@show hailstonelength(27); nothing
@show findmax([hailstonelength(i) for i in 1:100_000]); nothing


# Load Turing.
using Turing

# Load RDatasets.
using RDatasets

# Load StatsPlots for visualizations and diagnostics.
using StatsPlots

# Functionality for splitting and normalizing the data.
using MLDataUtils: shuffleobs, splitobs, rescale!

# We need a softmax function which is provided by NNlib.
using NNlib: softmax

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

# Hide the progress prompt while sampling.
Turing.setprogress!(false);

# Import the "iris" dataset.
data = RDatasets.dataset("datasets", "iris");

# Show twenty random rows.
data[rand(1:size(data, 1), 20), :]

# Recode the `Species` column.
species = ["setosa", "versicolor", "virginica"]
data[!, :Species_index] = indexin(data[!, :Species], species)

# Show twenty random rows of the new species columns
data[rand(1:size(data, 1), 20), [:Species, :Species_index]]

# Split our dataset 50%/50% into training/test sets.
trainset, testset = splitobs(shuffleobs(data), 0.5)

# Define features and target.
features = [:SepalLength, :SepalWidth, :PetalLength, :PetalWidth]
target = :Species_index

# Turing requires data in matrix and vector form.
train_features = Matrix(trainset[!, features])
test_features = Matrix(testset[!, features])
train_target = trainset[!, target]
test_target = testset[!, target]

# Standardize the features.
μ, σ = rescale!(train_features; obsdim = 1)
rescale!(test_features, μ, σ; obsdim = 1);

# Bayesian multinomial logistic regression
@model function logistic_regression(x, y, σ)
    n = size(x, 1)
    length(y) == n || throw(DimensionMismatch("number of observations in `x` and `y` is not equal"))

    # Priors of intercepts and coefficients.
    intercept_versicolor ~ Normal(0, σ)
    intercept_virginica ~ Normal(0, σ)
    coefficients_versicolor ~ MvNormal(4, σ)
    coefficients_virginica ~ MvNormal(4, σ)

    # Compute the likelihood of the observations.
    values_versicolor = intercept_versicolor .+ x * coefficients_versicolor
    values_virginica = intercept_virginica .+ x * coefficients_virginica
    for i in 1:n
        # the 0 corresponds to the base category `setosa`
        v = softmax([0, values_versicolor[i], values_virginica[i]])
        y[i] ~ Categorical(v)
    end
end;

chain = sample(logistic_regression(train_features, train_target, 1), HMC(0.05, 10), MCMCThreads(), 1500, 3)

plot(chain)

corner(
    chain, MCMCChains.namesingroup(chain, :coefficients_versicolor);
    label=[string(i) for i in 1:4]
)

corner(
    chain, MCMCChains.namesingroup(chain, :coefficients_virginica);
    label=[string(i) for i in 1:4]
)

function prediction(x::Matrix, chain)
    # Pull the means from each parameter's sampled values in the chain.
    intercept_versicolor = mean(chain, :intercept_versicolor)
    intercept_virginica = mean(chain, :intercept_virginica)
    coefficients_versicolor = [
        mean(chain, k) for k in
        MCMCChains.namesingroup(chain, :coefficients_versicolor)
    ]
    coefficients_virginica = [
        mean(chain, k) for k in
        MCMCChains.namesingroup(chain, :coefficients_virginica)
    ]

    # Compute the index of the species with the highest probability for each observation.
    values_versicolor = intercept_versicolor .+ x * coefficients_versicolor
    values_virginica = intercept_virginica .+ x * coefficients_virginica
    species_indices = [argmax((0, x, y)) for (x, y) in zip(values_versicolor, values_virginica)]

    return species_indices
end;

# Make the predictions.
predictions = prediction(test_features, chain)

# Calculate accuracy for our test set.
mean(predictions .== testset[!, :Species_index])

for s in 1:3
    rows = testset[!, :Species_index] .== s
    println("Number of `", species[s], "`: ", count(rows))
    println("Percentage of `", species[s], "` predicted correctly: ",
        mean(predictions[rows] .== testset[rows, :Species_index]))
end

if isdefined(Main, :TuringTutorials)
    Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder], WEAVE_ARGS[:file])
end


# This file contains some implementations of prox-operators

# L1 - PROX - OPERATOR
# Param : vector to
function L1prox(x, lambda)
    prox = zeros(size(x,1))
    for index in 1:size(x,1)
        if x[index] < -lambda
            prox[index] = x[index] + lamda
        elseif x[index] > lambda
            prox[index] = x[index] - lambda
        else
            prox[index] = 0
        end
    end
    return prox
end

# L2 - PROX - OPERATOR
# Param :
#   - x = vector to apply proxity
#   - λ = regularization parameter
function L2prox(x,lambda)
    # create empty prox vector
    prox = zeros(size(x,1))
    l2norm = norm(x,2)
    if l2norm > lambda
        return (1-lambda/l2norm)*x
    else
        return zeros(size(x,1))
    end
end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

# Semiring definitions
abstract type Semiring end

import Base: zero, one, +, *, /
import Base: isapprox, parse, reverse, convert, get
isapprox(a::T,b::T) where { T <: Semiring } = isapprox(a.x,b.x)
reversetype(::Type{S}) where {S <: Semiring} = S
reverse(a::S) where { S <: Semiring } = a
reverseback(a::S) where { S <: Semiring } = a
get(a::S) where { S <: Semiring } = a.x

# properties
"""
`isleft(::Type{W})`

Check if the semiring type `W` satisfies:

``\\forall a,b,c \\in \\mathbb{W} : c \\otimes(a \\oplus b) = c \\otimes a \\oplus c \\otimes b``
"""
isleft(::Type{W}) where {W} = false             # ∀ a,b,c: c*(a+b) = c*a + c*b

"""
`isright(::Type{W})`

Check if the semiring type `W` satisfies:

``\\forall a,b,c \\in \\mathbb{W} : c \\otimes(a \\oplus b) = a \\otimes c \\oplus b \\otimes c``
"""
isright(::Type{W}) where {W}  = false           # ∀ a,b,c: c*(a+b) = a*c + b*c

"""
`isweaklydivisible(::Type{W})`

Check if the semiring type `W` satisfies:

``\\forall a,b \\in \\mathbb{W} \\ \\text{s.t.} \\ a \\oplus b \\neq \\bar{0} \\ \\exists z : x = (x \\oplus y ) \\otimes z``
"""
isweaklydivisible(::Type{W}) where {W}  = false # a+b ≂̸ 0: ∃ z s.t. x = (x+y)*z

"""
`ispath(::Type{W})`

Check if the semiring type `W` satisfies:

``\\forall a,b \\in \\mathbb{W}: a \\oplus b = a \\lor a \\oplus b = b``
"""
ispath(::Type{W}) where {W}  = false            # ∀ a,b: a+b = a or a+b=b

"""
`isidempotent(::Type{W})`

Check if the semiring type `W` satisfies:

``\\forall a \\in \\mathbb{W}: a \\oplus a = a``
"""
isidempotent(::Type{W}) where {W} = false       # ∀ a: a+a = a

"""
`iscommulative(::Type{W})`

Check if the semiring type `W` satisfies:

``\\forall a,b \\in \\mathbb{W}: a \\otimes b = b \\otimes a``
"""
iscommulative(::Type{W}) where {W} = false      # ∀ a,b: a*b = b*a
iscomplete(::Type{W}) where {W} = false

Base.show(io::IO, T::Semiring) = Base.show(io, T.x)


module kmean
using Distributions
using Random
using Printf

function k_mean(X::AbstractMatrix{<:Real},
     k::Integer,
     w::Array{<:Real};
     maxiter::Integer=200, # maximun iteration
     tol::Real=1e-6 # tolearance of change at converage
     )
     d, n = size(X)
     (2 <= k < n) || throw(ArgumentError("k must be 2 <= k < n, k=$k given."))
     centers =  init_centers(X, k)
     # println(centers)
     centers_new = Array{Float64, 2}(undef, d, k)
     # println(centers)
     clustDist = zeros(n, 2)
     delta = 1.0
     iter_num = 1
     inertia = 0
     while delta > tol
         if iter_num > maxiter
             break
         end

         for i in range(1, length=n, step =1)
             distList = Float64[]
             for j in range(1, length=k, step=1)
                 append!(distList, distance(centers[:,j], X[:, i]))
             end
             # println(distList)
             dist, ind = findmin(distList)
             # print(ind)
             # * w[i]
             clustDist[i, :] = [dist, ind]
         end
         inertia = 0
         for j in range(1, length = k, step = 1)
              dInx = clustDist[:, 2].==j
              dist_t = clustDist[:, 1].==j
              inertia += mean(dist_t)
              # println(dInx)
              ptInCluster = X[:, dInx]
              # ptWeight = w[dInx]
              # ptWeight = ptWeight ./ sum(ptWeight)
              #* ptWeight
              centers_new[:,j] = mean(ptInCluster, dims=2)
              # println(centers_new[:,j])
         end
         # println(centers_new)
         # println(centers_new - centers)
         delta = sum((centers_new - centers).^2)
         centers = copy(centers_new)
         @printf("iter:\t%d\tdelta:\t%0.4f\n", iter_num, delta)
         # println("delta: $delta")
         iter_num += 1
     end
     # println(centers)
     cent = deepcopy(centers)
     inertia -= sum(var(cent, dims=2))/k
     # println(centers)
     return centers, inertia
end

function init_centers(X::AbstractMatrix{<:Real},
     k::Integer)
     d, n = size(X)

     # init_index = rand(1:n, k)
     # init_centers = X[:,init_index]
     init_index = Inf64[]
     index = rand(1:n)
     init_cluster = X[:,index]
     init_centers = Array{Float64, 2}(undef, d, k)
     distListTotal = Array{Float64, 1}(undef, n)
     init_centers[:,1] = init_cluster

     for i in 2:k
         sum_all = 0
         for j in 1:n
             distList = Float64[]
             for j in 1:i
                 append!(distList, distance(init_centers[:,i], X[:, j]))
             end
             dist, ind = findmin(distList)
             distListTotal[j] = dist
             sum_all += dist
         end
         sum_all *= rand()
         for j in 1:n
             di = distListTotal[k]
             sum_all -= di
             if sum_all > 0
                 continue
             end
             init_centers[:,i] = X[:, j]
             break
         end
     end
     return init_centers
end

function distance(v1, v2)
    sqrt(sum((v1 - v2).^2))
end

function k_mean_predict(X::AbstractMatrix{<:Real},
    centers::AbstractMatrix{<:Real})
    d, n = size(X)
    d2, k = size(centers)
    (d == d2) || throw(ArgumentError(
    "dims of data $d and center $d2 should be equal"))
    clustInd = Int[]
    for i in range(1, length=n, step =1)
        distList = Float64[]
        for j in range(1, length=k, step=1)
            append!(distList, distance(centers[:,j], X[:, i]))
        end
        # println(distList)
        dist, ind = findmin(distList)
        # print(ind)
        append!(clustInd, ind)
    end
    return clustInd
end

function K_mean(X::AbstractMatrix{<:Real},
     k::Integer,
     w::Array{<:Real};
     maxiter::Integer=100, # maximun iteration
     tol::Real=1e-4, # tolearance of change at converage
     random_state::Integer=10
     )
     d, n = size(X)
     centers_ls = Array{Float64, 3}(undef, random_state, d, k)
     inertia_ls = []
    for iter in 1:random_state
        centers, inertia = k_mean(X, k, w)
        if sum(isnan.(centers)) > 0
            continue
        end
        # println(centers_ls)
        # println(typeof(centers_ls))
        println(size(centers))
        centers_ls[iter,:,:] = centers
        append!(inertia_ls, inertia)
    end
    inertia_min, min_ind = findmin(inertia_ls)
    return centers_ls[min_ind,:,:]
end

end  # module kmean