#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 2, Zadanie 1.
#

# assert ARGS
err_mess = "provide two arguments: [64, 32] [up, down, max_to_min, min_to_max]"
if length(ARGS) != 2
    println(err_mess)
    exit(1)
end

# name the arguments
fl_str = ARGS[1]
mode = ARGS[2]

# parse the arguments
if !in(fl_str, ["64", "32"]) || !in(mode, ["up", "down", "max_to_min", "min_to_max"])
    println(err_mess)
    exit(1)
end

fl = Float64

# apply the first argument
if fl_str== "32"
    global fl = Float32
end
# end assert

# x₄ and x₅ were altered
x = [2.718281828, -3.141592654, 1.414213562, 0.577215664, 0.301029995]
y = [1486.2497, 878366.9879, -22.37492, 4773714.647, 0.000185049]

function calc_scalar_product(x, y)
    S = fl(0)
    # (a), (b)
    if in(mode, ["up", "down"])
        # change the direction of the range depending on the provided program argument
        for i in (mode == "down" ? (5:-1:1) : (1:5))
            S = fl(S + fl(x[i] * y[i]))
        end
    # (c), (d)
    elseif in(mode, ["max_to_min", "min_to_max"])
        # save partial products
        mid_positive::Array{fl} = []
        mid_negative::Array{fl} = []
        # iterate over the vectors
        for i in (1:5)
            out = fl(x[i] * y[i])
            if out < fl(0)
                push!(mid_negative, out)
            else
                push!(mid_positive, out)
            end
        end

        # sort the partial products
        sort!(mid_negative, rev=(mode == "max_to_min"))
        sort!(mid_positive, rev=(mode == "max_to_min"))

        # add the partial products
        S_positive = fl(0)
        S_negative = fl(0)
        for i in (1:length(mid_negative))
            S_negative = fl(S_negative + mid_negative[i])
        end
        for i in (1:length(mid_positive))
            S_positive = fl(S_positive + mid_positive[i])
        end
        # finally, add the partial sums
        S = fl(S_positive + S_negative)
    end
    return S
end

# print out the results
println(calc_scalar_product(x,y))


using NewtonsMethod
using Test

@testset "NewtonsMethod.jl" begin
    # Write your tests here.
    f_1(x) = x^2 - 3x + 2
    f_2(x) = sin(x)
    f_3(x) = log(x)
    f_4(x) = 1/x - 4
    f_5(x) = x+3
    f_6(x) = 3
    f_7(x) = 0
    f_8(x) = exp(x) - 5
    f_prime_1(x) = 2x - 3
    f_prime_2(x) = cos(x)
    f_prime_3(x) = 1/x
    f_prime_4(x) = -1/(x^2)
    f_prime_5(x) = 1
    f_prime_6(x) = 0
    f_prime_7(x) = 0
    f_prime_8(x) = exp(x)
    @test newtonroot(f_1, f_prime_1; x_0 = 0) ≈ 1
    @test newtonroot(f_2, f_prime_2; x_0 = 3) ≈ π
    @test newtonroot(f_3, f_prime_3; x_0 = 0.7) ≈ 1
    @test newtonroot(f_4, f_prime_4; x_0 = 0.4) ≈ 0.25
    @test newtonroot(f_5, f_prime_5; x_0 = 0) ≈ -3
    @test newtonroot(f_6, f_prime_6; x_0 = 1) === nothing
    @test newtonroot(f_7, f_prime_7; x_0 = 4) == 4
    @test newtonroot(f_8, f_prime_8; x_0 = 3) ≈ log(5)
    @test newtonroot(f_1; x_0 = 0) ≈ 1
    @test newtonroot(f_2; x_0 = 3) ≈ π
    @test newtonroot(f_3; x_0 = 0.7) ≈ 1
    @test newtonroot(f_4; x_0 = 0.4) ≈ 0.25
    @test newtonroot(f_5; x_0 = 0) ≈ -3
    @test newtonroot(f_6; x_0 = 1) === nothing
    @test newtonroot(f_7; x_0 = 4) == 4
    @test newtonroot(f_8; x_0 = 3) ≈ log(5)
end

@testset "non convergence" begin
    f(x) = x^2+2
    f_prime(x) = 2x
    @test newtonroot(f, f_prime; x_0 = 0) === nothing
    @test newtonroot(f, f_prime; x_0 = 0.5) === nothing
    @test newtonroot(f;x_0 = 0) === nothing
    @test newtonroot(f;x_0 = 0.5) === nothing
end

@testset "Big Float" begin
    f(x) = x^2 - 3x + 2
    f_prime(x) = 2x - 3
    x_0 = parse(BigFloat, "0.34")
    @test convert(Float64, newtonroot(f, f_prime; x_0 = x_0)) ≈ 1
    @test convert(Float64, newtonroot(f; x_0 = x_0)) ≈ 1
end

@testset "max iter" begin
    f(x) = exp(x) - 5
    f_prime(x) = exp(x)
    @test newtonroot(f, f_prime; x_0 = 0, maxiter = 5) === nothing
    @test newtonroot(f; x_0 = 0, maxiter = 5) === nothing
end

@testset "tolerance" begin
    f(x) = 3
    f_prime(x) = 3
    @test newtonroot(f, f_prime; x_0 = 4, tol = 3) == 4
    @test newtonroot(f; x_0 = 4, tol = 3) == 4
end


using StatsBase
using RDatasets
a = [123,4234,23423,1231231,1432432423,1341413413]
var(a)
var(a, 2)
std(a)
mean_and_var(a)
mean_and_std(a)
skewness(a)
kurtosis(a)
moment(a,3)
span(a)
variation(a)


using Optim

f(x) = -sin(x[1]^2/2 - x[2]^2/4 + 3) * cos(2x[1] + 1 - exp(x[2])) # Same as last section
res = optimize(f, [0, 0.5])

function df(x)
    # Same as last section
    a1 = x[1]^2/2 - x[2]^2/4 + 3
    a2 = 2x[1] + 1 - exp(x[2])
    b1 = cos(a1)*cos(a2)
    b2 = sin(a1)*sin(a2)
    return -[x[1]*b1 - 2b2, -x[2]/2*b1 + exp(x[2])*b2]
end
res = optimize(f, df, [0, 0.5]; inplace=false)

res = optimize(f, df, [0, 0.5], GradientDescent(); inplace=false)

res = optimize(f, [0, 0.5], GradientDescent(); autodiff=:forward)

res = optimize(f, [0, 0.5], Newton(); autodiff=:forward)

res = optimize(f, [0, 0.5], BFGS(); autodiff=:forward)


module XiCor

using Random

export xicor

"""
    xicor(X, Y[, break_ties_randomly=false[, rng=nothing]])

Computes the correlation ξ between X and Y.

Unlike most coefficients of correlation, ξ ranges from -0.5 to 1.

If there are duplicate values in X, then ties are
broken based on the order in which they are observed.
If the order of X is not random, then you should
set `break_ties_randomly` to `true` to avoid a biased
estimate. You can use the `rng` parameter to
deterministically break ties.

See _A new coefficient of correlation_ by Chatterjee.

[arXiv:1909.10140 [math.ST]](https://arxiv.org/abs/1909.10140)

# Examples
```julia-repl
julia> ξ = xicor(1:100, 1:100)
0.9702970297029703
julia> x = trunc.((1:100) ./ 10);  # create x with lots of duplicates
julia> y = rand(MersenneTwister(0), 100);
julia> ξ = xicor(x, y, true, MersenneTwister(0))
-0.01830183018301823
julia> ξ = xicor(x, y, true, MersenneTwister(42))
0.004500450045004545
```
"""
function xicor(X, Y, break_ties_randomly=false, rng=nothing)
    if break_ties_randomly
        if !isnothing(rng)
            index = randperm(rng, length(X))
        else
            index = randperm(length(X))
        end
        X = X[index]
        Y = Y[index]
    end
    n = length(X)
    Y = Y[sortperm(X)]  # how should offset arrays be handled?
    sorter = sortperm(Y)
    R = zeros(Int, n)  # R[i] is the number of j such that Y[j] ≤ Y[i]
    L = zeros(Int, n)  # L[i] is the number of j such that Y[j] ≥ Y[i]

    i = 1
    while i <= n
        curr = Y[sorter[i]]
        counter = 1
        i += 1  # look ahead for repeated Y values
        while i <= n && Y[sorter[i]] == curr
            counter += 1
            i += 1
        end
        i -= 1
        for j = i-counter+1:i  # fill in R and L values wherever Y == curr
            R[sorter[j]] = i
            L[sorter[j]] = (n-i) + counter
        end
        i += 1
    end

    1 - n * sum(abs.(diff(R))) / (2 * sum(L .* (n .- L)))
end

end # module


# Jonas August 1/2022

#=
This problem was asked by Google.

Given an array of integers where every integer occurs three times except for one integer, which only occurs once, find and return the non-duplicated integer.

For example, given [6, 1, 3, 3, 3, 6, 6], return 1. Given [13, 19, 13, 13], return 19.

Do this in O(N) time and O(1) space.
=#

#=
Assume that the array is correctly formed (one single among triples; no doubles, quadruples, etc.; only integers).

There are only at most (n-1)/3 + 1 unique numbers.

n = 3*m+1 is odd if m even and vice versa.

=#

#=
If y is the single entry, then the sum of the remaining entries is divisible by three.  Let x\y denote set x with y removed.

y is unique in in x => (sum of z s.t. z in x\y) mod 3 = 0

This is necessary condition for identifying y, but is it sufficient?  No!  Consider x=[13,19,13,13].  Try y=13.  Then sum of x\y is 19+13+13 = 45 which is divisible by 3.  Argh!

But maybe there's some function of the elements that can be used instead?

If a,b are integers, then cube root of aab is integer iff b = a?

Relevance Fermat's last theorem?
x^k + y^k = z^k has no positive integer solutions x,y,z for integer k>2.
Try k = 3 and cube entries?

=#

"Find that integer in input array x that has no match, while all other elements come in three, in no particular order.  Takes O(n) time, O(1) space, where n = length of the array."
function single_among_triples_fast(x::Vector{Int})::Int
    # TODO: allow zeros and negatives
    if any(x .== 0) @error("No zeroes allowed") end
    if any(x .< 0) @error("No negatives allowed") end
    p = 1
    for y in x
        p *= y
    end
    for y in x
        # TODO: CHECK that this always works
        c = (p/y)^(1/3)
        # TODO: what tolerance ensures correct results?
        if abs(c-round(c)) < sqrt(eps(c))    # is c approximately an integer?
            return y
        end
    end
    @error("This line should never be reached for well-formed input.")
end

"Find that integer in input array x that has no match, while all other elements come in three, in no particular order.  Takes O(n^2) time, O(1) space, where n = length of the array."
function single_among_triples_brute_force(x::Vector{Int})::Int
    for i in 1:length(x)
       matched = false
       for j in i+1:length(x)
           if x[j] == x[i]
	       matched = true
	       break
	   end
       end
       if !matched
           return x[i]
       end
    end
    @error("This line should never be reached for well-formed input.")
end

#=
Other False starts:

Not O(1) space:

- Convert to linked list.  Pick first element and scan for match and delete matches if found, then start again.  Must end up with unmatched number at beginning.

x[i[1]] is unique in x[:]
x[i[2]] == x[i[3]] == x[i[4]]
x[i[5]] == x[i[6]] == x[i[7]]
...

x[i[1]] + sum(j = 2:3:n) 3*x[i[j]]

Not O(n) time:

- Sort.  Triples will cluster, so linearly scan for the single.

=#

# TODO: Construct random inputs

using Test

for single_among_triples in [single_among_triples_brute_force, single_among_triples_fast]
  # @test @test single_among_triples([]) == ?  not defined? exception?
  @test single_among_triples([19]) == 19
  @test single_among_triples([13, 19, 13, 13]) == 19
  @test single_among_triples([6, 1, 3, 3, 3, 6, 6]) == 1
end


using Multicombinations
using Base.Test

# multichoose

@test multichoose(3,2) == 6
@test multichoose(3,0) == 1
@test multichoose(0,3) == 0
@test multichoose(0,0) == 1
@test multichoose(-1,-1) == 0

# multicombinations

@test collect(multicombinations("abc",3)) == ["aaa","aab","aac","abb","abc","acc","bbb","bbc","bcc","ccc"]
@test collect(multicombinations("åäö",3)) == ["ååå","ååä","ååö","åää","åäö","åöö","äää","ääö","äöö","ööö"]
@test collect(multicombinations("abc",0)) == collect(combinations("abc",0))
@test collect(multicombinations("",1)) == collect(combinations("",1))
@test collect(multicombinations("",0)) == collect(combinations("",0))
@test collect(multicombinations("",-1)) == collect(combinations("",-1))

# integersums

@test collect(integersums(3,3)) == {[3,0,0],[2,1,0],[2,0,1],[1,2,0],[1,1,1],[1,0,2],[0,3,0],[0,2,1],[0,1,2],[0,0,3]}
@test collect(integersums(3,0)) == {[0,0,0]}
@test collect(integersums(0,3)) == {}
@test collect(integersums(0,0)) == {[]}
@test collect(integersums(-1,-1)) == {}

# is2mc

@test map(is2mc,integersums(3,3)) == collect(multicombinations(1:3,3))
@test map(is2mc,integersums(3,0)) == collect(multicombinations(1:3,0))
@test map(is2mc,integersums(0,3)) == collect(multicombinations([],3))
@test map(is2mc,integersums(0,0)) == collect(multicombinations([],0))

# length

@test length(collect(multicombinations("abc",3))) == length(multicombinations("abc",3))
@test length(collect(multicombinations("åäö",3))) == length(multicombinations("åäö",3))
@test length(collect(integersums(3,3))) == length(integersums(3,3))


######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    kalmanFilter(data, ssmodel)

Description:
Apply Kalman filter to observed data.
Measurement Equation:
    y_{t} = H_{t} β_{t} + A z_{t} + e_{t} .
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t};
    e_{t} ~ i.i.d.N(0,R);
    v_{t} ~ i.i.d.N(0,Q);
    z_{t} ~ i.i.d.N(0,Z);
    E(e_t v_s') = 0.

Inputs:
- data      = observed data
- H         = measurement eq. state coef. matrix
- A         = measurement eq. exogenous coef. matrix
- F         = state eq. companion matrix
- μ         = state eq. intercept term
- R         = covariance matrix on measurement disturbance
- Q         = covariance matrix on state disturbance
- Z         = covariance matrix on predetermined var vector
"""
function kalmanFilter(data_y, ssmodel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Save number of observations
    num_obs = size(data_y)[1]

    # Empty filtered data matrices
    data_filtered_y = similar(data_y)
    data_filtered_β = zeros(num_obs, size(Q)[1])

    # Create empty lists for P_{t}, P_{t|t-1}
    Ptt = Any[]
    Pttlag = Any[]

    # Initialize β_pred and P_pred
    β_pred_laglag = inv(I - F) * μ
    P_pred_laglag = ones(size(Q)[1], size(Q)[1])

    for t = 1:num_obs
        # Save current obs of y
        y = data_y[t, :]

        # Prediction
        β_pred_lag = μ + F * β_pred_laglag
        P_pred_lag = F * P_pred_laglag * transpose(F) + Q
        y_pred_lag = H * β_pred_lag
        η_pred_lag = y - y_pred_lag
        f_pred_lag = H * P_pred_lag * transpose(H) + R

        # Save P_{t|t-1}
        push!(Pttlag, P_pred_lag)

        # Updating
        K = P_pred_lag * transpose(H) * pinv(f_pred_lag)
        β_pred = β_pred_lag + K * η_pred_lag
        P_pred = P_pred_lag - K * H * P_pred_lag

        # Save P_{t|t}
        push!(Ptt, P_pred)

        # Save data
        data_filtered_y[t, :] = y_pred_lag
        data_filtered_β[t, :] = β_pred

        # Lag the predictions
        β_pred_laglag = β_pred
        P_pred_laglag = P_pred
    end

    # Returned filtered series
    # for obs variable and state
    return data_filtered_y, data_filtered_β, Pttlag, Ptt
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    kalmanSmoother(data, ssmodel)

Description:
Apply Kalman smoother to observed data.
Measurement Equation:
    y_{t} = H_{t} β_{t} + A z_{t} + e_{t}.
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t};
    e_{t} ~ i.i.d.N(0,R);
    v_{t} ~ i.i.d.N(0,Q);
    z_{t} ~ i.i.d.N(0,Z);
    E(e_t v_s') = 0.

Inputs:
- data      = observed data
- H         = measurement eq. state coef. matrix
- A         = measurement eq. exogenous coef. matrix
- F         = state eq. companion matrix
- μ         = state eq. intercept term
- R         = covariance matrix on measurement disturbance
- Q         = covariance matrix on state disturbance
- Z         = covariance matrix on predetermined var vector
"""
function kalmanSmoother(data_y, ssmodel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Save number of observations
    num_obs = size(data_y)[1]

    # Empty filtered data matrices
    data_smoothed_y = similar(data_y)
    data_smoothed_β = zeros(num_obs, size(Q)[1])

    # Create empty list for P_{t|T}
    PtT = Any[]

    # Run Kalman filter
    data_filtered_y, data_filtered_β, Pttlag, Ptt = kalmanFilter(data_y, ssmodel)

    # Initialize β_{t+1|T} (β_{T|T})
    βtflagT = data_filtered_β[end]
    data_smoothed_β[end] = βtflagT

    # Initialize P_{t+1|T} (P_{T|T})
    Ptflag_T = Ptt[end]
    push!(PtT, Ptflag_T)

    # Initialize y_{t|T} (y_{T|T})
    data_smoothed_y[end] = data_filtered_y[end]

    # Run Kalman smoother
    for i = 1:(num_obs-1)

        # Retrieve β_{t|t}
        βtt = data_filtered_β[end-i]

        # Compute β_{t|T} using β_{t+1|T}, β_{t|t}, P_{t|t}, and P_{t+1|t}
        βtT = βtt +
              Ptt[end-i] * transpose(F) * inv(Pttlag[end-i+1]) *
              (βtflagT - F * βtt - μ)

        # Store β_{t+1|T} in smoothed data
        data_smoothed_β[end-i] = βtT

        # Set β_{t|T} as new β_{t+1|T} for next iteration
        βtflagT = βtT

        # Compute P_{t|T} using P_{t|t}, P_{t+1|t}, and P_{t+1|T}
        Pt_T = Ptt[end-i] +
               Ptt[end-i] * transpose(F) * inv(Pttlag[end-i+1]) *
               (Ptflag_T - Pttlag[end-i+1]) *
               transpose(Ptt[end-i] * transpose(F) * inv(Pttlag[end-i+1]))

        # Store P_{t|T}
        push!(PtT, Pt_T)

        # Set P_{t|T} as new P_{t+1|T} for next iteration
        Ptflag_T = Pt_T

        # Generate y_{t|T} (smoothed obs.)
        ytT = H * βtT

        # Store smoothed obs.
        data_smoothed_y[end-i] = ytT
    end

    # Flip P_{t|T} list
    PtT = reverse(PtT)

    # Returned filtered series
    # for obs variable and state
    return data_smoothed_y, data_smoothed_β, PtT
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    KNFactorSampler(data_y, ssmodel)

Description:
Draw a sample series of dynamic factor from conditional distribution in Ch 8, Kim & Nelson (1999).
Measurement Equation:
    y_{t} = H_{t} β_{t} + A z_{t} + e_{t}.
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t};
    e_{t} ~ i.i.d.N(0,R);
    v_{t} ~ i.i.d.N(0,Q);
    z_{t} ~ i.i.d.N(0,Z);
    E(e_t v_s') = 0.

Inputs:
- data      = observed data
- H         = measurement eq. state coef. matrix
- A         = measurement eq. exogenous coef. matrix
- F         = state eq. companion matrix
- μ         = state eq. intercept term
- R         = covariance matrix on measurement disturbance
- Q         = covariance matrix on state disturbance
- Z         = covariance matrix on predetermined var vector
"""
function KNFactorSampler(data_y, ssmodel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Run Kalman filter
    data_filtered_y, data_filtered_β, Pttlag, Ptt = kalmanFilter(data_y, ssmodel)

    # Create placeholders for factor distr.
    # mean vector and covariance matrix for all t
    β_t_mean = Any[]
    β_t_var = Any[]

    # Record number of time periods
    T = size(data_y)[1]

    # Create empty vector for factor realizations
    β_realized = similar(data_filtered_β)

    # Initialize β_realized
    push!(β_t_mean, data_filtered_β[T, :])
    push!(β_t_var, Ptt[T])
    β_realized[T, :] = sim_MvNormal_alt(β_t_mean[1], β_t_var[1])

    # Generate `β_t_mean` and `β_t_var`
    # for all time periods
    if isposdef(Q) == false
        ## IF Q IS SINGULAR

        # Determine number of rows
        # of state factor used
        num_use_rows = 0
        for i = 1:size(Q)[1]
            if Q[i, i] != 0
                num_use_rows += 1
            end
        end

        # Create modified F and Q matrices
        F_star = F[1:num_use_rows, :]
        Q_star = Q[1:num_use_rows, 1:num_use_rows]

        # Iteratively generate conditional draws
        # of state vector
        for j = 1:(T-1)

            # β_{t|t,β*_{t+1}}
            β_t_mean_temp = data_filtered_β[T-j, :]
            +Ptt[T-j] * transpose(F_star) * inv(F_star * Ptt[T-j] * transpose(F_star) + Q_star) * (β_realized[T+1-j, :][1:num_use_rows] - μ[1:num_use_rows] - F_star * data_filtered_β[T-j, :])
            push!(β_t_mean, β_t_mean_temp)

            # P_{t|t,β*_{t+1}}
            β_t_var_temp = Ptt[T-j] - Ptt[T-j] * transpose(F_star) * inv(F_star * Ptt[T-j] * transpose(F_star) + Q_star) * F_star * Ptt[T-j]
            push!(β_t_var, β_t_var_temp)

            # Draw new β_t
            β_realized[T-j, :] = sim_MvNormal_alt(β_t_mean[1+j], β_t_var[1+j])
        end
    else
        ## IF Q IS NOT SINGULAR (redundant, but potentially faster)
        for j = 1:(T-1)

            # β_{t|t,β_{t+1}}
            β_t_mean_temp = data_filtered_β[T-j, :]
            +Ptt[T-j] * transpose(F) * inv(F * Ptt[T-j] * transpose(F) + Q) * (β_realized[T+1-j, :] - μ - F * data_filtered_β[T-j, :])
            push!(β_t_mean, β_t_mean_temp)

            # P_{t|t,β_{t+1}}
            β_t_var_temp = Ptt[T-j] - Ptt[T-j] * transpose(F) * inv(F * Ptt[T-j] * transpose(F) + Q) * F * Ptt[T-j]
            push!(β_t_var, β_t_var_temp)

            # Draw new β_t
            β_realized[T-j, :] = sim_MvNormal_alt(β_t_mean[1+j], β_t_var[1+j])
        end
    end

    # Return sampled factor series
    # fot t = 1,...,T
    return β_realized
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    linearRegressionSampler(Y,X)

Description:
Estimate β and σ^2 in Y = Xβ + e, e ~ N(0,σ^2 I_T).
Generate samples of β and σ^2.

Inputs:
- Y     = Dependent data matrix
- X     = Independent data matrix
"""
function linearRegressionSampler(Y, X)

    # Save number of obs
    T = size(X)[1]

    # Initialize σ2
    σ2 = 1.0

    ##################################
    ##################################
    # Generate new β

    ## Prior parameters in N(β0,Σ0)
    β0 = zeros(size(X)[2])
    Σ0 = Matrix(I, size(β0)[1], size(β0)[1]) .* 1000.0

    ## Posterior parameters in N(β1,Σ1)
    β1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X) * (inv(Σ0) * β0 + inv(σ2) * transpose(X) * Y)
    β1 = vec(β1)
    Σ1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X)
    Σ1 = Hermitian(Σ1)

    ## Generate new β
    β = sim_MvNormal(β1, Σ1)

    ##################################
    ##################################
    # Update σ2

    ## Prior parameters in IG(ν0/2, δ0/2)
    ν0 = 0.002
    δ0 = 0.002

    ## Posterior parameters in IG(ν1/2, δ1/2)
    ν1 = ν0 + T
    δ1 = δ0 + norm(Y - X * β)^2

    ## Generate new σ2
    σ2 = rand(InverseGamma(ν1 / 2, δ1 / 2))

    ##################################
    ##################################
    # Return parameters
    return β, σ2
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    linearRegressionSamplerRestrictedVariance(Y, X, σ2)

Description:
Estimate β and σ^2 in Y = Xβ + e, e ~ N(0,σ^2 I_T),
where σ^2 is restricted to some chosen value.
Generate samples of β and σ^2.

Inputs:
- Y     = Dependent data matrix
- X     = Independent data matrix
- σ2    = Restricted error variance
"""
function linearRegressionSamplerRestrictedVariance(Y, X, σ2)

    # Save number of obs
    T = size(X)[1]

    ##################################
    ##################################
    # Generate new β

    ## Prior parameters in N(β0,Σ0)
    β0 = zeros(size(X)[2])
    Σ0 = Matrix(I, size(β0)[1], size(β0)[1]) .* 1000.0

    ## Posterior parameters in N(β1,Σ1)
    β1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X) * (inv(Σ0) * β0 + inv(σ2) * transpose(X) * Y)
    β1 = vec(β1)
    Σ1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X)
    Σ1 = Hermitian(Σ1)

    ## Generate new β
    β = sim_MvNormal(β1, Σ1)

    ##################################
    ##################################
    # Return parameters
    return β
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    autocorrErrorLinearRegressionSampler(Y, X, error_lag_num)

Description:
Estimate β, σ^2, and ϕ in Y = Xβ + e, e = Eϕ + ν, ν_t ~ i.i.d.N(0,σ^2).
Generate samples of β, σ^2, and ϕ.
Procedure descried in Section 7.4.2 in Kim & Nelson.

Inputs:
- Y             = Dependent data matrix
- X             = Independent data matrix
- error_lag_num = Number of lags in the disturbance DGP
"""

function autocorrErrorLinearRegressionSampler(Y, X, ϕold, σ2old, error_lag_num)

    # Save number of obs
    T = size(X)[1]

    # Initialize σ2
    ϕ = ϕold
    σ2 = σ2old

    ##################################
    ##################################
    # Generate β

    ## Prior parameters in N(b0,A0)
    b0 = zeros(size(X)[2])
    A0 = Matrix(I, size(b0)[1], size(b0)[1]) .* 1000.0

    ## Generate X^⋆
    X_star = similar(X[(1+length(ϕ)):end, :])
    for i = 1:size(X)[2] # iterate over variables in X
        x_temp = X[:, i]
        for p = 1:length(ϕ) # iterate over lag params in ϕ
            x_temp[(1+length(ϕ)):end, :] = x_temp[(1+length(ϕ)):end, :] - ϕ[p] .* lag(x_temp, p)[(1+length(ϕ)):end, :]
        end
        x_temp = x_temp[(1+length(ϕ)):end, :]
        X_star[:, i] = x_temp
    end

    ## Generate Y^⋆
    Y_star = similar(Y[(1+length(ϕ)):end, :])
    y_temp = Y
    for p = 1:length(ϕ)
        y_temp = Y
        y_temp = y_temp[(1+length(ϕ)):end, :] - ϕ[p] .* lag(y_temp, p)[(1+length(ϕ)):end, :]
    end
    Y_star = y_temp

    ## Posterior parameters in N(b1,A1)
    b1 = inv(inv(A0) + inv(σ2) * transpose(X_star) * X_star) * (inv(A0) * b0 + inv(σ2) * transpose(X_star) * Y_star)
    b1 = vec(b1)
    A1 = inv(inv(A0) + inv(σ2) * transpose(X_star) * X_star)
    A1 = Hermitian(A1)

    ## Draw new β
    β = sim_MvNormal(b1, A1)

    ##################################
    ##################################
    # Generate ϕ

    ## Prior parameters in N(c0,B0)
    c0 = zeros(size(ϕ)[1])
    B0 = Matrix(I, size(c0)[1], size(c0)[1]) .* 1000.0

    ## Generate e^⋆
    e_star = Y - X * β

    ## Generate E^⋆
    E_star = zeros(T, length(ϕ))
    E_star = E_star[(1+length(ϕ)):end, :]
    for i = 1:(size(E_star)[2]) # iterate over variables in X
        e_temp = lag(e_star, i)
        E_star[:, i] = e_temp[(1+length(ϕ)):end]
    end
    e_star = e_star[(1+length(ϕ)):end]

    ## Posterior parameters in N(c1,B1)
    c1 = inv(inv(B0) + inv(σ2) * transpose(E_star) * E_star) * (inv(B0) * c0 + inv(σ2) * transpose(E_star) * e_star)
    c1 = vec(c1)
    B1 = inv(inv(B0) + inv(σ2) * transpose(E_star) * E_star)
    B1 = Hermitian(B1)

    ## Generate new ϕ
    ###=
    ind = 0
    accept = 0
    ϕ = similar(ϕold)
    while accept == 0

        ind += 1

        ## Draw ψ
        ϕ = sim_MvNormal(c1, B1)

        ## Check for stationarity
        coef = [-reverse(vec(ϕ), dims = 1); 1]                      # check stationarity
        root = roots(Polynomial(reverse(coef)))
        rootmod = abs.(root)
        accept = min(rootmod...) >= 1.01

        ## If while loop goes on for too long
        if ind > 100
            ϕ = ϕold
            coef = [-reverse(vec(ϕ), dims = 1); 1]                      # check stationarity
            root = roots(Polynomial(reverse(coef)))
            rootmod = abs.(root)
            accept = min(rootmod...) >= 1.01
        end
    end
    ##=#

    #ϕ = sim_MvNormal(c1, B1)

    ##################################
    ##################################
    # Generate σ2

    ## Prior parameters in IG(ν0/2, δ0/2)
    ν0 = 0.002
    δ0 = 0.002

    ## Posterior parameters in IG(ν1/2, δ1/2)
    ν1 = ν0 + T
    δ1 = δ0 .+ norm(Y_star - X_star * β)^2
    δ1 = δ1[1, 1]

    ## Generate new σ2
    σ2 = rand(InverseGamma(ν1 / 2, δ1 / 2))

    ##################################
    ##################################
    # Return parameters
    return β, σ2, ϕ
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################


include("FractalDimensions.jl")

#S_0 = [[0.0,0.0], [1.0,0.0], [0.5,1.0]]
S_0 = [[0.0,0.0], [1.0,0.0], [0.5,2.0]]
#S_0 = [[0.0,0.0], [1.0,0.0], [0.0,1.0]]

f1(x) = (S_0[1]+x)/2
f2(x) = (S_0[2]+x)/2
f3(x) = (S_0[3]+x)/2

# WARNING This will use up your memory and slow down your computer for k>10
function sierpinsky_endpoints(k)
        s::Vector{Vector{Float64}} = []
        s1 = S_0
        for i in 1:k
            s2 = setdiff([f1.(s1); f2.(s1); f3.(s1)], s1)
            append!(s, s1)
            s1 = s2
        end
    return s
end

s = sierpinsky_endpoints(10)

using Plots
x = [u[1] for u in s]
y = [u[2] for u in s]

plotlyjs()
scatter(x,y, markersize=1, markerstrokealpha=0.0, legend=false)
savefig("sierpinsky.eps")

d, N, ε = box_counting_dimension(s, 9, 0.0)

loglog_regression(N,ε)

log(3)/log(2)

plot(-log.(ε), log.(N), xlabel="-ln ε", ylabel="ln N", marker=:d, legend=false)
savefig("boxdim_sierpinsky.eps")

C, r = correlation_dimension(s)

loglog_regression(C,r)

plot(log.(r), log.(C), xlabel="ln r", ylabel="ln C", marker=:d, legend=false)
savefig("cordim_sierpinsky.eps")


function acqr_ksmoother(a,xtt,Ptt,xtt1,Ptt1)
	#
	# Kalman smoother for model
	#
	# x_{t+1} = ax_{t} + w_{t}
	# y_{t}   = cx_{t} + v_{t}
	#
	# cov(w_{t},v_{t}) = [q 0;0 r]
	#
	# javier.cara@upm.es, 2020-05
	#

	nt = length(xtt)

	# allocation
	xtN = zeros(nt+1)
	PtN = zeros(nt+1)
	Pt1tN = zeros(nt)

	# values for t = nt+1
	xtN[nt+1] = xtt1[nt+1]
	PtN[nt+1] = Ptt1[nt+1]

	# values for t=nt
	xtN[nt] = xtt[nt]
	PtN[nt] = Ptt[nt]
	Pt1tN[nt] = a*Ptt[nt]

	# smother
	for t = nt-1:-1:1
		# Kalman Smoother matrix J
		Jt = (Ptt[t]*a) / Ptt1[t+1]

		xtN[t] = xtt[t] + Jt*( xtN[t+1] - xtt1[t+1] )
		PtN[t] = Ptt[t] + Jt^2*( PtN[t+1] - Ptt1[t+1] )

		Pt1tN[t] = PtN[t+1]*Jt
	end

	return xtN,PtN,Pt1tN

end


#####################
# Utility Functions #
#####################

"""
Kronecker Delta Function.

Returns 1 if inputs are equal returns 0 otherwise.

Arguments:
- `a::Real`: First input argument
- `b::Real`: Second input argument

Returns:
- `delta:Integer`: Kronecker delta result.
"""
function kron(a::Real, b::Real)
    return a == b ? 1 : 0
end

"""
Internal helper function to aid in denormalization of gravity field coefficients.

Provides method co compute the factorial ratio: (n-m)!/(n+m)!

Arguments:
- `n::Real`: Gravity model degree, n.
- `m::Real`: Gravity model order, m.

Returns:
- `p::Real`: Factorial product
"""
function facprod(n::Int, m::Int)
    p = 1.0

    for i in (n-m+1):(n+m+1) # Is this correct? Should it be (n-m+1):(n+m)
        p = p/i
    end

    return p
end

###########
# Gravity #
###########

export accel_point_mass
"""
Computes the acceleration of a satellite caused by a point-mass approximation
of the central body. Returns the acceleration vector of the satellite.

Assumes the satellite is much, much less massive than the central body.

Arguments:
- `r_sat::Array{<:Real, 1}`: satellite position in a commonn inertial frame [m]
- `r_body::Array{<:Real, 1}`: position of body in a commonn inertial frame [m]
- `GM::Array{<:Real, 1}`: gravitational coeffient of attracting body [m^3/s^2] Default: SatelliteDynamics.GM_EARTH)
(Default: SatelliteDynamics.GM_EARTH

Return:
- `a::Array{<:Real, 1}`: Acceleration in X, Y, and Z inertial directions [m/s^2]
"""
function accel_point_mass(r_sat::Array{<:Real, 1}, r_body::Array{<:Real, 1}, gm_body::Real=GM_EARTH)
    # Restrict inputs to position only
    r_sat  = r_sat[1:3]
    r_body = r_body[1:3]

    # Relative position vector of satellite w.r.t. the attraching body
    d = r_sat - r_body

    # Acceleration
    a = -gm_body * (d/norm(d)^3 + r_body/norm(r_body)^3)

    return a
end

"""
Computes the acceleration on a satellite caused by a point-mass approximation
of a massive body. Returns the acceleration vector of the satellite.

Arguments:
- `r_sat::Array{<:Real, 1}`: satellite position in the inertial frame [m]
- `GM::Array{<:Real, 1}`: gravitational coeffient of attracting body [m^3/s^2] Default: SatelliteDynamics.GM_EARTH)
(Default: SatelliteDynamics.GM_EARTH

Return:
- `a::Array{<:Real, 1}`: Acceleration in X, Y, and Z inertial directions [m/s^2]
"""
function accel_point_mass(x::Array{<:Real, 1}, gm_body::Real=GM_EARTH)
    # Restrict inputs to position only. Considered in body frame
    r  = x[1:3]

    # Acceleration
    a = -gm_body * r/norm(r)^3

    return a
end

"""
Compute the gravitational acceleration at a model given a spherical harmonic
gravity field model.

Arguments:
- `r::Array{<:Real, 1}`: Position of the point in the body (field) fixed frame. [m]
- `coef::Array{<:Real, 2}`: Gravity coefficients stored in dense matrix form. C_nm terns are stored along rows indexed from C_00 in coef[1, 1] to C_nm in coef[n+1, m+1]. S_nm terms are stored along matrix columns with S_n0 stored in coef[0, n+1]
- `n_max::Integer`: Maximum degree coefficient to use in expansion
- `m_max::Integer`: Maximum order coefficient to use in the expansion. Must be less than the degree.
- `r_ref::Real`: Reference distance of the gravity field.
- `GM::Real`: Gravitational constant of central body
- `normralized::Bool`: Whether the input gravity field coefficients are normalized coefficients (Default: true)

Returns:
- `a::Array{<:Real, 1}`: Acceleration in X, Y, and Z inertial directions [m/s^2]
"""
function spherical_harmonic_gravity(r::Array{<:Real, 1}, coef::Array{<:Real, 2}, n_max::Integer, m_max::Integer, r_ref::Real, GM::Real; normalized::Bool=true)
    # Intermediate computations
    r_sqr = dot(r, r)
    rho   = r_ref^2/r_sqr
    x0    = r_ref * r[1] / r_sqr
    y0    = r_ref * r[2] / r_sqr
    z0    = r_ref * r[3] / r_sqr

    # Initialize Intermediary Matrices
    V = zeros(eltype(r), n_max+2, n_max+2)
    W = zeros(eltype(r), n_max+2, n_max+2)

    # Calculate zonal terms V(n, 0). Set W(n,0)=0.0
    V[0+1, 0+1] = r_ref /sqrt(r_sqr)
    W[0+1, 0+1] = 0.0

    V[1+1, 0+1] = z0 * V[0+1, 0+1]
    W[1+1, 0+1] = 0.0

    for n in 2:(n_max+1)
        V[n+1, 0+1] = ((2*n-1)*z0*V[n+1-1, 0+1] - (n-1)*rho*V[n+1-2,0+1])/n
        W[n+1, 0+1] = 0.0
    end

    # Calculate tesseral and sectoral terms
    for m in 1:m_max+1
        # Calculate V(m,m) to V(n_max+1,m)
        V[m+1, m+1] = (2*m-1)*(x0*V[m+1-1, m+1-1] - y0*W[m+1-1, m+1-1])
        W[m+1, m+1] = (2*m-1)*(x0*W[m+1-1, m+1-1] + y0*V[m+1-1, m+1-1])

        if m <= m_max
            V[m+1+1, m+1] = (2*m+1)*z0*V[m+1,m+1]
            W[m+1+1, m+1] = (2*m+1)*z0*W[m+1,m+1]
        end

        for n in (m+2):(n_max+1)
            V[n+1,m+1] = ((2*n-1)*z0*V[n+1-1,m+1]-(n+m-1)*rho*V[n+1-2,m+1])/(n-m)
            W[n+1,m+1] = ((2*n-1)*z0*W[n+1-1,m+1]-(n+m-1)*rho*W[n+1-2,m+1])/(n-m)
        end
    end

    # Calculate accelerations
    ax = 0.0
    ay = 0.0
    az = 0.0

    for m in 0:m_max
        for n in m:n_max
            C = 0.0
            S = 0.0
            if m == 0
                # Denormalize Coeeficients
                if normalized
                    N = sqrt(2*n+1)
                    C = N*coef[n+1, 0+1]
                else
                    C = coef[n+1, 0+1]
                end

                ax -= C*V[n+1+1, 1+1]
                ay -= C*W[n+1+1, 1+1]
                az -= (n+1)*C*V[n+1+1, 0+1]

            else
                if normalized
                    N = sqrt((2 - kron(0,m))*(2*n+1)*facprod(n,m))
                    C = N*coef[n+1,   m+1]
                    S = N*coef[m+1-1, n+1]
                else
                    C = coef[n+1,   m+1]
                    S = coef[m+1-1, n+1]
                end

                fac =  0.5 * (n-m+1)*(n-m+2)
                ax += +0.5 * (-C * V[n+1+1, m+1+1] - S * W[n+1+1, m+1+1])
                      +fac * (+C * V[n+1+1, m+1-1] + S * W[n+1+1, m+1-1])
                ay += +0.5 * (-C * W[n+1+1, m+1+1] + S * V[n+1+1, m+1+1])
                      +fac * (-C * W[n+1+1, m+1-1] + S * V[n+1+1, m+1-1])
                az += (n-m+1)*(-C*V[n+1+1, m+1] - S*W[n+1+1, m+1])
            end
        end
    end

    a = (GM / (r_ref^2)) * [ax, ay, az]

    return a
end

export accel_gravity
"""
Computes the accleration caused by Earth gravity as modeled by a spherical
harmonic gravity field.

Arguments:
- `r_sat::Array{<:Real, 1}`: Satellite position in the inertial frame [m]
- `R_eci_ecef::Array{<:Real, 2}`: Rotation matrix transforming a vector from the inertial to body-fixed reference frames.
- `n_max::Integer`: Maximum degree coefficient to use in expansion
- `m_max::Integer`: Maximum order coefficient to use in the expansion. Must be less than the degree.

Return:
- `a::Array{<:Real, 1}`: Gravitational acceleration in X, Y, and Z inertial directions [m/s^2]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.56-68.
"""
function accel_gravity(x::Array{<:Real, 1}, R_eci_ecef::Array{<:Real, 2}, n_max::Int=20, m_max::Int=20)

    # Check Limits of Gravity Field
    if n_max > GRAVITY_MODEL.n_max
        error("Requested gravity model order $n_max is larger than the maximum order of the model maximum: ", GRAVITY_MODEL.n_max)
    end

    if m_max > GRAVITY_MODEL.m_max
        error("Requested gravity model degree $m_max is larger than the maximum degree of the model maximum: ", GRAVITY_MODEL.m_max)
    end

    # Gravitational parameters of primary body
    GM    = GRAVITY_MODEL.GM
    R_ref = GRAVITY_MODEL.R

    # Body-fixed position
    r_bf = R_eci_ecef * x[1:3]

    # Compute spherical harmonic acceleration
    a_ecef = spherical_harmonic_gravity(r_bf, GRAVITY_MODEL.data, n_max, m_max, R_ref, GM, normalized=GRAVITY_MODEL.normalized)

    # Inertial acceleration
    a_eci = R_eci_ecef' * a_ecef

    # Finished
    return a_eci
end

#########################
# Planetary Ephemerides #
#########################

export sun_position
"""
Compute the Sun's position in the EME2000 inertial frame through the use
of low-precision analytical functions.

Argument:
- `epc::Epoch`: Epoch

Returns:
- `r_sun::Array{<:Real, 1}`: Position vector of the Sun in the Earth-centered inertial fame.

Notes:
1. The EME2000 inertial frame is for most purposes equivalent to the GCRF frame.

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.70-73.
"""
function sun_position(epc::Epoch)
    # Constants
    mjd_tt  = mjd(epc, tsys="TT")      # MJD of epoch in TT
    epsilon = 23.43929111*pi/180.0     # Obliquity of J2000 ecliptic
    T       = (mjd_tt-MJD2000)/36525.0 # Julian cent. since J2000

    # Variables

    # Mean anomaly, ecliptic longitude and radius
    M = 2.0*pi * modf(0.9931267 + 99.9973583*T)[1]                 # [rad]
    L = 2.0*pi * modf(0.7859444 + M/(2.0*pi) + (6892.0*sin(M)+72.0*sin(2.0*M)) / 1296.0e3)[1] # [rad]
    r = 149.619e9 - 2.499e9*cos(M) - 0.021e9*cos(2*M)           # [m]

    # Equatorial position vector
    p_sun = Rx(-epsilon) * [r*cos(L), r*sin(L), 0.0]

    return p_sun
end

export moon_position
"""
Compute the Moon's position in the EME2000 inertial frame through the use
of low-precision analytical functions.

Argument:
- `epc::Epoch`: Epoch

Returns:
- `r_moon::Array{<:Real, 1}`: Position vector of the Moon in the Earth-centered inertial fame.

Notes:
1. The EME2000 inertial frame is for most purposes equivalent to the GCRF frame.

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.70-73.
"""
function moon_position(epc::Epoch)
    # Constants
    mjd_tt  = mjd(epc, tsys="TT")      # MJD of epoch in TT
    epsilon = 23.43929111*pi/180.0     # Obliquity of J2000 ecliptic
    T       = (mjd_tt-MJD2000)/36525.0 # Julian cent. since J2000

    # Mean elements of lunar orbit
    L_0 =     modf(0.606433 + 1336.851344*T)[1] # Mean longitude [rev] w.r.t. J2000 equinox
    l   = 2.0*pi*modf(0.374897 + 1325.552410*T)[1] # Moon's mean anomaly [rad]
    lp  = 2.0*pi*modf(0.993133 +   99.997361*T)[1] # Sun's mean anomaly [rad]
    D   = 2.0*pi*modf(0.827361 + 1236.853086*T)[1] # Diff. long. Moon-Sun [rad]
    F   = 2.0*pi*modf(0.259086 + 1342.227825*T)[1] # Argument of latitude

    # Ecliptic longitude (w.r.t. equinox of J2000)
    dL = + 22640*sin(l) - 4586*sin(l-2*D) + 2370*sin(2*D) +  769*sin(2*l)
         - 668*sin(lp) - 412*sin(2*F) - 212*sin(2*l-2*D) - 206*sin(l+lp-2*D)
         + 192*sin(l+2*D) - 165*sin(lp-2*D) - 125*sin(D) - 110*sin(l+lp)
         + 148*sin(l-lp) - 55*sin(2*F-2*D)

    L = 2.0*pi * modf(L_0 + dL/1296.0e3)[1]  # [rad]

    # Ecliptic latitude
    S  = F + (dL+412*sin(2*F)+541*sin(lp)) * AS2RAD
    h  = F-2*D
    N  = - 526*sin(h) + 44*sin(l+h) - 31*sin(-l+h) - 23*sin(lp+h)
         + 11*sin(-lp+h) - 25*sin(-2*l+F) + 21*sin(-l+F)
    B  = (18520.0*sin(S) + N) * AS2RAD   # [rad]

    # Distance [m]
    r = + 385000e3 - 20905e3*cos(l) - 3699e3*cos(2*D-l) - 2956e3*cos(2*D)
        - 570e3*cos(2*l) + 246e3*cos(2*l-2*D) - 205e3*cos(lp-2*D)
        - 171e3*cos(l+2*D) - 152e3*cos(l+lp-2*D)

    # Equatorial coordinates
    p_moon = Rx(-epsilon) * [r*cos(L)*cos(B), r*sin(L)*cos(B), r*sin(B)]

    return p_moon
end

######################
# Third-Body Gravity #
######################

export accel_thirdbody_sun
"""
Computes the acceleration of a satellite in the inertial frame due to the
gravitational attraction of the Sun.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]
- `r_sun::Array{<:Real, 1}`: Position of sun in inertial frame.

Return:
- `a::Array{<:Real, 1}`: Acceleration due to the Sun's gravity in the inertial frame [m/s^2]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.69-70.
"""
function accel_thirdbody_sun(x::Array{<:Real, 1}, r_sun::Array{<:Real, 1})
    # Acceleration due to sun point mass
    a_sun = accel_point_mass(x[1:3], r_sun, GM_SUN)

    return a_sun
end

function accel_thirdbody_sun(epc::Epoch, x::Array{<:Real, 1})
    # Compute solar position
    r_sun = sun_position(epc)

    # Acceleration due to sun point mass
    a_sun = accel_point_mass(x[1:3], r_sun, GM_SUN)

    return a_sun
end

export accel_thirdbody_moon
"""
Computes the acceleration of a satellite in the inertial frame due to the
gravitational attraction of the Moon.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]
- `r_moon::Array{<:Real, 1}`: Position of moon in inertial frame.

Returns:
- `a::Array{<:Real, 1}`: Acceleration due to the Moon's gravity in the inertial frame [m/s^2]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.69-70.
"""
function accel_thirdbody_moon(x::Array{<:Real, 1}, r_moon::Array{<:Real, 1})
    # Acceleration due to moon point mass
    a_moon = accel_point_mass(x[1:3], r_moon, GM_MOON)

    return a_moon
end

function accel_thirdbody_moon(epc::Epoch, x::Array{<:Real, 1})
    # Compute solar position
    r_moon = moon_position(epc)

    # Acceleration due to moon point mass
    a_moon = accel_point_mass(x[1:3], r_moon, GM_MOON)

    return a_moon
end

####################
# Atmospheric Drag #
####################

export density_harris_priester
"""
Computes the local density using the Harris-Priester density model.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]
- `r_sun::Array{<:Real, 1}`: Position of sun in inertial frame.

Returns:
- `rho:Float64`: Local atmospheric density [kg/m^3]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.89-91.
"""
function density_harris_priester(x::Array{<:Real, 1}, r_sun::Array{<:Real, 1})
    # Harris-Priester Constants
    hp_upper_limit =   1000.0          # Upper height limit [km]
    hp_lower_limit =    100.0          # Lower height limit [km]
    hp_ra_lag      = 0.523599          # Right ascension lag [rad]
    hp_n_prm       =        3          # Harris-Priester parameter
                                        # 2(6) low(high) inclination
    hp_N           = 50                # Number of coefficients

    # Height [km]
    hp_h = [100.0, 120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0,
            210.0, 220.0, 230.0, 240.0, 250.0, 260.0, 270.0, 280.0, 290.0, 300.0,
            320.0, 340.0, 360.0, 380.0, 400.0, 420.0, 440.0, 460.0, 480.0, 500.0,
            520.0, 540.0, 560.0, 580.0, 600.0, 620.0, 640.0, 660.0, 680.0, 700.0,
            720.0, 740.0, 760.0, 780.0, 800.0, 840.0, 880.0, 920.0, 960.0,1000.0]

    # Minimum density [g/km^3]
    hp_c_min = [4.974e+05, 2.490e+04, 8.377e+03, 3.899e+03, 2.122e+03, 1.263e+03,
                8.008e+02, 5.283e+02, 3.617e+02, 2.557e+02, 1.839e+02, 1.341e+02,
                9.949e+01, 7.488e+01, 5.709e+01, 4.403e+01, 3.430e+01, 2.697e+01,
                2.139e+01, 1.708e+01, 1.099e+01, 7.214e+00, 4.824e+00, 3.274e+00,
                2.249e+00, 1.558e+00, 1.091e+00, 7.701e-01, 5.474e-01, 3.916e-01,
                2.819e-01, 2.042e-01, 1.488e-01, 1.092e-01, 8.070e-02, 6.012e-02,
                4.519e-02, 3.430e-02, 2.632e-02, 2.043e-02, 1.607e-02, 1.281e-02,
                1.036e-02, 8.496e-03, 7.069e-03, 4.680e-03, 3.200e-03, 2.210e-03,
                1.560e-03, 1.150e-03]

    # Maximum density [g/km^3]
    hp_c_max = [4.974e+05, 2.490e+04, 8.710e+03, 4.059e+03, 2.215e+03, 1.344e+03,
                8.758e+02, 6.010e+02, 4.297e+02, 3.162e+02, 2.396e+02, 1.853e+02,
                1.455e+02, 1.157e+02, 9.308e+01, 7.555e+01, 6.182e+01, 5.095e+01,
                4.226e+01, 3.526e+01, 2.511e+01, 1.819e+01, 1.337e+01, 9.955e+00,
                7.492e+00, 5.684e+00, 4.355e+00, 3.362e+00, 2.612e+00, 2.042e+00,
                1.605e+00, 1.267e+00, 1.005e+00, 7.997e-01, 6.390e-01, 5.123e-01,
                4.121e-01, 3.325e-01, 2.691e-01, 2.185e-01, 1.779e-01, 1.452e-01,
                1.190e-01, 9.776e-02, 8.059e-02, 5.741e-02, 4.210e-02, 3.130e-02,
                2.360e-02, 1.810e-02]

    # Satellite height
    geod   = sECEFtoGEOD(x[1:3], use_degrees=true)
    height = geod[3]/1.0e3 # height in [km]

    # Exit with zero density outside height model limits
    if height > hp_upper_limit || height < hp_lower_limit
        return 0.0
    end

    # Sun right ascension, declination
    ra_sun  = atan( r_sun[2], r_sun[1] )
    dec_sun = atan( r_sun[3], sqrt( r_sun[1]^2 + r_sun[2]^2 ) )

    # Unit vector u towards the apex of the diurnal bulge
    # in inertial geocentric coordinates
    c_dec = cos(dec_sun)
    u     = [c_dec * cos(ra_sun + hp_ra_lag),
             c_dec * sin(ra_sun + hp_ra_lag),
             sin(dec_sun)]

    # Cosine of half angle between satellite position vector and
    # apex of diurnal bulge
    c_psi2 = 0.5 + 0.5 * dot(x[1:3], u)/norm(x[1:3])

    # Height index search and exponential density interpolation
    ih = 0                            # section index reset
    for i in 1:hp_N                   # loop over N_Coef height regimes
        if height >= hp_h[i] && height < hp_h[i+1]
            ih = i                    # ih identifies height section
            break
        end
    end

    h_min = ( hp_h[ih] - hp_h[ih+1] )/log( hp_c_min[ih+1]/hp_c_min[ih] )
    h_max = ( hp_h[ih] - hp_h[ih+1] )/log( hp_c_max[ih+1]/hp_c_max[ih] )

    d_min = hp_c_min[ih] * exp( (hp_h[ih]-height)/h_min )
    d_max = hp_c_max[ih] * exp( (hp_h[ih]-height)/h_max )

    # Density computation
    density = d_min + (d_max-d_min) * c_psi2^hp_n_prm

    # Convert from g/km^3 to kg/m^3
    density *= 1.0e-12

    # Finished
    return density
end

function density_harris_priester(epc::Epoch, x::Array{<:Real, 1})
    r_sun = sun_position(epc)
    return density_harris_priester(x, r_sun)
end

export accel_drag
"""
Computes the perturbing, non-conservative acceleration caused by atmospheric
drag assuming that the ballistic properties of the spacecraft are captured by
the coefficient of drag.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]
- `rho::Real`: atmospheric density [kg/m^3]
- `mass::Real`: Spacecraft mass [kg]
- `area::Real`: Wind-facing cross-sectional area [m^2]
- `Cd::Real`: coefficient of drag [dimensionless]
- `T::Array{<:Real, 2}`: Rotation matrix from the inertial to the true-of-date frame

Return:
- `a::Array{<:Real, 1}`: Acceleration due to drag in the X, Y, and Z inertial directions. [m/s^2]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.83-86.
"""
function accel_drag(x::Array{<:Real, 1}, rho::Real, mass::Real, area::Real, Cd::Real, T::Array{<:Real, 2})

    # Constants
    omega = [0, 0, OMEGA_EARTH]

    # Position and velocity in true-of-date system
    r_tod = T * x[1:3]
    v_tod = T * x[4:6]

    # Velocity relative to the Earth's atmosphere
    v_rel = v_tod - cross(omega, r_tod)
    v_abs = norm(v_rel)

    # Acceleration
    a_tod  = -0.5*Cd*(area/mass)*rho*v_abs*v_rel
    a_drag = T' * a_tod

    return a_drag
end

############################
# Solar Radiation Pressure #
############################

export eclipse_cylindrical
"""
Computes the illumination fraction of a satellite in Earth orbit using a
cylindrical Earth shadow model.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]
- `r_sun::Array{<:Real, 1}`: Position of sun in inertial frame.

Return:
- `nu::Float64`: Illumination fraction (0 <= nu <= 1). nu = 0 means spacecraft in complete shadow, nu = 1 mean spacecraft fully illuminated by sun.

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.80-83.
"""
function eclipse_cylindrical(x::Array{<:Real, 1}, r_sun::Array{<:Real, 1})
    # Satellite inertial position
    r = x[1:3]

    # Sun-direction unit-vector
    e_sun = r_sun / norm(r_sun)

    # Projection of spacecraft position
    s = dot(r, e_sun)

    # Compute illumination
    nu = 0.0
    if s/norm(s) >= 1.0 || norm(r - s*e_sun) > R_EARTH
        nu = 1.0
    end

    return nu
end

function eclipse_cylindrical(epc::Epoch, x::Array{<:Real, 1})
    r_sun = sun_position(epc)
    return eclipse_cylindrical(x, r_sun)
end

export eclipse_conical
"""
Computes the illumination fraction of a satellite in Earth orbit using a
conical Earth shadow model.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]
- `r_sun::Array{<:Real, 1}`: Position of sun in inertial frame.

Return:
- `nu::Float64`: Illumination fraction (0 <= nu <= 1). nu = 0 means spacecraft in complete shadow, nu = 1 mean spacecraft fully illuminated by sun.

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.80-83.
"""
function eclipse_conical(x::Array{<:Real, 1}, r_sun::Array{<:Real, 1})
    # Satellite inertial position
    r = x[1:3]

    # Occultation Geometry
    a = asin(R_SUN/norm(r_sun - r))
    b = asin(R_EARTH/norm(r))
    c = acos(dot(r, r_sun-r)/(norm(r)*norm(r_sun-r)) )

    e_sun = r_sun / norm(r_sun)

    # Test Occulation Conditions
    nu = 0.0
    if abs(a - b) < c && c < (a + b)
        # Partial occultation

        xx = (c^2 + a^2 - b^2)/(2*c)
        yy = sqrt(a^2 - xx^2)
        A  = a^2 * acos(xx/a) + b^2 * acos((c-xx)/b) - c * yy

        nu = 1 - A/(pi*a^2)
    elseif (a + b) <= c
        # No occultation
        nu = 1.0
    else
        # Full occultation
        nu = 0.0
    end

    return nu
end

function eclipse_conical(epc::Epoch, x::Array{<:Real, 1})
    r_sun = sun_position(epc)
    return eclipse_conical(x, r_sun)
end

export accel_srp
"""Computes the perturbing acceleration due to direct solar radiation
pressure assuming the reflecting surface is a flat plate pointed directly at
the Sun.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]

Returns:
- `a::Array{<:Real, 1}`: Satellite acceleration due to solar radiation pressure [m/s^2]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.77-79.
"""
function accel_srp(x::Array{<:Real, 1}, r_sun::Array{<:Real, 1}, mass::Real=0, area::Real=0, CR::Real=1.8, p0::Real=P_SUN, au::Real=AU)
    # Spacecraft position vector
    r = x[1:3]

    # Relative position vector of spacecraft w.r.t. Sun
    d = r - r_sun

    # Acceleration due to moon point mass
    a_srp = d * (CR*(area/mass)*p0*AU^2 / norm(d)^3)

    # Return
    return a_srp
end

##############
# Relativity #
##############

export accel_relativity
"""Computes perturbation accleration of a satellite in the Inertial frame
due to the combined effects of special and general relativity.

Arguments:
- `x::Array{<:Real, 1}`: Satellite Cartesean state in the inertial reference frame [m; m/s]

Returns:
- `a::Array{<:Real, 1}`: Satellite acceleration due to relativity. [m/s^2]

References:
1. O. Montenbruck, and E. Gill, _Satellite Orbits: Models, Methods and Applications_, 2012, p.110-112.
"""
function accel_relativity(x::Array{<:Real, 1})
    # Extract state variables
    r = x[1:3]
    v = x[4:6]

    # Intermediate computations
    norm_r = norm(r)
    r2     = norm_r^2

    norm_v = norm(v)
    v2     = norm_v^2

    c  = C_LIGHT
    c2 = c^2

    # Compute unit vectors
    er = r/norm_r
    ev = v/norm_v

    # Compute perturbation and return
    a_rel = GM_EARTH/r2 * ( (4*GM_EARTH/(c2*norm_r) - v2/c2)*er + 4*v2/c2*dot(er, ev)*ev)

    return a_rel
end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 37b018c4-34b0-11eb-234d-4fbeaa683873
using LinearAlgebra, SparseArrays, Plots, DifferentialEquations

# ╔═╡ baed3ace-34ad-11eb-0cf6-19357238b1de
md"## Heat equation with Dirichlet boundary conditions

The domain is given by the one-dimensional rod $\Omega = (0,L)$ and the time by $t \in (0,T).$ The rod has the physical properties
- thermal conductivity $\lambda$,
- specific heat capacity $c$,
- mass density $\rho$.

### Heat Equation

The diffusivity constant is given by $\alpha = \frac{\lambda}{c \rho}$. The heat equation is noted as

$\dot{\vartheta}(t,x) = \alpha \frac{\partial^2}{\partial x^2} \vartheta(t,x)$

for $(t,x) \in (0,T) \times \Omega$ with initial condition

$\vartheta(0,x) = \vartheta_{0}(x)$

for $x \in \overline{\Omega}$ and boundary conditions  $\qquad \vartheta(\cdot,0) = \vartheta(\cdot,L) = 0$.
"

# ╔═╡ 8f5d25dc-34ae-11eb-2aa7-39e601f648a1
md"#### Code: implementing physical constants"

# ╔═╡ 8f493cca-34ae-11eb-3e60-31ab68674568
# Length of rod
L = 0.5

# ╔═╡ 8f149f6a-34ae-11eb-01f2-5983e256a7de
# Thermal conductivity
λ = 45.0

# ╔═╡ 8efa0286-34ae-11eb-1952-e790d141e19b
# Specific heat capacity
c = 480.0

# ╔═╡ 8ecb6bba-34ae-11eb-2c16-819675412725
# Mass density
ρ = 7800.0

# ╔═╡ e6433430-34af-11eb-318d-2199461e5ec9
# Thermal diffusivity
α = λ/(c*ρ)

# ╔═╡ 37c91fae-34b0-11eb-0f82-639173602ac0
md"### Spatial approximation

The one-dimensional rod is discretized as a one-dimensional grid with N points and the finite discretization is noted by

$\Delta x = \frac{L}{N-1}.$

The resulting grid points have the position $x^{0} = 0$, $x^{n} = n~\Delta x$ and $x^{N-1} = L$. The second order derivative is approximated using the [Taylor series](https://en.wikipedia.org/wiki/Taylor_series) as

$\left. \frac{\partial^2 f(x)}{\partial x^2} \right\rvert_{x = \tilde{x}} \approx \frac{1}{\Delta x^2} \left[ f(\tilde{x} - \Delta x) - 2 f(\tilde{x}) + f(\tilde{x} + \Delta x)  \right].$

This [Finite Difference](https://en.wikipedia.org/wiki/Finite_difference_method) scheme is written in matrix-vector form as

$\frac{\partial^2 \vartheta(t,x)}{\partial x^2} \approx \frac{1}{\Delta x^2}
\begin{pmatrix}
-2 & 1 & 0 & \cdots &  & 0 \\
1 & -2 & 1 & 0  & \cdots & \vdots \\
0 & 1 & -2 & 1 & \ddots &  \\
 & & \ddots & \ddots & \ddots &  & \\
& & & 1 & -2 & 1 \\
& & &  & 1 & -2 \\
\end{pmatrix}
~
\begin{pmatrix}
\vartheta(t, x^{0}) \\
\vartheta(t, x^{1}) \\
\vdots \\
\vartheta(t, x^{N-1}) \\
\end{pmatrix}.$
"

# ╔═╡ 65049b1a-3549-11eb-33ee-9dfae609164d
# Number of grid elements
N = 101

# ╔═╡ a36fd440-3549-11eb-0e93-0d36244400ef
# Finite discretization
Δx = L/(N-1)

# ╔═╡ c18a976a-3549-11eb-2991-799f558e4493
# 1-dimensional grid
xspan = 0 : Δx : L

# ╔═╡ d7d5be46-3549-11eb-3a12-afab5ab4bce5
# Diffusion matrix
M = spdiagm(-1 => ones(N-1), 0 => -2*ones(N), 1 => ones(N-1));

# ╔═╡ 045066e2-354a-11eb-164b-09fc7095de47
Matrix(M)[1:5,1:5]

# ╔═╡ 37822ec8-34b0-11eb-175d-7745382c2955
md"### Boundary conditions

The Dirichlet boundary condition states that the value at both sides is fixed as

$\qquad \vartheta(\cdot,0) = \vartheta(\cdot, x^{0}) = 0$

on the left boundary and

$\vartheta(\cdot,L) = \vartheta(\cdot, x^{N-1})  = 0$

on the right boundary. This means, there is no dynamical (e.g. diffusive) behaviour at both boundaries and thus the first and the last row of the diffusion matrix M is set to 0. "

# ╔═╡ 3768b7e0-34b0-11eb-0772-f5f4c95cf8f5
# First row
M[1,1] = M[1,2] = 0;

# ╔═╡ 37529500-34b0-11eb-1ee4-b5837bea23c9
# Last row
M[end,end] = M[end,end-1] = 0;

# ╔═╡ 373652f0-34b0-11eb-2fea-4f14b3ee9140
md"### Heat equation as ODE

After the spatial approximation the heat equation has a form of

$\dot{\theta}(t) = \frac{\alpha}{\Delta x^2} ~ M ~ \theta(t)$

with M as diffusion matrix and $\quad \theta(t) = \left( \vartheta(t, x^{0}), \cdots, \vartheta(t, x^{N-1}) \right)^{\top}$. This is an ordinary differential equation (ODE) that can be solved with common solvers like forward Euler method or Runge-Kutta scheme."

# ╔═╡ 371c6d52-34b0-11eb-12ab-a51e10fe33c5
# Heat Equation as ODE
function heat_eq(dθ, θ, p, t)
    return dθ .= (1/Δx^2) * α * M * θ
end

# ╔═╡ 37013e12-34b0-11eb-3d23-9d43ca3f06a3
md"### Initial Conditions

The initial data of the original heat equation is assumed as $\quad \vartheta_{0}(x) := m \left[ L~x - x^2 \right]~$ with  $m > 0$. This function is approximated to gain the initial conditions of the ODE

$\theta(0) = \left( \vartheta_{0}(x^{0}), \cdots , \vartheta_{0}(x^{N-1}) \right)^{\top}.$
"

# ╔═╡ 78b7fa20-354c-11eb-3754-d5d515d0a2d5
#Initial conditions: θ(0) = m * (L * x - x^2)
function initial_data(x, params)
    # α = params[1] # unused here
    L = params[2]
    m = params[3]

    return m*( L * x - x^2)
end

# ╔═╡ 77906a42-354c-11eb-397e-9dca847ab964
# Initial heat distribution
θ₀ = zeros(N);

# ╔═╡ 93696cf4-354d-11eb-0af3-9b27a8186fbd
# Amplification of initial data
@bind amp html"<input type='range' min='0' max='100' step='1'>"

# ╔═╡ a91c8b80-354d-11eb-03ec-8f4386a69708
m = amp

# ╔═╡ 86f0f21c-354d-11eb-2a32-299183f9118b
# Parameters
param = [α, L, amp]

# ╔═╡ 666a8878-354d-11eb-22dc-e94a72ede22d
let

	for i = 1 : length(xspan)
		θ₀[i] = initial_data(xspan[i], param)
	end

	plot(xspan, θ₀, xlabel="Position x", ylabel="Temperature", legend=false)
end

# ╔═╡ cfc0e592-354d-11eb-2b06-458d198c38e5
md"### Numerical stability

Before the approximated heat equation can be simulated, the [numerical stability](https://en.wikipedia.org/wiki/Numerical_stability) has to be proven with the [von Neumann stability analysis](https://en.wikipedia.org/wiki/Von_Neumann_stability_analysis).

In case of the forward Euler method one finds the approach

$\vartheta(t^{k+1}, x^{n}) = \vartheta(t^{k}, x^{n}) + \alpha \frac{\Delta t}{\Delta x^2} \left( \vartheta(t^{k}, x^{n-1}) - 2 ~ \vartheta(t^{k}, x^{n}) + \vartheta(t^{k}, x^{n+1}) \right)$

or equivalent

$\vartheta(t^{k+1}, x^{n}) = \left(1 - 2 \alpha \frac{\Delta t}{\Delta x^2} \right) \vartheta(t^{k}, x^{n}) + \alpha \frac{\Delta t}{\Delta x^2}  \left( \vartheta(t^{k}, x^{n-1}) + \vartheta(t^{k}, x^{n+1}) \right).$

The sampling time $\Delta t$ has to be chosen to guarantee

$\left(1 - 2 \alpha \frac{\Delta t}{\Delta x^2} \right) > 0.$

Therefore, the upper limit of the sampling time is given by

$\Delta t < \frac{1}{2} \frac{\Delta x^2}{\alpha}.$
"

# ╔═╡ 25789968-3552-11eb-1e25-35c0d1b1c461
# Upper limit of sampling time
ul = 0.5 * Δx^2 / α

# ╔═╡ 56daee6e-3554-11eb-37cd-49ad4ac56805
# Sampling time
Δt = 0.8

# ╔═╡ 56bb225a-3554-11eb-0062-6d03b8ea772f
# time discretization
tspan = (0.0, 1000.0)

# ╔═╡ 569a20f0-3554-11eb-0479-f1e583620682
md"### Code: Solving the ODE"

# ╔═╡ 56790ed8-3554-11eb-3c51-b9b31a4c340e
# Solving the ODE
let
	prob = ODEProblem( heat_eq, θ₀, tspan ) # ODE Problem

	sol = solve(prob,Euler(),dt=Δt,progress=true, save_everystep=false,
				save_start=true) # Solving the ODE

	plot(xspan, sol.u[2], xlabel = "Position x", ylabel="Temperature", legend=false)
end

# ╔═╡ Cell order:
# ╟─baed3ace-34ad-11eb-0cf6-19357238b1de
# ╟─8f5d25dc-34ae-11eb-2aa7-39e601f648a1
# ╠═8f493cca-34ae-11eb-3e60-31ab68674568
# ╠═8f149f6a-34ae-11eb-01f2-5983e256a7de
# ╠═8efa0286-34ae-11eb-1952-e790d141e19b
# ╠═8ecb6bba-34ae-11eb-2c16-819675412725
# ╠═e6433430-34af-11eb-318d-2199461e5ec9
# ╟─37c91fae-34b0-11eb-0f82-639173602ac0
# ╠═37b018c4-34b0-11eb-234d-4fbeaa683873
# ╠═65049b1a-3549-11eb-33ee-9dfae609164d
# ╠═a36fd440-3549-11eb-0e93-0d36244400ef
# ╠═c18a976a-3549-11eb-2991-799f558e4493
# ╠═d7d5be46-3549-11eb-3a12-afab5ab4bce5
# ╠═045066e2-354a-11eb-164b-09fc7095de47
# ╟─37822ec8-34b0-11eb-175d-7745382c2955
# ╠═3768b7e0-34b0-11eb-0772-f5f4c95cf8f5
# ╠═37529500-34b0-11eb-1ee4-b5837bea23c9
# ╟─373652f0-34b0-11eb-2fea-4f14b3ee9140
# ╠═371c6d52-34b0-11eb-12ab-a51e10fe33c5
# ╟─37013e12-34b0-11eb-3d23-9d43ca3f06a3
# ╠═78b7fa20-354c-11eb-3754-d5d515d0a2d5
# ╠═77906a42-354c-11eb-397e-9dca847ab964
# ╠═93696cf4-354d-11eb-0af3-9b27a8186fbd
# ╠═a91c8b80-354d-11eb-03ec-8f4386a69708
# ╠═86f0f21c-354d-11eb-2a32-299183f9118b
# ╠═666a8878-354d-11eb-22dc-e94a72ede22d
# ╟─cfc0e592-354d-11eb-2b06-458d198c38e5
# ╠═25789968-3552-11eb-1e25-35c0d1b1c461
# ╠═56daee6e-3554-11eb-37cd-49ad4ac56805
# ╠═56bb225a-3554-11eb-0062-6d03b8ea772f
# ╟─569a20f0-3554-11eb-0479-f1e583620682
# ╠═56790ed8-3554-11eb-3c51-b9b31a4c340e


"""
Determine the null and alternate hypotheses under t-tests given:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
"""
function hypotheses(::Type{OneSampleTTest}, n::Real, δ::Real, σ::Real)
    df = n - 1
    ncp = sqrt(n) * δ / σ
    null = TDist(df)
    alt = NoncentralT(df, ncp)
    return null, alt
end

"""
Determine the null and alternate hypotheses under t-tests given:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
"""
function hypotheses(::Type{EqualVarianceTTest}, n::Real, δ::Real, σ::Real)
    df = 2 * n - 2
    ncp = sqrt(n / 2) * δ / σ
    null = TDist(df)
    alt = NoncentralT(df, ncp)
    return null, alt
end

"""
Determine the null and alternate hypotheses under t-tests given:

* `ns`: Per-group sample size(s)
* `δ`: Deviation from the null
* `σs`: Per-group standard deviations
"""
function hypotheses(::Type{EqualVarianceTTest}, ns::Tuple, δ::Real, σ::Real)
    n1, n2 = ns[1], ns[2]
    df = n1 + n2 - 2
    ncp = (1 / sqrt(1 / n1 + 1 / n2)) * δ / σ
    null = TDist(df)
    alt = NoncentralT(df, ncp)
    return null, alt
end

# @doc """
# Determine the values of t that lead to overestimates of the magnitude of δ.
# """ ->
thresholds(alt::NoncentralT) = -alt.λ, alt.λ


"""
    floatmod(x, y)

This is the equivalent of `x // y` for floating point.
"""
floatmod(x, y) = x - y * floor(x / y)

"""
    fibonacci_points(n)

Constructs lat-long points on a sphere from a Fibonacci sequence.

There are 2n + 1 points on the sphere. It returns a `(2,2n+1)` grid of lat-long.

These points are a good start on an evenly-spaced set of points on the
surface of a sphere. This algorithm is from González [^1]. The usual way
to use these points is to make them a little more even by adjustment
with a successive over-relaxation algorithm. That algorithm moves them
a little at each step due to a repulsive force among neighbors.
Then you make a Voronoi on the sphere and create triangles from this.
The Voronoi will be almost all hexagons, but there will be a few polygons
with 5 and 7 sides.

[^1]: González, Á. (2010). Measurement of areas on a sphere using Fibonacci
and latitude–longitude lattices. Mathematical Geosciences, 42(1), 49.
"""
function fibonacci_points(n)
    ϕ = (1 + sqrt(5)) / 2
    points = zeros(Float64, 2, 2n + 1)
    for i in -n:n
        points[1, i + n + 1] = asin(2i / (2n + 1)) * 180 / π
        lon = floatmod(i, ϕ) * 360 / ϕ
        if lon < -180
            points[2, i + n + 1] = lon + 360
        elseif lon >= 180
            points[2, i + n + 1] = lon - 360
        else
            points[2, i + n + 1] = lon
        end
    end
    points
end


export bisect, cross_edges

"""
`bisect(G::SimpleGraph)` partitions the vertex set of `G` using the
eigenvector associated with the second smallest eigenvalue of the
graph's Laplacian matrix (called `x` below).

This can be invoked as follows:

+ `bisect(G,"user",pivot)` splits the vertices `v` depending on
  `x[v] >= pivot` vs. `x[v] < pivot`.

+ `bisect(G,"zero")` is the same as `bisect(G,"user", 0.0)`.

+ `bisect(G,"median")` is equivalent to `bisect(G,"user",m)` where `m`
  is the median value of `x`.

+ `bisect(G,"equal")` creates a partition in which the two parts have sizes
  the differ by at most 1.

A plain call to `bisect(G)` is equivalent to `bisect(G,"zero")` (which
is the same as `bisect(G,"user", 0.0)`).
"""
function bisect(G::SimpleGraph,
                where::AbstractString="zero",
                pivot::Real=0.0
               )

    verbose = false

    if verbose
        println("G = $G")
        println("where = $where")
        println("pivot = $pivot")
    end

    T = vertex_type(G)
    VV = vlist(G)
    n  = NV(G)
    L  = laplace(G)
    # x  = collect( eig(L)[2][:,2] )
    x = collect( eigen(L).vectors[:,2])

    if verbose
        println(sort(x))
    end

    piv = 0.0

    if where=="equal"
        pairs = sort(collect(zip(x,VV)))
        vtcs  = [ p[2] for p in pairs ]

        if verbose
            println(pairs)
            println(vtcs)
        end

        mid = floor(Int, n/2)
        if verbose
            println("Equal partition of the vertex set: $mid and $(n-mid)")
        end
        A = Set{T}(vtcs[1:mid])
        B = Set{T}(vtcs[mid+1:end])
        return A,B
    end

    if where=="median"
        piv = median(x)
    elseif where=="user"
        piv = pivot
    elseif where=="zero"
        piv = 0.0
    else
        error("Unknown \"where\" specifier: $where")
    end

    if verbose
        println("pivot set to $piv")
    end

    A = Set{T}()
    B = Set{T}()

    for k=1:n
        v = VV[k]
        if x[k] >= piv
            push!(A,v)
        else
            push!(B,v)
        end
    end

    return A,B
end

# import IterTools.product

"""
`cross_edges(G::SimpleGraph,A,B)` returns the set of edges of `G` with
one end in `A` and one end in `B`. Here `A` and `B` are collections
of vertices of `G`.
"""
function cross_edges(G::SimpleGraph, A, B)

    AB = Base.Iterators.product(A,B)

    result = Set(filter(e -> has(G,e[1],e[2]), collect(AB)))
    return result
end


#=##############################################################################
# DESCRIPTION
    Linear algebra tools

# AUTHORSHIP
  * Author    : Eduardo J Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Oct 2021
  * License   : MIT License
=###############################################################################

dot(A, B) = sum(a*b for (a,b) in zip(A, B))
norm(A) = sqrt(mapreduce(x->x^2, +, A))

cross1(A, B) = A[2]*B[3] - A[3]*B[2]
cross2(A, B) = A[3]*B[1] - A[1]*B[3]
cross3(A, B) = A[1]*B[2] - A[2]*B[1]

function cross!(out, A, B)
    out[1] = cross1(A, B)
    out[2] = cross2(A, B)
    out[3] = cross3(A, B)
end

function cross(A, B)
    out = zero(A)
    cross!(out, A, B)
    return out
end


@test collect(partitions(4)) ==  Any[[4], [3,1], [2,2], [2,1,1], [1,1,1,1]]
@test collect(partitions(8,3)) == Any[[6,1,1], [5,2,1], [4,3,1], [4,2,2], [3,3,2]]
@test collect(partitions(8, 1)) == Any[[8]]
@test collect(partitions(8, 9)) == []
@test collect(partitions([1,2,3])) == Any[Any[[1,2,3]], Any[[1,2],[3]], Any[[1,3],[2]], Any[[1],[2,3]], Any[[1],[2],[3]]]
@test collect(partitions([1,2,3,4],3)) == Any[Any[[1,2],[3],[4]], Any[[1,3],[2],[4]], Any[[1],[2,3],[4]],
                                              Any[[1,4],[2],[3]], Any[[1],[2,4],[3]], Any[[1],[2],[3,4]]]
@test collect(partitions([1,2,3,4],1)) == Any[Any[[1, 2, 3, 4]]]
@test collect(partitions([1,2,3,4],5)) == []

@inferred first(partitions(4))
@inferred first(partitions(8,3))
@inferred first(partitions([1,2,3]))
@inferred first(partitions([1,2,3,4],3))

@test isa(collect(partitions(4)), Vector{Vector{Int}})
@test isa(collect(partitions(8,3)), Vector{Vector{Int}})
@test isa(collect(partitions([1,2,3])), Vector{Vector{Vector{Int}}})
@test isa(collect(partitions([1,2,3,4], 3)), Vector{Vector{Vector{Int}}})

@test length(partitions(0)) == 1
@test length(partitions(-1)) == 0
@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))

# integer_partitions
@test integer_partitions(0) == []
@test integer_partitions(5) == Any[[1, 1, 1, 1, 1], [2, 1, 1, 1], [2, 2, 1], [3, 1, 1], [3, 2], [4, 1], [5]]
@test_throws DomainError integer_partitions(-1)

@test_throws ArgumentError prevprod([2,3,5],Int128(typemax(Int))+1)
@test prevprod([2,3,5],30) == 30
@test prevprod([2,3,5],33) == 32

# noncrossing partitions
let nc4 = ncpartitions(4)
    @test nc4 == Any[Any[[1],[2],[3],[4]], Any[[1],[2],[3,4]], Any[[1],[2,3],[4]], Any[[1],[2,4],[3]], Any[[1],[2,3,4]],
                     Any[[1,2],[3],[4]], Any[[1,2],[3,4]],
                     Any[[1,3],[2],[4]], Any[[1,4],[2],[3]], Any[[1,4],[2,3]],
                     Any[[1,2,3],[4]], Any[[1,3,4],[2]], Any[[1,2,4],[3]],
                     Any[[1,2,3,4]]]
    @test length(nc4) == catalannum(4)
end


import Distributions: Uniform
import StaticArrays: SVector
import DynamicalSystems: DiscreteDynamicalSystem

"""
    eom_ikeda(u, p, t)

Equations of motion for a discrete two-dimensional Ikeda map system, adapted from [1]
by adding a noise term and allowing the influences from ``x \\to y`` (``c_{xy}``) and
from ``y \\to x`` (``c_{yx}``) to be adjusted. The difference equations are

```math
\\begin{aligned}
x(t+1) = 1 + \\mu(x(t) \\cos{(\\theta)} - c_{yx} y(t) \\sin{(\\theta)}) - min(\\dfrac{\\sigma \\xi_{t}^{(1)})}{(1-x)}, \\xi_{t}^{(2)} \\\\
y(t+1) = \\mu(y(t) \\cos{(\\theta)} - c_{xy} x(t) \\sin{(\\theta)}) - min(\\dfrac{\\sigma \\zeta_{t}^{(1)})}{(1-y)}, \\zeta_{t}^{(2)}
\\end{aligned}
```

## References

1. Cao, Liangyue, Alistair Mees, and Kevin Judd. "Modeling and predicting
    non-stationary time series." International Journal of Bifurcation and
    Chaos 7.08 (1997): 1823-1831.
"""
function eom_ikeda(u, p, t)
    x, y = u[1], u[2]
    c_xy, c_yx, a, b, c, r₁, r₂, σ = p[1], p[2], p[3], p[4], p[5], p[6], p[7], p[8]

    θ = a - b/(c + x^2 + y^2)
    μ = r₁*sin(t) - r₂
    d = Uniform(0.1, 0.4)

    dx = 1 + μ*(x*cos(θ) - c_yx*y*sin(θ)) - min(σ*rand(d)/(1-x), rand(d))
    dy = μ*(y*cos(θ) + c_xy*x*sin(θ)) -  min(σ*rand(d)/(1-y), rand(d))

    SVector{2}(dx, dy)
end

"""
    ikeda(; u₀ = rand(2), c_xy = 1.0, c_yx = 1.0, a = 0.8, b = 12, c = 0.9,
        r₁ = rand(Uniform(0.01, 0.3)), r₂ = rand(Uniform(0.01, 0.3)), σ = 0.05)

Initialise a discrete two-dimensional Ikeda map system, adapted from [1]
by adding a noise term and allowing the influences from ``x \\to y`` (``c_{xy}``) and
from ``y \\to x`` (``c_{yx}``) to be adjusted.

As a rule-of-thumb, if parameters `a`, `b`, and `c` are drawn from uniform
distributions on `[0.8, 1.5]`, `[10, 14]` and `[0.1, 0.9]`.

The difference equations are

```math
\\begin{aligned}
x(t+1) = 1 + \\mu(x \\cos{(\\theta)} - c_{yx} y \\sin{(\\theta)}) - min(\\dfrac{\\sigma \\xi_{t}^{(1)})}{(1-x)}, \\xi_{t}^{(2)} \\\\
y(t+1) = \\mu(y \\cos{(\\theta)} - c_{xy} x \\sin{(\\theta)}) - min(\\dfrac{\\sigma \\zeta_{t}^{(1)})}{(1-y)}, \\zeta_{t}^{(2)}
\\end{aligned}
```

## References

1. Cao, Liangyue, Alistair Mees, and Kevin Judd. "Modeling and predicting
    non-stationary time series." International Journal of Bifurcation and
    Chaos 7.08 (1997): 1823-1831.
"""
function ikeda(u₀, c_xy, c_yx, a, b, c, r₁, r₂, σ)
    p = [c_xy, c_yx, a, b, c, r₁, r₂, σ]
    DiscreteDynamicalSystem(eom_ikeda, u₀, p)
end

function ikeda(; u₀ = rand(2), c_xy = 1.0, c_yx = 1.0, a = 0.8, b = 12, c = 0.9,
        r₁ = rand(Uniform(0.01, 0.3)), r₂ = rand(Uniform(0.01, 0.3)), σ = 0.05)
    ikeda(u₀, c_xy, c_yx, a, b, c, r₁, r₂, σ)
end

export eom_ikeda, ikeda


using ExamplePackage_rosa
using Test, Polynomials, Random

@testset "ExamplePackage_rosa.jl" begin
    testfn(x)= (x-1)^3
    testfn_prime(x)=3*(x-1)^2
    testfn1(x)= exp(x)-1
    testfn_prime1(x)=exp(x)
    testfn2(x)= atan(x)
    testfn_prime2(x)= 1/((x)^2+1)

    Random.seed!(1234)

    #test for functions with derivative

    root1, conv, iter  = ExamplePackage_rosa.newtonroot(testfn,testfn_prime,x_0=0.0)
    root2, conv, iter  = ExamplePackage_rosa.newtonroot(testfn1,testfn_prime1,x_0=1.0)
    root3, conv, iter  = ExamplePackage_rosa.newtonroot(testfn1,testfn_prime1,x_0=-0.3)

    @test abs(root1-1.0)<1e-6
    @test abs(root2 - 0.0)<1e-6
    @test abs(root3 - 0.0)<1e-6

    #do the same for functions without derivative

    root1, conv, iter  = ExamplePackage_rosa.newtonroot(testfn,x_0=0.0)
    root2, conv, iter  = ExamplePackage_rosa.newtonroot(testfn1,x_0=1.0)
    root3, conv, iter  = ExamplePackage_rosa.newtonroot(testfn1,x_0=-.4)

    @test abs(root1-1.0)<1e-6
    @test abs(root2 - 0.0)<1e-6
    @test abs(root3 - 0.0)<1e-6

    #check maxiter works
    root1, conv, iter  = ExamplePackage_rosa.newtonroot(testfn,testfn_prime,x_0=0.5,  maxiter = 3)
    # if the function converged, return number of iteration, else return nothing
    @test iter == nothing ||iter <= 3
    #check for tolerance, increasing it to 1e-8
    root1, conv, iter  = ExamplePackage_rosa.newtonroot(testfn,testfn_prime,x_0=0.5,  tol = 1e-8)
    @test abs(root1-1.0)<1e-7

    #check lack iof convergence

    test_nc(x) = x^2+1

    rootnc, conv, iter  = ExamplePackage_rosa.newtonroot(test_nc,x_0=1.0)
    @test rootnc==nothing

    #check bigfloot
    x0 = BigFloat(0.53)
    root1, conv, iter  = ExamplePackage_rosa.newtonroot(testfn,x_0=x0)

    @test abs(root1-1.0)<1e-6

    #check for the polynomials roots

    p_1 = Polynomial([2, -5, 2])
    p_2 = Polynomial(rand(4))
    p_3 = convert(Polynomial,ChebyshevT(rand(4)))

    root1, conv, iter  = ExamplePackage_rosa.newtonroot(p_1,x_0=1.0)
    root2, conv, iter  = ExamplePackage_rosa.newtonroot(p_2,x_0=1.0)
    root3, conv, iter  = ExamplePackage_rosa.newtonroot(p_3,x_0=1.0)

    @test minimum(abs.(roots(p_1).-root1))<=1e-6
    @test minimum(abs.(roots(p_2).-root2))<=1e-6
    @test minimum(abs.(roots(p_3).-root3))<=1e-6

end


using FastGaussQuadrature
using Test, LinearAlgebra, Random, SpecialFunctions

include("test_gausschebyshev.jl")
include("test_gausslegendre.jl")
include("test_gaussjacobi.jl")
include("test_gaussradau.jl")
include("test_gausslobatto.jl")
include("test_besselroots.jl")
include("test_gausslaguerre.jl")
include("test_gausshermite.jl")


module Interpolation

export polyint

"""    polyint(xs, ys[, n])

Returns a function that performs polynomial interpolation across the `xs` and
`ys`.

`n`, if given, is the order of the interpolating polynomial; if the length of
the `xs` and `ys` is greater than `n+1`, then multiple interpolants are
returned.
"""
function polyint(xs, ys)
    n = length(xs)
    polyint(xs, ys, n-1)
end

function polyint(xs::Array{T, 1}, ys::Array{T, 1}, n) where T
    triangle = Array{T,1}[ys]
    for i in 2:n+1
        push!(triangle, zeros(T, length(triangle[i-1])-1))
    end

    full =  n == length(xs)-1

    function interpolate(x)
        for i in 2:n+1
            k = length(triangle[i])
            for j in 1:k
                triangle[i][j] = (x-xs[j+i-1])*triangle[i-1][j]/(xs[j]-xs[j+i-1]) + (x-xs[j])*triangle[i-1][j+1]/(xs[j+i-1]-xs[j])
            end
        end

        if full
            triangle[end][1]
        else
            Tuple(triangle[end])
        end
    end
end

end # module


import Pkg;Pkg.add("DifferentialEquations")
import Pkg;Pkg.add("Plots")
using DifferentialEquations
using Plots;gr()

# definition ode
f(y, t) = 3.0 * y + 2.0

# setting initial condiction
y0 = 1.0

# setting time step
tspan = (0.0, 1.0)

# solve of ODEProblem
prob = ODEProblem(f , y0 , tspan)
sol = solve(prob)

# plot
plot(
    sol,
    linewidth=5,
    title="Solution to the linear ODE with a thick line",
    xaxis="Time (t)",
    yaxis="y(t) (in micro.m)",
    label="My Thick Line!"
)

# exact solution
g(t) = 5//3 * exp.(3 * t) - 2//3

# comparison plot
plot!(g, lw=3, ls=:dash, label="True Solution!")


include("Utils.jl")

"""
File: MMAFunctions.jl
Authors: Reese Pathak, Stephen Boyd

This file defines a set of scalar-valued functions that take one or two vectors
as input and output a scalar as output. We support average, standard deviation,
root mean square, angle, distance, and correlation coefficient.
"""

module Functions
  using Utils
export avg, std, rms, angle, corrcoef, dist

# FUNCTION DEFINITIONS.

# Average
@mma_function avg(x) = mean(x)

# RMS value
@mma_function rms(x) = norm(x)/sqrt(length(x))

# Distance
@mma_function dist(x, y) = norm(x - y)

# Standard deviation
@mma_function std(x) = norm(demean(x))/sqrt(length(x))

# Angle
@mma_function angle(x, y) = acos(dot(x,y)/(norm(x)*norm(y)))

# Correlation coefficient
@mma_function corrcoef(x, y) = dot(demean(x), demean(y))/norm(demean(x))*norm(demean(y))

end


###################################
## Simple Variational Inference
## for 2D Gaussian

using PyPlot

function calc_KL(mu1, lambda1, mu2, lambda2)
    D = size(mu1, 1)
    px_lnqx = 0.5 * logdet(lambda2) - 0.5 * ((mu1 - mu2)' * lambda2 * (mu1 - mu2) + trace(lambda2 * inv(lambda1)))
    px_lnpx = 0.5 * logdet(lambda1) - 0.5 * D
    KL = - (px_lnqx - px_lnpx)
    return KL[1]
end

function plot_results(result, truth)
    N = size(result, 1)
    H = Int(ceil(sqrt(N)))
    W = Int(ceil(N / H))
    f, ax = subplots(H, W)
    for i in 1 : H
        for j in 1 : W
            n = (i - 1) * W + j
            if n <= N
                p = ax[i, j]
                p[:set_title](@sprintf("%d of %d", n, N))
                plot_gaussian(p, result[n][1], result[n][2], "b", "p(z)")
                plot_gaussian(p, truth[1], truth[2], "r", "q(z)")
                p[:set_xlim]([-1.5, 1.0])
                p[:set_ylim]([-0.6, 0.8])
            end
        end
    end
end

function plot_gaussian(p, Mu, Sigma, col, label)
    res = 100
    p[:plot](Mu[1], Mu[2], "x", color=col)

    F = eigfact(Sigma)
    vec = F[:vectors]
    val = F[:values]
    dw = 2*pi/res
    w = dw * (0 : res)

    c = 1.0
    a = sqrt(c * val[1])
    b = sqrt(c * val[2])
    P1 = a*cos(w)
    P2 = b*sin(w)
    P = Mu .+ vec'*vcat(P1', P2')
    p[:plot](P[1, :]', P[2, :]', "-", color=col, label=label)
end

function main()
    ## creat truth distribution
    D = 2 # dimension
    theta = 2.0*pi/12 # tilt
    A = reshape([cos(theta), -sin(theta),
                 sin(theta), cos(theta)],
                2, 2)
    mu = [0.0, 0.0]
    lambda = inv(A * inv(reshape([1,0,0,10], 2, 2)) * A')

    ## initialize
    mu_h = randn(D)
    lambda_h = zeros(D,D)

    ## main iteration
    max_iter = 10
    KL = Vector{Float64}()
    result = Vector()
    for i in 1 : max_iter
        ## update
        mu_h[1] = mu[1] - inv(lambda[1,1])*lambda[1,2] * (mu_h[2] - mu[2])
        lambda_h[1,1] = lambda[1,1]
        mu_h[2] = mu[2] - inv(lambda[2,2])*lambda[2,1] * (mu_h[1] - mu[1])
        lambda_h[2,2] = lambda[2,2]

        ## calculate KL divergeince
        push!(KL, calc_KL(mu_h, lambda_h, mu, lambda))
        #println(mu_h)
        #println(lambda_h)

        ## store the results
        push!(result, (deepcopy(mu_h), deepcopy(inv(lambda_h))))
    end

    ## visualize results
    plot_results(result, (mu, inv(lambda)))

    f, ax = subplots(1, 1)
    plot_gaussian(ax, mu, inv(lambda), "b", latexstring("\$p(z)\$"))
    plot_gaussian(ax, result[end][1], result[end][2], "r", latexstring("\$q(z)\$"))
    ax[:set_xlabel](latexstring("\$z_1\$"), fontsize=16)
    ax[:set_ylabel](latexstring("\$z_2\$"), fontsize=16)
    ax[:legend](fontsize=16)

    ## KL divergence
    f, ax = subplots(1)
    ax[:plot](1:max_iter, KL)
    ax[:set_ylabel]("KL divergence")
    ax[:set_xlabel]("iteration")
    show()
end

main()


# Bregman divergence

"""
Implements the Bregman divergence, a friendly introduction to which can be found
[here](http://mark.reid.name/blog/meet-the-bregman-divergences.html).
Bregman divergences are a minimal implementation of the "mean-minimizer" property.

It is assumed that the (convex differentiable) function F maps vectors (of any type or size) to real numbers.
The inner product used is `Base.dot`, but one can be passed in either by defining `inner` or by
passing in a keyword argument. If an analytic gradient isn't available, Julia offers a suite
of good automatic differentiation packages.

function evaluate(dist::Bregman, p::AbstractVector, q::AbstractVector)
"""
struct Bregman{T1 <: Function, T2 <: Function, T3 <: Function} <: PreMetric
    F::T1
    ∇::T2
    inner::T3
end

# Default costructor.
Bregman(F, ∇) =  Bregman(F, ∇, LinearAlgebra.dot)

# Evaluation fuction
function (dist::Bregman)(p, q)
    # Create cache vals.
    FP_val = dist.F(p)
    FQ_val = dist.F(q)
    DQ_val = dist.∇(q)
    p_size = length(p)
    # Check F codomain.
    if !(isa(FP_val, Real) && isa(FQ_val, Real))
        throw(ArgumentError("F Codomain Error: F doesn't map the vectors to real numbers"))
    end
    # Check vector size.
    if p_size != length(q)
        throw(DimensionMismatch("The vector p ($(size(p))) and q ($(size(q))) are different sizes."))
    end
    # Check gradient size.
    if length(DQ_val) != p_size
        throw(DimensionMismatch("The gradient result is not the same size as p and q"))
    end
    # Return the Bregman divergence.
    return FP_val - FQ_val - dist.inner(DQ_val, p .- q)
end

# Convenience function.
bregman(F, ∇, x, y; inner = LinearAlgebra.dot) = Bregman(F, ∇, inner)(x, y)


#=
Given an array of integers, return a new array such that each element at index i of the new array is the product of all the numbers in the original array except the one at i.

For example, if our input was [1, 2, 3, 4, 5], the expected output would be [120, 60, 40, 30, 24]. If our input was [3, 2, 1], the expected output would be [2, 3, 6].

Follow-up: what if you can't use division?
=#

using Test
include("Solutions/problem2_product_except_i.jl")
@test product_excepti([1, 2, 3, 4, 5]) == [120, 60, 40, 30, 24]
@test product_excepti([3, 2, 1]) == [2, 3, 6]


###Target functions
function sincos!(r::Matrix,paras::Vector,t::AbstractVector)
    p1 = 2pi*paras[1]
    p2 = 2pi*paras[2]
    @simd for i=1:length(t)
        @inbounds r[i,1] = sin(p1*t[i])
    end
    @simd for i=1:length(t)
        @inbounds r[i,2] = cos(p2*t[i])
    end
    r
end

sincos(paras::Vector,t::AbstractVector) = sincos!(zeros(eltype(t),length(t),2),paras,t)

@inline function _commonsincoscomponents(t::AbstractVector,measurementparas::Vector,variance::Vector,paraminit...)
    modelparameters = parameters([:a,:b],paraminit...)
    noisemodel = noise(:gaussian,variance)
    measurementdata = data(:array,t,applynoise!(noisemodel,sincos(measurementparas,t)))
    modelparameters,measurementdata,noisemodel
end

###Create a target function model
function sincosmodel(t::AbstractVector,measurementparas::Vector,variance::Vector,paraminit...)
    model(:target,_commonsincoscomponents(t,measurementparas,variance,paraminit...)...,sincos,t;name="Sine-Cosine Model")
end

###Create a target function model
function sincosmodel!(t::AbstractVector,measurementparas::Vector,variance::Vector,paraminit...)
    model(:target!,_commonsincoscomponents(t,measurementparas,variance,paraminit...)...,sincos!,t;name="Sine-Cosine In-Place Model")
end


# ---
# title: 780. Reaching Points
# id: problem780
# author: Indigo
# date: 2021-01-28
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/reaching-points/description/>
# hidden: true
# ---
#
# A move consists of taking a point `(x, y)` and transforming it to either `(x,
# x+y)` or `(x+y, y)`.
#
# Given a starting point `(sx, sy)` and a target point `(tx, ty)`, return `True`
# if and only if a sequence of moves exists to transform the point `(sx, sy)` to
# `(tx, ty)`. Otherwise, return `False`.
#
#
#
#     **Examples:**
#     Input: sx = 1, sy = 1, tx = 3, ty = 5
#     Output: True
#     Explanation:
#     One series of moves that transforms the starting point to the target is:
#     (1, 1) -> (1, 2)
#     (1, 2) -> (3, 2)
#     (3, 2) -> (3, 5)
#
#     Input: sx = 1, sy = 1, tx = 2, ty = 2
#     Output: False
#
#     Input: sx = 1, sy = 1, tx = 1, ty = 1
#     Output: True
#
#
#
# **Note:**
#
#   * `sx, sy, tx, ty` will all be integers in the range `[1, 10^9]`.
#
#
## @lc code=start
using LeetCode

function reaching_points(sx, sy, tx, ty)
    while tx >= sx && ty >= sy
        if tx == ty
            break
        end
        if tx > ty
            if ty > sy
                tx %= ty;
            else
                return (tx - sx) % ty == 0;
            end
        else
            if tx > sx
                ty %= tx;
            else
                return (ty - sy) % tx == 0;
            end
        end
    end
    tx == sx && ty == sy
end
## @lc code=end


"""
    cilc_weights(cij, a, b)

This function returns weights (a vector of the number of frequency channels) of the constrained ILC (CILC) method.

*Reference*: Equation (20) of Remazeilles, Delabrouille, Cardoso, MNRAS, 410, 2481 (2011)

# Arguments
- `cij::Array{<:AbstractFloat,2}`: symmetric covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `a::Array{<:AbstractFloat,1}`: vector of the frequency response, for the component to be extracted. E.g., `a = [1,...,1]` for CMB.
- `b::Array{<:AbstractFloat,1}`: vector of the frequency response, for the component to be nulled.
"""
function cilc_weights(
    cij::Array{T,2},
    a::Array{T,1},
    b::Array{T,1},
) where {T<:AbstractFloat}
    if size(cij)[1] ≠ size(cij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(cij)[1] ≠ length(a) || size(cij)[1] ≠ length(b)
        throw(DimensionMismatch("dimensions of the covariance matrix and the frequency response vector do not match"))
    elseif a ≈ b
        throw(ErrorException("vectors of the frequency response are too similar"))
    else
        M = Symmetric(cij)
    end
    x = M \ a # R^-1 a
    y = M \ b # R^-1 b
    w = (b'y * x - a'y * y) / (a'x * b'y - (a'y)^2) # ILC weights
end

"""
    cilc_weights(cijℓ, a, b[, ℓid=3])

This function returns weights (a `nν`-by-`nℓ` matrix) of the constrained ILC (CILC) method.

Here, `nν` is the number of frequency channels and `nℓ` is the number of elements in the relevant domain, e.g., multipoles, band-power bins, pixels, etc.

*Reference*: Equation (20) of Remazeilles, Delabrouille, Cardoso, MNRAS, 410, 2481 (2011)

# Arguments
- `cijℓ::Array{<:AbstractFloat,3}`: symmetric covariance matrix with the dimention of `(nℓ, nν, nν)`, `(nν, nℓ, nν)` or `(nν, nν, nℓ)` (default).
- `a::Array{<:AbstractFloat,1}`: vector of the frequency response, for the component to be extracted. E.g., `a = [1,...,1]` for CMB.
- `b::Array{<:AbstractFloat,1}`: vector of the frequency response, for the component to be nulled.

# Optional Arguments
- `ℓid::Integer=3`: location of the index for the `nℓ` domain. `ℓid=1` if `cijℓ[nℓ,nν,nν]`, `ℓid=2` if `cijℓ[nν,nℓ,nν]`, and `ℓid=3` (the default value) if `cijℓ[nν,nν,nℓ]`.
"""
function cilc_weights(
    cijℓ::Array{T,3},
    a::Array{T,1},
    b::Array{T,1},
    ℓid::Integer = 3,
) where {T<:AbstractFloat}
    if ℓid > 3 || ℓid < 1
        throw(DomainError(ℓid, "ℓid must be 1, 2, or 3"))
    end
    if (ℓid == 3 && size(cijℓ)[1] ≠ size(cijℓ)[2]) ||
       (ℓid == 2 && size(cijℓ)[1] ≠ size(cijℓ)[3]) ||
       (ℓid == 1 && size(cijℓ)[2] ≠ size(cijℓ)[3])
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    end
    nℓ = size(cijℓ)[ℓid]
    nν = ifelse(ℓid == 3, size(cijℓ)[1], size(cijℓ)[3])
    if length(a) ≠ nν || length(b) ≠ nν
        throw(DimensionMismatch("dimensions of the covariance matrix and the frequency response vector do not match"))
    elseif a ≈ b
        throw(ErrorException("vectors of the frequency response are too similar"))
    else
    end
    wℓ = zeros(nν, nℓ) # ILC weights
    for iℓ = 1:nℓ
        if ℓid == 3
            M = Symmetric(cijℓ[:, :, iℓ])
        elseif ℓid == 2
            M = Symmetric(cijℓ[:, iℓ, :])
        else
            M = Symmetric(cijℓ[iℓ, :, :])
        end
        x = M \ a # R^-1 a
        y = M \ b # R^-1 b
        wℓ[:, iℓ] = (b'y * x - a'y * y) / (a'x * b'y - (a'y)^2)
    end
    return wℓ
end


mutable struct Tikhonov{T} <: AbstractRegularization{T}
    α::T
end
Tikhonov() =  Tikhonov

get_regularization_matrix(regop::Tikhonov{T},N,ds) where {T} =   Matrix(I(N))

mutable struct SecondDerivative{T} <: AbstractRegularization{T}
    α::T
end
SecondDerivative()=SecondDerivative

function get_regularization_matrix(regop::SecondDerivative{T},N,ds) where {T}
    α=regop.α
    L=zeros( N+2,N)
    for j in 3:N
        L[j,j]= 1/(ds[j]*(ds[j]+ds[j-1]))#
        L[j,j-1]= -2*1/(ds[j-1]*(ds[j]))#
        L[j,j-2]= 1/(ds[j-1]*(ds[j]+ds[j-1]))#
    end

    # if regop.lower_bc
         L[1,1]=1/ds[1]
         L[2,1]=-2*1/(ds[1]*(ds[2]))
         L[2,2]=1/ds[2]*(ds[1]+ds[2])
    # end
    # if regop.upper_bc
         L[N+1,N-1]=1/ds[N-1]*(ds[N]+ds[N-1])
         L[N+1,N]=-2*1/((ds[N]*ds[N-1]))
         L[N+2,N]=1/ds[N]
    # end
    return L
    #L[:,end] .=0  #Do not regularize on the y offset
end


function nlp_cvx_106_010(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    # Test Goals:
    # - cos and sin expressions
    # - convex constraint, only in given domains
    # Variants
    #   010 - intersection point
    #   011 - intersection point

    model = Model(optimizer)

    @variable(model, -3 <= x <= 3)
    @variable(model, -1 <= y <= 1)

    @objective(model, Min, -x - y)
    @NLconstraint(model, sin(-x - 1.0) + x / 2 + 0.5 <= y)
    @NLconstraint(model, cos(x - 0.5) + x / 4 - 0.5 >= y)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, -1.8572155128552428, tol = objective_tol)
    return check_solution(
        [x, y],
        [1.369771397576555, 0.4874441152786876],
        tol = primal_tol,
    )
end


# Matern 1/2 isotropic covariance Function

"""
    Mat12Iso <: MaternISO

Isotropic Matern 1/2 kernel (covariance)
```math
k(x,x') = σ^2 \\exp(-|x-y|/ℓ)
```
with length scale ``ℓ`` and signal standard deviation ``σ``.
"""
mutable struct Mat12Iso{T} <: MaternIso where {T<:Real}
    "Length scale"
    ℓ::T
    "Signal variance"
    σ2::T
    "Priors for kernel parameters"
    priors::Array
end

"""
Matern 1/2 isotropic covariance Function

    Mat12Iso(ll::T, lσ::T)

# Arguments
  - `ll::Real`: length scale (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
"""
Mat12Iso(ll::T, lσ::T) where T = Mat12Iso{T}(exp(ll), exp(2 * lσ), [])

function set_params!(mat::Mat12Iso, hyp::AbstractVector)
    length(hyp) == 2 || throw(ArgumentError("Matern 1/2 covariance function has two parameters, received $(length(hyp))."))
    mat.ℓ, mat.σ2 = exp(hyp[1]), exp(2 * hyp[2])
end

get_params(mat::Mat12Iso{T}) where T = T[log(mat.ℓ), log(mat.σ2) / 2]
get_param_names(mat::Mat12Iso) = [:ll, :lσ]
num_params(mat::Mat12Iso) = 2

cov(mat::Mat12Iso, r::Number) = mat.σ2 * exp(-r / mat.ℓ)

@inline dk_dll(mat::Mat12Iso, r::Real) = r / mat.ℓ * cov(mat, r)


using YaoBase

export ShiftGate, shift

"""
    ShiftGate <: PrimitiveBlock

Phase shift gate.

# Definition

```math
\\begin{pmatrix}
1 & 0
0 & e^(im θ)
\\end{pmatrix}
```
"""
mutable struct ShiftGate{T <: Real} <: PrimitiveBlock{1}
    theta::T
end

"""
    shift(θ)

Create a [`ShiftGate`](@ref) with phase `θ`.

# Example

```jldoctest
julia> shift(0.1)
shift(0.1)
```
"""
shift(θ::AbstractFloat) = ShiftGate(θ)
shift(θ::Real) = shift(Float64(θ))
mat(::Type{T}, gate::ShiftGate) where {T <: Complex} = Diagonal(T[1.0, exp(im * gate.theta)])

cache_key(gate::ShiftGate) = gate.theta

# parametric interface
niparams(::Type{<:ShiftGate}) = 1
getiparams(x::ShiftGate) = x.theta
setiparams!(r::ShiftGate, param::Real) = (r.theta = param; r)

Base.adjoint(blk::ShiftGate) = ShiftGate(-blk.theta)
Base.copy(block::ShiftGate{T}) where T = ShiftGate{T}(block.theta)
Base.:(==)(lhs::ShiftGate, rhs::ShiftGate) = lhs.theta == rhs.theta
YaoBase.isunitary(r::ShiftGate) = true


# indicator of the L2 norm sphere with given radius

"""
  IndSphereL2(r::Real=1.0)

Returns the function `g = ind{x : ||x|| = r}`, for a real parameter `r > 0`.
"""

immutable IndSphereL2{R <: Real} <: IndicatorNonconvex
  r::R
  function IndSphereL2(r::R)
    if r <= 0
      error("parameter r must be positive")
    else
      new(r)
    end
  end
end

IndSphereL2{R <: Real}(r::R=1.0) = IndSphereL2{R}(r)

function (f::IndSphereL2){T <: RealOrComplex}(x::AbstractArray{T})
  if abs(vecnorm(x) - f.r)/f.r > 1e-14
    return +Inf
  end
  return 0.0
end

function prox!{T <: RealOrComplex}(f::IndSphereL2, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  normx = vecnorm(x)
  if normx > 0 # zero-zero?
    scal = f.r/normx
    for k in eachindex(x)
      y[k] = scal*x[k]
    end
  else
    normy = 0.0
    for k in eachindex(x)
      y[k] = randn()
      normy += y[k]*y[k]
    end
    normy = sqrt(normy)
    y[:] *= f.r/normy
  end
  return 0.0
end

fun_name(f::IndSphereL2) = "indicator of an L2 norm sphere"
fun_dom(f::IndSphereL2) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndSphereL2) = "x ↦ 0 if ||x|| = r, +∞ otherwise"
fun_params(f::IndSphereL2) = "r = $(f.r)"

function prox_naive{T <: RealOrComplex}(f::IndSphereL2, x::AbstractArray{T}, gamma::Real=1.0)
  normx = vecnorm(x)
  if normx > 0
    y = x*f.r/normx
  else
    y = randn(size(x))
    y *= f.r/vecnorm(y)
  end
  return y, 0.0
end