"""
Construct a basis for identity tangent space of differomorphisms, using the fourier
basis. The notation for the is in accordance with the project thesis, with "_"
 denoting the "flipping"  [bi(x, y), 0] -> [0, bi(y, x)]
"""

"""
TYPE 1 Basis Function: Sines
"""
function ξ(k)
    return x ->  [√2*sin(π*k*x[1]), 0.]
end

function ξx(k)
    return x -> √2π*k*cos(π*k*x[1])
end

function Dξ(k)
    return x -> [[√2π*k*cos(π*k*x[1]) 0.]; [0. 0.]]
end

"""
TYPE 2 Basis Function: Sines*Cosines
"""
function η(k, l)
    return x -> [2*sin(π*k*x[1])*cos(2π*l*x[2]), 0.]
end

function ηx(k, l)
    return x -> 2π*k*cos(π*k*x[1])*cos(2π*l*x[2])
end

function Dη(k, l)
    return x -> [[2π*k*cos(π*k*x[1])*cos(2π*l*x[2]) -4π*l*sin(π*k*x[1])*sin(2π*l*x[2])]; [0. 0.]]
end

"""
TYPE 3 Basis Function: Sines*Sines
"""
function φ(k, l)
    return x -> [2*sin(π*k*x[1])*sin(2π*l*x[2]), 0.]
end

function φx(k, l)
    return x -> 2π*k*cos(π*k*x[1])*sin(2π*l*x[2])
end

function Dφ(k, l)
    return x -> [[2π*k*cos(π*k*x[1])*sin(2π*l*x[2]) 4π*l*sin(π*k*x[1])*cos(2π*l*x[2])];[0. 0.]]
end

"""
***********************************
Now for the y-driection vector fields
*************************************
"""

"""
TYPE 1 Basis Function: Sines
"""
function _ξ(l)
    return x ->  [0., √2*sin(π*l*x[2])]
end

function _ξy(l)
    return x -> √2π*l*cos(π*l*x[2])
end

function _Dξ(l)
    return x -> [[0. 0.]; [0. √2π*l*cos(π*l*x[2])]]
end

"""
TYPE 2 Basis Function: Sines*Cosines
"""
function _η(l, k)
    return x -> [0., 2*sin(π*l*x[2])*cos(2π*k*x[1])]
end

function _ηy(l, k)
    return x -> 2π*l*cos(π*l*x[2])*cos(2π*k*x[1])
end

function _Dη(l, k)
    return x -> [[0. 0.]; [-4π*k*sin(π*l*x[2])*sin(2π*k*x[1]) 2π*l*cos(π*l*x[2])*cos(2π*k*x[1])]]
end

"""
TYPE 2 Basis Function: Sines*Sines
"""
function _φ(l, k)
    return x -> [0., 2*sin(π*l*x[2])*sin(2π*k*x[1])]
end

function _φy(l, k)
    return x -> 2π*l*cos(π*l*x[2])*sin(2π*k*x[1])
end

function _Dφ(l, k)
    return x -> [[0. 0.]; [4π*k*sin(π*l*x[2])*cos(2π*k*x[1]) 2π*l*cos(π*l*x[2])*sin(2π*k*x[1])]]
end

"""
Function for constructing bass functions with necessary derivatives
    * K::Int : Maximal Frequency of trigonometric functions

Returns:
    * B, DB, divB, H1, H2::Array{Function}
"""
function construct_basis(K::Int)
    N = K + 2K^2

    # Declare output function arrays
    B = Array{Function}(undef, 2N) # Basis functions
    DB = Array{Function}(undef, 2N) # Jacobians of basis functions
    divB = Array{Function}(undef, 2N) # Divergence of basis functions.

    # Loop through type 1 basis functions (only sines)
    for i in 1:K
        k = i

        # Basis Functions
        B[i] = ξ(k)
        B[N + i] = _ξ(k)

        # Jacobian Matrices
        DB[i] = Dξ(k)
        DB[N + i] = _Dξ(k)

        # Divergence
        divB[i] = ξx(k)
        divB[N + i] = _ξy(k)
    end

    for i in (K+1):(K+K^2)
        # Remap indices i -> j, k, l
        j = i - K
        k = ((j-1) ÷ K) + 1
        l = j - K * ((j-1) ÷ K)

        # Basis Functions
        B[i] = η(k, l)
        B[N + i] = _η(k, l)

        # Jacobian Matrices
        DB[i] = Dη(k, l)
        DB[N + i] = _Dη(k, l)

        # Divergence
        divB[i] = ηx(k, l)
        divB[N + i] = _ηy(k, l)
    end

    for i in (K + K^2 + 1):N
        # Remap indices i -> j, k, l
        j = i - (K + K^2)
        k = ((j-1) ÷ K) + 1
        l = j - K * ((j-1) ÷ K)

        # Basis Functions
        B[i] = φ(k, l)
        B[N + i] = _φ(k, l)

        # Jacobian Matrices
        DB[i] = Dφ(k, l)
        DB[N + i] = _Dφ(k, l)

        # Divergence
        divB[i] = φx(k, l)
        divB[N + i] = _φy(k, l)
    end

    return B, DB, divB
end


using LinearAlgebra
using KernelFunctions

"""
    correlation_kernel_Gram(B,R,unif_samples,all_samples,test_samples,k,sigma)

Compute Gram matrix of correlation kernel on test data.

# Arguments
- `B::Array{Float64,2}`: B matrix with same size as K.
- `R:Array{Float64,2}`: Cholesky s.t. K = R'R
- `unif_samples:Array{Float64,2}`: number of uniformly sampled points for correlation kernel approximaton
- `all_samples:Array{Float64,2}`: array containing all points used to write the representer
- `test_samples:Array{Float64,2}`: test samples for evaluating the Gram matrix
- `k:kernel`: RKHS kernel function; see KernelFunctions.
- `sigma:Float`: RKHS kernel bandwidth.

# Output
- `GramK::Array{Float64,2}`: Gram matrix of size |test_samples| x |test_samples|

"""
function correlation_kernel_Gram(B,R,unif_samples,all_samples,test_samples,k,sigma)

    F = cholesky(B).U; # B = F'F

    # construct cross kernel matrices
    x_m = (all_samples)'/(sigma);
    x_p = (unif_samples)'/(sigma);
    K_mp = kernelmatrix(k, x_m, x_p);

    x_t = (test_samples)'/(sigma);
    K_mt = kernelmatrix(k, x_m, x_t);

    T_1 = F*((R')\K_mp);
    p = size(K_mp,2);
    M = (1/p)*T_1*T_1' + I;

    T_2 = F*((R')\K_mt);

    GramK = T_2'*(M\T_2);

    GramK = 0.5*(GramK+GramK');

    return GramK;
end

"""
    likelihood_kernel_Gram(B,R,all_samples,test_samples,k,sigma)

Compute Gram matrix of likelihood kernel on test data.

# Arguments
- `B::Array`: B matrix with same size as K.
- `R:Array`: Cholesky s.t. K = R'R
- `all_samples:Array`: array containing all points used to write the representer
- `test_samples:Array`: test samples for evaluating the Gram matrix
- `k:kernel`: RKHS kernel function; see KernelFunctions.
- `sigma:Float`: RKHS kernel bandwidth.

# Output
- `GramA::Array{Float64,2}`: Gram matrix of size |test_samples| x |test_samples|

"""
function likelihood_kernel_Gram(B,R,all_sples,test_sples,k,sigma)

    F = cholesky(B).U; # B = F'F

    # construct cross kernel matrices
    x_m = (all_sples)'/(sigma);
    x_t = (test_sples)'/(sigma);
    K_mt = kernelmatrix(k, x_m, x_t);

    T = F*((R')\K_mt);
    GramA = T'*T;

    return GramA;
end


degtorad(deg) = deg*π/180

# adapted from Distances.jl
@inline function haversine(λ₁, φ₁, λ₂, φ₂)
    Δλ = λ₂ - λ₁  # longitudes
    Δφ = φ₂ - φ₁  # latitudes

    # haversine formula
    a = sin(Δφ/2)^2 + cos(φ₁)*cos(φ₂)*sin(Δλ/2)^2

    # # distance on the sphere
    2 * 6_371_000 * asin( min(√a, one(a)) ) # take care of floating point errors
end

function dist_pairwise(long, lat)
    dim = size(long, 1)
    dist = Array{Float64}(undef, dim, dim)
    @fastmath @inbounds for i ∈ axes(dist, 1), j ∈ 1:i
        dist[i, j] = haversine(long[i], lat[i], long[j], lat[j])
    end
    return Symmetric(dist, :L)
end

function dist_pointwise!(out, long_point, lat_point, long_vec, lat_vec)
    @fastmath @inbounds for i ∈ eachindex(out)
        out[i] = haversine(long_point, lat_point, long_vec[i], lat_vec[i])
    end
    return out
end

function fill_diag!(matrix, value)
    foreach(i -> matrix[i, i] = value, 1:size(matrix, 1))
end


#=
Euler discovered the remarkable quadratic formula:

n² + n + 41

It turns out that the formula will produce 40 primes for the consecutive values n = 0 to 39. However, when n = 40, 402 + 40 + 41 = 40(40 + 1) + 41 is divisible by 41, and certainly when n = 41, 41² + 41 + 41 is clearly divisible by 41.

The incredible formula  n² − 79n + 1601 was discovered, which produces 80 primes for the consecutive values n = 0 to 79. The product of the coefficients, −79 and 1601, is −126479.

Considering quadratics of the form:

n² + an + b, where |a| < 1000 and |b| < 1000

where |n| is the modulus/absolute value of n
e.g. |11| = 11 and |−4| = 4
Find the product of the coefficients, a and b, for the quadratic expression that produces the maximum number of primes for consecutive values of n, starting with n = 0.
=#

using Primes

function longestprimerun(a,b)
  i = 0
  while isprime((i^2)+(a*i)+b)
    i+= 1
  end
  i
end

function calc()
  (mina, maxa) = (-1000, 1000)
  (largesta, largestb) = (0,0)
  largest = 0
  for a in mina:maxa, b in mina:maxa
    primerun = longestprimerun(a,b)
    if primerun > largest
      (largesta, largestb) = (a,b)
      largest = primerun
    end
  end
  largesta * largestb
end
@time println(calc())


module CubicEquation

export Solver

type Solver
end

function (solver::Solver)(a::Real, b::Real, c::Real, d::Real)
    if a == zero(typeof(a))
        return Base.error("Can't solve. a=0 in a*x³ + b*x² + c*x + d =0")
    end
    f = (3c/a - (b/a)^2) / 3
    g = (2(b/a)^3 - (9*b*c/a^2) + 27d/a) / 27
    h = g^2 / 4 + f^3 / 27
    if f == zero(typeof(f)) && g == zero(typeof(g)) && h == zero(typeof(h))
        x1 = x2 = x3 = -cbrt(d/a)
        return [x1, x2, x3]
    elseif  h <= zero(typeof(h))
        return three_real_root(a, b, g, h)
    elseif h > zero(typeof(h))
        return one_real_root(a, b, g, h)
    end

    return Base.error("Can't solve when a=$(a) b=$(b) c=$(c) d=$(d)")
end

function three_real_root(a::Real, b::Real, g::Float64, h::Float64)
    i = (g^2 / 4 - h)^(1/2)
    j = cbrt(i)
    K = acos(-(g/2i))
    L = -j
    M = cos(K/3)
    N = sqrt(3) * sin(K/3)
    P = -(b/3a)
    x1 = 2j * cos(K/3) - (b/3a)
    x2 = L * (M + N) + P
    x3 = L * (M - N) + P
    return sort([x1, x2, x3])
end

function one_real_root(a::Real, b::Real, g::Float64, h::Float64)
    R = -g/2 + sqrt(h)
    S = cbrt(R)
    T = -g/2 - sqrt(h)
    U = cbrt(T)
    x1 = (S + U) - (b/3a)
    x2 = -(S + U)/2 - (b/3a) - im*(S-U)*sqrt(3) / 2
    x3 = -(S + U)/2 - (b/3a) + im*(S-U)*sqrt(3) / 2
    return Union{Real,Complex}[x1, x2, x3]
end

function (solver::Solver)(a::Real, b::Real, c::Real)
    if a == zero(typeof(a))
        return Base.error("Can't solve. a=0 in a*x² + b*x + c = 0")
    end

    # Discriminant for quadratic equation
    D = b^2 - 4*a*c

    if D >= zero(typeof(D))
        return [(-b - √(D)) / 2a, (-b + √(D)) / 2a]
    elseif D < zero(typeof(D))
        return [(-b - √(-D)*im) / 2a, (-b + √(-D)*im) / 2a]
    else
        return Base.error("Can't solve when a=$(a) b=$(b) c=$(c)")
    end
end

end #end module


# If an integer is passed down to the function, it means that the user is requesting an exact number of iterations.
function bisection_it(f,ansatz::Tuple{Float64,Float64}; N::Int64=100)
    (lower,upper) = ansatz
    mid = 0
    for _ in 1:N

    mid = (lower + upper)/2

        if f(lower)*f(mid) < 0
            upper = mid
        elseif f(lower)*f(mid) > 0
            lower = mid
        else
            return mid
        end

    end
    return mid
end

# f(x) = sin(x)
# bisection(f,(2.,4.))

#If then a Float is passed to the function, then it means it must mean an objective relative error.
function bisection_err(f,ansatz::Tuple{Float64,Float64}; error::Float64=0.01)
    mid = (lower + upper)/2;
    ea = 1; #Initialize
    #We do the first iteration, otherwise, there is no way to have a relative error to measure
    if f(lower)*f(mid) < 0
        upper = mid
    elseif f(lower)*f(mid) > 0
        lower = mid
    else
        return mid
    end
    #Now we start doing the iterations until it is done
    while ea > error
        last = mid #We store the last approximation
        mid = (lower + upper)/2
        ea = abs((mid-last)/mid);

        if f(lower)*f(mid) < 0
            upper = mid
        elseif f(lower)*f(mid) > 0
            lower = mid
        else
            return mid
        end

    end
    return mid
end


# ---
# title: 688. Knight Probability in Chessboard
# id: problem688
# author: Indigo
# date: 2021-06-29
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/knight-probability-in-chessboard/description/>
# hidden: true
# ---
#
# On an `N`x`N` chessboard, a knight starts at the `r`-th row and `c`-th column
# and attempts to make exactly `K` moves. The rows and columns are 0 indexed, so
# the top-left square is `(0, 0)`, and the bottom-right square is `(N-1, N-1)`.
#
# A chess knight has 8 possible moves it can make, as illustrated below. Each
# move is two squares in a cardinal direction, then one square in an orthogonal
# direction.
#
#
#
# ![](https://assets.leetcode.com/uploads/2018/10/12/knight.png)
#
#
#
# Each time the knight is to move, it chooses one of eight possible moves
# uniformly at random (even if the piece would go off the chessboard) and moves
# there.
#
# The knight continues moving until it has made exactly `K` moves or has moved
# off the chessboard. Return the probability that the knight remains on the
# board after it has stopped moving.
#
#
#
# **Example:**
#
#
#
#     Input: 3, 2, 0, 0
#     Output: 0.0625
#     Explanation: There are two moves (to (1,2), (2,1)) that will keep the knight on the board.
#     From each of those positions, there are also two moves that will keep the knight on the board.
#     The total probability the knight stays on the board is 0.0625.
#
#
#
#
# **Note:**
#
#   * `N` will be between 1 and 25.
#   * `K` will be between 0 and 100.
#   * The knight always initially starts on the board.
#
#
## @lc code=start
using LeetCode

function knight_probability(n::Int, k::Int, row::Int, col::Int)
    dp = fill(0, n, n, 2)
    dp[row + 1, col + 1, 1] = 1
    for i in 1:k
        dp1 = @view(dp[:, :, mod1(i, 2)])
        dp2 = @view(dp[:, :, mod1(i + 1, 2)])
        idcs = CartesianIndices(dp2)
        for I in idcs, hop in ((1, 2), (2, 1), (-1, 2), (2, -1), (1, -2), (-2, 1), (-1, -2), (-2, -1))
            new_I = I + CartesianIndex(hop)
            new_I ∈ idcs && (dp2[new_I] += dp1[I])
        end
        fill!(dp1, 0)
    end
    return sum(dp) / 8 ^ k
end
## @lc code=end


using LinearAlgebra, SimpleDifferentialOperators, DiffEqOperators, Test, Parameters

#parameters
params = @with_kw (
    μ = -0.1, # constant negative drift
    σ = 0.1,
    ρ = 0.05,
    M = 3, # size of grid (interior points)
    x̄ = range(0.0, 1.0, length = (M+2)),
    x = interiornodes(x̄), # i.e., x̄[2:end-1]
    S = 3.0
)
p = params();

#----------------------------------
#Testing For Negative Drifts
#----------------------------------
# payoff function
π(x) = x^2
# SimpleDifferentialOperators setup
function SDO_negative_drift(π, params)
    bc = (Reflecting(), Reflecting())
    Lₓ = params.μ*L₁₋bc(params.x̄, bc) + params.σ^2 / 2 * L₂bc(params.x̄, bc)
    L_bc = I * params.ρ - Lₓ
    v = L_bc \ π.(params.x);
    return (v = v, L_bc = L_bc, Lₓ = Lₓ, L₁₋bc = L₁₋bc(params.x̄, bc), L₂bc = L₂bc(params.x̄, bc))
end
# DiffEqOperators Setup
function DEO_negative_drift(π, params)
    dx = params.x[2] - params.x[1]
    # discretize L = ρ - μ D_x - σ^2 / 2 D_xx
    # subject to reflecting barriers at 0 and 1
    L1 = UpwindDifference(1,1,dx,params.M,1.)
    L2 = CenteredDifference(2,2,dx,params.M)
    Q = Neumann0BC(dx, 1)
    L₁₋bc = -1. .* Array(UpwindDifference(1,1,dx,params.M,-1.) * Q)[1]
    # Here Array(A::GhostDerivativeOperator) will return a tuple of the linear part
    # and the affine part of the operator A, hence we index Array(µ*L1*Q).
    # The operators in this example are purely linear, so we don't worry about Array(µ*L1*Q)[2]
    Lₓ = Array(params.μ*L1*Q)[1] + Array(params.σ^2/2*L2*Q)[1] # when scalar coefficient is properly implemented
    # Lₓ = Array(L1*Q)[1] + Array(params.σ^2/2*L2*Q)[1] # when scalar coefficient is properly implemented
    L_bc = I * params.ρ - Lₓ

    # solve the value function
    v = L_bc \ π.(params.x);
    return (v = v, L_bc = L_bc, Lₓ = Lₓ, L₁₋bc = L₁₋bc, L₂bc = Array(L2*Q)[1] )
end
@testset "Constant Negative Drifts" begin
    @test DEO_negative_drift(π, p).v ≈ SDO_negative_drift(π, p).v
    @test DEO_negative_drift(π, p).L_bc ≈ SDO_negative_drift(π, p).L_bc
    @test DEO_negative_drift(π, p).Lₓ ≈ SDO_negative_drift(π, p).Lₓ
    @test DEO_negative_drift(π, p).L₁₋bc ≈ SDO_negative_drift(π, p).L₁₋bc
    @test DEO_negative_drift(π, p).L₂bc ≈ SDO_negative_drift(π, p).L₂bc
end
p = params(μ = 0.1);
#----------------------------------
#Testing For Positive Drifts
#----------------------------------
# SimpleDifferentialOperators setup
function SDO_positive_drift(π, params)
    bc = (Reflecting(), Reflecting())
    Lₓ = params.μ*L₁₊bc(params.x̄, bc) + params.σ^2 / 2 * L₂bc(params.x̄, bc)
    L_bc = I * params.ρ - Lₓ
    v = L_bc \ π.(params.x);
    return (v = v, L_bc = L_bc, Lₓ = Lₓ, L₁₊bc = L₁₊bc(params.x̄, bc), L₂bc = L₂bc(params.x̄, bc))
end
# DiffEqOperators Setup
function DEO_positive_drift(π, params)
    dx = params.x[2] - params.x[1]
    # discretize L = ρ - μ D_x - σ^2 / 2 D_xx
    # subject to reflecting barriers at 0 and 1
    L1 = UpwindDifference(1,1,dx,params.M,1.0)
    L2 = CenteredDifference(2,2,dx,params.M)
    Q = Neumann0BC(dx, 1)
    # Here Array(A::GhostDerivativeOperator) will return a tuple of the linear part
    # and the affine part of the operator A, hence we index Array(µ*L1*Q).
    # The operators in this example are purely linear, so we don't worry about Array(µ*L1*Q)[2]
    Lₓ = Array(params.µ*L1*Q)[1] + Array(params.σ^2/2*L2*Q)[1]
    L_bc = I * params.ρ - Lₓ

    # solve the value function
    v = L_bc \ π.(params.x);
    return (v = v, L_bc = L_bc, Lₓ = Lₓ, L₁₊bc = Array(L1*Q)[1], L₂bc = Array(L2*Q)[1] )
end

@testset "Constant Positive Drifts" begin
    @test DEO_positive_drift(π, p).v ≈ SDO_positive_drift(π, p).v
    @test DEO_positive_drift(π, p).L_bc ≈ SDO_positive_drift(π, p).L_bc
    @test DEO_positive_drift(π, p).Lₓ ≈ SDO_positive_drift(π, p).Lₓ
    @test DEO_positive_drift(π, p).L₁₊bc ≈ SDO_positive_drift(π, p).L₁₊bc
    @test DEO_positive_drift(π, p).L₂bc ≈ SDO_positive_drift(π, p).L₂bc
end

#----------------------------------
#Testing For State Dependent Drifts
#----------------------------------
μ(x) = -x
# SimpleDifferentialOperators setup
function SDO_state_dependent_drift(π, μ, params)
    bc = (Reflecting(), Reflecting())
    L₁ = Diagonal(min.(μ.(params.x), 0.0)) * L₁₋bc(params.x̄, bc) + Diagonal(max.(μ.(params.x), 0.0)) * L₁₊bc(params.x̄, bc)
    Lₓ = L₁ - params.σ^2 / 2 * L₂bc(params.x̄, bc)
    L_bc = I * params.ρ - Lₓ
    v = L_bc \ π.(params.x)
    return (v = v, L_bc = L_bc,  Lₓ = Lₓ, L₁ = L₁, L₂bc = L₂bc(params.x̄, bc))
end
# DiffEqOperators Setup
function DEO_state_dependent_drift(π, μ, params)
    dx = params.x[2] - params.x[1]
    # discretize L = ρ - μ D_x - σ^2 / 2 D_xx
    # subject to reflecting barriers at 0 and 1
    drift = μ.(params.x)
    L1 = UpwindDifference(1,1,dx,params.M,drift)
    L2 = CenteredDifference(2,2,dx,params.M)
    Q = Neumann0BC(dx, 1)
    # Here Array(A::GhostDerivativeOperator) will return a tuple of the linear part
    # and the affine part of the operator A, hence we index Array(µ*L1*Q).
    # The operators in this example are purely linear, so we don't worry about Array(µ*L1*Q)[2]
    L₁ = Array(L1*Q)[1]
    Lₓ = L₁ - Array(params.σ^2/2*L2*Q)[1]
    L_bc = I * params.ρ - Lₓ

    # solve the value function
    v = L_bc \ π.(params.x);
    return (v = v, L_bc = L_bc, Lₓ = Lₓ, L₁ = L₁, L₂bc = Array(L2*Q)[1] )
end

@testset "State Dependent Drifts" begin
    @test DEO_state_dependent_drift(π, μ, p).v ≈ SDO_state_dependent_drift(π, μ, p).v
    @test DEO_state_dependent_drift(π, μ, p).L_bc ≈ SDO_state_dependent_drift(π, μ, p).L_bc
    @test DEO_state_dependent_drift(π, μ, p).Lₓ ≈ SDO_state_dependent_drift(π, μ, p).Lₓ
    @test DEO_state_dependent_drift(π, μ, p).L₁ ≈ SDO_state_dependent_drift(π, μ, p).L₁
    @test DEO_state_dependent_drift(π, μ, p).L₂bc ≈ SDO_state_dependent_drift(π, μ, p).L₂bc
end

#----------------------------------
#Testing For Absorbing BC
#----------------------------------

# payoff function
π(x) = x^2
# SimpleDifferentialOperators setup
function SDO_absorbing_bc(π, params)
    bc = (NonhomogeneousAbsorbing(params.S), Reflecting())
    Lₓbc = params.μ*L₁₋bc(params.x̄, bc) + params.σ^2 / 2 * L₂bc(params.x̄ , bc)
    L_bc = I * params.ρ - Lₓbc

    # construct the RHS with affine boundary
    π_star = π.(params.x) + L₀affine(params.x̄, π.(params.x), bc)

    # solve the value function
    v = L_bc \ π_star

    return (v = v, L_bc = L_bc, Lₓbc = Lₓbc, L₁₋bc = L₁₋bc(params.x̄, bc), L₂bc = L₂bc(params.x̄, bc),
        π_star = π_star, π = π.(params.x))
end

# DiffEqOperators Setup
function DEO_absorbing_bc(π, params)
    dx = params.x[2] - params.x[1]

    L1 = UpwindDifference(1,1,dx,params.M,params.μ)
    L2 = CenteredDifference(2,2,dx,params.M)
    # RobinBC(l::NTuple{3,T}, r::NTuple{3,T}, dx::T, order = 1)
    # The variables in l are [αl, βl, γl], and correspond to a BC of the form αl*u(0) + βl*u'(0) = γl
    # imposed on the lower index boundary. The variables in r are [αr, βr, γr],
    # and correspond to an analagous boundary on the higher index end.
    l = (1.0, 0.0, p.S)
    r = (0.0, 1.0, 0.0)
    Q = RobinBC(l, r, dx, 1)

    Lₓbc = Array(L1*Q)[1] + Array(params.σ^2/2*L2*Q)[1]
    L_bc = I * params.ρ - Lₓbc

    # solve the value function
    v = L_bc \ π.(params.x)
    return (v = v, L_bc = L_bc, Lₓbc = Lₓbc, L₁₋bc = Array(L1*Q)[1] ./ params.μ, L₂bc = Array(L2*Q)[1],
        π_star = π.(params.x), π = π.(params.x))
end

p = params(x̄ = range(0.0, 1.0, length = (p.M+2)))

@testset "Absorbing BC" begin
    @test DEO_absorbing_bc(π, p).v ≈ SDO_absorbing_bc(π, p).v
    @test DEO_absorbing_bc(π, p).L_bc ≈ SDO_absorbing_bc(π, p).L_bc
    @test DEO_absorbing_bc(π, p).Lₓbc ≈ SDO_absorbing_bc(π, p).Lₓbc
    @test DEO_absorbing_bc(π, p).L₁₋bc ≈ SDO_absorbing_bc(π, p).L₁₋bc
    @test DEO_absorbing_bc(π, p).L₂bc ≈ SDO_absorbing_bc(π, p).L₂bc
    @test DEO_absorbing_bc(π, p).π_star ≈ SDO_absorbing_bc(π, p).π_star
    @test DEO_absorbing_bc(π, p).π ≈ SDO_absorbing_bc(π, p).π
end

#----------------------------------
#Testing For Solving KFE
#----------------------------------

function SDO_Solve_KFE(params)
    # ξ values for mixed boundary conditions
    ξ_lb = ξ_ub = -2 * params.μ / params.σ^2
    # define the corresponding mixed boundary conditions
    # note that the direction on the lower bound is backward (default is forward)
    # as the drift μ is negative.
    bc = (Mixed(ξ = ξ_lb, direction = :backward), Mixed(ξ = ξ_ub))

    # use SimpleDifferentialOperators.jl to construct the operator on the interior
    L_KFE_with_drift = Array(-params.μ*L₁₊bc(params.x̄, bc) + params.σ^2 / 2 * L₂bc(params.x̄, bc))
    L_KFE_without = params.σ^2 / 2 * L₂bc(params.x̄, bc)
    return (L_KFE_with_drift = L_KFE_with_drift, L_KFE_without = L_KFE_without)
end

function SDO_Solve_KFE_forward(params)
    # ξ values for mixed boundary conditions
    ξ_lb = ξ_ub = -2 * params.μ / params.σ^2
    # define the corresponding mixed boundary conditions
    # but with forward finite difference derivatives at boundaries
    bc = (Mixed(ξ = ξ_lb, direction = :forward), Mixed(ξ = ξ_ub))

    # use SimpleDifferentialOperators.jl to construct the operator on the interior
    L_KFE_with_drift = Array(-params.μ*L₁₊bc(params.x̄, bc) + params.σ^2 / 2 * L₂bc(params.x̄, bc))
    L_KFE_without = params.σ^2 / 2 * L₂bc(params.x̄, bc)
    return (L_KFE_with_drift = L_KFE_with_drift, L_KFE_without = L_KFE_without)
end

function DEO_Solve_KFE(params)
    dx = params.x[2] - params.x[1]

    L1 = UpwindDifference(1,1,dx,params.M,-params.μ)
    # L2l = UpwindDifference(2,2,dx,params.M,
    #                          vcat(-1.,zeros(params.M-1)))
    # L2r = UpwindDifference(2,2,dx,params.M,
    #                          vcat(zeros(params.M-1),1.))
    L2 = CenteredDifference(2,2,dx,params.M)

    ξ_lb = ξ_ub = -2 * params.μ / params.σ^2

    l = (ξ_lb, 1., 0.)
    r = (ξ_ub, 1., 0.)
    Q = RobinBC(l, r, dx, 1)

    # use SimpleDifferentialOperators.jl to construct the operator on the interior
    # Only difference is handling of L2, we may need to upwind for boundary
    L_KFE_with_drift = Array(L1*Q)[1] + Array(params.σ^2/2 *L2*Q)[1]
    L_KFE_without = (params.σ^2 / 2) .* Array(L2 * Q)[1]
    # L_KFE_without = (params.σ^2 / 2) .* ([0., 1., 0.] .* Array(L2 * Q)[1] +
    #     -(Array(L2l * Q)) + Array(L2r * Q))

    return (L_KFE_with_drift = L_KFE_with_drift, L_KFE_without = L_KFE_without)
end

p = params(x̄ = range(0.0, 1.0, length = (p.M+2)))
@testset "Solving KFE" begin
    @test_broken DEO_Solve_KFE(p).L_KFE_with_drift ≈ SDO_Solve_KFE(p).L_KFE_with_drift # concretization of Robin conditions appears broken
    @test_broken DEO_Solve_KFE(p).L_KFE_without ≈ SDO_Solve_KFE(p).L_KFE_without # Hard to check what's wrong since KFE uses forward/backward at boundaries rather than central differences when w/mixed bcs
    @test DEO_Solve_KFE(p).L_KFE_with_drift ≈ SDO_Solve_KFE_forward(p).L_KFE_with_drift
    @test DEO_Solve_KFE(p).L_KFE_without ≈ SDO_Solve_KFE_forward(p).L_KFE_without
end

nothing


# an example on mixed integer programming
#
#   maximize x + 2 y + 5 z
#
#   s.t.  x + y + z <= 10
#         x + 2 y + z <= 15
#
#         x is continuous: 0 <= x <= 5
#         y is integer: 0 <= y <= 10
#         z is binary
#
# z = 1, y= 7, x = 0
using Xpress, Compat.Test
using Compat.SparseArrays

@testset "MIP 4" begin

    model = Xpress.Model( "mip_01", :maximize)

    add_cvar!(model, 1., 0., 5.)  # x
    add_ivar!(model, 2., 0, 10)   # y
    add_bvar!(model, 5.)          # z

    add_constr!(model, ones(3), '<', 10.)
    add_constr!(model, [1., 2., 1.], '<', 15.)

    @test Xpress.get_obj(model) == [1, 2, 5]
    @test Xpress.get_constrmatrix(model) == sparse([1. 1. 1.; 1. 2. 1.])
    @test Xpress.get_rhs(model) == [10, 15]
    @test Xpress.get_lb(model) == [0, 0, 0]
    @test Xpress.get_ub(model) == [5, 10, 1]
    @test Xpress.get_sense(model) == [:<=,:<=]

    optimize(model)

    ans = get_optiminfo(model)
    @test ans.status_mip == :mip_optimal

    @test isapprox(get_solution(model), [0, 7, 1]; atol = 1e-3)

    @test isapprox(get_objval(model), 19; atol = 1e-3)
end


import Voronoi

using Random
using Printf
include("Draw.jl")

function randf(start, finish, n)
   v = rand(n)
   return map(x -> start + x*(finish-start), v)
end
Random.seed!(100)

WIDTH = 100.0
HEIGHT = 100.0

function benchmarkLloyd(Z::Array{Tuple{Real, Real}, 1}, ϵ::Number)
   global WIDTH, HEIGHT

   points = deepcopy(Z)

   results = Array{Tuple{Real, Real, Real, Real}}([])

   times = @elapsed V, f, ∇f, ξ = Voronoi.Optimization.init(points, WIDTH, HEIGHT)
   evaluations = 1
   push!(results, (f, ξ, evaluations, times))

   while ξ > ϵ
      time = @elapsed V, f, ξ = Voronoi.Optimization.lloydIteration(points, V)

      evaluations += 1
      times += time

      push!(results, (f, ξ, evaluations, times))
   end

   return results
end

function benchmarkGradient(Z::Array{Tuple{Real, Real}, 1}, ϵ::Number)
   global WIDTH, HEIGHT

   points = deepcopy(Z)

   results = Array{Tuple{Real, Real, Real, Real}}([])

   times = @elapsed V, f, ∇f, ξ = Voronoi.Optimization.init(points, WIDTH, HEIGHT)
   evaluations = 1
   push!(results, (f, ξ, evaluations, times))
   μ = 0.5
   α = 0.01
   λ₀ = 0.01

   while ξ > ϵ
      d = ∇f

      time = @elapsed V, f, ∇f, ξ, evals = Voronoi.Optimization.lineSearch(points, d, α, λ₀, μ, ϵ, V, f, ∇f, ξ)

      evaluations += evals
      times += time

      push!(results, (f, ξ, evaluations, times))
   end

   return results
end

function benchmark(n, ϵ)
   global WIDTH, HEIGHT

   points = convert(Array{Tuple{Real, Real}}, collect(zip(randf(1, WIDTH-1, n), randf(1, HEIGHT-1, n))))

   lloyd_results = benchmarkLloyd(points, ϵ)
   gradient_results = benchmarkGradient(points, ϵ)

   lloyd_f = Array{Number}([])
   lloyd_ξ = Array{Number}([])
   lloyd_evals = Array{Number}([])
   lloyd_time = Array{Number}([])
   gradient_f = Array{Number}([])
   gradient_ξ = Array{Number}([])
   gradient_evals = Array{Number}([])
   gradient_time = Array{Number}([])

   i = 1
   t₀ = lloyd_results[1][4]
   for (fᵢ, ξᵢ, evalsᵢ, timeᵢ) in lloyd_results
      push!(lloyd_f, fᵢ)
      push!(lloyd_ξ, ξᵢ)
      push!(lloyd_evals, evalsᵢ)
      push!(lloyd_time, timeᵢ - t₀)
      i += 1
   end

   i = 1
   t₀ = gradient_results[1][4]
   for (fᵢ, ξᵢ, evalsᵢ, timeᵢ) in gradient_results
      push!(gradient_f, fᵢ)
      push!(gradient_ξ, ξᵢ)
      push!(gradient_evals, evalsᵢ)
      push!(gradient_time, timeᵢ - t₀)
      i += 1
   end

   #@printf("f = %.10g\t\t|∇f| = %.10g\t\ttime: %s seconds\n", f, ξ, time)
   #@printf("f = %.10g\t\t|∇f| = %.10g\t\ttime: %s seconds\n", f, ξ, time)

   Draw.init(0, 0)
   Draw.clear(string("n = ", n))
   Draw.plot(lloyd_time, lloyd_f, "xkcd:azure", "Algoritmo de Lloyd", "-")
   Draw.plot(gradient_time, gradient_f, "xkcd:tomato", "Método do Gradiente", ":")
   Draw.legend("Tempo (s)", "Valor da Função Objetivo")
   Draw.commit()
   println("Press Return to finish.")
   command = readline(stdin)

   Draw.init(0, 0)
   Draw.clear(string("n = ", n))
   Draw.plot(lloyd_time, lloyd_ξ, "xkcd:azure", "Algoritmo de Lloyd", "-")
   Draw.plot(gradient_time, gradient_ξ, "xkcd:tomato", "Método do Gradiente", ":")
   Draw.legend("Tempo (s)", "Norma do Gradiente")
   Draw.commit()
   println("Press Return to finish.")
   command = readline(stdin)
end

ϵ = 10^-2
benchmark(3, ϵ)
benchmark(30, ϵ)
benchmark(100, ϵ)
benchmark(300, ϵ)
benchmark(1000, ϵ)


# Numbers and symbols can't be simplified further
simplify(n::Number) = n
simplify(s::Symbol) = s

# Handles all lengths for ex.args
# Removes any 0's in a sum
function simplify_sum(ex::Expr)
    new_args = map(x -> simplify(x), filter(x -> x != 0, ex.args[2:end]))
    if length(new_args) == 0
        return 0
    # Special Case: simplify(:(+x)) == x
    elseif length(new_args) == 1
        return new_args[1]
    else
        unshift(new_args, :+)
        return Expr(:call, new_args, Any)
    end
end

# Assumes length(ex.args) == 3
# Removes any 0's in a subtraction
function simplify_subtraction(ex::Expr)
    new_args = map(x -> simplify(x), filter(x -> x != 0, ex.args[2:end]))
    if length(new_args) == 0
        return 0
    # Special Case: simplify(:(x - x)) == 0
    elseif length(new_args) == 2 && new_args[1] == new_args[2]
        return 0
    else
        unshift(new_args, :-)
        return Expr(:call, new_args, Any)
    end
end

# Handles all lengths for ex.args
# Removes any 1's in a product
function simplify_product(ex::Expr)
    new_args = map(x -> simplify(x), filter(x -> x != 1, ex.args[2:end]))
    if length(new_args) == 0
        return 1
    # Special Case: simplify(:(*x)) == x
    elseif length(new_args) == 1
        return new_args[1]
    # Special Case: simplify(:(x * y * z * 0)) == 0
    elseif any(new_args .== 0)
        return 0
    else
        unshift(new_args, :*)
        return Expr(:call, new_args, Any)
    end
end

# Assumes length(ex.args) == 3
function simplify_quotient(ex::Expr)
    new_args = map(x -> simplify(x), ex.args[2:end])
    # Special Case: simplify(:(x / 1)) == x
    if new_args[2] == 1
        return new_args[1]
    # Special Case: simplify(:(0 / x)) == 0
    elseif new_args[1] == 0
        return 0
    else
        unshift(new_args, :/)
        return Expr(:call, new_args, Any)
    end
end

# Assumes length(ex.args) == 3
function simplify_power(ex::Expr)
    new_args = map(x -> simplify(x), ex.args[2:end])
    # Special Case: simplify(:(x ^ 0)) == 1
    if new_args[2] == 0
        return 1
    # Special Case: simplify(:(x ^ 1)) == x
    elseif new_args[2] == 1
        return new_args[1]
    # Special Case: simplify(:(0 ^ x)) == 0
    elseif new_args[1] == 0
        return 0
    # Special Case: simplify(:(1 ^ x)) == 1
    elseif new_args[1] == 1
        return 1
    else
        unshift(new_args, :^)
        return Expr(:call, new_args, Any)
    end
end

# Lookup table of simplification rules
simplify_lookup = {
                    :+ => simplify_sum,
                    :- => simplify_subtraction,
                    :* => simplify_product,
                    :/ => simplify_quotient,
                    :^ => simplify_power
                  }

# The simplest form of an expression is a fixed point of simplify()
function simplify(ex::Expr)
    if ex.head == :call
        if all(map(a -> isa(a, Number), ex.args[2:end]))
            return eval(ex)
        end
        if has(simplify_lookup, ex.args[1])
            new_ex = simplify_lookup[ex.args[1]](ex)
            while new_ex != ex
                new_ex, ex = simplify(new_ex), new_ex
            end
            return new_ex
        else
          return ex
        end
    else
        return ex
    end
end

# The Constant Rule
# d/dx c = 0
differentiate(x::Number, target::Symbol) = 0

# The Symbol Rule
# d/dx x = 1
# d/dx y = 0
function differentiate(s::Symbol, target::Symbol)
    if s == target
        return 1
    else
        return 0
    end
end

# The Sum Rule for Unary and Binary +
# d/dx +(f) = +(d/dx f)
# d/dx (f + g) = d/dx f + d/dx g
function differentiate_sum(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :+
        error("Not a valid sum call: $(ex)")
    end
    n = length(ex.args)
    new_args = Array(Any, n)
    new_args[1] = :+
    for i in 2:n
        new_args[i] = differentiate(ex.args[i], target)
    end
    return Expr(:call, new_args, Any)
end

# The Subtraction Rule for Unary and Binary -
# d/dx -(f) = -(d/dx f)
# d/dx (f - g) = d/dx f - d/dx g
function differentiate_subtraction(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :-
        error("Not a valid subtraction call: $(ex)")
    end
    n = length(ex.args)
    new_args = Array(Any, n)
    new_args[1] = :-
    for i in 2:n
        new_args[i] = differentiate(ex.args[i], target)
    end
    return Expr(:call, new_args, Any)
end

# The Product Rule
# d/dx (f * g) = (d/dx f) * g + f * (d/dx g)
# d/dx (f * g * h) = (d/dx f) * g * h + f * (d/dx g) * h + ...
function differentiate_product(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :*
        error("Not a valid product call: $(ex)")
    end
    n = length(ex.args)
    res_args = Array(Any, n)
    res_args[1] = :+
    for i in 2:n
       new_args = Array(Any, n)
       new_args[1] = :*
       for j in 2:n
           if j == i
               new_args[j] = differentiate(ex.args[j], target)
           else
               new_args[j] = ex.args[j]
           end
       end
       res_args[i] = Expr(:call, new_args, Any)
    end
    return Expr(:call, res_args, Any)
end

# The Quotient Rule
# d/dx (f / g) = ((d/dx f) * g - f * (d/dx g)) / g^2
function differentiate_quotient(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :/
        error("Not a valid quotient call: $(ex)")
    end
    return Expr(:call,
                {
                    :/,
                    Expr(:call,
                         {
                            :-,
                            Expr(:call,
                                 {
                                    :*,
                                    differentiate(ex.args[2], target),
                                    ex.args[3]
                                 },
                                 Any),
                            Expr(:call,
                                 {
                                    :*,
                                    ex.args[2],
                                    differentiate(ex.args[3], target)
                                 },
                                 Any)
                         },
                         Any),
                    Expr(:call,
                         {
                            :^,
                            ex.args[3],
                            2
                         },
                         Any)
                },
                Any)
end

# The Power Rule:
# Case 1: x^n <=> /dx x^n = n * x^(n - 1)
# Case 2: x^x <=> d/dx x^x = x^x (log(x) + 1)
# Case 3: n^n <=> d/dx n^n = 0
# Case 4: n^x <=> n^x * log(n)
# TODO: Handle general case for things like sin(x)^2
function differentiate_power(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :^
        error("Not a valid power call: $(ex)")
    end
    if ex.args[2] == target && ex.args[3] != target
        return Expr(:call,
                    {
                        :*,
                        ex.args[3],
                        Expr(:call,
                             {
                                :^,
                                ex.args[2],
                                Expr(:call,
                                     {
                                         :-,
                                         ex.args[3],
                                         1
                                     },
                                     Any)
                             },
                             Any)
                    },
                    Any)
    elseif ex.args[2] == target && ex.args[3] == target
        return Expr(:call,
                    {
                        :*,
                        Expr(:call,
                             {
                                :^,
                                target,
                                target
                             },
                             Any),
                        Expr(:call,
                             {
                                :+,
                                Expr(:call,
                                     {
                                        :log,
                                        target
                                     },
                                     Any),
                                1
                             },
                             Any)
                    },
                    Any)
    elseif ex.args[2] != target && ex.args[3] != target
        return ex
    else
        return Expr(:call,
                    {
                        :*,
                        Expr(:call,
                             {
                                :^,
                                ex.args[2],
                                ex.args[3]
                             },
                             Any),
                        Expr(:call,
                             {
                                :log,
                                ex.args[2]
                             },
                             Any)
                    },
                    Any)
    end
end

# The Sin Rule:
# d/dx sin(x) = cos(x)
function differentiate_sin(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :sin
        error("Not a valid sin call: $(ex)")
    end
    return Expr(:call,
                {
                    :*,
                    Expr(:call,
                         {
                            :cos,
                            ex.args[2]
                         },
                         Any),
                    differentiate(ex.args[2], target)
                },
                Any)
end

# The Cos Rule:
# d/dx cos(x) = -sin(x)
function differentiate_cos(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :cos
        error("Not a valid cos call: $(ex)")
    end
    return Expr(:call,
                {
                    :*,
                    Expr(:call,
                         {
                            :-,
                            Expr(:call,
                                 {
                                    :sin,
                                    ex.args[2]
                                 },
                                 Any)
                         },
                         Any),
                    differentiate(ex.args[2], target)
                },
                Any)
end

# The Tan Rule:
# d/dx tan(x) = 1 + tan(x)^2
function differentiate_tan(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :tan
        error("Not a valid tan call: $(ex)")
    end
    return Expr(:call,
                {
                    :*,
                    Expr(:call,
                         {
                            :+,
                            1,
                            Expr(:call,
                                 {
                                    :^,
                                    Expr(:call,
                                         {
                                            :tan,
                                            ex.args[2]
                                         },
                                         Any),
                                    2
                                 },
                                 Any)
                         },
                         Any),
                    differentiate(ex.args[2], target)
                },
                Any)
end

# The Exp Rule:
# d/dx exp(x) = exp(x)
function differentiate_exp(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :exp
        error("Not a valid exp call: $(ex)")
    end
    return Expr(:call,
                {
                    :*,
                    Expr(:call,
                         {
                            :exp,
                            ex.args[2]
                         },
                         Any),
                    differentiate(ex.args[2], target)
                },
                Any)
end

# The Log Rule:
# d/dx log(x) = 1 / x
function differentiate_log(ex::Expr, target::Symbol)
    if ex.head != :call || ex.args[1] != :log
        error("Not a valid log call: $(ex)")
    end
    return Expr(:call,
                {
                    :*,
                    Expr(:call,
                         {
                            :/,
                            1,
                            ex.args[2]
                         },
                         Any),
                    differentiate(ex.args[2], target)
                },
                Any)
end

# Lookup table of differentation rules
differentiate_lookup = {
                          :+ => differentiate_sum,
                          :- => differentiate_subtraction,
                          :* => differentiate_product,
                          :/ => differentiate_quotient,
                          :^ => differentiate_power,
                          :sin => differentiate_sin,
                          :cos => differentiate_cos,
                          :tan => differentiate_tan,
                          :exp => differentiate_exp,
                          :log => differentiate_log
                       }

function differentiate(ex::Expr, target::Symbol)
    if ex.head == :call
        if has(differentiate_lookup, ex.args[1])
            return simplify(differentiate_lookup[ex.args[1]](ex, target))
        else
            error("Don't know how to differentiate $(ex.args[1])")
        end
    else
        return simplify(differentiate(ex.head))
    end
end
function differentiate(ex::Expr, targets::Vector{Symbol})
    n = length(targets)
    exprs = Array(Expr, n)
    for i in 1:n
        exprs[i] = differentiate(ex, targets[i])
    end
    return exprs
end

differentiate(ex::Expr) = differentiate(ex, :x)

function differentiate(s::String, target::Symbol)
    differentiate(parse(s)[1], target)
end
function differentiate(s::String, targets::Vector{Symbol})
    differentiate(parse(s)[1], targets)
end
function differentiate(s::String, target::String)
    differentiate(parse(s)[1], symbol(target))
end
function differentiate{T <: String}(s::String, targets::Vector{T})
    differentiate(parse(s)[1], map(target -> symbol(target), targets))
end
function differentiate(s::String)
    differentiate(parse(s)[1], :x)
end

# begin x = 1; eval(differentiate(:(sin(x)), :x)) end

# Full out differentation returns an immediately evaluable Julia function

# function derivative(ex::Expr, target::Symbol)
#     function f(x)
#         d_ex = differentiate(ex, target)
#         return eval(d_ex)
#     end
#     return f
# end

# function derivative(ex::Expr, target::Symbol, x::Any)
#     d_ex = differentiate(ex, target)
#     return eval(d_ex)
# end


# Helper functions for spherical geometry

using Rotations

"""
    c2s(x, y, z) -> [θ, ϕ]
    c2s(xyz::AbstractVector) -> [θ, ϕ]
    c2s(xyz::AbstractMatrix) -> [θs; ϕs]

Convert Cartesian vector `[x, y, z]` (or `xyz`) to spherical unit vector `[θ,
ϕ]`.  For AbstractMatrix `xyz`, each column is converted and the results are
reassmebled into a Matrix using `hcat`.
"""
function c2s(x,y,z)
    θ=atan(y,x)
    ϕ=atan(z, hypot(x,y))
    [θ, ϕ]
end

function c2s(xyz::AbstractVector{<:Real})
    c2s(xyz[1], xyz[2], xyz[3])
end

function c2s(xyz::AbstractMatrix{<:Real})
    @views reduce(hcat, c2s.(xyz[1,:], xyz[2,:], xyz[3,:]))
end

"""
    s2c(θ, ϕ) -> [x, y, z)
    s2c(θϕ::AbstractVector) -> [x, y, z]
    s2c(θϕ::AbstractMaxtrix) -> [xs; ys; zs]

Convert spherical unit vector `[θ, ϕ]` (or `θϕ`) to Cartesian unit vector `[x,
y, z]`.  For AbstractMatrix `θϕ`, each column is converted and the results are
reassmebled into a Matrix using `hcat`.
"""
function s2c(θ,ϕ)
    z, r = sincos(ϕ)
    y, x = sincos(θ) .* r
    [x, y, z]
end

function s2c(θϕ::AbstractVector{<:Real})
    s2c(θϕ[1], θϕ[2])
end

function s2c(θϕ::AbstractMatrix{<:Real})
    @views reduce(hcat, s2c.(θϕ[1,:], θϕ[2,:]))
end

"""
    beamrings(θ, ϕ; nrings=4, dϕ=deg2rad(10/3600)) -> Matrix

Compute the centers of beams that are arranged in concentric rings around the
spherical coordinates `(θ, ϕ)`, both in radians with `θ` being the longitudinal
coordinate and `ϕ` being the latitudinal coordinate.  The number of rings for
which beam centers are returned is given by `nrings`, which defaults to 4.  The
rings are radially separated by `dϕ` radians, which defaults to the equivalent
of 10 arcseconds.  The beams returned do not form a packed hexagon because that
is not possible on a sperical sky, but the number of beams per ring does follow
the hexagonal packing convention of `6n` beams for ring `n` (when `n>0`).  The
so-called boresight beam center at `(θ, ϕ)` is included as well, so 61 beam
centers are returned for `nrings==4`.  More generally,
`nbeams=6*sum(1:nrings)+1` beam centers are returned in a `2xnbeams` Matrix,
where the first row contains the `θ_beam` coordinates and the second row
contains the `ϕ_beam` coordinates.
"""
function beamrings(θ, ϕ; nrings=4, dϕ=deg2rad(10/3600))
    nbeams = 6 * sum(1:nrings) + 1
    θs = Vector{Float64}(undef, nbeams)
    ϕs = Vector{Float64}(undef, nbeams)
    θs[1] = 0.0
    ϕs[1] = π/2
    b=2
    for r in 1:nrings
        nθ = 6r
        dθ = 2π/nθ
        ϕr = π/2 - r * dϕ
        for i in 0:nθ-1
            θs[b] = i*dθ
            ϕs[b] = ϕr
            b += 1
        end
    end
    reduce(hcat, c2s.(Ref(RotZY(θ, π/2-ϕ)) .* s2c.(θs, ϕs)))
end


using LinearAlgebra
using Plots

datapoints=[2.4 420;
            2.0 350;
            2.1 310;
            1.8 280;
            1.3 75]
A = hcat(ones(size(datapoints,1)),datapoints[:,1])
b = datapoints[:,2]

x = A\b

line = t -> x[1] + x[2]*t

scatter(datapoints[:,1], datapoints[:,2], legend=:none)
plot!(line, legend=:none)
ylabel!("weight, in pounds")
xlabel!("length, in meters")


euler(f::Function, T::Number, t0::Int, t1::Int, h::Int) = collect(begin T += h * f(T); T end for t in t0:h:t1)

# Prints a series of arbitrary values in a tabular form, left aligned in cells with a given width
tabular(width, cells...) = println(join(map(s -> rpad(s, width), cells)))

# prints the table according to the task description for h=5 and 10 sec
for h in (5, 10)
    print("Step $h:\n\n")
    tabular(15, "Time", "Euler", "Analytic")
    t = 0
    for T in euler(y -> -0.07 * (y - 20.0), 100.0, 0, 100, h)
        tabular(15, t, round(T,6), round(20.0 + 80.0 * exp(-0.07t), 6))
        t += h
    end
    println()
end


export correlationsum_fixedmass

using SpecialFunctions, Random

"""
    correlationsum_fixedmass(data, max_j; metric = Euclidean(), M = length(data)) → rs, ys
A fixed mass algorithm for the calculation of the fractal dimension ``\\Delta`` according
to [^Grassberger1988] with `max_j` the maximum number of neighbours that
should be considered for the calculation, `M` defines the number of points
considered for the calculation, default is the whole data set.

Implements
```math
\\Delta \\overline{\\log r^{(j)}} \\sim Ψ(j) - \\log N
```
where `` \\Psi(j) = \\frac{\\text{d} \\log Γ(j)}{\\text{d} j}
``, `rs` = ``\\overline{\\log r^{(j)}}`` and `ys` = ``\\Psi(j) - \\log N``.

``\\Delta`` can be computed by using `linear_region(rs, ys)`.

[^Grassberger1988]: Peter Grassberger (1988) [Finite sample Corrections to Entropy and Dimension Estimates, Physics Letters A 128(6-7)](https://doi.org/10.1016/0375-9601(88)90193-4)
"""
function correlationsum_fixedmass(data, max_j; metric = Euclidean(), M = length(data), w = 0)
    N = length(data)
    M > N && throw(ArgumentError("The number of points used for the calculation `M` should not exceed the number of points provided."))
    # Transform the data into a tree.
    tree = searchstructure(KDTree, data, metric)
    # Calculate all nearest neighbours up to max_j.
    _, distances =
        bulksearch(
            tree,
            N == M ? data.data : data[view(randperm(N), 1:M)],
            NeighborNumber(max_j),
            Theiler(w),
        )
    # The epsilons define the left side of the equation
    ys = [digamma(j) - log(N) for j in 1:max_j]
    # Holds the mean value of the logarithms of the distances.
    rs = zeros(max_j)
    for dists in distances
        for j in 1:max_j
            rs[j] += log(dists[j])
        end
    end
    return rs ./ M, ys
end


@test factorial(7) == 5040
@test factorial(7,3) == 7*6*5*4
@test binomial(5,3) == 10
p = shuffle([1:1000])
@test isperm(p)
@test all(invperm(invperm(p)) .== p)
push!(p, 1)
@test !isperm(p)
a = randcycle(10)
@test ipermute!(permute!([1:10], a),a) == [1:10]
@test collect(combinations("abc",2)) == ["ab","ac","bc"]
@test collect(permutations("abc")) == ["abc","acb","bac","bca","cab","cba"]
@test collect(filter(x->(iseven(x[1])),permutations([1,2,3]))) == {[2,1,3],[2,3,1]}
@test collect(filter(x->(iseven(x[3])),permutations([1,2,3]))) == {[1,3,2],[3,1,2]}
@test collect(filter(x->(iseven(x[1])),combinations([1,2,3],2))) == {[2,3]}
@test collect(partitions(4)) ==  {[4], [3,1], [2,2], [2,1,1], [1,1,1,1]}
@test collect(partitions(8,3)) == {[6,1,1], [5,2,1], [4,3,1], [4,2,2], [3,3,2]}
@test collect(partitions([1,2,3])) == {{[1,2,3]}, {[1,2],[3]}, {[1,3],[2]}, {[1],[2,3]}, {[1],[2],[3]}}
@test collect(partitions([1,2,3,4],3)) == {{[1,2],[3],[4]}, {[1,3],[2],[4]}, {[1],[2,3],[4]}, {[1,4],[2],[3]}, {[1],[2,4],[3]},{[1],[2],[3,4]}}

@test length(collect(partitions(30))) == length(partitions(30))
@test length(collect(partitions(90,4))) == length(partitions(90,4))
@test length(collect(partitions('a':'h'))) == length(partitions('a':'h'))
@test length(collect(partitions('a':'h',5))) == length(partitions('a':'h',5))

for n = 0:7, k = 1:factorial(n)
    p = nthperm!([1:n], k)
    @test isperm(p)
    @test nthperm(p) == k
end

#Issue 6154
@test binomial(int32(34), int32(15)) == binomial(BigInt(34), BigInt(15)) == 1855967520
@test binomial(int64(67), int64(29)) == binomial(BigInt(67), BigInt(29)) == 7886597962249166160
@test binomial(int128(131), int128(62)) == binomial(BigInt(131), BigInt(62)) == 157311720980559117816198361912717812000


using OrdinaryDiffEq
using DiffEqCallbacks
using Surrogates
using DataFrames
using RCall
using Random
using Plots;

R"library(hmer)";

Random.seed!(123)
R"set.seed(123)";

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

tmax = 10000.0
tspan = (0.0,tmax)
δt = 1.0;

n_train = 50 # Number of training samples
n_test = 1000; # Number of test samples

# Parameters are β, γ
lb = [0.00005, 0.1]
ub = [0.001, 1.0];

N = 1000.0
u0 = [990.0,10.0,0.0]
p = [0.0005,0.25]
prob_ode = ODEProblem(sir_ode!,u0,tspan,p);

sampler = LatinHypercubeSample();

θ = Surrogates.sample(n_train,lb,ub,sampler);

logit = (x) -> log(x/(1-x))
invlogit = (x) -> exp(x)/(exp(x)+1.0)
cb_ss = TerminateSteadyState()
logit_final_size = function(z)
  prob = remake(prob_ode;p=z)
  sol = solve(prob, ROS34PW3(),callback=cb_ss)
  fsp = sol[end][3]/N
  logit(fsp)
end;

lfs = logit_final_size.(θ);

training_df = DataFrame(θ)
rename!(training_df,["b","g"])
training_df[!,:lfs] = lfs;

input_ranges = Dict(:b => [lb[1], ub[1]], :g => [lb[2], ub[2]])
output_names = ["lfs"];

@rput training_df
@rput output_names
@rput input_ranges;

R"emulator <- emulator_from_data(training_df, output_names, input_ranges)"

θ_test = sample(n_test,lb,ub,sampler)
lfs_test = logit_final_size.(θ_test);

test_df = DataFrame(θ_test)
rename!(test_df,["b","g"])
test_df[!,:lfs] = lfs_test
@rput test_df; # copy to R

R"lfs_test_pred <- list()"
R"lfs_test_pred$mean <- emulator$lfs$get_exp(test_df)"
R"lfs_test_pred$unc <- emulator$lfs$get_cov(test_df)"
@rget lfs_test_pred;

scatter(invlogit.(lfs_test),
        invlogit.(lfs_test_pred[:mean]),
        xlabel = "Model final size",
        ylabel = "Surrogate final size",
        legend = false,
        title = "Test set")

β_grid = collect(lb[1]:0.00001:ub[1])
θ_eval = [(βᵢ,0.25) for βᵢ in β_grid]
eval_df = DataFrame(θ_eval)
rename!(eval_df,["b","g"])
@rput eval_df
R"lfs_eval <- list()"
R"lfs_eval$mean <- emulator$lfs$get_exp(eval_df)"
R"lfs_eval$unc <- emulator$lfs$get_cov(eval_df)"
@rget lfs_eval
fs_eval = invlogit.(lfs_eval[:mean])
fs_eval_uc = invlogit.(lfs_eval[:mean] .+ 1.96 .* sqrt.(lfs_eval[:unc]))
fs_eval_lc = invlogit.(lfs_eval[:mean] .- 1.96 .* sqrt.(lfs_eval[:unc]))
plot(β_grid,
     fs_eval,
     xlabel = "Infectivity parameter, β",
     ylabel = "Final size",
     label = "Model")
plot!(β_grid,
      invlogit.(logit_final_size.(θ_eval)),
      ribbon = (fs_eval .- fs_eval_lc, fs_eval_uc - fs_eval),
      label = "Surrogate",
      legend = :right)

γ_grid = collect(lb[2]:0.001:ub[2])
θ_eval = [(0.001,γᵢ) for γᵢ in γ_grid]
eval_df = DataFrame(θ_eval)
rename!(eval_df,["b","g"])
@rput eval_df
R"lfs_eval <- list()"
R"lfs_eval$mean <- emulator$lfs$get_exp(eval_df)"
R"lfs_eval$unc <- emulator$lfs$get_cov(eval_df)"
@rget lfs_eval
fs_eval = invlogit.(lfs_eval[:mean])
fs_eval_uc = invlogit.(lfs_eval[:mean] .+ 1.96 .* sqrt.(lfs_eval[:unc]))
fs_eval_lc = invlogit.(lfs_eval[:mean] .- 1.96 .* sqrt.(lfs_eval[:unc]))
plot(γ_grid,
     fs_eval,
     xlabel = "Recovery rate, γ",
     ylabel = "Final size",
     label = "Model")
plot!(γ_grid,
      invlogit.(logit_final_size.(θ_eval)),
      ribbon = (fs_eval .- fs_eval_lc, fs_eval_uc - fs_eval),
      label = "Surrogate")

obs = logit_final_size(p)
obs_lc = logit(0.99*invlogit(obs))
obs_uc = logit(1.01*invlogit(obs))
target = Dict(:lfs => [obs_lc, obs_uc])
@rput target

R"new_points <- generate_new_runs(emulator, $n_test, target, method = 'lhs', cutoff = 3)"
@rget new_points;

l = @layout [a b]
pl1 = histogram(new_points[!,:b],legend=false,xlim=(lb[1],ub[1]),bins=lb[1]:0.00005:ub[1],title="NROY values for β")
vline!(pl1,[p[1]])
pl2 = histogram(new_points[!,:g],legend=false,xlim=(lb[2],ub[2]),bins=lb[2]:0.05:ub[2],title="NROY values for γ")
vline!(pl2,[p[2]])
plot(pl1, pl2, layout = l)


using Test, DayCounts, Dates
using XLSX
using Tables

xf = XLSX.readxlsx("yearfrac.xlsx")

function range_to_table(rng)
    Tables.table(rng[2:end,:]; header = Symbol.(rng[1,:]))
end

@testset "Excel basis 0" begin
    table = range_to_table(xf["basis0"][:])
    dc = DayCounts.Thirty360Excel()
    for r in Tables.rows(table)
        @test yearfrac(Date(r.startdate), Date(r.enddate), dc) == r.yearfrac
    end
end

@testset "Excel basis 1" begin
    table = range_to_table(xf["basis1"][:])
    dc = DayCounts.ActualActualExcel()
    for r in Tables.rows(table)
        @test yearfrac(Date(r.startdate), Date(r.enddate), dc) == r.yearfrac
    end
end

@testset "Excel basis 2" begin
    table = range_to_table(xf["basis2"][:])
    dc = DayCounts.Actual360()
    for r in Tables.rows(table)
        @test yearfrac(Date(r.startdate), Date(r.enddate), dc) == r.yearfrac
    end
end

@testset "Excel basis 3" begin
    table = range_to_table(xf["basis3"][:])
    dc = DayCounts.Actual365Fixed()
    for r in Tables.rows(table)
        @test yearfrac(Date(r.startdate), Date(r.enddate), dc) == r.yearfrac
    end
end

@testset "Excel basis 4" begin
    table = range_to_table(xf["basis4"][:])
    dc = DayCounts.ThirtyE360()
    for r in Tables.rows(table)
        @test yearfrac(Date(r.startdate), Date(r.enddate), dc) == r.yearfrac
    end
end


"""
    GaussLik <: Likelihood

Gaussian, a.k.a. Normal, likelihood
```math
p(y | f, σ) = 1 / √(2πσ²) \\exp(-(x - f)²/(2σ²)),
```
where standard deviation ``σ`` is a non-fixed hyperparameter and ``f`` is the latent
Gaussian process.
"""
mutable struct GaussLik <: Likelihood
    "Standard deviation"
    σ::Float64
    "Priors for likelihood parameters"
    priors::Array

    """
        GaussLik(lσ::Float64)

    Create `GaussLik` with standard deviation `exp(lσ)`.
    """
    GaussLik(lσ::Float64) = new(exp(lσ), [])
end

# log of probability density
function log_dens(gauss::GaussLik, f::VecF64, y::VecF64)
    return (-0.5 * log(2 * pi) - log(gauss.σ)) .- 0.5 * ((y - f) / gauss.σ).^2
end

# derivative of log pdf wrt latent function
function dlog_dens_df(gauss::GaussLik, f::VecF64, y::VecF64)
    return [(yi-fi)/gauss.σ^2 for (fi,yi) in zip(f,y)]
end

# derivative of log pdf wrt to parameters
function dlog_dens_dθ(gauss::GaussLik, f::VecF64, y::VecF64)
    return gauss.σ*[-1/gauss.σ + 1/gauss.σ^3*(yi-fi).^2 for (fi,yi) in zip(f,y)]
end

#mean and variance under likelihood
mean_lik(gauss::GaussLik, f::VecF64) = f
var_lik(gauss::GaussLik, f::VecF64) = ones(length(f))*gauss.σ^2

function set_params!(gauss::GaussLik, hyp::VecF64)
    length(hyp) == 1 || throw(ArgumentError("Gaussian/Normal likelihood has only one free parameter"))
    gauss.σ = exp(hyp[])
end

get_params(gauss::GaussLik) = Float64[log(gauss.σ)]
num_params(gauss::GaussLik) = 1

#Computes the predictive mean and variance
function predict_obs(gauss::GaussLik, fmean::VecF64, fvar::VecF64)
    return fmean, fvar + gauss.σ^2
end


export
    Line,
    get_polar_angle,
    get_distance,
    get_side

immutable Line
    C::VecE2
    θ::Float64
end
Line(A::VecE2, B::VecE2) = Line((A+B)/2, atan2(B - A))

Base.:-(line::Line, V::VecE2) = Line(line.C - V, line.θ)
Base.:+(line::Line, V::VecE2) = Line(line.C + V, line.θ)

get_polar_angle(line::Line) = line.θ
rot(line::Line, Δθ::Float64) = Line(line.C, line.θ+Δθ)
Base.rot180(line::Line) = rot(line, 1π)
Base.rotl90(line::Line) = rot(line,  π/2)
Base.rotr90(line::Line) = rot(line, -π/2)

"""
The distance between the line and the point P
"""
function get_distance(line::Line, P::VecE2)

    ab = polar(1.0, line.θ)
    pb = P - line.C

    denom = abs2(ab)
    if denom ≈ 0.0
        return 0.0
    end

    r = dot(ab, pb)/denom
    return abs(P - (line.C + r*ab))
end

"""
What side of the line you are on
Is -1 if on the left, 1 if on the right, and 0 if on the line
"""
function get_side(line::Line, p::VecE2, ε::Float64=1e-10)
    ab = polar(1.0, line.θ)
    signed_dist = ab.x*(p.y-line.C.y) - ab.y*(p.x-line.C.x)
    if abs(signed_dist) < ε
        return 0
    else
        return sign(signed_dist)
    end
end


"""
    frst(signal, α, p)

Computing the α order fractional sine transform of the input **signal**.

# Example

```julia-repl
julia> frst([1,2,3], 0.5, 2)
3-element Vector{ComplexF64}:
1.707106781186548 + 1.207106781186547im
1.9999999999999998 - 1.7071067811865481im
-1.1213203435596437 - 1.2071067811865468im
```

### References

```tex
@article{article,
author = {Pei, Soo-Chang and Yeh, Min-Hung},
year = {2001},
month = {07},
pages = {1198 - 1207},
title = {The discrete fractional cosine and sine transforms},
volume = {49},
journal = {Signal Processing, IEEE Transactions on},
doi = {10.1109/78.923302}
}
```
"""
function frst(signal, α, p)
    N = length(signal)
    @views signal = signal[:]
    p = min(max(2, p), N-1)
    E = dFRST(N,p)
    result = E *(exp.(-im*pi*α*collect(0:N-1)) .*(E' *signal))
    return result
end

function dFRST(N, p)
    N1 = 2*N+2
    d2 = [1, -2, 1]

    d_p = 1
    s = 0
    st = zeros(1, N1)

    for k = 1:floor(Int, p/2)
        if isa(d_p, Number)
            d_p = @. d2*d_p
        else
            d_p = conv(d2, d_p)
        end
        st[vcat(collect(N1-k+1:N1), collect(1:k+1))] = d_p
        st[1]=0
        temp=vcat(union(1, collect(1:k-1)), union(2,collect(1:k-1)))
        temp = temp[:] ./collect(1:2*k)
        s=s.+(-1)^(k-1)*prod(temp)*2*st
    end

    H = Toeplitz(s[:], s[:]) +diagm(real.(fft(s[:])))

    V = hcat(zeros(N), reverse(zeros(N, N)+I, dims=1), zeros(N), -(zeros(N, N)+I)) ./sqrt(2)

    Od = V*H*V'

    _, vo = eigen(Od)

    return reverse(reverse(vo, dims=2), dims=1)
end


using EngEconomics, Plots, ColorSchemes, Colors
plotly()

uwEngColorScheme = ColorScheme([parse(Colorant, "#D0B4E7"), parse(Colorant, "#FFFFFF"), parse(Colorant, "#57058B"), parse(Colorant, "#8100B4")])

F = 50
A_base = 7
i_base = 0.10
sensitivityVector = [.90, 0.95, 1.0, 1.05, 1.10]

baseN = log(F / A_base * i_base + 1) / log(1 + i_base)

N_annualSavings = log.(F ./ (A_base .* sensitivityVector) .* i_base .+ 1) ./ log(1 + i_base)
N_interestRate  = log.(F / A_base .* (sensitivityVector * i_base) .+ 1) ./ log.(1 .+ (i_base .* sensitivityVector))

plot(sensitivityVector, N_annualSavings, label="Annual Savings Sensitivity",
	xlabel="Sensitivity (+/- 5%, 10%)", ylabel="Number of Years",
	palette=cgrad(uwEngColorScheme))
plot!(sensitivityVector, N_interestRate, label="Interest Rate Sensitivity")


"""
perfect_square(number)

Check if a number is a perfect square or not.

# Example
```jula
perfect_square(9)   # returns True
perfect_square(16)  # returns True
perfect_square(1)   # returns True
perfect_square(0)   # returns True
perfect_square(10)  # returns False
```

Contributed By:- [Ashwani Rathee](https://github.com/ashwani-rathee)
"""
function perfect_square(number)
    return sqrt(number) * sqrt(number) == number
end
perfect_square(9)


using GeneralizedSampling
using Base.Test

#=
Various small functions are defined to interact with a Freq2Wave type.
=#

# ------------------------------------------------------------
# Creating matrices

J = 3
M = 2^(J+1)
N = 2^J

# Uniform samples, No Boundary
Usamples = grid(M, 0.5)
TUNB = Freq2Wave(Usamples, "haar", J)

# Uniform samples, with Boundary
wavename = "db3"
TUB = Freq2Wave(Usamples, wavename, J)

# Non-Uniform samples, No Boundary
K = N/2
NUsamples = N*rand(M) - K
TNUNB = Freq2Wave(NUsamples, "haar", J, K)

# Non-Uniform samples, with Boundary
TNUB = Freq2Wave(NUsamples, wavename, J, K)

# ------------------------------------------------------------
# Testing

@test isuniform(TUNB)
@test isuniform(TUB)
@test !isuniform(TNUNB)
@test !isuniform(TNUB)

@test !hasboundary(TUNB)
@test hasboundary(TUB)
@test !hasboundary(TNUNB)
@test hasboundary(TNUB)

@test dim(TUNB) == dim(TUB) == dim(TNUNB) == dim(TNUB) == 1

@test wscale(TUNB) == wscale(TUB) == wscale(TNUNB) == wscale(TNUB) == J

@test wsize(TUNB) == wsize(TUB) == wsize(TNUNB) == wsize(TNUB) == (N,)

@test size(TUNB) == size(TUB) == size(TNUNB) == size(TNUB) == (M,N)

@test van_moment(TUNB) == van_moment(TNUNB) == 1
@test van_moment(TUB) == van_moment(TNUB) == van_moment(wavename)


using ArrayFire
using BenchmarkTools

function pi_gpu(n)
    return 4 *
           sum(rand(AFArray{Float64}, n)^2 .+ rand(AFArray{Float64}, n)^2 .<= 1) ./
           n
end

@btime pi_gpu(10_000_000)

function pi_serial(n)
    inside = 0
    for i = 1:n
        x, y = rand(), rand()
        inside += (x^2 + y^2 <= 1)
    end
    return 4 * inside / n
end

@btime pi_serial(10_000_000)

#%% Mandelbrot fractal
using Plots
gr()

maxIterations = 500
gridSize = 2048
xlim = [-0.748766713922161, -0.748766707771757]
ylim = [0.123640844894862, 0.123640851045266]

x = range(xlim[1], xlim[2], length = gridSize)
y = range(ylim[1], ylim[2], length = gridSize)

xGrid = [i for i in x, j in y]
yGrid = [j for i in x, j in y]

z0 = xGrid + im * yGrid

function mandelbrotGPU(z0, maxIterations)
    z = z0
    count = ones(AFArray{Float32}, size(z))

    for n = 1:maxIterations
        z = z .* z .+ z0
        count = count + (abs(z) <= 2)
    end
    return sync(log(count))
end

count = Array(mandelbrotGPU(AFArray(z0), maxIterations))
heatmap(
    count,
    color = :inferno,
    grid = false,
    colorbar = :none,
    framestyle = :none,
    size = (2048, 2048),
)
savefig("mandelbrot.png")


# span.jl

@doc """
### findSpan

Determines the knot span index, i

**Inputs**

*  `u` : coordinate value (u,v,w is the coordinate space in 3D)
*  `U` : Knot vector
*  `k` : B-spline order
*  `nctl` : Number of control points

**Outputs**

*  `span` : Knot span index

SOURCE: The NURBS book 2nd Edition, Algorithm A2.1

"""->

function findSpan(u::Tffd, U::AbstractArray{Tffd,1}, k::Int, nctl::Int) where Tffd

  if u >= U[nctl+1]
    return nctl  # Special case when u = last term of knot vector
  elseif u < U[k]
    return k  # When u lies at the starting point of the curve
  end

  low = k-1
  high = nctl+1
  span = div(low + high, 2)
  # Do a binary search
  while u < U[span] || u >= U[span+1]
    if u < U[span]
      high = span
    else
      low = span
    end  # End if
    span = div(low+high, 2)
  end  # End while

  return span
end  # End function findSpan


## Linear algebra API for adjacency matrix
using LinearAlgebra

function adjacency_matrix(adj::AbstractMatrix, T::DataType=eltype(adj))
    m, n = size(adj)
    (m == n) || throw(DimensionMismatch("adjacency matrix is not a square matrix: ($m, $n)"))
    T.(adj)
end

"""
    degrees(g[, T; dir=:out])

Degree of each vertex. Return a vector which contains the degree of each vertex in graph `g`.

# Arguments
- `g`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `T`: result element type of degree vector; default is the element type of `g` (optional).
- `dir`: direction of degree; should be `:in`, `:out`, or `:both` (optional).

# Examples
```jldoctest
julia> using GeometricFlux

julia> m = [0 1 1; 1 0 0; 1 0 0];

julia> GeometricFlux.degrees(m)
3-element Array{Int64,1}:
 2
 1
 1

```
"""
function degrees(adj::AbstractMatrix, T::DataType=eltype(adj); dir::Symbol=:out)
    if issymmetric(adj)
        d = vec(sum(adj, dims=1))
    else
        if dir == :out
            d = vec(sum(adj, dims=1))
        elseif dir == :in
            d = vec(sum(adj, dims=2))
        elseif dir == :both
            d = vec(sum(adj, dims=1)) + vec(sum(adj, dims=2))
        else
            throw(DomainError(dir, "invalid argument, only accept :in, :out and :both"))
        end
    end
    d
end

"""
    degree_matrix(g[, T; dir=:out])

Degree matrix of graph `g`. Return a matrix which contains degrees of each vertex in its diagonal.
The values other than diagonal are zeros.

# Arguments
- `g`: should be a adjacency matrix, `FeaturedGraph`, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `T`: result element type of degree vector; default is the element type of `g` (optional).
- `dir`: direction of degree; should be `:in`, `:out`, or `:both` (optional).

# Examples
```jldoctest
julia> using GeometricFlux

julia> m = [0 1 1; 1 0 0; 1 0 0];

julia> GeometricFlux.degree_matrix(m)
3×3 SparseArrays.SparseMatrixCSC{Int64,Int64} with 3 stored entries:
  [1, 1]  =  2
  [2, 2]  =  1
  [3, 3]  =  1
```
"""
function degree_matrix(adj::AbstractMatrix, T::DataType=eltype(adj); dir::Symbol=:out)
    d = degrees(adj, T, dir=dir)
    return SparseMatrixCSC(T.(diagm(0=>d)))
end

"""
    inv_sqrt_degree_matrix(g[, T; dir=:out])

Inverse squared degree matrix of graph `g`. Return a matrix which contains inverse squared degrees of each vertex in its diagonal.
The values other than diagonal are zeros.

# Arguments
- `g`: should be a adjacency matrix, `FeaturedGraph`, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `T`: result element type of degree vector; default is the element type of `g` (optional).
- `dir`: direction of degree; should be `:in`, `:out`, or `:both` (optional).
"""
function inv_sqrt_degree_matrix(adj::AbstractMatrix, T::DataType=eltype(adj); dir::Symbol=:out)
    d  = inv.(sqrt.(degrees(adj, T, dir=dir)))
    return Diagonal(d)
end

"""
    laplacian_matrix(g[, T; dir=:out])

Laplacian matrix of graph `g`.

# Arguments
- `g`: should be a adjacency matrix, `FeaturedGraph`, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `T`: result element type of degree vector; default is the element type of `g` (optional).
- `dir`: direction of degree; should be `:in`, `:out`, or `:both` (optional).
"""
function laplacian_matrix(adj::AbstractMatrix, T::DataType=eltype(adj); dir::Symbol=:out)
    degree_matrix(adj, T, dir=dir) - SparseMatrixCSC(T.(adj))
end

"""
    normalized_laplacian(g[, T; selfloop=false])

Normalized Laplacian matrix of graph `g`.

# Arguments
- `g`: should be a adjacency matrix, `FeaturedGraph`, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `T`: result element type of degree vector; default is the element type of `g` (optional).
- `selfloop`: adding self loop while calculating the matrix (optional).
"""
function normalized_laplacian(adj::AbstractMatrix, T::DataType=eltype(adj); selfloop::Bool=false)
    selfloop && (adj += I)
    inv_sqrtD = inv_sqrt_degree_matrix(adj, T, dir=:both)
    T.(I - inv_sqrtD * adj * inv_sqrtD)
end

@doc raw"""
    scaled_laplacian(adj::AbstractMatrix[, T::DataType])

Scaled Laplacien matrix of graph `g`,
defined as ``\hat{L} = \frac{2}{\lambda_{max}} L - I`` where ``L`` is the normalized Laplacian matrix.

# Arguments
- `g`: should be a adjacency matrix, `FeaturedGraph`, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `T`: result element type of degree vector; default is the element type of `g` (optional).
"""
function scaled_laplacian(adj::AbstractMatrix, T::DataType=eltype(adj))
    T(2. / eigmax(adj)) * normalized_laplacian(adj, T) - I
end

function neighbors(adj::AbstractMatrix, T::DataType=eltype(adj))
    n = size(adj,1)
    @assert n == size(adj,2) "adjacency matrix is not a square matrix."
    A = (adj .!= zero(T))
    if !issymmetric(adj)
        A = A .| A'
    end
    indecies = collect(1:n)
    ne = Vector{Int}[indecies[view(A, :, i)] for i = 1:n]
    return ne
end


# This example solution is partially derived from the Julia base implementation of rational numbers (MIT license)
# https://github.com/JuliaLang/julia/blob/52bafeb981bac548afd2264edb518d8d86944dca/base/rational.jl
# https://github.com/JuliaLang/julia/blob/52bafeb981bac548afd2264edb518d8d86944dca/LICENSE.md

import Base: +, *, ^, -, /, <, <=, inv, show, promote_rule

struct RationalNumber{T<:Integer} <: Real
    num::T
    den::T

    # reduce to lowest terms
    function RationalNumber{T}(num::Integer, den::Integer) where T<:Integer
        num == den == zero(T) && throw(ArgumentError("zero($T)//zero($T) is not a valid rational number"))
        num2, den2 = sign(den) < 0 ? Base.divgcd(-num, -den) : Base.divgcd(num, den)
        new(num2, den2)
    end
end

RationalNumber(num::T, den::T) where {T<:Integer} = RationalNumber{T}(num, den)

numerator(x::RationalNumber) = x.num
denominator(x::RationalNumber) = x.den

RationalNumber{T}(x::Integer) where {T<:Integer} = RationalNumber{T}(convert(T, x), convert(T, 1))

promote_rule(::Type{RationalNumber{T}}, ::Type{S}) where {T<:Integer,S<:Integer} = RationalNumber{promote_type(T,S)}

show(io::IO, x::RationalNumber) = (show(io, numerator(x)); print(io, "//"); show(io, denominator(x)))

# Helper method
numden(x::RationalNumber) = numerator(x), denominator(x)

# Arithmetics
function +(x::RationalNumber, y::RationalNumber)
    a, b = numden(x)
    c, d = numden(y)

    RationalNumber(a * d + b * c, b * d)
end

-(x::RationalNumber, y::RationalNumber) = x + (-1) * y
-(x::RationalNumber) = -1 * x

function *(x::RationalNumber, y::RationalNumber)
    a, b = numden(x)
    c, d = numden(y)

    RationalNumber(a * c, b * d)
end

inv(x::RationalNumber) = RationalNumber(denominator(x), numerator(x))
/(x::RationalNumber, y::RationalNumber) = x * inv(y)

^(x::Real, y::RationalNumber) = x^(numerator(y) / denominator(y))

# Ordering
<(x::RationalNumber, y::RationalNumber) = x.den == y.den ? x.num < y.num : x.num * y.den < y.num * x.den
<=(x::RationalNumber, y::RationalNumber) = x.den == y.den ? x.num <= y.num : x.num * y.den <= y.num * x.den


# indicator of the Cartesian product of real binary sets

export IndBinary

"""
**Indicator of the product of binary sets**

    IndBinary(low, up)

Returns the indicator function of the set
```math
S = \\{ x : x_i = low_i\\ \\text{or}\\ x_i = up_i \\},
```
Parameters `low` and `up` can be either scalars or arrays of the same dimension as the space.
"""
struct IndBinary{T, S} <: ProximableFunction
    low::T
    high::S
end

is_set(f::IndBinary) = true

IndBinary() = IndBinary(0.0, 1.0)

IndBinary_low(f::IndBinary{<: Number, S}, i) where S = f.low
IndBinary_low(f::IndBinary{T, S}, i) where {T, S} = f.low[i]
IndBinary_high(f::IndBinary{T, <: Number}, i) where T = f.high
IndBinary_high(f::IndBinary{T, S}, i) where {T, S} = f.high[i]

function (f::IndBinary)(x::AbstractArray{C}) where {R <: Real, C <: RealOrComplex{R}}
    for k in eachindex(x)
        if x[k] != IndBinary_low(f, k) && x[k] != IndBinary_high(f, k)
            return R(Inf)
        end
    end
    return R(0)
end

function prox!(y::AbstractArray{T}, f::IndBinary, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: Union{R, Complex{R}}}
    for k in eachindex(x)
        low = IndBinary_low(f, k)
        high = IndBinary_high(f, k)
        if abs(x[k] - low) < abs(x[k] - high)
            y[k] = low
        else
            y[k] = high
        end
    end
    return R(0)
end

fun_name(f::IndBinary) = "indicator of binary array"
fun_dom(f::IndBinary) = "AbstractArray{Real}"

function prox_naive(f::IndBinary, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: Union{R, Complex{R}}}
    distlow = abs.(x .- f.low)
    disthigh = abs.(x .- f.high)
    indlow = distlow .< disthigh
    indhigh = distlow .>= disthigh
    y = f.low.*indlow + f.high.*indhigh
    return y, R(0)
end


import Distributions: Normal, pdf, cdf

###############################
# Type Definition/Constructors
###############################
mutable struct Grid
    points::Vector{Float64}
    weights::Vector{Float64}
    scale::Float64
    Grid(points, weights, scale) = sum(weights) ≈ scale ? new(points, weights, scale) : error("scaled weights do not sum up properly")
end

# Constructor utilizing a custom weight calculation function
function Grid(quadrature::Function,
              lower_bound::T, upper_bound::T,
              n_points::Int; scale::Float64 = 1.) where {T<:Real}
    grid, weights = quadrature(lower_bound, upper_bound, n_points)
    Grid(reshape(grid, size(grid)[1]), weights, scale)
end

####################
# Quadrature rules
####################
# Enforce that quadrature rules with additional keyword arguments get mapped to
# A pre-populated version of that rule that only accepts 3 arguments:
# Lower bound, upper bound, and number of points
function uniform_quadrature(lower_bound::T, upper_bound::T, n_points::Int;
                            scale::T = 1) where {T<:Real}
    grid = collect(range(lower_bound, stop = upper_bound, length = n_points))
    weights = fill(scale/n_points, n_points)
    return grid, weights
end

function uniform_quadrature(scale::T = 1) where {T<:Real}
    return (lb, ub, n) -> uniform_quadrature(lb, ub, n, scale = scale)
end

# NOTE: Should modify chebpts function to return cleaner types...
# Kind indicates the "kind" of the chebyshev polynomial grid
# The terminology is generally, "Chebyshev polynomials of the Nth kind"
function curtis_clenshaw_quadrature(lower_bound::T, upper_bound::T,
                                    n_points::Int; kind::Int = 2) where {T<:Real}
    grid, weights = chebpts(n_points, lower_bound, upper_bound, kind)
    return grid', dropdims(weights', dims=2)
end

function curtis_clenshaw_quadrature(kind::Int = 2)
    return (lb, ub, n) -> curtis_clenshaw_quadrature(lb, ub, n, kind = kind)
end

# Tauchen86 developed a method for "finding a discrete-valued Markov chain
# whose sample paths approximate well those of a vector autoregression"
# Hence the method below is the Tauchen86 method applied to a general AR(1).
function tauchen86(μ::AbstractFloat,ρ::AbstractFloat,σ::AbstractFloat,n::Int64,λ::AbstractFloat)
    #output is xgrid, xprob
    # x_t+1 = μ + ρ x_t + σ e_{t+1}, e_t+1 ∼ N(0,1)
    xhi = μ/(1-ρ) + λ*sqrt(σ^2/(1-ρ)^2)
    xlo = μ/(1-ρ) - λ*sqrt(σ^2/(1-ρ)^2)
    xgrid = zeros(n);
    xscale =(xhi-xlo)/(n-1)

    for i=1:n
        xgrid[i] = xlo + xscale*(i-1)
    end
    m=zeros(n-1);
    for i=1:n-1
        m[i] = (xgrid[i]+xgrid[i+1])/2
    end
    xprob = zeros(n,n) # xprob[i,j] = Pr(x_t+1=xgrid[j]|x_t=xgrid[i])
    for i=1:n # this is the state today
        normie = Normal(μ+ρ*xgrid[i],σ)
        normpdf(x) = pdf.(normie,x)
        normcdf(x) = cdf.(normie,x)
        for j=2:n-1
            xprob[i,j] = normcdf(m[j]) - normcdf(m[j-1])
        end
        xprob[i,1] = normcdf(m[1])
        xprob[i,n] = 1 - normcdf(m[n-1])
    end
    xprob = xprob./sum(xprob, dims = 2) # make sure the rows sum to 1
    return ( xgrid,xprob, xscale )
end

# However, Tauchen86 can also be applied to "degenerate" AR(1)s where ρ = 0.
# i.e. i.i.d noise with mean μ, std. σ, but the way the returned
# Markov transition matrix, xprob, is used in the code, assumes it to be
# 1-dimensional, since the second-dimension of the Markov transition matrix is trivial
# (since the process is i.i.d).
function tauchen86(μ::AbstractFloat, σ::AbstractFloat, n::Int64, λ::AbstractFloat)
    xgrid, xprob, xscale = tauchen86(μ, 0., σ, n, λ)
    return xgrid, vec(xprob[:, 1]), xscale
end

####################
# Grid-based utils
####################
function get_grid(m::AbstractDSGEModel, grid_name::Symbol)
    return m.grids[grid_name]
end

function quadrature_sum(x::Vector{T}, grid::Grid) where {T<:Real}
    return sum(grid.weights .* x .* grid.points)
end

function quadrature_sum(grid::Grid, x::Vector{T}) where {T<:Real}
    return sum(x, grid)
end

# # Defining arithmetic on/standard function evaluation of grids
# for op in (:(Base.:+),
           # :(Base.:-),
           # :(Base.:*),
           # :(Base.:/))
           # # Exponentiation also omitted because of broadcasting issues

    # @eval ($op)(g::Grid, x::Integer)           = ($op)(g.points, x)
    # @eval ($op)(g::Grid, x::Number)            = ($op)(g.points, x)
    # if op == :(Base.:/)
        # op = :(Base.:./) # Will not work because the syntax for broadcasting has been
                         # # updated to broadcast(/, ...)?
        # @eval ($op)(x::Integer, g::Grid)       = ($op)(x, g.points)
        # @eval ($op)(x::Number, g::Grid)        = ($op)(x, g.points)
    # end
# end

# for f in (:(Base.:-),
          # :(Base.:<),
          # :(Base.:>),
          # :(Base.:<=),
          # :(Base.:>=))

    # @eval ($f)(g::Grid) = ($f)(g.points)

    # # Similar issues with broadcasting...
    # if f != :(Base.:-)
        # @eval ($f)(g::Grid, x::Number) = ($f)(g.points, x)
    # end
# end


include("../src/HMat2d.jl");

function laplacian2d(nx,ny)
    M = speye(nx*ny);
    for x=1:nx
        for y=1:ny
            s = x+(y-1)*nx;
            M[s,s] = 4;
            if x > 1
                M[s,s-1] = -1;
            end
            if x < nx
                M[s,s+1] = -1;
            end
            if y > 1
                M[s,s-nx] = -1;
            end
            if y < ny
                M[s,s+nx] = -1;
            end
        end
    end
    return M;
end

EPS = 1e-6;
MaxRank = 2;
minn = 8;

for n = [16 32 64]
    A = full(laplacian2d(n,n));
    Z2C = Z2Cmapper([n,n],minn,reshape(1:n^2,n,n));
    A = full(A[Z2C,Z2C]);
    AH = HMat2dd2h(A, [n,n], [n,n], WEAK, [0,0], [0,0], 1, EPS, MaxRank, minn);
    BH = hcopy(AH);
    v = rand(n^2,5);
    tic();
    uhmat = hmatvec(AH,v);
    toc();
    tic();
    CH = hmul(AH,BH);
    toc();
    CD = hh2d(CH);
    @printf("N^2 = %4d^2, hmul error: %.3e\n",n,norm(CD-A*A)/norm(A*A));
end


using NumericalMethodsforEngineers
using Test

# write your own tests here

code_tests = [
  # Ch01 tests
  "ch01_nmex.jl",

  # Ch02 tests
  "ch02_lufac.jl",
  "ch02_ldlt.jl",
  #"ch02_SymmetricBandedMatrix.jl",
  #"ch02_SkylineMatrix.jl",
  "ch02_gaussian.jl",
  "ch02_penalized.jl",
  "ch02_jacobi.jl",
  "ch02_bicgstab.jl",                               # Initial version of bicgstab

  # Ch03 tests
  # Ch04 tests

  "ch04_symata.jl",

  # Ch05 tests

  # Ch06 tests
  "ch06_p6.2.jl",

  # Ch07 tests
  "ch07_ivp_ex7.3.jl",                           # Four direct methods
  "ch07_bvp_ex7.4a.jl",                        # Shooting method example 7.4a
  "ch07_bvp_ex7.4b.jl",                        # Shooting method example 7.4b

  # Symata Weighted Residual Method tests
  "wrm_with_using_symata.jl",
  "wrm_with_using_nmfe.jl",
  "../examples/ch07/WRM/Ex.7.13.jl",         # 1 point LagrangePolynomial
  "../examples/ch07/WRM/Ex.7.14.jl",     # 2 point LagrangePolynomial
  "../examples/ch07/WRM/Ex.7.15.jl",     # 1 point collocation
  "../examples/ch07/WRM/Ex.7.15b.jl",    # 2 point collocation
  "../examples/ch07/WRM/Ex.7.16.jl",     # 1 point subdomain
  "../examples/ch07/WRM/Ex.7.16b.jl",    # 2 point subdomain
  "../examples/ch07/WRM/Ex.7.17.jl",     # 1 point least-squares
  "../examples/ch07/WRM/Ex.7.17b.jl",    # 2 point least-squares
  "../examples/ch07/WRM/Ex.7.18.jl",     # 1 point galerkin
  "../examples/ch07/WRM/Ex.7.18b.jl",   # 2 point galerkin

  # Ch08 tests
  "ch08_ex01.jl",

  # Extra tests
  "vianello_symata.jl"                           # Vianello procedure

]

println("Running tests:")

@testset "NumericalMethodsforEngineers.jl" begin
  for my_test in code_tests
      println("\n  * $(my_test) *")
      include("$(my_test)")
  end
  println()
end


using Plots, LaTeXStrings, Statistics

function Deposition(;len, tot_time)
    surf = [0 for i=1:len]
    AvgList = [0.0 for i=1:tot_time]
    for n in 2:tot_time+1
        randsurf = rand(1:len)
        for i in randsurf
            index = FindLeast(surf, i, len)
            surf[index] += 1
        end
        AvgList[n-1] = mean(surf)
    end
    return AvgList
end

function sides(n, L)
    if n == L
        return n-1 , 1
    elseif n == 1
        return L , n+1
    else
        return n-1, n+1
    end
end

function FindLeast(surface, index_,L_surf)
    i1 , i2 = sides(index_, L_surf)
    lens = Dict(surface[i1]=>i1,
        surface[index_]=>index_,
        surface[i2]=>i2)
    minlen = min(surface[i1],surface[index_],surface[i2])
    return lens[minlen]
end

function Linear_fit(;Time, AvgList, tot_time)
    A = [hcat(Time[1:tot_time]) reshape(ones(tot_time), tot_time, 1)]
    b = reshape(AvgList[1:tot_time], tot_time, 1)
    line = (A \ b)
    return line
end

iternum = 1000
Parameters = Dict(:len => 300,
                    :tot_time => 50,)
allAvg = [ [0.0 for i in 1:Parameters[:tot_time]] for j = 1:iternum]
meanAvg = [0.0 for i in 1:Parameters[:tot_time]]
vars = [0.0 for i in 1:Parameters[:tot_time]]
for i in 1:iternum
    AvgList = Deposition(;Parameters...)
    allAvg[i] = AvgList
    meanAvg += AvgList
    print("\r$i")
end
meanAvg /= iternum
for i in 1:Parameters[:tot_time]
    vars[i] = std(hcat(allAvg...)[i,:])
end

Time = 0:Parameters[:tot_time]-1

Paraline = Dict(
                :Time => Time[1:end],
                :AvgList => meanAvg[1:end],
                :tot_time => 30)
Line = Linear_fit(;Paraline...)
X = Time[1]:Time[50]
Y = X .* Line[1] .+ Line[2]
Line[2]
plot(X,Y,c=:black,label = L"y = %$(round(Line[1],digits= 3))x + %$(round(Line[2],digits= 3))")
scatter!(Time[1:end], meanAvg[1:end],
    # xlims = (1, Parameters[:tot_time]),
    c = :steelblue,
    xlabel= L"Time",
    ylabel= L"H_{(t)}",
    title= L"~H_{(t)}-Time~\ (L = %$(Parameters[:len]))",
    label = L"Data\ point",
    yerror = vars)
    # legend = nothing)

savefig("../../Figs/Q1/H-t(L=$(Parameters[:len])).pdf")


# example from https://github.com/malmaud/TensorFlow.jl
using Distributions
using Printf

# Generate some synthetic data
x = randn(100, 50)
w = randn(50, 10)
y_prob = exp.(x*w)
y_prob ./= sum(y_prob,dims=2)

function draw(probs)
    y = zeros(size(probs))
    for i in 1:size(probs, 1)
        idx = rand(Categorical(probs[i, :]))
        y[i, idx] = 1
    end
    return y
end

y = draw(y_prob)

# Build the model
reset_default_graph()
sess = Session()

X = placeholder(Float64, shape=[nothing, 50])
Y_obs = placeholder(Float64, shape=[nothing, 10])

variable_scope("test_logisitic_model"; initializer=random_normal_initializer(0, .001), reuse=AUTO_REUSE) do
    global W = get_variable("W", shape=[50, 10], dtype=Float64)
    global B = get_variable("B", shape=[10], dtype=Float64)
end

Y=softmax(X*W + B)

Loss = -sum(log(Y).*Y_obs)
# adam = AdamOptimizer()
# minimize_op = minimize(adam, Loss)

# # Run training
# init(sess)
# for epoch in 1:1000
#     cur_loss, _ = run(sess, [Loss, minimize_op], feed_dict=Dict(X=>x, Y_obs=>y))
#     println(@sprintf("Current loss is %.2f.", cur_loss))
# end

function print_loss(loss_evaled)
    println(loss_evaled)
end

opt = ScipyOptimizerInterface(Loss, method="L-BFGS-B",options=Dict("maxiter"=> 100))
init(sess)
ScipyOptimizerMinimize(sess, opt, feed_dict=Dict(X=>x, Y_obs=>y), loss_callback=print_loss, fetches=[Loss])


using DifferentialEquations
using OrdinaryDiffEq
using Distributions
using MonteCarloMeasurements
using StatsBase
using Plots

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

δt = 1.0
tmax = 40.0
tspan = (0.0,tmax);

N = 1000.0;

n_samples = 1000; # Number of samples

p = [Particles(n_samples,Uniform(0.01,0.1)),
      Particles(n_samples,Uniform(5,20.0)),
      Particles(n_samples,Uniform(0.1,1.0))]

I₀=Particles(n_samples,Uniform(1.0,50.0))
u0 = [N-I₀,I₀,0.0]

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

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

s20 = sol_ode(20.0)

l = @layout [a b c]
binwidth = 50
pl1 = histogram(s20[1],bins=0:binwidth:N, title="S(20)", xlabel="S", ylabel="Frequency", color=:blue)
pl2 = histogram(s20[2],bins=0:binwidth:N, title="I(20)", xlabel="I", ylabel="Frequency", color=:red)
pl3 = histogram(s20[3],bins=0:binwidth:N, title="R(20)", xlabel="R", ylabel="Frequency", color=:green)
plot(pl1,pl2,pl3,layout=l,legend=false)

corkendall(hcat(Array(p),Array(I₀)),Array(s20))


export Cos

"""
`Cos([domainType=Float64::Type,] dim_in::Tuple)`

Creates a cosine non-linear operator with input dimensions `dim_in`:
```math
\\cos (\\mathbf{x} ).
```

"""
struct Cos{T,N} <: NonLinearOperator
	dim::NTuple{N,Int}
end

function Cos(DomainType::Type, DomainDim::NTuple{N,Int}) where {N}
	Cos{DomainType,N}(DomainDim)
end

Cos(DomainDim::NTuple{N,Int}) where {N} = Cos{Float64,N}(DomainDim)
Cos(DomainDim::Vararg{Int}) = Cos{Float64,length(DomainDim)}(DomainDim)

function mul!(y::AbstractArray{T,N}, L::Cos{T,N}, x::AbstractArray{T,N}) where {T,N}
	y .= cos.(x)
end

function mul!(y::AbstractArray,
              J::AdjointOperator{Jacobian{A,TT}},
              b::AbstractArray) where {T,N, A<: Cos{T,N}, TT <: AbstractArray{T,N}}
    L = J.A
    y .= -conj.(sin.(L.x)).*b
end

fun_name(L::Cos) = "cos"

size(L::Cos) = (L.dim, L.dim)

domainType(L::Cos{T,N}) where {T,N} = T
codomainType(L::Cos{T,N}) where {T,N} = T


@test factorial(7,3) == 7*6*5*4
@test_throws DomainError factorial(3,7)
@test_throws DomainError factorial(-3,-7)
@test_throws DomainError factorial(-7,-3)
#JuliaLang/julia#9943
@test factorial(big(100), (80)) == 1303995018204712451095685346159820800000
#JuliaLang/julia#9950
@test_throws OverflowError factorial(1000,80)

# derangement
@test derangement(4) == subfactorial(4) == 9
@test derangement(24) == parse(BigInt,"228250211305338670494289")

# partialderangement
@test partialderangement(7, 3) == 315
@test_throws DomainError partialderangement(8, 9)
@test_throws DomainError partialderangement(-8, 0)

# doublefactorial
@test doublefactorial(70) == parse(BigInt,"355044260642859198243475901411974413130137600000000")
@test_throws DomainError doublefactorial(-1)

# hyperfactorial
@test hyperfactorial(8) == parse(BigInt,"55696437941726556979200000")
@test hyperfactorial(0) == parse(BigInt,"1")
@test hyperfactorial(1) == parse(BigInt,"1")
@test hyperfactorial(2) == parse(BigInt,"4")

# multifactorial
@test multifactorial(40,2) == doublefactorial(40)
@test_throws DomainError multifactorial(-1,1)

# multinomial
@test multinomial(1,4,4,2) == 34650

# primorial
@test primorial(17) == 510510
@test_throws DomainError primorial(-1)


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

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

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

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

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

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

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

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

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


using Polynomials

function mypolyval(p::Poly,x::Number)
    s=p[0]
    t=one(x)
    for i=1:length(p)-1
        t*=x
        s+=p[i]*t
    end
    s
end

function myhorner(p::Poly,x::Number)
    s=p[end]
    for i=length(p)-2:-1:0
        s=s*x+p[i]
    end
    s
end

n=10001
pbig=Poly(rand(n))
x=0.12345;

src/Module A - Short Julia Course/myfunctions.jl


export AbLoss, MSE, MAD, LOGIT, altloss

abstract type AbLoss end
# Notation
#       y = Observed target
#       z = fit Σ_j ϕ(X, ω_j)
# Methods
#       Loss evaluation
#       Generalized Residual
valnderv(l::AbLoss, y, z) = l(y, z), derv(l, y, z)
altloss(l::AbLoss) =  (_...) -> NaN

################################################################################
type MSELoss <: AbLoss
end

const MSE = MSELoss()

(l::MSELoss)(y, z) = sum(abs2, z-y)/2
derv(::MSELoss, y, z) = z-y
derv2(::MSELoss, y, z) = 1.
function valnderv(::MSELoss, y, z)
  diff = z-y
  sum(abs2, diff)/2, diff
end

################################################################################
type MADLoss <: AbLoss
end

const MAD = MADLoss()

(l::MADLoss)(y, z) =  sum(abs, z-y)
derv(::MADLoss, y, z) = sign.(z-y)
derv2(::MADLoss, y, z) = 0.
function valnderv(::MADLoss, y, z)
    diff = z-y
    sum(abs, diff), sign.(diff)
end

################################################################################
type LogitLoss <: AbLoss
end

const LOGIT = LogitLoss()

log1plusexp(x) = (x<0.) ? log1p(exp(x)) : x+log1p(exp(-x))
sigmoid(x) = (x<0.) ? exp(x)/(1.+exp(x)) : 1./(1.+exp(-x))
(l::LogitLoss)(y, z) =  sum(log1plusexp.(z)) - y'*z
derv(::LogitLoss, y, z) = sigmoid.(z) - y
derv2(::LogitLoss, y, z) = (s=sigmoid.(z); (1-s).*s)
"""
bits(sigmoid(-64)) = "0011101000101001011010011101010001110011001000011110010011001100"
bits(1-(1-sigmoid(-64))) = "0000000000000000000000000000000000000000000000000000000000000000"
"""

misclassification(y, z, threshold=.5) = sum(y .!= (sigmoid.(z) .≥ threshold))
altloss(::LogitLoss) =  misclassification
################################################################################


# some kernels
epanechnikov(x::Tp) where {Tp<:AbstractFloat} = Tp( (abs(x) ≤  1.0) ? 0.75 * (1.0-x*x) : 0.0 )
gaussian(x::Tp) where {Tp<:AbstractFloat} = Tp( exp( -0.5 * x * x) * 0.3989422804014327 )
triweight(x::Tp) where {Tp<:AbstractFloat} = Tp( (abs(x) ≤ 1.0) ? 1.09375 * (1.0-x*x)^3  : 0.0 )
uniform(x::Tp) where {Tp<:AbstractFloat} = Tp( (abs(x) ≤ 1.0) ? 0.5 : 0.0 )
cosinus(x::Tp) where {Tp<:AbstractFloat} = Tp( (abs(x) ≤ 1.0) ?   0.7853981633974483    * cos( 1.5707963267948966 * x ) : 0.0 )
quartic(x::Tp) where {Tp<:AbstractFloat} = Tp( (abs(x) ≤ 1.0) ? 0.9375 * (1.0 - x*x)^2 : 0.0 )


module TestGaussNewton

using Test
using OptimizationAlgorithms
using OptimizationAlgorithms: GaussNewton, LevenbergMarquart, DecreasingStep, LevenbergMarquartSettings
using OptimizationAlgorithms: update!, optimize!

function test_problem(n)
    x = sort!(randn(n))
    g(a, b) = x -> sin(2π*a*x) + b*x^2
    a, b = randn(2)
    y = g(a, b).(x)
    return [a,b], x, y, g, function residual!(r, ab)
        @. r = y - g(ab...).(x)
    end
end

@testset "GaussNewton" begin
    n = 64
    ab0, x, y, g, r = test_problem(n)
    G = GaussNewton(r, ab0, y)
    σ = .1
    ab = ab0 + σ*randn(2)
    for i in 1:16
        update!(G, ab)
    end
    tol = 1e-8
    @test isapprox(ab, ab0, atol = tol)

    D = DecreasingStep(G, ab)
    ab = ab0 + σ*randn(2)
    for i in 1:16
        update!(D, ab)
    end
    tol = 1e-8
    @test isapprox(ab, ab0, atol = tol)
end

@testset "LevenbergMarquart" begin
    n = 64
    ab0, x, y, g, r = test_problem(n)
    LM = LevenbergMarquart(r, ab0, y)
    σ = .1
    ab = ab0 + σ*randn(2)
    tol = 1e-8
    # tesing out intitialization with very large lambda, min_decrease, and min_iter
    res = copy(y)
    stn = LevenbergMarquartSettings(max_iter = 16, min_decrease = 1e-10, min_resnorm = 1e-10, min_res = 0)
    optimize!(LM, ab, res, Val(false), stn, 1e-6, Val(true))
    @test isapprox(ab, ab0, atol = tol)
end

@testset "Geodesic LevenbergMarquart" begin
    n = 64
    ab0, x, y, g, r = test_problem(n)
    LM = LevenbergMarquart(r, ab0, y)
    σ = .1
    ab = ab0 + σ*randn(2)
    tol = 1e-8
    # tesing out intitialization with very large lambda, min_decrease, and min_iter
    res = copy(y)
    stn = LevenbergMarquartSettings(max_iter = 16, min_decrease = 1e-10, min_resnorm = 1e-10, min_res = 0)
    optimize!(LM, ab, res, Val(true), stn, 1e-6, 0.01, 0.75, Val(true))
    @test isapprox(ab, ab0, atol = tol)
end

end