# # NICE network
# For the definition of this network and concepts of normalizing flow,
# please refer this nice blog: https://lilianweng.github.io/lil-log/2018/10/13/flow-based-deep-generative-models.html,
# and the pytorch notebook: https://github.com/GiggleLiu/marburg/blob/master/notebooks/nice.ipynb

using NiLang, NiLang.AD
using LinearAlgebra
using DelimitedFiles
using Plots

# `include` the optimizer, you can find it under the `Adam.jl` file in the `examples/` folder.
include(NiLang.project_relative_path("examples", "Adam.jl"))

# ## Model definition
# First, define the single layer transformation and its behavior under `GVar` - the gradient wrapper.
struct NiceLayer{T}
    W1::Matrix{T}
    b1::Vector{T}
    W2::Matrix{T}
    b2::Vector{T}
    y1::Vector{T}
    y1a::Vector{T}
end

"""Apply a single NICE transformation."""
@i function nice_layer!(x::AbstractVector{T}, layer::NiceLayer{T},
                y!::AbstractVector{T}) where T
    @routine @invcheckoff begin
        i_affine!(layer.y1, layer.W1, layer.b1, x)
        @inbounds for i=1:length(layer.y1)
            if (layer.y1[i] > 0, ~)
                layer.y1a[i] += layer.y1[i]
            end
        end
    end
    i_affine!(y!, layer.W2, layer.b2, layer.y1a)
    ~@routine
    ## clean up accumulated rounding error, since this memory is reused.
    @safe layer.y1 .= zero(T)
end

# Here, in each layer, we use the information in `x` to update `y!`.
# During computing, we use the `y1` and `y1a` fields of the network as ancilla space,
# both of them can be uncomputed at the end of the function.
# However, we need to erase small numbers to make sure the rounding error does not accumulate.

# A nice network always transforms inputs reversibly.
# We update one half of `x!` a time, so that input and output memory space do not clash.
const NiceNetwork{T} = Vector{NiceLayer{T}}

"""Apply a the whole NICE network."""
@i function nice_network!(x!::AbstractVector{T}, network::NiceNetwork{T}) where T
    @invcheckoff for i=1:length(network)
        np ← length(x!)
        if (i%2==0, ~)
            @inbounds nice_layer!(x! |> subarray(np÷2+1:np), network[i], x! |> subarray(1:np÷2))
        else
            @inbounds nice_layer!(x! |> subarray(1:np÷2), network[i], x! |> subarray(np÷2+1:np))
        end
        np → length(x!)
    end
end

function random_nice_network(nparams::Int, nhidden::Int, nlayer::Int; scale=0.1)
    random_nice_network(Float64, nparams, nhidden, nlayer; scale=scale)
end

function random_nice_network(::Type{T}, nparams::Int, nhidden::Int, nlayer::Int; scale=0.1) where T
    nin = nparams÷2
    scale = T(scale)
    y1 = zeros(T, nhidden)
    NiceLayer{T}[NiceLayer(randn(T, nhidden, nin)*scale, randn(T, nhidden)*scale,
            randn(T, nin, nhidden)*scale, randn(T, nin)*scale, y1, zero(y1)) for _ = 1:nlayer]
end

# ## Parameter management

nparameters(n::NiceLayer) = length(n.W1) + length(n.b1) + length(n.W2) + length(n.b2)
nparameters(n::NiceNetwork) = sum(nparameters, n)

"""collect parameters in the `layer` into a vector `out`."""
function collect_params!(out, layer::NiceLayer)
    a, b, c, d = length(layer.W1), length(layer.b1), length(layer.W2), length(layer.b2)
    out[1:a] .= vec(layer.W1)
    out[a+1:a+b] .= layer.b1
    out[a+b+1:a+b+c] .= vec(layer.W2)
    out[a+b+c+1:end] .= layer.b2
    return out
end

"""dispatch vectorized parameters `out` into the `layer`."""
function dispatch_params!(layer::NiceLayer, out)
    a, b, c, d = length(layer.W1), length(layer.b1), length(layer.W2), length(layer.b2)
    vec(layer.W1) .= out[1:a]
    layer.b1 .= out[a+1:a+b]
    vec(layer.W2) .= out[a+b+1:a+b+c]
    layer.b2 .= out[a+b+c+1:end]
    return layer
end

function collect_params(n::NiceNetwork{T}) where T
    out = zeros(T, nparameters(n))
    k = 0
    for layer in n
        np = nparameters(layer)
        collect_params!(view(out, k+1:k+np), layer)
        k += np
    end
    return out
end

function dispatch_params!(network::NiceNetwork, out)
    k = 0
    for layer in network
        np = nparameters(layer)
        dispatch_params!(layer, view(out, k+1:k+np))
        k += np
    end
    return network
end

# ## Loss function

# To obtain the log-probability of a data.

@i function logp!(out!::T, x!::AbstractVector{T}, network::NiceNetwork{T}) where T
    (~nice_network!)(x!, network)
    @invcheckoff for i = 1:length(x!)
        @routine begin
            xsq ← zero(T)
            @inbounds xsq += x![i]^2
        end
        out! -= 0.5 * xsq
        ~@routine
    end
end

# The negative-log-likelihood loss function

@i function nice_nll!(out!::T, cum!::T, xs!::Matrix{T}, network::NiceNetwork{T}) where T
    @invcheckoff for i=1:size(xs!, 2)
        @inbounds logp!(cum!, xs! |> subarray(:,i), network)
    end
    out! -= cum!/(@const size(xs!, 2))
end

# ## Training

function train(x_data, model; num_epochs = 800)
    num_vars = size(x_data, 1)
    params = collect_params(model)
    optimizer = Adam(; lr=0.01)
    for epoch = 1:num_epochs
        loss, a, b, c = nice_nll!(0.0, 0.0, copy(x_data), model)
        if epoch % 50 == 1
            println("epoch = $epoch, loss = $loss")
            display(showmodel(x_data, model))
        end
        _, _, _, gmodel = (~nice_nll!)(GVar(loss, 1.0), GVar(a), GVar(b), GVar(c))
        g = grad.(collect_params(gmodel))
        update!(params, grad.(collect_params(gmodel)), optimizer)
        dispatch_params!(model, params)
    end
    return model
end

function showmodel(x_data, model; nsamples=2000)
    scatter(x_data[1,1:nsamples], x_data[2,1:nsamples]; xlims=(-5,5), ylims=(-5,5))
    zs = randn(2, nsamples)
    for i=1:nsamples
        nice_network!(view(zs, :, i), model)
    end
    scatter!(zs[1,:], zs[2,:])
end

# you can find the training data in `examples/` folder
x_data = Matrix(readdlm(NiLang.project_relative_path("examples", "train.dat"))')

import Random; Random.seed!(22)
model = random_nice_network(Float64, size(x_data, 1), 10, 4; scale=0.1)

# Before training, the distribution looks like
# ![before](../../asset/nice_before.png)
model = train(x_data, model; num_epochs=800)

# After training, the distribution looks like
# ![before](../../asset/nice_after.png)


immutable Outliers
    nsamples    :: Int64
    low_severe  :: Int64
    low_mild    :: Int64
    high_mild   :: Int64
    high_severe :: Int64

    function Outliers(nsamples::Int64,
		      low_severe::Int64, low_mild::Int64,
		      high_mild::Int64, high_severe::Int64)
        new(nsamples, low_severe, low_mild,
	    high_mild, high_severe)
    end

    function Outliers(n::Int64)
        @assert n > 0
    end
end

abstract OutlierEffect
immutable Unaffected <: OutlierEffect end
immutable Slight <: OutlierEffect end
immutable Moderate <: OutlierEffect end
immutable Severe <: OutlierEffect end

function outlier_effect(var_outlier::Float64)
    if     var_outlier < 0.01 return Unaffected()
    elseif var_outlier < 0.10 return Slight()
    elseif var_outlier < 0.50 return Moderate()
    else return Severe()
    end
end

immutable OutlierVariance
    effect :: OutlierEffect
    frac   :: Float64
end

function outlier_variance(mean_est::Float64, std_est::Float64, niter::Integer)
    mean_g  = 0.5 * (mean_est / niter)
    std_g = min(mean_g / 4.0, std_est / sqrt(niter))

    function var_out(x)
        ((niter - x) / niter) * std_est^2 - (niter - x) * std_g^2
    end

    function cmax(x)
        k   = (mean_est / niter) - x
        d   = k^2
        ad  = niter * d
        k0  = -niter * ad
        k1  = std_est^2 - (niter * std_g^2) + ad
        det = k1^2 - (4 * std_g^2 * k0)
        int(floor(-2 * k0 / (k1 * sqrt(det))))
    end
    var_out_min = min(var_out(1), min(cmax(0), cmax(mean_g))) / std_est^2
    OutlierVariance(outlier_effect(var_out_min), var_out_min)
end

immutable SampleAnalysis
    mean::Estimate
    std::Estimate
    outlier_variance::OutlierVariance
end

function classify_outliers{T<:Real}(samples::Vector{T})

    q1, q2, q3 = quantile(samples, [0.25, 0.5,  0.75])
    iqr = q3 - q1
    loS = q1 - (iqr * 3.0)
    loM = q1 - (iqr * 1.5)
    hiM = q3 + (iqr * 1.5)
    hiS = q3 + (iqr * 3.0)

    nsamples = length(samples)
    low_severe = 0
    low_mild = 0
    high_mild = 0
    high_severe = 0

    for s in samples
        if s <= loS && s < hiM
            low_severe += 1
        elseif s > loS && s <= loM
            low_mild += 1
        elseif s >= hiM && s < hiS
            high_mild += 1
        elseif s >= hiS && s > loM
            high_severe += 1
	end
    end
    return Outliers(nsamples,
		    low_severe,
		    low_mild,
		    high_mild,
		    high_severe)
end

function count_outliers(o::Outliers)
    return o.low_severe + o.low_mild +
	   o.high_mild + o.high_severe
end

function count_samples(o::Outliers)
    return o.nsamples
end

function note_outliers(o::Outliers)
    frac  = (n) -> 100.0 * (n / o.nsamples)
    check = (k,t,d) -> let fk = frac(k)
			if fk > t
                           @printf(" %d (%.1f%%) %s\n", k, fk, d)
                        end
                       end
    noutliers = count_outliers(o)
    if noutliers > 0
        @printf("found %d outliers among %d samples (%.1f%%)\n",
		noutliers, o.nsamples, frac(noutliers))
        check(o.low_severe, 0, "low severe")
        check(o.low_mild, 1, "low mild")
        check(o.high_mild, 1, "high mild")
        check(o.high_severe, 0, "high severe")
    end
end

function analyze_mean(sample::Vector{Float64}, niter::Integer)
    u = mean(sample)
    @printf("mean is %s (%d iterations)\n", time_str(u), niter)
    note_outliers(classify_outliers(sample))
    return u
end

function scale(sa::SampleAnalysis, factor::FloatingPoint)
    SampleAnalysis(sa.mean * factor,
                   sa.std  * factor,
                   sa.outlier_variance)
end

function analyze_sample(samples::Vector{Float64},
			tailq::Float64,
                        nresamples::Integer)

   est_mean = bootstrap_bca(samples, mean, tailq, nresamples)
   est_std  = bootstrap_bca(samples, std, tailq, nresamples)
   out_var  = outlier_variance(est_mean.point, est_std.point, length(samples))
   SampleAnalysis(est_mean, est_std, out_var)
end


@testset "1631.path-with-minimum-effort.jl" begin
    @test minimum_effort_path([[1, 2, 2], [3, 8, 2], [5, 3, 5]]) == 2
    @test minimum_effort_path([[1, 2, 3], [3, 8, 4], [5, 3, 5]]) == 1
    @test minimum_effort_path([
        [1, 2, 1, 1, 1], [1, 2, 1, 2, 1], [1, 2, 1, 2, 1], [1, 2, 1, 2, 1], [1, 1, 1, 2, 1]
    ]) == 0
    @test minimum_effort_path([[1, 10, 6, 7, 9, 10, 4, 9]]) == 9
end


 #map(abspath,readdir()) # get the file location and file names
 # create log file
using DelimitedFiles
writedlm("log_output_sigma1.txt",  " ")
writedlm("log_output_sigma10.txt", " ")
include("GB6.jl")
import GLM: fit, predict, LinearModel
using Distributions,SpecialFunctions
using LinearAlgebra, MultivariateStats
using StatsBase, Random, Distributed
using DelimitedFiles, DataStructures
using CSV
using DataFrames

Random.seed!(2019)
rng = 2019

# 1) Simulate n (= 500 )× p ( = 1~1000) training data and test data, where σ² = 1 and 10 separately
# 2) Plot the broken line chart

gbdt = GB6.GBDT(;
  sampling_rate = 0.75,                # Sampling rate
  learning_rate = 0.1,                # Learning rate
  num_iterations = 100,               # NuGBer of iterations
  tree_options = Dict(
      :nsubfeatures => 0,
      :min_samples_leaf => 1,
      :max_depth => 5))
gbl = GB6.GBLearner(gbdt, :regression)

######## σ² = 1 #########
temp_storedim = size(collect(10:1000),1)
fitval = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
rssfitted = []
predicted = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
rmsefitted = []
Ytestval = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
μ_testval = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
rmsefitted2 = []
s
for pval in collect(10:1000)
    # reproducibility
    Random.seed!(2019)
    rng = 2019

    ### simulate the training and testing data
    X, Y = GB6.test_data(500; p = pval, sigma = 1.0)
    X_test, Y_test = GB6.test_data(500; p = pval)
    μ = GB6.fried(X)
    μ_test = GB6.fried(X_test)

    println("The last simulated data:
        \nX:$(X[500, pval]) \nY:$(Y[500])
        \nX_test: $(X_test[500, pval]) \nY_test: $(Y_test[500])
        \nμ: $(μ[500]) \nμ_test: $(μ_test[500])\n\n")

    μ_testval[:, pval - 9] = μ_test
    Ytestval[:, pval - 9 ] = Y_test

    # train the model
    mymodel=GB6.fit!(gbl, X, Y)
    # get the fitted value
    myfit  = GB6.predict!(gbl, X)
    # store the fitted value
    fitval[:, pval - 9] = myfit

    # get the residual sum of squares
    rssfit = sqrt(mean(Y .- myfit) .^2)
    push!(rssfitted, rssfit)
    # get prediction
    mypred = GB6.predict!(gbl, X_test)

    println("The last pred is $(mypred[500])\n\n ")

    predicted[:, pval - 9] = mypred

    # get mse and root of mean squared error
    rmse = sqrt(mean((Y_test .- mypred ) .^2))

    println("The rmse b.w. Y_test and pred for p = $(pval) is $(rmse) ")

    push!(rmsefitted, rmse)

    # get the rmse(μ_test, mypred)
    rmse2 = sqrt(mean((μ_test .- mypred).^2))

    println("The rmse b.w. μ_test and pred for p = $(pval)  is $(rmse2).")

    push!(rmsefitted2, rmse2)
end

rssfitted
rmsefitted
rmsefitted2
Ytestval
μ_testval

fitval = DataFrame(fitval)
CSV.write("fitval_sigma1.csv", fitval)
rssfitted = DataFrame(rssfitted = rssfitted)
CSV.write("fss_sigma1.csv", rssfitted)
rmsefitted = DataFrame(rmsefitted = rmsefitted)
CSV.write("rmse_sigma1.csv", rmsefitted)
predicted = DataFrame(predicted)
CSV.write("pred_sigma1.csv", predicted)
rmsefitted2 = DataFrame(rmsefitted2 = rmsefitted2)
CSV.write("rmse2_sigma1.csv", rmsefitted2)
Ytestval = DataFrame(Ytestval)
CSV.write("Ytest_sigma1.csv", Ytestval)
μ_testval = DataFrame(μ_testval)
CSV.write("mu_test_sigma1.csv", μ_testval)
##### not runned yet

######## σ² = 10 #########

fitval_sigma10 = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
rssfitted_sigma10 = []
predicted_sigma10 = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
rmsefitted_sigma10 = []
Ytestval_sigma10 = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
μ_testval_sigma10 = Array{Union{Missing, Nothing, Int, Float64}}(nothing, 500, temp_storedim)
rmsefitted2_sigma10 = []
s
for pval in collect(10:1000)
    # reproducibility
    Random.seed!(2019)
    rng = 2019

    ### simulate the training and testing data
    X, Y = GB6.test_data(500; p = pval, sigma = sqrt(10.0))
    X_test, Y_test = GB6.test_data(500; p = pval, sigma = sqrt(10.0))
    μ = GB6.fried(X)
    μ_test = GB6.fried(X_test)

    println("Simulated data:
    \nX: $(X[500, pval])
    \nY: $(Y[500])
    \nX_test: $(X_test[500, pval])
    \nY_test: $(Y_test[500])
    \nμ: $(μ[500])
    \nμ_test: $(μ_test[500])")

    μ_testval_sigma10[:, pval - 9] = μ_test
    Ytestval_sigma10[:, pval - 9] = Y_test

    # train the model
    mymodel=GB6.fit!(gbl, X, Y)
    # get the fitted value
    myfit  = GB6.predict!(gbl, X)
    # store the fitted value
    fitval_sigma10[:, pval - 9] = myfit

    # get the residual sum of squares
    rssfit = sqrt(mean((Y .- myfit) .^2))
    push!(rssfitted_sigma10, rssfit)
    # get prediction
    mypred = GB6.predict!(gbl, X_test)

    println("The last pred is $(mypred[500])")

    predicted_sigma10[:, pval - 9] = mypred

    # get mse and root of mean squared error
    rmse = sqrt(mean((Y_test .- mypred ) .^2))
    println("The p = $(pval), and the rmse is $rmse.")
    push!(rmsefitted_sigma10, rmse)

    # get the rmse(Y_test, μ_test)
    rmse2 = sqrt(mean((mypred .- μ_test) .^2))
    println("The p = $(pval) dim, rmse2 is $rmse2.")
    push!(rmsefitted2_sigma10, rmse2)
end

fitval_sigma10
rssfitted_sigma10
predicted_sigma10
rmsefitted_sigma10
Ytestval_sigma10
μ_testval_sigma10
rmsefitted2_sigma10

#predicted_sigma10 = DataFrame(predicted_sigma10)
rmsefitted_sigma10 = DataFrame(rmsefitted_sigma10=rmsefitted_sigma10)

#CSV.write("fitval_sigma10.csv", fitval_sigma10)
#CSV.write("rssfitted_sigma10.csv", rssfitted_sigma10)
#CSV.write("predicted_sigma10.csv", predicted_sigma10)
#CSV.write("rmsefitted_sigma10.csv", rmsefitted_sigma10)

fitval_sigma10 = DataFrame(fitval_sigma10)
CSV.write("fitval_sigma10.csv", fitval)
rssfitted_sigma10 = DataFrame(rssfitted_sigma10 = rssfitted_sigma10)
CSV.write("fss_sigma10.csv", rssfitted_sigma10)

CSV.write("rmse_sigma10.csv", rmsefitted_sigma10)
predicted_sigma10 = DataFrame(predicted_sigma10)
CSV.write("pred_sigma10.csv", predicted_sigma10)
rmsefitted2_sigma10 = DataFrame(rmsefitted2_sigma10 = rmsefitted2_sigma10)
CSV.write("rmse2_sigma10.csv", rmsefitted2_sigma10)
Ytestval_sigma10 = DataFrame(Ytestval_sigma10)
CSV.write("Ytest_sigma10.csv", Ytestval_sigma10)
μ_testval_sigma10 = DataFrame(μ_testval_sigma10)
CSV.write("mu_test_sigma10.csv", μ_testval_sigma10)

# plot the rmsefitted for sigma = 1
using Gadfly

rmse_sigma1 = CSV.read("rmse_sigma1.csv")
rmse_sigma10 = CSV.read("rmse_sigma10.csv")

rmse2_sigma1 = CSV.read("rmse2_sigma1.csv")
rmse2_sigma10 = CSV.read("rmse2_sigma10.csv")

 # 991
i = collect(1:1:10)
ii = collect(20:10:100)
iii = collect(200:100:900)
ind = vcat(i, ii, iii)
rmse1_forplot = rmse_sigma1[ind,1]
rmse10_forplot = rmse_sigma10[ind,1]

p_1 = layer(x = ind, y = rmse1_forplot,
            Geom.point, Geom.line,
            Theme(default_color="red"))
p_2 = layer(x = ind, y = rmse10_forplot,
            Geom.point, Geom.line, Theme(default_color="blue"))
myplot1 = plot(p_1, p_2, Guide.xlabel("P"),
    Guide.ylabel("rmse"), Guide.title("Root Mean Squared Error"),
    Guide.manual_color_key("", ["σ² = 1", "σ² = 10"], ["red", "blue"]))

using Cairo, Fontconfig
draw(PNG("myplot1.png", 3inch, 3inch), myplot1)

rmse2_1_forplot = rmse2_sigma1[ind,1]
rmse2_10_forplot = rmse2_sigma10[ind,1]

pp_1 = layer(x = ind, y = rmse2_1_forplot,
            Geom.point, Geom.line,
            Theme(default_color="red"))
pp_2 = layer(x = ind, y = rmse2_10_forplot,
            Geom.point, Geom.line, Theme(default_color="blue"))
myplot2 = plot(pp_1, pp_2, Guide.xlabel("P"),
    Guide.ylabel("rmse"), Guide.title("Root Mean Squared Error"),
    Guide.manual_color_key("", ["σ² = 1", "σ² = 10"], ["red", "blue"]));

Gadfly.push_theme(style(line_width=1mm))

ppp_2 = layer(x = ind, y = rmse2_10_forplot,
            style(line_width=10cm),
            Geom.point, Geom.line, Theme(default_color="blue"))
pppp = plot(ppp_2, Guide.xlabel("P"),
    Guide.ylabel("rmse"), Guide.title("Root Mean Squared Error"),
    Guide.manual_color_key("", ["σ² = 1", "σ² = 10"], ["red", "blue"]));
draw(PNG("test.png", 15inch, 15inch), pppp)

Gadfly.pop_theme()

draw(PNG("myplot2.png", 15inch, 15inch), myplot2)

p = 10
x, y = GB6.test_data(500; p = 10, sigma = sqrt(10.0))
x_test, y_test = GB6.test_data(500; p = 10, sigma = sqrt(10.0))
mu = GB6.fried(x)
mu_test = GB6.fried(x_test)
x=DataFrame(x)
y=DataFrame(y=y)
x_test=DataFrame(x_test)
y_test=DataFrame(y_test=y_test)
mu = DataFrame(mu = mu)
mu_test = DataFrame(mu_test = mu_test)

## for comparing the data with bencnmark method, use CSV

CSV.write("_X.csv", x)
CSV.write("_Y.csv", y)
CSV.write("_X_test.csv", x_test)
CSV.write("_Y_test.csv", y_test)
CSV.write("_mu.csv", mu)
CSV.write("_mu_test.csv", mu_test)

for j in 1 : size(x, 2)
    x[:, j] = collect(skipmissing(x[:,j]))
end
x = convert(Matrix, x)
y = collect(skipmissing(y[:,1]))

for j in 1 : size(x_test, 2)
    x_test[:, j] = collect(skipmissing(x_test[:,j]))
end
x_test = convert(Matrix, x_test)
y_test = collect(skipmissing(y_test[:,1]))

mu_test = collect(skipmissing(mu_test[:,1]))

gbdt = GB6.GBDT(;
  sampling_rate = 0.75,                # Sampling rate
  learning_rate = 0.1,                # Learning rate
  num_iterations = 1,               # NuGBer of iterations
  tree_options = Dict(
      :nsubfeatures => 0,
      :min_samples_leaf => 1,
      :max_depth => 5))
gbl = GB6.GBLearner(gbdt, :regression)

hypers = GB6.Hypers()
hypers.σ = sqrt(1)
hypers.τ = 3 / 10 * var(y) / gbdt.num_iterations
mymodel=GB6.fit!(gbl, x, y, hypers)
# get the fitted value
myfit  = GB6.predict!(gbl, x)

# get the residual sum of squares
sqrt(mean(y .- myfit) .^2) # 0.0103

# get prediction
mypred = GB6.predict!(gbl, x_test)

# get mse and root of mean squared error
sqrt(mean((y_test .- mypred ) .^2)) #10.2

sqrt(mean((mypred .- mu_test ) .^2))

#### compare with R
# σ² = 10, σ = √10
xx = CSV.read("xx.csv")[:, 2:end]
yy = CSV.read("yy.csv")[:, 2]
xx_test = CSV.read("xx_test.csv")[:, 2:end]
yy_test = CSV.read("yy_test.csv")[:, 2]
μμ_test = CSV.read("mumu_test.csv")[:, 2]

for j in 1 : size(xx, 2)
    xx[:, j] = collect(skipmissing(xx[:,j]))
end
xx = convert(Matrix, xx)
yy = collect(skipmissing(yy[:,1]))

for j in 1 : size(xx_test, 2)
    xx_test[:, j] = collect(skipmissing(xx_test[:,j]))
end
xx_test = convert(Matrix, xx_test)
yy_test = collect(skipmissing(yy_test[:,1]))

μμ_test = collect(skipmissing(μμ_test[:,1]))

# set the same combination of parameters, though I did not use cross validation here
gbdt_new = GB6.GBDT(;
  sampling_rate = 0.8,                # Sampling rate
  learning_rate = 0.1,                # Learning rate
  num_iterations = 100,               # NuGBer of iterations
  tree_options = Dict(
      :nsubfeatures => 0,
      :min_samples_leaf => 5,
      :max_depth => 6))
gbl_new = GB6.GBLearner(gbdt_new, :regression)

mymodel2=GB6.fit!(gbl_new, xx, yy)
# get the fitted value
myfit2 = GB6.predict!(gbl_new, xx)

# get the residual sum of squares
sqrt(mean(yy .- myfit2) .^2) # 0.02446; 0.0045(σ = √10)

# get prediction
mypred2 = GB6.predict!(gbl_new, xx_test)

# get mse and root of mean squared error
mean((yy_test .- mypred2 ) .^2) # ; 14.1986(σ = √10)
sqrt(mean((yy_test .- mypred2 ) .^2)) #12.02; 3.7678(σ = √10)

sqrt(mean((mypred2 .- μμ_test ) .^2)) # 6.6 ; 1.93(σ = 1.9)


abstract type AbstractMonomialBasis{MT<:MP.AbstractMonomial, MV<:AbstractVector{MT}} <: AbstractPolynomialBasis end

Base.length(basis::AbstractMonomialBasis) = length(basis.monomials)
Base.copy(basis::AbstractMonomialBasis) = typeof(basis)(copy(basis.monomials))

MP.nvariables(basis::AbstractMonomialBasis) = MP.nvariables(basis.monomials)
MP.variables(basis::AbstractMonomialBasis) = MP.variables(basis.monomials)
MP.monomialtype(::Type{<:AbstractMonomialBasis{MT}}) where MT = MT

empty_basis(MB::Type{<:AbstractMonomialBasis{MT}}) where {MT} = MB(MP.emptymonovec(MT))
function maxdegree_basis(B::Type{<:AbstractMonomialBasis}, variables, maxdegree::Int)
    return B(MP.monomials(variables, 0:maxdegree))
end
function basis_covering_monomials(B::Type{<:AbstractMonomialBasis}, monos::AbstractVector{<:AbstractMonomial})
    return B(monos)
end

# The `i`th index of output is the index of occurence of `x[i]` in `y`,
# or `0` if it does not occur.
function multi_findsorted(x, y)
    I = zeros(Int, length(x))
    j = 1
    for i in eachindex(x)
        while j ≤ length(y) && x[i] < y[j]
            j += 1
        end
        if j ≤ length(y) && x[i] == y[j]
            I[i] = j
        end
    end
    return I
end

function merge_bases(basis1::MB, basis2::MB) where MB<:AbstractMonomialBasis
    monos = MP.mergemonovec([basis1.monomials, basis2.monomials])
    I1 = multi_findsorted(monos, basis1.monomials)
    I2 = multi_findsorted(monos, basis2.monomials)
    return MB(monos), I1, I2
end

"""
    struct MonomialBasis{MT<:MP.AbstractMonomial, MV<:AbstractVector{MT}} <: AbstractPolynomialBasis
        monomials::MV
    end

Monomial basis with the monomials of the vector `monomials`.
For instance, `MonomialBasis([1, x, y, x^2, x*y, y^2])` is the monomial basis
for the subspace of quadratic polynomials in the variables `x`, `y`.
"""
struct MonomialBasis{MT<:MP.AbstractMonomial, MV<:AbstractVector{MT}} <: AbstractMonomialBasis{MT, MV}
    monomials::MV
end
MonomialBasis(monomials::AbstractVector) = MonomialBasis(MP.monovec(monomials))

MP.polynomialtype(::Union{MonomialBasis{MT}, Type{<:MonomialBasis{MT}}}, T::Type) where MT = MP.polynomialtype(MT, T)
MP.polynomial(f::Function, mb::MonomialBasis) = MP.polynomial(f, mb.monomials)

MP.polynomial(Q::AbstractMatrix, mb::MonomialBasis, T::Type) = MP.polynomial(Q, mb.monomials, T)

function MP.coefficients(p, ::Type{<:MonomialBasis})
    return MP.coefficients(p)
end


### extirpolation.jl
#
# Copyright (C) 2016 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: periodogram, lomb scargle, extirpolation
#
# This file is a part of LombScargle.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# Utility functions used for implementing the fast algorithm proposed here:
# * Press, W. H., Rybicki, G. B. 1989, ApJ, 338, 277 (URL:
#   http://dx.doi.org/10.1086/167197,
#   Bibcode: http://adsabs.harvard.edu/abs/1989ApJ...338..277P)
# The actual implementation is adapted from Astropy project.  See
# stats/lombscargle/implementations/utils.py in Astropy source code.
#
### Code:

function add_at!{N1<:Number,N2<:Number,N3<:Number}(arr::AbstractVector{N1},
                                                   ind::AbstractVector{N2},
                                                   vals::AbstractVector{N3})
    for i in eachindex(ind)
        arr[mod(ind[i], length(arr)) + 1] += vals[i]
    end
end

function extirpolate{RE<:Real,NU<:Number}(X::AbstractVector{RE},
                                          Y::AbstractVector{NU},
                                          N::Integer=0, M::Integer=4)
    @assert length(X) == length(Y)
    x, y = collect(X), collect(Y)
    if N <= 0
        # Get the maximum of "X", `maximum' has a faster method for ranges.
        N = trunc(Int, maximum(X) + 0.5*M + 1)
    end
    # Now use legendre polynomial weights to populate the results array; This is
    # an efficient recursive implementation (See Press et al. 1989)
    result = zeros(NU, N)
    # first take care of the easy cases where x is an integer
    integers = find(isinteger, x)
    add_at!(result, trunc(Int, x[integers]), y[integers])
    deleteat!(x, integers)
    deleteat!(y, integers)
    # For each remaining x, find the index describing the extirpolation range.
    # i.e. ilo[i] < x[i] < ilo[i] + M with x[i] in the center, adjusted so that
    # the limits are within the range 0...N
    ilo = clamp(trunc(Int, x - div(M, 2)), 0, N - M)
    numerator = y .* [prod(x[j] - ilo[j] - (0:(M - 1))) for j in eachindex(x)]
    denominator = factorial(M - 1)
    @inbounds for j in 0:(M - 1)
        if j > 0
            denominator *= j/(j - M)
        end
        ind = ilo + (M - 1 - j)
        add_at!(result, ind, numerator ./ (denominator * (x .- ind)))
    end
    return result
end

function trig_sum{R1<:Real,R2<:Real}(t::AbstractVector{R1},
                                     h::AbstractVector{R2}, df::Real,
                                     N::Integer, f0::Real=0.0,
                                     freq_factor::Integer=1,
                                     oversampling::Integer=5, Mfft::Integer=4)
    @assert Mfft > 0
    df *= freq_factor
    f0 *= freq_factor
    @assert df > 0
    # `ifft' can take arrays of any length in input, but it's faster when the
    # length is exactly a power of 2.
    Nfft = nextpow2(N * oversampling)
    t0 = minimum(t)
    if f0 > 0
        h = h .* exp(2im * pi * f0 * (t - t0))
    end
    tnorm = mod(((t - t0) * Nfft * df), Nfft)
    grid = extirpolate(tnorm, h, Nfft, Mfft)
    fftgrid = Nfft * ifft(grid)[1:N]
    if t0 != 0
        f = f0 + df * (0:N - 1)
        fftgrid .*= exp(2im * pi * t0 * f)
    end
    C = real(fftgrid)
    S = imag(fftgrid)
    return C, S
end


using Test
using JuMP
using MultivariateMoments
using PolyJuMP
using DynamicPolynomials

function _zero_polynomial_test(
    optimizer::MOI.AbstractOptimizer,
    config::MOI.Test.Config,
    complex::Bool,
)
    atol = config.atol
    rtol = config.rtol

    model = _model(optimizer)

    @variable(model, α)
    @variable(model, β ≤ 1)
    @variable(model, γ)

    @polyvar x y
    z = complex ? 0 + 0im : 0
    cref = @constraint(model, (α - β) * x * y == z)
    p = γ * x * y
    cγ = @constraint(model, p in PolyJuMP.ZeroPoly())

    @objective(model, Max, α + γ)
    optimize!(model)

    @test termination_status(model) == MOI.OPTIMAL
    @test objective_value(model) ≈ 1.0 atol=atol rtol=rtol

    @test primal_status(model) == MOI.FEASIBLE_POINT
    @test value(α) ≈ 1.0 atol=atol rtol=rtol
    @test value(β) ≈ 1.0 atol=atol rtol=rtol
    @test value(γ) ≈ 0.0 atol=atol rtol=rtol
    @test value(UpperBoundRef(β)) ≈ 1.0 atol=atol rtol=rtol
    @test value(cref) isa MultivariatePolynomials.AbstractPolynomial{Float64}
    @test value(cref) ≈ 0.0 * x * y atol=atol rtol=rtol
    @test value(cγ) ≈ 0.0 * x * y atol=atol rtol=rtol

    @test dual_status(model) == MOI.FEASIBLE_POINT
    @test dual(UpperBoundRef(β)) ≈ -1.0 atol=atol rtol=rtol

    for μ in [dual(cref), moments(cref), dual(cγ), moments(cγ)]
        @test μ isa AbstractMeasure{Float64}
        @test length(moments(μ)) == 1
        @test moment_value(moments(μ)[1]) ≈ -1.0 atol=atol rtol=rtol
        @test monomial(moments(μ)[1]) == x*y
    end

    F = MOI.VectorAffineFunction{Float64}
    S = PolyJuMP.ZeroPolynomialSet{FullSpace,MB.MonomialBasis,Monomial{true},
                                   MonomialVector{true}}
    @test Set(MOI.get(model, MOI.ListOfConstraintTypesPresent())) == Set([
        (MOI.VariableIndex, MOI.LessThan{Float64}), (F, S), (MOI.VectorOfVariables, S)])
    @testset "Delete" begin
        test_delete_bridge(model, cref, 3, ((F, MOI.Zeros, 0),))
    end
end

function real_zero_polynomial_test(
    optimizer::MOI.AbstractOptimizer,
    config::MOI.Test.Config,
)
    return _zero_polynomial_test(optimizer, config, false)
end
function complex_zero_polynomial_test(
    optimizer::MOI.AbstractOptimizer,
    config::MOI.Test.Config,
)
    return _zero_polynomial_test(optimizer, config, false)
end

linear_tests["zero_polynomial"] = real_zero_polynomial_test
linear_tests["zero_polynomial"] = complex_zero_polynomial_test


module DimensionalPlotRecipes

using RecipesBase
using LinearAlgebra

# Splits a complex matrix to its real and complex parts
# Reals defaults solid, imaginary defaults dashed
# Label defaults are changed to match the real-imaginary reference / indexing
@recipe function f(x::AbstractArray{T},y::Array{Complex{T2}};
                   transformation=:split3D) where {T<:Real,T2<:Number}
  if transformation==:split2D
    A = real.(y)
    B = imag.(y)
    _y = [!iseven(i) ? A[:,i÷2+1] : B[:,i÷2] for i in 1:2size(A,2)]
    linestyle --> reshape([!iseven(i) ? :solid : :dash for i in 1:2size(A,2)],1,2size(A,2))
    label --> reshape([!iseven(i) ? "Re(y$(i÷2+1))" : "Im(y$(i÷2))" for i in 1:2size(A,2)],1,2size(A,2))
    seriescolor --> reshape([!iseven(i) ? i÷2+1 : i÷2 for i in 1:2size(A,2)],1,2size(A,2))
    retval = (x,_y)
  elseif transformation==:split3D
    A = real.(y)
    B = imag.(y)
    ylabel --> "Re(y)"
    zlabel --> "Im(y)"
    retval = (x,A,B)
  elseif transformation==:modulus
    _y = norm.(y)
    label --> reshape(["||y$(i)||" for i in 1:size(_y,2)],1,size(_y,2))
    retval = (x,_y)
  elseif transformation==:modulus2
    _y = abs2.(y)
    label --> reshape(["|y$(i)|^2" for i in 1:size(_y,2)],1,size(_y,2))
    retval = (x,_y)
  else
    error("Transformation unknown. Please use :split2D, :split3D, :modulus, or :modulus2")
  end
  (retval...,)
end

end # module


using Test

struct MyNumber
	x::Float64
end

for op = (:sin, :cos, :tan, :log, :csc, :sinpi, :cbrt, :mod2pi)
    eval(quote
        Base.$op(a::MyNumber) = MyNumber($op(a.x))
    end)
end

	#test using x = pi
	x = MyNumber(pi)

	@test cos(x) == MyNumber(-1.0)
	@time cos(x)
	# 0.025440 seconds (122.20 k allocations: 6.844 MiB, 11.69% gc time)
	#sin(x) should be 1.2246467991473532e-16
	@test sin(x) == MyNumber(1.2246467991473532e-16)
	@test tan(x) == MyNumber(-1.2246467991473532e-16)
	@test log(x) == MyNumber(1.1447298858494002)
	@time log(x)
	@test exp(x) == MyNumber(23.140692632779267)
	@test csc(x) == MyNumber(8.165619676597685e15)
	@time csc(x)
	@test sinpi(x) == MyNumber(-0.43030121700009194)
	@test cbrt(x) == MyNumber(1.4645918875615231)

#test using an integer = 1
struct MyNumber
	y::Float64
end

	@test cos(y) == MyNumber2(0.5403023058681398)
	@test sin(y) == MyNumber2(0.8414709848078965)
	@test tan(y) == MyNumber2(1.5574077246549023)
	@test log(y) == MyNumber2(0.0)
	@time log(y)
	@test exp(y) == MyNumber2(2.718281828459045)

#test


module MatrixImpute

using Distributed: pmap, @distributed, @everywhere
using LinearAlgebra: pinv, svd, norm, I, dot, diagm
using StatsBase: sample
using Arpack: svds
using DataFrames: allowmissing
using Statistics: mean

function Impute(A,k;method=:fastImpute,γ = 10e4, λ=0.1, B=nothing, lr= 1/64)

    dims = length(size(A))

    if dims == 2
        if method == :fastImpute
            Aopt = fastImpute(A,k,γ,lr,B, approx = true)
            return Aopt
        else
            error("Method not yet implemented!")
        end
    elseif dims ==3
        if method == :fastImputeT && B == nothing
            Aopt = fastImputeT(A,k,γ,lr)
            return Aopt
        elseif method == :fastImputeT
            error("Tensor with side information is not supported!")
        else
            error("Method not yet implemented!")
        end
    else
        error("Number of Dimensions Incorrect")
    end
end

function fastImpute(A,k,γ,lr,B; approx = true)
# Unknown values should be in missing for A
# A is the matrix that needs to be imputed, size n x m
# B is the feature selector, size m x p
# k is the maximum rank of the resultant matrix
# γ is the regularization parameter
# lr is the learning rate of the gradient descent function
    n = size(A)[1]
    m = size(A)[2]
    chunk = 1000
    if B == nothing
        p = m
        startchunk = collect(1:chunk:m)
        nchunks = length(startchunk)
        Uopt = zeros(Float64,n,k)
        Sopt = zeros(Float64,k,p)
        Aopt = zeros(Float64,n,p)
        if nchunks>1
            for sc in startchunk
                UoptT, SoptT = fastImputeInner(A[:,sc:min(sc+chunk-1,p)],k,γ,lr,B)
                Aopt[:,sc:min(sc+chunk-1,p)] = UoptT * SoptT
            end
            if approx
                return Aopt
            else
                output = svds(Aopt, nsv = k)
                Aopt = output[1].U*diagm(output[1].S)*output[1].Vt
                return Aopt
            end
        else
            Uopt, Sopt = fastImputeInner(A,k,γ,lr,B)
            return Uopt * Sopt
        end
    else
        if m != size(B)[1]
            error("Sizes of A and B must match")
        end
        Uopt, Sopt = fastImputeInner(A,k,γ,lr,B)
        return Uopt * Sopt
    end
end

function fastImputeInner(A,k,γ,lr,B)
    n = size(A)[1]
    m = size(A)[2]
    if B == nothing
        p = m
    elseif B != nothing
        p = size(B)[2]
    end
    M = sum(ismissing.(A))/(n*m)
    Uopt = zeros(n,k)
    W = zeros(n,m)
    for i = 1:n
        W[i,:] = ones(m) - convert(Vector{Float64},ismissing.(A[i,:]))
    end
    println("Preprocessing Complete")
    A[ismissing.(A)] .= 0
    if M<0.5 && B == nothing
        output = svds(A,nsv = k)
        S = diagm(output[1].S)*output[1].Vt
    else
        S = rand(k,p)
        S = S ./ norm(S)
    end
    objbest = Inf
    ∇obj = zeros(k,p)
    i = 0
    j = 0
    if B != nothing
        obj0, ∇obj0 = MatrixDeriv2(A,B,W,S,k,γ,M,j)
    else
        obj0, ∇obj0 = MatrixDeriv(A,W,S,k,γ,M,j)
    end
    noimprov = 0
    raised = false
    while i<50
        if (noimprov>=5) && (!raised)
            j = j + 1
            raised = true
            objbest = obj0
            noimprov = 0
        end
        ∇obj = ∇obj0 + i / (i+3) * ∇obj
        ∇tan = -∇obj + dot(∇obj,S) * S
        updategrad = ∇tan / norm(∇tan)
        S = S * cos(pi * lr) + updategrad*sin(pi * lr)
        S = S / norm(S)
        if B != nothing
            obj0, ∇obj0 = MatrixDeriv2(A,B,W,S,k,γ,M,j)
        else
            obj0, ∇obj0 = MatrixDeriv(A,W,S,k,γ,M,j)
        end
        if obj0 > objbest
            noimprov = noimprov + 1
            if (noimprov % 5 == 0) && (raised)
                raised = false
            end
        else
            noimprov = 0
            objbest = obj0
            raised = false
        end
        i = i + 1
    end
    Sopt = S
    if B == nothing
        X = Sopt'
    else
        X = B * Sopt'
    end
    println("Model Solved")
    function PopulateA(Xtemp,Atemp)
        return pinv(Xtemp'*Xtemp,1e-7)*(Xtemp'*Atemp)
    end
    Xtemp=Array{Float64}[X[W[i,:].==1,:] for i=1:n]
    Atemp=Array{Float64}[A[i,W[i,:].==1] for i=1:n]
    result=pmap(PopulateA,Xtemp,Atemp)
    for i=1:n
        Uopt[i,:]=result[i]
    end
    return Uopt, Sopt
end

function MatrixDeriv(A,W,S,k,γ,M,j)
    # Derivative Calculation for Matrix Completion without Side Information
    n = size(A)[1]
    m = convert(Int,size(A)[2])
    p = m
    nsquare = sqrt(n*m)
    mnew = m
    nnew = min(max(100,Int(round(nsquare * log(nsquare) * k * 2^j/(8 * mnew * (1-M))))),n)
    normfactor = 2 * mnew * nnew
    ∇obj = zeros(k,p)
    samplen=sample(1:n,nnew,replace=false)
    samplem=sample(1:m,mnew,replace=false)
    obj=0
    SmallInv = function (Xrow, Wrow, Arow)
        return Arow - Xrow * (inv(I / γ + Xrow' * Xrow) * (Xrow' * Arow))
    end
    Wpar = Array{Float64}[W[samplen[i],samplem] for i=1:nnew]
    Xpar = Array{Float64}[S[:,samplem[Wpar[i].==1]]' for i=1:nnew]
    Apar = Array{Float64}[A[samplen[i],samplem[Wpar[i].==1]] for i=1:nnew]
    objpar = map(SmallInv,Xpar,Wpar,Apar)
    obj = @distributed (+) for i = 1:nnew
        dot(Apar[i],objpar[i]) / normfactor
    end
    Btemp2 = zeros(k,nnew)
    for i = 1:nnew
        Btemp2[:,i] = Xpar[i]' * objpar[i]
        ∇obj[:,samplem[Wpar[i].==1]] += -2 * γ * Btemp2[:,i] * objpar[i]' / normfactor
    end
    return obj, ∇obj
end

function MatrixDeriv2(A,B,W,S,k,γ,M,j)
    # Derivative Calculation for Matrix Completion with Side Information
    p = size(B)[2]
    n = size(A)[1]
    m = convert(Int,size(A)[2])
    ∇obj = zeros(k,p)
    nsquare = sqrt(n * m)
    mnew = min(4*p,m)
    nnew = min(max(100,Int(round(nsquare * log(nsquare) * k * 2^j / (8 * mnew * (1-M))))),n)
    samplen = sample(1:n,nnew,replace=false)
    samplem = sample(1:m,mnew,replace=false)
    X = B[samplem,:]*S'
    obj = 0
    SmallInv = function (Xrow, Wrow, Arow)
        return Arow - Xrow * (inv(I / γ + Xrow' * Xrow) * (Xrow' * Arow))
    end
    Wpar = Array{Float64}[W[samplen[i],samplem] for i=1:nnew]
    Xpar = Array{Float64}[X[Wpar[i].==1,:] for i=1:nnew]
    Apar = Array{Float64}[A[samplen[i],samplem[Wpar[i].==1]] for i=1:nnew]
    objpar = map(SmallInv,Xpar,Wpar,Apar)
    obj = @distributed (+) for i = 1:nnew
        dot(Apar[i],objpar[i]) / (2*mnew*nnew)
    end
    Btemp1 = zeros(p,nnew)
    Btemp2 = zeros(k,nnew)
    for i = 1:nnew
        Btemp1[:,i] = B[samplem[Wpar[i].==1],:]' * objpar[i]
        Btemp2[:,i] = X[Wpar[i].==1,:]' * objpar[i]
    end
    ∇obj = @distributed (+) for i = 1:nnew
        -2 * γ * Btemp2[:,i] * Btemp1[:,i]' / (2*mnew*nnew)
    end
    return obj, ∇obj
end

function fastImputeT(A,k,γ,lr)
# Unknown values should be in missing for A
# A is the tensor that needs to be imputed, size n x m x p
# k is the maximum CP rank of the resultant matrix
# γ is the regularization parameter
# lr is the learning rate of the gradient descent function
    n,m,p = size(A)
    Amat = reshape(A,(n*m,p))
    W = ones(n,m,p)-ismissing.(A)
    Wmat = ones(n*m,p)-ismissing.(Amat)
    x = rand(n, k)
    y = rand(m, k)
    M = sum(ismissing.(A))/(n*m*p)

    z = [x; y]
    z = z ./ norm(z)
    x = z[1:n,:]
    y = z[n+1:end,:]
    i = 0
    j = 0
    println("Preprocessing Complete")
    obj0, ∇obj0x, ∇obj0y = TensorDeriv(Amat,Wmat,x,y,k,γ,M,j,size(A))

    ∇obj0 = [∇obj0x; ∇obj0y]
    ∇obj = zeros(m+n,k)

    objbest = Inf
    noimprov = 0
    raised = false
    while i<50
        if (noimprov>=5) && (!raised)
            j = j + 1
            raised = true
            objbest = obj0
            noimprov = 0
        end
        ∇obj = ∇obj0 + i / (i+3) * ∇obj
        ∇tan = - ∇obj + dot(∇obj,z) * z
        updategrad = ∇tan / norm(∇tan)
        z = z .* cos(pi * lr) + updategrad * sin(pi * lr)
        z = z ./ norm(z)
        x = z[1:n,:]
        y = z[n+1:end,:]
        obj0, ∇obj0x, ∇obj0y = TensorDeriv(Amat,Wmat,x,y,k,γ,M,j,size(A))
        ∇obj0 = [∇obj0x; ∇obj0y]
        if obj0 > objbest
            noimprov = noimprov + 1
            if (noimprov % 5 == 0) && (raised)
                raised = false
            end
        else
            noimprov = 0
            objbest = obj0
            raised = false
        end
        i = i + 1
    end
    X=zeros(m*n,k)
    function findbeta(Xrow, Arow)
        return inv(I/γ+Xrow'*Xrow)*Xrow'*Arow
    end
    for i=1:k
        X[:,i]=reshape(x[:,i]*y[:,i]',n*m)
    end
    function PopulateA(Xtemp,Atemp)
        return X*(pinv(Xtemp'*Xtemp,1e-7)*(Xtemp'*Atemp))
    end
    Aopt=zeros(n,m,p)
    Wpar=Array{Float64}[Wmat[:,i] for i=1:p]
    Xpar=Array{Float64}[X[Wpar[i].==1,:] for i=1:p]
    Apar=Array{Float64}[Amat[Wpar[i].==1,i] for i=1:p]
    result=map(PopulateA,Xpar,Apar)
    println("Model Solved")
    for i=1:p
        Aopt[:,:,i]=reshape(result[i],(n,m))
    end
    return Aopt
end

function TensorDeriv(Amat,Wmat,x,y,k,γ,M,j,dim)
    n,m,p = dim
    ncube = (n*m*p)^(1/3)
    nsamples = ncube * k * log(ncube)^2 * 4 * 2^j / (1-M)
    α = nsamples / (n*m*p)
    pnew = min(max(100, Int(round(α * p))), p)
    samplep = sample(1:p,pnew,replace=false)
    normfactor = n * m * pnew
    function SmallInv(Xrow, Wrow, Arow)
        return Arow- Xrow*(inv(I/γ+Xrow'*Xrow)*(Xrow'*Arow))
    end
    X = zeros(n*m,k)
    for i=1:k
        X[:,i] = reshape(x[:,i]*y[:,i]',n*m)
    end
    Wpar=Array{Float64}[Wmat[:,samplep[i]] for i=1:pnew]
    Xpar=Array{Float64}[X[Wpar[i].==1,:] for i=1:pnew]
    Apar=Array{Float64}[Amat[Wpar[i].==1,samplep[i]] for i=1:pnew]
    objpar=map(SmallInv,Xpar,Wpar,Apar)
    obj=@distributed (+) for i=1:pnew
        dot(Apar[i],objpar[i]) / normfactor
    end
    ∇objx=zeros(n,k)
    ∇objy=zeros(m,k)
    ∇obj=zeros(n*m,k)
    # Btemp1 = zeros(p, k)
    for i = 1:pnew
        # Btemp1[i, :] = objpar[i]' * Xpar[i]
        ∇obj[Wpar[i].==1,:] += -2 * γ * objpar[i] * (objpar[i]' * Xpar[i]) / normfactor
    end

    ∇obj = reshape(∇obj, (n,m,k))
    for i = 1:k
        ∇objx[:,i]=∇obj[:,:,i] * y[:,i]
    end
    for i = 1:k
        ∇objy[:,i]=∇obj[:,:,i]' * x[:,i]
    end

    return obj, ∇objx, ∇objy
end

export Impute

end # module


#!/usr/bin/env julia

using Base.Test
using Drs.DrsMathProgSolverInterface

A = [2 3 1 0 0 0; -3 2 0 1 0 0; 0 2 0 0 1 0; 2 1 0 0 0 1]

m = DrsMathProgModel(A)
Drs.DrsMathProgSolverInterface.DrsFindPotentialBasis!(m)

@test m.basis == [3, 4, 5, 6]
@test m.nonbasis == [1, 2]


function llhfit(fit::FitFunction{T, 1, NP}, h::Histogram)::Optim.MultivariateOptimizationResults where {T <: AbstractFloat, NP}
    first_bin::Int = StatsBase.binindex(h, first(fit.fitranges[1]))
    last_bin::Int  = StatsBase.binindex(h, last(fit.fitranges[1]))
    if first_bin < 1 first_bin = 1 end
    if (last_bin > length(h.weights)) last_bin = length(h.weights) end
    bin_centers::Vector{T} = StatsBase.midpoints(h.edges[1])[first_bin:last_bin]
    bin_widths::Vector{T} = [StatsBase.binvolume(h, StatsBase.binindex(h, mp)) for mp in bin_centers]
    counts::Vector{T} = h.weights[first_bin:last_bin]

    function log_likelihood(params)
        s = 0
        if in(0, in.(params, fit.parameter_bounds))
            return Inf
        end
        @inbounds for i in eachindex(counts)
            expected_counts = fit.model(bin_centers[i], params) * bin_widths[i]
            if isnan(expected_counts) || expected_counts < 0
                expected_counts = T(Inf)
            end
            s += -logpdf(Poisson(expected_counts), counts[i])
        end
        return s
    end
    optim_result = Optim.optimize( log_likelihood, fit.initial_parameters )

    uncertainties = try
        sqrt.(diag(inv(ForwardDiff.hessian(log_likelihood, Optim.minimizer(optim_result)))))
    catch e
        zeros(T, length(fit.initial_parameters))
    end

    _set_fit_backend_result!(fit, (optim_result, uncertainties))
    _set_fitted_parameters!(fit, optim_result.minimizer)
    _set_bin_widths!(fit, bin_widths)
    _set_bin_centers!(fit, bin_centers)
    _set_residuals!(fit, bin_centers, counts)
    _set_Χ²!(fit)
    optim_result
end

"""
    llhfit!(fit::FitFunction{T, 1, NP}, h::Histogram)::Nothing where {T <: AbstractFloat, NP}

Performs a log-likelihood fit of the model function `fit.model` and the initial parameters `fit.initial_parameters`
on the histogram `h` in the range `fit.fitranges[1]`. The determined parameters are stored in `fit.fitted_parameters`.

The likelihood for each individual bin is the Poission distribution.
"""
function llhfit!(fit::FitFunction{T, 1, NP}, h::Histogram) where {T <: AbstractFloat, NP}
    llhfit(fit, h)
    fit
end

get_fit_backend_result(fr::Tuple{<:Optim.MultivariateOptimizationResults,<:Any}) = fr[1]

_get_standard_deviations(f::FitFunction, fr::Tuple{<:Optim.MultivariateOptimizationResults,<:Any}) = fr[2]


# Author: Ashwani Rathee
#  Program to Check Whether Number is Even or Odd
println("Enter number to check if number is even or odd  : ")
number=readline()
number=parse(Int64,number)
check=number%2
if check == 1
    println("The Number is odd")
elseif check == 0
    println("The Number is even")
else
    println("Don't Know what this stuff is")
end


# This code by Tara Abrishami

export transitive_orientation

"""
`transitive_orientation(G)` finds a transitive orientation of the
simple graph `G`. The result is a `SimpleDigraph`. An error is raised
if `G` does not have a transitive orientation.
"""
function transitive_orientation(G::SimpleGraph)
  err_msg = "This graph does not have a transitive orientation"
  vertices = deepcopy(vlist(G))
  edges = deepcopy(elist(G))
  V = vertex_type(G)
  diredges = Tuple{V,V}[]
  D = SimpleDigraph{V}()
  while length(diredges) != 0 || length(edges) != 0
    if length(diredges) == 0
      e = popfirst!(edges)
      add!(D, e[1], e[2])
    else
      e = popfirst!(diredges)
    end
    v1 = e[1]
    v2 = e[2]
    vs = V[]
    append!(vs,G[v1])
    append!(vs,G[v2])
    for v in vs
      if has(G, v1, v) && !has(D, v1, v) && !has(G, v, v2)
        if has(D, v, v1)
          error(err_msg)
        end
        add!(D, v1, v)
        edg = (v1, v)
        swapedg = (v, v1)
        pushfirst!(diredges, edg)
        edges = filter(x -> x != edg && x != swapedg,edges)
      elseif has(G, v2, v) && !has(D, v, v2) && !has(G, v, v1)
        if has(D, v2, v)
          error(err_msg)
        end
        add!(D, v, v2)
        edg = (v, v2)
        swapedg = (v2, v)
        pushfirst!(diredges, edg)
        edges = filter(x -> x != edg && x != swapedg,edges)
      end
    end
  end
  if length(elist(D)) == length(elist(G))
    return D
  else
    error(err_msg)
  end
end

export num_trans_orientations

"""
`num_trans_orientations(G)` returns the number of transitive
orientations of the graph `G`.
"""
function num_trans_orientations(G2::SimpleGraph)
  if cache_check(G2,:num_trans_orientations)
    return cache_recall(G2,:num_trans_orientations)
  end
  G = deepcopy(G2)
  V = vertex_type(G)
  col = Dict{Tuple{V,V}, Int}()
  try
    col = makeColorClass(G)
  catch
    return 0
  end
  multiplexes = BigInt[]
  while length(elist(G)) != 0
    makeSimplex!(G, multiplexes, col)
  end
  ans = 1
  for m in multiplexes
    ans = ans*factorial(m)
  end
  cache_save(G2,:num_trans_orientations,ans)
  return ans
end

function makeSimplex!(G::SimpleGraph, multiplexes::Array, col::Dict)
  edge = elist(G)[1]
  vert1 = edge[1]
  vert2 = edge[2]
  colorSet = Set()
  T = vertex_type(G)
  simp = Set{T}()
  push!(simp, vert1)
  push!(simp, vert2)
  push!(colorSet, col[edge])
  for v in neighbors(G, vert1)
    colorSetTest = deepcopy(colorSet)
    isEdge = true
    for vo in simp
      if has(G, vo, v) && (((haskey(col, (v, vo)) && !(col[(v, vo)] in colorSetTest)) || (haskey(col, (vo,v)) &&!(col[(vo, v)] in colorSetTest))))
        if haskey(col, (v, vo))
          push!(colorSetTest, col[(v, vo)])
        else
          push!(colorSetTest, col[(vo, v)])
        end
      else
        isEdge = false
        break
      end
    end
    if isEdge == true
      push!(simp, v)
      colorSet = deepcopy(colorSetTest)
    end
  end
  pushfirst!(multiplexes, length(simp))
  for e in elist(G)
    if col[e] in colorSet
      delete!(G, e[1], e[2])
    end
  end
end

function makeColorClass(G1::SimpleGraph)
  err_msg = "error"
  G = deepcopy(G1)
  vertices = deepcopy(vlist(G))
  edge = deepcopy(elist(G))
  V = vertex_type(G)
  diredge = Tuple{V,V}[]
  D = SimpleDigraph{V}()
  E = SimpleGraph{V}()
  classes = SimpleGraph{V}[]
  while length(diredge) != 0 || length(edge) != 0
    if length(diredge) == 0
      if length(elist(E)) != 0
        pushfirst!(classes,E)
        E1 = SimpleGraph{V}()
        E = E1
      end
      e = popfirst!(edge)
      add!(D, e[1], e[2])
      add!(E, e[1], e[2])
    else
      e = popfirst!(diredge)
    end
    v1 = e[1]
    v2 = e[2]
    vs = V[]
    append!(vs,G[v1])
    append!(vs,G[v2])
    for v in vs
      if has(G, v1, v) && !has(D, v1, v) && !has(G, v, v2)
        if has(D, v, v1)
          error(err_msg)
        end
        add!(D, v1, v)
        add!(E, v1, v)
        edg = (v1, v)
        swapedg = (v, v1)
        pushfirst!(diredge, edg)
        edge = filter(x -> x != edg && x != swapedg,edge)
      elseif has(G, v2, v) && !has(D, v, v2) && !has(G, v, v1)
        if has(D, v2, v)
          error(err_msg)
        end
        add!(D, v, v2)
        add!(E, v, v2)
        edg = (v, v2)
        swapedg = (v2, v)
        pushfirst!(diredge, edg)
        edge = filter(x -> x != edg && x != swapedg,edge)
      end
    end
  end
  if length(elist(E)) != 0
    pushfirst!(classes,E)
  end
  col = Dict{Tuple{V,V}, Int}()
  i = 1
  for q in classes
    for x in elist(q)
      col[x] = i
    end
    i = i + 1
  end
  return col
end


function matrix_cone(S::Type{<:Union{MatrixConeTriangle,
                                     MOI.PositiveSemidefiniteConeTriangle}},
                     side_dimension)
    if iszero(side_dimension)
        # Some solvers such as Mosek does not support 0-dimensional PSD cone
        return EmptyCone()
    elseif isone(side_dimension)
        # PSD constraints on 1x1 matrices are equivalent to the nonnegativity
        # of the only entry.
        return MOI.Nonnegatives(1)
    elseif side_dimension == 2
        # PSD constraints on 2x2 matrices are SOC representable.
        return PositiveSemidefinite2x2ConeTriangle()
    else
        # PSD constraints on nxn matrices with n ≥ 3 is not SOC representable,
        # see [F18].
        #
        # [F18] Fawzi, Hamza
        # On representing the positive semidefinite cone using the second-order
        # cone.
        # Mathematical Programming (2018): 1-10.
        return S(side_dimension)
    end
end

"""
    struct DiagonallyDominantConeTriangle <: MatrixConeTriangle
        side_dimension::Int
    end

See Definition 4 of [AM17] for a precise definition of the last two items.

[AM17] Ahmadi, A. A. & Majumdar, A.
*DSOS and SDSOS Optimization: More Tractable Alternatives to Sum of Squares and Semidefinite Optimization*
ArXiv e-prints, **2017**.
"""
struct DiagonallyDominantConeTriangle <: MatrixConeTriangle
    side_dimension::Int
end

function matrix_cone(S::Type{DiagonallyDominantConeTriangle},
                     side_dimension)
    if iszero(side_dimension)
        return EmptyCone()
    elseif isone(side_dimension)
        return MOI.Nonnegatives(1)
    else
        # With `side_dimension` = 2, we want to avoid using SOC and only use LP
        return S(side_dimension)
    end
end

"""
    struct ScaledDiagonallyDominantConeTriangle <: MatrixConeTriangle
        side_dimension::Int
    end

See Definition 4 of [AM17] for a precise definition of the last two items.

[AM17] Ahmadi, A. A. & Majumdar, A.
*DSOS and SDSOS Optimization: More Tractable Alternatives to Sum of Squares and Semidefinite Optimization*
ArXiv e-prints, **2017**.
"""
struct ScaledDiagonallyDominantConeTriangle <: MatrixConeTriangle
    side_dimension::Int
end

function matrix_cone(S::Type{ScaledDiagonallyDominantConeTriangle},
                     side_dimension)
    if iszero(side_dimension)
        return EmptyCone()
    elseif isone(side_dimension)
        return MOI.Nonnegatives(side_dimension)
    elseif side_dimension == 2
        return PositiveSemidefinite2x2ConeTriangle()
    else
        return S(side_dimension)
    end
end

function side_dimension(set::Union{MatrixConeTriangle,
                                   MOI.PositiveSemidefiniteConeTriangle})
    return set.side_dimension
end

# isbits types, nothing to copy
function Base.copy(set::MatrixConeTriangle)
    return set
end

# TODO make PSDConeTriangle inherit from MatrixConeTriangle to remove the need
#      for this
function MOI.dimension(set::MatrixConeTriangle)
    return div(side_dimension(set) * (side_dimension(set) + 1), 2)
end

function MOIU.set_dot(x::Vector, y::Vector, set::MatrixConeTriangle)
    return MOIU.triangle_dot(x, y, side_dimension(set), 0)
end

function MOIU.dot_coefficients(a::Vector, set::MatrixConeTriangle)
    b = copy(a)
    MOIU.triangle_coefficients!(b, side_dimension(set), 0)
    return b
end


using SparseArrays, Test
using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

@testset "Diagonal element" begin
    k = 0
    for j in 1:10
        for i in 1:j
            k += 1
            @test MOIU.is_diagonal_vectorized_index(k) == (i == j)
        end
    end
    for side_dim in 1:10
        set = MOI.PositiveSemidefiniteConeTriangle(side_dim)
        vec_dim = MOI.dimension(set)
        @test MOIU.side_dimension_for_vectorized_dimension(vec_dim) == side_dim
    end
end

@testset "Side dimension" begin
    for side_dim in 1:10
        set = MOI.PositiveSemidefiniteConeTriangle(side_dim)
        vec_dim = MOI.dimension(set)
        @test MOIU.side_dimension_for_vectorized_dimension(vec_dim) == side_dim
    end
end

@testset "Constant" begin
    @test MOI.constant(MOI.EqualTo(3)) == 3
    @test MOI.constant(MOI.GreaterThan(6)) == 6
    @test MOI.constant(MOI.LessThan(2)) == 2
end

@testset "Shifts" begin
    @test MOIU.supports_shift_constant(MOI.EqualTo{Int})
    @test MOIU.shift_constant(MOI.EqualTo(3), 1) == MOI.EqualTo(4)
    @test MOIU.supports_shift_constant(MOI.GreaterThan{Int})
    @test MOIU.shift_constant(MOI.GreaterThan(6), -1) == MOI.GreaterThan(5)
    @test MOIU.supports_shift_constant(MOI.LessThan{Int})
    @test MOIU.shift_constant(MOI.LessThan(2), 2) == MOI.LessThan(4)
    @test MOIU.supports_shift_constant(MOI.Interval{Int})
    @test MOIU.shift_constant(MOI.Interval(-2, 3), 1) == MOI.Interval(-1, 4)
    @test MOIU.supports_shift_constant(MOI.ZeroOne) == false
    @test_throws MethodError MOIU.shift_constant(MOI.ZeroOne(), 1.0)
end

@testset "Dimension" begin
    @test MOI.dimension(MOI.EqualTo(3.0)) === 1
    @test MOI.dimension(MOI.Reals(8)) === 8
    @test MOI.dimension(MOI.NormInfinityCone(5)) === 5
    @test MOI.dimension(MOI.NormOneCone(5)) === 5
    @test MOI.dimension(MOI.DualExponentialCone()) === 3
    @test MOI.dimension(MOI.RelativeEntropyCone(5)) === 5
    @test MOI.dimension(MOI.NormSpectralCone(2, 3)) === 7
    @test MOI.dimension(MOI.NormNuclearCone(2, 3)) === 7
    @test MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(4)) === 10
    @test MOI.dimension(MOI.PositiveSemidefiniteConeSquare(5)) === 25
    @test MOI.dimension(MOI.RootDetConeTriangle(6)) === 22
    @test MOI.dimension(MOI.LogDetConeTriangle(6)) === 23
    @test MOI.dimension(MOI.RootDetConeSquare(4)) === 17
    @test MOI.dimension(MOI.LogDetConeSquare(4)) === 18
    @test MOI.dimension(MOI.SOS2(collect(1:6))) === 6
end

@testset "Set dot" begin
    vec = zeros(6)
    @test MOIU.set_dot(vec, vec, MOI.SecondOrderCone(6)) == 0
    @test MOIU.set_dot(vec, vec, MOI.PositiveSemidefiniteConeTriangle(3)) == 0
    @test MOIU.set_dot(vec, vec, MOI.PositiveSemidefiniteConeSquare(3)) == 0
    vec[5] = 1
    @test MOIU.set_dot(vec, vec, MOI.SecondOrderCone(6)) == 1
    @test MOIU.set_dot(vec, vec, MOI.PositiveSemidefiniteConeTriangle(3)) == 2
    @test MOIU.set_dot(vec, vec, MOI.PositiveSemidefiniteConeSquare(3)) == 1
    vec[5] = 0
    vec[3] = 1
    @test MOIU.set_dot(vec, vec, MOI.SecondOrderCone(6)) == 1
    @test MOIU.set_dot(vec, vec, MOI.PositiveSemidefiniteConeTriangle(3)) == 1
    @test MOIU.set_dot(vec, vec, MOI.PositiveSemidefiniteConeSquare(3)) == 1

    vec = zeros(7)
    @test MOIU.set_dot(vec, vec, MOI.RootDetConeTriangle(3)) == 0
    vec[5] = 1
    @test MOIU.set_dot(vec, vec, MOI.RootDetConeTriangle(3)) == 2
    vec = zeros(8)
    @test MOIU.set_dot(vec, vec, MOI.LogDetConeTriangle(3)) == 0
    vec[5] = 1
    @test MOIU.set_dot(vec, vec, MOI.LogDetConeTriangle(3)) == 1

    sp_vec = spzeros(6)
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.SecondOrderCone(6)) == 0
    @test MOIU.set_dot(
        sp_vec,
        sp_vec,
        MOI.PositiveSemidefiniteConeTriangle(3),
    ) == 0
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.PositiveSemidefiniteConeSquare(3)) ==
          0
    sp_vec[5] = 1
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.SecondOrderCone(6)) == 1
    @test MOIU.set_dot(
        sp_vec,
        sp_vec,
        MOI.PositiveSemidefiniteConeTriangle(3),
    ) == 2
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.PositiveSemidefiniteConeSquare(3)) ==
          1
    sp_vec[5] = 0
    sp_vec[3] = 1
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.SecondOrderCone(6)) == 1
    @test MOIU.set_dot(
        sp_vec,
        sp_vec,
        MOI.PositiveSemidefiniteConeTriangle(3),
    ) == 1
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.PositiveSemidefiniteConeSquare(3)) ==
          1

    sp_vec = spzeros(7)
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.RootDetConeTriangle(3)) == 0
    sp_vec[5] = 1
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.RootDetConeTriangle(3)) == 2
    sp_vec = spzeros(8)
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.LogDetConeTriangle(3)) == 0
    sp_vec[5] = 1
    @test MOIU.set_dot(sp_vec, sp_vec, MOI.LogDetConeTriangle(3)) == 1
end

@testset "dot coefficients" begin
    vec = zeros(6)
    @test MOIU.dot_coefficients(vec, MOI.SecondOrderCone(6)) == vec
    @test MOIU.dot_coefficients(vec, MOI.PositiveSemidefiniteConeTriangle(3)) ==
          vec
    vec[5] = 1
    @test MOIU.dot_coefficients(vec, MOI.SecondOrderCone(6)) == vec
    @test MOIU.dot_coefficients(vec, MOI.PositiveSemidefiniteConeTriangle(3)) ==
          vec ./ 2
    vec[5] = 0
    vec[3] = 1
    @test MOIU.dot_coefficients(vec, MOI.SecondOrderCone(6)) == vec
    @test MOIU.dot_coefficients(vec, MOI.PositiveSemidefiniteConeTriangle(3)) ==
          vec

    vec = zeros(7)
    @test MOIU.dot_coefficients(vec, MOI.RootDetConeTriangle(3)) == vec
    vec[5] = 1
    @test MOIU.dot_coefficients(vec, MOI.RootDetConeTriangle(3)) == vec ./ 2
    vec = zeros(8)
    @test MOIU.dot_coefficients(vec, MOI.LogDetConeTriangle(3)) == vec
    vec[5] = 1
    @test MOIU.dot_coefficients(vec, MOI.LogDetConeTriangle(3)) == vec

    sp_vec = spzeros(6)
    @test MOIU.dot_coefficients(sp_vec, MOI.SecondOrderCone(6)) == sp_vec
    @test MOIU.dot_coefficients(
        sp_vec,
        MOI.PositiveSemidefiniteConeTriangle(3),
    ) == sp_vec
    sp_vec[5] = 1
    @test MOIU.dot_coefficients(sp_vec, MOI.SecondOrderCone(6)) == sp_vec
    @test MOIU.dot_coefficients(
        sp_vec,
        MOI.PositiveSemidefiniteConeTriangle(3),
    ) == sp_vec ./ 2
    sp_vec[5] = 0
    sp_vec[3] = 1
    @test MOIU.dot_coefficients(sp_vec, MOI.SecondOrderCone(6)) == sp_vec
    @test MOIU.dot_coefficients(
        sp_vec,
        MOI.PositiveSemidefiniteConeTriangle(3),
    ) == sp_vec

    sp_vec = spzeros(7)
    @test MOIU.dot_coefficients(sp_vec, MOI.RootDetConeTriangle(3)) == sp_vec
    sp_vec[5] = 1
    @test MOIU.dot_coefficients(sp_vec, MOI.RootDetConeTriangle(3)) ==
          sp_vec ./ 2
    sp_vec = spzeros(8)
    @test MOIU.dot_coefficients(sp_vec, MOI.LogDetConeTriangle(3)) == sp_vec
    sp_vec[5] = 1
    @test MOIU.dot_coefficients(sp_vec, MOI.LogDetConeTriangle(3)) == sp_vec
end


#!/usr/bin/env bash
    #=
    exec julia --project="$(realpath $(dirname $(dirname $0)))" --color=yes --startup-file=no -e "include(popfirst!(ARGS))" \
    "${BASH_SOURCE[0]}" "$@"
    =#

"""
```julia
struct FinitePolynomial <: FiniteField
```

Has parameters `p`, which is an abstract polynomial, and `n` which is the modulus of the field under which the molynomial is defined.

---

```julia
FinitePolynomial(p::AbstractPolynomial, n::Int)
```

A constructor method for `FinitePolynomial`.  Takes in a polynomial `p` and a number `n`, and constructs a polynomial under modulo n.
"""
struct FinitePolynomial <: FiniteField
	p::AbstractPolynomial
	n::Int

	function FinitePolynomial(p::AbstractPolynomial, n::Int)
		p = Polynomial(mod.(p.coeffs, n))
		new(p, n)
	end
end

"""
```julia
mod(p::Polynomial, n::Int) -> Polynomial
```

Uses the `FinitePolynomial` constructor to return a polynomial `p` under modulus `n`.

Parameters:
  - `p::Polynomial`: The input polynomial.
  - `n::Int`: The modulus of the field.

Returns
  - Polynomial: A polynomial modulo n.
"""
Base.mod(p::Polynomial, n::Int) = FinitePolynomial(p, n).p

"""
```julia
Polynomial(A::Union{NTuple{N, T}, Vector{T}}, n::Int) -> Polynomial
```

Constructs a polynomial under modulo `n`.

Parameters:
  - `A::Union{Tuple, AbstractArray}`: The polynomial coefficients.
  - `n::Int`: The modulus of the field.

Returns
  - `Polynomial`: A polynomial modulo n.
"""
Polynomial(A::Union{NTuple{N, T}, Vector{T}}, n::Int) where {N, T} = mod(Polynomial(A), n)

"""
```julia
list_polys(n::Int, m::Int) -> Array
```

Lists all polynomials of degree less than to `n` under modulo `m`.

Parameters:
  - `n::Int`: Highest degree of polynomial.
  - `m::Int`: The modulus of the field.

Returns:
  - `Array`: An array of polynomials of degree less than n, under modulo m.
"""
function list_polys(n::Int, m::Int)
	return collect(Polynomial(collect(t)) for t in Iterators.product([0:(m-1) for i in 1:n]...))
end

"""
```julia
multiplication_table(degree::Int, modulo::Int) -> Matrix
```

Returns a table (matrix) of the multiplication of all combinations of polynomials for degree less than `degree`, under modulo `modulo`.

Parameters:
  - `degree::Int`: Highest degree of polynomial.
  - `modulo::Int`: The modulus of the field.

Returns:
  - `Matrix`: A multiplication table of all polynomials with degree less than n, under modulus.

---

### Examples

```julia
julia> julia> multiplication_table(2, 3) # multiplication table of all polynomials of degree less than 3 modulo 2
9×9 Array{Polynomial,2}:
 Polynomial(0)  Polynomial(0)        Polynomial(0)        …  Polynomial(0)                Polynomial(0)
 Polynomial(0)  Polynomial(1)        Polynomial(2)           Polynomial(1 + 2*x)          Polynomial(2 + 2*x)
 Polynomial(0)  Polynomial(2)        Polynomial(1)           Polynomial(2 + x)            Polynomial(1 + x)
 Polynomial(0)  Polynomial(x)        Polynomial(2*x)         Polynomial(x + 2*x^2)        Polynomial(2*x + 2*x^2)
 Polynomial(0)  Polynomial(1 + x)    Polynomial(2 + 2*x)     Polynomial(1 + 2*x^2)        Polynomial(2 + x + 2*x^2)
 Polynomial(0)  Polynomial(2 + x)    Polynomial(1 + 2*x)  …  Polynomial(2 + 2*x + 2*x^2)  Polynomial(1 + 2*x^2)
 Polynomial(0)  Polynomial(2*x)      Polynomial(x)           Polynomial(2*x + x^2)        Polynomial(x + x^2)
 Polynomial(0)  Polynomial(1 + 2*x)  Polynomial(2 + x)       Polynomial(1 + x + x^2)      Polynomial(2 + x^2)
 Polynomial(0)  Polynomial(2 + 2*x)  Polynomial(1 + x)       Polynomial(2 + x^2)          Polynomial(1 + 2*x + x^2)
```
"""
function multiplication_table(degree::Int, modulo::Int)
	polys = list_polys(degree, modulo)
	number_of_polys = length(polys)
	poly_matrix = Matrix{Polynomial}(undef, number_of_polys, number_of_polys)

	for i in 1:number_of_polys, j in 1:number_of_polys
		poly_matrix[i,j] = mod(polys[i] * polys[j], modulo)
	end

	return poly_matrix
end

function __list_span_inner(modulo::Int, u̲::Vector{T}...) where T <: Number
    n_vec = length(u̲)
    span = Vector{T}[zero(u̲[1])]

    for n in 1:n_vec
        new_span = copy(span)
        for v in span, λ in 1:(modulo - 1)
            w̲ = mod.(v + λ * u̲[n], modulo)
            if w̲ ∉ new_span
                push!(new_span, w̲)
            end
        end
        span = new_span
    end

	return span
end

"""
```julia
list_span(u̲::Vector, v̲::Vector,... modulo::Int) -> Array
```

Given any number of vectors `u̲`, `v̲`,... , prints all linear combinations of those vectors, modulo `modulo`.

Parameters:
  - `u̲::Vector`: One vector.
  - `v̲::Vector`: Another vector.
  - `...`: Other vectors.
  - `modulo::Int`: The modulus of the field.

Returns:
  - `Array`: All vectors in the span of u̲ and v̲, under modulo.

---

### Examples

```julia
julia> list_span([2, 1, 1], [1, 1, 1], 3) # list the span of two vectors modulo 3
9-element Array{Array{T,1} where T,1}:
 [0, 0, 0]
 [1, 1, 1]
 [2, 2, 2]
 [2, 1, 1]
 [0, 2, 2]
 [1, 0, 0]
 [1, 2, 2]
 [2, 0, 0]
 [0, 1, 1]
```
"""
list_span(args...) = __list_span_inner(last(args), args[1:end-1]...)

"""
```julia
islinear(C::Vector, modulo::Int; verbose::Bool = false) -> Bool
```

Determines whether a code `C` is a linear code (i.e., if it is closed under addition, scalar multiplication, and has the zero vector in it).

Parameters:
  - `C::Vector`: A code, typically consisting of multiple vectors or strings.
  - `modulo::Int`: The modulus of the field under which you are working.
  - `verbose::Bool` (kwarg): print the point at which C fails, if it does.

Returns:
  - `Bool`: Whether or not the code `C` is linear (true or false).

---

### Examples

```julia

julia> islinear([[0,0,0],[1,1,1],[1,0,1],[1,1,0]], 2) # checks whether a vector of vectors is linear/a subspace (modulo 2)
false
```
"""
function islinear(C::Vector, modulo::Int; verbose::Bool = false)
	allequal_length(C) || return false # not all codes are of the same length
	block_length = length(C[1])
	𝟎 = fill(0, block_length)

	if 𝟎 ∉ C
		verbose && println("The zero vector 0̲ is not in C.\n")
		return false # the zero vector is not in the code
	end

	for c̲ ∈ C
		for λ in 0:modulo-1
			if mod.(λ*c̲, modulo) ∉ C
				verbose && println(λ, " ⋅ ", c̲, " = ", mod.(λ*c̲, modulo), " ∉ C\n")
				return false # this code isn't closed under scalar multiplication
			end
		end

		for c̲ ∈ C, c̲′ ∈ C
			if c̲ ≠ c̲′
				if mod.(c̲ + c̲′, modulo) ∉ C
					verbose && println(c̲, " + ", c̲′, " = ", mod.(c̲ + c̲′, modulo), " ∉ C\n")
					return false # this code isn't closed under addition
				end
			end
		end
	end

	verbose && println()

	return true
end

"""
```julia
isirreducible(f::AbstractPolynomial, modulo::Int) -> Bool
```

Checks if a polynomial is irreducible.

Parameters:
  - `f::Polynomial`: The polynomial you need to check.
  - `modulo::Int`: The modulus under which you are working.

Returns:
  - `Bool`: Whether or not the polynomial is irreducible (true or false).

---

### Examples

```julia
julia> isirreducible(Polynomial([1, 1, 0, 0, 1]), 2) # is 1 + x + x^4 mod 2 irreducible?
true
```
"""
function isirreducible(f::Polynomial, modulo::Int)
	deg = length(f.coeffs) - 1
	f = mod(f, deg)
	polys = list_polys(deg, modulo)

	for a in polys, b in polys
		isequal(f, mod(a*b, modulo)) && return false
	end

	return true
end

"""
```julia
normal_form!(M::AbstractArray{Int}, n::Int) -> Matrix{Int}
```

Convert a matrix M into normal form under modulo n via Gauss-Jordan elimination.  *This directly changes the matrix M.  Use `normal_form` for a non-mutating version of this function.*

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.

Returns:
  - `Matrix{Int}`: A matrix in normal form from Gauss-Jordan elimination.

---

### Examples

```julia

julia> normal_form!([1 2 0 1 2 1 2; 2 2 2 0 1 1 1; 1 0 1 1 2 1 2; 0 1 0 1 1 2 2], 3) # computes rref colswap = false
4×7 Array{Int64,2}:
1  0  0  0  2  2  2
0  1  0  0  2  0  1
0  0  1  0  1  0  2
0  0  0  1  2  2  1
```
"""
function normal_form!(M::AbstractArray{Int}, n::Int)
	return rref!(M, n, colswap = false)
end

"""
```julia
normal_form(M::AbstractArray{Int}, n::Int) -> Matrix{Int}
```

Convert a matrix M into normal form under modulo n via Gauss-Jordan elimination.

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.

Returns:
  - `Matrix{Int}`: A matrix in normal form from Gauss-Jordan elimination.

---

### Examples

```julia

julia> normal_form([1 2 0 1 2 1 2; 2 2 2 0 1 1 1; 1 0 1 1 2 1 2; 0 1 0 1 1 2 2], 3) # computes rref colswap = false
4×7 Array{Int64,2}:
1  0  0  0  2  2  2
0  1  0  0  2  0  1
0  0  1  0  1  0  2
0  0  0  1  2  2  1
```
"""
function normal_form(M::AbstractArray{Int}, n::Int)
	normal_form!(copy(M), n)
end

"""
```julia
equivalent_code!(M::AbstractArray{Int}, n::Int) -> Matrix{Int}
```

Peforms Gauss-Jordan elimination on a matrix M, but allows for column swapping.  This constructs an "equivalent" matrix.  *This directly changes the matrix M.  Use `equivalent_code` for a non-mutating version of this function.*

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.

Returns:
  - `Matrix{Int}`: A which represents an "equivalent" code to that of the matrix M.

---

### Examples

```julia
julia> equivalent_code!([1 2 0 1 2 1 2; 2 2 2 0 1 1 1; 1 0 1 1 2 1 2; 0 1 0 1 1 2 2], 3) # computes rref colswap = true
4×7 Array{Int64,2}:
1  0  0  0  2  2  2
0  1  0  0  2  0  1
0  0  1  0  1  0  2
0  0  0  1  2  2  1
```
"""
function equivalent_code!(M::AbstractArray{Int}, n::Int)
	return rref!(M, n, colswap=true)
end

"""
```julia
equivalent_code(M::AbstractArray{Int}, n::Int) -> Matrix{Int}
```

Peforms Gauss-Jordan elimination on a matrix M, but allows for column swapping.

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.

Returns:
  - `Matrix{Int}`: A which represents an "equivalent" code to that of the matrix M.

---

### Examples

```julia
julia> equivalent_code([1 2 0 1 2 1 2; 2 2 2 0 1 1 1; 1 0 1 1 2 1 2; 0 1 0 1 1 2 2], 3) # computes rref colswap = true
4×7 Array{Int64,2}:
1  0  0  0  2  2  2
0  1  0  0  2  0  1
0  0  1  0  1  0  2
0  0  0  1  2  2  1
```
"""
function equivalent_code(M::AbstractArray{Int}, n::Int)
	equivalent_code!(copy(M), n)
end

"""
```julia
generator!(M::AbstractArray{Int}, n::Int; colswap::Bool = true) -> Matrix{Int}
```

Constructs a generator matrix of the code, depending on if you allow for column swapping or not.   This function uses `normal_form!` or `equivalent_code!`.  *This directly changes the matrix M.  Use `generator` for a non-mutating version of this function.*

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.
  - `colswap::Bool` (kwarg): A boolean flag indicating whether or not you allow for swapping of columns when constructing the generating matrix.

Returns:
  - `Matrix{Int}`: A generating matrix.
"""
function generator!(M::AbstractArray{Int}, n::Int; colswap::Bool = false)
	return ifelse(colswap, equivalent_code!(M, n), normal_form!(M, n))
end

"""
```julia
generator(M::AbstractArray{Int}, n::Int; colswap::Bool = true) -> Matrix{Int}
```

Constructs a generator matrix of the code, depending on if you allow for column swapping or not.  This function uses `normal_form` or `equivalent_code`.

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.
  - `colswap::Bool` (kwarg): A boolean flag indicating whether or not you allow for swapping of columns when constructing the generating matrix.

Returns:
  - `Matrix{Int}`: A generating matrix.
"""
function generator(M::AbstractArray{Int}, n::Int; colswap::Bool = true)
	generator!(copy(M), n, colswap=colswap)
end

"""
```julia
parity_check(M::AbstractArray{Int}, n::Int) -> Matrix{Int}
```

Constructs a parity check matrix.  This is calculated from taking the non-identity part of a matrix in normal form (or equivalent — see `generator`), transposing it, multiplying it by negative one, and appending to it an appropriate sized identity matrix.

Parameters:
  - `M::AbstractArray{Int}`: A matrix of Ints.
  - `n::Int`: The modulus of the finite field.

Returns:
  - `Matrix{Int}`: A parity check matrix.
"""
function parity_check(M::AbstractArray{Int}, n::Int)
	has_identity_on_left(M) || throw(error("This matrix is not in normal form.  Use normal_form or equivalent_code."))

	minus_Dᵀ = mod.(-one(Int) .* transpose(M[:, (size(M, 1) + 1):end]), n)
	return H = Int[minus_Dᵀ I(size(minus_Dᵀ, 1))]
end

"""
```julia
syndrome(v̲::Vector, Hᵀ::AbstractArray{Int}, n::Int) -> Matrix{Int}
```

Calculates the syndrome of a given word v̲ and a parity check matrix, transposed (Hᵀ), under modulo n.

Parameters:
  - `v̲::Vector`: A word in the code.
  - `Hᵀ::AbstractArray{Int}`: The transpose of a parity check matrix.

Returns:
  - `Vector`: The syndrome of a word in the code.

---

### Examples

```julia
julia> syndrome([0, 2, 1, 2, 0, 1, 0], [1 1 1; 1 0 2; 2 0 1; 1 1 2; 1 0 0; 0 1 0; 0 0 1], 3)
1×3 Array{Int64,2}:
0  0  0

julia> syndrome([0, 2, 1, 2, 0, 1, 0], transpose(parity_check([1 0 0 0 2 2 2; 0 1 0 0 2 0 1; 0 0 1 0 1 0 2; 0 0 0 1 2 2 1], 3)), 3)
1×3 Array{Int64,2}:
 0  0  0
```
"""
function syndrome(v̲::Vector, Hᵀ::AbstractArray{Int}, n::Int)
	return mod.(v̲' * Hᵀ, n)
end

"""
```julia
isincode(v̲::Vector, Hᵀ::AbstractArray{Int}, n::Int) -> Bool
```

If the syndrome of a code is the zero vector, then the word used to calculate the syndrome is in the code.

Parameters:
  - `v̲::Vector`: A word.
  - `Hᵀ::AbstractArray{Int}`: The transpose of a parity check matrix.

Returns:
  - `Bool`: If the word is in the code or not (true or false).

---

### Examples

```julia
julia> isincode([0, 2, 1, 2, 0, 1, 0], transpose(parity_check([1 0 0 0 2 2 2; 0 1 0 0 2 0 1; 0 0 1 0 1 0 2; 0 0 0 1 2 2 1], 3)), 3) # tests if the syndrome is equal to the zero vector, and is thus in the code
true
```
"""
function isincode(v̲::Vector, Hᵀ::AbstractArray{Int}, n::Int)
	return iszero(syndrome(v̲, Hᵀ, n))
end


using Test
using ImageCore
using IterTools
using ReferenceTests
using ImageDistances

# general distances should cover any combination of number_types and color_types unless it's special designed
include("testutils.jl")

include("hausdorff.jl")
include("metrics.jl")
include("ciede2000.jl")

nothing


using WoodburyMatrices
using Base.Test

srand(123)
n = 5

d = 1 .+ rand(n)
dl = -rand(n-1)
du = -rand(n-1)
v = randn(n)
B = randn(n,2)

U = randn(n,2)
V = randn(2,n)
C = randn(2,2)

for elty in (Float32, Float64, Complex64, Complex128, Int)
    if elty == Int
        srand(61516384)
        d = rand(1:100, n)
        dl = -rand(0:10, n-1)
        du = -rand(0:10, n-1)
        v = rand(1:100, n)
        B = rand(1:100, n, 2)

        # Woodbury
        U = rand(1:100, n, 2)
        V = rand(1:100, 2, n)
        C = rand(1:100, 2, 2)
    else
        d = convert(Vector{elty}, d)
        dl = convert(Vector{elty}, dl)
        du = convert(Vector{elty}, du)
        v = convert(Vector{elty}, v)
        B = convert(Matrix{elty}, B)
        U = convert(Matrix{elty}, U)
        V = convert(Matrix{elty}, V)
        C = convert(Matrix{elty}, C)
    end
    ε = eps(abs2(float(one(elty))))
    T = Tridiagonal(dl, d, du)
    @test size(T, 1) == n
    @test size(T) == (n, n)

    # Woodbury
    W = Woodbury(T, U, C, V)
    F = full(W)
    @test_approx_eq W*v F*v
    iFv = F\v
    @test norm(W\v - iFv)/norm(iFv) <= n*cond(F)*ε # Revisit. Condition number is wrong
    @test abs((det(W) - det(F))/det(F)) <= n*cond(F)*ε # Revisit. Condition number is wrong
    iWv = similar(iFv)
    if elty != Int
        iWv = A_ldiv_B!(W, copy(v))
        @test_approx_eq iWv iFv
    end
end

# Sparse U and V
n = 5
d = 1 .+ rand(n)
dl = -rand(n-1)
du = -rand(n-1)
T = Tridiagonal(dl, d, du)
U = sprandn(n,2,0.2)
V = sprandn(2,n,0.2)
C = randn(2,2)
W = Woodbury(T, U, C, V)
F = full(W)

v = randn(n)
ε = eps()
iFv = F\v
@test norm(W\v - iFv)/norm(iFv) <= n*cond(F)*ε

# Display
iob = IOBuffer()
show(iob, W)

# Vector-valued U and scalar-valued C
U = rand(n)
C = 2.3
V = rand(1,n)
W = Woodbury(T, U, C, V)
F = full(W)

v = randn(n)
ε = eps()
iFv = F\v
@test norm(W\v - iFv)/norm(iFv) <= n*cond(F)*ε

show(iob, W)

Wc = copy(W)

# Mismatched sizes
@test_throws DimensionMismatch Woodbury(rand(5,5),rand(5,2),rand(2,3),rand(3,5))
@test_throws DimensionMismatch Woodbury(rand(5,5),rand(5,2),rand(3,3),rand(2,5))

# spec builder
n = 10
num_nz = 3
for i in 1:5  # repeat 5 times
    args = [(rand(1:n), rand(1:n), rand()) for j in 1:num_nz]
    r, v, c = WoodburyMatrices.sparse_factors(Float64, n, args...)
    out = r*v*c

    by_hand = zeros(n, n)
    for (i, j, v) in args
        by_hand[i, j] = v
    end

    @test maxabs(out - by_hand) == 0.0
end

include("runtests_sym.jl")


#=------------------------------------------------------------------------------
   Functions typically associated with hypergraph operations, functions placed
   here should relate to how to operate on the abstract representation of the
   tensor. Note that these may be common operations in multi-linear algebra too.

  Operations Glossary
  -------------------
    * permute_tensor - ()
    * connected_components - ()
    * find_edge_incidence - ()

------------------------------------------------------------------------------=#

"""-----------------------------------------------------------------------------
    permute_tensor!(A,p)

  This function applies a permutation to the indices in the sparse tensor
defined by the array p. Iterates have the permutation applied and indices are
resorted.

Input:
------
* A - (SSSTensor)

    The tensor to apply the permutation to.
* p - (Array{Int,1})

    The permutation array, where p[i] = j implies that vertex j is mapped to
    index i.
Note:
-----

 Originally planned to be named permute!, but Base.permute! must be overloaded
-----------------------------------------------------------------------------"""
function permute_tensor!(A::SSSTensor,p::Array{Int,1})
  @assert length(p) == A.cubical_dimension
  @assert Set(1:length(p)) == Set(p)  #check for a proper permutation

  permuted_edges = Dict{Array{Int,1},AbstractFloat}()
  for (indices,val) in A.edges
    permuted_edges[sort(map(i->p[i],indices))] = val
  end
  A.edges = permuted_edges
end

"""-----------------------------------------------------------------------------
    connected_components(A)

  This function computes the connected components of the undirected hypergraph
associated with the tensor A. The algorithm runs a breadth first search until
all the vertices are found and returns an array with the component number for
each vertex, and an array which stores the visting order.

Input:
------
* A - (SSSTensor)

    The tensor to find the conncected components of.
* v0 - (Int)

    Optional starting vertex of bfs tree, default will be 1.
Output:
-------
* component_assignment - (Array{Int,1})

    The array storing the components each vertex belongs too.
* component_sizes - (Array{Int,1})

    An array which keeps track of the size of each component.
* visting_order - (Array{Int,1})

    The order the vertices were visited in, used for permuting into block
    diagonal structures.
-----------------------------------------------------------------------------"""
function connected_components(A::Ten,v0::Int = 1) where {Ten <: AbstractSSTen}
  n = A.cubical_dimension

  #initialize variables
  component_count = 1
  visited_vertex_count = 1

  unvisited_vertices = Set(1:n)
  q = Queue{Int}()
  enqueue!(q,v0)
  delete!(unvisited_vertices,v0)

  visiting_order = zeros(n)
  component_assignment = zeros(n)
  component_sizes = [0]

  edge_incidence = find_edge_incidence(A)

  while length(q.store) != 0

	v = dequeue!(q)

	#record information about v
	visiting_order[v] = visited_vertex_count
    component_assignment[v] = component_count
	component_sizes[component_count] += 1
	visited_vertex_count += 1

	#add neighbors to queue
	for (indices,_) in get(edge_incidence,v,[])
	  for neighbor in unique(indices)
	    if neighbor in unvisited_vertices
		  enqueue!(q,neighbor)
		  delete!(unvisited_vertices,neighbor)
		end
	  end

	end

	#check for diconnected components
 	if (length(q.store) == 0) && (length(unvisited_vertices) > 0)
 	  #add random unvisited vertex
	  enqueue!(q,pop!(unvisited_vertices))
 	  component_count += 1
	  push!(component_sizes,0)
     end
  end

  return component_assignment, component_sizes, visiting_order
end

"""-----------------------------------------------------------------------------
    get_largest_component(A,filepath)

  Computes the largest connected component (lcc) of the hypergraph. Will either
  compute it directly or load/save it it from/to an lcc file built from the
  filepath. The filepath's postfix is replaced from a .lcc of the formatting

  (largest_component_size::Int)\t(total vertex count::Int)\n
  (vertex_index_1::Int)\n
            ⋮
  (vertex_index_n::Int)\n

Inputs:
-------
* A - (SSSTensor)
  The hypergraph adjacency tensor to find the largest component of.
* filepath - (Optional String)
  The .ssten filepath associated with the tensor, if passed in, then .lcc files
  are created so that the largest connected component need not be computed
  multiple times.

Outputs:
--------
* subtensor - (SSSTensor)

  The subtensor associated with the largest component of the orignal hypergraph.
  Vertices are remapped so that they're indexed from 1:sizeof(lcc)
* lcc_indices - (Array{Int,1})

  Array with the indices of the vertices from the original graph which comprise
  the largest connected component.
-----------------------------------------------------------------------------"""
function get_largest_component(A::Ten,filepath::String="") where {Ten <: AbstractSSTen}

    if !isempty(filepath) #if path is passed in, look for lcc file
    	lcc_file = alterFilename(filepath,".lcc",keep_postfix=false)
    	if isfile(lcc_file)
    		open(lcc_file, "r") do f
				println("opened lcc file")
				lcc_size,full_size =
				    [parse(Int,elem) for elem in split(chomp(readline(f)),'\t')]
				if lcc_size == full_size
    				subtensor = A
					lcc_indices = 1:A.cubical_dimension
				else
					lcc_indices = Array{Int,1}(undef,lcc_size)
					for (i,line) in zip(1:lcc_size,eachline(f))
					  lcc_indices[i] = parse(Int,chomp(line))
					end
					subtensor = get_sub_tensor(A,lcc_indices,true)
				end
				return subtensor,lcc_indices
			end
		end
	end

	comps, comp_sizes, _ = connected_components(A)

	if length(comp_sizes) > 1
		largest_comp = findall(x->x==maximum(comp_sizes),comp_sizes)[1]
		lcc_indices = findall(x->x==largest_comp,comps)
	else
		#graph is connected, return self
		return A, 1:A.cubical_dimension
	end

	if !isempty(filepath) #save lcc file if one doesn't exist
		lcc_file = alterFilename(filepath,".lcc",keep_postfix=false)
		open(lcc_file,"w") do f
			header="$(length(lcc_indices))\t$(A.cubical_dimension)\n"
			write(f,header)
			if lcc_indices == A.cubical_dimension
				return
			end
			for v_i in lcc_indices
				write(f,"$(v_i)\n")
			end
		end
		println("made it")
	end
	get_sub_tensor(A,lcc_indices,true),lcc_indices
end

"""-----------------------------------------------------------------------------
    find_edge_incidence(A)

  This function creates a dictionary linking each vertex to each of the hyper
edges associated with that edge.

 Input:
 ------
* A - (SSSTensor)

    The tensor to find the hyper edge association of.

 Output:
 -------
* edge_incidence - (Dict{Int,Set{Tuple{Array{Int,1},AbstractFloat}},1})

    The dictionary which links all vertices to the hyper edges they're contained
    within.

TODO: May be good to create a version where you can request the incidence for a
particular vertex
-----------------------------------------------------------------------------"""
function find_edge_incidence(A::Ten) where {Ten <: AbstractSSTen}
  edge_incidence = Dict{Int,Set{Tuple{Array{Int,1},AbstractFloat}}}()

  for (indices,val) in A

	prev_v = -1
    for v in indices
	  if prev_v == v
	    continue
	  else
  	    if !haskey(edge_incidence,v)
	      edge_incidence[v] = Set([(indices,val)])
	    else
	      push!(edge_incidence[v],(indices,val))
	    end
	  end
	  prev_v = v
	end
  end

  return edge_incidence
end


"""
add_Y_diag!(Y,Ubar,Vbar,u,v,x,d)

Inplace computation of `Y[i,:] = Vbar'*u + Ubar'*v + d[i]*x`.
"""
function add_Y_diag!(Y,Ubar,Vbar,u,v,x,d)
    @inbounds for k = 1:size(Vbar,2)
        Y[k] = d*x[k]
        @inbounds for i = 1:size(Vbar,1)
            Y[k] += Vbar[i,k]*u[i] + Ubar[i,k]*v[i]
        end
    end
end
"""
add_Y_tril_diag!(Y,w,Ubar,x,d)

Inplace computation of `Y[i,:] = w'*Ubar + ds[i]*x`.
"""
function add_Y_tri_diag!(Y,w,Ubar,x,d)
    @inbounds for k = 1:size(Ubar,2)
        Y[k] = d*x[k]
        @inbounds for i = 1:size(Ubar,1)
            Y[k] += Ubar[i,k]*w[i]
        end
    end
end
"""
add_Y_tri!(Y,Wbar,u)

Inplace computation of `Y[i,:] = Wbar'*u`.
"""
function add_Y_tri!(Y,Wbar,u)
    for k = 1:size(Wbar,2)
        Y[k] = 0.0
        for i = 1:size(Wbar,1)
            Y[k] += Wbar[i,k]*u[i]
        end
    end
end

"""
add_Y!(Y,Ubar,Vbar,u,v)

Inplace computation of `Y[i,:] = Vbar'*u + Ubar'*v`.
"""
function add_Y!(Y,Ubar,Vbar,u,v)
    @inbounds for k = 1:size(Vbar,2)
        Y[k] = 0.0
        @inbounds for i = 1:size(Vbar,1)
            Y[k] += Vbar[i,k]*u[i] + Ubar[i,k]*v[i]
        end
    end
end

"""
add_inner_minus!(Ubar,u,x)

Inplace computation of `Ubar -= u .* x`.
"""
function add_inner_minus!(Ubar,u,x)
    for j=1:size(Ubar,2),i = 1:size(Ubar,1)
        Ubar[i,j] = Ubar[i,j] - u[i]*x[j]
    end
end
"""
add_inner_plus!(Vbar,v,x)

Inplace computation of `Vbar += v .* x`.
"""
function add_inner_plus!(Vbar,v,x)
    for j=1:size(Vbar,2),i = 1:size(Vbar,1)
        Vbar[i,j] = Vbar[i,j] + v[i]*x[j]
    end
end
"""
    add_outer_product!(wTw,w)

Inplace computation of `wTw = wTw + w*w'`
"""
function add_outer_product!(wTw,w)
    for j=1:size(wTw,2),i = 1:size(wTw,1)
        wTw[i,j] = wTw[i,j] + w[i]*w[j]
    end
end
"""
add_product!(Wbar,w,x)

Inplace computation of `Wbar  += w*x'`
"""
function add_product!(Wbar,w,x)
    for j=1:size(Wbar,2),i = 1:size(Wbar,1)
        Wbar[i,j] = Wbar[i,j] + w[i]*x[j]
    end
end


# Copyright 2015 Mitchell Kember. Subject to the MIT License.

using Gadfly
using Base.Test

# Plots the distribution of permutations for shuffling n cards with the custom
# shuffling function, using the given number of bins in the histogram. Returns
# the Gadfly plot object.
function distribution(n, samples, bins, shuffle)
  perms = Array(Int, samples)
  for i = 1:samples
    perms[i] = perm_index(shuffle([0:n-1]))
  end
  plot(
    x=perms, Geom.histogram(bincount=bins),
    yintercept=[samples/bins], Geom.hline(color="yellowgreen"),
    Guide.title(""), Guide.xlabel(""), Guide.ylabel(""),
    Scale.y_continuous(minvalue=0)
  )
end

# Saves the distribution as an SVG file with the given name.
function save_svg(dist, name)
  draw(SVG(name, 6inch, 4inch), dist)
end

# Returns a random number between a and b inclusive.
function next_rand(a, b)
  rand(a:b)
end

# Shuffles xs in-place using a bad algorithm.
function simple_shuffle!(xs)
  n = length(xs)
  for _ = 1:n
    a = next_rand(1, n)
    b = next_rand(1, n)
    xs[a], xs[b] = xs[b], xs[a]
  end
  xs
end

# Shuffles xs in-place using a good algorithm.
function knuth_shuffle!(xs)
  n = length(xs)
  for i = 1:n
    j = next_rand(i, n)
    xs[i], xs[j] = xs[j], xs[i]
  end
  xs
end

# Identifies the permutation of [0:n-1] that xs corresponds to, where n is the
# length of xs. Returns zero for ascending order and n!-1 for desceding order.
# Does not check that xs is a permutation.
function perm_index(xs)
  if isempty(xs)
    return 0
  end
  n = length(xs)
  available = ones(xs)
  index = 0
  f = factorial(n-1)
  for i = 1:n-1
    val = xs[i]
    pos = 0
    for j = 1:val
      pos += available[j]
    end
    index += pos * f
    available[val+1] = 0
    f /= n - i
  end
  index
end

# Verifies permutation indices for permutations of [0] through [0:3].
function test_perm_index()
  order = Any[
    [0],
    [0 1; 1 0],
    [0 1 2; 0 2 1; 1 0 2; 1 2 0; 2 0 1; 2 1 0],
    [0 1 2 3; 0 1 3 2; 0 2 1 3; 0 2 3 1; 0 3 1 2; 0 3 2 1;
     1 0 2 3; 1 0 3 2; 1 2 0 3; 1 2 3 0; 1 3 0 2; 1 3 2 0;
     2 0 1 2; 2 0 2 1; 2 1 0 3; 2 1 3 0; 2 3 0 1; 2 3 1 0;
     3 0 1 2; 3 0 2 1; 3 1 0 2; 3 1 2 0; 3 2 0 1; 3 2 1 0]
  ]
  for ps in order
    for i = 1:size(ps, 1)
      @test perm_index(ps[i,:]) == i - 1
    end
  end
end

# Plots the distributions for the simple shuffle and for the Knuth shuffle, and
# saves each to an SVG file.
function main()
  const n = 20
  const samples = 1_000_000
  const bins = 200
  save_svg(distribution(n, samples, bins, simple_shuffle!), "simple.svg")
  save_svg(distribution(n, samples, bins, knuth_shuffle!), "knuth.svg")
end


pos1 = Real2D(1.0,1.0)
pos2 = Real2D(1.0,1.0)
pos3 = Real2D(10.0,10.0)

@test (pos1 - pos2) == Real2D(0.0,0.0)
@test (pos1 + pos2) == Real2D(2.0,2.0)
@test pos1 == pos2

@test toroidal(Real2D(11.2,11.2),10.0,10.0) == Real2D(1.2,1.2)

@test toroidal(Real2D(11.2,11.2),10.0,10.0) != Real2D(1.3,1.3)

@test toroidal(Real2D(11.2,11.2),3.0,3.0) == Real2D(2.2,2.2)

@test distance(Real2D(1.0,1.0), Real2D(1.0,5.0), 10.0, 10.0, false) == 4.0

@test distance(Real2D(1.0,1.0), Real2D(1.0,5.0), 10.0, 10.0, true) == 4.0

@test distance(Real2D(1.0,1.0), Real2D(1.0,5.0), 4.0, 4.0, true) == 0.0


# This example is used to generate plots in the paper, "Reflection from a multi-species material and its transmitted effective wavenumber." Proc. R. Soc. (2018): 20170864.

# Everything related to ploting has been commented so that this example can run with requiring the Plots package. That way, this example can be checked everytime the package is modified.
using EffectiveWaves

# using LaTeXStrings
# using Plots
# height=500
 # pyplot(linewidth=3, size=(2.6*height,height), border=false)

 # Plots.scalefontsizes(1.8)

filename="fluid"
mediumname = "water"
## choose material

  # fluid
  medium = WaterDistilled
  # inclusion1 = SodiumNitrate # surfactant ?
  inclusion1 = Glycerol # surfactant ?
  inclusion2 = Hexadecane

  # medium = Hexadecane
  # # inclusion1 = SodiumNitrate # surfactant ?
  # inclusion1 = SodiumNitrate # surfactant ?
  # inclusion2 = Glycerol

  ωfactor = 4000;
  ωs = 20*ωfactor.*LinRange(real(medium.c/10000),real(medium.c),200) # k from 0 to 1

  volfrac = 0.22
  r1 = 0.1/ωfactor; vol1 = 0.11;
  r2 = 1.0/ωfactor; vol2 = volfrac - vol1

  sp1 = Specie(ρ=inclusion1.ρ, r=r1, c=inclusion1.c, volfrac = vol1)
  sp2 = Specie(ρ=inclusion2.ρ, r=r2, c=inclusion2.c, volfrac = volfrac-vol1)

  # True wavenumber
  kTs  = wavenumber_low_volfrac(ωs, medium, [sp1,sp2]);

  # Approximate Challis wavenumber
  kTCs = wavenumber_challis(ωs, medium, [sp1,sp2]; hankel_order=5);

  speed_arr = [ ωs./real(kTs), ωs./real(kTCs), 0 .*ωs .+ real(medium.c)]
  atten_arr = imag([kTs,kTCs])

  styles = [:solid :dashdot :dot]
  # labs = [L"Effective $k_{*}$" L"Approximate $k_{*C}$" mediumname];
  # xlabel = L"k a_S";
  ys_arr = speed_arr;
  xs = r1.*ωs ./ real(medium.c);
  # m =5;
  # p1= plot(xs, ys_arr, xlabel=xlabel, ylabel="sound speed (m/s)", labels=labs, line = styles
  #           , ylims = (minimum(minimum.(ys_arr))*0.995, maximum(maximum.(ys_arr))*1.005));

  ys_arr = atten_arr;
  # labs = [L"Effective $k_{*}$" L"Approximate $k_{*C}$"];
  # p2 = plot(xs, ys_arr, labels=labs, xlabel=xlabel, ylabel="attenuation (1/m)", line = styles
  #             , ylims = (minimum(minimum.(ys_arr))*0.995, maximum(maximum.(ys_arr))*1.005));
  # plot(p1,p2)
  # try mkdir("media") end
  # savefig("media/compare_$(filename)_large-w.png")
  # savefig("media/compare_$(filename)_large-w.pdf")
  #
  # Plots.scalefontsizes(1/1.8)


# some noise model (standard deviation)

"""
    LinearNoise

The noise increase linearly with the distance to the object

#Fields
- close_range: distance at which the error starts increasing
- min_noise: minimum error
- increase_rate: noise dimension/meter (e.g. position noise increase by 1m every meter)
"""
@with_kw struct LinearNoise <: AbstractNoiseModel
    close_range::Float64 = 2.0
    min_noise::Float64 = 0.5
    increase_rate::Float64 = 0.1
end

function noise(model::LinearNoise, dist::Float64)
    return max(model.min_noise, model.min_noise + dist*model.increase_rate)
end


using Flux.Losses: logitcrossentropy

# affine
Solaris.Affine(2, 1; isBias = true)
m = Solaris.Affine(2, 1; isBias = false)
m(randn(Float32, 2))

Chain(4, 4, tanh)
Solaris.dense_layer(4, 4; isBias = true)
Solaris.dense_layer(4, 4; isBias = false)

faf = Solaris.FastAffine(4, 4, tanh)
DiffEqFlux.paramlength(faf)
_p = initial_params(faf)
faf(randn(Float32, 4), _p)

nn = Chain(Dense(21, 21, tanh), Dense(21, 21))
sm = Shortcut(nn)
show(sm)
sm(rand(21))

# icnn
icnnl = Convex(4, 4, 1, identity; fw = randn, fb = zeros, precision = Float32)
icnnc = ICNN(4, 1, [10, 10], identity; fw = randn, fb = zeros, precision = Float32)
show(icnnl)
show(icnnc)
icnnl(randn(4))
icnnl(randn(4), randn(4))
icnnc(randn(4))

# fast icnn
fil = FastConvex(4, 4, 4)
fic = FastICNN(4, 1, [10, 10])
fil(rand(4), rand(4), initial_params(fil))
fic(rand(4), initial_params(fic))


using PtFEM

data = Dict(
  # Beam(ndim, nst, nxe, nip, finite_element(nod, nodof), axisymmetric)
  :struc_el => Beam(2, 1, 5, 1, :x, Line(2, 2), false),
  :properties => [
    1.924e4 0.2;
    1.924e4 0.6;
    1.924e4 1.0;
    1.924e4 1.4;
    1.924e4 1.8],
  :etype => [1, 2, 3, 4, 5],
  :x_coords => [0.0, 2.0, 4.0, 6.0, 8.0, 10.0],
  :loaded_nodes => [(1, [1.0 0.0])]
)

data |> display
println()

@time fem, dis_df, fm_df = p43(data)
println()

display(dis_df)
println()
display(fm_df)


using Test, YaoArrayRegister, YaoBlocks, BitBasis

function toffli(b::BitStr)
    t = @inbounds b[1] ⊻ (b[3] & b[2])
    return @inbounds bit_literal(t, b[2], b[3])
end

"""
    pshift(n::Int) -> Function

return a peridoc shift function.
"""
pshift(n::Int) = (b::Int, nbit::Int) -> mod(b+n, 1<<nbit)
pshift(n::Float64) = (b::Float64, nbit::Int) -> mod(b+n, 1)

@testset "test toffli" begin
    g = mathgate(toffli; nbits=3)
    check_truth(b1, b2) = apply!(ArrayReg(b1), g) == ArrayReg(b2)
    @test check_truth(bit"000", bit"000")
    @test check_truth(bit"001", bit"001")
    @test check_truth(bit"010", bit"010")
    @test check_truth(bit"011", bit"011")
    @test check_truth(bit"100", bit"100")
    @test check_truth(bit"101", bit"101")
    @test check_truth(bit"110", bit"111")
    @test check_truth(bit"111", bit"110")
end

@testset "test pshift" begin
    # bint
    nbits = 5
    ab = mathgate(pshift(3); nbits=nbits)
    mb = mathgate(pshift(-3); nbits=nbits)
    @test apply!(zero_state(nbits), ab) == product_state(nbits, 3)
    @test apply!(zero_state(nbits), mb) == product_state(nbits, 1<<nbits-3)
    @test isunitary(ab)

    ab = mathgate(pshift(3); nbits=nbits, bview=bint_r)
    @test apply!(zero_state(nbits), ab) == product_state(nbits, 0b11000)
    @test isunitary(ab)

    af = mathgate(pshift(0.5); nbits=nbits, bview=bfloat)
    @test isunitary(af)
    @test apply!(zero_state(nbits), af) == product_state(nbits, 1)

    # bfloat_r
    af = mathgate(pshift(0.5); nbits=nbits, bview=bfloat_r)
    @test isunitary(af)
    @test apply!(zero_state(nbits), af) == ArrayReg(bit"10000")
end