using Plots
using BenchmarkTools

#Problem 1(a)
function vfsolve(vnew, kgrid, tolerance, imax)
    α = 0.37
    δ = 0.1
    β = 0.95
    σ = 1.5

    v = vnew .+ 2*tolerance
    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1

    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while maximum(abs.(v - vnew)) > tolerance && i <= imax
        v = vnew;

        (vnew, cartesianindex) = findmax(u .+ β*v', dims = 2);
        i += 1;
    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex, M=i)
end

kupper = 5
klower = 0.05
n = 500
kgrid = collect(range(klower, stop = kupper, length = n))
(v, kprime, kprimeindex, M) = vfsolve(zeros(n), kgrid, 0.001, 1000);
M
#M=95 the output value function is close enough to the true value function
#1(b)
#define a new function that take imax as the only control for the number of iterations
function vfsolve1(vnew, kgrid, imax)
    α = 0.37
    δ = 0.1
    β = 0.95
    σ = 1.5

    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1
    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while i <= imax

        (vnew, cartesianindex) = findmax(u .+ β*vnew', dims = 2);
        i += 1;

    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex)
end

(v1, kprime1, kprimeindex1) = vfsolve1(zeros(n), kgrid, 1);
(v5, kprime5, kprimeindex5) = vfsolve1(zeros(n), kgrid, 5);
(v10, kprime10, kprimeindex10) = vfsolve1(zeros(n), kgrid, 10);
plot(kgrid, v1, label = "v1")
plot!(kgrid, v5, label = "v5")
plot!(kgrid, v10, label = "v10")
plot!(kgrid, v, label = "v")

#1(c)

#define a function to find the value in krgrid which is closest to the input x and return the index
function findindex(x)
    a=findmin(abs.(x .- kgrid), dims=1)
    return a[2]
end

#policy(k) returns the optimal k'
function policy(x)
    e=getindex(findall(kgrid .== x),1)
    return kprime[e]
end

# the function find and plot the time path of k
function findkpath(T, k_0)
    w = findindex(k_0)
    kpath = zeros(T)
    k = kgrid[w]
        for i in 1:T
            k = policy(k)
            kpath[i] = k
        end
    return kpath
end

# # the function find and plot the time path of k

function plotkpath(T, k_0)
w = findindex(k_0)
kpath = zeros(T)
k = kgrid[w]
    for i in 1:T
        k = policy(k)
        kpath[i] = k
    end
time = collect(range(1, T, length = T))
plot(time, kpath, label = "kpath")
end

findkpath(100, 0.5)
plotkpath(100, 0.5)

#the steady state level of capital is around 4
#1(d)
findkpath(100, 4.5)
plotkpath(100, 4.5)
#the steady state level of capital is also around 4
#(e)

#(e)(a)
function vfsolve_E(vnew, kgrid, tolerance, imax)
    α = 0.37
    δ = 0.1
    β = 0.95
    σ = 0.5

    v = vnew .+ 2*tolerance
    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1

    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while maximum(abs.(v - vnew)) > tolerance && i <= imax
        v = vnew;

        (vnew, cartesianindex) = findmax(u .+ β*v', dims = 2);
        i += 1;
    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex, n=i)
end
kupper = 5
klower = 0.05
n = 500
kgrid = collect(range(klower, stop = kupper, length = n))
(v_E, kprime_E, kprimeindex_E, n_E) = vfsolve_E(zeros(n), kgrid, 0.001, 1000)

#(e)(b)
function vfsolve1_E(vnew, kgrid, imax)
    α = 0.37
    δ = 0.1
    β = 0.95
    σ = 0.5

    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1
    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while i <= imax

        (vnew, cartesianindex) = findmax(u .+ β*vnew', dims = 2);
        i += 1;

    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex)
end

(v1_E, kprime1_E, kprimeindex1_E) = vfsolve1_E(zeros(n), kgrid, 1);
(v5_E, kprime5_E, kprimeindex5_E) = vfsolve1_E(zeros(n), kgrid, 5);
(v10_E, kprime10_E, kprimeindex10_E) = vfsolve1_E(zeros(n), kgrid, 10);
plot(kgrid, v1_E, label = "v1_E")
plot!(kgrid, v5_E, label = "v5_E")
plot!(kgrid, v10_E, label = "v10_E")
plot!(kgrid, v_E, label = "v_E")

#(e)(c)
function policy_E(x)
    e=getindex(findall(kgrid .== x),1)
    return kprime_E[e]
end

function findkpath_E(T, k_0)
    w = findindex(k_0)
    kpath = zeros(T)
    k = kgrid[w]
        for i in 1:T
            k = policy_E(k)
            kpath[i] = k
        end
    return kpath
end

function plotkpath_E(T, k_0)
    w = findindex(k_0)
    kpath = zeros(T)
    k = kgrid[w]
        for i in 1:T
            k = policy_E(k)
            kpath[i] = k
        end
    time = collect(range(1, T, length = T))
    plot(time, kpath, label = "kpath")
end

findkpath_E(100, 0.5)
plotkpath_E(100, 0.5)

#the steady state level of capital is around 4

#1(d)
findkpath_E(100, 4.5)
plotkpath_E(100, 4.5)

#the steady state level of capital is aslo around 4

#the steady state k is same as before but the capital converges faster
#1/σ is the eslasticity of intertemperal substitution, the bigger eslasticity means it is harder
# to subsitute consumption intertemperally, as a result, the consumer saves more resulting more faster capital accumulation


using Combinatorics

function solve(n::Vector{<:AbstractString}, pred::Vector{<:Function})
    rst = Vector{typeof(n)}(0)
    for candidate in permutations(n)
        if all(p(candidate) for p in predicates)
            push!(rst, candidate)
        end
    end
    return rst
end

Names = ["Baker", "Cooper", "Fletcher", "Miller", "Smith"]
predicates = [
    (s) -> last(s) != "Baker",
    (s) -> first(s) != "Cooper",
    (s) -> first(s) != "Fletcher" && last(s) != "Fletcher",
    (s) -> findfirst(s, "Miller") > findfirst(s, "Cooper"),
    (s) -> abs(findfirst(s, "Smith") - findfirst(s, "Fletcher")) != 1,
    (s) -> abs(findfirst(s, "Cooper") - findfirst(s, "Fletcher")) != 1]

solutions = solve(Names, predicates)
foreach(x -> println(join(x, ", ")), solutions)


using OmegaGrav
using Dierckx
using Tables, CSV
# %% Compute the matter power spectrum using CLASS
include("compute_pk_class.jl")
# %% Compute the Compton-y-weighted halo bias, `by`, at seven redshifts
# Reference: Equation (B1) of Chiang, Makiya, Ménard & Komatsu, arXiv:2006.14650
redshift = [0, 0.3, 0.5, 0.7, 1, 1.3, 1.5]
nred = length(redshift)
byt10MF, byt08MF = zeros(nred), zeros(nred)
for ired = 1:nred
   z = redshift[ired]
   # Define functions to return linear baryon+CDM power spectrum
   # Note: The CLASS code takes wavenumbers in units of 1/Mpc (no h) and
   # return power spectra in units of Mpc^3 (no 1/h^3).
   pkcb_class(kovh) = cosmo.pk_cb_lin(kovh * h0, z) * h0^3
   # Spline interpolate in log(k)
   lnk = log(1e-4):0.05:log(100)
   pkcb = Spline1D(lnk, pkcb_class.(exp.(lnk)))
   # %% Compute the Compton-y-weighted halo bias, `by`, for two halo multiplicity functions.
   byt08MF[ired] = by(x -> pkcb(log(x)), z, Ωm, Ωcb, αp = 0.12)
   byt10MF[ired] = by(x -> pkcb(log(x)), z, Ωm, Ωcb, t10MF = true) # not specifying `αp` sets `αp = 0.12`
end

#%% Save table data to by.csv
t = Tables.table([redshift byt08MF byt10MF])
header = ["z", "by(Tinker08)", "by(Tinker10)"]
CSV.write("by.csv", t, header = header)

# %% Clean CLASS (the equivalent of the struct_free() in the `main`
# of CLASS. This step is primordial when running in a loop over different
# cosmologies, as you will saturate your memory very fast if you ommit
# it.
cosmo.struct_cleanup()
# If you want to change completely the cosmology, you should also
# clean the arguments, otherwise, if you are simply running on a loop
# of different values for the same parameters, this step is not needed
cosmo.empty()


"""
   dependenton(command::String, logicset::LogicalCombo; verbose=true)

Checks whether variables on the left of the infix operator (⊥) depend for their
distribution on variables on the right of the operator.

#### Operators

* x ⊥ y: checks if the distribution of potential values of x does not change when y changes.
**Alternatively:** x independentof y

* x !⊥ y: checks if the distribution of potential values of x changes when y changes.
**Alternatively:** x dependenton y

#### Examples - API

```
julia> logicset = logicalparse("a, b, c ∈ 1:4; a > b", verbose=false)

julia> dependenton("a ⊥ c", logicset)
True: a is independent of c
true

julia> dependenton("a ⊥ b", logicset)
False: a is dependent on b
range( a | b == b ) =   {}
––––––––––––––––––– – –––––––
range( a | b == 1 ) = {2,3,4}
range( a | b == 2 ) =  {3,4}
range( a | b == 3 ) =   {4}
false

julia> dependenton("a !⊥ b", logicset)
a is dependent on b
true
```

#### Examples - REPL

```
abstractlogic> a, b, c ∈ 1:4; a > b [clear]
Activeset Already Empty

a, b, c ∈ 1:4            Feasible Outcomes: 64   Perceived Outcomes: 64 ✓        :1 4 2
a > b                    Feasible Outcomes: 24   Perceived Outcomes: 36 ✓        :4 2 1

abstractlogic> a ⊥ c
True: a is independent of c

abstractlogic> a ⊥ b
False: a is dependent on b
range( a | b == b ) =   {}
––––––––––––––––––– – –––––––
range( a | b == 1 ) = {2,3,4}
range( a | b == 2 ) =  {3,4}
range( a | b == 3 ) =   {4}

abstractlogic> a !⊥ b
True: a is dependent on b
```
"""
function dependenton(command::String, logicset::LogicalCombo; verbose=true)

    m = match(r"^(.*?)([⊥!]{1,2}|(dependenton|independentof))(.*?)$",command)
    left, operator, right  = m.captures[[1,2,4]] .|> strip

    ranges = range(logicset)

    rightranges = ranges[Symbol(right)]
    (length(rightranges) == 1) && return println("$left cannot vary in $right if right is constant!")

    leftranges = ranges[Symbol(left)]
    (length(rightranges) == 1) && return println("$left cannot vary in $right if left is constant!")

    leftsets =
      [range(logicalparse("$right = $v", logicset=logicset, verbose=false))[Symbol(left)]
      for v in rightranges ]

    independence = all(length.(leftsets) .== length(leftranges))
    indeptext = (independence ? "independent of" : "dependent on")

    (operator == "⊥")  && independence &&
      return (verbose && println("True: $left is $indeptext $right"); true)
    (operator == "!⊥") && !independence &&
      return (verbose && println("True: $left is $indeptext $right"); true)

   if verbose
       println("False: $left is $indeptext $right")
       txtout = "| " * join(["range( $left \\| $right == $right )", " = ", "{}"], " | ") * " | \n"
       txtout *= "| " * join(fill(":---:", 3), " | ") * " | \n"
       for i in 1:length(rightranges)
         txtout *= "| " * "range( $left \\| $right == $(rightranges[i]) ) " * " | = | {" *
           join(leftsets[i], ",") * "} | \n"
       end
       printmarkdown(txtout)
       println()
   end

   return false
end


using InsarLOS
using Test

@testset "InsarLOS.jl" begin

    dbfile = joinpath(@__DIR__, "S1A_IW_SLC__1SDV_20150316T005120_20150316T005148_005050_006570_8AB8.SAFE.db.1")
    lat = 31.0
    lon = -103.0
    xyz = [0.72732057375709702; 0.60360651252993913; -0.32659449018789627]
    @test isapprox(InsarLOS.calculate_los_xyz(lat, lon, dbfile=dbfile), xyz, atol=1e-3)
    enu = [0.572897473168904; 0.107232543160571; -0.812582098574662]
    @test isapprox(InsarLOS.rotate_xyz_to_enu(xyz, lat, lon), enu, atol=1e-3)

    # Combo function
    @test isapprox(InsarLOS.get_los_enu([lat, lon], dbfile=dbfile), enu, atol=1e-3)

    lat = 31.0
    lon = -102.0
    xyz = [0.78237765815518601; 0.55250630461773964; -0.28744039969575685]
    @test isapprox(InsarLOS.calculate_los_xyz(lat, lon, dbfile=dbfile), xyz, atol=1e-3)
    enu = [0.650408309211683; 0.115737820558811; -0.750715517490279]
    @test isapprox(InsarLOS.rotate_xyz_to_enu(xyz, lat, lon), enu, atol=1e-3)

    # InsarLOS._compute_xyz_ground(31., -102.)
    # xyzground = [-1.1378e6; -5.3529e6; 3.26624e6];

#     enumat = hcat([0.75734; 0.20000; -0.62163], [0.70140; 0.19368; -0.68559])
#     llmat = hcat([31.; -103.], [31.; -102.])
#     @show enumat, llmat
#     @test isapprox(InsarLOS.get_los_enu(llmat, dbfile=dbfile), enumat, atol=1e-3)
end


using LinearAlgebra
using StaticArrays
export factorial2, dist2, ipairs, rpairs, spairs, iindex

@inline factorial2(n::Int64) = prod(n:-2:1) # double factorial !!
@inline dist2(dx,dy,dz) = dx*dx+dy*dy+dz*dz # TODO: use hypot()^2?
@inline dist2(dxyz) = dot(dxyz,dxyz)
@inline dist2(xyz1,xyz2) = dist2(xyz1-xyz2)

ipairs(n::Int64) = ((i, j) for i = 1:n for j = 1:i)
rpairs(n::Int64) = ((i,j) for i in 1:n for j in 1:n) # rectangular option to old pairs
spairs(n::Int64) = ((i, j) for i = 1:n for j = 1:(i-1)) # subdiagonal option to old pairs

triangle(i::Int64) = div(i*(i+1),2)
triangle(i::Int64,j::Int64) = i<j ? triangle(j-1)+i : triangle(i-1)+j

@inline iiterator(n::Int64) = ((i,j,k,l) for (i,j) in ipairs(n) for (k,l) in ipairs(n) if triangle(i,j) <= triangle(k,l))

@inline iindex(i::Int64,j::Int64,k::Int64,l::Int64) = triangle(triangle(i,j),triangle(k,l))

struct Val{x} end
Val(x) = Val{x}()


count(vs::VarStats) = vs.n_obs
sum(vs::VarStats) = vs.sum
min(vs::VarStats) = vs.min
max(vs::VarStats) = vs.max
mean(vs::VarStats) = vs.sum / vs.n_obs
var(vs::VarStats) = vs.var
std(vs::VarStats) = sqrt(vs.var)

function updatestats!(vs::VarStats, value::Real)
    vs.n_obs += 1
    vs.sum += value
    if value < vs.min
        vs.min = value
    end
    if value > vs.max
        vs.max = value
    end
    if vs.n_obs > 2
        n = vs.n_obs
        vs.var = (n - 2) / (n - 1) * vs.n_obs + 1 / n * (value - mean(vs))^2
    end
end

function updatestats!{T<:Real}(varstats::Vector{VarStats}, values::Vector{T})
    for i=1:length(varstats)
        updatestats!(varstats[i], values[i])
    end
end

function aggregate!{N<:Number}(agg::StatsAggregator,
                               key::Any, values::Vector{N})
    if !haskey(agg.groups, key)
        # init new varstats group to the length of first incoming values
        agg.groups[key] = [VarStats() for i=1:length(values)]
    end
    updatestats!(agg.groups[key], values)
end

haskey(agg::StatsAggregator, key::Any) = haskey(agg.groups, key)
keys(agg::StatsAggregator) = keys(agg.groups)
values(agg::StatsAggregator) = values(agg.groups)
getindex(agg::StatsAggregator, k::Any) = getindex(agg.groups, k)


using Distributed
@everywhere using DistributedArrays

using Printf
import PyPlot
using BenchmarkTools

using PyCall
@pyimport pickle

#= @benchmark results for random_points_distributed(10000, 70, R32)
 nprocs()  mean time
---------------------
  1         4.035 s
  5         1.764 s
 15           605 ms
 30           313 ms
 40           272 ms
 60           268 ms
=#

# Calculate distance in meters between two points (ϕ1, λ1) and (ϕ2, λ2)
# on the Earth's surface using the haversine formula. ϕ denotes the latitude
# while λ denotes the longitude.
# See: http://www.movable-type.co.uk/scripts/latlong.html
function haversine_distance(ϕ1, λ1, ϕ2, λ2, R)
    Δϕ = ϕ2 - ϕ1
    Δλ = λ2 - λ1
    a = sinpi(Δϕ / 360.0f0)^2 + cospi(ϕ1 / 180.0f0) * cospi(ϕ2 / 180.0f0) * sinpi(Δλ / 360.0f0)^2
    # c = 2.0f0 * atan(√a, √(1-a))
    c = 2.0f0 * asin(min(1.0f0, √a))
    R*c
end

@everywhere function haversine_distance32(lat1::Float32, lon1::Float32, lat2::Float32, lon2::Float32, radius::Float32)
    c1 = cospi(lat1 / 180.0f0)
    c2 = cospi(lat2 / 180.0f0)
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    d1 = sinpi(dlat / 360.0f0)
    d2 = sinpi(dlon / 360.0f0)
    t = d2 * d2 * c1 * c2
    a = d1 * d1 + t
    c = 2.0f0 * asin(min(1.0f0, sqrt(a)))
    return radius * c
end

function pairwise_distance_hist!(lats, lons, dist)
    N = length(lats)
    for i in 1:N-1, j in (i+1):N
        bin = round(Int8, 10.0f0 * log10(haversine_distance32(lats[i], lons[i], lats[j], lons[j], R32)))
        @inbounds dist[bin] += 1
    end
end

@everywhere function pairwise_distance_histogram_1point(i, lats, lons, bins, R)
    N = length(lats)
    sub_hist = zeros(Int, bins)

    for j in (i+1):N
        bin = round(Int8, 10.0f0 * log10(max(1, haversine_distance32(lats[i], lons[i], lats[j], lons[j], R))))
        @inbounds sub_hist[bin] += 1
    end

    return sub_hist
end

function pairwise_distance_histogram_parallel(lats, lons, bins, R)
    N = length(lats)

    hist = zeros(Int, bins)
    # sub_hists = zeros(Int, bins, N)

    @DArray [pairwise_distance_histogram_1point(i, lats, lons, bins, R) for i in 1:(N-1)];
    # sum(sub_hists, dims=2)
end

function random_points(N, bins)
    lats = 20 .+ 20 .* rand(Float32, N)
    lons = 20 .+ 20 .* rand(Float32, N)

    dist = zeros(bins)

    pairwise_distance_hist!(lats, lons, dist)
end

function random_points_distributed(N, bins, R)
    lats = 20 .+ 20 .* rand(Float32, N)
    lons = 20 .+ 20 .* rand(Float32, N)
    pairwise_distance_histogram_parallel(lats, lons, bins, 6300.0f3)
end

function plot_pairwise_histogram(pickle_fpath, bins)
    R64, R32 = 6371.228e3, 6371.228f3 # average radius of the earth [m]

    ROCK, PAPER, SCISSOR = 1, 2, 3

    ROCK_COLOR = "red"
    PAPER_COLOR = "limegreen"
    SCISSOR_COLOR = "blue"

    f = open(pickle_fpath, "r")
    microbes = pickle.load(f)

    mlons, mlats, species = microbes[:, 1], microbes[:, 2], microbes[:, 3]

    mlons = convert(Array{Float32}, mlons)
    mlats = convert(Array{Float32}, mlats)
    species = convert(Array{Int8}, species)

    rock_mlons = mlons[species .== ROCK];
    rock_mlats = mlats[species .== ROCK];
    paper_mlons = mlons[species .== PAPER];
    paper_mlats = mlats[species .== PAPER];
    scissor_mlons = mlons[species .== SCISSOR];
    scissor_mlats = mlats[species .== SCISSOR];

    N = length(mlons)
    Nr = length(rock_mlons)
    Np = length(paper_mlons)
    Ns = length(scissor_mlons)

    @assert Nr + Np + Ns == N
    @show N, Nr, Np, Ns;

    # colors = Array{String,1}(undef, N); colors .= "";
    # colors[species .== ROCK] .= ROCK_COLOR;
    # colors[species .== PAPER] .= PAPER_COLOR;
    # colors[species .== SCISSOR] .= SCISSOR_COLOR;
    # @assert sum(colors .== "") == 0;

    @printf("Found %d workers.\n", nprocs())

    @printf("Computing rock distance histogram (Nr=%d)...", Nr)
    @time rock_distance_subhists = pairwise_distance_histogram_parallel(rock_mlats, rock_mlons, bins, R32)

    @printf("Computing paper distance histogram (Np=%d)...", Np)
    @time paper_distance_subhists = pairwise_distance_histogram_parallel(paper_mlats, paper_mlons, bins, R32)

    @printf("Computing scissor distance histogram (Ns=%d)...", Ns)
    @time scissor_distance_subhists = pairwise_distance_histogram_parallel(scissor_mlats, scissor_mlons, bins, R32)

    # @printf("Computing pairwise distance histogram (N=%d)...", N)
    # @time all_distance_histp = pairwise_distance_histogram_parallel(mlats, mlons, bins)

    rock_pdh = sum(rock_distance_subhists)
    paper_pdh = sum(paper_distance_subhists)
    scissor_pdh = sum(scissor_distance_subhists)

    max_initial_distance = haversine_distance32(25f0, -145f0, 35f0, -155f0, R32)
    interaction_distance = haversine_distance32(30f0, -150f0, 30f0, -150.01f0, R32)
    PyPlot.axvline(x=max_initial_distance, linestyle=":")
    PyPlot.axvline(x=interaction_distance, linestyle=":")

    length_m = 10 .^ (collect(0:bins-1) ./ 10)

    PyPlot.loglog(length_m, rock_pdh, label="Rocks")
    PyPlot.loglog(length_m, paper_pdh, label="Papers")
    PyPlot.loglog(length_m, scissor_pdh, label="Scissors")

    PyPlot.xlabel("Distance (m)");
    PyPlot.ylabel("Counts");
    PyPlot.xlim([1, 2e7])
    PyPlot.ylim([1, 4e9])

    PyPlot.legend()

    png_fpath = pickle_fpath * ".png"
    @printf("Saving pairwise distance histogram figure: %s\n", png_fpath)
    PyPlot.savefig(png_fpath, dpi=300, format="png", transparent=false)
    PyPlot.close("all")
end

# rps_fpath(p,h) = "/home/alir/nobackup/lagrangian_microbe_output/small_patch_490kp_p0.9_interactions/rps_microbe_species_p" * lpad(string(p), 4, "0") * "_h" * lpad(string(h), 3, "0") * ".pickle"
# rps_fpath(p,h) = "/home/alir/nobackup/lagrangian_microbe_output/small_patch_490kp_p0.55_interactions/rps_microbe_species_p" * lpad(string(p), 4, "0") * "_h" * lpad(string(h), 3, "0") * ".pickle"
rps_fpath(p,h) = "/home/alir/nobackup/lagrangian_microbe_output/small_patch_490kp_pRS0.51_interactions/rps_microbe_species_p" * lpad(string(p), 4, "0") * "_h" * lpad(string(h), 3, "0") * ".pickle"

function plot_multiple_pairwise_histogram()
    for p in 1:6
        plot_pairwise_histogram(rps_fpath(6*p, 0), 70)
    end
end


"""
    roc(scores, answers)

Compute Receiver operating characteristic (ROC) curve. Based on [scikit
implementation](github.com/scikit-learn/scikit-learn/blob/0fb307bf3/sklearn/metrics/_ranking.py)
to allow the comparison of systems created in both platforms.
"""
function roc(scores, answers)

	sorted_indices = sortperm(scores, rev = true)

	scores = scores[sorted_indices]
	answers = answers[sorted_indices]

	# Scores may have repeated points. Those are removed to prevent the curve
	# going straight up or right, as in these values both TPR and FPR grow
	unique_threshold_idxs = [findall(x -> x != 0, diff(scores)); length(scores)]

	# The true positive ratio can be evaluated by dividing the cumulative sum
	# of positive answers at each threshold by the total sum of positive answers.
	tp_sum = cumsum(answers)[unique_threshold_idxs]
	fp_sum = unique_threshold_idxs .- tp_sum
	return (fp_sum/fp_sum[end], tp_sum/tp_sum[end], scores[unique_threshold_idxs])

end

"""
    auc(x, y)

Evaluate the area under the given curve using the trapezoidal rule.
"""
auc(x, y) = sum(diff(x) .* (y[1:end-1] .+ y[2:end]) ./ 2)


# ===== Showcasing the main functionalities of WaveletsExt.jl =====

using Plots,
      Wavelets,
      WaveletsExt

# 1. Additional Redundant Wavelet Transform Methods ----------------------------------------
#    On top of maximal overlap transform implemented in Wavelets.jl, WaveletsExt.jl has
#    additional redundant transform functions through Autocorrelation Wavelet Transforms
#    (Beylkin, Saito), Stationary Wavelet Transforms (Nason, Silverman), and Shift Invariant
#    Wavelet Transforms (Cohen, Raz, Malah).

x = zeros(1<<8)
x[128] = 1
wt = wavelet(WT.db4)                    # Construct Daubechies 4-tap wavelet filter

# ----- Autocorrelation Wavelet Transforms -----
y = acdwt(x, wt)
p1 = wiggle(y) |> p -> plot!(p, yticks=1:9, title="Autocorrelation WT")

# ----- Stationary Wavelet Transforms -----
y = sdwt(x, wt)
p2 = wiggle(y) |> p -> plot!(p, yticks=1:9, title="Stationary WT")

# Combine and save plot
p = plot(p1, p2, layout=(1,2))
savefig(p, "transforms.png")

# 2. Best Basis Algorithms -----------------------------------------------------------------
#    Compared to Wavelets.jl, WaveletsExt.jl has an extended ability of catering toward
#    multiple signals at once via Joint Best Basis (JBB) and Least Statistically Dependent
#    Basis (LSDB). One may then use the `plot_tfbdry` function implemented in WaveletsExt.jl
#    to analyze the best basis subspace.

# Generate 100 noisy heavisine signals of length 2⁸
x = generatesignals(:heavisine, 8) |> x -> duplicatesignals(x, 100, 2, true, 0.5)
# Wavelet packet decomposition of all signals
xw = wpdall(x, wt, 6)

# ----- Joint Best Basis (JBB)
tree = bestbasistree(xw, JBB())
p1 = plot_tfbdry(tree, 6, nd_col=:green, ln_col=:black, bg_col=:white) |>
     p -> plot!(p, title="JBB")

# ----- Least Statistically Dependent Basis (LSDB)
tree = bestbasistree(xw, LSDB())
p2 = plot_tfbdry(tree, 6, nd_col=:green, ln_col=:black, bg_col=:white) |>
     p -> plot!(p, title="LSDB")

# Combine and save plot
p = plot(p1, p2, layout=(1,2))
savefig(p, "paper/bestbasis.png")

# 3. Denoising Algorithms ------------------------------------------------------------------
#    There are two functions available: `denoise` and `denoiseall`. The former denoises one
#    signal whereas the latter denoises multiple signals at once.

# Generate 6 circularly shifted original heavisine signals
x₀ = generatesignals(:heavisine, 8) |> x -> duplicatesignals(x, 6, 2, false)
# Generate 6 noisy versions of the original signals
x = generatesignals(:heavisine, 8) |> x -> duplicatesignals(x, 6, 2, true, 0.8)

# Decompose each noisy signal
xw = wpdall(x, wt)

# Get best basis tree from the decomposition of signals
bt = bestbasistree(xw, JBB())
# Get best basis coefficients based on best basis tree
y = getbasiscoefall(xw, bt)

# Denoise all signals based on computed best basis tree
x̂ = denoiseall(y, :wpt, wt, tree=bt)

# Plot results
xs = repeat([0,256],6) |> x -> reshape(x, (2,6))
ys = repeat(1:6, inner=2) |> x -> reshape(x, (2,6))
p1 = plot(title="Noisy Signals")
wiggle!(x₀, sc=0.7, FaceColor=:white, ZDir=:reverse)
wiggle!(x, sc=0.7, EdgeColor=:red, FaceColor=:white, ZDir=:reverse)
plot!(p1, xs, ys, lc=:black)

p2 = plot(title="Denoised Signals")
wiggle!(x₀, sc=0.7, FaceColor=:white, ZDir=:reverse)
wiggle!(x̂, sc=0.7, EdgeColor=:blue, FaceColor=:white, ZDir=:reverse)
plot!(p2, xs, ys, lc=:black)

# Combine and save plot
p = plot(p1, p2, layout=(1,2))
savefig(p, "paper/denoising.png")

# 4. Feature Extraction Algorithm For Signals ----------------------------------------------
#    The feature extraction algorithm implemented in WaveletsExt.jl is the Local
#    Discriminant Basis (LDB) (Saito, Coifman).

# Generate 100 signals for each class of cylinder-bell-funnel
X, y = generateclassdata(ClassData(:cbf, 100, 100, 100))
# View sample signals and how each class differs from one another
cylinder = wiggle(X[:,1:5], sc=0.3)
plot!(cylinder, yticks=1:5, title="Cylinder signals")
bell = wiggle(X[:,101:105], sc=0.3)
plot!(bell, yticks=1:5, title="Bell signals")
funnel = wiggle(X[:,201:205], sc=0.3)
plot!(funnel, yticks=1:5, title="Funnel signals")
p1 = plot(cylinder, bell, funnel, layout=(3,1))

# Define Local Discriminant Basis object (`n_features` can be tweaked depending on the
# number of desired features to be used as input into classification model)
wt = wavelet(WT.coif4)
ldb = LocalDiscriminantBasis(wt=wt,
                             max_dec_level=6,
                             dm=SymmetricRelativeEntropy(),
                             en=TimeFrequency(),
                             dp=BasisDiscriminantMeasure(),
                             top_k=10,
                             n_features=10)

# Fit and transform signals into features using LDB
X̂ = fit_transform(ldb, X, y)

# Plot the best basis for feature extraction
p2 = plot_tfbdry(ldb.tree, 6, nd_col=:green, ln_col=:black, bg_col=:white)
plot!(p2, title="Basis Selection using LDB")

p = plot(p1, p2, size=(600,300))
savefig(p, "paper/ldb.png")


# Standard rectlin layer: f(A,x,b)=rectlin(A*x+b)
# Note that the bias (a column vector) gets expanded here to match the dimension of A*x

FrectlinAXplusBias(A,X,b)=begin; a=A*X+b*ones(1,size(X,2)); return (max(a,0.),(a.>0)); end

function FrectlinAXplusBias_inplace(value,aux,A,X,b)
    a=A*X+b*ones(1,size(X,2))
    copy!(value,max(a,0.0))
    copy!(aux,(a.>0))
end

function DrectlinAXplusBias(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,X,b)
    if derivativeIDX==1
        axpy!(1.0,(grad_c.*faux_c)*X',grad_n)
    elseif derivativeIDX==2
        axpy!(1.0,A'*(grad_c.*faux_c),grad_n)
    elseif derivativeIDX==3
        axpy!(1.0,sum(grad_c.*faux_c,2),grad_n)
    end
end

if PROC=="GPU"
    function FrectlinAXplusBias(A::CudaArray,X::CudaArray,b::CudaArray)
        return(rectlin(FAXplusBias(A,X,b)[1]),FAXplusBias(A,X,b)[1])
        # TODO: FIX THIS UNNECESSARY DOUBLE COMPUTATION
    end

    function FrectlinAXplusBias_inplace(value,aux,A::CudaArray,X::CudaArray,b::CudaArray)
        FAXplusBias_inplace(value,aux,A,X,b)
        aux=CudaArray(Float64,size(value))
        copy!(aux,value)
        rectlin!(value,value)
    end

    function DrectlinAXplusBias(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,X::CudaArray,b::CudaArray)
        tmp=CudaArray(Float64,size(grad_c)); fill!(tmp,0.0)
        A_emult_Bg0!(grad_c,faux_c,tmp)
        if derivativeIDX==1
            gemm!('N','T',1.0,tmp,X,1.0,grad_n)
        elseif derivativeIDX==2
            gemm!('T','N',1.0,A,tmp,1.0,grad_n)
        elseif derivativeIDX==3
            ons=CudaArray(Float64,(size(X,2),1)); fill!(ons,1.0)
            gemm!('N','N',1.0,tmp,ons,1.0,grad_n)
            free(ons)
        end
        free(tmp)
    end

end

Derivative[FrectlinAXplusBias]=DrectlinAXplusBias
Inplace[FrectlinAXplusBias]=FrectlinAXplusBias_inplace

rectlinAXplusBias(A,X,b)=ADnode(FrectlinAXplusBias,[A X b])
export rectlinAXplusBias


using InfiniteOpt, Ipopt, Plots
using KNITRO

exponential_relaxation(z, δ) = exp(one(typeof(z)) - z / δ) - one(typeof(z)) - log(δ)
function relaxed_log_barrier(z; δ = 0.3f0)
    return max(z > δ ? -log(z) : exponential_relaxation(z, δ), zero(typeof(z)))
end
# Set the SEIR parameters
γ = 0.303
β = 0.727
N = 100
ξ = 0.2

# Set the domain information
i_max = 0.06
t0 = 0
tf = 50

# Set the intial condition values
e0 = 1 / N
i0 = 0
r0 = 0
s0 = 1 - 1 / N

model = InfiniteModel(KNITRO.Optimizer)

@infinite_parameter(model, t ∈ [t0, tf], num_supports = 201,
    derivative_method = OrthogonalCollocation(2))

@variable(model, s ≥ 0, Infinite(t))
@variable(model, e ≥ 0, Infinite(t))
@variable(model, i ≥ 0, Infinite(t))
@variable(model, r ≥ 0, Infinite(t))
@variable(model, 0.2 ≤ u ≤ 0.8, Infinite(t), start = 0.2)

@objective(model, Min, ∫(u, t))

# Define the initial conditions
@constraint(model, s(0) == s0)
@constraint(model, e(0) == e0)
@constraint(model, i(0) == i0)
@constraint(model, r(0) == r0)

# Define the SEIR equations
@constraint(model, s_constr, ∂(s, t) == -(1 - u) * β * s * i)
@constraint(model, e_constr, ∂(e, t) == (1 - u) * β * s * i - ξ * e)
@constraint(model, i_constr, ∂(i, t) == ξ * e - γ * i)
@constraint(model, r_constr, ∂(r, t) == γ * i)

# Define the infection rate limit
@constraint(model, imax_constr, i ≤ i_max)

print(model)

optimize!(model)

r_opt = value(r, ndarray = true) * 100 # make the population fractions into percentages
s_opt = value(s, ndarray = true) * 100
i_opt = value(i, ndarray = true) * 100
e_opt = value(e, ndarray = true) * 100
u_opt = value(u)
@show obj_opt = objective_value(model)
ts = value(t)

p = plot(ts, i_opt, label = "i(t)")
plot!(ts, e_opt, label = "e(t)")
ylabel!("Pop. (%)")
xlabel!("Time (Days)")
display(p)

p1 = plot(ts, u_opt, label = "u(t)")
xlabel!("Time (Days)")
ylabel!("Distancing Ratio")
display(p1)


#=
Interpolation Methods

=#

using OffsetArrays: OffsetVector

"""
private function, find index in array where x would be inserted for interpolation
between xvec[i] and xvec[i+1]
"""
function findindex(xvec, x)

    n = length(xvec)
    i = searchsortedlast(xvec, x)

    # this version allows extrapolation
    if i == 0
        i = 1
    elseif i == n
        i = n - 1
    end

    # this version prevents extrapolation
    # if i == 0
    #     throw(DomainError(x, "x falls below range of provided spline data"))
    # elseif i == n
    #     if x == xvec[n]
    #         i = n - 1
    #     else
    #         throw(DomainError(x, "x falls above range of provided spline data"))
    #     end
    # end

    return i
end

## ------ Akima Interpolation  ---------

# struct used internally
struct Akima{TX, TY, TCoeff}

    xdata::TX
    ydata::TY
    p0::Vector{TCoeff}
    p1::Vector{TCoeff}
    p2::Vector{TCoeff}
    p3::Vector{TCoeff}

end

"""
    Akima(xdata, ydata, delta_x=0.0)

Creates an akima spline at node points: `xdata`, `ydata`.  This is a 1D spline that avoids
overshooting issues common with many other polynomial splines resulting in a more
natural curve.  It also only depends on local points (`i-2`...`i+2`) allow for more efficient
computation.  `delta_x` is the half width of a smoothing interval used for the absolute
value function.  Set `delta_x=0` to recover the original akima spline.  The smoothing
is only useful if you want to differentiate xdata and ydata.  In many case the nodal points
are fixed so this is not needed.  Returns an akima spline object (Akima struct).
This function, only performs construction of the spline, not evaluation.
This is useful if you want to evaluate the same mesh at multiple different conditions.
A convenience method exists below to perform both in one shot.
"""
function Akima(xdata::AbstractVector{TFX}, ydata::AbstractVector{TFY}, delta_x=0.0) where {TFX, TFY}

    # setup
    eps = 1e-30
    n = length(xdata)
    TCoeff = promote_type(TFX, TFY)

    # compute segment slopes
    m = OffsetVector(zeros(TCoeff, n+3), -1:n+1)
    for i = 1:n-1
        m[i] = (ydata[i+1] - ydata[i]) / (xdata[i+1] - xdata[i])
    end

    # estimation for end points
    m[0] = 2.0*m[1] - m[2]
    m[-1] = 2.0*m[0] - m[1]
    m[n] = 2.0*m[n-1] - m[n-2]
    m[n+1] = 2.0*m[n] - m[n-1]

    # slope at points
    t = zeros(TCoeff, n)
    for i = 1:n
        m1 = m[i-2]
        m2 = m[i-1]
        m3 = m[i]
        m4 = m[i+1]
        w1 = abs_smooth(m4 - m3, delta_x)
        w2 = abs_smooth(m2 - m1, delta_x)
        if (w1 < eps && w2 < eps)
            t[i] = 0.5*(m2 + m3)  # special case to avoid divide by zero
        else
            t[i] = (w1*m2 + w2*m3) / (w1 + w2)
        end
    end

    # polynomial cofficients
    p0 = zeros(TCoeff, n-1)
    p1 = zeros(TCoeff, n-1)
    p2 = zeros(TCoeff, n-1)
    p3 = zeros(TCoeff, n-1)
    for i = 1:n-1
        dx = xdata[i+1] - xdata[i]
        t1 = t[i]
        t2 = t[i+1]
        p0[i] = ydata[i]
        p1[i] = t1
        p2[i] = (3.0*m[i] - 2.0*t1 - t2)/dx
        p3[i] = (t1 + t2 - 2.0*m[i])/dx^2
    end

    return Akima(xdata, ydata, p0, p1, p2, p3)
end

function (spline::Akima)(x::Number)

    j = findindex(spline.xdata, x)

    # evaluate polynomial
    dx = x - spline.xdata[j]
    y = spline.p0[j] + spline.p1[j]*dx + spline.p2[j]*dx^2 + spline.p3[j]*dx^3

    return y
end

(spline::Akima)(x::AbstractVector) = spline.(x)

"""
    akima(x, y, xpt)

A convenience method to perform construction and evaluation of the spline in one step.
See docstring for Akima for more details.

**Arguments**
- `x, y::Vector{Float}`: the node points
- `xpt::Vector{Float} or ::Float`: the evaluation point(s)

**Returns**
- `ypt::Vector{Float} or ::Float`: interpolated value(s) at xpt using akima spline.
"""
akima(x, y, xpt) = Akima(x, y)(xpt)

"""
    derivative(spline, x)

Computes the derivative of an Akima spline at x.

**Arguments**
- `spline::Akima}`: an Akima spline
- `x::Float`: the evaluation point(s)

**Returns**
- `dydx::Float`: derivative at x using akima spline.
"""
function derivative(spline::Akima, x)

    j = findindex(spline.xdata, x)

    # evaluate polynomial
    dx = x - spline.xdata[j]
    dydx = spline.p1[j] + 2*spline.p2[j]*dx + 3*spline.p3[j]*dx^2

    return dydx
end

"""
    gradient(spline, x)

Computes the gradient of a Akima spline at x.

**Arguments**
- `spline::Akima}`: an Akima spline
- `x::Vector{Float}`: the evaluation point(s)

**Returns**
- `dydx::Vector{Float}`: gradient at x using akima spline.
"""
gradient(spline::Akima, x) = derivative.(Ref(spline), x)

# ------ Linear Interpolation  ---------

"""
    linear(xdata, ydata, x::Number)

Linear interpolation.

**Arguments**
- `xdata::Vector{Float64}`: x data used in constructing interpolation
- `ydata::Vector{Float64}`: y data used in constructing interpolation
- `x::Float64`: point to evaluate spline at

**Returns**
- `y::Float64`: value at x using linear interpolation
"""
function linear(xdata, ydata, x::Number)

    i = findindex(xdata, x)

    # lienar interpolation
    eta = (x - xdata[i]) / (xdata[i+1] - xdata[i])
    y = ydata[i] + eta*(ydata[i+1] - ydata[i])

    return y
end

"""
    linear(xdata, ydata, x::AbstractVector)

Convenience function to perform linear interpolation at multiple points.
"""
linear(xdata, ydata, x::AbstractVector) = linear.(Ref(xdata), Ref(ydata), x)

"""
derivative of linear interpolation at `x::Number`
"""
function derivative(xdata, ydata, x)

    i = findindex(xdata, x)
    dydx = (ydata[i+1] - ydata[i]) / (xdata[i+1] - xdata[i])

    return dydx
end

"""
gradient of linear interpolation at `x::Vector`
"""
gradient(xdata, ydata, x) = derivative.(Ref(xdata), Ref(ydata), x)

# ------- higher order recursive 1D interpolation ------

"""
   interp2d(interp1d, xdata, ydata, fdata, xpt, ypt)

2D interpolation using recursive 1D interpolation.  This approach is likely less efficient than a more direct
2D interpolation method, especially one you can create separate creation from evaluation,
but it is generalizable to any spline approach and any dimension.

**Arguments**
- `interp1d`: any spline function of form: ypt = interp1d(xdata, ydata, xpt) where data are the known
    data(node) points and pt are the points where you want to evaluate the spline at.
- `xdata::Vector{Float}`, `ydata::Vector{Float}`: Define the 2D grid
- `fdata::Matrix{Float}`: where fdata[i, j] is the function value at xdata[i], ydata[j]
- `xpt::Vector{Float}`, `ypt::Vector{Float}`: the locations where you want to evaluate the spline

**Returns**
- `fhat::Matrix{Float}`: where fhat[i, j] is the estimate function value at xpt[i], ypt[j]
"""
function interp2d(interp1d, xdata, ydata, fdata, xpt, ypt)

    ny = length(ydata)
    nxpt = length(xpt)
    nypt = length(ypt)

    R = promote_type(eltype(xpt), eltype(ypt))

    yinterp = Array{R}(undef, ny, nxpt)
    output = Array{R}(undef, nxpt, nypt)

    for i = 1:ny
        yinterp[i, :] .= interp1d(xdata, fdata[:, i], xpt)
    end
    for i = 1:nxpt
        output[i, :] .= interp1d(ydata, yinterp[:, i], ypt)
    end

    return output
end

"""
    interp3d(interp1d, xdata, ydata, zdata, fdata, xpt, ypt, zpt)

Same as interp2d, except in three dimension.
"""
function interp3d(interp1d, xdata, ydata, zdata, fdata, xpt, ypt, zpt)

    nz = length(zdata)
    nxpt = length(xpt)
    nypt = length(ypt)
    nzpt = length(zpt)

    R = promote_type(eltype(xpt), eltype(ypt), eltype(zpt))
    zinterp = Array{R}(undef, nz, nxpt, nypt)
    output = Array{R}(undef, nxpt, nypt, nzpt)

    for i = 1:nz
        zinterp[i, :, :] .= interp2d(interp1d, xdata, ydata, fdata[:, :, i], xpt, ypt)
    end
    for j = 1:nypt
        for i = 1:nxpt
            output[i, j, :] .= interp1d(zdata, zinterp[:, i, j], zpt)
        end
    end

    return output
end

"""
    interp4d(interp1d, xdata, ydata, zdata, fdata, xpt, ypt, zpt)

Same as interp3d, except in four dimensions.
"""
function interp4d(interp1d, xdata, ydata, zdata, tdata, fdata, xpt, ypt, zpt, tpt)

    nt = length(tdata)
    nxpt = length(xpt)
    nypt = length(ypt)
    nzpt = length(zpt)
    ntpt = length(tpt)

    R = promote_type(eltype(xpt), eltype(ypt), eltype(zpt), eltype(tpt))
    tinterp = Array{R}(undef, nt, nxpt, nypt, nzpt)
    output = Array{R}(undef, nxpt, nypt, nzpt, ntpt)

    for i = 1:nt
        tinterp[i, :, :, :] .= interp3d(interp1d, xdata, ydata, zdata, fdata[:, :, :, i], xpt, ypt, zpt)
    end
    for k = 1:nzpt
        for j = 1:nypt
            for i = 1:nxpt
                output[i, j, k, :] .= interp1d(tdata, tinterp[:, i, j, k], tpt)
            end
        end
    end

    return output
end


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

function randlines(n=300, t=0.4)
	#n = 100 #1000 #1000 #20000
	#t = 0.4 #0.15 #0.4 #0.15
	V = zeros(Float64,2,2*n)
	EV = [zeros(Int64,2) for k=1:n]
	for k=1:n
		v1 = rand(Float64,2)
		v2 = rand(Float64,2)
		vm = (v1+v2)/2
		transl = rand(Float64,2)
		V[:,k] = (v1-vm)*t + transl
		V[:,n+k] = (v2-vm)*t + transl
		EV[k] = [k,n+k]
	end

	V = GL.normalize2(V)
	model = (V,EV)
	Sigma = Lar.spaceindex(model)

	model = V,EV;
	W,EW = Lar.fragmentlines(model);
	U,EVs = Lar.biconnectedComponent((W,EW::Lar.Cells));
	EV = convert(Lar.Cells, cat(EVs))
	V,FVs,EVs = Lar.arrange2D(U,EV)
end

# ////////////////////////////////////////////////////////////
# generation of 2D arrangement
V,FVs,EVs = randlines()

# native OpenGL visualization
GL.VIEW(GL.GLExplode(V,FVs,1.2,1.2,1.2,1));
GL.VIEW(GL.GLExplode(V,FVs,1.2,1.2,1.2,3,1));
GL.VIEW(GL.GLExplode(V,FVs,1.2,1.2,1.2,99,1));
GL.VIEW(GL.GLExplode(V,FVs,1.,1.,1.,99,1));
GL.VIEW(GL.GLExplode(V,EVs,1.2,1.2,1.2,1,1));


using DataDeps

register(DataDep(
    "spiral",
    """
    Dataset: spiral
    Credit: https://smorbieu.gitlab.io/generate-datasets-to-understand-some-clustering-algorithms-behavior/

    A simulated dataset of three noisy spirals.

    Observations: 1000
    Features:     2
    Classes:      3
    """,
    "script", # nothing to download
    "51d3cf8d223edcacb65db752ff3b0487722951b6e00cb566c3e41a574189c20e";
    fetch_method=function(unused, localdir)
        # Parameters
        N           = 1000
        max_A       = 600
        max_B       = 300
        max_C       = N - max_A - max_B
        max_radius  = 7.0
        x0          = -3.5
        y0          = 3.5
        angle_start = π / 8
        seed        = 1903

        # Simulate the data.
        rng = StableRNG(seed)
        target, x, y = Vector{Char}(undef, N), zeros(N), zeros(N)
        for i in 1:N
        if i ≤ max_A
            (class, k, n, θ) = ('A', i, max_A, angle_start)
            noise = 0.1
        elseif i ≤ max_A + max_B
            (class, k, n, θ) = ('B', i-max_A+1, max_B, angle_start + 2π/3)
            noise = 0.2
        else
            (class, k, n, θ) = ('C', i-max_A-max_B+1, max_C, angle_start + 4π/3)
            noise = 0.3
        end
        # Compute coordinates.
        angle = θ + π * k / n
        radius = max_radius * (1 - k / (n + n / 5))

        target[i] = class
        x[i] = x0 + radius*cos(angle) + noise*randn(rng)
        y[i] = y0 + radius*sin(angle) + noise*randn(rng)
        end
        local_file = joinpath(localdir, "data.csv")
        df = DataFrame(target=target, x1=x, x2=y)
        perm = Random.randperm(rng, size(df, 1))
        foreach(col -> permute!(col, perm), eachcol(df))
        CSV.write(local_file, df)
        return local_file
    end
))


#=
  From BLOG Swift example birthday.blog
  https://github.com/lileicc/swift/blob/master/example/birthday.blog
  Result
  """
  Distribution of values for exists Person x (exists Person y ((!(x = y) & (Birthday(x) = Birthday(y)))))
	true	0.5446927374301685
	false	0.4553072625698328
  """

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

@model function birthday(num_people=missing)
    # var day = _.range(365);
    num_days = 365

    if num_people === missing
        num_people ~ truncated(Poisson(20),1,num_days)
    end

    # Note: I call it zbirthdays so it's displayed last
    zbirthdays = Vector{Int}(undef, num_people)
    for d in 1:num_people
       zbirthdays[d] ~ DiscreteUniform(1,num_days)
    end

    #  is there a duplicate of birthdays?
    # hakank: I'm not sure if the is the best (or even good) way
    #         to model this...
    b = make_hash(zbirthdays) # Collect the birthdays
    s = sum(values(b) .> 1)
    if s == 0
        t ~ Bernoulli(0)
    else
        t ~ Bernoulli(1)
    end

end

model = birthday(23)
# model = birthday()

num_chns = 4
println("Prior distribution")
chns_prior = sample(model, Prior(), MCMCThreads(), 1000, num_chns)
display(chns_prior)
# display(plot(chns_prior))

println("\nPosterior distribution")
num_chns = 4

# chns = sample(model, MH(), MCMCThreads(), 10_000, num_chns)
# chns = sample(model, MH(), MCMCThreads(), 30_000, num_chns)

chns = sample(model, IS(), MCMCThreads(), 10_000, num_chns) # ORIG

# chns = sample(model, PG(20), MCMCThreads(), 1000, num_chns)

# chns = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chns)

## LoadError: MethodError: Cannot `convert` an object of type Float64 to an object of type Symbol
# chns = sample(model, Gibbs(NUTS(0.65,5,1,1,0,:lambda_1,:lambda_2), PG(20,:tau) ), 1000)
# chns = sample(model, Gibbs(HMC(0.1,5,:lambda_1,:lambda_2), PG(20,:tau) ), MCMCThreads(), 1000, num_chns)

# Error
# chns = sample(model, Gibbs(HMCDA(0.15,0.65,0.3,:lambda_1,:lambda_2), PG(20,:tau) ), MCMCThreads(), 1000, num_chns) # Error...

# Skipping this since it contains a lot of uninteresting zbirthdays[b]
# display(chns)
println("t:")
display(chns[:t])

# plot(chns)

if :num_people in chns.name_map.parameters
    println("num_people:")
    display(chns[:num_people])
end

# df = DataFrame(chns)
# display(df)


using LinearAlgebra, DiffEqOperators, Random, Test, BandedMatrices, SparseArrays

# 4th derivative, 4th order
function fourth_deriv_approx_stencil(N)
    A = zeros(N,N+2)
    A[1,1:8] = [3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3]
    A[2,1:8] = [2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6]

    A[N-1,N-5:end] = reverse([2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6], dims=2)
    A[N,N-5:end] = reverse([3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3], dims=2)

    for i in 3:N-2
        A[i,i-2:i+4] = [-1/6 2.0 -13/2 28/3 -13/2 2.0 -1/6]
    end
    return A
end

function second_deriv_fourth_approx_stencil(N)
    A = zeros(N,N+2)
    A[1,1:6] = [5/6 -15/12 -1/3 7/6 -6/12 5/60]
    A[N,N-3:end] = [1/12 -6/12 14/12 -4/12 -15/12 10/12]
    for i in 2:N-1
        A[i,i-1:i+3] = [-1/12 4/3 -5/2 4/3 -1/12]
    end
    return A
end

function second_derivative_stencil(N)
  A = zeros(N,N+2)
  for i in 1:N, j in 1:N+2
      (j-i==0 || j-i==2) && (A[i,j]=1)
      j-i==1 && (A[i,j]=-2)
  end
  A
end

@testset "Test Constructor, Multiplication, and Concretization" begin
    # Generate random parameters
    al = rand()
    bl = rand()
    cl = rand()

    ar = rand()
    br = rand()
    cr = rand()
    dx = rand()

    Q = RobinBC((al, bl, cl), (ar, br, cr), dx)
    N = 20
    L = CenteredDifference(4,4, 1.0, N)
    L2 = CenteredDifference(2,4, 1.0, N)

    function coeff_func(du,u,p,t)
        du .= u
    end

    cL = coeff_func*L
    coeffs = rand(N)
    DiffEqOperators.update_coefficients!(cL,coeffs,nothing,0.0)
    cA = cL*Q
    DiffEqOperators.update_coefficients!(cA,coeffs,nothing,0.0)
    @test cL.coefficients == coeffs == cA.L.coefficients

    A = L*Q
    c2A = coeff_func*A
    DiffEqOperators.update_coefficients!(c2A,coeffs,nothing,0.0)
    @test coeffs == c2A.L.coefficients

    # Test GhostDerivativeOperator constructor by *
    u = rand(N)
    A = L*Q
    # Test for consistency of GhostDerivativeOperator*u with L*(Q*u)
    @test A*u ≈ L*(Q*u)

    # Test for consistency of GhostDerivativeOperator*M with L*(Q*M)
    M = rand(N,10)
    LQM = zeros(N,10)
    for i in 1:10
        mul!(view(LQM,:,i), L, Q*M[:,i])
    end
    ghost_LQM = A*M
    @test ghost_LQM ≈ LQM

    u = rand(22)
    @test (L + L2) * u ≈ convert(AbstractMatrix,L + L2) * u ≈ (BandedMatrix(L) + BandedMatrix(L2)) * u

    # Test concretization
    @test Array(A)[1] ≈ (Array(L)*Array(Q,N)[1], Array(L)*Array(Q,N)[2])[1]
    @test Array(A)[2] ≈ (Array(L)*Array(Q,N)[1], Array(L)*Array(Q,N)[2])[2]
    @test SparseMatrixCSC(A)[1] ≈ (SparseMatrixCSC(L)*SparseMatrixCSC(Q,N)[1], SparseMatrixCSC(L)*SparseMatrixCSC(Q,N)[2])[1]
    @test SparseMatrixCSC(A)[2] ≈ (SparseMatrixCSC(L)*SparseMatrixCSC(Q,N)[1], SparseMatrixCSC(L)*SparseMatrixCSC(Q,N)[2])[2]
    @test sparse(A)[1] ≈ (sparse(L)*sparse(Q,N)[1], sparse(L)*sparse(Q,N)[2])[1]
    @test sparse(A)[2] ≈ (sparse(L)*sparse(Q,N)[1], sparse(L)*sparse(Q,N)[2])[2]
    # BandedMatrix not implemeted for boundary operator
    @test_broken BandedMatrix(A)[1] ≈ (BandedMatrix(L)*BandedMatrix(Q,N)[1], BandedMatrix(L)*BandedMatrix(Q,N)[2])[1]
    @test_broken BandedMatrix(A)[2] ≈ (BandedMatrix(L)*BandedMatrix(Q,N)[1], BandedMatrix(L)*BandedMatrix(Q,N)[2])[2]

    # Test that concretization works with multiplication
    u = rand(20)
    @test Array(A)[1]*u + Array(A)[2] ≈ L*(Q*u) ≈ A*u
    @test sparse(A)[1]*u + sparse(A)[2] ≈ L*(Q*u) ≈ A*u
end

@testset "Test Left Division L2 (second order)" begin
    # Test \ Homogenous BC
    # f(x) = -x^2 + x
    # f''(x) = -2.0
    # f'(0) = f'(1) = 0
    f(x) = -x^2 + x
    f2(x) = -2.0

    dx = 0.01
    x = dx:dx:1.0-dx
    N = length(x)

    L = CenteredDifference(2, 2, dx, N)
    Q = RobinBC((1.0, 0.0, 0.0), (1.0, 0.0, 0.0), dx)
    A = L*Q

    analytic_L = second_derivative_stencil(N) ./ dx^2
    analytic_QL = [transpose(zeros(N)); Diagonal(ones(N)); transpose(zeros(N))]
    analytic_AL = analytic_L*analytic_QL

    # No affine component to the this system
    analytic_f = analytic_AL \ f2.(x)
    ghost_f = A \ f2.(x)

    # Check that A\f2.(x) is consistent with analytic_AL \ f2.(x)
    @test analytic_f ≈ ghost_f

    # Additionally test that A\f2.(x) ≈ f.(x)
    @test f.(x) ≈ ghost_f ≈ analytic_f

    # Check ldiv!
    f_temp = zeros(N)
    ldiv!(f_temp, A, f2.(x))
    @test f_temp ≈ ghost_f ≈ analytic_f

    # Check that left division with matrices works
    ghost_fM = A \ [f2.(x) f2.(x)]
    analytic_fM = analytic_AL \ [f2.(x) f2.(x)]
    @test ghost_fM ≈ analytic_fM

    fM_temp = zeros(N,2)
    ldiv!(fM_temp, A, [f2.(x) f2.(x)])
    @test fM_temp ≈ analytic_fM

    # Test \ Inhomogenous BC
    # f(x) = -x^2 + x + 4.0
    # f''(x) = -2.0
    # f'(0) = f'(1) = 0
    f(x) = -x^2 + x + 4.0
    f2(x) = -2.0

    dx = 0.01
    x = dx:dx:1.0-dx
    N = length(x)

    L = CenteredDifference(2, 2, dx, N)
    Q = RobinBC((1.0, 0.0, 4.0), (1.0, 0.0, 4.0), dx)
    A = L*Q

    analytic_L = second_derivative_stencil(N) ./ dx^2
    analytic_QL = [transpose(zeros(N)); Diagonal(ones(N)); transpose(zeros(N))]
    analytic_Qb = [4.0; zeros(N); 4.0]
    analytic_AL = analytic_L*analytic_QL
    analytic_Ab = analytic_L*analytic_Qb

    analytic_f = analytic_AL \ (f2.(x) - analytic_Ab)
    ghost_f = A \ f2.(x)

    # Check that A\f2.(x) is consistent with analytic_AL \ f2.(x)
    @test analytic_f ≈ ghost_f

    # Additionally test that A\f2.(x) ≈ f.(x)
    @test f.(x) ≈ ghost_f ≈ analytic_f

    # Check ldiv!
    f_temp = zeros(N)
    ldiv!(f_temp, A, f2.(x))
    @test f_temp ≈ ghost_f ≈ analytic_f

    # Check \ for Matrix
    M2 = [f2.(x) 2.0*f2.(x) 10.0*f2.(x)]
    analytic_M = analytic_AL \ (M2 .- analytic_Ab)
    ghost_M = A \ M2
    @test analytic_M ≈ ghost_M

    # Check ldiv! for Matrix
    M_temp = zeros(N,3)
    ldiv!(M_temp, A, M2)
    @test M_temp ≈ analytic_M ≈ ghost_M

    # Additionally test that A\M2 ≈ [f, 2.0(f-4.0)+4.0, 10.0(f-4.0)+4.0]
    M = [f.(x) 2.0*(f.(x) .- 4.0).+4.0 10.0*(f.(x) .- 4.0).+4.0]
    @test M ≈ M_temp ≈ analytic_M ≈ ghost_M
end

@testset "Test Left Division L4 (fourth order)" begin

    # Test \ homogenous and inhomogenous BC
    dx = 0.01
    x = 0.01:dx:0.2
    N = length(x)
    u = sin.(x)

    L = CenteredDifference(4, 4, dx, N)
    Q = RobinBC((1.0, 0.0, sin(0.0)), (1.0, 0.0, sin(0.2+dx)), dx)
    A = L*Q

    analytic_L = fourth_deriv_approx_stencil(N) ./ dx^4
    analytic_QL = [transpose(zeros(N)); Diagonal(ones(N)); transpose(zeros(N))]
    analytic_AL = analytic_L*analytic_QL
    analytic_Qb = [zeros(N+1); sin(0.2+dx)]
    analytic_Ab = analytic_L*analytic_Qb

    analytic_u = analytic_AL \ (u - analytic_Ab)
    ghost_u = A \ u

    # Check that A\u.(x) is consistent with analytic_AL \ u.(x)
    @test analytic_u ≈ ghost_u

    # Check ldiv!
    u_temp = zeros(N)
    ldiv!(u_temp, A, u)
    @test u_temp ≈ ghost_u ≈ analytic_u

    # Check \ for Matrix
    M2 = [u 2.0*u 10.0*u]
    analytic_M = analytic_AL \ (M2 .- analytic_Ab)
    ghost_M = A \ M2
    @test analytic_M ≈ ghost_M

    # Check ldiv! for Matrix
    M_temp = zeros(N,3)
    ldiv!(M_temp, A, M2)
    @test M_temp ≈ ghost_M ≈ analytic_M
end


#=
# Steady flow through a duct with heat transfer
This notebook demonstrates the use of tools for computing steady compressible flow
through a constant-area duct with heat transfer, commonly known as **Rayleigh flow**.
=#

# ### Set up the module
using Gasdynamics1D
#-
using Plots
using LaTeXStrings

#=
#### Simple example
In a duct with helium, if the flow enters with stagnation temperature of 20 degrees C,
how much does the stagnation temperature change if we add heat 400 kJ/kg?
=#
q = HeatFlux(400u"kJ/kg")
T01 = StagnationTemperature(20u"°C")

# starting stagnation enthalpy
h01 = StagnationEnthalpy(T01,gas=He) # h01 = cp*T01

# add the heat
h02 = StagnationEnthalpy(h01 + q)

# calculate the final stagnation temperature
T02 = StagnationTemperature(h02,gas=He)  # T02 = h02/cp

# Report the final value in Celsius:
value(T02,u"°C")
# so the flow exiting the duct has stagnation temperature 97 C.

#=
The sonic state is used as a reference, similar to Fanno flow and isentropic flow.
All states share the same $T_{0}^*$, $p^*$, $u^*$, etc, so we can relate two points
via this reference. We have functions that allow us to do this. For example, to
find the ratio of the local stagnation temperature to its sonic reference value
at Mach number 0.5 in air:
=#
T0OverT0Star(MachNumber(0.5),RayleighFlow,gas=Air)

#=
Note that the argument `RayleighFlow` was used to designate this function's use.
If the Mach number is 1, then...
=#
T0OverT0Star(MachNumber(1),RayleighFlow,gas=Air)

#=
That is, $T_{0}$ is equal to $T_{0}^*$ when Mach number is 1, as it should be,
by definition.
=#
#=
Alternatively, if we know $T_0/T_0^*$, we can determine $M$. However, there are
sometimes two possible values of Mach number, as with Fanno flow and isentropic flow.
=#
Mrange = range(0.001,8,length=801)
T0_over_T0star = []
for M in Mrange
    push!(T0_over_T0star,value(T0OverT0Star(MachNumber(M),RayleighFlow,gas=Air)))
end
#-
plot(Mrange,T0_over_T0star,xlim=(0,8),ylim=(0,2),xticks=0:1:8,xlabel="Mach number",ylabel=L"T_0/T_0^*",legend=false)
plot!(Mrange,0.75*ones(length(Mrange)),style=:dash)

#=
For example, if $T_0/T_0^* = 0.75$, there are two possible Mach numbers, one subsonic
and one supersonic. The correct one depends on the circumstances.
=#
Msub, Msup = MachNumber(TemperatureRatio(0.75),RayleighFlow,gas=Air)
#-
Msub
#-
Msup

#=
### Example
Suppose the flow of air in a duct enters with velocity 75 m/s at a pressure of
150 kPa and temperature 300 K. We add heat 900 kJ/kg. Find

(a) The value of the Mach number at the entrance and exit

(b) The velocity, pressure, and temperature at the exit

(c) The maximum amount of heat we can add so that the flow is just choked
=#

# Set the given values
u1 = Velocity(75u"m/s")
p1 = Pressure(150u"kPa")
T1 = Temperature(300)
q = HeatFlux(900u"kJ/kg")

# We can immediately find the Mach number at the location 1, by finding $c_1$:
c1 = SoundSpeed(T1,gas=Air)
M1 = MachNumber(u1/c1)

# Now, use $M_1$ and $T_1$ to find $T_{01}$.
T01 = StagnationTemperature(T1,M1,Isentropic,gas=Air)

#=
Now let's focus on state 2. We will use the basic relation $h_{02} = h_{01} + q$
to determine $T_{02}$.
=#
h01 = StagnationEnthalpy(T01,gas=Air)
h02 = StagnationEnthalpy(h01+q)
T02 = StagnationTemperature(h02,gas=Air)

#=
To get the other states at location 2, we need the sonic reference conditions.
For example, we will find

$$T_{02}/T_{0}^*$$
to determine $M_2$. We get this by calculating

$$T_{0}^* = T_{01}/(T_{01}/T_{0}^*)$$
=#
T0star = StagnationTemperature(T01/T0OverT0Star(M1,RayleighFlow,gas=Air))
#-
T02_over_T0star = TemperatureRatio(T02/T0star)
#-
M2sub, M2sup = MachNumber(T02_over_T0star,RayleighFlow,gas=Air)
#=
The subsonic one is the only possible candidate, because the flow came in at
subsonic speed and can't have passed through Mach number 1. Thus
=#
M2 = M2sub
# We have increased the Mach number from 0.216 to 0.573 by adding heat.

# Now find the other reference states that we will use:
pstar = Pressure(p1/POverPStar(M1,RayleighFlow,gas=Air))
ustar = Velocity(u1/VOverVStar(M1,RayleighFlow,gas=Air))
Tstar = Temperature(T1/TOverTStar(M1,RayleighFlow,gas=Air))

#=
and now the values themselves. For example, we use $M_2$ to calculate $p_2/p^*$
and then calculate $p_2 = p^*(p_2/p^*)$.
=#
p2 = Pressure(pstar*POverPStar(M2,RayleighFlow,gas=Air))
#-
u2 = Velocity(ustar*VOverVStar(M2,RayleighFlow,gas=Air))
#-
T2 = Temperature(Tstar*TOverTStar(M2,RayleighFlow,gas=Air))
#=
Finally, let us calculate the maximum heat flux -- the heat flux that brings the
flow just to sonic conditions when starting at the given conditions at location 1.
This comes from

$$q_{max} = h_{0}^* - h_{01} = c_p (T_0^* - T_{01})$$

We have a convenience function for that, based on stagnation temperature:
=#
qmax = HeatFlux(T01,T0star,gas=Air)
#-
value(qmax,u"kJ/kg")
#=
Thus, we can add up to 1223 kJ/kg before the flow gets choked. Any more than that
will cause the entrance conditions to change.
=#


using Test
using LinearAlgebra
using PolynomialBasis
using Roots
using IntervalArithmetic
# using Revise
using ImplicitDomainQuadrature

IDQ = ImplicitDomainQuadrature

function allapprox(v1, v2; tol = 1e-14)
    @assert length(v1) == length(v2)
    flags = [isapprox(v1[i], v2[i], atol = tol) for i = 1:length(v1)]
    return all(flags)
end

f2(x) = x[1] - 0.5
P = InterpolatingPolynomial(1, 2, 2)
coeffs = [f2(P.basis.points[:, i]) for i = 1:size(P.basis.points)[2]]
update!(P, coeffs)
x0 = [1.0]
w0 = 3.0
p, w = IDQ.extend_edge_quadrature_to_surface_quadrature(
    P,
    x -> gradient(P, x),
    1,
    -1.0,
    1.0,
    x0,
    w0,
)
@test allapprox(p, [0.5, 1.0])
@test w ≈ w0

f2(x) = x[1]
P = InterpolatingPolynomial(1, 2, 2)
coeffs = [f2(P.basis.points[:, i]) for i = 1:size(P.basis.points)[2]]
update!(P, coeffs)
x0 = [-0.5 0.0 0.5]
w0 = [2.0, 3.0, 4.0]
@test_throws AssertionError IDQ.extend_edge_quadrature_to_surface_quadrature(
    P,
    x -> gradient(P, x),
    1,
    -1.0,
    1.0,
    x0,
    [2.0, 3.0, 4.0, 5.0],
)
surf_quad = IDQ.extend_edge_quadrature_to_surface_quadrature(
    P,
    x -> gradient(P, x),
    1,
    -1.0,
    1.0,
    x0,
    w0,
)
p = hcat([IDQ.extend(x0[i], 1, 0.0) for i = 1:3]...)
@test allapprox(surf_quad.points, p)
@test allapprox(surf_quad.weights, w0)

edgequad = IDQ.TemporaryQuadrature(x0, w0)
surf_quad = IDQ.extend_edge_quadrature_to_surface_quadrature(
    P,
    x -> gradient(P, x),
    1,
    Interval(-1.0, 1.0),
    edgequad,
)
@test allapprox(surf_quad.points, p)
@test allapprox(surf_quad.weights, w0)

f2(x) = x[2]
numqp = 5
quad1d = IDQ.ReferenceQuadratureRule(numqp)
xL, xR = [-1.0, -1.0], [1.0, 1.0]
P = InterpolatingPolynomial(1, 2, 2)
dP = InterpolatingPolynomial(2, 2, 2)
coeffs = [f2(P.basis.points[:, i]) for i = 1:size(P.basis.points)[2]]
update!(P, coeffs)
update_interpolating_gradient!(dP,P)
quad = IDQ.surface_quadrature(P,dP, xL, xR, numqp)
p = IDQ.extend([0.0], 1, quad1d.points)
@test allapprox(p, quad.points)
@test allapprox(quad1d.weights, quad.weights)

f2(x) = (x[2] - 0.5) * (x[2] + 0.75)
numqp = 5
quad1d = IDQ.ReferenceQuadratureRule(numqp)
xL, xR = [-1.0, -1.0], [1.0, 1.0]
P = InterpolatingPolynomial(1, 2, 2)
dP = InterpolatingPolynomial(2, 2, 2)
coeffs = [f2(P.basis.points[:, i]) for i = 1:size(P.basis.points)[2]]
update!(P, coeffs)
update_interpolating_gradient!(dP,P)
quad = IDQ.surface_quadrature(P, dP, xL, xR, numqp)

function integrate(f, quad)
    s = 0.0
    for (p, w) in quad
        s += f(p) * w
    end
    return s
end

f(x) = (x[2] ≈ -0.75 || x[2] ≈ 0.5) ? 1.0 : 0.0
s = integrate(f, quad)
@test s ≈ 4.0

function circle_distance_function(coords, center, radius)
    npts = size(coords)[2]
    return [norm(coords[:, i] - center) - radius for i = 1:npts]
end

radius = 0.5
center = [0.0, 0.0]
poly = InterpolatingPolynomial(1, 2, 3)
dpoly = InterpolatingPolynomial(2, 2, 3)
numqp = 5
xL, xR = [-1.0, -1.0], [1.0, 1.0]
coeffs = circle_distance_function(poly.basis.points, center, radius)
update!(poly, coeffs)
update_interpolating_gradient!(dpoly,poly)
squad = IDQ.surface_quadrature(poly,dpoly, xL,xR, numqp)

f(x) = isapprox(norm(x),radius,atol=1e-1) ? 1.0 : 0.0
s = integrate(f,squad)
@test isapprox(s,sum(squad.weights))
@test isapprox(s,pi,atol=1e-1)


using PkgSRA
using Plots, Random, Distributions, Interpolations
pyplot()
Random.seed!(11)
using BSON: @save
include("../scripts/paper_settings.jl");
# include("../../utils/sindy_turing_utils.jl")
# include("../../utils/sindy_statistics_utils.jl")
# include("../../utils/main_algorithm_utils.jl")

################################################################################
#####
##### SIR example
#####
include(EXAMPLE_FOLDERNAME*"example_sir.jl")

# Define the multivariate forcing function
num_ctr = 3;
    U_starts = rand(1, num_ctr) .* tspan[2]/2
    U_widths = 0.6;
    amplitude = 100.0
my_U_func_time2(t) = U_func_time(t, u0,
                        U_widths, U_starts,
                        F_dim=1,
                        amplitude=amplitude)

# Get data
sol = solve_sir_system(U_func_time=my_U_func_time2)
dat = Array(sol)
plot(ts,dat[1,:], label="S");
    plot!(ts,dat[2,:], label="I");
    plot!(ts,dat[3,:], label="R")
numerical_grad = numerical_derivative(dat, ts)
true_grad = core_dyn_true(dat)
plot(ts, true_grad[2,:]);plot!(ts, numerical_grad[2,:])

## Also get baseline true/ideal cases
# Uncontrolled
# dat_raw = Array(solve_sir_system())
# numerical_grad_raw = numerical_derivative(dat_raw, ts)
# true_grad_raw = zeros(size(dat))
# for i in 1:size(dat,2)
#     true_grad_raw[:,i] = sir_system(dat_raw[:,i], Float64.(p), [0])
# end

# True control signal
U_true = zeros(size(dat))
for (i, t) in enumerate(ts)
    U_true[:,i] = my_U_func_time2(t)
end
plot(ts,dat[1,:], label="S");
    plot!(ts,dat[2,:], label="I");
    plot!(ts,dat[3,:], label="R")
    plot!(ts,U_true', label="Control")

###
### Fit a naive model
###
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);

# Upgrade: Use AIC
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_unctr,best_criterion,all_criteria,all_models) =
    sindyc_ensemble(dat, numerical_grad, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,100),
                    use_clustering_minimization=true)
println("Naive SINDy model:")
    print_equations(sindy_unctr, digits=5)
println("True equations are:")
    print_equations(core_dyn_true, digits=5)
scatter(sum.(val_list), all_criteria)
    title!("AIC for various sparsities (naive model)")
    xlabel!("Number of nonzero terms")

# Integrate this (poor) model
# Callback aborts if it blows up
condition(u,t,integrator) = any(abs.(u).>1e4)
cb = DiscreteCallback(condition, terminate!)
prob_unctr = ODEProblem(sindy_unctr, u0, tspan, [0], callback=cb)
sindy_dat_unctr = Array(solve(prob_unctr, Tsit5(), saveat=ts));

# let d = sindy_dat_unctr, d2 = dat
#     i = 1;
#     plot(d[i,:], label="Naive model");#plot!(d[2,:]);plot!(d[3,:])
#     plot!(d2[i,:], label="Data");#plot!(d2[2,:]);plot!(d2[3,:])
# end

###
### Calculate the residuals
###
# NOTE: do not use Turing, for speed
sindy_grad1 = sindy_unctr(dat, 0)
residual = numerical_grad .- sindy_grad1
noise_guess = 2*abs(median(residual[1,:])) # Really should use Turing here

plot(residual[3,:], label="Residual")
    hline!([noise_guess], label="Noise line");
    hline!([-noise_guess], label="Noise line")
    # title!("Residual")

ctr_guess = process_residual(residual, noise_guess)

###
### Subsample based on residual
###
accepted_ind = subsample_using_residual(residual,
            noise_guess, min_length=4)

num_pts = 2000
start_ind = 1
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1]
dat_sub = dat[:,subsample_ind]
grad_sub = numerical_grad[:,subsample_ind]

# Plot the subsampled points
i = 2
    plot(dat[i,:])
    scatter!(subsample_ind, dat_sub[i,:], color=:blue)
    plot!(U_true[i,:])
    title!("Subsampled points for index $i")

# Any control signal in the subset?
# U_sub = U_true[:,subsample_ind]
# plot(U_sub')
#     title!("Control signals in the subsampled dataset")

# val_list = calc_permutations(6,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_sub,best_criterion,all_criteria,all_models,best_index) =
    sindyc_ensemble(dat_sub, grad_sub, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,2*noise_guess),
                    use_clustering_minimization=true)
println("Best index is $best_index:")
    print_equations(sindy_sub, digits=5)
println("True equations are:")
    print_equations(core_dyn_true, digits=5)
scatter(sum.(val_list), all_criteria)
    title!("AIC for various sparsities")
    xlabel!("Number of nonzero terms")

# Plot a function for the subsampled points
# f(S, I, R) = -0.00027*I.*I .- 0.0002*I.*R
# g(S, I, R) = -0.2*I
# i = 2
#     plot(g(dat[1,:], dat[2,:], dat[3,:]), label="True")
#     scatter!(subsample_ind, g(dat_sub[1,:], dat_sub[2,:], dat_sub[3,:]))
#     plot!(f(dat[1,:], dat[2,:], dat[3,:]), label="Found")
#     title!("Subsampled points for index $i")


function discrete_returns(series::Vector{Float64}) :: Vector{Float64}
    n = length(series)
    r = Vector{Float64}(undef, n-1)

    @inbounds for i in 2:n
        r[i-1] = series[i]/series[i-1] - 1
    end

    return r
end

function discrete_returns(series::Array{Float64, 2}) :: Array{Float64, 2}
    (n, num_series) = size(series)
    r = Array{Float64, 2}(undef, n-1, num_series)

    @inbounds for c in 1:num_series, i in 2:n
        r[i-1, c] = series[i, c]/series[i-1, c] - 1
    end

    return r
end

function log_returns(series::Vector{Float64}) :: Vector{Float64}
    n = length(series)
    r = Vector{Float64}(undef, n-1)

    @inbounds for i in 2:n
        r[i-1] = log(series[i]/series[i-1])
    end

    return r
end

function log_returns(series::Array{Float64, 2}) :: Array{Float64, 2}
    (n, num_series) = size(series)
    r = Array{Float64, 2}(undef, n-1, num_series)

    @inbounds for c in 1:num_series, i in 2:n
        r[i-1, c] = log(series[i, c]/series[i-1, c])
    end

    return r
end


@doc raw"""
    Point2

An object `p` of the type [`Point2`](@ref) is a point in the two-dimensional
Euclidean plane ``Ε^2``.

# Example

The following declaration creates two points with Cartesian coordinates.

```julia-repl
julia> p, q = Point2(), Point2(1.0, 2.0)
(PointC2(0, 0), PointC2(1, 2))
```
"""
Point2

"""
    Point2(::Origin = ORIGIN)

Introduces a variable `p` with Cartesian coordinates ``(0, 0)``.
"""
Point2(::Origin = ORIGIN)

"""
    Point2(x::Real, y::Real)

Introduces a point `p` initialized to ``(x, y)``.
"""
Point2(x::Real, y::Real) = Point2(convert(FT, x), convert(FT, y))

@doc raw"""
    Point2(x::Real, y::Real, hw::Real = 1)

Introduces a point `p` initialized to `(hx/hw,hy/hw)`.

!!! info "Precondition"

    `hw` ``≠`` `0`
"""
Point2(x::Real, y::Real, hw::Real) =
    Point2(convert(RT, x), convert(RT, y), convert(RT, hw))

"""
    Point2(wp::WeightedPoint2)

Introduces a point from a weighted point.
"""
Point2(wp::WeightedPoint2)

"""
    ==(p::Point2, q::Point2)

Test for equality.

Two points are equal, iff their ``x`` and ``y`` coordinates are equal. The point
can be compared with [`ORIGIN`](@ref).
"""
==(p::Point2, q::Point2)

"""
    hx(p::Point2)

Returns the homogeneous ``x`` coordinate.
"""
hx(p::Point2)

"""
    hy(p::Point2)

Returns the homogeneous ``y`` coordinate.
"""
hy(p::Point2)

"""
    hw(p::Point2)

Returns the homogenizing coordinate.
"""
hw(p::Point2)

"""
    x(p::Point2)

Returns the Cartesian ``x`` coordinate, that is [`hx()`](@ref)/[`hw()`](@ref).
"""
x(p::Point2)

"""
    y(p::Point2)

Returns the Cartesian ``y`` coordinate, that is [`hy()`](@ref)/[`hw()`](@ref).
"""
y(p::Point2)

@doc raw"""
    homogeneous(p::Point2, i::Integer)

Returns the `i`ᵗʰ homogeneous coordinate of `p`.

!!! info "Precondition"

    ``0 ≤ i ≤ 2``.
"""
homogeneous(p::Point2, i::Integer)

@doc raw"""
    cartesian(p::Point2, i::Integer)

Returns the `i`ᵗʰ Cartesian coordinate of `p`.

!!! info "Precondition"

    ``0 ≤ i ≤ 1``.
"""
cartesian(p::Point2, i::Integer)

"""
    dimension(p::Point2)

Returns the dimension (the constant 2).
"""
dimension(p::Point2)

"""
    bbox(p::Point2)

Returns a bounding box containing `p`.
"""
bbox(p::Point2)

"""
    transform(p::Point2, t::AffTransformation2)

Returns the point obtained by applying `t` on `p`.
"""
transform(p::Point2, t::AffTransformation2)

"""
    <(p::Point2, q::Point2)

Return `true` iff `p` is lexicographically smaller than `q`, i.e. either if
`x(p) < x(q)` or if `x(p) == x(q)` and `y(p) < y(q)`.
"""
<(p::Point2, q::Point2)

"""
    >(p::Point2, q::Point2)

Return `true` iff `p` is lexicographically greater than `q`.
"""
>(p::Point2, q::Point2)

"""
    <=(p::Point2, q::Point2)

Return `true` iff `p` is lexicographically smaller or equal to `q`.
"""
<=(p::Point2, q::Point2)

"""
    >=(p::Point2, q::Point2)

Return `true` iff `p` is lexicographically greater or equal to `q`.
"""
>=(p::Point2, q::Point2)

"""
    -(p::Point2, q::Point2)

Returns the difference vector between `p` and `q`.

You can substitute [`ORIGIN`](@ref) for either `p` or `q`, but not for both.
"""
-(p::Point2, q::Point2)


# # KissABC
# ## Usage guide
# The ingredients you need to use Approximate Bayesian Computation:
#
# 1. A simulation which depends on some parameters, able to generate datasets similar to your target dataset if parameters are tuned
# 1. A prior distribution over such parameters
# 1. A distance function to compare generated dataset to the true dataset
#
# We will start with a simple example, we have a dataset generated according to an Normal distribution whose parameters are unknown

tdata = randn(1000) .* 0.04 .+ 2;

# we are ofcourse able to simulate normal random numbers, so this constitutes our simulation

sim((μ, σ)) = randn(1000) .* σ .+ μ;

# The second ingredient is a prior over the parameters μ and σ

using KissABC
prior = Factored(Uniform(1, 3), Truncated(Normal(0, 0.1), 0, 100));

# we have chosen a uniform distribution over the interval [1,3] for μ and a normal distribution truncated over ℝ⁺ for σ.
#
# Now all that we need is a distance function to compare the true dataset to the simulated dataset, for this purpose comparing mean and std is optimal

function dist(x, y)
    d1 = mean(x) - mean(y)
    d2 = std(x) - std(y)
    hypot(d1, d2 * 50)
end

# Now we are all set, we can use `AIS` which is an Affine Invariant MC algorithm via the `sample` function, to simulate the posterior distribution for this model, inferring μ and σ
cost(x) = dist(tdata, sim(x))
approx_density = ApproxPosterior(prior, cost, 0.01)
res = sample(
    approx_density,
    AIS(50),
    500,
    discard_initial = 1000,
    ntransitions = 10,
    progress = true,
)

@show res

# You can also use Sequential Monte Carlo (SMC) to infer posterior parameters:

ressmc = smc(prior, cost, nparticles=500, epstol=0.01)

@show ressmc

# the parameters we chose are: a tolerance on distances equal to `0.1`, a number of samples equal to `2000`, the simulated posterior results are in `res`.
# We can now extract the inference results:

prsample = [rand(prior) for i = 1:2000] #some samples from the prior for comparison
μ_pr = getindex.(prsample, 1) # μ samples from the prior
σ_pr = getindex.(prsample, 2) # σ samples from the prior
μ_p = res[1].particles # μ samples from the posterior
σ_p = res[2].particles; # σ samples from the posterior

# and plotting prior and posterior side by side we get:

using Plots
a = stephist(
    μ_pr,
    xlims = (1, 3),
    xlabel = "μ prior",
    leg = false,
    lw = 2,
    normalize = true,
)
b = stephist(
    σ_pr,
    xlims = (0, 0.3),
    xlabel = "σ prior",
    leg = false,
    lw = 2,
    normalize = true,
)
ap = stephist(
    μ_p,
    xlims = (1, 3),
    xlabel = "μ posterior",
    leg = false,
    lw = 2,
    normalize = true,
)
bp = stephist(
    σ_p,
    xlims = (0, 0.3),
    xlabel = "σ posterior",
    leg = false,
    lw = 2,
    normalize = true,
)
plot(a, ap, b, bp)
savefig("inference.svg");
nothing; # hide

# ![inference_plot](inference.svg)
#
# we can see that the algorithm has correctly inferred both parameters, this exact recipe will work for much more complicated models and simulations, with some tuning.


using Test, LinearAlgebra, MPStates

@testset "Initialization of the Mps class" begin

"""Test whether a tensor product is left canonical."""
function is_left_canonical(A::Vector{Array{T, 3}}) where T
    output = true
    for i=1:length(A)-1
        Ai = A[i]
        rA = reshape(Ai, (size(Ai, 1)*size(Ai, 2), size(Ai, 3)))
        M = rA'*rA
        Id = Matrix{T}(I, size(Ai, 3), size(Ai, 3))
        output = isapprox(M, Id)
    end
    return output
end

"""Test whether a tensor product is right canonical."""
function is_right_canonical(A::Vector{Array{T, 3}}) where T
    output = true
    for i=length(A):-1:2
        Ai = A[i]
        rA = reshape(Ai, (size(Ai, 1), size(Ai, 2)*size(Ai, 3)))
        M = rA*rA'
        Id = Matrix{T}(I, size(Ai, 1), size(Ai, 1))
        output = isapprox(M, Id)
    end
    return output
end

@testset "make tensor state left & right canonical" begin
    # Test if state is initialized in left canonical form if it is normalized.
    rtest1 = MPStates.testMps("rtest1")
    @test is_left_canonical(rtest1.M)
    @test MPStates.norm(rtest1) ≈ 1.
    MPStates.make_right_canonical!(rtest1)
    @test is_right_canonical(rtest1.M)
    @test MPStates.norm(rtest1) ≈ 1.
end
end # @testset "Initialization of the Mps class"


### ============================================================================
### A type related to Proximal Gradient and Lasso
### ============================================================================

mutable struct ProximalGradient
    """
        Dynamical, functional, keeps references to detailed implementation
        functions.
    """

    g::Function
    gradient::Function
    prox::Function
    h::Function

    β::Float64  # For beta convex function.
    tol::Float64
    maxItr::Int64
    solutionDim::Tuple

    function ProximalGradient(tol::Float64, maxItr::Int64)
        this = new()
        this.tol = tol
        this.maxItr = maxItr
        return this
    end

end

function OptimizeProximalGradient(
        this::ProximalGradient,
        warm_start::Union{Array, Nothing}=nothing
    )::Matrix
    """
        Implement FISTA, Accelerated Proximal Gradient, copied from my HW.
    """
    tol = this.tol
    max_itr = this.maxItr

    if warm_start === nothing
        warm_start = zeros(this.solutionDim)
    end

    x  = warm_start
    Δx = Inf
    y  = x
    t  = 1
    ∇f = this.gradient(y)
    δ  = 1/this.β # stepsize
    xNew = similar(x)
    yNew = similar(y)

    while Δx >= tol && max_itr >= 1
        xNew = x - δ*∇f
        this.prox(xNew, δ)
        tNew = (1 + sqrt(1 + 4t^2))/2
        yNew .= xNew
        yNew += ((t - 1)/tNew)*(xNew - x)

        ∇f .= this.gradient(yNew)
        Δx = norm(xNew - x, 1)/norm(x, 1)

        t = tNew
        x = xNew
        y = yNew

        max_itr -= 1
    end
    if max_itr == 0
        println()
        Warn("Maximal iteration $(this.maxItr) reached for Proximal gradient. ")
        @assert norm(x, Inf) != Inf "Solution is infinite, something blowed up."
        Warn("It will be bumped up by 200% more and then we try again.")
        this.maxItr = 2*this.maxItr
        return OptimizeProximalGradient(this, warm_start)
    end

    return x

end

### ----------------------------------------------------------------------------

### ============================================================================
### Lasso Proximal Operator
### ============================================================================

function L1LassoProximal!(y::Union{Vector, Matrix}, t::Float64, λ::Float64)
    """
        Textbook definition of the L1 Lasso proximal operator.
    """

    return map!(y, y) do yi
        if yi > λ*t
            return yi - λ*t
        elseif yi < -λ*t
            return yi + λ*t
        else
            return 0
        end
    end

end

### ----------------------------------------------------------------------------

### ============================================================================
### The 2 Norm loss function with L1 regularization
### ============================================================================

function BuildPG2NormLasso(
        A::Matrix,
        b::Matrix,
        λ::Float64,
        tol::Float64=1e-5,
        max_itr::Int64=1000
    )::ProximalGradient
    """
        Given matrix A, vector b, and λ the regularization perameter,
        this will build an instance of ProximalGradient.

        * Use Lambdas to capture variables references
        * A factory methods
        ! Pass copy of matrices will be better.
    """

    @assert size(b, 2) == 1 "Expect vector b in the"*
    " shape of (n, 1) but it turns out to be $(size(b))"
    @assert size(b, 1) == size(A, 1) "Expect the size of the matrix to match "*
    "the vector but it turns out to be A is in $(size(A)), and b is $(size(b))"
    ATA = A'*A
    ATb = A'*b
    β = 4*opnorm(ATA)  # convexity from spectral norm.
    t = 1/β
    f(x) = norm(A*x - b)^2
    dg(x) = 2(ATA*x - ATb)

    # ======= build ==================================
    proxVectorized(y, t) = L1LassoProximal!(y, t, λ)
    this = ProximalGradient(tol, max_itr)
    this.g = f
    this.gradient = dg
    this.prox = proxVectorized
    this.β = β
    this.solutionDim = (size(A, 2), size(b, 2))
    this.h = (x) -> norm(x, 1)

    return this
end

function ChangeProximalGradientLassoλ!(this::ProximalGradient, λ::Float64)
    """
        Change the proximal operator for a different λ
    """
    this.prox = (y, t) -> L1LassoProximal!(y, t, λ)
end

# ------------------------------------------------------------------------------

### ============================================================================
### Lasso huber loss function and l1 norm regularization
### ============================================================================

function BuildPGHuberLasso(
    A::Matrix,
    B::Array,
    δ::Float64=1,
    λ::Float64=0
)
    # TODO: implement this

end

### ============================================================================
### Lasso Cross Entropy loss function and l1 norm regularization
### ============================================================================

function BuildPGEntropyLasso()

end


@doc raw"""
    K, γ, info = glover_mcfarlane(G::AbstractStateSpace{Continuous}, γ = 1.1)

Design a controller for `G` that maximizes the stability margin ϵ = 1/γ with normalized coprime factor uncertainty using the method of Glover and McFarlane
```
γ = 1/ϵ = ||[K;I] inv(I-G*K)*inv(M)||∞
G = inv(M + ΔM)*(N + ΔN)
```
γ is given as a relative factor above γmin and must be greater than 1, i.e., if γ = 1.1, the controller will be designed for γ = 1.1*γmin.

We want γmin (which is always ≥ 1) as small as possible, and we usually require that γmin is less than 4, corresponding to 25% allowed coprime uncertainty.

Performance modeling is incorporated in the design by calling `glover_mcfarlane` on the shaped system `W2*G*W1` and then forming the controller as `W1*K*W2`. Using this formulation, traditional loop shaping can be done on `W2*G*W1`.

# Example:
Example 9.3 from the reference below.
```julia
using RobustAndOptimalControl, ControlSystems, Plots, Test
G = tf(200, [10, 1])*tf(1, [0.05, 1])^2     |> ss
Gd = tf(100, [10, 1])                       |> ss
W1 = tf([1, 2], [1, 1e-6])                  |> ss
K, γ, info = glover_mcfarlane(G, 1.1; W1)
@test info.γmin ≈ 2.34 atol=0.005
Gcl = extended_gangoffour(G, K) # Form closed-loop system

bodeplot([G, info.Gs, G*K], lab=["G" "" "G scaled" "" "Loop transfer"]) |> display
bodeplot(Gcl, lab=["S" "KS" "PS" "T"], plotphase=false) |> display # Plot gang of four

plot( step(Gd*feedback(1, info.Gs), 3), lab="Initial controller")
plot!(step(Gd*feedback(1, G*K), 3), lab="Robustified") |> display

nyquistplot([info.Gs, G*K], ylims=(-2,1), xlims=(-2, 1),
    Ms_circles = 1.5,
    lab = ["Initial controller" "Robustified"],
    title = "Loop transfers with and without robustified controller"
    ) |> display
```

Ref: Sec 9.4.1 of Skogestad, "Multivariable Feedback Control: Analysis and Design"

# Extended help

Skogestad gives the following general advice:
1. Scale the plant outputs and inputs. This is very important for most design
    procedures and is sometimes forgotten. In general, scaling improves the
    conditioning of the design problem, it enables meaningful analysis to be made
    of the robustness properties of the feedback system in the frequency domain,
    and for loop-shaping it can simplify the selection of weights. There are a variety
    of methods available including normalization with respect to the magnitude of
    the maximum or average value of the signal in question. If one is to go straight to a design the following variation has
    proved useful in practice:
    (a) The outputs are scaled such that equal magnitudes of cross-coupling into each
        of the outputs is equally undesirable.
    (b) Each input is scaled by a given percentage (say 10%) of its expected range
        of operation. That is, the inputs are scaled to reflect the relative actuator
        capabilities. An example of this type of scaling is given in the aero-engine
        case study of Chapter 12.
2. Order the inputs and outputs so that the plant is as diagonal as possible. The
    relative gain array [`rga`](@ref) can be useful here. The purpose of this pseudo-diagonalization
    is to ease the design of the pre- and post-compensators which, for simplicity, will
    be chosen to be diagonal.

    Next, we discuss the selection of weights to obtain the shaped plant $G_s = W_2 G W_1$
    where $W_1 = W_p W_a W_g$
3. Select the elements of diagonal pre- and post-compensators $W_p$ and $W_2$ so that
    the singular values of $W_2 G W_p$ are desirable. This would normally mean high
    gain at low frequencies, roll-off rates of approximately 20 dB/decade (a slope of
    about 1) at the desired bandwidth(s), with higher rates at high frequencies. Some
    trial and error is involved here. $W_2$ is usually chosen as a constant, reflecting the
    relative importance of the outputs to be controlled and the other measurements
    being fed back to the controller. For example, if there are feedback measurements
    of two outputs to be controlled and a velocity signal, then $W_2$ might be chosen
    to be `diag([1, 1, 0.1])`, where 0.1 is in the velocity signal channel. $W_p$ contains the
    dynamic shaping. Integral action, for low frequency performance; phase-advance
    for reducing the roll-off rates at crossover, and phase-lag to increase the roll-off
    rates at high frequencies should all be placed in $W_p$ if desired. The weights should
    be chosen so that no unstable hidden modes are created in $G_s$.
5. Optional: Introduce an additional gain matrix $W_g$ cascaded with $W_a$ to provide
    control over actuator usage. $W_g$ is diagonal and is adjusted so that actuator rate
    limits are not exceeded for reference demands and typical disturbances on the
    scaled plant outputs. This requires some trial and error.

6. Robustly stabilize the shaped plant $G_s = W_2 G W_1$ , where $W_1 = W_p W_a W_g$,
    using `glover_mcfarlane`. First, the maximum stability
    margin $ϵ_{max} = 1/γ_{min}$ is calculated. If the margin is too small, $ϵmax < 0.25$, then go back and modify the weights. Otherwise, a γ-suboptimal controller is synthesized. There is usually no advantage to be gained by using the optimal controller. When $ϵ_{max}$ > 0.25
    (respectively $γ_{min}$ < 4) the design is usually successful. In this case, at least
    25% coprime factor uncertainty is allowed, and we also find that the shape of the
    open-loop singular values will not have changed much after robust stabilization.
    A small value of ϵmax indicates that the chosen singular value loop-shapes are
    incompatible with robust stability requirements. That the loop-shapes do not
    change much following robust stabilization if γ is small (ϵ large), is justified
    theoretically in McFarlane and Glover (1990).

7. Analyze the design and if all the specifications are not met make further
    modifications to the weights.
8. Implement the controller. The configuration shown in below has been found
    useful when compared with the conventional set up. This is because
    the references do not directly excite the dynamics of $K$, which can result in large amounts of overshoot (classical derivative kick). The constant prefilter ensures a steady-state gain of 1 between r and y, assuming integral action in $W_1$ or $G$ (note, the K returned by this function has opposite sign compared to that of Skogestad, so we use negative feedback here).

Anti-windup can be added to $W_1$ but putting $W_1$ on Hanus form after the synthesis, see [`hanus`](@ref).

```
       ┌─────────┐      ┌────────┐      ┌────────┐
    r  │         │    us│        │  u   │        │  y
   ───►│(K*W2)(0)├──+──►│   W1   ├─────►│   G    ├────┬──►
       │         │  │-  │        │      │        │    │
       └─────────┘  │   └────────┘      └────────┘    │
                    │                                 │
                    │                                 │
                    │   ┌────────┐      ┌────────┐    │
                    │   │        │  ys  │        │    │
                    └───┤   K    │◄─────┤   W2   │◄───┘
                        │        │      │        │
                        └────────┘      └────────┘
```
"""
function glover_mcfarlane(G::AbstractStateSpace{Continuous}, γ = 1.1; W1=1, W2=1)
    γ > 1 || throw(ArgumentError("γ must be greater than 1"))
    Gs = W2*G*W1
    A,B,C,D = ssdata(Gs)

    R = I + D*D'
    S = I + D'D
    Sl = lu(S)
    # arec(A, B, R, Q, S) solves A'X + XA - (XB+S)R^(-1)(B'X+S') + Q = 0

    Ā = A - B*(S\D'C)
    Z,_ = MatrixEquations.arec(Ā', C', R, B*(Sl\B'))
    X,_ = MatrixEquations.arec(Ā, B, S, C'*(R\C))

    γmin = sqrt(1 + ρ(X*Z))

    γ *= γmin

    L = (1-γ^2)*I + X*Z
    F = -(Sl\(D'C + B'X))
    BK = γ^2*(L'\Z)*C'
    AK = A + B*F + BK*(C + D*F)
    CK = B'X
    DK = -D'
    Ks = -ss(AK, BK, CK, DK)
    Gcl = extended_gangoffour(Gs, Ks)
    imargin, ω = hinfnorm2(Gcl)
    K = W2*Ks*W1
    Gcl = extended_gangoffour(G, K)
    K, γ, (; Gcl, margin = inv(imargin), ω, γmin, Ks, Gs)
end

"Spectral radius"
function ρ(X)
    e = eigvals(X)
    abs(e[end])
end

"""
    `Wh` = hanus(W)

Return `Wh` on Hanus form. `Wh` has twice the number of inputs, where the second half of the inputs are "actual inputs", e.g., potentially saturated. This is used to endow `W` with anti-windup protection.
`W` must have an invertable `D` matrix and be minimum phase.

Ref: Sec 9.4.5 of Skogestad, "Multivariable Feedback Control: Analysis and Design"
"""
function hanus(W)
    A,B,C,D = ssdata(W)
    nu = W.nu
    BD = B/D
    A2 = A - BD*C
    B2 = [0*I(nu) BD]
    D2 = [D 0*I(nu)]
    ss(A2, B2, C, D2)
end

"""
    extended_gangoffour(P, C)

Returns a single statespace system that maps
- `w1` reference or measurement noise
- `w2` load disturbance
to
- `z1` control error
- `z2` control input
```
      z1          z2
      ▲  ┌─────┐  ▲      ┌─────┐
      │  │     │  │      │     │
w1──+─┴─►│  C  ├──┴───+─►│  P  ├─┐
    │    │     │      │  │     │ │
    │    └─────┘      │  └─────┘ │
    │                 w2         │
    └────────────────────────────┘
```

The returned system has the transfer-function matrix
```math
\\begin{bmatrix}
I \\\\ C
\\end{bmatrix} (I + PC)^{-1} \\begin{bmatrix}
I & P
\\end{bmatrix}
```

The gang of four can be plotted like so
```julia
Gcl = extended_gangoffour(G, C) # Form closed-loop system
bodeplot(Gcl, lab=["S" "CS" "PS" "T"], plotphase=false) |> display # Plot gang of four
```
Note, the last output of Gcl is the negative of the `CS` and `PS` transfer functions from `gangoffour2`.
See [`glover_mcfarlane`](@ref) for an extended example. See also [`ncfmargin`](@ref).
"""
function extended_gangoffour(P, C)
    ny,nu = size(P)
    S = feedback(ss(I(ny+nu), P.timeevol), [0*I(ny) P; -C 0*I(nu)], pos_feedback=true)
    Gcl = S + cat(0*I(ny), -I(nu), dims=(1,2))
    Gcl
end

"""
    m, ω = ncfmargin(P, K)

Normalized coprime factor margin, defined has the inverse H∞ norm of
```math
\\begin{bmatrix}
I \\\\ K
\\end{bmatrix} (I + PK)^{-1} \\begin{bmatrix}
I & P
\\end{bmatrix}
```
A margin ≥ 0.25-0.3 is a reasonable for robustness.

If controller `K` stabilizes `P` with margin `m`, then `K` will also stabilize `P̃` if `nugap(P, P̃) < m`.

See also [`extended_gangoffour`](@ref), [`diskmargin`](@ref).
"""
function ncfmargin(P, K)
    Gcl = extended_gangoffour(P, K)
    im, w = hinfnorm2(Gcl)
    inv(im), w
end


function Linear_Spline(t::Vector,V::Vector)
  # t must be the time vector
  # V is any vector that you are interpolating

  # remove any repeating values
  M = Array(Bool,length(t)); M[1:length(t)] = false;
  for i in 1:length(t)-1
      if t[i]==t[i+1]
          M[i]=true
      else
          M[i]=false
      end
  end

  rm_idx = find(M)

  if (length(t)==length(rm_idx))
    error("No time has elapsed and there will be an issue with interpolation! -----> Do not even try to simulate this data!")
  end

  # initialize vetors
  t_new = Array(Float64,(length(t)-length(rm_idx)));
  V_new = Array(Float64,(length(t)-length(rm_idx)));

  q=1;
  for i in 1:length(t)
      if !M[i]
          t_new[q] = t[i];
          V_new[q] = V[i];
          q=q+1
      end
  end

  # make interpolant using Dierckx.jl
  #Spline1D(t_new,V_new,k=1)    # linear spline

  # make interpolant using Interpolations.jl
  knots = (t_new,)
  interpolate(knots,V_new,Gridded(Linear()))
end


"""
    GammaExpert(k, θ)

PDF:

```math
f(x; k, \\theta) = \\frac{x^{k-1} e^{-x/\\theta}}{\\Gamma(k) \\theta^k},
\\quad x > 0
```

See also: [Gamma Distribution](https://en.wikipedia.org/wiki/Gamma_distribution) (Wikipedia), shape-scale parameterization

"""
struct GammaExpert{T<:Real} <: NonZIContinuousExpert
    k::T
    θ::T
    GammaExpert{T}(k::T, θ::T) where {T<:Real} = new{T}(k, θ)
end

function GammaExpert(k::T, θ::T; check_args=true) where {T <: Real}
    check_args && @check_args(GammaExpert, k >= zero(k) && θ > zero(θ))
    return GammaExpert{T}(k, θ)
end

## Outer constructors
GammaExpert(k::Real, θ::Real) = GammaExpert(promote(k, θ)...)
GammaExpert(k::Integer, θ::Integer) = GammaExpert(float(k), float(θ))
GammaExpert() = GammaExpert(1.0, 1.0)

## Conversion
function convert(::Type{GammaExpert{T}}, k::S, θ::S) where {T <: Real, S <: Real}
    GammaExpert(T(k), T(θ))
end
function convert(::Type{GammaExpert{T}}, d::GammaExpert{S}) where {T <: Real, S <: Real}
    GammaExpert(T(d.k), T(d.θ), check_args=false)
end
copy(d::GammaExpert) = GammaExpert(d.k, d.θ, check_args=false)

## Loglikelihood of Expert
logpdf(d::GammaExpert, x...) = (d.k < 1 && x... <= 0.0) ? -Inf : Distributions.logpdf.(Distributions.Gamma(d.k, d.θ), x...)
pdf(d::GammaExpert, x...) = (d.k < 1 && x... <= 0.0) ? 0.0 : Distributions.pdf.(Distributions.Gamma(d.k, d.θ), x...)
logcdf(d::GammaExpert, x...) = (d.k < 1 && x... <= 0.0) ? -Inf : Distributions.logcdf.(Distributions.Gamma(d.k, d.θ), x...)
cdf(d::GammaExpert, x...) = (d.k < 1 && x... <= 0.0) ? 0.0 : Distributions.cdf.(Distributions.Gamma(d.k, d.θ), x...)

## expert_ll, etc
expert_ll_exact(d::GammaExpert, x::Real) = LRMoE.logpdf(d, x)
function expert_ll(d::GammaExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_ll = (yl == yu) ? logpdf.(d, yl) : logcdf.(d, yu) + log1mexp.(logcdf.(d, yl) - logcdf.(d, yu))
    expert_ll = (tu == 0.) ? -Inf : expert_ll
    return expert_ll
end
function expert_tn(d::GammaExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu))
    expert_tn = (tu == 0.) ? -Inf : expert_tn
    return expert_tn
end
function expert_tn_bar(d::GammaExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn_bar = (tl == tu) ? 0.0 : log1mexp.(logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu)))
    return expert_tn_bar
end

exposurize_expert(d::GammaExpert; exposure = 1) = d

## Parameters
params(d::GammaExpert) = (d.k, d.θ)
function params_init(y, d::GammaExpert)
    pos_idx = (y .> 0.0)
    μ, σ2 = mean(y[pos_idx]), var(y[pos_idx])
    θ_init = σ2/μ
    k_init = μ/θ_init
    if isnan(θ_init) || isnan(k_init)
        return GammaExpert()
    else
        return GammaExpert(k_init, θ_init)
    end
end

## KS stats for parameter initialization
function ks_distance(y, d::GammaExpert)
    p_zero = sum(y .== 0.0) / sum(y .>= 0.0)
    return max(abs(p_zero-0.0), (1-0.0)*HypothesisTests.ksstats(y[y .> 0.0], Distributions.Gamma(d.k, d.θ))[2])
end

## Simululation
sim_expert(d::GammaExpert) = Distributions.rand(Distributions.Gamma(d.k, d.θ), 1)[1]

## penalty
penalty_init(d::GammaExpert) = [2.0 10.0 2.0 10.0]
no_penalty_init(d::GammaExpert) = [1.0 Inf 1.0 Inf]
penalize(d::GammaExpert, p) = (p[1]-1)*log(d.k) - d.k/p[2] + (p[3]-1)*log(d.θ) - d.θ/p[4]

## statistics
mean(d::GammaExpert) = mean(Distributions.Gamma(d.k, d.θ))
var(d::GammaExpert) = var(Distributions.Gamma(d.k, d.θ))
quantile(d::GammaExpert, p) = quantile(Distributions.Gamma(d.k, d.θ), p)
lev(d::GammaExpert, u) = d.θ*d.k*gamma_inc(float(d.k+1), u/d.θ, 0)[1] + u*(1-gamma_inc(float(d.k), u/d.θ,0)[1])
excess(d::GammaExpert, u) = mean(d) - lev(d, u)

## Misc functions for E-Step
function _diff_dist_series(d::GammaExpert, yl, yu)
    return cdf.(Distributions.Gamma(d.k+1, d.θ), yu) - cdf.(Distributions.Gamma(d.k+1, d.θ), yl)
end

function _int_obs_Y(d::GammaExpert, yl, yu, expert_ll)
    if yl == yu
        return yl
    else
        return exp(-expert_ll) * ( _diff_dist_series(d, yl, yu) * d.k * d.θ )
    end
end

function _int_lat_Y(d::GammaExpert, tl, tu, expert_tn_bar)
    if tl == tu
        return tl
    else
        return exp(-expert_tn_bar) * ( (1 - _diff_dist_series(d, tl, tu)) * d.k * d.θ )
    end
end

function _int_logy_func(d::GammaExpert, x)
    return (iszero(x) || isinf(x)) ? 0.0 : log(x) * pdf.(d, x)
end

function _int_obs_logY_raw(d::GammaExpert, yl, yu)
    if yl == yu
        return log(yl) * pdf.(d, yl)
    else
        return quadgk.(x -> _int_logy_func(d, x), yl, yu, rtol = 1e-8)[1]
    end
end

function _int_lat_logY_raw(d::GammaExpert, tl, tu)
    return (tl==0 ? 0.0 : quadgk.(x -> _int_logy_func(d, x), 0.0, tl, rtol = 1e-8)[1]) + (isinf(tu) ? 0.0 : quadgk.(x -> _int_logy_func(d, x), tu, Inf, rtol = 1e-8)[1])
end

function _gamma_k_to_θ(k, sum_term_zkz, sum_term_zkzy;
                        penalty = true, pen_pararms_jk = [1.0 Inf 1.0 Inf])

    if penalty
        hyper1, hyper2 = pen_pararms_jk[3], pen_pararms_jk[4]
        if hyper1 == 1.0 && hyper2 == Inf
            # return sum(term_zkzy)[1] / (sum(term_zkz)[1] * k)
            return sum_term_zkzy / (sum_term_zkz * k)
        end
        quad1, quad2 = ( k * sum_term_zkz - (hyper1 - 1) )^2, (4/hyper2) * sum_term_zkzy
        return (hyper2/2) * ( (hyper1-1) - k*sum_term_zkz + sqrt(quad1 + quad2) )
    else
        return sum_term_zkzy / (sum_term_zkz * k)
    end
end

function _gamma_optim_k(logk,
                        sum_term_zkz, sum_term_zkzy, sum_term_zkzlogy;
                        penalty = true, pen_pararms_jk = [1.0 Inf 1.0 Inf])
    # Optimization in log scale for unconstrained computation
    k_tmp = exp(logk)

    θ_tmp = _gamma_k_to_θ(k_tmp, sum_term_zkz, sum_term_zkzy, penalty = penalty, pen_pararms_jk = pen_pararms_jk)

    obj = (k_tmp-1)*sum_term_zkzlogy - (1/θ_tmp)*sum_term_zkzy - (k_tmp*log(θ_tmp) + loggamma(k_tmp))*sum_term_zkz
    p = penalty ? (pen_pararms_jk[1]-1)*log(k_tmp) - k_tmp/pen_pararms_jk[2] + (pen_pararms_jk[3]-1)*log(θ_tmp) - θ_tmp/pen_pararms_jk[4] : 0.0
    return (obj + p) * (-1.0)
end

## EM: M-Step
function EM_M_expert(d::GammaExpert,
                     tl, yl, yu, tu,
                     exposure,
                    #  expert_ll_pos,
                    #  expert_tn_pos,
                    #  expert_tn_bar_pos,
                     z_e_obs, z_e_lat, k_e;
                     penalty = true, pen_pararms_jk = [1.0 Inf 1.0 Inf])

    # Further E-Step
    expert_ll_pos = expert_ll.(d, tl, yl, yu, tu)
    expert_tn_bar_pos = expert_tn_bar.(d, tl, yl, yu, tu)

    Y_e_obs = vec(_int_obs_Y.(d, yl, yu, expert_ll_pos))
    Y_e_lat = vec(_int_lat_Y.(d, tl, tu, expert_tn_bar_pos))
    nan2num(Y_e_obs, 0.0) # get rid of NaN
    nan2num(Y_e_lat, 0.0) # get rid of NaN

    yl_yu_unique = unique_bounds(yl, yu)
    int_obs_logY_tmp = _int_obs_logY_raw.(d, yl_yu_unique[:,1], yl_yu_unique[:,2])
    logY_e_obs = exp.(-expert_ll_pos) .* int_obs_logY_tmp[match_unique_bounds(hcat(vec(yl), vec(yu)), yl_yu_unique)]
    nan2num(logY_e_obs, 0.0) # get rid of NaN

    tl_tu_unique = unique_bounds(tl, tu)
    int_lat_logY_tmp = _int_lat_logY_raw.(d, tl_tu_unique[:,1], tl_tu_unique[:,2])
    logY_e_lat = exp.(-expert_tn_bar_pos) .* int_lat_logY_tmp[match_unique_bounds(hcat(vec(tl), vec(tu)), tl_tu_unique)]
    nan2num(logY_e_lat, 0.0) # get rid of NaN

    # Update parameters
    pos_idx = (yu .!= 0.0)
    term_zkz = z_e_obs[pos_idx] .+ (z_e_lat[pos_idx] .* k_e[pos_idx])
    term_zkz_Y = (z_e_obs[pos_idx] .* Y_e_obs[pos_idx]) .+ (z_e_lat[pos_idx] .* k_e[pos_idx] .* Y_e_lat[pos_idx])
    term_zkz_logY = (z_e_obs[pos_idx] .* logY_e_obs[pos_idx]) .+ (z_e_lat[pos_idx] .* k_e[pos_idx] .* logY_e_lat[pos_idx])

    logk_new = Optim.minimizer( Optim.optimize(x -> _gamma_optim_k(x, sum(term_zkz)[1], sum(term_zkz_Y)[1], sum(term_zkz_logY)[1],
                                               penalty = penalty, pen_pararms_jk = pen_pararms_jk),
                                               max(log(d.k)-2.0, 0.0), log(d.k)+2.0 )) # ,
                                               # log(d.k)-2.0, log(d.k)+2.0 )) # ,
                                               # GoldenSection() ))
                                               # , rel_tol = 1e-8) )
    k_new = exp(logk_new)
    θ_new = _gamma_k_to_θ(k_new, sum(term_zkz)[1], sum(term_zkz_Y)[1], penalty = penalty, pen_pararms_jk = pen_pararms_jk)

    # println("$k_new, $θ_new")
    return GammaExpert(k_new, θ_new)
end

## EM: M-Step, exact observations
function EM_M_expert_exact(d::GammaExpert,
                            ye, exposure,
                            z_e_obs;
                            penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    # Further E-Step
    Y_e_obs = ye
    logY_e_obs = log.(ye)
    # nan2num(logY_e_obs, 0.0) # get rid of NaN

    # Update parameters
    pos_idx = (ye .!= 0.0)
    term_zkz = z_e_obs[pos_idx]
    term_zkz_Y = z_e_obs[pos_idx] .* Y_e_obs[pos_idx]
    term_zkz_logY = z_e_obs[pos_idx] .* logY_e_obs[pos_idx]

    logk_new = Optim.minimizer( Optim.optimize(x -> _gamma_optim_k(x, sum(term_zkz)[1], sum(term_zkz_Y)[1], sum(term_zkz_logY)[1],
                                               penalty = penalty, pen_pararms_jk = pen_pararms_jk),
                                               max(log(d.k)-2.0, 0.0), log(d.k)+2.0 ))

    k_new = exp(logk_new)
    θ_new = _gamma_k_to_θ(k_new, sum(term_zkz)[1], sum(term_zkz_Y)[1], penalty = penalty, pen_pararms_jk = pen_pararms_jk)

    # println("$k_new, $θ_new")
    return GammaExpert(k_new, θ_new)
end