# S&P500指数の値動きからSPXLを計算するためのデータを作成する
using CSV, DataFrames, Dates, Plots
pyplot(fmt=:svg, fontfamily="meiryo") # plotの設定

# 正解値SPXL読み込み
dfmt = dateformat"yyyy/mm/dd HH:MM:SS"
spxl = DataFrame(CSV.File("./data/SPXL_20210430.csv", dateformat=dfmt, header=true))
dt = spxl.timestamp .- Dates.Hour(14) # Yahooで取得をしたデータは日本時間なので，USに変換する
spxl.timestamp = dt

# S&P500のDividendYield(配当利回り)の読み込み--------------------------------
dfmt = dateformat"yyyy/mm/dd"
sp500_yield = DataFrame(CSV.File("./data/S&P500_DividendYield.csv", dateformat=dfmt, header=true))
sort!(sp500_yield) # 昇順にソートsort

# SPYのdividendsは四半期単位なので，サンプリングをSPXLに合わせて前値ホールドする
spy_dint = copy(spxl)
select!(spy_dint, Not(:open))
rename!(spy_dint, "low" => "yield") # その年の配当金利回り
rename!(spy_dint, "high" => "dtb3") # 米国3ヶ月債金利
rename!(spy_dint, "close" => "spxl") # SPXL価格
rename!(spy_dint, "volume" => "spxl_pred") # 予測SPXL

d_temp = 1
for y in 1:length(sp500_yield.yield)
    println("y = " * string(y))
    for d in d_temp:length(spy_dint.timestamp)
        # println("d = " * string(d))
        yield = sp500_yield[y,"yield"] # 配当利回りを取得
        spy_dint.yield[d] = yield # 次のyearに該当するまで，前値を当てはめる
        if (Dates.Date(spy_dint.timestamp[d]) > sp500_yield.date[y])
            d_temp = d
            break
        end
    end
end

# DTB3(金利)の読み込み--------------------------------
dfmt = dateformat"yyyy/mm/dd"
dtb3 = DataFrame(CSV.File("./data/DTB3.csv", dateformat=dfmt, header=true))
dfmt = dateformat"yyyy-mm-dd"
dt = [DateTime(x, dfmt) for x in dtb3[:,"DATE"]]
dtb3.timestamp = dt
select!(dtb3, Not(:DATE))

# 値を入れていく
for y in 1:length(spy_dint.timestamp)
    for d in 1:length(dtb3.timestamp)
        if (Date(spy_dint.timestamp[y]) == Date(dtb3.timestamp[d]))
            spy_dint.dtb3[y] = dtb3.DTB3[d]
            # println("d = " * string(d))
            break
        end
        spy_dint.dtb3[y] = NaN
    end
end
spy_dint = spy_dint[ .!(isnan.(spy_dint.dtb3)), :] # NaNのある行の削除

# S&P500の読み込み--------------------------------
dfmt = dateformat"yyyy/mm/dd HH:MM:SS"
sp500 = DataFrame(CSV.File("./data/S&P500.csv", dateformat=dfmt, header=true))
dt = sp500.timestamp .- Dates.Hour(14) # Yahooで取得をしたデータは日本時間なので，USに変換する
sp500.timestamp = dt
select!(sp500, Not(:open))
select!(sp500, Not(:low))
select!(sp500, Not(:high))
select!(sp500, Not(:volume))
rename!(sp500, "close" => "sp500") # SPY価格

spy_dint = innerjoin(spy_dint, sp500, on=:timestamp)

# 疑似SPXLを計算する元データを保存
CSV.write("./data/spxl_prediction_data.csv", spy_dint)

# 疑似SPXLを計算-----------------------------------------------
# 1. 単純計算が実際のSPXLより大きいことを確認
dif = (spy_dint.sp500[2:end] - spy_dint.sp500[1:end - 1]) # 1日の値動き幅
dif_rate = dif ./ spy_dint.sp500[1:end - 1] # 値動き割合
dif_rate_y = dif_rate + spy_dint.yield[1:end - 1] ./ 365 # 配当込み1日値動き割合
dif_rate3 = dif_rate_y * 3
spxl_pred1 = ones(length(spy_dint.sp500), 1) * 4.25
for i = 2:length(spxl_pred1)
    spxl_pred1[i] = spxl_pred1[i - 1] * dif_rate3[i - 1] + spxl_pred1[i - 1]
end

plot(spy_dint.timestamp, spy_dint.spxl,
    xlabel="年月日", ylabel="価格[\$]", label="SPXL")
plot!(spy_dint.timestamp, spxl_pred1, label="SPXL(S&P500値動き3倍)")
savefig("SPXL_dif_rate3.png")

# 2. 0.95% 考慮でも実際のSPXLより大きいことを確認
dif_rate3_c095 = dif_rate3 .- (0.95 / 100 / 365)
spxl_pred2 = ones(length(spy_dint.sp500), 1) * 4.25
for i = 2:length(spxl_pred2)
    spxl_pred2[i] = spxl_pred2[i - 1] * dif_rate3_c095[i - 1] + spxl_pred2[i - 1]
end

plot(spy_dint.timestamp, spy_dint.spxl,
    xlabel="年月日", ylabel="価格[\$]", label="SPXL")
plot!(spy_dint.timestamp, spxl_pred2, label="SPXL(S&P500値動き3倍+コスト0.95%)")
savefig("SPXL_dif_rate3_c095.png")

spxl_err2 = (spxl_pred2 .- spy_dint.spxl) ./ spy_dint.spxl
plot(spy_dint.timestamp, spxl_err2,
    xlabel="年月日", ylabel="SPXL予測の乖離率",
    ylims=(-0.2, 0.4))
savefig("SPXL_dif_rate3_c095_err.png")

# 3. 計算式に割り当てる
dif_rate3_dtb3 = dif_rate3 - ( spy_dint.dtb3[1:end - 1] / 100 / 365 * 2)
dif_rate3_dtb3_c215 = dif_rate3_dtb3 .- (2.15 / 100 / 365)
spxl_pred3 = ones(length(spy_dint.sp500), 1) * 4.25
for i = 2:length(spxl_pred3)
    spxl_pred3[i] = spxl_pred3[i - 1] * dif_rate3_dtb3_c215[i - 1] + spxl_pred3[i - 1]
end

plot(spy_dint.timestamp, spy_dint.spxl,
    xlabel="年月日", ylabel="価格[\$]", label="SPXL")
plot!(spy_dint.timestamp, spxl_pred3, label="SPXL(S&P500値動き3倍+先物金利+コスト2.15%)")
savefig("SPXL_dif_rate3_dtb3_c215.png")

spxl_err3 = (spxl_pred3 .- spy_dint.spxl) ./ spy_dint.spxl
plot(spy_dint.timestamp, spxl_err3,
    xlabel="年月日", ylabel="SPXL予測の乖離率",
    ylims=(-0.2, 0.4))
savefig("SPXL_dif_rate3_dtb3_c215_err.png")

# 3.2 固定コストを上げる
dif_rate3_dtb3_c315 = dif_rate3_dtb3 .- (3.15 / 100 / 365)
spxl_pred4 = ones(length(spy_dint.sp500), 1) * 4.25
for i = 2:length(spxl_pred4)
    spxl_pred4[i] = spxl_pred4[i - 1] * dif_rate3_dtb3_c315[i - 1] + spxl_pred4[i - 1]
end

plot(spy_dint.timestamp, spy_dint.spxl,
    xlabel="年月日", ylabel="価格[\$]", label="SPXL")
plot!(spy_dint.timestamp, spxl_pred4, label="SPXL(S&P500値動き3倍+先物金利+コスト3.15%)")
savefig("SPXL_dif_rate3_dtb3_c315.png")

spxl_err4 = (spxl_pred4 .- spy_dint.spxl) ./ spy_dint.spxl
plot(spy_dint.timestamp, spxl_err4,
    xlabel="年月日", ylabel="SPXL予測の乖離率",
    ylims=(-0.2, 0.4))
savefig("SPXL_dif_rate3_dtb3_c315_err.png")

# 1954以降の疑似SPXLとS&P 500を計算-----------------------------------------------
# SP500, SP500利回り，米国3ヶ月債金利を用意
spxl_pred = sp500[sp500.timestamp.>Dates.Date("1954-1-4"), :]
spxl_pred.timestamp = Dates.Date.(spxl_pred.timestamp)
dtb3.timestamp = Dates.Date.(dtb3.timestamp)
spxl_pred = innerjoin(spxl_pred, dtb3, on=:timestamp)

spxl_pred.sp500_yield = ones(length(spxl_pred.timestamp))
d_temp = 1
for y in 1:length(sp500_yield.yield)
    println("y = " * string(y))
    for d in d_temp:length(spxl_pred.timestamp)
        # println("d = " * string(d))
        yield = sp500_yield[y,"yield"] # 配当利回りを取得
        spxl_pred.sp500_yield[d] = yield # 次のyearに該当するまで，前値を当てはめる
        if (Dates.Date(spxl_pred.timestamp[d]) > sp500_yield.date[y])
            d_temp = d
            break
        end
    end
end

# 1954年のSP500と同一金額からスタートした場合のSPXLを計算
spxl_pred.spxl = ones(length(spxl_pred.timestamp))*spxl_pred.sp500[1]

dif = (spxl_pred.sp500[2:end] - spxl_pred.sp500[1:end - 1]) # 1日の値動き幅
dif_rate = dif ./ spxl_pred.sp500[1:end - 1] # 値動き割合
dif_rate_y = dif_rate + spxl_pred.sp500_yield[1:end - 1] ./ 365 # 配当込み1日値動き割合
dif_rate3 = dif_rate_y * 3 # 配当込値動き割合の3倍
dif_rate3_dtb3 = dif_rate3 - ( spxl_pred.DTB3[1:end - 1] / 100 / 365 * 2) # 3ヶ月債の2倍の金利を引く
dif_rate3_dtb3_c315 = dif_rate3_dtb3 .- (3.15 / 100 / 365) # 固定コスト3.15%を上乗せ

for i = 2:length(spxl_pred.timestamp)
    spxl_pred.spxl[i] = spxl_pred.spxl[i - 1] * dif_rate3_dtb3_c315[i - 1] + spxl_pred.spxl[i - 1]
end

# 比較対象としてSP500の配当込み指数を計算．
spxl_pred.sp500_yidx = ones(length(spxl_pred.timestamp))*spxl_pred.sp500[1]
dif_rate_y_c = dif_rate_y .- (0.141 / 100 / 365) # 固定コスト0.141%を上乗せ

for i = 2:length(spxl_pred.timestamp)
    spxl_pred.sp500_yidx[i] = spxl_pred.sp500_yidx[i - 1] * dif_rate_y_c[i - 1] + spxl_pred.sp500_yidx[i - 1]
end

CSV.write("spxl_pred_1954.csv", spxl_pred)

# トータルリターンの計算
sp500_ret = spxl_pred.sp500_yidx[spxl_pred.timestamp.==Dates.Date("2020-7-1")]/spxl_pred.sp500_yidx[spxl_pred.timestamp.==Dates.Date("1954-7-1")]
spxl_ret = spxl_pred.spxl[spxl_pred.timestamp.==Dates.Date("2020-7-1")]/spxl_pred.spxl[spxl_pred.timestamp.==Dates.Date("1954-7-1")]

plot(spxl_pred.timestamp, spxl_pred.sp500_yidx,
    xlabel="年月日", ylabel="価格[\$]", label="仮想インデックス(S&P500配当込み指数+コスト0.141%)", yscale=:log10)
plot!(spxl_pred.timestamp, spxl_pred.spxl, label="仮想SPXL(S&P500値動き3倍+先物金利+コスト3.15%)", yscale=:log10)
savefig("SP500_SPXL_1954.png")


abstract type SamplingAlgorithm end

"""
GridSample{T}

T is the step dx for lb:dx:ub
"""
struct GridSample{T} <: SamplingAlgorithm
    dx::T
end

struct UniformSample <: SamplingAlgorithm end
struct SobolSample <: SamplingAlgorithm end
struct LatinHypercubeSample <: SamplingAlgorithm end

"""
LowDiscrepancySample{T}

T is the base for the sequence
"""
struct LowDiscrepancySample{T} <: SamplingAlgorithm
    base::T
end

struct RandomSample <: SamplingAlgorithm end

struct KroneckerSample{A,B} <: SamplingAlgorithm
    alpha::A
    s0::B
end

struct GoldenSample <: SamplingAlgorithm end

"""
sample(n,lb,ub,S::GridSample)

Returns a tuple containing numbers in a grid.

"""
function sample(n,lb,ub,S::GridSample)
    dx = S.dx
    if lb isa Number
        return vec(rand(lb:S.dx:ub,n))
    else
        d = length(lb)
        x = [[rand(lb[j]:dx[j]:ub[j]) for j = 1:d] for i in 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,lb,ub,::UniformRandom)

Returns a Tuple containing uniform random numbers.
"""
function sample(n,lb,ub,::UniformSample)
    if lb isa Number
        return rand(Uniform(lb,ub),n)
    else
        d = length(lb)
        x = [[rand(Uniform(lb[j],ub[j])) for j in 1:d] for i in 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,lb,ub,::SobolSampling)

Returns a Tuple containing Sobol sequences.
"""
function sample(n,lb,ub,::SobolSample)
    s = SobolSeq(lb,ub)
    skip(s,n)
    if lb isa Number
        return [next!(s)[1] for i = 1:n]
    else
        return Tuple.([next!(s) for i = 1:n])
    end
end

"""
sample(n,lb,ub,::LatinHypercube)

Returns a Tuple containing LatinHypercube sequences.
"""
function sample(n,lb,ub,::LatinHypercubeSample)
    d = length(lb)
    if lb isa Number
        x = vec(LHCoptim(n,d,1)[1])
        # x∈[0,n], so affine transform
        return @. (ub-lb) * x/(n) + lb
    else
        lib_out = float(LHCoptim(n,d,1)[1])
        # x∈[0,n], so affine transform column-wise
        @inbounds for c = 1:d
            lib_out[:,c] = (ub[c]-lb[c])*lib_out[:,c]/n .+ lb[c]
        end
        x = [lib_out[i,:] for i = 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,lb,ub,S::LowDiscrepancySample)

Low discrepancy sample:
- Dimension 1: Van der Corput sequence
- Dimension > 1: Halton sequence
If dimension d > 1, all bases must be coprime with each other.
"""
function sample(n,lb,ub,S::LowDiscrepancySample)
    d = length(lb)
    if d == 1
        #Van der Corput
        b = S.base
        x = zeros(Float32,n)
        for i = 1:n
            expansion = digits(i,base = b)
            L = length(expansion)
            val = zero(Float32)
            for k = 1:L
                val += expansion[k]*float(b)^(-(k-1)-1)
            end
            x[i] = val
        end
        # It is always defined on the unit interval, resizing:
        return @. (ub-lb) * x + lb
    else
        #Halton sequence
        x = zeros(Float32,n,d)
        for j = 1:d
            b = S.base[j]
            for i = 1:n
                val = zero(Float32)
                expansion = digits(i, base = b)
                L = length(expansion)
                val = zero(Float32)
                for k = 1:L
                    val += expansion[k]*float(b)^(-(k-1)-1)
                end
                x[i,j] = val
            end
        end
        #Resizing
        # x∈[0,1], so affine transform column-wise
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end

        y = [x[i,:] for i = 1:n]
        return Tuple.(y)
    end
end

"""
sample(n,d,D::Distribution)

Returns a Tuple containing numbers distributed as D
"""
function sample(n,d,D::Distribution)
    if d == 1
        return rand(D,n)
    else
        x = [[rand(D) for j in 1:d] for i in 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,d,K::KroneckerSample)

Returns a Tuple containing numbers following the Kronecker sample
"""
function sample(n,lb,ub,K::KroneckerSample)
    d = length(lb)
    alpha = K.alpha
    s0 = K.s0
    if d == 1
        x = zeros(n)
        @inbounds for i = 1:n
            x[i] = (s0+i*alpha)%1
        end
        return @. (ub-lb) * x + lb
    else
        x = zeros(n,d)
        @inbounds for j = 1:d
            for i = 1:n
                x[i,j] = (s0[j] + i*alpha[j])%i
            end
        end
        #Resizing
        # x∈[0,1], so affine transform column-wise
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end

        y = [x[i,:] for i = 1:n]
        return Tuple.(y)
    end
end

function sample(n,lb,ub,G::GoldenSample)
    d = length(lb)
    if d == 1
        x = zeros(n)
        g = (sqrt(5)+1)/2
        a = 1.0/g
        for i = 1:n
            x[i] = (0.5+a*i)%1
        end
        return @. (ub-lb) * x + lb
    else
        x = zeros(n,d)
        for j = 1:d
            #Approximate solution of x^(d+1) = x + 1, a simple newton is good enough
            y = 2.0
            for s = 1:10
                g = (1+y)^(1/(j+1))
            end
            a = 1.0/g
            for i = 1:n
                x[i,j] = (0.5+a*i)%1
            end
        end
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end
        y = [x[i,:] for i = 1:n]
        return Tuple.(y)
    end
end


# test typed version
using Plots, Revise, VortexDistributions

# include("typeversion.jl")
Exact()

# unitary?
v0 = [.2 .4 1]
w0 = PointVortex(v0)
@test rawData(w0) == v0

v1 = [.2 .4 1;0.7 1.5 -1;-.3 1.2 1]
w1 = PointVortex(v1)
@test rawData(w1) == v1

w3 = randPointVortex(1000)
v3 = rawData(w3)
@test v3 == rawData(w3)

# fast... really, really fast!
@time v3 = rawData(w3)
@time w3 = PointVortex(v3)

# Single vortex tests are reliable and don't require sorting
# (Phase cross talk will cause construction to wander)

@test foundNear(1)
@test foundNear(10)
@test foundNear(500)

# more tests below?
psi,vort = randVortexField(30); heatmap(psi.x,psi.y,abs2.(psi.ψ))

# Real data with masking examples
@load "./test/one_frame.jld2"
x0,y0 = 70,200

# circular
msk(x,y) = sqrt(x^2 + y^2) > 100
@. ψ1[msk(x-x0,y'-y0) ] = false
heatmap(x,y,abs2.(ψ1),transpose=true)
psi3 = Torus(ψ1,x,y)
vort3 = findvortices(psi3)
ψ1

# square
Lx,Ly = 100,100

xm = x[@. abs(x-x0)<Lx/2]
ym = y[@. abs(y-y0)<Ly/2]
psim = @. ψ1[abs(x-x0)<Lx/2,abs(y-y0)<Ly/2]

heatmap(xm,ym,angle.(psim),transpose=true)
psi2 = Torus(psim,xm,ym)
vort = findvortices(psi2)

scatter!([vort[1].xv],[vort[1].yv],alpha=0.4)
scatter!([vort[2].xv],[vort[2].yv],alpha=0.4)


using Test, SpaceExParser, HybridSystems, MathematicalSystems

include("unit_examples.jl")
include("unit_parse.jl")
include("unit_affine.jl")


using MVDA, Random, LinearAlgebra, StableRNGs
using StatsBase, Statistics
using ProgressMeter

# Compute an initial solution w/ regularization. Take this estimate as the ground truth.
function __init__(λ, problem, ϵ, gtol)
    MVDA.init!(MMSVD(), problem, ϵ, λ, verbose=false, maxiter=10^6, nesterov_threshold=0, gtol=gtol)
    B0 = copy(problem.coeff.all)
    return B0
end

# Solve for s-sparse solution.
function __solve__(s, problem, ϵ, gtol, dtol)
    MVDA.fit(MMSVD(), problem, ϵ, s, verbose=false, ninner=10^6, nouter=200, nesterov_threshold=0, gtol=gtol, dtol=dtol, rtol=0.0)
    B = problem.proj.all
    return B
end

# Check classification accuracy against X.
function __accuracy__(B, problem, targets, X)
    n = problem.n
    copyto!(problem.proj.all, B)
    labels = MVDA.classify(problem, X)
    return round(100 * (1 - sum(labels .== targets) / n), sigdigits=4)
end

# Check MSE against ground truth B0.
mse(B, B0) = mean( (B - B0) .^ 2 ) / dot(B0, B0)

# Check contingency table by comparing estimate x against ground truth y.
function discovery_metrics(x, y)
    TP = FP = TN = FN = 0
    for (xi, yi) in zip(eachrow(x), eachrow(y))
        nxi, nyi = norm(xi), norm(yi)
        TP += (nxi != 0) && (nyi != 0)
        FP += (nxi != 0) && (nyi == 0)
        TN += (nxi == 0) && (nyi == 0)
        FN += (nxi == 0) && (nyi != 0)
    end
    return (TP, FP, TN, FN)
end

function run(filename, λ)
    # Fix problem parameters.
    c = 10
    p = 50
    rng = StableRNG(1903)

    # Define range of values for various simulation settings.
    separation_levels = (1.0, 3.0)
    noise_levels = (1e0, 5e0, 1e1)
    dataset_sizes = (10, 40, 100)
    sparsity_levels = Float64[1 - k/p for k in p:-1:0]

    # Allocate arrays for interesting metrics.
    m = length(noise_levels)
    train_error = zeros(p+1, m)
    mean_squared_error = zeros(p+1, m)

    @info "Initializing ouptut file: $(filename).dat"
    open("$(filename).dat", "w+") do io
        write(io, "n,p,c,nclass,ncausal,separation,sigma,sparsity,error,MSE,TP,FP,TN,FN\n")
    end

    for samples_per_class in dataset_sizes, separation in separation_levels
        for (j, sigma) in enumerate(noise_levels)
            # Simulate ground truth with the given parameters.
            n = samples_per_class*c
            targets, X, B0 = MVDA.simulate_ground_truth(p, c, samples_per_class, d=separation, rng=rng, sigma=sigma)

            # Standardize ALL data based on the training set.
            F = StatsBase.fit(ZScoreTransform, X, dims=1)
            has_nan = any(isnan, F.scale) || any(isnan, F.mean)
            has_inf = any(isinf, F.scale) || any(isinf, F.mean)
            has_zero = any(iszero, F.scale)
            if has_nan
                error("Detected NaN in z-score.")
            elseif has_inf
                error("Detected Inf in z-score.")
            elseif has_zero
                for idx in eachindex(F.scale)
                    x = F.scale[idx]
                    F.scale[idx] = ifelse(iszero(x), one(x), x)
                end
            end
            StatsBase.transform!(F, X)

            # Figure out how to group samples together and create a problem instance.
            sortidx = sortperm(targets)
            problem = MVDAProblem(targets, X, intercept=false)

            # Fix hyperparameter controlling distance insensitivity.
            ϵ = MVDA.maximal_deadzone(problem)

            # Create closures specific to our problem instance.
            make_data_plot() = __make_data_plot__(targets, n, c, samples_per_class, ccp, colors)
            make_distance_plot(B) = __make_distance_plot__(B, problem, targets, ϵ, sortidx, ccp, colors)
            init() = __init__(λ, problem, ϵ, 1e-6)
            solve(s) = __solve__(s, problem, ϵ, 1e-3, 1e-3)
            accuracy(B) = __accuracy__(B, problem, targets, X)

            @showprogress "$(samples_per_class) per class / d = $(separation) / σ = $(sigma)... " for (i, s) in enumerate(sparsity_levels)
                # Estimate s-sparse coefficients.
                if s != 0
                    B = solve(s)
                else
                    B = init()
                end

                train_error[i,j] = accuracy(B)
                mean_squared_error[i,j] = mse(B, B0)
                TP, FP, TN, FN = discovery_metrics(B[1:p,:], B0[1:p,:])

                open("$(filename).dat", "a+") do io
                    write(io, "$(n),$(p),$(c),$(samples_per_class),$(c),$(separation),$(sigma),$(s),$(train_error[i,j]),$(mean_squared_error[i,j]),$(TP),$(FP),$(TN),$(FN)\n")
                end
            end
        end
    end
end

# Parse input arguments and run.
length(ARGS) < 2 && error("Script requires a filename (no extension required) and a value λ > 0 for regularization.")
filename = ARGS[1]
λ = parse(Float64, ARGS[2])
@info "Output file: $(filename)"
@info "λ value: $(λ)"
run(filename, λ)


using DifferentialEquations
using DataStructures
#here t denotes sensitivity, 0.3 is a constant which can be changed later after testing
eqn(i,p,s) =  i * 0.3
i0 = 0.5  #intial brightness
#=this would just generate a sol object from which any value of intensity corresponding to
a certain sensitivity value can be generated
eg :  sol(115)
=#
s_span = (0.0,10.0)
#I've taken a value of 1/3.0; might have to change this after testing
prob = ODEProblem(eqn,i0,s_span)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8, dt=0.1, save_everystep=true)
print(sol)
#=
On printing sol, you'll find 2 lists, one for t, one for u
here t is s
u is intensity
=#
arsol = Array(sol)
println(arsol)        #solution array

#############
#This part should logically be in  a different script. Will break down while integration
#############
#I've assumed my current intensity to be 1.0, hence i computed its index
searchsortedfirst(arsol,1.0) #result 12
#Assumed required intensity to be 7.0 (Note : intensity can't be more than 10); will have to discuss solution handling
searchsortedfirst(arsol,7.0) # result 34
 #now iterate through 12 to 34 and change brightness accordingly


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCell
include("../../test/useful_routines.jl")

function step_interface(levelset, levelsetcoeffs, mesh, levelsetspeed, dt)
    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    paddedmesh = CutCell.BoundaryPaddedMesh(cutmesh, 1)
    refseedpoints, spatialseedpoints, seedcellids =
        CutCell.seed_zero_levelset(2, levelset, levelsetcoeffs, cutmesh)
    paddedlevelset = CutCell.BoundaryPaddedLevelset(
        paddedmesh,
        refseedpoints,
        spatialseedpoints,
        seedcellids,
        levelset,
        levelsetcoeffs,
        cutmesh,
        1e-10,
    )
    return CutCell.step_first_order_levelset(paddedlevelset,levelsetspeed,dt)
end

function grid_size(mesh)
    w = CutCell.widths(mesh)
    nn = CutCell.nodes_per_mesh_side(mesh)

    return w ./ (nn .- 1)
end

function time_step_size(levelsetspeed,mesh;CFL=0.9)
    dx = minimum(grid_size(mesh))
    s = maximum(abs.(levelsetspeed))
    return CFL*dx/s
end

function run_time_steps(levelset,initialcondition,mesh,levelsetspeed,dt,nsteps)
    levelsetcoeffs = [copy(initialcondition) for i = 1:nsteps+1]
    for i = 1:nsteps
        levelsetcoeffs[i+1] = step_interface(levelset,levelsetcoeffs[i],mesh,levelsetspeed,dt)
    end
    return levelsetcoeffs
end

x0 = [0.0, 0.0]
L, W = 1.0, 1.0
nelmts = 10
numghostlayers = 1
polyorder = 2

xI = [0.05, 0.0]
normal = [1.0, 0.0]
tol = 1e-8
speed = 1.0
stoptime = 0.8

basis = TensorProductBasis(2, polyorder)
mesh = CutCell.Mesh(x0, [L, W], [nelmts, nelmts], basis)

levelset = InterpolatingPolynomial(1, basis)
initiallevelset =
    CutCell.levelset_coefficients(x -> plane_distance_function(x, normal, xI), mesh)
levelsetspeed = speed*ones(length(initiallevelset))

dt = time_step_size(levelsetspeed,mesh)
nsteps = ceil(Int,stoptime/dt)
levelsetcoeffs = run_time_steps(levelset,initiallevelset,mesh,levelsetspeed,stoptime)

actualstoptime = dt*nsteps
xT = xI + [actualstoptime*speed,0.0]

quad = tensor_product_quadrature(2,3)
err = uniform_mesh_L2_error(levelsetcoeffs[end]',x->plane_distance_function(x,normal,xT),basis,quad,mesh)
@test allapprox(err,[0.0],1e2eps())

function grid_range(mesh)
    x0 = CutCell.reference_corner(mesh)
    w = CutCell.widths(mesh)
    nfmside = CutCell.nodes_per_mesh_side(mesh)

    x = range(x0[1],stop=x0[1]+w[1],length=nfmside[1])
    y = range(x0[2],stop=x0[2]+w[2],length=nfmside[2])

    return x,y
end

using Plots
x,y = grid_range(mesh)
Z = reshape(levelsetcoeffs[1],length(y),:)
Plots.contour(x,y,Z,levels=-1:0.25:1,color="black",linewidth=2,legend=false)

anim = @animate for i = 1:length(levelsetcoeffs)
    Z = reshape(levelsetcoeffs[i],length(y),:)
    contour(x, y, Z, levels = -1:0.25:1,color="black",linewidth=2,legend=false)
end

gif(anim,fps=5)


using ModelingToolkit, StaticArrays, LinearAlgebra
using DiffEqBase
using Test

# Calculus
@parameters t σ ρ β
@variables x y z

eqs = [σ*(y-x),
       x*(ρ-z)-y,
       x*y - β*z]

simpexpr = [
   :(σ * (y - x))
   :(x * (ρ - z) - y)
   :(x * y - β * z)
   ]

for i in 1:3
   @test ModelingToolkit.simplified_expr.(eqs)[i] == simpexpr[i]
   @test ModelingToolkit.simplified_expr.(eqs)[i] == simpexpr[i]
end

∂ = ModelingToolkit.jacobian(eqs,[x,y,z])
for i in 1:3
    ∇ = ModelingToolkit.gradient(eqs[i],[x,y,z])
    @test isequal(∂[i,:],∇)
end
@test all(isequal.(ModelingToolkit.gradient(eqs[1],[x,y,z]),[σ * -1,σ,0]))
@test all(isequal.(ModelingToolkit.hessian(eqs[1],[x,y,z]),0))

# Function building

@parameters σ() ρ() β()
@variables x y z
eqs = [σ*(y-x),
       x*(ρ-z)-y,
       x*y - β*z]
f1,f2 = ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β])
f = eval(f1)
out = [1.0,2,3]
o1 = f([1.0,2,3],[1.0,2,3])
f = eval(f2)
f(out,[1.0,2,3],[1.0,2,3])
@test all(o1 .== out)

function test_worldage()
   @parameters σ() ρ() β()
   @variables x y z
   eqs = [σ*(y-x),
          x*(ρ-z)-y,
          x*y - β*z]
   f, f_iip = ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β],(),ModelingToolkit.simplified_expr,Val{false})
   out = [1.0,2,3]
   o1 = f([1.0,2,3],[1.0,2,3])
   f_iip(out,[1.0,2,3],[1.0,2,3])
end
test_worldage()

## No parameters
@variables x y z
eqs = [(y-x)^2,
       x*(x-z)-y,
       x*y - y*z]
f1,f2 = ModelingToolkit.build_function(eqs,[x,y,z])
f = eval(f1)
out = zeros(3)
o1 = f([1.0,2,3])
f = eval(f2)
f(out,[1.0,2,3])
@test all(out .== o1)

function test_worldage()
   @variables x y z
   eqs = [(y-x)^2,
          x*(x-z)-y,
          x*y - y*z]
   f, f_iip = ModelingToolkit.build_function(eqs,[x,y,z],(),(),ModelingToolkit.simplified_expr,Val{false})
   out = zeros(3)
   o1 = f([1.0,2,3])
   f_iip(out,[1.0,2,3])
end
test_worldage()


using FileIO
using Parameters

include("random_matrices.jl")
include("network_sims.jl")
include("LS_calculations.jl")

#------------------------------------------------------------------------------
# This is the part where we find which eta and network realization to

J_seeds = collect(5000:5019)
etas = collect(-1.0:0.1:1.0)

ind = parse(Int64, ARGS[1])
ps = [(eta, J_seed) for eta in etas, J_seed in J_seeds]
eta, J_seed = ps[ind]
#------------------------------------------------------------------------------

# Where the results will be saved
dir = "results1"

N = 1000
g = 10.0
s = 1.0

h0_seed = 6000
dh_seed = 7000
Q0_seed = 8000

J = generate_J(N, eta, seed=J_seed)
h = generate_nv(N, Float64, seed=h0_seed)
dh = generate_nv(N, Float64, seed=dh_seed)
Q = generate_ONB(N, Float64, seed=Q0_seed)

dt = 0.01
t_bound = 2000.0
t_ONS = 1.0
wait = 200.0

@time Lspec, max_LE = sim_lyapunov(J, g, h, dh, Q; s=s, dt=dt, t_bound=t_bound,
		t_ONS=t_ONS, wait=wait, return_max_LE=true, printstuff=false,
		p_factor=1e-3, scale_limit=100)

dimensionality = KY_attractor_dim(Lspec)
entropy_rate = KS_entropy_rate(Lspec)

fn = "Lspecs_E$(ind)_N=$N,g=$g,eta=$eta,J_seed=$J_seed.jld2"
out = Dict{String,Any}()
@pack! out = N, g, eta, s, J_seed, h0_seed, dh_seed, Q0_seed, dt, t_bound,
		t_ONS, wait, Lspec, max_LE, dimensionality, entropy_rate

save(joinpath(dir, fn), out)


#Laguerre breaks at 16
for i in 12:20
    println("n = $i ", all(L_rec(i) .< 1/eps()))
end

# Laguerre Coefficient Calculator
# Credit to [Peter Luschny Apr. 11, 2015](https://oeis.org/A021009)
function l(n,m)
    if 0 == n == m
        1
    elseif m == -1
        0
    elseif n < m
        0
    elseif n >= 1
        (n+m+1) * l(n-1, m) - l(n-1, m-1)
    end
end

LTable(n) = [l(i,j) for i in 0:n,j in 0:n]

horner_L12(x) = @evalpoly x 1.0 -12.0 33.0 -36.666666666666664 20.625 -6.6 1.2833333333333334 -0.15714285714285714 0.012276785714285714 -0.0006062610229276896 1.8187830687830687e-5 -3.0062530062530064e-7

# Hermite Coefficient Calculator
# Credit to [Paul Barry, Aug 28, 2005.](https://oeis.org/search?q=coefficient+triangle+hermite&sort=&language=&go=Search)
function h(n,k)

    if iseven(n-k) && (n-k >= 0)
        ( (-1)^((n-k)/2) * (2^k) * factorial(n)) / (factorial(k) * factorial((n-k)/2))
    else
        0
    end
end

Htable(n) = [h(i,j) for i in 0:n, j in 0:n]

horner_H12(x) = Base.Math.@evalpoly x 665280 0 -7983360 0 13305600 0 -7096320 0 1520640 0 -135168 0 4096

# Legendre Coefficient Calculator
# Credit to Ralf Stephan, Apr. 07. 2016 https://oeis.org/A008556
PTable(n) =  [binomial(2*(i-j),i-j) * binomial(i-j,j) for i in 0:n, j in 0:n]

horner_P12(x) = Base.Math.@evalpoly x

# Tshebyshev T First Kind Coefficient Calculator
# Credit to Micahel Somos, Aug. 08, 2011 https://oeis.org/A049310
function T(n,k)
    if k < 0 || k > n || (n + k) % 2 == 1
        return 0
    end
    return (-1)^((n+k)/2 +k) * binomial(Int((n+k)/2),k)
end

TTable(n) = Int[u(i,j) for i in 0:n, j in 0:n]

# Tshebyshev U Second Kind Coefficient Calculator

function u(n,m)
    if n < m || isodd(n+m)
        0
    else
        ((-1)^((n+m)/2+m))*(2^m)*binomial(Int((n+m)/2), m)
    end
end

UTable(n) = Int[u(i,j) for i in 0:n, j in 0:n]

# Consider adding the Bell and Bernoulli polynomials

# stirling numbers of the first Kind

s1(m,n) = sum([(-1)^k*binomial(n-1+k,n-m+k)*binomial(2*n-m,n-m-k)*s2(k,n-m-k) for k in 0:n-m])
row_s1(n) = [s1(i,n) for i in 1:n]

# stirling numbers of the second kind

s2(n,k) = (1/factorial(k))*sum([(-1)^i * binomial(k,i) * (k - i)^n for i in 0:k])

row_s2(n) = [s2(n,i) for i in 1:n]

s(n,m) = sum([(-1)^(n-j) * binomial(k,j) * factorial(j*y - 1 + n) / factorial(j*y-1) * y ^ (-k) / factorial(k) for j in 0:k])

function stirlings1(n::Int, k::Int, signed::Bool=false)
    if signed == true
        return (-1)^(n - k) * stirlings1(n, k)
    end

    if n < 0
        throw(DomainError(n, "n must be nonnegative"))
    elseif n == k == 0
        return 1
    elseif n == 0 || k == 0
        return 0
    elseif n == k
        return 1
    elseif k == 1
        return factorial(n-1)
    elseif k == n - 1
        return binomial(n, 2)
    elseif k == n - 2
        return div((3 * n - 1) * binomial(n, 3), 4)
    elseif k == n - 3
        return binomial(n, 2) * binomial(n, 4)
    end

    return (n - 1) * stirlings1(n - 1, k) + stirlings1(n - 1, k - 1)
end

function stirlings2(n::Int, k::Int)
    if n < 0
        throw(DomainError(n, "n must be nonnegative"))
    elseif n == k == 0
        return 1
    elseif n == 0 || k == 0
        return 0
    elseif k == n - 1
        return binomial(n, 2)
    elseif k == 2
        return 2^(n-1) - 1
    end

    return k * stirlings2(n - 1, k) + stirlings2(n - 1, k - 1)
end

# Straight from Julia codebase
macro evalpoly(z, p...)
    a = :($(esc(p[end])))
    b = :($(esc(p[end-1])))
    as = []
    for i = length(p)-2:-1:1
        ai = Symbol("a", i)
        push!(as, :($ai = $a))
        a = :(muladd(r, $ai, $b))
        b = :($(esc(p[i])) - s * $ai) # see issue #15985 on fused mul-subtract
    end
    ai = :a0
    push!(as, :($ai = $a))
    C = Expr(:block,
             :(x = real(tt)),
             :(y = imag(tt)),
             :(r = x + x),
             :(s = muladd(x, x, y*y)),
             as...,
             :(muladd($ai, tt, $b)))
    R = Expr(:macrocall, Symbol("@horner"), (), :tt, map(esc, p)...)
    :(let tt = $(esc(z))
          isa(tt, Complex) ? $C : $R
      end)
end

#P
d(p, α,n) = binomial(n + α,n)
b(p, α, β, n, m) = (n-m+1)*(α+β+n+m)
c(p, α, m) = 2m*(α + m)
f(p, x) = 1 - x

#Ceven
d(p, α, n) = (-1)^n * (poch(α,n)/factorial(n))
b(p, α, n, m) = 2*(n-m+1)*(α+n+m-1)
c(p, m) = m*(2m-1)
f(p, x) = x^2

#Codd
d(p, α, n) = (-1)^n * (poch(α,n+1)/factorial(n))*2x
b(p, α, β, n, m) = 2*(n-m+1)*(α+β+n+m)
c(p, m) = m*(2m+1)
f(p, x) = x^2

#Teven
d(p, n) = (-1)^n
b(p, n, m) = 2(n-m+1)*(n+m-1)
c(p, m) = m*(2m-1)
f(p, x) = x^2

#Todd
d(p, n, x) = (-1^n)*(2n+1)*x
b(p, n, m) = 2*(n-m+1)*(n+m)
c(p, m) = m*(2m+1)
f(p, x) = x^2

#Ueven
d(p, n) = (-1)^n
b(p, n, m) = 2*(n-m+1)*(n+m)
c(p, m) = m*(2m+1)
f(p, x) = x^2

#Uodd
d(p, n, x) = (-1)^n * 2(n+1)*x
b(p, n, m) = 2*(n-m+1)*(n+m)
c(p, m) = m*(2m + 1)
f(p, x) = x^2

#Peven
d(p, n) = (-1/4)^n * binomial(2n,n)
b(p, n, m) = (n - m + 1) * (2n + 2m -1)
c(p, m) = m*(2m-1)
f(x) = x^2

#Podd
d(p, n, x) = (-1/4)^n * binomial(2n+1,n) * (n+1) * x
b(p, n, m) = (n - m + 1) * (2n + 2m + 1)
c(p, m) = m * (2m + 1)
f(p, x) = x^2

#L
d(p, α, n) = binomial(n + α, n)
b(p, n, m) = (n - m + 1)
c(p, m, α) = m * (α + m)
f(p, x) = x

#Heven
d(p, n) = (-1) ^ n * (factorial(2n)/factorial(n))
b(p, n, m) = 2*(n - m + 1)
c(p, m) = m*(2m - 1)
f(p, x) = x^2

#Hodd
d(p, n, x) = (-1) ^ n * (factorial(2n +1)/factorial(n))*2x
b(p, n, m) = 2*(n - m + 1)
c(p, m) = m * (2m + 1)
f(p, x) = x^2

#interpolate an expression as an expression into another expression.
ex_new = Meta.quot(ex)
quote
    still_expression = $(esc(ex_new))
end


import Base.copy

export
    CostFunction,
    QuadraticQuatCost,
    RBCost,
    QuatLQRCost,
    SatDiffCost,
    ErrorQuadratic

#*********************************#
#       COST FUNCTION CLASS       #
#*********************************#

abstract type CostFunction end

#######################################################
#              COST FUNCTION INTERFACE                #
#######################################################

"""
$(TYPEDEF)
Cost function of the form
    1/2xₙᵀ Qf xₙ + qfᵀxₙ +  ∫ ( 1/2xᵀQx + 1/2uᵀRu + xᵀHu + qᵀx + rᵀu ) dt from 0 to tf
R must be positive definite, Q and Qf must be positive semidefinite

Constructor use any of the following constructors:
```julia
QuadraticCost(Q, R, H, q, r, c)
QuadraticCost(Q, R; H, q, r, c)
QuadraticCost(Q, q, c)
```
Any optional or omitted values will be set to zero(s).
"""
mutable struct QuadraticCost{TQ,TR,TH,Tq,Tr,T} <: CostFunction
    Q::TQ                 # Quadratic stage cost for states (n,n)
    R::TR                 # Quadratic stage cost for controls (m,m)
    H::TH                 # Quadratic Cross-coupling for state and controls (n,m)
    q::Tq                 # Linear term on states (n,)
    r::Tr                 # Linear term on controls (m,)
    c::T                                 # constant term
    function QuadraticCost(Q::TQ, R::TR, H::TH,
            q::Tq, r::Tr, c::T; checks=true) where {TQ,TR,TH,Tq,Tr,T}
        @assert size(Q,1) == length(q)
        @assert size(R,1) == length(r)
        @assert size(H) == (length(r), length(q))
        if checks
            if !isposdef(Array(R))
                @warn "R is not positive definite"
            end
            if !ispossemidef(Array(Q))
                err = ArgumentError("Q must be positive semi-definite")
                throw(err)
            end
        end
        new{TQ,TR,TH,Tq,Tr,T}(Q,R,H,q,r,c)
    end
end

state_dim(cost::QuadraticCost) = length(cost.q)
control_dim(cost::QuadraticCost) = length(cost.r)

# Constructors
function QuadraticCost(Q,R; H=similar(Q,size(R,1), size(Q,1)), q=zeros(size(Q,1)),
        r=zeros(size(R,1)), c=0.0, checks=true)
    QuadraticCost(Q,R,H,q,r,c, checks=checks)
end

function QuadraticCost(Q::Union{Diagonal{T,SVector{N,T}}, SMatrix{N,N,T}},
            R::Union{Diagonal{T,SVector{M,T}}, SMatrix{M,M,T}};
            H=(@SMatrix zeros(T,M,N)),
            q=(@SVector zeros(T,N)),
            r=(@SVector zeros(M)),
            c=0.0, checks=true) where {T,N,M}
    QuadraticCost(Q,R,H,q,r,c, checks=checks)
end

function QuadraticCost(Q,q,c; checks=true)
    QuadraticCost(Q,zeros(0,0),zeros(0,size(Q,1)),q,zeros(0),c,checks=checks)
end

"""
$(SIGNATURES)
Cost function of the form
``(x-x_f)^T Q (x_x_f) + u^T R u``
R must be positive definite, Q must be positive semidefinite
"""
function LQRCost(Q::AbstractArray, R::AbstractArray,
        xf::AbstractVector, uf=(@SVector zeros(size(R,1))); checks=true)
    H = @SMatrix zeros(size(R,1),size(Q,1))
    q = -Q*xf
    r = -R*uf
    c = 0.5*xf'*Q*xf + 0.5*uf'R*uf
    return QuadraticCost(Q, R, H, q, r, c, checks=checks)
end

"""
$(SIGNATURES)
Cost function of the form
``(x-x_f)^T Q (x-x_f)``
Q must be positive semidefinite
"""
function LQRCostTerminal(Qf::AbstractArray,xf::AbstractVector)
    qf = -Qf*xf
    cf = 0.5*xf'*Qf*xf
    return QuadraticCost(Qf,zeros(0,0),zeros(0,size(Qf,1)),qf,zeros(0),cf)
end

# Cost function methods
function stage_cost(cost::QuadraticCost, x::AbstractVector, u::AbstractVector)
    0.5*x'cost.Q*x + 0.5*u'*cost.R*u + cost.q'x + cost.r'u + cost.c #+ u'*cost.H*x
end

function stage_cost(cost::QuadraticCost, xN::AbstractVector{T}) where T
    0.5*xN'cost.Q*xN + cost.q'*xN + cost.c
end

function gradient!(E::AbstractExpansion, cost::QuadraticCost, x, u)
    E.x .= cost.Q*x + cost.q + cost.H'*u
    E.u .= cost.R*u + cost.r + cost.H*x
    return nothing
    mul!(E.x, cost.Q, x)
    E.x .+= cost.q
    mul!(E.x, cost.H', u, 1.0, 1.0)

    mul!(E.u, cost.R, u)
    E.u .+= cost.r
    mul!(E.u, cost.H, x, 1.0, 1.0)
    return nothing
end

function hessian!(E::AbstractExpansion, cost::QuadraticCost, x, u)
    E.xx .= cost.Q
    E.uu .= cost.R
    E.ux .= cost.H
    return nothing
end

# Additional Methods
function Base.show(io::IO, cost::QuadraticCost)
    print(io, "QuadraticCost{...}")
end

import Base: +
"Add two Quadratic costs of the same size"
function +(c1::QuadraticCost, c2::QuadraticCost)
    @assert state_dim(c1) == state_dim(c2)
    @assert control_dim(c1) == control_dim(c2)
    QuadraticCost(c1.Q + c2.Q, c1.R + c2.R, c1.H + c2.H,
                  c1.q + c2.q, c1.r + c2.r, c1.c + c2.c)
end

function copy(cost::QuadraticCost)
    return QuadraticCost(copy(cost.Q), copy(cost.R), copy(cost.H), copy(cost.q), copy(cost.r), copy(cost.c))
end

############################################################################################
#                                 QUADRATIC QUATERNION COST FUNCTION
############################################################################################

struct QuadraticQuatCost{T,N,M,N4} <: CostFunction
    Q::Diagonal{T,SVector{N,T}}
    R::Diagonal{T,SVector{M,T}}
    q::SVector{N,T}
    r::SVector{M,T}
    c::T
    w::T
    q_ref::SVector{4,T}
    q_ind::SVector{4,Int}
    Iq::SMatrix{N,4,T,N4}
    function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
            q::SVector{N,T}, r::SVector{M,T}, c::T, w::T,
            q_ref::SVector{4,T}, q_ind::SVector{4,Int}) where {T,N,M}
        Iq = @MMatrix zeros(N,4)
        for i = 1:4
            Iq[q_ind[i],i] = 1
        end
        Iq = SMatrix{N,4}(Iq)
        return new{T,N,M,N*4}(Q, R, q, r, c, w, q_ref, q_ind, Iq)
    end
end

state_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = N
control_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = M

function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}};
        q=(@SVector zeros(N)), r=(@SVector zeros(M)), c=zero(T), w=one(T),
        q_ref=(@SVector [1.0,0,0,0]), q_ind=(@SVector [4,5,6,7])) where {T,N,M}
    QuadraticQuatCost(Q, R, q, r, c, q_ref, q_ind)
end

function stage_cost(cost::QuadraticQuatCost, x::SVector, u::SVector)
    stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
end

function stage_cost(cost::QuadraticQuatCost, x::SVector)
    J = 0.5*x'cost.Q*x + cost.q'x + cost.c
    q = x[cost.q_ind]
    dq = cost.q_ref'q
    J += cost.w*min(1+dq, 1-dq)
end

function gradient(cost::QuadraticQuatCost{T,N,M}, x::SVector, u::SVector) where {T,N,M}
    Qx = cost.Q*x + cost.q
    q = x[cost.q_ind]
    dq = cost.q_ref'q
    if dq < 0
        Qx += cost.w*cost.Iq*cost.q_ref
    else
        Qx -= cost.w*cost.Iq*cost.q_ref
    end
    Qu = cost.R*u + cost.r
    return Qx, Qu
end

function hessian(cost::QuadraticQuatCost, x::SVector{N}, u::SVector{M}) where {N,M}
    Qxx = cost.Q
    Quu = cost.R
    Qux = @SMatrix zeros(M,N)
    return Qxx, Quu, Qux
end

function QuatLQRCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}}, xf,
        uf=(@SVector zeros(M)); w=one(T), quat_ind=(@SVector [4,5,6,7])) where {T,N,M}
    r = -R*uf
    q = -Q*xf
    c = 0.5*xf'Q*xf + 0.5*uf'R*uf
    q_ref = xf[quat_ind]
    return QuadraticQuatCost(Q, R, q, r, c, w, q_ref, quat_ind)
end

function change_dimension(cost::QuadraticQuatCost, n, m)
    n0,m0 = state_dim(cost), control_dim(cost)
    Q_diag = diag(cost.Q)
    R_diag = diag(cost.R)
    q = cost.q
    r = cost.r
    if n0 != n
        dn = n - n0  # assumes n > n0
        pad = @SVector zeros(dn)
        Q_diag = [Q_diag; pad]
        q = [q; pad]
    end
    if m0 != m
        dm = m - m0  # assumes m > m0
        pad = @SVector zeros(dm)
        R_diag = [R_diag; pad]
        r = [r; pad]
    end
    QuadraticQuatCost(Diagonal(Q_diag), Diagonal(R_diag), q, r, cost.c, cost.w,
        cost.q_ref, cost.q_ind)
end

function (+)(cost1::QuadraticQuatCost, cost2::QuadraticCost)
    @assert state_dim(cost1) == state_dim(cost2)
    @assert control_dim(cost1) == control_dim(cost2)
    @assert norm(cost2.H) ≈ 0
    QuadraticQuatCost(cost1.Q + cost2.Q, cost1.R + cost2.R,
        cost1.q + cost2.q, cost1.r + cost2.r, cost1.c + cost2.c,
        cost1.w, cost1.q_ref, cost1.q_ind)
end

(+)(cost1::QuadraticCost, cost2::QuadraticQuatCost) = cost2 + cost1

struct ErrorQuadratic{Rot,N,M} <: CostFunction
    model::RigidBody{Rot}
    Q::Diagonal{Float64,SVector{12,Float64}}
    R::Diagonal{Float64,SVector{M,Float64}}
    r::SVector{M,Float64}
    c::Float64
    x_ref::SVector{N,Float64}
    q_ind::SVector{4,Int}
end

state_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = N
control_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = M

function ErrorQuadratic(model::RigidBody{Rot}, Q::Diagonal{T,<:SVector{12}},
        R::Diagonal{T,<:SVector{M}},
        x_ref::SVector{N}, u_ref=(@SVector zeros(T,M)); r=(@SVector zeros(T,M)), c=zero(T),
        q_ind=(@SVector [4,5,6,7])) where {T,N,M,Rot}
    r += -R*u_ref
    c += 0.5*u_ref'R*u_ref
    return ErrorQuadratic{Rot,N,M}(model, Q, R, r, c, x_ref, q_ind)
end

function stage_cost(cost::ErrorQuadratic, x::SVector)
    dx = state_diff(cost.model, x, cost.x_ref)
    return 0.5*dx'cost.Q*dx + cost.c
end

function stage_cost(cost::ErrorQuadratic, x::SVector, u::SVector)
    stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
end

function cost_expansion(cost::ErrorQuadratic{Rot}, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M,Rot<:UnitQuaternion}
    x,u = state(z), control(z)
    model = cost.model
    Q = cost.Q
    q = orientation(model, x)
    q_ref = orientation(model, cost.x_ref)
    dq = SVector(q_ref\q)
    err = state_diff(model, x, cost.x_ref)
    dx = @SVector [err[1],  err[2],  err[3],
                    dq[1],   dq[2],   dq[3],   dq[4],
                   err[7],  err[8],  err[9],
                   err[10], err[11], err[12]]
    G = state_diff_jacobian(model, dx) # n × dn

    # Gradient
    dmap = inverse_map_jacobian(model, dx) # dn × n
    Qx = G'dmap'Q*err
    Qu = cost.R*u

    # Hessian
    ∇jac = inverse_map_∇jacobian(model, dx, Q*err)
    Qxx = G'dmap'Q*dmap*G + G'∇jac*G + ∇²differential(model, x, dmap'Q*err)
    Quu = cost.R
    Qux = @SMatrix zeros(M,N-1)
    return Qxx, Quu, Qux, Qx, Qu
end

function cost_expansion(cost::ErrorQuadratic, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M}
    x,u = state(z), control(z)
    model = cost.model
    q = orientation(model, x)
    q_ref = orientation(model, cost.x_ref)
    err = state_diff(model, x, cost.x_ref)
    dx = err
    G = state_diff_jacobian(model, dx) # n × n

    # Gradient
    dmap = inverse_map_jacobian(model, dx) # n × n
    Qx = G'dmap'cost.Q*err
    Qu = cost.R*u + cost.r

    # Hessian
    Qxx = G'dmap'cost.Q*dmap*G
    Quu = cost.R
    Qux = @SMatrix zeros(M,N)
    return Qxx, Quu, Qux, Qx, Qu
end

function change_dimension(cost::ErrorQuadratic, n, m)
    n0,m0 = state_dim(cost), control_dim(cost)
    Q_diag = diag(cost.Q)
    R_diag = diag(cost.R)
    r = cost.r
    if n0 != n
        dn = n - n0  # assumes n > n0
        pad = @SVector zeros(dn) # assume the new states don't have quaternions
        Q_diag = [Q_diag; pad]
    end
    if m0 != m
        dm = m - m0  # assumes m > m0
        pad = @SVector zeros(dm)
        R_diag = [R_diag; pad]
        r = [r; pad]
    end
    ErrorQuadratic(cost.model, Diagonal(Q_diag), Diagonal(R_diag), r, cost.c,
        cost.x_ref, cost.q_ind)
end

function (+)(cost1::ErrorQuadratic, cost2::QuadraticCost)
    @assert control_dim(cost1) == control_dim(cost2)
    @assert norm(cost2.H) ≈ 0
    @assert norm(cost2.q) ≈ 0
    if state_dim(cost2) == 13
        rm_quat = @SVector [1,2,3,4,5,6,8,9,10,11,12,13]
        Q2 = Diagonal(diag(cost2.Q)[rm_quat])
    else
        Q2 = cost2.Q
    end
    ErrorQuadratic(cost1.model, cost1.Q + Q2, cost1.R + cost2.R,
        cost1.r + cost2.r, cost1.c + cost2.c,
        cost1.x_ref, cost1.q_ind)
end

(+)(cost1::QuadraticCost, cost2::ErrorQuadratic) = cost2 + cost1

struct SatDiffCost{Rot} <: CostFunction
    model::RigidBody{Rot}
    Q1::Diagonal{Float64,SVector{3,Float64}}
    Q2::Diagonal{Float64,SVector{3,Float64}}
    R::Diagonal{Float64,SVector{3,Float64}}
    q_ref::SVector{4,Float64}
    ω_ref::SVector{3,Float64}
end

function stage_cost(cost::SatDiffCost{Rot}, x::SVector) where Rot
    ω = @SVector [x[1],x[2],x[3]]
    q = Dynamics.orientation(cost.model, x)
    q0 = Rot(UnitQuaternion(cost.q_ref))

    dω = ω - cost.ω_ref
    dq = q ⊖ q0
    return 0.5*(dq'cost.Q2*dq + dω'cost.Q1*dω)
end

function stage_cost(cost::SatDiffCost, x::SVector, u::SVector)
    J = stage_cost(cost, x) + 0.5*u'cost.R*u
end

function cost_expansion(cost::SatDiffCost{Rot}, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M,Rot<:UnitQuaternion}
    x,u = state(z), control(z)
    Q = cost.Q2  # cost for quaternion
    ω = @SVector [x[1],x[2],x[3]]
    q = @SVector [x[1],x[2],x[3],x[4]]
    q = Dynamics.orientation(cost.model, x)
    q0 = Rot(UnitQuaternion(cost.q_ref))

    dω = ω - cost.ω_ref
    dq = q0\q
    err = q ⊖ q0
    G = ∇differential(dq) #Lmult(dq)*Vmat()'

    # Gradient
    Qω = cost.Q1*ω
    dmap = inverse_map_jacobian(dq) #jacobian(CayleyMap,dq)
    Qq = G'dmap'Q*err
    Qx = [Qω; Qq]
    Qu = cost.R*u

    # Hessian
    Qωω = cost.Q1
    ∇jac = inverse_map_∇jacobian(dq, Q*err)
    Qqq = G'dmap'Q*dmap*G + G'∇jac*G
    Qxx = @SMatrix [
        Qωω[1,1] 0 0 0 0 0;
        0 Qωω[2,2] 0 0 0 0;
        0 0 Qωω[3,3] 0 0 0;
        0 0 0 Qqq[1,1] Qqq[1,2] Qqq[1,3];
        0 0 0 Qqq[2,1] Qqq[2,2] Qqq[2,3];
        0 0 0 Qqq[3,1] Qqq[3,2] Qqq[3,3];
    ]
    Quu = cost.R
    Qux = @SMatrix zeros(M,N-1)
    return Qxx, Quu, Qux, Qx, Qu
end

function cost_expansion(cost::SatDiffCost{Rot}, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M,Rot}
    x,u = state(z), control(z)
    Q = cost.Q2  # cost for quaternion
    ω = @SVector [x[1],x[2],x[3]]
    q = @SVector [x[1],x[2],x[3],x[4]]
    q = Dynamics.orientation(cost.model, x)
    q0 = Rot(UnitQuaternion(cost.q_ref))

    dω = ω - cost.ω_ref
    dq = q0\q
    err = q ⊖ q0
    G = ∇differential(dq) #Lmult(dq)*Vmat()'

    # Gradient
    Qω = cost.Q1*ω
    Qq = G'Q*err
    Qx = [Qω; Qq]
    Qu = cost.R*u

    # Hessian
    Qωω = cost.Q1
    ∇jac = ∇²differential(dq, Q*err)
    Qqq = G'Q*G + ∇jac
    Qxx = @SMatrix [
        Qωω[1,1] 0 0 0 0 0;
        0 Qωω[2,2] 0 0 0 0;
        0 0 Qωω[3,3] 0 0 0;
        0 0 0 Qqq[1,1] Qqq[1,2] Qqq[1,3];
        0 0 0 Qqq[2,1] Qqq[2,2] Qqq[2,3];
        0 0 0 Qqq[3,1] Qqq[3,2] Qqq[3,3];
    ]
    Quu = cost.R
    Qux = @SMatrix zeros(M,N)
    return Qxx, Quu, Qux, Qx, Qu
end

#
# """
# $(TYPEDEF)
# Cost function of the form
#     ℓf(xₙ) + ∫ ℓ(x,u) dt from 0 to tf
# """
# struct GenericCost <: CostFunction
#     ℓ::Function             # Stage cost
#     ℓf::Function            # Terminal cost
#     expansion::Function     # 2nd order Taylor Series Expansion of the form,  Q,R,H,q,r = expansion(x,u)
#     n::Int                  #                                                     Qf,qf = expansion(xN)
#     m::Int
# end
#
# """
# $(SIGNATURES)
# Create a Generic Cost, specifying the gradient and hessian of the cost function analytically
#
# # Arguments
# * hess: multiple-dispatch function of the form,
#     Q,R,H = hess(x,u) with sizes (n,n), (m,m), (m,n)
#     Qf = hess(xN) with size (n,n)
# * grad: multiple-dispatch function of the form,
#     q,r = grad(x,u) with sizes (n,), (m,)
#     qf = grad(x,u) with size (n,)
#
# """
# function GenericCost(ℓ::Function, ℓf::Function, grad::Function, hess::Function, n::Int, m::Int)
#     @warn "Use GenericCost with caution. It is untested and not likely to work"
#     function expansion(x::Vector{T},u::Vector{T}) where T
#         Q,R,H = hess(x,u)
#         q,r = grad(x,u)
#         return Q,R,H,q,r
#     end
#     expansion(xN) = hess(xN), grad(xN)
#     GenericCost(ℓ,ℓf, expansion, n,m)
# end
#
# """ $(TYPEDEF)
# This is an experimental cost function type that allows a cost function to be evaluated
#     on only a portion of the state and control. Right now, the implementation assumes
#     there is no coupling between the state and control.
#
# It should be noted that for `QuadraticCost`s it is likely more efficient to simply make a new
#     `QuadraticCost` that has zeros in the right places, and then add the cost functions together.
#
# # Constructor
# ```julia
# IndexedCost(cost, ix::UnitRange, iu::UnitRange)
# ```
# """
# struct IndexedCost{iX,iU,C} <: CostFunction
#     cost::C
# end
#
# function IndexedCost(cost::C, ix::UnitRange, iu::UnitRange) where C<:CostFunction
#     if C <: QuadraticCost
#         if norm(cost.H) != 0
#             throw(ErrorException("IndexedCost of functions with x-u coupling not implemented"))
#         end
#     else
#         @warn "IndexedCost will only work for costs without x-u coupling (Qux = 0)"
#     end
#     IndexedCost{ix,iu,C}(cost)
# end
#
# @generated function stage_cost(costfun::IndexedCost{iX,iU}, x::SVector{N}, u::SVector{M}) where {iX,iU,N,M}
#     ix = SVector{length(iX)}(iX)
#     iu = SVector{length(iU)}(iU)
#     quote
#         x0 = x[$ix]
#         u0 = u[$iu]
#         stage_cost(costfun.cost, x0, u0)
#     end
# end
#
# @generated function stage_cost(costfun::IndexedCost{iX,iU}, x::SVector{N}) where {iX,iU,N}
#     ix = SVector{length(iX)}(iX)
#     quote
#         x0 = x[$ix]
#         stage_cost(costfun.cost, x0)
#     end
# end
#
# @generated function gradient(costfun::IndexedCost{iX,iU}, x::SVector{N}, u::SVector{M}) where {iX,iU,N,M}
#     l1x = iX[1] - 1
#     l2x = N-iX[end]
#     l1u = iU[1] - 1
#     l2u = M-iU[end]
#     quote
#         x = x[$iX]
#         u = u[$iU]
#         Qx, Qu = gradient(costfun.cost, x, u)
#         Qx = [@SVector zeros($l1x); Qx; @SVector zeros($l2x)]
#         Qu = [@SVector zeros($l1u); Qu; @SVector zeros($l2u)]
#         return Qx, Qu
#     end
# end
#
#
# @generated function hessian(costfun::IndexedCost{iX,iU}, x::SVector{N}, u::SVector{M}) where {iX,iU,N,M}
#     l1x = iX[1] - 1
#     l2x = N-iX[end]
#     l1u = iU[1] - 1
#     l2u = M-iU[end]
#     quote
#         x = x[$iX]
#         u = u[$iU]
#         Qxx, Quu, Qux  = hessian(costfun.cost, x, u)
#         Qxx1 = Diagonal(@SVector zeros($l1x))
#         Qxx2 = Diagonal(@SVector zeros($l2x))
#         Quu1 = Diagonal(@SVector zeros($l1u))
#         Quu2 = Diagonal(@SVector zeros($l2u))
#
#         Qxx = blockdiag(Qxx1, Qxx, Qxx2)
#         Quu = blockdiag(Quu1, Quu, Quu2)
#         Qux = @SMatrix zeros(M,N)
#         Qxx, Quu, Qux
#     end
# end
#
function SparseArrays.blockdiag(Qs::Vararg{<:Diagonal})
    Diagonal(vcat(diag.(Qs)...))
end

function SparseArrays.blockdiag(Qs::Vararg{<:AbstractMatrix})
    # WARNING: this is slow and is only included as a fallback
    cat(Qs...,dims=(1,2))
end
#
# function change_dimension(cost::CostFunction,n,m)
#     n0,m0 = state_dim(cost), control_dim(cost)
#     ix = 1:n0
#     iu = 1:m0
#     IndexedCost(cost, ix, iu)
# end
#
function change_dimension(cost::QuadraticCost, n, m)
    n0,m0 = state_dim(cost), control_dim(cost)
    @assert n >= n0
    @assert m >= m0

    ix = 1:n0
    iu = 1:m0

    Q_ = Diagonal(@SVector zeros(n-n0))
    R_ = Diagonal(@SVector zeros(m-m0))
    H1 = @SMatrix zeros(m0, n-n0)
    H2 = @SMatrix zeros(m-m0, n)
    q_ = @SVector zeros(n-n0)
    r_ = @SVector zeros(m-m0)
    c = cost.c

    # Insert old values
    Q = blockdiag(cost.Q, Q_)
    R = blockdiag(cost.R, R_)
    H = [cost.H H1]
    H = [H; H2]
    q = [cost.q; q_]
    r = [cost.r; r_]
    QuadraticCost(Q,R,H,q,r,c,checks=false)
end


mutable struct MatrixNetwork{T}
    n::Int64 # number of columns/rows
    rp::Vector{Int64} # row pointers
    ci::Vector{Int64} # column indices
    vals::Vector{T} # corresponding values
end

function MatrixNetwork(A::SparseMatrixCSC{T,Int64}) where T
    At = copy(A')
    return MatrixNetwork(size(At,2),At.colptr,At.rowval,At.nzval)
end

function _second(x)
    return x[2]
end

MatrixNetwork(edges::Vector{Tuple{Int,Int}}, n::Int) =
    MatrixNetwork(map(first,edges),map(_second,edges), n)

MatrixNetwork(ei::Vector{Int64},ej::Vector{Int64}) =
    MatrixNetwork(ei,ej,max(maximum(ei),maximum(ej)))

function MatrixNetwork(ei::Vector{Int64},ej::Vector{Int64},n::Int64)
    At = sparse(ej,ei,true,n,n);
    return MatrixNetwork(size(At,2),At.colptr,At.rowval,At.nzval)
end

function _matrix_network_direct(A::SparseMatrixCSC{T,Int64}) where T
    return MatrixNetwork(size(A,2),A.colptr,A.rowval,A.nzval)
end

function _matrix_network_direct(A::SparseMatrixCSC{T,Int64},v) where T
    nzval = ones(typeof(v),length(A.nzval))
    return MatrixNetwork(size(A,2),A.colptr,A.rowval,nzval)
end

import SparseArrays.sparse, Base.size, Base.*, LinearAlgebra.mul!, Base.convert, Base.adjoint, Base.copy

"""
Return back an adjacency matrix representation
of the transpose. This requires no work.
"""
function sparse_transpose(A::MatrixNetwork{T}) where T
    return SparseMatrixCSC(A.n,A.n,A.rp,A.ci,A.vals)
end

adjoint(A::MatrixNetwork{T}) where T = Adjoint{Any,MatrixNetwork{T}}(A)

"""
Return back an adjacency matrix representation
of the current MatrixNetwork
"""
function sparse(A::MatrixNetwork{T}) where T
    return copy(sparse_transpose(A)')
end

function size(A::MatrixNetwork)
    return (A.n,A.n)
end

import Base.ndims
ndims(op::MatrixNetwork) = 2

function size(A::MatrixNetwork, dim::Integer)
    if dim == 1 || dim == 2
        return A.n
    elseif dim > 2
        return 1
    else
        throw(DomainError(dim))
    end
end

*(M::MatrixNetwork{T}, b::AbstractVector{S}) where {T,S} =
    mul!(Array{promote_type(T,S)}(undef,size(M,2)),sparse_transpose(M)',b)

mul!(output,M::MatrixNetwork,b) = mul!(output,sparse_transpose(M)',b)

*(M::Adjoint{<:Any,<:MatrixNetwork{T}}, b::AbstractVector{S}) where {T,S} =
    mul!(Array{promote_type(T,S)}(undef,size(M.parent,2)),sparse_transpose(M.parent),b)

mul!(output,M::Adjoint{<:Any,<:MatrixNetwork},b) = mul!(output,sparse_transpose(M.parent),b)

sparse(M::Adjoint{<:Any,<:MatrixNetwork}) = sparse_transpose(M.parent)
copy(M::Adjoint{<:Any,<:MatrixNetwork}) = sparse_transpose(M.parent)

"""
`is_empty`
==========

Return true if the graph is the empty graph and
false otherwise.

Functions
---------
-`is_empty(A::MatrixNetwork)`

Example
-------
~~~~
is_empty(MatrixNetwork(Int[],Int[],0))
is_empty(erdos_renyi_undirected(0,0))
~~~~
"""
is_empty(A::MatrixNetwork) = size(A,1) == 0

"""
`is_undirected`
===============

Check the matrix associated with a matrix network
for symmetry.

Input
-----
- `A`: a matrix network

Returns
-------
- `bool` with true indicating the network is undirected
    and the matrix is symmetric
"""
function is_undirected end

function is_undirected(A::MatrixNetwork)
   M = sparse_transpose(A)
   return issymmetric(M)
end

function is_undirected(A::SparseMatrixCSC)
   return issymmetric(A)
end

"""
`is_connected`
==============

Check the matrix associated with a matrix network
for (strong) connectivity

Usage
-----
- `is_connected(A)`

Input
-----
- `A`: a `MatrixNetwork` or `SparseMatrixCSC` class

Returns
-------
- `bool` with true indicating the matrix is strongly connected
and false indicating
"""
function is_connected end

function is_connected(A::Union{MatrixNetwork,SparseMatrixCSC})
    # this is equivalent to maximum with a default value of 0
    return mapreduce(identity, max, strong_components_map(A); init=0) == 1
end

"""
`empty_graph`
================

Returns an empty graph with n vertices and zero edges

Functions
---------
* `A = empty_graph(n)` generates an empty graph on n edges.

Example
-------
~~~~
is_connected(empty_graph(0))
is_connected(empty_graph(1))
~~~~
"""
function empty_graph end

function empty_graph(n::Integer=0)
    return MatrixNetwork(n,ones(Int64,n+1),Array{Int64}(undef, 0),Array{Float64}(undef, 0))
end

"""
`random_edge`
=============
Identify a random edge of a matrix network or sparse matrix.

Functions
---------
* `random_edge(A::MatrixNetwork) -> (ei,ej,ind)`
   gets a random edge/non-zero from the matrix

* `random_edge(A::SparseMatrixCSC) -> (ei,ej,ind)`
   gets a random non-zero from the matrix

Example
-------
~~~~
G = lollipop_graph(5,3)
# count the number of edges we randomly see between the regions
C = Dict{Symbol,Int}()
M = zeros(8,8)
for i=1:1000000
  ei,ej = MatrixNetwork.random_edge(G)
  M[ei,ej] += 1
  if 1 <= ei <= 5 && 1 <= ej <= 5
    C[:Stem] = get(C, :Stem, 0) + 1
  elseif 6 <= ei <= 10 && 6 <= ej <= 10
    C[:Pop] = get(C, :Pop, 0) + 1
  else
    C[:Bridge] = get(C, :Bridge, 0) + 1
  end
end
# 4 edges in stem, 3 edges in pop, 1 edge in bridge
@show C
@show M
~~~~
"""
function random_edge(A::MatrixNetwork)
    ind = rand(1:length(A.ci)) # the index
    ej = A.ci[ind]
    ei = searchsortedlast(A.rp, ind) # uses binary search for efficiency
    @assert ei <= A.n
    return (ei,ej,ind)
end
function random_edge(A::SparseMatrixCSC)
    ind = rand(1:length(A.rowval)) # the index
    ei = A.rowval[ind]
    ej = searchsortedlast(A.colptr, ind) # uses binary search for efficiency
    @assert ej <= A.n
    return (ei,ej,ind)
end

"""
    undirected_edges(A) -> srcs,dsts

Produce lists just for the undirected edges. This assumes you want only
those edges where (target >= source).
"""
function undirected_edges(A::MatrixNetwork)
    ei = Vector{Int64}()
    ej = Vector{Int64}()
    sizehint!(ei, div(A.rp[A.n+1],2))
    sizehint!(ej, div(A.rp[A.n+1],2))
    for i=1:A.n
        for nzi=A.rp[i]:A.rp[i+1]-1
            j = A.ci[nzi]
            if j >= i
                push!(ei, i)
                push!(ej, j)
            end
        end
    end
    return ei,ej
end

"""
    directed_edges(A) -> srcs,dsts

Produce lists just for all edges of the graph, including both sides for
undirected edges. This is essentially the same as findnz for a sparse matrix,
optimized not to return the values.
"""
function directed_edges(A::MatrixNetwork)
    ei = Vector{Int64}()
    ej = Vector{Int64}()
    sizehint!(ei, A.rp[A.n+1])
    sizehint!(ej, A.rp[A.n+1])
    for i=1:A.n
        for nzi=A.rp[i]:A.rp[i+1]-1
            j = A.ci[nzi]
            push!(ei, i)
            push!(ej, j)
        end
    end
    return ei,ej
end


using Plots, Statistics, Distributions
initial=[]
NList=[1000,10000,100000,1000000]
final_res=[]
for N in NList
    res=[]
    for i in 1:N
        num=rand(0:9)
        push!(initial, num)
    end
    for i in 2:N
        if initial[i-1]==4
            push!(res, initial[i])
        end
    end
    push!(final_res, res)
end
p1=histogram(final_res[1], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :red, label="N=1000")
p2=histogram(final_res[2], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :green,label="N=10000")
p3=histogram(final_res[3], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :yellow,label="N=100000")
p4=histogram(final_res[4], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :blue,label="N=1000000")
plot(p1,p2,p3,p4, layout=4,dpi=600,guidefontsize=9, legendfont=6, legend = :outertop)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter6\\6.2\\Figs\\TheFrequencyCruve.png")


using Test, BifurcationKit, ForwardDiff, RecursiveArrayTools, LinearAlgebra
const BK = BifurcationKit
N = 10
M = 5
δ = 1e-8
####################################################################################################
# test the jacobian of the multiple shooting functional using Linear flow
# TODO do example with A matrix and exp(At)
vf(x, p) = x
flow(x, p, t) = exp(t) .* x
dflow(x, p, dx, t) = (t = t, u = flow(x, p, t), du = exp(t) .* dx)
section(x, T) = dot(x[1:N], ones(N))
section(x::BorderedArray, T) = section(vec(x.u[:,:]), T)
par = nothing

fl = BK.Flow(vf, flow, dflow)

probSh = BK.ShootingProblem(M, fl,
	LinRange(0, 1, M+1) |> diff,
	section, false)

poguess = VectorOfArray([rand(N) for ii=1:M])
	po = BorderedArray(poguess, 1.)

dpoguess = VectorOfArray([rand(N) for ii=1:M])
	dpo = BorderedArray(dpoguess, 2.)

# use of AbstractArray structure
pov = vcat(vec(po.u), po.p)
dpov = vcat(vec(dpo.u), dpo.p)
resv = probSh(pov, par)

dresv = probSh(pov, par, dpov; δ = δ)
resad = ForwardDiff.derivative(t -> probSh(pov .+ t .* dpov, par), 0.)
@test norm(resad[1:end-1] - dresv[1:end-1], Inf) < 1e-14

# use of BorderedArray structure
res = probSh(po, par)
@test norm(vec(res.u[:,:]) - resv[1:end-1] , Inf) ≈ 0
@test norm(res.p - resv[end], Inf) ≈ 0

dres = probSh(po, par, dpo; δ = δ)
@test norm(dres.p - dresv[end], Inf) ≈ 0
@test norm(vec(dres.u[:,:]) - dresv[1:end-1], Inf) ≈ 0
####################################################################################################
# test the jacobian of the multiple shooting functional using nonLinear flow
vf(x, p) = x.^2
flow(x, p, t) = x ./ (1 .- t .* x)
dflow(x, p, dx, t) = (t = t, u = flow(x, p, t), du = dx ./ (1 .- t .* x).^2)

fl = BK.Flow(vf, flow, dflow)

probSh = BK.ShootingProblem(M, fl,
	LinRange(0,1,M+1) |> diff ,
	section, false)

poguess = VectorOfArray([rand(N) for ii=1:M])
	po = BorderedArray(poguess, 1.)

dpoguess = VectorOfArray([rand(N) for ii=1:M])
	dpo = BorderedArray(dpoguess, 2.)

# use of AbstractArray structure
pov = vcat(vec(po.u), po.p)
dpov = vcat(vec(dpo.u), dpo.p)
resv = probSh(pov, par)

dresv = probSh(pov, par, dpov; δ = δ)
resad = ForwardDiff.derivative(t -> probSh(pov .+ t .* dpov, par), 0.)
@test norm(resad[1:end-1] - dresv[1:end-1], Inf) < 1e-14

# use of BorderedArray structure
res = probSh(po, par)
@test norm(vec(res.u[:,:]) - resv[1:end-1] , Inf) ≈ 0
@test norm(res.p - resv[end], Inf) ≈ 0

dres = probSh(po, par, dpo; δ = δ)
@test norm(dres.p - dresv[end], Inf) ≈ 0
@test norm(vec(dres.u[:,:]) - dresv[1:end-1], Inf) ≈ 0
####################################################################################################
# test the hyperplane projections for Poincare Shooting
M = 3
normals = [rand(50) for ii=1:M]
for ii=1:M
	normals[ii] /= norm(normals[ii])
end
centers = [rand(50) for ii=1:M]

hyper = BK.SectionPS(normals, centers)

x = 1:50 |>collect .|> Float64 |> vec
x = rand(50)
xb = BK.R(hyper, x, 1)
# test
for ii=1:M
	x2 = BK.E(hyper, xb, ii)
	# test that x2 in Sigma2
	@test dot(x2 - centers[ii], normals[ii]) < 1e-14
	# that we have Rk∘Ek = Id and Ek∘Rk = IdΣ
	@test BK.R(hyper, x2, ii) ≈ xb
end

# test of the derivatives of E and R
dx = rand(50)
_out1 = (BK.R(hyper, x .+ δ .* dx, 1) - BK.R(hyper, x, 1)) ./ δ
_out2 = zero(_out1)
_out2 = BK.dR!(hyper, _out2, dx, 1)

@test norm(_out2 - _out1, Inf) < 1e-5

dx = rand(49)
_out1 = ForwardDiff.derivative(t -> BK.E(hyper, xb .+ t .* dx,1), 0.)
_out2 = BK.dE(hyper, dx, 1)
@test norm(_out2 - _out1, Inf) < 1e-12


"""
  Methods for the manipulation of geometric data.

  # AUTHORSHIP
    * Author    : Eduardo J Alvarez
    * Email     : Edo.AlvarezR@gmail.com
    * Created   : Aug 2017
    * License   : MIT License
"""
module GeometricTools

using Printf
using LinearAlgebra
using Requires
using Statistics
import PyCall
import Dierckx
import Roots
import QuadGK
import HDF5

const patch = PyCall.PyNULL()
const module_path = splitdir(@__FILE__)[1]      # Path to this module
                                                # Type of multidiscretize input
const multidisctype = Vector{Tuple{Float64,Int64,Float64,Bool}}

function __init__()
    copy!(patch, PyCall.pyimport_conda("matplotlib.patches", "patches"))
    @require PyPlot="d330b81b-6aea-500a-939a-2ce795aea3ee" include("GeometricTools_plotting.jl")
end

for header_name in ["vtk", "geometry", "misc", "gridabstract", "airfoil",
                    "surfacing", "plot3d", "vtk_shapes", "conics",
                    "statistics", "linearalgebra", "xdmf",
                    "DEPRECATED"]
    include("GeometricTools_"*header_name*".jl")
end

end # END OF MODULE


export jacobi, jacobi!, gauss_seidel, gauss_seidel!, sor, sor!, ssor, ssor!

import LinearAlgebra.SingularException
import Base: getindex, iterate

function check_diag(A::AbstractMatrix)
    for i = 1 : size(A, 1)
        if iszero(A[i,i])
            throw(SingularException(i))
        end
    end
end

"""
    jacobi(A, b) -> x

Same as [`jacobi!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
jacobi(A::AbstractMatrix, b; kwargs...) = jacobi!(zerox(A, b), A, b; kwargs...)

"""
    jacobi!(x, A::AbstractMatrix, b; maxiter=10) -> x

Performs exactly `maxiter` Jacobi iterations.

Allocates a single temporary vector and traverses `A` columnwise.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function jacobi!(x, A::AbstractMatrix, b; maxiter::Int=10)
    check_diag(A)
    iterable = DenseJacobiIterable(A, x, similar(x), b, maxiter)
    for _ = iterable end
    x
end

mutable struct DenseJacobiIterable{matT,vecT,solT,rhsT}
    A::matT
    x::solT
    next::vecT
    b::rhsT
    maxiter::Int
end

start(::DenseJacobiIterable) = 1
done(it::DenseJacobiIterable, iteration::Int) = iteration > it.maxiter
function iterate(j::DenseJacobiIterable, iteration::Int=start(j))
    if done(j, iteration) return nothing end

    n = size(j.A, 1)

    copyto!(j.next, j.b)

    # Computes next = b - (A - D)x
    for col = 1 : n
        @simd for row = 1 : col - 1
            @inbounds j.next[row] -= j.A[row, col] * j.x[col]
        end

        @simd for row = col + 1 : n
            @inbounds j.next[row] -= j.A[row, col] * j.x[col]
        end
    end

    # Computes x = D \ next
    for col = 1 : n
        @inbounds j.x[col] = j.next[col] / j.A[col, col]
    end

    nothing, iteration + 1
end

"""
    gauss_seidel(A, b) -> x

Same as [`gauss_seidel!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
gauss_seidel(A::AbstractMatrix, b; kwargs...) = gauss_seidel!(zerox(A, b), A, b; kwargs...)

"""
    gauss_seidel!(x, A::AbstractMatrix, b; maxiter=10) -> x

Performs exactly `maxiter` Gauss-Seidel iterations.

Works fully in-place and traverses `A` columnwise.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function gauss_seidel!(x, A::AbstractMatrix, b; maxiter::Int=10)
    check_diag(A)
    iterable = DenseGaussSeidelIterable(A, x, b, maxiter)
    for _ = iterable end
    x
end

mutable struct DenseGaussSeidelIterable{matT,solT,rhsT}
    A::matT
    x::solT
    b::rhsT
    maxiter::Int
end

start(::DenseGaussSeidelIterable) = 1
done(it::DenseGaussSeidelIterable, iteration::Int) = iteration > it.maxiter

function iterate(s::DenseGaussSeidelIterable, iteration::Int=start(s))
    if done(s, iteration) return nothing end

    n = size(s.A, 1)

    for col = 1 : n
        @simd for row = 1 : col - 1
            @inbounds s.x[row] -= s.A[row, col] * s.x[col]
        end

        s.x[col] = s.b[col]
    end

    for col = 1 : n
        @inbounds s.x[col] /= s.A[col, col]
        @simd for row = col + 1 : n
            @inbounds s.x[row] -= s.A[row, col] * s.x[col]
        end
    end

    nothing, iteration + 1
end

"""
    sor(A, b, ω::Real) -> x

Same as [`sor!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
sor(A::AbstractMatrix, b, ω::Real; kwargs...) =
    sor!(zerox(A, b), A, b, ω; kwargs...)

"""
    sor!(x, A::AbstractMatrix, b, ω::Real; maxiter=10) -> x

Performs exactly `maxiter` SOR iterations with relaxation parameter `ω`.

Allocates a single temporary vector and traverses `A` columnwise.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function sor!(x, A::AbstractMatrix, b, ω::Real; maxiter::Int=10)
    check_diag(A)
    iterable = DenseSORIterable(A, x, similar(x), b, ω, maxiter)
    for _ = iterable end
    x
end

mutable struct DenseSORIterable{matT,solT,vecT,rhsT,numT}
    A::matT
    x::solT
    tmp::vecT
    b::rhsT
    ω::numT
    maxiter::Int
end

start(::DenseSORIterable) = 1
done(it::DenseSORIterable, iteration::Int) = iteration > it.maxiter
function iterate(s::DenseSORIterable, iteration::Int=start(s))
    if done(s, iteration) return nothing end

    n = size(s.A, 1)

    for col = 1 : n
        @simd for row = 1 : col - 1
            @inbounds s.tmp[row] -= s.A[row, col] * s.x[col]
        end

        s.tmp[col] = s.b[col]
    end

    for col = 1 : n
        @inbounds s.x[col] += s.ω * (s.tmp[col] / s.A[col, col] - s.x[col])
        @simd for row = col + 1 : n
            @inbounds s.tmp[row] -= s.A[row, col] * s.x[col]
        end
    end

    nothing, iteration + 1
end

"""
    ssor(A, b, ω::Real) -> x

Same as [`ssor!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
ssor(A::AbstractMatrix, b, ω::Real; kwargs...) =
    ssor!(zerox(A, b), A, b, ω; kwargs...)

"""
    ssor!(x, A::AbstractMatrix, b, ω::Real; maxiter=10) -> x

Performs exactly `maxiter` SSOR iterations with relaxation parameter `ω`. Each iteration
is basically a forward *and* backward sweep of SOR.

Allocates a single temporary vector and traverses `A` columnwise.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function ssor!(x, A::AbstractMatrix, b, ω::Real; maxiter::Int=10)
    check_diag(A)
    iterable = DenseSSORIterable(A, x, similar(x), b, ω, maxiter)
    for _ = iterable end
    x
end

mutable struct DenseSSORIterable{matT,solT,vecT,rhsT,numT}
    A::matT
    x::solT
    tmp::vecT
    b::rhsT
    ω::numT
    maxiter::Int
end

start(::DenseSSORIterable) = 1
done(it::DenseSSORIterable, iteration::Int) = iteration > it.maxiter
function iterate(s::DenseSSORIterable, iteration::Int=start(s))
    if done(s, iteration) return nothing end

    n = size(s.A, 1)

    for col = 1 : n
        @simd for row = 1 : col - 1
            @inbounds s.tmp[row] -= s.A[row, col] * s.x[col]
        end

        s.tmp[col] = s.b[col]
    end

    for col = 1 : n
        @inbounds s.x[col] += s.ω * (s.tmp[col] / s.A[col, col] - s.x[col])
        @simd for row = col + 1 : n
            @inbounds s.tmp[row] -= s.A[row, col] * s.x[col]
        end
    end

    for col = n : -1 : 1
        s.tmp[col] = s.b[col]
        @simd for row = col + 1 : n
            @inbounds s.tmp[row] -= s.A[row, col] * s.x[col]
        end
    end

    for col = n : -1 : 1
        @simd for row = 1 : col - 1
            @inbounds s.tmp[row] -= s.A[row, col] * s.x[col]
        end

        @inbounds s.x[col] += s.ω * (s.tmp[col] / s.A[col, col] - s.x[col])
    end

    nothing, iteration + 1
end


#=

  The BUGS Book
  Example 3.4.1 Three coins (page 45)
  """
  Suppose I have 3 coins in my pocket. The coins may be either fair, biased 3:1 in
  favour of heads, or 3:1 in favour of tails, but I don't know how many of
  each type there are among the 3 coins. I randomly select 1 coin att toss it once,
  observing a head. What is the posterior distribution of the probability of a head?
  """

  Note that JAGS has the somewhat weird structure of dcat etc.

  ~/jags/bugs_book_3_4_1.R
  """
  data {
    y <- 1
  }

  model {

    y ~ dbern(theta.true)
    theta.true <- theta[coin]
    coin ~ dcat(p[])
    for(i in 1:3) {
      p[i] <- 1/3
      theta[i] <- 0.25*i
      coin.prob[i] <- equals(coin, i)
    }
  }

  Output:
               Mean     SD  Naive SE Time-series SE
coin         2.3320 0.7458 0.0026369      0.0026369
coin.prob[1] 0.1672 0.3732 0.0013194      0.0013194
coin.prob[2] 0.3335 0.4715 0.0016669      0.0016669
coin.prob[3] 0.4993 0.5000 0.0017678      0.0017590
theta.true   0.5830 0.1865 0.0006592      0.0006592
  """

  Summary Statistics
    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

          coin    2.3303    0.7518     0.0075    0.0101   4837.8060    1.0000      321.8124
      theta[1]    0.2500    0.0000     0.0000    0.0000         NaN       NaN           NaN
      theta[2]    0.5000    0.0000     0.0000    0.0000         NaN       NaN           NaN
      theta[3]    0.7500    0.0000     0.0000    0.0000         NaN       NaN           NaN
  coin_prob[1]    0.1720    0.3774     0.0038    0.0049   4883.5182    1.0000      324.8532
  coin_prob[2]    0.3257    0.4687     0.0047    0.0064   4998.4834    0.9999      332.5007
  coin_prob[3]    0.5023    0.5000     0.0050    0.0069   4914.2258    1.0000      326.8959
    theta_true    0.5826    0.1880     0.0019    0.0025   4837.8060    1.0000      321.8124
             y    1.0000    0.0000     0.0000    0.0000         NaN       NaN           NaN

  Distributions of variable coin (num:0)
  3.00000 =>    4996  (0.499600)
  2.00000 =>    3355  (0.335500)
  1.00000 =>    1649  (0.164900)

  Cf ~/webppl/bugs_book_3_4_1.wppl

=#

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

@model function bugs_book_3_4_1()
    n = 3
    coin ~ DiscreteUniform(1,3) #  Select one coin. Note: JAGS has 1-based indexing.
    p = 1/3
    coin_prob = tzeros(3)
    theta = tzeros(3)
    for i in 1:n
        theta[i] ~ Dirac(0.25*i)
        coin_prob[i] ~ Dirac(coin == i) # Is this is the coin
    end

    theta_true ~ Dirac(theta[coin])
    y ~ flip(theta_true)

    true ~ Dirac(y == 1) #  We observe a head.

end

model = bugs_book_3_4_1()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
chns = sample(model, SMC(10_000), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,5), 10_000)

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

show_var_dist_pct(chns, :coin)


module SIR

using RecordedArrays
using Random
using Test
using EvolutionaryModelingTools

# parameters
clock = ContinuousClock(300.0)
β = 0.001
γ = 0.01

# state
S = recorded(DynamicEntry, clock, 100)
I = recorded(DynamicEntry, clock, 1)
R = recorded(DynamicEntry, clock, 0)

# simulate with this package
@reaction infection begin
    β * S * I
    (S[] -= 1; I[] += 1)
end

@reaction recovery begin
    γ * I
    (I[] -= 1; R[] += 1)
end

ps = (; β, γ, S, I, R)
rs = (infection, recovery)

Random.seed!(1)
init!(clock)

ps_updated = gillespie(clock, ps, rs)[1]

# simulate manually
Random.seed!(1)
init!(clock)

for _ in clock
    infection_rate = β * S * I
    recovery_rate = γ * I

    summed = infection_rate + recovery_rate

    τ = -log(rand()) / summed # compute time interval

    increase!(clock, τ) # update current time

    # sample a reaction and adjust population size
    rn = rand() * summed
    if rn < infection_rate
        S[] -= 1
        I[] += 1
    else
        I[] -= 1
        R[] += 1
    end
end

for sym in (:S, :I, :R)
    @eval @test state($sym) == state(ps_updated.$sym)
    @eval @test getts(getentries($sym)) == getts(getentries(ps_updated.$sym))
    @eval @test getvs(getentries($sym)) == getvs(getentries(ps_updated.$sym))
end

end


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

using Markdown
using InteractiveUtils

# ╔═╡ 8c102f1e-e978-4ed3-8620-830c46bb566f
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ 0ae615d5-739f-456f-be76-d7d3307577bb
begin
	using LinearAlgebra
	using StatsBase
	using CSV
	using DataFrames
	using RDatasets
	using Plots
	using StatsPlots
	plotly(size=(680, 500))
end

# ╔═╡ 9bbf3d20-0b30-11ec-1917-ebba9c2ee5f0
md"
# Principal Component Analysis

JingYu Ning

**PCA** is a promissing **decorrelating dimension reduction.**
"

# ╔═╡ d67717f9-b5d6-4f84-ad80-31fe0c93b014
md"
## Let's start from a simple example

dataset: [Height and shoe size](https://osf.io/ja9dw/)

It is said that the length of ones foot is corelated with ones height...
"

# ╔═╡ a4d300fb-f712-4045-a8b8-7bf3b4b55ecc
md"
### Show data
"

# ╔═╡ 820f64a0-287d-40fa-b5f5-2cb35dc91658
SIG = "⚡";

# ╔═╡ 98d1834d-f61e-4b28-a6bc-9dec42a1bff8
begin
	SIG
	raw = CSV.read(joinpath(@__DIR__, "../data/HeightAndShoeSize.csv"), DataFrame)
	raw = raw[!, 2:end]
	data = copy(raw)
end

# ╔═╡ 53400761-545e-4f2a-ae75-80694a43dd92
@df data scatter(
	:height, :shoe_size, group=:sex,
	xlabel="height", ylabel="shoe size",
	legend=:topleft
)

# ╔═╡ 033440c0-efc2-42c4-918a-e0d966e27d3f
md"
### 1. Move the data to the origin
"

# ╔═╡ b3647ffa-d904-4bee-9327-6dec393702e5
begin
	for i in 2:3
		data[!, i] .-= mean(data[!, i])
		raw[!, i] .-= mean(raw[!, i])
	end
	@df data scatter(
		:height, :shoe_size, group=:sex,
		xlabel="height", ylabel="shoe size",
		legend=:topleft,
		xlim=(-40, 40), ylim=(-40, 40)
	)
end

# ╔═╡ 4f771fd7-84be-43f9-a9b2-78d40db4f91c
md"
### 2. SVD

$X = U \Sigma V^\dagger$
"

# ╔═╡ 6b6365a7-486a-4533-bd1e-d6c5726f3b35
u, s, v = svd(Matrix(data[!, 2:3])')

# ╔═╡ ddfc6838-cf88-4aa2-a95c-5d4c95a487ae
md"
### 3. Principal components

Choose the orthogonal basis `u` as the **principal components**
"

# ╔═╡ 66806d5b-55b7-4a4f-b592-2fe6304a7795
begin
	scale = 10
	p = @df data scatter(
		:height, :shoe_size, group=:sex,
		xlabel="height", ylabel="shoe size",
		legend=:topleft,
		xlim=(-40, 40), ylim=(-40, 40)
	)
	for i in 1:size(u, 1)
		pci = scale * u[:, i]
		plot!(p, [0, pci[1]], [0, pci[2]], label="pc$i", lw=5)
	end
	p
end

# ╔═╡ de3026a5-109e-458c-9e38-8e6ab43a81c6
md"
### 4. Apply the basis

$U^T X$
"

# ╔═╡ 5530ddd4-df5e-4308-9dfd-492f68db9f94
begin
	data[!, 2:3] .= (u' * Matrix(data[!, 2:3])')'
	@df data scatter(
		:height, :shoe_size, group=:sex,
		xlabel="pc1", ylabel="pc2",
		legend=:topleft,
		xlim=(-40, 40), ylim=(-40, 40)
	)
end

# ╔═╡ c7490656-5f5f-41a5-8a56-e3341880d855
md"
### 5. Rescale

$\Sigma^{-1} U^T X$
"

# ╔═╡ 89767b3c-63eb-4bdb-9f96-2289954908fa
begin
	data[!, 2:3] .= (inv(diagm(s)) * Matrix(data[!, 2:3])')'
	@df data scatter(
		:height, :shoe_size, group=:sex,
		xlim=(-0.5, 0.5), ylim=(-0.5, 0.5)
	)
end

# ╔═╡ 4ae5b880-1c74-4d37-b1eb-6f6f847b39a0
md"
$I = \Sigma^{-1} U^T X V$
"

# ╔═╡ 6a3f7794-3776-48cf-b188-ee892078c148
Matrix(data[!, 2:3])' * v

# ╔═╡ edcc9e1c-9cc0-4ead-ae6f-4c68bc509470
md"
### 6. Dimension reduction

We want to map those data into lower dimentions with less intormance loss.

Here, $U$ is a basis in the original high dimentional space. We will have least informance loss if project data into one of the axis of $U$, while have most loss if into another.

$U = [v_1, v_2]$
"

# ╔═╡ 18c0c325-c107-42e8-84a3-587580cc5772
md"
By choosing a new basis, we can therefore drop one dimention with little information lossed.
"

# ╔═╡ 88526e35-67c0-430a-a6ef-6454cdab7d2d
begin
	plot(size=(680, 200), yticks=[])
	scatter!((u'*Matrix(raw[!, 2:3])')[1, :], zeros(nrow(raw)), label="pc1")
	scatter!((u'*Matrix(raw[!, 2:3])')[2, :], ones(nrow(raw)), label="pc2")
end

# ╔═╡ 4d5c2e27-8a2c-486b-97bd-65ce4d108322
md"
And the singular value is the information of varience.
"

# ╔═╡ 5312759e-6c75-411b-b28f-9de6f6bfa65b
begin
	plot(size=(680, 200), yticks=[])
	scatter!(
		(inv(diagm(s))*u'*Matrix(raw[!, 2:3])')[1, :], zeros(nrow(raw)),
		label="pc1"
	)
	scatter!(
		(inv(diagm(s))*u'*Matrix(raw[!, 2:3])')[2, :], ones(nrow(raw)),
		label="pc2"
	)
end

# ╔═╡ d5fe909b-784a-443e-9e2e-c00c9e8cd611
md"
We constract a mapping that maps the data into `Z` space

$W = \Sigma^{-1}U^T$
"

# ╔═╡ e3e19d9e-767a-486b-8626-108c5f70a63b
w = inv(diagm(s)) * u'

# ╔═╡ d315c4c8-9945-409a-90b5-f08369665b25
md"The mapping is roughly the overall picture of the data."

# ╔═╡ 9658209d-57b5-4d68-a031-e15711d8cc92
md"
## Iris dataset
"

# ╔═╡ 7475abbd-6fdb-423f-8981-be4461df6648
iris = dataset("datasets", "iris")

# ╔═╡ 663406ed-49d2-439b-8272-69c0b29546da
begin
    cols = names(iris)

    function gen_pots()
        plots = Matrix{Plots.Plot}(undef, 4, 4)
        for i in 1:4, j in 1:4
            if i == j
                plots[i, j] = @df(iris, histogram(cols(i)))
            elseif i < j
                plots[i, j] = @df(iris, histogram2d(
                    cols(i), cols(j),
                    nbins=20, color=:coolwarm, colorbar=false
                ))
            else
                plots[i, j] =  @df(iris, scatter(cols(i), cols(j)))
            end

            (i==1) && (plots[i, j] = plot!(ylabel=cols[j])) # add x label
            (j==4) && (plots[i, j] = plot!(xlabel=cols[i])) # add y label
        end

        return plots
    end
end

# ╔═╡ f79ca361-a2ca-4e0e-b920-36936d01fe8b
plot(gen_pots()..., layout=(4, 4), legend=false)

# ╔═╡ 7a284458-83bb-48fb-863e-ff385a3ba27d
md"
### 1. Move the data to the origin
"

# ╔═╡ da5d583f-b9d8-4c96-892c-c3a018d83890
begin
	for i in 1:(ncol(iris)-1)
		iris[!, i] .-= mean(iris[!, i])
	end
	plot(gen_pots()..., layout=(4, 4), legend=false)
end

# ╔═╡ 221d3ebf-e734-4a7d-8811-d8874311aa21
md"
### 2. SVD
"

# ╔═╡ 58db4367-514b-4ce8-b994-a13d73af93e9
U, S, V = svd(Matrix(iris[!, 1:(ncol(iris)-1)])')

# ╔═╡ 4a0d265d-9c91-4095-899b-20e4c6765fbf
md"
### 3. Construct the mapping

$W = \Sigma^{-1}U^T$
"

# ╔═╡ 1227c477-2fa5-4929-b64b-aeaf5745c52d
W = inv(diagm(S)) * U'

# ╔═╡ c1b885d9-bb2b-486f-ad22-219e4ede6025
begin
	iris[!, 1:(ncol(iris)-1)] .= (W * Matrix(iris[!, 1:(ncol(iris)-1)])')'
	plot(gen_pots()..., layout=(4, 4), legend=false)
end

# ╔═╡ Cell order:
# ╟─9bbf3d20-0b30-11ec-1917-ebba9c2ee5f0
# ╟─8c102f1e-e978-4ed3-8620-830c46bb566f
# ╟─0ae615d5-739f-456f-be76-d7d3307577bb
# ╟─d67717f9-b5d6-4f84-ad80-31fe0c93b014
# ╟─a4d300fb-f712-4045-a8b8-7bf3b4b55ecc
# ╟─98d1834d-f61e-4b28-a6bc-9dec42a1bff8
# ╟─53400761-545e-4f2a-ae75-80694a43dd92
# ╠═820f64a0-287d-40fa-b5f5-2cb35dc91658
# ╟─033440c0-efc2-42c4-918a-e0d966e27d3f
# ╟─b3647ffa-d904-4bee-9327-6dec393702e5
# ╟─4f771fd7-84be-43f9-a9b2-78d40db4f91c
# ╟─6b6365a7-486a-4533-bd1e-d6c5726f3b35
# ╟─ddfc6838-cf88-4aa2-a95c-5d4c95a487ae
# ╟─66806d5b-55b7-4a4f-b592-2fe6304a7795
# ╟─de3026a5-109e-458c-9e38-8e6ab43a81c6
# ╟─5530ddd4-df5e-4308-9dfd-492f68db9f94
# ╟─c7490656-5f5f-41a5-8a56-e3341880d855
# ╟─89767b3c-63eb-4bdb-9f96-2289954908fa
# ╟─4ae5b880-1c74-4d37-b1eb-6f6f847b39a0
# ╟─6a3f7794-3776-48cf-b188-ee892078c148
# ╟─edcc9e1c-9cc0-4ead-ae6f-4c68bc509470
# ╟─18c0c325-c107-42e8-84a3-587580cc5772
# ╟─88526e35-67c0-430a-a6ef-6454cdab7d2d
# ╟─4d5c2e27-8a2c-486b-97bd-65ce4d108322
# ╟─5312759e-6c75-411b-b28f-9de6f6bfa65b
# ╟─d5fe909b-784a-443e-9e2e-c00c9e8cd611
# ╟─e3e19d9e-767a-486b-8626-108c5f70a63b
# ╟─d315c4c8-9945-409a-90b5-f08369665b25
# ╟─9658209d-57b5-4d68-a031-e15711d8cc92
# ╟─663406ed-49d2-439b-8272-69c0b29546da
# ╠═7475abbd-6fdb-423f-8981-be4461df6648
# ╟─f79ca361-a2ca-4e0e-b920-36936d01fe8b
# ╟─7a284458-83bb-48fb-863e-ff385a3ba27d
# ╟─da5d583f-b9d8-4c96-892c-c3a018d83890
# ╟─221d3ebf-e734-4a7d-8811-d8874311aa21
# ╟─58db4367-514b-4ce8-b994-a13d73af93e9
# ╟─4a0d265d-9c91-4095-899b-20e4c6765fbf
# ╟─1227c477-2fa5-4929-b64b-aeaf5745c52d
# ╟─c1b885d9-bb2b-486f-ad22-219e4ede6025


using JuMP, EAGO, Gurobi, CSV, DataFrames, Ipopt

m = Model(with_optimizer(
    EAGO.Optimizer,
    relaxed_optimizer = Gurobi.Optimizer(OutputFlag = 0),
))

# ----- Variables ----- #
x_Idx = Any[1, 2]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[1], -10.0)
JuMP.set_upper_bound(x[1], 10.0)
JuMP.set_lower_bound(x[2], -10.0)
JuMP.set_upper_bound(x[2], 10.0)

# ----- Objective ----- #
@NLobjective(
    m,
    Min,
    ((-5 + sqrt((x[1])^2)) * (-5 + sqrt((x[1])^2)) +
     (-5 + sqrt((x[2])^2)) * (-5 + sqrt((x[2])^2)))
)

JuMP.optimize!(m)
val = objective_value(m)
bnd = objective_bound(m)
gap = bnd - val
backend_opt = backend(m).optimizer.model.optimizer
last_relaxed_opt = backend_opt.relaxed_optimizer
run_time = backend_opt._run_time
println("run time: $run_time")


function [X,J,dXdx,dXdxI]=JacobiDG2(DG,F,Topo,Param)
ksi=DG.xwX;
eta=DG.xwY;
nX=DG.OrdPolyX+1;
nY=DG.OrdPolyY+1;
X=zeros(nX,nY,3);
dXdx=zeros(nX,nY,2,2);
dXdxI=zeros(nX,nY,2,2);
J=zeros(nX,nY);
for j=1:nY
  for i=1:nX
    X(i,j,1:2)=0.25*((1-ksi(i))*(1-eta(j))*F.P(1:2,1)...
                    +(1+ksi(i))*(1-eta(j))*F.P(1:2,2)...
                    +(1+ksi(i))*(1+eta(j))*F.P(1:2,3)...
                    +(1-ksi(i))*(1+eta(j))*F.P(1:2,4));
    [X(i,j,2),dXdx(i,j,2,2)]=Topo(X(i,j,1),X(i,j,2),Param);
  end
end

dXdx(:,:,1,1)=DG.DSX*X(:,:,1);
dXdx(:,:,2,1)=DG.DSX*X(:,:,2);
dXdx(:,:,1,2)=X(:,:,1)*DG.DSY';
dXdx(:,:,2,2)=X(:,:,2)*DG.DSY';
for j=1:nY
  for i=1:nX
    J(i,j)=det(reshape(dXdx(i,j,:,:),2,2));
    dXdxI(i,j,:,:)=inv(reshape(dXdx(i,j,:,:),2,2))*J(i,j);
  end
end
end


function fit_cumulative(cum,time)
    model(x,p) = @. (1 - exp(-p[1]*x)) * p[2]
    p0 = [0.5,1]
    fitting = curve_fit(model, cum, time, p0)
    # return DataFrame(Coef = fitting.param[1], Scale = fitting.param[2])
    # return (fitting.param...,)
    return fitting.param
end

ParetoModel(x,p) = @. (1 - ((p[1]/x)^p[2])) * p[3]
function fit_pareto(x,y)
    p0 = [0.5,0.5,0.5]
    fitting = curve_fit(ParetoModel, x, y, p0)
    return fitting.param
end

ExpoModel(x,p) = @. ((1 - exp(-p[1]*x)) * p[2]) + p[3]
function fit_ExpoModel(x,y)
    p0 = [0.5,0.5,0.5]
    fitting = curve_fit(ExpoModel, x, y, p0)
    return fitting.param
end

function MVTtangent(coef, P)
    # P is the point where the line pass through
    # cumulative to be tangential to
    c(x) = @. (1 - exp(-coef[1]*x))*coef[2]
    #derivative of the cum, equal to the slope
    d(x) = @. coef[1] * exp(-coef[1]*x) *coef[2]
    #solve using slope-intercept form of the line
    # y-y0 = m(x-x0)
    # x0 has to be a point on the curve so express it in this way
    # y-c(x0) = d(x0)(x-x0)
    # solve passing for the point P
    # yp -c(x0) = d(x0)(xp-x0)
    # yp -c(x0) - d(x0)(xp-x0) = 0
    # f(x) = P[2] - (1 - exp(-coef[1]*x))*coef[2] - (coef[1] * exp(-coef[1]*x))*coef[2]*(P[1]-x)
    f(x) = P[2] - c(x) - d(x)*(P[1]-x)
    x0 = Roots.find_zero(f,0) # search x0 in a reasonable range of trial time
    #find inercept using P and m = d(x0)
    # y = mx +q
    # P[2] = d(x0)*P[1] + q
    # q = -d(x0)*P[1] - P[2]
    q = -d(x0)*P[1] - P[2]
    # sol(x) = @. x * d(x0) + q
    # return DataFrame(OptLeave = x0, Slope = d(x0), Intercept = q)
    return (x0,d(x0),q)
end

GainModel(x,p) = @. (1 - exp(-p[1]*x)) *p[2]

# control = filter(r->r.Treatment == "Control", timedf)
# coef = fit_cumulative(control.Time, control.CumRew)
# control2 = filter(r->r.Treatment == "Control", timedf2)
# coef2 = fit_cumulative(control2.Time, control2.CumRew)
#
# xaxis = 0:0.5:35
# plot(xaxis,GainModel(xaxis,coef), label = "TrialTime", legend = :bottomright)
# plot!(xaxis,GainModel(xaxis,coef2), label = "ForagingTime", linecolor = :purple)


# Uses new formulation from limbdark paper.
include("cel_bulirsch.jl")

function s2(r::T,b::T) where {T <: Real}
# For now, just compute linear component:
Lambda1 = zero(typeof(b))
if b >= 1.0+r ||  r == 0.0
  # No occultation:
  Lambda1 = zero(b)  # Case 1
elseif b <= r-1.0
  # Full occultation:
  Lambda1 = zero(b)  # Case 11
else
  if b == 0
    Lambda1 = -2/3*sqrt(1.0-r^2)^3 # Case 10
  elseif b==r
    if r == 0.5
      Lambda1 = 1/3-4/(9pi) # Case 6
    elseif r < 0.5
      m = 4r^2
      Lambda1 = 1/3+2/(9pi)*cel_bulirsch(m,one(r),m-3,(1-m)*(2m-3)) # Case 5
    else
      m = 4r^2; minv = inv(m)
      Lambda1 = 1/3+1/(9pi*r)*cel_bulirsch(minv,one(r),m-3,1-m) # Case 7
    end
  else
    onembpr2 = 1-(b+r)^2; onembmr2=1-(b-r)^2; fourbr = 4b*r
    k2 = onembpr2/fourbr+1
    if (b+r) > 1.0 # k^2 < 1, Case 2, Case 8
      k2c = -onembpr2/fourbr; kc = sqrt(k2c)
      Eofk = cel_bulirsch(k2,kc,one(b),one(b),k2c) # Complete elliptic integral of second kind
      Lambda1 = onembmr2*(k2c*(b-r)*(b+r)*cel_bulirsch(k2,kc,(b-r)^2*k2c,zero(b),3one(b))
          -(3-6*r^2-2*b*r)*cel_bulirsch(k2,kc,one(b),one(b),zero(b))-fourbr*Eofk)/(9*pi*sqrt(b*r))
    elseif (b+r) < 1.0  # k^2 > 1, Case 3, Case 9
      k2inv = inv(k2); k2c =onembpr2/onembmr2; kc = sqrt(k2c)
      Eofk = cel_bulirsch(k2inv,kc,one(b),one(b),k2c) # Complete elliptic integral of second kind
      bmrdbpr = (b-r)/(b+r);
      mu = 3bmrdbpr/onembmr2
      p = bmrdbpr^2*onembpr2/onembmr2
      Lambda1 = 2*sqrt(onembmr2)*(onembpr2*cel_bulirsch(k2inv,kc,p,1.0+mu,p+mu)
             -(4-7r^2-b^2)*Eofk)/(9*pi)
    else
      # b+r = 1 or k^2=1, Case 4 (extending r up to 1)
      Lambda1 = 2/(3pi)*acos(1.-2.*r)-4/(9pi)*(3+2r-8r^2)*sqrt(r*b)-2/3*convert(typeof(b),r>.5)
    end
  end
end
flux = 1.0-1.5*Lambda1-convert(typeof(b),r>b)
return flux
end


using Convex, ECOS

function go()
    # A = [-1 -2  6  3 8;
    #      2  3 -2 -1 3]
    # x = Variable(2, Positive())

    A = [2 0 -2 0 3;
         0 5 -3 0 3;
         0 0 5 -1 8;
         0 -1 0 5 8]
    x = Variable(4, Positive())

    prob = maximize(sum(log((A' * x)[1:4])),
                    [sum(x) == 100,
                     (A' * x)[5] == 500
                     ])

    solve!(prob, ECOS.Optimizer)

    A' * round.(evaluate(x))
end

Int(prod(go()[1:4]))


include("../mathutils/Helpers.jl")
include("../termstructures/YieldCurve.jl")
include("../models/StochasticProcess.jl")

using Roots

# the actual data structure for the model
struct HullWhiteModel{T<:AbstractFloat} <: StochasticProcess
    yieldCurve::YieldCurve
    meanReversion::T
    volatilityTimes::Array{T}
    volatilityValues::Array{T}
    y_::Array{T}
end

# constructor

function HullWhiteModel(yieldCurve, meanReversion, volatilityTimes::Array, volatilityValues::Array)
    #
    t0 = 0.0
    y0 = 0.0
    y_ = zeros(size(volatilityTimes))
    for i = 1:size(y_)[1]
        y_[i] = (GPrime(meanReversion,t0,volatilityTimes[i])^2) * y0 +
            (volatilityValues[i]^2) *
            (1.0 - exp(-2*meanReversion*(volatilityTimes[i]-t0))) /
            (2.0 * meanReversion)
        t0 = volatilityTimes[i]
        y0 = y_[i]
    end
    return HullWhiteModel(yieldCurve,meanReversion,volatilityTimes,volatilityValues,y_)
end

# auxilliary functions/methods

G(meanReversion,t,T) =
    (1.0 - exp(-meanReversion*(T-t))) / meanReversion

GPrime(meanReversion,t,T) = exp(-meanReversion*(T-t))

G(self::HullWhiteModel,t,T) = G(self.meanReversion,t,T)

GPrime(self::HullWhiteModel,t,T) = GPrime(self.meanReversion,t,T)

function y(self::HullWhiteModel,t)
    idx = searchsortedfirst(self.volatilityTimes,t)  # first element of tupel
    t0 = (idx==1) ? 0.0 : self.volatilityTimes[idx-1]
    y0 = (idx==1) ? 0.0 : self.y_[idx-1]
    s1 = self.volatilityValues[min(idx,size(self.volatilityValues)[1])]  # flat extrapolation
    y1 = (GPrime(self,t0,t)^2) * y0 +
            s1^2 * (1.0 - exp(-2*self.meanReversion*(t-t0))) /
            (2.0 * self.meanReversion)
    return y1
end

function riskNeutralExpectationX(self::HullWhiteModel, t, xt, T)
    # E[x] = G'(t,T)x + \int_t^T G'(u,T)y(u)du
    f(u) = GPrime(self,u,T)*y(self,u)
    # use Simpson's rule to approximate integral, this should better be solved analytically
    integral = (T-t) / 6 * (f(t) + 4*f((t+T)/2) + f(T))
    return GPrime(self,t,T)*xt + integral
end

function sigma(self::HullWhiteModel,t)   # Todo test this method
    # find idx s.t. t[idx-1] < t <= t[idx]
    idx = searchsortedfirst(self.volatilityTimes,t)
    return self.volatilityValues[min(idx,size(self.volatilityValues)[1])]
end

# bond volatility is used in hybrid model vol adjuster

function zeroBondVolatility(self::HullWhiteModel, t, T)
    # we wrap scalar bond volatility into array to allow
    # for generalisation to multi-factor models
    return [ sigma(self,t)*G(self,t,T) ]
end

function zeroBondVolatilityPrime(self::HullWhiteModel, t, T)
    # we wrap scalar bond volatility into array to allow
    # for generalisation to multi-factor models
    return [ sigma(self,t)*GPrime(self,t,T) ]
end

# model methods

# conditional expectation in T-forward measure
function expectationX(self::HullWhiteModel, t, xt, T)
    return GPrime(self,t,T)*(xt + G(self,t,T)*y(self,t))
end

function varianceX(self::HullWhiteModel, t, T)
    return y(self,T) - GPrime(self,t,T)^2 * y(self,t)
end

function zeroBondPrice(self::HullWhiteModel, t, T, xt)
    G_ = G(self,t,T)
    return discount(self.yieldCurve,T) / discount(self.yieldCurve,t) *
        exp(-G_*xt - 0.5 * G_^2 * y(self,t) )
end

function zeroBondOption(self::HullWhiteModel, expiryTime, maturityTime, strikePrice, callOrPut)
    nu2 = G(self,expiryTime,maturityTime)^2 * y(self,expiryTime)
    P0  = discount(self.yieldCurve, expiryTime)
    P1  = discount(self.yieldCurve, maturityTime)
    return P0 * Black(strikePrice,P1/P0,sqrt(nu2),1.0,callOrPut)
end

function couponBondOption(self::HullWhiteModel, expiryTime, payTimes, cashFlows, strikePrice, callOrPut)
    function objective(x)
        bond = 0
        for i = 1:size(payTimes)[1]
            bond += cashFlows[i] * zeroBondPrice(self,expiryTime,payTimes[i],x)
        end
        return bond - strikePrice
    end
    xStar = find_zero(objective,(-1.0,1.0),Roots.Brent(),xatol=1.0e-8)
    bondOption = 0.0
    for i = 1:size(payTimes)[1]
        strike = zeroBondPrice(self,expiryTime,payTimes[i],xStar)
        bondOption += cashFlows[i] * zeroBondOption(self,expiryTime,payTimes[i],strike,callOrPut)
    end
    return bondOption
end

# stochastic process interface

function stateSize(self::HullWhiteModel)   # dimension of X(t)
    return 2      # [x, s], we also need for numeraire s = \int_0^t r dt
end

function factors(self::HullWhiteModel)   # dimension of W(t)
    return 1
end

function initialValues(self::HullWhiteModel)
    return [ 0.0, 0.0 ]
end

function zeroBond(self::HullWhiteModel, t, T, X, alias)
    return zeroBondPrice(self, t, T, X[1])
end

function numeraire(self::HullWhiteModel, t, X)
    return exp(X[2]) / discount(self.yieldCurve,t)
end

# evolve X(t0) -> X(t0+dt) using independent Brownian increments dW
# t0, dt are assumed float, X0, X1, dW are np.array
function evolve(self::HullWhiteModel, t0, X0, dt, dW, X1)
    x1 = riskNeutralExpectationX(self,t0,X0[1],t0+dt)
    # x1 = X0[0] + (self.y(t0) - self.meanReversion*X0[0])*dt
    nu = sqrt(varianceX(self,t0,t0+dt))
    x1 = x1 + nu*dW[1]
    # s1 = s0 + \int_t0^t0+dt x dt via Trapezoidal rule
    s1 = X0[2] + (X0[1] + x1) * dt / 2
    # gather results
    X1[1] = x1
    X1[2] = s1
    return nothing
end

# the short rate over an integration time period
# this is required for drift calculation in multi-asset and hybrid models
function shortRateOverPeriod(self::HullWhiteModel, t0, dt, X0, X1)
    B_d = discount(self.yieldCurve,t0) / discount(self.yieldCurve,t0 + dt)  # deterministic drift part for r_d
    x_av = 0.5 * (X0[1] + X1[1])
    return log(B_d) / dt + x_av
end


module nav_geo

export cosineSimilarity,deltaF1,deltaF2,destination,forwardAzmuth,haversine,speed

cosineSimilarity(a,b) = (a' * b)/(norm(a)*norm(b))

function deltaF1(args::Array{Float64})
   gradient = Array(Float64,2);

   r = args[1];
   theta = args[2];
   s = args[3];
   gradient[1] = -2*r*(s/(s + 1)*e^-(2*r/(s + 1))*(2 - 0.5*cos(theta + pi))  + 2*e^-(0.25*r^2)) + 1;
   gradient[2] = 0.5*s*e^-(2*r/(s + 1))*sin(theta);
   return gradient;
end

function deltaF2(args::Array{Float64})
   gradient = Array(Float64,2);

   r = args[1];
   theta = args[2];
   s = args[3];
   denom = abs(r -s*(2*cos(theta) + 5))
   dtheta = (2*s*sin(theta)*(r - s*(2*cos(theta) + 5)))/denom + r/4*sin(theta);
   dr = (r - 2*s*cos(theta) - 5*s)/denom - cos(theta)/4;
   gradient[1] = cos(theta)*dr - 1/r*sin(theta)*dtheta;
   gradient[2] = sin(theta)*dr + 1/r*cos(theta)*dtheta;
   return gradient;
end

function destination(args::Array{Float64})
   res = Array(Float64,2);
   srcLat = args[1]*pi/180;
   srcLong = args[2]*pi/180;
   azmuth = args[3]
   distanceInMeters = args[4];
   distanceInRadians = distanceInMeters/earthRadius(args[1]);
   destLat = asin(sin(srcLat)*cos(distanceInRadians) + cos(srcLat)*sin(distanceInRadians)*cos(azmuth));
   dlong = srcLong + atan2(sin(azmuth)*sin(distanceInRadians)*cos(srcLat),cos(distanceInRadians) - sin(srcLat)*sin(destLat));
   destLong = (dlong + 3*pi) % (2*pi) - pi;
   res[1] = destLat*180/pi;
   res[2] = destLong*180/pi;
   return res;
end

function earthRadius(lat)
   a =  6378137.0;
   b =  6356752.3;
   phi = lat*pi/180;
   return sqrt(((a^2*cos(phi))^2 + (b^2*sin(phi))^2)/((a*cos(phi))^2 + (b*sin(phi))^2));
end

function forwardAzmuth(args::Array{Float64})
   lat1 = args[1];
   long1 = args[2];
   lat2 = args[3];
   long2 = args[4];

   phi1 = lat1*pi/180;
   phi2 = lat2*pi/180;
   lam1 = long1*pi/180;
   lam2 = long2*pi/180;

   x = sin(lam2 - lam1)*cos(phi2);
   y = cos(phi1)*sin(phi2) - sin(phi1)*cos(phi2)*cos(lam2-lam1)
   return atan2(x,y) % (2*pi);
end

function haversine(args::Array{Float64})
   lat1 = args[1];
   long1 = args[2];
   lat2 = args[3];
   long2 = args[4];

   phi1 = lat1*pi/180;
   phi2 = lat2*pi/180;
   lam1 = long1*pi/180;
   lam2 = long2*pi/180;

   return 2 * earthRadius(lat1) * asin(sqrt(sin((phi2 - phi1)/2)^2 + cos(phi1) * cos(phi2) * sin((lam2 - lam1)/2)^2))
end

speed(x) = 7*(2/(1 + e^-(abs(x/10)^(5/4))) - 1)

end


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ 245de79e-260c-11ec-082f-bdef59a44c40
using ControlSystems

# ╔═╡ f5e08ada-5f12-45a2-ba56-62af6bbe0a0d
begin
	J = 2.0
	b = 0.04
	K = 1.0
	R = 0.08
	L = 1e-4

	# dc motor model
	s = tf("s")
	P = K/(s*((J*s + b)*(L*s + R) + K^2))

	# Create an array of closed loop systems for different values of kp
	Tcl = TransferFunction[kp*P/(1 + kp * P) for kp = [1, 5, 15]]

	# Closed loop step response
	stepplot(Tcl, label= ["kp = 1" "kp = 5" "kp = 15"])
end

# ╔═╡ 3f85cf01-bcd7-4395-84bf-2c5eb597f497

# ╔═╡ Cell order:
# ╠═245de79e-260c-11ec-082f-bdef59a44c40
# ╠═f5e08ada-5f12-45a2-ba56-62af6bbe0a0d
# ╠═3f85cf01-bcd7-4395-84bf-2c5eb597f497


"""
    getknn(Xtrain, X; k = 1, metric = "eucl")
Return the k nearest neighbors in Xtrain of each row of `X`.
* `Xtrain` : Training X-data.
* `X` : Query X-dta.
* `metric` : Type of distance used for the query.
    Possible values are "eucl" or "mahal".

The distances (not the squared distances) are also returned.
"""
function getknn(Xtrain, X; k = 1, metric = "eucl")
    Xtrain = ensure_mat(Xtrain)
    Xt = ensure_mat(X')
    n, p = size(Xtrain)
    k > n ? k = n : nothing
    if metric == "eucl"
        ztree = BruteTree(Xtrain', Euclidean())
        ind, d = knn(ztree, Xt, k, true)
    elseif metric == "mahal"
        S = Statistics.cov(Xtrain, corrected = false)
        if p == 1
            Uinv = inv(sqrt(S))
        else
            #S = S + Diagonal(1e-10 * ones(p))
            Uinv = LinearAlgebra.inv!(cholesky!(Hermitian(S)).U)
        end
        zXtrain = Xtrain * Uinv
        zX = X * Uinv
        ztree = BruteTree(zXtrain', Euclidean())
        # ztree = BruteTree(Xtraint, Mahalanobis(Sinv))
        # is very slow
        ind, d = knn(ztree, zX', k, true)    # ind and d = lists
    end
    #ind = reduce(hcat, ind)'
    (ind = ind, d = d)
end