module Clusterpath

using LinearAlgebra
using Random
using DataFrames
using Interpolations
using Distributed
using MultivariateStats
using Distributions
import Gaston

include("AVLTree.jl")

export clusterpath, plot_cluster, plot_path, assign_cluster
export clusterpath_pop, cond_mean_on_LR, find_truncation, find_split
export generate_mixture_normal

"""
    clusterpath!(y; α=0.1, jitter=1e-6, return_split=false, silence_warning=false)
Perform convex clustering via α-thresholded clusterpath algorithm to `y`.
To avoid ties, add a small perturbation (`jitter`) to each data, and then sort it.
Returns a dictionary of
* (1) idx: cluster indices assigned to each elements in `y`
* (2) lambda: λ values in each merge
* (3) alpha: solution corresponding to each observation where merge happened
* (4) n_cluster: number of clusters wrt merging algorithm
* (5) n_splits: number of total splits, in corresponding top-down algorithm
* (6) splits: sample split point (only get returned when `return_split=true`)
"""
function clusterpath(y; α=0.1, jitter=1e-6, return_split=false, silence_warning=false)
    n = length(y)

    # remove ties with small perturbation `jitter`.
    x = copy(y)
    Random.seed!(0)
    x += randn(n) * jitter

    ii = sortperm(sortperm(x))
    sort!(x)  # sort `x` in ascending order

    # Initialize variables
    n_cluster = n  # (n) number of clusters
    probability = 0

    ## linkage between clusters
    next = Array{Union{Missing,Int64},1}(2:n+1)
    next[n] = missing
    prev = Array{Union{Missing,Int64},1}(0:n-1)
    prev[1] = missing

    ## status of clusters
    cardinals = ones(Int64, n)  # (freq) cardinality(freq) of each clusters
    centroids = copy(x)
    distance = diff(centroids) ./ (cardinals[1:n_cluster-1] + cardinals[2:end])
    λ_old = 0
    idx_cluster = Array(1:n) # (ind) cluster indices of observations

    ## big merge tracker
    was_big = false  # true if last merge was big
    is_big = false   # true if current merge is big

    # store returns along the solution path
    alphas = Array{Array{Float64,1},1}([copy(centroids)[ii]])
    λs = Float64[λ_old]  # (lambda) lambdas
    idx_along_path = Array{Array{Int64,1},1}()#[Array(1:n)[ii]])
    n_cluster_along_path = Int[]#n]
    n_splits = Int64(0)
    splits = Float64[]
#    println(idx_cluster[ii])
#    println(round.(distance,digits=5))

    # tree initailization
    ## (tree of pairwise distance btw successive indices)
    tree = AVLTree{Float64}()
    for i in 1:n-1
        push!(tree, distance[i], i)
    end

    """
    function for status update at each lambda, except for alpha
    """
    function update_return!(; update_lambda=true)
        if update_lambda
            append!(λs, λ_old)
        end
        append!(idx_along_path, [idx_cluster[ii]])
        append!(n_cluster_along_path, n_cluster)
    end

    """
    function for status update just for alpha
    """
    function update_alpha!()
        walk = 1
        alpha = Array{Float64}(undef, n)
        while !ismissing(walk)
            if !ismissing(next[walk])
                #n_repeat = next[walk] - walk
                alpha[walk:next[walk]-1] .= centroids[walk] + λ_old *((n-next[walk]+1)-(walk-1))
            else
                alpha[walk:end] .= centroids[walk] + λ_old *(-(walk-1))
            end
            walk = next[walk]
        end
        append!(alphas, [alpha[ii]])
    end

    ## main loop
    while n_cluster > 1

        # find minimum distance
        pivot = minimum_node(tree.root)
        λ_old, i = pivot.data, pivot.idx  # λ of next merge (minimal distance)

        #println(cardinals[i], " ", cardinals[next[i]])
#        println((i, next[i]))

        # detect and handle big merge
        if min(cardinals[i], cardinals[next[i]]) > ceil(α * n)  # But the paper says ⌈n * α⌉
            #println(cardinals[i], " ", cardinals[next[i]], " ", i, " ", λ_old, " ", n)

            is_big = true
            probability = (cardinals[i] + cardinals[next[i]]) / n

            n_splits += 1
            if return_split
                a_1 = max((x[idx_cluster .== idx_cluster[i]])...)
                a_2 = min((x[idx_cluster .== idx_cluster[next[i]]])...)
                split = (cardinals[next[i]] * a_2 + cardinals[i] * a_1) / (cardinals[next[i]] + cardinals[i])
                append!(splits, split)
            end
        end

        if is_big || was_big
            if is_big
                was_big = true
                is_big = false
            else
                was_big = false
            end
            # update λ path
            update_return!()
        end

        # update status
        centroids[i] = (centroids[i] * cardinals[i] + centroids[next[i]] * cardinals[next[i]]) / (cardinals[i] + cardinals[next[i]])
        cardinals[i] += cardinals[next[i]]
        idx_cluster[idx_cluster .== next[i]] .= i
        #n_right[i] = n_right[next[i]]

        # update distance and tree
        delete!(tree, distance[i])

        if !ismissing(prev[i])
            delete!(tree, distance[prev[i]])
            d_prev = (centroids[i] - centroids[prev[i]]) / (cardinals[prev[i]] + cardinals[i])
            distance[prev[i]] = d_prev
            push!(tree, d_prev, prev[i])
        end
        if !ismissing(next[next[i]])
            delete!(tree, distance[next[i]])
            d_next = (centroids[next[next[i]]] - centroids[i]) / (cardinals[i] + cardinals[next[next[i]]])
            distance[i] = d_next
            push!(tree, d_next, i)
        end

        # update linkage
        next[i] = next[next[i]]
        if !ismissing(next[i])
            prev[next[i]] = i
        end

        n_cluster -= 1

        if is_big || was_big
            if is_big
                was_big = true
                is_big = false
            else
                was_big = false
            end
            # update λ path
            #update_return!()
            update_alpha!()
        end
    end

    # last update for status at n_cluster = 1
    update_return!(update_lambda=false)

    # if the last Big Merge has sum of probabilities < 50%,
    # consider the case as if no splits have ever occured
    if probability < 0.5
        if !silence_warning
            @warn "No clusters were found: the last big merge has probability $(probability) < 0.5"
        end
        idx_along_path = Array(1:n)
        n_cluster_along_path = n
        n_splits = 0
    end

    return Dict("lambda" => λs,
                "idx" => idx_along_path,
                "alpha" => alphas,
                "n_cluster" => n_cluster_along_path,
                "n_splits" => n_splits,
                "splits" => splits)

end

"""
    rep(x, times)
Julia equivalent of R's `rep`.
"""
function rep(x, times)
    vcat(fill.(x, times)...)
end

"""
    merge_dfs(df1, df2)
merge two solution path dataframes with `outerjoin()`
"""
function merge_dfs(df1, df2)
    outerjoin(df1, df2; on=[:obs, :lambda], makeunique=true)
end

"""
    cast_solution(x; α=0.1, silence_warning=false)

A generator function that casts a dataframe of solution path using `clusterpath()`. The function yields the solution path of each observation one at an iteration. This return will be used in `plot_path()` and `plot_cluster()`.

* `x`: the data for which we want to cluster.
* `α`: threshold for Big Merge Tracker procedure. `α` shoulbe be ≥ 0. The smaller the value of `α`, the more detailed the solution path will be. The larger the `α`, the faster the algorithm will run. (default: `0.1`)
"""
function cast_solution(x; α=0.1, silence_warning=false)

    n, p = try
             n, p = size(x)
             (n, p)
           catch e
             if isa(e, BoundsError)
               (length(x), 1)
             else
               throw(e)
             end
           end
#     println(n, " ", p)

    #lambdas = Array{Union{Nothing,Array{Float64,1}},1}(repeat([nothing], p))
    ccs = Array{Dict,1}(undef, p)
    lambdas = Array{Array{Float64,1},1}(undef, p)
    Threads.@threads for k=1:p
        cc = clusterpath(x[:, k]; α=α, silence_warning=silence_warning)
        ccs[k] = cc
        lambdas[k] = cc["lambda"]
    end

    lambdas = vcat(lambdas...)
    unique!(lambdas)
    sort!(lambdas)

    Channel{DataFrame}() do channel
        # report one obs at a time (for memory management)
        for i=1:n
            # initialize solution path matrix for ith obs
            solution = Array{Union{Float64, Missing}, 2}(missing, length(lambdas), p+2)
            solution[:, 1] .= i
            solution[:, 2] = lambdas

            # loop over dimensions
            for k=1:p
                cc = ccs[k]

                idx = [l in cc["lambda"] for l in lambdas]

                alpha_idx = 1
                for (lam_idx, is_in_lambda) in enumerate(idx)
                    if is_in_lambda
                        solution[lam_idx, k+2] = cc["alpha"][alpha_idx][i]
                        alpha_idx += 1
                    end
                end

#                 alpha = Float64[a[i] for a in cc["alpha"]]

#                 df = DataFrame(obs = repeat([i], n_path),
#                                lambda = lambdas[k],
#                                dim = alpha)#,
#                                #col = repeat(["dim$k"], n * n_path))

#                 if isnothing(solution)
#                     solution = df
#                 else
#                     solution = merge_dfs(solution, df)
#                 end

                # calculate cluster number based on n_splits
            end

            solution = DataFrame(solution)
            rename!(solution, vcat(["obs", "lambda"], ["dim$k" for k=1:p]))
#             sort!(solution, [:obs, :lambda])

            put!(channel, solution)
        end
    end
end

"""
    _prepare(solution_i; for_path::Bool)

prepare the return of `cast_solution()` for `plot_path()`, or `plot_cluster()`.
* `solution_i`: single solution path element of `cast_solution()`.
* `for_path`: if `true` then prepare for `plot_path()`, `false` then prepare for `plot_cluster()`.
"""
function _prepare!(solution_i; for_path::Bool)
    p = size(solution_i, 2) - 2
    n = length(unique(solution_i.obs))

    #unique!(solution_i, vcat(["obs"], ["dim$k" for k=1:p]))
    #Impute.interp!(solution_i)

    if !for_path
        nothing
        #solution_i[1:end-n_node-1, :] .= missing
        #dropmissing!(solution_i)
    end

    for j=1:p
        idx_missing = ismissing.(solution_i["dim$j"])
        #trunc = findlast(.~idx_missing)

        lambda_obs = solution_i[.~idx_missing, "lambda"]
        point_obs = Float64.(solution_i[.~idx_missing, "dim$j"])
        itp = interpolate((lambda_obs,), point_obs, Gridded(Linear()))
        etp = extrapolate(itp, Flat())

        solution_i[idx_missing, "dim$j"] = etp.(solution_i.lambda[idx_missing])
    end

    #Impute.locf!(solution_i)
end

"""
    plot_path(x; α=0.1, gt=ones(Int64, size(x,1)), force2dim=true, show::Bool, savefig=false, fname="path_plot", verbose=false)

plot clusterpath with the data(`x`).
If the dimension of `x` is greater than 2, only plot a combination of first two dimensions.

***`Gaston.jl` and `gnuplot` should be installed and on the PATH of your system. Install gnuplot [here](https://sourceforge.net/projects/gnuplot/files/gnuplot/).***

* `x`: the data to cluster.
* `α`: threshold for BMT-clusterpath
* `gt`: the ground truth labels.
* `force2dim`: force dim>2 data to 2-dim data (default: `true`)
* `show`: whether to show the plot in the notebook. **Highly recommended not to show if the number of samples is large.**
* `savefig`: whether to save the figure as a PNG file. (default: `false`)
* `fname`: image file name to be used when `savefig` is `true`. (default: `"path_plot"`)
"""
function plot_path(x; α=0.1, gt=ones(Int64, size(x,1)), force2dim=true, show::Bool, savefig=false, fname="path_plot", verbose=false)
    n_dim = length(size(x)) == 1 ? 1 : size(x, 2)
    if n_dim > 2 && !force2dim
        x = x[:, 1:3]
    elseif (n_dim > 2 && force2dim) || n_dim == 2
        x = x[:, 1:2]
    end
    solution = cast_solution(x, α=α, silence_warning=true)
#     n_plot = Int64(n_dim * (n_dim-1) / 2)
#     if n_plot > 6
#         n_plot = 6
#         n_dim = 4
#     end

    # initialize plot
#     if n_dim > 2
#         p = Plots.plot(layout=(1,n_plot), size=(400*n_plot, 400), legend=false)
#         counter = 0
#             for i=1:n_dim
#                 for j=(i+1):n_dim
#                     counter += 1
#                     Plots.plot!(p[counter], view(x, :, i), view(x, :, j), seriestype=:scatter, color=gt, alpha=.2)
#                 end
#             end
    if n_dim > 2 && !force2dim
        #p = Plots.plot(view(x, :, 1), view(x, :, 2),
        #               #series_annotations = text.(1:size(x, 1), :bottom),
        #               seriestype=:scatter, color=gt, alpha=.2, legend=false)
        clst = unique(gt)
        p = Gaston.scatter3(
                x[:, 1], x[:, 2], x[:, 3], supp=gt .% 8, variable="using 1:2:3:(4)",
                pointtype="ecircle", lw=0.75, lc="palette",
                Gaston.Axes(key="noautotitle",
                            grid=:on,
                            style="fill transparent solid 0.005;
                                   set style circle radius screen 0.005;
                                   set border 21;
                                   unset colorbox;
                                   set auto fix;
                                   set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
                                   set cbrange [0:7];
                                   set size square 1,1;
                                   set format y \"\"; set format x \"\"; set format z \"\"")
            )
#         p = Gaston.plot(
#                 x[:, 1], x[:, 2], z=x[:, 3], supp=gt.%8, variable="using 1:2:3:4", w=:circles, lw=0.75, lc="palette",
#                 Gaston.Axes(key="noautotitle",
#                             style="fill transparent solid 0.005;
#                                    set style circle radius screen 0.005;
#                                    set border 3;
#                                    unset colorbox;
#                                    set auto fix;
#                                    set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
#                                    set cbrange [0:7];
#                                    set size square 1,1;
#                                    set format y \"\"; set format x \"\"")
#             )
    elseif (n_dim > 2 && force2dim) || n_dim == 2
        #p = Plots.plot(view(x, :, 1), view(x, :, 2),
        #               #series_annotations = text.(1:size(x, 1), :bottom),
        #               seriestype=:scatter, color=gt, alpha=.2, legend=false)
        clst = unique(gt)
        p = Gaston.plot(
                x[:, 1], x[:, 2], supp=gt.%8, variable="using 1:2:3", w=:circles, lw=0.75, lc="palette",
                Gaston.Axes(key="noautotitle",
                            grid=:on,
                            style="fill transparent solid 0.005;
                                   set style circle radius screen 0.005;
                                   set border 3;
                                   unset colorbox;
                                   set auto fix;
                                   set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
                                   set cbrange [0:7];
                                   set size square 1,1;
                                   set format y \"\"; set format x \"\"")
            )
        #if length(clst) > 1
        #    for cl in clst[2:end]
        #        Gaston.plot!(x[gt .== cl, 1], x[gt .== cl, 2], w=:circles, lw=0.75)
        #    end
        #end
    else
        n = size(x, 1)
        #p = Plots.plot(zeros(n), x, seriestype=:scatter, color=gt, alpha=.2, legend=false)

        ## USE GASTON (GNUPLOT) FOR FASTER PLOTTING
        clst = unique(gt)
        p = Gaston.plot(
                zeros(n), x, supp=gt.%8, variable="using 1:2:3", w=:circles, lw=0.75, lc="palette",
                Gaston.Axes(key="noautotitle",
                            grid=:on,
                            style="fill transparent solid 0.005;
                                   set style circle radius screen 0.01;
                                   set border 3;
                                   unset colorbox;
                                   set auto fix;
                                   set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
                                   set cbrange [0:7];
                                   set size square 1,1;
                                   set format y \"\"; set format x \"\"")
            )
        #if length(clst) > 1
        #    for cl in clst[2:end]
        #        Gaston.plot!(zeros(n)[gt .== cl], x[gt .== cl], w=:circles, lw=0.75)
        #    end
        #end
    end

    # draw paths, one at a time
    i = 1
    for solution_i in solution
        _prepare!(solution_i; for_path=true)
        #println(solution_i)

#         if n_dim > 2
#             counter = 0
#             for i=1:n_dim
#                 for j=(i+1):n_dim
#                     counter += 1
#                     Plots.plot!(p[counter],
#                                 view(solution_i, "dim$i"), view(solution_i, "dim$j"),
#                                 seriestype=:path, color="black", alpha=0.25, legend=false)
#                 end
#             end
        if n_dim > 2 && !force2dim
            verbose && (i % 500 == 0) && print("$(i) ")

            #Plots.plot!(p, view(solution_i, "dim1"), view(solution_i, "dim2"),
            #            seriestype=:path, color="black", alpha=0.25, legend=false)
            Gaston.scatter3!(view(solution_i, "dim1"), view(solution_i, "dim2"), view(solution_i, "dim3"),
                            w=:lines, lw=0.05, lc="black")
#             Gaston.plot!(view(solution_i, "dim1"), view(solution_i, "dim2"), z=view(solution_i, "dim3"),
#                          w=:lines, lw=0.05, lc="black")
        elseif (n_dim > 2 && force2dim) || n_dim == 2
            verbose && (i % 500 == 0) && print("$(i) ")

            #Plots.plot!(p, view(solution_i, "dim1"), view(solution_i, "dim2"),
            #            seriestype=:path, color="black", alpha=0.25, legend=false)
            Gaston.plot!(view(solution_i, "dim1"), view(solution_i, "dim2"),
                         w=:lines, lw=0.05, lc="black")
        else
            verbose && (i % 500 == 0) && print("$(i) ")

            #Plots.plot!(p, view(solution_i, "lambda"), view(solution_i, "dim1"),
            #            seriestype=:path, color="black", alpha=0.25, legend=false)
            Gaston.plot!(view(solution_i, "lambda"), view(solution_i, "dim1"),
                         w=:lines, lw=0.1, lc="black")
        end

        i += 1
    end

    if savefig
        #png(p, fname*".png")
        Gaston.save(term = "png",
                    output= "$(fname).png",
                    size = "500,500",
                    background = "white")
    end

    if show
        return(p)
    end
end

"""
    assign_cluster(x; α=0.1, verbose=false)

assign cluster to each of the observations in `x`.
returns an array of length=size(x, 1) of cluster indices.

* `x`: data
* `α`: threshold for BMT-clusterpath
"""
function assign_cluster(x; α=0.1, verbose=false)

    n = size(x, 1)
    p = length(size(x)) == 1 ? 1 : size(x, 2)

    final = Array{Float64,2}(undef, n, p)

    clst = repeat([""], n)
    for i=1:p
        splits = clusterpath(x[:, i], α=α, return_split=true, silence_warning=true)["splits"]
        clst_dim = zeros(Int64, n)
        for s in splits
            clst_dim .+= Int64.(x[:, i]  .> s)
        end
        clst .*= string.(clst_dim)
    end

    DataFrames.levelcode.(DataFrames.CategoricalArray(clst))
end
# function assign_cluster!(x, solution; n_node=1, verbose=false)

#     n = size(x, 1)
#     p = length(size(x)) == 1 ? 1 : size(x, 2)

#     final = Array{Float64,2}(undef, n, p)
#     lambda = nothing

#     i = 1
#     for solution_i in solution
#         if verbose && (i % 500 == 0)
#             print("$(i) ")
#         end
#         _prepare!(solution_i; for_path=false, n_node=n_node)

#         #println(solution_i[end, 3:end])
#         final[i, :] = collect(solution_i[end-n_node, 3:end])
#         isnothing(lambda) && (lambda = solution_i.lambda[end-n_node])
#         i += 1
#     end

#     centroids = unique(final, dims=1)
#     n_centroids = size(centroids, 1)

#     idx_x = Array{Int64}(undef, n)

#     Threads.@threads for i=1:n_centroids
#         filter = vec(reduce(&, final .== centroids[i, :]', dims=2))
#         idx_x[filter] .= i
#     end

#     idx_x, lambda
# end

"""
    plot_cluster(x; α=0.1, n_node=1, show::Bool, savefig=false, fname="plot_clst", verbose=false)

Plots the scatter plot of the data `x` colored according to the cluster assigned by clusterpath algorithm.
If the dimension of `x` is greater than 2, perform PCA and plot two PCs.

***`Gaston.jl` and `gnuplot` should be installed and on the PATH of your system. Install gnuplot [here](https://sourceforge.net/projects/gnuplot/files/gnuplot/). ***

* `x`: data
* `α`: threshold for BMT-clusterpath
* `n_node`: if greater than `1`, will assign clusters from previous merge status. (default: `1`)
* `show`: whether to show the figure.
* `savefig`: whether to save the figure as a png file. (default: `false`)
* `fname`: file name to save if `savefig` is true. (default: `"plot_clst"`)
* `verbose`: print out current iteration. (default: `false`)
"""
function plot_cluster(x; α=0.1, force2dim=true, show::Bool, savefig=false, fname="plot_clst", verbose=false)
    clst = assign_cluster(x; α=α, verbose=verbose)

    #solution = cast_solution(x, α=α)
    #solution_1 = take!(solution)

    k = levels(clst)[end]

    p = size(x, 2)

    if p > 3
        force2dim && (pca = fit(PCA, Array(x'); maxoutdim=2))
        !force2dim && (pca = fit(PCA, Array(x'); maxoutdim=3))
        x = MultivariateStats.transform(pca, x')'
    elseif p == 3
        force2dim && (pca = fit(PCA, Array(x'); maxoutdim=2); x = MultivariateStats.transform(pca, x')')
    end

    # use Gaston for plotting
    #uniq_clusters = unique(clst)
    #cl = uniq_clusters[1]
    #idx = clst .== cl
    if p > 2 && !force2dim
        plt = Gaston.scatter3(
                x[:, 1], x[:, 2], x[:, 3], supp=clst .% 8, variable="using 1:2:3:(4)",
                pointtype="ecircle", lw=0.75, lc="palette",
                Gaston.Axes(key="noautotitle",
                            grid=:on,
                            style="fill transparent solid 0.005;
                                   set style circle radius screen 0.005;
                                   set border 21;
                                   unset colorbox;
                                   set auto fix;
                                   set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
                                   set cbrange [0:7];
                                   set size square 1,1;
                                   set format y \"\"; set format x \"\"; set format z \"\"")
            )
    elseif (p > 2 && force2dim) || p == 2
        plt = Gaston.plot(
            x[:, 1], x[:, 2], supp=clst .% 8, variable="using 1:2:3",
            w=:circles, lw=0.75, lc="palette",
            Gaston.Axes(key="noautotitle",
                        grid=:on,
                        xlabel="\"k=$(k)\"",
                        style="fill transparent solid 0.005;
                               set style circle radius screen 0.005;
                               set border 3;
                               unset colorbox;
                               set auto fix;
                               set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
                               set cbrange [0:7];
                               set size square 1,1;
                               set format y \"\"; set format x \"\"")
        )
    else
        plt = Gaston.plot(
            x, randn(length(x)) * 0.02, supp=clst .% 8, variable="using 1:2:3",
            w=:circles, lw=0.75, lc="palette",
            Gaston.Axes(key="noautotitle",
                        grid=:on,
                        xlabel="\"k=$(k)\"",
                        style="fill transparent solid 0.005;
                               set style circle radius screen 0.005;
                               set border 3;
                               set yrange [-1:1];
                               unset colorbox;
                               set auto fix;
                               set palette defined ( 0 '#E41A1C', 1 '#377EB8', 2 '#4DAF4A', 3 '#984EA3', 4 '#FF7F00', 5 '#d2d200', 6 '#A65628', 7 '#F781BF' );
                               set cbrange [0:7];
                               set size square 1,1;
                               set format y \"\"; set format x \"\"")
        )
    end

#     if length(uniq_clusters) > 1
#         for cl in uniq_clusters[2:end]
#             idx = clst .== cl
# #             println(sum(idx))
#             Gaston.plot!(x[idx, 1], x[idx, 2], w=:circles, lw=0.75)
#         end
#     end

    if savefig
        Gaston.save(term = "png",
                    output= "$(fname).png",
                    size = "500,500",
                    background = "white")
    end

#     plot(x[:, 1], x[:, 2], color=clst,
#          seriestype=:scatter, legend=false, alpha=.2,
#          title="lambda=$(round(λ, digits=3)), k=$(k)")
    if show
        return(plt)
    end
end

function count_final_cluster(x; α=0.1)
    levels(assign_cluster(x, α=α))[end]
end

# ---- population procedure ---- #

"""
    pnorm(x)

Julia equivalent of R's `pnorm`
"""
function pnorm(x)
    cdf(Normal(), x)
end

"""
    dnorm(x)

Julia equivalent of R's `dnorm`
"""
function dnorm(x)
    pdf(Normal(), x)
end

# """
#     cond_mean_on_LR(L::T, R::T, a1::T, a2::T, a3::T, μ1::T, μ2::T, μ3::T) where T <: Real

# Conditional mean on (L, R), defined as μ_{L,R} = (∫_L^R f(x) dx)^(-1) / (∫_L^R x f(x) dx)
# """
# function cond_mean_on_LR(L::T, R::T, a1::T, a2::T, a3::T, μ1::T, μ2::T, μ3::T) where T <: Real
#     pL = a1 * pnorm(L-μ1) + a2 * pnorm(L-μ2) + a3 * pnorm(L-μ3)
#     pR = a1 * pnorm(R-μ1) + a2 * pnorm(R-μ2) + a3 * pnorm(R-μ3)

#     ∫_L = -(a1 * dnorm(L-μ1) + a2 * dnorm(L-μ2) + a3 * dnorm(L-μ3)) +
#           a1 * μ1 * pnorm(L-μ1) + a2 * μ2 * pnorm(L-μ2) + a3 * μ3 * pnorm(L-μ3)
#     ∫_R = -(a1 * dnorm(R-μ1) + a2 * dnorm(R-μ2) + a3 * dnorm(R-μ3)) +
#           a1 * μ1 * pnorm(R-μ1) + a2 * μ2 * pnorm(R-μ2) + a3 * μ3 * pnorm(R-μ3)

#     μ_LR = (∫_R - ∫_L)/(pR - pL)

#     (pL, pR, ∫_L, ∫_R, μ_LR)
# end

"""
    cond_mean_on_LR(L, R, a1, a2, a3, μ1, μ2, μ3)

Conditional mean on (L, R), defined as μ_{L,R} = (∫_L^R f(x) dx)^(-1) / (∫_L^R x f(x) dx)
"""
function cond_mean_on_LR(L, R, a1, a2, a3, μ1, μ2, μ3)
    # (∫_l^r f(x) dx) = pR - pL
    pL = a1 * pnorm(L.-μ1) + a2 * pnorm(L.-μ2) #+ a3 * pnorm(L.-μ3)
    pR = a1 * pnorm(R.-μ1) + a2 * pnorm(R.-μ2) #+ a3 * pnorm(R.-μ3)

    # (∫_l^r x f(x) dx) = ∫_R - ∫_L
    ∫_L = -(a1 * dnorm(L.-μ1) + a2 * dnorm(L.-μ2)) + # + a3 * dnorm(L.-μ3)) +
          a1 * μ1 * pnorm(L.-μ1) + a2 * μ2 * pnorm(L.-μ2) #+ a3 * μ3 * pnorm(L.-μ3)
    ∫_R = -(a1 * dnorm(R-μ1) + a2 * dnorm(R-μ2)) + # + a3 * dnorm(R-μ3)) +
          a1 * μ1 * pnorm(R-μ1) + a2 * μ2 * pnorm(R-μ2) #+ a3 * μ3 * pnorm(R-μ3)
    ## adding a3/μ3 terms yields inaccurate result.

    # population conditional mean on (L, R)
    μ_LR = (∫_R .- ∫_L) ./ (pR .- pL)

    (pL, pR, ∫_L, ∫_R, μ_LR)
end

"""
    find_split(d, a1, a2, μ1, μ2, L, R;
               a3=0., μ3=0., find_deltas=false, find_split=!find_deltas)

Split the cluster (L, R) into subclusters that maximizes G_{L,R}.
Returns centered conditional mean (deltas) of the left and right subcluster
if `find_deltas` is true. Otherwise return the split point.
"""
function find_split(d, a1, a2, μ1, μ2, L, R;
                    a3=0., μ3=0., find_deltas=false, find_split=!find_deltas)
    # p, int, int_en evaluated at `d`
    ∫_d = -(a1 * dnorm(d.-μ1) + a2 * dnorm(d.-μ2)) + # + a3 * dnorm(d.-μ3)) +
          a1 * μ1 * pnorm(d.-μ1) + a2 * μ2 * pnorm(d.-μ2) #+ a3 * μ3 * pnorm(d.-μ3)
    p_d = a1 * pnorm(d.-μ1) + a2 * pnorm(d.-μ2) #+ a3 * pnorm(d.-μ3)

    pL, pR, ∫_L, ∫_R, μ_LR = cond_mean_on_LR(L, R, a1, a2, a3, μ1, μ2, μ3)

    # find split point with normlized G'_{L,R} (=h)(supp. p.32)
    h = μ_LR .* (p_d .- pL).^2 + (pL + pR .- 2p_d) .* (∫_d .- ∫_L) - d .* (p_d .- pL) .* (pR .- p_d)
    split = d[2:end][h[2:end] .* h[1:end-1] .< 0]

    if find_deltas && (sum(μ1 .< split .< μ2) == 0)
        return(1., 1.)
    end

    # split point
    s = max(split[μ1 .< split .< μ2]...)

    if find_split
        return(s)
    else
        _, _, _, _, δ_1 = cond_mean_on_LR(L, s, a1, a2, a3, μ1, μ2, μ3)
        δ_1 -= (L + s)/2  # centering

        _, _, _, _, δ_2 = cond_mean_on_LR(s, R, a1, a2, a3, μ1, μ2, μ3)
        δ_2 -= (s + R)/2  # centering

        return(δ_1, δ_2)
    end

end

"""
    find_truncation(a1, a2, μ1, μ2,
                    R_ub=(μ2 + 8), L_lb=(μ1 - 8),
                    tuning_par1=0.01, tuning_par2=0.005,
                    a3=0, μ3=0, σ=1)

Find population right-truncation points for bimodal Normal mixtures.
Inputs:
    a1, a2, a3: weights for Normal distributions. a1 + a2 + a3 should be == 1
    μ1, μ2, μ3: means for normals μ1 = -μ2

Possible returns:
    0: negligible result (not a symmetric problem)
    -1: no truncation
"""
function find_truncation(a2, μ2;
                         a1=1-a2, μ1=-μ2, R_ub=μ2+8., L_lb=μ1-8.,
                         tuning_par1=0.01, tuning_par2=0.005,
                         a3=0., μ3=0., σ=1.)

    r = Array(R_ub:(-tuning_par1):μ1)

    # if the bimodal distn is symmetric, and a1 == a2,
    # then truncate both ends equally
    if a1 == a2 && μ2 == -μ1
        for i=1:length(r)

            # truncate both left and right with the same amount
            R = r[i]
            L = -R

            d = Array(L:0.01:R)
            δ_1, δ_2 = find_split(d, a1, a2, μ1, μ2, L, R; find_deltas=true)

            # L^* is the smallest L that satisfies δ_1 < 0 < δ_2. (supp. p.30)
            # Thus returns right-trunctation point
            if δ_1 < 0 < δ_2
                return(R)
            end
        end
    end

    # possible cases:
    #   a1 != a2 && μ2 != -μ1
    #   a1 != a2 && μ2 == -μ1
    if a1 != a2
        L_old = -1000
        for i=1:length(r)
            R = r[i]
            L = Array{Float64,1}(L_old:tuning_par2:(R-0.05))

            pL, pR, ∫_L, ∫_R, μ_LR = cond_mean_on_LR(L, R, a1, a2, a3, μ1, μ2, μ3)

            # locate R_L using the equation μ_{L,R_L} - (L + R_L)/2.  (supp. p.30)
            temp = sign.(μ_LR .- (L .+ R)./2)

            if sum(temp[2:end] .* temp[1:end-1] .< 0) ≠ 0
                # this is L^* from (supp. p.28)
                L = L[2:end][temp[2:end] .* temp[1:end-1] .< 0][1]

                d = Array(L:0.01:R)
                δ_1, δ_2 = find_split(d, a1, a2, μ1, μ2, L, R; find_deltas=true)

                pL = a1 * pnorm(L-μ1) + a2 * pnorm(L-μ2) + a3 * pnorm(L-μ3)
                pR = a1 * pnorm(R-μ1) + a2 * pnorm(R-μ2) + a3 * pnorm(R-μ3)

                if sign((R-L)/(pR-pL) * (a1 * dnorm(L-μ1) + a2 * dnorm(L-μ2) + a3 * dnorm(L-μ3)) - 1) ≠ -1
                    println("🤔🤔 input is not good")
                end

                # L^* is the smallest L that satisfies δ_1 < 0 < δ_2. (supp. p.30)
                # Thus returns right-trunctation point
                if δ_1 < 0 < δ_2 && L_old ≤ L

                    # if L^* is achieved at i==1, it means there is no truncation.
                    if i == 1
                        return(-1)
                    end

                    return(R)
                end

                # increment L
                if L_old ≤ L
                    L_old = L
                end

                # if L > 0, it means roughly that L > m_ (m_≤0)  (supp. p.29 table and p.28)
                if L > 0
                    return(-1)
                end

            end
        end
    end

    # If a1 == a2 and μ1 != -μ2, then return 0
    # : do not need to assess in this case, since it can WLOG be reduced to a symmetric case)
    return(0)
end

"""
    clusterpath_pop(a2::T, μ2::T; a1=1-a2, μ1=-μ1, a3=0., μ3=0, σ=1) where T <: Real

Inputs:
    a2: weight on 2nd normal distribution
    μ2: mean of 2nd normal distribution
Returns:
    A dictionary of weights (`a1`, `a2`), means (`μ1`, `μ2`) and
    description of truncation (`L*`, `R*`) and split point (`s`).
"""
function clusterpath_pop(a2::T1, μ2::T2; a1=1-a2, μ1=-μ2, a3=0., μ3=0, σ=1.) where T1 <: Real where T2 <: Real
    # search right truncation point
    R = find_truncation(a2, μ2; a1=a1, μ1=μ1, a3=a3, μ3=μ3, σ=σ)

    # R == 0 or R == -1
    if R == 0 || R == -1
        d = Array(μ1:1e-04:μ2)
        # d[argmin(abs.(a1 * dnorm(d.-μ1) .- a2 * dnorm(d.-μ2)))]  ## ???????
        # d[argmin(a1 * dnorm(d.-μ1) .+ a2 * dnorm(d.-μ2))]

        return(Dict("a1" => a1, "a2" => a2,
                    "mu1" => μ1, "mu2" => μ2,
                    "s" => NaN, "L*" => NaN, "R*" => NaN))
    end

    # If R is a proper right-truncation point, then find split point
    L_lb = -3 * max(R, 3)  # lower bound of searching grid of L
    L_ub = R - 0.01  # upper bound
    L = Array(L_lb:5e-04:L_ub)

    _, _, _, _, μ_LR = cond_mean_on_LR(L, R, a1, a2, a3, μ1, μ2, μ3)
    δs = sign.(μ_LR .- (L .+ R)/2)

    # L^* is the smallest L that satisfies δ_1 < 0 < δ_2
    L = L[2:end][δs[2:end] .* δs[1:end-1] .< 0][1]

    d = Array(L:0.001:R)
    split = find_split(d, a1, a2, μ1, μ2, L, R)

    return(Dict("a1" => a1, "a2" => a2,
                "mu1" => μ1, "mu2" => μ2,
                "s" => split, "L*" => L, "R*" => R))
end

# ---- simulation data generation ---- #

"""
    generate_mixture_normal(n, m, p)

generate `n` observations from mixture of univariate normals
each with standard deviation 1 and mean parameters `m` and proportion `p`.
"""
function generate_mixture_normal(n::Int64, m::Array{T1,1}=[-20., 0., 20.], p::Array{T2,1}=ones(length(m))/length(m)) where T1 <: Real where T2 <: Real
    @assert sum(p) ≈ 1
    mixed_normal = Array{Float64}(undef, 0)  # placeholder for the resulting mixture
    n_mixtures = floor.(Int64, n * p)  # sample sizes for each univariate normals

    # leftover n's is assigned to the first distribution > sum(n_mixtures)
    n_mixtures[1] += (n - sum(n_mixtures))

    # sample from Normal(0, 1)
    for i=1:length(m)
        mixed_normal = [mixed_normal; randn(n_mixtures[i]) .+ m[i]]
    end

    mixed_normal
end

end  # end module


sparsity(B) = Float32(nnz(B)/length(B))
sparsity(B::Matrix) = 1.0f0
sparsity(B::Adjoint) = sparsity(B.parent)

function preconditioner(B, s, tol=1e-2)
    upsample(droptol!(sparse(inv(downsample2(B, s)))/s/s, tol), s)
end

#downsample(x, s = 10) = [mean(x[i:i+s-1,j:j+s-1]) for i in 1:s:size(x,1)-s, j in 1:s:size(x, 2)-s]
downsample(x, s = 10) = [mean(x[i:i+s-1,j:j+s-1]) for i in 1:s:size(x,1)-s+1, j in 1:s:size(x, 2)-s+1]
upsample(x, s = 10) = repeat(x, inner=(s, s))


m1 = fit(LinearModel, form, data)
λlm = dispersion(m1)

loss1 = RobustModels.L2Loss()
loss2 = RobustModels.TukeyLoss()
est1 = MEstimator(loss1)
est2 = MEstimator(loss2)

funcs = ( dof,
          dof_residual,
          confint,
          deviance,
          nulldeviance,
          loglikelihood,
          nullloglikelihood,
          dispersion,
          nobs,
          stderror,
          vcov,
          residuals,
          predict,
          response,
          weights,
          workingweights,
          fitted,
          predict,
          isfitted,
          islinear,
          leverage,
          modelmatrix,
          projectionmatrix,
          Estimator,
          scale
        )

@testset "linear: L2 estimator" begin
    println("\n\t\u25CF Estimator: L2")

    # OLS
    println(m1)
    println(" lm              : ", coef(m1))

    # Formula, dense and sparse entry  and methods :cg and :chol
    @testset "(type, method): ($(typeof(A)),\t$(method))" for (A, b) in ((form, data), (X, y), (sX, y)), method in (:cg, :chol)
        name  = if A==form; "formula" elseif A==X; "dense  " else "sparse " end
        name *= if method==:cg; ",  cg" else ",chol" end
        # use the dispersion from GLM to ensure that the loglikelihood is correct
        m = fit(RobustLinearModel, A, b, est1; method=method, verbose=false, initial_scale=λlm)
        β = copy(coef(m))
        println("rlm($name): ", β)
        @test_nowarn println(m)
        @test isapprox(coef(m1), β; rtol=1e-5)

        # refit
        refit!(m, y; verbose=false, initial_scale=:extrema)
        println("$m")
        @test all(coef(m) .== β)

        # make sure the methods for RobustLinearModel are well defined
        @testset "method: $(f)" for f in funcs
            if f in (Estimator, weights, workingweights, islinear, isfitted,
                     leverage, modelmatrix, projectionmatrix, scale)
                # method is not defined in GLM
                @test_nowarn f(m)
            else
                var = f(m1)
                robvar = f(m)
                if isa(var, Union{AbstractArray, Tuple})
                    if f != vcov
                        @test isapprox(var, robvar; rtol=1e-4)
                    end
                else
                    if f == dof
                        @test var == robvar + 1
                    elseif f == dof_residual
                        @test var == robvar
                    elseif f in (deviance, nulldeviance)
                        s = scale(m)
                        ## the deviance in RobustModels is the scaled deviance (divided by
                        ## the squared dispersion), whereas in GLM it is not.
                        @test isapprox(var, robvar * s^2; rtol=1e-4)
                    elseif f in (loglikelihood, nullloglikelihood)
                        ## TODO: should work
                        @test_broken isapprox(var, robvar; rtol=1e-4)
#                        @test isapprox(var + log(λ), (robvar + RobustModels.fullloglikelihood(r)) * s^2/λ^2 - log(RobustModels.estimator_norm(r.est)); rtol=1e-4)
#                        @test isapprox(var + log(dispersion(m1)), robvar + log(s); rtol=1e-4)
                    else
                        @test isapprox(var, robvar; rtol=1e-4)
                    end
                end
            end
        end

        # TauEstimator interface
        m3 = fit(RobustLinearModel, A, b, TauEstimator{TukeyLoss}(); method=method, initial_scale=λlm)
        @test_nowarn tauscale(m3)

        # later fit!
        m2 = fit(RobustLinearModel, A, b, est1; method=method, dofit=false)
        @test all(0 .== coef(m2))
        fit!(m2; verbose=false, initial_scale=:mad)
        @test all(β .== coef(m2))
    end
end


include("../src/ccsa.jl")
using LinearAlgebra
using CatViews
using BenchmarkTools
using Profile
using ProfileView
using PProf
a1=2;b1=0;a2=-1;b2=1;
function f(x)
    return [sqrt(x[2]),(a1*x[1]+b1)^3-x[2],(a2*x[1]+b2)^3-x[2]]
end
function ∇f(x) #(m+1)*n
    return [0 1/2/x[2]; 3*a1*(a1*x[1]+b1)^2 -1; 3*a2*(a2*x[1]+b2)^2 -1]
end
function f_and_∇f(x)
    return f(x),∇f(x)
end
n=2
m=2
σ=[30.0, 30.0] #n
ρ=[100.0, 100.0, 100.0]
x=[1.234, 5.678]
lb=[-Inf, 0.]
opt=CCSAState(n,m,f_and_∇f,ρ,σ,x)
opt.lb=lb
opt.xtol_rel = 1e-4
@time iters=optimize(opt)
ProfileView.@profview optimize(opt)


using miniWeather
using Test

@testset "test mass and energy conservation" begin
    nx = 100
    nz = 50
    sim_time = 400
    out_freq = 400
    weather_type = "thermal"

    model, grid = init(nx, nz, sim_time, weather_type);
    mass0, te0 = reductions(model, grid)
    run!(model, grid, out_freq, "output.nc")
    mass, te = reductions(model, grid)

    @test (mass - mass0) / mass0 < 1.0e-9
    @test (te - te0) / te0 < 4.5e-5

end


module ReactiveMPMessageTest

using Test
using ReactiveMP
using Distributions
using Random

import InteractiveUtils: methodswith
import Base: methods
import Base.Iterators: repeated, product
import ReactiveMP: materialize!
import ReactiveMP: mirrorlog, xtlog
import SpecialFunctions: loggamma

@testset "Message" begin

    @testset "Default methods" begin
        data    = PointMass(1)

        for clamped in [ true, false ], initial in [ true, false ]
            msg = Message(data, clamped, initial)
            @test getdata(msg)    === data
            @test is_clamped(msg) === clamped
            @test is_initial(msg) === initial
            @test materialize!(msg) === msg
            @test occursin("Message", repr(msg))
        end

    end

    @testset "multiply_messages" begin
        dist1 = NormalMeanVariance(randn(), rand())
        dist2 = NormalMeanVariance(randn(), rand())

        @test getdata(Message(dist1, false, false) * Message(dist2, false, false))    == prod(ProdAnalytical(), dist1, dist2)
        @test getdata(Message(dist2, false, false) * Message(dist1, false, false))    == prod(ProdAnalytical(), dist2, dist1)

        for (left_is_initial, right_is_initial) in product(repeated([ true, false ], 2)...)
            @test is_clamped(Message(dist1, true, left_is_initial) * Message(dist2, false, right_is_initial)) == false
            @test is_clamped(Message(dist1, false, left_is_initial) * Message(dist2, true, right_is_initial)) == false
            @test is_clamped(Message(dist1, true, left_is_initial) * Message(dist2, true, right_is_initial))  == true
            @test is_clamped(Message(dist2, true, left_is_initial) * Message(dist1, false, right_is_initial)) == false
            @test is_clamped(Message(dist2, false, left_is_initial) * Message(dist1, true, right_is_initial)) == false
            @test is_clamped(Message(dist2, true, left_is_initial) * Message(dist1, true, right_is_initial))  == true
        end

        for (left_is_clamped, right_is_clamped) in product(repeated([ true, false ], 2)...)
            @test is_initial(Message(dist1, left_is_clamped, true) * Message(dist2, right_is_clamped, true)) == !(left_is_clamped && right_is_clamped)
            @test is_initial(Message(dist2, left_is_clamped, true) * Message(dist1, right_is_clamped, true)) == !(left_is_clamped && right_is_clamped)
            @test is_initial(Message(dist1, left_is_clamped, false) * Message(dist2, right_is_clamped, false)) == false
            @test is_initial(Message(dist2, left_is_clamped, false) * Message(dist1, right_is_clamped, false)) == false
        end

        @test is_initial(Message(dist1, true, true) * Message(dist2, true, true)) == false
        @test is_initial(Message(dist1, true, true) * Message(dist2, true, false)) == false
        @test is_initial(Message(dist1, true, false) * Message(dist2, true, true)) == false
        @test is_initial(Message(dist1, false, true) * Message(dist2, true, false)) == true
        @test is_initial(Message(dist1, true, false) * Message(dist2, false, true)) == true
        @test is_initial(Message(dist2, true, true) * Message(dist1, true, true)) == false
        @test is_initial(Message(dist2, true, true) * Message(dist1, true, false)) == false
        @test is_initial(Message(dist2, true, false) * Message(dist1, true, true)) == false
        @test is_initial(Message(dist2, false, true) * Message(dist1, true, false)) == true
        @test is_initial(Message(dist2, true, false) * Message(dist1, false, true)) == true
    end

    @testset "Statistics" begin

        distributions = [
            PointMass(0.5),
            Gamma(10.0, 2.0),
            NormalMeanVariance(-10.0, 10.0),
            Wishart(4.0, [ 2.0 -0.5; -0.5 1.0 ]),
            MvNormalMeanPrecision([ 2.0, -1.0 ], [ 7.0 -1.0; -1.0 3.0 ]),
            Bernoulli(0.5),
            Categorical([ 0.8, 0.2 ])
        ]

        # Here we get all methods defined for a particular type of a distribution
        dists_methods = map(d -> methodswith(eval(nameof(typeof(d)))), distributions)

        methods_to_test = [
            Distributions.mean,
            Distributions.median,
            Distributions.mode,
            Distributions.shape,
            Distributions.scale,
            Distributions.rate,
            Distributions.var,
            Distributions.std,
            Distributions.cov,
            Distributions.invcov,
            Distributions.logdetcov,
            Distributions.entropy,
            Distributions.params,
            Base.precision,
            Base.length,
            Base.ndims,
            Base.size,
            mean_cov,
            mean_invcov,
            mean_precision,
            weightedmean_cov,
            weightedmean_invcov,
            weightedmean_precision,
            probvec,
            weightedmean
        ]

        for (distribution, distribution_methods) in zip(distributions, dists_methods), method in methods_to_test
            T       = typeof(distribution)
            message = Message(distribution, false, false)
            # Here we check that a specialised method for a particular type T exist
            ms = methods(method, (T, ))
            if !isempty(ms) && all(m -> m ∈ distribution_methods, ms)
                @test method(message) == method(distribution)
            end
        end

        fn_mean_functions = (inv, log, xtlog, mirrorlog, loggamma)

        for distribution in distributions, fn_mean in fn_mean_functions
            F       = typeof(fn_mean)
            T       = typeof(distribution)
            message = Message(distribution, false, false)
            # Here we check that a specialised method for a particular type T exist
            ms = methods(mean, (F, T, ), ReactiveMP)
            if !isempty(ms)
                @test mean(fn_mean, message) == mean(fn_mean, distribution)
            end
        end

        _getpoint(rng, distritubution) = _getpoint(rng, variate_form(distritubution), distritubution)
        _getpoint(rng, ::Type{ <: Univariate }, distribution) = 10rand(rng)
        _getpoint(rng, ::Type{ <: Multivariate }, distribution) = 10 .* rand(rng, 2)

        distributions2   = [
            Gamma(10.0, 2.0),
            NormalMeanVariance(-10.0, 1.0),
            MvNormalMeanPrecision([ 2.0, -1.0 ], [ 7.0 -1.0; -1.0 3.0 ]),
            Bernoulli(0.5),
            Categorical([ 0.8, 0.2 ])
        ]

        methods_to_test2 = [
            Distributions.pdf,
            Distributions.logpdf,
        ]

        rng = MersenneTwister(1234)

        for distribution in distributions2, method in methods_to_test2
            message = Message(distribution, false, false)

            for _ in 1:3
                point = _getpoint(rng, distribution)
                @test method(message, point) === method(distribution, point)
            end
        end

    end

end

end


#=
# 3D Poisson, Dirichlet bc
# Uses Dendro imported as a custom gen target
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("dendrotest");

useLog("dendrotestlog", level=3)

# default values (max_depth=6, wavelet_tol = 0.1, partition_tol = 0.3, solve_tol = 1e-6, max_iters = 100)
generateFor(DENDRO, params=(8, 0.001, 0.3, 0.000001, 500))

domain(3)
functionSpace(order=2)

u = variable("u")
testSymbol("v")

boundary(u, 1, DIRICHLET, "0")

# Write the weak form
coefficient("k1", "1.1")
coefficient("k2", "0.1")
coefficient("k3", "0.1")
coefficient("k4", "0.1")
coefficient("k5", "0.1")
f = coefficient("f", "-1.1*14*pi*pi*sin(3*pi*x)*sin(2*pi*y)*sin(pi*z) + 0.1*3*pi*cos(3*pi*x)*sin(2*pi*y)*sin(pi*z) + 0.1*2*pi*sin(3*pi*x)*cos(2*pi*y)*sin(pi*z) + 0.1*pi*sin(3*pi*x)*sin(2*pi*y)*cos(pi*z) - 0.1*sin(3*pi*x)*sin(2*pi*y)*sin(pi*z)")
weakForm(u, "-k1*dot(grad(u),grad(v)) + k2*deriv(u,1)*v + k3*deriv(u,2)*v + k4*deriv(u,3)*v - k5*u*v - f*v")

build_octree_with(f);
solve(u);

finalize_finch()


export TransfMatrix
export rotatex, rotatey, rotatez, rotate
export translate
export scalex, scaley, scalez, scale

import Base: *, ≈, show

################################################################################

"""
    struct TransfMatrix

A homogeneous transformation matrix. It contains the following fields:
- `m`
- `invm`

The multiplication `*` operation is defined on pairs of `TransMatrix`
objects. You can also use the approximate comparison operator `≈`
(useful for tests).
"""
struct TransfMatrix
    "4×4 matrix representing the transformation"
    m::StaticMatrix{4, 4, Float64}
    "4×4 matrix representing the inverse of the transformation"
    invm::StaticMatrix{4, 4, Float64}

    TransfMatrix() = new(SMatrix{4, 4}(Matrix{Float64}(I, 4, 4)),
                         SMatrix{4, 4}(Matrix{Float64}(I, 4, 4)))
    TransfMatrix(m) = new(m, inv(m))
    TransfMatrix(m, invm) = new(m, invm)
end

Base.:*(tr1::TransfMatrix, tr2::TransfMatrix) = TransfMatrix(tr1.m * tr2.m,
                                                             tr2.invm * tr1.invm)

Base.:≈(tr1::TransfMatrix, tr2::TransfMatrix) = (tr1.m ≈ tr2.m) && (tr1.invm ≈ tr2.invm)

Base.show(io::IO, tr::TransfMatrix) = println("TransfMatrix($(tr.m))")

################################################################################

"""
    identity_transform()

Return a `TransfMatrix` object that represents the identity transformation,
i.e., no transformation at all.
"""
function identity_transform()
    TransfMatrix()
end

################################################################################

"""
    translate(x, y, z)

Return a `TransfMatrix` object that represents a translation operation by
`x`, `y`, and `z` along the three axes.
"""
translate(x, y, z) = TransfMatrix(
    SMatrix{4, 4, Float64}([1 0 0 x;
                            0 1 0 y;
                            0 0 1 z;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([1 0 0 -x;
                            0 1 0 -y;
                            0 0 1 -z;
                            0 0 0 1])
)

"""
    translate(v::Vec3d)

Return a `TransfMatrix` object that represents a translation operator by
vector `v`.
"""
translate(v::Vec3d) = translate(v[1], v[2], v[3])

################################################################################

rotatex(angle) = TransfMatrix(
    SMatrix{4, 4, Float64}([1 0 0 0;
                            0 cos(angle) -sin(angle) 0;
                            0 sin(angle) cos(angle) 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([1 0 0 0;
                            0 cos(-angle) -sin(-angle) 0;
                            0 sin(-angle) cos(-angle) 0;
                            0 0 0 1])
)
rotatey(angle) = TransfMatrix(
    SMatrix{4, 4, Float64}([cos(angle) 0 sin(angle) 0;
                            0 1 0 0;
                            -sin(angle) 0 cos(angle) 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([cos(-angle) 0 sin(-angle) 0;
                            0 1 0 0;
                            -sin(-angle) 0 cos(-angle) 0;
                            0 0 0 1])
)
rotatez(angle) = TransfMatrix(
    SMatrix{4, 4, Float64}([cos(angle) -sin(angle) 0 0;
                            sin(angle) cos(angle) 0 0;
                            0 0 1 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([cos(-angle) -sin(-angle) 0 0;
                            sin(-angle) cos(-angle) 0 0;
                            0 0 1 0;
                            0 0 0 1])
)

"""
    rotatex(angle)
    rotatey(angle)
    rotatez(angle)

Return a `TransfMatrix` representing a rotation around one of the x, y, z axes.
The angle is expressed in radians.
"""
rotatex, rotatey, rotatez

"""
    rotate(angle, axis)

Return a `TransfMatrix` representing a rotation around an arbitrary angle. The
axis `axis` is assumed to have length one, and `angle` must be expressed in
radians.
"""
function rotate(angle, axis)
    cs, sn = cos(angle), sin(angle)
    ux, uy, uz = axis[1], axis[2], axis[3]
    m = SMatrix{4, 4, Float64}(
        cs + ux^2 * (1 - cs),
        uy * ux * (1 - cs) + uz * sn,
        uz * ux * (1 - cs) - uy * sn,
        0,
        ux * uy * (1 - cs) - uz * sn,
        cs + uy^2 * (1 - cs),
        uz * uy * (1 - cs) + ux * sn,
        0,
        ux * uz * (1 - cs) + uy * sn,
        uy * uz * (1 - cs) - ux * sn,
        cs + uz^2 * (1 - cs),
        0,
        0, 0, 0, 1,
    )

    TransfMatrix(m, inv(m))
end

################################################################################

scalex(f) = TransfMatrix(
    SMatrix{4, 4, Float64}([f 0 0 0;
                            0 1 0 0;
                            0 0 1 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([1/f 0 0 0;
                            0 1 0 0;
                            0 0 1 0;
                            0 0 0 1])
)

scaley(f) = TransfMatrix(
    SMatrix{4, 4, Float64}([1 0 0 0;
                            0 f 0 0;
                            0 0 1 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([1 0 0 0;
                            0 1/f 0 0;
                            0 0 1 0;
                            0 0 0 1])
)

scalez(f) = TransfMatrix(
    SMatrix{4, 4, Float64}([1 0 0 0;
                            0 1 0 0;
                            0 0 f 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([1 0 0 0;
                            0 1 0 0;
                            0 0 1/f 0;
                            0 0 0 1])
)

"""
    scalex(f)
    scaley(f)
    scalez(f)

Return a `TransfMatrix` object that represents a scaling operation along
one of the three axes `x`, `y`, or `z`. Use `scale` if you want a generic
scaling operation along the three axes.
"""
scalex, scaley, scalez

"""
    scale(fx, fy, fz)

Return a `TransfMatrix` object that represents a scaling operation along
the three axes `x`, `y`, and `z`. Use `scalex`, `scaley`, and `scalez` if
you are only interested in scaling along *one* of the three axes.
"""
scale(fx, fy, fz) = TransfMatrix(
    SMatrix{4, 4, Float64}([fx 0 0 0;
                            0 fy 0 0;
                            0 0 fz 0;
                            0 0 0 1]),
    SMatrix{4, 4, Float64}([1/fx 0 0 0;
                            0 1/fy 0 0;
                            0 0 1/fz 0;
                            0 0 0 1])
)

################################################################################

*(tr::TransfMatrix, v::Vec3d) = tr.m[1:3, 1:3] * v

function *(tr::TransfMatrix, r::Ray)

    # The origin point should be translated too
    neworig = vec3d(tr.invm[1, 1] * r.origin[1] +
                    tr.invm[1, 2] * r.origin[2] +
                    tr.invm[1, 3] * r.origin[3] +
                    tr.invm[1, 4],
                    tr.invm[2, 1] * r.origin[1] +
                    tr.invm[2, 2] * r.origin[2] +
                    tr.invm[2, 3] * r.origin[3] +
                    tr.invm[2, 4],
                    tr.invm[3, 1] * r.origin[1] +
                    tr.invm[3, 2] * r.origin[2] +
                    tr.invm[3, 3] * r.origin[3] +
                    tr.invm[3, 4])

    # Leave the translation off for the direction
    newdir = vec3d(tr.invm[1, 1] * r.dir[1] +
                    tr.invm[1, 2] * r.dir[2] +
                    tr.invm[1, 3] * r.dir[3],
                    tr.invm[2, 1] * r.dir[1] +
                    tr.invm[2, 2] * r.dir[2] +
                    tr.invm[2, 3] * r.dir[3],
                    tr.invm[3, 1] * r.dir[1] +
                    tr.invm[3, 2] * r.dir[2] +
                    tr.invm[3, 3] * r.dir[3])

    Ray(neworig, newdir)
end


"""
`SmoothedDisplacement <: ImageTransformation`

Description
============

Distorts the given image using a randomly (uniform) generated
`DisplacementField` of the given grid size. This field will be
streched over the given image and converted into a `DisplacementMesh`.
which in turn will morph the original image into a new image using
piecewise affine transformations.

In contrast to `RandomDisplacement`, the resulting vector field is
also smoothed using a gaussian filter with of parameter `sigma`.
This will result in a less chaotic displacement field and be much
more similar to an elastic distortion.

Usage
======

    SmoothedDisplacement(gridwidth, gridheight[, scale=0.2, sigma=2, iterations=1, static_border=true, normalize=true])

Arguments
==========

- **`gridwidth`** : The width of the displacment field. Effectively
specifies the number of vertices along the X dimension used as
landmarks during the deformation process

- **`gridheight`** : The height of the displacment field. Effectively
specifies the number of vertices along the Y dimension used as
landmarks during the deformation process

- **`scale`** : The scaling factor applied to all displacement
vectors in the field. This effectively defines the "strength" of the
deformation. There is no theoretical upper limit to this factor, but
a value somewhere between `0.01` and `1.0` seem to be the most reasonable
choices. Default to `0.2`.

- **`sigma`** : Sigma parameter of the gaussian filter.
This parameter effectively controls the strength of the smoothing.

- **`iterations`** : The number of times the smoothing operation
is applied to the `DisplacementField`. This is especially useful
if `static_border == true` because the border will be reset to
zero after each pass. Thus the displacement is a little less
aggressive towards the borders of the image than it is towards
its center.

- **`static_border`** : If true, the borders of the image will be
preserved. This effectively pins the outermost verticies on their
original position and the operation thus only deforms the inner
content of the image.

- **`normalize`** : If true, the field will be normalized its norm.
This will have the effect that the `scale` factor should be more or
less independent of the grid size.

Methods
========

- **`transform`** : Applies the transformation to the given Image
or set of images

Author(s)
==========

- Christof Stocker (Github: https://github.com/Evizero)

Examples
========

    using Augmentor
    using TestImages

    # load an example image
    img = testimage("lena")

    # Randomly distort the image using a smoothed 10x15 displacement field
    img_new = transform(SmoothedDisplacement(10, 15; scale = .1, sigma = 3), img)

see also
=========

`RandomDisplacement`

`uniform_displacement`, `ImageTransformation`, `transform`
"""
immutable SmoothedDisplacement <: ImageTransformation
    gridwidth::Int
    gridheight::Int
    scale::Float64
    sigma::Float64
    iterations::Int
    static_border::Bool
    normalize::Bool
end

function SmoothedDisplacement(gridwidth::Real, gridheight::Real; scale = .2, sigma = 2, iterations = 1, static_border::Bool = true, normalize::Bool = true)
    (gridwidth > 2 && gridheight > 2) || throw(ArgumentError("gridwidth and gridheight need to be greater than 2"))
    sigma > 0 || throw(ArgumentError("sigma needs to be greater than 0"))
    iterations > 0 || throw(ArgumentError("iterations needs to be greater than 0"))
    SmoothedDisplacement(Int(gridwidth), Int(gridheight), Float64(scale), Float64(sigma), Int(iterations), static_border, normalize)
end

Base.showcompact(io::IO, op::SmoothedDisplacement) = print(io, "Displace randomly with smoothed $(op.gridwidth)x$(op.gridheight) displacement field (σ = $(op.sigma)).")
multiplier(::SmoothedDisplacement) = 1 # not true, basically infinity...

function transform{T<:AbstractImage}(op::SmoothedDisplacement, img::T)
    df = gaussian_displacement(op.gridwidth, op.gridheight, op.scale, op.sigma, op.iterations, op.static_border, op.normalize)
    transform(df, img)::T
end


using Formatting: printfmt
using LinearAlgebra: Diagonal, Hermitian, eigen, norm, tr, diag, dot

using PyCall: pyimport
pyscf = pyimport("pyscf")
mp = pyimport("pyscf.mp")   # Had to import mp alone ??!
cc = pyimport("pyscf.cc") # Had to import mp alone ??!

import PySCF: scf_rhf, mymp2, short_mp2, ccsd

#########################################################################
# Benchmark own HF, MP₂, and CCSD against PySCF
function benchmark_h2o_sto3g()
   mol = pyscf.gto.Mole()
   mol.build(verbose = 0,
             basis = "sto-3g",
             #atom = pyscf_atom_from_xyz("../data/h2.xyz"),
             #atom = pyscf_atom_from_xyz("../data/acetaldehyde.xyz"),
             atom = """8 0  0.    0
                       1 0 -0.757 0.587
                       1 0  0.757 0.587""",
            )

   mf = pyscf.scf.RHF(mol)

   rhf = scf_rhf(mol,verbose=false)
   e_rhf_me, e, C = rhf.Emol, rhf.e, rhf.C
   e_rhf_pyscf    = mf.kernel()

   #mf.analyze(verbose=3,ncol=10, digits=9)

   println( "           E(HF)     ")
   println( "        -------------")
   printfmt(" ME:    {1:13.8f}  \n", e_rhf_me)
   printfmt(" PySCF: {1:13.8f}  \n", e_rhf_pyscf)

   eri   = mol.intor("cint2e_sph", aosym="s8")
   println()
   e_mp2_me1   = @time mymp2(e,C,eri,mol,alg=:ao2mo_noddy) # Naive loop  ao2mo with my MP₂
   e_mp2_me2   = @time mymp2(e,C,eri,mol,alg=:ao2mo_smart) # Smart loop  ao2mo with my MP₂
   e_mp2_me3   = @time short_mp2(e,C,eri,mol)              # @tensor     ao2mo with my MP₂
   e_mp2_pyscf = @time mp.MP2(mf).kernel()[1]              # pure PySCF MP₂ solution
   println()

   println( "           Ecorr(MP₂)        E(MP₂)   ")
   println( "          -----------   -------------")
   printfmt(" ME1:   {1:13.8f}   {2:13.8f}  \n", e_mp2_me1,  e_rhf_me   +e_mp2_me1)
   printfmt(" ME2:   {1:13.8f}   {2:13.8f}  \n", e_mp2_me2,  e_rhf_me   +e_mp2_me1)
   printfmt(" ME3:   {1:13.8f}   {2:13.8f}  \n", e_mp2_me3,  e_rhf_me   +e_mp2_me1)
   printfmt(" PySCF: {1:13.8f}   {2:13.8f}  \n", e_mp2_pyscf,e_rhf_pyscf+e_mp2_pyscf)
   println()

   @time e_ccsd_me, emp2 = ccsd(e,C,eri,mol,verbose=false)
   @time e_ccsd_pyscf    = cc.ccsd.CC(mf).kernel()[1]
   emp2 ≈ e_mp2_me3 || @warn "E_MP2 from CC is not accurate"
   println()

   println("            E(HF)          Ecorr(MP2)           E(MP2)         Ecorr(CCSD)  ")
   println("        --------------   --------------   ---------------   ---------------")
   printfmt(" ME:    {1:13.10f}    {2:13.10f}    {3:13.10f}    {4:13.10f}\n",
                    e_rhf_me,      e_mp2_me3,     e_rhf_me+e_mp2_me3,    e_ccsd_me)
   printfmt(" PySCF: {1:13.10f}    {2:13.10f}    {3:13.10f}    {4:13.10f}\n",
                     e_rhf_pyscf,  e_mp2_pyscf,e_rhf_pyscf+e_mp2_pyscf,e_ccsd_pyscf)
end

benchmark_h2o_sto3g()


############################################################################
# Find the ground state of the Ising Hamitlonian and get the magnetization #
############################################################################
using MatrixProductStates

let
    # Number of Spins
    N = 40
    # Bond Dimension
    D = 20
    # Physical dimension
    d = 2
    # Desired Accuracy
    acc = 1E-10
    # Maximum number of sweeps
    max_sweeps = 10

    # Coupling between nearest neigbours
    J = 1.0
    # Coupling to external field
    lambda = 1.0
    # Generate Ising Hamiltonian
    H = getIsingMPO(N, J, lambda)
    # Get the MPO for the total spin
    Sz = getTotalSpinMPO(N)

    # Run the variational ground state search
    E0, mps, num_of_sweeps = find_groundstate(H, D, d, acc, max_sweeps)

    # Print the results
    println("Groundstate energy:            ", E0)
    println("Groundstate energy density:    ", E0 / N)
    println("Number of sweeps:              ", num_of_sweeps)
    println("Total spin:                    ", expectation_value(mps, Sz))
    println("Entropy along center bond      ", compute_entropy(mps))

    nothing
end


export l₂loss

function l₂loss(𝐲̂, 𝐲)
    feature_dims = 2:(ndims(𝐲)-1)

    loss = sum(.√(sum(abs2, 𝐲̂-𝐲, dims=feature_dims)))
    y_norm = sum(.√(sum(abs2, 𝐲, dims=feature_dims)))

    return loss / y_norm
end


# full kth partial derivative
function ∇ϕ(x, k, A, y)
    res  = 0.0
    for j in 1:size(A)[1]
        res += ∇ϕs(x, k, j, A, y)
    end
    res
end

# full in  place gradient
function ∇ϕ!(G, x, X, y)
    for k in 1:length(x)
        G[k] = ∇ϕ(x, k, X, y)
    end
    G
end

println("Preprocessing data")
ξref = randn(d)
println("Distance before optimization: $(norm(ξref - ξtrue))")
results = optimize(x -> ϕ(x, X, y), (G, x) -> ∇ϕ!(G, x, X, y),  ξref)
println("Optimization success: $(results.ls_success)")
ξref = Optim.minimizer(results)
println("Distance after optimization: $(norm(ξref - ξtrue))")


using Multirate

# Filter tap parameters
Th              = Float64
Tx              = Float64
N𝜙              = 32
resamplerate    = 32/33.5
ƒcutoff         = 0.45
transitionwidth = 0.1
(hLen, β)       = kaiserlength( transitionwidth, samplerate = N𝜙 )
hLen            = ceil( Int, hLen/N𝜙 ) .* N𝜙
h               = Multirate.firdes( hLen, ƒcutoff, DSP.kaiser, samplerate = N𝜙, beta = β ) .* N𝜙
h               = convert( Vector{Th}, h )

# Filter instantiation
polyorder       = 5
x               = rand( Tx, 50  )
arbfilt         = FIRFilter( h, resamplerate, N𝜙 )
farrowfilt      = FIRFilter( h, resamplerate, N𝜙, polyorder )

# Results
yarb            = filt( arbfilt, x)
yfarrow         = filt( farrowfilt, x )
diff_arb_farrow = yarb.-yfarrow
maxError        = maximum( abs, diff_arb_farrow )
minError        = minimum( abs, diff_arb_farrow )

display( [ (1:length(yarb)) yarb yfarrow diff_arb_farrow] )
println()
println( "Max error: $maxError, Min error: $minError" )


using Knet, KnetModules
import KnetModules.forward

type Fire <: KnetModule
    squeeze::Conv
    expand1x1::Conv
    expand3x3::Conv
end

function Fire(inp::Int, sq::Int, e1x1::Int, e3x3::Int)
    squeeze   = Conv(1, 1, inp, sq)
    expand1x1 = Conv(1, 1, sq, e1x1)
    expand3x3 = Conv(3, 3, sq, e3x3; padding=1)
    return Fire(squeeze, expand1x1, expand3x3)
end

function forward(ctx, f::Fire, x)
    o = relu.(@mc f.squeeze(x))
    return ccat(relu.(@mc f.expand1x1(o)),
                relu.(@mc f.expand3x3(o)))
end

# FIXME: too many memory allocations
# FIXME: model is too slow in gpu, possibly because of this
function ccat(a1, a2)
    @assert size(a1, 1,2,4) == size(a2, 1,2,4)
    a1_ = mat(permutedims(a1, (1,2,4,3)))
    a2_ = mat(permutedims(a2, (1,2,4,3)))
    out = hcat(a1_, a2_)
    out = reshape(out, (size(a1, 1,2,4)...,
                        size(out, 2)))
    return permutedims(out, (1,2,4,3))
end

"""
# Usage
    s = SqueezeNet(;num_classes=1000, pdrop=0.5, trained=true)
    @run s(x)
"""
type SqueezeNet <: KnetModule
    features::Sequential
    last_conv::Conv
    pdrop
end

function SqueezeNet(;num_classes=1000, pdrop=0.5, trained=true)
    features = Sequential(
        Conv(3, 3, 3, 64; stride=2),
        ReLU(),
        MaxPool(;window=(3,3), stride=(2,2)),
        Fire(64, 16, 64, 64),
        Fire(128, 16, 64, 64),
        #Fire(128, 32, 128, 128),
        MaxPool(;window=(3,3), stride=(2,2)),
        Fire(128, 32, 128, 128),
        Fire(256, 32, 128, 128),
        MaxPool(;window=(3,3), stride=(2,2)),
        Fire(256, 48, 192, 192),
        Fire(384, 48, 192, 192),
        Fire(384, 64, 256, 256),
        Fire(512, 64, 256, 256)
    )
    last_conv = Conv(1, 1, 512, num_classes;
                     winit=(et, a...)->et(0.01) .* rand(et, a...))
    return SqueezeNet(features, last_conv, pdrop)
end

function forward(ctx, sn::SqueezeNet, x)
    o = @mc sn.features(x)
    o = dropout(o, sn.pdrop)
    o = @mc sn.last_conv(o)
    o = pool(o; window=size(o,1,2), mode=2)
    return mat(o)
end


module IntervalCensored

using Distributions
using Random
using SpecialFunctions
using HypergeometricFunctions
using ForwardDiff
using StaticArrays
using StatsFuns
import LinearAlgebra: dot
import Distributions: ccdf, cdf, logpdf, pdf, quantile,
       mean, rand, params, shape, scale

include("NonParametric.jl")
include("SimTools.jl")
include("calclp.jl")
include("./survdist/survdist.jl")
include("./survdist/GeneralizedGamma.jl")
include("./survdist/LogLogistic.jl")
include("MCEM.jl")

export calclp_ic, calclp_icrt, calclp_dic,
 make_ic, make_icrt, make_dic, make_dicrt,
 SurvIC, SurvICRT, SurvDIC, SurvDICRT,
 eqcdf, cdf, ccdf, pdf, logpdf, mean, shape, scale, params, rand,
 MCEMic, MCEMicrt, MCEMdic,
 GeneralizedGamma, LogLogistic

end


using QuantEcon
using Statistics
using LinearAlgebra

include(string(pwd(),"/src/gensys.jl"))
include(string(pwd(),"/hmc/diffs_estrut_v3.jl"))
include(string(pwd(),"/hmc/diff_kalman.jl"))

function solve_lyapunov_vec(A,Q)
    aux = (I - kron(A,A))
    vec_ans = \(aux,vec(Q))
    res = reshape(vec_ans,size(Q))
    return res
end

function log_like_dsge(par,data;kalman_tol = 1e-10,reduce=true)
    #order to par
    #alfa
    #beta
    #epsilon
    #theta
    #sig
    #sigma: this is the std dev of the innovation
    #phi
    #phi_pi
    #phi_y
    #rho_v

    #data will have t x p dimension: lines are periods p are variables

    alfa = par[1]
    bet = par[2]
    epsilon = par[3]
    theta = par[4]
    sig = par[5]
    #par 6 is coded bellow see line 56
    phi = par[7]
    phi_pi = par[8]
    phi_y = par[9]
    rho_v = par[10]

    THETA = (1-alfa)/(1-alfa+alfa*epsilon)
    lamb = (1-theta)*(1-bet*theta)/theta*THETA
    kappa = lamb*(sig+(phi+alfa)/(1-alfa))

    nobs = size(data,1)
    l = size(data,2)

    GAMMA_0 = [bet  0    0  0;
               1    sig  0  0;
               0    0    0  0;
               0    0    0  1]

    GAMMA_1 = [ 1       -kappa  0   0;
                0        sig    1   0;
               -phi_pi  -phi_y  1  -1;
                0        0      0   rho_v]

    PSI = [0; 0; 0; par[6]]

    PI = [bet  0;
          1    sig;
          0    0;
          0    0]

    p = size(GAMMA_1,1) #number of endogenous vars

    sol = gensys(GAMMA_0,GAMMA_1,PSI,PI; verbose = false)

    if sum(sol.eu) != 2
        return -Inf, repeat([0],length(par))
    end

    #Sig = zeros(p,p)
    #Sig[4,4] = par[6]

    G = zeros(1,p)
    G[1,2] = 1

    A = sol.Theta1
    B = sol.Theta2
    A[abs.(A) .< 10*eps()] .= 0
    B[abs.(B) .< 10*eps()] .= 0
    R = [0] .+ 1e-2
    Q = sol.Theta2*sol.Theta2'
    #Q_eig = eigen(Q)
    #Q = Q_eig.vectors*diagm(repeat([mean(Q_eig.values)],p))*Q_eig.vectors'

    kalman_res = Kalman(A,G,Q,R) #create a Kalman filter instance

    y_mean = mean(data,dims=1)
    y_var = var(data,dims=1)
    y_var = reshape(y_var,size(y_var,2))
    y_var = repeat([1],p)

    x_hat = repeat(y_mean,p) #initial mean of the state
    x_var = I(4) #diagm(repeat(y_var,p))#variance initial of state
    #x_var = x_var*x_var'
    set_state!(kalman_res,x_hat,x_var)

    fit = zeros(nobs,4)

    dA,dB = diff_mod(par)
    #dB = dB[1:4,:]
    dA = dA'*I
    dG = zeros(10,4)
    dQ = (kron(sol.Theta2,I(p)) + kron(I(p),sol.Theta2)*commutation_matrix(4,1))*dB
    dQ = dQ'
    dR = zeros(size(par,1),1)

    llh = zeros(nobs)
    dll = zeros(nobs,size(par,1))

    dx0 = zeros(p)
    S0 = zeros(p,p)

    g_x_hat_l = zeros(size(par,1),p)
    P = kalman_res.G*S0*kalman_res.G' + kalman_res.R
    d_Sl = zeros(10,p^2)#transpose(diff_S0(A,S0,dA,dQ)) #vec, not vech

    for j in 1:nobs
        med = kalman_res.cur_x_hat
        fit[j,:] = med
        varian = kalman_res.cur_sigma
        #println(det(varian))
        eta = data[j,:] - kalman_res.G*med #mean loglike
        P = kalman_res.G*varian*kalman_res.G' + kalman_res.R
        teste_cond = 1/cond(P)

        gain = (A*varian*G')*inv(P)

        d_p = transpose(diff_P(G,varian,dG,d_Sl,dR))
        d_gain = transpose(diff_K(A,P,G,varian,dA,dG,d_p,d_Sl))
        g_eta = transpose(grad_y(med,g_x_hat_l,G,dG))
        g_x_hat = transpose(grad_x_hat(med,dA,A,g_x_hat_l,eta,g_eta, gain,d_gain))

        p_inv = inv(P)

        if teste_cond < kalman_tol
            llh[j] = -500
            dll[j,:] .= 0
        #elseif det(P) <= 0
        #    llh[j] = -500
        else
            llh[j] = -(p*log(2*pi) + logdet(P) .+ eta'*p_inv*eta)/2
            dll[j,:] = -1/2*p_inv*d_p' - (eta'*p_inv)*g_eta' + 1/2*kron(eta'*p_inv,eta'*p_inv)*d_p'
            QuantEcon.update!(kalman_res,data[j,:]) #updating the kalman estimates
        end #end if
        #println(kalman_res.cur_sigma)
        #println(det(varian))
        #println(P)
        g_x_hat_l = g_x_hat
        d_Sl = transpose(diff_S(A,P,gain,varian,dA,d_p,d_gain,d_Sl,dQ))
    end #end for
    llh = llh[10:length(llh)]
    dll = dll[10:length(llh),:]
    if reduce == true
        return sum(llh),sum(dll,dims=1)
    else
        return llh,dll
    end
end


#==============================================================================#
#
#    Declarations of the main `types` used to define Ito diffusions
#    !! IMPORTANT !!
#    to define a diffusion one needs to define a struct and provide a drift
#    and a volatility coefficient functions:
#    `b`: denotes a drift
#    `σ`: denotes a volatility coefficient
#    `a`: (=:σσ') denotes a diffusivity coefficient
#    for linear diffusions
#    `B` and `β` jointly define the drift via `b(t,x):=Bₜx+βₜ`.
#
#==============================================================================#

"""
    DiffusionProcess{T,DP,DW,SS,EI}

Types inheriting from `DiffusionProcess` define Ito diffusions. `T` denotes the
datatype of each coordinate, `DP` the dimension of the stochastic process,
`DW` the dimension of the Wiener process, `SS` lists the state space
restrictions
"""
abstract type DiffusionProcess{T,DP,DW,SS} end

"""
    LinearDiffusion{T,DP,DW,SS} <: DiffusionProcess{T,DP,DW,SS}

Types inheriting from `LinearDiffusion` define a linear Ito-type diffusion, i.e.
solutions to stochastic differential equations of the form:
dXₜ = (BₜXₜ + βₜ)dt + σₜdWₜ, t∈[0,T], X₀=x₀.
"""
abstract type LinearDiffusion{T,DP,DW,SS} <: DiffusionProcess{T,DP,DW,SS} end

"""
    DiffusionDomain

Types inheriting from `DiffusionStateSpace` define the types of restrictions put
on the state-space of the stochastic process.
"""
abstract type DiffusionStateSpace end

const IndexedTime = Tuple{T,S} where {T<:Number,S<:Integer}

#------------------------------------------------------------------------------#
#       Names for functions defining drifts and volatility coefficients
#------------------------------------------------------------------------------#

"""
    b

Compute the drift of a diffusion out-of-place (should use StaticArrays).
"""
function b end

"""
    B

Compute matrix `B` of a linear diffusion out-of-place (should use StaticArrays).
"""
function B end

"""
    β

Compute vector `β` of a linear diffusion out-of-place (should use StaticArrays).
"""
function β end

"""
    σ

Compute the volatility coefficient of a diffusion out-of-place (should use
StaticArrays).
"""
function σ end

"""
    a

Compute the diffusion function σσ' of a diffusion out-of-place (should use
StaticArrays).
"""
function a end

"""
    b!

Compute the drift of a diffusion in-place (uses buffers to store temporary
results).
"""
function b! end

"""
    B!

Compute matrix `B` of a linear diffusion in-place (uses buffers to store
temporary results).
"""
function B! end

"""
    β!

Compute vector `β!` of a linear diffusion in-place (uses buffers to store
temporary results).
"""
function β! end

"""
    σ!

Compute the volatility coefficient of a diffusion in-place (uses buffers to
store temporary results).
"""
function σ! end

"""
    a!

Compute the diffusion function σσ' of a diffusion in-place (uses buffers to
store temporary results).
"""
function a! end

"""
    constdiff

Returns true if the diffusion coefficient does not depend on the state of the
process or time.
"""
function constdiff end


using GeometricIntegrators
using ChargedParticleDynamics.GuidingCenter4d.SymmetricSurface
using GeometricExamples

const Δt     = 5.
const ntime  = 2000
const nx     = 100
const ny     = 100
const nsave  = 4
const nplot  = 4
const run_id = "poincare_2nd_symmetric_vprk_dt5"

include("../guiding_center_4d_settings_vprk.jl")

tableau_list = get_tableau_list_vprk_projection()

pinv = guiding_center_4d_iode_poincare_invariant_2nd(Δt, nx, ny, ntime, nsave)

include("guiding_center_4d_symmetric_poincare_invariant_2nd.jl")


#tracywidom.m
#Algorithm 1.3 of Random Eigenvalues by Alan Edelman

#Theory:      Compute and plot the Tracy-Widom distribution
#See also:    RandomMatrices.TracyWidom(t) for tn=-8

##Parameters
t0=5.         # right endpoint
tn=-8.        # left  endpoint

##Theory: The differential equation solver
#Compute Tracy-Widom distribution
using ODE
include("gradient.jl")
function tracywidom(t0::Float64, tn::Float64)
    function deq(t::Float64, y::Vector{Float64})
        yout = [y[2]; t*y[1]+2y[1]^3; y[4]; y[1]^2]
    end

    y0=[airy(t0); airy(1, t0); 0; airy(t0)^2] # Initial conditions
    (t, y)=ode23(deq, [t0, tn], y0)           # Solve the ODE
    F2=exp(-y[:,3])                           # the distribution
    f2=gradient(F2, t)                        # the density
    return (t, f2)
end
(t, f2) = tracywidom(t0, tn)

## Plot
using Winston
p = FramedPlot()
add(p, Curve(t, f2, "linewidth", 2))
if isinteractive()
    Winston.display(p)
else
    file(p, "tracywidom.png")
end


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
using JuMP

# clnlbeam
# Based on AMPL model
# Copyright (C) 2001 Princeton University
# All Rights Reserved
#
# Permission to use, copy, modify, and distribute this software and
# its documentation for any purpose and without fee is hereby
# granted, provided that the above copyright notice appear in all
# copies and that the copyright notice and this
# permission notice appear in all supporting documentation.

#   Source:
#   H. Maurer and H.D. Mittelman,
#   "The non-linear beam via optimal control with bound state variables",
#   Optimal Control Applications and Methods 12, pp. 19-31, 1991.
let
    N     = 1000
    ni    = N
    h     = 1/ni
    alpha = 350

    m = Model()

    @variable(m, -1 <= t[1:(ni+1)] <= 1)
    @variable(m, -0.05 <= x[1:(ni+1)] <= 0.05)
    @variable(m, u[1:(ni+1)])

    @NLobjective(m, Min, sum{ 0.5*h*(u[i+1]^2 + u[i]^2) + 0.5*alpha*h*(cos(t[i+1]) + cos(t[i])), i = 1:ni})

    # cons1
    for i in 1:ni
        @NLconstraint(m, x[i+1] - x[i] - (0.5h)*(sin(t[i+1])+sin(t[i])) == 0)
    end
    # cons2
    for i in 1:ni
        @constraint(m, t[i+1] - t[i] - (0.5h)*u[i+1] - (0.5h)*u[i] == 0)
    end

    solve(m)

end


# using LinearAlgebra
# include("FieldCalc.jl")
# include("LitzACResist.jl")
# include("StrayCap.jl")
# include("InductanceCalc.jl")

"""
This function makes an elliptic solenoid and calculates the inductance, capacitance, and resistance
    The inputs are:
    N - Number of turns
    r₁ - one radius of the ellipse
    r₂ - Second radius of ellipse
    L - length of solenoid (axially)
    kwargs:
    StrandDiam - Conductor diamter(m)
    InsDiam = insulation diamter (m),
    ϵr_ins - relative permittivity of insulation
    NPts_Coil  - number of points that each wire loop will be discretized into,
"""
function MakeEllipSolenoid(
    N, #Number of turns
    r₁, #Radius 1 of ellipse
    r₂,#Radius 2 of ellipse
    L; #Length of solenoid in meters
    StrandDiam = 0.3e-3,
    InsDiam = 0.9*1e-3,
    ϵr_ins = 4,
    NumStrands = 1,
    BundleDiam = StrandDiam)

    FieldCentered = EvalField_Centered(N,r₁,r₂,L)

    MeanRad = (r₁ + r₂) / 2
    # ZCoord = [(wind_num - 1) / N * L for wind_num = 1:N]
    # RCoords = MeanRad .* ones(length(ZCoord))
    # Coords = hcat(RCoords[:], ZCoord[:])
    # L_Coil = CalcInductance(Coords)
    L_Coil = SimpleInduct(N,MeanRad,L)
    C_Coil = CL_DAE_Knight(
        L, #Coil length (axial)
        InsDiam, #Coil diamter
        N, #N turns
        1, #external permittivity
        ϵr_ins) #Internal permittivity
    SRF = 1 / (2 * π * sqrt(L_Coil * C_Coil))

    CoilProps, Rac, SkinDepth = LitzWire_ACRes(
        StrandDiam,
        N,
        L,
        MeanRad,
        NumStrands,
        BundleDiam;
        f = 75e3:1000:125e3,
        PlotOn=false)

    return L_Coil, C_Coil, SRF, FieldCentered,Rac
end

function LfromΦMat(ΦMat)

    N = size(ΦMat)[1]
    LMat = zeros(N,N)
    for i in 1:N
        for j in 1:N
            MagΦᵢ = sqrt(sum(ΦMat[i,i,:].^2))
            LMat[i,j] = sum(ΦMat[i,j,:].*ΦMat[i,i,:]) / MagΦᵢ
        end
    end

    return LMat
end

"""
This function makes an elliptic solenoid and calculates the inductance.
    The function calculates the flux through the center of each loop's cross-section
        Lₛ = Φ/I (self inductance = flux per amp)
        Lₘ₂₁ = (Φ₂₁ ⋅ Φ₂₂/|Φ₂₂|) /I
            -mutual inductance from the 1ˢᵗ wire on 2ⁿᵈ is flux from 1 in the 2ⁿᵈ's cross section
            that is aligned with the flux from the second on itself. All per ampere
        Lₘ₂₁ = Lₘ₁₂
        Total L = ∑ᵢ∑ⱼ(Lᵢⱼ)
    The inputs are:
    N - Number of turns
    r₁ - one radius of the ellipse
    r₂ - Second radius of ellipse
    L - length of solenoid (axially)
    kwargs:
    MeasLayers - Number of concentric layers to calc field
                    The field is calculated over concentric ellipses (linearly spaced)
                    This is the number of concentric test point layers
    PtsPerLayer -  Number of test points per layer
    NPts_Coil  - number of points that each wire loop will be discretized into,
"""
function EvalInduct_Biot(
    N, #Number of turns
    r₁, #Radius 1 of ellipse
    r₂,#Radius 2 of ellipse
    L; #Length of solenoid in meters
    MeasLayers = 55,#Number of concentric layers to calc field
    PtsPerLayer =  6,
    NPts_Coil = 500
    )

    A_All = zeros(MeasLayers*PtsPerLayer,3,N)
    Weights = zeros(MeasLayers*PtsPerLayer)
    NPtsPerSlice = MeasLayers*PtsPerLayer
    for wind_num in 1:N
        A_All[:,1,wind_num],A_All[:,2,wind_num],A_All[:,3,wind_num]  =
         MakeEllipTestPoints(r₁,r₂;Center = [0,0,(wind_num - 1) / N * L],NPts=PtsPerLayer,Layers = MeasLayers)

    end
    SliceArea = π*r₁*r₂
    Weights = sum(A_All[:,1:2,1].^2,dims=2) #A matrix that weights each value based on radius
    Weights = Weights / sum(Weights)*NPtsPerSlice  #normalizing

    ΦAll = zeros(N,N,3)

    for Wᵢ in 1:N
        for Aᵢ in 1:N
            DeconstructA =
                        [([A_All[i,1,Aᵢ],A_All[i,2,Aᵢ],A_All[i,3,Aᵢ]]) for i in 1:length(Weights)]
            Field3DSlice = [
                BiotSav(
                    MakeEllip(
                                r₁, #Radius 1 of ellipse
                                r₂;#Radius 2 of ellipse
                                Center = [0, 0, (Wᵢ - 1) / N * L], #The solenoid's axis is in Z
                                NPts = NPts_Coil, #How discretized the windings are
                                ),
                        TestLoc;
                        MinThreshold = 0.001,
                        )
                        for TestLoc in DeconstructA
                        ]

         Bweighted = vcat( [(Field3DSlice[i].*Weights[i]) for i in 1:length(Weights)]...)
          # Bweighted = vcat( [(Field3DSlice[i]) for i in 1:length(Weights)]...)

         ΦAll[Wᵢ,Aᵢ,:] =sum(Bweighted,dims=1).*SliceArea ./NPtsPerSlice
        # ΦAll[Wᵢ,Aᵢ,:] =sum(Bweighted,dims=1) ./NPtsPerSlice
        end
        println("$Wᵢ of $N windings calculated")
    end

    LMat = LfromΦMat(ΦAll)
    LTotal = sum(LMat)*1e6
end

"""
This function makes an elliptic solenoid and calculates the field at the center of the solenoid
        Lₘ₂₁ = Lₘ₁₂
        Total L = ∑ᵢ∑ⱼ(Lᵢⱼ)
    The inputs are:
    N - Number of turns
    r₁ - one radius of the ellipse
    r₂ - Second radius of ellipse
    L - length of solenoid (axially)
    kwargs:

    NPts_Coil  - number of points that each wire loop will be discretized into,
"""
function EvalField_Centered(
    N, #Number of turns
    r₁, #Radius 1 of ellipse
    r₂,#Radius 2 of ellipse
    L; #Length of solenoid in meters
    NPts_Coil=100
    )

            FieldCentered = [
                BiotSav(
                    MakeEllip(
                                r₁, #Radius 1 of ellipse
                                r₂;#Radius 2 of ellipse
                                Center = [0, 0, (Wᵢ - 1) / N * L], #The solenoid's axis is in Z
                                NPts = NPts_Coil, #How discretized the windings are
                                ),
                        [0, 0, L/2];
                        MinThreshold = 0.001,
                        )[3]
                        for Wᵢ in 1:N
                        ]

    return sum(FieldCentered)
end


using CGP
using Logging

inputs = Dict("cart_pole"=>4, "mountain_car"=>2, "pendulum"=>3)
outputs = Dict("cart_pole"=>2, "mountain_car"=>1, "pendulum"=>1)

function cart_pole(c::Chromosome)
    gravity = 9.81; m_cart = 1.0; m_pole = 0.1; tau = 0.02; force_mag = 10.0
    theta_threshold = 12 * 2 * pi / 360; x_threshold = 2.4; pole_length = 0.5
    m_total = m_cart + m_pole
    limits = [x_threshold, force_mag, theta_threshold, 2*pi]
    reward = 0
    state = 0.1*rand(4)-0.05
    x, x_dot, theta, theta_dot = state

    while ((-x_threshold <= x <= x_threshold) &&
           (-theta_threshold <= theta <= theta_threshold) &&
           reward < 100)
        # get action
        inps = [-1.0; max.(min.(state ./ limits, 1.0), -1.0)[:]]
        action = indmax(process(c, inps))
        force = force_mag * (2 * (action - 1) - 1)
        reward += 1

        # update
        temp = (force + m_pole * pole_length * theta_dot * theta_dot * sin(theta))/m_total
        theta_acc = ((gravity * sin(theta) - cos(theta) * temp) /
                     (pole_length * (4.0/3 - m_pole * cos(theta) * cos(theta) / m_total)))
        x_acc = temp - m_pole * pole_length * theta_acc * cos(theta) / m_total
        x += tau * x_dot
        x_dot += tau * x_acc
        theta += tau * theta_dot
        theta_dot += tau * theta_acc
        state = [x, x_dot, theta, theta_dot]
        # println(reward, " ", action, " ", state, " ", inps)
    end
    reward / 100
end

function mountain_car(c::Chromosome)
    min_x = -1.2; max_x = 0.6; max_v = 0.07
    goal_position = 0.45; power = 0.0015
    steps = 0
    x = -0.6 + 0.2 * rand(); v = 0.0
    best_x = x
    while ((x <= goal_position) && (steps < 100))
        inps = [0.0, 0.0, 2*((x-min_x)/(max_x-min_x))-1.0, v / max_v, 0.0]
        force = process(c, inps)[1]
        v = min(max(v + force * power - 0.0025 * cos(3 * x), -max_v), max_v)
        x = min(max(x + v, min_x), max_x)
        if (x == min_x) && (v < 0)
            v = 0
        end
        steps += 1
        best_x = max(best_x, x)
        # println(steps, " ", inps, " ", force, " ", x, " ", v, " ")
    end
    min(max(best_x / goal_position, 0.0), 1.0)
end

function pendulum(c::Chromosome)
    max_speed = 8; max_torque = 2.0; dt = 0.5;
    g = 10.; m = 1.; l = 1.

    th = 2*pi*rand() - pi
    thdot = 2*rand() - 1.0

    cost = 0.0; steps = 0;
    while steps < 100
        inps = [1.0, cos(th), sin(th), thdot / max_speed, 0.0]
        u = max_torque * process(c, inps)[2]

        angle = (((th+pi) % (2*pi)) - pi)
        cost -= angle^2 + (.1*thdot)^2 + 0.001*(u^2)

        thdot = min(max(thdot + (-3*g/(2*l) * sin(th +pi) + 3 ./(m*l^2)*u) * dt,
                        -max_speed), max_speed)
        th = th + thdot * dt
        steps += 1
        # println(steps, " ", inps, " ", u, " ", cost)
    end
    max(min((cost + 4000) / 3850, 1.0), 0.0)
    # min((-100-cost)/-100, 1.0)
end

function combined(c::Chromosome)
    score = cart_pole(c); CGP.reset!(c)
    score += mountain_car(c); CGP.reset!(c)
    score += pendulum(c)
    score
end

function repeat(c::Chromosome, func::Function)
    score = 0
    for i = 1:5
        score += func(c)
    end
    score/5
end

seed = 0
log = "log"
fitness = "pendulum"
if length(ARGS) > 0; seed = parse(Int64, ARGS[1]); end
if length(ARGS) > 1; log = ARGS[2]; end
if length(ARGS) > 2; fitness = ARGS[3]; end

# CGP.Config.init("cfg/base.yaml")
# CGP.Config.init("cfg/classic.yaml")
CGP.Config.init("cfg/test.yaml")

Logging.configure(filename=log, level=INFO)
nin = inputs[fitness]; nout = outputs[fitness]
fit = x->repeat(x, eval(parse(fitness)))

include("param_sweep.jl")
param_sweep()


# See https://en.wikipedia.org/wiki/Complex_number

expressions = [
    :(sqrt(-1+0im))
    :(exp(π * im) + 1)
    :(sin(im))
    :(acos(2+0im))
    :(asin(2+0im))
    :(atan(1+2im))
    :(log(-1+0im))
    :(log(-im))
]

for expr in expressions
    println("[] :julia-answer \"\"\"", escape_string(string(expr)), " = ", eval(expr), "\"\"\".")
end


# Implementation of Franz and Scholkopf
# A unifying view of Wiener and Volterra theory and polynomial
# kernel regression

abstract type VolterraKernel end

export AdaptiveKernel,
InhomogeneousPolyKernel,
GaussianKernel,
VovkKernel

immutable AdaptiveKernel <: VolterraKernel
    a::Vector{Float64}
end

(k::AdaptiveKernel)(h) = sum(k.a.^2.*[h^n for n in 0:length(k.a)-1])
(k::AdaptiveKernel)(x1,x2) = k(dot(x1,x2))
weights(k::AdaptiveKernel,n) = k.a[n+1]
weights(k::AdaptiveKernel) = k.a

immutable InhomogeneousPolyKernel <: VolterraKernel
    p
end

(k::InhomogeneousPolyKernel)(h) = (1+h)^k.p
(k::InhomogeneousPolyKernel)(x1,x2) = k(dot(x1,x2))
weights(k::InhomogeneousPolyKernel,n) = binomial(k.p,n)
weights(k::InhomogeneousPolyKernel) = [weights(k,n) for n in 0:p]

immutable GaussianKernel <: VolterraKernel
end

(k::GaussianKernel)(h) = exp(h)
(k::GaussianKernel)(x1,x2) = k(dot(x1,x2))
weights(k::GaussianKernel,n) = 1/factorial(n)

immutable VovkKernel <: VolterraKernel
    α
end

(k::VovkKernel)(h) = (1-h)^(-k.α)
(k::VovkKernel)(x1,x2) = k(dot(x1,x2))
weights(k::VovkKernel,n) = binomial(-k.α,n)*(-1)^n

## Now we need to do kernel regression

"""
Form the kernel matrix out of a data matrix.
"""
kernel_matrix(k::VolterraKernel,X1,X2) = map(k,X1'X2)
kernel_matrix(k::VolterraKernel,X) = kernel_matrix(k,X,X)

type VolterraModel <: DischargeModel
    M::Int
    k::VolterraKernel
    β::Vector{Float64}
    H::Matrix{Float64}
    Q::Vector{Float64}
end

VolterraModel(M,N,k=1) = VolterraModel(M,k,zeros(N),zeros(M,N),zeros(N))

"""
Fitting a Volterra Series Model simply requires inversion of the kernel matrix.
"""
function StatsBase.fit(M::VolterraModel,H::Matrix{Float64},Q::Vector{Float64})
    VolterraModel(M.M,M.k,kernel_matrix(M.k,H)\Q,H,Q)
end

function StatsBase.fit!(M::VolterraModel,H::Matrix{Float64},Q::Vector{Float64})
    M.β = kernel_matrix(M.k,H)\Q
    M.H = H
    M.Q = Q
    M
end

StatsBase.coef(M::VolterraModel) = M.β
StatsBase.nobs(M::VolterraModel) = length(M.Q)
StatsBase.model_response(M::VolterraModel) = M.Q
function StatsBase.predict(M::VolterraModel,H)
    vec(M.β'kernel_matrix(M.k,M.H,H))
end
StatsBase.predict(M::VolterraModel) = predict(M,M.H)
StatsBase.residuals(M::VolterraModel) = M.Q-predict(M)
StatsBase.residuals(M::VolterraModel,H,Q) = Q-predict(M,H)

type RegularizedVolterraModel<: DischargeModel
    M::Int
    k::VolterraKernel
    λ::Float64
    R::AbstractMatrix
    β::Vector{Float64}
    H::Matrix{Float64}
    Q::Vector{Float64}
end

RegularizedVolterraModel(M,N,k=1,λ=0.0) = RegularizedVolterraModel(M,InhomogeneousPolyKernel(k),λ,eye(N-M+1,N-M+1),zeros(N),zeros(M,N),zeros(N))

function StatsBase.fit(M::RegularizedVolterraModel,H::Matrix{Float64},Q::Vector{Float64})
    RegularizedVolterraModel(M.M,
                             M.k,
                             M.λ,
                             M.R,
                             (kernel_matrix(M.k,H)+M.λ*M.R)\Q,
                             H,
                             Q)
end

function StatsBase.fit!(M::RegularizedVolterraModel,H::Matrix{Float64},Q::Vector{Float64})
    M.β = (kernel_matrix(M.k,H)+M.λ*M.R)\Q
    M.H = H
    M.Q = Q
    M
end

StatsBase.coef(M::RegularizedVolterraModel) = M.β
StatsBase.nobs(M::RegularizedVolterraModel) = length(M.Q)
StatsBase.model_response(M::RegularizedVolterraModel) = M.Q
function StatsBase.predict(M::RegularizedVolterraModel,H)
    vec(M.β'kernel_matrix(M.k,M.H,H))
end
StatsBase.predict(M::RegularizedVolterraModel) = predict(M,M.H)
StatsBase.residuals(M::RegularizedVolterraModel) = M.Q-predict(M)
StatsBase.residuals(M::RegularizedVolterraModel,H,Q) = Q-predict(M,H)

"""
To get the Volterra operator, we have to calculate a design matrix.

This is Φn = [ϕn(x_1),ϕn(x_2),...]

where ϕn(x) collects all monomials of order n of the entries of x
"""
function ϕ(x,n)
    out = ones(ntuple(y->length(x),n))
    for I in CartesianRange(size(out))
        for i in 1:length(I)
            out[I]*=x[I[i]]
        end
    end
    vec(out)
end

function designmatrix(H,n)
    a,b = size(H)
    Φ = zeros(a^n,b)
    for i in 1:b
        Φ[:,i] = ϕ(H[:,i],n)
    end
    Φ
end

"""
Derive the Volterra operator of order n from a model
"""
operator(M::VolterraModel,n) = weights(M.k,n).*designmatrix(M.H,n)*M.β
operator(M::RegularizedVolterraModel,n) = weights(M.k,n).*designmatrix(M.H,n)*M.β


"""
    covdiff(xm, xd; include_diagonal=false)

Error of the covariance matrix.
"""
function covdiff(xm::AbstractArray, xd::AbstractArray; include_diagonal::Bool=false)
    @assert ndims(xm) == ndims(xd)
    @assert size(xm)[1:(end - 1)] == size(xd)[1:(end - 1)]
    Cm = cov(reshape(xm, :, size(xm)[end]); dims=2)
    Cd = cov(reshape(xd, :, size(xd)[end]); dims=2)
    if include_diagonal
        return norm(Cm - Cd) / length(Cm)
    else
        return norm(Cm - Diagonal(diag(Cm)) - Cd + Diagonal(diag(Cd))) / (length(Cm) - size(Cm, 1))
    end
end

"""
    adversarial_accuracy(xm, xd)

Adversarial accuracy indicator.
As defined in https://arxiv.org/abs/2105.13889.
"""
function adversarial_accuracy(xm::AbstractArray, xd::AbstractArray)

end


module q1

# function for calculating temperature and pressure
function tandp(ke, pe, vir, N)
    # Take the time average of the virial, ⟨∑∑r⃗∂V/∂r⟩
    virialavg = means(vir)

    # Find the temperatures
    T = 2ke/3N

    # Set Boltsmann constant
    kB = 1.3806e-23

    kBTs = kB*T

    p = 1 - virialavg./(6N*kBTs)

    return T, p
end


function initial_state_szegedy(sqrtstochastic::SparseMatrixCSC{<:Number})
   vec(sqrtstochastic)/sqrt(size(sqrtstochastic, 1))
end

function initial_state(qws::QWSearch{<:AbstractSzegedy})
   initial_state_szegedy(qws.model.sqrtstochastic)
end

function evolve(qwd::QWDynamics{<:AbstractSzegedy},
                state::AbstractVector{T}) where T<:Number
   _evolve_szegedy((qwd.parameters[:operators])..., state)
end

function _evolve_szegedy(operator1::AbstractMatrix{T},
                         operator2::AbstractMatrix{T},
                         state::AbstractVector{T}) where T<:Number
   operator2*(operator1*state)
end

function measure(::QWDynamics{<:AbstractSzegedy},
                 state::AbstractVector{<:Number})
   dim = floor(Int, sqrt(length(state)))
   vec(mapslices(sum, reshape((abs.(state)).^2, (dim, dim)), dims=[1]))
end

function measure(::QWDynamics{<:AbstractSzegedy},
                 state::AbstractVector{<:Number},
                 vertices::Vector{Int})
   dim = floor(Int, sqrt(length(state)))
   vec(mapslices(sum, abs.(reshape(state, (dim, dim))[:,vertices]).^2, dims=[1]))
end