using Modia
using Tables
using DataFrames
using StatsPlots
using BenchmarkTools

SIR = Model(
    equations = :[
        N = S + I + R
        der(S) = -β*c*I/N*S
        der(I) = β*c*I/N*S - γ*I
        der(R) = γ*I
    ]
);

δt = 0.1
tmax = 40.0;

u0 = Map(S = Var(init=990.0),
         I = Var(init=10.0),
         R = Var(init=0.0));

p = Map(β = 0.05,
        c = 10.0,
        γ = 0.25);

model = SIR | u0 | p;

sir = @instantiateModel(model);

simulate!(sir,Tsit5(),stopTime=tmax,interval=δt);

result = sir.result_x.u # extract result
result = hcat(result...) # concatenate vectors
result = result' # transpose
result = Tables.table(result) # convert to table
df_modia = DataFrame(result) # convert to DataFrame
rename!(df_modia,["S","I","R"]) # rename
df_modia[!,:t] = sir.result_x.t; # add in time

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

@benchmark simulate!(sir,Tsit5(),stopTime=tmax,interval=δt)


using Statistics, SpecialFunctions

"""
    tnvar(a, b)

Variance of the standard normal distribution truncated to the interval [a,b].
"""
function tnvar(a::Real, b::Real)
    if a == b
        return zero(middle(a, b))
    elseif a < b
        m1 = tnmean(a, b)
        m2 = √tnmom2(a, b)
        return (m2 - m1) * (m2 + m1)

        # m1 = tnmean(a, b)
        # @assert a ≤ m1 ≤ b
        # return tnmom2c(m1, a, b)
    else
        return oftype(middle(a, b), NaN)
    end
end

"""
    tnvar(a, b, μ, σ)

Variance of the truncated normal distribution, where μ, σ are the mean and
standard deviation of the untruncated distribution.
"""
function tnvar(a, b, μ, σ)
    α = (a - μ) / σ
    β = (b - μ) / σ
    return tnvar(α, β) * σ^2
end


module AnalyzeTimeseries

using FFTW
include("LinearRegression.jl")

function detrend(
    y :: Array{T, 1}
) where T <: AbstractFloat

N = length(y)
t = convert(Array{T}, 1:N |> collect)

# detrend
β = LinearRegression(t, y)

return y - (β[1] .+ β[2] * t)

end

function SpectralVariance(
    y :: Array{T, 1},
) where T <: AbstractFloat

N = length(y)
c = FFTW.fft(detrend(y))

return ((abs.(c)).^2)[2: 1+floor(Int, N / 2.0)]

end

EucNorm = x -> x / sum(x.^2)
end


function lin_int(grid,vals)
    function func(x)
        if x<grid[1]
            print("Error")
        elseif x>grid[n]
            print("Error")
        else index=searchsortedlast(grid,x)
            return (vals[index+1]-vals[index])/(grid[index+1]-grid[index])*(x-grid[index])+vals[index]
        end
    end
    return func
end


#
# Plot the equilibrium solutions for normalized CR3BP dynamics.
#

"""
Plot specified lagrange points in the rotating
reference frame of CR3BP system μ.
"""
function lagrangeplot(μ::Real, L=1:5; kwargs...)

    defaults = (; title  = "Nondimensional Lagrange Points",
                  xlabel = "X (DU)", ylabel="Y (DU)",
                  labels = ["Body 1" "Body 2" [string("L",i) for i ∈ L]...],
                  legend = :topleft, dpi=150)
    options  = merge(defaults, kwargs)

    lagrange_points = lagrange(μ)

    fig = scatter([-μ], [0]; markersize=10, markercolor=:lightblue, label="Body 1")
    scatter!(fig, [1-μ], [0]; markersize=6, markercolor=:gray, label="Body 2")

    colors = (:red, :orange, :tan, :cyan, :indigo)
    for point ∈ zip(lagrange_points, 1:length(lagrange_points))
        p, i = point
        scatter!(fig, [p[1]], [p[2]];
                markershape=:x, markercolor=colors[i], label=string("L", i))
    end

    scatter!(fig; options...)
    for i ∈ 1:min(length(fig.series_list), length(options.labels))
        fig.series_list[i].plotattributes[:label] = options.labels[i]
    end

    fig

end

"""
Plot specified lagrange points in the rotating
reference frame of CR3BP system μ.
"""
function lagrangeplot(sys::CircularRestrictedThreeBodySystem, L=1:5; normalize = true, exclude_z = true, kwargs...)

    defaults = (; title  = sys.name != "" ? "$(sys.name) Lagrange Points" : "Lagrange Plots",
                  xlabel = "X ($(string(normalized_length_unit(sys))))", ylabel="Y ($(string(normalized_length_unit(sys))))",
                  labels = [string("L",i) for i ∈ L],
                  legend = :topleft, dpi=150, aspect_ratio = 1)
    options  = merge(defaults, kwargs)

    μ = normalized_mass_parameter(sys)
    lagrange_points = lagrange(μ)

    colors = (:red, :orange, :tan, :cyan, :indigo)
    fig = scatter(; options...)
    for point ∈ zip(lagrange_points, 1:length(lagrange_points))
        p, i = point
        if !normalize
            p .= ustrip.(lengthunit(sys), p .* normalized_length_unit(sys))
        end
        if exclude_z
            scatter!(fig, [p[1]], [p[2]]; markershape=:x, markercolor=colors[i], markerstrokewidth=2, label=string("L", i))
        else
            scatter!(fig, [p[1]], [p[2], [p[3]]]; markershape=:x, markercolor=colors[i], markerstrokewidth=2, label=string("L", i))
        end
    end

    for i ∈ 1:min(length(fig.series_list), length(options.labels))
        fig.series_list[i].plotattributes[:label] = options.labels[i]
    end

    return fig

end

"""
Plot specified lagrange points in the rotating
reference frame of CR3BP system μ to the last figure.
"""
function lagrangeplot!(fig, sys::CircularRestrictedThreeBodySystem, L=1:5; normalize = true, exclude_z = true, kwargs...)

    defaults = (; )
    options  = merge(defaults, kwargs)

    μ = normalized_mass_parameter(sys)
    lagrange_points = lagrange(μ)

    colors = (:red, :orange, :tan, :cyan, :indigo)
    scatter!(fig; options...)
    for point ∈ zip(lagrange_points, 1:length(lagrange_points))
        p, i = point
        if !normalize
            p .= ustrip.(lengthunit(sys), p .* normalized_length_unit(sys))
        end
        if exclude_z
            scatter!(fig, [p[1]], [p[2]]; markershape=:x, markercolor=colors[i], markerstrokewidth=2, label=string("L", i))
        else
            scatter!(fig, [p[1]], [p[2], [p[3]]]; markershape=:x, markercolor=colors[i], markerstrokewidth=2, label=string("L", i))
        end
    end

    return fig

end


#=
This shows some data wrangling using the Card data
on returns to education
=#

##
using DelimitedFiles, CSV, DataFrames, DataFramesMeta
card = CSV.read("../Data/card.csv")
display(head(card))

##
# add some variables
card[:lnwage] = round.(log.(card[:wage]), digits=4) # round to save storage space
card[:expsq] = (card[:exper].^2)/100
card[:agesq] = card[:age].^2

##
# select the ones we want
cooked = @select(card, :lnwage, :educ, :exper, :expsq, :black, :south, :smsa, :age, :agesq, :nearc4 )
display(head(cooked))

##
# write as CSV
CSV.write("cooked.csv", cooked)

##
# write as plain ASCII
data = convert(Array{Float64}, cooked)
writedlm("cooked.txt", data)

# for binary write, see JLD.jl and Feather.jl


using CorrelationSpectroscopy
using Test

using Gtk  # https://github.com/JuliaPlots/VisualRegressionTests.jl/issues/30
using LinearAlgebra
using Plots
using Statistics
using VisualRegressionTests

"""
    demodata()

Calculate a demo dataset.

The demo dataset consists of fifteen signals at specific intervals, with peaks
at 10 and 20 rising in intensity, whereas peaks at 30 and 40 are decreasing in
intensity.
The example was inpired by this [one](https://en.wikipedia.org/wiki/Two-dimensional_correlation_analysis#Properties_of_2D_correlation_analysis>).

This function might go into the main code in the future if there is demand.

# Examples
```jldoctest
julia> x, A = demodata();

julia> length(x)
500

julia> size(A)
(15, 500)
```
"""
function demodata()
    n = 15
    x = range(0, 50, length=500)

    a₀ = 0.175
    A₁(n) = a₀ + 0.0125exp(-(n - 7)^2) + ((0.2375 - a₀) - exp(-0.8(n + 2.5)))
    A₂(n) = a₀ - 0.037exp(-0.2(n - 7)^2) + 1.025(0.2375 - a₀) / (1 + exp(-(n - 7)))
    A₃(n) = a₀*exp(-0.17(n -1)) + 0.0045(n - 1)
    A₄(n) = a₀ + 0.03exp(-0.3(n - 6)^2) - 1.6(0.2375 - a₀) / (1 + exp(-(n - 7)))

    σ = 2.0
    A = [
        A₁(t) * gaussian(xᵢ, 10, σ) +
        A₂(t) * gaussian(xᵢ, 20, σ) +
        A₃(t) * gaussian(xᵢ, 30, σ) +
        A₄(t) * gaussian(xᵢ, 40, σ) +
        0.02 * gaussian(xᵢ, 25, 10σ) +
        0.001 * randn()
        for t in 1:n, xᵢ in x
    ]

    return x, A
end
gaussian(x, μ=0.0, σ=1.0) = exp(-((x - μ) / σ)^2)

@testset "CorrelationSpectroscopy.jl" begin
    include("test_nodahilbert.jl")
    include("test_demodata.jl")
    include("test_synchronous.jl")
    include("test_asynchronous.jl")
    include("test_disrelation.jl")
end


# ---
# title: 454. 4Sum II
# id: problem454
# author: AquaIndigo
# date: 2020-11-27
# difficulty: Medium
# categories: Hash Table, Binary Search
# link: <https://leetcode.com/problems/4sum-ii/description/>
# hidden: true
# ---
#
# Given four lists A, B, C, D of integer values, compute how many tuples `(i, j,
# k, l)` there are such that `A[i] + B[j] + C[k] + D[l]` is zero.
#
# To make problem a bit easier, all A, B, C, D have same length of N where 0 ≤ N
# ≤ 500\. All integers are in the range of -228 to 228 \- 1 and the result is
# guaranteed to be at most 231 \- 1.
#
# **Example:**
#
#
#
#     Input:
#     A = [ 1, 2]
#     B = [-2,-1]
#     C = [-1, 2]
#     D = [ 0, 2]
#
#     Output:
#     2
#
#     Explanation:
#     The two tuples are:
#     1. (0, 0, 0, 1) -> A[0] + B[0] + C[0] + D[1] = 1 + (-2) + (-1) + 2 = 0
#     2. (1, 1, 0, 0) -> A[1] + B[1] + C[0] + D[0] = 2 + (-1) + (-1) + 0 = 0
#
#
#
#
#
## @lc code=start
using LeetCode

function four_sum_count(A::Vector{Int}, B::Vector{Int}, C::Vector{Int}, D::Vector{Int})::Int
    dic = counter([a + b for a in A for b in B])
    return sum(get(dic, -c - d, 0) for c in C for d in D)
end
## @lc code=end


# This set of functions is based entirely on the excellent
# scikit-learn package's method of calculating an ROC's area
# under the curve (AUC). The functions below produce the
# same output as scikit-learn.

function _binary_clf_curve(y_true, y_score)
    y_true = y_true .== 1       # make y_true a boolean vector
    desc_score_indices = sortperm(y_score, rev = true)

    y_score = y_score[desc_score_indices]
    y_true = y_true[desc_score_indices]

    distinct_value_indices = find(diff(y_score))
    threshold_idxs = push!(distinct_value_indices, length(y_score))

    tps = cumsum(y_true)[threshold_idxs]
    fps = threshold_idxs - tps
    return (fps, tps, y_score[threshold_idxs])
end

# y = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1, 1, 1]
# y_score = [0.3, 0.2, 0.3, 0.23, 0.5, 0.34, 0.45, 0.54, 0.6, 0.7, 0.8, 0.65, 0.5, 0.4, 0.3, 0.2, 0.6, 0.7, 0.5, 0.2, 0.1, 0.7, 0.2, 0.7, 0.4]

# _binary_clf_curve(y, y_score)

function roc_curve(y_true, y_score)
    fps, tps, thresholds = _binary_clf_curve(y_true, y_score)
    fpr = fps/fps[end]
    tpr = tps/tps[end]
    return (fpr, tpr, thresholds)
end

# roc_curve(y, y_score)

"""
Given 2 vectors, `x` and `y`, this function returns the indices that
sort the elements by `x`, with `y` breaking ties. See the example below.

julia> a = [2, 1, 3, 2]
julia> b = [3, 4, 1, 0]
julia> order = sortperm2(a, b)
4-element Array{Int64,1}:
 2
 4
 1
 3

julia> hcat(a[order], b[order]
4×2 Array{Int64,2}:
 1  4
 2  0
 2  3
 3  1
 """
function sortperm2(x, y; rev = false)
    n = length(x)
    no_ties = n == length(Set(x))
    if no_ties
        res = sortperm(x, rev = rev)
    else
        ord1 = sortperm(x, rev = rev)
        x_sorted = x[ord1]
        i = 1
        while i < n
            println(i)
            # println("x_i is $(x_sorted[i]) and x_(i+1) is $(x_sorted[i+1])")
            if x_sorted[i] == x_sorted[i+1]
                if rev && y[ord1][i] < y[ord1][i+1]
                    println("(1.) Switching $(y[ord1][i]) with $(y[ord1][i+1])")
                    ord1[i], ord1[i+1] = ord1[i+1], ord1[i]
                    i = i > 1 ? i - 1 : i
                    continue
                elseif !rev && y[ord1][i] > y[ord1][i+1]
                    println("(2.) Switching $(y[ord1][i]) with $(y[ord1][i+1])")
                    ord1[i], ord1[i+1] = ord1[i+1], ord1[i]
                    i = i > 1 ? i - 1 : i
                    continue
                end
            end
            i += 1
        end
        res = ord1
    end
    res
end

# a = [1, 5, 1, 4, 3, 4, 4, 3, 1, 4, 5, 3, 5]
# b = [9, 4, 0, 4, 0, 2, 1, 2, 1, 3, 2, 1, 1]
#
# ord = sortperm2(a, b, rev = true)
# hcat(a[ord], b[ord])
#
# ord = sortperm2(a, b)
# hcat(a[ord], b[ord])

function trapsum(y, x)
    d = diff(x)
    res = sum((d .* (y[2:end] + y[1:(end-1)]) ./ 2.0))
    res
end

# trapsum([1, 2, 3], [4, 6, 8])     # 8.0

function auc(x, y, reorder = false)
    direction = 1
    if reorder
        order = sortperm2(x, y)
        x, y = x[order], y[order]
    else
        dx = diff(x)
        if any(dx .<= 0)
            if all(dx .<= 0)
                direction = -1
            else
                error("Reordering is not turned on, and the x array is not increasing: $x")
            end
        end
    end
    area = direction * trapsum(y, x)
    return area
end

function roc_auc_score(y_true, y_score)
    if length(Set(y_true)) == 1
        warn("Only one class present in y_true.\n
              The AUC is not defined in that case; returning -Inf.")
        res = -Inf
    else
        fpr, tpr, thresholds = roc_curve(y_true, y_score)
        res = auc(fpr, tpr, true)
    end
    res
end

# roc_auc_score(y, y_score)


# [Julia package for function approximation](https://github.com/JuliaApproximation/ApproxFun.jl)

using ApproxFun

x = Fun(identity, 0..100)

g = sqrt(x)

g2 = g(2)

x = Fun(identity, 0.1 .. 1.)

w = Fun(identity,0.1 .. 1.)

T = Fun(identity,0.1 .. 1.)

ν = Fun(identity,0.1 .. 1.)

#kern = ν^4 * besselj(1, 2π*√(x*x+w*w/4) )^2

kern = Fun( x -> (ν^4 * besselj(1, 2π*√(x*x+w*w/4) )^2 / (x^2 + w^2/4) - T)^2 )


using JuMP
using DisjunctiveProgramming

m = Model()
@variable(m, -5 ≤ x ≤ 10)
@disjunction(
    m,
    begin
        exp(x) ≤ 2
        -3 ≤ x
    end,
    begin
        3 ≤ exp(x)
        5 ≤ x
    end,
    reformulation=:hull,
    name=:z
)
print(m)

# Feasibility
# Subject to
#  XOR(disj_z) : z[1] + z[2] == 1.0             <- XOR constraint
#  x_aggregation : x - x_z1 - x_z2 == 0.0       <- aggregation of disaggregated variables
#  x_z1_lb : -5 z[1] - x_z1 <= 0.0              <- lower-bound constraint on disaggregated variable x_z1 (x in 1st disjunct)
#  x_z1_ub : -10 z[1] + x_z1 <= 0.0             <- upper-bound constraint on disaggregated variable x_z1 (x in 1st disjunct)
#  x_z2_lb : -5 z[2] - x_z2 <= 0.0              <- lower-bound constraint on disaggregated variable x_z2 (x in 2nd disjunct)
#  x_z2_ub : -10 z[2] + x_z2 <= 0.0             <- upper-bound constraint on disaggregated variable x_z2 (x in 2nd disjunct)
#  x >= -5.0                                    <- lower-bound on x
#  x_z1 >= -5.0                                 <- lower-bound on x_z1 (disaggregated x in 1st disjunct)
#  x_z2 >= -5.0                                 <- lower-bound on x_z2 (disaggregated x in 2nd disjunct)
#  x <= 10.0                                    <- upper-bound on x
#  x_z1 <= 10.0                                 <- upper-bound on x_z1 (disaggregated x in 1st disjunct)
#  x_z2 <= 10.0                                 <- upper-bound on x_z2 (disaggregated x in 2nd disjunct)
#  z[1] >= 0.0                                  <- lower bound on binary
#  z[2] >= 0.0                                  <- lower bound on binary
#  z[1] <= 1.0                                  <- upper bound on binary
#  z[2] <= 1.0                                  <- upper bound on binary
#  z[1] binary                                  <- indicator variable (1st disjunct) is binary
#  z[2] binary                                  <- indicator variable (2nd disjunct) is binary
#  Perspective Functions:
#  (-1.0e-6 + -1.9999989999999999 * z[1]) + (1.0e-6 + 0.999999 * z[1]) * exp(x_z1 / (1.0e-6 + 0.999999 * z[1])) <= 0
#  (1.0000000000000002e-6 + 2.999999 * z[2]) + (-1.0e-6 + -0.999999 * z[2]) * exp(x_z2 / (1.0e-6 + 0.999999 * z[2])) <= 0
#  -1.0 * x_z1 + -3.0 * z[1] <= 0
#  -1.0 * x_z2 + 5.0 * z[2] <= 0


"""Struct that holds the quantised array as UInts with an additional
field for the min, max of the original range."""
struct LinQuantArray{T,N} <: AbstractArray{Unsigned,N}
    A::Array{T,N}       # array of UInts
    min::Float64        # offset min, max
    max::Float64
end

Base.size(QA::LinQuantArray) = size(QA.A)
Base.getindex(QA::LinQuantArray,i...) = getindex(QA.A,i...)
Base.eltype(Q::LinQuantArray{T,N}) where {T,N} = T

"""Quantise an array linearly into a LinQuantArray."""
function LinQuantization(::Type{T},A::AbstractArray) where {T<:Unsigned}
    all(isfinite.(A)) || throw(DomainError("Linear quantization only in (-∞,∞)"))

    Amin = Float64(minimum(A))              # minimum of value range
    Amax = Float64(maximum(A))              # maximum of value range

    if Amin == Amax
        Δ = 0.0                                 # set to zero for no range
    else
        Δ = (2^(sizeof(T)*8)-1)/(Amax-Amin)     # inverse spacing
    end

    Q = similar(A,T)                        # preallocate

    # map minimum to 0x0, maximum to 0xff...ff
    @inbounds for i in eachindex(Q)
        Q[i] = round((A[i]-Amin)*Δ)
    end

    return LinQuantArray{T,ndims(Q)}(Q,Amin,Amax)
end

# define 8, 16, 24 and 32 bit
LinQuant8Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt8,A)
LinQuant16Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt16,A)
LinQuant24Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt24,A)
LinQuant32Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt32,A)

"""De-quantise a LinQuantArray into floats."""
function Base.Array{T}(n::Integer,Q::LinQuantArray) where {T<:AbstractFloat}
    Qmin = Q.min                # min as Float64
    Qmax = Q.max                # max as Float64
    Δ = (Qmax-Qmin)/(2^n-1)     # linear spacing

    A = similar(Q,T)

    @inbounds for i in eachindex(A)
        # convert Q[i]::UInt to Float64 via *
        # then to T through =
        A[i] = Qmin + Q[i]*Δ
    end

    return A
end

# define default conversions for 8, 16, 24 and 32 bit
Base.Array{T}(Q::LinQuantArray{UInt8,N}) where {T,N} = Array{T}(8,Q)
Base.Array{T}(Q::LinQuantArray{UInt16,N}) where {T,N} = Array{T}(16,Q)
Base.Array{T}(Q::LinQuantArray{UInt24,N}) where {T,N} = Array{T}(24,Q)
Base.Array{T}(Q::LinQuantArray{UInt32,N}) where {T,N} = Array{T}(32,Q)

Base.Array(Q::LinQuantArray{UInt8,N}) where N = Array{Float32}(8,Q)
Base.Array(Q::LinQuantArray{UInt16,N}) where N = Array{Float32}(16,Q)
Base.Array(Q::LinQuantArray{UInt24,N}) where N = Array{Float32}(24,Q)
Base.Array(Q::LinQuantArray{UInt32,N}) where N = Array{Float64}(32,Q)

# one quantization per layer
"""Linear quantization independently for every element along dimension
dim in array A. Returns a Vector{LinQuantArray}."""
function LinQuantArray(::Type{TUInt},A::AbstractArray{T,N},dim::Int) where {TUInt,T,N}
    @assert dim <= N   "Can't quantize a $N-dimensional array in dim=$dim"
    n = size(A)[dim]
    L = Vector{LinQuantArray}(undef,n)
    t = [if j == dim 1 else Colon() end for j in 1:N]
    for i in 1:n
        t[dim] = i
        L[i] = LinQuantization(TUInt,A[t...])
    end
    return L
end

# for 8,16,24 and 32 bit
LinQuant8Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt8,A,dim)
LinQuant16Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt16,A,dim)
LinQuant24Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt24,A,dim)
LinQuant32Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt32,A,dim)

"""Undo the linear quantisation independently along one dimension, and returns
an array whereby the dimension always comes last. Hence, might be permuted compared
to the uncompressed array."""
function Base.Array{T}(L::Vector{LinQuantArray}) where T
    N = ndims(L[1])
    n = length(L)
    s = size(L[1])
    t = axes(L[1])
    A = Array{T,N+1}(undef,s...,length(L))
    for i in 1:n
        A[t...,i] = Array{T}(L[i])
    end
    return A
end


# Useful functions for checking proof details of the impossibility proof

"""
Builds the B matrix from the proof
"""
function Bmat(b1,b2)
    k = length(b1)
    @assert(isodd(k))
    r = round(Int64,(k+1)/2)
    b_half = [reverse(b2[r:end]); b1[r:end]]
    Db = diagm(b_half)
    Dk = zeros(k+1,k+1)
    for i = 1:r
        Dk[i,i] = k-i+1
        Dk[k-i+2,k-i+2] = k-i+1
    end
    row = collect(k:-1:0.0)
    r2 = collect(0.0:1:k)
    Top = repeat(row', r)
    Bottom = repeat(r2',r)
    R = [Top; Bottom]
    # @show size(Dk), size(Db), size(R), r, k

    B = Dk - Db*R

    return Dk, Db, R, B

end

"""
Computes primal variables corresponding to a complete hypergraph in the
case of using standard baseline scores for the proof.

These will often be optimal, but not if there is significant imbalance between
the two class sizes
"""
function Xstanard(k,n1,n2)
    n = n1+n2
    Xs = zeros(k+1)
    for i = 1:k+1
        t = i-1
        Xs[i] = binomial(n1,t)*binomial(n2,k-t)/binomial(n,k)
    end

    return Xs
end

"""
Optimal primal solution for but for asymptotic baselines
"""
function Xasym(k,alph)
    Xa = zeros(k+1)
    for i = 1:k+1
        t = i-1
        Xa[i] = alph^t*(1-alph)^(k-t)*binomial(k,t)
    end

    return Xa
end

"""
Use Gurobi to get the solution numerically
"""
function PrimalDualXY(gurobi_env,B,k)

    A = [B -ones(k+1,1)]
    p = [ones(1,k+1) 0]
    m = Model(optimizer_with_attributes(() -> Gurobi.Optimizer(gurobi_env),
        "OutputFlag" => 0, "FeasibilityTol"=>1e-9,"OptimalityTol"=>1e-9))
    @variable(m, x[1:k+2])
    @constraint(m, y, A*x .>= zeros(k+1,1))
    p = [ones(1,k+1) 0]
    @constraint(m, a, p*x .== 1)
    @constraint(m, z, x .>= zeros(k+2,1))

    @objective(m, Max, x[k+2])
    JuMP.optimize!(m)

    X = JuMP.value.(x)
    Y= JuMP.dual.(y)
    gam = X[end]

    return X[1:end-1], Y, gam
end

"""
Get the optimal dual variables for balanced classes.
"""
function BalancedDual(k)
    L = zeros(k+1)
    r = round(Int64,(k+1)/2)
    for i = 1:r
        L[i] = k/(k-i+1)
        L[k-i+2] = k/(k-i+1)
    end
    delta = sum(L)
    Ybal = L/delta
    return Ybal, delta
end

"""
Get optimal dual variables for any class size.
"""
function OptimalDual(b1,b2)
    k = length(b1)
    @assert(isodd(k))
    Ybal, delta = BalancedDual(k)
    r = round(Int64,(k+1)/2)
    s1 = 0
    s2 = 0
    for i = 1:r
        t = k-i+1
        s1 += b2[t]*(k/t - 1)
        s2 += (2-k/t)*b2[t]
    end

    y2k = 2/delta*s1/(1-s2)
    Ydual = zeros(2*r)
    Ydual[1] = y2k
    Ydual[k+1] = 2/delta - y2k

    for i = 2:r
        t = k-i+1
        Ydual[i] = 2/delta*(k/t-1) + (2-k/t)*y2k
        Ydual[k-i+2] = 2*k/(delta*t) - Ydual[i]
    end
    return Ydual
end

"""
For conceptual simplicity, store b and y values in a double-index
fashion
"""
function y_double(b1,b2)
    k = length(b1)
    b = zeros(2,k)
    y = zeros(2,k)
    @assert(isodd(k))

    Y = OptimalDual(b1,b2)
    r = round(Int64,(k+1)/2)

    for i = 1:r
        t = k-i+1
        y[2,t] = Y[i]
        y[1,t] = Y[k-i+2]
    end

    return y
end

# Convert from double indexing to linear
function tt(c,t)
    @assert(c == 1 || c == 2)

    if c == 2
        i = k-t+1
    else
        i = t+1
    end
    return i
end

# Now we have some code for the monotonic LP proof
function monotonic_lp_jump(b1,b2,k)
    """
    Maximize epsilon so that relative affinity scores satisfy
    h(t)/b(t) > epsilon + h(t-1)/b(t-1) for top r = (k+1)/2 scores
    """
    m = Model(optimizer_with_attributes(() -> Gurobi.Optimizer(gurobi_env), "OutputFlag" => 0))

    @variable(m, x[1:k+1])  # x[i+1] = proportion of hyperedges of type i
    @variable(m, y)         # variable we maximize to check beta value
    # @variable(m, d1)        # denominator for class 1 affinities
    # @variable(m, d2)        # denominator for class 2 affinities
    for i = 1:k+1
        @constraint(m, x[i] >= 0)
    end
    @constraint(m, 1 == sum(x[i] for i = 1:k+1))

    # Constraint for whether top J are a beta factor above baseline
    r = round(Int64,(k+1)/2)
    for t = r:k
        @constraint(m, (t)*x[t+1]/b1[t] - (t-1)*x[t]/b1[t-1]>= y)
        @constraint(m, (t)*x[k-t+1]/b2[t] - (t-1)*x[k-t+2]/b2[t-1] >= y)
    end

    @objective(m, Max, y)

    JuMP.optimize!(m)
    Y = JuMP.value.(y)
    X = JuMP.value.(x)

    return X,Y
end

function monotonic_LP_B(b1,b2,k)
    B = zeros(k+1,k+1)
    r = round(Int64,(k+1)/2)
    for i = 1:r
        t = k-i+1
        B[i,i] = t/(b2[t])
        B[i,i+1] = -(t-1)/(b2[t-1])
        B[k-i+2,k-i+2] = t/(b1[t])
        B[k-i+2,k-i+1] = -(t-1)/(b1[t-1])
    end
    return B
end

# Now we have some code for the monotonic LP proof
function monotonic_lp_jump_even_k(b1,b2,k)
    """
    Maximize epsilon so that relative affinity scores satisfy
    h(t)/b(t) > epsilon + h(t-1)/b(t-1) for top r = (k+1)/2 scores
    """
    m = Model(optimizer_with_attributes(() -> Gurobi.Optimizer(gurobi_env), "OutputFlag" => 0))

    @variable(m, x[1:k+1])  # x[i+1] = proportion of hyperedges of type i
    @variable(m, y)         # variable we maximize to check beta value
    for i = 1:k+1
        @constraint(m, x[i] >= 0)
    end
    @constraint(m, 1 == sum(x[i] for i = 1:k+1))

    # Constraint for whether top J are a beta factor above baseline
    half = round(Int64,k/2)
    for t = r:k
        @constraint(m, (t)*x[t+1]/b1[t] - (t-1)*x[t]/b1[t-1]>= y)
        @constraint(m, (t)*x[k-t+1]/b2[t] - (t-1)*x[k-t+2]/b2[t-1] >= y)
    end

    @objective(m, Max, y)

    JuMP.optimize!(m)
    Y = JuMP.value.(y)
    X = JuMP.value.(x)

    return X,Y
end


using UnicodePlots

scatterplot(randn(50), randn(50),
    title = "My Scatterplot")

myPlot = densityplot([1:100], randn(100), border=:dotted)

lineplot(rand(3), rand(3), title = "My Lineplot")


using DiffEqCallbacks
using OrdinaryDiffEq
using SimulationLogs
using Test

function lotka!(du, u, p, t)
    @log t
    @log x, y = u
    @log α, β, δ, γ = p

    @log total_population = x + y

    du[1] = α*x - β*x*y
    du[2] = δ*x*y - γ*y
end
function lotka(u, p, t)
    du = similar(u)
    lotka!(du, u, p , t)
    return du
end

lotka_sols = let
    p = [2/3, 4/3, 1, 1]
    u0 = [0.9, 0.9]
    tspan = (0.0, 10.0)

    callback = PresetTimeCallback(1:10, integrator->(integrator.p[end] += 0.01))
    prob_oop = ODEProblem(lotka, u0, tspan, p)
    prob_iip = ODEProblem(lotka!, u0, tspan, p)

    sol_oop = logged_solve(prob_oop, Tsit5())
    sol_iip = logged_solve(prob_iip, Tsit5())
    sol_oop_cb = logged_solve(prob_oop, Tsit5(); callback=callback)
    sol_iip_cb = logged_solve(prob_iip, Tsit5(); callback=callback)
    sol_oop_cbset = logged_solve(prob_oop, Tsit5(); callback=CallbackSet(callback))
    sol_iip_cbset = logged_solve(prob_iip, Tsit5(); callback=CallbackSet(callback))

    sols_iip = [sol_iip, sol_iip_cb, sol_iip_cb]
    sols_oop = [sol_oop, sol_oop_cb, sol_oop_cb]

    (iip=sols_iip, oop=sols_oop)
end

@testset "Lotka example" begin
    test_all(f, i=:) = all(f, lotka_sols.iip[i]) && all(f, lotka_sols.oop[i])
    test_cb(f) = test_all(f, 2:3)

    @testset "Logged values" begin
        @test test_all(sol->sol.log.γ[1]==1)
        @test test_all(sol->sol.log.x==sol[1,:])
        @test test_all(sol->sol.log.y==sol[2,:])
        @test test_all(sol->sol.log.t==sol.t)

        @test test_all(sol->all(sol.log.α .== sol.prob.p[1]))
        @test test_cb(sol->sol.log.γ[end] == 1.1)
    end

    this_sol = lotka_sols.iip[1]

    @testset "Indexing" begin
        @test this_sol[1] == this_sol.u[1]
        @test this_sol[:] == this_sol.u
        @test this_sol[:, :] == reduce(hcat, this_sol.u)
    end

    @testset "Interpolation" begin
        @test this_sol(10) ≈ this_sol.u[end]
    end

    @testset "Properties" begin
        @test this_sol.log isa SimulationLog
        @test this_sol.alg == this_sol.sol.alg
        @test Set(propertynames(this_sol)) == Set((:log, propertynames(this_sol.sol)...))
    end

    @testset "Printing" begin
        sol_string = @capture_out(print(this_sol))
        sol_sol_string = @capture_out(print(this_sol.sol))
        @test contains(sol_string, sol_sol_string)
    end
end


function missing_to_zeros!(X)
    for i in eachindex(X)
        if ismissing(X[i])
            X[i] = 0
        end
    end
    X
end

my_mean(μ, n, x)          = μ + (x - μ) / (n + 1), n + 1
my_mean(μ, n, x::Missing) = μ, n
function agg_mean(x)
    μ, n = zero(eltype(x)), 0
    @inbounds for xi in x
        μ, n = my_mean(μ, n, xi)
    end
    return μ
end

function missing_to_mean(X :: Matrix{Union{T, Missing}}) where T

    # If there are missing values here, then there are too many!
    # this causes quite a lot of allocations, but is still pretty fast
    m = [mean(skipmissing(view(X, :, i))) for i in 1:size(X, 2)]
    # much less allocations, but slow:
    # m = [agg_mean(view(X, :, i)) for i in 1:size(X, 2)]
    # this one is way to slow:
    # m = mapslices(x -> mean(skipmissing(x)), X, dims = 1)
    m2 = mean(m)

    X2 = Array{T}(undef, size(X))
    @inbounds for i in eachindex(X, X2)
        X2[i] = ismissing(X[i]) ? m2 : X[i]
    end

    return X2, m, m2
end

# eachslice requires Julia 1.1!
function missing_to_mean_slice(X :: Matrix{Union{T, Missing}}, dims = 1) where T

    X2 = Matrix{T}(undef, size(X))
    for (s1, s2) in zip(eachslice(X, dims = dims), eachslice(X, dims = dims))
        m = s1 |> skipmissing |> mean
        for i in eachindex(s1, s2)
            @inbounds s1[i] = ismissing(s2[i]) ? m : s2[i]
        end
    end

    X2
end

function remove_random!(
    x::Array{Union{T, Missing}},
    p = 0.05
) where {T <: AbstractFloat}
    for i in eachindex(x)
        if rand(T) < T(p)
            x[i] = missing
        end
    end
    x
end


#=

  xkcd's knapsack/subset-sum problem in Julia + ConstraintSolver.jl

  http://xkcd.com/287/

  Some amount (or none) of each dish should be ordered to give a total of exact 15.05

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
using Distributions
include("constraints_utils.jl")

function xkcd(prices=[215, 275, 335, 355, 420, 580],total=1505,max_val=10;all_solutions=false)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # "simplify"=>false,
                                                            "simplify"=>true,

                                                            "time_limit"=>14,
                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                        ))

    println("prices:$prices")
    println("total:$total   (max_val:$max_val)")
    len = length(prices)

    @variable(model, 0 <= x[1:len] <= max_val, Int)

    # @constraint(model, total == sum([x[i]*prices[i] for i in 1:len]))
    scalar_product(model,x,prices,total)
    # @constraint(model, total == sum(x.*prices))

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("\nnum_sols:$num_sols\n")
        # num_sols2 = JuMP.result_count(model) # same result
        # println("num_sols2:$num_sols2\n")

        # The solutions are in reverse order, i.e. the optimal is the first.
        for sol in 1:num_sols
            println("solution #$sol")
            # This works:
            xx = convert.(Integer,JuMP.value.(x,result=sol))
            println("x:$xx")
            println("used: ")
            println([(i,xx[i]) for i in 1:length(xx) if xx[i] > 0  ])
        end
    end
end

xkcd(;all_solutions=true)

println("\nA random instance")

n = 120
prices = rand(DiscreteUniform(100,1000),n)
s = sum(prices)
total = rand(DiscreteUniform(0,s))
max_val = rand(DiscreteUniform(1,100))
xkcd(prices,total,max_val)


using Utils

const Num = Int

function parse_file(f)
    collect(eachline(f))
end

function load()
    open(joinpath(@__DIR__, "18.txt"), "r") do f
        parse_file(f)
    end
end

isnodigit = !isdigit

is_paren(c) = c == '(' || c == ')'

function find_matching(s, from)
    @assert s[from] == '('
    level = 1
    while true
        cand = findnext(is_paren, s, from + 1)
        level -= s[cand] == ')' ? 1 : -1
        level == 0 && return cand
        from = cand
    end
end

function tokenize_l2r(expr)
    re = []
    cursor = 1
    while cursor <= length(expr)
        if expr[cursor] == '('
            close = find_matching(expr, cursor)
            push!(re, tokenize_l2r(expr[cursor+1:close-1]))
            cursor = close + 1
        elseif expr[cursor] == '+'
            push!(re, :plus)
            cursor += 1
        elseif expr[cursor] == '*'
            push!(re, :times)
            cursor += 1
        elseif isdigit(expr[cursor])
            after = something(findnext(isnodigit, expr, cursor), length(expr) + 1)
            push!(re, parse(Int, expr[cursor:after-1]))
            cursor = after
        else
            cursor += 1
        end
    end
    re
end

evaluate_l2r(term::Int) = term

function evaluate_l2r(terms::Array)
    acc = evaluate_l2r(terms[1])
    i = 2
    while i <= length(terms)
        if terms[i] == :plus
            acc += evaluate_l2r(terms[i + 1])
        elseif terms[i] == :times
            acc *= evaluate_l2r(terms[i + 1])
        end
        i += 2
    end
    acc
end

function reduce_stack(ops, operands)
    while true
        len = length(ops)
        if len == 0
            @assert length(operands) == 1
            return operands[1]
        elseif len >= 1 && ops[end] == :plus
            push!(operands, pop!(operands) + pop!(operands))
            pop!(ops)
        elseif len >= 2 && ops[end] == :times == ops[end - 1]
            keep = pop!(operands)
            push!(operands, pop!(operands) * pop!(operands))
            push!(operands, keep)
            pop!(ops)
        else
            push!(operands, pop!(operands) * pop!(operands))
            @assert pop!(ops) == :times
        end
    end
end

function evaluate(terms::Array)
    stack = Num[]
    for t in terms
        if isa(t, Num)
            push!(stack, t)
        else
            a = pop!(stack)
            b = pop!(stack)
            push!(stack, t == :+ ? a + b : a * b)
        end
    end

    only(stack)
end

function tokenize(expr)
    expr = replace(expr, " " => "")
    operands = Union{Num,Symbol}[]
    ops = Symbol[]

    cursor::Num = 1
    while cursor <= length(expr)
        if expr[cursor] == '('
            push!(ops, :open)
        elseif expr[cursor] == ')'
            while true
                top = pop!(ops)
                top == :open && break
                push!(operands, top)
            end
        elseif isdigit(expr[cursor])
            after = something(findnext(isnodigit, expr, cursor), length(expr) + 1)
            push!(operands, parse(Num, expr[cursor:after-1]))
            cursor = after
            continue
        elseif expr[cursor] == '+'
            push!(ops, :+)
        elseif expr[cursor] == '*'
            while !isempty(ops) && ops[end] == :+
                push!(operands, pop!(ops))
            end
            push!(ops, :*)
        end
        cursor += 1
    end

    while !isempty(ops)
        push!(operands, pop!(ops))
    end

    operands
end

function ex1(eqs)
    sum(eq -> tokenize_l2r(eq) |> evaluate_l2r, eqs)
end

function ex2(eqs)
    sum(eq -> tokenize(eq) |> evaluate, eqs)
end

eqs = load()

println("Sum of equations: ", ex1(eqs))
println("With precedence: ", ex2(eqs))

using Test

@testset "Operation Order" begin
@test find_matching("(3+5)", 1) == 5
@test find_matching("((3+5))", 1) == 7
@test find_matching("((3+5))", 1) == 7
@test find_matching("((3+5))", 2) == 6

@testset "example 1" begin
        @test ex1(parse_file(IOBuffer("1 + 2 * 3 + 4 * 5 + 6"))) == 71
        @test ex1(parse_file(IOBuffer("1 + (2 * 3) + (4 * (5 + 6))"))) == 51
        @test ex1(parse_file(IOBuffer("2 * 3 + (4 * 5)"))) == 26
        @test ex1(parse_file(IOBuffer("5 + (8 * 3 + 9 + 3 * 4 * 3)"))) == 437
        @test ex1(parse_file(IOBuffer("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))"))) == 12240
        @test ex1(parse_file(IOBuffer("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"))) == 13632
    end

    @testset "example 2" begin
        @test ex2(parse_file(IOBuffer("1 + (2 * 3) + (4 * (5 + 6))"))) == 51
        @test ex2(parse_file(IOBuffer("2 * 3 + (4 * 5)"))) == 46
        @test ex2(parse_file(IOBuffer("5 + (8 * 3 + 9 + 3 * 4 * 3)"))) == 1445
        @test ex2(parse_file(IOBuffer("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))"))) == 669060
        @test ex2(parse_file(IOBuffer("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"))) == 23340
    end

    @testset "results" begin
        @test ex1(eqs) == 12956356593940
        @test ex2(eqs) == 94240043727614
    end
end


 #=

Implementations of pcg that use BLAS when they can.
Computes error in matrix norm, if correct left hand side is present.

Started by Dan Spielman.
Contributors: Serban Stan

=#

"""
`pcg(mat, b, pre; tol, maxits, maxtime, verbose)` solves a symmetric linear system using preconditioner `pre`.
`pre` can be a function or a matrix.  If a matrix, a function to solve it is created with cholFact.
`tol` is set to 1e-6 by default,
`maxits` defaults to Inf
`maxtime` defaults to Inf.  It measures seconds.
`verbose` defaults to false
`lhs` defaults to the all zeros vector
"""

function pcgMatNorm{Tv}(mat, b::Array{Tv,1}, pre::Function, lhs::Array{Tv,1}; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    return pcgBLASMatNorm(mat, b, pre, lhs, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

#====================================
   PCG Code
=====================================#

# uses BLAS.  As fast as Matlab's pcg.
# set to use similar paramaters
function pcgBLASMatNorm{Tval}(mat, b::Array{Tval,1}, pre, lhs::Array{Tval,1}; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)

    n = size(mat,2)

    x = zeros(Tval,n)

    r = copy(b)
    z = pre(r)
    p = copy(z)

    rho = dot(r, z) # BLAS.dot does not seem faster

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        q = mat*p

        al = rho/dot(p, q)

        BLAS.axpy!(al,p,x)  # x = x + al * p
        BLAS.axpy!(-al,q,r)  # r -= al*q

        # This will take some time. Can we do something faster?
        if sqrt((lhs - x)' * mat * (lhs - x))[1] < tol
            break
        end

        # here is the top of the code in numerical templates

        z = pre(r)

        oldrho = rho
        rho = dot(z, r)

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho

        p = z + beta*p
        # BLAS.scal!(n,beta,p,1) # p *= beta
        # BLAS.axpy!(1.0,z,p) # p += z

        if (time() - t1) > maxtime
            if verbose
                println("PCG stopped at maxtime.")
            end
            break
        end

      end

    if verbose
        println("PCG stopped after: ", itcnt, " iterations with relative error ", sqrt((lhs - x)' * mat * (lhs - x))[1], ".")
    end

    return x
end


# Main functions used for work on MCL with Julia
# Berkeley 2015 - Maxime Legrand

function bern(k,l,p)
  A = rand(k,l)
  B = p*ones(k,l)
  C = A .<= B
  D = convert(Array{Float64,2}, C)
  return D
end

"Sample for 2-SBM"
function mcl_sample_dumb(a,b,p1,p2)
  n = a+b
  G = zeros(n,n)
  G[1:a,1:a] = bern(a,a,p1)
  G[a+1:n,a+1:n] = bern(b,b,p1)
  G[1:a,a+1:n] = bern(a,b,p2)
  G[a+1:n,1:a] = bern(b,a,p2)
  return G
end

"""Sample for SBM
Input : l, vector of the sizes of the wanted clusters
    p, vector of their intra-cluster connection probability
    p0, inter-cluster connection probability
Output : corresponding SBM sample matrix"""
function mcl_sample_sbm(l,p,p0)
  s = size(l,1)
  n = sum(l)
  G = zeros(n,n)
  l0 = cumsum([0; l])
  for i in 1:s
    G[l0[i]+1:l0[i+1],l0[i]+1:l0[i+1]] = bern(l[i],l[i],p[i])
    G[l0[i]+1:l0[i+1],l0[i+1]+1:n] = bern(l[i,1],n-l0[i+1],p0)
    G[l0[i+1]+1:n,l0[i]+1:l0[i+1]] = bern(n-l0[i+1],l[i],p0)
  end
  return G
end

function mcl_norm(A)
  n = length(A)
  As = sum(A,2)
  B = A ./ As
  return B
end

function mcl_expand(A,e)
  B = A^e
  return B
end

function mcl_inflate(A,r)
  B = A .^ r
  return B
end

"""MCL algorithm
Input : A, adjacency matrix
    e, expansion parameter
    r, inflation parameter
    p, number of iterations
Output : corresponding result of MCL algorithm"""
function mcl(A;e=2,r=3,p=20)
  B = mcl_norm(A)
  for i=1:p
    B = mcl_norm(mcl_inflate(mcl_expand(B,e),r))
  end
  return B
end

"""First non-zero
Input : v, vector
Output : first non-zero element of v"""
function fst_non_zero(v)
  n = length(v)
  r = n+1
  for i = 1:n
    if (v[i] != 0) && (i < r)
      r = i
    end
  end
  if r == n+1
    r = 0
  end
  return r
end

"""Element finding in a vector of vectors
Input : k, element to find
    v, vector of vectors
Output : index of the first vector containing k if found, 0 otherwise"""
function vvfind(k,v)
  for i in 1:length(v)
    if length(find(x -> x==k, v[i])) != 0
      return i
    end
  end
  return 0
end

"Returns the first element of each vector of a vectors vector"
function vvfirst(v)
  K = find(x -> length(x) != 0, v)
  if length(K) == 0
    return Any[]
  else
    r = [v[K[1]][1]]
    for k in K[2:end]
      push!(r, v[k][1])
    end
    return r
  end
end

"""Return a vector containing clusters
Input : G, typical output of MCL algorithm
    sort_v, set to true to sort vertices inside clusters
    sort_c, set to true to sort clusters by increasing size
    rev, set to true along with sort_c to sort by decreasing size"""
function mcl_clust(G; sort_v=false, sort_c=false, rev=false)
  # Main part
  n = size(G,1)
  r = Vector[]
  for i = 1:n
    k = fst_non_zero(G[i,:])
    s = vvfind(k,r)
    if s == 0
      if i == k
        push!(r, [i])
      else
        push!(r, [k,i])
      end
    else
      if i != k
        push!(r[s],i)
      end
    end
  end
  # Vertices sorting part
  if sort_v
    for k in 1:size(r,1)
      sort!(r[k])
    end
  end
  # Clusters sorting part
  if sort_c
    p = sortperm(map(length, r), rev=rev)
    r = r[p]
  end
  return r
end

"Return the number of spotted clusters"
function mcl_card(G)
  return length(mcl_clust(G))
end

"Heuristic approximation of the probability of recognizing two clusters generated with mcl_sample_dumb, of size a and b, with T iterations"
function mcl_prob_dumb(a,b,p1,p2,T;e=2,r=3,p=20)
  if a == 0 || b == 0
    return 0.
  else
    result = 0
    for i in 1:T
      G = mcl_sample_dumb(a,b,p1,p2)
      if mcl_card(mcl(G,e,r,p)) == 2
        result += 1
      end
    end
  return result/T
  end
end


using LinearAlgebra, Distributions
using Statistics: mean
export predictability

"""
    predictability(ds::DynamicalSystem; kwargs...) -> chaos_type, ν, C

Determine whether `ds` displays strongly chaotic, partially-predictable chaotic
or regular behaviour, using the method by Wernecke et al. described in[^Wernecke2017].

Return the type of the behavior, the cross-distance scaling coefficient `ν`
and the correlation coefficient `C`.
Typical values for `ν`, `C` and `chaos_type` are given in Table 2 of[^Wernecke2017]:

| `chaos_type` | `ν` | `C` |
|--------------|-----|-----|
|    `:SC`     |  0  |  0  |
|    `:PPC`    |  0  |  1  |
|    `:REG`    |  1  |  1  |

## Keyword Arguments
* `Ttr = 200` : Extra "transient" time to evolve the system before sampling from
   the trajectory. Should be `Int` for discrete systems.
* `T_sample = 1e4` : Time to evolve the system for taking samples. Should be
  `Int` for discrete systems.
* `n_samples = 500` : Number of samples to take for use in calculating statistics.
* `λ_max = lyapunov(ds, 5000)` : Value to use for largest Lyapunov exponent
  for finding the Lyapunov prediction time. If it is less than zero a regular
  result is returned immediatelly.
* `d_tol = 1e-3` : tolerance distance to use for calculating Lyapunov prediction time.
* `T_multiplier = 10` : Multiplier from the Lyapunov prediction time to the evaluation time.
* `T_max = Inf` : Maximum time at which to evaluate trajectory distance. If the internally
   computed evaluation time is larger than `T_max`, stop at `T_max` instead.
* `δ_range = 10.0 .^ (-9:-6)` : Range of initial condition perturbation distances
   to use to determine scaling `ν`.
* `diffeq...` : Keyword arguments propagated into `init` of DifferentialEquations.jl.
  See [`trajectory`](@ref) for examples. Only valid for continuous systems.

## Description
Samples points from a trajectory of the system to be used as initial conditions. Each of
these initial conditions is randomly perturbed by a distance `δ`, and the trajectories for
both the original and perturbed initial conditions are computed to the 'evaluation time'
`T`.

The average (over the samples) distance and cross-correlation coefficient
of the state at time `T` is
computed. This is repeated for a range of `δ` (defined by `δ_range`), and linear
regression is used to determine how the distance and cross-correlation scale with `δ`,
allowing for identification of chaos type.

The evaluation time `T` is calculated as `T = T_multiplier*Tλ`, where the Lyapunov
prediction time `Tλ = log(d_tol/δ)/λ_max`. This may be very large if the `λ_max` is small,
e.g. when the system is regular, so this internally computed time `T` can be overridden by
a smaller `T_max` set by the user.

## Performance Notes
For continuous systems, it is likely that the `maxiters` used by the integrators needs to
be increased, e.g. to 1e9. This is part of the `diffeq` kwargs.
In addition, be aware that this function does a *lot* of internal computations.
It is operating in a different speed than e.g. [`lyapunov`](@ref).

[^Wernecke2017]: Wernecke, H., Sándor, B. & Gros, C. *How to test for partially predictable chaos*. [Scientific Reports **7**, (2017)](https://www.nature.com/articles/s41598-017-01083-x).
"""
function predictability(ds::DynamicalSystem;
                        Ttr::Real = 200,
                        T_sample::Real = 1e4,
                        n_samples::Integer = 500,
                        λ_max::Real = lyapunov(ds, 5000),
                        d_tol::Real = 1e-3,
                        T_multiplier::Real = 10,
                        T_max::Real = Inf,
                        δ_range::AbstractArray = 10.0 .^ (-9:-6),
                        diffeq...
                        )

    λ_max < 0 && return :REG, 1.0, 1.0
    # Internal Constants
    ν_threshold = 0.5
    C_threshold = 0.5

    # Sample points from a single trajectory of the system
    samples = sample_trajectory(ds, Ttr, T_sample, n_samples; diffeq...)

    # Calculate the mean position and variance of the trajectory. ([1] pg. 5)
    # Using samples 'Monte Carlo' approach instead of direct integration
    μ = mean(samples)
    s² = mean(map(x->(x-μ)⋅(x-μ), samples))

    # Calculate cross-distance scaling and correlation scaling
    distances = Float64[] # Mean distances at time T for different δ
    correlations = Float64[] # Cross-correlation at time T for different δ
    p_integ = parallel_integrator(ds, samples[1:2]; diffeq...)
    for δ in δ_range
        # TODO: some kind of warning should be thrown for very large Tλ
        Tλ = log(d_tol/δ)/λ_max
        T = min(T_multiplier * Tλ, T_max)
        Σd = 0
        Σd² = 0
        for u in samples
            # Sample perturbation from surface of sphere
            n = rand(Normal(), size(u))
            n /= norm(n)
            û = u + δ*n
            # Update integrator with new initial conditions
            reinit!(p_integ, [u, û])
            # Simulate trajectory until T
            while p_integ.t < T
                step!(p_integ)
            end
            # Accumulate distance and square-distance
            d = norm(p_integ.u[1]-p_integ.u[2], 2)
            Σd  += d
            Σd² += d^2
        end
        # Calculate mean distance and square-distance
        d = Σd/length(samples)
        D = Σd²/length(samples)
        # Convert mean square-distance into cross-correlation
        C = 1 - D/2s²
        # Store mean distance and cross-correlation
        push!(distances, d)
        push!(correlations, C)
    end

    # Perform regression to check cross-distance scaling
    ν = slope(log.(δ_range), log.(distances))
    C = mean(correlations)

    # Determine chaotic nature of the system
    if ν > ν_threshold && C > C_threshold
        chaos_type = :REG
    elseif ν <= ν_threshold && C > C_threshold
        chaos_type = :PPC
    elseif ν <= ν_threshold && C ≤ C_threshold
        chaos_type = :SC
    else
        # Covers the case when ν > ν_threshold but C <= C_threshold
        chaos_type = :INDETERMINATE
    end

    return chaos_type, ν, C
end

function sample_trajectory(ds::ContinuousDynamicalSystem,
                           Ttr::Real, T_sample::Real,
                           n_samples::Real;
                           diffeq...)
    # Samples *approximately* `n_samples` points.
    β = T_sample/n_samples
    D_sample = Exponential(β)
    sample_trajectory(ds, Ttr, T_sample, D_sample; diffeq...)
end

function sample_trajectory(ds::DiscreteDynamicalSystem,
                           Ttr::Real, T_sample::Real,
                           n_samples::Real;
                           diffeq...)
    @assert n_samples < T_sample "discrete systems must satisfy n_samples < T_sample"
    # Samples *approximately* `n_samples` points.
    p = n_samples/T_sample
    D_sample = Geometric(p)
    sample_trajectory(ds, Ttr, T_sample, D_sample; diffeq...)
end

function sample_trajectory(ds::DynamicalSystem,
                           Ttr::Real, T_sample::Real,
                           D_sample::UnivariateDistribution;
                           diffeq...)
    # Simulate initial transient
    integ = integrator(ds; diffeq...)
    while integ.t < Ttr
        step!(integ)
    end

    # Time to the next sample is sampled from the distribution D_sample
    # e.g. Continuous systems: D_sample is Exponential distribution
    samples = typeof(integ.u)[]
    while integ.t < Ttr + T_sample
        step!(integ, rand(D_sample))
        push!(samples, integ.u)
    end
    samples
end


struct BalancedTernary <: Signed
    digits::Vector{Int8}
end
BalancedTernary() = zero(BalancedTernary)
BalancedTernary(n) = convert(BalancedTernary, n)

const sgn2chr = Dict{Int8,Char}(-1 => '-', 0 => '0', +1 => '+')
Base.show(io::IO, bt::BalancedTernary) = print(io, join(sgn2chr[x] for x in reverse(bt.digits)))
Base.copy(bt::BalancedTernary) = BalancedTernary(copy(bt.digits))
Base.zero(::Type{BalancedTernary}) = BalancedTernary(Int8[0])
Base.iszero(bt::BalancedTernary) = bt.digits == Int8[0]
Base.convert(::Type{T}, bt::BalancedTernary) where T<:Number = sum(3 ^ T(ex - 1) * s for (ex, s) in enumerate(bt.digits))
function Base.convert(::Type{BalancedTernary}, n::Signed)
    r = BalancedTernary(Int8[])
    if iszero(n) push!(r.digits, 0) end
    while n != 0
        if mod(n, 3) == 0
            push!(r.digits, 0)
            n = fld(n, 3)
        elseif mod(n, 3) == 1
            push!(r.digits, 1)
            n = fld(n, 3)
        else
            push!(r.digits, -1)
            n = fld(n + 1, 3)
        end
    end
    return r
end
const chr2sgn = Dict{Char,Int8}('-' => -1, '0' => 0, '+' => 1)
function Base.convert(::Type{BalancedTernary}, s::AbstractString)
    return BalancedTernary(getindex.(chr2sgn, collect(reverse(s))))
end

macro bt_str(s)
    convert(BalancedTernary, s)
end

const table = NTuple{2,Int8}[(0, -1), (1, -1), (-1, 0), (0, 0), (1, 0), (-1, 1), (0, 1)]
function _add(a::Vector{Int8}, b::Vector{Int8}, c::Int8=Int8(0))
    if isempty(a) || isempty(b)
        if c == 0 return isempty(a) ? b : a end
        return _add([c], isempty(a) ? b : a)
    else
        d, c = table[4 + (isempty(a) ? 0 : a[1]) + (isempty(b) ? 0 : b[1]) + c]
        r = _add(a[2:end], b[2:end], c)
        if !isempty(r) || d != 0
            return unshift!(r, d)
        else
            return r
        end
    end
end
function Base.:+(a::BalancedTernary, b::BalancedTernary)
    v = _add(a.digits, b.digits)
    return isempty(v) ? BalancedTernary(0) : BalancedTernary(v)
end
Base.:-(bt::BalancedTernary) = BalancedTernary(-bt.digits)
Base.:-(a::BalancedTernary, b::BalancedTernary) = a + (-b)
function _mul(a::Vector{Int8}, b::Vector{Int8})
    if isempty(a) || isempty(b)
        return Int8[]
    else
        if a[1] == -1 x = (-BalancedTernary(b)).digits
        elseif a[1] == 0 x = Int8[]
        elseif a[1] == 1 x = b end
        y = append!(Int8[0], _mul(a[2:end], b))
        return _add(x, y)
    end
end
function Base.:*(a::BalancedTernary, b::BalancedTernary)
    v = _mul(a.digits, b.digits)
    return isempty(v) ? BalancedTernary(0) : BalancedTernary(v)
end

a = bt"+-0++0+"
println("a: $(Int(a)), $a")
b = BalancedTernary(-436)
println("b: $(Int(b)), $b")
c = BalancedTernary("+-++-")
println("c: $(Int(c)), $c")
r = a * (b - c)
println("a * (b - c): $(Int(r)), $r")

@assert Int(r) == Int(a) * (Int(b) - Int(c))


SLS2 = RheoModelClass(
        # Model name
        name="SLS2",
        # Model parameters,
        p = [:G₀, :G₁, :η₁, :G₂, :η₂],
        # Relaxation modulus
        G = quote
              G₀ + G₁*exp(-t*G₁/η₁) + G₂*exp(-t*G₂/η₂)
            end,
        # Creep modulus
        J = quote
              tau1 = G₁/η₁
              tau2 = G₂/η₂
              Jbar(s) = (1/s^2)*(G₀/s+G₁*1/(s+tau1)+G₂*1/(s+tau2))
              invLaplace(s -> Jbar(s), t)
            end,
        # Network
        info= "
                                       ___
                 ___╱╲  ╱╲  ╱╲  ________| |_____
                 |    ╲╱  ╲╱  ╲╱       _|_|     |
                 |      G₂              η₂      |
                 |                     ___      |
              ___|__╱╲  ╱╲  ╱╲  ________| |_____|___
                 |    ╲╱  ╲╱  ╲╱       _|_|     |
                 |      G₁              η₁      |
                 |                              |
                 |__________╱╲  ╱╲  ╱╲  ________|
                              ╲╱  ╲╱  ╲╱
                                G₀
               "
        )

#
#  Generate array of parameter symbols and expression of the relaxation modulus
#  for the generalised maxwell model
#

function G_GSLS(n)

    params=[:k0]
    G_terms=[:(k0)]

    for i=1:n
        k=Symbol("k_"*string(i))
        η=Symbol("η_"*string(i))
        e=:($k * exp.( (- $k/$η) * t))

        params=vcat(params, [k,η])
        G_terms=vcat(G_terms,e)
    end
    G_Expr=Expr(:call,:+,G_terms...)

    return (params, G_Expr)
end

# (p,e)=G_GSLS(3)

# cd("/home/alexandre/.julia/dev/RHEOS/src/"); include("RHEOS.jl"); using Main.RHEOS; m=RheoModel(SLS2,(G₀=1., G₁=2., η₁=3., G₂=4., η₂=5.))


#=
# Builds a tiled ordering of a 2D or 3D grid
# Tiles are tiledim in size (except edge tiles which could be smaller)
# Grid has griddim size
# Returns the global order in which objects shall be indexed.
# example: 2D, griddim=(4,4), tiledim=(3,3)
#     13  14  15  16
#      7   8   9  12
#      4   5   6  11
#      1   2   3  10
# Invert=true will give  [1, 2, 3, 5, 6, 7 ...]   (encode?)
# Invert=false will give [1, 2, 3, 10, 4, 5, ...] (decode?)
=#
export reorder_grid_tiled, get_tiled_order

function tiled_order_2d(griddim, tiledim, invert = true)
    # griddim and tiledim are tuples of grid dimensions
    gridx = griddim[1];
    gridy = griddim[2];
    tilex = tiledim[1];
    tiley = tiledim[2];

    nnodes = gridx*gridy;
    tiled = zeros(Int64, nnodes); # The ordering

    # Count full tiles and renaining edge nodes.
    fullx = Int(floor(gridx / tilex));
    partialx = gridx - fullx*tilex;
    fully = Int(floor(gridy / tiley));
    partialy = gridy - fully*tiley;

    # Fill in ordering
    for ni=1:nnodes
        # Find destination tile index
        tx = Int(floor(mod(ni-1, gridx)/tilex)) + 1;
        ty = Int(floor((ni-1)/(gridx*tiley))) + 1;
    end

    tind = 0;
    for j = 1:(fully+1)
        ytill = j<=fully ? tiley : partialy;
        for i = 1:(fullx+1)
            xtill = i<=fullx ? tilex : partialx;

            # Add that tile's nodes one at a time
            for tj=1:ytill
                for ti=1:xtill
                    gind = ti + (i-1)*tilex + gridx*(tj + (j-1)*tiley - 1);
                    tind = tind + 1;
                    if invert
                        tiled[gind] = tind;
                    else
                        tiled[tind] = gind;
                    end
                end
            end

        end
    end

    return tiled;
end

function tiled_order_3d(griddim, tiledim, invert = true, tileorder=[])
    # griddim and tiledim are tuples of grid dimensions
    gridx = griddim[1];
    gridy = griddim[2];
    gridz = griddim[3];
    tilex = tiledim[1];
    tiley = tiledim[2];
    tilez = tiledim[3];

    nnodes = gridx*gridy*gridz;
    tiled = zeros(Int64, nnodes); # The ordering

    # Count full tiles and remaining edge nodes.
    fullx = Int(floor(gridx / tilex));
    partialx = gridx - fullx*tilex;
    fully = Int(floor(gridy / tiley));
    partialy = gridy - fully*tiley;
    fullz = Int(floor(gridz / tilez));
    partialz = gridz - fullz*tilez;

    # Fill in ordering
    if length(tileorder) == 0
        Ntiles = (fullx+1)*(fully+1)*(fullz+1);
        tileorder = 1:Ntiles;
    end

    tind = 0;
    for i=1:length(tileorder)
        # Get the tile coordinates in (x,y,z)
        tileind = tileorder[i];
        tx = Int(mod(tileind-1, fullx+1)) + 1;
        ty = Int(floor(mod(tileind-1, (fullx+1)*(fully+1))/(fullx+1))) + 1;
        tz = Int(floor((tileind-1)/((fullx+1)*(fully+1)))) + 1;
        xtill = tx<=fullx ? tilex : partialx;
        ytill = ty<=fully ? tiley : partialy;
        ztill = tz<=fullz ? tilez : partialz;

        # Add that tile's nodes one at a time
        for tk=1:ztill
            for tj=1:ytill
                for ti=1:xtill
                    gind = ti + (tx-1)*tilex + gridx*((tj + (ty-1)*tiley - 1) + gridy*(tk + (tz-1)*tilez - 1));
                    tind = tind + 1;
                    if invert
                        tiled[gind] = tind;
                    else
                        tiled[tind] = gind;
                    end
                end
            end
        end
    end

    return tiled;
end

function reorder_grid_tiled(grid, griddim, tiledim)
    dim = size(grid.allnodes,1);
    nel = size(grid.loc2glb,2);
    nfc = size(grid.face2glb,2);

    tiled = get_tiled_order(dim, griddim, tiledim, true);

    newnodes = zeros(size(grid.allnodes));
    newbdry = copy(grid.bdry);
    newloc2glb = copy(grid.loc2glb);
    newglbvertex = copy(grid.glbvertex);
    newface2glb = copy(grid.face2glb);
    newfaceVertex2glb = copy(grid.faceVertex2glb);

    for mi=1:length(tiled)
        for d=1:dim
            newnodes[d, tiled[mi]] = grid.allnodes[d,mi];
        end
    end

    for bid=1:length(newbdry)
        for i=1:length(newbdry[bid])
            newbdry[bid][i] = tiled[grid.bdry[bid][i]];
        end
    end

    for ei=1:nel
        for ni=1:size(newloc2glb,1)
            newloc2glb[ni,ei] = tiled[grid.loc2glb[ni,ei]];
        end
        for ni=1:size(newglbvertex,1)
            newglbvertex[ni,ei] = tiled[grid.glbvertex[ni,ei]];
        end
    end

    for fi=1:nfc
        for ni=1:size(newface2glb,1)
            newface2glb[ni,fi] = tiled[grid.face2glb[ni,fi]];
        end
        for ni=1:size(newfaceVertex2glb,1)
            newfaceVertex2glb[ni,fi] = tiled[grid.faceVertex2glb[ni,fi]];
        end
    end

    return Femshop.Grid(newnodes, newbdry, grid.bdryface, grid.bdrynorm, grid.bids, newloc2glb, newglbvertex, newface2glb, newfaceVertex2glb);
end

function get_tiled_order(dim, griddim, tiledim, invert = true)
    if dim == 3
        return tiled_order_3d(griddim, tiledim, invert);
    elseif dim == 2
        return tiled_order_2d(griddim, tiledim, invert);
    else
        return 1:griddim[1]; # 1D doesn't make sense
    end
end


#=

The robust control problem for a monopolist with adjustment costs.  The
inverse demand curve is:

  p_t = a_0 - a_1 y_t + d_t

where d_{t+1} = \rho d_t + \sigma_d w_{t+1} for w_t ~ N(0,1) and iid.
The period return function for the monopolist is

  r_t =  p_t y_t - gam (y_{t+1} - y_t)^2 / 2 - c y_t

The objective of the firm is E_t \sum_{t=0}^\infty \beta^t r_t

For the linear regulator, we take the state and control to be

    x_t = (1, y_t, d_t) and u_t = y_{t+1} - y_t

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-07-05

References
----------

Simple port of the file examples/robust_monopolist.py

http://quant-econ.net/robustness.html#application

=#
using QuantEcon
using Plots
pyplot()
using Grid

# model parameters
a_0     = 100
a_1     = 0.5
rho     = 0.9
sigma_d = 0.05
bet    = 0.95
c       = 2
gam   = 50.0
theta = 0.002
ac    = (a_0 - c) / 2.0

# Define LQ matrices
R = [0 ac    0
     ac -a_1 0.5
     0. 0.5  0]
R = -R  # For minimization
Q = [gam / 2.0]'
A = [1. 0. 0.
     0. 1. 0.
     0. 0. rho]
B = [0. 1. 0.]'
C = [0. 0. sigma_d]'

## Functions

function evaluate_policy(theta, F)
    rlq = RBLQ(Q, R, A, B, C, bet, theta)
    K_F, P_F, d_F, O_F, o_F = evaluate_F(rlq, F)
    x0 = [1.0 0.0 0.0]'
    value = - x0'*P_F*x0 - d_F
    entropy = x0'*O_F*x0 + o_F
    return value[1], entropy[1]  # return scalars
end

function value_and_entropy(emax, F, bw, grid_size=1000)
    if lowercase(bw) == "worst"
        thetas = 1 ./ linspace(1e-8, 1000, grid_size)
    else
        thetas = -1 ./ linspace(1e-8, 1000, grid_size)
    end

    data = Array{Float64}(grid_size, 2)

    for (i, theta) in enumerate(thetas)
        data[i, :] = collect(evaluate_policy(theta, F))
        if data[i, 2] >= emax  # stop at this entropy level
            data = data[1:i, :]
            break
        end
    end
    return data
end

## Main

# compute optimal rule
optimal_lq = LQ(Q, R, A, B, C, zero(B'A), bet=bet)
Po, Fo, Do = stationary_values(optimal_lq)

# compute robust rule for our theta
baseline_robust = RBLQ(Q, R, A, B, C, bet, theta)
Fb, Kb, Pb = robust_rule(baseline_robust)

# Check the positive definiteness of worst-case covariance matrix to
# ensure that theta exceeds the breakdown point
test_matrix = eye(size(Pb, 1)) - (C' * Pb * C ./ theta)[1]
eigenvals, eigenvecs = eig(test_matrix)
@assert all(eigenvals .>= 0)

emax = 1.6e6

# compute values and entropies
optimal_best_case = value_and_entropy(emax, Fo, "best")
robust_best_case = value_and_entropy(emax, Fb, "best")
optimal_worst_case = value_and_entropy(emax, Fo, "worst")
robust_worst_case = value_and_entropy(emax, Fb, "worst")

# we reverse order of "worst_case"s so values are ascending
data_pairs = ((optimal_best_case, optimal_worst_case),
              (robust_best_case, robust_worst_case))

egrid = linspace(0, emax, 100)
egrid_data = Array{Float64}[]
for data_pair in data_pairs
    for data in data_pair
        x, y = data[:, 2], data[:, 1]
        curve(z) = InterpIrregular(x, y, BCnearest, InterpLinear)[z]
        push!(egrid_data, curve(egrid))
    end
end
plot(egrid, egrid_data, color=[:red :red :blue :blue])
plot!(egrid, egrid_data[1], fillrange=egrid_data[2],
      fillcolor=:red, fillalpha=0.1, color=:red, legend=:none)
plot!(egrid, egrid_data[3], fillrange=egrid_data[4],
      fillcolor=:blue, fillalpha=0.1, color=:blue, legend=:none)
plot!(xlabel="Entropy", ylabel="Value")


using Nabla

# Generate the values required to compute a matrix quadratic form.
N = 5
B = randn(5, 5)
A = B.'B + UniformScaling(1e-6)

# Low-level API computation of derivatives.

# Construct some trackable data.
x_, y_ = randn(N), randn(N)
x, y = Leaf.(Tape(), (x_, y_))

# Compute the forward pass.
z = x.' * (A * y) # Temporary bracketting because we don't support RowVectors yet.

println("Output of the forward pass is:")
println(z)
println()
println("y is $(Nabla.unbox(z)).")
println()

# Get the reverse tape.
z̄ = ∇(z)
println("Output of reverse-pass is")
println(z̄)
println()

# Index into the reverse tape using x to get the gradient of `y` w.r.t. `x`.
x̄ = z̄[x]
println("Gradient of z w.r.t. x at $x_ is $x̄.")
println()

ȳ = z̄[y]
println("Gradient of z w.r.t. y at $y_ is $ȳ")

# (Current) High-Level API computation of derivatives. I will probably maintain this
# interface and extend is significantly as it is currently rather limited. It just returns
# a function which returns the gradient, which isn't really what you want.

# Define the function to be differentiated. Parameters w.r.t. which we want gradients must
# be arguments. Parameters that we don't want gradients w.r.t. should be passed in via a
# closure.
@unionise f(x::AbstractVector, y::AbstractVector) = x.'A * y

# Compute a function `∇f` which computes the derivative of `f` w.r.t. the inputs.
∇f = ∇(f)

# Compute the derivative of `f` w.r.t. `x` at `x_`. Result is currently a 1-Tuple. Might
# introduce a special case for unary functions where it just returns the result.
(x̄, ȳ) = ∇f(x_, y_)

@assert x̄ == z̄[x]
@assert ȳ == z̄[y]


abstract type AbstractMetric end

apply(::Type{M}, args...; kwargs...) where {M <: AbstractMetric} =
    apply(M, ConfusionMatrix(args...); kwargs...)

apply(::Type{M}, x::AbstractArray{<:ConfusionMatrix}; kwargs...) where {M <: AbstractMetric} =
    apply.(M, x; kwargs...)

macro metric(name)
    name_lw = Symbol(lowercase(string(name)))

    quote
        abstract type $(esc(name)) <: AbstractMetric end

        Base.@__doc__  function $(esc(name_lw))(args...; kwargs...)
            apply($(esc(name)), args...; kwargs...)
        end
    end
end

"""
    $(SIGNATURES)

Returns # true positive samples.
"""
@metric True_positive
apply(::Type{True_positive}, x::ConfusionMatrix) = x.tp

"""
    $(SIGNATURES)

Returns # true negative samples.
"""
@metric True_negative
apply(::Type{True_negative}, x::ConfusionMatrix) = x.tn

"""
    $(SIGNATURES)

Returns # false positive samples.
"""
@metric False_positive
apply(::Type{False_positive}, x::ConfusionMatrix) = x.fp

"""
    $(SIGNATURES)

Returns # false negative samples.
"""
@metric False_negative
apply(::Type{False_negative}, x::ConfusionMatrix) = x.fn

"""
    $(SIGNATURES)

Returns true positive rate `tp/p`.
Aliases: `sensitivity`,  `recall`, `hit_rate`.
"""
@metric True_positive_rate
apply(::Type{True_positive_rate}, x::ConfusionMatrix) = x.tp/x.p

const sensitivity = true_positive_rate
const recall      = true_positive_rate
const hit_rate    = true_positive_rate

"""
    $(SIGNATURES)

Returns true negative rate `tn/n`.
Aliases: `specificity`,  `selectivity`.
"""
@metric True_negative_rate
apply(::Type{True_negative_rate}, x::ConfusionMatrix) = x.tn/x.n

const specificity = true_negative_rate
const selectivity = true_negative_rate

"""
    $(SIGNATURES)

Returns false positive rate `fp/n`.
Aliases: `fall_out`, `type_I_error`.
"""
@metric False_positive_rate
apply(::Type{False_positive_rate}, x::ConfusionMatrix) = x.fp/x.n

const fall_out     = false_positive_rate
const type_I_error = false_positive_rate

"""
    $(SIGNATURES)

Returns false negative rate `fn/p`.
Aliases: `miss_rate`, `type_II_error`.
"""
@metric False_negative_rate
apply(::Type{False_negative_rate}, x::ConfusionMatrix) = x.fn/x.p

const miss_rate     = false_negative_rate
const type_II_error = false_negative_rate

"""
    $(SIGNATURES)

Returns precision `tp/(tp + fp)`.
Aliases: `positive_predictive_value`.
"""
@metric Precision
apply(::Type{Precision}, x::ConfusionMatrix) = (val = x.tp/(x.tp + x.fp); isnan(val) ? one(val) : val)

const positive_predictive_value = precision

"""
    $(SIGNATURES)

Returns negative predictive value `tn/(tn + fn)`.
"""
@metric Negative_predictive_value
apply(::Type{Negative_predictive_value}, x::ConfusionMatrix) = x.tn/(x.tn + x.fn)

"""
    $(SIGNATURES)

Returns false discovery rate `fp/(fp + tp)`.
"""
@metric False_discovery_rate
apply(::Type{False_discovery_rate}, x::ConfusionMatrix) = x.fp/(x.fp + x.tp)

"""
    $(SIGNATURES)

Returns false omission rate `fn/(fn + tn)`.
"""
@metric False_omission_rate
apply(::Type{False_omission_rate}, x::ConfusionMatrix) = x.fn/(x.fn + x.tn)

"""
    $(SIGNATURES)

Returns threat score `tp/(tp + fn + fp)`.
Aliases: `critical_success_index`.
"""
@metric Threat_score
apply(::Type{Threat_score}, x::ConfusionMatrix) = x.tp/(x.tp + x.fn + x.fp)

const critical_success_index = threat_score

"""
    $(SIGNATURES)

Returns accuracy `(tp + tn)/(p + n).
"""
@metric Accuracy
apply(::Type{Accuracy}, x::ConfusionMatrix) = (x.tp + x.tn)/(x.p + x.n)

"""
    $(SIGNATURES)

Returns balanced accuracy `(tpr + tnr)/2`.
"""
@metric Balanced_accuracy
apply(::Type{Balanced_accuracy}, x::ConfusionMatrix) =
    (true_positive_rate(x) + true_negative_rate(x))/2

"""
    $(SIGNATURES)

Returns error rate `1 - accuracy`.
"""
@metric Error_rate
apply(::Type{Error_rate}, x::ConfusionMatrix) = 1 - accuracy(x)

"""
    $(SIGNATURES)

Returns balanced error rate `1 - balanced_accuracy`.
"""
@metric Balanced_error_rate
apply(::Type{Balanced_error_rate}, x::ConfusionMatrix) = 1 - balanced_accuracy(x)

"""
    $(SIGNATURES)

Returns f1 score `2*precision*recall/(precision + recall)`.
"""
@metric F1_score
apply(::Type{F1_score}, x::ConfusionMatrix) =
    2*precision(x)*recall(x)/(precision(x) + recall(x))

"""
    $(SIGNATURES)

Returns fβ score `(1 + β^2)*precision*recall/(β^2*precision + recall)`.
"""
@metric Fβ_score
apply(::Type{Fβ_score}, x::ConfusionMatrix; β::Real = 1) =
    (1 + β^2)*precision(x)*recall(x)/(β^2*precision(x) + recall(x))

"""
    $(SIGNATURES)

Returns Matthews correlation coefficient `(tp*tn - fp*fn)/sqrt((tp+fp)(tp+fn)(tn+fp)(tn+fn))`.
Aliases: ` mcc`.
"""
@metric Matthews_correlation_coefficient
apply(::Type{Matthews_correlation_coefficient}, x::ConfusionMatrix) =
    (x.tp*x.tn - x.fp*x.fn)/sqrt((x.tp + x.fp)*(x.tp + x.fn)*(x.tn + x.fp)*(x.tn + x.fn))

const mcc = matthews_correlation_coefficient

"""
    $(SIGNATURES)

Returns quant `(fn + tn)/(p + n)`.
"""
@metric Quant
apply(::Type{Quant}, x::ConfusionMatrix) = (x.fn + x.tn)/(x.p + x.n)

"""
    $(SIGNATURES)

Returns topquant `1 - quant`.
"""
@metric Topquant
apply(::Type{Topquant}, x::ConfusionMatrix) = 1 - quant(x)

"""
    $(SIGNATURES)

Returns positive likelihood ratio `tpr/fpr`.
"""
@metric Positive_likelihood_ratio
apply(::Type{Positive_likelihood_ratio}, x::ConfusionMatrix) =
    true_positive_rate(x)/false_positive_rate(x)

"""
    $(SIGNATURES)

Returns negative likelihood ratio `fnr/tnr`.
"""
@metric Negative_likelihood_ratio
apply(::Type{Negative_likelihood_ratio}, x::ConfusionMatrix) =
    false_negative_rate(x)/true_negative_rate(x)

"""
    $(SIGNATURES)

Returns diagnostic odds ratio `tpr*tnr/(fpr*fnr)`.
"""
@metric Diagnostic_odds_ratio
apply(::Type{Diagnostic_odds_ratio}, x::ConfusionMatrix) =
    true_positive_rate(x)*true_negative_rate(x)/(false_positive_rate(x)*false_negative_rate(x))

"""
    $(SIGNATURES)

Returns prevalence `p/(p + n)`.
"""
@metric Prevalence
apply(::Type{Prevalence}, x::ConfusionMatrix) = x.p/(x.p + x.n)


using Revise
using StatsBase
using Statistics
using OnlineStats
using BenchmarkTools
using WeightedOnlineStats

x = (collect(reverse((1:10_000_000)*0.0000001)))
w = (collect((1:10_000_000)*0.0000001))
x1 = x[1:5_000_000]
x2 = (x[5_000_001:end])
w1 = w[1:5_000_000]
w2 = (w[5_000_001:end])

vw = WeightedVariance()
vw1 = WeightedVariance()
vw2 = WeightedVariance()

fit!(vw, x, w)
fit!(vw1, x1, w1)
fit!(vw2, x2, w2)

@code_warntype merge!(vw1, vw2)

@benchmark merge($vw1, $vw2)
#=
BenchmarkTools.Trial.1.1:
ll. 133-145 (w/o 135-138)
  memory estimate:  544 bytes
  allocs estimate:  7
  --------------
  minimum time:     3.140 μs (0.00% GC)
  median time:      3.257 μs (0.00% GC)
  mean time:        3.476 μs (1.52% GC)
  maximum time:     272.601 μs (97.37% GC)
  --------------
  samples:          10000
  evals/sample:     8

BenchmarkTools.Trial.1.2:
ll. 133-145 (w/o 139-141)
  memory estimate:  544 bytes
  allocs estimate:  7
  --------------
  minimum time:     3.141 μs (0.00% GC)
  median time:      3.206 μs (0.00% GC)
  mean time:        3.380 μs (1.54% GC)
  maximum time:     268.521 μs (97.16% GC)
  --------------
  samples:          10000
  evals/sample:     8

BenchmarkTools.Trial.2:
ll. 151-161
  memory estimate:  544 bytes
  allocs estimate:  7
  --------------
  minimum time:     3.136 μs (0.00% GC)
  median time:      3.229 μs (0.00% GC)
  mean time:        3.347 μs (1.53% GC)
  maximum time:     262.741 μs (97.26% GC)
  --------------
  samples:          10000
  evals/sample:     8
=#
x = (collect(reverse((1:10_000_000)*0.0000001)))
w = (collect((1:10_000_000)*0.0000001))
x1 = x[1:5_000_000]
x2 = BigFloat.(x[5_000_001:end])
w1 = w[1:5_000_000]
w2 = BigFloat.(w[5_000_001:end])

vw = WeightedVariance()
vw1 = WeightedVariance()
vw2 = WeightedVariance(BigFloat)

fit!(vw, x, w)
fit!(vw1, x1, w1)
fit!(vw2, x2, w2)

@code_warntype merge(vw1, vw2)
merge(vw1, vw2)
merge(vw2, vw1)
@benchmark merge($vw1, $vw2)
@benchmark merge($vw2, $vw1)
#=
BenchmarkTools.Trial.1.1:
ll. 133-145 (w/o 135-138)
  memory estimate:  2.28 KiB
  allocs estimate:  39
  --------------
  minimum time:     4.955 μs (0.00% GC)
  median time:      5.060 μs (0.00% GC)
  mean time:        5.603 μs (3.95% GC)
  maximum time:     346.025 μs (94.78% GC)
  --------------
  samples:          10000
  evals/sample:     7

BenchmarkTools.Trial.1.2:
ll. 133-145 (w/o 139-141)
  memory estimate:  2.06 KiB
  allocs estimate:  35
  --------------
  minimum time:     4.779 μs (0.00% GC)
  median time:      4.861 μs (0.00% GC)
  mean time:        5.366 μs (3.63% GC)
  maximum time:     348.297 μs (96.43% GC)
  --------------
  samples:          10000
  evals/sample:     7

BenchmarkTools.Trial.2:
ll. 151-161
  memory estimate:  2.50 KiB
  allocs estimate:  43
  --------------
  minimum time:     5.398 μs (0.00% GC)
  median time:      5.515 μs (0.00% GC)
  mean time:        6.042 μs (3.80% GC)
  maximum time:     426.032 μs (96.56% GC)
  --------------
  samples:          10000
  evals/sample:     6
=#

using Traceur

x = rand(100)
x2 = rand(100, 2)
w = rand(100)

@trace fit!(WeightedSum(), x, w)
@trace fit!(WeightedSum(Float32), x, w)

@trace fit!(WeightedMean(), x, w)
@trace fit!(WeightedMean(Float32), x, w)

@trace fit!(WeightedVariance(), x, w)
@trace fit!(WeightedVariance(Float32), x, w)

@trace fit!(WeightedCovMatrix(), x2, w)
@trace fit!(WeightedCovMatrix(Float32), x2, w)


## Kernel machines figures
using Plots
pyplot()
theme(:dark)

##
x1 = [1.0,5]
x2 = [2,-1.0]
y1 = 1.0
y2 = -1.0

v = [x1'; x2'] \ [y1,y2]

plot(framestyle=:origin,aspect_ratio=:equal,xlim=(-3,3),ylim=(-1.5,5.5))
scatter!([x1[1]],[x1[2]],label="")
annotate!([x1[1]],[x1[2]],text("+1 ",:white,:right,12))
scatter!([x2[1]],[x2[2]],label="")
annotate!([x2[1]],[x2[2]],text("-1 ",:white,:right,12))

##
plot!(size=(200,200))
savefig("kernel-interpolation-2d-setup.pdf")
##
scatter!([v[1]],[v[2]],label="")
plot!(v[1]*[-5,5], v[2]*[-5,5],label="")
annotate!([v[1]],[v[2]], text("v ",:white,:right,12))
savefig("kernel-interpolation-2d-soluiton.pdf")
##
x1 = [1.0,5,1]
x2 = [2,-1.0,-1]
y1 = 1.0
y2 = -1.0

v = [x1'; x2'] \ [y1,y2]

plot(aspect_ratio=:equal,xlim=(-3,3),ylim=(-1.5,5.5))
scatter!([x1[1]],[x1[2]],[x1[3]],label="")
annotate!([x1[1]],[x1[2]],[x1[3]],"+1 ")
scatter!([x2[1]],[x2[2]],[x2[3]],label="")
##
plot!(size=(200,200),dpi=300)
anim = @animate for i=1:180
  plot!(camera=(i,30))
end
gif(anim, "kernel-interpolation-3d-setup.gif")

##
#annotate!([x2[1]],[x2[2]],text("-1 ",:white,:right,12))
scatter!([v[1]],[v[2]],[v[3]], label="")
plot!(v[1]*[-5,5], v[2]*[-5,5], v[3]*[-5,5], label="",color=5)

## Other solutions
using LinearAlgebra
v = nullspace([x1'; x2'])
x1'*a
x1'*(a+v[:,1])

b = a+v[:,1]
scatter!([b[1]],[b[2]],[b[3]], label="")
plot!(b[1]*[-5,5], b[2]*[-5,5], b[3]*[-5,5], label="", color=8)

##
plot!(size=(200,200),dpi=300)
anim = @animate for i=1:180
  plot!(camera=(i,30))
end
gif(anim, "kernel-interpolation-3d-solution.gif")


struct Cartpole{T} <: AbstractModel
    mc::T
    mp::T
    l::T
    g::T
end

Cartpole() = Cartpole(1.0, 0.2, 0.5, 9.81)

function RobotDynamics.dynamics(model::Cartpole, x, u)
    mc = model.mc  # mass of the cart in kg (10)
    mp = model.mp   # mass of the pole (point mass at the end) in kg
    l = model.l   # length of the pole in m
    g = model.g  # gravity m/s^2

    q = x[ @SVector [1,2] ]
    qd = x[ @SVector [3,4] ]

    s = sin(q[2])
    c = cos(q[2])

    H = @SMatrix [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = @SMatrix [0 -mp*qd[2]*l*s; 0 0]
    G = @SVector [0, mp*g*l*s]
    B = @SVector [1, 0]

    qdd = -H\(C*qd + G - B*u[1])
    return [qd; qdd]
end

RobotDynamics.state_dim(::Cartpole) = 4
RobotDynamics.control_dim(::Cartpole) = 1
RobotDynamics.diffmethod(::Cartpole) = RobotDynamics.ForwardAD()


"""
    Kissler_VL_model.jl

Written by Carl Whitfield, University of Manchester, 2021

Only to be included by definitions.jl
"""

#===================== VL trajectory params  =====================#

#====> Kissler et al 2020 <====#
const beta_peak = 0.7
const alpha_peak = 3.5*beta_peak                      #Peak time from LOD distribution params
const beta_decay = 0.3
const alpha_decay_symp = 10.5*beta_decay              #Decay time to LOD distribution params
const alpha_decay_asymp = 6.7*beta_decay
const peakVL_mean = 7.533                             #Peak viral load distribution params
const peakVL_sd = 1.164
const VL_LOD_Kissler = log10(250) - 0.93733/3.60971   #Limit of Detection in Kissler study
const V_0 = 0.5255                  #from Ferretti (initial VL) (if using V0_model_opt)
const onset_frac = 0.38             #frac of peak time added before peak (if using PVT_model_opt)

#Distributions used to generate params
const PVLdist = truncated(Normal(peakVL_mean,peakVL_sd),VL_LOD_Kissler,Inf)
const PTdist = Gamma(alpha_peak,1.0/beta_peak)
const DTsympdist = Gamma(alpha_decay_symp,1.0/beta_decay)
const DTasympdist = Gamma(alpha_decay_asymp,1.0/beta_decay)

"""
    generate_VL_params_Kissler(Asymp::Bool)

Generate parameters for VL model.

## Arguments
`Asymp` = Whether individual is asymptomatic

## Returns:
`Float` = peak viral load (log10 copies/ml)
`Float` = peak viral load time (days since infection)
`Float` = growth rate (1/days)
`Float` = decay rate (1/days)
"""
function generate_VL_params_Kissler(Asymp::Bool)
    PVL = rand(PVLdist)  #peak viral load
    PT = rand(PTdist)    #time from LOD to peak

    r = log(10)*(PVL - VL_LOD_Kissler)/PT

    if onset_opt == V0_model_opt
        OT = log(10)*(VL_LOD_Kissler - V_0)/r    #time from infection to LOD
    elseif onset_opt == PVT_model_opt
        OT = onset_frac*PT
    end

    if Asymp
        DT = rand(DTasympdist)
    else
        DT = rand(DTsympdist)
    end

    tp = OT+PT
    d = log(10)*(PVL - VL_LOD_Kissler)/DT

    return PVL, tp, r, d
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

# Variogram functions using local anisotropies

function qmat(q,m)
  # Anisotropy matrix for Mahalanobis distance
  N = length(m)
  P = quat_to_dcm(q)[SOneTo(N),SOneTo(N)]
  Λ = Diagonal(SVector{N}(one(eltype(m))./m.^2))
  Q = P'*Λ*P
  # need to explicitly set Symmetric(Q), but then det(Q) will fail with StaticArrays
  # while not fixed, setting Symmetric only when calling Mahalanobis
  # https://github.com/JuliaArrays/StaticArrays.jl/issues/955
end

function mwvario(estimator, localpar)
  # get local Mahalanobis matrix
  Q = qmat(localpar[1],localpar[2])
  # get reference pars. and apply local anisotropy to given structures
  p = structures(estimator.γ)
  γs = map(p[3]) do γ
    Qs = Q ./ radii(γ.ball)' .^ 2
    γ = @set γ.ball = MetricBall(1.0, Mahalanobis(Symmetric(Qs)))
  end
  γl = NuggetEffect(p[1]) + sum(c*γ for (c, γ) in zip(p[2], γs))

  # return local estimator
  if typeof(estimator) <: SimpleKriging
    return SimpleKriging(γl, estimator.mean)
  elseif typeof(estimator) <: OrdinaryKriging
    return OrdinaryKriging(γl)
  end

end

function kcfill!(Γ, γ::Variogram, X, localaniso)
  n = length(X)
  @inbounds for j in 1:n
    xj = X[j]
    Qj = localaniso[j]
    for i=j+1:n
      xi = X[i]
      Qi = localaniso[i]
      Γ[i,j] = kccov(γ, xi, xj, Qi, Qj)
    end
    Γ[j,j] = sill(γ)
    for i=1:j-1
      Γ[i,j] = Γ[j,i]
    end
  end
end

function kccov(γ::Variogram, xi, xj, Qi::AbstractMatrix, Qj::AbstractMatrix)
  Qij = (Qi+Qj)/2

  ## modify variogram with average anisotropy matrix
  p = structures(γ)
  γs = map(p[3]) do γx
    Qs = Qij ./ radii(γx.ball)' .^ 2
    γx = @set γx.ball = MetricBall(1.0, Mahalanobis(Symmetric(Qs)))
  end
  γl = NuggetEffect(p[1]) + sum(c*γx for (c, γx) in zip(p[2], γs))

  Cij = (sill(γl)-γl(xi,xj))
  (det(Qi)^0.25)*(det(Qj)^0.25)*(det(Qij)^-0.5)*Cij
end


using ClassifierMetrics
using Random
Random.seed!(123)

# Generate synthetic data
labels = rand(Bool, 500);

# function to add noise
add_noise(label::Bool, λ=0.0) = label ? 1 - λ*rand() : λ*rand()
# simulate good classifier
good = add_noise.(labels, 0.6)
# simulate bad classifier
bad  = add_noise.(labels, 0.98)

roc_good = ROC(labels, good)
roc_bad  = ROC(labels,  bad)

pr_bad  = PR(labels,  bad)
pr_good  = PR(labels,  good)

det_bad  = DET(labels,  bad)
det_good  = DET(labels,  good)

println("ROCAUC (good detector): $(auc(roc_good)) ")
println("ROCAUC (bad  detector): $(auc(roc_bad )) ")

println("PRAUC (good detector): $(auc(pr_good)) ")
println("PRAUC (bad  detector): $(auc(pr_bad )) ")

using Plots
gr()

roc_plot = plot(roc_good, label="good");
plot!(roc_plot,roc_bad, label="bad", legend=:bottomright)

pr_plot = plot(pr_good);
plot!(pr_plot,pr_bad, legend=false)

det_plot = plot(det_good);
plot!(det_plot,det_bad, legend=false)

plot(roc_plot, pr_plot, det_plot)


module henon_heiles
using ReachabilityAnalysis
@taylorize function henon_heiles!(du, u, p, t)
    p₁, p₂, q₁, q₂ = u[1], u[2], u[3], u[4]
    du[1] = -q₁*(1 + 2q₂)
    du[2] = -q₂-(q₁^2 - q₂^2)
    du[3] = p₁
    du[4] = p₂
end

function model(X0)
    S = @system(x' = henon_heiles!(x), dim:4)
    return IVP(S, X0)
end
end # module


using SeisReconstruction
using SeisProcessing
using LinearAlgebra
using Test

d = SeisLinearEvents()

#randomly decimate 50% of traces

mask = rand(1,size(d,2));
perc=50
mask[(LinearIndices(mask .< perc/100))[findall(mask .< perc/100)]] .= 0;
mask[(LinearIndices(mask .>= perc/100))[findall(mask .>= perc/100)] ] .= 1;

ddec = Array{Float64}(undef,size(d));

for it = 1 : size(d,1)
    ddec[it,:] = d[it:it,:].*mask;
end

dmwni = SeisMWNI(ddec,dt=0.004,fmax=100,Niter_external=5,Niter_internal=10)

# test that quality factor is greater than 10 Decibels
quality_factor = 10*log10(norm(d[:],2)/norm(dmwni[:]-d[:],2))
println("Quality factor = ",quality_factor)
@test quality_factor > 5


@doc raw"""
First and second Poincaré invariant for a guiding center particle in an
axisymmetric magnetic field of the form ``B(x,y,z) = B_0 (1 + x^2 + y^2) e_z``.

The loop for the first invariant is initialized by
```math
q (\tau) = \begin{pmatrix}
r_x \cos (2\pi \tau) \\
r_y \sin (2\pi \tau) \\
z_0 + z_1 \sin (2\pi \tau) \\
u_0 + u_1 \cos (2\pi \tau) \\
\end{pmatrix}
```
with parameters
```math
B_0 = 1, \quad
r_x = 0.5, \quad
r_y = 0.3, \quad
z_0 = 0.0, \quad
z_1 = 0.1, \quad
u_0 = 0.5, \quad
u_1 = 0.05, \quad
\mu = 0.01 .
```

The surface for the second invariant is initialized by
```math
q (\tau) = \begin{pmatrix}
r_0 (\sigma - 0.5) \\
r_0 (\tau   - 0.5) \\
z_0 + z_1 \cos (2\pi \sigma) \cos (2\pi \tau) \\
u_0 + u_1 \sin (2\pi \sigma) \sin (2\pi \tau) \\
\end{pmatrix}
```
with parameters
```math
B_0 = 1, \quad
r_0 = 0.5, \quad
z_0 = 0.0, \quad
z_1 = 0.1, \quad
u_0 = 0.5, \quad
u_1 = 0.01, \quad
\mu = 0.01 .
```
"""
module SymmetricField

    import ElectromagneticFields.SymmetricQuadratic

    SymmetricQuadratic.@code() # inject magnetic field code

    μ_loop() = 1E-2
    μ_surface() = 1E-2

    function f_loop(s)
        rx = 0.5
        ry = 0.3
        z0 = 0.0
        z1 = 0.1
        u0 = 5E-1
        u1 = 5E-2

        xs = rx*cos(2π*s)
        ys = ry*sin(2π*s)
        zs = z0 + z1 * sin(2π*s)
        us = u0 + u1 * cos(2π*s)

        qs = [xs, ys, zs, us]

        return qs
    end

    function f_surface(s,t)
        r0 = 0.5
        z0 = 0.0
        z1 = 0.1
        u0 = 5E-1
        u1 = 1E-2

        x  = r0*(s-0.5)
        y  = r0*(t-0.5)
        z  = z0 + z1 * cos(2π*s) * cos(2π*t)
        u  = u0 + u1 * sin(2π*s) * sin(2π*t)

        q  = [x, y, z, u]

        return q
    end

    include("guiding_center_4d_common.jl")
    include("guiding_center_4d_equations.jl")
    include("guiding_center_4d_loop.jl")
    include("guiding_center_4d_surface.jl")
    include("guiding_center_4d_diagnostics.jl")

end


using Flux, DiffEqFlux, DifferentialEquations, Plots

## Setup DDE to optimize
function delay_lotka_volterra(du,u,h,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = (α - β*y)*h(p,t-0.1)[1]
  du[2] = dy = (δ*x - γ)*y
end

h(p,t) = ones(eltype(p),2)
prob = DDEProblem(delay_lotka_volterra,[1.0,1.0],h,(0.0,10.0),constant_lags=[0.1])
p = param([2.2, 1.0, 2.0, 0.4])
function predict_rd_dde()
  diffeq_rd(p,prob,MethodOfSteps(Tsit5()),saveat=0.1)[1,:]
end
loss_rd_dde() = sum(abs2,x-1 for x in predict_rd_dde())

# Optimize the parameters so the DDE's solution stays near 1

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function () #callback function to observe training
  display(loss_rd_dde())
  # using `remake` to re-create our `prob` with current parameters `p`
  display(plot(solve(remake(prob,p=Flux.data(p)),MethodOfSteps(Tsit5()),saveat=0.1),ylim=(0,6)))
end
# Display the DDE with the initial parameter values.
cb()
Flux.train!(loss_rd_dde, [p], data, opt, cb = cb)


module RBM
using Winston

export dream
export showWeights

function xp(x)
    return 1.0 ./ (1+exp(-x))
end

function vtohprob{T<:Real,P<:Integer}(rbmw::Matrix{T},vstate::Matrix{P})
    return xp(rbmw * vstate)
end

function htovprob{T<:Real,P<:Integer}(rbmw::Matrix{T},hstate::Matrix{P})
    return xp(rbmw' * hstate)
end

function gradient{T<:Real,P<:Integer}(vstate::Matrix{P},hstate::Matrix{T})
    hstate * vstate'/size(vstate)[2]
end

function bernoulli{T<:Real}(v::Vector{T})
    return map((x)->x>rand()?1:0,v)
end

function bernoulli{T<:Real}(m::Matrix{T})
    return map((x)->x>rand()?1:0,m)
end

function cd{T<:Real}(rbmw::Matrix{T},data::Matrix{T})
    vdata=bernoulli(data)
    h0=bernoulli(vtohprob(rbmw,vdata))
    vh0=gradient(vdata,h0)
    v1=bernoulli(htovprob(rbmw,h0))
    h1=bernoulli(vtohprob(rbmw,v1))
    vh1=gradient(v1,h1)
    return vh0-vh1
end

function showWeights{T<:Real}(W::Matrix{T},row,col)
    r,c = size(W)
    plots = []
    for i in 1:r
        plots=vcat(plots,Winston.imagesc((reshape(W[i,:],28,28))'))
    end

    mgrid=Winston.Table(row,col)
    c=1
    for i = 1:row
        for j=1:col
            mgrid[i,j] = plots[c]
            c+=1
        end
    end
    display(mgrid)
end

function dream{T<:Real}(rbmw::Matrix{T},data::Matrix{T})
    vdata=bernoulli(data)
    h0=bernoulli(vtohprob(rbmw,vdata))
    v1=bernoulli(htovprob(rbmw,h0))

    rw = size(vdata)[2]
    mgrid = Winston.Table(1,2)
    for i in 1:rw
        p1=Winston.imagesc(reshape(vdata[:,i],28,28)')
        p2=Winston.imagesc(reshape(v1[:,i],28,28)')
        mgrid[1,1]=p1
        mgrid[1,2]=p2
        display(mgrid)
    end
    return v1
end

function rbm{T<:Real,P<:Integer}(nhid::P,data::Matrix{T},lr::T,niter::P,mbatchsz::P,mom::T)
    psz,nsz=size(data)
    momentum=mom
    if nsz % mbatchsz != 0
        println("test cases must be multiple of batch size")
        return
    end
    model=(rand(nhid,psz)*2-1)*0.1
    momSpeed=zeros(nhid,psz)

    startMini=1
    for iter in 1:niter
        if mod(iter,50)==0
            println("iter:",iter," ",sqrt(mean(map((x)->x*x,model))))
        end
        mbatch=data[:,startMini:(startMini+mbatchsz-1)]
        startMini=(startMini+mbatchsz) % nsz
        gradient=cd(model,mbatch)
        momSpeed=momentum*momSpeed+gradient
        model=model+momSpeed*lr
    end
    return model
end
end

using RBM
#trainData=readdlm("./mnist_train.csv",',',Float64)
#trainLab=trainData[:,1]
#trainData=trainData[:,2:size(trainData)[2]]/255.0
#trainData=trainData'

testData=readdlm("./mnist_test.csv",',',Float64)
testLab=testData[:,1]
testData=testData[:,2:size(testData)[2]]/255.0
testData=testData'

hidSz=30
mbatchsz=100
#niter=20000
niter = 500
lr=0.01
mom=0.9

#W=RBM.rbm(hidSz,trainData,lr,niter,mbatchsz,mom)

W=readcsv("./weighnobias.csv")

RBM.dream(W,testData[:,1:50]);


using Test, Random, LinearAlgebra, YaoArrayRegister, YaoBlocks, YaoBase

A = rand(ComplexF64, 4, 4)
mg = matblock(A)
mg2 = copy(mg)
@test mg2 == mg
@test mat(ComplexF64, mg) ≈ A
@test_logs (:warn, "converting Complex{Float64} to eltype Complex{Float32}, consider create another matblock with eltype Complex{Float32}") mat(ComplexF32, mg)

mg2.mat[:, 2] .= 10
@test mg2 != mg
@test nqubits(mg) == 2
@test_throws DimensionMismatch matblock(randn(3,3))

reg = rand_state(2)
@test apply!(copy(reg), mg) |> statevec == mat(mg) * reg.state |> vec

@test matblock(X) == GeneralMatrixBlock(mat(X))

a = rand_unitary(2)
@test mat(matblock(a))' == mat(matblock(a)')


using ReservoirComputing
using ParameterizedFunctions
using OrdinaryDiffEq
using Plots

#rossler system parameters
u0 = [0.1,0.5,0.1]
tspan = (0.0,1000.0)
p = [0.5,2.0,4.0]

#define rossler system
function rossler(du,u,p,t)
    du[1] = -u[2]-u[3]
    du[2] = u[1] + p[1]*u[2]
    du[3] = p[2]+u[3]*(u[1]-p[3])
end

#solve system
prob = ODEProblem(rossler, u0, tspan, p)
sol = solve(prob, ABM54(), dt=0.02)
v = sol.u
data = Matrix(hcat(v...))

#model parameters
shift = 500
approx_res_size = 300
radius = 1.15
activation = tanh
degree = 4
sigma = 0.1
train_len = 5000
predict_len = 2500
beta = 0.0
alpha = 1.0
nla_type = NLADefault()
extended_states = false

#get data
train = data[:, shift:shift+train_len-1]
test = data[:, shift+train_len:shift+train_len+predict_len-1]

#create echo state network
esn = ESN(approx_res_size,
    train,
    degree,
    radius,
    activation = activation,
    sigma = sigma,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

#training and prediction
W_out = ESNtrain(esn, beta)
output = ESNpredict(esn, predict_len, W_out)

#plots and images
comp = plot(transpose(output),layout=(3,1), label="predicted")
comp = plot!(transpose(test),layout=(3,1), label="actual")
savefig(comp, "rossler_coord")
sf = plot(transpose(output)[:,1], transpose(output)[:,2], transpose(output)[:,3], label="predicted")
sf2 = plot(transpose(test)[:,1], transpose(test)[:,2], transpose(test)[:,3], label="actual")
final = plot(sf, sf2, layout=(1,2), label = ["predicted" "actual"])
savefig(final, "rossler_attractor")